From a53aeb18eb7a83d02ba9bcf2d025e292dd67100e Mon Sep 17 00:00:00 2001 From: mgwang Date: Wed, 4 Nov 2026 10:52:57 +0800 Subject: [PATCH] [Ascend] feat: add Ascend backend support for Triton 3.6.x --- CMakeLists.txt | 12 + cmake/FlagTreeOptions.cmake | 72 +- python/requirements-ascend.txt | 3 + python/setup_tools/setup_helper.py | 10 +- python/setup_tools/utils/__init__.py | 4 + python/setup_tools/utils/ascend.py | 16 + python/triton/_utils.py | 10 +- python/triton/compiler/compiler.py | 18 +- python/triton/extension/__init__.py | 4 + python/triton/knobs.py | 3 + python/triton/language/semantic.py | 4 +- python/triton/spec/ascend | 1 + third_party/ascend/CMakeLists.txt | 97 + third_party/ascend/Registrar.cc | 13 + third_party/ascend/ascend_ir.cc | 820 +++ third_party/ascend/backend/__init__.py | 54 + .../ascend/backend/backend_register.py | 341 + third_party/ascend/backend/compiler.py | 1157 ++++ third_party/ascend/backend/cpu_driver.py | 203 + third_party/ascend/backend/driver.py | 1058 ++++ .../ascend/backend/lib/libdevice.10.bc | Bin 0 -> 84984 bytes third_party/ascend/backend/name.conf | 1 + third_party/ascend/backend/npu_utils.cpp | 362 ++ .../ascend/backend/runtime/__init__.py | 37 + .../ascend/backend/runtime/autoparser.py | 939 +++ .../ascend/backend/runtime/autotuner.py | 1134 ++++ .../ascend/backend/runtime/tile_generator.py | 489 ++ third_party/ascend/backend/runtime/utils.py | 132 + .../backend/spec/include/flagtree_spec.h | 7 + .../spec/include/runtime/libentry/libentry.h | 77 + .../Dialect/Triton/IR/TritonAttrDefs.td | 180 + .../ascend/backend/spec/lib/CMakeLists.txt | 6 + .../backend/spec/lib/Dialect/CMakeLists.txt | 1 + .../spec/lib/Dialect/Triton/CMakeLists.txt | 1 + .../spec/lib/Dialect/Triton/IR/CMakeLists.txt | 22 + .../spec/lib/Dialect/Triton/IR/Traits.cpp | 276 + .../backend/spec/lib/runtime/CMakeLists.txt | 1 + .../spec/lib/runtime/libentry/CMakeLists.txt | 11 + .../spec/lib/runtime/libentry/libentry.cpp | 115 + third_party/ascend/backend/testing.py | 399 ++ third_party/ascend/backend/utils.py | 716 +++ third_party/ascend/bin/CMakeLists.txt | 218 + .../ascend/bin/RegisterTritonDialects.h | 192 + third_party/ascend/bin/triton-llvm-opt.cpp | 121 + third_party/ascend/bin/triton-lsp.cpp | 10 + third_party/ascend/bin/triton-mlir-opt.cpp | 45 + third_party/ascend/bin/triton-opt.cpp | 15 + third_party/ascend/bin/triton-reduce.cpp | 11 + .../ascend/bin/triton-tensor-layout.cpp | 234 + .../include/AutoBlockify/AutoBlockify.h | 119 + .../include/AutoBlockify/CMakeLists.txt | 3 + .../ascend/include/AutoBlockify/Passes.h | 37 + .../ascend/include/AutoBlockify/Passes.td | 21 + .../ascend/include/AutoBlockify/Utils.h | 65 + third_party/ascend/include/CMakeLists.txt | 11 + .../include/CommonIRToHIVM/CMakeLists.txt | 3 + .../ascend/include/CommonIRToHIVM/Passes.h | 49 + .../ascend/include/CommonIRToHIVM/Passes.td | 66 + .../CommonIRToHIVM/ReduceSumStrengthPass.h | 26 + .../ascend/include/Dialect/CMakeLists.txt | 1 + .../Dialect/TritonAscend/CMakeLists.txt | 1 + .../Dialect/TritonAscend/IR/CMakeLists.txt | 15 + .../TritonAscend/IR/TritonAscendAttrDefs.td | 25 + .../TritonAscend/IR/TritonAscendDialect.h | 34 + .../TritonAscend/IR/TritonAscendDialect.td | 31 + .../TritonAscend/IR/TritonAscendOps.td | 473 ++ .../AddControlFlowCondition.h | 71 + .../AddControlFlowCondition/CloneOps.h | 50 + .../AddControlFlowCondition/CreateIfOps.h | 64 + .../InitDependentMap.h | 49 + .../UpdateConditionInfo.h | 200 + .../AddControlFlowCondition/UpdateForOps.h | 61 + .../UpdateLoopIterTimes.h | 103 + .../AddControlFlowCondition/Utils.h | 54 + .../DynamicCVPipeline/AddDynamicCVPipeline.h | 57 + .../DynamicCVPipeline/AllocMultiCache.h | 51 + .../AddMultiBufferInnerScope.h | 64 + .../AddMultiBufferOuterScope.h | 128 + .../include/DynamicCVPipeline/CMakeLists.txt | 3 + .../Common/BufferCountManager.h | 43 + .../DynamicCVPipeline/Common/FlagIdManager.h | 65 + .../Common/MemoryEffectsTracker.h | 101 + .../include/DynamicCVPipeline/Common/Utils.h | 63 + .../ascend/include/DynamicCVPipeline/Passes.h | 37 + .../include/DynamicCVPipeline/Passes.td | 21 + .../PlanComputeBlock/Common.h | 48 + .../PlanComputeBlock/ComputeBlockIdManager.h | 73 + .../PlanComputeBlock/OpClassifier.h | 177 + .../PlanComputeBlock/Passes.h | 37 + .../PlanComputeBlock/PlanCubeBlockPass.h | 47 + .../PlanComputeBlock/ReorderOpsByBlockId.h | 44 + .../DynamicCVPipeline/PlanComputeBlockPass.h | 52 + .../AddMultiBufferToGMLoadInternal.h | 204 + .../AddMultiBufferToGMLoadPass.h | 70 + .../AddMultiBufferToGMLoadTypes.h | 85 + .../AsyncLoadHoistingPass.h | 50 + .../SeparateMemoryFromComputePass.h | 49 + .../SplitDataflow/AddBlockIdForControlOps.h | 51 + .../SplitDataflow/DataDependencyAnalysis.h | 136 + .../SplitDataflow/InterCoreTransferAndSync.h | 144 + .../SplitDataflow/MarkMainLoop.h | 51 + .../PreserveControlAttrsCanonicalize.h | 68 + .../SplitDataflow/SeparateCVScope.h | 49 + .../DynamicCVPipeline/SplitDataflowPass.h | 50 + .../include/TritonAffinityOpt/CMakeLists.txt | 3 + .../ascend/include/TritonAffinityOpt/DAG.h | 304 + .../ascend/include/TritonAffinityOpt/Passes.h | 47 + .../include/TritonAffinityOpt/Passes.td | 29 + .../include/TritonAffinityOpt/Utils.hpp | 64 + .../include/TritonToAnnotation/CMakeLists.txt | 3 + .../include/TritonToAnnotation/Passes.h | 43 + .../include/TritonToAnnotation/Passes.td | 17 + .../include/TritonToGraph/AliasAnalysis.h | 160 + .../include/TritonToGraph/CMakeLists.txt | 3 + .../include/TritonToGraph/ControlFlowGraph.h | 266 + .../TritonToGraph/ControlFlowGraphBuilder.h | 242 + .../include/TritonToGraph/DataflowGraph.h | 279 + .../include/TritonToGraph/GraphAnalysis.h | 402 ++ .../TritonToGraph/InterProceduralCFG.h | 108 + .../ascend/include/TritonToGraph/MemorySSA.h | 294 + .../include/TritonToGraph/MemorySsaBuilder.h | 175 + .../ascend/include/TritonToGraph/Passes.h | 44 + .../ascend/include/TritonToGraph/Passes.td | 45 + .../ProgramAxisDependenceAnalysis.h | 77 + .../TritonToGraph/ProgramGridTransform.h | 76 + .../ProgramMappingScalarSpecialization.h | 64 + .../ascend/include/TritonToGraph/tensor.h | 166 + .../include/TritonToHFusion/CMakeLists.txt | 3 + .../ascend/include/TritonToHFusion/Passes.h | 43 + .../ascend/include/TritonToHFusion/Passes.td | 17 + .../include/TritonToHIVM/CMakeLists.txt | 3 + .../ascend/include/TritonToHIVM/Passes.h | 43 + .../ascend/include/TritonToHIVM/Passes.td | 17 + .../include/TritonToLLVM/CMakeLists.txt | 3 + .../ascend/include/TritonToLLVM/Passes.h | 26 + .../ascend/include/TritonToLLVM/Passes.td | 17 + third_party/ascend/language/cann/__init__.py | 58 + .../language/cann/extension/__init__.py | 154 + .../ascend/language/cann/extension/_utils.py | 54 + .../ascend/language/cann/extension/aux_ops.py | 142 + .../ascend/language/cann/extension/builder.py | 85 + .../cann/extension/builtin_custom_ops.py | 220 + .../language/cann/extension/code_generator.py | 203 + .../ascend/language/cann/extension/core.py | 478 ++ .../language/cann/extension/custom_op.py | 400 ++ .../language/cann/extension/dispatch.py | 33 + .../language/cann/extension/math_ops.py | 54 + .../ascend/language/cann/extension/mem_ops.py | 521 ++ .../ascend/language/cann/extension/scope.py | 71 + .../language/cann/extension/semantic.py | 179 + .../ascend/language/cann/extension/vec_ops.py | 536 ++ third_party/ascend/language/cann/libdevice.py | 1004 +++ .../ascend/language/kernels/__init__.py | 28 + third_party/ascend/language/kernels/gather.py | 88 + .../ascend/lib/AutoBlockify/AutoBlockify.cpp | 363 ++ .../ascend/lib/AutoBlockify/CMakeLists.txt | 22 + .../lib/AutoBlockify/RewriteOperation.cpp | 512 ++ third_party/ascend/lib/AutoBlockify/Utils.cpp | 211 + third_party/ascend/lib/CMakeLists.txt | 11 + .../ascend/lib/CommonIRToHIVM/CMakeLists.txt | 20 + .../lib/CommonIRToHIVM/CommonIRToHIVM.cpp | 899 +++ .../CommonIRToHIVM/ReduceSumStrengthPass.cpp | 210 + third_party/ascend/lib/Dialect/CMakeLists.txt | 1 + .../lib/Dialect/TritonAscend/CMakeLists.txt | 1 + .../Dialect/TritonAscend/IR/CMakeLists.txt | 14 + .../TritonAscend/IR/TritonAscendAttrs.cpp | 15 + .../TritonAscend/IR/TritonAscendDialect.cpp | 44 + .../TritonAscend/IR/TritonAscendOps.cpp | 232 + .../AddControlFlowCondition.cpp | 159 + .../AddControlFlowCondition/CloneOps.cpp | 412 ++ .../AddControlFlowCondition/CreateIfOps.cpp | 363 ++ .../InitDependentMap.cpp | 295 + .../UpdateConditionInfo.cpp | 1482 +++++ .../AddControlFlowCondition/UpdateForOps.cpp | 456 ++ .../UpdateLoopIterTimes.cpp | 698 ++ .../AddControlFlowCondition/Utils.cpp | 163 + .../AddDynamicCVPipeline.cpp | 82 + .../lib/DynamicCVPipeline/AllocMultiCache.cpp | 71 + .../AddMultiBufferInnerScope.cpp | 1062 ++++ .../AddMultiBufferOuterScope.cpp | 1052 ++++ .../lib/DynamicCVPipeline/CMakeLists.txt | 43 + .../Common/BufferCountManager.cpp | 120 + .../Common/FlagIdManager.cpp | 67 + .../Common/MemoryEffectsTracker.cpp | 461 ++ .../lib/DynamicCVPipeline/Common/Utils.cpp | 67 + .../DynamicCVPipeline/PlanComputeBlock.cpp | 75 + .../PlanComputeBlock/Common.cpp | 70 + .../ComputeBlockIdManager.cpp | 148 + .../PlanComputeBlock/OpClassifier.cpp | 1387 ++++ .../PlanComputeBlock/PlanCubeBlock.cpp | 576 ++ .../PlanComputeBlock/PlanVectorBlockPass.cpp | 435 ++ .../PlanComputeBlock/ReorderOpsByBlockId.cpp | 400 ++ .../AddMultiBufferToGMLoad.cpp | 176 + .../AsyncLoadHoisting.cpp | 373 ++ .../DependencyAnalysis.cpp | 393 ++ .../LoopTransform.cpp | 805 +++ .../MultiBufferLoopBodyEmission.cpp | 900 +++ .../SeparateMemoryFromComputePass.cpp | 73 + .../lib/DynamicCVPipeline/SplitDataflow.cpp | 79 + .../SplitDataflow/AddBlockIdForControlOps.cpp | 83 + .../SplitDataflow/DataDependencyAnalysis.cpp | 489 ++ .../InterCoreTransferAndSync.cpp | 1048 +++ .../SplitDataflow/MarkMainLoop.cpp | 96 + .../PreserveControlAttrsCanonicalize.cpp | 164 + .../SplitDataflow/SeparateCVScope.cpp | 816 +++ .../lib/TritonAffinityOpt/CMakeLists.txt | 19 + .../ascend/lib/TritonAffinityOpt/DAG.cpp | 520 ++ .../lib/TritonAffinityOpt/DAGSSBuffer.cpp | 5601 +++++++++++++++++ .../ascend/lib/TritonAffinityOpt/DAGScope.cpp | 1104 ++++ .../ascend/lib/TritonAffinityOpt/DAGSync.cpp | 1635 +++++ .../lib/TritonToAnnotation/CMakeLists.txt | 15 + .../TritonToAnnotation/TritonToAnnotation.cpp | 76 + .../lib/TritonToGraph/AliasAnalysis.cpp | 248 + .../ascend/lib/TritonToGraph/CMakeLists.txt | 28 + .../lib/TritonToGraph/ControlFlowGraph.cpp | 1140 ++++ .../TritonToGraph/ControlFlowGraphBuilder.cpp | 903 +++ .../lib/TritonToGraph/DataflowGraph.cpp | 305 + .../lib/TritonToGraph/GraphAnalysis.cpp | 627 ++ .../lib/TritonToGraph/InterProceduralCFG.cpp | 351 ++ .../lib/TritonToGraph/MemorySsaBuilder.cpp | 625 ++ .../ascend/lib/TritonToGraph/Passes.cpp | 38 + .../ProgramAxisDependenceAnalysis.cpp | 496 ++ .../TritonToGraph/ProgramGridTransform.cpp | 345 + .../ProgramMappingScalarSpecialization.cpp | 266 + .../ascend/lib/TritonToHFusion/CMakeLists.txt | 14 + .../lib/TritonToHFusion/TritonToHFusion.cpp | 238 + .../ascend/lib/TritonToHIVM/CMakeLists.txt | 26 + .../ascend/lib/TritonToHIVM/TritonToHIVM.cpp | 199 + .../ascend/lib/TritonToLLVM/CMakeLists.txt | 13 + .../ascend/lib/TritonToLLVM/TritonToLLVM.cpp | 269 + third_party/ascend/python/src/ir.cc | 2124 +++++++ third_party/ascend/python/src/ir.h | 118 + third_party/ascend/python/src/main.cc | 67 + .../python/triton/extension/__init__.py | 20 + .../triton/extension/buffer/src/buffer_ir.cc | 183 + third_party/ascend/runtime/__init__.py | 35 + third_party/ascend/runtime/autoparser.py | 939 +++ third_party/ascend/runtime/autotuner.py | 1134 ++++ third_party/ascend/runtime/tile_generator.py | 489 ++ third_party/ascend/runtime/utils.py | 132 + .../spec/triton/extension/buffer/__init__.py | 20 + .../extension/buffer/language/__init__.py | 44 + .../extension/buffer/language/builder.py | 73 + .../triton/extension/buffer/language/core.py | 341 + .../extension/buffer/language/semantic.py | 128 + .../spec/triton/tools/get_ascend_devices.py | 49 + .../Dialect/Triton/IR/TritonAttrDefs.td | 180 + .../spec_cpp/lib/Dialect/Triton/IR/Traits.cpp | 276 + .../reduce-sum-strength-pipeline.mlir | 41 + .../CommonIRToHIVM/reduce-sum-strength.mlir | 178 + third_party/ascend/triton_ascend.cc | 462 ++ third_party/ascend/tutorials/01-vector-add.py | 101 + .../ascend/tutorials/02-fused-softmax.py | 124 + .../tutorials/03-matrix-multiplication.py | 199 + .../ascend/tutorials/04-low-memory-dropout.py | 137 + third_party/ascend/tutorials/05-layer-norm.py | 126 + .../ascend/tutorials/06-fused-attention.py | 351 ++ .../ascend/tutorials/07-extern-functions.py | 89 + .../ascend/tutorials/08-grouped-gemm.py | 281 + .../ascend/tutorials/09-persistent-matmul.py | 334 + third_party/ascend/unittest/CMakeLists.txt | 27 + .../AddMultiCache/Inner-scope.mlir | 573 ++ .../General/AutoBlockify/auto_blockify.mlir | 134 + .../General/DiscreteMaskAccess/atomic.mlir | 201 + .../General/DiscreteMaskAccess/loadstore.mlir | 66 + .../add_controlflow_condition_test.mlir | 169 + .../AddControlFlowCondition/clone-ops.mlir | 129 + .../create-if-ops.mlir | 146 + .../crossCoreDeps_test.mlir | 140 + .../multi_deps_add_condition_test.mlir | 210 + .../no_deps_add_condition_test.mlir | 112 + .../single_deps_add_condition_test.mlir | 138 + .../skip-control-flow.mlir | 117 + .../update-for-ops.mlir | 131 + .../updateLoopIterTimes2_test.mlir | 75 + .../updateLoopIterTimes_test.mlir | 73 + .../General/TritonAscendAllPass/conv1d.mlir | 160 + .../General/TritonToHFusion/fp_to_fp_rtz.mlir | 18 + .../General/TritonToHFusion/histogram.mlir | 9 + .../histogram_expect_fail.mlir | 16 + .../General/TritonToHFusion/mod.mlir | 11 + .../sync_block_op_conversion.mlir | 20 + .../General/TritonToLinalg/atomic_rmw.mlir | 109 + .../General/TritonToLinalg/clamp_fp8.mlir | 84 + .../General/TritonToLinalg/dotScaled_fp8.mlir | 62 + .../TritonToLinalg/hoist_broadcast.mlir | 76 + .../TritonToLinalg/hoist_broadcast2.mlir | 145 + .../General/TritonToLinalg/parse_custom.mlir | 69 + .../General/TritonToLinalg/parse_select.mlir | 37 + .../TritonToLinalg/test_implicit_permute.mlir | 224 + .../test_load_out_of_bound.mlir | 107 + .../TritonToLinalg/test_nd2nz_block_ptr.mlir | 20 + .../TritonToLinalg/test_parse_atomic_rmw.mlir | 41 + .../TritonToLinalg/test_parse_div.mlir | 45 + .../TritonToLinalg/test_parse_fptosi.mlir | 34 + .../test_parse_select_dense_const.mlir | 36 + .../TritonToLinalg/test_parse_sub.mlir | 34 + .../test_while_loop_index_tensor.mlir | 48 + .../TritonToStructured/CmpConverter.mlir | 24 + .../PromotePointerIterArgsPattern.mlir | 73 + .../TritonToStructured/SplatCmpConverter.mlir | 16 + .../General/TritonToStructured/parseCmp.mlir | 117 + .../TritonToStructured/parseConstant.mlir | 28 + .../TritonToStructured/parseMakeRange.mlir | 24 + .../General/TritonToStructured/parseRem.mlir | 81 + .../bubbleupoperation.mlir | 127 + .../TritonToUnstructure/if_simplifier.mlir | 45 + .../TritonToUnstructure/indirect_func.mlir | 141 + .../TritonToUnstructure/indirect_load.mlir | 42 + .../indirect_mem_access.mlir | 129 + .../TritonToUnstructure/indirect_store.mlir | 42 + .../TritonToUnstructure/nested_loop.mlir | 204 + .../General/TritonToUnstructure/splat.mlir | 14 + .../TritonToUnstructure/unstructure_mix.mlir | 82 + .../ascend/unittest/affine_map/affine_map.py | 40 + .../affine_map/affine_map_buffer_type_demo.py | 27 + .../affine_map_complex_expr_demo.py | 36 + .../affine_map_indexing_map_demo.py | 33 + .../affine_map/affine_map_parse_demo.py | 29 + .../unittest/autotune_ut/01-vector-add.py | 86 + .../unittest/autotune_ut/02-fused-softmax.py | 99 + .../unittest/autotune_ut/03-layer-norm.py | 133 + .../unittest/autotune_ut/04-libentry.py | 97 + .../autotune_ut/test_autotune_param_valid.py | 163 + .../unittest/autotune_ut/test_common.py | 103 + .../autotune_ut/test_customized_config.py | 91 + .../autotune_ut/test_do_bench_compat.py | 172 + .../autotune_ut/test_low_dim_axes_parse.py | 173 + .../unittest/autotune_ut/test_mask_parse.py | 162 + .../autotune_ut/test_no_tiling_axis_parse.py | 99 + .../autotune_ut/test_reduction_axes_parse.py | 157 + .../autotune_ut/test_split_axis_parse.py | 180 + .../autotune_ut/test_tiling_axis_parse.py | 123 + third_party/ascend/unittest/conftest.py | 38 + .../unittest/custom_op/builtin_ops_demo.py | 52 + .../unittest/custom_op/custom_op_demo.py | 118 + .../custom_op/custom_op_extra_buffer_demo.py | 115 + .../custom_op_indexing_map_complex_demo.py | 71 + .../custom_op_indexing_map_compose_demo.py | 73 + .../unittest/custom_op/test_gather_load.py | 32 + .../unittest/custom_op/test_index_select.py | 44 + .../unittest/generalization_cases/acc_util.py | 144 + .../unittest/generalization_cases/conftest.py | 33 + .../unittest/generalization_cases/test_abs.py | 157 + .../test_add_mindspore.py | 84 + .../generalization_cases/test_advance.py | 221 + .../unittest/generalization_cases/test_and.py | 161 + .../generalization_cases/test_argmax.py | 362 ++ .../generalization_cases/test_argmin.py | 360 ++ .../test_associative_scan.py | 523 ++ .../generalization_cases/test_atan.py | 62 + .../generalization_cases/test_atomic_add.py | 576 ++ .../generalization_cases/test_atomic_and.py | 562 ++ .../generalization_cases/test_atomic_cas.py | 484 ++ .../generalization_cases/test_atomic_max.py | 258 + .../generalization_cases/test_atomic_min.py | 264 + .../generalization_cases/test_atomic_or.py | 438 ++ .../generalization_cases/test_atomic_xchg.py | 434 ++ .../generalization_cases/test_atomic_xor.py | 441 ++ .../generalization_cases/test_broadcast.py | 324 + .../generalization_cases/test_broadcast_to.py | 327 + .../generalization_cases/test_cast.py | 391 ++ .../generalization_cases/test_cdiv.py | 162 + .../generalization_cases/test_ceil.py | 157 + .../generalization_cases/test_common.py | 343 + .../unittest/generalization_cases/test_cos.py | 157 + .../generalization_cases/test_count_dim0.py | 152 + .../generalization_cases/test_count_dim1.py | 151 + .../generalization_cases/test_cumprod.py | 254 + .../generalization_cases/test_cumsum.py | 253 + .../test_debug_barrier.py | 168 + .../generalization_cases/test_device_print.py | 106 + .../generalization_cases/test_div_rn.py | 154 + .../generalization_cases/test_dot_scaled.py | 459 ++ .../unittest/generalization_cases/test_eq.py | 128 + .../unittest/generalization_cases/test_erf.py | 172 + .../unittest/generalization_cases/test_exp.py | 155 + .../generalization_cases/test_exp2.py | 153 + .../generalization_cases/test_expand_dims.py | 223 + .../generalization_cases/test_fdiv.py | 155 + .../generalization_cases/test_flip.py | 29 + .../generalization_cases/test_full_op.py | 1096 ++++ .../generalization_cases/test_ge_op.py | 158 + .../generalization_cases/test_general_add.py | 438 ++ .../test_general_arange.py | 66 + .../generalization_cases/test_general_cat.py | 61 + .../test_general_clamp.py | 179 + .../generalization_cases/test_general_div.py | 139 + .../test_general_floor.py | 148 + .../test_general_floordiv.py | 163 + .../generalization_cases/test_general_fma.py | 162 + .../test_general_gather.py | 150 + .../test_general_interleave.py | 216 + .../generalization_cases/test_general_join.py | 229 + .../generalization_cases/test_general_log.py | 148 + .../generalization_cases/test_general_log2.py | 153 + .../test_general_maximum.py | 153 + .../test_general_minimum.py | 153 + .../generalization_cases/test_general_mul.py | 155 + .../test_general_ravel.py | 189 + .../test_general_reshape.py | 160 + .../test_general_rsqrt.py | 151 + .../test_general_sigmoid.py | 175 + .../generalization_cases/test_general_sin.py | 173 + .../test_general_softmax.py | 179 + .../test_general_split.py | 224 + .../generalization_cases/test_general_sub.py | 155 + .../test_general_tensor_descriptor.py | 268 + .../generalization_cases/test_general_view.py | 159 + .../generalization_cases/test_gt_op.py | 158 + .../generalization_cases/test_invalid_fp8.py | 55 + .../generalization_cases/test_invert.py | 166 + .../generalization_cases/test_le_op.py | 158 + .../generalization_cases/test_load_store.py | 168 + .../generalization_cases/test_log1p.py | 62 + .../test_logical_and_op.py | 157 + .../test_logical_or_op.py | 156 + .../generalization_cases/test_lshift_op.py | 183 + .../generalization_cases/test_lt_op.py | 158 + .../test_make_blkptr_matmul.py | 74 + .../test_make_block_ptr.py | 212 + .../generalization_cases/test_matmul.py | 186 + .../test_matmul_mindspore.py | 189 + .../unittest/generalization_cases/test_max.py | 313 + .../unittest/generalization_cases/test_min.py | 313 + .../unittest/generalization_cases/test_mod.py | 227 + .../test_multi-process.py | 106 + .../unittest/generalization_cases/test_ne.py | 128 + .../unittest/generalization_cases/test_neg.py | 158 + .../unittest/generalization_cases/test_not.py | 173 + .../unittest/generalization_cases/test_or.py | 159 + .../test_permute_1d_2d.py | 95 + .../generalization_cases/test_permute_3d.py | 101 + .../test_permute_4d_5d.py | 223 + .../generalization_cases/test_rand.py | 304 + .../generalization_cases/test_range.py | 187 + .../generalization_cases/test_reduce.py | 338 + .../generalization_cases/test_relu.py | 66 + .../generalization_cases/test_rshift_op.py | 184 + .../test_scalar_tensor.py | 82 + .../test_softmax_mindspore.py | 99 + .../generalization_cases/test_sort.py | 229 + .../generalization_cases/test_sqrt.py | 174 + .../generalization_cases/test_sqrt_rn.py | 174 + .../test_static_print_and_assert_op.py | 73 + .../unittest/generalization_cases/test_sum.py | 355 ++ .../generalization_cases/test_sum_dim0.py | 74 + .../generalization_cases/test_sum_dim1.py | 66 + .../generalization_cases/test_swizzle2d.py | 72 + .../unittest/generalization_cases/test_tan.py | 62 + .../generalization_cases/test_trans_1d_2d.py | 114 + .../generalization_cases/test_trans_3d.py | 144 + .../generalization_cases/test_trans_4d_5d.py | 223 + .../generalization_cases/test_umulhi.py | 147 + .../generalization_cases/test_where.py | 135 + .../unittest/generalization_cases/test_xor.py | 174 + .../generalization_cases/test_xorsum.py | 281 + .../generalization_cases/test_zeros_op.py | 534 ++ .../generalization_cases/test_zeroslike.py | 170 + third_party/ascend/unittest/kernels/README.md | 65 + .../ascend/unittest/kernels/common_kernel.py | 7 + .../ascend/unittest/kernels/test_common.py | 119 + .../unittest/kernels/test_triton_kernel.py | 73 + .../unittest/kernels/vllm/_triton_rope.py | 112 + ...ated_delta_rule_fwd_kernel_h_blockdim64.py | 162 + .../vllm/chunk_local_cumsum_scalar_kernel.py | 59 + .../vllm/chunk_scaled_dot_kkt_fwd_kernel.py | 67 + .../unittest/kernels/vllm/expand_kernel.py | 33 + .../kernels/vllm/fused_gdn_gating_kernel.py | 47 + .../fused_qkvzba_split_reshape_cat_kernel.py | 50 + .../kernels/vllm/l2norm_fwd_kernel2_loop.py | 19 + .../kernels/vllm/layer_norm_fwd_kernel.py | 79 + .../merge_16x16_to_32x32_inverse_kernel.py | 66 + .../merge_16x16_to_64x64_inverse_kernel.py | 133 + .../kernels/vllm/recompute_w_u_fwd_kernel.py | 62 + ...jection_greedy_sample_spec_len_1_triton.py | 42 + .../vllm/rejection_random_sample_kernel.py | 55 + .../vllm/sample_recovered_tokens_kernel.py | 77 + .../kernels/vllm/solve_tril_16x16_kernel.py | 108 + third_party/ascend/unittest/lit.cfg.py | 64 + .../ascend/unittest/lit.site.cfg.py.in | 23 + .../pytest_ut/attn_cp_triton_kernel_3d.py | 182 + .../pytest_ut/attn_cp_triton_kernel_3d_la.py | 100 + .../ascend/unittest/pytest_ut/conftest.py | 48 + .../unittest/pytest_ut/test_01_vector_add.py | 70 + .../pytest_ut/test_02_fused_softmax.py | 121 + .../test_03_matrix_multiplication.py | 211 + .../pytest_ut/test_04_low_memory_dropout.py | 127 + .../unittest/pytest_ut/test_05_layer_norm.py | 119 + .../pytest_ut/test_06_fused_attention.py | 340 + .../pytest_ut/test_07_extern_functions.py | 85 + .../pytest_ut/test_08_grouped_gemm.py | 287 + .../pytest_ut/test_09_persistent_matmul.py | 325 + .../pytest_ut/test_10_gather_sorted.py | 201 + .../unittest/pytest_ut/test_11_rab_time.py | 351 ++ .../pytest_ut/test_12_hstu_attention.py | 933 +++ ...test_13_matrix_multiplication_optimized.py | 216 + .../pytest_ut/test_14_accuracy_comparison.py | 144 + .../pytest_ut/test_15_demo_autotune.py | 88 + .../unittest/pytest_ut/test_16_profiler.py | 116 + .../pytest_ut/test_17_demo_libentry.py | 131 + .../unittest/pytest_ut/test_18_gather.py | 139 + .../unittest/pytest_ut/test_2d_permute.py | 54 + .../ascend/unittest/pytest_ut/test_3Dgrid.py | 91 + .../ascend/unittest/pytest_ut/test_abs.py | 69 + .../ascend/unittest/pytest_ut/test_abs_2.py | 62 + .../ascend/unittest/pytest_ut/test_acos.py | 87 + .../ascend/unittest/pytest_ut/test_acosh.py | 73 + .../ascend/unittest/pytest_ut/test_add.py | 91 + .../unittest/pytest_ut/test_address_check.py | 68 + .../ascend/unittest/pytest_ut/test_advance.py | 182 + .../unittest/pytest_ut/test_advance_ptr.py | 51 + .../pytest_ut/test_affine_map_binding.py | 114 + .../ascend/unittest/pytest_ut/test_alloc.py | 96 + .../ascend/unittest/pytest_ut/test_and.py | 61 + .../unittest/pytest_ut/test_annotations.py | 71 + .../ascend/unittest/pytest_ut/test_arange.py | 113 + .../ascend/unittest/pytest_ut/test_arch.py | 94 + .../ascend/unittest/pytest_ut/test_argmax.py | 65 + .../ascend/unittest/pytest_ut/test_argmin.py | 65 + .../unittest/pytest_ut/test_ascend_barrier.py | 75 + .../ascend/unittest/pytest_ut/test_asin.py | 90 + .../ascend/unittest/pytest_ut/test_asinh.py | 58 + .../ascend/unittest/pytest_ut/test_asm.py | 96 + .../unittest/pytest_ut/test_asm_scalar.py | 32 + .../pytest_ut/test_associative_scan.py | 179 + .../test_associative_scan_multi_input.py | 157 + .../ascend/unittest/pytest_ut/test_assume.py | 51 + .../ascend/unittest/pytest_ut/test_assume1.py | 32 + .../ascend/unittest/pytest_ut/test_atan.py | 101 + .../ascend/unittest/pytest_ut/test_atan2.py | 129 + .../ascend/unittest/pytest_ut/test_atanh.py | 77 + .../unittest/pytest_ut/test_atomic_add.py | 207 + .../unittest/pytest_ut/test_atomic_and.py | 68 + .../unittest/pytest_ut/test_atomic_cas.py | 190 + .../unittest/pytest_ut/test_atomic_max.py | 110 + .../unittest/pytest_ut/test_atomic_min.py | 94 + .../unittest/pytest_ut/test_atomic_or.py | 67 + .../pytest_ut/test_atomic_rmw_useanalysis.py | 70 + .../unittest/pytest_ut/test_atomic_xchg.py | 68 + .../unittest/pytest_ut/test_atomic_xor.py | 73 + .../ascend/unittest/pytest_ut/test_attn_cp.py | 500 ++ .../unittest/pytest_ut/test_bind_buffer.py | 67 + .../unittest/pytest_ut/test_block_ptr.py | 195 + .../unittest/pytest_ut/test_boundary_check.py | 275 + .../unittest/pytest_ut/test_broadcast.py | 57 + .../unittest/pytest_ut/test_broadcast_op.py | 54 + .../pytest_ut/test_cannonicalize_tl_where.py | 124 + .../unittest/pytest_ut/test_cast_full.py | 229 + .../ascend/unittest/pytest_ut/test_cat.py | 61 + .../ascend/unittest/pytest_ut/test_cat_dim.py | 119 + .../unittest/pytest_ut/test_cat_help_func.py | 728 +++ .../ascend/unittest/pytest_ut/test_cdiv.py | 113 + .../ascend/unittest/pytest_ut/test_ceil.py | 76 + .../unittest/pytest_ut/test_celoss_indices.py | 112 + .../ascend/unittest/pytest_ut/test_clamp.py | 62 + .../ascend/unittest/pytest_ut/test_common.py | 223 + .../unittest/pytest_ut/test_compile_hint.py | 58 + .../unittest/pytest_ut/test_complex_mask.py | 69 + .../ascend/unittest/pytest_ut/test_copy.py | 97 + .../unittest/pytest_ut/test_copysign.py | 108 + .../ascend/unittest/pytest_ut/test_cos.py | 99 + .../ascend/unittest/pytest_ut/test_cos_2.py | 60 + .../ascend/unittest/pytest_ut/test_cosh.py | 59 + .../ascend/unittest/pytest_ut/test_cumprod.py | 97 + .../ascend/unittest/pytest_ut/test_cumsum.py | 84 + .../ascend/unittest/pytest_ut/test_custom.py | 299 + .../unittest/pytest_ut/test_cyl_bessel_i0.py | 61 + .../unittest/pytest_ut/test_debug_barrier.py | 59 + .../pytest_ut/test_debug_triton_opt_cmd.py | 115 + .../unittest/pytest_ut/test_device_print.py | 158 + .../test_device_print_comprehensive.py | 122 + .../pytest_ut/test_device_print_script.py | 117 + .../pytest_ut/test_discrete_mask_atomic.py | 48 + .../pytest_ut/test_discrete_mask_loadstore.py | 408 ++ .../test_discrete_mask_tail_block_mte_oob.py | 238 + .../pytest_ut/test_discrete_overlap_mask.py | 235 + .../ascend/unittest/pytest_ut/test_div.py | 63 + .../ascend/unittest/pytest_ut/test_dot.py | 181 + .../unittest/pytest_ut/test_dot_scaled.py | 224 + .../unittest/pytest_ut/test_downgrade.py | 49 + .../pytest_ut/test_elementwise_ceil.py | 86 + .../pytest_ut/test_elementwise_clip.py | 88 + .../pytest_ut/test_elementwise_f2i.py | 137 + .../pytest_ut/test_elementwise_floor.py | 86 + .../pytest_ut/test_elementwise_i2f.py | 123 + .../pytest_ut/test_elementwise_round.py | 92 + .../ascend/unittest/pytest_ut/test_eq.py | 79 + .../ascend/unittest/pytest_ut/test_eq_2.py | 76 + .../ascend/unittest/pytest_ut/test_erfinv.py | 156 + .../ascend/unittest/pytest_ut/test_exp.py | 56 + .../ascend/unittest/pytest_ut/test_exp2.py | 59 + .../ascend/unittest/pytest_ut/test_exp_.py | 99 + .../unittest/pytest_ut/test_expand_dims.py | 74 + .../ascend/unittest/pytest_ut/test_expm1.py | 62 + .../unittest/pytest_ut/test_extract_slice.py | 60 + .../unittest/pytest_ut/test_fast_dividef.py | 59 + .../unittest/pytest_ut/test_fast_expf.py | 56 + .../ascend/unittest/pytest_ut/test_fdiv.py | 63 + .../ascend/unittest/pytest_ut/test_fixpipe.py | 84 + .../pytest_ut/test_fixpipe_tensor_dst.py | 169 + .../ascend/unittest/pytest_ut/test_flip.py | 28 + .../ascend/unittest/pytest_ut/test_floor.py | 59 + .../unittest/pytest_ut/test_floordiv.py | 72 + .../ascend/unittest/pytest_ut/test_for_ptr.py | 116 + .../ascend/unittest/pytest_ut/test_full.py | 95 + .../unittest/pytest_ut/test_fusedattention.py | 354 ++ .../ascend/unittest/pytest_ut/test_gamma.py | 86 + .../ascend/unittest/pytest_ut/test_gather.py | 73 + .../unittest/pytest_ut/test_gather_simd.py | 46 + .../ascend/unittest/pytest_ut/test_ge.py | 79 + .../ascend/unittest/pytest_ut/test_ge_2.py | 63 + .../ascend/unittest/pytest_ut/test_gelu.py | 99 + .../ascend/unittest/pytest_ut/test_gt.py | 61 + .../unittest/pytest_ut/test_hd_permute.py | 53 + .../unittest/pytest_ut/test_histogram.py | 69 + .../unittest/pytest_ut/test_hoistbroadcast.py | 220 + .../ascend/unittest/pytest_ut/test_hypot.py | 61 + .../unittest/pytest_ut/test_if_advance.py | 39 + .../ascend/unittest/pytest_ut/test_if_load.py | 82 + .../unittest/pytest_ut/test_if_tensor.py | 134 + .../pytest_ut/test_implicit_atomic.py | 140 + .../pytest_ut/test_implicit_permute.py | 1338 ++++ .../unittest/pytest_ut/test_index_select.py | 268 + .../pytest_ut/test_index_select_inductor.py | 91 + .../pytest_ut/test_indirect_atomic_add.py | 520 ++ .../pytest_ut/test_indirect_atomic_and.py | 517 ++ .../pytest_ut/test_indirect_atomic_cas.py | 407 ++ .../pytest_ut/test_indirect_atomic_max.py | 520 ++ .../pytest_ut/test_indirect_atomic_min.py | 520 ++ .../pytest_ut/test_indirect_atomic_or.py | 517 ++ .../pytest_ut/test_indirect_atomic_xchg.py | 520 ++ .../pytest_ut/test_indirect_atomic_xor.py | 517 ++ .../test_indirect_scalar_load_offset.py | 98 + .../unittest/pytest_ut/test_insert_slice.py | 64 + .../unittest/pytest_ut/test_interleave.py | 75 + .../pytest_ut/test_interleave_optimizaiton.py | 125 + .../ascend/unittest/pytest_ut/test_invert.py | 63 + .../unittest/pytest_ut/test_isfinited.py | 134 + .../ascend/unittest/pytest_ut/test_isnan.py | 107 + .../ascend/unittest/pytest_ut/test_join.py | 66 + .../ascend/unittest/pytest_ut/test_lanzcos.py | 246 + .../test_launcher_empty_signature.py | 37 + .../unittest/pytest_ut/test_layernorm.py | 126 + .../ascend/unittest/pytest_ut/test_ldst.py | 449 ++ .../ascend/unittest/pytest_ut/test_le.py | 79 + .../ascend/unittest/pytest_ut/test_lgamma.py | 114 + .../unittest/pytest_ut/test_linearize.py | 602 ++ .../pytest_ut/test_linearize_fallback.py | 86 + .../unittest/pytest_ut/test_linearize_mask.py | 168 + .../pytest_ut/test_linearize_mask_fallback.py | 129 + .../pytest_ut/test_linearize_permute.py | 517 ++ .../ascend/unittest/pytest_ut/test_load.py | 117 + .../unittest/pytest_ut/test_load_store.py | 151 + .../ascend/unittest/pytest_ut/test_log.py | 99 + .../ascend/unittest/pytest_ut/test_log10.py | 90 + .../ascend/unittest/pytest_ut/test_log1p.py | 60 + .../ascend/unittest/pytest_ut/test_log2.py | 58 + .../ascend/unittest/pytest_ut/test_log_2.py | 59 + .../unittest/pytest_ut/test_logical_and.py | 77 + .../unittest/pytest_ut/test_logical_or.py | 77 + .../ascend/unittest/pytest_ut/test_lshift.py | 99 + .../ascend/unittest/pytest_ut/test_lt.py | 60 + .../test_makeblockptr_negative_padding.py | 136 + .../pytest_ut/test_makeblockptr_permute.py | 115 + .../unittest/pytest_ut/test_max_constancy.py | 27 + .../unittest/pytest_ut/test_max_contiguous.py | 26 + .../unittest/pytest_ut/test_max_dim0.py | 130 + .../unittest/pytest_ut/test_max_dim1.py | 137 + .../pytest_ut/test_max_propagate_nan.py | 20 + .../unittest/pytest_ut/test_max_vector.py | 89 + .../ascend/unittest/pytest_ut/test_maximum.py | 61 + .../unittest/pytest_ut/test_mean_dim0.py | 145 + .../unittest/pytest_ut/test_mean_dim1.py | 145 + .../unittest/pytest_ut/test_mean_vector.py | 105 + .../unittest/pytest_ut/test_min_dim0.py | 152 + .../unittest/pytest_ut/test_min_dim1.py | 135 + .../unittest/pytest_ut/test_min_vector.py | 88 + .../ascend/unittest/pytest_ut/test_minimum.py | 61 + .../ascend/unittest/pytest_ut/test_mod.py | 80 + .../ascend/unittest/pytest_ut/test_mul.py | 61 + .../unittest/pytest_ut/test_mul_reduce.py | 55 + .../unittest/pytest_ut/test_multi_return.py | 599 ++ .../unittest/pytest_ut/test_multibuffer.py | 73 + .../unittest/pytest_ut/test_multiple_of.py | 26 + .../ascend/unittest/pytest_ut/test_ne.py | 76 + .../unittest/pytest_ut/test_nearbyint.py | 69 + .../ascend/unittest/pytest_ut/test_nearest.py | 116 + .../ascend/unittest/pytest_ut/test_neg.py | 58 + .../pytest_ut/test_negative_mask_dim.py | 48 + .../unittest/pytest_ut/test_nextafter.py | 63 + .../ascend/unittest/pytest_ut/test_not.py | 63 + .../unittest/pytest_ut/test_npu_indexing.py | 143 + .../unittest/pytest_ut/test_npu_indexing2.py | 92 + .../ascend/unittest/pytest_ut/test_or.py | 60 + .../pytest_ut/test_paged_kvcache_krope.py | 72 + .../unittest/pytest_ut/test_parallel.py | 100 + .../ascend/unittest/pytest_ut/test_permute.py | 127 + .../unittest/pytest_ut/test_permute_full.py | 171 + .../pytest_ut/test_permute_reshape.py | 91 + .../pytest_ut/test_permuted_boundary_check.py | 38 + .../unittest/pytest_ut/test_pointer_type.py | 49 + .../ascend/unittest/pytest_ut/test_pow.py | 121 + .../unittest/pytest_ut/test_precise_div.py | 60 + .../unittest/pytest_ut/test_precise_sqrt.py | 55 + .../ascend/unittest/pytest_ut/test_ptr_add.py | 64 + .../ascend/unittest/pytest_ut/test_rand.py | 88 + .../ascend/unittest/pytest_ut/test_range.py | 89 + .../ascend/unittest/pytest_ut/test_ravel.py | 71 + .../pytest_ut/test_reduce_count_vector.py | 151 + .../unittest/pytest_ut/test_reduce_maximum.py | 308 + .../unittest/pytest_ut/test_reduce_mean.py | 89 + ...reduce_min_4_keepdim_True_with_index_op.py | 96 + .../unittest/pytest_ut/test_reduce_minimum.py | 308 + .../unittest/pytest_ut/test_reduce_sum.py | 90 + .../ascend/unittest/pytest_ut/test_relu.py | 61 + .../ascend/unittest/pytest_ut/test_reshape.py | 70 + .../pytest_ut/test_resize_performance.py | 102 + .../ascend/unittest/pytest_ut/test_rint.py | 65 + .../unittest/pytest_ut/test_rms_norm.py | 142 + .../pytest_ut/test_rotary_embedding.py | 175 + .../unittest/pytest_ut/test_rotatry_gpt.py | 186 + .../pytest_ut/test_rotaty_embedding_gpt.py | 174 + .../ascend/unittest/pytest_ut/test_rshift.py | 99 + .../ascend/unittest/pytest_ut/test_rsqrt.py | 63 + .../unittest/pytest_ut/test_runtime_utils.py | 33 + .../unittest/pytest_ut/test_scalarPointer.py | 178 + .../unittest/pytest_ut/test_scalar_calc.py | 694 ++ .../ascend/unittest/pytest_ut/test_scope.py | 159 + .../unittest/pytest_ut/test_scope_buffer.py | 97 + .../pytest_ut/test_select_analysis.py | 438 ++ .../test_select_analysis_for_divsiop.py | 84 + .../test_select_analysis_for_invert.py | 153 + .../ascend/unittest/pytest_ut/test_sigmoid.py | 60 + .../ascend/unittest/pytest_ut/test_signbit.py | 81 + .../ascend/unittest/pytest_ut/test_silu.py | 99 + .../pytest_ut/test_simplify_iterargs.py | 78 + .../ascend/unittest/pytest_ut/test_sin.py | 99 + .../ascend/unittest/pytest_ut/test_sinh.py | 59 + .../unittest/pytest_ut/test_sink_broadcast.py | 89 + .../ascend/unittest/pytest_ut/test_softmax.py | 130 + .../ascend/unittest/pytest_ut/test_sort.py | 102 + .../ascend/unittest/pytest_ut/test_split.py | 72 + .../ascend/unittest/pytest_ut/test_sqrt.py | 99 + .../pytest_ut/test_static_print_and_assert.py | 187 + .../unittest/pytest_ut/test_store_scalar.py | 46 + .../unittest/pytest_ut/test_stride0_load.py | 66 + .../ascend/unittest/pytest_ut/test_strides.py | 140 + .../ascend/unittest/pytest_ut/test_sub.py | 62 + .../unittest/pytest_ut/test_sub_vec_id.py | 132 + .../unittest/pytest_ut/test_sub_vec_num.py | 130 + .../ascend/unittest/pytest_ut/test_subview.py | 89 + .../ascend/unittest/pytest_ut/test_sum.py | 124 + .../unittest/pytest_ut/test_sum_dim0.py | 133 + .../unittest/pytest_ut/test_sum_dim1.py | 133 + .../unittest/pytest_ut/test_sum_vector.py | 82 + .../ascend/unittest/pytest_ut/test_swap.py | 58 + .../ascend/unittest/pytest_ut/test_swiglu.py | 100 + .../unittest/pytest_ut/test_swizzle2d.py | 55 + .../unittest/pytest_ut/test_sync_block.py | 99 + .../unittest/pytest_ut/test_sync_block_all.py | 54 + .../ascend/unittest/pytest_ut/test_tan.py | 102 + .../ascend/unittest/pytest_ut/test_tanh.py | 59 + .../unittest/pytest_ut/test_template.py | 58 + .../pytest_ut/test_tensor_descriptor.py | 225 + .../unittest/pytest_ut/test_to_buffer.py | 54 + .../unittest/pytest_ut/test_to_tensor.py | 75 + .../pytest_ut/test_top2gating_argmax.py | 97 + .../ascend/unittest/pytest_ut/test_topk.py | 74 + .../unittest/pytest_ut/test_trans_3d.py | 80 + .../unittest/pytest_ut/test_triton_eq.py | 60 + .../unittest/pytest_ut/test_triton_le.py | 60 + .../unittest/pytest_ut/test_triton_lt.py | 60 + .../unittest/pytest_ut/test_triton_neq.py | 60 + .../test_triton_unified_attention.py | 469 ++ .../ascend/unittest/pytest_ut/test_trunc.py | 68 + .../ascend/unittest/pytest_ut/test_umulhi.py | 62 + .../pytest_ut/test_unlign_max_with_index.py | 115 + .../unittest/pytest_ut/test_unlign_sum.py | 64 + .../pytest_ut/test_unused_func_arg.py | 82 + .../unittest/pytest_ut/test_use_analysis.py | 74 + .../ascend/unittest/pytest_ut/test_utils.py | 48 + .../ascend/unittest/pytest_ut/test_view.py | 63 + .../unittest/pytest_ut/test_where_lt.py | 60 + .../unittest/pytest_ut/test_where_mask.py | 60 + .../unittest/pytest_ut/test_where_var.py | 62 + .../ascend/unittest/pytest_ut/test_xor.py | 97 + .../ascend/unittest/pytest_ut/test_xor_sum.py | 78 + .../ascend/unittest/pytest_ut/test_zeros.py | 88 + .../unittest/pytest_ut/test_zeroslike.py | 63 + third_party/ascend/unittest/run_prtest.sh | 170 + 791 files changed, 138071 insertions(+), 9 deletions(-) create mode 100644 python/requirements-ascend.txt create mode 100644 python/setup_tools/utils/ascend.py create mode 100644 python/triton/extension/__init__.py create mode 120000 python/triton/spec/ascend create mode 100644 third_party/ascend/CMakeLists.txt create mode 100644 third_party/ascend/Registrar.cc create mode 100644 third_party/ascend/ascend_ir.cc create mode 100644 third_party/ascend/backend/__init__.py create mode 100644 third_party/ascend/backend/backend_register.py create mode 100644 third_party/ascend/backend/compiler.py create mode 100644 third_party/ascend/backend/cpu_driver.py create mode 100644 third_party/ascend/backend/driver.py create mode 100644 third_party/ascend/backend/lib/libdevice.10.bc create mode 100644 third_party/ascend/backend/name.conf create mode 100644 third_party/ascend/backend/npu_utils.cpp create mode 100644 third_party/ascend/backend/runtime/__init__.py create mode 100644 third_party/ascend/backend/runtime/autoparser.py create mode 100644 third_party/ascend/backend/runtime/autotuner.py create mode 100644 third_party/ascend/backend/runtime/tile_generator.py create mode 100644 third_party/ascend/backend/runtime/utils.py create mode 100644 third_party/ascend/backend/spec/include/flagtree_spec.h create mode 100644 third_party/ascend/backend/spec/include/runtime/libentry/libentry.h create mode 100644 third_party/ascend/backend/spec/include/triton/Dialect/Triton/IR/TritonAttrDefs.td create mode 100644 third_party/ascend/backend/spec/lib/CMakeLists.txt create mode 100644 third_party/ascend/backend/spec/lib/Dialect/CMakeLists.txt create mode 100644 third_party/ascend/backend/spec/lib/Dialect/Triton/CMakeLists.txt create mode 100644 third_party/ascend/backend/spec/lib/Dialect/Triton/IR/CMakeLists.txt create mode 100644 third_party/ascend/backend/spec/lib/Dialect/Triton/IR/Traits.cpp create mode 100644 third_party/ascend/backend/spec/lib/runtime/CMakeLists.txt create mode 100644 third_party/ascend/backend/spec/lib/runtime/libentry/CMakeLists.txt create mode 100644 third_party/ascend/backend/spec/lib/runtime/libentry/libentry.cpp create mode 100644 third_party/ascend/backend/testing.py create mode 100644 third_party/ascend/backend/utils.py create mode 100644 third_party/ascend/bin/CMakeLists.txt create mode 100644 third_party/ascend/bin/RegisterTritonDialects.h create mode 100644 third_party/ascend/bin/triton-llvm-opt.cpp create mode 100644 third_party/ascend/bin/triton-lsp.cpp create mode 100644 third_party/ascend/bin/triton-mlir-opt.cpp create mode 100644 third_party/ascend/bin/triton-opt.cpp create mode 100644 third_party/ascend/bin/triton-reduce.cpp create mode 100644 third_party/ascend/bin/triton-tensor-layout.cpp create mode 100644 third_party/ascend/include/AutoBlockify/AutoBlockify.h create mode 100644 third_party/ascend/include/AutoBlockify/CMakeLists.txt create mode 100644 third_party/ascend/include/AutoBlockify/Passes.h create mode 100644 third_party/ascend/include/AutoBlockify/Passes.td create mode 100644 third_party/ascend/include/AutoBlockify/Utils.h create mode 100644 third_party/ascend/include/CMakeLists.txt create mode 100644 third_party/ascend/include/CommonIRToHIVM/CMakeLists.txt create mode 100644 third_party/ascend/include/CommonIRToHIVM/Passes.h create mode 100644 third_party/ascend/include/CommonIRToHIVM/Passes.td create mode 100644 third_party/ascend/include/CommonIRToHIVM/ReduceSumStrengthPass.h create mode 100644 third_party/ascend/include/Dialect/CMakeLists.txt create mode 100644 third_party/ascend/include/Dialect/TritonAscend/CMakeLists.txt create mode 100644 third_party/ascend/include/Dialect/TritonAscend/IR/CMakeLists.txt create mode 100644 third_party/ascend/include/Dialect/TritonAscend/IR/TritonAscendAttrDefs.td create mode 100644 third_party/ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.h create mode 100644 third_party/ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.td create mode 100644 third_party/ascend/include/Dialect/TritonAscend/IR/TritonAscendOps.td create mode 100644 third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/CloneOps.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/InitDependentMap.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/Utils.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/AddDynamicCVPipeline.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/AllocMultiCache.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/AllocMultiCache/AddMultiBufferOuterScope.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/CMakeLists.txt create mode 100644 third_party/ascend/include/DynamicCVPipeline/Common/BufferCountManager.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/Common/FlagIdManager.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/Common/MemoryEffectsTracker.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/Common/Utils.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/Passes.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/Passes.td create mode 100644 third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/OpClassifier.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/Passes.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlockPass.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ReorderOpsByBlockId.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/PlanComputeBlockPass.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadInternal.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadPass.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadTypes.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AsyncLoadHoistingPass.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromComputePass.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/SplitDataflow/AddBlockIdForControlOps.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/SplitDataflow/InterCoreTransferAndSync.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/SplitDataflow/MarkMainLoop.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/SplitDataflow/PreserveControlAttrsCanonicalize.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/SplitDataflow/SeparateCVScope.h create mode 100644 third_party/ascend/include/DynamicCVPipeline/SplitDataflowPass.h create mode 100644 third_party/ascend/include/TritonAffinityOpt/CMakeLists.txt create mode 100644 third_party/ascend/include/TritonAffinityOpt/DAG.h create mode 100644 third_party/ascend/include/TritonAffinityOpt/Passes.h create mode 100644 third_party/ascend/include/TritonAffinityOpt/Passes.td create mode 100644 third_party/ascend/include/TritonAffinityOpt/Utils.hpp create mode 100644 third_party/ascend/include/TritonToAnnotation/CMakeLists.txt create mode 100644 third_party/ascend/include/TritonToAnnotation/Passes.h create mode 100644 third_party/ascend/include/TritonToAnnotation/Passes.td create mode 100644 third_party/ascend/include/TritonToGraph/AliasAnalysis.h create mode 100644 third_party/ascend/include/TritonToGraph/CMakeLists.txt create mode 100644 third_party/ascend/include/TritonToGraph/ControlFlowGraph.h create mode 100644 third_party/ascend/include/TritonToGraph/ControlFlowGraphBuilder.h create mode 100644 third_party/ascend/include/TritonToGraph/DataflowGraph.h create mode 100644 third_party/ascend/include/TritonToGraph/GraphAnalysis.h create mode 100644 third_party/ascend/include/TritonToGraph/InterProceduralCFG.h create mode 100644 third_party/ascend/include/TritonToGraph/MemorySSA.h create mode 100644 third_party/ascend/include/TritonToGraph/MemorySsaBuilder.h create mode 100644 third_party/ascend/include/TritonToGraph/Passes.h create mode 100644 third_party/ascend/include/TritonToGraph/Passes.td create mode 100644 third_party/ascend/include/TritonToGraph/ProgramAxisDependenceAnalysis.h create mode 100644 third_party/ascend/include/TritonToGraph/ProgramGridTransform.h create mode 100644 third_party/ascend/include/TritonToGraph/ProgramMappingScalarSpecialization.h create mode 100644 third_party/ascend/include/TritonToGraph/tensor.h create mode 100644 third_party/ascend/include/TritonToHFusion/CMakeLists.txt create mode 100644 third_party/ascend/include/TritonToHFusion/Passes.h create mode 100644 third_party/ascend/include/TritonToHFusion/Passes.td create mode 100644 third_party/ascend/include/TritonToHIVM/CMakeLists.txt create mode 100644 third_party/ascend/include/TritonToHIVM/Passes.h create mode 100644 third_party/ascend/include/TritonToHIVM/Passes.td create mode 100644 third_party/ascend/include/TritonToLLVM/CMakeLists.txt create mode 100644 third_party/ascend/include/TritonToLLVM/Passes.h create mode 100644 third_party/ascend/include/TritonToLLVM/Passes.td create mode 100644 third_party/ascend/language/cann/__init__.py create mode 100644 third_party/ascend/language/cann/extension/__init__.py create mode 100644 third_party/ascend/language/cann/extension/_utils.py create mode 100644 third_party/ascend/language/cann/extension/aux_ops.py create mode 100644 third_party/ascend/language/cann/extension/builder.py create mode 100644 third_party/ascend/language/cann/extension/builtin_custom_ops.py create mode 100644 third_party/ascend/language/cann/extension/code_generator.py create mode 100644 third_party/ascend/language/cann/extension/core.py create mode 100644 third_party/ascend/language/cann/extension/custom_op.py create mode 100644 third_party/ascend/language/cann/extension/dispatch.py create mode 100644 third_party/ascend/language/cann/extension/math_ops.py create mode 100644 third_party/ascend/language/cann/extension/mem_ops.py create mode 100644 third_party/ascend/language/cann/extension/scope.py create mode 100644 third_party/ascend/language/cann/extension/semantic.py create mode 100644 third_party/ascend/language/cann/extension/vec_ops.py create mode 100644 third_party/ascend/language/cann/libdevice.py create mode 100644 third_party/ascend/language/kernels/__init__.py create mode 100644 third_party/ascend/language/kernels/gather.py create mode 100644 third_party/ascend/lib/AutoBlockify/AutoBlockify.cpp create mode 100644 third_party/ascend/lib/AutoBlockify/CMakeLists.txt create mode 100644 third_party/ascend/lib/AutoBlockify/RewriteOperation.cpp create mode 100644 third_party/ascend/lib/AutoBlockify/Utils.cpp create mode 100644 third_party/ascend/lib/CMakeLists.txt create mode 100644 third_party/ascend/lib/CommonIRToHIVM/CMakeLists.txt create mode 100644 third_party/ascend/lib/CommonIRToHIVM/CommonIRToHIVM.cpp create mode 100644 third_party/ascend/lib/CommonIRToHIVM/ReduceSumStrengthPass.cpp create mode 100644 third_party/ascend/lib/Dialect/CMakeLists.txt create mode 100644 third_party/ascend/lib/Dialect/TritonAscend/CMakeLists.txt create mode 100644 third_party/ascend/lib/Dialect/TritonAscend/IR/CMakeLists.txt create mode 100644 third_party/ascend/lib/Dialect/TritonAscend/IR/TritonAscendAttrs.cpp create mode 100644 third_party/ascend/lib/Dialect/TritonAscend/IR/TritonAscendDialect.cpp create mode 100644 third_party/ascend/lib/Dialect/TritonAscend/IR/TritonAscendOps.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/CloneOps.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/InitDependentMap.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/Utils.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/AddDynamicCVPipeline.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferOuterScope.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/CMakeLists.txt create mode 100644 third_party/ascend/lib/DynamicCVPipeline/Common/BufferCountManager.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/Common/FlagIdManager.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/Common/MemoryEffectsTracker.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/Common.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlock.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanVectorBlockPass.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ReorderOpsByBlockId.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoad.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/AsyncLoadHoisting.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/DependencyAnalysis.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/LoopTransform.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/MultiBufferLoopBodyEmission.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromComputePass.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/SplitDataflow.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/AddBlockIdForControlOps.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/InterCoreTransferAndSync.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/MarkMainLoop.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/PreserveControlAttrsCanonicalize.cpp create mode 100644 third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/SeparateCVScope.cpp create mode 100644 third_party/ascend/lib/TritonAffinityOpt/CMakeLists.txt create mode 100644 third_party/ascend/lib/TritonAffinityOpt/DAG.cpp create mode 100644 third_party/ascend/lib/TritonAffinityOpt/DAGSSBuffer.cpp create mode 100644 third_party/ascend/lib/TritonAffinityOpt/DAGScope.cpp create mode 100644 third_party/ascend/lib/TritonAffinityOpt/DAGSync.cpp create mode 100644 third_party/ascend/lib/TritonToAnnotation/CMakeLists.txt create mode 100644 third_party/ascend/lib/TritonToAnnotation/TritonToAnnotation.cpp create mode 100644 third_party/ascend/lib/TritonToGraph/AliasAnalysis.cpp create mode 100644 third_party/ascend/lib/TritonToGraph/CMakeLists.txt create mode 100644 third_party/ascend/lib/TritonToGraph/ControlFlowGraph.cpp create mode 100644 third_party/ascend/lib/TritonToGraph/ControlFlowGraphBuilder.cpp create mode 100644 third_party/ascend/lib/TritonToGraph/DataflowGraph.cpp create mode 100644 third_party/ascend/lib/TritonToGraph/GraphAnalysis.cpp create mode 100644 third_party/ascend/lib/TritonToGraph/InterProceduralCFG.cpp create mode 100644 third_party/ascend/lib/TritonToGraph/MemorySsaBuilder.cpp create mode 100644 third_party/ascend/lib/TritonToGraph/Passes.cpp create mode 100644 third_party/ascend/lib/TritonToGraph/ProgramAxisDependenceAnalysis.cpp create mode 100644 third_party/ascend/lib/TritonToGraph/ProgramGridTransform.cpp create mode 100644 third_party/ascend/lib/TritonToGraph/ProgramMappingScalarSpecialization.cpp create mode 100644 third_party/ascend/lib/TritonToHFusion/CMakeLists.txt create mode 100644 third_party/ascend/lib/TritonToHFusion/TritonToHFusion.cpp create mode 100644 third_party/ascend/lib/TritonToHIVM/CMakeLists.txt create mode 100644 third_party/ascend/lib/TritonToHIVM/TritonToHIVM.cpp create mode 100644 third_party/ascend/lib/TritonToLLVM/CMakeLists.txt create mode 100644 third_party/ascend/lib/TritonToLLVM/TritonToLLVM.cpp create mode 100644 third_party/ascend/python/src/ir.cc create mode 100644 third_party/ascend/python/src/ir.h create mode 100644 third_party/ascend/python/src/main.cc create mode 100644 third_party/ascend/python/triton/extension/__init__.py create mode 100644 third_party/ascend/python/triton/extension/buffer/src/buffer_ir.cc create mode 100644 third_party/ascend/runtime/__init__.py create mode 100644 third_party/ascend/runtime/autoparser.py create mode 100644 third_party/ascend/runtime/autotuner.py create mode 100644 third_party/ascend/runtime/tile_generator.py create mode 100644 third_party/ascend/runtime/utils.py create mode 100644 third_party/ascend/spec/triton/extension/buffer/__init__.py create mode 100644 third_party/ascend/spec/triton/extension/buffer/language/__init__.py create mode 100644 third_party/ascend/spec/triton/extension/buffer/language/builder.py create mode 100644 third_party/ascend/spec/triton/extension/buffer/language/core.py create mode 100644 third_party/ascend/spec/triton/extension/buffer/language/semantic.py create mode 100644 third_party/ascend/spec/triton/tools/get_ascend_devices.py create mode 100644 third_party/ascend/spec_cpp/include/triton/Dialect/Triton/IR/TritonAttrDefs.td create mode 100644 third_party/ascend/spec_cpp/lib/Dialect/Triton/IR/Traits.cpp create mode 100644 third_party/ascend/test/Conversion/CommonIRToHIVM/reduce-sum-strength-pipeline.mlir create mode 100644 third_party/ascend/test/Conversion/CommonIRToHIVM/reduce-sum-strength.mlir create mode 100644 third_party/ascend/triton_ascend.cc create mode 100644 third_party/ascend/tutorials/01-vector-add.py create mode 100644 third_party/ascend/tutorials/02-fused-softmax.py create mode 100644 third_party/ascend/tutorials/03-matrix-multiplication.py create mode 100644 third_party/ascend/tutorials/04-low-memory-dropout.py create mode 100644 third_party/ascend/tutorials/05-layer-norm.py create mode 100644 third_party/ascend/tutorials/06-fused-attention.py create mode 100644 third_party/ascend/tutorials/07-extern-functions.py create mode 100644 third_party/ascend/tutorials/08-grouped-gemm.py create mode 100644 third_party/ascend/tutorials/09-persistent-matmul.py create mode 100644 third_party/ascend/unittest/CMakeLists.txt create mode 100644 third_party/ascend/unittest/Conversion/General/AddDynamicCVPipeline/AddMultiCache/Inner-scope.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/AutoBlockify/auto_blockify.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/DiscreteMaskAccess/atomic.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/DiscreteMaskAccess/loadstore.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/add_controlflow_condition_test.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/clone-ops.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/create-if-ops.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/crossCoreDeps_test.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/multi_deps_add_condition_test.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/no_deps_add_condition_test.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/single_deps_add_condition_test.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/skip-control-flow.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/update-for-ops.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/updateLoopIterTimes2_test.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/updateLoopIterTimes_test.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonAscendAllPass/conv1d.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToHFusion/fp_to_fp_rtz.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToHFusion/histogram.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToHFusion/histogram_expect_fail.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToHFusion/mod.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToHIVM/sync_block_op_conversion.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToLinalg/atomic_rmw.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToLinalg/clamp_fp8.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToLinalg/dotScaled_fp8.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToLinalg/hoist_broadcast.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToLinalg/hoist_broadcast2.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToLinalg/parse_custom.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToLinalg/parse_select.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_implicit_permute.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_load_out_of_bound.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_nd2nz_block_ptr.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_atomic_rmw.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_div.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_fptosi.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_select_dense_const.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_sub.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_while_loop_index_tensor.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToStructured/CmpConverter.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToStructured/PromotePointerIterArgsPattern.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToStructured/SplatCmpConverter.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToStructured/parseCmp.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToStructured/parseConstant.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToStructured/parseMakeRange.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToStructured/parseRem.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToUnstructure/bubbleupoperation.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToUnstructure/if_simplifier.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToUnstructure/indirect_func.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToUnstructure/indirect_load.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToUnstructure/indirect_mem_access.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToUnstructure/indirect_store.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToUnstructure/nested_loop.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToUnstructure/splat.mlir create mode 100644 third_party/ascend/unittest/Conversion/General/TritonToUnstructure/unstructure_mix.mlir create mode 100755 third_party/ascend/unittest/affine_map/affine_map.py create mode 100755 third_party/ascend/unittest/affine_map/affine_map_buffer_type_demo.py create mode 100755 third_party/ascend/unittest/affine_map/affine_map_complex_expr_demo.py create mode 100755 third_party/ascend/unittest/affine_map/affine_map_indexing_map_demo.py create mode 100755 third_party/ascend/unittest/affine_map/affine_map_parse_demo.py create mode 100644 third_party/ascend/unittest/autotune_ut/01-vector-add.py create mode 100644 third_party/ascend/unittest/autotune_ut/02-fused-softmax.py create mode 100644 third_party/ascend/unittest/autotune_ut/03-layer-norm.py create mode 100644 third_party/ascend/unittest/autotune_ut/04-libentry.py create mode 100644 third_party/ascend/unittest/autotune_ut/test_autotune_param_valid.py create mode 100644 third_party/ascend/unittest/autotune_ut/test_common.py create mode 100644 third_party/ascend/unittest/autotune_ut/test_customized_config.py create mode 100644 third_party/ascend/unittest/autotune_ut/test_do_bench_compat.py create mode 100644 third_party/ascend/unittest/autotune_ut/test_low_dim_axes_parse.py create mode 100644 third_party/ascend/unittest/autotune_ut/test_mask_parse.py create mode 100644 third_party/ascend/unittest/autotune_ut/test_no_tiling_axis_parse.py create mode 100644 third_party/ascend/unittest/autotune_ut/test_reduction_axes_parse.py create mode 100644 third_party/ascend/unittest/autotune_ut/test_split_axis_parse.py create mode 100644 third_party/ascend/unittest/autotune_ut/test_tiling_axis_parse.py create mode 100644 third_party/ascend/unittest/conftest.py create mode 100755 third_party/ascend/unittest/custom_op/builtin_ops_demo.py create mode 100755 third_party/ascend/unittest/custom_op/custom_op_demo.py create mode 100755 third_party/ascend/unittest/custom_op/custom_op_extra_buffer_demo.py create mode 100755 third_party/ascend/unittest/custom_op/custom_op_indexing_map_complex_demo.py create mode 100755 third_party/ascend/unittest/custom_op/custom_op_indexing_map_compose_demo.py create mode 100755 third_party/ascend/unittest/custom_op/test_gather_load.py create mode 100644 third_party/ascend/unittest/custom_op/test_index_select.py create mode 100644 third_party/ascend/unittest/generalization_cases/acc_util.py create mode 100644 third_party/ascend/unittest/generalization_cases/conftest.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_abs.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_add_mindspore.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_advance.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_and.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_argmax.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_argmin.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_associative_scan.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_atan.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_atomic_add.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_atomic_and.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_atomic_cas.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_atomic_max.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_atomic_min.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_atomic_or.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_atomic_xchg.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_atomic_xor.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_broadcast.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_broadcast_to.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_cast.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_cdiv.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_ceil.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_common.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_cos.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_count_dim0.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_count_dim1.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_cumprod.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_cumsum.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_debug_barrier.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_device_print.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_div_rn.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_dot_scaled.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_eq.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_erf.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_exp.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_exp2.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_expand_dims.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_fdiv.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_flip.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_full_op.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_ge_op.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_add.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_arange.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_cat.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_clamp.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_div.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_floor.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_floordiv.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_fma.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_gather.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_interleave.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_join.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_log.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_log2.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_maximum.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_minimum.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_mul.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_ravel.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_reshape.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_rsqrt.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_sigmoid.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_sin.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_softmax.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_split.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_sub.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_tensor_descriptor.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_general_view.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_gt_op.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_invalid_fp8.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_invert.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_le_op.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_load_store.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_log1p.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_logical_and_op.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_logical_or_op.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_lshift_op.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_lt_op.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_make_blkptr_matmul.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_make_block_ptr.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_matmul.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_matmul_mindspore.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_max.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_min.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_mod.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_multi-process.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_ne.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_neg.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_not.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_or.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_permute_1d_2d.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_permute_3d.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_permute_4d_5d.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_rand.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_range.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_reduce.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_relu.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_rshift_op.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_scalar_tensor.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_softmax_mindspore.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_sort.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_sqrt.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_sqrt_rn.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_static_print_and_assert_op.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_sum.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_sum_dim0.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_sum_dim1.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_swizzle2d.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_tan.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_trans_1d_2d.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_trans_3d.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_trans_4d_5d.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_umulhi.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_where.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_xor.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_xorsum.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_zeros_op.py create mode 100644 third_party/ascend/unittest/generalization_cases/test_zeroslike.py create mode 100644 third_party/ascend/unittest/kernels/README.md create mode 100644 third_party/ascend/unittest/kernels/common_kernel.py create mode 100644 third_party/ascend/unittest/kernels/test_common.py create mode 100644 third_party/ascend/unittest/kernels/test_triton_kernel.py create mode 100644 third_party/ascend/unittest/kernels/vllm/_triton_rope.py create mode 100644 third_party/ascend/unittest/kernels/vllm/chunk_gated_delta_rule_fwd_kernel_h_blockdim64.py create mode 100644 third_party/ascend/unittest/kernels/vllm/chunk_local_cumsum_scalar_kernel.py create mode 100644 third_party/ascend/unittest/kernels/vllm/chunk_scaled_dot_kkt_fwd_kernel.py create mode 100644 third_party/ascend/unittest/kernels/vllm/expand_kernel.py create mode 100644 third_party/ascend/unittest/kernels/vllm/fused_gdn_gating_kernel.py create mode 100644 third_party/ascend/unittest/kernels/vllm/fused_qkvzba_split_reshape_cat_kernel.py create mode 100644 third_party/ascend/unittest/kernels/vllm/l2norm_fwd_kernel2_loop.py create mode 100644 third_party/ascend/unittest/kernels/vllm/layer_norm_fwd_kernel.py create mode 100644 third_party/ascend/unittest/kernels/vllm/merge_16x16_to_32x32_inverse_kernel.py create mode 100644 third_party/ascend/unittest/kernels/vllm/merge_16x16_to_64x64_inverse_kernel.py create mode 100644 third_party/ascend/unittest/kernels/vllm/recompute_w_u_fwd_kernel.py create mode 100644 third_party/ascend/unittest/kernels/vllm/rejection_greedy_sample_spec_len_1_triton.py create mode 100644 third_party/ascend/unittest/kernels/vllm/rejection_random_sample_kernel.py create mode 100644 third_party/ascend/unittest/kernels/vllm/sample_recovered_tokens_kernel.py create mode 100644 third_party/ascend/unittest/kernels/vllm/solve_tril_16x16_kernel.py create mode 100644 third_party/ascend/unittest/lit.cfg.py create mode 100644 third_party/ascend/unittest/lit.site.cfg.py.in create mode 100644 third_party/ascend/unittest/pytest_ut/attn_cp_triton_kernel_3d.py create mode 100644 third_party/ascend/unittest/pytest_ut/attn_cp_triton_kernel_3d_la.py create mode 100644 third_party/ascend/unittest/pytest_ut/conftest.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_01_vector_add.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_02_fused_softmax.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_03_matrix_multiplication.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_04_low_memory_dropout.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_05_layer_norm.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_06_fused_attention.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_07_extern_functions.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_08_grouped_gemm.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_09_persistent_matmul.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_10_gather_sorted.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_11_rab_time.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_12_hstu_attention.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_13_matrix_multiplication_optimized.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_14_accuracy_comparison.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_15_demo_autotune.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_16_profiler.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_17_demo_libentry.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_18_gather.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_2d_permute.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_3Dgrid.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_abs.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_abs_2.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_acos.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_acosh.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_add.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_address_check.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_advance.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_advance_ptr.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_affine_map_binding.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_alloc.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_and.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_annotations.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_arange.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_arch.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_argmax.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_argmin.py create mode 100755 third_party/ascend/unittest/pytest_ut/test_ascend_barrier.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_asin.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_asinh.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_asm.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_asm_scalar.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_associative_scan.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_associative_scan_multi_input.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_assume.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_assume1.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_atan.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_atan2.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_atanh.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_atomic_add.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_atomic_and.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_atomic_cas.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_atomic_max.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_atomic_min.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_atomic_or.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_atomic_rmw_useanalysis.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_atomic_xchg.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_atomic_xor.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_attn_cp.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_bind_buffer.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_block_ptr.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_boundary_check.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_broadcast.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_broadcast_op.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_cannonicalize_tl_where.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_cast_full.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_cat.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_cat_dim.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_cat_help_func.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_cdiv.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_ceil.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_celoss_indices.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_clamp.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_common.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_compile_hint.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_complex_mask.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_copy.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_copysign.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_cos.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_cos_2.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_cosh.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_cumprod.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_cumsum.py create mode 100755 third_party/ascend/unittest/pytest_ut/test_custom.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_cyl_bessel_i0.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_debug_barrier.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_debug_triton_opt_cmd.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_device_print.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_device_print_comprehensive.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_device_print_script.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_discrete_mask_atomic.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_discrete_mask_loadstore.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_discrete_mask_tail_block_mte_oob.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_discrete_overlap_mask.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_div.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_dot.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_dot_scaled.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_downgrade.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_elementwise_ceil.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_elementwise_clip.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_elementwise_f2i.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_elementwise_floor.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_elementwise_i2f.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_elementwise_round.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_eq.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_eq_2.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_erfinv.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_exp.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_exp2.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_exp_.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_expand_dims.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_expm1.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_extract_slice.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_fast_dividef.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_fast_expf.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_fdiv.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_fixpipe.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_fixpipe_tensor_dst.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_flip.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_floor.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_floordiv.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_for_ptr.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_full.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_fusedattention.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_gamma.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_gather.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_gather_simd.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_ge.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_ge_2.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_gelu.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_gt.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_hd_permute.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_histogram.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_hoistbroadcast.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_hypot.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_if_advance.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_if_load.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_if_tensor.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_implicit_atomic.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_implicit_permute.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_index_select.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_index_select_inductor.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_indirect_atomic_add.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_indirect_atomic_and.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_indirect_atomic_cas.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_indirect_atomic_max.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_indirect_atomic_min.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_indirect_atomic_or.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_indirect_atomic_xchg.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_indirect_atomic_xor.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_indirect_scalar_load_offset.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_insert_slice.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_interleave.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_interleave_optimizaiton.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_invert.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_isfinited.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_isnan.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_join.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_lanzcos.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_launcher_empty_signature.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_layernorm.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_ldst.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_le.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_lgamma.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_linearize.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_linearize_fallback.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_linearize_mask.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_linearize_mask_fallback.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_linearize_permute.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_load.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_load_store.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_log.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_log10.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_log1p.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_log2.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_log_2.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_logical_and.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_logical_or.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_lshift.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_lt.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_makeblockptr_negative_padding.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_makeblockptr_permute.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_max_constancy.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_max_contiguous.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_max_dim0.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_max_dim1.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_max_propagate_nan.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_max_vector.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_maximum.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_mean_dim0.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_mean_dim1.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_mean_vector.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_min_dim0.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_min_dim1.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_min_vector.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_minimum.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_mod.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_mul.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_mul_reduce.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_multi_return.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_multibuffer.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_multiple_of.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_ne.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_nearbyint.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_nearest.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_neg.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_negative_mask_dim.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_nextafter.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_not.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_npu_indexing.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_npu_indexing2.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_or.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_paged_kvcache_krope.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_parallel.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_permute.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_permute_full.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_permute_reshape.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_permuted_boundary_check.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_pointer_type.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_pow.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_precise_div.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_precise_sqrt.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_ptr_add.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_rand.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_range.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_ravel.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_reduce_count_vector.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_reduce_maximum.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_reduce_mean.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_reduce_min_4_keepdim_True_with_index_op.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_reduce_minimum.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_reduce_sum.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_relu.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_reshape.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_resize_performance.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_rint.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_rms_norm.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_rotary_embedding.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_rotatry_gpt.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_rotaty_embedding_gpt.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_rshift.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_rsqrt.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_runtime_utils.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_scalarPointer.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_scalar_calc.py create mode 100755 third_party/ascend/unittest/pytest_ut/test_scope.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_scope_buffer.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_select_analysis.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_select_analysis_for_divsiop.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_select_analysis_for_invert.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_sigmoid.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_signbit.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_silu.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_simplify_iterargs.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_sin.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_sinh.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_sink_broadcast.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_softmax.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_sort.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_split.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_sqrt.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_static_print_and_assert.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_store_scalar.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_stride0_load.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_strides.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_sub.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_sub_vec_id.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_sub_vec_num.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_subview.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_sum.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_sum_dim0.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_sum_dim1.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_sum_vector.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_swap.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_swiglu.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_swizzle2d.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_sync_block.py create mode 100755 third_party/ascend/unittest/pytest_ut/test_sync_block_all.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_tan.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_tanh.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_template.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_tensor_descriptor.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_to_buffer.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_to_tensor.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_top2gating_argmax.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_topk.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_trans_3d.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_triton_eq.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_triton_le.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_triton_lt.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_triton_neq.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_triton_unified_attention.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_trunc.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_umulhi.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_unlign_max_with_index.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_unlign_sum.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_unused_func_arg.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_use_analysis.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_utils.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_view.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_where_lt.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_where_mask.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_where_var.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_xor.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_xor_sum.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_zeros.py create mode 100644 third_party/ascend/unittest/pytest_ut/test_zeroslike.py create mode 100755 third_party/ascend/unittest/run_prtest.sh diff --git a/CMakeLists.txt b/CMakeLists.txt index b1db0614ba..f23a11fef2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -172,6 +172,12 @@ if(FLAGTREE_BACKEND AND EXISTS ${BACKEND_INCLUDE_DIR}) # flagtree: include directories include_directories(${PROJECT_SOURCE_DIR}/third_party/${FLAGTREE_BACKEND}/include) include_directories(${PROJECT_BINARY_DIR}/third_party/${FLAGTREE_BACKEND}/include) # Tablegen'd files + if(FLAGTREE_BACKEND STREQUAL "ascend") + # The ascend backend builds the core include/lib from its own CMakeLists, so + # only the search paths are needed here. + include_directories(${PROJECT_SOURCE_DIR}/include) + include_directories(${PROJECT_BINARY_DIR}/include) # Tablegen'd files + endif() else() include_directories(${PROJECT_SOURCE_DIR}/include) include_directories(${PROJECT_BINARY_DIR}/include) # Tablegen'd files @@ -212,6 +218,8 @@ if(TRITON_BUILD_PYTHON_MODULE) # FlagPrism: register component targets before backend codegen captures TRITON_LIBS. flagtree_add_flagprism_components() flagtree_configure_flir_dependency() + # Provide the BiShengIR targets before the backend subdirectory links them. + flagtree_configure_ascend_npu_ir_dependency() foreach(CODEGEN_BACKEND ${TRITON_CODEGEN_BACKENDS}) add_subdirectory(third_party/${CODEGEN_BACKEND}) @@ -257,6 +265,10 @@ if(TRITON_BUILD_PYTHON_MODULE) ${triton_plugins} # mlir + # Provides mlir::func::registerInlinerExtension, needed by the ascend + # backend's commonir-to-hivm lowering (the Inliner pass that follows it + # aborts if the func dialect's inliner interface is unregistered). + MLIRFuncInlinerExtension MLIRAMDGPUDialect MLIRNVVMDialect MLIRNVVMToLLVMIRTranslation diff --git a/cmake/FlagTreeOptions.cmake b/cmake/FlagTreeOptions.cmake index a5392d4428..47797ced76 100644 --- a/cmake/FlagTreeOptions.cmake +++ b/cmake/FlagTreeOptions.cmake @@ -45,6 +45,10 @@ macro(flagtree_configure_options) set(FLAGCX_ENABLED ON) endif() + # TLE DSA (ascend-only) needs the third_party/tle/dsa subtree, which this + # branch does not carry; keep it an explicit opt-in. + set(FLAGTREE_TLE_DSA "$ENV{FLAGTREE_TLE_DSA}") + set(FLAGTREE_TLE ON) if(FLAGTREE_BACKEND STREQUAL "xpu") set(FLAGTREE_TLE OFF) @@ -82,6 +86,42 @@ macro(flagtree_configure_options) set(CMAKE_CXX_COMPILER clang++-16) add_definitions(-D__NVIDIA__) add_definitions(-D__AMD__) + elseif(FLAGTREE_BACKEND STREQUAL "ascend") + set(CMAKE_C_COMPILER clang) + set(CMAKE_CXX_COMPILER clang++) + # The ascend backend consumes FLIR's incubated passes (TritonToLinalgIncubated). + set(FLIR_BUILD_INCUBATED ON CACHE BOOL "Build FLIR incubated passes" FORCE) + # AscendNPU-IR gates its sources on the LLVM major version. cmake/llvm-hash.txt + # pins LLVM 22 here. + # + # Deliberately leave LLVM_MAJOR_VERSION_21_COMPATIBLE unset: the ascend + # CMakeLists uses it to opt into linking MLIRMeshDialect, but LLVM 22 + # renamed that dialect to Shard (only libMLIRShardDialect.a ships in the + # pinned package), so turning it on would reference a missing target. + set(LLVM_MAJOR_VERSION_22_OR_HIGHER ON) + set(LLVM_MAJOR_VERSION_22_COMPATIBLE ON CACHE BOOL + "AscendNPU-IR build with llvm major version 22 or later" FORCE) + add_definitions(-D__LLVM_MAJOR_VERSION_22_COMPATIBLE__) + add_compile_options("-Wno-deprecated-declarations") + add_compile_options("-Wno-error=deprecated-declarations") + # TLE DSA is an opt-in extension requiring the third_party/tle/dsa subtree, + # which this branch does not carry. The ascend C++ sources guard every use + # of it with #ifdef __TLE_DSA__, so leaving this off compiles those paths + # out and drops the corresponding CMake dependencies. + if(FLAGTREE_TLE_DSA) + if(NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/third_party/tle/dsa/CMakeLists.txt") + message(FATAL_ERROR "FLAGTREE_TLE_DSA requires third_party/tle/dsa") + endif() + add_definitions(-D__TLE_DSA__) + endif() + # AscendNPU-IR: build only the standalone IR, and keep it out of the + # backend's own CMakeLists so the subdirectory is added exactly once. + set(BISHENGIR_ENABLE_A5_UNPUBLISHED_FEATURES ON) + set(BISHENGIR_BUILD_STANDALONE_IR_ONLY ON) + set(ASCENDNPU_IR_SRC_DIR ${CMAKE_CURRENT_SOURCE_DIR}/third_party/ascend/AscendNPU-IR) + set(ASCENDNPU_IR_BINARY_DIR ${CMAKE_CURRENT_BINARY_DIR}/third_party/ascend/AscendNPU-IR) + include_directories(${ASCENDNPU_IR_SRC_DIR}/bishengir/include) + include_directories(${ASCENDNPU_IR_BINARY_DIR}/bishengir/include) elseif(FLAGTREE_BACKEND STREQUAL "tsingmicro") set(CMAKE_C_COMPILER clang) set(CMAKE_CXX_COMPILER clang++) @@ -328,12 +368,12 @@ macro(flagtree_configure_flir_dependency) add_subdirectory(third_party/flir/lib/Dialect/CommonIR) endif() - if(FLAGTREE_BACKEND STREQUAL "tsingmicro") + if(FLAGTREE_BACKEND MATCHES "^(tsingmicro|ascend)$") if(NOT EXISTS "${PROJECT_SOURCE_DIR}/third_party/flir/CMakeLists.txt") message(FATAL_ERROR "The ${FLAGTREE_BACKEND} backend requires third_party/flir") endif() - # TsingMicro only consumes FLIR's C++ targets; do not build its Python/CPU plugin. + # These backends only consume FLIR's C++ targets; do not build its Python/CPU plugin. set(TRITON_SHARED_BUILD_CPU_BACKEND OFF) list(REMOVE_ITEM TRITON_CODEGEN_BACKENDS "flir") if(NOT TARGET TritonSharedUtils) @@ -346,6 +386,21 @@ macro(flagtree_configure_flir_dependency) endmacro() +# AscendNPU-IR provides the BiShengIR dialects the ascend backend links against. +# Added from the top level (not from the backend's own CMakeLists) so the source +# directory is only ever associated with one binary directory. +macro(flagtree_configure_ascend_npu_ir_dependency) + if(FLAGTREE_BACKEND STREQUAL "ascend") + if(NOT EXISTS "${ASCENDNPU_IR_SRC_DIR}/CMakeLists.txt") + message(FATAL_ERROR + "The ascend backend requires AscendNPU-IR at ${ASCENDNPU_IR_SRC_DIR}. " + "Run the Python package build to download third-party dependencies.") + endif() + add_subdirectory(${ASCENDNPU_IR_SRC_DIR} ${ASCENDNPU_IR_BINARY_DIR}) + endif() +endmacro() + + macro(flagtree_configure_tle_plugin append_tle_plugin) if(FLAGTREE_TLE) if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/third_party/tle/CMakeLists.txt") @@ -575,6 +630,12 @@ macro(flagtree_configure_tools_and_tests) flagtree_add_tle_generated_header_dependencies() endif() add_subdirectory(test) + elseif(FLAGTREE_BACKEND STREQUAL "ascend") + # The ascend backend adds its own bin/ (third_party/ascend/CMakeLists.txt), + # whose tools link TritonTestAnalysis/TritonTestDialect from test/lib. + # Add test/ so those libraries get defined; bin/ is not added here because + # the backend supplies its own. + add_subdirectory(test) elseif(FLAGTREE_BACKEND STREQUAL "iluvatar") option(FLAGTREE_ILUVATAR_BUILD_BIN "Build third_party/iluvatar/bin tools and lit tests" OFF) @@ -683,6 +744,13 @@ macro(flagtree_added_python_src) endforeach() add_library(triton SHARED ${_flagtree_python_sources}) + if(FLAGTREE_BACKEND STREQUAL "ascend") + # third_party/ascend/python/src/main.cc registers the buffer_ir submodule, + # so its implementation must be part of the triton module. + target_sources(triton PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/third_party/ascend/python/triton/extension/buffer/src/buffer_ir.cc) + endif() + if(FLAGTREE_BACKEND STREQUAL "xpu") target_sources(triton PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/third_party/xpu/python/src/mlir_pass_abi_shim.cc) add_dependencies(triton TritonAMDGPUTableGen TritonAMDGPUAttrDefsIncGen) diff --git a/python/requirements-ascend.txt b/python/requirements-ascend.txt new file mode 100644 index 0000000000..9d3f182fde --- /dev/null +++ b/python/requirements-ascend.txt @@ -0,0 +1,3 @@ +--extra-index-url https://ascend.devcloud.huaweicloud.com/cann/pypi/simple/ + +cann-shmem diff --git a/python/setup_tools/setup_helper.py b/python/setup_tools/setup_helper.py index 73df70aef7..1499dc3d3f 100644 --- a/python/setup_tools/setup_helper.py +++ b/python/setup_tools/setup_helper.py @@ -68,7 +68,8 @@ def get_long_description(): def init_backends(backend_installer): if flagtree_backend: - if flagtree_backend in ("aipu", "tsingmicro", "enflame", "rpu", "thrive", "sunrise", "tileir", "ppu", "amd"): + if flagtree_backend in ("aipu", "tsingmicro", "enflame", "rpu", "thrive", "sunrise", "tileir", "ppu", "amd", + "ascend"): backends = [ *backend_installer.copy(configs.default_backends + tuple(configs.extend_backends)), *backend_installer.copy_externals(), @@ -888,7 +889,12 @@ def uninstall_triton(): cache = FlagTreeCache() -download_flagtree_third_party("flir", condition=(flagtree_backend == "tsingmicro"), required=True) +download_flagtree_third_party("flir", condition=(flagtree_backend in ("tsingmicro", "ascend")), required=True) + +# The ascend backend links the BiShengIR dialects built from AscendNPU-IR. +if flagtree_backend == "ascend": + configs.activated_module.warn_stale_ascend_npu_ir_checkout(utils.flagtree_submodules["AscendNPU-IR"]) + download_flagtree_third_party("AscendNPU-IR", condition=True, required=True) ''' FlagCX is a third-party library adopted by the tle distributed system, refer to https://github.com/flagos-ai/FlagCX diff --git a/python/setup_tools/utils/__init__.py b/python/setup_tools/utils/__init__.py index 3931c53025..889419792f 100644 --- a/python/setup_tools/utils/__init__.py +++ b/python/setup_tools/utils/__init__.py @@ -60,6 +60,10 @@ def _register_submodule(self, submodule): }, {"name": "flir", "url": "https://github.com/flagos-ai/flir.git"}, {"name": "FlagPrism", "url": "https://github.com/flagos-ai/FlagPrism.git"}, + { + "name": "AscendNPU-IR", "url": "https://gitcode.com/Ascend/AscendNPU-IR.git", "relative_path": + "ascend/AscendNPU-IR", "commit_id": "bb13ef6120526aa3df6f950006af5ec0e585d2f6" + }, { "name": "flagcx", "url": "https://github.com/flagos-ai/FlagCX.git", "relative_path": "tle/third_party/flagcx", "commit_id": "c93438decaba8d935e7aa6ae8377e56fd3910b32" diff --git a/python/setup_tools/utils/ascend.py b/python/setup_tools/utils/ascend.py new file mode 100644 index 0000000000..e9c3c26c34 --- /dev/null +++ b/python/setup_tools/utils/ascend.py @@ -0,0 +1,16 @@ +import os +import subprocess + + +def warn_stale_ascend_npu_ir_checkout(module): + # DownloadManager reuses dst_path as-is when it already exists, so a + # checkout left on a different commit would silently shadow the pin + # registered for this build. + if not module.commit_id or not os.path.isdir(os.path.join(module.dst_path, ".git")): + return + head = subprocess.run(["git", "rev-parse", "HEAD"], cwd=module.dst_path, capture_output=True, text=True) + if head.returncode != 0 or head.stdout.strip().startswith(module.commit_id): + return + print(f"[WARNING] {module.dst_path} is at {head.stdout.strip()[:8]} but this build pins " + f"AscendNPU-IR to {module.commit_id[:8]}; remove the directory (or run " + f"'git checkout {module.commit_id}' inside it) before rebuilding") diff --git a/python/triton/_utils.py b/python/triton/_utils.py index 9f4ab1a48c..683d40cdd4 100644 --- a/python/triton/_utils.py +++ b/python/triton/_utils.py @@ -26,6 +26,11 @@ from typing import Any, Callable, TYPE_CHECKING, Union, List, Dict from triton._flagtree_spec import spec_func +from triton._flagtree_backend import FLAGTREE_BACKEND + +# The ascend backend lowers tensors through a SIMD pipeline that does not +# require power-of-2 block shapes (see validate_block_shape below). +_SKIP_POW2_BLOCK_SHAPE = FLAGTREE_BACKEND == "ascend" if TYPE_CHECKING: from .language import core @@ -80,7 +85,10 @@ def validate_block_shape(shape: List[int]): for i, d in enumerate(shape): if not isinstance(d, int): raise TypeError(f"Shape element {i} must have type `constexpr[int]`, got `constexpr[{type(d)}]") - if not is_power_of_two(d): + # Backends that lower through a SIMD pipeline (e.g. ascend) support + # non-power-of-2 block shapes; this check has no builder in scope, so + # gate it on the active backend. + if not is_power_of_two(d) and not _SKIP_POW2_BLOCK_SHAPE: raise ValueError(f"Shape element {i} must be a power of 2") numel *= d diff --git a/python/triton/compiler/compiler.py b/python/triton/compiler/compiler.py index 75fa1d4b18..d454b0e8fa 100644 --- a/python/triton/compiler/compiler.py +++ b/python/triton/compiler/compiler.py @@ -469,8 +469,12 @@ def __init__(self, src, metadata_group, hash): # stores the text of each level of IR that was generated during compilation asm_files = [Path(p) for c, p in metadata_group.items() if not c.endswith(".json")] binary_ext = backend.binary_ext + # A backend may emit more than one binary artifact (e.g. the ascend + # backend caches the MLIR bytecode stage as .mlirbc alongside .npubin); + # those must be read as bytes, not decoded as UTF-8 text. + binary_exts = set(getattr(backend, "binary_extensions", None) or {binary_ext}) self.asm = AsmDict({ - file.suffix[1:]: file.read_bytes() if file.suffix[1:] == binary_ext else file.read_text() + file.suffix[1:]: file.read_bytes() if file.suffix[1:] in binary_exts else file.read_text() for file in asm_files }) self.metadata_group = metadata_group @@ -516,8 +520,16 @@ def raise_(err): if knobs.runtime.kernel_load_start_hook is not None: knobs.runtime.kernel_load_start_hook(self.module, self.function, self.name, self.metadata_group, self.hash) # TODO: n_regs, n_spills should be metadata generated when calling `ptxas` - self.module, self.function, self.n_regs, self.n_spills, self.n_max_threads = kernel_driver.utils.load_binary( - self.name, self.kernel, self.metadata.shared, device) + # The ascend backend's load_binary needs the kernel's mix_mode (aiv/aic/mix) + # to pick the launch path, and keys the binary by the mangled kernel_name + # rather than the Python function name. + if hasattr(self.metadata, "mix_mode"): + self.module, self.function, self.n_regs, self.n_spills, self.n_max_threads = \ + kernel_driver.utils.load_binary(self.metadata.kernel_name, self.kernel, self.metadata.shared, device, + self.metadata.mix_mode) + else: + self.module, self.function, self.n_regs, self.n_spills, self.n_max_threads = \ + kernel_driver.utils.load_binary(self.name, self.kernel, self.metadata.shared, device) warp_size = self.metadata.target.warp_size if self.metadata.num_warps * warp_size > self.n_max_threads \ and self.metadata.target.backend != "maca": # flagtree metax diff --git a/python/triton/extension/__init__.py b/python/triton/extension/__init__.py new file mode 100644 index 0000000000..ac05b013db --- /dev/null +++ b/python/triton/extension/__init__.py @@ -0,0 +1,4 @@ +# flagtree backend path specialization +from triton._flagtree_spec import spec_path + +spec_path(__path__) diff --git a/python/triton/knobs.py b/python/triton/knobs.py index 49ad56768d..50e86894ac 100644 --- a/python/triton/knobs.py +++ b/python/triton/knobs.py @@ -450,6 +450,9 @@ class autotuning_knobs(base_knobs): cache: env_bool = env_bool("TRITON_CACHE_AUTOTUNING") print: env_bool = env_bool("TRITON_PRINT_AUTOTUNING") adjust_block_size: env_bool = env_bool("FLAGTREE_AABS", True) # flagtree aabs + # Opt-in to a backend-provided autotuner (consumed by the ascend runtime, + # which replaces triton.autotune when this is set). + use_backend_autotuning: env_bool = env_bool("FLAGTREE_USE_BACKEND_AUTOTUNING") class LaunchHook(Protocol): diff --git a/python/triton/language/semantic.py b/python/triton/language/semantic.py index 77dbd8fc1c..84b6bb0cd0 100644 --- a/python/triton/language/semantic.py +++ b/python/triton/language/semantic.py @@ -601,7 +601,9 @@ def arange(self, start: int, end: int, *, ret_ty: tl.block_type = None) -> Tenso if end <= start: raise ValueError("arange's end argument must be greater than the start argument") range = end - start - if (range & (range - 1)) != 0: + # The power-of-2 restriction only applies in SIMT mode; backends whose + # builder reports SIMD (e.g. ascend) accept arbitrary ranges. + if (range & (range - 1)) != 0 and getattr(self.builder, "is_simt_mode", lambda: True)(): raise ValueError("arange's range must be a power of 2") shape = [range] if ret_ty is None: diff --git a/python/triton/spec/ascend b/python/triton/spec/ascend new file mode 120000 index 0000000000..4233a96de7 --- /dev/null +++ b/python/triton/spec/ascend @@ -0,0 +1 @@ +../../../third_party/ascend/spec/ \ No newline at end of file diff --git a/third_party/ascend/CMakeLists.txt b/third_party/ascend/CMakeLists.txt new file mode 100644 index 0000000000..475a7be164 --- /dev/null +++ b/third_party/ascend/CMakeLists.txt @@ -0,0 +1,97 @@ +include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include) +include_directories(${CMAKE_CURRENT_BINARY_DIR}/include) +include_directories(${CMAKE_SOURCE_DIR}/third_party/flir/include) +include_directories(${CMAKE_BINARY_DIR}/third_party/flir/include) + +# list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake") + +# Triton Ascend is dependent on AscendNPU IR +# set(ASCENDNPU_IR_SRC_DIR "${PROJECT_SOURCE_DIR}/third_party/ascend/AscendNPU-IR") +# set(ASCENDNPU_IR_BINARY_DIR "${PROJECT_BINARY_DIR}/third_party/ascend/AscendNPU-IR") + +# set(BISHENGIR_ENABLE_A5_UNPUBLISHED_FEATURES ON) +# set(BISHENGIR_BUILD_STANDALONE_IR_ONLY ON) +# +# # Temporarily save and clear the RULE_LAUNCH_* attributes in the current directory to ensure AscendNPU-IR is not affected. +# get_property(_saved_launch_compile DIRECTORY PROPERTY RULE_LAUNCH_COMPILE) +# get_property(_saved_launch_link DIRECTORY PROPERTY RULE_LAUNCH_LINK) +# set_property(DIRECTORY PROPERTY RULE_LAUNCH_COMPILE "") +# set_property(DIRECTORY PROPERTY RULE_LAUNCH_LINK "") +# +# add_subdirectory(${ASCENDNPU_IR_SRC_DIR} ${ASCENDNPU_IR_BINARY_DIR}) +# +# # restore properties +# set_property(DIRECTORY PROPERTY RULE_LAUNCH_COMPILE ${_saved_launch_compile}) +# set_property(DIRECTORY PROPERTY RULE_LAUNCH_LINK ${_saved_launch_link}) +# +# include_directories(${ASCENDNPU_IR_SRC_DIR}/bishengir/include) +# include_directories(${ASCENDNPU_IR_BINARY_DIR}/bishengir/include) # Tablegen'd files + +add_subdirectory(backend/spec/lib) + +add_subdirectory(${PROJECT_SOURCE_DIR}/include + ${PROJECT_BINARY_DIR}/include) +add_subdirectory(${PROJECT_SOURCE_DIR}/lib + ${PROJECT_BINARY_DIR}/lib) +add_subdirectory(include) +add_subdirectory(lib) +add_subdirectory(bin) + +add_triton_library(Registrar Registrar.cc) + +if(LLVM_MAJOR_VERSION_21_COMPATIBLE) + set(_MLIRMeshDialect_LIB MLIRMeshDialect) +else() + set(_MLIRMeshDialect_LIB "") +endif() + +add_triton_plugin(TritonAscend + ${CMAKE_CURRENT_SOURCE_DIR}/triton_ascend.cc + ${CMAKE_CURRENT_SOURCE_DIR}/ascend_ir.cc + LINK_LIBS + TritonToLinalgIncubated + BiShengIRScopeDialect + BiShengIRHIVMDialect + ${_MLIRMeshDialect_LIB} +) +# AscendCustomOpsBitcode is built by the TLE DSA subtree; only require it when +# TLE DSA is part of the build. +if(FLAGTREE_TLE_DSA) + add_dependencies(TritonAscend AscendCustomOpsBitcode) +endif() + +target_include_directories(TritonAscend BEFORE PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/python/src) + +option(TRITON_ENABLE_COVERAGE_LLVM_COV "Enable code llvm-cov coverage tool for Ascend plugin " OFF) +if(TRITON_ENABLE_COVERAGE_LLVM_COV) + message(STATUS "Enabling llvm-cov coverage tool flags for TritonAscend") + target_compile_options(TritonAscend PRIVATE + -fprofile-arcs + -ftest-coverage + -O0 + -fprofile-update=atomic + --coverage + ) + target_link_options(TritonAscend PRIVATE + --coverage + -lgcov + ) + # branch coverage + target_compile_definitions(TritonAscend PRIVATE + COVERAGE_ENABLED=1 + ) +endif() + + +# To enable hitest coverage tool +if(TRITON_ENABLE_COVERAGE_HITEST) + set_target_properties(TritonAscend PROPERTIES + RULE_LAUNCH_COMPILE "hitestwrapper" + RULE_LAUNCH_LINK "hitestwrapper" + ) +endif() + +if(TRITON_BUILD_UT) + add_subdirectory(unittest) +endif() diff --git a/third_party/ascend/Registrar.cc b/third_party/ascend/Registrar.cc new file mode 100644 index 0000000000..5298142bbe --- /dev/null +++ b/third_party/ascend/Registrar.cc @@ -0,0 +1,13 @@ +#include "flagtree/Common/UnifiedHardware.h" + +class AscendUnifiedHardware : public mlir::flagtree::UnifiedHardware { +public: + bool getIncubatedTag() const override; +}; + +bool AscendUnifiedHardware::getIncubatedTag() const { return true; } + +std::unique_ptr +mlir::flagtree::createUnifiedHardwareManager() { + return std::make_unique(); +} diff --git a/third_party/ascend/ascend_ir.cc b/third_party/ascend/ascend_ir.cc new file mode 100644 index 0000000000..9eab9c7d8d --- /dev/null +++ b/third_party/ascend/ascend_ir.cc @@ -0,0 +1,820 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * Copyright 2018-2020 Philippe Tillet + * Copyright 2020-2022 OpenAI + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ir.h" +#include "pybind11/pybind11.h" +#include +#include + +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/Scope/IR/Scope.h" + +#include "mlir/AsmParser/AsmParser.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Arith/Utils/Utils.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/AffineExpr.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Types.h" +#include "mlir/Support/LLVM.h" +#include "llvm/IR/Instructions.h" + +using namespace mlir; +namespace py = pybind11; + +struct AscendNPUIROpBuilder : public TritonOpBuilder { + std::string target; + static constexpr char kTarget910_95[] = "Ascend910_95"; + static constexpr char kTarget950[] = "Ascend950"; + + explicit AscendNPUIROpBuilder(MLIRContext *context, std::string target = "") + : TritonOpBuilder(context), target(target) {} + + bool is_910_95() const { + // TODO: Use enum instead of strings after enabling HACC in satandalone + // build + constexpr size_t kLen910 = sizeof(kTarget910_95) - 1; + bool match_910 = target.size() >= kLen910 && + target.compare(0, kLen910, kTarget910_95) == 0; + + constexpr size_t kLen950 = sizeof(kTarget950) - 1; + bool match_950 = + target.size() >= kLen950 && target.compare(0, kLen950, kTarget950) == 0; + + return match_910 || match_950; + } +}; + +namespace { +MLIRContext *gDefaultAscendContext = nullptr; + +MLIRContext *resolveContext(const py::object &contextObj) { + if (!contextObj.is_none()) { + return &py::cast(contextObj); + } + if (gDefaultAscendContext) { + return gDefaultAscendContext; + } + throw std::invalid_argument( + "No default MLIR context. Pass context explicitly or call " + "ascend_ir.load_dialects(context) first."); +} + +struct ModeAndPipes { + hivm::SyncBlockModeAttr modeAttr = {}; + hivm::PipeAttr cubePipe = {}; + hivm::PipeAttr vectorPipe = {}; +}; + +hivm::TCoreTypeAttr GetCore(MLIRContext *ctx, llvm::StringRef opName, + llvm::StringRef sender) { + // Decide core type + hivm::TCoreTypeAttr core; + if (sender == "cube") { + if (opName == "sync_block_set") + core = hivm::TCoreTypeAttr::get(ctx, hivm::TCoreType::CUBE); + else + core = hivm::TCoreTypeAttr::get(ctx, hivm::TCoreType::VECTOR); + } else { + if (sender != "vector") { + throw std::runtime_error( + "sync_block_set/wait only supports 'cube' or 'vector' as sender"); + } + if (opName == "sync_block_set") + core = hivm::TCoreTypeAttr::get(ctx, hivm::TCoreType::VECTOR); + else + core = hivm::TCoreTypeAttr::get(ctx, hivm::TCoreType::CUBE); + } + + return core; +} + +void buildSyncBlockOp(AscendNPUIROpBuilder &self, const std::string &opName, + std::string &sender, std::string &receiver, Value id, + hivm::PIPE senderPipe, hivm::PIPE receiverPipe) { + auto *ctx = self.getBuilder().getContext(); + hivm::TCoreTypeAttr coreAttr = GetCore(ctx, opName, sender); + hivm::PipeAttr prodPipe = hivm::PipeAttr::get(ctx, senderPipe); + hivm::PipeAttr consPipe = hivm::PipeAttr::get(ctx, receiverPipe); + const size_t I64 = 64; + auto i64Ty = IntegerType::get(ctx, I64); + Value idI64 = id; + if (!id.getType().isInteger(I64)) { + idI64 = mlir::convertScalarToDtype(self.getBuilder(), id.getLoc(), id, + i64Ty, true); + } + if (opName == "sync_block_set") { + self.create(coreAttr, prodPipe, consPipe, idI64); + } else if (opName == "sync_block_wait") { + self.create(coreAttr, prodPipe, consPipe, idI64); + } else { + throw std::runtime_error("Unsupported operation name for SyncBlockOp"); + } +} + +ModeAndPipes GetSyncBlockModeAndPipes(MLIRContext *ctx, + const std::string &mode) { + hivm::SyncBlockModeAttr modeAttr = {}; + hivm::PipeAttr cubePipe = {}; + hivm::PipeAttr vectorPipe = {}; + + if (mode == "all_cube") { + modeAttr = hivm::SyncBlockModeAttr::get(ctx, hivm::SyncBlockMode::ALL_CUBE); + cubePipe = hivm::PipeAttr::get(ctx, hivm::PIPE::PIPE_ALL); + vectorPipe = hivm::PipeAttr{}; + } else if (mode == "all_vector") { + modeAttr = + hivm::SyncBlockModeAttr::get(ctx, hivm::SyncBlockMode::ALL_VECTOR); + cubePipe = hivm::PipeAttr{}; + vectorPipe = hivm::PipeAttr::get(ctx, hivm::PIPE::PIPE_ALL); + } else if (mode == "all") { + modeAttr = hivm::SyncBlockModeAttr::get(ctx, hivm::SyncBlockMode::ALL); + cubePipe = hivm::PipeAttr::get(ctx, hivm::PIPE::PIPE_ALL); + vectorPipe = hivm::PipeAttr::get(ctx, hivm::PIPE::PIPE_ALL); + } else if (mode == "all_sub_vector") { + modeAttr = + hivm::SyncBlockModeAttr::get(ctx, hivm::SyncBlockMode::ALL_SUB_VECTOR); + cubePipe = hivm::PipeAttr{}; + vectorPipe = hivm::PipeAttr::get(ctx, hivm::PIPE::PIPE_ALL); + } else { + llvm::report_fatal_error( + llvm::StringRef("Invalid sync-block mode: " + mode)); + } + return {modeAttr, cubePipe, vectorPipe}; +} +} // namespace + +void init_ascend_ir(py::module &&m) { + auto affineExprClass = + py::class_(m, "affine_expr", py::module_local()); + affineExprClass + .def("__str__", + [](AffineExpr self) { + std::string str; + llvm::raw_string_ostream os(str); + self.print(os); + return os.str(); + }) + .def("__repr__", + [](AffineExpr self) { + std::string str; + llvm::raw_string_ostream os(str); + self.print(os); + return ""; + }) + .def("is_symbolic_or_constant", &AffineExpr::isSymbolicOrConstant) + .def("is_pure_affine", &AffineExpr::isPureAffine) + .def("is_function_of_dim", &AffineExpr::isFunctionOfDim) + .def("compose", + [](AffineExpr self, AffineMap map) { return self.compose(map); }) + .def("get_largest_known_divisor", &AffineExpr::getLargestKnownDivisor) + .def("floordiv", [](AffineExpr self, + AffineExpr other) { return self.floorDiv(other); }) + .def("ceildiv", [](AffineExpr self, + AffineExpr other) { return self.ceilDiv(other); }) + .def("mod", + [](AffineExpr self, AffineExpr other) { return self % other; }) + .def("__hash__", + [](AffineExpr self) { + return py::int_(static_cast(mlir::hash_value(self))); + }) + .def("__eq__", [](AffineExpr lhs, AffineExpr rhs) { return lhs == rhs; }) + .def(py::self + py::self) + .def(py::self - py::self) + .def(py::self * py::self) + .def(py::self % py::self); + affineExprClass + .def_static( + "get_constant", + [](int64_t val, py::object contextObj) { + auto *context = resolveContext(contextObj); + return getAffineConstantExpr(val, context); + }, + py::arg("value"), py::arg("context") = py::none()) + .def_static( + "get_dim", + [](uint32_t pos, py::object contextObj) { + auto *context = resolveContext(contextObj); + return getAffineDimExpr(pos, context); + }, + py::arg("pos"), py::arg("context") = py::none()) + .def_static( + "get_symbol", + [](uint32_t pos, py::object contextObj) { + auto *context = resolveContext(contextObj); + return getAffineSymbolExpr(pos, context); + }, + py::arg("pos"), py::arg("context") = py::none()); + + py::class_(m, "affine_constant_expr", + py::module_local()) + .def("get_value", &AffineConstantExpr::getValue); + py::class_(m, "affine_dim_expr", + py::module_local()) + .def("get_position", &AffineDimExpr::getPosition); + py::class_(m, "affine_symbol_expr", + py::module_local()) + .def("get_position", &AffineSymbolExpr::getPosition); + py::class_(m, "affine_binary_op_expr", + py::module_local()) + .def("get_lhs", &AffineBinaryOpExpr::getLHS) + .def("get_rhs", &AffineBinaryOpExpr::getRHS); + + auto affineMapClass = + py::class_(m, "affine_map", py::module_local()); + affineMapClass + .def("__str__", + [](AffineMap &self) { + std::string str; + llvm::raw_string_ostream os(str); + self.print(os); + return os.str(); + }) + .def("__repr__", + [](AffineMap &self) { + std::string str; + llvm::raw_string_ostream os(str); + self.print(os); + return ""; + }) + .def("is_identity", &AffineMap::isIdentity) + .def("is_permutation", &AffineMap::isPermutation) + .def("get_num_dims", &AffineMap::getNumDims) + .def("get_num_symbols", &AffineMap::getNumSymbols) + .def("get_num_results", &AffineMap::getNumResults) + .def("is_empty", &AffineMap::isEmpty) + .def("is_single_constant", &AffineMap::isSingleConstant) + .def("is_constant", &AffineMap::isConstant) + .def("get_constant_result", + [](AffineMap &self) -> int64_t { + if (!self.isSingleConstant()) { + throw std::runtime_error( + "affine map is not a single constant map"); + } + return self.getSingleConstantResult(); + }) + .def("get_result", + [](AffineMap &self, uint32_t pos) { + if (pos >= self.getNumResults()) { + throw py::index_error("result index out of range"); + } + return self.getResult(pos); + }) + .def("get_sub_map", + [](AffineMap &self, const std::vector &resultPos) { + return self.getSubMap(resultPos); + }) + .def("replace", + [](AffineMap &self, AffineExpr expr, AffineExpr replacement, + uint32_t numResultDims, uint32_t numResultSymbols) { + return self.replace(expr, replacement, numResultDims, + numResultSymbols); + }) + .def("compose", + [](AffineMap &self, AffineMap map) { return self.compose(map); }) + .def("get_results", + [](AffineMap &self) -> std::vector { + auto results = self.getResults(); + return std::vector(results.begin(), results.end()); + }) + .def("__hash__", + [](AffineMap &self) { + return py::int_(static_cast(mlir::hash_value(self))); + }) + .def("__eq__", [](AffineMap &lhs, AffineMap &rhs) { return lhs == rhs; }) + .def("inverse_permutation", + [](AffineMap &self) -> py::object { + // Validate it's a permutation first + if (!self.isPermutation()) { + throw py::value_error( + "AffineMap must be a valid permutation to compute inverse"); + } + + // Returns AffineMap directly, not a pointer + AffineMap inverse = mlir::inversePermutation(self); + + // Check if result is valid (null AffineMap) + if (!inverse) { + throw py::value_error("Failed to compute inverse permutation"); + } + + return py::cast(inverse); + }) + .def("to_dict", [](AffineMap &self) -> py::dict { + py::list results; + for (AffineExpr result : self.getResults()) { + if (auto dimExpr = dyn_cast(result)) { + results.append(dimExpr.getPosition()); + } else { + std::string exprStr; + llvm::raw_string_ostream os(exprStr); + result.print(os); + results.append(py::str(exprStr)); + } + } + + py::dict ret; + ret["num_dims"] = self.getNumDims(); + ret["num_symbols"] = self.getNumSymbols(); + ret["results"] = std::move(results); + return ret; + }); + affineMapClass + .def_static( + "get", + [](int64_t numDims, int64_t numSymbols, const py::iterable &resultsIn, + py::object contextObj) -> AffineMap { + MLIRContext *context = nullptr; + if (numDims < 0 || numSymbols < 0) { + throw std::invalid_argument( + "num_dims and num_symbols must be non-negative"); + } + llvm::SmallVector results; + for (const auto &item : resultsIn) { + if (py::isinstance(item)) { + auto expr = py::cast(item); + if (!context) { + context = expr.getContext(); + } + results.push_back(expr); + continue; + } + if (py::isinstance(item)) { + if (!context) { + context = resolveContext(contextObj); + } + int64_t pos = py::cast(item); + if (pos < 0 || pos >= numDims) { + throw std::invalid_argument( + "result dim index is out of range for num_dims"); + } + results.push_back(getAffineDimExpr(pos, context)); + continue; + } + throw std::invalid_argument( + "results must contain affine_expr or int dim indices"); + } + if (!context) { + context = resolveContext(contextObj); + } + return AffineMap::get(numDims, numSymbols, results, context); + }, + py::arg("num_dims"), py::arg("num_symbols"), py::arg("result_dims"), + py::arg("context") = py::none()) + .def_static( + "get_identity", + [](int64_t numDims, py::object contextObj) -> AffineMap { + auto *context = resolveContext(contextObj); + if (numDims < 0) { + throw std::invalid_argument("num_dims must be non-negative"); + } + return AffineMap::getMultiDimIdentityMap(numDims, context); + }, + py::arg("num_dims"), py::arg("context") = py::none()) + .def_static( + "get_minor_identity", + [](int64_t dims, int64_t results, py::object contextObj) { + auto *context = resolveContext(contextObj); + if (dims < 0 || results < 0) { + throw std::invalid_argument("dims/results must be non-negative"); + } + return AffineMap::getMinorIdentityMap(dims, results, context); + }, + py::arg("dims"), py::arg("results"), py::arg("context") = py::none()) + .def_static( + "get_empty", + [](py::object contextObj) { + auto *context = resolveContext(contextObj); + return AffineMap::get(0, 0, {}, context); + }, + py::arg("context") = py::none()) + .def_static( + "get_permutation", + [](const std::vector &permutation, py::object contextObj) { + auto *context = resolveContext(contextObj); + return AffineMap::getPermutationMap(permutation, context); + }, + py::arg("permutation"), py::arg("context") = py::none()) + .def_static( + "get_constant", + [](int64_t value, py::object contextObj) { + auto *context = resolveContext(contextObj); + return AffineMap::getConstantMap(value, context); + }, + py::arg("value"), py::arg("context") = py::none()); + + py::enum_(m, "AddressSpace", py::module_local()) + .value("L1", hivm::AddressSpace::L1) + .value("UB", hivm::AddressSpace::UB) + .value("L0A", hivm::AddressSpace::L0A) + .value("L0B", hivm::AddressSpace::L0B) + .value("L0C", hivm::AddressSpace::L0C) + .export_values(); + + py::enum_(m, "CoreType", py::module_local()) + .value("CUBE", hivm::TCoreType::CUBE) + .value("VECTOR", hivm::TCoreType::VECTOR) + .value("CUBE_OR_VECTOR", hivm::TCoreType::CUBE_OR_VECTOR) + .value("CUBE_AND_VECTOR", hivm::TCoreType::CUBE_AND_VECTOR) + .export_values(); + + py::enum_(m, "PIPE", py::module_local()) + .value("PIPE_S", hivm::PIPE::PIPE_S) + .value("PIPE_V", hivm::PIPE::PIPE_V) + .value("PIPE_M", hivm::PIPE::PIPE_M) + .value("PIPE_MTE1", hivm::PIPE::PIPE_MTE1) + .value("PIPE_MTE2", hivm::PIPE::PIPE_MTE2) + .value("PIPE_MTE3", hivm::PIPE::PIPE_MTE3) + .value("PIPE_ALL", hivm::PIPE::PIPE_ALL) + .value("PIPE_FIX", hivm::PIPE::PIPE_FIX) + .export_values(); + + py::enum_(m, "MODE", py::module_local()) + .value("SIMD", hivm::VFMode::SIMD) + .value("SIMT", hivm::VFMode::SIMT) + .value("MIX", hivm::VFMode::MIX) + .export_values(); + + py::enum_(m, "IteratorType", py::module_local()) + .value("Parallel", hivm::IteratorType::kParallel) + .value("Broadcast", hivm::IteratorType::kBroadcast) + .value("Transpose", hivm::IteratorType::kTranspose) + .value("Reduction", hivm::IteratorType::kReduction) + .value("Interleave", hivm::IteratorType::kInterleave) + .value("Deinterleave", hivm::IteratorType::kDeinterleave) + .value("Inverse", hivm::IteratorType::kInverse) + .value("Pad", hivm::IteratorType::kPad) + .value("Concat", hivm::IteratorType::kConcat) + .value("Gather", hivm::IteratorType::kGather) + .value("Cumulative", hivm::IteratorType::kCumulative) + .value("Opaque", hivm::IteratorType::kOpaque) + .export_values(); + + py::enum_(m, "FixpipeDMAMode", py::module_local()) + .value("NZ2DN", hivm::FixpipeDMAMode::NZ2DN) + .value("NZ2ND", hivm::FixpipeDMAMode::NZ2ND) + .value("NZ2NZ", hivm::FixpipeDMAMode::NZ2NZ) + .export_values(); + + py::enum_(m, "FixpipeDualDstMode", + py::module_local()) + .value("NO_DUAL", hivm::FixpipeDualDstMode::NO_DUAL) + .value("COLUMN_SPLIT", hivm::FixpipeDualDstMode::COLUMN_SPLIT) + .value("ROW_SPLIT", hivm::FixpipeDualDstMode::ROW_SPLIT) + .export_values(); + + py::enum_(m, "FixpipePreQuantMode", + py::module_local()) + .value("NO_QUANT", hivm::FixpipePreQuantMode::NO_QUANT) + .value("F322BF16", hivm::FixpipePreQuantMode::F322BF16) + .value("F322F16", hivm::FixpipePreQuantMode::F322F16) + .value("S322I8", hivm::FixpipePreQuantMode::S322I8) + .export_values(); + + py::enum_(m, "FixpipePreReluMode", + py::module_local()) + .value("LEAKY_RELU", hivm::FixpipePreReluMode::LEAKY_RELU) + .value("NO_RELU", hivm::FixpipePreReluMode::NO_RELU) + .value("NORMAL_RELU", hivm::FixpipePreReluMode::NORMAL_RELU) + .value("P_RELU", hivm::FixpipePreReluMode::P_RELU) + .export_values(); + py::enum_(m, "DataLayout", py::module_local()) + .value("nZ", hivm::DataLayout::nZ) + .value("zN", hivm::DataLayout::zN) + .export_values(); + + m.def("load_dialects", [](MLIRContext &context) { + gDefaultAscendContext = &context; + DialectRegistry registry; + registry.insert(); + context.appendDialectRegistry(registry); + context.loadAllAvailableDialects(); + }); + + py::class_( + m, "ascendnpu_ir_builder", py::module_local(), py::dynamic_attr()) + .def(py::init(), py::arg("context"), + py::arg("target") = "") + .def("get_int_attr", + [](AscendNPUIROpBuilder &self, int64_t value) -> Attribute { + return IntegerAttr::get(self.getBuilder().getI64Type(), value); + }) + .def("get_core_type_attr", + [](AscendNPUIROpBuilder &self, + hivm::TCoreType core_type) -> Attribute { + return self.getBuilder().getAttr(core_type); + }) + .def("get_pipe_attr", + [](AscendNPUIROpBuilder &self, hivm::PIPE pipe) -> Attribute { + return self.getBuilder().getAttr(pipe); + }) + .def("get_vf_mode_attr", + [](AscendNPUIROpBuilder &self, hivm::VFMode mode) -> Attribute { + return self.getBuilder().getAttr(mode); + }) + .def("get_iterator_types_attr", + [](AscendNPUIROpBuilder &self, + const std::vector &array) { + auto attrs = llvm::to_vector( + llvm::map_range(array, [&self](hivm::IteratorType type) { + return cast( + self.getBuilder().getAttr(type)); + })); + return self.getBuilder().getArrayAttr(attrs); + }) + .def("get_t_core_type_attr_name", + [](AscendNPUIROpBuilder &self) -> std::string { + return hivm::TCoreTypeAttr::name.str(); + }) + .def("get_t_core_type_cube_attr", + [](AscendNPUIROpBuilder &self) -> Attribute { + return hivm::TCoreTypeAttr::get(self.getBuilder().getContext(), + hivm::TCoreType::CUBE); + }) + .def("get_t_core_type_vector_attr", + [](AscendNPUIROpBuilder &self) -> Attribute { + return hivm::TCoreTypeAttr::get(self.getBuilder().getContext(), + hivm::TCoreType::VECTOR); + }) + .def("parse_attr", + [](TritonOpBuilder &self, std::string value) -> Attribute { + auto *ctx = self.getBuilder().getContext(); + // Enable parsing of HACC attributes by allowing unregistered + // dialects. + ctx->allowUnregisteredDialects(); + return mlir::parseAttribute(value, ctx); + }) + .def("get_affine_map_attr", + [](AscendNPUIROpBuilder &self, AffineMap affineMap) -> Attribute { + return AffineMapAttr::get(affineMap); + }) + .def("get_affine_map_array_attr", + [](AscendNPUIROpBuilder &self, + const std::vector &affineMaps) -> Attribute { + auto *ctx = self.getBuilder().getContext(); + llvm::SmallVector attrs; + attrs.reserve(affineMaps.size()); + for (const auto &map : affineMaps) { + attrs.push_back(AffineMapAttr::get(map)); + } + return ArrayAttr::get(ctx, attrs); + }) + .def("get_buffer_ty_with_affine_map", + [](AscendNPUIROpBuilder &self, std::vector &shape, + Type &elementType, AffineMap affineMap, + const Attribute &memorySpace) -> Type { + auto layout = AffineMapAttr::get(affineMap); + return MemRefType::get(shape, elementType, layout, memorySpace); + }) + .def("create_fixpipe", + [](AscendNPUIROpBuilder &self, Value src, py::object dst_obj, + hivm::FixpipeDMAMode dma_mode, + hivm::FixpipeDualDstMode dual_dst_mode, + hivm::FixpipePreQuantMode pre_quant_mode, + hivm::FixpipePreReluMode pre_relu_mode) -> py::object { + if (!dyn_cast(src.getType())) { + llvm_unreachable("src is not of RankedTensorType"); + } + + auto *ctx = self.getBuilder().getContext(); + auto loc = self.getLastLoc(); + + Value dstValue; + bool needCreateDst = dst_obj.is_none(); + + if (needCreateDst) { + auto srcType = dyn_cast(src.getType()); + auto srcShape = srcType.getShape(); + + llvm::SmallVector dstShape(srcShape.begin(), + srcShape.end()); + + if (dual_dst_mode == hivm::FixpipeDualDstMode::ROW_SPLIT) { + if (dstShape.size() >= 1 && dstShape[0] > 0) { + dstShape[0] = dstShape[0] / 2; + } + } else if (dual_dst_mode == + hivm::FixpipeDualDstMode::COLUMN_SPLIT) { + if (dstShape.size() >= 2 && dstShape[1] > 0) { + dstShape[1] = dstShape[1] / 2; + } + } + + auto dstType = + RankedTensorType::get(dstShape, srcType.getElementType()); + auto emptyTensor = self.create( + dstType.getShape(), dstType.getElementType()); + dstValue = emptyTensor.getResult(); + } else { + dstValue = py::cast(dst_obj); + if (!dyn_cast(dstValue.getType())) { + llvm_unreachable("dst is not of ShapedType"); + } + } + + auto dma_mode_attr = + mlir::hivm::FixpipeDMAModeAttr::get(ctx, dma_mode); + auto dual_dst_mode_attr = + mlir::hivm::FixpipeDualDstModeAttr::get(ctx, dual_dst_mode); + auto pre_quant_mode_attr = + mlir::hivm::FixpipePreQuantModeAttr::get(ctx, pre_quant_mode); + auto pre_relu_mode_attr = + mlir::hivm::FixpipePreReluModeAttr::get(ctx, pre_relu_mode); + auto channel_split = BoolAttr::get(ctx, false); + if (needCreateDst) { + return py::cast( + self.create( + mlir::TypeRange{dstValue.getType()}, src, dstValue, + dma_mode_attr, dual_dst_mode_attr, + /*sub_block_idx=*/nullptr, pre_quant_mode_attr, + pre_relu_mode_attr, channel_split) + .getResult(0)); + + } else { + self.create(mlir::TypeRange{}, src, dstValue, + dma_mode_attr, dual_dst_mode_attr, + /*sub_block_idx=*/nullptr, + pre_quant_mode_attr, + pre_relu_mode_attr, channel_split); + return py::none(); + } + }) + .def("create_annotation_mark", + [](TritonOpBuilder &self, Value &ptr, const std::string &attrKey, + Attribute &attrVal) { + auto annotationOp = self.create(ptr); + annotationOp->setAttr(self.getBuilder().getStringAttr(attrKey), + attrVal); + }) + .def("create_bind_buffer", + [](TritonOpBuilder &self, Value &src, Value &alloc) -> void { + auto ctx = self.getBuilder().getContext(); + auto bind = StringAttr::get(ctx, "bind_buffer"); + self.create(src, ValueRange{alloc}, + ArrayAttr::get(ctx, bind)); + }) + .def("create_debug_barrier", + [](TritonOpBuilder &self, Value &ptr, const std::string &attrKey, + Attribute &attrVal) { + auto annotationOp = self.create(ptr); + annotationOp->setAttr(self.getBuilder().getStringAttr(attrKey), + attrVal); + }) + .def("create_custom_op", + [](AscendNPUIROpBuilder &self, const std::string &name, + const py::dict &attrs, const std::vector &ins, + const std::vector &outs, + const std::vector &arg_attrs) -> std::vector { + ValueRange inputs{ins}; + ValueRange outputs{outs}; + ValueRange temp_buffers{}; + TypeRange res_types{outputs}; + auto op = self.create(res_types, name, inputs, + outputs, temp_buffers); + for (auto &attr : attrs) { + std::string attr_name = py::cast(attr.first); + Attribute attr_value = py::cast(attr.second); + op->setAttr(attr_name, attr_value); + } + + SmallVector dictAttrs(arg_attrs.size()); + Attribute emptyDict = self.getBuilder().getDictionaryAttr({}); + for (const auto &[idx, attrs] : llvm::enumerate(arg_attrs)) { + if (idx >= op.getNumOperands()) + continue; + + if (attrs.is_none()) { + dictAttrs[idx] = emptyDict; + continue; + } + + llvm::SmallVector namedAttrs; + for (const auto &attr : attrs) { + std::string attr_name = py::cast(attr.first); + Attribute attr_value = py::cast(attr.second); + namedAttrs.push_back(NamedAttribute( + self.getBuilder().getStringAttr(attr_name), attr_value)); + } + + dictAttrs[idx] = self.getBuilder().getDictionaryAttr(namedAttrs); + } + + ArrayAttr arg_attrs_array = + self.getBuilder().getArrayAttr(dictAttrs); + op->setAttr("arg_attrs", arg_attrs_array); + + auto results = op->getResults(); + return std::vector(results.begin(), results.end()); + }) + .def("create_scope_op", + [](AscendNPUIROpBuilder &self, py::dict &scopeAttrs, + std::vector resultTypes) -> OpState { + llvm::SmallVector attrs; + for (auto item : scopeAttrs) { + std::string key = py::cast(item.first); + Attribute value = py::cast(item.second); + attrs.push_back( + NamedAttribute(self.getBuilder().getStringAttr(key), value)); + } + auto scopeOp = self.create(TypeRange(resultTypes)); + scopeOp->setAttrs(attrs); + return OpState(scopeOp); + }) + .def("scope_return", + [](AscendNPUIROpBuilder &self, + std::vector operands) -> OpState { + return self.create(ValueRange(operands)); + }) + .def("sync_block_set", + [](AscendNPUIROpBuilder &self, std::string &sender, + std::string &receiver, Value id, hivm::PIPE senderPipe, + hivm::PIPE receiverPipe) -> void { + buildSyncBlockOp(self, "sync_block_set", sender, receiver, id, + senderPipe, receiverPipe); + }) + .def("sync_block_wait", + [](AscendNPUIROpBuilder &self, std::string &sender, + std::string &receiver, Value id, hivm::PIPE senderPipe, + hivm::PIPE receiverPipe) -> void { + buildSyncBlockOp(self, "sync_block_wait", sender, receiver, id, + senderPipe, receiverPipe); + }) + .def("get_target_attribute", + [](AscendNPUIROpBuilder &self, + hivm::AddressSpace &addressSpace) -> Attribute { + return hivm::AddressSpaceAttr::get(self.getBuilder().getContext(), + addressSpace); + }) + .def("create_get_sub_vec_id", + [](AscendNPUIROpBuilder &self) -> Value { + auto subBlockIdxOp = self.create(); + auto moduleOp = subBlockIdxOp->getParentOfType(); + auto *ctx = self.getBuilder().getContext(); + // If user explicitly uses sub.block idx, add attribute to module. + // NPU compiler will parse this attribute and disable auto tile and + // bind subblock pass. + moduleOp->setAttr("hivm.disable_auto_tile_and_bind_subblock", + mlir::UnitAttr::get(ctx)); + return subBlockIdxOp; + }) + .def("sync_block_all", + [](AscendNPUIROpBuilder &self, std::string &mode, int id) -> void { + auto *ctx = self.getBuilder().getContext(); + auto [modeAttr, cubePipe, vectorPipe] = + GetSyncBlockModeAndPipes(ctx, mode); + mlir::IndexType indexType = mlir::IndexType::get(ctx); + mlir::IntegerAttr indexAttribute = + mlir::IntegerAttr::get(indexType, static_cast(id)); + self.create( + modeAttr, indexAttribute, mlir::Value{}, cubePipe, vectorPipe); + }) + .def("is_910_95", + [](AscendNPUIROpBuilder &self) -> bool { return self.is_910_95(); }) + .def("create_copy_buffer", + [](AscendNPUIROpBuilder &self, Value src, Value dst) { + self.create(mlir::TypeRange{}, src, dst); + }) + .def("create_copy_tensor", + [](AscendNPUIROpBuilder &self, Value src, Value dst) { + return self + .create(mlir::TypeRange{dst.getType()}, src, dst) + .getResult(0); + }) + .def("create_convert_layout", + [](AscendNPUIROpBuilder &self, Value src, Type memrefType) -> Value { + // src is a memref + // the layout is incorrect (temporarily) + auto *ctx = self.getBuilder().getContext(); + return self + .create( + memrefType, src, + hivm::DataLayoutAttr::get(ctx, hivm::DataLayout::ND), + hivm::DataLayoutAttr::get(ctx, hivm::DataLayout::ND)) + .getResult(); + }); +} diff --git a/third_party/ascend/backend/__init__.py b/third_party/ascend/backend/__init__.py new file mode 100644 index 0000000000..fe18edf6b2 --- /dev/null +++ b/third_party/ascend/backend/__init__.py @@ -0,0 +1,54 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +from triton._C.libtriton.ascend import ir as ascend_ir + +from .testing import do_bench_npu + + +def _apply_ascend_patch(): + from triton.compiler.code_generator import CodeGenerator + + if not getattr(CodeGenerator, "_ascend_patch_applied", False): + _original_cg_init = CodeGenerator.__init__ + + def _patched_cg_init(self, *args, **kwargs): + """ + Monkey Patch for Ascend: + Injects 'hacc.target' attribute into the module after initilization of module. + """ + _original_cg_init(self, *args, **kwargs) + options = self.builder.options + context = self.context + if hasattr(options, "arch") and options.arch: + try: + builder = ascend_ir.ascendnpu_ir_builder(context, options.arch) + + target_attr_str = f'#hacc.target<"{options.arch}">' + self.module.set_attr("hacc.target", builder.parse_attr(target_attr_str)) + except Exception as e: + logging.warning(f"[Ascend Patch] Failed to set hacc.target: {e}") + + CodeGenerator.__init__ = _patched_cg_init + CodeGenerator._ascend_patch_applied = True + + +__all__ = ["do_bench_npu"] diff --git a/third_party/ascend/backend/backend_register.py b/third_party/ascend/backend/backend_register.py new file mode 100644 index 0000000000..ca2d39c401 --- /dev/null +++ b/third_party/ascend/backend/backend_register.py @@ -0,0 +1,341 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import os +from typing import Callable, Dict + + +class BackendStrategyRegistry: + + def __init__(self): + self.strategies: Dict[str:Dict[str, Callable]] = {} + + def register(self, category: str, method: str): + + def decorator(func: Callable): + if category not in self.strategies: + self.strategies[category] = {} + if method in self.strategies[category]: + raise ValueError(f"Strategy {method} already registered") + self.strategies[category][method] = func + return func + + return decorator + + def execute_func(self, category, method, *args, **kwargs): + if category not in self.strategies: + raise ValueError(f"Strategy {category} not registered") + if method not in self.strategies[category]: + raise ValueError(f"Strategy {method} not registered") + return self.strategies[category][method](*args, **kwargs) + + def list_categories(self): + return list(self.strategies.keys()) + + def list_methods(self, category): + if category not in self.strategies: + raise ValueError(f"Strategy {category} not registered") + return list(self.strategies[category].keys()) + + +class _LazyBackendStrategyRegister: + + def __init__(self): + self._instance = None + + def _get_instance(self): + if self._instance is None: + self._instance = BackendStrategyRegistry() + return self._instance + + def register(self, *args, **kwargs): + return self._get_instance().register(*args, **kwargs) + + def execute_func(self, *args, **kwargs): + return self._get_instance().execute_func(*args, **kwargs) + + +backend_strategy_registry = _LazyBackendStrategyRegister() + + +@backend_strategy_registry.register("mindspore", "version_hash") +def version_hash(): + import mindspore + return [str(mindspore.version)] + + +@backend_strategy_registry.register("torch_npu", "version_hash") +def version_hash(): + import torch + import torch_npu + return [torch.version.git_version, torch_npu.version.git_version] + + +@backend_strategy_registry.register("mindspore", "cxx_abi") +def get_mindspore_cxx_abi(): + return 0 + + +@backend_strategy_registry.register("torch_npu", "cxx_abi") +def get_torch_cxx_abi(): + import torch + return 1 if torch._C._GLIBCXX_USE_CXX11_ABI else 0 + + +@backend_strategy_registry.register("mindspore", "type_convert") +def type_convert(): + import mindspore + import numpy as np + MINDSPORE_TO_NUMPY_DTYPE = { + mindspore.float32: np.float32, + mindspore.float64: np.float64, + mindspore.float16: np.float16, + mindspore.int8: np.int8, + mindspore.uint8: np.uint8, + mindspore.int16: np.int16, + mindspore.int32: np.int32, + mindspore.int64: np.int64, + mindspore.bool: np.bool_, + mindspore.complex64: np.complex64, + mindspore.complex128: np.complex128, + } + return MINDSPORE_TO_NUMPY_DTYPE + + +@backend_strategy_registry.register("torch_npu", "type_convert") +def type_convert(): + import torch + import numpy as np + TORCH_TO_NUMPY_DTYPE = { + torch.float32: np.float32, + torch.float64: np.float64, + torch.float16: np.float16, + torch.int8: np.int8, + torch.uint8: np.uint8, + torch.int16: np.int16, + torch.int32: np.int32, + torch.int64: np.int64, + torch.bool: np.bool_, + torch.complex64: np.complex64, + torch.complex128: np.complex128, + } + return TORCH_TO_NUMPY_DTYPE + + +@backend_strategy_registry.register("mindspore", "get_device_interface") +def get_device_interface(): + import mindspore + return mindspore + + +@backend_strategy_registry.register("torch_npu", "get_device_interface") +def get_device_interface(): + import torch + return torch.npu + + +@backend_strategy_registry.register("mindspore", "get_empty_tensor") +def get_empty_tensor(size): + import mindspore + return mindspore.mint.empty(size, dtype=mindspore.int32) + + +@backend_strategy_registry.register("torch_npu", "get_empty_tensor") +def get_empty_tensor(size): + import torch + return torch.empty(size, dtype=torch.int32, device='npu') + + +@backend_strategy_registry.register("mindspore", "get_tensor_params_shape") +def get_tensor_params_shape(*args): + import mindspore + tensor_params = [arg for arg in args if isinstance(arg, mindspore.Tensor)] + tensor_params_shape = [] + for t in tensor_params: + tensor_params_shape.append([s for s in t.shape]) + return tensor_params_shape + + +@backend_strategy_registry.register("torch_npu", "get_tensor_params_shape") +def get_tensor_params_shape(*args): + import torch + tensor_params = [arg for arg in args if isinstance(arg, torch.Tensor)] + tensor_params_shape = [] + for t in tensor_params: + tensor_params_shape.append([s for s in t.shape]) + return tensor_params_shape + + +@backend_strategy_registry.register("mindspore", "get_cc_cmd") +def get_cc_cmd(build_pch): + import mindspore + mindspore_path = os.path.dirname(os.path.realpath(mindspore.__file__)) + cc_cmd = [ + f"-I{mindspore_path}", + f"-I{os.path.join(mindspore_path, 'include/')}", + f"-I{os.path.join(mindspore_path, 'include/third_party')}", + f"-I{os.path.join(mindspore_path, 'include/third_party/robin_hood_hashing')}", + f"-I{os.path.join(mindspore_path, 'include/mindspore/core')}", + f"-I{os.path.join(mindspore_path, 'include/mindspore/core/include')}", + f"-I{os.path.join(mindspore_path, 'include/mindspore/core/mindrt/include')}", + f"-I{os.path.join(mindspore_path, 'include/mindspore/ccsrc')}", + f"-I{os.path.join(mindspore_path, 'include/mindspore/ccsrc/include')}", + f"-I{os.path.join(mindspore_path, 'include/mindspore/ops')}", + f"-I{os.path.join(mindspore_path, 'include/mindspore/ops/include')}", + f"-D_GLIBCXX_USE_CXX11_ABI={get_mindspore_cxx_abi()}", + "-DENABLE_FAST_HASH_TABLE=1", + ] + if not build_pch: + cc_cmd += [ + f"-L{os.path.join(mindspore_path, 'lib')}", + f"-lmindspore_pynative_utils", + ] + return cc_cmd + + +@backend_strategy_registry.register("torch_npu", "get_cc_cmd") +def get_cc_cmd(build_pch): + import torch + import torch_npu + torch_path = os.path.dirname(os.path.realpath(torch.__file__)) + torch_npu_path = os.path.dirname(os.path.realpath(torch_npu.__file__)) + cc_cmd = [ + f"-I{os.path.join(torch_path, 'include')}", + f"-I{os.path.join(torch_npu_path, 'include')}", + f"-D_GLIBCXX_USE_CXX11_ABI={get_torch_cxx_abi()}", + ] + if not build_pch: + cc_cmd += [ + f"-L{os.path.join(torch_npu_path, 'lib')}", + f"-ltorch_npu", + ] + return cc_cmd + + +@backend_strategy_registry.register("mindspore", "get_current_device") +def get_current_device(): + import mindspore + return mindspore.get_current_device().device_id + + +@backend_strategy_registry.register("torch_npu", "get_current_device") +def get_current_device(): + import torch + import torch_npu + return torch.npu.current_device() + + +@backend_strategy_registry.register("mindspore", "set_current_device") +def set_current_device(device_id): + import mindspore + return mindspore.set_device("Ascend", device_id) + + +@backend_strategy_registry.register("torch_npu", "set_current_device") +def set_current_device(device_id): + import torch + import torch_npu + return torch.npu.set_device(device_id) + + +@backend_strategy_registry.register("mindspore", "get_current_stream") +def get_current_stream(device): + import mindspore + try: + return mindspore.current_stream().stream_ptr() + except Exception: + return mindspore.current_stream().id + + +@backend_strategy_registry.register("torch_npu", "get_current_stream") +def get_current_stream(device): + import torch + import torch_npu + if device is None: + device = torch.npu.current_device() + if hasattr(torch_npu._C, "_npu_getCurrentRawStreamNoWait"): + from torch_npu._C import _npu_getCurrentRawStreamNoWait + return _npu_getCurrentRawStreamNoWait(device) + else: + from torch_npu._C import _npu_getCurrentRawStream + return _npu_getCurrentRawStream(device) + + +@backend_strategy_registry.register("mindspore", "header_file") +def header_file(enable_taskqueue): + return f'''#include "include/utils/device_manager_conf.h" +#include "include/runtime/hardware_abstract/device_context/device_context_manager.h" +#include "include/mindspore/ops/kernel/ascend/aclnn/pyboost_impl/aclnn_utils.h" +{'#include "include/pynative/utils/runtime/op_executor.h"' if {enable_taskqueue} else ''} +{'#include "include/runtime/pipeline/pipeline.h"' if {enable_taskqueue} else ''}''' + + +@backend_strategy_registry.register("torch_npu", "header_file") +def header_file(enable_taskqueue): + return f'''#include +#include +{'#include ' if {enable_taskqueue} else ''}''' + + +@backend_strategy_registry.register("mindspore", "allocate_memory") +def allocate_memory(size, stream): + return f'''auto work_ptr = std::make_shared(device_context, {size}, reinterpret_cast({stream})); + workspace_addr_ptr = work_ptr->ptr_;''' + + +@backend_strategy_registry.register("torch_npu", "allocate_memory") +def allocate_memory(size, stream): + return f"workspace_addr_ptr = const_cast(at::empty({size}, at::TensorOptions().device(at::kPrivateUse1).dtype(at::kByte)).storage().data());" + + +@backend_strategy_registry.register("mindspore", "allocate_sync_block_lock") +def allocate_sync_block_lock(size, stream): + return f'''auto sync_ptr = std::make_shared(device_context, {size}, reinterpret_cast({stream})); + syncBlockLock_ptr = work_ptr->ptr_;''' + + +@backend_strategy_registry.register("torch_npu", "allocate_sync_block_lock") +def allocate_sync_block_lock(size, stream): + return f"syncBlockLock_ptr = const_cast(at_npu::native::allocate_workspace({size}, {stream}).storage().data());" + + +@backend_strategy_registry.register("mindspore", "pre_launch") +def pre_launch(first_call): + if first_call: + return '''static auto device_context = mindspore::device::DeviceContextManager::GetInstance().GetOrCreateDeviceContext({mindspore::device::DeviceType::kAscend, mindspore::DeviceManagerConf::GetInstance()->device_id()}); + device_context->device_res_manager_->BindDeviceToCurrentThread(false);''' + else: + return '''device_context->device_res_manager_->BindDeviceToCurrentThread(false);''' + + +@backend_strategy_registry.register("torch_npu", "pre_launch") +def pre_launch(first_call): + return "" + + +@backend_strategy_registry.register("mindspore", "async_launch") +def async_launch(func): + return f'''mindspore::runtime::OpExecutor::DispatchLaunchTask({func});''' + + +@backend_strategy_registry.register("torch_npu", "async_launch") +def async_launch(func): + return f'''at_npu::native::OpCommand cmd; + cmd.Name(name.c_str()).SetCustomHandler({func}).Run();''' diff --git a/third_party/ascend/backend/compiler.py b/third_party/ascend/backend/compiler.py new file mode 100644 index 0000000000..7ec2efcef6 --- /dev/null +++ b/third_party/ascend/backend/compiler.py @@ -0,0 +1,1157 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import ctypes +import functools +import hashlib +import glob +import os +import re +import shlex +import subprocess +import tempfile +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType +from typing import Any, Dict, Optional, Tuple, Union + +from flagtree import _flagprism # FlagPrism +from triton._C.libtriton import ir, passes, ascend +from triton.backends.ascend.utils import ( + _check_bishengir_api_change, + _check_bishengir_able_save_ir, + _check_bishengir_is_regbased, + _enable_unpublished_feature, + _enable_print_ub_bits, + _enable_dump_memory_info, + _enable_msdebug, + _get_kernel_target, + _get_llvm_path, + _get_mlir_path, + _get_npucompiler_path, + _get_triton_adapter_opt_path, + _get_triton_mlir_opt_path, + _get_triton_opt_path, + _get_bishengir_opt_path, + _is_ascend_sanitizer_enabled, + _is_debug_line_info_disabled, + _is_auto_map_parallel_blocks_enabled, + downgrade_llir, + force_disable_ffts, + triton_enable_libdevice_simt, + get_cann_version_file_hash, +) +from triton.backends.ascend.driver import (NPUUtils) +from triton.backends.compiler import ( + BaseBackend, + GPUTarget, +) +from triton.runtime.cache import get_dump_manager +from triton.tools.get_ascend_devices import is_compile_on_910_95 + +# commonir: Environment variable to override compile options with a fixed custom set. +# When set to "1", metadata compile options are replaced with predefined values +# that disable most automatic optimizations (useful for debugging/validation). +_USE_CUSTOM_COMPILE_OPT = os.environ.get("USE_CUSTOM_COMPILE_OPT", None) or os.environ.get( + "_USE_CUSTOM_COMPILE_OPT", None) + + +# TODO: materialize the concrete min shape +def min_dot_size(target: GPUTarget): + return lambda lhsType, rhsType: (1, 1, 1) + + +def make_ttir(mod, metadata, opt): + if "hash" not in metadata: + metadata["hash"] = hashlib.sha256(f"{mod}-{metadata}".encode()).hexdigest() + # the same optimize pass for triton-ir as all other backends + pm = ir.pass_manager(mod.context) + pm.enable_debug() + passes.common.add_inliner(pm) + passes.ttir.add_combine(pm) + passes.common.add_canonicalizer(pm) + passes.ttir.add_reorder_broadcast(pm) + ascend.passes.ttir.add_lane_vectorize(pm) + passes.common.add_cse(pm) + # commonir: NOTE: LICM is intentionally omitted — it hoists tile.to_tensor above + # tile.copy in loops, breaking the read-after-write ordering required by + # Ascend's buffer semantics. This causes "operand does not dominate this use" + # errors in bishengir downstream. + # passes.common.add_licm(pm) + passes.common.add_symbol_dce(pm) + passes.ttir.add_loop_unroll(pm) + pm.run(mod) + if opt.debug: + dump_manager = get_dump_manager(metadata["hash"]) + print(f"Dumping intermediate results to {dump_manager.cache_dir}") + dump_manager.put(str(mod), "kernel.ttir.mlir", binary=False) + + return mod + + +def ttir_to_linalg(mod, metadata, opt, *, named_ops=False): + # use triton_adapter to lower Triton-MLIR to linalg + # Get Triton-MLIR as string + ttir_code = str(mod) + with tempfile.TemporaryDirectory() as tmpdir: + src_path = os.path.join(tmpdir, "kernel.ttir.mlir") + dst_path = os.path.join(tmpdir, "kernel.ttadapter.mlir") + Path(src_path).write_text(ttir_code) + triton_adapter_opt_path = _get_triton_adapter_opt_path() + + enable_nd2nz_on_vector = metadata["enable_nd2nz_on_vector"] + enable_select_analysis = metadata["enable_select_analysis"] + compile_on_910_95 = metadata["compile_on_910_95"] + force_simt_template = metadata["force_simt_template"] + enable_sync_block_lock = metadata["enable_sync_block_lock"] + enable_mask_fallback_conversion = metadata["enable_mask_fallback_conversion"] + optimize_dynamic_offset = metadata["optimize_dynamic_offset"] + auto_blockify_size = metadata["auto_blockify_size"] + if not _is_auto_map_parallel_blocks_enabled(): + auto_blockify_size = 1 + pm = ir.pass_manager(mod.context) + pm.enable_debug() + ascend.passes.ttir.add_auto_blockify(pm, auto_blockify_size) + if (metadata["add_auto_scheduling"]): + ascend.passes.ttir.add_dag_sync(pm) + ascend.passes.ttir.add_dag_scope(pm) + passes.common.add_cse(pm) + passes.common.add_canonicalizer(pm) + ascend.passes.ttir.add_dag_ssbuffer(pm) + passes.common.add_cse(pm) + passes.common.add_canonicalizer(pm) + + # commonir: lower tile.* ops and inline unresolved !tile.buf signatures. + ascend.passes.ttir.add_commonir_to_hivm(pm) + passes.common.add_inliner(pm) + passes.common.add_canonicalizer(pm) + + ascend.passes.ttir.add_triton_to_structure(pm, enable_mask_fallback_conversion, optimize_dynamic_offset) + ascend.passes.ttir.add_discrete_mask_access_conversion(pm, compile_on_910_95, force_simt_template, + enable_sync_block_lock) + ascend.passes.ttir.add_triton_to_annotation(pm) + ascend.passes.ttir.add_triton_to_unstructure(pm, compile_on_910_95, force_simt_template) + ascend.passes.ttir.add_triton_to_hivm(pm) + ascend.passes.ttir.add_triton_to_hfusion(pm) + ascend.passes.ttir.add_triton_to_llvm(pm) + ascend.passes.ttir.add_bubble_up_operation(pm) + ascend.passes.ttir.add_triton_to_structure(pm, enable_mask_fallback_conversion, optimize_dynamic_offset) + + ascend.passes.ttir.add_reduce_sum_strength(pm, metadata["enable_reduce_sum_strength"], 2) + ascend.passes.ttir.add_triton_to_linalg(pm, False, named_ops, enable_nd2nz_on_vector, enable_select_analysis, + compile_on_910_95) + + ascend.passes.ttir.add_commonir_to_hivm(pm) + passes.common.add_canonicalizer(pm) + passes.common.add_cse(pm) + passes.common.add_symbol_dce(pm) + + if metadata["enable_dynamic_cv_pipeline"]: + ascend.passes.ttir.add_dynamic_cv_pipeline(pm, compile_on_910_95) + + _val = metadata.get("intra_cache_num") + if _val is not None: + ascend.passes.ttir.set_buffer_count(0, _val) + + _val = metadata.get("inter_cache_num") + if _val is not None: + ascend.passes.ttir.set_buffer_count(1, _val) + + _val = metadata.get("load_cache_num") + if _val is not None: + ascend.passes.ttir.set_buffer_count(2, _val) + + if opt.debug: + # Print the equivalent triton-opt command line so the pass + # pipeline can be reproduced and debugged outside of Python. + cmd = [ + _get_triton_opt_path(), src_path, f"--pass-pipeline={pm.get_pipeline_str()}", "--mlir-print-debuginfo", + "-o", dst_path + ] + print(f"[DEBUG] cmd list: {shlex.join(cmd)}") + + pm.run(mod) + + # FlagPrism: publish the last structured IR before backend serialization. + _flagprism.emit_compiler_event( + phase="pre_backend_serialize", + ir_kind="ttadapter", + module=mod, + metadata=metadata, + ) + + if opt.debug: + dump_manager = get_dump_manager(metadata["hash"]) + dump_manager.put(str(mod), "kernel.ttadapter.mlir", binary=False) + + return str(mod) + + +def linalg_to_bc_by_triton_mlir_opt(linalg: str, metadata, opt): + """ + Convert Linalg IR to MLIR Bytecode format using triton-mlir-opt. + This function supports both MLIR and BishengIR ops. + Args: + linalg: Linalg IR in text format + metadata: Compilation metadata + opt: Compilation options + Returns: + Bytecode data as bytes (not file path, to avoid temp directory cleanup issues) + """ + with tempfile.TemporaryDirectory() as tmpdir: + ttadapter_path = os.path.join(tmpdir, "kernel.ttadapter.mlir") + bc_path = os.path.join(tmpdir, "kernel.mlirbc") + Path(ttadapter_path).write_text(linalg) + + triton_mlir_opt_path = _get_triton_mlir_opt_path() + + # The --emit-bytecode flag ensures output is in BC format. + # + # Pin the bytecode version: triton-mlir-opt is built against this tree's + # LLVM (22) and defaults to bytecode v6, but the consumer (bishengir-opt, + # from CANN, built against LLVM 19) cannot read the location encoding + # used from v5 onward and failed. v4 is the newest version both sides + # agree on once debug locations are present. Override with + # TRITON_ASCEND_BYTECODE_VERSION for bishengir-opt. + bytecode_version = os.environ.get("TRITON_ASCEND_BYTECODE_VERSION", "4") + subprocess.run([ + triton_mlir_opt_path, + ttadapter_path, + "--emit-bytecode", + f"--emit-bytecode-version={bytecode_version}", + "-o", + bc_path, + ], check=True, capture_output=True, text=True) + + # Read bytecode as binary before temp directory is cleaned up + with open(bc_path, "rb") as f: + bc_data = f.read() + + if opt.debug: + dump_manager = get_dump_manager(metadata["hash"]) + dump_manager.put(bc_data, "kernel.mlirbc", binary=True) + + return bc_data + + +def bc_to_linalg_by_bishengir_opt(bc_data: bytes, metadata, opt): + """ + Convert MLIR Bytecode to MLIR text format using bishengir-opt. + Args: + bc_data: Bytecode data as bytes + metadata: Compilation metadata + opt: Compilation options + Returns: + MLIR text as string + """ + with tempfile.TemporaryDirectory() as tmpdir: + bc_path = os.path.join(tmpdir, "kernel.mlirbc") + mlir_path = os.path.join(tmpdir, "kernel.mlir") + + # Write bytecode data to temporary file + with open(bc_path, "wb") as f: + f.write(bc_data) + + bishengir_opt_path, env = _get_bishengir_opt_path() + + subprocess.run([ + bishengir_opt_path, + bc_path, + "--mlir-print-debuginfo", + "-o", + mlir_path, + ], env=env, capture_output=True, check=True, text=True) + + # Read the generated MLIR text + linalg_text = Path(mlir_path).read_text() + + if opt.debug: + dump_manager = get_dump_manager(metadata["hash"]) + dump_manager.put(linalg_text, "kernel.mlir", binary=False) + + return linalg_text + + +def __get_metadata_attr_by_callback(lib, postfix: str, metadata, meta_key: str): + func_symbol = metadata["kernel_name"] + postfix + if hasattr(lib, func_symbol): + callback_func = getattr(lib, func_symbol) + callback_func.restype = ctypes.c_int64 + callback_func.argtypes = [] + metadata[meta_key] = callback_func() + + +def _parse_linalg_metadata(linalg: str, metadata: dict): + """ + Parse Linalg IR to extract metadata required for NPU compilation. + Extracts and updates the following fields in metadata: + - mix_mode + - kernel_name + - tensor_kinds + - shared (currently hardcoded) + - name (kernel_name) + + Additionally, removes the mix_mode attribute from the IR. + """ + # --- Regular expressions and examples --- + + DISABLE_AUTO_TILE_AND_BIND_SUBBLOCK_REGEX = r'hivm.disable_auto_tile_and_bind_subblock' + + # Example: mix_mode = "aiv" -> aiv + MIX_MODE_REGEX = r'mix_mode\s*=\s*"([^"]+)"' + + # Example: parallel_mode = "mix_simd_simt" -> mix_simd_simt + PARALLEL_MODE_REGEX = r'parallel_mode\s*=\s*"([^"]+)"' + + # Example: func.func @gather_sorted_kernel(%arg0: ...) -> gather_sorted_kernel + KERNEL_NAME_REGEX = r"func\.func\s+@(\w+)" + + # Example: %arg1: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32} -> ('1', '0') + TENSOR_KIND_REGEX = r'%arg(\d+):[^,)]*?\{[^}]*?tt\.tensor_kind\s*=\s*([^:\s}]+)\s*:[^}]*?\}' + + # Example: bitcode = "a.bc" + BITCODES_REGEX = r'bitcode\s*=\s*(?:"([^"]+)"|\'([^\']+)\'|(\w+))' + + # Note: Compiled Kernel requires to estimate size of shared memory to occupy + # Currently, NPU backend does not limit on shared memory + metadata["shared"] = 1 + # Force disable auto tile and bind subblock if attribute is present in module + metadata["auto_tile_and_bind_subblock"] = not re.search(DISABLE_AUTO_TILE_AND_BIND_SUBBLOCK_REGEX, linalg) + metadata["mix_mode"] = re.search(MIX_MODE_REGEX, linalg).group(1) + metadata["parallel_mode"] = re.search(PARALLEL_MODE_REGEX, linalg).group(1) + metadata["kernel_name"] = re.search(KERNEL_NAME_REGEX, linalg).group(1) + metadata["name"] = metadata["kernel_name"] + # Parse all tensor kinds from arguments + metadata["tensor_kinds"] = [int(kind) for _, kind in re.findall(TENSOR_KIND_REGEX, linalg)] + # init the ub bits of triton kernel for inductor autotune using + metadata["required_ub_bits"] = 0 + + # Parse all bitcode paths + bitcodes = re.findall(BITCODES_REGEX, linalg) + metadata["bitcodes"] = [val for group in bitcodes for val in group if val] + return linalg, metadata + + +def _parse_ttir_metadata(ttir: str, metadata: dict): + """ + Parse TTIR to extract metadata required for NPU compilation. + Extracts and updates the following fields in metadata: + - kernel_name + - shared (currently hardcoded) + """ + # --- Regular expressions and examples --- + # Example: tt.func @gather_sorted_kernel(%arg0: ...) -> gather_sorted_kernel + KERNEL_NAME_REGEX = r"tt\.func\spublic\s+@(\w+)" + + # Example: %arg1: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32} -> ('1', '0') + TENSOR_KIND_REGEX = r'%arg(\d+):[^,)]*?\{[^}]*?tt\.tensor_kind\s*=\s*([^:\s}]+)\s*:[^}]*?\}' + + # Note: Compiled Kernel requires to estimate size of shared memory to occupy + # Currently, NPU backend does not limit on shared memory + metadata["shared"] = 1 + # Note: Currently, for TTIR inputs, we only support vector kernels. + metadata["mix_mode"] = "aiv" + metadata["kernel_name"] = re.search(KERNEL_NAME_REGEX, ttir).group(1) + metadata["name"] = metadata["kernel_name"] + # Parse all tensor kinds from arguments + metadata["tensor_kinds"] = [int(kind) for _, kind in re.findall(TENSOR_KIND_REGEX, ttir)] + return metadata + + +def get_common_bishengir_compile_options(metadata): + bishengir_target = metadata['target'].arch + bishengir_target_opt = f"--target={bishengir_target}" + return [bishengir_target_opt] + + +def get_auto_bind_sub_block_option(metadata): + # auto_tile_and_bind_subblock is read from the module. + # enable_auto_bind_sub_block is set by the user and has a higher priority. + enable_auto_bind_sub_block = metadata["enable_auto_bind_sub_block"] + return (metadata["auto_tile_and_bind_subblock"] + if enable_auto_bind_sub_block is None else enable_auto_bind_sub_block) + + +def _save_npuir_debug_output(stdout_bytes: bytes, stderr_bytes: bytes, tmpdir: str, metadata_hash: str): + stdout = stdout_bytes.decode('utf-8') if stdout_bytes else '' + stderr = stderr_bytes.decode('utf-8') if stderr_bytes else '' + combined = stdout + stderr + if not combined.strip(): + combined = "No output captured." + output_path = os.path.join(tmpdir, "kernel.npuir.mlir") + with open(output_path, 'w', encoding='utf-8') as f: + f.write(combined) + + dump_manager = get_dump_manager(metadata_hash) + dump_manager.put(Path(output_path).read_text(encoding='utf-8'), "kernel.npuir.mlir", binary=False) + + +def linalg_to_bin_enable_npu_compile_910_95(linalg: str, metadata, opt): + linalg, metadata = _parse_linalg_metadata(linalg, metadata) + with tempfile.TemporaryDirectory() as tmpdir: + tmp_file_name = "kernel.mlir" if opt.use_bytecode else "kernel.ttadapter.mlir" + ttadapter_path = os.path.join(tmpdir, tmp_file_name) + Path(ttadapter_path).write_text(linalg) + bin_file = os.path.join(tmpdir, "kernel") + if _check_bishengir_api_change(): + bin_file_with_ext = "kernel.o" + else: + bin_file_with_ext = "kernel_reloc.o" + bin_path = os.path.join(tmpdir, bin_file_with_ext) + callback_path = os.path.join(tmpdir, "libkernel.so") + _compile_option_list = get_common_bishengir_compile_options(metadata) + + multibuffer = metadata.get("multibuffer") + num_stages = metadata.get("num_stages") + + if multibuffer is not None or num_stages is not None: + multi_buffer_value = True + if multibuffer is not None and not multibuffer: + multi_buffer_value = False + elif num_stages is not None and num_stages == 1: + multi_buffer_value = False + _compile_option_list += [ + f"--enable-auto-multi-buffer={multi_buffer_value}", + ] + + disable_tightly_coupled_buffer_reuse = metadata["disable_tightly_coupled_buffer_reuse"] + if disable_tightly_coupled_buffer_reuse: + _compile_option_list += ["--disable-tightly-coupled-buffer-reuse"] + + _compile_option_list += [ + f"--enable-auto-bind-sub-block={get_auto_bind_sub_block_option(metadata)}", + ] + + if force_disable_ffts(): + _compile_option_list += ["--disable-ffts"] + if _is_ascend_sanitizer_enabled(): + _compile_option_list += ["--enable-sanitizer=true"] + if not _is_debug_line_info_disabled(): + _compile_option_list += ["--enable-debug-info=true"] + + if _enable_print_ub_bits(): + _compile_option_list += ["--enable-print-memory-allocated-size"] + + enable_hivm_auto_cv_balance = metadata["enable_hivm_auto_cv_balance"] + if enable_hivm_auto_cv_balance is not None: + _compile_option_list += \ + [f"--enable-hivm-auto-cv-balance={enable_hivm_auto_cv_balance}"] + + sync_solver = metadata["sync_solver"] + if sync_solver is not None: + _compile_option_list += \ + [f"--enable-hivm-graph-sync-solver={sync_solver}"] + + unit_flag = metadata["unit_flag"] + if unit_flag is not None: + _compile_option_list += \ + [f"--enable-hivm-unit-flag-sync={unit_flag}"] + + inject_barrier_all = metadata["inject_barrier_all"] + if inject_barrier_all is not None: + _compile_option_list += \ + [f"--enable-hivm-inject-barrier-all-sync={inject_barrier_all}"] + + inject_block_all = metadata["inject_block_all"] + if inject_block_all is not None: + _compile_option_list += \ + [f"--enable-hivm-inject-block-all-sync={inject_block_all}"] + + limit_auto_multi_buffer_only_for_local_buffer = metadata["limit_auto_multi_buffer_only_for_local_buffer"] + if limit_auto_multi_buffer_only_for_local_buffer is not None: + _compile_option_list += \ + [f"--limit-auto-multi-buffer-only-for-local-buffer={limit_auto_multi_buffer_only_for_local_buffer}"] + + set_workspace_multibuffer = metadata["set_workspace_multibuffer"] + if set_workspace_multibuffer is not None: + _compile_option_list += \ + [f"--set-workspace-multibuffer={set_workspace_multibuffer}"] + + auto_multi_buffer = metadata["limit_auto_multi_buffer_of_local_buffer"] + if auto_multi_buffer is not None: + _compile_option_list += \ + [f"--limit-auto-multi-buffer-of-local-buffer={auto_multi_buffer}"] + + enable_mixed_cv = metadata["enable_mixed_cv"] + if enable_mixed_cv is not None: + _compile_option_list += \ + [f"--enable-mixed-cv={enable_mixed_cv}"] + + enable_cce_vf_auto_sync = metadata["enable_cce_vf_auto_sync"] + if enable_cce_vf_auto_sync is not None: + _compile_option_list += \ + [f"--append-bisheng-options=-mllvm --cce-vf-auto-sync={enable_cce_vf_auto_sync}"] + + enable_cce_vf_remove_membar = metadata["enable_cce_vf_remove_membar"] + if enable_cce_vf_remove_membar is not None: + _compile_option_list += \ + [f"--append-bisheng-options=-mllvm --cce-vf-remove-membar={enable_cce_vf_remove_membar}"] + + # TEMP: Allow TRITON_ENABLE_VF_FUSION env var to override metadata configuration + env_vf = os.getenv("TRITON_ENABLE_VF_FUSION") + enable_vf_fusion = (env_vf.lower() in ("true", "1", "yes") if env_vf is not None else metadata.get( + "enable_vf_fusion", False)) + if enable_vf_fusion: + _compile_option_list += ["--enable-vf-fusion"] + + enable_drop_unit_dims = metadata["enable_drop_unit_dims"] + if enable_drop_unit_dims is not None: + _compile_option_list += \ + [f"--enable-drop-unit-dims={enable_drop_unit_dims}"] + + enable_flatten = metadata["enable_flatten"] + if enable_flatten is not None: + _compile_option_list += \ + [f"--enable-flatten={enable_flatten}"] + + enable_auto_vectorize_v2 = metadata["enable_auto_vectorize_v2"] + if enable_auto_vectorize_v2 is not None: + _compile_option_list += \ + [f"--enable-auto-vectorize-v2={enable_auto_vectorize_v2}"] + auto_vectorize_v2_max_fused_ops_num = metadata["auto_vectorize_v2_max_fused_ops_num"] + if auto_vectorize_v2_max_fused_ops_num is not None: + _compile_option_list += \ + [f"--hfusion-max-fused-ops-in-auto-vectorize-v2={auto_vectorize_v2_max_fused_ops_num}"] + prevec_max_fused_ops_num = metadata["prevec_max_fused_ops_num"] + if prevec_max_fused_ops_num is not None: + _compile_option_list += \ + [f"--hfusion-max-fused-elementwise-ops={prevec_max_fused_ops_num}"] + + disable_auto_inject_block_sync = metadata["disable_auto_inject_block_sync"] + if disable_auto_inject_block_sync is not None: + _compile_option_list += \ + [f"--disable-auto-inject-block-sync={disable_auto_inject_block_sync}"] + + bitcodes = metadata["bitcodes"] + if bitcodes is not None: + for bitcode in bitcodes: + _compile_option_list += \ + [f"--link-aicore-bitcode={bitcode}"] + + disable_auto_cv_work_space_manage = metadata.get("disable_auto_cv_work_space_manage") + if disable_auto_cv_work_space_manage is not None: + _compile_option_list += \ + [f"--disable-auto-cv-work-space-manage={disable_auto_cv_work_space_manage}"] + + enable_legacy_insert_load_store_for_mix_cv = metadata.get("enable_legacy_insert_load_store_for_mix_cv") + if enable_legacy_insert_load_store_for_mix_cv is not None: + _compile_option_list += \ + [f"--enable-legacy-insert-load-store-for-mix-cv={enable_legacy_insert_load_store_for_mix_cv}"] + + enable_tuning_mode = metadata.get("enable_tuning_mode") + if enable_tuning_mode is not None: + _compile_option_list += \ + [f"--enable-tuning-mode={enable_tuning_mode}"] + + enable_auto_blockify = metadata["enable_auto_blockify"] + if _is_auto_map_parallel_blocks_enabled(): + if (enable_auto_blockify is None or enable_auto_blockify): + _compile_option_list += ["--enable-auto-blockify-loop"] + else: + if enable_auto_blockify: + _compile_option_list += ["--enable-auto-blockify-loop"] + + npu_compiler_path, env = _get_npucompiler_path() + if npu_compiler_path.endswith("bishengir-compile"): + _compile_option_list += [ + "--enable-hfusion-compile=true", + "--enable-triton-kernel-compile=true", + ] + bisheng_options = metadata["bisheng_options"] + if bisheng_options is not None: + _compile_option_list += [f"--append-bisheng-options={bisheng_options}"] + mix_mode = opt.mix_mode + if mix_mode in ["aic"]: + _compile_option_list += ["--disable-hfusion-vectorize=true"] + + if opt.debug: + _compile_option_list += ["--bishengir-print-ir-after=hivm-graph-sync-solver"] + + cmd_list = ([npu_compiler_path, ttadapter_path] + _compile_option_list + ["-o", bin_file]) + vf_merge_level = metadata["vf_merge_level"] + if vf_merge_level is not None and vf_merge_level != 1: + cmd_list += [f"--enable-vf-merge-level={vf_merge_level}"] + + hfusion_enable_multiple_consumer_fusion = metadata["hfusion_enable_multiple_consumer_fusion"] + if hfusion_enable_multiple_consumer_fusion: + cmd_list += [f"--hfusion-enable-multiple-consumer-fusion={hfusion_enable_multiple_consumer_fusion}"] + + if opt.debug: + print(f"[DEBUG] cmd_list: {' '.join(cmd_list)}") + + try: + ret = subprocess.run(cmd_list, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True) + except subprocess.CalledProcessError as e: + if opt.debug: + _save_npuir_debug_output(e.stdout, e.stderr, tmpdir, metadata["hash"]) + raise + + if opt.debug: + _save_npuir_debug_output(ret.stdout, ret.stderr, tmpdir, metadata["hash"]) + + stdout_str = ret.stdout.decode('utf-8') if ret.stdout else '' + match = re.search(r'UB\s+size\s*=\s*(\d+)\s*bits', stdout_str) + if match: + # get the ub bits of triton kernel from bisheng for inductor autotune using + metadata["required_ub_bits"] = int(match.group(1)) + + if not Path(bin_path).exists(): + error_msg = ret.stderr.decode('utf-8') if ret.stderr else '' + print(f"[DEBUG] {bin_path} is not found") + print(f"[DEBUG] Stderr:\n{error_msg}") + raise subprocess.CalledProcessError(ret.returncode, cmd_list, ret.stdout, ret.stderr) + + if Path(callback_path).is_file(): + lib = ctypes.CDLL(callback_path) + __get_metadata_attr_by_callback(lib, "_infer_task_type_function", metadata, "bs_task_type") + __get_metadata_attr_by_callback(lib, "_infer_workspace_shape_function", metadata, "workspace_size") + __get_metadata_attr_by_callback(lib, "_infer_sync_block_lock_num_function", metadata, "lock_num") + __get_metadata_attr_by_callback(lib, "_infer_sync_block_lock_init_function", metadata, "lock_init_val") + + return Path(bin_path).read_bytes() + + +def _compile_linalg_to_npu_bin(linalg: str, metadata, opt): + """Compile linalg IR to NPU binary with optional external IR replacement. + + When the environment variable TLE_REPLACE_IR_FILE is set, replaces the + intermediate linalg IR with the content of the specified file before + compilation. This enables injecting externally generated IR into the + compilation pipeline for debugging and validation. + + Delegates to the appropriate platform-specific compilation function + (910_95 or A2_A3) after the replacement. + """ + # commonir: disable several passes of bisheng compiler + # as we want to lift cv pipeline to kernel level + if _USE_CUSTOM_COMPILE_OPT is not None: + metadata["num_stages"] = 1 + metadata["multibuffer"] = False + metadata["enable_tuning_mode"] = True + metadata["enable_ubuf_saving"] = None + metadata["unit_flag"] = False + metadata["enable_auto_bind_sub_block"] = True + metadata["enable_hivm_auto_cv_balance"] = False + metadata["limit_auto_multi_buffer_only_for_local_buffer"] = False + metadata["disable_auto_inject_block_sync"] = True + metadata["disable_auto_cv_work_space_manage"] = True + + if opt.compile_on_910_95: + return linalg_to_bin_enable_npu_compile_910_95(linalg, metadata, opt) + else: + return linalg_to_bin_enable_npu_compile_A2_A3(linalg, metadata, opt) + + +def linalg_to_bin_enable_npu_compile_A2_A3(linalg: str, metadata, opt): + linalg, metadata = _parse_linalg_metadata(linalg, metadata) + with tempfile.TemporaryDirectory() as tmpdir: + tmp_file_name = "kernel.mlir" if opt.use_bytecode else "kernel.ttadapter.mlir" + ttadapter_path = os.path.join(tmpdir, tmp_file_name) + Path(ttadapter_path).write_text(linalg) + bin_file = os.path.join(tmpdir, "kernel") + if _check_bishengir_api_change(): + bin_file_with_ext = "kernel.o" + else: + bin_file_with_ext = "kernel_reloc.o" + if _check_bishengir_is_regbased(): + bishengir_hivm_opt = "--reg-based=true" + else: + bishengir_hivm_opt = "--enable-hivm-compile=true" + bin_path = os.path.join(tmpdir, bin_file_with_ext) + callback_path = os.path.join(tmpdir, "libkernel.so") + _compile_option_list = [ + f"--target={NPUUtils().get_arch()}", + ] + + multibuffer = metadata.get("multibuffer") + num_stages = metadata.get("num_stages") + + if multibuffer is not None or num_stages is not None: + multi_buffer_value = True + if multibuffer is not None and not multibuffer: + multi_buffer_value = False + elif num_stages is not None and num_stages == 1: + multi_buffer_value = False + _compile_option_list.append(f"--enable-auto-multi-buffer={multi_buffer_value}") + + enable_ubuf_saving = metadata["enable_ubuf_saving"] + if enable_ubuf_saving is not None: + _compile_option_list += [ + f"--enable-ubuf-saving={enable_ubuf_saving}", + ] + + enable_preload = metadata["enable_preload"] + if enable_preload is not None: + _compile_option_list += [ + f"--enable-preload={enable_preload}", + ] + + _compile_option_list += [ + f"--enable-auto-bind-sub-block={get_auto_bind_sub_block_option(metadata)}", + ] + + if _is_ascend_sanitizer_enabled(): + _compile_option_list += ["--enable-sanitizer=true"] + if not _is_debug_line_info_disabled(): + _compile_option_list += ["--enable-debug-info=true"] + + if _enable_print_ub_bits(): + _compile_option_list += ["--enable-print-memory-allocated-size"] + + if _enable_dump_memory_info(): + _compile_option_list += ["--enable-memory-display=true"] + + if _enable_msdebug(): + _compile_option_list += ["--enable-ms-debug=true"] + + enable_hivm_auto_cv_balance = metadata["enable_hivm_auto_cv_balance"] + if enable_hivm_auto_cv_balance is not None: + _compile_option_list += \ + [f"--enable-hivm-auto-cv-balance={enable_hivm_auto_cv_balance}"] + + sync_solver = metadata["sync_solver"] + if sync_solver is not None: + _compile_option_list += [ + f"--enable-hivm-graph-sync-solver={sync_solver}", + f"--enable-hivm-cross-core-gss={sync_solver}", + ] + + unit_flag = metadata["unit_flag"] + if unit_flag is not None: + _compile_option_list += \ + [f"--enable-hivm-unit-flag-sync={unit_flag}"] + + enable_drop_unit_dims = metadata["enable_drop_unit_dims"] + if enable_drop_unit_dims is not None: + _compile_option_list += \ + [f"--enable-drop-unit-dims={enable_drop_unit_dims}"] + + enable_flatten = metadata["enable_flatten"] + if enable_flatten is not None: + _compile_option_list += \ + [f"--enable-flatten={enable_flatten}"] + + enable_auto_vectorize_v2 = metadata["enable_auto_vectorize_v2"] + if enable_auto_vectorize_v2 is not None: + _compile_option_list += \ + [f"--enable-auto-vectorize-v2={enable_auto_vectorize_v2}"] + + inject_barrier_all = metadata["inject_barrier_all"] + if inject_barrier_all is not None: + _compile_option_list += \ + [f"--enable-hivm-inject-barrier-all-sync={inject_barrier_all}"] + + inject_block_all = metadata["inject_block_all"] + if inject_block_all is not None: + _compile_option_list += \ + [f"--enable-hivm-inject-block-all-sync={inject_block_all}"] + + limit_auto_multi_buffer_only_for_local_buffer = metadata["limit_auto_multi_buffer_only_for_local_buffer"] + if limit_auto_multi_buffer_only_for_local_buffer is not None: + _compile_option_list += \ + [f"--limit-auto-multi-buffer-only-for-local-buffer={limit_auto_multi_buffer_only_for_local_buffer}"] + + set_workspace_multibuffer = metadata["set_workspace_multibuffer"] + if set_workspace_multibuffer is not None: + _compile_option_list += \ + [f"--set-workspace-multibuffer={set_workspace_multibuffer}"] + + tile_mix_vector_loop = metadata["tile_mix_vector_loop"] + if tile_mix_vector_loop is not None: + _compile_option_list += \ + [f"--tile-mix-vector-loop={tile_mix_vector_loop}"] + + tile_mix_cube_loop = metadata["tile_mix_cube_loop"] + if tile_mix_cube_loop is not None: + _compile_option_list += \ + [f"--tile-mix-cube-loop={tile_mix_cube_loop}"] + + auto_multi_buffer = metadata["limit_auto_multi_buffer_of_local_buffer"] + if auto_multi_buffer is not None: + _compile_option_list += \ + [f"--limit-auto-multi-buffer-of-local-buffer={auto_multi_buffer}"] + + disable_auto_inject_block_sync = metadata["disable_auto_inject_block_sync"] + if disable_auto_inject_block_sync is not None: + _compile_option_list += \ + [f"--disable-auto-inject-block-sync={disable_auto_inject_block_sync}"] + + bitcodes = metadata["bitcodes"] + if bitcodes is not None: + for bitcode in bitcodes: + _compile_option_list += \ + [f"--link-aicore-bitcode={bitcode}"] + + _compile_option_list += [f"--link-aicore-bitcode={get_libdevice()}"] + + disable_size_align_for_cast = metadata["disable_size_align_for_cast"] + if disable_size_align_for_cast is not None: + _compile_option_list += \ + [f"--disable-size-align-for-cast={disable_size_align_for_cast}"] + + disable_auto_cv_work_space_manage = metadata.get("disable_auto_cv_work_space_manage") + if disable_auto_cv_work_space_manage is not None: + _compile_option_list += \ + [f"--disable-auto-cv-work-space-manage={disable_auto_cv_work_space_manage}"] + enable_legacy_insert_load_store_for_mix_cv = metadata.get("enable_legacy_insert_load_store_for_mix_cv") + if enable_legacy_insert_load_store_for_mix_cv is not None: + _compile_option_list += \ + [f"--enable-legacy-insert-load-store-for-mix-cv={enable_legacy_insert_load_store_for_mix_cv}"] + enable_tuning_mode = metadata.get("enable_tuning_mode") + if enable_tuning_mode is not None: + _compile_option_list += \ + [f"--enable-tuning-mode={enable_tuning_mode}"] + + if _is_auto_map_parallel_blocks_enabled(): + _compile_option_list += ["--enable-auto-blockify-loop"] + npu_compiler_path, env = _get_npucompiler_path() + if npu_compiler_path.endswith("bishengir-compile"): + _compile_option_list += [ + "--enable-hfusion-compile=true", + bishengir_hivm_opt, + "--enable-triton-kernel-compile=true", + ] + + if opt.debug: + _compile_option_list += ["--mlir-print-ir-after-failure"] + _compile_option_list += ["--bishengir-print-ir-after=hivm-graph-sync-solver"] + cmd_list = ([npu_compiler_path, ttadapter_path] + _compile_option_list + ["-o", bin_file]) + if opt.debug: + print(f"[DEBUG] cmd_list: {' '.join(cmd_list)}") + + import shutil + shutil.copy2(ttadapter_path, "/tmp/debug_kernel_input.mlir") + + try: + ret = subprocess.run(cmd_list, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True) + except subprocess.CalledProcessError as e: + if opt.debug: + _save_npuir_debug_output(e.stdout, e.stderr, tmpdir, metadata["hash"]) + raise + + if opt.debug: + _save_npuir_debug_output(ret.stdout, ret.stderr, tmpdir, metadata["hash"]) + + if "--mlir-print-ir-after-all" in _compile_option_list: + dump_path = os.environ.get("BISHENGIR_DUMP_PATH", os.path.join(tmpdir, "bishengir_pass_dump.log")) + with open(dump_path, "w") as f: + f.write(ret.stderr.decode("utf-8", errors="replace")) + print(f"[bishengir] Pass IR dump written to: {dump_path}") + + stdout_str = ret.stdout.decode('utf-8') if ret.stdout else '' + match = re.search(r'UB\s+size\s*=\s*(\d+)\s*bits', stdout_str) + if match: + metadata["required_ub_bits"] = int(match.group(1)) + + if not Path(bin_path).exists(): + error_msg = ret.stderr.decode('utf-8') if ret.stderr else '' + print(f"[DEBUG] {bin_path} is not found") + print(f"[DEBUG] Stderr:\n{error_msg}") + raise subprocess.CalledProcessError(ret.returncode, cmd_list, ret.stdout, ret.stderr) + + if Path(callback_path).is_file(): + lib = ctypes.CDLL(callback_path) + __get_metadata_attr_by_callback(lib, "_infer_task_type_function", metadata, "bs_task_type") + __get_metadata_attr_by_callback(lib, "_infer_workspace_shape_function", metadata, "workspace_size") + __get_metadata_attr_by_callback(lib, "_infer_sync_block_lock_num_function", metadata, "lock_num") + __get_metadata_attr_by_callback(lib, "_infer_sync_block_lock_init_function", metadata, "lock_init_val") + + return Path(bin_path).read_bytes() + + +def get_libdevice(): + current = os.path.dirname(__file__) + return os.path.join(current, "lib/libdevice.10.bc") + + +@dataclass(frozen=True) +class NPUOptions: + debug: bool = False + # FlagPrism: instrumentation mode participates in the option hash. + instrumentation_mode: str = "" + sanitize_overflow: bool = True + llvm_version: int = 15 + kernel_name: str = "triton_" + arch: str = "" + + cluster_dims: tuple = (1, 1, 1) + num_warps: int = 32 + num_ctas: int = 1 + num_stages: int = 1 if is_compile_on_910_95 else 2 + warp_size: int = 32 + num_buffers_warp_spec: int = 0 + num_consumer_groups: int = 0 + reg_dec_producer: int = 0 + reg_inc_consumer: int = 0 + + auto_blockify_size: int = 1 + enable_auto_blockify: bool = None + compile_on_910_95: bool = is_compile_on_910_95 + optimize_dynamic_offset: bool = False + enable_mask_fallback_conversion: bool = False + enable_warp_specialization: bool = False + enable_nd2nz_on_vector: bool = False + enable_persistent: bool = False + optimize_epilogue: bool = False + enable_fp_fusion: bool = True + # Reduce sum strength: rewrites an add-reduce (e.g. tl.sum) on a + # large power-of-two reduce dim into "pre-add the two halves with vector + # adds, then reduce the halved tensor". This reorders floating-point + # addition, so at low precision (f16/bf16, ~10/8 mantissa bits) the reduce + # result may differ from the original summation order by 1~2 ulp. Enable + # only when such rounding-order changes are acceptable. Off by default. + enable_reduce_sum_strength: bool = False + allow_fp8e4nv: bool = False + auto_tile_and_bind_subblock: bool = True + supported_fp8_dtypes: Tuple[str] = ("fp8e5", "fp8e4b15", "fp8e4nv", "fp8e4b8", "fp8e5b16") + deprecated_fp8_dtypes: Tuple[str] = () + vf_merge_level: int = 1 + default_dot_input_precision: str = "ieee" + allowed_dot_input_precisions: Tuple[str] = ("ieee", "hf32") + max_num_imprecise_acc_default: int = 0 + extern_libs: dict = None + bisheng_options: str = "-cce-link-aicore-ll-module " + get_libdevice() + + multibuffer: bool = not is_compile_on_910_95 + enable_ubuf_saving: bool = None + enable_preload: bool = None + enable_auto_bind_sub_block: bool = None + disable_tightly_coupled_buffer_reuse: bool = False + disable_auto_cv_work_space_manage: bool = None + enable_legacy_insert_load_store_for_mix_cv: bool = None + enable_select_analysis: bool = True + enable_hivm_auto_cv_balance: bool = None + sync_solver: bool = None + unit_flag: bool = None + enable_cce_vf_auto_sync: bool = None + enable_cce_vf_remove_membar: bool = None + enable_drop_unit_dims: bool = None + enable_flatten: bool = None + enable_auto_vectorize_v2: bool = None + auto_vectorize_v2_max_fused_ops_num: int = None + prevec_max_fused_ops_num: int = None + inject_barrier_all: bool = None + inject_block_all: bool = None + disable_size_align_for_cast: bool = None + limit_auto_multi_buffer_only_for_local_buffer: bool = None + limit_auto_multi_buffer_of_local_buffer: str = None + set_workspace_multibuffer: int = None + tile_mix_vector_loop: int = None + tile_mix_cube_loop: int = None + disable_auto_inject_block_sync: bool = None + enable_mixed_cv: bool = None + enable_vf_fusion: bool = False + # todo: this code will be removed in version 530. + add_auto_scheduling: bool = False + enable_dynamic_cv_pipeline: bool = False + hfusion_enable_multiple_consumer_fusion: bool = False + + stream: int = None + parallel_mode: str = "simd" + force_simt_only: bool = False + force_simt_template: bool = False + enable_sync_block_lock: bool = False + # only take effect on the simt-only & simd-simt-mix scenarios + shared_mem_dynamic_size: int = None + # enable_bishengir_simt_optimization is passed as + # -enable-bishengir-simt-optimization flag to bishengir-compile. + enable_bishengir_simt_optimization: int = 000 + # compile_mode: "simd" (default), "unstructured_in_simt", "simt_only" + # When compile_mode is provided, it automatically sets other fields + compile_mode: str = "simd" + mix_mode: str = "" + simt_stack_limit: int = None + # use_bytecode: + # If True, the compilation flow is: + # Linalg IR → MLIR Bytecode (via triton-mlir-opt) + # → LLIR (via bishengir-opt) + # → Binary (via bishengir-compile) + # + # If False, the compilation flow is: + # Linalg IR → LLIR → Binary (via bishengir-compile directly) + use_bytecode: bool = True + # take effect on the reorder instruction pattern for SIMT. The pattern is disabled by default. + enable_simt_reorder_instruction: bool = False + # disable simt fma optimization to get high precision + disable_fma: bool = False + + def __post_init__(self): + # Parse compile_mode and set related fields + if self.compile_mode == "simd": + object.__setattr__(self, "parallel_mode", "simd") + elif self.compile_mode == "unstructured_in_simt": + # For historical compatibility reasons, force_simt_template will still be used. + object.__setattr__(self, "force_simt_template", True) + elif self.compile_mode == "simt_only": + object.__setattr__(self, "force_simt_only", True) + object.__setattr__(self, "parallel_mode", "simt") + + if self.force_simt_only: + if self.shared_mem_dynamic_size is None: + object.__setattr__(self, "shared_mem_dynamic_size", 122880) + else: + object.__setattr__(self, "shared_mem_dynamic_size", 221184) + + def hash(self): + key = "_".join([f"{name}-{val}" for name, val in self.__dict__.items()]) + key = "_".join([key, get_cann_version_file_hash()]) + return hashlib.sha256(key.encode("utf-8")).hexdigest() + + +def ttir_to_npubin(mod, metadata, opt): + # Get Triton-MLIR as string + ttir_code = str(mod) + metadata = _parse_ttir_metadata(ttir_code, metadata) + with tempfile.TemporaryDirectory() as tmpdir: + # prepare input + src_path = os.path.join(tmpdir, "kernel.ttir.mlir") + Path(src_path).write_text(ttir_code) + # prepare output + bin_file = os.path.join(tmpdir, "kernel") + bin_path = os.path.join(tmpdir, "kernel.o") + # build compile options + _compile_option_list = get_common_bishengir_compile_options(metadata) + if opt.force_simt_only: + _compile_option_list += ["--enable-hivm-compile=false"] + _compile_option_list += ["--enable-triton-ir-compile"] + _compile_option_list += ["--pure-simt"] + _compile_option_list += [f"--num-warps={opt.num_warps}"] + _compile_option_list += [f"--threads-per-warp={opt.warp_size}"] + if opt.enable_bishengir_simt_optimization != 000: + _compile_option_list += [ + f"--enable-bishengir-simt-optimization={opt.enable_bishengir_simt_optimization}" + ] + if opt.simt_stack_limit: + _compile_option_list += [f"--simt-stack-limit={opt.simt_stack_limit}"] + if opt.shared_mem_dynamic_size is not None: + _compile_option_list += [f"--shared-mem-dynamic-size={opt.shared_mem_dynamic_size}"] + if opt.enable_simt_reorder_instruction: + _compile_option_list += ["--enable-simt-reorder-instruction=true"] + if opt.disable_fma: + _compile_option_list += [f"--disable-fma"] + enable_libdevice_simt = triton_enable_libdevice_simt() + if (enable_libdevice_simt): + bisheng_options = metadata["bisheng_options"] + if bisheng_options is not None: + _compile_option_list += [f"--append-bisheng-options={bisheng_options}"] + + # Enable SIMT auto-blockify if user opted in, or if the env var is + # set and the user didn't explicitly opt out (matches the SIMD path + # at line ~541). + enable_auto_blockify = opt.enable_auto_blockify + if _is_auto_map_parallel_blocks_enabled(): + if enable_auto_blockify is None or enable_auto_blockify: + _compile_option_list += ["--enable-auto-blockify-loop"] + else: + if enable_auto_blockify: + _compile_option_list += ["--enable-auto-blockify-loop"] + + npu_compiler_path, env = _get_npucompiler_path() + cmd_list = ([npu_compiler_path, src_path] + _compile_option_list + ["-o", bin_file]) + ret = subprocess.run(cmd_list, env=env, capture_output=True, check=True) + if not Path(bin_path).exists(): + error_msg = ret.stderr.decode('utf-8') + print(f"[DEBUG] {bin_path} is not found") + print(f"[DEBUG] Stderr:\n{error_msg}") + raise subprocess.CalledProcessError(ret.returncode, cmd_list, ret.stdout, ret.stderr) + return Path(bin_path).read_bytes() + + +class AscendBackend(BaseBackend): + + @staticmethod + def supports_target(target: GPUTarget): + return target.backend == "npu" + + def __init__(self, target: GPUTarget) -> None: + super().__init__(target) + if target.backend == "npu": + self.binary_ext = "npubin" + # Include all binary file extensions (mlirbc is used in bytecode mode) + self.binary_extensions = {"npubin", "mlirbc"} + + def parse_options(self, opts) -> Any: + # TODO: get available targets when building options? + if self.target.backend == "npu": + args = {k: opts[k] for k in NPUOptions.__dataclass_fields__.keys() if k in opts} + args.setdefault("arch", self.target.arch) + options = NPUOptions(**args) + else: + raise NotImplementedError(f"Backend '{self.target.backend}' is not supported. " + "Please ensure the target backend is set to 'npu'.") + return options + + def pack_metadata(self, metadata): + # collect necessary metadata to launch kernels + # TORCHINDUCTOR_UNIQUE_KERNEL_NAMES=1 could set unique name. + # Get this name as the kernel_name to CANN runtime. + # kernel_name is unique to Ascend backend and should not be public. + # CANN runtime limits the length of kernel name <= 50. + # Considering '\n' is appended, thus the real kernel name <= 49. + KERNEL_NAME_MAX_LEN = 49 + kernel_name_orig = metadata.kernel_name + if len(kernel_name_orig) > KERNEL_NAME_MAX_LEN: + kernel_name = kernel_name_orig[-KERNEL_NAME_MAX_LEN:] + else: + kernel_name = kernel_name_orig + return { + "kernel_name": kernel_name, + "hash": metadata.hash, + "debug": metadata.debug, + "tensor_kinds": metadata.tensor_kinds, + } + + def get_codegen_implementation(self, options): + # Note: a dict of functions is required to generate vendor-specific code piecies + # e.g. convert custom types like fp8e4b15 + from triton.backends.ascend import _apply_ascend_patch + _apply_ascend_patch() + codegen_fns = {"min_dot_size": min_dot_size(self.target)} + return codegen_fns + + def load_dialects(self, ctx): + ascend.load_dialects(ctx) + + def add_stages(self, stages, options, language): + if self.target.backend == "npu": + stages["ttir"] = lambda src, metadata: make_ttir(src, metadata, options) + if options.force_simt_only: + stages["npubin"] = (lambda src, metadata: ttir_to_npubin(src, metadata, options)) + return + stages["ttadapter"] = lambda src, metadata: ttir_to_linalg(src, metadata, options, named_ops=True) + # Support BC mode: convert Linalg IR to Bytecode format, then back to MLIR + if options.use_bytecode: + # Step 1: Convert Linalg IR to Bytecode using triton-mlir-opt + stages["mlirbc"] = lambda src, metadata: linalg_to_bc_by_triton_mlir_opt(src, metadata, options) + # Step 2: Convert Bytecode back to MLIR text using bishengir-opt + stages["bcmlir"] = lambda src, metadata: bc_to_linalg_by_bishengir_opt(src, metadata, options) + stages["npubin"] = (lambda src, metadata: _compile_linalg_to_npu_bin(src, metadata, options)) + else: + raise NotImplementedError(f"Backend '{self.target.backend}' is not supported. " + "Please ensure the target backend is set to 'npu'.") + + @functools.lru_cache() + def hash(self): + # TODO fetch compiler version + version_key = self.target + return str(version_key) + + def get_module_map(self) -> Dict[str, ModuleType]: + return {} diff --git a/third_party/ascend/backend/cpu_driver.py b/third_party/ascend/backend/cpu_driver.py new file mode 100644 index 0000000000..d429ffb6a2 --- /dev/null +++ b/third_party/ascend/backend/cpu_driver.py @@ -0,0 +1,203 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +from triton.runtime.cache import get_cache_manager, get_dump_manager +from pathlib import Path +import tempfile +import os +import sysconfig +import subprocess +import importlib +from triton.backends.ascend.utils import _get_llvm_path + + +# TODO: temporarily fake CPUUtils class +class CPUUtils(object): + + def __new__(cls): + if not hasattr(cls, 'instance'): + cls.instance = super(CPUUtils, cls).__new__(cls) + return cls.instance + + def __init__(self): + pass + + def get_device_properties(self, device): + # temperoarily added properties to avoid triton-compiler complain + # fetch available memory at runtime + return {"max_shared_mem": 1} + + def load_binary(self, name, kernel, shared, device, mix_mode=None): + # TODO (temperoarily fake function) load a binary from binary object to device + # return value are: (mod, funcptr/handle, n_regs, n_spills) + return None, kernel, 0, 0 + + +class CPULauncher(object): + + def __init__(self, src, metadata): + kernel_name = metadata.name.split()[0] + signature = src.signature + constants = src.constants + launcher_src = generate_cpu_wrapper_src(constants, signature, kernel_name) + self.launch = compile_module(launcher_src) + + def __call__(self, *args, **kwargs): + self.launch(*args, **kwargs) + + +class CPUDriver: + + def __init__(self): + self.utils = CPUUtils() + self.launcher_cls = CPULauncher + super().__init__() + + def get_current_target(self): + # TODO: do we rely on CPU arch? + return ("cpu", "arm-64") + + def get_current_device(self): + """ + Get current device + """ + # TODO: dummy device-getter for cpu backend + return 0 + + def set_current_device(self, device): + """ + Set current device as the given device + """ + # TODO: dummy device-setter for cpu backend + return + + def get_current_stream(self, device): + """ + Get stream for current device + """ + # TODO: dummy stream api for cpu backend. + return 0 + + +# the template is from triton-adapter HEAD. Wrapping the generated kernel assembly into a python module +def generate_cpu_wrapper_src(constants, signature, kernel_name): + + def _ty_to_cpp(ty): + if ty[0] == '*': + return "void*" + return { + "i1": "int32_t", + "i8": "int8_t", + "i16": "int16_t", + "i32": "int32_t", + "i64": "int64_t", + "u32": "uint32_t", + "u64": "uint64_t", + "fp16": "float", + "bf16": "float", + "fp32": "float", + "f32": "float", + "fp64": "double", + }[ty] + + def _extracted_ty(ty): + if ty[0] == '*': + return "PyObject*" + return { + 'i1': 'int32_t', + 'i32': 'int32_t', + 'i64': 'int64_t', + 'u32': 'uint32_t', + 'u64': 'uint64_t', + 'fp16': 'float', + 'bf16': 'float', + 'fp32': 'float', + 'f32': 'float', + 'fp64': 'double', + }[ty] + + def _format_of(ty): + return { + "PyObject*": "O", + "float": "f", + "double": "d", + "long": "l", + "uint32_t": "I", + "int32_t": "i", + "uint64_t": "K", + "int64_t": "L", + }[ty] + + def _generate_launcher(constants, signature, kernel_name): + arg_decls = ', '.join(f"{_ty_to_cpp(ty)} arg{i}" for i, ty in signature.items()) + format = "iiiOOO" + ''.join([_format_of(_extracted_ty(ty)) for ty in signature.values()]) + # to be filled + return f""" + """ + + launcher_src = _generate_launcher(constants, signature, kernel_name) + return launcher_src + + +def compile_module(launcher_src): + # This function was renamed and made public in Python 3.10 + if hasattr(sysconfig, 'get_default_scheme'): + scheme = sysconfig.get_default_scheme() + else: + scheme = sysconfig._get_default_scheme() + # 'posix_local' is a custom scheme on Debian. However, starting Python 3.10, the default install + # path changes to include 'local'. This change is required to use triton with system-wide python. + if scheme == 'posix_local': + scheme = 'posix_prefix' + py_include_dir = sysconfig.get_paths(scheme=scheme)["include"] + + def launch(gridX, gridY, gridZ, stream, cu_function, packed_metadata, launch_metadata, launch_enter_hook, + launch_exit_hook, *args): + kernel_name = packed_metadata["kernel_name"] + cache = get_cache_manager(packed_metadata["hash"]) + filename = f"{kernel_name}_cpu_launcher.so" + cache_path = cache.get_file(filename) + if cache_path is None: + asm_src = cu_function + with tempfile.TemporaryDirectory() as tmpdir: + asm_src_path = os.path.join(tmpdir, "kernel.s") + launcher_src_path = os.path.join(tmpdir, "main.cxx") + if packed_metadata["debug"]: + dump_manager = get_dump_manager(packed_metadata["hash"]) + dump_manager.put(launcher_src, "kernel_cpu_launcher.cxx", binary=False) + so_path = os.path.join(tmpdir, "kernel.so") + Path(asm_src_path).write_bytes(asm_src) + Path(launcher_src_path).write_text(launcher_src) + # Compile it together. + subprocess.check_call([ + _get_llvm_path("bin", "clang++"), launcher_src_path, asm_src_path, f"-I{py_include_dir}", + f"-I{Path(__file__).resolve().parent}", "-shared", "-fPIC", "-o", so_path + ]) + + with open(so_path, "rb") as f: + cache_path = cache.put(f.read(), filename, binary=True) + + # Load and launch the compiled kernel. + spec = importlib.util.spec_from_file_location("__triton_adapter_ref_cpu_kernel_launcher", cache_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.launch(gridX, gridY, gridZ, launch_enter_hook, launch_exit_hook, packed_metadata, *args) + + return launch diff --git a/third_party/ascend/backend/driver.py b/third_party/ascend/backend/driver.py new file mode 100644 index 0000000000..50b757b5e1 --- /dev/null +++ b/third_party/ascend/backend/driver.py @@ -0,0 +1,1058 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +from pathlib import Path +import tempfile +from dataclasses import dataclass +import os +import os.path +import re +import subprocess +import sysconfig +from typing import Optional +import functools +import hashlib + +from flagtree import _flagprism # FlagPrism +from triton.runtime.cache import get_cache_manager, get_dump_manager +from triton.backends.driver import DriverBase +from triton.backends.compiler import GPUTarget +from triton.backends.ascend.utils import (_precompile_npu_hash, _precompile_npu_ext, _build_npu_ext, _check_cxx11_abi, + convert_sigtype_to_int, _is_auto_map_parallel_blocks_enabled, + get_ascend_arch_from_env, is_ffts_supported, force_disable_ffts, + get_backend_func) + + +class NPUUtils(object): + + def __new__(cls): + if not hasattr(cls, 'instance'): + cls.instance = super(NPUUtils, cls).__new__(cls) + return cls.instance + + def __init__(self): + dirname = os.path.dirname(os.path.realpath(__file__)) + src_path = os.path.join(dirname, "npu_utils.cpp") + src = Path(src_path).read_text() + key = hashlib.md5(src.encode("utf-8")).hexdigest() + cache = get_cache_manager(key) + fname = "npu_utils.so" + cache_path = cache.get_file(fname) + if cache_path is None: + with tempfile.TemporaryDirectory() as tmpdir: + tmp_src_path = os.path.join(tmpdir, "npu_utils.cpp") + with open(tmp_src_path, "w") as f: + f.write(src) + so = _build_npu_ext("npu_utils", None, tmp_src_path) + with open(so, "rb") as f: + cache_path = cache.put(f.read(), fname, binary=True) + import importlib.util + spec = importlib.util.spec_from_file_location("npu_utils", cache_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + self.npu_utils_mod = mod + # setup for remote run + env_arch = get_ascend_arch_from_env() + + def load_binary(self, name, kernel, shared, device, mix_mode): + return self.npu_utils_mod.load_kernel_binary(name, kernel, shared, device, mix_mode) + + @functools.lru_cache() + def get_device_properties(self, device): + # temperoarily added "max_shared_mem" properties to avoid triton-compiler complain + # fetch available memory at runtime + num_aic = self.get_aicore_num() + num_aiv = num_aic * 2 + return {"max_shared_mem": 1, "num_aicore": num_aic, "num_vectorcore": num_aiv} + + @functools.lru_cache() + def get_arch(self): + # temporarily return empty arch descriptor + return self.npu_utils_mod.get_arch() + + @functools.lru_cache() + def get_aicore_num(self): + # temporarily return empty arch descriptor + return self.npu_utils_mod.get_aicore_num() + + @functools.lru_cache() + def get_aivector_core_num(self): + return self.get_device_properties("npu")["num_vectorcore"] + + +@dataclass(frozen=True) +class _DebuggerHiddenArgABI: + """FlagPrism: generated-launcher fragments for one hidden pointer.""" + + parse_format: str = "" + extracted_declaration: str = "" + parse_argument: str = "" + launch_declaration: str = "" + struct_field: str = "" + struct_value: str = "" + call_argument: str = "" + + @classmethod + def from_metadata(cls, metadata): + if not bool(getattr(metadata, "debug_launch_hidden_arg", False)): + return cls() + return cls( + parse_format="K", + extracted_declaration="uint64_t _debugHiddenArg = 0;", + parse_argument=", &_debugHiddenArg", + launch_declaration=", uint64_t debugHiddenArg", + struct_field="void* debug_hidden_arg __attribute__((aligned(8)));", + struct_value="reinterpret_cast(debugHiddenArg),", + call_argument=", _debugHiddenArg", + ) + + +class NPULauncher(object): + + def __init__(self, src, metadata): + self.compile_only = os.getenv("TRITON_COMPILE_ONLY", 'false').lower() in ('true', '1') + self.enable_msprof_register_tensor = os.getenv("TRITON_REGISTER_TENSOR_MSPROF", + 'false').lower() in ('true', '1') + debug_mode = metadata.debug + header_src = generate_npu_header_src() + constants = src.constants if hasattr(src, "constants") else dict() + cst_key = lambda i: src.fn.arg_names.index(i) if isinstance(i, str) else i + constants = {cst_key(key): value for key, value in constants.items()} + signature = {cst_key(key): value for key, value in src.signature.items()} + wrapper_src = make_launcher(constants, signature, metadata) + so_launcher_path = make_npu_launcher_stub(header_src, wrapper_src, metadata.debug) + self.metadata = metadata # FlagPrism + # setup for remote run + # TODO: use a var to pack all vars required to run on a remote machine + self.mix_mode = metadata.mix_mode + self.shared = metadata.shared + # initialize launcher + import importlib.util + spec = importlib.util.spec_from_file_location("__triton_launcher", so_launcher_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + self.launch = getattr(mod, "launch") + + def __call__(self, *args, **kwargs): + if self.compile_only: + cache_manager = get_cache_manager(args[5]['hash']) + print("[INFO]: skip running kernel") + print(f"[INFO]: The compiled kernel cache is in {cache_manager.cache_dir}") + if self.enable_msprof_register_tensor: + tensor_params_shape = get_backend_func("get_tensor_params_shape", *args) + # args[5] must be the packed metadata. + # Check the launch wrapper in which PyArg_ParseTuple specifies the ordering of args + args[5]['tensor_params_shape'] = tensor_params_shape + else: + if self.compile_only: + return + debug_enabled = bool(getattr(self.metadata, "debug_enabled", False)) + if debug_enabled: + # FlagPrism: use the backend-neutral launch gateway; this + # backend only supplies its name and launch state. + with _flagprism.debugger_launch_context( + "ascend", + self.metadata, + args[:3], + args[3], + args[6], + args[9:], + ) as hidden_args: + profiler_registered = self.launch(*args, *hidden_args, **kwargs) + else: + profiler_registered = self.launch(*args, **kwargs) + import triton + triton.backends.ascend.utils.TRITON_PROFILER_REGISTERED = True if profiler_registered == 1 else False + + +class NPUDriver(DriverBase): + + def __init__(self): + self.utils = NPUUtils() + self.launcher_cls = NPULauncher + super().__init__() + + @classmethod + def is_active(cls): + + def test_npucompiler(): + from triton.backends.ascend.utils import _get_bisheng_path + npucompiler = _get_bisheng_path() + targets = subprocess.check_output([npucompiler, "-print-targets"]).decode().strip().split() + return "hiipu64" in targets + + try: + return test_npucompiler() + except Exception as e_npucompiler: + import warnings + red = "\x1b[31;20m" + reset = "\x1b[0m" + warnings.warn(red + str(e_npucompiler) + reset) + return False + + def map_python_to_cpp_type(self, ty: str) -> str: + return ty_to_cpp(ty) + + def get_current_target(self): + backend = "npu" + env_target = get_ascend_arch_from_env() + if env_target: + arch = env_target + else: + arch = self.utils.get_arch() + warp_size = 0 + return GPUTarget(backend, arch, warp_size) + + def get_current_device(self): + """ + Get current device + """ + return get_backend_func("get_current_device") + + def get_active_torch_device(self): + import torch + import torch_npu + return torch.device("npu", self.get_current_device()) + + def set_current_device(self, device): + """ + Set current device as the given device + """ + return get_backend_func("set_current_device", device) + + def get_current_stream(self, device: Optional[int] = None) -> int: + """ + Get stream for current device + """ + # According to torch_npu, the content of a torch.npu.Stream is essentilly an rtStream_t + # TODO: use CANN API instead of torchnpu + return get_backend_func("get_current_stream", device) + + def get_benchmarker(self): + from triton.testing import do_bench + return do_bench + + def get_device_interface(self): + return get_backend_func("get_device_interface") + + def get_empty_cache_for_benchmark(self): + cache_size = 192 * 1024 * 1024 + return get_backend_func("get_empty_tensor", cache_size // 4) + + def clear_cache(self, cache): + cache.zero_() + + +def _precompile_npu_ext_with_lock(header_src, enable_precompile): + import fcntl + precompile_hash = _precompile_npu_hash(header_src) + cache = get_cache_manager(precompile_hash) + gch_path = cache.get_file("precompiled.h.gch") + header_path = cache.get_file("precompiled.h") + if enable_precompile: + if header_path is not None and gch_path is not None: + return header_path + else: + if header_path is not None: + return header_path + cache_dir = os.getenv("TRITON_CACHE_DIR", "").strip() + lock_path = os.path.join(cache_dir, f"{precompile_hash}.lock") + with open(lock_path, "a+") as f: + try: + fcntl.flock(f, fcntl.LOCK_EX) + header_path = cache.get_file("precompiled.h") + if enable_precompile: + gch_path = cache.get_file("precompiled.h.gch") + if header_path is not None and gch_path is not None: + return header_path + else: + if header_path is not None: + return header_path + header_path = cache.put(header_src, "precompiled.h", binary=False) + if not enable_precompile: + return header_path + src_dir = os.path.dirname(header_path) + gch_path = os.path.join(src_dir, "precompiled.h.gch") + _precompile_npu_ext(header_path, gch_path) + return header_path + finally: + fcntl.flock(f, fcntl.LOCK_UN) + + +def make_npu_launcher_stub(header_src, wrapper_src, debug=False): + """ + Generate the launcher stub to launch the kernel + """ + enable_precompile = not os.getenv("TRITON_DISABLE_PRECOMPILE", 'false').lower() in ('true', '1') + # if precompile header file and its gch file not exist, do precompile + header_path = _precompile_npu_ext_with_lock(header_src, enable_precompile) + assert header_path is not None, "the precompiled.h path is empty." + + # try to get cached file + so_cache_key = hashlib.sha256(wrapper_src.encode("utf-8")).hexdigest() + so_cache_manager = get_cache_manager(so_cache_key) + # append the cxx11_abi value to the launcher name to avoid + # linking to a launcher with wrong cxx11_abi. + use_cxx11_abi = _check_cxx11_abi() + name = f"launcher_cxx11abi{use_cxx11_abi}" + suffix = sysconfig.get_config_var('EXT_SUFFIX') + so_name = f"{name}{suffix}" + + if debug: + dump_manager = get_dump_manager(so_cache_key) + if header_path is not None: + print(f"Dumping precompiled.h to {dump_manager.cache_dir}") + dump_manager.put(header_src, "precompiled.h", binary=False) + print(f"Dumping {name}.cxx to {dump_manager.cache_dir}") + dump_manager.put(wrapper_src, f"{name}.cxx", binary=False) + + cache_path = so_cache_manager.get_file(so_name) + if cache_path is not None: + return cache_path + + kernel_launcher_type = "torch" + + with tempfile.TemporaryDirectory() as tmpdir: + src_path = os.path.join(tmpdir, f"{name}.cxx") + with open(src_path, "w") as f: + f.write(wrapper_src) + so_path = _build_npu_ext(name, header_path, src_path, kernel_launcher=kernel_launcher_type, + precompile=enable_precompile) + if debug: + with open(so_path, "rb") as f: + dump_manager.put(f.read(), so_name, binary=True) + with open(so_path, "rb") as f: + so_cache_path = so_cache_manager.put(f.read(), so_name, binary=True) + return so_cache_path + + +def extract_device_print_code_from_cann(): + from triton.backends.ascend.utils import _get_bisheng_path + ccec_compiler_bin_folder, _ = os.path.split(os.path.realpath(_get_bisheng_path())) + ccec_compiler_folder, _ = os.path.split(ccec_compiler_bin_folder) + clang_version = os.listdir(os.path.join(ccec_compiler_folder, "lib/clang/"))[0] + ccelib_path = os.path.join(ccec_compiler_folder, f"lib/clang/{clang_version}/include/ccelib") + + def read_header(header_path): + with open(os.path.join(ccelib_path, header_path), 'r') as f: + code = f.read() + + # remove all #include "..." + lines = code.splitlines() + purged_lines = [] + for line in lines: + normalized_line = ' '.join(line.split()) + if not normalized_line.startswith('#include "'): + purged_lines.append(line) + code = '\n'.join(purged_lines) + + # remove [aicore] functions + aicore_positions = [] + for m in re.finditer('\[aicore\]', code): + aicore_positions.append(m.start()) + + def find_aicore_function_span(src, pos): + for i in range(pos - 1, -1, -1): + if src[i] == '}': # this relies on that all [aicore] functions come after normal functions + left = i + 1 + break + n = len(src) + brace_nest = 0 + for j in range(pos, n, 1): + if src[j] == '{': + brace_nest += 1 + elif src[j] == '}': + brace_nest -= 1 + if brace_nest == 0: + right = j + break + return left, right + + new_code = '' + segment_start = 0 + for pos in aicore_positions: + left, right = find_aicore_function_span(code, pos) + new_code += code[segment_start:left] + segment_start = right + 1 + new_code += code[segment_start:] + + # remove __gm__ and rename macros + new_code = new_code.replace('__gm__', ' ') + new_code = new_code.replace('__CCELIB_RT_ERROR_NONE', 'RT_ERROR_NONE') + new_code = new_code.replace('__CCELIB_RT_MEMORY_HBM', 'RT_MEMORY_HBM') + new_code = new_code.replace('__CCELIB_RT_MEMCPY_HOST_TO_DEVICE', 'RT_MEMCPY_HOST_TO_DEVICE') + new_code = new_code.replace('__CCELIB_RT_MEMCPY_DEVICE_TO_HOST', 'RT_MEMCPY_DEVICE_TO_HOST') + return new_code + + # the following headers should be included in this order + return '\n'.join([ + read_header('common/common_impl.h'), + read_header('internal/debug_tunnel/payload.h'), + read_header('internal/debug_tunnel/payload_impl.h'), + read_header('internal/debug_tunnel/tunnel.h'), + read_header('internal/debug_tunnel/tunnel_impl.h') + ]) + + +def generate_npu_header_src(): + enable_taskqueue = os.getenv("TRITON_ENABLE_TASKQUEUE", 'true').lower() in ('true', '1') + return f""" +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * Copyright 2018-2020 Philippe Tillet + * Copyright 2020-2022 OpenAI + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ +#ifndef TRITON_NPU_HEADERS +#define TRITON_NPU_HEADERS + +#include +#include +#include +#include +#include +#include +#include "runtime/runtime/rt.h" +#include +{get_backend_func("header_file", enable_taskqueue)} + +#endif +""" + + +# ------------------------ +# Launcher +# ------------------------ + + +def ty_to_cpp(ty): + if ty[0] == '*': + return "void*" + return { + "i1": "int32_t", + "i8": "int8_t", + "i16": "int16_t", + "i32": "int32_t", + "i64": "int64_t", + "u1": "uint32_t", + "u8": "uint8_t", + "u16": "uint16_t", + "u32": "uint32_t", + "u64": "uint64_t", + "fp16": "float", + "bf16": "float", + "fp32": "float", + "f32": "float", + "fp64": "double", + }[ty] + + +# the template is from triton-adapter HEAD. Wrapping the generated kernel binary into a python module +def make_launcher(constants, signature, metadata): + import os + workspace_size = int(metadata.workspace_size) \ + if hasattr(metadata, 'workspace_size') else -1 + lock_init_value = int(metadata.lock_init_value) \ + if hasattr(metadata, 'lock_init_value') else 0 + lock_num = int(metadata.lock_num) \ + if hasattr(metadata, 'lock_num') else -1 + bs_task_type = metadata.bs_task_type if hasattr(metadata, 'bs_task_type') else 0 + mix_mode = metadata.mix_mode + compile_on_910_95 = metadata.compile_on_910_95 + parallel_mode = metadata.parallel_mode + enable_simt = ("simt" in parallel_mode) or metadata.force_simt_only + + def _serialize_signature(sig): + if isinstance(sig, tuple): + return ','.join(map(_serialize_signature, sig)) + return sig + + def _extracted_type(ty): + if isinstance(ty, tuple): + val = ','.join(map(_extracted_type, ty)) + return f"[{val}]" + if ty[0] == '*': + return "PyObject*" + if ty in ("constexpr"): + return "PyObject*" + return ty_to_cpp(ty) + + def format_of(ty): + if isinstance(ty, tuple): + val = ''.join(map(format_of, ty)) + return f"({val})" + if ty[0] == '*': + return "O" + if ty in ("constexpr"): + return "O" + if ty == "void*": + return "O" + return { + "float": "f", + "double": "d", + "long": "l", + "int8_t": "b", + "int16_t": "h", + "int32_t": "i", + "int64_t": "L", + "uint8_t": "B", + "uint16_t": "H", + "uint32_t": "I", + "uint64_t": "K", + }[ty_to_cpp(ty)] + + def _format_of_msprof_task_type_ratio(bs_task_type, mix_mode): + # Default fallback based on mix_mode + default_task_type = "MSPROF_GE_TASK_TYPE_AIV" if mix_mode == "aiv" else "MSPROF_GE_TASK_TYPE_AI_CORE" + + if not bs_task_type: + return default_task_type, 0 + + task_type_num, mix_block_dim_ratio = divmod(int(bs_task_type), 10) + task_type_map = { + 1: "MSPROF_GE_TASK_TYPE_AIV", + 2: "MSPROF_GE_TASK_TYPE_AI_CORE", + 3: "MSPROF_GE_TASK_TYPE_MIX_AIC", + 4: "MSPROF_GE_TASK_TYPE_MIX_AIV", + } + + task_type = task_type_map.get(task_type_num, default_task_type) + return task_type, mix_block_dim_ratio + + """ + args: + int gridX, gridY, gridZ; + rtStream_t stream; + const void *functon; + PyObject* packed_metadata, *launch_metadata; + PyObject* launch_enter_hook, *launch_exit_hook; + *args_expand + """ + debug_abi = _DebuggerHiddenArgABI.from_metadata(metadata) # FlagPrism + + args_format = ''.join([format_of(ty) for ty in signature.values()]) + format = "iiiKKOOOO" + args_format + debug_abi.parse_format # FlagPrism + signature = ','.join(map(_serialize_signature, signature.values())) + signature = list(filter(bool, signature.split(','))) + signature = {i: s for i, s in enumerate(signature)} + args_list = ', ' + ', '.join(f"&_arg{i}" for i, ty in signature.items()) if len(signature) > 0 else '' + args_list += debug_abi.parse_argument # FlagPrism + # Record the end of regular arguments; + # subsequent arguments are architecture-specific descriptors. + arg_decls = ', '.join(f"{ty_to_cpp(ty)} arg{i}" for i, ty in signature.items() if ty != "constexpr") + internal_args_list = [] + for i, ty in signature.items(): + if ty[0] == "*": + internal_args_list.append(f"ptr_info{i}.dev_ptr") + elif ty != "constexpr": + internal_args_list.append(f"_arg{i}") + + # generate glue code + newline = '\n ' + ptr_decls = [ + f"DevicePtrInfo ptr_info{i} = getPointer(_arg{i}, {i}); if (!ptr_info{i}.valid) return NULL;" + for i, ty in signature.items() + if ty[0] == "*" + ] + + grid_info = {'X': 'i32', 'Y': 'i32', 'Z': 'i32'} + # TODO: automatically check if gather load ops are used. + + arch = get_ascend_arch_from_env() + target_support_ffts = is_ffts_supported(arch) and (not force_disable_ffts()) + enable_device_print = os.getenv("TRITON_DEVICE_PRINT", 'false').lower() in ('true', '1') + enable_taskqueue = os.getenv("TRITON_ENABLE_TASKQUEUE", 'true').lower() in ('true', '1') + enable_grid_warn_print = os.getenv("TRITON_GRID_WARN_PRINT", 'false').lower() in ('true', '1') + # Per-kernel metadata wins; fall back to the env var when unset. + enable_auto_map_parallel_blocks = getattr(metadata, "enable_auto_blockify", None) + if enable_auto_map_parallel_blocks is None: + enable_auto_map_parallel_blocks = _is_auto_map_parallel_blocks_enabled() + npu_utils = NPUUtils() + num_physical_blocks = npu_utils.get_aivector_core_num() if mix_mode == "aiv" else npu_utils.get_aicore_num() + task_type, mix_block_dim_ratio = _format_of_msprof_task_type_ratio(bs_task_type, mix_mode) + is_mix_task_type = "true" if ("MIX" in task_type) else "false" + LINE_CHANGE_CHAR = chr(10) # it is \n + alloc_success_code = 'return 1;' + sync_lock_fail_code = 'fprintf(stderr, "Error: syncBlockLock allocation failed\\n"); return;' + workspace_fail_code = 'fprintf(stderr, "Error: workspace allocation failed\\n"); return;' + + cpp_device_pointer = """ +typedef struct _DevicePtrInfo { + void *dev_ptr; + bool valid; +} DevicePtrInfo; + +static inline DevicePtrInfo getPointer(PyObject *obj, int idx) { + DevicePtrInfo ptr_info; + ptr_info.dev_ptr = 0; + ptr_info.valid = true; + if (PyLong_Check(obj)) { + ptr_info.dev_ptr = reinterpret_cast(PyLong_AsUnsignedLongLong(obj)); + return ptr_info; + } + if (obj == Py_None) { + // valid nullptr + return ptr_info; + } + PyObject *ptr = PyObject_GetAttrString(obj, "data_ptr"); + if(ptr){ + PyObject *empty_tuple = PyTuple_New(0); + PyObject *ret = PyObject_Call(ptr, empty_tuple, NULL); + Py_DECREF(empty_tuple); + Py_DECREF(ptr); + if (!PyLong_Check(ret)) { + PyErr_SetString(PyExc_TypeError, "data_ptr method of Pointer object must return 64-bit int"); + ptr_info.valid = false; + return ptr_info; + } + ptr_info.dev_ptr = reinterpret_cast(PyLong_AsUnsignedLongLong(ret)); + if(!ptr_info.dev_ptr) + return ptr_info; + aclrtPtrAttributes attributes; + aclError status = aclrtPointerGetAttributes(ptr_info.dev_ptr, &attributes); + + if (status == ACL_SUCCESS) { + if (attributes.location.type != ACL_MEM_LOCATION_TYPE_DEVICE && attributes.location.type != 4) { + Py_DECREF(ret); + PyErr_Format(PyExc_ValueError, + "Pointer argument (at %d) cannot be accessed from Triton (cpu tensor?)", idx); + ptr_info.valid = false; + return ptr_info; + } + } else { + Py_DECREF(ret); + PyErr_Format(PyExc_RuntimeError, + "Failed to query pointer attributes at argument %d. " + "Error code: %d. This may indicate invalid memory address " + "or NPU device error.", + idx, status); + ptr_info.valid = false; + return ptr_info; + } + Py_DECREF(ret); + return ptr_info; + } + PyErr_SetString(PyExc_TypeError, "Pointer argument must be either uint64 or have data_ptr method"); + ptr_info.valid = false; + return ptr_info; +} +""" + + cpp_msprof_extern = """ +extern "C" { + typedef int (* callback)(unsigned int type, void* data, unsigned int len); + extern int MsprofReportApi(unsigned int agingFlag, const MsprofApi *api); + extern unsigned long int MsprofSysCycleTime(); + extern int MsprofRegisterCallback(unsigned int moduleId, callback handle); + static unsigned int __MsprofFlagL0 = 0; + static unsigned int __MsprofFlagL1 = 0; + + int ProfCtrlHandle(unsigned int CtrlType, void* CtrlData, unsigned int DataLen) { + if ((CtrlData == nullptr) || (DataLen == 0U)) { + return 1; + } + + if (CtrlType == 1) { + MsprofCommandHandle* handle = (MsprofCommandHandle *)(CtrlData); + if (handle->type >= 6) // 6 is not used here + return 1; + if (handle->type == 1) { // init - 0 , start - 1 + __MsprofFlagL0 = ((0x00000800ULL & handle->profSwitch) == 0x00000800ULL) ? 1 : 0; + __MsprofFlagL1 = ((0x00000002ULL & handle->profSwitch) == 0x00000002ULL) ? 1 : 0; + } + } + return 0; + } +} +""" + + cpp_msprof_callback = """ + MsprofRegisterCallback(8, ProfCtrlHandle); // 8 - CCE defined in msprof headerfile slog.h +""" + + cpp_msprof_call_before_launch = """ + unsigned long int beginTime = 0; + unsigned long int endTime = 0; + unsigned long int opNameHashID = 0; + unsigned int threadId = 0; + char* _kernelName = const_cast(name.c_str()); + size_t length = name.length(); + if (__MsprofFlagL0 || __MsprofFlagL1) + { + beginTime = MsprofSysCycleTime(); + } +""" + + cpp_msprof_call_after_launch = f""" + if (__MsprofFlagL0 || __MsprofFlagL1) + {{ + endTime = MsprofSysCycleTime(); + opNameHashID = MsprofGetHashId(_kernelName, length); + threadId = (unsigned int)(syscall(SYS_gettid)); + MsprofApi info; + info.level = MSPROF_REPORT_NODE_LEVEL; + info.magicNumber = 0x5a5a; //MSPROF_REPORT_DATA_MAGIC_NUM + info.type = MSPROF_REPORT_NODE_LAUNCH_TYPE; + info.threadId = threadId; + info.reserve = 0; + info.beginTime = beginTime; + info.endTime = endTime; + info.itemId = opNameHashID; + MsprofReportApi(false, &info); + }} + if (__MsprofFlagL1) + {{ + MsprofCompactInfo nodeBasicInfo; + nodeBasicInfo.level = MSPROF_REPORT_NODE_LEVEL; + nodeBasicInfo.magicNumber = 0x5a5a; //MSPROF_REPORT_DATA_MAGIC_NUM + nodeBasicInfo.type = MSPROF_REPORT_NODE_BASIC_INFO_TYPE; + nodeBasicInfo.threadId = threadId; + nodeBasicInfo.timeStamp = endTime; + nodeBasicInfo.data.nodeBasicInfo.opName = opNameHashID; + nodeBasicInfo.data.nodeBasicInfo.opType = opNameHashID; + nodeBasicInfo.data.nodeBasicInfo.taskType = {task_type}; + nodeBasicInfo.data.nodeBasicInfo.blockDim = nodeBasicBlockDim; + MsprofReportCompactInfo(0, static_cast(&nodeBasicInfo), sizeof(MsprofCompactInfo)); + + // 'mix' kernel need to report the ctxID + if ({is_mix_task_type} > 0) {{ + MsprofAdditionalInfo info; + info.level = MSPROF_REPORT_NODE_LEVEL; + info.type = MSPROF_REPORT_NODE_CONTEXT_ID_INFO_TYPE; + info.threadId = threadId; + info.timeStamp = endTime; + MsprofContextIdInfo ctxId; + ctxId.opName = opNameHashID; + ctxId.ctxIdNum = 1; + for (uint32_t i = 0; i < ctxId.ctxIdNum; i++) {{ + ctxId.ctxIds[i] = i; + }} + size_t copyLen = sizeof(MsprofContextIdInfo); + if (copyLen > MSPROF_ADDTIONAL_INFO_DATA_LENGTH) {{ + copyLen = MSPROF_ADDTIONAL_INFO_DATA_LENGTH; + }} + memcpy(info.data, &ctxId, copyLen); + MsprofReportAdditionalInfo(false, static_cast(&info), sizeof(MsprofAdditionalInfo)); + }} + + // Report tensor info + int max_tensors_num = tensorShapes.size() < MSPROF_GE_TENSOR_DATA_NUM ? tensorShapes.size() : MSPROF_GE_TENSOR_DATA_NUM; + MsprofAdditionalInfo tensorInfo; + tensorInfo.level = MSPROF_REPORT_NODE_LEVEL; + tensorInfo.type = MSPROF_REPORT_NODE_TENSOR_INFO_TYPE; + tensorInfo.threadId = threadId; + tensorInfo.timeStamp = endTime; + auto profTensorData = reinterpret_cast(tensorInfo.data); + profTensorData->opName = opNameHashID; + int tensorCount = 0; + int dataTypes[MSPROF_GE_TENSOR_DATA_NUM]; + if (tensorShapes.size() > 0) {{ + {LINE_CHANGE_CHAR.join( + f'dataTypes[{i}] = {convert_sigtype_to_int(ty[1:])};' + for i, ty in signature.items() + if ty.startswith("*") and i < 5 + )} + }} + for (int i = 0; i < tensorShapes.size() && tensorCount < MSPROF_GE_TENSOR_DATA_NUM; i++) {{ + auto fillTensorData = [&](int index, int tensorType) {{ + profTensorData->tensorData[index].tensorType = tensorType; + profTensorData->tensorData[index].format = 2; // GeDataFormat: ND = 2 + profTensorData->tensorData[index].dataType = dataTypes[i]; + int nDim = tensorShapes[i].size(); + nDim = nDim < MSPROF_GE_TENSOR_DATA_SHAPE_LEN ? nDim : MSPROF_GE_TENSOR_DATA_SHAPE_LEN; + for (int j = 0; j < nDim; j++) {{ + profTensorData->tensorData[index].shape[j] = tensorShapes[i][j]; + }} + for (int j = nDim; j < MSPROF_GE_TENSOR_DATA_SHAPE_LEN; j++) {{ + profTensorData->tensorData[index].shape[j] = 0; + }} + }}; + int tensorType = (i < tensorKinds.size()) ? tensorKinds[i] : 0; // DeFault tensor type is input + if (tensorType == TENSOR_KIND_INPUT || tensorType == TENSOR_KIND_INPUT_OUTPUT) {{ + fillTensorData(tensorCount, MSPROF_GE_TENSOR_TYPE_INPUT); + tensorCount++; + }} + if ((tensorType == TENSOR_KIND_OUTPUT || tensorType == TENSOR_KIND_INPUT_OUTPUT) && tensorCount < MSPROF_GE_TENSOR_DATA_NUM){{ + fillTensorData(tensorCount, MSPROF_GE_TENSOR_TYPE_OUTPUT); + tensorCount++; + }} + }} + profTensorData->tensorNum = tensorCount; + MsprofReportAdditionalInfo(false, static_cast(&tensorInfo), sizeof(MsprofAdditionalInfo)); + }} +""" + + cpp_kernel_launch = f""" + ret = rtKernelLaunch(func, blockNum, static_cast(&args), sizeof(args), NULL, stream); +""" + if compile_on_910_95 and enable_simt: + cpp_kernel_launch = f""" + rtArgsEx_t argsInfo = {{}}; + argsInfo.args = static_cast(&args); + argsInfo.argsSize = sizeof(args); + rtTaskCfgInfo_t cfgInfo = {{}}; + cfgInfo.localMemorySize = {metadata.shared_mem_dynamic_size}; + ret = rtKernelLaunchWithFlagV2(func, blockNum, &argsInfo, NULL, stream, 0, &cfgInfo); +""" + + precompile_headers = f""" +#include "precompiled.h" +""" + + return f""" +{precompile_headers} +{'#define __CCE_ENABLE_PRINT__' if enable_device_print else ''} +{extract_device_print_code_from_cann() if enable_device_print else ''} +#define PY_SSIZE_T_CLEAN +{'#define ENABLE_GRID_WARN_PRINT' if enable_grid_warn_print else ''} +#define TENSOR_KIND_INPUT 0 +#define TENSOR_KIND_OUTPUT 1 +#define TENSOR_KIND_INPUT_OUTPUT 2 + +{cpp_msprof_extern} + +{cpp_device_pointer} + +// FlagPrism: append the optional debugger pointer to the generated launch ABI. +static void _launch(const char* kernelName, const void* func, rtStream_t stream, int gridX, int gridY, int gridZ, std::vector> &tensorShapes, std::vector &tensorKinds{', ' + arg_decls if len(signature) > 0 else ''}{debug_abi.launch_declaration}) {{ + // only 1D parallelization is supported for NPU + // Pointer type becomes flattend 1-D Memref tuple: base_ptr, data_ptr, offset, shape, stride + // base_ptr offset shape and stride are not used, arbitrarily set for now + std::string name = ""; + name.append(kernelName); + void *workspace_addr_ptr = NULL; + uint32_t blockNum4Workspace = gridX * gridY * gridZ; + {get_backend_func("pre_launch", True)} + {f''' + // Widen before multiplication to avoid wrapping workspace sizes at 4 GiB. + uint64_t totalWorkSpaceSize = static_cast({workspace_size}) * blockNum4Workspace; + {get_backend_func("allocate_memory", "totalWorkSpaceSize", "stream")} + ''' if workspace_size > 0 else ''} + {'auto launch_call = [=]() -> rtError_t' if enable_taskqueue else ''} {{ + {get_backend_func("pre_launch", False)} + uint32_t blockNum = gridX * gridY * gridZ; + + #ifdef ENABLE_GRID_WARN_PRINT + static bool warned = false; + if (!warned && blockNum > (uint32_t){num_physical_blocks}) {{ + printf("WARNING: Grid %u > physical limit {num_physical_blocks}, performance maybe reduced.\\n",blockNum); + warned = true; + }} + #endif + {'blockNum = std::min(blockNum, (uint32_t)' + str(num_physical_blocks) + ');' if enable_auto_map_parallel_blocks else ''} + // set mixBlockNumRation for nodeBasicBlockDim for msprof report + uint32_t mixBlockNumRation = {mix_block_dim_ratio}; + uint32_t nodeBasicBlockDim = (mixBlockNumRation << 16) + blockNum; + + {'cce::internal::DebugTunnelData *DTData = cce::internal::DebugTunnel::Open(blockNum);' if enable_device_print else ''} + rtError_t ret = RT_ERROR_NONE; + {'void *ffts_addr = NULL; uint32_t ffts_len; ret = rtGetC2cCtrlAddr((uint64_t*)&ffts_addr, &ffts_len);' if target_support_ffts else ''} + {'if (ret != RT_ERROR_NONE) return ret;' if (target_support_ffts and enable_taskqueue) else 'if (ret != RT_ERROR_NONE) return;' if (target_support_ffts and (not enable_taskqueue)) else ''} + // stub argument for workspace + void *syncBlockLock_ptr = NULL; + uint16_t ModuleId = 0; + {f''' + uint64_t syncBlockLockSize = {lock_num} * sizeof(int64_t); + {get_backend_func("allocate_sync_block_lock", "syncBlockLockSize", "stream")} + if (!syncBlockLock_ptr) {{ + {alloc_success_code if enable_taskqueue else sync_lock_fail_code} + }} + std::vector lockInitData({lock_num}, {lock_init_value}); + ret = rtMemcpy( + syncBlockLock_ptr, syncBlockLockSize, + reinterpret_cast(lockInitData.data()), syncBlockLockSize, + RT_MEMCPY_HOST_TO_DEVICE + ); + if (ret != RT_ERROR_NONE) {{ + return {'ret' if enable_taskqueue else ''}; + }} + ''' if lock_num > 0 else ''} + {'if (ret != RT_ERROR_NONE) return ret;' if (workspace_size > 0 and enable_taskqueue) else 'if (ret != RT_ERROR_NONE) return;' if (workspace_size > 0 and not enable_taskqueue) else ''} + struct __attribute__((packed)) {{ + {'void* ffts_addr __attribute__((aligned(8)));' if target_support_ffts else ''} + {'void* syncBlockLock __attribute__((aligned(8)));' if not metadata.force_simt_only else ''} + {'void* workspace_addr __attribute__((aligned(8)));' if not metadata.force_simt_only else ''} + {' '.join(f'{ty_to_cpp(ty)} arg{i} __attribute__((aligned({4 if ty[0] != "*" and ty[-2:] != "64" else 8})));' for i, ty in signature.items() if i not in constants and ty != "constexpr")} + // FlagPrism: optional debugger control pointer; empty in normal launches. + {debug_abi.struct_field} + {' '.join(f'{ty_to_cpp(ty)} grid{mark} __attribute__((aligned(4)));' for mark, ty in grid_info.items())} + {'void* DTData __attribute__((aligned(8)));' if enable_device_print else ''} + }} args = {{ + {'static_cast(ffts_addr),' if target_support_ffts else ''} + {('static_cast(syncBlockLock_ptr),' if lock_num > 0 else 'nullptr,') if not metadata.force_simt_only else ''} + {('static_cast(workspace_addr_ptr),' if workspace_size > 0 else 'nullptr,') if not metadata.force_simt_only else ''} + {(lambda _rt: (', '.join(_rt) + ',') if _rt else '')( + [f'static_cast<{ty_to_cpp(ty)}>(arg{i})' for i, ty in signature.items() if i not in constants and ty != "constexpr"] + )} + {debug_abi.struct_value} // FlagPrism + {', '.join(f'static_cast<{ty_to_cpp(ty)}>(grid{mark})' for mark, ty in grid_info.items())} + {', static_cast(DTData)' if enable_device_print else ''} + }}; + {cpp_msprof_call_before_launch} + {cpp_kernel_launch} + {'void *&stream_ref = const_cast(stream);' if enable_device_print else ''} + {'cce::internal::DebugTunnel::Close(DTData, stream_ref);' if enable_device_print else ''} + {cpp_msprof_call_after_launch} + {'return ret;' if enable_taskqueue else 'ret = rtStreamSynchronize(stream);'} + }}; + {f'''{get_backend_func("async_launch", "launch_call") if enable_taskqueue else ''}'''} + return; +}} + +// Extract tensor shape from PyObject +static std::vector _get_tensor_shape(PyObject *tensor) {{ + std::vector shape; + + // Early return if tensor is None or null + if (!tensor || tensor == Py_None) {{ + return shape; + }} + + // Calling tensor.size() + PyObject* size_result = PyObject_CallMethod(tensor, "size", NULL); + if (!size_result) {{ + return shape; + }} + // Using PySequence_Fast to improve access efficiency + PyObject* seq = PySequence_Fast(size_result, "Expected a sequence from tensor.size()"); + if (seq) {{ + Py_ssize_t len = PySequence_Fast_GET_SIZE(seq); + PyObject** items = PySequence_Fast_ITEMS(seq); + for (Py_ssize_t i = 0; i < len; ++i) {{ + PyObject* dim = items[i]; + if (PyLong_Check(dim)) {{ + shape.push_back(PyLong_AsLong(dim)); + }} + }} + }} + Py_DECREF(seq); + Py_DECREF(size_result); + return shape; +}} + +static PyObject* launch(PyObject* self, PyObject* args) {{ + int gridX, gridY, gridZ; + rtStream_t stream; + const void *function; + PyObject *packedMetadata = NULL; + PyObject *launch_metadata = NULL; + PyObject *launch_enter_hook = NULL; + PyObject *launch_exit_hook = NULL; + std::vector> tensorShapes; + + // FlagPrism: decode the optional hidden argument after user arguments. + {debug_abi.extracted_declaration} + {newline.join([f"{_extracted_type(ty)} _arg{i};" for i, ty in signature.items()])} + if(!PyArg_ParseTuple( + args, \"{format}\", + &gridX, &gridY, &gridZ, &stream, &function, + &packedMetadata, &launch_metadata, + &launch_enter_hook, &launch_exit_hook{args_list})) {{ + return NULL; + }} + if (__MsprofFlagL1) + {{ + { + LINE_CHANGE_CHAR.join( + f"{{ auto tmp = _get_tensor_shape(_arg{i}); if (!tmp.empty()) tensorShapes.push_back(tmp); }}" + for i, ty in signature.items() if ty[0] == "*" + ) + } + }} + + if (launch_enter_hook != Py_None){{ + PyObject* args = Py_BuildValue("(O)", launch_metadata); + PyObject* ret = PyObject_CallObject(launch_enter_hook, args); + Py_DECREF(args); + if (!ret) + return NULL; + }} + + // get kernel_name + PyObject *kernelNameObj = PyDict_GetItemString(packedMetadata, "kernel_name"); + const char *kernelName = PyUnicode_AsUTF8(kernelNameObj); + // get tensor_kinds + std::vector tensorKinds; + PyObject *tensorKindList = PyDict_GetItemString(packedMetadata, "tensor_kinds"); + if (tensorKindList) {{ + int size = PyObject_Size(tensorKindList); + for (int i = 0; i < size; i++) {{ + PyObject *kind = PySequence_GetItem(tensorKindList, i); + tensorKinds.push_back(PyLong_AsLong(kind)); + }} + }} + + // raise exception asap + {newline.join(ptr_decls)} + // FlagPrism: forward the optional pointer after all user arguments. + _launch(kernelName, function, stream, gridX, gridY, gridZ, tensorShapes, tensorKinds{', ' + ', '.join(internal_args_list) if len(internal_args_list) > 0 else ''}{debug_abi.call_argument}); + if (PyErr_Occurred()) {{ + return NULL; + }} + if(launch_exit_hook != Py_None){{ + PyObject* args = Py_BuildValue("(O)", launch_metadata); + PyObject* ret = PyObject_CallObject(launch_exit_hook, args); + Py_DECREF(args); + if (!ret) + return NULL; + }} + Py_RETURN_NONE; +}} + +static PyMethodDef ModuleMethods[] = {{ + {{"launch", launch, METH_VARARGS, "Entry point for all kernels with this signature"}}, + {{NULL, NULL, 0, NULL}} // sentinel +}}; + +static struct PyModuleDef ModuleDef = {{ + PyModuleDef_HEAD_INIT, + \"__triton_launcher\", + NULL, //documentation + -1, //size + ModuleMethods +}}; + +PyMODINIT_FUNC PyInit___triton_launcher(void) {{ + PyObject *m = PyModule_Create(&ModuleDef); + if(m == NULL) {{ + return NULL; + }} + PyModule_AddFunctions(m, ModuleMethods); + {cpp_msprof_callback} + return m; +}} +""" diff --git a/third_party/ascend/backend/lib/libdevice.10.bc b/third_party/ascend/backend/lib/libdevice.10.bc new file mode 100644 index 0000000000000000000000000000000000000000..3c01c89bb74a92536ca555032af75c6583f2528e GIT binary patch literal 84984 zcmdSC4OoT3~ z{_G=s<~!%ybKUoKUFUV4=XKvb>2}YDlZ=9(6a-++ z1>w*fK?n_s6N2!1DE_J4ckQ*cDrT$nNmHJcRcJ>mHPNw(rx&|SnL+8#hNidZlZ>ev zQ&vdQ;;LEBmbkkzA}V+5ZO;ZRZPmPH6zwyr#G|sL#TBM;@#Cr!w>RYP6%V)KxKE0L za5w%YjjRe57tavl?ip2UiW`kz(POkb_g5DzUYakCT3#lEJ*Kiw3p%0sRO7MA;xx*U zpK3(qeY>styMzsj`U2H5*=}pXE>Tlv2{P2WMa{&7P;)<+u|#W>MJ)_S59uKED_aPt!hCh9E|yJU)V()L>Ai>7vt|X zzh577&>h<@t5=DF^`(pk3%(hnmgQ^vQvsfwbG78g$?}cjD{he05 z%%dJ&l`gyxHny;zXDMi0V}&P#ewipUwySL>O~sT}dupS7N~0{T*?wP>d|FL;n($?h z?n{SGIHkMPrCaOKU2y8UJ-W38;TJ5r-X7gQ+g0CI>&{!lFA5%+waGWNQ8o<+@yOHm zmrt#;4{x-mw#)Amdd2W>aFE3CUcK&<2HjdA{GUC#UWe`zYxq}U_*zl-O@Z!H2aoo3 zweEjl87nBJb(Bx7u@B!ZPwS97br-sH7sT*O7M(}+v0iwv(<;iXA;as>^F-8xrzwOS zkk^M#uB-SbWx0s6^T>MCE`{U&H)pr&*}U)3_8o%46YiQWjB4+j#u{+|LyOOt(ykTDfQrBODu#y6!xKDE(3 zt;Uzy?wi_bAMUr`S1q4bFBjB8x9+QU6-H6F5Z+U*`_idf+pg>8sB14z>%c);(%S8* zyM6az=(hT%w%e!dmQQOgPiu6ic9c)4wPRopZ?oTb5tgNbzSj|sA-&e3>v4p8rqp1V zm8b3Zr3%BVX9z=GM&Z6rYibZD%V#_vwe@|gHOQ%g|C$Xyp~K%H71mmdHY3q(sgS>D zQ5;OPSt^1vM00C}_@c#hFfmSZx8iVx^+kQc5Y>#8cWa)UXe)&a6qcxjsIo-)y+ z<*v|2G!pIY%UcHm`w*yTb_R!NNMx$)V+rwWGR zvo*2$k)tZK(PNW>R#sT1{MV=dyz=x*L!UeTuimEqjlR$SkZAt3-)oZ7Y&nIP8mLM zvTRpw$=9iZfCZ5!gs)y3r2XxT&$a|wpnxO_?<-6X9(|_w_J)o8<>xZI8<1TB8}`QZ zaC`*ErYy+Vum^U=;#_cSD&U=Y;==-{-9w` zy;UD9xNvMMjyKSiJEjbg!WkE(z!i)98@`V%z=sE}SmD32Z*N$a_IEA;9X#;Eu5`bT zhW+r0ZJmUM)x+jy!@gPl6NFDcXxOoHcXs7Jrr{$vX&1jB(6ARah_`p;eVC|_}VXKnnMw_GN5 z;LS_-rxU`T0;k311ZMGzD}-f&fY#j7yvq|y5ud`L3AoGy18MZIG`M?cE=U92IxOst z=AEg%vw089mESRzM*mJ$Arm|})XnBSdq%k9>OHwyC=>(>>~O||Gi#FTj^>RVZoQ*< zFNWPQYo709=W>dRINX3cS|>ktd*?3utL%B1Pz{rt&3otEomcO->+igJy(ykMnwMlQ zK=a^IUHsa>q?n|8`S#8|6RSRc``kZPQ@-r>&W$|oPP~1>{`G?E#XH8{hU%4qkc_Xy z)wA9k_})#r17SR?fgfpIn58~-^Lx*yORLONm&7|p%aa@R3vQpbpWat>$6($clX3f8 z``!Dxk)eTx@S6iua*KNJ?JniN|9i3w*cBI=NA2cdeq)OTbM1$I z*gCBW-NDgwb1=Uan}NB8*DlBd-BA?I?!Zec{so$MrdhfycqvkoT)|W4?`+=uukUPL z(sy@Wy>pN2=ik1tI;K~R69kOOTblQm5gC!vAa(I;1HJZ5_{$;Ea%H6^UA;5ktLG{O z4_=D+7iiu%r4p`yuWx#)jpb!`tUKfE-aAIiiG7yQR>90;Nk3fovvX=kiBgm&Nm6+5 zmv~Ez{-(w#;&;VHQ@BRJ>KXXqpqif+5~fgiaFqi;T>6iV*x4Es)HNs}La;6_w|2jN zx8O?9{{t&+Qf#dlZOE?P9enr9!wP%E;eY;f@`N*a2Z}c6MlYIIlt1)=zis^J=nU_l z_A95nx#6W@D<K;p@pmNy_S`7S^4IX58ixw{iY+^WQQ;0 z_x8L#X2GKkhn{$4+TK6U{xGWbn`?*Oe((C!Ydg2@+IaO>g|Dtzzd7=mooCk_oisH4 zxsPTCzxDR?nyrT9jSKp|UUXsEzH3{qzx3ljy>R8p@8{+3SH9=}{KcU?|M;(*$IM?h z{g?IZ{!1a1-u^!--}=7rgQ4vYz3P4B!1df$$8Y$>qxyR`r9$|S$|At0R0v;Ec@*$3DunN;I04EOy`qqy26zc! zfNnxKz*Yo^qzxg8HiQ`35Ju3Z7%-YPgfX-sjHeA@0&O+`CeenFLL0&~+7PDGrbKS5 z&H&^PG698z8Gs_fOn{e=4LC)Z1L!8q1zaJ_1K2hKasfGn`G7*gj{rr4hXGzfKHwB# zA)uSE2ylh)D8RM}uoRF(SOzF0EC&=3Rsg(&HGoruwSaEII=~gedVp;+pahUZ*a#>j zYyuP!HUqqb?SNB+#{u1hCjeInI{>y)xy?}lmZ_>OMtUf2;EeE0k}kkaD~b*0avLIY+Ir12h0Yv6LP5h z8nA#0p^(ZefK^loMO1zR*g}QirSdA^C>6pfD*pvIONG!)MR0vn7{1I@K3ck`SZ4LvL(T1>!HiWgb zAr#T(2w)>^2wP}Fc$_u_FKyZYhiF4MN*lrn+7M3B<^_a&4dDcB2&ZUs4seDx zgtN3EoTCk)n>HT;F3^T>i8h2UXhXO{n~wnB(1vi8HiT=mA=sXP&3Qm3z)6@5Xdui1 z^bm4ra{(})HiQMVAuOZ~p^!G80+!K+u!=T>wX`7=(dIM2M%oaz(1!3hZ3tf4Tm~GX z4dEzl2q$PmI7OQ;0B2}JI7=JCIoc4qY4auE0&NJFXhZmdHiRp*`3mq2Z3tIsL%2p8 zf^7$EdI6aLCt)_AfiMTqL&%}cH-P!HAuON`VIge@g|zt>U>R))t7t=5OB+HFZN3F; zqzz#UZ3vIkhTx@5Kj08;2uEo{I6)i2DcW2EoS_ZjENuwqXhZ0x&G&!{v>{xg4dDyg z5U$YXI^Y}H5U$dOaE&$uTeaL~QcTe+?0{54CLoV66HrQ+4X7n#0h$PN0G)(+fIdPF zz@!8`1V|;!2jmeR29y#O0BQ-30GbF30iA?L0eyr*fJp^d0!Sq+1LP4_07?m~0JVhG zfF{CPKqp~6ppQ@lFsT6>0I7tHfIPxxKq+Ahpq8)=&_s9~&`H<<=p%RmCN01RNF^Ks z70JVf|fF{D@fKI{=Kp(*iFqr{9Kq}!7Adhea zP)axos3jZ+G!ae!ItgzB`Us~0rWn9!Kq}!3Adm1Kppz%&BzH6WGn4Iq#3EufTe6;Mm~ z4$wrn2IwSQ2lNqa&*C_v0CqqsArp{CmWJU|~I2VfcvcnFY6 zm=DM!JPar$ECAFJ9sx8F76Lj6j{^D#g#gptfF*!b!ZJV}VFjR+unJI1SPf_*tOaxu z)&u$oMF7(nzy?4nVIv@quo+NF*aD~}Yy&hA9tU(1b^!VaUVte9-~*%*4gvBAM*yXS zqkvk%aX=H{1fY}fHlUAi3Sb%!I1NZ8oB`w!-UE~p&H`!)9{`#N=K!6Aj{tpyZh+}t zz$buI!UaGc;WI!f;S!*h@HwD~@CBfg@D-qsa0OtR0Qee^O85qlNB9;{O1KKBC42{H zB3uJ>60QUK2)5^NoMeC z@(2$DN(l=9wS-3iO@xJjPQs&rK0+bDG#Ri2kV;qv$Rn%(loD0}Y6+_WO@y_8PQrRX zAE5|fN&##Dq!Km)@(7y&rGzbjTEaF!6X9_{Ct(MmkKhHE(f~d{D&Y_yk8lJ~N;nFr zB^(Dd5l#R)32y`X2&Vw1X@Jv!RKgiR9^pMeDd8-jmhb_fiEs|kN%#oRN9YEa?gxAV zNF`hVQX%wHVH{;tLLry|Qb3hVg^^Vn6-HLMR0#Q07-E%BA#A6@SgVE# zp^gfJu0|?^W-5%nI;ar*R2YW!QX%wHDFhf*PzYv#6qqGbVT6`OWf>rs3L&4$3P1@J z!geaF05wzybyQXZ8mSPPsjLNbP$BrKtOxW`A@oxz0vOd$2xfp3;3ZRGq?blzBOsRw zA)m@-KnWGXb}Cx{HB<<7RJH*csSui}JPzofLhw`B0qCVd=%?ZZ7`0FcW`GnZCR1U= zm`3FgAeRaupUM$H2^GS2Dn|h|R0wrcjsqI05Spo+0CZ3x_^G@N=%qsFr*aBl41+>2 z1Ehd6nF=G%G%9BRxl{=GRNe!WP$6umau!fSg-}Q313)7cLNk?ffDS4IKb4OFy;KPO zRJs919Tb8YAO)_;R2adgQMmxfr9#N3@)@9n3Sm2yOMn_GggPpp0~)CinyGvN=%7OI zQ~3(eONG!+2w;qYLNEiQ1RDbt14^h6wo};xsG&lrqp}UqNQKZ$<#9j<6@s724nQvzLO&HRz-Wd- zFaxASBbf>_jx;KV0J&5M`BaVoN~jRFQ#lH#p+cymavab|h0sjp1fYWo!B6FFKra04ae=rov1rjmkHGTq=ZoD&GQ1s1UYOxeBPE zLa3wi9iWj4p_$4xKnE3qpUQPWFBL*R73N>YQBVkGfRu6@s72M}S@`gnlaB0OP$-2xfqkI44tK2AxLb0w9+PA)m@;fD$T%?NlxS zYN!zEsC*7+q(W$>@&%xS3c*k1D?l$5LO+!&0OJHG1T#QNz>}#k6HlY^4Iq~aA)m^( zfD$T%?NqJ;YN!zEsC)-#q(W$>at+Wyh2W=h9nedK&`*W=yD=FG!3>ZR`D7~0=+me$ zqtB&6$fv^8zJv;4I~C^mHB<<7RG92HQXw=`Vbc^uL!1p`Qxt0LDpB2xfp( z5J;xN%0LBP1vOL%byQe1Xrw}Dro!q$2Ni;!3d;z+R0#c4 z3IWE+PzYv#R9Hx+!U{tgm1TfjDujG0D*z=_2-~Tw0@P3;)KOUtXrw}Drm_~$L51L_ zvL4V&h0sr>2w+TsLNEiQ0!1>op3@D*O*iK~&poR*ej>{@2>n#N0Am^yf*BwcI+Cfd;*mz>5FnQdA)m?-KnWGXb}B~!HB<<7 zRE`50sSui}oB(uCA^54h4d|sp=%;cDV4MbpU;RFy0S^UZJILTC4 z$w{N~4Iq~aA)m^(fD$T%?NqJ;YN!zEsC)-#q(W$>at+Wyh2W=h9nedK&_8`j(%~Sb zps?v>HhqOvuSkSHvfzH83^^bb&_ZxfQ3CR)5S&!hfKn=iYAV5iS}KGFDxrWTDui|_ zVSr95gdQs4fIccm*b)VgZ#1GYg${}~8AL6`WQyIAXs2kTm@8T4&@z*vnPQ1#DTy;F zj-ps2Sw2Y1*%T8fHcFO~m_;#}Vuxh;5H06WOrzK@8-n#v+mWpQogmRffL zACT(OedEzxw#XN^;<2!O<~HfKztDUv&fMaTIaa88x0=tO1fE2h(&S66v8Ptar`Gr; z;7j|n~G^W9OM-7z1c0kfi~Hk40mEFWIsOWkds z(&kHR@=n1+IUVw}wpoouctWl~U4#pmR~h}mVcq3|@GrVlJtCeb3&#Uq|8nTo;Sn)B z`eF(Hx=V#eh0Z&4eU`gEzFcTN=#J^^j=CHd^J!d6OI!?k@M7GE<>E_Io9)w@~tfEwrWXf1{cuP`5~t94h{ zrgZsm>87El+N>@0c+@YZ%NF~g*Svhx5(^)CyKD)^K`vO9#Bp?OcBxZREAd5JFwEvg z*W2;*630)2Z|ZJe+I~DdCl|Jeg0KtYyEt~1Xs^;n2Nx&J67O>F-n;iyM!i0MXw`ebV67Pp`OAf2!0495V|OAR!!8- z?P;@E)02Hd*eR_%BV(fCkk*^%E2wcN+g)w?Q4E&J8|RwdceiPW<|dJ+D6VUu}nWkH5TTo$+IN zxkX=|+u3>f{rq=YkNGdo2tJJS+zbDQJM)FmkU}A7N1jC?+$;PR+Gt5Dj^`lTn&+&= z_~y`@((Wt!)O3zi2*R$ln*;u?Gvy?fcBFfyq|r#lvQJ=FSlg|BEIG zr*MlkZ`@Y5Nso8Jf#LDI@-#TJY9H@$yf zTV+8(K_fe?GJeJv_CS#8b*){mO-TqD)|RAAxA;BrOPza)Plvhh%1*JIwD?8KQfJWX zGm}%!>5aY;pF{5pJ7`ypS+h;C`E-~qF*`-8O!SM{OPytZp1J*2qh+dSz1uFxUsw37 z>NpLavQ2NmmlleGfZOJUA&Oa2J8>Jf&}h^f?h=%mDu<}lXf!!`{*HIJ@S~aJuqx3l zXzOk{+By?$O%6v}&FRj$<dj zKI!IaZ}8&{GyA+~+#xj1*y1dAsFzpE45CCjPp&|en%?4oI%qrumuYeDB1ih<`1 zhA3>(wc~imsCn#-$KFVPapH=T!)nG%xPQCVW}J0ZCTvbF_3rzrqDc1gnAB4D%O8&tnEKb|6qLQoJf;=i!`knSynRMLYIM(2C@#KWiSY!R|kA5HcpQ>ichEHGRzXb+k z`j&rw?5hkPasO6hh8rDsVTADKcznZqiX!jJV-iYjzv6e<%K2UHUk!eX^j%i|j@CQf z`-e;4`RFa*$?@7~R2Y=DtYCeP*l3g)l=iG>WBOR5eTtq>i3`FNOW@gx;i`!iIR?%| z!(wZ{J$|e_sWL=`XRV2!tW;}XDU^@Y2X!@TCu;URU7oBu z(;zcwWX{6+LRDNz(Sf$Ez=zjq-?9Eg`>rq{Pd;s;C~vM0>uq=WLO--l%zOFi@^4i# zWnrsJqiydBidgXS(>_zX&1y3xi9zKlP43X$b5=fiuWnVndcf+-F4alRY$X) zq&)ZVHPMZpus5}fx>S#-k~Mg@T}AOz5g)O~D%SC0;-^p)gbDaNMDc($p4jik-NJ5z zOVoV4!!koC6RWQKOfpTG*l);@-z8k&L%9w&<|UheIT>8mmNW`?eAbKJ1p7#ZFDx_y zg1uU6i&ku!3t!Y#thz9wcy10nq%!&d-xw=O2V_gktWSZr_3F#hokz>c6)sh zCQZbxcJyI`$K$g<$#V@ZJ|{F~_p`qEfnz2N9P>Zh?)8}lj{EOyrSmc>gcL=kyIM9= zTPfQsGigWCrQXq2wjrYiT_qP_zqPzRO7o{dJGN_x!dqpNRfwih58e3u5A<0N-q3Gt zFOSi(M-XOd<1C8{ggA|AVWh^S7inB52G%DW+v!6O8SLhqmXdP4Hp+P8STWLf+oiF_ zr+2VEbx2<39?8Mj=vmb8Dm5ddvBkA1U0?8Wg-7}HqI!>V$fAaeVYYnF#bD#2u6UJP zn_nHL3ex85FDVU+;ggjS((mUMiz+x!QJbqrF5?;Sn{XXC@4$ESVZ%mVpX)PahMaUF zyiIzRyXLw;2V|?68CPR7=o4`vePQ!B8q@XW-iuDV#qFJ~pPQn14&!o_FLTkfi?OH_+O`OdJhDPx8-%M^C86TH42=KKUBZFEC#HBD&h^3LaJ)3;3b!AlYtFzi@V4~bx6YdzV+@0g zbFDsF@wdS-hG)!h?0`?<_<>K+=SW`T!h_$-FO0ARS-p~1`93tG=6BZDLOR&?r|CLu z)!Nsh61E;Ji%Pi9ueV@nuhw{-<(_J9g3k*7qLiynheZr&ix>sI3mr4tATDVUlDxtB z4Huohk#28-e6b=WZM*wdzA&2-2CasrhRe#K_{0{))U0Z4hvEsL&*lny!m6)yt&z@a zjnv+`(zxV8^B;WRb@dVj{`)6=DRhija1gT?C0}`PgkjPxXT*UeS1wUp+gf%hBV(`O zVL8mYEE7(CWb(%T{>t|ICh zds0?Ay5{d`jDSIgjj`R?WeMpF&X(j^XbR#eAvX z;=q#6lP>+^AN{Fm@{H;upIq9uL=ivy^!|Z98yI$b485NRzcR#-IU_&JVVNPm{_Bu3 zugSLhrIeO6H;gLCd zV_P6b;k>-3B4t*2WRQ73CdF?J+uW~2E2jGOYg!WIk?6V%lj4AY%b5uWWzB3*Y!&L- zXGggAn9Np9nJqKI{oXl!d!%f?^|Ep~X4cOYuSM<0#l5UthOV~7i5Y4OBHuchDLbVO z52Pe%Y?i>~d_LQ<#;J;G<+WdlzIE6!$m(TD$E^m8G?xmG%ehtFVDHH!BZg8zgWGb> zqgtzPi*QuF%FnE~WT!NEhbdAN)iSK%Ka%YZG->ou@9?zO6-TT}*^tAV(enQ(NH}86 z^n@QSI zxG>4PQ#S5mf_tLF;IU1df~z|>`MC}06Izvj7o+;i@&#S7Ewxqou^6?xJm0D-elh{0 z(c^-fIFEjijf+1$(P8uKOqw>PuPnbOv{(k?ZZWE=Jl~=#mQ86@UJ!rU7lbPlgVbWy zJd4B7r~Cw8<4&~{bfHxj*s5G3`!d6djy3#6RT6Z6Yw#${a<_@!6RCvt>QV2ag{}6o zX}++>+I8@2);V5lrX#x|*@`j`BKB1g=uE3}57U$kJ}C)^rU~KA~BY=xPa4 zRY%&MOkeT_tTSNJejQk0ME{xIO z_l6w)mvX{!fWSl=ea)I5re}z#r1x_e0*3f}GZb z@=CF4$=N9d4?FEbw7$kx5b3G%a@gdK zGOSr*$#TF|HCneHvSLoPJc$0-YBgu-+zSouVeWi`HKb*cNVhawzs(lCQsrG#pi1)V zjY0Qs5#+uw*)SI#(;TMoC6wo&6+?&BBAzW)RH`0Ea)Bu+TN>gXQP7|gon<31^m^1B zdU0OEn)=rjC!F@rROnD~tapo3_;wG+?tYHllM(qQkNaeV1>?+|Fa+bQ-j|d(rq!Kq zQ5VZ%T9q5cs7`mj6E9+0m0QH0o|j>qY0>?u*%)W_K}nM^oM(v-Z73Vps{Fke)lin- z9#ZVSr&ak!G0IbxuMaJrJr={covU64hjUhZEQVQCd625jX@4(-t(wMRmR5mpH~wT^ z>X^o|x&j@BSrvzwG@g4}3h_PeG;`iWf4%Arj^_t&9nZ>=^wieiFR?;R{UHqTQz+hy&AkJ$xdZwT#T;pj;<@zyv=)EAM@mxuAq`y3j*JeEM!=H$ zo0UI*6T|i&*S9{gM3LI|^w%>p{=&M)_cFYRHCvSONMG1b%C#xhxzaGpMnwVF;kXBJ zhHsSeb!{62hXaBh3$hYIwqbv!L*=eBlS2DjP!B>ny6LdEm z+iZ)!XA7zU(Mx5jP)%Dz!jXD6Lcig-*a;6<{8rTxF+1f|cT|DkMhzjUJ-y2270X7d zc9}*YV_Kze85q;|7%F6@0pEKlBAsEKwBWwcqObQx7doRVWJ~pz4emCKu1W5OD!HI@ z->Z4W(i!wxAWqyQT}Ov>9r2$Zxbr&74XCnIU3W+m8L8cathPX!<+v_Gp0|VL7yTis zJ?cnU4g@&Y5kw+>$!JST!@RvEe(AyWvuEP+ysh7#d0F`!&f18x)@Kf`O!qHcFqw{A zU*N(JQ$^$#BKj@B(DiLlytlE)W_?}p_GptXCEkQ}VN@1FnI~gJd1Om z>-XRz0w-H_ylj`-fUNMiRT&pohA^mBb}&BPCDYUis;#o<%V8Vp;}t2EEDzkmD-A6+ zjDLOB@MOf&;}cf2Lu_iOu$9|m>pa0qbf;`sc8bSq$S?4el$-Ua6CLCf?@E@%rfbYY zLP}@JPlsi>L|e!hXAnGlx=S_2WI5+h94xVq5Piz9?38|Yju=tlDg2%5bpraR6*K14;(dkY&zeg7DSJn$ruCiEzRpmVyktUxen=KR=s&|Uviu=rE_MwNRkVmTX zJ+jL>3(0!ie(ww2`gf4u&p!LYKHaD=bfq+J5E@~g%#`|<^X_KMYR$O%i_cF^KwQHx zm9YxO{dH-+{BsBTuF~C>WJ2G$%$Up;MIc`&STHjwt3aC)ADno+00ZCmtm1X;48th< zW#vO3M^vD_u`R92vVjpfDcswhZ*4WW2VO4=cMFIBVg_G7AC9;}-#T{i^`UTg0bVD# z6Wb#+J7a~FVv`qMnt6MnC zvTSaOai)~vDbhVcpayVzp4wQ0{F^y_yjFC%g`g;7txPsh5BT@C$m`>^7TC&fwY@dx z;YxbrxF(HwYYpJv+ex(o%dK|oz`?vr4!+TEQqEiM!u*gYp0lEKX}4-lhwap7r--(J zIUUHKGo?Ns$g6%>tC1Jo#iUwFZS6`_dedxFjS&3uYs0=8I1qBhUb`33#q`{PpyJ=f z6o|ELfmnMb@}?-oS-k#@6nmHA^%$J~;x5DM+T(EYyS-8H%1E)Ssu2^9J4{$8>qA-= z!#C7h6An9(OT)7>(P}BI4r1Q4D?uMsWo~}Gs>~Pd@;J-u z2S4flkqII5zR1HxWy3sS>x40i5}U|UmeLzO%QOnhuG@GN5sR2rTp`X@Odbg0=P61H z)G3OROs3Y}5eQfAwp~`fjL>?UH@4MjcW5+!2|{A0cw9g_{J38Gy7qyHQA0bGPlyBS z#9wxnhj_Wb7&K2!*YsQRyFw7|^ve>DwZVzcqu)l|kA&`aI=;KR@%RqM@fP0x@sJV`(Fj5!@bn9WP`H z)6_Ngf`=6)b}^6@)|3abLYIzNVFgxRWQF0=m=)HCA}h4$kQM5Y2$%2FEk|5D+xwa? zGT@bjof*~PrdQ>fy6W)z{NhD9R)#eXIWtl(CdelH+mFhzz?95)#-@25|Gpj@6Vp*9 zurF6esV*Qcd(WCFsNXFK%hmT3&Ur8Rf5fQE<@uiQV%zjq<@d-9%JYSYV%x0N(O+Sb z$hf(1HWFdu0n40DtI|F4bj3@Sh+^67R^{lpsM?BrCtf_zsvH{^b)X{uPvKbjORx?# z6k^o4Mm)bcELZWd!|k2qZ$G49tA30K&Nyr@zJ2HNyUzwsiW^gbbw;t6;nhn+@jdY) z7j#-P#qfw-VJQb}a}inDk04X9c?};}l{wB$e|o8YNOAg)ag}u~ zcUK)dI8-~!QBb{LSYxz5ux4Igkjc9 zV~BDhQfPz~-%ily=y4b+_E+LBa^h6x#C6_~INdu1J4N*hT>77swX2Rg}7geY>#%=8HC9jxwCc5n%NaHD+b4KSG0L2 z`YnH9RGm79eIMpUn=<9Pw&1e59yNwb0glm??6uAmRO%Mf3FTH{vsxB2LQ%9@k15@t z#=INB7|NkW-84qyO^|U?*^;Iq379_W^oBAqN>Ob4;jiF$w61}z^Q3vLv+%6*_BmXK3+aUwW48JYipPGP3s^ST!_R2;->ngJ^wYEf8q10+sW%Vw0JvpY{WjbeCJo15f@vu^!-_vzV z-DYH7`S|jPM(+#m;P-KG{1l3T|3=tgj2~FXez@bH@kmRDCPW%<{a6_p3ovf6Fm5s7 zA}y9YmWW8D{8LxdlW-kFW+$7re8 z24^3sn3s~B%K$ND>3pn?3%tSZlS!7DmgBBbzOljcQVJH}&9<5Ns4*MS&cXdxj`th$XK*hC3qE5S?6l%M94PpeC-OWP ziYX7C#Qtk_2o~kZSr%=HqRb`B8>snikdB)z&C6ln4EBrTj;uYyI7=235m@}MVuDVW zG9UY@n(H>%NNxr9TPfc~i*B3j;Er6e-d=2JM?hwXL63!Zm zY?bWyxP%t5zDpmiFPdSBN6CF@L&k(%jn;Bk(BhIMj?hW7q^?Io;sqFm~?)QMw`=(l)hYl-iA!-M1)oz;}&PRj7w4Aa=K6|{7_J%-s^D(G6uisk(ugb z>Qb5cqSD!LI?SopS1PvI%$KpG`Kw!KW}4|d9(di*%Gcs-=e*c1<%W7wn@nRiUN*c^ zKSr_b`Is){YxSGYF$vqD$Ch+aR>Q-Jr)CLE9h~Kd)r&31>l3_dx4T-^Hm%vz<6h)Z zKfT@B;t7_aV6d&ijTP&a*D~(TLOL5cDER^U*d3>@mGil=09ID+HWP;ZE0oar%_ z=BWvzqRkCbJ{+iB{`)z=GUgD)Y^hKGZ7dhqcN?fVi}5fYm;>U#9FP|W=782aKIBA% z=nKonj#=_-B(=UU-4Ja`TDFwmX4o+L9TWXx=D}duyp&lvSk9G)5>OH>FdugM=1)SY zd32q=r$jzO`MRQY?JyimKHhE8Xl98jrrAS>$YIre6N2Pvg#%1L*h!*CyEYh!1 z;T?ShZ!S2zK2}5O=NiO@CHbTuN#_^n56&^1-~YBYRu!m?^%Yl$Kd(uAHTrT$=H|qL?D-D zpNh8D1UdImhZtFakS-K!V0ksz$(fJkpty4sYA`r|b0)0wasJN0`LB`A-#laB{GIe3 zafSzr9zbw!GwH{aIxEzt!-8Tl!^Dbn=~%9g9CBMK&sbbi-D}6dAbIr7t0=&+#gQ z+agTA)tAVIVb8C}jb$TxH!|9WRo)W0N#V1}@V0!sEJ0|{TPqw@?g2l|c`>+ta_$HA zBL)&NyhiQOTfacRO|~ye`*T|ZOR&<|bK>t1MV90XMuIbHvaW_pzF5lP2IH$%>{6Z? zF)*tTcOgrb7D#8bG#AU>msU;GX0OX`T*%g2H!iAB46wa;Z0Y~(!Pu^Q9otbq{{0)s zz|Q~u#GmnE?#m^3hkO6mEnNO@+;73S4NI5sB66u+QA@YlUzxe|KEK6TfOLFiq0?QD z^h1wL2p>`@i04FAz6Bv!iWqCkFT(j z{a_CD;?5owO3z7-Aq@&wNZACfXww8w5Qi&Ta4mR&urF22>%yNxQ4rSS?-0c;>-6Ay z`tkHz`|!qI;T}wZ%iqU;jcfnqPm-<&!fWg?HN5-Sc+@ZNTYd5`DKomh+%y%NSoi+q z$RAT?=>PWG-?nB*>-u0tGgT1AVt;A>m|0jrv)Luth0b(To=lO*Ze@vnr(?MT)vYyd z#GrYIrRieI_#U=~ECqD`~}91 z`a$nMR2I`K^`oKF9XBcx6+V5llp9A>$*rLk^#?M;IE} z#9dEHR9~ffEDNyL9m6i53X7WTbc)|s@5A36W#f8SWt*=ugtKb!95tR*FQZ)6EDt4+4a{0ICSR-%izM3*?pt)*5&jKcHWssj86=9 zo-%6Z2&~t~_N5!6YRZs6|5jwE$KB?tp!|l=T3f>kMTuL8YAnl-Xb+a<3!zEaNTjN>z_jDH^d(w)N&jUOCfzlQnX0Q;Sx z^WcQGhE9 zA?=yeF6Fp=oA;I@4vjPXFr+>0EPp|VkXD@IU3^@-P8YlYe;*sLJDcc zG=#Lav_MFkHgKrHkoGh~TJ5;h$AsW#Cu9tUv;+I`|30Rzmgdfb(%j*Fw0Gt|Hx_kS z2zC+o7@DOz0=-|bN?I+Mef%eH`TvTq^o4aIg33V5P`wZf|0XFN)Cb#&TNYIZVuGKc z8Wo1O>IY&1Y{Z}tTiAf;_*Bf#WB-IL2oFL6pJyjB#8_b{SKEq33o$yqj&iYt-)@;9`*Em!{9=d zno}zq^tN-2%Lj+%uZQf69vqrqWE83fhh|z%t<93v;~$<;WE&2L+m3l^|AY?LJQGZG{ZTR;NH{ZC!Ip}bs4viT6(i53iZp_(|oc_T%JAa{Z$6yTNt(JDUEH@L88zB?! z;c58pS8W07e(Y@{>OpkLc&KlW;o_@PKKsz*{@Vv%T;5}7N{ju+hbG(Q!(aZlAqTr* zed>X|?Jg_se-%ZWlxo+^xH`fprppdWaoezEy zsIOw=#Ue5~#XxBrTqPW{zq+CfFsdg{!9!$|bu2+pw+;>j6%267P zJCM(-mpdY|Q?eN>xXK{LU?#Mqlq|16*iar1r~f40uS$3?+6^Jc5|IC{Fcr_d z)75^FSrpPX;A*$$)D61Y`HvcZGdNZ*L>(GCIJB<89P9Fi2K^!BA4J?fF#%_?X9#bC z`MnTzK#m(qqGbz3kMd3Nr;3tvA?kctTSO|HaeG0Vk}J&}lc+Jr&J0)|DU>|YRfdt?%%B|mECtyc@gfv zKj8lFyPRN~?5{pRGsk>q=v3)-!>fFq_pozbT9@*2y{RGK{w@3H{-<^+r|#SArTf3< zHurzlS^lGNxPRO-vGyBRUO3%bN^@Q{v*aY z9xDj0osb`J|NpCc#xJBfyFi+=vEbY413uLB=;p8OF;xG zlV)vyU{XR1EC^C^BgRWH@^hGNm{ZqL_L|ZD?f9ox5AOJ5Q$KJ+9w*Th5_MkG=ox z2c5os`@AJS5ur=x*k5AcG6}*)=GT}+PLB%TI1Kk~Ww~B5x5$=Z#PB9BY;ab)v2=4M z#D#d-UZR}j;SPUZrAkSobpfMh7+L1qBAvKzP0?Xhc1N7f!0tnPuLDzIdUoi}@^OAa zwNVJd_Fks-VTJQ(T&*Wa%W3pP;0XOV0xAj^+C!eTpBt-4%yMF%Vx3{myVlHl^{FK| z+Wm&{Lv??3%1t#vnx(wKYoP&)URycd#SHC&lo;qvb|9wEAC zJ`{tLd*cIZ2KT_Lt`FYV#C!X&JoKi_uzp`W9Jl;QLlSw{`3ZlHOT9fAlNORdYF!@~ zuVuHDxul!o2S$3JHgR2jvkH67E(J6H=Su?@4d!;~SJdqf;zWOo^`SnnK71H7xIQ%9 zSRa}~rv}!CX3vfF;gwi@kuA1Bl~uB|=jrL1_|yzlM%^j(Q+)6u;ndUK?t2g}{jM6WTs9f^X&f$*C+ts8q}l$Uaac7vGYGvQy`m+#<+4DxoO3OJ5_dqtS6hftWhDTm-OBL#xYu>_uIUjT{=hntG z)@T$}#$$&R?(hw?ZI*++7UNRC&=%&*z!Z7+Lb{Dc69PtEFv_0qP}Uh*6J(Pdw)P{* zIoP&Lo`$_EA@+27cYAn>%RR+qRq@`jNLhZd=I*;>a&E*<4IH5dM{vx?9#2>0c}Hb9 zmVkUzAogTDH_2hRXzR_x(W)|@-f7RmH2l}iX;|RY+&wXHpy`1FiE(uDmUcsFSKvJ9D;C8$ug&8|Wocv$)8j+T=VF7#IV5UeomTfEcm zb4PVlEOrbLosNc?9Xo%azdS$H9aXXG!{>)!R#I3cr!2M={dm?YEO9txj1ACtX_kJU8T>W3~ew;F+`q{|mK%mwasc=e4RACEQgF8D*!xmG=go4mQFWu4uY zQ0_q|K93v2wb__>^D`ErzYB1GQDzwKEz4vm=k#IsZ=f3U{B+IGf!hVnOxOI#*)o{l z?#yro|5Q;t8za9BxAEX!2i~t{c@cZneAvd${h2a+lP(&+!lR``xF z_+MvSoe$g%)maZS(Z#mIW?cO9M4P|1Arg*1$@J7rjIjw{F!1HsB#e?D28$%8a z8w|IuWF)F>sK2(WN^QgAReIF>6zd99V-9&27N8U{7m=kkkWt5sDl3cFIotGCXMVh* zWSu^wxLiL5x1@Wtjfy0hC!;GOC_#Vl#&a#F(4T|%f?L45ojd2cAm2xW`a&QS+|<#(gL$ZBMVphWD-ap#<(M43z1%Ww2Ijylx##1+GtLsQ&}k ze^6un-{b{%?7RuWU*P*2ZukAo|4ZMGIlURR7>*~=kGoSu3udcReJsJpkyR0^pAVl= z0C%yqT8%s0s@1p|Jx7tEEw&Ago#%Hps2?9J>@Co2)0-0q^U;-t9=CHg0&n&}O2QM? zwxk5yLGI~8EFlf7=Wt7jG_dB5!?(>FcXMEMgkTik(}bm>!JolfN<6rw#1nGoEhWz{ zQv9|awc22Jo!A;ymR}S)(3Z%RH;Xfow-jfdoFThWocUEwS9&#~veuwQV(3$NVhnC0 zq1HS5O=pw|dmX34KGP%NJZ{Y@P(NAUIzu*m*S+QyI@wFZYs>v&OEBUr3rhOyg})lQ zb_t;s4GiS}?uzOw`>>039H~H)JQtU?FUVBa9a`!}9p~>j?BAt8-itbpY?8lPs@tSF zoLQ*j;B|wGfqce;&bLv=>8>|b1OoZH7|2ILx_95^=NZUAN<^eS^r0MoK@0X3e>Tq z!FSYg2A&BFY!(F4W@O40e{222x&PCNgZlQ^v4`t%e{Rin>y>c$!Uy3C6EJoynb>cM zl!EmJ4Hk|0hK=45`~O4Sn};`1y>H{0bdsifk}gRXXj+z1D3Bshc9QOtEg&dMsVQY& z(*m+XrW<9`qOBqpO%bqSfv^b*V$ubQRt!kBYQcaY1q&A03bMTSnard`KiBv3e&0WS zKQ1p}k~1?WGtYUp`+lBdgv+jl7W9la5F!Mp1;-DS)c4^?Q$r=MwrW8p-Uwg*8-q(W z6krMzV8^WhdchpXkwXFI5WenkAUR$ml&qW@JK~#{p#b9oNgl-%auo{D2n7gd-s(k7 z(*U&|FW}Oa)_Aq}M@PV!-Jq)RU&w#)Q@qTX8;KbMP3oAL|0usqUKbYo<^|Ej8ws+D z=9-im72<3&6IgOjyp*Bfi(PDXr3jtnF#WnwvOZa=8e_}OTq7l!O{qJn$0?%2DO0@N zJ;&szUHY1RSB`vjkN=|g9OWwnwS@IL2sGh#`0>{>=)L_P&VlpWpzgUZ{e@~5C+DBb zi|pK~%;BvUDb;=wXW%(I~cCZ3~v z?4%e=szWuyL5ghS>o#2B^VY4tu6g4CN8bAQF{iJKb@Mv(9g!JTJrSIGM=oR|VhO($ znj-3ZFDnT!RzO$1C^r5RC95s2-0LjPv$$GiIb5!5Nq{yKzAPl?JoIDESjfmQeBikO zGSp?a5quD2^~_sIrTkWhZ=4L+z26@Aq-7Od;gf!?{Rp`WZ8uc-eX4Q)m+)LnFuJ!D zp2NKG+)qe%>asOspe-H3Q}>$RXT0!;f~va_;i01JA~PpERoQ_h@gAJ;XuMm6$8D~3 z4-uZv`S?bsMEcM^U0#>`3a?*_68ynkdH!e$U8c2v2 z1OFFU>i;QEah_c-jvRIOHo7g_nf)UGD4c=0XZ!!*P;hHC@s2D-il|Z~yC8*&3Vz9AL%Xj7cJGd60wnMnH1iHYBHA;_3b{QOYt2=PbY2){5Cvn9g$zOwOk_%`)D}f8ULR~X zd;grrzu5bgQ4aS0g=`0V|2HFB&_me!jlfX9aIp8S1Fmjnx;KbE6DaC=_Pz#n7jq=f z)HQL8cR_PYWo8)&vQ=psr1*L^dR;+?yM#d%8X=cjd#WRzWTRCEI0PndyINmpw0K<> zC~{q<^rOePaBu8dFr)J3bkfs&X+b6TuI$nDQ%)I-UnaHof6`J3#$4v9sTOLTileHa@516$os63NpvCv9+I18mck?; zJhsW=B-A3fbVMxz<+XF(H+I*z2wW;iSE1(!n|E>K>IdVq7yNFCLvRr&ku(hK4n(@# zaX>1r|5x=Vx@V+}7K@YPIXEgBPtC@xTps_b(^r>x=^S~Q zbJ9E>bA?tc)f6EGW+e;wqyMmgs{@$Gq6+Q>&5tgMePIDFWK?Bj0mt$Sn55_(3%Dy; zz@+r6bpbEo7jTo1Tfp7O0w#mnF5qrt0lT*@;Hjh}m@MGJ%F8-$EZ``zfWg^j7jQNy zzxl@k2IoTWv4ZK*?1-z3iV3)yd7jUIx0S_PxILoJX0ayC)3;2%D zzZdWzZUN^w7H}uVxq!Qte(&|zv4FM1v496Rh#p`8!-YYlc^1s|KFo&yuz=;usG*Hs z;iY5&gDTz=j*^>l|G9uk zFl46FEbNtBg$+680KFl9I0-%!$Errf98S=L$eTp!1Bd9J3~@q|qCs-F zdoT95QcjMM;~gBsC~9)2svOxKV?N3D^kEd``In}XY)|t)vOUu(y?!2r#F5K8qN91@ z@X>U%LEo%cVVxEAn3lOj(90NQlT4gu4&HW3{$3O_O;%zNO}8qfxYWHwE^#F9BKAc8 zgJ~|^>2AqonopQc`jdv~NS~9saKf3KEPgT9V{@Brp5Bt%yS`TgQ#3ztl9m{I{7F{{ z1hK51n#Sgbd*QPevbN7W{#o_?zcu*|yI#X(nxSX7lmhW0P%QgjRM-EjBY@vg=`~GF zs!fzMQaFKyFz9M=tYL;;oTUu}H(ZDB;(Ei0jV0-xI7hliidyivEin0UD=9GA8xzsO zFqobo$Cx}2G4vKiNRUzK#3#iL%I*mVacj6wa9mG%DN-mZo_AI!&aSOA#7TYdGL$+j zPBr_laawQGD-!^*)`dj^`DkE)GCz8_6lvM`WD{1`Zqch3h=+Zjl$E3xiazeb)={P# z&J4zzVU?(P33)T~?v(=Ksm!0PD66`RRI;IR2;NMy4|y}0RCX@Wl6Z z@&r{k%|0uv$xU9pOW9INJFd|2e6qH4Dijm40If>t;2h!{UOM;6~2gBF#RL(=^U z=rBY5BzqJS)1cGp4f52+#g!9iq1aE7s#Yw|X!c!d9c>W}Elr}4+&Kh7s?N9N>#IIR zgC;Zf-lk;Uw<8E$JL<4B%9C=^3f5np-8aW9iw^T>4>~pLPm2{7C0S z6F>aG8Ce-We10}PhadiaM~WYQzjKN$=sNwDPS4%q%rJXEMHHh1jcHAoTe^-o`q=|& z<(q)CsS21gf@!|WIs%*Wr1(cZ z9M^Xa`!&#CIEg6b9=Dc@p8h{}ifm*R_jNSOwbGMY#sl>%Y7v~{a;5kON6NP=B zm@u#_sjbiC^W^6e-s*^O&EQh%flhYDE>{0?45GS`GO;51$KLERKi6rWvC?ZpH_{yE zJgHMWP>_dM%F|;vnpdoI>2I-k9VD0Tjz+p%U^iB!4(IHUHdsg}4Q)+=uj_Xjk+%1IqIh?lDpPhLY$gD3+nn3h2y-0_zggR>#SZe!gq;RP+YV+9#1o zmI~Kxq0h?up%nxzj@XquDyA_`rPnYPQPh?g@dJaR-X&GelRD&MIeQ|aYCioBU0P}Ns;;HUL{@n(W|Pq788%bAQfGYGgAHG#aDQzB&m zc{7X2n{jI(Zzia+R1lyt2q!d(rX7fB$aB0IqMB0_uD{~VES^sk?`P+?6K!lyQHVpm zS7)}q8G*v_X0o(+Gw!pW~hFhJJP02qJ4~_ApepjXq#52 zD9lea2P0}J5-*CKU9-4S3Sl^$9^@i+lQan08>l(9(nM(=-f(^*mwdv%w>=W=9A3G0 zjnf5lit>QFtn3u0i}#MDKjVjgs;EVDk&pB#aoGCW2%o-?}ZTFZZNTIrIVjc-Nlt2o* z@J}CGPkd|}@v#lDLGuC|HD7<`xkoj7vUmkCGmh-GU#w_=FH9h#T%7kZk6&-s={BrL@T~O64Zv>(3|B3soUOjZCTv zKPOrTp1i=DiS*=qi`dwA>FAe+@`D$A2S;WT++GKa&ZSiy6(?_pvd}jH2DmDgGWI?E zkjMeQldKYCg*1BMvPv9U--dNgHIrE_ymnR zcpi?;=@ZyD4tDqiCmY*h6RS)Pc~UhKEEZ+Rl@`}{8A1gR&f>ZlD17qP=?LG3M0vp#X_bYchBbtWzm5Qy6KH)U$MZ$fDM%FeyGk-o2V9|MHdf02s1n2 zesH=_2hl5>oAGm)$wTSqCKC?Oh`P?r0LG&TIEV_!{ z^^9?F=7R(N*)+mGg>DGsgr!(6R6y?C+D$QtIawM%Tx#XNz9vsF@Vf2GTv-XPtkmiB zX@g4#FC8K9F!!&Y1S5M(oxfK~G zw*D(Wrf>kBBl(lA5(pf6(Sa=?MaAnyM{ksZMQoVU={H$HMdyQ8lnt3(%Q=TjY&Bh~12e@?l$E_7p^zLkmp=P_# zk84Y^D^3b-I6Aw8`6gyTd}gG3o)#TZ!J+R-CvSBjoms@_bqEj_xw9m=k-EIB)Mj_? z&c(AOASp_#gwYqJRTU`BtE~VvzoSbhYObJ6Eh*^9x|!YKW}0Gb=q5)WCJO3nz|4}v z0`W|nFiqpJg6iJ1*rgxK)YE&A&?u(kXPh3(b{t0?CC6ud8@@-;QI9U-V*=t}Rt+Hz z<`M8$hBZz(0ex#7_@aY}q@O8|uvSog^vht?hB!jgyJAulRPN~4mZsO&LDMVn22HQf zo2BVdqqs|PG(Bx^TdX-d+kI0|KjBr6(PCFNYa!q;r5AIzcGt0ND+;oX2#mIf<&%>6qRtw1-HXSvW1Ws+|uWybzB;m1J;JvVcLyqse$}6Z8h1 zsXv$+Hc>^CmK29l0(azUUa^JHK6pa zwWAnGoAnMlC&t-HE*N?{eT_-cVg)fiO8b^XcA62kKJ80Uj;HO4R-h2iNTXHs=I_#` z#W~$*#{)k`o5YmkffE|O2?pPNMI6DVn%xmpp`8PogL>!oitY`2ua^g zAClJ_x~UW4IIJb!^K)2`;*Y3lZwNMZqDj{^)<`u6OZ&e=vT?B(uZnlj=Or@7XJ4@% zh`6fGA`!?s2tM5I8B>TSYko3;iY<#)qJ)AarYR%wRB8j9PlFpyr! zSyZq^#Fa89vDHachz1F7OXHfAqKTjl)^;^ow6M_>^Z_g0f9x`YP-&-7T)}p6)F|C2 zNJhj=v&LMsxUMCNmoWAKDnDMRDD%+6&0e^|p@q zHVnnD=(g7&uZd<-eQWwN4Zbyq7hsP_TdPsNFj2h@D8JPtXk0zcd&a~c72snV`M5d` zq>A`EHF=BPRYlagz;Fc+)sH^6enycgm>DeDVQAEbySraie#hc&1o)pU*0Eisa@m%% zOhCtl-`G9gV%r(8vM6P9%Z^$gCnun=e~@Ol8}EQpTQb~OStOG9Tup?n;5iYNXRaLt zF!-|q5*6~@zLL#nTQJLt+tf`V9FMt*{ z4iXN}*#=XW(0KS<0{VCN(Rn($HnY5fzrgePA6UHqt+uZXhnEy@{8PS{BZkyFV#o*C zd|tWHsf(3fV9p6byE)I=u9kWl?0EhnU# z!aQ#g5&|^7%>z~bW9V1!p)LjR;JMhaVvIeHQB-MMPb>1$l&iF! zi>2hFbn;Q)1kqG=l`>ZElOYifDv{bm3|coEv}EBh{c2?F%DSdn%C0O^P6=DPI?jl(h-DSiY8S#I=SAEILt2URH`TfLf;h6TCWHs z6YC9eS!5W9!QCU-Z+=Ef<%$-^&XRTKQTaOXNC^qF^s*}Td{j(lJU+J}t zE$zq?9o;LsGJ;qyweJe+h^VhcZLF7?L+A$(gnpm{p}+7CgkDdqSA3cJ6K$u|lX?e2 z@6GB7;^S5NmRsMky8o>I#vqz+o^Bzv7}K<<#Xy62bO)AtPcxt?@QT6{-E)Q}FTiQ+ zpLCT#@H(qUW;!t{XMD`}(_*tdb~jRNvHIY0RCVpS69^XFwBr?xF^Q_52X6&(NFSA!e~^iTiuCUufR^{&hx1Sba5XaU)02p( zuRr<)^&KZ~9xWqe*qFd8tpw-RD9wSFHFI=XB)JGt5K)=-s}J-8NR#-gI)SXUQM{J3 zr{Z`&8olxy?L6ZCE343+gLa}6}edEZoJQ#PS3p_}+1U$l*HseVZtaW6X+ zDK#sYhISQNm8Gh@v#e-ge^vEvA+hG9n>b41NH_795`ZtwSHHYiIJic1AFao8QVl*? z)Q}p{BbA~tC)Mned6c zitQ#o;bq*X=JZtDhZz(#0db^QMPcTzkhYo$2`EpEAG1ZWq#~upZ8r3xmF*@zp6wJ@pv2GcPb zuGOk^Dc<_&>zyxYMA56ZE@fT|@m@>0#MF>#xk2QYgevKJUjo(Tf>RN;TnL0-I>R2V zH)fACTHN}UDjKUU8G|;&8c-4s^hjGXOj%_mgW{Ek*}wP6#WU4n*qHKM^v$BY_qD<7 zECCM8ZUMpUQsxoNE^w2UB~8?edD6tyQ~x4O6e=8~iIp5_Vj)Nq04T>vl_l}*QQIHH z7%7oZQId2?SzkRxuP`uE8f9$C49^FE>eNrpcIvTp(5yzA+cZ^=hRv~>k9|=Mvlaa$ z^*xy{1)uejOm~n$Wzfu%#p6kT_wI$HzxzB^gBNG{GbCmtMK~WIqh@CtW2FV4E_;ESao+PsC@7(J(;=&jHQwJAR}S5?8~(#x)KtkBzI?-B$jY9VK2T$Oq8S zY%Mw4khtic3({!uUTBi|1Yg7@aCj4ZM3L6&iA{l6swx}Rkn|XhtSZ-ELNX=_UrN@J zU9lvN>Hg&~Pm&y*jsEF0iOc^loB8Y_!4%r~;1{ht{;46#Hdfrp;Y*S{NO)3B;TUne zkZ^Km%v+eg7=L%8`}NE5j}6&rb=Q)fs%qDn=JaJQFBSInT&!`G3m&?zF=4HR68|KG z*5(Uk3Zrm_IX)@>ygH9-=DE{^&DsjB@-LN>(4N)3;M`_wpVY8Yqf+6C%on1FZc>Fu zxdp}}3XQlknJ6|taLg-}%MYY`@b$%)V^)Q+`2mYdh%-Onz^&?L5eM=RU3n{8qqrxm zm7lO`s9Cpk-$YL^AmD* zr8++jqna|>OW08@CKnI}v3&dnz`(tWl8teH8Ebo`lI=!zq1}~XR~%c;OiDJUJwOui zaOGg6c7L4T1e2g00+Y~U^Zr-K#+D>0nkt|u2f@}`tI+dR)Bem3%=-5&At^7I1f!3` zBv^+KlVI_2n1ucm$u|#$Nyr%rlTf(>>g=7NFbO%8^mQNJBuoonM(d^U-LglM4JHU> zO4r9K3okuHpEtxM$)S7=gRRRuQx7-i!6Zy0o~y54s&Z@YS#+R)Nzl9ulVFqdPKYq& zrWVSx7&I&H7N#)grrK~Y97NWzcTbQ#5{XF|F@~6g*q4Su6}{j4?Zw%XYD9A}I4{)} zki|@{5iJZ;Sn^V(imWrEIFk_92|NbsgF$J;BviVHkw=LMAhQ~SMglrr71_t6XwQg< zO?f5yAl@XTjwAPR&%k|{LF%o^j%^S@)E|uB&6R+SrR*|#%&*O5-O*&CP5YT za!o!Z_g0Pw5RK1YqMcpC&+Nk*rkg98_b#KDlD_@e@x#4`H8=@Y1U) zXR*$SMq zaL1)V#Ndh343AAdEZt^?;i*m0ufguKnml~HN=%4#M(P_d;i?2;)n8K#?bo@i0JWo( z;P!LKic{uljpgaNfbEGDI3Gdm#^vx(UvZZrF^G>rE3|W(yzV&Yb*x>eEzn}?3=ljD zw87+WD19p9`3!dDq&)*42ywu_oe zl@?dFylC1E(gkIeU1Zp?^c-Pfyu^_j_`a^l^q8)Z`yOV3>7!>PJJyicM$lmewK?i9 z*EX?@psOP@7cNI^)1iXR%oU`(n3+3fCH0dtjqadZ%vBv`dYZLj*IKy+%Bx~@Y!H-K zf-x>Oko;{LTXA3P+?8hDqV{pNb=2#ryG$x&3~em<#G(n~9F*dQ!>6V9(#`EA1$L2~ z@}f*{cU9{oyhz5}jiy&CmnFR_W|dm|!3SOE)&l+;`qnJd){%Z@KA#%L&>#nC%gg8UW8K+`8 z5MCgzO=+w=vp8JZCk%c@>!6fUqle^OR zS+oy-rA|RqJb$IliPA0n@SU<`=kVFIFF$;G5S73W-=3)3QV~TV@uS;dPBnTTBvn!@ z#-E>z-uv&O_x@GQ9>_QOldcj7Nd2@g?bxrE3vdMX}MGKhnjOxGiY=q;s0~1)7FRS@xGE#9 z(RIFMj7jYBN*SZ=m<{i9!r3U74AHCZi{Rce=%eX5p6#Z#))euCN^2%<;EV*~97oz3 z7zw&C2*9)spTU$1(d&(_GsIB@9MdIG+1WA|i95H?=2TKThMhXRKC&qK;W9hPWppnW zCYTjvaP7w}2a2lmxFQZ4pPiRZ1Ut=IGuWmwzPX*L^RIUIyi6TA#nfiY@z?Ryx(bu$?XgTDw@7wK(K7g@16YzCp*Qoes(Tu&$LM zJ)$zTgI>7Av9PA>9^5&_!8!dKd3EtcbYrYR#4p1;i5fu4dc#RgwFrBSkgXax5y@L; z@%g*8HNj?AtIyw@R`OpiVKSX7`Mo3^dx;PFTSOhh+FoMsI=lBR{+yMbk8FG?=aJJ0 zSp0K&<=37H8z&&MnAPlUAWnaqM{z&QQZ>BwQ9RV*kUK9gwC$sK?06LDK~OPs`5fCL z%GG1Ch=?^F@j_tYEQ-(KK>1pP_c_5Nyf=x6v+$X!ZG3+aYs;5@h2?xk68%8%Zqgd@ zIol7}!Hbb>qh$WwxJ}2N>w2&#wx#Jun#9~ZPY@f$;#@QbeT-sj6AAba{QA`$kem*p z9@R+QshO*g=zVF>be?JZW zlhY4AL>JtU!DcN(93b_xx_toevOQTKz_c~1@F6VlW~;JFq}O0=6jRD%F`6Rz!Gd>@ zngsChX)T2yDYiz41yumdk^qX?cp-JNjKWtpF;mrIY>FP-D5^};iA%KRSIelWIYxmP zb@N-q#?e7Wi)c2Yb!`w?4p-nPi{aZHR|DOI)5gQt@HeALFquNqO^T%{kAXO^;97qQeo z`fbhn1WoW6pqL(b&y&Th?2?-1bWh^k7BMK75%(i3W}`dFT}>hbDtky~6R}`zc_oD2 zmp&vTJ?3ih8E5ZUeks{sL!~;b^&)=>MI4ZB!q)W%goK0l|D#y5Su;~0>q_r;Bg<+H zzzM#`;;LhZeuJHS95&%D*1<5g1rlKzex-?LyU(LN%Oh8ImuAkkatvR0pf0e<23%LX zOXvBYd<`J4&EL=I4n&4H>YNd`t6dyd4j+dn6;Apnpkub zd05SN$8g*0T|UHZ?;pF;0qqnvq~FNOpcgmn)5j$0e?C40!}`oQwFd`hckh8J_~Z7c z190%>Q{pOE*I)1X>t|BVqdZ+ z!W;Kb0kYKc^Y_04s(RRKHPqPkTj#Oqs^&vgBo5ZT@z>`#w`swYA1*4blkd&}oBheN z56+`C#rt2VC3yO4zPuTRMz`>S3Q6X%EkApuZoo?T>V^~q=cHhA*{9xZz2c`6&!`~_+ zJ{#iuksF-c0MO}F=$%>PlN`Iffon zx_MSJWtB%n3+m*kY(`Qt#e@`EDIIZ)r02Vn6MshTN7{qIzmnw*BWgihl=mKPnljRy zi+F@I^4B&dXLmq^r0qBkdk?7GUID%xlaBEoz9lO77_ZYKuKKHyX`nq*jle?9iW+*R zw8=k)Dd^4aBA8T9+cSKzo#ROOV0x0IvyA}W6c&J~xS)FI)I?-`avspj>DRTRcrDj6 z{AxVXhg06-KZxGZ7(WXAY7KF}kqjCj-#)Rm@>!2sw_9G)ftlc^hBT^&Rdp%#B0 zGnIN(RP?GoLXw@XLxE7XYF0pol+T8)U#}0Co*BoRhb?aDI80J4^;=2ytgM)YoTJMU zQG_@%EJh3(Wf3#mXmu08yuG87ah|}n&Wt7`2pvI97Y{3VpTA^Wf+1O|Fl)yNI73tC z_CSG_`h6P&t_LWwtE&zzSKrfx@Lo~`qjJS;m#pBOqg!W^XvKxr^B)aUniBC$ycUak zGh3^K-fwn~-=ku*2@^d#PkS=ey!mCSR$hZW_lfAJow?~Ao+(a1BwXSRiNy)y;<@~EV9$ti8TUk`Juta6o< zs3c#XR_5|-)!X4UNvj&1r2);c0NKn3X8+t>f%&FClApjpEs= zaXi8H#&DH%ofusz4_ik@@Sei`;p6FbbS4!X&wB#nYxLdZWDvH@xF>O#UbA?f%8v(s zi$|zxU|EAxXNGgXz3$Aac!Rc`^Ag2h=Te)=@51KQTXS8#r|LbXncSxf#IrQJs6%35 zaN$@baZe9Mf72|!E(oafnqC@?=t3220PSmQ-UhfmaUnHY!-5^*ZR4-vT1nKwaA*AG z|E&~%-F9aycCN?#=8*aAvXl9lI%9s#t|l_SgJgdDv%2CHMB;EBncps{J9Tiq z%q6)SH^1y|BrbDDxuBsNJ3s3_%+DoqQ!JTZm3ytcD3;8xsDjLIl}WtX_>>7U3ZvbH zktGPw_J%7O){_9u8X59|X*fyaH19wNR&g95m^Bg}?ubb2GOUv&@6vCrP|L#1u_>mC zp*ypacVDvM8v$J&MIxZR zhRdj^OZM*#vzxLy3v6?<(0>2uOmaW>?M{Q5-Lm^+g?I6v|?W)f6kqk3O4TJPU$o> zXJ2cM#Q!OstAqYJT(O_lPsP?c1CZ_89_8zYBUUzI0v657GyySg-J<1>8CUmsN zg>J}1)DkC7)+~2BZ4Y@t~b${aXpBZwKnitc6P_wk>n0eEn629x%j0R zR?C{m8BM~Ttd`Yyx6!ixNz|`xwQNFC-mlO`R}Fr@wrbfYUPQ}gpKaB$t2!Fb{@SW# z*Tj=Ur-p=HZ`99mZ8{adS+{dO$TvQ+=nq=CFxbYJ~~mi6zE$Fz{&rACaC8?`mQyD_du zeMOHlVNavAr_OFE$WeP6Ts`dTJSx48Nrxgg+<$AFdYCVJ-K^Yt8Rp zP!f#4bNkyr4E8)EeVFZSk_DH_f_KP*d&`0=y&eb$KmiqccY$No;L!;6oHHOdzR(4- zrF{n3svK}4&QS3jvSl9%*=qJ_BC_?I$kr{pw>2y|)i$*E@MfQni{2!%WgQ6FdUDQts9cW}y9 zadKq0ow2NJJt1e*0~01;bF)ty*~%Mh_4IG!kR2JD;3YalKQ~YOP@JCr|a3h_9La+BfC@9Pi;L$8*1l$G?TLBe5c)vbz1%w-rPV5jnlOR_i& z9(nR})jQhc61zvHGkJD@ydkb{Z?4D2Yi^9!W@U&*lC8b}YMeffdlXE7tlq{w;?nW7 zbNH;lzT@~Yn}3j9;K%IDqx*1|96!60mcQimNh!LxqkL;o|2chjKe1Sj?4Ln=iA>CD#i6+SWIdyYIgRAÐiq*8*Wwy`p|>;A2Clx`EWm?&t#F$ zEq=`5>GVK;R`c6gy}2>dCnuZvOJ02CVB7$1jQ@g2h96TrdANxmb0fgFk-y||ktm@* z|05QPgg(w5rcd9wKr33Lss7xLI;@d*72-!xtHm9)-YO1g@L%qhPA76tx%fw^k?$Y) zdjhTJ=25rYpHgyj%}YxZB%;3!_qtYuwVu8ARk`!*rd`gn$FJ<}Jo{SO2*dE1H2?cn z^}p;_;>*r2rc4^R1#*^Bm{8c4>>%216wOVsc(rKeG=7Q~-<59Vr`Xbmp20uPH8LiU zo8qq567>A=Tkpzh&p9`(uX{}rw4FcuSus76|AA{Gt^C=dslyHY*?CIe5`K925A+|4 zlQi6u+EQOR)~fajt4>TVqtnW(rkLsE!?{zARcYryZ+C;Bx6CM%XQz}LPBf5W*eHBb zII~b*k&?P=H2X<+d{Tx_gu_F`V1;2!4}4<5C$YmDUVj;%1WJ*oEcZuqrV~{o=!mLf zVqyu?ID_(q8kD_CQkcHIunSj}Xko%3IwP}Y%0tgX_BsC~?cmN|C3Z#vAkXgJZr zBlJB0K{ggc3)|vHvRYUjkW@Q9E&A0B)>LTWmAjyYwNr=|E`62P!trU;kmVqkT_v|) zo`Bo;wUXP9nFK975@BzHzKj$UY!ZzkT9|r`*TPwyo5!n%s{@=`*nfNE>grcnEnG#; zXsjB=YGHLi8!fzRrgnB)EnGLVdN#CGM&`KLtyZdFEB@%! zQw=R_Q#iD+G+=6&Gmg45iPgdrQh%JxFO>Z+YT@z6@~kV!?}{c`c*ItEq|s$mefFrb zicv;pv~H+1H%}F?Bltw_6LC!DS5n6kX?}ay7VQoRar@SmVy#SUss?=}U=CQmQ-RQI zQj&*SO}LcwdVKFPX=j3=`}y!prr=x@lh&}Y>|HiVG6^QSbI4ASo9lGc&q`-&i>m<` z@5~Ddcu$ zd~B)TxSixo;%49sFq34d9jY+oYHMS``7vp2&k1eR*&f?cBW0t86<5| zlhK+sxm;q&82Gv~ZE_>TibH4GB=&>%@U*5hBp9X|W@KfM83Z;5Dob>hsJJQhdCts9 z!s~qIL=jW{9v2uad40SrL|!GDCw1~d=B}&5#?mGQ2-}7^dV@hw9p#dCHP4wmx&Jzw zJXw+1u7qDKnGa=F(UR6nELMouu)j<5s>XL#x_mC>y)fM)v(!C{5--vIB+605rkNJ$ z1NyZ)R{eCsKGafh)5%+^U2vH%-d-L-ALDyxU-F_Ba0>Fs(s<)!?ls(*BD3?uS1+ZH zExG*H48w*`c?!5G0(>^@yxLyPqFGH9ce3|mYc$(jyd6RCYIuw-+4~`AIj=oQFkH{- zkxA`87e&oW{i1Li8)wXdpsi^4M3m;bk_g&bCJ2AaRBKBouY@8hc^|@(n-3lr+(AKb zl5uhxwO|8v42r06H%aImBZ??<0k4R1Kc|NLE*2`bz7Iw8dm$u!^XKoKYg4!zv;{a2rLmWDRqo zts>gBM)V~XtU0sjM5`jY5KI)&&%d-PqQ4^QCx6+hi0&nkL#HBAFT(0F1?`0*$}Qm( zQD0Y95t)O^TNTlx*EmIFIW@9d1ox`)^tyjxDcRJ6EW{E>tXj17jMYl{9rE0FC#V%p zaTl|)X3pIHHUB)iDgTkv@erpQ7qM4CuI{p;vn0{mk%nEpX?PdtyCeLSF^(4eXxB=PPyxK)gRXW z4Hs$DKuc!!9j*#SZ@uoW%+KO&o5b@~Rrx(#vwONOw12~Na}*6{0B@O`ubPEEw1kjE zkzQwJ0zB#oc+*=8Zc{ixAk$U&R_jeZv=Kxne<(q83?2m0S;j&>1ISqxS29>dE)~|9 z46e;3@Xv>A6B`8?{2&(aPuv_b*z83HA0vacF)UV8K(4%k?YGv844#*W!7gXXU}Z2F zys`tCGI;C=GI%o?-0vkacvfdJjqk|do!iJX#`|Kh zel8i@Fp>;DK?Y0ulfnHW$lzOK@YU^P@R#i16=d+L(PVH78C)CtW2Fe2+pO zU;4N&!NkttGw{0FRA~RNfF`o(@mds>I!GU6SM5E*>aIGjcP43Y0fv`^LnGWB`$f(t z_)$e?4w`zGl1r?0bhE01DePc&+Y@ z+8e)(BIHj-$`Wtu5SPcWggGQu-Sh;ZooHXfu8~$>i_;8 z(kuD+{+kU-X7GbQ36m(|gRfSB`L^dJkH0IF%>FO_M6;5TPuAW9!^iN(nz`6j;MuPm zmtqq4me1V{%FTy67PR2Bho3$N(x!W>+#5)ncGSDGaF>?rp&J2+=_>nf0AzP?*})-w z6K!8CdO|MT^W^I{Hp}dDN52pJs^Ob^*gmf0=bzd#__pH;i-CSeITQ%H%?oP14Up6cR()NM+A~3k$m+!xf zckp|9C+L{n?klPQ$Q27M5Vx?IGE}VeZ=&55m<~e(@HdI>=*< zLmo{W97=Yc1Rz)Afp=q?dEzr*fvV&Cr{H;<2)U9@p2u~`JE$Q$wy~)J&jU25FL1Zt zhW>(EL$+H&v4jMbUJ>bX)NBiB5rCz$~UxkcUBG{b_lZ7dRr; zbtZXarZdTU#sXpehOQ4lp}u2$;VZsOa?sY{jcz>|#R&qZsR^96 z88%#^Cg6xG7zlp&?0c`lh%Qik+@Q1M!Z!9EW>lwF2`b{z6(?#ygPCoEr%o|cj%XIW zfy?wKmq7<#S(9i0b%-8GQqR57b~#8oN2gyGq^b&6!(F|Gs^$rXd8NK!ESVDs&t24pXL;?|X#J{TX*h%o65($I4Vti1EFa_c!$fVW zK#@hICX4BB<%UYHRZ<1=7PjeNX@ypa;FHXnY{+DSqKBKoM9Ti%bt0%^(%@}ceHY6q zuUa`aF3DZ&_6t1bDX035h?k_Rd@l}N7Jm7Dm8<;2(MBu&1G&jCBxi~(s4c=Ve!FD< z}X+@ummJsgkrewHO1E)zFmKP<%?G$wa zmyTHuf*slD#XyH0qVGjFsP7Tv&UTu#NTyd1-1$x@!e^uSpsKdlPHQNghtB>G=RX+8 z(;U0l>ed67q1F>|3M~g%s40>l5I1Xs>FfP6t((Aa3k{(|S-f%qxcdd8gF)_rzZ)N{ zn+@xrO@+2bQ_&w8lQ)-; zSmx!Khk5I3=ABhd@4q+!C0(!K;M_V`SiI;m`PTU4-9 zobZ?OZ=JwmLL7}lKH=Z74Y0T*jsS}_ac9tjy0-cv7VTFC4`S-1F)Lyl?4nHI$xH(= zbv^BP$=Cu7a&tAi2kxGSj~7?|S_4$98K_QP12kA=&!$*oRz{WJ>ye!KwR>T@Er7;i zE_kslVRH;o2OeKai3MJTqSUVniOu^8&9cjw?q$~Q`p-?`FEsbE?YJnh!^WHe^t(Z; zI^gIJ%VeFaRUl~uc&ra&o8Jin-I;RwgT?IaU@;HG{EXh$+4&v9{BE%GyFunh^()i%)Bk1? zU)7Z47``!X#DO^?wo7GD<3~O%L%|&y`PL~wR?Fl6OQm5!*4HDi=K68i;%zf zwjlQTBB%3K>Py8w=eK_!;MRsCy!T}S3*RwZd~6YQ>DcxX*0P)-G( z`Cc`&6(LU;?cfOCG0fBqH87rM+H!;+ceL3&+c`J=dlHXBE-zF!y6vNQ?syd72nz>I z_|qi5ZFjvT@Mt2icTFWW77jk5sv>>u6QS)_*Txx!ghNlJsS{Bk8Iwr1JT}DXVhOj) zr$t&~OoTyeA=sNKelVn|%7^5Q1~D*n7wqQ0p zd6%5g?4x3{lT=C@dpbTiNu~3B(~bqJH^JtnTwtF z^pQV-XnCrBsd6PJMYLv&^+QLcdkxBour~C7s zj=`3`a@bSVD2F}G8CH_(w5NS0v-Y%0X1h>6V|3*Hi}v)GzqpM3UBsR~(90hQyr1el zo}1j!44{K{3!Nda7-pQN@wpw?45YG$-8cGT;V6}V(U-_$r{Ja1DTT9C9>CYKdtN@L) z$l;M4jt4Z7TLd%ukjB9D%W{}eW;EokS+16aWLJqgCt&Y@%&)t#&`1@a5oQ}TV(SbG zfkr4_KLU+NMi`=ykt}~3m#-mM3o??@fTW^vY0Rh%)Bt28Zxeu;89*|U*0H=5uS%hY z{0V^i#M`+2pJblBO9?b$eHp58WCMXls{A1Gxp^ar6=w$V&ojJp&CBEtPAl%eKJsqv zSk{W$$r<)|2Q(sm0UDWV9nv-oGm=I>V#V_>v|8~#opb$?TQia| ziR92}#ci(;EABf48A)R?XTdJ`)=nrBW@KPrL)GuLm&)O1zeno%O7U6N0 zARxpp%P+l5>RbY50HCfV;#+;IVh%!G;30HNP;TBg;tZsxlhd_FtrK1VsGm+Wz8u5l z(f!NZtRp#n=|3mV@G_@?*9XZOE^rp%qgQC>@J|9`rt@Qpf0otrV|M4$eYs24ebR~H zFPS$dMgOuBpf29y1gPuw^grfL0BY+hOm~(pV1X*Y#hkCXBwrkNIo{J&2#x-gXo}CU0)Cg}>yrNfj>-2z)$OLe>|-9>?4!q z&a-dtcAi~#yu0&k(ZUhZ*PH-#@rnM={Ru#g`Nf!Z69G_HPTHFIPk`G0q-fR>eu~%L zmzMBTeB76w$v@7MGV4Tciq#7f4E*r>@5^j;PJsIU!sJP%{MqPFHH-g&HIdXL?(9xk z!;Spen?rm}{P2!H(w$y+0@Uv{d^D8+)ZZp0o2y43&KrG9lzbv#w53`*7i()i1ii*! z4iHW*sd|M4pBLa0J3fh>oVp}~{bV6NX~8GilS2YtWj`_Clfq1MK*i*S-^myjptdNX zg{u;#&=IxoK@H9pLk;c<1McTiIs;l*tys(g)XNeN(YrE}3s^0Dp;{xdo}#9dWeOC~ z!diu2D$&AI^ag0*dt0G}(Z+<;!n&ZOhKXsV85@{4poR4W7oy%ITKMuzUJJ*6NDWB` zE|g1dpH6V0WtYkAFJo=Tr8Pk-=r;vH2^8h05iLx;$7|um5ru&1bU{unTzoNd`tF&m z7A_-aG$*{q#!|YVHnCLoQtkU~webC=)$ikp?#lH0v{eh&1QIQL_NP`Y{4pL+NvjtA zB#9h4wQ%!Xk|$8?fEKRV#%bXzVpa`T zUo%p(bfwl!`I?}>HYYU?rhca^j6oq}%_R4yygXOV;k0o$4r z=Rg<7e?|PiIpv8gKyCf{TokpiVWZ7~8_tD`aIjsuc%JLW1ULLtyZsc>ZRPoa+d&r> zEX1R-1?^)2YHDFfY1Z60!VGaw13z)tL4xJ_V01Hq*4SEm=K%OYOjEckOmPPg6AtG1?|X%&${ zN@?d;vzDSa{DBB5i50h#b& z2Z{gX2+mjP2eb2nlo;DPq{I|_yj4o7)1{Iu@&taoQNsC?eb!lDr1}4pE)~yUuR?sm zJ?Zkjl=p#jk3*&Ihbi$zrTIg^H63=PG*ab^W-8Zz$bq1B1$S`fl_$_TVcAJWhMq zKSr2)i-D-EdIEsh?s0S-UYlLk-n#vp&#Oe&k`B>-W!|{*I-AG40%6;_O4`0-=)F@Y zAjx%|gTE7ux$73(huoP~KCtt;i^!lUaIRS z-q5uDS61#G61lrk+MboWnxr;zw<|;1nATSAnv{Mg@MdaGk9pK8cg-t_++o*Fr`(;( zP#J&yvsLaMD5ds?D_iBxjJkEZO|Y|4?)TX}4qGXN32rRB6! z&FPVM`Jc0S+P_fiC2vu$uSK>J8pFbPi@BWMs;YI)7~}{p1s)@_Dm7^>FTumbFV+Ey z-qtFOdJqk3P7~3yr^|n%ysxKw+$eRwNr`W1KX-Krja_9~YY#rzW$S5ChSG9wBv@8= zq0kD(PAc%7WVtZ0**2*qWp<%xq2w>^8?h;c9~H{4cG*eycaO|m>RW}7Z0AqN7J1(L zKka=9SX0;bZ$iRJm^30tn4_YgBC~=+EeZ1|Qnl795QZoeFoZ#zXc9nzqNOcbtf;BQ zp{-V0=a~Q^L!jtYYmr*YwO(rdQ?zQi*6Le(=j;&d?Y;N^fA@Xgd*Ao?YFEzh?7h!E zYpt{Q+S6LRA-5{#JX48!lhV2OW1iWNo7QNa>BGQPP}-4?=lpK zuZL5_Uct{1x>0^nl@rPiCZTRNs?r%o#m=TfnlhseXaH3d_Xs*;!I0w0WodEJMsyCf z|8f!TGAk#9$&=wIs#wjh5HnadwPH2&nnGI&okp!mmXmW%)ya&ZgJK2*I}pQxRD~Tn zt6uLs0oH+dC;2cl^!o#-RIz3g6B_MW)!V4=6xvfn8i|z19-U0xG(19n+tw|cUy-ak zoazn--rE@y##^FEXW=HpL2^25?xQXlFVG(~`bZMjD+WQ2qe4{+7YR zL|)XF%cXn!_D?^VxfwD7jAPX(s~{hgcGQ;>7frc@-nxfG0S>L*OMToyyku!WO$ z18~xAiY;2b7&ryq5d*tY0>R%-b0mcZ#uL^|fNe6W9qeuIqX15DY~UTFc}?M4*slGk zJZ3bKUfg6KqzvA+pt6k>UE$JKZJ(v`Y@VeC9%ZM=#I}PeST*A&f@&6-`YHDqRUp;e z66fUus_CcHv`C8s-$T;4*#afQm^yB9*idTa>m{jben~2*W=$$mO-(pf&8jrLYF4H( zhSaH`nrd8|NTizbiStlx@@+viXMH&lRP&?lNHwdl&&gmYQWdteQz5B4#{ORfB4-v87eB@ zm@HeOnj7r?pqix_@=IyBf}omR5?VDYOm+;b<~}>DnpLg&+4_4U1Na8gWeVm(<~%}o`3E~;qxymBI1*aHpgH_AYgkxpJ|xojPk zXN%WK)UDu(=4iu?`a?;R!M4wtyj~2HDVy9kt-2M=x1ERwwdEgTg^;_3!5-Mxk4e=0 zV-ipZtpqE?TUa5QR_GO?eg$I)Ue3s-OM=uh3q7_DwxEMTjJE@Y$n2Yp?6^HhA$lWr zVui@kE5vE4Aqy0OxqDb4ezbEtWqWLeP9fe#3NbhtB2YqG#oTSE5N9c7;RuCjeJf4{ zUOy~*TNk`AP>8A(SRw8jD8#-ailp7bH3kZCkt$h3{a~OFf{zh1S|WCVLQJxw6(a6q zokBF$bPX#++LuHjdRt@ut{3AJ`#&he@l5$+T?LU@mrg6h&DLGR3gK;!6{5E_|Eazt zU;CF7qJt7A?FEI9yRu>w5w#`($I7RlE1z0#!q=LFxK`}quHM^ev=gnDo$+l~yZ&s1 z&iZ}7M*$M(@7xWBjzoNb32!$Tv4?%&fMVAOcPrbzqKtM^tL=uTN z5)QASWfiZ_6uZ4XgZN;dDWLN|hMv>-nK}}z`L2vWsLUsOr^39Z@G~a3V<~wn@H61| z;M7J7GY{yp9MFOMkONt1Rjx z8rb2`)CQyTbDJ2`U@GR(RO*oV*fGk(EhUGZH6Btm9eTt*EN34c3O@WS`0%o%Bgp%R ze)H)MXBF}zn)AbcbgYLc_66!!eRQ%53WL@C>q8V;E$eCf=q+`VvA}-Ba!bEG_b0o9 zZxMZXN)?}{?vZl#UC^jM7Y;AFe6flLJ03KLYtB-yEvx);M4hZJD(b#ev<=tk`ZKDj zjhbptk7z366T zI>{%UwwihYWxX=8tVs^sSM7%~hkeHEk_JA@>^(~C2{#AB{XCyYd9W^}ISksTJiO%i zUg}sM=ctj{m+>I(UgWOQe*J_V(8c-_5FJ`Q=d@)_?+FExYqmWc7TkS{(yhlZj-r0P zf720|tjU6gRM?Hwyw~F_YDhEep`T|Rc@&A3WO}oplOK?Uw@`R!bO8?5*PZT~mf>*+ zltcK)bQ>rK^XjxEpTqv4SL|WIzwwAX3j0ABkTrcvttsuKHYqy66$# zsPsEZKWp91py~D29r1r++6#XEbz@lk5u=`0<~F!lLn4bE+S^I&K!!8 zCJW9g$W#uz1v-UayZOxUIQ^&tJx*7g?MXdD_HU=P3;!^iE^0ZudChlHC z-R3TjnV$WoCZqb|T^wNA^DIoum7MK5SLq7k{@%kq;7+TT*#gf&OS{t(?N0mJo&IIq zY2The&riUUrF+o1`e^VirMoCTR~wB_5i0(UY)&uFUoC4Y8&E6w0*SP-wJQqDQn^$P za(bNXi4GF>HZVIZYqV2N+G9haryZ7Rep8;mc|qU_hJ@I&>KxQYbzb_=r$?Y&l>`}w z-t*L$fHSHZ?U?BfcdQ5P-Kar(*mLm&``-~y^Pq34iZ zMEz60Sqt&lyU64CV|c{YB{2AT%7C&v3MwBRA7#Ao7e9FxM z`{gm^3=hvA=Lu@1@`hp8Jdgjve7AvV${8M^U&}$v7$4-LDQDVo=q1A59G!z6fzQN= z`U5We@_F2*kA&cX%QbiGXbzeCdz15r$sws@QWCD9FF1hs*g>w_T0SYoN58e)dN~I# z6uRav*O#-fL_fx>M-`y;z=4?W|g@$pxMbIkHB9Ls&G_mO_VTUK(BtB zc4SK&s)*O5ZP^@ATr;WiSx0lBEvfQ@fYM(%AZL7#w#%sfAQ!#t%-MD3v(#GpJqH}Q84k7N(bv~;f-3co9(3dh?Z~4q zt&U@-lSfbbPBEr7M;b=j=_{+{2XlrqK6Z2&&iKu0^>D`g?=+VBr&$+qqAS1B*fPn| z(szp7t3qiku5piClm&i#-Gj7sF@IPPQgN6LoK>VY)_n@zIr0kIRG?k@SD5I;)7X<& zc(*W)t$kd2STUOc{&DlbUXW@uc-99N}_~ z;oR4*wQ)W9&^{jOSa8g7c!HrHDhCq`>l}wC7#`nK)P1c;057g&4(<@ADS@-rad?73 zdK)pL%Ay;_lV=?1@ud2;Zai81P0bel1jF&4$awNlN6hDrq*b}#^oQ|e`8s)P0j?m7 zC$Fuh$CI^Nx`)S;cb)Kf@=!Po=;&(cw#q-YR1ybb`@YY#sQ;zPZWR*9ms%_V;WaEWm0c@P~P-p;0-k3qiL) zje>jPvXC0(q=aVq_6T}cqdqh@$R8UNs?IS!PIs&yJHlRFM=&UUb3_jPq~BqKVwRKM zpa^hc%zNet2E{W+FevK3Lk7jQj@w{R1b+w`Qv7W%)GKKSd-mI#>yHeI!S9Gc@%W)) z@H^m$V7A_%Sfu5csyH*8utD+KU8H_@w%o=Bg>g!z4+aI=9y(Z2>UWBo7wnsVVzVo*F2OnadirI>l_w$EIsp)P|tHTd-z}`s*I{&RqM-&5qKw zmD7gf-;@{eRI7QnNfY|9Ggnwhn$UNPyx8B9fH;euGRKsxrn|yzKXO8bGaYbtfJ`9Z8mI!_tG#9yoj`qyf2@hcCnbev zp^1~v`M{&VNjGZ-5v9a`Y5cM4k=JZAYc~6!H&CW*_L&*8!LRxu4lF)Z&i-^P+DFbD zaL^106CT`CBoE$WnlFp{NrBd%zHq`0uzqGvxT^CvIsRv6Z^}RIc9l{Ujvm6Fwrvk-gRCS zahq=af@5@s^EanqL+80OuivR=9rWe2lQW|Z#0Dpsw#3cx|F3n#7&xr{!-~{nrO}h0aRt$?xF_Hw7P+NNnsaiZ{!Dj=8ZU;* z=Y0J@U($8;;K+?}0d6JCHu&(F18_RhPdt0_yR&<{W}g_G{V45*tmO)~^A!Q7&7+T-kG@oKFSQ`mhq<9j%B^ybDOIN?|9fKMf2U#@4+NlJ}7Iu3ZS6deaVfeuIZjzh-*?}$+1 zl_+sFJ{)-*AIhAI68Dy&#LrRUx(O(;EEp08Ring}W6?pwya`b9_w!NWo=uRLNNqnKhLT+GMpZ45-|r&5?X-k4FF zvG>tj?1^+J1!r-=S-fpgmL=^dOJf4gvcwk);zNAAwFhTuT7iW%4`*?)Ls?p`pe#Y} z;4JliC<|{t&f;T%Dt0X$XIY4|9K>0wuj8kj!C87|;HMnH72A%pbY-F}A7$b!E3n(& z%3CN)@jaX+!5?MmI*YSB##!oC;iqiCS^Tkf1^kGg62n2S-!&U$QQyE7JA;qy9xFgu zgeP#8UPtsUs-JO|^*GD@IVej~AI@Sw8okS%LX<_J#aY%np)B0TILl`^OWQn@rT!_d z*lnEU*xM+}_rK#TPn}VgN+I_6YJt7Kw#A_=ab-%RMRTmsQ>ypkEaULkaoP=K5rm;E ztXFXsQ9Q~r_#V!36lb})4`=biSw5M9vZ!NlmUTEw-)lI_9>{_Y{=uJEz+fx^I+pz! z_M2=Czv&2Z;Ahskf`I}BztF`D?wW;9#5Mr%Ta?a4`XKy%?h}k$tBUAo?=3HM$6=py zQrK|!*H^7#8G_@&I+14fJG7>Dy?X;k**c;Gbxy8 z4sWyw=gT#Y5V}hcA+>C_QK1!YcSTi(H2G6aIhv2-RfBs^!4xp5lc*;5;$!St><={M zaEJg7UwOr&X(ET<-kOp-mL2X22sBcQtk}76a-~yXApHb3%s{voTNy7ohU2M1vN)TA z5aFHc$v(#z6lu&#JX}txdqqLsEfa%JFc!1}vk{-)FyQmJ{|KKoz~JA3&*&&8@VNx2 z`u_=DLkK-RuTPcajo|eHSdpkl)@FE|ywOgsXUBu`_S?q@LRl)Ux0m~lx8j;@$nfp77K6he+_F8d*)C20I;_FR_%wOV5| zdbC}@c9)p@ypM!`9W~BwA$Ri1dzBwCk3-A{u>jg6EYBD$&$0g#dBOo{i0MGbN>hnuULvU$q^&Wg$`;b{t3>is z5cyRg`PqW}ys-RI*D6jTZq%Fm3!GqY$bEO#du&(ndd&fh-m$HzULoy6W-L1eH9=@W zt^z`{z%t(`O^2lyD~2fv;7(;9WT-_7(~`Psk;t?SH{Xaa@t7~M|5DF&9tK)oFn8hB zLIBaME?(4PzXe>-=bt@=$o=*N5)U4=CXnPi2rb7$P7qM2kUwd|WB#NK_cLt6dw9|$ zX;2-CLbI&)NAI1xh5gb-L5ap#8S8HqEMwY2>mBP~A`H}8`t6`xC3ZE1A_j=W7SuEb z_{mKL*7${121#ve0(#=upxjr`!3(XB-5TZRs(@@D5?kbyw*kfx4QN_F=qNj&R@Z6> zi;bOi4T&EN`1$62JQYiBRx;h{ggmb73$FQex2i9gSkIPXe#ZZ^zR&{ob^^8`Bby8u z8H5>0`2Fwqg$w?KkHLm~Y!9USLKJ`p)|hRmlfiR)c$pH++z>06OIR5wZ8YAM;6?33 zpY4^1oA48~x(c)F%e9jhA%M9QW%QS0`jrWdZCyhKj`;rSf#S zP#t@;v$biw@g!u*cF{5FN(|G@Tm$*$m{WQ-kf7!ST#p%LZp0&nM{`)0#= zwL$l>D}d4bovsiqaS-!*pv)V2qr@bYetmc(qB~%M^Z(wfH^W{Ab_c zzt9_BBkZT1_635hsQ+Gnv-D5;o0TK_8>~Oj-{|zGf!3cqo&J=LBl@!k0!Co{>6Wry zgLT=Sx@wrX1FdO%m;HLyIviR?4fwac)T$({kKjXTqFul7huGHIo-0l6r1c89Qdgt? zs2ABqna48Q|>OpS>4#XEx*1gQh&( ziNV_lQ79R88gtx!xO_WP5%n1c8~H%@{e$ZeT@WN37ziA3K$9wPqkAL< z?ti`kuhC#2gIh(IVwngl$8*|wTuls|LWUSUVdywnJI_3q7pky~iiA0i_C;3keq$Es zcH{N%qP565{TsY0VbW?gxEM%X8OMV0c?ry~WCmM%eaqt`8+y!9^ET9)?gG7bfj)~( zHRKh4ZEWKO`}tBuUB;_m0$Kq(^`ymVa0kVM_4?Ke9F=Qe?xig-B9LQT!Zoqu>nB%y z#?%WTK!!W`L~v$*9?0D^jeEx{6WD=yf2i@QQHwxgB3?w#E6J4NAlLUBE4x#Q3k`uEZq-nFFOE0QZG3H0r3hN;Hj8%r`N+(IZGvH#;sHR+a*X&djdL_xLa+9wxWZ(TcU@^91|k}%3q z8|x~7`U;@F@|Wt{V$p``%apYRfzrSxPODH9d&(VuXs`qgp_QPnwctg@0U7R0h!&M( zG^3uWwrUt}3*6N4a|+P0Y#l#Qyt047&#++d1GNThwU1mKpiU0|kcX{09+n5$%mp60 z5A)CqT9hOKX!al)yvZPpt*#kqvqlFd(K$!A2zeO?o49IoE92b|EP#+0nbo9^sx=e2ZZC)3ts33 zh5|Q;m@*NPfjwoSgx3TIB5K44utHHK%r!I2J1pHAEjnS3%iJib2qN%S!s1hKEc`rb z9xOEXkBf(W6)_Da%=!LtuR)7odDE7zmoh$01iPYd9G#X^N< zN>>6Np(!-BrMp2{SrfsydJml*V~uf0ePFy3Yq~FlPiu>@P3-;AIJpiR?_g1J^IpFU zvkVDIy8cq@9<(?5WGOQ-iGs7|N7>rrO#3e5ECE=ju-gSi3-&k95V#2M=QY|v;CTf@ zH#D#^KEqIJ7?g2Ac)16=!DtHlKoBqHVP5F$d*lY~U&|WYK@E*YUcFv;8O5<63&?#l z$lVX*UJs23xILid{>rEhWH!ekVaH&(Tf)dMI!-V5B@2Cfk``gPFInXEZ5+s5=!@mv zxE#ZfGUkV5=dB8l7!&iII!!>anchG0l8v5TeH3fd*j@Bya{@ z@Qao9fD1uEu8wucx2CL6itCz4=@BNi|iD9+zlb}0Vl;8XhZ_lJ9p7G*Gs z;v375-iCl{aJMjDP<5#PXn(&5Sod&&nC_$e=WfXr!Q4|ZBnU>4=w@Mb;d{9XKv~c| zE_@$eJEqeF7Alfn&($>gv*&zUDZY{yBsV$A20tAq!5`G^abOM6+X5zzr{LtoUo7B+ z65J0AtLHF$xlOwILx@tSzpl;rZ`D8a|4{wmo7LCfy?6nM+Oees;{8Gxcj)k7_oo8Z zw4m8hI5R|I(-L3k(ZvE{^XT3MeV6}C7u}Sn_>cr-(gHGZps-BPJGqDZ_SCsxnFPR4 zchC)UKX9vFCgJlTi1I5KUGlzOCQGV(d!!4nOj4K7GEt|2Oe#Spa3bn5$i&$OS*Y+% z{#1>FEjvfVu&EaIz>gI?N(QmOOqk2mRL)zxCio`9NgNXEXqRAmacxnw8~M2{+AWku zuSdIwN<$s$jl+5z*g3_#1MHk2hK++X42A;Dff4Fa^0*e3f&;<2!~zCGXY>4DEbm5( z$G{#N7L(Ip4E+FmoGt=z zh?iBHlnEv`GYr{y;Qu6$wE%AK@M)KNLHAiX<*2!dWz9@DtpHZS$~F{n9=w4^)jRO{ z-Z1M`#V(L{DMXyjj;KXZ@=d=8Y&KTFrUDa%*1OQiUhBVmb1h0+X{Y@%N z1`9+k=kW}KuH6#DJt7+?`2D{x8#*SPE^Zke19CJKyp`&THwxZ5NiI-rO=c*2o!(*U ztUzR;F#be7@w9xrLZEN9*2|}O5|K|fmXA~~pYb3ch+@jNF_4cCFZu{Av(~fP!dQ z4jUqlznUthVR0Zm{714l3vwa$z5vqv16iOkGO=CxJZ~df!D(hYkC_30X8p!{&&Rw@ z`)cjGh(8%Fhy}%g1Mi_Q+TIIb*N3^W;OSBq{dY{&C(#>q!Qqf^&;ae`U;gZyu`bow z$w>>~n+d@(0Xtb5w9(`*7S-TV1gzzq;~)umRjJQ{iH&$L7ag0?fD`<&)eEnn(AUmZ zt=>Jn)Fw3D%d!^cs6=q2*NWq0<(jEf+SNluEd~XH_8fFm5depL7YN?VY2b5BlNd%Y zgNdxAa7&nv`A{f;(j)>HQh8BY6N$jJlDQ&~ZO7zB`}2cyG`98lYqC%{jA^k})W9)* zvZv?L*fM)ErhVM>Jw^7K%*#0?B%RxQhivfGjze4+k1Mk$=0NGC;pEAXymCPjy}j?< zkUEJ34I|6cS{rQT}l`*PJ9@W_@p%P>7~Ty#;Kyv)XGR9bv!+Tx9AOPkXY zC`odtM6ywm(JaZPR>(qE6l`3vws}PnmDv!Q*}O5ctvR!k%DNhwbz@^zcXQTlD!VT< z`|ifa~|w7jf1%gspEbTi$+|u~)(#5LMr7fii+>OcM8zriZ87&*L zxn;8OvI14v+Lp2+Zbd_QMYF1+t)-%qyX9*5mK&-q-7Q;gbGP<|Z@sJ9daq^c0Qa4T z;qN?Az4Nr?o#)&tQAAZ}c~xX-KDu4`dM^ z6qJ9kw)KM|lbVKzn&$GFw$_?XlU-LMcHJo7)!n-5w#n|kh~0O~ci(H>Jz%ouVZ@#% z<$IpC?s;xfD~hZQ-BcUdRvT-&FFtbL;!XROw(UzW-Jcw}U$SX`M%(^u(*v@|0|lE7 ztZh3`WO}F}@=)`pLv3w`I!zB>jXZo~)8X#6!?#U8?~DBW?xxT0wS7KddgNi`ktdsu zJZ(Gj-1L|z>R9OJW0CF0V$F`nM;%|h`S{ZI;|XRblA}&YHlN67Kap*AQWkZxVDrhf z?I(-O&NW1xYuEf-HeY<& ze(||klPJ0=w4y0e+Z1blDL(qr;)+X4wU-jio0FrPB^AvX+U9KY%d+Up1r?XqYA+X= zYa61q%@x`Pf{{Pqkk?H@_x| zxfZ(RT4cwySc~iNG1nJwxxTdHdVn<_1C3HrT8DM$ZJz#TJMIp#Nko0zX0`|kK~ zjjhsp|M-;}JEhH+<5gGemA0409~3$%?YB>8QA|}j9-q+vV4BkTrwQEmUQv$UH_2^p zrgCEEb%f-E0nQqO$^W`b~Dw*N?Eo|j!s&| zQ+2}GR-S(JJ?24)r>*?06Yn!yB=c-ithGh5 zL4Ro=ry*F~v}9Es^Mq^EhrI(6k!Uh18va4|h4;rGd_DYI1OKj}iC_5J2IB%43XM_F_CBN(7_hJOX-$KZb=6DE2b2to`xD^AA_KOgTaIQOpu^MnI zzzHSzX1W|g5su2K1{~by>nXwa(>S37KL$8j+aQ$SjWkXu!LI^t3F!$X_#ME}9t=VW z{u^Bmp#(Ptd7$SKN^n~mCzRlBG)^eN{b`&~f``&Lp#)z9IN}eX1W%!HLJ6J&IFc8k z1TUg-LJ7W^#t9{OHH{NW@O?B+D8av^aY6}x5pYz0LJ8ghI9h)pl;Gb2j`&O{!SB=Q z2_^V5x_m+jZVDd zp#Mwe+W2wE};bfgvJRa_$eAEl;D?XoKS+_031D+P=ep3aY70Hgf53rf*Tv_ z>PaZUt!SK3f={GzLJ95zII0h!1Q*dbp#+boaY6~6Oyh(STt?%B5?o2+gcAHc8Yh(C zdug0df`17(k~^UUZ>DiV3EoYYLny)T(Kw+5e+oDs$|scIme469oKS*K1sv%!p#+}; zIO?hiC3q-}6H4#}fWvU7cxdKio!qEqf>}G@$N^pSmb~3S- zd7zuX(pYb=mEufEq9k*bIB#ZZ zTAnyP(|ZPjd@z)lo+Zr;$&kF3=k1@KotL&kCeBPuS(TleC(h1`NRCeS4haeI_RE5+ z^xRx=VxBlTGwnZDZd##qRd!l--ha4UpH&h`ZW<|7_dXtpi7T@ti79I`6O+?&bJH>t z)4dWUQtz1_Qn5^&<(>F%CG(+6$y}S5?(378S&-$CnJ!7oOV3L40G`S6JfvPD(X{MT zzznj@^wV8PS4mS)A%<7vUhAKp$P1VHyvee(g5jIGE0Qx;oH~39UFdUxREF>13zA2i zY;Xr(7~UCvb$HQLG6Pm*u3F(^a0&Y4jigFm5v~oc5OaV>I%bVXq$>jpP6B)or>1&H zNEx{~GDB9&Ub>c~re_&EL$(q!55J`DQqO>qH}RGJUamMVUnb5=gO62T|CF@!OoL15 zs?`Qp@QJQ8xE$6B@0mKak%)8i5>wL)(o@q6(&fqWvs3=THOU1T5@l-*5=gREr5awM zQb#g2J54N0UNchnQdUXVpeh;U1YwPkjyN~n@HTx(>LE^91+vn!g1*$>g||19yqLZw z_o3*K%)yuBE);Eem%b!-QAG@IlWYC`m$IX(q4fIuFG&W~&G0@j589jIS0mRXe|B17 zo>-EXW~k_f7)iHsjdYcvYLaVm|K+xhX!OfCNuMT@q-Pi4ioKO4%Z4Tr7$Kp5!aGWk zyrk(_;=Gla>B)&8PD$npaaNXCk|;||Tv1qvnylPa+2YLfyfujhbfc$BV(_}~fhEfF zUj__S45u(uK@5!075K8H= z5K3WK8Jxnf4E#;~0Wi#hQy4Nuk|lomGX&rdDGUq6SquwBz_6Mj=?!0jJPocumiP)c zfzY==^A4g+^B|p^AvDtur-D{me;fa%@!TwNq5cMG1yFP9J;TGxW5&w#bZP!fzv(N} mbJH`^!PJ + +#include +#include +#include +#include +#include +#include +#include + +#include "runtime/runtime/rt.h" + +// Use map to differentiate same name functions from different binary +static std::unordered_map registered_names; +static std::unordered_map> func_stubs; + +static std::tuple registerKernel(const char *name, + const void *data, + size_t data_size, int device, + const char *kernel_mode_str) { + rtError_t rtRet; + + rtDevBinary_t devbin; + devbin.data = data; + devbin.length = data_size; + const std::string kernel_mode{kernel_mode_str}; + if (kernel_mode == "aiv") + devbin.magic = RT_DEV_BINARY_MAGIC_ELF_AIVEC; + else + devbin.magic = RT_DEV_BINARY_MAGIC_ELF; + devbin.version = 0; + + rtRet = rtSetDevice(device); + if (rtRet != RT_ERROR_NONE) { + printf("rtSetDevice failed, 0x%x\n", rtRet); + return {nullptr, nullptr}; + } + + void *devbinHandle = nullptr; + rtRet = rtDevBinaryRegister(&devbin, &devbinHandle); + if (rtRet != RT_ERROR_NONE) { + printf("rtDevBinaryRegister failed, 0x%x\n", rtRet); + return {nullptr, nullptr}; + } + + std::string stubName = name; + stubName += "_" + std::to_string(registered_names[name]); + registered_names[name]++; + auto registered = func_stubs.emplace(stubName, std::make_unique(0)); + void *func_stub_handle = registered.first->second.get(); + rtRet = rtFunctionRegister(devbinHandle, func_stub_handle, stubName.c_str(), + (void *)name, 0); + if (rtRet != RT_ERROR_NONE) { + printf("rtFunctionRegister failed(stubName = %s), 0x%x\n", stubName.c_str(), + rtRet); + return {nullptr, nullptr}; + } + + return std::make_tuple(devbinHandle, func_stub_handle); +} + +static PyObject *loadKernelBinary(PyObject *self, PyObject *args) { + const char *name; // kernel name + const char *data; // binary pointer + Py_ssize_t data_size; // binary size + int shared; // shared_memory(meaningless now) + int device; // device ID + const char *kernel_mode; // kernel mode + int32_t n_max_threads = 0; + + if (!PyArg_ParseTuple(args, "ss#iis", &name, &data, &data_size, &shared, + &device, &kernel_mode)) { + return nullptr; + } + + auto [module_handle, func_handle] = + registerKernel(name, data, data_size, device, kernel_mode); + + uint64_t mod = reinterpret_cast(module_handle); + uint64_t func = reinterpret_cast(func_handle); + if (PyErr_Occurred()) { + return nullptr; + } + + return Py_BuildValue("(KKiii)", mod, func, 0, 0, n_max_threads); +} + +static PyObject *getArch(PyObject *self, PyObject *args) { + char name[64] = {'\0'}; + + rtError_t rtRet = rtGetSocVersion(name, 64); + + if (rtRet != RT_ERROR_NONE) { + printf("rtGetSocVersion failed, 0x%x", rtRet); + return nullptr; + } + if (PyErr_Occurred()) { + return nullptr; + } + return Py_BuildValue("s", name); +} + +static PyObject *getAiCoreNum(PyObject *self, PyObject *args) { + uint32_t aiCoreCnt; + + rtError_t rtRet = rtGetAiCoreCount(&aiCoreCnt); + + if (rtRet != RT_ERROR_NONE) { + printf("rtGetAiCoreCount failed, 0x%x", rtRet); + return nullptr; + } + if (PyErr_Occurred()) { + return nullptr; + } + return Py_BuildValue("I", aiCoreCnt); +} + +static PyObject *createStream(PyObject *self, PyObject *args) { + rtStream_t stream; + + rtError_t rtRet = rtStreamCreate(&stream, 0); + + if (rtRet != RT_ERROR_NONE) { + printf("rtStreamCreate failed, 0x%x", rtRet); + return nullptr; + } + if (PyErr_Occurred()) { + return nullptr; + } + uint64_t stream_uint64 = reinterpret_cast(stream); + PyObject *result = Py_BuildValue("K", stream_uint64); + + if (result == nullptr) { + rtStreamDestroy(stream); + } + + return result; +} + +/** + * Read binary data from a file into a vector. + * + * @param filename Path to the binary file + * @return Vector of floats read from the file + * @throws std::runtime_error if file cannot be opened or read + */ +std::vector readDataFromBinaryFile(const std::string &filename) { + std::ifstream file(filename, std::ios::binary); + if (!file.is_open()) { + throw std::runtime_error("Failed to open file: " + filename); + } + + file.seekg(0, std::ios::end); + const size_t fileSize = file.tellg(); + file.seekg(0, std::ios::beg); + + // const size_t count = fileSize / sizeof(float); + // if (fileSize % num_bytes_in_elem != 0) { + // throw std::runtime_error("File size is not a multiple of float size"); + // } + + // Read the data into a vector + std::vector data(fileSize); + file.read(data.data(), fileSize); + + // Check if the read was successful + if (!file) { + throw std::runtime_error("Failed to read entire file"); + } + + return data; +} + +static PyObject *readDataFromBinaryFileWrapper(PyObject *self, PyObject *args) { + const char *filename; + uint64_t arr_ptr; + if (!PyArg_ParseTuple(args, "sK", &filename, &arr_ptr)) { + return nullptr; + } + + try { + std::vector data = readDataFromBinaryFile(filename); + char *arr = reinterpret_cast(arr_ptr); + std::copy(data.begin(), data.end(), arr); + return Py_None; + } catch (const std::exception &e) { + PyErr_SetString(PyExc_RuntimeError, e.what()); + return nullptr; + } +} + +void writeDataToBinaryFile(const std::string &filename, const char *data, + size_t num_bytes) { + std::ofstream file(filename, std::ios::binary); + if (!file.is_open()) { + throw std::runtime_error("Failed to open file: " + filename); + } + + file.write(data, num_bytes); + + if (!file) { + throw std::runtime_error("Failed to write to file"); + } +} + +static PyObject *writeDataToBinaryFileWrapper(PyObject *self, PyObject *args) { + const char *filename; + uint64_t arr_ptr; + size_t num_bytes; + + if (!PyArg_ParseTuple(args, "sKn", &filename, &arr_ptr, &num_bytes)) { + return nullptr; + } + + try { + const char *data = reinterpret_cast(arr_ptr); + writeDataToBinaryFile(filename, data, num_bytes); + return Py_None; + } catch (const std::exception &e) { + PyErr_SetString(PyExc_RuntimeError, e.what()); + return nullptr; + } +} + +static PyObject *allocateHostMemory(PyObject *self, PyObject *args) { + uint64_t num_bytes; + if (!PyArg_ParseTuple(args, "K", &num_bytes)) { + return nullptr; + } + + void *host_ptr = nullptr; + rtError_t error = rtMallocHost(&host_ptr, num_bytes, RT_MEMORY_HOST); + if (error != RT_ERROR_NONE) { + PyErr_Format(PyExc_RuntimeError, + "rtMallocHost failed with error code: 0x%x", error); + return nullptr; + } + + PyObject *result = Py_BuildValue("K", (uint64_t)host_ptr); + + if (result == nullptr) { + rtFreeHost(host_ptr); + } + + return result; +} + +static PyObject *allocateDeviceMemory(PyObject *self, PyObject *args) { + uint64_t num_bytes; + if (!PyArg_ParseTuple(args, "K", &num_bytes)) { + return nullptr; + } + + void *device_ptr = nullptr; + rtError_t error = rtMalloc(&device_ptr, num_bytes, RT_MEMORY_HBM, 0); + if (error != RT_ERROR_NONE) { + PyErr_Format(PyExc_RuntimeError, "rtMalloc failed with error code: 0x%x", + error); + return nullptr; + } + + PyObject *result = Py_BuildValue("K", (uint64_t)device_ptr); + + if (result == nullptr) { + rtFree(device_ptr); + } + + return result; +} + +static PyObject *copyMemory(PyObject *self, PyObject *args) { + uint64_t dst_ptr; + uint64_t src_ptr; + size_t count; + const char *direction_str; + rtMemcpyKind_t copy_direction; + + if (!PyArg_ParseTuple(args, "KKns", &dst_ptr, &src_ptr, &count, + &direction_str)) { + return nullptr; + } + + if (strcmp(direction_str, "H2D") == 0) { + copy_direction = RT_MEMCPY_HOST_TO_DEVICE; + } else if (strcmp(direction_str, "D2H") == 0) { + copy_direction = RT_MEMCPY_DEVICE_TO_HOST; + } else { + PyErr_SetString(PyExc_ValueError, + "Invalid copy direction. Must be 'H2D' or 'D2H'."); + return nullptr; + } + + void *dst = (void *)dst_ptr; + void *src = (void *)src_ptr; + + rtError_t error = rtMemcpy(dst, count, src, count, copy_direction); + if (error != RT_ERROR_NONE) { + PyErr_Format(PyExc_RuntimeError, "rtMemcpy failed with error code: 0x%x", + error); + return nullptr; + } + + Py_INCREF(Py_None); + return Py_None; +} + +static PyMethodDef NpuUtilsMethods[] = { + {"load_kernel_binary", loadKernelBinary, METH_VARARGS, + "Load NPU kernel binary into NPU driver"}, + {"get_arch", getArch, METH_VARARGS, "Get soc version of NPU"}, + // sentinel + {"get_aicore_num", getAiCoreNum, METH_VARARGS, "Get the number of AI core"}, + {"create_stream", createStream, METH_VARARGS, "Create a stream"}, + {"read_data_from_file", readDataFromBinaryFileWrapper, METH_VARARGS, + "Read binary file into the array already allocated"}, + {"write_data_to_file", writeDataToBinaryFileWrapper, METH_VARARGS, + "Write an array to a binary file"}, + {"allocate_device_memory", allocateDeviceMemory, METH_VARARGS, + "Allocate device memory"}, + {"allocate_host_memory", allocateHostMemory, METH_VARARGS, + "Allocate host memory"}, + {"copy_memory", copyMemory, METH_VARARGS, + "Copy data between host and device"}, + {nullptr, nullptr, 0, nullptr}}; + +static PyModuleDef ModuleDef = { + PyModuleDef_HEAD_INIT, "npu_utils", + "Utilities for fetching NPU device info and preparing kernel binary", -1, + NpuUtilsMethods}; + +PyMODINIT_FUNC PyInit_npu_utils(void) { + PyObject *m = PyModule_Create(&ModuleDef); + if (m == nullptr) { + return nullptr; + } + + PyModule_AddFunctions(m, NpuUtilsMethods); + return m; +} diff --git a/third_party/ascend/backend/runtime/__init__.py b/third_party/ascend/backend/runtime/__init__.py new file mode 100644 index 0000000000..ec8b005dcf --- /dev/null +++ b/third_party/ascend/backend/runtime/__init__.py @@ -0,0 +1,37 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +from .autotuner import get_max_configs, max_autotune + + +def _patch_autotune(): + try: + import triton + if not triton.knobs.autotuning.use_backend_autotuning: + return + except ImportError: + return + + from .autotuner import autotune + + triton.autotune = autotune + + +_patch_autotune() diff --git a/third_party/ascend/backend/runtime/autoparser.py b/third_party/ascend/backend/runtime/autoparser.py new file mode 100644 index 0000000000..c1925b138e --- /dev/null +++ b/third_party/ascend/backend/runtime/autoparser.py @@ -0,0 +1,939 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import ast +from typing import Dict, List, Union + + +class AutoParser(ast.NodeVisitor): + """ + Base class for parsing triton dsl kernel code using AST analysis. + + Provides common functionality for traversing the AST (abstract syntax tree) + of a triton dsl kernel function and identifying specific elements in the code. + + Subclassed should implement specific parsing logic by overriding the relevant + node visit methods. + """ + + def __init__(self, func_ast: ast.AST): + self.func_ast = func_ast + + def parse(self): + self.visit(self.func_ast) + + def contains_target_var(self, node, var): + """ + Recursively checks if a given AST node or its children contain a reference + to the specified variable. + + :param node: the AST node to check + :type node: ast.AST + :param var: the variable name to search for + :type var: str + :return: True if the variable is found, False otherwise + """ + if isinstance(node, ast.Name) and node.id == var: + return True + for _, value in ast.iter_fields(node): + if isinstance(value, list): + for item in value: + if self.contains_target_var(item, var): + return True + elif isinstance(value, ast.AST): + if self.contains_target_var(value, var): + return True + return False + + +class AxesKeyParser(AutoParser): + """ + A parser for extracting axis information from a given function's AST. + This class is designed to handle specific patterns in the function's code to + determine the axis associated with a given variable. It is particularly useful + for parsing triton DSL kernel code and identifying axis information. + It recursively processes assignment nodes and lessthan nodes to obtain the axes + corresponding to the specified var in the given function. + """ + + def __init__(self, func_ast: ast.AST, keys: Dict[str, str]): + super().__init__(func_ast) + self.keys = keys + self.checked_vars = list() + + def get_axis(self, var: str, node=None): + """ + Traverse the AST using the provided variable name and mask-based less-than + operations to obtain the corresponding axis name. + + :param var: the variable name to get the corresponding axis. + :type var: str + """ + if var in self.checked_vars: + return None + axis = None + if not node: + node = self.func_ast + for child_node in ast.walk(node): + # handle compare node + if isinstance(child_node, ast.Compare): + axis = self.handle_lt_node(var, child_node) + elif isinstance(child_node, ast.Assign): + axis = self.handle_assign_node(var, child_node) + + elif isinstance(child_node, ast.BinOp) and \ + isinstance(child_node.op, ast.BitAnd): + + axis = self.handle_lt_node(var, child_node.left) + if axis is None: + axis = self.handle_lt_node(var, child_node.right) + + if axis is not None: + return axis + self.checked_vars.append(var) + return None + + def handle_assign_node(self, var, node): + if not isinstance(node, ast.Assign) or not isinstance(node.targets, list): + return None + if len(node.targets) != 1 or not isinstance(node.targets[0], ast.Name): + return None + + target = node.targets[0].id + if target in self.checked_vars: + return None + # Prevent cyclic assignment. + if var == target or not self.contains_target_var(node.value, var): + return None + + axis = self.get_axis(var, node.value) + if axis: + return axis + + axis = self.get_axis(target) + return axis + + def handle_lt_node(self, var, node): + if not isinstance(node, ast.Compare) or not isinstance(node.ops, list): + return None + if len(node.ops) != 1 or not isinstance(node.ops[0], ast.Lt): + return None + if not isinstance(node.comparators, list) or len(node.comparators) != 1: + return None + if not isinstance(node.left, ast.Name) or var != node.left.id: + return None + + comparator = node.comparators[0] + if not isinstance(comparator, ast.Name) and \ + not (isinstance(comparator, ast.Call) and \ + isinstance(comparator.func, ast.Name) and \ + comparator.func.id == 'min'): + return None + + for k, v in self.keys.items(): + if self.contains_target_var(comparator, v): + return k + return None + + +class SplitAxesParser(AxesKeyParser): + """ + Extracts the split axis parameters from triton kernel code. The parsing is based on the + `tl.program_id` statement. This class identifies potential split axes by analyzing the usage + of the `tl.program_id` variable in the program and its multiplication operations with other + variables(currently supporting scenarios where multiplication is either direct or indirect via + intermediate variables). It then filters these candidates based on a list of candidate parameters + (parameters not provided by the user). After that, it confirms the split axis corresponding to + the current parameter using mask comparison and the `keys` passed in `autotune`. + + Note: + 1. Split axis parameters must be multiplied with `tl.program_id`. + 2. Without mask comparision, it is impossible to confirm the exact split axis, which would lead + to parameter parsing failure. (eg. mask = offsets < n_elements) + 3. The identified split axes are limited to the list of candidated parameters, ensuring that + only those parameters that can be dynamically adjusted through the autotune process are considered. + """ + + def __init__(self, func_ast: ast.AST, keys: Dict[str, str], candidates_params: List[str]): + """ + :param func_ast: Abstract syntax tree of the triton kernel function + :type func_ast: ast.AST + :param keys: a dict of axis name: argument name, used to confirm the split axis corresponding to + the split axis parameters. + :type keys: Dict[str, str] + :param candidates_params: a list of parameters names that were not provided by the user when calling + triton kernel function. The parser will only consider these parameters as potential split axis + parameters. + :type candidates_params: List[str] + """ + super().__init__(func_ast, keys) + self.split_axes = dict() + self.program_id_vars = list() + self.program_id_var_dims = dict() + self.num_programs_var_dims = dict() + self.grid_stride_tiling_only = dict() + # axis_name -> program_id axis dim + self.split_axis_pid_dims = dict() + # axis_name -> program_id axis dim (includes axes inferred without split params) + self.axis_pid_dims = dict() + self.candidates_params = candidates_params + + def parse(self) -> Dict[str, str]: + super().parse() + return self.split_axes + + def visit_Assign(self, node): + pid_dim = self._get_program_id_dim(node.value) + if pid_dim is not None: + if (len(node.targets) == 1 and isinstance(node.targets[0], ast.Name) + and node.targets[0].id not in self.program_id_vars): + self.program_id_vars.append(node.targets[0].id) + self.program_id_var_dims[node.targets[0].id] = pid_dim + num_programs_dim = self._get_num_programs_dim(node.value) + if num_programs_dim is not None: + if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): + self.num_programs_var_dims[node.targets[0].id] = num_programs_dim + self.generic_visit(node) + + def visit_BinOp(self, node): + if isinstance(node.op, ast.Mult): + split_axes_val = None + split_axis_pid_dim = None + if isinstance(node.left, ast.Name) and node.left.id in self.program_id_vars: + if isinstance(node.right, ast.Name): + split_axes_val = node.right.id + split_axis_pid_dim = self.program_id_var_dims.get(node.left.id) + elif isinstance(node.left, ast.Call) and isinstance(node.left.func, ast.Attribute): + if (isinstance(node.left.func.value, ast.Name) and node.left.func.value.id == "tl" + and node.left.func.attr == "program_id"): + if isinstance(node.right, ast.Name): + split_axes_val = node.right.id + split_axis_pid_dim = self._get_program_id_dim(node.left) + + if isinstance(node.right, ast.Name) and node.right.id in self.program_id_vars: + if isinstance(node.left, ast.Name): + split_axes_val = node.left.id + split_axis_pid_dim = self.program_id_var_dims.get(node.right.id) + elif isinstance(node.right, ast.Call) and isinstance(node.right.func, ast.Attribute): + if (isinstance(node.right.func.value, ast.Name) and node.right.func.value.id == "tl" + and node.right.func.attr == "program_id"): + if isinstance(node.left, ast.Name): + split_axes_val = node.left.id + split_axis_pid_dim = self._get_program_id_dim(node.right) + + if split_axes_val in self.candidates_params and \ + split_axes_val not in self.split_axes.values(): + split_axes_key = self.get_axis(split_axes_val) + if split_axes_key and not self._is_tiling_only_split(split_axes_key, split_axes_val): + self.split_axes[split_axes_key] = split_axes_val + if split_axis_pid_dim is not None: + self._record_axis_pid_dim(split_axes_key, split_axis_pid_dim) + self.generic_visit(node) + + def visit_For(self, node): + if not isinstance(node.iter, ast.Call): + self.generic_visit(node) + return + + iter_fn = node.iter.func + is_range = isinstance(iter_fn, ast.Name) and iter_fn.id == "range" + is_tl_range = (isinstance(iter_fn, ast.Attribute) and isinstance(iter_fn.value, ast.Name) + and iter_fn.value.id == "tl" and iter_fn.attr == "range") + if not (is_range or is_tl_range): + self.generic_visit(node) + return + + if len(node.iter.args) == 0: + self.generic_visit(node) + return + + start = node.iter.args[0] if len(node.iter.args) >= 2 else None + stop = node.iter.args[1] if len(node.iter.args) >= 2 else node.iter.args[0] + pid_dim = self._extract_pid_dim_from_expr(start) + axis = self._axis_from_expr(stop) + if axis is not None and pid_dim is not None: + self._record_axis_pid_dim(axis, pid_dim) + if len(node.iter.args) >= 3: + step = node.iter.args[2] + loop_tiling_only_param = self._extract_grid_stride_split_param(start, step, pid_dim) + if loop_tiling_only_param is not None: + self._mark_tiling_only_param(axis, loop_tiling_only_param) + + self.generic_visit(node) + + def _get_program_id_dim(self, node): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and isinstance( + node.func.value, ast.Name) and node.func.value.id == "tl" and node.func.attr == "program_id"): + return None + + axis_dim = 0 + if len(node.args) > 0: + if isinstance(node.args[0], ast.Constant) and isinstance(node.args[0].value, int): + axis_dim = node.args[0].value + else: + return None + + for kw in node.keywords: + if kw.arg == "axis": + if isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, int): + axis_dim = kw.value.value + else: + return None + break + return axis_dim + + def _get_num_programs_dim(self, node): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and isinstance( + node.func.value, ast.Name) and node.func.value.id == "tl" and node.func.attr == "num_programs"): + return None + + axis_dim = 0 + if len(node.args) > 0: + if isinstance(node.args[0], ast.Constant) and isinstance(node.args[0].value, int): + axis_dim = node.args[0].value + else: + return None + + for kw in node.keywords: + if kw.arg == "axis": + if isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, int): + axis_dim = kw.value.value + else: + return None + break + return axis_dim + + def _extract_pid_dim_from_expr(self, node): + if node is None: + return None + for child in ast.walk(node): + if isinstance(child, ast.Name) and child.id in self.program_id_var_dims: + return self.program_id_var_dims[child.id] + pid_dim = self._get_program_id_dim(child) + if pid_dim is not None: + return pid_dim + return None + + def _contains_pid_dim(self, node, pid_dim): + if node is None: + return False + for child in ast.walk(node): + if isinstance(child, ast.Name): + if self.program_id_var_dims.get(child.id, None) == pid_dim: + return True + if self._get_program_id_dim(child) == pid_dim: + return True + return False + + def _contains_num_programs_dim(self, node, pid_dim): + if node is None: + return False + for child in ast.walk(node): + if isinstance(child, ast.Name): + if self.num_programs_var_dims.get(child.id, None) == pid_dim: + return True + if self._get_num_programs_dim(child) == pid_dim: + return True + return False + + def _is_candidate_name(self, node, candidate_name): + return (isinstance(node, ast.Name) and node.id == candidate_name and candidate_name in self.candidates_params) + + def _extract_pid_multiplied_candidate(self, node, pid_dim): + if node is None: + return None + candidates = set() + for child in ast.walk(node): + if not isinstance(child, ast.BinOp) or not isinstance(child.op, ast.Mult): + continue + left = child.left + right = child.right + if isinstance(left, ast.Name) and left.id in self.candidates_params and \ + self._contains_pid_dim(right, pid_dim): + candidates.add(left.id) + if isinstance(right, ast.Name) and right.id in self.candidates_params and \ + self._contains_pid_dim(left, pid_dim): + candidates.add(right.id) + if len(candidates) == 1: + return next(iter(candidates)) + return None + + def _contains_num_programs_multiplied_candidate(self, node, candidate_name, pid_dim): + if node is None: + return False + for child in ast.walk(node): + if not isinstance(child, ast.BinOp) or not isinstance(child.op, ast.Mult): + continue + if self._is_candidate_name(child.left, candidate_name): + if self._contains_num_programs_dim(child.right, pid_dim): + return True + if self._is_candidate_name(child.right, candidate_name): + if self._contains_num_programs_dim(child.left, pid_dim): + return True + return False + + def _extract_grid_stride_split_param(self, start, step, pid_dim): + if start is None or step is None: + return None + candidate_name = self._extract_pid_multiplied_candidate(start, pid_dim) + if candidate_name is None: + return None + if self._contains_num_programs_multiplied_candidate(step, candidate_name, pid_dim): + return candidate_name + return None + + def _mark_tiling_only_param(self, axis, candidate_name): + self.grid_stride_tiling_only.setdefault(axis, set()).add(candidate_name) + if self.split_axes.get(axis, None) == candidate_name: + del self.split_axes[axis] + self.split_axis_pid_dims.pop(axis, None) + + def _is_tiling_only_split(self, axis, candidate_name): + return candidate_name in self.grid_stride_tiling_only.get(axis, set()) + + def _axis_from_expr(self, node): + if node is None: + return None + for k, v in self.keys.items(): + if self.contains_target_var(node, v): + return k + return None + + def _record_axis_pid_dim(self, axis, pid_dim): + self.axis_pid_dims[axis] = pid_dim + if axis in self.split_axes: + self.split_axis_pid_dims[axis] = pid_dim + + +class TilingAxesParser(AxesKeyParser): + """ + Extracts the tiling axis parameters from triton kernel code. The parsing is based on the + `tl.arange`, `tl.range` and `range()` statement. This class identifies potential tiling axes by analyzing + the usage of the `range` and `tl.range` within `for` loop in the program. Common parameters + between `range()` or `tl.range` and `tl.arange` are extracted. It then filters these candidates based on a + list of candidate parameters (parameters not provided by the user). After that, it confirms the + tiling axis corresponding to the current parameter using mask comparison and the `keys` passed + in `autotune`. + + Note: + 1. Tiling axis parameters must be calculated within the `tl.arange` function and the `for` loop + using `tl.range`. + 2. Without mask comparision, it is impossible to confirm the exact tiling axis, which would lead + to parameter parsing failure. (eg. mask = offsets < n_elements). + 3. The identified tiling axes are limited to the list of candidated parameters, ensuring that + only those parameters that can be dynamically adjusted through the autotune process are considered. + """ + + def __init__(self, func_ast: ast.AST, keys: Dict[str, str], candidates_params: List[str]): + """ + :param func_ast: Abstract syntax tree of the triton kernel function + :type func_ast: ast.AST + :param keys: a dict of axis name: argument name, used to confirm the tiling axis corresponding to + the tiling axis parameters. + :type keys: Dict[str, str] + :param candidates_params: a list of parameters names that were not provided by the user when calling + triton kernel function. The parser will only consider these parameters as potential tiling axis + parameters. + :type candidates_params: List[str] + """ + super().__init__(func_ast, keys) + self.tiling_axes = dict() + self.candidates_params = candidates_params + self.candidates_params_for_loop = list() + + def parse(self) -> Dict[str, str]: + super().parse() + return self.tiling_axes + + def visit_For(self, node): + if isinstance(node.iter, ast.Call) and len(node.iter.args) == 3: + step_expr = node.iter.args[2] + for_loop_param = self._extract_unique_candidate(step_expr) + if (for_loop_param is not None and for_loop_param not in self.candidates_params_for_loop): + self.candidates_params_for_loop.append(for_loop_param) + self.generic_visit(node) + + def visit_Assign(self, node): + if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): + # handle FloorDiv + if isinstance(node.value, ast.BinOp) and isinstance(node.value.op, ast.FloorDiv): + denominator = node.value.right + denominator_param = self._extract_unique_candidate(denominator) + if denominator_param is not None and \ + denominator_param not in self.candidates_params_for_loop: + self.candidates_params_for_loop.append(denominator_param) + self.visit(self.func_ast) + + tiling_axes_val = self.get_tiling_axes_val(node.value) + if tiling_axes_val is not None and \ + tiling_axes_val in self.candidates_params_for_loop: + tiling_axes_key = self.get_axis(tiling_axes_val) + if tiling_axes_key: + self.tiling_axes[tiling_axes_key] = tiling_axes_val + self.generic_visit(node) + + def get_tiling_axes_val(self, node): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): + if node.func.attr == 'arange' and \ + isinstance(node.func.value, ast.Name) and \ + node.func.value.id == 'tl': + if isinstance(node.args, list) and len(node.args) == 2: + for param in self.candidates_params_for_loop: + if self.contains_target_var(node.args[1], param): + return param + + for _, value in ast.iter_fields(node): + if isinstance(value, list): + for item in value: + val = self.get_tiling_axes_val(item) + if val: + return val + elif isinstance(value, ast.AST): + val = self.get_tiling_axes_val(value) + if val: + return val + return None + + def _extract_unique_candidate(self, expr): + """ + Extract a unique tiling candidate from an expression. + Return None when no candidate or ambiguous (more than one candidate) appears. + """ + if expr is None: + return None + candidates = [param for param in self.candidates_params if self.contains_target_var(expr, param)] + if len(candidates) == 1: + return candidates[0] + return None + + +class ReductionAxesParser(AxesKeyParser): + """ + Extracts the reduction axis from triton kernel code. The parsing is based on the + reduction function (eg. tl.max, tl.min, tl.sum, ...). This class identifies the + dimensions of reduction operations by analyzing the reduction function calls in + the program. After that, It confirms the reduction axis corresponding to the current + parameter using mask comparison and the keys passed in autotune. + + Note: + 1. The call to the reduction function must start with 'tl', meaning the function must + be a function from triton.language + 2. It's preferable to specify the reduction axis dimension in the reduction function + using keyword arguments(eg. axis=xxx). Otherwise, specifying it via positional + arguments may lead to errors. + 3. Mask comparison must be performed on the potential reduction axis length parameters, + and the comparison parameters or target parameters of the comparison expression + must be sliced. Otherwise, the correspondence between dimensions and axes cannot + be confirmed, which will lead to failure in parsing the reduction axis. + 4. The identified reduction axes are limited to the candidate list provided in the keys. + """ + + def __init__(self, func_ast: ast.AST, keys: Dict[str, str]): + """ + :param func_ast: Abstract syntax tree of the triton kernel function + :type func_ast: ast.AST + :param keys: a dict of axis name: argument name, used to confirm the reduction axis. + :type keys: Dict[str, str] + """ + super().__init__(func_ast, keys) + self.reduction_axes = list() + self.reduction_func = ('sum', 'xor_sum', 'max', 'min', 'argmax', 'argmin') # tl.xxx + self.ndim = 1 + + def parse(self) -> List[str]: + super().parse() + return self.reduction_axes + + def visit_Assign(self, node): + self._scan_subscripts(node.value) + self.generic_visit(node) + + def _scan_subscripts(self, node): + if isinstance(node, ast.Subscript): + ndim = self._get_subscripts_ndim(node) + if ndim > self.ndim: + self.ndim = ndim + + for child in ast.iter_child_nodes(node): + self._scan_subscripts(child) + + def _get_subscripts_ndim(self, subscript_node): + slice_node = subscript_node.slice + + if isinstance(slice_node, ast.Tuple): + # e.g. [:, None] -> Tuple(elts=[Slice(), Constant(None)]) + return len(slice_node.elts) + elif isinstance(slice_node, (ast.Slice, ast.Constant, ast.Name, ast.UnaryOp, ast.BinOp)): + # e.g. [0], [:], [i], [-1], [i+1] + return 1 + else: + # Fallback: treat as 1D + return 1 + + def visit_Call(self, node): + if not isinstance(node.func, ast.Attribute): + return + func = node.func + if not isinstance(func.value, ast.Name) or func.value.id != 'tl': + self.generic_visit(node) + return + if func.attr not in self.reduction_func: + return + + axis_dim = None + args = node.args + if len(args) == 1: + # Axis passed as keyword argument + for keyword in node.keywords: + if keyword.arg == 'axis': + axis_dim = self.get_axis_dim(keyword.value) + break + + elif len(args) == 2: + # Axis passed as positional argument. Check the second param + axis_dim = self.get_axis_dim(args[1]) + + else: + raise ValueError("Reduction funtions args error") + + if axis_dim is not None: + reduction_axis = self.get_axis(axis_dim) + if reduction_axis and reduction_axis not in self.reduction_axes: + self.reduction_axes.append(reduction_axis) + + def get_axis_dim(self, node): + if isinstance(node, ast.Constant): + axis_dim = node.value + elif isinstance(node, ast.UnaryOp) and \ + isinstance(node.op, ast.USub): + operand = node.operand + if isinstance(operand, ast.Constant): + axis_dim = self.ndim - operand.value + else: + raise ValueError(f"Reduction function axis error, got: {ast.dump(node)}") + + if not isinstance(axis_dim, int): + raise ValueError("Reduction function axis must be an integer, " + f"got {type(node.value).__name__}: {node.value}") + return axis_dim + + def get_axis(self, axis_dim: int): + """ + Override the parent class method to accept an integer axis dimension + instead of a string. + + :param axis_dim: + :type axis_dim: int + """ + if axis_dim in self.checked_vars: + return None + self.checked_vars.append(axis_dim) + for node in ast.walk(self.func_ast): + if not isinstance(node, ast.Assign): + continue + reduction_axis = self.handle_assign_node(axis_dim, node) + if reduction_axis: + return reduction_axis + return None + + def handle_assign_node(self, axis_dim: int, node): + if not isinstance(node.value, ast.Compare): + return None + + # only support less than + if len(node.value.ops) != 1 or not isinstance(node.value.ops[0], ast.Lt): + return None + + target_axis_len = None + for axis_len in self.keys.values(): + if self.contains_target_var(node.value, axis_len): + target_axis_len = axis_len + break + if not target_axis_len: + return None + + # handel compare left var + if isinstance(node.value.left, ast.Name): + if self.check_compare_left(node.value.left.id, axis_dim): + reduction_axis = next((k for k, v in self.keys.items() if target_axis_len == v), None) + if reduction_axis and reduction_axis not in self.reduction_axes: + return reduction_axis + # handel compare target var + if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): + if self.check_compare_target(node.targets[0].id, axis_dim): + reduction_axis = next((k for k, v in self.keys.items() if target_axis_len == v), None) + if reduction_axis and reduction_axis not in self.reduction_axes: + return reduction_axis + return None + + def check_compare_left(self, var, axis_dim): + for node in ast.walk(self.func_ast): + if not isinstance(node, ast.Assign): + continue + if len(node.targets) != 1 or \ + not isinstance(node.targets[0], ast.Name) or \ + node.targets[0].id != var: + continue + if self.is_current_dim_slice(node.value, axis_dim): + return True + return False + + def check_compare_target(self, var, axis_dim, node=None): + if not node: + node = self.func_ast + for _, value in ast.iter_fields(node): + if isinstance(value, list): + for item in value: + if self.check_compare_target(var, axis_dim, item): + return True + elif isinstance(value, ast.AST): + if isinstance(value, ast.Subscript): + if not isinstance(value.value, ast.Name) or value.value.id != var: + continue + if self.is_current_dim_slice(value, axis_dim): + return True + else: + if self.check_compare_target(var, axis_dim, value): + return True + return False + + def is_current_dim_slice(self, node, dim): + for node in ast.walk(node): + if not isinstance(node, ast.Subscript) or not isinstance(node.slice, ast.Tuple): + continue + elts = node.slice.elts + if len(elts) != 0 and isinstance(elts[dim], ast.Slice): + return True + return False + + +class LowDimsAxesParser(AxesKeyParser): + """ + Extracts the low dimensions axis from triton kernel code. The parsing is based on the + `tl.arange` statement. This class identifies low dimensions axis by analyzing the usage + of the `tl.arange` in the program and extracts the variables computed by `tl.arange` and + their associated operations. Then it checks if these variables are involved in slicing + operations to determine dimension expansion and filters out variables that are expanded + in non-lowest dimensions. After that, it compares the extracted variables with the provided + `keys` to map them to specific low-dimensional axis. + + Note: + 1. low dimensions axis must be calculated within the `tl.arange` function and involved in + slicing operations to be identified. + 2. Without mask comparision, it is impossible to confirm the exact low dimensions axis, which + would lead to parameter parsing failure. (eg. mask = offsets < n_elements). + """ + + def __init__(self, func_ast: ast.AST, keys: Dict[str, str]): + """ + :param func_ast: Abstract syntax tree of the triton kernel function + :type func_ast: ast.AST + :param keys: a dict of axis name: argument name, used to confirm the low-dimensional axis. + :type keys: Dict[str, str] + """ + super().__init__(func_ast, keys) + self.low_dims_axis = list() + self.keys = keys + self.checked_slice_vars = list() + + def parse(self) -> List[str]: + super().parse() + return self.low_dims_axis + + def visit_Assign(self, node): + if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): + tl_arange_node = self.get_tl_arange_node(node) + low_dims_axis = None + if isinstance(tl_arange_node, ast.Call): + partin_other_slice = [False] + if self.is_partin_low_dim_slice(node.targets[0].id, partin_other_slice): + low_dims_axis = self.get_axis(node.targets[0].id) + elif not partin_other_slice[0]: + low_dims_axis = self.get_axis(node.targets[0].id) + elif isinstance(tl_arange_node, ast.Subscript) and \ + self.is_low_dim_slice(tl_arange_node, [False]): + low_dims_axis = self.get_axis(node.targets[0].id) + + if low_dims_axis and low_dims_axis not in self.low_dims_axis: + self.low_dims_axis.append(low_dims_axis) + self.generic_visit(node) + + def get_tl_arange_node(self, node): + for _, value in ast.iter_fields(node): + if isinstance(value, list): + for item in value: + if self.is_tl_arange_call(item): + return item + node = self.get_tl_arange_node(item) + if node: + return node + elif isinstance(value, ast.AST): + if self.is_tl_arange_call(value): + return value + node = self.get_tl_arange_node(value) + if node: + return node + return None + + def is_tl_arange_call(self, node): + """ + Checks if the given AST node is a call to `tl.arange` or a subscript of `tl.arange`. + It supports direct calls to `tl.arange` and subscripts of `tl.arange`, such as + `tl.arange()[None, :]` + """ + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): + if node.func.attr == 'arange' and \ + isinstance(node.func.value, ast.Name) and \ + node.func.value.id == 'tl': + return True + elif isinstance(node, ast.Subscript): + return self.is_tl_arange_call(node.value) + return False + + def is_low_dim_slice(self, node: ast.Subscript, partin_other_slice): + if not isinstance(node.slice, ast.Tuple) or not isinstance(node.slice.elts, list): + return False + elts = node.slice.elts + if len(elts) != 0 and not isinstance(elts[-1], ast.Slice): + partin_other_slice[0] = True + return False + return True + + def is_partin_low_dim_slice(self, var, partin_other_slice, node=None): + if not node: + node = self.func_ast + for child_node in ast.walk(node): + if isinstance(child_node, ast.Subscript) and isinstance(child_node.value, ast.Name): + if var == child_node.value.id and self.is_low_dim_slice(child_node, partin_other_slice): + return True + elif isinstance(child_node, ast.Assign): + if len(child_node.targets) == 1 and \ + isinstance(child_node.targets[0], ast.Name) and \ + var != child_node.targets[0].id: # Prevent cyclic assignment. + if not self.contains_target_var(child_node.value, var): + continue + target_var = child_node.targets[0].id + if target_var in self.checked_slice_vars: + continue + + if self.is_partin_low_dim_slice(var, partin_other_slice, child_node.value): + return True + if self.is_partin_low_dim_slice(target_var, partin_other_slice): + return True + + self.checked_slice_vars.append(var) + return False + + +class PtrNumsParser(AutoParser): + """ + Counts the number of pointer parameters from triton kernel code. The parsing of pointer-type + parameters is determined based on whether these parameters participate in memory access + statements such as `tl.load` and `tl.store`. + First, all input parameters in the kernel function are parsed, and then recursively, all variables + involved in the computation of each input parameter are identified. + If an input parameter directly participates in the computation of the first argument of `tl.load` + or `tl.store`, or if an intermediate variable computed from this input parameter indirectly + participates in the computation of the first argument of `tl.load` or `tl.store`, then this + parameter is considered a pointer-type parameter. + + Note: + 1. Variables modified with `tl.constexpr` are not pointer-type variables and will not be + further parsed. + 2. Only memory access statementes where the input parameter is directly involved or indirectly + involved through one level of computation are counted. Intermediate variables computed from + the input parameter through two or more levels of computation are not counted. + """ + + def __init__(self, func_ast: ast.AST, keys: Dict[str, str], miss_params: List[str]): + """ + :param func_ast: Abstract syntax tree of the triton kernel function + :type func_ast: ast.AST + :param keys: a dict of axis name: argument name, used to exclude potential ptr params. + :type keys: Dict[str, str] + :param miss_params: a list of parameters names that were not provided by the user when calling triton + kernel function. + :type miss_params: List[str] + """ + super().__init__(func_ast) + self.checked_vars = list() + self.ptr_nums = 0 + self.ptr_params = list() + self.keys = keys + self.miss_params = miss_params + self.constexpr_params = list() + + def parse(self): + super().parse() + return self.ptr_nums, self.ptr_params + + def visit_FunctionDef(self, node): + if isinstance(node.args, ast.arguments): + for arg in node.args.args: + if not isinstance(arg, ast.arg): + continue + + if isinstance(arg.annotation, ast.Attribute): + # var modified by tl.constexpr are not pointer type var, passed + is_tl = isinstance(arg.annotation.value, ast.Name) and \ + arg.annotation.value.id == 'tl' + if is_tl and arg.annotation.attr == 'constexpr': + if arg.arg not in self.constexpr_params: + self.constexpr_params.append(arg.arg) + continue + + if self.is_in_addr_calc(arg.arg) and arg.arg not in self.keys.values(): + self.ptr_params.append(arg.arg) + self.ptr_nums += 1 + + for miss_param in self.miss_params: + if miss_param not in self.constexpr_params: + print(f"[WARNING] The parameter '{miss_param}' needs to be declared as tl.constexpr!") + self.generic_visit(node) + + def is_in_addr_calc(self, var): + for node in ast.walk(self.func_ast): + if isinstance(node, ast.Call): + if isinstance(node.func, ast.Attribute) and \ + isinstance(node.func.value, ast.Name): + if node.func.value.id == "tl" and \ + (node.func.attr == "load" or node.func.attr == "store"): + if [arg for arg in node.args if self.contains_target_var(arg, var)]: + return True + + elif isinstance(node, ast.Assign): + if len(node.targets) == 1 and \ + isinstance(node.targets[0], ast.Name) and \ + var != node.targets[0].id: # Prevent cyclic assignment. + target_var = node.targets[0].id + if target_var in self.checked_vars: + continue + if isinstance(node.value, ast.BinOp) and \ + isinstance(node.value.op, ast.Add): + if isinstance(node.value.left, ast.Name) and \ + node.value.left.id == var: + if self.is_in_addr_calc(node.targets[0].id): + return True + elif isinstance(node.value.right, ast.Name) and \ + node.value.right.id == var: + if self.is_in_addr_calc(node.targets[0].id): + return True + self.checked_vars.append(var) + return False diff --git a/third_party/ascend/backend/runtime/autotuner.py b/third_party/ascend/backend/runtime/autotuner.py new file mode 100644 index 0000000000..a5a2ad1625 --- /dev/null +++ b/third_party/ascend/backend/runtime/autotuner.py @@ -0,0 +1,1134 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +from __future__ import annotations + +import builtins +import copy +import functools +import ast +import os +import time +from concurrent.futures import ThreadPoolExecutor +import itertools +from typing import Dict, List + +from torch import Tensor + +import triton +from triton.runtime.autotuner import Autotuner, Config +from triton.tools.get_ascend_devices import is_compile_on_910_95 + +from .autoparser import (LowDimsAxesParser, PtrNumsParser, ReductionAxesParser, SplitAxesParser, TilingAxesParser) +from .utils import get_byte_per_numel, is_valid_axis_name, valid_axis_names + + +class AutoTilingTuner(Autotuner): + """ + Automatic generateing candidate tiling configs and evaluating their performance to get the best config. + """ + + def __init__( + self, + fn, + arg_names, + configs, + key, + reset_to_zero, + restore_value, + pre_hook=None, + post_hook=None, + prune_configs_by: Dict = None, + warmup=None, + rep=None, + use_cuda_graph=False, + do_bench=None, + auto_profile_dir=None, + hints=None, + ): + """ + :param key: a list of argument name, where the change of arguments in value will triger re-generating candidates configs and evaluating. + The parameters in the list will be assigned axis names in sequence, with the axis name being in + {'x','y','z','w','v','t','rx','ry','rz','rw','rv','rt}, where the prefix 'r' means a reduction axis. + Only the axis name in this param should add perfix 'r' if it's a reduction axis. + :type key: List[str] + """ + super().__init__( + fn, + arg_names, + configs, + key, + reset_to_zero, + restore_value, + pre_hook, + post_hook, + prune_configs_by, + warmup, + rep, + use_cuda_graph, + do_bench, + ) + self.user_defined_do_bench = do_bench is not None + if not hints: + self.hints = {} + else: + self.hints = hints + split_params = self.hints.get("split_params", None) + tiling_params = self.hints.get("tiling_params", None) + low_dim_axes = self.hints.get("low_dim_axes", None) + reduction_axes = self.hints.get("reduction_axes", None) + self._init_axis_params( + key, + split_params, + tiling_params, + low_dim_axes, + reduction_axes, + ) + + self.auto_gen_config = not configs or self.hints.get("auto_gen_config", False) + self.gen_configs = [] # generated configs from TileGenerator + self.auto_profile_dir = auto_profile_dir + if not configs: + self.user_configs = [] + else: + self.user_configs = configs + self.is_simt_mode = False + self.simt_stack_limit = 8192 + self.user_specified_warps = None + self.user_specified_multibuffer = None + self.default_multibuffer = not is_compile_on_910_95 + self.print_autotuning = os.getenv("TRITON_PRINT_AUTOTUNING", None) == "1" + # Compile kernels in parallel by default for triton.runtime.JITFunction, + # but not for others, e.g., LibEntry, since it's not compatible with AsyncCompileMode + self.compile_parallel = (isinstance(self.fn, triton.runtime.JITFunction) + and os.getenv("TRITON_AUTOTUNE_PARALLEL_COMPILE", "1") == "1") + + def _expand_simt_num_warps_configs(self, base_configs: List[Config]) -> List[Config]: + _default_cand_num_warps = [8, 16, 32, 64] + cand_num_warps = (_default_cand_num_warps if self.user_specified_warps is None else [self.user_specified_warps]) + + simt_configs = [] + for base_cfg in base_configs: + for num_warps in cand_num_warps: + new_cfg = copy.deepcopy(base_cfg) + new_cfg.num_warps = num_warps + simt_configs.append(new_cfg) + + if self.print_autotuning: + print(f"Triton autotuning: Expanded to {len(simt_configs)} SIMT configs (with warps: {cand_num_warps})") + return simt_configs + + def _expand_simd_multibuffer_configs(self, base_configs: List[Config]) -> List[Config]: + if self.user_specified_multibuffer is not None: + if self.print_autotuning: + print("Triton autotuning: Skip SIMD multibuffer expansion because user " + f"specified multibuffer={self.user_specified_multibuffer}") + return base_configs + + opposite_default_multibuffer = not self.default_multibuffer + simd_configs = [] + for base_cfg in base_configs: + simd_configs.append(base_cfg) + new_cfg = copy.deepcopy(base_cfg) + new_cfg.kwargs["multibuffer"] = opposite_default_multibuffer + simd_configs.append(new_cfg) + + if self.print_autotuning: + print("Triton autotuning: Expanded to " + f"{len(simd_configs)} SIMD configs (toggle multibuffer={opposite_default_multibuffer})") + return simd_configs + + def _init_axis_params(self, key, split_params, tiling_params, low_dim_axes, reduction_axes): + if isinstance(key, list): + if (split_params or tiling_params or low_dim_axes or reduction_axes): + raise ValueError( + "If any axis-related parameters (split_params, tiling_params, low_dim_axes, reduction_axes)" + " are provided, 'key' must be a dict, not a list.") + if len(key) > len(valid_axis_names): + raise ValueError("Number of parameters exceeds the number of available axes.") + self.keys = {axis: param for axis, param in zip(valid_axis_names, key)} + elif isinstance(key, dict): + if not set(key.keys()).issubset(set(valid_axis_names)): + raise ValueError("All keys in 'key' must be valid axis names. Got unexpected keys.") + self.keys = key + if any([split_params, tiling_params, low_dim_axes, reduction_axes]) is None: + raise ValueError( + "If 'key' is a dict, all axis-related parameters (split_params, tiling_params, low_dim_axes," + " reduction_axes) must be provided.") + if not isinstance(split_params, dict): + raise ValueError("split_params must be a dict, got: {}".format(type(split_params))) + if not isinstance(tiling_params, dict): + raise ValueError("tiling_params must be a dict, got: {}".format(type(tiling_params))) + if not isinstance(low_dim_axes, list): + raise ValueError("low_dim_axes must be a list, got: {}".format(type(low_dim_axes))) + if not isinstance(reduction_axes, list): + raise ValueError("reduction_axes must be a list, got: {}".format(type(reduction_axes))) + + used_axes = set(split_params.keys()).union( + tiling_params.keys(), + low_dim_axes, + reduction_axes, + ) + if not used_axes.issubset(self.keys.keys()): + raise ValueError( + "The following axes are used but not present in the 'key': {}".format(used_axes - + set(self.keys.keys()))) + + self.split_params = split_params + self.all_split_params = {} + self.fixed_split_params = {} + self.tiling_params = tiling_params + self.low_dim_axes = low_dim_axes + self.reduction_axes = reduction_axes + self.fixed_grid_dims = set() + self.fixed_grid_dim_values = {} + self.split_axis_pid_dims = {} + self.axis_pid_dims = {} + self.dual_reduction = False + self.persistent_reduction = False + self.num_buffers = -1 + + def _autoparse_axis_params(self, all_args): + miss_params = [arg for arg in self.arg_names if arg not in all_args.keys()] + # parse pointer params nums + if self.num_buffers == -1: + self.num_buffers = self._autoparse_ptr_nums(all_args) + + # parse autotiling axes + # reduction axis must be parsed before other axes. it will alter the key + if not self.reduction_axes: + self.reduction_axes = self._autoparse_reduction_axes() + if len(self.reduction_axes) >= 2: + self.dual_reduction = True + + if not self.low_dim_axes: + self.low_dim_axes = self._autoparse_low_dim_axes() + + if len(self.reduction_axes) == 1: + reduction_axis = self.reduction_axes[0] + reduction_param = self.keys.get(reduction_axis, None) + reduction_numel = all_args.get(reduction_param, float("inf")) + persistent_threshold = self._get_persistent_reduction_threshold(reduction_axis) + if reduction_numel <= persistent_threshold: + self.persistent_reduction = True + + if not self.split_params: + all_split_params = self._autoparse_split_params(self._get_constexpr_candidates()) + self.all_split_params = dict(all_split_params) + self.fixed_split_params = {} + self.fixed_grid_dim_values = self._get_fixed_grid_dim_values( + all_args.get("grid", None), + all_args, + ) + self.fixed_grid_dims = set(self.fixed_grid_dim_values.keys()) + + fixed_grid_axes = {axis for axis, pid_dim in self.axis_pid_dims.items() if pid_dim in self.fixed_grid_dims} + + # Only missing constexpr params are tunable, and fixed-grid axes + # should not be tuned on split. + self.split_params = { + axis: param + for axis, param in all_split_params.items() + if param in miss_params and axis not in fixed_grid_axes + } + + # Fixed split is inferred only from fixed grid dims. + for axis, pid_dim in self.axis_pid_dims.items(): + if pid_dim not in self.fixed_grid_dims: + continue + core_num = self.fixed_grid_dim_values.get(pid_dim, 0) + axis_len_name = self.keys.get(axis, None) + axis_len = all_args.get(axis_len_name, None) + if not isinstance(core_num, int) or core_num <= 0: + continue + if not isinstance(axis_len, int) or axis_len <= 0: + continue + + self.fixed_split_params[axis] = (axis_len + core_num - 1) // core_num + elif not self.axis_pid_dims: + # When split axes are provided by hints, parse axis->program_id mapping + # independently for fixed-grid semantics and diagnostics. + self._autoparse_axis_pid_dims() + miss_params = [arg for arg in miss_params if arg not in self.split_params.values()] + if not self.tiling_params: + self.tiling_params = self._autoparse_tiling_params(miss_params) + miss_params = [arg for arg in miss_params if arg not in self.tiling_params.values()] + if miss_params: + raise ValueError(f"Missing required arguments: {miss_params}. " + f"These arguments must be explicitly provided and cannot be automatically tuned. " + f"Please ensure that these arguments are passed when calling the function.") + + def _gen_tile_configs(self, kv_dict: Dict[str, int], dtype: torch.dtype) -> List[Config]: + from .tile_generator import KernelMeta, TileGenerator + + axis_sizes = {} + for k, v in kv_dict.items(): + if not is_valid_axis_name(k): + continue + if not isinstance(v, int): + raise ValueError(f"Not supported dim type: {type(v)}, `int` is the only supported type") + axis_sizes[k] = v + + kernel_meta = KernelMeta( + axis_sizes, + self.split_params, + self.fixed_split_params, + self.tiling_params, + self.low_dim_axes, + dtype, + self.persistent_reduction, + self.dual_reduction, + self.num_buffers, + self.is_simt_mode, + ) + tile_gen = TileGenerator(kernel_meta=kernel_meta) + tile_gen.descend_split_tiling() + + self.gen_configs.clear() + self.gen_configs = tile_gen.configs + + if self.is_simt_mode: + self.gen_configs = self._expand_simt_num_warps_configs(self.gen_configs) + else: + self.gen_configs = self._expand_simd_multibuffer_configs(self.gen_configs) + + if len(self.gen_configs) == 0: + print("[WARNING] The generated candidate tiling configs are empty based on provided parameters!") + + if self.print_autotuning: + print("Generated configs number: {}".format(len(self.gen_configs))) + + def generate_key_and_configs(self, *args, **kwargs): + self.nargs = dict(zip(self.arg_names, args)) + self.is_simt_mode = kwargs.get('force_simt_only', False) + if 'num_warps' in kwargs and kwargs['num_warps'] is not None: + self.user_specified_warps = kwargs['num_warps'] + else: + self.user_specified_warps = None + if 'multibuffer' in kwargs and kwargs['multibuffer'] is not None: + self.user_specified_multibuffer = kwargs['multibuffer'] + else: + self.user_specified_multibuffer = None + + # generate key + all_args = {**self.nargs, **kwargs} + _args = {k: v for (k, v) in all_args.items() if k in self.arg_names} + key = [_args[v] for _, v in self.keys.items() if v in _args] + + # Currently, we use the dtype with maximum byte length + dtype = None + for _, arg in _args.items(): + if hasattr(arg, "dtype"): + key.append(str(arg.dtype)) + dtype = (arg.dtype if get_byte_per_numel(arg.dtype) >= get_byte_per_numel(dtype) else dtype) + if dtype is None: + raise NotImplementedError("Not support for non-Tensor inputs") + + key = tuple(key) + if key not in self.cache: + if self.auto_gen_config: + self._autoparse_axis_params(all_args) + _kv_dict = {k: _args[v] for k, v in self.keys.items() if v in _args} + self._gen_tile_configs(_kv_dict, dtype) + if len(self.gen_configs) == 0 and len(self.user_configs) == 0: + self.configs = [ + Config( + {}, + num_warps=4, + num_stages=2, + num_ctas=1, + num_buffers_warp_spec=0, + num_consumer_groups=0, + reg_dec_producer=0, + reg_inc_consumer=0, + ) + ] + else: + self.configs = self.gen_configs + self.user_configs + return key + + def run(self, *args, **kwargs): + key = self.generate_key_and_configs(*args, **kwargs) + if self.is_simt_mode and kwargs.get('simt_stack_limit', None) is None: + kwargs['simt_stack_limit'] = self.simt_stack_limit + used_cached_result = True + if key not in self.cache: + # prune configs + pruned_configs = self.prune_configs(kwargs) + if len(pruned_configs) > 1: + used_cached_result = False + bench_start = time.time() + timings = self._batch_bench(*args, configs=pruned_configs, **kwargs) + bench_end = time.time() + self.bench_time = bench_end - bench_start + self.cache[key] = builtins.min(timings, key=timings.get) + full_nargs = {**self.nargs, **kwargs, **self.cache[key].all_kwargs()} + self.pre_hook(full_nargs, reset_only=True) + self.configs_timings = timings + config = self.cache[key] + else: + config = pruned_configs[0] + else: + config = self.cache[key] + + self.best_config = config + if self.print_autotuning and not used_cached_result: + print(f"Triton autotuning for function {self.base_fn.__name__} finished after " + f"{self.bench_time:.2f}s; best config selected: {self.best_config};") + + if not used_cached_result and self.auto_profile_dir is not None: + self._profile(*args, config=self.best_config, **kwargs) + if config.pre_hook is not None: + full_nargs = {**self.nargs, **kwargs, **config.all_kwargs()} + config.pre_hook(full_nargs) + final_kwargs = dict(config.all_kwargs(), **kwargs) + ret = self.fn.run( + *args, + **final_kwargs, + ) + self.nargs = None + return ret + + def _batch_bench(self, *args, configs, **kwargs): + from triton.compiler.errors import CompileTimeAssertionFailure, MLIRCompilationError + from triton.runtime.errors import OutOfResources + + kernels_call = {config: self._make_kernel_call(*args, config=config, **kwargs) for config in configs} + run_fns = {} + exc = None + exc_stack = "" + + if self.compile_parallel: + import psutil + + max_workers = min(psutil.cpu_count(logical=False) // 2, len(kernels_call)) + future_kernels = [] + try: + with ( + ThreadPoolExecutor(max_workers=max_workers) as executor, + triton.AsyncCompileMode(executor), + ): + for config, fn in kernels_call.items(): + future_kernels.append((config, fn(warmup=True))) + + for config, fut in future_kernels: + try: + if hasattr(fut, "result"): + fut = fut.result() + run_fns[config] = functools.partial(kernels_call[config], warmup=False) + except (CompileTimeAssertionFailure, MLIRCompilationError) as e: + import traceback + exc_stack = traceback.format_exc() + exc = e + except Exception as e: + # ignore exception from __exit__() of AsyncCompileMode + triton.runtime._async_compile.active_mode.set(None) + else: + for config, fn in kernels_call.items(): + try: + fn(warmup=False) + run_fns[config] = functools.partial(fn, warmup=False) + except (CompileTimeAssertionFailure, MLIRCompilationError, OutOfResources) as e: + import traceback + exc_stack = traceback.format_exc() + exc = e + + if len(run_fns) == 0: + raise RuntimeError(f"No valid triton configs. {type(exc).__name__}: {exc} \nStack trace: {exc_stack}") + + if len(run_fns) == 1: + # we ignore expensive profiling method when only single config is left + return {config: self.do_bench(fn, quantiles=(0.5, 0.2, 0.8)) for config, fn in run_fns.items()} + + use_profiling = os.getenv("TRITON_BENCH_METHOD", "default").lower() == "npu" + # Respect user-provided benchmarkers even when NPU profiling mode is enabled. + use_npu_profiling = use_profiling and not self.user_defined_do_bench + if use_npu_profiling: + from ..testing import do_bench_npu + + time_cost = do_bench_npu(list(run_fns.values()), clear_l2_cache=False) + assert len(time_cost) == len(run_fns) + return {config: cost for config, cost in zip(run_fns.keys(), time_cost)} + else: + return {config: self.do_bench(fn, quantiles=(0.5, 0.2, 0.8)) for config, fn in run_fns.items()} + + def _make_kernel_call(self, *args, config, **meta): + # check for conflicts, i.e. meta-parameters both provided + # as kwargs and by the autotuner + conflicts = meta.keys() & config.kwargs.keys() + if conflicts: + raise ValueError(f"Conflicting meta-parameters: {', '.join(conflicts)}." + " Make sure that you don't re-define auto-tuned symbols.") + # augment meta-parameters with tunable ones + current = dict(meta, **config.all_kwargs()) + full_nargs = {**self.nargs, **current} + + def kernel_call(warmup): + if config.pre_hook: + config.pre_hook(full_nargs) + self.pre_hook(full_nargs) + try: + current.update({"warmup": warmup}) + res = self.fn.run( + *args, + **current, + ) + if warmup: + return res + except Exception as e: + try: + self.post_hook(full_nargs, exception=e) + finally: + # Throw exception raised by `self.fn.run` + raise + + self.post_hook(full_nargs, exception=None) + + return kernel_call + + def warmup(self, *args, **kwargs): + _ = self.generate_key_and_configs(*args, **kwargs) + pruned_configs = self.prune_configs(kwargs) + ret = [] + if self.compile_parallel: + import psutil + + max_workers = min(psutil.cpu_count(logical=False) // 2, len(pruned_configs)) + with ( + ThreadPoolExecutor(max_workers=max_workers) as executor, + triton.AsyncCompileMode(executor), + ): + for config in pruned_configs: + ret.append(self.fn.warmup(*args, **kwargs, **config.all_kwargs())) + else: + for config in pruned_configs: + ret.append(self.fn.warmup(*args, **kwargs, **config.all_kwargs())) + self.nargs = None + return ret + + def _profile(self, *args, config, **meta): + from ..testing import do_bench_npu + + kernel_call = self._make_kernel_call(*args, config=config, **meta) + fn = functools.partial(kernel_call, warmup=False) + do_bench_npu(fn, prof_dir=self.auto_profile_dir, keep_res=True) + + def _autoparse_split_params(self, candidates_params: List[str]) -> Dict[str, str]: + """ + Extracts the split axis parameters from triton kernel code. + """ + func_ast = self.fn.parse() + parser = SplitAxesParser(func_ast, self.keys, candidates_params) + split_axes = parser.parse() + self.split_axis_pid_dims = dict(getattr(parser, "split_axis_pid_dims", {})) + self.axis_pid_dims = dict(getattr(parser, "axis_pid_dims", {})) + if self.print_autotuning: + print(f"Ascend autotuning parse split axes: {split_axes}, " + f"split axis pid dims: {self.split_axis_pid_dims}, " + f"axis pid dims: {self.axis_pid_dims}") + return split_axes + + def _autoparse_axis_pid_dims(self) -> Dict[str, int]: + """ + Extract axis -> program_id dim mapping without relying on split-parameter + classification, so fixed-grid semantics can always consume it. + """ + func_ast = self.fn.parse() + parser = SplitAxesParser( + func_ast, + self.keys, + self._get_constexpr_candidates(), + ) + _ = parser.parse() + self.axis_pid_dims = dict(getattr(parser, "axis_pid_dims", {})) + self.split_axis_pid_dims = dict(getattr(parser, "split_axis_pid_dims", {})) + if self.print_autotuning: + print("Ascend autotuning parse axis pid dims (independent): " + f"{self.axis_pid_dims}") + return self.axis_pid_dims + + def _get_constexpr_candidates(self) -> List[str]: + """ + Returns all constexpr parameter names from the kernel function definition. + """ + func_ast = self.fn.parse() + constexpr_names = [] + for node in ast.walk(func_ast): + if not isinstance(node, ast.FunctionDef): + continue + if not isinstance(node.args, ast.arguments): + continue + for arg in node.args.args: + if not isinstance(arg, ast.arg): + continue + ann = arg.annotation + if (isinstance(ann, ast.Attribute) and isinstance(ann.value, ast.Name) and ann.value.id == "tl" + and ann.attr == "constexpr"): + constexpr_names.append(arg.arg) + break + return constexpr_names + + def _get_fixed_grid_dim_values(self, grid, all_args: Dict[str, object] = None) -> Dict[int, int]: + """ + Returns fixed grid dim -> value. + - Static tuple/list grid: direct extraction + - Callable grid: infer fixed dims by perturbing missing constexpr params + """ + if grid is None: + return {} + if callable(grid): + return self._infer_fixed_dims_from_callable_grid(grid, all_args or {}) + return self._extract_fixed_grid_dims(grid) + + def _extract_fixed_grid_dims(self, grid) -> Dict[int, int]: + if isinstance(grid, int): + grid = (grid, ) + if not isinstance(grid, (tuple, list)): + return {} + fixed_dims = {} + for idx, dim in enumerate(grid): + if isinstance(dim, int) and dim > 0: + fixed_dims[idx] = dim + return fixed_dims + + def _normalize_grid_tuple(self, grid_out): + if isinstance(grid_out, int): + return (grid_out, ) + if isinstance(grid_out, (tuple, list)): + return tuple(grid_out) + return None + + def _infer_fixed_dims_from_callable_grid(self, grid_fn, all_args: Dict[str, object]) -> Dict[int, int]: + constexpr_candidates = self._get_constexpr_candidates() + base_meta = dict(all_args or {}) + + # Fill missing constexpr with stable probe defaults so grid(meta) can execute. + for name in constexpr_candidates: + if name not in base_meta: + base_meta[name] = 128 + + try: + base_grid_raw = grid_fn(dict(base_meta)) + except Exception: + return {} + + base_grid = self._normalize_grid_tuple(base_grid_raw) + if base_grid is None: + return {} + + dynamic_dims = set() + # Missing constexpr are tunable candidates. + tunable_probe_names = [name for name in constexpr_candidates if name not in (all_args or {})] + probe_values = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512] + + for name in tunable_probe_names: + baseline = base_meta.get(name, 128) + for probe in probe_values: + if probe == baseline: + continue + probe_meta = dict(base_meta) + probe_meta[name] = probe + try: + probe_grid_raw = grid_fn(probe_meta) + except Exception: + continue + probe_grid = self._normalize_grid_tuple(probe_grid_raw) + if probe_grid is None: + continue + if len(probe_grid) != len(base_grid): + dynamic_dims.update(range(min(len(probe_grid), len(base_grid)))) + continue + for idx, (base_dim, probe_dim) in enumerate(zip(base_grid, probe_grid)): + if not (isinstance(base_dim, int) and isinstance(probe_dim, int)): + dynamic_dims.add(idx) + continue + if base_dim != probe_dim: + dynamic_dims.add(idx) + + fixed_dims = {} + for idx, dim in enumerate(base_grid): + if idx in dynamic_dims: + continue + if isinstance(dim, int) and dim > 0: + fixed_dims[idx] = dim + return fixed_dims + + def _autoparse_tiling_params(self, candidates_params: List[str]) -> Dict[str, str]: + """ + Extracts the tiling axis parameters from triton kernel code. + """ + func_ast = self.fn.parse() + parser = TilingAxesParser(func_ast, self.keys, candidates_params) + tiling_axes = parser.parse() + if self.print_autotuning: + print(f"Ascend autotuning parse tiling axes: {tiling_axes}") + return tiling_axes + + def _autoparse_reduction_axes(self) -> List[str]: + """ + Extracts the reduction axis parameters from triton kernel code. + """ + func_ast = self.fn.parse() + parser = ReductionAxesParser(func_ast, self.keys) + reduction_axes = parser.parse() + for axis in reduction_axes: + self.keys[f"r{axis}"] = self.keys.pop(axis) + reduction_axes = [f"r{axis}" for axis in reduction_axes] + + if self.print_autotuning: + print(f"Ascend autotuning parse keys: {self.keys} \n" + f"Ascend autotuning parse reduction axes: {reduction_axes}") + return reduction_axes + + def _autoparse_low_dim_axes(self) -> List[str]: + """ + Extracts the low dimension axis from triton kernel code. + """ + func_ast = self.fn.parse() + parser = LowDimsAxesParser(func_ast, self.keys) + low_dim_axes = parser.parse() + if len(low_dim_axes) < 1: + if self.print_autotuning: + print("[WARNING] Failed to parse low-dimensional axes, fallback to empty low_dim_axes.") + return [] + if self.print_autotuning: + print(f"Ascend autotuning parse low dimensional axes: {low_dim_axes}") + return low_dim_axes + + def _autoparse_ptr_nums(self, all_args: dict) -> int: + """ + Counts the number of pointer parameters from triton kernel code. + """ + ptr_nums = 0 + ptr_params = list() + for k, v in all_args.items(): + if isinstance(v, Tensor): + ptr_nums += 1 + ptr_params.append(k) + + if self.print_autotuning: + print(f"Ascend autotuning parse pointer params: {ptr_params}, pointer nums: {ptr_nums}") + return ptr_nums + + def _get_persistent_reduction_threshold(self, reduction_axis: str) -> int: + # Keep this heuristic aligned with inductor-style policy: + # inner reduction axis uses a larger threshold than other axes. + if self.low_dim_axes and reduction_axis == self.low_dim_axes[0]: + return 1024 + return 64 + + +def autotune(configs, key, prune_configs_by=None, reset_to_zero=None, restore_value=None, pre_hook=None, post_hook=None, + warmup=None, rep=None, use_cuda_graph=False, do_bench=None, *, auto_prof_dir=None, hints=None): + """ + Decorator for auto-tuning a :code:`triton.jit`'d function. + + .. highlight:: python + .. code-block:: python + + @triton.autotune(configs=[ + triton.Config(kwargs={'BLOCK_SIZE': 128}, num_warps=4), + triton.Config(kwargs={'BLOCK_SIZE': 1024}, num_warps=8), + ], + key=['x_size'] # the two above configs will be evaluated anytime + # the value of x_size changes + ) + @triton.jit + def kernel(x_ptr, x_size, **META): + BLOCK_SIZE = META['BLOCK_SIZE'] + :note: When all the configurations are evaluated, the kernel will run multiple times. + This means that whatever value the kernel updates will be updated multiple times. + To avoid this undesired behavior, you can use the `reset_to_zero` argument, which + resets the value of the provided tensor to `zero` before running any configuration. + + If the environment variable :code:`TRITON_PRINT_AUTOTUNING` is set to + :code:`"1"`, Triton will print a message to stdout after autotuning each + kernel, including the time spent autotuning and the best configuration. + + :param configs: a list of :code:`triton.Config` objects + :type configs: list[triton.Config] + :param key: a list of argument names whose change in value will trigger the evaluation of all provided configs. + :type key: list[str] + :param prune_configs_by: a dict of functions that are used to prune configs, fields: + 'perf_model': performance model used to predicate running time with different configs, returns running time + 'top_k': number of configs to bench + 'early_config_prune'(optional): a function used to do early prune (eg, num_stages). It takes configs:List[Config] as its input, and returns pruned configs. + :param reset_to_zero: a list of argument names whose value will be reset to zero before evaluating any configs. + :type reset_to_zero: list[str] + :param restore_value: a list of argument names whose value will be restored after evaluating any configs. + :type restore_value: list[str] + :param pre_hook: a function that will be called before the kernel is called. + This overrides the default pre_hook used for 'reset_to_zero' and 'restore_value'. + 'kwargs': a dict of all arguments passed to the kernel. + 'reset_only': a boolean indicating whether the pre_hook is called to reset the values only, without a corresponding post_hook. + :type pre_hook: lambda args, reset_only + :param post_hook: a function that will be called after the kernel is called. + This overrides the default post_hook used for 'restore_value'. + 'kwargs': a dict of all arguments passed to the kernel. + 'exception': the exception raised by the kernel in case of a compilation or runtime error. + :type post_hook: lambda args, exception + :param warmup: warmup time (in ms) to pass to benchmarking (deprecated). + :type warmup: int + :param rep: repetition time (in ms) to pass to benchmarking (deprecated). + :type rep: int + :param do_bench: a benchmark function to measure the time of each run. + :type do_bench: lambda fn, quantiles + :param auto_prof_dir: the specified directory to store the profiling results of the best config. + If this parameter is None or the best config is retrieved from cache, the profiling process will be ignored. + :type auto_prof_dir: str + :param hints: a dict of autotune hint auguments passed to AutoTilingTuner. + """ + + def decorator(fn): + return AutoTilingTuner(fn, fn.arg_names, configs, key, reset_to_zero, restore_value, pre_hook=pre_hook, + post_hook=post_hook, prune_configs_by=prune_configs_by, warmup=warmup, rep=rep, + use_cuda_graph=use_cuda_graph, do_bench=do_bench, auto_profile_dir=auto_prof_dir, + hints=hints) + + return decorator + + +_ALL_PARAMS = { + "num_stages", + "unit_flag", + "limit_auto_multi_buffer_only_for_local_buffer", + "limit_auto_multi_buffer_of_local_buffer", + "set_workspace_multibuffer", + "enable_hivm_auto_cv_balance", + "tile_mix_vector_loop", + "tile_mix_cube_loop", + "enable_ubuf_saving", +} + +_DEFAULTS = { + "num_stages": [2], + "unit_flag": [False], + "limit_auto_multi_buffer_only_for_local_buffer": [False], + "limit_auto_multi_buffer_of_local_buffer": ["no-l0c"], + "set_workspace_multibuffer": [2, 4], + "enable_hivm_auto_cv_balance": [True], + "tile_mix_vector_loop": [2, 4], + "tile_mix_cube_loop": [2, 4], + "enable_ubuf_saving": [True], +} + +_VALID_VALUES = { + "num_stages": [1, 2], + "limit_auto_multi_buffer_of_local_buffer": ["no-limit", "no-l0c"], + "set_workspace_multibuffer": [2, 4], + "tile_mix_vector_loop": [2, 4, 8], + "tile_mix_cube_loop": [2, 4, 8], +} + +_CUBE_PARAMS = {"num_stages", "unit_flag", "limit_auto_multi_buffer_of_local_buffer"} + +_MIXCV_PARAMS = { + "num_stages", + "unit_flag", + "limit_auto_multi_buffer_only_for_local_buffer", + "limit_auto_multi_buffer_of_local_buffer", + "set_workspace_multibuffer", + "enable_hivm_auto_cv_balance", + "tile_mix_vector_loop", + "tile_mix_cube_loop", + "enable_ubuf_saving", +} + +_VECTOR_PARAMS = { + "num_stages", + "enable_ubuf_saving", +} + + +def _check_boolean_list(val, param_name): + return isinstance(val, (list, tuple)) and len(val) > 0 and all(isinstance(x, bool) for x in val) + + +def _check_string_in_set(val, valid_set, param_name): + return isinstance(val, (list, tuple)) and len(val) > 0 and all(v in valid_set for v in val) + + +def _check_int_in_set(val, valid_set, param_name): + return isinstance(val, (list, tuple)) and len(val) > 0 and all(isinstance(v, int) and v in valid_set for v in val) + + +_VALIDATION_RULES = { + "num_stages": { + "desc": f"must be one or more of: {_VALID_VALUES['num_stages']}", "check": + lambda val, p: _check_int_in_set(val, _VALID_VALUES['num_stages'], p) + }, + "unit_flag": {"desc": "must be non-empty list/tuple of boolean values", "check": _check_boolean_list}, + "limit_auto_multi_buffer_only_for_local_buffer": + {"desc": "must be non-empty list/tuple of boolean values", "check": _check_boolean_list}, + "limit_auto_multi_buffer_of_local_buffer": { + "desc": f"must be one or more of: {_VALID_VALUES['limit_auto_multi_buffer_of_local_buffer']}", "check": + lambda val, p: _check_string_in_set(val, _VALID_VALUES['limit_auto_multi_buffer_of_local_buffer'], p) + }, + "set_workspace_multibuffer": { + "desc": f"must be one or more of: {_VALID_VALUES['set_workspace_multibuffer']}", "check": + lambda val, p: _check_int_in_set(val, _VALID_VALUES['set_workspace_multibuffer'], p) + }, + "enable_hivm_auto_cv_balance": + {"desc": "must be non-empty list/tuple of boolean values", "check": _check_boolean_list}, + "tile_mix_vector_loop": { + "desc": f"must be one or more of: {_VALID_VALUES['tile_mix_vector_loop']}", "check": + lambda val, p: _check_int_in_set(val, _VALID_VALUES['tile_mix_vector_loop'], p) + }, + "tile_mix_cube_loop": { + "desc": f"must be one or more of: {_VALID_VALUES['tile_mix_cube_loop']}", "check": + lambda val, p: _check_int_in_set(val, _VALID_VALUES['tile_mix_cube_loop'], p) + }, + "enable_ubuf_saving": {"desc": "must be non-empty list/tuple of boolean values", "check": _check_boolean_list}, +} + + +class BaseAutotuner: + """ + Base class for generating auto-tuning configurations without block dimensions. + Users must provide fixed dimension parameters when calling the kernel. + """ + + def __init__(self, operator_name, supported_params, default_params, validation_rules): + self.operator_name = operator_name + self.supported_params = supported_params + self.default_params = default_params + self.validation_rules = validation_rules + + def validate_parameters(self, **kwargs): + # Check for unsupported parameters + invalid_params = [k for k in kwargs.keys() if k not in _ALL_PARAMS] + if invalid_params: + print(f"[ERROR] Invalid parameters for {self.operator_name}: {invalid_params}") + return False + + for param, rule in self.validation_rules.items(): + if param in kwargs: + if not rule["check"](kwargs[param], param): + print(f"[ERROR] Invalid value for '{param}' in {self.operator_name}: {kwargs[param]}") + print(f" Expected: {rule['desc']}") + return False + return True + + def get_configs(self, **kwargs): + """ + Generate a list of Config objects. + Each parameter must be provided as a list (even for a single value). + The function produces the Cartesian product of all parameter lists. + - num_stages: each value will be set as Config.num_stages (not placed in kwargs) + - other parameters: each value will be placed in Config.kwargs + Returns a list of Config objects. + """ + if not self.validate_parameters(**kwargs): + return [] + + # Collect parameter values, using defaults for missing ones + param_values = {} + for p in sorted(self.supported_params): + if p in kwargs: + param_values[p] = kwargs[p] + else: + param_values[p] = self.default_params.get(p, [None]) + + keys = list(param_values.keys()) + values = list(param_values.values()) + combos = list(itertools.product(*values)) + + configs = [] + for combo in combos: + config_kwargs = {} + num_stages_val = None + for i, pname in enumerate(keys): + val = combo[i] + if pname == "num_stages": + num_stages_val = val + else: + config_kwargs[pname] = val + + configs.append(Config(kwargs=config_kwargs, num_stages=num_stages_val if num_stages_val is not None else 2)) + return configs + + +CubeAutotuner = BaseAutotuner(operator_name="cube", supported_params=_CUBE_PARAMS, default_params=_DEFAULTS, + validation_rules=_VALIDATION_RULES) + +MixcvAutotuner = BaseAutotuner(operator_name="mixcv", supported_params=_MIXCV_PARAMS, default_params=_DEFAULTS, + validation_rules=_VALIDATION_RULES) + +VectorAutotuner = BaseAutotuner(operator_name="vector", supported_params=_VECTOR_PARAMS, default_params=_DEFAULTS, + validation_rules=_VALIDATION_RULES) + + +def get_autotune_cube_config(**kwargs: Any) -> List[triton.Config]: + """ + Generate autotune configuration for the cube operator. + Supported parameters: num_stages, unit_flag, limit_auto_multi_buffer_of_local_buffer. + """ + import triton + return CubeAutotuner.get_configs(**kwargs) + + +def get_autotune_cv_config(**kwargs: Any) -> List[triton.Config]: + """ + Generate autotune configuration for the mixcv operator. + Supported parameters: num_stages, unit_flag, limit_auto_multi_buffer_only_for_local_buffer, + limit_auto_multi_buffer_of_local_buffer, set_workspace_multibuffer, + enable_hivm_auto_cv_balance, tile_mix_vector_loop, tile_mix_cube_loop, enable_ubuf_saving + """ + import triton + return MixcvAutotuner.get_configs(**kwargs) + + +def get_autotune_vector_config(**kwargs: Any) -> List[triton.Config]: + """ + Generate autotune configuration for the vector operator. + Supported parameters: num_stages, enable_ubuf_saving + """ + import triton + return VectorAutotuner.get_configs(**kwargs) + + +def get_max_configs(config, kernel_type="mixcv", **kwargs): + """ + Expand a single base Config by combining it with tuning parameters. + + :param config: A triton.Config object serving as the base. + :param kernel_type: Operator type, one of "cube", "mixcv", "vector". Default "mixcv". + :param kwargs: Tuning parameters, each provided as a list (e.g., enable_hivm_auto_cv_balance=[True, False]). + If a parameter is not provided, its value is taken from the base config (if present) + or from the defaults. + :return: List of expanded Config objects. + """ + # Determine the set of parameters supported by the current kernel_type + if kernel_type == "cube": + supported = _CUBE_PARAMS + elif kernel_type == "vector": + supported = _VECTOR_PARAMS + else: + supported = _MIXCV_PARAMS + + # Warn about unsupported parameters provided in kwargs + unsupported = [k for k in kwargs if k not in supported and k in _ALL_PARAMS] + if unsupported: + print( + f"[WARNING] The following parameters are not supported for kernel_type '{kernel_type}': {unsupported}. They will be ignored." + ) + + # Build value lists for each parameter (priority: kwargs > base config > defaults) + param_values = {} + base_kwargs = config.kwargs + base_num_stages = config.num_stages + + for param in sorted(supported): + if param in kwargs: + # User-provided list via tuning_params takes precedence + val_list = kwargs[param] + elif param == "num_stages": + # num_stages is an attribute of Config, not part of kwargs. + # If not provided in tuning_params, use the value from the base config as a fixed single-element list. + val_list = [base_num_stages] + elif param in base_kwargs: + # Parameter present in base config's kwargs -> fix to that single value + val_list = [base_kwargs[param]] + else: + # Otherwise fall back to defaults + val_list = _DEFAULTS.get(param, [None]) + + # Validate the value list + if param in _VALIDATION_RULES: + rule = _VALIDATION_RULES[param] + if not rule["check"](val_list, param): + raise ValueError(f"Invalid value for '{param}': {val_list}. Expected: {rule['desc']}") + param_values[param] = val_list + + # Cartesian product of all parameter lists + keys = list(param_values.keys()) + values = list(param_values.values()) + combos = list(itertools.product(*values)) + + new_configs = [] + for combo in combos: + # Start with a copy of the original config's kwargs + new_kwargs = config.kwargs.copy() + num_stages_val = None + + for i, pname in enumerate(keys): + val = combo[i] + if pname == "num_stages": + num_stages_val = val + else: + # Overwrite or add the parameter to kwargs + new_kwargs[pname] = val + + new_config = Config(kwargs=new_kwargs, num_warps=config.num_warps, + num_stages=num_stages_val if num_stages_val is not None else config.num_stages, + num_ctas=config.num_ctas, num_buffers_warp_spec=config.num_buffers_warp_spec, + num_consumer_groups=config.num_consumer_groups, reg_dec_producer=config.reg_dec_producer, + reg_inc_consumer=config.reg_inc_consumer, maxnreg=config.maxnreg, pre_hook=config.pre_hook) + new_configs.append(new_config) + + return new_configs + + +def max_autotune(configs, key, kernel_type="mixcv", prune_configs_by=None, reset_to_zero=None, restore_value=None, + pre_hook=None, post_hook=None, warmup=None, rep=None, use_cuda_graph=False, do_bench=None, + **tuning_params): + """ + Decorator that expands each base Config with tuning parameters before auto-tuning. + + Usage is similar to @triton.autotune, but allows automatic expansion of + additional tuning parameters (e.g., enable_hivm_auto_cv_balance, tile_mix_vector_loop, ...) + for each provided base configuration. + + :param configs: List of base triton.Config objects. + :param key: List of argument names whose change triggers re-tuning. + :param kernel_type: Operator type, one of "cube", "mixcv", "vector". Default "mixcv". + :param prune_configs_by: Same as in autotune. + :param reset_to_zero: Same as in autotune. + :param restore_value: Same as in autotune. + :param pre_hook: Same as in autotune. + :param post_hook: Same as in autotune. + :param warmup: Deprecated. + :param rep: Deprecated. + :param use_cuda_graph: Deprecated. + :param do_bench: Same as in autotune. + :param tuning_params: Additional tuning parameters as keyword arguments. + Each value must be a list; the Cartesian product of these lists + will be combined with each base config. + """ + + def decorator(fn): + if not configs or len(configs) == 0: + raise ValueError("[max_autotune] The argument 'configs' cannot be empty. " + "Please provide at least one base config. ") + # Expand each base config with the provided tuning parameters + expanded_configs = [] + for cfg in configs: + expanded = get_max_configs(cfg, kernel_type=kernel_type, **tuning_params) + expanded_configs.extend(expanded) + + # Call the original autotune decorator with the expanded configs + return autotune(configs=expanded_configs, key=key, prune_configs_by=prune_configs_by, + reset_to_zero=reset_to_zero, restore_value=restore_value, pre_hook=pre_hook, + post_hook=post_hook, warmup=warmup, rep=rep, use_cuda_graph=use_cuda_graph, + do_bench=do_bench)(fn) + + return decorator diff --git a/third_party/ascend/backend/runtime/tile_generator.py b/third_party/ascend/backend/runtime/tile_generator.py new file mode 100644 index 0000000000..a4e2a9a83c --- /dev/null +++ b/third_party/ascend/backend/runtime/tile_generator.py @@ -0,0 +1,489 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +from __future__ import annotations + +import functools +import sys +from dataclasses import dataclass +from typing import ( + Dict, + List, + Tuple, +) + +from triton.runtime.autotuner import Config + +from .utils import ( + get_byte_per_numel, + next_power_of_2, + num_vector_core, + ub_size_in_kbytes, + rf_size_in_kbytes, +) + + +@dataclass +class AxisInfo: + name: str + index: int + length: int + + prefix: str = "" + split_name: str = "" + tiling_name: str = "" + is_split_axis: bool = False + is_tunable_split_axis: bool = False + is_tiling_axis: bool = False + fixed_split_size: int = 0 + + @property + def is_reduction(self): + return self.prefix == "r" + + +class KernelMeta: + + def __init__( + self, + axis_sizes: Dict[str, int], + split_params: Dict[str, str], + fixed_split_params: Dict[str, int], + tiling_params: Dict[str, str], + low_dims: List[str], + dtype: torch.dtype, + persistent_reduction: bool, + dual_reduction: bool, + num_buffers: int, + is_simt_mode: bool, + ): + """ + :param split_params: a dict of axis name: argument name, the argument is an adjustable parameter in a split axis, such as 'XBLOCK'. + The axis name must be in key's axis names. Do not add prefix 'r' before the axis name. + This param can be empty. Note that the auto tiling feature will be disabled when the split_params and tiling_params are both empty. + The split axis can usually be identified according to `tl.program_id()` expression. + :type split_params: Dict[str, str] + :param tiling_params: a dict of axis name: argument name, the argument is an adjustable parameter in a tiling axis, such as 'XBLOCK_SUB'. + The axis name must be in key's axis names. Do not add prefix 'r' before the axis name. + This param can be empty. Note that the auto tiling feature will be disabled when the split_params and tiling_params are both empty. + The tiling axis can usually be identified according to `tl.arange()` expression. + :type tiling_params: Dict[str, str] + :param low_dims: a list of axis name in which the corresponding axis is low dim aixs. + The axis name must be in key's axis names. Do not add prefix 'r' before the axis name. + :type low_dims: List[str] + :param dual_reduction: performing reduction on more than one axis. + :param persistent_reduction: there is no splitting in reduction axis. + """ + self._validate_axis(axis_sizes, split_params, fixed_split_params, tiling_params, low_dims) + + axis_dict = {} + idx = 0 + for name, length in axis_sizes.items(): + prefix = "" + if name.startswith("r"): + prefix = "r" + + is_tunable_split_axis = name in split_params + fixed_split_size = fixed_split_params.get(name, 0) + is_split_axis = is_tunable_split_axis or fixed_split_size > 0 + is_tiling_axis = name in tiling_params + split_name = "" if not is_tunable_split_axis else split_params[name] + tiling_name = "" if name not in tiling_params else tiling_params[name] + + axis_dict[name] = AxisInfo( + name=name, + index=idx, + length=length, + prefix=prefix, + split_name=split_name, + tiling_name=tiling_name, + is_split_axis=is_split_axis, + is_tunable_split_axis=is_tunable_split_axis, + is_tiling_axis=is_tiling_axis, + fixed_split_size=fixed_split_size, + ) + idx += 1 + + self.axis_info = list(axis_dict.values()) + self.split_axis = [x for x in axis_dict.values() if x.is_split_axis] + self.tunable_split_axis = [x for x in axis_dict.values() if x.is_tunable_split_axis] + self.tiling_axis = [x for x in axis_dict.values() if x.is_tiling_axis] + self.low_dims_axis = [x for x in axis_dict.values() if x.name in low_dims] + self.dtype = dtype + self.persistent_reduction = persistent_reduction + self.dual_reduction = dual_reduction + self.num_buffers = num_buffers + self.is_simt_mode = is_simt_mode + + @classmethod + def _validate_axis( + cls, + axis_sizes: Dict[str, int], + split_params: Dict[str, str], + fixed_split_params: Dict[str, int], + tiling_params: Dict[str, str], + low_dims: List[str], + ) -> None: + for axis_name in axis_sizes.keys(): + if axis_name.startswith("r") and len(axis_name) == 1: + raise ValueError("The name of a reduction axis is empty!") + + def check_keys(params: List[str], context="parameter"): + for k in params: + if k not in axis_sizes and ("r" + k) not in axis_sizes: + raise KeyError(f"{context} '{k}' not found in known axes: {axis_sizes.keys()}") + + check_keys(split_params.keys(), "split axis") + check_keys(fixed_split_params.keys(), "fixed split axis") + check_keys(tiling_params.keys(), "tiling axis") + check_keys(low_dims, "low dim axis") + + +@dataclass +class BlockInfo: + block_name: str # e.g., XBLOCK + sub_block_name: str # e.g., XBLOCK_SUB + block_size: int + sub_block_size: int + + +""" +Generate possible candidate tiling configs for benchmarking +""" + + +class TileGenerator: + num_warps = 1 + num_stages = 1 + + def __init__(self, kernel_meta: KernelMeta): + self.kernel_meta = kernel_meta + self.persistent_reduction = self.kernel_meta.persistent_reduction + self.dual_reduction = self.kernel_meta.dual_reduction + + self.blocks = self.init_blocks_info(kernel_meta) + self.numels = [axis.length for axis in kernel_meta.axis_info] + self.candidate_blocks = [] + self.configs = [] + self.dtype_bytes = get_byte_per_numel(kernel_meta.dtype) + + self.num_buffers = 3 if kernel_meta.num_buffers == 0 else min(kernel_meta.num_buffers, 3) + self.is_simt_mode = kernel_meta.is_simt_mode + local_mem_size = (rf_size_in_kbytes if self.is_simt_mode else ub_size_in_kbytes) + self.max_numel_threshold = local_mem_size * 1024 // self.dtype_bytes // self.num_buffers + self.max_total_numel = functools.reduce(lambda x, y: x * y, [x.block_size + for x in self.blocks]) if self.blocks else 1 + self.small_kernel = self.max_total_numel < 128 * 1024 + self.tiny_kernel = self.max_total_numel <= 32 * 1024 + self.stop_numel = min(1024 // self.dtype_bytes, self.max_total_numel // + (num_vector_core * 2)) if self.small_kernel else 1024 // self.dtype_bytes + self.max_programs_num = 65535 + self.tiny_program_threshold = num_vector_core // 8 + self.tiny_per_program_cap = 1 + self.tiny_low_program_hist = {p: 0 for p in range(1, self.tiny_program_threshold + 1)} + self.tiny_low_program_active = False + self.tiny_low_program_tile_floor = 0 + + @classmethod + def init_blocks_info(cls, kernel_meta: KernelMeta) -> List[BlockInfo]: + blocks = [] + for axis in kernel_meta.axis_info: + block_name = axis.split_name + sub_block_name = axis.tiling_name + block_size = axis.fixed_split_size if axis.fixed_split_size > 0 else axis.length + sub_block_size = block_size + blocks.append(BlockInfo(block_name, sub_block_name, block_size, sub_block_size)) + + return blocks + + @classmethod + def get_key_from_dict(cls, kwargs: Dict[str, int]): + return tuple(sorted(kwargs.items())) + + def calcu_last_split_blocks(self, axis_idx): + splits = 1 + for x in self.kernel_meta.split_axis: + if x.index != axis_idx: + splits = splits * ( + (self.numels[x.index] + self.blocks[x.index].block_size - 1) // self.blocks[x.index].block_size) + else: + break + + last_splits = num_vector_core // splits + last_splits = max(1, last_splits) + last_blocks = (self.numels[axis_idx] + last_splits - 1) // last_splits + return last_blocks + + def aligned_numel(self, numel, align_bytes=32): + if self.is_simt_mode: + return next_power_of_2(numel) + + align_numel = align_bytes // self.dtype_bytes + if numel <= align_numel: + return numel + return ((numel + align_numel - 1) // align_numel) * align_numel + + def calculate_tile_numel(self): + tile_numel = 1 + for axis in self.kernel_meta.axis_info: + if axis.is_tiling_axis: + tile_numel *= self.blocks[axis.index].sub_block_size + else: + tile_numel *= self.blocks[axis.index].block_size + + return tile_numel + + def fill_config(self, cfg, candi_block): + for axis in self.kernel_meta.axis_info: + if not (axis.is_split_axis or axis.is_tiling_axis): + continue + block_info = self.blocks[axis.index] + if axis.is_split_axis: + curr_numel = candi_block[axis.index] + if not axis.is_tiling_axis: + curr_numel = self.aligned_numel(curr_numel) + if block_info.block_name: + cfg[block_info.block_name] = curr_numel + if axis.is_tiling_axis: + tiling_numel = self.aligned_numel(block_info.sub_block_size) + cfg[block_info.sub_block_name] = (tiling_numel if self.is_simt_mode else min( + tiling_numel, candi_block[axis.index])) + + def find_config(self, cfg): + for config_var in self.configs: + if config_var.kwargs == cfg: + return True + return False + + def _try_add_tiny_low_program_config(self, total_programs): + if (not self.tiny_kernel or total_programs < 1 or total_programs > self.tiny_program_threshold): + return + + if self.tiny_low_program_hist.get(total_programs, 0) >= self.tiny_per_program_cap: + return + + candi_block = tuple([x.block_size for x in self.blocks]) + if self.add_to_configs(list(candi_block)): + if candi_block not in self.candidate_blocks: + self.candidate_blocks.append(candi_block) + if not self.tiny_low_program_active: + self.tiny_low_program_active = True + self.tiny_low_program_tile_floor = self.calculate_tile_numel() + self.tiny_low_program_hist[total_programs] = (self.tiny_low_program_hist.get(total_programs, 0) + 1) + + def _calc_total_programs(self, candi_block=None): + grids = [] + for axis in self.kernel_meta.split_axis: + numel = self.numels[axis.index] + block_size = (self.blocks[axis.index].block_size if candi_block is None else candi_block[axis.index]) + programs = (numel + block_size - 1) // block_size + grids.append(programs) + + total_programs = functools.reduce(lambda x, y: x * y, grids) if grids else 1 + return total_programs + + def add_to_configs(self, candi_block): + newcfg = {} + self.fill_config(newcfg, candi_block) + tile_numel = self.calculate_tile_numel() + stop_numel_threshold = 0 if len(self.configs) < 10 or self.small_kernel else self.stop_numel + 100 + if self.tiny_low_program_active and self.tiny_low_program_tile_floor > 0: + total_programs = self._calc_total_programs(candi_block) + program_threshold = self.tiny_program_threshold if self.small_kernel else num_vector_core // 2 + if total_programs <= program_threshold: + tiny_low_program_threshold = max(self.stop_numel, self.tiny_low_program_tile_floor // 2) + stop_numel_threshold = max(stop_numel_threshold, tiny_low_program_threshold) + if (tile_numel <= self.max_numel_threshold and tile_numel >= stop_numel_threshold + and not self.find_config(newcfg)): + self.configs.append(Config(newcfg, num_warps=1, num_stages=1)) + return True + return False + + def desecnd_all_low_dims_with_all_blocks(self): + restore_sub_blocks = {} + for axis in self.kernel_meta.low_dims_axis: + restore_sub_blocks[axis.index] = self.blocks[axis.index].sub_block_size + self.descend_all_low_dims() + for axis in self.kernel_meta.low_dims_axis: + self.blocks[axis.index].sub_block_size = restore_sub_blocks[axis.index] + + def descend_one_axis(self, axis_idx: int, is_split=False): + + def calc_total_programs(): + grids = [] + for axis in self.kernel_meta.split_axis: + numel = self.numels[axis.index] + block_size = self.blocks[axis.index].block_size + programs = (numel + block_size - 1) // block_size + grids.append(programs) + + total_programs = functools.reduce(lambda x, y: x * y, grids) if grids else 1 + return total_programs + + reached_stop_numel = False + slow_decend_split = False + num_vector_core_tile = num_vector_core + max_programs_num = num_vector_core_tile if self.kernel_meta.tiling_axis else self.max_programs_num + if not is_split and len(self.candidate_blocks) == 0: + self.candidate_blocks.append(tuple([x.block_size for x in self.blocks])) + + axis = self.kernel_meta.axis_info[axis_idx] + while True: + for candi_block in self.candidate_blocks: + if self.add_to_configs(candi_block): + self.desecnd_all_low_dims_with_all_blocks() + + # tile numel reached threshold + tile_numel = self.calculate_tile_numel() + if tile_numel <= self.stop_numel: + if self.add_to_configs([x.block_size for x in self.blocks]): + self.desecnd_all_low_dims_with_all_blocks() + reached_stop_numel = True + break + + numel = (self.blocks[axis_idx].block_size if is_split else self.blocks[axis_idx].sub_block_size) + if numel == 1: + if self.add_to_configs([x.block_size for x in self.blocks]): + self.desecnd_all_low_dims_with_all_blocks() + break + + if is_split: + if self.persistent_reduction and axis.is_reduction: + reached_stop_numel = True + break + total_programs = calc_total_programs() + if total_programs > num_vector_core_tile: + if len(self.configs) == 0: + num_vector_core_tile = max_programs_num + slow_decend_split = (total_programs > num_vector_core_tile // 2) + if total_programs > num_vector_core_tile: + last_blocks = self.calcu_last_split_blocks(axis_idx) + if last_blocks != self.blocks[axis_idx].block_size: + self.blocks[axis_idx].block_size = last_blocks + self.candidate_blocks.append(tuple([x.block_size for x in self.blocks])) + break + + program_threshold = self.tiny_program_threshold if self.small_kernel else num_vector_core // 2 + if self.tiny_kernel and total_programs <= program_threshold: + self._try_add_tiny_low_program_config(total_programs) + if total_programs > program_threshold or self.dual_reduction: + if len(self.candidate_blocks) > 2: + self.candidate_blocks.pop(0) + self.candidate_blocks.append(tuple([x.block_size for x in self.blocks])) + if self.small_kernel: + self.add_to_configs(list(tuple([x.block_size for x in self.blocks]))) + slow_decend_split = (total_programs > num_vector_core_tile // 2) + + if not slow_decend_split: + self.blocks[axis_idx].block_size = (numel + 1) // 2 + else: + step = (numel + 3) // 4 if (numel + 3) // 4 > 1 else 1 + self.blocks[axis_idx].block_size = numel - step + self.blocks[axis_idx].sub_block_size = self.blocks[axis_idx].block_size + total_programs = calc_total_programs() + if self.blocks[axis_idx].block_size == 1 and (total_programs > program_threshold + or self.dual_reduction): + self.candidate_blocks.append(tuple([x.block_size for x in self.blocks])) + else: + if numel >= 32: + self.blocks[axis_idx].sub_block_size = next_power_of_2(numel // 2) + else: + self.blocks[axis_idx].sub_block_size = numel - 1 + return reached_stop_numel + + def descend_all_low_dims(self): + low_dim_numels = [self.blocks[x.index].sub_block_size for x in self.kernel_meta.low_dims_axis] + if not low_dim_numels: + return False + + def descend_all_axis(min_numel): + + for axis in self.kernel_meta.low_dims_axis: + if axis.is_reduction and self.persistent_reduction: + continue + + numel = self.blocks[axis.index].sub_block_size + if numel == 1: + continue + if min_numel > 1 and abs(numel - min_numel) / min_numel < 0.2: + continue + if numel >= 128: + self.blocks[axis.index].sub_block_size = next_power_of_2(numel // 2) + else: + numel = self.blocks[axis.index].sub_block_size + numel = numel // 2 + self.blocks[axis.index].sub_block_size = min(self.aligned_numel(numel), next_power_of_2(numel)) + + if len(self.candidate_blocks) == 0: + # means there is no split axis and tiling_not_low_dim axis + # so we need to init the candidates_blk_sizes + self.candidate_blocks.append(tuple([x.block_size for x in self.blocks])) + + count = 0 + tile_numel = self.calculate_tile_numel() + while tile_numel > self.stop_numel and count < 100: + count += 1 + tile_numel = self.calculate_tile_numel() + for candi_block in self.candidate_blocks: + self.add_to_configs(candi_block) + min_numel = min(low_dim_numels) + descend_all_axis(min_numel) + new_tile_numel = self.calculate_tile_numel() + if tile_numel == new_tile_numel: + descend_all_axis(0) + + return tile_numel < self.stop_numel + + def descend_split_tiling(self): + + tiling_not_low_dims = [x for x in self.kernel_meta.tiling_axis if x not in self.kernel_meta.low_dims_axis] + + def descend_split_axis(): + for axis in self.kernel_meta.tunable_split_axis: + if self.descend_one_axis(axis.index, is_split=True): + return True + + return self.calculate_tile_numel() <= self.stop_numel + + def descend_tiling_not_low_dims(): + for axis in tiling_not_low_dims: + if axis.is_reduction and self.persistent_reduction: + continue + + if self.descend_one_axis(axis.index): + return True + return self.calculate_tile_numel() <= self.stop_numel + + while True: + # descend split axis + if descend_split_axis(): + break + if len(self.candidate_blocks) > 0: + candi_block = self.candidate_blocks[0] + for i, blk_size in enumerate(candi_block): + self.blocks[i].sub_block_size = blk_size + # descend tiling but not low dims + if descend_tiling_not_low_dims(): + break + # descend low dims, need to descend all axis at the same time + self.descend_all_low_dims() + break diff --git a/third_party/ascend/backend/runtime/utils.py b/third_party/ascend/backend/runtime/utils.py new file mode 100644 index 0000000000..fadeeec98d --- /dev/null +++ b/third_party/ascend/backend/runtime/utils.py @@ -0,0 +1,132 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch + +_cached_params = None + + +def _init_npu_params(): + global _cached_params + if _cached_params is not None: + return _cached_params + + from triton.runtime.driver import driver + + target = driver.active.get_current_target() + device = driver.active.get_current_device() + prop = driver.active.utils.get_device_properties(device) + + num_cube_core = prop["num_aicore"] + num_vector_core = prop["num_aicore"] + ub_size_in_kbytes = 192 + rf_size_in_kbytes = None + + ASCEND_VARIANTS = ["Ascend910B", "Ascend910_93", "Ascend910_95", "Ascend950"] + if any(variant in target.arch for variant in ASCEND_VARIANTS): + num_vector_core = num_cube_core * 2 + + if target.arch.startswith("Ascend910_95") or target.arch.startswith("Ascend950"): + ub_size_in_kbytes = 256 + rf_size_in_kbytes = 128 + + _cached_params = { + 'target': target, + 'device': device, + 'prop': prop, + 'num_cube_core': num_cube_core, + 'num_vector_core': num_vector_core, + 'ub_size_in_kbytes': ub_size_in_kbytes, + 'rf_size_in_kbytes': rf_size_in_kbytes, + } + return _cached_params + + +def __getattr__(name): + if name in [ + 'target', 'device', 'prop', 'num_cube_core', 'num_vector_core', 'ub_size_in_kbytes', 'rf_size_in_kbytes' + ]: + return _init_npu_params()[name] + raise AttributeError(f"module '{__name__}' has no attribute '{name}'") + + +# wrapper npu 32 bytes align, get and pass unalign info to triton meta +# then autotune choose tiling param and send them to bishengIR +byte_per_numel = { + torch.float32: 4, # torch.float32 or torch.float + torch.float64: 8, # torch.float64 or torch.double + torch.float16: 2, # torch.float16 or torch.half + torch.bfloat16: 2, # torch.bfloat16 + torch.int32: 4, # torch.int32 or torch.int + torch.int64: 8, # torch.int64 or torch.long + torch.int16: 2, # torch.int16 or torch.short + torch.int8: 1, # torch.int8 + torch.uint8: 1, # torch.uint8 + torch.bool: 1, # torch.bool + torch.complex32: 4, # torch.complex32 (not yet available in PyTorch as of the latest stable release) + torch.complex64: 8, # torch.complex64 + torch.complex128: 16, # torch.complex128 +} + +# Some PyTorch versions expose extra fp8 dtypes. Register them when available. +for fp8_dtype_name in ( + "float8_e4m3fn", + "float8_e4m3fnuz", + "float8_e5m2", + "float8_e5m2fnuz", +): + fp8_dtype = getattr(torch, fp8_dtype_name, None) + if fp8_dtype is not None: + byte_per_numel[fp8_dtype] = 1 + +valid_axis_names = [ + "x", + "y", + "z", + "w", + "v", + "t", +] + + +def get_byte_per_numel(dtype: torch.dtype) -> int: + return 1 if dtype is None else byte_per_numel[dtype] + + +def is_valid_axis_name(name: str) -> bool: + if name.startswith("r"): + return name[1:] in valid_axis_names + return name in valid_axis_names + + +# move to an appropriate place, currently duplicated with triton.__init__.py +def next_power_of_2(n: int): + """Return the smallest power of 2 greater than or equal to n""" + n -= 1 + n |= n >> 1 + n |= n >> 2 + n |= n >> 4 + n |= n >> 8 + n |= n >> 16 + n |= n >> 32 + n += 1 + return n diff --git a/third_party/ascend/backend/spec/include/flagtree_spec.h b/third_party/ascend/backend/spec/include/flagtree_spec.h new file mode 100644 index 0000000000..7854bd49c0 --- /dev/null +++ b/third_party/ascend/backend/spec/include/flagtree_spec.h @@ -0,0 +1,7 @@ +#ifndef ASCEND_FLAGTREE_SPEC_H_ +#define ASCEND_FLAGTREE_SPEC_H_ + +#define FLAGTREE_SPEC_Dialect_Triton_IR_Ops_cpp +#define FLAGTREE_SPEC_Dialect_Triton_IR_Traits_cpp + +#endif // ASCEND_FLAGTREE_SPEC_H_ diff --git a/third_party/ascend/backend/spec/include/runtime/libentry/libentry.h b/third_party/ascend/backend/spec/include/runtime/libentry/libentry.h new file mode 100644 index 0000000000..2a9ba371bd --- /dev/null +++ b/third_party/ascend/backend/spec/include/runtime/libentry/libentry.h @@ -0,0 +1,77 @@ +/* + * Copyright 2018-2020 Philippe Tillet + * Copyright 2020-2022 OpenAI + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef LIBENTRY_H +#define LIBENTRY_H + +#include +#include +#include +#include + +#include +#include + +namespace py = pybind11; + +using KeyType = py::tuple; + +namespace libentry { + +class ArgProcessor { +public: + ArgProcessor(int div) : divisibility_(div) {}; + + void + classifyArguments(const py::list &args, const py::dict &kwargs, + const py::list &jit_params, + const std::unordered_set &specialize_indices, + const std::unordered_set &do_not_specialize_indices); + + KeyType generateKey(); + + py::list getKArgs(); + +private: + py::list spec_args_; // specialize args + py::list dns_args_; // do not specialize args + py::list const_args_; // constexpr args + py::list k_args_; // kernel args + int divisibility_; // 对齐 +}; + +} // namespace libentry + +PYBIND11_MODULE(libentryC, m) { + py::class_(m, "ArgProcessor") + .def(py::init()) + .def("classify_arguments", &libentry::ArgProcessor::classifyArguments, + py::arg("args"), py::arg("kwargs"), py::arg("jit_params"), + py::arg("specialize_indices"), py::arg("do_not_specialize_indices"), + "classify arguments") + .def("get_k_args", &libentry::ArgProcessor::getKArgs, "get kernel") + .def("generate_key", &libentry::ArgProcessor::generateKey, + "generate kernel cache key"); +} + +#endif diff --git a/third_party/ascend/backend/spec/include/triton/Dialect/Triton/IR/TritonAttrDefs.td b/third_party/ascend/backend/spec/include/triton/Dialect/Triton/IR/TritonAttrDefs.td new file mode 100644 index 0000000000..51a6008f20 --- /dev/null +++ b/third_party/ascend/backend/spec/include/triton/Dialect/Triton/IR/TritonAttrDefs.td @@ -0,0 +1,180 @@ +/* + * Copyright 2018-2020 Philippe Tillet + * Copyright 2020-2022 OpenAI + * Copyright 2025- FlagOS Contributors + * + * Permission is hereby granted, free of charge, to any person obtaining + * a copy of this software and associated documentation files + * (the "Software"), to deal in the Software without restriction, + * including without limitation the rights to use, copy, modify, merge, + * publish, distribute, sublicense, and/or sell copies of the Software, + * and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be + * included in all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, + * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF + * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. + * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY + * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, + * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE + * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ + +#ifndef TRITON_ATTR_DEFS +#define TRITON_ATTR_DEFS + +include "mlir/IR/EnumAttr.td" + +// Attributes for LoadOp and StoreOp +def TT_CacheModifierAttr : I32EnumAttr< + "CacheModifier", "", + [ + I32EnumAttrCase<"NONE", 1, "none">, + I32EnumAttrCase<"CA", 2, "ca">, + I32EnumAttrCase<"CG", 3, "cg">, + I32EnumAttrCase<"WB", 4, "wb">, + I32EnumAttrCase<"CS", 5, "cs">, + I32EnumAttrCase<"WT", 6, "wt">, + I32EnumAttrCase<"CV", 7, "cv">, + ]> { + let cppNamespace = "::mlir::triton"; +} + +def TT_MemSemanticAttr : I32EnumAttr< + "MemSemantic", "", + [ + I32EnumAttrCase<"RELAXED", 1, "relaxed">, + I32EnumAttrCase<"ACQUIRE", 2, "acquire">, + I32EnumAttrCase<"RELEASE", 3, "release">, + I32EnumAttrCase<"ACQUIRE_RELEASE", 4, "acq_rel">, + ]> { + let cppNamespace = "::mlir::triton"; +} + +def TT_EvictionPolicyAttr : I32EnumAttr< + "EvictionPolicy", "", + [ + I32EnumAttrCase<"NORMAL", 1, "evict_normal">, + I32EnumAttrCase<"EVICT_FIRST", 2, "evict_first">, + I32EnumAttrCase<"EVICT_LAST", 3, "evict_last"> + ]> { + let cppNamespace = "::mlir::triton"; +} + +def TT_PaddingOptionAttr : I32EnumAttr< + "PaddingOption", "", + [ + I32EnumAttrCase<"PAD_ZERO", 1, "zero">, + // We can not set the string value to "NAN" because it is a keyword in C++ + I32EnumAttrCase<"PAD_NAN", 2, "nan"> + ]> { + let cppNamespace = "::mlir::triton"; +} + +// atomic +def TT_AtomicRMWAttr : I32EnumAttr< + "RMWOp", "", + [ + I32EnumAttrCase<"AND", 1, "and">, + I32EnumAttrCase<"OR", 2, "or">, + I32EnumAttrCase<"XOR", 3, "xor">, + I32EnumAttrCase<"ADD", 4, "add">, + I32EnumAttrCase<"FADD", 5, "fadd">, + I32EnumAttrCase<"MAX", 6, "max">, + I32EnumAttrCase<"MIN", 7, "min">, + I32EnumAttrCase<"UMAX", 8, "umax">, + I32EnumAttrCase<"UMIN", 9, "umin">, + I32EnumAttrCase<"XCHG", 10, "exch"> + ]> { + let cppNamespace = "::mlir::triton"; +} + +def TT_DescriptorReduceKindAttr : I32EnumAttr< + "DescriptorReduceKind", "", + [ + I32EnumAttrCase<"ADD", 1, "add">, + I32EnumAttrCase<"MIN", 2, "min">, + I32EnumAttrCase<"MAX", 3, "max">, + I32EnumAttrCase<"INC", 4, "inc">, + I32EnumAttrCase<"DEC", 5, "dec">, + I32EnumAttrCase<"AND", 6, "and">, + I32EnumAttrCase<"OR", 7, "or">, + I32EnumAttrCase<"XOR", 8, "xor">, + ]> { + let cppNamespace = "::mlir::triton"; +} + +def TT_MemSyncScopeAttr : I32EnumAttr< + "MemSyncScope", "", + [ + I32EnumAttrCase<"GPU", 1, "gpu">, + I32EnumAttrCase<"CTA", 2, "cta">, + I32EnumAttrCase<"SYSTEM", 3, "sys">, + ]> { + let cppNamespace = "::mlir::triton"; +} + +// Program ID dimensions. +def TT_ProgramDim : I32EnumAttr< + "ProgramIDDim", "", + [ + I32EnumAttrCase<"X", 0, "x">, + I32EnumAttrCase<"Y", 1, "y">, + I32EnumAttrCase<"Z", 2, "z">, + ]> { + let cppNamespace = "::mlir::triton"; +} + +// Rounding mode. +def TT_RoundingModeAttr : I32EnumAttr< + "RoundingMode", "", + [ + I32EnumAttrCase<"RTZ", 0, "rtz">, + I32EnumAttrCase<"RTNE", 1, "rtne">, + ]> { + let cppNamespace = "::mlir::triton"; +} + +// PropagateNan. +def TT_PropagateNanAttr : I32EnumAttr< + "PropagateNan", "", + [ + I32EnumAttrCase<"NONE", 0, "none">, + I32EnumAttrCase<"ALL", 0xFFFF, "all">, + ]> { + let cppNamespace = "::mlir::triton"; +} + +// InputPrecision +def TT_InputPrecisionAttr : I32EnumAttr< + "InputPrecision", "", + [ + I32EnumAttrCase<"TF32", 0, "tf32">, + I32EnumAttrCase<"TF32x3", 1, "tf32x3">, + I32EnumAttrCase<"IEEE", 2, "ieee">, + I32EnumAttrCase<"HF32", 3, "hf32">, + I32EnumAttrCase<"BF16x3", 4, "bf16x3">, + I32EnumAttrCase<"BF16x6", 5, "bf16x6"> + ]>{ + let cppNamespace = "::mlir::triton"; +} + +// Type for ScaleDotElemType kind of floats. +def TT_ScaleDotElemTypeAttr : I32EnumAttr< + "ScaleDotElemType", "", + [ + I32EnumAttrCase<"E4M3", 0, "e4m3">, + I32EnumAttrCase<"E5M2", 1, "e5m2">, + I32EnumAttrCase<"E2M3", 2, "e2m3">, + I32EnumAttrCase<"E3M2", 3, "e3m2">, + I32EnumAttrCase<"E2M1", 4, "e2m1">, + I32EnumAttrCase<"BF16", 5, "bf16">, + I32EnumAttrCase<"FP16", 6, "fp16"> + ]>{ + let cppNamespace = "::mlir::triton"; +} + +#endif diff --git a/third_party/ascend/backend/spec/lib/CMakeLists.txt b/third_party/ascend/backend/spec/lib/CMakeLists.txt new file mode 100644 index 0000000000..4e0b706748 --- /dev/null +++ b/third_party/ascend/backend/spec/lib/CMakeLists.txt @@ -0,0 +1,6 @@ +# The Dialect/ subtree held a v3.5-style spec library (FlagTree_ascend_TritonIR) +# that replaced core Triton sources. On this branch the core files carry no +# FLAGTREE_SPEC_* guards, so such a library can never suppress them; source +# specialization goes through third_party/ascend/spec_cpp instead, which +# FlagTreeBackendSpec.cmake substitutes into the owning core target. +add_subdirectory(runtime) diff --git a/third_party/ascend/backend/spec/lib/Dialect/CMakeLists.txt b/third_party/ascend/backend/spec/lib/Dialect/CMakeLists.txt new file mode 100644 index 0000000000..5e601271ec --- /dev/null +++ b/third_party/ascend/backend/spec/lib/Dialect/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(Triton) diff --git a/third_party/ascend/backend/spec/lib/Dialect/Triton/CMakeLists.txt b/third_party/ascend/backend/spec/lib/Dialect/Triton/CMakeLists.txt new file mode 100644 index 0000000000..f33061b2d8 --- /dev/null +++ b/third_party/ascend/backend/spec/lib/Dialect/Triton/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(IR) diff --git a/third_party/ascend/backend/spec/lib/Dialect/Triton/IR/CMakeLists.txt b/third_party/ascend/backend/spec/lib/Dialect/Triton/IR/CMakeLists.txt new file mode 100644 index 0000000000..2c9abe302b --- /dev/null +++ b/third_party/ascend/backend/spec/lib/Dialect/Triton/IR/CMakeLists.txt @@ -0,0 +1,22 @@ +# Only Traits.cpp is specialized: it relaxes the power-of-two tensor size check +# (same change triton-ascend applies to the core file). The Ops.cpp and +# TritonOps.td overrides were 3.5-era snapshots whose deltas are all superseded +# by the 3.6 core, so they were dropped and the core versions are used. +add_triton_library(FlagTree_ascend_TritonIR + Traits.cpp + + DEPENDS + TritonTableGen + TritonCanonicalizeIncGen + + LINK_LIBS PUBLIC + MLIRIR + MLIRArithDialect + MLIRMathDialect + MLIRSCFDialect +) + +target_include_directories(FlagTree_ascend_TritonIR PRIVATE + ${CMAKE_CURRENT_BINARY_DIR} + ${CMAKE_BINARY_DIR}/lib/Dialect/Triton/IR +) diff --git a/third_party/ascend/backend/spec/lib/Dialect/Triton/IR/Traits.cpp b/third_party/ascend/backend/spec/lib/Dialect/Triton/IR/Traits.cpp new file mode 100644 index 0000000000..81d4ffcdbc --- /dev/null +++ b/third_party/ascend/backend/spec/lib/Dialect/Triton/IR/Traits.cpp @@ -0,0 +1,276 @@ +/* + * Copyright 2018-2020 Philippe Tillet + * Copyright 2020-2022 OpenAI + * Copyright 2025- FlagOS Contributors + * + * Permission is hereby granted, free of charge, to any person obtaining + * a copy of this software and associated documentation files + * (the "Software"), to deal in the Software without restriction, + * including without limitation the rights to use, copy, modify, merge, + * publish, distribute, sublicense, and/or sell copies of the Software, + * and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be + * included in all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, + * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF + * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. + * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY + * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, + * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE + * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ + +#include "triton/Dialect/Triton/IR/Traits.h" + +#include + +#include "mlir/IR/TypeUtilities.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Types.h" +#include "triton/Dialect/Triton/IR/Utility.h" +#include "triton/Dialect/TritonGPU/IR/Types.h" +#include "llvm/Support/ErrorHandling.h" + +using namespace mlir; +using namespace mlir::triton::gpu; + +LogicalResult OpTrait::impl::verifyEquivalentType(Type typeA, Type typeB) { + auto memdescA = dyn_cast(typeA); + auto memdescB = dyn_cast(typeB); + if (memdescA || memdescB) { + if (!memdescA || !memdescB) + return failure(); + if (memdescA.getShape() != memdescB.getShape()) + return failure(); + if (memdescA.getAllocShape() != memdescB.getAllocShape()) + return failure(); + if (memdescA.getElementType() != memdescB.getElementType()) + return failure(); + if (memdescA.getMemorySpace() != memdescB.getMemorySpace()) + return failure(); + if (memdescA.getMutableMemory() != memdescB.getMutableMemory()) + return failure(); + + Attribute encodingA = memdescA.getEncoding(); + Attribute encodingB = memdescB.getEncoding(); + if (encodingA == encodingB) + return success(); + if (static_cast(encodingA) != static_cast(encodingB)) + return failure(); + + auto layoutInterface = + cast(&encodingA.getDialect()); + return layoutInterface->verifyLayoutsAreEqual(memdescA.getShape(), + encodingA, encodingB, {}); + } + auto tensorTypeA = dyn_cast(typeA); + auto tensorTypeB = dyn_cast(typeB); + if (!(bool(tensorTypeA) && bool(tensorTypeB))) + return typeA == typeB ? success() : failure(); + auto encodingA = tensorTypeA.getEncoding(); + auto encodingB = tensorTypeB.getEncoding(); + auto shapeA = tensorTypeA.getShape(); + auto shapeB = tensorTypeB.getShape(); + if (shapeA != shapeB) + return failure(); + if (tensorTypeA.getElementType() != tensorTypeB.getElementType()) + return failure(); + // If there's no encoding or the encodings are the same + if (encodingA == encodingB) + return success(); + if (bool(encodingA) != bool(encodingB)) + return failure(); + + return cast(&encodingA.getDialect()) + ->verifyLayoutsAreEqual(shapeA, encodingA, encodingB, {}); +} + +static LogicalResult verifySameEncoding(Type typeA, Type typeB, + bool allowTensorPointerType) { + // TODO(Keren): the allowTensorPointerType argument is a hack to allow. + // The type checking code is kind of a mess with the current design. + auto getEncoding = [=](Type type) -> Attribute { + Attribute ret; + if (auto tensorType = dyn_cast(type)) { + ret = tensorType.getEncoding(); + } + if (!allowTensorPointerType) { + assert(!triton::isTensorPointerType(type)); + } + return ret; + }; + auto encodingA = getEncoding(typeA); + auto encodingB = getEncoding(typeB); + if (!encodingA || !encodingB) + return success(); + return encodingA == encodingB ? success() : failure(); +} + +LogicalResult +OpTrait::impl::verifySameOperandsEncoding(Operation *op, + bool allowTensorPointerType) { + if (failed(verifyAtLeastNOperands(op, 1))) + return failure(); + + auto type = op->getOperand(0).getType(); + for (auto opType : llvm::drop_begin(op->getOperandTypes(), 1)) + if (failed(verifySameEncoding(opType, type, allowTensorPointerType))) + return op->emitOpError() << "requires the same encoding for all operands"; + + return success(); +} + +LogicalResult OpTrait::impl::verifySameOperandsAndResultEncoding( + Operation *op, bool allowTensorPointerType) { + if (op->getNumOperands() == 0) + return success(); + + if (failed(verifyAtLeastNOperands(op, 1)) || + failed(verifyAtLeastNResults(op, 1))) + return failure(); + + auto type = op->getOperand(0).getType(); + for (auto resultType : op->getResultTypes()) + if (failed(verifySameEncoding(resultType, type, allowTensorPointerType))) + return op->emitOpError() + << "requires the same encoding for all operands and results"; + + return verifySameOperandsEncoding(op, allowTensorPointerType); +} + +LogicalResult OpTrait::impl::verifyTensorSize(Operation *op) { + for (auto opType : op->getOperandTypes()) { + if (auto tensorType = dyn_cast(opType)) { + int64_t numElements = 1; + for (int64_t s : tensorType.getShape()) + numElements *= s; + if (numElements > maxTensorNumElements) + return op->emitError("Maximum allowed number of elements is ") + << maxTensorNumElements << ", but " << *op + << " has more than that"; + // FIXME:patched triton community + // if ((numElements & (numElements - 1)) != 0) + // return op->emitError("Number of elements must be power-of-two, but ") + // << *op << " doesn't follow the rule (" << numElements << ")" + // << " elements"; + } + } + for (auto opType : op->getResultTypes()) { + if (auto tensorType = dyn_cast(opType)) { + int64_t numElements = 1; + for (int64_t s : tensorType.getShape()) + numElements *= s; + if (numElements > maxTensorNumElements) + return op->emitError("Maximum allowed number of elements is ") + << maxTensorNumElements << ", but " << *op + << " has more than that"; + // FIXME:patched triton community + // if ((numElements & (numElements - 1)) != 0) + // return op->emitError("Number of elements must be power-of-two, but ") + // << *op << " doesn't follow the rule (" << numElements << ")" + // << " elements"; + } + } + return success(); +} + +// Check that the Triton layouts on op's operands and return types are valid. +// For example, we check that the number of warps per block in a Triton GPU +// blocked layout matches that of its module. +// +// It's a little weird to check these properties of a layout only when the +// layout is used in an op, since most of the properties don't actually depend +// on the op. They do depend on the *module*, though, and a layout is attached +// to a module only by virtue of being used in one of the module's ops. +LogicalResult OpTrait::impl::verifyTensorLayouts(Operation *op) { + auto checkLayout = [&](Value val, auto makeErr) -> LogicalResult { + // Only ranked tensors can have layouts. + auto rankedTy = dyn_cast(val.getType()); + if (!rankedTy) + return success(); + + mlir::Attribute layout = rankedTy.getEncoding(); + if (!layout) + return success(); + + Dialect &dialect = layout.getDialect(); + auto verifyLayoutInterface = + dyn_cast(&dialect); + if (verifyLayoutInterface) { + return verifyLayoutInterface->verifyTensorLayout(layout, rankedTy, op, + makeErr); + } + + return success(); + }; + + for (size_t i = 0; i < op->getNumOperands(); i++) { + auto operand = op->getOperand(i); + auto err = checkLayout(operand, [&]() { + // Stringify the operand using `printAsOperand`. This prints e.g. "%42" + // rather than the full definition. + std::string operandStr; + llvm::raw_string_ostream os(operandStr); + // If we don't assume verified, dump() will recursively call this + // function! + operand.printAsOperand(os, OpPrintingFlags().assumeVerified()); + + return op->emitError("Operand ") + << i << " (" << operand << ") has an invalid layout: "; + }); + if (!err.succeeded()) + return err; + } + + for (size_t i = 0; i < op->getNumResults(); i++) { + auto result = op->getResult(i); + auto err = checkLayout(result, [&]() { + if (op->getNumResults() == 1) { + return op->emitError("Result has an invalid layout: "); + } else { + return op->emitError("Result ") << i << " has an invalid layout: "; + } + }); + if (!err.succeeded()) + return err; + } + + return success(); +} + +static ArrayRef getTypeShape(Type type) { + auto rankedType = dyn_cast(type); + if (auto ptrType = dyn_cast(type)) + rankedType = dyn_cast(ptrType.getPointeeType()); + return rankedType ? rankedType.getShape() : ArrayRef(); +} + +LogicalResult OpTrait::impl::verifySameLoadStoreOperandsShape(Operation *op) { + if (failed(verifyAtLeastNOperands(op, 1))) + return failure(); + + auto firstOperandShape = getTypeShape(op->getOperand(0).getType()); + for (auto type : llvm::drop_begin(op->getOperandTypes(), 1)) + if (failed(verifyCompatibleShape(getTypeShape(type), firstOperandShape))) + return op->emitOpError() << "requires the same shape for all operands"; + + return success(); +} + +LogicalResult +OpTrait::impl::verifySameLoadStoreOperandsAndResultShape(Operation *op) { + if (failed(verifyAtLeastNOperands(op, 1)) || + failed(verifyAtLeastNResults(op, 1))) + return failure(); + + auto firstOperandShape = getTypeShape(op->getOperand(0).getType()); + for (auto type : op->getResultTypes()) + if (failed(verifyCompatibleShape(getTypeShape(type), firstOperandShape))) + return op->emitOpError() + << "requires the same shape for all operands and results"; + + return verifySameLoadStoreOperandsShape(op); +} diff --git a/third_party/ascend/backend/spec/lib/runtime/CMakeLists.txt b/third_party/ascend/backend/spec/lib/runtime/CMakeLists.txt new file mode 100644 index 0000000000..4bfdd0d51d --- /dev/null +++ b/third_party/ascend/backend/spec/lib/runtime/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(libentry) diff --git a/third_party/ascend/backend/spec/lib/runtime/libentry/CMakeLists.txt b/third_party/ascend/backend/spec/lib/runtime/libentry/CMakeLists.txt new file mode 100644 index 0000000000..cebf7227a5 --- /dev/null +++ b/third_party/ascend/backend/spec/lib/runtime/libentry/CMakeLists.txt @@ -0,0 +1,11 @@ +add_library(entryC SHARED + libentry.cpp +) +target_compile_options(entryC PRIVATE -fexceptions -frtti) +# libentry.cpp includes "runtime/libentry/libentry.h", which lives in this +# backend's spec include tree. Scope the path to this target: the global +# BACKEND_SPEC_INCLUDE_DIR points at spec_cpp/include, which is the source +# override tree consumed by FlagTreeBackendSpec.cmake. +target_include_directories(entryC PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/../../../include +) diff --git a/third_party/ascend/backend/spec/lib/runtime/libentry/libentry.cpp b/third_party/ascend/backend/spec/lib/runtime/libentry/libentry.cpp new file mode 100644 index 0000000000..1e2a9071e9 --- /dev/null +++ b/third_party/ascend/backend/spec/lib/runtime/libentry/libentry.cpp @@ -0,0 +1,115 @@ +/* + * Copyright 2018-2020 Philippe Tillet + * Copyright 2020-2022 OpenAI + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "runtime/libentry/libentry.h" + +using namespace libentry; + +void libentry::ArgProcessor::classifyArguments( + const py::list &args, const py::dict &kwargs, const py::list &jit_params, + const std::unordered_set &specialize_indices, + const std::unordered_set &do_not_specialize_indices) { + for (size_t i = 0; i < args.size(); ++i) { + if (specialize_indices.count(i)) { + k_args_.append(args[i]); + spec_args_.append(args[i]); + } else if (do_not_specialize_indices.count(i)) { + k_args_.append(args[i]); + dns_args_.append(args[i]); + } else { + const_args_.append(args[i]); + } + } + + for (size_t i = args.size(); i < jit_params.size(); ++i) { + const py::object ¶m = jit_params[i]; + py::object val; + + if (kwargs.contains(param.attr("name"))) { + val = kwargs[param.attr("name")]; + } else if (py::hasattr(param, "default") && + !param.attr("default").is_none()) { + val = param.attr("default"); + } else { + continue; + } + + if (param.attr("is_constexpr").cast()) { + const_args_.append(val); + } else if (param.attr("do_not_specialize").cast()) { + dns_args_.append(val); + k_args_.append(val); + } else { + spec_args_.append(val); + k_args_.append(val); + } + } +} + +KeyType libentry::ArgProcessor::generateKey() { + auto is_tensor = [](py::handle x) { return py::hasattr(x, "data_ptr"); }; + auto is_int = [](py::handle x) { return py::isinstance(x); }; + + py::list spec_key; + for (auto arg : spec_args_) { + if (is_tensor(arg)) { + auto dtype = arg.attr("dtype"); + uintptr_t data_ptr = arg.attr("data_ptr")().cast(); + bool aligned = (data_ptr & (divisibility_ - 1)) == 0; + spec_key.append(py::make_tuple(dtype, aligned)); + } else { + spec_key.append(py::make_tuple(py::type::of(arg), arg)); + } + } + + py::list dns_key; + for (auto arg : dns_args_) { + if (is_tensor(arg)) { + dns_key.append(arg.attr("dtype")); + } else if (!is_int(arg)) { + dns_key.append(py::type::of(arg)); + } else { + int64_t val = arg.cast(); + if (val >= -0x80000000LL && val <= 0x7FFFFFFFLL) { + dns_key.append(py::str("i32")); + } else if (val >= 0 && val <= 0xFFFFFFFFFFFFFFFFLL) { + dns_key.append(py::str("u64")); + } else { + dns_key.append(py::str("i64")); + } + } + } + + py::list result; + auto list_append = [&](const py::list &src) { + for (auto handle : src) { + result.append(handle); + } + }; + list_append(spec_key); + list_append(dns_key); + list_append(const_args_); + return result; +} + +py::list libentry::ArgProcessor::getKArgs() { return k_args_; } diff --git a/third_party/ascend/backend/testing.py b/third_party/ascend/backend/testing.py new file mode 100644 index 0000000000..04089750ce --- /dev/null +++ b/third_party/ascend/backend/testing.py @@ -0,0 +1,399 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import builtins +import multiprocessing +import os +from datetime import datetime, timezone +from typing import Optional +import fnmatch + +import triton.runtime as runtime +from triton.knobs import cache + + +class ProfilerResultMismatchError(RuntimeError): + + def __init__(self, target_kernel_name: str, expected_rows: int, actual_rows: int): + self.target_kernel_name = target_kernel_name + self.expected_rows = expected_rows + self.actual_rows = actual_rows + super().__init__( + "Profiler rows filtered by target kernel name do not match the expected count. " + f"target_kernel_name={target_kernel_name!r}, expected_rows={expected_rows}, actual_rows={actual_rows}") + + +def do_bench_npu_profiler( + funcs, + warmup=5, + active=30, + clear_l2_cache=False, + prof_dir=None, + keep_res=False, + target_kernel_name: Optional[str] = None, + _raise_on_mismatch: bool = False, +): + import torch + import torch_npu + if not isinstance(funcs, list): + funcs = [funcs] + + # warmup kernel + for fn in funcs: + fn() + torch.npu.synchronize() + + experimental_config = torch_npu.profiler._ExperimentalConfig( + aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, + l2_cache=False, + data_simplification=False, + ) + + if prof_dir is not None: + torch_path = prof_dir + else: + process = multiprocessing.current_process() + pid = process.pid + process_name = process.name + timestamp = datetime.now(tz=timezone.utc).strftime("%Y%m%d_%H%M%S") + base_path = cache.get_triton_dir("profile_results") + torch_path = os.path.join(base_path, f"prof_{timestamp}_{process_name}-{pid}") + + if clear_l2_cache: + buffer = runtime.driver.active.get_empty_cache_for_benchmark() + buffer = buffer.float() # to avoid type cast + buffer.sum() + torch.npu.synchronize() # shake out of any npu error + + total = warmup + active + with torch_npu.profiler.profile( + activities=[torch_npu.profiler.ProfilerActivity.NPU], + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(torch_path), + record_shapes=False, + profile_memory=False, + with_stack=False, + with_flops=False, + with_modules=False, + experimental_config=experimental_config, + ) as prof: + for fn in funcs: + for _ in builtins.range(total): + if clear_l2_cache: + buffer.sum() # use buffer read to clear l2 cache + torch.npu.synchronize() + fn() + torch.npu.synchronize() + if clear_l2_cache: + del buffer + + try: + return _collect_prof_result( + torch_path, + funcs, + warmup, + active, + target_kernel_name=target_kernel_name, + clear_l2_cache=clear_l2_cache, + _raise_on_mismatch=_raise_on_mismatch, + ) + finally: + _rm_dic(keep_res, torch_path) + + +def _rm_dic(keep_res, torch_path): + if keep_res: + return + import shutil + + if os.path.exists(torch_path): + shutil.rmtree(torch_path) + + +def _collect_prof_result( + base_dir: str, + funcs, + num_warmup: int, + num_active: int, + target_kernel_name: Optional[str] = None, + clear_l2_cache: bool = False, + _raise_on_mismatch: bool = False, +): + """ + Collect kernel performance from task_time*.csv or kernel_details.csv, returned in millisecond. + Uses task_time*.csv by default. If target_kernel_name is provided, scans for kernel_details.csv as a fallback for accuracy. + The first `num_warmup` rows of each function are warmup data and will be ignored, the next `num_active` rows will be averaged. + + :param base_dir: the profiler path + :type base_dir: str + :param funcs: a list of Callable being profiled + :type funcs: List[Callable] + :param num_warmup: warmup count in task_time*.csv or kernel_details.csv of each fn + :type num_warmup: int + :param num_active: active count in task_time*.csv or kernel_details.csv of each fn + :type num_active: int + :param target_kernel_name: target triton kernel name reported by profiler + :type target_kernel_name: Optional[str] + """ + + import numpy as np + import pandas as pd + use_task_time = (target_kernel_name is None) + kernel_details_file = None + for root, _, files in os.walk(base_dir): + for file in files: + if use_task_time and fnmatch.fnmatch(file, "task_time*.csv"): + kernel_details_file = os.path.join(root, file) + break + elif not use_task_time and file == "kernel_details.csv": + kernel_details_file = os.path.join(root, file) + break + num_funcs = len(funcs) + + def _error(msg: str): + print(f"[Error] {msg}") + if num_funcs == 1: + return float("inf") + return [float("inf")] * num_funcs + + if kernel_details_file is None: + return _error( + f"No profiling data found under {base_dir}. The profiler may have failed to collect device tasks or stopped abnormally." + ) + + df = pd.read_csv(kernel_details_file, keep_default_na=False) + if use_task_time: + # The first and last lines of the task_time*.csv file are PROFILING_DISABLE, which should be deleted. + df = df[1:-1] + + if df.empty: + print("[WARNING] No profiling data on the device side.") + if num_funcs == 1: + return float("0") + return [float("0")] * num_funcs + + if use_task_time: + col_time = "task_time(us)" + filter_cond = (not clear_l2_cache) | ~df["kernel_name"].str.contains(r"^ReduceSum", case=False, na=False) + else: + col_time = "Duration(us)" + filter_cond = (not clear_l2_cache) | ~df["Type"].str.contains(r"^ReduceSum$", case=False, na=False) + # filter out l2 cache clearing operation + filter_df = df[filter_cond] + if target_kernel_name is not None: + filter_df = filter_df[filter_df["Name"] == target_kernel_name] + + expected_rows = num_funcs * (num_warmup + num_active) + actual_rows = len(filter_df) + + mul = 1 + if num_funcs == 1: + if actual_rows % expected_rows != 0: + if target_kernel_name is not None and _raise_on_mismatch: + raise ProfilerResultMismatchError(target_kernel_name, expected_rows, actual_rows) + return _error( + f"Expected the actual row count to be a multiple of {expected_rows}, but got {actual_rows}. The expected row count and the actual row count do not match." + ) + mul = actual_rows // expected_rows + num_warmup = num_warmup * mul + num_active = num_active * mul + else: + if actual_rows != expected_rows: + print( + "[WARNING] Passing a list of functions where a function may contain multiple kernels or launch no kernel is not supported and may lead to incorrect results." + ) + + time_cost = [0] * num_funcs + for func_idx in np.arange(0, num_funcs): + for active_index in np.arange(0, num_active): + row_index = func_idx * (num_warmup + num_active) + num_warmup + active_index + time_cost[func_idx] += filter_df.iloc[row_index][col_time] + time_cost = [x * mul / num_active / 1e3 for x in time_cost] + + if num_funcs == 1: + return time_cost[0] + else: + return time_cost + + +try: + from mspti import KernelMonitor +except ImportError: + KernelMonitor = None + +try: + from triton.backends.ascend.utils import is_cann_version_at_least + CANN_VERSION_AVAILABLE = True +except ImportError: + CANN_VERSION_AVAILABLE = False + print("[WARNING] triton.backends.ascend.utils not found. CANN version check skipped.") + + +# If the CANN version is earlier than 9.1.0, it needs to set libmspti.so in LD_PRELOAD to use mspti. +def do_bench_npu_mspti( + funcs, + warmup=5, + active=30, + clear_l2_cache=False, + target_kernel_name: Optional[str] = None, +): + import torch + import torch_npu + if not isinstance(funcs, list): + funcs = [funcs] + + for fn in funcs: + fn() + torch.npu.synchronize() + + if clear_l2_cache: + buffer = runtime.driver.active.get_empty_cache_for_benchmark() + else: + buffer = None + + all_kernel_durations = [] + + def callback(data): + if clear_l2_cache and ('zero' in data.name.lower() or 'zeroslike' in data.name.lower()): + return + if target_kernel_name is not None and target_kernel_name not in data.name: + return + all_kernel_durations.append(data.end - data.start) + + monitor = KernelMonitor() + torch.npu.synchronize() + + monitor.start(callback) + + try: + total = warmup + active + for fn in funcs: + for _ in builtins.range(total): + if clear_l2_cache: + buffer.zero_() + fn() + finally: + torch.npu.synchronize() + monitor.stop() + + num_funcs = len(funcs) + duration_per_kernel = [] + + expected_rows = num_funcs * total + actual_rows = len(all_kernel_durations) + if actual_rows < expected_rows: + if num_funcs == 1: + return float("inf") + return [float("inf")] * num_funcs + + mul = 1 + if num_funcs == 1: + if actual_rows % expected_rows != 0: + return float("inf") + mul = actual_rows // expected_rows + warmup = warmup * mul + total = actual_rows + else: + if actual_rows != expected_rows: + print( + "[WARNING] Passing a list of functions containing multiple kernels is not fully supported and may lead to inaccurate results." + ) + + current_idx = 0 + for i in range(num_funcs): + current_func_records = all_kernel_durations[current_idx:current_idx + total] + current_idx += total + current_active_records = current_func_records[warmup:total] + avg_time = sum(current_active_records) * mul / len(current_active_records) + avg_time_ms = avg_time / 1000000.0 + duration_per_kernel.append(avg_time_ms) + + if num_funcs == 1: + return duration_per_kernel[0] + else: + return duration_per_kernel + + +def _check_mspti_env(): + if CANN_VERSION_AVAILABLE: + if not is_cann_version_at_least(9, 1, 0): + if 'libmspti.so' not in os.getenv('LD_PRELOAD', ''): + print("[WARNING] libmspti.so not set in LD_PRELOAD, please set libmspti.so in LD_PRELOAD to use mspti.") + + +def do_bench_npu( + funcs, + warmup=5, + active=30, + clear_l2_cache=False, + prof_dir=None, + keep_res=False, + target_kernel_name: Optional[str] = None, +): + """ + Benchmark the runtime of the provided function on NPU. + this function utilizes NPU profiling tools (mspti or torch_npu.profiler) to capture pure Device-side kernel execution time. + By default, it returns the mean runtime in milliseconds of the provided function(s) based on `active` iterations. + + :param funcs: Function (or list of functions) to benchmark. If a list is provided, returns a list of mean runtimes. + :type funcs: Callable or List[Callable] + :param warmup: Warmup iterations. Runs the function `warmup` times before actual timing to stabilize performance. Defaults to 5. + :type warmup: int + :param active: Active iterations to record for timing. The final result is the average over these iterations. Defaults to 30. + :type active: int + :param clear_l2_cache: Whether to clear the L2 cache before each function execution. Defaults to False. + :type clear_l2_cache: bool, optional + :param prof_dir: Directory to save profiler results. If None, defaults to a temporary directory under triton cache. If specified, forces fallback to torch_npu.profiler. + :type prof_dir: str, optional + :param keep_res: Whether to keep the raw profiler result files (e.g., CSVs) after parsing. Defaults to False. If True, forces fallback to torch_npu.profiler. + :type keep_res: bool, optional + :param target_kernel_name: Specific NPU kernel name to filter and benchmark. If None, benchmarks the entire function. If specified, returns the execution time of only that kernel. Defaults to None. Only used when falling back to torch_npu.profiler. + :type target_kernel_name: str, optional + """ + import math + import os + mspti_available = True + if KernelMonitor is None: + mspti_available = False + print(f"[WARNING] mspti package not found. Falling back to torch_npu.profiler.") + _check_mspti_env() + + if not isinstance(funcs, list): + funcs = [funcs] + use_autotune = False + else: + use_autotune = os.getenv("TRITON_BENCH_METHOD", "default").lower() == "npu" + + results = None + need_fallback = True + force_fallback = (prof_dir is not None) or keep_res or (target_kernel_name is not None and not use_autotune) + + if mspti_available and not force_fallback: + try: + results = do_bench_npu_mspti(funcs, warmup, active, clear_l2_cache, target_kernel_name) + first_val = results[0] if isinstance(results, list) else results + if not math.isinf(first_val): + need_fallback = False + except Exception: + pass + if need_fallback: + results = do_bench_npu_profiler(funcs, warmup, active, clear_l2_cache, prof_dir, keep_res, target_kernel_name, + use_autotune) + return results diff --git a/third_party/ascend/backend/utils.py b/third_party/ascend/backend/utils.py new file mode 100644 index 0000000000..fc9399151f --- /dev/null +++ b/third_party/ascend/backend/utils.py @@ -0,0 +1,716 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import functools +import hashlib +import os +import re +import shutil +import subprocess +import sysconfig +from pathlib import Path +import logging +import platform +from triton.tools.get_ascend_devices import is_compile_on_910_95 +from triton.backends.ascend.backend_register import backend_strategy_registry + +import pybind11 + +backend_policy = None + + +def get_backend_func(name, *args, **kwargs): + global backend_policy + if backend_policy is None: + backend_policy_env = os.getenv("TRITON_BACKEND", "default").lower() + if backend_policy_env == "torch_npu" or backend_policy_env == "mindspore": + backend_policy = backend_policy_env + if backend_policy is None: + try: + import torch + import torch_npu + backend_policy = "torch_npu" + except ImportError: + backend_policy = "mindspore" + return backend_strategy_registry.execute_func(backend_policy, name, *args, **kwargs) + + +def get_logger(logger_name, logger_level_str): + ''' + ''' + logging_level_mapping = { + "DEBUG": logging.DEBUG, + "INFO": logging.INFO, + "WARNING": logging.WARNING, + "ERROR": logging.ERROR, + "CRITICAL": logging.CRITICAL, + } + logger = logging.getLogger(logger_name) + logger.propagate = False + formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') + console_handler = logging.StreamHandler() + console_handler.setFormatter(formatter) + logger.setLevel(logging_level_mapping.get(logger_level_str.upper(), "INFO")) + logger.addHandler(console_handler) + return logger + + +def downgrade_llir(llir): + llir = _downgrade_mem_attrs(llir) + llir = _downgrade_stacksaverestore_intrinsics(llir) + return llir + + +def _downgrade_mem_attrs(llir: str): + memory_pattern = r"memory\([^()]*\)" + + def replace_mem_attr(m): + attrs = m[0][7:-1].split(",") + if len(attrs) == 0: + return "readnone" + loc_map = {"argmem": 1, "inaccessiblemem": 2, "other": 4} + loc_attr = 0 + rw_map = {"readwrite": 3, "write": 2, "read": 1, "none": 0} + rw_attr = 0 + for attr_pair in attrs: + pair = attr_pair.split(":") + assert len(pair) <= 2 + if len(pair) == 1: + rw = rw_map[pair[0].strip()] + loc = loc_map["other"] # all location + else: + rw = rw_map[pair[1].strip()] + loc_str = pair[0].strip() + if loc_str == "argmem" or loc_str == "inaccessiblemem": + loc = loc_map[loc_str] + else: + loc = loc_map["other"] + if rw > 0: + loc_attr = loc_attr | loc + rw_attr = rw_attr | rw + rev_rw_map = {0: "readnone", 1: "readonly", 2: "writeonly"} + if rw_attr in rev_rw_map: + rw_attr_str = rev_rw_map[rw_attr] + else: + rw_attr_str = "" + rev_loc_map = { + 1: "argmemonly", + 2: "inaccessiblememonly", + 3: "inaccessiblemem_or_argmemonly", + } + if loc_attr in rev_loc_map: + loc_attr_str = rev_loc_map[loc_attr] + else: + loc_attr_str = "" + return rw_attr_str + " " + loc_attr_str + + return re.sub(memory_pattern, replace_mem_attr, llir) + + +def _downgrade_stacksaverestore_intrinsics(llir: str): + llir = re.sub(r"llvm\.stacksave\.\w+", "llvm.stacksave", llir) + llir = re.sub(r"llvm\.stackrestore\.\w+", "llvm.stackrestore", llir) + return llir + + +def _get_triton_adapter_opt_path() -> str: + path = os.path.dirname(__file__) + path = os.path.join(path, "triton-adapter-opt") + return path + + +def _get_mlir_path(path: str, *paths) -> str: + root_path = os.getenv("MLIR_ROOT", "") + if root_path == "": + raise EnvironmentError("MLIR_ROOT is not set.") + return os.path.join(root_path, path, *paths) + + +def _get_llvm_path(path: str, *paths) -> str: + root_path = os.getenv("LLVM_ROOT", "") + if root_path == "": + raise EnvironmentError("LLVM_ROOT is not set.") + return os.path.join(root_path, path, *paths) + + +def _get_tool_path(tool_name: str) -> str: + """ + Get the path to a Triton tool binary. + Search order: + 1. Installed package location (triton/_C/) + 2. TRITON_BUILD_DIR environment variable + 3. System PATH + """ + try: + import triton._C.libtriton as libtriton + tool_path = os.path.join(os.path.dirname(libtriton.__file__), tool_name) + if os.path.exists(tool_path) and os.access(tool_path, os.X_OK): + return tool_path + except (ImportError, AttributeError): + pass + + build_path = os.getenv("TRITON_BUILD_DIR", "") + if build_path: + tool_path = os.path.join(build_path, "bin", tool_name) + if os.path.exists(tool_path) and os.access(tool_path, os.X_OK): + return tool_path + + tool_path = shutil.which(tool_name) + if tool_path: + return tool_path + + raise EnvironmentError(f"Could not find {tool_name} tool. " + f"It should be installed in triton/_C/ directory or available in PATH.") + + +def _get_triton_opt_path() -> str: + """ + Get the path to triton-opt tool. + This tool is used to convert ttir to ttadapter. + """ + return _get_tool_path("triton-opt") + + +def _get_triton_mlir_opt_path() -> str: + """ + Get the path to triton-mlir-opt tool. + This tool is used to convert MLIR to Bytecode format, supporting both + MLIR native ops and AscendNPU-IR custom ops. + """ + return _get_tool_path("triton-mlir-opt") + + +def _get_bishengir_opt_path() -> str: + ascend_dir = os.path.dirname(os.path.abspath(__file__)) + env = os.environ.copy() + bishengir_opt_path = os.path.join(ascend_dir, "bishengir", "bin", "bishengir-opt") + if os.path.exists(bishengir_opt_path): + npuir_env_path = os.path.dirname(bishengir_opt_path) + env["PATH"] = npuir_env_path + ":" + env["PATH"] + else: + bishengir_opt_path = shutil.which("bishengir-opt") + if bishengir_opt_path is None: + bishengir_opt_root = os.getenv("TRITON_NPU_COMPILER_PATH", "") + if bishengir_opt_root is None: + raise EnvironmentError("Couldn't find executable bishengir-opt or TRITON_NPU_COMPILER_PATH") + bishengir_opt_path = os.path.join(bishengir_opt_root, "bishengir-opt") + return bishengir_opt_path, env + + +def _get_npucompiler_path() -> str: + ascend_dir = os.path.dirname(os.path.abspath(__file__)) + env = os.environ.copy() + npu_compiler_path = os.path.join(ascend_dir, "bishengir", "bin", "bishengir-compile") + if os.path.exists(npu_compiler_path): + npuir_env_path = os.path.dirname(npu_compiler_path) + env["PATH"] = npuir_env_path + ":" + env["PATH"] + else: + npu_compiler_path = shutil.which("bishengir-compile") + if npu_compiler_path is None: + npu_compiler_root = os.getenv("TRITON_NPU_COMPILER_PATH", None) + if npu_compiler_root is None: + raise EnvironmentError("Couldn't find executable bishengir-compile or TRITON_NPU_COMPILER_PATH.") + npu_compiler_path = os.path.join(npu_compiler_root, "npuc") + return npu_compiler_path, env + + +def _get_bisheng_path() -> str: + bisheng_path = shutil.which("bisheng") + if bisheng_path is None: + npu_compiler_root = os.getenv("TRITON_NPU_COMPILER_PATH", None) + if npu_compiler_root is None: + raise EnvironmentError("Couldn't find executable bisheng or TRITON_NPU_COMPILER_PATH") + bisheng_path = os.path.join(npu_compiler_root, "ccec") + return bisheng_path + + +def _is_valid_bishengir_path(path: str) -> bool: + if not path or not isinstance(path, str): + return False + if os.path.basename(path) != "bishengir-compile": + return False + if not os.path.isfile(path) or not os.access(path, os.X_OK): + return False + return True + + +# grep bishengir-compile's option limit-auto-multi-buffer-buffer to check +# if bishengir-compile is a newer version which does not generate kernel_reloc.o +# any more. +def _check_bishengir_api_change() -> bool: + bishengir_path, _ = _get_npucompiler_path() + if not _is_valid_bishengir_path(bishengir_path): + print(f"ERROR: Invalid bishengir path format: {bishengir_path}") + return False + try: + result = subprocess.run( + [bishengir_path, "--help"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + if result.returncode == 0 and 'limit-auto-multi-buffer-buffer' in result.stdout: + # bishengir-compile is newer version + return True + else: + # bishengir-compile is older version + return False + except Exception as e: + print(f"ERROR: {e}") + return False + + +def _check_bishengir_is_regbased() -> bool: + bishengir_path, _ = _get_npucompiler_path() + if not _is_valid_bishengir_path(bishengir_path): + print(f"ERROR: Invalid bishengir path format: {bishengir_path}") + return False + try: + result = subprocess.run( + [bishengir_path, "--help"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + if result.returncode == 0 and 'reg-based' in result.stdout: + # bishengir-compile is regbased version + return True + else: + # bishengir-compile is membased version + return False + except Exception as e: + print(f"ERROR: {e}") + return False + + +@functools.lru_cache(None) +def _get_ascend_path() -> Path: + path = os.getenv("ASCEND_HOME_PATH", "") + if path == "": + raise EnvironmentError("ASCEND_HOME_PATH is not set, source /set_env.sh first") + return Path(path) + + +def _is_ascend_sanitizer_enabled() -> bool: + return os.getenv("TRITON_ENABLE_SANITIZER", "false").lower() in ("true", "1") + + +def _is_debug_line_info_disabled() -> bool: + return os.getenv("TRITON_DISABLE_LINE_INFO", "true").lower() in ("true", "1") + + +def _is_auto_map_parallel_blocks_enabled() -> bool: + return os.getenv("TRITON_ALL_BLOCKS_PARALLEL", "false").lower() in ("true", "1") + + +def _enable_unpublished_feature() -> bool: + return os.getenv("ENABLE_UNPUBLISHED_FEATURE", "false").lower() in ("true", "1") + + +def _enable_print_ub_bits() -> bool: + return os.getenv("ENABLE_PRINT_UB_BITS", "false").lower() in ("true", "1") + + +def _enable_dump_memory_info() -> bool: + return os.getenv("TRITON_MEMORY_DISPLAY", "false").lower() in ("true", "1") + + +def _enable_msdebug() -> bool: + return os.getenv("LLVM_EXTRACT_DI_LOCAL_VARIABLES", "false").lower() in ("true", "1") + + +def _get_cxx(): + cxx = os.environ.get("CC") + if cxx is None: + clangxx = shutil.which("clang++") + gxx = shutil.which("g++") + cxx = clangxx if clangxx is not None else gxx + if cxx is None: + raise RuntimeError("Failed to find C++ compiler") + return cxx + + +def _get_cxx_precompiled(header_path): + cc_cmd = [] + cxx = os.environ.get("CC") + if cxx is None: + clangxx = shutil.which("clang++") + gxx = shutil.which("g++") + if clangxx is not None: + cc_cmd += [clangxx, "-include", header_path] + elif gxx is not None: + cc_cmd += [gxx] + else: + raise RuntimeError("Failed to find C++ compiler") + else: + cc_cmd += [cxx] + return cc_cmd + + +def _precompile_npu_hash(header_src): + import sys + cxx = _get_cxx() + py_version = sys.version + asc_path = _get_ascend_path().name + version_txt = [header_src, cxx, py_version, asc_path] + version_txt += get_backend_func("version_hash") + hash_txt = hashlib.sha256("_".join(version_txt).encode("utf-8")).hexdigest() + return hash_txt + + +def _precompile_npu_ext(header_path, gch_path): + cxx = _get_cxx() + cc_cmd = [cxx, "-x", "c++-header", header_path] + # disable all warnings + cc_cmd += [f"-w"] + # find the python library + if hasattr(sysconfig, "get_default_scheme"): + scheme = sysconfig.get_default_scheme() + else: + scheme = sysconfig._get_default_scheme() + # 'posix_local' is a custom scheme on Debian. However, starting Python 3.10, the default install + # path changes to include 'local'. This change is required to use triton with system-wide python. + if scheme == "posix_local": + scheme = "posix_prefix" + py_include_dir = sysconfig.get_paths(scheme=scheme)["include"] + cc_cmd += [f"-I{py_include_dir}"] + # device_print.h + cc_cmd += [f"-I{os.path.dirname(os.path.realpath(__file__))}"] + # find the ascend library + asc_path = _get_ascend_path() + + rt_path = os.path.join(asc_path, "include/experiment/runtime/runtime/rt.h") + if not os.path.exists(rt_path): + cc_cmd += [ + f"-I{os.path.join(asc_path, 'pkg_inc')}", + f"-I{os.path.join(asc_path, 'pkg_inc/profiling')}", + ] + + cc_cmd += [ + f"-I{os.path.join(asc_path, 'include')}", + f"-I{os.path.join(asc_path, 'include/experiment')}", + f"-I{os.path.join(asc_path, 'include/experiment/msprof')}", + f"-I{pybind11.get_include()}", + ] + + cc_cmd += get_backend_func("get_cc_cmd", build_pch=True) + + cc_cmd += ["-std=c++17", "-shared", "-fPIC", "-o", gch_path] + + result = subprocess.run(cc_cmd, capture_output=True, text=True) + + if result.returncode == 0: + return header_path + else: + raise RuntimeError(f"Failed to compile {gch_path}, error: {result.stderr},cmd={cc_cmd}") + + +def _build_npu_ext(obj_name: str, header_path, src_path, *, kernel_launcher="torch", precompile=False) -> str: + suffix = sysconfig.get_config_var("EXT_SUFFIX") + src_dir = os.path.dirname(src_path) + so_path = os.path.join(src_dir, f"{obj_name}{suffix}") + if precompile: + cc_cmd = _get_cxx_precompiled(header_path) + cc_cmd += [src_path] + else: + cxx = _get_cxx() + cc_cmd = [cxx, src_path] + # disable all warnings + cc_cmd += [f"-w"] + # find the python library + if hasattr(sysconfig, "get_default_scheme"): + scheme = sysconfig.get_default_scheme() + else: + scheme = sysconfig._get_default_scheme() + # 'posix_local' is a custom scheme on Debian. However, starting Python 3.10, the default install + # path changes to include 'local'. This change is required to use triton with system-wide python. + if scheme == "posix_local": + scheme = "posix_prefix" + py_include_dir = sysconfig.get_paths(scheme=scheme)["include"] + cc_cmd += [f"-I{py_include_dir}"] + # device_print.h + cc_cmd += [f"-I{os.path.dirname(os.path.realpath(__file__))}"] + # find the ascend library + asc_path = _get_ascend_path() + if header_path is not None: + cc_cmd += [f"-I{os.path.dirname(header_path)}"] + + rt_path = os.path.join(asc_path, "include/experiment/runtime/runtime/rt.h") + if not os.path.exists(rt_path): + cc_cmd += [ + f"-I{os.path.join(asc_path, 'pkg_inc')}", + f"-I{os.path.join(asc_path, 'pkg_inc/profiling')}", + ] + + cc_cmd += [ + f"-I{os.path.join(asc_path, 'include')}", + f"-I{os.path.join(asc_path, 'include/experiment')}", + f"-I{os.path.join(asc_path, 'include/experiment/msprof')}", + f"-I{pybind11.get_include()}", + f"-L{os.path.join(asc_path, 'lib64')}", + "-lruntime", + "-lascendcl", + ] + # FIXME: check why this condition works wrong in parall scene + if kernel_launcher == "torch": + cc_cmd += get_backend_func("get_cc_cmd", build_pch=False) + + cc_cmd += ["-std=c++17", "-shared", "-fPIC", "-Winvalid-pch", "-o", so_path] + + result = subprocess.run(cc_cmd, capture_output=True, text=True) + + if result.returncode == 0: + return so_path + else: + if "precompiled.h.gch" in result.stderr: + # only for clang++, when precompile invalid, fallback to normal compile + return _build_npu_ext(obj_name, header_path, src_path, precompile=False) + else: + raise RuntimeError(f"Failed to compile {src_path}, error: {result.stderr},cmd={cc_cmd}") + + +def _get_kernel_target(metadata: dict): + if "target" not in metadata: + raise Exception("No target provided!") + sub_target = metadata["target"].arch + assert isinstance(sub_target, str) + if sub_target.startswith("Ascend910B"): + mix_mode = metadata["mix_mode"] + if mix_mode.lower().strip("_").startswith("aiv"): + return "ascend_910b_vec", "c220-vec", "aiv" + elif mix_mode.lower().strip("_").startswith("aic"): + return "ascend_910b_cube", "c220-cube", "aic" + else: + return "ascend_910b", "c220", "mix" + elif sub_target.startswith("Ascend910"): + return "ascend_910", "c100", "mix" + else: + raise NotImplementedError(f"NPU subtarget {sub_target} not supported yet") + + +def _check_cxx11_abi(): + return get_backend_func("cxx_abi") + + +def convert_sigtype_to_int(sigty: str): + MAP_SIGTYPE_TO_INT = { + # Boolean + "i1": 12, # BOOL + # Integer types + "i4": 29, # INT4 + "i8": 2, # INT8 + "i16": 6, # INT16 + "i32": 3, # INT32 + "i64": 9, # INT64 + # Unsigned integer types + "u1": 30, # UINT1 + "u8": 4, # UINT8 + "u16": 7, # UINT16 + "u32": 8, # UINT32 + "u64": 10, # UINT64 + # Floating point types + "fp16": 1, # FLOAT16 + "bf16": 27, # DT_BF16 + "fp32": 0, # FLOAT + "fp64": 11, # DOUBLE + "fp8e5": 35, # FLOAT8_E5M2 + "fp8e4nv": 36, # FLOAT8_E4M3FN + } + if sigty not in MAP_SIGTYPE_TO_INT: + raise ValueError(f"Unsupported data type: {sigty}") + + return MAP_SIGTYPE_TO_INT[sigty] + + +def convert_dtype_to_numpy(dtype): + return get_backend_func("type_convert")[dtype] + + +def _check_bishengir_able_save_ir() -> bool: + bishengir_path, _ = _get_npucompiler_path() + if not _is_valid_bishengir_path(bishengir_path): + print(f"ERROR: Invalid bishengir path format: {bishengir_path}") + return False + try: + result = subprocess.run( + [bishengir_path, "--help"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + if result.returncode == 0 and 'save-linked-ir' in result.stdout: + return True + else: + return False + except Exception as e: + print(f"ERROR: {e}") + return False + + +def get_ascend_arch_from_env(): + arch = os.getenv("TRITON_ASCEND_ARCH", "") + if arch == "": + # User does not set arch by ENV. Thus directly return. + return arch + # User sets arch by ENV. Thus we need to check if it is supported. + valid_arch_list = [ + "Ascend910B1", + "Ascend910B2", + "Ascend910B3", + "Ascend910B4", + "Ascend910_9362", + "Ascend910_9372", + "Ascend910_9381", + "Ascend910_9382", + "Ascend910_9391", + "Ascend910_9392", + "Ascend310B1", + "Ascend310B2", + "Ascend310B3", + "Ascend310B4", + "Ascend910_9579", + "Ascend910_9581", + "Ascend910_9589", + "Ascend910_9599", + ] + is_valid = arch in valid_arch_list + if not is_valid: + valid_arch_str = ", ".join(valid_arch_list) + raise ValueError(f"TRITON_ASCEND_ARCH = {arch} is invalid!" + f"Candidates are [{valid_arch_str}]") + return arch + + +def is_ffts_supported(arch: str): + ''' + Cases: + - empty str: User does not specify arch, thus it runs on 910B/910D both of which support ffts. Return True. + - Ascend310B4: 310B4 does not support ffts. Return False. + - Ascend910_95*: 910_95 does not support ffts. Return False. + - Other arch: 910B/910D supports ffts. Return True. + ''' + if is_compile_on_910_95: + return False + if arch in ["Ascend910A", "Ascend310B4"]: + return False + return True + + +def force_disable_ffts(): + ''' + ''' + if is_compile_on_910_95: + return True + disable_ffts = os.getenv("TRITON_DISABLE_FFTS", "false").lower() in ("true", "1") + return disable_ffts + + +def triton_support_ffts(): + arch = get_ascend_arch_from_env() + return is_ffts_supported(arch) and (not force_disable_ffts()) + + +def triton_enable_libdevice_simt(): + enable_libdevice_simt = os.getenv("TRITON_ENABLE_LIBDEVICE_SIMT", False) + return enable_libdevice_simt + + +def _parse_cann_version(line: str): + m = re.search(r'(\d+)\.(\d+)(?:\.(\d+))?', line) + if m: + major = int(m.group(1)) + minor = int(m.group(2)) + patch = int(m.group(3)) if m.group(3) is not None else 0 + return (major, minor, patch) + return None + + +def _find_cann_version_file(): + ascend_path = str(_get_ascend_path()) + arch = get_machine_arch() + candidates = [ + os.path.join(ascend_path, arch + "-linux", "ascend_toolkit_install.info"), + os.path.join(ascend_path, arch + "-linux", "ascend_all_cann_install.info"), + ] + for p in candidates: + if os.path.exists(p): + return p + return None + + +@functools.lru_cache(None) +def get_cann_version(): + try: + version_file = _find_cann_version_file() + if version_file is None: + return None + with open(version_file, "r", encoding="utf-8", errors="ignore") as f: + for line in f: + line = line.strip() + if "version" in line.lower(): + parsed = _parse_cann_version(line) + if parsed is not None: + return parsed + return None + except Exception: + raise EnvironmentError("Could not parse CANN version file") + + +def is_cann_version_at_least(major: int, minor: int = 0, patch: int = 0) -> bool: + v = get_cann_version() + if v is None: + return False + return v >= (major, minor, patch) + + +def get_cann_version_file_hash(): + ascend_path = _get_ascend_path() + arch = get_machine_arch() + cann_version_file_path = os.path.join(ascend_path, arch + "-linux", "ascend_toolkit_install.info") + if not os.path.exists(cann_version_file_path): + cann_version_file_path = os.path.join(ascend_path, arch + "-linux", "ascend_all_cann_install.info") + return get_file_hash256(cann_version_file_path) + + +def get_file_hash256(file_path): + sha256 = hashlib.sha256() + try: + with open(file_path, "rb") as f: + for byte_block in iter(lambda: f.read(4096), b""): + sha256.update(byte_block) + return sha256.hexdigest() + except FileNotFoundError as e: + raise FileNotFoundError(f"Can't find file {os.path.basename(file_path)}") from e + + +def get_machine_arch(): + ARCHITECTURE_ALIASES = { + "x86_64": "x86_64", + "amd64": "x86_64", + "i386": "x86_64", + "i686": "x86_64", + "arm64": "aarch64", + "aarch64": "aarch64", + "armv7l": "aarch64", + "armv8l": "aarch64", + "arm": "aarch64", + } + system_arch = platform.machine() + return ARCHITECTURE_ALIASES[system_arch] diff --git a/third_party/ascend/bin/CMakeLists.txt b/third_party/ascend/bin/CMakeLists.txt new file mode 100644 index 0000000000..627d5d6b11 --- /dev/null +++ b/third_party/ascend/bin/CMakeLists.txt @@ -0,0 +1,218 @@ +get_property(triton_libs GLOBAL PROPERTY TRITON_LIBS) + +include_directories( + ${PROJECT_SOURCE_DIR}/third_party/ascend/AscendNPU-IR/bishengir/include + ${PROJECT_BINARY_DIR}/third_party/ascend/AscendNPU-IR/bishengir/include + ) + +add_llvm_executable(triton-opt triton-opt.cpp PARTIAL_SOURCES_INTENDED) + +# TODO: what's this? +llvm_update_compile_flags(triton-opt) +target_link_libraries(triton-opt PRIVATE + ${triton_libs} + # Ascend backend libraries + DiscreteMaskAccessConversion + TritonToAnnotation + TritonToHFusion + TritonToHIVM + TritonToLinalgIncubated + TritonToLLVM + TritonToStructuredIncubated + TritonToUnstructureIncubated + # BiShengIR dialects + BiShengIRAnnotationDialect + BiShengIRHACCDialect + BiShengIRHFusionDialect + BiShengIRHIVMDialect + BiShengIRMathExtDialect + BiShengIRMemRefExtDialect + BiShengIRScopeDialect + BiShengIRSymbolDialect + BiShengIRTensorDialect + # tests + TritonTestAnalysis + TritonTestDialect + TritonAMDGPUTestAnalysis + TritonTestProton + # MLIR core + MLIROptLib + MLIRPass + MLIRRegisterAllDialects + MLIRRegisterAllPasses + MLIRTransforms +) + +mlir_check_all_link_libraries(triton-opt) + +add_llvm_executable(triton-reduce triton-reduce.cpp PARTIAL_SOURCES_INTENDED) +mlir_check_all_link_libraries(triton-reduce) + +llvm_update_compile_flags(triton-reduce) +target_link_libraries(triton-reduce PRIVATE + ${triton_libs} + # Ascend backend libraries + DiscreteMaskAccessConversion + TritonToAnnotation + TritonToHFusion + TritonToHIVM + TritonToLinalgIncubated + TritonToLLVM + TritonToStructuredIncubated + TritonToUnstructureIncubated + # BiShengIR dialects + BiShengIRAnnotationDialect + BiShengIRHACCDialect + BiShengIRHFusionDialect + BiShengIRHIVMDialect + BiShengIRMathExtDialect + BiShengIRMemRefExtDialect + BiShengIRScopeDialect + BiShengIRSymbolDialect + BiShengIRTensorDialect + # tests + TritonTestAnalysis + TritonTestDialect + TritonAMDGPUTestAnalysis + TritonTestProton + # MLIR core + MLIRReduceLib + MLIRPass + MLIRRegisterAllDialects + MLIRRegisterAllPasses + MLIRTransforms +) + +mlir_check_all_link_libraries(triton-reduce) + +add_llvm_executable(triton-lsp triton-lsp.cpp PARTIAL_SOURCES_INTENDED) + +llvm_update_compile_flags(triton-lsp) +target_link_libraries(triton-lsp PRIVATE + ${triton_libs} + # Ascend backend libraries + DiscreteMaskAccessConversion + TritonToAnnotation + TritonToHFusion + TritonToHIVM + TritonToLinalgIncubated + TritonToLLVM + TritonToStructuredIncubated + TritonToUnstructureIncubated + # BiShengIR dialects + BiShengIRAnnotationDialect + BiShengIRHACCDialect + BiShengIRHFusionDialect + BiShengIRHIVMDialect + BiShengIRMathExtDialect + BiShengIRMemRefExtDialect + BiShengIRScopeDialect + BiShengIRSymbolDialect + BiShengIRTensorDialect + # tests + TritonTestAnalysis + TritonTestDialect + TritonAMDGPUTestAnalysis + TritonTestProton + # MLIR core + MLIRLspServerLib + MLIRPass + MLIRRegisterAllDialects + MLIRRegisterAllPasses + MLIRTransforms +) + +mlir_check_all_link_libraries(triton-lsp) + + +add_llvm_executable(triton-llvm-opt + triton-llvm-opt.cpp + + PARTIAL_SOURCES_INTENDED + DEPENDS + intrinsics_gen + SUPPORT_PLUGINS + ) +target_link_libraries(triton-llvm-opt PRIVATE + TritonLLVMIR + + LLVMAnalysis + LLVMCore + LLVMSupport + LLVMOption + LLVMCodeGen + ) +export_executable_symbols_for_plugins(triton-llvm-opt) + + +add_llvm_executable(triton-tensor-layout triton-tensor-layout.cpp PARTIAL_SOURCES_INTENDED) +target_link_libraries(triton-tensor-layout PRIVATE + ${triton_libs} + # Ascend backend libraries + DiscreteMaskAccessConversion + TritonToAnnotation + TritonToHFusion + TritonToHIVM + TritonToLinalgIncubated + TritonToLLVM + TritonToStructuredIncubated + TritonToUnstructureIncubated + # BiShengIR dialects + BiShengIRAnnotationDialect + BiShengIRHACCDialect + BiShengIRHFusionDialect + BiShengIRHIVMDialect + BiShengIRMathExtDialect + BiShengIRMemRefExtDialect + BiShengIRScopeDialect + BiShengIRSymbolDialect + BiShengIRTensorDialect + # tests + TritonTestAnalysis + TritonTestDialect + TritonTestProton + TritonAMDGPUTestAnalysis + MLIRRegisterAllDialects + MLIRRegisterAllPasses + MLIRTransforms + ) + +# triton-mlir-opt: Triton-Ascend optimizer combining MLIR and BishengIR dialects +add_llvm_executable(triton-mlir-opt triton-mlir-opt.cpp PARTIAL_SOURCES_INTENDED) +llvm_update_compile_flags(triton-mlir-opt) + +target_link_libraries(triton-mlir-opt PRIVATE + ${triton_libs} + # MLIR core + MLIROptLib + MLIRPass + MLIRRegisterAllDialects + MLIRRegisterAllPasses + MLIRTransforms + # Core BishengIR dialect libraries + BiShengIRAnnotationDialect + BiShengIRHACCDialect + BiShengIRHFusionDialect + BiShengIRHIVMDialect + BiShengIRMathExtDialect + BiShengIRMemRefExtDialect + BiShengIRScopeDialect + BiShengIRSymbolDialect + BiShengIRTensorDialect +) + +mlir_check_all_link_libraries(triton-mlir-opt) + +if(DEFINED CMAKE_LIBRARY_OUTPUT_DIRECTORY) + set(_ascend_tool_targets + triton-opt + triton-mlir-opt + ) + foreach(_tool IN LISTS _ascend_tool_targets) + if(TARGET ${_tool}) + set_target_properties(${_tool} PROPERTIES + RUNTIME_OUTPUT_DIRECTORY "${CMAKE_LIBRARY_OUTPUT_DIRECTORY}" + ) + endif() + endforeach() +endif() diff --git a/third_party/ascend/bin/RegisterTritonDialects.h b/third_party/ascend/bin/RegisterTritonDialects.h new file mode 100644 index 0000000000..50b7f0f165 --- /dev/null +++ b/third_party/ascend/bin/RegisterTritonDialects.h @@ -0,0 +1,192 @@ +#pragma once +#include "ascend/include/AutoBlockify/Passes.h" +#include "ascend/include/CommonIRToHIVM/Passes.h" +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition.h" +#include "ascend/include/TritonToAnnotation/Passes.h" +#include "ascend/include/TritonToHFusion/Passes.h" +#include "ascend/include/TritonToHIVM/Passes.h" +#include "ascend/include/TritonToLLVM/Passes.h" +#include "incubated/Conversion/DiscreteMaskAccessConversion/Passes.h" +#include "incubated/Conversion/TritonToLinalgIncubated/Passes.h" + +#include "ascend/include/TritonToGraph/Passes.h" +#include "incubated/Conversion/TritonToStructuredIncubated/Passes.h" +#include "incubated/Conversion/TritonToUnstructureIncubated/Passes.h" + +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "bishengir/Dialect/HACC/IR/HACC.h" +#include "bishengir/Dialect/HFusion/IR/HFusion.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/Scope/IR/Scope.h" + +#include "amd/include/Dialect/TritonAMDGPU/IR/Dialect.h" +#include "amd/include/TritonAMDGPUTransforms/Passes.h" +#include "ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.h" +#include "nvidia/include/Dialect/NVGPU/IR/Dialect.h" +#include "nvidia/include/Dialect/NVWS/IR/Dialect.h" +#ifdef __FLAGPRISM__ +#include "FlagPrism/Profiler/Dialect/include/Integration/Registration.h" +#else +#include "proton/Dialect/include/Conversion/ProtonGPUToLLVM/Passes.h" +#include "proton/Dialect/include/Conversion/ProtonGPUToLLVM/ProtonAMDGPUToLLVM/Passes.h" +#include "proton/Dialect/include/Conversion/ProtonGPUToLLVM/ProtonNvidiaGPUToLLVM/Passes.h" +#include "proton/Dialect/include/Conversion/ProtonToProtonGPU/Passes.h" +#include "proton/Dialect/include/Dialect/Proton/IR/Dialect.h" +#include "proton/Dialect/include/Dialect/ProtonGPU/IR/Dialect.h" +#include "proton/Dialect/include/Dialect/ProtonGPU/Transforms/Passes.h" +#endif +#include "triton/Dialect/Gluon/Transforms/Passes.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "triton/Dialect/TritonInstrument/IR/Dialect.h" +#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" + +// Below headers will allow registration to ROCm passes +#include "TritonAMDGPUToLLVM/Passes.h" +#include "TritonAMDGPUTransforms/Passes.h" +#include "TritonAMDGPUTransforms/TritonGPUConversion.h" + +#include "triton/Dialect/Triton/Transforms/Passes.h" +#include "triton/Dialect/TritonGPU/Transforms/Passes.h" +#include "triton/Dialect/TritonInstrument/Transforms/Passes.h" +#include "triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h" + +#include "nvidia/hopper/include/Transforms/Passes.h" +#include "nvidia/include/Dialect/NVWS/Transforms/Passes.h" +#include "nvidia/include/NVGPUToLLVM/Passes.h" +#include "nvidia/include/TritonNVIDIAGPUToLLVM/Passes.h" +#include "triton/Conversion/TritonGPUToLLVM/Passes.h" +#include "triton/Conversion/TritonToTritonGPU/Passes.h" +#include "triton/Target/LLVMIR/Passes.h" + +#include "mlir/Dialect/LLVMIR/NVVMDialect.h" +#include "mlir/Dialect/LLVMIR/ROCDLDialect.h" +#include "mlir/InitAllPasses.h" + +namespace mlir { +namespace test { +void registerTestAliasPass(); +void registerTestAlignmentPass(); +void registerAMDTestAlignmentPass(); +void registerTestAllocationPass(); +void registerTestMembarPass(); +void registerTestAMDGPUMembarPass(); +void registerTestTritonAMDGPURangeAnalysis(); +void registerTestLoopPeelingPass(); +namespace proton { +void registerTestScopeIdAllocationPass(); +} // namespace proton +} // namespace test +} // namespace mlir + +inline void registerTritonDialects(mlir::DialectRegistry ®istry) { + mlir::registerAllPasses(); + mlir::triton::registerTritonPasses(); + mlir::triton::gpu::registerTritonGPUPasses(); + mlir::triton::nvidia_gpu::registerTritonNvidiaGPUPasses(); + mlir::triton::instrument::registerTritonInstrumentPasses(); + mlir::triton::gluon::registerGluonPasses(); + mlir::test::registerTestAliasPass(); + mlir::test::registerTestAlignmentPass(); + mlir::test::registerAMDTestAlignmentPass(); + mlir::test::registerTestAllocationPass(); + mlir::test::registerTestMembarPass(); + mlir::test::registerTestLoopPeelingPass(); + mlir::test::registerTestAMDGPUMembarPass(); + mlir::test::registerTestTritonAMDGPURangeAnalysis(); + mlir::triton::registerConvertTritonToTritonGPUPass(); + mlir::triton::registerRelayoutTritonGPUPass(); + mlir::triton::gpu::registerAllocateSharedMemoryPass(); + mlir::triton::gpu::registerTritonGPUAllocateWarpGroups(); + mlir::triton::gpu::registerTritonGPUGlobalScratchAllocationPass(); + mlir::triton::registerConvertWarpSpecializeToLLVM(); + mlir::triton::registerConvertTritonGPUToLLVMPass(); + mlir::triton::registerConvertNVGPUToLLVMPass(); + mlir::triton::registerAllocateSharedMemoryNvPass(); + mlir::triton::Incubated::registerTritonToLinalgIncubatedPasses(); + + mlir::triton::registerDiscreteMaskAccessConversion(); + mlir::triton::registerTritonToStructuredIncubatedPasses(); + mlir::triton::registerTritonToAnnotationPasses(); + mlir::triton::registerTritonToUnstructureIncubatedPasses(); + mlir::triton::registerTritonToHIVMPasses(); + mlir::triton::registerTritonToHFusionPasses(); + mlir::triton::registerTritonToLLVMPasses(); + mlir::triton::registerAutoBlockifyPasses(); + mlir::triton::registerCommonIRToHIVMPasses(); + mlir::triton::cfg::registerTritonToGraphPasses(); + mlir::triton::registerBubbleUpOperationPass(); + + mlir::registerLLVMDIScope(); + + // TritonAMDGPUToLLVM passes + mlir::triton::registerAllocateAMDGPUSharedMemory(); + mlir::triton::registerConvertTritonAMDGPUToLLVM(); + mlir::triton::registerConvertBuiltinFuncToLLVM(); + mlir::triton::registerOptimizeAMDLDSUsage(); + + // TritonAMDGPUTransforms passes + mlir::registerTritonAMDGPUAccelerateMatmul(); + mlir::registerTritonAMDGPUOptimizeEpilogue(); + mlir::registerTritonAMDGPUHoistLayoutConversions(); + mlir::registerTritonAMDGPUReorderInstructions(); + mlir::registerTritonAMDGPUBlockPingpong(); + // Renamed from registerTritonAMDGPUStreamPipeline; ScheduleLoops was split out. + mlir::registerTritonAMDGPUPipeline(); + mlir::registerTritonAMDGPUScheduleLoops(); + mlir::registerTritonAMDGPUCanonicalizePointers(); + mlir::registerTritonAMDGPUConvertToBufferOps(); + mlir::registerTritonAMDGPUInThreadTranspose(); + mlir::registerTritonAMDGPUCoalesceAsyncCopy(); + mlir::registerTritonAMDGPUUpdateAsyncWaitCount(); + mlir::triton::registerTritonAMDGPUInsertInstructionSchedHints(); + mlir::triton::registerTritonAMDGPULowerInstructionSchedHints(); + mlir::registerTritonAMDFoldTrueCmpI(); + // mlir::triton::amdgpu::registerTritonAMDGPUOptimizeDotOperands(); + + // NVWS passes + mlir::triton::registerNVWSTransformsPasses(); + + // NVGPU transform passes + mlir::registerNVHopperTransformsPasses(); + + // Proton passes +#ifdef __FLAGPRISM__ + mlir::triton::proton::registerFlagTreeProtonTestPasses(); + mlir::triton::proton::registerFlagTreeProtonPassesAndDialects(registry); +#else + mlir::test::proton::registerTestScopeIdAllocationPass(); + mlir::triton::proton::registerConvertProtonToProtonGPU(); + mlir::triton::proton::gpu::registerConvertProtonNvidiaGPUToLLVM(); + mlir::triton::proton::gpu::registerConvertProtonAMDGPUToLLVM(); + mlir::triton::proton::gpu::registerAllocateProtonSharedMemoryPass(); + mlir::triton::proton::gpu::registerAllocateProtonGlobalScratchBufferPass(); + mlir::triton::proton::gpu::registerScheduleBufferStorePass(); + mlir::triton::proton::gpu::registerAddSchedBarriersPass(); +#endif + + // DynamicCVPipeline passes + mlir::triton::registerAddControlFlowConditionPasses(); + + registry.insert< + mlir::triton::TritonDialect, mlir::cf::ControlFlowDialect, + mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect, + mlir::triton::gpu::TritonGPUDialect, + mlir::triton::instrument::TritonInstrumentDialect, + mlir::math::MathDialect, mlir::arith::ArithDialect, mlir::scf::SCFDialect, + mlir::tensor::TensorDialect, mlir::gpu::GPUDialect, + mlir::LLVM::LLVMDialect, mlir::NVVM::NVVMDialect, + mlir::triton::nvgpu::NVGPUDialect, mlir::triton::nvws::NVWSDialect, + mlir::triton::amdgpu::TritonAMDGPUDialect, +#ifndef __FLAGPRISM__ + mlir::triton::proton::ProtonDialect, + mlir::triton::proton::gpu::ProtonGPUDialect, +#endif + mlir::ROCDL::ROCDLDialect, mlir::triton::gluon::GluonDialect, + mlir::triton::ascend::TritonAscendDialect, mlir::hivm::HIVMDialect, + mlir::scope::ScopeDialect, mlir::hacc::HACCDialect, + mlir::annotation::AnnotationDialect, mlir::hfusion::HFusionDialect, + mlir::tensor::TensorDialect, mlir::linalg::LinalgDialect, + mlir::memref::MemRefDialect, mlir::bufferization::BufferizationDialect, + mlir::func::FuncDialect>(); +} diff --git a/third_party/ascend/bin/triton-llvm-opt.cpp b/third_party/ascend/bin/triton-llvm-opt.cpp new file mode 100644 index 0000000000..3beeeabdc1 --- /dev/null +++ b/third_party/ascend/bin/triton-llvm-opt.cpp @@ -0,0 +1,121 @@ +/// Trimmed down clone of llvm opt to be able to test triton custom llvm ir +/// passes. +#include "lib/Target/LLVMIR/LLVMPasses.h" +#include "llvm/CodeGen/CommandFlags.h" +#include "llvm/IR/Constants.h" +#include "llvm/IR/DataLayout.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/Module.h" +#include "llvm/IR/Verifier.h" +#include "llvm/IRReader/IRReader.h" +#include "llvm/Passes/PassBuilder.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/Error.h" +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/InitLLVM.h" +#include "llvm/Support/SourceMgr.h" +#include "llvm/Support/SystemUtils.h" +#include "llvm/Support/ToolOutputFile.h" +#include "llvm/TargetParser/Triple.h" +#include + +using namespace llvm; + +static cl::opt InputFilename(cl::Positional, + cl::desc(""), + cl::init("-"), + cl::value_desc("filename")); + +static cl::opt OutputFilename("o", + cl::desc("Override output filename"), + cl::value_desc("filename")); + +static cl::opt ClDataLayout("data-layout", + cl::desc("data layout string to use"), + cl::value_desc("layout-string"), + cl::init("")); +static cl::opt + TargetTriple("mtriple", cl::desc("Override target triple for module")); + +static cl::opt + BreakStructPhiNodes("break-struct-phi-nodes", + llvm::cl::desc("run pass to break phi struct"), + cl::init(false)); + +namespace { +static std::function makeOptimizingPipeline() { + return [](Module *m) -> Error { + PipelineTuningOptions tuningOptions; + PassBuilder pb(nullptr, tuningOptions); + + LoopAnalysisManager lam; + FunctionAnalysisManager fam; + CGSCCAnalysisManager cgam; + ModuleAnalysisManager mam; + pb.registerModuleAnalyses(mam); + pb.registerCGSCCAnalyses(cgam); + pb.registerFunctionAnalyses(fam); + pb.registerLoopAnalyses(lam); + pb.crossRegisterProxies(lam, fam, cgam, mam); + + ModulePassManager mpm; + llvm::FunctionPassManager fpm; + if (BreakStructPhiNodes) + fpm.addPass(BreakStructPhiNodesPass()); + mpm.addPass(createModuleToFunctionPassAdaptor(std::move(fpm))); + mpm.run(*m, mam); + return Error::success(); + }; +} +} // namespace + +int main(int argc, char **argv) { + InitLLVM X(argc, argv); + cl::ParseCommandLineOptions( + argc, argv, "llvm .bc -> .bc modular optimizer and analysis printer\n"); + + LLVMContext Context; + SMDiagnostic Err; + + // Load the input module... + auto SetDataLayout = [](StringRef, StringRef) -> std::optional { + if (ClDataLayout.empty()) + return std::nullopt; + return ClDataLayout; + }; + std::unique_ptr M; + M = parseIRFile(InputFilename, Err, Context, ParserCallbacks(SetDataLayout)); + if (!M) { + Err.print(argv[0], errs()); + return 1; + } + // If we are supposed to override the target triple or data layout, do so now. + if (!TargetTriple.empty()) + M->setTargetTriple(Triple(Triple::normalize(TargetTriple))); + auto optPipeline = makeOptimizingPipeline(); + if (auto err = optPipeline(M.get())) { + llvm::errs() << "Failed to optimize LLVM IR " << err << "\n"; + } + + if (verifyModule(*M, &errs())) { + errs() << argv[0] << ": " << InputFilename + << ": error: input module is broken!\n"; + return 1; + } + + // Write to standard output. + std::unique_ptr Out; + // Default to standard output. + if (OutputFilename.empty()) + OutputFilename = "-"; + std::error_code EC; + sys::fs::OpenFlags Flags = sys::fs::OF_TextWithCRLF; + Out.reset(new ToolOutputFile(OutputFilename, EC, Flags)); + if (EC) { + errs() << EC.message() << '\n'; + return 1; + } + Out->os() << *M << "\n"; + Out->keep(); + return 0; +} diff --git a/third_party/ascend/bin/triton-lsp.cpp b/third_party/ascend/bin/triton-lsp.cpp new file mode 100644 index 0000000000..f95036dc6c --- /dev/null +++ b/third_party/ascend/bin/triton-lsp.cpp @@ -0,0 +1,10 @@ +#include "./RegisterTritonDialects.h" + +#include "mlir/Tools/mlir-lsp-server/MlirLspServerMain.h" + +int main(int argc, char **argv) { + mlir::DialectRegistry registry; + registerTritonDialects(registry); + + return mlir::failed(mlir::MlirLspServerMain(argc, argv, registry)); +} diff --git a/third_party/ascend/bin/triton-mlir-opt.cpp b/third_party/ascend/bin/triton-mlir-opt.cpp new file mode 100644 index 0000000000..7658424a02 --- /dev/null +++ b/third_party/ascend/bin/triton-mlir-opt.cpp @@ -0,0 +1,45 @@ +//===- triton-mlir-opt.cpp - Triton-Ascend Optimizer Driver -----------*- C++ +//-*-===// +// +// Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +//===----------------------------------------------------------------------===// +// +// Main entry function for triton-mlir-opt built as standalone binary. +// This tool combines MLIR dialects and BishengIR dialects. +// +//===----------------------------------------------------------------------===// + +#include "./RegisterTritonDialects.h" + +// BishengIR includes +#include "bishengir/InitAllDialects.h" + +#include "mlir/InitAllDialects.h" +#include "mlir/InitAllExtensions.h" +#include "mlir/InitAllPasses.h" +#include "mlir/Tools/mlir-opt/MlirOptMain.h" +#include "llvm/Support/InitLLVM.h" + +int main(int argc, char **argv) { + // Register all MLIR dialects and passes + mlir::DialectRegistry registry; + mlir::registerAllDialects(registry); + + // Register BishengIR dialects and passes + bishengir::registerAllDialects(registry); + + return mlir::asMainReturnCode(mlir::MlirOptMain( + argc, argv, "Triton-Ascend optimizer driver\n", registry)); +} diff --git a/third_party/ascend/bin/triton-opt.cpp b/third_party/ascend/bin/triton-opt.cpp new file mode 100644 index 0000000000..91ed8b79ae --- /dev/null +++ b/third_party/ascend/bin/triton-opt.cpp @@ -0,0 +1,15 @@ +#include "./RegisterTritonDialects.h" +#include "ascend/include/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.h" + +#include "mlir/Tools/mlir-opt/MlirOptMain.h" + +int main(int argc, char **argv) { + mlir::DialectRegistry registry; + registerTritonDialects(registry); + + // Register AddMultiBufferInnerScope pass + mlir::triton::registerAddMultiBufferInnerScopePasses(); + + return mlir::asMainReturnCode(mlir::MlirOptMain( + argc, argv, "Triton (GPU) optimizer driver\n", registry)); +} diff --git a/third_party/ascend/bin/triton-reduce.cpp b/third_party/ascend/bin/triton-reduce.cpp new file mode 100644 index 0000000000..8235f8fc8c --- /dev/null +++ b/third_party/ascend/bin/triton-reduce.cpp @@ -0,0 +1,11 @@ +#include "./RegisterTritonDialects.h" + +#include "mlir/Tools/mlir-reduce/MlirReduceMain.h" + +int main(int argc, char **argv) { + mlir::DialectRegistry registry; + registerTritonDialects(registry); + + mlir::MLIRContext context(registry); + return mlir::failed(mlir::mlirReduceMain(argc, argv, context)); +} diff --git a/third_party/ascend/bin/triton-tensor-layout.cpp b/third_party/ascend/bin/triton-tensor-layout.cpp new file mode 100644 index 0000000000..7b82fad6cc --- /dev/null +++ b/third_party/ascend/bin/triton-tensor-layout.cpp @@ -0,0 +1,234 @@ +#include "RegisterTritonDialects.h" + +#include "mlir/AsmParser/AsmParser.h" +#include "mlir/AsmParser/AsmParserState.h" +#include "mlir/IR/MLIRContext.h" + +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" + +#include "llvm/Support/CommandLine.h" +#include "llvm/Support/ErrorOr.h" +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/MemoryBuffer.h" +#include "llvm/Support/SourceMgr.h" +#include "llvm/Support/raw_ostream.h" + +using namespace llvm; +using namespace mlir; + +// A CLI tool to print the layout of a tensor. +// +// clang-format off +// Example usage: +// +// triton-tensor-layout -l "#ttg.nvidia_mma<{versionMajor = 3, versionMinor = 0, warpsPerCTA = [8, 1], CTAsPerCGA = [1, 1], CTASplitNum = [1, 1], CTAOrder = [1, 0], instrShape = [16, 256, 32]}>" -t "tensor<128x256xf16>" +// +// triton-tensor-layout -i input.mlir -t "tensor<1x128x128xf16>" -o output.txt +// +// triton-tensor-layout -i input.mlir -t "tensor<1x128x128xf16>" -o output.txt -alias-names="blocked,mma" -use-hw-view +// +// An input file usually looks like: +// ''' +// #mma = #ttg.amd_mfma<{versionMajor = 2, versionMinor = 0, warpsPerCTA = [1, 1, 8], instrShape = [32, 32], isTransposed = false}> +// #blocked = #ttg.blocked<{sizePerThread = [1, 8, 1], threadsPerWarp = [1, 16, 4], warpsPerCTA = [1, 1, 8], order = [0, 1, 2]}> +// ''' +// clang-format on + +//===--------------------------------------------------------------------===// +// CLI options +//===--------------------------------------------------------------------===// + +cl::OptionCategory PrinterCategory("Available Print Options", + "Options for the tensor layout printing."); + +static cl::opt InputFile( + "i", cl::desc("File that contains the tensor data layout attributes"), + cl::init(""), cl::value_desc("filename"), cl::cat(PrinterCategory)); + +static cl::opt + OutputFile("o", cl::desc("Output file to write the layout into"), + cl::init(""), cl::value_desc("filename"), + cl::cat(PrinterCategory)); + +static cl::opt + DataLayoutStr("l", cl::desc("Tensor data layout attribute in string"), + cl::value_desc("layout-string"), cl::init(""), + cl::cat(PrinterCategory)); + +static cl::list + AliasName("alias-names", + cl::desc("A list of alias names (separated by comma) of the " + "layout attributes in the input file"), + cl::value_desc("name1,name2,name3,..."), cl::CommaSeparated, + cl::ZeroOrMore, cl::cat(PrinterCategory)); + +static cl::opt UseHWPointOfView( + "use-hw-view", + llvm::cl::desc( + "Print the layout in hardware point of view. This means the output is " + "from the warp's perspective. Otherwise, the output is from the " + "tensor's perspective (e.g., each element maps to xxx thread)."), + cl::init(false), cl::cat(PrinterCategory)); + +static cl::opt TensorStr( + "t", cl::desc("Tensor shape and element type (e.g., tensor<2x2xf32>)"), + cl::init(""), cl::value_desc("tensor-type"), cl::cat(PrinterCategory)); + +//===--------------------------------------------------------------------===// +// Helper functions +//===--------------------------------------------------------------------===// + +static LogicalResult layoutPrint(RankedTensorType tensorType, raw_ostream &os) { + // DistributedEncodingTrait and SharedEncodingTrait implements the + // toLinearLayout interface. + mlir::Attribute layout = tensorType.getEncoding(); + if (isa(layout)) { + os << triton::gpu::getLayoutStr(tensorType, UseHWPointOfView); + return success(); + } + + llvm::errs() << "Unsupported tensor layout attribute: " + << tensorType.getEncoding() << "\n"; + return failure(); +} + +static LogicalResult printLayoutFromFile(MLIRContext *context, + StringRef filename, + ArrayRef names, + TensorType tensorTy, + raw_string_ostream &ss) { + if (filename.empty()) + return success(); + + llvm::ErrorOr> fileOrErr = + llvm::MemoryBuffer::getFileOrSTDIN(filename); + if (std::error_code ec = fileOrErr.getError()) { + llvm::errs() << "Could not open input file: " << ec.message() << "\n"; + return failure(); + } + + llvm::SourceMgr sourceMgr; + sourceMgr.AddNewSourceBuffer(std::move(*fileOrErr), llvm::SMLoc()); + ParserConfig config(context); + auto asmState = AsmParserState(); + + Block parsedIR; + if (failed(parseAsmSourceFile(sourceMgr, &parsedIR, config, &asmState))) { + llvm::errs() << "Fail to parse the input file: " << filename << "\n"; + return failure(); + } + + auto printLambda = [&](StringRef name, mlir::Attribute attr) { + ss << "Print layout attribute: #" << name << " = " << attr << "\n"; + + auto rankedTensorTy = RankedTensorType::get( + tensorTy.getShape(), tensorTy.getElementType(), attr); + + return layoutPrint(rankedTensorTy, ss); + }; + + if (names.empty()) + // If no alias name is given, we print all layout attributes in the file. + for (const auto &def : asmState.getAttributeAliasDefs()) { + if (failed(printLambda(def.name, def.value))) + return failure(); + } + else { + // Print the layout attributes with the given alias names. + for (const auto &alias : names) { + auto def = asmState.getAttributeAliasDef(alias); + if (!def) { + llvm::errs() << "Can't find the layout attribute: " << alias << "\n"; + return failure(); + } + + if (failed(printLambda(alias, def->value))) + return failure(); + + ss << "\n"; + } + } + + return success(); +} + +static LogicalResult printLayoutFromString(MLIRContext *context, + StringRef layoutAttrStr, + TensorType tensorTy, + raw_string_ostream &ss) { + if (layoutAttrStr.empty()) + return success(); + + mlir::Attribute layout = parseAttribute(layoutAttrStr, context); + if (!layout) { + llvm::errs() << "Invalid layout attribute: " << layoutAttrStr << "\n"; + return failure(); + } + + auto rankedTensorTy = RankedTensorType::get( + tensorTy.getShape(), tensorTy.getElementType(), layout); + + ss << "Print layout attribute: " << layout << "\n"; + + return layoutPrint(rankedTensorTy, ss); +} + +//===--------------------------------------------------------------------===// +// Main entry point +//===--------------------------------------------------------------------===// + +int main(int argc, char **argv) { + cl::HideUnrelatedOptions(PrinterCategory); + cl::ParseCommandLineOptions(argc, argv, "tensor layout printer\n"); + + DialectRegistry registry; + registerTritonDialects(registry); + + MLIRContext ctx(registry); + ctx.loadAllAvailableDialects(); + + if (TensorStr.empty()) { + llvm::errs() << "Must specify the tensor type argument\n"; + return 1; + } + + mlir::Type parsedTy = parseType(TensorStr, &ctx); + if (!parsedTy) { + llvm::errs() << "Fail to parse the tensor type argument: " << TensorStr + << "\n"; + return 1; + } + + TensorType tensorType = dyn_cast(parsedTy); + if (!tensorType) { + llvm::errs() << "Invalid tensor type argument: " << TensorStr << "\n"; + return 1; + } + + std::string storage; + raw_string_ostream ss(storage); + + if (failed(printLayoutFromFile(&ctx, InputFile, AliasName, tensorType, ss))) + return 1; + + if (failed(printLayoutFromString(&ctx, DataLayoutStr, tensorType, ss))) + return 1; + + if (OutputFile.empty()) { + llvm::outs() << ss.str(); + } else { + std::error_code ec; + llvm::raw_fd_ostream outFs(OutputFile, ec, llvm::sys::fs::OF_Text); + if (ec) { + llvm::errs() << "Error: " << ec.message() << " : unable to open " + << OutputFile << " for output\n"; + return 1; + } + outFs << ss.str(); + outFs.close(); + } + + return 0; +} diff --git a/third_party/ascend/include/AutoBlockify/AutoBlockify.h b/third_party/ascend/include/AutoBlockify/AutoBlockify.h new file mode 100644 index 0000000000..3ce8159ef9 --- /dev/null +++ b/third_party/ascend/include/AutoBlockify/AutoBlockify.h @@ -0,0 +1,119 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#pragma once + +#include "mlir/Pass/Pass.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "mlir/IR/PatternMatch.h" + +#define GEN_PASS_DECL_AUTOBLOCKIFY +#include "ascend/include/AutoBlockify/Passes.h.inc" + +#define GEN_PASS_DEF_AUTOBLOCKIFY +#include "ascend/include/AutoBlockify/Passes.h.inc" + +namespace mlir { +namespace triton { + +std::unique_ptr> +createAutoBlockifyPass(const AutoBlockifyOptions &options = {}); + +} // namespace triton +} // namespace mlir + +using namespace mlir; +using namespace triton; + +class PropagateUnrealizedCastDown + : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + + explicit PropagateUnrealizedCastDown(MLIRContext *context, + Value logicalBlockId, + Value logicalBlockNum, + int autoBlockifySize); + + LogicalResult matchAndRewrite(UnrealizedConversionCastOp op, + PatternRewriter &rewriter) const override; + +private: + void handleBlockifyLoop(scf::ForOp blockifyLoop, Operation *op, + PatternRewriter &rewriter) const; + void rewriteSplat(UnrealizedConversionCastOp op, triton::SplatOp splatOp, + PatternRewriter &rewriter) const; + void rewriteExpandDims(UnrealizedConversionCastOp op, + triton::ExpandDimsOp expandDimsOp, + PatternRewriter &rewriter) const; + void rewriteReduce(UnrealizedConversionCastOp op, triton::ReduceOp reduceOp, + PatternRewriter &rewriter) const; + void rewriteScan(UnrealizedConversionCastOp op, triton::ScanOp scanOp, + PatternRewriter &rewriter) const; + void rewriteLoad(UnrealizedConversionCastOp op, triton::LoadOp loadOp, + PatternRewriter &rewriter) const; + void rewriteStore(UnrealizedConversionCastOp op, triton::StoreOp storeOp, + PatternRewriter &rewriter) const; + void rewriteAtomicRMW(UnrealizedConversionCastOp op, + triton::AtomicRMWOp atomicRMWOp, + PatternRewriter &rewriter) const; + void rewriteAssert(UnrealizedConversionCastOp op, triton::AssertOp assertOp, + PatternRewriter &rewriter) const; + void rewriteExtractSlice(UnrealizedConversionCastOp op, + tensor::ExtractSliceOp extractSliceOp, + PatternRewriter &rewriter) const; + void rewriteInsertSlice(UnrealizedConversionCastOp op, + tensor::InsertSliceOp insertSliceOp, + PatternRewriter &rewriter) const; + void rewriteWhile(UnrealizedConversionCastOp op, scf::WhileOp whileOp, + PatternRewriter &rewriter) const; + void rewriteLoop(UnrealizedConversionCastOp op, LoopLikeOpInterface loopOp, + PatternRewriter &rewriter) const; + void rewriteIf(UnrealizedConversionCastOp &op, scf::IfOp ifOp, + ArrayRef indices, PatternRewriter &rewriter) const; + void rewriteYield(UnrealizedConversionCastOp &op, scf::YieldOp yieldOp, + PatternRewriter &rewriter) const; + void rewriteCondition(UnrealizedConversionCastOp op, + scf::ConditionOp conditionOp, + PatternRewriter &rewriter) const; + void rewriteGeneraleOp(UnrealizedConversionCastOp op, Operation *generalOp, + PatternRewriter &rewriter) const; + + Value logicalBlockId; + Value logicalBlockNum; + int autoBlockifySize; +}; + +class AutoBlockifyPass : public ::impl::AutoBlockifyBase { +public: + explicit AutoBlockifyPass(const AutoBlockifyOptions &options); + void runOnOperation() override; + +private: + bool checkBlockifiable(Value v); + void preProcess(triton::FuncOp func); + + DenseSet checkedValues; + Value logicalBlockId; + Value logicalBlockNum; +}; diff --git a/third_party/ascend/include/AutoBlockify/CMakeLists.txt b/third_party/ascend/include/AutoBlockify/CMakeLists.txt new file mode 100644 index 0000000000..ca4cf9f552 --- /dev/null +++ b/third_party/ascend/include/AutoBlockify/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name AutoBlockify) +add_public_tablegen_target(AutoBlockifyPassIncGen) diff --git a/third_party/ascend/include/AutoBlockify/Passes.h b/third_party/ascend/include/AutoBlockify/Passes.h new file mode 100644 index 0000000000..7d5147ef92 --- /dev/null +++ b/third_party/ascend/include/AutoBlockify/Passes.h @@ -0,0 +1,37 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_AUTO_BLOCKIFY_PASSES_H +#define TRITON_ADAPTER_AUTO_BLOCKIFY_PASSES_H + +#include "AutoBlockify.h" + +namespace mlir { +namespace triton { + +#define GEN_PASS_REGISTRATION +#include "ascend/include/AutoBlockify/Passes.h.inc" + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_AUTO_BLOCKIFY_PASSES_H diff --git a/third_party/ascend/include/AutoBlockify/Passes.td b/third_party/ascend/include/AutoBlockify/Passes.td new file mode 100644 index 0000000000..7d9f1a80a3 --- /dev/null +++ b/third_party/ascend/include/AutoBlockify/Passes.td @@ -0,0 +1,21 @@ +#ifndef AUTO_BLOCKIFY_PASSES +#define AUTO_BLOCKIFY_PASSES + +include "mlir/Pass/PassBase.td" + +def AutoBlockify : Pass<"auto-blockify", "mlir::ModuleOp"> { + let summary = "Apply auto blockify v2"; + let constructor = "triton::createAutoBlockifyPass()"; + let dependentDialects = [ + "mlir::arith::ArithDialect", + "mlir::tensor::TensorDialect", + "mlir::triton::TritonDialect" + ]; + let options = [ + Option<"autoBlockifySize", "auto-blockify-size", "int", "1", + "Apply auto blockify v2 when TRITON_ALL_BLOCKS_PARALLEL is 1." + "Expand highest dimension with blockify size"> + ]; +} + +#endif // AUTO_BLOCKIFY_PASSES diff --git a/third_party/ascend/include/AutoBlockify/Utils.h b/third_party/ascend/include/AutoBlockify/Utils.h new file mode 100644 index 0000000000..385fa51a10 --- /dev/null +++ b/third_party/ascend/include/AutoBlockify/Utils.h @@ -0,0 +1,65 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#pragma once + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +using namespace mlir; +using namespace triton; + +constexpr llvm::StringLiteral autoBlockifySizeAttr = "auto_blockify_size"; +constexpr llvm::StringLiteral logicalBlockIdAttr = "logical_block_id"; +constexpr llvm::StringLiteral autoBlockifyLoopAttr = "auto_blockify_loop"; +constexpr llvm::StringLiteral autoBlockifyRegionOpAttr = + "auto_blockify_region_op"; + +RankedTensorType getExpandedType(Type type, UnrealizedConversionCastOp op); + +Value rewriteValue(Value value, UnrealizedConversionCastOp op, + OpBuilder &builder); + +void replaceValue(Operation *newOp, Operation *oldOp, Value newMask, + RewriterBase &rewriter, + ArrayRef replaceIndices = {}); + +Value createMask(Value mask, Value uccMask, ArrayRef targetShape, + RewriterBase &rewriter); + +void mapRegionIterArg(IRMapping &mapping, ValueRange oldArgs, + ValueRange newArgs, ArrayRef indices, Value mask, + OpBuilder &builder); + +void mapYieldedValue(IRMapping &mapping, scf::YieldOp yieldOp, + ArrayRef indices, UnrealizedConversionCastOp op, + OpBuilder &builder); + +Operation *createBlockifyLoop(Operation *targetOp, + UnrealizedConversionCastOp op, + Value logicalBlockId, Value logicalBlockNum, + int autoBlockifySize, RewriterBase &rewriter); + +std::optional getBlockifyLoop(Operation *op); diff --git a/third_party/ascend/include/CMakeLists.txt b/third_party/ascend/include/CMakeLists.txt new file mode 100644 index 0000000000..7925432eed --- /dev/null +++ b/third_party/ascend/include/CMakeLists.txt @@ -0,0 +1,11 @@ +add_subdirectory(Dialect) +add_subdirectory(TritonToAnnotation) +add_subdirectory(TritonToHFusion) +add_subdirectory(TritonToHIVM) +add_subdirectory(TritonToGraph) +add_subdirectory(TritonToLLVM) +add_subdirectory(CommonIRToHIVM) +add_subdirectory(AutoBlockify) +# todo: this code will be removed in version 530. +add_subdirectory(TritonAffinityOpt) +add_subdirectory(DynamicCVPipeline) diff --git a/third_party/ascend/include/CommonIRToHIVM/CMakeLists.txt b/third_party/ascend/include/CommonIRToHIVM/CMakeLists.txt new file mode 100644 index 0000000000..1543d859b3 --- /dev/null +++ b/third_party/ascend/include/CommonIRToHIVM/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name CommonIRToHIVM) +add_public_tablegen_target(CommonIRToHIVMConversionPassIncGen) diff --git a/third_party/ascend/include/CommonIRToHIVM/Passes.h b/third_party/ascend/include/CommonIRToHIVM/Passes.h new file mode 100644 index 0000000000..a1e5be1179 --- /dev/null +++ b/third_party/ascend/include/CommonIRToHIVM/Passes.h @@ -0,0 +1,49 @@ +/* + * Copyright 2018-2020 Philippe Tillet + * Copyright 2020-2022 OpenAI + * Copyright 2025- FlagOS Contributors + * + * Permission is hereby granted, free of charge, to any person obtaining + * a copy of this software and associated documentation files + * (the "Software"), to deal in the Software without restriction, + * including without limitation the rights to use, copy, modify, merge, + * publish, distribute, sublicense, and/or sell copies of the Software, + * and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be + * included in all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, + * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF + * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. + * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY + * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, + * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE + * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_TILE_IR_TO_HIVM_CONVERSION_PASSES_H +#define TRITON_ADAPTER_TILE_IR_TO_HIVM_CONVERSION_PASSES_H + +#include "mlir/Pass/Pass.h" + +namespace mlir { +// Forward declarations. +class ModuleOp; + +namespace triton { + +/// Creates a pass to convert CommonIR dialect to HIVM dialect. +std::unique_ptr> createCommonIRToHIVMPass(); + +/// Creates a pass to reduce sum strength for profitable add reductions. +std::unique_ptr> createReduceSumStrengthPass(); + +#define GEN_PASS_REGISTRATION +#include "ascend/include/CommonIRToHIVM/Passes.h.inc" + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_TILE_IR_TO_HIVM_CONVERSION_PASSES_H diff --git a/third_party/ascend/include/CommonIRToHIVM/Passes.td b/third_party/ascend/include/CommonIRToHIVM/Passes.td new file mode 100644 index 0000000000..dedfb46258 --- /dev/null +++ b/third_party/ascend/include/CommonIRToHIVM/Passes.td @@ -0,0 +1,66 @@ +/* + * Copyright 2018-2020 Philippe Tillet + * Copyright 2020-2022 OpenAI + * Copyright 2025- FlagOS Contributors + * + * Permission is hereby granted, free of charge, to any person obtaining + * a copy of this software and associated documentation files + * (the "Software"), to deal in the Software without restriction, + * including without limitation the rights to use, copy, modify, merge, + * publish, distribute, sublicense, and/or sell copies of the Software, + * and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be + * included in all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, + * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF + * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. + * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY + * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, + * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE + * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ + +#ifndef TILE_IR_TO_HIVM_CONVERSION_PASSES +#define TILE_IR_TO_HIVM_CONVERSION_PASSES + +include "mlir/Pass/PassBase.td" + +def CommonIRToHIVM : Pass<"commonir-to-hivm", "mlir::ModuleOp"> { + let summary = "Convert CommonIR dialect to HIVM dialect"; + let constructor = "triton::createCommonIRToHIVMPass()"; + let dependentDialects = [ + "hivm::HIVMDialect", + "triton::tile::CommonIRDialect", + "mlir::bufferization::BufferizationDialect", + "mlir::memref::MemRefDialect", + "mlir::tensor::TensorDialect", + "mlir::triton::TritonDialect", + ]; +} + +def ReduceSumStrength : Pass<"reduce-sum-strength", "mlir::ModuleOp"> { + let summary = "Reduce sum strength for profitable tl.sum-style add reductions"; + let description = [{ + This pass rewrites profitable add reductions by pre-adding two equal halves + of the reduced dimension and reducing the half-sized tensor. It is owned by + the Ascend ttir-to-ttadapter pipeline for backends where a full-width reduction + is more expensive than an elementwise add plus a smaller reduction. + }]; + let constructor = "triton::createReduceSumStrengthPass()"; + + let dependentDialects = ["mlir::arith::ArithDialect", + "mlir::tensor::TensorDialect", + "mlir::triton::TritonDialect"]; + + let options = [ + Option<"enable", "enable", "bool", /*default*/"true", + "Enable the rewrite">, + Option<"splitFactor", "split-factor", "int32_t", /*default*/"2", + "Reduction-axis split factor; currently only 2 is supported"> + ]; +} + +#endif // TILE_IR_TO_HIVM_CONVERSION_PASSES diff --git a/third_party/ascend/include/CommonIRToHIVM/ReduceSumStrengthPass.h b/third_party/ascend/include/CommonIRToHIVM/ReduceSumStrengthPass.h new file mode 100644 index 0000000000..67a7bd94f0 --- /dev/null +++ b/third_party/ascend/include/CommonIRToHIVM/ReduceSumStrengthPass.h @@ -0,0 +1,26 @@ +/* + * Copyright 2026- Xcoresigma Technology Co., Ltd + */ + +#ifndef TRITON_ADAPTER_COMMONIR_TO_HIVM_REDUCESUMSTRENGTHPASS_H +#define TRITON_ADAPTER_COMMONIR_TO_HIVM_REDUCESUMSTRENGTHPASS_H + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +#include +#include +#include + +namespace mlir { +namespace triton { + +std::unique_ptr> createReduceSumStrengthPass(); + +std::unique_ptr> +createReduceSumStrengthPass(bool enable, int32_t splitFactor); + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_COMMONIR_TO_HIVM_REDUCESUMSTRENGTHPASS_H diff --git a/third_party/ascend/include/Dialect/CMakeLists.txt b/third_party/ascend/include/Dialect/CMakeLists.txt new file mode 100644 index 0000000000..b7e6595642 --- /dev/null +++ b/third_party/ascend/include/Dialect/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(TritonAscend) diff --git a/third_party/ascend/include/Dialect/TritonAscend/CMakeLists.txt b/third_party/ascend/include/Dialect/TritonAscend/CMakeLists.txt new file mode 100644 index 0000000000..f33061b2d8 --- /dev/null +++ b/third_party/ascend/include/Dialect/TritonAscend/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(IR) diff --git a/third_party/ascend/include/Dialect/TritonAscend/IR/CMakeLists.txt b/third_party/ascend/include/Dialect/TritonAscend/IR/CMakeLists.txt new file mode 100644 index 0000000000..488e9132f6 --- /dev/null +++ b/third_party/ascend/include/Dialect/TritonAscend/IR/CMakeLists.txt @@ -0,0 +1,15 @@ +set(MLIR_BINARY_DIR ${CMAKE_BINARY_DIR}) + +set(LLVM_TARGET_DEFINITIONS TritonAscendOps.td) +mlir_tablegen(TritonAscendDialect.h.inc -gen-dialect-decls -dialect=ascend) +mlir_tablegen(TritonAscendDialect.cpp.inc -gen-dialect-defs -dialect=ascend) +mlir_tablegen(TritonAscendOps.h.inc -gen-op-decls) +mlir_tablegen(TritonAscendOps.cpp.inc -gen-op-defs) +add_mlir_doc(TritonAscendDialect TritonAscendDialect dialects/ -gen-dialect-doc) +add_mlir_doc(TritonAscendOps TritonAscendOps dialects/ -gen-op-doc) +add_public_tablegen_target(TritonAscendTableGen) + +set(LLVM_TARGET_DEFINITIONS TritonAscendAttrDefs.td) +mlir_tablegen(TritonAscendOpsAttrDefs.h.inc -gen-attrdef-decls) +mlir_tablegen(TritonAscendOpsAttrDefs.cpp.inc -gen-attrdef-defs) +add_public_tablegen_target(TritonAscendAttrDefsIncGen) diff --git a/third_party/ascend/include/Dialect/TritonAscend/IR/TritonAscendAttrDefs.td b/third_party/ascend/include/Dialect/TritonAscend/IR/TritonAscendAttrDefs.td new file mode 100644 index 0000000000..7ce5eefb80 --- /dev/null +++ b/third_party/ascend/include/Dialect/TritonAscend/IR/TritonAscendAttrDefs.td @@ -0,0 +1,25 @@ +//===-- TritonAscendAttrDefs.td - dialect attributes def. ----*- tablegen -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef TRITON_ASCEND_ATTRDEFS +#define TRITON_ASCEND_ATTRDEFS + +include "ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.td" + +include "mlir/IR/AttrTypeBase.td" +include "mlir/IR/EnumAttr.td" + +class TritonAscend_Attr traits = []> + : AttrDef { + let mnemonic = attrMnemonic; + let cppNamespace = "::mlir::triton::ascend"; +} + + + +#endif // TRITON_ASCEND_ATTRDEFS diff --git a/third_party/ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.h b/third_party/ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.h new file mode 100644 index 0000000000..b2902d2c20 --- /dev/null +++ b/third_party/ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.h @@ -0,0 +1,34 @@ +//===- TritonAscendDialect.h - MLIR TritonAscend dialect --------------*- C++ +//-*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// This file defines the TritonAscend dialect in MLIR, containing Ascend +// operations. +// +//===----------------------------------------------------------------------===// + +#ifndef TRITON_DIALECT_ASCEND_DIALECT_H +#define TRITON_DIALECT_ASCEND_DIALECT_H + +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/IR/Dialect.h" +#include "mlir/IR/OpDefinition.h" + +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.h.inc" + +#define GET_ATTRDEF_CLASSES +#include "ascend/include/Dialect/TritonAscend/IR/TritonAscendOpsAttrDefs.h.inc" + +#define GET_OP_CLASSES +#include "ascend/include/Dialect/TritonAscend/IR/TritonAscendOps.h.inc" + +namespace mlir::triton::ascend {} // namespace mlir::triton::ascend + +#endif // TRITON_DIALECT_ASCEND_DIALECT_H diff --git a/third_party/ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.td b/third_party/ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.td new file mode 100644 index 0000000000..a7314eff6d --- /dev/null +++ b/third_party/ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.td @@ -0,0 +1,31 @@ +//===-- TritonAscendDialect.td - dialect op definitions -------*- tablegen -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef TRITON_ASCEND_DIALECT +#define TRITON_ASCEND_DIALECT + +include "mlir/IR/OpBase.td" + +def TritonAscend_Dialect : Dialect { + let name = "ascend"; + let cppNamespace = "::mlir::triton::ascend"; + let summary = "The TritonAscend dialect in Triton."; + + let description = [{ + TritonAscend is a dialect for representing operations on Ascend NPUs. + }]; + + let dependentDialects = [ + "mlir::LLVM::LLVMDialect", + "triton::TritonDialect", + ]; + + let extraClassDeclaration = [{}]; +} + +#endif // TRITON_ASCEND_DIALECT diff --git a/third_party/ascend/include/Dialect/TritonAscend/IR/TritonAscendOps.td b/third_party/ascend/include/Dialect/TritonAscend/IR/TritonAscendOps.td new file mode 100644 index 0000000000..f98d132975 --- /dev/null +++ b/third_party/ascend/include/Dialect/TritonAscend/IR/TritonAscendOps.td @@ -0,0 +1,473 @@ +//===-- TritonAscendOps.td - TritonAscend op definitions ---------*- tablegen -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// This is the TritonAscend IR operation definition file. +// +//===----------------------------------------------------------------------===// + +#ifndef TRITON_ASCEND_OPS +#define TRITON_ASCEND_OPS + +include "mlir/IR/OpBase.td" +include "mlir/IR/EnumAttr.td" +include "mlir/Dialect/LLVMIR/LLVMTypes.td" +include "mlir/Interfaces/SideEffectInterfaces.td" +include "mlir/IR/OpAsmInterface.td" +include "mlir/Interfaces/InferTypeOpInterface.td" // SameOperandsAndResultType + +include "triton/Dialect/Triton/IR/TritonAttrDefs.td" +include "triton/Dialect/Triton/IR/TritonTypes.td" +include "triton/Dialect/Triton/IR/TritonInterfaces.td" + +include "ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.td" +include "ascend/include/Dialect/TritonAscend/IR/TritonAscendAttrDefs.td" + +//===----------------------------------------------------------------------===// +// TritonAscend op definitions +//===----------------------------------------------------------------------===// + +class TT_Ascend_Op traits = []> : + Op; + + +// +// Interfaces +// +def GlobalMemory : Resource<"::mlir::triton::GlobalMemory">; + + +// +// Annotation Op +// +def AnnotationOp : TT_Ascend_Op<"annotation", [Pure, MemoryEffects<[MemWrite]>]> { + let summary = "Annotate a tensor with key-value attribute pairs"; + let description = [{ + `ascend.annotation` operation can be used to annotate a tensor with + key-value attribute pairs. + + Example: + ```mlir + ascend.annotation %target {key : val} + ``` + }]; + let arguments = (ins TT_Tensor:$src); + let assemblyFormat = [{ + $src attr-dict `:` type($src) + }]; +} + + +// +// Mod Op +// +def ModOp : TT_Ascend_Op<"mod", [Pure]> { + let summary = "Mod operation (%) of input tensors."; + let description = [{ + Performs element-wise division with remainder of input tensors. + }]; + + let arguments = (ins TT_Tensor:$lhs, TT_Tensor:$rhs); + let results = (outs TT_Tensor:$result); + + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) type($rhs) `->` type($result)"; +} + + +// +// IndexPut Op +// +def IndexPutOp : TT_Ascend_Op<"index_put", [ + MemoryEffects<[MemWrite]>, + SameVariadicOperandSize, +]> { + let summary = "Scatter store to a tensor pointer with embedding semantics"; + + let description = [{ + Index put values from a tensor into a destination tensor. + + The operation takes: + - ptr: pointer type, the destination tensor pointer (in GM) + - index: tensor, a index to scatter (in UB) + - value: tensor, a value to store (in UB) + - dim: int32, the dimension to scatter along + - index_boundary: int64, the upper boundary for index values + - end_offset: tuple of int, the offsets of each dimension for the end of the scatter region + - start_offset: tuple of int, the offsets of each dimension for the start of the scatter region + - dst_stride: tuple of int, the stride of each dimension of destination tensor + + + Constraints: + - `ptr` and `value` must have the same rank. + - `ptr.dtype` only supports `float16`, `bfloat16`, `float32` currently. + - `index` must be an integer tensor. If `index.rank` != 1, it will be reshaped to 1D. + - `index.numel` must equal `value.shape[dim]`. + - `value` support 2~5D tensors. + - `dim` must be valid (0 <= dim < rank(value) - 1). + }]; + + let arguments = ( + ins TT_Ptr:$ptr, + TT_Tensor:$index, + TT_Tensor:$value, + TT_Int:$dim, + TT_Int:$indexBoundary, + Variadic>:$endOffset, + Variadic>:$startOffset, + Variadic>:$dstStride + ); + + let assemblyFormat = [{ + $ptr `:` type($ptr) `,` $index `:` type($index) `,` + $value `:` type($value) `,` $dim `:` type($dim) `,` $indexBoundary `:` type($indexBoundary) `,` + `[` $endOffset `:` type($endOffset) `]` `,` `[` $startOffset `:` type($startOffset) `]` `,` + `[` $dstStride `:` type($dstStride) `]` + attr-dict + }]; +} + +// +// GatherOutToUb Op +// +def GatherOutToUbOp : TT_Ascend_Op<"gather_out_to_ub", [ + DeclareOpInterfaceMethods, + AttrSizedOperandSegments, +]> { + let summary = "Gather load from a tensor pointer with the embedding semantics"; + + let description = [{ + Gather from a source tensor in Global Memory (GM) to Unified Buffer (UB) + along a specified dimension with out-of-bound handling. + + The operation takes: + - src: pointer type, the source tensor pointer (in GM) + - index: tensor, a tensor to gather (in UB) + - index_boundary: int64, the upper boundary for index values + - dim: int32, the dimension to gather along + - src_stride: tuple of int64, the stride of each dimension of src tensor + - end_offset: tuple of int32, the end offsets of each dimension for index tensor + - start_offset: tuple of int32, the start offsets of each dimension for index tensor + - other(Optional): scalar value, the default value when index is out of boundary (in UB) + + Returns: + a tensor, with the same shape as `index.shape` (in UB) + + Constraints: + - `src` and `index` must have the same rank. + - `src.dtype` only supports `float16`, `bfloat16`, `float32` currently. + - `index` must be an integer tensor, with rank between 1 and 5. + - `dim` must be valid (0 <= dim < rank(index)). + - `other` must be a scalar value. + - For every dimension `i` not equal to `dim`, `index.size[i]` <= `src.size[i]`. + - The output shape is the same as `index.shape`. If `index` is None, \ + the output tensor will be an empty tensor with the same shape as `index`. + }]; + + let arguments = ( + ins TT_Ptr:$src, + TT_Tensor:$index, + TT_Int:$indexBoundary, + TT_Int:$dim, + Variadic>:$srcStride, + Variadic>:$endOffset, + Variadic>:$startOffset, + Optional:$other + ); + + let results = (outs TT_Tensor:$result); + + let assemblyFormat = [{ + $src `:` type($src) `,` $index `:` type($index) `,` + $indexBoundary `:` type($indexBoundary) `,` $dim `:` type($dim) `,` + `[` $srcStride `:` type($srcStride) `]` `,` `[` $endOffset `:` type($endOffset) `]` `,` + `[` $startOffset `:` type($startOffset) `]` (`,` $other^ `:` type($other))? + attr-dict `->` type($result) + }]; +} + +// +// ScatterUbToOut Op +// +def ScatterUbToOutOp : TT_Ascend_Op<"scatter_ub_to_out", [ + MemoryEffects<[MemWrite]>, + SameVariadicOperandSize, +]> { + let summary = "scatter store from a tensor pointer with the embedding semantics"; + + let description = [{ + Scatter a tile from Unified Buffer (UB) into a destination tensor in Global Memory (GM) + along a specified dimension, with index-boundary checking. + + The operation takes: + - ptr: pointer type, the destination tensor pointer (in GM) + - value: tensor, a tile value to store (in UB) + - index: tensor, a tile index to scatter (in UB) + - index_boundary: int, the upper boundary for index values + - dim: int, the dimension to scatter along + - dst_stride: tuple of int, the stride of each dimension of destination tensor + - end_offset: tuple of int32, the end offsets of each dimension for index tensor + - start_offset: tuple of int32, the start offsets of each dimension for index tensor + + Constraints: + - `ptr` and `index` must have the same rank. + - `ptr.dtype` only supports `float16`, `bfloat16`, `float32` currently. + - `index` must be an integer tensor, with rank between 1 and 5. + - `dim` must be valid (0 <= dim < rank(index)). + - For every dimension `i` not equal to `dim`, `index.size[i]` <= `ptr.size[i]`. + - The output shape is the same as `index.shape`. If `index` is None, \ + the output tensor will be an empty tensor with the same shape as `index`. + }]; + + let arguments = ( + ins TT_Ptr:$ptr, + TT_Tensor:$value, + TT_Tensor:$index, + TT_Int:$indexBoundary, + TT_Int:$dim, + Variadic>:$dstStride, + Variadic>:$endOffset, + Variadic>:$startOffset + ); + + let assemblyFormat = [{ + $ptr `:` type($ptr) `,` $value `:` type($value) `,` `,` $index `:` type($index) `,` + $indexBoundary `:` type($indexBoundary) `,` $dim `:` type($dim) `,` + `[` $dstStride `:` type($dstStride) `]` `,` `[` $endOffset `:` type($endOffset) `]` `,` + `[` $startOffset `:` type($startOffset) `]` + attr-dict + }]; +} + + +// +// IndexSelectSimd Op +// +def IndexSelectSimdOp : TT_Ascend_Op<"index_select_simd", [ + MemoryEffects<[MemRead]>, + DeclareOpInterfaceMethods, + AttrSizedOperandSegments +]> { + let summary = "Index select SIMD operation from global memory"; + + let description = [{ + Index select operation (SIMD version) that loads data from multiple indices along a + specified dimension. The operation selects data from GM and loads them + as tiles directly to UB with zero-copy semantics. + + The operation takes: + - src: Source pointer (in GM) + - index: 1D tensor of indices to select (already in UB) + - dim: The dimension along which to select + - src_shape: Complete shape of the source tensor + - src_offset: Starting offset for reading + - read_shape: Size to read (tile shape) + + Constraints: + - read_shape[dim] must be -1 + - src_offset[dim] can be -1 (will be ignored) + }]; + + let arguments = ( + ins + TT_PtrLike:$src, + TT_IntTensor:$index, + I32Attr:$dim, + Variadic:$src_shape, + Variadic:$src_offset, + DenseI32ArrayAttr:$read_shape + ); + + let results = (outs TT_Tensor:$result); + + let assemblyFormat = [{ + $src `,` $index `,` $dim `,` + `[` $src_shape `]` `,` + `[` $src_offset `]` `,` + $read_shape + attr-dict `:` type($src) `,` type($index) `->` type($result) + }]; +} + + +// +// Built-in: IndirectLoad Op +// +def IndirectLoadOp : TT_Ascend_Op<"indirect_load", [ + DeclareOpInterfaceMethods, + AttrSizedOperandSegments +]> { + let summary = "Built-in: indirect load from global memory using per-element offsets with optional mask/other"; + + let description = [{ + Built-in operation emitted by the compiler for unstructured (discrete) memory + accesses.These are not written directly in the user IR. + + Load values from global memory based on per-element offsets. If `mask` + is provided, false lanes return `other`. + + The operation takes: + - src: Source pointer + - offsets: Tensor of per-element offsets (relative to `src`) for accessing source memory + - mask (optional): if mask[idx] is false, do not load the data at address pointer[idx] + - other (optional): if mask[idx] is false, return other[idx] + }]; + + let arguments = ( + ins TT_Ptr:$src, + TT_IntTensor:$offsets, + Optional:$mask, + Optional:$other + ); + + let results = (outs TT_Tensor:$result); + + let assemblyFormat = [{ + $src `:` type($src) `,` + $offsets `:` type($offsets) + (`,` $mask^ `:` type($mask))? + (`,` $other^ `:` type($other))? + attr-dict `->` type($result) + }]; + +} + + +// +// Built-in: IndirectStore Op +// +def IndirectStoreOp : TT_Ascend_Op<"indirect_store", [ + MemoryEffects<[MemWrite]> +]> { + let summary = "Built-in: indirect store from UB using per-element offsets with optional mask/other"; + + let description = [{ + Built-in operation emitted by the compiler for unstructured (discrete) memory + accesses.These are not written directly in the user IR. + + Store values from UB based to GM on per-element offsets. + + The operation takes: + - src: Source pointer + - offsets: Tensor of per-element offsets (relative to `src`) for accessing source memory + - value: The tensor of elements to be stored + - mask (optional): If mask[idx] is false, do not store value[idx] at pointer[idx] + }]; + + let arguments = ( + ins TT_Ptr:$src, + TT_IntTensor:$offsets, + TT_Type:$value, + Optional:$mask + ); + + let assemblyFormat = [{ + $src `:` type($src) `,` + $offsets `:` type($offsets) `,` + $value `:` type($value) + (`,` $mask^ `:` type($mask))? + attr-dict + }]; + +} + +// +// Custom Op +// +def CustomOp : TT_Ascend_Op<"custom", [Pure, MemoryEffects<[MemWrite]>]> { + let summary = "self-defined custom operation"; + let description = [{ + `ascend.custom` triton custom op is designed to pass self-defined custom operation. + + Example: + ```ascend.custom {str_args = ["sync_block_wait", "cube"]} + ``` + }]; + let arguments = (ins StrAttr:$op_name, ArrayAttr:$str_args, Variadic:$args); + + let assemblyFormat = "$op_name attr-dict ($args^ `:` type($args))?"; +} + +def FlipOp : TT_Ascend_Op<"flip", [ + NoMemoryEffect, + DeclareOpInterfaceMethods +]> { + let summary = "Reverse a tensor along a given dimension"; + let description = [{ + Reverses the elements of the input tensor along the specified dimension. + The output tensor has the same shape and element type as the input. + }]; + + let arguments = (ins + TT_Tensor:$src, // Input tensor + I64Attr:$dim // Dimension to flip along + ); + + let results = (outs + TT_Tensor:$flipped // Flipped values + ); + + let assemblyFormat = + "$src `,` $dim attr-dict `:` type($src) `->` type($flipped)"; +} + +def SortOp : TT_Ascend_Op<"sort", [ + NoMemoryEffect, + DeclareOpInterfaceMethods +]> { + let summary = "Sorts a tensor along a given dimension and returns sorted values."; + let description = [{ + Sorts the elements of the input tensor along the specified dimension. + Returns one tensor: + The sorted tensor (same shape and element type as input). + }]; + + let arguments = (ins + TT_Tensor:$src, // Input tensor + I64Attr:$dim, // Dimension to sort along + BoolAttr:$descending // Sort order + ); + + let results = (outs + TT_Tensor:$sorted // Sorted values + ); + + let assemblyFormat = "$src `,` $dim `,` $descending attr-dict `:` type($src) `->` type($sorted)"; +} + +// +// Conv Op +// +def Conv1dOp : TT_Ascend_Op<"conv1d", [ + Pure, + DeclareOpInterfaceMethods +]> { + let summary = "1D convolution operation"; + let description = [{ + Performs a 1D convolution operation on the input tensor. + }]; + + let arguments = (ins + TT_Tensor:$input, + TT_Tensor:$weight, + Optional:$bias, + I64Attr:$stride, + I64Attr:$padding_size, + I64Attr:$dilation, + I64Attr:$groups + ); + + let results = (outs TT_Tensor:$result); + + let assemblyFormat = [{ + `(` operands `)` + attr-dict `:` functional-type(operands, $result) + }]; + + let hasVerifier = 1; +} + +#endif // TRITON_ASCEND_OPS diff --git a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition.h b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition.h new file mode 100644 index 0000000000..ebe1f84e53 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition.h @@ -0,0 +1,71 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_ADD_CONTROLFLOW_CONDITION_PASS_H +#define TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_ADD_CONTROLFLOW_CONDITION_PASS_H + +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Linalg/TransformOps/DialectExtension.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/SmallVector.h" + +namespace mlir { +namespace triton { + +struct ControlFlowConditionInfo { + llvm::DenseMap> blockCounters; + llvm::DenseMap blockCounterNums; + llvm::DenseMap> innerDepConds; + + llvm::DenseMap> crossCoreDependentMap; + llvm::DenseMap>> + intraCoreDependentMap; + + // unique counter value for each ifblock + llvm::DenseMap cntArgs; +}; + +class AddControlFlowConditionPass + : public PassWrapper> { +public: + AddControlFlowConditionPass() = default; + + void runOnOperation() override; + + void getDependentDialects(DialectRegistry ®istry) const override { + registry.insert(); + } + + llvm::StringRef getArgument() const override { + return "add-control-flow-condition"; + } +}; + +std::unique_ptr> createAddControlFlowConditionPass(); + +void registerAddControlFlowConditionPasses(); +} // namespace triton +} // namespace mlir +#endif // TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_ADD_CONTROLFLOW_CONDITION_PASS_H diff --git a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/CloneOps.h b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/CloneOps.h new file mode 100644 index 0000000000..056944c504 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/CloneOps.h @@ -0,0 +1,50 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ASCEND_SSBUF_CLONE_OPS_FOR_CONTROL_FLOW_H +#define TRITON_ASCEND_SSBUF_CLONE_OPS_FOR_CONTROL_FLOW_H +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/DialectRegistry.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace triton { + +class CloneOpsPass : public PassWrapper> { +public: + CloneOpsPass() = default; + + void runOnOperation() override; + + LogicalResult validateBlockIdsConsecutive(ModuleOp module); + LogicalResult cloneOpsInMainLoop(scf::ForOp forOp); + LogicalResult cleanupClonedOpsInMainLoop(scf::ForOp forOp); + + llvm::StringRef getArgument() const override { return "clone-ops"; } +}; + +std::unique_ptr> createCloneOpsPass(); + +} // namespace triton +} // namespace mlir +#endif // TRITON_ASCEND_SSBUF_CLONE_OPS_FOR_CONTROL_FLOW_H diff --git a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.h b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.h new file mode 100644 index 0000000000..6a59ae9de3 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.h @@ -0,0 +1,64 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ASCEND_SSBUF_CREATE_IF_OPS_FOR_CONTROL_FLOW_H +#define TRITON_ASCEND_SSBUF_CREATE_IF_OPS_FOR_CONTROL_FLOW_H +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition.h" + +namespace mlir { +namespace triton { + +class CreateIfOpsPass + : public PassWrapper> { +public: + CreateIfOpsPass() = default; + + void runOnOperation() override; + + void setConditionInfo(ControlFlowConditionInfo *info) { this->info = info; } + + LogicalResult computeYieldValues( + scf::ForOp forOp, + const llvm::DenseMap> &blockOps, + llvm::DenseMap> &thenYieldValues, + llvm::DenseMap> &elseYieldValues); + + LogicalResult createIfInMainLoop( + scf::ForOp forOp, + const llvm::DenseMap> &blockOps, + const llvm::DenseMap> &thenYieldValues, + const llvm::DenseMap> &elseYieldValues); + + llvm::StringRef getArgument() const override { return "create-if-ops"; } + +private: + ControlFlowConditionInfo *info = nullptr; +}; + +std::unique_ptr> createCreateIfOpsPass(); + +} // namespace triton +} // namespace mlir +#endif // TRITON_ASCEND_SSBUF_CREATE_IF_OPS_FOR_CONTROL_FLOW_H diff --git a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/InitDependentMap.h b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/InitDependentMap.h new file mode 100644 index 0000000000..3559a281c8 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/InitDependentMap.h @@ -0,0 +1,49 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ASCEND_SSBUF_INITDEPENDENTMAPFORCONTROLFLOW_H +#define TRITON_ASCEND_SSBUF_INITDEPENDENTMAPFORCONTROLFLOW_H + +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace triton { + +class InitDependentMapPass + : public PassWrapper> { +public: + InitDependentMapPass() = default; + void runOnOperation() override; + void setConditionInfo(ControlFlowConditionInfo *info) { this->info = info; } + +private: + ControlFlowConditionInfo *info = nullptr; +}; + +std::unique_ptr> createInitDependentMapPass(); + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ASCEND_SSBUF_INITDEPENDENTMAPFORCONTROLFLOW_H diff --git a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.h b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.h new file mode 100644 index 0000000000..cd3f3a187e --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.h @@ -0,0 +1,200 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_UPDATE_CONDITION_INFO_H +#define TRITON_ADAPTER_UPDATE_CONDITION_INFO_H + +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Linalg/TransformOps/DialectExtension.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Operation.h" +#include "mlir/Pass/Pass.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/SmallVector.h" + +namespace mlir { +namespace triton { +enum class VarUpdateType { INC, DEC }; + +struct OutputGroupInfo { + SmallVector outputs; + SmallVector inputVars; +}; + +class UpdateConditionInfoPass + : public PassWrapper> { +public: + UpdateConditionInfoPass() = default; + + void runOnOperation() override; + + void setConditionInfo(ControlFlowConditionInfo *info) { this->info = info; } + +private: + SmallVector> allocSSBuffer(ModuleOp module); + + int updateIfConds(ModuleOp module, + SmallVector> ssbufferPtrs); + + void updateForIterTimes(ModuleOp module); + + scf::ForOp extendForOpIterationCount(scf::ForOp oldForOp, int ifCount, + int requiredBuffers, int x, + IRMapping &mapper, + SmallVector &ifOpsInThisFor); + + void replaceForOpCounterInIfOps(SmallVector ifOpsInThisFor, + scf::ForOp oldForOp, scf::ForOp newForOp, + IRMapping &mapper); + + Value getVarValue(scf::ForOp forOp, int varIndex); + + void collectDependencyBuffers( + scf::ForOp forOp, + DenseMap>> &crossCoreBuffers, + DenseMap>> &intraCoreBuffers); + + DenseMap>> + extendCrossCoreBuffersWithEquivalentValues( + ModuleOp module, + DenseMap>> crossCoreBuffers); + + int buildIdxToVarMap(scf::ForOp forOp, + const DenseMap>> + &intraCoreBuffers, + DenseMap &idxToVar); + + int getInputOutputValues( + scf::IfOp ifOp, + DenseMap>> crossCoreBuffers, + DenseMap>> intraCoreBuffers, + SmallVector &crossCoreInputValues, + SmallVector &crossCoreOutputValues, + SmallVector &intraCoreInputValues, + SmallVector &intraCoreOutputValues); + + int buildOutputGroups( + SmallVector &intraCoreOutputValues, + DenseMap>> &intraCoreBuffers, + DenseMap &idxToVar, + SmallVector &outputGroups); + + void collectIntraCoreInputConditions( + OpBuilder &builder, Location loc, SmallVector &intraCoreInputValues, + DenseMap &idxToVar, SmallVector &conditions, + DenseSet &usedVarsSet, + DenseMap &varUpdateTypes); + + int collectIntraCoreOutputConditions( + OpBuilder &builder, Location loc, + DenseMap>> &intraCoreBuffers, + SmallVector &intraCoreOutputValues, DenseMap &idxToVar, + SmallVector &conditions, DenseSet &usedVarsSet, + DenseMap &varUpdateTypes); + + SmallVector buildNewIfResultTypes(scf::IfOp oldIfOp, bool hasCounter, + Value counter); + + void collectYieldOperands(Block &block, Operation *&yieldOp, + SmallVector &yieldOperands); + + void populateNewThenBlock(scf::IfOp newIfOp, Block &oldThenBlock, + Operation *oldThenYieldOp, + ArrayRef oldYieldOperands, + DenseMap &varUpdateTypes, + bool hasCounter, Value counter, Value step); + + void populateNewElseBlock(scf::IfOp newIfOp, scf::IfOp oldIfOp, + bool needsYield, bool oldHasElse, bool hasCounter, + Value counter); + + scf::IfOp + createNewIfOpWithBlocks(scf::IfOp oldIfOp, Value combinedCond, + DenseMap &varUpdateTypes, + bool hasCounter, Value counter, Value step); + + int setIntraCoreCondition( + ModuleOp module, scf::IfOp ifOp, + DenseMap>> &intraCoreBuffers, + SmallVector &intraCoreInputIndices, + SmallVector &intraCoreOutputIndices, DenseMap &idxToVar, + DenseMap &varUpdateTypes, Value &intraCoreCond); + + void updateControlVarToLatestValue(scf::IfOp newIfOp, scf::IfOp oldIfOp, + bool hasCounter, Value counter); + + int updateForOpYield(scf::ForOp forOp); + + int combineConditions(ModuleOp module, Value crossCoreCond, + Value intraCoreCond, scf::IfOp ifOp, scf::ForOp forOp, + size_t &usedCounterNum, + DenseMap &varUpdateTypes); + + int setCrossCoreCondition( + SmallVector crossCoreInputValues, + SmallVector crossCoreOutputValues, + DenseMap>> &crossCoreBuffers, + scf::IfOp ifOp, SmallVector> ssbufferPtrs, + Value &crossCoreCond); + + // Helper function to get pointer based on core type + Value getSSBufferPtr(bool isAIC, int groupIdx, int ptrSetIdx, + DenseMap &precomputedPtrs, + SmallVector> ssbufferPtrs); + + // Compute pointers for VECTOR core SSBuffer + DenseMap + computeVectorSSBufferPtrs(OpBuilder &builder, Location loc, + Operation *scopeOp, + SmallVector crossCoreInputValues, + SmallVector crossCoreOutputValues); + + // Part 2: Add cross-core conditions + Value addCrossCoreConditions( + OpBuilder &builder, Location loc, SmallVector crossCoreInputValues, + SmallVector crossCoreOutputValues, + DenseMap>> &crossCoreBuffers, + bool isAIC, Value zeroConst, DenseMap &precomputedPtrs, + SmallVector> ssbufferPtrs); + + // Part 3: Update control variables in then block + void updateCrossCoreControlVars(OpBuilder &builder, Location loc, + scf::IfOp ifOp, + SmallVector crossCoreInputValues, + SmallVector crossCoreOutputValues, + bool isAIC, Value oneConst, + DenseMap &precomputedPtrs, + SmallVector> ssbufferPtrs); + + DenseMap controlVarToLatestValue; + SmallVector currentUsedVars; + ControlFlowConditionInfo *info = nullptr; +}; + +std::unique_ptr> createUpdateConditionInfoPass(); +} // namespace triton +} // namespace mlir +#endif // TRITON_ADAPTER_UPDATE_CONDITION_INFO_H diff --git a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.h b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.h new file mode 100644 index 0000000000..5ca39897f4 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.h @@ -0,0 +1,61 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ASCEND_SSBUF_UPDATE_FOR_OPS_FOR_CONTROL_FLOW_H +#define TRITON_ASCEND_SSBUF_UPDATE_FOR_OPS_FOR_CONTROL_FLOW_H +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition.h" + +namespace mlir { +namespace triton { + +class UpdateForOpsPass + : public PassWrapper> { +public: + UpdateForOpsPass() = default; + + void runOnOperation() override; + + void setConditionInfo(ControlFlowConditionInfo *info) { this->info = info; } + + llvm::StringRef getArgument() const override { return "update-for-ops"; } + +private: + LogicalResult + addBlockCountersAndInnerDepConds(ModuleOp module, + ControlFlowConditionInfo *info); + + LogicalResult deriveBlockCountersFromIfOps(ModuleOp module, + ControlFlowConditionInfo *info); + + LogicalResult insertInterCorePipeS(ModuleOp module); + + ControlFlowConditionInfo *info = nullptr; +}; + +std::unique_ptr> createUpdateForOpsPass(); + +} // namespace triton +} // namespace mlir +#endif // TRITON_ASCEND_SSBUF_UPDATE_FOR_OPS_FOR_CONTROL_FLOW_H diff --git a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.h b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.h new file mode 100644 index 0000000000..0346f011f5 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.h @@ -0,0 +1,103 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_UPDATE_LOOP_ITER_TIMES_H +#define TRITON_ADAPTER_UPDATE_LOOP_ITER_TIMES_H + +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/IRMapping.h" +#include "mlir/Pass/Pass.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" + +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition.h" + +namespace mlir { +namespace triton { + +// Loop iteration times info +struct IterationTimesInfo { + int ifCount = 0; // Number of if operations + int requiredBuffers = 1; // Number of required buffers + int x = 1; // Number of producer buffers + SmallVector ifOpsInThisFor; // If operations in this loop +}; + +class UpdateLoopIterTimesPass + : public PassWrapper> { +public: + UpdateLoopIterTimesPass() = default; + void runOnOperation() override; + void setConditionInfo(ControlFlowConditionInfo *info) { this->info = info; } + +private: + ControlFlowConditionInfo *info = nullptr; + + int GetMainLoopIdToLoopOpMap(ModuleOp module, + DenseMap> &cmap, + DenseMap> &vmap); + + int ComputeMainLoopTimes(DenseMap> &loopMap, + DenseMap &infoMap); + + int UpdateForLoopIteration( + DenseMap> &cmap, + DenseMap> &vmap, + DenseMap &infoMap); + + int collectForOpsAndUpdateMax( + DenseMap> &map, int id, + SmallVector &allForOps, int &maxIfCount, + int &maxRequiredBuffers, int &maxX, + DenseMap &infoMap); + + int replaceForOpCounterInIfOps(); + + // Calculate factor = requiredBuffers / x + std::pair calculateFactor(scf::ForOp forOp); + + // Extend for loop iteration count + scf::ForOp extendForOpIterationCount(scf::ForOp oldForOp, int ifCount, + int requiredBuffers, int x, + IRMapping &mapper, + SmallVector &ifOpsInThisFor); + + Value computeNewLoopUpperBound(OpBuilder &builder, Location loc, + scf::ForOp forOp, int ifCount, + int requiredBuffers, int x); + + scf::ForOp cloneForOpWithNewUpperBound(OpBuilder &builder, Location loc, + scf::ForOp oldForOp, + Value newUpperBound, + IRMapping &mapper); + + int updateCntArgsAfterClone(scf::ForOp oldForOp, IRMapping &mapper, + SmallVector &ifOpsInThisFor); +}; + +std::unique_ptr> createUpdateLoopIterTimesPass(); + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_UPDATE_LOOP_ITER_TIMES_H diff --git a/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/Utils.h b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/Utils.h new file mode 100644 index 0000000000..6cebf41191 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/AddControlFlowCondition/Utils.h @@ -0,0 +1,54 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_UTILS_H +#define TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_UTILS_H + +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/SmallVector.h" + +namespace mlir { +namespace triton { + +// Collect all nested ops within an operation's regions +LogicalResult collectAllNestedOps(Operation *op, + llvm::DenseSet ®ionOps); + +// Group operations by their block_id attribute +LogicalResult +collectOpsByBlockId(scf::ForOp forOp, + llvm::DenseMap> &blockOps); + +// Topological sort of operations based on operand dependencies +LogicalResult topologicalSort(llvm::DenseSet &ops, + llvm::DenseMap *opOrder, + SmallVectorImpl &sorted); + +LogicalResult topologicalSort(SmallVector &ops); + +// Get block_ids in order of appearance in for loop body +SmallVector getBlockIdsInOrder(scf::ForOp forOp); +} // namespace triton +} // namespace mlir +#endif // TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_UTILS_H diff --git a/third_party/ascend/include/DynamicCVPipeline/AddDynamicCVPipeline.h b/third_party/ascend/include/DynamicCVPipeline/AddDynamicCVPipeline.h new file mode 100644 index 0000000000..bf2e67a00a --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/AddDynamicCVPipeline.h @@ -0,0 +1,57 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_ADD_DYNAMIC_CVPIPELINE_PASSES_H +#define TRITON_ADAPTER_ADD_DYNAMIC_CVPIPELINE_PASSES_H + +#include "mlir/Pass/Pass.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#define GEN_PASS_DECL_ADDDYNAMICCVPIPELINE +#include "ascend/include/DynamicCVPipeline/Passes.h.inc" + +#define GEN_PASS_DEF_ADDDYNAMICCVPIPELINE +#include "ascend/include/DynamicCVPipeline/Passes.h.inc" + +extern bool compileOn91095Flag; + +namespace mlir { +namespace triton { +std::unique_ptr> +createAddDynamicCVPipelinePass(const AddDynamicCVPipelineOptions &options = {}); +} // namespace triton +} // namespace mlir + +namespace { +using namespace mlir; +using namespace triton; + +class AddDynamicCVPipelinePass + : public ::impl::AddDynamicCVPipelineBase { +public: + explicit AddDynamicCVPipelinePass(const AddDynamicCVPipelineOptions &options); + void runOnOperation() override; +}; + +} // namespace + +#endif // TRITON_ADAPTER_ADD_DYNAMIC_CVPIPELINE_PASSES_H diff --git a/third_party/ascend/include/DynamicCVPipeline/AllocMultiCache.h b/third_party/ascend/include/DynamicCVPipeline/AllocMultiCache.h new file mode 100644 index 0000000000..23a06ca13d --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/AllocMultiCache.h @@ -0,0 +1,51 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_ALLOC_MULTI_CACHE_PASS_H +#define TRITON_ADAPTER_ALLOC_MULTI_CACHE_PASS_H + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace triton { + +// AllocMultiCachePass for allocating multi-cache +class AllocMultiCachePass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AllocMultiCachePass) + + // Constructor + AllocMultiCachePass() = default; + + // Run the pass + void runOnOperation() override; +}; + +// Create the pass +std::unique_ptr> createAllocMultiCachePass(); + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_ALLOC_MULTI_CACHE_PASS_H diff --git a/third_party/ascend/include/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.h b/third_party/ascend/include/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.h new file mode 100644 index 0000000000..7bc8b02355 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.h @@ -0,0 +1,64 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_ALLOC_MULTI_CACHE_ADD_MULTI_BUFFER_INNER_SCOPE_PASS_H +#define TRITON_ADAPTER_ALLOC_MULTI_CACHE_ADD_MULTI_BUFFER_INNER_SCOPE_PASS_H + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace triton { + +// AddMultiBufferInnerScopePass for adding multi-buffer optimization within +// inner scope +class AddMultiBufferInnerScopePass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AddMultiBufferInnerScopePass) + + // Constructor + AddMultiBufferInnerScopePass() = default; + + // Pass argument + StringRef getArgument() const override { + return "add_multi_buffer_inner_scope"; + } + + // Dependent dialects + void getDependentDialects(DialectRegistry ®istry) const override; + + // Run the pass + void runOnOperation() override; +}; + +// Create the pass +std::unique_ptr> createAddMultiBufferInnerScopePass(); + +// Register the pass +void registerAddMultiBufferInnerScopePasses(); + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_ALLOC_MULTI_CACHE_ADD_MULTI_BUFFER_INNER_SCOPE_PASS_H diff --git a/third_party/ascend/include/DynamicCVPipeline/AllocMultiCache/AddMultiBufferOuterScope.h b/third_party/ascend/include/DynamicCVPipeline/AllocMultiCache/AddMultiBufferOuterScope.h new file mode 100644 index 0000000000..254c28789d --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/AllocMultiCache/AddMultiBufferOuterScope.h @@ -0,0 +1,128 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_ALLOC_MULTI_CACHE_ADD_MULTI_BUFFER_OUTER_SCOPE_PASS_H +#define TRITON_ADAPTER_ALLOC_MULTI_CACHE_ADD_MULTI_BUFFER_OUTER_SCOPE_PASS_H + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/DialectRegistry.h" +#include "mlir/Pass/Pass.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallSet.h" + +namespace mlir { +namespace triton { + +// Data structures + +/// Buffer alloc pair: {allocOp, markOp} +struct BufferAllocPair { + Operation *allocOp = nullptr; + Operation *markOp = nullptr; +}; + +/// Transfer operation chain for sender or receiver side +struct TransferOpChain { + Operation *waitOp = nullptr; + Operation *transferOp = + nullptr; // fixpipe / hir.copy / memory_space_cast / convert_layout + Operation *toTensorOp = + nullptr; // bufferization.to_tensor (memory_space_cast scenario only) + Operation *setOp = nullptr; +}; + +/// Buffer alloc info: {sender, receiver} +struct BufferAllocInfo { + BufferAllocPair sender; + BufferAllocPair receiver; +}; + +/// Extra sync info for sync ops whose parent does NOT have main_loop attribute +struct ExtraSyncInfo { + Operation *setOp = nullptr; + Operation *waitOp = nullptr; +}; + +/// Transfer chain info: {senderChain, receiverChain} +struct TransferChainInfo { + TransferOpChain sender; + TransferOpChain receiver; +}; + +/// Complete transfer group information +struct TransferGroupInfo { + int tid = -1; + int originalFlag = -1; + int outputFlag = -1; + bool isCtoV = false; // true=C→V, false=V→C + + BufferAllocPair senderBuf; + BufferAllocPair receiverBuf; + + TransferOpChain senderChain; + TransferOpChain receiverChain; + + // Input/output buffer values + Value senderInputBuffer; + Value senderOutputBuffer; + Value receiverInputBuffer; + Value receiverOutputBuffer; + + // TCB ID shared across all 4 buffers in the same transfer group + int tcbId = -1; + + // Extra sync positions for output flag synchronization + Operation *extraSyncSetOp = nullptr; + Operation *extraSyncWaitOp = nullptr; +}; + +/// AddMultiBufferOuterScopePass for adding outer (CV inter-core) multi-buffer +/// optimization +class AddMultiBufferOuterScopePass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AddMultiBufferOuterScopePass) + + // Constructor + AddMultiBufferOuterScopePass() = default; + + // Pass argument + StringRef getArgument() const override { + return "add_multi_buffer_outer_scope"; + } + + // Run the pass + void runOnOperation() override; + + // Get dependent dialects + void getDependentDialects(DialectRegistry ®istry) const override; +}; + +std::unique_ptr> createAddMultiBufferOuterScopePass(); + +void registerAddMultiBufferOuterScopePasses(); + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_ALLOC_MULTI_CACHE_ADD_MULTI_BUFFER_OUTER_SCOPE_PASS_H diff --git a/third_party/ascend/include/DynamicCVPipeline/CMakeLists.txt b/third_party/ascend/include/DynamicCVPipeline/CMakeLists.txt new file mode 100644 index 0000000000..c2514cc0a2 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name AddDynamicCVPipeline) +add_public_tablegen_target(AddDynamicCVPipelinePassIncGen) diff --git a/third_party/ascend/include/DynamicCVPipeline/Common/BufferCountManager.h b/third_party/ascend/include/DynamicCVPipeline/Common/BufferCountManager.h new file mode 100644 index 0000000000..8debd03e9f --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/Common/BufferCountManager.h @@ -0,0 +1,43 @@ +#ifndef TRITON_DYNAMIC_CV_PIPELINE_ADDMULTIBUFFERCONTROL_BUFFER_COUNT_MANAGER_H +#define TRITON_DYNAMIC_CV_PIPELINE_ADDMULTIBUFFERCONTROL_BUFFER_COUNT_MANAGER_H + +#include "mlir/IR/Value.h" +#include "llvm/ADT/DenseMap.h" +#include + +namespace mlir { + +class Operation; + +namespace triton { + +class BufferCountManager { +public: + static BufferCountManager &getInstance(); + + enum class DepType { IntraCore, InterCore, LoadStore }; + + void setBufferCount(DepType type, int count); + + void + buildBufferCountMap(llvm::DenseMap> &depValueMap, + llvm::DenseMap &bufferCountMap, DepType type); + + int getBufferCountByType(DepType type) const; + +private: + BufferCountManager(); + BufferCountManager(const BufferCountManager &) = delete; + BufferCountManager &operator=(const BufferCountManager &) = delete; + + int intraBufferCount_; + int interCoreBufferCount_; + int loadStoreBufferCount_; +}; + +#define BUFFER_COUNT (BufferCountManager::getInstance()) + +} // namespace triton +} // namespace mlir + +#endif // TRITON_DYNAMIC_CV_PIPELINE_ADDMULTIBUFFERCONTROL_BUFFER_COUNT_MANAGER_H diff --git a/third_party/ascend/include/DynamicCVPipeline/Common/FlagIdManager.h b/third_party/ascend/include/DynamicCVPipeline/Common/FlagIdManager.h new file mode 100644 index 0000000000..f5874ecbc0 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/Common/FlagIdManager.h @@ -0,0 +1,65 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_COMMON_FLAG_ID_MANAGER_H +#define TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_COMMON_FLAG_ID_MANAGER_H + +#include "mlir/IR/BuiltinOps.h" +#include + +namespace mlir { +namespace triton { + +// Flag ID Manager for synchronization between cores +// Purpose: Globally scan and allocate unique flag_id for synchronization +class FlagIdManager { +public: + // Maximum flag count (hardware limitation) + static constexpr int MAX_FLAG_ID = 14; + static constexpr int INVALID_FLAG_ID = -1; + + // Constructor: initialize with Module for scanning + FlagIdManager(ModuleOp module); + + // Acquire an available ID: try to reuse existing flags first (conservative + // analysis), if cannot reuse, then increment and allocate. insertionPoint: + // the position where sync operation will be inserted, used for linear + // comparison in reuse analysis. Can be nullptr. + int acquireId(Operation *insertionPoint); + +private: + // Scan existing Flag IDs + // Traverse Module to find all existing flag_id to prevent duplicate + // allocation + void scanExistingFlags(ModuleOp module); + + // Currently allocated maximum ID + int64_t currentMaxId = 0; + + // Save module for reuse analysis + ModuleOp module; +}; + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_COMMON_FLAG_ID_MANAGER_H diff --git a/third_party/ascend/include/DynamicCVPipeline/Common/MemoryEffectsTracker.h b/third_party/ascend/include/DynamicCVPipeline/Common/MemoryEffectsTracker.h new file mode 100644 index 0000000000..3df3ed6b6d --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/Common/MemoryEffectsTracker.h @@ -0,0 +1,101 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_COMMON_MEMORY_EFFECTS_TRACKER_H +#define TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_COMMON_MEMORY_EFFECTS_TRACKER_H + +#include "mlir/Analysis/AliasAnalysis.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/Value.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" + +namespace mlir { +namespace CVPipeline { +constexpr int INIT_SIZE = 4; +class MemoryDependenceGraph { +public: + MemoryDependenceGraph(Operation *root, AliasAnalysis &aa); + + ArrayRef getMemDefs(Operation *op) const; + ArrayRef getMemUsers(Operation *op) const; + + ArrayRef getExecBefore(Operation *op) const; + ArrayRef getExecAfter(Operation *op) const; + +private: + struct MemSlot { + Value memref; + Operation *lastWriter = nullptr; + Operation *dataSource = nullptr; + SmallPtrSet pendingReads; + explicit MemSlot(Value v) : memref(v) {} + }; + + struct Snapshot { + SmallVector states; + }; + + void analyzeOp(Operation *op); + void analyzeRegionsOf(Operation *op); + + SmallVector collectOuterEffects(Operation *op, + bool &unknown); + + SmallVector findAliasSlots(Value v); + ArrayRef + resolveAliasSlots(Value v, DenseMap> &cache); + MemSlot *getOrCreateSlot(Value v); + + void collectPreds(ArrayRef effects, + bool unknown, SmallVectorImpl &defsOut, + SmallVectorImpl &predsOut); + + void applyEffects(Operation *op, + ArrayRef effects, + bool unknown); + + Snapshot takeSnapshot() const; + void restoreSnapshot(Snapshot &&snap); + + void recordEdges(Operation *op, ArrayRef defs, + ArrayRef preds); + AliasResult queryAlias(Value lhs, Value rhs); + + Operation *root; + AliasAnalysis &aa; + + // Data dependency + DenseMap> memDefs; + DenseMap> memUsers; + + // Execute order + DenseMap> execBefore; + DenseMap> execAfter; + + SmallVector> slots; + DenseMap valueToSlot; +}; + +} // namespace CVPipeline +} // namespace mlir + +#endif // TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_COMMON_MEMORY_EFFECTS_TRACKER_H diff --git a/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h b/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h new file mode 100644 index 0000000000..47bde37294 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h @@ -0,0 +1,63 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef ADD_AUTO_SCHEDULING_COMMON_UTILS_H +#define ADD_AUTO_SCHEDULING_COMMON_UTILS_H +#include "mlir/IR/Operation.h" +#include "mlir/IR/Value.h" +#include "llvm/ADT/StringRef.h" +#include + +namespace mlir { +namespace CVPipeline { + +inline constexpr llvm::StringLiteral kCoreType = "ssbuffer.core_type"; +inline constexpr llvm::StringLiteral kBlockId = "ssbuffer.block_id"; + +enum CoreType { + UNDETERMINED = 0, + VECTOR_ONLY = 1 << 0, + CUBE_ONLY = 1 << 1, + CUBE_AND_VECTOR = VECTOR_ONLY | CUBE_ONLY, +}; + +inline constexpr CoreType fromStrCoreType(std::string_view s) { + if (s == "VECTOR") { + return CoreType::VECTOR_ONLY; + } + if (s == "CUBE") { + return CoreType::CUBE_ONLY; + } + + return CoreType::UNDETERMINED; +} + +// Functions for managing core types +CoreType getOpCoreType(Operation *op); + +llvm::LogicalResult verifyOpBlockId(Operation *op); +std::optional getOpBlockId(Operation *op); + +} // namespace CVPipeline +} // namespace mlir + +#endif diff --git a/third_party/ascend/include/DynamicCVPipeline/Passes.h b/third_party/ascend/include/DynamicCVPipeline/Passes.h new file mode 100644 index 0000000000..b60274fd14 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/Passes.h @@ -0,0 +1,37 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_PASSES_H +#define TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_PASSES_H + +#include "AddDynamicCVPipeline.h" + +namespace mlir { +namespace triton { + +#define GEN_PASS_REGISTRATION +#include "ascend/include/DynamicCVPipeline/Passes.h.inc" + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_PASSES_H diff --git a/third_party/ascend/include/DynamicCVPipeline/Passes.td b/third_party/ascend/include/DynamicCVPipeline/Passes.td new file mode 100644 index 0000000000..9450000817 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/Passes.td @@ -0,0 +1,21 @@ +/* + * Copyright (c) Huawei Technologies Co. + * Licensed under the MIT license. + */ + +#ifndef TRITON_DYNAMIC_CV_PIPELINE_PASSES +#define TRITON_DYNAMIC_CV_PIPELINE_PASSES + +include "mlir/Pass/PassBase.td" + +def AddDynamicCVPipeline : Pass<"add_dynamic_cv_pipeline", "mlir::ModuleOp"> { + let summary = "Add dynamic cv pipeline based on ssbufer"; + let constructor = "triton::createAddDynamicCVPipelinePass()"; + let options = [ + Option<"compileOn91095", "compile-on-910-95", + "bool", /*default*/"false", + "compile on 910_95"> + ]; +} + +#endif // TRITON_DYNAMIC_CV_PIPELINE_PASSES diff --git a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h new file mode 100644 index 0000000000..94e2d9338e --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h @@ -0,0 +1,48 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_COMMON_H +#define TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_COMMON_H + +#include "ascend/include/DynamicCVPipeline/Common/MemoryEffectsTracker.h" +#include "ascend/include/DynamicCVPipeline/Common/Utils.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/Operation.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" + +namespace mlir { +namespace CVPipeline { + +inline bool isCubeOp(Operation *op) { + return getOpCoreType(op) == CoreType::CUBE_ONLY; +} +Operation *getAncestorInBlock(Operation *inner, Block *block); +void initializeIndegreeForBlock(Block *block, + llvm::DenseMap &indegree, + const MemoryDependenceGraph &memGraph); + +} // namespace CVPipeline +} // namespace mlir + +#endif // TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_COMMON_H diff --git a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h new file mode 100644 index 0000000000..b9e562b317 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h @@ -0,0 +1,73 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_COMPUTE_BLOCK_ID_MANAGER_H +#define TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_COMPUTE_BLOCK_ID_MANAGER_H + +#include "mlir/IR/Operation.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/LogicalResult.h" +#include + +namespace mlir { +namespace CVPipeline { + +/** + * the class is to promise CUBEID and VECTORID are unified. + */ +class ComputeBlockIdManager { +public: + static ComputeBlockIdManager &getInstance() { + static ComputeBlockIdManager instance; + return instance; + } + + unsigned int getNextId(); + + bool isSameBlock(Operation *a, Operation *b); + bool isWholeCubeReady(Operation *seedOp, + llvm::DenseMap &indegree); + + llvm::LogicalResult markOpBlockId(Operation *op, int blockId); + + llvm::SmallVector getOpsByBlockId(int blockId); + int getBlockIdByOp(Operation *op); + llvm::LogicalResult markOpsWithNewId(llvm::SmallVectorImpl &ops); + void reset(); + +private: + ComputeBlockIdManager() : cntComputeBlockId(0) {} + llvm::LogicalResult record(Operation *op, int blockId); + + unsigned int cntComputeBlockId; + llvm::DenseMap> blockIdToOps; + llvm::DenseMap opToBlockId; + mutable std::mutex managerMutex; + const int blockIdWidth = 32; +}; + +} // namespace CVPipeline +} // namespace mlir + +#endif // TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_COMPUTE_BLOCK_ID_MANAGER_H diff --git a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/OpClassifier.h b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/OpClassifier.h new file mode 100644 index 0000000000..458267e866 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/OpClassifier.h @@ -0,0 +1,177 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_OP_CLASSIFIER_H +#define TRITON_ADAPTER_OP_CLASSIFIER_H + +#include "ascend/include/DynamicCVPipeline/Common/MemoryEffectsTracker.h" +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "mlir/Analysis/AliasAnalysis.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Pass/PassManager.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +namespace mlir { +namespace triton { + +// Core type enumeration for operations +enum OpCoreType { + OP_UNDETERMINED = 0, + OP_CUBE_ONLY = 1, + OP_VECTOR_ONLY = 2, + OP_CUBE_AND_VECTOR = 3 +}; + +// OpClassifierPass for categorizing operations as CUBE or VECTOR +class OpClassifierPass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OpClassifierPass) + + // Constructor + OpClassifierPass() = default; + + // Run the pass + void runOnOperation() override; + + // Return the pass argument name + static constexpr ::llvm::StringRef getArgumentName() { + return "op-classifier"; + } + ::llvm::StringRef getArgument() const override { return "op-classifier"; } + ::llvm::StringRef getDescription() const override { + return "Classify operations as CUBE or VECTOR for dynamic CV pipeline"; + } + ::llvm::StringRef getName() const override { return "OpClassifierPass"; } + +private: + // Map from operation to its core type + llvm::DenseMap opCoreTypes; + + // All operations in the module + llvm::SmallVector allOps; + + // Seed operations for CUBE upstream propagation + llvm::SmallVector cubeSeeds; + + std::shared_ptr aliasAnalysis; + std::shared_ptr memDepGraph; + + // Mark an operation as CUBE + void markCube(Operation *op); + + // Pattern matching for CUBE operations + int patternMatchCUBE(); + + // Upstream pattern matching helpers + void matchToTensorPattern(Operation *def); + void matchTransposePattern(Operation *def); + void matchFillPattern(Operation *def); + + // Downstream pattern matching helpers + void matchStorePattern(Operation *user); + void matchExtractSlicePattern(Operation *user); + void matchMaterializePattern(Operation *user); + + // Propagate CUBE core type upstream + int propagateCubeUpstream(); + + // Get upstream operations based on both SSA and memory dependencies + void + getUpstreamOpsWithMemoryDeps(Operation *cur, + llvm::SmallVectorImpl &upstreamOps); + + // Step 3: Mark remaining operations as VECTOR + int markRemainingAsVector(); + + // Step 4: Propagate VECTOR core type upstream + int propagateVectorUpstream(); + + // Initialize the pass + void initializePass(ModuleOp module); + + // Get the core type of an operation + OpCoreType getCoreType(Operation *op) const; + + // Set the core type of an operation + void setCoreType(Operation *op, OpCoreType coreType); + + // Helper: join core types into comma-separated string + std::string joinCoreTypes(const std::vector &coreTypes); + + // Helper: check if propagation should stop at this operation for yield + bool shouldStopPropagationForYield(Operation *op, OpCoreType targetCoreType); + + // Helper: propagate core_type upward for yield operand + void propagateCoreTypeUpwardForYield(Operation *startOp, + OpCoreType targetCoreType); + + // Step 5: Handle else region yield of scf.if by extracting core_type from + // then region yield + bool handleYieldFromElseRegion(std::vector &coreTypes, + unsigned operandIndex, + Operation *thenYieldForElse, Value &operand); + + // Step 6: Handle CUBE_AND_VECTOR operations + int handleCubeAndVector(); + + // Helper: split CUBE_AND_VECTOR operation into CUBE and VECTOR versions + void splitOperationForCubeAndVector( + Operation *op, llvm::DenseSet &processedOps, + llvm::DenseMap &opToVectorClone); + + // Step 7: Process SCF yield results + int handleSCFYield(); + + // Helper: process a single yield operation + void processYieldOperation(Operation *op, Operation *thenYieldForElse); + + // Helper: Mark fill operations as CUBE when their output buffer is CUBE + void markFillOpsAsCube(); + + // Step 7: Pre-legalize matmul (before initializePass) + int preLegalizeMatmul(); + + // Helper: bulk delete operations and clean up tracking structures + void bulkDeleteOps(llvm::SmallVectorImpl &opsToDelete); + + // Step 8: Stamp core type info to IR + int stampToIR(); + + // Helper: describe operation for logging + std::string describeOp(Operation *op) const; +}; + +// Create the pass +std::unique_ptr> createOpClassifierPass(); +void registerAddOpClassifierPasses(); +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_OP_CLASSIFIER_H diff --git a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/Passes.h b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/Passes.h new file mode 100644 index 0000000000..e85d832821 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/Passes.h @@ -0,0 +1,37 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_PASSES_H +#define TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_PASSES_H + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace triton { + +std::unique_ptr> createPlanVectorBlockPass(); + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_PASSES_H diff --git a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlockPass.h b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlockPass.h new file mode 100644 index 0000000000..fa43fd4453 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlockPass.h @@ -0,0 +1,47 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITION_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_PLAN_CUBE_BLOCK_PASS_H +#define TRITION_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_PLAN_CUBE_BLOCK_PASS_H + +#include + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace triton { + +class PlanCubeBlockPass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PlanCubeBlockPass); + + PlanCubeBlockPass() = default; + void runOnOperation() override; +}; + +std::unique_ptr> createPlanCubeBlockPass(); +} // namespace triton +} // namespace mlir + +#endif // TRITION_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_PLAN_CUBE_BLOCK_PASS_H diff --git a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ReorderOpsByBlockId.h b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ReorderOpsByBlockId.h new file mode 100644 index 0000000000..4aee2ac289 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ReorderOpsByBlockId.h @@ -0,0 +1,44 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITION_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_REORDER_OPS_BY_BLOCKID_H +#define TRITION_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_REORDER_OPS_BY_BLOCKID_H + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir::triton { + +class ReorderOpsByBlockIdPass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ReorderOpsByBlockIdPass); + + ReorderOpsByBlockIdPass() = default; + void runOnOperation() override; +}; + +std::unique_ptr> createReorderOpsByBlockIdPass(); + +} // namespace mlir::triton + +#endif // TRITION_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_REORDER_OPS_BY_BLOCKID_H diff --git a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlockPass.h b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlockPass.h new file mode 100644 index 0000000000..67c4008804 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlockPass.h @@ -0,0 +1,52 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_COMPUTE_BLOCK_PASS_H +#define TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_COMPUTE_BLOCK_PASS_H + +#include "mlir/Dialect/Linalg/TransformOps/DialectExtension.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace triton { + +// PlanComputeBlockPass for partitioning operations into compute blocks +class PlanComputeBlockPass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PlanComputeBlockPass) + + // Constructor + PlanComputeBlockPass() = default; + + // Run the pass + void runOnOperation() override; +}; + +// Create the pass +std::unique_ptr> createPlanComputeBlockPass(); + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_COMPUTE_BLOCK_PASS_H diff --git a/third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadInternal.h b/third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadInternal.h new file mode 100644 index 0000000000..d6e92f02c2 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadInternal.h @@ -0,0 +1,204 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_ADD_MULTI_BUFFER_TO_GMLOAD_INTERNAL_H +#define TRITON_ADAPTER_ADD_MULTI_BUFFER_TO_GMLOAD_INTERNAL_H + +#include "ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadTypes.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/IRMapping.h" +#include "mlir/IR/Operation.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "mlir/Support/LogicalResult.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/SetVector.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Debug.h" + +using namespace mlir; + +// ============================================================================ +// Constants +// ============================================================================ + +static constexpr const char *DEBUG_TYPE = "add-multi-buffer-to-gm-load"; +#define LOG_DEBUG(...) \ + LLVM_DEBUG(llvm::dbgs() << " [" << DEBUG_TYPE << "] " << __VA_ARGS__) + +static constexpr llvm::StringLiteral kBufferableMarker("gm_load_bufferable"); +static constexpr llvm::StringLiteral kBlockIdAttr("ssbuffer.block_id"); +static constexpr llvm::StringLiteral kLoadStoreAttr("ssbuffer.load_store"); +static constexpr llvm::StringLiteral + kGMLoadMultiBufferStage("gm-load multi-buffer"); +static constexpr unsigned kInitialRefCapacity = + 8; // Tuned for typical load chain depth +static constexpr unsigned kForInductionVarArgIndex = 0; +static constexpr unsigned kForBodyIterArgOffset = 1; +// Each load group carries `depth` slot flags plus producer/consumer counters. +// This count is independent of the configured buffer depth. +static constexpr int kLoadGroupCounterIterArgCount = 2; +static constexpr int kLoadGroupProducerCounterOffset = 0; +static constexpr int kLoadGroupConsumerCounterOffset = 1; + +// ============================================================================ +// Data structures (debug/utility only, types are in +// AddMultiBufferToGMLoadTypes.h) +// ============================================================================ + +namespace gmload { + +/// Reuses existing arith.constant ops found in the IR rather than creating +/// duplicates. Pre-populated by scanning a block; get-or-create semantics. +/// Suitable for untagged constants only (no ssbuffer.block_id attribution). +struct ConstantCache { + llvm::DenseMap table; + + void scan(Block *block) { + for (auto &op : *block) + if (auto cst = dyn_cast(op)) + table.try_emplace(cst.getValue(), cst.getResult()); + } + + Value get(OpBuilder &builder, Location loc, mlir::TypedAttr attr) { + auto [cacheIt, inserted] = table.try_emplace(attr, Value{}); + if (!inserted) + return cacheIt->second; + cacheIt->second = builder.create(loc, attr).getResult(); + return cacheIt->second; + } + + Value getFalse(OpBuilder &builder, Location loc) { + return get(builder, loc, builder.getBoolAttr(false)); + } + + Value getIndex(OpBuilder &builder, Location loc, int64_t value) { + return get(builder, loc, + cast(builder.getIndexAttr(value))); + } +}; + +// ============================================================================ +// Function declarations: Dependency analysis +// ============================================================================ + +Operation *getAncestorInBlock(Operation *op, Block *scopeBlock); +void backwardTrace(Operation *op, Block *scopeBlock, + llvm::SetVector &visited, + SmallVectorImpl &worklist); +void backwardTraceFromWorklist(llvm::SetVector &visited, + SmallVectorImpl &worklist, + Block *scopeBlock); +SmallVector +forwardTraceFromAllocs(llvm::SetVector &visited, + Block *scopeBlock); +void captureRegionFreeVars(llvm::SetVector &visited, + SmallVectorImpl &worklist, + Block *scopeBlock); +SmallVector computeLoadChain(Operation *markedOp, + Block *scopeBlock); +memref::AllocOp findBackingAlloc(Operation *markedOp, + ArrayRef chain); +SmallVector collectMarkedOps(ModuleOp module); +std::optional getConstantTripCount(scf::ForOp forOp); +void groupByEnclosingForOp(SmallVector &markedOps, + SmallVectorImpl &contexts); +void collectChainAndMarkedOps(const SmallVector &loads, + llvm::DenseSet &allChainOps, + llvm::DenseSet &markedOps); +bool isResultUsedExternally(Value result, Operation *yieldOp, + const llvm::DenseSet &allChainOps, + Block *forBody); +SmallVector +collectOperandsIncludingRegions(Operation *op); +llvm::DenseSet +computeSkipInConsumer(const SmallVector &loads, Block *forBody); + +// ============================================================================ +// Function declarations: Loop transformation utilities +// ============================================================================ + +bool isLoopInvariant(Value value, scf::ForOp forOp); +bool getLinearIterArgDelta(Value iterArg, Value yieldVal, scf::ForOp forOp, + Value &delta); +Value castIndexTo(OpBuilder &builder, Location loc, Value value, + Type targetType); +Value castToIndex(OpBuilder &builder, Location loc, Value value); +void tagWithBlockId(Value value, IntegerAttr blockId); +void allocateBufferSlots(OpBuilder &builder, Location loc, scf::ForOp forOp, + LoadGroup &group); +ExtendedForInfo buildExtendedFor(OpBuilder &builder, Location loc, + scf::ForOp forOp, ArrayRef groups, + ConstantCache &cache); +bool isDeadIterArg(scf::ForOp forOp, unsigned iterArgIdx); +FailureOr pruneDeadIterArgs(OpBuilder &builder, scf::ForOp forOp, + unsigned candidateCount); +void eraseDeadBodyOps(Block *body); +bool isAvailableForClone(Value value, Block *oldBody, const IRMapping &mapping); +bool areOperandsAvailableForClone(Operation *op, Block *oldBody, + const IRMapping &mapping); +std::optional getBlockId(Operation *op); +bool sameRunBlockId(const BodyRun &run, std::optional bid); +SmallVector collectBodyRuns(Block *body); +int findRunContainingOp(ArrayRef runs, Operation *target); +int findFirstRunWithBlockId(ArrayRef runs, int32_t blockId); +void logRepeatedTopLevelBlockRuns(Block *body, llvm::StringRef stage); +bool isAncestorBlock(Block *ancestor, Block *descendant); +void deduplicateConstants(ModuleOp module); +LogicalResult +applyMultiBufferToForLoop(ForBufferCtx &context, + const llvm::DenseSet &allCtxForOps); + +// ============================================================================ +// Function declarations: Consumer/Producer emission +// ============================================================================ + +Value emitLoadSlotSelection(OpBuilder &builder, Location loc, + const MarkedLoad &load, ArrayRef slotBufs, + Value target, int depth, IRMapping &mapping, + ArrayRef slotConsts); +Value emitSlotFlagClear(OpBuilder &builder, Location loc, Value target, + Value slotConst, Value flagNext, Value falseVal, + IntegerAttr blockId); +FailureOr> emitProducerFillForSlot( + OpBuilder &builder, Location loc, Block *oldBody, Block *newBody, + const LoadGroup &group, int slotIdx, Value flagArg, + Value currentProducerCounter, Value tripCount, Value falseVal, + Value loopLowerBound, Value loopStep, Value trueFlag, Value indexOne, + const SmallVector &iterArgDeltas, scf::ForOp origForOp); +LogicalResult emitMultiBufferLoopBody( + OpBuilder &builder, Location loc, SmallVectorImpl &groups, + ExtendedForInfo &info, const SmallVector &groupIndexOneVals, + const SmallVector> &allGroupSlotConsts, + const SmallVector &allGroupDepthConsts, + const llvm::DenseSet &allCtxForOps, Value tripCount, + Value loopLowerBound, Value loopStep, + const SmallVector &groupTrueFlagVals, + const SmallVector &iterArgDeltas, scf::ForOp origForOp); + +} // namespace gmload + +#endif // TRITON_ADAPTER_ADD_MULTI_BUFFER_TO_GMLOAD_INTERNAL_H diff --git a/third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadPass.h b/third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadPass.h new file mode 100644 index 0000000000..2bca3fe8d2 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadPass.h @@ -0,0 +1,70 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_ADD_MULTI_BUFFER_TO_GMLOAD_PASS_H +#define TRITON_ADAPTER_ADD_MULTI_BUFFER_TO_GMLOAD_PASS_H + +#include "ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadTypes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Support/LogicalResult.h" +#include "llvm/ADT/DenseSet.h" + +namespace mlir { +namespace triton { + +// Rewrites GM load operations inside scf.for loops into a multi-buffer +// producer/consumer structure to overlap memory transfers with compute. +class AddMultiBufferToGMLoadPass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AddMultiBufferToGMLoadPass) + + StringRef getArgument() const override { return "gm-load-multi-buffer"; } + void runOnOperation() override; + +private: + // Step 1: Collect marked ops and group by enclosing forOp + void collectAndGroupMarkedOps(); + + // Step 2: Sort contexts inner-first + void sortContextsInnerFirst(); + + // Step 3: Transform each for loop with multi-buffer logic + LogicalResult applyMultiBufferToGMLoadLoops(); + + // Step 4: Cleanup transformed IR + void cleanupTransformedIR(); + + // Shared state between steps + llvm::SmallVector markedOps_; + llvm::SmallVector contexts_; + llvm::DenseSet allCtxForOps_; +}; + +// Create the pass +std::unique_ptr> createAddMultiBufferToGMLoadPass(); + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_ADD_MULTI_BUFFER_TO_GMLOAD_PASS_H diff --git a/third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadTypes.h b/third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadTypes.h new file mode 100644 index 0000000000..d55b6ef349 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadTypes.h @@ -0,0 +1,85 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_ADD_MULTI_BUFFER_TO_GMLOAD_TYPES_H +#define TRITON_ADAPTER_ADD_MULTI_BUFFER_TO_GMLOAD_TYPES_H + +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/IRMapping.h" +#include "mlir/IR/Operation.h" +#include "llvm/ADT/SmallVector.h" + +namespace gmload { + +// One marked op and the topo-sorted dependency chain needed to reproduce it. +struct MarkedLoad { + mlir::Operation *markedOp = nullptr; + llvm::SmallVector chain; + mlir::memref::AllocOp allocOp; +}; + +// Each marked load gets its own independent LoadGroup with its own +// producer/consumer counter pair. +struct LoadGroup { + llvm::SmallVector loads; + int depth = 0; + llvm::SmallVector> + bufSlots; // bufSlots[slot][load] + llvm::SmallVector mergedChain; +}; + +// Transformation context for one scf.for. +struct ForBufferCtx { + mlir::scf::ForOp forOp; + llvm::SmallVector groups; +}; + +// Iter-arg handles for one LoadGroup inside the extended scf.for. +struct GroupIterArgs { + llvm::SmallVector flagArgs; + mlir::Value prodCounter; + mlir::Value consCounter; +}; + +// Output of buildExtendedFor: new loop + per-group iter_arg handles. +struct ExtendedForInfo { + mlir::scf::ForOp newForOp; + mlir::Block *oldBody = nullptr; + mlir::Block *newBody = nullptr; + mlir::IRMapping mapping; + llvm::SmallVector groupArgs; + mlir::Value falseVal; + int numOrig = 0; + int depth = 0; +}; + +// One consecutive top-level run in an old loop body. +struct BodyRun { + bool hasBlockId = false; + int32_t blockId = 0; + llvm::SmallVector ops; +}; + +} // namespace gmload + +#endif // TRITON_ADAPTER_ADD_MULTI_BUFFER_TO_GMLOAD_TYPES_H diff --git a/third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AsyncLoadHoistingPass.h b/third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AsyncLoadHoistingPass.h new file mode 100644 index 0000000000..f1a1aa6db1 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AsyncLoadHoistingPass.h @@ -0,0 +1,50 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_ASYNC_LOAD_HOISTING_PASS_H +#define TRITON_ADAPTER_ASYNC_LOAD_HOISTING_PASS_H + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace triton { + +class AsyncLoadHoistingPass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AsyncLoadHoistingPass) + + AsyncLoadHoistingPass() = default; + + StringRef getArgument() const override { return "async-load-hoisting"; } + + void runOnOperation() override; +}; + +std::unique_ptr> createAsyncLoadHoistingPass(); +void registerAsyncLoadHoistingPasses(); + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_ASYNC_LOAD_HOISTING_PASS_H diff --git a/third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromComputePass.h b/third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromComputePass.h new file mode 100644 index 0000000000..406088340e --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/SeparateMemoryFromComputePass.h @@ -0,0 +1,49 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_SEPARATE_MEMORY_FROM_COMPUTE_PASS_H +#define TRITON_ADAPTER_SEPARATE_MEMORY_FROM_COMPUTE_PASS_H + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace triton { + +/// Sub-pipeline pass that separates memory from compute operations. +class SeparateMemoryFromComputePass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SeparateMemoryFromComputePass) + + SeparateMemoryFromComputePass() = default; + + void runOnOperation() override; +}; + +std::unique_ptr> createSeparateMemoryFromComputePass(); + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_SEPARATE_MEMORY_FROM_COMPUTE_PASS_H diff --git a/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/AddBlockIdForControlOps.h b/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/AddBlockIdForControlOps.h new file mode 100644 index 0000000000..27d22418a5 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/AddBlockIdForControlOps.h @@ -0,0 +1,51 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_ADD_BLOCK_ID_FOR_CONTROL_OPS_H +#define TRITON_ADAPTER_ADD_BLOCK_ID_FOR_CONTROL_OPS_H + +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace triton { +// Define pass +// Pass for adding block_id attributes to control flow operations +class AddBlockIdForControlOpsPass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AddBlockIdForControlOpsPass) + + AddBlockIdForControlOpsPass() = default; + + // Run the pass + void runOnOperation() override; + +private: +}; + +std::unique_ptr> createAddBlockIdForControlOpsPass(); +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_ADD_BLOCK_ID_FOR_CONTROL_OPS_H diff --git a/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.h b/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.h new file mode 100644 index 0000000000..761b21deb6 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.h @@ -0,0 +1,136 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_DATA_DEPENDENCY_ANALYSIS_H +#define TRITON_ADAPTER_DATA_DEPENDENCY_ANALYSIS_H + +#include "mlir/Analysis/DataFlowFramework.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringRef.h" + +namespace mlir { +namespace triton { +enum class DependencyType { VectorToCube, CubeToVector }; + +struct BlockInfo { + int blockId; + bool isCube; + bool isControl; + llvm::SmallVector inputs; + llvm::SmallVector outputs; + llvm::SmallVector Operations; +}; + +struct DependencyInfo { + DependencyType type; + mlir::Value value; + + int producerBlockId; + int consumerBlockId; + int iniProducerBlockId; + int iniConsumerBlockId; +}; + +class DataDependencyInfo { +public: + explicit DataDependencyInfo(mlir::Operation *op) {} + + bool isValid() const { return valid; } + + // for MLIR Analysis framework + bool isInvalidated(const mlir::AnalysisManager::PreservedAnalyses &pa) { + return false; + } + + static bool classof(const DataDependencyInfo *info) { return true; } + + // get analyze result + llvm::DenseMap &getBlockInfoMap() { return blockInfoMap; } + llvm::SmallVector &getV2CDependencies() { + return v2cDependencies; + } + llvm::SmallVector &getC2VDependencies() { + return c2vDependencies; + } + llvm::SmallVector &getMemoryDependencies() { + return memoryDependencies; + } + + void setValid(bool v) { valid = v; } + +private: + bool valid = false; + llvm::DenseMap blockInfoMap; + llvm::SmallVector v2cDependencies; + llvm::SmallVector c2vDependencies; + llvm::SmallVector c2cDependencies; + llvm::SmallVector memoryDependencies; +}; + +// Define pass +// Pass for analyzing data dependencies between Vector and Cube blocks +class DataDependencyAnalysisPass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DataDependencyAnalysisPass) + + DataDependencyAnalysisPass() = default; + + // Run the pass + void runOnOperation(); + +private: + void createBlockInfoMap(DataDependencyInfo &info); + void collectBlockInfo(DataDependencyInfo &info, int blockId, + llvm::SmallVector &ops); + + void collectDepInfo(mlir::Value depvalue, DependencyType dependencyType, + llvm::SmallVector &dependencies, + int iniProdId, int iniConsId, DataDependencyInfo &info); + void analyzeExternalInputs(DataDependencyInfo &info); + void analyzeExternalOutputs(DataDependencyInfo &info); + + void analyzeMemoryEffect(DataDependencyInfo &info); + std::pair findCommonLevelBlockIds(DataDependencyInfo &info, + int producerBlockId, + int consumerBlockId); + + bool isControlFlowOp(mlir::Operation *op); + + mlir::ModuleOp module; +}; + +std::unique_ptr> createDataDependencyAnalysisPass(); + +// Helper: Get BlockId +int getSsbufferBlockId(Operation *op); +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_DATA_DEPENDENCY_ANALYSIS_H diff --git a/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/InterCoreTransferAndSync.h b/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/InterCoreTransferAndSync.h new file mode 100644 index 0000000000..4d989ed007 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/InterCoreTransferAndSync.h @@ -0,0 +1,144 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_INTER_CORE_TRANSFER_AND_SYNC_H +#define TRITON_ADAPTER_INTER_CORE_TRANSFER_AND_SYNC_H + +#include "ascend/include/DynamicCVPipeline/Common/FlagIdManager.h" +#include "ascend/include/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Pass/PassManager.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +namespace mlir { +namespace triton { +// Define pass +class InterCoreTransferAndSyncPass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InterCoreTransferAndSyncPass) + + InterCoreTransferAndSyncPass() = default; + + // Declare dependent dialects + void getDependentDialects(DialectRegistry ®istry) const override; + + // Run the pass + void runOnOperation() override; + +private: + mlir::ModuleOp module; + int transferIndex = 0; + int markAllocIndex = 0; + + llvm::DenseMap vecValueMapping; + llvm::DenseMap cubeValueMapping; + + mlir::LogicalResult processDependencies(FlagIdManager &flagManager); + mlir::LogicalResult + handleVectorToCube(mlir::OpBuilder &builder, DependencyInfo &dep, + llvm::DenseMap vecvalueMapping, + llvm::DenseMap cubeValueMapping, + FlagIdManager &flagManager); + mlir::LogicalResult + handleCubeToVector(mlir::OpBuilder &builder, DependencyInfo &dep, + llvm::DenseMap cubeValueMapping, + FlagIdManager &flagManager); + mlir::LogicalResult + handleMemoryDependency(mlir::OpBuilder &builder, DependencyInfo &dep, + size_t depIndex, + llvm::SmallVector memDependencies, + FlagIdManager &flagManager); + + std::pair + getBlockStartEnd(int blockId, mlir::ModuleOp module); + bool + isOuterLayerDependency(size_t depIndex, mlir::Operation *currProdEnd, + mlir::Operation *currConsStart, + llvm::SmallVector &memDependencies); + + SmallVector computeExpectedShape(mlir::Value value); + bool isShapeExpected(mlir::Value value, + llvm::SmallVector &expectedShape); + mlir::Value normalizeIfNeeded(mlir::OpBuilder &builder, DependencyInfo &dep, + mlir::Location loc, mlir::Value origValue, + llvm::SmallVector expectedShape, + int originBlockId); + void Nd2NzNormalize(mlir::OpBuilder &builder, DependencyInfo &dep, + mlir::Location loc); + void rewriteMatmulWithNewShape(mlir::OpBuilder &builder, + mlir::Operation *matmulOp, mlir::Location loc); + void rewriteTransposeWithNewShape(mlir::OpBuilder &builder, + mlir::Operation *transposeOp, + mlir::Location loc); + llvm::DenseMap getVecValueMapping() { + return vecValueMapping; + } + llvm::DenseMap getCubeValueMapping() { + return cubeValueMapping; + } + + mlir::Operation *annotateTightlyCoupledBuffer(mlir::OpBuilder &builder, + mlir::Operation *allocOp, + mlir::Location loc); + mlir::Operation *findMainLoopforTransfer(mlir::Operation *endOp, + mlir::Operation *startOp); + std::pair + createTransferAllocs(mlir::OpBuilder &builder, mlir::Location loc, + llvm::ArrayRef shape, mlir::Type elemType, + hivm::AddressSpace addrSpace, mlir::Operation *prodEndOp, + mlir::Operation *consStartOp, int prodBlockId, + int consBlockId, llvm::StringRef prodTag, + llvm::StringRef consTag, int transferIndex); + mlir::Operation * + insertVectorToCubeTransfer(mlir::OpBuilder &builder, mlir::Value srcValue, + mlir::Value normalizedValue, + mlir::Operation *vectorEndOp, + mlir::Operation *cubeStartOp, mlir::Location loc, + int transferIndex, int iniConsumerId); + mlir::Operation * + insertCubeToVectorTransfer(mlir::OpBuilder &builder, mlir::Value srcValue, + mlir::Operation *cubeEndOp, + mlir::Operation *vectorStartOp, mlir::Location loc, + int transferIndex, int iniConsumerId); + void insertInterCoreSync(mlir::OpBuilder &builder, + mlir::Operation *transferOp, + mlir::Operation *consumerStartOp, + mlir::Operation *consumerEndOp, int flag, + mlir::Location loc, int transferIndex); + void insertPipeSSync(mlir::OpBuilder &builder, mlir::Operation *producerOp, + mlir::Operation *consumerOp, int flag, + mlir::Location loc, bool isCubeToVector); +}; + +std::unique_ptr> createInterCoreTransferAndSyncPass(); +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_INTER_CORE_TRANSFER_AND_SYNC_H diff --git a/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/MarkMainLoop.h b/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/MarkMainLoop.h new file mode 100644 index 0000000000..383ed32691 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/MarkMainLoop.h @@ -0,0 +1,51 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_MARK_MAIN_LOOP_H +#define TRITON_ADAPTER_MARK_MAIN_LOOP_H + +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace triton { +// Define pass +// Pass for marking the main computation loop in the module +class MarkMainLoopPass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MarkMainLoopPass) + + MarkMainLoopPass() = default; + + // Run the pass + void runOnOperation() override; + +private: +}; + +std::unique_ptr> createMarkMainLoopPass(); +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_MARK_MAIN_LOOP_H diff --git a/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/PreserveControlAttrsCanonicalize.h b/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/PreserveControlAttrsCanonicalize.h new file mode 100644 index 0000000000..eea6b3ec60 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/PreserveControlAttrsCanonicalize.h @@ -0,0 +1,68 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_PRESERVE_CONTROL_ATTRS_CANONICALIZE_H +#define TRITON_ADAPTER_PRESERVE_CONTROL_ATTRS_CANONICALIZE_H + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace triton { + +class PreserveControlAttrsCanonicalizePass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID( + PreserveControlAttrsCanonicalizePass) + + PreserveControlAttrsCanonicalizePass() = default; + + void runOnOperation() override; + + static constexpr ::llvm::StringRef getArgumentName() { + return "preserve-control-attrs-canonicalize"; + } + + ::llvm::StringRef getArgument() const override { + return "preserve-control-attrs-canonicalize"; + } + + ::llvm::StringRef getDescription() const override { + return "Canonicalize while preserving selected control-flow attrs"; + } + + ::llvm::StringRef getName() const override { + return "PreserveControlAttrsCanonicalizePass"; + } +}; + +std::unique_ptr> +createPreserveControlAttrsCanonicalizePass(); + +void registerPreserveControlAttrsCanonicalizePasses(); + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_PRESERVE_CONTROL_ATTRS_CANONICALIZE_H diff --git a/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/SeparateCVScope.h b/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/SeparateCVScope.h new file mode 100644 index 0000000000..0ae78d33c2 --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/SplitDataflow/SeparateCVScope.h @@ -0,0 +1,49 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_SEPARATE_CV_SCOPE_H +#define TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_SEPARATE_CV_SCOPE_H + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace triton { + +class SeparateCVScopePass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SeparateCVScopePass) + + SeparateCVScopePass() = default; + + void getDependentDialects(DialectRegistry ®istry) const override; + + void runOnOperation() override; +}; + +std::unique_ptr> createSeparateCVScopePass(); + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_SEPARATE_CV_SCOPE_H diff --git a/third_party/ascend/include/DynamicCVPipeline/SplitDataflowPass.h b/third_party/ascend/include/DynamicCVPipeline/SplitDataflowPass.h new file mode 100644 index 0000000000..79e07cae0e --- /dev/null +++ b/third_party/ascend/include/DynamicCVPipeline/SplitDataflowPass.h @@ -0,0 +1,50 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_SPLIT_DATAFLOW_PASS_H +#define TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_SPLIT_DATAFLOW_PASS_H + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace triton { + +class SplitDataflowPass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SplitDataflowPass) + + // Constructor + SplitDataflowPass() = default; + + // Run the pass + void runOnOperation() override; +}; + +// Create the pass +std::unique_ptr> createSplitDataflowPass(); + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_SPLIT_DATAFLOW_PASS_H diff --git a/third_party/ascend/include/TritonAffinityOpt/CMakeLists.txt b/third_party/ascend/include/TritonAffinityOpt/CMakeLists.txt new file mode 100644 index 0000000000..c6193d1f5d --- /dev/null +++ b/third_party/ascend/include/TritonAffinityOpt/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name TritonAffinityOpt) +add_public_tablegen_target(TritonAffinityOptConversionPassIncGen) diff --git a/third_party/ascend/include/TritonAffinityOpt/DAG.h b/third_party/ascend/include/TritonAffinityOpt/DAG.h new file mode 100644 index 0000000000..ff2262fd76 --- /dev/null +++ b/third_party/ascend/include/TritonAffinityOpt/DAG.h @@ -0,0 +1,304 @@ +#ifndef AffinityDAGDEF +#define AffinityDAGDEF +#include "Utils.hpp" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/Value.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/TinyPtrVector.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/ErrorHandling.h" +#include +#include +#include +#include +#include +#include + +namespace mlir { +namespace AffinityDAG { + +enum class OpAbility { + PREFER_VECTOR = 1 << 0, + CUBE_ONLY = 1 << 1, + CUBE_AND_VECTOR = PREFER_VECTOR | CUBE_ONLY + +}; + +enum CoreType { + UNDETERMINED = 0, + VECTOR_ONLY = 1 << 0, + CUBE_ONLY = 1 << 1, + CUBE_AND_VECTOR = VECTOR_ONLY | CUBE_ONLY +}; + +inline constexpr CoreType toCoreType(OpAbility ct) { + using U = std::underlying_type_t; + return static_cast(static_cast(ct)); +} + +constexpr inline CoreType operator|(CoreType lhs, CoreType rhs) { + return enumOp(std::bit_or<>(), lhs, rhs); +} + +inline CoreType operator&(CoreType lhs, CoreType rhs) { + return enumOp(std::bit_and<>(), lhs, rhs); +} + +inline bool intersects(CoreType lhs, CoreType rhs) { + return (lhs & rhs) != CoreType::UNDETERMINED; +} + +inline CoreType operator&(OpAbility lhs, CoreType rhs) { + return toCoreType(lhs) & rhs; +} + +inline CoreType operator!(CoreType ct) { + CoreType newCt = UNDETERMINED; + if ((ct & CoreType::CUBE_ONLY) == UNDETERMINED) { + newCt = newCt | CoreType::CUBE_ONLY; + } + + if ((ct & CoreType::VECTOR_ONLY) == UNDETERMINED) { + newCt = newCt | CoreType::VECTOR_ONLY; + } + + return newCt; +} + +inline hivm::TCoreType toHivm(CoreType ct) { + switch (ct) { + case UNDETERMINED: + return hivm::TCoreType::CUBE_OR_VECTOR; + case CUBE_ONLY: + return hivm::TCoreType::CUBE; + case VECTOR_ONLY: + return hivm::TCoreType::VECTOR; + case CUBE_AND_VECTOR: + return hivm::TCoreType::CUBE_AND_VECTOR; + default: + llvm_unreachable("Invalid CoreType that cannot convert to hivm"); + } +} + +inline bool intersects(OpAbility lhs, CoreType rhs) { + return (lhs & rhs) != CoreType::UNDETERMINED; +} + +inline bool exactlyOneType(CoreType ct) { + return (ct == CUBE_ONLY) || (ct == VECTOR_ONLY); +} + +const char *literalCoreType(CoreType ct); + +class MoveOnly { +protected: + MoveOnly() = default; + ~MoveOnly() = default; + + MoveOnly(const MoveOnly &) = delete; + MoveOnly &operator=(const MoveOnly &) = delete; + + MoveOnly(MoveOnly &&) = default; + MoveOnly &operator=(MoveOnly &&) = default; +}; + +class Node; +class OpNode; +class ValueNode; + +ValueNode *getDataSource(OpNode *op); + +class Graph : MoveOnly { +public: + using OpMapRaw = llvm::DenseMap>; + using ValueMapRaw = llvm::DenseMap>; + using OpMap = std::shared_ptr; + using ValueMap = std::shared_ptr; + + Graph(Block *block, Graph *parent = nullptr, OpMap opMap = nullptr, + ValueMap valueMap = nullptr, bool inheritParent = true); + + static std::unique_ptr fromMultiBlockFunc(triton::FuncOp funcOp); + + OpMapRaw &getOpMap() const { return *opMap; } + + ValueMapRaw &getValueMap() const { return *valueMap; } + + // [DEBUG] start + std::unique_ptr> legacyOpMap = nullptr; + std::unique_ptr> legacyValueTypes = nullptr; + + inline llvm::DenseMap &getOpMapLegacy() { + if (!legacyOpMap) { + legacyOpMap = + std::move(std::make_unique>()); + for (auto &[key, val] : *opMap) { + (*legacyOpMap)[key] = val.get(); + } + } + + return *legacyOpMap; + } + + llvm::DenseMap &getValueTypes(); + + // [DEBUG] end + +private: + friend class Node; + friend class OpNode; + OpMap opMap; + ValueMap valueMap; + Block *block; + Graph *parent; + OpNode *terminator = nullptr; + size_t opCount = 0; + llvm::SmallVector blockArgs; +}; + +class Node : MoveOnly { +protected: + friend class Graph; + friend class ValueNode; + bool isUpstreamOfCubeMem = false; + virtual CoreType absorbImpl() = 0; + llvm::SmallVector outputs; + +public: + CoreType isOnPrivate = UNDETERMINED; + + enum NodeKind { NK_Op, NK_Value }; + + inline CoreType isOn() const { return isOnPrivate; } + + bool absorb() { + auto newCoreType = absorbImpl(); + auto changed = newCoreType != isOnPrivate; + isOnPrivate = newCoreType; + + return changed; + }; + + virtual llvm::SmallVector getAffected() const = 0; + virtual OpNode *getSourceOpNode() = 0; + + ArrayRef getOutputs() const { return outputs; } + + CoreType absorbCommon(); + +private: + const NodeKind kind; + +public: + NodeKind getKind() const { return kind; } + +protected: + Node(NodeKind kind) : kind(kind) {} +}; + +class OpNode : public Node { + friend class Graph; + friend class ValueNode; + llvm::SmallVector inputs; + llvm::SmallVector subgraphs; + virtual CoreType absorbImpl() override; + +public: + Operation *op; + + OpNode(Operation *op, Graph *graph); + OpAbility canRunOn() const; + inline ArrayRef getInputs() const { return inputs; } + + static bool classof(const Node *node) { return node->getKind() == NK_Op; } + + virtual llvm::SmallVector getAffected() const override { + llvm::SmallVector result(inputs.begin(), inputs.end()); + result.append(outputs.begin(), outputs.end()); + + return result; + } + + virtual OpNode *getSourceOpNode() override { return this; } +}; + +class ValueNode : public Node { + friend class Graph; + friend class OpNode; + virtual CoreType absorbImpl() override; + +public: + Node *source = nullptr; + Value value; + // ValueNode(OpResult value); + // ValueNode(BlockArgument value); + + ValueNode(Value value) : Node(NK_Value), value(value) {}; + virtual OpNode *getSourceOpNode() override { + if (!source) { + return nullptr; + } + + return source->getSourceOpNode(); + } + static bool classof(const Node *node) { return node->getKind() == NK_Value; } + + virtual llvm::SmallVector getAffected() const override { + llvm::SmallVector result(outputs.begin(), outputs.end()); + if (source) + result.push_back(source); + + return result; + } +}; + +class GraphManager { +private: + llvm::DenseMap> graphs; + +public: + static GraphManager &getInstance() { + static GraphManager instance; + return instance; + } + + void registerGraph(llvm::StringRef funcName, + std::shared_ptr graph) { + graphs[funcName] = graph; + } + + AffinityDAG::Graph *getGraph(llvm::StringRef funcName) { + auto it = graphs.find(funcName); + return it != graphs.end() ? it->second.get() : nullptr; + } + + void removeGraph(llvm::StringRef funcName) { graphs.erase(funcName); } +}; + +inline llvm::DenseMap &Graph::getValueTypes() { + static std::mutex mtx; + std::lock_guard lock(mtx); + if (!legacyValueTypes) { + legacyValueTypes = + std::move(std::make_unique>()); + for (auto &[key, val] : *valueMap) { + llvm::dbgs() << key << "\n"; + llvm::dbgs().flush(); + (*legacyValueTypes)[key] = val.get()->isOn(); + } + } + + return *legacyValueTypes; +} + +} // namespace AffinityDAG +} // namespace mlir +#endif diff --git a/third_party/ascend/include/TritonAffinityOpt/Passes.h b/third_party/ascend/include/TritonAffinityOpt/Passes.h new file mode 100644 index 0000000000..f58c7563bc --- /dev/null +++ b/third_party/ascend/include/TritonAffinityOpt/Passes.h @@ -0,0 +1,47 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_TRITON_AFFINITY_OPTIMIZATION_PASSES_H +#define TRITON_ADAPTER_TRITON_AFFINITY_OPTIMIZATION_PASSES_H + +#include "mlir/Pass/Pass.h" + +namespace mlir { +// Forward declarations. +class ModuleOp; + +namespace triton { + +/// Creates a pass to convert Triton dialect to Annotation dialect. +std::unique_ptr> createDAGSSBufferPass(); + +std::unique_ptr> createDAGSyncPass(); + +std::unique_ptr> createDAGScopePass(); + +#define GEN_PASS_REGISTRATION +#include "ascend/include/TritonAffinityOpt/Passes.h.inc" + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_TRITON_AFFINITY_OPTIMIZATION_PASSES_H diff --git a/third_party/ascend/include/TritonAffinityOpt/Passes.td b/third_party/ascend/include/TritonAffinityOpt/Passes.td new file mode 100644 index 0000000000..f12de8444e --- /dev/null +++ b/third_party/ascend/include/TritonAffinityOpt/Passes.td @@ -0,0 +1,29 @@ +/* + * Copyright (c) Huawei Technologies Co. + * Licensed under the MIT license. + */ + +#ifndef TRITON_AFFINITY_OPTIMIZATION_PASSES +#define TRITON_AFFINITY_OPTIMIZATION_PASSES + +include "mlir/Pass/PassBase.td" + +def DAGSSBuffer : Pass<"dag-ssbuf", "mlir::ModuleOp"> { + let summary = "Convert vector operations to shared storage buffer operations"; + let constructor = "triton::createDAGSSBufferPass()"; + let dependentDialects = ["hivm::HIVMDialect", "bufferization::BufferizationDialect", "scope::ScopeDialect", "annotation::AnnotationDialect"]; +} + +def DAGScope : Pass<"dag-scope", "mlir::ModuleOp"> { + let summary = "Convert native triton code to NPU-affine code"; + let constructor = "triton::createDAGScopePass()"; + let dependentDialects = ["hivm::HIVMDialect", "bufferization::BufferizationDialect", "scope::ScopeDialect", "annotation::AnnotationDialect"]; +} + +def DAGSync : Pass<"dag-sync", "mlir::ModuleOp"> { + let summary = "DAG sync"; + let constructor = "triton::createDAGSyncPass()"; + let dependentDialects = ["hivm::HIVMDialect", "bufferization::BufferizationDialect", "annotation::AnnotationDialect"]; +} + +#endif // TRITON_AFFINITY_OPTIMIZATION_PASSES diff --git a/third_party/ascend/include/TritonAffinityOpt/Utils.hpp b/third_party/ascend/include/TritonAffinityOpt/Utils.hpp new file mode 100644 index 0000000000..6957108258 --- /dev/null +++ b/third_party/ascend/include/TritonAffinityOpt/Utils.hpp @@ -0,0 +1,64 @@ +#ifndef TRITON_AFFINITY_UTILS_HPP +#define TRITON_AFFINITY_UTILS_HPP + +#include +#include +#include +#include + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/STLExtras.h" + +namespace mlir::AffinityDAG { + +template +constexpr inline T enumOp(F &&func, T lhs, T rhs) { + static_assert(std::is_enum_v, "T must be an enum type"); + + using U = std::underlying_type_t; + + return static_cast(std::invoke(std::forward(func), static_cast(lhs), + static_cast(rhs))); +} + +// since we do not have llvm::set_intersects in this version... +template bool intersects(S1Ty &s1, S2Ty &s2) { + if (s1.size() > s2.size()) { + return intersects(s2, s1); + } + + return llvm::any_of(s1, [&](auto e) { return s2.count(e); }); +} + +/** + * @returns the pointer to the value wrapped by the pointer if the key is in the + * map, otherwise nullptr + * + * @safety The user is responsible for checking the nullity of the retured + * pointer; the lifespan of the pointer is valid as long as the value belongs to + * the map + */ +template +inline auto getFromSmartPtr(MapTy &map, const typename MapTy::key_type &key) + -> decltype(map.find(key)->second.get()) { + auto it = map.find(key); + return it != map.end() ? it->second.get() : nullptr; +} + +/** + * @returns the pointer to the value if the key is in the map, otherwise nullptr + * + * @safety The user is responsible for checking the nullity of the retured + * pointer; the lifespan of the pointer is valid as long as the the map is not + * modified + */ +template +inline auto getPtr(MapTy &map, const typename MapTy::key_type &key) + -> decltype(map.find(key)->second) * { + auto it = map.find(key); + return it != map.end() ? &it->second : nullptr; +} + +} // namespace mlir::AffinityDAG + +#endif diff --git a/third_party/ascend/include/TritonToAnnotation/CMakeLists.txt b/third_party/ascend/include/TritonToAnnotation/CMakeLists.txt new file mode 100644 index 0000000000..d8cdc5dcbb --- /dev/null +++ b/third_party/ascend/include/TritonToAnnotation/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name TritonToAnnotation) +add_public_tablegen_target(TritonToAnnotationConversionPassIncGen) diff --git a/third_party/ascend/include/TritonToAnnotation/Passes.h b/third_party/ascend/include/TritonToAnnotation/Passes.h new file mode 100644 index 0000000000..2899e4c008 --- /dev/null +++ b/third_party/ascend/include/TritonToAnnotation/Passes.h @@ -0,0 +1,43 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_TRITON_TO_ANNOTATION_CONVERSION_PASSES_H +#define TRITON_ADAPTER_TRITON_TO_ANNOTATION_CONVERSION_PASSES_H + +#include "mlir/Pass/Pass.h" + +namespace mlir { +// Forward declarations. +class ModuleOp; + +namespace triton { + +/// Creates a pass to convert Triton dialect to Annotation dialect. +std::unique_ptr> createTritonToAnnotationPass(); + +#define GEN_PASS_REGISTRATION +#include "ascend/include/TritonToAnnotation/Passes.h.inc" + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_TRITON_TO_ANNOTATION_CONVERSION_PASSES_H diff --git a/third_party/ascend/include/TritonToAnnotation/Passes.td b/third_party/ascend/include/TritonToAnnotation/Passes.td new file mode 100644 index 0000000000..871de7f498 --- /dev/null +++ b/third_party/ascend/include/TritonToAnnotation/Passes.td @@ -0,0 +1,17 @@ +/* + * Copyright (c) Huawei Technologies Co. + * Licensed under the MIT license. + */ + +#ifndef TRITON_TO_ANNOTATION_CONVERSION_PASSES +#define TRITON_TO_ANNOTATION_CONVERSION_PASSES + +include "mlir/Pass/PassBase.td" + +def TritonToAnnotation : Pass<"triton-to-annotation", "mlir::ModuleOp"> { + let summary = "Convert Triton to Annotation dialect"; + let constructor = "triton::createTritonToAnnotationPass()"; + let dependentDialects = ["annotation::AnnotationDialect"]; +} + +#endif // TRITON_TO_ANNOTATION_CONVERSION_PASSES diff --git a/third_party/ascend/include/TritonToGraph/AliasAnalysis.h b/third_party/ascend/include/TritonToGraph/AliasAnalysis.h new file mode 100644 index 0000000000..2f885ff7eb --- /dev/null +++ b/third_party/ascend/include/TritonToGraph/AliasAnalysis.h @@ -0,0 +1,160 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_TO_CFG_ALIAS_ANALYSIS_H +#define TRITON_TO_CFG_ALIAS_ANALYSIS_H + +#include "TritonToGraph/tensor.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/TypeUtilities.h" +#include "mlir/IR/Types.h" +#include "mlir/IR/Value.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Types.h" +#include "llvm/ADT/DenseMap.h" + +namespace mlir { +namespace triton { +namespace cfg { + +class ControlFlowGraph; + +// AliasAnalysis - Alias分析和指针跟踪 +// 用于分析pointer的alias关系，跟踪全局内存对象（如gm_obj）的别名链 +class AliasAnalysis { +public: + // 为整个CFG分析pointer别名 + void analyzePointerAliases(ControlFlowGraph &cfg); + + // 获取Value的base pointer（递归查找真实的基础指针） + Value getBasePointer(Value ptr) const; + + // 获取Value对应的TensorObject + TensorObject *getTensorObject(Value value) const { + auto it = baseTensorMap.find(value); + return (it != baseTensorMap.end()) ? it->second : nullptr; + } + + // 判断两个指针是否可能指向同一tensor + bool mayAlias(Value ptr1, Value ptr2) const { + return getBasePointer(ptr1) == getBasePointer(ptr2); + } + + // 添加alias关系 + void addAlias(Value ptr, Value base, TensorObject *tensor) { + aliasMap[ptr] = base; + baseTensorMap[ptr] = tensor; + } + + // 判断是否是指针类型 + static bool isPointerType(Type type) { + if (auto ptrType = + mlir::dyn_cast(getElementTypeOrSelf(type))) { + return true; + } + return false; + } + + // 判断是否是tensor pointer类型（指向tensor的指针） + static bool isTensorPointerType(Type type) { + if (auto ptrType = mlir::dyn_cast(type)) { + Type pointeeType = ptrType.getPointeeType(); + return mlir::isa(pointeeType); + } + return false; + } + + // 判断是否是标量指针类型（指向标量的指针） + static bool isScalarPointerType(Type type) { + if (auto ptrType = mlir::dyn_cast(type)) { + Type pointeeType = ptrType.getPointeeType(); + return !mlir::isa(pointeeType); + } + return false; + } + + // 判断是否是全局内存类型 + static bool isGlobalMemoryType(Type type) { + // 在Triton中，全局内存通常由!tt.ptr类型表示 + // 可以进一步根据地址空间等信息判断 + if (auto ptrType = mlir::dyn_cast(type)) { + // 地址空间1通常是全局内存 + // 需要根据具体硬件进行调整 + return ptrType.getAddressSpace() == 1 || ptrType.getAddressSpace() == 0; + } + return false; + } + + // 获取所有tracked的base pointer + const DenseMap &getAliasMap() const { return aliasMap; } + + // 获取所有tracked的tensor对象 + const DenseMap &getBaseTensorMap() const { + return baseTensorMap; + } + + // 打印所有alias信息 + void print(llvm::raw_ostream &os) const { + os << "=== Alias Analysis Result ===\n"; + os << "Total tracked aliases: " << aliasMap.size() << "\n"; + + for (const auto &entry : aliasMap) { + Value ptr = entry.first; + Value base = entry.second; + TensorObject *tensor = baseTensorMap.lookup(ptr); + + os << " " << ptr << " -> " << base; + if (tensor) { + os << " [" << tensor->getName() << "]"; + } + os << "\n"; + } + } + +private: + // 分析addptr操作 + void analyzeAddPtrOp(mlir::triton::AddPtrOp addptrOp); + + // 分析make_tensor_ptr操作 + void analyzeMakeTensorPtrOp(mlir::triton::MakeTensorPtrOp op); + + // 分析load操作 + void analyzeLoadOp(mlir::triton::LoadOp loadOp); + + // 分析store操作 + void analyzeStoreOp(mlir::triton::StoreOp storeOp); + + // 分析broadcast操作 + void analyzeBroadcastOp(mlir::triton::BroadcastOp broadcastOp); + + // 分析splat操作 + void analyzeSplatOp(mlir::triton::SplatOp splatOp); + + DenseMap aliasMap; // ptr -> base ptr + DenseMap baseTensorMap; // value -> tensor +}; + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_CFG_ALIAS_ANALYSIS_H diff --git a/third_party/ascend/include/TritonToGraph/CMakeLists.txt b/third_party/ascend/include/TritonToGraph/CMakeLists.txt new file mode 100644 index 0000000000..5ffe7ae0cb --- /dev/null +++ b/third_party/ascend/include/TritonToGraph/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name TritonToGraph) +add_public_tablegen_target(TritonToGraphPassIncGen) diff --git a/third_party/ascend/include/TritonToGraph/ControlFlowGraph.h b/third_party/ascend/include/TritonToGraph/ControlFlowGraph.h new file mode 100644 index 0000000000..53e32ebb11 --- /dev/null +++ b/third_party/ascend/include/TritonToGraph/ControlFlowGraph.h @@ -0,0 +1,266 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_TO_CFG_CONTROL_FLOW_GRAPH_H +#define TRITON_TO_CFG_CONTROL_FLOW_GRAPH_H + +#include "TritonToGraph/MemorySSA.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/Region.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/MapVector.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/raw_ostream.h" + +#include +#include +#include +#include + +namespace mlir { +namespace triton { +namespace cfg { + +// 前向声明 +class BasicBlock; +class ControlFlowGraph; + +// Instruction 结构体：表示一个指令 +class Instruction { +public: + Instruction(size_t id, Operation *op, BasicBlock *parentBlock) + : id(id), operation(op), parentBlock(parentBlock), memorySSAInfo() {} + + // 获取基本信息 + size_t getId() const { return id; } + Operation *getOperation() const { return operation; } + BasicBlock *getParentBlock() const { return parentBlock; } + + // 检查是否有子图（内部区域） + bool hasSubGraph() const { return subGraph != nullptr; } + ControlFlowGraph *getSubGraph() const { return subGraph.get(); } + void setSubGraph(std::unique_ptr graph) { + subGraph = std::move(graph); + } + + // 获取指令的字符串表示 + std::string getAsString() const; + + // 打印 + void print(raw_ostream &os, unsigned indent = 0) const; + void dump() const; + + // Memory SSA信息 + MemorySSAInfo &getMemorySSAInfo() { return memorySSAInfo; } + const MemorySSAInfo &getMemorySSAInfo() const { return memorySSAInfo; } + +private: + size_t id; // 唯一ID + Operation *operation; // 对应的 MLIR Operation + BasicBlock *parentBlock; // 所属的 BasicBlock + std::unique_ptr + subGraph; // 子图（用于 reduce 等有内部区域的操作） + MemorySSAInfo memorySSAInfo; // Memory SSA信息（用于tensor/pointer分析） +}; + +// 基本块类型 +enum class BlockType { + NORMAL, // 普通块（包含多个指令） + ENTRY, // 函数入口块 + EXIT, // 函数出口块 + IF_COND, // if 条件判断块（包含单个 scf.if 指令） + FOR_COND, // for 循环头块（包含单个 scf.for 指令） + WHILE_COND, // while 条件块（包含单个 scf.while 指令） + COND_BR, // cf.cond_br 条件分支块（包含单个 cf.cond_br 指令） + BR, // cf.br 无条件跳转块（包含单个 cf.br 指令） + LOOP_BODY, // 循环体块 + LOOP_EXIT, // 循环出口块 +}; + +// 基本块节点 +class BasicBlock { +public: + BasicBlock(size_t id, BlockType type, BasicBlock *parentStructure = nullptr) + : id(id), type(type), parentStructure(parentStructure) {} + + // 获取基本信息 + size_t getId() const { return id; } + BlockType getType() const { return type; } + void setType(BlockType t) { type = t; } + + // 外层结构（如果是嵌套在 loop/if 中） + BasicBlock *getParentStructure() const { return parentStructure; } + void setParentStructure(BasicBlock *parent) { parentStructure = parent; } + + // 对应控制流结构的出口块（用于 IF_COND/FOR_COND/WHILE_COND） + // 指向 if/for/while 结束后到达的基本块 + BasicBlock *getExitBlock() const { return exitBlock; } + void setExitBlock(BasicBlock *exit) { exitBlock = exit; } + + // Instruction 操作 + void addInstruction(std::unique_ptr inst); + Instruction *getInstruction(size_t idx) const; + size_t getNumInstructions() const { return instructions.size(); } + const SmallVector> &getInstructions() const { + return instructions; + } + + // 检查最后一条指令是否为 ReturnOp + bool endsWithReturnOp() const; + + // 获取后继和前驱 + ArrayRef getSuccessors() const { return successors; } + ArrayRef getPredecessors() const { return predecessors; } + + // 边的操作 + void addSuccessor(BasicBlock *succ); + void addPredecessor(BasicBlock *pred); + + size_t getNumSuccessors() const { return successors.size(); } + size_t getNumPredecessors() const { return predecessors.size(); } + + // 获取名称 + std::string getName() const; + + // 获取类型字符串 + StringRef getTypeString() const; + + // 打印 + void print(raw_ostream &os) const; + void dump() const; + + // 导出为 JSON（用于网页可视化） + void exportToJSON(raw_ostream &os, unsigned indent = 0) const; + +private: + size_t id; // 唯一ID + BlockType type; // 块类型 + BasicBlock *parentStructure; // 外层结构（loop/if）的 basic block + BasicBlock *exitBlock = + nullptr; // 对应控制流结构的出口块（用于 IF_COND/FOR_COND/WHILE_COND） + SmallVector> instructions; // 指令列表 + SmallVector successors; // 后继块指针列表 + SmallVector predecessors; // 前驱块指针列表 +}; + +// 控制流图 +class ControlFlowGraph { +public: + explicit ControlFlowGraph(triton::FuncOp func); + ~ControlFlowGraph(); + + // 获取函数 + triton::FuncOp getFunction() const { return function; } + + // 基本块操作 + BasicBlock *createBasicBlock(BlockType type, + BasicBlock *parentStructure = nullptr); + + BasicBlock *getBasicBlock(size_t id) { + if (id < basicBlocks.size()) + return basicBlocks[id].get(); + return nullptr; + } + const BasicBlock *getBasicBlock(size_t id) const { + if (id < basicBlocks.size()) + return basicBlocks[id].get(); + return nullptr; + } + + size_t getNumBlocks() const { return basicBlocks.size(); } + + // 获取入口和出口块 + BasicBlock *getEntryBlock() { return entryBlock; } + const BasicBlock *getEntryBlock() const { return entryBlock; } + BasicBlock *getExitBlock() { return exitBlock; } + const BasicBlock *getExitBlock() const { return exitBlock; } + + void setEntryBlock(BasicBlock *bb) { entryBlock = bb; } + void setExitBlock(BasicBlock *bb) { exitBlock = bb; } + + // 添加边 + void addEdge(BasicBlock *from, BasicBlock *to); + + // 遍历 + using BlockVisitor = llvm::function_ref; + void traverse(BlockVisitor visitor); + + // Operation 到 Instruction 的查询 + Instruction *getInstruction(Operation *op) const { + auto it = opToInstructionMap.find(op); + return (it != opToInstructionMap.end()) ? it->second : nullptr; + } + + // 添加 Operation 到 Instruction 的映射 + void addOpToInstruction(Operation *op, Instruction *inst) { + opToInstructionMap[op] = inst; + } + + bool isBackEdge(BasicBlock *from, BasicBlock *to) const; + + // 结构化搜索 API + // 从起始块开始，沿着 successor 顺序向下结构化搜索 + // NORMAL 块：遍历每个 Instruction，调用 callback 传入 Operation* + // IF_COND/FOR_COND/WHILE_COND 块：递归访问每个 successor + // 遇到 exitBlock 时停止 + using OperationVisitor = llvm::function_ref; + void searchNormalBlock(BasicBlock *block, OperationVisitor callback) const; + void searchCondBlock(BasicBlock *block, OperationVisitor callback) const; + void searchBlock(BasicBlock *block, OperationVisitor callback) const; + + // 打印 + void print(raw_ostream &os) const; + void dump() const; + + // 导出为 DOT 格式 + void exportDOT(raw_ostream &os) const; + + // 导出到文件 + llvm::Error exportToFile(StringRef filename) const; + + // 导出为 HTML（网页可视化） + llvm::Error exportToHTML(StringRef filename) const; + + // 导出为 JSON + void exportToJSON(raw_ostream &os) const; + +private: + triton::FuncOp function; // 所属函数 + SmallVector> basicBlocks; // 基本块节点 + BasicBlock *entryBlock = nullptr; // 入口块 + BasicBlock *exitBlock = nullptr; // 出口块 + size_t nextBlockId = 0; // 下一个块ID + + // Operation 到 Instruction 的映射（支持快速查询） + std::unordered_map opToInstructionMap; +}; + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_CFG_CONTROL_FLOW_GRAPH_H diff --git a/third_party/ascend/include/TritonToGraph/ControlFlowGraphBuilder.h b/third_party/ascend/include/TritonToGraph/ControlFlowGraphBuilder.h new file mode 100644 index 0000000000..c76ddf2072 --- /dev/null +++ b/third_party/ascend/include/TritonToGraph/ControlFlowGraphBuilder.h @@ -0,0 +1,242 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_TO_CFG_CONTROL_FLOW_GRAPH_BUILDER_H +#define TRITON_TO_CFG_CONTROL_FLOW_GRAPH_BUILDER_H + +#include "TritonToGraph/ControlFlowGraph.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Value.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Pass/PassRegistry.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include +#include + +namespace mlir { +namespace triton { +namespace cfg { + +// 构建控制流图的 Pass +class BuildCFGPass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(BuildCFGPass) + + BuildCFGPass() = default; + + StringRef getArgument() const override { return "build-cfg"; } + StringRef getDescription() const override { + return "Build Control Flow Graph from TTIR"; + } + + void runOnOperation() override; + + // Pass 选项 + std::string outputDir = "."; + + // 获取依赖的方言 + void getDependentDialects(DialectRegistry ®istry) const override { + registry.insert(); + } + +protected: + // 构建单个函数的 CFG + std::unique_ptr buildForFunction(triton::FuncOp func); +}; + +// 添加数据结构定义 +// IF_COND 的 yield value 和 result value 的对应关系 +struct IfYieldResultMapping { + // then 分支的 yield values（来自 scf.yield 操作的操作数） + SmallVector thenYieldValues; + // else 分支的 yield values（如果有 else 分支） + SmallVector elseYieldValues; + // if 操作的 result values + SmallVector resultValues; + // 对应关系: resultValues[i] 对应 thenYieldValues[i] 或 elseYieldValues[i] +}; + +// FOR_COND 的 yield value 和 iter args value 的对应关系 +struct ForYieldIterArgMapping { + // yield 操作的 values（来自循环体末尾的 scf.yield） + SmallVector yieldValues; + // iter_args（循环初始参数，对应 for 操作的 iter_args） + SmallVector iterArgValues; + // for 操作的 result values + SmallVector resultValues; + // 对应关系: iterArgValues[i] 在循环体中使用时被更新，yieldValues[i] 是新的值 + // resultValues[i] 对应最后一次迭代的 yieldValues[i] +}; + +// COND_BR 的 true/false 分支信息 +struct CondBranchMapping { + // true 分支的目标块参数值（来自 cf.cond_br 的 trueOperands） + SmallVector trueOperands; + // false 分支的目标块参数值（来自 cf.cond_br 的 falseOperands） + SmallVector falseOperands; + // 条件值 + Value condition; + // true 分支目标块 + Block *trueDest; + // false 分支目标块 + Block *falseDest; +}; + +// BR 的分支信息 +struct BranchMapping { + // 目标块参数值（来自 cf.br 的 destOperands） + SmallVector destOperands; + // 目标块 + Block *dest; +}; + +// 独立的 CFG 构建器类（用于非 Pass 场景） +class ControlFlowGraphBuilder { +public: + // 为函数构建 CFG + std::unique_ptr build(triton::FuncOp func); + + // 为模块构建所有函数的 CFG + std::vector> + buildForModule(ModuleOp module); + + // 处理一个 region，返回该 region 的入口块和出口块 + struct RegionBlocks { + cfg::BasicBlock *entryBlock; + cfg::BasicBlock *exitBlock; + }; + + RegionBlocks buildForRegion(Region ®ion, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *entryBlock, + cfg::BasicBlock *parentStructure = nullptr); + + // 处理 block 中的操作，返回最后处理的基本块 + cfg::BasicBlock *processBlock(Block &block, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *currentBB, + cfg::BasicBlock *parentStructure = nullptr); + + // 处理 scf.if 操作，返回 if 后面的基本块 + cfg::BasicBlock *handleIfOp(scf::IfOp ifOp, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *currentBB, + cfg::BasicBlock *parentStructure = nullptr); + + // 处理 scf.for 操作，返回 for 后面的基本块 + cfg::BasicBlock *handleForOp(scf::ForOp forOp, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *currentBB, + cfg::BasicBlock *parentStructure = nullptr); + + // 处理 scf.while 操作，返回 while 后面的基本块 + cfg::BasicBlock *handleWhileOp(scf::WhileOp whileOp, + cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *currentBB, + cfg::BasicBlock *parentStructure = nullptr); + + // 处理 cf.cond_br 操作，返回条件分支后面的基本块 + cfg::BasicBlock * + handleCondBranchOp(cf::CondBranchOp condBrOp, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *currentBB, + cfg::BasicBlock *parentStructure = nullptr); + + // 处理 cf.br 操作，返回无条件跳转后面的基本块 + cfg::BasicBlock *handleBranchOp(cf::BranchOp brOp, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *currentBB, + cfg::BasicBlock *parentStructure = nullptr); + + // 创建一个新的指令并添加到 basic block + cfg::Instruction *createInstruction(Operation *op, + cfg::BasicBlock *parentBlock, + cfg::ControlFlowGraph &cfg); + + // 1. 快速收集所有的 IF_COND 基本块 + // 遍历 CFG 中所有基本块，返回类型为 IF_COND 的基本块列表 + SmallVector + collectIfCondBlocks(cfg::ControlFlowGraph &cfg); + + // 2. 快速收集所有的 FOR_COND 基本块 + // 遍历 CFG 中所有基本块，返回类型为 FOR_COND 的基本块列表 + SmallVector + collectForCondBlocks(cfg::ControlFlowGraph &cfg); + + // 3. 获取 IF_COND 对应的 yield value 和 result value 的对应关系 + // 参数: IF_COND 类型的基本块 + // 返回: IfYieldResultMapping 结构体，包含 then/else 的 yield values 和 result + // values + std::optional + getIfYieldResultMapping(cfg::BasicBlock *ifCondBB); + + // 4. 获取 FOR_COND 对应的 yield value 和 iter args value 的对应关系 + // 参数: FOR_COND 类型的基本块 + // 返回: ForYieldIterArgMapping 结构体，包含 yield values、iter_args 和 result + // values + std::optional + getForYieldIterArgMapping(cfg::BasicBlock *forCondBB); + + // 5. 快速收集所有的 COND_BR 基本块 + // 遍历 CFG 中所有基本块，返回类型为 COND_BR 的基本块列表 + SmallVector + collectCondBrBlocks(cfg::ControlFlowGraph &cfg); + + // 6. 获取 COND_BR 对应的条件分支信息 + // 参数: COND_BR 类型的基本块 + // 返回: CondBranchMapping 结构体，包含条件、目标块和参数信息 + std::optional + getCondBranchMapping(cfg::BasicBlock *condBrBB); + + // 7. 快速收集所有的 BR 基本块 + // 遍历 CFG 中所有基本块，返回类型为 BR 的基本块列表 + SmallVector collectBrBlocks(cfg::ControlFlowGraph &cfg); + + // 8. 获取 BR 对应的分支信息 + // 参数: BR 类型的基本块 + // 返回: BranchMapping 结构体，包含目标块和参数信息 + std::optional getBranchMapping(cfg::BasicBlock *brBB); + + // 获取下一个指令 ID + size_t getNextInstructionId() { return nextInstructionId++; } + + // MLIR Block 到 CFG BasicBlock 的映射（用于处理 cf.cond_br 等跳转指令） + DenseMap blockToBasicBlockMap; + + // 获取或创建 Block 对应的 BasicBlock + cfg::BasicBlock * + getOrCreateBasicBlockForBlock(Block *block, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *parentStructure = nullptr); + + // 注册 Block 到 BasicBlock 的映射 + void registerBlockMapping(Block *mlirBlock, cfg::BasicBlock *cfgBlock); + +private: + size_t nextInstructionId = 0; // 下一个指令 ID +}; + +// 创建 Pass 的工厂函数 +std::unique_ptr> createBuildCFGPass(); + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_CFG_CONTROL_FLOW_GRAPH_BUILDER_H diff --git a/third_party/ascend/include/TritonToGraph/DataflowGraph.h b/third_party/ascend/include/TritonToGraph/DataflowGraph.h new file mode 100644 index 0000000000..56b697c646 --- /dev/null +++ b/third_party/ascend/include/TritonToGraph/DataflowGraph.h @@ -0,0 +1,279 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_TO_CFG_DATAFLOW_GRAPH_H +#define TRITON_TO_CFG_DATAFLOW_GRAPH_H + +#include "TritonToGraph/AliasAnalysis.h" +#include "TritonToGraph/ControlFlowGraph.h" +#include "TritonToGraph/MemorySSA.h" +#include "TritonToGraph/MemorySsaBuilder.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/Value.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Casting.h" +#include + +namespace mlir { +namespace triton { +namespace cfg { + +// 前向声明结果类 +class DataFlowResult; +class MemorySSAResult; +class SSAResult; + +// DataFlowInfo - 统一的数据流信息 +class DataFlowInfo { +public: + // 为函数入口创建参数定义 + void createParameterDefinitions(triton::FuncOp func); + + // Memory SSA接口 + MemorySSADef *getMemoryDefinition(Value value) const; + void addMemoryDefinition(Value value, MemorySSADef *def); + + SmallVector getMemoryUses(Value value) const; + void addMemoryUse(Value value, const MemorySSAUse &use); + + void removeMemoryDefinition(Value value); + void clearMemoryUses(Value value); + + // 传统SSA接口（复用MLIR原生功能） + Operation *getSSADefinition(Value value) const { + return value.getDefiningOp(); + } + + SmallVector getSSAUses(Value value) const { + SmallVector result; + for (OpOperand &use : value.getUses()) { + result.push_back(&use); + } + return result; + } + + // 循环Phi接口 + void addPhi(Value value, const PhiInfo &phiInfo) { Phis[value] = phiInfo; } + + PhiInfo &getPhi(Value value) { return Phis[value]; } + + bool hasPhi(Value value) const { return Phis.count(value) > 0; } + + // 统一查询接口 - 返回unique_ptr，使用LLVM RTTI进行类型判断 + std::unique_ptr queryDataFlow(Value value) const; + + // 查询某个定义的所有使用 + SmallVector getUses(MemorySSADef *def) const; + + // 查询某个操作的memory使用 + SmallVector getUsesByUserOp(Operation *userOp) const; + + // 遍历接口 + void + forEachDefinition(llvm::function_ref func) const; + void forEachUse(llvm::function_ref func) const; + + // 获取所有Memory SSA definitions + const DenseMap &getMemoryDefinitions() const { + return memoryDefinitions; + } + + // 获取所有循环phi信息 + const DenseMap &getPhis() const { return Phis; } + + // 构建def-use缓存 + void buildDefUseCache() const; + + // 打印信息（调试用） + void print(llvm::raw_ostream &os) const; + + // 导出到JSON + void exportToJSON(llvm::raw_ostream &os) const; + +private: + // Memory SSA映射 + DenseMap memoryDefinitions; + DenseMap> memoryUses; + + // Loop Phi映射 + DenseMap Phis; + + // Use-Def映射缓存（def -> uses） + mutable DenseMap> defUseCache; + mutable bool defUseCacheValid = false; + + void invalidateDefUseCache() { + defUseCacheValid = false; + defUseCache.clear(); + } +}; + +// DataFlowResult - 数据流查询结果的基类 +// 使用LLVM RTTI系统，支持isa<>和dyn_cast<>进行类型判断 +class DataFlowResult { +public: + enum class Kind { + MemorySSA, // Memory SSA结果（tensor/pointer） + SSA, // 传统SSA结果（标量） + NONE // 无数据流信息 + }; + + DataFlowResult(Kind kind, Operation *originOp) + : kind(kind), originOp(originOp) {} + virtual ~DataFlowResult() = default; + + Kind getKind() const { return kind; } + Operation *getOriginOp() const { return originOp; } + + SmallVector &getUses() { return uses; } + const SmallVector &getUses() const { return uses; } + + std::optional &getPhi() { return Phi; } + const std::optional &getPhi() const { return Phi; } + + // LLVM RTTI支持 + static bool classof(const DataFlowResult *) { return true; } + +protected: + Kind kind; + Operation *originOp; + SmallVector uses; // 所有uses + std::optional Phi; // Phi信息（如果有） +}; + +// MemorySSAResult - Memory SSA的结果 +class MemorySSAResult : public DataFlowResult { +public: + MemorySSAResult(Operation *originOp, MemorySSADef *definition) + : DataFlowResult(Kind::MemorySSA, originOp), definition(definition) {} + + MemorySSADef *getDefinition() const { return definition; } + + // LLVM RTTI支持 + static bool classof(const DataFlowResult *result) { + return result->getKind() == Kind::MemorySSA; + } + +private: + MemorySSADef *definition; // MEMORY_SSA时使用 +}; + +// SSAResult - 传统SSA的结果 +class SSAResult : public DataFlowResult { +public: + SSAResult(Operation *originOp, Operation *ssaDefinition) + : DataFlowResult(Kind::SSA, originOp), ssaDefinition(ssaDefinition) {} + + Operation *getSSADefinition() const { return ssaDefinition; } + + // LLVM RTTI支持 + static bool classof(const DataFlowResult *result) { + return result->getKind() == Kind::SSA; + } + +private: + Operation *ssaDefinition; // SSA时使用 +}; + +// NoneResult - 无数据流信息的结果 +class NoneResult : public DataFlowResult { +public: + NoneResult() : DataFlowResult(Kind::NONE, nullptr) {} + + // LLVM RTTI支持 + static bool classof(const DataFlowResult *result) { + return result->getKind() == Kind::NONE; + } +}; + +// DataFlowGraph - 数据流图 +class DataFlowGraph { +public: + explicit DataFlowGraph(ControlFlowGraph &cfg) : cfg(cfg) {} + + ~DataFlowGraph() = default; + + // 构建完整的数据流信息 + void build(); + + // 查询Value的数据流信息（使用LLVM RTTI判断具体类型） + std::unique_ptr queryDataFlow(Value value) const { + return dataFlowInfo.queryDataFlow(value); + } + + // 获取所有Memory SSA definitions + SmallVector getAllDefinitions() const { + SmallVector result; + for (const auto &entry : dataFlowInfo.getMemoryDefinitions()) { + result.push_back(entry.second); + } + return result; + } + + // 获取definition的所有uses + SmallVector getUses(MemorySSADef *def) const { + return dataFlowInfo.getUses(def); + } + + // 获取操作的所有uses + SmallVector getUsesByUserOp(Operation *userOp) const { + return dataFlowInfo.getUsesByUserOp(userOp); + } + + // 获取CFG + ControlFlowGraph &getCFG() { return cfg; } + const ControlFlowGraph &getCFG() const { return cfg; } + + // 获取DataFlowInfo + DataFlowInfo &getDataFlowInfo() { return dataFlowInfo; } + const DataFlowInfo &getDataFlowInfo() const { return dataFlowInfo; } + + // 导出数据流信息到JSON + void exportToJSON(llvm::raw_ostream &os) const; + + // 导出def-use链到DOT格式 + void exportDefUseToDOT(llvm::raw_ostream &os) const; + + // 打印所有数据流信息（调试用） + void print(llvm::raw_ostream &os) const; + void dump() const; + +private: + ControlFlowGraph &cfg; + + // 组件 + std::unique_ptr aliasAnalysis; + std::unique_ptr memorySSABuilder; + + // 数据流信息 + DataFlowInfo dataFlowInfo; + + // 构建def-use图 + void buildDefUseGraph(); +}; + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_CFG_DATAFLOW_GRAPH_H diff --git a/third_party/ascend/include/TritonToGraph/GraphAnalysis.h b/third_party/ascend/include/TritonToGraph/GraphAnalysis.h new file mode 100644 index 0000000000..6652ab1ad3 --- /dev/null +++ b/third_party/ascend/include/TritonToGraph/GraphAnalysis.h @@ -0,0 +1,402 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + */ + +#ifndef TRITON_TO_GRAPH_GRAPH_ANALYSIS_H +#define TRITON_TO_GRAPH_GRAPH_ANALYSIS_H + +#include "TritonToGraph/ControlFlowGraph.h" +#include "TritonToGraph/DataflowGraph.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/Value.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/SmallVector.h" +#include + +namespace mlir { +namespace triton { +namespace cfg { + +//===----------------------------------------------------------------------===// +// Traversal Context - 遍历时维护上下文 +//===----------------------------------------------------------------------===// + +struct TraversalContext { + // 当前所处的嵌套结构栈（从外到内） + SmallVector structureStack; + + // 当前深度 + int depth = 0; + + void push(BasicBlock *structure) { structureStack.push_back(structure); } + void pop() { + if (!structureStack.empty()) + structureStack.pop_back(); + } + BasicBlock *currentStructure() const { + return structureStack.empty() ? nullptr : structureStack.back(); + } + bool isEmpty() const { return structureStack.empty(); } +}; + +/// CFGTraversalBase - CFG 遍历基类 +/// 用户继承此类，实现 preVisit/postVisit 回调 +/// 可在子类中添加自定义状态字段 +class CFGTraversalBase { +public: + virtual ~CFGTraversalBase() = default; + + // 访问基本块前调用，返回 false 则跳过此块 + virtual bool preVisitBlock(BasicBlock *block, TraversalContext &ctx); + + // 访问基本块后调用 + virtual void postVisitBlock(BasicBlock *block, TraversalContext &ctx) {} + + // 访问指令后调用 + virtual void VisitInstruction(Instruction *inst, TraversalContext &ctx) {} + + // 遇到控制流结构（for/if/while）时调用 + virtual void onEnterStructure(BasicBlock *structure, TraversalContext &ctx) {} + virtual void onExitStructure(BasicBlock *structure, TraversalContext &ctx) {} + + // 遇到回边时调用 + virtual void onBackEdge(BasicBlock *from, BasicBlock *to, + TraversalContext &ctx) {} +}; + +/// CFGTraverser - CFG 遍历器（双向支持） +class CFGTraverser { +public: + explicit CFGTraverser(ControlFlowGraph &cfg) : cfg(cfg) {} + + //===----------------------------------------------------------------------=== + // Forward Traversal (沿后继节点向下遍历) + //===----------------------------------------------------------------------=== + + /// 从入口开始的 DFS + void dfsForward(CFGTraversalBase &visitor); + + /// 从指定块开始的 DFS + void dfsForward(BasicBlock *start, CFGTraversalBase &visitor); + + /// BFS 遍历 + void bfsForward(CFGTraversalBase &visitor); + void bfsForward(BasicBlock *start, CFGTraversalBase &visitor); + + //===----------------------------------------------------------------------=== + // Backward Traversal (沿前驱节点向上遍历) + //===----------------------------------------------------------------------=== + + /// 从指定块反向 DFS（沿前驱遍历） + void dfsBackward(BasicBlock *start, CFGTraversalBase &visitor); + + /// 从指定块反向 BFS + void bfsBackward(BasicBlock *start, CFGTraversalBase &visitor); + +private: + ControlFlowGraph &cfg; + + void dfsForwardImpl(BasicBlock *block, DenseSet &visited, + TraversalContext &ctx, CFGTraversalBase &visitor); + void dfsBackwardImpl(BasicBlock *block, DenseSet &visited, + TraversalContext &ctx, CFGTraversalBase &visitor); +}; + +//===----------------------------------------------------------------------===// +// Curly Recursive Template Pattern for DFG Traversal +//===----------------------------------------------------------------------===// + +/// DFGTraversalBase - DFG 遍历基类 +class DFGTraversalBase { +public: + virtual ~DFGTTraversalBase() = default; + + // 访问定义前调用（反向遍历 value -> def） + virtual bool VisitDef(Value value, Operation *defOp, int depth); + + // 访问使用前调用（正向遍历 value -> use） + virtual bool VisitUse(Value value, OpOperand *use, int depth); + + // 遇到 phi/iter_arg 时调用 + virtual void onPhi(Value phiValue, const PhiInfo &phiInfo, int depth) {} +}; + +/// DFGTraverser - DFG 遍历器（双向 SSA/MemorySSA） +class DFGTraverser { +public: + struct Options { + bool useMemorySSA = false; // false=传统 SSA, true=Memory SSA + bool followPhi = true; // 是否跨越 phi/iter_arg + int maxDepth = -1; // -1 = 无限制 + DenseSet stopOps; // 遇到停止的操作 + }; + + explicit DFGTraverser(DataFlowGraph &dfg) : dfg(dfg) {} + + //===----------------------------------------------------------------------=== + // Backward Traversal (value -> definitions) + //===----------------------------------------------------------------------=== + + /// 从 value 开始反向追踪所有定义 + void dfsBackward(Value seed, DFGTraversalBase &visitor, + const Options &opts = {}); + + /// 多起点反向 DFS + void dfsBackward(ArrayRef seeds, DFGTraversalBase &visitor, + const Options &opts = {}); + + /// BFS 反向追踪 + void bfsBackward(Value seed, DFGTraversalBase &visitor, + const Options &opts = {}); + + //===----------------------------------------------------------------------=== + // Forward Traversal (value -> uses) + //===----------------------------------------------------------------------=== + + /// 从 value 开始正向追踪所有使用 + void dfsForward(Value seed, DFGTraversalBase &visitor, + const Options &opts = {}); + + /// 多起点正向 DFS + void dfsForward(ArrayRef seeds, DFGTraversalBase &visitor, + const Options &opts = {}); + + /// BFS 正向追踪 + void bfsForward(Value seed, DFGTraversalBase &visitor, + const Options &opts = {}); + + //===----------------------------------------------------------------------=== + // Bidirectional + //===----------------------------------------------------------------------=== + + /// 从 value 双向遍历（先 backward 到根，再 forward 到所有 uses） + void traverseBidirectional(Value seed, DFGTraversalBase &visitor, + const Options &opts = {}); + +private: + DataFlowGraph &dfg; + + void dfsBackwardImpl(Value value, DFGTraversalBase &visitor, + DenseSet &visited, const Options &opts, + int depth); + void dfsForwardImpl(Value value, DFGTraversalBase &visitor, + DenseSet &visited, const Options &opts, + int depth); +}; + +//===----------------------------------------------------------------------===// +// Region Abstraction and Analysis +//===----------------------------------------------------------------------===// + +/// Region - 指令集合（替代原始代码中的 SmallVector） +class Region { +public: + explicit Region(StringRef name = "") : name_(name.str()) {} + + void add(Instruction *inst); + void add(Operation *op, ControlFlowGraph &cfg); + void addAll(ArrayRef insts); + + bool contains(Instruction *inst) const; + bool contains(Operation *op) const; + + void remove(Instruction *inst); + void clear(); + + size_t size() const { return instSet_.size(); } + bool empty() const { return instSet_.empty(); } + + // 获取按块内顺序排序的指令列表 + SmallVector orderedInstructions() const; + + // 获取所有操作 + SmallVector operations() const; + + StringRef name() const { return name_; } + void setName(StringRef name) { name_ = name.str(); } + + // 迭代器支持 + auto begin() const { return instSet_.begin(); } + auto end() const { return instSet_.end(); } + +private: + std::string name_; + DenseSet instSet_; +}; + +/// RegionAnalyzer - Region 分析器 +class RegionAnalyzer { +public: + struct Dependency { + enum Type { DATA, CONTROL }; + Type type; + Value value; + Instruction *from; + Instruction *to; + }; + + struct ExternalDeps { + // 外部定义 -> 内部使用 + struct Input { + Value value; + Instruction *externalDef; // null 表示函数参数 + SmallVector internalUses; + }; + // 内部定义 -> 外部使用 + struct Output { + Value value; + Instruction *internalDef; + SmallVector externalUses; + }; + SmallVector inputs; + SmallVector outputs; + }; + + explicit RegionAnalyzer(DataFlowGraph &dfg, ControlFlowGraph &cfg) + : dfg(dfg), cfg(cfg) {} + + // 检查 region 之间是否有依赖 + bool hasDependency(const Region &from, const Region &to) const; + + // 获取两个 region 间的所有依赖 + SmallVector getDependencies(const Region &from, + const Region &to) const; + + // 分析 region 的外部依赖 + ExternalDeps analyzeExternalDeps(const Region ®ion) const; + + // 检查依赖是否为循环依赖（region A 依赖 B，B 又依赖 A） + bool isCyclicDependency(const Region &a, const Region &b) const; + +private: + DataFlowGraph &dfg; + ControlFlowGraph &cfg; +}; + +//===----------------------------------------------------------------------===// +// Program Slicing +//===----------------------------------------------------------------------===// + +/// SliceCriterion - 切片准则 +struct SliceCriterion { + enum Direction { BACKWARD, FORWARD, BIDIRECTIONAL }; + SmallVector seeds; + Direction dir = BACKWARD; + DFGTraverser::Options dfgOpts; +}; + +/// ProgramSlice - 程序切片 +class ProgramSlice { +public: + void add(Instruction *inst) { instructions_.insert(inst); } + void addAll(const Region ®ion); + + bool contains(Instruction *inst) const { + return instructions_.contains(inst); + } + + size_t size() const { return instructions_.size(); } + bool empty() const { return instructions_.empty(); } + + // 获取入口点（没有前驱在切片中） + SmallVector entryPoints(DataFlowGraph &dfg) const; + + // 获取出口点（没有后继在切片中） + SmallVector exitPoints(DataFlowGraph &dfg) const; + + // 集合操作 + void merge(const ProgramSlice &other); + void intersect(const ProgramSlice &other); + void subtract(const ProgramSlice &other); + + // 转换为 Region + Region toRegion(StringRef name = "") const; + + // 迭代器 + auto begin() const { return instructions_.begin(); } + auto end() const { return instructions_.end(); } + +private: + DenseSet instructions_; +}; + +/// ProgramSlicer - 程序切片器 +class ProgramSlicer { +public: + ProgramSlicer(DataFlowGraph &dfg, ControlFlowGraph &cfg) + : dfg(dfg), cfg(cfg) {} + + // 计算切片 + ProgramSlice compute(const SliceCriterion &criterion); + + // 从 yield values 计算切片（常用场景） + ProgramSlice sliceFromYields(ArrayRef yields, + SliceCriterion::Direction dir); + + // 多切片操作 + static ProgramSlice merge(ArrayRef slices); + static ProgramSlice intersect(ArrayRef slices); + + // 检查切片间依赖 + struct SliceDependency { + const ProgramSlice *from; + const ProgramSlice *to; + SmallVector values; + }; + SmallVector + computeDependencies(ArrayRef slices); + +private: + DataFlowGraph &dfg; + ControlFlowGraph &cfg; +}; + +//===----------------------------------------------------------------------===// +// Region Absorption +//===----------------------------------------------------------------------===// + +/// AbsorptionPolicy - 吸收策略 +struct AbsorptionPolicy { + enum Direction { UPSTREAM, DOWNSTREAM, BOTH }; + Direction dir = BOTH; + int maxDepth = -1; + bool crossRegionBoundary = false; + DenseSet stopOps; + std::function shouldStop = nullptr; +}; + +/// RegionAbsorber - Region 吸收器 +class RegionAbsorber { +public: + RegionAbsorber(DataFlowGraph &dfg, ControlFlowGraph &cfg) + : dfg(dfg), cfg(cfg) {} + + // 从种子指令开始吸收 + void absorb(Region ®ion, ArrayRef seeds, + const AbsorptionPolicy &policy); + + // 从 value 的 def/use 链吸收 + void absorbFromValue(Region ®ion, Value value, + const AbsorptionPolicy &policy); + + // 吸收直到遇到边界 + void absorbUntilBoundary(Region ®ion, ArrayRef seeds, + std::function isBoundary); + +private: + DataFlowGraph &dfg; + ControlFlowGraph &cfg; + + void absorbUpstream(Region ®ion, Instruction *inst, + const AbsorptionPolicy &policy, + DenseSet &visited, int depth); + void absorbDownstream(Region ®ion, Instruction *inst, + const AbsorptionPolicy &policy, + DenseSet &visited, int depth); +}; + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_GRAPH_GRAPH_ANALYSIS_H diff --git a/third_party/ascend/include/TritonToGraph/InterProceduralCFG.h b/third_party/ascend/include/TritonToGraph/InterProceduralCFG.h new file mode 100644 index 0000000000..c4c8bf1a25 --- /dev/null +++ b/third_party/ascend/include/TritonToGraph/InterProceduralCFG.h @@ -0,0 +1,108 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_TO_CFG_INTER_PROCEDURAL_CFG_H +#define TRITON_TO_CFG_INTER_PROCEDURAL_CFG_H + +#include "TritonToGraph/ControlFlowGraph.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Operation.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/raw_ostream.h" + +#include +#include + +namespace mlir { +namespace triton { +namespace cfg { + +// 过程间控制流图 (ICFG) +class InterProceduralCFG { +public: + explicit InterProceduralCFG(ModuleOp module); + ~InterProceduralCFG(); + + // 为每个函数构建CFG + void build(); + + // 获取函数的CFG + ControlFlowGraph *getFunctionCFG(triton::FuncOp func); + const ControlFlowGraph *getFunctionCFG(triton::FuncOp func) const; + ControlFlowGraph *getFunctionCFG(StringRef funcName); + + // 在调用点连接调用者和被调用者的CFG + struct CallSite { + Operation *callOp; // 调用操作 + triton::FuncOp caller; // 调用者函数 + triton::FuncOp callee; // 被调用者函数 + BasicBlock *callBlock; // 调用点的基本块 + Instruction *callInst; // 调用指令 + }; + + // 获取所有调用点 + const SmallVector &getCallSites() const { return callSites; } + + // 连接ICFG中的调用边 + void connectCallGraph(); + + // 查询函数调用关系 + SmallVector getCallees(triton::FuncOp caller) const; + SmallVector getCallers(triton::FuncOp callee) const; + + // 全局可达性分析 + void computeReachability(); + bool isReachable(triton::FuncOp from, triton::FuncOp to) const; + + // 可视化 + void dumpToDot(const std::string &filename) const; + void print(raw_ostream &os) const; + + // 导出到 HTML（包含所有函数） + llvm::Error exportToHTML(const std::string &filename) const; + +private: + ModuleOp module; + + // 函数到CFG的映射 + DenseMap> functionCFGs; + + // 调用点列表 + SmallVector callSites; + + // 调用图 (函数级别) + DenseMap> callGraph; + DenseMap> reverseCallGraph; + + // 可达性矩阵 + DenseMap> reachability; +}; + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_CFG_INTER_PROCEDURAL_CFG_H diff --git a/third_party/ascend/include/TritonToGraph/MemorySSA.h b/third_party/ascend/include/TritonToGraph/MemorySSA.h new file mode 100644 index 0000000000..26f553b2fa --- /dev/null +++ b/third_party/ascend/include/TritonToGraph/MemorySSA.h @@ -0,0 +1,294 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_TO_CFG_MEMORY_SSA_H +#define TRITON_TO_CFG_MEMORY_SSA_H + +#include "TritonToGraph/tensor.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/Value.h" +#include "llvm/ADT/SmallVector.h" + +namespace mlir { +namespace triton { +namespace cfg { + +// Forward declarations +class MemorySSADef; +class MemorySSAUse; + +class BasicBlock; +class ControlFlowGraph; + +// Memory SSA Definition - 表示tensor/pointer的定义 +class MemorySSADef { +public: + // 构造函数 + MemorySSADef(TensorObject *tensor, Operation *defOp, unsigned version = 0) + : tensor(tensor), defOp(defOp), version(version) {} + + // 获取tensor对象 + TensorObject *getTensor() const { return tensor; } + + // 获取创建该definition的操作 + // - 对于入参，返回nullptr + // - 对于其他操作，返回对应的Operation指针 + Operation *getDefOp() const { return defOp; } + + // 获取版本号（函数内全局唯一序数） + unsigned getVersion() const { return version; } + + // 获取唯一标识: tensor_name,version 或 tensor_name,param + std::string getId() const { + return tensor->getName() + "," + + (defOp ? std::to_string(version) : std::to_string(0)); + } + + // 判断是否是入参（函数参数） + bool isParameter() const { return defOp == nullptr; } + + // 判断是否是phi节点（控制流合并） + bool isPhi() const { + return defOp && (isa(defOp) || isa(defOp) || + isa(defOp)); + } + + // 打印信息 + void print(llvm::raw_ostream &os) const { + os << "Definition[" << getId() << ", tensor=" << tensor->getName(); + if (defOp) { + os << ", op=" << defOp->getName(); + } else { + os << ", param"; + } + os << "]"; + } + +private: + TensorObject *tensor; // 对应的tensor对象 + Operation *defOp; // 创建该definition的操作 + unsigned version; // 版本号（函数内全局递增） +}; + +// Memory SSA Use - 表示tensor/pointer的使用 +class MemorySSAUse { +public: + // 构造函数 + MemorySSAUse(MemorySSADef *definition, Operation *userOp, unsigned operandIdx) + : definition(definition), userOp(userOp), operandIdx(operandIdx) { + // 缓存value以提高查询性能 + operandValue = userOp->getOperand(operandIdx); + } + + // 获取使用的definition + MemorySSADef *getDefinition() const { return definition; } + + // 获取使用该definition的操作 + Operation *getUserOp() const { return userOp; } + + // 获取operand序号 + unsigned getOperandIdx() const { return operandIdx; } + + // 获取Value（从userOp的operand） + Value getValue() const { return operandValue; } + + // 获取用户操作的名称 + std::string getUserOpName() const { + return userOp->getName().getStringRef().str(); + } + + // 打印信息 + void print(llvm::raw_ostream &os) const { + os << "Use["; + if (definition) { + os << definition->getId(); + } else { + os << "null"; + } + os << " in " << userOp->getName() << ", operand #" << operandIdx << "]"; + } + +private: + MemorySSADef *definition; // 使用的definition + Operation *userOp; // 使用该definition的操作 + unsigned operandIdx; // operand序号 + Value operandValue; // 缓存的operand value +}; + +// PhiInfo - 循环Phi信息 +struct PhiInfo { + // Phi类型 + enum Type { + ITER_ARG, // scf.for的iter_arg + IF_RESULT, // scf.if的result + WHILE_ARG // scf.while的arg + }; + + Type type; + BasicBlock *loopHeader; // 循环头基本块 + + // Phi值的来源 + struct { + MemorySSADef *initialValue; // 初始值（初始iteration） + MemorySSADef *yieldValue; // yield的值（后续iteration） + } comingFrom; + + // 是否是第一次迭代 + bool isInitial() const { return comingFrom.yieldValue == nullptr; } + + // 获取当前definition（根据上下文决定） + MemorySSADef *getCurrentDefinition(int iteration) const { + return (iteration == 0) ? comingFrom.initialValue : comingFrom.yieldValue; + } +}; + +// MemorySSAInfo - 指令的Memory SSA信息 +struct MemorySSAInfo { + // 指令的operands使用的definitions + SmallVector uses; + + // 指令的results创建的definitions + SmallVector definitions; + + // Alias信息（仅对pointer相关操作） + struct AliasInfo { + Value aliasee; // 别名的源value + TensorObject *baseTensor; // 对应的tensor对象 + }; + std::optional aliasInfo; + + // 快速查询接口 + bool hasDefinition(Value value) const { + // 通过 getDefiningOp() 获取定义该 value 的操作 + Operation *defOp = value.getDefiningOp(); + if (!defOp) + return false; + + // 获取该操作的所有 results + auto results = defOp->getResults(); + for (size_t i = 0; i < definitions.size() && i < results.size(); ++i) { + if (results[i] == value) { + return definitions[i] != nullptr; + } + } + return false; + } + + MemorySSADef *getDefinition(Value value) const { + // 查找该value在results中的索引 + for (auto result : llvm::enumerate(value.getDefiningOp()->getResults())) { + if (result.value() == value) { + size_t idx = result.index(); + if (idx < definitions.size()) { + return definitions[idx]; + } + break; + } + } + return nullptr; + } + + bool hasUse(Value value) const { + for (const MemorySSAUse &use : uses) { + if (use.getValue() == value) { + return true; + } + } + return false; + } + + SmallVector getUses(Value value) const { + SmallVector result; + for (const MemorySSAUse &use : uses) { + if (use.getValue() == value) { + result.push_back(use); + } + } + return result; + } + + // 判断是否创建了新的definition + bool hasNewDefinitions() const { return !definitions.empty(); } + + // 判断是否是写入操作 + bool isMemoryWriter() const { + for (MemorySSADef *def : definitions) { + if (def && !def->isParameter()) { + return true; + } + } + return false; + } + + // 遍历定义 + void forEachDefinition(llvm::function_ref func) const { + for (MemorySSADef *def : definitions) { + if (def) + func(def); + } + } + + void forEachUse(llvm::function_ref func) const { + for (const MemorySSAUse &use : uses) { + if (use.getDefinition()) + func(use); + } + } + + // 清空所有信息 + void clear() { + uses.clear(); + definitions.clear(); + aliasInfo.reset(); + } + + // 打印信息 + void print(llvm::raw_ostream &os) const { + os << "MemorySSAInfo[\n"; + os << " Uses: " << uses.size() << "\n"; + for (const MemorySSAUse &use : uses) { + os << " "; + use.print(os); + os << "\n"; + } + os << " Definitions: " << definitions.size() << "\n"; + for (MemorySSADef *def : definitions) { + if (def) { + os << " "; + def->print(os); + os << "\n"; + } + } + if (aliasInfo) { + os << " Alias: " << aliasInfo->aliasee << " -> " + << aliasInfo->baseTensor->getName() << "\n"; + } + os << "]"; + } +}; + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_CFG_MEMORY_SSA_H diff --git a/third_party/ascend/include/TritonToGraph/MemorySsaBuilder.h b/third_party/ascend/include/TritonToGraph/MemorySsaBuilder.h new file mode 100644 index 0000000000..0d237425b4 --- /dev/null +++ b/third_party/ascend/include/TritonToGraph/MemorySsaBuilder.h @@ -0,0 +1,175 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_TO_CFG_MEMORY_SSA_BUILDER_H +#define TRITON_TO_CFG_MEMORY_SSA_BUILDER_H + +#include "TritonToGraph/AliasAnalysis.h" +#include "TritonToGraph/MemorySSA.h" +#include "mlir/IR/Types.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" + +namespace mlir { +namespace triton { +namespace cfg { + +// 前向声明 +class DataFlowInfo; +class ControlFlowGraph; +class Instruction; + +// MemorySSABuilder - Memory SSA构建器 +// 构建整个CFG的Memory SSA信息，包括创建definitions、uses、处理控制流phi节点 +class MemorySSABuilder { +public: + MemorySSABuilder(ControlFlowGraph &cfg, AliasAnalysis &aliasAnalysis, + DataFlowInfo &dataFlowInfo) + : cfg(cfg), aliasAnalysis(aliasAnalysis), dataFlowInfo(dataFlowInfo) {} + + ~MemorySSABuilder(); + + // 构建整个CFG的Memory SSA + void build(); + +private: + // 处理单个BasicBlock + void processBasicBlock(BasicBlock *bb); + + // 处理单个指令 + void processInstruction(Instruction *inst); + + // 处理scf.if的phi节点 + void processIfOp(scf::IfOp ifOp, Instruction *inst, BasicBlock *thenEntryBB, + BasicBlock *elseEntryBB); + + // 处理scf.for的iter_args + void processForOp(scf::ForOp forOp, Instruction *inst, + BasicBlock *loopBodyEntryBB); + + // 处理scf.while的args + void processWhileOp(scf::WhileOp whileOp, Instruction *inst, + BasicBlock *beforeEntryBB, BasicBlock *afterEntryBB); + + // 判断是否是tensor类型 + bool isTensorType(Type type) const { + return mlir::isa(type) || + mlir::isa(type); + } + + // 根据操作创建tensor对象 + TensorObject *createTensorObject(Operation *op); + + // 创建tensor definition + MemorySSADef *createDefinition(TensorObject *tensor, Operation *op); + + // 创建use + MemorySSAUse createUse(MemorySSADef *def, Operation *userOp, + unsigned operandIdx); + + // 判断是否是入参 + bool isParameter(Operation *op) const { + return op == nullptr; // 入参的defOp为nullptr + } + + // 判断是否是返回新Tensor的操作（根据返回值类型判断，排除load） + bool isTensorWriter(Operation *op) const; + + // 判断是否是修改内存的操作（有副作用） + bool isMemoryWriter(Operation *op) const { + // 只有：tt.store（写入内存） + return isa(op); + } + + // 判断是否是修改读取的操作（有副作用） + bool isMemoryReader(Operation *op) const { + // 只有：tt.store（写入内存） + return isa(op); + } + + // 判断是否是创建指针的操作 + bool isPointerOp(Operation *op) const { + // 返回指针类型：addptr（偏移指针）、make_tensor_ptr（创建张量指针） + if (isa(op)) + return true; + + if (isPointerBroadcastOrSplat(op)) + return true; + + return false; + } + + bool isPointerBroadcastOrSplat(mlir::Operation *op) const; + + // 创建函数的参数定义 + void createParameterDefinitions(); + + // 获取或创建tensor对象 + TensorObject *getOrCreateTensorObject(Value value); + + // 获取操作的字符串名称（用于生成tensor名称） + std::string getOpName(Operation *op); + + // 成员变量 + ControlFlowGraph &cfg; + AliasAnalysis &aliasAnalysis; + DataFlowInfo &dataFlowInfo; + // size_t nextVersionId; // 下一个可用的版本号（用于MemorySSADef） + // size_t nextTensorId; // 下一个可用的tensor ID（用于TensorObject命名） + + std::map nextVersion; + std::map nextTensor; + + // 所有创建的definitions（用于内存管理） + SmallVector allDefinitions; + + // tensor对象缓存 + DenseMap tensorObjectCache; +}; + +// MemorySSABuilderHelper - 辅助函数 +namespace MemorySSABuilderHelper { +// 获取操作的结果类型 +Type getResultType(Operation *op, unsigned resultIdx); + +// 获取Value的形状信息 +SmallVector getShapeFromValue(Value value); + +// 判断两个shape是否相同 +bool shapesEqual(ArrayRef shape1, ArrayRef shape2); + +// 获取scf.if的region yield操作 +Operation *getYieldOp(Region ®ion); + +// 为tensor创建唯一名称 +std::string createUniqueTensorName(StringRef prefix, size_t id); + +// 判断是否需要为操作创建新版本 +bool shouldCreateNewVersion(Operation *op, MemorySSADef *currentDef); +} // namespace MemorySSABuilderHelper + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_CFG_MEMORY_SSA_BUILDER_H diff --git a/third_party/ascend/include/TritonToGraph/Passes.h b/third_party/ascend/include/TritonToGraph/Passes.h new file mode 100644 index 0000000000..1f2e522af3 --- /dev/null +++ b/third_party/ascend/include/TritonToGraph/Passes.h @@ -0,0 +1,44 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_TO_CFG_PASSES_H +#define TRITON_TO_CFG_PASSES_H + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace triton { +namespace cfg { + +// 创建 BuildCFG pass 的工厂函数 +std::unique_ptr> createBuildCFGPass(); + +// 注册所有 CFG 相关的 passes +#define GEN_PASS_REGISTRATION +#include "ascend/include/TritonToGraph/Passes.h.inc" + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_CFG_PASSES_H diff --git a/third_party/ascend/include/TritonToGraph/Passes.td b/third_party/ascend/include/TritonToGraph/Passes.td new file mode 100644 index 0000000000..787fd542c7 --- /dev/null +++ b/third_party/ascend/include/TritonToGraph/Passes.td @@ -0,0 +1,45 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_TO_CFG_PASSES +#define TRITON_TO_CFG_PASSES + +include "mlir/Pass/PassBase.td" + +def BuildCFG : Pass<"build-cfg", "mlir::ModuleOp"> { + let summary = "Build Control Flow Graph from TTIR"; + let description = [{ + This pass analyzes the TTIR and builds a Control Flow Graph (CFG) + representation. It exports the CFG to multiple formats (text, DOT, JSON, HTML) + for visualization and analysis. + + The output files will be named as _cfg. and placed in + the directory specified by the output-dir option. + }]; + let constructor = "mlir::triton::cfg::createBuildCFGPass()"; + let options = [ + Option<"output-dir", "outputDir", "std::string", "\".\"", + "Directory where CFG files will be generated">, + ]; +} + +#endif // TRITON_TO_CFG_PASSES diff --git a/third_party/ascend/include/TritonToGraph/ProgramAxisDependenceAnalysis.h b/third_party/ascend/include/TritonToGraph/ProgramAxisDependenceAnalysis.h new file mode 100644 index 0000000000..52cdece65b --- /dev/null +++ b/third_party/ascend/include/TritonToGraph/ProgramAxisDependenceAnalysis.h @@ -0,0 +1,77 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + */ + +#ifndef TRITON_TO_GRAPH_PROGRAM_AXIS_DEPENDENCE_ANALYSIS_H +#define TRITON_TO_GRAPH_PROGRAM_AXIS_DEPENDENCE_ANALYSIS_H + +#include "mlir/IR/Operation.h" +#include "mlir/IR/Value.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/SmallVector.h" + +#include +#include + +namespace mlir { +namespace triton { +namespace cfg { + +enum class StoreAddressIndependence : uint8_t { + NotProgramDependent, + ProvenDisjoint, + CanonicalDynamicStride, + Unknown, +}; + +struct StoreAddressDependence { + Operation *store = nullptr; + bool pointerDependsOnAxis = false; + StoreAddressIndependence independence = + StoreAddressIndependence::NotProgramDependent; +}; + +struct ProgramAxisDependence { + int32_t axis = 0; + SmallVector programIds; + SmallVector dependenceClosure; + SmallVector stores; + SmallVector reductionAxes; + + bool escapes = false; + bool hasSideEffects = false; + bool hasUnsupportedSideEffects = false; + bool readsNumPrograms = false; + + bool hasProgramId() const { return !programIds.empty(); } + bool hasOnlyDisjointStoreAddresses() const; + bool hasOnlyLiftableStoreAddressesForProgramMapping() const; + bool isIndependentAxisTransformCandidate() const; + bool isProgramMappingTransformCandidate() const; +}; + +class ProgramAxisDependenceAnalysis { +public: + explicit ProgramAxisDependenceAnalysis(triton::FuncOp function); + + const ProgramAxisDependence &get(int32_t axis) const; + +private: + std::array axes; +}; + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif diff --git a/third_party/ascend/include/TritonToGraph/ProgramGridTransform.h b/third_party/ascend/include/TritonToGraph/ProgramGridTransform.h new file mode 100644 index 0000000000..357e87048e --- /dev/null +++ b/third_party/ascend/include/TritonToGraph/ProgramGridTransform.h @@ -0,0 +1,76 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_TO_GRAPH_PROGRAM_GRID_TRANSFORM_H +#define TRITON_TO_GRAPH_PROGRAM_GRID_TRANSFORM_H + +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Support/LLVM.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringRef.h" + +#include + +namespace mlir { +namespace triton { +class FuncOp; +namespace cfg { + +inline constexpr llvm::StringLiteral kProgramGridTransformsAttr = + "hacc.program_grid_transforms"; +inline constexpr int64_t kProgramGridTransformsVersion = 2; + +struct ProgramGridTransform { + int32_t order = 0; + int32_t axis = 0; + int64_t factor = 1; + bool persistentCoverage = false; + bool gridStrideAbiVerified = false; +}; + +struct ProgramGridTransformContract { + int64_t version = kProgramGridTransformsVersion; + SmallVector transforms; +}; + +LogicalResult addProgramGridHiddenExtentArguments(triton::FuncOp function); +bool hasProgramGridHiddenExtentArguments(triton::FuncOp function); + +LogicalResult commitProgramGridFunctionFromSandbox(triton::FuncOp destination, + triton::FuncOp source); + +FailureOr +parseProgramGridTransformContract(Attribute attribute); + +DictionaryAttr serializeProgramGridTransformContract( + MLIRContext *context, const ProgramGridTransformContract &contract); + +LogicalResult +setProgramGridTransformContract(ModuleOp module, + const ProgramGridTransformContract &contract); + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif diff --git a/third_party/ascend/include/TritonToGraph/ProgramMappingScalarSpecialization.h b/third_party/ascend/include/TritonToGraph/ProgramMappingScalarSpecialization.h new file mode 100644 index 0000000000..7bca0044bd --- /dev/null +++ b/third_party/ascend/include/TritonToGraph/ProgramMappingScalarSpecialization.h @@ -0,0 +1,64 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + */ + +#ifndef TRITON_TO_GRAPH_PROGRAM_MAPPING_SCALAR_SPECIALIZATION_H +#define TRITON_TO_GRAPH_PROGRAM_MAPPING_SCALAR_SPECIALIZATION_H + +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Support/LLVM.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringRef.h" + +#include +#include + +namespace mlir { +namespace triton { +namespace cfg { + +inline constexpr llvm::StringLiteral kProgramMappingScalarSpecializationAttr = + "hacc.program_mapping_scalar_specialization"; +inline constexpr int64_t kProgramMappingScalarSpecializationLegacyVersion = 1; +inline constexpr int64_t kProgramMappingScalarSpecializationVersion = 2; + +struct ProgramMappingScalarArgument { + uint32_t index = 0; + std::string name; + int64_t value = 0; +}; + +struct ProgramMappingScalarSpecialization { + int64_t version = kProgramMappingScalarSpecializationVersion; + SmallVector arguments; +}; + +FailureOr +parseProgramMappingScalarSpecialization(Attribute attribute); + +DictionaryAttr serializeProgramMappingScalarSpecialization( + MLIRContext *context, + const ProgramMappingScalarSpecialization &specialization); + +LogicalResult setProgramMappingScalarSpecialization( + ModuleOp module, const ProgramMappingScalarSpecialization &specialization); +void clearProgramMappingScalarSpecialization(ModuleOp module); + +LogicalResult applyProgramMappingScalarSpecialization(ModuleOp module); + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif diff --git a/third_party/ascend/include/TritonToGraph/tensor.h b/third_party/ascend/include/TritonToGraph/tensor.h new file mode 100644 index 0000000000..49a9e58df8 --- /dev/null +++ b/third_party/ascend/include/TritonToGraph/tensor.h @@ -0,0 +1,166 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_TO_CFG_TENSOR_H +#define TRITON_TO_CFG_TENSOR_H + +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Types.h" +#include "triton/Dialect/Triton/IR/Types.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringRef.h" + +namespace mlir { +namespace triton { +namespace cfg { + +// Compute类型（用于指令分类） +enum class ComputeType { + CUBE, // shape为2维 + VECTOR, // shape为1维 + SCALAR // 其他标量 +}; + +// TensorObject - Tensor对象定义 +class TensorObject { +public: + // Tensor类型分类 + enum class TensorKind { + GLOBAL_MEMORY, // 全局内存tensor（如gm_obj） + L2, // L2 Cache + L1, // L1 Cache + UB // Unified Buffer + }; + + // 构造函数 + TensorObject(StringRef name, ArrayRef shape, Type type, + Type elementType, TensorKind kind = TensorKind::GLOBAL_MEMORY) + : name(name.str()), shape(shape.begin(), shape.end()), type(type), + elementType(elementType), kind(kind) {} + + // 获取tensor名称 + const std::string &getName() const { return name; } + + // 设置tensor名称 + void setName(StringRef newName) { name = newName.str(); } + + // 获取shape + ArrayRef getShape() const { return shape; } + + // 获取MLIR类型 + Type getType() const { return type; } + + // 获取tensor种类 + TensorKind getKind() const { return kind; } + + // 设置tensor种类 + void setKind(TensorKind newKind) { kind = newKind; } + + // 获取元素数据类型 + Type getElementType() const { return elementType; } + + // 设置元素数据类型 + void setElementType(Type newElementType) { elementType = newElementType; } + + // 获取维度数 + size_t getRank() const { return shape.size(); } + + // 获取tensor大小（元素总数） + int64_t getSize() const { + int64_t size = 1; + for (int64_t dim : shape) { + size *= dim; + } + return size; + } + + // 判断是否是一维tensor + bool isVector() const { return shape.size() == 1; } + + // 判断是否是二维tensor + bool isMatrix() const { return shape.size() == 2; } + + // 判断是否是标量 + bool isScalar() const { return shape.empty(); } + + // 打印信息 + void print(llvm::raw_ostream &os) const { + os << "Tensor[" << name << ", shape=["; + for (size_t i = 0; i < shape.size(); ++i) { + if (i > 0) + os << ", "; + os << shape[i]; + } + os << "], element=" << elementType << ", kind=" << getKindString() << "]"; + } + + // 获取kind的字符串表示 + std::string getKindString() const { + switch (kind) { + case TensorKind::GLOBAL_MEMORY: + return "GLOBAL_MEMORY"; + case TensorKind::L2: + return "L2"; + case TensorKind::L1: + return "L1"; + case TensorKind::UB: + return "UB"; + } + return "UNKNOWN"; + } + +private: + std::string name; // Tensor名称，如"gm_obj_0" + SmallVector shape; + Type type; // 完整类型（如tensor<64x64xf32>） + Type elementType; // 元素数据类型（如f32, i8, f16等） + TensorKind kind; +}; + +// 从类型中提取shape和element type +inline void extractShapeAndElementType(Type type, + SmallVectorImpl &shape, + Type &elementType) { + if (auto rankedType = mlir::dyn_cast(type)) { + shape.append(rankedType.getShape().begin(), rankedType.getShape().end()); + elementType = rankedType.getElementType(); + } else if (auto ptrType = mlir::dyn_cast(type)) { + Type pointeeType = ptrType.getPointeeType(); + if (auto rankedType = mlir::dyn_cast(pointeeType)) { + shape.append(rankedType.getShape().begin(), rankedType.getShape().end()); + elementType = rankedType.getElementType(); + } else { + // 标量指针 + elementType = pointeeType; + } + } else { + // 默认值 + elementType = type; + } +} + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_CFG_TENSOR_H diff --git a/third_party/ascend/include/TritonToHFusion/CMakeLists.txt b/third_party/ascend/include/TritonToHFusion/CMakeLists.txt new file mode 100644 index 0000000000..dc6dfc3906 --- /dev/null +++ b/third_party/ascend/include/TritonToHFusion/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name TritonToHFusion) +add_public_tablegen_target(TritonToHFusionConversionPassIncGen) diff --git a/third_party/ascend/include/TritonToHFusion/Passes.h b/third_party/ascend/include/TritonToHFusion/Passes.h new file mode 100644 index 0000000000..1ae954d323 --- /dev/null +++ b/third_party/ascend/include/TritonToHFusion/Passes.h @@ -0,0 +1,43 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_TRITON_TO_HFUSION_CONVERSION_PASSES_H +#define TRITON_ADAPTER_TRITON_TO_HFUSION_CONVERSION_PASSES_H + +#include "mlir/Pass/Pass.h" + +namespace mlir { +// Forward declarations. +class ModuleOp; + +namespace triton { + +/// Creates a pass to convert Triton dialect to HFusion dialect. +std::unique_ptr> createTritonToHFusionPass(); + +#define GEN_PASS_REGISTRATION +#include "ascend/include/TritonToHFusion/Passes.h.inc" + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_TRITON_TO_HFUSION_CONVERSION_PASSES_H diff --git a/third_party/ascend/include/TritonToHFusion/Passes.td b/third_party/ascend/include/TritonToHFusion/Passes.td new file mode 100644 index 0000000000..9903d97a12 --- /dev/null +++ b/third_party/ascend/include/TritonToHFusion/Passes.td @@ -0,0 +1,17 @@ +/* + * Copyright (c) Huawei Technologies Co. + * Licensed under the MIT license. + */ + +#ifndef TRITON_TO_HFUSION_CONVERSION_PASSES +#define TRITON_TO_HFUSION_CONVERSION_PASSES + +include "mlir/Pass/PassBase.td" + +def TritonToHFusion : Pass<"triton-to-hfusion", "mlir::ModuleOp"> { + let summary = "Convert Triton to HFusion dialect"; + let constructor = "triton::createTritonToHFusionPass()"; + let dependentDialects = ["hfusion::HFusionDialect"]; +} + +#endif // TRITON_TO_HFUSION_CONVERSION_PASSES diff --git a/third_party/ascend/include/TritonToHIVM/CMakeLists.txt b/third_party/ascend/include/TritonToHIVM/CMakeLists.txt new file mode 100644 index 0000000000..112189a75f --- /dev/null +++ b/third_party/ascend/include/TritonToHIVM/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name TritonToHIVM) +add_public_tablegen_target(TritonToHIVMConversionPassIncGen) diff --git a/third_party/ascend/include/TritonToHIVM/Passes.h b/third_party/ascend/include/TritonToHIVM/Passes.h new file mode 100644 index 0000000000..6d55efdb73 --- /dev/null +++ b/third_party/ascend/include/TritonToHIVM/Passes.h @@ -0,0 +1,43 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#ifndef TRITON_ADAPTER_TRITON_TO_HIVM_CONVERSION_PASSES_H +#define TRITON_ADAPTER_TRITON_TO_HIVM_CONVERSION_PASSES_H + +#include "mlir/Pass/Pass.h" + +namespace mlir { +// Forward declarations. +class ModuleOp; + +namespace triton { + +/// Creates a pass to convert Triton dialect to HIVM dialect. +std::unique_ptr> createTritonToHIVMPass(); + +#define GEN_PASS_REGISTRATION +#include "ascend/include/TritonToHIVM/Passes.h.inc" + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_TRITON_TO_HIVM_CONVERSION_PASSES_H diff --git a/third_party/ascend/include/TritonToHIVM/Passes.td b/third_party/ascend/include/TritonToHIVM/Passes.td new file mode 100644 index 0000000000..7aa37c5136 --- /dev/null +++ b/third_party/ascend/include/TritonToHIVM/Passes.td @@ -0,0 +1,17 @@ +/* + * Copyright (c) Huawei Technologies Co. + * Licensed under the MIT license. + */ + +#ifndef TRITON_TO_HIVM_CONVERSION_PASSES +#define TRITON_TO_HIVM_CONVERSION_PASSES + +include "mlir/Pass/PassBase.td" + +def TritonToHIVM : Pass<"triton-to-hivm", "mlir::ModuleOp"> { + let summary = "Convert Triton to HIVM dialect"; + let constructor = "triton::createTritonToHIVMPass()"; + let dependentDialects = ["hivm::HIVMDialect"]; +} + +#endif // TRITON_TO_HIVM_CONVERSION_PASSES diff --git a/third_party/ascend/include/TritonToLLVM/CMakeLists.txt b/third_party/ascend/include/TritonToLLVM/CMakeLists.txt new file mode 100644 index 0000000000..fb5a9b2952 --- /dev/null +++ b/third_party/ascend/include/TritonToLLVM/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name TritonToLLVM) +add_public_tablegen_target(TritonToLLVMConversionPassIncGen) diff --git a/third_party/ascend/include/TritonToLLVM/Passes.h b/third_party/ascend/include/TritonToLLVM/Passes.h new file mode 100644 index 0000000000..db8ecec567 --- /dev/null +++ b/third_party/ascend/include/TritonToLLVM/Passes.h @@ -0,0 +1,26 @@ +/* + * Copyright (c) Huawei Technologies Co. + * Licensed under the MIT license. + */ + +#ifndef TRITON_ADAPTER_TRITON_TO_LLVM_CONVERSION_PASSES_H +#define TRITON_ADAPTER_TRITON_TO_LLVM_CONVERSION_PASSES_H + +#include "mlir/Pass/Pass.h" + +namespace mlir { +// Forward declarations. +class ModuleOp; + +namespace triton { + +/// Creates a pass to convert Triton dialect to LLVM dialect. +std::unique_ptr> createTritonToLLVMPass(); + +#define GEN_PASS_REGISTRATION +#include "ascend/include/TritonToLLVM/Passes.h.inc" + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_TRITON_TO_LLVM_CONVERSION_PASSES_H diff --git a/third_party/ascend/include/TritonToLLVM/Passes.td b/third_party/ascend/include/TritonToLLVM/Passes.td new file mode 100644 index 0000000000..fa39f39f65 --- /dev/null +++ b/third_party/ascend/include/TritonToLLVM/Passes.td @@ -0,0 +1,17 @@ +/* + * Copyright (c) Huawei Technologies Co. + * Licensed under the MIT license. + */ + +#ifndef TRITON_TO_LLVM_CONVERSION_PASSES +#define TRITON_TO_LLVM_CONVERSION_PASSES + +include "mlir/Pass/PassBase.td" + +def TritonToLLVM : Pass<"triton-to-llvm", "mlir::ModuleOp"> { + let summary = "Convert Triton to LLVM dialect"; + let constructor = "triton::createTritonToLLVMPass()"; + let dependentDialects = ["LLVM::LLVMDialect", "tensor::TensorDialect", "arith::ArithDialect"]; +} + +#endif // TRITON_TO_LLVM_CONVERSION_PASSES diff --git a/third_party/ascend/language/cann/__init__.py b/third_party/ascend/language/cann/__init__.py new file mode 100644 index 0000000000..883dee2a91 --- /dev/null +++ b/third_party/ascend/language/cann/__init__.py @@ -0,0 +1,58 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +from triton.language import math +from triton.backends.ascend.utils import triton_enable_libdevice_simt + +from . import libdevice +from . import extension + +extension.parallel = extension.aux_ops.parallel +if not triton_enable_libdevice_simt(): + libdevice.atan2 = extension.math_ops.atan2 +libdevice.isfinited = extension.math_ops.isfinited +libdevice.finitef = extension.math_ops.finitef +libdevice.flip = extension.flip + +libdevice.umulhi = math.umulhi +libdevice.exp = math.exp +libdevice.exp2 = math.exp2 +libdevice.log = math.log +libdevice.log2 = math.log2 +libdevice.cos = math.cos +libdevice.sin = math.sin +libdevice.sqrt = math.sqrt +libdevice.sqrt_rn = math.sqrt_rn +libdevice.rsqrt = math.rsqrt +libdevice.div_rn = math.div_rn +libdevice.erf = math.erf +libdevice.floor = math.floor +libdevice.ceil = math.ceil +libdevice.fdiv = math.fdiv +libdevice.fma = math.fma +libdevice.abs = math.abs +math.tanh = libdevice.tanh + +__all__ = ["libdevice", "extension"] + +# flagtree language extension +from triton._flagtree_spec import bind_language_extension_symbols_to_tl + +bind_language_extension_symbols_to_tl(extension) diff --git a/third_party/ascend/language/cann/extension/__init__.py b/third_party/ascend/language/cann/extension/__init__.py new file mode 100644 index 0000000000..541e5c849b --- /dev/null +++ b/third_party/ascend/language/cann/extension/__init__.py @@ -0,0 +1,154 @@ +from triton.tools.get_ascend_devices import is_compile_on_910_95 +from triton._C.libtriton.ascend import ir as _ascend_ir + +# MLIR affine bindings (same objects as triton._C.libtriton.ascend.ir). +affine_expr = _ascend_ir.affine_expr +affine_constant_expr = _ascend_ir.affine_constant_expr +affine_dim_expr = _ascend_ir.affine_dim_expr +affine_symbol_expr = _ascend_ir.affine_symbol_expr +affine_binary_op_expr = _ascend_ir.affine_binary_op_expr +affine_map = _ascend_ir.affine_map + +AffineExpr = affine_expr +AffineConstantExpr = affine_constant_expr +AffineDimExpr = affine_dim_expr +AffineSymbolExpr = affine_symbol_expr +AffineBinaryOpExpr = affine_binary_op_expr +AffineMap = affine_map + +from .core import ( + ascend_address_space, + builtin, + CORE, + copy_from_ub_to_l1, + copy, + debug_barrier, + fixpipe, + FixpipeDMAMode, + FixpipeDualDstMode, + FixpipePreQuantMode, + FixpipePreReluMode, + int64, + is_builtin, + MODE, + PIPE, + IteratorType, + sub_vec_id, + sub_vec_num, + sync_block_all, + sync_block_set, + sync_block_wait, + SYNC_IN_VF, + conv1d, +) + +from .scope import scope + +from .custom_op import ( + custom, + custom_object_semantic, + custom_semantic, + register_custom_op, +) + +from . import builtin_custom_ops + +from .math_ops import (atan2, isfinited, finitef) + +from .aux_ops import ( + parallel, + compile_hint, + multibuffer, +) + +from .vec_ops import ( + insert_slice, + extract_slice, + get_element, + sort, + flip, + cast, +) + +from .mem_ops import ( + index_put, + gather_out_to_ub, + scatter_ub_to_out, + index_select_simd, +) + +__all__ = [ + # core + "builtin", + "copy_from_ub_to_l1", + "copy", + "CORE", + "debug_barrier", + "fixpipe", + "FixpipeDMAMode", + "FixpipeDualDstMode", + "FixpipePreQuantMode", + "FixpipePreReluMode", + "int64", + "is_builtin", + "MODE", + "PIPE", + "IteratorType", + "sub_vec_id", + "sub_vec_num", + "sync_block_all", + "SYNC_IN_VF", + "conv1d", + + # address space + "ascend_address_space", + + # ascend IR affine (MLIR) + "affine_expr", + "affine_constant_expr", + "affine_dim_expr", + "affine_symbol_expr", + "affine_binary_op_expr", + "affine_map", + "AffineExpr", + "AffineConstantExpr", + "AffineDimExpr", + "AffineSymbolExpr", + "AffineBinaryOpExpr", + "AffineMap", + + # scope + "scope", + + # custom op + "custom", + "custom_object_semantic", + "custom_semantic", + "register_custom_op", + + # math ops + "atan2", + "isfinited", + "finitef", + + # aux ops + "sync_block_set", + "sync_block_wait", + "parallel", + "compile_hint", + "multibuffer", + + # vec ops + "insert_slice", + "extract_slice", + "get_element", + "sort", + "flip", + "cast", + + # mem ops + "index_put", + "gather_out_to_ub", + "scatter_ub_to_out", + "index_select_simd", +] diff --git a/third_party/ascend/language/cann/extension/_utils.py b/third_party/ascend/language/cann/extension/_utils.py new file mode 100644 index 0000000000..97824ae9c9 --- /dev/null +++ b/third_party/ascend/language/cann/extension/_utils.py @@ -0,0 +1,54 @@ +import triton.language.core as tl +from triton._C.libtriton import ir + + +def custom_op(builder: ir.builder, op_name: str, **kwargs): + if op_name == "sync_block_all": + return builder.create_custom_op_for_inter_core_sync(op_name, kwargs["mode"], kwargs["event_id"]) + + elif op_name == "sync_block_set": + return builder.create_custom_op_for_inter_core_sync(op_name, kwargs["sender"], kwargs["event_id"]) + + elif op_name == "sync_block_wait": + return builder.create_custom_op_for_inter_core_sync(op_name, kwargs["sender"], kwargs["event_id"]) + + raise ValueError(f"Unsupported custom op: {op_name}") + + +def _is_int_like_elem(x) -> bool: + """Accept int / tl.constexpr(int) / tl.tensor(int*).""" + if isinstance(x, int): + return True + if isinstance(x, tl.constexpr): + # constexpr value should be python int + return isinstance(x.value, int) + if isinstance(x, tl.tensor): + # Offsets/strides must be integer typed (i32/i64 etc.) + return x.dtype.is_int() + return False + + +def _assert_int_like_tuple(name: str, xs): + assert isinstance(xs, (list, tuple, tl.tuple)), f"{name} should be a tuple/list, but got {type(xs)}" + assert all(_is_int_like_elem(x) for x in xs), f"{name} should be integer" + + +def _convert_elem_to_ir_value(builder, elem, require_i64): + if isinstance(elem, int): + elem = tl.constexpr(elem) + if isinstance(elem, tl.constexpr): + if require_i64: + assert -2**63 <= elem.value < 2**63, f"Block pointers only support 64 bit `shape/strides`, " \ + f"got a value {elem.value} which is out of the range" + return builder.get_int64(elem.value) + else: + assert -2**31 <= elem.value < 2**31, f"Block pointers only support 32 bit `offsets/block_shape`, " \ + f"got a value {elem.value} which is out of the range" + return builder.get_int32(elem.value) + elif isinstance(elem, tl.tensor): + if require_i64: + return builder.create_int_cast(elem.handle, builder.get_int64_ty(), elem.dtype.is_int_signed()) + else: + return builder.create_int_cast(elem.handle, builder.get_int32_ty(), elem.dtype.is_int_signed()) + else: + assert False, f"Unsupported element type in shape/strides/offsets: {type(elem)}" diff --git a/third_party/ascend/language/cann/extension/aux_ops.py b/third_party/ascend/language/cann/extension/aux_ops.py new file mode 100644 index 0000000000..79b035bdbb --- /dev/null +++ b/third_party/ascend/language/cann/extension/aux_ops.py @@ -0,0 +1,142 @@ +import triton.language as tl +from triton.language import core +from triton.language.core import (_unwrap_if_constexpr, _tensor_member_fn, _unwrap_iterable, builtin, constexpr, dtype, + tensor, check_bit_width, _unwrap_if_constexpr, range) + +from typing import Optional, Tuple, List, overload, Union +from triton._C.libtriton import ir +from ._utils import custom_op + + +@_tensor_member_fn +@builtin +def sync_block_all(mode, event_id, _semantic=None): + import warnings + + warnings.warn( + ("This method would be deprecated. Use al.sync_block_all instead."), + DeprecationWarning, + stacklevel=1, + ) + mode = _unwrap_if_constexpr(mode) + event_id = _unwrap_if_constexpr(event_id) + assert isinstance(mode, str), f"mode: {mode} is not string" + assert isinstance(event_id, int) and (event_id >= 0) and (event_id < 16), f"event_id: {event_id} should be 0 ~ 15" + assert mode == "all_cube" or mode == "all_vector" or mode == "all", f"ERROR: mode = {mode}, only supports all_cube/all_vector/all" + custom_op(_semantic.builder, "sync_block_all", mode=mode, event_id=event_id) + + +@_tensor_member_fn +@builtin +def sync_block_set(sender, receiver, event_id, _semantic=None): + import warnings + + warnings.warn( + ("This method would be deprecated. Use al.sync_block_set instead."), + DeprecationWarning, + stacklevel=1, + ) + sender = _unwrap_if_constexpr(sender) + receiver = _unwrap_if_constexpr(receiver) + event_id = _unwrap_if_constexpr(event_id) + assert isinstance(sender, str) and (sender == "cube" + or sender == "vector"), f"ERROR: sender = {sender}, only supports cube/vector" + assert isinstance(receiver, str) and (receiver == "cube" or receiver + == "vector"), f"ERROR: receiver = {receiver}, only supports cube/vector" + assert isinstance(event_id, int) and (event_id >= 0) and (event_id < 16), f"event_id: {event_id} should be 0 ~ 15" + if sender == receiver: + raise ValueError(f'Unexpected pair: {sender} -> {receiver}, only supports cube -> vector or vector -> cube') + custom_op(_semantic.builder, "sync_block_set", sender=sender, event_id=event_id) + + +@_tensor_member_fn +@builtin +def sync_block_wait(sender, receiver, event_id, _semantic=None): + import warnings + + warnings.warn( + ("This method would be deprecated. Use al.sync_block_wait instead."), + DeprecationWarning, + stacklevel=1, + ) + sender = _unwrap_if_constexpr(sender) + receiver = _unwrap_if_constexpr(receiver) + event_id = _unwrap_if_constexpr(event_id) + assert isinstance(sender, str) and (sender == "cube" + or sender == "vector"), f"ERROR: sender = {sender}, only supports cube/vector" + assert isinstance(receiver, str) and (receiver == "cube" or receiver + == "vector"), f"ERROR: receiver = {receiver}, only supports cube/vector" + assert isinstance(event_id, int) and (event_id >= 0) and (event_id < 16), f"event_id: {event_id} should be 0 ~ 15" + if sender == receiver: + raise ValueError(f'Unexpected pair: {sender} -> {receiver}, only supports cube -> vector or vector -> cube') + custom_op(_semantic.builder, "sync_block_wait", sender=sender, event_id=event_id) + + +class parallel(range): + """ + Iterator that counts upward forever, with parallel execution semantics. + + This is a special iterator used to implement similar semantics to Python's :code:`range` in the context of + :code:`triton.jit` functions. In addition, it allows user to pass extra attributes to the compiler. + :param bind_sub_block: Tells the compiler if multiple vector cores participate in the loop. + This is used in the mixed cube-vector kernel on 910B. The number of vector cores is determined by the number of + iteration in this loop. Currently on 910B, max 2 vector cores could be used. + """ + + def __init__(self, arg1, arg2=None, step=None, num_stages=None, loop_unroll_factor=None, + bind_sub_block: bool = False): + super().__init__(arg1, arg2, step, num_stages, loop_unroll_factor) + self.bind_sub_block = bind_sub_block + + +def compile_hint_impl(ptr: tensor, hint_name: str, hint_val, builder: ir.builder): + # simt mode does not support hint annotations + # FIXME: is_simt_mode + # if builder.is_simt_mode(): + # return + # Check isinstance(hint_val, bool) first to handle False explicitly + if isinstance(hint_val, bool): + hint_val = builder.get_bool_attr(hint_val) + elif not hint_val: + hint_val = builder.get_unit_attr() + elif isinstance(hint_val, int): + hint_val = builder.get_int32_attr(hint_val) + elif isinstance(hint_val, core.constexpr): + hint_val = builder.get_string_attr(hint_val.value) + elif isinstance(hint_val, (list, core.tuple)): + # only support i64 array attr for now + hint_val = builder.get_i64_array_attr(hint_val) + else: + raise ValueError(f"Unsupported hint value type: {type(hint_val)}") + builder.create_annotation_mark(ptr.handle, hint_name, hint_val) + + +@builtin +def compile_hint(ptr, hint_name, hint_val=None, _semantic=None): + # simt mode does not support hint annotations + if _semantic.builder.is_simt_mode(): + return + + def _unwrap(val): + return _unwrap_if_constexpr(val) if val else val + + hint_name = _unwrap_if_constexpr(hint_name) + assert isinstance(hint_name, str), f"hint name: {hint_name} is not string" + if isinstance(hint_val, (list, core.tuple)): + hint_val = [_unwrap(val) for val in hint_val] + else: + hint_val = _unwrap(hint_val) + hint_val = _unwrap_if_constexpr(hint_val) if hint_val else hint_val + compile_hint_impl(ptr, hint_name, hint_val, _semantic.builder) + + +@builtin +def multibuffer(src: tensor, size, _semantic=None): + """ + Set multi_buffer for an existing tensor + :src: tensor set to bufferize multiple time + :size: number of copies + """ + buffer_size = _unwrap_if_constexpr(size) + assert isinstance(buffer_size, int) and buffer_size == 2, f"only support bufferize equals 2" + compile_hint_impl(src, "hivm.multi_buffer", buffer_size, _semantic.builder) diff --git a/third_party/ascend/language/cann/extension/builder.py b/third_party/ascend/language/cann/extension/builder.py new file mode 100644 index 0000000000..8cf699f63a --- /dev/null +++ b/third_party/ascend/language/cann/extension/builder.py @@ -0,0 +1,85 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Ascend-specific builder utilities for code generation. +""" + +__all__ = [ + "create_builder_method_wrapper", + "attach_builder_methods", + "setup_unified_builder", +] + + +def create_builder_method_wrapper(main_builder, delegate_builder, method_name): + """ + Create a wrapper that delegates a method call to another builder while + synchronizing insertion points and locations. + """ + delegate_method = getattr(delegate_builder, method_name) + + def wrapper(*args, **kwargs): + saved_ip = main_builder.get_insertion_point() + saved_loc = main_builder.get_loc() + delegate_builder.restore_insertion_point(saved_ip) + if saved_loc: + delegate_builder.set_loc(saved_loc) + result = delegate_method(*args, **kwargs) + main_builder.restore_insertion_point(saved_ip) + if saved_loc: + main_builder.set_loc(saved_loc) + return result + + wrapper.__name__ = method_name + wrapper.__doc__ = getattr(delegate_method, '__doc__', None) + return wrapper + + +def attach_builder_methods(main_builder, delegate_builder, method_names): + """Attach multiple methods from a delegate builder to the main builder.""" + for method_name in method_names: + wrapper = create_builder_method_wrapper(main_builder, delegate_builder, method_name) + setattr(main_builder, method_name, wrapper) + + +def setup_unified_builder(main_builder, ascend_builder): + """Set up a unified builder interface by attaching methods from specialized builders.""" + main_builder._ascend_builder = ascend_builder + ascend_methods = [ + 'create_scope_op', + 'scope_return', + 'get_t_core_type_attr_name', + 'get_t_core_type_cube_attr', + 'get_t_core_type_vector_attr', + 'get_target_attribute', + 'create_get_sub_vec_id', + 'create_copy_buffer', + 'create_copy_tensor', + 'create_fixpipe', + 'create_annotation_mark', + 'create_bind_buffer', + 'create_debug_barrier', + 'is_910_95', + "sync_block_set", + "sync_block_wait", + "create_convert_layout", + 'sync_block_all', + ] + attach_builder_methods(main_builder, ascend_builder, ascend_methods) diff --git a/third_party/ascend/language/cann/extension/builtin_custom_ops.py b/third_party/ascend/language/cann/extension/builtin_custom_ops.py new file mode 100644 index 0000000000..49ac7bed52 --- /dev/null +++ b/third_party/ascend/language/cann/extension/builtin_custom_ops.py @@ -0,0 +1,220 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton.language.core as tl +from .custom_op import register_custom_op +from .core import CORE, PIPE, MODE +from ._utils import _is_int_like_elem, _assert_int_like_tuple + + +@register_custom_op +class _index_select: + """ + This operation gathers values from the src GM tensor into the out UB tensor + at positions with offsets specified by the index UB tensor along the specified + dimension using a SIMT template. This operation supports 2D–5D. + + Arguments: + - src: pointer type, the source tensor pointer (in GM) + - index: tensor, a tensor to gather (in UB) + - dim: int, the dimension to gather along + - bound: int, the upper boundary for index + - end_offset: tuple of int, the end offsets of each dimension for index tensor + - start_offset: tuple of int, the start offsets of each dimension for src tensor + - src_stride: tuple of int, the stride of each dimension of src tensor + - other(Optional): scalar value, the default value when index is out of boundary (in UB) + - out: the output tensor (in UB) + + Note: + - Supported source ranks: 2D ~ 5D. + - Supported index ranks: 1D or 2D. + - `dim` must be valid (0 <= dim < source ranks). + + Reference formula: + Index select operation for different tensor ranks: + 1. 2D index gather (0 <= dim <= 1) + 1.1 dim = 0, index_rank = 1, src_rank = 2, out_rank = 2 + index_shape = (Ai,) + end_offset = (Ai_end, B_end) + start_offset = (0, B_begin) + out[i][0:B_end-B_begin] = src[index[i]][B_begin:B_end] + 1.2 dim = 0, index_rank = 2, src_rank = 2, out_rank = 3 + index_shape = (Ai, Aj) + end_offset = (Ai_end, Aj_end, B_end) + start_offset = (0, B_begin) + out[i][j][0:B_end-B_begin] = src[index[i][j]][B_begin:B_end] + 2. 3D index gather (0 <= dim <= 2) + 2.1 dim = 0, index_rank = 2, src_rank = 3, out_rank = 4 + index_shape = (Ai, Aj) + end_offset = (Ai_end, Aj_end, B_end, C_end) + start_offset = (0, B_begin, C_begin) + out[i][j][0:B_end-B_begin][0:C_end-C_begin] = src[index[i][j]][B_begin:B_end][C_begin:C_end] + and so on. + """ + name = '__builtin_index_select' + core = CORE.VECTOR + pipe = PIPE.PIPE_V + mode = MODE.SIMT + + def __init__(self, src, index, dim, bound: tl.int64, end_offset, start_offset, src_stride, other=None, out=None): + assert src.type.is_ptr() or src.dtype.is_ptr(), f"src should be a pointer, but got {src.type}" + assert index.dtype.is_int(), "index should be integer tensor" + src_rank = len(src_stride) + idx_rank = len(index.shape) + assert 2 <= src_rank <= 5, f"src rank should in [2, 5], but got {src_rank}" + assert 1 <= idx_rank <= 2, f"index rank should in [1, 2], but got {idx_rank}" + assert _is_int_like_elem(dim), "dim should be an integer" + assert _is_int_like_elem(bound), "bound should be an integer" + assert 0 <= dim < src_rank, f"dim should in [0, {src_rank - 1}], but got {dim}" + assert len(start_offset) == len(src_stride), "start_offset and src_stride should have same size" + assert len(end_offset) == idx_rank + len( + start_offset) - 1, "len(end_offset) should be equal to index rank + len(start_offset) - 1" + + _assert_int_like_tuple("end_offset", end_offset) + _assert_int_like_tuple("start_offset", start_offset) + _assert_int_like_tuple("src_stride", src_stride) + + assert out, "out is required" + assert out.dtype == src.dtype.element_ty, "out should have same dtype as src" + + # use index type for end_offset, start_offset and src_stride. + self.arg_type['end_offset'] = index.dtype + self.arg_type['start_offset'] = index.dtype + self.arg_type['src_stride'] = index.dtype + self.extra_attr = f"src_stride_len={len(src_stride)}" + + +@register_custom_op +class _index_put: + """This operation assigns values from the value UB tensor to the dst GM buffer + at positions with offsets specified by the index UB tensor along the specified + scatter dimension with SIMT template. This operation supports 2D-5D. + + Arguments: + - dst: the destination tensor pointer (in GM) + - index: the index tensor (in UB) + - value: the value tensor to be put (in UB) + - dim: the dimension on which index is applied + - bound: upper bound of index + - dst_shape: the shape of destination tensor + - dst_offset: the offset of each dimension in destination tensor + - dst_stride: the stride of each dimension in destination tensor + """ + name = '__builtin_index_put' + core = CORE.VECTOR + pipe = PIPE.PIPE_V + mode = MODE.SIMT + + def __init__(self, dst, index, value, dim, bound: tl.int64, dst_shape, dst_offset, dst_stride): + assert dst.type.is_ptr() or dst.dtype.is_ptr(), f"dst should be a pointer, but got {dst.type}" + assert index.dtype.is_int(), "index should be integer tensor" + value_rank = len(value.shape) + assert 2 <= value_rank <= 5, f"value rank should in [2, 5], but got {value_rank}" + assert isinstance(dim, int), "dim should be an integer" + assert isinstance(bound, int), "bound should be an integer" + assert 0 <= dim < value_rank - 1, f"dim should in [0, {value_rank - 1}), but got {dim}" + assert len(dst_shape) == len(dst_offset), "dst_shape and dst_offset should have same size" + assert len(dst_shape) == len(dst_stride), "dst_shape and dst_stride should have same size" + assert all(isinstance(x, int) for x in dst_shape), "dst_shape should all be integer" + assert all(isinstance(x, int) for x in dst_offset), "dst_offset should all be integer" + assert all(isinstance(x, int) for x in dst_stride), "dst_stride should all be integer" + + # use index type for dst_shape, dst_offset and dst_stride. + self.arg_type['dst_shape'] = index.dtype + self.arg_type['dst_offset'] = index.dtype + self.arg_type['dst_stride'] = index.dtype + + +@register_custom_op +class _gather_load: + """This operation takes a source memory GM buffer and a UB tensor of index, + and produces an output UB tensor by gathering elements from the source + at the index positions with offsets. This operation supports 1D-5D. + + Arguments: + - src: pointer to the source memory GM buffer + - index: UB tensor that specifying the position in the src + - bound: upper bound of index + - dim: the dimension to gather along + - src_stride: the stride of the source tensor + - index_shape: the shape of the index tensor + - offsets: the offsets of each dimension for index tensor + - out: the gathered UB tensor, with the same shape as index.shape + """ + name = '__builtin_gather_load' + core = CORE.VECTOR + pipe = PIPE.PIPE_V + mode = MODE.SIMT + + def __init__(self, src, index, bound: tl.int64, dim, src_stride: tl.int64, index_shape, offsets, out=None): + assert src.type.is_ptr() or src.dtype.is_ptr(), f"src should be a pointer, but got {src.type}" + assert index.dtype.is_int(), "index should be an integer tensor" + assert isinstance(bound, int), "bound should be an integer" + assert isinstance(dim, int), "dim should be an integer" + idx_rank = len(index.shape) + assert 1 <= idx_rank <= 5, f"index rank should in [1, 5], but got {idx_rank}" + assert 0 <= dim < idx_rank, f"dim should in [0, {idx_rank}), but got {dim}" + assert len(src_stride) == idx_rank, f"src_stride size should be {idx_rank}" + assert len(index_shape) == idx_rank, f"index_shape size should be {idx_rank}" + assert len(offsets) == idx_rank, f"offsets size should be {idx_rank}" + assert all(isinstance(x, int) for x in src_stride), "src_stride should all be integer" + assert all(isinstance(x, int) for x in index_shape), "index_shape should all be integer" + assert all(isinstance(x, int) for x in offsets), "offsets should all be integer" + assert out, "out is required" + assert out.shape == index.shape, "Output should have same shape as index" + + +@register_custom_op +class _scatter_store: + """This operation assigns values from the UB tensor to the dst GM buffer at positions with + offsets specified by the UB index tensor with SIMT template. this operation supports 2D-5D. + + Arguments: + - dst: the pointer of destination tensor GM memory buffer + - value: value tensor from UB to store + - index: index tensor from UB specifying positions in the destination tensor + - bound: upper bound of index + - dim: dimension along which the assignment operation is performed + - dst_stride: the strides of destination + - index_shape: the shape of index tensor + - offsets: the offsets for each dims of the index + """ + name = '__builtin_scatter_store' + core = CORE.VECTOR + pipe = PIPE.PIPE_V + mode = MODE.SIMT + + def __init__(self, dst, value, index, bound: tl.int64, dim, dst_stride: tl.int64, index_shape, offsets): + assert dst.type.is_ptr() or dst.dtype.is_ptr(), f"dst should be a pointer, but got {dst.type}" + assert index.dtype.is_int(), "index should be an integer tensor" + assert isinstance(value, tl.tensor), "value should be a tensor" + assert isinstance(bound, int), "bound should be an integer" + assert isinstance(dim, int), "dim should be an integer" + idx_rank = len(index.shape) + assert 1 <= idx_rank <= 5, f"index rank should in [1, 5], but got {idx_rank}" + assert 0 <= dim < idx_rank, f"dim should in [0, {idx_rank}), but got {dim}" + assert len(dst_stride) == idx_rank, f"dst_stride size should be {idx_rank}" + assert len(index_shape) == idx_rank, f"index_shape size should be {idx_rank}" + assert len(offsets) == idx_rank, f"offsets size should be {idx_rank}" + assert all(isinstance(x, int) for x in dst_stride), "dst_stride should all be integer" + assert all(isinstance(x, int) for x in index_shape), "index_shape should all be integer" + assert all(isinstance(x, int) for x in offsets), "offsets should all be integer" diff --git a/third_party/ascend/language/cann/extension/code_generator.py b/third_party/ascend/language/cann/extension/code_generator.py new file mode 100644 index 0000000000..b49f0636d0 --- /dev/null +++ b/third_party/ascend/language/cann/extension/code_generator.py @@ -0,0 +1,203 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Ascend-specific code generation handlers for 'with' statement context managers. +""" + +__all__ = ["handle_scope_with", "mangle_ty"] +import ast + + +def mangle_ty(ty): + """ + Replacement implementation for triton.compiler.code_generator.mangle_ty. + + This is registered via ASCEND_WITH_DISPATCH["mangle_ty"] and picked up by + triton.compiler.code_generator through its global WITH_DISPATCH table. + """ + # Lazy imports to avoid circular dependencies at module import time. + from triton import language + from triton.extension.buffer.language import core as bl + # commonir: TLE buffer types use the same mangle format as extension buffers. + from triton.experimental.tle.language.dsa.types import buffer_type as tle_buffer_type + + # Buffer types are Python-side dtypes; handle them first. + if isinstance(ty, (bl.buffer_type, tle_buffer_type)): + elt = mangle_ty(ty.element_ty) + shape = "_".join(map(str, ty.shape)) + return f"B{elt}S{shape}S" + + if ty.is_ptr(): + return "P" + mangle_ty(ty.element_ty) + if ty.is_int(): + SIGNED = language.dtype.SIGNEDNESS.SIGNED + prefix = "i" if ty.int_signedness == SIGNED else "u" + return prefix + str(ty.int_bitwidth) + if ty.is_floating(): + return str(ty) + if ty.is_block(): + elt = mangle_ty(ty.scalar) + shape = "_".join(map(str, ty.shape)) + return f"{elt}S{shape}S" + if ty.is_void(): + return "V" + raise TypeError(f"Unsupported type {ty}") + + +def _extract_scope_attributes(context_expr): + """Extract attributes from scope(...) call.""" + scope_attrs = {} + for keyword in context_expr.keywords: + if isinstance(keyword.value, ast.Constant): + scope_attrs[keyword.arg] = keyword.value.value + return scope_attrs + + +def _py_value_to_mlir_attr(builder, value): + """Convert Python value to MLIR attribute.""" + attr_creators = { + str: lambda v: builder.get_string_attr(v), + bool: lambda v: builder.get_bool_attr(v), + int: lambda v: builder.get_int32_attr(v), + list: lambda v: builder.get_i64_array_attr(v), + } + creator = attr_creators.get(type(value)) + return creator(value) if creator else value + + +def _handle_core_mode_attr(builder, core_mode): + """Handle core_mode attribute conversion.""" + if core_mode not in ("cube", "vector"): + return {} + return { + builder.get_t_core_type_attr_name(): + (builder.get_t_core_type_cube_attr() if core_mode == "cube" else builder.get_t_core_type_vector_attr()) + } + + +def _build_mlir_attrs_from_scope_attrs(builder, scope_attrs): + """Convert Python scope attributes to MLIR attributes. + + Args: + builder: The IR builder + scope_attrs: Dict of scope attributes (e.g., {'core_mode': 'vector', 'noinline': True}) + + Returns: + Dict of MLIR attributes + """ + mlir_attrs = {"noinline": builder.get_unit_attr()} + for k, v in scope_attrs.items(): + if k == "core_mode": + mlir_attrs.update(_handle_core_mode_attr(builder, v)) + elif k == "noinline": + if not v: + mlir_attrs.pop("noinline") + elif k == "disable_auto_sync": + if v: + mlir_attrs["hivm.disable_auto_sync"] = _py_value_to_mlir_attr(builder, v) + else: + mlir_attrs[k] = _py_value_to_mlir_attr(builder, v) + return mlir_attrs + + +def _verify_loop_carried_variable(_is_triton_value, _is_triton_tensor, name, loop_val, live_val): + """Verify that loop-carried variable types are consistent.""" + assert _is_triton_value(loop_val), f'cannot reassign constxpr {name} in the loop' + assert _is_triton_value(live_val), f'cannot reasign constexpr {name} in the loop' + assert type(loop_val) == type(live_val), f'Loop carried variable {name} changed type' + assert not _is_triton_tensor(loop_val) or loop_val.type == live_val.type, \ + f'Loop-carried variable {name} has initial type {live_val.type} '\ + f'but is re-assigned to {loop_val.type} in loop! '\ + f'Please make sure that the type stays consistent.' + + +def _reconstruct_value_from_ir(language, entry_block_arg, ret_type): + """Reconstruct a tensor value from IR.""" + return language.core.tensor(entry_block_arg, ret_type) + + +def handle_scope_with(generator, node): + """ + Handle 'with scope(...)' statements by creating a scope.scope operation. + + This creates a scope.scope operation with a region for the scope block. + Uses SSA threading to properly handle variables modified inside the scope. + + Args: + generator: The CodeGenerator instance + node: AST node for the with statement + """ + # Lazy imports to avoid circular dependency + from triton import language + from triton.compiler.code_generator import enter_sub_region, _is_triton_value, _is_triton_tensor + + context_expr = node.items[0].context_expr + scope_attrs = _extract_scope_attributes(context_expr) + + with enter_sub_region(generator) as sr: + liveins, _ = sr + ip, last_loc = generator._get_insertion_point_and_loc() + + # This implementation is similar to visit_while + dummy = generator.builder.create_block() + generator.builder.set_insertion_point_to_start(dummy) + generator.visit_compound_statement(node.body) + scope_defs = generator.local_defs + dummy.erase() + + # Verify and get return type of the scope.scope + # (variables that exist in parent scope AND are modified in scope) + names = [] + ret_types = [] + for name in scope_defs: + scope_val = scope_defs[name] + ret_types.append(scope_val.type) + names.append(name) + if name in liveins: + live_val = liveins[name] + _verify_loop_carried_variable(_is_triton_value, _is_triton_tensor, name, scope_val, live_val) + + # Convert Python primitives to MLIR attributes + mlir_attrs = _build_mlir_attrs_from_scope_attrs(generator.builder, scope_attrs) + + # Create scope operation with operands (values from outside) + generator._set_insertion_point_and_loc(ip, last_loc) + scope_op = generator.builder.create_scope_op(mlir_attrs, [ty.to_ir(generator.builder) for ty in ret_types]) + + # Create the entry block with arguments matching the operands + entry_block = generator.builder.create_block_with_parent(scope_op.get_region(0), []) + generator.builder.set_insertion_point_to_start(entry_block) + + # Initialize the scope's symbol table with liveins + generator.lscope = liveins.copy() + generator.visit_compound_statement(node.body) + generator.builder.set_insertion_point_to_end(entry_block) + reconstructed_values = [] + + for i in range(len(names)): + # generator.lscope[names[i]] is already a tensor, just get its IR handle + reconstructed_values.append(generator.lscope[names[i]].handle) + generator.builder.scope_return(reconstructed_values) + + # After exiting enter_sub_region, update symbol table with results + # Convert IR values back to tensor objects + for i, name in enumerate(names): + generator.set_value(name, _reconstruct_value_from_ir(language, scope_op.get_result(i), ret_types[i])) + return None diff --git a/third_party/ascend/language/cann/extension/core.py b/third_party/ascend/language/cann/extension/core.py new file mode 100644 index 0000000000..426e133b26 --- /dev/null +++ b/third_party/ascend/language/cann/extension/core.py @@ -0,0 +1,478 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +__all__ = [ + "ascend_address_space", "builtin", "CORE", "copy_from_ub_to_l1", "copy", "debug_barrier", "fixpipe", + "FixpipeDMAMode", "FixpipeDualDstMode", "FixpipePreQuantMode", "FixpipePreReluMode", "int64", "is_builtin", "MODE", + "PIPE", "IteratorType", "sub_vec_id", "sub_vec_num", "sync_block_all", "sync_block_set", "sync_block_wait", + "SYNC_IN_VF", "conv1d" +] + +import enum +from typing import TypeVar, List, Union +from functools import wraps + +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir +import triton.language.core as tl + +import triton.extension.buffer.language as bl +from triton.language.core import _unwrap_if_constexpr +from triton.backends.ascend.driver import NPUUtils + +from . import semantic as semantic + +PIPE = semantic.PIPE + +T = TypeVar("T") + +TRITON_BUILTIN = "__triton_builtin__" +ASCEND_BUILTIN = "__ascend_builtin__" + + +def builtin(fn: T) -> T: + """Mark a function as a buffer language builtin.""" + assert callable(fn) + + @wraps(fn) + def wrapper(*args, **kwargs): + if "_semantic" not in kwargs or kwargs["_semantic"] is None: + raise ValueError("Did you forget to add @triton.jit ? " + "(`_semantic` argument must be provided outside of JIT functions.)") + return fn(*args, **kwargs) + + # also set triton_builtin to true so that CodeGenerator will recognize this function + setattr(wrapper, TRITON_BUILTIN, True) + setattr(wrapper, ASCEND_BUILTIN, True) + + return wrapper + + +def is_builtin(fn) -> bool: + """Is this a registered ascend language builtin function?""" + return getattr(fn, ASCEND_BUILTIN, False) + + +class int64(int): + """ + For custom op, python int argument will be converted to int32 by default, + if a device-side int64 is required, you can pass an al.int64(x) to it. + """ + + def __new__(cls, value): + obj = int.__new__(cls, value) + obj.type = tl.int64 + return obj + + +class CORE(enum.Enum): + VECTOR = ascend_ir.CoreType.VECTOR + CUBE = ascend_ir.CoreType.CUBE + CUBE_OR_VECTOR = ascend_ir.CoreType.CUBE_OR_VECTOR + CUBE_AND_VECTOR = ascend_ir.CoreType.CUBE_AND_VECTOR + + +class PIPE(enum.Enum): + PIPE_S = ascend_ir.PIPE.PIPE_S + PIPE_V = ascend_ir.PIPE.PIPE_V + PIPE_M = ascend_ir.PIPE.PIPE_M + PIPE_MTE1 = ascend_ir.PIPE.PIPE_MTE1 + PIPE_MTE2 = ascend_ir.PIPE.PIPE_MTE2 + PIPE_MTE3 = ascend_ir.PIPE.PIPE_MTE3 + PIPE_ALL = ascend_ir.PIPE.PIPE_ALL + PIPE_FIX = ascend_ir.PIPE.PIPE_FIX + + +class MODE(enum.Enum): + SIMD = ascend_ir.MODE.SIMD + SIMT = ascend_ir.MODE.SIMT + MIX = ascend_ir.MODE.MIX + + +class IteratorType(enum.Enum): + Parallel = ascend_ir.IteratorType.Parallel + Broadcast = ascend_ir.IteratorType.Broadcast + Transpose = ascend_ir.IteratorType.Transpose + Reduction = ascend_ir.IteratorType.Reduction + Interleave = ascend_ir.IteratorType.Interleave + Deinterleave = ascend_ir.IteratorType.Deinterleave + Inverse = ascend_ir.IteratorType.Inverse + Pad = ascend_ir.IteratorType.Pad + Concat = ascend_ir.IteratorType.Concat + Gather = ascend_ir.IteratorType.Gather + Cumulative = ascend_ir.IteratorType.Cumulative + Opaque = ascend_ir.IteratorType.Opaque + + +class ascend_address_space_base(bl.address_space): + + def __init__(self, address_space_value: ascend_ir.AddressSpace) -> None: + super().__init__() + self.real_address_space = address_space_value + + def to_ir(self, builder: ir.builder) -> ir.attribute: + return builder.get_target_attribute(self.real_address_space) + + +class ascend_address_space_group: + + def __init__(self): + for k, v in {k: v + for k, v in ascend_ir.AddressSpace.__dict__.items() + if isinstance(v, ascend_ir.AddressSpace)}.items(): + setattr(self, k, ascend_address_space_base(v)) + + +ascend_address_space = ascend_address_space_group() + + +@builtin +def sub_vec_id(_semantic=None) -> tl.tensor: + """ + Get the Vector Core index on the AI Core. + """ + return semantic.sub_vec_id(_semantic) + + +@builtin +def copy_from_ub_to_l1(src: Union[tl.tensor, bl.buffer], dst: Union[tl.tensor, bl.buffer], _semantic: None) -> None: + """ + Copies data from the Unified Buffer (UB) to the L1 Buffer. + + :param src: The source data located in the Unified Buffer. + :type src: tl.tensor | bl.buffer + :param dst: The destination buffer located in L1 memory. + :type dst: tl.tensor | bl.buffer + """ + from warnings import warn + warn("copy_from_ub_to_l1 is deprecated, please use copy instead.") + return semantic.copy_from_ub_to_l1(src, dst, _semantic) + + +@builtin +def copy(src: Union[tl.tensor, bl.buffer], dst: Union[tl.tensor, bl.buffer], _semantic: None) -> None: + """ + Copies data from the Unified Buffer (UB) to the Unified Buffer (UB) or L1 Buffer. + + :param src: The source data located in the Unified Buffer. + :type src: tl.tensor | bl.buffer + :param dst: The destination buffer located Unified Buffer (UB) or L1 memory. + :type dst: tl.tensor | bl.buffer + """ + return semantic.copy(src, dst, _semantic) + + +def create_sync_block(sender, receiver, event_id, is_set: bool, sender_pipe=None, receiver_pipe=None, _semantic=None): + sender = _unwrap_if_constexpr(sender) + receiver = _unwrap_if_constexpr(receiver) + assert isinstance(sender, str) and (sender == "cube" + or sender == "vector"), f"ERROR: sender = {sender}, only supports cube/vector" + assert isinstance(receiver, str) and (receiver == "cube" or receiver + == "vector"), f"ERROR: receiver = {receiver}, only supports cube/vector" + if isinstance(event_id, int): + assert (event_id >= 0) and (event_id < 16), f"event_id: {event_id} should be 0 ~ 15" + if sender == receiver: + raise ValueError(f'Unexpected pair: {sender} -> {receiver}, only supports cube -> vector or vector -> cube') + if sender_pipe is None and receiver_pipe is None: + if sender == "cube": + sender_pipe = PIPE.PIPE_FIX + receiver_pipe = PIPE.PIPE_MTE2 + if sender == "vector": + sender_pipe = PIPE.PIPE_MTE3 + receiver_pipe = PIPE.PIPE_MTE2 + if not isinstance(sender_pipe, PIPE) or not isinstance(receiver_pipe, PIPE): + raise TypeError("sender_pipe and receiver_pipe must be instances of PIPE enum") + if is_set: + return semantic.create_sync_block_set(sender, receiver, event_id, sender_pipe, receiver_pipe, _semantic) + return semantic.create_sync_block_wait(sender, receiver, event_id, sender_pipe, receiver_pipe, _semantic) + + +@builtin +def sync_block_set(sender, receiver, event_id, sender_pipe=None, receiver_pipe=None, _semantic=None): + return create_sync_block(sender, receiver, event_id, True, sender_pipe, receiver_pipe, _semantic) + + +@builtin +def sync_block_wait(sender, receiver, event_id, sender_pipe=None, receiver_pipe=None, _semantic=None): + return create_sync_block(sender, receiver, event_id, False, sender_pipe, receiver_pipe, _semantic) + + +@builtin +def sync_block_all(mode, event_id, _semantic=None): + mode = _unwrap_if_constexpr(mode) + event_id = _unwrap_if_constexpr(event_id) + assert isinstance(mode, str), f"mode: {mode} is not string" + assert isinstance(event_id, int) and (event_id >= 0) and (event_id < 16), f"event_id: {event_id} should be 0 ~ 15" + assert mode in ("all_cube", "all_vector", "all", + "all_sub_vector"), f"ERROR: mode = {mode}, only supports all_cube/all_vector/all/all_sub_vector" + _semantic.builder.sync_block_all(mode, event_id) + + +class FixpipeDMAMode(enum.Enum): + NZ2DN = ascend_ir.FixpipeDMAMode.NZ2DN + NZ2ND = ascend_ir.FixpipeDMAMode.NZ2ND + NZ2NZ = ascend_ir.FixpipeDMAMode.NZ2NZ + + +class FixpipeDualDstMode(enum.Enum): + NO_DUAL = ascend_ir.FixpipeDualDstMode.NO_DUAL + COLUMN_SPLIT = ascend_ir.FixpipeDualDstMode.COLUMN_SPLIT + ROW_SPLIT = ascend_ir.FixpipeDualDstMode.ROW_SPLIT + + +class FixpipePreQuantMode(enum.Enum): + NO_QUANT = ascend_ir.FixpipePreQuantMode.NO_QUANT + F322BF16 = ascend_ir.FixpipePreQuantMode.F322BF16 + F322F16 = ascend_ir.FixpipePreQuantMode.F322F16 + S322I8 = ascend_ir.FixpipePreQuantMode.S322I8 + + +class FixpipePreReluMode(enum.Enum): + LEAKY_RELU = ascend_ir.FixpipePreReluMode.LEAKY_RELU + NO_RELU = ascend_ir.FixpipePreReluMode.NO_RELU + NORMAL_RELU = ascend_ir.FixpipePreReluMode.NORMAL_RELU + P_RELU = ascend_ir.FixpipePreReluMode.P_RELU + + +@builtin +def fixpipe( + src: tl.tensor, + dst: Union[None, bl.buffer] = None, + dma_mode: FixpipeDMAMode = FixpipeDMAMode.NZ2ND, + dual_dst_mode: FixpipeDualDstMode = FixpipeDualDstMode.NO_DUAL, + _semantic=None, +) -> Union[None, tl.tensor]: + """ + Directly store a tensor on L0C to a local buffer via fixpipe. + Fixpipe is pipeline that performing data movement from L0C to other memory hierarchies. + Currently support: + - L0C to UB (for Ascend910_95 sereies) + + :param src: the source tensor, Must be located in the l0C memory region. + :type src: tl.tensor + :param dst: The destination buffer in UB. If None, an empty tensor will be created. + Must be buffer type in UB if provided. + :type dst: bl.buffer | None + :param dma_mode: DMA transfer mode, "nz2nd" enables NZ to ND layout transformation + :type dma_mode: FixpipeDMAMode + :param dual_dst_mode: Dual destination mode for split operations + :type dual_dst_mode: FixpipeDualDstMode + :return: If dst is None, returns the created tensor; otherwise returns None. + :rtype: tl.tensor | None + """ + if not _semantic.builder.is_910_95(): + raise RuntimeError("this feature is only supported on Ascend910_95") + if not isinstance(src, tl.tensor): + raise TypeError("src is not of tensor type") + + if dst is not None: + if not isinstance(dst, bl.buffer): + raise TypeError("dst must be buffer type or None") + if dst.space != ascend_address_space.UB: + raise TypeError("dst's AddressSpace must be UB") + if len(dst.shape) == 2 and (dst.type.element_ty == tl.float32 or dst.type.element_ty == tl.int32): + N = dst.shape[1] + if N % 8 != 0: + raise ValueError("32b Fixpipe last dim must be aligned to 8") + if (dma_mode != FixpipeDMAMode.NZ2ND) and (N % 16 != 0): + raise ValueError("32b non-NZ2ND Fixpipe last dim must be aligned to 16") + if (dual_dst_mode == FixpipeDualDstMode.COLUMN_SPLIT) and (N % 32 != 0): + raise ValueError("32b Column split dual Fixpipe last dim must be aligned to 32") + M = dst.shape[0] + if (dma_mode == FixpipeDMAMode.NZ2DN) and (M % 8 != 0): + raise ValueError("32b NZ2DN Fixpipe first dim must be aligned to 8") + dst16bits = (dst.type.element_ty == tl.float16 or dst.type.element_ty == tl.int16 + or dst.type.element_ty == tl.bfloat16) + if len(dst.shape) == 2 and dst16bits: + N = dst.shape[1] + if N % 16 != 0: + raise ValueError("16b Fixpipe last dim must be aligned to 16") + M = dst.shape[0] + if (dma_mode == FixpipeDMAMode.NZ2DN) and (M % 16 != 0): + raise ValueError("16b NZ2DN Fixpipe first dim must be aligned to 16") + + return semantic.fixpipe( + src, + dst, + dma_mode, + dual_dst_mode, + FixpipePreQuantMode.NO_QUANT, + FixpipePreReluMode.NO_RELU, + _semantic, + ) + + +class SYNC_IN_VF(enum.Enum): + VV_ALL = enum.auto() + VST_VLD = enum.auto() + VLD_VST = enum.auto() + VST_VST = enum.auto() + VS_ALL = enum.auto() + VST_LD = enum.auto() + VLD_ST = enum.auto() + VST_ST = enum.auto() + SV_ALL = enum.auto() + ST_VLD = enum.auto() + LD_VST = enum.auto() + ST_VST = enum.auto() + + +@builtin +def debug_barrier( + sync_mode: SYNC_IN_VF, + _semantic=None, +) -> None: + return semantic.debug_barrier(sync_mode.name, _semantic) + + +@builtin +def sub_vec_num(_semantic=None) -> tl.constexpr: + """ + Get the Vector Core Num on one AI Core. + """ + npuUtils = NPUUtils() + cube_num = npuUtils.get_aivector_core_num() + vector_num = npuUtils.get_aicore_num() + const_val = cube_num // vector_num + return tl.constexpr(const_val) + + +@builtin +def conv1d(input: tl.tensor, weight: tl.tensor, bias: tl.tensor = None, stride=None, padding_size=None, dilation=None, + groups=None, _semantic=None) -> tl.tensor: + """ + Applies a 1D convolution over an input signal. + + :param input: Input tensor of shape (N, C_in, L_in) or (C_in, L_in). N is a batch size, C denotes a number of channels, L is a length of signal sequence. + :type input: tensor + :param weight: Weight tensor of shape (C_out, C_in // groups, kernel_size). + :type weight: tensor + :param bias: Bias tensor of shape (C_out) or None. Default: None. + :type bias: tensor or None + :param stride: The stride of the convolution kernel. Can be an int or a 1-element tuple. + :type stride: int or Tuple[int] + :param padding_size: Padding added to both sides of the input. Can be an int, a 1-element tuple, or a string. Can be a string {'valid', 'same'}, single number or a one-element tuple. + ``padding_size='valid'`` is the same as no padding. + ``padding_size='same'`` pads the input so the output has the same shape as the input. However, this mode doesn't support any stride values other than 1. + :type padding_size: int, Tuple[int], or str + :param dilation: The spacing between kernel elements. Can be an int or a 1-element tuple. + :type dilation: int or Tuple[int] + :param groups: Number of blocked connections from input to output channels. + :type groups: int + + **Example:** + + .. code-block:: python + + @triton.jit + def conv_kernel(input_ptr, weight_ptr, bias_ptr, output_ptr, N, C, L, K, BLOCK_SIZE: tl.constexpr): + # Load a tile of input and weight + input_block = tl.load(input_ptr + ...) + weight_block = tl.load(weight_ptr + ...) + + # Perform 1D convolution + # Using default stride=1, padding_size=0, dilation=1, groups=1 + conv_output = al.conv1d( + input_block, + weight_block, + bias=None, + stride=1, + padding_size=0, + dilation=1, + groups=1, + ) + + # Store the result + tl.store(output_ptr + ..., conv_output) + + :return: The output tensor of shape (N, C_out, L_out). + :rtype: tensor + """ + + stride = _unwrap_if_constexpr(stride) + padding_size = _unwrap_if_constexpr(padding_size) + dilation = _unwrap_if_constexpr(dilation) + groups = _unwrap_if_constexpr(groups) + + # Set default value + stride = stride if stride is not None else 1 + padding_size = padding_size if padding_size is not None else 0 + dilation = dilation if dilation is not None else 1 + groups = groups if groups is not None else 1 + + if type(bias).__name__ == 'constexpr': + bias = getattr(bias, 'value', bias) + if bias is not None: + assert len(bias.shape) == 1, f"bias must be a 1D tensor (C_out), got {len(bias.shape)}D" + assert isinstance(groups, int), f"groups must be an integer, got {groups}" + + def _check_and_normalize_1d_param(param, name): + if param is None: + return None + if isinstance(param, (list, tuple)): + assert len(param) == 1, f"{name} must be an integer or a 1-element tuple, got {param}" + return param[0] + assert isinstance(param, int), f"{name} must be an integer or a 1-element tuple, got {type(param)}" + return param + + stride = _check_and_normalize_1d_param(stride, 'stride') + dilation = _check_and_normalize_1d_param(dilation, 'dilation') + + is_batched = len(input.shape) == 3 + L_in = input.shape[-1] + K = weight.shape[2] + + if isinstance(padding_size, str): + assert padding_size in ['same', 'valid'], f"padding_size string must be 'same' or 'valid', got '{padding_size}'" + else: + padding_size = _check_and_normalize_1d_param(padding_size, 'padding_size') + if isinstance(padding_size, str): + if padding_size == 'valid': + padding_size_int = 0 + elif padding_size == 'same': + #The case where padding_needed is an odd number needs to be handled. + if stride != 1: + raise ValueError("padding_size='same' is only supported when stride=1") + padding_needed = (L_in - 1) * stride + dilation * (K - 1) + 1 - L_in + padding_size_int = padding_needed // 2 + else: + padding_size_int = padding_size if padding_size is not None else 0 + + assert len(input.shape) in [2, 3], f"input must be a 2D (C, L) or 3D (N, C, L) tensor, got {len(input.shape)}D" + assert len( + weight.shape) == 3, f"weight must be a 3D tensor (C_out, C_in // groups, kernel_size), got {len(weight.shape)}D" + + # Create output type + C_in = input.shape[-2] if is_batched else input.shape[0] + C_out = weight.shape[0] + + L_in_val = _unwrap_if_constexpr(input.shape[-1]) + K_val = _unwrap_if_constexpr(weight.shape[2]) + + calculation_result = (L_in_val + 2 * padding_size_int - dilation * (K_val - 1) - 1) / stride + 1 + L_out_val = -int(-calculation_result) + if is_batched: + output_shape = [input.shape[0], C_out, L_out_val] + else: + output_shape = [C_out, L_out_val] + + return semantic.conv1d(input, weight, bias, stride, padding_size_int, dilation, groups, output_shape, _semantic) diff --git a/third_party/ascend/language/cann/extension/custom_op.py b/third_party/ascend/language/cann/extension/custom_op.py new file mode 100644 index 0000000000..75ed166f5e --- /dev/null +++ b/third_party/ascend/language/cann/extension/custom_op.py @@ -0,0 +1,400 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +__all__ = ["custom", "custom_semantic", "register_custom_op"] + +import inspect +import types +import typing +import itertools +import triton.language.core as tl +from . import core + +# Registry for custom op, mapping name to its configuration. +_custom_op_registry = {} + + +def _get_op_class(name): + # Try to get op class in _custom_op_registry. + op_class = _custom_op_registry.get(name) + if op_class is None: + # Allow bulitin custom ops used without registry. + assert name.startswith('__builtin_'), f"Custom Op '{name}' not registered." + # Return a dummy op class for builtin custom op. + op_class = type( + "_builtin_custom_op", (object, ), { + "name": name, + "core": core.CORE.VECTOR, + "pipe": core.PIPE.PIPE_V, + "mode": core.MODE.SIMT, + "signature": inspect.signature(object), + }) + return op_class + + +def _unwrap_constexpr(arg): + if isinstance(arg, tl.constexpr): + return arg.value + if isinstance(arg, (tuple, tl.tuple)): + return tuple(_unwrap_constexpr(x) for x in arg) + if isinstance(arg, list): + return [_unwrap_constexpr(x) for x in arg] + if isinstance(arg, dict): + return {k: _unwrap_constexpr(v) for k, v in arg.items()} + return arg + + +def _to_value(value, _semantic=None, ty=None): + # Try to use 'type' attribute if ty not set. + ty = getattr(value, 'type', ty) if ty is None else ty + if isinstance(value, tl.tensor): + if not value.type.is_block() and isinstance(ty, tl.dtype) and value.type != ty: + # For a scalar variable, if its type is not the expected one + # that specified by type hint 'ty', insert a cast for it. + return _semantic.cast(value, ty).handle + return value.handle + if isinstance(value, bool): + return _semantic.builder.get_int1(value) + if isinstance(value, int): + if isinstance(ty, tl.dtype): + if ty.is_int64(): + return _semantic.builder.get_int64(value) + if ty.is_uint64(): + return _semantic.builder.get_uint64(value) + if ty.is_int32(): + return _semantic.builder.get_int32(value) + if ty.is_uint32(): + return _semantic.builder.get_uint32(value) + if ty.is_int16(): + return _semantic.builder.get_int16(value) + if ty.is_uint16(): + return _semantic.builder.get_uint16(value) + if ty.is_int8(): + return _semantic.builder.get_int8(value) + if ty.is_uint8(): + return _semantic.builder.get_uint8(value) + # default int32 + return _semantic.builder.get_int32(value) + if isinstance(value, float): + if isinstance(ty, tl.dtype): + if ty.is_fp64(): + return _semantic.builder.get_fp64(value) + if ty.is_fp32(): + return _semantic.builder.get_fp32(value) + if ty.is_fp16(): + return _semantic.builder.get_fp16(value) + if ty.is_bf16(): + return _semantic.builder.get_bf16(value) + # default float32 + return _semantic.builder.get_fp32(value) + if isinstance(value, tl.constexpr): + return _to_value(value.value, _semantic) + raise TypeError(f"Unsupported argument type {value} : {type(value)}") + + +def _to_operands(args, _semantic=None): + operands = [] + for value in args: + if value is None: + continue + if isinstance(value, (list, tuple, tl.tuple)): + for item in value: + operands.append(_to_value(item, _semantic)) + else: + operands.append(_to_value(value, _semantic)) + return operands + + +def _get_element_type(ty): + if isinstance(ty, types.GenericAlias): + return typing.get_args(ty)[0] + return ty + + +def _args_to_operands(op, _semantic, args, kwargs): + if not op.signature.parameters: + # Without parameters in signature, use the actual parameter order. + return _to_operands(itertools.chain(args, kwargs.values()), _semantic) + + # Convert arguments to operands according the signature. + operands = [] + bind = op.signature.bind(*args, **kwargs) + for param in op.signature.parameters.values(): + value = bind.arguments.get(param.name) + if value is None: + continue + ty = op.arg_type.get(param.name, param.annotation) + if isinstance(value, (list, tuple, tl.tuple)): + ty = _get_element_type(ty) + for item in value: + operands.append(_to_value(item, _semantic, ty)) + else: + operands.append(_to_value(value, _semantic, ty)) + return operands + + +def _bind_op_arguments(op, args, kwargs): + if not op.signature.parameters: + return None + return op.signature.bind(*args, **kwargs) + + +def _make_align_dim_attrs(op, builder, arg_attrs): + # Find op argument by name using op.align_dim's key + # We want to return a dict mapping for each align_dim key -> int attribute for the actual bound argument value. + name = 'align_dim' + if not hasattr(op, name): + return + + # To find argument indices matching each align_dim key, check the op.signature parameters + # and map align_dim key (argument name) to its index position. + align_arg_indices = {} + if hasattr(op, "signature"): + param_names = list(op.signature.parameters.keys()) + for arg_name in op.align_dim.keys(): + if arg_name in param_names: + align_arg_indices[arg_name] = param_names.index(arg_name) + + for arg, align_val in op.align_dim.items(): + if isinstance(arg, str) and arg in align_arg_indices: + arg_attrs[align_arg_indices[arg]] = {name: builder.get_int_attr(align_val)} + print(arg_attrs[align_arg_indices[arg]]) + elif isinstance(arg, int): + arg_attrs[arg] = {name: builder.get_int_attr(align_val)} + print(arg_attrs[arg]) + else: + assert False, f"{name}'s keys should be string or int" + + +def _make_arg_attrs(op, builder): + num_args = len(op.signature.parameters) if hasattr(op, "signature") else 0 + arg_attrs = [{} for _ in range(num_args)] + + _make_align_dim_attrs(op, builder, arg_attrs) + return arg_attrs + + +def _add_optional_attr(op, name, builder, attrs): + if hasattr(op, name): + attrs[name] = builder.get_string_attr(getattr(op, name)) + + +def _add_bitcode_attr(op, builder, attrs): + name = 'bitcode' + if not hasattr(op, name): + return + + from pathlib import Path + bitcode = Path(getattr(op, name)) + assert bitcode.exists(), f"Provided bitcode ({name}) not exist" + attrs[name] = builder.get_string_attr(str(bitcode.absolute())) + + +def _add_optional_extra_buffer_attr(op, builder, attrs): + name = 'extra_buffers' + if not hasattr(op, name): + return + + extra_buffers = getattr(op, name) + if isinstance(extra_buffers, tuple): + extra_buffers = [extra_buffers] + + extra_buffer_types, extra_buffer_sizes = zip(*extra_buffers) + attrs[name + "_types"] = builder.get_type_array_attr([ty.to_ir(builder) for ty in extra_buffer_types]) + attrs[name + "_sizes"] = builder.get_i64_array_attr(list(extra_buffer_sizes)) + + +def _add_optional_indexing_map_attr(op, builder, attrs): + # Optional indexing map attribute: + # `indexing_map` should be an iterable of al.affine_map (MLIR AffineMap) objects. + name = 'indexing_map' + if not hasattr(op, name): + return + + indexing_map = getattr(op, name) + attrs[name] = builder.get_affine_map_array_attr(indexing_map) + + +def _add_optional_iterator_types_attr(op, builder, attrs): + name = 'iterator_types' + if not hasattr(op, name): + return + + attrs[name] = builder.get_iterator_types_attr([iterator_type.value for iterator_type in getattr(op, name)]) + + +def _make_attrs(op, builder): + attrs = { + 'hivm.tcore_type': builder.get_core_type_attr(op.core.value), + 'hivm.pipe': builder.get_pipe_attr(op.pipe.value), + 'hivm.vf_mode': builder.get_vf_mode_attr(op.mode.value), + } + + if not op.name.startswith('__builtin_'): + assert hasattr(op, 'symbol'), f"Non builtin custom op, symbol is required." + assert hasattr(op, 'bitcode'), f"Non builtin custom op, bitcode path is required." + + # Add bit code path attribute, formalize to abosulte path. + _add_bitcode_attr(op, builder, attrs) + + _add_optional_indexing_map_attr(op, builder, attrs) + _add_optional_iterator_types_attr(op, builder, attrs) + + _add_optional_extra_buffer_attr(op, builder, attrs) + + _add_optional_attr(op, 'symbol', builder, attrs) + _add_optional_attr(op, 'source', builder, attrs) + _add_optional_attr(op, 'compile', builder, attrs) + # Extra attributes can be added here, such as op.extra_attr="attr_a=xx" + _add_optional_attr(op, 'extra_attr', builder, attrs) + + return attrs + + +def _to_result(res, res_types): + assert (len(res) == len(res_types)) + n_res = len(res) + if n_res == 0: + return None + if n_res == 1: + return tl.tensor(res[0], res_types[0]) + return tl.tuple(tl.tensor(res[i], res_types[i]) for i in range(n_res)) + + +def _init_op(op_class, *args, **kwargs): + op = op_class.__new__(op_class) + # Add arg_type dict to support dynamic argument type specifying. + setattr(op, 'arg_type', {}) + if op_class.signature.parameters: + # Init with arguments validate. + op_class.__init__(op, *args, **kwargs) + return op + + +def custom_semantic(name: str, *args, _semantic=None, **kwargs): + name = _unwrap_constexpr(name) + # Get op class according the name. + op_class = _get_op_class(name) + # Convert constexpr to value in arguments. + args = _unwrap_constexpr(args) + kwargs = _unwrap_constexpr(kwargs) + # Create op instance from op class with the arguments. + op = _init_op(op_class, *args, **kwargs) + # Prepare inputs and outputs operands. + out = kwargs.pop('out', []) + outs = out if isinstance(out, (list, tuple, tl.tuple)) else [out] + outputs = _to_operands(outs, _semantic) + inputs = _args_to_operands(op, _semantic, args, kwargs) + builder = getattr(_semantic.builder, '_ascend_builder') + # Setup attributes. + attrs = _make_attrs(op, builder) + arg_attrs = _make_arg_attrs(op, builder) + # Build IR for the custom op. + res = builder.create_custom_op(name, attrs, inputs, outputs, arg_attrs) + # Results with same types as outputs. + res_types = [out.type for out in outs] + return _to_result(res, res_types) + + +def custom_object_semantic(op, args, output_indices=(), _semantic=None): + """Create a custom op from a pre-built op-like object. + + Unlike custom_semantic, this entry takes an already-constructed object + exposing the custom-op attributes (core / pipe / mode / symbol / bitcode, + e.g. a @dialect(name="cann") CANNJITFunction) instead of a registry name. + ``args`` is the full operand list in call order; ``output_indices`` marks + the output/aliased operands by position. Operands are regrouped into + hivm.custom's inputs/outputs ranges, so the exported extern function is + expected to take the non-output arguments in order first and the output + arguments last. + """ + args = _unwrap_constexpr(list(args)) + output_indices = set(_unwrap_constexpr(list(output_indices))) + outs = [value for index, value in enumerate(args) if index in output_indices] + ins = [value for index, value in enumerate(args) if index not in output_indices] + outputs = _to_operands(outs, _semantic) + inputs = _to_operands(ins, _semantic) + builder = getattr(_semantic.builder, '_ascend_builder') + attrs = _make_attrs(op, builder) + res = builder.create_custom_op(op.name, attrs, inputs, outputs, []) + res_types = [out.type for out in outs] + return _to_result(res, res_types) + + +@core.builtin +def custom(name: str, *args, _semantic=None, **kwargs): + """Invoke a custom operation with the given name and arguments.""" + return custom_semantic(name, *args, _semantic=_semantic, **kwargs) + + +def register_custom_op(op): + """Register a custom operation so that we can invoke it using al.custom().""" + assert inspect.isclass(op), "@register_custom_op should decorate on a class." + # Use class name if name not set. + if not hasattr(op, 'name'): + setattr(op, 'name', op.__name__) + # The op name should not be used. + assert op.name not in _custom_op_registry, f"Custom op name '{op.name}' already used." + + # Check required core, pipe, mode fields. + assert hasattr(op, 'core'), "'core' field is required." + assert hasattr(op, 'pipe'), "'pipe' field is required." + assert hasattr(op, 'mode'), "'mode' field is required." + assert isinstance(op.core, core.CORE), "Invalid 'core' field, CORE type is required." + assert isinstance(op.pipe, core.PIPE), "Invalid 'pipe' field, PIPE type is required." + assert isinstance(op.mode, core.MODE), "Invalid 'mode' field, MODE type is required." + # Retrieve arguments signature from __init__ method and save it. + signature = inspect.signature(op) + setattr(op, 'signature', signature) + # Register the custom op configuration. + _custom_op_registry[op.name] = op + return op + + +_dtype_cname_dict = { + 'int1': 'bool', + 'int8': 'int8_t', + 'int16': 'int16_t', + 'int32': 'int32_t', + 'int64': 'int64_t', + 'uint8': 'uint8_t', + 'uint16': 'uint16_t', + 'uint32': 'uint32_t', + 'uint64': 'uint64_t', + 'fp16': 'half', + 'bf16': 'bfloat16_t', + 'fp32': 'float', + 'fp64': 'double', + 'fp8e5': 'float8_e5m2_t', + 'fp8e4nv': 'float8_e4m3_t', + # other float8 types are not supported yet, + # such as 'fp8e4b8', 'fp8e4b15', 'fp8e5b16'. +} + + +def _cname(self): + """Return the corresponding C name of the given tl.dtype""" + return _dtype_cname_dict.get(self.name, self.name) + + +# Add 'cname' property to tl.dtype class. +tl.dtype.cname = property(_cname, None) diff --git a/third_party/ascend/language/cann/extension/dispatch.py b/third_party/ascend/language/cann/extension/dispatch.py new file mode 100644 index 0000000000..d91abdee36 --- /dev/null +++ b/third_party/ascend/language/cann/extension/dispatch.py @@ -0,0 +1,33 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Dispatch table for Ascend-specific 'with' statement context managers. +""" + +from .scope import scope +from .code_generator import handle_scope_with, mangle_ty + +__all__ = ["ASCEND_WITH_DISPATCH"] + +# Registry of 'with' statement handlers for Ascend extension +ASCEND_WITH_DISPATCH = { + scope: handle_scope_with, + "mangle_ty": mangle_ty, +} diff --git a/third_party/ascend/language/cann/extension/math_ops.py b/third_party/ascend/language/cann/extension/math_ops.py new file mode 100644 index 0000000000..fa198ea386 --- /dev/null +++ b/third_party/ascend/language/cann/extension/math_ops.py @@ -0,0 +1,54 @@ +from math import pi as math_pi +from triton.language import core, math +from triton.language.core import float32, int1 +from ..libdevice import atan, isnan, isinf +from triton.runtime.jit import jit + +pi = core.constexpr(math_pi) +half_pi = core.constexpr(0.5 * math_pi) + + +@core._tensor_member_fn +@jit +@math._add_math_2arg_docstr("atan2") +def atan2(y, x): + _is_int8_type_x: core.constexpr = x.dtype.is_int8() + core.static_assert(not _is_int8_type_x, f"Expected dtype fp16/fp32/bf16, but got int8 or int1") + _is_int8_type_y: core.constexpr = y.dtype.is_int8() + core.static_assert(not _is_int8_type_y, f"Expected dtype fp16/fp32/bf16, but got int8 or int1") + _is_floating_type_x: core.constexpr = x.dtype.is_floating() + core.static_assert(_is_floating_type_x == True, f"Expected dtype fp16/fp32/bf16, but got {core.constexpr(x.dtype)}") + _is_floating_type_y: core.constexpr = y.dtype.is_floating() + core.static_assert(_is_floating_type_y == True, f"Expected dtype fp16/fp32/bf16, but got {core.constexpr(y.dtype)}") + base = core.where(x == 0, 0.0, atan(y.to(core.dtype("fp32")) / x.to(core.dtype("fp32")))) + base = core.where((x == 0) & (y > 0), half_pi, base) + base = core.where((x == 0) & (y < 0), -half_pi, base) + + add_pi = core.where((x < 0) & (y >= 0), pi, 0.0) + sub_pi = core.where((x < 0) & (y < 0), -pi, 0.0) + return (base + add_pi + sub_pi).to(x.dtype) + + +@core._tensor_member_fn +@jit +@math._add_math_1arg_docstr("isfinited") +def isfinited(x): + _is_int8_type: core.constexpr = x.dtype.is_int8() + core.static_assert(not _is_int8_type, f"Expected dtype fp16/fp32/bf16, but got int8 or int1") + _is_floating_type: core.constexpr = x.dtype.is_floating() + core.static_assert(_is_floating_type == True, f"Expected dtype fp16/fp32/bf16, but got {core.constexpr(x.dtype)}") + nan_mask = isnan(x) + inf_mask = isinf(x) + return (~nan_mask & ~inf_mask).to(int1) + + +@core._tensor_member_fn +@jit +@math._add_math_1arg_docstr("finitef") +def finitef(x): + _is_int8_type: core.constexpr = x.dtype.is_int8() + core.static_assert(not _is_int8_type, f"finitef only supports float32, but got int8 or int1") + core.static_assert(x.dtype == float32, f"finitef only supports float32, but got {core.constexpr(x.dtype)}") + nan_mask = isnan(x) + inf_mask = isinf(x) + return (~nan_mask & ~inf_mask).to(int1) diff --git a/third_party/ascend/language/cann/extension/mem_ops.py b/third_party/ascend/language/cann/extension/mem_ops.py new file mode 100644 index 0000000000..b9a71d1c9b --- /dev/null +++ b/third_party/ascend/language/cann/extension/mem_ops.py @@ -0,0 +1,521 @@ +import numbers +import triton.language as tl +from triton.language.core import ( + _unwrap_if_constexpr, + _tensor_member_fn, + _unwrap_iterable, + builtin, + constexpr, + dtype, + tensor, + check_bit_width, + _unwrap_if_constexpr, +) + +from typing import Optional, Tuple, List, overload, Union +from triton._C.libtriton import ir + +from ._utils import _convert_elem_to_ir_value + + +@_tensor_member_fn +@builtin +def index_put(ptr: tensor, index: tensor, value: tensor, dim: int, index_boundary: int, end_offset: tuple, + start_offset: tuple, dst_stride: tuple, _semantic=None): + """ + Index put values from a tensor into a destination tensor. + + Index put operation for different tensor ranks: + 1. 2D index scatter (0 <= dim < 1): + 1.1 dim = 0 + out[index[i]][start_offset[1]:end_offset[1]] = value[i][0:end_offset[1]-start_offset[1]] + 2. 3D index scatter (0 <= dim < 2): + 2.1 dim = 0 + out[index[i]][start_offset[1]:end_offset[1]][start_offset[2]:end_offset[2]] + = value[i][0:end_offset[1]-start_offset[1]][0:end_offset[2]-start_offset[2]] + 2.2 dim = 1 + out[start_offset[0]:end_offset[0]][index[j]][start_offset[2]:end_offset[2]] + = value[0:end_offset[0]-start_offset[0]][j][0:end_offset[2]-start_offset[2]] + + + :param ptr: pointer type, the destination tensor pointer (in GM) + :param index: tensor, a index to scatter (in UB) + :param value: tensor, a value to store (in UB) + :param dim: int32, the dimension to scatter along + :param index_boundary: int64, the upper boundary for index values + :param end_offset: tuple of int, the offsets of each dimension for the end of the scatter region + :param start_offset: tuple of int, the offsets of each dimension for the start of the scatter region + :param dst_stride: tuple of int, the stride of each dimension of destination tensor + + Constraints + *********** + - `ptr` and `value` must have the same rank. + - `ptr.dtype` only supports `float16`, `bfloat16`, `float32` currently. + - `index` must be an integer tensor. If `index.rank` != 1, it will be reshaped to 1D. + - `index.numel` must equal `value.shape[dim]`. + - `value` support 2~5D tensors. + - `dim` must be valid (0 <= dim < rank(value) - 1). + + Example + ******* + .. code-block:: python + + import torch + import triton + import triton.language as tl + from triton.language.extra.cann.extension import index_put + + @triton.jit + def simple_index_put_kernel(value_ptr, index_ptr, dst_ptr): + # index tile shape: [2] + index_local = tl.arange(0, 2) + x1_local = tl.arange(0, 2)[None, :] # shape=(1,2) + + index_tile = tl.load(index_ptr + index_local) + value_tile = tl.load(value_ptr + index_local[:, None]*2 + x1_local) + + index_put( + ptr=dst_ptr, + index=index_tile, + value=value_tile, + dim=0, + index_boundary=4, + end_offset=(2, 2), + start_offset=(0, 0), + dst_stride=(2, 1) + ) + + dst = torch.zeros((4,2), device='npu', dtype=torch.float32) + value = torch.tensor([[1.,2.], [3.,4.]], device='npu') + index = torch.tensor([2, 0], device='npu') + + simple_index_put_kernel[(1,)](value, index, dst) + print("IndexPut result:", dst) # ref:[[3.,4.], [0.,0.], [1.,2.], [0.,0.]] + """ + + def index_put_impl(ptr: tl.tensor, index: tl.tensor, value: tl.tensor, dim: int, index_boundary: int, + end_offset: Tuple, start_offset: Tuple, dst_stride: Tuple, _builder: ir.builder): + assert index.dtype.is_int(), "index must be an integer tensor" + if not ptr.dtype.element_ty.is_floating(): + raise ValueError(f"Expected dtype fp16/fp32/bf16, but got {ptr.dtype.element_ty}") + if not isinstance(dim, int): + raise ValueError("dim must be of type tl.constexpr") + + v_rank = len(value.shape) + idx_rank = len(index.shape) + if v_rank < 2 or v_rank > 5: + raise ValueError(f"value rank must be in [2, 5], got value rank={v_rank}") + if dim < 0 or dim >= v_rank - 1: + raise ValueError(f"dim must satisfy 0<=dim 5: + raise ValueError(f"index rank must be in [1, 5], got rank={idx_rank}") + if dim < 0 or dim >= idx_rank: + raise ValueError(f"dim must satisfy 0<=dim 5: + raise ValueError(f"index rank must be in [1, 5], got rank={idx_rank}") + if dim < 0 or dim >= idx_rank: + raise ValueError(f"dim must satisfy 0<=dim 0 + + dim = _unwrap_if_constexpr(dim) + index_boundary = _unwrap_if_constexpr(index_boundary) + value = _unwrap_if_constexpr(value) + + if not _is_ranked_tensor(value) or isinstance(value, constexpr): + element_ty = ptr.type.scalar.element_ty + value = _semantic.full(index.shape, value, element_ty) + return scatter_ub_to_out_impl(ptr, value, index, index_boundary, dim, dst_stride, end_offset, start_offset, + _semantic.builder) + + +@_tensor_member_fn +@builtin +def index_select_simd(src, dim, index, src_shape, src_offset, read_shape, _semantic=None) -> tensor: + """ + Parallel index_select operation from Global Memory to Unified Buffer (SIMD version). + + Selects data from multiple indices along a specified dimension and loads + them as tiles from GM directly to UB with zero-copy semantics. + + :param src: Source tensor pointer (in GM) + :type src: tensor (pointer type) + :param dim: The dimension along which to select indices + :type dim: int or constexpr + :param index: 1D tensor of indices to select (in UB) + :type index: tensor + :param src_shape: Complete shape of the source tensor (can be int or tensor) + :type src_shape: List[Union[int, tensor]] + :param src_offset: Starting offset for reading (can be int or tensor) + :type src_offset: List[Union[int, tensor]] + :param read_shape: Size to read (tile shape, can be int or tensor) + :type read_shape: List[Union[int, tensor]] + + **Constraints:** + + - ``read_shape[dim]`` must be ``-1`` + - ``src_offset[dim]`` can be ``-1`` (will be ignored) + - Boundary handling: ``src_offset + read_shape > src_shape`` automatically + truncates to ``src_shape`` boundary + - Does not check if ``index`` contains out-of-bounds values + + **Example:** + + .. code-block:: python + + @triton.jit + def kernel(src_ptr, output_ptr, indices_ptr, M, N, D, ...): + # Load indices (e.g., [5, 10, 15, 20]) + indices = tl.load(indices_ptr + tl.arange(0, 4)) + + # Example 1: Static shapes (constants) + # Index select from dimension 1 + # src: [8, 100, 256], index_select at dim=1 + # Read: [4, ?, 128] starting from [4, ?, 128] + result = extension.index_select_simd( + src_ptr, + dim=1, + index=indices, + src_shape=[8, 100, 256], + src_offset=[4, -1, 128], + read_shape=[4, -1, 128] + ) + # result shape: [4, 4, 128] + + # Example 2: Dynamic shapes (variables) + result2 = extension.index_select_simd( + src_ptr, + dim=1, + index=indices, + src_shape=[M, N, D], + src_offset=[4, -1, 128], + read_shape=[4, -1, 128] + ) + + tl.store(output_ptr + ..., result) + + :return: Result tensor in UB with shape where ``dim`` is replaced + by the length of ``index`` + :rtype: tensor + """ + + def index_select_simd_impl(src: tl.tensor, dim: int, index: tl.tensor, src_shape: List[Union[int, tl.tensor]], + src_offset: List[Union[int, tl.tensor]], read_shape: List[Union[int, tl.tensor]], + _builder: ir.builder) -> tl.tensor: + # Validate inputs + ndim = len(src_shape) + assert len(src_offset) == ndim, \ + f"src_offset length {len(src_offset)} must match src_shape length {ndim}" + assert len(read_shape) == ndim, \ + f"read_shape length {len(read_shape)} must match src_shape length {ndim}" + assert 0 <= dim < ndim, \ + f"dim={dim} must be in range [0, {ndim})" + assert len(index.shape) == 1, \ + f"index must be 1D tensor, got {len(index.shape)}D" + assert dim < ndim - 1, \ + f"index_select_simd cannot support trailing dimension as dim={dim}, ndim={ndim}" + # Handle both tensor and int offsets (for interpreter mode) + newsrc_shape = [] + for s in src_shape: + if isinstance(s, tensor): + newsrc_shape.append(s.handle) + elif isinstance(s, int): + # For interpreter mode: keep as int + newsrc_shape.append(s) + else: + newsrc_shape.append(s.handle if hasattr(s, 'handle') else s) + newsrc_offset = [] + for s in src_offset: + if isinstance(s, tensor): + newsrc_offset.append(s.handle) + elif isinstance(s, int): + # For interpreter mode: keep as int + newsrc_offset.append(s) + else: + newsrc_offset.append(s.handle if hasattr(s, 'handle') else s) + + # Create output type + return_shape = [index.shape[0] if i == dim else read_shape[i] for i in range(ndim)] + element_ty = src.type.element_ty + output_ty = tl.block_type(element_ty, return_shape) + out = _builder.create_index_select_simd(src.handle, index.handle, dim, newsrc_shape, newsrc_offset, read_shape, + return_shape) + return tl.tensor(out, output_ty) + + dim = _unwrap_if_constexpr(dim) + + # Process shape parameters: convert constexpr to values, keep tensors as-is + def process_param(val): + """Convert constexpr to value, keep tensor or int as-is""" + if isinstance(val, tensor): + return val + else: + return _unwrap_if_constexpr(val) + + newsrc_shape = [_semantic.to_tensor(o) if isinstance(o, constexpr) else o for o in src_shape] + newsrc_offset = [_semantic.to_tensor(o) if isinstance(o, constexpr) else o for o in src_offset] + assert len(index.shape) == 1, "index must be a 1D tensor" + + return index_select_simd_impl(src, dim, index, newsrc_shape, newsrc_offset, read_shape, _semantic.builder) diff --git a/third_party/ascend/language/cann/extension/scope.py b/third_party/ascend/language/cann/extension/scope.py new file mode 100644 index 0000000000..302177a7b2 --- /dev/null +++ b/third_party/ascend/language/cann/extension/scope.py @@ -0,0 +1,71 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +__all__ = ["scope"] + +from triton.language.core import _unwrap_if_constexpr + + +class scope: + """ + Context manager for entering and exiting a scope, where operations within a scope shares some common characteristics. + + Example: + ```python + import triton.language.extra.cann.extension as extension + + @triton.jit + def kernel(x_ptr, y_ptr, N): + # specify annotation + with extension.scope(feature_a=True): + a = tl.load(x_ptr) + b = tl.load(y_ptr) + result = tl.dot(a, b) + ``` + + Reserved keywords: + - `core_mode`: Allows explicitly specify which core type should be used for operations within a code block, helping the compiler generate appropriate code for cube or vector cores. + """ + + def __init__(self, core_mode: str, _builder=None, _semantic=None, **kwargs): + """ + :param core_mode: Either "cube" or "vector" to specify the core type + :param _builder: Internal builder object (set by code_generator) + :param _semantic: Internal semantic object (set by code_generator) + :param kwargs: Additional internal parameters + """ + # Convert constexpr to value if not being called from code generator + self.core_mode = _unwrap_if_constexpr(core_mode) if _builder is None else core_mode + self._builder = _builder + self._semantic = _semantic + + # Validate core_mode + if self.core_mode not in ("cube", "vector"): + raise ValueError(f'core_mode must be "cube" or "vector", got {self.core_mode}') + + def __enter__(self): + if self._builder is None: + raise RuntimeError("scope can only be used inside a Triton kernel") + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + return False diff --git a/third_party/ascend/language/cann/extension/semantic.py b/third_party/ascend/language/cann/extension/semantic.py new file mode 100644 index 0000000000..424cf40315 --- /dev/null +++ b/third_party/ascend/language/cann/extension/semantic.py @@ -0,0 +1,179 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +__all__ = [ + "fixpipe", + "create_address_space", +] + +import enum +from typing import (TypeVar, List, Union) + +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir +import triton.language.core as tl +import triton.language.extra.cann.extension as al +import triton.extension.buffer.language as bl + +T = TypeVar('T') + + +def create_address_space(address_space: ascend_ir.AddressSpace, + builder: ascend_ir.ascendnpu_ir_builder) -> ir.attribute: + return builder.get_target_attribute(address_space) + + +class PIPE(enum.Enum): + PIPE_S = ascend_ir.PIPE.PIPE_S + PIPE_V = ascend_ir.PIPE.PIPE_V + PIPE_M = ascend_ir.PIPE.PIPE_M + PIPE_MTE1 = ascend_ir.PIPE.PIPE_MTE1 + PIPE_MTE2 = ascend_ir.PIPE.PIPE_MTE2 + PIPE_MTE3 = ascend_ir.PIPE.PIPE_MTE3 + PIPE_ALL = ascend_ir.PIPE.PIPE_ALL + PIPE_FIX = ascend_ir.PIPE.PIPE_FIX + + +def create_sync_block_set(sender, receiver, event_id, sender_pipe: PIPE, receiver_pipe: PIPE, _semantic=None): + if isinstance(event_id, int): + _semantic.builder.sync_block_set(sender, receiver, + _semantic.to_tensor(tl.constexpr(event_id)).handle, sender_pipe.value, + receiver_pipe.value) + elif isinstance(event_id, tl.constexpr): + _semantic.builder.sync_block_set(sender, receiver, + _semantic.to_tensor(event_id).handle, sender_pipe.value, receiver_pipe.value) + else: + _semantic.builder.sync_block_set(sender, receiver, event_id.handle, sender_pipe.value, receiver_pipe.value) + + +def create_sync_block_wait(sender, receiver, event_id, sender_pipe: PIPE, receiver_pipe: PIPE, _semantic=None): + if isinstance(event_id, int): + _semantic.builder.sync_block_wait(sender, receiver, + _semantic.to_tensor(tl.constexpr(event_id)).handle, sender_pipe.value, + receiver_pipe.value) + elif isinstance(event_id, tl.constexpr): + _semantic.builder.sync_block_wait(sender, receiver, + _semantic.to_tensor(event_id).handle, sender_pipe.value, receiver_pipe.value) + else: + _semantic.builder.sync_block_wait(sender, receiver, event_id.handle, sender_pipe.value, receiver_pipe.value) + + +def sub_vec_id(_semantic=None) -> tl.tensor: + return tl.tensor(_semantic.builder._ascend_builder.create_get_sub_vec_id(), tl.int64) + + +def copy_from_ub_to_l1(src: Union[tl.tensor, bl.buffer], dst: Union[tl.tensor, bl.buffer], _semantic=None): + if not _semantic.builder.is_910_95(): + raise RuntimeError("this feature is only supported on Ascend910_95") + if isinstance(src, tl.tensor) or isinstance(dst, tl.tensor): + raise TypeError("tensor not support yet") + if src.shape != dst.shape: + raise TypeError("src and dst must have same shape") + if src.dtype != dst.dtype: + raise TypeError("src and dst need to have the same type") + if isinstance(src, bl.buffer) and isinstance(dst, bl.buffer): + if src.space != al.ascend_address_space.UB: + raise TypeError("src's AddressSpace must be UB") + if dst.space != al.ascend_address_space.L1: + raise TypeError("dst's AddressSpace must be L1") + _semantic.builder.create_copy_buffer(src.handle, dst.handle) + else: + raise TypeError("src and dst must be tl.tensor or bl.buffer") + + +def copy(src: Union[tl.tensor, bl.buffer], dst: Union[tl.tensor, bl.buffer], _semantic=None): + if not _semantic.builder.is_910_95(): + raise RuntimeError("this feature is only supported on Ascend910_95") + if isinstance(src, tl.tensor) or isinstance(dst, tl.tensor): + raise TypeError("tensor not support yet") + if src.shape != dst.shape: + raise TypeError("src and dst must have same shape") + if src.dtype != dst.dtype: + raise TypeError("src and dst need to have the same type") + if isinstance(src, bl.buffer) and isinstance(dst, bl.buffer): + if src.space != al.ascend_address_space.UB: + raise TypeError("src's AddressSpace must be UB") + if dst.space not in (al.ascend_address_space.L1, al.ascend_address_space.UB): + raise TypeError("dst's AddressSpace must be UB or L1") + _semantic.builder.create_copy_buffer(src.handle, dst.handle) + else: + raise TypeError("src and dst must be tl.tensor or bl.buffer") + + +def fixpipe( + src: tl.tensor, + dst, + dma_mode, + dual_dst_mode, + pre_quant_mode, + pre_relu_mode, + _semantic=None, +): + if dst is None: + result = _semantic.builder._ascend_builder.create_fixpipe( + src.handle, + dst, + dma_mode.value, + dual_dst_mode.value, + pre_quant_mode.value, + pre_relu_mode.value, + ) + if dual_dst_mode == al.FixpipeDualDstMode.ROW_SPLIT: + new_shape = list(src.type.shape) + if len(new_shape) >= 1 and new_shape[0] > 0: + new_shape[0] = new_shape[0] // 2 + new_type = tl.block_type(src.type.element_ty, new_shape) + return tl.tensor(result, new_type) + elif dual_dst_mode == al.FixpipeDualDstMode.COLUMN_SPLIT: + new_shape = list(src.type.shape) + if len(new_shape) >= 2 and new_shape[1] > 0: + new_shape[1] = new_shape[1] // 2 + new_type = tl.block_type(src.type.element_ty, new_shape) + return tl.tensor(result, new_type) + else: + return tl.tensor(result, src.type) + else: + _semantic.builder._ascend_builder.create_fixpipe( + src.handle, + dst.handle, + dma_mode.value, + dual_dst_mode.value, + pre_quant_mode.value, + pre_relu_mode.value, + ) + + +def debug_barrier(sync_mode: str, _semantic=None) -> None: + target = tl.tensor(_semantic.builder.get_int64(0), tl.int64) + attr = _semantic.builder.get_string_attr(sync_mode) + _semantic.builder.create_debug_barrier(target.handle, "SYNC_IN_VF", attr) + + +def conv1d(input_tensor: tl.tensor, weight_tensor: tl.tensor, bias: Union[tl.tensor, None], stride: int, + padding_size: int, dilation: int, groups: int, output_shape, _semantic=None) -> tl.tensor: + bias_handle = None if bias is None else bias.handle + + element_ty = input_tensor.type.element_ty + output_ty = tl.block_type(element_ty, output_shape) + out = _semantic.builder.create_conv1d(input_tensor.handle, weight_tensor.handle, bias_handle, stride, padding_size, + dilation, groups, output_ty.to_ir(_semantic.builder)) + return tl.tensor(out, output_ty) diff --git a/third_party/ascend/language/cann/extension/vec_ops.py b/third_party/ascend/language/cann/extension/vec_ops.py new file mode 100644 index 0000000000..321fa345f2 --- /dev/null +++ b/third_party/ascend/language/cann/extension/vec_ops.py @@ -0,0 +1,536 @@ +# insert_slice +# extract_slice +# get_element +# sort +# flip +# gather + +import triton.language as tl +from triton.language import semantic, core, standard +from triton.language.core import (_unwrap_if_constexpr, _tensor_member_fn, _unwrap_iterable, builtin, constexpr, dtype, + tensor, check_bit_width, _unwrap_if_constexpr, range) + +from . import is_compile_on_910_95 +from .aux_ops import compile_hint_impl + +from typing import Optional, Tuple, List, overload +from triton._C.libtriton import ir + + +@_tensor_member_fn +@builtin +def insert_slice(ful, sub, offsets, sizes, strides, _semantic=None, _generator=None) -> tensor: + """ + Insert a tensor to another tensor as specified by the operation’s offsets, sizes and strides arguments. + + :param ful: The tensor to receive tensor. + :type ful: Tensor + :param sub: The tensor to be inserted. + :type sub: Tensor + :param offsets: + :type offsets: tuple of ints + :param sizes: + :type sizes: tuple of ints + :param strides: + :type strides: tuple of ints + """ + + def insert_slice_impl(ful: tensor, sub: tensor, offsets: List[tensor], sizes: List[int], strides: List[int], + builder: ir.builder) -> tensor: + assert (len(ful.shape) == len(offsets)) + assert (len(ful.shape) == len(sizes)) + assert (len(ful.shape) == len(strides)) + assert (all([s >= 1 for s in sizes])) + assert (all([s >= 0 for s in strides])) + # Handle both tensor and int offsets (for interpreter mode) + new_offsets = [] + for o in offsets: + if isinstance(o, tensor): + new_offsets.append(o.handle) + elif isinstance(o, int): + # For interpreter mode: keep as int + new_offsets.append(o) + else: + new_offsets.append(o.handle if hasattr(o, 'handle') else o) + ret_type = tl.block_type(ful.type.scalar, ful.shape) + out = builder.create_insert_slice(ful.handle, sub.handle, new_offsets, sizes, strides) + return tensor(out, ret_type) + + assert len(ful.shape) > 0 + assert len(ful.shape) == len(sub.shape) + new_offsets = [_semantic.to_tensor(o) if isinstance(o, constexpr) else o for o in offsets] + out = insert_slice_impl(ful, sub, new_offsets, sizes, strides, _semantic.builder) + return out + + +@_tensor_member_fn +@builtin +def extract_slice(ful, offsets, sizes, strides, _semantic=None, _generator=None) -> tensor: + """ + Extract a tensor from another tensor as specified by the operation’s offsets, sizes and strides arguments. + + :param ful: The tensor to split. + :type ful: Tensor + :param offsets: + :type offsets: tuple of ints + :param sizes: + :type sizes: tuple of ints + :param strides: + :type strides: tuple of ints + """ + + def extract_slice_impl(ful: tensor, offsets: List[tensor], sizes: List[int], strides: List[int], + builder: ir.builder) -> tensor: + assert (len(ful.shape) == len(offsets)) + assert (len(ful.shape) == len(sizes)) + assert (len(ful.shape) == len(strides)) + assert (all([s >= 1 for s in sizes])) + assert (all([s >= 0 for s in strides])) + # Handle both tensor and int offsets (for interpreter mode) + new_offsets = [] + for o in offsets: + if isinstance(o, tensor): + new_offsets.append(o.handle) + elif isinstance(o, int): + # For interpreter mode: keep as int + new_offsets.append(o) + else: + new_offsets.append(o.handle if hasattr(o, 'handle') else o) + ret_type = tl.block_type(ful.type.scalar, sizes) + out = builder.create_extract_slice(ful.handle, new_offsets, sizes, strides) + return tensor(out, ret_type) + + assert len(ful.shape) > 0 + new_offsets = [_semantic.to_tensor(o) if isinstance(o, constexpr) else o for o in offsets] + sub = extract_slice_impl(ful, new_offsets, sizes, strides, _semantic.builder) + return sub + + +@_tensor_member_fn +@builtin +def get_element(src, indice, _semantic=None, _generator=None): + """ + get_element op reads a ranked tensor and returns one element as specified by the given indices. + The result of the op is a value with the same type as the elements of the tensor. + The arity of indices must match the rank of the accessed value. + + :param src: The tensor to be accessed. + :type src: Tensor + :param indice: + :type indice: tuple of ints + """ + + def get_element_impl(src: tensor, indice: List[tensor], builder: ir.builder): + if len(src.shape) != len(indice): + raise ValueError("Indice's rank must be equal to src tensor's rank") + + # Handle both tensor and int indices (for interpreter mode) + new_indice = [] + for i in indice: + if isinstance(i, tensor): + new_indice.append(i.handle) + elif isinstance(i, int): + # For interpreter mode: convert int to TensorHandle + new_indice.append(i) + else: + # Try to use .handle attribute if available + new_indice.append(i.handle if hasattr(i, 'handle') else i) + + result = builder.create_extract_scalar(src.handle, new_indice) + return _semantic.wrap_tensor(result, src.type.scalar, None) + + assert len(src.shape) > 0 + new_indice = [_semantic.to_tensor(i) if isinstance(i, constexpr) else i for i in indice] + return get_element_impl(src, new_indice, _semantic.builder) + + +@builtin +def flip(ptr, dim=-1, _semantic=None, _generator=None): + + def flip_impl(ptr: tensor, dim: int, builder: ir.builder, generator=None): + """ + Flips a tensor `ptr` along the dimension `dim`. + + :param ptr: the first input tensor + :type ptr: tensor + :param dim: the dimension to flip along + :type dim: int + :param generator: the code generator (required for reduce operations) + :type generator: generator object + """ + + def _get_flip_dim(dim, shape): + dim = _unwrap_if_constexpr(dim) + shape = _unwrap_if_constexpr(shape) + if dim is None: + dim = len(shape) - 1 + if dim < 0: # flip doesn't work if dim < 0 because the xor-swap for loop will start/end at the wrong index + dim += len(shape) + return constexpr(dim) + + def _log2(i: core.constexpr): + log2 = 0 + n = core.constexpr(i).value + while n > 1: + n >>= 1 + log2 += 1 + return core.constexpr(log2) + + def flip_simd(ptr: tensor, dim: int, builder: ir.builder): + """ + Triton flip operation for simd + + Args: + ptr: tensor, input tensor + dim: int, dimension to flip (can be negative, normalized here) + builder: ir.builder, underlying IR builder + Returns: + flipped: tensor, same type and shape as input + """ + + shape = getattr(ptr, "shape", None) + if shape is None or shape == (): + shape = getattr(getattr(ptr, "type", None), "shape", None) + + rank = None + if shape is not None: + try: + rank = len(shape) + except Exception: + rank = len(list(shape)) + + if rank is not None: + if rank < 1: + raise ValueError("ascend.flip requires tensor rank >= 1") + norm_dim = dim if dim >= 0 else dim + rank + if not (0 <= norm_dim < rank): + raise ValueError(f"ascend.flip got invalid dim={dim} for shape {tuple(shape)}") + dim = norm_dim + else: + if dim < 0: + raise ValueError("ascend.flip with unknown rank requires non-negative dim") + + flipped_vals = builder.create_flip(ptr.handle, dim) + flipped = tensor(flipped_vals, type=ptr.type) + return flipped + + # If compile_mode is not simt, use the simd implementation + if not builder.is_simt_mode(): + return flip_simd(ptr, dim, builder) + core.static_assert(-len(ptr.shape) <= dim and dim < len(ptr.shape), _builder=builder) + _dim: core.constexpr = _get_flip_dim(dim, ptr.shape) + core.static_assert(standard._is_power_of_two(ptr.shape[_dim]), _builder=builder) + steps: core.constexpr = _log2(ptr.shape[_dim]) + # If steps is 0, return the original tensor + if steps == 0: + return ptr + # reshape the swap dimension to (2, 2, ..., 2) + idtype = core.get_int_dtype(bitwidth=ptr.dtype.primitive_bitwidth, signed=True) + y = core.reshape( + ptr.to(idtype, bitcast=True, _builder=builder), + ptr.shape.__getitem__(slice(None, _dim)) + [2] * steps + ptr.shape.__getitem__(slice(_dim + 1, None)), + _builder=builder) + for i in static_range(steps): + y = y.__xor__(standard.xor_sum(y, _dim + i, True, _builder=builder, _generator=generator), _builder=builder) + ptr = core.reshape(y, ptr.shape, _builder=builder).to(ptr.dtype, bitcast=True, _builder=builder) + return ptr + + try: + dim = int(dim.value) if hasattr(dim, "value") else int(dim) + except Exception as e: + raise TypeError(f"dim must be an integer (or tl.constexpr int), got {dim!r}") from e + + dim = len(ptr.shape) - 1 if dim == -1 else dim + return flip_impl(ptr, dim, _semantic.builder, _generator) + + +class static_range: + """ + Iterator for non-JIT Python functions that need to iterate over constexpr values. + This is used in functions like flip that are called during compilation. + """ + + def __init__(self, arg1, arg2=None, step=None): + if step is None: + self.step = core.constexpr(1) + else: + self.step = step + if arg2 is None: + self.start = core.constexpr(0) + self.end = arg1 + else: + self.start = arg1 + self.end = arg2 + + def __iter__(self): + # Extract actual values from constexpr objects for iteration + start_val = core._unwrap_if_constexpr(self.start) + end_val = core._unwrap_if_constexpr(self.end) + step_val = core._unwrap_if_constexpr(self.step) + # Store as regular Python integers for iteration + self._current = start_val + self._end = end_val + self._step = step_val + return self + + def __next__(self): + if self._current >= self._end: + raise StopIteration + value = self._current + self._current += self._step + return value + + +@builtin +def sort(ptr, dim=-1, descending=False, _semantic=None): + """ + sort the input tensor along 'dim' + + param: + ptr: tensor, input tensor + dim: int or tl.constexpr[int], dimension to sort + descending: bool or tl.constexpr[bool], the result is descending or not + _builder: ir.builder + return: + values: tensor, the sorted tensor + """ + + def sort_impl(ptr: tensor, dim: int, descending, builder: ir.builder): + allowed_types = { + tl.int8, tl.int16, tl.bfloat16, tl.float16, tl.float32, tl.int32, tl.int64, tl.float8e4nv, tl.float8e5 + } + base_ty = ptr.type.scalar if hasattr(ptr.type, "scalar") else ptr.type + if base_ty not in allowed_types: + raise TypeError( + f"ascend.sort only supports int8, int16, bfloat16, float16, float32, int32, int64, float8e4nv, float8e5" + f"but got {ptr.type}") + + shape = getattr(ptr, "shape", None) + if shape is None or shape == (): + shape = getattr(getattr(ptr, "type", None), "shape", None) + + rank = None + if shape is not None: + try: + rank = len(shape) + except Exception: + rank = len(list(shape)) + + if rank is not None: + if rank < 1: + raise ValueError("ascend.sort requires tensor rank >= 1") + last_dim = rank - 1 + norm_dim = dim if dim >= 0 else dim + rank + if norm_dim != last_dim: + raise ValueError(f"ascend.sort only supports sorting along the last dimension " + f"(dim={last_dim} or -1) for shape {tuple(shape)}, but got dim={dim}") + dim = last_dim + else: + if dim != -1: + raise ValueError("ascend.sort only supports the last dimension; when rank is unknown " + "you must pass dim=-1") + + if hasattr(descending, "value"): + descending = bool(descending.value) + else: + descending = bool(descending) + + sorted_vals = builder.create_sort(ptr.handle, dim, descending) + + values = tensor(sorted_vals, type=ptr.type) + + return values + + try: + dim = int(dim.value) if hasattr(dim, "value") else int(dim) + except Exception as e: + raise TypeError(f"dim must be an integer (or tl.constexpr int), got {dim!r}. Error: {str(e)}") from e + + if hasattr(descending, "value"): + descending = bool(descending.value) + else: + descending = bool(descending) + + ret = sort_impl(ptr, dim, descending, _semantic.builder) + # interpreter mode not support compile_hint overflow_mode, direct return + from triton.runtime.interpreter import InterpreterBuilder + if isinstance(_semantic.builder, InterpreterBuilder): + return ret + base_ty = ptr.type.scalar if hasattr(ptr.type, "scalar") else ptr.type + if base_ty.is_int8() or base_ty.is_int16(): + compile_hint_impl(ret, "overflow_mode", constexpr("saturate"), _semantic.builder) + return ret + + +def ascend_cast_impl(input: tensor, dst_ty: dtype, _semantic=None, fp_downcast_rounding: Optional[str] = None, + overflow_mode: Optional[str] = None) -> tensor: + src_ty = input.type + if isinstance(dst_ty, tl.constexpr): + dst_ty = dst_ty.value + if isinstance(fp_downcast_rounding, tl.constexpr): + fp_downcast_rounding = fp_downcast_rounding.value + if src_ty.is_block(): + dst_ty = tl.block_type(dst_ty.scalar, input.type.get_block_shapes()) + if src_ty == dst_ty: + return input + + src_sca_ty = src_ty.scalar + dst_sca_ty = dst_ty.scalar + if src_sca_ty == dst_sca_ty: + return input + + # For fp downcasting default rounding mode should be RTNE, for all other conversions it should + # not be set + fp_downcast_rounding = _semantic._str_to_rounding_mode(fp_downcast_rounding) + use_custom_rounding = False + if dst_sca_ty.is_floating() and src_sca_ty.is_floating( + ) and dst_sca_ty.primitive_bitwidth < src_sca_ty.primitive_bitwidth: + if fp_downcast_rounding is None: fp_downcast_rounding = ir.ROUNDING_MODE.RTNE + elif fp_downcast_rounding != ir.ROUNDING_MODE.RTNE: use_custom_rounding = True + else: + if fp_downcast_rounding is not None: + raise ValueError("fp_downcast_rounding should be set only for truncating fp conversions. " + "Source scalar type is " + str(src_sca_ty) + " and destination type is " + str(dst_sca_ty)) + if not is_compile_on_910_95: + if (src_sca_ty.is_fp8() or dst_sca_ty.is_fp8()) or (src_sca_ty.is_fp64() or dst_sca_ty.is_fp64()): + raise ValueError("[fp8, fp64] is unsupported on Ascend for now." + "Source scalar type is " + str(src_sca_ty) + " and destination type is " + str(dst_sca_ty)) + if (src_sca_ty.is_fp8e4b15() or dst_sca_ty.is_fp8e4b15()): + assert _semantic.builder.codegen_fns.get( + "convert_custom_types") is not None, "target doesn't provide conversion for this type." + return _semantic.builder.codegen_fns["convert_custom_types"](input, dst_ty, fp_downcast_rounding, + _semantic=_semantic) + # Casting with customized floating types involved: fp8 <=> bf16, fp16, fp32, fp64 + # and non-default rounding modes for downcasting + if (src_sca_ty.is_fp8() and dst_sca_ty.is_floating()) or \ + (src_sca_ty.is_floating() and dst_sca_ty.is_fp8()) or \ + use_custom_rounding: + return tensor( + _semantic.builder.create_fp_to_fp(input.handle, dst_ty.to_ir(_semantic.builder), fp_downcast_rounding), + dst_ty) + + # bf16 <=> (not fp32) + if (src_sca_ty.is_fp16() and not dst_sca_ty.is_fp32()) or \ + (src_sca_ty.is_bf16() and not dst_sca_ty.is_fp32()): + return ascend_cast_impl(ascend_cast_impl(input, tl.float32, _semantic), dst_sca_ty, _semantic) + + # Standard floating types' casting: truncation + # fp64 => fp32, fp16, bf16 + # fp32 => fp16, bf16 + truncate_fp = src_sca_ty.is_floating() and \ + dst_sca_ty.is_floating() and \ + src_sca_ty.primitive_bitwidth > dst_sca_ty.primitive_bitwidth + if truncate_fp: + return tensor(_semantic.builder.create_fp_trunc(input.handle, dst_ty.to_ir(_semantic.builder)), dst_ty) + + # Standard floating types' casting: extension + # fp32 => fp64 + # fp16 => fp32, fp64 + # bf16 => fp32, fp64 + ext_fp = src_sca_ty.is_floating() and \ + dst_sca_ty.is_floating() and \ + src_sca_ty.primitive_bitwidth < dst_sca_ty.primitive_bitwidth + if ext_fp: + return tensor(_semantic.builder.create_fp_ext(input.handle, dst_ty.to_ir(_semantic.builder)), dst_ty) + + # Casting between integer types + if src_sca_ty.is_int() and dst_sca_ty.is_int() and \ + (src_sca_ty.int_bitwidth != dst_sca_ty.int_bitwidth or src_sca_ty.int_signedness != dst_sca_ty.int_signedness): + sign_extend = src_sca_ty.is_int_signed() and not src_sca_ty.is_bool() + if dst_sca_ty.is_bool(): + ty = input.dtype.to_ir(_semantic.builder) + _0 = tensor(_semantic.builder.get_null_value(ty), input.dtype) + return _semantic.not_equal(input, _0) + elif overflow_mode == "saturate" and \ + (src_sca_ty.is_int_unsigned() or dst_sca_ty.is_int_unsigned()) and \ + src_sca_ty.int_bitwidth >= dst_sca_ty.int_bitwidth: + if is_compile_on_910_95: + result = tensor( + _semantic.builder.create_int_cast(input.handle, dst_ty.to_ir(_semantic.builder), sign_extend), + dst_ty) + compile_hint_impl(result, "saturate_src_unsigned", src_sca_ty.is_int_unsigned(), _semantic.builder) + compile_hint_impl(result, "saturate_dst_unsigned", dst_sca_ty.is_int_unsigned(), _semantic.builder) + return result + else: + return ascend_cast_impl(ascend_cast_impl(input, tl.float32, _semantic), dst_sca_ty, _semantic) + return tensor(_semantic.builder.create_int_cast(input.handle, dst_ty.to_ir(_semantic.builder), sign_extend), + dst_ty) + + # Casting standard floating types to integer types + if src_sca_ty.is_standard_floating() and dst_sca_ty.is_int(): + if dst_sca_ty.is_bool(): + ty = input.dtype.to_ir(_semantic.builder) + _0 = tensor(_semantic.builder.get_null_value(ty), input.dtype) + return _semantic.not_equal(input, _0) + elif dst_sca_ty.is_int_signed(): + return tensor(_semantic.builder.create_fp_to_si(input.handle, dst_ty.to_ir(_semantic.builder)), dst_ty) + else: + return tensor(_semantic.builder.create_fp_to_ui(input.handle, dst_ty.to_ir(_semantic.builder)), dst_ty) + + # Casting integer types to standard floating types + if src_sca_ty.is_int() and dst_sca_ty.is_standard_floating(): + if src_sca_ty.is_bool() or not src_sca_ty.is_int_signed(): + return tensor(_semantic.builder.create_ui_to_fp(input.handle, dst_ty.to_ir(_semantic.builder)), dst_ty) + else: + return tensor(_semantic.builder.create_si_to_fp(input.handle, dst_ty.to_ir(_semantic.builder)), dst_ty) + + # Casting pointer types to integer types + if src_sca_ty.is_ptr() and dst_sca_ty.is_int(): + bitwidth = dst_sca_ty.int_bitwidth + if bitwidth == 64: + return tensor(_semantic.builder.create_ptr_to_int(input.handle, dst_ty.to_ir(_semantic.builder)), dst_ty) + if bitwidth == 1: + return _semantic.not_equal(ascend_cast_impl(input, tl.int64, _semantic), + tensor(_semantic.builder.get_int64(0), tl.int64)) + + # Casting integer types to pointer types + if src_sca_ty.is_int() and dst_sca_ty.is_ptr(): + return tensor(_semantic.builder.create_int_to_ptr(input.handle, dst_ty.to_ir(_semantic.builder)), dst_ty) + + # Casting pointer types to pointer types + if src_sca_ty.is_ptr() and dst_sca_ty.is_ptr(): + return tensor(_semantic.builder.create_bitcast(input.handle, dst_ty.to_ir(_semantic.builder)), dst_ty) + + assert False, f'cannot cast {input} to {dst_ty}' + + +@_tensor_member_fn +@builtin +def cast(input, dtype: dtype, fp_downcast_rounding: Optional[str] = None, bitcast: bool = False, + overflow_mode: Optional[str] = None, _semantic=None): + """ + Casts a tensor to the given :code:`dtype`. + + :param dtype: The target data type. + :type dtype: dtype + :param fp_downcast_rounding: The rounding mode for downcasting + floating-point values. This parameter is only used when self is a + floating-point tensor and dtype is a floating-point type with a + smaller bitwidth. Supported values are :code:`"rtne"` (round to + nearest, ties to even) and :code:`"rtz"` (round towards zero). + :type fp_downcast_rounding: str, optional + :param bitcast: If true, the tensor is bitcasted to the given + :code:`dtype`, instead of being numerically casted. + :type bitcast: bool, optional + :param overflow_mode: When overflow_mode is not set or is "trunc", + truncation (cut-off) will be used to handle overflow. When + overflow_mode is "sautrate", the maximum value of the data type + will be used to handle overflow. + :type overflow_mode: string, optional + """ + overflow_modes = ["trunc", "saturate"] + input = _semantic.to_tensor(input) + if isinstance(bitcast, constexpr): + bitcast = bitcast.value + if bitcast: + return _semantic.bitcast(input, dtype) + ret = ascend_cast_impl(input, dtype, _semantic, fp_downcast_rounding, overflow_mode) + if overflow_mode is not None: + if overflow_mode in overflow_modes: + from triton.runtime.interpreter import InterpreterBuilder + if isinstance(_semantic.builder, InterpreterBuilder): + overflow_mode = constexpr(overflow_mode) + compile_hint_impl(ret, "overflow_mode", overflow_mode, _semantic.builder) + else: + raise ValueError(f"Unknown overflow_mode:{overflow_mode} is found.") + return ret diff --git a/third_party/ascend/language/cann/libdevice.py b/third_party/ascend/language/cann/libdevice.py new file mode 100644 index 0000000000..2961c2c54f --- /dev/null +++ b/third_party/ascend/language/cann/libdevice.py @@ -0,0 +1,1004 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +from math import pi as math_pi +from triton.language import core, math, semantic +from triton._C.libtriton import ir +from triton.runtime.jit import jit +from triton.backends.ascend.utils import get_ascend_arch_from_env, triton_enable_libdevice_simt +from triton.tools.get_ascend_devices import is_compile_on_910_95 +import triton + + +@core.extern +def reciprocal(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__hmf_recipf", core.dtype("fp32")), + (core.dtype("fp16"), ): ("__hmf_recipDh", core.dtype("fp16")), + }, is_pure=True, _semantic=_semantic) + + +@core.extern +def log1p(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__hmf_log1pf", core.dtype("fp32")), + (core.dtype("fp16"), ): ("__hmf_log1pDh", core.dtype("fp16")), + }, is_pure=True, _semantic=_semantic) + + +@core.extern +def relu(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__hmf_reluf", core.dtype("fp32")), + (core.dtype("fp16"), ): ("__hmf_reluDh", core.dtype("fp16")), + }, is_pure=True, _semantic=_semantic) + + +@core.extern +def isinf(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__hmf_isinf", core.dtype("int1")), + (core.dtype("fp16"), ): ("__hmf_isinf", core.dtype("int1")), + (core.dtype("bf16"), ): ("__hmf_isinf", core.dtype("int1")), + }, is_pure=True, _semantic=_semantic) + + +@core.extern +def tan(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__hmf_tanf", core.dtype("fp32")), + (core.dtype("fp16"), ): ("__hmf_tanDh", core.dtype("fp16")), + }, is_pure=True, _semantic=_semantic) + + +@core.extern +def atan(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__hmf_atanf", core.dtype("fp32")), + (core.dtype("fp16"), ): ("__hmf_atanDh", core.dtype("fp16")), + }, is_pure=True, _semantic=_semantic) + + +@core.extern +def tanh(arg0, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + original_dtype = arg0.dtype + if original_dtype == core.dtype("bf16"): + arg0 = _semantic.cast(arg0, core.float32) + + if triton_enable_libdevice_simt() and is_compile_on_910_95: + dispatch = { + (core.dtype("fp32"), ): ("__hmf_tanh_fp32", core.dtype("fp32")), + } + else: + dispatch = { + (core.dtype("fp32"), ): ("__hmf_tanhf", core.dtype("fp32")), + (core.dtype("fp16"), ): ("__hmf_tanhDh", core.dtype("fp16")), + } + + res = core.extern_elementwise("", "", [arg0], dispatch, is_pure=True, _semantic=_semantic) + if original_dtype == core.dtype("bf16"): + return _semantic.cast(res, core.dtype("bf16")) + return res + + +@core.extern +def ilogb(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__hmf_ilogbf", core.dtype("fp32")), + (core.dtype("fp16"), ): ("__hmf_ilogbDh", core.dtype("fp16")), + }, is_pure=True, _semantic=_semantic) + + +@core.extern +def ldexp(arg0, arg1, _semantic=None): + return core.extern_elementwise( + "", "", [arg0, arg1], { + (core.dtype("fp32"), core.dtype("int32")): ("__hmf_ldexpf", core.dtype("fp32")), + (core.dtype("fp16"), core.dtype("int32")): ("__hmf_ldexpDh", core.dtype("fp16")), + }, is_pure=True, _semantic=_semantic) + + +@core.extern +def pow(arg0, arg1, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + return core.extern_elementwise("", "", [arg0, arg1], { + (core.dtype("fp32"), core.dtype("fp32")): ("__hmf_pow_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + else: + return core.extern_elementwise( + "", "", [arg0, arg1], { + (core.dtype("fp32"), core.dtype("fp32")): ("__hmf_powf", core.dtype("fp32")), + (core.dtype("fp16"), core.dtype("fp16")): ("__hmf_powDh", core.dtype("fp16")), + (core.dtype("bf16"), core.dtype("bf16")): ("__hmf_powDb", core.dtype("bf16")), + }, is_pure=True, _semantic=_semantic) + + +@core.extern +def isnan(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__hmf_isnan", core.dtype("int1")), + (core.dtype("fp16"), ): ("__hmf_isnan", core.dtype("int1")), + (core.dtype("bf16"), ): ("__hmf_isnan", core.dtype("int1")), + }, is_pure=True, _semantic=_semantic) + + +@core.extern +def div_rz(arg0, arg1, _semantic=None): + return core.extern_elementwise("", "", [arg0, arg1], { + (core.dtype("fp32"), core.dtype("fp32")): ("__hmf_div_rz_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + + +@core.builtin +def fast_dividef(arg0, arg1, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + arg1 = _semantic.to_tensor(arg1) + ret = _semantic.fdiv(arg0, arg1, False) + return ret + + +@core.builtin +def fast_expf(arg0, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + ret = core.tensor(_semantic.builder.create_exp(arg0.handle), arg0.type) + return ret + + +@core.extern +def fmod(arg0, arg1, _semantic=None): + return core.extern_elementwise("", "", [arg0, arg1], { + (core.dtype("fp32"), core.dtype("fp32")): ("__hmf_fmod_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + + +@core.extern +def float_as_int(arg0, _semantic=None): + return core.extern_elementwise("", "", [arg0], { + (core.dtype("fp32"), ): ("__hmf_float_as_int_fp32", core.dtype("int32")), + }, is_pure=True, _semantic=_semantic) + + +@core.extern +def atan2(arg0, arg1, _semantic=None): + if arg0.dtype == core.dtype("bf16") or arg1.dtype == core.dtype("bf16"): + core.static_print("extern livdevice.atan2 for dtype bf16 is unspported for now.") + core.static_assert(False) + return core.extern_elementwise( + "", "", [arg0, arg1], { + (core.dtype("fp16"), core.dtype("fp16")): ("__hmf_atan2_fp16", core.dtype("fp16")), + (core.dtype("fp32"), core.dtype("fp32")): ("__hmf_atan2_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + + +@core.extern +def round(arg0, _semantic=None): + return core.extern_elementwise("", "", [arg0], { + (core.dtype("fp32"), ): ("__hmf_roundf", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + + +@core.builtin +@math._check_dtype(dtypes=["bf16", "fp16", "fp32"]) +@math._add_math_1arg_docstr("acos") +def acos(arg0: core.tensor, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + if arg0.dtype == core.dtype("bf16"): + core.static_print("extern livdevice.acos for dtype bf16 is unspported for now.") + core.static_assert(False) + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp16"), ): ("__hmf_acos_fp16", core.dtype("fp16")), + (core.dtype("fp32"), ): ("__hmf_acos_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + else: + pi = 3.1415926536 + pi_half = 1.5707963268 + sqrt2 = 1.4142135624 + eps = 1e-8 + + # |x| < 0.5, acos(x) = pi/2 - [x + x*x²*(0.1666667 + x²*(0.075 + x²*(0.0446429 + 0.0303810*x²))] + arg0 = _semantic.to_tensor(arg0) + abs_x = math.abs(arg0, _semantic=_semantic) + dtype = arg0.dtype + arg0_2 = _semantic.mul(arg0, arg0, True) + arg0_4 = _semantic.mul(arg0_2, arg0_2, True) + arg0_6 = _semantic.mul(arg0_4, arg0_2, True) + arg0_8 = _semantic.mul(arg0_6, arg0_2, True) + arg0_10 = _semantic.mul(arg0_8, arg0_2, True) + poly = _semantic.add(1.0, _semantic.mul(0.166667, arg0_2, True), True) + poly = _semantic.add(poly, _semantic.mul(0.075, arg0_4, True), True) + poly = _semantic.add(poly, _semantic.mul(0.044643, arg0_6, True), True) + poly = _semantic.add(poly, _semantic.mul(0.030380, arg0_8, True), True) + poly = _semantic.add(poly, _semantic.mul(0.022372, arg0_10, True), True) + acos_center = _semantic.sub(pi_half, _semantic.mul(arg0, poly, True), True) + + # 0.5<|x|<0.9, acos(x) = 2*arctan(t), t=sqrt((1-abs_x)/(1+abs_x)) + numerator_mid = _semantic.sub(1.0, abs_x, True) + denom_mid = _semantic.add(1.0, abs_x, True) + div_mid = _semantic.truediv(numerator_mid, denom_mid) + t_mid = math.sqrt(div_mid, _semantic=_semantic) + t2_mid = _semantic.mul(t_mid, t_mid, True) + t4_mid = _semantic.mul(t2_mid, t2_mid, True) + t6_mid = _semantic.mul(t4_mid, t2_mid, True) + + poly_mid1 = _semantic.mul(0.1065976, t2_mid, True) + poly_mid2 = _semantic.add(-0.1420890, poly_mid1, True) + poly_mid3 = _semantic.mul(poly_mid2, t2_mid, True) + poly_mid4 = _semantic.add(0.1999341, poly_mid3, True) + poly_mid5 = _semantic.mul(poly_mid4, t2_mid, True) + poly_mid6 = _semantic.add(-0.3333310, poly_mid5, True) + poly_mid = _semantic.add(1.0, _semantic.mul(poly_mid6, t2_mid, True), True) + arctan_t = _semantic.mul(t_mid, poly_mid, True) + acos_mid = _semantic.mul(2.0, arctan_t, True) + is_neg_mid = _semantic.less_than(arg0, 0.0) + acos_mid_signed = _semantic.where(is_neg_mid, _semantic.sub(pi, acos_mid, True), acos_mid) + + is_center = _semantic.less_than(abs_x, 0.6) + res_mid_boundary = _semantic.where(is_center, acos_center, acos_mid_signed) + return res_mid_boundary + + +@core.builtin +@math._check_dtype(dtypes=["bf16", "fp16", "fp32"]) +@math._add_math_1arg_docstr("sinh") +def sinh(arg0: core.tensor, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + if arg0.dtype == core.dtype("bf16"): + core.static_print("extern livdevice.sinh for dtype bf16 is unspported for now.") + core.static_assert(False) + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp16"), ): ("__hmf_sinh_fp16", core.dtype("fp16")), + (core.dtype("fp32"), ): ("__hmf_sinh_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + else: + arg0 = _semantic.to_tensor(arg0) + exp0 = core.tensor(_semantic.builder.create_exp(arg0.handle), arg0.type) + exp1 = _semantic.truediv(1.0, exp0) + tmp = _semantic.sub(exp0, exp1, True) + ret = _semantic.truediv(tmp, 2.0) + return ret + + +@core.builtin +@math._check_dtype(dtypes=["bf16", "fp16", "fp32"]) +@math._add_math_1arg_docstr("cosh") +def cosh(arg0: core.tensor, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + if arg0.dtype == core.dtype("bf16"): + core.static_print("extern livdevice.cosh for dtype bf16 is unspported for now.") + core.static_assert(False) + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp16"), ): ("__hmf_cosh_fp16", core.dtype("fp16")), + (core.dtype("fp32"), ): ("__hmf_cosh_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + else: + arg0 = _semantic.to_tensor(arg0) + exp0 = core.tensor(_semantic.builder.create_exp(arg0.handle), arg0.type) + exp1 = _semantic.truediv(1.0, exp0) + tmp = _semantic.add(exp0, exp1, True) + ret = _semantic.truediv(tmp, 2.0) + return ret + + +@core.builtin +@math._check_dtype(dtypes=["bf16", "fp16", "fp32"]) +@math._add_math_1arg_docstr("acosh") +def acosh(arg0: core.tensor, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + if arg0.dtype == core.dtype("bf16"): + core.static_print("extern livdevice.acosh for dtype bf16 is unspported for now.") + core.static_assert(False) + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp16"), ): ("__hmf_acosh_fp16", core.dtype("fp16")), + (core.dtype("fp32"), ): ("__hmf_acosh_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + else: + arg0 = _semantic.to_tensor(arg0) + tmp = _semantic.sub(_semantic.mul(arg0, arg0, True), 1.0, True) + sqrt_res = core.tensor(_semantic.builder.create_sqrt(tmp.handle), tmp.type) + sum_res = _semantic.add(arg0, sqrt_res, True) + return core.tensor(_semantic.builder.create_log(sum_res.handle), sum_res.type) + + +@core.builtin +@math._check_dtype(dtypes=["bf16", "fp16", "fp32"]) +@math._add_math_1arg_docstr("asinh") +def asinh(arg0: core.tensor, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + if arg0.dtype == core.dtype("bf16"): + core.static_print("extern livdevice.asinh for dtype bf16 is unspported for now.") + core.static_assert(False) + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp16"), ): ("__hmf_asinh_fp16", core.dtype("fp16")), + (core.dtype("fp32"), ): ("__hmf_asinh_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + else: + arg0 = _semantic.to_tensor(arg0) + tmp = _semantic.add(_semantic.mul(arg0, arg0, True), 1.0, True) + sqrt_res = core.tensor(_semantic.builder.create_sqrt(tmp.handle), tmp.type) + sum_res = _semantic.add(arg0, sqrt_res, True) + return core.tensor(_semantic.builder.create_log(sum_res.handle), sum_res.type) + + +@core.builtin +@math._check_dtype(dtypes=["bf16", "fp16", "fp32"]) +@math._add_math_1arg_docstr("atanh") +def atanh(arg0: core.tensor, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + if arg0.dtype == core.dtype("bf16"): + core.static_print("extern livdevice.atanh for dtype bf16 is unspported for now.") + core.static_assert(False) + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp16"), ): ("__hmf_atanh_fp16", core.dtype("fp16")), + (core.dtype("fp32"), ): ("__hmf_atanh_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + else: + arg0 = _semantic.to_tensor(arg0) + a = _semantic.add(1.0, arg0, True) + b = _semantic.sub(1.0, arg0, True) + lna = core.tensor(_semantic.builder.create_log(a.handle), a.type) + lnb = core.tensor(_semantic.builder.create_log(b.handle), b.type) + tmp = _semantic.sub(lna, lnb, True) + return _semantic.mul(tmp, 0.5, True) + + +@core.builtin +@math._check_dtype(dtypes=["bf16", "fp16", "fp32"]) +@math._add_math_1arg_docstr("expm1") +def expm1(arg0: core.tensor, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + if arg0.dtype == core.dtype("bf16"): + core.static_print("extern livdevice.expm1 for dtype bf16 is unspported for now.") + core.static_assert(False) + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp16"), ): ("__hmf_expm1_fp16", core.dtype("fp16")), + (core.dtype("fp32"), ): ("__hmf_expm1_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + else: + arg0 = _semantic.to_tensor(arg0) + tmp = core.tensor(_semantic.builder.create_exp(arg0.handle), arg0.type) + return _semantic.sub(tmp, 1, True) + + +@core.builtin +@math._check_dtype(dtypes=["fp16", "fp32"]) +@math._add_math_2arg_docstr("nextafter") +def nextafter(arg0: core.tensor, arg1: core.tensor, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + return core.extern_elementwise( + "", "", [arg0, arg1], { + (core.dtype("fp16"), core.dtype("fp16")): ("__hmf_nextafter_fp16", core.dtype("fp16")), + (core.dtype("fp32"), core.dtype("fp32")): ("__hmf_nextafter_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + else: + x = _semantic.to_tensor(arg0) + y = _semantic.to_tensor(arg1) + dtype_map = {"bf16": core.int16, "fp16": core.int16, "fp32": core.int32} + min_pos_bit = {"bf16": 0x0001, "fp16": 0x0001, "fp32": 0x00000001} + max_neg_bit = {"bf16": 0x8001, "fp16": 0x8001, "fp32": 0x80000001} + int_type = dtype_map[x.type.scalar.name] + x_eq_y = _semantic.equal(x, y) + x_gt_0 = _semantic.greater_than(x, 0) + y_gt_x = _semantic.greater_than(y, x) + next_neg = _semantic.xor_(x_gt_0, y_gt_x) + next_pos = _semantic.not_(next_neg) + + p1 = _semantic.full(x.shape, 1, int_type) + n1 = _semantic.full(x.shape, -1, int_type) + dir_xy = _semantic.where(next_pos, p1, n1) + x_abs = math.abs(x, _semantic=_semantic) + x_is_0 = _semantic.equal(x_abs, 0) + + min_pos = _semantic.full(x.shape, min_pos_bit[x.type.scalar.name], int_type) + max_neg = _semantic.full(x.shape, max_neg_bit[x.type.scalar.name], int_type) + min_pos = _semantic.bitcast(min_pos, x.dtype) + max_neg = _semantic.bitcast(max_neg, x.dtype) + bits_x = _semantic.bitcast(x, int_type) + bits_next = _semantic.add(bits_x, dir_xy, True) + next_val = _semantic.bitcast(bits_next, x.dtype) + + need_min_pos = _semantic.logical_and(x_is_0, next_pos) + need_max_neg = _semantic.logical_and(x_is_0, next_neg) + next_val = _semantic.where(need_min_pos, min_pos, next_val) + next_val = _semantic.where(need_max_neg, max_neg, next_val) + return _semantic.where(x_eq_y, x, next_val) + + +@core.builtin +@math._check_dtype(dtypes=["bf16", "fp16", "fp32"]) +@math._add_math_2arg_docstr("hypot(Euclidean Distance)") +def hypot(arg0: core.tensor, arg1: core.tensor, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + if arg0.dtype == core.dtype("bf16"): + core.static_print("extern livdevice.hypot for dtype bf16 is unspported for now.") + core.static_assert(False) + return core.extern_elementwise( + "", "", [arg0, arg1], { + (core.dtype("fp16"), core.dtype("fp16")): ("__hmf_hypot_fp16", core.dtype("fp16")), + (core.dtype("fp32"), core.dtype("fp32")): ("__hmf_hypot_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + else: + arg0 = _semantic.to_tensor(arg0) + arg1 = _semantic.to_tensor(arg1) + x2 = _semantic.mul(arg0, arg0, True) + y2 = _semantic.mul(arg1, arg1, True) + sum_res = _semantic.add(x2, y2, True) + return core.tensor(_semantic.builder.create_sqrt(sum_res.handle), sum_res.type) + + +# This function is derived from the Cephes Math Library release 2.8: June, 2000 +# https://netlib.org/cephes/ +# Copyright (c) 1984, 1987, 2000 by Stephen L. Moshier +# All rights reserved. +@core.builtin +@math._check_dtype(dtypes=["fp16", "fp32"]) +@math._add_math_2arg_docstr("besseli0 (Modified Bessel function of the first kind, order 0).") +def cyl_bessel_i0(arg0: core.tensor, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + if arg0.dtype == core.dtype("fp16"): + core.static_print("extern livdevice.cyl_bessel_i0 for dtype bf16 is unspported for now.") + core.static_assert(False) + return core.extern_elementwise("", "", [arg0], { + (core.dtype("fp32"), ): ("__hmf_cyl_bessel_i0_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + else: + param1 = [ + -4.41534164647933937950e-18, + +3.33079451882223809783e-17, + -2.43127984654795469359e-16, + +1.71539128555513303061e-15, + -1.16853328779934516808e-14, + +7.67618549860493561688e-14, + -4.85644678311192946090e-13, + +2.95505266312963983461e-12, + -1.72682629144155570723e-11, + +9.67580903537323691224e-11, + -5.18979560163526290666e-10, + +2.65982372468238665035e-09, + -1.30002500998624804212e-08, + +6.04699502254191894932e-08, + -2.67079385394061173391e-07, + +1.11738753912010371815e-06, + -4.41673835845875056359e-06, + +1.64484480707288970893e-05, + -5.75419501008210370398e-05, + +1.88502885095841655729e-04, + -5.76375574538582365885e-04, + +1.63947561694133579842e-03, + -4.32430999505057594430e-03, + +1.05464603945949983183e-02, + -2.37374148058994688156e-02, + +4.93052842396707084878e-02, + -9.49010970480476444210e-02, + +1.71620901522208775349e-01, + -3.04682672343198398683e-01, + +6.76795274409476084995e-01, + ] + param2 = [ + -7.23318048787475395456e-18, + -4.83050448594418207126e-18, + +4.46562142029675999901e-17, + +3.46122286769746109310e-17, + -2.82762398051658348494e-16, + -3.42548561967721913462e-16, + +1.77256013305652638360e-15, + +3.81168066935262242075e-15, + -9.55484669882830764870e-15, + -4.15056934728722208663e-14, + +1.54008621752140982691e-14, + +3.85277838274214270114e-13, + +7.18012445138366623367e-13, + -1.79417853150680611778e-12, + -1.32158118404477131188e-11, + -3.14991652796324136454e-11, + +1.18891471078464383424e-11, + +4.94060238822496958910e-10, + +3.39623202570838634515e-09, + +2.26666899049817806459e-08, + +2.04891858946906374183e-07, + +2.89137052083475648297e-06, + +6.88975834691682398426e-05, + +3.36911647825569408990e-03, + +8.04490411014108831608e-01, + ] + arg0 = _semantic.to_tensor(arg0) + abs_x = core.tensor(_semantic.builder.create_fabs(arg0.handle), arg0.type) + x_a = _semantic.sub(_semantic.mul(abs_x, 0.5, True), 2.0, True) + a_n_2 = 0 + a_n_1 = 0 + a_n = param1[0] + for i in range(1, 30): + a_n_2 = a_n_1 + a_n_1 = a_n + a_n = _semantic.sub(_semantic.mul(x_a, a_n_1, True), a_n_2, True) + a_n = _semantic.add(a_n, param1[i], True) + + f_32 = _semantic.full(abs_x.shape, 32.0, abs_x.type.scalar) + x_b = _semantic.sub(_semantic.fdiv(f_32, abs_x, True), 2.0, True) + b_n_2 = 0 + b_n_1 = 0 + b_n = param2[0] + for i in range(1, 25): + b_n_2 = b_n_1 + b_n_1 = b_n + b_n = _semantic.sub(_semantic.mul(x_b, b_n_1, True), b_n_2, True) + b_n = _semantic.add(b_n, param2[i], True) + + half_exp = _semantic.mul(core.tensor(_semantic.builder.create_exp(abs_x.handle), abs_x.type), 0.5, True) + res_a = _semantic.mul(half_exp, _semantic.sub(a_n, a_n_2, True), True) + res_b = _semantic.fdiv(_semantic.mul(half_exp, _semantic.sub(b_n, b_n_2, True), True), \ + core.tensor(_semantic.builder.create_sqrt(abs_x.handle), abs_x.type), True) + cond = _semantic.less_equal(abs_x, 8.0) + res = _semantic.where(cond, res_a, res_b) + return res + + +@core.extern +@math._check_dtype(dtypes=["fp16", "fp32"]) +def signbit(arg0, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp16"), ): ("__hmf_signbit_fp16", core.dtype("int32")), + (core.dtype("fp32"), ): ("__hmf_signbit_fp32", core.dtype("int32")), + }, is_pure=True, _semantic=_semantic) + else: + arg0_scalar_ty = arg0.type.scalar + if arg0_scalar_ty == core.float32: + int_ty = core.int32 + else: # arg0 type: float16 / bfloat16 + int_ty = core.int16 + + arg0 = _semantic.to_tensor(arg0) + int_tensor = _semantic.bitcast(arg0, int_ty) + if int_ty == core.int32: + shift = 31 + elif int_ty == core.int16: + shift = 15 + + shift = _semantic.full(arg0.shape, shift, int_ty) + sign_bit_tensor = _semantic.lshr(int_tensor, shift) + sign_bit_tensor = _semantic.and_(sign_bit_tensor, _semantic.full(arg0.shape, 1, int_ty)) + return _semantic.equal(sign_bit_tensor, 1) + + +# Note: +# For inputs x very close to ±1 (criterion: 1 - |x| < 1.1e-4), erfinv(x) → ±∞ and the +# inverse error function becomes extremely sensitive to tiny changes in x. The asymptotic +# behavior includes terms like sqrt(-ln(1-|x|)), so tiny relative changes in (1-|x|) map +# to large absolute changes in erfinv, leading to numerical instability and loss of precision, +# resulting in deviations from the reference results. +@core.extern +@math._check_dtype(dtypes=["fp32"]) +def erfinv(arg0, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + return core.extern_elementwise("", "", [arg0], { + (core.dtype("fp32"), ): ("__hmf_erfinv_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + else: + arg0_scalar_ty = arg0.type.scalar + arg0 = _semantic.to_tensor(arg0) + + inv_sqrt_pi_times_2 = _semantic.full(arg0.shape, 1.128379167, arg0_scalar_ty).handle # 2 / sqrt(pi) + coeff_low_numerator = [-0.140543331, 0.914624893, -1.645349621, 0.886226899] + coeff_low_denominator = [0.012229801, -0.329097515, 1.442710462, -2.118377725, 1.0] + coeff_high_numerator = [1.641345311, 3.429567803, -1.624906493, -1.970840454] + coeff_high_denominator = [1.6370678, 3.5438892, 1.0] + + # low cal + arg0_squared = _semantic.builder.create_fmul(arg0.handle, arg0.handle) + numerator_low_range = _semantic.full(arg0.shape, coeff_low_numerator[0], arg0_scalar_ty).handle + for i in range(1, len(coeff_low_numerator)): + numerator_low_range = _semantic.builder.create_fma( + numerator_low_range, arg0_squared, + _semantic.full(arg0.shape, coeff_low_numerator[i], arg0_scalar_ty).handle) + + denominator_low_range = _semantic.full(arg0.shape, coeff_low_denominator[0], arg0_scalar_ty).handle + for i in range(1, len(coeff_low_denominator)): + denominator_low_range = _semantic.builder.create_fma( + denominator_low_range, arg0_squared, + _semantic.full(arg0.shape, coeff_low_denominator[i], arg0_scalar_ty).handle) + + low_res = _semantic.builder.create_fmul( + arg0.handle, _semantic.builder.create_fdiv(numerator_low_range, denominator_low_range)) + + # high cal + arg0_erf_trans = _semantic.builder.create_sqrt( # (log2-log(1-|arg0|))^1/2 + _semantic.builder.create_fmul( + _semantic.full(arg0.shape, -1, arg0_scalar_ty).handle, + _semantic.builder.create_log( + _semantic.builder.create_fdiv( + _semantic.builder.create_fsub( + _semantic.full(arg0.shape, 1, arg0_scalar_ty).handle, + _semantic.builder.create_fabs(arg0.handle)), + _semantic.full(arg0.shape, 2, arg0_scalar_ty).handle)))) + numerator_high_range = _semantic.full(arg0.shape, coeff_high_numerator[0], arg0_scalar_ty).handle + for i in range(1, len(coeff_high_numerator)): + numerator_high_range = _semantic.builder.create_fma( + numerator_high_range, arg0_erf_trans, + _semantic.full(arg0.shape, coeff_high_numerator[i], arg0_scalar_ty).handle) + + denominator_high_range = _semantic.full(arg0.shape, coeff_high_denominator[0], arg0_scalar_ty).handle + for i in range(1, len(coeff_high_denominator)): + denominator_high_range = _semantic.builder.create_fma( + denominator_high_range, arg0_erf_trans, + _semantic.full(arg0.shape, coeff_high_denominator[i], arg0_scalar_ty).handle) + + high_res = _semantic.builder.create_fdiv(numerator_high_range, denominator_high_range) + high_res = _semantic.mul( + _semantic.where(signbit(arg0, _semantic=_semantic), _semantic.full(arg0.shape, -1, arg0_scalar_ty), + _semantic.full(arg0.shape, 1, arg0_scalar_ty)), core.tensor(high_res, arg0.type), True).handle + + for _ in range(2): + low_res = _semantic.builder.create_fsub( + low_res, + _semantic.builder.create_fdiv( + _semantic.builder.create_fsub(_semantic.builder.create_erf(low_res), arg0.handle), + _semantic.builder.create_fmul( + inv_sqrt_pi_times_2, + _semantic.builder.create_exp( + _semantic.builder.create_fmul( + _semantic.full(arg0.shape, -1, arg0_scalar_ty).handle, + _semantic.builder.create_fmul(low_res, low_res)))))) + + high_res = _semantic.builder.create_fsub( + high_res, + _semantic.builder.create_fdiv( + _semantic.builder.create_fsub(_semantic.builder.create_erf(high_res), arg0.handle), + _semantic.builder.create_fmul( + inv_sqrt_pi_times_2, + _semantic.builder.create_exp( + _semantic.builder.create_fmul( + _semantic.full(arg0.shape, -1, arg0_scalar_ty).handle, + _semantic.builder.create_fmul(high_res, high_res)))))) + + arg0_abs = core.tensor(_semantic.builder.create_fabs(arg0.handle), arg0.type) + # Check if |arg0| > 1 + arg0_over = _semantic.greater_than(arg0_abs, _semantic.full(arg0.shape, 1, arg0_scalar_ty)) + nan_tensor = _semantic.full(arg0.shape, float("nan"), arg0_scalar_ty) + # Check if |arg0| = 1 + arg0_equal1 = _semantic.equal(arg0_abs, _semantic.full(arg0.shape, 1, arg0_scalar_ty)) + pos_inf_tensor = _semantic.full(arg0.shape, float("inf"), arg0_scalar_ty) + neg_inf_tensor = _semantic.full(arg0.shape, float("-inf"), arg0_scalar_ty) + inf_res = _semantic.where(signbit(arg0, _semantic=_semantic), neg_inf_tensor, pos_inf_tensor) + # Check if |arg0| >= 0.7 + arg0_high = _semantic.greater_equal(arg0_abs, _semantic.full(arg0.shape, 0.7, arg0_scalar_ty)) + + return _semantic.where( + arg0_equal1, inf_res, + _semantic.where(arg0_over, nan_tensor, + _semantic.where(arg0_high, core.tensor(high_res, arg0.type), core.tensor(low_res, arg0.type)))) + + +# Note: +# The gamma function is implemented using the reflection formula for negative inputs: +# gamma(x) = pi / (sin(pi * x) * gamma(1 - x)). For inputs x close to a negative integer +# (e.g., -1, -2, ... ), criterion: x = -1 ± 0.66e-3, x = -2 ± 1.30e-3, x = -3 ± 2.30e-3, ... +# The denominator sin(pi * x) approaches zero, leading to numerical instability and loss +# of precision. Resulting in deviations from the reference results; +# Similar issues occur near other negative integers. +@core.extern +@math._check_dtype(dtypes=["fp32"]) +def gamma(arg0, _semantic=None): + arg0_scalar_ty = arg0.type.scalar + arg0 = _semantic.to_tensor(arg0) + pi_tensor = _semantic.full(arg0.shape, math_pi, arg0_scalar_ty).handle + sqrt_2pi_tensor = _semantic.full(arg0.shape, 2.506628275, arg0_scalar_ty).handle # sqrt(2*pi) + lanczos_coeff = [ + 676.5203681218851, -1259.1392167224028, 771.32342877765313, -176.61502916214059, 12.507343278686905, + -0.13857109526572012, 9.9843695780195716e-6, 1.5056327351493116e-7 + ] + condition = _semantic.less_than(arg0, 0.5) # 1 - x = x -> x = 0.5 + reflect_arg0 = _semantic.where(condition, _semantic.sub(1, arg0, True), arg0) + + x = _semantic.full(arg0.shape, 0.99999999999980993, arg0_scalar_ty) + for i in range(0, len(lanczos_coeff)): + x = _semantic.add( + x, + _semantic.fdiv(_semantic.full(arg0.shape, lanczos_coeff[i], arg0_scalar_ty), + _semantic.add(reflect_arg0, i, True), True), True) + t = _semantic.add(reflect_arg0, 6.5, True) + + gamma_res = _semantic.builder.create_fmul( + _semantic.builder.create_fmul(sqrt_2pi_tensor, + pow(t, _semantic.sub(reflect_arg0, 0.5, True), _semantic=_semantic).handle), + _semantic.builder.create_fmul( + x.handle, + _semantic.builder.create_exp( + _semantic.builder.create_fmul(t.handle, + _semantic.full(arg0.shape, -1, arg0_scalar_ty).handle)))) + + gamma_res_reflect = _semantic.builder.create_fdiv( + _semantic.builder.create_fdiv(pi_tensor, gamma_res), + _semantic.builder.create_sin(_semantic.builder.create_fmul(pi_tensor, arg0.handle))) + + is_neg_int = _semantic.logical_and(_semantic.equal(math.floor(arg0, _semantic=_semantic), arg0), + _semantic.less_than(arg0, 0)) + pos_inf_tensor = _semantic.full(arg0.shape, float('inf'), arg0_scalar_ty) + neg_inf_tensor = _semantic.full(arg0.shape, float('-inf'), arg0_scalar_ty) + gamma_res_reflect = _semantic.where(is_neg_int, pos_inf_tensor, core.tensor(gamma_res_reflect, arg0.type)) + + res = _semantic.where(condition, gamma_res_reflect, core.tensor(gamma_res, arg0.type)) + is_pos_inf_input = _semantic.equal(arg0, pos_inf_tensor) + is_neg_inf_input = _semantic.equal(arg0, neg_inf_tensor) + + return _semantic.where(is_pos_inf_input, pos_inf_tensor, _semantic.where(is_neg_inf_input, neg_inf_tensor, res)) + + +# Note: +# The lgamma function computes the natural logarithm of the absolute value of the gamma function. +# Since it uses gamma(x) internally, it inherits the same numerical instability near negative integers: +# For inputs x close to a negative integer (e.g., -1, -2, ...), criterion: x = -1 ± 5.75e-5, +# x = -2 ± 1.39e-6, ..., the computation involves log(|pi / (sin(pi * x) * gamma(1 - x))|). +# As sin(pi * x) approaches zero near negative integers, this leads to numerical instability and loss +# of precision, resulting in deviations from the reference results. +# Similar issues occur near other negative integers. +@core.extern +@math._check_dtype(dtypes=["fp32"]) +def lgamma(arg0, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + return core.extern_elementwise("", "", [arg0], { + (core.dtype("fp32"), ): ("__hmf_lgamma_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + else: + arg0_scalar_ty = arg0.type.scalar + arg0 = _semantic.to_tensor(arg0) + + inf_tensor = _semantic.full(arg0.shape, float('inf'), arg0_scalar_ty) + is_inf = _semantic.equal(core.tensor(_semantic.builder.create_fabs(arg0.handle), arg0.type), inf_tensor) + gamma_res = _semantic.builder.create_fabs(gamma(arg0, _semantic=_semantic).handle) + lgamma_res = _semantic.builder.create_log(gamma_res) + + return _semantic.where(is_inf, inf_tensor, core.tensor(lgamma_res, arg0.type)) + + +@core.builtin +@math._check_dtype(dtypes=[ + "fp32", +]) +@math._add_math_1arg_docstr("trunc") +def trunc(arg0: core.tensor, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp16"), ): ("__hmf_trunc_fp16", core.dtype("fp16")), + (core.dtype("fp32"), ): ("__hmf_trunc_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + else: + """ + Truncate the input to the nearest integer toward zero. + + For positive numbers, this is equivalent to floor(x). + For negative numbers, this is equivalent to ceil(x). + + Special cases: + - trunc(±0) returns ±0. + - trunc(±inf) returns ±inf. + - trunc(NaN) returns NaN. + """ + arg0 = _semantic.to_tensor(arg0) + + zero = _semantic.full(arg0.shape, 0.0, arg0.type.scalar) + condition = _semantic.greater_equal(arg0, zero) + + floor_result = core.tensor(_semantic.builder.create_floor(arg0.handle), arg0.type) + ceil_result = core.tensor(_semantic.builder.create_ceil(arg0.handle), arg0.type) + + return _semantic.where(condition, floor_result, ceil_result) + + +@core.builtin +@math._check_dtype(dtypes=[ + "fp32", +]) +@math._add_math_1arg_docstr("nearbyint") +def nearbyint(arg0: core.tensor, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + return core.extern_elementwise("", "", [arg0], { + (core.dtype("fp32"), ): ("__hmf_nearbyint_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + else: + """ + Round argument x to an integer value in floating-point format. + + Uses the current rounding mode (round-to-nearest-even, aka banker's rounding). + """ + arg0 = _semantic.to_tensor(arg0) + + half = _semantic.full(arg0.shape, 0.5, arg0.type.scalar) + + positive_adjust = _semantic.add(arg0, half, True) + negative_adjust = _semantic.sub(arg0, half, True) + + positive_result = core.tensor(_semantic.builder.create_floor(positive_adjust.handle), arg0.type) + negative_result = core.tensor(_semantic.builder.create_ceil(negative_adjust.handle), arg0.type) + + zero = _semantic.full(arg0.shape, 0.0, arg0.type.scalar) + is_positive = _semantic.greater_equal(arg0, zero) + basic_round = _semantic.where(is_positive, positive_result, negative_result) + + # Banker's rounding special treatment: For values exactly in the middle, round to the nearest even number. + fractional = _semantic.sub(arg0, basic_round, True) + abs_fractional = core.tensor(_semantic.builder.create_fabs(fractional.handle), fractional.type) + + is_half = _semantic.equal(abs_fractional, half) + + two = _semantic.full(arg0.shape, 2.0, arg0.type.scalar) + + half_value = math.fdiv(basic_round, two, _semantic=_semantic) + half_floor = core.tensor(_semantic.builder.create_floor(half_value.handle), half_value.type) + double_half = _semantic.mul(half_floor, two, True) + + is_even = _semantic.equal(basic_round, double_half) + + adjustment = _semantic.where(is_positive, _semantic.full(arg0.shape, -1.0, arg0.type.scalar), + _semantic.full(arg0.shape, 1.0, arg0.type.scalar)) + + banker_result = _semantic.where( + is_even, + basic_round, + _semantic.add(basic_round, adjustment, True), + ) + # Final result: Use banker's rounding for cases exactly at 0.5, otherwise use basic rounding. + return _semantic.where(is_half, banker_result, basic_round) + + +@core.builtin +@math._check_dtype(dtypes=[ + "fp32", +]) +@math._add_math_1arg_docstr("arcsine") +def asin(arg0: core.tensor, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp16"), ): ("__hmf_asin_fp16", core.dtype("fp16")), + (core.dtype("fp32"), ): ("__hmf_asin_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + else: + """ + Calculate the principal value of the arc sine of the input argument x. + + Returns result in radians, in the interval [-π/2, +π/2] for x inside [-1, +1]. + Returns NaN for x outside [-1, +1]. + """ + arg0 = _semantic.to_tensor(arg0) + + # asin(x) = π/2 - acos(x) + half_pi = _semantic.full(arg0.shape, 1.5707963267948966, arg0.type.scalar) # π/2 + acos_val = acos(arg0, _semantic=_semantic) + return _semantic.sub(half_pi, acos_val, True) + + +@core.builtin +@math._check_dtype(dtypes=[ + "fp32", +]) +@math._add_math_1arg_docstr("base-10 logarithm") +def log10(arg0: core.tensor, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + return core.extern_elementwise("", "", [arg0], { + (core.dtype("fp32"), ): ("__hmf_log10_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + else: + """ + Calculate the base 10 logarithm of the input argument x. + + Returns NaN for x < 0, -inf for x = 0, and +0 for x = 1. + log10(x) = log(x) / log(10) + """ + arg0 = _semantic.to_tensor(arg0) + + log_val = math.log(arg0, _semantic=_semantic) + log10_const = _semantic.full(arg0.shape, 2.302585092994046, arg0.type.scalar) + + return math.fdiv(log_val, log10_const, _semantic=_semantic) + + +@core.builtin +@math._check_dtype(dtypes=[ + "fp32", +]) +@math._add_math_2arg_docstr("copysign") +def copysign(arg0: core.tensor, arg1: core.tensor, _semantic=None): + if triton_enable_libdevice_simt() and is_compile_on_910_95: + return core.extern_elementwise("", "", [arg0, arg1], { + (core.dtype("fp32"), core.dtype("fp32")): ("__hmf_copysign_fp32", core.dtype("fp32")), + }, is_pure=True, _semantic=_semantic) + else: + """ + Create a floating-point value with the magnitude of x and the sign of y. + """ + x = _semantic.to_tensor(arg0) + y = _semantic.to_tensor(arg1) + + magnitude = core.tensor(_semantic.builder.create_fabs(x.handle), x.type) + + zero = _semantic.full(y.shape, 0.0, y.type.scalar) + one = _semantic.full(y.shape, 1.0, y.type.scalar) + + is_zero = _semantic.equal(y, zero) + y_reciprocal = math.fdiv(one, y, _semantic=_semantic) + is_negative_reciprocal = _semantic.less_than(y_reciprocal, zero) + is_negative_zero = _semantic.and_(is_zero, is_negative_reciprocal) + + is_negative_nonzero = _semantic.less_than(y, zero) + is_negative = _semantic.or_(is_negative_zero, is_negative_nonzero) + + neg_magnitude = _semantic.mul(magnitude, _semantic.full(magnitude.shape, -1.0, magnitude.type.scalar), True) + + return _semantic.where(is_negative, neg_magnitude, magnitude) + + +if get_ascend_arch_from_env() == "Ascend910_9589": + # if we have hardware support + @core.extern + def rint(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__hmf_rint", core.dtype("fp32")), + (core.dtype("fp16"), ): ("__hmf_rint", core.dtype("fp16")), + (core.dtype("bf16"), ): ("__hmf_rint", core.dtype("bf16")), + }, is_pure=True, _semantic=_semantic) +else: + + @core.builtin + @math._check_dtype(dtypes=["fp16", "fp32", "bf16"]) + @math._add_math_1arg_docstr("rint") + def rint(arg0: core.tensor, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + + floor_x = math.floor(arg0, _semantic=_semantic) + fractional = _semantic.sub(arg0, floor_x, True) + + half = _semantic.full(arg0.shape, 0.5, arg0.type.scalar) + eps = _semantic.full(arg0.shape, 1e-8, arg0.type.scalar) + is_half = _semantic.less_than(math.abs(_semantic.sub(fractional, half, True), _semantic=_semantic), eps) + + floor_int = floor_x.to(core.int32, _semantic=_semantic) if hasattr(floor_x, "to") else _semantic.cast( + floor_x, core.int32) + two_i32 = _semantic.full(arg0.shape, 2, core.int32) + is_even = _semantic.equal(_semantic.mod(floor_int, two_i32), _semantic.full(arg0.shape, 0, core.int32)) + + zero = _semantic.full(arg0.shape, 0.0, arg0.type.scalar) + is_pos = _semantic.greater_equal(arg0, zero) + + round_pos = math.floor(_semantic.add(arg0, half, True), _semantic=_semantic) + round_neg = math.ceil(_semantic.sub(arg0, half, True), _semantic=_semantic) + normal_round = _semantic.where(is_pos, round_pos, round_neg) + + half_round = _semantic.where(is_even, floor_x, _semantic.add(floor_x, 1.0, True)) + + return _semantic.where(is_half, half_round, normal_round) diff --git a/third_party/ascend/language/kernels/__init__.py b/third_party/ascend/language/kernels/__init__.py new file mode 100644 index 0000000000..bd828caca2 --- /dev/null +++ b/third_party/ascend/language/kernels/__init__.py @@ -0,0 +1,28 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Optimized Triton kernels for Ascend NPU. + +This module provides high-level, optimized kernels written in Triton for common operations. +These are different from the low-level builtin extensions in `cann.extension`. +""" +__all__ = ["gather_2d_simd"] + +from .gather import gather_2d_simd diff --git a/third_party/ascend/language/kernels/gather.py b/third_party/ascend/language/kernels/gather.py new file mode 100644 index 0000000000..56d1ce0160 --- /dev/null +++ b/third_party/ascend/language/kernels/gather.py @@ -0,0 +1,88 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Gather kernel optimized for Ascend NPU. +""" +__all__ = ["gather_2d_simd"] + +import triton +import triton.language as tl +from triton.language.core import constexpr + + +@triton.jit +def gather_2d_simd(src_ptr, idx_ptr, out_ptr, M: constexpr, N: constexpr, K: constexpr, XBLOCK: constexpr, + XBLOCK_SUB: constexpr): + """ + 2D gather kernel for axis=1 (tail axis) with SIMD-style vectorization. + + This kernel is optimized for Ascend NPU architecture with focus on: + - Vectorized memory access using XBLOCK_SUB + - Efficient global memory (GM) access patterns + - Suitable for cases where N and K are not extremely large + + Args: + src_ptr: [M, N] source tensor in GM (Global Memory) + idx_ptr: [M, K] indices tensor in GM + out_ptr: [M, K] output tensor in GM + M: batch dimension size + N: source dimension size (gather from this dimension) + K: output dimension size (number of indices per batch) + XBLOCK: outer block size for M dimension (for program distribution) + XBLOCK_SUB: inner block size for M dimension (for SIMD vectorization) + + Example: + import torch + import triton + from third_party.ascend.language.kernels import gather_2d_simd + + M, N, K = 128, 256, 64 + src = torch.randn(M, N, device='npu') + indices = torch.randint(0, N, (M, K), dtype=torch.int32, device='npu') + output = torch.empty((M, K), dtype=src.dtype, device='npu') + + grid = (triton.cdiv(M, 32),) + gather_2d_simd[grid](src, indices, output, M, N, K, + XBLOCK=32, XBLOCK_SUB=4) + """ + pid = tl.program_id(0) + m_start = pid * XBLOCK + m_end = min(m_start + XBLOCK, M) + m_base = tl.arange(0, XBLOCK_SUB) + + # Process multiple rows at once using XBLOCK_SUB for vectorization + for m_tile_start in range(m_start, m_end, XBLOCK_SUB): + # M dimension offsets: [XBLOCK_SUB] + m_offs = m_tile_start + m_base + m_mask = m_offs < M + + # Load indices: [XBLOCK_SUB, K] + k_offs = tl.arange(0, K) + idx_tile = tl.load(idx_ptr + m_offs[:, None] * K + k_offs[None, :]) + + # Load source data: [XBLOCK_SUB, N] + n_offs = tl.arange(0, N) + src_tile = tl.load(src_ptr + m_offs[:, None] * N + n_offs[None, :]) + + # Gather operation along axis=1 + gathered_values = tl.gather(src_tile, idx_tile, axis=1) + + # Store results + tl.store(out_ptr + m_offs[:, None] * K + k_offs[None, :], gathered_values, mask=m_mask[:, None]) diff --git a/third_party/ascend/lib/AutoBlockify/AutoBlockify.cpp b/third_party/ascend/lib/AutoBlockify/AutoBlockify.cpp new file mode 100644 index 0000000000..076ad5083b --- /dev/null +++ b/third_party/ascend/lib/AutoBlockify/AutoBlockify.cpp @@ -0,0 +1,363 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "AutoBlockify/AutoBlockify.h" +#include "AutoBlockify/Utils.h" +#include "Dialect/TritonAscend/IR/TritonAscendDialect.h" +#include "incubated/Conversion/UtilsIncubated/Utils.h" + +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Transforms/Passes.h" + +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "auto-blockify" + +using namespace mlir; +using namespace triton; + +PropagateUnrealizedCastDown::PropagateUnrealizedCastDown(MLIRContext *context, + Value logicalBlockId, + Value logicalBlockNum, + int autoBlockifySize) + : OpRewritePattern(context), + logicalBlockId(logicalBlockId), logicalBlockNum(logicalBlockNum), + autoBlockifySize(autoBlockifySize) {} + +LogicalResult +PropagateUnrealizedCastDown::matchAndRewrite(UnrealizedConversionCastOp op, + PatternRewriter &rewriter) const { + if (op.getInputs().size() != 2) + return failure(); + auto funcOp = op->getParentOfType(); + auto input = op.getInputs()[0]; + auto res = op->getResult(0); + SmallPtrSet users(op->user_begin(), op->user_end()); + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "Handling UnrealizedConversionCastOp:\n" << op << "\n"; + os << "Users:\n"; + for (auto *user : users) + os << *user << "\n"; + }); + for (auto *user : users) { + PatternRewriter::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(user); + if (auto uccOp = dyn_cast(user)) { + if (uccOp->getResultTypes()[0] != input.getType()) { + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << *user << "\n"; + }); + return op.emitError("UnrealizedConversionCastOp cannot be resolved\n"); + } + rewriter.replaceOp(user, input); + } else if (auto blockifyLoop = getBlockifyLoop(user)) { + handleBlockifyLoop(blockifyLoop.value(), user, rewriter); + } else if (auto splatOp = dyn_cast(user)) { + rewriteSplat(op, splatOp, rewriter); + } else if (auto expandDimsOp = dyn_cast(user)) { + rewriteExpandDims(op, expandDimsOp, rewriter); + } else if (auto reduceOp = dyn_cast(user)) { + rewriteReduce(op, reduceOp, rewriter); + } else if (auto scanOp = dyn_cast(user)) { + rewriteScan(op, scanOp, rewriter); + } else if (auto loadOp = dyn_cast(user)) { + rewriteLoad(op, loadOp, rewriter); + } else if (auto storeOp = dyn_cast(user)) { + rewriteStore(op, storeOp, rewriter); + } else if (auto atomicRMWOp = dyn_cast(user)) { + rewriteAtomicRMW(op, atomicRMWOp, rewriter); + } else if (auto assertOp = dyn_cast(user)) { + rewriteAssert(op, assertOp, rewriter); + } else if (auto extractSliceOp = dyn_cast(user)) { + rewriteExtractSlice(op, extractSliceOp, rewriter); + } else if (auto insertSliceOp = dyn_cast(user)) { + rewriteInsertSlice(op, insertSliceOp, rewriter); + } else if (auto whileOp = dyn_cast(user)) { + rewriteWhile(op, whileOp, rewriter); + } else if (auto loopOp = dyn_cast(user)) { + rewriteLoop(op, loopOp, rewriter); + } else if (auto yieldOp = dyn_cast(user)) { + rewriteYield(op, yieldOp, rewriter); + } else if (auto conditionOp = dyn_cast(user)) { + rewriteCondition(op, conditionOp, rewriter); + } else if (user->hasTrait() || + isa(user)) { + rewriteGeneraleOp(op, user, rewriter); + } else if (isa(user)) { + auto *newOp = + createBlockifyLoop(user, op, logicalBlockId, logicalBlockNum, + autoBlockifySize, rewriter); + rewriter.setInsertionPoint(newOp); + handleBlockifyLoop(*getBlockifyLoop(newOp), newOp, rewriter); + } else { + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "Unhandled Op\n" << *user << "\n"; + }); + llvm_unreachable("Unhandled operation"); + } + } + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "After successful conversion\n"; + os << funcOp << "\n"; + }); + rewriter.eraseOp(op); + return success(); +} + +AutoBlockifyPass::AutoBlockifyPass(const AutoBlockifyOptions &options) + : AutoBlockifyBase(options) {} + +bool AutoBlockifyPass::checkBlockifiable(Value v) { + if (!checkedValues.insert(v).second) + return true; + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "Checking blockifiable:\n" << v << "\n"; + }); + for (auto &use : v.getUses()) { + auto *user = use.getOwner(); + auto opNum = use.getOperandNumber(); + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "User:\n" << *user << "\n"; + }); + if (isa( + user) || + llvm::any_of(user->getOperandTypes(), isTensorPtrType)) + return false; + if (auto ifOp = dyn_cast(user)) { + user->setAttr(autoBlockifyRegionOpAttr, UnitAttr::get(v.getContext())); + return true; + } else if (auto whileOp = dyn_cast(user)) { + if (!checkBlockifiable(whileOp.getBeforeArguments()[opNum])) + return false; + } else if (auto loopOp = dyn_cast(user)) { + auto regionIterArg = loopOp.getTiedLoopRegionIterArg(&use); + auto loopResult = loopOp.getTiedLoopResult(&use); + if (!regionIterArg || !loopResult) { + user->setAttr(autoBlockifyRegionOpAttr, UnitAttr::get(v.getContext())); + return true; + } + if (!checkBlockifiable(regionIterArg) || !checkBlockifiable(loopResult)) + return false; + } else if (auto conditionOp = dyn_cast(user)) { + auto whileOp = cast(user->getParentOp()); + if (opNum == 0) { + whileOp->setAttr(autoBlockifyRegionOpAttr, + UnitAttr::get(v.getContext())); + return true; + } + if (!checkBlockifiable(whileOp.getAfterArguments()[opNum - 1]) || + !checkBlockifiable(whileOp->getResult(opNum - 1))) + return false; + } else if (auto conditionOp = dyn_cast(user)) { + if (auto loopOp = dyn_cast(user->getParentOp()); + loopOp && !checkBlockifiable(loopOp.getInits()[opNum])) + return false; + } else { + for (auto res : user->getResults()) { + if (!checkBlockifiable(res)) + return false; + } + } + } + return true; +} + +void AutoBlockifyPass::preProcess(triton::FuncOp func) { + IRRewriter rewriter(func.getContext()); + rewriter.setInsertionPointToStart(&func.getBody().front()); + auto loc = rewriter.getUnknownLoc(); + // Get logical block num + auto xNum = + rewriter.create(loc, triton::ProgramIDDim::X); + auto yNum = + rewriter.create(loc, triton::ProgramIDDim::Y); + auto zNum = + rewriter.create(loc, triton::ProgramIDDim::Z); + auto yzNum = rewriter.create(loc, yNum, zNum); + logicalBlockNum = rewriter.create(loc, yzNum, xNum); + + // Get logical block id + auto xDim = + rewriter.create(loc, triton::ProgramIDDim::X); + auto yDim = + rewriter.create(loc, triton::ProgramIDDim::Y); + auto zDim = + rewriter.create(loc, triton::ProgramIDDim::Z); + xDim->setAttr(logicalBlockIdAttr, rewriter.getUnitAttr()); + yDim->setAttr(logicalBlockIdAttr, rewriter.getUnitAttr()); + zDim->setAttr(logicalBlockIdAttr, rewriter.getUnitAttr()); + auto xFlatten = rewriter.create(loc, xDim, yzNum); + auto yFlatten = rewriter.create(loc, yDim, zNum); + logicalBlockId = rewriter.create(loc, xFlatten, yFlatten); + logicalBlockId = rewriter.create(loc, logicalBlockId, zDim); + + // get blockified block id + auto blockifyTensorType = + RankedTensorType::get({autoBlockifySize}, rewriter.getI32Type()); + auto blockfyRange = rewriter.create( + loc, blockifyTensorType, 0, autoBlockifySize); + auto splatedLogicalBlockId = rewriter.create( + loc, blockfyRange.getType(), logicalBlockId); + Value blockifiedId = + rewriter.create(loc, splatedLogicalBlockId, blockfyRange); + + // get mask + auto splatedBlockNum = rewriter.create( + loc, blockfyRange.getType(), logicalBlockNum); + auto upperboundMask = rewriter.create( + loc, arith::CmpIPredicate::slt, blockifiedId, splatedBlockNum); + auto splatedZero = rewriter.create( + loc, DenseElementsAttr::get(blockifyTensorType, + rewriter.getI32IntegerAttr(0))); + auto lowerboundMask = rewriter.create( + loc, arith::CmpIPredicate::sge, blockifiedId, splatedZero); + Value blockifiedIdMask = + rewriter.create(loc, upperboundMask, lowerboundMask); + + blockifiedId = rewriter + .create( + loc, logicalBlockId.getType(), + ValueRange({blockifiedId, blockifiedIdMask})) + ->getResult(0); + + // replace program id to be computed from blockified id + SmallVector toReplace; + func.walk([&](triton::GetProgramIdOp id) { + if (id->hasAttr(logicalBlockIdAttr)) + return; + toReplace.push_back(id); + }); + for (auto id : toReplace) { + rewriter.setInsertionPoint(id); + Value newId; + if (id.getAxis() == triton::ProgramIDDim::X) { + newId = rewriter.create(id.getLoc(), blockifiedId, yzNum); + newId = rewriter.create(id.getLoc(), newId, xNum); + } else if (id.getAxis() == triton::ProgramIDDim::Y) { + newId = rewriter.create(id.getLoc(), blockifiedId, zNum); + newId = rewriter.create(id.getLoc(), newId, yNum); + } else { + newId = rewriter.create(id.getLoc(), blockifiedId, zNum); + } + rewriter.replaceOp(id, newId); + } + + // Create for loop for region ops + func.walk([&](Operation *op) { + if (op->hasAttr(autoBlockifyRegionOpAttr)) { + auto *newOp = createBlockifyLoop( + op, blockifiedId.getDefiningOp(), + logicalBlockId, logicalBlockNum, autoBlockifySize, rewriter); + newOp->removeAttr(autoBlockifyRegionOpAttr); + return WalkResult::skip(); + } + return WalkResult::advance(); + }); +} + +void AutoBlockifyPass::runOnOperation() { + if (autoBlockifySize == 1) + return; + ModuleOp moduleOp = getOperation(); + if (autoBlockifySize <= 0) { + moduleOp->emitWarning("[AutoBlockify V2] AutoBlockifySize cannot be " + "negative integer, skipping."); + return signalPassFailure(); + } + + MLIRContext *ctx = &getContext(); + + moduleOp.walk([&](triton::FuncOp func) { + LogicalResult result = success(); + func.walk([&](triton::GetProgramIdOp id) { + if (!checkBlockifiable(id.getResult())) { + result = failure(); + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + if (failed(result)) { + func->emitWarning("Cannot apply auto blockify"); + return WalkResult::skip(); + } + preProcess(func); + + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "After preprocess:\n" << func << "\n"; + }); + + RewritePatternSet patterns(ctx); + patterns.add( + ctx, logicalBlockId, logicalBlockNum, autoBlockifySize); + + if (failed(applyPatternsGreedily(func, std::move(patterns)))) { + moduleOp->emitError("failed to apply Patterns"); + signalPassFailure(); + return WalkResult::interrupt(); + } + + IRRewriter rewriter(ctx); + func->walk([&](UnrealizedConversionCastOp op) { + rewriter.setInsertionPoint(op); + auto input = op.getInputs()[0]; + auto resType = cast(op->getResultTypes()[0]); + if (auto constantOp = input.getDefiningOp()) { + Attribute val = constantOp.getValue(); + if (auto denseAttr = dyn_cast(val)) + val = denseAttr.getSplatValue(); + rewriter.replaceOpWithNewOp( + op, DenseElementsAttr::get(resType, val)); + } else if (auto tensorType = + dyn_cast(input.getType())) { + input = rewriter.create(input.getLoc(), input, 0); + rewriter.replaceOpWithNewOp(op, resType, input); + } else { + rewriter.replaceOpWithNewOp(op, resType, input); + } + }); + func->setAttr(autoBlockifySizeAttr, + rewriter.getI32IntegerAttr(autoBlockifySize)); + return WalkResult::skip(); + }); + + PassManager pm(&getContext(), moduleOp.getOperationName()); + pm.addPass(createCSEPass()); + pm.addPass(createCanonicalizerPass()); + if (failed(runPipeline(pm, moduleOp))) { + signalPassFailure(); + } +} + +std::unique_ptr> +triton::createAutoBlockifyPass(const AutoBlockifyOptions &options) { + return std::make_unique(options); +} diff --git a/third_party/ascend/lib/AutoBlockify/CMakeLists.txt b/third_party/ascend/lib/AutoBlockify/CMakeLists.txt new file mode 100644 index 0000000000..20ffce4753 --- /dev/null +++ b/third_party/ascend/lib/AutoBlockify/CMakeLists.txt @@ -0,0 +1,22 @@ +add_triton_library(AutoBlockify + AutoBlockify.cpp + RewriteOperation.cpp + Utils.cpp + + DEPENDS + AutoBlockifyPassIncGen + + LINK_LIBS PUBLIC + MLIRArithDialect + MLIRDialectUtils + MLIRIR + MLIRMathDialect + MLIRPass + MLIRTensorDialect + TritonIR + TritonTransforms + TritonAnalysis + MLIRTransforms + MLIRSupport + MLIRSCFTransforms +) diff --git a/third_party/ascend/lib/AutoBlockify/RewriteOperation.cpp b/third_party/ascend/lib/AutoBlockify/RewriteOperation.cpp new file mode 100644 index 0000000000..9b5a4bac5b --- /dev/null +++ b/third_party/ascend/lib/AutoBlockify/RewriteOperation.cpp @@ -0,0 +1,512 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "AutoBlockify/AutoBlockify.h" +#include "AutoBlockify/Utils.h" +#include "incubated/Conversion/UtilsIncubated/Utils.h" + +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "auto-blockify-rewrite-operation" + +using namespace mlir; +using namespace triton; + +void PropagateUnrealizedCastDown::handleBlockifyLoop( + scf::ForOp blockifyLoop, Operation *op, PatternRewriter &rewriter) const { + SmallVector newOperands; + for (auto opr : op->getOperands()) { + auto uccOp = opr.getDefiningOp(); + if (!uccOp) { + newOperands.push_back(opr); + continue; + } + auto input = uccOp.getInputs()[0]; + auto tensorType = cast(input.getType()); + Value newOperand; + if (tensorType.getRank() > 1) { + SmallVector offsets(tensorType.getRank(), + rewriter.getIndexAttr(0)); + SmallVector sizes(1, rewriter.getIndexAttr(1)); + SmallVector strides(tensorType.getRank(), + rewriter.getIndexAttr(1)); + offsets[0] = blockifyLoop.getInductionVar(); + for (auto dim : llvm::drop_begin(tensorType.getShape())) + sizes.push_back(rewriter.getIndexAttr(dim)); + newOperand = rewriter.create( + input.getLoc(), cast(opr.getType()), input, offsets, + sizes, strides); + } else { + newOperand = rewriter.create( + input.getLoc(), input, ValueRange{blockifyLoop.getInductionVar()}); + if (isa(opr.getType())) { + newOperand = rewriter.create( + input.getLoc(), rewriter.getIndexType(), newOperand); + } + } + newOperands.push_back(newOperand); + } + rewriter.modifyOpInPlace(op, [&]() { op->setOperands(newOperands); }); +} + +void PropagateUnrealizedCastDown::rewriteGeneraleOp( + UnrealizedConversionCastOp op, Operation *generalOp, + PatternRewriter &rewriter) const { + auto input = op.getInputs()[0]; + auto mask = op.getInputs()[1]; + auto res = op->getResult(0); + auto inputType = cast(input.getType()); + SmallVector newOperands; + SmallVector newResults; + SmallVector newResultTypes; + + for (auto operand : generalOp->getOperands()) + newOperands.push_back(rewriteValue(operand, op, rewriter)); + for (auto resType : generalOp->getResultTypes()) { + newResultTypes.push_back(getExpandedType(resType, op)); + } + auto *newOp = + rewriter.create(generalOp->getLoc(), generalOp->getName().getIdentifier(), + newOperands, newResultTypes, generalOp->getAttrs()); + replaceValue(newOp, generalOp, mask, rewriter); +} + +void PropagateUnrealizedCastDown::rewriteSplat( + UnrealizedConversionCastOp op, triton::SplatOp splatOp, + PatternRewriter &rewriter) const { + auto input = op.getInputs()[0]; + auto mask = op.getInputs()[1]; + auto resType = cast(splatOp.getResult().getType()); + auto curShape = + llvm::to_vector(cast(input.getType()).getShape()); + auto splatedShape = resType.getShape(); + for (auto dim : splatedShape) { + input = rewriter.create(input.getLoc(), input, + curShape.size()); + curShape.push_back(dim); + input = rewriter.create( + input.getLoc(), + RankedTensorType::get(curShape, getElementTypeOrSelf(input)), input); + } + replaceValue(input.getDefiningOp(), splatOp, mask, rewriter); +} + +void PropagateUnrealizedCastDown::rewriteExpandDims( + UnrealizedConversionCastOp op, triton::ExpandDimsOp expandDimsOp, + PatternRewriter &rewriter) const { + auto input = op.getInputs()[0]; + auto mask = op.getInputs()[1]; + auto newOp = rewriter.create( + expandDimsOp.getLoc(), input, expandDimsOp.getAxis() + 1); + for (auto attr : expandDimsOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + replaceValue(newOp, expandDimsOp, mask, rewriter); +} + +void PropagateUnrealizedCastDown::rewriteReduce( + UnrealizedConversionCastOp op, triton::ReduceOp reduceOp, + PatternRewriter &rewriter) const { + auto mask = op.getInputs()[1]; + auto srcs = llvm::map_to_vector(reduceOp.getSrcs(), [&](Value src) { + return rewriteValue(src, op, rewriter); + }); + auto newOp = rewriter.create(reduceOp.getLoc(), srcs, + reduceOp.getAxis() + 1); + auto &newCombineOp = newOp.getCombineOp(); + rewriter.cloneRegionBefore(reduceOp.getCombineOp(), newCombineOp, + newCombineOp.end()); + for (auto attr : reduceOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + replaceValue(newOp, reduceOp, mask, rewriter); +} + +void PropagateUnrealizedCastDown::rewriteScan(UnrealizedConversionCastOp op, + triton::ScanOp scanOp, + PatternRewriter &rewriter) const { + auto mask = op.getInputs()[1]; + auto srcs = llvm::map_to_vector(scanOp.getSrcs(), [&](Value src) { + return rewriteValue(src, op, rewriter); + }); + auto newOp = rewriter.create( + scanOp.getLoc(), srcs, scanOp.getAxis() + 1, scanOp.getReverse()); + auto &newCombineOp = newOp.getCombineOp(); + rewriter.cloneRegionBefore(scanOp.getCombineOp(), newCombineOp, + newCombineOp.end()); + for (auto attr : scanOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + replaceValue(newOp, scanOp, mask, rewriter); +} + +void PropagateUnrealizedCastDown::rewriteLoad(UnrealizedConversionCastOp op, + triton::LoadOp loadOp, + PatternRewriter &rewriter) const { + auto uccMask = op.getInputs()[1]; + auto ptr = rewriteValue(loadOp.getPtr(), op, rewriter); + auto other = rewriteValue(loadOp.getOther(), op, rewriter); + auto mask = rewriteValue(loadOp.getMask(), op, rewriter); + auto res = loadOp.getResult(); + auto resType = getExpandedType(res.getType(), op); + if (!other) { + other = rewriter.create( + rewriter.getUnknownLoc(), + DenseElementsAttr::get( + resType, rewriter.getZeroAttr(getElementTypeOrSelf(res)))); + } + mask = createMask(mask, uccMask, resType.getShape(), rewriter); + auto boundaryCheck = llvm::map_to_vector(loadOp.getBoundaryCheck(), + [](int32_t idx) { return idx + 1; }); + auto newOp = rewriter.create( + loadOp.getLoc(), ptr, mask, other, boundaryCheck, loadOp.getPadding(), + loadOp.getCache(), loadOp.getEvict(), loadOp.getIsVolatile()); + for (auto attr : loadOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + replaceValue(newOp, loadOp, uccMask, rewriter); +} + +void PropagateUnrealizedCastDown::rewriteStore( + UnrealizedConversionCastOp op, triton::StoreOp storeOp, + PatternRewriter &rewriter) const { + auto uccMask = op.getInputs()[1]; + auto ptr = rewriteValue(storeOp.getPtr(), op, rewriter); + auto value = rewriteValue(storeOp.getValue(), op, rewriter); + auto mask = rewriteValue(storeOp.getMask(), op, rewriter); + auto ptrShape = cast(ptr.getType()).getShape(); + mask = createMask(mask, uccMask, ptrShape, rewriter); + auto boundaryCheck = llvm::map_to_vector(storeOp.getBoundaryCheck(), + [](int32_t idx) { return idx + 1; }); + auto newOp = rewriter.create( + storeOp.getLoc(), ptr, value, mask, boundaryCheck, storeOp.getCache(), + storeOp.getEvict()); + for (auto attr : storeOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + rewriter.replaceOp(storeOp, newOp); +} + +void PropagateUnrealizedCastDown::rewriteAtomicRMW( + UnrealizedConversionCastOp op, triton::AtomicRMWOp atomicRMWOp, + PatternRewriter &rewriter) const { + auto uccMask = op.getInputs()[1]; + auto ptr = rewriteValue(atomicRMWOp.getPtr(), op, rewriter); + auto val = rewriteValue(atomicRMWOp.getVal(), op, rewriter); + auto mask = rewriteValue(atomicRMWOp.getMask(), op, rewriter); + auto resType = getExpandedType(atomicRMWOp.getResult().getType(), op); + mask = createMask(mask, uccMask, resType.getShape(), rewriter); + auto newOp = rewriter.create( + atomicRMWOp.getLoc(), resType, atomicRMWOp.getAtomicRmwOp(), ptr, val, + mask, atomicRMWOp.getSem(), atomicRMWOp.getScope()); + for (auto attr : atomicRMWOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + replaceValue(newOp, atomicRMWOp, uccMask, rewriter); +} + +void PropagateUnrealizedCastDown::rewriteAssert( + UnrealizedConversionCastOp op, triton::AssertOp assertOp, + PatternRewriter &rewriter) const { + auto input = op.getInputs()[0]; + auto mask = op.getInputs()[1]; + auto inputShape = cast(input.getType()).getShape(); + auto conditionType = cast(mask.getType()); + auto oneAttr = rewriter.getIntegerAttr(getElementTypeOrSelf(mask), 1); + auto one = rewriter.create( + mask.getLoc(), DenseElementsAttr::get(conditionType, oneAttr)); + Value condition = rewriter.create(input.getLoc(), mask, one); + condition = createMask(nullptr, condition, inputShape, rewriter); + condition = + rewriter.create(condition.getLoc(), condition, input); + auto newOp = rewriter.create(assertOp.getLoc(), condition, + assertOp.getMessage()); + for (auto attr : assertOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + rewriter.replaceOp(assertOp, newOp); +} + +void PropagateUnrealizedCastDown::rewriteExtractSlice( + UnrealizedConversionCastOp op, tensor::ExtractSliceOp extractSliceOp, + PatternRewriter &rewriter) const { + auto mask = op.getInputs()[1]; + auto src = rewriteValue(extractSliceOp.getSource(), op, rewriter); + auto offsets = llvm::to_vector(extractSliceOp.getMixedOffsets()); + auto sizes = llvm::to_vector(extractSliceOp.getMixedSizes()); + auto strides = llvm::to_vector(extractSliceOp.getMixedStrides()); + auto srcType = cast(src.getType()); + offsets.insert(offsets.begin(), rewriter.getIndexAttr(0)); + sizes.insert(sizes.begin(), rewriter.getIndexAttr(srcType.getShape()[0])); + strides.insert(strides.begin(), rewriter.getIndexAttr(1)); + auto newOp = rewriter.create( + extractSliceOp.getLoc(), src, offsets, sizes, strides); + auto newMask = rewriter.create( + mask.getLoc(), mask, offsets, sizes, strides); + for (auto attr : extractSliceOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + replaceValue(newOp, extractSliceOp, newMask, rewriter); +} + +void PropagateUnrealizedCastDown::rewriteInsertSlice( + UnrealizedConversionCastOp op, tensor::InsertSliceOp insertSliceOp, + PatternRewriter &rewriter) const { + auto mask = op.getInputs()[1]; + auto src = rewriteValue(insertSliceOp.getSource(), op, rewriter); + auto dst = rewriteValue(insertSliceOp.getDest(), op, rewriter); + auto offsets = llvm::to_vector(insertSliceOp.getMixedOffsets()); + auto sizes = llvm::to_vector(insertSliceOp.getMixedSizes()); + auto strides = llvm::to_vector(insertSliceOp.getMixedStrides()); + auto srcType = cast(src.getType()); + offsets.insert(offsets.begin(), rewriter.getIndexAttr(0)); + sizes.insert(sizes.begin(), rewriter.getIndexAttr(srcType.getShape()[0])); + strides.insert(strides.begin(), rewriter.getIndexAttr(1)); + auto newOp = rewriter.create( + insertSliceOp.getLoc(), src, dst, offsets, sizes, strides); + for (auto attr : insertSliceOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + replaceValue(newOp, insertSliceOp, mask, rewriter); +} + +void PropagateUnrealizedCastDown::rewriteWhile( + UnrealizedConversionCastOp op, scf::WhileOp whileOp, + PatternRewriter &rewriter) const { + auto input = op.getInputs()[0]; + auto mask = op.getInputs()[1]; + auto res = op->getResult(0); + SmallVector indices; + SmallVector newInits; + IRMapping mapping; + for (auto [idx, init] : llvm::enumerate(whileOp.getInits())) { + if (init == res) { + indices.push_back(idx); + newInits.push_back(input); + } else { + newInits.push_back(init); + } + } + auto newOp = rewriter.create( + whileOp.getLoc(), whileOp->getResultTypes(), newInits, + [&](OpBuilder &b, Location loc, ValueRange args) { + mapRegionIterArg(mapping, whileOp.getBeforeArguments(), args, indices, + mask, b); + for (auto &bodyOp : *whileOp.getBeforeBody()) + b.clone(bodyOp, mapping); + }, + [&](OpBuilder &b, Location loc, ValueRange args) { + mapRegionIterArg(mapping, whileOp.getAfterArguments(), args, {}, mask, + b); + for (auto &bodyOp : whileOp.getAfterBody()->without_terminator()) + b.clone(bodyOp, mapping); + auto yieldOp = + cast(whileOp.getAfterBody()->getTerminator()); + mapYieldedValue(mapping, yieldOp, indices, op, b); + }); + for (auto attr : whileOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + rewriter.replaceOp(whileOp, newOp); +} + +void PropagateUnrealizedCastDown::rewriteLoop(UnrealizedConversionCastOp op, + LoopLikeOpInterface loopOp, + PatternRewriter &rewriter) const { + auto input = op.getInputs()[0]; + auto mask = op.getInputs()[1]; + auto res = op->getResult(0); + SmallVector indices; + SmallVector newInits; + IRMapping mapping; + for (auto [idx, init] : llvm::enumerate(loopOp.getInits())) { + if (init == res) { + indices.push_back(idx); + newInits.push_back(input); + } else { + newInits.push_back(init); + } + } + LoopLikeOpInterface newOp; + if (auto forOp = dyn_cast(loopOp.getOperation())) { + newOp = rewriter.create( + forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(), + forOp.getStep(), newInits, + [&](OpBuilder &b, Location loc, Value iv, ValueRange args) { + mapping.map(forOp.getInductionVar(), iv); + mapRegionIterArg(mapping, forOp.getRegionIterArgs(), args, indices, + mask, b); + for (auto &bodyOp : forOp.getBody()->without_terminator()) + b.clone(bodyOp, mapping); + auto yieldOp = cast(forOp.getBody()->getTerminator()); + mapYieldedValue(mapping, yieldOp, indices, op, b); + }); + for (auto attr : forOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + } else { + llvm_unreachable("Unhandled loopOp"); + } + replaceValue(newOp, loopOp, mask, rewriter, indices); +} + +void PropagateUnrealizedCastDown::rewriteIf(UnrealizedConversionCastOp &op, + scf::IfOp ifOp, + ArrayRef indices, + PatternRewriter &rewriter) const { + IRMapping mapping; + auto mask = op.getInputs()[1]; + auto thenBlockBuilder = [&](OpBuilder &b, Location loc) { + for (auto &bodyOp : *ifOp.thenBlock()) + b.clone(bodyOp, mapping); + }; + auto elseBlockBuilder = [&](OpBuilder &b, Location loc) { + for (auto &bodyOp : *ifOp.elseBlock()) + b.clone(bodyOp, mapping); + }; + scf::IfOp newOp; + if (ifOp.elseBlock()) { + newOp = rewriter.create(ifOp.getLoc(), ifOp.getCondition(), + thenBlockBuilder, elseBlockBuilder); + } else { + newOp = rewriter.create(ifOp.getLoc(), ifOp.getCondition(), + thenBlockBuilder, nullptr); + } + for (auto attr : ifOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + if (mapping.contains(op)) + op = cast(mapping.lookup(op)); + replaceValue(newOp, ifOp, mask, rewriter, indices); +} + +void PropagateUnrealizedCastDown::rewriteYield( + UnrealizedConversionCastOp &op, scf::YieldOp yieldOp, + PatternRewriter &rewriter) const { + auto input = op.getInputs()[0]; + auto mask = op.getInputs()[1]; + auto res = op->getResult(0); + SmallVector indices; + auto newOperands = llvm::to_vector(yieldOp.getOperands()); + for (auto [idx, opr] : llvm::enumerate(newOperands)) { + if (opr == res) + indices.push_back(idx); + } + if (auto loopOp = dyn_cast(yieldOp->getParentOp())) { + auto uccOp = rewriter.create( + op.getLoc(), res.getType(), ValueRange({input})); + for (auto curIdx : indices) + newOperands[curIdx] = uccOp->getResult(0); + auto newOp = rewriter.create(yieldOp.getLoc(), newOperands); + for (auto attr : yieldOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + rewriter.replaceOp(yieldOp, newOp); + rewriter.setInsertionPoint(loopOp); + for (auto curIdx : indices) { + auto &initArg = loopOp.getInitsMutable()[curIdx]; + auto initVal = initArg.get(); + uccOp = rewriter.create( + initVal.getLoc(), input.getType(), ValueRange({initVal})); + uccOp = rewriter.create( + initVal.getLoc(), initVal.getType(), + ValueRange({uccOp->getResult(0), mask})); + rewriter.modifyOpInPlace(loopOp, + [&]() { initArg.set(uccOp->getResult(0)); }); + } + } else if (auto ifOp = dyn_cast(yieldOp->getParentOp())) { + for (auto curIdx : indices) + newOperands[curIdx] = input; + auto newOp = rewriter.create(yieldOp.getLoc(), newOperands); + for (auto attr : yieldOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + rewriter.replaceOp(yieldOp, newOp); + yieldOp = ifOp.thenYield() == yieldOp ? ifOp.elseYield() : ifOp.thenYield(); + if (yieldOp) { + rewriter.setInsertionPoint(yieldOp); + newOperands = llvm::to_vector(yieldOp.getOperands()); + for (auto curIdx : indices) { + auto uccOp = rewriter.create( + op.getLoc(), input.getType(), ValueRange({newOperands[curIdx]})); + newOperands[curIdx] = uccOp->getResult(0); + } + rewriter.replaceOpWithNewOp(yieldOp, newOperands); + } + rewriter.setInsertionPoint(ifOp); + rewriteIf(op, ifOp, indices, rewriter); + } +} + +void PropagateUnrealizedCastDown::rewriteCondition( + UnrealizedConversionCastOp op, scf::ConditionOp conditionOp, + PatternRewriter &rewriter) const { + auto whileOp = cast(conditionOp->getParentOp()); + auto input = op.getInputs()[0]; + auto mask = op.getInputs()[1]; + auto res = op->getResult(0); + int64_t curIdx = -1; + auto args = llvm::to_vector(conditionOp.getArgs()); + for (auto [idx, opr] : llvm::enumerate(args)) { + if (opr == res) + curIdx = idx; + } + args[curIdx] = input; + auto newOp = rewriter.create( + conditionOp.getLoc(), conditionOp.getCondition(), args); + for (auto attr : conditionOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + rewriter.replaceOp(conditionOp, newOp); + + res = whileOp->getResult(curIdx); + auto oldResType = res.getType(); + auto newResType = getExpandedType(oldResType, op); + rewriter.modifyOpInPlace(whileOp, [&]() { res.setType(newResType); }); + rewriter.setInsertionPointAfter(whileOp); + auto newUccOp = rewriter.create( + res.getLoc(), oldResType, ValueRange({res, mask})); + rewriter.replaceAllUsesExcept(res, newUccOp->getResult(0), newUccOp); + auto arg = whileOp.getAfterArguments()[curIdx]; + auto oldArgType = arg.getType(); + auto newArgType = getExpandedType(oldArgType, op); + rewriter.modifyOpInPlace(whileOp, [&]() { arg.setType(newArgType); }); + rewriter.setInsertionPointToStart(whileOp.getAfterBody()); + newUccOp = rewriter.create( + arg.getLoc(), oldArgType, ValueRange({arg, mask})); + rewriter.replaceAllUsesExcept(arg, newUccOp->getResult(0), newUccOp); +} diff --git a/third_party/ascend/lib/AutoBlockify/Utils.cpp b/third_party/ascend/lib/AutoBlockify/Utils.cpp new file mode 100644 index 0000000000..3c9ea90f76 --- /dev/null +++ b/third_party/ascend/lib/AutoBlockify/Utils.cpp @@ -0,0 +1,211 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "AutoBlockify/Utils.h" +#include "incubated/Conversion/UtilsIncubated/Utils.h" + +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "auto-blockify-utils" + +using namespace mlir; +using namespace triton; + +RankedTensorType getExpandedType(Type type, UnrealizedConversionCastOp op) { + auto target = op.getInputs()[0]; + auto targetType = cast(target.getType()); + SmallVector targetShape{targetType.getShape()[0]}; + if (auto valueType = dyn_cast(type)) { + targetShape.append(valueType.getShape().begin(), + valueType.getShape().end()); + } + return RankedTensorType::get(targetShape, getElementTypeOrSelf(type)); +} + +Value rewriteValue(Value value, UnrealizedConversionCastOp op, + OpBuilder &builder) { + if (value == nullptr) + return nullptr; + if (value == op->getResult(0)) + return op.getInputs()[0]; + return builder + .create( + value.getLoc(), getExpandedType(value.getType(), op), value) + ->getResult(0); +} + +void replaceValue(Operation *newOp, Operation *oldOp, Value newMask, + RewriterBase &rewriter, ArrayRef replaceIndices) { + int64_t idx = 0; + for (auto [res, oldRes] : + llvm::zip_equal(newOp->getResults(), oldOp->getResults())) { + if (replaceIndices.empty() || + llvm::find(replaceIndices, idx) != replaceIndices.end()) { + auto resType = res.getType(); + auto newUccOp = rewriter.create( + newOp->getLoc(), oldRes.getType(), ValueRange({res, newMask})); + rewriter.replaceAllUsesExcept(oldRes, newUccOp->getResult(0), newUccOp); + } else { + rewriter.replaceAllUsesWith(oldRes, res); + } + idx++; + } + rewriter.eraseOp(oldOp); +} + +Value createMask(Value mask, Value uccMask, ArrayRef targetShape, + RewriterBase &rewriter) { + SmallVector curShape{targetShape[0]}; + for (auto [idx, dim] : llvm::drop_begin(llvm::enumerate(targetShape))) { + curShape.push_back(dim); + uccMask = + rewriter.create(uccMask.getLoc(), uccMask, idx); + uccMask = rewriter.create( + uccMask.getLoc(), + RankedTensorType::get(curShape, getElementTypeOrSelf(uccMask)), + uccMask); + } + if (mask) { + mask = rewriter.create(mask.getLoc(), mask, uccMask); + } else { + mask = uccMask; + } + return mask; +} + +void mapRegionIterArg(IRMapping &mapping, ValueRange oldArgs, + ValueRange newArgs, ArrayRef indices, Value mask, + OpBuilder &builder) { + auto newArgIter = newArgs.begin(); + for (auto [idx, oldArg] : llvm::enumerate(oldArgs)) { + if (llvm::find(indices, idx) != indices.end()) { + auto newUccOp = builder.create( + oldArg.getLoc(), oldArg.getType(), ValueRange({*newArgIter, mask})); + mapping.map(oldArg, newUccOp->getResult(0)); + } else { + mapping.map(oldArg, *newArgIter); + } + ++newArgIter; + } +} + +void mapYieldedValue(IRMapping &mapping, scf::YieldOp yieldOp, + ArrayRef indices, UnrealizedConversionCastOp op, + OpBuilder &builder) { + SmallVector newOperands; + for (auto [idx, operand] : llvm::enumerate(yieldOp.getOperands())) { + operand = mapping.lookup(operand); + if (llvm::find(indices, idx) != indices.end()) + newOperands.push_back(rewriteValue(operand, op, builder)); + else + newOperands.push_back(operand); + } + builder.create(yieldOp.getLoc(), newOperands); +} + +Operation *createBlockifyLoop(Operation *targetOp, + UnrealizedConversionCastOp op, + Value logicalBlockId, Value logicalBlockNum, + int autoBlockifySize, RewriterBase &rewriter) { + auto loc = targetOp->getLoc(); + rewriter.setInsertionPoint(targetOp); + auto initVal = + rewriter.create(loc, rewriter.getIndexAttr(0)); + auto stepVal = + rewriter.create(loc, rewriter.getIndexAttr(1)); + auto blockifySizeVal = rewriter.create( + loc, rewriter.getIndexAttr(autoBlockifySize)); + Value upperBound = + rewriter.create(loc, logicalBlockNum, logicalBlockId); + auto i32Zero = + rewriter.create(loc, rewriter.getI32IntegerAttr(0)); + upperBound = rewriter.create(loc, upperBound, i32Zero); + upperBound = rewriter.create(loc, rewriter.getIndexType(), + upperBound); + upperBound = + rewriter.create(loc, upperBound, blockifySizeVal); + SmallVector inits; + if (auto loopOp = dyn_cast(targetOp)) { + inits = llvm::map_to_vector(loopOp.getInits(), + [&rewriter, &op](Value v) -> Value { + return rewriteValue(v, op, rewriter); + }); + } else { + auto resultTypes = + llvm::map_to_vector(targetOp->getResultTypes(), [&op](Type type) { + return getExpandedType(type, op); + }); + inits = + llvm::map_to_vector(resultTypes, [&rewriter, &loc](Type type) -> Value { + auto tensorType = cast(type); + return rewriter.create(loc, tensorType.getShape(), + tensorType.getElementType()); + }); + } + auto mask = op.getInputs()[1]; + Operation *newOp; + auto blockifyLoop = rewriter.create( + loc, initVal, upperBound, stepVal, inits, + [&](OpBuilder &b, Location loc, Value iv, ValueRange args) { + newOp = b.clone(*targetOp); + + SmallVector newResults; + for (auto [arg, res] : llvm::zip_equal(args, newOp->getResults())) { + auto tensorType = cast(arg.getType()); + auto rank = tensorType.getRank(); + Value newRes; + if (rank > 1) { + SmallVector offsets(tensorType.getRank(), + b.getIndexAttr(0)); + SmallVector sizes(1, b.getIndexAttr(1)); + SmallVector strides(tensorType.getRank(), + b.getIndexAttr(1)); + offsets[0] = iv; + for (auto dim : llvm::drop_begin(tensorType.getShape())) + sizes.push_back(b.getIndexAttr(dim)); + newRes = b.create(loc, res, arg, offsets, + sizes, strides); + } else { + newRes = b.create(loc, res, arg, ValueRange{iv}); + } + newResults.push_back(newRes); + } + b.create(loc, newResults); + }); + + replaceValue(blockifyLoop, targetOp, mask, rewriter); + blockifyLoop->setAttr(autoBlockifyLoopAttr, rewriter.getUnitAttr()); + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "After creating blockify loop:\n" << blockifyLoop << "\n"; + }); + return newOp; +} + +std::optional getBlockifyLoop(Operation *op) { + while (auto forOp = op->getParentOfType()) { + if (forOp->hasAttr(autoBlockifyLoopAttr)) + return forOp; + op = forOp; + } + return std::nullopt; +} diff --git a/third_party/ascend/lib/CMakeLists.txt b/third_party/ascend/lib/CMakeLists.txt new file mode 100644 index 0000000000..124c9b4c1d --- /dev/null +++ b/third_party/ascend/lib/CMakeLists.txt @@ -0,0 +1,11 @@ +add_subdirectory(AutoBlockify) +add_subdirectory(Dialect) +add_subdirectory(TritonToAnnotation) +add_subdirectory(TritonToHFusion) +add_subdirectory(TritonToHIVM) +add_subdirectory(TritonToGraph) +add_subdirectory(TritonToLLVM) +add_subdirectory(CommonIRToHIVM) +add_subdirectory(DynamicCVPipeline) +# todo: this code will be removed in version 530. +add_subdirectory(TritonAffinityOpt) diff --git a/third_party/ascend/lib/CommonIRToHIVM/CMakeLists.txt b/third_party/ascend/lib/CommonIRToHIVM/CMakeLists.txt new file mode 100644 index 0000000000..333e9797ab --- /dev/null +++ b/third_party/ascend/lib/CommonIRToHIVM/CMakeLists.txt @@ -0,0 +1,20 @@ +add_triton_library(CommonIRToHIVM + CommonIRToHIVM.cpp + ReduceSumStrengthPass.cpp + + DEPENDS + CommonIRToHIVMConversionPassIncGen + + LINK_LIBS + CommonIRIR + BiShengIRHIVMDialect + MLIRIR + MLIRPass + MLIRTransforms + MLIRSupport + TritonIR + MLIRMemRefDialect + MLIRTensorDialect + MLIRBufferizationDialect + MLIRArithDialect +) diff --git a/third_party/ascend/lib/CommonIRToHIVM/CommonIRToHIVM.cpp b/third_party/ascend/lib/CommonIRToHIVM/CommonIRToHIVM.cpp new file mode 100644 index 0000000000..902e642e10 --- /dev/null +++ b/third_party/ascend/lib/CommonIRToHIVM/CommonIRToHIVM.cpp @@ -0,0 +1,899 @@ +/* + * Copyright 2018-2020 Philippe Tillet + * Copyright 2020-2022 OpenAI + * Copyright 2025- FlagOS Contributors + * + * Permission is hereby granted, free of charge, to any person obtaining + * a copy of this software and associated documentation files + * (the "Software"), to deal in the Software without restriction, + * including without limitation the rights to use, copy, modify, merge, + * publish, distribute, sublicense, and/or sell copies of the Software, + * and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be + * included in all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, + * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF + * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. + * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY + * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, + * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE + * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ + +//===----------------------------------------------------------------------===// +// CommonIRToHIVM — Lowers CommonIR dialect ops to HIVM (Ascend NPU) dialect +// ops. +// +// Strategy: greedy rewrite runs patterns iteratively to fixed point. +// Step 1: tile.alloc → memref.alloc (produces memref with +// #hivm.address_space) Step 2: tile.to_tensor → RAUW: replace all uses of +// result with src operand, +// then erase. The src is now memref (from Step 1), consumers get +// memref instead of tensor, bridged by UnrealizedConversionCast. +// Step 3: tile.copy → hivm.copy (tile.* src, on-chip DMA) +// or memref.copy (!tt.ptr src, GM↔local DMA; the ptr is +// bridged to memref via UnrealizedConversionCast) +// tile.store_tensor → hivm.copy. +// Step 4: tile.load/store → hivm.load/store +// Step 5: sync ops → hivm sync ops +//===----------------------------------------------------------------------===// + +#include "ascend/include/CommonIRToHIVM/Passes.h" + +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "mlir-ext/Dialect/CommonIR/IR/CommonIRDialect.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/BuiltinDialect.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Types.h" +#include "llvm/Support/LogicalResult.h" + +using namespace mlir; +namespace tile = mlir::triton::tile; +namespace hivm = mlir::hivm; + +namespace mlir { +namespace triton { +#define GEN_PASS_DEF_COMMONIRTOHIVM +#include "ascend/include/CommonIRToHIVM/Passes.h.inc" +} // namespace triton +} // namespace mlir + +// ============================================================================= +// Common Helpers +// ============================================================================= + +static bool isOneToOne(UnrealizedConversionCastOp op) { + return op.getInputs().size() == 1 && op->getNumResults() == 1; +} + +/// Return true if \p op is an annotation.mark operation. +static bool isAnnotationMark(Operation *op) { + return op->getName().getStringRef() == "annotation.mark"; +} + +static Value castDefaultMemrefToGM(Value value, Location loc, + PatternRewriter &rewriter) { + auto memrefTy = dyn_cast(value.getType()); + if (!memrefTy || memrefTy.getMemorySpace()) + return value; + + auto gmSpace = hivm::AddressSpaceAttr::get(rewriter.getContext(), + hivm::AddressSpace::GM); + auto gmTy = MemRefType::get(memrefTy.getShape(), memrefTy.getElementType(), + memrefTy.getLayout(), gmSpace); + return rewriter.create(loc, gmTy, value); +} + +static hivm::AddressSpace mapMemSpaceToHIVM(tile::MemorySpace tileSpace) { + switch (tileSpace) { + case tile::MemorySpace::GM: + return hivm::AddressSpace::GM; + case tile::MemorySpace::L1: + return hivm::AddressSpace::L1; + case tile::MemorySpace::L0A: + return hivm::AddressSpace::L0A; + case tile::MemorySpace::L0B: + return hivm::AddressSpace::L0B; + case tile::MemorySpace::L0C: + return hivm::AddressSpace::L0C; + case tile::MemorySpace::UB: + return hivm::AddressSpace::UB; + // Generic GPU memory spaces (Global/Shared/Local/Register) have no HIVM + // counterpart and should never reach this lowering pass. + case tile::MemorySpace::Global: + case tile::MemorySpace::Shared: + case tile::MemorySpace::Local: + case tile::MemorySpace::Register: + llvm_unreachable( + "unsupported generic memory space in CommonIR-to-HIVM lowering"); + } + llvm_unreachable("unknown CommonIR memory space"); +} + +static MemRefType convertBufToMemRef(tile::BufType bufTy) { + auto shape = bufTy.getShape(); + auto elemTy = bufTy.getElementType(); + auto space = bufTy.getMemorySpace(); + auto *ctx = elemTy.getContext(); + return MemRefType::get( + llvm::SmallVector(shape), elemTy, MemRefLayoutAttrInterface{}, + hivm::AddressSpaceAttr::get(ctx, mapMemSpaceToHIVM(space))); +} + +static MemRefType convertTensorToMemRef(tile::TensorType tensorTy) { + auto shape = tensorTy.getShape(); + auto elemTy = tensorTy.getElementType(); + auto space = tensorTy.getMemorySpace(); + auto *ctx = elemTy.getContext(); + return MemRefType::get( + llvm::SmallVector(shape), elemTy, MemRefLayoutAttrInterface{}, + hivm::AddressSpaceAttr::get(ctx, mapMemSpaceToHIVM(space))); +} + +/// Get the underlying memref value, converting tile/builtin tensor types to +/// memref. +static Value getAsMemRef(Value val, PatternRewriter &rewriter) { + if (isa(val.getType())) + return val; + // Look through UnrealizedConversionCast + if (auto cast = val.getDefiningOp()) { + if (cast->getNumResults() == 1 && cast->getNumOperands() == 1) { + auto inner = cast->getOperand(0); + if (isa(inner.getType())) + return inner; + } + } + // Convert tile/builtin types to memref + Type targetTy; + if (auto bufTy = dyn_cast(val.getType())) + targetTy = convertBufToMemRef(bufTy); + else if (auto tensorTy = dyn_cast(val.getType())) + targetTy = convertTensorToMemRef(tensorTy); + else if (auto rankedTy = dyn_cast(val.getType())) { + auto elemTy = rankedTy.getElementType(); + if (auto ptrElemTy = dyn_cast(elemTy)) { + auto gmSpace = hivm::AddressSpaceAttr::get(val.getType().getContext(), + hivm::AddressSpace::GM); + targetTy = + MemRefType::get(rankedTy.getShape(), ptrElemTy.getPointeeType(), + MemRefLayoutAttrInterface{}, gmSpace); + } else { + targetTy = MemRefType::get(rankedTy.getShape(), elemTy); + } + } else if (auto ptrTy = dyn_cast(val.getType())) { + // tt.ptr> → memref<...> + auto pointeeTy = ptrTy.getPointeeType(); + auto gmSpace = hivm::AddressSpaceAttr::get(val.getType().getContext(), + hivm::AddressSpace::GM); + if (auto rankedTy = dyn_cast(pointeeTy)) + targetTy = MemRefType::get(rankedTy.getShape(), rankedTy.getElementType(), + MemRefLayoutAttrInterface{}, gmSpace); + else + targetTy = MemRefType::get({ShapedType::kDynamic}, pointeeTy, + MemRefLayoutAttrInterface{}, gmSpace); + } else + return val; + return rewriter + .create(val.getLoc(), targetTy, val) + ->getResult(0); +} + +static bool isTensorOfPointer(Type ty) { + auto rankedTy = dyn_cast(ty); + return rankedTy && isa(rankedTy.getElementType()); +} + +static void lowerScatteredLoad(tile::CopyOp op, Value ptr, Value buf, + PatternRewriter &rewriter) { + Location loc = op.getLoc(); + Value loaded = rewriter.create( + loc, ptr, Value(), Value(), ArrayRef{}, + std::optional(), triton::CacheModifier::NONE, + triton::EvictionPolicy::NORMAL, false); + + auto ty = cast(loaded.getType()); +#ifndef __LLVM_MAJOR_VERSION_22_COMPATIBLE__ + Value mem = rewriter.create( + loc, MemRefType::get(ty.getShape(), ty.getElementType()), loaded); +#else + Value mem = rewriter.create( + loc, MemRefType::get(ty.getShape(), ty.getElementType()), loaded); +#endif + + rewriter.create(loc, TypeRange{}, mem, + getAsMemRef(buf, rewriter)); + rewriter.eraseOp(op); +} + +static void lowerScatteredStore(tile::CopyOp op, Value buf, Value ptr, + PatternRewriter &rewriter) { + Location loc = op.getLoc(); + Value mem = getAsMemRef(buf, rewriter); + auto m = cast(mem.getType()); + + Value t = rewriter.create( + loc, RankedTensorType::get(m.getShape(), m.getElementType()), mem, true, + false); + + rewriter.replaceOpWithNewOp( + op, ptr, t, Value(), ArrayRef{}, triton::CacheModifier::NONE, + triton::EvictionPolicy::NORMAL); +} + +static bool isTritonPointerLike(Type ty) { + if (isa(ty)) + return true; + if (auto rankedTy = dyn_cast(ty)) + return isa(rankedTy.getElementType()); + return false; +} + +static hivm::PIPE mapPipe(int64_t tilePipe) { + switch (static_cast(tilePipe)) { + case tile::Pipe::PIPE_M: + return hivm::PIPE::PIPE_M; + case tile::Pipe::PIPE_V: + return hivm::PIPE::PIPE_V; + case tile::Pipe::PIPE_MTE1: + return hivm::PIPE::PIPE_MTE1; + case tile::Pipe::PIPE_MTE2: + return hivm::PIPE::PIPE_MTE2; + case tile::Pipe::PIPE_MTE3: + return hivm::PIPE::PIPE_MTE3; + case tile::Pipe::PIPE_FIX: + return hivm::PIPE::PIPE_FIX; + case tile::Pipe::PIPE_S: + return hivm::PIPE::PIPE_S; + } + llvm_unreachable("unknown CommonIR pipe"); +} + +static hivm::EVENT mapEvent(int64_t tileEvent) { + return static_cast(tileEvent); +} + +// ============================================================================= +// Step 1: tile.alloc → memref.alloc + #hivm.address_space +// ============================================================================= +struct TileAllocToMemRef : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(tile::AllocOp op, + PatternRewriter &rewriter) const final { + auto bufTy = op.getResult().getType(); + rewriter.replaceOpWithNewOp(op, convertBufToMemRef(bufTy)); + return success(); + } +}; + +// ============================================================================= +// step 1.5 +// ============================================================================= + +struct TileSubviewToMemrefSubview : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(tile::SubViewOp op, + PatternRewriter &rewriter) const override { + // 1. Convert the source buffer and deduce the result type + Value sourceMemRef = getAsMemRef(op.getOperand(0), rewriter); +#ifndef __LLVM_MAJOR_VERSION_22_COMPATIBLE__ + auto sourceMemRefType = sourceMemRef.getType().cast(); +#else + auto sourceMemRefType = mlir::cast(sourceMemRef.getType()); +#endif + + // 2. Map dynamic offsets directly to OpFoldResult + SmallVector mixedOffsets = llvm::to_vector<4>(llvm::map_range( + op.getOffsets(), [](Value v) { return OpFoldResult(v); })); + + // 3. Unpack static sizes and strides into OpFoldResult arrays + SmallVector mixedSizes = getAsOpFoldResult(op.getSizesAttr()); + SmallVector mixedStrides = + getAsOpFoldResult(op.getStridesAttr()); + + auto targetShape = op.getType().getShape(); +#ifndef __LLVM_MAJOR_VERSION_22_COMPATIBLE__ + auto inferredType = memref::SubViewOp::inferRankReducedResultType( + targetShape, sourceMemRefType, mixedOffsets, + mixedSizes, mixedStrides) + .cast(); +#else + auto inferredType = + mlir::cast(memref::SubViewOp::inferRankReducedResultType( + targetShape, sourceMemRefType, mixedOffsets, mixedSizes, + mixedStrides)); +#endif + + // 4. Replace the old op with the standard memref.subview + auto memrefSubview = rewriter.create( + op.getLoc(), inferredType, sourceMemRef, mixedOffsets, mixedSizes, + mixedStrides); + + rewriter.replaceOp(op, memrefSubview.getResult()); + + return success(); + } +}; + +// ============================================================================= +// Step 2: tile.to_tensor → UnrealizedConversionCast (memref → tensor) +// ============================================================================= +struct TileToTensorEliminate : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(tile::ToTensorOp op, + PatternRewriter &rewriter) const final { + // The src is now memref (from Step 1). The result type is tensor<>. + // Bridge the gap with UnrealizedConversionCast: memref → tensor. + Value src = op.getOperand(); + auto resultTy = op.getResult().getType(); + if (op->getAttr("writable")) { + // The caller (tle.dsa.to_tensor defaults to writable=True) expects an + // in-place view: custom ops use such a tensor as a DPS out and write + // the results back into the source buffer. Lower it to a writable + // bufferization.to_tensor; a plain cast would make One-Shot + // Bufferization copy the out into a fresh UB allocation that + // hivm-plan-memory cannot plan, which ends up as a `call @malloc` in + // the device code (AICore link failure) and a broken dataflow. + rewriter.replaceOpWithNewOp( + op, cast(resultTy), src, + /*restrict=*/true, /*writable=*/true); + return success(); + } + if (src.getType() != resultTy) { + auto cast = rewriter.create(op.getLoc(), + resultTy, src); + rewriter.replaceOp(op, cast->getResult(0)); + } else { + rewriter.replaceOp(op, src); + } + return success(); + } +}; + +// ============================================================================= +// Step 3: tile.copy → hivm.copy / memref.copy +// - tile.* source → hivm.copy (on-chip buffer DMA) +// - !tt.ptr source → memref.copy (GM→local DMA; later passes lower it) +// ============================================================================= +struct TileCopyToHIVM : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(tile::CopyOp op, + PatternRewriter &rewriter) const final { + Value src = op.getOperand(0); + Value dst = op.getOperand(1); + + if (isTensorOfPointer(src.getType())) { + lowerScatteredLoad(op, src, dst, rewriter); + return success(); + } + if (isTensorOfPointer(dst.getType())) { + lowerScatteredStore(op, src, dst, rewriter); + return success(); + } + + Value srcMem = getAsMemRef(src, rewriter); + Value dstMem = getAsMemRef(dst, rewriter); + if (isTritonPointerLike(src.getType())) { + rewriter.replaceOpWithNewOp(op, srcMem, dstMem); + return success(); + } + rewriter.create(op.getLoc(), TypeRange{}, srcMem, dstMem); + rewriter.eraseOp(op); + return success(); + } +}; + +struct TileStoreTensorToHIVM : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(tile::StoreTensorOp op, + PatternRewriter &rewriter) const final { + Value dstMem = getAsMemRef(op.getOperand(1), rewriter); + Value srcMem = op.getOperand(0); + if (!isa(srcMem.getType())) { + auto tensorTy = dyn_cast(srcMem.getType()); + if (!tensorTy) + return failure(); + auto memrefTy = + MemRefType::get(tensorTy.getShape(), tensorTy.getElementType()); +#ifndef __LLVM_MAJOR_VERSION_22_COMPATIBLE__ + srcMem = rewriter.create(op.getLoc(), memrefTy, + srcMem); +#else + srcMem = rewriter.create(op.getLoc(), memrefTy, + srcMem); +#endif + } + + rewriter.create(op.getLoc(), TypeRange{}, srcMem, dstMem); + rewriter.eraseOp(op); + return success(); + } +}; + +// ============================================================================= +// Step 4: tile.load → hivm.load +// ============================================================================= +struct TileLoadToHIVM : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(tile::LoadOp op, + PatternRewriter &rewriter) const final { + auto resultTy = op.getResult().getType(); + if (auto t = dyn_cast(resultTy)) { + auto memrefTy = convertTensorToMemRef(t); + rewriter.replaceOpWithNewOp(op, memrefTy, op.getSrc()); + return success(); + } + return failure(); + } +}; + +// ============================================================================= +// Step 4: tile.store → hivm.store +// ============================================================================= +struct TileStoreToHIVM : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(tile::StoreOp op, + PatternRewriter &rewriter) const final { + Value src = getAsMemRef(op.getOperand(0), rewriter); + rewriter.replaceOpWithNewOp(op, Type(), src, + op.getOperand(1)); + return success(); + } +}; + +// ============================================================================= +// Step 5: Sync ops +// ============================================================================= +struct TileSetFlagToHIVM : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(tile::SetFlagOp op, + PatternRewriter &rewriter) const final { + auto *ctx = op->getContext(); + rewriter.replaceOpWithNewOp( + op, + hivm::PipeAttr::get(ctx, + mapPipe(static_cast(op.getProducer()))), + hivm::PipeAttr::get(ctx, + mapPipe(static_cast(op.getConsumer()))), + hivm::EventAttr::get(ctx, + mapEvent(static_cast(op.getEvent()))), + Value()); + return success(); + } +}; + +struct TileWaitFlagToHIVM : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(tile::WaitFlagOp op, + PatternRewriter &rewriter) const final { + auto *ctx = op->getContext(); + rewriter.replaceOpWithNewOp( + op, + hivm::PipeAttr::get(ctx, + mapPipe(static_cast(op.getProducer()))), + hivm::PipeAttr::get(ctx, + mapPipe(static_cast(op.getConsumer()))), + hivm::EventAttr::get(ctx, + mapEvent(static_cast(op.getEvent()))), + Value()); + return success(); + } +}; + +struct TilePipeBarrierToHIVM : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(tile::PipeBarrierOp op, + PatternRewriter &rewriter) const final { + auto *ctx = op->getContext(); + rewriter.replaceOpWithNewOp( + op, + hivm::PipeAttr::get(ctx, mapPipe(static_cast(op.getPipe())))); + return success(); + } +}; + +struct TileCubeWaitToHIVM : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(tile::CubeWaitOp op, + PatternRewriter &rewriter) const final { + auto *ctx = op->getContext(); + rewriter.replaceOpWithNewOp( + op, hivm::TCoreTypeAttr::get(ctx, hivm::TCoreType::VECTOR), + hivm::PipeAttr::get(ctx, hivm::PIPE::PIPE_FIX), + hivm::PipeAttr::get(ctx, hivm::PIPE::PIPE_MTE3), + OpFoldResult(rewriter.getIndexAttr(0))); + return success(); + } +}; + +struct TileGmOffsetToHIVM : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(tile::GmOffsetOp op, + PatternRewriter &rewriter) const final { + auto loc = op.getLoc(); + Value result = op.getBase(); + auto indices = op.getIndices(); + auto strides = op.getStrides(); + for (size_t i = 0; i < indices.size(); ++i) { + Value off = rewriter.create(loc, indices[i], strides[i]); + result = rewriter.create(loc, result, off); + } + rewriter.replaceOp(op, result); + return success(); + } +}; + +//===----------------------------------------------------------------------===// +// Pattern A: !tt.ptr> -> memref<> +//===----------------------------------------------------------------------===// +struct LowerPtrTensorCastToMemref + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(UnrealizedConversionCastOp op, + PatternRewriter &rewriter) const override { + if (!isOneToOne(op)) + return failure(); + Value in = op.getInputs().front(); + auto ptrTy = dyn_cast(in.getType()); + if (!ptrTy) + return failure(); + auto tensorPointee = dyn_cast(ptrTy.getPointeeType()); + if (!tensorPointee) + return failure(); + auto outMemref = dyn_cast(op->getResult(0).getType()); + if (!outMemref) + return failure(); + if (outMemref.getShape() != tensorPointee.getShape() || + outMemref.getElementType() != tensorPointee.getElementType()) + return failure(); + + Location loc = op.getLoc(); + // %t = tt.load %ptr : !tt.ptr> (tensor-pointer overload) + Value loaded = rewriter.create( + loc, in, /*boundaryCheck=*/ArrayRef{}, + /*padding=*/std::optional(), + triton::CacheModifier::NONE, triton::EvictionPolicy::NORMAL, + /*isVolatile=*/false); + // %m = bufferization.to_memref %t +#ifndef __LLVM_MAJOR_VERSION_22_COMPATIBLE__ + Value asMemref = + rewriter.create(loc, outMemref, loaded); +#else + Value asMemref = + rewriter.create(loc, outMemref, loaded); +#endif + rewriter.replaceOp(op, asMemref); + return success(); + } +}; + +//===----------------------------------------------------------------------===// +// Pattern B: memref<..., #space> -> tensor<...> +//===----------------------------------------------------------------------===// +struct LowerMemrefCastToTensor + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(UnrealizedConversionCastOp op, + PatternRewriter &rewriter) const override { + if (!isOneToOne(op)) + return failure(); + Value in = op.getInputs().front(); + auto memrefTy = dyn_cast(in.getType()); + if (!memrefTy) + return failure(); + auto tensorOut = dyn_cast(op->getResult(0).getType()); + if (!tensorOut) + return failure(); + if (memrefTy.getShape() != tensorOut.getShape() || + memrefTy.getElementType() != tensorOut.getElementType()) + return failure(); + + Location loc = op.getLoc(); + // %t = bufferization.to_tensor %m restrict + Value asTensor = rewriter.create( + loc, tensorOut, in, /*restrict=*/true, /*writable=*/false); + rewriter.replaceOp(op, asTensor); + return success(); + } +}; + +//===----------------------------------------------------------------------===// +// Pattern C (utility): simplify chains of one-to-one unrealized casts whose +// outer source type equals the outer result type. This catches cases like +// %a = unrealized_conversion_cast %x : memref -> tensor +// %b = unrealized_conversion_cast %a : tensor -> memref +// which arise after the linalg pass partially converts surrounding ops. +//===----------------------------------------------------------------------===// +struct FoldRoundtripCast : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(UnrealizedConversionCastOp op, + PatternRewriter &rewriter) const override { + if (!isOneToOne(op)) + return failure(); + Value in = op.getInputs().front(); + auto prev = in.getDefiningOp(); + if (!prev || !isOneToOne(prev)) + return failure(); + Value origin = prev.getInputs().front(); + if (origin.getType() != op->getResult(0).getType()) + return failure(); + rewriter.replaceOp(op, origin); + return success(); + } +}; + +//===----------------------------------------------------------------------===// +// Rewrite pattern: fold staging alloc + copy1 + to_tensor + copy2 +//===----------------------------------------------------------------------===// +struct FoldStagingCopyPattern : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(memref::CopyOp copyOp, + PatternRewriter &rewriter) const override { + Value stage = copyOp.getSource(); + auto stageCast = stage.getDefiningOp(); + if (stageCast) + stage = stageCast.getSource(); + Value dst = copyOp.getTarget(); + + // ① Destination must have an explicit memory space (on-chip buffer). + auto dstType = dyn_cast(dst.getType()); + if (!dstType || !dstType.getMemorySpace()) + return failure(); + + // ② Source must be a staging memref.alloc in default/GM address space. + auto stageAlloc = stage.getDefiningOp(); + if (!stageAlloc) + return failure(); + auto stageType = stageAlloc.getType(); + if (auto stageSpace = stageType.getMemorySpace()) { + auto stageAddr = dyn_cast(stageSpace); + if (!stageAddr || stageAddr.getAddressSpace() != hivm::AddressSpace::GM) + return failure(); + } + + // ③ Enumerate the source copy: memref.copy %src, %stage. + // Collect all users of %stage to verify the expected pattern. + memref::CopyOp srcCopy; + Operation *stageMark = nullptr; + bufferization::ToTensorOp toTensor; + Operation *tensorMark = nullptr; + SmallVector otherUses; + + for (Operation *user : stage.getUsers()) { + if (auto c = dyn_cast(user)) { + if (c.getTarget() == stage) { + if (srcCopy) + return failure(); // multiple incoming copies — ambiguous + srcCopy = c; + } else if (c == copyOp.getOperation()) { + continue; // the outgoing copy we are folding + } else { + otherUses.push_back(user); + } + } else if (isAnnotationMark(user)) { + if (stageMark) + otherUses.push_back(user); + else + stageMark = user; + } else if (stageCast && user == stageCast.getOperation()) { + continue; + } else if (auto tt = dyn_cast(user)) { + if (toTensor) + otherUses.push_back(user); + else + toTensor = tt; + } else { + otherUses.push_back(user); + } + } + + // srcCopy is mandatory; toTensor is optional (Flavour A has it, B doesn't). + if (!srcCopy || !otherUses.empty()) + return failure(); + + // Collect annotation.mark users of the tensor value (only if toTensor + // exists — Flavour A path). + if (toTensor) { + Value tensorVal = toTensor.getResult(); + for (Operation *user : tensorVal.getUsers()) { + if (isAnnotationMark(user) && !tensorMark) + tensorMark = user; + } + } + + Location loc = copyOp.getLoc(); + Value src = castDefaultMemrefToGM(srcCopy.getSource(), loc, rewriter); + Value dstCbuf = dst; + + // ④ Create the merged copy: memref.copy %src, %dst(#space) + rewriter.create(loc, src, dstCbuf); + + // ⑤ If there was an annotation.mark on %stage, move it to %dst. + if (stageMark) { + OperationState state(loc, "annotation.mark"); + state.addOperands(dstCbuf); + for (auto &namedAttr : stageMark->getAttrs()) + state.addAttribute(namedAttr.getName(), namedAttr.getValue()); + rewriter.create(state); + } + + // ⑥ For Flavour A: drop the address space via memory_space_cast and + // re-create the to_tensor pointing at the cbuf alloc. + // For Flavour B: nothing to do — the tensor path never existed. + if (toTensor) { + auto genericType = + MemRefType::get(dstType.getShape(), dstType.getElementType()); + Value cast = + rewriter.create(loc, genericType, dstCbuf); + Value tensor = rewriter.create( + toTensor.getLoc(), toTensor.getType(), cast, toTensor.getRestrict(), + toTensor.getWritable()); + + // ⑦ If there was an annotation.mark on the tensor value, re-create it. + if (tensorMark) { + OperationState state(loc, "annotation.mark"); + state.addOperands(tensor); + for (auto &namedAttr : tensorMark->getAttrs()) + state.addAttribute(namedAttr.getName(), namedAttr.getValue()); + rewriter.create(state); + } + + rewriter.replaceOp(toTensor, tensor); + } + + // ⑧ Erase the dead ops. + rewriter.eraseOp(copyOp); + rewriter.eraseOp(srcCopy); + if (stageCast) + rewriter.eraseOp(stageCast); + if (stageMark) + rewriter.eraseOp(stageMark); + if (tensorMark) + rewriter.eraseOp(tensorMark); + + return success(); + } +}; + +// ============================================================================= +// Pass +// ============================================================================= +namespace { +struct CommonIRToHIVMPass + : public mlir::triton::impl::CommonIRToHIVMBase { + void runOnOperation() override; +}; +} // namespace + +void CommonIRToHIVMPass::runOnOperation() { + auto module = getOperation(); + + // Use greedy rewrite to iteratively apply patterns to fixed point. + // This avoids the type-conversion complexity of the dialect conversion + // framework: tile.alloc → memref.alloc replaces !tile.buf with memref, + // then tile.copy (now seeing memref operands) → hivm.copy, etc. + +// apply pattern separately to ensure their relative order +// template lambda is a C++ 20 feature, so we'll have to use MACRO +#define APPLY_REWRITE_PATTERN(...) \ + { \ + RewritePatternSet patterns(&getContext()); \ + patterns.add<__VA_ARGS__>(&getContext()); \ + if (failed(applyPatternsAndFoldGreedily(module, std::move(patterns)))) { \ + return signalPassFailure(); \ + } \ + } + + APPLY_REWRITE_PATTERN(TileAllocToMemRef); + APPLY_REWRITE_PATTERN(TileSubviewToMemrefSubview); + APPLY_REWRITE_PATTERN(TileToTensorEliminate); + APPLY_REWRITE_PATTERN(TileCopyToHIVM); + APPLY_REWRITE_PATTERN(TileStoreTensorToHIVM); + APPLY_REWRITE_PATTERN(TileLoadToHIVM); + APPLY_REWRITE_PATTERN(TileStoreToHIVM); + APPLY_REWRITE_PATTERN(TileSetFlagToHIVM); + APPLY_REWRITE_PATTERN(TileWaitFlagToHIVM); + APPLY_REWRITE_PATTERN(TilePipeBarrierToHIVM); + APPLY_REWRITE_PATTERN(TileCubeWaitToHIVM); + APPLY_REWRITE_PATTERN(TileGmOffsetToHIVM); + + // Step N: After all tile ops are lowered, convert !tile.buf types remaining + // in tt.func signatures and tt.call ops (these arise when tile.alloc results + // are passed across function boundaries). + module.walk([](triton::FuncOp funcOp) { + auto funcType = funcOp.getFunctionType(); + bool changed = false; + + // Convert input types + SmallVector newInputs; + for (auto ty : funcType.getInputs()) { + if (auto bufTy = dyn_cast(ty)) { + newInputs.push_back(convertBufToMemRef(bufTy)); + changed = true; + } else { + newInputs.push_back(ty); + } + } + + // Convert result types + SmallVector newResults; + for (auto ty : funcType.getResults()) { + if (auto bufTy = dyn_cast(ty)) { + newResults.push_back(convertBufToMemRef(bufTy)); + changed = true; + } else { + newResults.push_back(ty); + } + } + + if (!changed) + return; + + // Update function type + auto newFuncType = + FunctionType::get(funcOp.getContext(), newInputs, newResults); + funcOp.setFunctionType(newFuncType); + + // Update block argument types + if (!funcOp.empty()) { + Block &entry = funcOp.front(); + for (unsigned i = 0; i < entry.getNumArguments(); ++i) { + if (auto bufTy = + dyn_cast(entry.getArgument(i).getType())) { + entry.getArgument(i).setType(convertBufToMemRef(bufTy)); + } + } + } + }); + + // Update tt.call result types to match new function signatures. + module.walk([](triton::CallOp callOp) { + for (unsigned i = 0; i < callOp->getNumResults(); ++i) { + if (auto bufTy = + dyn_cast(callOp->getResult(i).getType())) { + callOp->getResult(i).setType(convertBufToMemRef(bufTy)); + } + } + }); + + APPLY_REWRITE_PATTERN(LowerPtrTensorCastToMemref); + APPLY_REWRITE_PATTERN(LowerMemrefCastToTensor); + APPLY_REWRITE_PATTERN(FoldRoundtripCast); + + // Finally, fold any residual A<->B unrealized cast chains. + SmallVector casts; + module->walk([&](UnrealizedConversionCastOp c) { casts.push_back(c); }); + reconcileUnrealizedCasts(casts); + + APPLY_REWRITE_PATTERN(FoldStagingCopyPattern); +} + +std::unique_ptr> +mlir::triton::createCommonIRToHIVMPass() { + return std::make_unique(); +} diff --git a/third_party/ascend/lib/CommonIRToHIVM/ReduceSumStrengthPass.cpp b/third_party/ascend/lib/CommonIRToHIVM/ReduceSumStrengthPass.cpp new file mode 100644 index 0000000000..378c34ad5f --- /dev/null +++ b/third_party/ascend/lib/CommonIRToHIVM/ReduceSumStrengthPass.cpp @@ -0,0 +1,210 @@ +/* + * Copyright 2026- Xcoresigma Technology Co., Ltd + */ + +#include "ascend/include/CommonIRToHIVM/ReduceSumStrengthPass.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Support/LogicalResult.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallVector.h" + +namespace mlir { +namespace triton { +#define GEN_PASS_DECL_REDUCESUMSTRENGTH +#define GEN_PASS_DEF_REDUCESUMSTRENGTH +#include "ascend/include/CommonIRToHIVM/Passes.h.inc" +} // namespace triton +} // namespace mlir + +using namespace mlir; +using namespace mlir::triton; + +namespace { + +enum class StrategyDType { F32, F16, BF16, I32, Unsupported }; + +StrategyDType getStrategyDType(Type elementType) { + if (isa(elementType)) + return StrategyDType::F32; + if (isa(elementType)) + return StrategyDType::F16; + if (isa(elementType)) + return StrategyDType::BF16; + if (auto intType = dyn_cast(elementType)) { + if (intType.getWidth() == 32) + return StrategyDType::I32; + } + return StrategyDType::Unsupported; +} + +bool isAddCombiner(ReduceOp reduceOp, Type elementType) { + Region &combineOp = reduceOp.getCombineOp(); + if (!combineOp.hasOneBlock() || combineOp.front().getOperations().size() != 2) + return false; + + Operation *combiner = reduceOp.getSingleCombiner(); + if (!combiner) + return false; + + if (isa(elementType)) + return isa(combiner); + if (isa(elementType)) + return isa(combiner); + return false; +} + +struct ReductionStrategy { + StrategyDType dtype; + int64_t minProfitableReduceDim; + int64_t maxProfitableReduceDim; +}; + +const ReductionStrategy reductionStrategies[] = { + {StrategyDType::F32, 256, 16384}, + {StrategyDType::F16, 512, 8192}, + {StrategyDType::BF16, 512, 512}, + {StrategyDType::I32, 0, -1}, +}; + +bool isProfitable(Type elementType, int64_t reduceDim) { + StrategyDType dtype = getStrategyDType(elementType); + for (const ReductionStrategy &entry : reductionStrategies) { + if (entry.dtype != dtype) + continue; + return reduceDim >= entry.minProfitableReduceDim && + reduceDim <= entry.maxProfitableReduceDim && + llvm::isPowerOf2_64(reduceDim); + } + return false; +} + +class ReduceSumStrengthPattern : public OpRewritePattern { +public: + explicit ReduceSumStrengthPattern(MLIRContext *context) + : OpRewritePattern(context) {} + + LogicalResult matchAndRewrite(ReduceOp reduceOp, + PatternRewriter &rewriter) const override { + if (reduceOp.getNumOperands() != 1 || reduceOp->getNumResults() != 1) + return failure(); + + auto inputType = + dyn_cast(reduceOp.getOperand(0).getType()); + if (!inputType || !inputType.hasStaticShape()) + return failure(); + + Type elementType = inputType.getElementType(); + if (!isAddCombiner(reduceOp, elementType)) + return failure(); + + int64_t axis = reduceOp.getAxis(); + if (axis < 0 || axis >= inputType.getRank()) + return failure(); + + ArrayRef shape = inputType.getShape(); + int64_t reduceDim = shape[axis]; + if (reduceDim <= 1 || reduceDim % 2 != 0) + return failure(); + + if (!isProfitable(elementType, reduceDim)) + return failure(); + + SmallVector halfShape(shape.begin(), shape.end()); + halfShape[axis] = reduceDim / 2; + auto halfType = + RankedTensorType::get(halfShape, elementType, inputType.getEncoding()); + + SmallVector loOffsets(inputType.getRank(), + rewriter.getIndexAttr(0)); + SmallVector hiOffsets(inputType.getRank(), + rewriter.getIndexAttr(0)); + SmallVector sizes; + SmallVector strides(inputType.getRank(), + rewriter.getIndexAttr(1)); + sizes.reserve(inputType.getRank()); + for (int64_t dim = 0; dim < inputType.getRank(); ++dim) + sizes.push_back(rewriter.getIndexAttr(halfShape[dim])); + hiOffsets[axis] = rewriter.getIndexAttr(reduceDim / 2); + + Location loc = reduceOp.getLoc(); + Value input = reduceOp.getOperand(0); + auto lo = rewriter.create( + loc, halfType, input, loOffsets, sizes, strides); + auto hi = rewriter.create( + loc, halfType, input, hiOffsets, sizes, strides); + + Value preAdded; + if (isa(elementType)) + preAdded = rewriter.create(loc, lo, hi); + else if (isa(elementType)) + preAdded = rewriter.create(loc, lo, hi); + else + return failure(); + + auto newReduce = rewriter.create(loc, ValueRange{preAdded}, axis); + for (auto attr : reduceOp->getAttrs()) { + if (!newReduce->hasAttr(attr.getName())) + newReduce->setAttr(attr.getName(), attr.getValue()); + } + + Region &newCombineOp = newReduce.getCombineOp(); + rewriter.cloneRegionBefore(reduceOp.getCombineOp(), newCombineOp, + newCombineOp.end()); + rewriter.replaceOp(reduceOp, newReduce.getResults()); + return success(); + } +}; + +class ReduceSumStrengthPass + : public mlir::triton::impl::ReduceSumStrengthBase { +public: + using mlir::triton::impl::ReduceSumStrengthBase< + ReduceSumStrengthPass>::ReduceSumStrengthBase; + + void runOnOperation() override { + if (!enable) + return; + if (splitFactor != 2) { + getOperation().emitError( + "reduce-sum-strength only supports split-factor=2"); + signalPassFailure(); + return; + } + + MLIRContext *context = &getContext(); + RewritePatternSet patterns(context); + patterns.add(context); + + if (applyPatternsGreedily(getOperation(), std::move(patterns)).failed()) + signalPassFailure(); + } +}; + +} // namespace + +namespace mlir { +namespace triton { + +std::unique_ptr> createReduceSumStrengthPass() { + return std::make_unique(); +} + +std::unique_ptr> +createReduceSumStrengthPass(bool enable, int32_t splitFactor) { + ReduceSumStrengthOptions opts; + opts.enable = enable; + opts.splitFactor = splitFactor; + return std::make_unique(opts); +} + +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/Dialect/CMakeLists.txt b/third_party/ascend/lib/Dialect/CMakeLists.txt new file mode 100644 index 0000000000..b7e6595642 --- /dev/null +++ b/third_party/ascend/lib/Dialect/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(TritonAscend) diff --git a/third_party/ascend/lib/Dialect/TritonAscend/CMakeLists.txt b/third_party/ascend/lib/Dialect/TritonAscend/CMakeLists.txt new file mode 100644 index 0000000000..f33061b2d8 --- /dev/null +++ b/third_party/ascend/lib/Dialect/TritonAscend/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(IR) diff --git a/third_party/ascend/lib/Dialect/TritonAscend/IR/CMakeLists.txt b/third_party/ascend/lib/Dialect/TritonAscend/IR/CMakeLists.txt new file mode 100644 index 0000000000..07ad62ddb1 --- /dev/null +++ b/third_party/ascend/lib/Dialect/TritonAscend/IR/CMakeLists.txt @@ -0,0 +1,14 @@ +add_triton_library(TritonAscendIR + TritonAscendAttrs.cpp + TritonAscendDialect.cpp + TritonAscendOps.cpp + + DEPENDS + TritonAscendTableGen + TritonAscendAttrDefsIncGen + + LINK_LIBS PUBLIC + MLIRIR + MLIRLLVMDialect + TritonIR +) diff --git a/third_party/ascend/lib/Dialect/TritonAscend/IR/TritonAscendAttrs.cpp b/third_party/ascend/lib/Dialect/TritonAscend/IR/TritonAscendAttrs.cpp new file mode 100644 index 0000000000..b8424380d9 --- /dev/null +++ b/third_party/ascend/lib/Dialect/TritonAscend/IR/TritonAscendAttrs.cpp @@ -0,0 +1,15 @@ +//===- TritonAscendAttrs.cpp - TritonAscend Attributes Definition +//--------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.h" + +#include "llvm/ADT/SmallSet.h" +#include "llvm/ADT/TypeSwitch.h" + +namespace mlir::triton::ascend {} // namespace mlir::triton::ascend diff --git a/third_party/ascend/lib/Dialect/TritonAscend/IR/TritonAscendDialect.cpp b/third_party/ascend/lib/Dialect/TritonAscend/IR/TritonAscendDialect.cpp new file mode 100644 index 0000000000..b9fe553cf2 --- /dev/null +++ b/third_party/ascend/lib/Dialect/TritonAscend/IR/TritonAscendDialect.cpp @@ -0,0 +1,44 @@ +//===- TritonAscendDialect.cpp - TritonAscend Dialect registration +//--------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.h" + +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/SPIRV/IR/TargetAndABI.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/DialectImplementation.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/Operation.h" + +#include "llvm/ADT/StringExtras.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/AsmParser/Parser.h" +#include "llvm/IR/Function.h" +#include "llvm/Support/SourceMgr.h" + +using namespace mlir; +using namespace mlir::triton::ascend; + +void TritonAscendDialect::initialize() { + addOperations< +#define GET_OP_LIST +#include "ascend/include/Dialect/TritonAscend/IR/TritonAscendOps.cpp.inc" + >(); + addAttributes< +#define GET_ATTRDEF_LIST +#include "ascend/include/Dialect/TritonAscend/IR/TritonAscendOpsAttrDefs.cpp.inc" + >(); +} + +#include "ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.cpp.inc" +#define GET_ATTRDEF_CLASSES +#include "ascend/include/Dialect/TritonAscend/IR/TritonAscendOpsAttrDefs.cpp.inc" +#define GET_OP_CLASSES +#include "ascend/include/Dialect/TritonAscend/IR/TritonAscendOps.cpp.inc" diff --git a/third_party/ascend/lib/Dialect/TritonAscend/IR/TritonAscendOps.cpp b/third_party/ascend/lib/Dialect/TritonAscend/IR/TritonAscendOps.cpp new file mode 100644 index 0000000000..d734e01337 --- /dev/null +++ b/third_party/ascend/lib/Dialect/TritonAscend/IR/TritonAscendOps.cpp @@ -0,0 +1,232 @@ +//===- TritonAscendOps.cpp - TritonAscend dialect operations +//--------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.h" +#include "mlir/Dialect/SPIRV/IR/TargetAndABI.h" +#include "mlir/Dialect/Utils/StaticValueUtils.h" +#include "triton/Tools/Sys/GetEnv.hpp" +#include "llvm/ADT/STLExtras.h" +#include + +using namespace mlir; +using namespace mlir::triton; + +namespace mlir::triton::ascend { + +void GatherOutToUbOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSrcMutable(), + triton::GlobalMemory::get()); +} + +void IndirectLoadOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSrcMutable(), + triton::GlobalMemory::get()); +} + +//-- IndexSelectSimdOp -- +LogicalResult IndexSelectSimdOp::inferReturnTypes( + MLIRContext *context, std::optional location, ValueRange operands, + DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + + // Get operands using adaptor + IndexSelectSimdOpAdaptor adaptor(operands, attributes, properties, regions); + + // Get element type from src pointer + Type elemType; + if (auto ptrType = + dyn_cast(adaptor.getSrc().getType())) { + elemType = ptrType.getPointeeType(); + } else { + return failure(); + } + + // Get index shape to determine the size of dim + auto indicesType = dyn_cast(adaptor.getIndex().getType()); + if (!indicesType) + return failure(); + int64_t numIndices = indicesType.getShape()[0]; + + // Use adaptor to get attributes - this is the compatible way + int32_t dim = adaptor.getDim(); + auto readShapeAttr = adaptor.getReadShape(); + + // Build result shape: read_shape but with dim replaced by numIndices + SmallVector resultShape; + for (size_t i = 0; i < readShapeAttr.size(); ++i) { + if (i == static_cast(dim)) { + resultShape.push_back(numIndices); + } else { + resultShape.push_back(readShapeAttr[i]); + } + } + + // Create result tensor type + inferredReturnTypes.push_back(RankedTensorType::get(resultShape, elemType)); + + return success(); +} + +// FlipOp +LogicalResult +FlipOp::inferReturnTypes(MLIRContext *context, std::optional location, + ValueRange operands, DictionaryAttr attributes, + OpaqueProperties properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + auto inputTy = dyn_cast(operands[0].getType()); + if (!inputTy) { + if (location) + return emitOptionalError(location, + "expected ranked tensor for flip input"); + return failure(); + } + inferredReturnTypes.push_back(inputTy); + return success(); +} + +//-- SortOp -- +LogicalResult +SortOp::inferReturnTypes(MLIRContext *context, std::optional location, + ValueRange operands, DictionaryAttr attributes, + OpaqueProperties properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + if (operands.size() != 1) { + return emitOptionalError(location, + "expected exactly one operand for SortOp"); + } + + if (!isa(operands[0].getType())) { + return emitOptionalError(location, + "operand must be a ranked tensor type for SortOp"); + } + + Value src = operands[0]; + auto srcTy = cast(src.getType()); + auto srcShape = srcTy.getShape(); + auto srcEnc = srcTy.getEncoding(); + + if (srcShape.empty()) { + return emitOptionalError(location, "input tensor must have rank >= 1"); + } + + Type sortedTy = + RankedTensorType::get(srcShape, srcTy.getElementType(), srcEnc); + + inferredReturnTypes.push_back(sortedTy); + + return success(); +} + +//-- Conv1dOp -- +LogicalResult Conv1dOp::verify() { + auto inputType = dyn_cast(getInput().getType()); + auto weightType = dyn_cast(getWeight().getType()); + + constexpr int64_t dim2 = 2; + constexpr int64_t dim3 = 3; + auto inputRank = inputType.getShape().size(); + if (inputRank != dim2 && inputRank != dim3) { + return emitOpError("input tensor must be 2D or 3D, but got rank ") + << inputRank; + } + if (weightType.getShape().size() != dim3) { + return emitOpError("weight tensor must be 3D, but got rank ") + << weightType.getShape().size(); + } + + Value biasValue = getBias(); + if (biasValue) { + auto biasType = dyn_cast(biasValue.getType()); + if (!biasType || biasType.getRank() != 1) { + return emitOpError("bias must be a 1D ranked tensor"); + } + if (biasType.getElementType() != inputType.getElementType()) { + return emitOpError( + "bias must have the same element type as input and weight"); + } + if (biasType.getShape()[0] != weightType.getShape()[0]) { + return emitOpError( + "bias size must match weight's output channel dimension"); + } + } + + int64_t C_in = inputType.getShape()[inputRank - 2]; + int64_t C_out = weightType.getShape()[0]; + int64_t weight_C_in = weightType.getShape()[1]; + + if (C_in % getGroups() != 0) { + return emitOpError("input channels must be divisible by groups"); + } + if (C_out % getGroups() != 0) { + return emitOpError("output channels must be divisible by groups"); + } + if (weight_C_in != C_in / getGroups()) { + return emitOpError( + "weight's input channel dimension must equal input channels / groups"); + } + + return success(); +} + +LogicalResult Conv1dOp::inferReturnTypes( + MLIRContext *context, std::optional location, ValueRange operands, + DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + Conv1dOpAdaptor adaptor(operands, attributes, properties, regions); + + auto inputType = dyn_cast(adaptor.getInput().getType()); + auto weightType = dyn_cast(adaptor.getWeight().getType()); + if (!inputType || !weightType) { + return failure(); + } + + ArrayRef inputShape = inputType.getShape(); + bool isBatched = inputShape.size() == 3; + if (isBatched) { + int64_t N = inputShape[0]; + } + int64_t C_in = inputShape[isBatched ? 1 : 0]; + int64_t L_in = inputShape[isBatched ? 2 : 1]; + + ArrayRef weightShape = weightType.getShape(); + int64_t C_out = weightShape[0]; + int64_t K = weightShape[2]; + + int64_t stride = adaptor.getStride(); + int64_t padding_size = adaptor.getPaddingSize(); + int64_t dilation = adaptor.getDilation(); + + if (stride == 0) { + return failure(); + } + double l_out_double = + static_cast(L_in + 2 * padding_size - dilation * (K - 1) - 1) / + stride + + 1; + int64_t L_out = static_cast(std::floor(l_out_double)); + + constexpr int64_t dim3 = 3; + SmallVector outputShape; + if (isBatched) { + outputShape.push_back(inputShape[0]); + } + outputShape.push_back(weightShape[0]); + outputShape.push_back(L_out); + + Type elementType = inputType.getElementType(); + auto returnType = RankedTensorType::get(outputShape, elementType); + inferredReturnTypes.push_back(returnType); + return success(); +} + +} // namespace mlir::triton::ascend diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition.cpp new file mode 100644 index 0000000000..dd4ef2d397 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition.cpp @@ -0,0 +1,159 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition.h" +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/CloneOps.h" +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.h" +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/InitDependentMap.h" +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.h" +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.h" +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/Scope/IR/Scope.h" +#include "mlir/Pass/PassManager.h" +#include "llvm/Support/Debug.h" + +static constexpr const char *DEBUG_TYPE = "AddControlFlowCondition"; +#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ") +#define LDBG(X) LLVM_DEBUG(DBGS() << (X) << "\n") + +using namespace mlir; +using namespace triton; + +// Check if the module should be skipped for control flow condition processing +static bool shouldSkipControlFlowCondition(ModuleOp module) { + int cubeScopeCount = 0; + int vectorScopeCount = 0; + + module.walk([&](scope::ScopeOp scopeOp) { + auto attr = scopeOp->getAttrOfType("hivm.tcore_type"); + if (!attr) { + return; + } + if (attr.getTcoretype() == hivm::TCoreType::CUBE) { + ++cubeScopeCount; + } else if (attr.getTcoretype() == hivm::TCoreType::VECTOR) { + ++vectorScopeCount; + } + }); + + // If either CUBE or VECTOR scope is missing, skip processing + if (cubeScopeCount == 0 || vectorScopeCount == 0) { + LDBG("CUBE or VECTOR scope missing, skip processing."); + return true; + } + + // Check if scopeOp has ssbuffer.skip + bool hasSkipAttr = false; + module.walk([&](scope::ScopeOp scopeOp) { + if (scopeOp->hasAttr("ssbuffer.skip")) { + hasSkipAttr = true; + } + }); + if (hasSkipAttr) { + LDBG("scopeOp has ssbuffer.skip, skip processing."); + return true; + } + + // Only skip if ALL forOps lack main_loop attr + bool hasMainLoopForOp = false; + module.walk([&](scf::ForOp forOp) { + if (forOp->hasAttr("ssbuffer.main_loop")) { + hasMainLoopForOp = true; + } + }); + if (!hasMainLoopForOp) { + LDBG("All forOps lack ssbuffer.main_loop, skip processing."); + return true; + } + + return false; +} + +void AddControlFlowConditionPass::runOnOperation() { + ModuleOp module = getOperation(); + + LDBG("Enter add controlflow condition pass."); + + if (shouldSkipControlFlowCondition(module)) { + return; + } + + PassManager pm(&getContext(), module.getOperationName()); + ControlFlowConditionInfo info; + + // Step0: Initialize crossCoreDependentMap and intraCoreDependentMap + std::unique_ptr initDependentMapPass( + new InitDependentMapPass()); + initDependentMapPass->setConditionInfo(&info); + pm.addPass(std::move(initDependentMapPass)); + + // Step1: Clone ops in vector/cube to ensure that each block_id has its own + // ops without sharing + pm.addPass(createCloneOpsPass()); + + // Step2: Create if ops based on block_id + std::unique_ptr createIfOpsPass(new CreateIfOpsPass()); + createIfOpsPass->setConditionInfo(&info); + pm.addPass(std::move(createIfOpsPass)); + + // Step3: Update for ops with block counters and inner dependency conditions, + // and insert PIPE_S inter-core synchronization + std::unique_ptr updateForOpsPass(new UpdateForOpsPass()); + updateForOpsPass->setConditionInfo(&info); + pm.addPass(std::move(updateForOpsPass)); + + // Step4:Update the conditions of ifOp based on the intraCoreDependentMap and + // crossCoreDependentMap + auto updatePass = std::make_unique(); + updatePass->setConditionInfo(&info); + pm.addPass(std::move(updatePass)); + + // Step5: Update for loop iteration times based on intraCoreDependentMap + std::unique_ptr updateLoopIterTimesPass( + new UpdateLoopIterTimesPass()); + updateLoopIterTimesPass->setConditionInfo(&info); + pm.addPass(std::move(updateLoopIterTimesPass)); + + if (failed(runPipeline(pm, module))) { + module->emitError() << "[" << DEBUG_TYPE << "] Pass failed!"; + signalPassFailure(); + } + + LDBG("Exit add controlflow condition pass."); +} + +namespace mlir { +namespace triton { + +std::unique_ptr> createAddControlFlowConditionPass() { + return std::make_unique(); +} + +void registerAddControlFlowConditionPasses() { + registerPass(createCloneOpsPass); + registerPass(createCreateIfOpsPass); + registerPass(createUpdateForOpsPass); + registerPass(createAddControlFlowConditionPass); +} +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/CloneOps.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/CloneOps.cpp new file mode 100644 index 0000000000..325cd78dd1 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/CloneOps.cpp @@ -0,0 +1,412 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/CloneOps.h" +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/Utils.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/Scope/IR/Scope.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/IRMapping.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/Support/Debug.h" + +static constexpr const char *DEBUG_TYPE = "CloneOps"; +#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ") +#define LDBG(...) \ + LLVM_DEBUG({ \ + DBGS(); \ + llvm::outs() << __VA_ARGS__; \ + llvm::outs() << "\n"; \ + }) + +using namespace mlir; +using namespace triton; +using namespace hivm; + +// Update op operands using value mapping, skip yield values of forOp +static LogicalResult +updateCloneMapping(Operation *op, llvm::DenseMap &valueMap, + const llvm::DenseSet &yieldValues) { + if (!op) { + return failure(); + } + + for (OpOperand &operand : op->getOpOperands()) { + // Only skip if this operand is a yield value from the main_loop forOp + // Nested yield ops (in scf.if/scf.for) should have their operands updated + Value v = operand.get(); + if (yieldValues.contains(v)) { + continue; + } + + auto it = valueMap.find(v); + if (it != valueMap.end()) { + if (it->second.getType() != v.getType()) { + LDBG("[Error]: type mismatch in value mapping: " + << v.getType() << " vs " << it->second.getType() << "\n"); + return failure(); + } + operand.set(it->second); + } + } + + for (Region ®ion : op->getRegions()) { + for (Block &block : region) { + for (Operation &nestedOp : block) { + if (failed(updateCloneMapping(&nestedOp, valueMap, yieldValues))) { + return failure(); + } + } + } + } + + return success(); +} + +// Clone a single op with IRMapping +static Operation *cloneOpWithMapping(Operation *op, OpBuilder &builder, + llvm::DenseMap &valueMap) { + IRMapping mapper; + for (auto result : op->getResults()) { + if (valueMap.count(result)) { + mapper.map(result, valueMap[result]); + } + } + + Operation *cloned = builder.clone(*op, mapper); + for (auto it : llvm::zip(op->getResults(), cloned->getResults())) { + valueMap[std::get<0>(it)] = std::get<1>(it); + } + + return cloned; +} + +// Clone ops for a single block in vector/cube mode +static LogicalResult +cloneOpsForBlock(int curId, SmallVector &curOps, + const SmallVector &earlierIds, + const llvm::DenseMap> &blockOps, + scf::ForOp forOp) { + if (curOps.empty() || earlierIds.empty()) { + return success(); + } + + // Collect all ops from earlier blocks to clone + SmallVector toClone; + for (int eid : earlierIds) { + llvm::append_range(toClone, blockOps.lookup(eid)); + } + if (toClone.empty()) { + return success(); + } + + llvm::DenseMap valueMap; + SmallVector clonedOps; + OpBuilder builder(curOps.front()); + + for (Operation *op : toClone) { + Operation *cloned = cloneOpWithMapping(op, builder, valueMap); + cloned->setAttr("ssbuffer.block_id", builder.getI32IntegerAttr(curId)); + if (auto origBlockIdAttr = + op->getAttrOfType("ssbuffer.block_id")) { + cloned->setAttr("ssbuffer.clone", origBlockIdAttr); + } + clonedOps.push_back(cloned); + } + + curOps.insert(curOps.begin(), clonedOps.begin(), clonedOps.end()); + + llvm::DenseSet yieldValues; + if (auto yieldOp = dyn_cast(forOp.getBody()->getTerminator())) { + for (Value operand : yieldOp.getOperands()) { + yieldValues.insert(operand); + } + } + + for (Operation *op : curOps) { + if (failed(updateCloneMapping(op, valueMap, yieldValues))) { + return failure(); + } + } + + if (failed(topologicalSort(curOps))) { + return failure(); + } + + return success(); +} + +// Clone ops for all blocks in main loop +LogicalResult CloneOpsPass::cloneOpsInMainLoop(scf::ForOp forOp) { + llvm::DenseMap> blockOps; + if (failed(collectOpsByBlockId(forOp, blockOps))) { + return failure(); + } + + SmallVector idsInOrder = getBlockIdsInOrder(forOp); + + for (int i = idsInOrder.size() - 1; i >= 0; --i) { + int curId = idsInOrder[i]; + SmallVector earlierIds(idsInOrder.begin(), idsInOrder.begin() + i); + + if (failed(cloneOpsForBlock(curId, blockOps[curId], earlierIds, blockOps, + forOp))) { + return failure(); + } + } + + return success(); +} + +// Check if an op should be erased during cleanup (for cube) +static bool shouldEraseOpForCube(Operation *op) { + if (auto ifOp = dyn_cast(op)) { + if (ifOp->hasAttr("ssbuffer.cross_buffer") || + ifOp->hasAttr("ssbuffer.intra_buffer") || + ifOp->hasAttr("ssbuffer.load_store")) { + return true; + } + } + + if (isa(op) || isa(op) || + isa(op)) { + return true; + } + + // Copy and fill ops: erase if operand not used elsewhere + if (isa(op) || + (isa(op) && op->getNumResults() == 0)) { + if (op->getNumOperands() > 1) { + Value secondOperand = op->getOperand(1); + bool usedByOtherOp = + llvm::any_of(secondOperand.getUsers(), + [&](Operation *user) { return user != op; }); + if (!usedByOtherOp) { + return true; + } + } + } + + // Erase ops with no used results + return llvm::none_of(op->getResults(), + [](auto result) { return !result.use_empty(); }); +} + +// Check if an op should be erased (for vector) +static bool shouldEraseOpForVector(Operation *op) { + if (auto ifOp = dyn_cast(op)) { + if (ifOp->hasAttr("ssbuffer.cross_buffer") || + ifOp->hasAttr("ssbuffer.intra_buffer") || + ifOp->hasAttr("ssbuffer.load_store")) { + return true; + } + } + + return llvm::none_of(op->getResults(), + [](auto result) { return !result.use_empty(); }); +} + +// Cleanup for cloned ops in a forOp +static LogicalResult +cleanupClonedOps(scf::ForOp forOp, + llvm::DenseMap> &blockOps, + const SmallVector &idsInOrder, bool isCube) { + for (int i = idsInOrder.size() - 1; i >= 0; --i) { + auto &curOps = blockOps[idsInOrder[i]]; + if (curOps.empty()) { + continue; + } + + // Find last index of cloned ops + int startIdx = -1; + for (int j = curOps.size() - 1; j >= 0; --j) { + if (curOps[j]->hasAttr("ssbuffer.clone")) { + startIdx = j; + break; + } + } + if (startIdx < 0) { + continue; + } + + // Erase cloned ops that are not needed + for (int j = startIdx; j >= 0; --j) { + Operation *op = curOps[j]; + if (!op->hasAttr("ssbuffer.clone")) { + break; + } + bool shouldErase = + isCube ? shouldEraseOpForCube(op) : shouldEraseOpForVector(op); + if (shouldErase) { + op->erase(); + } + } + } + + // Check whether new forOp is valid after cleanup + for (Operation &op : forOp.getBody()->without_terminator()) { + if (op.hasAttr("ssbuffer.clone")) { + if (isa(op) || isa(op) || + isa(op)) { + LDBG("[ERROR]: Cloned sync/fixpipe op should have been erased: " + << op.getName() << "\n"); + return failure(); + } + } + } + + return success(); +} + +// Cleanup cloned ops for a single main loop +LogicalResult CloneOpsPass::cleanupClonedOpsInMainLoop(scf::ForOp forOp) { + ModuleOp module = getOperation(); + scope::ScopeOp scopeOp = forOp->getParentOfType(); + if (!scopeOp) { + return success(); + } + + auto attr = scopeOp->getAttrOfType("hivm.tcore_type"); + if (!attr) { + return success(); + } + + bool isCube = (attr == hivm::TCoreTypeAttr::get(module.getContext(), + hivm::TCoreType::CUBE)); + + llvm::DenseMap> blockOps; + if (failed(collectOpsByBlockId(forOp, blockOps))) { + return failure(); + } + + SmallVector idsInOrder = getBlockIdsInOrder(forOp); + if (failed(cleanupClonedOps(forOp, blockOps, idsInOrder, isCube))) { + return failure(); + } + + return success(); +} + +// Validate that each block_id's ops form contiguous ranges (not interleaved +// with other ids) e.g., [1,1,2,2] is valid, but [1,2,1,2] is invalid +static bool areBlockIdsConsecutive(scf::ForOp forOp) { + SmallVector idsInOrder; + for (Operation &op : forOp.getBody()->without_terminator()) { + auto blockIdAttr = op.getAttrOfType("ssbuffer.block_id"); + if (!blockIdAttr) { + LDBG("[ERROR]: Op missing ssbuffer.block_id: " << op.getName() << "\n"); + return false; + } + + idsInOrder.push_back(blockIdAttr.getInt()); + } + + // Check that each block_id forms a contiguous range + for (size_t i = 0; i < idsInOrder.size();) { + int currentId = idsInOrder[i]; + size_t j = i; + + while (j < idsInOrder.size() && idsInOrder[j] == currentId) { + ++j; + } + + for (size_t k = j; k < idsInOrder.size(); ++k) { + if (idsInOrder[k] == currentId) { + LDBG("[ERROR]: block_id: " << currentId << " is interleaved\n"); + return false; + } + } + + i = j; + } + + return true; +} + +LogicalResult CloneOpsPass::validateBlockIdsConsecutive(ModuleOp module) { + WalkResult result = module.walk([&](Operation *op) -> WalkResult { + if (!op->hasAttr("ssbuffer.main_loop")) { + return WalkResult::advance(); + } + auto forOp = dyn_cast(op); + if (!forOp) { + LDBG("[Error]: op with ssbuffer.main_loop is not a scf::ForOp\n"); + return WalkResult::interrupt(); + } + if (!areBlockIdsConsecutive(forOp)) { + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + if (result.wasInterrupted()) + return failure(); + + return success(); +} + +void CloneOpsPass::runOnOperation() { + ModuleOp module = getOperation(); + + LDBG("before cloneOps:\n" << module << "\n"); + + // Validate block_ids are consecutive before cloning + if (failed(validateBlockIdsConsecutive(module))) + return; + + // Clone ops in vector/cube to ensure that each block_id has its own + // ops without sharing + module.walk([&](Operation *op) -> WalkResult { + if (!op->hasAttr("ssbuffer.main_loop")) { + return WalkResult::advance(); + } + auto forOp = dyn_cast(op); + if (!forOp) { + LDBG("[Error]: op with ssbuffer.main_loop is not a scf::ForOp\n"); + signalPassFailure(); + return WalkResult::interrupt(); + } + + if (failed(cloneOpsInMainLoop(forOp))) { + signalPassFailure(); + return WalkResult::interrupt(); + } + + if (failed(cleanupClonedOpsInMainLoop(forOp))) { + signalPassFailure(); + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + + LDBG("after cloneOps:\n" << module << "\n"); +} + +namespace mlir { +namespace triton { + +std::unique_ptr> createCloneOpsPass() { + return std::make_unique(); +} + +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.cpp new file mode 100644 index 0000000000..84e2b2cc40 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.cpp @@ -0,0 +1,363 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/CreateIfOps.h" +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/Utils.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/Scope/IR/Scope.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/IRMapping.h" +#include "llvm/Support/Debug.h" + +static constexpr const char *DEBUG_TYPE = "CreateIfOps"; +#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ") +#define LDBG(...) \ + LLVM_DEBUG({ \ + DBGS(); \ + llvm::outs() << __VA_ARGS__; \ + llvm::outs() << "\n"; \ + }) + +using namespace llvm; +using namespace mlir; +using namespace triton; + +// Check if a value is used outside the given region ops +static bool isUsedOutsideRegion(Value v, + const llvm::DenseSet ®ionOps) { + for (OpOperand &use : v.getUses()) { + if (!regionOps.contains(use.getOwner())) { + return true; + } + } + + return false; +} + +// Find the iteration argument in main loop that corresponds to the given value +static Value findIterArgInMainLoop(Value v, mlir::Type t) { + for (Operation *user : v.getUsers()) { + auto yieldOp = dyn_cast(user); + if (!yieldOp) { + continue; + } + + auto forOp = dyn_cast(yieldOp->getParentOp()); + if (!forOp) { + continue; + } + + for (auto [idx, operand] : llvm::enumerate(yieldOp.getOperands())) { + if (operand.getAsOpaquePointer() == v.getAsOpaquePointer()) { + Value iterArg = forOp.getRegionIterArgs()[idx]; + if (iterArg.getType() == t) { + return iterArg; + } + } + } + } + + LDBG("[Error]: else yield value not found in forOp iter_args: " << v << "\n"); + return nullptr; +} + +// Replace uses of old values that are outside the ifOp with the new values from +// ifOp +static LogicalResult replaceExternalIfOpUses(scf::IfOp ifOp, + ArrayRef oldYieldValues) { + for (size_t i = 0; i < oldYieldValues.size(); ++i) { + Value oldVal = oldYieldValues[i]; + + if (!oldVal) { + LDBG("[Error]: oldVal is null at index " << i << "\n"); + return failure(); + } + + if (i >= ifOp.getNumResults()) { + LDBG("[Error]: index " << i << " exceeds ifOp results count " + << ifOp.getNumResults() << "\n"); + return failure(); + } + + Value newVal = ifOp.getResult(i); + if (oldVal.getType() != newVal.getType()) { + LDBG("[Error]: type mismatch at index " << i << ": " << oldVal.getType() + << " vs " << newVal.getType() + << "\n"); + return failure(); + } + + SmallVector usesToReplace; + + for (OpOperand &use : llvm::make_early_inc_range(oldVal.getUses())) { + Operation *user = use.getOwner(); + + // Skip uses inside the ifOp + if (ifOp->isAncestor(user)) { + continue; + } + + // Skip uses after the ifOp in the same block + if (user->getBlock() == ifOp->getBlock() && + !ifOp->isBeforeInBlock(user)) { + continue; + } + + usesToReplace.push_back(&use); + } + + for (OpOperand *use : usesToReplace) { + use->set(newVal); + } + } + + return success(); +} + +// Compute yield values for each block: values that need to be yielded from the +// if +LogicalResult CreateIfOpsPass::computeYieldValues( + scf::ForOp forOp, + const llvm::DenseMap> &blockOps, + llvm::DenseMap> &thenYieldValues, + llvm::DenseMap> &elseYieldValues) { + for (auto &p : blockOps) { + int id = p.first; + const SmallVector &ops = p.second; + + // Collect all ops in this block including nested ones + llvm::DenseSet regionOps; + for (Operation *op : ops) { + if (failed(collectAllNestedOps(op, regionOps))) { + return failure(); + } + } + + // Find values used outside the region + SmallVector yieldValues; + for (Operation *op : ops) { + for (Value res : op->getResults()) { + if (isUsedOutsideRegion(res, regionOps)) { + yieldValues.push_back(res); + } + } + } + + thenYieldValues[id] = yieldValues; + + // For else branch, use iteration args from main loop + elseYieldValues[id].clear(); + elseYieldValues[id].reserve(yieldValues.size()); + for (Value v : yieldValues) { + Value foundArg = findIterArgInMainLoop(v, v.getType()); + if (!foundArg) { + return failure(); + } else { + elseYieldValues[id].push_back(foundArg); + } + } + } + + return success(); +} + +// Create result types from yield values +static SmallVector +getResultTypes(const SmallVector &values) { + SmallVector types; + + for (Value v : values) { + types.push_back(v.getType()); + } + + return types; +} + +// Create ifOp for a single block +static scf::IfOp createIfOpForBlock(OpBuilder &builder, Location loc, + int blockId, + const SmallVector &thenValues, + const SmallVector &elseValues) { + bool needsYield = !thenValues.empty(); + // Check size consistency + if (needsYield && thenValues.size() != elseValues.size()) { + LDBG("[Error]: then/else yield count mismatch: " + << thenValues.size() << " vs " << elseValues.size() << "\n"); + return scf::IfOp(); + } + + // Check type consistency + if (needsYield) { + for (size_t i = 0; i < thenValues.size(); ++i) { + if (thenValues[i].getType() != elseValues[i].getType()) { + LDBG("[Error]: then/else yield type mismatch at index " + << i << ": " << thenValues[i].getType() << " vs " + << elseValues[i].getType() << "\n"); + return scf::IfOp(); + } + } + } + + SmallVector resultTypes = getResultTypes(thenValues); + + Value trueVal = builder.create(loc, builder.getI1Type(), + builder.getBoolAttr(true)); + + scf::IfOp ifOp; + if (needsYield) { + ifOp = builder.create(loc, resultTypes, trueVal, true); + } else { + ifOp = builder.create(loc, TypeRange{}, trueVal, false); + } + + ifOp->setAttr("ssbuffer.if", builder.getI32IntegerAttr(blockId)); + + return ifOp; +} + +// Move ops to then branch and create yield +static LogicalResult moveOpsToThenBranch(scf::IfOp ifOp, + SmallVector &ops, + const SmallVector &thenValues, + const SmallVector &elseValues, + Location loc) { + if (ops.empty() && !thenValues.empty()) { + LDBG("[Error]: moving empty ops but thenValues not empty\n"); + return failure(); + } + + Block &thenBlock = ifOp.getThenRegion().front(); + + for (Operation *op : llvm::reverse(ops)) { + op->moveBefore(&thenBlock, thenBlock.begin()); + } + + if (!thenValues.empty()) { + OpBuilder thenBuilder(&thenBlock, thenBlock.end()); + thenBuilder.create(loc, thenValues); + + Block &elseBlock = ifOp.getElseRegion().front(); + OpBuilder elseBuilder(&elseBlock, elseBlock.end()); + elseBuilder.create(loc, elseValues); + } + + return success(); +} + +// Create if ops (scf.if %true) for each block_id in the main loop +LogicalResult CreateIfOpsPass::createIfInMainLoop( + scf::ForOp forOp, + const llvm::DenseMap> &blockOps, + const llvm::DenseMap> &thenYieldValues, + const llvm::DenseMap> &elseYieldValues) { + SmallVector ids = getBlockIdsInOrder(forOp); + + for (int id : ids) { + const SmallVector &ops = blockOps.lookup(id); + if (ops.empty()) { + continue; + } + + OpBuilder builder(ops.front()); + Location loc = ops.front()->getLoc(); + + scf::IfOp ifOp = + createIfOpForBlock(builder, loc, id, thenYieldValues.lookup(id), + elseYieldValues.lookup(id)); + if (!ifOp) { + return failure(); + } + + // Clone ops list since we'll modify it + SmallVector opsToMove = ops; + if (failed(moveOpsToThenBranch(ifOp, opsToMove, thenYieldValues.lookup(id), + elseYieldValues.lookup(id), loc))) { + return failure(); + } + + if (failed(replaceExternalIfOpUses(ifOp, thenYieldValues.lookup(id)))) { + return failure(); + } + } + + return success(); +} + +void CreateIfOpsPass::runOnOperation() { + ModuleOp module = getOperation(); + + LDBG("before createIfOps:\n" << module << "\n"); + + module.walk([&](Operation *op) -> WalkResult { + if (!op->hasAttr("ssbuffer.main_loop")) { + return WalkResult::advance(); + } + auto forOp = dyn_cast(op); + if (!forOp) { + LDBG("[Error]: op with ssbuffer.main_loop is not a scf::ForOp\n"); + signalPassFailure(); + return WalkResult::interrupt(); + } + + // Create if ops (scf.if %true) by block_id + llvm::DenseMap> blockOps; + if (failed(collectOpsByBlockId(forOp, blockOps))) { + signalPassFailure(); + return WalkResult::interrupt(); + } + + if (info) { + info->blockCounterNums[forOp] = blockOps.size(); + } + + llvm::DenseMap> thenYieldValues; + llvm::DenseMap> elseYieldValues; + + if (failed(computeYieldValues(forOp, blockOps, thenYieldValues, + elseYieldValues))) { + signalPassFailure(); + return WalkResult::interrupt(); + } + + if (failed(createIfInMainLoop(forOp, blockOps, thenYieldValues, + elseYieldValues))) { + signalPassFailure(); + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + + LDBG("after createIfOps:\n" << module << "\n"); +} + +namespace mlir { +namespace triton { + +std::unique_ptr> createCreateIfOpsPass() { + return std::make_unique(); +} + +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/InitDependentMap.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/InitDependentMap.cpp new file mode 100644 index 0000000000..32705e41f6 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/InitDependentMap.cpp @@ -0,0 +1,295 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/InitDependentMap.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Debug.h" + +// Role in dependency attribute: ssbuffer.crossDeps/intraDeps = [groupId, +// roleId] role: 1=producer, 0=consumer +static const int producerId = 1; +static const int consumerId = 0; +static constexpr const char *DEBUG_TYPE = "InitDependentMap"; +#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ") +#define LDBG(...) LLVM_DEBUG(DBGS() << __VA_ARGS__ << "\n") + +using namespace mlir; +using namespace triton; + +// Function: Check if a consumer op is inside a given mainLoop +// but not inside any nested mainloop, and push to consumers if true +// Input: Consumer operation, target mainLoop forOp, reference to consumers +// vector Output: consumers - push consumer to this vector if it's inside +// mainLoop (not in nested mainloop) Return: 0 for success (consumer pushed), -1 +// for failure (not in mainLoop or in nested mainLoop) +static int isConsumerInMainLoop(Operation *consumer, scf::ForOp mainLoop, + SmallVector &consumers) { + Operation *current = consumer->getParentOp(); + + // Traverse up the parent chain until we reach the top (nullptr) + while (current != nullptr) { + if (auto forOp = dyn_cast(current)) { + if (forOp->hasAttr("ssbuffer.main_loop") && forOp != mainLoop) { + // comsumer Op not in the current mainloop + return 0; + } + } + // If we reach the target mainLoop, consumer is inside it + if (current == mainLoop) { + consumers.push_back(consumer); + return 0; + } + current = current->getParentOp(); + } + + LDBG("Can not find the consumer's mainloop!"); + return -1; +} + +// Function: Collect ops with dependency attributes, grouped by group ID +// Input: Root operation to traverse (module or forOp), attribute name +// Output: depsByGroup - Ops grouped by group ID, format: group -> [(op, role), +// ...] +// Attribute format: [group, role], role: 1=producer, 0=consumer +// Return: 0 for success, -1 for failure +static int +collectDepsByGroup(Operation *rootOp, const char *attrName, + llvm::DenseMap>> + &depsByGroup) { + // Attribute format: {ssbuffer.crossDeps/intraDeps = [group, role]} + int ret = 0; + int depSize = 2; + + rootOp->walk([&](Operation *op) { + auto depsAttr = op->getAttrOfType(attrName); + if (!depsAttr) + return; + + if (depsAttr.size() < depSize) { + LDBG("format of dependency attribute error!"); + ret = -1; + return; + } + + int group = cast(depsAttr[0]).getInt(); + int role = cast(depsAttr[1]).getInt(); + depsByGroup[group].push_back({op, role}); + }); + + return ret; +} + +// Function: Build mapping from consumer to producer +// Input: Ops grouped by group ID, format: group -> [(op, role), ...] +// role: 1=producer, 0=consumer +// mainLoop: if not nullptr, only include consumers inside this mainLoop +// Output: result - Mapping from consumer Value to list of producer Values +// Return: 0 for success, -1 for failure +static int buildProducerConsumerMapping( + llvm::DenseMap>> &depsByGroup, + llvm::DenseMap> &result, + scf::ForOp mainLoop = nullptr) { + for (auto &groupEntry : depsByGroup) { + auto &ops = groupEntry.second; + + // Collect all producers and consumers in this group + SmallVector producers; + SmallVector consumers; + + for (auto &opRole : ops) { + Operation *op = opRole.first; + int role = opRole.second; + if (role == producerId) { + producers.push_back(op); + } else if (role == consumerId) { + // For intra-core mapping, only include consumers inside mainLoop + if (mainLoop != nullptr) { + if (isConsumerInMainLoop(op, mainLoop, consumers) != 0) { + LDBG("isConsumerInMainLoop failed"); + return -1; + } + } else { + consumers.push_back(op); + } + } else { + LDBG("Get error role id in dependency attribute: OP: " + << *op << ", role: " << role); + return -1; + } + } + + // Skip if no consumers (for intra-core mapping with mainLoop filter) + if (mainLoop != nullptr && consumers.empty()) + continue; + + // For each consumer, build mapping to all producers + for (Operation *consumer : consumers) { + for (Value consumerResult : consumer->getResults()) { + SmallVector producerValues; + for (Operation *producer : producers) { + for (Value producerResult : producer->getResults()) { + producerValues.push_back(producerResult); + } + } + result[consumerResult] = producerValues; + } + } + } + + return 0; +} + +// Initialize crossCoreDependentMap (cross-core data dependency) +// Find ops with ssbuffer.crossDeps attribute +// Attribute value is a list: [group, role], role: 1=producer, 0=consumer +// Map key is consumer, value is list of all producers in the same group +// Return: 0 for success, -1 for failure +int initCrossCoreDependentMap(ModuleOp module, ControlFlowConditionInfo *info) { + llvm::DenseMap>> + crossDepsByGroup; + if (collectDepsByGroup(module, "ssbuffer.crossDeps", crossDepsByGroup) != 0) { + LDBG("collectDepsByGroup on crossDeps Failed!"); + return -1; + } + + llvm::DenseMap> crossDepsMap; + if (buildProducerConsumerMapping(crossDepsByGroup, crossDepsMap) != 0) { + LDBG("buildProducerConsumerMapping on crossDeps Failed!"); + return -1; + } + info->crossCoreDependentMap = crossDepsMap; + return 0; +} + +// Initialize intraCoreDependentMap (intra-core data dependency) +// Find forOp with ssbuffer.main_loop attribute +// Collect all intra-core deps from module (producers may be outside the loop) +// For each mainLoop, filter consumers that are inside it (not in nested +// mainloops) Return: 0 for success, -1 for failure +int initIntraCoreDependentMap(ModuleOp module, ControlFlowConditionInfo *info) { + // Collect all intra-core deps from the entire module + llvm::DenseMap>> + allIntraDepsByGroup; + if (collectDepsByGroup(module, "ssbuffer.intraDeps", allIntraDepsByGroup) != + 0) { + LDBG("collectDepsByGroup on intraDeps Failed!"); + return -1; + } + + // For each mainLoop, build mapping with consumers inside it + int ret = 0; + module.walk([&](Operation *op) { + if (!op->hasAttr("ssbuffer.main_loop")) + return; + auto forOp = dyn_cast(op); + if (!forOp) { + LDBG("Do not support other mainloop except forOp!"); + ret = -1; + return; + } + + llvm::DenseMap> depMap; + if (buildProducerConsumerMapping(allIntraDepsByGroup, depMap, forOp) != 0) { + LDBG("buildProducerConsumerMapping on intraDeps Failed!"); + ret = -1; + return; + } + + // Only insert if there are dependencies for this mainLoop + if (!depMap.empty()) { + info->intraCoreDependentMap[forOp] = depMap; + } + }); + return ret; +} + +// Print all dependent maps for verification +static void printDependentMaps(ControlFlowConditionInfo *info) { + // Print crossCoreDependentMap + LDBG("crossCoreDependentMap size: " << info->crossCoreDependentMap.size()); + LDBG("crossCoreDependentMap contents:"); + for (auto &entry : info->crossCoreDependentMap) { + Value consumer = entry.first; + SmallVector &producers = entry.second; + LDBG(" Consumer: " << consumer + << " (producers count: " << producers.size() << ")"); + for (Value producer : producers) { + LDBG(" Producer: " << producer); + } + } + + // Print intraCoreDependentMap + LDBG("intraCoreDependentMap size: " << info->intraCoreDependentMap.size()); + LDBG("intraCoreDependentMap contents:"); + for (auto &forEntry : info->intraCoreDependentMap) { + scf::ForOp forOp = forEntry.first; + auto &depMap = forEntry.second; + LDBG(" ForOp (depMap size: " << depMap.size() << "):"); + LDBG(" "); + forOp->print(llvm::dbgs(), OpPrintingFlags().skipRegions()); + + for (auto &entry : depMap) { + Value consumer = entry.first; + SmallVector &producers = entry.second; + LDBG(" Consumer: " << consumer << " (producers count: " + << producers.size() << ")"); + for (Value producer : producers) { + LDBG(" Producer: " << producer); + } + } + } +} + +void InitDependentMapPass::runOnOperation() { + ModuleOp module = getOperation(); + LDBG("Enter InitDependentMap pass."); + + // Step 1: Initialize crossCoreDependentMap + if (initCrossCoreDependentMap(module, info) != 0) { + LDBG("initCrossCoreDependentMap failed!"); + signalPassFailure(); + return; + } + + // Step 2: Initialize intraCoreDependentMap + if (initIntraCoreDependentMap(module, info) != 0) { + LDBG("initIntraCoreDependentMap failed!"); + signalPassFailure(); + return; + } + + // Print all dependent maps for verification + printDependentMaps(info); + + LDBG("Exit InitDependentMap pass."); +} + +namespace mlir { +namespace triton { +std::unique_ptr> createInitDependentMapPass() { + return std::make_unique(); +} +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.cpp new file mode 100644 index 0000000000..40b9abbdd9 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.cpp @@ -0,0 +1,1482 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateConditionInfo.h" +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/HIVM/IR/HIVMImpl.h" +#include "bishengir/Dialect/HIVM/IR/HIVMInterfaces.h" +#include "bishengir/Dialect/HIVM/Transforms/Passes.h" +#include "bishengir/Dialect/HIVM/Utils/Utils.h" +#include "bishengir/Dialect/Scope/IR/Scope.h" +#include "llvm/ADT/APFloat.h" +#include "llvm/Support/Debug.h" + +static constexpr const char *DEBUG_TYPE = "UpdateConditionInfoPass"; +static constexpr const char *SSBUFFER_Main_LOOP = "ssbuffer.main_loop"; +static constexpr const char *SSBUFFER_IF = "ssbuffer.if"; +static constexpr int SSBUF_ADDR_SPACE = 11; +static constexpr int ADDR_INT_TYPE = 64; +static constexpr int CONST_INT_TYPE = 32; +static constexpr int VECTOR_SSBUF_OFFSET = 1024; +static constexpr int VALUE_SSBUF_OFFSET = 4; +static constexpr int UPDATE_CONDITION_INFO_SUCCESS = 0; +static constexpr int UPDATE_CONDITION_INFO_FAILED = -1; + +#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ") +#define LDBG(X) LLVM_DEBUG(DBGS() << (X) << "\n") +using namespace mlir; +using namespace triton; +using namespace hivm; + +// Allocate the SSBuffer pointer +SmallVector> +UpdateConditionInfoPass::allocSSBuffer(ModuleOp module) { + OpBuilder builder(module.getContext()); + auto i64Type = builder.getIntegerType(ADDR_INT_TYPE); + auto i32Type = builder.getIntegerType(CONST_INT_TYPE); + auto ptrType = + mlir::LLVM::LLVMPointerType::get(builder.getContext(), SSBUF_ADDR_SPACE); + + // alloc 2 group of ssbuffer pointers: + // Core Vector 0: allocate ssbuffer address: 0, 4, 8, ... + // Core Vector 1: allocate ssbuffer address: 1024, 1028, 1032, ... + SmallVector> ssbufferPtrs; + SmallVector ssbufferVec0Ptrs; + SmallVector ssbufferVec1Ptrs; + int numBuffers = info->crossCoreDependentMap.size(); + if (numBuffers == 0) { + LDBG("crossCoreDependentMap is empty!"); + return ssbufferPtrs; + } + + module->walk([&](Operation *op) { + if (auto scopeOp = dyn_cast(op)) { + builder.setInsertionPoint(scopeOp); + auto zeroConst = builder.create( + scopeOp->getLoc(), i32Type, builder.getIntegerAttr(i32Type, 0)); + + for (int i = 0; i < numBuffers; i++) { + auto addr0Attr = + builder.getIntegerAttr(i64Type, i * VALUE_SSBUF_OFFSET); + auto addr1Attr = builder.getIntegerAttr( + i64Type, VECTOR_SSBUF_OFFSET + i * VALUE_SSBUF_OFFSET); + + auto addr0Const = builder.create( + scopeOp->getLoc(), i64Type, addr0Attr); + auto addr1Const = builder.create( + scopeOp->getLoc(), i64Type, addr1Attr); + + auto ptr0 = builder.create( + scopeOp->getLoc(), ptrType, addr0Const.getResult()); + auto ptr1 = builder.create( + scopeOp->getLoc(), ptrType, addr1Const.getResult()); + + builder.create(scopeOp->getLoc(), zeroConst, ptr0); + builder.create(scopeOp->getLoc(), zeroConst, ptr1); + + ssbufferVec0Ptrs.push_back(ptr0.getResult()); + ssbufferVec1Ptrs.push_back(ptr1.getResult()); + } + return mlir::WalkResult::interrupt(); + } + return mlir::WalkResult::advance(); + }); + + ssbufferPtrs.push_back(ssbufferVec0Ptrs); + ssbufferPtrs.push_back(ssbufferVec1Ptrs); + return ssbufferPtrs; +} + +// Collect dependency buffer +void UpdateConditionInfoPass::collectDependencyBuffers( + scf::ForOp forOp, + DenseMap>> &crossCoreBuffers, + DenseMap>> &intraCoreBuffers) { + int crossCoreIdx = 0; + for (auto &entry : info->crossCoreDependentMap) { + crossCoreBuffers[crossCoreIdx][entry.first] = entry.second; + crossCoreIdx++; + } + + if (info->intraCoreDependentMap.count(forOp)) { + auto &forOpDeps = info->intraCoreDependentMap[forOp]; + int intraCoreIdx = 0; + for (auto &entry : forOpDeps) { + intraCoreBuffers[intraCoreIdx][entry.first] = entry.second; + intraCoreIdx++; + } + } +} + +// Helper: Find the tcb group id that contains value v +// Returns the group id if found, -1 otherwise +static int +findTcbGroupId(Value v, + DenseMap> &tightlyCoupledBufferGroups) { + for (auto &tcbEntry : tightlyCoupledBufferGroups) { + if (llvm::is_contained(tcbEntry.second, v)) { + return tcbEntry.first; + } + } + return UPDATE_CONDITION_INFO_FAILED; +} + +// Helper: Add all equivalent values from tcbValues to values (excluding v +// itself) +int addEquivalentValues(Value v, SmallVector &tcbValues, + SmallVector &values) { + int ret = -1; + for (Value equivValue : tcbValues) { + if (equivValue != v && !llvm::is_contained(values, equivValue)) { + ret = 0; + values.push_back(equivValue); + } + } + return ret; +} + +// Buffers in cube/vector scope are different value +// scope { +// %alloc_5 = memref.alloc() +// annotation.mark %alloc_5 {hivm.tightly_coupled_buffer = +// #hivm.tightly_coupled_buffer<1>} fixpipe ins() outs(alloc_5) +// } {CUBE} +// scope { +// %alloc_6 = memref.alloc() +// bufferization.to_tensor %alloc_6 +// annotation.mark %alloc_6 {hivm.tightly_coupled_buffer = +// #hivm.tightly_coupled_buffer<1>} +// } {VECTOR} +// % alloc_5 and alloc_6 have the same tightly_coupled_buffer id indicates they +// are the same buffer crossCoreBuffers only include producer buffers in one +// scope: {cosumer: {%alloc_6}, ...} This function is to extend crossCoreBuffers +// to include producer buffers in another scope: {cosumer: {%alloc_6, %alloc_5}, +// ...} +DenseMap>> +UpdateConditionInfoPass::extendCrossCoreBuffersWithEquivalentValues( + ModuleOp module, + DenseMap>> crossCoreBuffers) { + // Error map to return when error occurs + DenseMap>> errorMap; + errorMap[-1] = DenseMap>(); + + // copy from crossCoreBuffers + DenseMap>> extendedCrossCoreBuffers; + for (auto &entry : crossCoreBuffers) { + int groupIdx = entry.first; + for (auto &entry2 : entry.second) { + extendedCrossCoreBuffers[groupIdx][entry2.first] = entry2.second; + } + } + + // Get the buffers have the same tightly_coupled_buffer id + int ret = 0; + DenseMap> tightlyCoupledBufferGroups; + WalkResult walkResult = module.walk([&](Operation *op) -> WalkResult { + if (isa(op)) { + if (auto tcbAttr = op->getAttrOfType( + "hivm.tightly_coupled_buffer")) { + auto id = tcbAttr.getId(); + if (id.has_value()) { + int tcb = id.value(); + Value markedValue = op->getOperand(0); + tightlyCoupledBufferGroups[tcb].push_back(markedValue); + } else { + ret = -1; + LDBG("hivm.tightly_coupled_buffer Attribute has no id!"); + return WalkResult::interrupt(); + } + } + } + return WalkResult::advance(); + }); + if (ret == -1) { + return errorMap; + } + + // Extend crossCoreBuffers to include producer buffers in another scope + for (auto &entry : extendedCrossCoreBuffers) { + int groupIdx = entry.first; + for (auto &deps : entry.second) { + SmallVector &producers = deps.second; + for (Value buffer : producers) { + int tcbGroupId = findTcbGroupId(buffer, tightlyCoupledBufferGroups); + if (tcbGroupId == -1) { + LDBG("Can not find tightly_coupled_buffer id"); + return errorMap; + } + if (addEquivalentValues(buffer, tightlyCoupledBufferGroups[tcbGroupId], + producers) == -1) { + LDBG("Can not find the crossCore Buffer from another scope"); + return errorMap; + } + } + } + } + return extendedCrossCoreBuffers; +} + +int UpdateConditionInfoPass::buildIdxToVarMap( + scf::ForOp forOp, + const DenseMap>> &intraCoreBuffers, + DenseMap &idxToVar) { + int varIdx = 0; + int iterArgNum = static_cast(forOp.getNumRegionIterArgs()); + + const auto &innerDepIndices = info->innerDepConds[forOp]; + if (innerDepIndices.size() < intraCoreBuffers.size()) { + LLVM_DEBUG(llvm::dbgs() + << "Not enough inner dependency condition indices: assigned " + << innerDepIndices.size() << ", expected " + << intraCoreBuffers.size() << "\n"); + return UPDATE_CONDITION_INFO_FAILED; + } + + for (const auto &entry : intraCoreBuffers) { + int idx = entry.first; + + int argIdx = innerDepIndices[varIdx]; + if (argIdx < 0 || argIdx >= iterArgNum) { + LLVM_DEBUG(llvm::dbgs() + << "Invalid inner dependency arg index: " << argIdx + << ", iter args " << iterArgNum << "\n"); + return UPDATE_CONDITION_INFO_FAILED; + } + + idxToVar[idx] = forOp.getRegionIterArgs()[argIdx]; + LLVM_DEBUG(llvm::dbgs() << "Assign intraCore buffer group " << idx + << " to iter arg index " << argIdx << "\n"); + varIdx++; + } + + LLVM_DEBUG(llvm::dbgs() << "Assigned " << idxToVar.size() + << " intraCore condition variables.\n"); + return UPDATE_CONDITION_INFO_SUCCESS; +} + +// getInputOutputValues - Analyze the input/output buffer groups used in a +// single ifOp This function traverses all operations within ifOp, identifies +// FixpipeOp, CopyOp and MaterializeInDestinationOp, and extracts cross-core and +// intra-core buffer group indices from their operands and yield values. + +// Data structure description +// - crossCoreBuffers: {groupIdx -> (consumer -> [producers])} +// - intraCoreBuffers: {groupIdx -> (consumer -> [producers])} +// - crossCoreInputValues: list of cross-core group indices read by this ifOp +// - crossCoreOutputValues: list of cross-core group indices written by this +// ifOp +// - intraCoreInputValues: list of intra-core group indices read by this ifOp +// - intraCoreOutputValues: list of intra-core group indices written by this +// ifOp + +// Processing flow +// 1. Build reverse mapping table from Value to groupIdx +// 2. Traverse Fixpipe/Copy/MaterializeInDestination ops in ifOp, collect +// input/output groups +// 3. Traverse ifOp.thenYield() operands, collect yield output groups +// 4. Deduplicate and output four groups of index values +int UpdateConditionInfoPass::getInputOutputValues( + scf::IfOp ifOp, + DenseMap>> crossCoreBuffers, + DenseMap>> intraCoreBuffers, + SmallVector &crossCoreInputValues, + SmallVector &crossCoreOutputValues, + SmallVector &intraCoreInputValues, + SmallVector &intraCoreOutputValues) { + DenseSet crossCoreInputSet; + DenseSet crossCoreOutputSet; + DenseSet intraCoreInputSet; + DenseSet intraCoreOutputSet; + + // Get the crossCoreBufferToGroup: {bufferValue, groupIdx} + DenseMap crossCoreBufferToGroup; + DenseMap intraCoreInputToGroup; + DenseMap> intraCoreOutputToGroups; + + for (auto &entry : crossCoreBuffers) { + int groupIdx = entry.first; + for (auto &entry2 : entry.second) { + for (Value v : entry2.second) { + crossCoreBufferToGroup[v] = groupIdx; + } + } + } + + // Get the intraCoreInputToGroup: {bufferValue, groupIdx} + for (auto &entry : intraCoreBuffers) { + int groupIdx = entry.first; + for (auto &entry2 : entry.second) { + Value input = entry2.first; + SmallVector outputs = entry2.second; + intraCoreInputToGroup[input] = groupIdx; + for (Value output : outputs) { + intraCoreOutputToGroups[output].push_back(groupIdx); + } + } + } + + // collect input/output groups + ifOp.walk([&](Operation *op) { + if (op == ifOp) + return WalkResult::advance(); + + bool isFixpipeOrCopy = + dyn_cast(op) || dyn_cast(op); + bool isBufferizationWrite = + dyn_cast(op); + // Op is FixpipeOp/CopyOp/BufferizationWriteOp + // they have two operand, operand 0(ins) is input, operand 1(outs) is output + if (isFixpipeOrCopy || isBufferizationWrite) { + Value insVal = op->getOperands()[0]; + if (crossCoreBufferToGroup.count(insVal)) { + crossCoreInputSet.insert(crossCoreBufferToGroup[insVal]); + } else if (intraCoreInputToGroup.count(insVal)) { + intraCoreInputSet.insert(intraCoreInputToGroup[insVal]); + } + + Value outsVal = op->getOperands()[1]; + if (crossCoreBufferToGroup.count(outsVal)) { + crossCoreOutputSet.insert(crossCoreBufferToGroup[outsVal]); + } else if (intraCoreOutputToGroups.count(outsVal)) { + for (int idx : intraCoreOutputToGroups[outsVal]) { + intraCoreOutputSet.insert(idx); + } + } + return WalkResult::advance(); + } else { + // Op is normal, operand is input + for (Value operand : op->getOperands()) { + if (crossCoreBufferToGroup.count(operand)) + crossCoreInputSet.insert(crossCoreBufferToGroup[operand]); + if (intraCoreInputToGroup.count(operand)) + intraCoreInputSet.insert(intraCoreInputToGroup[operand]); + } + } + return WalkResult::advance(); + }); + + // operands in yield op are output + scf::YieldOp thenYield = ifOp.thenYield(); + for (Value yieldVal : thenYield.getOperands()) { + if (crossCoreBufferToGroup.count(yieldVal)) + crossCoreOutputSet.insert(crossCoreBufferToGroup[yieldVal]); + if (intraCoreOutputToGroups.count(yieldVal)) { + for (int idx : intraCoreOutputToGroups[yieldVal]) { + intraCoreOutputSet.insert(idx); + } + } + } + + crossCoreInputValues.assign(crossCoreInputSet.begin(), + crossCoreInputSet.end()); + crossCoreOutputValues.assign(crossCoreOutputSet.begin(), + crossCoreOutputSet.end()); + intraCoreInputValues.assign(intraCoreInputSet.begin(), + intraCoreInputSet.end()); + intraCoreOutputValues.assign(intraCoreOutputSet.begin(), + intraCoreOutputSet.end()); + LDBG("==== Cross Core & Intra Core Values ===="); + LLVM_DEBUG(llvm::dbgs() << "crossCoreInputValues: "); + for (int val : crossCoreInputValues) { + LLVM_DEBUG(llvm::dbgs() << val << " "); + } + LLVM_DEBUG(llvm::dbgs() << "\n"); + + LLVM_DEBUG(llvm::dbgs() << "crossCoreOutputValues: "); + for (int val : crossCoreOutputValues) { + LLVM_DEBUG(llvm::dbgs() << val << " "); + } + LLVM_DEBUG(llvm::dbgs() << "\n"); + + LLVM_DEBUG(llvm::dbgs() << "intraCoreInputValues: "); + for (int val : intraCoreInputValues) { + LLVM_DEBUG(llvm::dbgs() << val << " "); + } + LLVM_DEBUG(llvm::dbgs() << "\n"); + + LLVM_DEBUG(llvm::dbgs() << "intraCoreOutputValues: "); + for (int val : intraCoreOutputValues) { + LLVM_DEBUG(llvm::dbgs() << val << " "); + } + LLVM_DEBUG(llvm::dbgs() << "\n"); + + return UPDATE_CONDITION_INFO_SUCCESS; +} + +Value UpdateConditionInfoPass::getVarValue(scf::ForOp forOp, int varIndex) { + if (!info->innerDepConds.count(forOp)) + return Value(); + SmallVector &innerDepIndices = info->innerDepConds[forOp]; + if (varIndex < (int)innerDepIndices.size()) { + int argIdx = innerDepIndices[varIndex]; + return forOp.getRegionIterArgs()[argIdx]; + } + return Value(); +} + +// Build the information of the producer group. +int UpdateConditionInfoPass::buildOutputGroups( + SmallVector &intraCoreOutputValues, + DenseMap>> &intraCoreBuffers, + DenseMap &idxToVar, + SmallVector &outputGroups) { + outputGroups.clear(); + + for (int idx : intraCoreOutputValues) { + auto bufferIt = intraCoreBuffers.find(idx); + if (bufferIt == intraCoreBuffers.end()) { + LLVM_DEBUG(llvm::dbgs() << "Failed to build output groups: no buffer " + "entry for intraCore output group " + << idx << ".\n"); + return UPDATE_CONDITION_INFO_FAILED; + } + + auto varIt = idxToVar.find(idx); + if (varIt == idxToVar.end()) { + LLVM_DEBUG(llvm::dbgs() << "Failed to build output groups: no control " + "variable for intraCore output group " + << idx << ".\n"); + return UPDATE_CONDITION_INFO_FAILED; + } + Value var = varIt->second; + + for (auto &entry : bufferIt->second) { + SmallVector &outputs = entry.second; + if (outputs.empty()) + continue; + + bool flag = true; + for (auto &outputGroup : outputGroups) { + if (outputGroup.outputs == outputs) { + outputGroup.inputVars.push_back(var); + flag = false; + break; + } + } + if (flag) { + OutputGroupInfo groupInfo; + groupInfo.outputs = outputs; + groupInfo.inputVars.push_back(var); + outputGroups.push_back(groupInfo); + } + } + } + + LLVM_DEBUG(llvm::dbgs() << "Built " << outputGroups.size() + << " intraCore output groups.\n"); + for (size_t i = 0; i < outputGroups.size(); ++i) { + auto &group = outputGroups[i]; + LLVM_DEBUG(llvm::dbgs() << "buildOutputGroups: Input Vars (Consumer): "); + for (Value var : group.inputVars) { + LLVM_DEBUG(llvm::dbgs() << var << " "); + } + LLVM_DEBUG(llvm::dbgs() << "\n"); + LLVM_DEBUG(llvm::dbgs() << "buildOutputGroups: Output Vars (Producer): "); + + for (Value output : group.outputs) { + LLVM_DEBUG(llvm::dbgs() << output << " "); + } + + LLVM_DEBUG(llvm::dbgs() << "\n"); + } + return UPDATE_CONDITION_INFO_SUCCESS; +} + +// Select corresponding SSBuffer ptr based on ifblock running on vector/cube +// core +Value UpdateConditionInfoPass::getSSBufferPtr( + bool isAIC, int groupIdx, int ptrSetIdx, + DenseMap &VectorSSBufferPtrs, + SmallVector> ssbufferPtrs) { + if (isAIC) { + return ssbufferPtrs[ptrSetIdx][groupIdx]; + } else { + return VectorSSBufferPtrs[groupIdx]; + } +} + +// Compute pointers for VECTOR core SSBuffer +DenseMap UpdateConditionInfoPass::computeVectorSSBufferPtrs( + OpBuilder &builder, Location loc, Operation *scopeOp, + SmallVector crossCoreInputValues, + SmallVector crossCoreOutputValues) { + // Collect all unique group indices + SmallVector allGroupIndices; + DenseSet uniqueIndices; + for (int idx : crossCoreInputValues) { + if (uniqueIndices.insert(idx).second) { + allGroupIndices.push_back(idx); + } + } + for (int idx : crossCoreOutputValues) { + if (uniqueIndices.insert(idx).second) { + allGroupIndices.push_back(idx); + } + } + + DenseMap vectorSSBufferPtrs; + + builder.setInsertionPointToStart(&scopeOp->getRegion(0).front()); + int vec1Offset = 1024; + Value vec1OffsetValue = builder.create( + loc, VECTOR_SSBUF_OFFSET, ADDR_INT_TYPE); + auto subIdOp = builder.create( + loc, builder.getIntegerType(ADDR_INT_TYPE)); + Value ssbAddrOffset = + builder.create(loc, subIdOp, vec1OffsetValue); + + for (int groupIdx : allGroupIndices) { + auto ssbBaseAddr = builder.create( + loc, groupIdx * VALUE_SSBUF_OFFSET, ADDR_INT_TYPE); + auto ssbAddr = + builder.create(loc, ssbBaseAddr, ssbAddrOffset); + Value ptr = builder.create( + loc, LLVM::LLVMPointerType::get(builder.getContext(), SSBUF_ADDR_SPACE), + ssbAddr.getResult()); + vectorSSBufferPtrs[groupIdx] = ptr; + } + + return vectorSSBufferPtrs; +} + +// Part 2: Add cross-core conditions +Value UpdateConditionInfoPass::addCrossCoreConditions( + OpBuilder &builder, Location loc, SmallVector crossCoreInputValues, + SmallVector crossCoreOutputValues, + DenseMap>> &crossCoreBuffers, + bool isAIC, Value zeroConst, DenseMap &VectorSSBufferPtrs, + SmallVector> ssbufferPtrs) { + Value conditions = nullptr; + + auto combineCondition = [&](Value newCond) { + if (conditions) { + conditions = builder.create(loc, conditions, newCond); + } else { + conditions = newCond; + } + }; + + for (int inputGroupIdx : crossCoreInputValues) { + Value cond = nullptr; + if (isAIC) { + Value vec0Value = builder.create( + loc, builder.getI32Type(), + getSSBufferPtr(isAIC, inputGroupIdx, 0, VectorSSBufferPtrs, + ssbufferPtrs)); + Value vec1Value = builder.create( + loc, builder.getI32Type(), + getSSBufferPtr(isAIC, inputGroupIdx, 1, VectorSSBufferPtrs, + ssbufferPtrs)); + Value vec0Cond = builder.create( + loc, arith::CmpIPredicate::sgt, vec0Value, zeroConst); + Value vec1Cond = builder.create( + loc, arith::CmpIPredicate::sgt, vec1Value, zeroConst); + cond = builder.create(loc, vec0Cond, vec1Cond); + } else { + Value value = builder.create( + loc, builder.getI32Type(), + getSSBufferPtr(isAIC, inputGroupIdx, 0, VectorSSBufferPtrs, + ssbufferPtrs)); + cond = builder.create(loc, arith::CmpIPredicate::sgt, + value, zeroConst); + } + combineCondition(cond); + } + + for (int outputGroupIdx : crossCoreOutputValues) { + int outputCount = 0; + for (auto &entry : crossCoreBuffers[outputGroupIdx]) { + outputCount += entry.second.size(); + } + Value bufferNum = + builder.create(loc, outputCount, CONST_INT_TYPE); + Value cond = nullptr; + if (isAIC) { + Value vec0Value = builder.create( + loc, builder.getI32Type(), + getSSBufferPtr(isAIC, outputGroupIdx, 0, VectorSSBufferPtrs, + ssbufferPtrs)); + Value vec1Value = builder.create( + loc, builder.getI32Type(), + getSSBufferPtr(isAIC, outputGroupIdx, 1, VectorSSBufferPtrs, + ssbufferPtrs)); + Value vec0Cond = builder.create( + loc, arith::CmpIPredicate::slt, vec0Value, bufferNum); + Value vec1Cond = builder.create( + loc, arith::CmpIPredicate::slt, vec1Value, bufferNum); + cond = builder.create(loc, vec0Cond, vec1Cond); + } else { + Value value = builder.create( + loc, builder.getI32Type(), + getSSBufferPtr(isAIC, outputGroupIdx, 0, VectorSSBufferPtrs, + ssbufferPtrs)); + cond = builder.create(loc, arith::CmpIPredicate::slt, + value, bufferNum); + } + combineCondition(cond); + } + + return conditions; +} + +// Part 3: Update control variables in then block +void UpdateConditionInfoPass::updateCrossCoreControlVars( + OpBuilder &builder, Location loc, scf::IfOp ifOp, + SmallVector crossCoreInputValues, + SmallVector crossCoreOutputValues, bool isAIC, Value oneConst, + DenseMap &VectorSSBufferPtrs, + SmallVector> ssbufferPtrs) { + Block *thenBlock = &ifOp.getThenRegion().front(); + auto yieldOp = cast(thenBlock->getTerminator()); + builder.setInsertionPoint(yieldOp); + + for (int inputGroupIdx : crossCoreInputValues) { + if (isAIC) { + Value vec0Value = builder.create( + loc, builder.getI32Type(), + getSSBufferPtr(isAIC, inputGroupIdx, 0, VectorSSBufferPtrs, + ssbufferPtrs)); + Value vec1Value = builder.create( + loc, builder.getI32Type(), + getSSBufferPtr(isAIC, inputGroupIdx, 1, VectorSSBufferPtrs, + ssbufferPtrs)); + Value vec0NewValue = + builder.create(loc, vec0Value, oneConst); + Value vec1NewValue = + builder.create(loc, vec1Value, oneConst); + builder.create(loc, vec0NewValue, + getSSBufferPtr(isAIC, inputGroupIdx, 0, + VectorSSBufferPtrs, + ssbufferPtrs)); + builder.create(loc, vec1NewValue, + getSSBufferPtr(isAIC, inputGroupIdx, 1, + VectorSSBufferPtrs, + ssbufferPtrs)); + } else { + Value value = builder.create( + loc, builder.getI32Type(), + getSSBufferPtr(isAIC, inputGroupIdx, 0, VectorSSBufferPtrs, + ssbufferPtrs)); + Value newValue = builder.create(loc, value, oneConst); + builder.create(loc, newValue, + getSSBufferPtr(isAIC, inputGroupIdx, 0, + VectorSSBufferPtrs, + ssbufferPtrs)); + } + } + + for (int outputGroupIdx : crossCoreOutputValues) { + if (isAIC) { + Value vec0Value = builder.create( + loc, builder.getI32Type(), + getSSBufferPtr(isAIC, outputGroupIdx, 0, VectorSSBufferPtrs, + ssbufferPtrs)); + Value vec1Value = builder.create( + loc, builder.getI32Type(), + getSSBufferPtr(isAIC, outputGroupIdx, 1, VectorSSBufferPtrs, + ssbufferPtrs)); + Value vec0NewValue = + builder.create(loc, vec0Value, oneConst); + Value vec1NewValue = + builder.create(loc, vec1Value, oneConst); + builder.create(loc, vec0NewValue, + getSSBufferPtr(isAIC, outputGroupIdx, 0, + VectorSSBufferPtrs, + ssbufferPtrs)); + builder.create(loc, vec1NewValue, + getSSBufferPtr(isAIC, outputGroupIdx, 1, + VectorSSBufferPtrs, + ssbufferPtrs)); + } else { + Value value = builder.create( + loc, builder.getI32Type(), + getSSBufferPtr(isAIC, outputGroupIdx, 0, VectorSSBufferPtrs, + ssbufferPtrs)); + Value newValue = builder.create(loc, value, oneConst); + builder.create(loc, newValue, + getSSBufferPtr(isAIC, outputGroupIdx, 0, + VectorSSBufferPtrs, + ssbufferPtrs)); + } + } +} + +// Set the crossCore condition +int UpdateConditionInfoPass::setCrossCoreCondition( + SmallVector crossCoreInputValues, + SmallVector crossCoreOutputValues, + DenseMap>> &crossCoreBuffers, + scf::IfOp ifOp, SmallVector> ssbufferPtrs, + Value &crossCoreCond) { + OpBuilder builder(ifOp); + Location loc = ifOp.getLoc(); + + // ========== Part 1: Preparation ========== + // Determine whether the current ifblock is on cube or vector core + auto aiCAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + auto aivAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::VECTOR); + bool isAIC = false; + bool isAIV = false; + mlir::Operation *parentOp = ifOp->getParentOp(); + mlir::Operation *scopeOp = nullptr; + while (parentOp) { + if (dyn_cast(parentOp)) { + scopeOp = parentOp; + break; + } + parentOp = parentOp->getParentOp(); + } + if (scopeOp && scopeOp->hasAttr("hivm.tcore_type")) { + auto attr = scopeOp->getAttr("hivm.tcore_type"); + if (attr == aiCAttr) { + isAIC = true; + } else if (attr == aivAttr) { + isAIV = true; + } else { + LDBG("scope block has invalid tcore_type attribute"); + return UPDATE_CONDITION_INFO_FAILED; + } + } else { + LDBG("ifblock not in a correct scope block"); + return UPDATE_CONDITION_INFO_FAILED; + } + + Value zeroConst = + builder.create(loc, 0, CONST_INT_TYPE); + Value oneConst = builder.create(loc, 1, CONST_INT_TYPE); + + // If ifblock is on vector core, compute the required SSBuffer ptrs for vector + // side + DenseMap VectorSSBufferPtrs; + if (!isAIC) { + VectorSSBufferPtrs = computeVectorSSBufferPtrs( + builder, loc, scopeOp, crossCoreInputValues, crossCoreOutputValues); + } + + builder.setInsertionPoint(ifOp); + + // ========== Part 2: Add cross-core conditions ========== + crossCoreCond = addCrossCoreConditions( + builder, loc, crossCoreInputValues, crossCoreOutputValues, + crossCoreBuffers, isAIC, zeroConst, VectorSSBufferPtrs, ssbufferPtrs); + + // ========== Part 3: Update control variables ========== + updateCrossCoreControlVars(builder, loc, ifOp, crossCoreInputValues, + crossCoreOutputValues, isAIC, oneConst, + VectorSSBufferPtrs, ssbufferPtrs); + + return UPDATE_CONDITION_INFO_SUCCESS; +} + +// Collect the conditions for intra-core consumer values. +void UpdateConditionInfoPass::collectIntraCoreInputConditions( + OpBuilder &builder, Location loc, SmallVector &intraCoreInputValues, + DenseMap &idxToVar, SmallVector &conditions, + DenseSet &usedVarsSet, + DenseMap &varUpdateTypes) { + if (intraCoreInputValues.empty()) { + LLVM_DEBUG(llvm::dbgs() << "No intraCore input conditions to collect.\n"); + return; + } + + size_t beforeConditionNum = conditions.size(); + Value zeroConst = + builder.create(loc, 0, CONST_INT_TYPE); + for (int idx : intraCoreInputValues) { + auto varIt = idxToVar.find(idx); + if (varIt == idxToVar.end()) { + LLVM_DEBUG(llvm::dbgs() << "Skip intraCore input group " << idx + << ": no control variable.\n"); + continue; + } + + Value var = varIt->second; + Value varToUse = var; + auto latestIt = controlVarToLatestValue.find(var); + if (latestIt != controlVarToLatestValue.end()) { + varToUse = latestIt->second; + } + + Value cond = builder.create(loc, arith::CmpIPredicate::sgt, + varToUse, zeroConst); + conditions.push_back(cond); + usedVarsSet.insert(var); + varUpdateTypes[var] = VarUpdateType::DEC; + LLVM_DEBUG(llvm::dbgs() + << "Add intraCore input condition for group " << idx << ".\n"); + } + LLVM_DEBUG(llvm::dbgs() << "Collected " + << (conditions.size() - beforeConditionNum) + << " intraCore input conditions.\n"); +} + +// Collect the conditions for intra-core producer values. +int UpdateConditionInfoPass::collectIntraCoreOutputConditions( + OpBuilder &builder, Location loc, + DenseMap>> &intraCoreBuffers, + SmallVector &intraCoreOutputValues, DenseMap &idxToVar, + SmallVector &conditions, DenseSet &usedVarsSet, + DenseMap &varUpdateTypes) { + if (intraCoreOutputValues.empty()) { + LLVM_DEBUG(llvm::dbgs() << "No intraCore output conditions to collect.\n"); + return UPDATE_CONDITION_INFO_SUCCESS; + } + + size_t beforeConditionNum = conditions.size(); + SmallVector outputGroups; + if (buildOutputGroups(intraCoreOutputValues, intraCoreBuffers, idxToVar, + outputGroups) == UPDATE_CONDITION_INFO_FAILED) { + return UPDATE_CONDITION_INFO_FAILED; + } + for (auto &group : outputGroups) { + int size = group.outputs.size(); + Value limitVal = + builder.create(loc, size, CONST_INT_TYPE); + for (Value var : group.inputVars) { + Value varToUse = var; + auto latestIt = controlVarToLatestValue.find(var); + if (latestIt != controlVarToLatestValue.end()) { + varToUse = latestIt->second; + } + + Value cond = builder.create(loc, arith::CmpIPredicate::slt, + varToUse, limitVal); + conditions.push_back(cond); + usedVarsSet.insert(var); + varUpdateTypes[var] = VarUpdateType::INC; + LLVM_DEBUG(llvm::dbgs() + << "Add intraCore output condition with producer limit " + << size << ".\n"); + } + } + LLVM_DEBUG(llvm::dbgs() << "Collected " + << (conditions.size() - beforeConditionNum) + << " intraCore output conditions.\n"); + return UPDATE_CONDITION_INFO_SUCCESS; +} + +// Set the intraCore condition. +int UpdateConditionInfoPass::setIntraCoreCondition( + ModuleOp module, scf::IfOp ifOp, + DenseMap>> &intraCoreBuffers, + SmallVector &intraCoreInputValues, + SmallVector &intraCoreOutputValues, DenseMap &idxToVar, + DenseMap &varUpdateTypes, Value &intraCoreCond) { + LDBG("Enter set intraCore condition."); + intraCoreCond = Value(); + OpBuilder builder(ifOp.getContext()); + builder.setInsertionPoint(ifOp); + Location loc = ifOp.getLoc(); + + SmallVector conditions; + DenseSet usedVarsSet; + LLVM_DEBUG(llvm::dbgs() << "Collect intraCore conditions: inputs " + << intraCoreInputValues.size() << ", outputs " + << intraCoreOutputValues.size() << "\n"); + // Collect the conditions for intra-core consumer values. + collectIntraCoreInputConditions(builder, loc, intraCoreInputValues, idxToVar, + conditions, usedVarsSet, varUpdateTypes); + // Collect the conditions for intra-core producer values. + if (collectIntraCoreOutputConditions( + builder, loc, intraCoreBuffers, intraCoreOutputValues, idxToVar, + conditions, usedVarsSet, + varUpdateTypes) == UPDATE_CONDITION_INFO_FAILED) { + return UPDATE_CONDITION_INFO_FAILED; + } + + if (!conditions.empty()) { + intraCoreCond = conditions[0]; + for (size_t i = 1; i < conditions.size(); ++i) { + intraCoreCond = + builder.create(loc, intraCoreCond, conditions[i]); + } + } + + currentUsedVars.clear(); + for (Value var : usedVarsSet) { + currentUsedVars.push_back(var); + } + LLVM_DEBUG(llvm::dbgs() << "Built " << conditions.size() + << " intraCore conditions using " + << currentUsedVars.size() << " control variables.\n"); + + LDBG("Exit set intraCore condition."); + return UPDATE_CONDITION_INFO_SUCCESS; +} + +// Update the mapping of control variables to their latest values +void UpdateConditionInfoPass::updateControlVarToLatestValue(scf::IfOp newIfOp, + scf::IfOp oldIfOp, + bool hasCounter, + Value counter) { + if (currentUsedVars.empty() && !hasCounter) { + LLVM_DEBUG(llvm::dbgs() + << "No control variable latest values to update.\n"); + return; + } + + size_t origResultCount = oldIfOp.getNumResults(); + + for (size_t i = 0; i < currentUsedVars.size(); ++i) { + Value var = currentUsedVars[i]; + Value newValue = newIfOp.getResult(origResultCount + i); + controlVarToLatestValue[var] = newValue; + LLVM_DEBUG(llvm::dbgs() + << "Record latest intraCore control value at result index " + << (origResultCount + i) << ".\n"); + } + + if (hasCounter) { + size_t counterResultIdx = origResultCount + currentUsedVars.size(); + Value newCounterValue = newIfOp.getResult(counterResultIdx); + controlVarToLatestValue[counter] = newCounterValue; + LLVM_DEBUG(llvm::dbgs() << "Record latest counter value at result index " + << counterResultIdx << ".\n"); + } + LLVM_DEBUG(llvm::dbgs() << "[DEBUG] controlVarToLatestValue size: " + << controlVarToLatestValue.size() << "\n"); + + for (auto &entry : controlVarToLatestValue) { + LLVM_DEBUG(llvm::dbgs() << "[DEBUG] key = " << entry.first + << " --> new value = " << entry.second << "\n"); + } +} + +// Update the yield in the forOp +int UpdateConditionInfoPass::updateForOpYield(scf::ForOp forOp) { + LDBG("Enter update forOp yield "); + if (controlVarToLatestValue.empty()) { + LLVM_DEBUG(llvm::dbgs() << "Failed to update forOp yield: no latest " + "control variable values.\n"); + return UPDATE_CONDITION_INFO_FAILED; + } + + Location loc = forOp.getLoc(); + Block *forBody = forOp.getBody(); + auto yieldOp = dyn_cast(forBody->getTerminator()); + if (!yieldOp) { + LLVM_DEBUG( + llvm::dbgs() + << "Failed to update forOp yield: terminator is not scf.yield.\n"); + return UPDATE_CONDITION_INFO_FAILED; + } + + SmallVector newYieldOperands(yieldOp.getOperands().begin(), + yieldOp.getOperands().end()); + if (newYieldOperands.size() != forOp.getNumRegionIterArgs()) { + LLVM_DEBUG(llvm::dbgs() << "Failed to update forOp yield: yield operands " + << newYieldOperands.size() << ", iter args " + << forOp.getNumRegionIterArgs() << "\n"); + return UPDATE_CONDITION_INFO_FAILED; + } + + DenseMap iterArgToIndex; + for (unsigned j = 0; j < forOp.getNumRegionIterArgs(); ++j) { + iterArgToIndex[forOp.getRegionIterArgs()[j]] = j; + } + + for (auto &entry : controlVarToLatestValue) { + Value origVar = entry.first; + Value latestValue = entry.second; + auto it = iterArgToIndex.find(origVar); + if (it == iterArgToIndex.end()) { + LLVM_DEBUG(llvm::dbgs() << "Failed to update forOp yield: control " + "variable is not a region iter arg.\n"); + return UPDATE_CONDITION_INFO_FAILED; + } + newYieldOperands[it->second] = latestValue; + LLVM_DEBUG(llvm::dbgs() + << "Update forOp yield operand index " << it->second << "\n"); + } + + OpBuilder yieldBuilder(yieldOp); + yieldBuilder.create(loc, newYieldOperands); + yieldOp.erase(); + LLVM_DEBUG(llvm::dbgs() << "Updated forOp yield with " + << controlVarToLatestValue.size() + << " latest control values.\n"); + LDBG("Exit update forOp yield "); + return UPDATE_CONDITION_INFO_SUCCESS; +} + +SmallVector +UpdateConditionInfoPass::buildNewIfResultTypes(scf::IfOp oldIfOp, + bool hasCounter, Value counter) { + SmallVector resultTypes; + for (Value result : oldIfOp.getResults()) { + resultTypes.push_back(result.getType()); + } + for (Value var : currentUsedVars) { + resultTypes.push_back(var.getType()); + } + if (hasCounter) { + resultTypes.push_back(counter.getType()); + } + LLVM_DEBUG(llvm::dbgs() << "Build new if result types: old results " + << oldIfOp.getNumResults() << ", control vars " + << currentUsedVars.size() << ", has counter " + << hasCounter << ".\n"); + return resultTypes; +} + +void UpdateConditionInfoPass::collectYieldOperands( + Block &block, Operation *&yieldOp, SmallVector &yieldOperands) { + yieldOp = nullptr; + yieldOperands.clear(); + if (block.empty()) { + LLVM_DEBUG(llvm::dbgs() << "Collect yield operands: block is empty.\n"); + return; + } + + Operation *lastOp = &block.back(); + if (!isa(lastOp)) { + LLVM_DEBUG( + llvm::dbgs() + << "Collect yield operands: block terminator is not scf.yield.\n"); + return; + } + + yieldOp = lastOp; + auto scfYieldOp = cast(lastOp); + yieldOperands.assign(scfYieldOp.getOperands().begin(), + scfYieldOp.getOperands().end()); + LLVM_DEBUG(llvm::dbgs() << "Collected " << yieldOperands.size() + << " yield operands.\n"); +} + +void UpdateConditionInfoPass::populateNewThenBlock( + scf::IfOp newIfOp, Block &oldThenBlock, Operation *oldThenYieldOp, + ArrayRef oldYieldOperands, + DenseMap &varUpdateTypes, bool hasCounter, + Value counter, Value step) { + Location loc = newIfOp.getLoc(); + Block &newThenBlock = newIfOp.getThenRegion().front(); + for (Operation &op : llvm::make_early_inc_range(oldThenBlock)) { + if (&op != oldThenYieldOp) { + op.moveBefore(&newThenBlock, newThenBlock.end()); + } + } + LLVM_DEBUG(llvm::dbgs() << "Populate new then block with " + << oldYieldOperands.size() + << " original yield operands.\n"); + + OpBuilder thenBuilder(&newThenBlock, newThenBlock.end()); + SmallVector thenYieldOperands(oldYieldOperands.begin(), + oldYieldOperands.end()); + if (!currentUsedVars.empty()) { + Value one = + thenBuilder.create(loc, 1, CONST_INT_TYPE); + for (Value var : currentUsedVars) { + Value varToUse = var; + auto latestIt = controlVarToLatestValue.find(var); + if (latestIt != controlVarToLatestValue.end()) { + varToUse = latestIt->second; + } + + Value yieldVal = varToUse; + auto it = varUpdateTypes.find(var); + if (it != varUpdateTypes.end()) { + if (it->second == VarUpdateType::DEC) { + yieldVal = thenBuilder.create(loc, varToUse, one); + } else if (it->second == VarUpdateType::INC) { + yieldVal = thenBuilder.create(loc, varToUse, one); + } + } + thenYieldOperands.push_back(yieldVal); + } + } + + if (hasCounter) { + Value newCounter = thenBuilder.create(loc, counter, step); + thenYieldOperands.push_back(newCounter); + LLVM_DEBUG(llvm::dbgs() << "Append updated counter to then yield.\n"); + } + + LLVM_DEBUG(llvm::dbgs() << "Create then yield with " + << thenYieldOperands.size() << " operands.\n"); + thenBuilder.create(loc, thenYieldOperands); +} + +void UpdateConditionInfoPass::populateNewElseBlock( + scf::IfOp newIfOp, scf::IfOp oldIfOp, bool needsYield, bool oldHasElse, + bool hasCounter, Value counter) { + if (!needsYield && !oldHasElse) { + LLVM_DEBUG(llvm::dbgs() << "Skip populating else block: no yield needed " + "and old if has no else.\n"); + return; + } + + Location loc = newIfOp.getLoc(); + Block &newElseBlock = newIfOp.getElseRegion().front(); + SmallVector oldElseYieldOperands; + Operation *oldElseYieldOp = nullptr; + + if (oldHasElse) { + Block &oldElseBlock = oldIfOp.getElseRegion().front(); + collectYieldOperands(oldElseBlock, oldElseYieldOp, oldElseYieldOperands); + for (Operation &op : llvm::make_early_inc_range(oldElseBlock)) { + if (&op != oldElseYieldOp) { + op.moveBefore(&newElseBlock, newElseBlock.end()); + } + } + LLVM_DEBUG(llvm::dbgs() + << "Moved old else block ops and collected " + << oldElseYieldOperands.size() << " old else yield operands.\n"); + } + + if (needsYield) { + OpBuilder elseBuilder(&newElseBlock, newElseBlock.end()); + SmallVector elseYieldOperands; + for (Value operand : oldElseYieldOperands) { + Value newOperand = operand; + auto it = controlVarToLatestValue.find(operand); + if (it != controlVarToLatestValue.end()) { + newOperand = it->second; + } + elseYieldOperands.push_back(newOperand); + } + + for (Value var : currentUsedVars) { + Value varToUse = var; + auto it = controlVarToLatestValue.find(var); + if (it != controlVarToLatestValue.end()) { + varToUse = it->second; + } + elseYieldOperands.push_back(varToUse); + } + + if (hasCounter) { + Value counterToUse = counter; + auto it = controlVarToLatestValue.find(counter); + if (it != controlVarToLatestValue.end()) { + counterToUse = it->second; + } + elseYieldOperands.push_back(counterToUse); + } + + LLVM_DEBUG(llvm::dbgs() << "Create else yield with " + << elseYieldOperands.size() << " operands.\n"); + elseBuilder.create(loc, elseYieldOperands); + } else if (oldElseYieldOp) { + oldElseYieldOp->erase(); + LLVM_DEBUG( + llvm::dbgs() + << "Erase old else yield because new if does not need yield values.\n"); + } +} + +// Create new IfOp with new then and else blocks. +scf::IfOp UpdateConditionInfoPass::createNewIfOpWithBlocks( + scf::IfOp oldIfOp, Value combinedCond, + DenseMap &varUpdateTypes, bool hasCounter, + Value counter, Value step) { + Location loc = oldIfOp.getLoc(); + OpBuilder builder(oldIfOp); + + bool needsYield = !currentUsedVars.empty() || hasCounter; + bool oldHasElse = oldIfOp.getElseRegion().hasOneBlock(); + LLVM_DEBUG(llvm::dbgs() << "Create replacement if op: needs yield " + << needsYield << ", old has else " << oldHasElse + << ", current used vars " << currentUsedVars.size() + << ".\n"); + + Block &oldThenBlock = oldIfOp.getThenRegion().front(); + Operation *oldThenYieldOp = nullptr; + SmallVector oldYieldOperands; + collectYieldOperands(oldThenBlock, oldThenYieldOp, oldYieldOperands); + SmallVector resultTypes = + buildNewIfResultTypes(oldIfOp, hasCounter, counter); + scf::IfOp newIfOp = + builder.create(loc, resultTypes, combinedCond, true); + LLVM_DEBUG(llvm::dbgs() << "Created replacement if op with " + << resultTypes.size() << " results.\n"); + + for (auto &attr : oldIfOp->getAttrs()) { + newIfOp->setAttr(attr.getName(), attr.getValue()); + } + + populateNewThenBlock(newIfOp, oldThenBlock, oldThenYieldOp, oldYieldOperands, + varUpdateTypes, hasCounter, counter, step); + populateNewElseBlock(newIfOp, oldIfOp, needsYield, oldHasElse, hasCounter, + counter); + + for (size_t i = 0; i < oldIfOp.getNumResults(); ++i) { + oldIfOp.getResult(i).replaceAllUsesWith(newIfOp.getResult(i)); + } + LLVM_DEBUG(llvm::dbgs() << "Replaced " << oldIfOp.getNumResults() + << " old if results.\n"); + + return newIfOp; +} + +// Combine the three conditions: crossCore condition + intraCore condition + +// counter condition +int UpdateConditionInfoPass::combineConditions( + ModuleOp module, Value crossCoreCond, Value intraCoreCond, scf::IfOp ifOp, + scf::ForOp forOp, size_t &usedCounterNum, + DenseMap &varUpdateTypes) { + Location loc = ifOp.getLoc(); + SmallVector validConditions; + Value counter; + bool hasCounter = false; + + if (crossCoreCond) { + validConditions.push_back(crossCoreCond); + } + if (intraCoreCond) { + validConditions.push_back(intraCoreCond); + } + + if (!info->blockCounters.count(forOp)) { + LLVM_DEBUG(llvm::dbgs() << "Missing block counters for forOp.\n"); + return UPDATE_CONDITION_INFO_FAILED; + } + + SmallVector &counterIndices = info->blockCounters[forOp]; + + if (info->cntArgs.count(ifOp)) { + counter = info->cntArgs[ifOp]; + hasCounter = true; + } else { + if (usedCounterNum >= counterIndices.size()) { + LLVM_DEBUG(llvm::dbgs() + << "Not enough counters for ssbuffer if ops: used " + << usedCounterNum << ", counters " << counterIndices.size() + << "\n"); + return UPDATE_CONDITION_INFO_FAILED; + } + + int argIdx = counterIndices[usedCounterNum]; + int iterArgNum = static_cast(forOp.getNumRegionIterArgs()); + if (argIdx < 0 || argIdx >= iterArgNum) { + LLVM_DEBUG(llvm::dbgs() << "Invalid counter arg index: " << argIdx + << ", iter args " << iterArgNum << "\n"); + return UPDATE_CONDITION_INFO_FAILED; + } + + counter = forOp.getRegionIterArgs()[argIdx]; + hasCounter = true; + info->cntArgs[ifOp] = counter; + usedCounterNum++; + LLVM_DEBUG(llvm::dbgs() << "Assign counter iter arg index " << argIdx + << " to ssbuffer if op.\n"); + } + + LLVM_DEBUG(llvm::dbgs() << "this ifop used counter is: " << counter << "\n"); + if (hasCounter) { + OpBuilder builder(ifOp); + Value upperBound = forOp.getUpperBound(); + Value counterToUse = counter; + auto latestIt = controlVarToLatestValue.find(counter); + if (latestIt != controlVarToLatestValue.end()) { + counterToUse = latestIt->second; + } + Value counterCond = builder.create( + loc, arith::CmpIPredicate::slt, counterToUse, upperBound); + validConditions.push_back(counterCond); + } + + if (validConditions.empty()) { + LLVM_DEBUG(llvm::dbgs() + << "Failed to build any condition for ssbuffer if op.\n"); + return UPDATE_CONDITION_INFO_FAILED; + } + + LLVM_DEBUG(llvm::dbgs() << "Combine " << validConditions.size() + << " conditions for ssbuffer if op.\n"); + OpBuilder builder(ifOp); + Value combinedCond = validConditions[0]; + for (size_t i = 1; i < validConditions.size(); ++i) { + combinedCond = + builder.create(loc, combinedCond, validConditions[i]); + } + + scf::IfOp newIfOp = createNewIfOpWithBlocks( + ifOp, combinedCond, varUpdateTypes, hasCounter, counter, forOp.getStep()); + + if (hasCounter) { + info->cntArgs.erase(ifOp); + info->cntArgs[newIfOp] = counter; + } + + updateControlVarToLatestValue(newIfOp, ifOp, hasCounter, counter); + + ifOp.erase(); + return UPDATE_CONDITION_INFO_SUCCESS; +} + +// Update the conditions of ifOp. +int UpdateConditionInfoPass::updateIfConds( + ModuleOp module, SmallVector> ssbufferPtrs) { + // Walk the forOp in the module to update the conditions of ifOp + SmallVector mainLoopForOps; + WalkResult walkResult = module.walk([&](Operation *op) -> WalkResult { + if (!op->hasAttr(SSBUFFER_Main_LOOP)) { + return WalkResult::advance(); + } + + auto forOp = dyn_cast(op); + if (!forOp) { + LLVM_DEBUG(llvm::dbgs() << "Found unsupported main loop op: " + << op->getName() << "\n"); + return WalkResult::interrupt(); + } + mainLoopForOps.push_back(forOp); + return WalkResult::advance(); + }); + if (walkResult.wasInterrupted()) { + return UPDATE_CONDITION_INFO_FAILED; + } + for (scf::ForOp forOp : mainLoopForOps) { + controlVarToLatestValue.clear(); + + DenseMap>> crossCoreBuffers; + DenseMap>> intraCoreBuffers; + // Step1:Collect the dependency buffer info of this forOp + collectDependencyBuffers(forOp, crossCoreBuffers, intraCoreBuffers); + if (crossCoreBuffers.empty() && intraCoreBuffers.empty()) { + LDBG("crossCoreBuffers and intraCoreBuffers are both empty!"); + return UPDATE_CONDITION_INFO_FAILED; + } + + DenseMap>> + extendedCrossCoreBuffers = extendCrossCoreBuffersWithEquivalentValues( + module, crossCoreBuffers); + if (extendedCrossCoreBuffers.count(-1)) { + LDBG("extendCrossCoreBuffersWithEquivalentValues failed!"); + return UPDATE_CONDITION_INFO_FAILED; + } + // Step2:Assign a variable to each inputValue of this forOp + DenseMap idxToVar; + if (buildIdxToVarMap(forOp, intraCoreBuffers, idxToVar) == + UPDATE_CONDITION_INFO_FAILED) { + return UPDATE_CONDITION_INFO_FAILED; + } + size_t usedCounterNum = 0; + SmallVector ifOps; + WalkResult ifWalkResult = forOp.walk([&](Operation *op) -> WalkResult { + if (!op->hasAttr(SSBUFFER_IF)) { + return WalkResult::advance(); + } + + auto ifOp = dyn_cast(op); + if (!ifOp) { + LLVM_DEBUG(llvm::dbgs() << "Found unsupported ssbuffer if op: " + << op->getName() << "\n"); + return WalkResult::interrupt(); + } + + ifOps.push_back(ifOp); + return WalkResult::advance(); + }); + if (ifWalkResult.wasInterrupted()) { + return UPDATE_CONDITION_INFO_FAILED; + } + auto counterIt = info->blockCounters.find(forOp); + if (counterIt == info->blockCounters.end()) { + LLVM_DEBUG(llvm::dbgs() << "Failed to assign counters for ssbuffer if " + "ops: no counters for forOp.\n"); + return UPDATE_CONDITION_INFO_FAILED; + } + + size_t counterNum = counterIt->second.size(); + if (ifOps.size() > counterNum) { + LLVM_DEBUG(llvm::dbgs() + << "Failed to assign counters for all ssbuffer if ops: if ops " + << ifOps.size() << ", counters " << counterNum << "\n"); + return UPDATE_CONDITION_INFO_FAILED; + } + // Update the conditions of ifOp in this forOp. + for (scf::IfOp ifOp : ifOps) { + // Walk the ifOp in this forOp to update the conditions of ifOp + SmallVector crossCoreInputValues; + SmallVector crossCoreOutputValues; + SmallVector intraCoreInputValues; + SmallVector intraCoreOutputValues; + + if (getInputOutputValues(ifOp, extendedCrossCoreBuffers, intraCoreBuffers, + crossCoreInputValues, crossCoreOutputValues, + intraCoreInputValues, + intraCoreOutputValues) != 0) { + LDBG("getInputOutputValues failed!"); + return UPDATE_CONDITION_INFO_FAILED; + } + + // Step3:Set the crossCore condition + Value crossCoreCond; + if (setCrossCoreCondition(crossCoreInputValues, crossCoreOutputValues, + crossCoreBuffers, ifOp, ssbufferPtrs, + crossCoreCond) != 0) { + LDBG("setCrossCoreCondition failed!"); + return UPDATE_CONDITION_INFO_FAILED; + } + // Step4:Set the intraCore condition + DenseMap varUpdateTypes; + Value intraCoreCond; + if (setIntraCoreCondition(module, ifOp, intraCoreBuffers, + intraCoreInputValues, intraCoreOutputValues, + idxToVar, varUpdateTypes, intraCoreCond) == + UPDATE_CONDITION_INFO_FAILED) { + return UPDATE_CONDITION_INFO_FAILED; + } + // Step5:Combine the three conditions: crossCore condition + intraCore + // condition + counter condition + if (combineConditions(module, crossCoreCond, intraCoreCond, ifOp, forOp, + usedCounterNum, + varUpdateTypes) == UPDATE_CONDITION_INFO_FAILED) { + return UPDATE_CONDITION_INFO_FAILED; + } + } + // Step6:Update the yield variable of the forOp + if (updateForOpYield(forOp) == UPDATE_CONDITION_INFO_FAILED) { + return UPDATE_CONDITION_INFO_FAILED; + } + } + return UPDATE_CONDITION_INFO_SUCCESS; +} + +void UpdateConditionInfoPass::runOnOperation() { + ModuleOp module = getOperation(); + + LDBG("Enter UpdateConditionInfo pass.\n"); + // Step1:Init the ssbufferPtrs + SmallVector> ssbufferPtrs = allocSSBuffer(module); + + // Step2:Update the conditions of ifOp based on the intraCoreDependentMap and + // crossCoreDependentMap + int updateResult = updateIfConds(module, ssbufferPtrs); + if (updateResult == UPDATE_CONDITION_INFO_FAILED) { + signalPassFailure(); + return; + } + LDBG("Exit UpdateConditionInfo pass.\n"); +} + +namespace mlir { +namespace triton { +std::unique_ptr> createUpdateConditionInfoPass() { + return std::make_unique(); +} +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.cpp new file mode 100644 index 0000000000..354ef38f8f --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.cpp @@ -0,0 +1,456 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateForOps.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/Scope/IR/Scope.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/Support/Debug.h" + +static constexpr const char *DEBUG_TYPE = "UpdateForOps"; +static constexpr int kPipeSFlagId = 15; +#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ") +#define LDBG(...) \ + LLVM_DEBUG({ \ + DBGS(); \ + llvm::outs() << __VA_ARGS__; \ + llvm::outs() << "\n"; \ + }) + +using namespace llvm; +using namespace mlir; +using namespace triton; +using namespace hivm; + +// Replace old block arguments with new ones +static LogicalResult replaceBlockArguments(Block *oldBlock, Block *newBlock) { + if (!oldBlock || !newBlock) { + LDBG("[Error]: oldBlock or newBlock is null\n"); + return failure(); + } + + unsigned totalArgs = oldBlock->getNumArguments(); + + for (unsigned i = 0; i < totalArgs; ++i) { + oldBlock->getArgument(i).replaceAllUsesWith(newBlock->getArgument(i)); + } + + return success(); +} + +// Collect forOps that need processing based on info +static SmallVector +collectForOpsToProcess(ModuleOp module, + const llvm::DenseMap &numInfo) { + SmallVector forOps; + + module.walk([&](scf::ForOp forOp) { + if (numInfo.count(forOp)) { + forOps.push_back(forOp); + } + }); + + return forOps; +} + +// Create new yield operands: original yield ops + extra args from new block +static SmallVector createNewYieldOperands(scf::YieldOp oldYield, + unsigned oldNumArgs, + Block *newBlock, + int numExtraArgs) { + SmallVector newYieldOperands; + + for (unsigned i = 0; i < oldNumArgs; ++i) { + newYieldOperands.push_back(oldYield.getOperand(i)); + } + + for (int i = 0; i < numExtraArgs; ++i) { + newYieldOperands.push_back(newBlock->getArgument(1 + oldNumArgs + i)); + } + + return newYieldOperands; +} + +// Derive block counters from ssbuffer.if attributes when info is not +// pre-populated +LogicalResult +UpdateForOpsPass::deriveBlockCountersFromIfOps(ModuleOp module, + ControlFlowConditionInfo *info) { + if (!info) { + LDBG("[Error]: info is null\n"); + return failure(); + } + + module.walk([&](Operation *op) -> WalkResult { + if (!op->hasAttr("ssbuffer.main_loop")) { + return WalkResult::advance(); + } + auto forOp = dyn_cast(op); + if (!forOp) { + LDBG("[Error]: op with ssbuffer.main_loop is not a scf::ForOp\n"); + return WalkResult::interrupt(); + } + + // Collect unique ssbuffer.if values in this for op's body + llvm::DenseSet ifBlockIds; + + forOp.walk([&](Operation *innerOp) { + if (auto ifAttr = innerOp->getAttrOfType("ssbuffer.if")) { + ifBlockIds.insert(ifAttr.getInt()); + } + }); + + if (!ifBlockIds.empty()) { + info->blockCounterNums[forOp] = ifBlockIds.size(); + } + return WalkResult::advance(); + }); + + return success(); +} + +// Create new for op with extra iter args and migrate body +static scf::ForOp +createForOpAndMigrateBody(scf::ForOp oldForOp, int numExtraArgs, + const SmallVector &extraInitArgs) { + if (numExtraArgs < 0) { + LDBG("[Error]: invalid numExtraArgs " << numExtraArgs << "\n"); + return scf::ForOp(); + } + if (numExtraArgs == 0) + return oldForOp; + if ((int)extraInitArgs.size() != numExtraArgs) { + LDBG("[Error]: extraInitArgs size " << extraInitArgs.size() + << " != numExtraArgs " << numExtraArgs + << "\n"); + return scf::ForOp(); + } + + OpBuilder builder(oldForOp); + // Create new for op with extra iter args + SmallVector newInitArgs(oldForOp.getInitArgs().begin(), + oldForOp.getInitArgs().end()); + llvm::append_range(newInitArgs, extraInitArgs); + + scf::ForOp newForOp = builder.create( + oldForOp.getLoc(), oldForOp.getLowerBound(), oldForOp.getUpperBound(), + oldForOp.getStep(), newInitArgs); + + for (auto &attr : oldForOp->getAttrs()) + newForOp->setAttr(attr.getName(), attr.getValue()); + + // Migrate body + Block *oldBlock = oldForOp.getBody(); + Block *newBlock = newForOp.getBody(); + + if (failed(replaceBlockArguments(oldBlock, newBlock))) + return scf::ForOp(); + + for (Operation &op : + llvm::make_early_inc_range(oldBlock->without_terminator())) + op.moveBefore(newBlock, newBlock->end()); + + auto oldYield = cast(oldBlock->getTerminator()); + SmallVector newYieldOperands = createNewYieldOperands( + oldYield, oldForOp.getNumRegionIterArgs(), newBlock, numExtraArgs); + + builder.setInsertionPointToEnd(newBlock); + builder.create(newForOp.getLoc(), newYieldOperands); + oldYield.erase(); + + return newForOp; +} + +static LogicalResult replaceForOpUsesAndErase(scf::ForOp oldForOp, + scf::ForOp newForOp) { + if (oldForOp.getNumResults() > 0) { + SmallVector newResults; + for (unsigned i = 0; i < oldForOp.getNumResults(); ++i) { + if (oldForOp.getResult(i).getType() != newForOp.getResult(i).getType()) { + LDBG("[Error]: result type mismatch at index " << i << "\n"); + return failure(); + } + newResults.push_back(newForOp.getResult(i)); + } + oldForOp.replaceAllUsesWith(newResults); + } + + oldForOp.erase(); + return success(); +} + +LogicalResult extendForOpWithExtraArgs(scf::ForOp oldForOp, + ControlFlowConditionInfo *info) { + int numBlockCounters = info->blockCounterNums[oldForOp]; + int numInnerDepConds = info->intraCoreDependentMap[oldForOp].size(); + int totalExtraArgs = numBlockCounters + numInnerDepConds; + if (totalExtraArgs == 0) { + return success(); + } + + OpBuilder builder(oldForOp); + SmallVector extraInitArgs; + for (int i = 0; i < numBlockCounters; ++i) + extraInitArgs.push_back(oldForOp.getLowerBound()); + for (int i = 0; i < numInnerDepConds; ++i) + extraInitArgs.push_back(builder.create( + oldForOp.getLoc(), builder.getI32Type(), builder.getI32IntegerAttr(0))); + + scf::ForOp newForOp = + createForOpAndMigrateBody(oldForOp, totalExtraArgs, extraInitArgs); + if (!newForOp) { + return failure(); + } + + unsigned baseIdx = oldForOp.getNumRegionIterArgs(); + if (numBlockCounters > 0) { + SmallVector indices; + for (int j = 0; j < numBlockCounters; ++j) + indices.push_back(baseIdx + j); + info->blockCounters.erase(oldForOp); + info->blockCounters[newForOp] = indices; + } + + if (numInnerDepConds > 0) { + SmallVector indices; + for (int j = 0; j < numInnerDepConds; ++j) + indices.push_back(baseIdx + numBlockCounters + j); + info->innerDepConds.erase(oldForOp); + info->innerDepConds[newForOp] = indices; + } + + if (info->intraCoreDependentMap.count(oldForOp)) { + info->intraCoreDependentMap[newForOp] = + info->intraCoreDependentMap[oldForOp]; + info->intraCoreDependentMap.erase(oldForOp); + } + + return replaceForOpUsesAndErase(oldForOp, newForOp); +} + +// Add block counter and inner dependency condition iter args to for ops +LogicalResult UpdateForOpsPass::addBlockCountersAndInnerDepConds( + ModuleOp module, ControlFlowConditionInfo *info) { + llvm::DenseSet allForOps; + for (auto &p : info->blockCounterNums) { + if (p.second < 0) { + LDBG("[Error]: invalid blockCounterNum " << p.second << "\n"); + return failure(); + } + allForOps.insert(p.first); + } + for (auto &entry : info->intraCoreDependentMap) + allForOps.insert(entry.first); + + for (scf::ForOp oldForOp : allForOps) { + if (failed(extendForOpWithExtraArgs(oldForOp, info))) + return failure(); + } + return success(); +} + +// Insert sync ops for a single forOp +static WalkResult insertSyncOpsForForOp(scf::ForOp forOp, Block *forBody, + hivm::TCoreTypeAttr coreType, + PipeAttr setPipe, PipeAttr waitPipe, + int waitFlagId, int setFlagId) { + Operation *forTerminator = forBody->getTerminator(); + if (!forTerminator) { + return WalkResult::interrupt(); + } + + Location loc = forOp.getLoc(); + + // Insert wait at for loop start + OpBuilder insertionBuilder(&forBody->front()); + auto waitFlagAttr = insertionBuilder.getIntegerAttr( + insertionBuilder.getI64Type(), waitFlagId); + insertionBuilder.create(loc, coreType, setPipe, waitPipe, + waitFlagAttr); + + // Insert set before yield + OpBuilder setBuilder(forTerminator); + auto setFlagAttr = + setBuilder.getIntegerAttr(setBuilder.getI64Type(), setFlagId); + setBuilder.setInsertionPoint(forTerminator); + setBuilder.create(loc, coreType, setPipe, waitPipe, + setFlagAttr); + + return WalkResult::advance(); +} + +// Insert sync ops for a single scopeOp +static WalkResult insertSyncOpsForCube(scope::ScopeOp scopeOp, + hivm::TCoreTypeAttr coreType, + PipeAttr setPipe, PipeAttr waitPipe, + int waitFlagId, int setFlagId) { + Block &scopeBlock = scopeOp.getRegion().front(); + Operation *scopeTerminator = scopeBlock.getTerminator(); + if (!scopeTerminator) { + return WalkResult::interrupt(); + } + + OpBuilder scopeBuilder(scopeTerminator); + auto scopeFlagAttr = + scopeBuilder.getIntegerAttr(scopeBuilder.getI64Type(), waitFlagId); + scopeBuilder.setInsertionPoint(scopeTerminator); + scopeBuilder.create(scopeTerminator->getLoc(), coreType, + setPipe, waitPipe, scopeFlagAttr); + + WalkResult innerResult = scopeOp.walk([&](scf::ForOp forOp) -> WalkResult { + if (!forOp->hasAttr("ssbuffer.main_loop")) { + return WalkResult::advance(); + } + Block &forBody = forOp.getRegion().front(); + return insertSyncOpsForForOp(forOp, &forBody, coreType, setPipe, waitPipe, + waitFlagId, setFlagId); + }); + if (innerResult.wasInterrupted()) { + return WalkResult::interrupt(); + } + + return WalkResult::advance(); +} + +// Insert sync ops for a single scopeOp (vector variant: inserts SyncBlockSetOp +// at scope start) +static WalkResult insertSyncOpsForVector(scope::ScopeOp scopeOp, + hivm::TCoreTypeAttr coreType, + PipeAttr setPipe, PipeAttr waitPipe, + int waitFlagId, int setFlagId) { + Block &scopeBlock = scopeOp.getRegion().front(); + OpBuilder builder(&scopeBlock, scopeBlock.begin()); + auto scopeFlagAttr = builder.getIntegerAttr(builder.getI64Type(), setFlagId); + builder.create(scopeOp.getLoc(), coreType, setPipe, waitPipe, + scopeFlagAttr); + + WalkResult innerResult = scopeOp.walk([&](scf::ForOp forOp) -> WalkResult { + if (!forOp->hasAttr("ssbuffer.main_loop")) { + return WalkResult::advance(); + } + Block &forBody = forOp.getRegion().front(); + return insertSyncOpsForForOp(forOp, &forBody, coreType, setPipe, waitPipe, + waitFlagId, setFlagId); + }); + if (innerResult.wasInterrupted()) { + return WalkResult::interrupt(); + } + + return WalkResult::advance(); +} + +// Insert inter-core PIPE_S synchronization for cube cores +static LogicalResult insertInterCorePipeSForCube(ModuleOp module) { + auto cubeCoreType = + hivm::TCoreTypeAttr::get(module.getContext(), hivm::TCoreType::CUBE); + auto setPipeType = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + auto waitPipeType = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + + WalkResult result = module.walk([&](scope::ScopeOp scopeOp) -> WalkResult { + auto attr = scopeOp->getAttrOfType("hivm.tcore_type"); + if (!attr || attr != cubeCoreType) { + return WalkResult::advance(); + } + + return insertSyncOpsForCube(scopeOp, cubeCoreType, setPipeType, + waitPipeType, kPipeSFlagId, kPipeSFlagId); + }); + + return result.wasInterrupted() ? failure() : success(); +} + +// Insert inter-core PIPE_S synchronization for vector cores +static LogicalResult insertInterCorePipeSForVector(ModuleOp module) { + auto vectorCoreType = + hivm::TCoreTypeAttr::get(module.getContext(), hivm::TCoreType::VECTOR); + auto setPipeType = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + auto waitPipeType = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + + WalkResult result = module.walk([&](scope::ScopeOp scopeOp) -> WalkResult { + auto attr = scopeOp->getAttrOfType("hivm.tcore_type"); + if (!attr || attr != vectorCoreType) { + return WalkResult::advance(); + } + + return insertSyncOpsForVector(scopeOp, vectorCoreType, setPipeType, + waitPipeType, kPipeSFlagId, kPipeSFlagId); + }); + + return result.wasInterrupted() ? failure() : success(); +} + +LogicalResult UpdateForOpsPass::insertInterCorePipeS(ModuleOp module) { + if (failed(insertInterCorePipeSForCube(module))) { + return failure(); + } + + if (failed(insertInterCorePipeSForVector(module))) { + return failure(); + } + + return success(); +} + +void UpdateForOpsPass::runOnOperation() { + ModuleOp module = getOperation(); + + LDBG("before updateForOps:\n" << module << "\n"); + + // Use provided info, or create a local one if not available + ControlFlowConditionInfo localInfo; + ControlFlowConditionInfo *infoToUse = info ? info : &localInfo; + + // Derive block counters from ssbuffer.if if blockCounterNums is empty + if (infoToUse->blockCounterNums.empty()) { + if (failed(deriveBlockCountersFromIfOps(module, infoToUse))) { + signalPassFailure(); + return; + } + } + + // Update for ops iter_args for block counters and inner dependency conditions + if (infoToUse && (!infoToUse->blockCounterNums.empty() || + !infoToUse->intraCoreDependentMap.empty())) + if (failed(addBlockCountersAndInnerDepConds(module, infoToUse))) { + signalPassFailure(); + return; + } + + // Insert PIPE_S inter-core synchronization + if (failed(insertInterCorePipeS(module))) { + signalPassFailure(); + return; + } + + LDBG("after updateForOps:\n" << module << "\n"); +} + +namespace mlir { +namespace triton { + +std::unique_ptr> createUpdateForOpsPass() { + return std::make_unique(); +} + +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.cpp new file mode 100644 index 0000000000..6d3d25e105 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.cpp @@ -0,0 +1,698 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/UpdateLoopIterTimes.h" +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/Scope/IR/Scope.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/Support/Debug.h" + +static constexpr const char *DEBUG_TYPE = "UpdateLoopIterTimes"; +#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ") +#define LDBG(...) LLVM_DEBUG(DBGS() << __VA_ARGS__ << "\n") + +using namespace llvm; +using namespace mlir; +using namespace triton; +using namespace hivm; + +// Find the IfOp index that an operation belongs to +// Uses isAncestor to check if the operation is within the IfOp's region +static int findIfOpIndexInList(Operation *op, SmallVector &ifOps, + llvm::DenseMap &ifOpIndex) { + for (scf::IfOp ifOp : ifOps) { + if (ifOp->isAncestor(op)) { + return ifOpIndex[ifOp.getOperation()]; + } + } + return -1; +} + +// calculateFactor computes the buffer factor needed for loop iteration +// extension Core idea: if producer ifOp (n) and consumer ifOp (m) have +// dependency, we need (m - n + 1) buffers to support the loop extension Returns +// {maxRequiredBuffers, x} for computing: ceil(original_iter_count * +// requiredBuffers / x) + ifCount +std::pair UpdateLoopIterTimesPass::calculateFactor(scf::ForOp forOp) { + int maxRequiredBuffers = 1; + int maxX = 1; + + // Step1: Collect all IfOps with ssbuffer.if attribute in this for loop and + // build index Index represents execution order (smaller index = earlier + // execution) + SmallVector ifOps; + DenseMap ifOpIndex; + int index = 1; + + int ret = 0; + forOp.walk([&](Operation *op) { + if (op->hasAttr("ssbuffer.if")) { + auto ifOp = dyn_cast(op); + if (!ifOp) { + ret = 1; + LDBG("ssbuffer.if attribute is not allocated on ifOp!"); + return WalkResult::interrupt(); + } + ifOps.push_back(ifOp); + ifOpIndex[ifOp.getOperation()] = index++; + } + return WalkResult::advance(); + }); + if (ret == -1) { + return {-1, -1}; + } + + if (ifOps.empty()) { + LDBG("mainloop do not contains ifblocks!"); + return {-1, -1}; + } + + // No dependency, multiple factor = 1 + if (!info->intraCoreDependentMap.count(forOp)) { + return {1, 1}; + } + auto &deps = info->intraCoreDependentMap[forOp]; + if (deps.empty()) { + return {1, 1}; + } + + // Step2: Iterate all dependencies and calculate required buffer count + for (auto &entry : deps) { + Value consumerResult = entry.first; // Consumer result value + SmallVector producerBuffers = entry.second; // Producer buffer list + int x = producerBuffers.size(); // Producer buffer count + + // Get the defining operation of consumer result + Operation *consumerDefOp = consumerResult.getDefiningOp(); + if (!consumerDefOp) { + LDBG("consumerResult do not have the defOp!"); + return {-1, -1}; + } + + // Find the IfOp index that consumer belongs to (m) + int m = findIfOpIndexInList(consumerDefOp, ifOps, ifOpIndex); + if (m == -1) { + LDBG("Can not find the consumerDefOp in any ifOps!\nconsumerDefOp: " + << *consumerDefOp); + return {-1, -1}; + } + + // Step3: Find the IfOp index that producer belongs to (n) + // Producer is the operation that uses producer buffer, i.e., + // bufferization.materialize_in_destination This op is the production + // behavior op, with producer buffer as its input + if (producerBuffers.empty()) { + LDBG("consumer do not have the producerBuffers!"); + return {-1, -1}; + } + + int producerIfOpIndex = -1; + for (Value buffer : producerBuffers) { + for (Operation *user : buffer.getUsers()) { + if (isa(user) || + isa(user)) { + producerIfOpIndex = findIfOpIndexInList(user, ifOps, ifOpIndex); + if (producerIfOpIndex == -1) { + LDBG("Can not find the producerBuffers in any " + "ifOps\nproducerBuffers: " + << *user); + return {-1, -1}; + } + break; + } + } + // any producerBuffer found in a ifblock means all producerBuffers in that + // ifblock + if (producerIfOpIndex != -1) { + break; + } + } + + // If cannot find the IfOp producer belongs to, skip this dependency + if (producerIfOpIndex == -1) { + LDBG("All producerBuffers are not found in any ifOps"); + return {-1, -1}; + } + + int n = producerIfOpIndex; + + // Step4: If consumer is after producer (m > n), calculate required buffer + // count m - n + 1 represents the buffer count needed to cover this distance + if (m <= n) { + LDBG("producer is after the comsumer!"); + return {-1, -1}; + } + int requiredBuffers = m - n + 1; + + // Step5: Update max value using fraction comparison to avoid precision + // issues Comparing requiredBuffers/maxX vs maxRequiredBuffers/x is + // equivalent to comparing requiredBuffers * maxX vs maxRequiredBuffers * x + if (requiredBuffers * maxX > maxRequiredBuffers * x) { + maxRequiredBuffers = requiredBuffers; + maxX = x; + } + } + + return {maxRequiredBuffers, maxX}; +} + +// Part 1: Compute new loop upper bound and create arithmetic ops +Value UpdateLoopIterTimesPass::computeNewLoopUpperBound( + OpBuilder &builder, Location loc, scf::ForOp forOp, int ifCount, + int requiredBuffers, int x) { + Value originalLowerBound = forOp.getLowerBound(); + Value originalUpperBound = forOp.getUpperBound(); + Value originalStep = forOp.getStep(); + Type ubType = originalStep.getType(); + + // Helper lambda to create constant value based on type + auto createConstant = [&](int val) -> Value { + if (ubType.isIndex()) { + return builder.create(loc, val); + } else if (auto intType = dyn_cast(ubType)) { + return builder.create(loc, intType, val); + } else { + auto indexVal = builder.create(loc, val); + return builder.create(loc, ubType, indexVal); + } + }; + + // Helper lambda to create ceil division op based on type + auto createCeilDiv = [&](Value lhs, Value rhs) -> Value { + if (ubType.isIndex()) { + return builder.create(loc, lhs, rhs); + } else if (auto intType = dyn_cast(ubType)) { + if (intType.isSigned()) { + return builder.create(loc, lhs, rhs); + } else { + return builder.create(loc, lhs, rhs); + } + } else { + return builder.create(loc, lhs, rhs); + } + }; + + // Create constant values + Value ifCountValue = createConstant(ifCount); + Value requiredBuffersValue = createConstant(requiredBuffers); + Value xValue = createConstant(x); + + // Calculate: rangeDiff = upperBound - lowerBound + Value rangeDiff = builder.create(loc, originalUpperBound, + originalLowerBound); + + // Calculate: iterCount = ceil(rangeDiff / step) + Value iterCount = createCeilDiv(rangeDiff, originalStep); + + // Calculate: scaledIterCount = iterCount * requiredBuffers + Value scaledIterCount = + builder.create(loc, iterCount, requiredBuffersValue); + + // Calculate: ceiledScaledIterCount = ceil(scaledIterCount / x) + Value ceiledScaledIterCount = createCeilDiv(scaledIterCount, xValue); + + // Calculate: newIterCount = ceiledScaledIterCount + ifCount + Value newIterCount = + builder.create(loc, ceiledScaledIterCount, ifCountValue); + + // Calculate: totalSteps = step * newIterCount + Value totalSteps = + builder.create(loc, originalStep, newIterCount); + + // Calculate: newUpperBound = lowerBound + totalSteps + Value newUpperBound = + builder.create(loc, originalLowerBound, totalSteps); + + return newUpperBound; +} + +// Part 2: Clone for loop with new upper bound +scf::ForOp UpdateLoopIterTimesPass::cloneForOpWithNewUpperBound( + OpBuilder &builder, Location loc, scf::ForOp oldForOp, Value newUpperBound, + IRMapping &mapper) { + Value originalLowerBound = oldForOp.getLowerBound(); + Value originalStep = oldForOp.getStep(); + + SmallVector newInitArgs(oldForOp.getInitArgs().begin(), + oldForOp.getInitArgs().end()); + + auto newForOp = builder.create( + loc, originalLowerBound, newUpperBound, originalStep, newInitArgs); + + // Copy attributes from old forOp + for (auto &attr : oldForOp->getAttrs()) { + newForOp->setAttr(attr.getName(), attr.getValue()); + } + + // Map induction variable and iter args + mapper.map(oldForOp.getInductionVar(), newForOp.getInductionVar()); + + for (auto [oldArg, newArg] : + llvm::zip(oldForOp.getRegionIterArgs(), newForOp.getRegionIterArgs())) { + mapper.map(oldArg, newArg); + } + + Block *oldBlock = oldForOp.getBody(); + Block *newBlock = newForOp.getBody(); + + // Replace block arguments + unsigned totalArgs = oldBlock->getNumArguments(); + for (unsigned i = 0; i < totalArgs; ++i) { + BlockArgument oldArg = oldBlock->getArgument(i); + BlockArgument newArg = newBlock->getArgument(i); + oldArg.replaceAllUsesWith(newArg); + } + + Operation *oldTerminator = oldBlock->getTerminator(); + + // Clone operations in the loop body + builder.setInsertionPointToStart(newBlock); + for (Operation &op : + llvm::make_early_inc_range(oldBlock->without_terminator())) { + builder.clone(op, mapper); + } + + // Create new yield operation + auto oldYield = cast(oldTerminator); + SmallVector newYieldOperands; + for (unsigned i = 0; i < oldYield.getNumOperands(); ++i) { + newYieldOperands.push_back(mapper.lookupOrDefault(oldYield.getOperand(i))); + } + + builder.setInsertionPointToEnd(newBlock); + builder.create(loc, newYieldOperands); + + // Handle results replacement + unsigned numOriginalResults = oldForOp.getNumResults(); + if (numOriginalResults > 0) { + SmallVector originalResults; + for (unsigned i = 0; i < numOriginalResults; ++i) { + originalResults.push_back(newForOp.getResult(i)); + } + oldForOp.replaceAllUsesWith(originalResults); + } + + return newForOp; +} + +// Part 3: Update cntArgs mapping after cloning +int UpdateLoopIterTimesPass::updateCntArgsAfterClone( + scf::ForOp oldForOp, IRMapping &mapper, + SmallVector &ifOpsInThisFor) { + // Update cntArgs for each old ifOp + for (scf::IfOp oldIfOp : ifOpsInThisFor) { + scf::IfOp newIfOp = dyn_cast(mapper.lookupOrDefault(oldIfOp)); + if (!newIfOp) + continue; + + Value oldCntVal = info->cntArgs[oldIfOp]; + Value newCntVal = mapper.lookupOrDefault(oldCntVal); + LDBG("updating CntArgs After Clone..."); + LDBG("oldCntVal: " << oldCntVal); + LDBG("newCntVal: " << newCntVal); + + // Remove old entry and add new entry + info->cntArgs.erase(oldIfOp); + info->cntArgs[newIfOp] = newCntVal; + } + return 0; +} + +// Extend for loop iteration count +scf::ForOp UpdateLoopIterTimesPass::extendForOpIterationCount( + scf::ForOp oldForOp, int ifCount, int requiredBuffers, int x, + IRMapping &mapper, SmallVector &ifOpsInThisFor) { + OpBuilder builder(oldForOp); + Location loc = oldForOp.getLoc(); + + // Part 1: Compute new loop upper bound + Value newUpperBound = computeNewLoopUpperBound(builder, loc, oldForOp, + ifCount, requiredBuffers, x); + if (!newUpperBound) { + LDBG("computeNewLoopUpperBound failed!"); + return nullptr; + } + + // Part 2: Clone for loop with new upper bound + scf::ForOp newForOp = cloneForOpWithNewUpperBound(builder, loc, oldForOp, + newUpperBound, mapper); + if (!newForOp) { + LDBG("cloneForOpWithNewUpperBound failed!"); + return nullptr; + } else { + LDBG("cloneForOpWithNewUpperBound Success!"); + LDBG("new ForOp: " << newForOp); + LDBG("old ForOp: " << oldForOp); + } + + // Part 3: Update cntArgs mapping + if (updateCntArgsAfterClone(oldForOp, mapper, ifOpsInThisFor) != 0) { + LDBG("updateCntArgsAfterClone failed!"); + return nullptr; + } + + return newForOp; +} + +// step4: Replace loop counter by if blocks' counter +// Traverse each mainloop, find ifOp with ssbuffer.if attribute inside, +// and replace the mainloop's induction variable with the counter in cntArgs +int UpdateLoopIterTimesPass::replaceForOpCounterInIfOps() { + int ret = 0; + // Traverse all mainloops in the module + getOperation().walk([&](Operation *op) { + if (op->hasAttr("ssbuffer.main_loop")) { + auto forOp = dyn_cast(op); + if (!forOp) { + ret = -1; + LDBG("Do not support other mainloop except forOp!"); + return WalkResult::interrupt(); + } + + Value indVar = forOp.getInductionVar(); + + // Find all ifOps with ssbuffer.if attribute inside this mainloop + forOp.walk([&](scf::IfOp ifOp) { + if (ifOp->hasAttr("ssbuffer.if")) { + if (!info->cntArgs.count(ifOp)) { + LDBG("ifblock has no counter in cntArgs"); + ret = -1; + return WalkResult::interrupt(); + } + + Value cntVal = info->cntArgs[ifOp]; + + // Replace all uses of induction variable in the ifOp with cntVal + ifOp.walk([&](Operation *op) { + for (OpOperand &operand : op->getOpOperands()) { + if (operand.get() == indVar) { + operand.set(cntVal); + } + } + }); + } + return mlir::WalkResult::advance(); + }); + } + return mlir::WalkResult::advance(); + }); + return ret; +} + +// step1: Get mainloop id to loop operation mapping, separated into cube and +// vector +int UpdateLoopIterTimesPass::GetMainLoopIdToLoopOpMap( + ModuleOp module, DenseMap> &cmap, + DenseMap> &vmap) { + cmap.clear(); + vmap.clear(); + int ret = 0; + + module.walk([&](scope::ScopeOp scopeOp) { + // Determine if it's CUBE or VECTOR + bool isCube = false; + bool isVector = false; + if (scopeOp->hasAttr("hivm.tcore_type")) { + auto attr = scopeOp->getAttr("hivm.tcore_type"); + auto aiCAttr = hivm::TCoreTypeAttr::get(scopeOp->getContext(), + hivm::TCoreType::CUBE); + auto aiVAttr = hivm::TCoreTypeAttr::get(scopeOp->getContext(), + hivm::TCoreType::VECTOR); + if (attr == aiCAttr) { + isCube = true; + } else if (attr == aiVAttr) { + isVector = true; + } + } + + if (!(isCube || isVector)) { + ret = -1; + LDBG("mlir do not processed by split mix kernel!"); + return mlir::WalkResult::interrupt(); + } + + // Walk for loops inside the scope + scopeOp.walk([&](Operation *op) { + if (op->hasAttr("ssbuffer.main_loop")) { + auto forOp = dyn_cast(op); + if (!forOp) { + ret = -1; + LDBG("do not surpport other loop op temprarily!"); + return mlir::WalkResult::interrupt(); + } + + auto mainLoopId = + forOp->getAttrOfType("ssbuffer.main_loop"); + if (mainLoopId) { + int id = mainLoopId.getInt(); + if (isCube) { + cmap[id].push_back(forOp.getOperation()); + } else if (isVector) { + vmap[id].push_back(forOp.getOperation()); + } + } + } + return mlir::WalkResult::advance(); + }); + return mlir::WalkResult::advance(); + }); + + return ret; +} + +// Input: loopMap： {mainloop_id: [loopOp1, loopOp2], ...} +// Output: infoMap: {loopOp: extraIterationTimesInfo, ...} +int UpdateLoopIterTimesPass::ComputeMainLoopTimes( + DenseMap> &loopMap, + DenseMap &infoMap) { + for (auto &entry : loopMap) { + for (Operation *loopOp : entry.second) { + scf::ForOp forOp = dyn_cast(loopOp); + if (!forOp) { + LDBG("currently only support forOp!"); + return -1; + } + + // Collect if operations in this loop + IterationTimesInfo iterInfo; + if (info->cntArgs.empty()) { + LDBG("cntArgs is empty, no ifblock is contained!"); + return -1; + } + for (auto &[ifOp, cntVal] : info->cntArgs) { + if (ifOp->hasAttr("ssbuffer.if")) { + auto parentOp = ifOp->getParentOp(); + if (parentOp->hasAttr("ssbuffer.main_loop") && + isa(parentOp)) { + if (parentOp == forOp.getOperation()) { + iterInfo.ifCount++; + iterInfo.ifOpsInThisFor.push_back(ifOp); + } + } else { + LDBG("Get wrong ifblock structure: ifblock's parentOp is not " + "mainloop!"); + LDBG("ifblock Op: " << ifOp); + LDBG("Parent Op: " << *parentOp); + return -1; + } + } else { + LDBG("ifOp in cntArgs does not contain ssbuffer.if Attribute\n ifOp: " + << ifOp); + return -1; + } + } + + // Calculate factor + auto [requiredBuffers, x] = calculateFactor(forOp); + if (requiredBuffers == -1 || x == -1) { + LDBG("calculateFactor failed!"); + return -1; + } + iterInfo.requiredBuffers = requiredBuffers; + iterInfo.x = x; + infoMap[loopOp] = iterInfo; + } + } + return 0; +} + +int UpdateLoopIterTimesPass::collectForOpsAndUpdateMax( + DenseMap> &map, int id, + SmallVector &allForOps, int &maxIfCount, + int &maxRequiredBuffers, int &maxX, + DenseMap &infoMap) { + if (map.count(id)) { + for (Operation *loopOp : map[id]) { + allForOps.push_back(loopOp); + if (infoMap.count(loopOp)) { + IterationTimesInfo &iterInfo = infoMap[loopOp]; + if (iterInfo.ifCount > maxIfCount) { + maxIfCount = iterInfo.ifCount; + } + // requiredBuffers/x > maxRequiredBuffers/maxX = + // requiredBuffers * maxX > maxRequiredBuffers * x + if (iterInfo.requiredBuffers * maxX > maxRequiredBuffers * iterInfo.x) { + maxRequiredBuffers = iterInfo.requiredBuffers; + maxX = iterInfo.x; + } + } else { + LDBG("mainloop Op: " << *loopOp << " do not include in infoMap!"); + return -1; + } + } + } + return 0; +} + +int UpdateLoopIterTimesPass::UpdateForLoopIteration( + DenseMap> &cmap, + DenseMap> &vmap, + DenseMap &infoMap) { + int ret = 0; + // Collect all mainloop ids + DenseSet allIds; + for (auto &entry : cmap) { + allIds.insert(entry.first); + } + for (auto &entry : vmap) { + allIds.insert(entry.first); + } + + // For each mainloop id, calculate max values and update all related loops + SmallVector allForOps; + for (int id : allIds) { + int maxIfCount = 0; + int maxRequiredBuffers = 1; + int maxX = 1; + + // Collect all loops with same id from cmap and vmap + SmallVector sameIdForOps; + ret = collectForOpsAndUpdateMax(cmap, id, sameIdForOps, maxIfCount, + maxRequiredBuffers, maxX, infoMap); + if (ret != 0) + return -1; + ret = collectForOpsAndUpdateMax(vmap, id, sameIdForOps, maxIfCount, + maxRequiredBuffers, maxX, infoMap); + if (ret != 0) + return -1; + + if (maxIfCount == 0) { + LDBG("no ifblock in mainloop!"); + return -1; + } + + // Update all loops with this id using the same max values + for (Operation *loopOp : sameIdForOps) { + scf::ForOp oldForOp = dyn_cast(loopOp); + if (!oldForOp) { + LDBG("do not surpport other loop op except forOp!"); + return -1; + } + + IterationTimesInfo &iterInfo = infoMap[loopOp]; + IRMapping mapper; + scf::ForOp newForOp = + extendForOpIterationCount(oldForOp, maxIfCount, maxRequiredBuffers, + maxX, mapper, iterInfo.ifOpsInThisFor); + if (!newForOp) { + LDBG("extendForOpIterationCount failed!"); + return -1; + } + allForOps.push_back(loopOp); + } + } + + // Delete old forOp after cntArgs has been updated in + // extendForOpIterationCount + for (Operation *loopOp : allForOps) { + if (!loopOp) { + LDBG("erasing error: loopOp is nullptr, there are nested mainloop!"); + return -1; + } + loopOp->erase(); + } + return 0; +} + +void UpdateLoopIterTimesPass::runOnOperation() { + ModuleOp module = getOperation(); + + LDBG("before updateloopitertimes:\n" << module); + LDBG("\nEnter UpdateLoopIterTimesPass!"); + + int ret = 0; + // step1: Get mainloop id to loop operation mapping, separated into cube and + // vector + DenseMap> cmap; + DenseMap> vmap; + ret = GetMainLoopIdToLoopOpMap(module, cmap, vmap); + if (ret != 0) { + LDBG("\nGetMainLoopIdToLoopOpMap Failed!"); + signalPassFailure(); + return; + } + + // step2: Calculate info for each loop operation, store into the same + // iterationTimesinfoMap + DenseMap infoMap; + ret = ComputeMainLoopTimes(cmap, infoMap); + if (ret != 0) { + LDBG("\nComputeMainLoopTimes from cube Failed!"); + signalPassFailure(); + return; + } + ret = ComputeMainLoopTimes(vmap, infoMap); + if (ret != 0) { + LDBG("\nComputeMainLoopTimes from vector Failed!"); + signalPassFailure(); + return; + } + + // step3: Update loop iteration count (process loop operations with same id + // from both cmap and vmap) + ret = UpdateForLoopIteration(cmap, vmap, infoMap); + if (ret != 0) { + LDBG("Update ForLoop Iteration Failed!\n"); + signalPassFailure(); + } + LDBG("after UpdateForLoopIteration:\n" << module); + + // step4: Replace loop counter by if blocks' counter + if (replaceForOpCounterInIfOps() != 0) { + LDBG("replaceForOpCounterInIfOps Failed!\n"); + signalPassFailure(); + } + + LDBG("after updateloopitertimes:\n" << module); + LDBG("\nExit UpdateLoopIterTimes pass."); +} + +namespace mlir { +namespace triton { +std::unique_ptr> createUpdateLoopIterTimesPass() { + return std::make_unique(); +} +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/Utils.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/Utils.cpp new file mode 100644 index 0000000000..6d8fe5aaaf --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/AddControlFlowCondition/Utils.cpp @@ -0,0 +1,163 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition/Utils.h" + +using namespace mlir; +using namespace llvm; + +// Collect all nested ops within an operation's regions +LogicalResult +triton::collectAllNestedOps(Operation *op, + llvm::DenseSet ®ionOps) { + if (!op) { + return failure(); + } + + if (regionOps.contains(op)) { + return success(); + } + + regionOps.insert(op); + for (Region ®ion : op->getRegions()) { + for (Block &block : region) { + for (Operation &nestedOp : block) { + if (failed(triton::collectAllNestedOps(&nestedOp, regionOps))) { + return failure(); + } + } + } + } + + return success(); +} + +// Group operations by their block_id attribute +LogicalResult triton::collectOpsByBlockId( + scf::ForOp forOp, llvm::DenseMap> &blockOps) { + for (Operation &op : forOp.getBody()->without_terminator()) { + if (auto attr = op.getAttrOfType("ssbuffer.block_id")) { + blockOps[attr.getInt()].push_back(&op); + } else { + return failure(); + } + } + + return success(); +} + +// DFS for topological sort - returns failure if cycle detected +static LogicalResult +dfsTopologicalSort(Operation *op, llvm::DenseSet &visited, + llvm::DenseSet &inStack, + const llvm::DenseSet &ops, + llvm::DenseMap *opOrder, + SmallVectorImpl &sorted) { + if (!op) { + return success(); + } + if (inStack.contains(op)) { + return failure(); + } + if (visited.contains(op)) { + return success(); + } + + visited.insert(op); + inStack.insert(op); + + SmallVector deps; + for (Value operand : op->getOperands()) { + if (Operation *def = operand.getDefiningOp()) { + if (ops.contains(def)) { + deps.push_back(def); + } + } + } + + if (opOrder && !opOrder->empty()) { + llvm::sort(deps, [&](Operation *a, Operation *b) { + return (*opOrder)[a] < (*opOrder)[b]; + }); + } + + for (Operation *dep : deps) { + if (failed( + dfsTopologicalSort(dep, visited, inStack, ops, opOrder, sorted))) { + return failure(); + } + } + + inStack.erase(op); + sorted.push_back(op); + return success(); +} + +// Topological sort of operations based on operand dependencies +LogicalResult triton::topologicalSort(llvm::DenseSet &ops, + llvm::DenseMap *opOrder, + SmallVectorImpl &sorted) { + llvm::DenseSet visited; + llvm::DenseSet inStack; + + SmallVector opList(ops.begin(), ops.end()); + if (opOrder && !opOrder->empty()) { + llvm::sort(opList, [&](Operation *a, Operation *b) { + return (*opOrder)[a] < (*opOrder)[b]; + }); + } + + for (Operation *op : opList) { + if (failed( + dfsTopologicalSort(op, visited, inStack, ops, opOrder, sorted))) { + return failure(); + } + } + return success(); +} + +LogicalResult triton::topologicalSort(SmallVector &ops) { + llvm::DenseSet opSet(ops.begin(), ops.end()); + SmallVector sorted; + + if (failed(triton::topologicalSort(opSet, nullptr, sorted))) { + return failure(); + } + ops.assign(sorted.begin(), sorted.end()); + return success(); +} + +// Get block_ids in order of appearance in for loop body +SmallVector triton::getBlockIdsInOrder(scf::ForOp forOp) { + SmallVector idsInOrder; + llvm::DenseSet seenIds; + + for (Operation &op : forOp.getBody()->without_terminator()) { + if (auto blockIdAttr = op.getAttrOfType("ssbuffer.block_id")) { + int id = blockIdAttr.getInt(); + if (seenIds.insert(id).second) { + idsInOrder.push_back(id); + } + } + } + return idsInOrder; +} diff --git a/third_party/ascend/lib/DynamicCVPipeline/AddDynamicCVPipeline.cpp b/third_party/ascend/lib/DynamicCVPipeline/AddDynamicCVPipeline.cpp new file mode 100644 index 0000000000..4fbfb3b732 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/AddDynamicCVPipeline.cpp @@ -0,0 +1,82 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/AddControlFlowCondition.h" +#include "ascend/include/DynamicCVPipeline/AllocMultiCache.h" +#include "ascend/include/DynamicCVPipeline/Passes.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlockPass.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlockPass.h" +#include "ascend/include/DynamicCVPipeline/SeparateMemoryFromComputePass.h" +#include "ascend/include/DynamicCVPipeline/SplitDataflowPass.h" +#include "mlir/Pass/PassManager.h" +#include "llvm/Support/Debug.h" + +static constexpr const char *DEBUG_TYPE = "AddDynamicCVPipeline"; +#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ") +#define LDBG(X) LLVM_DEBUG(DBGS() << (X) << "\n") + +namespace mlir { +namespace triton { +#define GEN_PASS_DEF_ADDDYNAMICCVPIPELINE +#include "ascend/include/DynamicCVPipeline/Passes.h.inc" +} // namespace triton +} // namespace mlir + +using namespace mlir; + +AddDynamicCVPipelinePass::AddDynamicCVPipelinePass( + const AddDynamicCVPipelineOptions &options) + : AddDynamicCVPipelineBase(options) {} + +void AddDynamicCVPipelinePass::runOnOperation() { + auto moduleOp = getOperation(); + compileOn91095Flag = this->compileOn91095; + + LDBG("Enter pass"); + + if (!compileOn91095Flag) { + llvm::errs() << "Add-dynamic-cv-pipeline is only supported on 91095 now.\n"; + return; + } + + PassManager pm(&getContext(), moduleOp.getOperationName()); + + // todo: add related passes. + pm.addPass(createPlanComputeBlockPass()); + pm.addPass(createSplitDataflowPass()); + pm.addPass(createSeparateMemoryFromComputePass()); + pm.addPass(createAllocMultiCachePass()); + pm.addPass(createAddControlFlowConditionPass()); + + if (failed(runPipeline(pm, getOperation()))) { + moduleOp->emitError() << "[" << DEBUG_TYPE << "] Pass failed!"; + signalPassFailure(); + } + + LDBG("Process successfully"); +} + +std::unique_ptr> +mlir::triton::createAddDynamicCVPipelinePass( + const AddDynamicCVPipelineOptions &options) { + return std::make_unique(options); +} diff --git a/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache.cpp b/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache.cpp new file mode 100644 index 0000000000..72ca23d34a --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache.cpp @@ -0,0 +1,71 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/AllocMultiCache.h" +#include "ascend/include/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.h" +#include "ascend/include/DynamicCVPipeline/AllocMultiCache/AddMultiBufferOuterScope.h" + +#include "mlir/Pass/PassManager.h" +#include "llvm/Support/Debug.h" + +static constexpr const char *DEBUG_TYPE = "AllocMultiCache"; +#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ") +#define LDBG(...) \ + LLVM_DEBUG({ \ + DBGS(); \ + llvm::outs() << __VA_ARGS__; \ + llvm::outs() << "\n"; \ + }) + +using namespace mlir; +using namespace triton; + +// Run the pass +void AllocMultiCachePass::runOnOperation() { + ModuleOp module = getOperation(); + OpPassManager pm(module.getOperationName()); + LDBG("Enter pass."); + LDBG("before innerscope:\n" << module << "\n"); + // Step 1:Inner multibuffer + pm.addPass(createAddMultiBufferInnerScopePass()); + + // Step 2: Outer multibuffer + pm.addPass(createAddMultiBufferOuterScopePass()); + + if (failed(runPipeline(pm, module))) { + module->emitError() << "[" << DEBUG_TYPE << "] Pass failed!"; + signalPassFailure(); + } + + LDBG("Process successfully"); + LDBG(llvm::StringRef("after innerscope:\n") << module << "\n"); +} + +namespace mlir { +namespace triton { + +std::unique_ptr> createAllocMultiCachePass() { + return std::make_unique(); +} + +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.cpp b/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.cpp new file mode 100644 index 0000000000..c1f918860f --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.cpp @@ -0,0 +1,1062 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/AllocMultiCache/AddMultiBufferInnerScope.h" +#include "ascend/include/DynamicCVPipeline/Common/BufferCountManager.h" + +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/Scope/IR/Scope.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Pass/PassManager.h" +#include "llvm/Support/Debug.h" +#include + +static constexpr const char *DEBUG_TYPE = "AddMultiBufferInnerScope"; +#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ") +#define LDBG(X) LLVM_DEBUG(DBGS() << (X) << "\n") + +using namespace mlir; +using namespace hivm; +using namespace annotation; +using namespace triton; + +using BufferPair = std::pair; +using BufferMap = DenseMap>; + +// Buffer count constants +constexpr int kBufferCountOne = 1; + +// Attribute name constants +static constexpr const char *kAttrBlockId = "ssbuffer.block_id"; +static constexpr const char *kAttrMainLoop = "ssbuffer.main_loop"; + +namespace mlir { +namespace triton { + +// Check if forOp has main_loop attribute +static bool hasMainLoopAttr(scf::ForOp forOp) { + if (forOp->hasAttr(kAttrMainLoop)) { + return true; + } + if (auto *term = forOp.getBody()->getTerminator()) + return term->hasAttr(kAttrMainLoop); + return false; +} + +// Collect main_loop forOps in a single block +static int collectMainLoopsInBlock(Block &block, + SmallVector &mainLoopForOps) { + int count = 0; + for (Operation &op : block) { + if (auto forOp = dyn_cast(&op)) { + if (hasMainLoopAttr(forOp)) { + mainLoopForOps.push_back(forOp); + count++; + } + } + } + return count; +} + +// Recursively collect main_loop forOps, returns count of collected items +static int +collectMainLoopsRecursively(Region ®ion, + SmallVector &mainLoopForOps) { + int totalCount = 0; + for (Block &block : region) { + totalCount += collectMainLoopsInBlock(block, mainLoopForOps); + for (Operation &op : block) { + for (auto &nestedRegion : op.getRegions()) + totalCount += collectMainLoopsRecursively(nestedRegion, mainLoopForOps); + } + } + return totalCount; +} + +struct InnerBlockInfo { + Value blockId; + SmallVector ops; +}; + +// Get ssbuffer block_id attribute from op, returns INT_MIN if not found +static int getSsbufferId(Operation *op) { + if (auto idAttr = op->getAttrOfType(kAttrBlockId)) + return idAttr.getInt(); + return INT_MIN; +} + +void collectNestedOps(Block *block, SmallVector &ops) { + for (auto &op : *block) { + ops.push_back(&op); + for (auto ®ion : op.getRegions()) + for (auto &innerBlock : region) + collectNestedOps(&innerBlock, ops); + } +} + +// Get priority of forOp (lower value means higher priority) +// Priority order: main_loop (1) > block_id (2) > iter_args (3) > none (0) +// This is used to select the most relevant main loop when multiple candidates +// exist +static int getForOpPriority(scf::ForOp f) { + constexpr int priorityMainLoop = 1; + constexpr int priorityBlockId = 2; + constexpr int priorityIterArgs = 3; + + // Check if forOp itself has main_loop attribute + bool hasMainloop = f->hasAttr(kAttrMainLoop); + bool bodyHasMainloop = false; + bool bodyHasBlockId = false; + + // Check terminator for main_loop and block_id attributes + if (auto *term = f.getBody()->getTerminator()) { + bodyHasMainloop = term->hasAttr(kAttrMainLoop); + bodyHasBlockId = term->getAttrOfType(kAttrBlockId) != nullptr; + } + + bool opHasBlockId = f->getAttrOfType(kAttrBlockId) != nullptr; + bool hasIterArgs = f.getNumResults() > 0 || !f.getInitArgs().empty(); + + if (hasMainloop || bodyHasMainloop) { + return priorityMainLoop; + } + if (opHasBlockId || bodyHasBlockId) { + return priorityIterArgs; + } + if (hasIterArgs) { + return priorityIterArgs; + } + return 0; +} + +scf::ForOp findMainloopInScope(scope::ScopeOp scope) { + SmallVector allOps; + collectNestedOps(&scope.getBodyRegion().front(), allOps); + + scf::ForOp mainLoopForOp; + int bestPriority = INT_MAX; + + for (Operation *op : allOps) { + auto f = dyn_cast(op); + if (!f) + continue; + + int priority = getForOpPriority(f); + if (priority > 0 && priority < bestPriority) { + mainLoopForOp = f; + bestPriority = priority; + } + } + return mainLoopForOp; +} + +// Collect a single dependency value to depValueMap +static void collectDepValue(Value operand, Block *body, int currentBlockId, + DenseMap &outputToBlockId, + DenseMap> &depValueMap, + Value groupKey) { + if (auto barg = dyn_cast(operand)) { + if (barg.getOwner() == body && + !llvm::is_contained(depValueMap[groupKey], barg)) + depValueMap[groupKey].push_back(barg); + return; + } + + if (outputToBlockId.count(operand) && + outputToBlockId[operand] != currentBlockId && + !llvm::is_contained(depValueMap[groupKey], operand)) + depValueMap[groupKey].push_back(operand); +} + +// Recursively find nested main_loop +static scf::ForOp findNestedMainloopInForOp(scf::ForOp forOp) { + SmallVector allOps; + collectNestedOps(forOp.getBody(), allOps); + + for (Operation *op : allOps) { + auto nestedFor = dyn_cast(op); + if (!nestedFor) + continue; + if (nestedFor->hasAttr(kAttrMainLoop)) + return nestedFor; + } + return {}; +} + +bool isInsideMainLoopForOp(Operation *op) { + Operation *parent = op->getParentOp(); + if (!parent) { + return false; + } + if (auto forOp = dyn_cast(parent)) { + return forOp->hasAttr(kAttrMainLoop); + } + return false; +} + +bool isInsideMainLoopForOpTraverse(Operation *op) { + Operation *parent = op->getParentOp(); + while (parent) { + if (auto forOp = dyn_cast(parent)) { + if (forOp->hasAttr(kAttrMainLoop)) { + return true; + } + } + parent = parent->getParentOp(); + } + return false; +} + +// Collect all ops with ssbuffer.id from allOps, grouped by id +// Returns 0=success, -1=invalid negative block ID from upstream pass +static int +groupOpsBySsbufferId(SmallVector &allOps, + llvm::MapVector> &opsById) { + llvm::MapVector opsByValue; + for (Operation *op : allOps) { + int id = getSsbufferId(op); + if (id == INT_MIN) + continue; // Attribute doesn't exist, skip + for (auto res : op->getResults()) + opsByValue[res] = op; + } + for (auto &p : opsByValue) { + int id = getSsbufferId(p.second); + if (id == INT_MIN) + continue; + opsById[id].push_back(p.second); + } + return 0; +} + +// Returns 0=success (including normal skip when blocks empty), -1=invalid +// negative block ID +static int +collectInnerBlockInfo(scf::ForOp forOp, DenseMap &blocks, + DenseMap> &depValueMap) { + depValueMap.clear(); + Block *body = forOp.getBody(); + if (!body) + return 0; + + SmallVector allOps; + collectNestedOps(body, allOps); + + llvm::MapVector> opsById; + if (groupOpsBySsbufferId(allOps, opsById) != 0) + return -1; + if (opsById.empty()) + return 0; + + // Build mapping from output to block id + DenseMap outputToBlockId; + for (auto &p : opsById) + for (Operation *op : p.second) + for (auto res : op->getResults()) + outputToBlockId[res] = p.first; + + // Collect dependency values for each block + for (auto &p : opsById) { + Value groupKey = p.second.front()->getResult(0); + InnerBlockInfo bi; + bi.blockId = groupKey; + bi.ops = p.second; + blocks[groupKey] = bi; + + for (Operation *op : bi.ops) + for (Value operand : op->getOperands()) + collectDepValue(operand, body, p.first, outputToBlockId, depValueMap, + groupKey); + } + + return 0; +} + +DenseMap> +buildDepUserMap(DenseMap &blocks) { + DenseMap> depUserMap; + for (auto &p : blocks) + for (Operation *op : p.second.ops) + for (Value operand : op->getOperands()) + depUserMap[operand].push_back(op); + return depUserMap; +} + +SmallVector +collectBufferValues(DenseMap> &depValueMap) { + SmallVector valueList; + SmallVector seenOps; + + for (auto &p : depValueMap) { + for (Value depVal : p.second) { + Operation *op = depVal.getDefiningOp(); + if (!op || llvm::is_contained(seenOps, op)) + continue; + seenOps.push_back(op); + + auto shapedType = dyn_cast(depVal.getType()); + if (!shapedType) + continue; + + valueList.push_back(depVal); + } + } + + return valueList; +} + +SmallVector +collectScalarDeps(DenseMap> &depValueMap, + DenseMap> &depUserMap) { + SmallVector scalarValueList; + + for (auto &p : depValueMap) { + for (Value depVal : p.second) { + if (isa(depVal)) + continue; + + Operation *depDefinedOp = depVal.getDefiningOp(); + if (!depDefinedOp) + continue; + + if (isa(depVal.getType())) + continue; + + auto userIt = depUserMap.find(depVal); + if (userIt == depUserMap.end()) + continue; + + int producerId = getSsbufferId(depDefinedOp); + if (producerId < 0) + continue; + + SmallVector depUsers = userIt->second; + bool hasCrossBlockUser = false; + for (Operation *depUser : depUsers) { + int userId = getSsbufferId(depUser); + if (userId < 0 || userId != producerId) { + hasCrossBlockUser = true; + break; + } + } + + if (hasCrossBlockUser) + scalarValueList.push_back(depVal); + } + } + + return scalarValueList; +} + +// Compute iteration index: (iv - lb) / step, used for buffer selection in +// double buffering +static Value getIterCount(OpBuilder &builder, mlir::scf::ForOp forOp, + Location loc, SmallVector *newOps, + int blockId = -1) { + auto i32Type = builder.getI32Type(); + Value iv = forOp.getInductionVar(); + Value lb = forOp.getLowerBound(); + Value step = forOp.getStep(); + Type ivType = iv.getType(); + + // Check if lower bound is a constant zero + bool lbIsZero = false; + if (auto constOp = lb.getDefiningOp()) { + if (auto intAttr = dyn_cast(constOp.getValue())) + lbIsZero = (intAttr.getInt() == 0); + } + + Value iterIdx; + if (lbIsZero) { + // Optimization: if lb is 0, use iv directly (or iv/step if step != 1) + bool stepIsOne = false; + if (auto constOp = step.getDefiningOp()) + if (auto intAttr = dyn_cast(constOp.getValue())) + stepIsOne = intAttr.getInt() == 1; + if (stepIsOne) { + iterIdx = iv; + } else { + iterIdx = builder.create(loc, iv, step); + if (newOps) + newOps->push_back(iterIdx.getDefiningOp()); + if (blockId >= 0) { + iterIdx.getDefiningOp()->setAttr(kAttrBlockId, + builder.getI32IntegerAttr(blockId)); + } + } + } else { + // General case: (iv - lb) / step + Value diff = builder.create(loc, iv, lb); + iterIdx = builder.create(loc, diff, step); + if (newOps) { + newOps->push_back(diff.getDefiningOp()); + newOps->push_back(iterIdx.getDefiningOp()); + } + if (blockId >= 0) { + diff.getDefiningOp()->setAttr(kAttrBlockId, + builder.getI32IntegerAttr(blockId)); + iterIdx.getDefiningOp()->setAttr(kAttrBlockId, + builder.getI32IntegerAttr(blockId)); + } + } + + // Cast to i32 if necessary + if (ivType == i32Type) + return iterIdx; + + Value result; + constexpr int bits32 = 32; + if (ivType.isIndex()) { + result = builder.create(loc, i32Type, iterIdx); + } else if (auto intType = dyn_cast(ivType)) { + // Extend or truncate integer types to i32 + if (intType.getWidth() < bits32) + result = builder.create(loc, i32Type, iterIdx); + else if (intType.getWidth() > bits32) + result = builder.create(loc, i32Type, iterIdx); + else + return iterIdx; + } else { + result = builder.create(loc, i32Type, iterIdx); + } + if (newOps) + newOps->push_back(result.getDefiningOp()); + if (blockId >= 0) { + result.getDefiningOp()->setAttr(kAttrBlockId, + builder.getI32IntegerAttr(blockId)); + } + return result; +} + +// Build if-else chain for buffer selection: if (idx==0) -> buf[0] else ... else +// -> buf[N-1] +static int +buildIfChain(OpBuilder &builder, Location loc, Value indexVal, + SmallVector &buffers, SmallVector &newOps, + SmallVector &outIfOps, + function_ref createOpFn, + function_ref yieldFn, + std::optional resultTypes = std::nullopt, + int blockId = -1) { + int N = buffers.size(); + auto types = resultTypes.value_or(mlir::TypeRange{}); + + // Create condition: index == 0 + Value zero = builder.create(loc, 0, 32); + Value firstCond = builder.create( + loc, mlir::arith::CmpIPredicate::eq, indexVal, zero); + auto firstIf = + builder.create(loc, types, firstCond, true, true); + + newOps.push_back(zero.getDefiningOp()); + newOps.push_back(firstCond.getDefiningOp()); + newOps.push_back(firstIf); + outIfOps.push_back(firstIf); + + // Tag counter operations with block_id + if (blockId >= 0) { + zero.getDefiningOp()->setAttr(kAttrBlockId, + builder.getI32IntegerAttr(blockId)); + firstCond.getDefiningOp()->setAttr(kAttrBlockId, + builder.getI32IntegerAttr(blockId)); + } + + // Then branch: use buffer[0] + builder.setInsertionPointToStart(&firstIf.getThenRegion().front()); + Operation *op0 = createOpFn(builder, loc, buffers[0].second); + if (!op0) { + return -1; + } + newOps.push_back(op0); + if (yieldFn) { + builder.create(loc, yieldFn(builder, loc, op0)); + } else { + builder.create(loc); + } + + // Build nested else-if chain for buffer[1] to buffer[N-2] + mlir::Block *currentElseBlock = &firstIf.getElseRegion().front(); + for (int i = 1; i < N - 1; ++i) { + builder.setInsertionPointToStart(currentElseBlock); + Value iVal = builder.create(loc, i, 32); + Value cond = builder.create( + loc, mlir::arith::CmpIPredicate::eq, indexVal, iVal); + auto nestedIf = + builder.create(loc, types, cond, true, true); + if (!nestedIf) { + return -1; + } + + newOps.push_back(iVal.getDefiningOp()); + newOps.push_back(cond.getDefiningOp()); + newOps.push_back(nestedIf); + outIfOps.push_back(nestedIf); + + // Tag counter operations with block_id + if (blockId >= 0) { + iVal.getDefiningOp()->setAttr(kAttrBlockId, + builder.getI32IntegerAttr(blockId)); + cond.getDefiningOp()->setAttr(kAttrBlockId, + builder.getI32IntegerAttr(blockId)); + } + + // Then branch: use buffer[i] + builder.setInsertionPointToStart(&nestedIf.getThenRegion().front()); + Operation *op = createOpFn(builder, loc, buffers[i].second); + if (!op) { + return -1; + } + newOps.push_back(op); + if (yieldFn) { + builder.create(loc, yieldFn(builder, loc, op)); + } else { + builder.create(loc); + } + + currentElseBlock = &nestedIf.getElseRegion().front(); + } + + // Final else branch: use buffer[N-1] + builder.setInsertionPointToStart(currentElseBlock); + Operation *opLast = createOpFn(builder, loc, buffers[N - 1].second); + if (!opLast) { + return -1; + } + newOps.push_back(opLast); + if (yieldFn) { + builder.create(loc, yieldFn(builder, loc, opLast)); + } else { + builder.create(loc); + } + + builder.setInsertionPointAfter(firstIf); + return 0; +} + +// Compute buffer index: iterCount % N +static Value computeBufferIndex(OpBuilder &builder, mlir::scf::ForOp forOp, + Location loc, int N, + SmallVector *newOps, + int blockId = -1) { + Value iterCount = getIterCount(builder, forOp, loc, newOps, blockId); + Value Nval = builder.create(loc, N, 32); + Value bufIdx = builder.create(loc, iterCount, Nval); + if (newOps) { + newOps->push_back(Nval.getDefiningOp()); + newOps->push_back(bufIdx.getDefiningOp()); + } + // Tag counter operations with block_id + if (blockId >= 0) { + MLIRContext *ctx = builder.getContext(); + Nval.getDefiningOp()->setAttr(kAttrBlockId, + builder.getI32IntegerAttr(blockId)); + bufIdx.getDefiningOp()->setAttr(kAttrBlockId, + builder.getI32IntegerAttr(blockId)); + } + return bufIdx; +} + +static SmallVector +insertProducerLogic(OpBuilder &builder, Value depVal, + SmallVector &buffers, mlir::scf::ForOp forOp) { + SmallVector newOps; + int N = buffers.size(); + Location loc = depVal.getLoc(); + // Single buffer producer logic + if (N == kBufferCountOne) { + Operation *producerOp = builder.create( + loc, mlir::TypeRange{}, depVal, buffers[0].second); + if (!producerOp) + return newOps; + newOps.push_back(producerOp); + return newOps; + } + + Value bufIdx = computeBufferIndex(builder, forOp, loc, N, &newOps); + SmallVector dummyOutIfOps; + if (buildIfChain( + builder, loc, bufIdx, buffers, newOps, dummyOutIfOps, + [&](OpBuilder &b, Location l, Value buffer) -> Operation * { + return b.create(l, mlir::TypeRange{}, depVal, buffer); + }, + nullptr) != 0) { + return {}; + } + return newOps; +} + +// Handle consumer when N==1 (directly return buffer) +static Operation *handleSingleBufferConsumer(OpBuilder &builder, Location loc, + SmallVector &buffers) { + auto memrefType = mlir::cast(buffers[0].second.getType()); + auto tensorType = mlir::RankedTensorType::get(memrefType.getShape(), + memrefType.getElementType()); + return builder.create( + loc, tensorType, buffers[0].second, + mlir::UnitAttr::get(builder.getContext()), + mlir::UnitAttr::get(builder.getContext())); +} + +// Helper function to create ToTensorOp +static mlir::bufferization::ToTensorOp createToTensorOp(OpBuilder &builder, + Location loc, + mlir::Type tensorType, + Value buffer) { + return builder.create( + loc, tensorType, buffer, mlir::UnitAttr::get(builder.getContext()), + mlir::UnitAttr::get(builder.getContext())); +} + +static int insertConsumerLogic(OpBuilder &builder, Value depVal, + SmallVector &buffers, + mlir::scf::ForOp forOp, + SmallVector &outIfOps, + int groupId = -1, int blockId = -1) { + SmallVector newOps; + int N = buffers.size(); + Location loc = builder.getInsertionPoint()->getLoc(); + + if (N == kBufferCountOne) { + Operation *consumerOp = handleSingleBufferConsumer(builder, loc, buffers); + outIfOps.push_back(consumerOp); + if (groupId >= 0) { + consumerOp->setAttr("ssbuffer.intraDeps", + builder.getI32ArrayAttr({groupId, 0})); + } + return 0; + } + + Value readIdx = computeBufferIndex(builder, forOp, loc, N, &newOps, blockId); + auto memrefType = mlir::cast(buffers[0].second.getType()); + auto tensorType = mlir::RankedTensorType::get(memrefType.getShape(), + memrefType.getElementType()); + mlir::TypeRange resultTypes(tensorType); + int ret = buildIfChain( + builder, loc, readIdx, buffers, newOps, outIfOps, + [&](OpBuilder &b, Location l, Value buffer) -> Operation * { + return createToTensorOp(b, l, tensorType, buffer); + }, + [&](OpBuilder &b, Location l, Operation *op) -> Value { + return cast(op).getResult(); + }, + resultTypes, blockId); + if (ret != 0) { + return ret; + } + if (groupId >= 0 && !outIfOps.empty()) { + outIfOps.front()->setAttr("ssbuffer.intraDeps", + builder.getI32ArrayAttr({groupId, 0})); + } + return 0; +} + +static void addBlockAttrForOps(SmallVector &newOps, int blockId, + OpBuilder &builder) { + auto attr = builder.getI32IntegerAttr(blockId); + for (auto *op : newOps) + op->setAttr(kAttrBlockId, attr); +} + +// Add dep_mark attribute to operation +static void addDepMarkAttr(Operation *op, int depMark, OpBuilder &builder) { + if (auto existingAttr = + op->getAttrOfType("ssbuffer.dep_mark")) { + SmallVector marks; + for (auto attr : existingAttr) + marks.push_back(cast(attr).getInt()); + marks.push_back(depMark); + op->setAttr("ssbuffer.dep_mark", builder.getI32ArrayAttr(marks)); + } else { + op->setAttr("ssbuffer.dep_mark", builder.getI32ArrayAttr({depMark})); + } +} + +// Add ssbuffer.intra_buffer attribute to buffer operations +// Only tag scf::IfOp (multi-buffer), materialize_in_destination/to_tensor +// (single buffer) +static void addIntraBufferAttr(SmallVector &ops, + OpBuilder &builder) { + for (auto *op : ops) { + if (isa(op) || isa(op) || + isa(op)) { + op->setAttr("ssbuffer.intra_buffer", builder.getUnitAttr()); + } + } +} + +// Collect cross-block user operations +static SmallVector +collectCrossBlockUsers(Value depVal, int producerId, + DenseMap> &depUserMap) { + SmallVector crossBlockUsers; + + auto userIt = depUserMap.find(depVal); + if (userIt == depUserMap.end()) + return crossBlockUsers; + + for (Operation *depUser : userIt->second) { + int userId = getSsbufferId(depUser); + if ((userId < 0 || userId != producerId) && + isInsideMainLoopForOpTraverse(depUser)) + crossBlockUsers.push_back(depUser); + } + return crossBlockUsers; +} + +static void +markScalarDeps(SmallVector &scalarValueList, + DenseMap> &depUserMap, + OpBuilder &builder, int startDepMark) { + int nextDepMark = startDepMark; + + for (Value depVal : scalarValueList) { + Operation *depDefinedOp = depVal.getDefiningOp(); + if (!depDefinedOp) + continue; + + if (!isInsideMainLoopForOp(depDefinedOp)) + continue; + + int producerId = getSsbufferId(depDefinedOp); + if (producerId < 0) + continue; + auto crossBlockUsers = + collectCrossBlockUsers(depVal, producerId, depUserMap); + if (crossBlockUsers.empty()) + continue; + + int depMark = nextDepMark++; + // Add depmark to producer + addDepMarkAttr(depDefinedOp, depMark, builder); + // Add depmark to consumer + for (Operation *depUser : crossBlockUsers) { + addDepMarkAttr(depUser, depMark, builder); + } + } +} + +// Process producer and consumer for a single dependency value +static int processDepVal(Value depVal, mlir::scf::ForOp mainLoopForOp, + BufferMap &bufferMap, + DenseMap> &depUserMap, + OpBuilder &globalBuilder, int producerId, + int groupId) { + Operation *depDefinedOp = depVal.getDefiningOp(); + if (!depDefinedOp) + return 0; + + SmallVector &buffers = bufferMap[depVal]; + + auto userIt = depUserMap.find(depVal); + if (userIt == depUserMap.end()) + return 0; + SmallVector depUsers = userIt->second; + + // Create producer + OpBuilder producedBuffers(mainLoopForOp.getContext()); + producedBuffers.setInsertionPointAfter(depDefinedOp); + SmallVector producerNewOps = + insertProducerLogic(producedBuffers, depVal, buffers, mainLoopForOp); + addBlockAttrForOps(producerNewOps, producerId, globalBuilder); + // Tag producer: N > 1 only tag scf.if, N == 1 tag materialize_in_destination + if (buffers.size() > kBufferCountOne) { + for (auto *op : producerNewOps) { + if (isa(op)) { + op->setAttr("ssbuffer.intra_buffer", globalBuilder.getUnitAttr()); + } + } + } else { + addIntraBufferAttr(producerNewOps, globalBuilder); + } + + // Process each consumer + for (Operation *depUser : depUsers) { + int userBlockId = getSsbufferId(depUser); + if (userBlockId < 0 || userBlockId == producerId) + continue; + + OpBuilder consumedBuilder(mainLoopForOp.getContext()); + consumedBuilder.setInsertionPoint(depUser); + SmallVector resultIfOps; + int ret = + insertConsumerLogic(consumedBuilder, depVal, buffers, mainLoopForOp, + resultIfOps, groupId, userBlockId); + if (ret != 0) + return -1; + + if (resultIfOps.empty()) + continue; + // Tag consumer with block_id + addBlockAttrForOps(resultIfOps, userBlockId, globalBuilder); + // Tag consumer: N > 1 only tag scf.if, N == 1 tag to_tensor + if (buffers.size() > kBufferCountOne) { + for (auto *op : resultIfOps) { + if (isa(op)) { + op->setAttr("ssbuffer.intra_buffer", globalBuilder.getUnitAttr()); + } + } + } else { + addIntraBufferAttr(resultIfOps, globalBuilder); + } + + Operation *resultIf = resultIfOps.back(); + Value selectedBuffer = resultIf->getResult(0); + + for (OpOperand &use : depUser->getOpOperands()) { + if (use.get() == depVal) + use.set(selectedBuffer); + } + } + return 0; +} + +// Process cross-block tensor dependencies for double buffering +static int processTensorDependencies( + mlir::scf::ForOp mainLoopForOp, DenseMap &blocks, + DenseMap> &depValueMap, + DenseMap> &depUserMap, BufferMap &bufferMap, + OpBuilder &globalBuilder, int &groupId) { + SmallVector seenOps; + + for (auto &blockPair : blocks) { + Value blockKey = blockPair.first; + auto depIt = depValueMap.find(blockKey); + if (depIt == depValueMap.end()) + continue; + + SmallVector &depValues = depIt->second; + + for (Value depVal : depValues) { + // Skip if already processed + if (llvm::is_contained(seenOps, depVal.getDefiningOp())) + continue; + seenOps.push_back(depVal.getDefiningOp()); + + // Validate dependency value (skip BlockArgument, null definingOp, + // non-ShapedType) + if (isa(depVal) || !depVal.getDefiningOp() || + !isa(depVal.getType())) + continue; + + auto userIt = depUserMap.find(depVal); + if (userIt == depUserMap.end()) + continue; + SmallVector depUsers = userIt->second; + + int producerId = getSsbufferId(depVal.getDefiningOp()); + if (producerId < 0) + continue; // Skip if producer has no ssbuffer attribute + + // Check if all users are in the same block + bool allUsersSameBlock = true; + for (Operation *depUser : depUsers) { + int userId = getSsbufferId(depUser); + if (userId < 0 || userId != producerId) { + allUsersSameBlock = false; + break; + } + } + if (allUsersSameBlock) + continue; + + // Process cross-block dependency with double buffering + if (processDepVal(depVal, mainLoopForOp, bufferMap, depUserMap, + globalBuilder, producerId, groupId) != 0) + return -1; + groupId++; + } + } + return 0; +} + +static BufferMap insertBuffersBeforeFor(mlir::scf::ForOp forOp, + SmallVector &valueList, + OpBuilder &builder, int groupId) { + BufferMap bufferMap; + Block *parentBlock = forOp->getBlock(); + OpBuilder insertedBuffers(builder.getContext()); + insertedBuffers.setInsertionPoint(parentBlock, forOp->getIterator()); + + using BufferCountManager = mlir::triton::BufferCountManager; + int bufNum = BufferCountManager::getInstance().getBufferCountByType( + BufferCountManager::DepType::IntraCore); + + for (Value depVal : valueList) { + ShapedType shapedType = cast(depVal.getType()); + Type elemType = shapedType.getElementType(); + AddressSpace addrSpace = AddressSpace::UB; + + SmallVector buffers; + for (int i = 0; i < bufNum; ++i) { + MemRefType memrefType = MemRefType::get( + shapedType.getShape(), elemType, MemRefLayoutAttrInterface{}, + AddressSpaceAttr::get(insertedBuffers.getContext(), addrSpace)); + + auto allocOp = + insertedBuffers.create(forOp.getLoc(), memrefType); + + auto genericType = MemRefType::get(shapedType.getShape(), elemType, + MemRefLayoutAttrInterface{}, 0u); + + auto casted = insertedBuffers.create( + forOp.getLoc(), genericType, allocOp.getResult()); + + casted->setAttr("ssbuffer.intraDeps", + insertedBuffers.getI32ArrayAttr({groupId, 1})); + + buffers.push_back({casted.getResult(), casted.getResult()}); + } + + bufferMap[depVal] = buffers; + groupId++; + } + + return bufferMap; +} + +static bool +hasMemrefDepValue(DenseMap> &depValueMap) { + for (auto &p : depValueMap) { + for (Value depVal : p.second) { + if (isa(depVal.getType())) + return true; + } + } + return false; +} + +static int addInnerMultiBuffer(mlir::scf::ForOp mainLoopForOp, + OpBuilder &builder, scope::ScopeOp vectorScope, + int &groupId) { + DenseMap blocks; + DenseMap> depValueMap; + if (collectInnerBlockInfo(mainLoopForOp, blocks, depValueMap) != 0) + return -1; + + if (blocks.empty()) + return -1; + + // If memref-type dependency values exist, skip double buffer processing + if (hasMemrefDepValue(depValueMap)) { + vectorScope->setAttr("ssbuffer.skip", builder.getUnitAttr()); + return 0; + } + + auto depUserMap = buildDepUserMap(blocks); + + auto valueList = collectBufferValues(depValueMap); + auto bufferMap = + insertBuffersBeforeFor(mainLoopForOp, valueList, builder, groupId); + + auto scalarValueList = collectScalarDeps(depValueMap, depUserMap); + + OpBuilder globalBuilder(mainLoopForOp.getContext()); + markScalarDeps(scalarValueList, depUserMap, globalBuilder, 1); + + if (processTensorDependencies(mainLoopForOp, blocks, depValueMap, depUserMap, + bufferMap, globalBuilder, groupId) != 0) { + return -1; + } + + return 0; +} + +void AddMultiBufferInnerScopePass::getDependentDialects( + DialectRegistry ®istry) const { + registry + .insert(); +} + +void AddMultiBufferInnerScopePass::runOnOperation() { + auto module = getOperation(); + OpBuilder builder(module.getContext()); + + LDBG("Enter pass."); + + module.walk([&](scope::ScopeOp scope) -> WalkResult { + // Step 1: Check if scope has coreType attribute + auto coreTypeAttr = + scope->getAttrOfType(hivm::TCoreTypeAttr::name); + if (!coreTypeAttr) + return WalkResult::advance(); + + // Step 2: Check if core type is VECTOR + hivm::TCoreType coreType = coreTypeAttr.getTcoretype(); + if (coreType != hivm::TCoreType::VECTOR) { + LDBG("Not vector scope"); + return WalkResult::advance(); + } + + // Step 3: Collect all forOps with main_loop attribute + SmallVector mainLoopForOps; + int foundCount = + collectMainLoopsRecursively(scope.getBodyRegion(), mainLoopForOps); + if (foundCount < 0) { + LDBG("collectMainLoopsRecursively failed"); + signalPassFailure(); + return WalkResult::interrupt(); + } + if (foundCount == 0) + return WalkResult::advance(); + + // Step 4: Process each main_loop forOp + int groupId = 0; + for (scf::ForOp mainLoopForOp : mainLoopForOps) { + scf::ForOp nestedMainloop = findNestedMainloopInForOp(mainLoopForOp); + if (nestedMainloop) { + LDBG("Nested main_loop found, this is not allowed"); + signalPassFailure(); + return WalkResult::interrupt(); + } + if (addInnerMultiBuffer(mainLoopForOp, builder, scope, groupId) != 0) { + LDBG("addInnerMultiBuffer failed"); + signalPassFailure(); + return WalkResult::interrupt(); + } + } + + return WalkResult::advance(); + }); + + LDBG("Process successfully"); +} + +std::unique_ptr> createAddMultiBufferInnerScopePass() { + return std::make_unique(); +} + +void registerAddMultiBufferInnerScopePasses() { + registerPass([]() -> std::unique_ptr { + return createAddMultiBufferInnerScopePass(); + }); +} + +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferOuterScope.cpp b/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferOuterScope.cpp new file mode 100644 index 0000000000..9f8ab10db6 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/AllocMultiCache/AddMultiBufferOuterScope.cpp @@ -0,0 +1,1052 @@ +#include "ascend/include/DynamicCVPipeline/AllocMultiCache/AddMultiBufferOuterScope.h" + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallSet.h" +#include "llvm/Support/Debug.h" + +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/Scope/IR/Scope.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" + +#include "ascend/include/DynamicCVPipeline/Common/BufferCountManager.h" +#include "ascend/include/DynamicCVPipeline/Common/FlagIdManager.h" + +static constexpr const char *DEBUG_TYPE = "AddMultiBufferOuterScope"; +#define LDBG(...) \ + LLVM_DEBUG(llvm::dbgs() << " [" << DEBUG_TYPE << "] " << __VA_ARGS__) + +using namespace mlir; +using namespace triton; +using namespace hivm; + +namespace mlir { +namespace triton { + +// Maximum number of flag allocation attempts per transfer group +static constexpr int kMaxFlagAttempts = 16; + +// --- Attribute helpers --- + +static int getFlagFromSyncOp(Operation *op) { + if (auto attr = op->getAttrOfType("flag_id")) { + return attr.getInt(); + } + if (auto attr = op->getAttrOfType("static_flag_id")) { + return attr.getInt(); + } + if (auto attr = op->getAttrOfType("flag")) { + return attr.getInt(); + } + return -1; +} + +static int getBlockId(Operation *op) { + if (auto attr = op->getAttrOfType("ssbuffer.block_id")) + return attr.getInt(); + return -1; +} + +static int getTransferId(Operation *op) { + if (auto attr = op->getAttrOfType("ssbuffer.transfer_id")) + return attr.getInt(); + return -1; +} + +// --- Address space helpers --- + +static bool isInVectorScope(Operation *op) { + auto scopeOp = op->getParentOfType(); + if (!scopeOp) { + return false; + } + if (auto tcoreAttr = scopeOp->getAttrOfType("hivm.tcore_type")) + return tcoreAttr.getTcoretype() == TCoreType::VECTOR; + return false; +} + +// --- main_loop attribute helpers --- + +/// Check if forOp (or its terminator) has ssbuffer.main_loop attribute +static bool forOpHasMainLoopAttr(scf::ForOp forOp) { + if (forOp->hasAttr("ssbuffer.main_loop")) { + return true; + } + Operation *terminator = forOp.getBody()->getTerminator(); + return terminator && terminator->hasAttr("ssbuffer.main_loop"); +} + +/// Check if a sync op's direct parent has ssbuffer.main_loop attribute +static bool parentOpHasMainLoopAttr(Operation *syncOp) { + if (!syncOp) { + return false; + } + Operation *parent = syncOp->getParentOp(); + if (!parent) { + return false; + } + if (auto forOp = dyn_cast(parent)) { + return forOpHasMainLoopAttr(forOp); + } + return false; +} + +// --- Operation search helpers --- + +/// Find sync op with a specific flag, searching forward or backward in a block +static Operation *findSyncOpWithFlag(Block *block, Operation *start, int flag, + bool forward, bool wantWait) { + if (!block) { + return nullptr; + } + auto it = start->getIterator(); + if (forward) { + for (auto e = block->end(); it != e; ++it) { + Operation *op = &*it; + if (!(isa(op) || isa(op))) { + continue; + } + if (getFlagFromSyncOp(op) != flag) { + continue; + } + if (wantWait && isa(op)) { + return op; + } + if (!wantWait && isa(op)) { + return op; + } + } + } else { + if (it == block->begin()) { + return nullptr; + } + do { + --it; + Operation *op = &*it; + if (!(isa(op) || isa(op))) { + continue; + } + if (getFlagFromSyncOp(op) != flag) { + continue; + } + if (wantWait && isa(op)) { + return op; + } + if (!wantWait && isa(op)) { + return op; + } + } while (it != block->begin()); + } + return nullptr; +} + +/// Find the to_tensor op after a given op in the same block +static Operation *findToTensorAfter(Block *block, Operation *start) { + if (!block) { + return nullptr; + } + auto it = start->getIterator(); + for (auto e = block->end(); it != e; ++it) { + if (isa(&*it)) { + return &*it; + } + } + return nullptr; +} + +// ============================================================================ +// Step 1: Collect transfer group info +// ============================================================================ + +/// Collect all ops with ssbuffer.transfer_id, grouped by transfer_id +static int +collectOpsByTransferId(ModuleOp module, + DenseMap> &opsByTid) { + module.walk([&](Operation *op) { + if (!op->hasAttr("ssbuffer.transfer_id")) { + return; + } + int tid = getTransferId(op); + if (tid >= 0) { + opsByTid[tid].push_back(op); + } + }); + LDBG("Collected " << opsByTid.size() << " transfer groups"); + + for (auto &p : opsByTid) { + LDBG(" tid=" << p.first << " has " << p.second.size() << " ops"); + DenseMap blockIdCount; + for (auto *op : p.second) { + int bid = getBlockId(op); + blockIdCount[bid]++; + } + for (auto &bp : blockIdCount) { + LDBG(" block_id=" << bp.first << ": " << bp.second << " ops"); + } + } + return 0; +} + +/// Collect alloc/mark pairs (independent of block_id and main_loop) +static int collectBufferAllocs(const SmallVector &ops, + BufferAllocInfo &info) { + SmallVector allocs; + SmallVector marks; + + for (Operation *op : ops) { + if (isa(op)) { + allocs.push_back(op); + } else if (isa(op)) { + marks.push_back(op); + } + } + + LDBG("collectBufferAllocs: allocs=" << allocs.size() + << ", marks=" << marks.size()); + + // Pair in order: sender first, receiver second + if (!allocs.empty()) { + info.sender.allocOp = allocs[0]; + } + if (allocs.size() > 1) { + info.receiver.allocOp = allocs[1]; + } + if (!marks.empty()) { + info.sender.markOp = marks[0]; + } + if (marks.size() > 1) { + info.receiver.markOp = marks[1]; + } + + return 0; +} + +/// Collect extra sync ops (parent has no main_loop), paired by flag +static int collectExtraSync(const SmallVector &ops, + int originalFlag, ExtraSyncInfo &info) { + SmallVector extraSets; + SmallVector extraWaits; + + for (Operation *op : ops) { + if (!(isa(op) || isa(op))) { + continue; + } + + bool hasMainLoop = parentOpHasMainLoopAttr(op); + LDBG("sync op: flag=" << getFlagFromSyncOp(op) + << ", block_id=" << getBlockId(op) + << ", parentHasMainLoop=" << hasMainLoop); + + if (!hasMainLoop) { + if (isa(op)) { + extraSets.push_back(op); + } else if (isa(op)) { + extraWaits.push_back(op); + } + } + } + + // Match by flag + for (auto *setOp : extraSets) { + if (getFlagFromSyncOp(setOp) != originalFlag) { + continue; + } + for (auto *waitOp : extraWaits) { + if (getFlagFromSyncOp(waitOp) != originalFlag) { + continue; + } + info.setOp = setOp; + info.waitOp = waitOp; + LDBG("Extra sync pair: set(flag=" << originalFlag + << ", block_id=" << getBlockId(setOp) + << "), wait(flag=" << originalFlag + << ", block_id=" << getBlockId(waitOp)); + return 0; + } + } + + // Fallback: use first available pair if exact match not found + if (!extraSets.empty() && !extraWaits.empty()) { + info.setOp = extraSets.front(); + info.waitOp = extraWaits.front(); + } + + return 0; +} + +/// Collect transfer chain ops (parent has main_loop) +static int collectTransferChains(const SmallVector &ops, + int originalFlag, TransferChainInfo &info) { + for (Operation *op : ops) { + if ((isa(op) || isa(op)) || + !op->getBlock()) { + continue; + } + if (!parentOpHasMainLoopAttr(op)) { + continue; + } + + Block *block = op->getBlock(); + + if (isa(op)) { + info.sender.transferOp = op; + info.sender.waitOp = + findSyncOpWithFlag(block, op, originalFlag, false, true); + info.sender.setOp = + findSyncOpWithFlag(block, op, originalFlag, true, false); + LDBG("Sender chain (CUBE): fixpipe, flag=" << originalFlag); + } else if (isa(op)) { + info.sender.transferOp = op; + info.sender.waitOp = + findSyncOpWithFlag(block, op, originalFlag, false, true); + info.sender.setOp = + findSyncOpWithFlag(block, op, originalFlag, true, false); + LDBG("Sender chain (VECTOR): hir.copy, flag=" << originalFlag); + } else if (isa(op) && isInVectorScope(op)) { + info.receiver.transferOp = op; + info.receiver.waitOp = + findSyncOpWithFlag(block, op, originalFlag, false, true); + info.receiver.setOp = + findSyncOpWithFlag(block, op, originalFlag, true, false); + info.receiver.toTensorOp = findToTensorAfter(block, op); + LDBG("Receiver chain (VECTOR): memory_space_cast, flag=" << originalFlag); + } else if (isa(op)) { + info.receiver.transferOp = op; + info.receiver.waitOp = + findSyncOpWithFlag(block, op, originalFlag, false, true); + info.receiver.setOp = + findSyncOpWithFlag(block, op, originalFlag, true, false); + LDBG("Receiver chain (CUBE): convert_layout, flag=" << originalFlag); + } + } + + return 0; +} + +/// Build TransferGroupInfo for a single transfer_id +static int buildTransferGroupData(int tid, const SmallVector &ops, + FlagIdManager &flagIdMgr, + TransferGroupInfo &info) { + info.tid = tid; + + LDBG("Building group tid=" << tid << ", ops=" << ops.size()); + + // 1. Collect buffer alloc/mark pairs + BufferAllocInfo bufInfo; + if (collectBufferAllocs(ops, bufInfo)) { + return -1; + } + info.senderBuf = bufInfo.sender; + info.receiverBuf = bufInfo.receiver; + LDBG("Sender buffer: " << (info.senderBuf.allocOp ? "alloc" : "none") << " + " + << (info.senderBuf.markOp ? "mark" : "none")); + LDBG("Receiver buffer: " << (info.receiverBuf.allocOp ? "alloc" : "none") + << " + " + << (info.receiverBuf.markOp ? "mark" : "none")); + + // 2. Determine original flag + for (Operation *op : ops) { + if ((isa(op) || isa(op))) { + int f = getFlagFromSyncOp(op); + if (f >= 0) { + info.originalFlag = f; + break; + } + } + } + + // 3. Collect extra sync (parent has no main_loop) + ExtraSyncInfo extraInfo; + if (collectExtraSync(ops, info.originalFlag, extraInfo)) { + return -1; + } + info.extraSyncSetOp = extraInfo.setOp; + info.extraSyncWaitOp = extraInfo.waitOp; + if (extraInfo.setOp && extraInfo.waitOp) { + LDBG("Extra sync: set(block_id=" << getBlockId(extraInfo.setOp) + << "), wait(block_id=" + << getBlockId(extraInfo.waitOp)); + } else { + LDBG("Extra sync: not found"); + } + + // 4. Collect transfer chain (parent has main_loop) + TransferChainInfo chainInfo; + if (collectTransferChains(ops, info.originalFlag, chainInfo)) { + return -1; + } + info.senderChain = chainInfo.sender; + info.receiverChain = chainInfo.receiver; + + // 5. Determine direction + if (info.senderChain.transferOp) { + if (isa(info.senderChain.transferOp)) { + info.isCtoV = true; + } else if (isa(info.senderChain.transferOp)) { + info.isCtoV = false; + } + } + + // For C→V transfer, sender uses receiver's buffer (the second alloc) + if (info.isCtoV && info.senderBuf.allocOp && info.receiverBuf.allocOp) { + LDBG("C→V transfer: swapping sender/receiver buffers"); + std::swap(info.senderBuf, info.receiverBuf); + } + + // 6. Acquire output flag + for (int attempt = 0; attempt < kMaxFlagAttempts; ++attempt) { + int64_t pf = flagIdMgr.acquireId(nullptr); + if (pf == FlagIdManager::INVALID_FLAG_ID) { + break; + } + if (pf != info.originalFlag) { + info.outputFlag = static_cast(pf); + break; + } + } + + if (info.senderChain.transferOp || info.receiverChain.transferOp) { + LDBG("Direction: " << (info.isCtoV ? "C→V" : "V→C") + << ", flag=" << info.originalFlag + << ", outputFlag=" << info.outputFlag); + } + + return 0; +} + +/// Collect TransferGroupInfo for all transfer groups +static int collectTransferGroupData( + ModuleOp module, DenseMap> &opsByTid, + FlagIdManager &flagIdMgr, DenseMap &groups) { + for (auto &p : opsByTid) { + TransferGroupInfo info; + if (buildTransferGroupData(p.first, p.second, flagIdMgr, info)) { + continue; + } + if ((info.senderChain.transferOp || info.receiverChain.transferOp) && + info.outputFlag >= 0) { + groups[p.first] = info; + } + } + + // Output flag reuse: groups with same (originalFlag, direction) share an + // output flag + std::map, int> outputFlagByKey; + for (auto &p : groups) { + auto &g = p.second; + auto key = std::make_pair(g.originalFlag, g.isCtoV); + auto it = outputFlagByKey.find(key); + if (it != outputFlagByKey.end()) { + g.outputFlag = it->second; + LDBG("Group tid=" << g.tid << " reuses outputFlag=" << g.outputFlag + << " (shared originalFlag=" << g.originalFlag << ")"); + } else { + outputFlagByKey[key] = g.outputFlag; + LDBG("Group tid=" << g.tid + << " gets new shared outputFlag=" << g.outputFlag + << " for originalFlag=" << g.originalFlag); + } + } + + LDBG("=== Step 1 Summary ==="); + LDBG("Transfer groups: " << groups.size()); + for (auto &p : groups) { + LDBG("Group tid=" << p.first + << ", dir=" << (p.second.isCtoV ? "C→V" : "V→C") + << ", flag=" << p.second.originalFlag + << ", outputFlag=" << p.second.outputFlag); + if (p.second.senderChain.transferOp) + LDBG(" Sender: " + << p.second.senderChain.transferOp->getName().getStringRef()); + if (p.second.receiverChain.transferOp) + LDBG(" Receiver: " + << p.second.receiverChain.transferOp->getName().getStringRef()); + } + + return 0; +} + +// ============================================================================ +// Step 2: Create output buffers +// ============================================================================ + +static constexpr int kMaxTcbSearch = 100; + +static int allocateNewTcbId(int startFrom, std::set &usedTcbIds) { + for (int id = startFrom; id < kMaxTcbSearch; ++id) { + if (!usedTcbIds.count(id)) { + usedTcbIds.insert(id); + return id; + } + } + return -1; +} + +/// Create an output buffer for an input/output buffer pair +static int createOutputBufferPair(Operation *inputAllocOp, int tid, int tcbId, + Value &inputBuffer, Value &outputBuffer, + OpBuilder &builder, bool isSender) { + if (!inputAllocOp) { + return -1; + } + + Location loc = builder.getUnknownLoc(); + + inputBuffer = inputAllocOp->getResult(0); + auto memRefType = dyn_cast(inputBuffer.getType()); + if (!memRefType) { + return -1; + } + + int origBlockId = getBlockId(inputAllocOp); + int outputBlockId = origBlockId; + + builder.setInsertionPointAfter(inputAllocOp); + auto outputAlloc = builder.create(loc, memRefType); + outputAlloc->setAttr("ssbuffer.block_id", + builder.getI32IntegerAttr(outputBlockId)); + outputAlloc->setAttr("ssbuffer.transfer_id", builder.getI32IntegerAttr(tid)); + outputBuffer = outputAlloc.getResult(); + + if (!isSender) { + outputAlloc->setAttr("ssbuffer.crossDeps", + builder.getArrayAttr({builder.getI32IntegerAttr(tid), + builder.getI32IntegerAttr(1)})); + } + + auto outputMark = builder.create(loc, outputBuffer); + outputMark->setAttr("effects", builder.getStrArrayAttr({"write", "read"})); + outputMark->setAttr("ssbuffer.block_id", + builder.getI32IntegerAttr(outputBlockId)); + outputMark->setAttr("ssbuffer.transfer_id", builder.getI32IntegerAttr(tid)); + outputMark->setAttr( + "hivm.tightly_coupled_buffer", + hivm::HIVMTightlyCoupledBufferAttr::get(builder.getContext(), tcbId)); + LDBG("Created " << (isSender ? "sender" : "receiver") + << " output buffer: block_id=" << outputBlockId + << ", tcb_id=" << tcbId); + return 0; +} + +static constexpr unsigned kBits32 = 32; + +static int attachSsbufferTags(Operation *op, int blockId, int transferId) { + MLIRContext *ctx = op->getContext(); + op->setAttr("ssbuffer.block_id", + IntegerAttr::get(IntegerType::get(ctx, kBits32), blockId)); + op->setAttr("ssbuffer.transfer_id", + IntegerAttr::get(IntegerType::get(ctx, kBits32), transferId)); + return 0; +} + +static hivm::SyncBlockSetOp createOutputSyncSetOp(Operation *origSetOp, + int outputFlag, int tid, + OpBuilder &builder) { + auto setOp = cast(origSetOp); + builder.setInsertionPointAfter(origSetOp); + auto newSetOp = builder.create( + setOp.getLoc(), setOp.getTcoreType(), setOp.getTpipe(), setOp.getPipe(), + builder.getI64IntegerAttr(outputFlag)); + attachSsbufferTags(newSetOp.getOperation(), getBlockId(setOp), tid); + return newSetOp; +} + +static hivm::SyncBlockWaitOp createOutputSyncWaitOp(Operation *origWaitOp, + int outputFlag, int tid, + OpBuilder &builder) { + auto waitOp = cast(origWaitOp); + builder.setInsertionPointAfter(origWaitOp); + auto newWaitOp = builder.create( + waitOp.getLoc(), waitOp.getTcoreType(), waitOp.getTpipe(), + waitOp.getPipe(), builder.getI64IntegerAttr(outputFlag)); + attachSsbufferTags(newWaitOp.getOperation(), getBlockId(waitOp), tid); + return newWaitOp; +} + +/// Create output buffer for a single transfer group, with output flag sync ops +static int createOutputBufferForGroup(TransferGroupInfo &g, + OpBuilder &builder) { + if (createOutputBufferPair(g.senderBuf.allocOp, g.tid, g.tcbId, + g.senderInputBuffer, g.senderOutputBuffer, builder, + true)) { + return -1; + } + if (createOutputBufferPair(g.receiverBuf.allocOp, g.tid, g.tcbId, + g.receiverInputBuffer, g.receiverOutputBuffer, + builder, false)) { + return -1; + } + // Insert output sync set at extra_sync position + if (g.extraSyncSetOp) { + createOutputSyncSetOp(g.extraSyncSetOp, g.outputFlag, g.tid, builder); + LDBG("Created output sync set with flag=" << g.outputFlag << " at block_id=" + << getBlockId(g.extraSyncSetOp) + << " (sender scope)"); + } + + // Insert output sync wait at extra_sync position + Operation *outputWaitInsertOp = + g.extraSyncWaitOp ? g.extraSyncWaitOp : g.receiverChain.waitOp; + if (outputWaitInsertOp) { + createOutputSyncWaitOp(outputWaitInsertOp, g.outputFlag, g.tid, builder); + LDBG("Created output sync wait with flag=" + << g.outputFlag << " at block_id=" << getBlockId(outputWaitInsertOp) + << " (receiver scope)"); + } + return 0; +} + +/// Create output buffers for all transfer groups +static int createOutputBuffers(DenseMap &groups, + ModuleOp module) { + OpBuilder builder(module.getContext()); + std::set usedTcbIds; + + // Collect existing tcb ids + module.walk([&](Operation *op) { + if (auto tcbAttr = op->getAttrOfType( + "hivm.tightly_coupled_buffer")) { + auto id = tcbAttr.getId(); + if (id.has_value()) { + LDBG("Found mark op with tcb_id=" << id.value()); + usedTcbIds.insert(id.value()); + } + } + }); + + LDBG("=== Step 2: Creating output buffers ==="); + { + std::string ids; + llvm::raw_string_ostream os(ids); + for (int id : usedTcbIds) + os << id << " "; + LDBG("Collected existing tcb_ids: " << ids); + } + + int maxExistingTcbId = usedTcbIds.empty() ? 0 : *usedTcbIds.rbegin(); + LDBG("Max existing tcb_id: " << maxExistingTcbId); + + int nextTcbId = maxExistingTcbId + 1; + + for (auto &p : groups) { + TransferGroupInfo &g = p.second; + LDBG("Group tid=" << g.tid << " (" << (g.isCtoV ? "C→V" : "V→C") << ")"); + + g.tcbId = allocateNewTcbId(nextTcbId, usedTcbIds); + LDBG("Allocated tcb_id=" << g.tcbId); + + nextTcbId = g.tcbId + 1; + + createOutputBufferForGroup(g, builder); + } + return 0; +} + +/// Tag consumer-side alloc and transferOp with crossDeps marks +static int addConsumerCrossDepsTags(TransferGroupInfo &g, ModuleOp module) { + bool consumerIsVector = g.isCtoV; + auto &consumerBuf = consumerIsVector ? g.receiverBuf : g.senderBuf; + auto &consumerChain = consumerIsVector ? g.receiverChain : g.senderChain; + + OpBuilder builder(module.getContext()); + + if (consumerBuf.allocOp) { + consumerBuf.allocOp->setAttr( + "ssbuffer.crossDeps", + builder.getArrayAttr( + {builder.getI32IntegerAttr(g.tid), builder.getI32IntegerAttr(1)})); + } + if (consumerChain.transferOp) { + consumerChain.transferOp->setAttr( + "ssbuffer.crossDeps", + builder.getArrayAttr( + {builder.getI32IntegerAttr(g.tid), builder.getI32IntegerAttr(0)})); + } + return 0; +} + +// ============================================================================ +// Step 3: Add polling control flow +// ============================================================================ + +/// Set ssbuffer tags on an op +static int setSsbufferTags(Operation *op, OpBuilder &builder, int blockId, + int tid) { + op->setAttr("ssbuffer.block_id", builder.getI32IntegerAttr(blockId)); + op->setAttr("ssbuffer.transfer_id", builder.getI32IntegerAttr(tid)); + return 0; +} + +/// Create polling condition: (iter / step) % 2 == 0 (true=input, false=output) +static Value createPollingCondition(scf::ForOp forOp, OpBuilder &builder, + int blockId, int tid) { + Location loc = forOp.getLoc(); + Value iterVar = forOp.getInductionVar(); + Value step = forOp.getStep(); + + auto divOp = builder.create(loc, iterVar, step); + setSsbufferTags(divOp.getOperation(), builder, blockId, tid); + + Type counterType = divOp.getResult().getType(); + int bitWidth = counterType.getIntOrFloatBitWidth(); + auto c2Val = builder.create(loc, 2, bitWidth); + setSsbufferTags(c2Val.getOperation(), builder, blockId, tid); + auto remOp = + builder.create(loc, divOp.getResult(), c2Val.getResult()); + setSsbufferTags(remOp.getOperation(), builder, blockId, tid); + + auto c0Val = builder.create(loc, 0, bitWidth); + auto cmpOp = builder.create( + loc, arith::CmpIPredicate::eq, remOp.getResult(), c0Val.getResult()); + setSsbufferTags(cmpOp.getOperation(), builder, blockId, tid); + setSsbufferTags(c0Val.getOperation(), builder, blockId, tid); + + return cmpOp.getResult(); +} + +/// Wrap a sync op (wait/set) in scf.if: then=clone original, else=create +/// alternative +template +static Operation *wrapSyncOpWithScfIf( + Operation *op, Value cond, int outputFlag, OpBuilder &builder, + std::function createAltFn) { + static_assert(std::is_same::value || + std::is_same::value, + "OpTy must be SyncBlockWaitOp or SyncBlockSetOp"); + + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPoint(op); + + Location loc = op->getLoc(); + auto ifOp = builder.create(loc, TypeRange{}, cond, + true /* withElseRegion */); + ifOp->setAttr("ssbuffer.block_id", builder.getI32IntegerAttr(getBlockId(op))); + ifOp->setAttr("ssbuffer.cross_buffer", builder.getI32IntegerAttr(1)); + + // then branch: clone original op + auto thenBuilder = ifOp.getThenBodyBuilder(); + Operation *cloned = thenBuilder.clone(*op); + + // else branch: create alternative op + auto elseBuilder = ifOp.getElseBodyBuilder(); + Operation *altOp = createAltFn(elseBuilder, loc); + + // Copy ssbuffer tags + int bid = getBlockId(op); + int tid = getTransferId(op); + if (bid >= 0) { + cloned->setAttr("ssbuffer.block_id", builder.getI32IntegerAttr(bid)); + altOp->setAttr("ssbuffer.block_id", builder.getI32IntegerAttr(bid)); + } + if (tid >= 0) { + cloned->setAttr("ssbuffer.transfer_id", builder.getI32IntegerAttr(tid)); + altOp->setAttr("ssbuffer.transfer_id", builder.getI32IntegerAttr(tid)); + } + + op->replaceAllUsesWith(ifOp.getOperation()); + op->erase(); + return ifOp.getOperation(); +} + +/// Wrap a transfer op (with external uses) in scf.if with yield +static Operation *wrapTransferOpWithScfIfYield(Operation *transferOp, + Value cond, Value inputBuffer, + Value outputBuffer, int bid, + int tid, bool isProducer, + OpBuilder &builder) { + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPoint(transferOp); + + Location loc = transferOp->getLoc(); + auto ifOp = builder.create(loc, transferOp->getResultTypes(), cond, + true /* withElseRegion */); + + // then branch: use inputBuffer + { + auto thenBuilder = ifOp.getThenBodyBuilder(); + IRMapping inputMap; + if (transferOp->getNumOperands() > 0) { + inputMap.map(transferOp->getOperand(transferOp->getNumOperands() - 1), + inputBuffer); + } + Operation *cloned = thenBuilder.clone(*transferOp, inputMap); + thenBuilder.create(loc, cloned->getResults()); + } + + // else branch: use outputBuffer + { + auto elseBuilder = ifOp.getElseBodyBuilder(); + IRMapping outputMap; + if (transferOp->getNumOperands() > 0) { + outputMap.map(transferOp->getOperand(transferOp->getNumOperands() - 1), + outputBuffer); + } + Operation *cloned = elseBuilder.clone(*transferOp, outputMap); + elseBuilder.create(loc, cloned->getResults()); + } + + // Tag the ifOp + ifOp->setAttr("ssbuffer.block_id", builder.getI32IntegerAttr(bid)); + ifOp->setAttr("ssbuffer.transfer_id", builder.getI32IntegerAttr(tid)); + ifOp->setAttr("ssbuffer.cross_buffer", builder.getI32IntegerAttr(1)); + if (!isProducer) { + ifOp->setAttr("ssbuffer.crossDeps", + builder.getArrayAttr({builder.getI32IntegerAttr(tid), + builder.getI32IntegerAttr(0)})); + } + + // Replace all uses of the original transferOp + for (auto [oldResult, newResult] : + llvm::zip_equal(transferOp->getResults(), ifOp->getResults())) { + oldResult.replaceAllUsesWith(newResult); + } + transferOp->erase(); + return ifOp.getOperation(); +} + +/// Wrap a transfer op (no external uses) in scf.if without yield +static Operation *wrapTransferOpWithScfIfSimple(Operation *transferOp, + Value cond, Value inputBuffer, + Value outputBuffer, int bid, + int tid, bool isProducer, + OpBuilder &builder) { + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPoint(transferOp); + + Location loc = transferOp->getLoc(); + auto ifOp = builder.create(loc, TypeRange{}, cond, + true /* withElseRegion */); + + // then branch: clone directly + { + auto thenBuilder = ifOp.getThenBodyBuilder(); + thenBuilder.clone(*transferOp); + } + + // else branch: use outputBuffer + { + auto elseBuilder = ifOp.getElseBodyBuilder(); + IRMapping outputMap; + if (transferOp->getNumOperands() > 0) { + outputMap.map(transferOp->getOperand(transferOp->getNumOperands() - 1), + outputBuffer); + } + elseBuilder.clone(*transferOp, outputMap); + } + + // Tag the ifOp + ifOp->setAttr("ssbuffer.block_id", builder.getI32IntegerAttr(bid)); + ifOp->setAttr("ssbuffer.transfer_id", builder.getI32IntegerAttr(tid)); + ifOp->setAttr("ssbuffer.cross_buffer", builder.getI32IntegerAttr(1)); + if (!isProducer) { + ifOp->setAttr("ssbuffer.crossDeps", + builder.getArrayAttr({builder.getI32IntegerAttr(tid), + builder.getI32IntegerAttr(0)})); + } + + transferOp->erase(); + return ifOp.getOperation(); +} + +/// Process polling for a sender or receiver transfer chain +static int processTransferChain(TransferOpChain &chain, Value cond, + Value inputBuffer, Value outputBuffer, + int outputFlag, bool isProducer, + OpBuilder &builder) { + if (!chain.waitOp) { + return -1; + } + + Location loc = chain.waitOp->getLoc(); + + // 1. Wrap waitOp in polling if + chain.waitOp = wrapSyncOpWithScfIf( + chain.waitOp, cond, outputFlag, builder, + [&](OpBuilder &b, Location l) -> Operation * { + auto waitOp = cast(chain.waitOp); + return b + .create(l, waitOp.getTcoreType(), + waitOp.getTpipe(), waitOp.getPipe(), + b.getI64IntegerAttr(outputFlag)) + .getOperation(); + }); + + // 2. Wrap transferOp in polling if (then=use inputBuffer, else=use + // outputBuffer) + if (chain.transferOp) { + int bid = getBlockId(chain.transferOp); + int tid = getTransferId(chain.transferOp); + bool hasExternalUses = !chain.transferOp->getResults().empty() && + !chain.transferOp->getResult(0).getUses().empty(); + + LDBG("transferOp: " << chain.transferOp->getName() + << ", hasExternalUses=" << hasExternalUses); + + chain.transferOp = + hasExternalUses + ? wrapTransferOpWithScfIfYield(chain.transferOp, cond, inputBuffer, + outputBuffer, bid, tid, isProducer, + builder) + : wrapTransferOpWithScfIfSimple(chain.transferOp, cond, inputBuffer, + outputBuffer, bid, tid, isProducer, + builder); + } + + // 3. Wrap setOp in polling if + if (chain.setOp) { + chain.setOp = wrapSyncOpWithScfIf( + chain.setOp, cond, outputFlag, builder, + [&](OpBuilder &b, Location l) -> Operation * { + auto setOp = cast(chain.setOp); + return b + .create(l, setOp.getTcoreType(), + setOp.getTpipe(), setOp.getPipe(), + b.getI64IntegerAttr(outputFlag)) + .getOperation(); + }); + } + return 0; +} + +/// Add polling control flow for all transfer groups +static int addPollingControlFlow(DenseMap &groups) { + for (auto &p : groups) { + TransferGroupInfo &g = p.second; + + // Get sender's scf.for + Operation *senderWaitParent = g.senderChain.waitOp->getParentOp(); + scf::ForOp senderForOp = cast(senderWaitParent); + + int senderBid = getBlockId(g.senderChain.waitOp); + int senderTid = getTransferId(g.senderChain.waitOp); + + // Insert polling condition at sender waitOp's position + OpBuilder senderCondBuilderForInsert(senderForOp.getBody(), + Block::iterator(g.senderChain.waitOp)); + Value senderCond = createPollingCondition( + senderForOp, senderCondBuilderForInsert, senderBid, senderTid); + OpBuilder senderBuilder(senderForOp.getBody()->getTerminator()); + + // Process sender chain (isProducer=true) + if (processTransferChain(g.senderChain, senderCond, g.senderInputBuffer, + g.senderOutputBuffer, g.outputFlag, true, + senderBuilder) != 0) { + return -1; + } + + // Process receiver chain (may use different scf.for) (isProducer=false) + if (g.receiverChain.waitOp) { + Operation *receiverWaitParent = g.receiverChain.waitOp->getParentOp(); + + if (receiverWaitParent == senderWaitParent) { + // Use the same cond + if (processTransferChain(g.receiverChain, senderCond, + g.receiverInputBuffer, g.receiverOutputBuffer, + g.outputFlag, false, senderBuilder) != 0) { + return -1; + } + } else { + // Receiver uses a different scf.for, create new cond + scf::ForOp receiverForOp = cast(receiverWaitParent); + int receiverBid = getBlockId(g.receiverChain.waitOp); + int receiverTid = getTransferId(g.receiverChain.waitOp); + OpBuilder receiverCondBuilderForInsert( + receiverForOp.getBody(), Block::iterator(g.receiverChain.waitOp)); + Value receiverCond = + createPollingCondition(receiverForOp, receiverCondBuilderForInsert, + receiverBid, receiverTid); + OpBuilder receiverBuilder(receiverForOp.getBody()->getTerminator()); + if (processTransferChain(g.receiverChain, receiverCond, + g.receiverInputBuffer, g.receiverOutputBuffer, + g.outputFlag, false, receiverBuilder) != 0) { + return -1; + } + } + } + } + return 0; +} + +// ============================================================================ +// Pass entry point +// ============================================================================ + +void AddMultiBufferOuterScopePass::runOnOperation() { + ModuleOp module = getOperation(); + LDBG("============================================================"); + LDBG("[AddMultiBufferOuterScope] ENTER"); + LDBG("============================================================"); + + // Step 1: Collect transfer group information + LDBG("[Step 1/3] Start: transfer group collection"); + FlagIdManager flagIdMgr(module); + DenseMap> opsByTid; + collectOpsByTransferId(module, opsByTid); + DenseMap groups; + if (collectTransferGroupData(module, opsByTid, flagIdMgr, groups)) { + LDBG("[Step 1/3] FAILED: no valid transfer groups found"); + signalPassFailure(); + return; + } + LDBG("[Step 1/3] Done: " << groups.size() << " transfer groups"); + + int interCoreBufNum = BufferCountManager::getInstance().getBufferCountByType( + BufferCountManager::DepType::InterCore); + bool isDoubleBuf = (interCoreBufNum > 1); + LDBG("[BufferCount] interCoreBufNum=" << interCoreBufNum + << " doubleBuf=" << isDoubleBuf); + + // Tag consumer-side alloc and transferOp with crossDeps (both modes) + for (auto &p : groups) + addConsumerCrossDepsTags(p.second, module); + + if (isDoubleBuf) { + LDBG("[Step 2/3] Start: output buffer creation"); + if (createOutputBuffers(groups, module)) { + LDBG("[Step 2/3] FAILED: output buffer creation failed"); + signalPassFailure(); + return; + } + LDBG("[Step 2/3] Done"); + + LDBG("[Step 3/3] Start: polling control flow"); + if (addPollingControlFlow(groups)) { + LDBG("[Step 3/3] FAILED: polling control flow failed"); + signalPassFailure(); + return; + } + LDBG("[Step 3/3] Done"); + } else { + LDBG("[Step 2-3] Skipped (single-buffer mode)"); + } + + LDBG("============================================================"); + LDBG("[AddMultiBufferOuterScope] EXIT successfully"); + LDBG("============================================================"); +} + +std::unique_ptr> createAddMultiBufferOuterScopePass() { + return std::make_unique(); +} + +void AddMultiBufferOuterScopePass::getDependentDialects( + DialectRegistry ®istry) const { + registry + .insert(); +} + +void registerAddMultiBufferOuterScopePasses() { + registerPass([]() -> std::unique_ptr { + return createAddMultiBufferOuterScopePass(); + }); +} + +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/CMakeLists.txt b/third_party/ascend/lib/DynamicCVPipeline/CMakeLists.txt new file mode 100644 index 0000000000..0f6a8474dd --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/CMakeLists.txt @@ -0,0 +1,43 @@ +file(GLOB PLAN_COMPUTE_BLOCK_PASS_SRC + "PlanComputeBlock/*.cpp" +) +file(GLOB COMMON_SRC + "Common/*.cpp" +) +file(GLOB SPLIT_DATAFLOW_SRC_FILES CONFIGURE_DEPENDS + "SplitDataflow/*.cpp" +) +file(GLOB SEPARATE_MEMORY_FROM_COMPUTE_SRC_FILES CONFIGURE_DEPENDS + "SeparateMemoryFromCompute/*.cpp" +) + +file(GLOB ALLOC_MULTI_CACHE_SRCS_FILES CONFIGURE_DEPENDS + "AllocMultiCache/*.cpp" +) + +file(GLOB ADD_CONTROLFLOW_CONDITION_SRC_FILES CONFIGURE_DEPENDS + "AddControlFlowCondition/*.cpp" +) + +add_triton_library(DynamicCVPipeline + AddDynamicCVPipeline.cpp + PlanComputeBlock.cpp + ${PLAN_COMPUTE_BLOCK_PASS_SRC} + ${COMMON_SRC} + SplitDataflow.cpp + SeparateMemoryFromComputePass.cpp + AllocMultiCache.cpp + AddControlFlowCondition.cpp + ${PLAN_COMPUTE_BLOCK_SRC_FILES} + ${SPLIT_DATAFLOW_SRC_FILES} + ${SEPARATE_MEMORY_FROM_COMPUTE_SRC_FILES} + ${ALLOC_MULTI_CACHE_SRCS_FILES} + ${ADD_CONTROLFLOW_CONDITION_SRC_FILES} + + DEPENDS + AddDynamicCVPipelinePassIncGen + + LINK_LIBS + BiShengIRHIVMDialect + BiShengIRScopeDialect +) diff --git a/third_party/ascend/lib/DynamicCVPipeline/Common/BufferCountManager.cpp b/third_party/ascend/lib/DynamicCVPipeline/Common/BufferCountManager.cpp new file mode 100644 index 0000000000..7884fad59c --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/Common/BufferCountManager.cpp @@ -0,0 +1,120 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/Common/BufferCountManager.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/raw_ostream.h" + +static constexpr const char *DEBUG_TYPE = "BufferCountManager"; +#define LOG_DEBUG(...) \ + LLVM_DEBUG(llvm::dbgs() << " [" << DEBUG_TYPE << "] " << __VA_ARGS__) + +constexpr int kDefaultIntraBufferCount = 2; +constexpr int kDefaultInterBufferCount = 1; +constexpr int kDefaultLoadBufferCount = 2; +constexpr int kBufferCountWarningThreshold = 3; + +namespace mlir { +namespace triton { + +BufferCountManager &BufferCountManager::getInstance() { + static BufferCountManager instance; + return instance; +} + +BufferCountManager::BufferCountManager() + : intraBufferCount_(kDefaultIntraBufferCount), + interCoreBufferCount_(kDefaultInterBufferCount), + loadStoreBufferCount_(kDefaultLoadBufferCount) { + LOG_DEBUG("Default initialized: " + << "IntraBufferCount=" << intraBufferCount_ + << ", InterBufferCount=" << interCoreBufferCount_ + << ", LoadBufferCount=" << loadStoreBufferCount_); +} + +void BufferCountManager::setBufferCount(DepType type, int count) { + if (count <= 0) { + LOG_DEBUG("Invalid buffer count: " << count << " (must be > 0)"); + return; + } + if (count >= kBufferCountWarningThreshold) { + LOG_DEBUG("Warning: buffer count " << count + << " >= " << kBufferCountWarningThreshold + << " is not recommended"); + } + switch (type) { + case DepType::IntraCore: + intraBufferCount_ = count; + LOG_DEBUG("IntraBufferCount set to " << count); + break; + case DepType::InterCore: + interCoreBufferCount_ = count; + LOG_DEBUG("InterBufferCount set to " << count); + break; + case DepType::LoadStore: + loadStoreBufferCount_ = count; + LOG_DEBUG("LoadBufferCount set to " << count); + break; + default: + LOG_DEBUG("Unknown DepType: " << static_cast(type)); + break; + } +} + +void BufferCountManager::buildBufferCountMap( + llvm::DenseMap> &depValueMap, + llvm::DenseMap &bufferCountMap, DepType type) { + int bufCount = getBufferCountByType(type); + + for (auto &p : depValueMap) { + for (Value depVal : p.second) { + if (isa(depVal) || !depVal.getDefiningOp()) + continue; + bufferCountMap[depVal] = bufCount; + } + } +} + +int BufferCountManager::getBufferCountByType(DepType type) const { + int count = 1; + switch (type) { + case DepType::IntraCore: + count = intraBufferCount_; + break; + case DepType::InterCore: + count = interCoreBufferCount_; + break; + case DepType::LoadStore: + count = loadStoreBufferCount_; + break; + default: + LOG_DEBUG("Unknown DepType: " << static_cast(type)); + break; + } + LOG_DEBUG("getBufferCountByType(" + << static_cast(type) << ") = " << count + << " (IntraCore=0, InterCore=1, LoadStore=2)"); + return count; +} + +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/Common/FlagIdManager.cpp b/third_party/ascend/lib/DynamicCVPipeline/Common/FlagIdManager.cpp new file mode 100644 index 0000000000..d91475dee5 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/Common/FlagIdManager.cpp @@ -0,0 +1,67 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/Common/FlagIdManager.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "mlir/IR/Operation.h" +#include "llvm/Support/Debug.h" + +using namespace mlir; +static constexpr const char *DEBUG_TYPE = "flag-id-manager"; +#define LOG_DEBUG(...) \ + LLVM_DEBUG(llvm::dbgs() << " [" << DEBUG_TYPE << "] " << __VA_ARGS__) + +using namespace mlir::triton; +using namespace hivm; + +static constexpr int kMaxOps = 1024; + +FlagIdManager::FlagIdManager(ModuleOp module) { + this->module = module; + scanExistingFlags(module); + LOG_DEBUG("FlagIdManager: Initialized with max_id = " << currentMaxId + << "\n"); +} + +void FlagIdManager::scanExistingFlags(ModuleOp module) { + module.walk([&](Operation *op) { + if (isa(op) || isa(op)) { + int flag = -1; + if (auto intAttr = op->getAttrOfType("static_flag_id")) { + flag = (int)intAttr.getInt(); + } else if (auto intAttr = op->getAttrOfType("flag")) { + flag = (int)intAttr.getInt(); + } + if (flag >= 0) { + currentMaxId = std::max(currentMaxId, (int64_t)flag); + } + } + }); +} + +int FlagIdManager::acquireId(Operation *insertionPoint) { + if (currentMaxId < MAX_FLAG_ID) { + return ++currentMaxId; + } + + return INVALID_FLAG_ID; +} diff --git a/third_party/ascend/lib/DynamicCVPipeline/Common/MemoryEffectsTracker.cpp b/third_party/ascend/lib/DynamicCVPipeline/Common/MemoryEffectsTracker.cpp new file mode 100644 index 0000000000..91227b171a --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/Common/MemoryEffectsTracker.cpp @@ -0,0 +1,461 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +// Memory side-effect analysis: tracks per-location read/write effects across +// ops to identify implicit ordering constraints (RAW, WAR, WAW) and expose +// them as scheduling edges (memDefs/memUsers, execBefore/execAfter). +// +// Each memory location is a MemSlot recording its last writer and pending +// readers. Two-level alias resolution determines which slots an effect touches: +// 1. Exact: value V found directly in valueToSlot. +// 2. AA query: AliasAnalysis asked whether a slot’s memref may-alias V. +// Root BlockArguments reject MayAlias (rejectMayAlias=true) to avoid +// false positives from conservative AA on function parameters. +// +// Unknown ops (no SideEffect interface) act as full barriers: they depend on +// all prior writers/readers and become the sole writer for every slot. + +#include "ascend/include/DynamicCVPipeline/Common/MemoryEffectsTracker.h" +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/Region.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "mlir/Interfaces/ViewLikeInterface.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SetVector.h" +#include "llvm/Support/Debug.h" + +using namespace mlir; +static constexpr const char *DEBUG_TYPE = "memory-effects-tracker"; +#define LOG_DEBUG(...) \ + LLVM_DEBUG(llvm::dbgs() << " [" << DEBUG_TYPE << "] " << __VA_ARGS__) + +using namespace mlir::CVPipeline; + +namespace { + +bool isDefinedInside(Value v, Operation *op) { + if (!v || !op) { + return false; + } + + if (auto arg = dyn_cast(v)) { + Operation *parent = arg.getOwner()->getParentOp(); + return parent && op->isAncestor(parent); + } + + Operation *defOp = v.getDefiningOp(); + if (!defOp) { + return false; + } + + return op->isProperAncestor(defOp); +} + +} // namespace + +MemoryDependenceGraph::MemoryDependenceGraph(Operation *root, AliasAnalysis &aa) + : root(root), aa(aa) { + if (!root) { + return; + } + analyzeOp(root); + slots.clear(); + valueToSlot.clear(); +} + +ArrayRef MemoryDependenceGraph::getMemDefs(Operation *op) const { + auto it = memDefs.find(op); + if (it == memDefs.end()) { + return {}; + } + return it->second; +} + +ArrayRef MemoryDependenceGraph::getMemUsers(Operation *op) const { + auto it = memUsers.find(op); + if (it == memUsers.end()) + return {}; + return it->second; +} + +ArrayRef +MemoryDependenceGraph::getExecBefore(Operation *op) const { + auto it = execBefore.find(op); + if (it == execBefore.end()) + return {}; + return it->second; +} + +ArrayRef MemoryDependenceGraph::getExecAfter(Operation *op) const { + auto it = execAfter.find(op); + if (it == execAfter.end()) + return {}; + return it->second; +} + +void MemoryDependenceGraph::analyzeOp(Operation *op) { + if (!op) { + return; + } + + // The graph root is treated as a pure container: we don't query its own + // effects against itself, only recurse into its regions. + if (op == root) { + analyzeRegionsOf(op); + return; + } + + LOG_DEBUG("Analyzing op " << *op << "\n"); + // Step 1: collect outside memory effects + bool unknown = false; + SmallVector effects = + collectOuterEffects(op, unknown); + + // Step 2: collect defs(WAR) and preds(WAR, WAW, RAW) dependency. + SmallVector defs; + SmallVector preds; + collectPreds(effects, unknown, defs, preds); + LOG_DEBUG("Defs: \n"); + for (auto def : defs) { + LOG_DEBUG(*def << "\n"); + } + LOG_DEBUG("Preds: \n"); + for (auto pred : preds) { + LOG_DEBUG(*pred << "\n"); + } + + // Step 3: extract edges from defs and preds to the graph. + recordEdges(op, defs, preds); + + // Step 4: recursively search regions. + if (op->getNumRegions() > 0) { + analyzeRegionsOf(op); + } + + // Step 5: store effects into slots. + applyEffects(op, effects, unknown); +} + +void MemoryDependenceGraph::analyzeRegionsOf(Operation *op) { + // Whether to isolated slots. + // For example: scf.for should keeps slots; funOp should use new slots. + const bool isolated = + op != root && op->hasTrait(); + + for (Region ®ion : op->getRegions()) { + // Snapshot is used for quickly restoration. + Snapshot snap = takeSnapshot(); + if (isolated) { + slots.clear(); + valueToSlot.clear(); + } + + for (Block &block : region) { + for (BlockArgument arg : block.getArguments()) { + if (isa(arg.getType())) { + auto *slot = getOrCreateSlot(arg); + } + } + for (Operation &inner : block) { + analyzeOp(&inner); + } + } + + restoreSnapshot(std::move(snap)); + } +} + +SmallVector +MemoryDependenceGraph::collectOuterEffects(Operation *op, bool &unknown) { + unknown = false; + + std::optional> raw = + getEffectsRecursively(op); + if (!raw) { + unknown = true; + return {}; + } + + SmallVector filtered; + filtered.reserve(raw->size()); + for (auto &e : *raw) { + if (isDefinedInside(e.getValue(), op)) { + continue; + } + filtered.push_back(e); + } + return filtered; +} + +AliasResult MemoryDependenceGraph::queryAlias(Value lhs, Value rhs) { + auto isFuncEntryArg = [](const Value &val) -> bool { + auto arg = llvm::dyn_cast(val); + return arg && arg.getOwner()->isEntryBlock(); + }; + auto getSource = [](Value val) -> Value { + while (auto viewLike = val.getDefiningOp()) { + val = viewLike.getViewSource(); + } + return val; + }; + if (isFuncEntryArg(getSource(lhs)) && isFuncEntryArg(getSource(rhs))) { + return lhs == rhs ? AliasResult::MustAlias : AliasResult::NoAlias; + } + return aa.alias(lhs, rhs); +} + +SmallVector +MemoryDependenceGraph::findAliasSlots(Value v) { + SmallVector result; + + if (!v) { // unknown location: conservatively alias with every slot + result.reserve(slots.size()); + for (auto &slot : slots) { + result.push_back(slot.get()); + } + return result; + } + + SmallPtrSet seen; + + auto it = valueToSlot.find(v); + if (it != valueToSlot.end()) { + result.push_back(it->second); + seen.insert(it->second); + } + + for (auto &slot : slots) { + MemSlot *raw = slot.get(); + if (seen.contains(raw) || !raw->memref) { + continue; + } + auto aliasResult = queryAlias(raw->memref, v); + if (aliasResult != AliasResult::NoAlias) { + result.push_back(raw); + seen.insert(raw); + } + } + + return result; +} + +MemoryDependenceGraph::MemSlot * +MemoryDependenceGraph::getOrCreateSlot(Value v) { + if (!v) { + return nullptr; + } + auto it = valueToSlot.find(v); + if (it != valueToSlot.end()) { + return it->second; + } + auto slot = std::make_unique(v); + MemSlot *raw = slot.get(); + slots.push_back(std::move(slot)); + valueToSlot[v] = raw; + return raw; +} + +ArrayRef +MemoryDependenceGraph::resolveAliasSlots( + Value v, DenseMap> &cache) { + auto it = cache.find(v); + if (it != cache.end()) { + return it->second; + } + auto [ins, _] = cache.try_emplace(v, findAliasSlots(v)); + return ins->second; +} + +void MemoryDependenceGraph::collectPreds( + ArrayRef effects, bool unknown, + SmallVectorImpl &defsOut, + SmallVectorImpl &predsOut) { + llvm::SmallSetVector defs; + llvm::SmallSetVector preds; + + // defs collects lastWriter only on reads (RAW); preds collects all ordering + // deps (RAW/WAR/WAW). + auto addFromSlot = [&](MemSlot *slot, bool isWriteLike) { + if (slot->dataSource && !isWriteLike) { + defs.insert(slot->dataSource); + } + if (slot->lastWriter) { + preds.insert( + slot->lastWriter); // XAW: we write/read what lastWriter wrote + } + if (isWriteLike) { + for (Operation *r : slot->pendingReads) { + preds.insert(r); // WAR: we will overwrite data still needed by readers + } + } + }; + + // The unknown op is considered as a barrier. + if (unknown) { + for (auto &slot : slots) { + addFromSlot(slot.get(), true); + } + predsOut.assign(preds.begin(), preds.end()); + return; + } + + DenseMap> cache; + for (const auto &e : effects) { + Value v = e.getValue(); + if (isa(e.getEffect())) { + for (MemSlot *s : resolveAliasSlots(v, cache)) { + addFromSlot(s, false); + } + } else if (isa(e.getEffect()) || + isa(e.getEffect())) { + for (MemSlot *s : resolveAliasSlots(v, cache)) { + addFromSlot(s, true); + } + } + } + + defsOut.assign(defs.begin(), defs.end()); + predsOut.assign(preds.begin(), preds.end()); +} + +void MemoryDependenceGraph::applyEffects( + Operation *op, ArrayRef effects, + bool unknown) { + // Conservative strategy. Unknown op is to treat as barrier. + if (unknown) { + for (auto &slot : slots) { + slot->lastWriter = op; + slot->pendingReads.clear(); + } + return; + } + + for (auto result : op->getOpResults()) { + if (isa(result.getType())) { + if (MemSlot *s = getOrCreateSlot(result)) { + s->dataSource = op; + s->lastWriter = op; + s->pendingReads.clear(); + } + } + } + + DenseMap> cache; + + for (const auto &e : effects) { + // Reads first: prevents a self-aliasing op from clearing its own pending + // read. + if (isa(e.getEffect())) { + for (MemSlot *s : resolveAliasSlots(e.getValue(), cache)) { + s->pendingReads.insert(op); + } + } + } + + for (const auto &e : effects) { + Value v = e.getValue(); + if (isa(e.getEffect())) { + if (MemSlot *s = getOrCreateSlot(v)) { + s->dataSource = op; + s->lastWriter = op; + s->pendingReads.clear(); + } + } else if (isa(e.getEffect())) { + for (MemSlot *s : resolveAliasSlots(v, cache)) { + // annotation::MarkOp, no data produced, is marked with Mem::Write + if (!isa(op)) { + s->dataSource = op; + } + s->lastWriter = op; + s->pendingReads.clear(); + } + } else if (isa(e.getEffect())) { + SmallPtrSet toRemove; + if (!v) { + LOG_DEBUG("Free of unknown value: conservatively drop all slots."); + for (auto &slot : slots) { + toRemove.insert(slot.get()); + } + } else { + auto aliasSlots = findAliasSlots(v); + for (auto *slot : aliasSlots) { + toRemove.insert(slot); + } + } + llvm::erase_if(slots, [&](std::unique_ptr &sp) { + if (toRemove.contains(sp.get())) { + if (sp->memref) { + valueToSlot.erase(sp->memref); + } + return true; + } + return false; + }); + cache.clear(); + } + } +} + +MemoryDependenceGraph::Snapshot MemoryDependenceGraph::takeSnapshot() const { + Snapshot snap; + snap.states.reserve(slots.size()); + for (const auto &slot : slots) { + snap.states.push_back(*slot); + } + return snap; +} + +void MemoryDependenceGraph::restoreSnapshot(Snapshot &&snap) { + slots.clear(); + valueToSlot.clear(); + slots.reserve(snap.states.size()); + for (auto &s : snap.states) { + auto slot = std::make_unique(std::move(s)); + if (slot->memref) { + valueToSlot[slot->memref] = slot.get(); + } + slots.push_back(std::move(slot)); + } +} + +void MemoryDependenceGraph::recordEdges(Operation *op, + ArrayRef defs, + ArrayRef preds) { + if (!defs.empty()) { + auto &defList = memDefs[op]; + defList.assign(defs.begin(), defs.end()); + for (Operation *p : defs) { + memUsers[p].push_back(op); + } + } + + if (!preds.empty()) { + auto &execBeforeList = execBefore[op]; + execBeforeList.assign(preds.begin(), preds.end()); + for (Operation *p : preds) { + execAfter[p].push_back(op); + } + } +} diff --git a/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp b/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp new file mode 100644 index 0000000000..5d14d108b7 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp @@ -0,0 +1,67 @@ +#include +#include +#include + +#include "llvm/IR/PatternMatch.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/FormatVariadic.h" +#include "llvm/Support/LogicalResult.h" + +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypes.h" + +#include "ascend/include/DynamicCVPipeline/Common/Utils.h" +namespace mlir { +namespace CVPipeline { + +CoreType getOpCoreType(Operation *op) { + if (!op) { + return CoreType::UNDETERMINED; + } + if (auto a = op->getAttrOfType(kCoreType)) { + return fromStrCoreType(a.getValue()); + } + return CoreType::UNDETERMINED; +} + +llvm::LogicalResult verifyOpBlockId(Operation *op) { + if (!op) { + assert(false && "Op is nullptr, please check calling function"); + + // return failure to signal disabling of CV dynamic pipeline in release mode + return llvm::failure(); + } + + auto blockId = op->getAttrOfType(kBlockId); + if (blockId && blockId.getInt() < 0) { + std::string_view errorPass = "previous passes"; + auto diag = op->emitError() + << "block id should not be negative! Please report to "; + switch (getOpCoreType(op)) { + case CoreType::CUBE_ONLY: + diag << "PlanCubePass"; + break; + case CoreType::VECTOR_ONLY: + diag << "PlanVectorPass"; + break; + default: + diag << "previous passes"; + } + return llvm::failure(); + } + + return llvm::success(); +} + +std::optional getOpBlockId(Operation *op) { + auto blockIdAttr = op->getAttrOfType(kBlockId); + if (!blockIdAttr) { + return std::nullopt; + } + + return blockIdAttr.getInt(); +} + +} // namespace CVPipeline +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock.cpp new file mode 100644 index 0000000000..a7de496a2b --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock.cpp @@ -0,0 +1,75 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "llvm/Support/Debug.h" + +#include "mlir/Pass/PassManager.h" + +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/OpClassifier.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/Passes.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/ReorderOpsByBlockId.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlockPass.h" + +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlockPass.h" + +using namespace mlir; +using namespace triton; + +static constexpr const char *DEBUG_TYPE = "plan-compute-block"; +#define LOG_DEBUG(...) \ + LLVM_DEBUG(llvm::dbgs() << " [" << DEBUG_TYPE << "] " << __VA_ARGS__) + +// Run the pass +void PlanComputeBlockPass::runOnOperation() { + ModuleOp module = getOperation(); + OpPassManager pm(module.getOperationName()); + CVPipeline::ComputeBlockIdManager::getInstance().reset(); + LOG_DEBUG("Enter pass.\n"); + + // Step 1: Run OpClassifierPass to classify operations + pm.addPass(createOpClassifierPass()); + + // Step 2: Partition compute blocks for core_type=cube + pm.addPass(createPlanCubeBlockPass()); + + // Step 3: Partition compute blocks for core_type=vector + pm.addPass(createPlanVectorBlockPass()); + + // Step 4: Reorder + pm.addPass(createReorderOpsByBlockIdPass()); + + if (failed(runPipeline(pm, module))) { + module->emitError() << "[" << DEBUG_TYPE << "] Pass failed!"; + signalPassFailure(); + } + + LOG_DEBUG("Process successfully\n"); +} + +namespace mlir { +namespace triton { +std::unique_ptr> createPlanComputeBlockPass() { + return std::make_unique(); +} +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/Common.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/Common.cpp new file mode 100644 index 0000000000..82739bcbd1 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/Common.cpp @@ -0,0 +1,70 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h" + +namespace mlir { +namespace CVPipeline { + +void initializeIndegreeForBlock(Block *block, + llvm::DenseMap &indegree, + const MemoryDependenceGraph &memGraph) { + auto &bm = ComputeBlockIdManager::getInstance(); + + block->walk([&](Operation *op) { + if (op->getBlock() != block) { + return; + } + indegree[op] = 0; + // We need to consider op itself && op's region-contained ops. + op->walk([&](Operation *nestedOp) { + for (auto inValue : nestedOp->getOperands()) { + if (auto defOp = inValue.getDefiningOp()) { + if (defOp->getBlock() == block && !bm.isSameBlock(defOp, op)) { + indegree[op]++; + } + } + } + + for (auto memDepUser : memGraph.getExecBefore(nestedOp)) { + if (memDepUser->getBlock() == block && + !bm.isSameBlock(memDepUser, op)) { + indegree[op]++; + } + } + }); + }); +} + +Operation *getAncestorInBlock(Operation *inner, Block *block) { + Operation *cur = inner; + while (cur) { + if (cur->getBlock() == block) { + return cur; + } + cur = cur->getParentOp(); + } + return nullptr; +} + +} // namespace CVPipeline +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp new file mode 100644 index 0000000000..e4d469d0bf --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp @@ -0,0 +1,148 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/LogicalResult.h" + +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypes.h" + +#include "ascend/include/DynamicCVPipeline/Common/Utils.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h" + +namespace mlir { +namespace CVPipeline { + +bool ComputeBlockIdManager::isWholeCubeReady( + Operation *seedOp, llvm::DenseMap &indegree) { + auto id = getBlockIdByOp(seedOp); + if (id == -1) { + return (indegree[seedOp] == 0); + } + auto cubeBlock = getOpsByBlockId(id); + for (auto op : cubeBlock) { + if (!indegree.contains(op)) { + continue; + } + if (indegree[op] != 0) { + return false; + } + } + return true; +} + +bool ComputeBlockIdManager::isSameBlock(Operation *a, Operation *b) { + if (getBlockIdByOp(a) == getBlockIdByOp(b) && getBlockIdByOp(a) != -1) { + return true; + } + return false; +} + +unsigned int ComputeBlockIdManager::getNextId() { + std::lock_guard lock(managerMutex); + return cntComputeBlockId++; +} + +llvm::LogicalResult ComputeBlockIdManager::markOpBlockId(Operation *op, + int blockId) { + if (blockId < 0) { + op->emitError("marking block id as negative"); + return llvm::failure(); + } + if (!op) { + return llvm::success(); + } + MLIRContext *ctx = op->getContext(); + op->setAttr(kBlockId, + IntegerAttr::get(IntegerType::get(ctx, blockIdWidth), blockId)); + return record(op, blockId); +} + +llvm::SmallVector +ComputeBlockIdManager::getOpsByBlockId(int blockId) { + if (blockId == -1) { + return {}; + } + std::lock_guard lock(managerMutex); + auto it = blockIdToOps.find(blockId); + if (it == blockIdToOps.end()) { + return {}; + } + return llvm::SmallVector(it->second.begin(), it->second.end()); +} + +int ComputeBlockIdManager::getBlockIdByOp(Operation *op) { + std::lock_guard lock(managerMutex); + auto it = opToBlockId.find(op); + if (it != opToBlockId.end()) { + return it->second; + } + return -1; +} + +llvm::LogicalResult ComputeBlockIdManager::markOpsWithNewId( + llvm::SmallVectorImpl &ops) { + if (ops.empty()) { + return llvm::success(); + } + int id = static_cast(getNextId()); + for (Operation *op : ops) { + if (llvm::failed(markOpBlockId(op, id))) { + return llvm::failure(); + } + } + return llvm::success(); +} + +void ComputeBlockIdManager::reset() { + std::lock_guard lock(managerMutex); + cntComputeBlockId = 0; + blockIdToOps.clear(); + opToBlockId.clear(); +} + +llvm::LogicalResult ComputeBlockIdManager::record(Operation *op, int blockId) { + if (!op || blockId < 0) { + return llvm::success(); + } + std::lock_guard lock(managerMutex); + auto itOld = opToBlockId.find(op); + if (itOld != opToBlockId.end()) { + llvm::errs() << "Error: Operation already has a block id. Op: " << *op + << ", old block id: " << itOld->second + << ", new block id: " << blockId << "\n"; + return llvm::failure(); + } + + opToBlockId[op] = blockId; + auto &vec = blockIdToOps[blockId]; + for (Operation *existing : vec) { + if (existing == op) { + return llvm::success(); + } + } + vec.push_back(op); + return llvm::success(); +} + +} // namespace CVPipeline +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp new file mode 100644 index 0000000000..4a0103c6a4 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp @@ -0,0 +1,1387 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/OpClassifier.h" + +#include "llvm/ADT/DenseSet.h" + +#include + +using namespace mlir; +static constexpr const char *DEBUG_TYPE = "op-classifier"; +#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ") +#define LDBG(X) LLVM_DEBUG(DBGS() << (X) << "\n") +using namespace mlir::triton; + +namespace { + +bool isInsideNestedLinalgRegion(Operation *op) { + for (Operation *parent = op ? op->getParentOp() : nullptr; parent; + parent = parent->getParentOp()) { + if (isa(parent->getDialect())) { + return true; + } + } + return false; +} + +} // namespace + +// Helper: describe operation for logging +std::string OpClassifierPass::describeOp(Operation *op) const { + if (!op) + return ""; + std::string result; + llvm::raw_string_ostream os(result); + os << op->getName().getStringRef(); + return result; +} + +// Helper: convert OpCoreType to string for IR attribute +std::string coreTypeToString(OpCoreType ct) { + switch (ct) { + case OP_CUBE_ONLY: + return "CUBE"; + case OP_VECTOR_ONLY: + return "VECTOR"; + case OP_CUBE_AND_VECTOR: + return "CUBE_AND_VECTOR"; + default: + return "UNDETERMINED"; + } +} + +namespace { + +// Maximum number of core types in a comma-separated core type string (e.g., +// "CUBE,VECTOR") +constexpr size_t kMaxCoreTypeParts = 4; + +// Minimum number of inputs for linalg.matmul (A and B matrices) +constexpr size_t kMinMatmulInputs = 2; + +} // namespace + +// Helper: parse OpCoreType from string (handles both "CUBE" and "CUBE,VECTOR" +// formats) For comma-separated multi-value strings, returns the i-th component +// if index is provided, otherwise returns the first component. +OpCoreType parseCoreTypeFromString(const std::string &coreTypeStr, + size_t index = 0) { + llvm::StringRef ref(coreTypeStr); + llvm::SmallVector parts; + ref.split(parts, ','); + if (parts.size() > 1 && index < parts.size()) { + // Multi-value: pick i-th component + return parseCoreTypeFromString(parts[index].str()); + } + StringRef singleValue = parts[0].trim(); + if (singleValue == "CUBE") { + return OP_CUBE_ONLY; + } else if (singleValue == "VECTOR") { + return OP_VECTOR_ONLY; + } else if (singleValue == "CUBE_AND_VECTOR") { + return OP_CUBE_AND_VECTOR; + } + return OP_UNDETERMINED; +} + +// Initialize the pass +void OpClassifierPass::initializePass(ModuleOp module) { + // Reset data structures + opCoreTypes.clear(); + allOps.clear(); + cubeSeeds.clear(); + + // Collect all operations + module.walk([&](Operation *op) { + if (isa(op)) { + return; // Skip module, function, and return ops + } + allOps.push_back(op); + }); + + LLVM_DEBUG(DBGS() << "Collected: " << allOps.size() << " ops\n"); + + // Initialize all ops to UNDETERMINED + for (Operation *op : allOps) { + opCoreTypes[op] = OP_UNDETERMINED; + } +} + +// Mark an operation as CUBE +void OpClassifierPass::markCube(Operation *op) { + if (op && opCoreTypes.count(op)) { + opCoreTypes[op] = static_cast(opCoreTypes[op] | OP_CUBE_ONLY); + LLVM_DEBUG(DBGS() << "CUBE: " << *op << "\n"); + } +} + +// ============================================================================ +// Pattern: to_tensor → matmul (Upstream) +// ============================================================================ +// Matches cases where matmul's input comes from bufferization.to_tensor. +// to_tensor converts a memref to a tensor, typically used for loading matrix +// data from memory. IR Example +// %memref = memref.alloc() : memref<1024x1024xf32> +// %tensor = bufferization.to_tensor %memref : memref<1024x1024xf32> -> +// tensor<1024x1024xf32> %result = linalg.matmul ins(%tensor, %tensor_b) +// outs(%init) +// Matching Logic +// 1. Check if matmul's operand defining op is bufferization.to_tensor +// 2. If matched, mark to_tensor as CUBE and add to cubeSeeds +// 3. Also trace back to_tensor's memref source and mark its defining op as +// CUBE +// Purpose: Ensure matmul's input data loading executes on CUBE for efficient +// matrix data transfer. +// ============================================================================ +void OpClassifierPass::matchToTensorPattern(Operation *def) { + auto toTensorOp = dyn_cast(def); + if (!toTensorOp) + return; + + markCube(toTensorOp); + cubeSeeds.push_back(toTensorOp); + + // Also mark the memref allocation as CUBE + Value memref = toTensorOp.getBuffer(); + if (Operation *memrefDef = memref.getDefiningOp()) { + markCube(memrefDef); + cubeSeeds.push_back(memrefDef); + } +} + +// ============================================================================ +// Pattern: transpose → matmul (Upstream) +// ============================================================================ +// Matches cases where matmul's input comes from linalg.transpose. +// transpose performs matrix transposition, commonly used as preprocessing for +// matrix multiplication inputs. +// IR Example +// %input = tensor.empty() : tensor<1024x1024xf32> +// %transposed = linalg.transpose ins(%input: tensor<1024x1024xf32>) +// outs(%out: tensor<1024x1024xf32>) +// permutation = [1, 0] +// %result = linalg.matmul ins(%transposed, %tensor_b) outs(%init) +// Matching Logic +// 1. Check if matmul's operand defining op is linalg.transpose +// 2. Mark transpose itself as CUBE (transpose operation is suitable for CUBE +// execution) +// 3. Only mark transpose's input as CUBE seed when its operands come from +// - bufferization dialect ops (e.g., to_tensor, clone), excluding +// AllocTensorOp +// - tensor.empty op +// 4. Check both transpose's operands and DpsInits (output initial values), +// either one +// satisfying the condition is sufficient +// Special Handling +// When transpose's input comes from compute ops (e.g., arith.truncf), +// transpose itself remains CUBE, but its upstream chain stays VECTOR to avoid +// unnecessary CUBE expansion. +// Purpose: transpose is a common preprocessing op for matmul; executing on CUBE +// leverages matrix transfer capabilities. +// ============================================================================ +void OpClassifierPass::matchTransposePattern(Operation *def) { + auto transposeOp = dyn_cast(def); + if (!transposeOp) + return; + + markCube(transposeOp); + + // Helper lambda to check if an operand's defining op qualifies for CUBE seed + auto shouldMarkCubeSeed = [](Operation *opDef) -> bool { + if (!opDef) + return false; + return (isa(opDef->getDialect()) && + !isa(opDef)) || + isa(opDef); + }; + + // Check input tensor + auto operands = transposeOp->getOperands(); + for (const auto &op : operands) { + if (shouldMarkCubeSeed(op.getDefiningOp())) { + markCube(op.getDefiningOp()); + cubeSeeds.push_back(op.getDefiningOp()); + break; // No need to check other operands, one is enough to seed the + // transpose as CUBE + } + } + + // Check outs (DpsInits) + auto outs = transposeOp.getDpsInits(); + for (const auto &out : outs) { + if (shouldMarkCubeSeed(out.getDefiningOp())) { + markCube(out.getDefiningOp()); + cubeSeeds.push_back(out.getDefiningOp()); + break; + } + } +} + +// ============================================================================ +// Pattern: linalg.fill → matmul (Upstream) +// ============================================================================ +// Matches cases where matmul's output initial value comes from linalg.fill. +// fill initializes the output matrix (typically to 0). +// IR Example +// %value = arith.constant 0.0 : f32 +// %out = tensor.empty() : tensor<1024x1024xf32> +// %init = linalg.fill ins(%value: f32) outs(%out: tensor<1024x1024xf32>) +// %result = linalg.matmul ins(%a, %b) outs(%init) +// Matching Logic +// 1. Check if matmul's operand defining op is linalg.fill +// 2. If matched, mark fill as CUBE and add to cubeSeeds +// Purpose: fill operation initializes matmul's output buffer; executing on CUBE +// leverages efficient data filling capabilities. +// ============================================================================ +void OpClassifierPass::matchFillPattern(Operation *def) { + auto fillOp = dyn_cast(def); + if (!fillOp) + return; + + markCube(fillOp); + cubeSeeds.push_back(fillOp); +} + +// ============================================================================ +// Pattern: matmul → hivm.hir.store (Downstream) +// ============================================================================ +// Matches cases where matmul's output is directly stored to memory. +// hivm.store is the HIVM dialect's store operation, writing tensor to memory. +// IR Example +// %result = linalg.matmul ins(%a, %b) outs(%init) +// hivm.store %result, %memref[%offset] : tensor<1024x1024xf32> to memref +// Matching Logic +// 1. Check if matmul result's user is hivm.store +// 2. If matched, mark store as CUBE and add to cubeSeeds +// Purpose: Ensure matmul result's store operation executes on CUBE for +// efficient matrix data write-back. +// ============================================================================ +void OpClassifierPass::matchStorePattern(Operation *user) { + if (!isa(user)) + return; + + markCube(user); + cubeSeeds.push_back(user); +} + +// ============================================================================ +// Pattern: matmul → tensor.extract_slice → hivm.store (Downstream) +// ============================================================================ +// Matches cases where matmul's output is sliced before being stored. +// extract_slice extracts a sub-region (slice) of a tensor, commonly used after +// tiled matrix computation for result extraction. +// IR Example +// %result = linalg.matmul ins(%a, %b) outs(%init) : tensor<1024x1024xf32> +// %slice = tensor.extract_slice %result[0, 0][256, 256][1, 1] : +// tensor<1024x1024xf32> to tensor<256x256xf32> hivm.store %slice, +// %memref[%offset] : tensor<256x256xf32> to memref +// Matching Logic +// 1. Check if matmul result's user is tensor.extract_slice +// 2. If matched, mark extract_slice as CUBE and add to cubeSeeds +// 3. Further check all users of extract_slice and mark hivm.store as CUBE +// Purpose: After tiled matrix computation, slice and store operations on CUBE +// improve local data transfer efficiency. +// ============================================================================ +void OpClassifierPass::matchExtractSlicePattern(Operation *user) { + auto extractSliceOp = dyn_cast(user); + if (!extractSliceOp) + return; + + markCube(extractSliceOp); + cubeSeeds.push_back(extractSliceOp); + // Also mark downstream hivm.hir.store as CUBE + for (Operation *sliceUser : extractSliceOp->getUsers()) { + if (isa(sliceUser)) { + markCube(sliceUser); + cubeSeeds.push_back(sliceUser); + } + } +} + +// ============================================================================ +// Pattern: matmul → materialize_in_destination (Downstream) +// ============================================================================ +// Matches cases where matmul's output is written to a destination buffer via +// materialize_in_destination. This operation materializes tensor results into a +// pre-allocated destination location. +// IR Example +// %result = linalg.matmul ins(%a, %b) outs(%init) : tensor<1024x1024xf32> +// %memref = memref.alloc() : memref<1024x1024xf32> +// bufferization.materialize_in_destination %result in %memref +// : tensor<1024x1024xf32> to memref<1024x1024xf32> +// Matching Logic +// 1. Check if matmul result's user is +// bufferization.materialize_in_destination +// 2. If matched, mark it as CUBE and add to cubeSeeds +// Purpose: matmul result materialization on CUBE leverages efficient matrix +// data transfer capabilities. +// ============================================================================ +void OpClassifierPass::matchMaterializePattern(Operation *user) { + if (!isa(user)) + return; + + markCube(user); + cubeSeeds.push_back(user); +} + +// Pattern matching for CUBE operations +int OpClassifierPass::patternMatchCUBE() { + LDBG("--- Step 1: pattern match --->\n"); + + for (Operation *op : allOps) { + if (!isa(op)) + continue; + + // matmul always CUBE + opCoreTypes[op] = OP_CUBE_ONLY; + LLVM_DEBUG(DBGS() << "CUBE (matmul): " << *op << "\n"); + + // ---- Upstream pattern matching ---- + for (Value operand : op->getOperands()) { + Operation *def = operand.getDefiningOp(); + if (!def) + continue; + + matchToTensorPattern(def); + matchTransposePattern(def); + matchFillPattern(def); + } + + // ---- Downstream pattern matching ---- + for (Value result : op->getResults()) { + for (Operation *user : result.getUsers()) { + matchStorePattern(user); + matchExtractSlicePattern(user); + matchMaterializePattern(user); + } + } + } + + LLVM_DEBUG(DBGS() << "seeds: " << cubeSeeds.size() << " to_tensor(s)\n"); + for (Operation *seed : cubeSeeds) { + LLVM_DEBUG(DBGS() << "CUBE: " << *seed << "\n"); + } + + return 0; +} + +// Helper function to get upstream operations based on both SSA and memory +// dependencies +void OpClassifierPass::getUpstreamOpsWithMemoryDeps( + Operation *cur, llvm::SmallVectorImpl &upstreamOps) { + // Collect SSA dependencies (direct operands) + for (Value operand : cur->getOperands()) { + Operation *def = operand.getDefiningOp(); + if (def && isa(def) && + opCoreTypes[def] == OP_CUBE_ONLY) { + continue; + } + if (def && def != cur) { + LLVM_DEBUG(DBGS() << "push op: " << *def << "\n"); + upstreamOps.push_back(def); + } + } + if (!isa(cur)) { + // Only consider memory dependencies for non-to_tensor ops, + // as to_tensor already has SSA deps to its memref source + return; + } + // Collect memory dependencies + if (memDepGraph) { + // Get operations that define memory used by current op + for (Operation *memDef : memDepGraph->getMemDefs(cur)) { + LLVM_DEBUG(DBGS() << "memDef: cur " << *cur << " -> memDef " << *memDef + << "\n"); + if (isa(memDef) && + opCoreTypes[memDef] == OP_UNDETERMINED) { + LLVM_DEBUG(DBGS() << "push op: " << *memDef << "\n"); + upstreamOps.push_back(memDef); + } + } + } +} + +// Propagate CUBE core type upstream +int OpClassifierPass::propagateCubeUpstream() { + LLVM_DEBUG(DBGS() << "--- Step 2: CUBE upstream BFS --->\n"); + llvm::DenseSet cubeVisited; + std::queue cubeQueue; + + for (Operation *seed : cubeSeeds) { + cubeVisited.insert(seed); + cubeQueue.push(seed); + } + + while (!cubeQueue.empty()) { + Operation *cur = cubeQueue.front(); + cubeQueue.pop(); + LLVM_DEBUG(DBGS() << "cur: " << *cur << "\n"); + + // Get upstream operations considering both SSA and memory dependencies + llvm::SmallVector upstreamOps; + getUpstreamOpsWithMemoryDeps(cur, upstreamOps); + + for (Operation *def : upstreamOps) { + if (!def || cubeVisited.count(def) || isa(def)) + continue; + + // Skip operations inside linalg block (internal values) + // But don't skip the linalg op itself + if (isInsideNestedLinalgRegion(def)) { + LLVM_DEBUG(DBGS() << "skip " << def->getName().getStringRef() + << ": inside linalg block\n"); + continue; + } + + cubeVisited.insert(def); + + LLVM_DEBUG(DBGS() << "\tcolor-cube: " << def->getName().getStringRef() + << "\n"); + opCoreTypes[def] = OP_CUBE_ONLY; + cubeQueue.push(def); + } + } + + // Handle fill ops: mark as CUBE only if linalg.fill output buffer is CUBE + markFillOpsAsCube(); + + return 0; +} + +// ============================================================================ +// Step 3: Mark remaining operations as VECTOR +// ============================================================================ +// After Step 2 (CUBE propagation), any operation still marked as UNDETERMINED +// is classified as VECTOR. This establishes the default that non-CUBE +// operations should execute on the VECTOR core. +// ============================================================================ +int OpClassifierPass::markRemainingAsVector() { + LLVM_DEBUG(DBGS() << "--- Step 3: default VECTOR --->\n"); + for (Operation *op : allOps) { + if (isa(op)) { + auto markOp = cast(op); + Value src = markOp.getSrc(); + if (Operation *srcDef = src.getDefiningOp()) { + opCoreTypes[op] = getCoreType(srcDef); + } else { + opCoreTypes[op] = OP_VECTOR_ONLY; + } + continue; + } + + if (opCoreTypes[op] == OP_UNDETERMINED && !isa(op)) { + opCoreTypes[op] = OP_VECTOR_ONLY; + } + } + return 0; +} + +// ============================================================================ +// Helper: Mark fill operations as CUBE when their output buffer is CUBE +// ============================================================================ +// This function handles a special case where linalg.fill's result feeds into +// a CUBE operation. In such cases, the fill itself should be marked CUBE. +// Case 1: fill's outs operand is defined by a CUBE_ONLY operation (not +// tensor.empty) +// IR example +// %empty = tensor.empty() : tensor<1024x1024xf32> // CUBE +// %filled = linalg.fill ins(%value) outs(%empty) // should be CUBE +// Here %filled's outs comes from a CUBE op, so fill is CUBE +// Case 2: fill's outs is a BlockArgument (scf.for/scf.if iter_arg) that is CUBE +// IR example +// scf.for ... iter_args(%arg = %filled) { // %arg is CUBE +// ... +// } +// Here the iter_arg is CUBE, so fill is CUBE +// ============================================================================ +void OpClassifierPass::markFillOpsAsCube() { + LLVM_DEBUG(DBGS() << "--- Marking fill ops as CUBE when outs is CUBE --->\n"); + + for (Operation *op : allOps) { + // Only process fill ops that are still UNDETERMINED + if (!isa(op) || opCoreTypes[op] != OP_UNDETERMINED) { + continue; + } + + auto fillOp = cast(op); + Value outs = fillOp.getDpsInits()[0]; + Operation *outsDef = outs.getDefiningOp(); + + bool outsIsCube = false; + + // Case 1: outs is defined by a CUBE_ONLY operation (not tensor.empty) + // tensor.empty is a special case - it doesn't determine CUBE/VECTOR + if (outsDef && opCoreTypes[outsDef] == OP_CUBE_ONLY && + !isa(outsDef)) { + outsIsCube = true; + LLVM_DEBUG(DBGS() << "\tfill outs defined by CUBE op: " + << outsDef->getName().getStringRef() << "\n"); + } else if (!outsDef) { // Case 2: outs is a BlockArgument (scf.for/scf.if + // iter_arg) + auto blockArg = dyn_cast(outs); + if (blockArg) { + Operation *parentOp = blockArg.getOwner()->getParentOp(); + // Check if it's an scf.for or scf.if iter_arg that is CUBE + if ((isa(parentOp) || isa(parentOp)) && + opCoreTypes[parentOp] == OP_CUBE_ONLY) { + outsIsCube = true; + LLVM_DEBUG(DBGS() << "\tfill outs is CUBE iter_arg of: " + << parentOp->getName().getStringRef() << "\n"); + } + } + } + + if (outsIsCube) { + LLVM_DEBUG(DBGS() << "\tfill-cube (outs is CUBE): " << *op << "\n"); + opCoreTypes[op] = OP_CUBE_ONLY; + } + } +} + +// ============================================================================ +// Step 4: Propagate VECTOR core type upstream +// ============================================================================ +// After Step 3, propagate VECTOR classification to upstream operations. +// BFS traversal from all VECTOR operations, marking their upstream dependencies +// as VECTOR unless they are CUBE-only operations (e.g., matmul). +// Skips builtin.module and func.func operations. +// ============================================================================ +int OpClassifierPass::propagateVectorUpstream() { + LLVM_DEBUG(DBGS() << "--- Step 4: VECTOR upstream BFS --->\n"); + llvm::DenseSet vecVisited; + std::queue vecQueue; + + for (Operation *op : allOps) { + if (opCoreTypes[op] == OP_VECTOR_ONLY) { + vecVisited.insert(op); + vecQueue.push(op); + } + } + + while (!vecQueue.empty()) { + Operation *cur = vecQueue.front(); + vecQueue.pop(); + // Skip builtin.module, func.func, and func.return operations + if (isa(cur) || + (isa(cur) && + opCoreTypes[cur] == OP_CUBE_ONLY)) { + continue; + } + LLVM_DEBUG(DBGS() << "vec-def: " << *cur << "\n"); + + // Get upstream operations considering both SSA and memory dependencies + llvm::SmallVector upstreamOps; + getUpstreamOpsWithMemoryDeps(cur, upstreamOps); + + for (Operation *def : upstreamOps) { + if (!def || vecVisited.count(def)) + continue; + + // Skip operations that should not be marked VECTOR: + // - matmul: never mark matmul as vector (CUBE-only operation) + if (isa(def)) { + LLVM_DEBUG(DBGS() << "skip " << def->getName().getStringRef() + << ": should not be marked VECTOR\n"); + continue; + } + + vecVisited.insert(def); + + LLVM_DEBUG(DBGS() << "\tcolor-vec: " << def->getName().getStringRef() + << "\n"); + opCoreTypes[def] = + static_cast(opCoreTypes[def] | OP_VECTOR_ONLY); + vecQueue.push(def); + } + } + + return 0; +} + +// ============================================================================ +// Helper Functions +// ============================================================================ + +OpCoreType OpClassifierPass::getCoreType(Operation *op) const { + auto it = opCoreTypes.find(op); + if (it != opCoreTypes.end()) { + return it->second; + } + return OP_VECTOR_ONLY; +} + +void OpClassifierPass::setCoreType(Operation *op, OpCoreType coreType) { + opCoreTypes[op] = coreType; +} + +std::string +OpClassifierPass::joinCoreTypes(const std::vector &coreTypes) { + std::string coreTypeStr; + for (size_t i = 0; i < coreTypes.size(); ++i) { + if (i > 0) { + coreTypeStr += ", "; + } + coreTypeStr += coreTypeToString(coreTypes[i]); + } + return coreTypeStr; +} + +// ============================================================================ +// Step 5: SCF Yield Results Handling +// ============================================================================ + +bool OpClassifierPass::shouldStopPropagationForYield( + Operation *op, OpCoreType targetCoreType) { + // Stop at linalg.matmul + if (isa(op)) { + return true; + } + + // Stop at hivm.hir.store + if (isa(op)) { + return true; + } + + // Stop at tensor-type operations that are already VECTOR_ONLY when targeting + // CUBE + if (targetCoreType == OP_CUBE_ONLY) { + if (op->getNumResults() > 0) { + for (Value result : op->getResults()) { + if (isa(result.getType())) { + OpCoreType currentType = getCoreType(op); + if (currentType == OP_VECTOR_ONLY) { + return true; + } + } + } + } + } + + // Stop at scf dialect operations + if (isa(op->getDialect())) { + return true; + } + + // Stop at func.func or builtin.module + if (isa(op)) { + return true; + } + + return false; +} + +void OpClassifierPass::propagateCoreTypeUpwardForYield( + Operation *startOp, OpCoreType targetCoreType) { + if (!startOp) + return; + + llvm::DenseSet visited; + std::queue queue; + queue.push(startOp); + visited.insert(startOp); + + while (!queue.empty()) { + Operation *cur = queue.front(); + queue.pop(); + + if (!cur || !cur->getBlock()) + continue; + + if (shouldStopPropagationForYield(cur, targetCoreType)) + continue; + + opCoreTypes[cur] = targetCoreType; + + // Propagate to operands (upstream) + for (Value operand : cur->getOperands()) { + if (Operation *defOp = operand.getDefiningOp()) { + if (!visited.count(defOp) && defOp != cur) { + visited.insert(defOp); + queue.push(defOp); + } + } + } + + // Also propagate to memory dependencies + if (memDepGraph) { + for (Operation *memDef : memDepGraph->getMemDefs(cur)) { + if (!visited.count(memDef) && memDef != cur) { + visited.insert(memDef); + queue.push(memDef); + } + } + } + } +} + +// handleYieldFromElseRegion: Extracts core_type from then region yield for else +// region processing. For scf.if with else region, when processing the else +// region's yield terminator, we need to use the then region's yield's core_type +// attribute as the authoritative source. If the then yield has a +// comma-separated multi-value type string (e.g., "CUBE,VECTOR"), this function +// extracts the i-th component for the operand at index i. +// Parameters +// - coreTypes: Vector to append the extracted core type to +// - operandIndex: Index of the operand in the yield (for extracting +// multi-value types) +// - thenYieldForElse: The then region's yield operation (from scf.if) +// - operand: The current operand value from the else region's yield +// Returns: true if core_type was successfully extracted and propagated, false +// otherwise. +bool OpClassifierPass::handleYieldFromElseRegion( + std::vector &coreTypes, unsigned operandIndex, + Operation *thenYieldForElse, Value &operand) { + // Only handle if thenYieldForElse is provided and is a scf.yield + if (!thenYieldForElse || !isa(thenYieldForElse)) { + return false; + } + + // Get the core_type attribute from then region yield + auto thenCoreTypeAttr = thenYieldForElse->getAttr("ssbuffer.core_type"); + if (!thenCoreTypeAttr) { + return false; + } + + StringAttr thenCoreTypeStrAttr = dyn_cast(thenCoreTypeAttr); + if (!thenCoreTypeStrAttr) { + return false; + } + + // Extract the i-th component from comma-separated core type string + std::string thenCoreTypeStr = thenCoreTypeStrAttr.getValue().str(); + OpCoreType coreTypeToUse = + parseCoreTypeFromString(thenCoreTypeStr, operandIndex); + coreTypes.push_back(coreTypeToUse); + + // Propagate the determined core_type upstream to the defining operation + if (Operation *defOp = operand.getDefiningOp()) { + propagateCoreTypeUpwardForYield(defOp, coreTypeToUse); + } + + return true; +} + +// ============================================================================ +// processYieldOperation: Set core_type on scf.yield and its parent scf op +// ============================================================================ +// For each operand of scf.yield: +// - If thenYieldForElse is provided (else region of scf.if): use then yield's +// per-operand core_type as reference to determine this operand's type, and +// propagate that type upstream. +// - Otherwise (then region of scf.if, or scf.for): look at the operand's +// defining op's core_type. If it is OP_CUBE_ONLY -> "CUBE", else -> +// "VECTOR". BlockArguments without a defining op (e.g., iter_arg block +// argument) default to VECTOR. +// After collecting all operand types, set ssbuffer.core_type on both the yield +// op and its parent scf operation (scf.if or scf.for). +void OpClassifierPass::processYieldOperation(Operation *op, + Operation *thenYieldForElse) { + std::vector coreTypes; + + // Get the parent SCF operation (scf.if or scf.for) + Operation *parentOp = op->getParentOp(); + + // Process each yield operand in order + if (op->getNumOperands() > 0) { + for (unsigned i = 0; i < op->getNumOperands(); ++i) { + Value operand = op->getOperand(i); + // --------------------------------------------------------------- + // Case 1: else region yield of scf.if + // --------------------------------------------------------------- + // Use then region yield's core_type attribute to determine type. + // If then yield has a comma-separated multi-value type string (e.g. + // "CUBE,VECTOR"), extract the i-th component for this operand. + if (handleYieldFromElseRegion(coreTypes, i, thenYieldForElse, operand)) { + continue; + } + + // --------------------------------------------------------------- + // Case 2: then region yield of scf.if, or scf.for yield + // --------------------------------------------------------------- + // Determine type from the operand's defining operation: + // - OP_CUBE_ONLY -> yield type is "CUBE" + // - anything else (OP_UNDETERMINED, OP_VECTOR_ONLY) -> "VECTOR" + // If the operand is a BlockArgument without a defining op (e.g., an + // iter_arg that is passed through without modification), default to + // VECTOR since no compute op is associated with it. + if (Operation *def = operand.getDefiningOp()) { + OpCoreType ct = getCoreType(def); + if (ct == OP_CUBE_ONLY) { + coreTypes.push_back(OP_CUBE_ONLY); + } else { + coreTypes.push_back(OP_VECTOR_ONLY); + } + } else { + // BlockArgument (e.g., iter_arg passed through scf.for) -> VECTOR + coreTypes.push_back(OP_VECTOR_ONLY); + } + } + } + + // Write ssbuffer.core_type attribute onto the yield op and its parent scf op + if (!coreTypes.empty()) { + std::string coreTypeStr = joinCoreTypes(coreTypes); + op->setAttr("ssbuffer.core_type", + StringAttr::get(op->getContext(), coreTypeStr)); + + OpCoreType opCt = (coreTypeStr.find("CUBE") != std::string::npos) + ? OP_CUBE_ONLY + : OP_VECTOR_ONLY; + opCoreTypes[op] = opCt; + + if (parentOp && llvm::isa(parentOp->getDialect())) { + parentOp->setAttr("ssbuffer.core_type", + StringAttr::get(parentOp->getContext(), coreTypeStr)); + opCoreTypes[parentOp] = opCt; + } + } +} + +// ============================================================================ +// Step 5: Process SCF Yield Operations +// ============================================================================ +// Entry point for yield processing. Walks all scf ops and sets core_type on +// their yield terminators and on the ops themselves. +// Processing order +// 1. All scf.if ops: process then-region yield first (no thenYieldForElse), +// then process else-region yield (passing thenYieldForElse for type +// reference). +// 2. Remaining scf.yield ops not inside scf.if (e.g., scf.for terminators). +// After this step, every scf.yield and its parent scf.if/scf.for carry a +// ssbuffer.core_type attribute reflecting the core_type of their result +// values (derived from the defining ops of each yield operand). +int OpClassifierPass::handleSCFYield() { + // Collect all scf.if operations + llvm::SmallVector ifOps; + for (Operation *op : allOps) { + if (auto ifOp = dyn_cast(op)) { + ifOps.push_back(ifOp); + } + } + + // Process scf.if operations: first then region, then else region + for (scf::IfOp ifOp : ifOps) { + Operation *thenYield = nullptr; + if (!ifOp.getThenRegion().empty()) { + thenYield = ifOp.getThenRegion().back().getTerminator(); + } + + if (isa(thenYield)) { + processYieldOperation(thenYield, nullptr); + } + + if (!ifOp.getElseRegion().empty()) { + Operation *elseYield = ifOp.getElseRegion().back().getTerminator(); + if (isa(elseYield)) { + processYieldOperation(elseYield, thenYield); + } + } + } + + // Process remaining scf.yield operations (not inside scf.if, e.g. scf.for) + for (Operation *op : allOps) { + if (!isa(op)) + continue; + + Operation *parentOp = op->getParentOp(); + if (dyn_cast_or_null(parentOp)) + continue; + + processYieldOperation(op, nullptr); + } + + return 0; +} + +// ============================================================================ +// Step 6: CUBE_AND_VECTOR Operation Handling +// ============================================================================ +// Problem: an operation (e.g., linalg.fill) is used by both CUBE users +// (linalg.matmul) and VECTOR users (arith.addf) simultaneously. +// Both cores cannot share the same op instance. +// Solution: clone the op into two versions. +// - Original op → reclassified as OP_CUBE_ONLY → CUBE users keep using it +// - Cloned op → classified as OP_VECTOR_ONLY → VECTOR users switched to it +// Example +// shared_fill (CUBE_AND_VECTOR) +// ↑ ↑ +// matmul(CUBE) arith.addf(VECTOR) +// shared_fill (CUBE_ONLY) → matmul (CUBE) +// shared_fill_vector_clone → arith.addf (VECTOR) +// (VECTOR_ONLY) +// This step runs after step 4 (VECTOR propagation) and before step 6 +// (SCF yield processing). The do-while loop handles cases where splitting +// an op produces new CUBE_AND_VECTOR ops among its operands. +// ============================================================================ + +void OpClassifierPass::splitOperationForCubeAndVector( + Operation *op, llvm::DenseSet &processedOps, + llvm::DenseMap &opToVectorClone) { + // Guard against double-processing the same op + if (processedOps.count(op)) + return; + processedOps.insert(op); + + // Only split ops explicitly marked CUBE_AND_VECTOR + if (opCoreTypes[op] != OP_CUBE_AND_VECTOR) + return; + + // ------------------------------------------------------------------ + // Phase 1: Recursively handle CUBE_AND_VECTOR operands first + // ------------------------------------------------------------------ + // An operand of this op might itself be CUBE_AND_VECTOR (shared upstream). + // We must split the operand's chain before we can split this op, because + // when we build the operand mapping below, we need the VECTOR clone of + // the operand to be available. Process operands depth-first. + for (Value operand : op->getOperands()) { + if (Operation *def = operand.getDefiningOp()) { + if (opCoreTypes[def] == OP_CUBE_AND_VECTOR) { + splitOperationForCubeAndVector(def, processedOps, opToVectorClone); + } + } + } + + // ------------------------------------------------------------------ + // Phase 2: Build IRMapping — redirect operands to VECTOR clones + // ------------------------------------------------------------------ + // When cloning this op for VECTOR users, any operand that already has + // a VECTOR clone (produced by a recursive call above) must be replaced + // with that clone so the cloned op uses the VECTOR-side values. + IRMapping mapping; + for (Value operand : op->getOperands()) { + if (Operation *def = operand.getDefiningOp()) { + auto it = opToVectorClone.find(def); + if (it != opToVectorClone.end()) { + // Operand has a VECTOR clone — use it in the cloned op + Operation *vectorClone = it->second; + for (unsigned i = 0; i < def->getNumResults(); ++i) { + if (def->getResult(i) == operand) { + mapping.map(operand, vectorClone->getResult(i)); + break; + } + } + } + } + } + + // ------------------------------------------------------------------ + // Phase 3: Clone the op for VECTOR users + // ------------------------------------------------------------------ + OpBuilder builder(op); + Operation *vectorOp = builder.clone(*op, mapping); + + // Reclassify: original → CUBE_ONLY, clone → VECTOR_ONLY + opCoreTypes[op] = OP_CUBE_ONLY; + opCoreTypes[vectorOp] = OP_VECTOR_ONLY; + opToVectorClone[op] = vectorOp; // record for callers' operand mapping + allOps.push_back(vectorOp); // track new op so it gets core_type stamped + + // ------------------------------------------------------------------ + // Phase 4: Redirect VECTOR-only users to the cloned result + // ------------------------------------------------------------------ + // VECTOR users of the original result must now use the cloned result. + // CUBE users still use the original. We iterate through all result uses + // and switch VECTOR-only users over to vectorOp's corresponding result. + for (unsigned i = 0; i < op->getNumResults(); ++i) { + Value originalResult = op->getResult(i); + Value vectorResult = vectorOp->getResult(i); + + llvm::SmallVector usesToUpdate; + for (OpOperand &use : originalResult.getUses()) { + Operation *user = use.getOwner(); + if (user && user != vectorOp && getCoreType(user) == OP_VECTOR_ONLY) { + usesToUpdate.push_back(&use); + } + } + + for (OpOperand *use : usesToUpdate) { + use->set(vectorResult); // redirect to VECTOR clone + } + } +} + +// ============================================================================ +// handleCubeAndVector: entry point for CUBE_AND_VECTOR splitting +// ============================================================================ +// Repeatedly scans for CUBE_AND_VECTOR ops and splits them until no more +// remain. The do-while loop is needed because splitOperationForCubeAndVector +// is called recursively on operands; each recursive call may turn a +// CUBE_AND_VECTOR operand into CUBE_ONLY + VECTOR_ONLY (via the same +// splitOperationForCubeAndVector), which can in turn enable more splits at +// the current level. + +int OpClassifierPass::handleCubeAndVector() { + bool changed = true; + do { + changed = false; + + // Collect all ops still marked CUBE_AND_VECTOR + llvm::SmallVector cubeAndVectorOps; + for (Operation *op : allOps) { + if (opCoreTypes[op] == OP_CUBE_AND_VECTOR) { + cubeAndVectorOps.push_back(op); + } + } + + if (cubeAndVectorOps.empty()) + break; + + llvm::DenseMap opToVectorClone; + llvm::DenseSet processedOps; + + for (Operation *op : cubeAndVectorOps) { + if (processedOps.count(op)) + continue; + splitOperationForCubeAndVector(op, processedOps, opToVectorClone); + changed = true; + } + } while (changed); + + return 0; +} + +// ============================================================================ +// Step 7: Stamp Core Type to IR +// ============================================================================ + +int OpClassifierPass::stampToIR() { + for (Operation *op : allOps) { + if (!op || !op->getBlock()) + continue; + + auto it = opCoreTypes.find(op); + if (it == opCoreTypes.end()) + continue; + OpCoreType coreType = it->second; + + // Skip scf dialect operations + if (llvm::isa(op->getDialect())) + continue; + + // Skip linalg operations' internal block operations + Operation *parent = op->getParentOp(); + bool isInsideLinalgBlock = false; + while (parent) { + if (llvm::isa(parent->getDialect())) { + isInsideLinalgBlock = true; + break; + } + parent = parent->getParentOp(); + } + if (isInsideLinalgBlock) + continue; + + // Skip builtin.module, func.func, and tt.func operations + if (isa(op)) + continue; + + op->setAttr("ssbuffer.core_type", + StringAttr::get(op->getContext(), coreTypeToString(coreType))); + } + + return 0; +} + +// ============================================================================ +// Pre-legalize matmul: Replace matmul with zero-filled accumulator + add +// ============================================================================ +// matmul(A, B, C) -> matmul(A, B, zero) + add(result, C) +// - zero = linalg.fill(tensor.empty(const 0), const 0) +// - add(result, C) adds the original outs value to the matmul result +// This effectively changes a*b+c to result=a*b+0, add(result, c) +// This must be done before initializePass so that the classification pass +// sees the modified IR structure. +// ============================================================================ + +// Helper: bulk delete operations and clean up tracking structures +void OpClassifierPass::bulkDeleteOps( + llvm::SmallVectorImpl &opsToDelete) { + llvm::DenseSet deletedOps(opsToDelete.begin(), + opsToDelete.end()); + for (Operation *op : opsToDelete) { + if (!op || !op->getBlock()) { + continue; + } + if (op->use_empty()) { + auto it = std::find(allOps.begin(), allOps.end(), op); + if (it != allOps.end()) { + allOps.erase(it); + } + opCoreTypes.erase(op); + op->erase(); + } + } + + // Remove any stale references from the classifier map + for (Operation *op : deletedOps) { + opCoreTypes.erase(op); + } +} + +int OpClassifierPass::preLegalizeMatmul() { + LLVM_DEBUG(DBGS() << "--- Pre-legalizing matmul operations --->\n"); + + // Collect all linalg::MatmulOp from the module + llvm::SmallVector matmulOps; + getOperation().walk( + [&](linalg::MatmulOp matmulOp) { matmulOps.push_back(matmulOp); }); + + if (matmulOps.empty()) { + LLVM_DEBUG(DBGS() << "\tNo matmul operations found\n"); + return 0; + } + + LLVM_DEBUG(DBGS() << "\tFound " << matmulOps.size() + << " matmul operations\n"); + + // deferredDelete pattern: collect ops to delete, erase after iteration + // completes avoids iterator invalidation when erasing during the loop + llvm::SmallVector opsToDelete; + + for (linalg::MatmulOp matmulOp : matmulOps) { + if (!matmulOp || !matmulOp->getBlock()) { + continue; + } + + if (matmulOp.getNumResults() == 0) { + continue; + } + + auto outputs = matmulOp.getDpsInits(); + if (outputs.empty()) { + continue; + } + Value outsValue = outputs[0]; + if (!outsValue) { + continue; + } + + Value matmulResult = matmulOp.getResult(0); + + // Get element type from result tensor + auto rankedTensorType = dyn_cast(matmulResult.getType()); + if (!rankedTensorType) { + LLVM_DEBUG(DBGS() << "matmul result is not a ranked tensor, skipping\n"); + continue; + } + Type elemType = rankedTensorType.getElementType(); + + mlir::OpBuilder builder(matmulOp); + Location loc = matmulOp.getLoc(); + + // [Step 1] Create tensor.empty for the new accumulator tensor + // Same shape and type as original matmul output + SmallVector dynamicSizes; + for (int64_t i = 0; i < rankedTensorType.getRank(); ++i) { + if (rankedTensorType.isDynamicDim(i)) { + dynamicSizes.push_back( + builder.create(loc, outsValue, i)); + } + } + auto emptyOp = builder.create( + loc, rankedTensorType.getShape(), rankedTensorType.getElementType(), + dynamicSizes); + + // [Step 2] Create zero constant based on element type + // Supports both floating-point (arith.constant float) and integer types + Value zeroValue; + if (auto floatType = dyn_cast(elemType)) { + APFloat zeroAPFloat = APFloat::getZero(floatType.getFloatSemantics()); + zeroValue = + builder.create(loc, floatType, zeroAPFloat) + .getResult(); + } else if (auto intType = dyn_cast(elemType)) { + zeroValue = + builder.create(loc, intType, 0).getResult(); + } else { + LLVM_DEBUG( + DBGS() << "matmul element type is not float or int, skipping\n"); + continue; + } + + // [Step 3] Use linalg.fill to populate empty tensor with zero -> zero + // accumulator + auto fillOp = builder.create( + loc, ValueRange{zeroValue}, ValueRange{emptyOp.getResult()}); + + // [Step 4] Get matmul's two input matrices A and B + auto inputs = matmulOp.getDpsInputs(); + if (inputs.size() < kMinMatmulInputs) { + LLVM_DEBUG(DBGS() << "matmul has insufficient inputs, skipping\n"); + continue; + } + Value a = inputs[0]; + Value b = inputs[1]; + + // [Step 5] Create new matmul using zero-filled tensor as accumulator + // New matmul runs entirely on CUBE with no VECTOR dependency + auto newMatmul = builder.create( + loc, ValueRange{a, b}, ValueRange{fillOp.getResult(0)}); + + // Copy attributes from original matmul to new matmul (skipping internal + // attrs) + for (auto attr : matmulOp->getAttrs()) { + StringRef attrName = attr.getName().getValue(); + if (attrName == "operandSegmentSizes" || attrName == "res_attrs" || + attrName == "arg_attrs") { + continue; + } + newMatmul->setAttr(attr.getName(), attr.getValue()); + } + + // [Step 6] Create add: add(new_matmul_result, outs_value) + // This is the "c" in a*b+c, added after the matmul result + Operation *addOp; + if (isa(elemType)) { + addOp = + builder.create(loc, newMatmul.getResult(0), outsValue) + .getOperation(); + } else { + addOp = + builder.create(loc, newMatmul.getResult(0), outsValue) + .getOperation(); + } + + // Move addOp to right after newMatmul + addOp->moveAfter(newMatmul.getOperation()); + + // [Step 7] Get add result and replace uses of original matmul result + Value addResult = addOp->getResult(0); + + // Replace all uses of original matmulResult with addResult, EXCEPT addOp's + // own operand + for (auto &use : llvm::make_early_inc_range(matmulResult.getUses())) { + if (use.getOwner() == addOp) { + continue; + } + use.set(addResult); + } + + // [Step 8] Mark new operations as CUBE_ONLY + opCoreTypes[newMatmul.getOperation()] = OP_CUBE_ONLY; + opCoreTypes[emptyOp.getOperation()] = OP_CUBE_ONLY; + opCoreTypes[zeroValue.getDefiningOp()] = OP_CUBE_ONLY; + opCoreTypes[fillOp.getOperation()] = OP_CUBE_ONLY; + opCoreTypes[addOp] = OP_CUBE_ONLY; + + // [Step 9] Mark original matmul for deferred deletion + opsToDelete.push_back(matmulOp.getOperation()); + + LLVM_DEBUG(DBGS() << "[OpClassifier] Transformed matmul to zero-filled " + "accumulator + add\n"); + } + + // Step 10: Delete all marked original matmuls in bulk + bulkDeleteOps(opsToDelete); + + LLVM_DEBUG(DBGS() << "--- Pre-legalize matmul complete --->\n"); + return 0; +} + +// Run the pass +void OpClassifierPass::runOnOperation() { + ModuleOp module = getOperation(); + + LLVM_DEBUG(DBGS() << "\n--- Before Plan Compute Block --->\n"); + LLVM_DEBUG(DBGS() << module << "\n"); + + LLVM_DEBUG(DBGS() << "\n--- Running OpClassifierPass --->\n"); + + // Pre-legalize matmul: transform matmul to zero-filled accumulator + add + // BEFORE initializePass This transforms: matmul(A, B, C) -> matmul(A, B, + // zero) + add(result, C) where zero = linalg.fill(tensor.empty(const 0), + // const 0) This effectively changes a*b+c to result=a*b+0, add(result, c) + if (preLegalizeMatmul() != 0) { + signalPassFailure(); + return; + } + + // Initialize memory dependence graph for tracking memory side effects + + aliasAnalysis = std::make_shared(module); + memDepGraph = std::make_shared( + module, *aliasAnalysis); + + // Initialize the pass + initializePass(module); + + // Step 1: Pattern match around each linalg.matmul to find CUBE seeds + if (patternMatchCUBE() != 0) { + signalPassFailure(); + return; + } + + // Step 2: CUBE upstream BFS from seed loads + if (propagateCubeUpstream() != 0) { + signalPassFailure(); + return; + } + + // Step 3: Mark remaining operations as VECTOR + if (markRemainingAsVector() != 0) { + signalPassFailure(); + return; + } + + // Step 4: VECTOR upstream BFS + if (propagateVectorUpstream() != 0) { + signalPassFailure(); + return; + } + + // Step 5: Handle CUBE_AND_VECTOR operations + if (handleCubeAndVector() != 0) { + signalPassFailure(); + return; + } + + // Step 6: Process SCF yield results + if (handleSCFYield() != 0) { + signalPassFailure(); + return; + } + + // Step 7: Stamp to IR + if (stampToIR() != 0) { + signalPassFailure(); + return; + } + + LLVM_DEBUG(DBGS() << "\n--- OpClassifierPass: end --->\n"); +} + +// Create the pass +namespace mlir { +namespace triton { + +std::unique_ptr> createOpClassifierPass() { + return std::make_unique(); +} + +void registerAddOpClassifierPasses() { + registerPass( + []() -> std::unique_ptr { return createOpClassifierPass(); }); +} + +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlock.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlock.cpp new file mode 100644 index 0000000000..057f0503fc --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlock.cpp @@ -0,0 +1,576 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include +#include + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/LogicalResult.h" +#include "llvm/Support/raw_ostream.h" + +#include "mlir/Analysis/AliasAnalysis.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Operation.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Support/WalkResult.h" + +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlockPass.h" + +#include "ascend/include/DynamicCVPipeline/Common/MemoryEffectsTracker.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h" + +using namespace mlir; +using namespace triton; +using namespace CVPipeline; + +static constexpr const char *DEBUG_TYPE = "PlanCubeBlock"; +#define LOG_DEBUG(...) \ + LLVM_DEBUG(llvm::dbgs() << " [" << DEBUG_TYPE << "] " << __VA_ARGS__) + +static bool isMatmulOp(Operation *op) { return isa(op); } + +namespace { + +class SeedRegionPlanner { + Operation *seed; + Block *block; + const MemoryDependenceGraph &memGraph; + llvm::DenseSet &assigned; + llvm::SmallVectorImpl &group; + bool willCreateCycle(Operation *op); + bool isEligible(Operation *op); + bool tryAddToGroup(Operation *op); + void addSourcesToGroup(); + void addUsersToGroup(); + +public: + SeedRegionPlanner(Operation *seed, Block *block, + const MemoryDependenceGraph &memGraph, + llvm::DenseSet &assigned, + llvm::SmallVectorImpl &group) + : seed(seed), block(block), memGraph(memGraph), assigned(assigned), + group(group) { + group.push_back(seed); + } + + void run(); +}; + +} // namespace + +namespace { + +struct DependencyCycleDetector { + llvm::DenseSet &okSet; + llvm::DenseSet visited; + const MemoryDependenceGraph &memGraph; + Block *block; + void clear() { visited.clear(); } + bool operator()(Operation *cur); + bool dfs(Operation *cur) { return (*this)(cur); }; + + DependencyCycleDetector(Block *block, const MemoryDependenceGraph &memGraph, + llvm::DenseSet &okSet) + : block(block), memGraph(memGraph), okSet(okSet) {} +}; + +} // namespace + +bool DependencyCycleDetector::operator()(Operation *cur) { + if (okSet.contains(cur)) { + return true; + } + if (!visited.insert(cur).second) { + return false; + } + + SmallVector allusers; + allusers.append(cur->getUsers().begin(), cur->getUsers().end()); + for (auto *memUser : memGraph.getExecAfter(cur)) { + allusers.push_back(memUser); + } + for (auto *user : allusers) { + auto *userInBlock = getAncestorInBlock(user, block); + auto &bm = ComputeBlockIdManager::getInstance(); + if (bm.getBlockIdByOp(userInBlock) == -1) { + if (dfs(userInBlock)) { + return true; + } + } else { + for (auto *nx : bm.getOpsByBlockId(bm.getBlockIdByOp(userInBlock))) { + if (dfs(nx)) { + return true; + } + } + } + } + return false; +} + +bool SeedRegionPlanner::willCreateCycle(Operation *op) { + auto *block = op->getBlock(); + llvm::DenseSet okSet(group.begin(), group.end()); + okSet.insert(op); + + DependencyCycleDetector dfs = {block, memGraph, okSet}; + + // DFS from every result in okSet + for (mlir::Operation *okOp : okSet) { + SmallVector allusers; + allusers.append(okOp->getUsers().begin(), okOp->getUsers().end()); + for (auto *memUser : memGraph.getExecAfter(okOp)) { + allusers.push_back(memUser); + } + for (auto *user : allusers) { + auto *userInBlock = getAncestorInBlock(user, block); + if (okSet.contains(userInBlock)) { + continue; + } + auto &bm = ComputeBlockIdManager::getInstance(); + if (bm.getBlockIdByOp(userInBlock) == -1) { + dfs.clear(); + if (dfs(userInBlock)) { + return true; + } + continue; + } + auto opsUsedBlockId = bm.getOpsByBlockId(bm.getBlockIdByOp(userInBlock)); + for (auto *userOp : opsUsedBlockId) { + dfs.clear(); + if (dfs(userOp)) { + return true; + } + } + } + } + return false; +} + +/** + * Checks if an operation is eligible to be added to a Cube group. + * Eligibility depends on it being a cube op, not yet assigned, not a compute + * op, and not creating a cycle in the dependence graph. + */ +bool SeedRegionPlanner::isEligible(Operation *op) { + if (!isCubeOp(op) || assigned.contains(op) || isMatmulOp(op)) { + return false; + } + return !willCreateCycle(op); +} + +bool SeedRegionPlanner::tryAddToGroup(Operation *op) { + if (!op || llvm::is_contained(group, op) || op->getBlock() != block || + !isEligible(op)) { + return false; + } + group.push_back(op); + return true; +} + +void SeedRegionPlanner::addSourcesToGroup() { + size_t head = 0; + while (head < group.size()) { + Operation *currOp = group[head++]; + + // Check data operands + for (Value iop : currOp->getOperands()) { + if (auto *def = iop.getDefiningOp()) { + tryAddToGroup(def); + } + // Check loop-carried dependencies (SCF ForOp iter_args) + if (auto barg = dyn_cast(iop)) { + if (barg.getOwner() == block && isa(block->getParentOp()) && + barg.getArgNumber() > 0) { + auto *yieldOp = barg.getOwner()->getTerminator(); + if (auto *yieldedValDef = yieldOp->getOperand(barg.getArgNumber() - 1) + .getDefiningOp()) { + tryAddToGroup(yieldedValDef); + } + } + } + } + + // Check memory dependencies (RAW/WAW/WAR) + for (auto *def : memGraph.getMemDefs(currOp)) { + tryAddToGroup(def); + } + } +} + +void SeedRegionPlanner::addUsersToGroup() { + llvm::SmallVector queue{seed}; + llvm::DenseSet forwardVisited; + forwardVisited.insert(seed); + + unsigned qIdx = 0; + while (qIdx < queue.size()) { + Operation *currOp = queue[qIdx++]; + SmallVector allUsers; + for (auto *u : currOp->getUsers()) { + allUsers.push_back(u); + } + + for (auto *u : memGraph.getMemUsers(currOp)) { + allUsers.push_back(u); + } + + for (auto *userOp : allUsers) { + auto *userInBlock = getAncestorInBlock(userOp, block); + if (tryAddToGroup(userInBlock)) { + queue.push_back(userInBlock); + } + } + } +} + +/** + * Performs a BFS to expand a group from a seed operation (usually a Dot/Compute + * op). It explores both operands (backward) and users (forward). + */ +void SeedRegionPlanner::run() { + addSourcesToGroup(); + addUsersToGroup(); +} + +namespace { + +/** + * Processes remaining unassigned cube operations by following the topological + * order. + */ +class TopologicalPartitionPlanner { + Block *block; + unsigned nonAssignedCubeCnt = 0; + llvm::DenseMap indegree; + llvm::DenseSet &assigned; + const MemoryDependenceGraph &memGraph; + llvm::DenseSet newassigned; + llvm::DenseSet bypassVisited; + std::queue queue; + + void removeNonCubeOpsRecursively(Operation *op); + llvm::LogicalResult removeReadyNonCubeOps(); + bool shouldSkip(Operation *op) { + return !isCubeOp(op) || assigned.contains(op); + }; + bool canExpandTo(Operation *op); + void dumpQueueAndIndegreeInfo(); + llvm::LogicalResult populateQueueWithReadyOps(); + llvm::SmallVector createNewGroupFromQueue(); + +public: + TopologicalPartitionPlanner(Block *block, + llvm::DenseSet &assigned, + const MemoryDependenceGraph &memGraph) + : block(block), assigned(assigned), memGraph(memGraph) { + initializeIndegreeForBlock(block, indegree, memGraph); + + block->walk([&](Operation *op) { + if (op->getBlock() == block && isCubeOp(op) && !assigned.contains(op)) { + nonAssignedCubeCnt++; + } + }); + } + + llvm::LogicalResult run(); +}; + +} // namespace + +// Recursively bypass non-cube ops: decrement indegree of users; collect +// newly-exposed cube ops +void TopologicalPartitionPlanner::removeNonCubeOpsRecursively(Operation *op) { + LOG_DEBUG("\tRemoved non-cube:" << *op << "\n"); + bypassVisited.insert(op); + auto *block = op->getBlock(); + SmallVector allusers; + allusers.append(op->getUsers().begin(), op->getUsers().end()); + auto &bm = ComputeBlockIdManager::getInstance(); + for (auto *memUser : memGraph.getExecAfter(op)) { + allusers.push_back(memUser); + } + for (auto *user : allusers) { + auto *userInBlock = getAncestorInBlock(user, block); + if (!userInBlock || !indegree.contains(userInBlock) || + ComputeBlockIdManager::getInstance().isSameBlock(userInBlock, op)) { + continue; + } + LOG_DEBUG("Sub indegree to " + << *userInBlock << " from " << *op + << "new degree = " << indegree[userInBlock] - 1 << "\n"); + indegree[userInBlock]--; + if (!bm.isWholeCubeReady(userInBlock, indegree) || + bypassVisited.contains(userInBlock) || !shouldSkip(userInBlock)) { + continue; + } + auto &bm = ComputeBlockIdManager::getInstance(); + auto blockId = bm.getBlockIdByOp(userInBlock); + if (blockId == -1) { + removeNonCubeOpsRecursively(userInBlock); + continue; + } + for (auto *passop : bm.getOpsByBlockId(blockId)) { + if (!bypassVisited.contains(passop)) { + removeNonCubeOpsRecursively(passop); + } + } + } +} + +static bool mapsAreDiff(const llvm::DenseMap &a, + const llvm::DenseMap &b) { + if (a.size() != b.size()) { + return true; + } + return llvm::any_of(a, [&b](std::pair aIter) { + auto bIter = b.find(aIter.first); + return bIter == b.end() || bIter->second != aIter.second; + }); +} + +/** + * Logic to bypass non-cube operations that are ready to be executed. + * This unblocks downstream cube operations in the topological sort. + */ +llvm::LogicalResult TopologicalPartitionPlanner::removeReadyNonCubeOps() { + auto &bm = ComputeBlockIdManager::getInstance(); + auto indegreeBefore = indegree; + size_t beforeVisitedSize = bypassVisited.size(); + for (auto &p : indegree) { + Operation *op = p.first; + if (shouldSkip(op) && bm.isWholeCubeReady(op, indegree) && + !bypassVisited.contains(op)) { + int blockId = bm.getBlockIdByOp(op); + if (blockId == -1) { + removeNonCubeOpsRecursively(op); + } else { + for (auto *passOp : bm.getOpsByBlockId(blockId)) { + if (!bypassVisited.contains(passOp)) { + removeNonCubeOpsRecursively(passOp); + } + } + } + } + } + if (!mapsAreDiff(indegreeBefore, indegree) && + beforeVisitedSize == bypassVisited.size()) { + if (Operation *parentOp = block->getParentOp()) { + parentOp->emitError("PlanCubeBlock cannot make progress while scheduling " + "cube operations"); + } + dumpQueueAndIndegreeInfo(); + return llvm::failure(); + } + return llvm::success(); +} + +// Expansion condition: op must be CUBE_ONLY, indegree == 0 and all its +// dependency ops are CUBE_ONLY +bool TopologicalPartitionPlanner::canExpandTo(Operation *op) { + if (!isCubeOp(op) || assigned.contains(op)) { + return false; + } + auto it = indegree.find(op); + return it->second == 0; +} + +// Encountered error. Need to print failure reason, so no need for LLVM_DEBUG +void TopologicalPartitionPlanner::dumpQueueAndIndegreeInfo() { + // simply print a debug header in a new line + auto errs = []() -> llvm::raw_ostream & { + return llvm::errs() << "\n[" << DEBUG_TYPE << "] "; + }; + + errs() << "failed to make progress while planning cube blocks."; + errs() << "remaining cube count: " << nonAssignedCubeCnt; + + errs() << "ready queue"; + if (queue.empty()) { + llvm::errs() << " is empty."; + } else { + llvm::errs() << ":\n"; + while (!queue.empty()) { + Operation *op = queue.front(); + queue.pop(); + llvm::errs() << " " << *op << "\n"; + } + } + + errs() << "remaining unassigned cube ops:\n"; + bool foundRemainingCube = false; + for (auto &p : indegree) { + Operation *op = p.first; + if (!op || op->getBlock() != block || !CVPipeline::isCubeOp(op) || + assigned.contains(op) || newassigned.contains(op)) { + continue; + } + foundRemainingCube = true; + llvm::errs() << " indegree=" << p.second << " op=" << *op << "\n"; + } + if (!foundRemainingCube) { + llvm::errs() << " \n"; + } +} + +llvm::LogicalResult TopologicalPartitionPlanner::populateQueueWithReadyOps() { + for (auto [op, indegree] : indegree) { + if (indegree < 0) { + op->emitError("Indegree cannot be negative"); + return llvm::failure(); + } + if (indegree == 0 && !newassigned.contains(op) && isCubeOp(op) && + !assigned.contains(op)) { + queue.push(op); + } + } + return llvm::success(); +} + +llvm::SmallVector +TopologicalPartitionPlanner::createNewGroupFromQueue() { + llvm::SmallVector group; + while (!queue.empty()) { + auto *currOp = queue.front(); + queue.pop(); + + newassigned.insert(currOp); + group.push_back(currOp); + nonAssignedCubeCnt--; + + for (auto *user : llvm::concat( + currOp->getUsers(), memGraph.getExecAfter(currOp))) { + auto *userInBlock = getAncestorInBlock(user, block); + if (userInBlock && !newassigned.contains(userInBlock)) { + auto &userInDegree = indegree[userInBlock]; + userInDegree--; + LOG_DEBUG("Sub indegree to " << *userInBlock << " from " << *currOp + << "new degree = " << userInDegree + << "\n"); + if (canExpandTo(userInBlock)) { + queue.push(userInBlock); + } + } + } + } + return group; +} + +llvm::LogicalResult TopologicalPartitionPlanner::run() { + auto &bm = ComputeBlockIdManager::getInstance(); + while (nonAssignedCubeCnt > 0) { + if (failed(populateQueueWithReadyOps())) { + return llvm::failure(); + } + + if (queue.empty()) { + if (failed(removeReadyNonCubeOps())) { + return llvm::failure(); + } + continue; + } + + auto group = createNewGroupFromQueue(); + if (llvm::failed(bm.markOpsWithNewId(group))) { + return llvm::failure(); + } + } + + return llvm::success(); +} + +static SmallVector collectMatmulOps(Block *block) { + SmallVector computeOps; + for (Operation &op : *block) { + if (isMatmulOp(&op)) { + computeOps.push_back(&op); + } + } + return computeOps; +} + +/** + * Main entry point: Process a single block by grouping operations into + * execution blocks using BFS and topological traversal. + */ +static llvm::LogicalResult +processBlockWithCubeBFS(Block *block, const MemoryDependenceGraph &memGraph) { + llvm::DenseSet assigned; + auto allDots = collectMatmulOps(block); + + // Phase 1: Add helper ops (transpose, load/store, ptr etc.) to cube block of + // related matmul + for (auto *dot : allDots) { + if (assigned.contains(dot)) { + continue; + } + + llvm::SmallVector newGroup; + SeedRegionPlanner regionPlanner{dot, block, memGraph, assigned, newGroup}; + regionPlanner.run(); + + for (auto *op : newGroup) { + assigned.insert(op); + } + if (llvm::failed( + ComputeBlockIdManager::getInstance().markOpsWithNewId(newGroup))) { + return llvm::failure(); + } + } + + // Phase 2: Handle remaining Cube Ops following Topo order + TopologicalPartitionPlanner topoPlanner{block, assigned, memGraph}; + return topoPlanner.run(); +} + +void mlir::triton::PlanCubeBlockPass::runOnOperation() { + LOG_DEBUG( + "\n--- Step 2: Partitioning compute blocks for cube operations --->\n"); + auto moduleOp = getOperation(); + auto &aa = getAnalysis(); + auto memGraph = MemoryDependenceGraph(moduleOp, aa); + + // We do not need to skip linalg blocks since they do not have core types and + // do not contain matmul + auto result = moduleOp.walk([&](Block *block) { + if (llvm::failed(processBlockWithCubeBFS(block, memGraph))) { + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + if (result.wasInterrupted()) { + signalPassFailure(); + } + LOG_DEBUG("\n--- Step 2: end --->\n"); +} + +std::unique_ptr> +mlir::triton::createPlanCubeBlockPass() { + return std::make_unique(); +} diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanVectorBlockPass.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanVectorBlockPass.cpp new file mode 100644 index 0000000000..a14c59460b --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanVectorBlockPass.cpp @@ -0,0 +1,435 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/Common/Utils.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/Passes.h" +#include "mlir/Analysis/AliasAnalysis.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/Operation.h" +#include "mlir/Support/LLVM.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SetVector.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/LogicalResult.h" +#include "llvm/Support/raw_ostream.h" +#include +#include + +static constexpr const char *DEBUG_TYPE = "plan-vector-block"; +#define LOG_DEBUG(...) \ + LLVM_DEBUG(llvm::dbgs() << " [" << DEBUG_TYPE << "] " << __VA_ARGS__) + +using namespace mlir; +using namespace triton; + +namespace mlir { +namespace triton { +class PlanVectorBlockPass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PlanVectorBlockPass) + + PlanVectorBlockPass() = default; + void runOnOperation() override; + + llvm::StringRef getArgument() const final { return "plan-vector-block"; } +}; + +bool isFusableOp(Operation *op) { + if (CVPipeline::getOpCoreType(op) == CVPipeline::CoreType::VECTOR_ONLY) { + if (isa(op)) { + // pass control ops like scf::ForOp/scf::IfOp/scf::WhileOp + return false; + } + if (op == op->getBlock()->getTerminator()) { + return false; + } + return true; + } + return false; +} + +void passAndCollectCandidates( + Operation *nowOp, DenseMap &indegree, + SmallVector &candidates, DenseMap &visited, + const CVPipeline::MemoryDependenceGraph &memGraph) { + auto &bm = CVPipeline::ComputeBlockIdManager::getInstance(); + LOG_DEBUG("Bypassing non-fusable op " << *nowOp << "\nnow candidates size: " + << candidates.size() << "\n"); + auto block = nowOp->getBlock(); + SmallVector allusers; + allusers.append(nowOp->getUsers().begin(), nowOp->getUsers().end()); + for (auto memUser : memGraph.getExecAfter(nowOp)) { + allusers.push_back(memUser); + } + for (auto user : allusers) { + auto userInBlock = CVPipeline::getAncestorInBlock(user, block); + if (!userInBlock) { + continue; + } + if (!bm.isSameBlock(userInBlock, nowOp)) { + indegree[userInBlock]--; + } + + if (bm.isWholeCubeReady(userInBlock, indegree)) { + if (!isFusableOp(userInBlock) && !visited[userInBlock]) { + if (bm.getBlockIdByOp(userInBlock) == -1) { + visited[userInBlock] = + true; // mark as fused to avoid duplicate bypass + passAndCollectCandidates(userInBlock, indegree, candidates, visited, + memGraph); + } else { + for (auto cubeop : + bm.getOpsByBlockId(bm.getBlockIdByOp(userInBlock))) { + if (!visited[cubeop]) { + visited[cubeop] = true; + passAndCollectCandidates(cubeop, indegree, candidates, visited, + memGraph); + } + } + } + } else if (isFusableOp(userInBlock) && !visited[userInBlock]) { + visited[userInBlock] = true; + candidates.push_back(userInBlock); + } + } + } + + LOG_DEBUG("After bypassing, candidates size: " << candidates.size() << "\n"); +} + +void byPassNonFusable(DenseMap &indegree, + SmallVector &candidates, + DenseMap &visited, + const CVPipeline::MemoryDependenceGraph &memGraph) { + // for every non-fusable candidates, bypass it. + for (auto &[op, degree] : indegree) { + auto &bm = CVPipeline::ComputeBlockIdManager::getInstance(); + if (bm.isWholeCubeReady(op, indegree) && !isFusableOp(op) && !visited[op]) { + if (bm.getBlockIdByOp(op) == -1) { + visited[op] = true; // mark as fused to avoid duplicate bypass + passAndCollectCandidates(op, indegree, candidates, visited, memGraph); + } else { + for (auto cubeop : bm.getOpsByBlockId(bm.getBlockIdByOp(op))) { + if (!visited[cubeop]) { + visited[cubeop] = true; + passAndCollectCandidates(cubeop, indegree, candidates, visited, + memGraph); + } + } + } + } + } +} + +void updateCandidates(Operation *nextFused, + SmallVector &candidates, + DenseMap &indegree, + DenseMap &visited, + const CVPipeline::MemoryDependenceGraph &memGraph) { + // 1. Already fuse with nextFused, so remove it from candidates + auto &bm = CVPipeline::ComputeBlockIdManager::getInstance(); + for (auto it = candidates.begin(); it != candidates.end(); it++) { + if (*it == nextFused) { + it = candidates.erase(it); + break; + } + } + + // 2. Add new candidates whose indegree becomes 0 after fusing nextFused. + auto block = nextFused->getBlock(); + SmallVector allusers; + allusers.append(nextFused->getUsers().begin(), nextFused->getUsers().end()); + for (auto memUser : memGraph.getExecAfter(nextFused)) { + allusers.push_back(memUser); + } + for (auto user : allusers) { + auto userInBlock = CVPipeline::getAncestorInBlock(user, block); + if (!userInBlock) { + continue; + } + if (!visited[userInBlock]) { + indegree[userInBlock]--; + if (indegree[userInBlock] == 0 && isFusableOp(userInBlock)) { + visited[userInBlock] = true; + candidates.push_back(userInBlock); + } + } + } +} + +void findCandidates(DenseMap &indegree, + SmallVector &candidates, + DenseMap &visited, + const CVPipeline::MemoryDependenceGraph &memGraph) { + // 1. if no candidate, try to bypass non-fusable + LOG_DEBUG("Finding source ops............\n"); + if (candidates.empty()) { + LOG_DEBUG("No candidates available, try bypass\n"); + byPassNonFusable(indegree, candidates, visited, memGraph); + } + // 2. find candidates whose indegree is 0 and not visited, add them to + // candidates and mark visited + for (auto &[op, degree] : indegree) { + if (degree == 0 && isFusableOp(op) && !visited[op]) { + visited[op] = true; + candidates.push_back(op); + } + } + LOG_DEBUG("end finding source ops............\n"); +} + +static SmallVector +findOpsAdjacentToCube(Block *block, const SmallVector &fuseGroup, + DenseMap &visited, + const CVPipeline::MemoryDependenceGraph &memGraph) { + SmallVector toProcess; + for (Operation *op : fuseGroup) { + SmallVector allUsers; + allUsers.append(op->getUsers().begin(), op->getUsers().end()); + for (auto memUser : memGraph.getExecAfter(op)) { + allUsers.push_back(memUser); + } + + for (auto user : allUsers) { + auto userInBlock = CVPipeline::getAncestorInBlock(user, block); + if (!userInBlock || userInBlock == block->getTerminator()) { + continue; + } + if (!isFusableOp(userInBlock) && !visited[userInBlock]) { + toProcess.push_back(op); + } + } + } + return toProcess; +} + +static int getLoopCarriedArgIndex(Value operand, Block *block) { + auto barg = dyn_cast(operand); + if (!barg || barg.getOwner() != block || + !isa(block->getParentOp())) { + return -1; + } + unsigned argIdx = barg.getArgNumber(); + if (argIdx == 0) { + return -1; + } + return argIdx; +} + +static SetVector +collectKeepOps(Block *block, SmallVector toProcess, + const SmallVector &fuseGroup, + const CVPipeline::MemoryDependenceGraph &memGraph) { + SetVector keepOps; + while (!toProcess.empty()) { + Operation *op = toProcess.front(); + toProcess.erase(toProcess.begin()); + if (keepOps.contains(op)) { + continue; + } + keepOps.insert(op); + + // Add all operands to process + for (auto operand : op->getOperands()) { + if (auto defOp = operand.getDefiningOp()) { + if (!keepOps.contains(defOp) && llvm::is_contained(fuseGroup, defOp)) { + toProcess.push_back(defOp); + } + continue; + } + + // Loop-carried dependency: block argument -> yielded value + int argIdx = getLoopCarriedArgIndex(operand, block); + if (argIdx == -1) { + continue; + } + auto barg = cast(operand); + auto *yieldOp = barg.getOwner()->getTerminator(); + auto *yieldedDef = yieldOp->getOperand(argIdx - 1).getDefiningOp(); + if (!keepOps.contains(yieldedDef) && + llvm::is_contained(fuseGroup, yieldedDef)) { + toProcess.push_back(yieldedDef); + } + } + + // Memory dependency + for (auto memDef : memGraph.getExecBefore(op)) { + if (!keepOps.contains(memDef) && llvm::is_contained(fuseGroup, memDef)) { + toProcess.push_back(memDef); + } + } + } + return keepOps; +} + +static void evictAndRestoreState( + Block *block, const SetVector &keepOps, + SmallVector &fuseGroup, DenseMap &visited, + SmallVector &candidates, DenseMap &indegree, + const CVPipeline::MemoryDependenceGraph &memGraph) { + // 1. Collect ops to remove + SmallVector toRemove; + for (Operation *op : fuseGroup) { + if (!keepOps.contains(op)) { + toRemove.push_back(op); + } + } + fuseGroup.assign(keepOps.begin(), keepOps.end()); + + // 2. Restore indegree for successors and reset visited for removed ops + for (Operation *op : toRemove) { + SmallVector allUsers; + allUsers.append(op->getUsers().begin(), op->getUsers().end()); + for (auto memUser : memGraph.getExecAfter(op)) { + allUsers.push_back(memUser); + } + + for (auto user : allUsers) { + if (auto userInBlock = CVPipeline::getAncestorInBlock(user, block)) { + indegree[userInBlock]++; + } + } + visited[op] = false; + } + + // 3. Reset visited for current candidates so they can re-enter + for (auto cand : candidates) { + visited[cand] = false; + } + + // 4. Rebuild candidates from (old candidates + removed ops), zero indegree + // only + SmallVector pool(candidates.begin(), candidates.end()); + pool.append(toRemove.begin(), toRemove.end()); + candidates.clear(); + // Need to find source again + for (Operation *op : pool) { + if (indegree[op] == 0 && !visited[op]) { + visited[op] = true; + candidates.push_back(op); + } + } +} + +void refineFuseGroup(Block *block, SmallVector &nowFuseGroup, + DenseMap &visited, + SmallVector &candidates, + DenseMap &indegree, + const CVPipeline::MemoryDependenceGraph &memGraph) { + // 1.Find ops in fuse group whose next node is a non-fusable (CUBE-only) op + auto toProcess = + findOpsAdjacentToCube(block, nowFuseGroup, visited, memGraph); + // 2. early return: if cannot find one node which next node is CUBE, need to + // check if return or not; + if (toProcess.empty()) { + findCandidates(indegree, candidates, visited, memGraph); + if (candidates.empty()) { + return; + } + } + + // 3. Collect keepOps transitively (data + memory + loop-carried deps) + auto keepOps = collectKeepOps(block, toProcess, nowFuseGroup, memGraph); + + // 4. Remove non-kept ops from fuseGroup and restore BFS state + evictAndRestoreState(block, keepOps, nowFuseGroup, visited, candidates, + indegree, memGraph); + LOG_DEBUG("After cutting, kept " << keepOps.size() << "\n"); +} + +// Main function to plan vector block id for one block +llvm::LogicalResult +planVectorBlockId(Block *block, + const CVPipeline::MemoryDependenceGraph &memGraph) { + // 1. topo initialize + llvm::DenseMap indegree; + llvm::SmallVector queue; + llvm::DenseMap visited; // has been visited in search + CVPipeline::initializeIndegreeForBlock(block, indegree, memGraph); + + // 2. initialize visited and find initial candidates + block->walk([&](Operation *op) { + if (op->getBlock() == block) { + visited[op] = false; + if (isFusableOp(op) && indegree[op] == 0) { + visited[op] = true; + queue.push_back(op); + } + } + }); + findCandidates(indegree, queue, visited, memGraph); + + // 3. find fuse group follow topo order + llvm::SmallVector nowFuseGroup; + while (!queue.empty()) { + auto nextFused = queue.front(); + if (nextFused) { + // fused one && update candidates + nowFuseGroup.push_back(nextFused); + updateCandidates(nextFused, queue, indegree, visited, memGraph); + } + if (queue.empty() || nextFused == nullptr) { + // finish one group, assign block id and start next iteration + // Cut error operations before assigning block id + refineFuseGroup(block, nowFuseGroup, visited, queue, indegree, memGraph); + + auto &bm = CVPipeline::ComputeBlockIdManager::getInstance(); + if (llvm::failed(bm.markOpsWithNewId(nowFuseGroup))) { + return llvm::failure(); + } + nowFuseGroup.clear(); + // reset queue + findCandidates(indegree, queue, visited, memGraph); + } + } + return llvm::success(); +} + +void PlanVectorBlockPass::runOnOperation() { + ModuleOp module = getOperation(); + // 1. Build memory dependence graph + auto &aliasAnalysis = getAnalysis(); + CVPipeline::MemoryDependenceGraph memDepGraph(module, aliasAnalysis); + + // 2. search blocks in topo order and assign block id for each block + llvm::SmallVector blocks; + module.walk([&](Block *block) { + if (llvm::failed(planVectorBlockId(block, memDepGraph))) { + module.emitError() << "[" << DEBUG_TYPE + << "] Failed to plan vector block id for block"; + signalPassFailure(); + return; + } + }); +} + +std::unique_ptr> createPlanVectorBlockPass() { + return std::make_unique(); +} + +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ReorderOpsByBlockId.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ReorderOpsByBlockId.cpp new file mode 100644 index 0000000000..5ed3d6026e --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ReorderOpsByBlockId.cpp @@ -0,0 +1,400 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include +#include + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/iterator.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/Error.h" +#include "llvm/Support/LogicalResult.h" +#include "llvm/Support/raw_ostream.h" + +#include "mlir/Analysis/AliasAnalysis.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/Operation.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Support/WalkResult.h" + +#include "ascend/include/DynamicCVPipeline/Common/MemoryEffectsTracker.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/ReorderOpsByBlockId.h" + +#include "ascend/include/DynamicCVPipeline/Common/Utils.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h" +#include "incubated/Conversion/TritonToUnstructureIncubated/OffsetAnalysis.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +using namespace mlir; +static constexpr const char *DEBUG_TYPE = "ReorderOpsByBlockIdPass"; +#define LOG_DEBUG(...) \ + LLVM_DEBUG(llvm::dbgs() << " [" << DEBUG_TYPE << "] " << __VA_ARGS__) + +using namespace triton; +using namespace CVPipeline; + +namespace { + +// A dependency DAG of both SSA and memory of the ops +struct BlockOpGraph { + Block *block; + ArrayRef ops; + DenseMap opIndex; // op → position in ops + DenseMap> preds; // op → its defs + DenseMap> succs; // op → its uses + BlockOpGraph(ArrayRef allOps, Block *block, + const MemoryDependenceGraph &memGraph); +}; + +// Helper class to manage edges in OpGraph, mainly to reduce congitive +// complexity of the build function +struct EdgeHelper { + BlockOpGraph &graph; + DenseSet> seen; + Block *block; + + // find the ancestor directly in the block, and in opIndex; return nullptr if + // either fails + Operation *resolveToBlockOp(Operation *op); + + template + void addEdge(Operation *pred, Operation *succ); + + template + void addEdgeToUser(Operation *op, Operation *user) { + if (graph.opIndex.contains(user)) { + return; // same-level use, already covered by the def-side loop + } + Operation *ancestor = resolveToBlockOp(user); + addEdge(op, ancestor); + }; + + EdgeHelper(BlockOpGraph &g, Block *block) : graph(g), block(block) {}; +}; + +} // namespace + +Operation *EdgeHelper::resolveToBlockOp(Operation *op) { + if (graph.opIndex.contains(op)) { + return op; + } + Operation *ancestor = getAncestorInBlock(op, block); + if (!ancestor || !graph.opIndex.contains(ancestor)) { + return nullptr; + } + return ancestor; +} + +template +void EdgeHelper::addEdge(Operation *pred, Operation *succ) { + if (!pred || !succ || pred == succ) { + return; + } + if (seen.insert({pred, succ}).second) { + LOG_DEBUG("Adding " << (IsMemory ? "memory " : "") << "edge from " << *pred + << " to " << *succ << "\n"); + graph.succs[pred].push_back(succ); + graph.preds[succ].push_back(pred); + } +}; + +BlockOpGraph::BlockOpGraph(ArrayRef allOps, Block *block, + const MemoryDependenceGraph &memGraph) + : block(block), ops(allOps) { + for (unsigned i = 0; i < allOps.size(); ++i) { + opIndex[allOps[i]] = i; + preds[allOps[i]]; // ensure every node has an entry + succs[allOps[i]]; + } + + EdgeHelper edges(*this, block); + + for (Operation *op : allOps) { + LOG_DEBUG("Processing op: " << *op << "\n"); + // Edges from operand defs (including defs nested inside other ops). + for (Value const operand : op->getOperands()) { + Operation *defOp = operand.getDefiningOp(); + if (!defOp) { + continue; + } + Operation *def = edges.resolveToBlockOp(defOp); + edges.addEdge(def, op); + } + + // Edges from uses that live inside nested regions of another block-level + // op. + for (Value const result : op->getResults()) { + for (Operation *user : result.getUsers()) { + edges.addEdgeToUser(op, user); + } + } + + for (auto *memDef : memGraph.getExecBefore(op)) { + Operation *def = edges.resolveToBlockOp(memDef); + edges.addEdge(def, op); + } + + for (auto *memUser : memGraph.getExecAfter(op)) { + edges.addEdgeToUser(op, memUser); + } + } +} + +static llvm::FailureOr> +collectBlockIds(ArrayRef allOps) { + DenseMap opBlockId; + for (Operation *op : allOps) { + if (llvm::failed(verifyOpBlockId(op))) { + return llvm::failure(); + } + auto blockIdAttrRes = getOpBlockId(op); + int64_t blockId = blockIdAttrRes.has_value() + ? blockIdAttrRes.value() + : ComputeBlockIdManager::getInstance().getNextId(); + opBlockId[op] = blockId; + } + return opBlockId; +} + +namespace { + +// Helper structure to hold the group-level graph data. +struct GroupAdjacencyGraph { + Block *block; + SmallVector groupIds; + SmallVector> succs; + SmallVector inDeg; + GroupAdjacencyGraph(const BlockOpGraph &g, + const DenseMap &opBlockId); + llvm::FailureOr> computeTopologicalOrder(); +}; + +} // namespace + +/** + * Step 1: Build the group-level dependency graph from operator-level edges. + * Maps individual operations to their respective groups and identifies + * dependencies between those groups. + */ +GroupAdjacencyGraph::GroupAdjacencyGraph( + const BlockOpGraph &g, const DenseMap &opBlockId) + : block(g.block) { + // 1. Collect distinct group IDs while preserving the first-appearance order. + DenseSet seenIds; + for (Operation *op : g.ops) { + int id = opBlockId.at(op); + if (seenIds.insert(id).second) { + groupIds.push_back(id); + } + } + + unsigned n = groupIds.size(); + succs.resize(n); + inDeg.assign(n, 0); + + // Map group ID to its index in the groupIds vector for fast lookup. + DenseMap groupPos; + for (unsigned i = 0; i < n; ++i) { + groupPos[groupIds[i]] = i; + } + + // 2. Build group-level edges. Use a set to avoid duplicate edges between + // groups. + DenseSet> addedEdges; + for (Operation *op : g.ops) { + unsigned fromIdx = groupPos[opBlockId.at(op)]; + + for (Operation *succ : g.succs.at(op)) { + unsigned toIdx = groupPos[opBlockId.at(succ)]; + // Ignore intra-group dependencies and duplicate inter-group edges. + if (fromIdx != toIdx && addedEdges.insert({fromIdx, toIdx}).second) { + succs[fromIdx].push_back(toIdx); + inDeg[toIdx]++; + } + } + } + + // Logging the constructed group graph. + LOG_DEBUG("Group-level edges:\n"); + for (unsigned i = 0; i < n; ++i) { + LOG_DEBUG(" Group " << groupIds[i] << " -> "); + for (unsigned succIdx : succs[i]) { + LOG_DEBUG(groupIds[succIdx] << " "); + } + LOG_DEBUG("\n"); + } +} + +/** + * Step 2: Perform a topological sort (Kahn's Algorithm) on the group graph. + * Returns the group IDs in an order that satisfies all dependencies. + */ +llvm::FailureOr> +GroupAdjacencyGraph::computeTopologicalOrder() { + SmallVector result; + SmallVector ready; // Nodes with in-degree 0. + unsigned n = groupIds.size(); + + for (unsigned i = 0; i < n; ++i) { + if (inDeg[i] == 0) { + ready.push_back(i); + } + } + + while (!ready.empty()) { + auto cur = ready.pop_back_val(); + + result.push_back(groupIds[cur]); + + for (unsigned succIdx : succs[cur]) { + if (--inDeg[succIdx] == 0) { + ready.push_back(succIdx); + } + } + } + + LOG_DEBUG("Group order: "); + for (int id : result) { + LOG_DEBUG(id << " "); + } + LOG_DEBUG("\n"); + + if (result.size() == n) { + return result; + } + Operation *op = block->getParentOp(); + constexpr std::string_view kErrorPrefix = + "Failed to compute topological order for "; + if (!op) { + llvm::errs() << kErrorPrefix + << "an unknown block that is not contained in an op"; + return llvm::failure(); + } + size_t regionIdx = 0; + bool found = false; + for (auto [i, region] : llvm::enumerate(op->getRegions())) { + for (auto &possibleBlock : region.getBlocks()) { + if (&possibleBlock == block) { + regionIdx = i; + } + } + } + op->emitError(kErrorPrefix) << "block in region " << regionIdx; + return llvm::failure(); +} + +// Stable sort ops based on their group orders +static llvm::FailureOr> +buildReorderedOps(const BlockOpGraph &graph, + const DenseMap &opBlockId) { + SmallVector reordered; + GroupAdjacencyGraph adjacencyGraph{graph, opBlockId}; + auto groupOrderResult = adjacencyGraph.computeTopologicalOrder(); + if (llvm::failed(groupOrderResult)) { + return llvm::failure(); + } + + for (int const blockId : groupOrderResult.value()) { + for (Operation *op : graph.ops) { + if (opBlockId.at(op) == blockId) { + reordered.push_back(op); + } + } + } + return reordered; +} + +// Reorder the ops in the mlir representation +static void applyReorder(Block &block, ArrayRef reordered) { + Operation *terminator = + block.mightHaveTerminator() ? block.getTerminator() : nullptr; + for (Operation *op : reordered) { + op->moveBefore(&block, block.end()); + } + + if (terminator) { + terminator->moveBefore(&block, block.end()); + } +} + +static llvm::LogicalResult +reorderOpsInBlock(Block &block, const MemoryDependenceGraph &memGraph) { + const auto allOps = + llvm::to_vector(llvm::make_pointer_range(block.without_terminator())); + + const BlockOpGraph graph{allOps, &block, memGraph}; + llvm::FailureOr> opBlockIdOpt = + collectBlockIds(allOps); + if (failed(opBlockIdOpt)) { + return failure(); + } + + auto &opBlockId = *opBlockIdOpt; + LOG_DEBUG("Initial opBlockIds:\n"); + for (Operation *op : allOps) { + LOG_DEBUG(" Op: " << *op << ", opBlockId = " << opBlockId[op] << "\n"); + } + + const auto reorderedRes = buildReorderedOps(graph, opBlockId); + if (failed(reorderedRes)) { + return failure(); + } + + applyReorder(block, reorderedRes.value()); + + return llvm::success(); +} + +void ReorderOpsByBlockIdPass::runOnOperation() { + LOG_DEBUG("\n=== Pass: TuningOpSeq ===\n"); + OpBuilder const builder(&getContext()); + + auto moduleOp = getOperation(); + auto &aa = getAnalysis(); + auto memGraph = MemoryDependenceGraph(moduleOp, aa); + moduleOp.walk([&](Block *block) { + auto *parentOp = block->getParentOp(); + if (!parentOp || + // whitelist ops to reorder + !(isa(parentOp) || + isa(parentOp->getDialect()))) { + return WalkResult::skip(); + } + if (llvm::failed(reorderOpsInBlock(*block, memGraph))) { + signalPassFailure(); + } + return WalkResult::advance(); + }); + + LOG_DEBUG("=== Pass TuningOpSeq complete ===\n"); + ComputeBlockIdManager::getInstance().reset(); +} + +std::unique_ptr> +mlir::triton::createReorderOpsByBlockIdPass() { + return std::make_unique(); +} diff --git a/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoad.cpp b/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoad.cpp new file mode 100644 index 0000000000..d92624dfd2 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoad.cpp @@ -0,0 +1,176 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadInternal.h" +#include "ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadPass.h" + +using namespace mlir; +using namespace triton; +using namespace gmload; + +// Default multi-buffer depth for GM load operations. +static constexpr int kDefaultBufferDepth = 2; + +// ============================================================================ +// Step functions +// ============================================================================ + +void AddMultiBufferToGMLoadPass::collectAndGroupMarkedOps() { + auto module = getOperation(); + + // Scan all IR for marked ops (generic, no loop dependency). + markedOps_ = collectMarkedOps(module); + if (markedOps_.empty()) { + LOG_DEBUG("No marked loads found, nothing to transform\n"); + return; + } + LOG_DEBUG("Marked loads collected, start transformation\n"); + + // Group by enclosing scf::ForOp. + groupByEnclosingForOp(markedOps_, contexts_); + + // Apply depth policy: skip loops whose compile-time trip count is too small + // to benefit, then record the slot count on each group. + int depth = kDefaultBufferDepth; + llvm::erase_if(contexts_, [depth](const ForBufferCtx &context) { + if (auto tripCount = getConstantTripCount(context.forOp)) + return *tripCount <= depth; + return false; + }); + for (auto &context : contexts_) + for (auto &group : context.groups) + group.depth = depth; + + if (contexts_.empty()) + LOG_DEBUG("No bufferable loops found\n"); +} + +void AddMultiBufferToGMLoadPass::sortContextsInnerFirst() { + // Sort contexts inner-first so that inner loops are transformed + // before the outer loops that contain them. + auto getNestingDepth = [](Operation *loopOp) { + unsigned nestingDepth = 0; + for (Operation *parentOp = loopOp->getParentOp(); parentOp; + parentOp = parentOp->getParentOp()) + ++nestingDepth; + return nestingDepth; + }; + + llvm::sort(contexts_, [&getNestingDepth](const ForBufferCtx &leftContext, + const ForBufferCtx &rightContext) { + Operation *leftLoopOp = + const_cast(leftContext.forOp).getOperation(); + Operation *rightLoopOp = + const_cast(rightContext.forOp).getOperation(); + unsigned leftDepth = getNestingDepth(leftLoopOp); + unsigned rightDepth = getNestingDepth(rightLoopOp); + if (leftDepth != rightDepth) + return leftDepth > rightDepth; + if (leftLoopOp->getBlock() == rightLoopOp->getBlock()) + return leftLoopOp->isBeforeInBlock(rightLoopOp); + return false; + }); +} + +LogicalResult AddMultiBufferToGMLoadPass::applyMultiBufferToGMLoadLoops() { + for (auto &context : contexts_) + allCtxForOps_.insert(context.forOp.getOperation()); + + // Process inner loops before outer loops. + for (auto &context : contexts_) + if (failed(applyMultiBufferToForLoop(context, allCtxForOps_))) + return failure(); + + return success(); +} + +void AddMultiBufferToGMLoadPass::cleanupTransformedIR() { + auto module = getOperation(); + + // Deduplicate untagged constants that inner-loop processing created + // before outer-loop processing could provide the dominating equivalents. + deduplicateConstants(module); + + // Erase replaced original for ops. + llvm::DenseSet nestedForOps; + for (auto &context : contexts_) { + Operation *parentOp = context.forOp->getParentOp(); + while (parentOp) { + if (allCtxForOps_.contains(parentOp)) { + nestedForOps.insert(context.forOp.getOperation()); + break; + } + parentOp = parentOp->getParentOp(); + } + } + + for (auto &context : llvm::reverse(contexts_)) { + if (nestedForOps.contains(context.forOp.getOperation())) + continue; + context.forOp.erase(); + } +} + +// ============================================================================ +// Pass entry point +// ============================================================================ + +void AddMultiBufferToGMLoadPass::runOnOperation() { + auto module = getOperation(); + LOG_DEBUG("Enter add-multi-buffer-to-gm-load pass\n"); + LOG_DEBUG("Before add-multi-buffer-to-gm-load:\n" << module << "\n"); + + // Step 1: Collect marked ops and group by enclosing forOp + markedOps_.clear(); + contexts_.clear(); + collectAndGroupMarkedOps(); + if (contexts_.empty()) + return; + + // Step 2: Sort contexts inner-first + sortContextsInnerFirst(); + + // Step 3: Transform each for loop with multi-buffer logic + allCtxForOps_.clear(); + if (failed(applyMultiBufferToGMLoadLoops())) { + module.emitError() << "[" << DEBUG_TYPE + << "] Step 3 applyMultiBufferToGMLoadLoops failed"; + signalPassFailure(); + return; + } + + // Step 4: Cleanup transformed IR + cleanupTransformedIR(); + + LOG_DEBUG("After add-multi-buffer-to-gm-load:\n" << module << "\n"); + LOG_DEBUG("Process successfully\n"); +} + +namespace mlir { +namespace triton { + +std::unique_ptr> createAddMultiBufferToGMLoadPass() { + return std::make_unique(); +} + +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/AsyncLoadHoisting.cpp b/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/AsyncLoadHoisting.cpp new file mode 100644 index 0000000000..e6ab7a79a0 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/AsyncLoadHoisting.cpp @@ -0,0 +1,373 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AsyncLoadHoistingPass.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/Value.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "mlir/Pass/Pass.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "async-load-hoisting" + +using namespace mlir; +using namespace triton; + +namespace mlir { +namespace triton { + +// Check if op is a mask processing operation (broadcast, cmpf, select, etc.) +// Mask values are not directly fed into Cube +static bool isMaskProcessingOp(Operation *op) { + return mlir::isa(op); +} + +// Recursively check if all downstream users of op are mask processing +// operations +static bool +isOnlyUsedByMaskProcessingOpsRecursively(Operation *op, + llvm::DenseSet &visited) { + if (visited.count(op)) { + return true; // Already visited + } + visited.insert(op); + + if (op->getNumResults() == 0) { + return false; + } + + bool hasAnyUsers = false; + for (auto result : op->getResults()) { + for (auto *user : result.getUsers()) { + hasAnyUsers = true; + if (!isMaskProcessingOp(user) && + !isOnlyUsedByMaskProcessingOpsRecursively(user, visited)) { + return false; + } + } + } + if (!hasAnyUsers) { + return false; + } + return true; +} + +// Check if all downstream users of op are mask processing operations +static bool isOnlyUsedByMaskProcessingOps(Operation *op) { + llvm::DenseSet visited; + return isOnlyUsedByMaskProcessingOpsRecursively(op, visited); +} + +static bool isLoadOp(Operation *op) { + auto effectInterface = dyn_cast(op); + if (!effectInterface) { + return false; + } + + // Skip memref.copy (has exactly 2 operands: src and dst) + if (mlir::isa(op)) { + return false; + } + + // Skip memref.load with all scalar indices + if (mlir::isa(op)) { + bool allScalarIndices = true; + for (size_t i = 1; i < op->getNumOperands(); ++i) { + auto *defOp = op->getOperand(i).getDefiningOp(); + if (!defOp || !llvm::isa(defOp)) { + allScalarIndices = false; + break; + } + } + if (allScalarIndices && op->getNumOperands() > 1) { + return false; + } + } + + if (isOnlyUsedByMaskProcessingOps(op)) { + return false; + } + + return effectInterface.hasEffect(); +} + +// Check if a Value's computation chain is linear and predictable. +// Returns false if the chain contains remsi (modulo), load (indirect +// addressing), or complex loops. +static bool isLinearAndPredictable(mlir::Value value, + llvm::DenseSet &visited) { + if (value.getDefiningOp()) + return true; + if (mlir::isa(value)) + return true; + if (!visited.insert(value).second) + return true; + + mlir::Operation *defOp = value.getDefiningOp(); + if (!defOp) + return true; + + // Block nonlinear ops: memref.load (indirect addressing), scf.while/for + // (complex control flow) + if (mlir::isa( + defOp)) { + return false; + } + + // Recursively check operands + for (mlir::Value operand : defOp->getOperands()) { + mlir::Operation *operandDefOp = operand.getDefiningOp(); + if (!operandDefOp || + operandDefOp->getParentRegion() != defOp->getParentRegion()) { + return true; // Cross-region boundary, treat as linear + } + if (!isLinearAndPredictable(operand, visited)) { + return false; + } + } + + return true; +} + +// Check if all reinterpret_cast offsets in chain are linearly predictable +static bool +filterNonLinearReinterpretCast(const std::vector &chain) { + for (mlir::Operation *op : chain) { + auto castOp = llvm::dyn_cast(op); + if (!castOp) + continue; + + llvm::DenseSet visited; + + for (auto offset : castOp.getOffsets()) { + if (!isLinearAndPredictable(offset, visited)) { + return false; + } + } + } + return true; +} + +// Get ssbuffer.block_id attribute value, returns -1 if not present +static int64_t getBlockId(Operation *op) { + if (auto blockIdAttr = op->getAttrOfType("ssbuffer.block_id")) { + return blockIdAttr.getInt(); + } + return -1; +} + +static void collectAddressGenerationOps(Operation *op, int64_t expectedBlockId, + llvm::DenseSet &visited, + std::vector &chain) { + for (auto operand : op->getOperands()) { + auto *defOp = operand.getDefiningOp(); + + if (!defOp || visited.count(defOp) || llvm::isa(defOp)) { + continue; + } + + if (isLoadOp(defOp)) + continue; + + if (mlir::isa(operand.getType())) { + for (auto *user : operand.getUsers()) { + if (auto copyOp = llvm::dyn_cast(user); + copyOp && copyOp.getTarget() == operand && + visited.insert(copyOp).second) { + chain.push_back(copyOp); + collectAddressGenerationOps(copyOp, expectedBlockId, visited, chain); + continue; + } + if (mlir::isa(user) && visited.insert(user).second) { + chain.push_back(user); + for (auto result : user->getResults()) { + for (auto *subviewUser : result.getUsers()) { + if (auto copyOp = llvm::dyn_cast(subviewUser); + visited.insert(copyOp).second) { + chain.push_back(copyOp); + collectAddressGenerationOps(copyOp, expectedBlockId, visited, + chain); + } + } + } + collectAddressGenerationOps(user, expectedBlockId, visited, chain); + } + } + } + + visited.insert(defOp); + chain.push_back(defOp); + + collectAddressGenerationOps(defOp, expectedBlockId, visited, chain); + } +} +// Returns full chain and filtered chain (only ops with same block_id) +static std::pair, std::vector> +getAddressGenerationChainWithFilter(Operation *loadOp) { + std::vector chain; + chain.push_back(loadOp); + + llvm::DenseSet visitedOp; + visitedOp.insert(loadOp); + + int64_t expectedBlockId = getBlockId(loadOp); + + collectAddressGenerationOps(loadOp, expectedBlockId, visitedOp, chain); + + if (chain.size() > 1) { + auto *block = loadOp->getBlock(); + auto sortStartIter = chain.begin() + 1; + std::stable_sort(sortStartIter, chain.end(), + [block](Operation *a, Operation *b) { + for (auto &op : *block) { + if (&op == a) + return true; + if (&op == b) + return false; + } + return false; + }); + auto loadOpIter = chain.begin(); + std::rotate(loadOpIter, loadOpIter + 1, chain.end()); + } + + std::vector filteredChain; + for (Operation *op : chain) { + if (getBlockId(op) == expectedBlockId) { + filteredChain.push_back(op); + } + } + + return {chain, filteredChain}; +} + +// Check if chain contains any block argument +static bool chainContainsBlockArg(Operation *loadOp) { + auto [fullChain, filteredChain] = getAddressGenerationChainWithFilter(loadOp); + + for (Operation *op : fullChain) { + for (auto operand : op->getOperands()) { + if (!operand.getDefiningOp()) { + return true; + } + } + } + return false; +} + +// Process all load ops with the same block_id +static void scanAndHoistBlock(llvm::SmallVector &cache) { + for (Operation *opPtr : cache) { + bool isLoad = isLoadOp(opPtr); + bool hasBlockArg = chainContainsBlockArg(opPtr); + if (isLoad) { + auto [fullChain, filteredChain] = + getAddressGenerationChainWithFilter(opPtr); + (void)fullChain; + bool passNonLinear = filterNonLinearReinterpretCast(filteredChain); + if (passNonLinear && hasBlockArg) { + opPtr->setAttr("gm_load_bufferable", + UnitAttr::get(opPtr->getContext())); + } + } + } +} + +// Traverse all ops in region, grouped by ssbuffer.block_id +static void asyncLoadHoistingImpl(Region ®ion) { + llvm::SmallVector cache; + + for (Block &block : llvm::make_early_inc_range(region)) { + for (Operation &op : llvm::make_early_inc_range(block)) { + int64_t opBlockId = getBlockId(&op); + + if (op.getNumRegions() > 0) { + for (Region &childRegion : op.getRegions()) { + if (childRegion.empty()) + continue; + asyncLoadHoistingImpl(childRegion); + } + continue; + } + + if (cache.empty()) { + cache.push_back(&op); + continue; + } + + int64_t cacheBlockId = getBlockId(cache.front()); + if (opBlockId == cacheBlockId && opBlockId != -1) { + cache.push_back(&op); + } else { + scanAndHoistBlock(cache); + cache.clear(); + cache.push_back(&op); + } + } + } + + if (!cache.empty()) { + scanAndHoistBlock(cache); + cache.clear(); + } +} + +} // namespace triton +} // namespace mlir + +void AsyncLoadHoistingPass::runOnOperation() { + auto module = getOperation(); + + LLVM_DEBUG({ + llvm::dbgs() << "[async-load-hoisting] Before AsyncLoadHoistingPass:\n" + << module << "\n"; + }); + + module->walk( + [&](func::FuncOp func) { asyncLoadHoistingImpl(func.getBody()); }); + + LLVM_DEBUG({ + llvm::dbgs() << "[async-load-hoisting] After AsyncLoadHoistingPass:\n" + << module << "\n"; + }); +} + +std::unique_ptr> +mlir::triton::createAsyncLoadHoistingPass() { + return std::make_unique(); +} + +void mlir::triton::registerAsyncLoadHoistingPasses() { + registerPass([]() -> std::unique_ptr { + return createAsyncLoadHoistingPass(); + }); +} diff --git a/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/DependencyAnalysis.cpp b/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/DependencyAnalysis.cpp new file mode 100644 index 0000000000..9d2e35d856 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/DependencyAnalysis.cpp @@ -0,0 +1,393 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadInternal.h" + +namespace gmload { + +// ============================================================================ +// Analysis: dependency chain computation +// ============================================================================ + +/// Given an op that may be nested inside a region (e.g., inside scf.if), +/// walk up the parent chain and return the ancestor op that lives directly +/// in `scopeBlock`. Returns nullptr if no such ancestor exists. +Operation *getAncestorInBlock(Operation *nestedOp, Block *scopeBlock) { + while (nestedOp) { + if (nestedOp->getBlock() == scopeBlock) + return nestedOp; + nestedOp = nestedOp->getParentOp(); + } + return nullptr; +} + +/// Add `op` and all its transitive SSA operand-producers within `scopeBlock` +/// to `visited`. Newly discovered ops are also appended to `worklist`. +void backwardTrace(Operation *sourceOp, Block *scopeBlock, + llvm::SetVector &visited, + SmallVectorImpl &worklist) { + for (Value operand : sourceOp->getOperands()) { + auto *defOp = operand.getDefiningOp(); + if (!defOp || defOp->getBlock() != scopeBlock) + continue; + if (visited.insert(defOp)) + worklist.push_back(defOp); + } +} + +/// Process worklist with backward SSA trace until empty. +/// Used in Phase 1 and Phase 3 of computeLoadChain. +void backwardTraceFromWorklist(llvm::SetVector &visited, + SmallVectorImpl &worklist, + Block *scopeBlock) { + while (!worklist.empty()) { + Operation *currentOp = worklist.pop_back_val(); + backwardTrace(currentOp, scopeBlock, visited, worklist); + } +} + +/// Forward trace from every memref::AllocOp in visited set. +/// Returns newly discovered ops that use alloc results. +/// Used in Phase 2 of computeLoadChain. +SmallVector +forwardTraceFromAllocs(llvm::SetVector &visited, + Block *scopeBlock) { + SmallVector valueWorklist; + for (Operation *currentOp : visited) { + if (isa(currentOp)) { + for (Value result : currentOp->getResults()) + valueWorklist.push_back(result); + } + } + + SmallVector newlyDiscoveredOps; + while (!valueWorklist.empty()) { + Value currentValue = valueWorklist.pop_back_val(); + for (Operation *user : currentValue.getUsers()) { + Operation *ancestor = getAncestorInBlock(user, scopeBlock); + if (!ancestor) + continue; + if (!visited.insert(ancestor)) + continue; + newlyDiscoveredOps.push_back(ancestor); + for (Value result : ancestor->getResults()) + valueWorklist.push_back(result); + } + } + return newlyDiscoveredOps; +} + +/// Capture free variables used in regions of ops in visited set. +/// Performs backward trace for each discovered free variable definition. +/// Used in Phase 4 of computeLoadChain. +void captureRegionFreeVars(llvm::SetVector &visited, + SmallVectorImpl &worklist, + Block *scopeBlock) { + SmallVector freeVarWorklist; + SmallVector initialChain(visited.begin(), visited.end()); + + // First pass: collect free variables from all region-containing ops + for (Operation *chainOp : initialChain) { + if (chainOp->getNumRegions() == 0) + continue; + chainOp->walk([&](Operation *nested) { + if (nested == chainOp) + return; + for (Value operand : nested->getOperands()) { + auto *defOp = operand.getDefiningOp(); + if (!defOp || defOp->getBlock() != scopeBlock) + continue; + if (visited.insert(defOp)) + freeVarWorklist.push_back(defOp); + } + }); + } + + // Second pass: backward trace and recursively capture nested region free vars + worklist.assign(freeVarWorklist.begin(), freeVarWorklist.end()); + while (!worklist.empty()) { + Operation *currentOp = worklist.pop_back_val(); + backwardTrace(currentOp, scopeBlock, visited, worklist); + if (currentOp->getNumRegions() == 0) + continue; + currentOp->walk([&](Operation *nested) { + if (nested == currentOp) + return; + for (Value operand : nested->getOperands()) { + auto *defOp = operand.getDefiningOp(); + if (!defOp || defOp->getBlock() != scopeBlock) + continue; + if (visited.insert(defOp)) { + freeVarWorklist.push_back(defOp); + worklist.push_back(defOp); + } + } + }); + } +} + +/// Compute the full dependency chain of `markedOp` within `scopeBlock`. +/// Phase 1 Backward SSA trace from markedOp. +/// Phase 2 Forward trace from every memref::AllocOp found in Phase 1. +/// Phase 3 Backward SSA trace from newly discovered ops. +/// Phase 4 Capture free variables of chain-op regions. +SmallVector computeLoadChain(Operation *markedOp, + Block *scopeBlock) { + llvm::SetVector visited; + SmallVector worklist; + + // Phase 1: backward SSA trace from markedOp + visited.insert(markedOp); + worklist.push_back(markedOp); + backwardTraceFromWorklist(visited, worklist, scopeBlock); + + // Phase 2: forward trace from every memref.alloc in the chain + SmallVector newlyDiscoveredOps = + forwardTraceFromAllocs(visited, scopeBlock); + + // Phase 3: backward SSA trace from newly discovered ops + worklist.assign(newlyDiscoveredOps.begin(), newlyDiscoveredOps.end()); + backwardTraceFromWorklist(visited, worklist, scopeBlock); + + // Phase 4: capture free variables of chain-op regions + captureRegionFreeVars(visited, worklist, scopeBlock); + + // Sort by block position (topological order) + SmallVector sorted(visited.begin(), visited.end()); + llvm::sort(sorted, [](Operation *leftOp, Operation *rightOp) { + return leftOp->isBeforeInBlock(rightOp); + }); + return sorted; +} + +/// Find the memref::AllocOp in the chain that backs the marked op. +memref::AllocOp findBackingAlloc(Operation *markedOp, + ArrayRef chain) { + for (Operation *chainOp : chain) { + if (auto alloc = dyn_cast(chainOp)) { + if (markedOp->getNumOperands() > 0 && + markedOp->getOperand(0).getDefiningOp() == alloc) + return alloc; + } + } + for (Operation *chainOp : chain) { + if (auto alloc = dyn_cast(chainOp)) + return alloc; + } + return nullptr; +} + +// ============================================================================ +// Step 1a: Generic marked-op collection (scan all IR) +// ============================================================================ + +/// Scan the entire module IR for ops carrying the `kBufferableMarker` +/// attribute and compute the dependency chain for each. +SmallVector collectMarkedOps(ModuleOp module) { + SmallVector results; + + module->walk([&](Operation *op) { + if (!op->hasAttr(kBufferableMarker)) + return; + if (op->getNumResults() == 0) + return; + + Block *scope = op->getBlock(); + auto chain = computeLoadChain(op, scope); + + auto alloc = findBackingAlloc(op, chain); + if (!alloc) + return; + + results.push_back({op, std::move(chain), alloc}); + }); + + return results; +} + +// ============================================================================ +// Step 1b: Group by enclosing loop construct +// ============================================================================ + +/// Try to determine the loop trip count at compile time. +std::optional getConstantTripCount(scf::ForOp forOp) { + auto lowerBoundConst = + forOp.getLowerBound().getDefiningOp(); + auto upperBoundConst = + forOp.getUpperBound().getDefiningOp(); + auto stepConst = forOp.getStep().getDefiningOp(); + if (!lowerBoundConst || !upperBoundConst || !stepConst) + return std::nullopt; + + auto lowerBoundAttr = dyn_cast(lowerBoundConst.getValue()); + auto upperBoundAttr = dyn_cast(upperBoundConst.getValue()); + auto stepAttr = dyn_cast(stepConst.getValue()); + if (!lowerBoundAttr || !upperBoundAttr || !stepAttr) + return std::nullopt; + + int64_t lowerBoundValue = lowerBoundAttr.getInt(); + int64_t upperBoundValue = upperBoundAttr.getInt(); + int64_t stepValue = stepAttr.getInt(); + if (stepValue < 1) + return std::nullopt; + + // stepValue >= 1 guaranteed by check above + int64_t tripCount = + (upperBoundValue - lowerBoundValue + stepValue - 1) / stepValue; + return tripCount > 0 ? std::optional(tripCount) : std::nullopt; +} + +/// Group marked loads by their nearest enclosing scf::ForOp. +/// Each marked load becomes its own independent LoadGroup so that every GM +/// load has its own producer flags, counters, and slot targets. Shared +/// address-computation ops are handled correctly by the consumer skip analysis +/// without forcing multiple loads into the same multi-buffer state machine. +void groupByEnclosingForOp(SmallVector &markedOps, + SmallVectorImpl &contexts) { + // Bucket loads by the nearest enclosing for loop. + llvm::DenseMap> forBuckets; + SmallVector forOpsInOrder; + for (size_t markedIndex = 0; markedIndex < markedOps.size(); ++markedIndex) { + auto forOp = markedOps[markedIndex].markedOp->getParentOfType(); + if (!forOp) + continue; + auto [bucketIt, inserted] = forBuckets.try_emplace(forOp.getOperation()); + if (inserted) + forOpsInOrder.push_back(forOp); + bucketIt->second.push_back(markedIndex); + } + + // Each marked load becomes its own independent LoadGroup. + for (auto forOp : forOpsInOrder) { + auto &indices = forBuckets[forOp.getOperation()]; + + ForBufferCtx context; + context.forOp = forOp; + for (size_t markedIndex : indices) { + LoadGroup group; + group.loads.push_back(std::move(markedOps[markedIndex])); + if (!group.loads.empty()) + group.mergedChain = group.loads[0].chain; + context.groups.push_back(std::move(group)); + } + contexts.push_back(std::move(context)); + } +} + +// ============================================================================ +// Consumer skip analysis +// ============================================================================ + +/// Collect all chain ops and marked ops from the given loads. +void collectChainAndMarkedOps(const SmallVector &loads, + llvm::DenseSet &allChainOps, + llvm::DenseSet &markedOps) { + for (const auto &markedLoad : loads) { + for (Operation *chainOp : markedLoad.chain) + allChainOps.insert(chainOp); + markedOps.insert(markedLoad.markedOp); + } +} + +/// Check if a result value is used by yield or by non-chain ops in forBody. +bool isResultUsedExternally(Value result, Operation *yieldOp, + const llvm::DenseSet &allChainOps, + Block *forBody) { + // Check if result is yielded + for (Value yieldedValue : yieldOp->getOperands()) { + if (yieldedValue == result) + return true; + } + // Check if result is used by non-chain ops + for (Operation *user : result.getUsers()) { + Operation *ancestor = getAncestorInBlock(user, forBody); + if (ancestor && !allChainOps.contains(ancestor)) + return true; + } + return false; +} + +/// Collect all operands from an op, including those from nested regions. +SmallVector +collectOperandsIncludingRegions(Operation *op) { + SmallVector refs; + for (Value operand : op->getOperands()) + refs.push_back(operand); + if (op->getNumRegions() > 0) { + op->walk([&](Operation *nested) { + if (nested == op) + return; + for (Value operand : nested->getOperands()) + refs.push_back(operand); + }); + } + return refs; +} + +/// Return the set of ops that must be skipped (not cloned) in the consumer +llvm::DenseSet +computeSkipInConsumer(const SmallVector &loads, Block *forBody) { + llvm::DenseSet allChainOps; + llvm::DenseSet markedOps; + collectChainAndMarkedOps(loads, allChainOps, markedOps); + + // Phase 1: chain ops whose results are directly used by non-chain ops or + // yielded. These must be cloned so non-chain consumers have valid SSA defs. + llvm::DenseSet needed; + Operation *yieldOp = forBody->getTerminator(); + for (Operation *chainOp : allChainOps) { + if (markedOps.contains(chainOp)) + continue; + bool isNeeded = llvm::any_of(chainOp->getResults(), [&](Value result) { + return isResultUsedExternally(result, yieldOp, allChainOps, forBody); + }); + if (isNeeded) + needed.insert(chainOp); + } + + // Phase 2: backward propagation through operands and region free variables. + SmallVector worklist(needed.begin(), needed.end()); + while (!worklist.empty()) { + Operation *currentOp = worklist.pop_back_val(); + SmallVector operands = + collectOperandsIncludingRegions(currentOp); + for (Value operand : operands) { + auto *defOp = operand.getDefiningOp(); + if (!defOp || defOp->getBlock() != forBody) + continue; + if (!allChainOps.contains(defOp) || markedOps.contains(defOp)) + continue; + if (needed.insert(defOp).second) + worklist.push_back(defOp); + } + } + + // Compute skip set as difference: allChainOps - needed + llvm::DenseSet skip; + for (Operation *chainOp : allChainOps) { + if (!needed.contains(chainOp)) + skip.insert(chainOp); + } + return skip; +} + +} // namespace gmload diff --git a/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/LoopTransform.cpp b/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/LoopTransform.cpp new file mode 100644 index 0000000000..6a4cc0c493 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/LoopTransform.cpp @@ -0,0 +1,805 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadInternal.h" + +namespace gmload { + +// ============================================================================ +// Dead iter_arg elimination +// ============================================================================ + +/// Return true if the iter_arg at `iterArgIdx` in `forOp` is dead +/// (a) its loop result has no uses outside the loop, AND +/// (b) its block argument is only used in a chain of side-effect-free ops +/// whose results ultimately feed only the loop yield. +bool isDeadIterArg(scf::ForOp forOp, unsigned iterArgIdx) { + if (!forOp.getResult(iterArgIdx).use_empty()) { + return false; + } + + Value blockArg = + forOp.getBody()->getArgument(iterArgIdx + kForBodyIterArgOffset); + Block *forBody = forOp.getBody(); + Operation *yieldOp = forBody->getTerminator(); + + llvm::SmallSetVector frontier; + frontier.insert(blockArg); + + while (!frontier.empty()) { + Value currentValue = frontier.pop_back_val(); + for (OpOperand &use : currentValue.getUses()) { + Operation *user = use.getOwner(); + Operation *ancestor = getAncestorInBlock(user, forBody); + if (!ancestor) { + continue; + } + if (ancestor == yieldOp) { + continue; + } + if (!mlir::isMemoryEffectFree(ancestor)) { + return false; + } + for (Value result : ancestor->getResults()) { + frontier.insert(result); + } + } + } + return true; +} + +llvm::DenseSet collectDeadIterArgs(scf::ForOp forOp, + unsigned candidateCount) { + llvm::DenseSet deadSet; + for (unsigned iterArgIdx = 0; iterArgIdx < candidateCount; ++iterArgIdx) { + if (isDeadIterArg(forOp, iterArgIdx)) { + deadSet.insert(iterArgIdx); + } + } + return deadSet; +} + +SmallVector +collectLiveIterArgInits(scf::ForOp forOp, + const llvm::DenseSet &deadSet) { + unsigned numIter = forOp.getNumResults(); + SmallVector liveInits; + liveInits.reserve(numIter - deadSet.size()); + for (unsigned iterArgIdx = 0; iterArgIdx < numIter; ++iterArgIdx) { + if (!deadSet.contains(iterArgIdx)) { + liveInits.push_back(forOp.getInitArgs()[iterArgIdx]); + } + } + return liveInits; +} + +scf::ForOp createForWithLiveIterArgs(OpBuilder &builder, scf::ForOp forOp, + ArrayRef liveInits) { + builder.setInsertionPoint(forOp); + auto newFor = builder.create( + forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(), + forOp.getStep(), liveInits); + newFor->setAttrs(forOp->getAttrs()); + return newFor; +} + +IRMapping buildPrunedIterArgMapping(scf::ForOp oldForOp, scf::ForOp newForOp, + const llvm::DenseSet &deadSet) { + unsigned numIter = oldForOp.getNumResults(); + Block *oldBody = oldForOp.getBody(); + Block *newBody = newForOp.getBody(); + + IRMapping mapping; + mapping.map(oldBody->getArgument(kForInductionVarArgIndex), + newForOp.getInductionVar()); + unsigned liveArgIdx = 0; + for (unsigned iterArgIdx = 0; iterArgIdx < numIter; ++iterArgIdx) { + Value oldArg = oldBody->getArgument(iterArgIdx + kForBodyIterArgOffset); + if (deadSet.contains(iterArgIdx)) { + mapping.map(oldArg, oldForOp.getInitArgs()[iterArgIdx]); + continue; + } + mapping.map(oldArg, + newBody->getArgument(liveArgIdx++ + kForBodyIterArgOffset)); + } + return mapping; +} + +void clonePrunedForBody(OpBuilder &builder, scf::ForOp oldForOp, + scf::ForOp newForOp, IRMapping &mapping) { + builder.setInsertionPointToStart(newForOp.getBody()); + for (auto &op : oldForOp.getBody()->without_terminator()) { + builder.clone(op, mapping); + } +} + +SmallVector +collectLiveYieldOperands(scf::ForOp forOp, + const llvm::DenseSet &deadSet, + IRMapping &mapping) { + unsigned numIter = forOp.getNumResults(); + auto oldYield = cast(forOp.getBody()->getTerminator()); + SmallVector liveYield; + liveYield.reserve(numIter - deadSet.size()); + for (unsigned iterArgIdx = 0; iterArgIdx < numIter; ++iterArgIdx) { + if (!deadSet.contains(iterArgIdx)) { + liveYield.push_back( + mapping.lookupOrDefault(oldYield.getOperand(iterArgIdx))); + } + } + return liveYield; +} + +void replaceLiveForResults(scf::ForOp oldForOp, scf::ForOp newForOp, + const llvm::DenseSet &deadSet) { + unsigned numIter = oldForOp.getNumResults(); + unsigned liveArgIdx = 0; + for (unsigned iterArgIdx = 0; iterArgIdx < numIter; ++iterArgIdx) { + if (!deadSet.contains(iterArgIdx)) { + oldForOp.getResult(iterArgIdx) + .replaceAllUsesWith(newForOp.getResult(liveArgIdx++)); + } + } +} + +/// Rebuild `forOp` dropping all iter_args that satisfy `isDeadIterArg`. +/// Only the first `candidateCount` iter_args are considered for pruning. +/// Returns the pruned ForOp, or `forOp` itself when nothing was pruned. +FailureOr pruneDeadIterArgs(OpBuilder &builder, scf::ForOp forOp, + unsigned candidateCount) { + unsigned numIter = forOp.getNumResults(); + if (candidateCount > numIter) { + forOp.emitError() << "[" << DEBUG_TYPE << "] candidate iter_arg count " + << candidateCount << " exceeds total iter_arg count " + << numIter; + return failure(); + } + + llvm::DenseSet deadSet = collectDeadIterArgs(forOp, candidateCount); + if (deadSet.empty()) { + return forOp; + } + + SmallVector liveInits = collectLiveIterArgInits(forOp, deadSet); + scf::ForOp newFor = createForWithLiveIterArgs(builder, forOp, liveInits); + IRMapping mapping = buildPrunedIterArgMapping(forOp, newFor, deadSet); + + clonePrunedForBody(builder, forOp, newFor, mapping); + SmallVector liveYield = + collectLiveYieldOperands(forOp, deadSet, mapping); + builder.create(forOp.getBody()->getTerminator()->getLoc(), + liveYield); + + replaceLiveForResults(forOp, newFor, deadSet); + forOp.erase(); + return newFor; +} + +/// Erase side-effect-free ops in `body` whose results are entirely unused. +/// Propagates: erasing an op may expose its operand-defining ops as newly dead. +/// Only considers ops directly in `body` (not nested regions). +void eraseDeadBodyOps(Block *body) { + SmallVector worklist; + llvm::SmallDenseSet queued; + auto enqueueIfDead = [&body, &queued, &worklist](Operation *candidateOp) { + if (!candidateOp || candidateOp->getBlock() != body || + candidateOp == body->getTerminator()) { + return; + } + if (!queued.insert(candidateOp).second) { + return; + } + worklist.push_back(candidateOp); + }; + + for (auto &op : *body) { + if (&op == body->getTerminator()) { + continue; + } + if (mlir::isMemoryEffectFree(&op) && + llvm::all_of(op.getResults(), + [](Value result) { return result.use_empty(); })) { + enqueueIfDead(&op); + } + } + + while (!worklist.empty()) { + Operation *currentOp = worklist.pop_back_val(); + queued.erase(currentOp); + if (!mlir::isMemoryEffectFree(currentOp)) { + continue; + } + if (!llvm::all_of(currentOp->getResults(), + [](Value result) { return result.use_empty(); })) { + continue; + } + + SmallVector operandDefs; + for (Value operand : currentOp->getOperands()) { + auto *defOp = operand.getDefiningOp(); + if (defOp && defOp->getBlock() == body) { + operandDefs.push_back(defOp); + } + } + currentOp->erase(); + + for (Operation *defOp : operandDefs) { + if (mlir::isMemoryEffectFree(defOp) && + llvm::all_of(defOp->getResults(), + [](Value result) { return result.use_empty(); })) { + enqueueIfDead(defOp); + } + } + } +} + +// ============================================================================ +// Producer address projection helpers +// ============================================================================ + +/// Return true if `value` is defined outside `forOp` (i.e. loop-invariant). +bool isLoopInvariant(Value value, scf::ForOp forOp) { + if (auto blockArg = dyn_cast(value)) { + return blockArg.getOwner() != forOp.getBody(); + } + Operation *defOp = value.getDefiningOp(); + Operation *parent = defOp->getParentOp(); + while (parent) { + if (parent == forOp.getOperation()) { + return false; + } + parent = parent->getParentOp(); + } + return true; +} + +/// If the yield value for iter-arg `iterArg` is of the form +/// arith.addi(iterArg, delta) where `delta` is loop-invariant, +/// return true and set `delta`. +bool getLinearIterArgDelta(Value iterArg, Value yieldVal, scf::ForOp forOp, + Value &delta) { + auto addOp = yieldVal.getDefiningOp(); + if (!addOp) { + return false; + } + Value candidate; + if (addOp.getLhs() == iterArg) { + candidate = addOp.getRhs(); + } else if (addOp.getRhs() == iterArg) { + candidate = addOp.getLhs(); + } else { + return false; + } + if (!isLoopInvariant(candidate, forOp)) { + return false; + } + delta = candidate; + return true; +} + +Value castIndexTo(OpBuilder &builder, Location loc, Value value, + Type targetType) { + if (value.getType() == targetType) { + return value; + } + return builder.create(loc, targetType, value); +} + +Value castToIndex(OpBuilder &builder, Location loc, Value value) { + if (value.getType().isIndex()) { + return value; + } + return builder.create(loc, builder.getIndexType(), value); +} + +// ============================================================================ +// Block-id tagging helper +// ============================================================================ + +/// Set kBlockIdAttr on `value`'s defining op if both the op and the attribute +/// exist. +void tagWithBlockId(Value value, IntegerAttr blockId) { + if (blockId) { + if (auto *defOp = value.getDefiningOp()) { + defOp->setAttr(kBlockIdAttr, blockId); + } + } +} + +// ============================================================================ +// Core transformation helpers +// ============================================================================ + +/// Allocate one memref slot per (depth × load) pair before the loop. +void allocateBufferSlots(OpBuilder &builder, Location loc, scf::ForOp forOp, + LoadGroup &group) { + builder.setInsertionPoint(forOp); + int depth = group.depth; + int numLoads = static_cast(group.loads.size()); + group.bufSlots.resize(depth); + for (int slotIdx = 0; slotIdx < depth; ++slotIdx) { + group.bufSlots[slotIdx].resize(numLoads); + for (int loadIdx = 0; loadIdx < numLoads; ++loadIdx) { + auto newAlloc = builder.create( + loc, group.loads[loadIdx].allocOp.getType()); + newAlloc->setAttrs( + group.loads[loadIdx].allocOp.getOperation()->getAttrs()); + group.bufSlots[slotIdx][loadIdx] = newAlloc; + } + } +} + +/// Build a new scf.for whose iter_args cover all groups: +/// [original..., (flags_g[depth], prodCounter_g, consCounter_g) for each g]. +/// Returns metadata including per-group iter_arg handles. +ExtendedForInfo buildExtendedFor(OpBuilder &builder, Location loc, + scf::ForOp forOp, ArrayRef groups, + ConstantCache &cache) { + builder.setInsertionPoint(forOp); + int numOrig = static_cast(forOp.getInitArgs().size()); + int depth = groups.empty() ? 0 : groups[0].depth; + + Value falseVal = cache.getFalse(builder, loc); + Value zeroIndex = cache.getIndex(builder, loc, 0); + + SmallVector inits; + inits.reserve(numOrig + static_cast(groups.size()) * + (depth + kLoadGroupCounterIterArgCount)); + for (Value initArg : forOp.getInitArgs()) { + inits.push_back(initArg); + } + for (auto &group : groups) { + for (int slotIdx = 0; slotIdx < group.depth; ++slotIdx) { + inits.push_back(falseVal); + } + inits.push_back(zeroIndex); // prodCounter + inits.push_back(zeroIndex); // consCounter + } + + auto newFor = + builder.create(loc, forOp.getLowerBound(), + forOp.getUpperBound(), forOp.getStep(), inits); + newFor->setAttrs(forOp->getAttrs()); + Block *oldBody = forOp.getBody(); + Block *newBody = newFor.getBody(); + + IRMapping mapping; + mapping.map(oldBody->getArgument(kForInductionVarArgIndex), + newBody->getArgument(kForInductionVarArgIndex)); + for (int initArgIdx = 0; initArgIdx < numOrig; ++initArgIdx) { + mapping.map(oldBody->getArgument(initArgIdx + kForBodyIterArgOffset), + newBody->getArgument(initArgIdx + kForBodyIterArgOffset)); + } + + if (!newBody->empty()) { + newBody->getTerminator()->erase(); + } + builder.setInsertionPointToEnd(newBody); + + SmallVector groupArgs; + int argOffset = numOrig; + for (auto &group : groups) { + GroupIterArgs groupIterArgs; + for (int slotIdx = 0; slotIdx < group.depth; ++slotIdx) { + groupIterArgs.flagArgs.push_back( + newBody->getArgument(kForBodyIterArgOffset + argOffset + slotIdx)); + } + groupIterArgs.prodCounter = + newBody->getArgument(kForBodyIterArgOffset + argOffset + group.depth + + kLoadGroupProducerCounterOffset); + groupIterArgs.consCounter = + newBody->getArgument(kForBodyIterArgOffset + argOffset + group.depth + + kLoadGroupConsumerCounterOffset); + argOffset += group.depth + kLoadGroupCounterIterArgCount; + groupArgs.push_back(std::move(groupIterArgs)); + } + + return {newFor, oldBody, newBody, std::move(mapping), std::move(groupArgs), + falseVal, numOrig, depth}; +} + +// ============================================================================ +// Clone availability checks +// ============================================================================ + +/// Return true if `value` can be used while cloning into the new body with the +/// current mapping. Values defined outside the old body still dominate; values +/// defined in the old body must already have been cloned and mapped. +bool isAvailableForClone(Value value, Block *oldBody, + const IRMapping &mapping) { + if (auto blockArg = dyn_cast(value)) { + return blockArg.getOwner() != oldBody || + static_cast(mapping.lookupOrNull(value)); + } + + Operation *defOp = value.getDefiningOp(); + if (!defOp || defOp->getBlock() != oldBody) { + return true; + } + return static_cast(mapping.lookupOrNull(value)); +} + +/// Check all direct and region-captured operands before cloning an op out of +/// the old loop body. +bool areOperandsAvailableForClone(Operation *op, Block *oldBody, + const IRMapping &mapping) { + for (Value operand : collectOperandsIncludingRegions(op)) { + if (!isAvailableForClone(operand, oldBody, mapping)) { + return false; + } + } + return true; +} + +// ============================================================================ +// Body-run utilities +// ============================================================================ + +std::optional getBlockId(Operation *op) { + if (auto bid = op->getAttrOfType(kBlockIdAttr)) { + return static_cast(bid.getInt()); + } + return std::nullopt; +} + +bool sameRunBlockId(const BodyRun &run, std::optional bid) { + if (run.hasBlockId != bid.has_value()) { + return false; + } + return !run.hasBlockId || run.blockId == *bid; +} + +SmallVector collectBodyRuns(Block *body) { + SmallVector runs; + for (auto &op : body->without_terminator()) { + std::optional bid = getBlockId(&op); + if (runs.empty() || !sameRunBlockId(runs.back(), bid)) { + BodyRun run; + if (bid) { + run.hasBlockId = true; + run.blockId = *bid; + } + runs.push_back(std::move(run)); + } + runs.back().ops.push_back(&op); + } + return runs; +} + +int findRunContainingOp(ArrayRef runs, Operation *target) { + for (size_t runIdx = 0; runIdx < runs.size(); ++runIdx) { + for (Operation *op : runs[runIdx].ops) { + if (op == target) { + return static_cast(runIdx); + } + } + } + return -1; +} + +int findFirstRunWithBlockId(ArrayRef runs, int32_t blockId) { + for (size_t runIdx = 0; runIdx < runs.size(); ++runIdx) { + if (runs[runIdx].hasBlockId && runs[runIdx].blockId == blockId) { + return static_cast(runIdx); + } + } + return -1; +} + +void logRepeatedTopLevelBlockRuns(Block *body, llvm::StringRef stage) { + llvm::DenseSet closed; + bool hasCurrent = false; + int32_t current = 0; + + for (auto &op : body->without_terminator()) { + std::optional bid = getBlockId(&op); + if (!bid) { + if (hasCurrent) { + closed.insert(current); + } + hasCurrent = false; + continue; + } + + if (hasCurrent && current == *bid) { + continue; + } + + if (hasCurrent) { + closed.insert(current); + } + + if (closed.contains(*bid)) { + LOG_DEBUG("non-consecutive top-level ssbuffer.block_id run after " + << stage << ": block_id=" << *bid << "\n"); + } + current = *bid; + hasCurrent = true; + } +} + +// ============================================================================ +// Core transformation +// ============================================================================ + +struct LoopInvariantValues { + Value lowerBound; + Value step; + Value tripCount; +}; + +struct GroupInvariantValues { + SmallVector trueFlags; + SmallVector indexOnes; + SmallVector> slotConsts; + SmallVector depthConsts; +}; + +void allocateBufferSlotsForGroups(OpBuilder &builder, Location loc, + scf::ForOp forOp, + SmallVectorImpl &groups) { + for (auto &group : groups) { + allocateBufferSlots(builder, loc, forOp, group); + } +} + +LoopInvariantValues emitLoopInvariantValues(OpBuilder &builder, Location loc, + scf::ForOp forOp) { + IntegerAttr forBlockId = forOp->getAttrOfType(kBlockIdAttr); + Value lowerBound = forOp.getLowerBound(); + Value step = forOp.getStep(); + Value upperBound = forOp.getUpperBound(); + + Value lowerBoundIndex = castToIndex(builder, loc, lowerBound); + if (lowerBoundIndex != lowerBound) { + tagWithBlockId(lowerBoundIndex, forBlockId); + } + Value stepIndex = castToIndex(builder, loc, step); + if (stepIndex != step) { + tagWithBlockId(stepIndex, forBlockId); + } + Value upperBoundIndex = castToIndex(builder, loc, upperBound); + if (upperBoundIndex != upperBound) { + tagWithBlockId(upperBoundIndex, forBlockId); + } + + Value range = + builder.create(loc, upperBoundIndex, lowerBoundIndex); + tagWithBlockId(range, forBlockId); + Value tripCount = builder.create(loc, range, stepIndex); + tagWithBlockId(tripCount, forBlockId); + return {lowerBound, step, tripCount}; +} + +GroupInvariantValues emitGroupInvariantValues(OpBuilder &builder, Location loc, + ArrayRef groups) { + int numGroups = static_cast(groups.size()); + GroupInvariantValues values; + values.trueFlags.resize(numGroups); + values.indexOnes.resize(numGroups); + values.slotConsts.resize(numGroups); + values.depthConsts.resize(numGroups); + + for (int groupIdx = 0; groupIdx < numGroups; ++groupIdx) { + const auto &group = groups[groupIdx]; + IntegerAttr blockId = + group.loads.empty() + ? IntegerAttr{} + : group.loads[0].markedOp->getAttrOfType(kBlockIdAttr); + int groupDepth = group.depth; + values.trueFlags[groupIdx] = + builder.create(loc, builder.getBoolAttr(true)); + tagWithBlockId(values.trueFlags[groupIdx], blockId); + values.indexOnes[groupIdx] = builder.create(loc, 1); + tagWithBlockId(values.indexOnes[groupIdx], blockId); + + values.slotConsts[groupIdx].resize(groupDepth); + for (int slot = 0; slot < groupDepth; ++slot) { + values.slotConsts[groupIdx][slot] = + builder.create(loc, slot); + tagWithBlockId(values.slotConsts[groupIdx][slot], blockId); + } + + values.depthConsts[groupIdx] = + builder.create(loc, groupDepth); + tagWithBlockId(values.depthConsts[groupIdx], blockId); + } + + return values; +} + +FailureOr> +collectLinearIterArgDeltas(const ExtendedForInfo &info, scf::ForOp forOp) { + if (info.numOrig < 0) { + forOp.emitError() << "[" << DEBUG_TYPE + << "] original iter_arg count is negative: " + << info.numOrig; + return failure(); + } + auto numOrig = static_cast(info.numOrig); + if (info.oldBody->getNumArguments() < numOrig + kForBodyIterArgOffset) { + forOp.emitError() << "[" << DEBUG_TYPE << "] old loop body has " + << info.oldBody->getNumArguments() + << " arguments, expected at least " + << (numOrig + kForBodyIterArgOffset); + return failure(); + } + + auto oldYieldOp = cast(info.oldBody->getTerminator()); + if (oldYieldOp->getNumOperands() < numOrig) { + forOp.emitError() << "[" << DEBUG_TYPE << "] old loop yield has " + << oldYieldOp->getNumOperands() + << " operands, expected at least " << numOrig; + return failure(); + } + + SmallVector iterArgDeltas(info.numOrig, Value{}); + for (unsigned iterArgIdx = 0; iterArgIdx < numOrig; ++iterArgIdx) { + Value delta; + if (getLinearIterArgDelta( + info.oldBody->getArgument(iterArgIdx + kForBodyIterArgOffset), + oldYieldOp->getOperand(iterArgIdx), forOp, delta)) { + iterArgDeltas[iterArgIdx] = delta; + } + } + return iterArgDeltas; +} + +LogicalResult replaceOriginalForResults(scf::ForOp oldForOp, + scf::ForOp newForOp, int numOrig) { + if (numOrig < 0) { + oldForOp.emitError() << "[" << DEBUG_TYPE + << "] original result count is negative: " << numOrig; + return failure(); + } + + auto numOrigResults = static_cast(numOrig); + if (oldForOp.getNumResults() < numOrigResults || + newForOp.getNumResults() < numOrigResults) { + oldForOp.emitError() << "[" << DEBUG_TYPE << "] cannot replace " + << numOrigResults << " loop results; old loop has " + << oldForOp.getNumResults() << ", new loop has " + << newForOp.getNumResults(); + return failure(); + } + + for (unsigned resultIdx = 0; resultIdx < numOrigResults; ++resultIdx) { + oldForOp.getResult(resultIdx).replaceAllUsesWith( + newForOp.getResult(resultIdx)); + } + return success(); +} + +LogicalResult finalizeMultiBufferLoop(OpBuilder &builder, scf::ForOp newForOp, + int numOrig) { + if (numOrig < 0) { + newForOp.emitError() << "[" << DEBUG_TYPE + << "] original iter_arg count is negative: " + << numOrig; + return failure(); + } + + FailureOr finalFor = + pruneDeadIterArgs(builder, newForOp, static_cast(numOrig)); + if (failed(finalFor)) { + return failure(); + } + + eraseDeadBodyOps(finalFor->getBody()); + + // Debug-only diagnostic: keep this as a log instead of asserting so the + // pass can still dump the transformed IR for investigation. + logRepeatedTopLevelBlockRuns(finalFor->getBody(), kGMLoadMultiBufferStage); + return success(); +} + +/// Apply GM-load multi-buffer rewriting to one scf.for. +/// allCtxForOps carries the set of all forOps being transformed so that inner +/// forOps (already handled) are not cloned again into the consumer body. +LogicalResult +applyMultiBufferToForLoop(ForBufferCtx &context, + const llvm::DenseSet &allCtxForOps) { + scf::ForOp forOp = context.forOp; + Location loc = forOp.getLoc(); + OpBuilder builder(forOp); + + allocateBufferSlotsForGroups(builder, loc, forOp, context.groups); + + // Scan parent block once so buildExtendedFor can reuse existing constants. + ConstantCache cache; + cache.scan(forOp->getBlock()); + + builder.setInsertionPoint(forOp); + LoopInvariantValues loopValues = emitLoopInvariantValues(builder, loc, forOp); + GroupInvariantValues groupValues = + emitGroupInvariantValues(builder, loc, context.groups); + + auto info = buildExtendedFor(builder, loc, forOp, context.groups, cache); + FailureOr> iterArgDeltas = + collectLinearIterArgDeltas(info, forOp); + if (failed(iterArgDeltas)) { + return failure(); + } + + LogicalResult emitBodyStatus = emitMultiBufferLoopBody( + builder, loc, context.groups, info, groupValues.indexOnes, + groupValues.slotConsts, groupValues.depthConsts, allCtxForOps, + loopValues.tripCount, loopValues.lowerBound, loopValues.step, + groupValues.trueFlags, *iterArgDeltas, forOp); + if (failed(emitBodyStatus)) { + return failure(); + } + + if (failed(replaceOriginalForResults(forOp, info.newForOp, info.numOrig))) { + return failure(); + } + if (failed(finalizeMultiBufferLoop(builder, info.newForOp, info.numOrig))) { + return failure(); + } + return success(); +} + +// ============================================================================ +// Cleanup utilities +// ============================================================================ + +bool isAncestorBlock(Block *ancestor, Block *descendant) { + for (Block *currentBlock = descendant; currentBlock;) { + if (currentBlock == ancestor) { + return true; + } + auto *parentOp = currentBlock->getParentOp(); + if (!parentOp) { + break; + } + currentBlock = parentOp->getBlock(); + } + return false; +} + +void deduplicateConstants(ModuleOp module) { + llvm::DenseMap canonical; + SmallVector toErase; + + module->walk([&](arith::ConstantOp cst) { + if (cst->hasAttr(kBlockIdAttr)) { + return mlir::WalkResult::advance(); + } + auto [canonicalIt, inserted] = + canonical.try_emplace(cst.getValue(), cst.getResult()); + if (!inserted && canonicalIt->second != cst.getResult()) { + Block *canonBlock = canonicalIt->second.getParentBlock(); + Block *thisBlock = cst->getBlock(); + if (canonBlock == thisBlock || isAncestorBlock(canonBlock, thisBlock)) { + cst.getResult().replaceAllUsesWith(canonicalIt->second); + toErase.push_back(cst); + } + } + return mlir::WalkResult::advance(); + }); + + for (Operation *op : llvm::reverse(toErase)) { + op->erase(); + } +} + +} // namespace gmload diff --git a/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/MultiBufferLoopBodyEmission.cpp b/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/MultiBufferLoopBodyEmission.cpp new file mode 100644 index 0000000000..90e8a56d69 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromCompute/MultiBufferLoopBodyEmission.cpp @@ -0,0 +1,900 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadInternal.h" + +namespace gmload { + +// ============================================================================ +// Producer emission helpers +// ============================================================================ + +static bool canSkipBecauseAllResultsMapped(Operation *op, + const IRMapping &mapping) { + if (op->getNumResults() == 0) { + return false; + } + return llvm::all_of(op->getResults(), [&](Value result) { + return static_cast(mapping.lookupOrNull(result)); + }); +} + +static LogicalResult +materializeProducerOp(OpBuilder &builder, Operation *op, Block *oldBody, + const llvm::DenseSet &producerSourceOps, + IRMapping &mapping, + llvm::DenseSet &materializedOps, + llvm::DenseSet &activeOps); + +static LogicalResult +materializeProducerValue(OpBuilder &builder, Value value, Block *oldBody, + const llvm::DenseSet &producerSourceOps, + IRMapping &mapping, + llvm::DenseSet &materializedOps, + llvm::DenseSet &activeOps) { + if (mapping.lookupOrNull(value)) { + return success(); + } + + if (auto blockArg = dyn_cast(value)) { + if (blockArg.getOwner() == oldBody) { + return failure(); + } + return success(); + } + + Operation *defOp = value.getDefiningOp(); + if (!defOp) { + return success(); + } + + if (producerSourceOps.contains(defOp) || defOp->getBlock() == oldBody) { + return materializeProducerOp(builder, defOp, oldBody, producerSourceOps, + mapping, materializedOps, activeOps); + } + + return success(); +} + +static LogicalResult +materializeProducerOp(OpBuilder &builder, Operation *op, Block *oldBody, + const llvm::DenseSet &producerSourceOps, + IRMapping &mapping, + llvm::DenseSet &materializedOps, + llvm::DenseSet &activeOps) { + if (materializedOps.contains(op) || + canSkipBecauseAllResultsMapped(op, mapping)) { + return success(); + } + + if (activeOps.contains(op)) { + return failure(); + } + + if (!producerSourceOps.contains(op) && op->getBlock() == oldBody) { + if (!mlir::isMemoryEffectFree(op) || op->getNumRegions() != 0) { + return failure(); + } + } + + activeOps.insert(op); + for (Value operand : collectOperandsIncludingRegions(op)) { + LogicalResult materializeStatus = + materializeProducerValue(builder, operand, oldBody, producerSourceOps, + mapping, materializedOps, activeOps); + if (failed(materializeStatus)) { + activeOps.erase(op); + return failure(); + } + } + activeOps.erase(op); + + builder.clone(*op, mapping); + materializedOps.insert(op); + return success(); +} + +static IntegerAttr getGroupBlockId(const LoadGroup &group) { + return group.loads.empty() + ? IntegerAttr{} + : group.loads[0].markedOp->getAttrOfType( + kBlockIdAttr); +} + +static Value createProducerFillCondition(OpBuilder &builder, Location loc, + Value flagArg, + Value currentProducerCounter, + Value tripCount, Value falseVal, + IntegerAttr blockId) { + Value flagEmpty = builder.create(loc, arith::CmpIPredicate::eq, + flagArg, falseVal); + tagWithBlockId(flagEmpty, blockId); + Value prodLtTripCount = builder.create( + loc, arith::CmpIPredicate::ult, currentProducerCounter, tripCount); + tagWithBlockId(prodLtTripCount, blockId); + Value cond = builder.create(loc, flagEmpty, prodLtTripCount); + tagWithBlockId(cond, blockId); + return cond; +} + +static scf::IfOp createProducerFillIf(OpBuilder &builder, Location loc, + Value cond, IntegerAttr blockId) { + auto ifOp = builder.create( + loc, TypeRange{builder.getI1Type(), builder.getIndexType()}, cond, true); + if (blockId) + ifOp->setAttr(kBlockIdAttr, blockId); + ifOp->setAttr(kLoadStoreAttr, builder.getUnitAttr()); + return ifOp; +} + +struct ProducerLoopPosition { + Type inductionVarType; + Value projectedInductionVar; + Value producerCounterInInductionType; +}; + +static ProducerLoopPosition +mapProducerLoopPosition(OpBuilder &builder, Location loc, Block *oldBody, + Value loopLowerBound, Value loopStep, + Value currentProducerCounter, IntegerAttr blockId, + IRMapping &prodMapping) { + Value oldInductionVar = oldBody->getArgument(kForInductionVarArgIndex); + Type inductionVarType = oldInductionVar.getType(); + + Value lowerBound = + castIndexTo(builder, loc, loopLowerBound, inductionVarType); + if (lowerBound != loopLowerBound) + tagWithBlockId(lowerBound, blockId); + Value step = castIndexTo(builder, loc, loopStep, inductionVarType); + if (step != loopStep) + tagWithBlockId(step, blockId); + Value producerCounterInInductionType = + castIndexTo(builder, loc, currentProducerCounter, inductionVarType); + tagWithBlockId(producerCounterInInductionType, blockId); + Value inductionOffset = + builder.create(loc, producerCounterInInductionType, step); + tagWithBlockId(inductionOffset, blockId); + Value projectedInductionVar = + builder.create(loc, lowerBound, inductionOffset); + tagWithBlockId(projectedInductionVar, blockId); + prodMapping.map(oldInductionVar, projectedInductionVar); + + return {inductionVarType, projectedInductionVar, + producerCounterInInductionType}; +} + +static LogicalResult +validateProducerIterArgInputs(scf::ForOp origForOp, Block *oldBody, + Block *newBody, ArrayRef iterArgDeltas) { + int numOrig = static_cast(origForOp.getInitArgs().size()); + if (numOrig < 0) { + origForOp.emitError() << "[" << DEBUG_TYPE + << "] original iter_arg count is negative: " + << numOrig; + return failure(); + } + + auto numOrigArgs = static_cast(numOrig); + if (iterArgDeltas.size() < numOrigArgs) { + origForOp.emitError() << "[" << DEBUG_TYPE << "] iter_arg delta count " + << iterArgDeltas.size() + << " is smaller than original iter_arg count " + << numOrigArgs; + return failure(); + } + if (oldBody->getNumArguments() < numOrigArgs + kForBodyIterArgOffset || + newBody->getNumArguments() < numOrigArgs + kForBodyIterArgOffset) { + origForOp.emitError() + << "[" << DEBUG_TYPE + << "] loop body argument count is smaller than expected " + << (numOrigArgs + kForBodyIterArgOffset); + return failure(); + } + return success(); +} + +static LogicalResult +mapProducerIterArgs(OpBuilder &builder, Location loc, Block *oldBody, + Block *newBody, scf::ForOp origForOp, Value loopLowerBound, + Value loopStep, Value currentProducerCounter, + const SmallVector &iterArgDeltas, + const ProducerLoopPosition &position, IntegerAttr blockId, + IRMapping &prodMapping) { + if (failed(validateProducerIterArgInputs(origForOp, oldBody, newBody, + iterArgDeltas))) { + return failure(); + } + + // producerCounterInInductionType is already cast to inductionVarType; cache + // casts by type to avoid redundant index_cast ops when multiple iter_args + // share the same type. + llvm::DenseMap producerCounterByType; + producerCounterByType[position.inductionVarType] = + position.producerCounterInInductionType; + + int numOrig = static_cast(origForOp.getInitArgs().size()); + for (int iterArgIdx = 0; iterArgIdx < numOrig; ++iterArgIdx) { + Value oldArg = oldBody->getArgument(iterArgIdx + kForBodyIterArgOffset); + Value delta = iterArgDeltas[iterArgIdx]; + if (!delta) { + prodMapping.map(oldArg, + newBody->getArgument(iterArgIdx + kForBodyIterArgOffset)); + continue; + } + + Type argType = oldArg.getType(); + Value initArg = origForOp.getInitArgs()[iterArgIdx]; + // If initArg == lower bound and delta == step this iter_arg tracks the + // induction variable exactly; reuse the projected IV. + if (initArg == loopLowerBound && delta == loopStep && + argType == position.inductionVarType) { + prodMapping.map(oldArg, position.projectedInductionVar); + continue; + } + + auto [counterIt, inserted] = + producerCounterByType.try_emplace(argType, Value{}); + if (inserted) { + counterIt->second = + castIndexTo(builder, loc, currentProducerCounter, argType); + tagWithBlockId(counterIt->second, blockId); + } + Value argMul = builder.create(loc, counterIt->second, delta); + tagWithBlockId(argMul, blockId); + Value argAdd = builder.create(loc, initArg, argMul); + tagWithBlockId(argAdd, blockId); + prodMapping.map(oldArg, argAdd); + } + return success(); +} + +static void mapProducerSlotBuffers(const LoadGroup &group, int slotIdx, + IRMapping &prodMapping) { + for (int loadIdx = 0; loadIdx < static_cast(group.loads.size()); + ++loadIdx) + prodMapping.map(group.loads[loadIdx].allocOp->getResult(0), + group.bufSlots[slotIdx][loadIdx]); +} + +static LogicalResult materializeProducerChain(OpBuilder &builder, + Block *oldBody, + const LoadGroup &group, + IRMapping &prodMapping) { + llvm::DenseSet skipSet; + for (const MarkedLoad &load : group.loads) { + skipSet.insert(load.allocOp); + skipSet.insert(load.markedOp); + } + + llvm::DenseSet producerSourceOps(group.mergedChain.begin(), + group.mergedChain.end()); + llvm::DenseSet materializedOps; + llvm::DenseSet activeOps; + for (Operation *op : group.mergedChain) { + if (skipSet.contains(op)) { + continue; + } + if (failed(materializeProducerOp(builder, op, oldBody, producerSourceOps, + prodMapping, materializedOps, activeOps))) + return failure(); + } + return success(); +} + +static LogicalResult emitProducerFillThenBlock( + OpBuilder &builder, Location loc, scf::IfOp ifOp, Block *oldBody, + Block *newBody, const LoadGroup &group, int slotIdx, + Value currentProducerCounter, Value loopLowerBound, Value loopStep, + Value trueFlag, Value indexOne, const SmallVector &iterArgDeltas, + scf::ForOp origForOp, IntegerAttr blockId) { + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPointToStart(ifOp.thenBlock()); + + IRMapping prodMapping; + ProducerLoopPosition position = + mapProducerLoopPosition(builder, loc, oldBody, loopLowerBound, loopStep, + currentProducerCounter, blockId, prodMapping); + LogicalResult mapIterArgStatus = mapProducerIterArgs( + builder, loc, oldBody, newBody, origForOp, loopLowerBound, loopStep, + currentProducerCounter, iterArgDeltas, position, blockId, prodMapping); + if (failed(mapIterArgStatus)) { + return failure(); + } + mapProducerSlotBuffers(group, slotIdx, prodMapping); + if (failed(materializeProducerChain(builder, oldBody, group, prodMapping))) { + return failure(); + } + + Value prodNext = + builder.create(loc, currentProducerCounter, indexOne); + tagWithBlockId(prodNext, blockId); + builder.create(loc, ValueRange{trueFlag, prodNext}); + return success(); +} + +static void emitProducerFillElseBlock(OpBuilder &builder, Location loc, + scf::IfOp ifOp, Value flagArg, + Value currentProducerCounter) { + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPointToStart(ifOp.elseBlock()); + builder.create(loc, + ValueRange{flagArg, currentProducerCounter}); +} + +/// Emit the producer scf.if for one buffer slot. Returns {flagNext, +/// prodCounterNext}. +FailureOr> emitProducerFillForSlot( + OpBuilder &builder, Location loc, Block *oldBody, Block *newBody, + const LoadGroup &group, int slotIdx, Value flagArg, + Value currentProducerCounter, Value tripCount, Value falseVal, + Value loopLowerBound, Value loopStep, Value trueFlag, Value indexOne, + const SmallVector &iterArgDeltas, scf::ForOp origForOp) { + IntegerAttr blockId = getGroupBlockId(group); + Value cond = + createProducerFillCondition(builder, loc, flagArg, currentProducerCounter, + tripCount, falseVal, blockId); + scf::IfOp ifOp = createProducerFillIf(builder, loc, cond, blockId); + LogicalResult thenStatus = emitProducerFillThenBlock( + builder, loc, ifOp, oldBody, newBody, group, slotIdx, + currentProducerCounter, loopLowerBound, loopStep, trueFlag, indexOne, + iterArgDeltas, origForOp, blockId); + if (failed(thenStatus)) { + return failure(); + } + emitProducerFillElseBlock(builder, loc, ifOp, flagArg, + currentProducerCounter); + + return std::make_pair(ifOp.getResult(0), ifOp.getResult(1)); +} + +// ============================================================================ +// Consumer emission helpers +// ============================================================================ + +/// Emit slot-selection for one load: one ToTensorOp per slot, followed by a +/// (depth-1)-level comparison/select chain that picks the slot matching +/// `target`. Sets mapping[markedOp->result] = selected value. +Value emitLoadSlotSelection(OpBuilder &builder, Location loc, + const MarkedLoad &load, ArrayRef slotBufs, + Value target, int depth, IRMapping &mapping, + ArrayRef slotConsts) { + IntegerAttr blockId = load.markedOp->getAttrOfType(kBlockIdAttr); + Type tensorTy = load.markedOp->getResult(0).getType(); + + SmallVector slotVals(depth); + for (int slotIdx = 0; slotIdx < depth; ++slotIdx) { + auto toTensor = builder.create( + loc, tensorTy, slotBufs[slotIdx], true, true); + if (blockId) + toTensor->setAttr(kBlockIdAttr, blockId); + slotVals[slotIdx] = toTensor; + } + + Value result = slotVals[depth - 1]; + for (int slotIdx = depth - 2; slotIdx >= 0; --slotIdx) { + Value eq = builder.create(loc, arith::CmpIPredicate::eq, + target, slotConsts[slotIdx]); + tagWithBlockId(eq, blockId); + result = + builder.create(loc, eq, slotVals[slotIdx], result); + tagWithBlockId(result, blockId); + } + + mapping.map(load.markedOp->getResult(0), result); + return result; +} + +/// Emit the scf.if that clears slot `slotIdx`'s flag when it is the consumed +/// slot (`target == slotIdx`). Returns the updated flag value. +Value emitSlotFlagClear(OpBuilder &builder, Location loc, Value target, + Value slotConst, Value flagNext, Value falseVal, + IntegerAttr blockId) { + Value isConsumed = builder.create( + loc, arith::CmpIPredicate::eq, target, slotConst); + tagWithBlockId(isConsumed, blockId); + + auto ifClear = builder.create(loc, TypeRange{builder.getI1Type()}, + isConsumed, true); + if (blockId) + ifClear->setAttr(kBlockIdAttr, blockId); + ifClear->setAttr(kLoadStoreAttr, builder.getUnitAttr()); + { + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPointToStart(ifClear.thenBlock()); + builder.create(loc, ValueRange{falseVal}); + } + { + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPointToStart(ifClear.elseBlock()); + builder.create(loc, ValueRange{flagNext}); + } + return ifClear.getResult(0); +} + +// ============================================================================ + +namespace { + +struct MultiBufferLoopBodyState { + SmallVector> flagFinals; + SmallVector consCounterNexts; + SmallVector prodCounterFinals; + SmallVector> flagNexts; + SmallVector targets; + SmallVector prefixEmitted; + SmallVector finalEmitted; + + explicit MultiBufferLoopBodyState(int numGroups) + : flagFinals(numGroups), consCounterNexts(numGroups), + prodCounterFinals(numGroups), flagNexts(numGroups), targets(numGroups), + prefixEmitted(numGroups, 0), finalEmitted(numGroups, 0) {} +}; + +static void +computeConsumerCloneSets(ArrayRef groups, Block *oldBody, + llvm::DenseSet &allChainOps, + llvm::DenseSet &retainedChainOpSet) { + // A retained chain op is a chain op that must still be cloned into the + // consumer body because non-chain body code uses it. Retained ops are + // emitted in their original top-level run, not in the load group that + // discovered them. + for (const LoadGroup &group : groups) { + llvm::DenseSet skipSet = + computeSkipInConsumer(group.loads, oldBody); + for (Operation *op : group.mergedChain) { + allChainOps.insert(op); + if (!skipSet.contains(op)) + retainedChainOpSet.insert(op); + } + } +} + +static bool +shouldSkipOldBodyOp(Operation *op, + const llvm::DenseSet &allCtxForOps, + const llvm::DenseSet &allChainOps, + const llvm::DenseSet &retainedChainOpSet) { + if (allCtxForOps.contains(op)) { + return true; + } + return allChainOps.contains(op) && !retainedChainOpSet.contains(op); +} + +struct ProducerFillState { + SmallVector flagNexts; + Value producerCounter; +}; + +static FailureOr +emitProducerFillsForGroup(OpBuilder &builder, Location loc, LoadGroup &group, + int groupIdx, ExtendedForInfo &info, Block *oldBody, + Value tripCount, Value loopLowerBound, Value loopStep, + const SmallVector &groupTrueFlagVals, + const SmallVector &groupIndexOneVals, + const SmallVector &iterArgDeltas, + scf::ForOp origForOp) { + int groupDepth = group.depth; + Value currentProducerCounter = info.groupArgs[groupIdx].prodCounter; + SmallVector flagNexts(groupDepth); + for (int slotIdx = 0; slotIdx < groupDepth; ++slotIdx) { + FailureOr> producerResult = emitProducerFillForSlot( + builder, loc, oldBody, info.newBody, group, slotIdx, + info.groupArgs[groupIdx].flagArgs[slotIdx], currentProducerCounter, + tripCount, info.falseVal, loopLowerBound, loopStep, + groupTrueFlagVals[groupIdx], groupIndexOneVals[groupIdx], iterArgDeltas, + origForOp); + if (failed(producerResult)) { + return failure(); + } + + auto [flagNext, prodCounterNext] = *producerResult; + flagNexts[slotIdx] = flagNext; + currentProducerCounter = prodCounterNext; + } + + return ProducerFillState{std::move(flagNexts), currentProducerCounter}; +} + +static Value emitConsumerSlotTarget(OpBuilder &builder, Location loc, + ExtendedForInfo &info, int groupIdx, + Value groupDepthConst, + IntegerAttr blockId) { + Value target = builder.create( + loc, info.groupArgs[groupIdx].consCounter, groupDepthConst); + tagWithBlockId(target, blockId); + return target; +} + +static void emitLoadSlotSelectionsForGroup( + OpBuilder &builder, Location loc, LoadGroup &group, int groupIdx, + int groupDepth, Value target, ExtendedForInfo &info, + const SmallVector> &allGroupSlotConsts) { + for (int loadIdx = 0; loadIdx < static_cast(group.loads.size()); + ++loadIdx) { + SmallVector slotBufs(groupDepth); + for (int slotIdx = 0; slotIdx < groupDepth; ++slotIdx) + slotBufs[slotIdx] = group.bufSlots[slotIdx][loadIdx]; + emitLoadSlotSelection(builder, loc, group.loads[loadIdx], slotBufs, target, + groupDepth, info.mapping, + allGroupSlotConsts[groupIdx]); + } +} + +static LogicalResult validateGroupEmissionInputs( + ArrayRef groups, ExtendedForInfo &info, int groupIdx, + ArrayRef groupTrueFlagVals, ArrayRef groupIndexOneVals, + ArrayRef allGroupDepthConsts, + ArrayRef> allGroupSlotConsts, scf::ForOp origForOp) { + if (groupIdx < 0 || static_cast(groupIdx) >= groups.size() || + static_cast(groupIdx) >= info.groupArgs.size() || + static_cast(groupIdx) >= groupTrueFlagVals.size() || + static_cast(groupIdx) >= groupIndexOneVals.size() || + static_cast(groupIdx) >= allGroupDepthConsts.size() || + static_cast(groupIdx) >= allGroupSlotConsts.size()) { + origForOp.emitError() << "[" << DEBUG_TYPE << "] load group index " + << groupIdx + << " is out of range for multi-buffer emission"; + return failure(); + } + return success(); +} + +static LogicalResult emitProducerAndSelectSlots( + OpBuilder &builder, Location loc, SmallVectorImpl &groups, + int groupIdx, ExtendedForInfo &info, Block *oldBody, Value tripCount, + Value loopLowerBound, Value loopStep, + const SmallVector &groupTrueFlagVals, + const SmallVector &groupIndexOneVals, + const SmallVector &iterArgDeltas, scf::ForOp origForOp, + const SmallVector &allGroupDepthConsts, + const SmallVector> &allGroupSlotConsts, + MultiBufferLoopBodyState &state) { + LogicalResult validateStatus = validateGroupEmissionInputs( + groups, info, groupIdx, groupTrueFlagVals, groupIndexOneVals, + allGroupDepthConsts, allGroupSlotConsts, origForOp); + if (failed(validateStatus)) { + return failure(); + } + if (static_cast(groupIdx) >= state.prefixEmitted.size()) { + origForOp.emitError() << "[" << DEBUG_TYPE << "] load group index " + << groupIdx << " exceeds multi-buffer state size " + << state.prefixEmitted.size(); + return failure(); + } + if (state.prefixEmitted[groupIdx]) { + return success(); + } + + LoadGroup &group = groups[groupIdx]; + int groupDepth = group.depth; + IntegerAttr blockId = getGroupBlockId(group); + + FailureOr fillState = emitProducerFillsForGroup( + builder, loc, group, groupIdx, info, oldBody, tripCount, loopLowerBound, + loopStep, groupTrueFlagVals, groupIndexOneVals, iterArgDeltas, origForOp); + if (failed(fillState)) { + return failure(); + } + + state.flagNexts[groupIdx] = std::move(fillState->flagNexts); + state.prodCounterFinals[groupIdx] = fillState->producerCounter; + Value target = emitConsumerSlotTarget(builder, loc, info, groupIdx, + allGroupDepthConsts[groupIdx], blockId); + state.targets[groupIdx] = target; + emitLoadSlotSelectionsForGroup(builder, loc, group, groupIdx, groupDepth, + target, info, allGroupSlotConsts); + + state.prefixEmitted[groupIdx] = 1; + return success(); +} + +static LogicalResult emitConsumerReleaseForGroup( + OpBuilder &builder, Location loc, SmallVectorImpl &groups, + int groupIdx, ExtendedForInfo &info, + const SmallVector &groupIndexOneVals, + const SmallVector> &allGroupSlotConsts, + MultiBufferLoopBodyState &state) { + if (state.finalEmitted[groupIdx]) { + return success(); + } + + LoadGroup &group = groups[groupIdx]; + int groupDepth = group.depth; + IntegerAttr blockId = getGroupBlockId(group); + + state.flagFinals[groupIdx].resize(groupDepth); + for (int slotIdx = 0; slotIdx < groupDepth; ++slotIdx) { + state.flagFinals[groupIdx][slotIdx] = emitSlotFlagClear( + builder, loc, state.targets[groupIdx], + allGroupSlotConsts[groupIdx][slotIdx], + state.flagNexts[groupIdx][slotIdx], info.falseVal, blockId); + } + + state.consCounterNexts[groupIdx] = builder.create( + loc, info.groupArgs[groupIdx].consCounter, groupIndexOneVals[groupIdx]); + tagWithBlockId(state.consCounterNexts[groupIdx], blockId); + state.finalEmitted[groupIdx] = 1; + return success(); +} + +static LogicalResult emitConsumerReleaseAfterPrefix( + OpBuilder &builder, Location loc, SmallVectorImpl &groups, + int groupIdx, ExtendedForInfo &info, Block *oldBody, Value tripCount, + Value loopLowerBound, Value loopStep, + const SmallVector &groupTrueFlagVals, + const SmallVector &groupIndexOneVals, + const SmallVector &iterArgDeltas, scf::ForOp origForOp, + const SmallVector &allGroupDepthConsts, + const SmallVector> &allGroupSlotConsts, + MultiBufferLoopBodyState &state) { + if (!state.prefixEmitted[groupIdx]) { + LogicalResult producerStatus = emitProducerAndSelectSlots( + builder, loc, groups, groupIdx, info, oldBody, tripCount, + loopLowerBound, loopStep, groupTrueFlagVals, groupIndexOneVals, + iterArgDeltas, origForOp, allGroupDepthConsts, allGroupSlotConsts, + state); + if (failed(producerStatus)) { + return failure(); + } + } + + return emitConsumerReleaseForGroup(builder, loc, groups, groupIdx, info, + groupIndexOneVals, allGroupSlotConsts, + state); +} + +static LogicalResult collectGroupsByOwnerOp( + ArrayRef groups, Block *oldBody, + llvm::DenseMap> &groupsByOwnerOp) { + for (int groupIdx = 0; groupIdx < static_cast(groups.size()); + ++groupIdx) { + if (groups[groupIdx].loads.empty()) { + continue; + } + + Operation *owner = + getAncestorInBlock(groups[groupIdx].loads[0].markedOp, oldBody); + if (!owner) { + return failure(); + } + groupsByOwnerOp[owner].push_back(groupIdx); + } + return success(); +} + +static LogicalResult emitProducerHooksForOwnerGroups( + OpBuilder &builder, Location loc, ArrayRef ownerGroups, + SmallVectorImpl &groups, ExtendedForInfo &info, Block *oldBody, + Value tripCount, Value loopLowerBound, Value loopStep, + const SmallVector &groupTrueFlagVals, + const SmallVector &groupIndexOneVals, + const SmallVector &iterArgDeltas, scf::ForOp origForOp, + const SmallVector &allGroupDepthConsts, + const SmallVector> &allGroupSlotConsts, + MultiBufferLoopBodyState &state) { + for (int groupIdx : ownerGroups) { + LogicalResult producerStatus = emitProducerAndSelectSlots( + builder, loc, groups, groupIdx, info, oldBody, tripCount, + loopLowerBound, loopStep, groupTrueFlagVals, groupIndexOneVals, + iterArgDeltas, origForOp, allGroupDepthConsts, allGroupSlotConsts, + state); + if (failed(producerStatus)) { + return failure(); + } + } + return success(); +} + +static LogicalResult emitConsumerReleaseHooksForOwnerGroups( + OpBuilder &builder, Location loc, ArrayRef ownerGroups, + SmallVectorImpl &groups, ExtendedForInfo &info, Block *oldBody, + Value tripCount, Value loopLowerBound, Value loopStep, + const SmallVector &groupTrueFlagVals, + const SmallVector &groupIndexOneVals, + const SmallVector &iterArgDeltas, scf::ForOp origForOp, + const SmallVector &allGroupDepthConsts, + const SmallVector> &allGroupSlotConsts, + MultiBufferLoopBodyState &state) { + for (int groupIdx : ownerGroups) { + LogicalResult releaseStatus = emitConsumerReleaseAfterPrefix( + builder, loc, groups, groupIdx, info, oldBody, tripCount, + loopLowerBound, loopStep, groupTrueFlagVals, groupIndexOneVals, + iterArgDeltas, origForOp, allGroupDepthConsts, allGroupSlotConsts, + state); + if (failed(releaseStatus)) { + return failure(); + } + } + return success(); +} + +static LogicalResult +cloneOldBodyOpIfNeeded(OpBuilder &builder, Operation *op, Block *oldBody, + IRMapping &mapping, + const llvm::DenseSet &allCtxForOps, + const llvm::DenseSet &allChainOps, + const llvm::DenseSet &retainedChainOpSet) { + if (shouldSkipOldBodyOp(op, allCtxForOps, allChainOps, retainedChainOpSet)) { + return success(); + } + if (!areOperandsAvailableForClone(op, oldBody, mapping)) { + return failure(); + } + builder.clone(*op, mapping); + return success(); +} + +static LogicalResult emitOldBodyWithMultiBufferHooks( + OpBuilder &builder, Location loc, SmallVectorImpl &groups, + ExtendedForInfo &info, + const llvm::DenseMap> &groupsByOwnerOp, + const llvm::DenseSet &allCtxForOps, + const llvm::DenseSet &allChainOps, + const llvm::DenseSet &retainedChainOpSet, Value tripCount, + Value loopLowerBound, Value loopStep, + const SmallVector &groupTrueFlagVals, + const SmallVector &groupIndexOneVals, + const SmallVector &iterArgDeltas, scf::ForOp origForOp, + const SmallVector &allGroupDepthConsts, + const SmallVector> &allGroupSlotConsts, + MultiBufferLoopBodyState &state) { + Block *oldBody = info.oldBody; + for (auto &op : oldBody->without_terminator()) { + auto ownerIt = groupsByOwnerOp.find(&op); + if (ownerIt != groupsByOwnerOp.end()) { + LogicalResult producerHookStatus = emitProducerHooksForOwnerGroups( + builder, loc, ownerIt->second, groups, info, oldBody, tripCount, + loopLowerBound, loopStep, groupTrueFlagVals, groupIndexOneVals, + iterArgDeltas, origForOp, allGroupDepthConsts, allGroupSlotConsts, + state); + if (failed(producerHookStatus)) { + return failure(); + } + } + + LogicalResult cloneStatus = + cloneOldBodyOpIfNeeded(builder, &op, oldBody, info.mapping, + allCtxForOps, allChainOps, retainedChainOpSet); + if (failed(cloneStatus)) { + return failure(); + } + + if (ownerIt != groupsByOwnerOp.end()) { + LogicalResult releaseHookStatus = emitConsumerReleaseHooksForOwnerGroups( + builder, loc, ownerIt->second, groups, info, oldBody, tripCount, + loopLowerBound, loopStep, groupTrueFlagVals, groupIndexOneVals, + iterArgDeltas, origForOp, allGroupDepthConsts, allGroupSlotConsts, + state); + if (failed(releaseHookStatus)) { + return failure(); + } + } + } + return success(); +} + +static LogicalResult emitUnownedGroupsAtBodyEnd( + OpBuilder &builder, Location loc, SmallVectorImpl &groups, + ExtendedForInfo &info, Block *oldBody, Value tripCount, + Value loopLowerBound, Value loopStep, + const SmallVector &groupTrueFlagVals, + const SmallVector &groupIndexOneVals, + const SmallVector &iterArgDeltas, scf::ForOp origForOp, + const SmallVector &allGroupDepthConsts, + const SmallVector> &allGroupSlotConsts, + MultiBufferLoopBodyState &state) { + for (int groupIdx = 0; groupIdx < static_cast(groups.size()); + ++groupIdx) { + if (state.prefixEmitted[groupIdx]) { + continue; + } + + LogicalResult producerStatus = emitProducerAndSelectSlots( + builder, loc, groups, groupIdx, info, oldBody, tripCount, + loopLowerBound, loopStep, groupTrueFlagVals, groupIndexOneVals, + iterArgDeltas, origForOp, allGroupDepthConsts, allGroupSlotConsts, + state); + if (failed(producerStatus)) { + return failure(); + } + LogicalResult releaseStatus = emitConsumerReleaseAfterPrefix( + builder, loc, groups, groupIdx, info, oldBody, tripCount, + loopLowerBound, loopStep, groupTrueFlagVals, groupIndexOneVals, + iterArgDeltas, origForOp, allGroupDepthConsts, allGroupSlotConsts, + state); + if (failed(releaseStatus)) { + return failure(); + } + } + return success(); +} + +static void emitCombinedYield(OpBuilder &builder, Location loc, + ArrayRef groups, ExtendedForInfo &info, + const MultiBufferLoopBodyState &state) { + auto oldYield = cast(info.oldBody->getTerminator()); + SmallVector yieldVals; + yieldVals.reserve(info.numOrig + + static_cast(groups.size()) * + (info.depth + kLoadGroupCounterIterArgCount)); + for (Value yieldedValue : oldYield.getOperands()) { + yieldVals.push_back(info.mapping.lookupOrDefault(yieldedValue)); + } + for (int groupIdx = 0; groupIdx < static_cast(groups.size()); + ++groupIdx) { + for (int slotIdx = 0; slotIdx < groups[groupIdx].depth; ++slotIdx) { + yieldVals.push_back(state.flagFinals[groupIdx][slotIdx]); + } + yieldVals.push_back(state.prodCounterFinals[groupIdx]); + yieldVals.push_back(state.consCounterNexts[groupIdx]); + } + builder.create(loc, yieldVals); +} + +} // namespace + +/// Emit the full loop body with producer prefetch, consumer slot selection, +/// cloned compute, slot release, and a combined scf.yield. +LogicalResult emitMultiBufferLoopBody( + OpBuilder &builder, Location loc, SmallVectorImpl &groups, + ExtendedForInfo &info, const SmallVector &groupIndexOneVals, + const SmallVector> &allGroupSlotConsts, + const SmallVector &allGroupDepthConsts, + const llvm::DenseSet &allCtxForOps, Value tripCount, + Value loopLowerBound, Value loopStep, + const SmallVector &groupTrueFlagVals, + const SmallVector &iterArgDeltas, scf::ForOp origForOp) { + int numGroups = static_cast(groups.size()); + Block *oldBody = info.oldBody; + + llvm::DenseSet allChainOps; + llvm::DenseSet retainedChainOpSet; + computeConsumerCloneSets(groups, oldBody, allChainOps, retainedChainOpSet); + + MultiBufferLoopBodyState state(numGroups); + llvm::DenseMap> groupsByOwnerOp; + if (failed(collectGroupsByOwnerOp(groups, oldBody, groupsByOwnerOp))) { + return failure(); + } + + LogicalResult bodyStatus = emitOldBodyWithMultiBufferHooks( + builder, loc, groups, info, groupsByOwnerOp, allCtxForOps, allChainOps, + retainedChainOpSet, tripCount, loopLowerBound, loopStep, + groupTrueFlagVals, groupIndexOneVals, iterArgDeltas, origForOp, + allGroupDepthConsts, allGroupSlotConsts, state); + if (failed(bodyStatus)) { + return failure(); + } + + LogicalResult unownedStatus = emitUnownedGroupsAtBodyEnd( + builder, loc, groups, info, oldBody, tripCount, loopLowerBound, loopStep, + groupTrueFlagVals, groupIndexOneVals, iterArgDeltas, origForOp, + allGroupDepthConsts, allGroupSlotConsts, state); + if (failed(unownedStatus)) { + return failure(); + } + + emitCombinedYield(builder, loc, groups, info, state); + return success(); +} + +} // namespace gmload diff --git a/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromComputePass.cpp b/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromComputePass.cpp new file mode 100644 index 0000000000..4c33ea6d6e --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/SeparateMemoryFromComputePass.cpp @@ -0,0 +1,73 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/SeparateMemoryFromComputePass.h" +#include "ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AddMultiBufferToGMLoadPass.h" +#include "ascend/include/DynamicCVPipeline/SeparateMemoryFromCompute/AsyncLoadHoistingPass.h" +#include "mlir/Pass/PassManager.h" +#include "llvm/Support/Debug.h" + +static constexpr const char *DEBUG_TYPE = "separate-memory-from-compute"; +#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ") +#define LDBG(X) LLVM_DEBUG(DBGS() << (X) << "\n") + +using namespace mlir; +using namespace triton; + +static constexpr int kDefaultBufferDepth = 2; + +void SeparateMemoryFromComputePass::runOnOperation() { + ModuleOp module = getOperation(); + + int depth = kDefaultBufferDepth; + + if (depth <= 1) { + LDBG("Buffer depth <= 1, skip multi-buffer transformation"); + return; + } + + OpPassManager pm(module.getOperationName()); + LDBG("Enter SeparateMemoryFromCompute pass"); + + // Step 1: Hoist memory operations out of compute blocks + pm.addPass(createAsyncLoadHoistingPass()); + + // Step 2: Apply multi-buffering to memory operations + pm.addPass(createAddMultiBufferToGMLoadPass()); + + if (failed(runPipeline(pm, module))) { + module->emitError() << "[" << DEBUG_TYPE << "] Pass failed!"; + signalPassFailure(); + } + + LDBG("Process successfully"); +} + +namespace mlir { +namespace triton { + +std::unique_ptr> createSeparateMemoryFromComputePass() { + return std::make_unique(); +} + +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow.cpp b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow.cpp new file mode 100644 index 0000000000..865e0fa5b3 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow.cpp @@ -0,0 +1,79 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/SplitDataflow/AddBlockIdForControlOps.h" +#include "ascend/include/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.h" +#include "ascend/include/DynamicCVPipeline/SplitDataflow/InterCoreTransferAndSync.h" +#include "ascend/include/DynamicCVPipeline/SplitDataflow/MarkMainLoop.h" +#include "ascend/include/DynamicCVPipeline/SplitDataflow/PreserveControlAttrsCanonicalize.h" +#include "ascend/include/DynamicCVPipeline/SplitDataflow/SeparateCVScope.h" +#include "ascend/include/DynamicCVPipeline/SplitDataflowPass.h" +#include "mlir/Pass/PassManager.h" +#include "llvm/Support/Debug.h" + +static constexpr const char *DEBUG_TYPE = "SplitDataflow"; +#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ") +#define LDBG(X) LLVM_DEBUG(DBGS() << (X) << "\n") + +using namespace mlir; +using namespace triton; + +// Run the pass +void SplitDataflowPass::runOnOperation() { + ModuleOp module = getOperation(); + OpPassManager pm(module.getOperationName()); + LDBG("Enter pass."); + + // Step 1: Add block_id for control flow operations + pm.addPass(createAddBlockIdForControlOpsPass()); + + // Step 2: Analyze data dependencies between Vector and Cube blocks + pm.addPass(createDataDependencyAnalysisPass()); + + // Step 3: Run InterCoreTransferAndSync + pm.addPass(createInterCoreTransferAndSyncPass()); + + // Step 4: Mark the main computation loop + pm.addPass(createMarkMainLoopPass()); + + // Step 5: Run SeparateCVScope + pm.addPass(createSeparateCVScopePass()); + + // Step 6: Canonicalize to preserve control flow attributes + pm.addPass(createPreserveControlAttrsCanonicalizePass()); + + if (failed(runPipeline(pm, module))) { + module->emitError() << "[" << DEBUG_TYPE << "] Pass failed!"; + signalPassFailure(); + } + + LDBG("Process successfully"); +} + +namespace mlir { +namespace triton { + +std::unique_ptr> createSplitDataflowPass() { + return std::make_unique(); +} +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/AddBlockIdForControlOps.cpp b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/AddBlockIdForControlOps.cpp new file mode 100644 index 0000000000..6c488e165c --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/AddBlockIdForControlOps.cpp @@ -0,0 +1,83 @@ + +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/SplitDataflow/AddBlockIdForControlOps.h" +#include "llvm/Support/Debug.h" + +using namespace mlir; + +static constexpr const char *DEBUG_TYPE = "add-block-id-for-control-ops"; +#define LOG_DEBUG(...) \ + LLVM_DEBUG(llvm::dbgs() << " [" << DEBUG_TYPE << "] " << __VA_ARGS__) + +using namespace mlir::triton; + +// Pass Entry Point +void AddBlockIdForControlOpsPass::runOnOperation() { + LOG_DEBUG("\n--- enter AddBlockIdForControlOpsPass --->\n"); + ModuleOp module = getOperation(); + + // Step 1: find the max block_id + int maxBlockId = -1; + module.walk([&](Operation *op) { + if (auto attr = op->getAttrOfType("ssbuffer.block_id")) { + int currentId = attr.getInt(); + if (currentId > maxBlockId) { + maxBlockId = currentId; + } + } + }); + + LOG_DEBUG("maxBlockId: " << maxBlockId << "\n"); + + // Step 2: add block_id for control flow ops + module.walk([&](Operation *op) { + // skip op with block_id + if (op->getAttrOfType("ssbuffer.block_id")) { + return; + } + + if (isa(op)) { + maxBlockId++; + static constexpr int kIntegerBitWidth = 32; + op->setAttr("ssbuffer.block_id", + IntegerAttr::get( + IntegerType::get(module.getContext(), kIntegerBitWidth), + maxBlockId)); + LOG_DEBUG("Added block_id " << maxBlockId << " to " << *op << "\n"); + } + }); + + LOG_DEBUG("\n--- exit AddBlockIdForControlOpsPass --->\n"); +} + +// Create the pass +namespace mlir { +namespace triton { + +std::unique_ptr> createAddBlockIdForControlOpsPass() { + return std::make_unique(); +} + +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.cpp b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.cpp new file mode 100644 index 0000000000..deb6e0f7bf --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.cpp @@ -0,0 +1,489 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.h" +#include "ascend/include/DynamicCVPipeline/Common/MemoryEffectsTracker.h" + +#include "mlir/Analysis/AliasAnalysis.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Dominance.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/OperationSupport.h" +#include "mlir/IR/Value.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Debug.h" + +using namespace mlir; + +static constexpr const char *DEBUG_TYPE = "data-dependency-analysis"; +#define LOG_DEBUG(...) \ + LLVM_DEBUG(llvm::dbgs() << " [" << DEBUG_TYPE << "] " << __VA_ARGS__) + +using namespace mlir::triton; +using namespace mlir::CVPipeline; + +// Attribute name constants +static constexpr const char *kBlockIdAttr = "ssbuffer.block_id"; +static constexpr const char *kCoreTypeAttr = "ssbuffer.core_type"; +static constexpr const char *kTransferIdAttr = "ssbuffer.transfer_id"; + +// Helper: ssbuffer.core_type +llvm::StringRef getSsbufferCoreType(Operation *op) { + if (auto attr = op->getAttrOfType(kCoreTypeAttr)) { + return attr.getValue(); + } + return ""; +} + +// Helper: Get CoreType from op and index +llvm::StringRef getCoreTypeWithIndex(Operation *op, int index) { + llvm::StringRef typeStr = getSsbufferCoreType(op); + if (typeStr.contains(", ")) { + llvm::SmallVector types; + typeStr.split(types, ", ", -1, false); + if (index < types.size()) { + return types[index].trim(); + } + LOG_DEBUG("Warning: Core type string has multiple types but value is not " + "an OpResult or index out of range.\n"); + return ""; + } + + return typeStr; +} + +// Helper: Check if operation is control flow +bool DataDependencyAnalysisPass::isControlFlowOp(mlir::Operation *op) { + if (!op) + return false; + return isa(op) || isa(op) || isa(op) || + isa(op); +} + +// Helper: Build and record BlockInfo +void DataDependencyAnalysisPass::collectBlockInfo( + DataDependencyInfo &info, int blockId, + llvm::SmallVector &ops) { + if (ops.empty()) { + LOG_DEBUG("Warning: Block ID " << blockId << " has no operations.\n"); + return; + } + + BlockInfo blockInfo; + blockInfo.blockId = blockId; + blockInfo.isCube = false; + + // In cases with one or more core_types + // as long as there is a cube, it is necessary to check the dataflow. + StringRef coreType = getSsbufferCoreType(ops[0]); + if (coreType.contains("CUBE")) { + blockInfo.isCube = true; + } + + blockInfo.isControl = false; + if (isControlFlowOp(ops[0])) { + blockInfo.isControl = true; + } + + llvm::DenseSet opSet(ops.begin(), ops.end()); + + for (auto *op : ops) { + blockInfo.Operations.push_back(op); + for (auto operand : op->getOperands()) { + mlir::Operation *defOp = operand.getDefiningOp(); + // If defOp is not null and defOp is not in current ops set, it's an + // external input + if (!defOp || opSet.find(defOp) == opSet.end()) { + blockInfo.inputs.push_back(operand); + } + } + for (auto result : op->getResults()) { + // If any user is not in the current ops set, it's an external output + bool hasExternalUser = false; + for (mlir::Operation *user : result.getUsers()) { + if (opSet.find(user) == opSet.end()) { + hasExternalUser = true; + break; + } + } + if (hasExternalUser) { + blockInfo.outputs.push_back(result); + } + } + } + + info.getBlockInfoMap()[blockInfo.blockId] = blockInfo; + + LOG_DEBUG("Block_ID=" << blockInfo.blockId << "Processed!\n"); +} + +// Block Information Collection +void DataDependencyAnalysisPass::createBlockInfoMap(DataDependencyInfo &info) { + int currentId = -2; + static constexpr int startCurrId = -2; + llvm::SmallVector currentOps; + + module.walk([&](mlir::Operation *op) { + int opBlockId = getSsbufferBlockId(op); + if (opBlockId != -1) { + // When the id changes, the block ends && Exclude the initial state + if (opBlockId != currentId && currentId != startCurrId) { + collectBlockInfo(info, currentId, currentOps); + currentOps.clear(); + } + currentId = opBlockId; + currentOps.push_back(op); + } + }); + // Process the last group + if (!currentOps.empty()) { + collectBlockInfo(info, currentId, currentOps); + } +} + +void DataDependencyAnalysisPass::collectDepInfo( + mlir::Value depvalue, DependencyType dependencyType, + llvm::SmallVector &dependencies, int iniProdId, + int iniConsId, DataDependencyInfo &info) { + DependencyInfo depInfo; + depInfo.type = dependencyType; + depInfo.value = depvalue; + LOG_DEBUG("try finding common level block IDs\n"); + depInfo.iniProducerBlockId = iniProdId; + depInfo.iniConsumerBlockId = iniConsId; + std::pair commonLevelIds = + findCommonLevelBlockIds(info, iniProdId, iniConsId); + assert( + commonLevelIds.first != -1 && commonLevelIds.second != -1 && + "Could not find common level block IDs for producer and consumer blocks"); + depInfo.producerBlockId = commonLevelIds.first; + depInfo.consumerBlockId = commonLevelIds.second; + + dependencies.push_back(depInfo); +} + +// Analyze V->C +void DataDependencyAnalysisPass::analyzeExternalInputs( + DataDependencyInfo &info) { + auto &blockInfoMap = info.getBlockInfoMap(); + auto &v2cDependencies = info.getV2CDependencies(); + + LOG_DEBUG("Analyzing external inputs for Cube blocks...\n"); + for (auto &[id, blockInfo] : blockInfoMap) { + if (!blockInfo.isCube || blockInfo.isControl || blockInfo.inputs.empty()) + continue; + LOG_DEBUG("Analyzing external inputs for Cube Block ID: " << id << "\n"); + for (mlir::Value input : blockInfo.inputs) { + // Check if input is a func.func blockarg. + if (auto arg = mlir::dyn_cast(input)) { + LOG_DEBUG("Warning: [v->c] Input value is a function parameter.\n"); + continue; + } + // Check if input is a value which can be produced by CUBE + if (!dyn_cast(input.getType())) { + LOG_DEBUG("Warning: [v->c] Input value is not TensorType\n"); + continue; + } + if (isa(input.getDefiningOp())) { + LOG_DEBUG("Warning: [v->c] Input value is defined by " + "tensor::EmptyOp/linalg::FillOp.\n"); + continue; + } + + Operation *defOp = input.getDefiningOp(); + auto defReuslt = dyn_cast(input); + auto coreType = getCoreTypeWithIndex( + defOp, defReuslt ? defReuslt.getResultNumber() : 0); + if (coreType == "") { + LOG_DEBUG("Warning: [v->c] Input value has no core type attribute.\n"); + continue; + } + + // Case 1: Cube -> C->C special case + if (coreType == "CUBE") { + continue; + } + // Case 2: Vector -> V->C dependency + if (coreType == "VECTOR") { + LOG_DEBUG("Found external input with VECTOR core type: " << input + << "\n"); + auto producerId = getSsbufferBlockId(input.getDefiningOp()); + if (producerId == -1) { + LOG_DEBUG( + "Warning: [v->c] Producer block ID not found for input value.\n"); + continue; + } + collectDepInfo(input, DependencyType::VectorToCube, v2cDependencies, + producerId, blockInfo.blockId, info); + } + } + } + LOG_DEBUG("External input analysis complete.\n"); +} + +// Analyze C->V +void DataDependencyAnalysisPass::analyzeExternalOutputs( + DataDependencyInfo &info) { + auto &blockInfoMap = info.getBlockInfoMap(); + auto &c2vDependencies = info.getC2VDependencies(); + + LOG_DEBUG("Analyzing external outputs for Cube blocks...\n"); + for (auto &[id, blockInfo] : blockInfoMap) { + if (!blockInfo.isCube || blockInfo.outputs.empty()) + continue; + + for (mlir::Value output : blockInfo.outputs) { + // Check if input is a value which can be produced by VECTOR + if (!dyn_cast(output.getType())) { + LOG_DEBUG("Warning: ExternalOutput is not TensorType\n"); + continue; + } + if (isa(output.getDefiningOp())) { + LOG_DEBUG("Warning: [c->v] output value is defined by " + "tensor::EmptyOp/linalg::FillOp.\n"); + continue; + } + + auto opResult = dyn_cast(output); + unsigned resultIndex = opResult.getResultNumber(); + StringRef resultCoreType = + getCoreTypeWithIndex(output.getDefiningOp(), resultIndex); + if (resultCoreType != "CUBE") { + continue; + } + // Check who is using this output + for (mlir::Operation *user : output.getUsers()) { + int outputIndex = 0; + if (isControlFlowOp(user)) { + for (unsigned i = 0; i < user->getNumOperands(); ++i) { + if (user->getOperand(i) == output) { + outputIndex = i; + break; + } + } + } + auto userCoreType = getCoreTypeWithIndex(user, outputIndex); + if ((userCoreType == "")) { + LOG_DEBUG( + "Warning: [c->v] Input value has no core type attribute.\n"); + continue; + } + if (userCoreType == "VECTOR") { + LOG_DEBUG("Found external output used by VECTOR core type: " << output + << "\n"); + auto consumerId = getSsbufferBlockId(user); + if (consumerId == -1) { + LOG_DEBUG("Warning: [c->v] Consumer block ID not found for user " + "operation.\n"); + continue; + } + collectDepInfo(output, DependencyType::CubeToVector, c2vDependencies, + blockInfo.blockId, consumerId, info); + } + // If user belongs to Cube block, this C->C dependency was handled + // in the Input analysis phase, so here we only handle C->V. + } + } + } + LOG_DEBUG("External output analysis complete.\n"); +} + +void DataDependencyAnalysisPass::analyzeMemoryEffect(DataDependencyInfo &info) { + auto &memoryDependencies = info.getMemoryDependencies(); + LOG_DEBUG("\n=== start mem dep analysis ===\n"); + + auto &aliasAnalysis = getAnalysis(); + MemoryDependenceGraph memDepGraph(module, aliasAnalysis); + + module.walk([&](mlir::Operation *op) { + int currBlockId = getSsbufferBlockId(op); + llvm::StringRef currCoreType = getSsbufferCoreType(op); + if (currBlockId == -1 || currCoreType.empty()) { + return; + } + + for (mlir::Operation *predOp : memDepGraph.getExecBefore(op)) { + int predBlockId = getSsbufferBlockId(predOp); + llvm::StringRef predCoreType = getSsbufferCoreType(predOp); + if (predBlockId == -1 || predCoreType == currCoreType || + predCoreType.empty()) { + continue; + } + + auto [producerBlockId, consumerBlockId] = + findCommonLevelBlockIds(info, predBlockId, currBlockId); + assert(producerBlockId != -1 && consumerBlockId != -1 && + "Could not find common level block IDs for producer and consumer " + "blocks"); + if (producerBlockId == consumerBlockId) { + continue; + } + + DependencyInfo depInfo; + + if (predCoreType == "CUBE") { + depInfo.type = DependencyType::CubeToVector; + } else if (predCoreType == "VECTOR") { + depInfo.type = DependencyType::VectorToCube; + } + depInfo.producerBlockId = producerBlockId; + depInfo.consumerBlockId = consumerBlockId; + depInfo.iniProducerBlockId = predBlockId; + depInfo.iniConsumerBlockId = currBlockId; + + memoryDependencies.push_back(depInfo); + LOG_DEBUG("=mem dep analysis= " << "producer Block: " << predBlockId + << "consumer Block: " << currBlockId + << "\n"); + } + }); + LOG_DEBUG("=== mem dep analysis complete ===\n"); +} + +// Producer/Consumer Hierarchy Analysis +std::pair DataDependencyAnalysisPass::findCommonLevelBlockIds( + DataDependencyInfo &info, int producerBlockId, int consumerBlockId) { + auto &blockInfoMap = info.getBlockInfoMap(); + + LOG_DEBUG("start findCommonLevelBlockIds...\n"); + + // Step 1: Get corresponding BlockInfo from Map + auto pIt = blockInfoMap.find(producerBlockId); + auto cIt = blockInfoMap.find(consumerBlockId); + // Defensive programming: if corresponding Block info not found, return + // original ID or error code + if (pIt == blockInfoMap.end() || cIt == blockInfoMap.end()) { + return {producerBlockId, consumerBlockId}; + } + + BlockInfo &pInfo = pIt->second; + BlockInfo &cInfo = cIt->second; + + // Take the first operation of each Block as representative to check hierarchy + // (Assumes all operations in a Block are closely related in hierarchy) + mlir::Operation *producerOp = pInfo.Operations[0]; + mlir::Operation *consumerOp = cInfo.Operations[0]; + + mlir::Block *pBlock = producerOp->getBlock(); + mlir::Block *cBlock = consumerOp->getBlock(); + + // Case 1: In the same MLIR Block + if (pBlock == cBlock) { + return {producerBlockId, consumerBlockId}; + } + + // Case 2: In different Blocks, find Lowest Common Ancestor (LCA) + // Step 1: Collect producer's ancestor chain + llvm::SmallVector pAncestors; + pAncestors.push_back(producerOp); + mlir::Operation *current = producerOp->getParentOp(); + while (current) { + pAncestors.push_back(current); + current = current->getParentOp(); + } + + // Step 2: Walk up consumer's ancestors, using current and before for rolling + mlir::Operation *before = consumerOp; // Initialize as consumerOp itself + current = consumerOp; // Initialize as parent + + while (current) { + // --- Found common ancestor --- + auto it = std::find(pAncestors.begin(), pAncestors.end(), current); + if (it != pAncestors.end()) { + size_t pIndex = std::distance(pAncestors.begin(), it); + if (pIndex == 0) { + break; + } + mlir::Operation *pPrevOp = pAncestors[pIndex - 1]; + int pPrevId = getSsbufferBlockId(pPrevOp); + int cPrevId = getSsbufferBlockId(before); + if (pPrevId == -1) { + LOG_DEBUG("Warning: Producer ancestor operation has no block ID " + "attribute.\n"); + } + if (cPrevId == -1) { + LOG_DEBUG("Warning: Consumer ancestor operation has no block ID " + "attribute.\n"); + } + return {pPrevId, cPrevId}; + } + + // Rolling: continue upward + before = current; + current = current->getParentOp(); + } + LOG_DEBUG("Warning: No common ancestor found for producer block " + << producerBlockId << " and consumer block " << consumerBlockId + << "\n"); + return {-1, -1}; +} + +void DataDependencyAnalysisPass::runOnOperation() { + LOG_DEBUG("\n--- enter DataDependencyAnalysisPass --->\n"); + module = getOperation(); + + auto &info = getAnalysis(); + + // Step 1: Collect block information (populate blockInfoMap) + createBlockInfoMap(info); + + // Step 2: Analyze dependencies (populate v2c, c2v lists) + analyzeExternalInputs(info); + analyzeExternalOutputs(info); + + // Step 3: Analyze memory dependencies (PIPE_S sync) + analyzeMemoryEffect(info); + + info.setValid(true); + + LOG_DEBUG("DataDependencyAnalysisPass: Analysis complete.\n"); + LOG_DEBUG(" V->C dependencies: " << info.getV2CDependencies().size() + << "\n"); + LOG_DEBUG(" C->V dependencies: " << info.getC2VDependencies().size() + << "\n"); + LOG_DEBUG(" Memory dependencies: " << info.getMemoryDependencies().size() + << "\n"); + + LOG_DEBUG("\n--- exit DataDependencyAnalysisPass --->\n"); +} + +// Create the pass +namespace mlir { +namespace triton { +std::unique_ptr> createDataDependencyAnalysisPass() { + return std::make_unique(); +} + +// Helper: Get BlockId +int getSsbufferBlockId(Operation *op) { + if (auto attr = op->getAttrOfType(kBlockIdAttr)) { + return attr.getInt(); + } + return -1; +} +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/InterCoreTransferAndSync.cpp b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/InterCoreTransferAndSync.cpp new file mode 100644 index 0000000000..28b019381f --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/InterCoreTransferAndSync.cpp @@ -0,0 +1,1048 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/SplitDataflow/InterCoreTransferAndSync.h" + +#include +#include + +#include "ascend/include/DynamicCVPipeline/Common/FlagIdManager.h" +#include "ascend/include/DynamicCVPipeline/SplitDataflow/DataDependencyAnalysis.h" + +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/HIVM/IR/HIVMImpl.h" +#include "bishengir/Dialect/HIVM/IR/HIVMInterfaces.h" +#include "incubated/Conversion/UtilsIncubated/Utils.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/Dominance.h" +#include "mlir/IR/Operation.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Debug.h" + +using namespace mlir; + +static constexpr const char *DEBUG_TYPE = "inter-core-transfer-and-sync"; +#define LOG_DEBUG(...) \ + LLVM_DEBUG(llvm::dbgs() << " [" << DEBUG_TYPE << "] " << __VA_ARGS__) + +using namespace mlir::triton; +using namespace hivm; + +// Attribute name constants +static constexpr const char *kBlockIdAttr = "ssbuffer.block_id"; +static constexpr const char *kCoreTypeAttr = "ssbuffer.core_type"; +static constexpr const char *kTransferIdAttr = "ssbuffer.transfer_id"; + +static constexpr int kIntegerBitWidth = 32; +static constexpr int NzDimWidth = 16; + +static uint64_t getElemBytesForAlign(Type t) { + static constexpr uint64_t kBitsPerByte = 8; + if (auto ft = dyn_cast(t)) { + return (uint64_t)((ft.getWidth() + kBitsPerByte - 1) / kBitsPerByte); + } + if (auto it = dyn_cast(t)) { + return (uint64_t)((it.getWidth() + kBitsPerByte - 1) / kBitsPerByte); + } + if (isa(t)) { + return 8ULL; + } + if (auto ct = dyn_cast(t)) { + return 2ULL * getElemBytesForAlign(ct.getElementType()); + } + return 0ULL; +} + +static uint64_t getBlockElemsFor32BAlign(Type elemType) { + constexpr uint64_t kAlignBytes = 32; + uint64_t elemBytes = getElemBytesForAlign(elemType); + if (elemBytes == 0) { + return -1; + } + if (elemBytes < 0 || kAlignBytes % elemBytes != 0) { + return -1; + } + if (elemBytes >= kAlignBytes) { + return 1; + } + return kAlignBytes / elemBytes; +} + +static void attachCommonTags(Operation *op, int blockId, StringRef coreType) { + MLIRContext *ctx = op->getContext(); + op->setAttr( + kBlockIdAttr, + IntegerAttr::get(IntegerType::get(ctx, kIntegerBitWidth), blockId)); + op->setAttr(kCoreTypeAttr, StringAttr::get(ctx, coreType)); +} + +static void attachTransferTags(Operation *op, int blockId, StringRef coreType, + int transferId) { + MLIRContext *ctx = op->getContext(); + op->setAttr( + kBlockIdAttr, + IntegerAttr::get(IntegerType::get(ctx, kIntegerBitWidth), blockId)); + op->setAttr(kCoreTypeAttr, StringAttr::get(ctx, coreType)); + op->setAttr( + kTransferIdAttr, + IntegerAttr::get(IntegerType::get(ctx, kIntegerBitWidth), transferId)); +} + +// Block Start/End Operation Retrieval +std::pair +InterCoreTransferAndSyncPass::getBlockStartEnd(int targetId, + mlir::ModuleOp module) { + mlir::Operation *knownOpInBlock = nullptr; + module.walk([&](mlir::Operation *op) { + if (knownOpInBlock) { + return; + } + if (getSsbufferBlockId(op) == targetId) { + knownOpInBlock = op; + } + }); + + if (!knownOpInBlock) { + return {nullptr, nullptr}; + } + + mlir::Block *block = knownOpInBlock->getBlock(); + if (!block) + return {nullptr, nullptr}; + + mlir::Operation *start = nullptr; + mlir::Operation *end = nullptr; + + // Iterate through all operations in the current block + for (Operation &op : *block) { + int blockId = getSsbufferBlockId(&op); + if (blockId == -1) { + continue; + } + + if (!start) { + if (targetId == blockId) { + start = &op; + end = &op; + } + } else { + if (targetId == blockId) { + end = &op; + } else { + break; + } + } + } + return {start, end}; +} + +bool InterCoreTransferAndSyncPass::isOuterLayerDependency( + size_t depIndex, mlir::Operation *currProdEnd, + mlir::Operation *currConsStart, + llvm::SmallVector &memDependencies) { + if (!currProdEnd || !currConsStart) { + return false; + } + mlir::Block *currBlock = currProdEnd->getBlock(); + if (currBlock != currConsStart->getBlock()) { + return false; + } + for (size_t i = 0; i < memDependencies.size(); ++i) { + if (i == depIndex) { + continue; + } + auto &otherDep = memDependencies[i]; + + auto [otherProdStart, otherProdEnd] = + getBlockStartEnd(otherDep.producerBlockId, module); + auto [otherConsStart, otherConsEnd] = + getBlockStartEnd(otherDep.consumerBlockId, module); + + if (!otherProdEnd || !otherConsStart) { + continue; + } + + if (otherProdEnd->getBlock() != currBlock || + otherConsStart->getBlock() != currBlock) { + continue; + } + + // otherProdEnd is before currProdEnd + // AND currConsStart is before otherConsStart + bool isOtherInsideCurrent = !otherProdEnd->isBeforeInBlock(currProdEnd) && + !currConsStart->isBeforeInBlock(otherConsStart); + + if (otherProdEnd == currProdEnd && otherConsStart == currConsStart) { + if (i < depIndex) { + // if otherDep has smaller index, current dep is outer layer and can be + // skipped + return true; + } + } else if (isOtherInsideCurrent) { + return true; + } + } + + return false; +} + +// Nd2NzNormalizer +SmallVector +InterCoreTransferAndSyncPass::computeExpectedShape(Value value) { + auto tensorTy = dyn_cast(value.getType()); + static constexpr int NdShapeLength = 2; + assert(tensorTy && tensorTy.getRank() == NdShapeLength && + "source shape is not 2-dim!"); + + int64_t M = tensorTy.getDimSize(0); + int64_t N = tensorTy.getDimSize(1); + + // Compute bit width & Nwidth + int64_t nWidth = getBlockElemsFor32BAlign(tensorTy.getElementType()); + if (nWidth == 0) { + LOG_DEBUG("Unsupported element type for 32B alignment.\n"); + return {M, N}; + } + // Calculate newM / newN using the formula + int64_t blM = (M + NzDimWidth - 1) / NzDimWidth; + int64_t newM = blM * NzDimWidth; + + int64_t blN = (N + nWidth - 1) / nWidth; + int64_t newN = blN * nWidth; + + return {newM, newN}; // Return 2D shape +} + +bool InterCoreTransferAndSyncPass::isShapeExpected( + Value value, SmallVector &expectedShape) { + auto tensorTy = dyn_cast(value.getType()); + ArrayRef currShape = tensorTy.getShape(); + return currShape.equals(expectedShape); +} + +void InterCoreTransferAndSyncPass::rewriteMatmulWithNewShape( + OpBuilder &builder, Operation *matmulOp, Location loc) { + int matmulOpBlockId = getSsbufferBlockId(matmulOp); + + Value lhs = matmulOp->getOperands()[0]; + Value rhs = matmulOp->getOperands()[1]; + Value acc = matmulOp->getOperands()[2]; + Value originalResult = matmulOp->getResult(0); + auto lhsType = dyn_cast(lhs.getType()); + auto rhsType = dyn_cast(rhs.getType()); + auto accType = dyn_cast(acc.getType()); + auto resType = dyn_cast(originalResult.getType()); + ArrayRef accshape = accType.getShape(); + ArrayRef resshape = resType.getShape(); + SmallVector expectedShape = {lhsType.getShape()[0], + rhsType.getShape()[1]}; + auto expectedType = + RankedTensorType::get(expectedShape, resType.getElementType()); + + builder.setInsertionPoint(matmulOp); + + auto floatElemTy = cast(resType.getElementType()); + auto zeroConstOp = builder.create( + loc, floatElemTy, APFloat::getZero(floatElemTy.getFloatSemantics())); + auto tensorEmptyOp = builder.create( + loc, expectedShape, resType.getElementType()); + auto linalgFillOp = builder.create( + loc, zeroConstOp.getResult(), tensorEmptyOp.getResult()); + + attachCommonTags(zeroConstOp, matmulOpBlockId, "CUBE"); + attachCommonTags(tensorEmptyOp, matmulOpBlockId, "CUBE"); + attachCommonTags(linalgFillOp, matmulOpBlockId, "CUBE"); + + Value newAccResult = linalgFillOp->getResult(0); + + builder.setInsertionPointAfter(matmulOp); + static constexpr int accIndex = 2; + matmulOp->setOperand(accIndex, newAccResult); + matmulOp->getResult(0).setType(expectedType); + auto newMatmulOp = dyn_cast(matmulOp); + Value newMatmulResult = newMatmulOp->getResult(0); + LOG_DEBUG("newmatmulOp" << newMatmulOp << "\n"); + SmallVector offsets = {builder.getIndexAttr(0), + builder.getIndexAttr(0)}; + SmallVector strides = {builder.getIndexAttr(1), + builder.getIndexAttr(1)}; + SmallVector sizes = { + builder.getIndexAttr(accType.getShape()[0]), + builder.getIndexAttr(accType.getShape()[1])}; + auto extractSliceOp = builder.create( + loc, newMatmulResult, offsets, sizes, strides); + attachCommonTags(extractSliceOp, matmulOpBlockId, "CUBE"); + + originalResult.replaceUsesWithIf( + extractSliceOp.getResult(), [&](OpOperand &use) { + return use.getOwner() != extractSliceOp.getOperation(); + }); + LOG_DEBUG("cubeValueMapping[originalResult]" << originalResult << "\n"); + LOG_DEBUG("cubeValueMapping[originalResult]extractSliceOp.getResult() " + << extractSliceOp.getResult() << "\n"); + cubeValueMapping[originalResult] = extractSliceOp.getResult(); +} + +void InterCoreTransferAndSyncPass::rewriteTransposeWithNewShape( + OpBuilder &builder, Operation *transposeOp, Location loc) { + Value inputvalue = transposeOp->getOperands()[0]; + Value outputvalue = transposeOp->getOperands()[0]; + + auto inputTy = dyn_cast(inputvalue.getType()); + Type elemType = inputTy.getElementType(); + SmallVector newOutputShape = {inputTy.getShape()[1], + inputTy.getShape()[0]}; + auto expectedType = RankedTensorType::get(newOutputShape, elemType); + // Create new empty tensor with new shape + auto tensorEmptyOp = + builder.create(loc, newOutputShape, elemType); + attachCommonTags(tensorEmptyOp, getSsbufferBlockId(transposeOp), "CUBE"); + Value transposeOpResult = transposeOp->getResult(0); + transposeOp->setOperand(1, tensorEmptyOp.getResult()); + transposeOp->getResult(0).setType(expectedType); +} + +// padding v->c tensor +mlir::Value InterCoreTransferAndSyncPass::normalizeIfNeeded( + OpBuilder &builder, DependencyInfo &dep, Location loc, + mlir::Value origValue, SmallVector expectedShape, + int originBlockId) { + auto origTensorType = dyn_cast(origValue.getType()); + + int64_t iniM = origTensorType.getDimSize(0); + int64_t iniN = origTensorType.getDimSize(1); + Type elemType = origTensorType.getElementType(); + + builder.setInsertionPointAfter(origValue.getDefiningOp()); + + auto floatElemTy = cast(elemType); + auto zeroConstOp = builder.create( + loc, floatElemTy, APFloat::getZero(floatElemTy.getFloatSemantics())); + auto tensorEmptyOp = + builder.create(loc, expectedShape, elemType); + auto linalgFillOp = builder.create( + loc, zeroConstOp.getResult(), tensorEmptyOp.getResult()); + SmallVector offsets = {builder.getIndexAttr(0), + builder.getIndexAttr(0)}; + SmallVector insertsizes = {builder.getIndexAttr(iniM), + builder.getIndexAttr(iniN)}; + SmallVector strides = {builder.getIndexAttr(1), + builder.getIndexAttr(1)}; + auto tensorInsertSliceOp = builder.create( + loc, origValue, linalgFillOp->getResult(0), offsets, insertsizes, + strides); + + attachCommonTags(zeroConstOp, originBlockId, "VECTOR"); + attachCommonTags(tensorEmptyOp, originBlockId, "VECTOR"); + attachCommonTags(linalgFillOp, originBlockId, "VECTOR"); + attachCommonTags(tensorInsertSliceOp, originBlockId, "VECTOR"); + int cId = dep.iniConsumerBlockId; + LOG_DEBUG("int cId = dep.iniConsumerBlockId;" << cId << "\n"); + for (Operation *user : origValue.getUsers()) { + LOG_DEBUG(*user << "\n"); + int userBlockId = getSsbufferBlockId(user); + LOG_DEBUG("int userBlockId = getSsbufferBlockId(user);" << userBlockId + << "\n"); + if (userBlockId == -1 || userBlockId != cId) { + continue; + } + user->replaceUsesOfWith(origValue, tensorInsertSliceOp.getResult()); + if (auto matmulOp = dyn_cast(user)) { + rewriteMatmulWithNewShape(builder, matmulOp, loc); + continue; + } + if (auto transposeOp = dyn_cast(user)) { + LOG_DEBUG("before rewriteTransposeWithNewShape\n"); + rewriteTransposeWithNewShape(builder, transposeOp, loc); + LOG_DEBUG("after rewriteTransposeWithNewShape\n"); + for (Operation *transposeuser : transposeOp->getUsers()) { + auto matmulOp = dyn_cast(transposeuser); + if (matmulOp && getSsbufferBlockId(matmulOp) == cId) { + rewriteMatmulWithNewShape(builder, matmulOp, loc); + } + } + } + } + cubeValueMapping[origValue] = tensorInsertSliceOp.getResult(); + return tensorInsertSliceOp.getResult(); +} + +void InterCoreTransferAndSyncPass::Nd2NzNormalize(OpBuilder &builder, + DependencyInfo &dep, + Location loc) { + Value origValue = dep.value; + Value newValue = origValue; + // Step 0: Check if this Value has already been processed + auto it = vecValueMapping.find(origValue); + if (it != vecValueMapping.end()) { + return; + } + // Step 1: Compute expected shape + SmallVector expectedShape = computeExpectedShape(origValue); + int originBlockId = getSsbufferBlockId(origValue.getDefiningOp()); + // Step 2: If shapes match, return original value + if (!isShapeExpected(origValue, expectedShape)) { + newValue = normalizeIfNeeded(builder, dep, loc, origValue, expectedShape, + originBlockId); + } + // Step 3: insert nd2nz + auto srcTensorType = cast(newValue.getType()); + int64_t M = srcTensorType.getDimSize(0); + int64_t N = srcTensorType.getDimSize(1); + Type elemType = srcTensorType.getElementType(); + + int64_t blk = getBlockElemsFor32BAlign(elemType); + if (blk == 0) { + LOG_DEBUG("Invalid block size.\n"); + return; + } + + SmallVector shape3D = {M, N / blk, blk}; + SmallVector shapeTrans = {N / blk, M, blk}; + SmallVector shapeFinal = {N / blk, M / NzDimWidth, NzDimWidth, blk}; + + auto type3D = RankedTensorType::get(shape3D, elemType); + auto typeTrans = RankedTensorType::get(shapeTrans, elemType); + auto typeFinal = RankedTensorType::get(shapeFinal, elemType); + builder.setInsertionPointAfter(newValue.getDefiningOp()); + auto reshape3Dcst = + builder.create(loc, builder.getI64TensorAttr(shape3D)); + auto reshape3DOp = + builder.create(loc, type3D, newValue, reshape3Dcst); + + auto emptyTrans = builder.create(loc, shapeTrans, elemType); + SmallVector transposeOrder = {1, 0, 2}; + auto transposeOp = builder.create( + loc, reshape3DOp.getResult(), emptyTrans.getResult(), transposeOrder); + auto reshape4Dcst = builder.create( + loc, builder.getI64TensorAttr(shapeFinal)); + auto reshape4DOp = builder.create( + loc, typeFinal, transposeOp->getResult(0), reshape4Dcst); + + attachCommonTags(reshape3Dcst, originBlockId, "VECTOR"); + attachCommonTags(reshape3DOp, originBlockId, "VECTOR"); + attachCommonTags(emptyTrans, originBlockId, "VECTOR"); + attachCommonTags(transposeOp, originBlockId, "VECTOR"); + attachCommonTags(reshape4Dcst, originBlockId, "VECTOR"); + attachCommonTags(reshape4DOp, originBlockId, "VECTOR"); + LOG_DEBUG("[reshape3DOp]: " << *reshape3DOp << "\n"); + LOG_DEBUG("[transposeOp]: " << *transposeOp << "\n"); + LOG_DEBUG("[reshape4DOp]: " << *reshape4DOp << "\n"); + vecValueMapping[origValue] = reshape4DOp.getResult(); +} + +// mark memref.alloc +mlir::Operation *InterCoreTransferAndSyncPass::annotateTightlyCoupledBuffer( + OpBuilder &builder, Operation *allocOp, Location loc) { + builder.setInsertionPointAfter(allocOp); + auto markAllocOp = + builder.create(loc, allocOp->getResult(0)); + auto writeAttr = builder.getStringAttr("write"); + auto readAttr = builder.getStringAttr("read"); + auto effectsAttr = builder.getArrayAttr({writeAttr, readAttr}); + markAllocOp->setAttr("effects", effectsAttr); + markAllocOp->setAttr( + hivm::HIVMTightlyCoupledBufferAttr::name, + HIVMTightlyCoupledBufferAttr::get(builder.getContext(), markAllocIndex)); + return markAllocOp; +} + +// find the insert point for memref.alloc +Operation * +InterCoreTransferAndSyncPass::findMainLoopforTransfer(Operation *endOp, + Operation *startOp) { + Operation *lca = endOp->getParentOp(); + assert(lca == startOp->getParentOp() && + "startOp and endOp are not in the same parent block, which is " + "unexpected."); + Operation *current = lca; + while (current) { + if (isa(current)) { + return current; + } + current = current->getParentOp(); + } + return nullptr; +} + +std::pair +InterCoreTransferAndSyncPass::createTransferAllocs( + OpBuilder &builder, Location loc, ArrayRef shape, Type elemType, + hivm::AddressSpace addrSpace, Operation *prodEndOp, Operation *consStartOp, + int prodBlockId, int consBlockId, StringRef prodTag, StringRef consTag, + int transferIndex) { + auto addressSpaceAttr = builder.getAttr(addrSpace); + auto allocType = MemRefType::get(shape, elemType, nullptr, addressSpaceAttr); + + Operation *prodAllocOp = nullptr; + Operation *consAllocOp = nullptr; + + Operation *mainLoopOp = findMainLoopforTransfer(prodEndOp, consStartOp); + + if (mainLoopOp) { + builder.setInsertionPoint(mainLoopOp); + prodAllocOp = builder.create(loc, allocType); + auto markProdOp = annotateTightlyCoupledBuffer(builder, prodAllocOp, loc); + consAllocOp = builder.create(loc, allocType); + auto markConsOp = annotateTightlyCoupledBuffer(builder, consAllocOp, loc); + + int loopBlockId = getSsbufferBlockId(mainLoopOp); + attachTransferTags(prodAllocOp, loopBlockId, prodTag, transferIndex); + attachTransferTags(consAllocOp, loopBlockId, consTag, transferIndex); + attachTransferTags(markProdOp, loopBlockId, prodTag, transferIndex); + attachTransferTags(markConsOp, loopBlockId, consTag, transferIndex); + + builder.setInsertionPointAfter(prodEndOp); + } else { + builder.setInsertionPoint(consStartOp); + consAllocOp = builder.create(loc, allocType); + auto markConsOp = annotateTightlyCoupledBuffer(builder, consAllocOp, loc); + + builder.setInsertionPointAfter(prodEndOp); + prodAllocOp = builder.create(loc, allocType); + auto markProdOp = annotateTightlyCoupledBuffer(builder, prodAllocOp, loc); + + attachTransferTags(prodAllocOp, prodBlockId, prodTag, transferIndex); + attachTransferTags(consAllocOp, consBlockId, consTag, transferIndex); + attachTransferTags(markProdOp, prodBlockId, prodTag, transferIndex); + attachTransferTags(markConsOp, consBlockId, consTag, transferIndex); + } + markAllocIndex++; + + return {prodAllocOp, consAllocOp}; +} + +Operation *InterCoreTransferAndSyncPass::insertVectorToCubeTransfer( + OpBuilder &builder, Value srcValue, Value normalizedValue, + Operation *vectorEndOp, Operation *cubeStartOp, Location loc, + int transferIndex, int iniConsumerId) { + LOG_DEBUG("Inserting [Vector->Cube] transfer for value: " << srcValue + << "\n"); + // Step 1: Get input information (2D tensor: MxN) + auto srcTensorType = cast(srcValue.getType()); + auto normalizedTensorType = cast(normalizedValue.getType()); + Type elemType = srcTensorType.getElementType(); + + int vecBlockId = getSsbufferBlockId(vectorEndOp); + int cubeBlockId = getSsbufferBlockId(cubeStartOp); + + auto [vecAllocOp, cubeAllocOp] = createTransferAllocs( + builder, loc, normalizedTensorType.getShape(), elemType, + hivm::AddressSpace::L1, vectorEndOp, cubeStartOp, vecBlockId, cubeBlockId, + "VECTOR", "CUBE", transferIndex); + + auto copyOp = builder.create( + loc, mlir::TypeRange{}, normalizedValue, vecAllocOp->getResult(0)); + + attachTransferTags(copyOp, vecBlockId, "VECTOR", transferIndex); + + LOG_DEBUG("[copyOp]: " << *copyOp << "\n"); + + builder.setInsertionPoint(cubeStartOp); + + auto nzLayout = + hivm::DataLayoutAttr::get(builder.getContext(), hivm::DataLayout::nZ); + auto ndLayout = + hivm::DataLayoutAttr::get(builder.getContext(), hivm::DataLayout::ND); + auto cbufaddressSpaceAttr = + builder.getAttr(hivm::AddressSpace::L1); + auto newAllocType = MemRefType::get(srcTensorType.getShape(), elemType, + nullptr, cbufaddressSpaceAttr); + auto convertLayoutOp = builder.create( + loc, newAllocType, cubeAllocOp->getResult(0), + nzLayout, // srcLayout + ndLayout // dstLayout + ); + auto plainMemrefType = MemRefType::get(srcTensorType.getShape(), elemType); + auto memspaceCastOp = builder.create( + loc, plainMemrefType, convertLayoutOp.getResult()); + auto toTensorOp = builder.create( + loc, srcTensorType, memspaceCastOp.getResult(), true, true); + + attachTransferTags(convertLayoutOp, cubeBlockId, "CUBE", transferIndex); + attachTransferTags(memspaceCastOp, cubeBlockId, "CUBE", transferIndex); + attachTransferTags(toTensorOp, cubeBlockId, "CUBE", transferIndex); + LOG_DEBUG("[toTensorOp]: " << *toTensorOp << "\n"); + + for (Operation *user : srcValue.getUsers()) { + int userBlockId = getSsbufferBlockId(user); + if (userBlockId == iniConsumerId) { + user->replaceUsesOfWith(srcValue, toTensorOp.getResult()); + } + } + return copyOp; +} + +Operation *InterCoreTransferAndSyncPass::insertCubeToVectorTransfer( + OpBuilder &builder, Value srcValue, Operation *cubeEndOp, + Operation *vectorStartOp, Location loc, int transferIndex, + int iniConsumerId) { + LOG_DEBUG("Inserting [Cube->Vector] transfer for value: " << srcValue + << "\n"); + auto srcTensorType = cast(srcValue.getType()); + int64_t M = srcTensorType.getDimSize(0); + int64_t N = srcTensorType.getDimSize(1); + Type elemType = srcTensorType.getElementType(); + + int cubeBlockId = getSsbufferBlockId(srcValue.getDefiningOp()); + int vecBlockId = getSsbufferBlockId(vectorStartOp); + + auto [cubeAllocOp, vecAllocOp] = createTransferAllocs( + builder, loc, {M, N}, elemType, hivm::AddressSpace::UB, cubeEndOp, + vectorStartOp, cubeBlockId, vecBlockId, "CUBE", "VECTOR", transferIndex); + + FixpipeDMAModeAttr dmaModeAttr = + FixpipeDMAModeAttr::get(builder.getContext(), FixpipeDMAMode::NZ2ND); + auto fixpipeOp = builder.create( + loc, mlir::TypeRange{}, // No return value + srcValue, // src + cubeAllocOp->getResult(0), // dst +#ifndef __LLVM_MAJOR_VERSION_22_COMPATIBLE__ + mlir::ValueRange{}, dmaModeAttr, nullptr, nullptr, nullptr, nullptr, + mlir::ArrayAttr{}); +#else + // This unit_flag_cond overload takes every parameter explicitly (no + // defaults). sub_block_idx and c0_pad_en were added, and quant_scale now + // precedes unit_flag_mode. + mlir::ValueRange{}, // unit_flag_cond + dmaModeAttr, // dma_mode + nullptr, // dual_dst_mode + nullptr, // sub_block_idx + nullptr, // pre_quant + nullptr, // pre_relu + nullptr, // channel_split + nullptr, // c0_pad_en + mlir::Value{}, // quant_scale + mlir::ArrayAttr{}, // unit_flag_mode + mlir::IntegerAttr{}); // unit_flag_group_id +#endif + attachTransferTags(fixpipeOp, cubeBlockId, "CUBE", transferIndex); + LOG_DEBUG("[fixpipeOp]: " << *fixpipeOp << "\n"); + + // Vector side: memspace_cast + to_tensor + builder.setInsertionPoint(vectorStartOp); + + auto plainMemrefType = MemRefType::get({M, N}, elemType); + auto memspaceCastOp = builder.create( + loc, plainMemrefType, vecAllocOp->getResult(0)); + + auto toTensorOp = builder.create( + loc, srcTensorType, memspaceCastOp.getResult(), true, true); + + attachTransferTags(memspaceCastOp, vecBlockId, "VECTOR", transferIndex); + attachTransferTags(toTensorOp, vecBlockId, "VECTOR", transferIndex); + LOG_DEBUG("[toTensorOp]: " << *toTensorOp << "\n"); + + for (Operation *user : srcValue.getUsers()) { + int userBlockId = getSsbufferBlockId(user); + if (userBlockId == iniConsumerId) { + user->replaceUsesOfWith(srcValue, toTensorOp.getResult()); + } + } + return fixpipeOp; +} + +void InterCoreTransferAndSyncPass::insertInterCoreSync( + OpBuilder &builder, Operation *transferOp, Operation *consumerStartOp, + Operation *consumerEndOp, int flag, Location loc, int transferIndex) { + LOG_DEBUG("Inserting inter-core synchronization for transferOp: " + << *transferOp << "\n"); + auto cubeCoreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + auto vecCoreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::VECTOR); + auto pipeFixAttr = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_FIX); + auto pipeVAttr = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_V); + auto pipeMte3Attr = + PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_MTE3); + auto pipeMte1Attr = + PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_MTE1); + auto pipeMAttr = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_M); + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + + int producerBlockId = getSsbufferBlockId(transferOp); + int consumerBlockId = getSsbufferBlockId(consumerStartOp); + + Operation *mainLoopOp = findMainLoopforTransfer(transferOp, consumerStartOp); + + if (dyn_cast(transferOp)) { + builder.setInsertionPointAfter(transferOp); + auto setOpForRead = builder.create( + loc, cubeCoreAttr, pipeFixAttr, pipeVAttr, flagId); + attachTransferTags(setOpForRead, producerBlockId, "CUBE", transferIndex); + builder.setInsertionPoint(consumerStartOp); + auto waitOpForRead = builder.create( + loc, vecCoreAttr, pipeFixAttr, pipeVAttr, flagId); + attachTransferTags(waitOpForRead, consumerBlockId, "VECTOR", transferIndex); + + if (mainLoopOp) { + builder.setInsertionPoint(transferOp); + auto waitOpForWrite = builder.create( + loc, cubeCoreAttr, pipeVAttr, pipeFixAttr, flagId); + attachTransferTags(waitOpForWrite, producerBlockId, "CUBE", + transferIndex); + builder.setInsertionPointAfter(consumerEndOp); + auto setOpForWrite = builder.create( + loc, vecCoreAttr, pipeVAttr, pipeFixAttr, flagId); + attachTransferTags(setOpForWrite, consumerBlockId, "VECTOR", + transferIndex); + + builder.setInsertionPoint(mainLoopOp); + auto setOpForStart = builder.create( + loc, vecCoreAttr, pipeVAttr, pipeFixAttr, flagId); + builder.setInsertionPointAfter(mainLoopOp); + auto waitOpForEnd = builder.create( + loc, cubeCoreAttr, pipeVAttr, pipeFixAttr, flagId); + + int startEndBlockId = getSsbufferBlockId(mainLoopOp); + attachTransferTags(setOpForStart, startEndBlockId, "VECTOR", + transferIndex); + attachTransferTags(waitOpForEnd, startEndBlockId, "CUBE", transferIndex); + } + return; + } else if (dyn_cast(transferOp)) { + builder.setInsertionPointAfter(transferOp); + auto setOpForRead = builder.create( + loc, vecCoreAttr, pipeMte3Attr, pipeMte1Attr, flagId); + attachTransferTags(setOpForRead, producerBlockId, "VECTOR", transferIndex); + builder.setInsertionPoint(consumerStartOp); + auto waitOpForRead = builder.create( + loc, cubeCoreAttr, pipeMte3Attr, pipeMte1Attr, flagId); + attachTransferTags(waitOpForRead, consumerBlockId, "CUBE", transferIndex); + + if (mainLoopOp) { + builder.setInsertionPoint(transferOp); + auto waitOpForWrite = builder.create( + loc, vecCoreAttr, pipeMAttr, pipeMte3Attr, flagId); + attachTransferTags(waitOpForWrite, producerBlockId, "VECTOR", + transferIndex); + + builder.setInsertionPointAfter(consumerEndOp); + auto setOpForWrite = builder.create( + loc, cubeCoreAttr, pipeMAttr, pipeMte3Attr, flagId); + attachTransferTags(setOpForWrite, consumerBlockId, "CUBE", transferIndex); + + builder.setInsertionPoint(mainLoopOp); + auto setOpForStart = builder.create( + loc, cubeCoreAttr, pipeMAttr, pipeMte3Attr, flagId); + builder.setInsertionPointAfter(mainLoopOp); + auto waitOpForEnd = builder.create( + loc, vecCoreAttr, pipeMAttr, pipeMte3Attr, flagId); + + int startEndBlockId = getSsbufferBlockId(mainLoopOp); + attachTransferTags(setOpForStart, startEndBlockId, "CUBE", transferIndex); + attachTransferTags(waitOpForEnd, startEndBlockId, "VECTOR", + transferIndex); + } + return; + } +} + +void InterCoreTransferAndSyncPass::insertPipeSSync(OpBuilder &builder, + Operation *producerOp, + Operation *consumerOp, + int flag, Location loc, + bool isCubeToVector) { + LOG_DEBUG("Inserting PIPE_S sync: " + << (isCubeToVector ? "CUBE->VECTOR" : "VECTOR->CUBE") + << ", flag = " << flag << "\n"); + + // CUBE -> VECTOR: srcPipe = PIPE_FIX, srcCoreType = CUBE, dstCoreType = + // VECTOR VECTOR -> CUBE: srcPipe = PIPE_MTE2, srcCoreType = VECTOR, + // dstCoreType = CUBE + hivm::PIPE srcPipe = + isCubeToVector ? hivm::PIPE::PIPE_FIX : hivm::PIPE::PIPE_MTE2; + hivm::TCoreType srcCoreType = + isCubeToVector ? hivm::TCoreType::CUBE : hivm::TCoreType::VECTOR; + hivm::TCoreType dstCoreType = + isCubeToVector ? hivm::TCoreType::VECTOR : hivm::TCoreType::CUBE; + hivm::PIPE dstPipe = hivm::PIPE::PIPE_S; + + auto srcCoreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), srcCoreType); + auto dstCoreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), dstCoreType); + auto srcPipeAttr = PipeAttr::get(builder.getContext(), srcPipe); + auto dstPipeAttr = PipeAttr::get(builder.getContext(), dstPipe); + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + + builder.setInsertionPointAfter(producerOp); + auto setOp = builder.create(loc, srcCoreAttr, srcPipeAttr, + dstPipeAttr, flagId); + + builder.setInsertionPoint(consumerOp); + auto waitOp = builder.create(loc, dstCoreAttr, srcPipeAttr, + dstPipeAttr, flagId); + + int prodBlockId = getSsbufferBlockId(producerOp); + int consBlockId = getSsbufferBlockId(consumerOp); + if (prodBlockId != -1) { + StringRef prodCoreType = isCubeToVector ? "CUBE" : "VECTOR"; + attachCommonTags(setOp, prodBlockId, prodCoreType); + } + if (consBlockId != -1) { + StringRef consCoreType = isCubeToVector ? "VECTOR" : "CUBE"; + attachCommonTags(waitOp, consBlockId, consCoreType); + } + + LOG_DEBUG("[PIPE_S setOp]: " << *setOp << "\n"); + LOG_DEBUG("[PIPE_S waitOp]: " << *waitOp << "\n"); +} + +// V->C Transfer Logic +LogicalResult InterCoreTransferAndSyncPass::handleVectorToCube( + OpBuilder &builder, DependencyInfo &dep, + llvm::DenseMap vecvalueMapping, + llvm::DenseMap cubeValueMapping, + FlagIdManager &flagManager) { + mlir::Value srcValue = dep.value; + auto it = cubeValueMapping.find(srcValue); + if (it != cubeValueMapping.end()) { + srcValue = it->second; + } + Location loc = dep.value.getLoc(); + // Step 1: Shape normalization (automatically insert slice) + Value normalizedVal = vecvalueMapping[dep.value]; + + // Get start/end operations for V/C blocks + auto [prodStart, prodEnd] = getBlockStartEnd(dep.producerBlockId, module); + auto [consStart, consEnd] = getBlockStartEnd(dep.consumerBlockId, module); + + Operation *transferOp = insertVectorToCubeTransfer( + builder, srcValue, normalizedVal, prodEnd, consStart, loc, transferIndex, + dep.iniConsumerBlockId); + + int flagId = flagManager.acquireId(prodStart); + assert(flagId != -1 && "No available flag_id!!"); + auto [newProdStart, newProdEnd] = + getBlockStartEnd(dep.producerBlockId, module); + auto [newConsStart, newConsEnd] = + getBlockStartEnd(dep.consumerBlockId, module); + insertInterCoreSync(builder, transferOp, newConsStart, newConsEnd, flagId, + loc, transferIndex); + + transferIndex++; + LOG_DEBUG("Inserted V->C transfer and sync: block " + << dep.producerBlockId << " -> block " << dep.consumerBlockId + << "\n"); + return success(); +} + +// C->V Transfer Logic +LogicalResult InterCoreTransferAndSyncPass::handleCubeToVector( + OpBuilder &builder, DependencyInfo &dep, + llvm::DenseMap cubeValueMapping, + FlagIdManager &flagManager) { + mlir::Value srcValue = dep.value; + auto it = cubeValueMapping.find(srcValue); + if (it != cubeValueMapping.end()) { + srcValue = it->second; + } + Location loc = srcValue.getLoc(); + auto [prodStart, prodEnd] = + getBlockStartEnd(dep.producerBlockId, module); // C Block + auto [consStart, consEnd] = + getBlockStartEnd(dep.consumerBlockId, module); // V Block + LOG_DEBUG("[newProdStart]" << *prodStart << "\n"); + LOG_DEBUG("[newProdEnd]" << *prodEnd << "\n"); + LOG_DEBUG("[newConsStart]" << *consStart << "\n"); + LOG_DEBUG("[newConsEnd]" << *consEnd << "\n"); + Operation *transferOp = + insertCubeToVectorTransfer(builder, srcValue, prodEnd, consStart, loc, + transferIndex, dep.iniConsumerBlockId); + + auto [newProdStart, newProdEnd] = + getBlockStartEnd(dep.producerBlockId, module); // C Block + auto [newConsStart, newConsEnd] = + getBlockStartEnd(dep.consumerBlockId, module); // V Block + int flagId = flagManager.acquireId(newProdStart); + assert(flagId != -1 && "No available flag_id!!"); + insertInterCoreSync(builder, transferOp, newConsStart, newConsEnd, flagId, + loc, transferIndex); + + transferIndex++; + LOG_DEBUG("Inserted C->V transfer and sync: block " + << dep.producerBlockId << " -> block " << dep.consumerBlockId + << "\n"); + return success(); +} + +// PIPE_S Memory Dependency +LogicalResult InterCoreTransferAndSyncPass::handleMemoryDependency( + OpBuilder &builder, DependencyInfo &dep, size_t depIndex, + llvm::SmallVector memDependencies, + FlagIdManager &flagManager) { + LOG_DEBUG("Handling PIPE_S memory dependency...\n"); + + // Get producer and consumer block start/end operations + auto [prodStart, prodEnd] = getBlockStartEnd(dep.producerBlockId, module); + auto [consStart, consEnd] = getBlockStartEnd(dep.consumerBlockId, module); + + if (!prodStart || !prodEnd || !consStart || !consEnd) { + LOG_DEBUG("[ERROR] Failed to get block start/end operations.\n"); + return failure(); + } + + if (isOuterLayerDependency(depIndex, prodEnd, consStart, memDependencies)) { + LOG_DEBUG("[PIPE_S] Skipping outer layer dependency: block " + << dep.producerBlockId << " -> block " << dep.consumerBlockId + << "\n"); + return success(); + } + + // Get flag ID + int flagId = flagManager.acquireId(prodStart); + assert(flagId != -1 && "No available flag_id!!"); + + // Determine sync direction: CUBE->VECTOR or VECTOR->CUBE + bool isCubeToVector = (dep.type == DependencyType::CubeToVector); + + // Get location info + Location loc = prodEnd->getLoc(); + + // Insert PIPE_S sync + insertPipeSSync(builder, prodEnd, consStart, flagId, loc, isCubeToVector); + + transferIndex++; + + LOG_DEBUG("Inserted PIPE_S sync: block " + << dep.producerBlockId << " -> block " << dep.consumerBlockId + << ", flagId = " << flagId << "\n"); + + return success(); +} + +// Main Processing +LogicalResult +InterCoreTransferAndSyncPass::processDependencies(FlagIdManager &flagManager) { + LOG_DEBUG("Starting InterCoreTransferAndSyncPass processDependencies...\n"); + OpBuilder builder(module.getContext()); + + auto &info = getAnalysis(); + if (!info.isValid()) { + LOG_DEBUG("Error: Data dependency analysis failed.\n"); + return failure(); + } + + llvm::SmallVector &V2CDependencies = + info.getV2CDependencies(); + LOG_DEBUG("[DEBUG] V2CDependencies size: " << V2CDependencies.size() << "\n"); + for (size_t i = 0; i < V2CDependencies.size(); ++i) { + auto &dep = V2CDependencies[i]; + LOG_DEBUG("[V2C-" << i << "] producerBlockId = " << dep.producerBlockId + << ", consumerBlockId = " << dep.consumerBlockId + << ", iniProducerBlockId = " << dep.iniProducerBlockId + << ", iniConsumerBlockId = " << dep.iniConsumerBlockId + << ", value = " << dep.value << "\n"); + } + // Step 1: Handle V->C dependencies + for (auto &dep : V2CDependencies) { + Location loc = dep.value.getLoc(); + Nd2NzNormalize(builder, dep, loc); + } + llvm::DenseMap vecvalueMapping = + getVecValueMapping(); + llvm::DenseMap cubevalueMapping = + getCubeValueMapping(); + for (auto &dep : V2CDependencies) { + LOG_DEBUG("[V->C] producerBlockId = " << dep.producerBlockId + << ", consumerBlockId = " + << dep.consumerBlockId << "\n"); + if (failed(handleVectorToCube(builder, dep, vecvalueMapping, + cubevalueMapping, flagManager))) { + LOG_DEBUG("[ERROR] V->C failed! producerBlockId = " + << dep.producerBlockId + << ", consumerBlockId = " << dep.consumerBlockId << "\n"); + return failure(); + } + } + LOG_DEBUG("Completed V->C transfers and syncs.\n"); + + llvm::SmallVector &C2VDependencies = + info.getC2VDependencies(); + LOG_DEBUG("[DEBUG] C2VDependencies size: " << C2VDependencies.size() << "\n"); + // Step 2: Handle C->V dependencies + for (auto &dep : C2VDependencies) { + LOG_DEBUG("[C->V] producerBlockId = " << dep.producerBlockId + << ", consumerBlockId = " + << dep.consumerBlockId << "\n"); + if (failed( + handleCubeToVector(builder, dep, cubevalueMapping, flagManager))) { + LOG_DEBUG("[ERROR] C->V failed! producerBlockId = " + << dep.producerBlockId + << ", consumerBlockId = " << dep.consumerBlockId << "\n"); + return failure(); + } + } + LOG_DEBUG("Completed C->V transfers and syncs.\n"); + LOG_DEBUG("=====================================================\n"); + LOG_DEBUG("InterCoreTransferAndSyncPass success!\n"); + + return success(); +} + +// Declare dependent dialects +void InterCoreTransferAndSyncPass::getDependentDialects( + DialectRegistry ®istry) const { + registry.insert(); +} + +// Pass Entry Point +void InterCoreTransferAndSyncPass::runOnOperation() { + LOG_DEBUG("\n--- enter InterCoreTransferAndSyncPass --->\n"); + module = getOperation(); + + // Phase 1: Initialize FlagIdManager as local variable + FlagIdManager flagManager(module); + + // Phase 2: Execute transfer and sync insertion + if (failed(processDependencies(flagManager))) { + signalPassFailure(); + LOG_DEBUG("Error: Inter-core transfer and sync failed.\n"); + return; + } + + LOG_DEBUG("Module after InterCoreTransferAndSyncPass:\n" << module << "\n"); + + LOG_DEBUG("--- exit InterCoreTransferAndSyncPass --->\n"); +} + +// Create the pass +namespace mlir { +namespace triton { +std::unique_ptr> createInterCoreTransferAndSyncPass() { + return std::make_unique(); +} +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/MarkMainLoop.cpp b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/MarkMainLoop.cpp new file mode 100644 index 0000000000..e478855a21 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/MarkMainLoop.cpp @@ -0,0 +1,96 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/SplitDataflow/MarkMainLoop.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "mlir/IR/Operation.h" +#include "llvm/Support/Debug.h" + +using namespace mlir; + +static constexpr const char *DEBUG_TYPE = "mark-main-loop"; +#define LOG_DEBUG(...) \ + LLVM_DEBUG(llvm::dbgs() << " [" << DEBUG_TYPE << "] " << __VA_ARGS__) + +using namespace mlir::triton; + +// Pass Entry Point +void MarkMainLoopPass::runOnOperation() { + LOG_DEBUG("\n--- enter MarkMainLoopPass --->\n"); + ModuleOp module = getOperation(); + + int mainLoopIdCounter = 0; + SmallVector mainLoops; + + // Find all candidate main loops + module.walk([&](Operation *op) { + if (isa(op)) { + if (auto forOp = op->getParentOfType()) { + mainLoops.push_back(forOp); + } + } + }); + + for (scf::ForOp forOp : mainLoops) { + if (!forOp->hasAttr("ssbuffer.main_loop")) { + // Add attribute with integer value (current counter ID) + forOp->setAttr( + "ssbuffer.main_loop", + Builder(module.getContext()).getI32IntegerAttr(mainLoopIdCounter)); + mainLoopIdCounter++; + } + } + + // Remove main_loop attribute from outer loops if nested loops both have it + // Keep only the innermost main_loop + SmallVector allMainLoops; + module.walk([&](scf::ForOp forOp) { + if (forOp->hasAttr("ssbuffer.main_loop")) { + allMainLoops.push_back(forOp); + } + }); + + for (scf::ForOp forOp : allMainLoops) { + // Check if there's any nested for loop with main_loop attribute + bool hasNestedMainLoop = false; + forOp.walk([&](scf::ForOp nestedForOp) { + if (nestedForOp != forOp && nestedForOp->hasAttr("ssbuffer.main_loop")) { + hasNestedMainLoop = true; + } + }); + // Remove attribute from outer loop if inner loop also has it + if (hasNestedMainLoop) { + forOp->removeAttr("ssbuffer.main_loop"); + } + } + + LOG_DEBUG("--- exit MarkMainLoopPass --->\n"); +} + +// Create the pass +namespace mlir { +namespace triton { +std::unique_ptr> createMarkMainLoopPass() { + return std::make_unique(); +} +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/PreserveControlAttrsCanonicalize.cpp b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/PreserveControlAttrsCanonicalize.cpp new file mode 100644 index 0000000000..d0dea9b487 --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/PreserveControlAttrsCanonicalize.cpp @@ -0,0 +1,164 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/SplitDataflow/PreserveControlAttrsCanonicalize.h" + +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/PassRegistry.h" +#include "mlir/Rewrite/FrozenRewritePatternSet.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "llvm/ADT/SetVector.h" +#include "llvm/Support/Debug.h" + +using namespace mlir; + +static constexpr const char *DEBUG_TYPE = "PreserveControlAttrsCanonicalize"; +#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ") + +template static void logDebug(const Args &...args) { + LLVM_DEBUG({ + auto &debugStream = llvm::dbgs(); + debugStream << '[' << DEBUG_TYPE << "] "; + (debugStream << ... << args); + debugStream << "\n"; + }); +} + +namespace { + +static bool isTrackedControlFlowOp(Operation *op) { + return isa(op); +} + +static bool canTransferAttrs(Operation *from, Operation *to) { + return from && to && from != to && isTrackedControlFlowOp(from) && + isTrackedControlFlowOp(to) && from->getName() == to->getName(); +} + +class PreserveControlAttrsListener : public RewriterBase::Listener { +public: + void notifyOperationInserted(Operation *op, OpBuilder::InsertPoint) override { + recentInserts.insert(op); + } + + void notifyOperationErased(Operation *op) override { + recentInserts.remove(op); + } + + void notifyOperationReplaced(Operation *op, Operation *newOp) override { + transferAttrs(op, newOp); + } + + void notifyOperationReplaced(Operation *op, ValueRange values) override { + if (Operation *newOp = findReplacementOp(op, values)) { + transferAttrs(op, newOp); + } + } + +private: + Operation *findReplacementOp(Operation *oldOp, + ValueRange replacements) const { + if (!isTrackedControlFlowOp(oldOp)) { + return nullptr; + } + + for (Value value : replacements) { + Operation *defOp = value.getDefiningOp(); + if (defOp && recentInserts.contains(defOp) && + canTransferAttrs(oldOp, defOp)) { + return defOp; + } + } + + for (Operation *candidate : llvm::reverse(recentInserts.getArrayRef())) { + if (canTransferAttrs(oldOp, candidate)) { + return candidate; + } + } + return nullptr; + } + + static void transferAttrs(Operation *from, Operation *to) { + if (!canTransferAttrs(from, to)) { + return; + } + + for (NamedAttribute attr : from->getAttrs()) { + if (to->hasAttr(attr.getName())) { + continue; + } + to->setAttr(attr.getName(), attr.getValue()); + logDebug("carry attr '", attr.getName(), "' from ", + from->getName().getStringRef(), " to ", + to->getName().getStringRef()); + } + } + + llvm::SetVector recentInserts; +}; + +static void populateCanonicalizationPatterns(MLIRContext *ctx, + RewritePatternSet &patterns) { + for (Dialect *dialect : ctx->getLoadedDialects()) { + dialect->getCanonicalizationPatterns(patterns); + } + + for (RegisteredOperationName opName : ctx->getRegisteredOperations()) { + opName.getCanonicalizationPatterns(patterns, ctx); + } +} + +} // namespace + +void mlir::triton::PreserveControlAttrsCanonicalizePass::runOnOperation() { + RewritePatternSet patterns(&getContext()); + populateCanonicalizationPatterns(&getContext(), patterns); + + PreserveControlAttrsListener listener; + GreedyRewriteConfig config; + config.setListener(&listener); + + if (failed(applyPatternsGreedily(getOperation(), + FrozenRewritePatternSet(std::move(patterns)), + config))) { + getOperation()->emitError("PreserveControlAttrsCanonicalizePass failed"); + signalPassFailure(); + } +} + +namespace mlir { +namespace triton { + +std::unique_ptr> +createPreserveControlAttrsCanonicalizePass() { + return std::make_unique(); +} + +void registerPreserveControlAttrsCanonicalizePasses() { + registerPass([]() -> std::unique_ptr { + return createPreserveControlAttrsCanonicalizePass(); + }); +} + +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/SeparateCVScope.cpp b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/SeparateCVScope.cpp new file mode 100644 index 0000000000..ec9b40225a --- /dev/null +++ b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/SeparateCVScope.cpp @@ -0,0 +1,816 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "ascend/include/DynamicCVPipeline/SplitDataflow/SeparateCVScope.h" + +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/Scope/IR/Scope.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/IRMapping.h" +#include "mlir/Pass/Pass.h" + +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Debug.h" +#include + +using namespace mlir; + +using namespace mlir::triton; + +static constexpr const char *DEBUG_TYPE = "SeparateCVScope"; +#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ") + +template static void logDebug(const Args &...args) { + LLVM_DEBUG({ + auto &debugStream = llvm::dbgs(); + debugStream << '[' << DEBUG_TYPE << "] "; + (debugStream << ... << args); + debugStream << "\n"; + }); +} + +static constexpr unsigned kForOpOperandPrefixCount = 3; +static constexpr unsigned kSeenValuesCapacity = 16; + +static void debugDumpOperation(StringRef prefix, Operation *op) { + LLVM_DEBUG({ + DBGS() << prefix << "\n"; + op->print(llvm::dbgs()); + llvm::dbgs() << "\n"; + }); +} + +struct CoreTypeInfo { + SmallVector resultTypes; + + StringRef getResultType(size_t index) const { + if (index < resultTypes.size()) { + return resultTypes[index]; + } + return resultTypes.front(); + } + + bool hasResultForScope(StringRef scopeType) const { + for (StringRef t : resultTypes) { + if (t == scopeType) { + return true; + } + } + return false; + } + + bool allResultsMatchScope(StringRef scopeType, unsigned numResults) const { + if (numResults == 0) { + return getResultType(0) == scopeType; + } + for (unsigned i = 0; i < numResults; ++i) { + if (getResultType(i) != scopeType) { + return false; + } + } + return true; + } +}; + +static std::optional parseCoreTypeInfo(Operation *op) { + auto attr = op->getAttrOfType("ssbuffer.core_type"); + if (!attr) { + return std::nullopt; + } + + CoreTypeInfo info; + StringRef raw = attr.getValue(); + SmallVector parts; + raw.split(parts, ',', -1, false); + + if (parts.empty()) { + return std::nullopt; + } + + for (auto &part : parts) { + part = part.trim(); + } + + info.resultTypes.assign(parts.begin(), parts.end()); + return info; +} + +struct ScopeMatchInfo { + bool matches = false; + bool allResultsMatch = false; +}; + +static ScopeMatchInfo getScopeMatchInfo(Operation *op, StringRef scopeType) { + auto info = parseCoreTypeInfo(op); + if (!info) { + return {}; + } + + unsigned numResults = op->getNumResults(); + bool matches = numResults == 0 ? info->getResultType(0) == scopeType + : info->hasResultForScope(scopeType); + return {matches, info->allResultsMatchScope(scopeType, numResults)}; +} + +static bool matchesScope(Operation *op, StringRef scopeType) { + return getScopeMatchInfo(op, scopeType).matches; +} + +static bool isOpAlive(Operation *op) { return op && op->getBlock(); } + +static Value buildNeutralValue(OpBuilder &builder, Value oldOperand, + Location loc, StringRef scopeType) { + if (!oldOperand) { + return Value(); + } + + Type type = oldOperand.getType(); + Operation *createdOp = nullptr; + + if (auto memrefTy = dyn_cast(type)) { + createdOp = builder.create(loc, memrefTy); + } else if (auto shapedTy = dyn_cast(type)) { + if (shapedTy.hasStaticShape() && !isa(type)) { + if (Attribute elemZero = builder.getZeroAttr(shapedTy.getElementType())) { + auto zeroDense = DenseElementsAttr::get(shapedTy, elemZero); + auto typedAttr = mlir::cast(zeroDense); + createdOp = builder.create(loc, type, typedAttr); + } + } + } else if (Attribute zeroAttr = builder.getZeroAttr(type)) { + if (auto typedZero = dyn_cast(zeroAttr)) { + createdOp = builder.create(loc, type, typedZero); + } + } + + return createdOp ? createdOp->getResult(0) : Value(); +} + +static FailureOr> +createTwoFullScopes(func::FuncOp funcOp) { + if (!llvm::hasSingleElement(funcOp.getBody())) { + logDebug("createTwoFullScopes failed for func '", funcOp.getName(), + "': expected single block"); + funcOp.emitError("SeparateCVScope only supports single-block functions"); + return failure(); + } + + Block &lastBlock = funcOp.getBody().back(); + SmallVector opsToMove; + for (Operation &op : lastBlock.without_terminator()) { + opsToMove.push_back(&op); + } + + if (opsToMove.empty()) { + return std::make_pair(scope::ScopeOp(), scope::ScopeOp()); + } + + Operation *lastOpToMove = opsToMove.back(); + OpBuilder builder(&lastBlock, ++lastOpToMove->getIterator()); + + auto vecScope = + builder.create(builder.getUnknownLoc(), ArrayRef{}); + vecScope.getBodyRegion().emplaceBlock(); + + Block *vecBlock = &vecScope.getBodyRegion().front(); + OpBuilder vecBuilder(vecBlock, vecBlock->end()); + + for (Operation *op : opsToMove) { + op->remove(); + vecBuilder.insert(op); + } + vecBuilder.create(builder.getUnknownLoc()); + + builder.setInsertionPointAfter(vecScope); + IRMapping mapping; + auto cloned = builder.clone(*vecScope.getOperation(), mapping); + auto cubeScope = cast(cloned); + + auto vecAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::VECTOR); + auto cubeAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + + vecScope->setAttr(hivm::TCoreTypeAttr::name, vecAttr); + cubeScope->setAttr(hivm::TCoreTypeAttr::name, cubeAttr); + + return std::make_pair(vecScope, cubeScope); +} + +static bool isPureForScope(Operation *op, StringRef scopeType) { + ScopeMatchInfo matchInfo = getScopeMatchInfo(op, scopeType); + if (!matchInfo.matches || !matchInfo.allResultsMatch) { + return false; + } + + for (Region ®ion : op->getRegions()) { + for (Block &block : region) { + for (Operation &nested : block) { + if (nested.hasTrait()) { + continue; + } + if (nested.getNumRegions() > 0) { + if (!isPureForScope(&nested, scopeType)) { + return false; + } + continue; + } + if (!getScopeMatchInfo(&nested, scopeType).matches) { + return false; + } + } + } + } + return true; +} + +enum class PendingActionKind { EraseDirectly, NormalizeControlFlow }; + +struct PendingAction { + Operation *op = nullptr; + PendingActionKind kind; +}; + +static bool isControlFlowOp(Operation *op) { + return isa(op); +} + +static void collectActionForOp(Operation *op, StringRef scopeType, + SmallVector &actions) { + ScopeMatchInfo matchInfo = getScopeMatchInfo(op, scopeType); + if (!matchInfo.matches) { + actions.push_back({op, PendingActionKind::EraseDirectly}); + } else if (!matchInfo.allResultsMatch) { + actions.push_back({op, PendingActionKind::NormalizeControlFlow}); + } +} + +static SmallVector collectActionsInRegion(Region ®ion, + StringRef scopeType) { + SmallVector actions; + for (Block &block : region) { + for (Operation &op : block) { + if (op.hasTrait()) + continue; + + if (op.getNumRegions() > 0) { + if (isControlFlowOp(&op)) { + if (!isPureForScope(&op, scopeType)) + actions.push_back({&op, PendingActionKind::NormalizeControlFlow}); + } else if (!matchesScope(&op, scopeType)) { + actions.push_back({&op, PendingActionKind::EraseDirectly}); + } + continue; + } + collectActionForOp(&op, scopeType, actions); + } + } + return actions; +} + +static bool hasLiveUsers(Operation *op) { + for (Value result : op->getResults()) { + for (OpOperand &use : result.getUses()) { + if (isOpAlive(use.getOwner())) { + return true; + } + } + } + return false; +} + +static bool needsLoopCarryPreserve(Operation *owner, unsigned slotIndex, + StringRef scopeType); + +static bool canSkipForwardingUse(OpOperand &use, StringRef scopeType) { + Operation *user = use.getOwner(); + auto info = parseCoreTypeInfo(user); + if (!info) { + return false; + } + + if (auto forOp = dyn_cast(user)) { + unsigned operandIndex = use.getOperandNumber(); + if (operandIndex < kForOpOperandPrefixCount) { + return false; + } + unsigned resultIndex = operandIndex - kForOpOperandPrefixCount; + return resultIndex < forOp.getNumResults() && + info->getResultType(resultIndex) != scopeType; + } + + if (auto whileOp = dyn_cast(user)) { + unsigned slotIdx = use.getOperandNumber(); + if (slotIdx >= whileOp.getNumResults()) { + return false; + } + if (info->getResultType(slotIdx) == scopeType) { + return false; + } + if (needsLoopCarryPreserve(whileOp, slotIdx, scopeType)) { + return false; + } + return true; + } + return false; +} + +static Operation *findLiveUser(Value value, StringRef scopeType) { + for (OpOperand &use : value.getUses()) { + Operation *user = use.getOwner(); + if (isOpAlive(user) && !canSkipForwardingUse(use, scopeType)) { + return user; + } + } + return nullptr; +} + +static Operation *findNonTermUser(Value value, StringRef scopeType) { + for (OpOperand &use : value.getUses()) { + Operation *user = use.getOwner(); + if (isOpAlive(user) && !canSkipForwardingUse(use, scopeType) && + !user->hasTrait()) { + return user; + } + } + return nullptr; +} + +enum class UseCheckResult { Skip, Active, Continue }; + +static UseCheckResult checkConditionUse(OpOperand &use, Operation *owner, + unsigned slotIndex, + StringRef scopeType) { + auto conditionOp = dyn_cast(use.getOwner()); + if (!conditionOp) { + return UseCheckResult::Skip; + } + + unsigned idx = use.getOperandNumber(); + if (conditionOp->getParentOp() != owner || idx == 0 || idx - 1 != slotIndex) { + Operation *parentOp = conditionOp->getParentOp(); + if (parentOp && !matchesScope(parentOp, scopeType)) { + return UseCheckResult::Continue; + } + return UseCheckResult::Active; + } + return UseCheckResult::Continue; +} + +static UseCheckResult checkYieldUse(OpOperand &use, Operation *owner, + unsigned slotIndex, StringRef scopeType, + SmallVector &worklist) { + auto yieldOp = dyn_cast(use.getOwner()); + if (!yieldOp) { + return UseCheckResult::Skip; + } + + Operation *yieldOwner = yieldOp->getParentOp(); + unsigned idx = use.getOperandNumber(); + + if (yieldOwner == owner) { + if (idx != slotIndex && idx < yieldOwner->getNumResults()) { + worklist.push_back(yieldOwner->getResult(idx)); + } + return UseCheckResult::Continue; + } + + if (idx < yieldOwner->getNumResults()) { + worklist.push_back(yieldOwner->getResult(idx)); + return UseCheckResult::Continue; + } + + if (matchesScope(yieldOwner, scopeType) || isControlFlowOp(yieldOwner)) { + return UseCheckResult::Active; + } + return UseCheckResult::Continue; +} + +static UseCheckResult checkGeneralUse(OpOperand &use, StringRef scopeType, + SmallVector &worklist) { + Operation *user = use.getOwner(); + if (user->getNumResults() == 0) { + if (matchesScope(user, scopeType) || isControlFlowOp(user)) { + return UseCheckResult::Active; + } + return UseCheckResult::Continue; + } + for (Value result : user->getResults()) { + worklist.push_back(result); + } + return UseCheckResult::Continue; +} + +static bool slotHasActiveUse(Value startValue, Operation *owner, + unsigned slotIndex, StringRef scopeType) { + SmallVector worklist{startValue}; + llvm::SmallPtrSet seenValues; + + while (!worklist.empty()) { + Value value = worklist.pop_back_val(); + if (!value || !seenValues.insert(value.getAsOpaquePointer()).second) { + continue; + } + + for (OpOperand &use : value.getUses()) { + if (!isOpAlive(use.getOwner())) { + continue; + } + + UseCheckResult result; + if ((result = checkConditionUse(use, owner, slotIndex, scopeType)) != + UseCheckResult::Skip || + (result = checkYieldUse(use, owner, slotIndex, scopeType, + worklist)) != UseCheckResult::Skip || + (result = checkGeneralUse(use, scopeType, worklist)) != + UseCheckResult::Skip) { + if (result == UseCheckResult::Active) { + return true; + } + } + } + } + return false; +} + +static bool needsLoopCarryPreserve(Operation *owner, unsigned slotIndex, + StringRef scopeType) { + if (auto forOp = dyn_cast(owner)) { + Block &body = forOp.getRegion().front(); + unsigned iterArgIndex = slotIndex + 1; + if (iterArgIndex >= body.getNumArguments()) { + return false; + } + return slotHasActiveUse(body.getArgument(iterArgIndex), owner, slotIndex, + scopeType); + } + + if (auto whileOp = dyn_cast(owner)) { + Block &before = whileOp.getBefore().front(); + Block &after = whileOp.getAfter().front(); + if (slotIndex < before.getNumArguments() && + slotHasActiveUse(before.getArgument(slotIndex), owner, slotIndex, + scopeType)) { + return true; + } + if (slotIndex < after.getNumArguments() && + slotHasActiveUse(after.getArgument(slotIndex), owner, slotIndex, + scopeType)) { + return true; + } + } + return false; +} + +static LogicalResult neutralizeYieldInRegion(Operation *op, + const CoreTypeInfo &info, + StringRef scopeType, + Location loc) { + if (op->getNumRegions() == 0) { + return success(); + } + + for (Region ®ion : op->getRegions()) { + if (region.empty()) { + continue; + } + + for (Block &block : region) { + auto yieldOp = dyn_cast(block.getTerminator()); + if (!yieldOp) { + continue; + } + + OpBuilder builder(yieldOp); + for (unsigned i = 0; i < yieldOp.getNumOperands(); ++i) { + if (info.getResultType(i) == scopeType) { + continue; + } + if (needsLoopCarryPreserve(op, i, scopeType)) { + continue; + } + + Value oldOperand = yieldOp.getOperand(i); + if (i < op->getNumResults()) { + if (Operation *resultUser = + findLiveUser(op->getResult(i), scopeType)) { + yieldOp.emitWarning() + << "skip neutralizing yield operand #" << i << " for scope " + << scopeType << " because parent result #" << i + << " still has live user '" + << resultUser->getName().getStringRef() << "'"; + continue; + } + } + + Value replacement = + buildNeutralValue(builder, oldOperand, loc, scopeType); + if (!replacement) { + logDebug("neutralizeYieldInRegion failed for op '", + op->getName().getStringRef(), "' at operand #", i, + " in scope ", scopeType); + return failure(); + } + yieldOp.setOperand(i, replacement); + } + } + } + return success(); +} + +static LogicalResult neutralizeTerminatorUses(Operation *op, + const CoreTypeInfo &info, + StringRef scopeType) { + for (unsigned i = 0; i < op->getNumResults(); ++i) { + if (info.getResultType(i) == scopeType) { + continue; + } + + Value result = op->getResult(i); + if (Operation *extraUser = findNonTermUser(result, scopeType)) { + op->emitWarning() << "skip neutralizing result #" << i << " for scope " + << scopeType + << " because the value still has live user '" + << extraUser->getName().getStringRef() << "'"; + continue; + } + + SmallVector usesToNeutralize; + for (OpOperand &use : llvm::make_early_inc_range(result.getUses())) { + if (use.getOwner()->hasTrait()) { + usesToNeutralize.push_back(&use); + } + } + + for (OpOperand *use : usesToNeutralize) { + OpBuilder builder(use->getOwner()); + Value replacement = + buildNeutralValue(builder, result, op->getLoc(), scopeType); + if (!replacement) { + logDebug("neutralizeTerminatorUses failed for op '", + op->getName().getStringRef(), "' at result #", i, " in scope ", + scopeType); + return failure(); + } + use->set(replacement); + } + } + return success(); +} + +static LogicalResult executeActions(SmallVector &actions, + StringRef scopeType); + +static LogicalResult normalizeRegionOp(Operation *op, StringRef scopeType) { + auto infoOpt = parseCoreTypeInfo(op); + if (!infoOpt) { + for (Region ®ion : op->getRegions()) { + SmallVector nestedActions = + collectActionsInRegion(region, scopeType); + if (failed(executeActions(nestedActions, scopeType))) { + return failure(); + } + } + return success(); + } + + auto info = *infoOpt; + Location loc = op->getLoc(); + + debugDumpOperation("before normalizeRegionOp", op); + + if (op->getNumRegions() > 0) { + if (failed(neutralizeYieldInRegion(op, info, scopeType, loc))) { + logDebug("normalizeRegionOp failed while neutralizing yields for op '", + op->getName().getStringRef(), "' in scope ", scopeType); + return failure(); + } + + for (Region ®ion : op->getRegions()) { + SmallVector nestedActions = + collectActionsInRegion(region, scopeType); + if (failed(executeActions(nestedActions, scopeType))) { + logDebug( + "normalizeRegionOp failed while executing nested actions for op '", + op->getName().getStringRef(), "' in scope ", scopeType); + return failure(); + } + } + } + + debugDumpOperation("after normalizeRegionOp", op); + return success(); +} + +static LogicalResult normalizeNonRegionOp(Operation *op, StringRef scopeType) { + auto infoOpt = parseCoreTypeInfo(op); + if (!infoOpt) { + logDebug("normalizeNonRegionOp failed for op '", + op->getName().getStringRef(), + "': missing ssbuffer.core_type in scope ", scopeType); + op->emitError("missing ssbuffer.core_type"); + return failure(); + } + + debugDumpOperation("before normalizeNonRegionOp", op); + + LogicalResult result = neutralizeTerminatorUses(op, *infoOpt, scopeType); + if (succeeded(result)) { + debugDumpOperation("after normalizeNonRegionOp", op); + } + return result; +} + +static bool regionHasScopeContent(Region ®ion, StringRef scopeType) { + for (Block &block : region) { + for (Operation &op : block) { + if (op.hasTrait()) { + continue; + } + + if (op.getNumRegions() > 0) { + if (!isControlFlowOp(&op) && matchesScope(&op, scopeType)) { + return true; + } + + for (Region &subRegion : op.getRegions()) { + if (regionHasScopeContent(subRegion, scopeType)) { + return true; + } + } + continue; + } + + if (matchesScope(&op, scopeType)) { + return true; + } + } + } + return false; +} + +static bool isNormalizedDeadShell(Operation *op, StringRef scopeType) { + if (!op || op->getNumRegions() == 0) { + return false; + } + + for (Value result : op->getResults()) { + if (findLiveUser(result, scopeType)) { + return false; + } + } + + for (Region ®ion : op->getRegions()) { + if (regionHasScopeContent(region, scopeType)) { + return false; + } + } + return true; +} + +static LogicalResult executeActions(SmallVector &actions, + StringRef scopeType) { + for (auto it = actions.rbegin(); it != actions.rend(); ++it) { + Operation *op = it->op; + if (!isOpAlive(op)) { + continue; + } + + switch (it->kind) { + case PendingActionKind::EraseDirectly: + if (!hasLiveUsers(op)) { + op->erase(); + } + break; + case PendingActionKind::NormalizeControlFlow: + if (op->getNumRegions() > 0) { + if (failed(normalizeRegionOp(op, scopeType))) { + logDebug("failed to normalize region op '", + op->getName().getStringRef(), "' for scope ", scopeType); + return failure(); + } + } else { + if (failed(normalizeNonRegionOp(op, scopeType))) { + logDebug("failed to normalize non-region op '", + op->getName().getStringRef(), "' for scope ", scopeType); + return failure(); + } + } + if (isOpAlive(op) && op->getNumRegions() > 0 && + isNormalizedDeadShell(op, scopeType)) { + logDebug("erasing dead shell after normalize: '", + op->getName().getStringRef(), "' in scope ", scopeType); + op->erase(); + } + break; + } + } + return success(); +} + +static void removeSsbufferAttrs(Operation *op) { + op->removeAttr("ssbuffer.core_type"); +} + +static void cleanupSsbufferAttrs(Operation *rootOp) { + removeSsbufferAttrs(rootOp); + rootOp->walk([](Operation *op) { removeSsbufferAttrs(op); }); +} + +static LogicalResult separateScopes(func::FuncOp funcOp) { + debugDumpOperation("before SeparateCVScope on func", funcOp.getOperation()); + + FailureOr> scopes = + createTwoFullScopes(funcOp); + if (failed(scopes)) { + logDebug("SeparateCVScope failed to create VECTOR/CUBE scopes for func '", + funcOp.getName(), "'"); + return failure(); + } + + auto vecScope = scopes->first; + auto cubeScope = scopes->second; + + auto vecActions = collectActionsInRegion(vecScope.getRegion(), "VECTOR"); + auto cubeActions = collectActionsInRegion(cubeScope.getRegion(), "CUBE"); + if (failed(executeActions(vecActions, "VECTOR")) || + failed(executeActions(cubeActions, "CUBE"))) { + logDebug("SeparateCVScope failed while executing actions for func '", + funcOp.getName(), "'"); + return failure(); + } + + cleanupSsbufferAttrs(funcOp); + + debugDumpOperation("after SeparateCVScope on func", funcOp.getOperation()); + return success(); +} + +// Declare dependent dialects +void mlir::triton::SeparateCVScopePass::getDependentDialects( + DialectRegistry ®istry) const { + registry.insert(); +} + +void mlir::triton::SeparateCVScopePass::runOnOperation() { + auto module = getOperation(); + SmallVector funcOps; + module.walk([&](func::FuncOp funcOp) { funcOps.push_back(funcOp); }); + + debugDumpOperation("before SeparateCVScopePass", module.getOperation()); + + for (auto funcOp : funcOps) { + if (funcOp.getBody().empty()) { + continue; + } + if (failed(separateScopes(funcOp))) { + logDebug("SeparateCVScopePass failed on func '", funcOp.getName(), "'"); + signalPassFailure(); + return; + } + } + + debugDumpOperation("after SeparateCVScopePass", module.getOperation()); +} + +namespace mlir { +namespace triton { + +std::unique_ptr> createSeparateCVScopePass() { + return std::make_unique(); +} + +void registerSeparateCVScopePasses() { + registerPass([]() -> std::unique_ptr { + return createSeparateCVScopePass(); + }); +} + +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/TritonAffinityOpt/CMakeLists.txt b/third_party/ascend/lib/TritonAffinityOpt/CMakeLists.txt new file mode 100644 index 0000000000..925ca52d92 --- /dev/null +++ b/third_party/ascend/lib/TritonAffinityOpt/CMakeLists.txt @@ -0,0 +1,19 @@ +add_triton_library(TritonAffinityOpt + DAGSSBuffer.cpp + DAG.cpp + DAGSync.cpp + DAGScope.cpp + + DEPENDS + TritonAffinityOptConversionPassIncGen + + LINK_LIBS + BiShengIRHIVMDialect + BiShengIRScopeDialect + MLIRIR + MLIRPass + MLIRTransforms + MLIRSupport + TritonIR + MLIRSCFDialect +) diff --git a/third_party/ascend/lib/TritonAffinityOpt/DAG.cpp b/third_party/ascend/lib/TritonAffinityOpt/DAG.cpp new file mode 100644 index 0000000000..d97030d331 --- /dev/null +++ b/third_party/ascend/lib/TritonAffinityOpt/DAG.cpp @@ -0,0 +1,520 @@ +#include "TritonAffinityOpt/DAG.h" +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/OperationSupport.h" +#include "mlir/IR/TypeUtilities.h" +#include "mlir/IR/Value.h" +#include "mlir/IR/ValueRange.h" +#include "mlir/Interfaces/ControlFlowInterfaces.h" +// LLVM 22 removed mlir/Interfaces/CopyOpInterface.h (upstream PR #157711); +// AscendNPU-IR vendors the interface under bishengir/Interfaces. +#include "bishengir/Interfaces/CopyOpInterface.h" +#include "mlir/Interfaces/LoopLikeInterface.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Types.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/FormatVariadic.h" +#include "llvm/Support/raw_ostream.h" +#include +#include +#include +#include +#include +#include + +namespace mlir { +namespace AffinityDAG { + +const auto printFlags = + OpPrintingFlags().enableDebugInfo(true, true).skipRegions(); + +const char *literalCoreType(CoreType ct) { + switch (ct) { + case VECTOR_ONLY: + return "VECTOR_ONLY"; + case CUBE_ONLY: + return "CUBE_ONLY"; + case CUBE_AND_VECTOR: + return "CUBE_AND_VECTOR"; + case UNDETERMINED: + return "UNDETERMINED"; + } + return "Unknown"; +} + +bool opIsScf(Operation *op) { + if (!llvm::isa(op->getDialect())) + return false; + return true; +} + +Graph::Graph(Block *block, Graph *parent, OpMap opMap, ValueMap valueMap, + bool inheritParent) + : block(block), parent(parent), opMap(opMap), valueMap(valueMap) { + + if (parent && inheritParent) { + if (!this->opMap) { + this->opMap = parent->opMap; + } + + if (!this->valueMap) { + this->valueMap = parent->valueMap; + } + } + + if (!this->opMap) { + this->opMap = std::make_shared(); + } + + if (!this->valueMap) { + this->valueMap = std::make_shared(); + } + + for (auto blockArg : block->getArguments()) { + (*this->valueMap)[blockArg] = std::make_unique(blockArg); + blockArgs.push_back((*this->valueMap)[blockArg].get()); + } + + for (auto &opRef : block->getOperations()) { + opCount += 1; + auto op = &opRef; + auto opNodeUnique = std::make_unique(op, this); + auto opNode = opNodeUnique.get(); + (*this->opMap)[op] = std::move(opNodeUnique); + + if (block->mightHaveTerminator() && op == block->getTerminator()) { + terminator = opNode; + } + + for (auto &subgraph : opNode->subgraphs) { + opCount += subgraph.opCount; + } + } +}; + +bool valueIsScalar(Value value) { + auto type = value.getType(); + + if (type.isIntOrIndexOrFloat()) { + return true; + } + + if (auto tensorType = llvm::dyn_cast(type)) { + return tensorType.getRank() == 0; + } + + if (auto _ = llvm::dyn_cast(type)) { + return true; + } + + return false; +} + +bool valueIsTensorOfPtr(Value value) { + auto type = value.getType(); + if (auto tensorType = llvm::dyn_cast(type)) { + auto elementType = tensorType.getElementType(); + if (llvm::isa(elementType)) { + return true; + } + } + + return false; +} + +OpAbility OpNode::canRunOn() const { + if (opIsScf(op)) { + return OpAbility::CUBE_AND_VECTOR; + } + return llvm::TypeSwitch(op) + .Case([](auto) { return OpAbility::CUBE_ONLY; }) + .Case([](auto) { return OpAbility::CUBE_AND_VECTOR; }) + .Case([](arith::SelectOp op) { + // when cond is vector, selectOp should be vector, otherwise scalar + return (valueIsScalar(op.getCondition()) ? OpAbility::CUBE_AND_VECTOR + : OpAbility::PREFER_VECTOR); + }) + .Default([](Operation *op) { + auto isVector = false; + for (auto operand : op->getOperands()) { + if (!valueIsScalar(operand)) { + // if (valueIsTensorOfPtr(operand)) { + // return SCALAR; + // } + isVector = true; + } + } + + for (auto result : op->getResults()) { + if (!valueIsScalar(result)) { + // if (valueIsTensorOfPtr(result)) { + // return SCALAR; + // } + isVector = true; + } + } + + if (isVector) { + return OpAbility::PREFER_VECTOR; + } + + return OpAbility::CUBE_AND_VECTOR; + }); +} + +OpNode::OpNode(Operation *op, Graph *graph) : Node(Node::NK_Op), op(op) { + if (op == nullptr) { + return; + } + + llvm::outs() << op << "\n"; + + auto &valueMap = *graph->valueMap.get(); + auto &opMap = *graph->opMap.get(); + for (const auto operand : op->getOperands()) { + auto valueNode = valueMap.at(operand).get(); + valueNode->outputs.push_back(this); + inputs.push_back(valueNode); + } + + for (const auto &result : op->getResults()) { + auto valueNodeUnique = std::make_unique(result); + auto valueNode = valueNodeUnique.get(); + valueMap[result] = std::move(valueNodeUnique); + valueNode->source = this; + outputs.push_back(valueNode); + } + + // if (!op->hasTrait()) { + // llvm::dbgs() << "Not building subgraph because op is not SingleBlock: " + // << op << '\n'; return; + // } + + if (auto branchOp = llvm::dyn_cast(op)) { + + OpNode *terminator = nullptr; + llvm::SmallVector, 2> validRegions; + + for (auto ®ion : branchOp->getRegions()) { + if (region.getBlocks().empty()) + continue; + subgraphs.emplace_back(®ion.getBlocks().front(), graph); + validRegions.emplace_back(region, subgraphs.back()); + } + + for (auto [region, subgraph] : validRegions) { + SmallVector succRegions; + + branchOp.getSuccessorRegions(region, succRegions); + if (auto currTerminator = dyn_cast( + subgraph.terminator->op)) { + for (auto &succ : succRegions) { + auto forwardedVal = currTerminator.getSuccessorOperands(succ); + if (succ.isParent()) { + // Step1: first yield to parent -> results: double direction + if (!terminator && subgraph.terminator) { + terminator = subgraph.terminator; + for (auto [forwardedVal, resultNode] : + llvm::zip_equal(forwardedVal, outputs)) { + auto resultValueNode = llvm::dyn_cast(resultNode); + assert(resultValueNode && + "Output of a OpNode should be ValueNode!"); + auto forwardedNode = valueMap[forwardedVal].get(); + resultValueNode->source = forwardedNode; + forwardedNode->outputs.push_back(resultNode); + } + } + + } else { + // Step2: Region terminator -> Succ Operands + auto succRegion = succ.getSuccessor(); + + for (auto [operand, succInput] : + llvm::zip_equal(forwardedVal, succ.getSuccessorInputs())) { + auto forwardedNode = valueMap[operand].get(); + auto succNode = valueMap[succInput].get(); + forwardedNode->outputs.push_back(succNode); + succNode->source = forwardedNode; + } + } + } + } + } + + if (auto loopOp = llvm::dyn_cast(op)) { + // Step3: inits->iter_args (single directional) (should be handled in step + // 2: ) last terminator -> iter_args (bidirectional) + for (auto [init, iterArgVal] : + llvm::zip_equal(loopOp.getInits(), loopOp.getRegionIterArgs())) { + auto &initNode = valueMap[init]; + auto &iterArgNode = valueMap[iterArgVal]; + initNode->outputs.push_back(iterArgNode.get()); + } + // for(auto [init, iterArgVal, yieldNode] : + // llvm::zip_equal(loopOp.getInits(), loopOp.getRegionIterArgs(), + // terminator->outputs)) { + // auto& initNode = valueMap[init]; + // auto& iterArgNode = valueMap[iterArgVal]; + // initNode->outputs.push_back(iterArgNode.get()); + // yieldNode->outputs.push_back(iterArgNode.get()); + // iterArgNode->source = yieldNode; + // } + } + } +} + +// llvm::SmallVector getWriteOperandPriority(OpNode* op) { + +// llvm::SmallVector result(op->getInputs()); + +// auto getPriority = [](ValueNode* node) { +// auto typ = getElementTypeOrSelf(node->value); +// if (typ.isInteger(1)) { +// return 2; +// } +// if (llvm::isa(typ)) { +// return 1; +// } +// return 0; +// }; + +// std::stable_sort(result.begin(), result.end(), [&](ValueNode* a, ValueNode* +// b) { +// return getPriority(a) < getPriority(b); +// }); + +// return result; +// } + +ValueNode *getWriteDataSource(OpNode *op) { + auto inputRange = op->getInputs(); + for (auto node : inputRange.drop_front()) { + auto typ = getElementTypeOrSelf(node->value); + if (!typ.isInteger(1)) { + return node; + } + }; + + return nullptr; +} + +enum class MemPolicy { NONE, READ, WRITE }; + +CoreType Node::absorbCommon() { + + auto sourceNode = getSourceOpNode(); + auto op = sourceNode ? sourceNode->op : nullptr; + + if (!sourceNode || !op) { + CoreType newCoreType = isOnPrivate; + for (auto output : outputs) { + newCoreType = newCoreType | output->isOn(); + isUpstreamOfCubeMem = isUpstreamOfCubeMem || output->isUpstreamOfCubeMem; + } + return newCoreType; + } + + CoreType newCoreType = sourceNode->isOn(); + + OpAbility ability = sourceNode->canRunOn(); + + if (ability == OpAbility::CUBE_ONLY) { + return CUBE_ONLY; + } + + auto memIface = llvm::dyn_cast(op); + auto memPolicy = MemPolicy::NONE; + + if (memIface) { + // Possible improvements: Determine the policy to use based on shapes, + // inputs and outputs, etc + if (memIface.hasEffect()) { + memPolicy = MemPolicy::WRITE; + } else if (memIface.hasEffect()) { + memPolicy = MemPolicy::READ; + } + } + + if (memPolicy == MemPolicy::WRITE) { + if (auto data = getWriteDataSource(sourceNode)) { + auto currCt = data->isOn(); + if (exactlyOneType(currCt)) { + if (currCt == CUBE_ONLY) { + isUpstreamOfCubeMem = true; + } + return currCt; + } + } + + // data is not cube_only + return VECTOR_ONLY; + } + + for (auto output : outputs) { + switch (output->isOn()) { + case CUBE_AND_VECTOR: + newCoreType = newCoreType | VECTOR_ONLY; + // not breaking the switch because we need to handle cube + case CUBE_ONLY: + if (ability != OpAbility::PREFER_VECTOR || output->isUpstreamOfCubeMem || + memPolicy == MemPolicy::READ) { + isUpstreamOfCubeMem = + (isUpstreamOfCubeMem || output->isUpstreamOfCubeMem || + memPolicy == MemPolicy::READ); + newCoreType = newCoreType | CUBE_ONLY; + } + break; + case VECTOR_ONLY: + newCoreType = newCoreType | VECTOR_ONLY; + default: // UNDETERMINED, skip + break; + }; + } + + return newCoreType; +} + +CoreType OpNode::absorbImpl() { + if (opIsScf(op)) { + return CUBE_AND_VECTOR; + } + + auto newCoreType = absorbCommon(); + + // if (canRunOn() == OpAbility::CUBE_AND_VECTOR) { + // for (auto input : inputs) { + // newCoreType = newCoreType | input->isOn(); + // } + // } + + return newCoreType; +} + +CoreType ValueNode::absorbImpl() { return absorbCommon(); } + +std::unique_ptr Graph::fromMultiBlockFunc(triton::FuncOp funcOp) { + + auto dummyBlock = new Block(); + auto dummyGraph = std::make_unique(dummyBlock); + auto dummyNode = std::make_unique(nullptr, dummyGraph.get()); + size_t opCount = 0; + + for (auto &block : funcOp.getBody()) { + auto &subgraph = + dummyNode->subgraphs.emplace_back(&block, dummyGraph.get()); + opCount += subgraph.opCount; + } + + auto &opMap = *dummyGraph->opMap.get(); + auto &valueMap = *dummyGraph->valueMap.get(); + + llvm::SmallVector nodes; + nodes.reserve(opMap.size() + valueMap.size()); + + for (auto &[_, node] : opMap) { + if (node.get()) + nodes.push_back(node.get()); + } + + for (auto &[_, node] : valueMap) { + if (node.get()) + nodes.push_back(node.get()); + } + + auto diffuse = [&]() { + // Not sure if determinism is required + llvm::SmallSetVector worklist(nodes.begin(), nodes.end()); + + size_t threshold = worklist.size() * 5; + + for (size_t i = 0; i < threshold; i++) { + if (worklist.empty()) { + break; + } + + auto node = worklist.pop_back_val(); + + if (node->absorb()) { + auto affected = node->getAffected(); + worklist.insert(affected.begin(), affected.end()); + } + } + }; + + diffuse(); + + for (auto node : nodes) { + if (node->isOn() == UNDETERMINED) { + node->isOnPrivate = VECTOR_ONLY; + } + } + + diffuse(); + + OpPrintingFlags flags; + flags.skipRegions(); + + for (auto [idx, node] : llvm::enumerate(nodes)) { + llvm::TypeSwitch(node) + .Case([&, idx = idx](OpNode *node) { + if (node->op) { + llvm::dbgs() << llvm::formatv( + "\n\n====== OpNode on: {1} @ {0} ======\n", node->op, + literalCoreType(node->isOn())); + node->op->print(llvm::dbgs(), flags); + llvm::dbgs() << "\nAbility: " + << literalCoreType(toCoreType(node->canRunOn())); + llvm::dbgs() << llvm::formatv("\n====== {0} ======\n", node->op); + } + }) + .Case([&, idx = idx](ValueNode *node) { + if (node->value) { + llvm::dbgs() << llvm::formatv( + "\n\n====== ValueNode on {1} @ {0} ======\n", node->value, + literalCoreType(node->isOn())); + node->value.print(llvm::dbgs(), flags); + llvm::dbgs() << llvm::formatv("\n====== {0} ======\n", node->value); + } + }); + // if (auto opNode = llvm::dyn_cast(node)) { + // if (auto forOp = llvm::dyn_cast_if_present(opNode->op)) { + // llvm::dbgs() << "\n==== ForOp ====\n"; + // llvm::dbgs() << forOp << "\n"; + // llvm::dbgs() << "\n---- IterArgs ----\n"; + // for(auto iterArg : forOp.getRegionIterArgs()) { + // auto& valueNode = valueMap[iterArg]; + // llvm::dbgs() << llvm::formatv( + // "{0}: {1} upstream: {2} definingOp: {3} \n", + // iterArg.getArgNumber(), + // literalCoreType(valueNode->isOn()), + // literalCoreType(valueNode->source->isOn()), + // valueNode->getSourceOp()->op + // ); + // } + // llvm::dbgs() << "\n---- Results ----\n"; + // for(auto result : forOp.getResults()) { + // llvm::dbgs() << result.getResultNumber() << ' ' << + // literalCoreType(valueMap[result]->isOn()) << '\n'; + // } + // } + // } + } + + return dummyGraph; +}; + +} // namespace AffinityDAG +} // namespace mlir diff --git a/third_party/ascend/lib/TritonAffinityOpt/DAGSSBuffer.cpp b/third_party/ascend/lib/TritonAffinityOpt/DAGSSBuffer.cpp new file mode 100644 index 0000000000..389f0d9a25 --- /dev/null +++ b/third_party/ascend/lib/TritonAffinityOpt/DAGSSBuffer.cpp @@ -0,0 +1,5601 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "TritonAffinityOpt/Passes.h" + +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/HIVM/IR/HIVMImpl.h" +#include "bishengir/Dialect/HIVM/IR/HIVMInterfaces.h" +#include "bishengir/Dialect/HIVM/Transforms/Passes.h" +#include "bishengir/Dialect/HIVM/Utils/Utils.h" +#include "bishengir/Dialect/Scope/IR/Scope.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Linalg/Transforms/TilingInterfaceImpl.h" +#include "mlir/Dialect/Linalg/Utils/Utils.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/Operation.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "incubated/Conversion/UtilsIncubated/Utils.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/SmallVector.h" +#include + +// #include "mlir/Pass/Pass.h" +// #include "mlir/Pass/PassManager.h" + +// #include "mlir/Transforms/Canonicalizer.h" +// #include "mlir/Support/LogicalResult.h" + +namespace mlir { +namespace triton { +#define GEN_PASS_DEF_DAGSSBUFFER +#include "ascend/include/TritonAffinityOpt/Passes.h.inc" +} // namespace triton +} // namespace mlir + +using namespace mlir; +using namespace hivm; + +namespace { +struct DAGSSBufferPass + : public mlir::triton::impl::DAGSSBufferBase { + void runOnOperation() override; + void getDependentDialects(DialectRegistry ®istry) const override { + registry.insert(); + registry.insert(); + } +}; +} // namespace + +void ControlSsbufV2(ModuleOp module) { + mlir::OpBuilder builder(module.getContext()); + // 用于记录已经处理过的scope.scope操作 + llvm::DenseSet processedScopes; + + auto aiCAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + int cubeControlIndex = 15; + int vectorControlIndex = 14; + + llvm::DenseSet processedScopes2; + module->walk([&](SyncBlockWaitOp op) { + auto pipeS = hivm::PipeAttr::get(op->getContext(), hivm::PIPE::PIPE_S); + + // 向上查找父scope.scope操作 + mlir::Operation *parentOp = op->getParentOp(); + mlir::Operation *scopeOp = nullptr; + mlir::Operation *forOp = nullptr; + + // 向上遍历查找scope.scope操作 + while (parentOp) { + if (dyn_cast(parentOp)) { + scopeOp = parentOp; + break; + } + parentOp = parentOp->getParentOp(); + } + parentOp = op->getParentOp(); + while (parentOp) { + if (dyn_cast(parentOp)) { + forOp = parentOp; + break; + } + parentOp = parentOp->getParentOp(); + } + // 如果没有找到scope.scope操作，则跳过 + if (!scopeOp) { + return; + } + if (!forOp) { + return; + } + + // 如果该scope已经处理过，则跳过 + if (processedScopes2.count(forOp) > 0) + return; + + // 标记该scope为已处理 + processedScopes2.insert(forOp); + }); + bool firstSet = true; + bool firstWait = true; + for (auto forOp : processedScopes2) { + mlir::Operation *parentOp = forOp->getParentOp(); + mlir::Operation *scopeOp = nullptr; + + // 向上遍历查找scope.scope操作 + while (parentOp) { + if (dyn_cast(parentOp)) { + scopeOp = parentOp; + break; + } + parentOp = parentOp->getParentOp(); + } + bool isAIC = false; + // 1. 先检查操作是否有这个属性 + + if (scopeOp->hasAttr("hivm.tcore_type")) { + auto attr = scopeOp->getAttr("hivm.tcore_type"); + if (attr == aiCAttr) { + isAIC = true; + } + } + + if (isAIC) { + // 在for循环的开头插入代码 + builder.setInsertionPoint(scopeOp); + // %ssb_ready_addr = llvm.mlir.constant(0 : i64) : i64 + auto i64Type = builder.getIntegerType(64); + auto i32Type = builder.getIntegerType(32); + + builder.setInsertionPointToStart(&forOp->getRegion(0).front()); + // %ssb_ready_addr = llvm.mlir.constant(0 : i64) : i64 + // add sync_block_wait + auto coreAttr = + hivm::TCoreTypeAttr::get(module.getContext(), hivm::TCoreType::CUBE); + auto setPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + auto waitPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + auto flagId = + builder.getIntegerAttr(builder.getI64Type(), vectorControlIndex); + builder.create(forOp->getLoc(), coreAttr, setPipe, + waitPipe, flagId); + + // 在循环末尾（yield之前）插入代码 + auto &loopBody = forOp->getRegion(0).front(); + // 找到循环体的terminator（应该是yield操作） + auto *terminator = loopBody.getTerminator(); + builder.setInsertionPoint(terminator); + + // add sync_block_set + coreAttr = + hivm::TCoreTypeAttr::get(module.getContext(), hivm::TCoreType::CUBE); + setPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + waitPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + flagId = builder.getIntegerAttr(builder.getI64Type(), cubeControlIndex); + builder.create(forOp->getLoc(), coreAttr, setPipe, + waitPipe, flagId); + + if (firstWait) { + auto &scopeBlock = scopeOp->getRegion(0).front(); + auto *scope_terminator = scopeBlock.getTerminator(); + builder.setInsertionPoint(scope_terminator); + // add sync_block_wait + coreAttr = hivm::TCoreTypeAttr::get(module.getContext(), + hivm::TCoreType::CUBE); + setPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + waitPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + flagId = + builder.getIntegerAttr(builder.getI64Type(), vectorControlIndex); + builder.create(forOp->getLoc(), coreAttr, setPipe, + waitPipe, flagId); + firstWait = false; + } + } else { + // 1. 在scopeop的开头插入代码 + // 假设scopeOp是一个具有区域的操作，我们获取其第一个块 + if (firstSet) { + auto &scopeBlock = scopeOp->getRegion(0).front(); + builder.setInsertionPointToStart(&scopeBlock); + + // add sync_block_wait + auto coreAttr = hivm::TCoreTypeAttr::get(module.getContext(), + hivm::TCoreType::VECTOR); + auto setPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + auto waitPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + auto flagId = + builder.getIntegerAttr(builder.getI64Type(), vectorControlIndex); + builder.create(forOp->getLoc(), coreAttr, setPipe, + waitPipe, flagId); + firstSet = false; + } + + auto i64Type = builder.getIntegerType(64); + auto i32Type = builder.getIntegerType(32); + + // 创建需要的常量 + auto c32ConstAttr = mlir::IntegerAttr::get(i64Type, 32); + auto c32ConstOp = builder.create( + scopeOp->getLoc(), i64Type, c32ConstAttr); + + auto c0i64ConstAttr = mlir::IntegerAttr::get(i64Type, 0); + auto c0i64ConstOp = builder.create( + scopeOp->getLoc(), i64Type, c0i64ConstAttr); + + auto c0i32ConstAttr = mlir::IntegerAttr::get(i32Type, 0); + auto c0i32ConstOp = builder.create( + scopeOp->getLoc(), i32Type, c0i32ConstAttr); + + auto c1i32ConstAttr = mlir::IntegerAttr::get(i32Type, 1); + auto c1i32ConstOp = builder.create( + scopeOp->getLoc(), i32Type, c1i32ConstAttr); + + // %sub_id = hivm.hir.get_sub_block_idx -> i64 + // 这里假设有一个getSubBlockIdxOp操作 + auto subIdOp = + builder.create(scopeOp->getLoc(), i64Type); + + // %ssb_addr_offset = arith.muli %sub_id, %c32_i64 : i64 + auto ssbAddrOffsetOp = builder.create( + scopeOp->getLoc(), subIdOp.getResult(), c32ConstOp.getResult()); + + // %ssb_addr = arith.addi %ssb_addr_offset, %c32_i64 : i64 + auto ssbAddrOp = builder.create( + scopeOp->getLoc(), ssbAddrOffsetOp.getResult(), + c32ConstOp.getResult()); + + // %vec_id = arith.cmpi eq, %sub_id, %c0_i64 : i64 + auto vecIdOp = builder.create( + scopeOp->getLoc(), mlir::arith::CmpIPredicate::eq, + subIdOp.getResult(), c0i64ConstOp.getResult()); + + // 2. 在parentop的开头插入代码 + builder.setInsertionPointToStart(&forOp->getRegion(0).front()); + + // add sync_block_wait + auto coreAttr = hivm::TCoreTypeAttr::get(module.getContext(), + hivm::TCoreType::VECTOR); + auto setPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + auto waitPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + auto flagId = + builder.getIntegerAttr(builder.getI64Type(), cubeControlIndex); + builder.create(forOp->getLoc(), coreAttr, setPipe, + waitPipe, flagId); + + // 在循环末尾（yield之前）插入代码 + auto &loopBody = forOp->getRegion(0).front(); + // 找到循环体的terminator（应该是yield操作） + auto *terminator = loopBody.getTerminator(); + builder.setInsertionPoint(terminator); + + // add sync_block_wait + coreAttr = hivm::TCoreTypeAttr::get(module.getContext(), + hivm::TCoreType::VECTOR); + setPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + waitPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + flagId = builder.getIntegerAttr(builder.getI64Type(), vectorControlIndex); + builder.create(forOp->getLoc(), coreAttr, setPipe, + waitPipe, flagId); + } + } + + auto i64Type = builder.getIntegerType(64); + auto i32Type = builder.getIntegerType(32); + auto initPtrType = mlir::LLVM::LLVMPointerType::get(builder.getContext(), 11); + SmallVector scopeOps; + module->walk([&](mlir::Operation *op) { + // 检查是否为目标操作 + if (auto scopeOp = dyn_cast(op)) { + scopeOps.push_back(scopeOp); + } + }); + if (!scopeOps.empty()) { + auto scopeOp = scopeOps[0]; + builder.setInsertionPoint(scopeOp); + auto c0i64ConstAttr = mlir::IntegerAttr::get(i64Type, 0); + auto c0i64ConstOp = builder.create( + scopeOp->getLoc(), i64Type, c0i64ConstAttr); + auto c32i64ConstAttr = mlir::IntegerAttr::get(i64Type, 32); + auto c32i64ConstOp = builder.create( + scopeOp->getLoc(), i64Type, c32i64ConstAttr); + auto c64i64ConstAttr = mlir::IntegerAttr::get(i64Type, 64); + auto c64i64ConstOp = builder.create( + scopeOp->getLoc(), i64Type, c64i64ConstAttr); + auto c96i64ConstAttr = mlir::IntegerAttr::get(i64Type, 96); + auto c96i64ConstOp = builder.create( + scopeOp->getLoc(), i64Type, c96i64ConstAttr); + auto c0i32ConstAttr = mlir::IntegerAttr::get(i32Type, 0); + auto c0i32ConstOp = builder.create( + scopeOp->getLoc(), i32Type, c0i32ConstAttr); + + auto c0initInttoptrOp = builder.create( + scopeOp->getLoc(), initPtrType, c0i64ConstOp.getResult()); + auto c32initInttoptrOp = builder.create( + scopeOp->getLoc(), initPtrType, c32i64ConstOp.getResult()); + auto c64initInttoptrOp = builder.create( + scopeOp->getLoc(), initPtrType, c64i64ConstOp.getResult()); + auto c96initInttoptrOp = builder.create( + scopeOp->getLoc(), initPtrType, c96i64ConstOp.getResult()); + + builder.create(scopeOp->getLoc(), c0i32ConstOp, + c0initInttoptrOp); + builder.create(scopeOp->getLoc(), c0i32ConstOp, + c32initInttoptrOp); + builder.create(scopeOp->getLoc(), c0i32ConstOp, + c64initInttoptrOp); + builder.create(scopeOp->getLoc(), c0i32ConstOp, + c96initInttoptrOp); + } +} + +scf::ForOp transformLoop(scf::ForOp forOp, OpBuilder &builder) { + + // 1. 获取原始循环的信息 + Value originalLowerBound = forOp.getLowerBound(); + Value originalUpperBound = forOp.getUpperBound(); + Value originalStep = forOp.getStep(); + SmallVector iterArgs; + for (auto arg : forOp.getInitArgs()) { + iterArgs.push_back(arg); + } + auto yields = forOp.getBody()->getTerminator(); + + // 2. 检查循环体中是否有特定操作 + int hasTargetOps = 0; + forOp.walk([&](Operation *op) { + if (auto ifOp = dyn_cast(op)) { + if (ifOp->hasAttr("ssbuffer")) { + hasTargetOps++; + } + } + }); + // 3. 如果存在目标操作，在迭代参数中添加计数器 + Value counterInit = nullptr; + mlir::Operation *parentOp = forOp->getParentOp(); + mlir::Operation *scopeOp = nullptr; + // 向上遍历查找scope.scope操作 + while (parentOp) { + if (dyn_cast(parentOp)) { + scopeOp = parentOp; + break; + } + parentOp = parentOp->getParentOp(); + } + + builder.setInsertionPoint(scopeOp); + for (int i = 0; i < hasTargetOps; i++) { + Location loc = forOp.getLoc(); + auto argType = originalLowerBound.getType(); + + // 添加到迭代参数列表 + iterArgs.push_back(originalLowerBound); + } + // 2. 创建新的上界：originalUpperBound * 2 + Location loc = forOp.getLoc(); + Type ubType = originalStep.getType(); + builder.setInsertionPoint(forOp); + + int count = 0; + for (auto &op : forOp.getBody()->getOperations()) { + if (auto ifOp = dyn_cast(op)) { + auto parentOp = ifOp->getParentOp(); + if (parentOp == forOp && ifOp->hasAttr("ssbuffer")) { + count++; + } + } + } + + Value two; + if (ubType.isIndex()) { + two = builder.create(loc, count - 1); + } else if (auto intType = dyn_cast(ubType)) { + // 对于整数类型，创建相应类型的常数2 + two = builder.create(loc, intType, count - 1); + } else { + // 其他类型可能需要特殊处理 + llvm::errs() << "Warning: Unexpected type for upper bound: " << ubType + << "\n"; + // 尝试创建索引类型的2然后转换 + auto indexTwo = builder.create(loc, count - 1); + two = builder.create(loc, ubType, indexTwo); + } + + auto steps = builder.create(forOp.getLoc(), originalStep, two); + + auto nowUpperBound = + builder.create(forOp.getLoc(), originalUpperBound, steps); + + // 3. Create a new for loop + auto newForOp = + builder.create(forOp.getLoc(), originalLowerBound, + nowUpperBound, originalStep, iterArgs); + + // 4. 设置IR映射表，将旧循环的变量映射到新循环 + IRMapping mapper; + + // 映射迭代变量 + mapper.map(forOp.getInductionVar(), newForOp.getInductionVar()); + + // 映射迭代参数 + for (auto [oldArg, newArg] : + llvm::zip(forOp.getRegionIterArgs(), newForOp.getRegionIterArgs())) { + mapper.map(oldArg, newArg); + } + + SmallVector newCounterArgs; + for (int i = forOp.getRegionIterArgs().size(); + i < newForOp.getRegionIterArgs().size(); i++) { + newCounterArgs.push_back(newForOp.getRegionIterArgs()[i]); + } + // 5. 克隆循环体内容到新循环 + auto &newLoopBody = *newForOp.getBody(); + builder.setInsertionPointToStart(&newLoopBody); + + for (auto &op : forOp.getBody()->without_terminator()) { + builder.clone(op, mapper); + } + + // 6. 克隆yield操作 + if (auto yieldOp = dyn_cast(yields)) { + SmallVector newYieldOperands; + for (auto operand : yieldOp.getOperands()) { + newYieldOperands.push_back(mapper.lookupOrDefault(operand)); + } + if (hasTargetOps != 0) { + for (auto currentCounter : newCounterArgs) { + // 将更新后的计数器添加到yield操作数中 + newYieldOperands.push_back(currentCounter); + } + } + builder.create(yieldOp.getLoc(), newYieldOperands); + } + + // 7. 替换原循环的结果 + if (hasTargetOps != 0) { + // 新循环有额外的计数器结果，但原循环没有对应结果 + // 我们可以选择只替换原循环对应的结果，或者忽略计数器结果 + unsigned numOriginalResults = forOp.getNumResults(); + SmallVector originalResults; + for (unsigned i = 0; i < numOriginalResults; i++) { + originalResults.push_back(newForOp.getResult(i)); + } + forOp.replaceAllUsesWith(originalResults); + } else { + forOp.replaceAllUsesWith(newForOp.getResults()); + } + + // 8. 删除原循环 + forOp.erase(); + return newForOp; +} + +// Find the first occurrence of convert_layout or fixpipe operation after the +// specified operation +Operation * +findFirstTargetOpAfterWait(SyncBlockWaitOp waitOp, + SmallVector &excludedValues) { + bool startSearching = false; + + for (Operation &op : waitOp->getBlock()->getOperations()) { + // meet waitop, start searching + if (&op == waitOp) { + startSearching = true; + continue; + } + + // have not meet waitOp, skip + if (!startSearching) { + continue; + } + + Operation *res = nullptr; + + if (isa(op)) { + res = op.getOperands()[0].getDefiningOp(); + } else if (isa(op)) { + res = op.getOperands()[1].getDefiningOp(); + } else if (isa(op)) { + res = op.getOperands()[1].getDefiningOp(); + } else if (isa(op)) { + res = op.getOperands()[0].getDefiningOp(); + } + + // get res + if (res) { + // in excludedValues → skip，search next one + if (llvm::is_contained(excludedValues, res)) { + continue; + } + // not in excludedValues → note and return + excludedValues.push_back(res); + return res; + } + } + + return nullptr; +} + +void getWaitType(std::string CoreType, scf::ForOp forOp, + SmallVector &waitTypes, + SmallVector &allocTypes) { + auto scalarWaitPipe = PipeAttr::get(forOp.getContext(), hivm::PIPE::PIPE_S); + auto cubeWaitPipe = PipeAttr::get(forOp.getContext(), hivm::PIPE::PIPE_FIX); + auto vectorWaitPipe = + PipeAttr::get(forOp.getContext(), hivm::PIPE::PIPE_MTE3); + SmallVector excludedValues; + forOp.walk([&](Operation *op) { + if (auto waitOp = dyn_cast(op)) { + auto parentOp = op->getParentOp(); + if (isa(parentOp) && parentOp->hasAttr("ssbuffer")) { + auto ifOp = dyn_cast(parentOp); + if (forOp == ifOp->getParentOp()) { + auto waitPipe = waitOp.getPipe(); + auto setPipe = waitOp.getTpipe(); + if (waitPipe == scalarWaitPipe || setPipe == scalarWaitPipe) { + bool hasBackward = waitOp->hasAttr("ssbuf.backward"); + if (hasBackward) { + waitTypes.push_back(0); + allocTypes.push_back(waitOp); + } else { + waitTypes.push_back(1); + allocTypes.push_back(waitOp); + } + } else { + if ((waitPipe == cubeWaitPipe && CoreType == "cube") || + (waitPipe == vectorWaitPipe && CoreType == "vector")) { + auto allocOp = findFirstTargetOpAfterWait(waitOp, excludedValues); + waitTypes.push_back(0); + allocTypes.push_back(allocOp); + } else { + auto allocOp = findFirstTargetOpAfterWait(waitOp, excludedValues); + waitTypes.push_back(1); + allocTypes.push_back(allocOp); + } + } + } + } + } + }); +} + +DenseMap getCounterOffset(scf::ForOp forOp) { + int i = 0; + DenseMap bufferMap; + auto scalarWaitPipe = PipeAttr::get(forOp.getContext(), hivm::PIPE::PIPE_S); + forOp.walk([&](Operation *op) { + bufferMap[i] = 0; + auto ifOp = dyn_cast(op); + if (ifOp && ifOp->hasAttr("ssbuffer") && ifOp->getParentOp() == forOp) { + ifOp.walk([&](Operation *op) { + if (auto waitOp = dyn_cast(op)) { + if (auto waitIfOp = dyn_cast(op->getParentOp())) { + if (waitIfOp == ifOp) { + auto waitPipe = waitOp.getPipe(); + auto setPipe = waitOp.getTpipe(); + if ((waitPipe != scalarWaitPipe || setPipe != scalarWaitPipe)) { + bufferMap[i]++; + } + } + } + } + }); + i++; + } + }); + return bufferMap; +} + +SmallVector addBufValLoop(scf::ForOp forOp, + DenseMap VecBitMap, + DenseMap CubeBitMap, + OpBuilder &builder) { + auto aiCAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + bool isAIC = false; + // 向上查找父scope.scope操作 + mlir::Operation *parentOp = forOp->getParentOp(); + mlir::Operation *scopeOp = nullptr; + // 向上遍历查找scope.scope操作 + while (parentOp) { + if (dyn_cast(parentOp)) { + scopeOp = parentOp; + break; + } + parentOp = parentOp->getParentOp(); + } + if (scopeOp->hasAttr("hivm.tcore_type")) { + auto attr = scopeOp->getAttr("hivm.tcore_type"); + if (attr == aiCAttr) { + isAIC = true; + } + } + auto bufferMap = getCounterOffset(forOp); + SmallVector buf_vals; + SmallVector if_conditions; + builder.setInsertionPointToStart(&scopeOp->getRegion(0).front()); + + // 1. 提取并处理end值 + Value startValue = forOp.getLowerBound(); + Value endValue = forOp.getUpperBound(); + // 2. 提取并处理step值 + Value stepValue = forOp.getStep(); + builder.setInsertionPoint(forOp); + Location loc = forOp.getLoc(); + int count = 0; + for (auto &op : forOp.getBody()->getOperations()) { + if (auto ifOp = dyn_cast(op)) { + auto parentOp = ifOp->getParentOp(); + if (parentOp == forOp && ifOp->hasAttr("ssbuffer")) { + count++; + } + } + } + + Value two; + Type ubType = stepValue.getType(); + if (ubType.isIndex()) { + two = builder.create(loc, count - 1); + } else if (auto intType = dyn_cast(ubType)) { + // 对于整数类型，创建相应类型的常数2 + two = builder.create(loc, intType, count - 1); + } else { + // 其他类型可能需要特殊处理 + llvm::errs() << "Warning: Unexpected type for upper bound: " << ubType + << "\n"; + // 尝试创建索引类型的2然后转换 + auto indexTwo = builder.create(loc, count - 1); + two = builder.create(loc, ubType, indexTwo); + } + + auto steps = builder.create(forOp.getLoc(), endValue.getType(), + stepValue, two); + + auto subLoopValue = builder.create( + forOp.getLoc(), endValue.getType(), endValue, steps); + + SmallVector WaitType; + SmallVector AllocType; + SmallVector bufferPtrs; + if (isAIC) { + builder.setInsertionPointToStart(&forOp->getRegion(0).front()); + // 创建常量32和64 + Value c0 = + builder.create(forOp.getLoc(), 0, 32 // 值32，64位 + ); + Value c32 = builder.create(forOp.getLoc(), 32, + 64 // 值32，64位 + ); + Value c64 = builder.create(forOp.getLoc(), 64, + 64 // 值64，64位 + ); + // 创建inttoptr操作 + Value ssb_vec0_ptr = builder.create( + forOp.getLoc(), + LLVM::LLVMPointerType::get(builder.getContext(), 11), // 地址空间11 + c32); + Value ssb_vec1_ptr = builder.create( + forOp.getLoc(), + LLVM::LLVMPointerType::get(builder.getContext(), 11), // 地址空间11 + c64); + bufferPtrs.push_back(ssb_vec0_ptr); + bufferPtrs.push_back(ssb_vec1_ptr); + // 创建load操作 + Value status_vec0 = builder.create( + forOp.getLoc(), builder.getI32Type(), ssb_vec0_ptr); + + Value status_vec1 = builder.create( + forOp.getLoc(), builder.getI32Type(), ssb_vec1_ptr); + + getWaitType("cube", forOp, WaitType, AllocType); + + for (auto i = 0; i < WaitType.size(); i++) { + auto correnspondAlloc = CubeBitMap[AllocType[i]]; + auto i32ConstAttr = + mlir::IntegerAttr::get(builder.getI32Type(), 1 << correnspondAlloc); + auto buf_constant_set = builder.create( + scopeOp->getLoc(), builder.getI32Type(), i32ConstAttr); + Value bufi_vec0_val = builder.create( + forOp.getLoc(), status_vec0, buf_constant_set); + Value bufi_vec1_val = builder.create( + forOp.getLoc(), status_vec1, buf_constant_set); + Value flag_bufi_vec0; + Value flag_bufi_vec1; + // 创建比较操作 + if (WaitType[i] == 0) { + flag_bufi_vec0 = builder.create( + forOp.getLoc(), arith::CmpIPredicate::eq, bufi_vec0_val, c0); + flag_bufi_vec1 = builder.create( + forOp.getLoc(), arith::CmpIPredicate::eq, bufi_vec1_val, c0); + } else { + flag_bufi_vec0 = builder.create( + forOp.getLoc(), arith::CmpIPredicate::eq, bufi_vec0_val, + buf_constant_set); + flag_bufi_vec1 = builder.create( + forOp.getLoc(), arith::CmpIPredicate::eq, bufi_vec1_val, + buf_constant_set); + } + // 创建最终的and操作 + Value bufi_val = builder.create( + forOp.getLoc(), flag_bufi_vec0, flag_bufi_vec1); + buf_vals.push_back(bufi_val); + } + + } else { + builder.setInsertionPointToStart(&scopeOp->getRegion(0).front()); + Value c0 = + builder.create(forOp.getLoc(), 0, 32 // 值32，64位 + ); + auto i64Type = builder.getIntegerType(64); + // %sub_id = hivm.hir.get_sub_block_idx -> i64 + // 这里假设有一个getSubBlockIdxOp操作 + Value subId = builder.create(scopeOp->getLoc(), i64Type); + auto i64ConstAttr = mlir::IntegerAttr::get(i64Type, 32); + Value cstOffset = builder.create( + scopeOp->getLoc(), i64Type, i64ConstAttr); + Value ssbAddrOffset = + builder.create(scopeOp->getLoc(), subId, cstOffset); + Value ssbAddr = builder.create(scopeOp->getLoc(), + ssbAddrOffset, cstOffset); + builder.setInsertionPointToStart(&forOp->getRegion(0).front()); + // 创建inttoptr操作 + Value ssb_cube_ptr = builder.create( + forOp.getLoc(), + LLVM::LLVMPointerType::get(builder.getContext(), 11), // 地址空间11 + ssbAddr); + bufferPtrs.push_back(ssb_cube_ptr); + // 创建load操作 + Value status_cube = builder.create( + forOp.getLoc(), builder.getI32Type(), ssb_cube_ptr); + + getWaitType("vector", forOp, WaitType, AllocType); + for (auto i = 0; i < WaitType.size(); i++) { + auto correnspondAlloc = VecBitMap[AllocType[i]]; + auto i32ConstAttr = + mlir::IntegerAttr::get(builder.getI32Type(), 1 << correnspondAlloc); + auto buf_constant_set = builder.create( + scopeOp->getLoc(), builder.getI32Type(), i32ConstAttr); + Value bufi_cube_val = builder.create( + forOp.getLoc(), status_cube, buf_constant_set); + + Value flag_bufi_cube; + // 创建比较操作 + if (WaitType[i] == 0) { + flag_bufi_cube = builder.create( + forOp.getLoc(), arith::CmpIPredicate::eq, bufi_cube_val, c0); + } else { + flag_bufi_cube = builder.create( + forOp.getLoc(), arith::CmpIPredicate::eq, bufi_cube_val, + buf_constant_set); + } + buf_vals.push_back(flag_bufi_cube); + } + } + int bufIdx = 0; + int groupIdx = 0; + + for (const auto &pair : bufferMap) { + if (bufferMap[groupIdx] == 0) { + continue; + } + + // 获取对应的region迭代参数 + Value cnti = builder.create( + forOp.getLoc(), arith::CmpIPredicate::slt, + forOp.getRegionIterArgs()[forOp.getRegionIterArgs().size() - + (bufferMap.size() - 1 - groupIdx)], + subLoopValue); + + // 计算该组中所有buffer值的AND + Value finalBufVal = buf_vals[bufIdx]; + for (int count = 1; count < bufferMap[groupIdx]; count++) { + finalBufVal = builder.create(forOp.getLoc(), finalBufVal, + buf_vals[bufIdx + count]); + } + + auto cond = + builder.create(forOp.getLoc(), finalBufVal, cnti); + if_conditions.push_back(cond); + + // 更新索引 + bufIdx += bufferMap[groupIdx]; + groupIdx++; + } + int ifIndex = 0; + int acc = 0; + int bufferBit = 0; + for (int i = 0; i < CubeBitMap.size(); i++) { + bufferBit += (1 << i); + } + forOp.getBody()->walk([&](Operation *op) { + auto ifOp = dyn_cast(op); + if (ifOp && ifOp->hasAttr("ssbuffer")) { + // 获取then区域 + Block *thenBlock = &ifOp.getThenRegion().front(); + + // 找到then区域中的yield操作 + Operation *yieldOp = nullptr; + for (auto &op : *thenBlock) { + if (isa(op)) { + yieldOp = &op; + break; + } + } + if (yieldOp) { + builder.setInsertionPoint(yieldOp); + + if (isAIC) { + // 创建插入的语句 + // %status_v2 = llvm.load %ssb_ptr : !llvm.ptr<11> -> i32 + Value status_v2_0 = builder.create( + yieldOp->getLoc(), + builder.getIntegerType(32), // i32类型 + bufferPtrs[0] // 假设ssb_ptr已在作用域中定义 + ); + Value status_v2_1 = builder.create( + yieldOp->getLoc(), + builder.getIntegerType(32), // i32类型 + bufferPtrs[1] // 假设ssb_ptr已在作用域中定义 + ); + Value buf_val_new_0 = status_v2_0; + Value buf_val_new_1 = status_v2_1; + auto bufferNum = bufferMap[ifIndex]; + for (int i = 0; i < bufferNum; i++) { + if (WaitType[acc + i] == 0) { + auto correnspondAlloc = CubeBitMap[AllocType[acc + i]]; + auto i32ConstAttr = mlir::IntegerAttr::get(builder.getI32Type(), + 1 << correnspondAlloc); + auto buf_constant_set = builder.create( + scopeOp->getLoc(), builder.getI32Type(), i32ConstAttr); + buf_val_new_0 = builder.create( + yieldOp->getLoc(), buf_val_new_0, + buf_constant_set // 假设buf3_clear已在作用域中定义 + ); + buf_val_new_1 = builder.create( + yieldOp->getLoc(), buf_val_new_1, + buf_constant_set // 假设buf3_clear已在作用域中定义 + ); + } else { + auto correnspondAlloc = CubeBitMap[AllocType[acc + i]]; + int bitPos = correnspondAlloc; + int basePattern = bufferBit; + int finalValue = basePattern ^ (1 << bitPos); + auto i32ConstAttr = + mlir::IntegerAttr::get(builder.getI32Type(), finalValue); + auto buf_constant_set = builder.create( + scopeOp->getLoc(), builder.getI32Type(), i32ConstAttr); + buf_val_new_0 = builder.create( + yieldOp->getLoc(), buf_val_new_0, + buf_constant_set // 假设buf3_clear已在作用域中定义 + ); + buf_val_new_1 = builder.create( + yieldOp->getLoc(), buf_val_new_1, + buf_constant_set // 假设buf3_clear已在作用域中定义 + ); + } + } + acc += bufferNum; + builder.create(yieldOp->getLoc(), buf_val_new_0, + bufferPtrs[0]); + builder.create(yieldOp->getLoc(), buf_val_new_1, + bufferPtrs[1]); + + } else { + // 创建插入的语句 + // %status_v2 = llvm.load %ssb_ptr : !llvm.ptr<11> -> i32 + Value status_v2 = builder.create( + yieldOp->getLoc(), + builder.getIntegerType(32), // i32类型 + bufferPtrs[0] // 假设ssb_ptr已在作用域中定义 + ); + Value buf_val_new = status_v2; + auto bufferNum = bufferMap[ifIndex]; + for (int i = 0; i < bufferNum; i++) { + if (WaitType[acc + i] == 0) { + auto correnspondAlloc = VecBitMap[AllocType[acc + i]]; + auto i32ConstAttr = mlir::IntegerAttr::get(builder.getI32Type(), + 1 << correnspondAlloc); + auto buf_constant_set = builder.create( + scopeOp->getLoc(), builder.getI32Type(), i32ConstAttr); + buf_val_new = builder.create( + yieldOp->getLoc(), buf_val_new, + buf_constant_set // 假设buf3_clear已在作用域中定义 + ); + } else { + auto correnspondAlloc = VecBitMap[AllocType[acc + i]]; + int bitPos = correnspondAlloc; + int basePattern = bufferBit; + int finalValue = basePattern ^ (1 << bitPos); + auto i32ConstAttr = + mlir::IntegerAttr::get(builder.getI32Type(), finalValue); + auto buf_constant_set = builder.create( + scopeOp->getLoc(), builder.getI32Type(), i32ConstAttr); + buf_val_new = builder.create( + yieldOp->getLoc(), buf_val_new, + buf_constant_set // 假设buf3_clear已在作用域中定义 + ); + } + } + acc += bufferNum; + builder.create(yieldOp->getLoc(), buf_val_new, + bufferPtrs[0]); + } + ifIndex++; + } + } + }); + + return if_conditions; +} + +void ReplaceIf(scf::ForOp forOp, SmallVector conditions, + SmallVector &opsToErase, + DenseMap &ifArgMap, OpBuilder &builder, + ModuleOp moduleOp) { + SmallVector ifToProcess; + llvm::outs() << "enter replaceif\n"; + Value step = forOp.getStep(); + auto aiCAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + forOp.getBody()->walk([&](Operation *op) { + auto ifOp = dyn_cast(op); + if (ifOp && ifOp->hasAttr("ssbuffer") && forOp == ifOp->getParentOp()) { + ifToProcess.push_back(ifOp); + } + }); + + IRMapping IRMap; + for (int i = 0; i < ifToProcess.size(); i++) { + auto ifOp = ifToProcess[i]; + auto parentOp = ifOp->getParentOp(); + auto loc = ifOp.getLoc(); + // 获取for循环的iterargs（迭代参数） + auto iterArgs = forOp.getRegionIterArgs(); + if (iterArgs.size() < conditions.size()) { + return; + } + auto thenYieldOp = + dyn_cast(ifOp.getThenRegion().front().getTerminator()); + SmallVector thenResults; + if (thenYieldOp) { + // 如果已有返回值，保留它们 + for (auto result : thenYieldOp.getResults()) { + thenResults.push_back(result); + } + } + // 创建新的else区域，返回两个迭代参数 + SmallVector elseResults; + scf::YieldOp elseYieldOp = nullptr; + bool hasElse = false; + if (!ifOp.getElseRegion().empty()) { + elseYieldOp = + dyn_cast(ifOp.getElseRegion().front().getTerminator()); + hasElse = true; + } + if (elseYieldOp) { + for (auto result : elseYieldOp.getResults()) { + elseResults.push_back(result); + } + } + // 获取最后两个迭代参数 + Value iterArgMinus = iterArgs[iterArgs.size() - (conditions.size() - i)]; + // 创建新的then区域，返回两个迭代参数 + thenResults.push_back(iterArgMinus); + elseResults.push_back(iterArgMinus); + + // 保存原有的操作，以便后续克隆 + SmallVector thenOps; + for (auto &op : ifOp.getThenRegion().front()) { + thenOps.push_back(&op); + } + + SmallVector elseOps; + if (!ifOp.getElseRegion().empty()) { + for (auto &op : ifOp.getElseRegion().front()) { + elseOps.push_back(&op); + } + } + SmallVector resultTypes; + for (auto val : thenResults) { + resultTypes.push_back(val.getType()); + } + // 创建新的scf.if操作 + builder.setInsertionPoint(ifOp); + auto newIfOp = builder.create(loc, resultTypes, conditions[i], + /*withElseRegion=*/true); + newIfOp->setAttr("ssbuffer", builder.getUnitAttr()); + // 处理then区域 + auto &newThenBlock = newIfOp.getThenRegion().front(); + builder.setInsertionPointToStart(&newThenBlock); + + // 克隆then区域的操作 + for (auto op : thenOps) { + if (auto yieldOp = dyn_cast(op)) { + // 处理yield的操作数映射 + SmallVector mappedOperands; + for (auto operand : yieldOp->getOperands()) { + mappedOperands.push_back(IRMap.lookupOrDefault(operand)); + } + // 获取最后两个迭代参数 + Value iterArgMinus = + iterArgs[iterArgs.size() - (conditions.size() - i)]; + + // %ssb_addr = arith.addi %ssb_addr_offset, %c32_i64 : i64 + auto AddIOp = builder.create(forOp->getLoc(), + iterArgMinus, step); + // 这里加个add1 + mappedOperands.push_back(AddIOp); + builder.create(loc, mappedOperands); + } else { + auto newOp = builder.clone(*op, IRMap); + IRMap.map(op->getResults(), newOp->getResults()); + } + } + + // 处理else区域 + auto &newElseBlock = newIfOp.getElseRegion().front(); + builder.setInsertionPointToStart(&newElseBlock); + // 克隆else区域的操作 + if (hasElse) { + for (auto op : elseOps) { + if (auto yieldOp = dyn_cast(op)) { + // 处理yield的操作数映射 + SmallVector mappedOperands; + for (auto operand : yieldOp->getOperands()) { + mappedOperands.push_back(IRMap.lookupOrDefault(operand)); + } + Value iterArgMinus = + iterArgs[iterArgs.size() - (conditions.size() - i)]; + mappedOperands.push_back(iterArgMinus); + builder.create(loc, mappedOperands); + } else { + auto newOp = builder.clone(*op, IRMap); + IRMap.map(op->getResults(), newOp->getResults()); + } + } + } else { + SmallVector cntOperands; + cntOperands.push_back(iterArgMinus); + builder.create(loc, cntOperands); + } + + // 替换原有if操作的使用 + // 首先，将原if操作的结果替换为新if操作的对应结果 + for (unsigned j = 0; j < ifOp.getNumResults(); ++j) { + ifOp.getResult(j).replaceAllUsesWith(newIfOp.getResult(j)); + } + // 获取新if操作所在的块 + Block *newIfBlock = ifOp->getBlock(); + // 在for循环体内替换迭代参数的使用 + forOp.getBody()->walk([&](Operation *op) { + // 检查操作是否与新ifOp在同一个块中 + Block *opBlock = op->getBlock(); + if (opBlock != newIfBlock) { + // 不在同一个块中，跳过 + return; + } + if (op->isBeforeInBlock(newIfOp)) { + return; // 只处理if操作之后的use + } + for (unsigned j = 0; j < op->getNumOperands(); ++j) { + for (auto argIndex = 0; argIndex < conditions.size(); argIndex++) { + // 获取最后两个迭代参数 + Value iterArgMinus = + iterArgs[iterArgs.size() - (conditions.size() - i)]; + if (op->getOperand(j) == iterArgMinus) { + op->setOperand(j, + newIfOp.getResults()[newIfOp.getNumResults() - 1]); + } + } + } + }); + + // // 删除原有的if操作 + opsToErase.push_back(ifOp); + if (ifArgMap.find(newIfOp) == ifArgMap.end()) { + ifArgMap[newIfOp] = iterArgMinus; + } + } +} + +int getNestingDepth(scf::ForOp forOp) { + int depth = 0; + Operation *op = forOp.getOperation(); + while (op) { + if (op->getDialect() && op->getDialect()->getNamespace() == "scf") { + ++depth; + } + op = op->getParentOp(); + } + return depth; +} + +void printDenseMap(const mlir::DenseMap &Map) { + for (const auto &pair : Map) { + auto op = pair.first; + auto bitValue = pair.second; + + // ✅ 正确打印 Operation 的方法：print(llvm::outs()) + llvm::outs() << "Operation: \n"; + llvm::outs().flush(); + llvm::outs() << "Operation Name: " << *op << "\n"; + + // 打印对应的值 + llvm::outs() << " | Bit Value: " << bitValue << " | allocmap\n"; + llvm::outs().flush(); + llvm::outs() << "------------------------------------------------\n\n"; + llvm::outs().flush(); + } + llvm::outs().flush(); +} + +void getAllocBit(ModuleOp module, DenseMap &VecBitMap, + DenseMap &CubeBitMap, OpBuilder builder) { + auto aiCAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + auto scalarWaitPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + auto cubeWaitPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_FIX); + auto vectorWaitPipe = + PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_MTE3); + + SmallVector scopeOpToEdit; + module.walk( + [&](scope::ScopeOp scopeOp) { scopeOpToEdit.push_back(scopeOp); }); + + int cubeAcc = 0; + int vecAcc = 0; + for (auto scopeOp : scopeOpToEdit) { + SmallVector excludedValues; + if (!scopeOp->hasAttr("hivm.tcore_type")) + continue; + + auto attr = scopeOp->getAttr("hivm.tcore_type"); + if (attr == aiCAttr) { + // CUBE 类型：分配 cubeAcc 自增 ID + scopeOp.walk([&](SyncBlockWaitOp waitOp) { + auto parentOp = waitOp->getParentOp(); + if (isa(parentOp) && parentOp->hasAttr("ssbuffer")) { + auto waitPipe = waitOp.getPipe(); + if (waitPipe != scalarWaitPipe) { + auto allocOp = findFirstTargetOpAfterWait(waitOp, excludedValues); + if (VecBitMap.count(allocOp)) { + CubeBitMap[allocOp] = VecBitMap[allocOp]; + cubeAcc++; + } else { + CubeBitMap[allocOp] = cubeAcc++; + } + } + } + }); + } else { + // VEC 类型：分配 vecAcc 自增 ID + scopeOp.walk([&](SyncBlockWaitOp waitOp) { + auto parentOp = waitOp->getParentOp(); + if (isa(parentOp) && parentOp->hasAttr("ssbuffer")) { + auto waitPipe = waitOp.getPipe(); + if (waitPipe != scalarWaitPipe) { + auto allocOp = findFirstTargetOpAfterWait(waitOp, excludedValues); + if (!VecBitMap.count(allocOp)) { + VecBitMap[allocOp] = vecAcc++; + } + } + } + }); + } + } + + // 记录【自增ID的最大值】，用于后面偏移计算 + const int maxVecId = vecAcc - 2; + const int maxCubeId = cubeAcc - 2; + // ======================== 阶段 2：分配【偏移ID】给 + // waitOp（绝对不重复）======================== + for (auto scopeOp : scopeOpToEdit) { + if (!scopeOp->hasAttr("hivm.tcore_type")) + continue; + + auto attr = scopeOp->getAttr("hivm.tcore_type"); + if (attr == aiCAttr) { + // CUBE：waitOp = maxCubeId + 1 + flagid + scopeOp.walk([&](SyncBlockWaitOp waitOp) { + auto parentOp = waitOp->getParentOp(); + if (isa(parentOp) && parentOp->hasAttr("ssbuffer")) { + auto setPipe = waitOp.getTpipe(); + auto waitPipe = waitOp.getPipe(); + if (setPipe != scalarWaitPipe && waitPipe == scalarWaitPipe) { + uint32_t flagid = 0; + if (auto attr = + waitOp->getAttrOfType("ssbuf.flagid")) { + flagid = attr.getUInt(); + } + // 核心修改：偏移分配，不再写死 16 + CubeBitMap[waitOp] = + (maxCubeId >= 0 ? (maxCubeId + 1) : 0) + flagid; + } + } + }); + } else { + // VEC：waitOp = maxVecId + 1 + flagid + scopeOp.walk([&](SyncBlockWaitOp waitOp) { + auto parentOp = waitOp->getParentOp(); + if (isa(parentOp) && parentOp->hasAttr("ssbuffer")) { + auto setPipe = waitOp.getTpipe(); + auto waitPipe = waitOp.getPipe(); + if (setPipe != scalarWaitPipe && waitPipe == scalarWaitPipe) { + uint32_t flagid = 0; + if (auto attr = + waitOp->getAttrOfType("ssbuf.flagid")) { + flagid = attr.getUInt(); + } + // 核心修改：偏移分配 + VecBitMap[waitOp] = (maxVecId >= 0 ? (maxVecId + 1) : 0) + flagid; + } + } + }); + } + } +} + +void modifyForIterargDeps(scf::ForOp forOp, + DenseMap ifCounters) { + Value iterArg = forOp.getInductionVar(); + + for (Operation &op : forOp.getBody()->without_terminator()) { + if (auto ifOp = dyn_cast(op)) { + if (ifCounters.find(ifOp) != ifCounters.end()) { + Value counter = ifCounters[ifOp]; + + ifOp.walk([&](Operation *opInIf) { + for (auto [i, operand] : llvm::enumerate(opInIf->getOperands())) { + if (operand == iterArg) { + opInIf->setOperand(i, counter); + } + } + }); + } + } + } +} + +void FlowSssbuf(ModuleOp module) { + mlir::OpBuilder builder(module.getContext()); + // 收集所有需要转换的循环 + SmallVector targetLoops; + llvm::outs() << "enter flowsssbuf\n\n"; + module.walk([&](Operation *op) { + if (auto forOp = dyn_cast(op)) { + // 检查循环是否包含特定的 sync_block_set 操作 + bool hasSyncBlockSet = false; + forOp.walk([&](Operation *op) { + if (isa(op)) { + if (auto ifOp = dyn_cast(op->getParentOp())) { + if (forOp == ifOp->getParentOp() && ifOp->hasAttr("ssbuffer")) { + hasSyncBlockSet = true; + } + } + } + }); + + if (hasSyncBlockSet) { + if (llvm::find(targetLoops, forOp) == targetLoops.end()) { + targetLoops.push_back(forOp); + } + } + } + }); + llvm::outs() << "enter flowsssbuf\n\n"; + + SmallVector transformLoops; + // 转换每个目标循环 + for (scf::ForOp forOp : targetLoops) { + auto newforOp = transformLoop(forOp, builder); + } + + module.walk([&](Operation *op) { + if (auto forOp = dyn_cast(op)) { + // 检查循环是否包含特定的 sync_block_set 操作 + bool hasSyncBlockSet = false; + forOp.walk([&](Operation *op) { + if (isa(op)) { + if (auto ifOp = dyn_cast(op->getParentOp())) { + if (forOp == ifOp->getParentOp() && ifOp->hasAttr("ssbuffer")) { + hasSyncBlockSet = true; + } + } + } + }); + + if (hasSyncBlockSet) { + if (llvm::find(transformLoops, forOp) == transformLoops.end()) { + transformLoops.push_back(forOp); + } + } + } + }); + + llvm::sort(transformLoops, [](scf::ForOp a, scf::ForOp b) { + return getNestingDepth(a) > getNestingDepth(b); + }); + DenseMap VecBitMap; + DenseMap CubeBitMap; + getAllocBit(module, VecBitMap, CubeBitMap, builder); + printDenseMap(CubeBitMap); + printDenseMap(VecBitMap); + SmallVector opsToErase; + for (scf::ForOp forOp : transformLoops) { + DenseMap ifArgMap; + llvm::outs() << "before replaceif\n"; + auto bufvals = addBufValLoop(forOp, VecBitMap, CubeBitMap, builder); + ReplaceIf(forOp, bufvals, opsToErase, ifArgMap, builder, module); + llvm::outs() << "after replaceif\n"; + for (const auto &pair : ifArgMap) { + auto val = pair.first; + auto bitValue = pair.second; + llvm::outs() << val << " " << bitValue << " ifargmrp\n\n\n"; + llvm::outs().flush(); + } + + modifyForIterargDeps(forOp, ifArgMap); + } + for (auto op : opsToErase) { + op->erase(); + } +} + +bool isTransOp(mlir::Operation *op) { + auto fixpipeOp = dyn_cast(op); + if (fixpipeOp) + return true; + + auto copyOp = dyn_cast(op); + if (!copyOp) + return false; + else { + + Value copySrc = copyOp.getODSOperands(0).front(); + MemRefType copySrcTy = dyn_cast(copySrc.getType()); + auto SrcAddrSpace = + dyn_cast_or_null(copySrcTy.getMemorySpace()); + bool isSrcUbSpace = + SrcAddrSpace.getAddressSpace() == hivm::AddressSpace::UB; + + Value copyDst = copyOp.getODSOperands(1).front(); + MemRefType copyDstTy = dyn_cast(copyDst.getType()); + auto DstAddrSpace = + dyn_cast_or_null(copyDstTy.getMemorySpace()); + bool isDstCbufSpace = + DstAddrSpace.getAddressSpace() == hivm::AddressSpace::L1; + + return isSrcUbSpace && isDstCbufSpace; + } +} + +void FindAndMarkBuffer(ModuleOp module) { + OpBuilder builder(module.getContext()); + unsigned int BufferIdx = 0; + Type idxType = builder.getI32Type(); + StringAttr setFlagAttr = builder.getStringAttr("Set flag"); + StringAttr waitFlagAttr = builder.getStringAttr("Wait flag"); + IntegerAttr idxAttr = builder.getI32IntegerAttr(BufferIdx); + + module.walk([&](mlir::Operation *op) { + if (isTransOp(op)) { + llvm::outs() << "Buffer idx" << BufferIdx << "\n"; + llvm::outs() << "Trans Op" << *op << "\n"; + Value SharedBuffer; + if (auto fixpipeOp = dyn_cast(op)) { + SharedBuffer = fixpipeOp.getODSOperands(1).front(); + } else { + auto copyOp = dyn_cast(op); + SharedBuffer = copyOp.getODSOperands(1).front(); + } + llvm::outs() << "SharedBuffer" << SharedBuffer << "\n"; + + if (!SharedBuffer) { + op->emitWarning("fixpipe op has empty output operand!"); + return; + } + + // 在Buffer的生产op后set flag标记，在Buffer消费op前增加wait flag标记 + op->setAttr("Buffer idx", builder.getI32IntegerAttr(BufferIdx)); + op->setAttr("Wait Flag", builder.getI32IntegerAttr(0)); + op->setAttr("Set Flag", builder.getI32IntegerAttr(1)); + + for (Operation *consumerOp : SharedBuffer.getUsers()) { + if (consumerOp == op) + continue; + if (!consumerOp) + continue; + + llvm::outs() << "consumerOp: " << *consumerOp << "\n"; + + consumerOp->setAttr("Buffer idx", builder.getI32IntegerAttr(BufferIdx)); + consumerOp->setAttr("Wait Flag", builder.getI32IntegerAttr(0)); + } + BufferIdx++; + } + }); +} + +// 结构体存 wait-set 区块信息 +struct WaitSetRegion { + Operation *waitOp; + Operation *lastSetOp; + SmallVector opsToMove; + bool hasCopyOrFixpipe = false; +}; + +struct MergedRegion { + SmallVector regions; + SmallVector opsToMove; + SmallVector yieldValues; + SmallVector resultTypes; +}; + +void MoveIterArgUsersIntoIf(scf::ForOp forOp, + SmallVector &mergedRegions) { + + Block &body = forOp.getRegion().front(); + + // iter_arg -> mergedRegion index + DenseMap iterArgToRegion; + + for (int r = 0; r < mergedRegions.size(); ++r) { + MergedRegion &mr = mergedRegions[r]; + + for (Operation *op : mr.opsToMove) { + for (Value v : op->getOperands()) { + if (auto barg = mlir::dyn_cast(v)) { + if (barg.getOwner() == &body) { + iterArgToRegion.try_emplace(barg, r); + } + } + } + } + } + + if (iterArgToRegion.empty()) + return; + + // 找最后一个 mergedRegion 的最后一个 op + Operation *lastOp = nullptr; + for (MergedRegion &mr : mergedRegions) + lastOp = mr.opsToMove.back(); + + if (!lastOp) + return; + + DenseMap opIndex; + int idx = 0; + for (Operation &op : body) + opIndex[&op] = idx++; + + int startIdx = opIndex[lastOp] + 1; + + // 扫描 for body 尾部 op + for (Operation &op : body) { + if (opIndex[&op] < startIdx) + continue; + + llvm::SmallDenseSet usedRegions; + for (Value v : op.getOperands()) { + if (auto barg = mlir::dyn_cast(v)) { + auto it = iterArgToRegion.find(barg); + if (it != iterArgToRegion.end()) + usedRegions.insert(it->second); + } + } + + // 必须且只能依赖一个 mergedRegion + if (usedRegions.size() != 1) + continue; + + int target = *usedRegions.begin(); + + mergedRegions[target].opsToMove.push_back(&op); + } +} + +void ComputeYieldForMergedRegion(MergedRegion &mr, Block &body) { + + mr.yieldValues.clear(); + mr.resultTypes.clear(); + + SmallPtrSet inRegion(mr.opsToMove.begin(), + mr.opsToMove.end()); + + for (Operation *op : mr.opsToMove) { + for (Value res : op->getResults()) { + bool usedOutside = false; + + for (OpOperand &use : res.getUses()) { + Operation *user = use.getOwner(); + + // 不在同一个 for body，交给外层处理（通常不会出现） + if (user->getBlock() != &body) + continue; + + // 只要有一个 use 在 region 外，就必须 yield + if (!inRegion.contains(user)) { + usedOutside = true; + break; + } + } + + if (usedOutside) { + mr.yieldValues.push_back(res); + mr.resultTypes.push_back(res.getType()); + } + } + } +} + +static void ComputeYieldForMergedRegionV2(MergedRegion &mr, Block &body) { + + mr.yieldValues.clear(); + mr.resultTypes.clear(); + + // 当前 region 内的 ops + SmallPtrSet inRegion(mr.opsToMove.begin(), + mr.opsToMove.end()); + + for (Operation *op : mr.opsToMove) { + for (Value res : op->getResults()) { + + bool usedOutside = false; + + for (OpOperand &use : res.getUses()) { + Operation *user = use.getOwner(); + + // 如果使用在 region 内部 op，跳过 + if (inRegion.contains(user)) + continue; + + // 使用在 region 外部，包括嵌套 region 内部的 block + usedOutside = true; + break; + } + + if (usedOutside) { + mr.yieldValues.push_back(res); + mr.resultTypes.push_back(res.getType()); + } + } + } +} + +static void ComputeYieldForMergedRegionV3(MergedRegion &mr) { + mr.yieldValues.clear(); + mr.resultTypes.clear(); + + // 用 DenseSet 暂存当前 region 的所有 ops + DenseSet regionOps(mr.opsToMove.begin(), mr.opsToMove.end()); + + for (Operation *op : mr.opsToMove) { + for (Value res : op->getResults()) { + + bool needsYield = false; + + for (OpOperand &use : res.getUses()) { + Operation *user = use.getOwner(); + + // 如果 user 不在当前 region，则需要 yield + if (!regionOps.contains(user)) { + needsYield = true; + break; + } + } + + if (needsYield) { + mr.yieldValues.push_back(res); + mr.resultTypes.push_back(res.getType()); + } + } + } +} + +// 递归收集 op 和它所有 region 内的 ops +static void CollectAllNestedOps(Operation *op, + DenseSet ®ionOps) { + if (!op) + return; + + if (regionOps.contains(op)) + return; // 已经收集过 + + regionOps.insert(op); + + // 遍历所有 region，递归收集 + for (Region ®ion : op->getRegions()) { + for (Block &block : region) { + for (Operation &nestedOp : block) { + CollectAllNestedOps(&nestedOp, regionOps); + } + } + } +} + +static void ComputeYieldForMergedRegionV4(MergedRegion &mr) { + mr.yieldValues.clear(); + mr.resultTypes.clear(); + + // 用 DenseSet 暂存当前 region 的所有 ops + // 初始 DenseSet: 顶层 opsToMove + DenseSet regionOps; + for (Operation *op : mr.opsToMove) { + CollectAllNestedOps(op, regionOps); // 完整展开嵌套 + } + + for (Operation *op : mr.opsToMove) { + for (Value res : op->getResults()) { + + bool needsYield = false; + + for (OpOperand &use : res.getUses()) { + Operation *user = use.getOwner(); + + // 如果 user 不在当前 region，则需要 yield + if (!regionOps.contains(user)) { + needsYield = true; + break; + } + } + + if (needsYield) { + mr.yieldValues.push_back(res); + mr.resultTypes.push_back(res.getType()); + } + } + } +} + +int findTargetRegion(Operation *startOp, Block &body, + DenseMap &opToRegion) { + + SmallVector worklist{startOp}; + SmallPtrSet visited; + + while (!worklist.empty()) { + Operation *op = worklist.pop_back_val(); + if (!visited.insert(op).second) + continue; + + auto it = opToRegion.find(op); + if (it != opToRegion.end()) + return it->second; + + for (Value operand : op->getOperands()) { + if (isa(operand)) + continue; + + Operation *defOp = operand.getDefiningOp(); + if (defOp && defOp->getBlock() == &body) + worklist.push_back(defOp); + } + } + + return -1; +} + +void greedyAbsorbToRegion(Operation *startOp, int regionIdx, int lowerBound, + Block &body, DenseMap &opIndex, + DenseMap &opToRegion, + SmallVector &mergedRegions) { + + auto &mr = mergedRegions[regionIdx]; + + SmallVector worklist; + SmallPtrSet visited(mr.opsToMove.begin(), + mr.opsToMove.end()); + + // 先把 startOp 本身吸收（如果还没被吸收） + if (!opToRegion.count(startOp)) { + mr.opsToMove.push_back(startOp); + opToRegion[startOp] = regionIdx; + visited.insert(startOp); + } + + worklist.push_back(startOp); + + while (!worklist.empty()) { + Operation *op = worklist.pop_back_val(); + + for (Value operand : op->getOperands()) { + if (isa(operand)) + continue; + + Operation *defOp = operand.getDefiningOp(); + if (!defOp || defOp->getBlock() != &body) + continue; + + int defIdx = opIndex[defOp]; + + // 超过前一个 region 的末尾 + if (defIdx < lowerBound) + continue; + + auto it = opToRegion.find(defOp); + + // 不能跨到其他 region + if (it != opToRegion.end() && it->second != regionIdx) + continue; + + // 去重 + if (!visited.insert(defOp).second) + continue; + + // 吸收 defOp + mr.opsToMove.push_back(defOp); + opToRegion[defOp] = regionIdx; + worklist.push_back(defOp); + } + } +} + +SmallVector +getOperationInput(Operation *op, SmallVector dependValues, + DenseMap>> + &collectDepValueMap) { + // Analyse each Op's input + DenseSet opInput; + if (isa(op) || isa(op)) { + SmallVector regionBlocks; + if (auto ifOp = dyn_cast(op)) { + regionBlocks.push_back(&(ifOp.getThenRegion().front())); + regionBlocks.push_back(&(ifOp.getElseRegion().front())); + } else { + auto forOp = dyn_cast(op); + regionBlocks.push_back(forOp.getBody()); + } + + // recursively walk scf op + for (Block *curBlock : regionBlocks) { + for (auto &curOp : *curBlock) { + for (auto operand : + getOperationInput(&curOp, dependValues, collectDepValueMap)) { + Operation *defOp; + if (auto blockArg = dyn_cast(operand)) { + Block *ownerBlock = blockArg.getOwner(); + defOp = ownerBlock->getParentOp(); + } else { + defOp = operand.getDefiningOp(); + } + Block *defBlock = defOp->getBlock(); + + if (!(defOp == op || llvm::is_contained(regionBlocks, defBlock))) { + opInput.insert(operand); + } + } + } + } + SmallVector retVector(opInput.begin(), opInput.end()); + return retVector; + } else { + SmallVector operands = op->getOperands(); + // store ifresult value that will be replaced + for (auto operand : operands) { + if (llvm::is_contained(dependValues, operand)) { + if (collectDepValueMap.find(operand) != collectDepValueMap.end()) { + collectDepValueMap[operand].second.push_back(op); + } else { + SmallVector userOps; + userOps.push_back(op); + collectDepValueMap[operand] = {operand, userOps}; + } + } + } + return operands; + } +} + +SmallVector collectDepValuesCalculation( + DenseSet forRegionOps, DenseSet regionOps, + Operation *op, SmallVector dependValues, + DenseMap>> + &collectDepValueMap) { + DenseSet collectOps; + std::deque opStack; + bool flag = false; + + opStack.push_back(op); + while (opStack.size()) { + Operation *curOp = opStack.front(); + opStack.pop_front(); + + for (auto operand : + getOperationInput(curOp, dependValues, collectDepValueMap)) { + if (llvm::is_contained(dependValues, operand)) { + flag = true; + } + + Operation *parentOp = operand.getDefiningOp(); + if (llvm::is_contained(regionOps, parentOp)) { + opStack.push_back(parentOp); + continue; + } else if (llvm::is_contained(forRegionOps, parentOp)) { + opStack.push_back(parentOp); + collectOps.insert(parentOp); + } + } + } + + if (flag) { + SmallVector retVector(collectOps.begin(), collectOps.end()); + return retVector; + } else { + collectDepValueMap.clear(); + SmallVector emptyVector; + emptyVector.clear(); + return emptyVector; + } +} + +void copyOpsToMergedRegion( + scf::ForOp forOp, SmallVector collectOps, + MergedRegion &mergedRegion, + DenseMap>> + &collectDepValueMap) { + Block *forBodyBlock = forOp.getBody(); + OpBuilder builder(forOp); + SmallVector clonedOps; + IRMapping mapper; + + // copy calculation of ifreult value related to load/store op + int cnt = 0; + for (Operation &origOp : forBodyBlock->without_terminator()) { + if (cnt >= collectOps.size()) + break; + + if (llvm::is_contained(collectOps, &origOp)) { + builder.setInsertionPointAfter(&origOp); + + Operation *clonedOp = (&origOp)->clone(mapper); + builder.insert(clonedOp); + mapper.map(&origOp, clonedOp); + + clonedOps.push_back(clonedOp); + cnt++; + + // replace the ifresult value by new cloned op's result + SmallVector results = origOp.getResults(); + for (auto [idx, result] : llvm::enumerate(origOp.getResults())) { + if (collectDepValueMap.find(result) != collectDepValueMap.end()) { + collectDepValueMap[result].first = clonedOp->getResult(idx); + } + } + } + } + + DenseSet mergedRegionOps; + for (Operation *op : mergedRegion.opsToMove) { + CollectAllNestedOps(op, mergedRegionOps); + } + + // replace the ifresult value by new cloned op's result + for (Operation *op : mergedRegionOps) { + for (auto [idx, operand] : llvm::enumerate(op->getOperands())) { + if (collectDepValueMap.find(operand) != collectDepValueMap.end()) { + op->setOperand(idx, collectDepValueMap[operand].first); + } + } + } + + // update MergedRegion + clonedOps.append(mergedRegion.opsToMove); + mergedRegion.opsToMove = clonedOps; +} + +void copyLoadCalculation(scf::ForOp forOp, SmallVector dependValues, + SmallVector &mergedRegions) { + mlir::Operation *parentOp = forOp->getParentOp(); + mlir::Operation *scopeOp = nullptr; + while (parentOp) { + if (dyn_cast(parentOp)) { + scopeOp = parentOp; + break; + } + parentOp = parentOp->getParentOp(); + } + auto coreTypeAttr = + scopeOp->getAttrOfType(hivm::TCoreTypeAttr::name); + // only process the vector core + if (coreTypeAttr.getTcoretype() == hivm::TCoreType::CUBE) { + return; + } + + // recursively collect all op in forOp + DenseSet forRegionOps; + for (Operation &op : forOp.getBody()->without_terminator()) { + CollectAllNestedOps(&op, forRegionOps); + } + + for (MergedRegion &mr : mergedRegions) { + DenseSet regionOps; + for (Operation *op : mr.opsToMove) { + CollectAllNestedOps(op, regionOps); + } + + for (Operation *op : regionOps) { + if (isa(op) || isa(op)) { + // recusively check that whether load/store op's operands originated + // from if results + DenseMap>> + collectDepValueMap; + SmallVector collectOps = collectDepValuesCalculation( + forRegionOps, regionOps, op, dependValues, collectDepValueMap); + copyOpsToMergedRegion(forOp, collectOps, mr, collectDepValueMap); + } + } + } +} + +// 以 forOp 的 yield value 为中心 +// 决定它应该归属哪个 mergedRegion, 然后再向前吸 operand +void ExpandMergedRegionOpsForAIV(scf::ForOp forOp, + SmallVector &mergedRegions) { + + Block &body = forOp.getRegion().front(); + + // 记录 block 中 op 顺序 + DenseMap opIndex; + int idx = 0; + for (Operation &op : body) + opIndex[&op] = idx++; + + // 建立 op -> region 映射 + DenseMap opToRegion; + for (int r = 0; r < mergedRegions.size(); ++r) + for (Operation *op : mergedRegions[r].opsToMove) + opToRegion[op] = r; + + // 取 scf.yield + auto yieldOp = cast(body.getTerminator()); + + // 依次处理每个 yield value（按编号顺序） + for (Value yv : yieldOp.getOperands()) { + + Operation *defOp = yv.getDefiningOp(); + if (!defOp || defOp->getBlock() != &body) + continue; + + int targetRegion = -1; + + // 如果已经在 region 中 + auto it = opToRegion.find(defOp); + if (it != opToRegion.end()) { + targetRegion = it->second; + } else { + // 否则向前搜索确定归属 + targetRegion = findTargetRegion(defOp, body, opToRegion); + } + + if (targetRegion == -1) + continue; + + // 计算边界 lowerBound + int lowerBound = 0; + + if (targetRegion > 0) { + Operation *prevLast = mergedRegions[targetRegion - 1].opsToMove.back(); + lowerBound = opIndex[prevLast] + 1; + } + + // 真正贪心吸收 + greedyAbsorbToRegion(defOp, targetRegion, lowerBound, body, opIndex, + opToRegion, mergedRegions); + } + + // 每个 region 内按 block 顺序排序 + for (auto &mr : mergedRegions) { + llvm::sort(mr.opsToMove, [&](Operation *a, Operation *b) { + return opIndex[a] < opIndex[b]; + }); + } +} + +// 以 mergedRegion 为中心, 向前吸 operand +void ExpandMergedRegionOpsForAIC(scf::ForOp forOp, + SmallVector &mergedRegions) { + Block &body = forOp.getRegion().front(); + + // 记录每个 mergedRegion 的起始 op index + DenseMap opIndex; + int idx = 0; + for (Operation &op : body) { + opIndex[&op] = idx++; + } + + for (int r = 0; r < mergedRegions.size(); ++r) { + MergedRegion &mr = const_cast(mergedRegions[r]); + + // 本 mergedRegion 的最早 op + Operation *firstOp = mr.opsToMove.front(); + int lowerBound = 0; + + // 边界: 前一个 mergedRegion 的最后一个 op + if (r > 0) { + Operation *prevLast = mergedRegions[r - 1].opsToMove.back(); + lowerBound = opIndex[prevLast] + 1; + } + + SmallVector worklist(mr.opsToMove.begin(), mr.opsToMove.end()); + SmallPtrSet visited(mr.opsToMove.begin(), + mr.opsToMove.end()); + + while (!worklist.empty()) { + Operation *op = worklist.pop_back_val(); + + // 往前吸收operand + for (Value operand : op->getOperands()) { + // BlockArgument + if (mlir::isa(operand)) + continue; + + Operation *defOp = operand.getDefiningOp(); + if (!defOp) + continue; + + // 不在 for body + if (defOp->getBlock() != &body) + continue; + + int defIdx = opIndex[defOp]; + + // 超出允许向前吸收的边界 + if (defIdx < lowerBound) + continue; + + // 已经在 opsToMove + if (!visited.insert(defOp).second) + continue; + + // 吸收这个 defOp + mr.opsToMove.push_back(defOp); + worklist.push_back(defOp); + } + } + + // 最后按原 block 顺序排序 + llvm::sort(mr.opsToMove, [&](Operation *a, Operation *b) { + return opIndex[a] < opIndex[b]; + }); + } +} + +static void pullInRegionDependencies(Operation *regionOp, int regionId, + DenseMap &opToRegion, + Block &body) { + + SmallVector worklist; + + // 先把 region 内的 op 放进去 + for (Region ®ion : regionOp->getRegions()) + for (Block &block : region) + for (Operation &inner : block) + worklist.push_back(&inner); + + SmallPtrSet visited; + + while (!worklist.empty()) { + Operation *innerOp = worklist.pop_back_val(); + + if (!visited.insert(innerOp).second) + continue; + + // operand 的 defining op + for (Value operand : innerOp->getOperands()) { + + Operation *def = operand.getDefiningOp(); + if (!def) + continue; + + if (def->getBlock() != &body) + continue; + + if (!opToRegion.count(def)) { + + opToRegion[def] = regionId; + + // 如果 def 也是 region-op，继续扩展 + if (def->getNumRegions() > 0) + worklist.push_back(def); + } + } + + // 继续遍历 region + for (Region &r : innerOp->getRegions()) + for (Block &b : r) + for (Operation &child : b) + worklist.push_back(&child); + } +} + +// BFS 查找某个 op 最早被哪个 region 使用 +static int findEarliestRegion(Operation *startOp, + const DenseMap &seedRegionMap, + Block &body) { + + SmallVector worklist{startOp}; + SmallPtrSet visited; + int earliestRegion = -1; + + while (!worklist.empty()) { + Operation *op = worklist.pop_back_val(); + + if (!visited.insert(op).second) + continue; + + for (Value result : op->getResults()) { + for (OpOperand &use : result.getUses()) { + Operation *user = use.getOwner(); + + if (user->getBlock() != &body) + continue; + + auto it = seedRegionMap.find(user); + if (it != seedRegionMap.end()) { + int region = it->second; + if (earliestRegion == -1 || region < earliestRegion) + earliestRegion = region; + } else { + worklist.push_back(user); + } + } + } + } + + return earliestRegion; +} + +void ExpandMergedRegionOpsForAll(scf::ForOp forOp, + SmallVector &mergedRegions) { + + Block &body = forOp.getRegion().front(); + + // block 内 op 顺序 + DenseMap opIndex; + int idx = 0; + for (Operation &op : body) + opIndex[&op] = idx++; + + // seed region map + DenseMap seedRegionMap; + for (int r = 0; r < mergedRegions.size(); r++) { + for (Operation *op : mergedRegions[r].opsToMove) { + seedRegionMap[op] = r; + } + } + + // 最终 op -> region + DenseMap opToRegion = seedRegionMap; + + // ---------- Step1 顺序扫描 ---------- + for (Operation &op : body) { + + if (isa(&op)) + continue; + + if (opToRegion.count(&op)) + continue; + + int region = findEarliestRegion(&op, seedRegionMap, body); + + if (region != -1) + opToRegion[&op] = region; + } + + // ---------- Step2 region-op 依赖补全 ---------- + for (Operation &op : body) { + + auto it = opToRegion.find(&op); + if (it == opToRegion.end()) + continue; + + if (op.getNumRegions() == 0) + continue; + + pullInRegionDependencies(&op, it->second, opToRegion, body); + } + + // ---------- Step3 append op ---------- + SmallPtrSet seen; + + for (Operation &op : body) { + + auto it = opToRegion.find(&op); + if (it == opToRegion.end()) + continue; + + if (!seen.insert(&op).second) + continue; + + int region = it->second; + mergedRegions[region].opsToMove.push_back(&op); + } + + // ---------- Step4 排序 ---------- + for (auto &mr : mergedRegions) { + + llvm::sort(mr.opsToMove, [&](Operation *a, Operation *b) { + return opIndex[a] < opIndex[b]; + }); + } +} + +void ExpandMergedRegionOpsByInput(scf::ForOp forOp, + SmallVector &mergedRegions) { + + Block &body = forOp.getRegion().front(); + + // block 内 op 顺序 + DenseMap opIndex; + int idx = 0; + for (Operation &op : body) + opIndex[&op] = idx++; + + // seed region map + DenseMap seedRegionMap; + for (int r = 0; r < mergedRegions.size(); r++) { + for (Operation *op : mergedRegions[r].opsToMove) { + seedRegionMap[op] = r; + } + } + + // 最终 op -> region + DenseMap opToRegion = seedRegionMap; + + // ---------- Step1 顺序扫描 ---------- + for (Operation &op : body) { + + if (isa(&op)) + continue; + + if (opToRegion.count(&op)) + continue; + + int region = findEarliestRegion(&op, seedRegionMap, body); + + if (region != -1) + opToRegion[&op] = region; + } + + // ---------- Step2 region-op 依赖补全 ---------- + for (Operation &op : body) { + + auto it = opToRegion.find(&op); + if (it == opToRegion.end()) + continue; + + if (op.getNumRegions() == 0) + continue; + + pullInRegionDependencies(&op, it->second, opToRegion, body); + } + + // ---------- Step3 append op ---------- + SmallPtrSet seen; + + for (Operation &op : body) { + + auto it = opToRegion.find(&op); + if (it == opToRegion.end()) + continue; + + if (!seen.insert(&op).second) + continue; + + int region = it->second; + mergedRegions[region].opsToMove.push_back(&op); + } + + // ---------- Step4 排序 ---------- + for (auto &mr : mergedRegions) { + + llvm::sort(mr.opsToMove, [&](Operation *a, Operation *b) { + return opIndex[a] < opIndex[b]; + }); + } +} + +static void +ExpandMergedRegionOpsByOutput(scf::ForOp forOp, + SmallVector &mergedRegions) { + + Block &body = forOp.getRegion().front(); + + // block 顺序（保持 IR 顺序） + DenseMap opOrder; + int idx = 0; + for (Operation &op : body) + opOrder[&op] = idx++; + + for (auto &merged : mergedRegions) { + + // 收集 region 当前产生的 value + SmallPtrSet regionValues; + + for (Operation *op : merged.opsToMove) + for (Value res : op->getResults()) + regionValues.insert(res); + + bool changed = true; + + while (changed) { + changed = false; + + for (Operation &op : body) { + + if (isa(op) || isa(op)) + continue; + + if (llvm::is_contained(merged.opsToMove, &op)) + continue; + + bool depends = false; + + for (Value operand : op.getOperands()) { + if (regionValues.contains(operand)) { + depends = true; + break; + } + } + + if (!depends) + continue; + + // 加入 region + merged.opsToMove.push_back(&op); + + // 更新 regionValues + for (Value res : op.getResults()) + regionValues.insert(res); + + changed = true; + } + } + + // 排序保持原 block 顺序 + llvm::sort(merged.opsToMove, [&](Operation *a, Operation *b) { + return opOrder[a] < opOrder[b]; + }); + } +} + +static void MoveIndependentOpsIntoIf(scf::ForOp forOp, + SmallVector &mergedRegions) { + + Block &body = forOp.getRegion().front(); + + // 记录哪些 op 已经在 region 里 + SmallPtrSet alreadyAssigned; + + for (auto &mr : mergedRegions) + for (Operation *op : mr.opsToMove) + alreadyAssigned.insert(op); + + // 记录 iter_arg -> region + DenseMap iterArgToRegion; + + for (int r = 0; r < mergedRegions.size(); r++) { + for (Operation *op : mergedRegions[r].opsToMove) { + + for (Value operand : op->getOperands()) { + + if (auto barg = mlir::dyn_cast(operand)) { + + if (barg.getOwner() == &body) + iterArgToRegion[barg] = r; + } + } + } + } + + // block 顺序 + DenseMap opIndex; + int idx = 0; + for (Operation &op : body) + opIndex[&op] = idx++; + + // 扫描所有 op + for (Operation &op : body) { + + if (isa(op) || isa(op)) + continue; + + if (alreadyAssigned.contains(&op)) + continue; + + int targetRegion = -1; + + // 看 operand 是否来自 iter_arg + for (Value operand : op.getOperands()) { + + if (auto barg = mlir::dyn_cast(operand)) { + + if (barg.getOwner() != &body) + continue; + + auto it = iterArgToRegion.find(barg); + if (it != iterArgToRegion.end()) { + + targetRegion = it->second; + break; + } + } + } + + if (targetRegion == -1) + continue; + + mergedRegions[targetRegion].opsToMove.push_back(&op); + alreadyAssigned.insert(&op); + } + + // 排序保持 block 顺序 + for (auto &mr : mergedRegions) { + + llvm::sort(mr.opsToMove, [&](Operation *a, Operation *b) { + return opIndex[a] < opIndex[b]; + }); + } +} + +// 暴力包裹 +static void +ExpandMergedRegionOpsGreedyMaximum(scf::ForOp forOp, + SmallVector &mergedRegions) { + + Block &body = forOp.getRegion().front(); + + // 记录哪些 op 已经属于 region + DenseSet regionOps; + + for (auto ®ion : mergedRegions) + for (Operation *op : region.opsToMove) + regionOps.insert(op); + + // block op 列表 + SmallVector ops; + for (Operation &op : body) + ops.push_back(&op); + + DenseMap opIndex; + for (int i = 0; i < ops.size(); i++) + opIndex[ops[i]] = i; + + for (auto ®ion : mergedRegions) { + + if (region.opsToMove.empty()) + continue; + + // 找到 region 在 block 中的范围 + int start = ops.size(); + int end = -1; + + for (Operation *op : region.opsToMove) { + int idx = opIndex[op]; + start = std::min(start, idx); + end = std::max(end, idx); + } + + SmallVector newOps; + + // ---------- backward 扩展 ---------- + for (int i = start - 1; i >= 0; i--) { + Operation *op = ops[i]; + + if (isa(op)) + break; + + if (regionOps.contains(op)) + break; + + newOps.push_back(op); + } + + // ---------- forward 扩展 ---------- + for (int i = end + 1; i < ops.size(); i++) { + Operation *op = ops[i]; + + if (isa(op)) + break; + + if (regionOps.contains(op)) + break; + + newOps.push_back(op); + } + + // 加入 region + for (Operation *op : newOps) { + region.opsToMove.push_back(op); + regionOps.insert(op); + } + } + + // 最后保持 block 顺序 + for (auto ®ion : mergedRegions) { + + llvm::sort(region.opsToMove, [](Operation *a, Operation *b) { + return a->isBeforeInBlock(b); + }); + + region.opsToMove.erase( + std::unique(region.opsToMove.begin(), region.opsToMove.end()), + region.opsToMove.end()); + } +} + +static void CollectForYieldRelatedOps(scf::ForOp forOp, + SmallVector &mergedRegions, + DenseSet &yieldRelatedOps) { + + Block &body = forOp.getRegion().front(); + + // 已经属于 region 的 op + DenseSet regionOps; + for (auto ®ion : mergedRegions) + for (Operation *op : region.opsToMove) + regionOps.insert(op); + + auto yield = cast(body.getTerminator()); + + SmallVector worklist; + DenseSet visited; + + // 初始化 worklist + for (Value v : yield.getOperands()) + worklist.push_back(v); + + while (!worklist.empty()) { + Value v = worklist.pop_back_val(); + + if (!visited.insert(v).second) + continue; + + Operation *def = v.getDefiningOp(); + if (!def) + continue; + + // 只处理 for body 内的 op + if (def->getBlock() != &body) + continue; + + // 已经在 region 内 + if (regionOps.contains(def)) + continue; + + // 记录 + if (yieldRelatedOps.insert(def).second) { + + // 继续向上找依赖 + for (Value operand : def->getOperands()) + worklist.push_back(operand); + } + } +} + +// 贪心吸收region前后的op +static void +ExpandMergedRegionOpsGreedy(scf::ForOp forOp, + SmallVector &mergedRegions, + DenseSet &skipOps) { + + Block &body = forOp.getRegion().front(); + + // 记录哪些 op 已经属于 region + DenseSet regionOps; + for (auto ®ion : mergedRegions) + for (Operation *op : region.opsToMove) + regionOps.insert(op); + + // block op 列表 + SmallVector ops; + for (Operation &op : body) + ops.push_back(&op); + + // op -> index + DenseMap opIndex; + for (int i = 0; i < ops.size(); i++) + opIndex[ops[i]] = i; + + for (auto ®ion : mergedRegions) { + + if (region.opsToMove.empty()) + continue; + + // 找到 region 在 block 中的范围 + int start = ops.size(); + int end = -1; + + for (Operation *op : region.opsToMove) { + int idx = opIndex[op]; + start = std::min(start, idx); + end = std::max(end, idx); + } + + SmallVector newOps; + + // ---------- backward 扩展 ---------- + for (int i = start - 1; i >= 0; i--) { + Operation *op = ops[i]; + + // block terminator + if (isa(op)) + break; + + // 遇到其他 region 的 op + if (regionOps.contains(op)) + break; + + // yield 关联 op，跳过但继续扫描 + if (skipOps.contains(op)) + continue; + + newOps.push_back(op); + } + + // ---------- forward 扩展 ---------- + for (int i = end + 1; i < ops.size(); i++) { + Operation *op = ops[i]; + + // block terminator + if (isa(op)) + break; + + // 遇到其他 region 的 op + if (regionOps.contains(op)) + break; + + // yield 关联 op，跳过 + if (skipOps.contains(op)) + continue; + + newOps.push_back(op); + } + + // 加入 region + for (Operation *op : newOps) { + region.opsToMove.push_back(op); + regionOps.insert(op); + } + } + + // 最后保持 block 顺序 + for (auto ®ion : mergedRegions) { + + llvm::sort(region.opsToMove, [](Operation *a, Operation *b) { + return a->isBeforeInBlock(b); + }); + + region.opsToMove.erase( + std::unique(region.opsToMove.begin(), region.opsToMove.end()), + region.opsToMove.end()); + } +} + +// 贪心吸收region前面的op +static void +ExpandMergedRegionOpsGreedyV2(scf::ForOp forOp, + SmallVector &mergedRegions, + DenseSet &skipOps) { + + Block &body = forOp.getRegion().front(); + + // block op 列表 + SmallVector ops; + for (Operation &op : body) + ops.push_back(&op); + + // op -> index + DenseMap opIndex; + for (int i = 0; i < ops.size(); i++) + opIndex[ops[i]] = i; + + // 记录哪些 op 已经属于 region + DenseSet regionOps; + for (auto ®ion : mergedRegions) + for (Operation *op : region.opsToMove) + regionOps.insert(op); + + for (int r = 0; r < mergedRegions.size(); r++) { + + auto ®ion = mergedRegions[r]; + if (region.opsToMove.empty()) + continue; + + // ---------- 当前 region block 范围 ---------- + int start = ops.size(); + int end = -1; + + for (Operation *op : region.opsToMove) { + int idx = opIndex[op]; + start = std::min(start, idx); + end = std::max(end, idx); + } + + // ---------- 前一个 region 的末尾 ---------- + int prevEnd = -1; + + if (r > 0 && !mergedRegions[r - 1].opsToMove.empty()) { + for (Operation *op : mergedRegions[r - 1].opsToMove) { + prevEnd = std::max(prevEnd, opIndex[op]); + } + } + + SmallVector newOps; + + // ---------- backward expand ---------- + for (int i = start - 1; i > prevEnd; i--) { + + Operation *op = ops[i]; + + // terminator + if (isa(op)) + break; + + // 已属于 region + if (regionOps.contains(op)) + break; + + // yield chain op + if (skipOps.contains(op)) + continue; + + newOps.push_back(op); + } + + // ---------- 加入 region ---------- + for (Operation *op : newOps) { + region.opsToMove.push_back(op); + regionOps.insert(op); + } + } + + // ---------- 保持 block 顺序 ---------- + for (auto ®ion : mergedRegions) { + + llvm::sort(region.opsToMove, [](Operation *a, Operation *b) { + return a->isBeforeInBlock(b); + }); + + region.opsToMove.erase( + std::unique(region.opsToMove.begin(), region.opsToMove.end()), + region.opsToMove.end()); + } +} + +// 贪心吸收region前面的op +static void +ExpandMergedRegionOpsGreedyV2ForAIC(scf::ForOp forOp, + SmallVector &mergedRegions) { + + Block &body = forOp.getRegion().front(); + + // block op 列表 + SmallVector ops; + for (Operation &op : body) + ops.push_back(&op); + + // op -> index + DenseMap opIndex; + for (int i = 0; i < ops.size(); i++) + opIndex[ops[i]] = i; + + // 记录哪些 op 已经属于 region + DenseSet regionOps; + for (auto ®ion : mergedRegions) + for (Operation *op : region.opsToMove) + regionOps.insert(op); + + for (int r = 0; r < mergedRegions.size(); r++) { + + auto ®ion = mergedRegions[r]; + if (region.opsToMove.empty()) + continue; + + // ---------- 当前 region block 范围 ---------- + int start = ops.size(); + int end = -1; + + for (Operation *op : region.opsToMove) { + int idx = opIndex[op]; + start = std::min(start, idx); + end = std::max(end, idx); + } + + // ---------- 前一个 region 的末尾 ---------- + int prevEnd = -1; + + if (r > 0 && !mergedRegions[r - 1].opsToMove.empty()) { + for (Operation *op : mergedRegions[r - 1].opsToMove) { + prevEnd = std::max(prevEnd, opIndex[op]); + } + } + + SmallVector newOps; + + // ---------- backward expand ---------- + for (int i = start - 1; i > prevEnd; i--) { + + Operation *op = ops[i]; + + // terminator + if (isa(op)) + break; + + // 已属于 region + if (regionOps.contains(op)) + break; + + newOps.push_back(op); + } + + // ---------- 加入 region ---------- + for (Operation *op : newOps) { + region.opsToMove.push_back(op); + regionOps.insert(op); + } + } + + // ---------- 保持 block 顺序 ---------- + for (auto ®ion : mergedRegions) { + + llvm::sort(region.opsToMove, [](Operation *a, Operation *b) { + return a->isBeforeInBlock(b); + }); + + region.opsToMove.erase( + std::unique(region.opsToMove.begin(), region.opsToMove.end()), + region.opsToMove.end()); + } +} + +static void MoveForYieldOpIntoRegion(scf::ForOp forOp, + DenseSet &yieldRelatedOps, + SmallVector &mergedRegions) { + + DenseMap opToRegion; + + for (int i = 0; i < mergedRegions.size(); i++) + for (Operation *op : mergedRegions[i].opsToMove) + opToRegion[op] = i; + + auto yield = cast(forOp.getBody()->getTerminator()); + + for (int i = 0; i < yield.getNumOperands(); i++) { + + Value iterArg = forOp.getRegionIterArgs()[i]; + Value yieldVal = yield.getOperand(i); + + Operation *def = yieldVal.getDefiningOp(); + if (!def) + continue; + + if (!yieldRelatedOps.contains(def)) + continue; + + int targetRegion = -1; + + for (Operation *user : iterArg.getUsers()) { + + if (opToRegion.count(user)) { + targetRegion = opToRegion[user]; + break; + } + } + + if (targetRegion == -1) + continue; + + SmallVector stack; + stack.push_back(def); + + while (!stack.empty()) { + Operation *op = stack.pop_back_val(); + + if (!yieldRelatedOps.contains(op)) + continue; + + mergedRegions[targetRegion].opsToMove.push_back(op); + + yieldRelatedOps.erase(op); + + for (Value operand : op->getOperands()) { + if (Operation *dep = operand.getDefiningOp()) + stack.push_back(dep); + } + } + } + + for (auto ®ion : mergedRegions) { + + llvm::sort(region.opsToMove, [](Operation *a, Operation *b) { + return a->isBeforeInBlock(b); + }); + + region.opsToMove.erase( + std::unique(region.opsToMove.begin(), region.opsToMove.end()), + region.opsToMove.end()); + } +} + +static void +MoveRemainingYieldOpsToPrevRegion(scf::ForOp forOp, + DenseSet &yieldRelatedOps, + SmallVector &mergedRegions) { + + if (yieldRelatedOps.empty()) + return; + + Block &body = forOp.getRegion().front(); + + // op -> region index + DenseMap opToRegion; + for (int i = 0; i < mergedRegions.size(); i++) + for (Operation *op : mergedRegions[i].opsToMove) + opToRegion[op] = i; + + // block 顺序 + SmallVector ops; + for (Operation &op : body) + ops.push_back(&op); + + DenseMap opIndex; + for (int i = 0; i < ops.size(); i++) + opIndex[ops[i]] = i; + + for (Operation *op : yieldRelatedOps) { + + if (op->getBlock() != &body) + continue; + + int idx = opIndex[op]; + + int targetRegion = -1; + + // 向前找最近的 region + for (int i = idx - 1; i >= 0; i--) { + Operation *prev = ops[i]; + + if (opToRegion.count(prev)) { + targetRegion = opToRegion[prev]; + break; + } + } + + if (targetRegion == -1) + continue; + + mergedRegions[targetRegion].opsToMove.push_back(op); + } + + // 排序 + 去重 + for (auto ®ion : mergedRegions) { + + llvm::sort(region.opsToMove, [](Operation *a, Operation *b) { + return a->isBeforeInBlock(b); + }); + + region.opsToMove.erase( + std::unique(region.opsToMove.begin(), region.opsToMove.end()), + region.opsToMove.end()); + } +} + +static void MoveIndependentOpsIntoRegionBackwardV2( + scf::ForOp forOp, SmallVector &mergedRegions) { + + Block &body = forOp.getRegion().front(); + SmallVector ops; + for (Operation &op : body) + ops.push_back(&op); + + DenseMap opToRegion; + for (int i = 0; i < mergedRegions.size(); i++) + for (Operation *op : mergedRegions[i].opsToMove) + opToRegion[op] = i; + + // ----------- 收集移动计划 ----------- + DenseMap movePlan; + + for (int i = 0; i < mergedRegions.size(); i++) { + MergedRegion ®ion = mergedRegions[i]; + if (region.opsToMove.empty()) + continue; + + Operation *firstOp = region.opsToMove.front(); + Operation *lastOp = region.opsToMove.back(); + auto itFirst = std::find(ops.begin(), ops.end(), firstOp); + auto itLast = std::find(ops.begin(), ops.end(), lastOp); + if (itFirst == ops.end() || itLast == ops.end()) + continue; + + int startIdx = std::distance(ops.begin(), itFirst); + int endIdx = std::distance(ops.begin(), itLast); + + // ----------- 收集 wait-set 区间 ----------- + SmallVector> waitIntervals; + bool inWait = false; + int begin = -1; + for (int j = startIdx; j <= endIdx; j++) { + Operation *op = ops[j]; + if (op->getName().getStringRef().contains("sync_block_wait")) { + inWait = true; + begin = j + 1; + continue; + } + if (op->getName().getStringRef().contains("sync_block_set") && inWait) { + inWait = false; + waitIntervals.push_back({begin, j - 1}); + } + } + auto isInWaitSet = [&](int idx) { + for (auto &p : waitIntervals) + if (idx >= p.first && idx <= p.second) + return true; + return false; + }; + + // ----------- 从后往前扫描 region 内的 op ----------- + for (int j = endIdx; j >= startIdx; j--) { + Operation *op = ops[j]; + if (isa(op) || isInWaitSet(j)) + continue; + + // ---------- operand 是否依赖本 region ---------- + bool dependCurrentRegion = false; + for (Value operand : op->getOperands()) { + Operation *def = operand.getDefiningOp(); + if (!def) + continue; + if (std::find(region.opsToMove.begin(), region.opsToMove.end(), def) != + region.opsToMove.end()) { + dependCurrentRegion = true; + break; + } + } + if (dependCurrentRegion) + continue; + + // ---------- 当前 region 后续是否使用 ---------- + bool usedLaterInSameRegion = false; + for (Value result : op->getResults()) + for (Operation *user : result.getUsers()) + if (std::find(region.opsToMove.begin(), region.opsToMove.end(), + user) != region.opsToMove.end() && + std::find(region.opsToMove.begin(), region.opsToMove.end(), op) < + std::find(region.opsToMove.begin(), region.opsToMove.end(), + user)) { + usedLaterInSameRegion = true; + break; + } + if (usedLaterInSameRegion) + continue; + + // ---------- 找使用该 op 的后续 region ---------- + int targetRegion = -1; + for (int k = i + 1; k < mergedRegions.size(); ++k) { + for (Operation *candidate : mergedRegions[k].opsToMove) + for (Value operand : candidate->getOperands()) + if (operand.getDefiningOp() == op) { + targetRegion = k; + break; + } + if (targetRegion != -1) + break; + if (targetRegion != -1) + break; + } + if (targetRegion == -1) + continue; + + movePlan[op] = targetRegion; + // llvm::outs() << "MJ: plan move " << *op + // << " -> region " << targetRegion << "\n"; + } + } + + // ----------- 统一应用移动 ----------- + for (auto &it : movePlan) { + Operation *op = it.first; + int targetRegionIdx = it.second; + MergedRegion &targetRegion = mergedRegions[targetRegionIdx]; + // 更新数据结构 + targetRegion.opsToMove.push_back(op); + + llvm::outs() << "MJ: move " << *op << " -> region " << targetRegionIdx + << "\n"; + } + + // ----------- 更新原 region 的 opsToMove ----------- + for (int i = 0; i < mergedRegions.size(); ++i) { + MergedRegion ®ion = mergedRegions[i]; + SmallVector newOps; + for (Operation *op : region.opsToMove) { + auto it = movePlan.find(op); + if (it == movePlan.end() || it->second == i) { + // 没有移动计划，或者移动的目标就是自己，保留 + newOps.push_back(op); + } + } + region.opsToMove.swap(newOps); + } + + // ----------- 排序 + 去重 ----------- + for (auto ®ion : mergedRegions) { + llvm::sort(region.opsToMove, [](Operation *a, Operation *b) { + return a->isBeforeInBlock(b); + }); + region.opsToMove.erase( + std::unique(region.opsToMove.begin(), region.opsToMove.end()), + region.opsToMove.end()); + } +} + +// // debug: 如果一个forop的第一个region的最后3条op是%27 = tt.expand_dims %25#1 +// {axis = 1 : i32} : tensor<64xf32> -> tensor<64x1xf32> +// %28 = tt.broadcast %27 : tensor<64x1xf32> -> tensor<64x128xf32> +// %29 = arith.mulf %arg10, %28 : tensor<64x128xf32> +// 直接放到第2个region里 +static void TempChange(scf::ForOp forOp, + SmallVector &mergedRegions) { + + if (mergedRegions.size() < 2) + return; + + auto &srcRegion = mergedRegions[0]; + auto &dstRegion = mergedRegions[1]; + + if (srcRegion.opsToMove.size() < 3) + return; + + Operation *op1 = srcRegion.opsToMove[srcRegion.opsToMove.size() - 3]; + Operation *op2 = srcRegion.opsToMove[srcRegion.opsToMove.size() - 2]; + Operation *op3 = srcRegion.opsToMove[srcRegion.opsToMove.size() - 1]; + + // ---------- pattern 匹配 ---------- + if (!op1->getName().getStringRef().contains("tt.expand_dims")) + return; + + if (!op2->getName().getStringRef().contains("tt.broadcast")) + return; + + if (!op3->getName().getStringRef().contains("arith.mulf")) + return; + + llvm::outs() << "TempChange triggered\n"; + + SmallVector opsToMove = {op1, op2, op3}; + + // ---------- 移动到 region2 末尾 ---------- + for (Operation *op : opsToMove) { + dstRegion.opsToMove.push_back(op); + llvm::outs() << "TempChange move: " << *op << "\n"; + } + + // ---------- 从 region1 删除 ---------- + srcRegion.opsToMove.resize(srcRegion.opsToMove.size() - 3); + + // ---------- 排序 ---------- + for (auto ®ion : mergedRegions) { + llvm::sort(region.opsToMove, [](Operation *a, Operation *b) { + return a->isBeforeInBlock(b); + }); + + region.opsToMove.erase( + std::unique(region.opsToMove.begin(), region.opsToMove.end()), + region.opsToMove.end()); + } +} + +static void sortOperationsByDataFlow(llvm::SmallVector &ops) { + llvm::DenseSet visited; + llvm::SmallVector result; + + std::function dfs = [&](Operation *op) { + if (!visited.insert(op).second) + return; + + for (Value operand : op->getOperands()) { + if (Operation *def = operand.getDefiningOp()) { + if (llvm::is_contained(ops, def)) + dfs(def); + } + } + + result.push_back(op); + }; + + for (Operation *op : ops) + dfs(op); + + ops.assign(result.begin(), result.end()); +} + +static void rewriteOperandsRecursively(Operation *op, + DenseMap &valueMap) { + + // 1 rewrite 当前 op 的 operands + for (OpOperand &operand : op->getOpOperands()) { + Value v = operand.get(); + auto it = valueMap.find(v); + if (it != valueMap.end()) + operand.set(it->second); + } + + // 2 递归进入 region + for (Region ®ion : op->getRegions()) { + for (Block &block : region) { + for (Operation &nestedOp : block) { + rewriteOperandsRecursively(&nestedOp, valueMap); + } + } + } +} + +static void CopyOpsToAfterwardRegions( + SmallVector &mergedRegions, + DenseMap &yieldMap, + DenseMap &cloneAndOriYieldMap, + SmallVector &copiedForOps) { + + if (mergedRegions.size() <= 1) + return; + + // 先整理一个 set，方便判断哪些 op 是 yield defining op + DenseSet yieldDefOps; + for (auto &it : yieldMap) + yieldDefOps.insert(it.second); + + // 倒序遍历 region + for (int i = mergedRegions.size() - 1; i >= 0; --i) { + MergedRegion &curRegion = mergedRegions[i]; + + DenseMap valueMap; + SmallVector clonedOps; + + // 遍历前面的 region + for (int k = 0; k < i; ++k) { + MergedRegion &prevRegion = mergedRegions[k]; + + int waitSetLevel = 0; + + for (Operation *op : prevRegion.opsToMove) { + + if (isa(op)) { + waitSetLevel++; + continue; + } + + if (isa(op)) { + waitSetLevel = std::max(waitSetLevel - 1, 0); + continue; + } + + if (waitSetLevel > 0) + continue; + + if (isa(op)) + continue; + + IRMapping mapper; + + for (auto result : op->getResults()) + if (valueMap.count(result)) + mapper.map(result, valueMap[result]); + + Operation *insertPoint = + curRegion.opsToMove.empty() ? nullptr : curRegion.opsToMove.front(); + + OpBuilder builder(insertPoint ? insertPoint : op); + + Operation *cloned = builder.clone(*op, mapper); + + // 记录 result mapping + for (auto it : llvm::zip(op->getResults(), cloned->getResults())) + valueMap[std::get<0>(it)] = std::get<1>(it); + + // 如果这个 op 是 yield defining op，记录 clone -> original + if (yieldDefOps.contains(op)) { + cloneAndOriYieldMap[cloned] = op; + } + + // 记录copy的for op + if (auto forOp = dyn_cast(cloned)) { + copiedForOps.push_back(forOp); + } + + clonedOps.push_back(cloned); + } + } + + // 插入到当前 region 开头 + curRegion.opsToMove.insert(curRegion.opsToMove.begin(), clonedOps.begin(), + clonedOps.end()); + + // rebuild SSA + for (Operation *op : curRegion.opsToMove) { + rewriteOperandsRecursively(op, valueMap); + } + + // 排序保证拓扑顺序 + sortOperationsByDataFlow(curRegion.opsToMove); + } +} + +/// 记录 forOp 的 yield value 与其原始生成的 op 的映射 +static void GetYieldMap(scf::ForOp forOp, + DenseMap &yieldMap) { + yieldMap.clear(); + + // 取 forOp body 的 scf.yield + auto yieldOp = dyn_cast(forOp.getBody()->getTerminator()); + if (!yieldOp) + return; + + for (Value yieldVal : yieldOp.getOperands()) { + // 获取生成 yieldVal 的原始 op + Operation *defOp = yieldVal.getDefiningOp(); + + // 对 block arg（可能是 iter_arg）没有 definingOp 的情况，可以跳过或直接记录 + // nullptr + if (!defOp) + continue; + + yieldMap[yieldVal] = defOp; + } +} + +static Value findIterArgForAIC(Value v, scf::ForOp forOp) { + while (true) { + if (auto arg = dyn_cast(v)) { + if (arg.getOwner() == forOp.getBody()) + return v; + return Value(); + } + + Operation *def = v.getDefiningOp(); + if (!def) + return Value(); + + if (def->getNumOperands() == 0) + return Value(); + + v = def->getOperand(0); + } +} + +static Operation * +findCloneOfYieldOp(Operation *oriYieldOp, + DenseMap &cloneAndOriYieldMap, + MergedRegion ®ion) { + + for (Operation *op : region.opsToMove) { + auto it = cloneAndOriYieldMap.find(op); + if (it != cloneAndOriYieldMap.end() && it->second == oriYieldOp) + return op; + } + return nullptr; +} + +static void RebuildForYielValuesForAIC( + scf::ForOp forOp, SmallVector &mergedRegions, + DenseMap &yieldMap, + DenseMap &cloneAndOriYieldMap) { + + auto yieldOp = cast(forOp.getBody()->getTerminator()); + + for (MergedRegion ®ion : mergedRegions) { + + triton::DotOp dotOp = nullptr; + + for (Operation *op : region.opsToMove) { + if (auto d = dyn_cast(op)) { + dotOp = d; + break; + } + } + + if (!dotOp) + continue; + + // 处理 dot operand + for (Value operand : dotOp->getOperands()) { + + Value iterArg = findIterArgForAIC(operand, forOp); + if (!iterArg) + continue; + + auto arg = cast(iterArg); + int idx = arg.getArgNumber(); + + if (idx >= yieldOp.getNumOperands()) + continue; + + Value oriYieldValue = yieldOp.getOperand(idx); + + auto it = yieldMap.find(oriYieldValue); + if (it == yieldMap.end()) + continue; + + Operation *oriYieldOp = it->second; + + Operation *cloneOp = + findCloneOfYieldOp(oriYieldOp, cloneAndOriYieldMap, region); + + if (!cloneOp) + continue; + + yieldOp.setOperand(idx, cloneOp->getResult(0)); + } + } +} + +void ExpandMergedRegionOps(scf::ForOp forOp, + SmallVector &mergedRegions, + SmallVector &copiedForOps) { + bool isInAIV = false; + auto scopeOp = forOp->getParentOfType(); + if (!scopeOp) + return; + + auto coreTypeAttr = + scopeOp->getAttrOfType(hivm::TCoreTypeAttr::name); + + if (coreTypeAttr.getTcoretype() == hivm::TCoreType::VECTOR) { + isInAIV = true; + } + + if (isInAIV) { + DenseSet yieldRelatedOps; + + // 1 收集 yield 相关 op + CollectForYieldRelatedOps(forOp, mergedRegions, yieldRelatedOps); + + // 2 greedy 扩展 + // ExpandMergedRegionOpsGreedy(forOp, mergedRegions, yieldRelatedOps); + ExpandMergedRegionOpsGreedyV2(forOp, mergedRegions, yieldRelatedOps); + + // 3 与前面wait-set region独立的op应该被放入后面的关联的region + MoveIndependentOpsIntoRegionBackwardV2(forOp, mergedRegions); + + // 4 根据 iter_arg 使用位置放入 region + MoveForYieldOpIntoRegion(forOp, yieldRelatedOps, mergedRegions); + + // 5 剩余 yield chain 放入前一个 region + MoveRemainingYieldOpsToPrevRegion(forOp, yieldRelatedOps, mergedRegions); + } else { // AIC单独处理, 避免出现CUBE内的tensor变量依赖 + // 用Map记录原始的for yield op的的映射 + DenseMap yieldMap; + GetYieldMap(forOp, yieldMap); + + llvm::outs() << "YieldMap:\n"; + for (auto it : yieldMap) { + llvm::outs() << *(it.second) << "\n"; + } + + // 2 greedy 扩展, yield value后续处理 + ExpandMergedRegionOpsGreedyV2ForAIC(forOp, mergedRegions); + + // 复制当前region的除tt.dot、以及[wait - + // set]之间的op到后续的所有MergedRegion 倒序实现 + // 记录clone和original的yield对应op的map + DenseMap cloneAndOriYieldMap; + CopyOpsToAfterwardRegions(mergedRegions, yieldMap, cloneAndOriYieldMap, + copiedForOps); + + // 4 + // 先确定每个MergedRegion的tt.dot的operand的来源是for的哪个iter_arg(递归查找), + // 假设为%arg0, 依据yieldMap可以得到oriYield 遍历当前MergedRegion的所有op, + // 确定哪条op对应的cloneAndOriYieldMap的second是oriYield, 假设为%45 + // 最后替换for yield op对应位置的operand为%45 + RebuildForYielValuesForAIC(forOp, mergedRegions, yieldMap, + cloneAndOriYieldMap); + } +} + +void MergeWaitSetRegions(SmallVector ®ions, + SmallVector &merged) { + for (int i = 0; i < regions.size();) { + MergedRegion mr; + mr.regions.push_back(®ions[i]); + mr.opsToMove.append(regions[i].opsToMove); + + int j = i; + while (!regions[j].hasCopyOrFixpipe && j + 1 < regions.size()) { + j++; + mr.regions.push_back(®ions[j]); + mr.opsToMove.append(regions[j].opsToMove); + } + + merged.push_back(std::move(mr)); + i = j + 1; + } + + for (MergedRegion &mr : merged) { + SmallPtrSet regionValues; + SmallPtrSet opSet; + + for (Operation *op : mr.opsToMove) + opSet.insert(op); + + for (Operation *op : mr.opsToMove) { + for (Value v : op->getResults()) { + bool usedOutside = false; + for (OpOperand &use : v.getUses()) { + Operation *user = use.getOwner(); + if (!opSet.contains(user) && user->getBlock() == op->getBlock()) { + usedOutside = true; + break; + } + } + if (usedOutside) { + mr.yieldValues.push_back(v); + mr.resultTypes.push_back(v.getType()); + } + } + } + } +} + +void GetBlockInfos(SmallVector ®ions, Block &body) { + for (auto it = body.begin(); it != body.end();) { + Operation *op = &*it; + + auto waitOp = dyn_cast(op); + if (!waitOp) { + it++; + continue; + } + + auto pipeS = hivm::PipeAttr::get(op->getContext(), hivm::PIPE::PIPE_S); + if (auto syncWait = dyn_cast(op)) { + if ((syncWait.getTpipe() == pipeS || syncWait.getPipe() == pipeS) && + (std::next(it) != body.end())) { + if (isa(&*std::next(it)) || + isa(&*std::next(it))) + return; + } + } + + Operation *lastSetOp = nullptr; + + // 扫描到下一个 wait, 收集所有 set + auto curIt = std::next(it); + auto endIt = curIt; + int setOpCount = 0; + SmallVector opsInRegion; + for (; curIt != body.end(); ++curIt) { + Operation *curOp = &*curIt; + if (isa(curOp) && setOpCount >= 1) + break; + if (isa(curOp)) { + setOpCount++; + endIt = curIt; // setop的位置 + lastSetOp = curOp; // 最后一个 set + } + } + + if (!lastSetOp) { + it = curIt; + continue; + } // 没有 set, 不包 + + // 收集 [wait, ..., lastSet] 之间的 ops + bool hasCopyOrFixpipe = false; + for (auto it2 = it; it2 != std::next(endIt); ++it2) { + Operation *curOp = &*it2; + opsInRegion.push_back(curOp); + if (isa(curOp) || isa(curOp)) { + hasCopyOrFixpipe = true; + } + } + + it = endIt++; + regions.push_back({waitOp, lastSetOp, opsInRegion, hasCopyOrFixpipe}); + } +} + +Value findIterArg(Value v, Type t) { + SmallVector worklist = {v}; + SmallPtrSet visited; + + while (!worklist.empty()) { + Value cur = worklist.front(); + worklist.erase(worklist.begin()); + if (!visited.insert(cur).second) + continue; + + // 匹配scf.for原始迭代参数, 直接返回 + if (auto b = mlir::dyn_cast(cur)) { + auto forOp = mlir::dyn_cast(b.getOwner()->getParentOp()); + if (forOp && b.getType() == t) { + for (Value iterArg : forOp.getRegionIterArgs()) { + if (iterArg.getAsOpaquePointer() == b.getAsOpaquePointer()) { + return b; + } + } + } + } + + Operation *defOp = cur.getDefiningOp(); + if (!defOp) + continue; + + // 核心逻辑：如果当前值是scf.if的结果 + // 进入then块找源头 + if (auto ifOp = mlir::dyn_cast(defOp)) { + Block &thenBlock = ifOp.getThenRegion().front(); + // 找到then块最后一个op（scf.yield） + // 取其operands（即ifOp结果的源头值） + for (auto &innerOp : llvm::reverse(thenBlock)) { + if (auto yieldOp = mlir::dyn_cast(&innerOp)) { + // 按索引匹配: cur是ifOp的第n个结果, 取yieldOp的第n个operand + for (auto [idx, res] : llvm::enumerate(ifOp.getResults())) { + if (res.getAsOpaquePointer() == cur.getAsOpaquePointer()) { + Value srcVal = yieldOp.getOperand(idx); + if (!visited.count(srcVal)) + worklist.push_back(srcVal); + break; + } + } + break; // 找到yield即退出, 无需遍历其他op + } + } + } else { + // 非if结果值 + // 正常往前追溯operands + for (Value operand : defOp->getOperands()) { + if (!visited.count(operand)) + worklist.push_back(operand); + } + } + } + + llvm::outs() << "未找到迭代参数, 返回原值: "; + v.print(llvm::outs()); + llvm::outs() << "\n"; + return v; +} + +// 如果 v 最终被 scf.for 的 yield 使用 +// → 返回对应的 forOp 的 iter_arg +// 如果 v 只是流向后面的 wait-set region / 其他 op +// → 直接返回原值 v +Value findIterArgForAll(Value v, Type t) { + for (Operation *user : v.getUsers()) { + + if (auto yieldOp = dyn_cast(user)) { + + if (auto forOp = dyn_cast(yieldOp->getParentOp())) { + + for (auto [idx, operand] : llvm::enumerate(yieldOp.getOperands())) { + + if (operand.getAsOpaquePointer() == v.getAsOpaquePointer()) { + + Value iterArg = forOp.getRegionIterArgs()[idx]; + + if (iterArg.getType() == t) + return iterArg; + } + } + } + } + } + + return v; +} + +void FindDependValues(SmallVector &dependValues, + SmallVector mergedRegions) { + dependValues.clear(); + for (auto &curMR : mergedRegions) { + for (Value yieldValue : curMR.yieldValues) { + // llvm::outs() << "yieldValue: "<< yieldValue << "\n"; + // 遍历当前区域的yieldValue的所有user OP，判断是否存在依赖关系 + for (OpOperand &use : yieldValue.getUses()) { + Operation *userOp = use.getOwner(); + + // llvm::outs() << "userOp: "<< *userOp << "\n"; + bool isUserInOtherRegion = false; + for (auto &otherMR : mergedRegions) { + // 跳过当前区域，只检查yieldValue是否被其他区域使用 + if (&otherMR == &curMR) + continue; + + // 只要有一个 userOp在 otherMR 的 opsToMove + // 列表中，就认为是dependValue llvm::outs() << "judge comtain\n"; for + // (size_t k = 0; k < otherMR.opsToMove.size(); k++) { + // llvm::outs() << "otherMR op: " << *(otherMR.opsToMove[k]) << + // "\n"; + // } + // llvm::outs() << "otherMR end\n"; + + // if (llvm::is_contained(otherMR.opsToMove, userOp)) { + // isUserInOtherRegion = true; + // llvm::outs() << "is_contained\n"; + // break; + // } + + // 用 DenseSet 暂存当前 region 的所有 ops + // 初始 DenseSet: 顶层 opsToMove + DenseSet otherOps; + for (Operation *op : otherMR.opsToMove) { + CollectAllNestedOps(op, otherOps); // 完整展开嵌套 + } + if (otherOps.contains(userOp)) { + isUserInOtherRegion = true; + break; + } + } + + // 无重复的添加依赖变量 + if (isUserInOtherRegion) { + if (!llvm::is_contained(dependValues, yieldValue)) { + dependValues.push_back(yieldValue); + } + break; + } + } + } + } +} + +void UpdateMergedRegionsWithNewForOp(SmallVector &mergedRegions, + IRMapping &mapper) { + for (auto &mr : mergedRegions) { + // WaitSetRegion 后续已经不使用了，直接释放，否则会出现野指针 + SmallVector newRegions; + newRegions.clear(); + mr.regions = newRegions; + // // 更新 opsToMove 列表 + // llvm::outs() << "before \n"; + // for (auto &op : mr.opsToMove) { + // llvm::outs() << "opsToMove: " << op << ", " << *op << '\n'; + // } + SmallVector newOpsToMove; + newOpsToMove.clear(); + for (Operation *op : mr.opsToMove) { + if (op) { + Operation *newOp = mapper.lookupOrNull(op); + newOpsToMove.push_back(newOp); + } + } + mr.opsToMove = newOpsToMove; + // llvm::outs() << "after \n"; + // for (auto &op : mr.opsToMove) { + // llvm::outs() << "opsToMove: " << op << ", " << *op << '\n'; + // } + // 更新 yieldValues 列表 + SmallVector newYieldValues; + newYieldValues.clear(); + for (Value v : mr.yieldValues) { + if (v) { + newYieldValues.push_back(mapper.lookupOrNull(v)); + } + } + mr.yieldValues = newYieldValues; + // resultTypes 是type 类型，无需更新 + } +} + +void AddArgsForDependValues(scf::ForOp forOp, SmallVector &dependValues, + SmallVector &mergedRegions, + ModuleOp module) { + OpBuilder moduleBuilder(module.getContext()); + SmallVector valueTypes; + valueTypes.clear(); + + if (dependValues.empty()) { + return; + } else { + for (Value v : dependValues) { + Type valueType = v.getType(); + valueTypes.push_back(valueType); + } + } + + // 为每个 dependValue 创建一个初始值（可能不存在相同shape和type的常量tensor） + SmallVector initTensors; + initTensors.clear(); + module.walk([&](Operation *op) { + if (auto constOp = dyn_cast(op)) { + moduleBuilder.setInsertionPoint(constOp); + for (Type valueType : valueTypes) { + auto tensorType = dyn_cast(valueType); + triton::PointerType ptrType; + ptrType = + (tensorType) + ? dyn_cast(tensorType.getElementType()) + : dyn_cast(valueType); + if (ptrType) { + // 如果依赖变量是一个ptr类型 + // 1. 创建 i64 0 + // 2. cast 成 !tt.ptr<...> + Value zero = moduleBuilder.create( + constOp.getLoc(), 0, 64); + Value ptrValue = moduleBuilder.create( + constOp.getLoc(), ptrType, zero); + if (tensorType) { + // 3. splat 成 tensor<...x!tt.ptr<...>> + Value ptrTensor = moduleBuilder.create( + constOp.getLoc(), tensorType, ptrValue); + initTensors.push_back(ptrTensor); + } else { + initTensors.push_back(ptrValue); + } + } else if (auto memrefType = dyn_cast(valueType)) { + // 如果中间变量是一个memref类型，为iterarg创建一个 alloc = memref + // 仅支持#hivm.address_space，对于#hivm.address_space，不存在 + // copy cbuf to cbuf 行为 + auto spaceAttr = + cast(memrefType.getMemorySpace()); + if (spaceAttr && + spaceAttr.getAddressSpace() == hivm::AddressSpace::L1) { + llvm::dbgs() << "AddArgsForDependValues: dependValue type is a " + "memref hivm::AddressSpace::L1 type!!!\n"; + return mlir::WalkResult::interrupt(); + } else { + mlir::Value alloc = moduleBuilder.create( + constOp.getLoc(), memrefType); + initTensors.push_back(alloc); + } + } else { + // 非 ptr 类型创建零值常量 + auto zeroAttr = moduleBuilder.getZeroAttr(valueType); + Value zeroTensor = moduleBuilder.create( + constOp.getLoc(), zeroAttr); + initTensors.push_back(zeroTensor); + } + } + return mlir::WalkResult::interrupt(); + } + return mlir::WalkResult::advance(); + }); + + auto initArgs = forOp.getInitArgs(); + + // 构建新的初始化参数列表 + SmallVector newInitArgs(initArgs.begin(), initArgs.end()); + // 添加 dependValue 的初始化参数 + for (Value initTensor : initTensors) { + newInitArgs.push_back(initTensor); + } + + // 获取原循环的边界和步长 + Value lb = forOp.getLowerBound(); + Value ub = forOp.getUpperBound(); + Value step = forOp.getStep(); + + // 创建新的 ForOp，插入点位于原操作之前 + OpBuilder builder(forOp); + auto newForOp = + builder.create(forOp.getLoc(), lb, ub, step, newInitArgs); + + // 获取新循环的 region 块（已自动包含循环索引和迭代参数） + Block &newBlock = newForOp.getRegion().front(); + Block &oldBlock = forOp.getRegion().front(); + + // 建立块参数的映射：原块参数 -> 新块参数 + IRMapping mapper; + for (unsigned i = 0; i < oldBlock.getNumArguments(); ++i) { + mapper.map(oldBlock.getArgument(i), newBlock.getArgument(i)); + } + // 将原循环体中的操作（不包括终结符）克隆到新块中 + // 同时按照顺序克隆新的 dependValues + SmallVector newDependValues = dependValues; + int cnt = 0; + builder.setInsertionPointToStart(&newBlock); + for (auto &op : oldBlock) { + auto newOp = builder.clone(op, mapper); + // dependValue 的定义OP 可能有多个 result + for (size_t i = 0; i < dependValues.size(); i++) { + Operation *defineOp = dependValues[i].getDefiningOp(); + if (defineOp == &op) { + unsigned int index = cast(dependValues[i]).getResultNumber(); + newDependValues[i] = newOp->getResult(index); + cnt++; + break; + } + } + } + // 判断是否找到了所有的 dependValue + if (newDependValues.size() != cnt) { + llvm::outs() << "can not find the depend value! \n"; + return; + } + dependValues = newDependValues; + + // 更新 mergedRegions 中的 op 为新的for循环的 op + UpdateMergedRegionsWithNewForOp(mergedRegions, mapper); + + // 创建新的循环 yield 操作：原操作数 + dependValues + auto oldYield = cast(newBlock.getTerminator()); + SmallVector newYieldOps(oldYield.getOperands()); + // 按顺序增加找到的 dependvalue + for (Value v : newDependValues) { + newYieldOps.push_back(v); + } + builder.setInsertionPointToEnd(&newBlock); + builder.create(oldYield.getLoc(), newYieldOps); + oldYield.erase(); + + // 将原 forOp 的所有使用替换为新 forOp + int oldResultNum = forOp->getResults().size(); + for (auto it : llvm::zip(forOp->getResults(), + newForOp->getResults().take_front(oldResultNum))) { + std::get<0>(it).replaceAllUsesWith(std::get<1>(it)); + } + forOp.erase(); +} + +void ComputeElseYieldValues(MergedRegion mergedRegion, + SmallVector &elseYieldValues, + SmallVector dependValues) { + int idx = 0; + for (Value v : mergedRegion.yieldValues) { + Type yieldType = mergedRegion.resultTypes[idx]; + elseYieldValues.push_back(findIterArg(v, yieldType)); + idx++; + } +} + +void ComputeElseYieldValuesV2(MergedRegion mergedRegion, + SmallVector &elseYieldValues, + SmallVector dependValues) { + // 对于yieldValues，其中的 yield value 一定是被 for op yield + // 所引用，或者被其他 region 所使用 + auto forOp = dyn_cast( + mergedRegion.yieldValues[0].getDefiningOp()->getBlock()->getParentOp()); + if (!forOp) { + llvm::outs() << "define op's parent is not ForOp \n"; + return; + } + auto iterArgs = forOp.getRegionIterArgs(); + auto forYieldValues = forOp.getYieldedValues(); + + // 新增的与 dependvalue 相关的 initarg + // 是接在原本for循环args后面，数量与dependvalue数量相等 + int baseDependIdx = iterArgs.size() - dependValues.size(); + + int idx = 0; + for (Value v : mergedRegion.yieldValues) { + Type yieldType = mergedRegion.resultTypes[idx]; + // yieldValue 中是dependvalue 的情况下 + // else yield value 使用对应的新增 iterargs + if (llvm::is_contained(dependValues, v)) { + int dependIdx = 0; + for (; dependIdx < dependValues.size(); dependIdx++) { + if (v == dependValues[dependIdx]) { + break; + } + } + // llvm::outs()<<"v2for:"< newYieldValues; + SmallVector newResultTypes; + + SmallPtrSet seen; + + for (auto [idx, v] : llvm::enumerate(region.yieldValues)) { + if (seen.insert(v).second) { + newYieldValues.push_back(v); + newResultTypes.push_back(region.resultTypes[idx]); + } + } + + region.yieldValues.swap(newYieldValues); + region.resultTypes.swap(newResultTypes); +} + +static void replaceExternalIfOpUses(scf::IfOp ifOp, + ArrayRef oldYieldValues) { + + for (size_t i = 0; i < oldYieldValues.size(); ++i) { + Value oldVal = oldYieldValues[i]; + Value newVal = ifOp.getResult(i); + + SmallVector usesToReplace; + + for (OpOperand &use : llvm::make_early_inc_range(oldVal.getUses())) { + + Operation *user = use.getOwner(); + + // 跳过 ifOp 内部的使用（then / else region） + if (ifOp->isAncestor(user)) + continue; + + // 只替换 ifOp 之后的使用 + if (user->getBlock() == ifOp->getBlock()) { + if (!ifOp->isBeforeInBlock(user)) + continue; + } + + usesToReplace.push_back(&use); + } + + for (OpOperand *use : usesToReplace) + use->set(newVal); + } +} + +void CreateIfOps(SmallVector &mergedRegions, + SmallVector dependValues) { + for (auto ®ion : mergedRegions) { + + // 去重yieldvalues + RemoveRedundantYieldValues(region); + + Operation *insertPt = region.opsToMove.front(); + OpBuilder builder(insertPt); + Location loc = insertPt->getLoc(); + Value cond = builder.create(loc, builder.getI1Type(), + builder.getBoolAttr(true)); + + bool needsYield = !region.yieldValues.empty(); + scf::IfOp ifOp; + if (needsYield) + ifOp = builder.create(loc, region.resultTypes, cond, true); + else + ifOp = builder.create(loc, TypeRange{}, cond, false); + + // 加标记 + ifOp->setAttr("ssbuffer", builder.getUnitAttr()); + + // 获取if yield value 在 else块 返回值 + SmallVector elseYieldValues; + + llvm::outs() << "before ComputeElseYieldValuesV2" << "\n"; + if (needsYield) { + // ComputeElseYieldValues(region, elseYieldValues, dependValues); + ComputeElseYieldValuesV2(region, elseYieldValues, dependValues); + } + + llvm::outs() << "after ComputeElseYieldValuesV2" << "\n"; + // 将op移进then块 + Block &thenBlock = ifOp.getThenRegion().front(); + for (Operation *m : llvm::reverse(region.opsToMove)) { + m->moveBefore(&thenBlock, thenBlock.begin()); + } + + // 创建 then/else yield + if (needsYield) { + OpBuilder thenBuilder(builder.getContext()); + thenBuilder.setInsertionPointToEnd(&thenBlock); + thenBuilder.create(loc, region.yieldValues); + + // else block + Block &elseBlock = ifOp.getElseRegion().front(); + OpBuilder elseBuilder(&elseBlock, elseBlock.end()); + elseBuilder.create(loc, elseYieldValues); + + // 替换外部使用 + + replaceExternalIfOpUses(ifOp, region.yieldValues); + + // 旧的逻辑 + // Block *block = ifOp->getBlock(); + // auto ifIt = Block::iterator(ifOp); + + // for (size_t i = 0; i < region.yieldValues.size(); ++i) { + // Value oldVal = region.yieldValues[i]; + // Value newVal = ifOp.getResult(i); + + // SmallVector usesToReplace; + + // for (OpOperand &use : llvm::make_early_inc_range(oldVal.getUses())) { + // Operation *user = use.getOwner(); + // // 同一个 block, user 必须在 ifOp 之后, 不能在 ifOp 内部（then / + // else） if (user->getBlock() != ifOp->getBlock() || + // !ifOp->isBeforeInBlock(user) || user->getParentOp() == ifOp) + // continue; + // usesToReplace.push_back(&use); + // } + + // for (OpOperand *use : usesToReplace) + // use->set(newVal); + // } + } + + llvm::outs() << "Create ifOp: " << *ifOp << "\n"; + } +} + +void CreateIfOpsOrigin(SmallVector &mergedRegions) { + for (auto ®ion : mergedRegions) { + + // 去重yieldvalues + RemoveRedundantYieldValues(region); + + Operation *insertPt = region.opsToMove.front(); + OpBuilder builder(insertPt); + Location loc = insertPt->getLoc(); + Value cond = builder.create(loc, builder.getI1Type(), + builder.getBoolAttr(true)); + + bool needsYield = !region.yieldValues.empty(); + scf::IfOp ifOp; + if (needsYield) + ifOp = builder.create(loc, region.resultTypes, cond, true); + else + ifOp = builder.create(loc, TypeRange{}, cond, false); + + // 加标记 + ifOp->setAttr("ssbuffer", builder.getUnitAttr()); + + // 将op移进then块 + Block &thenBlock = ifOp.getThenRegion().front(); + for (Operation *m : llvm::reverse(region.opsToMove)) { + m->moveBefore(&thenBlock, thenBlock.begin()); + } + + // 创建 then/else yield + if (needsYield) { + OpBuilder thenBuilder(builder.getContext()); + thenBuilder.setInsertionPointToEnd(&thenBlock); + thenBuilder.create(loc, region.yieldValues); + + // else block + SmallVector elseYieldValues; + int idx = 0; + for (Value v : region.yieldValues) { + Type yieldType = region.resultTypes[idx]; + elseYieldValues.push_back(findIterArgForAll(v, yieldType)); + idx++; + } + Block &elseBlock = ifOp.getElseRegion().front(); + OpBuilder elseBuilder(&elseBlock, elseBlock.end()); + elseBuilder.create(loc, elseYieldValues); + + // 替换外部使用 + Block *block = ifOp->getBlock(); + auto ifIt = Block::iterator(ifOp); + + for (size_t i = 0; i < region.yieldValues.size(); ++i) { + Value oldVal = region.yieldValues[i]; + Value newVal = ifOp.getResult(i); + + SmallVector usesToReplace; + + for (OpOperand &use : llvm::make_early_inc_range(oldVal.getUses())) { + Operation *user = use.getOwner(); + // 同一个 block, user 必须在 ifOp 之后, 不能在 ifOp 内部（then / + // else） + if (user->getBlock() != ifOp->getBlock() || + !ifOp->isBeforeInBlock(user) || user->getParentOp() == ifOp) + continue; + usesToReplace.push_back(&use); + } + + for (OpOperand *use : usesToReplace) + use->set(newVal); + } + } + + llvm::outs() << "Create ifOp: " << *ifOp << "\n"; + } +} + +void AddIfCondition(ModuleOp module) { + SmallVector copiedForOps; + SmallVector forOpList; + SmallVector, 1> regionList; + + module.walk([&](scf::ForOp forOp) { + Block &body = forOp.getRegion().front(); + SmallVector regions; + + // 获取基本的wait-set分块信息 + GetBlockInfos(regions, body); + + SmallVector mergedRegions; + // 合并wait-set块, 依据copyop / fixpipeop合并 + MergeWaitSetRegions(regions, mergedRegions); + + // 扩展if包裹的op范围 + // AIV、AIC处理有区别 + ExpandMergedRegionOps(forOp, mergedRegions, copiedForOps); + + // 处理forop的末尾对于iter_arg的自增操作, 如tt.advance, 移进对应的if op + MoveIterArgUsersIntoIf(forOp, mergedRegions); + + // 获取if yield的value, 并更新if内op的user为yield value + for (MergedRegion &mr : mergedRegions) { + // ComputeYieldForMergedRegion(mr, body); + ComputeYieldForMergedRegionV4(mr); + } + + // // 创建最终的if op + // CreateIfOpsOrigin(mergedRegions); + // }); + + forOpList.push_back(forOp); + regionList.push_back(mergedRegions); + }); + + llvm::outs() << "CopyForOp:\n"; + for (auto op : copiedForOps) { + llvm::outs() << *op << "\n"; + } + + SmallVector tmpOps; + for (auto copiedOp : copiedForOps) { + Block &body = copiedOp.getRegion().front(); + SmallVector regions; + + // 获取基本的wait-set分块信息 + GetBlockInfos(regions, body); + + SmallVector mergedRegions; + // 合并wait-set块, 依据copyop / fixpipeop合并 + MergeWaitSetRegions(regions, mergedRegions); + + // 扩展if包裹的op范围 + // AIV、AIC处理有区别 + ExpandMergedRegionOps(copiedOp, mergedRegions, tmpOps); + + // 处理forop的末尾对于iter_arg的自增操作, 如tt.advance, 移进对应的if op + MoveIterArgUsersIntoIf(copiedOp, mergedRegions); + + // 获取if yield的value, 并更新if内op的user为yield value + for (MergedRegion &mr : mergedRegions) { + // ComputeYieldForMergedRegion(mr, body); + ComputeYieldForMergedRegionV4(mr); + } + + // // 创建最终的if op + // CreateIfOpsOrigin(mergedRegions); + // }); + + forOpList.push_back(copiedOp); + regionList.push_back(mergedRegions); + } + + for (size_t i = 0; i < forOpList.size(); ++i) { + scf::ForOp oldForOp = forOpList[i]; + SmallVector newMergedRegions = regionList[i]; + + // 找到所有的VV或CC依赖 + SmallVector dependValues; + llvm::outs() << "FindDependValues! \n "; + FindDependValues(dependValues, newMergedRegions); + + if (dependValues.size() != 0) { + copyLoadCalculation(oldForOp, dependValues, newMergedRegions); + + // repeat previous operations + for (MergedRegion &mr : newMergedRegions) { + mr.yieldValues.clear(); + mr.resultTypes.clear(); + ComputeYieldForMergedRegionV4(mr); + } + FindDependValues(dependValues, newMergedRegions); + } + + // 如果存在VV或CC依赖，更新ForOp添加新的对应args + if (dependValues.size() != 0) { + AddArgsForDependValues(oldForOp, dependValues, newMergedRegions, module); + } + + // 创建最终的if op + llvm::outs() << "before create if ops" << '\n'; + CreateIfOps(newMergedRegions, dependValues); + } +} + +void ChangeAdvanceOpForm(ModuleOp module) { + module.walk([&](scf::ForOp forOp) { + Block &body = forOp.getRegion().front(); + constexpr int num = 8; + SmallVector ifOps; + for (Operation &op : body) + if (auto ifOp = dyn_cast(&op)) + ifOps.push_back(ifOp); + + for (scf::IfOp ifOp : ifOps) { + // 找 then region 中的 advance + triton::AdvanceOp advanceOp; + for (Operation &thenOp : ifOp.getThenRegion().front()) { + if (auto adv = dyn_cast(thenOp)) { + advanceOp = adv; + break; + } + } + if (!advanceOp) + continue; + + // base 必须是 for的iter_arg + Value base = advanceOp.getPtr(); + auto barg = dyn_cast(base); + if (!barg || barg.getOwner() != &body) + continue; + + // yield 去掉 advance 的返回值 + auto thenYield = + cast(ifOp.getThenRegion().front().getTerminator()); + auto elseYield = + cast(ifOp.getElseRegion().front().getTerminator()); + + int advanceIdx = -1; + for (auto it : llvm::enumerate(thenYield.getOperands())) { + if (it.value() == advanceOp.getResult()) { + advanceIdx = it.index(); + break; + } + } + + if (advanceIdx == -1) + continue; + + // 删除 advance + SmallVector thenOps(thenYield.getOperands().begin(), + thenYield.getOperands().end()); + SmallVector elseOps(elseYield.getOperands().begin(), + elseYield.getOperands().end()); + + thenOps.erase(thenOps.begin() + advanceIdx); + elseOps.erase(elseOps.begin() + advanceIdx); + + thenYield->setOperands(thenOps); + elseYield->setOperands(elseOps); + + // 重建 ifOp（去掉 advance 对应的 result） + OpBuilder ifBuilder(ifOp); + ifBuilder.setInsertionPoint(ifOp); + + // 构造新的 result types + SmallVector newResultTypes; + for (int i = 0; i < ifOp.getNumResults(); ++i) { + if (i != advanceIdx) + newResultTypes.push_back(ifOp.getResult(i).getType()); + } + + // 创建新的 if + auto newIf = ifBuilder.create(ifOp.getLoc(), newResultTypes, + ifOp.getCondition(), + /*withElseRegion=*/true); + newIf->setAttr("ssbuffer", ifBuilder.getUnitAttr()); + // 把已经修改过 yield 的 region 搬过去 + newIf.getThenRegion().takeBody(ifOp.getThenRegion()); + newIf.getElseRegion().takeBody(ifOp.getElseRegion()); + + // 替换if result的user + int newIdx = 0; + for (int oldIdx = 0; oldIdx < ifOp.getNumResults(); ++oldIdx) { + if (oldIdx == advanceIdx) + continue; + ifOp.getResult(oldIdx).replaceAllUsesWith(newIf.getResult(newIdx++)); + } + + OpBuilder builder(newIf); + builder.setInsertionPointAfter(newIf); + + Value flag = newIf.getCondition(); + + SmallVector newOffsets; + for (Value off : advanceOp.getOffsets()) { + auto intTy = cast(off.getType()); + auto zero = builder.create(newIf.getLoc(), 0, + intTy.getWidth()); + auto sel = + builder.create(newIf.getLoc(), flag, off, zero); + newOffsets.push_back(sel); + } + + auto newAdvance = builder.create( + newIf.getLoc(), base.getType(), base, newOffsets); + + // 原 if 的 advance result 的 users，接到 newAdvance + ifOp.getResult(advanceIdx).replaceAllUsesWith(newAdvance.getResult()); + + // 删除旧的ifOp和advance + advanceOp.erase(); + ifOp.erase(); + } + }); +} + +void processRedudantIf(ModuleOp module) { + SmallVector forOps; + llvm::outs() << module << " wwwww\n\n\n"; + module.walk([&](scf::ForOp forOp) { + auto initArgs = forOp.getInitArgs(); + if (initArgs.size() == 5) { + forOps.push_back(forOp); + } + }); + + for (auto forOp : forOps) { + auto initArgs = forOp.getInitArgs(); + Value newInit = initArgs[2]; + + // 构建新的初始化参数列表 + SmallVector newInitArgs(initArgs.begin(), initArgs.end()); + newInitArgs.push_back(newInit); + + // 获取原循环的边界和步长 + Value lb = forOp.getLowerBound(); + Value ub = forOp.getUpperBound(); + Value step = forOp.getStep(); + + // 创建新的 ForOp，插入点位于原操作之前 + OpBuilder builder(forOp); + auto newForOp = + builder.create(forOp.getLoc(), lb, ub, step, newInitArgs); + + // 获取新循环的 region 块（已自动包含循环索引和迭代参数） + Block &newBlock = newForOp.getRegion().front(); + Block &oldBlock = forOp.getRegion().front(); + + // 建立块参数的映射：原块参数 -> 新块参数（前6个对应） + IRMapping mapper; + for (unsigned i = 0; i < oldBlock.getNumArguments(); ++i) { + mapper.map(oldBlock.getArgument(i), newBlock.getArgument(i)); + } + // 将原循环体中的操作（不包括终结符）克隆到新块中 + builder.setInsertionPointToStart(&newBlock); + for (auto &op : oldBlock) { + auto newOp = builder.clone(op, mapper); + } + + // 在新块中查找第一个 scf::IfOp（即原代码中的第一个 if） + scf::IfOp firstIfOp = nullptr; + for (auto &op : newBlock.getOperations()) { + if (auto ifOp = dyn_cast(&op)) { + firstIfOp = ifOp; + break; + } + } + assert(firstIfOp && "Expected at least one if op in the loop body"); + + // 修改第一个 if 的 else 分支的 yield 操作： + // 将其第二个操作数（索引1）从原来的 %arg9 改为新迭代参数（新块参数索引6） + Block &elseBlock = firstIfOp.getElseRegion().front(); + auto elseYield = cast(elseBlock.getTerminator()); + SmallVector newElseYieldOps(elseYield.getOperands()); + newElseYieldOps[1] = newBlock.getArgument(6); // 新迭代参数 + builder.setInsertionPoint(elseYield); + builder.create(elseYield.getLoc(), newElseYieldOps); + elseYield->erase(); + + // 创建新的循环 yield 操作：原5个操作数 + 第一个 if 的第二个结果 + auto oldYield = cast(newBlock.getTerminator()); + SmallVector newYieldOps(oldYield.getOperands()); + newYieldOps.push_back(firstIfOp.getResult(1)); // 第一个 if 的第二个结果 + builder.setInsertionPointToEnd(&newBlock); + builder.create(oldYield.getLoc(), newYieldOps); + oldYield.erase(); + + // 将原 forOp 的所有使用替换为新 forOp 的前5个结果 + for (auto it : + llvm::zip(forOp->getResults(), newForOp->getResults().take_front(5))) { + std::get<0>(it).replaceAllUsesWith(std::get<1>(it)); + } + } + for (auto forOp : forOps) { + forOp.erase(); + } +} +// 针对依赖变量，对原本的for op增加double buffer相关的迭代参数 +scf::ForOp addDoubleBuffForArgs(ModuleOp module, SmallVector uniqueDeps, + int bufferNum) { + mlir::OpBuilder builder(module.getContext()); + SmallVector depValueForIdxs; + + // ========== 找到scf.if所在的scf::ForOp ========== + if (!isa(uniqueDeps[0].getDefiningOp()->getParentOp())) { + llvm::errs() << "Error: parent op of scf.if is not scf.for"; + } + scf::ForOp forOp = + dyn_cast(uniqueDeps[0].getDefiningOp()->getParentOp()); + + for (Value dependencyValue : uniqueDeps) { + // ========== 步骤1：验证目标Value是scf.if的返回值，并找到对应的scf::IfOp + // ========== + Operation *ifOp = dependencyValue.getDefiningOp(); + if (!ifOp || !isa(ifOp)) { + llvm::errs() << "Error: 目标Value不是scf.if的返回值\n"; + return nullptr; + } + scf::IfOp targetIfOp = dyn_cast(ifOp); + + // 确认当前Value是scf.if的第几个返回值 + int64_t depValueIdx = -1; + for (auto [idx, result] : llvm::enumerate(targetIfOp.getResults())) { + if (result == dependencyValue) { + depValueIdx = idx; + break; + } + } + + // ========== 步骤2：找到%38#2关联的scf.for迭代参数以及索引 ========== + // %38#2对应scf.if else分支yield的第2个操作数 → 即%arg10 + Operation *elseYield = targetIfOp.elseYield(); + Value dependencyArg = elseYield->getOperand( + depValueIdx); // depValueIdx=2，对应else yield的第2个参数 + + int64_t depValueForIdx = -1; + for (auto [idx, result] : llvm::enumerate(forOp.getRegionIterArgs())) { + if (result == dependencyArg) { + depValueForIdx = idx; + break; + } + } + depValueForIdxs.push_back(depValueForIdx); + llvm::outs() << "depValueForIdx: " << depValueForIdx << '\n'; + } + + llvm::outs() << "oldFor: " << forOp << '\n'; + + // 获取原始循环的信息 + Value originalLowerBound = forOp.getLowerBound(); + Value originalUpperBound = forOp.getUpperBound(); + Value originalStep = forOp.getStep(); + SmallVector originalInitArgs = forOp.getInitArgs(); + SmallVector iterArgs; + for (auto arg : originalInitArgs) { + iterArgs.push_back(arg); + } + auto yields = forOp.getBody()->getTerminator(); + + // 创建计数器初始零值 + Value counterInit = nullptr; + mlir::Operation *parentOp = forOp->getParentOp(); + mlir::Operation *scopeOp = nullptr; + // 向上遍历查找scope.scope操作 + while (parentOp) { + if (dyn_cast(parentOp)) { + scopeOp = parentOp; + break; + } + parentOp = parentOp->getParentOp(); + } + + builder.setInsertionPoint(scopeOp); + Location loc = forOp.getLoc(); + auto boundType = originalLowerBound.getType(); + counterInit = builder.create(loc, boundType, 0); + + // 添加和depValueForIdxs相同的迭代参数和计数器 + for (int64_t idx : depValueForIdxs) { + for (int i = 0; i < bufferNum - 1; i++) { + iterArgs.push_back(originalInitArgs[idx]); + } + + // 在迭代参数中添加计数器 + for (int i = 0; i < 2; i++) { + iterArgs.push_back(counterInit); + } + } + + builder.setInsertionPoint(forOp); + // 创建新的for循环 + auto newForOp = + builder.create(forOp.getLoc(), originalLowerBound, + originalUpperBound, originalStep, iterArgs); + + // 设置IR映射表，将旧循环的变量映射到新循环 + IRMapping mapper; + + // 映射迭代变量 + mapper.map(forOp.getInductionVar(), newForOp.getInductionVar()); + + // 映射迭代参数 + for (auto [oldArg, newArg] : + llvm::zip(forOp.getRegionIterArgs(), newForOp.getRegionIterArgs())) { + mapper.map(oldArg, newArg); + } + + SmallVector newArgs; + for (int i = forOp.getRegionIterArgs().size(); + i < newForOp.getRegionIterArgs().size(); i++) { + newArgs.push_back(newForOp.getRegionIterArgs()[i]); + } + // 克隆循环体内容到新循环 + auto &newLoopBody = *newForOp.getBody(); + builder.setInsertionPointToStart(&newLoopBody); + + for (auto &op : forOp.getBody()->without_terminator()) { + builder.clone(op, mapper); + } + + // 克隆yield操作 + if (auto yieldOp = dyn_cast(yields)) { + SmallVector newYieldOperands; + for (auto operand : yieldOp.getOperands()) { + newYieldOperands.push_back(mapper.lookupOrDefault(operand)); + } + // 将新增的迭代参数添加到yield操作数中 + for (auto currentCounter : newArgs) { + newYieldOperands.push_back(currentCounter); + } + builder.create(yieldOp.getLoc(), newYieldOperands); + } + + // 替换原循环的结果 + unsigned numOriginalResults = forOp.getNumResults(); + SmallVector originalResults; + for (unsigned i = 0; i < numOriginalResults; i++) { + originalResults.push_back(newForOp.getResult(i)); + } + forOp.replaceAllUsesWith(originalResults); + + // 8. 删除原循环 + forOp.erase(); + + llvm::outs() << "for op erased!\n"; + return newForOp; +} + +static bool isScalarValue(Value value) { + if (!value) + return false; + Type type = value.getType(); + + return type.isIntOrIndex() || mlir::isa(type) || + mlir::isa(type) || + mlir::isa(type); +} + +SmallVector buildNBufferProducer(OpBuilder &builder, Location loc, + Value frontCnt, Value newDepVal, + ArrayRef buffs, + ArrayRef constants) { + // N-buffer producer: determines which buffer is written to newDepVal based on + // frontCnt % N + const int N = buffs.size(); + SmallVector results; + + // idx = frontCnt % N + Value bufferIndex = + builder.create(loc, frontCnt, constants[N]); + + // 1. buffer0: handle the first buffer separately + Value isBuffer0 = builder.create(loc, arith::CmpIPredicate::eq, + bufferIndex, constants[0]); + + auto dstShapedType = mlir::dyn_cast(newDepVal.getType()); + Value mask; + if (isScalarValue(isBuffer0) == 0) { + auto maskType = + RankedTensorType::get(dstShapedType.getShape(), isBuffer0.getType()); + mask = builder.create(loc, maskType, isBuffer0); + llvm::outs() << "build NBufferProducer debug15\n"; + llvm::outs().flush(); + } else { + mask = isBuffer0; + } + Value newBuff0 = + builder.create(loc, mask, newDepVal, buffs[0]); + + results.push_back(newBuff0); + + // 2. Double-buffer specialization (when N == 2, a direct select is + // sufficient) + if (N == 2) { + + Value newBuff1 = + builder.create(loc, mask, buffs[1], newDepVal); + + auto nextCnt = builder.create(loc, frontCnt, constants[1]); + + results.push_back(newBuff1); + results.push_back(nextCnt.getResult()); + + return results; + } + + // 3. Build the root IF: when idx == 0, + // use the first buffer; otherwise enter the nestedIf chain to use other + // buffers + SmallVector resultTypes; + for (int i = 1; i < N; ++i) + resultTypes.push_back(buffs[i].getType()); + + auto rootIf = builder.create(loc, resultTypes, isBuffer0, true); + + // ---- THEN: buffers are directly forwarded ---- + { + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPointToStart(&rootIf.getThenRegion().front()); + + SmallVector unchangedBuffers(buffs.begin() + 1, buffs.end()); + + builder.create(loc, unchangedBuffers); + } + + // 4. Construct the nested-if chain, updating one buffer at each level + Block *currentElseBlock = &rootIf.getElseRegion().front(); + + scf::IfOp parentIf = rootIf; + + for (int i = 1; i < N - 1; ++i) { + + builder.setInsertionPointToStart(currentElseBlock); + + // Check whether the current buffer is selected + Value isCurrent = builder.create( + loc, arith::CmpIPredicate::eq, bufferIndex, constants[i]); + + // Update buffer[i] + dstShapedType = mlir::dyn_cast(newDepVal.getType()); + auto maskType = + RankedTensorType::get(dstShapedType.getShape(), isCurrent.getType()); + mask = builder.create(loc, maskType, isCurrent); + Value updatedBuffer = + builder.create(loc, mask, newDepVal, buffs[i]); + + // If this is the last level: directly yield both buffers + if (i == N - 2) { + + dstShapedType = mlir::dyn_cast(newDepVal.getType()); + maskType = + RankedTensorType::get(dstShapedType.getShape(), isCurrent.getType()); + mask = builder.create(loc, maskType, isCurrent); + Value lastBuffer = + builder.create(loc, mask, buffs[N - 1], newDepVal); + + builder.create(loc, ValueRange{updatedBuffer, lastBuffer}); + + break; + } + + // Create the next nested if + SmallVector subResultTypes; + for (int j = i + 1; j < N; ++j) + subResultTypes.push_back(buffs[j].getType()); + + auto nextIf = + builder.create(loc, subResultTypes, isCurrent, true); + + // THEN: forward the remaining buffers + { + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPointToStart(&nextIf.getThenRegion().front()); + + SmallVector remainingBuffers(buffs.begin() + i + 1, buffs.end()); + + builder.create(loc, remainingBuffers); + } + + // Update the else yield + builder.setInsertionPointToEnd(&parentIf.getElseRegion().front()); + + SmallVector yields; + yields.push_back(updatedBuffer); + yields.append(nextIf.getResults().begin(), nextIf.getResults().end()); + + builder.create(loc, yields); + + parentIf = nextIf; + currentElseBlock = &nextIf.getElseRegion().front(); + } + + // 5. Update the frontCnt counter + builder.setInsertionPointAfter(rootIf); + + auto nextCnt = builder.create(loc, frontCnt, constants[1]); + + // Collect results + results.append(rootIf.getResults().begin(), rootIf.getResults().end()); + + results.push_back(nextCnt.getResult()); + + return results; +} + +SmallVector buildNBufferConsumer(OpBuilder &builder, Location loc, + Value postCnt, ArrayRef oldBuffs, + ArrayRef constants) { + // Consumer: selects which buffer to read based on postCnt % N + const int bufferNum = oldBuffs.size(); + SmallVector results; + + // idx = postCnt % N + Value bufferIndex = + builder.create(loc, postCnt, constants[bufferNum]); + + Value isBuffer0 = builder.create(loc, arith::CmpIPredicate::eq, + bufferIndex, constants[0]); + auto dstShapedType = mlir::dyn_cast(oldBuffs[0].getType()); + Value mask; + if (isScalarValue(isBuffer0) == 0) { + auto maskType = + RankedTensorType::get(dstShapedType.getShape(), isBuffer0.getType()); + mask = builder.create(loc, maskType, isBuffer0); + } else { + mask = isBuffer0; + } + + // 1. Double-buffer specialization (avoid generating scf.if) + if (bufferNum == 2) { + Value selected = + builder.create(loc, mask, oldBuffs[0], oldBuffs[1]); + auto nextCnt = builder.create(loc, postCnt, constants[1]); + + results.push_back(selected); + results.push_back(nextCnt); + + return results; + } + + // 2. Build the root IF: + // when idx == 0, use the first buffer; otherwise enter the nestedIf chain to + // use other buffers + SmallVector resultTypes{oldBuffs[0].getType()}; + + auto rootIf = builder.create(loc, resultTypes, isBuffer0, true); + + // ---- THEN: directly return buffer0 ---- + { + builder.setInsertionPointToStart(&rootIf.getThenRegion().front()); + + builder.create(loc, oldBuffs[0]); + } + + // 3. Construct the nested-if chain + Block *currentElse = &rootIf.getElseRegion().front(); + + for (int i = 1; i < bufferNum - 2; ++i) { + + builder.setInsertionPointToStart(currentElse); + + Value isCurrent = builder.create( + loc, arith::CmpIPredicate::eq, bufferIndex, constants[i]); + + auto nestedIf = builder.create( + loc, TypeRange{oldBuffs[0].getType()}, isCurrent, true); + + // THEN → return the current buffer + { + builder.setInsertionPointToStart(&nestedIf.getThenRegion().front()); + + builder.create(loc, oldBuffs[i]); + } + + // ELSE → yield nested result + builder.setInsertionPointToEnd(currentElse); + builder.create(loc, nestedIf.getResult(0)); + + // Enter the next else branch + currentElse = &nestedIf.getElseRegion().front(); + } + + // 4. Final level (use select to finish) + builder.setInsertionPointToStart(currentElse); + + int last = bufferNum - 2; + + Value isLast = builder.create(loc, arith::CmpIPredicate::eq, + bufferIndex, constants[last]); + + auto maskType = RankedTensorType::get({}, isLast.getType()); + dstShapedType = mlir::dyn_cast(oldBuffs[last].getType()); + maskType = RankedTensorType::get(dstShapedType.getShape(), isLast.getType()); + mask = builder.create(loc, maskType, isLast); + + Value finalSelect = builder.create(loc, mask, oldBuffs[last], + oldBuffs[last + 1]); + + builder.create(loc, finalSelect); + + // rootIf result = selected buffer + results.push_back(rootIf.getResult(0)); + + // 5. Update the postCnt counter + builder.setInsertionPointAfter(rootIf); + + auto nextCnt = builder.create(loc, postCnt, constants[1]); + + results.push_back(nextCnt); + + return results; +} + +void replaceDepsMap(scf::IfOp oldIfOp, scf::IfOp newIfOp, + SmallVector &newDeps, bool isFront, + DenseMap> &newIfResultDeps) { + mlir::IRMapping valueMap; + + // old result -> new result + for (unsigned i = 0; i < oldIfOp.getNumResults(); ++i) { + valueMap.map(oldIfOp.getResult(i), newIfOp.getResult(i)); + } + + if (isFront) { + for (int i = 0; i < newDeps.size(); i++) { + Value v = newDeps[i]; + if (valueMap.contains(v)) + newDeps[i] = valueMap.lookup(v); + } + } + + // rewrite deps in-place + for (auto &it : newIfResultDeps) { + auto &deps = it.second; + + for (auto &value : deps) { + if (auto mapped = valueMap.lookupOrNull(value)) + value = mapped; + } + } +} + +scf::IfOp addResultsForFrontIfOp( + scf::IfOp frontIfOp, OpBuilder builder, int bufferNum, Value depValue, + SmallVector constants, SmallVector buffs, Value frontCnt, + Value postCnt, SmallVector &extraResultIndices, + SmallVector &newDeps, + DenseMap> &newIfResultDeps) { + OpBuilder::InsertionGuard guard(builder); + + Location loc = frontIfOp.getLoc(); + Value cond = frontIfOp.getCondition(); + + auto &oldThenBlock = frontIfOp.getThenRegion().front(); + auto &oldElseBlock = frontIfOp.getElseRegion().front(); + + // New result types = old results + extra buffers + counter + SmallVector newResultTypes(frontIfOp.getResultTypes().begin(), + frontIfOp.getResultTypes().end()); + + for (int i = 1; i < bufferNum; ++i) + newResultTypes.push_back(buffs[i].getType()); + + newResultTypes.push_back(frontCnt.getType()); + + unsigned oldNumResults = frontIfOp.getNumResults(); + + // Create new IfOp + builder.setInsertionPoint(frontIfOp); + auto newIfOp = + builder.create(loc, newResultTypes, cond, /*hasElse=*/true); + + SmallVector bufferIndices(bufferNum); + SmallVector newBuffs; + int frontCntIndex = -1; + + // THEN region + { + mlir::IRMapping mapping; + Block &newThenBlock = newIfOp.getThenRegion().front(); + + builder.setInsertionPointToStart(&newThenBlock); + + // Clone original then body + for (auto &op : oldThenBlock.without_terminator()) + builder.clone(op, mapping); + + // Update dependency value position inf ifOp results + auto result = dyn_cast(depValue); + if (!result) { + llvm::outs() << "depValue is not a result Value!\n"; + return nullptr; + } + + int depIdx = result.getResultNumber(); + Value depYieldValue = frontIfOp.thenYield()->getOperand(depIdx); + + Value newDepVal = mapping.contains(depYieldValue) + ? mapping.lookup(depYieldValue) + : depYieldValue; + + builder.setInsertionPointAfter(newDepVal.getDefiningOp()); + + // Create N buffer + SmallVector produced = buildNBufferProducer( + builder, loc, frontCnt, newDepVal, buffs, constants); + + // Last value in newBuffs is the counter + newBuffs.append(produced.begin(), produced.end() - 1); + + // Rebuild new yield + SmallVector thenOperands; + + for (Value v : oldThenBlock.getTerminator()->getOperands()) { + Value mapped = mapping.lookupOrDefault(v); + + // Replace first buffer + if (mapped == newDepVal) { + thenOperands.push_back(newBuffs[0]); + bufferIndices[0] = thenOperands.size() - 1; + } else { + thenOperands.push_back(mapped); + } + } + + // Replace other buffer + for (int i = 1; i < bufferNum; ++i) { + thenOperands.push_back(newBuffs[i]); + bufferIndices[i] = thenOperands.size() - 1; + } + + // Add counter + thenOperands.push_back(produced.back()); + frontCntIndex = thenOperands.size() - 1; + + builder.setInsertionPointToEnd(&newThenBlock); + builder.create(loc, thenOperands); + + // record new result indices + for (int idx : bufferIndices) + extraResultIndices.push_back(idx); + + extraResultIndices.push_back(frontCntIndex); + } + + // ELSE region + { + mlir::IRMapping mapping; + Block &newElseBlock = newIfOp.getElseRegion().front(); + + builder.setInsertionPointToStart(&newElseBlock); + + // Clone original else body + for (auto &op : oldElseBlock.without_terminator()) + builder.clone(op, mapping); + + builder.setInsertionPointToEnd(&newElseBlock); + + SmallVector elseOperands; + + for (Value v : oldElseBlock.getTerminator()->getOperands()) + elseOperands.push_back(mapping.lookupOrDefault(v)); + + // Add buffer + for (int i = 1; i < bufferNum; ++i) + elseOperands.push_back(buffs[i]); + + // Add counter + elseOperands.push_back(frontCnt); + + builder.create(loc, elseOperands); + } + + // Update dependency value + replaceDepsMap(frontIfOp, newIfOp, newDeps, true, newIfResultDeps); + + // Replace old ifOp + frontIfOp.replaceAllUsesWith(newIfOp.getResults().take_front(oldNumResults)); + + frontIfOp.erase(); + + return newIfOp; +} + +scf::IfOp addResultsForPostIfOp( + scf::IfOp postIfOp, scf::IfOp newfrontIfOp, OpBuilder builder, + int bufferNum, Value newDepValue, SmallVector constants, + SmallVector buffs, Value frontCnt, Value postCnt, + SmallVector &extraResultIndices, SmallVector &newDeps, + DenseMap> &newIfResultDeps) { + // 1. Parse the extra result indices produced by frontIf (added buffers and + // counters) + SmallVector bufferIndices(extraResultIndices.begin(), + extraResultIndices.end() - 1); + int frontCntIndex = extraResultIndices[bufferNum]; + + Location ifLoc = postIfOp.getLoc(); + Value cond = postIfOp.getCondition(); + + auto &oldThenBlock = postIfOp.getThenRegion().front(); + auto &oldElseBlock = postIfOp.getElseRegion().front(); + + // 2. Create a new IfOp (add a new postCnt result) + SmallVector newResultTypes(postIfOp.getResultTypes().begin(), + postIfOp.getResultTypes().end()); + newResultTypes.push_back(postCnt.getType()); + + builder.setInsertionPoint(postIfOp); + auto newIfOp = builder.create(ifLoc, newResultTypes, cond, + /*hasElse=*/true); + + mlir::IRMapping mapping; + + // 3. THEN region: clone the original logic, insert the multibuffer consumer + // and update dependency buffers + auto &newThenBlock = newIfOp.getThenRegion().front(); + builder.setInsertionPointToStart(&newThenBlock); + + // clone then body + for (auto &op : oldThenBlock.without_terminator()) + builder.clone(op, mapping); + builder.setInsertionPointToStart(&newThenBlock); + + // Find dependency uses that need to be replaced (located inside the current + // IfOp) + SmallVector replaceUses; + for (auto &use : newDepValue.getUses()) { + if (newIfOp == dyn_cast(use.getOwner()->getParentOp())) { + replaceUses.push_back(&use); + } + } + + // Collect buffers produced by frontIf + SmallVector oldBuffers; + for (int i = 0; i < bufferIndices.size(); ++i) + oldBuffers.push_back(newfrontIfOp.getResult(bufferIndices[i])); + + // Multibuffer consumer caculation + SmallVector consumerResults = + buildNBufferConsumer(builder, ifLoc, postCnt, oldBuffers, constants); + + Value selectedBuffer = consumerResults[0]; + Value nextPostCnt = consumerResults[1]; + + // Replace dependent buffer + for (auto *usePtr : replaceUses) { + usePtr->set(selectedBuffer); + } + + // Create then yield + SmallVector thenOperands; + for (auto v : oldThenBlock.getTerminator()->getOperands()) + thenOperands.push_back(mapping.lookupOrDefault(v)); + + int postCntIndex = thenOperands.size(); + thenOperands.push_back(nextPostCnt); + + builder.setInsertionPointToEnd(&newThenBlock); + builder.create(ifLoc, thenOperands); + extraResultIndices.push_back(postCntIndex); + + // 4. ELSE region：forward counter directly + auto &newElseBlock = newIfOp.getElseRegion().front(); + + for (auto &op : oldElseBlock.without_terminator()) + builder.clone(op, mapping); + + builder.setInsertionPointToEnd(&newElseBlock); + + SmallVector elseOperands; + for (auto v : oldElseBlock.getTerminator()->getOperands()) + elseOperands.push_back(mapping.lookupOrDefault(v)); + + elseOperands.push_back(postCnt); + + builder.create(ifLoc, elseOperands); + + // 5. Replace old ifOp with new one + auto oldNumResults = postIfOp.getNumResults(); + + // Update depency value + replaceDepsMap(postIfOp, newIfOp, newDeps, false, newIfResultDeps); + + postIfOp.replaceAllUsesWith(newIfOp.getResults().take_front(oldNumResults)); + + postIfOp.erase(); + + return newIfOp; +} + +void addMultiBuffCaculate(ModuleOp module, SmallVector newUniqueDeps, + DenseMap> &ifResultDeps, + scf::ForOp &newForOp, int bufferNum) { + + // ============================================================ + // Overall Idea + // + // For each dependency Value: + // 1. Find the front IfOp that produces it + // 2. Add multi-buffer results to the front IfOp + // 3. Find the post IfOp that consumes the result and extend it accordingly + // 4. Update the for-loop yield so that buffer states are correctly propagated + // ============================================================ + + OpBuilder builder(module.getContext()); + int processedDepCount = 0; + + SmallVector postIfOps; + newForOp.walk([&](scf::IfOp postIfOp) { postIfOps.push_back(postIfOp); }); + for (auto postIfOp : postIfOps) { + if (!ifResultDeps.count(postIfOp)) { + continue; + } + auto newDeps = ifResultDeps[postIfOp]; + for (int depValueIdx = 0; depValueIdx < newDeps.size(); depValueIdx++) { + Value depValue = newDeps[depValueIdx]; + + // Step 1. Locate the front IfOp that produces depValue + Operation *defOp = depValue.getDefiningOp(); + if (!defOp || !isa(defOp)) { + llvm::outs() << "Error: depValue is not produced by scf.if\n"; + break; + } + + scf::IfOp frontIfOp = cast(defOp); + + // Position of depValue in the IfOp results + auto result = dyn_cast(depValue); + if (!result) { + llvm::outs() << "depValue is not an OpResult!\n"; + return; + } + + int64_t depResultIndex = result.getResultNumber(); + + // Position of depValue in the IfOp results + Value depYieldValue = frontIfOp.thenYield()->getOperand(depResultIndex); + + // Step 2. Find the multi-buffer position in the ForOp + int64_t extraArgBaseIdx = + newForOp.getRegionIterArgs().size() - + (2 + bufferNum - 1) * (newUniqueDeps.size() - processedDepCount++); + + // Collect all buffers + SmallVector buffers; + + // buffer0 来自 else yield + buffers.push_back(frontIfOp.elseYield()->getOperand(depResultIndex)); + + // Other buffers come from for iter args + for (int i = 1; i < bufferNum; ++i) { + buffers.push_back( + newForOp.getRegionIterArgs()[extraArgBaseIdx + i - 1]); + } + + // Two counters + Value frontCnt = + newForOp.getRegionIterArgs()[extraArgBaseIdx + bufferNum - 1]; + Value postCnt = newForOp.getRegionIterArgs()[extraArgBaseIdx + bufferNum]; + + // Step 3. Create constants (0 ~ bufferNum) for rem / cmp buffer selection + // logic + SmallVector constants; + builder.setInsertionPoint(frontIfOp); + + auto dataType = frontCnt.getType(); + for (int i = 0; i <= bufferNum; ++i) { + constants.push_back(builder.create( + frontIfOp.getLoc(), dataType, builder.getIntegerAttr(dataType, i))); + } + + // Record the positions of newly added results in the IfOp + SmallVector extraResultIndices(bufferNum + 1); + extraResultIndices.clear(); + + // Step 4. Extend the front IfOp + scf::IfOp newFrontIfOp = addResultsForFrontIfOp( + frontIfOp, builder, bufferNum, depValue, constants, buffers, frontCnt, + postCnt, extraResultIndices, newDeps, ifResultDeps); + + // buffer result indices + SmallVector bufferResultIndices(extraResultIndices.begin(), + extraResultIndices.end() - 1); + + int frontCntResultIndex = extraResultIndices[bufferNum]; + + Value newDepValue = newFrontIfOp.getResult(depResultIndex); + + // Step 5. Find the post IfOp that consumes the dependency value + scf::IfOp postIfOp = nullptr; + + for (auto &use : newDepValue.getUses()) { + if (auto candidate = + dyn_cast(use.getOwner()->getParentOp())) { + postIfOp = candidate; + break; + } + } + + if (!postIfOp) { + llvm::outs() << "Error: no consuming IfOp found.\n"; + return; + } + + // Step 6. Extend the post IfOp + + scf::IfOp newPostIfOp = addResultsForPostIfOp( + postIfOp, newFrontIfOp, builder, bufferNum, newDepValue, constants, + buffers, frontCnt, postCnt, extraResultIndices, newDeps, + ifResultDeps); + + llvm::outs() << "after addResultsForPostIfOp.\n"; + + int postCntResultIndex = extraResultIndices.back(); + + // Step 7. Update the ForOp yield (buffer propagation) + auto forYield = cast(newForOp.getBody()->getTerminator()); + + // Update buffer1 ~ bufferN + for (int i = 1; i < bufferNum; ++i) { + + int yieldIdx = extraArgBaseIdx + (i - 1); + + if (yieldIdx < forYield->getNumOperands() && + bufferResultIndices[i] < newFrontIfOp.getNumResults()) { + + forYield->setOperand(yieldIdx, + newFrontIfOp.getResult(bufferResultIndices[i])); + + llvm::outs() << "Replaced yield operand " << yieldIdx << "\n"; + } else { + llvm::errs() << "Warning: index out of range\n"; + } + } + + // Step 8. Update frontCnt + OpOperand *frontCntYieldUse = nullptr; + + for (auto &use : frontCnt.getUses()) { + if (isa(use.getOwner()) && + newForOp == use.getOwner()->getParentOp()) { + frontCntYieldUse = &use; + break; + } + } + + frontCntYieldUse->set(newFrontIfOp.getResult(frontCntResultIndex)); + + // Step 9. Update postCnt + OpOperand *postCntYieldUse = nullptr; + + for (auto &use : postCnt.getUses()) { + if (isa(use.getOwner()) && + newForOp == use.getOwner()->getParentOp()) { + postCntYieldUse = &use; + break; + } + } + + postCntYieldUse->set(newPostIfOp.getResult(postCntResultIndex)); + } + } + + llvm::outs() << "multibuffer end!\n"; +} + +// Compute the nesting level of an ifOp within the specified forOp +static int computeIfLevel(scf::IfOp ifOp, scf::ForOp rootForOp) { + int level = 1; + + Operation *parent = ifOp->getParentOp(); + + while (parent && parent != rootForOp.getOperation()) { + if (isa(parent)) + level++; + + parent = parent->getParentOp(); + } + + return level; +} + +int assignIfOpLevels(scf::ForOp forOp) { + SmallVector targetIfOps; + int maxLevel = 0; + // Collect all ifOp assigned with ssbuffer tag + forOp.walk([&](scf::IfOp ifOp) { + if (ifOp->hasAttr("ssbuffer")) { + targetIfOps.push_back(ifOp); + } + }); + + // Caculate buffer levels + for (auto ifOp : targetIfOps) { + int level = computeIfLevel(ifOp, forOp); + maxLevel = std::max(level, maxLevel); + Builder builder(ifOp.getContext()); + ifOp->setAttr("ssbuffer.level", builder.getI32IntegerAttr(level)); + } + return maxLevel; +} + +static bool hasSSBufferIf(scf::ForOp forOp) { + bool found = false; + + forOp.walk([&](scf::IfOp ifOp) { + if (ifOp->hasAttr("ssbuffer")) { + found = true; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + + return found; +} + +static bool hasAncestorSSBufferFor(scf::ForOp forOp) { + Operation *parent = forOp->getParentOp(); + + while (parent) { + if (auto parentFor = dyn_cast(parent)) { + if (hasSSBufferIf(parentFor)) + return true; + } + parent = parent->getParentOp(); + } + + return false; +} + +static bool hasAncestorRootFor(scf::ForOp forOp) { + Operation *parent = forOp->getParentOp(); + + while (parent) { + if (auto parentFor = dyn_cast(parent)) { + if (hasSSBufferIf(parentFor)) + return true; + } + parent = parent->getParentOp(); + } + return false; +} + +SmallVector +collectIfInfo(scf::ForOp &curForOp, + DenseMap> &ifDeps, int level) { + // Find all dependency variables based on the inputs and outputs of ifOp + SmallVector allDeps; + DenseSet producedValues; + scf::ForOp newForOp = nullptr; + curForOp.walk([&](scf::IfOp ifOp) { + auto attr = ifOp->getAttrOfType("ssbuffer.level"); + // No level or level mismatch → continue searching + if (!attr || attr.getInt() != level) + return WalkResult::advance(); + + // Levels match → check the direct parent + if (auto parentFor = dyn_cast(ifOp->getParentOp())) { + newForOp = parentFor; // 更新 + } + + // Stop walking regardless of whether the parent is a for-loop + return WalkResult::interrupt(); + }); + + if (newForOp) + curForOp = newForOp; + + // Step 1: Collect first to preserve order + SmallVector ifOps; + curForOp.walk([&](scf::IfOp ifOp) { + auto curLevel = ifOp->getAttrOfType("ssbuffer.level"); + if (!curLevel || curLevel.getInt() != level) { + return WalkResult::advance(); + } + ifOps.push_back(ifOp); + return WalkResult::advance(); + }); + llvm::outs() << "ifOps:" << ifOps.size() << "\n"; + + int miniDepNum = 2; + if (ifOps.size() < miniDepNum) { + return allDeps; + } + // Step 2: Process in order + for (auto ifOp : ifOps) { + llvm::outs() << "ifOp->getOperands():" << ifOp->getOperands().size() + << "\n"; + SmallVector deps; + if (producedValues.empty()) { + llvm::outs() << "producedValues为空!" << "\n"; + } + + // inputs + Region &thenRegion = ifOp.getThenRegion(); + for (Operation &op : thenRegion.front()) { + for (Value operand : op.getOperands()) { + for (Value v : producedValues) { + if (operand == v && !llvm::is_contained(deps, operand)) { + deps.push_back(operand); + } + } + } + } + + // outputs + for (Value result : ifOp.getResults()) { + producedValues.insert(result); + } + + if (!deps.empty()) { + ifDeps[ifOp] = deps; + allDeps.append(deps.begin(), deps.end()); + } + } + llvm::outs().flush(); + return allDeps; +} + +bool isCube(scope::ScopeOp scope) { + bool ret = false; + scope.walk([&](Operation *op) { + if (isa(op)) { + ret = true; + } + }); + return ret; +} + +// Traverse each Vector scope, find the outer ForOp, and process internal IfOps +void WalkAIVNestedForAndProcess( + ModuleOp module, DenseMap> &ifResultDeps, + int bufferNum) { + if (bufferNum < 2) { + return; + } + + module.walk([&](scope::ScopeOp scope) { + if (isCube(scope)) { + return; + } + + // Traverse ForOps inside the Cube scope (outer loops) + SmallVector targetFors; + + scope.walk([&](scf::ForOp forOp) { + // Must contain an ssbuffer if + if (!hasSSBufferIf(forOp)) + return WalkResult::advance(); + + // Skip if an ancestor is already the root + if (hasAncestorRootFor(forOp)) + return WalkResult::advance(); + + // Find rootForOp + targetFors.push_back(forOp); + + return WalkResult::advance(); + }); + llvm::outs() << "targetFors: " << targetFors.size(); + int maxLevels; + for (auto outerFor : targetFors) { + ifResultDeps.clear(); + scf::ForOp currentFor = outerFor; + maxLevels = assignIfOpLevels(currentFor); + for (int level = 1; level <= maxLevels; level++) { + auto uniqueDeps = collectIfInfo(currentFor, ifResultDeps, level); + llvm::outs() << "maxLevels:" << maxLevels << "\n"; + if (uniqueDeps.empty()) { + continue; + } + llvm::outs() << "uniqueDeps:" << uniqueDeps.size() << "\n"; + auto newForOp = addDoubleBuffForArgs(module, uniqueDeps, bufferNum); + DenseMap> newIfResultDeps; + auto uniqueList = collectIfInfo(newForOp, newIfResultDeps, level); + addMultiBuffCaculate(module, uniqueList, newIfResultDeps, newForOp, + bufferNum); + } + } + }); +} + +void DAGSSBufferPass::runOnOperation() { + auto module = getOperation(); + + AddIfCondition(module); + + FlowSssbuf(module); + ControlSsbufV2(module); + + // advance不能出现在if里, 规避处理 + ChangeAdvanceOpForm(module); + + DenseMap> ifResultDeps; + WalkAIVNestedForAndProcess(module, ifResultDeps, 2); + + return; +} + +std::unique_ptr> mlir::triton::createDAGSSBufferPass() { + return std::make_unique(); +} diff --git a/third_party/ascend/lib/TritonAffinityOpt/DAGScope.cpp b/third_party/ascend/lib/TritonAffinityOpt/DAGScope.cpp new file mode 100644 index 0000000000..6d44a94200 --- /dev/null +++ b/third_party/ascend/lib/TritonAffinityOpt/DAGScope.cpp @@ -0,0 +1,1104 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "TritonAffinityOpt/Passes.h" + +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/HIVM/IR/HIVMImpl.h" +#include "bishengir/Dialect/HIVM/IR/HIVMInterfaces.h" +#include "bishengir/Dialect/HIVM/Transforms/Passes.h" +#include "bishengir/Dialect/HIVM/Utils/Utils.h" +#include "bishengir/Dialect/Scope/IR/Scope.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/Operation.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "incubated/Conversion/UtilsIncubated/Utils.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/SmallVector.h" +#include + +#include "TritonAffinityOpt/DAG.h" + +namespace mlir { +namespace triton { +#define GEN_PASS_DEF_DAGSCOPE +#include "ascend/include/TritonAffinityOpt/Passes.h.inc" +} // namespace triton +} // namespace mlir + +using namespace mlir; +using namespace hivm; + +namespace { +struct DAGScopePass : public mlir::triton::impl::DAGScopeBase { + void runOnOperation() override; +}; +} // namespace + +static std::pair +encapsulateWithScope(triton::FuncOp funcOp) { + Block &entryBlock = funcOp.getBody().front(); + Block &lastBlock = funcOp.getBody().back(); + Operation *terminator = lastBlock.getTerminator(); + + // 辅助函数：判断操作是否应该被跳过 + auto shouldSkipOp = [](Operation *op) -> bool { + return isa(op) || isa(op) || + isa(op); + }; + + // 第三步：准备要移动的操作列表（按顺序） + SmallVector opsToMove; + DenseMap opOrder; + int order = 0; + + // 记录原始顺序并收集需要移动的操作 + for (Operation &op : lastBlock.without_terminator()) { + opOrder[&op] = order++; + if (!shouldSkipOp(&op)) { + opsToMove.push_back(&op); + } + } + + // 按原始顺序排序 + std::sort( + opsToMove.begin(), opsToMove.end(), + [&](Operation *a, Operation *b) { return opOrder[a] < opOrder[b]; }); + + if (opsToMove.empty()) { + return std::make_pair(nullptr, nullptr); + } + + // 第四步：创建scope操作并移动操作 + Operation *lastOpToMove = opsToMove.back(); + OpBuilder builder(&lastBlock, ++lastOpToMove->getIterator()); + + // 创建第一个scope + auto scopeOp = builder.create(builder.getUnknownLoc(), + llvm::ArrayRef{}); + scopeOp.getBodyRegion().emplaceBlock(); + Block *scopeBody = &scopeOp.getBodyRegion().front(); + + // 移动操作到scope中 + OpBuilder scopeBuilder(scopeBody, scopeBody->end()); + DenseMap valueMapping; + + for (Operation *op : opsToMove) { + SmallVector originalResults = op->getResults(); + op->remove(); + scopeBuilder.insert(op); + + // 更新值的映射 + for (size_t i = 0; i < originalResults.size(); ++i) { + valueMapping[originalResults[i]] = op->getResult(i); + } + } + + // 添加return操作 + scopeBuilder.create(builder.getUnknownLoc()); + + // 创建第二个scope（如果需要） + scopeBuilder.setInsertionPointAfter(scopeOp); + auto newScopeOp = scopeBuilder.create( + builder.getUnknownLoc(), llvm::ArrayRef{}); + newScopeOp.getRegion().emplaceBlock(); + + OpBuilder newScopeBuilder(&newScopeOp.getRegion().front(), + newScopeOp.getRegion().front().begin()); + newScopeBuilder.create(scopeOp->getLoc()); + + // 设置属性 + auto vecAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::VECTOR); + auto aicAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + + scopeOp->setAttr(hivm::TCoreTypeAttr::name, vecAttr); + newScopeOp->setAttr(hivm::TCoreTypeAttr::name, aicAttr); + + return std::make_pair(scopeOp, newScopeOp); +} + +struct OpMoveInfo { + Operation *op; + Operation *targetParent; // 目标父操作（nullptr表示aicScope本身） +}; + +// 递归遍历函数 - 优化版本 +void collectOpsToMove(Operation *op, AffinityDAG::Graph &graph, + Operation *parentFor, + llvm::SmallVector &aivToMove, + llvm::SmallVector &cubeToMove) { + // 检查当前操作是否需要移动 + bool needsMoveAiv = false; + bool needsMoveCube = false; + auto &valueTypes = graph.getValueTypes(); + // 检查结果类型 + int i = 0; + for (auto res : op->getResults()) { + i++; + if (AffinityDAG::intersects(valueTypes[res], + AffinityDAG::CoreType::VECTOR_ONLY)) { + needsMoveAiv = true; + } + if (AffinityDAG::intersects(valueTypes[res], + AffinityDAG::CoreType::CUBE_ONLY)) { + needsMoveCube = true; + } + } + + if (isa(op)) { + auto res = op->getOperand(0); + if (AffinityDAG::intersects(valueTypes[res], + AffinityDAG::CoreType::VECTOR_ONLY)) { + needsMoveAiv = true; + } + if (AffinityDAG::intersects(valueTypes[res], + AffinityDAG::CoreType::CUBE_ONLY)) { + needsMoveCube = true; + } + } + // 检查特定操作类型 + if (isa(op)) { + needsMoveAiv = true; + } + + // 检查特定操作类型 + if (isa(op)) { + needsMoveCube = true; + } + + // 检查特定操作类型 + if (isa(op) || isa(op) || isa(op)) { + needsMoveAiv = true; + needsMoveCube = true; + } + + if (isa(op)) { + if (auto storeOp = dyn_cast(op)) { + // 获取所有操作数列表 + auto operands = storeOp.getOperands(); + bool typeMatched = false; + + // 按顺序检查第1个、第0个、第2个操作数 + std::vector checkOrder = {1, 0, 2}; + for (size_t idx : checkOrder) { + // 先判断操作数索引是否有效，避免越界访问 + if (idx >= operands.size()) { + continue; + } + auto operand = operands[idx]; + auto coreType = valueTypes[operand]; + + if (coreType == AffinityDAG::CoreType::VECTOR_ONLY) { + needsMoveAiv = true; + typeMatched = true; + } else if (coreType == AffinityDAG::CoreType::CUBE_ONLY) { + needsMoveCube = true; + typeMatched = true; + } + } + // 所有指定操作数都不匹配时，执行原else逻辑 + if (!typeMatched) { + needsMoveAiv = true; + needsMoveCube = true; + } + } + } + + if (isa(op)) { + if (auto assertOp = dyn_cast(op)) { + // 获取所有操作数列表 + auto operand = assertOp.getCondition(); + + auto coreType = valueTypes[operand]; + if (coreType == AffinityDAG::CoreType::VECTOR_ONLY) { + needsMoveAiv = true; + } else if (coreType == AffinityDAG::CoreType::CUBE_ONLY) { + needsMoveCube = true; + } else { + needsMoveAiv = true; + needsMoveCube = true; + } + } + } + + // 检查 Sync 操作的 tcore_type 属性 + if ((isa(op) || isa(op))) { + mlir::OpBuilder builder(op); + auto coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + if (op->getAttr("tcore_type") == coreAttr) { + needsMoveCube = true; + } else { + needsMoveAiv = true; + } + } + + // 如果不需要移动，直接返回 + if (!needsMoveAiv && !needsMoveCube) { + llvm::outs() << "Unsupport Op: " << *op << " \n"; + } + + // 处理 for 循环 + if (auto forOp = dyn_cast(op)) { + // 确定父级 for 循环 + Operation *targetParent = parentFor != nullptr ? parentFor : nullptr; + aivToMove.push_back({op, targetParent}); + cubeToMove.push_back({op, targetParent}); + + // 递归处理循环体 + for (auto &block : forOp.getRegion()) { + for (auto &innerOp : block) { + collectOpsToMove(&innerOp, graph, forOp, aivToMove, cubeToMove); + } + } + } else if (auto ifOp = dyn_cast(op)) { + // 确定父级 for 循环 + Operation *targetParent = parentFor != nullptr ? parentFor : nullptr; + aivToMove.push_back({op, targetParent}); + cubeToMove.push_back({op, targetParent}); + + // 递归处理循环体 + for (auto &block : ifOp.getThenRegion()) { + for (auto &innerOp : block) { + collectOpsToMove(&innerOp, graph, ifOp, aivToMove, cubeToMove); + } + } + + // 检查并遍历IfOp的else分支（如果存在） + for (auto &block : ifOp.getElseRegion()) { + for (auto &innerOp : block) { + collectOpsToMove(&innerOp, graph, ifOp, aivToMove, cubeToMove); + } + } + } else { + if (needsMoveAiv) { + // 处理其他操作 + aivToMove.push_back({op, parentFor}); + } + if (needsMoveCube) { + cubeToMove.push_back({op, parentFor}); + } + } +} + +mlir::Block *getBlockByIndex(mlir::Region ®ion, int blockIndex) { + // 边界校验：索引非法时返回nullptr + if (blockIndex < 0) + return nullptr; + + int currentIdx = 0; + for (auto &block : region) { + if (currentIdx == blockIndex) { + return █ // 找到对应索引的Block，直接返回 + } + currentIdx++; + } + // 索引越界时返回nullptr + return nullptr; +} + +void processOperationToMove( + const OpMoveInfo &info, + llvm::DenseMap &parentMap, + mlir::OpBuilder &builder, mlir::IRMapping &mapper, mlir::Block *aivBlock, + mlir::Operation *terminator, AffinityDAG::Graph &graph, int MoveType) { + // llvm::outs()<<*info.op<<" ssss\n\n\n"; + // llvm::outs().flush(); + // 获取原始Block信息并计算索引 + mlir::Block *originalBlock = info.op->getBlock(); + int originalRegionIndex = -1; + int originalBlockIndex = -1; + int blockCounter = 0; + auto &valueTypes = graph.getValueTypes(); + if (originalBlock) { + mlir::Operation *parentOp = info.op->getParentOp(); // 原始父操作 + if (parentOp) { // 确保父操作存在 + // 老版本MLIR用 getParent() 替代 getParentRegion()，返回值就是Region* + mlir::Region *blockBelongsToRegion = originalBlock->getParent(); + int regionCounter = 0; + for (auto ®ion : parentOp->getRegions()) { // 遍历父操作的所有region + // 直接对比指针，判断当前region是否是block所属的region + if (®ion == blockBelongsToRegion) { + originalRegionIndex = regionCounter; + break; + } + regionCounter++; + } + } + } + + if (originalBlock) { + for (auto &block : originalBlock->getParent()->getBlocks()) { + if (&block == originalBlock) { + originalBlockIndex = blockCounter; + break; + } + blockCounter++; + } + } + + if (originalBlockIndex == -1) { + originalBlockIndex = 0; + } + if (originalRegionIndex == -1) { + originalRegionIndex = 0; + } + + // 处理 scf::ForOp 类型操作 + if (mlir::isa(info.op)) { + auto forOp = mlir::cast(info.op); + + auto getMapped = [&](mlir::Value v) { return mapper.lookupOrDefault(v); }; + auto inputs = forOp.getInitArgs(); + auto outputs = forOp.getResults(); + + // 分离需要移动到aivScope的参数 + llvm::SmallVector aivInputs; + llvm::DenseMap aivInputsMap; + int aivIndex = 1; + + for (int i = 0; i < inputs.size(); ++i) { + if (valueTypes[outputs[i]] != MoveType) { + aivInputs.push_back(inputs[i]); + aivInputsMap[i + 1] = aivIndex; + aivIndex++; + } + } + + // 创建新的for循环 + auto aivForOp = builder.create( + forOp.getLoc(), getMapped(forOp.getLowerBound()), + getMapped(forOp.getUpperBound()), getMapped(forOp.getStep()), + llvm::to_vector(llvm::map_range(aivInputs, getMapped))); + + // 清空循环体 + if (!aivForOp.getBody()->empty()) { + aivForOp.getBody()->getTerminator()->erase(); + } + + // 处理原始循环的yield操作 + auto oldBody = forOp.getBody(); + auto oldYield = + mlir::dyn_cast(oldBody->getTerminator()); + assert(oldYield && "scf::ForOp must have a yield terminator"); + + llvm::SmallVector aivYieldOperands; + for (int i = 0; i < inputs.size(); ++i) { + if (valueTypes[outputs[i]] != MoveType) { + aivYieldOperands.push_back(oldYield.getOperand(i)); + } + } + + // 映射循环参数 + auto oldBodyArgs = forOp.getBody()->getArguments(); + auto aivBodyArgs = aivForOp.getBody()->getArguments(); + + for (auto it = aivInputsMap.begin(); it != aivInputsMap.end(); ++it) { + int oldInputIndex = it->first; + int mappedNewIndex = it->second; + mapper.map(oldBodyArgs[oldInputIndex], aivBodyArgs[mappedNewIndex]); + mapper.map((*info.op).getResults()[oldInputIndex - 1], + aivForOp->getResults()[mappedNewIndex - 1]); + } + mapper.map(oldBodyArgs[0], aivBodyArgs[0]); + + // 将新循环移动到目标位置 + if (info.targetParent == nullptr) { + mlir::Block *targetBlock = aivBlock; + if (terminator) { + aivForOp->moveBefore(terminator); + } else { + aivForOp->moveBefore(targetBlock, targetBlock->end()); + } + parentMap[forOp] = aivForOp; + } else { + auto targetParent = parentMap[info.targetParent]; + auto ®ion = targetParent->getRegion(originalRegionIndex); + + if (region.empty()) { + region.push_back(new mlir::Block()); + } + + mlir::Block *targetBlock = getBlockByIndex(region, originalBlockIndex); + if (targetBlock) { + aivForOp->moveBefore(targetBlock, targetBlock->end()); + parentMap[forOp] = aivForOp; + } else { + llvm::outs() << "Can't find block by index\n"; + } + } + } + + // 处理 scf::YieldOp 类型操作 + else if (mlir::isa(info.op)) { + auto yieldOp = mlir::cast(info.op); + + // 处理父节点为 scf::ForOp 的情况 + if (auto parentForOp = + mlir::dyn_cast(info.targetParent)) { + auto it = parentMap.find(parentForOp); + if (it == parentMap.end()) { + return; + } + auto targetOp = it->second; + auto newForOp = mlir::cast(targetOp); + + auto oldInputs = parentForOp.getInitArgs(); + auto oldOutputs = parentForOp.getResults(); + auto oldYieldOperands = yieldOp.getOperands(); + + llvm::SmallVector newYieldOperands; + for (int i = 0; i < oldInputs.size(); ++i) { + if (valueTypes[oldOutputs[i]] != MoveType) { + mlir::Value oldOperand = oldYieldOperands[i]; + mlir::Value newOperand = mapper.lookupOrDefault(oldOperand); + newYieldOperands.push_back(newOperand); + } + } + + auto newYieldOp = builder.create(yieldOp.getLoc(), + newYieldOperands); + auto ®ion = newForOp->getRegion(0); + mlir::Block *targetBlock = ®ion.front(); + newYieldOp->moveBefore(targetBlock, targetBlock->end()); + } + // 处理父节点为 scf::IfOp 的情况 + else if (auto parentIfOp = + mlir::dyn_cast(info.targetParent)) { + auto it = parentMap.find(parentIfOp); + if (it == parentMap.end()) { + return; + } + auto targetOp = it->second; + auto newIfOp = mlir::cast(targetOp); + + auto oldInputs = parentIfOp.getResults(); + auto oldOutputs = parentIfOp.getResults(); + auto oldYieldOperands = yieldOp.getOperands(); + + llvm::SmallVector newYieldOperands; + for (int i = 0; i < oldInputs.size(); ++i) { + if (valueTypes[oldOutputs[i]] != MoveType) { + mlir::Value oldOperand = oldYieldOperands[i]; + mlir::Value newOperand = mapper.lookupOrDefault(oldOperand); + newYieldOperands.push_back(newOperand); + } + } + + auto ®ion = newIfOp->getRegion(originalRegionIndex); + auto newYieldOp = builder.create(yieldOp.getLoc(), + newYieldOperands); + mlir::Block *targetBlock = getBlockByIndex(region, originalBlockIndex); + if (targetBlock) { + newYieldOp->moveBefore(targetBlock, targetBlock->end()); + } else { + llvm::outs() << "Can't find block by index\n"; + } + } + } + + // 处理 scf::IfOp 类型操作 + else if (mlir::isa(info.op)) { + auto ifOp = mlir::cast(info.op); + + auto getMapped = [&](mlir::Value v) { return mapper.lookupOrDefault(v); }; + mlir::Value condition = ifOp.getCondition(); + + // 分离需要移动到aivScope的结果 + llvm::SmallVector aivResults; + llvm::SmallVector aivResultTypes; + llvm::DenseMap aivResultMap; + int aivResultIndex = 0; + + for (int i = 0; i < ifOp.getNumResults(); ++i) { + mlir::Value result = ifOp.getResult(i); + if (valueTypes[result] != MoveType) { + aivResults.push_back(result); + aivResultTypes.push_back(result.getType()); + aivResultMap[i] = aivResultIndex; + aivResultIndex++; + } + } + + // 创建新的if操作 + auto aivIfOp = builder.create( + ifOp.getLoc(), aivResultTypes, getMapped(condition)); + + // 映射if操作结果 + for (auto &[oldIdx, newIdx] : aivResultMap) { + mapper.map(ifOp.getResult(oldIdx), aivIfOp.getResult(newIdx)); + } + + // 初始化then和else区域 + mlir::Region &thenRegion = aivIfOp.getThenRegion(); + mlir::Block *thenBlock = new mlir::Block(); + thenRegion.push_back(thenBlock); + + mlir::Region &elseRegion = ifOp.getElseRegion(); + if (!elseRegion.empty()) { + mlir::Region &elseRegion = aivIfOp.getElseRegion(); + mlir::Block *elseBlock = new mlir::Block(); + elseRegion.push_back(elseBlock); + } + + // 将新if操作移动到目标位置 + if (info.targetParent == nullptr) { + mlir::Block *targetBlock = aivBlock; + if (terminator) { + aivIfOp->moveBefore(terminator); + } else { + aivIfOp->moveBefore(targetBlock, targetBlock->end()); + } + parentMap[ifOp] = aivIfOp; + } else { + auto ®ion = + parentMap[info.targetParent]->getRegion(originalRegionIndex); + if (region.empty()) { + region.push_back(new mlir::Block()); + } + + mlir::Block *targetBlock = getBlockByIndex(region, originalBlockIndex); + if (targetBlock) { + aivIfOp->moveBefore(targetBlock, targetBlock->end()); + parentMap[ifOp] = aivIfOp; + } else { + llvm::outs() << "Can't find block by index\n"; + } + } + } + + // 处理其他类型操作（克隆） + else { + auto clonedOp = builder.clone(*info.op, mapper); + auto numberRes = clonedOp->getNumResults(); + for (auto i = 0; i < numberRes; i++) { + mapper.map((*info.op).getResults()[i], clonedOp->getResults()[i]); + } + + if (info.targetParent == nullptr) { + mlir::Block *targetBlock = aivBlock; + clonedOp->moveBefore(terminator); + parentMap[info.op] = clonedOp; + } else { + auto parentIt = parentMap.find(info.targetParent); + auto mappedParentOp = parentIt->second; + auto ®ion = mappedParentOp->getRegion(originalRegionIndex); + + if (region.empty()) { + region.push_back(new mlir::Block()); + } + + mlir::Block *targetBlock = getBlockByIndex(region, originalBlockIndex); + if (targetBlock) { + clonedOp->moveBefore(targetBlock, targetBlock->end()); + } else { + llvm::outs() << "Can't find block by index\n"; + } + } + } +} + +static void SplitScope(triton::FuncOp funcOp, AffinityDAG::Graph &graph, + Operation *aivScope, Operation *aicScope, + ModuleOp module) { + llvm::SmallVector aivToMove; + llvm::SmallVector cubeToMove; + for (auto &block : aivScope->getRegion(0)) { + for (auto &op : block) { + collectOpsToMove(&op, graph, nullptr, aivToMove, cubeToMove); + } + } + mlir::IRMapping aivmapper; + mlir::OpBuilder builder(aivScope); + llvm::DenseMap aivparentMap; + + // 第二遍：实际移动操作 + // 先移动for循环 + mlir::Block *aivBlock = + &aivScope->getRegion(0).front(); // 或者使用合适的block + SmallVector deleteOp; + auto *terminator = aivBlock->getTerminator(); + // 如果操作已被使用，直接跳过 + llvm::SmallVector + aivUsedOp; // 改为函数内静态，保持原有逻辑 + for (const auto &info : aivToMove) { + if (std::find(aivUsedOp.begin(), aivUsedOp.end(), info.op) != + aivUsedOp.end()) { + return; + } + aivUsedOp.push_back(info.op); + processOperationToMove(info, aivparentMap, builder, aivmapper, aivBlock, + terminator, graph, AffinityDAG::CoreType::CUBE_ONLY); + } + + llvm::DenseMap aicparentMap; + mlir::IRMapping aicmapper; + mlir::Block *aicBlock = + &aicScope->getRegion(0).front(); // 或者使用合适的block + terminator = aicBlock->getTerminator(); + llvm::SmallVector + aicUsedOp; // 改为函数内静态，保持原有逻辑 + for (const auto &info : cubeToMove) { + if (std::find(aicUsedOp.begin(), aicUsedOp.end(), info.op) != + aicUsedOp.end()) { + return; + } + aicUsedOp.push_back(info.op); + processOperationToMove(info, aicparentMap, builder, aicmapper, aicBlock, + terminator, graph, + AffinityDAG::CoreType::VECTOR_ONLY); + } + + for (const auto &info : aivToMove) { + if (std::find(deleteOp.begin(), deleteOp.end(), info.op) == + deleteOp.end()) { + deleteOp.push_back(info.op); + } + } + for (const auto &info : cubeToMove) { + if (std::find(deleteOp.begin(), deleteOp.end(), info.op) == + deleteOp.end()) { + deleteOp.push_back(info.op); + } + } + + // llvm::outs() << "\n" << module<<" ====== ddd ====== \n\n\n"; + // llvm::outs().flush(); + for (auto it = deleteOp.rbegin(); it != deleteOp.rend(); ++it) { + (*it)->erase(); // 解引用反向迭代器，调用 erase 方法 + } + return; +} + +/// 创建setop +static hivm::SyncBlockSetOp +createSyncBlockSetOp(OpBuilder &builder, Location loc, hivm::TCoreType coreType, + hivm::PIPE setPipeEnum, hivm::PIPE waitPipeEnum, + int64_t flag) { + MLIRContext *ctx = builder.getContext(); + auto coreAttr = hivm::TCoreTypeAttr::get(ctx, coreType); + auto setPipe = hivm::PipeAttr::get(ctx, setPipeEnum); + auto waitPipe = hivm::PipeAttr::get(ctx, waitPipeEnum); + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + return builder.create(loc, coreAttr, setPipe, waitPipe, + flagId); +} + +/// 创建waitop +static hivm::SyncBlockWaitOp +createSyncBlockWaitOp(OpBuilder &builder, Location loc, + hivm::TCoreType coreType, hivm::PIPE setPipeEnum, + hivm::PIPE waitPipeEnum, int64_t flag) { + MLIRContext *ctx = builder.getContext(); + auto coreAttr = hivm::TCoreTypeAttr::get(ctx, coreType); + auto setPipe = hivm::PipeAttr::get(ctx, setPipeEnum); + auto waitPipe = hivm::PipeAttr::get(ctx, waitPipeEnum); + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + return builder.create(loc, coreAttr, setPipe, waitPipe, + flagId); +} + +// 在scope return前插入wait +static void insertWaitBeforeFinalReturn(Region *region, OpBuilder &builder, + int64_t flag, bool coretypebool) { + for (Block &block : *region) { + if (auto returnOp = + dyn_cast_or_null(block.getTerminator())) { + builder.setInsertionPoint(returnOp); + if (coretypebool) { + createSyncBlockWaitOp(builder, returnOp->getLoc(), + hivm::TCoreType::CUBE, hivm::PIPE::PIPE_V, + hivm::PIPE::PIPE_FIX, flag); + return; + } else { + createSyncBlockWaitOp(builder, returnOp->getLoc(), + hivm::TCoreType::VECTOR, hivm::PIPE::PIPE_M, + hivm::PIPE::PIPE_MTE3, flag); + return; + } + } + } +} + +/// 在scope内起始位置加上set +static void insertSetAtRegionStart(Region *region, OpBuilder &builder, + int64_t flag, bool coretypebool) { + if (!region->empty()) { + Block &entry = region->front(); + Location loc = entry.empty() ? region->getParentOp()->getLoc() + : entry.front().getLoc(); + builder.setInsertionPointToStart(&entry); + if (coretypebool) { + createSyncBlockSetOp(builder, loc, hivm::TCoreType::VECTOR, + hivm::PIPE::PIPE_V, hivm::PIPE::PIPE_FIX, flag); + } else { + createSyncBlockSetOp(builder, loc, hivm::TCoreType::CUBE, + hivm::PIPE::PIPE_M, hivm::PIPE::PIPE_MTE3, flag); + } + } +} + +static Operation *findNextSyncBlockSetAfter(Operation *startOp) { + Block *block = startOp->getBlock(); + auto it = ++startOp->getIterator(); + for (; it != block->end(); ++it) { + if (isa(*it)) + return &*it; + } + return nullptr; +} + +static hivm::SyncBlockWaitOp findWaitOpInRegionWithFlag(Region *region, + int64_t flag) { + hivm::SyncBlockWaitOp result; + region->walk([&](hivm::SyncBlockWaitOp op) { + auto flagAttr = op->getAttrOfType("static_flag_id"); + if (flagAttr && flagAttr.getInt() == flag) { + result = op; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + return result; +} + +static Operation *findInsertionPointAfterWaitForAIV(Operation *waitOp) { + Block *block = waitOp->getBlock(); + auto it = ++waitOp->getIterator(); + + for (; it != block->end(); ++it) { + if (isa(*it) || isa(*it)) { + break; + } + } + + while (it != block->begin()) { + auto prevIt = std::prev(it); + if (isa(*prevIt)) { + it = prevIt; + } else { + break; + } + } + + return &*it; +} + +static Operation *findInsertionPointAfterWaitForAIC(Operation *waitOp) { + Block *block = waitOp->getBlock(); + auto it = ++waitOp->getIterator(); + for (; it != block->end(); ++it) { + if (auto fixpipe = dyn_cast(*it)) { + if (it != block->begin()) { + auto prev = std::prev(it); + if (isa(*prev)) + return &*prev; + } + return &*it; + } + if (isa(*it)) + return &*it; + } + return nullptr; +} + +// 查找 FixpipeOp 下一行的 sync_block_set 操作的 flag 值 +static int findFixPipeFlagSafe(hivm::FixpipeOp fixpipeOp) { + mlir::Operation *fixpipeOperation = fixpipeOp.getOperation(); + if (!fixpipeOperation || !fixpipeOperation->getBlock()) { + return -1; + } + + // 获取 FixpipeOp 的迭代器 + auto it = ++fixpipeOperation->getIterator(); + + // 遍历后续操作直到找到 sync_block_set + while (it != fixpipeOperation->getBlock()->end()) { + mlir::Operation &op = *it++; + + if (op.getName().getStringRef() == "hivm.hir.sync_block_set") { + auto staticFlagAttr = + op.getAttrOfType("static_flag_id"); + return staticFlagAttr.getInt(); + break; + } + } + + return -1; +} + +/// cube处理逻辑 +static void processFixpipeOpsInAIC(Region *aicRegion, Region *aivRegion) { + + MLIRContext *ctx = aicRegion->getContext(); + OpBuilder builder(ctx); + SmallVector fixpipes; + aicRegion->walk([&](hivm::FixpipeOp op) { fixpipes.push_back(op); }); + + for (auto fixpipeOp : fixpipes) { + + auto newflag = findFixPipeFlagSafe(fixpipeOp); + // 1. 在 FixpipeOp 前插 Wait + builder.setInsertionPoint(fixpipeOp); + createSyncBlockWaitOp(builder, fixpipeOp->getLoc(), hivm::TCoreType::CUBE, + hivm::PIPE::PIPE_V, hivm::PIPE::PIPE_FIX, newflag); + bool coretypebool = true; + + // 2. 在 aicRegion 末尾 Return 前插 Wait + insertWaitBeforeFinalReturn(aicRegion, builder, newflag, coretypebool); + + // 3. 在 aivRegion 开头插 Set + insertSetAtRegionStart(aivRegion, builder, newflag, coretypebool); + + // 4. 在 aicRegion 向后找 SyncBlockSetOp + if (auto *nextSetOp = findNextSyncBlockSetAfter(fixpipeOp)) { + auto setFlagAttr = + nextSetOp->getAttrOfType("static_flag_id"); + // 调试：打印set + // llvm::dbgs() << "aicnextSetOp:"; + // nextSetOp->dump(); + if (!setFlagAttr) { + llvm::dbgs() << "AIC can not find setop in aic\n"; + continue; + } + int64_t setflag = setFlagAttr.getInt(); + + // 5. 在 aivRegion 中找 flag=setflag 的 WaitOp + auto targetWait = findWaitOpInRegionWithFlag(aivRegion, setflag); + if (!targetWait) { + llvm::dbgs() << "AIC can not find waitop in aiv\n"; + continue; + } + + // 调试：打印wait + // llvm::dbgs() << "aictargetWait:"; + // llvm::dbgs() << targetWait << "\n"; + + // 6. 从该 Wait 向下找 ToMemrefOp 或 Yield，插 Set(newflag) + if (auto *insertPt = findInsertionPointAfterWaitForAIV(targetWait)) { + builder.setInsertionPoint(insertPt); + createSyncBlockSetOp(builder, fixpipeOp->getLoc(), + hivm::TCoreType::VECTOR, hivm::PIPE::PIPE_V, + hivm::PIPE::PIPE_FIX, newflag); + } + } + } +} + +// 查找 copyOp 下一行的 sync_block_set 操作的 flag 值 +static int findCopyFlagSafe(bufferization::ToBufferOp toMemrefOp) { + mlir::Operation *toMemrefOperation = toMemrefOp.getOperation(); + if (!toMemrefOperation || !toMemrefOperation->getBlock()) { + return -1; + } + + // 获取 copyOp 的迭代器 + auto it = ++toMemrefOperation->getIterator(); + + // 遍历后续操作直到找到 sync_block_set + while (it != toMemrefOperation->getBlock()->end()) { + mlir::Operation &op = *it++; + + if (op.getName().getStringRef() == "hivm.hir.sync_block_set") { + auto staticFlagAttr = + op.getAttrOfType("static_flag_id"); + return staticFlagAttr.getInt(); + break; + } + } + + return -1; +} +/// vector处理逻辑 +static void processToMemrefOpsInAIV(Region *aivRegion, Region *aicRegion) { + + MLIRContext *ctx = aivRegion->getContext(); + OpBuilder builder(ctx); + SmallVector toMemrefs; + aivRegion->walk( + [&](bufferization::ToBufferOp op) { toMemrefs.push_back(op); }); + + for (auto toMemrefOp : toMemrefs) { + auto newflag = findCopyFlagSafe(toMemrefOp); + + // 1. 在 ToMemrefOp 前插 Wait + builder.setInsertionPoint(toMemrefOp); + createSyncBlockWaitOp(builder, toMemrefOp->getLoc(), + hivm::TCoreType::VECTOR, hivm::PIPE::PIPE_M, + hivm::PIPE::PIPE_MTE3, newflag); + bool coretypebool = false; + + // 2. 在 aivRegion 末尾 Return 前插 Wait + insertWaitBeforeFinalReturn(aivRegion, builder, newflag, coretypebool); + + // 3. 在 aicRegion 开头插 Set + insertSetAtRegionStart(aicRegion, builder, newflag, coretypebool); + + // 4. 在 aivRegion 向后找 SyncBlockSetOp + if (auto *nextSetOp = findNextSyncBlockSetAfter(toMemrefOp)) { + auto setFlagAttr = + nextSetOp->getAttrOfType("static_flag_id"); + // 调试：打印set及其所有attribute + // llvm::dbgs() << "aivnextSetOp:"; + // nextSetOp->dump(); + // llvm::dbgs() << "Attributes:\n"; + // for (auto namedAttr : nextSetOp->getAttrs()) { + // llvm::dbgs() << " " << namedAttr.getName() << " = "; + // namedAttr.getValue().print(llvm::dbgs()); + // llvm::dbgs() << "\n"; + // } + if (!setFlagAttr) { + llvm::dbgs() << "AIV can not find setop in aiv\n"; + continue; + } + int64_t setflag = setFlagAttr.getInt(); + + // 5. 在 aicRegion 中找 flag=setflag 的 WaitOp + auto targetWait = findWaitOpInRegionWithFlag(aicRegion, setflag); + + if (!targetWait) { + llvm::dbgs() << "AIV can not find waitop in aic\n"; + continue; + } + + // 调试：打印wait + // llvm::dbgs() << "aivtargetWait:"; + // llvm::dbgs() << targetWait << "\n"; + + // 6. 从该 Wait 向下找 Fixpipe 前 Wait 或 Yield，插 Set(newflag) + if (auto *insertPt = findInsertionPointAfterWaitForAIC(targetWait)) { + builder.setInsertionPoint(insertPt); + createSyncBlockSetOp(builder, toMemrefOp->getLoc(), + hivm::TCoreType::CUBE, hivm::PIPE::PIPE_M, + hivm::PIPE::PIPE_MTE3, newflag); + } + } + } +} + +/// 同步点增强 +void addSyncOpsForBufferWait(ModuleOp module) { + for (auto funcOp : + llvm::make_early_inc_range(module.getOps())) { + if (funcOp.getBody().empty()) { + continue; + } + + Region *aicRegion = nullptr; + Region *aivRegion = nullptr; + + funcOp.walk([&](scope::ScopeOp scopeOp) { + auto coreTypeAttr = scopeOp->getAttrOfType( + hivm::TCoreTypeAttr::name); + if (!coreTypeAttr) + return; + + if (coreTypeAttr.getTcoretype() == hivm::TCoreType::CUBE) { + aicRegion = &scopeOp.getRegion(); + } + if (coreTypeAttr.getTcoretype() == hivm::TCoreType::VECTOR) { + aivRegion = &scopeOp.getRegion(); + } + }); + + if (!aicRegion || !aivRegion) { + continue; + } + + processFixpipeOpsInAIC(aicRegion, aivRegion); + processToMemrefOpsInAIV(aivRegion, aicRegion); + } +} + +void DAGScopePass::runOnOperation() { + auto module = getOperation(); + // llvm::outs()<())) { + // skip invalid function + if (funcOp.getBody().empty()) { + continue; + } + + // 收集所有 memref.alloc 操作 + llvm::SmallVector allocOps; + + // 遍历函数中的所有操作（包括嵌套区域中的操作） + funcOp.walk([&](mlir::Operation *op) { + if (mlir::isa(op)) { + allocOps.push_back(op); + } + }); + + mlir::Block &entryBlock = funcOp.getBody().front(); + mlir::Block::iterator insertPos = entryBlock.begin(); + + // 将 alloc 操作移动到函数的最前面 + for (mlir::Operation *allocOp : allocOps) { + // 如果 alloc 操作已经是最前面的操作，跳过 + if (allocOp->getBlock() == &entryBlock && + allocOp->isBeforeInBlock(&*insertPos)) { + continue; + } + + // 将 alloc 操作移动到指定位置 + allocOp->moveBefore(&entryBlock, insertPos); + } + + auto funcName = funcOp.getName(); + auto *graph_ptr = + AffinityDAG::GraphManager::getInstance().getGraph(funcName); + if (!graph_ptr) { + continue; + } + auto &main_graph = *graph_ptr; + + auto ScopeList = encapsulateWithScope(funcOp); + auto aivScope = ScopeList.first; // 第一个元素 + auto aicScope = ScopeList.second; // 第二个元素 + + SplitScope(funcOp, main_graph, aivScope, aicScope, module); + } + + addSyncOpsForBufferWait(module); + // llvm::outs()<> mlir::triton::createDAGScopePass() { + return std::make_unique(); +} diff --git a/third_party/ascend/lib/TritonAffinityOpt/DAGSync.cpp b/third_party/ascend/lib/TritonAffinityOpt/DAGSync.cpp new file mode 100644 index 0000000000..0a1ca83ff0 --- /dev/null +++ b/third_party/ascend/lib/TritonAffinityOpt/DAGSync.cpp @@ -0,0 +1,1635 @@ +#include "TritonAffinityOpt/Passes.h" + +#include "TritonAffinityOpt/Utils.hpp" +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/HIVM/IR/HIVMImpl.h" +#include "bishengir/Dialect/HIVM/IR/HIVMInterfaces.h" +#include "bishengir/Dialect/HIVM/Transforms/Passes.h" +#include "bishengir/Dialect/HIVM/Utils/Utils.h" +#include "bishengir/Dialect/Scope/IR/Scope.h" +#include "mlir/Analysis/AliasAnalysis.h" +#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h" +#include "mlir/Analysis/DataFlowFramework.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Dominance.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/Value.h" +#include "mlir/Interfaces/LoopLikeInterface.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallBitVector.h" +#include "llvm/Support/Casting.h" + +#include "incubated/Conversion/UtilsIncubated/Utils.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/FormatVariadic.h" +#include +#include +#include + +#include "TritonAffinityOpt/DAG.h" +#include "triton/Dialect/Triton/IR/Types.h" + +namespace mlir { +namespace triton { +#define GEN_PASS_DEF_DAGSYNC +#include "ascend/include/TritonAffinityOpt/Passes.h.inc" +} // namespace triton +} // namespace mlir + +#define DEBUG_TYPE "triton-affinity-dag-sync" + +// 使用 DAG 命名空间 +using namespace mlir; +using namespace hivm; +using namespace AffinityDAG; + +constexpr size_t MAX_FLAG_ID = 14; +llvm::DenseMap *valueTypes; +// 修改类声明，将数据搬运逻辑集成到同步插入中 +namespace { +struct DAGSyncPass : public mlir::triton::impl::DAGSyncBase { + void runOnOperation() override; + +private: + // 原有的辅助函数 + CoreType getNodeDeviceType(OpNode *node, + llvm::DenseMap *valueTypes); + bool needVectorCubeSync(CoreType src, CoreType dst); + + // 修改后的同步插入函数，包含数据搬运 + void insertSyncAndMovement(mlir::Operation *srcOp, mlir::Operation *dstOp, + CoreType srcType, CoreType dstType, + mlir::OpBuilder &builder, int flag, + llvm::DenseMap *valueMap, + Graph &mainGraph); + + // 新增：处理跨 block 的同步和数据搬运 + void insertSyncAndMovementForCrossBlock( + mlir::Operation *srcOp, mlir::Operation *dstOp, CoreType srcType, + CoreType dstType, mlir::OpBuilder &builder, int flag, + bool dstIsInnerBlock, llvm::DenseMap *valueMap, + Graph &mainGraph); + + // 新增：处理 scf.for 循环迭代参数的同步 + void processScfForSync(mlir::scf::ForOp forOp, Node *forNode, + llvm::DenseMap *valueTypes, + mlir::OpBuilder &builder, int &flag); + + // 数据搬运相关的辅助函数 + void insertCubeToVectorDataMovement(mlir::Operation *srcOp, + mlir::Operation *dstOp, + mlir::Value srcResult, + mlir::OpBuilder &builder, + mlir::Location loc, mlir::Value iterArgs); + + void + insertVectorToCubeDataMovement(mlir::Operation *srcOp, mlir::Operation *dstOp, + Operation *posOp, mlir::Value srcResult, + mlir::OpBuilder &builder, mlir::Location loc, + llvm::DenseMap *valueMap); + + // 获取或创建合适的 memref.alloc + mlir::Value getOrCreateAllocation(mlir::Operation *op, mlir::Type tensorType, + hivm::AddressSpace addressSpace, + mlir::OpBuilder &builder, + mlir::Location loc); + + // 获取 tensor 的形状和元素类型 + mlir::RankedTensorType getTensorType(mlir::Value tensorValue); + + // 替换 dstOp 中使用 srcResult 的操作数 + void replaceOperandWithNewValue(mlir::Operation *dstOp, mlir::Value oldValue, + mlir::Value newValue); + + // Find sync position + Operation *FindLastestPosition(Operation *srcOp, Graph &mainGraph, + OpBuilder &builder); + Operation *FindEarliestPosition(Operation *dstOp, Graph &mainGraph, + OpBuilder &builder); +}; +} // namespace + +void DAGSyncPass::processScfForSync( + mlir::scf::ForOp forOp, Node *forNode, + llvm::DenseMap *valueTypes, mlir::OpBuilder &builder, + int &flag) { + + mlir::Block *loopBody = forOp.getBody(); + mlir::scf::YieldOp yieldOp = nullptr; + for (mlir::Operation &op : *loopBody) { + if (auto yield = mlir::dyn_cast(&op)) { + yieldOp = yield; + break; + } + } + Location loc = forOp.getLoc(); + + for (int i = 0; i < forOp.getInitArgs().size(); i++) { + mlir::BlockArgument iterArg = loopBody->getArgument(i + 1); + // 找到首次使用 + mlir::Operation *firstUser = nullptr; + + for (mlir::Operation &op : *loopBody) { + // 跳过 yield 操作 + if (mlir::isa(&op)) { + continue; + } + + // 检查是否使用该迭代参数 + bool usesIterArg = false; + for (mlir::Value operand : op.getOperands()) { + if (operand == iterArg) { + usesIterArg = true; + break; + } + } + + if (usesIterArg) { + firstUser = &op; + break; + } + } + // map 内找到对应的iterType，iterType由首次在loop内使用到的op定义 + if (!firstUser) { + continue; + } + CoreType iterType = CoreType::CUBE_AND_VECTOR; + if (valueTypes->find(firstUser->getResult(0)) != valueTypes->end()) { + iterType = valueTypes->find(firstUser->getResult(0))->second; + } + + // 获取对应yield + mlir::Value yieldOperand = yieldOp->getOperand(i); + CoreType yieldType = CoreType::CUBE_AND_VECTOR; + if (valueTypes->find(yieldOperand) != valueTypes->end()) { + yieldType = valueTypes->find(yieldOperand)->second; + } + mlir::Operation *yieldDefiningOp = yieldOperand.getDefiningOp(); + + if (yieldType == CoreType::CUBE_ONLY && iterType == CoreType::VECTOR_ONLY) { + + // 2. 插入同步指令 + auto coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + auto setPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_FIX); + auto waitPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_V); + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + + // set 在 yieldDefiningOp 后 + builder.setInsertionPointAfter(yieldDefiningOp); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + + mlir::Value srcResult = yieldDefiningOp->getResult(0); + + // // 1. 插入数据搬运 + insertCubeToVectorDataMovement(yieldDefiningOp, firstUser, srcResult, + builder, loc, iterArg); + + // wait 在 firstUser 前 + builder.setInsertionPoint(firstUser); + coreAttr = hivm::TCoreTypeAttr::get(builder.getContext(), + hivm::TCoreType::VECTOR); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + // llvm::outs() << "yieldOp" << yieldDefiningOp << "iterargs" << firstUser + // << "\n"; llvm::outs() << "Inserted CUBE->VECTOR sync and data movement + // (flag=" << flag << ")\n"; + } + // VECTOR -> CUBE + else if (yieldType == CoreType::VECTOR_ONLY && + iterType == CoreType::CUBE_ONLY) { + + // 2. 插入同步指令 + auto coreAttr = hivm::TCoreTypeAttr::get(builder.getContext(), + hivm::TCoreType::VECTOR); + auto setPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_MTE3); + auto waitPipe = + PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_MTE1); + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + + // set 在 yieldDefiningOp 后 + builder.setInsertionPointAfter(yieldDefiningOp); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + + // 1. 插入数据搬运 + // insertVectorToCubeDataMovement(yieldDefiningOp, firstUser, srcResult, + // builder, loc, iterArg); + + // wait 在 firstUser 前 + builder.setInsertionPoint(firstUser); + coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + // llvm::outs() << "yieldOp" << yieldDefiningOp << "iterargs" << firstUser + // << "\n"; llvm::outs() << "Inserted VECTOR->CUBE sync and data movement + // (flag=" << flag << ")\n"; + } + } +} + +// 获取节点的设备类型 +CoreType DAGSyncPass::getNodeDeviceType( + OpNode *node, llvm::DenseMap *valueTypes) { + if (!node || !node->op) { + return CoreType::CUBE_AND_VECTOR; + } + + // 尝试从节点的结果中获取设备类型 + // 通常使用第一个结果来代表节点的设备类型 + if (node->op->getNumResults() > 0) { + mlir::Value result = node->op->getResult(0); + auto it = valueTypes->find(result); + if (it != valueTypes->end()) { + return it->second; + } + } + + // 如果没有找到，检查操作数 + // for (mlir::Value operand : node->op->getOperands()) { + // auto it = valueTypes->find(operand); + // if (it != valueTypes->end()) { + // return it->second; + // } + // } + + return CoreType::CUBE_AND_VECTOR; // 默认 +} + +// 判断是否需要vector<->cube同步 +bool DAGSyncPass::needVectorCubeSync(CoreType src, CoreType dst) { + return (src == CoreType::VECTOR_ONLY && dst == CoreType::CUBE_ONLY) || + (src == CoreType::CUBE_ONLY && dst == CoreType::VECTOR_ONLY); +} + +// 获取 tensor 类型 +mlir::RankedTensorType DAGSyncPass::getTensorType(mlir::Value tensorValue) { + if (auto tensorType = + dyn_cast(tensorValue.getType())) { + return tensorType; + } + return nullptr; +} + +// 替换操作数 +void DAGSyncPass::replaceOperandWithNewValue(mlir::Operation *dstOp, + mlir::Value oldValue, + mlir::Value newValue) { + for (unsigned i = 0; i < dstOp->getNumOperands(); ++i) { + if (dstOp->getOperand(i) == oldValue) { + dstOp->setOperand(i, newValue); + // llvm::outs() << "Replaced operand " << i << " of " << + // dstOp->getName().getStringRef() + // << " with new value\n"; + } + } +} + +// 修改 getOrCreateAllocation 函数，将 alloc 提到函数最外层 +mlir::Value DAGSyncPass::getOrCreateAllocation(mlir::Operation *op, + mlir::Type tensorType, + hivm::AddressSpace addressSpace, + mlir::OpBuilder &builder, + mlir::Location loc) { + auto rankedTensorType = cast(tensorType); + auto elementType = rankedTensorType.getElementType(); + auto shape = rankedTensorType.getShape(); + + auto addressSpaceAttr = + hivm::AddressSpaceAttr::get(builder.getContext(), addressSpace); + auto memrefType = mlir::MemRefType::get(shape, elementType, + /*layout=*/nullptr, addressSpaceAttr); + + // 查找是否已经存在相同类型的 allocation（在函数的 entry block 中） + mlir::Operation *funcOp = op; + while (funcOp && !mlir::isa(funcOp)) { + funcOp = funcOp->getParentOp(); + } + + if (auto func = mlir::dyn_cast(funcOp)) { + // 在函数的 entry block 中查找现有的 allocation + mlir::Block &entryBlock = func.getBody().front(); + // for (auto& blockOp : entryBlock) { + // if (auto allocOp = mlir::dyn_cast(&blockOp)) { + // if (allocOp.getType() == memrefType) { + // // 找到匹配的 allocation，直接复用 + // llvm::outs() << "Reusing existing allocation: " << allocOp << + // "\n"; return allocOp.getResult(); + // } + // } + // } + + // 没有找到现有的 allocation，在函数开头创建新的 + builder.setInsertionPointToStart(&entryBlock); + return builder.create(loc, memrefType); + } + + // 如果没有找到函数，回退到原逻辑 + builder.setInsertionPoint(op); + return builder.create(loc, memrefType); +} + +// 插入 CUBE -> VECTOR 数据搬运 +void DAGSyncPass::insertCubeToVectorDataMovement( + mlir::Operation *srcOp, mlir::Operation *dstOp, mlir::Value srcResult, + mlir::OpBuilder &builder, mlir::Location loc, mlir::Value iterArgs) { + auto srcTensorType = getTensorType(srcResult); + if (!srcTensorType) { + return; + } + + // 1. 在 srcOp 之后创建 UB 空间的 memref.alloc + builder.setInsertionPointAfter(srcOp); + mlir::Value ubAlloc = getOrCreateAllocation( + srcOp, srcTensorType, hivm::AddressSpace::UB, builder, loc); + + // 2. 创建 fixpipe 指令 + builder.setInsertionPointAfter(srcOp); + FixpipeDMAModeAttr dmaModeAttr = + FixpipeDMAModeAttr::get(builder.getContext(), FixpipeDMAMode::NZ2ND); + + auto fixpipeOp = + builder.create(loc, mlir::TypeRange{}, // 没有返回值 + srcResult, // src + ubAlloc, // dst +#ifndef __LLVM_MAJOR_VERSION_22_COMPATIBLE__ + /*unit_flag_cond=*/mlir::ValueRange{}, + /*dma_mode=*/dmaModeAttr, + /*dual_dst_mode=*/nullptr, + /*pre_quant=*/nullptr, + /*pre_relu=*/nullptr, + /*channel_split=*/nullptr, + /*unit_flag_mode=*/mlir::ArrayAttr{}); +#else + // All parameters are explicit in this + // overload; sub_block_idx and c0_pad_en + // were added and quant_scale moved ahead + // of unit_flag_mode. + /*unit_flag_cond=*/mlir::ValueRange{}, + /*dma_mode=*/dmaModeAttr, + /*dual_dst_mode=*/nullptr, + /*sub_block_idx=*/nullptr, + /*pre_quant=*/nullptr, + /*pre_relu=*/nullptr, + /*channel_split=*/nullptr, + /*c0_pad_en=*/nullptr, + /*quant_scale=*/mlir::Value{}, + /*unit_flag_mode=*/mlir::ArrayAttr{}, + /*unit_flag_group_id=*/ + mlir::IntegerAttr{}); +#endif + + llvm::outs() << "Inserted fixpipe after " << srcOp->getName().getStringRef() + << " for CUBE->VECTOR data movement\n"; + + // 3. 在 dstOp 前创建 memory_space_cast 和 to_tensor + builder.setInsertionPoint(dstOp); + + // memory_space_cast（如果需要） + mlir::Value plainMemref = ubAlloc; + auto memrefType = cast(ubAlloc.getType()); + if (memrefType.getMemorySpace()) { + auto plainMemrefType = mlir::MemRefType::get(memrefType.getShape(), + memrefType.getElementType()); + plainMemref = builder.create( + loc, plainMemrefType, ubAlloc); + (*valueTypes)[plainMemref] = CoreType::VECTOR_ONLY; + } + + // 4. 创建 to_tensor + auto toTensorOp = builder.create( + loc, + srcTensorType, // 原始的 tensor 类型 + plainMemref, + /*restrict=*/true, + /*writable=*/true); + (*valueTypes)[toTensorOp.getResult()] = CoreType::VECTOR_ONLY; + + // 5. 替换 dstOp 的操作数 + if (!iterArgs) { + replaceOperandWithNewValue(dstOp, srcResult, toTensorOp.getResult()); + } else { + replaceOperandWithNewValue(dstOp, iterArgs, toTensorOp.getResult()); + } +} + +static uint64_t getElemBytesForAlign(Type t) { + if (auto ft = dyn_cast(t)) + return (uint64_t)((ft.getWidth() + 7) / 8); + if (auto it = dyn_cast(t)) + return (uint64_t)((it.getWidth() + 7) / 8); + if (isa(t)) + return 8ULL; + if (auto ct = dyn_cast(t)) + return 2ULL * getElemBytesForAlign(ct.getElementType()); + return 0ULL; +} + +static FailureOr getBlockElemsFor32BAlign(Type elemType) { + constexpr uint64_t kAlignBytes = 32; + uint64_t elemBytes = getElemBytesForAlign(elemType); + if (elemBytes <= 0) + return failure(); + if (elemBytes >= kAlignBytes) + return 1; + if (kAlignBytes % elemBytes != 0) + return failure(); + return kAlignBytes / elemBytes; +} + +static std::optional> +newCbubAllocShape(memref::AllocOp allocOp) { + auto type = dyn_cast(allocOp.getType()); + // 仅支持静态 2D MemRef + if (!type || type.getRank() != 2) + return std::nullopt; + + auto shape = type.getShape(); + int64_t M = shape[0]; + int64_t N = shape[1]; + auto elemType = type.getElementType(); + auto blkOr = getBlockElemsFor32BAlign(elemType); + int64_t blk = (int64_t)*blkOr; + // 必须是静态且 16 对齐 + if (ShapedType::isDynamic(M) || ShapedType::isDynamic(N)) + return std::nullopt; + if (M % 16 != 0) + return std::nullopt; + + // 新 shape: (N/16, M/16, 16, 16) + SmallVector newShape = {N / blk, M / 16, 16, blk}; + + return newShape; +} + +// 修改 VECTOR->CUBE 数据搬运函数 +void DAGSyncPass::insertVectorToCubeDataMovement( + mlir::Operation *srcOp, mlir::Operation *dstOp, Operation *posOp, + mlir::Value srcResult, mlir::OpBuilder &builder, mlir::Location loc, + llvm::DenseMap *valueMap) { + auto srcTensorType = getTensorType(srcResult); + if (!srcTensorType) { + return; + } + if (isa(srcOp) && isa(dstOp)) { + return; + } + + // 1. 在 srcOp 之后创建 UB 空间的 memref.alloc（用于 to_memref） + builder.setInsertionPointAfter(srcOp); + + // 首先创建 UB 空间的 memref type + auto ubSpaceAttr = + hivm::AddressSpaceAttr::get(builder.getContext(), hivm::AddressSpace::UB); + auto ubMemrefType = mlir::MemRefType::get(srcTensorType.getShape(), + srcTensorType.getElementType(), + /*layout=*/nullptr, ubSpaceAttr); + + // 创建 bufferization.to_memref + if (srcOp->getBlock() == dstOp->getBlock()) { + builder.setInsertionPoint(posOp); + } + auto toBufferOp = + builder.create(loc, ubMemrefType, srcResult); + + // 2. 创建 CBUF 空间的 memref.alloc（用于 copy 的目标） + mlir::Value cbufAllocOld = getOrCreateAllocation( + srcOp, srcTensorType, hivm::AddressSpace::L1, builder, loc); + auto cbufShape = *newCbubAllocShape( + dyn_cast(cbufAllocOld.getDefiningOp())); + // 获取旧的memref类型并创建新的类型 + auto oldType = dyn_cast(cbufAllocOld.getType()); + + // 获取新的维度数量 + unsigned newRank = cbufShape.size(); + + // 方法1：创建新的恒等布局映射 + AffineMap identityMap = builder.getMultiDimIdentityMap(newRank); + MemRefLayoutAttrInterface layout = AffineMapAttr::get(identityMap); + + // 方法2：如果旧类型有布局，尝试调整它（更安全的选择） + // 先检查旧类型是否有布局 + if (auto oldLayout = oldType.getLayout()) { + if (auto affineMapAttr = dyn_cast(oldLayout)) { + // 如果旧布局是AffineMap，尝试创建新的恒等映射 + // 因为维度改变，旧的affine map可能不再有效 + layout = AffineMapAttr::get(identityMap); + } else { + // 对于其他类型的布局，可能需要特殊处理 + layout = oldLayout; + } + } + + // 创建新的alloc类型 + auto newAllocType = MemRefType::get(cbufShape, oldType.getElementType(), + layout, // 使用新创建的布局 + oldType.getMemorySpace()); + + builder.setInsertionPoint(cbufAllocOld.getDefiningOp()); + // 创建新的alloc操作 + auto cbufAlloc = builder.create( + cbufAllocOld.getDefiningOp()->getLoc(), newAllocType); + + builder.setInsertionPointAfter(toBufferOp); + // 3. 创建 copy 指令（src 是 ub memref，dst 是 cbuf memref） + auto copyOp = + builder.create(loc, mlir::TypeRange{}, // 没有返回值 + toBufferOp.getResult(), // src (memref in UB) + cbufAlloc // dst (memref in CBUF) + ); + + // llvm::outs() << "Inserted copy after " << srcOp->getName().getStringRef() + // << " for VECTOR->CUBE data movement\n"; + + // 4. 在 dstOp 前创建 convert_layout + builder.setInsertionPoint(dstOp); + auto ndLayout = + hivm::DataLayoutAttr::get(builder.getContext(), hivm::DataLayout::ND); + // 创建 convert_layout + auto convertLayoutOp = builder.create( + loc, + cbufAllocOld.getType(), // 输出类型与输入相同 + cbufAlloc, + ndLayout, // srcLayout + ndLayout // dstLayout + ); + (*valueTypes)[convertLayoutOp.getResult()] = CoreType::CUBE_ONLY; + + // 5. 创建 memory_space_cast + auto cbufMemrefType = cast(convertLayoutOp.getType()); + auto plainMemrefType = mlir::MemRefType::get(cbufMemrefType.getShape(), + cbufMemrefType.getElementType()); + + auto memspaceCastOp = builder.create( + loc, plainMemrefType, convertLayoutOp.getResult()); + (*valueTypes)[memspaceCastOp.getResult()] = CoreType::CUBE_ONLY; + + // 6. 创建 to_tensor + auto toTensorOp = builder.create( + loc, + srcTensorType, // 原始的 tensor 类型 + memspaceCastOp.getResult(), + /*restrict=*/true, + /*writable=*/true); + (*valueTypes)[toTensorOp.getResult()] = CoreType::CUBE_ONLY; + + // 7. 替换 dstOp 的操作数 + replaceOperandWithNewValue(dstOp, srcResult, toTensorOp.getResult()); +} + +Operation *DAGSyncPass::FindLastestPosition(Operation *srcOp, Graph &mainGraph, + OpBuilder &builder) { + Operation *insertPos = nullptr; + auto opMap = mainGraph.getOpMapLegacy(); + auto valueTypes = &mainGraph.getValueTypes(); + // Find the first cube-dependent vector core operation. + for (auto nextOp = srcOp->getNextNode(); nextOp != nullptr; + nextOp = nextOp->getNextNode()) { + auto nextType = getNodeDeviceType(opMap[nextOp], valueTypes); + if (nextType == CoreType::CUBE_ONLY) + continue; + // No memref ops in IR yet; directly tracing operands + for (auto operand : nextOp->getOperands()) { + auto defOp = operand.getDefiningOp(); + auto defType = getNodeDeviceType(opMap[defOp], valueTypes); + if (defType == CoreType::CUBE_ONLY) { + // To prevent UB overflow, we need to break the dependency at the point + // where the result shape is minimized + // — i.e., trace upward to find the first broadcast. + for (auto prevOp = nextOp->getPrevNode(); + prevOp != nullptr && prevOp != srcOp; + prevOp = prevOp->getPrevNode()) { + if (isa(prevOp)) { + if (prevOp->getPrevNode() && + isa(prevOp->getPrevNode())) { + return prevOp->getPrevNode(); + } + return prevOp; + } + } + // Can't find the result shape is minimized + return nextOp; + } + } + + // Once meet SyncBlockWaitOp, return now! + if (auto waitOp = dyn_cast(nextOp)) { + if (waitOp.getTcoreType() == + hivm::TCoreTypeAttr::get(builder.getContext(), + hivm::TCoreType::VECTOR)) { + return nextOp; + } + } + insertPos = nextOp; + } + return insertPos; +} + +Operation *DAGSyncPass::FindEarliestPosition(Operation *dstOp, Graph &mainGraph, + OpBuilder &builder) { + auto insertPos = dstOp; + auto opMap = mainGraph.getOpMapLegacy(); + auto valueTypes = &mainGraph.getValueTypes(); + for (auto prevOp = dstOp->getPrevNode(); prevOp != nullptr; + prevOp = prevOp->getPrevNode()) { + if (dstOp->getBlock() != prevOp->getBlock()) + continue; + // Once meet SyncBlockSetOp, return now! + if (auto waitOp = dyn_cast(prevOp)) { + if (waitOp.getTcoreType() == + hivm::TCoreTypeAttr::get(builder.getContext(), + hivm::TCoreType::VECTOR)) { + return insertPos; + } + } + insertPos = prevOp; + } + return insertPos; +} + +// 主要的同步和数据搬运插入函数 +void DAGSyncPass::insertSyncAndMovement( + mlir::Operation *srcOp, mlir::Operation *dstOp, CoreType srcType, + CoreType dstType, mlir::OpBuilder &builder, int flag, + llvm::DenseMap *valueMap, Graph &mainGraph) { + mlir::Location loc = srcOp->getLoc(); + // 保存当前的插入点 + mlir::OpBuilder::InsertionGuard guard(builder); + + // 检查是否是跨 block + mlir::Block *srcBlock = srcOp->getBlock(); + mlir::Block *dstBlock = dstOp->getBlock(); + bool sameBlock = (srcBlock == dstBlock); + + if (!sameBlock) { + // 检查是否是外层到内层的依赖 + bool dstIsInnerBlock = false; + mlir::Operation *dstParentOp = dstBlock->getParentOp(); + while (dstParentOp) { + if (dstParentOp->getBlock() == srcBlock) { + dstIsInnerBlock = true; + break; + } + if (dstParentOp->getBlock()) { + dstParentOp = dstParentOp->getBlock()->getParentOp(); + } else { + break; + } + } + + if (dstIsInnerBlock) { + insertSyncAndMovementForCrossBlock(srcOp, dstOp, srcType, dstType, + builder, flag, true, valueMap, + mainGraph); + return; + } + } + + // 同一 block 内的处理 + // 获取 srcOp 的输出（假设第一个结果） + if (srcOp->getNumResults() == 0) { + return; + } + mlir::Value srcResult = srcOp->getResult(0); + + // CUBE -> VECTOR + if (srcType == CoreType::CUBE_ONLY && dstType == CoreType::VECTOR_ONLY) { + + // 2. 插入同步指令 + auto coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + auto setPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_FIX); + auto waitPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_V); + auto lastSetPipe = + PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_MTE3); + auto lastWaitPipe = + PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_MTE1); + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + auto flagAddId = builder.getIntegerAttr(builder.getI64Type(), flag * 2); + auto lastFlagAddId = + builder.getIntegerAttr(builder.getI64Type(), (flag - 1) * 2); + + // set 在 srcOp 后 + builder.setInsertionPointAfter(srcOp); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + + // wait 在 dstOp 前 + + auto posOp = FindEarliestPosition(dstOp, mainGraph, builder); + builder.setInsertionPoint(posOp); + coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::VECTOR); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + + // 1. 插入数据搬运 + insertCubeToVectorDataMovement(srcOp, dstOp, srcResult, builder, loc, + nullptr); + + // llvm::outs() << "Inserted CUBE->VECTOR sync and data movement (flag=" << + // flag << ")\n"; + } + // VECTOR -> CUBE + else if (srcType == CoreType::VECTOR_ONLY && dstType == CoreType::CUBE_ONLY) { + + // 2. 插入同步指令 + auto coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::VECTOR); + auto setPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_MTE3); + auto waitPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_MTE1); + auto lastSetPipe = + PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_FIX); + auto lastWaitPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_V); + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + auto flagAddId = builder.getIntegerAttr(builder.getI64Type(), flag * 2); + auto lastFlagAddId = + builder.getIntegerAttr(builder.getI64Type(), (flag - 1) * 2); + + // set 在 srcOp 后 + // builder.setInsertionPointAfter(srcOp); + auto posOp = FindLastestPosition(srcOp, mainGraph, builder); + if (posOp) { + builder.setInsertionPoint(posOp); + } else { + builder.setInsertionPointAfter(srcOp); + } + auto setOp = builder.create(loc, coreAttr, setPipe, + waitPipe, flagId); + + // wait 在 dstOp 前 + builder.setInsertionPoint(dstOp); + coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + + // 1. 插入数据搬运 + insertVectorToCubeDataMovement(srcOp, dstOp, setOp, srcResult, builder, loc, + valueMap); + + // llvm::outs() << "Inserted VECTOR->CUBE sync and data movement (flag=" << + // flag << ")\n"; + } +} + +// 跨 block 的同步和数据搬运 +void DAGSyncPass::insertSyncAndMovementForCrossBlock( + mlir::Operation *srcOp, mlir::Operation *dstOp, CoreType srcType, + CoreType dstType, mlir::OpBuilder &builder, int flag, bool dstIsInnerBlock, + llvm::DenseMap *valueMap, Graph &mainGraph) { + if (!dstIsInnerBlock) { + insertSyncAndMovement(srcOp, dstOp, srcType, dstType, builder, flag, + valueMap, mainGraph); + return; + } + + mlir::Location loc = srcOp->getLoc(); + mlir::Block *dstBlock = dstOp->getBlock(); + + // 获取 srcOp 的输出 + if (srcOp->getNumResults() == 0) { + return; + } + mlir::Value srcResult = srcOp->getResult(0); + + // CUBE -> VECTOR + if (srcType == CoreType::CUBE_ONLY && dstType == CoreType::VECTOR_ONLY) { + + // 2. 插入同步指令（跨 block 特殊处理） + auto coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + auto setPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_FIX); + auto waitPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_V); + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + + // set 在 srcOp 后（外层） + builder.setInsertionPointAfter(srcOp); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + + // 1. 插入数据搬运（同 block 内逻辑） + insertCubeToVectorDataMovement(srcOp, dstOp, srcResult, builder, loc, + nullptr); + + // wait 在内层 block 入口前 + mlir::Operation *parentOp = dstBlock->getParentOp(); + while (srcOp->getBlock() != parentOp->getBlock()) { + parentOp = parentOp->getBlock()->getParentOp(); + } + if (parentOp) { + builder.setInsertionPoint(parentOp); + coreAttr = hivm::TCoreTypeAttr::get(builder.getContext(), + hivm::TCoreType::VECTOR); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + } else { + builder.setInsertionPoint(dstOp); + coreAttr = hivm::TCoreTypeAttr::get(builder.getContext(), + hivm::TCoreType::VECTOR); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + } + + } + // VECTOR -> CUBE + else if (srcType == CoreType::VECTOR_ONLY && dstType == CoreType::CUBE_ONLY) { + + // 2. 插入同步指令（跨 block 特殊处理） + auto coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::VECTOR); + auto setPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_MTE3); + auto waitPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_MTE1); + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + + // set 在 srcOp 后（外层） + builder.setInsertionPointAfter(srcOp); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + + // 1. 插入数据搬运（同 block 内逻辑） + insertVectorToCubeDataMovement(srcOp, dstOp, srcOp, srcResult, builder, loc, + valueMap); + + // wait 在内层 block 入口前 + mlir::Operation *parentOp = dstBlock->getParentOp(); + if (parentOp) { + builder.setInsertionPoint(parentOp); + coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + } else { + builder.setInsertionPoint(dstOp); + coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + } + + // llvm::outs() << "Inserted cross-block VECTOR->CUBE sync and data movement + // (flag=" << flag << ")\n"; + } +} + +void LegalizeDot(triton::FuncOp funcOp) { + mlir::OpBuilder builder(funcOp); + funcOp.walk([&](triton::DotOp dotOp) { + // 获取dot操作的输入 + Value a = dotOp.getOperands()[0]; + Value b = dotOp.getOperands()[1]; + Value c = dotOp.getOperands()[2]; // 累加器参数 + + // 检查累加器是否为全零常量 + bool isZeroAccumulator = false; + + // 检查是否直接是arith.constant 0 + if (auto constantOp = c.getDefiningOp()) { + if (auto denseAttr = dyn_cast(constantOp.getValue())) { + if (denseAttr.isSplat() && + denseAttr.getSplatValue().getValueAsDouble() == 0.0) { + isZeroAccumulator = true; + } + } + } + + if (!isZeroAccumulator) { + // 创建新的零累加器 + Location loc = dotOp.getLoc(); + auto resultType = dotOp.getResult().getType(); + + Value originalResult = dotOp.getResult(); + builder.setInsertionPoint(dotOp); + // 创建全零张量 + auto zeroAttr = DenseElementsAttr::get( + dyn_cast(resultType), APFloat(0.0f)); + auto zeroConstant = builder.create(loc, zeroAttr); + + // 创建新的dot操作，使用零作为累加器 + auto newDot = + builder.create(loc, resultType, a, b, zeroConstant); + + // 创建加法操作，将新的dot结果与原来的累加器c相加 + auto addOp = builder.create(loc, newDot, c); + + // 用addOp替换原来的dotOp + originalResult.replaceAllUsesWith(addOp.getResult()); + + // 删除原dotOp（如果它没有其他用途） + if (dotOp.use_empty()) { + dotOp.erase(); + } + } + }); +} + +static void rewriteCopyChainForCbub(hivm::CopyOp copyOp, + ArrayRef newShape, + OpBuilder &builder) { + + // 获取 copy 的输入（ins），应为 to_memref 的结果 + Value insVal = copyOp.getOperands()[0]; + auto toBufferOp = insVal.getDefiningOp(); + if (!toBufferOp) + return; + + Value inputTensor = toBufferOp.getTensor(); + auto inputTensorType = dyn_cast(inputTensor.getType()); + if (!inputTensorType || inputTensorType.getRank() != 2) + return; + + // blk = 32/位宽 + // 中间 reshape 形状：[M/16, 16, N/ blk, blk] + int64_t M = inputTensorType.getShape()[0]; + int64_t N = inputTensorType.getShape()[1]; + auto elemType = inputTensorType.getElementType(); + auto blkOr = getBlockElemsFor32BAlign(elemType); + int64_t blk = (int64_t)*blkOr; + SmallVector intermediateShape3D = {M, N / blk, blk}; + SmallVector intermediateShapetrans = {N / blk, M, blk}; + auto elementType = inputTensorType.getElementType(); + auto interTensor3DType = + RankedTensorType::get(intermediateShape3D, elementType); + auto interTensortransType = + RankedTensorType::get(intermediateShapetrans, elementType); + + auto finalTensorType = RankedTensorType::get(newShape, elementType); + + auto loc = inputTensor.getLoc(); + + // Set insertion point before copyOp (or toBufferOp) + auto tensorOp = inputTensor.getDefiningOp(); + builder.setInsertionPointAfter(tensorOp); + + // 插入 triton.reshape 将 2D tensor 展开为 3D + auto reshape3DOp = + builder.create(loc, interTensor3DType, inputTensor); + (*valueTypes)[reshape3DOp.getResult()] = CoreType::VECTOR_ONLY; + + // nark tiling dim for reshapeop + auto markOp3d = builder.create(loc, reshape3DOp); + auto tilingDimAttr3d = builder.getDictionaryAttr(SmallVector{ + NamedAttribute(builder.getStringAttr("1"), builder.getIndexAttr(1))}); + markOp3d->setAttr("tiling_dim_mapping", tilingDimAttr3d); + + // 插入 triton.trans 调整维度顺序 Insert tt.trans {order = [1, 0, 2]} + SmallVector order = {1, 0, 2}; + auto orderAttr = + builder.getDenseI32ArrayAttr(order); // OpBuilder supports this + auto transOp = builder.create( + loc, interTensortransType, reshape3DOp.getResult(), orderAttr); + (*valueTypes)[transOp.getResult()] = CoreType::VECTOR_ONLY; + + // 插入 triton.reshape 将 3D tensor 展开为 4D + auto reshape4DOp = builder.create(loc, finalTensorType, + transOp.getResult()); + (*valueTypes)[reshape4DOp.getResult()] = CoreType::VECTOR_ONLY; + + // nark tiling dim for reshapeop + auto markOp4d = builder.create(loc, reshape4DOp); + auto tilingDimAttr4d = builder.getDictionaryAttr(SmallVector{ + NamedAttribute(builder.getStringAttr("1"), builder.getIndexAttr(1))}); + markOp4d->setAttr("tiling_dim_mapping", tilingDimAttr4d); + + // Create new to_memref + builder.setInsertionPoint(toBufferOp); + auto oldBufferType = cast(toBufferOp.getResult().getType()); + auto newMemRefType = MemRefType::get(newShape, elementType, mlir::AffineMap{}, + oldBufferType.getMemorySpace()); + auto newtoBufferOp = builder.create( + toBufferOp.getLoc(), newMemRefType, reshape4DOp.getResult()); + (*valueTypes)[newtoBufferOp.getResult()] = CoreType::VECTOR_ONLY; + + // Create NEW copyOp (replacing the old one) + builder.setInsertionPoint(copyOp); + auto resultTypes = copyOp->getResultTypes(); + auto newCopyOp = + builder.create(copyOp.getLoc(), + resultTypes, // TypeRange + reshape4DOp.getResult(), // src (ins) + copyOp.getOperands()[1] // dst (outs) + ); + + // 替换 uses 并清理旧 op + copyOp.replaceAllUsesWith(newCopyOp); + copyOp.erase(); + toBufferOp.erase(); + + return; +} + +bool valueIsPtrOrShapedPtr(Value value) { + auto type = value.getType(); + if (llvm::isa(type)) { + return true; + } + + if (auto tensorType = llvm::dyn_cast(type)) { + return llvm::isa(tensorType.getElementType()); + } + + return false; +} + +using AliasSet = llvm::SmallBitVector; +using AliasInfo = llvm::DenseMap; +constexpr size_t EXPECTED_MAX_ROOT_PTR_COUNT = 16; + +/** + * Reworked alias analysis inplace of triton::SharedMemoryAliasAnalysis or + * mlir::AliasAnalysis where both failed to establish aliases from function + * arguments (i.e. global memory): + * 1. triton::SharedMemoryAliasAnalysis - does not reach global memory + * 2. mlir::AliasAnalysis - marks all function arguments as MayAlias or + * PartialAlias; we may be able to get around by marking all arguments as + * noalias, but it may affect other analyses... + */ +AliasInfo getAlias(triton::FuncOp funcOp, Graph::ValueMapRaw &valueMap) { + AliasInfo aliasInfo; + + struct WorkNode { + Value value; + Value cause; + }; + + llvm::SmallVector worklist; + + auto addDownstreamToWorklist = [&worklist](ValueNode *node) { + if (node == nullptr) { + return; + } + for (auto *output : node->getOutputs()) { + if (auto *opNode = dyn_cast(output)) { + for (auto result : opNode->op->getResults()) { + if (valueIsPtrOrShapedPtr(result)) { + worklist.push_back({result, node->value}); + } + } + } else if (auto *valueNode = dyn_cast(output)) { + worklist.push_back({valueNode->value, node->value}); + } + } + }; + + SmallVector rootPtrs; + + for (auto arg : funcOp.getBody().getArguments()) { + if (valueIsPtrOrShapedPtr(arg)) { + rootPtrs.push_back(arg); + } + } + funcOp.walk([&](Operation *op) { + // If an op generates ptr from nowhere (e.g. tt.empty()), then we consider + // it as a root + bool generatesRootPtr = + !llvm::any_of(op->getOperands(), valueIsPtrOrShapedPtr); + if (generatesRootPtr) { + for (auto result : op->getResults()) { + if (valueIsPtrOrShapedPtr(result)) { + rootPtrs.push_back(result); + } + } + } + }); + + size_t numRootPtrs = rootPtrs.size(); + + for (auto [index, value] : llvm::enumerate(rootPtrs)) { + auto &bitset = aliasInfo[value]; + bitset.resize(numRootPtrs); + bitset[index] = true; + + addDownstreamToWorklist(getFromSmartPtr(valueMap, value)); + } + + while (!worklist.empty()) { + auto [value, upstream] = worklist.pop_back_val(); + if (!valueIsPtrOrShapedPtr(value)) { + continue; + } + auto node = getFromSmartPtr(valueMap, value); + if (!node) { + continue; + } + + /** + * Safety: + * 1. aliasInfo[value] : may insert into densemap and cause resize + * 2. aliasInfo[upstream] : steady reference without any insertion into the + * densemap, will not cause resize, since upstream is in the map + * 3. after aliasInfo[upstream] : currAlias is still valid + */ + auto &currAlias = aliasInfo[value]; + auto &causeAlias = aliasInfo[upstream]; + + if (!causeAlias.test(currAlias)) { + continue; + } + + currAlias |= causeAlias; + addDownstreamToWorklist(node); + } + return aliasInfo; +} + +/** + * @returns whether two values share any root ptr + */ +bool mayAlias(AliasInfo &aliasInfo, Value valueA, Value valueB) { + auto aliasesA = getPtr(aliasInfo, valueA); + auto aliasesB = getPtr(aliasInfo, valueB); + if (!valueIsPtrOrShapedPtr(valueA) || !valueIsPtrOrShapedPtr(valueB) || + aliasesA == nullptr || aliasesB == nullptr) { + return false; + } + + return aliasesA->anyCommon(*aliasesB); +} + +template +OpTy createBlockSync(OpBuilder builder, hivm::TCoreType coreType, + hivm::PIPE srcPipe, hivm::PIPE dstPipe, int flag, + Operation *cause) { + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + auto coreAttr = hivm::TCoreTypeAttr::get(builder.getContext(), coreType); + auto setPipe = PipeAttr::get(builder.getContext(), srcPipe); + auto waitPipe = PipeAttr::get(builder.getContext(), dstPipe); + return builder.create(cause->getLoc(), coreAttr, setPipe, waitPipe, + flagId); +} + +const size_t MAX_EXPECTED_PARENTS_COUNT = 8; +std::optional> +findAncestorCommonBlock(mlir::Operation *opA, mlir::Operation *opB) { + if (opA->getBlock() == opB->getBlock()) { + return std::make_pair(opA, opB); + } + + // record all ancestors of opA + llvm::SmallPtrSet ancestorsA; + mlir::Operation *curr = opA; + while (curr) { + ancestorsA.insert(curr); + curr = curr->getParentOp(); + } + + // find the last ancestor of opB which is also the ancestor of opA + mlir::Operation *commonAncOp = nullptr; + curr = opB; + while (curr) { + if (ancestorsA.count(curr)) { + commonAncOp = curr; + break; + } + curr = curr->getParentOp(); + } + + if (!commonAncOp) { + return std::nullopt; + } + + // find the ancestors in the given region + for (mlir::Region ®ion : commonAncOp->getRegions()) { + for (mlir::Block &block : region) { + auto *ancA = block.findAncestorOpInBlock(*opA); + auto *ancB = block.findAncestorOpInBlock(*opB); + if (ancA && ancB) { + return std::make_pair(ancA, ancB); + } + } + } + return std::nullopt; +} + +struct SyncCandidate { + CoreType srcCoreType; + Operation *setCause; + Operation *setAfter; + Operation *waitCause; + Operation *waitBefore; +}; + +// setOp, waitOp +std::pair +createBlockSyncBetween(OpBuilder builder, hivm::PIPE srcPipe, + hivm::PIPE dstPipe, SyncCandidate candidate, int flag) { + auto srcCoreType = toHivm(candidate.srcCoreType); + auto dstCoreType = toHivm(!candidate.srcCoreType); + + builder.setInsertionPointAfter(candidate.setAfter); + auto setOp = createBlockSync( + builder, srcCoreType, srcPipe, dstPipe, flag, candidate.setCause); + builder.setInsertionPoint(candidate.waitBefore); + auto waitOp = createBlockSync( + builder, dstCoreType, srcPipe, dstPipe, flag, candidate.waitCause); + return {setOp, waitOp}; +}; + +/** + * @returns the source/destination of a memory op + * + * Strangely, triton memory ops do not register their pointers to memory + * effects, so we need some special cases as fallback here + */ +Value getMemoryAddress(Operation *op, MemoryEffects::EffectInstance effect) { + if (mlir::Value v = effect.getValue()) + return v; + + return TypeSwitch(op) + .Case([](auto op) { return op.getPtr(); }) + .Default([](auto) { return nullptr; }); +} + +void addMemEffectsSync(triton::FuncOp funcOp, Graph *graph, OpBuilder &builder, + int &syncFlag) { + DominanceInfo domInfo(funcOp); + PostDominanceInfo postDomInfo(funcOp); + + // [(node, EffectInstance, LinearisationPt)] + llvm::SmallVector> memOps; + + // [(setAfter, waitBefore, srcOP, dstOp)][CoreType] + llvm::SmallVector candidates; + llvm::SmallVector backwardCandidates; + + auto aliasInfo = getAlias(funcOp, graph->getValueMap()); + + funcOp.walk([&](MemoryEffectOpInterface memIface) { + auto *op = memIface.getOperation(); + if (llvm::isa(op)) { + return; + } + + auto *currNode = graph->getOpMap()[op].get(); + SmallVector effects; + + memIface.getEffects(effects); + + for (auto &effect : effects) { + if (!isa(effect.getEffect())) { + continue; + } + auto currVal = getMemoryAddress(op, effect); + if (!currVal) { + op->emitWarning("Mem effect src/dst is unknown!"); + continue; + } + memOps.emplace_back(currNode, effect); + bool isWrite = isa(effect.getEffect()); + for (auto &[prevNode, prevEffect] : memOps) { + auto prevVal = getMemoryAddress(prevNode->op, prevEffect); + if ((isa(prevEffect.getEffect()) || isWrite) && + mayAlias(aliasInfo, prevVal, currVal) && + prevNode->isOn() != + currNode->isOn() // write is forced on single core type, so we + // are safe to judge based on whether the core + // types are different + ) { + CoreType srcCoreType = isWrite ? !currNode->isOn() : prevNode->isOn(); + CoreType dstCoreType = !srcCoreType; + auto opPair = findAncestorCommonBlock(prevNode->op, currNode->op); + if (!opPair.has_value()) { + op->emitWarning(llvm::formatv( + "Unable to find ancestors in common block with {0}\n", + *prevNode->op)); + continue; + } + auto [setAfter, waitBefore] = opPair.value(); + if (setAfter == waitBefore) { + continue; + } + candidates.push_back(SyncCandidate{srcCoreType, prevNode->op, + setAfter, op, waitBefore}); + + if (op->getParentOfType()) { // need to insert + // backward sync + backwardCandidates.push_back(SyncCandidate{ + dstCoreType, op, waitBefore, prevNode->op, setAfter}); + } + } + } + } + }); + + /** + * Adds block sync operations for the candidate + * Specially, for backward syncs, we also perform a kick-start for first wait + * and clean-up to avoid affecting possible future syncs + */ + auto addBlockSyncCommon = [&builder, &syncFlag, + &aliasInfo](SyncCandidate cand) { + LLVM_DEBUG(llvm::dbgs() << "\n\n=== Insert sync between ===\n" + << *cand.setAfter << "\n" + << "Cause: " << *cand.setCause << "\n"); + + LLVM_DEBUG(llvm::dbgs() << "----------" << "\n" + << *cand.waitBefore << "\n" + << "Cause: " << *cand.waitCause << "\n"); + + LLVM_DEBUG(llvm::dbgs() << "=== Insert Sync End ===\n\n"); + + auto srcPipe = cand.srcCoreType == CoreType::CUBE_ONLY + ? hivm::PIPE::PIPE_FIX + : hivm::PIPE::PIPE_MTE2; + auto dstPipe = hivm::PIPE::PIPE_S; + auto [forwardSet, forwardWait] = createBlockSyncBetween( + builder, srcPipe, dstPipe, cand, syncFlag % MAX_FLAG_ID); + forwardSet->setAttr("ssbuf.flagid", builder.getUI32IntegerAttr(syncFlag)); + forwardWait->setAttr("ssbuf.flagid", builder.getUI32IntegerAttr(syncFlag)); + + auto loopAncestor = cand.waitBefore->getParentOfType(); + if (loopAncestor) { // backward sync, need to kick-start and clean-up; only + // need to ensure the first and last one, because we + // ensure once a flag is consumed, it will be set by + // construction (wait-set in the same block) + while (auto newAncestor = + loopAncestor->getParentOfType()) { + loopAncestor = newAncestor; + } + + auto srcPipe = cand.srcCoreType != CoreType::CUBE_ONLY + ? hivm::PIPE::PIPE_FIX + : hivm::PIPE::PIPE_MTE2; + auto backendFlag = syncFlag; + syncFlag++; + auto backwardCandidate = + SyncCandidate{!cand.srcCoreType, cand.waitCause, cand.waitBefore, + cand.setCause, cand.setAfter}; + auto [backwardSet, backwardWait] = createBlockSyncBetween( + builder, srcPipe, dstPipe, backwardCandidate, syncFlag % MAX_FLAG_ID); + + backwardSet->setAttr("ssbuf.backward", builder.getUnitAttr()); + backwardWait->setAttr("ssbuf.backward", builder.getUnitAttr()); + backwardSet->setAttr("ssbuf.flagid", + builder.getUI32IntegerAttr(backendFlag)); + backwardWait->setAttr("ssbuf.flagid", + builder.getUI32IntegerAttr(backendFlag)); + + builder.setInsertionPoint(loopAncestor); + auto kickstart = createBlockSync( + builder, toHivm(backwardCandidate.srcCoreType), srcPipe, dstPipe, + syncFlag, backwardCandidate.setCause); + + builder.setInsertionPointAfter(loopAncestor); + auto cleanup = createBlockSync( + builder, toHivm(!backwardCandidate.srcCoreType), srcPipe, dstPipe, + syncFlag, backwardCandidate.waitCause); + } + + syncFlag++; + }; + + if (candidates.empty()) { + return; + } + + auto setAfterDominate = [&domInfo](Operation *a, Operation *b) { + if (domInfo.dominates(a, b)) { + return true; + } + if (domInfo.dominates(b, a)) { + return false; + } + if (a->isAncestor(b)) { + return false; + } + if (b->isAncestor(a)) { + return true; + } + return false; + }; + + auto waitBeforePostDominate = [&postDomInfo](Operation *a, Operation *b) { + if (postDomInfo.postDominates(a, b)) { + return true; + } + if (postDomInfo.postDominates(b, a)) { + return false; + } + if (a->isAncestor(b)) { + return true; + } + if (b->isAncestor(a)) { + return false; + } + return false; + }; + + /** + * Sorts the candidates by (setAfter, waitBefore), and then removes redundant + * sync candidates that fully contains another + */ + auto selectCandidatesToInsertSync = + [&setAfterDominate, &waitBeforePostDominate, + &addBlockSyncCommon](llvm::SmallVectorImpl &candidates) { + llvm::sort( + candidates, [&](const SyncCandidate &a, const SyncCandidate &b) { + if (a.setAfter != b.setAfter) { + return setAfterDominate(a.setAfter, b.setAfter); + } + + if (a.waitBefore != b.waitBefore) { + return waitBeforePostDominate(a.waitBefore, b.waitBefore); + } + + return false; + }); + + for (auto [i, cand] : llvm::enumerate(candidates)) { + bool shouldInsert = true; + for (auto otherCand : ArrayRef(candidates).drop_front(i + 1)) { + bool duplicated = (cand.waitBefore == otherCand.waitBefore && + cand.setAfter == otherCand.setAfter && + cand.srcCoreType == otherCand.srcCoreType); + bool containsOther = + (cand.srcCoreType == otherCand.srcCoreType && + setAfterDominate(cand.setAfter, otherCand.setAfter) && + waitBeforePostDominate(cand.waitBefore, otherCand.waitBefore)); + if (duplicated || containsOther) { + shouldInsert = false; + break; + } + } + + if (shouldInsert) { + addBlockSyncCommon(cand); + } + } + }; + + selectCandidatesToInsertSync(candidates); +} + +void DAGSyncPass::runOnOperation() { + auto module = getOperation(); + mlir::OpBuilder builder(&getContext()); + + // 遍历所有函数 + for (auto funcOp : + llvm::make_early_inc_range(module.getOps())) { + // 跳过无效函数 + LegalizeDot(funcOp); + if (funcOp.getBody().empty()) { + continue; + } + + // llvm::outs() << "\n====================================\n"; + // llvm::outs() << "处理函数: " << funcOp.getName() << "\n"; + // llvm::outs() << "====================================\n"; + + auto unique_graph = Graph::fromMultiBlockFunc(funcOp); + std::shared_ptr shared_graph = std::move(unique_graph); + auto &main_graph = *shared_graph; + + auto funcName = funcOp.getName(); + + // 获取 DAG 图的映射 + auto opMapRaw = main_graph.getOpMapLegacy(); + valueTypes = &main_graph.getValueTypes(); + auto *opMap = &opMapRaw; + for (const auto &pair : *opMap) { + Operation *op = pair.first; // 键：Operation 指针 + Node *node = pair.second; // 值：Node 指针 + + // 打印指针地址（最直接的方式） + // llvm::outs() << "Operation*: " << *op + // << "\n"; + for (auto res : op->getResults()) { + llvm::outs() << "Value: " << (*valueTypes)[res] << "\n"; + } + } + + if (!opMap) { + llvm::errs() << "Warning: Failed to create DAG graph for function " + << funcOp.getName() << "\n"; + continue; + } + + // 用于避免重复插入同步 + llvm::DenseSet> + processedPairs; + int syncFlag = 1; + addMemEffectsSync(funcOp, shared_graph.get(), builder, syncFlag); + + // 3. 使用 walk 遍历函数中的所有操作 + funcOp.walk([&](mlir::Operation *op) { + // 查找当前操作对应的 Node + auto nodeIt = opMap->find(op); + if (nodeIt == opMap->end()) { + // 这个操作不在 entry block 的 DAG 图中 + // 可能是嵌套在控制流内部的操作 + return; + } + + OpNode *currentNode = nodeIt->second; + + // 检查是否是 scf.for 操作 + if (auto forOp = mlir::dyn_cast(op)) { + // 处理 scf.for 循环的特殊同步逻辑 + int temp = syncFlag % MAX_FLAG_ID; + processScfForSync(forOp, currentNode, valueTypes, builder, temp); + } + + // 获取当前节点的设备类型 + CoreType currentType = getNodeDeviceType(currentNode, valueTypes); + + // 打印操作信息（可选） + // if (!llvm::isa(op->getDialect())) { + // llvm::outs() << "操作: " << *op + // << " 设备类型: " + // << (currentType == CoreType::VECTOR_ONLY ? "VECTOR" : + // currentType == CoreType::CUBE_ONLY ? "CUBE" : + // "SCALAR") + // << "\n"; + // } + + // 4. 遍历当前节点的所有输入节点 + for (ValueNode *inputValNode : currentNode->getInputs()) { + auto inputOp = inputValNode->value.getDefiningOp(); + if (!inputOp && opMap->contains(inputOp)) { + continue; + } + + auto inputNode = (*opMap)[inputOp]; + + // 获取输入节点的设备类型 + CoreType inputType = getNodeDeviceType(inputNode, valueTypes); + + // 5. 判断是否需要插入同步和数据搬运 + if (needVectorCubeSync(inputType, currentType)) { + // 检查是否已经处理过这对操作 + auto opPair = std::make_pair(inputOp, op); + if (processedPairs.insert(opPair).second) { + // 插入同步和数据搬运指令 + // 检查是否是跨 block 的依赖 + mlir::Block *srcBlock = inputOp->getBlock(); + mlir::Block *dstBlock = op->getBlock(); + + if (srcBlock == dstBlock) { + // 同一 block 内 + insertSyncAndMovement(inputOp, op, inputType, currentType, + builder, syncFlag % 14, valueTypes, + main_graph); + syncFlag++; + } else { + // 跨 block，判断是否是外层到内层 + llvm::outs() << "#########\n"; + bool dstIsInnerBlock = false; + mlir::Operation *dstParentOp = dstBlock->getParentOp(); + + // 向上查找，看 dstBlock 是否在 srcBlock 的区域内 + while (dstParentOp) { + if (dstParentOp->getBlock() == srcBlock) { + dstIsInnerBlock = true; + break; + } + if (dstParentOp->getBlock()) { + dstParentOp = dstParentOp->getBlock()->getParentOp(); + } else { + break; + } + } + if (dstIsInnerBlock) { + + insertSyncAndMovementForCrossBlock( + inputOp, op, inputType, currentType, builder, syncFlag % 14, + dstIsInnerBlock, valueTypes, main_graph); + syncFlag++; + } + } + } + } + } + }); + + // llvm::outs() << "\n函数 " << funcOp.getName() << " 统计:\n"; + // llvm::outs() << " - 插入的总同步操作数: " << syncFlag << "\n"; + funcOp.walk([&](hivm::CopyOp copyOp) { + llvm::outs() << copyOp << " sss\n\n\n\n"; + rewriteCopyChainForCbub( + copyOp, + dyn_cast(copyOp.getOperands()[1].getType()).getShape(), + builder); + }); + GraphManager::getInstance().registerGraph(funcName, shared_graph); + } + + // llvm::outs()<> mlir::triton::createDAGSyncPass() { + return std::make_unique(); +} diff --git a/third_party/ascend/lib/TritonToAnnotation/CMakeLists.txt b/third_party/ascend/lib/TritonToAnnotation/CMakeLists.txt new file mode 100644 index 0000000000..b8f47ced6b --- /dev/null +++ b/third_party/ascend/lib/TritonToAnnotation/CMakeLists.txt @@ -0,0 +1,15 @@ +add_triton_library(TritonToAnnotation + TritonToAnnotation.cpp + + DEPENDS + TritonToAnnotationConversionPassIncGen + + LINK_LIBS + BiShengIRAnnotationDialect + BiShengIRDialectUtils + MLIRIR + MLIRPass + MLIRTransforms + MLIRSupport + TritonIR +) diff --git a/third_party/ascend/lib/TritonToAnnotation/TritonToAnnotation.cpp b/third_party/ascend/lib/TritonToAnnotation/TritonToAnnotation.cpp new file mode 100644 index 0000000000..b59f8b3598 --- /dev/null +++ b/third_party/ascend/lib/TritonToAnnotation/TritonToAnnotation.cpp @@ -0,0 +1,76 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "TritonToAnnotation/Passes.h" + +#include "Dialect/TritonAscend/IR/TritonAscendDialect.h" +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +namespace mlir { +namespace triton { +#define GEN_PASS_DEF_TRITONTOANNOTATION +#include "ascend/include/TritonToAnnotation/Passes.h.inc" +} // namespace triton +} // namespace mlir + +using namespace mlir; + +namespace { +struct TritonToAnnotationPass + : public mlir::triton::impl::TritonToAnnotationBase< + TritonToAnnotationPass> { + void runOnOperation() override; +}; +} // namespace + +struct TritonAnnotationConversionPattern + : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(mlir::triton::ascend::AnnotationOp op, + PatternRewriter &rewriter) const final { + auto markOp = rewriter.create(op.getLoc(), op.getSrc()); + // Forward all annotations. + markOp->setAttrs(op->getAttrs()); + rewriter.eraseOp(op); + return success(); + } +}; + +void TritonToAnnotationPass::runOnOperation() { + auto module = getOperation(); + ConversionTarget target(getContext()); + target.addLegalDialect(); + + RewritePatternSet patterns(&getContext()); + patterns.add(patterns.getContext()); + if (failed(applyPartialConversion(module, target, std::move(patterns)))) { + signalPassFailure(); + } +} + +std::unique_ptr> +mlir::triton::createTritonToAnnotationPass() { + return std::make_unique(); +} diff --git a/third_party/ascend/lib/TritonToGraph/AliasAnalysis.cpp b/third_party/ascend/lib/TritonToGraph/AliasAnalysis.cpp new file mode 100644 index 0000000000..4040278d06 --- /dev/null +++ b/third_party/ascend/lib/TritonToGraph/AliasAnalysis.cpp @@ -0,0 +1,248 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "TritonToGraph/AliasAnalysis.h" +#include "TritonToGraph/ControlFlowGraph.h" +#include "TritonToGraph/MemorySSA.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "triton/Dialect/Triton/IR/OpInterfaces.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/raw_ostream.h" + +#define DEBUG_TYPE "alias-analysis" + +using namespace mlir; +using namespace triton; +using namespace cfg; + +//===----------------------------------------------------------------------===// +// AliasAnalysis Implementation +//===----------------------------------------------------------------------===// + +void AliasAnalysis::analyzePointerAliases(ControlFlowGraph &cfg) { + LLVM_DEBUG(llvm::dbgs() << "=== Starting Pointer Alias Analysis ===\n"); + + // 步骤1: 识别全局内存参数 + triton::FuncOp func = cfg.getFunction(); + + LLVM_DEBUG(llvm::dbgs() << "Analyzing function: " << func.getName() << "\n"); + + for (BlockArgument arg : func.getArguments()) { + Type argType = arg.getType(); + + if (isPointerType(argType) && isGlobalMemoryType(argType)) { + // 为入参创建tensor对象 + std::string paramName = "param_" + std::to_string(arg.getArgNumber()); + + SmallVector shape; + Type elementType; + + // 使用辅助函数提取shape和element type + extractShapeAndElementType(argType, shape, elementType); + + TensorObject *tensor = + new TensorObject(paramName, shape, argType, elementType, + TensorObject::TensorKind::GLOBAL_MEMORY); + + // 记录alias关系 [param, param, tensor] + addAlias(arg, arg, tensor); + + LLVM_DEBUG(llvm::dbgs() + << " Found global memory parameter " << arg.getArgNumber() + << ": " << tensor->getName() << "\n"); + } + } + + LLVM_DEBUG(llvm::dbgs() << "Found " << aliasMap.size() + << " global memory parameters\n"); + + // 步骤2: 分析alias操作 + size_t aliasOpsFound = 0; + cfg.traverse([&](BasicBlock &bb) { + for (const auto &instPtr : bb.getInstructions()) { + const Instruction *inst = instPtr.get(); + Operation *op = inst->getOperation(); + + if (!op) + continue; + + if (auto addptrOp = dyn_cast(op)) { + analyzeAddPtrOp(addptrOp); + aliasOpsFound++; + } else if (auto makeTensorPtrOp = dyn_cast(op)) { + analyzeMakeTensorPtrOp(makeTensorPtrOp); + aliasOpsFound++; + } else if (auto loadOp = dyn_cast(op)) { + analyzeLoadOp(loadOp); + } else if (auto storeOp = dyn_cast(op)) { + analyzeStoreOp(storeOp); + } else if (auto broadcastOp = dyn_cast(op)) { + analyzeBroadcastOp(broadcastOp); + aliasOpsFound++; + } else if (auto splatOp = dyn_cast(op)) { + analyzeSplatOp(splatOp); + aliasOpsFound++; + } + } + }); + + LLVM_DEBUG( + llvm::dbgs() << "=== Alias Analysis Complete ===\n" + << "Total tracked aliases: " << aliasMap.size() << "\n" + << "Alias operations found: " << aliasOpsFound << "\n"); +} + +Value AliasAnalysis::getBasePointer(Value ptr) const { + // 递归查找base pointer + auto it = aliasMap.find(ptr); + if (it == aliasMap.end()) { + // 没有找到，返回自身（可能是原始base pointer） + return ptr; + } + + Value base = it->second; + + // 如果base不是自己，继续递归查找 + if (base != ptr) { + return getBasePointer(base); + } + + return base; +} + +void AliasAnalysis::analyzeAddPtrOp(mlir::triton::AddPtrOp addptrOp) { + // %new = tt.addptr %ptr, %offset + Value ptr = addptrOp.getPtr(); + Value result = addptrOp.getResult(); + + // result是ptr的alias，指向同一个tensor + Value basePtr = getBasePointer(ptr); + TensorObject *tensor = getTensorObject(basePtr); + + if (tensor) { + addAlias(result, basePtr, tensor); + + LLVM_DEBUG(llvm::dbgs() << " AddPtr: " << result << " -> " << basePtr + << " [" << tensor->getName() << "]\n"); + } +} + +void AliasAnalysis::analyzeMakeTensorPtrOp(mlir::triton::MakeTensorPtrOp op) { + // %tensor_ptr = tt.make_tensor_ptr %base, ... + Value basePtr = op.getBase(); + Value result = op.getResult(); + + TensorObject *tensor = getTensorObject(basePtr); + + if (tensor) { + addAlias(result, basePtr, tensor); + + LLVM_DEBUG(llvm::dbgs() << " MakeTensorPtr: " << result << " -> " + << basePtr << " [" << tensor->getName() << "]\n"); + } +} + +void AliasAnalysis::analyzeLoadOp(mlir::triton::LoadOp loadOp) { + // tt.load操作通常不改变alias关系， + // 但可以验证load的ptr是否被正确跟踪 + + Value ptr = loadOp.getPtr(); + + // 验证ptr是否被alias analysis跟踪 + TensorObject *tensor = getTensorObject(ptr); + + if (tensor) { + LLVM_DEBUG(llvm::dbgs() + << " Load from tracked tensor: " << tensor->getName() << "\n"); + } else { + LLVM_DEBUG(llvm::dbgs() << " Load from untracked pointer\n"); + } +} + +void AliasAnalysis::analyzeStoreOp(mlir::triton::StoreOp storeOp) { + // tt.store操作不改变alias关系， + // 但可以验证store的目标是否被正确跟踪 + + Value ptr = storeOp.getPtr(); + + // 验证ptr是否被alias analysis跟踪 + TensorObject *tensor = getTensorObject(ptr); + + if (tensor) { + LLVM_DEBUG(llvm::dbgs() + << " Store to tracked tensor: " << tensor->getName() << "\n"); + } else { + LLVM_DEBUG(llvm::dbgs() << " Store to untracked pointer\n"); + } +} + +void AliasAnalysis::analyzeBroadcastOp(mlir::triton::BroadcastOp broadcastOp) { + // %out = tt.broadcast %src : tensor> -> tensor> + // 输入和输出都是指针tensor，输出tensor的每个元素与输入tensor对应行/列元素是别名关系 + + Value src = broadcastOp.getSrc(); + Value result = broadcastOp.getResult(); + + // 只处理元素类型为指针的情况 + Type srcElemType = getElementTypeOrSelf(src.getType()); + if (!isPointerType(srcElemType)) + return; + + // 查找输入tensor关联的TensorObject + // 注意：src本身是指针tensor，需要获取其base pointer对应的tensor + Value srcBasePtr = getBasePointer(src); + TensorObject *tensor = getTensorObject(srcBasePtr); + + if (tensor) { + // broadcast操作保持指针值不变（只是复制到更多位置）， + // 因此result与src指向相同的底层tensor对象 + addAlias(result, srcBasePtr, tensor); + + LLVM_DEBUG(llvm::dbgs() << " Broadcast: " << result << " -> " << srcBasePtr + << " [" << tensor->getName() << "]\n"); + } +} + +void AliasAnalysis::analyzeSplatOp(mlir::triton::SplatOp splatOp) { + // %out = tt.splat %src : !tt.ptr -> tensor> + // 将标量指针广播到tensor的每个位置，所有输出元素都是输入指针的别名 + + Value src = splatOp.getSrc(); // 标量指针 + Value result = splatOp.getResult(); // 指针tensor + + // 确认输入是指针类型 + if (!isPointerType(src.getType())) + return; + + // 获取标量指针的base pointer和对应的TensorObject + Value basePtr = getBasePointer(src); + TensorObject *tensor = getTensorObject(basePtr); + + if (tensor) { + // splat操作将同一个指针值复制到tensor的每个元素， + // 因此result tensor中所有元素都与src是别名关系（指向同一内存对象） + addAlias(result, basePtr, tensor); + + LLVM_DEBUG(llvm::dbgs() << " Splat: " << result << " -> " << basePtr + << " [" << tensor->getName() << "]\n"); + } +} diff --git a/third_party/ascend/lib/TritonToGraph/CMakeLists.txt b/third_party/ascend/lib/TritonToGraph/CMakeLists.txt new file mode 100644 index 0000000000..d5e6a201a2 --- /dev/null +++ b/third_party/ascend/lib/TritonToGraph/CMakeLists.txt @@ -0,0 +1,28 @@ +add_triton_library(TritonToGraph + AliasAnalysis.cpp + ControlFlowGraph.cpp + ControlFlowGraphBuilder.cpp + DataflowGraph.cpp + MemorySsaBuilder.cpp + Passes.cpp + # Required by ascend_ir.cc's program-axis / grid-mapping bindings. + ProgramAxisDependenceAnalysis.cpp + ProgramGridTransform.cpp + ProgramMappingScalarSpecialization.cpp + + DEPENDS + TritonToGraphPassIncGen + + LINK_LIBS PUBLIC + MLIRArithDialect + MLIRControlFlowDialect + MLIRDialectUtils + MLIRIR + MLIRFuncDialect + MLIRPass + MLIRSCFDialect + MLIRSupport + MLIRTransforms + TritonIR + TritonTransforms +) diff --git a/third_party/ascend/lib/TritonToGraph/ControlFlowGraph.cpp b/third_party/ascend/lib/TritonToGraph/ControlFlowGraph.cpp new file mode 100644 index 0000000000..f63de7a109 --- /dev/null +++ b/third_party/ascend/lib/TritonToGraph/ControlFlowGraph.cpp @@ -0,0 +1,1140 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "TritonToGraph/ControlFlowGraph.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Types.h" + +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/Path.h" +#include "llvm/Support/raw_ostream.h" + +#include +#include +#include +#include + +using namespace mlir; +using namespace triton; +using namespace cfg; + +//===----------------------------------------------------------------------===// +// Instruction +//===----------------------------------------------------------------------===// + +std::string Instruction::getAsString() const { + if (!operation) + return ""; + std::string str; + llvm::raw_string_ostream os(str); + operation->print(os); + return str; +} + +void Instruction::print(raw_ostream &os, unsigned indent) const { + std::string instStr = getAsString(); + os << "Inst[" << id << "]: "; + + if (operation) { + BlockType parentType = parentBlock->getType(); + if (parentType == BlockType::IF_COND || parentType == BlockType::FOR_COND || + parentType == BlockType::WHILE_COND || + parentType == BlockType::COND_BR || parentType == BlockType::BR) { + size_t newlinePos = instStr.find('\n'); + if (newlinePos != std::string::npos) { + instStr = instStr.substr(0, newlinePos) + " ..."; + } + } + os.indent(indent) << instStr; + } else { + os << "\n"; + } +} + +void Instruction::dump() const { print(llvm::errs()); } + +//===----------------------------------------------------------------------===// +// BasicBlock +//===----------------------------------------------------------------------===// + +void BasicBlock::addInstruction(std::unique_ptr inst) { + instructions.push_back(std::move(inst)); +} + +Instruction *BasicBlock::getInstruction(size_t idx) const { + if (idx < instructions.size()) { + return instructions[idx].get(); + } + return nullptr; +} + +void BasicBlock::addSuccessor(BasicBlock *succ) { + if (!succ) + return; + // 避免重复添加 + for (auto *s : successors) { + if (s == succ) + return; + } + successors.push_back(succ); + succ->addPredecessor(this); +} + +void BasicBlock::addPredecessor(BasicBlock *pred) { + if (!pred) + return; + // 避免重复添加 + for (auto *p : predecessors) { + if (p == pred) + return; + } + predecessors.push_back(pred); +} + +std::string BasicBlock::getName() const { + std::string name = "BB"; + name += std::to_string(id); + return name; +} + +bool BasicBlock::endsWithReturnOp() const { + // 空块检查 + if (instructions.empty()) { + return false; + } + + // 获取最后一条指令 + const Instruction *lastInst = instructions.back().get(); + if (!lastInst) { + return false; + } + + // 获取对应的 Operation + Operation *op = lastInst->getOperation(); + if (!op) { + return false; + } + + // 检查是否为 triton::ReturnOp + return isa(op); +} + +StringRef BasicBlock::getTypeString() const { + switch (type) { + case BlockType::NORMAL: + return "NORMAL"; + case BlockType::ENTRY: + return "ENTRY"; + case BlockType::EXIT: + return "EXIT"; + case BlockType::IF_COND: + return "IF_COND"; + case BlockType::FOR_COND: + return "FOR_COND"; + case BlockType::WHILE_COND: + return "WHILE_COND"; + case BlockType::COND_BR: + return "COND_BR"; + case BlockType::BR: + return "BR"; + case BlockType::LOOP_BODY: + return "LOOP_BODY"; + case BlockType::LOOP_EXIT: + return "LOOP_EXIT"; + } + return "UNKNOWN"; +} + +void BasicBlock::print(raw_ostream &os) const { + os << "============================================================\n"; + os << "BasicBlock " << getName() << " [" << getTypeString() << "]\n"; + if (parentStructure) { + os << " Parent Structure: " << parentStructure->getName() << "\n"; + } + + // 打印前驱 + if (!predecessors.empty()) { + os << " Predecessors: ["; + for (size_t i = 0; i < predecessors.size(); ++i) { + if (i > 0) + os << ", "; + os << predecessors[i]->getName(); + } + os << "]\n"; + } + + // 打印后继 + if (!successors.empty()) { + os << " Successors: ["; + for (size_t i = 0; i < successors.size(); ++i) { + if (i > 0) + os << ", "; + os << successors[i]->getName(); + } + os << "]\n"; + } + + // 打印指令 - COND节点只打印第一行 + os << " Instructions (" << instructions.size() << "):\n"; + for (const auto &inst : instructions) { + std::string instStr = inst->getAsString(); + + // 对COND节点截断：只保留第一行 + if (type == BlockType::IF_COND || type == BlockType::FOR_COND || + type == BlockType::WHILE_COND) { + size_t newlinePos = instStr.find('\n'); + if (newlinePos != std::string::npos) { + instStr = instStr.substr(0, newlinePos) + " ..."; + } + } + + os.indent(4) << "Inst[" << inst->getId() << "]: " << instStr << "\n"; + } + os << "\n"; +} + +void BasicBlock::dump() const { print(llvm::errs()); } + +void BasicBlock::exportToJSON(raw_ostream &os, unsigned indent) const { + std::string ind(indent, ' '); + os << ind << "{\n"; + os << ind << " \"id\": " << id << ",\n"; + os << ind << " \"name\": \"" << getName() << "\",\n"; + os << ind << " \"type\": \"" << getTypeString() << "\",\n"; + os << ind << " \"parentStructure\": " + << (parentStructure ? std::to_string(parentStructure->getId()) : "null") + << ",\n"; + + // 前驱 + os << ind << " \"predecessors\": ["; + for (size_t i = 0; i < predecessors.size(); ++i) { + if (i > 0) + os << ", "; + os << predecessors[i]->getId(); + } + os << "],\n"; + + // 后继 + os << ind << " \"successors\": ["; + for (size_t i = 0; i < successors.size(); ++i) { + if (i > 0) + os << ", "; + os << successors[i]->getId(); + } + os << "],\n"; + + // 指令 + os << ind << " \"instructions\": [\n"; + for (size_t i = 0; i < instructions.size(); ++i) { + os << ind << " {\n"; + os << ind << " \"id\": " << instructions[i]->getId() << ",\n"; + // 转义字符串用于 JSON + std::string instStr = instructions[i]->getAsString(); + + // COND节点只取第一行 + if (type == BlockType::IF_COND || type == BlockType::FOR_COND || + type == BlockType::WHILE_COND) { + size_t newlinePos = instStr.find('\n'); + if (newlinePos != std::string::npos) { + instStr = instStr.substr(0, newlinePos) + " ..."; + } + } + + // 简单的 JSON 字符串转义 + std::string escaped; + for (char c : instStr) { + if (c == '"') + escaped += "\\\""; + else if (c == '\\') + escaped += "\\\\"; + else if (c == '\n') + escaped += "\\n"; + else if (c == '\r') + escaped += "\\r"; + else if (c == '\t') + escaped += "\\t"; + else if ((unsigned char)c < 0x20) { + char buf[8]; + snprintf(buf, sizeof(buf), "\\u%04x", c); + escaped += buf; + } else { + escaped += c; + } + } + os << ind << " \"operation\": \"" << escaped << "\"\n"; + os << ind << " }"; + if (i < instructions.size() - 1) + os << ","; + os << "\n"; + } + os << ind << " ]\n"; + os << ind << "}"; +} + +//===----------------------------------------------------------------------===// +// ControlFlowGraph +//===----------------------------------------------------------------------===// + +ControlFlowGraph::ControlFlowGraph(triton::FuncOp func) : function(func) {} + +ControlFlowGraph::~ControlFlowGraph() = default; + +BasicBlock *ControlFlowGraph::createBasicBlock(BlockType type, + BasicBlock *parentStructure) { + auto bb = std::make_unique(nextBlockId++, type, parentStructure); + BasicBlock *bbPtr = bb.get(); + basicBlocks.push_back(std::move(bb)); + + // 自动设置入口/出口块 + if (type == BlockType::ENTRY) + entryBlock = bbPtr; + if (type == BlockType::EXIT) + exitBlock = bbPtr; + + return bbPtr; +} + +void ControlFlowGraph::addEdge(BasicBlock *from, BasicBlock *to) { + if (!from || !to) + return; + from->addSuccessor(to); +} + +// 判断是否为回边（从循环体回到循环头） +bool ControlFlowGraph::isBackEdge(BasicBlock *from, BasicBlock *to) const { + // 回边定义：指向 FOR_COND/WHILE_COND 且 from 是该循环的后代 + if (to->getType() != BlockType::FOR_COND && + to->getType() != BlockType::WHILE_COND) + return false; + + // 检查 from 是否属于以 to 为头的循环 + // 方法：检查 from 的 parentStructure 链是否包含 to + BasicBlock *current = from; + while (current) { + if (current == to) + return true; + current = current->getParentStructure(); + } + return false; +} + +void ControlFlowGraph::searchNormalBlock(BasicBlock *block, + OperationVisitor callback) const { + for (const auto &inst : block->getInstructions()) { + if (Operation *op = inst->getOperation()) { + callback(op); + } + } +} + +void ControlFlowGraph::searchCondBlock(BasicBlock *block, + OperationVisitor callback) const { + if (!block) + return; + + // 获取该 Cond block 对应的 exit block（停止条件） + BasicBlock *exitBlock = block->getExitBlock(); + + // 遍历 Cond block 的每一个 successor + for (BasicBlock *succ : block->getSuccessors()) { + if (!succ) + continue; + + // 从每个 successor 开始，沿着 block 链条向下遍历 + // 使用队列进行 BFS 遍历该分支 + std::deque workList; + workList.push_back(succ); + + while (!workList.empty()) { + BasicBlock *current = workList.front(); + workList.pop_front(); + + // 如果碰到 exit block，停止该分支的遍历 + if (current == exitBlock) + break; + + // 对当前 block 调用 searchBlock + searchBlock(current, callback); + + if (current->getType() == BlockType::NORMAL) { + // 将后继加入队列（排除回边） + for (BasicBlock *next : current->getSuccessors()) { + // 跳过回边避免无限循环 + if (isBackEdge(current, next)) + continue; + workList.push_back(next); + } + } else if (block->getType() == BlockType::IF_COND || + block->getType() == BlockType::FOR_COND || + block->getType() == BlockType::WHILE_COND) { + BasicBlock *next = block->getExitBlock(); + if (isBackEdge(current, next)) + continue; + workList.push_back(next); + } + } + } +} + +void ControlFlowGraph::searchBlock(BasicBlock *block, + OperationVisitor callback) const { + if (block->getType() == BlockType::NORMAL) { + searchNormalBlock(block, callback); + } else if (block->getType() == BlockType::IF_COND || + block->getType() == BlockType::FOR_COND || + block->getType() == BlockType::WHILE_COND) { + searchCondBlock(block, callback); + } +} + +void ControlFlowGraph::traverse(BlockVisitor visitor) { + // 从入口块开始进行拓扑排序 + std::vector topoOrder; + DenseSet visited; + + // 使用DFS进行拓扑排序 + std::function dfs = [&](BasicBlock *bb) { + if (!bb || visited.contains(bb)) + return; + + visited.insert(bb); + + // 递归访问所有后继块 + for (BasicBlock *succ : bb->getSuccessors()) { + dfs(succ); + } + + // 在后序位置添加，最后需要反转 + topoOrder.push_back(bb); + }; + + // 从入口块开始DFS + dfs(entryBlock); + + // 反转得到拓扑序 + std::reverse(topoOrder.begin(), topoOrder.end()); + + // 按照拓扑序遍历 + for (BasicBlock *bb : topoOrder) { + visitor(*bb); + } +} + +void ControlFlowGraph::print(raw_ostream &os) const { + auto funcName = const_cast(function).getName(); + os << "=================================================================\n"; + os << "Control Flow Graph for function '" + << (funcName.empty() ? "unnamed" : funcName) << "'\n"; + os << "Number of blocks: " << basicBlocks.size() << "\n"; + os << "=================================================================\n\n"; + + for (const auto &bb : basicBlocks) { + bb->print(os); + } +} + +void ControlFlowGraph::dump() const { print(llvm::errs()); } + +void ControlFlowGraph::exportDOT(raw_ostream &os) const { + auto funcName = const_cast(function).getName(); + std::string funcNameStr = funcName.empty() ? "unnamed" : funcName.str(); + + // 清理函数名用于 DOT 标识符 + std::string cleanFuncName; + for (char c : funcNameStr) { + if (isalnum(c) || c == '_') + cleanFuncName += c; + else + cleanFuncName += '_'; + } + + os << "digraph CFG_" << cleanFuncName << " {\n"; + os << " label=\"CFG for " << funcNameStr << "\";\n"; + os << " labelloc=t;\n"; + os << " rankdir=TB;\n"; + os << " splines=true;\n"; // 使用曲线边 + os << " overlap=false;\n"; // 防止节点重叠 + os << " nodesep=0.6;\n"; // 节点水平间距 + os << " ranksep=1.2;\n"; // 层间距 + os << " fontsize=12;\n\n"; + + // 设置节点样式 + for (const auto &bb : basicBlocks) { + os << " \"" << bb->getName() << "\" ["; + os << "label=\"" << bb->getName() << "\\n(" << bb->getTypeString() << ")"; + + // 指令摘要：COND节点只取第一行，其他节点最多3条 + size_t numInsts = bb->getNumInstructions(); + if (numInsts > 0) { + os << "\\n"; + + // COND节点只显示第一条指令的第一行 + if (bb->getType() == BlockType::IF_COND || + bb->getType() == BlockType::FOR_COND || + bb->getType() == BlockType::WHILE_COND) { + std::string instStr = bb->getInstruction(0)->getAsString(); + size_t newlinePos = instStr.find('\n'); + if (newlinePos != std::string::npos) { + instStr = instStr.substr(0, newlinePos); + } + if (instStr.length() > 40) + instStr = instStr.substr(0, 37) + "..."; + + // 转义 + std::string escaped; + for (char c : instStr) { + if (c == '"') + escaped += "\\\""; + else if (c == '\\') + escaped += "\\\\"; + else if (c == '\n') + escaped += "\\n"; + else + escaped += c; + } + os << escaped; + if (numInsts > 1) + os << "\\n... (" << (numInsts - 1) << " more)"; + } else { + // 其他节点显示最多3条 + for (size_t i = 0; i < std::min(numInsts, (size_t)3); ++i) { + std::string instStr = bb->getInstruction(i)->getAsString(); + // 只取第一行 + size_t newlinePos = instStr.find('\n'); + if (newlinePos != std::string::npos) { + instStr = instStr.substr(0, newlinePos) + "..."; + } + if (instStr.length() > 40) + instStr = instStr.substr(0, 37) + "..."; + + std::string escaped; + for (char c : instStr) { + if (c == '"') + escaped += "\\\""; + else if (c == '\\') + escaped += "\\\\"; + else if (c == '\n') + escaped += "\\n"; + else + escaped += c; + } + os << escaped << "\\n"; + } + if (numInsts > 3) + os << "... (" << (numInsts - 3) << " more)\\n"; + } + } + + os << "\", "; + + // 形状和颜色 + switch (bb->getType()) { + case BlockType::ENTRY: + os << "style=filled, fillcolor=lightgreen, shape=ellipse"; + break; + case BlockType::EXIT: + os << "style=filled, fillcolor=lightcoral, shape=ellipse"; + break; + case BlockType::IF_COND: + os << "style=filled, fillcolor=lightyellow, shape=diamond"; + break; + case BlockType::FOR_COND: + case BlockType::WHILE_COND: + os << "style=filled, fillcolor=lightblue, shape=diamond"; + break; + case BlockType::LOOP_BODY: + os << "style=filled, fillcolor=lightcyan, shape=box"; + break; + case BlockType::LOOP_EXIT: + os << "style=filled, fillcolor=lightpink, shape=box"; + break; + default: + os << "shape=box"; + break; + } + os << ", fontsize=10];\n"; + } + + os << "\n"; + + // 输出边：回边红色，其他绿色 + for (const auto &bb : basicBlocks) { + for (auto *succ : bb->getSuccessors()) { + bool isBack = isBackEdge(bb.get(), succ); + os << " \"" << bb->getName() << "\" -> \"" << succ->getName() << "\""; + os << " [color=" << (isBack ? "red" : "green"); + os << ", penwidth=" << (isBack ? "2.5" : "1.5") << ""; + if (isBack) { + os << ", style=dashed"; // 回边用虚线 + } + os << "];\n"; + } + } + + os << "}\n"; +} + +llvm::Error ControlFlowGraph::exportToFile(StringRef filename) const { + std::error_code ec; + llvm::raw_fd_ostream os(filename, ec, llvm::sys::fs::OF_Text); + + if (ec) { + return llvm::createStringError(llvm::inconvertibleErrorCode(), + "Failed to open file: " + filename); + } + + // 判断文件扩展名 + if (filename.ends_with(".dot")) { + exportDOT(os); + } else if (filename.ends_with(".json")) { + exportToJSON(os); + } else { + print(os); + } + + os.close(); + return llvm::Error::success(); +} + +void ControlFlowGraph::exportToJSON(raw_ostream &os) const { + auto funcName = const_cast(function).getName(); + std::string funcNameStr = funcName.empty() ? "unnamed" : funcName.str(); + + os << "{\n"; + os << " \"functionName\": \"" << funcNameStr << "\",\n"; + os << " \"numBlocks\": " << basicBlocks.size() << ",\n"; + os << " \"blocks\": [\n"; + + for (size_t i = 0; i < basicBlocks.size(); ++i) { + basicBlocks[i]->exportToJSON(os, 4); + if (i < basicBlocks.size() - 1) + os << ","; + os << "\n"; + } + + // 添加边信息（带类型标记） + os << " ],\n"; + os << " \"edges\": [\n"; + + bool first = true; + for (const auto &bb : basicBlocks) { + for (auto *succ : bb->getSuccessors()) { + if (!first) + os << ",\n"; + first = false; + + bool isBack = isBackEdge(bb.get(), succ); + os << " {\"from\": " << bb->getId() << ", \"to\": " << succ->getId() + << ", \"fromName\": \"" << bb->getName() << "\"" + << ", \"toName\": \"" << succ->getName() << "\"" + << ", \"type\": \"" << (isBack ? "back" : "normal") << "\"}"; + } + } + os << "\n ]\n"; + os << "}\n"; +} + +llvm::Error ControlFlowGraph::exportToHTML(StringRef filename) const { + std::error_code ec; + llvm::raw_fd_ostream os(filename, ec, llvm::sys::fs::OF_Text); + + if (ec) { + return llvm::createStringError(llvm::inconvertibleErrorCode(), + "Failed to open file: " + filename); + } + + auto funcName = const_cast(function).getName(); + std::string funcNameStr = funcName.empty() ? "unnamed" : funcName.str(); + + os << R"html( + + + + CFG for )html" + << funcNameStr << R"html( + + + + + + +
+ +
+
+
正常控制流 (向下)
+
循环回边 (向上)
+
+ + + + +)html"; + + os.close(); + return llvm::Error::success(); +} diff --git a/third_party/ascend/lib/TritonToGraph/ControlFlowGraphBuilder.cpp b/third_party/ascend/lib/TritonToGraph/ControlFlowGraphBuilder.cpp new file mode 100644 index 0000000000..30038d70f6 --- /dev/null +++ b/third_party/ascend/lib/TritonToGraph/ControlFlowGraphBuilder.cpp @@ -0,0 +1,903 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "TritonToGraph/ControlFlowGraphBuilder.h" +#include "TritonToGraph/DataflowGraph.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "llvm/Support/Debug.h" +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/Path.h" +#include "llvm/Support/raw_ostream.h" + +#define DEBUG_TYPE "build-cfg" + +using namespace mlir; +using namespace mlir::triton; +using namespace cfg; + +//===----------------------------------------------------------------------===// +// BuildCFGPass Implementation +//===----------------------------------------------------------------------===// + +void BuildCFGPass::runOnOperation() { + auto module = getOperation(); + llvm::errs() << "Building CFG for module\n"; + + // 获取输出目录（从命令行选项，默认为当前目录） + // TableGen 生成的选项成员变量名为 outputDir + std::string outputDir = this->outputDir; + + // 创建输出目录 + llvm::SmallString<128> outputPath(outputDir); + llvm::sys::fs::create_directories(outputPath); + + llvm::errs() << "CFG output directory: " << outputPath << "\n"; + + // 遍历模块中的所有 Triton 函数 (tt.func) + for (triton::FuncOp func : module.getOps()) { + llvm::errs() << "Processing function: " << func.getName() << "\n"; + + auto cfg = buildForFunction(func); + if (!cfg) { + func.emitError() << "Failed to build CFG for function"; + signalPassFailure(); + return; + } + + // 打印 CFG 到标准输出 + // cfg->print(llvm::outs()); + + // 导出到文件 + std::string baseName = func.getName().str(); + + // 导出文本格式 + llvm::SmallString<128> textPath(outputPath); + llvm::sys::path::append(textPath, baseName + "_cfg.txt"); + if (auto err = cfg->exportToFile(textPath)) { + llvm::errs() << "Failed to export CFG to " << textPath << "\n"; + } else { + llvm::errs() << "Exported CFG to " << textPath << "\n"; + } + + // 导出 DOT 格式 + llvm::SmallString<128> dotPath(outputPath); + llvm::sys::path::append(dotPath, baseName + "_cfg.dot"); + if (auto err = cfg->exportToFile(dotPath)) { + llvm::errs() << "Failed to export CFG to " << dotPath << "\n"; + } else { + llvm::errs() << "Exported CFG to " << dotPath << "\n"; + } + + // 导出 JSON 格式 + llvm::SmallString<128> jsonPath(outputPath); + llvm::sys::path::append(jsonPath, baseName + "_cfg.json"); + if (auto err = cfg->exportToFile(jsonPath)) { + llvm::errs() << "Failed to export CFG to " << jsonPath << "\n"; + } else { + llvm::errs() << "Exported CFG to " << jsonPath << "\n"; + } + + // 导出 HTML 格式（网页可视化） + llvm::SmallString<128> htmlPath(outputPath); + llvm::sys::path::append(htmlPath, baseName + "_cfg.html"); + if (auto err = cfg->exportToHTML(htmlPath)) { + llvm::errs() << "Failed to export CFG to " << htmlPath << "\n"; + } else { + llvm::errs() << "Exported CFG to " << htmlPath << "\n"; + } + + // 构建 DataFlowGraph（包含 Memory SSA 分析） + llvm::errs() << " Building DataFlowGraph with Memory SSA...\n"; + DataFlowGraph dataFlowGraph(*cfg); + dataFlowGraph.build(); + + // 导出 DataFlowGraph + std::error_code ec; + llvm::SmallString<128> dataflowPath(outputPath); + llvm::sys::path::append(dataflowPath, baseName + "_dataflow.json"); + llvm::raw_fd_ostream dfOs(dataflowPath, ec); + if (!ec) { + dataFlowGraph.exportToJSON(dfOs); + llvm::errs() << " Exported DataFlowGraph to " << dataflowPath << "\n"; + } + } +} + +std::unique_ptr +BuildCFGPass::buildForFunction(triton::FuncOp func) { + ControlFlowGraphBuilder cfgBuilder; + return cfgBuilder.build(func); +} + +ControlFlowGraphBuilder::RegionBlocks ControlFlowGraphBuilder::buildForRegion( + Region ®ion, cfg::ControlFlowGraph &cfg, cfg::BasicBlock *entryBlock, + cfg::BasicBlock *parentStructure) { + cfg::BasicBlock *currentBB = entryBlock; + cfg::BasicBlock *lastBlock = entryBlock; + + // 首先为 region 中的所有 block 创建对应的 BasicBlock 映射 + // 这样可以确保在处理 cf.cond_br 等跳转指令时目标块已存在 + for (Block &block : region) { + if (!blockToBasicBlockMap.count(&block)) { + auto *bb = cfg.createBasicBlock(cfg::BlockType::NORMAL, parentStructure); + registerBlockMapping(&block, bb); + } + } + + // 遍历 region 中的所有 block + for (Block &block : region) { + // 获取该 block 对应的 BasicBlock + cfg::BasicBlock *blockBB = blockToBasicBlockMap[&block]; + + // 如果是第一个 block，合并到 entryBlock + if (blockBB == blockToBasicBlockMap.lookup(®ion.front())) { + // 将 entryBlock 的指令移动到 blockBB（或者反过来） + // 这里简化处理：使用 entryBlock 继续处理 + currentBB = processBlock(block, cfg, currentBB, parentStructure); + } else { + // 确保从上一个 block 的结尾连接到这个 block + if (lastBlock && lastBlock != blockBB) { + // 检查是否已经有边连接 + bool hasEdge = false; + for (auto *succ : lastBlock->getSuccessors()) { + if (succ == blockBB) { + hasEdge = true; + break; + } + } + if (!hasEdge && !lastBlock->endsWithReturnOp()) { + cfg.addEdge(lastBlock, blockBB); + } + } + currentBB = processBlock(block, cfg, blockBB, parentStructure); + } + + if (currentBB) { + lastBlock = currentBB; + if (lastBlock->endsWithReturnOp()) { + cfg.addEdge(lastBlock, cfg.getExitBlock()); + } + } + } + + return {entryBlock, lastBlock}; +} + +cfg::BasicBlock * +ControlFlowGraphBuilder::processBlock(Block &block, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *currentBB, + cfg::BasicBlock *parentStructure) { + if (!currentBB) + return nullptr; + + // 遍历 block 中的所有操作 + for (Operation &op : block) { + // 检查是否是控制流操作 + if (isa(op)) { + // 为 if 条件创建单独的 basic block + auto *ifCondBB = + cfg.createBasicBlock(BlockType::IF_COND, parentStructure); + + // 将当前 if 指令添加到 ifCondBB + createInstruction(&op, ifCondBB, cfg); + + // 连接当前块到 if 条件块 + cfg.addEdge(currentBB, ifCondBB); + + // 处理 if 操作，返回 if 后面的块 + currentBB = + handleIfOp(cast(op), cfg, ifCondBB, parentStructure); + } else if (isa(op)) { + // 为 for 条件创建单独的 basic block + auto *forCondBB = + cfg.createBasicBlock(BlockType::FOR_COND, parentStructure); + + // 将当前 for 指令添加到 forCondBB + createInstruction(&op, forCondBB, cfg); + + // 连接当前块到 for 条件块 + cfg.addEdge(currentBB, forCondBB); + + // 处理 for 操作，返回 for 后面的块 + currentBB = + handleForOp(cast(op), cfg, forCondBB, parentStructure); + } else if (isa(op)) { + // 为 while 条件创建单独的 basic block + auto *whileCondBB = + cfg.createBasicBlock(BlockType::WHILE_COND, parentStructure); + + // 将当前 while 指令添加到 whileCondBB + createInstruction(&op, whileCondBB, cfg); + + // 连接当前块到 while 条件块 + cfg.addEdge(currentBB, whileCondBB); + + // 处理 while 操作，返回 while 后面的块 + currentBB = handleWhileOp(cast(op), cfg, whileCondBB, + parentStructure); + } else if (isa(op)) { + // yield 操作：创建指令并继续（后续由循环处理逻辑连接） + createInstruction(&op, currentBB, cfg); + } else if (isa(op)) { + // condition 操作（while 循环条件）：创建指令并继续 + createInstruction(&op, currentBB, cfg); + } else if (isa(op)) { + // cf.cond_br 条件分支 - 创建专门的 COND_BR 块并处理 + auto *condBrBB = + cfg.createBasicBlock(BlockType::COND_BR, parentStructure); + + // 将 cond_br 指令添加到 condBrBB + createInstruction(&op, condBrBB, cfg); + + // 连接当前块到 COND_BR 块 + cfg.addEdge(currentBB, condBrBB); + + // 处理 cond_br 操作，返回后续的基本块 + currentBB = handleCondBranchOp(cast(op), cfg, condBrBB, + parentStructure); + } else if (isa(op)) { + // cf.br 无条件跳转 - 创建专门的 BR 块并处理 + auto *brBB = cfg.createBasicBlock(BlockType::BR, parentStructure); + + // 将 br 指令添加到 brBB + createInstruction(&op, brBB, cfg); + + // 连接当前块到 BR 块 + cfg.addEdge(currentBB, brBB); + + // 处理 br 操作 + currentBB = + handleBranchOp(cast(op), cfg, brBB, parentStructure); + } else if (isa(op)) { + // return 操作 + createInstruction(&op, currentBB, cfg); + } else if (op.getNumRegions() > 0) { + // 有内部区域的 Triton 操作 (如 tt.reduce, tt.scan 等) + // 先创建指令 + auto *inst = createInstruction(&op, currentBB, cfg); + + // 为该操作创建子图 + auto subGraph = + std::make_unique(cfg.getFunction()); + auto *subEntry = subGraph->createBasicBlock(cfg::BlockType::ENTRY); + subGraph->setEntryBlock(subEntry); + auto *subExit = subGraph->createBasicBlock(cfg::BlockType::EXIT); + subGraph->setExitBlock(subExit); + + // 用于子图中生成唯一指令 ID + size_t subInstId = 0; + + // 遍历所有区域 + for (size_t regionIdx = 0; regionIdx < op.getNumRegions(); ++regionIdx) { + Region ®ion = op.getRegion(regionIdx); + if (!region.empty()) { + // 为每个区域构建 CFG + cfg::BasicBlock *regionEntryBB = nullptr; + cfg::BasicBlock *regionLastBB = nullptr; + + for (Block ®ionBlock : region) { + auto *bb = subGraph->createBasicBlock(cfg::BlockType::NORMAL); + if (!regionEntryBB) + regionEntryBB = bb; + + // 将区域中的操作添加到子图的基本块 + for (Operation ®ionOp : regionBlock) { + auto regionInst = std::make_unique( + subInstId++, ®ionOp, bb); + cfg::Instruction *instPtr = regionInst.get(); + bb->addInstruction(std::move(regionInst)); + + // 添加到op到instruction的映射 + subGraph->addOpToInstruction(®ionOp, instPtr); + } + + // 连接基本块 + if (regionLastBB) { + subGraph->addEdge(regionLastBB, bb); + } + regionLastBB = bb; + } + + // 连接区域入口到子图入口 + if (regionEntryBB) { + subGraph->addEdge(subEntry, regionEntryBB); + } + // 连接区域出口到子图出口 + if (regionLastBB) { + subGraph->addEdge(regionLastBB, subExit); + } + } + } + + // 设置子图 + inst->setSubGraph(std::move(subGraph)); + } else { + // 普通操作，直接添加到当前 basic block + createInstruction(&op, currentBB, cfg); + } + } + + return currentBB; +} + +cfg::BasicBlock * +ControlFlowGraphBuilder::handleIfOp(scf::IfOp ifOp, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *ifCondBB, + cfg::BasicBlock *parentStructure) { + // 创建 if 后面的汇合块 + auto *mergeBB = cfg.createBasicBlock(BlockType::NORMAL, parentStructure); + + // 设置 ifCondBB 的出口块为 mergeBB + ifCondBB->setExitBlock(mergeBB); + + // 处理 then 分支 + cfg::BasicBlock *thenExitBB = nullptr; + if (!ifOp.getThenRegion().empty()) { + // 创建 then 区域的入口块 + auto *thenEntryBB = cfg.createBasicBlock(BlockType::NORMAL, ifCondBB); + cfg.addEdge(ifCondBB, thenEntryBB); + + // 构建 then 区域的 CFG + auto result = + buildForRegion(ifOp.getThenRegion(), cfg, thenEntryBB, ifCondBB); + thenExitBB = result.exitBlock; + } + + // 处理 else 分支 + cfg::BasicBlock *elseExitBB = nullptr; + bool hasElse = !ifOp.getElseRegion().empty(); + + if (hasElse) { + // 创建 else 区域的入口块 + auto *elseEntryBB = cfg.createBasicBlock(BlockType::NORMAL, ifCondBB); + cfg.addEdge(ifCondBB, elseEntryBB); + + // 构建 else 区域的 CFG + auto result = + buildForRegion(ifOp.getElseRegion(), cfg, elseEntryBB, ifCondBB); + elseExitBB = result.exitBlock; + } + + // 连接 then 分支到汇合块 + if (thenExitBB) { + cfg.addEdge(thenExitBB, mergeBB); + } else { + // 空的 then 分支，直接从 ifCondBB 连接到 mergeBB + cfg.addEdge(ifCondBB, mergeBB); + } + + // 连接 else 分支到汇合块 + if (hasElse) { + if (elseExitBB) { + cfg.addEdge(elseExitBB, mergeBB); + } else { + // 空的 else 分支，直接从 ifCondBB 连接到 mergeBB + cfg.addEdge(ifCondBB, mergeBB); + } + } else { + // 没有 else 分支，ifCondBB 直接连接到 mergeBB（else 路径） + cfg.addEdge(ifCondBB, mergeBB); + } + + return mergeBB; +} + +cfg::BasicBlock *ControlFlowGraphBuilder::handleForOp( + scf::ForOp forOp, cfg::ControlFlowGraph &cfg, cfg::BasicBlock *forCondBB, + cfg::BasicBlock *parentStructure) { + // 创建循环体入口块 + auto *loopBodyEntryBB = cfg.createBasicBlock(BlockType::LOOP_BODY, forCondBB); + cfg.addEdge(forCondBB, loopBodyEntryBB); + + // 创建循环出口块 + auto *loopExitBB = + cfg.createBasicBlock(BlockType::LOOP_EXIT, parentStructure); + + // 设置 forCondBB 的出口块为 loopExitBB + forCondBB->setExitBlock(loopExitBB); + + // 构建循环体的 CFG + auto result = + buildForRegion(forOp.getRegion(), cfg, loopBodyEntryBB, forCondBB); + + // 循环体结束需要回到循环头（通过 yield 操作） + if (result.exitBlock) { + cfg.addEdge(result.exitBlock, forCondBB); + } + + // 循环出口 + cfg.addEdge(forCondBB, loopExitBB); + + return loopExitBB; +} + +cfg::BasicBlock *ControlFlowGraphBuilder::handleWhileOp( + scf::WhileOp whileOp, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *whileCondBB, cfg::BasicBlock *parentStructure) { + // while 操作有两个区域：before（条件）和 after（循环体） + // 控制流： + // whileCondBB (包含 scf.while 指令) + // ↓ + // beforeEntryBB (条件计算区域) + // ↓ + // scf.condition 分支：真 → afterEntryBB, 假 → loopExitBB + // ↓ + // afterEntryBB (循环体) + // ↓ + // scf.yield + // ↓ + // whileCondBB (回到 while 头，重新进入 before) + + // 创建 before 区域的入口块（条件计算） + auto *beforeEntryBB = cfg.createBasicBlock(BlockType::LOOP_BODY, whileCondBB); + cfg.addEdge(whileCondBB, beforeEntryBB); + + // 创建 after 区域的入口块（循环体） + auto *afterEntryBB = cfg.createBasicBlock(BlockType::LOOP_BODY, whileCondBB); + + // 创建循环出口块 + auto *loopExitBB = + cfg.createBasicBlock(BlockType::LOOP_EXIT, parentStructure); + + // 设置 whileCondBB 的出口块为 loopExitBB + whileCondBB->setExitBlock(loopExitBB); + + // 构建 before 区域的 CFG（条件计算区域） + // before 区域以一个 scf.condition 操作结束 + auto beforeResult = + buildForRegion(whileOp.getBefore(), cfg, beforeEntryBB, whileCondBB); + + // 构建 after 区域的 CFG（循环体区域） + // after 区域以 scf.yield 结束，yield 的参数会传递给 before 区域的参数 + auto afterResult = + buildForRegion(whileOp.getAfter(), cfg, afterEntryBB, whileCondBB); + + // 处理 before 区域结束后的分支 + // before 区域应该以一个 scf.condition 操作结束 + // 该操作决定是进入 after 区域还是退出循环 + if (beforeResult.exitBlock) { + // 从 before 出口连接到 after 入口（条件为真时） + cfg.addEdge(beforeResult.exitBlock, afterEntryBB); + + // 从 before 出口连接到循环出口（条件为假时） + // 注意：在实际的 scf.condition 中，条件为假会直接退出循环 + cfg.addEdge(beforeResult.exitBlock, loopExitBB); + } + + // after 区域结束后回到 whileCondBB（重新进入 before 区域进行条件检查） + if (afterResult.exitBlock) { + cfg.addEdge(afterResult.exitBlock, whileCondBB); + } + + return loopExitBB; +} + +cfg::Instruction *ControlFlowGraphBuilder::createInstruction( + Operation *op, cfg::BasicBlock *parentBlock, cfg::ControlFlowGraph &cfg) { + if (!op || !parentBlock) + return nullptr; + + auto inst = std::make_unique(getNextInstructionId(), op, + parentBlock); + cfg::Instruction *instPtr = inst.get(); + parentBlock->addInstruction(std::move(inst)); + + // 将 Operation 到 Instruction 的映射添加到 CFG + cfg.addOpToInstruction(op, instPtr); + + return instPtr; +} + +//===----------------------------------------------------------------------===// +// ControlFlowGraphBuilder Implementation +//===----------------------------------------------------------------------===// + +std::unique_ptr +ControlFlowGraphBuilder::build(triton::FuncOp func) { + auto cfg = std::make_unique(func); + + if (func.getBody().empty()) { + // 空函数，只创建入口和出口 + auto *entry = cfg->createBasicBlock(BlockType::ENTRY); + auto *exit = cfg->createBasicBlock(BlockType::EXIT); + cfg->addEdge(entry, exit); + return cfg; + } + + // 创建入口块 + auto *entryBlock = cfg->createBasicBlock(BlockType::ENTRY); + cfg->setEntryBlock(entryBlock); + + // 创建出口块 + auto *exitBlock = cfg->createBasicBlock(BlockType::EXIT); + cfg->setExitBlock(exitBlock); + + // 为函数体构建 CFG + auto result = buildForRegion(func.getBody(), *cfg, entryBlock, nullptr); + + // 连接函数体出口到函数出口 + if (result.exitBlock) { + cfg->addEdge(result.exitBlock, exitBlock); + } + + return cfg; +} + +std::vector> +ControlFlowGraphBuilder::buildForModule(ModuleOp module) { + std::vector> cfgs; + + for (triton::FuncOp func : module.getOps()) { + auto cfg = build(func); + if (cfg) { + cfgs.push_back(std::move(cfg)); + } + } + + return cfgs; +} + +cfg::BasicBlock *ControlFlowGraphBuilder::handleCondBranchOp( + cf::CondBranchOp condBrOp, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *condBrBB, cfg::BasicBlock *parentStructure) { + // 创建汇合块（用于 cond_br 之后的代码） + auto *mergeBB = cfg.createBasicBlock(BlockType::NORMAL, parentStructure); + + // 设置 condBrBB 的出口块为 mergeBB + condBrBB->setExitBlock(mergeBB); + + // 获取条件值 + Value condition = condBrOp.getCondition(); + + // 获取 true 分支的目标块和参数 + Block *trueDest = condBrOp.getTrueDest(); + SmallVector trueOperands(condBrOp.getTrueDestOperands()); + + // 获取 false 分支的目标块和参数 + Block *falseDest = condBrOp.getFalseDest(); + SmallVector falseOperands(condBrOp.getFalseDestOperands()); + + LLVM_DEBUG(llvm::dbgs() << " CondBr: condition=" << condition << "\n"); + LLVM_DEBUG(llvm::dbgs() << " True dest: " << trueDest << "\n"); + LLVM_DEBUG(llvm::dbgs() << " False dest: " << falseDest << "\n"); + + // 为 true 分支创建入口块（如果目标块还没有对应的 BasicBlock） + cfg::BasicBlock *trueEntryBB = + getOrCreateBasicBlockForBlock(trueDest, cfg, parentStructure); + + // 为 false 分支创建入口块 + cfg::BasicBlock *falseEntryBB = + getOrCreateBasicBlockForBlock(falseDest, cfg, parentStructure); + + // 连接 COND_BR 块到两个分支 + cfg.addEdge(condBrBB, trueEntryBB); + cfg.addEdge(condBrBB, falseEntryBB); + + // 存储分支信息到指令的 MemorySSAInfo 中（用于后续查询） + if (condBrBB->getNumInstructions() > 0) { + cfg::Instruction *inst = condBrBB->getInstruction(0); + // 可以通过 inst->getMemorySSAInfo() 存储额外信息 + } + + // 返回汇合块，后续代码将在此块中继续 + return mergeBB; +} + +cfg::BasicBlock *ControlFlowGraphBuilder::handleBranchOp( + cf::BranchOp brOp, cfg::ControlFlowGraph &cfg, cfg::BasicBlock *brBB, + cfg::BasicBlock *parentStructure) { + // 无条件跳转没有汇合块，直接连接到目标块 + + // 获取目标块和参数 + Block *dest = brOp.getDest(); + SmallVector destOperands(brOp.getDestOperands()); + + LLVM_DEBUG(llvm::dbgs() << " Br: unconditional branch\n"); + LLVM_DEBUG(llvm::dbgs() << " Dest: " << dest << "\n"); + + // 获取或创建目标块对应的 BasicBlock + cfg::BasicBlock *destBB = + getOrCreateBasicBlockForBlock(dest, cfg, parentStructure); + + // 连接 BR 块到目标块 + cfg.addEdge(brBB, destBB); + + // 无条件跳转没有后续代码，返回 nullptr 表示当前路径结束 + return nullptr; +} + +cfg::BasicBlock *ControlFlowGraphBuilder::getOrCreateBasicBlockForBlock( + Block *block, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *parentStructure) { + // 检查是否已经有对应的 BasicBlock + auto it = blockToBasicBlockMap.find(block); + if (it != blockToBasicBlockMap.end()) { + return it->second; + } + + // 创建新的 BasicBlock + auto *bb = cfg.createBasicBlock(BlockType::NORMAL, parentStructure); + + // 注册映射关系 + registerBlockMapping(block, bb); + + return bb; +} + +void ControlFlowGraphBuilder::registerBlockMapping(Block *mlirBlock, + cfg::BasicBlock *cfgBlock) { + blockToBasicBlockMap[mlirBlock] = cfgBlock; +} + +SmallVector +ControlFlowGraphBuilder::collectCondBrBlocks(cfg::ControlFlowGraph &cfg) { + SmallVector condBrBlocks; + + // 遍历 CFG 中的所有基本块 + for (size_t i = 0; i < cfg.getNumBlocks(); ++i) { + cfg::BasicBlock *bb = cfg.getBasicBlock(i); + if (bb && bb->getType() == BlockType::COND_BR) { + condBrBlocks.push_back(bb); + } + } + + return condBrBlocks; +} + +std::optional +ControlFlowGraphBuilder::getCondBranchMapping(cfg::BasicBlock *condBrBB) { + // 验证输入基本块类型 + if (!condBrBB || condBrBB->getType() != BlockType::COND_BR) { + return std::nullopt; + } + + // 获取 COND_BR 块中的指令（应该包含 cf.cond_br 操作） + if (condBrBB->getNumInstructions() == 0) { + return std::nullopt; + } + + cfg::Instruction *inst = condBrBB->getInstruction(0); + Operation *op = inst->getOperation(); + + // 确保是 cf.cond_br 操作 + auto condBrOp = dyn_cast(op); + if (!condBrOp) { + return std::nullopt; + } + + CondBranchMapping mapping; + + // 收集条件值 + mapping.condition = condBrOp.getCondition(); + + // 收集 true 分支信息 + mapping.trueDest = condBrOp.getTrueDest(); + for (Value operand : condBrOp.getTrueDestOperands()) { + mapping.trueOperands.push_back(operand); + } + + // 收集 false 分支信息 + mapping.falseDest = condBrOp.getFalseDest(); + for (Value operand : condBrOp.getFalseDestOperands()) { + mapping.falseOperands.push_back(operand); + } + + return mapping; +} + +SmallVector +ControlFlowGraphBuilder::collectBrBlocks(cfg::ControlFlowGraph &cfg) { + SmallVector brBlocks; + + // 遍历 CFG 中的所有基本块 + for (size_t i = 0; i < cfg.getNumBlocks(); ++i) { + cfg::BasicBlock *bb = cfg.getBasicBlock(i); + if (bb && bb->getType() == BlockType::BR) { + brBlocks.push_back(bb); + } + } + + return brBlocks; +} + +std::optional +ControlFlowGraphBuilder::getBranchMapping(cfg::BasicBlock *brBB) { + // 验证输入基本块类型 + if (!brBB || brBB->getType() != BlockType::BR) { + return std::nullopt; + } + + // 获取 BR 块中的指令（应该包含 cf.br 操作） + if (brBB->getNumInstructions() == 0) { + return std::nullopt; + } + + cfg::Instruction *inst = brBB->getInstruction(0); + Operation *op = inst->getOperation(); + + // 确保是 cf.br 操作 + auto brOp = dyn_cast(op); + if (!brOp) { + return std::nullopt; + } + + BranchMapping mapping; + + // 收集目标块信息 + mapping.dest = brOp.getDest(); + for (Value operand : brOp.getDestOperands()) { + mapping.destOperands.push_back(operand); + } + + return mapping; +} + +std::unique_ptr> +mlir::triton::cfg::createBuildCFGPass() { + return std::unique_ptr>(new BuildCFGPass()); +} + +//===----------------------------------------------------------------------===// +// ControlFlowGraphBuilderWithQueries Implementation +//===----------------------------------------------------------------------===// + +SmallVector +ControlFlowGraphBuilder::collectIfCondBlocks(cfg::ControlFlowGraph &cfg) { + SmallVector ifCondBlocks; + + // 遍历 CFG 中的所有基本块 + for (size_t i = 0; i < cfg.getNumBlocks(); ++i) { + cfg::BasicBlock *bb = cfg.getBasicBlock(i); + if (bb && bb->getType() == BlockType::IF_COND) { + ifCondBlocks.push_back(bb); + } + } + + return ifCondBlocks; +} + +SmallVector +ControlFlowGraphBuilder::collectForCondBlocks(cfg::ControlFlowGraph &cfg) { + SmallVector forCondBlocks; + + // 遍历 CFG 中的所有基本块 + for (size_t i = 0; i < cfg.getNumBlocks(); ++i) { + cfg::BasicBlock *bb = cfg.getBasicBlock(i); + if (bb && bb->getType() == BlockType::FOR_COND) { + forCondBlocks.push_back(bb); + } + } + + return forCondBlocks; +} + +std::optional +ControlFlowGraphBuilder::getIfYieldResultMapping(cfg::BasicBlock *ifCondBB) { + // 验证输入基本块类型 + if (!ifCondBB || ifCondBB->getType() != BlockType::IF_COND) { + return std::nullopt; + } + + // 获取 IF_COND 块中的指令（应该包含 scf.if 操作） + if (ifCondBB->getNumInstructions() == 0) { + return std::nullopt; + } + + cfg::Instruction *inst = ifCondBB->getInstruction(0); + Operation *op = inst->getOperation(); + + // 确保是 scf.if 操作 + auto ifOp = dyn_cast(op); + if (!ifOp) { + return std::nullopt; + } + + IfYieldResultMapping mapping; + + // 收集 result values + for (Value result : ifOp.getResults()) { + mapping.resultValues.push_back(result); + } + + // 处理 then 分支的 yield values + if (!ifOp.getThenRegion().empty()) { + Block &thenBlock = ifOp.getThenRegion().back(); + // 查找 then 区域末尾的 scf.yield 操作 + for (Operation &thenOp : thenBlock) { + if (auto yieldOp = dyn_cast(thenOp)) { + for (Value operand : yieldOp.getOperands()) { + mapping.thenYieldValues.push_back(operand); + } + break; + } + } + } + + // 处理 else 分支的 yield values（如果有） + if (!ifOp.getElseRegion().empty()) { + Block &elseBlock = ifOp.getElseRegion().back(); + // 查找 else 区域末尾的 scf.yield 操作 + for (Operation &elseOp : elseBlock) { + if (auto yieldOp = dyn_cast(elseOp)) { + for (Value operand : yieldOp.getOperands()) { + mapping.elseYieldValues.push_back(operand); + } + break; + } + } + } + + return mapping; +} + +std::optional +ControlFlowGraphBuilder::getForYieldIterArgMapping(cfg::BasicBlock *forCondBB) { + // 验证输入基本块类型 + if (!forCondBB || forCondBB->getType() != BlockType::FOR_COND) { + return std::nullopt; + } + + // 获取 FOR_COND 块中的指令（应该包含 scf.for 操作） + if (forCondBB->getNumInstructions() == 0) { + return std::nullopt; + } + + cfg::Instruction *inst = forCondBB->getInstruction(0); + Operation *op = inst->getOperation(); + + // 确保是 scf.for 操作 + auto forOp = dyn_cast(op); + if (!forOp) { + return std::nullopt; + } + + ForYieldIterArgMapping mapping; + + // 收集 iter_args（循环初始参数） + for (Value iterArg : forOp.getRegionIterArgs()) { + mapping.iterArgValues.push_back(iterArg); + } + + // 收集 result values + for (Value result : forOp.getResults()) { + mapping.resultValues.push_back(result); + } + + // 处理循环体的 yield values + if (!forOp.getRegion().empty()) { + Block &loopBlock = forOp.getRegion().back(); + // 查找循环体末尾的 scf.yield 操作 + for (Operation &loopOp : loopBlock) { + if (auto yieldOp = dyn_cast(loopOp)) { + for (Value operand : yieldOp.getOperands()) { + mapping.yieldValues.push_back(operand); + } + break; + } + } + } + + return mapping; +} diff --git a/third_party/ascend/lib/TritonToGraph/DataflowGraph.cpp b/third_party/ascend/lib/TritonToGraph/DataflowGraph.cpp new file mode 100644 index 0000000000..11972cd753 --- /dev/null +++ b/third_party/ascend/lib/TritonToGraph/DataflowGraph.cpp @@ -0,0 +1,305 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "TritonToGraph/DataflowGraph.h" +#include "TritonToGraph/AliasAnalysis.h" +#include "TritonToGraph/ControlFlowGraph.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/raw_ostream.h" + +#define DEBUG_TYPE "dataflow-graph" + +using namespace mlir; +using namespace triton; +using namespace cfg; + +//===----------------------------------------------------------------------===// +// DataFlowInfo +//===----------------------------------------------------------------------===// + +// Memory SSA相关接口 +MemorySSADef *DataFlowInfo::getMemoryDefinition(Value value) const { + auto it = memoryDefinitions.find(value); + return (it != memoryDefinitions.end()) ? it->second : nullptr; +} + +void DataFlowInfo::addMemoryDefinition(Value value, MemorySSADef *def) { + memoryDefinitions[value] = def; + // 确保没有旧的uses冲突 + memoryUses.erase(value); +} + +SmallVector DataFlowInfo::getMemoryUses(Value value) const { + auto it = memoryUses.find(value); + if (it != memoryUses.end()) { + return it->second; + } + return SmallVector(); +} + +void DataFlowInfo::addMemoryUse(Value value, const MemorySSAUse &use) { + memoryUses[value].push_back(use); +} + +void DataFlowInfo::removeMemoryDefinition(Value value) { + memoryDefinitions.erase(value); + memoryUses.erase(value); + invalidateDefUseCache(); +} + +void DataFlowInfo::clearMemoryUses(Value value) { + memoryUses[value].clear(); + invalidateDefUseCache(); +} + +std::unique_ptr DataFlowInfo::queryDataFlow(Value value) const { + // 1. 优先查询Memory SSA + if (MemorySSADef *def = getMemoryDefinition(value)) { + auto result = std::make_unique(def->getDefOp(), def); + result->getUses() = getSSAUses(value); + return result; + } + + // 2. 查询传统SSA + if (Operation *defOp = value.getDefiningOp()) { + auto result = std::make_unique(defOp, defOp); + result->getUses() = getSSAUses(value); + return result; + } + + // 3. 入参 + auto result = std::make_unique(nullptr, nullptr); + result->getUses() = getSSAUses(value); + return result; +} + +SmallVector DataFlowInfo::getUses(MemorySSADef *def) const { + SmallVector result; + + // 遍历所有uses，查找使用该definition的 + for (const auto &entry : memoryUses) { + for (const MemorySSAUse &use : entry.second) { + if (use.getDefinition() == def) { + result.push_back(use); + } + } + } + + return result; +} + +SmallVector +DataFlowInfo::getUsesByUserOp(Operation *userOp) const { + SmallVector result; + + for (const auto &entry : memoryUses) { + for (const MemorySSAUse &use : entry.second) { + if (use.getUserOp() == userOp) { + result.push_back(use); + } + } + } + + return result; +} + +void DataFlowInfo::buildDefUseCache() const { + if (defUseCacheValid) + return; + + defUseCache.clear(); + + for (const auto &entry : memoryUses) { + for (const MemorySSAUse &use : entry.second) { + MemorySSADef *def = use.getDefinition(); + if (def) { + defUseCache[def].push_back(use); + } + } + } + + defUseCacheValid = true; +} + +void DataFlowInfo::forEachDefinition( + llvm::function_ref func) const { + for (const auto &entry : memoryDefinitions) { + func(entry.first, entry.second); + } +} + +void DataFlowInfo::forEachUse( + llvm::function_ref func) const { + for (const auto &entry : memoryUses) { + for (const MemorySSAUse &use : entry.second) { + func(use); + } + } +} + +void DataFlowInfo::print(llvm::raw_ostream &os) const { + os << "=== Data Flow Information ===" << "\n"; + + os << "Memory Definitions: " << memoryDefinitions.size() << "\n"; + for (const auto &entry : memoryDefinitions) { + os << " " << entry.first << " -> "; + entry.second->print(os); + os << "\n"; + } + + os << "Memory Uses: " << "\n"; + for (const auto &entry : memoryUses) { + os << " " << entry.first << ": "; + for (const MemorySSAUse &use : entry.second) { + os << "[" << use.getDefinition()->getId() << "] "; + } + os << "\n"; + } + + os << "Phis: " << Phis.size() << "\n"; + for (const auto &entry : Phis) { + os << " " << entry.first << ": "; + switch (entry.second.type) { + case PhiInfo::ITER_ARG: + os << "ITER_ARG"; + break; + case PhiInfo::IF_RESULT: + os << "IF_RESULT"; + break; + case PhiInfo::WHILE_ARG: + os << "WHILE_ARG"; + break; + } + os << "\n"; + } +} + +void DataFlowInfo::exportToJSON(llvm::raw_ostream &os) const { + os << "{\n"; + os << " \"memoryDefinitions\": {\n"; + bool first = true; + for (const auto &entry : memoryDefinitions) { + if (!first) + os << ",\n"; + first = false; + os << " \"" << entry.first << "\": {\n"; + os << " \"id\": \"" << entry.second->getId() << "\",\n"; + os << " \"tensor\": \"" << entry.second->getTensor()->getName() + << "\",\n"; + os << " \"version\": " << entry.second->getVersion(); + os << "\n }"; + } + os << "\n },\n"; + + os << " \"Phis\": {\n"; + first = true; + for (const auto &entry : Phis) { + if (!first) + os << ",\n"; + first = false; + os << " \"" << entry.first << "\": {\n"; + os << " \"type\": " << entry.second.type << "\n"; + os << " }"; + } + os << "\n }\n"; + os << "}\n"; +} + +//===----------------------------------------------------------------------===// +// DataFlowGraph +//===----------------------------------------------------------------------===// + +void DataFlowGraph::build() { + LLVM_DEBUG(llvm::dbgs() << "=== Starting Data Flow Graph Build ===" << "\n"); + + // 步骤1: 构建Alias分析 + aliasAnalysis = std::make_unique(); + aliasAnalysis->analyzePointerAliases(cfg); + // aliasAnalysis->print(llvm::outs()); + + LLVM_DEBUG(llvm::dbgs() << "Alias analysis complete" << "\n"); + + // 步骤2: 构建Memory SSA + memorySSABuilder = + std::make_unique(cfg, *aliasAnalysis, dataFlowInfo); + memorySSABuilder->build(); + + LLVM_DEBUG(llvm::dbgs() << "Memory SSA build complete" << "\n"); + + // 步骤3: 构建def-use图 + buildDefUseGraph(); + + LLVM_DEBUG(llvm::dbgs() << "=== Data Flow Graph Build Complete ===" << "\n"); +} + +void DataFlowGraph::buildDefUseGraph() { + // 构建def-use图（在DataFlowInfo中实现） + dataFlowInfo.buildDefUseCache(); + + LLVM_DEBUG(llvm::dbgs() << "Def-use graph built" << "\n"); +} + +void DataFlowGraph::print(llvm::raw_ostream &os) const { + os << "=== Data Flow Graph ===" << "\n"; + dataFlowInfo.print(os); +} + +void DataFlowGraph::dump() const { print(llvm::errs()); } + +void DataFlowGraph::exportToJSON(llvm::raw_ostream &os) const { + auto funcOp = cfg.getFunction(); + auto funcName = funcOp.getName(); + std::string funcNameStr = funcName.empty() ? "unnamed" : funcName.str(); + + os << "{\n"; + os << " \"function\": \"" << funcNameStr << "\"," << "\n"; + os << " \"dataFlow\": "; + dataFlowInfo.exportToJSON(os); + os << "}\n"; +} + +void DataFlowGraph::exportDefUseToDOT(llvm::raw_ostream &os) const { + os << "digraph DefUseGraph {\n"; + os << " rankdir=TB;\n"; + os << " node [shape=box];\n\n"; + + // 遍历所有definitions并导出 + size_t nodeId = 0; + DenseMap defToNode; + + cfg.traverse([&](const BasicBlock &bb) { + for (const auto &instPtr : bb.getInstructions()) { + const Instruction *inst = instPtr.get(); + const MemorySSAInfo &ssaInfo = inst->getMemorySSAInfo(); + + for (MemorySSADef *def : ssaInfo.definitions) { + if (def && !defToNode.count(def)) { + defToNode[def] = nodeId++; + os << " node_" << defToNode[def] << " [label=\"" << def->getId() + << "\"];\n"; + } + } + } + }); + + os << "}\n"; +} diff --git a/third_party/ascend/lib/TritonToGraph/GraphAnalysis.cpp b/third_party/ascend/lib/TritonToGraph/GraphAnalysis.cpp new file mode 100644 index 0000000000..6b2cc3be72 --- /dev/null +++ b/third_party/ascend/lib/TritonToGraph/GraphAnalysis.cpp @@ -0,0 +1,627 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + */ + +#include "TritonToGraph/GraphAnalysis.h" +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "graph-analysis" + +using namespace mlir; +using namespace triton; +using namespace cfg; + +//===----------------------------------------------------------------------===// +// CFGTraverser Implementation +//===----------------------------------------------------------------------===// + +void CFGTraverser::dfsForward(CFGTraversalBase &visitor) { + DenseSet visited; + TraversalContext ctx; + dfsForwardImpl(cfg.getEntryBlock(), visited, ctx, visitor); +} + +void CFGTraverser::dfsForward(BasicBlock *start, CFGTraversalBase &visitor) { + DenseSet visited; + TraversalContext ctx; + dfsForwardImpl(start, visited, ctx, visitor); +} + +void CFGTraverser::dfsForwardImpl(BasicBlock *block, + DenseSet &visited, + TraversalContext &ctx, + CFGTraversalBase &visitor) { + if (!block || visited.contains(block)) + return; + + visited.insert(block); + + // pre-visit block + visitor.preVisitBlock(block, ctx); + + // visit instructions in block + for (auto &instPtr : block->getInstructions()) { + Instruction *inst = instPtr.get(); + + // check for structure entry + if (inst->hasSubGraph()) { + visitor.onEnterStructure(block, ctx); + ctx.push(block); + } + + visitor.VisitInstruction(inst, ctx); + + // check for structure exit (last instruction in structure block) + if (inst->hasSubGraph()) { + ctx.pop(); + visitor.onExitStructure(block, ctx); + } + } + + // visit successors + for (BasicBlock *succ : block->getSuccessors()) { + if (visited.contains(succ)) + continue; + + // check for back edge + if (cfg.isBackEdge(block, succ)) { + visitor.onBackEdge(block, succ, ctx); + } + + dfsForwardImpl(succ, visited, ctx, visitor); + } + + visitor.postVisitBlock(block, ctx); +} + +void CFGTraverser::dfsBackward(BasicBlock *start, CFGTraversalBase &visitor) { + DenseSet visited; + TraversalContext ctx; + dfsBackwardImpl(start, visited, ctx, visitor); +} + +void CFGTraverser::dfsBackwardImpl(BasicBlock *block, + DenseSet &visited, + TraversalContext &ctx, + CFGTraversalBase &visitor) { + if (!block || visited.contains(block)) + return; + + visited.insert(block); + + visitor.preVisitBlock(block, ctx); + + // visit instructions in reverse order + auto &insts = block->getInstructions(); + for (auto it = insts.rbegin(); it != insts.rend(); ++it) { + Instruction *inst = it->get(); + visitor.VisitInstruction(inst, ctx); + } + + // visit predecessors + for (BasicBlock *pred : block->getPredecessors()) { + if (cfg.isBackEdge(pred, block)) { + visitor.onBackEdge(pred, block, ctx); + } + dfsBackwardImpl(pred, visited, ctx, visitor); + } + + visitor.postVisitBlock(block, ctx); +} + +void CFGTraverser::bfsForward(CFGTraversalBase &visitor) { + bfsForward(cfg.getEntryBlock(), visitor); +} + +void CFGTraverser::bfsForward(BasicBlock *start, CFGTraversalBase &visitor) { + DenseSet visited; + SmallVector> worklist; + + worklist.push_back({start, TraversalContext()}); + visited.insert(start); + + while (!worklist.empty()) { + auto [block, ctx] = worklist.pop_back_val(); + + preVisitBlock(block, const_cast(ctx)); + + for (auto &instPtr : block->getInstructions()) { + Instruction *inst = instPtr.get(); + visitor.VisitInstruction(inst, const_cast(ctx)); + } + + visitor.postVisitBlock(block, const_cast(ctx)); + + for (BasicBlock *succ : block->getSuccessors()) { + if (!visited.contains(succ)) { + visited.insert(succ); + worklist.push_back({succ, ctx}); + } + } + } +} + +void CFGTraverser::bfsBackward(BasicBlock *start, CFGTraversalBase &visitor) { + DenseSet visited; + SmallVector> worklist; + + worklist.push_back({start, TraversalContext()}); + visited.insert(start); + + while (!worklist.empty()) { + auto [block, ctx] = worklist.pop_back_val(); + + visitor.preVisitBlock(block, const_cast(ctx)) + + for (auto it = insts.rbegin(); it != insts.rend(); ++it) { + Instruction *inst = it->get(); + visitor.VisitInstruction(inst, ctx); + } + + visitor.postVisitBlock(block, const_cast(ctx)); + + for (BasicBlock *pred : block->getPredecessors()) { + if (!visited.contains(pred)) { + visited.insert(pred); + worklist.push_back({pred, ctx}); + } + } + } +} + +//===----------------------------------------------------------------------===// +// DFGTraverser Implementation +//===----------------------------------------------------------------------===// + +void DFGTraverser::dfsBackward(Value seed, DFGTraversalBase &visitor, + const Options &opts) { + DenseSet visited; + dfsBackwardImpl(seed, visitor, visited, pts, 0); +} + +void DFGTraverser::dfsBackward(ArrayRef seeds, DFGTraversalBase &visitor, + const Options &opts) { + DenseSet visited; + for (Value seed : seeds) { + dfsBackwardImpl(seed, visitor, visited, opts, 0); + } +} + +void DFGTraverser::dfsBackwardImpl(Value value, DFGTraversalBase &visitor, + DenseSet &visited, + const Options &opts, int depth) { + if (opts.maxDepth >= 0 && depth > opts.maxDepth) + return; + + // get definition + Operation *defOp = nullptr; + if (opts.useMemorySSA) { + auto result = dfg.queryDataFlow(value); + if (auto *memResult = dyn_cast(result.get())) { + defOp = memResult->getDefinition()->getDefOp(); + } + } else { + defOp = value.getDefiningOp(); + } + + if (!defOp) + return; + + if (visited.contains(defOp)) + return; + + if (opts.stopOps.contains(defOp)) + return; + + visited.insert(defOp); + + visitor.VisitDef(value, defOp, depth); + + // recursively visit operands + for (Value operand : defOp->getOperands()) { + dfsBackwardImpl(operand, visitor, visited, opts, depth + 1); + } + + // handle phi/iter_arg + if (opts.followPhi) { + auto &dataFlowInfo = dfg.getDataFlowInfo(); + if (dataFlowInfo.hasPhi(value)) { + auto &phiInfo = dataFlowInfo.getPhi(value); + visitor.onPhi(value, phiInfo, depth); + } + } +} + +void DFGTraverser::dfsForward(Value seed, DFGTraversalBase &visitor, + const Options &opts) { + DenseSet visited; + dfsForwardImpl(seed, visitor, visited, ctx, opts, 0); +} + +void DFGTraverser::dfsForwardImpl(Value value, DFGTraversalBase &visitor, + DenseSet &visited, + const Options &opts, int depth) { + if (opts.maxDepth >= 0 && depth > opts.maxDepth) + return; + + SmallVector uses; + if (opts.useMemorySSA) { + uses = dfg.getDataFlowInfo().getSSAUses(value); + } else { + for (OpOperand &use : value.getUses()) { + uses.push_back(&use); + } + } + + for (OpOperand *use : uses) { + Operation *userOp = use->getOwner(); + + if (visited.contains(userOp)) + continue; + + if (opts.stopOps.contains(userOp)) + continue; + + visited.insert(userOp); + + visitor.VisitUse(value, use, depth); + + // visit results of this operation + for (Value result : userOp->getResults()) { + dfsForwardImpl(result, visitor, visited, opts, depth + 1); + } + } +} + +//===----------------------------------------------------------------------===// +// Region Implementation +//===----------------------------------------------------------------------===// + +void Region::add(Instruction *inst) { instSet_.insert(inst); } + +void Region::add(Operation *op, ControlFlowGraph &cfg) { + if (Instruction *inst = cfg.getInstruction(op)) { + add(inst); + } +} + +void Region::addAll(ArrayRef insts) { + for (Instruction *inst : insts) { + add(inst); + } +} + +bool Region::contains(Instruction *inst) const { + return instSet_.contains(inst); +} + +bool Region::contains(Operation *op) const { + // note: requires cfg to be available - caller must ensure this + // use the version that takes cfg as parameter for proper lookup + return false; +} + +void Region::remove(Instruction *inst) { instSet_.erase(inst); } + +void Region::clear() { instSet_.clear(); } + +SmallVector Region::orderedInstructions() const { + SmallVector result(instSet_.begin(), instSet_.end()); + // sort by block id then instruction index + llvm::sort(result, [](Instruction *a, Instruction *b) { + auto *bbA = a->getParentBlock(); + auto *bbB = b->getParentBlock(); + if (bbA != bbB) + return bbA->getId() < bbB->getId(); + // within same block, find position + // this is approximate - exact order requires full block scan + return a->getId() < b->getId(); + }); + return result; +} + +SmallVector Region::operations() const { + SmallVector result; + for (Instruction *inst : instSet_) { + if (Operation *op = inst->getOperation()) { + result.push_back(op); + } + } + return result; +} + +//===----------------------------------------------------------------------===// +// RegionAnalyzer Implementation +//===----------------------------------------------------------------------===// + +bool RegionAnalyzer::hasDependency(const Region &from, const Region &to) const { + auto deps = getDependencies(from, to); + return !deps.empty(); +} + +SmallVector +RegionAnalyzer::getDependencies(const Region &from, const Region &to) const { + SmallVector deps; + + // check for data dependencies: from defines, to uses + for (Instruction *fromInst : from) { + for (Value result : fromInst->getOperation()->getResults()) { + for (OpOperand &use : result.getUses()) { + Operation *userOp = use.getOwner(); + if (Instruction *toInst = cfg.getInstruction(userOp)) { + if (to.contains(toInst)) { + deps.push_back({Dependency::DATA, result, fromInst, toInst}); + } + } + } + } + } + + return deps; +} + +RegionAnalyzer::ExternalDeps +RegionAnalyzer::analyzeExternalDeps(const Region ®ion) const { + ExternalDeps result; + + // find inputs: external definitions used inside region + for (Instruction *inst : region) { + for (Value operand : inst->getOperation()->getOperands()) { + Operation *defOp = operand.getDefiningOp(); + if (!defOp) { + // block argument - treat as external input + bool alreadyTracked = false; + for (auto &input : result.inputs) { + if (input.value == operand) { + input.internalUses.push_back(inst); + alreadyTracked = true; + break; + } + } + if (!alreadyTracked) { + result.inputs.push_back({operand, nullptr, {inst}}); + } + } else if (Instruction *defInst = cfg.getInstruction(defOp)) { + if (!region.contains(defInst)) { + // external definition + bool alreadyTracked = false; + for (auto &input : result.inputs) { + if (input.value == operand) { + input.internalUses.push_back(inst); + alreadyTracked = true; + break; + } + } + if (!alreadyTracked) { + result.inputs.push_back({operand, defInst, {inst}}); + } + } + } + } + } + + // find outputs: internal definitions used outside region + for (Instruction *inst : region) { + for (Value result : inst->getOperation()->getResults()) { + SmallVector externalUses; + for (OpOperand &use : result.getUses()) { + Operation *userOp = use.getOwner(); + if (Instruction *userInst = cfg.getInstruction(userOp)) { + if (!region.contains(userInst)) { + externalUses.push_back(userInst); + } + } + } + if (!externalUses.empty()) { + result.outputs.push_back({result, inst, externalUses}); + } + } + } + + return result; +} + +//===----------------------------------------------------------------------===// +// ProgramSlicer Implementation +//===----------------------------------------------------------------------===// + +ProgramSlice ProgramSlicer::compute(const SliceCriterion &criterion) { + ProgramSlice slice; + + DFGTraverser dfgTraverser(dfg); + + for (Value seed : criterion.seeds) { + class SliceBuilder : public DFGTraversalBase { + public: + SliceBuilder(ProgramSlice &slice, ControlFlowGraph &cfg) + : slice(slice), cfg(cfg) {} + + bool VisitDef(Value value, Operation *defOp, int depth) override { + if (Instruction *inst = cfg.getInstruction(defOp)) { + slice.add(inst); + } + return true; + } + + bool VisitUse(Value value, OpOperand *use, int depth) override { + Operation *userOp = use->getOwner(); + if (Instruction *inst = cfg.getInstruction(userOp)) { + slice.add(inst); + } + return true; + } + + private: + ProgramSlice &slice; + ControlFlowGraph &cfg; + }; + + SliceBuilder builder(slice, cfg); + + switch (criterion.dir) { + case SliceCriterion::BACKWARD: + dfgTraverser.dfsBackward(seed, builder, criterion.dfgOpts); + break; + case SliceCriterion::FORWARD: + dfgTraverser.dfsForward(seed, builder, criterion.dfgOpts); + break; + case SliceCriterion::BIDIRECTIONAL: + dfgTraverser.dfsBackward(seed, builder, criterion.dfgOpts); + dfgTraverser.dfsForward(seed, builder, criterion.dfgOpts); + break; + } + } + + return slice; +} + +ProgramSlice ProgramSlicer::sliceFromYields(ArrayRef yields, + SliceCriterion::Direction dir) { + SliceCriterion criterion; + criterion.seeds = SmallVector(yields); + criterion.dir = dir; + return compute(criterion); +} + +void ProgramSlice::merge(const ProgramSlice &other) { + for (Instruction *inst : other) { + instructions_.insert(inst); + } +} + +void ProgramSlice::intersect(const ProgramSlice &other) { + DenseSet toRemove; + for (Instruction *inst : instructions_) { + if (!other.contains(inst)) { + toRemove.insert(inst); + } + } + for (Instruction *inst : toRemove) { + instructions_.erase(inst); + } +} + +void ProgramSlice::subtract(const ProgramSlice &other) { + for (Instruction *inst : other) { + instructions_.erase(inst); + } +} + +Region ProgramSlice::toRegion(StringRef name) const { + Region region(name); + for (Instruction *inst : instructions_) { + region.add(inst); + } + return region; +} + +//===----------------------------------------------------------------------===// +// RegionAbsorber Implementation +//===----------------------------------------------------------------------===// + +void RegionAbsorber::absorb(Region ®ion, ArrayRef seeds, + const AbsorptionPolicy &policy) { + DenseSet visited; + + for (Instruction *seed : seeds) { + region.add(seed); + + if (policy.dir == AbsorptionPolicy::UPSTREAM || + policy.dir == AbsorptionPolicy::BOTH) { + absorbUpstream(region, seed, policy, visited, 0); + } + + if (policy.dir == AbsorptionPolicy::DOWNSTREAM || + policy.dir == AbsorptionPolicy::BOTH) { + absorbDownstream(region, seed, policy, visited, 0); + } + } +} + +void RegionAbsorber::absorbUpstream(Region ®ion, Instruction *inst, + const AbsorptionPolicy &policy, + DenseSet &visited, + int depth) { + if (policy.maxDepth >= 0 && depth > policy.maxDepth) + return; + + if (visited.contains(inst)) + return; + visited.insert(inst); + + if (policy.shouldStop && policy.shouldStop(inst)) + return; + + if (policy.stopOps.contains(inst->getOperation())) + return; + + // add to region + region.add(inst); + + // visit operands + for (Value operand : inst->getOperation()->getOperands()) { + if (Operation *defOp = operand.getDefiningOp()) { + if (Instruction *defInst = cfg.getInstruction(defOp)) { + if (!policy.crossRegionBoundary && region.contains(defInst)) + continue; + + absorbUpstream(region, defInst, policy, visited, depth + 1); + } + } + } +} + +void RegionAbsorber::absorbDownstream(Region ®ion, Instruction *inst, + const AbsorptionPolicy &policy, + DenseSet &visited, + int depth) { + if (policy.maxDepth >= 0 && depth > policy.maxDepth) + return; + + if (visited.contains(inst)) + return; + visited.insert(inst); + + if (policy.shouldStop && policy.shouldStop(inst)) + return; + + if (policy.stopOps.contains(inst->getOperation())) + return; + + // add to region + region.add(inst); + + // visit uses + for (Value result : inst->getOperation()->getResults()) { + for (OpOperand &use : result.getUses()) { + Operation *userOp = use.getOwner(); + if (Instruction *userInst = cfg.getInstruction(userOp)) { + if (!policy.crossRegionBoundary && region.contains(userInst)) + continue; + + absorbDownstream(region, userInst, policy, visited, depth + 1); + } + } + } +} + +void RegionAbsorber::absorbFromValue(Region ®ion, Value value, + const AbsorptionPolicy &policy) { + Operation *defOp = value.getDefiningOp(); + if (!defOp) + return; + + Instruction *inst = cfg.getInstruction(defOp); + if (!inst) + return; + + absorb(region, {inst}, policy); +} + +void RegionAbsorber::absorbUntilBoundary( + Region ®ion, ArrayRef seeds, + std::function isBoundary) { + AbsorptionPolicy policy; + policy.shouldStop = isBoundary; + absorb(region, seeds, policy); +} diff --git a/third_party/ascend/lib/TritonToGraph/InterProceduralCFG.cpp b/third_party/ascend/lib/TritonToGraph/InterProceduralCFG.cpp new file mode 100644 index 0000000000..e9f8320e48 --- /dev/null +++ b/third_party/ascend/lib/TritonToGraph/InterProceduralCFG.cpp @@ -0,0 +1,351 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "TritonToGraph/InterProceduralCFG.h" +#include "TritonToGraph/ControlFlowGraphBuilder.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/raw_ostream.h" + +#define DEBUG_TYPE "icfg" + +using namespace mlir; +using namespace mlir::triton; + +namespace mlir { +namespace triton { +namespace cfg { + +InterProceduralCFG::InterProceduralCFG(ModuleOp module) : module(module) {} + +InterProceduralCFG::~InterProceduralCFG() = default; + +void InterProceduralCFG::build() { + LLVM_DEBUG(llvm::dbgs() << "Building ICFG for module\n"); + + ControlFlowGraphBuilder builder; + + // 遍历模块中的所有函数 + for (triton::FuncOp func : module.getOps()) { + LLVM_DEBUG(llvm::dbgs() + << "Building CFG for function: " << func.getName() << "\n"); + + // 为每个函数构建CFG + auto cfg = builder.build(func); + functionCFGs[func] = std::move(cfg); + } +} + +ControlFlowGraph *InterProceduralCFG::getFunctionCFG(triton::FuncOp func) { + auto it = functionCFGs.find(func); + if (it != functionCFGs.end()) { + return it->second.get(); + } + return nullptr; +} + +const ControlFlowGraph * +InterProceduralCFG::getFunctionCFG(triton::FuncOp func) const { + auto it = functionCFGs.find(func); + if (it != functionCFGs.end()) { + return it->second.get(); + } + return nullptr; +} + +ControlFlowGraph *InterProceduralCFG::getFunctionCFG(StringRef funcName) { + for (auto &[func, cfg] : functionCFGs) { + if (func.getName() == funcName) { + return cfg.get(); + } + } + return nullptr; +} + +void InterProceduralCFG::connectCallGraph() { + LLVM_DEBUG(llvm::dbgs() << "Connecting call graph\n"); + + // 遍历所有操作，查找调用点 + module.walk([&](Operation *op) { + // 检查是否是调用操作 + if (auto callInterface = dyn_cast(op)) { + // 尝试获取被调用函数 + auto callable = callInterface.getCallableForCallee(); + + // 从符号引用获取函数名 + StringRef calleeName; + if (auto symbolRef = callable.dyn_cast()) { + calleeName = symbolRef.getRootReference(); + } else if (auto flatSymbolRef = + op->getAttrOfType("callee")) { + calleeName = flatSymbolRef.getValue(); + } + + if (!calleeName.empty()) { + // 在模块中查找被调用函数 + if (auto calleeFunc = module.lookupSymbol(calleeName)) { + // 获取调用者的函数 + Operation *parentOp = op->getParentOp(); + while (parentOp && !isa(parentOp)) { + parentOp = parentOp->getParentOp(); + } + + if (auto callerFunc = dyn_cast(parentOp)) { + auto callerCFG = getFunctionCFG(callerFunc); + + if (callerCFG) { + // 找到调用点的基本块 + BasicBlock *callBlock = nullptr; + Instruction *callInst = nullptr; + + // 遍历查找包含该操作的 basic block 和 instruction + for (size_t i = 0; i < callerCFG->getNumBlocks(); ++i) { + auto *bb = callerCFG->getBasicBlock(i); + for (size_t j = 0; j < bb->getNumInstructions(); ++j) { + auto *inst = bb->getInstruction(j); + if (inst->getOperation() == op) { + callBlock = bb; + callInst = inst; + break; + } + } + if (callBlock) + break; + } + + if (callBlock) { + // 创建调用点记录 + CallSite callSite; + callSite.callOp = op; + callSite.caller = callerFunc; + callSite.callee = calleeFunc; + callSite.callBlock = callBlock; + callSite.callInst = callInst; + + callSites.push_back(callSite); + + // 添加到调用图 + callGraph[callerFunc].push_back(calleeFunc); + reverseCallGraph[calleeFunc].push_back(callerFunc); + + LLVM_DEBUG(llvm::dbgs() + << "Found call: " << callerFunc.getName() << " -> " + << calleeFunc.getName() << "\n"); + } + } + } + } + } + } + }); +} + +SmallVector +InterProceduralCFG::getCallees(triton::FuncOp caller) const { + auto it = callGraph.find(caller); + if (it != callGraph.end()) { + return it->second; + } + return SmallVector(); +} + +SmallVector +InterProceduralCFG::getCallers(triton::FuncOp callee) const { + auto it = reverseCallGraph.find(callee); + if (it != reverseCallGraph.end()) { + return it->second; + } + return SmallVector(); +} + +void InterProceduralCFG::computeReachability() { + LLVM_DEBUG(llvm::dbgs() << "Computing reachability\n"); + + // 使用简单的递归DFS计算可达性 + std::function &)> dfs = + [&](triton::FuncOp func, DenseSet &visited) { + if (visited.contains(func)) { + return; + } + visited.insert(func); + + for (auto callee : getCallees(func)) { + reachability[func].insert(callee); + dfs(callee, reachability[func]); + } + }; + + // 对每个函数计算可达性 + for (auto &[func, cfg] : functionCFGs) { + reachability[func].clear(); + dfs(func, reachability[func]); + } +} + +bool InterProceduralCFG::isReachable(triton::FuncOp from, + triton::FuncOp to) const { + auto it = reachability.find(from); + if (it != reachability.end()) { + return it->second.contains(to); + } + return false; +} + +void InterProceduralCFG::dumpToDot(const std::string &filename) const { + std::error_code ec; + llvm::raw_fd_ostream file(filename, ec); + + if (ec) { + llvm::errs() << "Failed to open file: " << filename << "\n"; + return; + } + + file << "digraph ICFG {\n"; + file << " node [shape=box];\n\n"; + + // 输出所有函数节点 + for (const auto &[func, cfg] : functionCFGs) { + StringRef name = const_cast(func).getName(); + file << " \"" << name << "\" [label=\"" << name << "\"];\n"; + } + + file << "\n"; + + // 输出调用边 + for (const auto &[caller, callees] : callGraph) { + StringRef callerName = const_cast(caller).getName(); + for (auto callee : callees) { + StringRef calleeName = const_cast(callee).getName(); + file << " \"" << callerName << "\" -> \"" << calleeName + << "\" [label=\"call\"];\n"; + } + } + + file << "}\n"; +} + +void InterProceduralCFG::print(raw_ostream &os) const { + os << "InterProcedural Control Flow Graph:\n"; + os << "Functions: " << functionCFGs.size() << "\n"; + os << "CallSites: " << callSites.size() << "\n"; + + os << "\nCall Graph:\n"; + for (const auto &[caller, callees] : callGraph) { + StringRef callerName = const_cast(caller).getName(); + os << " " << callerName << " -> "; + for (auto callee : callees) { + StringRef calleeName = const_cast(callee).getName(); + os << calleeName << " "; + } + os << "\n"; + } +} + +llvm::Error +InterProceduralCFG::exportToHTML(const std::string &filename) const { + std::error_code ec; + llvm::raw_fd_ostream os(filename, ec, llvm::sys::fs::OF_Text); + + if (ec) { + return llvm::createStringError(llvm::inconvertibleErrorCode(), + "Failed to open file: " + filename); + } + + os << R"html( + + + + Inter-Procedural CFG + + + + +

Inter-Procedural CFG

+
+ + + +)html"; + + os.close(); + return llvm::Error::success(); +} + +} // namespace cfg +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/TritonToGraph/MemorySsaBuilder.cpp b/third_party/ascend/lib/TritonToGraph/MemorySsaBuilder.cpp new file mode 100644 index 0000000000..fb5011e8aa --- /dev/null +++ b/third_party/ascend/lib/TritonToGraph/MemorySsaBuilder.cpp @@ -0,0 +1,625 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "TritonToGraph/MemorySsaBuilder.h" +#include "TritonToGraph/ControlFlowGraph.h" +#include "TritonToGraph/DataflowGraph.h" +#include "TritonToGraph/tensor.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "triton/Dialect/Triton/IR/OpInterfaces.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/raw_ostream.h" + +#define DEBUG_TYPE "memory-ssa-builder" + +using namespace mlir; +using namespace triton; +using namespace cfg; + +//===----------------------------------------------------------------------===// +// MemorySSABuilder +//===----------------------------------------------------------------------===// + +MemorySSABuilder::~MemorySSABuilder() { + // 清理所有创建的tensor objects（如果有动态分配的） + // 这里假设TensorObject由外部管理生命周期 +} + +void MemorySSABuilder::build() { + LLVM_DEBUG(llvm::dbgs() << "=== Starting Memory SSA Build ===\n"); + + // 步骤1: 拓扑排序，确定处理顺序 + std::vector topoOrder; + { + // 简单的拓扑排序实现 + DenseSet visited; + std::function dfs = [&](BasicBlock *bb) { + if (visited.contains(bb)) + return; + visited.insert(bb); + for (BasicBlock *succ : bb->getSuccessors()) { + dfs(succ); + } + topoOrder.push_back(bb); + }; + + // 从入口块开始DFS + dfs(cfg.getEntryBlock()); + std::reverse(topoOrder.begin(), topoOrder.end()); + } + + LLVM_DEBUG(llvm::dbgs() << "Topological order: " << topoOrder.size() + << " blocks\n"); + + // 步骤2: 初始化函数的参数 + createParameterDefinitions(); + + LLVM_DEBUG(llvm::dbgs() << "Created parameter definitions\n"); + + // 步骤3: 按拓扑序遍历每个BasicBlock + for (BasicBlock *bb : topoOrder) { + LLVM_DEBUG(llvm::dbgs() << "Processing BB" << bb->getId() << "\n"); + + // 处理block + processBasicBlock(bb); + } + + LLVM_DEBUG(llvm::dbgs() << "=== Memory SSA Build Complete ===\n" + << "Processed blocks: " << topoOrder.size() << "\n" + << "Total definitions: " << allDefinitions.size() + << "\n"); +} + +void MemorySSABuilder::createParameterDefinitions() { + triton::FuncOp func = cfg.getFunction(); + + LLVM_DEBUG(llvm::dbgs() << "Creating parameter definitions for " + << func.getName() << "\n"); + + // 遍历函数参数 + for (BlockArgument arg : func.getArguments()) { + Type argType = arg.getType(); + + // 检查是否是我们关心的类型 + if (isTensorType(argType)) { + // 创建参数名称 + std::string paramName = "param_" + std::to_string(arg.getArgNumber()); + + // 如果是指针类型，从aliasAnalysis获取tensor对象 + TensorObject *tensor = nullptr; + if (aliasAnalysis.isPointerType(argType)) { + tensor = aliasAnalysis.getTensorObject(arg); + } + + // 如果没有找到tensor对象，创建一个新的 + if (!tensor) { + // 从类型推断shape和element type + SmallVector shape; + Type elementType; + + extractShapeAndElementType(argType, shape, elementType); + + tensor = new TensorObject(paramName, shape, argType, elementType, + TensorObject::TensorKind::GLOBAL_MEMORY); + } + + tensor->print(llvm::outs()); + llvm::outs() << "\n"; + + // 为入参创建definition + MemorySSADef *def = createDefinition(tensor, nullptr); + + // 记录到dataFlowInfo + dataFlowInfo.addMemoryDefinition(arg, def); + + LLVM_DEBUG(llvm::dbgs() + << " Created parameter definition: " << def->getId() << "\n"); + } + } +} + +void MemorySSABuilder::processBasicBlock(BasicBlock *bb) { + if (!bb) + return; + + // 处理block内的所有指令 + for (auto &instPtr : bb->getInstructions()) { + Instruction *inst = instPtr.get(); + inst->print(llvm::outs()); + llvm::outs() << "\n"; + processInstruction(inst); + MemorySSAInfo &ssaInfo = inst->getMemorySSAInfo(); + ssaInfo.print(llvm::outs()); + llvm::outs() << "\n"; + } +} + +void MemorySSABuilder::processInstruction(Instruction *inst) { + if (!inst) + return; + + Operation *op = inst->getOperation(); + if (!op) + return; + + LLVM_DEBUG(llvm::dbgs() << "Processing: " << op->getName() << "\n"); + + MemorySSAInfo &ssaInfo = inst->getMemorySSAInfo(); + + // 1. 处理operands：创建uses + LLVM_DEBUG(llvm::dbgs() << " Processing operands...\n"); + + for (OpOperand &operand : op->getOpOperands()) { + Value operandValue = operand.get(); + unsigned operandIdx = operand.getOperandNumber(); + + // 检查是否是tensor或pointer类型 + if (isTensorType(operandValue.getType()) || + aliasAnalysis.isPointerType(operandValue.getType())) { + + // 查找operand的definition + MemorySSADef *def = dataFlowInfo.getMemoryDefinition(operandValue); + + if (def) { + // 创建use + MemorySSAUse use = createUse(def, op, operandIdx); + ssaInfo.uses.push_back(use); + + // 记录到全局map + dataFlowInfo.addMemoryUse(operandValue, use); + + LLVM_DEBUG(llvm::dbgs() + << " MemorySSAUse: " << def->getId() << " in " + << op->getName() << " [operand #" << operandIdx << "]\n"); + } + } + } + + // 2. 按操作类型处理：store、load、tensor writer、pointer op + LLVM_DEBUG(llvm::dbgs() << " Processing by operation type...\n"); + + // 对store操作（内存写入）- 为第一个operand指向的tensor创建新definition + if (isMemoryWriter(op)) { + LLVM_DEBUG(llvm::dbgs() << " Memory writer: " << op->getName() << "\n"); + + // store的第一个operand是ptr（第二个是value） + // 针对Triton算子场景做了简化，此处不是流敏感分析 + if (op->getNumOperands() >= 2) { + Value ptr = op->getOperand(0); + if (isTensorType(ptr.getType()) || + aliasAnalysis.isPointerType(ptr.getType())) { + // 获取ptr当前的definition（修改前的状态） + MemorySSADef *oldDef = dataFlowInfo.getMemoryDefinition(ptr); + if (oldDef) { + // store会修改内存，为同一个tensor创建新的definition + TensorObject *tensor = oldDef->getTensor(); + MemorySSADef *newDef = createDefinition(tensor, op); + + // store之后，ptr指向的内存状态改变，更新definition + dataFlowInfo.addMemoryDefinition(ptr, newDef); + + LLVM_DEBUG(llvm::dbgs() << " Store from: " << oldDef->getId() + << " to: " << newDef->getId() << "\n"); + } + } + } + } + // 对load操作（内存读取）- 第一个operand的definition作为result的definition + else if (isMemoryReader(op)) { + LLVM_DEBUG(llvm::dbgs() << " Memory reader: " << op->getName() << "\n"); + + if (op->getNumOperands() > 0 && op->getNumResults() > 0) { + Value ptr = op->getOperand(0); + Value result = op->getResult(0); + + // 获取ptr的definition + MemorySSADef *ptrDef = dataFlowInfo.getMemoryDefinition(ptr); + if (ptrDef) { + ssaInfo.definitions.push_back(ptrDef); + dataFlowInfo.addMemoryDefinition(result, ptrDef); + + LLVM_DEBUG(llvm::dbgs() + << " Load creates new def: " << ptrDef->getId() + << " from " << ptrDef->getId() << "\n"); + } + } + } + // 对返回新Tensor的操作（tensor writer） + else if (isTensorWriter(op)) { + LLVM_DEBUG(llvm::dbgs() << " Tensor writer: " << op->getName() << "\n"); + + for (Value result : op->getResults()) { + Type resultType = result.getType(); + + // 检查是否是tensor类型 + if (isTensorType(resultType)) { + // 创建tensor对象 + TensorObject *tensor = createTensorObject(op); + + // Tensor writer：创建新definition + MemorySSADef *newDef = createDefinition(tensor, op); + ssaInfo.definitions.push_back(newDef); + dataFlowInfo.addMemoryDefinition(result, newDef); + + LLVM_DEBUG(llvm::dbgs() << " Tensor definition: " << newDef->getId() + << " for " << result << "\n"); + } + } + } + // 对pointer操作（addptr, make_tensor_ptr, [broadcast, splat]） + else if (isPointerOp(op)) { + LLVM_DEBUG(llvm::dbgs() + << " Pointer operation: " << op->getName() << "\n"); + + for (Value result : op->getResults()) { + Type resultType = result.getType(); + + // 检查是否是pointer类型 + if (aliasAnalysis.isPointerType(resultType)) { + // Pointer op：复用base pointer的definition（alias） + Value basePtr = aliasAnalysis.getBasePointer(result); + MemorySSADef *baseDef = dataFlowInfo.getMemoryDefinition(basePtr); + + if (baseDef) { + ssaInfo.definitions.push_back(baseDef); + dataFlowInfo.addMemoryDefinition(result, baseDef); + + LLVM_DEBUG(llvm::dbgs() << " Pointer alias: " << baseDef->getId() + << " for " << result << "\n"); + } + } + } + } + + // 3. 特殊处理控制流操作 + if (auto ifOp = dyn_cast(op)) { + processIfOp(ifOp, inst, nullptr, nullptr); + } else if (auto forOp = dyn_cast(op)) { + processForOp(forOp, inst, nullptr); + } else if (auto whileOp = dyn_cast(op)) { + // processWhileOp(whileOp, inst, nullptr, nullptr); + } + + LLVM_DEBUG(llvm::dbgs() << " Done\n"); +} + +void MemorySSABuilder::processIfOp(scf::IfOp ifOp, Instruction *inst, + BasicBlock *thenEntryBB, + BasicBlock *elseEntryBB) { + // 实现scf.if的phi节点处理 + // 从then和else区域收集yield的values，创建phi definitions + + LLVM_DEBUG(llvm::dbgs() << "Processing IfOp: " << ifOp << "\n"); + + // 获取if指令 + // scf::IfOp ifOp = + // cast(ifCondBB->getInstruction(0)->getOperation()); + + // 为每个result创建phi节点 + for (size_t i = 0; i < ifOp.getNumResults(); ++i) { + Value ifResult = ifOp.getResult(i); + Type resultType = ifResult.getType(); + + // 只处理tensor/pointer类型 + if (!isTensorType(resultType) && !aliasAnalysis.isPointerType(resultType)) { + continue; + } + + // 从then区域获取yield的value + Operation *thenYield = + MemorySSABuilderHelper::getYieldOp(ifOp.getThenRegion()); + Value thenValue = thenYield ? thenYield->getOperand(i) : Value(); + MemorySSADef *thenDef = + thenValue ? dataFlowInfo.getMemoryDefinition(thenValue) : nullptr; + + // 从else区域获取yield的value + Operation *elseYield = + MemorySSABuilderHelper::getYieldOp(ifOp.getElseRegion()); + Value elseValue = elseYield ? elseYield->getOperand(i) : Value(); + MemorySSADef *elseDef = + elseValue ? dataFlowInfo.getMemoryDefinition(elseValue) : nullptr; + + // 如果then和else都返回相同的definition，可以直接使用 + if (thenDef && elseDef && thenDef == elseDef) { + dataFlowInfo.addMemoryDefinition(ifResult, thenDef); + + LLVM_DEBUG(llvm::dbgs() + << " If result #" << i + << " uses same definition: " << thenDef->getId() << "\n"); + continue; + } + + // 创建phi definition + if (thenDef || elseDef) { + TensorObject *tensor = thenDef ? thenDef->getTensor() + : elseDef ? elseDef->getTensor() + : nullptr; + + if (tensor) { + std::string phiName = "phi_" + std::to_string(ifOp.getNumResults()) + + "_" + std::to_string(i); + TensorObject *phiTensor = + new TensorObject(phiName, tensor->getShape(), resultType, + tensor->getElementType(), tensor->getKind()); + + MemorySSADef *phiDef = createDefinition(phiTensor, ifOp.getOperation()); + + // 记录if result的definition + dataFlowInfo.addMemoryDefinition(ifResult, phiDef); + + // 为phi节点的operands创建uses + if (thenDef && thenValue) { + MemorySSAUse thenUse(thenDef, ifOp.getOperation(), /*operandIdx=*/i); + dataFlowInfo.addMemoryUse(thenValue, thenUse); + } + if (elseDef && elseValue) { + MemorySSAUse elseUse(elseDef, ifOp.getOperation(), /*operandIdx=*/i); + dataFlowInfo.addMemoryUse(elseValue, elseUse); + } + + // 创建phi信息 + PhiInfo phiInfo; + phiInfo.type = PhiInfo::IF_RESULT; + phiInfo.loopHeader = nullptr; + phiInfo.comingFrom.initialValue = thenDef; + phiInfo.comingFrom.yieldValue = elseDef; + + dataFlowInfo.addPhi(ifResult, phiInfo); + + LLVM_DEBUG(llvm::dbgs() << " Phi: " << phiDef->getId() + << " for if result #" << i << "\n"); + } + } + } +} + +void MemorySSABuilder::processForOp(scf::ForOp forOp, Instruction *inst, + BasicBlock *loopBodyEntryBB) { + // 实现scf.for的iter_args处理 + // iter_args在所有迭代中共享同一个definition + + LLVM_DEBUG(llvm::dbgs() << "Processing ForOp: " << forOp << "\n"); + + unsigned numIterArgs = forOp.getInitArgs().size(); + + for (unsigned i = 0; i < numIterArgs; i++) { + Value iterArg = forOp.getRegionIterArg(i); + Value initValue = forOp.getInitArgs()[i]; + + // 查找initValue的definition + MemorySSADef *initDef = dataFlowInfo.getMemoryDefinition(initValue); + + if (initDef) { + // iter_arg使用同一个definition，不创建新版本 + dataFlowInfo.addMemoryDefinition(iterArg, initDef); + + // 创建PhiInfo（用于跟踪循环依赖） + PhiInfo phiInfo; + phiInfo.type = PhiInfo::ITER_ARG; + phiInfo.loopHeader = nullptr; // 需要根据CFG确定 + phiInfo.comingFrom.initialValue = initDef; + phiInfo.comingFrom.yieldValue = nullptr; // 将在处理yield时更新 + + dataFlowInfo.addPhi(iterArg, phiInfo); + + // 为initValue创建use + MemorySSAUse initUse(initDef, forOp.getOperation(), /*operandIdx=*/i); + dataFlowInfo.addMemoryUse(initValue, initUse); + + LLVM_DEBUG(llvm::dbgs() << " IterArg #" << i << ": " << iterArg << " -> " + << initDef->getId() << "\n"); + } + } + + // 处理yield操作（在循环体中） + Operation *yieldOp = MemorySSABuilderHelper::getYieldOp(forOp.getRegion()); + if (yieldOp) { + for (unsigned i = 0; i < yieldOp->getNumOperands(); i++) { + Value yieldedValue = yieldOp->getOperand(i); + Value iterArg = forOp.getRegionIterArg(i); + + MemorySSADef *yieldedDef = dataFlowInfo.getMemoryDefinition(yieldedValue); + + if (yieldedDef) { + // 更新PhiInfo + if (PhiInfo *phiInfo = &dataFlowInfo.getPhi(iterArg)) { + phiInfo->comingFrom.yieldValue = yieldedDef; + + // 为yieldedValue创建use + MemorySSAUse yieldUse(yieldedDef, yieldOp, /*operandIdx=*/i); + dataFlowInfo.addMemoryUse(yieldedValue, yieldUse); + + LLVM_DEBUG(llvm::dbgs() << " Yield #" << i << ": " << yieldedValue + << " updates iter_arg\n"); + } + } + } + } +} + +MemorySSADef *MemorySSABuilder::createDefinition(TensorObject *tensor, + Operation *op) { + unsigned version = isParameter(op) ? 0 : ++nextVersion[tensor]; + auto *def = new MemorySSADef(tensor, op, version); + allDefinitions.push_back(def); + return def; +} + +MemorySSAUse MemorySSABuilder::createUse(MemorySSADef *def, Operation *userOp, + unsigned operandIdx) { + return MemorySSAUse(def, userOp, operandIdx); +} + +TensorObject *MemorySSABuilder::createTensorObject(Operation *op) { + if (!op) + return nullptr; + + // 根据操作创建tensor对象，使用独立的Tensor ID确保唯一性 + std::string name = getOpName(op); + Type resultType = op->getResultTypes().front(); + + SmallVector shape; + Type elementType; + + // 使用Tensor.h中的辅助函数提取shape和element type + extractShapeAndElementType(resultType, shape, elementType); + + // 设置默认的kind（可以根据操作类型推断） + TensorObject::TensorKind kind = TensorObject::TensorKind::GLOBAL_MEMORY; + + auto *tensor = new TensorObject(name, shape, resultType, elementType, kind); + + // 缓存tensor对象 + if (!op->getResults().empty()) { + tensorObjectCache[op->getResult(0)] = tensor; + } + + return tensor; +} + +std::string MemorySSABuilder::getOpName(Operation *op) { + if (!op) + return "unknown"; + + // 基于操作类型和独立的Tensor ID生成名称 + std::string opName = op->getName().getStringRef().str(); + std::replace(opName.begin(), opName.end(), '.', '_'); + + // 使用独立的Tensor ID生成器，确保唯一性 + return opName + "_tensor_" + std::to_string(++nextTensor[opName]); +} + +//===----------------------------------------------------------------------===// +// MemorySSABuilderHelper +//===----------------------------------------------------------------------===// + +namespace mlir { +namespace triton { +namespace cfg { +namespace MemorySSABuilderHelper { + +Type getResultType(Operation *op, unsigned resultIdx) { + if (!op || resultIdx >= op->getNumResults()) + return Type(); + return op->getResultTypes()[resultIdx]; +} + +SmallVector getShapeFromValue(Value value) { + Type type = value.getType(); + SmallVector shape; + + if (auto rankedType = mlir::dyn_cast(type)) { + shape.append(rankedType.getShape().begin(), rankedType.getShape().end()); + } + + return shape; +} + +bool shapesEqual(ArrayRef shape1, ArrayRef shape2) { + if (shape1.size() != shape2.size()) + return false; + return std::equal(shape1.begin(), shape1.end(), shape2.begin()); +} + +Operation *getYieldOp(Region ®ion) { + if (region.empty()) + return nullptr; + + Block &block = region.back(); + if (block.empty()) + return nullptr; + + Operation &lastOp = block.back(); + if (isa(&lastOp)) { + return &lastOp; + } + + return nullptr; +} + +std::string createUniqueTensorName(StringRef prefix, size_t id) { + return prefix.str() + "_" + std::to_string(id); +} + +bool shouldCreateNewVersion(Operation *op, MemorySSADef *currentDef) { + if (!op || !currentDef) + return true; + + // 如果操作会修改tensor内容，则创建新版本 + // 例如：tt.store、tt.trans等 + if (isa(op)) + return true; + if (isa(op)) + return true; + + // 其他操作可能复用当前definition + return false; +} + +} // namespace MemorySSABuilderHelper +} // namespace cfg +} // namespace triton +} // namespace mlir + +bool MemorySSABuilder::isPointerBroadcastOrSplat(mlir::Operation *op) const { + if (auto broadcastOp = mlir::dyn_cast(op)) { + Type elemType = getElementTypeOrSelf(broadcastOp.getResult().getType()); + return mlir::isa(elemType); + } + + // 处理 SplatOp: 检查第一个 operand (src) 是否是指针 + if (auto splatOp = mlir::dyn_cast(op)) { + return mlir::isa(splatOp.getSrc().getType()); + } + + return false; +} + +// 判断是否是返回新Tensor的操作（根据返回值类型判断，排除load） +bool MemorySSABuilder::isTensorWriter(Operation *op) const { + if (!op || op->getNumResults() == 0) + return false; + + // 检查返回值类型 + for (Value result : op->getResults()) { + Type resultType = result.getType(); + // 如果是RankedTensorType（不是指针），则是TensorWriter + if (mlir::isa(resultType)) { + // 但排除load（虽然load返回tensor，但它是从内存读取，不是"写入"或"创建"） + if (mlir::isa(op) || mlir::isa(op)) + return false; + + if (mlir::isa(op)) + return false; + + if (isPointerBroadcastOrSplat(op)) + return false; + + if (auto ptrType = mlir::dyn_cast( + getElementTypeOrSelf(resultType))) { + return false; + } + return true; + } + } + return false; +} diff --git a/third_party/ascend/lib/TritonToGraph/Passes.cpp b/third_party/ascend/lib/TritonToGraph/Passes.cpp new file mode 100644 index 0000000000..801a568e6f --- /dev/null +++ b/third_party/ascend/lib/TritonToGraph/Passes.cpp @@ -0,0 +1,38 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "TritonToGraph/Passes.h" +#include "TritonToGraph/ControlFlowGraphBuilder.h" + +#include "mlir/Pass/PassRegistry.h" + +#define GEN_PASS_REGISTRATION +#include "ascend/include/TritonToGraph/Passes.h.inc" + +namespace mlir { +namespace triton { + +// registerTritonToCFGPasses() 由 Passes.h.inc 生成 +// createBuildCFGPass() 的实现需要在 ControlFlowGraphBuilder.cpp 中 + +} // namespace triton +} // namespace mlir diff --git a/third_party/ascend/lib/TritonToGraph/ProgramAxisDependenceAnalysis.cpp b/third_party/ascend/lib/TritonToGraph/ProgramAxisDependenceAnalysis.cpp new file mode 100644 index 0000000000..71b61a4afa --- /dev/null +++ b/third_party/ascend/lib/TritonToGraph/ProgramAxisDependenceAnalysis.cpp @@ -0,0 +1,496 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + */ + +#include "TritonToGraph/ProgramAxisDependenceAnalysis.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Matchers.h" +#include "mlir/Interfaces/CallInterfaces.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallPtrSet.h" + +#include +#include +#include +#include +#include +#include + +using namespace mlir; +using namespace mlir::triton; +using namespace mlir::triton::cfg; + +namespace { + +struct OffsetForm { + bool valid = true; + bool symbolicStride = false; + int64_t programCoefficient = 0; + int64_t laneMin = 0; + int64_t laneMax = 0; +}; + +bool addNoOverflow(int64_t lhs, int64_t rhs, int64_t &result) { + if ((rhs > 0 && lhs > std::numeric_limits::max() - rhs) || + (rhs < 0 && lhs < std::numeric_limits::min() - rhs)) + return false; + result = lhs + rhs; + return true; +} + +bool multiplyNoOverflow(int64_t lhs, int64_t rhs, int64_t &result) { + return !__builtin_mul_overflow(lhs, rhs, &result); +} + +std::optional getConstantInt(Value value) { + APInt constant; + if (!matchPattern(value, m_ConstantInt(&constant))) + return std::nullopt; + return constant.getSExtValue(); +} + +std::optional getUniformConstantInt(Value value) { + if (std::optional constant = getConstantInt(value)) + return constant; + DenseElementsAttr elements; + if (!matchPattern(value, m_Constant(&elements)) || !elements.isSplat() || + !isa(elements.getElementType())) + return std::nullopt; + return elements.getSplatValue().getSExtValue(); +} + +OffsetForm invalidOffsetForm() { + OffsetForm form; + form.valid = false; + return form; +} + +OffsetForm addOffsetForms(const OffsetForm &lhs, const OffsetForm &rhs, + bool subtractRhs) { + if (!lhs.valid || !rhs.valid) + return invalidOffsetForm(); + if (subtractRhs && rhs.symbolicStride) + return invalidOffsetForm(); + if ((lhs.symbolicStride && + (rhs.symbolicStride || rhs.programCoefficient != 0)) || + (rhs.symbolicStride && lhs.programCoefficient != 0)) + return invalidOffsetForm(); + const int64_t rhsCoefficient = + subtractRhs ? -rhs.programCoefficient : rhs.programCoefficient; + const int64_t rhsMin = subtractRhs ? -rhs.laneMax : rhs.laneMin; + const int64_t rhsMax = subtractRhs ? -rhs.laneMin : rhs.laneMax; + OffsetForm result; + result.symbolicStride = lhs.symbolicStride || rhs.symbolicStride; + if (!addNoOverflow(lhs.programCoefficient, rhsCoefficient, + result.programCoefficient) || + !addNoOverflow(lhs.laneMin, rhsMin, result.laneMin) || + !addNoOverflow(lhs.laneMax, rhsMax, result.laneMax)) + return invalidOffsetForm(); + return result; +} + +OffsetForm scaleOffsetForm(const OffsetForm &input, int64_t factor) { + if (!input.valid) + return invalidOffsetForm(); + if (input.symbolicStride && factor != 1) + return invalidOffsetForm(); + OffsetForm result; + result.symbolicStride = input.symbolicStride; + if (!multiplyNoOverflow(input.programCoefficient, factor, + result.programCoefficient) || + !multiplyNoOverflow(input.laneMin, factor, result.laneMin) || + !multiplyNoOverflow(input.laneMax, factor, result.laneMax)) + return invalidOffsetForm(); + if (factor < 0) + std::swap(result.laneMin, result.laneMax); + return result; +} + +bool valueDependsOnAxis(Value value, int32_t targetAxis, DenseSet &seen); + +bool isRuntimeIntegerEntryArgument(Value value) { + if (auto argument = dyn_cast(value)) { + auto integer = dyn_cast(argument.getType()); + return integer && argument.getOwner() && + isa(argument.getOwner()->getParentOp()); + } + if (auto splat = value.getDefiningOp()) + return isRuntimeIntegerEntryArgument(splat.getSrc()); + if (auto broadcast = value.getDefiningOp()) + return isRuntimeIntegerEntryArgument(broadcast.getSrc()); + if (auto expand = value.getDefiningOp()) + return isRuntimeIntegerEntryArgument(expand.getSrc()); + return false; +} + +OffsetForm analyzeOffset(Value value, int32_t targetAxis, + DenseSet &visited) { + if (!visited.insert(value).second) + return invalidOffsetForm(); + + auto finish = [&](OffsetForm form) { + visited.erase(value); + return form; + }; + + if (auto pid = value.getDefiningOp()) { + if (pid.getAxisAsInt() != targetAxis) + return finish(invalidOffsetForm()); + OffsetForm form; + form.programCoefficient = 1; + return finish(form); + } + if (value.getDefiningOp()) + return finish(invalidOffsetForm()); + if (std::optional constant = getConstantInt(value)) { + OffsetForm form; + form.laneMin = *constant; + form.laneMax = *constant; + return finish(form); + } + if (auto range = value.getDefiningOp()) { + if (range.getEnd() <= range.getStart()) + return finish(invalidOffsetForm()); + OffsetForm form; + form.laneMin = range.getStart(); + form.laneMax = range.getEnd() - 1; + return finish(form); + } + DenseSet dependenceSeen; + if (!valueDependsOnAxis(value, targetAxis, dependenceSeen)) + return finish(OffsetForm{}); + if (auto splat = value.getDefiningOp()) + return finish(analyzeOffset(splat.getSrc(), targetAxis, visited)); + if (auto broadcast = value.getDefiningOp()) + return finish(analyzeOffset(broadcast.getSrc(), targetAxis, visited)); + if (auto expand = value.getDefiningOp()) + return finish(analyzeOffset(expand.getSrc(), targetAxis, visited)); + if (auto add = value.getDefiningOp()) { + OffsetForm lhs = analyzeOffset(add.getLhs(), targetAxis, visited); + OffsetForm rhs = analyzeOffset(add.getRhs(), targetAxis, visited); + return finish(addOffsetForms(lhs, rhs, false)); + } + if (auto subtract = value.getDefiningOp()) { + OffsetForm lhs = analyzeOffset(subtract.getLhs(), targetAxis, visited); + OffsetForm rhs = analyzeOffset(subtract.getRhs(), targetAxis, visited); + return finish(addOffsetForms(lhs, rhs, true)); + } + if (auto multiply = value.getDefiningOp()) { + if (std::optional lhs = getConstantInt(multiply.getLhs())) + return finish(scaleOffsetForm( + analyzeOffset(multiply.getRhs(), targetAxis, visited), *lhs)); + if (std::optional rhs = getConstantInt(multiply.getRhs())) + return finish(scaleOffsetForm( + analyzeOffset(multiply.getLhs(), targetAxis, visited), *rhs)); + DenseSet lhsSeen; + DenseSet rhsSeen; + const bool lhsDepends = + valueDependsOnAxis(multiply.getLhs(), targetAxis, lhsSeen); + const bool rhsDepends = + valueDependsOnAxis(multiply.getRhs(), targetAxis, rhsSeen); + if (lhsDepends == rhsDepends) + return finish(invalidOffsetForm()); + Value pidTerm = lhsDepends ? multiply.getLhs() : multiply.getRhs(); + Value strideTerm = lhsDepends ? multiply.getRhs() : multiply.getLhs(); + DenseSet pidSeen; + OffsetForm pidForm = analyzeOffset(pidTerm, targetAxis, pidSeen); + if (!pidForm.valid || pidForm.symbolicStride || + pidForm.programCoefficient != 1 || pidForm.laneMin != 0 || + pidForm.laneMax != 0 || !isRuntimeIntegerEntryArgument(strideTerm)) + return finish(invalidOffsetForm()); + OffsetForm form; + form.symbolicStride = true; + return finish(form); + } + return finish(invalidOffsetForm()); +} + +bool valueDependsOnAxis(Value value, int32_t targetAxis, + DenseSet &seen) { + if (!seen.insert(value).second) + return false; + if (auto pid = value.getDefiningOp()) + return pid.getAxisAsInt() == targetAxis; + Operation *definition = value.getDefiningOp(); + if (!definition) + return false; + return llvm::any_of(definition->getOperands(), [&](Value operand) { + return valueDependsOnAxis(operand, targetAxis, seen); + }); +} + +OffsetForm analyzePointerOffset(Value pointer, int32_t targetAxis, + DenseSet &visited) { + if (!visited.insert(pointer).second) + return invalidOffsetForm(); + auto finish = [&](OffsetForm form) { + visited.erase(pointer); + return form; + }; + + if (auto addPtr = pointer.getDefiningOp()) { + OffsetForm base = + analyzePointerOffset(addPtr.getPtr(), targetAxis, visited); + DenseSet offsetVisited; + OffsetForm offset = + analyzeOffset(addPtr.getOffset(), targetAxis, offsetVisited); + return finish(addOffsetForms(base, offset, false)); + } + if (auto splat = pointer.getDefiningOp()) + return finish(analyzePointerOffset(splat.getSrc(), targetAxis, visited)); + if (auto broadcast = pointer.getDefiningOp()) + return finish( + analyzePointerOffset(broadcast.getSrc(), targetAxis, visited)); + if (auto expand = pointer.getDefiningOp()) + return finish(analyzePointerOffset(expand.getSrc(), targetAxis, visited)); + + DenseSet dependenceSeen; + return finish(valueDependsOnAxis(pointer, targetAxis, dependenceSeen) + ? invalidOffsetForm() + : OffsetForm{}); +} + +std::optional +getConjunctiveOffsetUpperBound(Value mask, Value offset, + DenseSet &visited) { + if (!mask || !visited.insert(mask).second) + return std::nullopt; + if (auto compare = mask.getDefiningOp()) { + if (compare.getPredicate() == arith::CmpIPredicate::slt && + compare.getLhs() == offset) { + std::optional bound = getUniformConstantInt(compare.getRhs()); + if (bound && *bound > 0) + return bound; + } + return std::nullopt; + } + if (auto conjunction = mask.getDefiningOp()) { + std::optional lhs = + getConjunctiveOffsetUpperBound(conjunction.getLhs(), offset, visited); + std::optional rhs = + getConjunctiveOffsetUpperBound(conjunction.getRhs(), offset, visited); + if (lhs && rhs) + return std::min(*lhs, *rhs); + return lhs ? lhs : rhs; + } + return std::nullopt; +} + +Value peelPointerShapeOps(Value pointer) { + while (true) { + if (auto splat = pointer.getDefiningOp()) { + pointer = splat.getSrc(); + continue; + } + if (auto broadcast = pointer.getDefiningOp()) { + pointer = broadcast.getSrc(); + continue; + } + if (auto expand = pointer.getDefiningOp()) { + pointer = expand.getSrc(); + continue; + } + return pointer; + } +} + +OffsetForm analyzeStorePointerOffset(triton::StoreOp store, + int32_t targetAxis) { + Value pointer = peelPointerShapeOps(store.getPtr()); + auto finalAddPtr = pointer.getDefiningOp(); + if (!finalAddPtr) { + DenseSet pointerSeen; + return analyzePointerOffset(store.getPtr(), targetAxis, pointerSeen); + } + + DenseSet baseSeen; + OffsetForm base = + analyzePointerOffset(finalAddPtr.getPtr(), targetAxis, baseSeen); + DenseSet offsetSeen; + OffsetForm offset = + analyzeOffset(finalAddPtr.getOffset(), targetAxis, offsetSeen); + if (!base.valid || !offset.valid) + return invalidOffsetForm(); + + DenseSet maskSeen; + std::optional exclusiveBound = getConjunctiveOffsetUpperBound( + store.getMask(), finalAddPtr.getOffset(), maskSeen); + if (exclusiveBound && *exclusiveBound > offset.laneMin) { + const int64_t maskedMax = *exclusiveBound - 1; + if (maskedMax < offset.laneMax) + offset.laneMax = maskedMax; + } + return addOffsetForms(base, offset, false); +} + +StoreAddressIndependence classifyStoreAddress(triton::StoreOp store, + int32_t targetAxis, + bool &pointerDependsOnAxis) { + DenseSet dependencySeen; + pointerDependsOnAxis = + valueDependsOnAxis(store.getPtr(), targetAxis, dependencySeen); + if (!pointerDependsOnAxis) + return StoreAddressIndependence::NotProgramDependent; + + OffsetForm offset = analyzeStorePointerOffset(store, targetAxis); + if (!offset.valid || + (!offset.symbolicStride && offset.programCoefficient == 0) || + offset.laneMax < offset.laneMin) + return StoreAddressIndependence::Unknown; + + if (offset.symbolicStride) + return StoreAddressIndependence::CanonicalDynamicStride; + + const uint64_t span = static_cast(offset.laneMax) - + static_cast(offset.laneMin); + const uint64_t coefficient = + offset.programCoefficient < 0 + ? static_cast(-(offset.programCoefficient + 1)) + 1 + : static_cast(offset.programCoefficient); + return coefficient > span ? StoreAddressIndependence::ProvenDisjoint + : StoreAddressIndependence::Unknown; +} + +bool isControlFlowEscape(Operation *operation) { + if (isa(operation)) + return true; + StringRef name = operation->getName().getStringRef(); + return name == "cf.br" || name == "cf.cond_br" || name.ends_with(".yield"); +} + +bool hasUnsupportedEffect(Operation *operation) { + if (isa(operation)) + return false; + if (isa(operation)) + return true; + StringRef name = operation->getName().getStringRef(); + if (name.contains("atomic") || name.contains("barrier") || + name.contains("print")) + return true; + auto effects = dyn_cast(operation); + if (!effects) + return false; + SmallVector effectInstances; + effects.getEffects(effectInstances); + return llvm::any_of(effectInstances, [](const auto &effect) { + return isa(effect.getEffect()); + }); +} + +void recordOperationFacts(Operation *operation, ProgramAxisDependence &info, + DenseSet &seenStores) { + if (auto reduce = dyn_cast(operation)) + info.reductionAxes.push_back(reduce.getAxis()); + if (auto scan = dyn_cast(operation)) + info.reductionAxes.push_back(scan.getAxis()); + + if (isControlFlowEscape(operation)) + info.escapes = true; + if (auto store = dyn_cast(operation)) { + info.hasSideEffects = true; + if (seenStores.insert(operation).second) { + bool pointerDependsOnAxis = false; + StoreAddressIndependence independence = + classifyStoreAddress(store, info.axis, pointerDependsOnAxis); + info.stores.push_back({operation, pointerDependsOnAxis, independence}); + } + return; + } + if (hasUnsupportedEffect(operation)) { + info.hasSideEffects = true; + info.hasUnsupportedSideEffects = true; + } +} + +void buildClosure(ProgramAxisDependence &info) { + DenseSet visitedValues; + DenseSet visitedOperations; + DenseSet seenStores; + SmallVector worklist(info.programIds.begin(), info.programIds.end()); + + while (!worklist.empty()) { + Value value = worklist.pop_back_val(); + if (!visitedValues.insert(value).second) + continue; + for (OpOperand &use : value.getUses()) { + Operation *operation = use.getOwner(); + if (!visitedOperations.insert(operation).second) + continue; + info.dependenceClosure.push_back(operation); + recordOperationFacts(operation, info, seenStores); + for (Value result : operation->getResults()) + worklist.push_back(result); + } + } +} + +} // namespace + +bool ProgramAxisDependence::hasOnlyDisjointStoreAddresses() const { + return !stores.empty() && llvm::all_of(stores, [](const auto &store) { + return store.pointerDependsOnAxis && + store.independence == StoreAddressIndependence::ProvenDisjoint; + }); +} + +bool ProgramAxisDependence::hasOnlyLiftableStoreAddressesForProgramMapping() + const { + return !stores.empty() && llvm::all_of(stores, [](const auto &store) { + return store.pointerDependsOnAxis && + (store.independence == StoreAddressIndependence::ProvenDisjoint || + store.independence == + StoreAddressIndependence::CanonicalDynamicStride); + }); +} + +bool ProgramAxisDependence::isIndependentAxisTransformCandidate() const { + return hasProgramId() && !escapes && !hasUnsupportedSideEffects && + !readsNumPrograms && hasOnlyDisjointStoreAddresses(); +} + +bool ProgramAxisDependence::isProgramMappingTransformCandidate() const { + return hasProgramId() && !escapes && !hasUnsupportedSideEffects && + !readsNumPrograms && hasOnlyLiftableStoreAddressesForProgramMapping(); +} + +ProgramAxisDependenceAnalysis::ProgramAxisDependenceAnalysis( + triton::FuncOp function) { + for (int32_t axis = 0; axis < 3; ++axis) + axes[axis].axis = axis; + + function.walk([&](triton::GetProgramIdOp pid) { + const int32_t axis = pid.getAxisAsInt(); + if (axis >= 0 && axis < static_cast(axes.size())) + axes[axis].programIds.push_back(pid.getResult()); + }); + function.walk([&](triton::GetNumProgramsOp programs) { + const int32_t axis = programs.getAxisAsInt(); + if (axis >= 0 && axis < static_cast(axes.size())) + axes[axis].readsNumPrograms = true; + }); + for (ProgramAxisDependence &axis : axes) + buildClosure(axis); +} + +const ProgramAxisDependence & +ProgramAxisDependenceAnalysis::get(int32_t axis) const { + assert(axis >= 0 && axis < static_cast(axes.size()) && + "program axis must be x, y, or z"); + return axes[axis]; +} diff --git a/third_party/ascend/lib/TritonToGraph/ProgramGridTransform.cpp b/third_party/ascend/lib/TritonToGraph/ProgramGridTransform.cpp new file mode 100644 index 0000000000..af51dd0862 --- /dev/null +++ b/third_party/ascend/lib/TritonToGraph/ProgramGridTransform.cpp @@ -0,0 +1,345 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "TritonToGraph/ProgramGridTransform.h" + +#include "mlir/IR/Builders.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/STLExtras.h" + +#include +#include + +using namespace mlir; +using namespace mlir::triton::cfg; + +namespace { + +constexpr llvm::StringLiteral kVersion = "version"; +constexpr llvm::StringLiteral kExtentSource = "extent_source"; +constexpr llvm::StringLiteral kHiddenExtentAxes = "hidden_extent_axes"; +constexpr llvm::StringLiteral kHiddenArgumentOrder = "hidden_argument_order"; +constexpr llvm::StringLiteral kHiddenArgumentTypes = "hidden_argument_types"; +constexpr llvm::StringLiteral kTransforms = "transforms"; +constexpr llvm::StringLiteral kOrder = "order"; +constexpr llvm::StringLiteral kKind = "kind"; +constexpr llvm::StringLiteral kAxis = "axis"; +constexpr llvm::StringLiteral kFactor = "factor"; +constexpr llvm::StringLiteral kPersistentCoverage = "persistent_coverage"; +constexpr llvm::StringLiteral kGridStrideAbiVerified = + "grid_stride_abi_verified"; +constexpr llvm::StringLiteral kCeilDiv = "ceil_div"; +constexpr llvm::StringLiteral kRuntimeOriginalGrid = "runtime_original_grid"; +constexpr llvm::StringLiteral kOriginalGridX = "originalGridX"; +constexpr llvm::StringLiteral kOriginalGridY = "originalGridY"; +constexpr llvm::StringLiteral kI32 = "i32"; +constexpr llvm::StringLiteral kHiddenXName = "__hacc_original_grid_x"; +constexpr llvm::StringLiteral kHiddenYName = "__hacc_original_grid_y"; + +bool hasExactKeys(DictionaryAttr dictionary, ArrayRef keys) { + if (dictionary.size() != keys.size()) + return false; + return llvm::all_of(keys, + [&](llvm::StringRef key) { return dictionary.get(key); }); +} + +std::optional getInteger(DictionaryAttr dictionary, + llvm::StringRef name) { + auto value = dyn_cast_or_null(dictionary.get(name)); + if (!value) + return std::nullopt; + return value.getInt(); +} + +std::optional getBoolean(DictionaryAttr dictionary, + llvm::StringRef name) { + auto value = dyn_cast_or_null(dictionary.get(name)); + if (!value) + return std::nullopt; + return value.getValue(); +} + +bool hasExactIntegerArray(Attribute attribute, ArrayRef expected) { + auto values = dyn_cast_or_null(attribute); + if (!values || values.size() != expected.size()) + return false; + for (auto [value, wanted] : llvm::zip(values, expected)) { + auto integer = dyn_cast(value); + if (!integer || integer.getInt() != wanted) + return false; + } + return true; +} + +bool hasExactStringArray(Attribute attribute, + ArrayRef expected) { + auto values = dyn_cast_or_null(attribute); + if (!values || values.size() != expected.size()) + return false; + for (auto [value, wanted] : llvm::zip(values, expected)) { + auto string = dyn_cast(value); + if (!string || string.getValue() != wanted) + return false; + } + return true; +} + +bool isPublicEntry(triton::FuncOp function) { + auto visibility = function->getAttrOfType("sym_visibility"); + return !visibility || visibility.getValue() == "public"; +} + +bool hasNamedI32Argument(BlockArgument argument, llvm::StringRef name) { + auto integer = dyn_cast(argument.getType()); + auto location = dyn_cast(argument.getLoc()); + return integer && integer.getWidth() == 32 && location && + location.getName().getValue() == name; +} + +bool isIAT16(const ProgramGridTransform &transform) { + return transform.axis == 1 && transform.factor == 16 && + !transform.persistentCoverage && !transform.gridStrideAbiVerified; +} + +bool isPTSM4(const ProgramGridTransform &transform) { + return transform.axis == 0 && transform.factor == 4 && + transform.persistentCoverage && transform.gridStrideAbiVerified; +} + +bool isPTSM64(const ProgramGridTransform &transform) { + return transform.axis == 0 && transform.factor == 64 && + transform.persistentCoverage && transform.gridStrideAbiVerified; +} + +bool isDotRowTransform(const ProgramGridTransform &transform) { + return transform.axis == 0 && !transform.persistentCoverage && + !transform.gridStrideAbiVerified && + (transform.factor == 2 || transform.factor == 4 || + transform.factor == 8 || transform.factor == 16 || + transform.factor == 32); +} + +bool isSupportedSequence(ArrayRef transforms) { + if (transforms.size() == 1) + return isIAT16(transforms.front()) || isPTSM64(transforms.front()) || + isDotRowTransform(transforms.front()); + return transforms.size() == 2 && isIAT16(transforms[0]) && + isPTSM4(transforms[1]); +} + +} // namespace + +bool mlir::triton::cfg::hasProgramGridHiddenExtentArguments( + triton::FuncOp function) { + if (function.getNumArguments() < 2) + return false; + const unsigned xIndex = function.getNumArguments() - 2; + return hasNamedI32Argument(function.getArgument(xIndex), kHiddenXName) && + hasNamedI32Argument(function.getArgument(xIndex + 1), kHiddenYName); +} + +LogicalResult mlir::triton::cfg::addProgramGridHiddenExtentArguments( + triton::FuncOp function) { + if (hasProgramGridHiddenExtentArguments(function)) + return success(); + for (BlockArgument argument : function.getArguments()) { + auto location = dyn_cast(argument.getLoc()); + if (location && (location.getName().getValue() == kHiddenXName || + location.getName().getValue() == kHiddenYName)) + return failure(); + } + + MLIRContext *context = function.getContext(); + SmallVector originalArgAttrs; + const bool hasArgAttrs = static_cast(function.getAllArgAttrs()); + if (hasArgAttrs) { + function.getAllArgAttrs(originalArgAttrs); + if (originalArgAttrs.size() != function.getNumArguments()) + return failure(); + } + + Type i32 = IntegerType::get(context, 32); + DictionaryAttr hiddenArgAttrs = + hasArgAttrs ? DictionaryAttr::get(context) : DictionaryAttr(); + if (failed(function.insertArgument( + function.getNumArguments(), i32, hiddenArgAttrs, + NameLoc::get(StringAttr::get(context, kHiddenXName), + function.getLoc())))) + return failure(); + if (failed(function.insertArgument( + function.getNumArguments(), i32, hiddenArgAttrs, + NameLoc::get(StringAttr::get(context, kHiddenYName), + function.getLoc())))) + return failure(); + + if (hasArgAttrs) { + originalArgAttrs.push_back(DictionaryAttr::get(context)); + originalArgAttrs.push_back(DictionaryAttr::get(context)); + function.setAllArgAttrs(originalArgAttrs); + } else { + function->removeAttr("arg_attrs"); + } + return success(); +} + +LogicalResult mlir::triton::cfg::commitProgramGridFunctionFromSandbox( + triton::FuncOp destination, triton::FuncOp source) { + if (source.getNumArguments() != source.getFunctionType().getNumInputs() || + !hasProgramGridHiddenExtentArguments(source)) + return failure(); + + SmallVector sourceArgAttrs; + const bool sourceHasArgAttrs = static_cast(source.getAllArgAttrs()); + if (sourceHasArgAttrs) { + source.getAllArgAttrs(sourceArgAttrs); + if (sourceArgAttrs.size() != source.getNumArguments()) + return failure(); + } + + destination.setFunctionType(source.getFunctionType()); + if (sourceHasArgAttrs) + destination.setAllArgAttrs(sourceArgAttrs); + else + destination->removeAttr("arg_attrs"); + destination.getRegion().takeBody(source.getRegion()); + return success(); +} + +FailureOr +mlir::triton::cfg::parseProgramGridTransformContract(Attribute attribute) { + auto contract = dyn_cast_or_null(attribute); + if (!contract) + return failure(); + if (!hasExactKeys(contract, + {kVersion, kExtentSource, kHiddenExtentAxes, + kHiddenArgumentOrder, kHiddenArgumentTypes, kTransforms})) + return failure(); + + std::optional version = getInteger(contract, kVersion); + auto extentSource = dyn_cast_or_null(contract.get(kExtentSource)); + constexpr std::array kAxes = {0, 1}; + constexpr std::array kArgumentOrder = {kOriginalGridX, + kOriginalGridY}; + constexpr std::array kArgumentTypes = {kI32, kI32}; + if (!version || *version != kProgramGridTransformsVersion || !extentSource || + extentSource.getValue() != kRuntimeOriginalGrid || + !hasExactIntegerArray(contract.get(kHiddenExtentAxes), kAxes) || + !hasExactStringArray(contract.get(kHiddenArgumentOrder), + kArgumentOrder) || + !hasExactStringArray(contract.get(kHiddenArgumentTypes), kArgumentTypes)) + return failure(); + + auto transforms = dyn_cast_or_null(contract.get(kTransforms)); + if (!transforms || transforms.empty() || transforms.size() > 2) + return failure(); + + ProgramGridTransformContract parsed; + parsed.version = *version; + bool seenAxis[2] = {false, false}; + for (auto [expectedOrder, rawTransform] : llvm::enumerate(transforms)) { + auto transform = dyn_cast(rawTransform); + if (!transform || + !hasExactKeys(transform, {kOrder, kKind, kAxis, kFactor, + kPersistentCoverage, kGridStrideAbiVerified})) + return failure(); + std::optional order = getInteger(transform, kOrder); + auto kind = dyn_cast_or_null(transform.get(kKind)); + std::optional axis = getInteger(transform, kAxis); + std::optional factor = getInteger(transform, kFactor); + std::optional persistent = getBoolean(transform, kPersistentCoverage); + std::optional gridStride = + getBoolean(transform, kGridStrideAbiVerified); + if (!order || !kind || !axis || !factor || !persistent || !gridStride || + *order != static_cast(expectedOrder) || + kind.getValue() != kCeilDiv || *axis < 0 || *axis > 1 || *factor < 2 || + *persistent != *gridStride || seenAxis[*axis]) + return failure(); + seenAxis[*axis] = true; + parsed.transforms.push_back(ProgramGridTransform{ + static_cast(*order), static_cast(*axis), *factor, + *persistent, *gridStride}); + } + if (!isSupportedSequence(parsed.transforms)) + return failure(); + return parsed; +} + +DictionaryAttr mlir::triton::cfg::serializeProgramGridTransformContract( + MLIRContext *context, const ProgramGridTransformContract &contract) { + Builder builder(context); + SmallVector transforms; + transforms.reserve(contract.transforms.size()); + for (const ProgramGridTransform &transform : contract.transforms) { + transforms.push_back(DictionaryAttr::get( + context, + {{builder.getStringAttr(kOrder), + builder.getI64IntegerAttr(transform.order)}, + {builder.getStringAttr(kKind), builder.getStringAttr(kCeilDiv)}, + {builder.getStringAttr(kAxis), + builder.getI64IntegerAttr(transform.axis)}, + {builder.getStringAttr(kFactor), + builder.getI64IntegerAttr(transform.factor)}, + {builder.getStringAttr(kPersistentCoverage), + builder.getBoolAttr(transform.persistentCoverage)}, + {builder.getStringAttr(kGridStrideAbiVerified), + builder.getBoolAttr(transform.gridStrideAbiVerified)}})); + } + return DictionaryAttr::get( + context, + {{builder.getStringAttr(kVersion), + builder.getI64IntegerAttr(contract.version)}, + {builder.getStringAttr(kExtentSource), + builder.getStringAttr(kRuntimeOriginalGrid)}, + {builder.getStringAttr(kHiddenExtentAxes), + ArrayAttr::get(context, {builder.getI64IntegerAttr(0), + builder.getI64IntegerAttr(1)})}, + {builder.getStringAttr(kHiddenArgumentOrder), + ArrayAttr::get(context, {builder.getStringAttr(kOriginalGridX), + builder.getStringAttr(kOriginalGridY)})}, + {builder.getStringAttr(kHiddenArgumentTypes), + ArrayAttr::get(context, {builder.getStringAttr(kI32), + builder.getStringAttr(kI32)})}, + {builder.getStringAttr(kTransforms), + ArrayAttr::get(context, transforms)}}); +} + +LogicalResult mlir::triton::cfg::setProgramGridTransformContract( + ModuleOp module, const ProgramGridTransformContract &contract) { + DictionaryAttr serialized = + serializeProgramGridTransformContract(module.getContext(), contract); + if (failed(parseProgramGridTransformContract(serialized))) + return failure(); + + triton::FuncOp entry; + for (triton::FuncOp function : module.getOps()) { + if (!isPublicEntry(function)) + continue; + if (entry) + return failure(); + entry = function; + } + if (!entry || !hasProgramGridHiddenExtentArguments(entry)) + return failure(); + module->setAttr(kProgramGridTransformsAttr, serialized); + return success(); +} diff --git a/third_party/ascend/lib/TritonToGraph/ProgramMappingScalarSpecialization.cpp b/third_party/ascend/lib/TritonToGraph/ProgramMappingScalarSpecialization.cpp new file mode 100644 index 0000000000..ec6391a060 --- /dev/null +++ b/third_party/ascend/lib/TritonToGraph/ProgramMappingScalarSpecialization.cpp @@ -0,0 +1,266 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + */ + +#include "TritonToGraph/ProgramMappingScalarSpecialization.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/IR/Builders.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/STLExtras.h" + +#include +#include + +using namespace mlir; +using namespace mlir::triton::cfg; + +namespace { + +constexpr llvm::StringLiteral kVersion = "version"; +constexpr llvm::StringLiteral kArguments = "arguments"; +constexpr llvm::StringLiteral kIndex = "index"; +constexpr llvm::StringLiteral kName = "name"; +constexpr llvm::StringLiteral kValue = "value"; + +bool hasExactKeys(DictionaryAttr dictionary, ArrayRef keys) { + if (dictionary.size() != keys.size()) + return false; + return llvm::all_of(keys, + [&](llvm::StringRef key) { return dictionary.get(key); }); +} + +std::optional getInteger(DictionaryAttr dictionary, + llvm::StringRef name) { + auto value = dyn_cast_or_null(dictionary.get(name)); + if (!value) + return std::nullopt; + return value.getInt(); +} + +std::optional getString(DictionaryAttr dictionary, + llvm::StringRef name) { + auto value = dyn_cast_or_null(dictionary.get(name)); + if (!value || value.getValue().empty()) + return std::nullopt; + return value.getValue(); +} + +void setAttrOnFunctions(ModuleOp module, llvm::StringRef name, + Attribute attribute) { + module.walk([&](Operation *operation) { + if (operation->getName().getStringRef() == "tt.func") + operation->setAttr(name, attribute); + }); +} + +void removeAttrFromFunctions(ModuleOp module, llvm::StringRef name) { + module.walk([&](Operation *operation) { + if (operation->getName().getStringRef() == "tt.func") + operation->removeAttr(name); + }); +} + +bool isPublicEntry(triton::FuncOp function) { + auto visibility = function->getAttrOfType("sym_visibility"); + return !visibility || visibility.getValue() == "public"; +} + +bool sameScalarSpecialization(const ProgramMappingScalarSpecialization &lhs, + const ProgramMappingScalarSpecialization &rhs) { + if (lhs.version != rhs.version || + lhs.arguments.size() != rhs.arguments.size()) + return false; + for (unsigned index = 0; index < lhs.arguments.size(); ++index) { + const ProgramMappingScalarArgument &left = lhs.arguments[index]; + const ProgramMappingScalarArgument &right = rhs.arguments[index]; + if (left.index != right.index || left.name != right.name || + left.value != right.value) + return false; + } + return true; +} + +bool isRepresentableInIntegerType(int64_t value, IntegerType type) { + const unsigned width = type.getWidth(); + if (width == 0) + return false; + if (width >= 64) + return true; + const int64_t minimum = -(int64_t{1} << (width - 1)); + const int64_t maximum = (int64_t{1} << (width - 1)) - 1; + return value >= minimum && value <= maximum; +} + +std::optional findEntryArgumentByName(triton::FuncOp function, + llvm::StringRef name) { + std::optional found; + for (unsigned index = 0; index < function.getNumArguments(); ++index) { + BlockArgument argument = function.getArgument(index); + auto namedLocation = dyn_cast(argument.getLoc()); + if (!namedLocation || namedLocation.getName().getValue() != name) + continue; + if (found) + return std::nullopt; + found = argument; + } + return found; +} + +} // namespace + +FailureOr +mlir::triton::cfg::parseProgramMappingScalarSpecialization( + Attribute attribute) { + auto dictionary = dyn_cast_or_null(attribute); + if (!dictionary || !hasExactKeys(dictionary, {kVersion, kArguments})) + return failure(); + + std::optional version = getInteger(dictionary, kVersion); + auto arguments = dyn_cast_or_null(dictionary.get(kArguments)); + if (!version || + (*version != kProgramMappingScalarSpecializationLegacyVersion && + *version != kProgramMappingScalarSpecializationVersion) || + !arguments || arguments.empty()) + return failure(); + + ProgramMappingScalarSpecialization parsed; + parsed.version = *version; + std::optional previousIndex; + for (Attribute attribute : arguments) { + auto argument = dyn_cast(attribute); + const bool isLegacy = + *version == kProgramMappingScalarSpecializationLegacyVersion; + if (!argument || (isLegacy && !hasExactKeys(argument, {kIndex, kValue})) || + (!isLegacy && !hasExactKeys(argument, {kIndex, kName, kValue}))) + return failure(); + std::optional index = getInteger(argument, kIndex); + std::optional value = getInteger(argument, kValue); + std::optional name = + isLegacy ? std::optional(llvm::StringRef()) + : getString(argument, kName); + if (!index || !value || *index < 0 || + *index > std::numeric_limits::max() || !name || + (previousIndex && *index <= *previousIndex)) + return failure(); + parsed.arguments.push_back(ProgramMappingScalarArgument{ + static_cast(*index), name->str(), *value}); + previousIndex = *index; + } + return parsed; +} + +DictionaryAttr mlir::triton::cfg::serializeProgramMappingScalarSpecialization( + MLIRContext *context, + const ProgramMappingScalarSpecialization &specialization) { + Builder builder(context); + SmallVector arguments; + arguments.reserve(specialization.arguments.size()); + for (const ProgramMappingScalarArgument &argument : + specialization.arguments) { + SmallVector fields; + fields.emplace_back(builder.getStringAttr(kIndex), + builder.getI64IntegerAttr(argument.index)); + if (specialization.version == kProgramMappingScalarSpecializationVersion) + fields.emplace_back(builder.getStringAttr(kName), + builder.getStringAttr(argument.name)); + fields.emplace_back(builder.getStringAttr(kValue), + builder.getI64IntegerAttr(argument.value)); + arguments.push_back(DictionaryAttr::get(context, fields)); + } + return DictionaryAttr::get( + context, {{builder.getStringAttr(kVersion), + builder.getI64IntegerAttr(specialization.version)}, + {builder.getStringAttr(kArguments), + ArrayAttr::get(context, arguments)}}); +} + +LogicalResult mlir::triton::cfg::setProgramMappingScalarSpecialization( + ModuleOp module, const ProgramMappingScalarSpecialization &specialization) { + DictionaryAttr serialized = serializeProgramMappingScalarSpecialization( + module.getContext(), specialization); + if (failed(parseProgramMappingScalarSpecialization(serialized))) + return failure(); + module->setAttr(kProgramMappingScalarSpecializationAttr, serialized); + setAttrOnFunctions(module, kProgramMappingScalarSpecializationAttr, + serialized); + return success(); +} + +void mlir::triton::cfg::clearProgramMappingScalarSpecialization( + ModuleOp module) { + module->removeAttr(kProgramMappingScalarSpecializationAttr); + removeAttrFromFunctions(module, kProgramMappingScalarSpecializationAttr); +} + +LogicalResult +mlir::triton::cfg::applyProgramMappingScalarSpecialization(ModuleOp module) { + Attribute attribute = + module->getAttr(kProgramMappingScalarSpecializationAttr); + if (!attribute) { + clearProgramMappingScalarSpecialization(module); + return success(); + } + + FailureOr specialization = + parseProgramMappingScalarSpecialization(attribute); + if (failed(specialization)) { + clearProgramMappingScalarSpecialization(module); + return failure(); + } + + for (triton::FuncOp function : module.getOps()) { + Attribute functionAttribute = + function->getAttr(kProgramMappingScalarSpecializationAttr); + if (functionAttribute) { + FailureOr functionSpecialization = + parseProgramMappingScalarSpecialization(functionAttribute); + if (failed(functionSpecialization) || + !sameScalarSpecialization(*specialization, *functionSpecialization)) { + clearProgramMappingScalarSpecialization(module); + return failure(); + } + } + if (!isPublicEntry(function) || function->getNumRegions() != 1 || + function.getBody().empty()) + continue; + + Block &entry = function.getBody().front(); + OpBuilder builder(function.getContext()); + builder.setInsertionPointToStart(&entry); + for (const ProgramMappingScalarArgument &argument : + specialization->arguments) { + std::optional blockArgument; + if (specialization->version == + kProgramMappingScalarSpecializationLegacyVersion) { + if (argument.index < function.getNumArguments()) + blockArgument = function.getArgument(argument.index); + } else { + blockArgument = findEntryArgumentByName(function, argument.name); + } + if (!blockArgument) + continue; + auto integerType = dyn_cast(blockArgument->getType()); + if (!integerType || + !isRepresentableInIntegerType(argument.value, integerType)) + continue; + Value constant = builder.create( + function.getLoc(), argument.value, integerType.getWidth()); + blockArgument->replaceAllUsesWith(constant); + } + } + + clearProgramMappingScalarSpecialization(module); + return success(); +} diff --git a/third_party/ascend/lib/TritonToHFusion/CMakeLists.txt b/third_party/ascend/lib/TritonToHFusion/CMakeLists.txt new file mode 100644 index 0000000000..f69c1bfd55 --- /dev/null +++ b/third_party/ascend/lib/TritonToHFusion/CMakeLists.txt @@ -0,0 +1,14 @@ +add_triton_library(TritonToHFusion + TritonToHFusion.cpp + + DEPENDS + TritonToHFusionConversionPassIncGen + + LINK_LIBS + BiShengIRHFusionDialect + MLIRIR + MLIRPass + MLIRTransforms + MLIRSupport + TritonIR +) diff --git a/third_party/ascend/lib/TritonToHFusion/TritonToHFusion.cpp b/third_party/ascend/lib/TritonToHFusion/TritonToHFusion.cpp new file mode 100644 index 0000000000..5ea7683931 --- /dev/null +++ b/third_party/ascend/lib/TritonToHFusion/TritonToHFusion.cpp @@ -0,0 +1,238 @@ +/* + * Copyright (c) Huawei Technologies Co. + * Licensed under the MIT license. + */ + +#include "TritonToHFusion/Passes.h" + +#include "Dialect/TritonAscend/IR/TritonAscendDialect.h" +#include "bishengir/Dialect/HFusion/IR/HFusion.h" +#include "bishengir/Dialect/HFusion/IR/HFusionImpl.h" +#include "bishengir/Dialect/Tensor/IR/TensorImpl.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/Attributes.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/LogicalResult.h" + +namespace mlir { +namespace triton { +#define GEN_PASS_DEF_TRITONTOHFUSION +#include "ascend/include/TritonToHFusion/Passes.h.inc" +} // namespace triton +} // namespace mlir + +using namespace mlir; +using namespace hfusion; + +namespace { +struct TritonModToHFusionConversion : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::ascend::ModOp op, + PatternRewriter &rewriter) const final { + auto lhsType = dyn_cast(op.getLhs().getType()); + auto rhsType = dyn_cast(op.getRhs().getType()); + if (!lhsType || !rhsType) { + return failure(); + } + + auto emptyTensor = rewriter.create( + op.getLoc(), lhsType.getShape(), lhsType.getElementType()); + auto newOp = + hfusion::createBinaryOp( + rewriter, op.getLoc(), hfusion::BinaryFn::mod, + ValueRange({op.getLhs(), op.getRhs()}), + ValueRange({emptyTensor.getResult()})); + + rewriter.replaceOp(op, newOp->getResult(0)); + return success(); + } +}; + +struct TritonHistogramToHFusionConversion + : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::HistogramOp op, + PatternRewriter &rewriter) const final { + auto loc = op.getLoc(); + Value input = op.getSrc(); + auto resultType = op.getResult().getType(); + + int64_t numBins = 256; // 256 is default fallback. + if (auto rankedTy = dyn_cast(resultType)) + if (rankedTy.hasStaticShape() && rankedTy.getNumElements() > 0) + numBins = rankedTy.getNumElements(); + + auto numBinsAttr = rewriter.getI64IntegerAttr(numBins); + + auto newOp = rewriter.create(loc, resultType, input, + numBinsAttr, Value()); + + rewriter.replaceOp(op, newOp.getResult()); + return success(); + } +}; + +struct TritonFpToFpToHFusionConversion : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::FpToFpOp op, + PatternRewriter &rewriter) const final { + auto loc = op.getLoc(); + Value input = op.getSrc(); + auto resultType = op.getResult().getType(); + + // Only handle float-to-float conversions with non-RTNE rounding modes + // RTNE (default) rounding is handled by TritonToLinalg pass using + // arith.truncf/extf + auto srcType = cast(input.getType()); + auto dstType = cast(resultType); + if (!srcType.getElementType().isIntOrFloat() || + !dstType.getElementType().isIntOrFloat()) { + return failure(); + } + + // Check if this has a non-RTNE rounding mode + auto roundingMode = op.getRounding(); + if (!roundingMode.has_value() || + roundingMode.value() == triton::RoundingMode::RTNE) { + // RTNE or no rounding mode specified: let TritonToLinalg handle it + return failure(); + } + + // Map non-RTNE rounding modes to HFusion rounding mode + hfusion::RoundMode hfusionRoundMode; + switch (roundingMode.value()) { + case triton::RoundingMode::RTZ: + hfusionRoundMode = hfusion::RoundMode::TRUNC; + break; + default: + return op.emitError("Unsupported rounding mode for HFusion conversion"); + } + // Note: Only RTZ (and potential future non-RTNE modes) reach here + + // Get or create destination tensor (destination-style) + SmallVector dsts; + if (failed(tensor::getOrCreateDestinations(rewriter, loc, op, dsts))) + return failure(); + + // Create the HFusion cast operation with round_mode attribute + auto roundModeAttr = + hfusion::RoundModeAttr::get(rewriter.getContext(), hfusionRoundMode); + auto modeAttr = rewriter.getNamedAttr("mode", roundModeAttr); + + rewriter.replaceOpWithNewOp( + op, ValueRange{input}, ValueRange{dsts}, ArrayRef{modeAttr}); + + return success(); + } +}; + +struct TritonConv1dToHFusionConversion + : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::ascend::Conv1dOp op, + PatternRewriter &rewriter) const final { + auto loc = op.getLoc(); + + Value input = op.getInput(); + Value weight = op.getWeight(); + Value biasValue = op.getBias(); + int64_t stride = op.getStride(); + int64_t padding_size = op.getPaddingSize(); + int64_t dilation = op.getDilation(); + int64_t groups = op.getGroups(); + + auto inputType = mlir::cast(input.getType()); + auto weightType = mlir::cast(weight.getType()); + if (!inputType.hasStaticShape() || !weightType.hasStaticShape()) { + return failure(); + } + + ArrayRef inputShape = inputType.getShape(); + ArrayRef weightShape = weightType.getShape(); + + bool isBatched = inputShape.size() == 3; + int64_t N; + if (isBatched) + N = inputShape[0]; + int64_t L_in = inputShape[isBatched ? 2 : 1]; + int64_t C_out = weightShape[0]; + int64_t kernel_size = weightShape[2]; + + if (stride == 0) { + return failure(); + } + int64_t L_out = + (L_in + 2 * padding_size - dilation * (kernel_size - 1) - 1) / stride + + 1; + + auto resultType = mlir::cast(op.getResult().getType()); + Type resultElementType = resultType.getElementType(); + + constexpr int64_t dim2 = 2; + constexpr int64_t dim3 = 3; + Value initTensor; + if (isBatched) { + SmallVector outputShape{N, C_out, L_out}; + initTensor = + rewriter.create(loc, outputShape, resultElementType); + } else { + SmallVector outputShape{C_out, L_out}; + initTensor = + rewriter.create(loc, outputShape, resultElementType); + } + + SmallVector ins; + ins.push_back(input); + ins.push_back(weight); + if (biasValue) { + ins.push_back(biasValue); + } + + auto newOp = rewriter.create( + loc, ins, initTensor, stride, padding_size, dilation, groups); + + rewriter.replaceOp(op, newOp.getResult()); + + return success(); + } +}; +} // namespace + +namespace { +struct TritonToHFusionPass + : public mlir::triton::impl::TritonToHFusionBase { + void runOnOperation() override; +}; +} // namespace + +void TritonToHFusionPass::runOnOperation() { + auto module = getOperation(); + + // Use greedy pattern rewriter for simpler pattern matching + // Patterns decide themselves whether to convert (via returning + // success/failure) + RewritePatternSet patterns(&getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + + // Apply patterns with greedy rewriting + // This allows patterns to return failure() without causing pass failure + if (failed(applyPatternsGreedily(module, std::move(patterns)))) { + signalPassFailure(); + } +} + +std::unique_ptr> +mlir::triton::createTritonToHFusionPass() { + return std::make_unique(); +} diff --git a/third_party/ascend/lib/TritonToHIVM/CMakeLists.txt b/third_party/ascend/lib/TritonToHIVM/CMakeLists.txt new file mode 100644 index 0000000000..535e47f1fd --- /dev/null +++ b/third_party/ascend/lib/TritonToHIVM/CMakeLists.txt @@ -0,0 +1,26 @@ +# TLE DSA support is optional. TritonToHIVM.cpp guards every use of it with +# #ifdef __TLE_DSA__, so only depend on the TLE targets when that subtree is +# part of the build. +set(_TritonToHIVM_TLE_DEPENDS "") +set(_TritonToHIVM_TLE_LINK_LIBS "") +if(FLAGTREE_TLE_DSA) + set(_TritonToHIVM_TLE_DEPENDS TleDSATableGen) + set(_TritonToHIVM_TLE_LINK_LIBS TleDSAToHIVM) +endif() + +add_triton_library(TritonToHIVM + TritonToHIVM.cpp + + DEPENDS + TritonToHIVMConversionPassIncGen + ${_TritonToHIVM_TLE_DEPENDS} + + LINK_LIBS + BiShengIRHIVMDialect + MLIRIR + MLIRPass + MLIRTransforms + MLIRSupport + TritonIR + ${_TritonToHIVM_TLE_LINK_LIBS} +) diff --git a/third_party/ascend/lib/TritonToHIVM/TritonToHIVM.cpp b/third_party/ascend/lib/TritonToHIVM/TritonToHIVM.cpp new file mode 100644 index 0000000000..b01dcfae7e --- /dev/null +++ b/third_party/ascend/lib/TritonToHIVM/TritonToHIVM.cpp @@ -0,0 +1,199 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "TritonToHIVM/Passes.h" + +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "mlir/IR/Attributes.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/LogicalResult.h" + +#include "ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.h" +#ifdef __TLE_DSA__ +#include "tle/dsa/dialect/include/Conversion/TleToHIVM/TleDistributedOpToHIVMConverter.h" +#endif +namespace mlir { +namespace triton { +#define GEN_PASS_DEF_TRITONTOHIVM +#include "ascend/include/TritonToHIVM/Passes.h.inc" +} // namespace triton +} // namespace mlir + +using namespace mlir; +using namespace hivm; + +namespace { + +#ifdef __TLE_DSA__ +constexpr int patternBenefitPrioritizeOverLLVMConversions = 10; +#endif +struct CoreAndPipes { + TCoreTypeAttr core; + PipeAttr producer; + PipeAttr consumer; +}; + +static LogicalResult EmitUnknownOpError(Operation *op, llvm::StringRef opName) { + op->emitError("Unknown custom operation: ") << opName; + return failure(); +} + +static void CreateSyncBlock(PatternRewriter &rewriter, Location loc, + MLIRContext *ctx, Operation *op, int64_t id, + hivm::SyncBlockMode mode, PipeAttr pipe1, + PipeAttr pipe2) { + auto syncMode = hivm::SyncBlockModeAttr::get(ctx, mode); + auto newOp = rewriter.create( + loc, syncMode, rewriter.getI16IntegerAttr(id), Value{}, pipe1, pipe2); + rewriter.replaceOp(op, newOp); +} + +static CoreAndPipes GetCoreAndPipes(MLIRContext *ctx, llvm::StringRef opName, + llvm::StringRef sender) { + // Step 1: Decide pipes + PipeAttr producer; + PipeAttr consumer = PipeAttr::get(ctx, PIPE::PIPE_MTE2); + + if (sender == "cube") { + producer = PipeAttr::get(ctx, PIPE::PIPE_FIX); + } else { + producer = PipeAttr::get(ctx, PIPE::PIPE_MTE3); + } + + // Step 2: Decide core type + TCoreTypeAttr core; + if (sender == "cube") { + if (opName == "sync_block_set") + core = TCoreTypeAttr::get(ctx, TCoreType::CUBE); + else + core = TCoreTypeAttr::get(ctx, TCoreType::VECTOR); + } else { + if (opName == "sync_block_set") + core = TCoreTypeAttr::get(ctx, TCoreType::VECTOR); + else + core = TCoreTypeAttr::get(ctx, TCoreType::CUBE); + } + + return {core, producer, consumer}; +} + +} // end anonymous namespace +namespace { +struct TritonToHIVMPass + : public mlir::triton::impl::TritonToHIVMBase { + void runOnOperation() override; +}; +} // namespace + +struct TritonCustomOpToHIVMSyncOpConversion + : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::ascend::CustomOp op, + PatternRewriter &rewriter) const final { + auto *ctx = op->getContext(); + auto loc = op->getLoc(); + auto args = op.getStrArgs(); + auto argAttr = dyn_cast(args[0]); + auto id = dyn_cast(args[1]).getInt(); + llvm::StringRef opName = op.getOpName(); + llvm::StringRef arg = argAttr.getValue(); + + if (opName == "sync_block_all") { + if (arg == "all_cube") { + CreateSyncBlock(rewriter, loc, ctx, op, id, + hivm::SyncBlockMode::ALL_CUBE, + PipeAttr::get(ctx, PIPE::PIPE_FIX), hivm::PipeAttr{}); + } else if (arg == "all_vector") { + CreateSyncBlock(rewriter, loc, ctx, op, id, + hivm::SyncBlockMode::ALL_VECTOR, hivm::PipeAttr{}, + PipeAttr::get(ctx, PIPE::PIPE_MTE3)); + } else if (arg == "all") { + CreateSyncBlock(rewriter, loc, ctx, op, id, hivm::SyncBlockMode::ALL, + PipeAttr::get(ctx, PIPE::PIPE_FIX), + PipeAttr::get(ctx, PIPE::PIPE_MTE3)); + } else { + return EmitUnknownOpError(op, opName); + } + return success(); + } + + if (opName == "sync_block_set") { + auto [coreAttr, prodPipe, consPipe] = GetCoreAndPipes(ctx, opName, arg); + rewriter.replaceOp(op, rewriter.create( + loc, coreAttr, prodPipe, consPipe, + rewriter.getIndexAttr(id))); + return success(); + } + + if (opName == "sync_block_wait") { + auto [coreAttr, prodPipe, consPipe] = GetCoreAndPipes(ctx, opName, arg); + rewriter.replaceOp(op, rewriter.create( + loc, coreAttr, prodPipe, consPipe, + rewriter.getIndexAttr(id))); + return success(); + } + + return EmitUnknownOpError(op, opName); + } +}; + +void TritonToHIVMPass::runOnOperation() { + auto module = getOperation(); + ConversionTarget target(getContext()); + target.addLegalDialect(); + +#ifdef __TLE_DSA__ + // CUBE_AND_VECTOR core type; otherwise VECTOR only. + bool existDot = false; + module.walk([&](triton::DotOp dotOp) { + existDot = true; + return WalkResult::interrupt(); + }); + module.walk([&](triton::DotScaledOp dotScaledOp) { + existDot = true; + return WalkResult::interrupt(); + }); +#endif + + RewritePatternSet patterns(&getContext()); + +#ifdef __TLE_DSA__ + // TLE distributed ops + mlir::triton::tle::populateTleDistributedOpToHIVMConversionPatterns( + patterns, patternBenefitPrioritizeOverLLVMConversions, existDot); +#endif + + patterns.add(patterns.getContext()); + if (failed(applyPartialConversion(module, target, std::move(patterns)))) { + signalPassFailure(); + } +} + +std::unique_ptr> +mlir::triton::createTritonToHIVMPass() { + return std::make_unique(); +} diff --git a/third_party/ascend/lib/TritonToLLVM/CMakeLists.txt b/third_party/ascend/lib/TritonToLLVM/CMakeLists.txt new file mode 100644 index 0000000000..4ec7329118 --- /dev/null +++ b/third_party/ascend/lib/TritonToLLVM/CMakeLists.txt @@ -0,0 +1,13 @@ +add_triton_library(TritonToLLVM + TritonToLLVM.cpp + + DEPENDS + TritonToLLVMConversionPassIncGen + + LINK_LIBS + MLIRIR + MLIRPass + MLIRTransforms + MLIRSupport + TritonIR +) diff --git a/third_party/ascend/lib/TritonToLLVM/TritonToLLVM.cpp b/third_party/ascend/lib/TritonToLLVM/TritonToLLVM.cpp new file mode 100644 index 0000000000..3eee21de82 --- /dev/null +++ b/third_party/ascend/lib/TritonToLLVM/TritonToLLVM.cpp @@ -0,0 +1,269 @@ +#include "TritonToLLVM/Passes.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/Attributes.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/LogicalResult.h" + +namespace mlir { +namespace triton { +#define GEN_PASS_DEF_TRITONTOLLVM +#include "ascend/include/TritonToLLVM/Passes.h.inc" +} // namespace triton +} // namespace mlir + +using namespace mlir; + +namespace { +struct TritonToLLVMPass + : public mlir::triton::impl::TritonToLLVMBase { + void runOnOperation() override; +}; +} // namespace + +namespace { + +static Type getElementType(Value value) { + auto type = value.getType(); + if (auto tensorType = dyn_cast(type)) + return tensorType.getElementType(); + return type; +} + +static int64_t getTensorNumElements(Value tensor) { + auto type = mlir::cast(tensor.getType()); + return type.getNumElements(); +} + +static Value getInt32Value(RewriterBase &rewriter, Location loc, int val) { + Type ty = rewriter.getI32Type(); + return rewriter.create(loc, ty, + rewriter.getIntegerAttr(ty, val)); +} + +// If operand size is smaller than 32 bits, pack in groups of 32 bits. +SmallVector packOperands(mlir::triton::ElementwiseInlineAsmOp op, + const SmallVector> &operands, + RewriterBase &rewriter, Location loc) { + SmallVector packedOperands; + unsigned numPackedElements = op.getPackedElement(); + for (int i = 0, e = op.getNumOperands(); i < e; i++) { + Type elemTy = getElementType(op.getOperand(i)); + unsigned bitWidth = + elemTy.isIntOrFloat() ? elemTy.getIntOrFloatBitWidth() : 64; + unsigned numElementPerReg = std::max(32 / bitWidth, 1u); + numElementPerReg = std::min(numElementPerReg, numPackedElements); + for (int j = 0; j < numPackedElements; j += numElementPerReg) { + if (numElementPerReg == 1) { + packedOperands.push_back(operands[j][i]); + continue; + } + Type t = VectorType::get(numElementPerReg, elemTy); + Value packed = rewriter.create(loc, t); + for (int k = 0; k < numElementPerReg; k++) { + packed = rewriter.create( + loc, packed, operands[j + k][i], getInt32Value(rewriter, loc, k)); + } + packedOperands.push_back(packed); + } + } + return packedOperands; +} + +static SmallVector unpackElements(Location loc, Value packedValues, + RewriterBase &rewriter) { + auto type = mlir::cast(packedValues.getType()); + auto elementType = type.getElementType(); + auto shape = type.getShape(); + + int64_t numElements = type.getNumElements(); + + SmallVector result; + for (int64_t linearIdx = 0; linearIdx < numElements; linearIdx++) { + SmallVector indexes(shape.size()); + int64_t remaining = linearIdx; + for (int64_t dim = shape.size() - 1; dim >= 0; dim--) { + indexes[dim] = + rewriter.create(loc, remaining % shape[dim]); + remaining /= shape[dim]; + } + Value extracted = rewriter.create(loc, elementType, + packedValues, indexes); + result.push_back(extracted); + } + + return result; +} + +static SmallVector> +createDestOps(triton::ElementwiseInlineAsmOp op, RewriterBase &rewriter, + const SmallVector> operands, Location loc) { + auto ctx = op->getContext(); + if (operands.size() % op.getPackedElement() != 0) + llvm::report_fatal_error("Inline asm op has more packed elements than " + "number of elements per thread."); + + // Pack elems smaller than 32 bits into 32-bit registers. + SmallVector packedOperands = packOperands(op, operands, rewriter, loc); + + // Types returned by the LLVM asm op. If there's more than one, they'll be + // wrapped in a type tuple. + SmallVector asmRetTypes; + for (auto result : op.getResult()) { + auto ty = getElementType(result); + + // Pack return elements into 32-bits. + unsigned bitWidth = ty.isIntOrFloat() ? ty.getIntOrFloatBitWidth() : 64; + unsigned numElemsPerReg = + std::min(std::max(32 / bitWidth, 1u), op.getPackedElement()); + assert(op.getPackedElement() % numElemsPerReg == 0); + if (numElemsPerReg > 1) { + ty = VectorType::get(numElemsPerReg, ty); + } + for (unsigned i = 0; i < op.getPackedElement() / numElemsPerReg; i++) { + asmRetTypes.push_back(ty); + } + } + Type asmRetType = asmRetTypes.size() > 1 + ? LLVM::LLVMStructType::getLiteral(ctx, asmRetTypes) + : asmRetTypes[0]; + + Value asmResults = + rewriter + .create( + loc, asmRetType, + /*operands=*/packedOperands, + /*asm_string=*/op.getAsmString(), + /*constraints=*/op.getConstraints(), + /*has_side_effects=*/!op.getPure(), + /*is_align_stack=*/false, + /*tail_call_kind=*/LLVM::tailcallkind::TailCallKind::None, + /*asm_dialect=*/ + LLVM::AsmDialectAttr::get(rewriter.getContext(), + LLVM::AsmDialect::AD_ATT), + /*operand_attrs=*/ArrayAttr()) + ->getResult(0); + + // asmResults is a flat struct; pack its values into + // [return_value][op.getPackedElement()]. + SmallVector> ret(op->getNumResults()); + int structIdx = 0; + for (int i = 0; i < op->getNumResults(); i++) { + for (int j = 0; j < op.getPackedElement(); j++) { + Value val; + if (asmRetTypes.size() > 1) { + val = + rewriter.create(loc, asmResults, structIdx++); + } else { + val = asmResults; + } + if (auto vectorTy = dyn_cast(val.getType())) { + for (int k = 0; k < vectorTy.getNumElements(); k++) { + ret[i].push_back(rewriter.create( + loc, val, getInt32Value(rewriter, loc, k))); + } + j += vectorTy.getNumElements() - 1; + } else { + ret[i].push_back(val); + } + } + } + return ret; +} + +static LogicalResult processScalarInlineAsm(triton::ElementwiseInlineAsmOp op, + PatternRewriter &rewriter) { + Location loc = op.getLoc(); + + auto outsWrapped = createDestOps(op, rewriter, {}, loc); + + SmallVector outs; + for (const auto &resWrapped : outsWrapped) { + outs.push_back(resWrapped[0]); + } + rewriter.replaceOp(op, outs); + + return success(); +} + +static LogicalResult processVectorInlineAsm(triton::ElementwiseInlineAsmOp op, + PatternRewriter &rewriter) { + Location loc = op.getLoc(); + + SmallVector> unpackedOperands; + for (auto operand : op.getOperands()) { + auto unpackedOperand = unpackElements(loc, operand, rewriter); + unpackedOperands.push_back(unpackedOperand); + } + + int64_t resultLength = getTensorNumElements(op->getResult(0)); + if (resultLength % op.getPackedElement()) { + op.emitError("Result tensor should be diveded to pack"); + return failure(); + } + + SmallVector> unpackedResults(op->getNumResults()); + for (int64_t i = 0; i < resultLength; i += op.getPackedElement()) { + // Block of elements to process with one call to the inline asm. This is + // ordered opposite `unpackedResults`: The outer dim is + // op.getPackedElement(), and the inner dim is the operand. + SmallVector> block(op.getPackedElement()); + for (auto &os : unpackedOperands) { + for (int j = 0; j < op.getPackedElement(); j++) { + block[j].push_back(os[i + j]); + } + } + auto cur = createDestOps(op, rewriter, block, loc); + assert(cur.size() == unpackedResults.size()); + for (unsigned j = 0; j < cur.size(); j++) { + unpackedResults[j].insert(unpackedResults[j].end(), cur[j].begin(), + cur[j].end()); + } + } + // Reorder and pack the results. + SmallVector outs; + for (int i = 0; i < unpackedResults.size(); i++) { + outs.push_back(rewriter.create( + loc, op->getResult(i).getType(), unpackedResults[i])); + } + rewriter.replaceOp(op, outs); + + return success(); +} + +} // namespace + +struct ElementwiseInlineAsmOpConversion + : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::ElementwiseInlineAsmOp op, + PatternRewriter &rewriter) const final { + return op.getOperands().empty() ? processScalarInlineAsm(op, rewriter) + : processVectorInlineAsm(op, rewriter); + } +}; + +void TritonToLLVMPass::runOnOperation() { + auto module = getOperation(); + ConversionTarget target(getContext()); + target.addLegalDialect(); + + RewritePatternSet patterns(&getContext()); + patterns.add(patterns.getContext()); + if (failed(applyPartialConversion(module, target, std::move(patterns)))) { + signalPassFailure(); + } +} + +std::unique_ptr> +mlir::triton::createTritonToLLVMPass() { + return std::make_unique(); +} diff --git a/third_party/ascend/python/src/ir.cc b/third_party/ascend/python/src/ir.cc new file mode 100644 index 0000000000..949d21aea0 --- /dev/null +++ b/third_party/ascend/python/src/ir.cc @@ -0,0 +1,2124 @@ +#include "./ir.h" + +#include +#include +#include +#include +#include + +#include "mlir/Bytecode/BytecodeWriter.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +// Needed by the CommonIR lowering path, which is not TLE DSA specific. +#include "mlir/Dialect/Func/Extensions/InlinerExtension.h" +#include "mlir/Dialect/Index/IR/IndexDialect.h" +// The CANN vec_ops extension builds tensor.extract_slice / insert_slice +// directly, so the tensor dialect must be loaded in the builder's context. +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Dialect/LLVMIR/LLVMAttrs.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/UB/IR/UBOps.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Diagnostics.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/Verifier.h" +#include "mlir/Parser/Parser.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Support/FileUtilities.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h" +#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h" +#include "mlir/Transforms/LocationSnapshot.h" + +#include "tle/dialect/include/IR/Dialect.h" +#include "triton/Conversion/TritonGPUToLLVM/Utility.h" +#include "triton/Dialect/Gluon/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Types.h" +#include "triton/Dialect/Triton/IR/Utility.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "triton/Dialect/TritonInstrument/IR/Dialect.h" +#include "triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h" +#include "triton/Tools/Sys/GetEnv.hpp" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/SourceMgr.h" + +void setAsyncTaskIds(mlir::Operation *op, + llvm::ArrayRef asyncTaskIds) { + llvm::SmallVector sortedAsyncTaskIds(asyncTaskIds.begin(), + asyncTaskIds.end()); + sort(sortedAsyncTaskIds); + auto i32Ty = IntegerType::get(op->getContext(), 32); + auto size = static_cast(sortedAsyncTaskIds.size()); + auto vecTy = VectorType::get(size, i32Ty); + op->setAttr("async_task_id", + DenseI32ArrayAttr::get(op->getContext(), sortedAsyncTaskIds)); +} +namespace { + +namespace py = pybind11; +using namespace mlir; +using namespace triton; +namespace tt = triton; +namespace ttg = triton::gpu; +namespace ttng = triton::nvidia_gpu; + +llvm::raw_fd_ostream &mlir_dumps() { + std::error_code EC; + static llvm::raw_fd_ostream S(::triton::tools::getStrEnv("MLIR_DUMP_PATH"), + EC, llvm::sys::fs::CD_CreateAlways); + assert(!EC); + return S; +} + +llvm::raw_ostream &mlir_dumps_or_dbgs() { + if (!::triton::tools::getStrEnv("MLIR_DUMP_PATH").empty()) { + return mlir_dumps(); + } else { + return llvm::dbgs(); + } +} + +// Function to parse a comma-separated string into a vector of C-style strings +llvm::SmallVector +parseCommaSeparatedValues(const std::string &input, + llvm::SmallVector &storage) { + llvm::SmallVector split; + llvm::SmallVector result; + StringRef(input.c_str()).split(split, ','); + llvm::transform(split, std::back_inserter(result), [&storage](StringRef str) { + // StringRefs are not always null-terminated. + // The purpose for this storage pattern is to + // produce a collection of C-strings that are. + storage.push_back(str.str()); + return storage.back().c_str(); + }); + return result; +} + +// Run the pass manager under a source manager diagnostic handler, which +// enables emitted MLIR diagnostics to directly reference Python source +// code. This diagnostic handler supports filtering diagnostic info by +// severity levels. +struct TritonSourceMgrDiagnosticHandler : public SourceMgrDiagnosticHandler { + TritonSourceMgrDiagnosticHandler(MLIRContext *ctx, + DiagnosticSeverity minSeverity) + : SourceMgrDiagnosticHandler(sourceMgr, ctx, llvm::errs()) { + setHandler([this, minSeverity](Diagnostic &diag) { + auto severity = diag.getSeverity(); + switch (severity) { + case DiagnosticSeverity::Error: + break; + case DiagnosticSeverity::Warning: + if (minSeverity == DiagnosticSeverity::Error) + return success(); + break; + case DiagnosticSeverity::Remark: + if (minSeverity == DiagnosticSeverity::Error || + minSeverity == DiagnosticSeverity::Warning) + return success(); + break; + case DiagnosticSeverity::Note: + // notes are handled somewhere else. + return failure(); + default: + llvm_unreachable("Unknown diagnostic severity"); + } + emitDiagnostic(diag); + return success(); + }); + } + + llvm::SourceMgr sourceMgr; +}; + +TritonSourceMgrDiagnosticHandler +setupTritonDiagnosticHandler(MLIRContext *context) { + bool showOperations = false, showStacktraces = false, showRemarks = false, + showWarnings = false; + + if (auto enableDiagnostics = + triton::tools::getStrEnv("MLIR_ENABLE_DIAGNOSTICS"); + !enableDiagnostics.empty()) { + llvm::SmallVector storage; + parseCommaSeparatedValues(enableDiagnostics, storage); + for (auto &str : storage) { + if (str == "warnings") { + showWarnings = true; + } else if (str == "remarks") { + showRemarks = true; + } else if (str == "stacktraces") { + showStacktraces = true; + } else if (str == "operations") { + showOperations = true; + } + // we show errors by default, so no need to set it + } + } + + DiagnosticSeverity minSeverity = + showWarnings ? DiagnosticSeverity::Warning : DiagnosticSeverity::Error; + minSeverity = showRemarks ? DiagnosticSeverity::Remark : minSeverity; + + context->printOpOnDiagnostic(showOperations); + context->printStackTraceOnDiagnostic(showStacktraces); + if (showStacktraces) { + context->disableMultithreading(); + } + + return TritonSourceMgrDiagnosticHandler(context, minSeverity); +} + +std::string locationToString(Location loc) { + std::string str; + llvm::raw_string_ostream os(str); + loc.print(os); + os.flush(); // Make sure all the content is dumped into the 'str' string + return str; +} + +void outputWarning(Location loc, const std::string &msg) { + std::string locStr = locationToString(loc); + + PyErr_WarnEx(PyExc_UserWarning, (locStr + ": " + msg).c_str(), + /*stack_level=*/2); +} + +// Allow dump a reproducer in the console on crash. +struct ConsoleReproducerStream : public mlir::ReproducerStream { + ~ConsoleReproducerStream() override {} + + StringRef description() override { + return "std::errs, please share the reproducer above with Triton project."; + } + raw_ostream &os() override { return llvm::errs(); } +}; + +ReproducerStreamFactory makeConsoleReproducer() { + return [](std::string &error) -> std::unique_ptr { + return std::make_unique(); + }; +} + +OpPrintingFlags getOpPrintingFlags() { + auto printingFlags = OpPrintingFlags(); + printingFlags.enableDebugInfo(); + printingFlags.printNameLocAsPrefix(true); + return printingFlags; +} + +py::list getTensorDescMetadata(ModuleOp &mod) { + py::list result; + triton::FuncOp kernelFunc; + mod.walk([&](triton::FuncOp func) { + if (triton::isKernel(func)) { + kernelFunc = func; + return WalkResult::interrupt(); + } + return WalkResult::skip(); + }); + assert(kernelFunc); + + for (auto [i, argTy] : llvm::enumerate(kernelFunc.getArgumentTypes())) { + auto descTy = dyn_cast(argTy); + if (!descTy) + continue; + + auto blockType = descTy.getBlockType(); + auto encoding = blockType.getEncoding(); + auto mmaEncoding = dyn_cast(encoding); + auto swizzle = ttng::getTMASwizzleMode(nullptr, descTy); + auto elemType = ttng::getTMAElementType(nullptr, descTy); + assert(swizzle.has_value()); + assert(elemType.has_value()); + auto blockSize = ttng::getTMABlockShape(blockType, /*packedSize=*/false); + py::dict metadata; + metadata["swizzle"] = *swizzle; + metadata["elem_size"] = descTy.getBlockType().getElementTypeBitWidth() / 8; + metadata["elem_type"] = *elemType; + metadata["block_size"] = + std::vector(blockSize.begin(), blockSize.end()); + metadata["fp4_padded"] = mmaEncoding && mmaEncoding.getFp4Padded(); + result.append(std::move(metadata)); + } + return result; +} + +} // anonymous namespace + +/*****************************************************************************/ +/* Python bindings for ir */ +/*****************************************************************************/ + +namespace ir { + +// Pointer to the TritonOpBuilder class, used to register IR ops for third-party +// dialects. +static py::class_ *builderClassPtr = nullptr; +py::class_ *getBuilderClass() { return builderClassPtr; } + +} // namespace ir + +void init_triton_ir(py::module &&m) { + using ret = py::return_value_policy; + using namespace pybind11::literals; + + py::enum_(m, "PADDING_OPTION", py::module_local()) + .value("PAD_ZERO", PaddingOption::PAD_ZERO) + .value("PAD_NAN", PaddingOption::PAD_NAN) + .export_values(); + + py::enum_(m, "CACHE_MODIFIER", py::module_local()) + .value("NONE", CacheModifier::NONE) + .value("CA", CacheModifier::CA) + .value("CG", CacheModifier::CG) + .value("WB", CacheModifier::WB) + .value("CS", CacheModifier::CS) + .value("WT", CacheModifier::WT) + .value("CV", CacheModifier::CV) + .export_values(); + + py::enum_(m, "MEM_SEMANTIC", py::module_local()) + .value("ACQUIRE_RELEASE", MemSemantic::ACQUIRE_RELEASE) + .value("ACQUIRE", MemSemantic::ACQUIRE) + .value("RELEASE", MemSemantic::RELEASE) + .value("RELAXED", MemSemantic::RELAXED) + .export_values(); + + py::enum_(m, "MEM_SYNC_SCOPE", py::module_local()) + .value("GPU", MemSyncScope::GPU) + .value("CTA", MemSyncScope::CTA) + .value("SYSTEM", MemSyncScope::SYSTEM) + .export_values(); + + py::enum_(m, "EVICTION_POLICY", py::module_local()) + .value("NORMAL", EvictionPolicy::NORMAL) + .value("EVICT_FIRST", EvictionPolicy::EVICT_FIRST) + .value("EVICT_LAST", EvictionPolicy::EVICT_LAST) + .export_values(); + + py::enum_(m, "ATOMIC_OP", py::module_local()) + .value("ADD", RMWOp::ADD) + .value("FADD", RMWOp::FADD) + .value("AND", RMWOp::AND) + .value("OR", RMWOp::OR) + .value("XOR", RMWOp::XOR) + .value("XCHG", RMWOp::XCHG) + .value("MAX", RMWOp::MAX) + .value("MIN", RMWOp::MIN) + .value("UMIN", RMWOp::UMIN) + .value("UMAX", RMWOp::UMAX); + + py::enum_(m, "DESCRIPTOR_REDUCE_KIND", + py::module_local()) + .value("ADD", DescriptorReduceKind::ADD) + .value("AND", DescriptorReduceKind::AND) + .value("OR", DescriptorReduceKind::OR) + .value("XOR", DescriptorReduceKind::XOR) + .value("MAX", DescriptorReduceKind::MAX) + .value("MIN", DescriptorReduceKind::MIN) + .value("INC", DescriptorReduceKind::INC) + .value("DEC", DescriptorReduceKind::DEC); + + py::enum_(m, "ROUNDING_MODE", py::module_local()) + .value("RTZ", RoundingMode::RTZ) + .value("RTNE", RoundingMode::RTNE); + + py::enum_(m, "PROPAGATE_NAN", py::module_local()) + .value("NONE", PropagateNan::NONE) + .value("ALL", PropagateNan::ALL); + + py::enum_(m, "INPUT_PRECISION", py::module_local()) + .value("TF32", InputPrecision::TF32) + .value("TF32x3", InputPrecision::TF32x3) + .value("IEEE", InputPrecision::IEEE) + .value("HF32", InputPrecision::HF32) + .export_values(); + + py::enum_(m, "ScaleDotElemTypeTY", py::module_local()) + .value("E4M3", ScaleDotElemType::E4M3) + .value("E5M2", ScaleDotElemType::E5M2) + .value("E2M3", ScaleDotElemType::E2M3) + .value("E3M2", ScaleDotElemType::E3M2) + .value("E2M1", ScaleDotElemType::E2M1) + .value("BF16", ScaleDotElemType::BF16) + .value("FP16", ScaleDotElemType::FP16) + .export_values(); + + py::class_(m, "context", py::module_local()) + .def(py::init<>()) + .def( + "__enter__", [](MLIRContext &self) -> MLIRContext & { return self; }, + py::return_value_policy::reference) + .def("__exit__", + [](MLIRContext &, py::object, py::object, py::object) -> bool { + // Keep context alive for the duration of the scope. + return false; + }) + .def("printOpOnDiagnostic", + [](MLIRContext &self, bool v) { self.printOpOnDiagnostic(v); }) + .def("printStackTraceOnDiagnostic", + [](MLIRContext &self, bool v) { + self.printStackTraceOnDiagnostic(v); + }) + .def("disable_multithreading", + [](MLIRContext &self) { self.disableMultithreading(); }); + + py::class_(m, "source_mgr_diag", + py::module_local()) + .def(py::init()); + + m.def("load_dialects", [](MLIRContext &context) { + DialectRegistry registry; + registry.insert< + TritonDialect, ::mlir::triton::gpu::TritonGPUDialect, + ::mlir::triton::instrument::TritonInstrumentDialect, math::MathDialect, + arith::ArithDialect, scf::SCFDialect, ::mlir::gpu::GPUDialect, + cf::ControlFlowDialect, LLVM::LLVMDialect, mlir::ub::UBDialect, + mlir::triton::gluon::GluonDialect, mlir::triton::tle::TleDialect, + // func is required by commonir-to-hivm's inliner extension and index + // by the CommonIR lowering; neither is TLE DSA specific. + ::mlir::index::IndexDialect, func::FuncDialect, + // tensor is required by the CANN vec_ops extract_slice/insert_slice + // builders; without it MLIR aborts the process with + // "Building op `tensor.extract_slice` but it isn't known". + ::mlir::tensor::TensorDialect>(); + mlir::LLVM::registerInlinerInterface(registry); + registerBuiltinDialectTranslation(registry); + registerLLVMDialectTranslation(registry); + mlir::LLVM::registerInlinerInterface(registry); + // CommonIR lowering needs the func inliner interface registered. The + // commonir-to-hivm pass runs unconditionally (see backend/compiler.py), + // and the Inliner pass right after it aborts without this extension, so + // it must not be gated on __TLE_DSA__. + func::registerInlinerExtension(registry); + context.appendDialectRegistry(registry); + context.loadAllAvailableDialects(); + }); + + py::class_(m, "type", py::module_local()) + .def("is_integer", + [](Type &self, unsigned width) { return self.isInteger(width); }) + .def("is_fp16", &Type::isF16) + .def("__eq__", + [](Type &self, py::object &other) { + Type *other_ty = py::cast(other); + return (other_ty != nullptr) && (*other_ty == self); + }) + .def("__ne__", + [](Type &self, py::object &other) { + Type *other_ty = py::cast(other); + return (other_ty == nullptr) || (*other_ty != self); + }) + .def("__str__", [](Type &self) { + std::string str; + llvm::raw_string_ostream os(str); + self.print(os); + return os.str(); + }); + + py::class_(m, "function_type", py::module_local()) + .def("param_types", [](FunctionType &self) { + return std::vector(self.getInputs().begin(), + self.getInputs().end()); + }); + + py::class_(m, "location", py::module_local()) + .def("__str__", [](Location &self) { + std::string str; + llvm::raw_string_ostream os(str); + self.print(os); + return os.str(); + }); + + py::class_(m, "value", py::module_local()) + .def(py::init<>()) + .def("set_attr", + [](Value &self, std::string &name, Attribute &attr) -> void { + if (Operation *definingOp = self.getDefiningOp()) + definingOp->setAttr(name, attr); + else { + auto arg = mlir::cast(self); + int id = arg.getArgNumber(); + std::string attrName = name + "_arg" + std::to_string(id); + Block *owner = arg.getOwner(); + if (owner->isEntryBlock() && + !isa(owner->getParentOp())) { + owner->getParentOp()->setAttr(attrName, attr); + } + } + }) + .def("get_context", &Value::getContext) + .def("replace_all_uses_with", + [](Value &self, Value &newValue) { + self.replaceAllUsesWith(newValue); + }) + .def("get_type", &Value::getType) + .def("id", + [](Value &self) { + // The Value is identified by and compared with + // other Values via the underlying ValueImpl + return (uint64_t)self.getImpl(); + }) + .def("set_loc", + [](Value &self, Location loc) { return self.setLoc(loc); }) + .def("get_loc", [](Value &self) { return self.getLoc(); }); + + py::class_(m, "op_result", py::module_local()); + + py::class_(m, "block_argument", py::module_local()); + + py::class_(m, "region", py::module_local()) + .def("get_parent_region", &Region::getParentRegion, ret::reference) + .def("size", [](Region &self) { return self.getBlocks().size(); }) + .def("empty", &Region::empty) + .def("id", [](Region &self) { return (uint64_t)&self; }) + .def("push_back", + [](Region &self, Block *block) { self.push_back(block); }) + .def("push_front", + [](Region &self, Block *block) { self.push_front(block); }); + + py::class_(m, "block", py::module_local()) + .def("arg", + [](Block &self, int index) -> BlockArgument { + if (index >= self.getNumArguments()) + throw pybind11::index_error("Block argument index out of range"); + return self.getArgument(index); + }) + .def("add_argument", + [](Block &self, Type ty) { + auto loc = UnknownLoc::get(ty.getContext()); + self.addArgument(ty, loc); + }) + .def("add_argument_at", [](Block &self, Type ty, + Location loc) { self.addArgument(ty, loc); }) + .def("get_num_arguments", &Block::getNumArguments) + .def("get_argument", &Block::getArgument) + .def("dump", &Block::dump) + .def("move_before", + [](Block &self, Block &dst) { self.moveBefore(&dst); }) + .def("insert_before", &Block::insertBefore) + .def("get_parent", &Block::getParent, ret::reference) + .def("merge_block_before", + [](Block &self, Block &dst) { + // ref: RewriterBase::mergeBlocks() + if (self.getNumArguments() != 0) + throw std::runtime_error( + "This block has arguments, don't merge"); + dst.getOperations().splice(dst.begin(), self.getOperations()); + self.dropAllUses(); + self.erase(); + }) + .def("replace_use_in_block_with", + [](Block &self, Value &v, Value &newVal) { + v.replaceUsesWithIf(newVal, [&](OpOperand &operand) { + Operation *user = operand.getOwner(); + Block *currentBlock = user->getBlock(); + while (currentBlock) { + if (currentBlock == &self) + return true; + // Move up one level + currentBlock = + currentBlock->getParent()->getParentOp()->getBlock(); + } + return false; + }); + }) + .def("__str__", + [](Block &self) { + std::string str; + llvm::raw_string_ostream os(str); + self.print(os); + return str; + }) + .def("has_terminator", + [](Block &self) { + return !self.empty() && + self.back().hasTrait(); + }) + .def("has_return", + [](Block &self) { + return !self.empty() && + self.back().hasTrait(); + }) + .def("erase", [](Block &self) { self.erase(); }) + .def("id", [](Block &self) { return (uint64_t)&self; }); + + py::class_(m, "attribute", py::module_local()); + py::class_(m, "integer_attr", py::module_local()); + py::class_(m, "bool_attr", py::module_local()); + py::class_(m, "unit_attr", py::module_local()); + py::class_(m, "str_attr", py::module_local()); + py::class_(m, "array_attr", py::module_local()); + + // Ops + py::class_(m, "OpState", py::module_local()) + .def("set_attr", + [](OpState &self, std::string &name, Attribute &attr) -> void { + self->setAttr(name, attr); + }) + .def("get_num_results", + [](OpState &self) -> unsigned { return self->getNumResults(); }) + .def("get_result", + [](OpState &self, unsigned idx) -> Value { + if (idx >= self->getNumResults()) + throw pybind11::index_error("Op result index out of range"); + return self->getResult(idx); + }) + .def( + "get_region", + [](OpState &self, unsigned idx) -> Region & { + if (idx >= self->getNumRegions()) + throw pybind11::index_error("Op region index out of range"); + return self->getRegion(idx); + }, + ret::reference) + .def( + "get_body", + [](scf::ForOp &self, unsigned idx) -> Block * { + if (idx >= self->getNumRegions()) + throw pybind11::index_error("Op region index out of range"); + return self.getBody(idx); + }, + ret::reference) + .def("dump", [](OpState &self) { self->dump(); }) + .def("__str__", + [](OpState &self) -> std::string { + std::string str; + llvm::raw_string_ostream os(str); + auto printingFlags = getOpPrintingFlags(); + self->print(os, printingFlags); + return str; + }) + .def("str_nodebug", + [](OpState &self) -> std::string { + std::string str; + llvm::raw_string_ostream os(str); + self->print(os); + return str; + }) + .def("append_operand", + [](OpState &self, Value &val) { + self->insertOperands(self->getNumOperands(), val); + }) + .def("verify", + [](OpState &self) -> bool { + return succeeded(verify(self.getOperation())); + }) + .def("get_operation", [](OpState &self) { return self.getOperation(); }); + + // scf Ops + py::class_(m, "ForOp", py::module_local()) + .def("get_induction_var", &scf::ForOp::getInductionVar); + + py::class_(m, "IfOp", py::module_local()) + .def("get_then_block", &scf::IfOp::thenBlock, ret::reference) + .def("get_else_block", &scf::IfOp::elseBlock, ret::reference) + .def("get_then_yield", &scf::IfOp::thenYield) + .def("get_else_yield", &scf::IfOp::elseYield); + py::class_(m, "YieldOp", py::module_local()); + py::class_(m, "WhileOp", py::module_local()) + .def("get_before", &scf::WhileOp::getBefore, ret::reference) + .def("get_after", &scf::WhileOp::getAfter, ret::reference); + py::class_(m, "ConditionOp", py::module_local()); + + py::class_>( + m, "operation", py::module_local()) + .def("get_name", + [](Operation &self) { + llvm::StringRef opName = self.getName().getStringRef(); + return opName.str(); + }) + .def("get_num_operands", &Operation::getNumOperands) + .def("get_operand", &Operation::getOperand) + .def("get_num_results", &Operation::getNumResults) + .def("get_result", &Operation::getResult) + .def("get_num_regions", &Operation::getNumRegions) + .def("get_region", &Operation::getRegion, ret::reference) + .def("get_block", &Operation::getBlock, ret::reference) + .def("get_str_attr", + [](Operation &self, const std::string &name) -> py::object { + auto ret = self.getAttrOfType(name); + if (!ret) + return py::none(); + return py::str(ret.getValue().str()); + }) + .def("get_bool_attr", + [](Operation &self, const std::string &name) -> py::object { + auto ret = self.getAttrOfType(name); + if (!ret) + return py::none(); + return py::bool_(ret.getValue()); + }) + .def("get_flat_symbol_ref_attr", + [](Operation &self, const std::string &name) -> py::object { + auto ret = self.getAttrOfType(name); + if (!ret) + return py::none(); + return py::str(ret.getValue().str()); + }); + + // dynamic_attr is used to transfer ownership of the MLIR context to the + // module + py::class_(m, "module", py::module_local(), + py::dynamic_attr()) + .def("dump", &ModuleOp::dump) + .def("str", + [](ModuleOp &self) -> std::string { + std::string str; + llvm::raw_string_ostream os(str); + auto printingFlags = getOpPrintingFlags(); + self.print(os, printingFlags); + return str; + }) + .def("push_back", + [](ModuleOp &self, FuncOp &funcOp) -> void { + self.push_back(funcOp); + }) + .def("get_entry_func_name", + [](ModuleOp &self) -> std::string { + for (auto &op : self.getOps()) { + if (auto func = dyn_cast(op)) { + if (triton::isKernel(func)) + return func.getName().str(); + } + } + return ""; + }) + .def("has_function", + [](ModuleOp &self, std::string &funcName) -> bool { + if (self.lookupSymbol(funcName)) + return true; + return false; + }) + .def("get_function", + [](ModuleOp &self, std::string &funcName) -> FuncOp { + return self.lookupSymbol(funcName); + }) + /* + * def ty_to_cpp(ty) is the consumer of this function. + * If the type is a ptr it expects ty[0] == '*', else the type itself. + */ + + .def("get_function_signature", + [](ModuleOp &self, FuncOp &func) -> std::vector { + std::vector strVec; + + auto type = func.getFunctionType(); + unsigned numArgs = type.getNumInputs(); + for (unsigned i = 0; i != numArgs; ++i) { + std::string tempType; + llvm::raw_string_ostream os(tempType); + + auto ty = type.getInput(i); + if (auto attributes = func.getCallableArgAttrs()) { + Attribute attr = attributes[i]; + // Check for tt.nv_tma_desc = 1 + if (auto dAttr = dyn_cast(attr)) { + if (dAttr.contains("tt.nv_tma_desc")) { + strVec.push_back("nvTmaDesc"); + continue; + } + } + } + if (auto ptrType = dyn_cast(ty)) { + auto pType = ptrType.getPointeeType(); + os << "*"; + pType.print(os); + } else { + ty.print(os); + } + strVec.push_back(tempType); + } + return strVec; + }) + .def("get_int_attr", + [](ModuleOp &self, std::string name) -> py::object { + auto ret = self->getAttrOfType(name); + if (!ret) + return py::none(); + return py::int_(ret.getInt()); + }) + .def("get_tensordesc_metadata", getTensorDescMetadata) + .def("create_location_snapshot", + [](ModuleOp &self, const std::string &fileName) -> void { + auto printingFlags = getOpPrintingFlags(); + if (failed(generateLocationsFromIR(fileName, self, printingFlags))) + throw std::runtime_error("Failed to create location snapshot"); + }) + .def("walk", + [](ModuleOp &self, const std::function &fn) { + self.walk(fn); + }) + .def("verify_with_diagnostics", [](ModuleOp &self) { + TritonSourceMgrDiagnosticHandler handler = + setupTritonDiagnosticHandler(self.getContext()); + return succeeded(verify(self.getOperation())); + }); + + m.def("make_attr", [](const std::vector &values, MLIRContext &context) { + return mlir::cast(DenseIntElementsAttr::get( + RankedTensorType::get({static_cast(values.size())}, + IntegerType::get(&context, 32)), + values)); + }); + + m.def( + "parse_mlir_module", + [](const std::string &inputFilename, MLIRContext &context) { + // parse module + OwningOpRef module = + parseSourceFile(inputFilename, &context); + if (!module) + throw std::runtime_error("Parse MLIR file failed."); + return module->clone(); + }, + ret::take_ownership); + + py::class_(m, "function", py::module_local()) + // .def_property_readonly("attrs", &ir::function::attrs) + // .def("add_attr", &ir::function::add_attr); + .def("args", + [](FuncOp &self, unsigned idx) -> BlockArgument { + if (idx >= self.getNumArguments()) + throw pybind11::index_error( + "Function argument index out of range"); + return self.getArgument(idx); + }) + .def("get_num_args", &FuncOp::getNumArguments) + .def( + "add_entry_block", + [](FuncOp &self) -> Block * { return self.addEntryBlock(); }, + ret::reference) + .def( + "set_arg_attr", + [](FuncOp &self, int arg_no, const std::string &name, int val) { + if (arg_no >= self.getNumArguments()) + throw pybind11::index_error( + "Function argument index out of range"); + // set arg attributes "name" to value "val" + auto attrTy = IntegerType::get(self.getContext(), 32); + self.setArgAttr(arg_no, name, IntegerAttr::get(attrTy, val)); + }, + ret::reference) + // .def("has_attr", &::FuncOp::hasAttr) + .def("finalize", [](FuncOp &self) -> void {}) + .def_property_readonly("type", &FuncOp::getFunctionType) + .def("reset_type", &FuncOp::setType); + + py::class_(m, "op_builder", py::module_local(), + py::dynamic_attr()) + .def(py::init()); + + py::class_(m, "InsertPoint", py::module_local()); + + static py::class_ builderClass( + m, "builder", py::module_local(), py::dynamic_attr()); + ir::builderClassPtr = &builderClass; + builderClass + .def(py::init(), py::arg("context"), + py::arg("compile_mode") = "simd", + "Create a TritonOpBuilder with optional compile_mode (simt or simd, " + "default: simd)") + .def("is_simt_mode", &TritonOpBuilder::isSimtMode, + "Check if the compile mode is simt") + .def("get_op_builder", &TritonOpBuilder::getBuilder, ret::reference) + // getters + .def("create_module", + [](TritonOpBuilder &self) -> ModuleOp { + return self.create(); + }) + // insertion block/point + .def("set_insertion_point_to_start", + [](TritonOpBuilder &self, Block &block) -> void { + self.setInsertionPointToStart(block); + }) + .def("set_insertion_point_to_end", + [](TritonOpBuilder &self, Block &block) { + self.setInsertionPointToEnd(block); + }) + .def("set_insertion_point_after", + [](TritonOpBuilder &self, Operation &op) { + self.setInsertionPointAfter(op); + }) + .def( + "get_insertion_block", + [](TritonOpBuilder &self) -> Block * { + return self.getBuilder().getInsertionBlock(); + }, + ret::reference) + .def("get_insertion_point", + [](TritonOpBuilder &self) { + return self.getBuilder().saveInsertionPoint(); + }) + .def("restore_insertion_point", + [](TritonOpBuilder &self, OpBuilder::InsertPoint pt) { + self.restoreInsertionPoint(pt); + }) + // Attr + .def( + "get_unit_attr", + [](TritonOpBuilder &self) { return self.getBuilder().getUnitAttr(); }) + .def("get_bool_attr", + [](TritonOpBuilder &self, bool value) { + return self.getBuilder().getBoolAttr(value); + }) + .def("get_int32_attr", + [](TritonOpBuilder &self, int32_t value) { + return self.getBuilder().getI32IntegerAttr(value); + }) + .def("get_string_attr", + [](TritonOpBuilder &self, std::string value) -> Attribute { + return self.getBuilder().getStringAttr(value); + }) + .def("get_i64_array_attr", + [](TritonOpBuilder &self, const std::vector &array) { + return self.getBuilder().getI64ArrayAttr(array); + }) + .def("get_type_array_attr", + [](TritonOpBuilder &self, const std::vector &array) { + return self.getBuilder().getTypeArrayAttr(array); + }) + .def("get_disable_loop_licm_attr", + [](TritonOpBuilder &self) -> Attribute { + auto licmAttr = + LLVM::LoopLICMAttr::get(self.getBuilder().getContext(), + self.getBuilder().getBoolAttr(true), + self.getBuilder().getBoolAttr(true)); + mlir::LLVM::LoopAnnotationAttr la = + mlir::LLVM::LoopAnnotationAttr::get( + self.getBuilder().getContext(), {}, {}, {}, {}, {}, + licmAttr, {}, {}, {}, {}, {}, {}, {}, {}, {}); + return la; + }) + // Use arith.ConstantOp to create constants + // Constants + .def("get_int1", + [](TritonOpBuilder &self, bool v) -> Value { + return Value(self.create( + self.getBuilder().getI1Type(), v)); + }) + .def("get_int8", + [](TritonOpBuilder &self, int64_t v) -> Value { + return Value(self.create( + self.getBuilder().getI8Type(), v)); + }) + .def("get_int16", + [](TritonOpBuilder &self, int64_t v) -> Value { + return Value(self.create( + self.getBuilder().getI16Type(), v)); + }) + .def("get_int32", + [](TritonOpBuilder &self, int64_t v) -> Value { + return Value(self.create( + self.getBuilder().getI32Type(), v)); + }) + .def("get_int64", + [](TritonOpBuilder &self, int64_t v) -> Value { + return Value(self.create( + self.getBuilder().getI64Type(), v)); + }) + .def("get_uint8", + [](TritonOpBuilder &self, uint64_t v) -> Value { + return Value(self.create( + self.getBuilder().getI8Type(), v)); + }) + .def("get_uint16", + [](TritonOpBuilder &self, uint64_t v) -> Value { + return Value(self.create( + self.getBuilder().getI16Type(), v)); + }) + .def("get_uint32", + [](TritonOpBuilder &self, uint64_t v) -> Value { + return Value(self.create( + self.getBuilder().getI32Type(), v)); + }) + .def("get_uint64", + [](TritonOpBuilder &self, uint64_t v) -> Value { + return Value(self.create( + self.getBuilder().getI64Type(), v)); + }) + .def("get_bf16", + [](TritonOpBuilder &self, float v) -> Value { + auto type = self.getBuilder().getBF16Type(); + return self.create( + type, APFloat(type.getFloatSemantics(), std::to_string(v))); + }) + .def("get_fp16", + [](TritonOpBuilder &self, float v) -> Value { + return self.create( + self.getBuilder().getF16FloatAttr(v)); + }) + .def("get_fp32", + [](TritonOpBuilder &self, float v) -> Value { + return self.create( + self.getBuilder().getF32FloatAttr(v)); + }) + .def("get_fp64", + [](TritonOpBuilder &self, double v) -> Value { + return self.create( + self.getBuilder().getF64FloatAttr(v)); + }) + .def("get_fp8e4nv", + [](TritonOpBuilder &self, double v) -> Value { + return self.create(FloatAttr::get( + Float8E4M3FNType::get((self.getBuilder().getContext())), v)); + }) + .def("get_fp8e5", + [](TritonOpBuilder &self, double v) -> Value { + return self.create(FloatAttr::get( + Float8E5M2Type::get(self.getBuilder().getContext()), v)); + }) + .def("get_null_value", + [](TritonOpBuilder &self, Type type) -> Value { + if (auto floatTy = dyn_cast(type)) + return self.create( + floatTy, APFloat(floatTy.getFloatSemantics(), 0)); + else if (auto intTy = dyn_cast(type)) + return self.create(intTy, 0); + else + throw std::runtime_error("Not implemented"); + }) + .def("get_all_ones_value", + [](TritonOpBuilder &self, Type type) -> Value { + uint64_t val = 0xFFFFFFFFFFFFFFFF; + if (auto intTy = dyn_cast(type)) + return self.create(intTy, val); + else + throw std::runtime_error("Not implemented"); + }) + + // Types + .def("get_void_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getNoneType(); + }) + .def("get_int1_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getI1Type(); + }) // or ret::copy? + .def("get_int8_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getI8Type(); + }) + .def("get_int16_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getType(16); + }) + .def("get_int32_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getI32Type(); + }) + .def("get_int64_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getI64Type(); + }) + .def("get_fp8e4nv_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getType(); + }) + .def("get_fp8e4b8_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getType(); + }) + .def("get_fp8e4b15_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getI8Type(); + }) + .def("get_fp8e5_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getType(); + }) + .def("get_fp8e5b16_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getType(); + }) + .def("get_half_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getF16Type(); + }) + .def("get_bf16_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getBF16Type(); + }) + .def("get_float_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getF32Type(); + }) + .def("get_double_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getF64Type(); + }) + .def("get_ptr_ty", + [](TritonOpBuilder &self, Type &type, int addrSpace) -> Type { + return PointerType::get(type, addrSpace); + }) + .def("get_block_ty", + [](TritonOpBuilder &self, Type &elementType, + std::vector &shape) -> Type { + return RankedTensorType::get(shape, elementType); + }) + .def("get_buffer_ty", + [](TritonOpBuilder &self, std::vector &shape, + Type &elementType, const Attribute &memorySpace) -> Type { + return MemRefType::get(shape, elementType, + MemRefLayoutAttrInterface{}, memorySpace); + }) + .def("get_buffer_ty_with_strides", + [](TritonOpBuilder &self, std::vector &shape, + Type &elementType, const std::vector &strides, + const Attribute &memorySpace) -> Type { + // create a layout with strides, using dynamic offset + auto layout = StridedLayoutAttr::get( + self.getBuilder().getContext(), ShapedType::kDynamic, strides); + return MemRefType::get(shape, elementType, layout, memorySpace); + }) + .def("get_function_ty", + [](TritonOpBuilder &self, std::vector inTypes, + std::vector outTypes) -> Type { + return self.getBuilder().getFunctionType(inTypes, outTypes); + }) + // locs + .def("set_loc", + [](TritonOpBuilder &self, Location loc) { self.setLastLoc(loc); }) + .def("set_loc", + [](TritonOpBuilder &self, std::string name) { + auto nameAttr = StringAttr::get(self.getContext(), name); + auto loc = NameLoc::get(nameAttr); + self.setLastLoc(loc); + }) + .def("create_loc", + [](TritonOpBuilder &self, const std::string &fileName, int line, + int column) -> Location { + return mlir::FileLineColLoc::get(self.getContext(), fileName, line, + column); + }) + .def( + "create_name_loc", + [](TritonOpBuilder &self, std::string name, + std::optional childLoc) -> Location { + auto nameAttr = StringAttr::get(self.getContext(), name); + if (childLoc) + return NameLoc::get(nameAttr, *childLoc); + return NameLoc::get(nameAttr); + }, + py::arg("name"), py::arg("child_loc") = py::none()) + .def("set_loc", + [](TritonOpBuilder &self, const std::string &fileName, int line, + int column) { self.setLastLoc(fileName, line, column); }) + .def("get_loc", + [](TritonOpBuilder &self) -> Location { return self.getLastLoc(); }) + + // Ops + .def("get_or_insert_function", + [](TritonOpBuilder &self, ModuleOp &module, std::string &funcName, + Type &funcType, std::string &visibility, + bool noinline) -> FuncOp { + if (Operation *funcOperation = module.lookupSymbol(funcName)) + return llvm::dyn_cast(funcOperation); + if (auto funcTy = dyn_cast(funcType)) { + llvm::SmallVector attrs = { + NamedAttribute( + self.getBuilder().getStringAttr("sym_visibility"), + self.getBuilder().getStringAttr(visibility)), + NamedAttribute(self.getBuilder().getStringAttr("noinline"), + self.getBuilder().getBoolAttr(noinline))}; + return self.create(funcName, funcTy, attrs); + } + throw std::invalid_argument("invalid function type"); + }) + .def( + "create_block", + [](TritonOpBuilder &self) -> Block * { + Region *parent = self.getBuilder().getBlock()->getParent(); + return self.getBuilder().createBlock(parent); + }, + ret::reference) + .def( + "create_block_with_parent", + [](TritonOpBuilder &self, Region &parent, + std::vector &argTypes) -> Block * { + // TODO: update arg loc + auto loc = self.getBuilder().getUnknownLoc(); + llvm::SmallVector argLocs(argTypes.size(), loc); + return self.getBuilder().createBlock(&parent, {}, argTypes, + argLocs); + }, + ret::reference) + .def( + "new_block", + [](TritonOpBuilder &self) -> Block * { return new Block(); }, + ret::reference) + // Function + .def("ret", + [](TritonOpBuilder &self, std::vector &vals) -> OpState { + return self.create(vals); + }) + .def("call", + [](TritonOpBuilder &self, FuncOp &func, std::vector &args) + -> OpState { return self.create(func, args); }) + // Unstructured control flow + .def("create_cond_branch", + [](TritonOpBuilder &self, Value condition, Block *trueDest, + Block *falseDest) -> OpState { + return self.create(condition, trueDest, + falseDest); + }) + .def("create_branch", + [](TritonOpBuilder &self, Block *dest, std::vector &args) + -> OpState { return self.create(dest, args); }) + // Structured control flow + .def("create_for_op", + [](TritonOpBuilder &self, Value &lb, Value &ub, Value &step, + std::vector &initArgs) -> scf::ForOp { + return self.create(lb, ub, step, initArgs); + }) + .def("create_if_op", + [](TritonOpBuilder &self, std::vector &retTypes, + Value &condition, bool withElse) -> scf::IfOp { + return self.create(retTypes, condition, withElse); + }) + .def("create_yield_op", + [](TritonOpBuilder &self, std::vector &yields) + -> scf::YieldOp { return self.create(yields); }) + .def("create_while_op", + [](TritonOpBuilder &self, std::vector &retTypes, + std::vector &initArgs) -> scf::WhileOp { + return self.create(retTypes, initArgs); + }) + .def("create_condition_op", + [](TritonOpBuilder &self, Value &cond, + std::vector &args) -> scf::ConditionOp { + return self.create(cond, args); + }) + + // miscellaneous + .def("create_make_range", + [](TritonOpBuilder &self, Type retTy, int start, int end) -> Value { + return self.create(retTy, start, end); + }) + + // Cast instructions + // Conversions for custom FP types (FP8 and non-standard rounding modes) + .def("create_fp_to_fp", + [](TritonOpBuilder &self, Value &src, Type &dstType, + std::optional roundingMode) -> Value { + if (roundingMode.has_value()) + return self.create( + dstType, src, + RoundingModeAttr::get(self.getBuilder().getContext(), + roundingMode.value())); + else + return self.create(dstType, src); + }) + // Conversions for standard LLVM builtin types + .def("create_bitcast", + [](TritonOpBuilder &self, Value &src, Type &dstType) -> Value { + return self.create(dstType, src); + }) + .def("create_si_to_fp", + [](TritonOpBuilder &self, Value &src, Type &dstType) -> Value { + return self.create(dstType, src); + }) + .def("create_ui_to_fp", + [](TritonOpBuilder &self, Value &src, Type &dstType) -> Value { + return self.create(dstType, src); + }) + .def("create_fp_to_si", + [](TritonOpBuilder &self, Value &src, Type &dstType) -> Value { + return self.create(dstType, src); + }) + .def("create_fp_to_ui", + [](TritonOpBuilder &self, Value &src, Type &dstType) -> Value { + return self.create(dstType, src); + }) + .def("create_fp_ext", + [](TritonOpBuilder &self, Value &src, Type &dstType) -> Value { + return self.create(dstType, src); + }) + .def("create_fp_trunc", + [](TritonOpBuilder &self, Value &src, Type &dstType) -> Value { + return self.create(dstType, src); + }) + .def("create_int_cast", + [](TritonOpBuilder &self, Value &src, Type &dstType, + bool isSigned) -> Value { + // get element type if necessary + Type srcType = src.getType(); + auto srcTensorType = dyn_cast(srcType); + auto dstTensorType = dyn_cast(dstType); + Type srcEltType = srcType; + Type dstEltType = dstType; + if (dstTensorType && srcTensorType) { + dstEltType = dstTensorType.getElementType(); + srcEltType = srcTensorType.getElementType(); + } + unsigned srcWidth = srcEltType.getIntOrFloatBitWidth(); + unsigned dstWidth = dstEltType.getIntOrFloatBitWidth(); + if (srcWidth == dstWidth) + return self.create(dstType, src); + else if (srcWidth > dstWidth) + return self.create(dstType, src); + else if (isSigned) + return self.create(dstType, src); + else + return self.create(dstType, src); + }) + .def("create_fmul", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_fdiv", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_frem", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_fadd", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_fsub", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_mul", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_umulhi", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_sdiv", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_udiv", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_srem", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_urem", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_add", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_sub", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + .def("create_fma", + [](TritonOpBuilder &self, Value &a, Value &b, Value &c) -> Value { + return Value(self.create(a, b, c)); + }) + .def("create_shl", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + .def("create_lshr", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + .def("create_ashr", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + .def("create_minsi", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + .def("create_minui", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + // minimumf follows the torch.minimum convention and returns NaN if either + // operand is NaN + .def("create_minimumf", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + // minnumf follows the torch.fmin convention and returns the non-NaN + // operand + .def("create_minnumf", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + .def("create_maxsi", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + .def("create_maxui", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + // maximumf follows the torch.maximum convention and returns NaN if either + // operand is NaN + .def("create_maximumf", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + // maxnumf follows the torch.fmax convention and returns the non-NaN + // operand + .def("create_maxnumf", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + .def("create_clampf", + [](TritonOpBuilder &self, Value &input, Value &min, Value &max, + PropagateNan propagateNan) -> Value { + return Value(self.create(input, min, max, propagateNan)); + }) + .def("create_precise_sqrt", + [](TritonOpBuilder &self, Value &input) -> Value { + return Value(self.create(input)); + }) + .def("create_precise_divf", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + // AddPtr (similar to GEP) + .def("create_addptr", + [](TritonOpBuilder &self, Value &ptr, Value &offset) -> Value { + return self.create(ptr.getType(), ptr, offset); + }) + // Comparison (int) + .def("create_icmpSLE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::sle, lhs, + rhs); + }) + .def("create_icmpSLT", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::slt, lhs, + rhs); + }) + .def("create_icmpSGE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::sge, lhs, + rhs); + }) + .def("create_icmpSGT", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::sgt, lhs, + rhs); + }) + .def("create_icmpULE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::ule, lhs, + rhs); + }) + .def("create_icmpULT", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::ult, lhs, + rhs); + }) + .def("create_icmpUGE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::uge, lhs, + rhs); + }) + .def("create_icmpUGT", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::ugt, lhs, + rhs); + }) + .def("create_icmpEQ", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::eq, lhs, + rhs); + }) + .def("create_icmpNE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::ne, lhs, + rhs); + }) + // Comparison (float) + .def("create_fcmpOLT", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::OLT, lhs, + rhs); + }) + .def("create_fcmpOGT", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::OGT, lhs, + rhs); + }) + .def("create_fcmpOLE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::OLE, lhs, + rhs); + }) + .def("create_fcmpOGE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::OGE, lhs, + rhs); + }) + .def("create_fcmpOEQ", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::OEQ, lhs, + rhs); + }) + .def("create_fcmpONE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::ONE, lhs, + rhs); + }) + .def("create_fcmpULT", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::ULT, lhs, + rhs); + }) + .def("create_fcmpUGT", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::UGT, lhs, + rhs); + }) + .def("create_fcmpULE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::ULE, lhs, + rhs); + }) + .def("create_fcmpUGE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::UGE, lhs, + rhs); + }) + .def("create_fcmpUEQ", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::UEQ, lhs, + rhs); + }) + .def("create_fcmpUNE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::UNE, lhs, + rhs); + }) + // // Logical + .def("create_and", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_xor", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_or", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + // Input/Output + .def("create_load", + [](TritonOpBuilder &self, Value &ptrs, CacheModifier cacheModifier, + EvictionPolicy evictionPolicy, bool isVolatile, + std::optional flagtree_hints) -> Value { + // flagtree hints + auto hintsAttr = + flagtree_hints + ? mlir::StringAttr::get(self.getContext(), *flagtree_hints) + : mlir::StringAttr::get(self.getContext(), ""); + return self.create(ptrs, cacheModifier, evictionPolicy, + isVolatile, hintsAttr); + }) + .def("create_store", + [](TritonOpBuilder &self, Value &ptrs, Value &value, + CacheModifier cacheModifier, + EvictionPolicy evictionPolicy) -> void { + self.create(ptrs, value, cacheModifier, evictionPolicy); + }) + .def("create_tensor_pointer_load", + [](TritonOpBuilder &self, Value &ptr, + std::vector &boundaryCheck, + std::optional paddingOption, + CacheModifier cacheModifier, EvictionPolicy evictionPolicy, + bool isVolatile, + std::optional flagtree_hints) -> Value { + // flagtree hints + auto hintsAttr = + flagtree_hints + ? mlir::StringAttr::get(self.getContext(), *flagtree_hints) + : mlir::StringAttr::get(self.getContext(), ""); + return self.create(ptr, boundaryCheck, paddingOption, + cacheModifier, evictionPolicy, + isVolatile, hintsAttr); + }) + .def("create_tensor_pointer_store", + [](TritonOpBuilder &self, Value &ptr, Value &val, + std::vector &boundaryCheck, CacheModifier cacheModifier, + EvictionPolicy evictionPolicy) -> void { + self.create(ptr, val, boundaryCheck, cacheModifier, + evictionPolicy); + }) + .def("create_masked_load", + [](TritonOpBuilder &self, Value &ptrs, Value &mask, + std::optional &other, CacheModifier cacheModifier, + EvictionPolicy evictionPolicy, bool isVolatile, + std::optional flagtree_hints) -> Value { + // flagtree hints + auto hintsAttr = + flagtree_hints + ? mlir::StringAttr::get(self.getContext(), *flagtree_hints) + : mlir::StringAttr::get(self.getContext(), ""); + return self.create(ptrs, mask, other.value_or(Value()), + cacheModifier, evictionPolicy, + isVolatile, hintsAttr); + }) + .def("create_masked_store", + [](TritonOpBuilder &self, Value &ptrs, Value &val, Value &mask, + CacheModifier cacheModifier, + EvictionPolicy evictionPolicy) -> void { + self.create(ptrs, val, mask, cacheModifier, + evictionPolicy); + }) + .def("create_tensor_descriptor_type", + [](TritonOpBuilder &self, Type blockTy, bool isSigned) -> Type { + auto ctx = self.getContext(); + return triton::TensorDescType::get( + ctx, cast(blockTy), isSigned); + }) + .def("create_descriptor_load", + [](TritonOpBuilder &self, Value desc, std::vector &indices, + CacheModifier cacheModifier, + EvictionPolicy evictionPolicy) -> Value { + auto descTy = cast(desc.getType()); + auto resTy = descTy.getSignlessBlockType(); + return self.create( + resTy, desc, indices, cacheModifier, evictionPolicy); + }) + .def("create_descriptor_gather", + [](TritonOpBuilder &self, Value desc, Value x_indices, Value y_index, + Type type) -> Value { + return self.create(type, desc, x_indices, + y_index); + }) + .def("create_descriptor_store", + [](TritonOpBuilder &self, Value desc, Value value, + std::vector &indices) -> void { + self.create(desc, value, indices); + }) + .def("create_descriptor_reduce", + [](TritonOpBuilder &self, DescriptorReduceKind kind, Value desc, + Value value, std::vector &indices) -> void { + self.create(kind, desc, value, indices); + }) + .def("create_descriptor_scatter", + [](TritonOpBuilder &self, Value desc, Value value, Value x_indices, + Value y_index) -> void { + self.create(desc, x_indices, y_index, value); + }) + .def("create_reshape", + [](TritonOpBuilder &self, Value &arg, std::vector &shape, + bool allowReorder) -> Value { + return self.create(shape, arg, allowReorder); + }) + .def("create_expand_dims", + [](TritonOpBuilder &self, Value &arg, int axis) -> Value { + return self.create(arg, axis); + }) + .def("create_cat", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + auto lhsType = dyn_cast(lhs.getType()); + auto rhsType = dyn_cast(rhs.getType()); + if (!(lhsType.getShape().size() == 1 && + rhsType.getShape().size() == 1)) + throw std::invalid_argument( + "shape not supported by cat. Expecting rank-1 inputs"); + std::vector shape{lhsType.getShape()[0] + + rhsType.getShape()[0]}; + return self.create(lhsType.clone(shape), lhs, rhs); + }) + .def("create_join", + [](TritonOpBuilder &self, Value &a, Value &b) -> Value { + return self.create(a, b); + }) + .def("create_split", + [](TritonOpBuilder &self, Value &a) -> std::vector { + auto op = self.create(a); + return std::vector(op->result_begin(), op->result_end()); + }) + // Implements tl.trans and tl.permute. + .def("create_trans", + [](TritonOpBuilder &self, Value &arg, std::vector &order) + -> Value { return self.create(arg, order); }) + .def("create_broadcast", + [](TritonOpBuilder &self, Value &arg, + std::vector &shape) -> Value { + if (auto argType = dyn_cast(arg.getType())) + return self.createOrFold(argType.clone(shape), arg); + throw std::invalid_argument( + "arg is not of RankedTensorType, use create_splat"); + }) + .def("create_splat", + [](TritonOpBuilder &self, Type &retTy, Value &arg) -> Value { + return self.createOrFold(retTy, arg); + }) + .def("create_unsplat", + [](TritonOpBuilder &self, Value &arg) -> Value { + return self.createOrFold(arg); + }) + // // atomic + .def("create_atomic_cas", + [](TritonOpBuilder &self, Value &ptr, Value &cmp, Value &val, + MemSemantic sem, MemSyncScope scope) -> Value { + Type dstType; + if (auto srcTensorType = + dyn_cast(ptr.getType())) { + Type dstElemType = + cast(srcTensorType.getElementType()) + .getPointeeType(); + dstType = srcTensorType.clone(dstElemType); + } else { + auto ptrType = cast(getElementTypeOrSelf(ptr)); + dstType = ptrType.getPointeeType(); + } + return self.create(dstType, ptr, cmp, val, sem, + scope); + }) + .def("create_atomic_rmw", + [](TritonOpBuilder &self, RMWOp rmwOp, Value &ptr, Value &val, + Value &mask, MemSemantic sem, MemSyncScope scope) -> Value { + Type dstType; + if (auto srcTensorType = + dyn_cast(ptr.getType())) { + Type dstElemType = + cast(srcTensorType.getElementType()) + .getPointeeType(); + dstType = srcTensorType.clone(dstElemType); + } else { + auto ptrType = cast(getElementTypeOrSelf(ptr)); + dstType = ptrType.getPointeeType(); + } + return self.create(dstType, rmwOp, ptr, val, mask, + sem, scope); + }) + // External + .def("create_extern_elementwise", + [](TritonOpBuilder &self, const std::string &libName, + const std::string &libPath, const std::string &symbol, + std::vector &argList, Type retType, bool isPure) -> Value { + return self.create(retType, argList, libName, + libPath, symbol, isPure); + }) + // Built-in instruction + .def("create_get_program_id", + [](TritonOpBuilder &self, int axis) -> Value { + if (axis < 0 || axis > 3) + throw pybind11::index_error("program_id must be in [0,3]"); + return self.create(axis); + }) + .def("create_get_num_programs", + [](TritonOpBuilder &self, int axis) -> Value { + if (axis < 0 || axis > 3) + throw pybind11::index_error("program_id must be in [0,3]"); + return self.create(axis); + }) + .def("create_dot", + [](TritonOpBuilder &self, mlir::Value &a, mlir::Value &b, + mlir::Value &c, InputPrecision inputPrecision, + int maxNumImpreciseAcc) -> mlir::Value { + return self.create(c.getType(), a, b, c, inputPrecision, + maxNumImpreciseAcc); + }) + .def("create_dot_scaled", + [](TritonOpBuilder &self, mlir::Value &lhs, + std::optional &lhs_scale, + ScaleDotElemType lhs_format, mlir::Value &rhs, + std::optional &rhs_scale, + ScaleDotElemType rhs_format, bool fast_math, bool lhs_k_pack, + bool rhs_k_pack, mlir::Value &c) -> mlir::Value { + return self.create( + c.getType(), lhs, rhs, c, lhs_scale.value_or(Value()), + rhs_scale.value_or(Value()), lhs_format, rhs_format, fast_math, + lhs_k_pack, rhs_k_pack); + }) + .def("create_floor", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_ceil", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_exp", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_exp2", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_cos", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_sin", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_log", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_log2", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_erf", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_sqrt", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_rsqrt", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_fabs", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_iabs", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_reduce", + [](TritonOpBuilder &self, std::vector operands, int axis) + -> OpState { return self.create(operands, axis); }) + .def("create_reduce_ret", + [](TritonOpBuilder &self, py::args args) -> OpState { + llvm::SmallVector return_values; + for (const auto &arg : args) { + return_values.push_back(py::cast(arg)); + } + return self.create(return_values); + }) + .def("create_scan", + [](TritonOpBuilder &self, std::vector operands, int axis, + bool reverse) -> OpState { + return self.create(operands, axis, reverse); + }) + .def("create_scan_ret", + [](TritonOpBuilder &self, py::args args) -> OpState { + llvm::SmallVector return_values; + for (const auto &arg : args) { + return_values.push_back(py::cast(arg)); + } + return self.create(return_values); + }) + .def("create_map_elementwise", + [](TritonOpBuilder &self, std::vector inputs, + std::vector returnTys, int pack) -> OpState { + return self.create(returnTys, inputs, pack); + }) + .def("create_map_elementwise_ret", + [](TritonOpBuilder &self, std::vector returnVals) -> OpState { + return self.create(returnVals); + }) + .def("create_ptr_to_int", + [](TritonOpBuilder &self, Value &val, Type &type) -> Value { + return self.create(type, val); + }) + .def("create_int_to_ptr", + [](TritonOpBuilder &self, Value &val, Type &type) -> Value { + return self.create(type, val); + }) + .def("create_select", + [](TritonOpBuilder &self, Value &condition, Value &trueValue, + Value &falseValue) -> Value { + return self.create(condition, trueValue, + falseValue); + }) + .def("create_inline_asm", + [](TritonOpBuilder &self, const std::string &inlineAsm, + const std::string &constraints, const std::vector &values, + const std::vector &types, bool isPure, + int pack) -> OpState { + return self.create( + types, inlineAsm, constraints, isPure, pack, values); + }) + .def("create_print", + [](TritonOpBuilder &self, const std::string &prefix, bool hex, + const std::vector &values, + const std::vector &isSigned) -> void { + auto prefixAttr = StringAttr::get(self.getBuilder().getContext(), + llvm::StringRef(prefix)); + self.create(prefixAttr, hex, values, isSigned); + }) + .def("create_assert", + [](TritonOpBuilder &self, Value &condition, + const std::string &message) -> void { + auto messageAttr = StringAttr::get(self.getBuilder().getContext(), + llvm::StringRef(message)); + self.create(condition, messageAttr); + }) + .def("create_assume", + [](TritonOpBuilder &self, Value &condition) { + self.create(condition); + }) + .def("create_poison", + [](TritonOpBuilder &self, Type &type) -> Value { + return self.create(type); + }) + .def("create_histogram", + [](TritonOpBuilder &self, Value operand, int numBins, + std::optional mask) -> Value { + if (!mask) { + return self.create( + RankedTensorType::get( + {static_cast(numBins)}, + IntegerType::get(operand.getContext(), 32)), + operand); + } else { + return self.create( + RankedTensorType::get( + {static_cast(numBins)}, + IntegerType::get(operand.getContext(), 32)), + operand, *mask); + } + }) + .def("create_gather", + [](TritonOpBuilder &self, Value src, Value indices, int axis) + -> Value { return self.create(src, indices, axis); }) + // Force GPU barrier + .def("create_barrier", + [](TritonOpBuilder &self) { self.create(); }) + // Make a block pointer (tensor pointer in Triton IR) + .def("create_make_block_ptr", + [](TritonOpBuilder &self, Value &base, std::vector &shape, + std::vector &strides, std::vector &offsets, + std::vector &tensorShape, + std::vector &order) -> Value { + return self.create(base, shape, strides, offsets, + tensorShape, order); + }) + // Advance a block pointer + .def("create_advance", + [](TritonOpBuilder &self, Value &ptr, + std::vector &offsets) -> Value { + return self.create(ptr.getType(), ptr, offsets); + }) + // Make a tensor descriptor + .def("create_make_tensor_descriptor", + [](TritonOpBuilder &self, Value &base, std::vector &shape, + std::vector &strides, std::vector &tensorShape, + bool isSignedInteger, PaddingOption paddingOption) -> Value { + return self.create(base, shape, strides, + tensorShape, isSignedInteger, + paddingOption); + }); + + py::class_(m, "pass_manager", py::module_local()) + .def(py::init()) + .def("enable_debug", + [](PassManager &self) -> bool { + auto *context = self.getContext(); + bool haveDump = ::triton::tools::getBoolEnv("MLIR_ENABLE_DUMP"); + std::string funcToDump; + if (!haveDump) { + funcToDump = triton::tools::getStrEnv("MLIR_ENABLE_DUMP"); + bool isEnvValueBool = + triton::tools::isEnvValueBool(funcToDump).has_value(); + if (!funcToDump.empty() && !isEnvValueBool) + haveDump = true; + } + if (haveDump) { + context->disableMultithreading(); + auto printingFlags = getOpPrintingFlags(); + auto printAlways = [funcToDump](Pass *, Operation *op) -> bool { + if (funcToDump.empty()) + return true; + if (auto mod = dyn_cast(op)) { + return mod.lookupSymbol(funcToDump); + } + if (auto func = dyn_cast(op)) { + return SymbolTable::getSymbolName(func).getValue() == + funcToDump; + } + + return false; + }; + self.enableIRPrinting( + /*shouldPrintBeforePass=*/printAlways, + /*shouldPrintAfterPass=*/printAlways, + /*printModuleScope=*/true, + /*printAfterOnlyOnChange=*/false, + /*printAfterOnlyOnFailure*/ true, mlir_dumps_or_dbgs(), + printingFlags); + } + return haveDump; + }) + .def("get_pipeline_str", + [](PassManager &self) { + std::string str; + llvm::raw_string_ostream os(str); + self.printAsTextualPipeline(os); + return str; + }) + .def( + "run", + [](PassManager &self, ModuleOp &mod) { + // TODO: maybe dump module to file and print error for better + // diagnostics + + auto *context = mod.getContext(); + if (::triton::tools::getBoolEnv("MLIR_DISABLE_MULTITHREADING")) + context->disableMultithreading(); + + auto reproducerPath = + triton::tools::getStrEnv("TRITON_REPRODUCER_PATH"); + if (!reproducerPath.empty()) { + auto anchorName = self.getOpAnchorName(); + auto passes = self.getPasses(); + Operation *op = mod.getOperation(); + // Save a reproducer for the current pass manager invocation + // immediately. + makeReproducer(anchorName, passes, op, reproducerPath); + // But if the pass manager crashes, attempt to generate a local + // reproducer instead. + context->disableMultithreading(); + self.enableCrashReproducerGeneration(reproducerPath, + /*genLocalReproducer=*/true); + } else { + self.enableCrashReproducerGeneration(makeConsoleReproducer()); + } + + if (triton::tools::getBoolEnv("TRITON_ENABLE_LLVM_DEBUG")) { + ::llvm::DebugFlag = true; + } + + if (auto debugOnly = + triton::tools::getStrEnv("TRITON_LLVM_DEBUG_ONLY"); + !debugOnly.empty()) { + llvm::SmallVector storage; + llvm::SmallVector debugTypes = + parseCommaSeparatedValues(debugOnly, storage); + ::llvm::DebugFlag = true; + using namespace llvm; + setCurrentDebugTypes(debugTypes.data(), debugTypes.size()); + } + + bool haveTiming = ::triton::tools::getBoolEnv("MLIR_ENABLE_TIMING"); + if (haveTiming) { + self.enableTiming(); + } + + TritonSourceMgrDiagnosticHandler diagHandler = + setupTritonDiagnosticHandler(context); + if (failed(self.run(mod.getOperation()))) + throw std::runtime_error("PassManager::run failed"); + }, + py::call_guard()); +} + +bool str_eq_ignore_case(const char *s1, const char *s2, int n) { + for (int i = 0; i < n; ++i) { + if (tolower(s1[i]) != s2[i]) + return false; + } + return true; +} + +int strlen_max(const char *str, int max) { + for (int i = 0; i <= max; ++i) { + if (str[i] == '\0') { + return i; + } + } + return 0; +} + +bool is_truthy(char *str) { + int len = strlen_max(str, 4); + switch (len) { + case 1: + return str[0] == '1' || tolower(str[0]) == 'y'; + case 2: + return str_eq_ignore_case(str, "on", len); + case 3: + return str_eq_ignore_case(str, "yes", len); + case 4: + return str_eq_ignore_case(str, "true", len); + default: + return false; + } +} + +PyObject *py_getenv(PyObject *self, PyObject *const *args, Py_ssize_t nargs) { + if (!(nargs == 1 || nargs == 2)) { + PyErr_SetString(PyExc_TypeError, "getenv expected 1 or 2 arguments"); + return NULL; + } + PyObject *name = args[0]; + PyObject *default_val = nargs == 2 ? args[1] : Py_None; + if (!PyUnicode_CheckExact(name)) { + PyErr_SetString(PyExc_TypeError, "name must be a string"); + return NULL; + } + char *env_val = getenv(PyUnicode_AsUTF8(name)); + if (!env_val) { + Py_INCREF(default_val); + return default_val; + } + return PyUnicode_FromString(env_val); +} + +PyObject *py_getenv_bool(PyObject *self, PyObject *const *args, + Py_ssize_t nargs) { + if (nargs != 2) { + PyErr_SetString(PyExc_TypeError, "getenv_bool expected 2 arguments"); + return NULL; + } + PyObject *name = args[0]; + PyObject *default_val = args[1]; + if (!PyUnicode_CheckExact(name)) { + PyErr_SetString(PyExc_TypeError, "name must be a string"); + return NULL; + } + char *env_val = getenv(PyUnicode_AsUTF8(name)); + PyObject *res = default_val; + if (env_val) { + res = is_truthy(env_val) ? Py_True : Py_False; + } + Py_INCREF(res); + return res; +} + +static PyMethodDef ModuleMethods[] = { + {"getenv", (PyCFunction)py_getenv, METH_FASTCALL, NULL}, + {"getenv_bool", (PyCFunction)py_getenv_bool, METH_FASTCALL, NULL}, + {NULL, NULL, 0, NULL} // sentinel +}; + +void init_triton_env_vars(py::module &m) { + m.def("get_cache_invalidating_env_vars", + []() -> std::map { + std::map ret; + for (const auto &envVar : CACHE_INVALIDATING_ENV_VARS) { + auto strVal = triton::tools::getStrEnv(envVar); + if (strVal.empty()) + continue; + auto boolV = triton::tools::isEnvValueBool(strVal); + if (boolV.has_value()) + ret[envVar] = boolV.value() ? "true" : "false"; + else + ret[envVar] = strVal; + } + return ret; + }); + PyModule_AddFunctions(m.ptr(), ModuleMethods); +} diff --git a/third_party/ascend/python/src/ir.h b/third_party/ascend/python/src/ir.h new file mode 100644 index 0000000000..baba1a196d --- /dev/null +++ b/third_party/ascend/python/src/ir.h @@ -0,0 +1,118 @@ +// FIXME: +#pragma once +#include "mlir/IR/Builders.h" +#include "triton/Tools/Sys/GetEnv.hpp" +#include "llvm/ADT/ArrayRef.h" +#include +// Guarded like the core python/src/ir.h: TLE's non-Python libraries compile +// with -U__TLE__ and LLVM's -fno-exceptions/-fno-rtti, which pybind11 headers +// cannot tolerate. +#ifdef __TLE__ +#include +#include +#include +#endif +typedef int AsyncTaskId; +void setAsyncTaskIds(mlir::Operation *op, + llvm::ArrayRef asyncTaskIds); +#ifdef __TLE__ +namespace py = pybind11; +#endif + +using namespace mlir; +using namespace triton; + +// A custom op builder that keeps track of the last location +class TritonOpBuilder { +public: + TritonOpBuilder(mlir::MLIRContext *context, + const std::string &compile_mode = "simd") { + builder = std::make_unique(context); + lastLoc = std::make_unique(builder->getUnknownLoc()); + this->compile_mode = compile_mode; + } + + OpBuilder &getBuilder() { return *builder; } + mlir::MLIRContext *getContext() { return builder->getContext(); } + + bool isLineInfoEnabled() { return lineInfoEnabled; } + + bool isSimtMode() const { return compile_mode == "simt"; } + + void setLastLoc(Location loc) { + if (lineInfoEnabled) + lastLoc = std::make_unique(loc); + } + + void setLastLoc(const std::string &fileName, int line, int column) { + auto context = builder->getContext(); + setLastLoc(FileLineColLoc::get(context, fileName, line, column)); + } + + Location getLastLoc() { + assert(lastLoc); + return *lastLoc; + } + + void setInsertionPointToStart(Block &block) { + if (!block.empty()) + setLastLoc(block.begin()->getLoc()); + else + setLastLoc(builder->getUnknownLoc()); + builder->setInsertionPointToStart(&block); + } + + void setInsertionPointToEnd(Block &block) { + if (!block.empty()) + setLastLoc(block.back().getLoc()); + else + setLastLoc(builder->getUnknownLoc()); + builder->setInsertionPointToEnd(&block); + } + + void setInsertionPointAfter(Operation &op) { + setLastLoc(op.getLoc()); + builder->setInsertionPointAfter(&op); + } + + void restoreInsertionPoint(OpBuilder::InsertPoint pt) { + if (pt.isSet() && pt.getPoint() != pt.getBlock()->end()) + setLastLoc(pt.getPoint()->getLoc()); + else + setLastLoc(builder->getUnknownLoc()); + builder->restoreInsertionPoint(pt); + } + + template OpTy create(Args &&...args) { + auto loc = getLastLoc(); + return builder->create(loc, std::forward(args)...); + } + + // Overload to create or fold a single result operation. + template + std::enable_if_t(), Value> + createOrFold(Args &&...args) { + auto loc = getLastLoc(); + return builder->createOrFold(loc, std::forward(args)...); + } + + // Overload to create or fold a zero result operation. + template + std::enable_if_t(), OpTy> + createOrFold(Args &&...args) { + auto loc = getLastLoc(); + return builder->createOrFold(loc, std::forward(args)...); + } + +private: + std::unique_ptr builder; + std::unique_ptr lastLoc; + bool lineInfoEnabled = !triton::tools::getBoolEnv("TRITON_DISABLE_LINE_INFO"); + std::string compile_mode; +}; + +#ifdef __TLE__ +namespace ir { +extern py::class_ *getBuilderClass(); +} // namespace ir +#endif diff --git a/third_party/ascend/python/src/main.cc b/third_party/ascend/python/src/main.cc new file mode 100644 index 0000000000..eb416a316f --- /dev/null +++ b/third_party/ascend/python/src/main.cc @@ -0,0 +1,67 @@ +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/Signals.h" +#include + +namespace py = pybind11; + +#define FOR_EACH_1(MACRO, X) MACRO(X) +#define FOR_EACH_2(MACRO, X, ...) MACRO(X) FOR_EACH_1(MACRO, __VA_ARGS__) +#define FOR_EACH_3(MACRO, X, ...) MACRO(X) FOR_EACH_2(MACRO, __VA_ARGS__) +#define FOR_EACH_4(MACRO, X, ...) MACRO(X) FOR_EACH_3(MACRO, __VA_ARGS__) +#define FOR_EACH_5(MACRO, X, ...) MACRO(X) FOR_EACH_4(MACRO, __VA_ARGS__) +#define FOR_EACH_6(MACRO, X, ...) MACRO(X) FOR_EACH_5(MACRO, __VA_ARGS__) +#define FOR_EACH_7(MACRO, X, ...) MACRO(X) FOR_EACH_6(MACRO, __VA_ARGS__) +#define FOR_EACH_8(MACRO, X, ...) MACRO(X) FOR_EACH_7(MACRO, __VA_ARGS__) + +#define FOR_EACH_NARG(...) FOR_EACH_NARG_(__VA_ARGS__, FOR_EACH_RSEQ_N()) +#define FOR_EACH_NARG_(...) FOR_EACH_ARG_N(__VA_ARGS__) +#define FOR_EACH_ARG_N(_1, _2, _3, _4, _5, _6, _7, _8, N, ...) N +#define FOR_EACH_RSEQ_N() 8, 7, 6, 5, 4, 3, 2, 1, 0 + +#define CONCATENATE(x, y) CONCATENATE1(x, y) +#define CONCATENATE1(x, y) x##y + +#define FOR_EACH(MACRO, ...) \ + CONCATENATE(FOR_EACH_, FOR_EACH_NARG_HELPER(__VA_ARGS__))(MACRO, __VA_ARGS__) +#define FOR_EACH_NARG_HELPER(...) FOR_EACH_NARG(__VA_ARGS__) + +// New macro to remove parentheses +#define REMOVE_PARENS(...) __VA_ARGS__ + +// Intermediate macro to ensure correct expansion +#define FOR_EACH_P_INTERMEDIATE(MACRO, ...) FOR_EACH(MACRO, __VA_ARGS__) + +// Modified FOR_EACH to handle parentheses +#define FOR_EACH_P(MACRO, ARGS_WITH_PARENS) \ + FOR_EACH_P_INTERMEDIATE(MACRO, REMOVE_PARENS ARGS_WITH_PARENS) + +#define DECLARE_BACKEND(name) void init_triton_##name(pybind11::module &&m); + +#define INIT_BACKEND(name) init_triton_##name(m.def_submodule(#name)); + +void init_triton_env_vars(pybind11::module &m); +void init_triton_ir(pybind11::module &&m); +void init_buffer_ir(pybind11::module &&m); +void init_triton_llvm(pybind11::module &&m); +void init_triton_interpreter(pybind11::module &&m); +void init_triton_passes(pybind11::module &&m); +void init_triton_stacktrace_hook(pybind11::module &m); +void init_gluon_ir(pybind11::module &&m); +void init_linear_layout(pybind11::module &&m); +void init_native_specialize(pybind11::module &m); +FOR_EACH_P(DECLARE_BACKEND, TRITON_BACKENDS_TUPLE) + +PYBIND11_MODULE(libtriton, m) { + m.doc() = "Python bindings to the C++ Triton API"; + init_triton_stacktrace_hook(m); + init_triton_env_vars(m); + init_native_specialize(m); + init_triton_ir(m.def_submodule("ir")); + init_buffer_ir(m.def_submodule("buffer_ir")); + init_triton_passes(m.def_submodule("passes")); + init_triton_interpreter(m.def_submodule("interpreter")); + init_triton_llvm(m.def_submodule("llvm")); + init_linear_layout(m.def_submodule("linear_layout")); + init_gluon_ir(m.def_submodule("gluon_ir")); + FOR_EACH_P(INIT_BACKEND, TRITON_BACKENDS_TUPLE) +} diff --git a/third_party/ascend/python/triton/extension/__init__.py b/third_party/ascend/python/triton/extension/__init__.py new file mode 100644 index 0000000000..3c2c15fac8 --- /dev/null +++ b/third_party/ascend/python/triton/extension/__init__.py @@ -0,0 +1,20 @@ +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. diff --git a/third_party/ascend/python/triton/extension/buffer/src/buffer_ir.cc b/third_party/ascend/python/triton/extension/buffer/src/buffer_ir.cc new file mode 100644 index 0000000000..9ff1e4df97 --- /dev/null +++ b/third_party/ascend/python/triton/extension/buffer/src/buffer_ir.cc @@ -0,0 +1,183 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. + * Copyright 2018-2020 Philippe Tillet + * Copyright 2020-2022 OpenAI + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include +#include + +#include "ir.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Types.h" + +using namespace mlir; +namespace py = pybind11; + +constexpr unsigned kIntegerAttrBitWidth = 64; + +struct BufferOpBuilder : public TritonOpBuilder {}; + +void init_buffer_ir(py::module &&m) { + m.def("load_dialects", [](MLIRContext &context) { + DialectRegistry registry; + registry.insert(); + registry.insert(); + context.appendDialectRegistry(registry); + context.loadAllAvailableDialects(); + }); + + py::class_( + m, "buffer_builder", py::module_local(), py::dynamic_attr()) + .def(py::init()) + .def("get_null_attr", [](BufferOpBuilder &self) { return Attribute(); }) + .def("get_str_array_attr", + [](BufferOpBuilder &self, + const std::vector &array) -> ArrayAttr { + auto strRefVec = to_vector(llvm::map_range( + array, [](const auto &s) { return llvm::StringRef(s); })); + return self.getBuilder().getStrArrayAttr( + llvm::ArrayRef{strRefVec}); + }) + .def("alloc", + [](BufferOpBuilder &self, Type memrefType) -> Value { + return self.create( + mlir::cast(memrefType)); + }) + .def("to_buffer", + [](BufferOpBuilder &self, Value &src, + const Attribute &addressSpace) -> Value { + auto tensorType = dyn_cast(src.getType()); + if (!tensorType) { + llvm::report_fatal_error("to_buffer: src must be tensor type"); + } + auto memrefType = MemRefType::get(tensorType.getShape(), + tensorType.getElementType(), + MemRefLayoutAttrInterface{}); + // TODO: We need to add a pass before OneShotBufferize to generate + // MemorySpaceCastOp + Operation *memref = + self.create(memrefType, src); + if (addressSpace) { + memref = self.create( + MemRefType::get(memrefType.getShape(), + memrefType.getElementType(), + memrefType.getLayout(), addressSpace), + memref->getResult(0)); + } + return memref->getResult(0); + }) + .def("to_tensor", + [](BufferOpBuilder &self, Value &src, bool writable) -> Value { + const auto &memrefType = mlir::cast(src.getType()); + auto tensorType = mlir::RankedTensorType::get( + memrefType.getShape(), memrefType.getElementType()); + auto hasAddressSpace = memrefType.getMemorySpace(); + if (hasAddressSpace) { + MemRefType targetType = MemRefType::get( + memrefType.getShape(), memrefType.getElementType(), + memrefType.getLayout()); + return self.create( + tensorType, + self.create(targetType, src), + true ? mlir::UnitAttr::get(self.getContext()) : nullptr, + writable ? mlir::UnitAttr::get(self.getContext()) : nullptr); + } + return self.create( + tensorType, src, + true ? mlir::UnitAttr::get(self.getContext()) : nullptr, + writable ? mlir::UnitAttr::get(self.getContext()) : nullptr); + }) + .def("subview", + [](BufferOpBuilder &self, Value source, std::vector &offsets, + const std::vector &sizes, + const std::vector &strides) -> Value { + SmallVector mixedOffsets; + auto *context = self.getBuilder().getContext(); + auto &builder = self.getBuilder(); + + // Get source memref type for validation + auto sourceType = mlir::cast(source.getType()); + int64_t rank = sourceType.getRank(); + // Verify the number of parameters + if (offsets.size() != rank || sizes.size() != rank || + strides.size() != rank) { + throw std::runtime_error("Number of offsets, sizes, and strides " + "must match memref rank"); + } + + for (const auto &offset : offsets) { + auto indexType = builder.getIndexType(); + if (offset.getType() != indexType) { + Value offset_val = + self.create(indexType, offset); + mixedOffsets.push_back(offset_val); + } else { + mixedOffsets.push_back(offset); + } + } + + SmallVector mixedSizes; + SmallVector mixedStrides; + for (int64_t i = 0; i < rank; ++i) { + int64_t size = sizes[i]; + int64_t stride = strides[i]; + int64_t srcDim = sourceType.getDimSize(i); + + // verify sizes cannot be negative or zero + if (size <= 0) { + throw std::runtime_error("Expected sizes to be positive"); + } + + // verify strides cannot be negative or zero + if (stride <= 0) { + throw std::runtime_error("Expected strides to be positive"); + } + + // getDimSize() returns -1 (ShapedType::kDynamic) for dynamic + // dimensions + if (!ShapedType::isDynamic(srcDim)) { + // verify the subview size does not exceed the source dimension + if (size > srcDim) { + throw std::runtime_error( + "Subview size cannot exceed source dimension size"); + } + + // verify strides cannot exceed the source dimension size + if (stride > srcDim) { + throw std::runtime_error( + "Stride cannot exceed source dimension size"); + } + } + + mixedSizes.push_back(IntegerAttr::get( + IntegerType::get(context, kIntegerAttrBitWidth), size)); + mixedStrides.push_back(IntegerAttr::get( + IntegerType::get(context, kIntegerAttrBitWidth), stride)); + } + + return self.create(source, mixedOffsets, + mixedSizes, mixedStrides); + }); +} diff --git a/third_party/ascend/runtime/__init__.py b/third_party/ascend/runtime/__init__.py new file mode 100644 index 0000000000..ababac7a81 --- /dev/null +++ b/third_party/ascend/runtime/__init__.py @@ -0,0 +1,35 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +from .autotuner import get_max_configs, max_autotune + + +def _patch_autotune(): + try: + import triton + except ImportError: + return + + from .autotuner import autotune + + triton.autotune = autotune + + +_patch_autotune() diff --git a/third_party/ascend/runtime/autoparser.py b/third_party/ascend/runtime/autoparser.py new file mode 100644 index 0000000000..c1925b138e --- /dev/null +++ b/third_party/ascend/runtime/autoparser.py @@ -0,0 +1,939 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import ast +from typing import Dict, List, Union + + +class AutoParser(ast.NodeVisitor): + """ + Base class for parsing triton dsl kernel code using AST analysis. + + Provides common functionality for traversing the AST (abstract syntax tree) + of a triton dsl kernel function and identifying specific elements in the code. + + Subclassed should implement specific parsing logic by overriding the relevant + node visit methods. + """ + + def __init__(self, func_ast: ast.AST): + self.func_ast = func_ast + + def parse(self): + self.visit(self.func_ast) + + def contains_target_var(self, node, var): + """ + Recursively checks if a given AST node or its children contain a reference + to the specified variable. + + :param node: the AST node to check + :type node: ast.AST + :param var: the variable name to search for + :type var: str + :return: True if the variable is found, False otherwise + """ + if isinstance(node, ast.Name) and node.id == var: + return True + for _, value in ast.iter_fields(node): + if isinstance(value, list): + for item in value: + if self.contains_target_var(item, var): + return True + elif isinstance(value, ast.AST): + if self.contains_target_var(value, var): + return True + return False + + +class AxesKeyParser(AutoParser): + """ + A parser for extracting axis information from a given function's AST. + This class is designed to handle specific patterns in the function's code to + determine the axis associated with a given variable. It is particularly useful + for parsing triton DSL kernel code and identifying axis information. + It recursively processes assignment nodes and lessthan nodes to obtain the axes + corresponding to the specified var in the given function. + """ + + def __init__(self, func_ast: ast.AST, keys: Dict[str, str]): + super().__init__(func_ast) + self.keys = keys + self.checked_vars = list() + + def get_axis(self, var: str, node=None): + """ + Traverse the AST using the provided variable name and mask-based less-than + operations to obtain the corresponding axis name. + + :param var: the variable name to get the corresponding axis. + :type var: str + """ + if var in self.checked_vars: + return None + axis = None + if not node: + node = self.func_ast + for child_node in ast.walk(node): + # handle compare node + if isinstance(child_node, ast.Compare): + axis = self.handle_lt_node(var, child_node) + elif isinstance(child_node, ast.Assign): + axis = self.handle_assign_node(var, child_node) + + elif isinstance(child_node, ast.BinOp) and \ + isinstance(child_node.op, ast.BitAnd): + + axis = self.handle_lt_node(var, child_node.left) + if axis is None: + axis = self.handle_lt_node(var, child_node.right) + + if axis is not None: + return axis + self.checked_vars.append(var) + return None + + def handle_assign_node(self, var, node): + if not isinstance(node, ast.Assign) or not isinstance(node.targets, list): + return None + if len(node.targets) != 1 or not isinstance(node.targets[0], ast.Name): + return None + + target = node.targets[0].id + if target in self.checked_vars: + return None + # Prevent cyclic assignment. + if var == target or not self.contains_target_var(node.value, var): + return None + + axis = self.get_axis(var, node.value) + if axis: + return axis + + axis = self.get_axis(target) + return axis + + def handle_lt_node(self, var, node): + if not isinstance(node, ast.Compare) or not isinstance(node.ops, list): + return None + if len(node.ops) != 1 or not isinstance(node.ops[0], ast.Lt): + return None + if not isinstance(node.comparators, list) or len(node.comparators) != 1: + return None + if not isinstance(node.left, ast.Name) or var != node.left.id: + return None + + comparator = node.comparators[0] + if not isinstance(comparator, ast.Name) and \ + not (isinstance(comparator, ast.Call) and \ + isinstance(comparator.func, ast.Name) and \ + comparator.func.id == 'min'): + return None + + for k, v in self.keys.items(): + if self.contains_target_var(comparator, v): + return k + return None + + +class SplitAxesParser(AxesKeyParser): + """ + Extracts the split axis parameters from triton kernel code. The parsing is based on the + `tl.program_id` statement. This class identifies potential split axes by analyzing the usage + of the `tl.program_id` variable in the program and its multiplication operations with other + variables(currently supporting scenarios where multiplication is either direct or indirect via + intermediate variables). It then filters these candidates based on a list of candidate parameters + (parameters not provided by the user). After that, it confirms the split axis corresponding to + the current parameter using mask comparison and the `keys` passed in `autotune`. + + Note: + 1. Split axis parameters must be multiplied with `tl.program_id`. + 2. Without mask comparision, it is impossible to confirm the exact split axis, which would lead + to parameter parsing failure. (eg. mask = offsets < n_elements) + 3. The identified split axes are limited to the list of candidated parameters, ensuring that + only those parameters that can be dynamically adjusted through the autotune process are considered. + """ + + def __init__(self, func_ast: ast.AST, keys: Dict[str, str], candidates_params: List[str]): + """ + :param func_ast: Abstract syntax tree of the triton kernel function + :type func_ast: ast.AST + :param keys: a dict of axis name: argument name, used to confirm the split axis corresponding to + the split axis parameters. + :type keys: Dict[str, str] + :param candidates_params: a list of parameters names that were not provided by the user when calling + triton kernel function. The parser will only consider these parameters as potential split axis + parameters. + :type candidates_params: List[str] + """ + super().__init__(func_ast, keys) + self.split_axes = dict() + self.program_id_vars = list() + self.program_id_var_dims = dict() + self.num_programs_var_dims = dict() + self.grid_stride_tiling_only = dict() + # axis_name -> program_id axis dim + self.split_axis_pid_dims = dict() + # axis_name -> program_id axis dim (includes axes inferred without split params) + self.axis_pid_dims = dict() + self.candidates_params = candidates_params + + def parse(self) -> Dict[str, str]: + super().parse() + return self.split_axes + + def visit_Assign(self, node): + pid_dim = self._get_program_id_dim(node.value) + if pid_dim is not None: + if (len(node.targets) == 1 and isinstance(node.targets[0], ast.Name) + and node.targets[0].id not in self.program_id_vars): + self.program_id_vars.append(node.targets[0].id) + self.program_id_var_dims[node.targets[0].id] = pid_dim + num_programs_dim = self._get_num_programs_dim(node.value) + if num_programs_dim is not None: + if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): + self.num_programs_var_dims[node.targets[0].id] = num_programs_dim + self.generic_visit(node) + + def visit_BinOp(self, node): + if isinstance(node.op, ast.Mult): + split_axes_val = None + split_axis_pid_dim = None + if isinstance(node.left, ast.Name) and node.left.id in self.program_id_vars: + if isinstance(node.right, ast.Name): + split_axes_val = node.right.id + split_axis_pid_dim = self.program_id_var_dims.get(node.left.id) + elif isinstance(node.left, ast.Call) and isinstance(node.left.func, ast.Attribute): + if (isinstance(node.left.func.value, ast.Name) and node.left.func.value.id == "tl" + and node.left.func.attr == "program_id"): + if isinstance(node.right, ast.Name): + split_axes_val = node.right.id + split_axis_pid_dim = self._get_program_id_dim(node.left) + + if isinstance(node.right, ast.Name) and node.right.id in self.program_id_vars: + if isinstance(node.left, ast.Name): + split_axes_val = node.left.id + split_axis_pid_dim = self.program_id_var_dims.get(node.right.id) + elif isinstance(node.right, ast.Call) and isinstance(node.right.func, ast.Attribute): + if (isinstance(node.right.func.value, ast.Name) and node.right.func.value.id == "tl" + and node.right.func.attr == "program_id"): + if isinstance(node.left, ast.Name): + split_axes_val = node.left.id + split_axis_pid_dim = self._get_program_id_dim(node.right) + + if split_axes_val in self.candidates_params and \ + split_axes_val not in self.split_axes.values(): + split_axes_key = self.get_axis(split_axes_val) + if split_axes_key and not self._is_tiling_only_split(split_axes_key, split_axes_val): + self.split_axes[split_axes_key] = split_axes_val + if split_axis_pid_dim is not None: + self._record_axis_pid_dim(split_axes_key, split_axis_pid_dim) + self.generic_visit(node) + + def visit_For(self, node): + if not isinstance(node.iter, ast.Call): + self.generic_visit(node) + return + + iter_fn = node.iter.func + is_range = isinstance(iter_fn, ast.Name) and iter_fn.id == "range" + is_tl_range = (isinstance(iter_fn, ast.Attribute) and isinstance(iter_fn.value, ast.Name) + and iter_fn.value.id == "tl" and iter_fn.attr == "range") + if not (is_range or is_tl_range): + self.generic_visit(node) + return + + if len(node.iter.args) == 0: + self.generic_visit(node) + return + + start = node.iter.args[0] if len(node.iter.args) >= 2 else None + stop = node.iter.args[1] if len(node.iter.args) >= 2 else node.iter.args[0] + pid_dim = self._extract_pid_dim_from_expr(start) + axis = self._axis_from_expr(stop) + if axis is not None and pid_dim is not None: + self._record_axis_pid_dim(axis, pid_dim) + if len(node.iter.args) >= 3: + step = node.iter.args[2] + loop_tiling_only_param = self._extract_grid_stride_split_param(start, step, pid_dim) + if loop_tiling_only_param is not None: + self._mark_tiling_only_param(axis, loop_tiling_only_param) + + self.generic_visit(node) + + def _get_program_id_dim(self, node): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and isinstance( + node.func.value, ast.Name) and node.func.value.id == "tl" and node.func.attr == "program_id"): + return None + + axis_dim = 0 + if len(node.args) > 0: + if isinstance(node.args[0], ast.Constant) and isinstance(node.args[0].value, int): + axis_dim = node.args[0].value + else: + return None + + for kw in node.keywords: + if kw.arg == "axis": + if isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, int): + axis_dim = kw.value.value + else: + return None + break + return axis_dim + + def _get_num_programs_dim(self, node): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and isinstance( + node.func.value, ast.Name) and node.func.value.id == "tl" and node.func.attr == "num_programs"): + return None + + axis_dim = 0 + if len(node.args) > 0: + if isinstance(node.args[0], ast.Constant) and isinstance(node.args[0].value, int): + axis_dim = node.args[0].value + else: + return None + + for kw in node.keywords: + if kw.arg == "axis": + if isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, int): + axis_dim = kw.value.value + else: + return None + break + return axis_dim + + def _extract_pid_dim_from_expr(self, node): + if node is None: + return None + for child in ast.walk(node): + if isinstance(child, ast.Name) and child.id in self.program_id_var_dims: + return self.program_id_var_dims[child.id] + pid_dim = self._get_program_id_dim(child) + if pid_dim is not None: + return pid_dim + return None + + def _contains_pid_dim(self, node, pid_dim): + if node is None: + return False + for child in ast.walk(node): + if isinstance(child, ast.Name): + if self.program_id_var_dims.get(child.id, None) == pid_dim: + return True + if self._get_program_id_dim(child) == pid_dim: + return True + return False + + def _contains_num_programs_dim(self, node, pid_dim): + if node is None: + return False + for child in ast.walk(node): + if isinstance(child, ast.Name): + if self.num_programs_var_dims.get(child.id, None) == pid_dim: + return True + if self._get_num_programs_dim(child) == pid_dim: + return True + return False + + def _is_candidate_name(self, node, candidate_name): + return (isinstance(node, ast.Name) and node.id == candidate_name and candidate_name in self.candidates_params) + + def _extract_pid_multiplied_candidate(self, node, pid_dim): + if node is None: + return None + candidates = set() + for child in ast.walk(node): + if not isinstance(child, ast.BinOp) or not isinstance(child.op, ast.Mult): + continue + left = child.left + right = child.right + if isinstance(left, ast.Name) and left.id in self.candidates_params and \ + self._contains_pid_dim(right, pid_dim): + candidates.add(left.id) + if isinstance(right, ast.Name) and right.id in self.candidates_params and \ + self._contains_pid_dim(left, pid_dim): + candidates.add(right.id) + if len(candidates) == 1: + return next(iter(candidates)) + return None + + def _contains_num_programs_multiplied_candidate(self, node, candidate_name, pid_dim): + if node is None: + return False + for child in ast.walk(node): + if not isinstance(child, ast.BinOp) or not isinstance(child.op, ast.Mult): + continue + if self._is_candidate_name(child.left, candidate_name): + if self._contains_num_programs_dim(child.right, pid_dim): + return True + if self._is_candidate_name(child.right, candidate_name): + if self._contains_num_programs_dim(child.left, pid_dim): + return True + return False + + def _extract_grid_stride_split_param(self, start, step, pid_dim): + if start is None or step is None: + return None + candidate_name = self._extract_pid_multiplied_candidate(start, pid_dim) + if candidate_name is None: + return None + if self._contains_num_programs_multiplied_candidate(step, candidate_name, pid_dim): + return candidate_name + return None + + def _mark_tiling_only_param(self, axis, candidate_name): + self.grid_stride_tiling_only.setdefault(axis, set()).add(candidate_name) + if self.split_axes.get(axis, None) == candidate_name: + del self.split_axes[axis] + self.split_axis_pid_dims.pop(axis, None) + + def _is_tiling_only_split(self, axis, candidate_name): + return candidate_name in self.grid_stride_tiling_only.get(axis, set()) + + def _axis_from_expr(self, node): + if node is None: + return None + for k, v in self.keys.items(): + if self.contains_target_var(node, v): + return k + return None + + def _record_axis_pid_dim(self, axis, pid_dim): + self.axis_pid_dims[axis] = pid_dim + if axis in self.split_axes: + self.split_axis_pid_dims[axis] = pid_dim + + +class TilingAxesParser(AxesKeyParser): + """ + Extracts the tiling axis parameters from triton kernel code. The parsing is based on the + `tl.arange`, `tl.range` and `range()` statement. This class identifies potential tiling axes by analyzing + the usage of the `range` and `tl.range` within `for` loop in the program. Common parameters + between `range()` or `tl.range` and `tl.arange` are extracted. It then filters these candidates based on a + list of candidate parameters (parameters not provided by the user). After that, it confirms the + tiling axis corresponding to the current parameter using mask comparison and the `keys` passed + in `autotune`. + + Note: + 1. Tiling axis parameters must be calculated within the `tl.arange` function and the `for` loop + using `tl.range`. + 2. Without mask comparision, it is impossible to confirm the exact tiling axis, which would lead + to parameter parsing failure. (eg. mask = offsets < n_elements). + 3. The identified tiling axes are limited to the list of candidated parameters, ensuring that + only those parameters that can be dynamically adjusted through the autotune process are considered. + """ + + def __init__(self, func_ast: ast.AST, keys: Dict[str, str], candidates_params: List[str]): + """ + :param func_ast: Abstract syntax tree of the triton kernel function + :type func_ast: ast.AST + :param keys: a dict of axis name: argument name, used to confirm the tiling axis corresponding to + the tiling axis parameters. + :type keys: Dict[str, str] + :param candidates_params: a list of parameters names that were not provided by the user when calling + triton kernel function. The parser will only consider these parameters as potential tiling axis + parameters. + :type candidates_params: List[str] + """ + super().__init__(func_ast, keys) + self.tiling_axes = dict() + self.candidates_params = candidates_params + self.candidates_params_for_loop = list() + + def parse(self) -> Dict[str, str]: + super().parse() + return self.tiling_axes + + def visit_For(self, node): + if isinstance(node.iter, ast.Call) and len(node.iter.args) == 3: + step_expr = node.iter.args[2] + for_loop_param = self._extract_unique_candidate(step_expr) + if (for_loop_param is not None and for_loop_param not in self.candidates_params_for_loop): + self.candidates_params_for_loop.append(for_loop_param) + self.generic_visit(node) + + def visit_Assign(self, node): + if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): + # handle FloorDiv + if isinstance(node.value, ast.BinOp) and isinstance(node.value.op, ast.FloorDiv): + denominator = node.value.right + denominator_param = self._extract_unique_candidate(denominator) + if denominator_param is not None and \ + denominator_param not in self.candidates_params_for_loop: + self.candidates_params_for_loop.append(denominator_param) + self.visit(self.func_ast) + + tiling_axes_val = self.get_tiling_axes_val(node.value) + if tiling_axes_val is not None and \ + tiling_axes_val in self.candidates_params_for_loop: + tiling_axes_key = self.get_axis(tiling_axes_val) + if tiling_axes_key: + self.tiling_axes[tiling_axes_key] = tiling_axes_val + self.generic_visit(node) + + def get_tiling_axes_val(self, node): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): + if node.func.attr == 'arange' and \ + isinstance(node.func.value, ast.Name) and \ + node.func.value.id == 'tl': + if isinstance(node.args, list) and len(node.args) == 2: + for param in self.candidates_params_for_loop: + if self.contains_target_var(node.args[1], param): + return param + + for _, value in ast.iter_fields(node): + if isinstance(value, list): + for item in value: + val = self.get_tiling_axes_val(item) + if val: + return val + elif isinstance(value, ast.AST): + val = self.get_tiling_axes_val(value) + if val: + return val + return None + + def _extract_unique_candidate(self, expr): + """ + Extract a unique tiling candidate from an expression. + Return None when no candidate or ambiguous (more than one candidate) appears. + """ + if expr is None: + return None + candidates = [param for param in self.candidates_params if self.contains_target_var(expr, param)] + if len(candidates) == 1: + return candidates[0] + return None + + +class ReductionAxesParser(AxesKeyParser): + """ + Extracts the reduction axis from triton kernel code. The parsing is based on the + reduction function (eg. tl.max, tl.min, tl.sum, ...). This class identifies the + dimensions of reduction operations by analyzing the reduction function calls in + the program. After that, It confirms the reduction axis corresponding to the current + parameter using mask comparison and the keys passed in autotune. + + Note: + 1. The call to the reduction function must start with 'tl', meaning the function must + be a function from triton.language + 2. It's preferable to specify the reduction axis dimension in the reduction function + using keyword arguments(eg. axis=xxx). Otherwise, specifying it via positional + arguments may lead to errors. + 3. Mask comparison must be performed on the potential reduction axis length parameters, + and the comparison parameters or target parameters of the comparison expression + must be sliced. Otherwise, the correspondence between dimensions and axes cannot + be confirmed, which will lead to failure in parsing the reduction axis. + 4. The identified reduction axes are limited to the candidate list provided in the keys. + """ + + def __init__(self, func_ast: ast.AST, keys: Dict[str, str]): + """ + :param func_ast: Abstract syntax tree of the triton kernel function + :type func_ast: ast.AST + :param keys: a dict of axis name: argument name, used to confirm the reduction axis. + :type keys: Dict[str, str] + """ + super().__init__(func_ast, keys) + self.reduction_axes = list() + self.reduction_func = ('sum', 'xor_sum', 'max', 'min', 'argmax', 'argmin') # tl.xxx + self.ndim = 1 + + def parse(self) -> List[str]: + super().parse() + return self.reduction_axes + + def visit_Assign(self, node): + self._scan_subscripts(node.value) + self.generic_visit(node) + + def _scan_subscripts(self, node): + if isinstance(node, ast.Subscript): + ndim = self._get_subscripts_ndim(node) + if ndim > self.ndim: + self.ndim = ndim + + for child in ast.iter_child_nodes(node): + self._scan_subscripts(child) + + def _get_subscripts_ndim(self, subscript_node): + slice_node = subscript_node.slice + + if isinstance(slice_node, ast.Tuple): + # e.g. [:, None] -> Tuple(elts=[Slice(), Constant(None)]) + return len(slice_node.elts) + elif isinstance(slice_node, (ast.Slice, ast.Constant, ast.Name, ast.UnaryOp, ast.BinOp)): + # e.g. [0], [:], [i], [-1], [i+1] + return 1 + else: + # Fallback: treat as 1D + return 1 + + def visit_Call(self, node): + if not isinstance(node.func, ast.Attribute): + return + func = node.func + if not isinstance(func.value, ast.Name) or func.value.id != 'tl': + self.generic_visit(node) + return + if func.attr not in self.reduction_func: + return + + axis_dim = None + args = node.args + if len(args) == 1: + # Axis passed as keyword argument + for keyword in node.keywords: + if keyword.arg == 'axis': + axis_dim = self.get_axis_dim(keyword.value) + break + + elif len(args) == 2: + # Axis passed as positional argument. Check the second param + axis_dim = self.get_axis_dim(args[1]) + + else: + raise ValueError("Reduction funtions args error") + + if axis_dim is not None: + reduction_axis = self.get_axis(axis_dim) + if reduction_axis and reduction_axis not in self.reduction_axes: + self.reduction_axes.append(reduction_axis) + + def get_axis_dim(self, node): + if isinstance(node, ast.Constant): + axis_dim = node.value + elif isinstance(node, ast.UnaryOp) and \ + isinstance(node.op, ast.USub): + operand = node.operand + if isinstance(operand, ast.Constant): + axis_dim = self.ndim - operand.value + else: + raise ValueError(f"Reduction function axis error, got: {ast.dump(node)}") + + if not isinstance(axis_dim, int): + raise ValueError("Reduction function axis must be an integer, " + f"got {type(node.value).__name__}: {node.value}") + return axis_dim + + def get_axis(self, axis_dim: int): + """ + Override the parent class method to accept an integer axis dimension + instead of a string. + + :param axis_dim: + :type axis_dim: int + """ + if axis_dim in self.checked_vars: + return None + self.checked_vars.append(axis_dim) + for node in ast.walk(self.func_ast): + if not isinstance(node, ast.Assign): + continue + reduction_axis = self.handle_assign_node(axis_dim, node) + if reduction_axis: + return reduction_axis + return None + + def handle_assign_node(self, axis_dim: int, node): + if not isinstance(node.value, ast.Compare): + return None + + # only support less than + if len(node.value.ops) != 1 or not isinstance(node.value.ops[0], ast.Lt): + return None + + target_axis_len = None + for axis_len in self.keys.values(): + if self.contains_target_var(node.value, axis_len): + target_axis_len = axis_len + break + if not target_axis_len: + return None + + # handel compare left var + if isinstance(node.value.left, ast.Name): + if self.check_compare_left(node.value.left.id, axis_dim): + reduction_axis = next((k for k, v in self.keys.items() if target_axis_len == v), None) + if reduction_axis and reduction_axis not in self.reduction_axes: + return reduction_axis + # handel compare target var + if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): + if self.check_compare_target(node.targets[0].id, axis_dim): + reduction_axis = next((k for k, v in self.keys.items() if target_axis_len == v), None) + if reduction_axis and reduction_axis not in self.reduction_axes: + return reduction_axis + return None + + def check_compare_left(self, var, axis_dim): + for node in ast.walk(self.func_ast): + if not isinstance(node, ast.Assign): + continue + if len(node.targets) != 1 or \ + not isinstance(node.targets[0], ast.Name) or \ + node.targets[0].id != var: + continue + if self.is_current_dim_slice(node.value, axis_dim): + return True + return False + + def check_compare_target(self, var, axis_dim, node=None): + if not node: + node = self.func_ast + for _, value in ast.iter_fields(node): + if isinstance(value, list): + for item in value: + if self.check_compare_target(var, axis_dim, item): + return True + elif isinstance(value, ast.AST): + if isinstance(value, ast.Subscript): + if not isinstance(value.value, ast.Name) or value.value.id != var: + continue + if self.is_current_dim_slice(value, axis_dim): + return True + else: + if self.check_compare_target(var, axis_dim, value): + return True + return False + + def is_current_dim_slice(self, node, dim): + for node in ast.walk(node): + if not isinstance(node, ast.Subscript) or not isinstance(node.slice, ast.Tuple): + continue + elts = node.slice.elts + if len(elts) != 0 and isinstance(elts[dim], ast.Slice): + return True + return False + + +class LowDimsAxesParser(AxesKeyParser): + """ + Extracts the low dimensions axis from triton kernel code. The parsing is based on the + `tl.arange` statement. This class identifies low dimensions axis by analyzing the usage + of the `tl.arange` in the program and extracts the variables computed by `tl.arange` and + their associated operations. Then it checks if these variables are involved in slicing + operations to determine dimension expansion and filters out variables that are expanded + in non-lowest dimensions. After that, it compares the extracted variables with the provided + `keys` to map them to specific low-dimensional axis. + + Note: + 1. low dimensions axis must be calculated within the `tl.arange` function and involved in + slicing operations to be identified. + 2. Without mask comparision, it is impossible to confirm the exact low dimensions axis, which + would lead to parameter parsing failure. (eg. mask = offsets < n_elements). + """ + + def __init__(self, func_ast: ast.AST, keys: Dict[str, str]): + """ + :param func_ast: Abstract syntax tree of the triton kernel function + :type func_ast: ast.AST + :param keys: a dict of axis name: argument name, used to confirm the low-dimensional axis. + :type keys: Dict[str, str] + """ + super().__init__(func_ast, keys) + self.low_dims_axis = list() + self.keys = keys + self.checked_slice_vars = list() + + def parse(self) -> List[str]: + super().parse() + return self.low_dims_axis + + def visit_Assign(self, node): + if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): + tl_arange_node = self.get_tl_arange_node(node) + low_dims_axis = None + if isinstance(tl_arange_node, ast.Call): + partin_other_slice = [False] + if self.is_partin_low_dim_slice(node.targets[0].id, partin_other_slice): + low_dims_axis = self.get_axis(node.targets[0].id) + elif not partin_other_slice[0]: + low_dims_axis = self.get_axis(node.targets[0].id) + elif isinstance(tl_arange_node, ast.Subscript) and \ + self.is_low_dim_slice(tl_arange_node, [False]): + low_dims_axis = self.get_axis(node.targets[0].id) + + if low_dims_axis and low_dims_axis not in self.low_dims_axis: + self.low_dims_axis.append(low_dims_axis) + self.generic_visit(node) + + def get_tl_arange_node(self, node): + for _, value in ast.iter_fields(node): + if isinstance(value, list): + for item in value: + if self.is_tl_arange_call(item): + return item + node = self.get_tl_arange_node(item) + if node: + return node + elif isinstance(value, ast.AST): + if self.is_tl_arange_call(value): + return value + node = self.get_tl_arange_node(value) + if node: + return node + return None + + def is_tl_arange_call(self, node): + """ + Checks if the given AST node is a call to `tl.arange` or a subscript of `tl.arange`. + It supports direct calls to `tl.arange` and subscripts of `tl.arange`, such as + `tl.arange()[None, :]` + """ + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): + if node.func.attr == 'arange' and \ + isinstance(node.func.value, ast.Name) and \ + node.func.value.id == 'tl': + return True + elif isinstance(node, ast.Subscript): + return self.is_tl_arange_call(node.value) + return False + + def is_low_dim_slice(self, node: ast.Subscript, partin_other_slice): + if not isinstance(node.slice, ast.Tuple) or not isinstance(node.slice.elts, list): + return False + elts = node.slice.elts + if len(elts) != 0 and not isinstance(elts[-1], ast.Slice): + partin_other_slice[0] = True + return False + return True + + def is_partin_low_dim_slice(self, var, partin_other_slice, node=None): + if not node: + node = self.func_ast + for child_node in ast.walk(node): + if isinstance(child_node, ast.Subscript) and isinstance(child_node.value, ast.Name): + if var == child_node.value.id and self.is_low_dim_slice(child_node, partin_other_slice): + return True + elif isinstance(child_node, ast.Assign): + if len(child_node.targets) == 1 and \ + isinstance(child_node.targets[0], ast.Name) and \ + var != child_node.targets[0].id: # Prevent cyclic assignment. + if not self.contains_target_var(child_node.value, var): + continue + target_var = child_node.targets[0].id + if target_var in self.checked_slice_vars: + continue + + if self.is_partin_low_dim_slice(var, partin_other_slice, child_node.value): + return True + if self.is_partin_low_dim_slice(target_var, partin_other_slice): + return True + + self.checked_slice_vars.append(var) + return False + + +class PtrNumsParser(AutoParser): + """ + Counts the number of pointer parameters from triton kernel code. The parsing of pointer-type + parameters is determined based on whether these parameters participate in memory access + statements such as `tl.load` and `tl.store`. + First, all input parameters in the kernel function are parsed, and then recursively, all variables + involved in the computation of each input parameter are identified. + If an input parameter directly participates in the computation of the first argument of `tl.load` + or `tl.store`, or if an intermediate variable computed from this input parameter indirectly + participates in the computation of the first argument of `tl.load` or `tl.store`, then this + parameter is considered a pointer-type parameter. + + Note: + 1. Variables modified with `tl.constexpr` are not pointer-type variables and will not be + further parsed. + 2. Only memory access statementes where the input parameter is directly involved or indirectly + involved through one level of computation are counted. Intermediate variables computed from + the input parameter through two or more levels of computation are not counted. + """ + + def __init__(self, func_ast: ast.AST, keys: Dict[str, str], miss_params: List[str]): + """ + :param func_ast: Abstract syntax tree of the triton kernel function + :type func_ast: ast.AST + :param keys: a dict of axis name: argument name, used to exclude potential ptr params. + :type keys: Dict[str, str] + :param miss_params: a list of parameters names that were not provided by the user when calling triton + kernel function. + :type miss_params: List[str] + """ + super().__init__(func_ast) + self.checked_vars = list() + self.ptr_nums = 0 + self.ptr_params = list() + self.keys = keys + self.miss_params = miss_params + self.constexpr_params = list() + + def parse(self): + super().parse() + return self.ptr_nums, self.ptr_params + + def visit_FunctionDef(self, node): + if isinstance(node.args, ast.arguments): + for arg in node.args.args: + if not isinstance(arg, ast.arg): + continue + + if isinstance(arg.annotation, ast.Attribute): + # var modified by tl.constexpr are not pointer type var, passed + is_tl = isinstance(arg.annotation.value, ast.Name) and \ + arg.annotation.value.id == 'tl' + if is_tl and arg.annotation.attr == 'constexpr': + if arg.arg not in self.constexpr_params: + self.constexpr_params.append(arg.arg) + continue + + if self.is_in_addr_calc(arg.arg) and arg.arg not in self.keys.values(): + self.ptr_params.append(arg.arg) + self.ptr_nums += 1 + + for miss_param in self.miss_params: + if miss_param not in self.constexpr_params: + print(f"[WARNING] The parameter '{miss_param}' needs to be declared as tl.constexpr!") + self.generic_visit(node) + + def is_in_addr_calc(self, var): + for node in ast.walk(self.func_ast): + if isinstance(node, ast.Call): + if isinstance(node.func, ast.Attribute) and \ + isinstance(node.func.value, ast.Name): + if node.func.value.id == "tl" and \ + (node.func.attr == "load" or node.func.attr == "store"): + if [arg for arg in node.args if self.contains_target_var(arg, var)]: + return True + + elif isinstance(node, ast.Assign): + if len(node.targets) == 1 and \ + isinstance(node.targets[0], ast.Name) and \ + var != node.targets[0].id: # Prevent cyclic assignment. + target_var = node.targets[0].id + if target_var in self.checked_vars: + continue + if isinstance(node.value, ast.BinOp) and \ + isinstance(node.value.op, ast.Add): + if isinstance(node.value.left, ast.Name) and \ + node.value.left.id == var: + if self.is_in_addr_calc(node.targets[0].id): + return True + elif isinstance(node.value.right, ast.Name) and \ + node.value.right.id == var: + if self.is_in_addr_calc(node.targets[0].id): + return True + self.checked_vars.append(var) + return False diff --git a/third_party/ascend/runtime/autotuner.py b/third_party/ascend/runtime/autotuner.py new file mode 100644 index 0000000000..a5a2ad1625 --- /dev/null +++ b/third_party/ascend/runtime/autotuner.py @@ -0,0 +1,1134 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +from __future__ import annotations + +import builtins +import copy +import functools +import ast +import os +import time +from concurrent.futures import ThreadPoolExecutor +import itertools +from typing import Dict, List + +from torch import Tensor + +import triton +from triton.runtime.autotuner import Autotuner, Config +from triton.tools.get_ascend_devices import is_compile_on_910_95 + +from .autoparser import (LowDimsAxesParser, PtrNumsParser, ReductionAxesParser, SplitAxesParser, TilingAxesParser) +from .utils import get_byte_per_numel, is_valid_axis_name, valid_axis_names + + +class AutoTilingTuner(Autotuner): + """ + Automatic generateing candidate tiling configs and evaluating their performance to get the best config. + """ + + def __init__( + self, + fn, + arg_names, + configs, + key, + reset_to_zero, + restore_value, + pre_hook=None, + post_hook=None, + prune_configs_by: Dict = None, + warmup=None, + rep=None, + use_cuda_graph=False, + do_bench=None, + auto_profile_dir=None, + hints=None, + ): + """ + :param key: a list of argument name, where the change of arguments in value will triger re-generating candidates configs and evaluating. + The parameters in the list will be assigned axis names in sequence, with the axis name being in + {'x','y','z','w','v','t','rx','ry','rz','rw','rv','rt}, where the prefix 'r' means a reduction axis. + Only the axis name in this param should add perfix 'r' if it's a reduction axis. + :type key: List[str] + """ + super().__init__( + fn, + arg_names, + configs, + key, + reset_to_zero, + restore_value, + pre_hook, + post_hook, + prune_configs_by, + warmup, + rep, + use_cuda_graph, + do_bench, + ) + self.user_defined_do_bench = do_bench is not None + if not hints: + self.hints = {} + else: + self.hints = hints + split_params = self.hints.get("split_params", None) + tiling_params = self.hints.get("tiling_params", None) + low_dim_axes = self.hints.get("low_dim_axes", None) + reduction_axes = self.hints.get("reduction_axes", None) + self._init_axis_params( + key, + split_params, + tiling_params, + low_dim_axes, + reduction_axes, + ) + + self.auto_gen_config = not configs or self.hints.get("auto_gen_config", False) + self.gen_configs = [] # generated configs from TileGenerator + self.auto_profile_dir = auto_profile_dir + if not configs: + self.user_configs = [] + else: + self.user_configs = configs + self.is_simt_mode = False + self.simt_stack_limit = 8192 + self.user_specified_warps = None + self.user_specified_multibuffer = None + self.default_multibuffer = not is_compile_on_910_95 + self.print_autotuning = os.getenv("TRITON_PRINT_AUTOTUNING", None) == "1" + # Compile kernels in parallel by default for triton.runtime.JITFunction, + # but not for others, e.g., LibEntry, since it's not compatible with AsyncCompileMode + self.compile_parallel = (isinstance(self.fn, triton.runtime.JITFunction) + and os.getenv("TRITON_AUTOTUNE_PARALLEL_COMPILE", "1") == "1") + + def _expand_simt_num_warps_configs(self, base_configs: List[Config]) -> List[Config]: + _default_cand_num_warps = [8, 16, 32, 64] + cand_num_warps = (_default_cand_num_warps if self.user_specified_warps is None else [self.user_specified_warps]) + + simt_configs = [] + for base_cfg in base_configs: + for num_warps in cand_num_warps: + new_cfg = copy.deepcopy(base_cfg) + new_cfg.num_warps = num_warps + simt_configs.append(new_cfg) + + if self.print_autotuning: + print(f"Triton autotuning: Expanded to {len(simt_configs)} SIMT configs (with warps: {cand_num_warps})") + return simt_configs + + def _expand_simd_multibuffer_configs(self, base_configs: List[Config]) -> List[Config]: + if self.user_specified_multibuffer is not None: + if self.print_autotuning: + print("Triton autotuning: Skip SIMD multibuffer expansion because user " + f"specified multibuffer={self.user_specified_multibuffer}") + return base_configs + + opposite_default_multibuffer = not self.default_multibuffer + simd_configs = [] + for base_cfg in base_configs: + simd_configs.append(base_cfg) + new_cfg = copy.deepcopy(base_cfg) + new_cfg.kwargs["multibuffer"] = opposite_default_multibuffer + simd_configs.append(new_cfg) + + if self.print_autotuning: + print("Triton autotuning: Expanded to " + f"{len(simd_configs)} SIMD configs (toggle multibuffer={opposite_default_multibuffer})") + return simd_configs + + def _init_axis_params(self, key, split_params, tiling_params, low_dim_axes, reduction_axes): + if isinstance(key, list): + if (split_params or tiling_params or low_dim_axes or reduction_axes): + raise ValueError( + "If any axis-related parameters (split_params, tiling_params, low_dim_axes, reduction_axes)" + " are provided, 'key' must be a dict, not a list.") + if len(key) > len(valid_axis_names): + raise ValueError("Number of parameters exceeds the number of available axes.") + self.keys = {axis: param for axis, param in zip(valid_axis_names, key)} + elif isinstance(key, dict): + if not set(key.keys()).issubset(set(valid_axis_names)): + raise ValueError("All keys in 'key' must be valid axis names. Got unexpected keys.") + self.keys = key + if any([split_params, tiling_params, low_dim_axes, reduction_axes]) is None: + raise ValueError( + "If 'key' is a dict, all axis-related parameters (split_params, tiling_params, low_dim_axes," + " reduction_axes) must be provided.") + if not isinstance(split_params, dict): + raise ValueError("split_params must be a dict, got: {}".format(type(split_params))) + if not isinstance(tiling_params, dict): + raise ValueError("tiling_params must be a dict, got: {}".format(type(tiling_params))) + if not isinstance(low_dim_axes, list): + raise ValueError("low_dim_axes must be a list, got: {}".format(type(low_dim_axes))) + if not isinstance(reduction_axes, list): + raise ValueError("reduction_axes must be a list, got: {}".format(type(reduction_axes))) + + used_axes = set(split_params.keys()).union( + tiling_params.keys(), + low_dim_axes, + reduction_axes, + ) + if not used_axes.issubset(self.keys.keys()): + raise ValueError( + "The following axes are used but not present in the 'key': {}".format(used_axes - + set(self.keys.keys()))) + + self.split_params = split_params + self.all_split_params = {} + self.fixed_split_params = {} + self.tiling_params = tiling_params + self.low_dim_axes = low_dim_axes + self.reduction_axes = reduction_axes + self.fixed_grid_dims = set() + self.fixed_grid_dim_values = {} + self.split_axis_pid_dims = {} + self.axis_pid_dims = {} + self.dual_reduction = False + self.persistent_reduction = False + self.num_buffers = -1 + + def _autoparse_axis_params(self, all_args): + miss_params = [arg for arg in self.arg_names if arg not in all_args.keys()] + # parse pointer params nums + if self.num_buffers == -1: + self.num_buffers = self._autoparse_ptr_nums(all_args) + + # parse autotiling axes + # reduction axis must be parsed before other axes. it will alter the key + if not self.reduction_axes: + self.reduction_axes = self._autoparse_reduction_axes() + if len(self.reduction_axes) >= 2: + self.dual_reduction = True + + if not self.low_dim_axes: + self.low_dim_axes = self._autoparse_low_dim_axes() + + if len(self.reduction_axes) == 1: + reduction_axis = self.reduction_axes[0] + reduction_param = self.keys.get(reduction_axis, None) + reduction_numel = all_args.get(reduction_param, float("inf")) + persistent_threshold = self._get_persistent_reduction_threshold(reduction_axis) + if reduction_numel <= persistent_threshold: + self.persistent_reduction = True + + if not self.split_params: + all_split_params = self._autoparse_split_params(self._get_constexpr_candidates()) + self.all_split_params = dict(all_split_params) + self.fixed_split_params = {} + self.fixed_grid_dim_values = self._get_fixed_grid_dim_values( + all_args.get("grid", None), + all_args, + ) + self.fixed_grid_dims = set(self.fixed_grid_dim_values.keys()) + + fixed_grid_axes = {axis for axis, pid_dim in self.axis_pid_dims.items() if pid_dim in self.fixed_grid_dims} + + # Only missing constexpr params are tunable, and fixed-grid axes + # should not be tuned on split. + self.split_params = { + axis: param + for axis, param in all_split_params.items() + if param in miss_params and axis not in fixed_grid_axes + } + + # Fixed split is inferred only from fixed grid dims. + for axis, pid_dim in self.axis_pid_dims.items(): + if pid_dim not in self.fixed_grid_dims: + continue + core_num = self.fixed_grid_dim_values.get(pid_dim, 0) + axis_len_name = self.keys.get(axis, None) + axis_len = all_args.get(axis_len_name, None) + if not isinstance(core_num, int) or core_num <= 0: + continue + if not isinstance(axis_len, int) or axis_len <= 0: + continue + + self.fixed_split_params[axis] = (axis_len + core_num - 1) // core_num + elif not self.axis_pid_dims: + # When split axes are provided by hints, parse axis->program_id mapping + # independently for fixed-grid semantics and diagnostics. + self._autoparse_axis_pid_dims() + miss_params = [arg for arg in miss_params if arg not in self.split_params.values()] + if not self.tiling_params: + self.tiling_params = self._autoparse_tiling_params(miss_params) + miss_params = [arg for arg in miss_params if arg not in self.tiling_params.values()] + if miss_params: + raise ValueError(f"Missing required arguments: {miss_params}. " + f"These arguments must be explicitly provided and cannot be automatically tuned. " + f"Please ensure that these arguments are passed when calling the function.") + + def _gen_tile_configs(self, kv_dict: Dict[str, int], dtype: torch.dtype) -> List[Config]: + from .tile_generator import KernelMeta, TileGenerator + + axis_sizes = {} + for k, v in kv_dict.items(): + if not is_valid_axis_name(k): + continue + if not isinstance(v, int): + raise ValueError(f"Not supported dim type: {type(v)}, `int` is the only supported type") + axis_sizes[k] = v + + kernel_meta = KernelMeta( + axis_sizes, + self.split_params, + self.fixed_split_params, + self.tiling_params, + self.low_dim_axes, + dtype, + self.persistent_reduction, + self.dual_reduction, + self.num_buffers, + self.is_simt_mode, + ) + tile_gen = TileGenerator(kernel_meta=kernel_meta) + tile_gen.descend_split_tiling() + + self.gen_configs.clear() + self.gen_configs = tile_gen.configs + + if self.is_simt_mode: + self.gen_configs = self._expand_simt_num_warps_configs(self.gen_configs) + else: + self.gen_configs = self._expand_simd_multibuffer_configs(self.gen_configs) + + if len(self.gen_configs) == 0: + print("[WARNING] The generated candidate tiling configs are empty based on provided parameters!") + + if self.print_autotuning: + print("Generated configs number: {}".format(len(self.gen_configs))) + + def generate_key_and_configs(self, *args, **kwargs): + self.nargs = dict(zip(self.arg_names, args)) + self.is_simt_mode = kwargs.get('force_simt_only', False) + if 'num_warps' in kwargs and kwargs['num_warps'] is not None: + self.user_specified_warps = kwargs['num_warps'] + else: + self.user_specified_warps = None + if 'multibuffer' in kwargs and kwargs['multibuffer'] is not None: + self.user_specified_multibuffer = kwargs['multibuffer'] + else: + self.user_specified_multibuffer = None + + # generate key + all_args = {**self.nargs, **kwargs} + _args = {k: v for (k, v) in all_args.items() if k in self.arg_names} + key = [_args[v] for _, v in self.keys.items() if v in _args] + + # Currently, we use the dtype with maximum byte length + dtype = None + for _, arg in _args.items(): + if hasattr(arg, "dtype"): + key.append(str(arg.dtype)) + dtype = (arg.dtype if get_byte_per_numel(arg.dtype) >= get_byte_per_numel(dtype) else dtype) + if dtype is None: + raise NotImplementedError("Not support for non-Tensor inputs") + + key = tuple(key) + if key not in self.cache: + if self.auto_gen_config: + self._autoparse_axis_params(all_args) + _kv_dict = {k: _args[v] for k, v in self.keys.items() if v in _args} + self._gen_tile_configs(_kv_dict, dtype) + if len(self.gen_configs) == 0 and len(self.user_configs) == 0: + self.configs = [ + Config( + {}, + num_warps=4, + num_stages=2, + num_ctas=1, + num_buffers_warp_spec=0, + num_consumer_groups=0, + reg_dec_producer=0, + reg_inc_consumer=0, + ) + ] + else: + self.configs = self.gen_configs + self.user_configs + return key + + def run(self, *args, **kwargs): + key = self.generate_key_and_configs(*args, **kwargs) + if self.is_simt_mode and kwargs.get('simt_stack_limit', None) is None: + kwargs['simt_stack_limit'] = self.simt_stack_limit + used_cached_result = True + if key not in self.cache: + # prune configs + pruned_configs = self.prune_configs(kwargs) + if len(pruned_configs) > 1: + used_cached_result = False + bench_start = time.time() + timings = self._batch_bench(*args, configs=pruned_configs, **kwargs) + bench_end = time.time() + self.bench_time = bench_end - bench_start + self.cache[key] = builtins.min(timings, key=timings.get) + full_nargs = {**self.nargs, **kwargs, **self.cache[key].all_kwargs()} + self.pre_hook(full_nargs, reset_only=True) + self.configs_timings = timings + config = self.cache[key] + else: + config = pruned_configs[0] + else: + config = self.cache[key] + + self.best_config = config + if self.print_autotuning and not used_cached_result: + print(f"Triton autotuning for function {self.base_fn.__name__} finished after " + f"{self.bench_time:.2f}s; best config selected: {self.best_config};") + + if not used_cached_result and self.auto_profile_dir is not None: + self._profile(*args, config=self.best_config, **kwargs) + if config.pre_hook is not None: + full_nargs = {**self.nargs, **kwargs, **config.all_kwargs()} + config.pre_hook(full_nargs) + final_kwargs = dict(config.all_kwargs(), **kwargs) + ret = self.fn.run( + *args, + **final_kwargs, + ) + self.nargs = None + return ret + + def _batch_bench(self, *args, configs, **kwargs): + from triton.compiler.errors import CompileTimeAssertionFailure, MLIRCompilationError + from triton.runtime.errors import OutOfResources + + kernels_call = {config: self._make_kernel_call(*args, config=config, **kwargs) for config in configs} + run_fns = {} + exc = None + exc_stack = "" + + if self.compile_parallel: + import psutil + + max_workers = min(psutil.cpu_count(logical=False) // 2, len(kernels_call)) + future_kernels = [] + try: + with ( + ThreadPoolExecutor(max_workers=max_workers) as executor, + triton.AsyncCompileMode(executor), + ): + for config, fn in kernels_call.items(): + future_kernels.append((config, fn(warmup=True))) + + for config, fut in future_kernels: + try: + if hasattr(fut, "result"): + fut = fut.result() + run_fns[config] = functools.partial(kernels_call[config], warmup=False) + except (CompileTimeAssertionFailure, MLIRCompilationError) as e: + import traceback + exc_stack = traceback.format_exc() + exc = e + except Exception as e: + # ignore exception from __exit__() of AsyncCompileMode + triton.runtime._async_compile.active_mode.set(None) + else: + for config, fn in kernels_call.items(): + try: + fn(warmup=False) + run_fns[config] = functools.partial(fn, warmup=False) + except (CompileTimeAssertionFailure, MLIRCompilationError, OutOfResources) as e: + import traceback + exc_stack = traceback.format_exc() + exc = e + + if len(run_fns) == 0: + raise RuntimeError(f"No valid triton configs. {type(exc).__name__}: {exc} \nStack trace: {exc_stack}") + + if len(run_fns) == 1: + # we ignore expensive profiling method when only single config is left + return {config: self.do_bench(fn, quantiles=(0.5, 0.2, 0.8)) for config, fn in run_fns.items()} + + use_profiling = os.getenv("TRITON_BENCH_METHOD", "default").lower() == "npu" + # Respect user-provided benchmarkers even when NPU profiling mode is enabled. + use_npu_profiling = use_profiling and not self.user_defined_do_bench + if use_npu_profiling: + from ..testing import do_bench_npu + + time_cost = do_bench_npu(list(run_fns.values()), clear_l2_cache=False) + assert len(time_cost) == len(run_fns) + return {config: cost for config, cost in zip(run_fns.keys(), time_cost)} + else: + return {config: self.do_bench(fn, quantiles=(0.5, 0.2, 0.8)) for config, fn in run_fns.items()} + + def _make_kernel_call(self, *args, config, **meta): + # check for conflicts, i.e. meta-parameters both provided + # as kwargs and by the autotuner + conflicts = meta.keys() & config.kwargs.keys() + if conflicts: + raise ValueError(f"Conflicting meta-parameters: {', '.join(conflicts)}." + " Make sure that you don't re-define auto-tuned symbols.") + # augment meta-parameters with tunable ones + current = dict(meta, **config.all_kwargs()) + full_nargs = {**self.nargs, **current} + + def kernel_call(warmup): + if config.pre_hook: + config.pre_hook(full_nargs) + self.pre_hook(full_nargs) + try: + current.update({"warmup": warmup}) + res = self.fn.run( + *args, + **current, + ) + if warmup: + return res + except Exception as e: + try: + self.post_hook(full_nargs, exception=e) + finally: + # Throw exception raised by `self.fn.run` + raise + + self.post_hook(full_nargs, exception=None) + + return kernel_call + + def warmup(self, *args, **kwargs): + _ = self.generate_key_and_configs(*args, **kwargs) + pruned_configs = self.prune_configs(kwargs) + ret = [] + if self.compile_parallel: + import psutil + + max_workers = min(psutil.cpu_count(logical=False) // 2, len(pruned_configs)) + with ( + ThreadPoolExecutor(max_workers=max_workers) as executor, + triton.AsyncCompileMode(executor), + ): + for config in pruned_configs: + ret.append(self.fn.warmup(*args, **kwargs, **config.all_kwargs())) + else: + for config in pruned_configs: + ret.append(self.fn.warmup(*args, **kwargs, **config.all_kwargs())) + self.nargs = None + return ret + + def _profile(self, *args, config, **meta): + from ..testing import do_bench_npu + + kernel_call = self._make_kernel_call(*args, config=config, **meta) + fn = functools.partial(kernel_call, warmup=False) + do_bench_npu(fn, prof_dir=self.auto_profile_dir, keep_res=True) + + def _autoparse_split_params(self, candidates_params: List[str]) -> Dict[str, str]: + """ + Extracts the split axis parameters from triton kernel code. + """ + func_ast = self.fn.parse() + parser = SplitAxesParser(func_ast, self.keys, candidates_params) + split_axes = parser.parse() + self.split_axis_pid_dims = dict(getattr(parser, "split_axis_pid_dims", {})) + self.axis_pid_dims = dict(getattr(parser, "axis_pid_dims", {})) + if self.print_autotuning: + print(f"Ascend autotuning parse split axes: {split_axes}, " + f"split axis pid dims: {self.split_axis_pid_dims}, " + f"axis pid dims: {self.axis_pid_dims}") + return split_axes + + def _autoparse_axis_pid_dims(self) -> Dict[str, int]: + """ + Extract axis -> program_id dim mapping without relying on split-parameter + classification, so fixed-grid semantics can always consume it. + """ + func_ast = self.fn.parse() + parser = SplitAxesParser( + func_ast, + self.keys, + self._get_constexpr_candidates(), + ) + _ = parser.parse() + self.axis_pid_dims = dict(getattr(parser, "axis_pid_dims", {})) + self.split_axis_pid_dims = dict(getattr(parser, "split_axis_pid_dims", {})) + if self.print_autotuning: + print("Ascend autotuning parse axis pid dims (independent): " + f"{self.axis_pid_dims}") + return self.axis_pid_dims + + def _get_constexpr_candidates(self) -> List[str]: + """ + Returns all constexpr parameter names from the kernel function definition. + """ + func_ast = self.fn.parse() + constexpr_names = [] + for node in ast.walk(func_ast): + if not isinstance(node, ast.FunctionDef): + continue + if not isinstance(node.args, ast.arguments): + continue + for arg in node.args.args: + if not isinstance(arg, ast.arg): + continue + ann = arg.annotation + if (isinstance(ann, ast.Attribute) and isinstance(ann.value, ast.Name) and ann.value.id == "tl" + and ann.attr == "constexpr"): + constexpr_names.append(arg.arg) + break + return constexpr_names + + def _get_fixed_grid_dim_values(self, grid, all_args: Dict[str, object] = None) -> Dict[int, int]: + """ + Returns fixed grid dim -> value. + - Static tuple/list grid: direct extraction + - Callable grid: infer fixed dims by perturbing missing constexpr params + """ + if grid is None: + return {} + if callable(grid): + return self._infer_fixed_dims_from_callable_grid(grid, all_args or {}) + return self._extract_fixed_grid_dims(grid) + + def _extract_fixed_grid_dims(self, grid) -> Dict[int, int]: + if isinstance(grid, int): + grid = (grid, ) + if not isinstance(grid, (tuple, list)): + return {} + fixed_dims = {} + for idx, dim in enumerate(grid): + if isinstance(dim, int) and dim > 0: + fixed_dims[idx] = dim + return fixed_dims + + def _normalize_grid_tuple(self, grid_out): + if isinstance(grid_out, int): + return (grid_out, ) + if isinstance(grid_out, (tuple, list)): + return tuple(grid_out) + return None + + def _infer_fixed_dims_from_callable_grid(self, grid_fn, all_args: Dict[str, object]) -> Dict[int, int]: + constexpr_candidates = self._get_constexpr_candidates() + base_meta = dict(all_args or {}) + + # Fill missing constexpr with stable probe defaults so grid(meta) can execute. + for name in constexpr_candidates: + if name not in base_meta: + base_meta[name] = 128 + + try: + base_grid_raw = grid_fn(dict(base_meta)) + except Exception: + return {} + + base_grid = self._normalize_grid_tuple(base_grid_raw) + if base_grid is None: + return {} + + dynamic_dims = set() + # Missing constexpr are tunable candidates. + tunable_probe_names = [name for name in constexpr_candidates if name not in (all_args or {})] + probe_values = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512] + + for name in tunable_probe_names: + baseline = base_meta.get(name, 128) + for probe in probe_values: + if probe == baseline: + continue + probe_meta = dict(base_meta) + probe_meta[name] = probe + try: + probe_grid_raw = grid_fn(probe_meta) + except Exception: + continue + probe_grid = self._normalize_grid_tuple(probe_grid_raw) + if probe_grid is None: + continue + if len(probe_grid) != len(base_grid): + dynamic_dims.update(range(min(len(probe_grid), len(base_grid)))) + continue + for idx, (base_dim, probe_dim) in enumerate(zip(base_grid, probe_grid)): + if not (isinstance(base_dim, int) and isinstance(probe_dim, int)): + dynamic_dims.add(idx) + continue + if base_dim != probe_dim: + dynamic_dims.add(idx) + + fixed_dims = {} + for idx, dim in enumerate(base_grid): + if idx in dynamic_dims: + continue + if isinstance(dim, int) and dim > 0: + fixed_dims[idx] = dim + return fixed_dims + + def _autoparse_tiling_params(self, candidates_params: List[str]) -> Dict[str, str]: + """ + Extracts the tiling axis parameters from triton kernel code. + """ + func_ast = self.fn.parse() + parser = TilingAxesParser(func_ast, self.keys, candidates_params) + tiling_axes = parser.parse() + if self.print_autotuning: + print(f"Ascend autotuning parse tiling axes: {tiling_axes}") + return tiling_axes + + def _autoparse_reduction_axes(self) -> List[str]: + """ + Extracts the reduction axis parameters from triton kernel code. + """ + func_ast = self.fn.parse() + parser = ReductionAxesParser(func_ast, self.keys) + reduction_axes = parser.parse() + for axis in reduction_axes: + self.keys[f"r{axis}"] = self.keys.pop(axis) + reduction_axes = [f"r{axis}" for axis in reduction_axes] + + if self.print_autotuning: + print(f"Ascend autotuning parse keys: {self.keys} \n" + f"Ascend autotuning parse reduction axes: {reduction_axes}") + return reduction_axes + + def _autoparse_low_dim_axes(self) -> List[str]: + """ + Extracts the low dimension axis from triton kernel code. + """ + func_ast = self.fn.parse() + parser = LowDimsAxesParser(func_ast, self.keys) + low_dim_axes = parser.parse() + if len(low_dim_axes) < 1: + if self.print_autotuning: + print("[WARNING] Failed to parse low-dimensional axes, fallback to empty low_dim_axes.") + return [] + if self.print_autotuning: + print(f"Ascend autotuning parse low dimensional axes: {low_dim_axes}") + return low_dim_axes + + def _autoparse_ptr_nums(self, all_args: dict) -> int: + """ + Counts the number of pointer parameters from triton kernel code. + """ + ptr_nums = 0 + ptr_params = list() + for k, v in all_args.items(): + if isinstance(v, Tensor): + ptr_nums += 1 + ptr_params.append(k) + + if self.print_autotuning: + print(f"Ascend autotuning parse pointer params: {ptr_params}, pointer nums: {ptr_nums}") + return ptr_nums + + def _get_persistent_reduction_threshold(self, reduction_axis: str) -> int: + # Keep this heuristic aligned with inductor-style policy: + # inner reduction axis uses a larger threshold than other axes. + if self.low_dim_axes and reduction_axis == self.low_dim_axes[0]: + return 1024 + return 64 + + +def autotune(configs, key, prune_configs_by=None, reset_to_zero=None, restore_value=None, pre_hook=None, post_hook=None, + warmup=None, rep=None, use_cuda_graph=False, do_bench=None, *, auto_prof_dir=None, hints=None): + """ + Decorator for auto-tuning a :code:`triton.jit`'d function. + + .. highlight:: python + .. code-block:: python + + @triton.autotune(configs=[ + triton.Config(kwargs={'BLOCK_SIZE': 128}, num_warps=4), + triton.Config(kwargs={'BLOCK_SIZE': 1024}, num_warps=8), + ], + key=['x_size'] # the two above configs will be evaluated anytime + # the value of x_size changes + ) + @triton.jit + def kernel(x_ptr, x_size, **META): + BLOCK_SIZE = META['BLOCK_SIZE'] + :note: When all the configurations are evaluated, the kernel will run multiple times. + This means that whatever value the kernel updates will be updated multiple times. + To avoid this undesired behavior, you can use the `reset_to_zero` argument, which + resets the value of the provided tensor to `zero` before running any configuration. + + If the environment variable :code:`TRITON_PRINT_AUTOTUNING` is set to + :code:`"1"`, Triton will print a message to stdout after autotuning each + kernel, including the time spent autotuning and the best configuration. + + :param configs: a list of :code:`triton.Config` objects + :type configs: list[triton.Config] + :param key: a list of argument names whose change in value will trigger the evaluation of all provided configs. + :type key: list[str] + :param prune_configs_by: a dict of functions that are used to prune configs, fields: + 'perf_model': performance model used to predicate running time with different configs, returns running time + 'top_k': number of configs to bench + 'early_config_prune'(optional): a function used to do early prune (eg, num_stages). It takes configs:List[Config] as its input, and returns pruned configs. + :param reset_to_zero: a list of argument names whose value will be reset to zero before evaluating any configs. + :type reset_to_zero: list[str] + :param restore_value: a list of argument names whose value will be restored after evaluating any configs. + :type restore_value: list[str] + :param pre_hook: a function that will be called before the kernel is called. + This overrides the default pre_hook used for 'reset_to_zero' and 'restore_value'. + 'kwargs': a dict of all arguments passed to the kernel. + 'reset_only': a boolean indicating whether the pre_hook is called to reset the values only, without a corresponding post_hook. + :type pre_hook: lambda args, reset_only + :param post_hook: a function that will be called after the kernel is called. + This overrides the default post_hook used for 'restore_value'. + 'kwargs': a dict of all arguments passed to the kernel. + 'exception': the exception raised by the kernel in case of a compilation or runtime error. + :type post_hook: lambda args, exception + :param warmup: warmup time (in ms) to pass to benchmarking (deprecated). + :type warmup: int + :param rep: repetition time (in ms) to pass to benchmarking (deprecated). + :type rep: int + :param do_bench: a benchmark function to measure the time of each run. + :type do_bench: lambda fn, quantiles + :param auto_prof_dir: the specified directory to store the profiling results of the best config. + If this parameter is None or the best config is retrieved from cache, the profiling process will be ignored. + :type auto_prof_dir: str + :param hints: a dict of autotune hint auguments passed to AutoTilingTuner. + """ + + def decorator(fn): + return AutoTilingTuner(fn, fn.arg_names, configs, key, reset_to_zero, restore_value, pre_hook=pre_hook, + post_hook=post_hook, prune_configs_by=prune_configs_by, warmup=warmup, rep=rep, + use_cuda_graph=use_cuda_graph, do_bench=do_bench, auto_profile_dir=auto_prof_dir, + hints=hints) + + return decorator + + +_ALL_PARAMS = { + "num_stages", + "unit_flag", + "limit_auto_multi_buffer_only_for_local_buffer", + "limit_auto_multi_buffer_of_local_buffer", + "set_workspace_multibuffer", + "enable_hivm_auto_cv_balance", + "tile_mix_vector_loop", + "tile_mix_cube_loop", + "enable_ubuf_saving", +} + +_DEFAULTS = { + "num_stages": [2], + "unit_flag": [False], + "limit_auto_multi_buffer_only_for_local_buffer": [False], + "limit_auto_multi_buffer_of_local_buffer": ["no-l0c"], + "set_workspace_multibuffer": [2, 4], + "enable_hivm_auto_cv_balance": [True], + "tile_mix_vector_loop": [2, 4], + "tile_mix_cube_loop": [2, 4], + "enable_ubuf_saving": [True], +} + +_VALID_VALUES = { + "num_stages": [1, 2], + "limit_auto_multi_buffer_of_local_buffer": ["no-limit", "no-l0c"], + "set_workspace_multibuffer": [2, 4], + "tile_mix_vector_loop": [2, 4, 8], + "tile_mix_cube_loop": [2, 4, 8], +} + +_CUBE_PARAMS = {"num_stages", "unit_flag", "limit_auto_multi_buffer_of_local_buffer"} + +_MIXCV_PARAMS = { + "num_stages", + "unit_flag", + "limit_auto_multi_buffer_only_for_local_buffer", + "limit_auto_multi_buffer_of_local_buffer", + "set_workspace_multibuffer", + "enable_hivm_auto_cv_balance", + "tile_mix_vector_loop", + "tile_mix_cube_loop", + "enable_ubuf_saving", +} + +_VECTOR_PARAMS = { + "num_stages", + "enable_ubuf_saving", +} + + +def _check_boolean_list(val, param_name): + return isinstance(val, (list, tuple)) and len(val) > 0 and all(isinstance(x, bool) for x in val) + + +def _check_string_in_set(val, valid_set, param_name): + return isinstance(val, (list, tuple)) and len(val) > 0 and all(v in valid_set for v in val) + + +def _check_int_in_set(val, valid_set, param_name): + return isinstance(val, (list, tuple)) and len(val) > 0 and all(isinstance(v, int) and v in valid_set for v in val) + + +_VALIDATION_RULES = { + "num_stages": { + "desc": f"must be one or more of: {_VALID_VALUES['num_stages']}", "check": + lambda val, p: _check_int_in_set(val, _VALID_VALUES['num_stages'], p) + }, + "unit_flag": {"desc": "must be non-empty list/tuple of boolean values", "check": _check_boolean_list}, + "limit_auto_multi_buffer_only_for_local_buffer": + {"desc": "must be non-empty list/tuple of boolean values", "check": _check_boolean_list}, + "limit_auto_multi_buffer_of_local_buffer": { + "desc": f"must be one or more of: {_VALID_VALUES['limit_auto_multi_buffer_of_local_buffer']}", "check": + lambda val, p: _check_string_in_set(val, _VALID_VALUES['limit_auto_multi_buffer_of_local_buffer'], p) + }, + "set_workspace_multibuffer": { + "desc": f"must be one or more of: {_VALID_VALUES['set_workspace_multibuffer']}", "check": + lambda val, p: _check_int_in_set(val, _VALID_VALUES['set_workspace_multibuffer'], p) + }, + "enable_hivm_auto_cv_balance": + {"desc": "must be non-empty list/tuple of boolean values", "check": _check_boolean_list}, + "tile_mix_vector_loop": { + "desc": f"must be one or more of: {_VALID_VALUES['tile_mix_vector_loop']}", "check": + lambda val, p: _check_int_in_set(val, _VALID_VALUES['tile_mix_vector_loop'], p) + }, + "tile_mix_cube_loop": { + "desc": f"must be one or more of: {_VALID_VALUES['tile_mix_cube_loop']}", "check": + lambda val, p: _check_int_in_set(val, _VALID_VALUES['tile_mix_cube_loop'], p) + }, + "enable_ubuf_saving": {"desc": "must be non-empty list/tuple of boolean values", "check": _check_boolean_list}, +} + + +class BaseAutotuner: + """ + Base class for generating auto-tuning configurations without block dimensions. + Users must provide fixed dimension parameters when calling the kernel. + """ + + def __init__(self, operator_name, supported_params, default_params, validation_rules): + self.operator_name = operator_name + self.supported_params = supported_params + self.default_params = default_params + self.validation_rules = validation_rules + + def validate_parameters(self, **kwargs): + # Check for unsupported parameters + invalid_params = [k for k in kwargs.keys() if k not in _ALL_PARAMS] + if invalid_params: + print(f"[ERROR] Invalid parameters for {self.operator_name}: {invalid_params}") + return False + + for param, rule in self.validation_rules.items(): + if param in kwargs: + if not rule["check"](kwargs[param], param): + print(f"[ERROR] Invalid value for '{param}' in {self.operator_name}: {kwargs[param]}") + print(f" Expected: {rule['desc']}") + return False + return True + + def get_configs(self, **kwargs): + """ + Generate a list of Config objects. + Each parameter must be provided as a list (even for a single value). + The function produces the Cartesian product of all parameter lists. + - num_stages: each value will be set as Config.num_stages (not placed in kwargs) + - other parameters: each value will be placed in Config.kwargs + Returns a list of Config objects. + """ + if not self.validate_parameters(**kwargs): + return [] + + # Collect parameter values, using defaults for missing ones + param_values = {} + for p in sorted(self.supported_params): + if p in kwargs: + param_values[p] = kwargs[p] + else: + param_values[p] = self.default_params.get(p, [None]) + + keys = list(param_values.keys()) + values = list(param_values.values()) + combos = list(itertools.product(*values)) + + configs = [] + for combo in combos: + config_kwargs = {} + num_stages_val = None + for i, pname in enumerate(keys): + val = combo[i] + if pname == "num_stages": + num_stages_val = val + else: + config_kwargs[pname] = val + + configs.append(Config(kwargs=config_kwargs, num_stages=num_stages_val if num_stages_val is not None else 2)) + return configs + + +CubeAutotuner = BaseAutotuner(operator_name="cube", supported_params=_CUBE_PARAMS, default_params=_DEFAULTS, + validation_rules=_VALIDATION_RULES) + +MixcvAutotuner = BaseAutotuner(operator_name="mixcv", supported_params=_MIXCV_PARAMS, default_params=_DEFAULTS, + validation_rules=_VALIDATION_RULES) + +VectorAutotuner = BaseAutotuner(operator_name="vector", supported_params=_VECTOR_PARAMS, default_params=_DEFAULTS, + validation_rules=_VALIDATION_RULES) + + +def get_autotune_cube_config(**kwargs: Any) -> List[triton.Config]: + """ + Generate autotune configuration for the cube operator. + Supported parameters: num_stages, unit_flag, limit_auto_multi_buffer_of_local_buffer. + """ + import triton + return CubeAutotuner.get_configs(**kwargs) + + +def get_autotune_cv_config(**kwargs: Any) -> List[triton.Config]: + """ + Generate autotune configuration for the mixcv operator. + Supported parameters: num_stages, unit_flag, limit_auto_multi_buffer_only_for_local_buffer, + limit_auto_multi_buffer_of_local_buffer, set_workspace_multibuffer, + enable_hivm_auto_cv_balance, tile_mix_vector_loop, tile_mix_cube_loop, enable_ubuf_saving + """ + import triton + return MixcvAutotuner.get_configs(**kwargs) + + +def get_autotune_vector_config(**kwargs: Any) -> List[triton.Config]: + """ + Generate autotune configuration for the vector operator. + Supported parameters: num_stages, enable_ubuf_saving + """ + import triton + return VectorAutotuner.get_configs(**kwargs) + + +def get_max_configs(config, kernel_type="mixcv", **kwargs): + """ + Expand a single base Config by combining it with tuning parameters. + + :param config: A triton.Config object serving as the base. + :param kernel_type: Operator type, one of "cube", "mixcv", "vector". Default "mixcv". + :param kwargs: Tuning parameters, each provided as a list (e.g., enable_hivm_auto_cv_balance=[True, False]). + If a parameter is not provided, its value is taken from the base config (if present) + or from the defaults. + :return: List of expanded Config objects. + """ + # Determine the set of parameters supported by the current kernel_type + if kernel_type == "cube": + supported = _CUBE_PARAMS + elif kernel_type == "vector": + supported = _VECTOR_PARAMS + else: + supported = _MIXCV_PARAMS + + # Warn about unsupported parameters provided in kwargs + unsupported = [k for k in kwargs if k not in supported and k in _ALL_PARAMS] + if unsupported: + print( + f"[WARNING] The following parameters are not supported for kernel_type '{kernel_type}': {unsupported}. They will be ignored." + ) + + # Build value lists for each parameter (priority: kwargs > base config > defaults) + param_values = {} + base_kwargs = config.kwargs + base_num_stages = config.num_stages + + for param in sorted(supported): + if param in kwargs: + # User-provided list via tuning_params takes precedence + val_list = kwargs[param] + elif param == "num_stages": + # num_stages is an attribute of Config, not part of kwargs. + # If not provided in tuning_params, use the value from the base config as a fixed single-element list. + val_list = [base_num_stages] + elif param in base_kwargs: + # Parameter present in base config's kwargs -> fix to that single value + val_list = [base_kwargs[param]] + else: + # Otherwise fall back to defaults + val_list = _DEFAULTS.get(param, [None]) + + # Validate the value list + if param in _VALIDATION_RULES: + rule = _VALIDATION_RULES[param] + if not rule["check"](val_list, param): + raise ValueError(f"Invalid value for '{param}': {val_list}. Expected: {rule['desc']}") + param_values[param] = val_list + + # Cartesian product of all parameter lists + keys = list(param_values.keys()) + values = list(param_values.values()) + combos = list(itertools.product(*values)) + + new_configs = [] + for combo in combos: + # Start with a copy of the original config's kwargs + new_kwargs = config.kwargs.copy() + num_stages_val = None + + for i, pname in enumerate(keys): + val = combo[i] + if pname == "num_stages": + num_stages_val = val + else: + # Overwrite or add the parameter to kwargs + new_kwargs[pname] = val + + new_config = Config(kwargs=new_kwargs, num_warps=config.num_warps, + num_stages=num_stages_val if num_stages_val is not None else config.num_stages, + num_ctas=config.num_ctas, num_buffers_warp_spec=config.num_buffers_warp_spec, + num_consumer_groups=config.num_consumer_groups, reg_dec_producer=config.reg_dec_producer, + reg_inc_consumer=config.reg_inc_consumer, maxnreg=config.maxnreg, pre_hook=config.pre_hook) + new_configs.append(new_config) + + return new_configs + + +def max_autotune(configs, key, kernel_type="mixcv", prune_configs_by=None, reset_to_zero=None, restore_value=None, + pre_hook=None, post_hook=None, warmup=None, rep=None, use_cuda_graph=False, do_bench=None, + **tuning_params): + """ + Decorator that expands each base Config with tuning parameters before auto-tuning. + + Usage is similar to @triton.autotune, but allows automatic expansion of + additional tuning parameters (e.g., enable_hivm_auto_cv_balance, tile_mix_vector_loop, ...) + for each provided base configuration. + + :param configs: List of base triton.Config objects. + :param key: List of argument names whose change triggers re-tuning. + :param kernel_type: Operator type, one of "cube", "mixcv", "vector". Default "mixcv". + :param prune_configs_by: Same as in autotune. + :param reset_to_zero: Same as in autotune. + :param restore_value: Same as in autotune. + :param pre_hook: Same as in autotune. + :param post_hook: Same as in autotune. + :param warmup: Deprecated. + :param rep: Deprecated. + :param use_cuda_graph: Deprecated. + :param do_bench: Same as in autotune. + :param tuning_params: Additional tuning parameters as keyword arguments. + Each value must be a list; the Cartesian product of these lists + will be combined with each base config. + """ + + def decorator(fn): + if not configs or len(configs) == 0: + raise ValueError("[max_autotune] The argument 'configs' cannot be empty. " + "Please provide at least one base config. ") + # Expand each base config with the provided tuning parameters + expanded_configs = [] + for cfg in configs: + expanded = get_max_configs(cfg, kernel_type=kernel_type, **tuning_params) + expanded_configs.extend(expanded) + + # Call the original autotune decorator with the expanded configs + return autotune(configs=expanded_configs, key=key, prune_configs_by=prune_configs_by, + reset_to_zero=reset_to_zero, restore_value=restore_value, pre_hook=pre_hook, + post_hook=post_hook, warmup=warmup, rep=rep, use_cuda_graph=use_cuda_graph, + do_bench=do_bench)(fn) + + return decorator diff --git a/third_party/ascend/runtime/tile_generator.py b/third_party/ascend/runtime/tile_generator.py new file mode 100644 index 0000000000..a4e2a9a83c --- /dev/null +++ b/third_party/ascend/runtime/tile_generator.py @@ -0,0 +1,489 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +from __future__ import annotations + +import functools +import sys +from dataclasses import dataclass +from typing import ( + Dict, + List, + Tuple, +) + +from triton.runtime.autotuner import Config + +from .utils import ( + get_byte_per_numel, + next_power_of_2, + num_vector_core, + ub_size_in_kbytes, + rf_size_in_kbytes, +) + + +@dataclass +class AxisInfo: + name: str + index: int + length: int + + prefix: str = "" + split_name: str = "" + tiling_name: str = "" + is_split_axis: bool = False + is_tunable_split_axis: bool = False + is_tiling_axis: bool = False + fixed_split_size: int = 0 + + @property + def is_reduction(self): + return self.prefix == "r" + + +class KernelMeta: + + def __init__( + self, + axis_sizes: Dict[str, int], + split_params: Dict[str, str], + fixed_split_params: Dict[str, int], + tiling_params: Dict[str, str], + low_dims: List[str], + dtype: torch.dtype, + persistent_reduction: bool, + dual_reduction: bool, + num_buffers: int, + is_simt_mode: bool, + ): + """ + :param split_params: a dict of axis name: argument name, the argument is an adjustable parameter in a split axis, such as 'XBLOCK'. + The axis name must be in key's axis names. Do not add prefix 'r' before the axis name. + This param can be empty. Note that the auto tiling feature will be disabled when the split_params and tiling_params are both empty. + The split axis can usually be identified according to `tl.program_id()` expression. + :type split_params: Dict[str, str] + :param tiling_params: a dict of axis name: argument name, the argument is an adjustable parameter in a tiling axis, such as 'XBLOCK_SUB'. + The axis name must be in key's axis names. Do not add prefix 'r' before the axis name. + This param can be empty. Note that the auto tiling feature will be disabled when the split_params and tiling_params are both empty. + The tiling axis can usually be identified according to `tl.arange()` expression. + :type tiling_params: Dict[str, str] + :param low_dims: a list of axis name in which the corresponding axis is low dim aixs. + The axis name must be in key's axis names. Do not add prefix 'r' before the axis name. + :type low_dims: List[str] + :param dual_reduction: performing reduction on more than one axis. + :param persistent_reduction: there is no splitting in reduction axis. + """ + self._validate_axis(axis_sizes, split_params, fixed_split_params, tiling_params, low_dims) + + axis_dict = {} + idx = 0 + for name, length in axis_sizes.items(): + prefix = "" + if name.startswith("r"): + prefix = "r" + + is_tunable_split_axis = name in split_params + fixed_split_size = fixed_split_params.get(name, 0) + is_split_axis = is_tunable_split_axis or fixed_split_size > 0 + is_tiling_axis = name in tiling_params + split_name = "" if not is_tunable_split_axis else split_params[name] + tiling_name = "" if name not in tiling_params else tiling_params[name] + + axis_dict[name] = AxisInfo( + name=name, + index=idx, + length=length, + prefix=prefix, + split_name=split_name, + tiling_name=tiling_name, + is_split_axis=is_split_axis, + is_tunable_split_axis=is_tunable_split_axis, + is_tiling_axis=is_tiling_axis, + fixed_split_size=fixed_split_size, + ) + idx += 1 + + self.axis_info = list(axis_dict.values()) + self.split_axis = [x for x in axis_dict.values() if x.is_split_axis] + self.tunable_split_axis = [x for x in axis_dict.values() if x.is_tunable_split_axis] + self.tiling_axis = [x for x in axis_dict.values() if x.is_tiling_axis] + self.low_dims_axis = [x for x in axis_dict.values() if x.name in low_dims] + self.dtype = dtype + self.persistent_reduction = persistent_reduction + self.dual_reduction = dual_reduction + self.num_buffers = num_buffers + self.is_simt_mode = is_simt_mode + + @classmethod + def _validate_axis( + cls, + axis_sizes: Dict[str, int], + split_params: Dict[str, str], + fixed_split_params: Dict[str, int], + tiling_params: Dict[str, str], + low_dims: List[str], + ) -> None: + for axis_name in axis_sizes.keys(): + if axis_name.startswith("r") and len(axis_name) == 1: + raise ValueError("The name of a reduction axis is empty!") + + def check_keys(params: List[str], context="parameter"): + for k in params: + if k not in axis_sizes and ("r" + k) not in axis_sizes: + raise KeyError(f"{context} '{k}' not found in known axes: {axis_sizes.keys()}") + + check_keys(split_params.keys(), "split axis") + check_keys(fixed_split_params.keys(), "fixed split axis") + check_keys(tiling_params.keys(), "tiling axis") + check_keys(low_dims, "low dim axis") + + +@dataclass +class BlockInfo: + block_name: str # e.g., XBLOCK + sub_block_name: str # e.g., XBLOCK_SUB + block_size: int + sub_block_size: int + + +""" +Generate possible candidate tiling configs for benchmarking +""" + + +class TileGenerator: + num_warps = 1 + num_stages = 1 + + def __init__(self, kernel_meta: KernelMeta): + self.kernel_meta = kernel_meta + self.persistent_reduction = self.kernel_meta.persistent_reduction + self.dual_reduction = self.kernel_meta.dual_reduction + + self.blocks = self.init_blocks_info(kernel_meta) + self.numels = [axis.length for axis in kernel_meta.axis_info] + self.candidate_blocks = [] + self.configs = [] + self.dtype_bytes = get_byte_per_numel(kernel_meta.dtype) + + self.num_buffers = 3 if kernel_meta.num_buffers == 0 else min(kernel_meta.num_buffers, 3) + self.is_simt_mode = kernel_meta.is_simt_mode + local_mem_size = (rf_size_in_kbytes if self.is_simt_mode else ub_size_in_kbytes) + self.max_numel_threshold = local_mem_size * 1024 // self.dtype_bytes // self.num_buffers + self.max_total_numel = functools.reduce(lambda x, y: x * y, [x.block_size + for x in self.blocks]) if self.blocks else 1 + self.small_kernel = self.max_total_numel < 128 * 1024 + self.tiny_kernel = self.max_total_numel <= 32 * 1024 + self.stop_numel = min(1024 // self.dtype_bytes, self.max_total_numel // + (num_vector_core * 2)) if self.small_kernel else 1024 // self.dtype_bytes + self.max_programs_num = 65535 + self.tiny_program_threshold = num_vector_core // 8 + self.tiny_per_program_cap = 1 + self.tiny_low_program_hist = {p: 0 for p in range(1, self.tiny_program_threshold + 1)} + self.tiny_low_program_active = False + self.tiny_low_program_tile_floor = 0 + + @classmethod + def init_blocks_info(cls, kernel_meta: KernelMeta) -> List[BlockInfo]: + blocks = [] + for axis in kernel_meta.axis_info: + block_name = axis.split_name + sub_block_name = axis.tiling_name + block_size = axis.fixed_split_size if axis.fixed_split_size > 0 else axis.length + sub_block_size = block_size + blocks.append(BlockInfo(block_name, sub_block_name, block_size, sub_block_size)) + + return blocks + + @classmethod + def get_key_from_dict(cls, kwargs: Dict[str, int]): + return tuple(sorted(kwargs.items())) + + def calcu_last_split_blocks(self, axis_idx): + splits = 1 + for x in self.kernel_meta.split_axis: + if x.index != axis_idx: + splits = splits * ( + (self.numels[x.index] + self.blocks[x.index].block_size - 1) // self.blocks[x.index].block_size) + else: + break + + last_splits = num_vector_core // splits + last_splits = max(1, last_splits) + last_blocks = (self.numels[axis_idx] + last_splits - 1) // last_splits + return last_blocks + + def aligned_numel(self, numel, align_bytes=32): + if self.is_simt_mode: + return next_power_of_2(numel) + + align_numel = align_bytes // self.dtype_bytes + if numel <= align_numel: + return numel + return ((numel + align_numel - 1) // align_numel) * align_numel + + def calculate_tile_numel(self): + tile_numel = 1 + for axis in self.kernel_meta.axis_info: + if axis.is_tiling_axis: + tile_numel *= self.blocks[axis.index].sub_block_size + else: + tile_numel *= self.blocks[axis.index].block_size + + return tile_numel + + def fill_config(self, cfg, candi_block): + for axis in self.kernel_meta.axis_info: + if not (axis.is_split_axis or axis.is_tiling_axis): + continue + block_info = self.blocks[axis.index] + if axis.is_split_axis: + curr_numel = candi_block[axis.index] + if not axis.is_tiling_axis: + curr_numel = self.aligned_numel(curr_numel) + if block_info.block_name: + cfg[block_info.block_name] = curr_numel + if axis.is_tiling_axis: + tiling_numel = self.aligned_numel(block_info.sub_block_size) + cfg[block_info.sub_block_name] = (tiling_numel if self.is_simt_mode else min( + tiling_numel, candi_block[axis.index])) + + def find_config(self, cfg): + for config_var in self.configs: + if config_var.kwargs == cfg: + return True + return False + + def _try_add_tiny_low_program_config(self, total_programs): + if (not self.tiny_kernel or total_programs < 1 or total_programs > self.tiny_program_threshold): + return + + if self.tiny_low_program_hist.get(total_programs, 0) >= self.tiny_per_program_cap: + return + + candi_block = tuple([x.block_size for x in self.blocks]) + if self.add_to_configs(list(candi_block)): + if candi_block not in self.candidate_blocks: + self.candidate_blocks.append(candi_block) + if not self.tiny_low_program_active: + self.tiny_low_program_active = True + self.tiny_low_program_tile_floor = self.calculate_tile_numel() + self.tiny_low_program_hist[total_programs] = (self.tiny_low_program_hist.get(total_programs, 0) + 1) + + def _calc_total_programs(self, candi_block=None): + grids = [] + for axis in self.kernel_meta.split_axis: + numel = self.numels[axis.index] + block_size = (self.blocks[axis.index].block_size if candi_block is None else candi_block[axis.index]) + programs = (numel + block_size - 1) // block_size + grids.append(programs) + + total_programs = functools.reduce(lambda x, y: x * y, grids) if grids else 1 + return total_programs + + def add_to_configs(self, candi_block): + newcfg = {} + self.fill_config(newcfg, candi_block) + tile_numel = self.calculate_tile_numel() + stop_numel_threshold = 0 if len(self.configs) < 10 or self.small_kernel else self.stop_numel + 100 + if self.tiny_low_program_active and self.tiny_low_program_tile_floor > 0: + total_programs = self._calc_total_programs(candi_block) + program_threshold = self.tiny_program_threshold if self.small_kernel else num_vector_core // 2 + if total_programs <= program_threshold: + tiny_low_program_threshold = max(self.stop_numel, self.tiny_low_program_tile_floor // 2) + stop_numel_threshold = max(stop_numel_threshold, tiny_low_program_threshold) + if (tile_numel <= self.max_numel_threshold and tile_numel >= stop_numel_threshold + and not self.find_config(newcfg)): + self.configs.append(Config(newcfg, num_warps=1, num_stages=1)) + return True + return False + + def desecnd_all_low_dims_with_all_blocks(self): + restore_sub_blocks = {} + for axis in self.kernel_meta.low_dims_axis: + restore_sub_blocks[axis.index] = self.blocks[axis.index].sub_block_size + self.descend_all_low_dims() + for axis in self.kernel_meta.low_dims_axis: + self.blocks[axis.index].sub_block_size = restore_sub_blocks[axis.index] + + def descend_one_axis(self, axis_idx: int, is_split=False): + + def calc_total_programs(): + grids = [] + for axis in self.kernel_meta.split_axis: + numel = self.numels[axis.index] + block_size = self.blocks[axis.index].block_size + programs = (numel + block_size - 1) // block_size + grids.append(programs) + + total_programs = functools.reduce(lambda x, y: x * y, grids) if grids else 1 + return total_programs + + reached_stop_numel = False + slow_decend_split = False + num_vector_core_tile = num_vector_core + max_programs_num = num_vector_core_tile if self.kernel_meta.tiling_axis else self.max_programs_num + if not is_split and len(self.candidate_blocks) == 0: + self.candidate_blocks.append(tuple([x.block_size for x in self.blocks])) + + axis = self.kernel_meta.axis_info[axis_idx] + while True: + for candi_block in self.candidate_blocks: + if self.add_to_configs(candi_block): + self.desecnd_all_low_dims_with_all_blocks() + + # tile numel reached threshold + tile_numel = self.calculate_tile_numel() + if tile_numel <= self.stop_numel: + if self.add_to_configs([x.block_size for x in self.blocks]): + self.desecnd_all_low_dims_with_all_blocks() + reached_stop_numel = True + break + + numel = (self.blocks[axis_idx].block_size if is_split else self.blocks[axis_idx].sub_block_size) + if numel == 1: + if self.add_to_configs([x.block_size for x in self.blocks]): + self.desecnd_all_low_dims_with_all_blocks() + break + + if is_split: + if self.persistent_reduction and axis.is_reduction: + reached_stop_numel = True + break + total_programs = calc_total_programs() + if total_programs > num_vector_core_tile: + if len(self.configs) == 0: + num_vector_core_tile = max_programs_num + slow_decend_split = (total_programs > num_vector_core_tile // 2) + if total_programs > num_vector_core_tile: + last_blocks = self.calcu_last_split_blocks(axis_idx) + if last_blocks != self.blocks[axis_idx].block_size: + self.blocks[axis_idx].block_size = last_blocks + self.candidate_blocks.append(tuple([x.block_size for x in self.blocks])) + break + + program_threshold = self.tiny_program_threshold if self.small_kernel else num_vector_core // 2 + if self.tiny_kernel and total_programs <= program_threshold: + self._try_add_tiny_low_program_config(total_programs) + if total_programs > program_threshold or self.dual_reduction: + if len(self.candidate_blocks) > 2: + self.candidate_blocks.pop(0) + self.candidate_blocks.append(tuple([x.block_size for x in self.blocks])) + if self.small_kernel: + self.add_to_configs(list(tuple([x.block_size for x in self.blocks]))) + slow_decend_split = (total_programs > num_vector_core_tile // 2) + + if not slow_decend_split: + self.blocks[axis_idx].block_size = (numel + 1) // 2 + else: + step = (numel + 3) // 4 if (numel + 3) // 4 > 1 else 1 + self.blocks[axis_idx].block_size = numel - step + self.blocks[axis_idx].sub_block_size = self.blocks[axis_idx].block_size + total_programs = calc_total_programs() + if self.blocks[axis_idx].block_size == 1 and (total_programs > program_threshold + or self.dual_reduction): + self.candidate_blocks.append(tuple([x.block_size for x in self.blocks])) + else: + if numel >= 32: + self.blocks[axis_idx].sub_block_size = next_power_of_2(numel // 2) + else: + self.blocks[axis_idx].sub_block_size = numel - 1 + return reached_stop_numel + + def descend_all_low_dims(self): + low_dim_numels = [self.blocks[x.index].sub_block_size for x in self.kernel_meta.low_dims_axis] + if not low_dim_numels: + return False + + def descend_all_axis(min_numel): + + for axis in self.kernel_meta.low_dims_axis: + if axis.is_reduction and self.persistent_reduction: + continue + + numel = self.blocks[axis.index].sub_block_size + if numel == 1: + continue + if min_numel > 1 and abs(numel - min_numel) / min_numel < 0.2: + continue + if numel >= 128: + self.blocks[axis.index].sub_block_size = next_power_of_2(numel // 2) + else: + numel = self.blocks[axis.index].sub_block_size + numel = numel // 2 + self.blocks[axis.index].sub_block_size = min(self.aligned_numel(numel), next_power_of_2(numel)) + + if len(self.candidate_blocks) == 0: + # means there is no split axis and tiling_not_low_dim axis + # so we need to init the candidates_blk_sizes + self.candidate_blocks.append(tuple([x.block_size for x in self.blocks])) + + count = 0 + tile_numel = self.calculate_tile_numel() + while tile_numel > self.stop_numel and count < 100: + count += 1 + tile_numel = self.calculate_tile_numel() + for candi_block in self.candidate_blocks: + self.add_to_configs(candi_block) + min_numel = min(low_dim_numels) + descend_all_axis(min_numel) + new_tile_numel = self.calculate_tile_numel() + if tile_numel == new_tile_numel: + descend_all_axis(0) + + return tile_numel < self.stop_numel + + def descend_split_tiling(self): + + tiling_not_low_dims = [x for x in self.kernel_meta.tiling_axis if x not in self.kernel_meta.low_dims_axis] + + def descend_split_axis(): + for axis in self.kernel_meta.tunable_split_axis: + if self.descend_one_axis(axis.index, is_split=True): + return True + + return self.calculate_tile_numel() <= self.stop_numel + + def descend_tiling_not_low_dims(): + for axis in tiling_not_low_dims: + if axis.is_reduction and self.persistent_reduction: + continue + + if self.descend_one_axis(axis.index): + return True + return self.calculate_tile_numel() <= self.stop_numel + + while True: + # descend split axis + if descend_split_axis(): + break + if len(self.candidate_blocks) > 0: + candi_block = self.candidate_blocks[0] + for i, blk_size in enumerate(candi_block): + self.blocks[i].sub_block_size = blk_size + # descend tiling but not low dims + if descend_tiling_not_low_dims(): + break + # descend low dims, need to descend all axis at the same time + self.descend_all_low_dims() + break diff --git a/third_party/ascend/runtime/utils.py b/third_party/ascend/runtime/utils.py new file mode 100644 index 0000000000..fadeeec98d --- /dev/null +++ b/third_party/ascend/runtime/utils.py @@ -0,0 +1,132 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch + +_cached_params = None + + +def _init_npu_params(): + global _cached_params + if _cached_params is not None: + return _cached_params + + from triton.runtime.driver import driver + + target = driver.active.get_current_target() + device = driver.active.get_current_device() + prop = driver.active.utils.get_device_properties(device) + + num_cube_core = prop["num_aicore"] + num_vector_core = prop["num_aicore"] + ub_size_in_kbytes = 192 + rf_size_in_kbytes = None + + ASCEND_VARIANTS = ["Ascend910B", "Ascend910_93", "Ascend910_95", "Ascend950"] + if any(variant in target.arch for variant in ASCEND_VARIANTS): + num_vector_core = num_cube_core * 2 + + if target.arch.startswith("Ascend910_95") or target.arch.startswith("Ascend950"): + ub_size_in_kbytes = 256 + rf_size_in_kbytes = 128 + + _cached_params = { + 'target': target, + 'device': device, + 'prop': prop, + 'num_cube_core': num_cube_core, + 'num_vector_core': num_vector_core, + 'ub_size_in_kbytes': ub_size_in_kbytes, + 'rf_size_in_kbytes': rf_size_in_kbytes, + } + return _cached_params + + +def __getattr__(name): + if name in [ + 'target', 'device', 'prop', 'num_cube_core', 'num_vector_core', 'ub_size_in_kbytes', 'rf_size_in_kbytes' + ]: + return _init_npu_params()[name] + raise AttributeError(f"module '{__name__}' has no attribute '{name}'") + + +# wrapper npu 32 bytes align, get and pass unalign info to triton meta +# then autotune choose tiling param and send them to bishengIR +byte_per_numel = { + torch.float32: 4, # torch.float32 or torch.float + torch.float64: 8, # torch.float64 or torch.double + torch.float16: 2, # torch.float16 or torch.half + torch.bfloat16: 2, # torch.bfloat16 + torch.int32: 4, # torch.int32 or torch.int + torch.int64: 8, # torch.int64 or torch.long + torch.int16: 2, # torch.int16 or torch.short + torch.int8: 1, # torch.int8 + torch.uint8: 1, # torch.uint8 + torch.bool: 1, # torch.bool + torch.complex32: 4, # torch.complex32 (not yet available in PyTorch as of the latest stable release) + torch.complex64: 8, # torch.complex64 + torch.complex128: 16, # torch.complex128 +} + +# Some PyTorch versions expose extra fp8 dtypes. Register them when available. +for fp8_dtype_name in ( + "float8_e4m3fn", + "float8_e4m3fnuz", + "float8_e5m2", + "float8_e5m2fnuz", +): + fp8_dtype = getattr(torch, fp8_dtype_name, None) + if fp8_dtype is not None: + byte_per_numel[fp8_dtype] = 1 + +valid_axis_names = [ + "x", + "y", + "z", + "w", + "v", + "t", +] + + +def get_byte_per_numel(dtype: torch.dtype) -> int: + return 1 if dtype is None else byte_per_numel[dtype] + + +def is_valid_axis_name(name: str) -> bool: + if name.startswith("r"): + return name[1:] in valid_axis_names + return name in valid_axis_names + + +# move to an appropriate place, currently duplicated with triton.__init__.py +def next_power_of_2(n: int): + """Return the smallest power of 2 greater than or equal to n""" + n -= 1 + n |= n >> 1 + n |= n >> 2 + n |= n >> 4 + n |= n >> 8 + n |= n >> 16 + n |= n >> 32 + n += 1 + return n diff --git a/third_party/ascend/spec/triton/extension/buffer/__init__.py b/third_party/ascend/spec/triton/extension/buffer/__init__.py new file mode 100644 index 0000000000..3c2c15fac8 --- /dev/null +++ b/third_party/ascend/spec/triton/extension/buffer/__init__.py @@ -0,0 +1,20 @@ +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. diff --git a/third_party/ascend/spec/triton/extension/buffer/language/__init__.py b/third_party/ascend/spec/triton/extension/buffer/language/__init__.py new file mode 100644 index 0000000000..86af01bdb8 --- /dev/null +++ b/third_party/ascend/spec/triton/extension/buffer/language/__init__.py @@ -0,0 +1,44 @@ +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +__all__ = [ + # core + "builtin", + "is_builtin", + + # buffer + "buffer", + + # base address space + "address_space", + + # alloc + "alloc", + + # to_buffer + "to_buffer", + + # to_tensor + "to_tensor", + "subview", +] + +from .core import builtin, is_builtin, address_space, buffer, alloc, to_buffer, to_tensor, subview diff --git a/third_party/ascend/spec/triton/extension/buffer/language/builder.py b/third_party/ascend/spec/triton/extension/buffer/language/builder.py new file mode 100644 index 0000000000..2257080bbd --- /dev/null +++ b/third_party/ascend/spec/triton/extension/buffer/language/builder.py @@ -0,0 +1,73 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +buffer-specific builder utilities for code generation. +""" + +__all__ = [ + "create_builder_method_wrapper_with_buffer_builder", + "attach_builder_methods_with_buffer_builder", + "setup_unified_builder_with_buffer_builder", +] + + +def create_builder_method_wrapper_with_buffer_builder(main_builder, delegate_builder, method_name): + """ + Create a wrapper that delegates a method call to another builder while + synchronizing insertion points and locations. + """ + delegate_method = getattr(delegate_builder, method_name) + + def wrapper(*args, **kwargs): + saved_ip = main_builder.get_insertion_point() + saved_loc = main_builder.get_loc() + delegate_builder.restore_insertion_point(saved_ip) + if saved_loc: + delegate_builder.set_loc(saved_loc) + result = delegate_method(*args, **kwargs) + main_builder.restore_insertion_point(saved_ip) + if saved_loc: + main_builder.set_loc(saved_loc) + return result + + wrapper.__name__ = method_name + wrapper.__doc__ = getattr(delegate_method, '__doc__', None) + return wrapper + + +def attach_builder_methods_with_buffer_builder(main_builder, delegate_builder, method_names): + """Attach multiple methods from a delegate builder to the main builder.""" + for method_name in method_names: + wrapper = create_builder_method_wrapper_with_buffer_builder(main_builder, delegate_builder, method_name) + setattr(main_builder, method_name, wrapper) + + +def setup_unified_builder_with_buffer_builder(main_builder, buffer_builder): + """Set up a unified builder interface by attaching methods from specialized builders.""" + main_builder._buffer_builder = buffer_builder + buffer_methods = [ + 'get_null_attr', + 'get_str_array_attr', + 'alloc', + 'to_buffer', + 'to_tensor', + 'subview', + ] + attach_builder_methods_with_buffer_builder(main_builder, buffer_builder, buffer_methods) diff --git a/third_party/ascend/spec/triton/extension/buffer/language/core.py b/third_party/ascend/spec/triton/extension/buffer/language/core.py new file mode 100644 index 0000000000..d170561813 --- /dev/null +++ b/third_party/ascend/spec/triton/extension/buffer/language/core.py @@ -0,0 +1,341 @@ +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +__all__ = [ + "address_space", + "buffer_type", + "subview", + "alloc", + "buffer", + "to_buffer", + "to_tensor", +] + +import importlib +from typing import TypeVar, List +from functools import wraps + +from triton._C.libtriton import ir +import triton.language.core as tl + +T = TypeVar("T") + +TRITON_BUILTIN = "__triton_builtin__" +BUFFER_BUILTIN = "__buffer_builtin__" + + +def builtin(fn: T) -> T: + """Mark a function as a buffer language builtin.""" + assert callable(fn) + + @wraps(fn) + def wrapper(*args, **kwargs): + if "_semantic" not in kwargs or kwargs["_semantic"] is None: + raise ValueError("Did you forget to add @triton.jit ? " + "(`_semantic` argument must be provided outside of JIT functions.)") + return fn(*args, **kwargs) + + # also set triton_builtin to true so that CodeGenerator will recognize this function + setattr(wrapper, TRITON_BUILTIN, True) + setattr(wrapper, BUFFER_BUILTIN, True) + + return wrapper + + +def is_builtin(fn) -> bool: + """Is this a registered buffer language builtin function?""" + return getattr(fn, BUFFER_BUILTIN, False) + + +class address_space: + """Represents a buffer's address space. + + The :code:`address_space` of a buffer is a target-specific attribute. + """ + + def to_ir(self, builder: ir.builder) -> ir.type: + raise NotImplementedError("Abstract address_space cannot be converted to ir") + + +class buffer_type(tl.dtype): + + def __init__(self, element_ty: tl.dtype, shape: List, space: address_space = None, strides: List = None): + self.element_ty = element_ty + self.shape = shape if isinstance(shape, list) else list(shape) + self.space = space + self.strides = strides if strides is not None else [] + self.name = self._make_name() + + def _make_name(self): + res = '' + + def to_ir(self, builder: ir.builder) -> ir.type: + element_ty_ir = self.element_ty.to_ir(builder) + addr_space_attr = self.space.to_ir(builder) if self.space else builder.get_null_attr() + + # use the method with strides if strides is not empty + if self.strides: + return builder.get_buffer_ty_with_strides(self.shape, element_ty_ir, self.strides, addr_space_attr) + else: + return builder.get_buffer_ty(self.shape, element_ty_ir, addr_space_attr) + + def __str__(self): + return self.name + + def __repr__(self): + return self.__str__() + + def __eq__(self, other) -> bool: + if not isinstance(other, buffer_type): + return False + return (self.element_ty == other.element_ty and self.shape == other.shape and self.space == other.space + and self.strides == other.strides) + + def __ne__(self, other) -> bool: + return not self.__eq__(other) + + @property + def scalar(self): + return self.element_ty + + def mangle(self) -> str: + elt = self.element_ty.mangle() + shape = "_".join(map(str, self.shape)) + return f"B{elt}S{shape}S" + + def _unflatten_ir(self, handles: List[ir.value], cursor: int): + return buffer(handles[cursor], self), cursor + 1 + + +# ----------------------- +# buffer +# ----------------------- + + +class buffer(tl.base_value): + """Represents a region of memory. + + :code:`buffer` is the fundamental data structure for Triton programs using + the buffer language extension. Most functions in + :py:mod:`triton.extension.buffer.language` operate on and return buffers. + + Most of the named member functions here are duplicates of the free functions + in :code:`triton.language`. For example, :code:`triton.language.sqrt(x)` is + equivalent to :code:`x.sqrt()`. + + .. rubric:: Constructors + .. + For some reason Sphinx includes __init__ before printing the full table + of methods. Not what I want, but I can't figure out how to fix it. Give + it its own section so it looks intentional. :) + """ + + def __init__(self, handle, buffer_ty: buffer_type): + """Not called by user code.""" + super().__init__() + self.handle = handle + self.type = buffer_ty + self.dtype = buffer_ty.element_ty.scalar + self.shape = buffer_ty.shape + self.space = buffer_ty.space + self.strides = buffer_ty.strides + + def _flatten_ir(self, handles: List[ir.value]) -> None: + handles.append(self.handle) + + def __str__(self) -> str: + # ex. "<16x32xfloat32, address_space>" + res = '<' + 'x'.join(str(s) for s in self.shape) + 'x' + str(self.dtype) + if self.space: + res += ', ' + str(self.space) + return res + '>' + + @builtin + def subview(self, offsets: List[tl.constexpr], sizes: List[tl.constexpr], strides: List[tl.constexpr], + _semantic=None) -> 'buffer': + return subview(self, offsets, sizes, strides, _semantic=_semantic) + + @builtin + def to_tensor(self, writable=True, target_shape=None, _semantic=None): + """Convert this buffer to a tl.tensor""" + return to_tensor(self, writable=writable, target_shape=target_shape, _semantic=_semantic) + + +semantic = importlib.import_module(".semantic", package=__package__) + + +@builtin +def alloc(etype: tl.dtype, shape: List[tl.constexpr], _address_space: address_space = None, is_mem_unique: bool = False, + _semantic=None) -> buffer: + """ + Allocates a region of local memory with the specified shape and type. + + :param etype: the element type of the buffer. + :type etype: tl.dtype + :param shape: A list of non-negative integers representing the shape of the buffer. + :type shape: List[tl.constexpr] + :param _address_space: (Optional) backend-specific local memory address space + :type _address_space: bl.address_space + """ + return semantic.alloc(etype, shape, _address_space, is_mem_unique, _semantic.builder) + + +@builtin +def to_buffer(tensor: tl.tensor, space: address_space = None, bind_buffer: buffer = None, _semantic=None) -> buffer: + """ + Convert a tensor to a buffer. + + :param tensor: the tensor to convert. + :type tensor: tl.tensor + :param space: the address space for the buffer (optional). + :type space: address_space + """ + return semantic.to_buffer(tensor, space, bind_buffer, _semantic.builder) + + +@builtin +def to_tensor(memref: buffer, writable: bool = True, target_shape=None, _semantic=None) -> tl.tensor: + """ + Create a tl.tensor from a bl.buffer. + + :param memref: the input bl.buffer object. + :memref type: bl.buffer + :param writable: If set true, the resultant tensor is considered "writable" during bufferization. + :type writable: bool + """ + return semantic.to_tensor(memref, writable, _semantic.builder, target_shape=target_shape) + + +def check_subview(src, offsets, sizes, strides): + """ + Check data of subview methods which the data length and the offset value must be 32-byte aligned. + + The conditions for checking data are as follows: + 1. offset value must be 32-bytes aligned. + 2. all strides must be 1. + 3. the first point's offset in the second row of the last dimension must be 32-bytes aligned. + + For instance, the following example fails to satisfy the specified criteria. + %subview = memref.subview %arg0[1, 1][4, 4][2, 2] + : memref<8x8xf32, strided<[8, 1], offset: 0>> to + memref<4x4xf32, strided<[16, 2], offset: 9>> + offsets = [8, 8] | sizes = [4, 4] | strides = [2, 2] + result_offset = 9 + second_row_start_offset = 25 + The scene will be go wrong because the follow conditions are not meet. + 1) result_offset is not 32-bytes aligned. + 2) strides = [2, 2], not all strides are equal to 1. + 3) second_row_start_offset are not 32-bytes aligned. + """ + bytes_per_block = 32 + bits_per_byte = 8 + base_byte = bytes_per_block // (src.dtype.primitive_bitwidth // bits_per_byte) + result_strides = [] + result_offset = 0 + second_row_start_offset = 0 + length = len(strides) + src_strides = [1] * length + if length == 1: + if offset[0] % base_byte != 0: # noqa: F821 + raise TypeError("all strides should be 1 and the offset value should be 32-bytes aligned.") + return + for i in range(length - 2, -1, -1): + src_strides[i] = src_strides[i + 1] * src.shape[i + 1] + for i in range(0, length): + if isinstance(offsets[i], tl.tensor): + return + result_strides.append(src_strides[i] * strides[i]) + result_offset = result_offset + offsets[i] * src_strides[i] + second_row_start_offset = result_offset + src_strides[-2] * strides[-2] + is_unaligned = False + if sizes[1] > 1: + is_unaligned = second_row_start_offset % base_byte != 0 + stride_1 = all(s == 1 for s in strides) + is_unaligned = result_offset % base_byte != 0 or is_unaligned or not stride_1 + if is_unaligned: + raise TypeError("all strides should be 1 and the offset value should be 32-bytes aligned.") + + +@builtin +def subview(src: buffer, offsets: List[tl.constexpr], sizes: List[tl.constexpr], strides: List[tl.constexpr], + _semantic=None) -> buffer: + ''' + Creates a subview of the source buffer with the specified offsets, sizes, and strides. + + :param src: The source buffer to create a subview from. + :type src: buffer + :param offsets: A list of non-negative integers representing the offsets in each dimension. + :type offsets: List[tl.constexpr] + :param sizes: A list of non-negative integers representing the sizes in each dimension. + :type sizes: List[tl.constexpr] + :param strides: A list of non-negative integers representing the strides in each dimension. + :type strides: List[tl.constexpr] + :return: A new buffer representing the subview of the source buffer. + :rtype: buffer + ''' + # Validate that sizes and strides contain only constexpr values + new_sizes = [] + for i, size in enumerate(sizes): + if isinstance(size, int): + # Convert regular integers to constexpr + new_sizes.append(tl.constexpr(size)) + elif isinstance(size, tl.constexpr): + new_sizes.append(size) + else: + raise TypeError(f"sizes[{i}] must be constexpr, got {type(size).__name__}") + + new_strides = [] + for i, stride in enumerate(strides): + if isinstance(stride, int): + # Convert regular integers to constexpr + new_strides.append(tl.constexpr(stride)) + elif isinstance(stride, tl.constexpr): + new_strides.append(stride) + else: + raise TypeError(f"strides[{i}] must be constexpr, got {type(stride).__name__}") + + check_offsets = [] + new_offsets = [] + for offset in offsets: + if isinstance(offset, tl.constexpr): + # Check that constexpr offset values cannot be negative + if offset < 0: + raise ValueError(f"Offset value must be non-negative, got {offset}") + new_offsets.append(_semantic.to_tensor(offset)) + check_offsets.append(offset) + elif isinstance(offset, int): + # Convert regular integers to constexpr and then to tensor + if offset < 0: + raise ValueError(f"Offset value must be non-negative, got {offset}") + new_offsets.append(_semantic.to_tensor(tl.constexpr(offset))) + check_offsets.append(tl.constexpr(offset)) + else: + # Assume it's already a tensor + new_offsets.append(offset) + check_offsets.append(offset) + + check_subview(src, check_offsets, new_sizes, new_strides) + return semantic.subview(src, new_offsets, new_sizes, new_strides, _semantic.builder) diff --git a/third_party/ascend/spec/triton/extension/buffer/language/semantic.py b/third_party/ascend/spec/triton/extension/buffer/language/semantic.py new file mode 100644 index 0000000000..dcbabb709e --- /dev/null +++ b/third_party/ascend/spec/triton/extension/buffer/language/semantic.py @@ -0,0 +1,128 @@ +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +from typing import (TypeVar, List) + +from triton._C.libtriton import ir +import triton.language.core as tl + +from . import core as bl + +T = TypeVar('T') + + +def alloc(etype: tl.dtype, shape: List[tl.constexpr], address_space: bl.address_space, is_mem_unique, + builder: ir.builder) -> bl.buffer: + shape = tl._unwrap_shape(shape) + if etype == tl.int1: + raise TypeError("Unsupported alloc int1 type") + if not isinstance(shape, (tl.tuple, list)): + raise TypeError("shape must be list/tuple") + etype = tl._unwrap_if_constexpr(etype) + address_space = tl._unwrap_if_constexpr(address_space) + element_ty_ir = etype.to_ir(builder) + addr_space_attr = (address_space.to_ir(builder) if address_space else builder.get_null_attr()) + memref_ty = builder.get_buffer_ty(shape, element_ty_ir, addr_space_attr) + handle = builder.alloc(memref_ty) + if is_mem_unique: + builder.create_annotation_mark(handle, "mem_unique", builder.get_unit_attr()) + builder.create_annotation_mark(handle, "effects", builder.get_str_array_attr(["write", "read"])) + + buffer_ty = bl.buffer_type(element_ty=etype, shape=shape, space=address_space) + return bl.buffer(handle, buffer_ty) + + +def to_buffer( + tensor: tl.tensor, + address_space: bl.address_space, + bind_buffer: bl.buffer, + builder: ir.builder, +) -> bl.buffer: + if not isinstance(tensor.shape, (tl.tuple, list)) or not tensor.shape: + raise TypeError("scalar type cannot be converted to buffer") + if isinstance(bind_buffer, bl.buffer): + builder.create_bind_buffer(tensor.handle, bind_buffer.handle) + return bind_buffer + if bind_buffer is not None: + raise ValueError("bind_buffer must be a buffer or None") + address_space = tl._unwrap_if_constexpr(address_space) + addr_space_attr = (address_space.to_ir(builder) if address_space else builder.get_null_attr()) + handle = builder.to_buffer(tensor.handle, addr_space_attr) + buffer_ty = bl.buffer_type(element_ty=tensor.dtype, shape=tensor.shape, space=address_space) + return bl.buffer(handle, buffer_ty) + + +def to_tensor(memref: bl.buffer, writable: bool, builder: ir.builder, target_shape=None) -> tl.tensor: + if not isinstance(memref, bl.buffer): + raise TypeError("memref must be bl.buffer") + + need_convert_layout = False + shape = memref.shape + if target_shape: + need_convert_layout = True + shape = tl._unwrap_shape(target_shape) + assert shape != memref.shape, "target shape is the same as source shape" + if not isinstance(shape, (tl.tuple, list)): + raise TypeError("shape must be list/tuple") + tensor_type = tl.block_type(memref.dtype, shape) + + memref_value = memref.handle + if need_convert_layout: + buffer_ty = bl.buffer_type( + element_ty=memref.dtype, + shape=shape, + space=memref.space, + ) + memref_value = builder.create_convert_layout(memref_value, buffer_ty.to_ir(builder)) + + return tl.tensor(builder.to_tensor(memref_value, writable), tensor_type) + + +def subview(src: bl.buffer, offsets: List[tl.tensor], sizes: List[tl.constexpr], strides: List[tl.constexpr], + builder: ir.builder) -> bl.buffer: + + new_offsets = [offset.handle for offset in offsets] + sizes_int = tl._unwrap_shape(sizes) + strides_int = tl._unwrap_shape(strides) + + result_handle = builder.subview(src.handle, new_offsets, sizes_int, strides_int) + + # calculate the memory layout strides of the source buffer + if src.strides: + # use the strides of the source buffer + src_memory_strides = src.strides + else: + # calculate the default row-major strides + src_memory_strides = [] + stride = 1 + for dim_size in reversed(src.shape): + if dim_size < 0: + raise ValueError("Cannot compute strides for buffer with dynamic dimensions") + src_memory_strides.insert(0, stride) + stride *= dim_size + + result_memory_strides = [] + for src_stride, subview_stride in zip(src_memory_strides, strides_int): + result_memory_strides.append(src_stride * subview_stride) + + # create buffer_type with strides + buffer_ty = bl.buffer_type(element_ty=src.dtype, shape=sizes_int, space=src.space, strides=result_memory_strides) + return bl.buffer(result_handle, buffer_ty) diff --git a/third_party/ascend/spec/triton/tools/get_ascend_devices.py b/third_party/ascend/spec/triton/tools/get_ascend_devices.py new file mode 100644 index 0000000000..f6ba1c9c01 --- /dev/null +++ b/third_party/ascend/spec/triton/tools/get_ascend_devices.py @@ -0,0 +1,49 @@ +import os +import glob +import logging +import subprocess + +logger = logging.getLogger(__name__) + + +def get_ascend_devices(): + devices = [] + pci_path = '/sys/bus/pci/devices/*' + + for dev in glob.glob(pci_path): + try: + vendor_path = os.path.join(dev, 'vendor') + device_path = os.path.join(dev, 'device') + + if os.path.exists(vendor_path): + with open(vendor_path, 'r') as f: + vendor = f.read().strip() + + if vendor == "0x19e5" and os.path.exists(device_path): + with open(device_path, 'r') as f: + device = f.read().strip() + devices.append(device) + except (IOError, OSError) as e: + logger.warning(f"can not fetch device {dev}: {e}") + continue + + return devices + + +def check_npu_smi_device(): + try: + result = subprocess.run(["npu-smi", "info"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, + shell=False, timeout=100) + if result.returncode == 0: + output = result.stdout.lower() + return "ascend910_95" in output or "ascend950" in output or "910_958b" in output + return False + except Exception: + logger.warning("can not use command: npu-smi info") + return False + + +ascend_devices = get_ascend_devices() +pci_condition = any("0xd806" in dev for dev in ascend_devices) +npu_smi_condition = check_npu_smi_device() +is_compile_on_910_95 = pci_condition or npu_smi_condition diff --git a/third_party/ascend/spec_cpp/include/triton/Dialect/Triton/IR/TritonAttrDefs.td b/third_party/ascend/spec_cpp/include/triton/Dialect/Triton/IR/TritonAttrDefs.td new file mode 100644 index 0000000000..51a6008f20 --- /dev/null +++ b/third_party/ascend/spec_cpp/include/triton/Dialect/Triton/IR/TritonAttrDefs.td @@ -0,0 +1,180 @@ +/* + * Copyright 2018-2020 Philippe Tillet + * Copyright 2020-2022 OpenAI + * Copyright 2025- FlagOS Contributors + * + * Permission is hereby granted, free of charge, to any person obtaining + * a copy of this software and associated documentation files + * (the "Software"), to deal in the Software without restriction, + * including without limitation the rights to use, copy, modify, merge, + * publish, distribute, sublicense, and/or sell copies of the Software, + * and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be + * included in all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, + * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF + * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. + * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY + * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, + * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE + * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ + +#ifndef TRITON_ATTR_DEFS +#define TRITON_ATTR_DEFS + +include "mlir/IR/EnumAttr.td" + +// Attributes for LoadOp and StoreOp +def TT_CacheModifierAttr : I32EnumAttr< + "CacheModifier", "", + [ + I32EnumAttrCase<"NONE", 1, "none">, + I32EnumAttrCase<"CA", 2, "ca">, + I32EnumAttrCase<"CG", 3, "cg">, + I32EnumAttrCase<"WB", 4, "wb">, + I32EnumAttrCase<"CS", 5, "cs">, + I32EnumAttrCase<"WT", 6, "wt">, + I32EnumAttrCase<"CV", 7, "cv">, + ]> { + let cppNamespace = "::mlir::triton"; +} + +def TT_MemSemanticAttr : I32EnumAttr< + "MemSemantic", "", + [ + I32EnumAttrCase<"RELAXED", 1, "relaxed">, + I32EnumAttrCase<"ACQUIRE", 2, "acquire">, + I32EnumAttrCase<"RELEASE", 3, "release">, + I32EnumAttrCase<"ACQUIRE_RELEASE", 4, "acq_rel">, + ]> { + let cppNamespace = "::mlir::triton"; +} + +def TT_EvictionPolicyAttr : I32EnumAttr< + "EvictionPolicy", "", + [ + I32EnumAttrCase<"NORMAL", 1, "evict_normal">, + I32EnumAttrCase<"EVICT_FIRST", 2, "evict_first">, + I32EnumAttrCase<"EVICT_LAST", 3, "evict_last"> + ]> { + let cppNamespace = "::mlir::triton"; +} + +def TT_PaddingOptionAttr : I32EnumAttr< + "PaddingOption", "", + [ + I32EnumAttrCase<"PAD_ZERO", 1, "zero">, + // We can not set the string value to "NAN" because it is a keyword in C++ + I32EnumAttrCase<"PAD_NAN", 2, "nan"> + ]> { + let cppNamespace = "::mlir::triton"; +} + +// atomic +def TT_AtomicRMWAttr : I32EnumAttr< + "RMWOp", "", + [ + I32EnumAttrCase<"AND", 1, "and">, + I32EnumAttrCase<"OR", 2, "or">, + I32EnumAttrCase<"XOR", 3, "xor">, + I32EnumAttrCase<"ADD", 4, "add">, + I32EnumAttrCase<"FADD", 5, "fadd">, + I32EnumAttrCase<"MAX", 6, "max">, + I32EnumAttrCase<"MIN", 7, "min">, + I32EnumAttrCase<"UMAX", 8, "umax">, + I32EnumAttrCase<"UMIN", 9, "umin">, + I32EnumAttrCase<"XCHG", 10, "exch"> + ]> { + let cppNamespace = "::mlir::triton"; +} + +def TT_DescriptorReduceKindAttr : I32EnumAttr< + "DescriptorReduceKind", "", + [ + I32EnumAttrCase<"ADD", 1, "add">, + I32EnumAttrCase<"MIN", 2, "min">, + I32EnumAttrCase<"MAX", 3, "max">, + I32EnumAttrCase<"INC", 4, "inc">, + I32EnumAttrCase<"DEC", 5, "dec">, + I32EnumAttrCase<"AND", 6, "and">, + I32EnumAttrCase<"OR", 7, "or">, + I32EnumAttrCase<"XOR", 8, "xor">, + ]> { + let cppNamespace = "::mlir::triton"; +} + +def TT_MemSyncScopeAttr : I32EnumAttr< + "MemSyncScope", "", + [ + I32EnumAttrCase<"GPU", 1, "gpu">, + I32EnumAttrCase<"CTA", 2, "cta">, + I32EnumAttrCase<"SYSTEM", 3, "sys">, + ]> { + let cppNamespace = "::mlir::triton"; +} + +// Program ID dimensions. +def TT_ProgramDim : I32EnumAttr< + "ProgramIDDim", "", + [ + I32EnumAttrCase<"X", 0, "x">, + I32EnumAttrCase<"Y", 1, "y">, + I32EnumAttrCase<"Z", 2, "z">, + ]> { + let cppNamespace = "::mlir::triton"; +} + +// Rounding mode. +def TT_RoundingModeAttr : I32EnumAttr< + "RoundingMode", "", + [ + I32EnumAttrCase<"RTZ", 0, "rtz">, + I32EnumAttrCase<"RTNE", 1, "rtne">, + ]> { + let cppNamespace = "::mlir::triton"; +} + +// PropagateNan. +def TT_PropagateNanAttr : I32EnumAttr< + "PropagateNan", "", + [ + I32EnumAttrCase<"NONE", 0, "none">, + I32EnumAttrCase<"ALL", 0xFFFF, "all">, + ]> { + let cppNamespace = "::mlir::triton"; +} + +// InputPrecision +def TT_InputPrecisionAttr : I32EnumAttr< + "InputPrecision", "", + [ + I32EnumAttrCase<"TF32", 0, "tf32">, + I32EnumAttrCase<"TF32x3", 1, "tf32x3">, + I32EnumAttrCase<"IEEE", 2, "ieee">, + I32EnumAttrCase<"HF32", 3, "hf32">, + I32EnumAttrCase<"BF16x3", 4, "bf16x3">, + I32EnumAttrCase<"BF16x6", 5, "bf16x6"> + ]>{ + let cppNamespace = "::mlir::triton"; +} + +// Type for ScaleDotElemType kind of floats. +def TT_ScaleDotElemTypeAttr : I32EnumAttr< + "ScaleDotElemType", "", + [ + I32EnumAttrCase<"E4M3", 0, "e4m3">, + I32EnumAttrCase<"E5M2", 1, "e5m2">, + I32EnumAttrCase<"E2M3", 2, "e2m3">, + I32EnumAttrCase<"E3M2", 3, "e3m2">, + I32EnumAttrCase<"E2M1", 4, "e2m1">, + I32EnumAttrCase<"BF16", 5, "bf16">, + I32EnumAttrCase<"FP16", 6, "fp16"> + ]>{ + let cppNamespace = "::mlir::triton"; +} + +#endif diff --git a/third_party/ascend/spec_cpp/lib/Dialect/Triton/IR/Traits.cpp b/third_party/ascend/spec_cpp/lib/Dialect/Triton/IR/Traits.cpp new file mode 100644 index 0000000000..81d4ffcdbc --- /dev/null +++ b/third_party/ascend/spec_cpp/lib/Dialect/Triton/IR/Traits.cpp @@ -0,0 +1,276 @@ +/* + * Copyright 2018-2020 Philippe Tillet + * Copyright 2020-2022 OpenAI + * Copyright 2025- FlagOS Contributors + * + * Permission is hereby granted, free of charge, to any person obtaining + * a copy of this software and associated documentation files + * (the "Software"), to deal in the Software without restriction, + * including without limitation the rights to use, copy, modify, merge, + * publish, distribute, sublicense, and/or sell copies of the Software, + * and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be + * included in all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, + * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF + * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. + * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY + * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, + * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE + * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ + +#include "triton/Dialect/Triton/IR/Traits.h" + +#include + +#include "mlir/IR/TypeUtilities.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Types.h" +#include "triton/Dialect/Triton/IR/Utility.h" +#include "triton/Dialect/TritonGPU/IR/Types.h" +#include "llvm/Support/ErrorHandling.h" + +using namespace mlir; +using namespace mlir::triton::gpu; + +LogicalResult OpTrait::impl::verifyEquivalentType(Type typeA, Type typeB) { + auto memdescA = dyn_cast(typeA); + auto memdescB = dyn_cast(typeB); + if (memdescA || memdescB) { + if (!memdescA || !memdescB) + return failure(); + if (memdescA.getShape() != memdescB.getShape()) + return failure(); + if (memdescA.getAllocShape() != memdescB.getAllocShape()) + return failure(); + if (memdescA.getElementType() != memdescB.getElementType()) + return failure(); + if (memdescA.getMemorySpace() != memdescB.getMemorySpace()) + return failure(); + if (memdescA.getMutableMemory() != memdescB.getMutableMemory()) + return failure(); + + Attribute encodingA = memdescA.getEncoding(); + Attribute encodingB = memdescB.getEncoding(); + if (encodingA == encodingB) + return success(); + if (static_cast(encodingA) != static_cast(encodingB)) + return failure(); + + auto layoutInterface = + cast(&encodingA.getDialect()); + return layoutInterface->verifyLayoutsAreEqual(memdescA.getShape(), + encodingA, encodingB, {}); + } + auto tensorTypeA = dyn_cast(typeA); + auto tensorTypeB = dyn_cast(typeB); + if (!(bool(tensorTypeA) && bool(tensorTypeB))) + return typeA == typeB ? success() : failure(); + auto encodingA = tensorTypeA.getEncoding(); + auto encodingB = tensorTypeB.getEncoding(); + auto shapeA = tensorTypeA.getShape(); + auto shapeB = tensorTypeB.getShape(); + if (shapeA != shapeB) + return failure(); + if (tensorTypeA.getElementType() != tensorTypeB.getElementType()) + return failure(); + // If there's no encoding or the encodings are the same + if (encodingA == encodingB) + return success(); + if (bool(encodingA) != bool(encodingB)) + return failure(); + + return cast(&encodingA.getDialect()) + ->verifyLayoutsAreEqual(shapeA, encodingA, encodingB, {}); +} + +static LogicalResult verifySameEncoding(Type typeA, Type typeB, + bool allowTensorPointerType) { + // TODO(Keren): the allowTensorPointerType argument is a hack to allow. + // The type checking code is kind of a mess with the current design. + auto getEncoding = [=](Type type) -> Attribute { + Attribute ret; + if (auto tensorType = dyn_cast(type)) { + ret = tensorType.getEncoding(); + } + if (!allowTensorPointerType) { + assert(!triton::isTensorPointerType(type)); + } + return ret; + }; + auto encodingA = getEncoding(typeA); + auto encodingB = getEncoding(typeB); + if (!encodingA || !encodingB) + return success(); + return encodingA == encodingB ? success() : failure(); +} + +LogicalResult +OpTrait::impl::verifySameOperandsEncoding(Operation *op, + bool allowTensorPointerType) { + if (failed(verifyAtLeastNOperands(op, 1))) + return failure(); + + auto type = op->getOperand(0).getType(); + for (auto opType : llvm::drop_begin(op->getOperandTypes(), 1)) + if (failed(verifySameEncoding(opType, type, allowTensorPointerType))) + return op->emitOpError() << "requires the same encoding for all operands"; + + return success(); +} + +LogicalResult OpTrait::impl::verifySameOperandsAndResultEncoding( + Operation *op, bool allowTensorPointerType) { + if (op->getNumOperands() == 0) + return success(); + + if (failed(verifyAtLeastNOperands(op, 1)) || + failed(verifyAtLeastNResults(op, 1))) + return failure(); + + auto type = op->getOperand(0).getType(); + for (auto resultType : op->getResultTypes()) + if (failed(verifySameEncoding(resultType, type, allowTensorPointerType))) + return op->emitOpError() + << "requires the same encoding for all operands and results"; + + return verifySameOperandsEncoding(op, allowTensorPointerType); +} + +LogicalResult OpTrait::impl::verifyTensorSize(Operation *op) { + for (auto opType : op->getOperandTypes()) { + if (auto tensorType = dyn_cast(opType)) { + int64_t numElements = 1; + for (int64_t s : tensorType.getShape()) + numElements *= s; + if (numElements > maxTensorNumElements) + return op->emitError("Maximum allowed number of elements is ") + << maxTensorNumElements << ", but " << *op + << " has more than that"; + // FIXME:patched triton community + // if ((numElements & (numElements - 1)) != 0) + // return op->emitError("Number of elements must be power-of-two, but ") + // << *op << " doesn't follow the rule (" << numElements << ")" + // << " elements"; + } + } + for (auto opType : op->getResultTypes()) { + if (auto tensorType = dyn_cast(opType)) { + int64_t numElements = 1; + for (int64_t s : tensorType.getShape()) + numElements *= s; + if (numElements > maxTensorNumElements) + return op->emitError("Maximum allowed number of elements is ") + << maxTensorNumElements << ", but " << *op + << " has more than that"; + // FIXME:patched triton community + // if ((numElements & (numElements - 1)) != 0) + // return op->emitError("Number of elements must be power-of-two, but ") + // << *op << " doesn't follow the rule (" << numElements << ")" + // << " elements"; + } + } + return success(); +} + +// Check that the Triton layouts on op's operands and return types are valid. +// For example, we check that the number of warps per block in a Triton GPU +// blocked layout matches that of its module. +// +// It's a little weird to check these properties of a layout only when the +// layout is used in an op, since most of the properties don't actually depend +// on the op. They do depend on the *module*, though, and a layout is attached +// to a module only by virtue of being used in one of the module's ops. +LogicalResult OpTrait::impl::verifyTensorLayouts(Operation *op) { + auto checkLayout = [&](Value val, auto makeErr) -> LogicalResult { + // Only ranked tensors can have layouts. + auto rankedTy = dyn_cast(val.getType()); + if (!rankedTy) + return success(); + + mlir::Attribute layout = rankedTy.getEncoding(); + if (!layout) + return success(); + + Dialect &dialect = layout.getDialect(); + auto verifyLayoutInterface = + dyn_cast(&dialect); + if (verifyLayoutInterface) { + return verifyLayoutInterface->verifyTensorLayout(layout, rankedTy, op, + makeErr); + } + + return success(); + }; + + for (size_t i = 0; i < op->getNumOperands(); i++) { + auto operand = op->getOperand(i); + auto err = checkLayout(operand, [&]() { + // Stringify the operand using `printAsOperand`. This prints e.g. "%42" + // rather than the full definition. + std::string operandStr; + llvm::raw_string_ostream os(operandStr); + // If we don't assume verified, dump() will recursively call this + // function! + operand.printAsOperand(os, OpPrintingFlags().assumeVerified()); + + return op->emitError("Operand ") + << i << " (" << operand << ") has an invalid layout: "; + }); + if (!err.succeeded()) + return err; + } + + for (size_t i = 0; i < op->getNumResults(); i++) { + auto result = op->getResult(i); + auto err = checkLayout(result, [&]() { + if (op->getNumResults() == 1) { + return op->emitError("Result has an invalid layout: "); + } else { + return op->emitError("Result ") << i << " has an invalid layout: "; + } + }); + if (!err.succeeded()) + return err; + } + + return success(); +} + +static ArrayRef getTypeShape(Type type) { + auto rankedType = dyn_cast(type); + if (auto ptrType = dyn_cast(type)) + rankedType = dyn_cast(ptrType.getPointeeType()); + return rankedType ? rankedType.getShape() : ArrayRef(); +} + +LogicalResult OpTrait::impl::verifySameLoadStoreOperandsShape(Operation *op) { + if (failed(verifyAtLeastNOperands(op, 1))) + return failure(); + + auto firstOperandShape = getTypeShape(op->getOperand(0).getType()); + for (auto type : llvm::drop_begin(op->getOperandTypes(), 1)) + if (failed(verifyCompatibleShape(getTypeShape(type), firstOperandShape))) + return op->emitOpError() << "requires the same shape for all operands"; + + return success(); +} + +LogicalResult +OpTrait::impl::verifySameLoadStoreOperandsAndResultShape(Operation *op) { + if (failed(verifyAtLeastNOperands(op, 1)) || + failed(verifyAtLeastNResults(op, 1))) + return failure(); + + auto firstOperandShape = getTypeShape(op->getOperand(0).getType()); + for (auto type : op->getResultTypes()) + if (failed(verifyCompatibleShape(getTypeShape(type), firstOperandShape))) + return op->emitOpError() + << "requires the same shape for all operands and results"; + + return verifySameLoadStoreOperandsShape(op); +} diff --git a/third_party/ascend/test/Conversion/CommonIRToHIVM/reduce-sum-strength-pipeline.mlir b/third_party/ascend/test/Conversion/CommonIRToHIVM/reduce-sum-strength-pipeline.mlir new file mode 100644 index 0000000000..91fe496c47 --- /dev/null +++ b/third_party/ascend/test/Conversion/CommonIRToHIVM/reduce-sum-strength-pipeline.mlir @@ -0,0 +1,41 @@ +// Copyright 2026- Xcoresigma Technology Co., Ltd + +// RUN: triton-opt %s -reduce-sum-strength -triton-to-linalg-incubated='global-kernel=false' | FileCheck %s + +module { + tt.func @kernel(%afloat : !tt.ptr, %res : !tt.ptr) -> () { + %0 = tt.make_range {end = 512 : i32, start = 0 : i32} : tensor<512xi32> + %c256 = arith.constant 256 : i32 + %ct256 = tt.splat %c256 : i32 -> tensor<512xi32> + %ws = arith.muli %ct256, %0 : tensor<512xi32> + %1 = tt.expand_dims %ws {axis = 1 : i32} : tensor<512xi32> -> tensor<512x1xi32> + %moff = tt.broadcast %1 : tensor<512x1xi32> -> tensor<512x256xi32> + %3 = tt.make_range {end = 256 : i32, start = 0 : i32} : tensor<256xi32> + %4 = tt.expand_dims %3 {axis = 0 : i32} : tensor<256xi32> -> tensor<1x256xi32> + %koff = tt.broadcast %4 : tensor<1x256xi32> -> tensor<512x256xi32> + %mkoff = arith.addi %moff, %koff : tensor<512x256xi32> + %8 = tt.splat %afloat : !tt.ptr -> tensor<512x256x!tt.ptr> + %9 = tt.addptr %8, %mkoff : tensor<512x256x!tt.ptr>, tensor<512x256xi32> + %18 = tt.splat %res : !tt.ptr -> tensor<256x!tt.ptr> + %19 = tt.addptr %18, %3 : tensor<256x!tt.ptr>, tensor<256xi32> + %afm = tt.load %9 : tensor<512x256x!tt.ptr> + %5 = "tt.reduce"(%afm) ({ + ^bb0(%arg5: f32, %arg6: f32): + %21 = arith.addf %arg5, %arg6 : f32 + tt.reduce.return %21 : f32 + }) {axis = 0 : i32} : (tensor<512x256xf32>) -> tensor<256xf32> + tt.store %19, %5 : tensor<256x!tt.ptr> + tt.return + } +} + +// CHECK-LABEL: func.func @kernel +// 512 -> 256 -> 128 along the reduced axis, final linalg.reduce on 128x256. +// CHECK: %[[LO:.*]] = tensor.extract_slice {{.*}}[0, 0] [256, 256] [1, 1] : tensor<512x256xf32> to tensor<256x256xf32> +// CHECK: %[[HI:.*]] = tensor.extract_slice {{.*}}[256, 0] [256, 256] [1, 1] : tensor<512x256xf32> to tensor<256x256xf32> +// CHECK: linalg.generic {{.*}} ins(%[[LO]], %[[HI]] : tensor<256x256xf32>, tensor<256x256xf32>) +// CHECK: tensor.extract_slice {{.*}}[0, 0] [128, 256] [1, 1] : tensor<256x256xf32> to tensor<128x256xf32> +// CHECK: tensor.extract_slice {{.*}}[128, 0] [128, 256] [1, 1] : tensor<256x256xf32> to tensor<128x256xf32> +// CHECK: linalg.generic {{.*}} ins({{.*}} : tensor<128x256xf32>, tensor<128x256xf32>) +// CHECK: linalg.reduce ins({{.*}} : tensor<128x256xf32>) {{.*}} dimensions = [0] +// CHECK: bufferization.materialize_in_destination diff --git a/third_party/ascend/test/Conversion/CommonIRToHIVM/reduce-sum-strength.mlir b/third_party/ascend/test/Conversion/CommonIRToHIVM/reduce-sum-strength.mlir new file mode 100644 index 0000000000..becd35e3eb --- /dev/null +++ b/third_party/ascend/test/Conversion/CommonIRToHIVM/reduce-sum-strength.mlir @@ -0,0 +1,178 @@ +// Copyright 2026- Xcoresigma Technology Co., Ltd + +// RUN: triton-opt %s -split-input-file -reduce-sum-strength | FileCheck %s + +module { + // CHECK-LABEL: tt.func @rewrite_f32_512 + tt.func @rewrite_f32_512(%arg0: tensor<512xf32>) -> f32 { + // 512 -> 256 -> 128, final reduce on 128. + // CHECK: %[[LO:.*]] = tensor.extract_slice %arg0[0] [256] [1] : tensor<512xf32> to tensor<256xf32> + // CHECK: %[[HI:.*]] = tensor.extract_slice %arg0[256] [256] [1] : tensor<512xf32> to tensor<256xf32> + // CHECK: %[[ADD:.*]] = arith.addf %[[LO]], %[[HI]] : tensor<256xf32> + // CHECK: tensor.extract_slice %[[ADD]][0] [128] [1] : tensor<256xf32> to tensor<128xf32> + // CHECK: tensor.extract_slice %[[ADD]][128] [128] [1] : tensor<256xf32> to tensor<128xf32> + // CHECK: arith.addf {{.*}} : tensor<128xf32> + // CHECK: "tt.reduce"({{.*}}) <{axis = 0 : i32}> ({ + // CHECK: }) : (tensor<128xf32>) -> f32 + %0 = "tt.reduce"(%arg0) <{axis = 0 : i32}> ({ + ^bb0(%arg1: f32, %arg2: f32): + %1 = arith.addf %arg1, %arg2 : f32 + tt.reduce.return %1 : f32 + }) : (tensor<512xf32>) -> f32 + tt.return %0 : f32 + } +} + +// ----- + +module { + // CHECK-LABEL: tt.func @rewrite_f16_8192 + tt.func @rewrite_f16_8192(%arg0: tensor<8192xf16>) -> f16 { + // 8192 -> 4096 -> 2048 -> 1024 -> 512 -> 256, final reduce on 256. + // CHECK: tensor.extract_slice %arg0[0] [4096] [1] : tensor<8192xf16> to tensor<4096xf16> + // CHECK: tensor.extract_slice %arg0[4096] [4096] [1] : tensor<8192xf16> to tensor<4096xf16> + // CHECK: arith.addf {{.*}} : tensor<4096xf16> + // CHECK: tensor.extract_slice {{.*}}[0] [2048] [1] : tensor<4096xf16> to tensor<2048xf16> + // CHECK: tensor.extract_slice {{.*}}[2048] [2048] [1] : tensor<4096xf16> to tensor<2048xf16> + // CHECK: arith.addf {{.*}} : tensor<2048xf16> + // CHECK: tensor.extract_slice {{.*}}[0] [1024] [1] : tensor<2048xf16> to tensor<1024xf16> + // CHECK: tensor.extract_slice {{.*}}[1024] [1024] [1] : tensor<2048xf16> to tensor<1024xf16> + // CHECK: arith.addf {{.*}} : tensor<1024xf16> + // CHECK: tensor.extract_slice {{.*}}[0] [512] [1] : tensor<1024xf16> to tensor<512xf16> + // CHECK: tensor.extract_slice {{.*}}[512] [512] [1] : tensor<1024xf16> to tensor<512xf16> + // CHECK: arith.addf {{.*}} : tensor<512xf16> + // CHECK: tensor.extract_slice {{.*}}[0] [256] [1] : tensor<512xf16> to tensor<256xf16> + // CHECK: tensor.extract_slice {{.*}}[256] [256] [1] : tensor<512xf16> to tensor<256xf16> + // CHECK: arith.addf {{.*}} : tensor<256xf16> + // CHECK: "tt.reduce" + // CHECK: }) : (tensor<256xf16>) -> f16 + %0 = "tt.reduce"(%arg0) <{axis = 0 : i32}> ({ + ^bb0(%arg1: f16, %arg2: f16): + %1 = arith.addf %arg1, %arg2 : f16 + tt.reduce.return %1 : f16 + }) : (tensor<8192xf16>) -> f16 + tt.return %0 : f16 + } +} + +// ----- + +module { + // CHECK-LABEL: tt.func @rewrite_bf16_512 + tt.func @rewrite_bf16_512(%arg0: tensor<512xbf16>) -> bf16 { + // 512 -> 256, final reduce on 256. + // CHECK: tensor.extract_slice %arg0[0] [256] [1] : tensor<512xbf16> to tensor<256xbf16> + // CHECK: tensor.extract_slice %arg0[256] [256] [1] : tensor<512xbf16> to tensor<256xbf16> + // CHECK: arith.addf {{.*}} : tensor<256xbf16> + // CHECK: "tt.reduce" + // CHECK: }) : (tensor<256xbf16>) -> bf16 + %0 = "tt.reduce"(%arg0) <{axis = 0 : i32}> ({ + ^bb0(%arg1: bf16, %arg2: bf16): + %1 = arith.addf %arg1, %arg2 : bf16 + tt.reduce.return %1 : bf16 + }) : (tensor<512xbf16>) -> bf16 + tt.return %0 : bf16 + } +} + +// ----- + +module { + // CHECK-LABEL: tt.func @rewrite_axis1 + tt.func @rewrite_axis1(%arg0: tensor<4x512xf32>) -> tensor<4xf32> { + // 512 -> 256 -> 128 along axis 1, final reduce on 4x128. + // CHECK: tensor.extract_slice %arg0[0, 0] [4, 256] [1, 1] : tensor<4x512xf32> to tensor<4x256xf32> + // CHECK: tensor.extract_slice %arg0[0, 256] [4, 256] [1, 1] : tensor<4x512xf32> to tensor<4x256xf32> + // CHECK: arith.addf {{.*}} : tensor<4x256xf32> + // CHECK: tensor.extract_slice {{.*}}[0, 0] [4, 128] [1, 1] : tensor<4x256xf32> to tensor<4x128xf32> + // CHECK: tensor.extract_slice {{.*}}[0, 128] [4, 128] [1, 1] : tensor<4x256xf32> to tensor<4x128xf32> + // CHECK: arith.addf {{.*}} : tensor<4x128xf32> + // CHECK: "tt.reduce" + // CHECK: }) : (tensor<4x128xf32>) -> tensor<4xf32> + %0 = "tt.reduce"(%arg0) <{axis = 1 : i32}> ({ + ^bb0(%arg1: f32, %arg2: f32): + %1 = arith.addf %arg1, %arg2 : f32 + tt.reduce.return %1 : f32 + }) : (tensor<4x512xf32>) -> tensor<4xf32> + tt.return %0 : tensor<4xf32> + } +} + +// ----- + +module { + // CHECK-LABEL: tt.func @rewrite_f32_1024 + tt.func @rewrite_f32_1024(%arg0: tensor<1024xf32>) -> f32 { + // 1024 -> 512 -> 256 -> 128, final reduce on 128. + // CHECK: tensor.extract_slice %arg0[0] [512] [1] : tensor<1024xf32> to tensor<512xf32> + // CHECK: tensor.extract_slice %arg0[512] [512] [1] : tensor<1024xf32> to tensor<512xf32> + // CHECK: arith.addf {{.*}} : tensor<512xf32> + // CHECK: tensor.extract_slice {{.*}}[0] [256] [1] : tensor<512xf32> to tensor<256xf32> + // CHECK: tensor.extract_slice {{.*}}[256] [256] [1] : tensor<512xf32> to tensor<256xf32> + // CHECK: arith.addf {{.*}} : tensor<256xf32> + // CHECK: tensor.extract_slice {{.*}}[0] [128] [1] : tensor<256xf32> to tensor<128xf32> + // CHECK: tensor.extract_slice {{.*}}[128] [128] [1] : tensor<256xf32> to tensor<128xf32> + // CHECK: arith.addf {{.*}} : tensor<128xf32> + // CHECK: "tt.reduce" + // CHECK: }) : (tensor<128xf32>) -> f32 + %0 = "tt.reduce"(%arg0) <{axis = 0 : i32}> ({ + ^bb0(%arg1: f32, %arg2: f32): + %1 = arith.addf %arg1, %arg2 : f32 + tt.reduce.return %1 : f32 + }) : (tensor<1024xf32>) -> f32 + tt.return %0 : f32 + } +} + +// ----- + +module { + // CHECK-LABEL: tt.func @skip_f32_8200 + tt.func @skip_f32_8200(%arg0: tensor<8200xf32>) -> f32 { + // Even and inside the f32 range, but not a power of two: must not be + // rewritten, otherwise the greedy splitting degenerates to a + // non-power-of-two tail (8200 -> 4100 -> 2050 -> 1025). + // CHECK-NOT: tensor.extract_slice + // CHECK: "tt.reduce"(%arg0) <{axis = 0 : i32}> + // CHECK: }) : (tensor<8200xf32>) -> f32 + %0 = "tt.reduce"(%arg0) <{axis = 0 : i32}> ({ + ^bb0(%arg1: f32, %arg2: f32): + %1 = arith.addf %arg1, %arg2 : f32 + tt.reduce.return %1 : f32 + }) : (tensor<8200xf32>) -> f32 + tt.return %0 : f32 + } +} + +// ----- + +module { + // CHECK-LABEL: tt.func @skip_i32_512 + tt.func @skip_i32_512(%arg0: tensor<512xi32>) -> i32 { + // CHECK-NOT: tensor.extract_slice + // CHECK: "tt.reduce"(%arg0) <{axis = 0 : i32}> + // CHECK: }) : (tensor<512xi32>) -> i32 + %0 = "tt.reduce"(%arg0) <{axis = 0 : i32}> ({ + ^bb0(%arg1: i32, %arg2: i32): + %1 = arith.addi %arg1, %arg2 : i32 + tt.reduce.return %1 : i32 + }) : (tensor<512xi32>) -> i32 + tt.return %0 : i32 + } +} + +// ----- + +module { + // CHECK-LABEL: tt.func @skip_non_add_reduce + tt.func @skip_non_add_reduce(%arg0: tensor<512xf32>) -> f32 { + // CHECK-NOT: tensor.extract_slice + // CHECK: arith.maximumf + %0 = "tt.reduce"(%arg0) <{axis = 0 : i32}> ({ + ^bb0(%arg1: f32, %arg2: f32): + %1 = arith.maximumf %arg1, %arg2 : f32 + tt.reduce.return %1 : f32 + }) : (tensor<512xf32>) -> f32 + tt.return %0 : f32 + } +} diff --git a/third_party/ascend/triton_ascend.cc b/third_party/ascend/triton_ascend.cc new file mode 100644 index 0000000000..9402a41b2a --- /dev/null +++ b/third_party/ascend/triton_ascend.cc @@ -0,0 +1,462 @@ +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Linalg/Passes.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Tools/mlir-opt/MlirOptMain.h" + +#include "ascend/include/AutoBlockify/Passes.h" +#include "ascend/include/CommonIRToHIVM/ReduceSumStrengthPass.h" +#include "ascend/include/CommonIRToHIVM/Passes.h" +#include "ascend/include/Dialect/TritonAscend/IR/TritonAscendDialect.h" +#include "ascend/include/TritonToAnnotation/Passes.h" +#include "ascend/include/TritonToHFusion/Passes.h" +#include "ascend/include/TritonToHIVM/Passes.h" +#include "ascend/include/TritonToLLVM/Passes.h" +#include "incubated/Conversion/DiscreteMaskAccessConversion/Passes.h" +#include "incubated/Conversion/TritonToLinalgIncubated/Passes.h" +#include "incubated/Conversion/TritonToStructuredIncubated/Passes.h" +#include "incubated/Conversion/TritonToUnstructureIncubated/Passes.h" +#include "mlir-ext/Dialect/CommonIR/Transforms/Passes.h" + +#include "ascend/include/DynamicCVPipeline/Common/BufferCountManager.h" +#include "ascend/include/DynamicCVPipeline/Passes.h" +// todo: this code will be removed in version 530. +#include "ascend/include/TritonAffinityOpt/Passes.h" + +#include "ir.h" // TritonOpBuilder +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include + +#ifdef __TLE_DSA__ +#include "llvm/Config/llvm-config.h" +#endif + +namespace py = pybind11; +using namespace ir; +using namespace mlir; + +void init_triton_ascend_ir(py::module &&m) { + auto *builder_cls = ir::getBuilderClass(); + builder_cls + ->def("create_extract_scalar", + [](TritonOpBuilder &self, Value &src, + std::vector &indices) -> Value { + llvm::SmallVector arg_indices; + for (const auto &i : indices) { + auto iTy = i.getType(); + if (!iTy.isIndex()) { + auto v = self.create( + self.getBuilder().getIndexType(), i); + arg_indices.push_back(v); + } else { + arg_indices.push_back(i); + } + } + auto ret = self.create(src, arg_indices); + return ret; + }) + .def("create_extract_slice", + [](TritonOpBuilder &self, Value &ful, std::vector &offs_vec, + std::vector &sizs_vec, std::vector &strd_vec) -> Value { + llvm::SmallVector offsets; + llvm::SmallVector staticOffsets; + for (const auto &o : offs_vec) { + auto oTy = o.getType(); + if (!oTy.isIndex()) { + auto v = self.create( + self.getBuilder().getIndexType(), o); + offsets.push_back(v); + } else { + offsets.push_back(o); + } + staticOffsets.push_back(ShapedType::kDynamic); + } + llvm::SmallVector sizes; + llvm::SmallVector staticSizes; + llvm::SmallVector retSizes; + for (const auto &s : sizs_vec) { + // auto v = self.create(s); + // sizes.push_back(v); + staticSizes.push_back(s); + retSizes.push_back(s); + } + llvm::SmallVector strides; + llvm::SmallVector staticStrides; + for (const auto &s : strd_vec) { + auto v = self.create(s); + strides.push_back(v); + staticStrides.push_back(ShapedType::kDynamic); + } + auto retTy = RankedTensorType::get( + retSizes, + cast(ful.getType()).getElementType()); + + return self.create( + retTy, ful, offsets, sizes, strides, staticOffsets, + staticSizes, staticStrides); + }) + .def("create_insert_slice", + [](TritonOpBuilder &self, Value &ful, Value &sub, + std::vector &offs_vec, std::vector &sizs_vec, + std::vector &strd_vec) -> Value { + llvm::SmallVector offsets; + llvm::SmallVector staticOffsets; + for (const auto &o : offs_vec) { + auto oTy = o.getType(); + if (!oTy.isIndex()) { + auto v = self.create( + self.getBuilder().getIndexType(), o); + offsets.push_back(v); + } else { + offsets.push_back(o); + } + staticOffsets.push_back(ShapedType::kDynamic); + } + llvm::SmallVector sizes; + llvm::SmallVector staticSizes; + llvm::SmallVector retSizes; + for (const auto &s : sizs_vec) { + // auto v = self.create(s); + // sizes.push_back(v); + staticSizes.push_back(s); + retSizes.push_back(s); + } + llvm::SmallVector strides; + llvm::SmallVector staticStrides; + for (const auto &s : strd_vec) { + auto v = self.create(s); + strides.push_back(v); + staticStrides.push_back(ShapedType::kDynamic); + } + auto retTy = RankedTensorType::get( + retSizes, + cast(ful.getType()).getElementType()); + auto ret = self.create( + sub, ful, offsets, sizes, strides, staticOffsets, staticSizes, + staticStrides); + return ret; + }) + .def("create_custom_op_for_inter_core_sync", + [](TritonOpBuilder &self, std::string &op_name, + std::string &mode_or_sender, int id) -> void { + auto args = self.getBuilder().getArrayAttr( + {self.getBuilder().getStringAttr(mode_or_sender), + self.getBuilder().getI32IntegerAttr(id)}); + self.create(op_name, args, ValueRange()); + }) + .def("create_index_select_simd", + [](TritonOpBuilder &self, Value &src, Value &index, int32_t dim, + std::vector &srcShape, std::vector &srcOffset, + std::vector &readShape, + std::vector &returnShape) -> Value { + auto &builder = self.getBuilder(); + auto loc = self.getLastLoc(); + + // Get element type from source pointer + Type elemType; + if (auto ptrTy = dyn_cast(src.getType())) { + elemType = ptrTy.getPointeeType(); + } else { + llvm::report_fatal_error( + "index_select_simd: src must be pointer type"); + } + + // Create return tensor type + llvm::SmallVector retShape; + for (const auto &s : returnShape) { + retShape.push_back(s); + } + auto retTensorType = RankedTensorType::get(retShape, elemType); + + // Convert srcShape and srcOffset values to index type if needed + llvm::SmallVector srcShapeIndex; + for (auto val : srcShape) { + if (!val.getType().isIndex()) { + val = self.create(builder.getIndexType(), + val); + } + srcShapeIndex.push_back(val); + } + + llvm::SmallVector srcOffsetIndex; + for (auto val : srcOffset) { + if (!val.getType().isIndex()) { + val = self.create(builder.getIndexType(), + val); + } + srcOffsetIndex.push_back(val); + } + + // Create attributes + auto dimAttr = builder.getI32IntegerAttr(dim); + auto readShapeAttr = builder.getDenseI32ArrayAttr(readShape); + + // Create the IndexSelectSimdOp + // Parameter order must match TritonOps.td definition: + // src, index, dim, src_shape, src_offset, read_shape + auto indexSelectSimdOp = + builder.create( + loc, + retTensorType, // result type + src, // src pointer + index, // index tensor + dimAttr, // dim attribute + srcShapeIndex, // src_shape (variadic, index type) + srcOffsetIndex, // src_offset (variadic, index type) + readShapeAttr // read_shape attribute + ); + + return indexSelectSimdOp.getResult(); + }) + .def("create_index_put", + [](TritonOpBuilder &self, Value &ptr, Value &index, Value &value, + const int32_t dim, const int64_t indexBoundary, + std::vector &endOffset, std::vector &startOffset, + std::vector &dstStride) -> void { + // dim need to be i32 type + auto dimI32Ty = self.getBuilder().getI32Type(); + auto dim_val = self.create( + dimI32Ty, static_cast(dim)); + // indexBoundary need to be i64 type + auto BoundI64Ty = self.getBuilder().getI64Type(); + auto bound_val = self.create( + BoundI64Ty, static_cast(indexBoundary)); + + self.create(ptr, index, value, dim_val, + bound_val, endOffset, + startOffset, dstStride); + }) + .def("create_gather_out_to_ub", + [](TritonOpBuilder &self, Value &src, Value &index, + const int64_t indexBoundary, const int32_t dim, + std::vector &srcStride, std::vector &endOffset, + std::vector &startOffset, + std::optional &other) -> Value { + auto elemTy = cast(src.getType()).getPointeeType(); + auto idxTy = cast(index.getType()); + auto idxShape = idxTy.getShape(); + std::vector retShape(idxShape.begin(), idxShape.end()); + auto resType = RankedTensorType::get(retShape, elemTy); + + // indexBoundary need to be i64 type + auto BoundI64Ty = self.getBuilder().getI64Type(); + auto bound_val = self.create( + BoundI64Ty, static_cast(indexBoundary)); + // dim need to be i32 type + auto dimI32Ty = self.getBuilder().getI32Type(); + auto dim_val = self.create( + dimI32Ty, static_cast(dim)); + return self.create( + resType, src, index, bound_val, dim_val, srcStride, endOffset, + startOffset, other.value_or(Value())); + }) + .def("create_scatter_ub_to_out", + [](TritonOpBuilder &self, Value &ptr, Value &value, Value &index, + const int64_t indexBoundary, const int32_t dim, + std::vector &dstStride, std::vector &endOffset, + std::vector &startOffset) -> void { + auto idxTy = cast(index.getType()); + + // indexBoundary need to be i64 type + auto BoundI64Ty = self.getBuilder().getI64Type(); + auto bound_val = self.create( + BoundI64Ty, static_cast(indexBoundary)); + // dim need to be i32 type + auto dimI32Ty = self.getBuilder().getI32Type(); + auto dim_val = self.create( + dimI32Ty, static_cast(dim)); + + self.create( + ptr, value, index, bound_val, dim_val, dstStride, endOffset, + startOffset); + }) + // conv1d operation + .def( + "create_conv1d", + [](TritonOpBuilder &self, Value input, Value weight, py::object bias, + int64_t stride, int64_t padding_size, int64_t dilation, + int64_t groups, Type output_type) -> Value { + Value biasValue; + if (!bias.is_none()) { + biasValue = bias.cast(); + } else { + biasValue = Value(); + } + auto &builder = self.getBuilder(); + auto strideAttr = builder.getI64IntegerAttr(stride); + auto paddingSizeAttr = builder.getI64IntegerAttr(padding_size); + auto dilationAttr = builder.getI64IntegerAttr(dilation); + auto groupsAttr = builder.getI64IntegerAttr(groups); + auto op = self.create( + output_type, input, weight, biasValue, strideAttr, + paddingSizeAttr, dilationAttr, groupsAttr); + return op.getResult(); + }, + py::arg("input"), py::arg("weight"), py::arg("bias"), + py::arg("stride"), py::arg("padding_size"), py::arg("dilation"), + py::arg("groups"), py::arg("output_type")) + // Add sort + .def("create_sort", + [](TritonOpBuilder &self, Value src, int64_t dim, + bool descending) -> Value { + auto &builder = self.getBuilder(); + auto loc = self.getLastLoc(); + + auto dimAttr = builder.getI64IntegerAttr(dim); + auto descendingAttr = builder.getBoolAttr(descending); + + auto op = builder.create(loc, src, dimAttr, + descendingAttr); + + return op->getResult(0); + }) + // Add flip + .def("create_flip", + [](TritonOpBuilder &self, Value src, int64_t dim) -> Value { + auto &builder = self.getBuilder(); + auto loc = self.getLastLoc(); + + auto dimAttr = builder.getI64IntegerAttr(dim); + + auto op = + builder.create(loc, src, dimAttr); + + return op->getResult(0); + }) + .def("create_annotation", + [](TritonOpBuilder &self, Value &ptr, const std::string &attrKey, + Attribute &attrVal) { + auto annotationOp = self.create(ptr); + annotationOp->setAttr(self.getBuilder().getStringAttr(attrKey), + attrVal); + }); +} + +void init_triton_ascend_passes_ttir(py::module &&m) { + m.def("add_lane_vectorize", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::tile::createTritonLaneVectorize()); + }); + + m.def("add_auto_blockify", [](mlir::PassManager &pm, int autoBlockifySize) { + AutoBlockifyOptions opts; + opts.autoBlockifySize = autoBlockifySize; + pm.addPass(mlir::triton::createAutoBlockifyPass(opts)); + }); + + m.def("add_reduce_sum_strength", [](mlir::PassManager &pm, bool enable, + int32_t splitFactor) { + pm.addPass(mlir::triton::createReduceSumStrengthPass(enable, splitFactor)); + }); + + m.def("add_triton_to_structure", + [](mlir::PassManager &pm, bool enableMaskFallbackConversion, + bool optimizeDynamicOffset) { + pm.addPass(mlir::triton::createTritonToStructuredIncubatedPass( + enableMaskFallbackConversion, optimizeDynamicOffset)); + }); + + m.def("add_triton_to_annotation", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::createTritonToAnnotationPass()); + }); + + m.def("add_triton_to_linalg", + [](mlir::PassManager &pm, bool globalKernel, bool namedOps, + bool enableNd2nzOnVector, bool enableSelectAnalysis, + bool compileOn91095) { + pm.addPass(mlir::triton::Incubated::createTritonToLinalgIncubatedPass( + globalKernel, namedOps, enableNd2nzOnVector, enableSelectAnalysis, + compileOn91095)); + }); + + m.def("add_triton_to_unstructure", + [](mlir::PassManager &pm, bool compileOn91095, bool forceSimtTemplate) { + TritonToUnstructureIncubatedOptions opts; + opts.compileOn91095 = compileOn91095; + opts.forceSimtTemplate = forceSimtTemplate; + pm.addPass( + mlir::triton::createTritonToUnstructureIncubatedPass(opts)); + }); + + m.def("add_triton_to_hfusion", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::createTritonToHFusionPass()); + }); + + m.def("add_discrete_mask_access_conversion", [](mlir::PassManager &pm, + bool compileOn91095, + bool forceSimtTemplate, + bool enableSyncBlockLock) { + DiscreteMaskAccessConversionOptions opts; + opts.compileOn91095 = compileOn91095; + opts.forceSimtTemplate = forceSimtTemplate; + opts.enableSyncBlockLock = enableSyncBlockLock; + pm.addPass(mlir::triton::createDiscreteMaskAccessConversionPass(opts)); + }); + + m.def("add_triton_to_hivm", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::createTritonToHIVMPass()); + }); + + m.def("add_commonir_to_hivm", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::createCommonIRToHIVMPass()); + }); + + m.def("add_triton_to_llvm", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::createTritonToLLVMPass()); + }); + + m.def("add_bubble_up_operation", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::createBubbleUpOperationPass()); + }); + + m.def("add_dynamic_cv_pipeline", + [](mlir::PassManager &pm, bool compileOn91095) { + AddDynamicCVPipelineOptions opts; + opts.compileOn91095 = compileOn91095; + pm.addPass(mlir::triton::createAddDynamicCVPipelinePass(opts)); + }); + + // todo: this code will be removed in version 530. + m.def("add_dag_sync", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::createDAGSyncPass()); + }); + + m.def("add_dag_scope", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::createDAGScopePass()); + }); + + m.def("add_dag_ssbuffer", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::createDAGSSBufferPass()); + }); + + m.def("set_buffer_count", [](int type, int count) { + auto depType = static_cast(type); + mlir::triton::BufferCountManager::getInstance().setBufferCount(depType, + count); + }); +} + +// Forward declaration for ascend_ir bindings (defined in ascend_ir.cc) +void init_ascend_ir(py::module &&m); + +void init_triton_ascend(py::module &&m) { + auto passes = m.def_submodule("passes"); + // load dialects + m.def("load_dialects", [](mlir::MLIRContext &context) { + mlir::DialectRegistry registry; + registry.insert(); + context.appendDialectRegistry(registry); + context.loadAllAvailableDialects(); + }); + + init_triton_ascend_passes_ttir(passes.def_submodule("ttir")); + init_triton_ascend_ir(m.def_submodule("ascend_ir")); + + // Initialize ascend IR bindings (ascendnpu_ir_builder, scope/hivm dialects) + init_ascend_ir(m.def_submodule("ir")); +} diff --git a/third_party/ascend/tutorials/01-vector-add.py b/third_party/ascend/tutorials/01-vector-add.py new file mode 100644 index 0000000000..e59935f456 --- /dev/null +++ b/third_party/ascend/tutorials/01-vector-add.py @@ -0,0 +1,101 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Vector Addition +=============== + +In this tutorial, you will write a simple vector addition using Triton. + +In doing so, you will learn about: + +* The basic programming model of Triton. + +* The `triton.jit` decorator, which is used to define Triton kernels. + +* The best practices for validating and benchmarking your custom ops against native reference implementations. + +""" + +# %% +# Compute Kernel +# -------------- + +import torch +import torch_npu + +import triton +import triton.language as tl + + +@triton.jit +def add_kernel(x_ptr, # *Pointer* to first input vector. + y_ptr, # *Pointer* to second input vector. + output_ptr, # *Pointer* to output vector. + n_elements, # Size of the vector. + BLOCK_SIZE: tl.constexpr, # Number of elements each program should process. + # NOTE: `constexpr` so it can be used as a shape value. + ): + # There are multiple 'programs' processing different data. We identify which program + # we are here: + pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. + # This program will process inputs that are offset from the initial data. + # For instance, if you had a vector of length 256 and block_size of 64, the programs + # would each access the elements [0:64, 64:128, 128:192, 192:256]. + # Note that offsets is a list of pointers: + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + # Create a mask to guard memory operations against out-of-bounds accesses. + mask = offsets < n_elements + # Load x and y from DRAM, masking out any extra elements in case the input is not a + # multiple of the block size. + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + # Write x + y back to DRAM. + tl.store(output_ptr + offsets, output, mask=mask) + + +# %% +# Let's also declare a helper function to (1) allocate the `z` tensor +# and (2) enqueue the above kernel with appropriate grid/block sizes: + + +def add(x: torch.Tensor, y: torch.Tensor): + output = torch.empty_like(x) + n_elements = output.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024) + return output + + +# %% +# We can now use the above function to compute the element-wise sum of two `torch.tensor` objects and test its correctness: +torch.manual_seed(0) +size = 98432 +x = torch.rand(size, device='npu') +y = torch.rand(size, device='npu') +output_torch = x + y +output_triton = add(x, y) +print(output_torch) +print(output_triton) +print(f'The maximum difference between torch and triton is ' + f'{torch.max(torch.abs(output_torch - output_triton))}') diff --git a/third_party/ascend/tutorials/02-fused-softmax.py b/third_party/ascend/tutorials/02-fused-softmax.py new file mode 100644 index 0000000000..e476061a40 --- /dev/null +++ b/third_party/ascend/tutorials/02-fused-softmax.py @@ -0,0 +1,124 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Fused Softmax +============= +""" + +import torch +import torch_npu +import triton +import triton.language as tl +from triton.runtime import driver + + +def naive_softmax(x): + """Compute row-wise softmax of X using native pytorch + + We subtract the maximum element in order to avoid overflows. Softmax is invariant to + this shift. + """ + # read MN elements ; write M elements + x_max = x.max(dim=1)[0] + # read MN + M elements ; write MN elements + z = x - x_max[:, None] + # read MN elements ; write MN elements + numerator = torch.exp(z) + # read MN elements ; write M elements + denominator = numerator.sum(dim=1) + # read MN + M elements ; write MN elements + ret = numerator / denominator[:, None] + # in total: read 5MN + 2M elements ; wrote 3MN + 2M elements + return ret + + +@triton.jit +def softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride, n_rows, n_cols, + BLOCK_SIZE: tl.constexpr): + # starting row of the program + row_start = tl.program_id(0) + row_step = tl.num_programs(0) + for row_idx in tl.range(row_start, n_rows, row_step): + # The stride represents how much we need to increase the pointer to advance 1 row + row_start_ptr = input_ptr + row_idx * input_row_stride + # The block size is the next power of two greater than n_cols, so we can fit each + # row in a single block + col_offsets = tl.arange(0, BLOCK_SIZE) + input_ptrs = row_start_ptr + col_offsets + # Load the row into SRAM, using a mask since BLOCK_SIZE may be > than n_cols + mask = col_offsets < n_cols + row = tl.load(input_ptrs, mask=mask, other=-float('inf')) + # Subtract maximum for numerical stability + row_minus_max = row - tl.max(row, axis=0) + + numerator = tl.exp(row_minus_max) + denominator = tl.sum(numerator, axis=0) + softmax_output = numerator / denominator + # Write back output to DRAM + output_row_start_ptr = output_ptr + row_idx * output_row_stride + output_ptrs = output_row_start_ptr + col_offsets + tl.store(output_ptrs, softmax_output, mask=mask) + + +target = triton.runtime.driver.active.get_current_target() +kernels = {} + + +def softmax(x, stream): + n_rows, n_cols = x.shape + + # The block size of each loop iteration is the smallest power of two greater than the number of columns in `x` + BLOCK_SIZE = triton.next_power_of_2(n_cols) + + # Allocate output + y = torch.empty_like(x) + + # pre-compile kernel to get register usage and compute thread occupancy. + kernel, num_programs = kernels.get(BLOCK_SIZE, (None, 0)) + if kernel is None: + num_programs = 32 + kernel = softmax_kernel + kernels[BLOCK_SIZE] = (kernel, num_programs) + + num_programs = min(num_programs, n_rows) + + # Create a number of persistent programs. + kernel[(num_programs, 1, 1)](y, x, x.stride(0), y.stride(0), n_rows, n_cols, BLOCK_SIZE) + return y + + +# %% +# Unit Test +# --------- + +# %% +# We make sure that we test our kernel on a matrix with an irregular number of rows and columns. +# This will allow us to verify that our padding mechanism works. +device = torch.npu.current_device() +stream = torch.npu.current_stream(device).npu_stream +torch.manual_seed(0) +x = torch.randn(1823, 781, device='npu') +y_triton = softmax(x, stream) +y_torch = torch.softmax(x, axis=1) +assert torch.allclose(y_triton, y_torch), (y_triton, y_torch) +print(y_triton) +print(y_torch) +print(f'The maximum difference between torch and triton is ' + f'{torch.max(torch.abs(y_triton-y_torch))}') diff --git a/third_party/ascend/tutorials/03-matrix-multiplication.py b/third_party/ascend/tutorials/03-matrix-multiplication.py new file mode 100644 index 0000000000..3b63adb66b --- /dev/null +++ b/third_party/ascend/tutorials/03-matrix-multiplication.py @@ -0,0 +1,199 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Matrix Multiplication +=============== +""" + +import triton +import triton.language as tl +import torch +import torch_npu +import triton.language.extra.cann.extension as extension + +DEV = "npu" + + +def get_autotune_config(): + return [ + triton.Config({"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 128}), + triton.Config({"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64}), + ] + + +@triton.autotune( + configs=get_autotune_config(), + key=["M", "N", "K"], +) +@triton.jit +def matmul_kernel( + # Pointers to matrices + a_ptr, b_ptr, c_ptr, + # Matrix dimensions + M, N, K, + # The stride variables represent how much to increase the ptr by when moving by 1 + # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr` + # by to get the element one row down (A has M rows). + stride_am, stride_ak, # + stride_bk, stride_bn, # + stride_cm, stride_cn, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, # + ACTIVATION: tl.constexpr, # +): + """Kernel for computing the matmul C = A x B. + A has shape (M, K), B has shape (K, N) and C has shape (M, N) + """ + GROUP_SIZE_M: tl.constexpr = 1 + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + # See above `L2 Cache Optimizations` section for details. + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # ---------------------------------------------------------- + # Create pointers for the first blocks of A and B. + # We will advance this pointer as we move in the K direction + # and accumulate + # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers + # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers + # See above `Pointer Arithmetic` section for details + offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs_base = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs_base = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + msk_m = offs_am < M + msk_n = offs_bn < N + # ----------------------------------------------------------- + # Iterate to compute a block of the C matrix. + # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block + # of fp32 values for higher accuracy. + # `accumulator` will be converted back to fp16 after the loop. + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + # Load the next block of A and B, generate a mask by checking the K dimension. + # If it is out of bounds, set it to 0. + a_ptrs = a_ptrs_base + k * BLOCK_SIZE_K * stride_ak + b_ptrs = b_ptrs_base + k * BLOCK_SIZE_K * stride_bk + a = tl.load( + a_ptrs, + mask=msk_m[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + b = tl.load( + b_ptrs, + mask=msk_n[None, :] & (offs_k[:, None] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + # We accumulate along the K dimension. + accumulator = tl.dot(a, b, accumulator) + # You can fuse arbitrary activation functions here + # while the accumulator is still in FP32! + # Original vector operations + # # ----------------------------------------------------------- + # # Write back the block of the output matrix C with masks. + # Comment out the following lines to enable split the workload to two vector cores + SUB_BLK_M: tl.constexpr = BLOCK_SIZE_M // 2 + for s in extension.parallel(0, 2, bind_sub_block=True): + vec_sub_blk = extension.extract_slice(accumulator, (s * SUB_BLK_M, 0), (SUB_BLK_M, BLOCK_SIZE_N), (1, 1)) + if ACTIVATION == "leaky_relu_custom": + vec_sub_blk = leaky_relu_custom(vec_sub_blk) + c_sub_blk = vec_sub_blk.to(tl.float16) + # ----------------------------------------------------------- + # Write back the block of the output matrix C with masks. + offs_cm = pid_m * BLOCK_SIZE_M + s * SUB_BLK_M + tl.arange(0, SUB_BLK_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c_sub_blk, mask=c_mask) + + +# We can fuse `leaky_relu_custom` by providing it as an `ACTIVATION` meta-parameter in `matmul_kernel`. +@triton.jit +def leaky_relu_custom(x): + return tl.where(x >= 0, x, 0.01 * x) + 1.0 + + +def torch_matmul(a, b, activation=""): + c = torch.matmul(a, b) + if activation == "leaky_relu_custom": + c = torch.where(c >= 0, c, 0.01 * c) + 1.0 + return c + + +# %% +# We can now create a convenience wrapper function that only takes two input tensors, +# and (1) checks any shape constraint; (2) allocates the output; (3) launches the above kernel. + + +def matmul(a, b, activation=""): + # Check constraints. + assert a.shape[1] == b.shape[0], "Incompatible dimensions" + assert a.is_contiguous(), "Matrix A must be contiguous" + M, K = a.shape + K, N = b.shape + # Allocates output. + c = torch.empty((M, N), device=a.device, dtype=torch.float16) + + # 1D launch kernel where each block gets its own program. + + def grid(META): + return (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), ) + + matmul_kernel[grid]( + a, b, c, # + M, N, K, # + a.stride(0), a.stride(1), # + b.stride(0), b.stride(1), # + c.stride(0), c.stride(1), # + ACTIVATION=activation, # + ) + return c + + +# %% +# Unit Test +# --------- +# +# We can test our custom matrix multiplication operation against a native torch implementation (i.e., cuBLAS). +def test(): + activation = "leaky_relu_custom" + torch.manual_seed(0) + a = torch.randn((512, 512), device=DEV, dtype=torch.float16) + b = torch.randn((512, 512), device=DEV, dtype=torch.float16) + triton_output = matmul(a, b, activation) + torch_output = torch_matmul(a, b, activation) + print(f"triton_output_with_fp16_inputs={triton_output}") + print(f"torch_output_with_fp16_inputs={torch_output}") + torch.testing.assert_close(triton_output, torch_output, atol=1e-3, rtol=1e-3) + print("Passed") + + +if __name__ == "__main__": + test() diff --git a/third_party/ascend/tutorials/04-low-memory-dropout.py b/third_party/ascend/tutorials/04-low-memory-dropout.py new file mode 100644 index 0000000000..2c5570a0f4 --- /dev/null +++ b/third_party/ascend/tutorials/04-low-memory-dropout.py @@ -0,0 +1,137 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Low-Memory Dropout +================== +""" + +import tabulate +import torch +import torch_npu + +import triton +import triton.language as tl + +DEV = "npu" + + +@triton.jit +def _dropout( + x_ptr, # pointer to the input + x_keep_ptr, # pointer to a mask of 0s and 1s + output_ptr, # pointer to the output + n_elements, # number of elements in the `x` tensor + p, # probability that an element of `x` is changed to zero + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + # Load data + x = tl.load(x_ptr + offsets, mask=mask) + x_keep = tl.load(x_keep_ptr + offsets, mask=mask) + # The line below is the crucial part, described in the paragraph above! + output = tl.where(x_keep != 0, x / (1 - p), 0.0) + # Write-back output + tl.store(output_ptr + offsets, output, mask=mask) + + +def dropout(x, x_keep, p): + output = torch.empty_like(x) + assert x.is_contiguous() + n_elements = x.numel() + + def grid(meta): + return (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + + _dropout[grid](x, x_keep, output, n_elements, p, BLOCK_SIZE=1024) + return output + + +@triton.jit +def _seeded_dropout( + x_ptr, + output_ptr, + n_elements, + p, + seed, + BLOCK_SIZE: tl.constexpr, +): + # compute memory offsets of elements handled by this instance + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + # load data from x + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + # randomly prune it + random = tl.rand(seed, offsets) + x_keep = random > p + # write-back + output = tl.where(x_keep, x / (1 - p), 0.0) + tl.store(output_ptr + offsets, output, mask=mask) + + +def seeded_dropout(x, p, seed): + output = torch.empty_like(x) + assert x.is_contiguous() + n_elements = x.numel() + + def grid(meta): + return (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + + _seeded_dropout[grid](x, output, n_elements, p, seed, BLOCK_SIZE=1024) + return output + + +def test(): + # Input tensor + x = torch.randn(size=(10, ), device=DEV) + # Dropout mask + p = 0.5 + x_keep = (torch.rand(size=(10, ), device=DEV) > p).to(torch.int32) + # + output = dropout(x, x_keep=x_keep, p=p) + print(tabulate.tabulate([ + ["input"] + x.tolist(), + ["keep mask"] + x_keep.tolist(), + ["output"] + output.tolist(), + ])) + + x = torch.randn(size=(10, ), device=DEV) + # Compare this to the baseline - dropout mask is never instantiated! + output = seeded_dropout(x, p=0.5, seed=123) + output2 = seeded_dropout(x, p=0.5, seed=123) + output3 = seeded_dropout(x, p=0.5, seed=512) + + print( + tabulate.tabulate([ + ["input"] + x.tolist(), + ["output (seed = 123)"] + output.tolist(), + ["output (seed = 123)"] + output2.tolist(), + ["output (seed = 512)"] + output3.tolist(), + ])) + + +if __name__ == "__main__": + test() diff --git a/third_party/ascend/tutorials/05-layer-norm.py b/third_party/ascend/tutorials/05-layer-norm.py new file mode 100644 index 0000000000..8af05fc81f --- /dev/null +++ b/third_party/ascend/tutorials/05-layer-norm.py @@ -0,0 +1,126 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Layer Normalization +============= +""" + +import pytest +import torch +import triton +import triton.language as tl +import torch_npu + + +@triton.jit +def _layer_norm_fwd_fused( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride, # how much to increase the pointer when moving by 1 row + N, # number of columns in X + eps, # epsilon to avoid division by zero + BLOCK_SIZE: tl.constexpr, +): + # Map the program id to the row of X and Y it should compute. + row = tl.program_id(0) + Y += row * stride + X += row * stride + # Compute mean + mean = 0 + _mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + a = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32) + _mean += a + mean = tl.sum(_mean, axis=0) / N + # Compute variance + _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + x = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32) + x = tl.where(cols < N, x - mean, 0.) + _var += x * x + var = tl.sum(_var, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + # Write mean / rstd + tl.store(Mean + row, mean) + tl.store(Rstd + row, rstd) + # Normalize and apply linear transformation + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + mask = cols < N + w = tl.load(W + cols, mask=mask) + b = tl.load(B + cols, mask=mask) + x = tl.load(X + cols, mask=mask, other=0.).to(tl.float32) + x_hat = (x - mean) * rstd + y = x_hat * w + b + # Write output + tl.store(Y + cols, y, mask=mask) + + +@torch.inference_mode() +def layer_norm(x, normalized_shape, weight, bias, eps=1e-5): + # allocate output + y = torch.empty_like(x) + # reshape input data into 2D tensor + x_arg = x.reshape(-1, x.shape[-1]) + M, N = x_arg.shape + mean = torch.empty((M, ), dtype=torch.float32, device=x.device) + rstd = torch.empty((M, ), dtype=torch.float32, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + # heuristics for number of warps + num_warps = min(max(BLOCK_SIZE // 256, 1), 8) + # enqueue kernel + kernel = _layer_norm_fwd_fused[(M, )]( # + x_arg, y, weight, bias, mean, rstd, # + x_arg.stride(0), N, eps, # + BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_ctas=1) + return y + + +def _layer_norm(M, N, dtype, eps=1e-5, device='npu'): + # create data + x_shape = (M, N) + w_shape = (x_shape[-1], ) + weight = torch.rand(w_shape, dtype=dtype, device=device, requires_grad=True) + bias = torch.rand(w_shape, dtype=dtype, device=device, requires_grad=True) + x = -2.3 + 0.5 * torch.randn(x_shape, dtype=dtype, device=device) + dy = .1 * torch.randn_like(x) + x.requires_grad_(True) + # forward pass + y_tri = layer_norm(x, w_shape, weight, bias, eps) + y_ref = torch.nn.functional.layer_norm(x, w_shape, weight, bias, eps).to(dtype) + # compare + assert torch.allclose(y_tri, y_ref, atol=1e-2, rtol=0) + print(f"y_tri: {y_tri}") + print(f"y_ref: {y_ref}") + print(f"Layer Normalization {M},{N} {dtype} PASSED!") + + +if __name__ == "__main__": + _layer_norm(128, 128, torch.float16) + _layer_norm(128, 128, torch.bfloat16) + _layer_norm(128, 128, torch.float32) diff --git a/third_party/ascend/tutorials/06-fused-attention.py b/third_party/ascend/tutorials/06-fused-attention.py new file mode 100644 index 0000000000..3eed90789b --- /dev/null +++ b/third_party/ascend/tutorials/06-fused-attention.py @@ -0,0 +1,351 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Fused Attention +=============== + +This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao (https://tridao.me/publications/flash2/flash2.pdf) + +Credits: OpenAI kernel team + +Extra Credits: + +* Original flash attention paper (https://arxiv.org/abs/2205.14135) +* Rabe and Staats (https://arxiv.org/pdf/2112.05682v2.pdf) + +""" + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +DEVICE = "npu" + + +@triton.jit +def _attn_fwd_inner(acc_ptr, l_i, m_i, q, # Accumulator, local l, local m, query vector + K_block_ptr, V_block_ptr, # Key and value block pointers for current stage + start_m, qk_scale, # Starting position of current query block, qk scale factor + BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr, BLOCK_N: tl.constexpr, # Block size constants + STAGE: tl.constexpr, offs_m: tl.constexpr, + offs_n: tl.constexpr, # Current stage flag, m and n offset indices + N_CTX: tl.constexpr, + fp8_v: tl.constexpr): # Total context length, whether to enable FP8 for value precision + # Set the processing range [lo, hi) for the current stage (in column block units) + # Causal attention, as the name implies, restricts the flow of information during computation, + # only allowing the model to see the current and previous positions. + # In other words, the output at the current position can only depend on the input at or before this position, + # and cannot access information from future positions. + # Causal attention ensures sequential order and prevents "leakage of future information." + # But the following logic will also be triggered + if STAGE == 1: + # Stage 1: process all tokens before the query block + tl.static_assert(BLOCK_M >= BLOCK_N) + lo, hi = 0, start_m * BLOCK_M + elif STAGE == 2: + # Stage 2: process the current query block + tl.static_assert(BLOCK_M >= BLOCK_N) + lo, hi = start_m * BLOCK_M, (start_m + 1) * BLOCK_M + lo = tl.multiple_of(lo, BLOCK_M) # Align starting position + # causal = False (no need for masking) + else: + lo, hi = 0, N_CTX # Process the entire context + + # Adjust K and V block pointers to the starting position `lo` + K_block_ptr = tl.advance(K_block_ptr, (lo, 0)) # K is [HEAD_DIM, N_CTX], shift along the second dim by lo + V_block_ptr = tl.advance(V_block_ptr, (lo, 0)) # V is [N_CTX, HEAD_DIM], shift along the first dim by lo + + # Index mapping for the accumulator , used for slicing when HEAD_DIM >= 256 + row = tl.arange(0, BLOCK_M)[:, None] + col_head_dim = tl.arange(0, HEAD_DIM)[None, :] + block2d_acc = row * HEAD_DIM + col_head_dim + + # Iterate over all k, v blocks in the current stage and accumulate the output + for start_n in range(lo, hi, BLOCK_N): # Process BLOCK_N columns at a time + start_n = tl.multiple_of(start_n, BLOCK_N) # Align column start position + # -- Compute qk ---- + k = tl.load(K_block_ptr) + # Modify K + trans_k = tl.trans(k) + qk = tl.dot(q, trans_k) + # Apply causal mask for STAGE 2 + if STAGE == 2: + mask = offs_m[:, None] >= (start_n + offs_n[None, :]) # Construct upper triangular mask + qk = qk * qk_scale + tl.where(mask, 0, -1.0e6) # Set invalid positions to -∞ + m_ij = tl.maximum(m_i, tl.max(qk, 1)) # Update m_ij = max(m_i, max(qk)) + qk -= m_ij[:, None] # Subtract max for softmax stability + else: + qk = qk * qk_scale + m_ij = tl.maximum(m_i, tl.max(qk, 1)) # Scaled max + qk = qk - m_ij[:, None] # Stabilize + + # Softmax weights p = exp(qk) + p = tl.math.exp(qk) + + # Convert softmax weight type depending on FP8 usage + if fp8_v: + p_cast = p.to(tl.float8e5) # Convert to FP8 format (save memory) + else: + p_cast = p.to(k.dtype) + + v = tl.load(V_block_ptr) # Load corresponding V block + pv = tl.dot(p_cast, v) + l_ij = tl.sum(p, 1) # Softmax denominator (sum of each row) + # -- Update m_i and l_i + alpha = tl.math.exp(m_i - m_ij) # Update factor: exp difference between old and new max + l_i = l_i * alpha + l_ij # Update softmax denominator + # -- Update output accumulator -- + if HEAD_DIM < 256: + acc_ptr = acc_ptr * alpha[:, None] + acc_ptr = tl.dot(p_cast, v, acc_ptr) + else: + # 1. Load current slice of accumulator + acc = tl.load(acc_ptr + block2d_acc) + # 2. Update in slices (split by 1/4 of BLOCK_M to avoid ub overflow) + for i in range(4): + # Calculate start/end rows for current slice + offset = i * (BLOCK_M // 4) + # Extract slice data + acc_i = tl.extract_slice(acc, (offset, 0), (BLOCK_M // 4, HEAD_DIM), (1, 1)) + alpha_i = tl.extract_slice(alpha, [offset], [BLOCK_M // 4], [1]) + pv_i = tl.extract_slice(pv, (offset, 0), (BLOCK_M // 4, HEAD_DIM), (1, 1)) + # Incrementally update slice: acc = acc * alpha + pv + acc_i = acc_i * alpha_i[:, None] + pv_i + # Write updated slice back to accumulator + acc = tl.insert_slice(acc, acc_i, (offset, 0), (BLOCK_M // 4, HEAD_DIM), (1, 1)) + # 3. updated accumulator + tl.store(acc_ptr + block2d_acc, acc) + + m_i = m_ij # Update current block max + # Advance V and K block pointers to next BLOCK_N range + V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0)) + K_block_ptr = tl.advance(K_block_ptr, (BLOCK_N, 0)) + # Return accumulated output acc_ptr, softmax denominator l_i, and max value m_i + return acc_ptr, l_i, m_i + + +@triton.jit +def _attn_fwd(Q, K, V, M, Out, acc, sm_scale, stride_qz: tl.constexpr, stride_qh: tl.constexpr, stride_qm: tl.constexpr, + stride_qk: tl.constexpr, stride_kz: tl.constexpr, stride_kh: tl.constexpr, stride_kn: tl.constexpr, + stride_kk: tl.constexpr, stride_vz: tl.constexpr, stride_vh: tl.constexpr, stride_vn: tl.constexpr, + stride_vk: tl.constexpr, stride_oz: tl.constexpr, stride_oh: tl.constexpr, stride_om: tl.constexpr, + stride_on: tl.constexpr, Z: tl.constexpr, H: tl.constexpr, N_CTX: tl.constexpr, HEAD_DIM: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, STAGE: tl.constexpr): + # Total number of blocks in sequence dimension (M) + NUM_BLOCKS_M = N_CTX // BLOCK_M + # Total tasks = number of sequence blocks × batch size (Z) × number of attention heads (H) + NUM_BLOCKS = NUM_BLOCKS_M * Z * H + + # Current M-dimension block index + pid = tl.program_id(0) + + for block_idx in range(pid, NUM_BLOCKS, 20): + task_hz_idx = block_idx // NUM_BLOCKS_M + task_m_idx = block_idx % NUM_BLOCKS_M + off_z = task_hz_idx // H + off_h = task_hz_idx % H + qvk_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh + # Create block pointers for Q, K, V, Output + Q_block_ptr = tl.make_block_ptr( + base=Q + qvk_offset, + shape=(N_CTX, HEAD_DIM), + strides=(stride_qm, stride_qk), + offsets=(task_m_idx * BLOCK_M, 0), + block_shape=(BLOCK_M, HEAD_DIM), + order=(1, 0), + ) + V_block_ptr = tl.make_block_ptr( + base=V + qvk_offset, + shape=(N_CTX, HEAD_DIM), + strides=(stride_vn, stride_vk), + offsets=(0, 0), + block_shape=(BLOCK_N, HEAD_DIM), + order=(1, 0), + ) + K_block_ptr = tl.make_block_ptr( + base=K + qvk_offset, + shape=(N_CTX, HEAD_DIM), + strides=(stride_kn, stride_kk), + offsets=(0, 0), + block_shape=(BLOCK_N, HEAD_DIM), + order=(1, 0), + ) + O_block_ptr = tl.make_block_ptr( + base=Out + qvk_offset, + shape=(N_CTX, HEAD_DIM), + strides=(stride_om, stride_on), + offsets=(task_m_idx * BLOCK_M, 0), + block_shape=(BLOCK_M, HEAD_DIM), + order=(1, 0), + ) + # Initialize offsets + offs_m = task_m_idx * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0 + + # Initialize accumulator + if HEAD_DIM < 256: + acc_ptr = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) + else: + acc_offset = (off_z.to(tl.int64) * stride_qz // stride_qm * HEAD_DIM + + off_h.to(tl.int64) * stride_qh // stride_qm * HEAD_DIM + task_m_idx * BLOCK_M * HEAD_DIM) + acc_ptr = acc + acc_offset + + # load q: it will stay in SRAM throughout + q = tl.load(Q_block_ptr) + + # stage 1: off-band + # For causal = True, STAGE = 3 and _attn_fwd_inner gets 1 as its STAGE + # For causal = False, STAGE = 1, and _attn_fwd_inner gets 3 as its STAGE + if STAGE & 1: + acc_ptr, l_i, m_i = _attn_fwd_inner(acc_ptr, l_i, m_i, q, K_block_ptr, V_block_ptr, # + task_m_idx, sm_scale, # + BLOCK_M, HEAD_DIM, BLOCK_N, # + 4 - STAGE, offs_m, offs_n, N_CTX, V.dtype.element_ty == tl.float8e5 # + ) + # stage 2: on-band + if STAGE & 2: + # barrier makes it easier for compielr to schedule the + # two loops independently + acc_ptr, l_i, m_i = _attn_fwd_inner(acc_ptr, l_i, m_i, q, K_block_ptr, V_block_ptr, # + task_m_idx, sm_scale, # + BLOCK_M, HEAD_DIM, BLOCK_N, # + 2, offs_m, offs_n, N_CTX, V.dtype.element_ty == tl.float8e5 # + ) + + m_i += tl.math.log(l_i) + if HEAD_DIM < 256: + accumulator = acc_ptr / l_i[:, None] + else: + row = tl.arange(0, BLOCK_M)[:, None] + col_head_dim = tl.arange(0, HEAD_DIM)[None, :] + block2d_acc = row * HEAD_DIM + col_head_dim + accumulator = tl.load(acc_ptr + block2d_acc) + accumulator = accumulator / l_i[:, None] + + m_ptrs = M + task_hz_idx * N_CTX + offs_m + + tl.store(m_ptrs, m_i) + tl.store(O_block_ptr, accumulator.to(Out.type.element_ty)) + + +class _attention(torch.autograd.Function): + + @staticmethod + def forward(ctx, q, k, v, causal, sm_scale, BM, BN): + """ + Forward computation interface: + Args: + ctx: Context object + q: Query tensor (Q), shape [Z, H, N_CTX, HEAD_DIM] + k: Key tensor (K), shape [Z, H, N_CTX, HEAD_DIM] + v: Value tensor (V), shape [Z, H, N_CTX, HEAD_DIM] + causal: Whether to enable causal attention + sm_scale: Scaling factor for QK product + BM: Q block size (BLOCK_M) + BN: K/V block size (BLOCK_N) + Returns: + out: Attention output tensor, shape [Z, H, N_CTX, HEAD_DIM] + """ + # shape constraints + HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1] + # when v is in float8_e5m2 it is transposed. + HEAD_DIM_V = v.shape[-1] + assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V + assert HEAD_DIM_K in {16, 32, 64, 128, 256} + + out = torch.empty_like(q) + stage = 3 if causal else 1 + extra_kern_args = {} + + # Number of NPU cores (adjust based on hardware) + num_cores = 20 + acc = torch.zeros((q.shape[0], q.shape[1], q.shape[2], HEAD_DIM_K), dtype=torch.float32, device=q.device) + M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32) + + _attn_fwd[(num_cores, )](q, k, v, M, out, acc, sm_scale, q.stride(0), q.stride(1), q.stride(2), q.stride(3), + k.stride(0), k.stride(1), k.stride(2), k.stride(3), v.stride(0), v.stride(1), + v.stride(2), v.stride(3), out.stride(0), out.stride(1), out.stride(2), out.stride(3), + q.shape[0], q.shape[1], N_CTX=q.shape[2], HEAD_DIM=HEAD_DIM_K, BLOCK_M=BM, BLOCK_N=BN, + STAGE=stage, **extra_kern_args) + + ctx.save_for_backward(q, k, v, out, M) + ctx.sm_scale = sm_scale + ctx.HEAD_DIM = HEAD_DIM_K + ctx.causal = causal + return out + + +attention = _attention.apply + + +@pytest.mark.parametrize("Z, H, N_CTX, HEAD_DIM, causal, dtype, BM, BN", [ + (1, 1, 128, 128, False, torch.float16, 32, 128), + (1, 1, 128, 128, False, torch.bfloat16, 64, 128), + (1, 2, 256, 256, False, torch.bfloat16, 32, 256), + (2, 2, 128, 256, False, torch.float16, 64, 128), + (4, 32, 64, 64, False, torch.float16, 32, 64), + (4, 32, 1024, 64, False, torch.bfloat16, 64, 128), + (4, 32, 4096, 64, False, torch.float16, 128, 128), +]) +def test_op(Z, H, N_CTX, HEAD_DIM, causal, dtype, BM, BN): + # Filter out non-integer cases; N_CTX must be divisible by BM and BN, and HEAD_DIM must be divisible by 16. + if N_CTX % BM != 0 or N_CTX % BN != 0 or HEAD_DIM % 16 != 0: + pytest.skip("Skipping non-divisible case") + + torch.manual_seed(20) + q = (torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device=DEVICE).normal_(mean=0.0, std=0.5).requires_grad_()) + k = (torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device=DEVICE).normal_(mean=0.0, std=0.5).requires_grad_()) + v = (torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device=DEVICE).normal_(mean=0.0, std=0.5).requires_grad_()) + + sm_scale = 0.5 + + tri_out = attention(q, k, v, causal, sm_scale, BM, BN) + ref_out = torch_npu.npu_fusion_attention( + q, + k, + v, + H, + padding_mask=None, + atten_mask=None, + scale=sm_scale, + keep_prob=1.0, + input_layout="BNSD", + pre_tockens=65535, + next_tockens=65535, + sparse_mode=0, + )[0] + + torch.testing.assert_close(ref_out, tri_out, atol=1e-2, rtol=1e-2, equal_nan=True) + print(f"[PASSED] Attention shape:({Z}, {H}, {N_CTX}, {HEAD_DIM}), BM: {BM}, BN: {BN}, dtype: {dtype}") + + +if __name__ == "__main__": + test_op(1, 1, 128, 128, causal=False, dtype=torch.float16, BM=32, BN=128) + test_op(1, 1, 128, 128, causal=False, dtype=torch.bfloat16, BM=64, BN=128) + test_op(1, 2, 256, 256, causal=False, dtype=torch.bfloat16, BM=32, BN=256) + test_op(2, 2, 128, 256, causal=False, dtype=torch.float16, BM=64, BN=128) + test_op(4, 32, 64, 64, causal=False, dtype=torch.float16, BM=32, BN=64) + test_op(4, 32, 1024, 64, causal=False, dtype=torch.bfloat16, BM=64, BN=128) + test_op(4, 32, 4096, 64, causal=False, dtype=torch.float16, BM=128, BN=128) diff --git a/third_party/ascend/tutorials/07-extern-functions.py b/third_party/ascend/tutorials/07-extern-functions.py new file mode 100644 index 0000000000..5748a57bc5 --- /dev/null +++ b/third_party/ascend/tutorials/07-extern-functions.py @@ -0,0 +1,89 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Libdevice (`tl.extra.libdevice`) function +============================== +""" +import inspect +import os +from pathlib import Path + +import torch +import torch_npu + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +from triton.backends.ascend.compiler import get_libdevice + +DEV = "npu" + + +@triton.jit +def asin_kernel( + x_ptr, + y_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + x = libdevice.asin(x) + tl.store(y_ptr + offsets, x, mask=mask) + + +def test(): + torch.manual_seed(0) + size = 98432 + x = torch.rand(size, device=DEV) + output_triton = torch.zeros(size, device=DEV) + output_torch = torch.asin(x) + assert x.device.type == DEV and output_triton.device.type == DEV + n_elements = output_torch.numel() + + def grid(meta): + return (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + + asin_kernel[grid](x, output_triton, n_elements, BLOCK_SIZE=1024) + print(output_torch) + print(output_triton) + print(f'The maximum difference between torch and triton is ' + f'{torch.max(torch.abs(output_torch - output_triton))}') + + current_file = inspect.getfile(inspect.currentframe()) + current_dir = Path(os.path.dirname(os.path.abspath(current_file))) + extern_libs = {'libdevice': get_libdevice()} + + output_triton = torch.empty_like(x) + asin_kernel[grid](x, output_triton, n_elements, BLOCK_SIZE=1024, extern_libs=extern_libs) + torch.testing.assert_close(output_torch, output_triton, rtol=1e-4, atol=1e-4) + print(output_torch) + print(output_triton) + print(f'The maximum difference between torch and triton is ' + f'{torch.max(torch.abs(output_torch - output_triton))}') + + +if __name__ == "__main__": + test() diff --git a/third_party/ascend/tutorials/08-grouped-gemm.py b/third_party/ascend/tutorials/08-grouped-gemm.py new file mode 100644 index 0000000000..96739be81d --- /dev/null +++ b/third_party/ascend/tutorials/08-grouped-gemm.py @@ -0,0 +1,281 @@ +# Copyright (c) 2023 NVIDIA Corporation & Affiliates. All rights reserved. +# Copyright (c) Huawei Technologies Co., Ltd. 2025. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Group GEMM +============================ +""" + +import torch +import torch_npu + +import triton +import triton.language as tl +import triton.runtime.driver as driver + +DEV = "npu" + + +def get_npu_properties(): + device = torch.npu.current_device() + return driver.active.utils.get_device_properties(device) + + +NUM_CORES = get_npu_properties()["num_aicore"] + + +@triton.autotune( + configs=[ + triton.Config({ + 'BLOCK_SIZE_M': 128, + 'BLOCK_SIZE_N': 128, + 'BLOCK_SIZE_K': 32, + 'NUM_SM': NUM_CORES, + }), + triton.Config({ + 'BLOCK_SIZE_M': 128, + 'BLOCK_SIZE_N': 128, + 'BLOCK_SIZE_K': 32, + 'NUM_SM': NUM_CORES, + }), + triton.Config({ + 'BLOCK_SIZE_M': 64, + 'BLOCK_SIZE_N': 64, + 'BLOCK_SIZE_K': 32, + 'NUM_SM': NUM_CORES, + }), + triton.Config({ + 'BLOCK_SIZE_M': 64, + 'BLOCK_SIZE_N': 64, + 'BLOCK_SIZE_K': 32, + 'NUM_SM': NUM_CORES, + }), + ], + key=['group_size'], +) +@triton.jit +def grouped_matmul_kernel( + group_a_ptrs, + group_b_ptrs, + group_c_ptrs, + group_gemm_sizes, + g_lds, + group_size, + NUM_SM: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, +): + tile_idx = tl.program_id(0) + last_problem_end = 0 + for g in range(group_size): + gm = tl.load(group_gemm_sizes + g * 3) + gn = tl.load(group_gemm_sizes + g * 3 + 1) + gk = tl.load(group_gemm_sizes + g * 3 + 2) + num_m_tiles = tl.cdiv(gm, BLOCK_SIZE_M) + num_n_tiles = tl.cdiv(gn, BLOCK_SIZE_N) + num_tiles = num_m_tiles * num_n_tiles + while (tile_idx >= last_problem_end and tile_idx < last_problem_end + num_tiles): + k = gk + lda = tl.load(g_lds + g * 3) + ldb = tl.load(g_lds + g * 3 + 1) + ldc = tl.load(g_lds + g * 3 + 2) + a_ptr = tl.load(group_a_ptrs + g).to(tl.pointer_type(tl.float16)) + b_ptr = tl.load(group_b_ptrs + g).to(tl.pointer_type(tl.float16)) + c_ptr = tl.load(group_c_ptrs + g).to(tl.pointer_type(tl.float16)) + tile_idx_in_gemm = tile_idx - last_problem_end + tile_m_idx = tile_idx_in_gemm // num_n_tiles + tile_n_idx = tile_idx_in_gemm % num_n_tiles + + offs_am = tile_m_idx * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_bn = tile_n_idx * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + offs_am[:, None] * lda + offs_k[None, :] + b_ptrs = b_ptr + offs_k[:, None] * ldb + offs_bn[None, :] + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for _ in range(0, tl.cdiv(k, BLOCK_SIZE_K)): + tl.multiple_of(a_ptrs, [16, 16]) + tl.multiple_of(b_ptrs, [16, 16]) + a = tl.load(a_ptrs) + b = tl.load(b_ptrs) + accumulator += tl.dot(a, b) + a_ptrs += BLOCK_SIZE_K + b_ptrs += BLOCK_SIZE_K * ldb + c = accumulator.to(tl.float16) + + offs_cm = tile_m_idx * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = tile_n_idx * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + ldc * offs_cm[:, None] + offs_cn[None, :] + + tl.store(c_ptrs, c) + tile_idx += NUM_SM + + last_problem_end = last_problem_end + num_tiles + + +def group_gemm_fn(group_A, group_B): + device = torch.device(DEV) + assert len(group_A) == len(group_B) + group_size = len(group_A) + + A_addrs = [] + B_addrs = [] + C_addrs = [] + g_sizes = [] + g_lds = [] + group_C = [] + for i in range(group_size): + A = group_A[i] + B = group_B[i] + assert A.shape[1] == B.shape[0] + M, K = A.shape + K, N = B.shape + C = torch.empty((M, N), device=device, dtype=A.dtype) + group_C.append(C) + A_addrs.append(A.data_ptr()) + B_addrs.append(B.data_ptr()) + C_addrs.append(C.data_ptr()) + g_sizes += [M, N, K] + g_lds += [A.stride(0), B.stride(0), C.stride(0)] + + d_a_ptrs = torch.tensor(A_addrs, device=device) + d_b_ptrs = torch.tensor(B_addrs, device=device) + d_c_ptrs = torch.tensor(C_addrs, device=device) + d_g_sizes = torch.tensor(g_sizes, dtype=torch.int32, device=device) + d_g_lds = torch.tensor(g_lds, dtype=torch.int32, device=device) + + def grid(meta): + return (meta['NUM_SM'], ) + + grouped_matmul_kernel[grid]( + d_a_ptrs, + d_b_ptrs, + d_c_ptrs, + d_g_sizes, + d_g_lds, + group_size, + ) + + return group_C + + +def test(): + group_m = [1024, 512, 256, 128] + group_n = [1024, 512, 256, 128] + group_k = [1024, 512, 256, 128] + group_A = [] + group_B = [] + assert len(group_m) == len(group_n) + assert len(group_n) == len(group_k) + group_size = len(group_m) + for i in range(group_size): + M = group_m[i] + N = group_n[i] + K = group_k[i] + A = torch.rand((M, K), device=DEV, dtype=torch.float16) + B = torch.rand((K, N), device=DEV, dtype=torch.float16) + group_A.append(A) + group_B.append(B) + + tri_out = group_gemm_fn(group_A, group_B) + ref_out = [torch.matmul(a, b) for a, b in zip(group_A, group_B)] + for i in range(group_size): + torch.testing.assert_close(ref_out[i], tri_out[i], atol=1e-2, rtol=1e-3) + print("Passed") + + +def triton_perf_fn(a_ptrs, b_ptrs, c_ptrs, sizes, lds, group_size): + + def grid(meta): + return (meta['NUM_SM'], ) + + grouped_matmul_kernel[grid]( + a_ptrs, + b_ptrs, + c_ptrs, + sizes, + lds, + group_size, + ) + + +def torch_perf_fn(group_A, group_B): + for a, b in zip(group_A, group_B): + torch.matmul(a, b) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=['N'], + x_vals=[2**i for i in range(7, 11)], + line_arg='provider', + line_vals=['torch', 'triton'], + line_names=["Torch", "Triton"], + styles=[('green', '-'), ('blue', '-')], + ylabel="runtime(ms)", + plot_name="group-gemm-performance", + args={}, + )) +def benchmark(N, provider): + group_size = 4 + group_A = [] + group_B = [] + A_addrs = [] + B_addrs = [] + C_addrs = [] + g_sizes = [] + g_lds = [] + group_C = [] + for _ in range(group_size): + A = torch.rand((N, N), device=DEV, dtype=torch.float16) + B = torch.rand((N, N), device=DEV, dtype=torch.float16) + C = torch.empty((N, N), device=DEV, dtype=torch.float16) + group_A.append(A) + group_B.append(B) + group_C.append(C) + A_addrs.append(A.data_ptr()) + B_addrs.append(B.data_ptr()) + C_addrs.append(C.data_ptr()) + g_sizes += [N, N, N] + g_lds += [N, N, N] + + d_a_ptrs = torch.tensor(A_addrs, device=DEV) + d_b_ptrs = torch.tensor(B_addrs, device=DEV) + d_c_ptrs = torch.tensor(C_addrs, device=DEV) + d_g_sizes = torch.tensor(g_sizes, dtype=torch.int32, device=DEV) + d_g_lds = torch.tensor(g_lds, dtype=torch.int32, device=DEV) + + quantiles = [0.5, 0.2, 0.8] + + def bench_torch(): + torch_perf_fn(group_A, group_B) + + def bench_triton(): + triton_perf_fn(d_a_ptrs, d_b_ptrs, d_c_ptrs, d_g_sizes, d_g_lds, group_size) + + if provider == 'torch': + ms, min_ms, max_ms = triton.testing.do_bench(bench_torch, quantiles=quantiles) + if provider == 'triton': + ms, min_ms, max_ms = triton.testing.do_bench(bench_triton, quantiles=quantiles) + return ms, max_ms, min_ms + + +if __name__ == "__main__": + test() diff --git a/third_party/ascend/tutorials/09-persistent-matmul.py b/third_party/ascend/tutorials/09-persistent-matmul.py new file mode 100644 index 0000000000..ef5bb8e8f7 --- /dev/null +++ b/third_party/ascend/tutorials/09-persistent-matmul.py @@ -0,0 +1,334 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Persistent Matmul +===================== +""" + +import argparse +import time + +import torch +import torch_npu +import triton +import triton.language as tl +import triton.runtime.driver as driver + +DEV = "npu" +DTYPE = torch.float16 + + +def get_npu_properties(): + device = torch.npu.current_device() + return driver.active.utils.get_device_properties(device) + + +def get_num_compute_cores(): + return get_npu_properties()["num_aicore"] + + +def _matmul_launch_metadata(grid, kernel, args): + ret = {} + M, N, K = args["M"], args["N"], args["K"] + bytes_per_elem = args["c_ptr"].element_size() + ret["name"] = f"{kernel.name} [M={M}, N={N}, K={K}]" + ret[f"flops{bytes_per_elem * 8}"] = 2.0 * M * N * K + ret["bytes"] = bytes_per_elem * (M * K + N * K + M * N) + return ret + + +@triton.jit(launch_metadata=_matmul_launch_metadata) +def matmul_kernel( + a_ptr, + b_ptr, + c_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + start_m = pid_m * BLOCK_SIZE_M + start_n = pid_n * BLOCK_SIZE_N + + offs_am = start_m + tl.arange(0, BLOCK_SIZE_M) + offs_bn = start_n + tl.arange(0, BLOCK_SIZE_N) + offs_am = tl.where(offs_am < M, offs_am, 0) + offs_bn = tl.where(offs_bn < N, offs_bn, 0) + + offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M) + offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + accumulator = tl.dot(a, b, accumulator) + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + c = accumulator.to(tl.float16) + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +@triton.jit(launch_metadata=_matmul_launch_metadata) +def matmul_kernel_persistent( + a_ptr, + b_ptr, + c_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_SMS: tl.constexpr, +): + start_pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + k_tiles = tl.cdiv(K, BLOCK_SIZE_K) + num_tiles = num_pid_m * num_pid_n + + tiles_per_sm = num_tiles // NUM_SMS + if start_pid < num_tiles % NUM_SMS: + tiles_per_sm += 1 + + tile_id = start_pid - NUM_SMS + ki = -1 + offs_k_for_mask = tl.arange(0, BLOCK_SIZE_K) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + + pid_m = 0 + pid_n = 0 + offs_am = tl.arange(0, BLOCK_SIZE_M) + offs_bn = tl.arange(0, BLOCK_SIZE_N) + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + for _ in range(0, k_tiles * tiles_per_sm): + ki = tl.where(ki == k_tiles - 1, 0, ki + 1) + if ki == 0: + tile_id += NUM_SMS + group_id = tile_id // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (tile_id % group_size_m) + pid_n = (tile_id % num_pid_in_group) // group_size_m + + start_m = pid_m * BLOCK_SIZE_M + start_n = pid_n * BLOCK_SIZE_N + offs_am = start_m + tl.arange(0, BLOCK_SIZE_M) + offs_bn = start_n + tl.arange(0, BLOCK_SIZE_N) + offs_am = tl.where(offs_am < M, offs_am, 0) + offs_bn = tl.where(offs_bn < N, offs_bn, 0) + offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M) + offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N) + + offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + a = tl.load(a_ptrs, mask=offs_k_for_mask[None, :] < K - ki * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k_for_mask[:, None] < K - ki * BLOCK_SIZE_K, other=0.0) + accumulator = tl.dot(a, b, accumulator) + + if ki == k_tiles - 1: + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + c = accumulator.to(tl.float16) + tl.store(c_ptrs, c, mask=c_mask) + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + +def get_configs(dtype): + return { + torch.float16: { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, # Origin is 256, ub overflow + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_stages": 3, + "num_warps": 8, + } + }[dtype] + + +def matmul(a, b): + assert a.shape[1] == b.shape[0], "Incompatible dimensions" + assert a.dtype == b.dtype, "Incompatible dtypes" + M, K = a.shape + _, N = b.shape + configs = get_configs(a.dtype) + c = torch.empty((M, N), device=a.device, dtype=a.dtype) + + def grid(meta): + return (triton.cdiv(M, meta["BLOCK_SIZE_M"]) * triton.cdiv(N, meta["BLOCK_SIZE_N"]), ) + + matmul_kernel[grid]( + a, + b, + c, + M, + N, + K, + a.stride(0), + a.stride(1), + b.stride(0), + b.stride(1), + c.stride(0), + c.stride(1), + BLOCK_SIZE_M=configs["BLOCK_SIZE_M"], + BLOCK_SIZE_N=configs["BLOCK_SIZE_N"], + BLOCK_SIZE_K=configs["BLOCK_SIZE_K"], + GROUP_SIZE_M=configs["GROUP_SIZE_M"], + num_stages=configs["num_stages"], + num_warps=configs["num_warps"], + ) + return c + + +def matmul_persistent(a, b): + assert a.shape[1] == b.shape[0], "Incompatible dimensions" + assert a.dtype == b.dtype, "Incompatible dtypes" + num_sms = get_num_compute_cores() + M, K = a.shape + _, N = b.shape + configs = get_configs(a.dtype) + c = torch.empty((M, N), device=a.device, dtype=a.dtype) + + def grid(meta): + return (min(num_sms, triton.cdiv(M, meta["BLOCK_SIZE_M"]) * triton.cdiv(N, meta["BLOCK_SIZE_N"])), ) + + matmul_kernel_persistent[grid]( + a, + b, + c, + M, + N, + K, + a.stride(0), + a.stride(1), + b.stride(0), + b.stride(1), + c.stride(0), + c.stride(1), + BLOCK_SIZE_M=configs["BLOCK_SIZE_M"], + BLOCK_SIZE_N=configs["BLOCK_SIZE_N"], + BLOCK_SIZE_K=configs["BLOCK_SIZE_K"], + GROUP_SIZE_M=configs["GROUP_SIZE_M"], + NUM_SMS=num_sms, + num_stages=configs["num_stages"], + num_warps=configs["num_warps"], + ) + return c + + +def torch_matmul(a, b): + return torch.matmul(a, b) + + +def bench(K, reps=10): + M = 8192 + N = 8192 + a = torch.randn((M, K), device=DEV, dtype=DTYPE) + b = torch.randn((K, N), device=DEV, dtype=DTYPE) + + for _ in range(reps): + _ = torch_matmul(a, b) + time.sleep(0.01) + for _ in range(reps): + _ = matmul(a, b) + time.sleep(0.01) + for _ in range(reps): + _ = matmul_persistent(a, b) + time.sleep(0.01) + + +def validate(M, N, K): + a = torch.randn((M, K), device=DEV, dtype=DTYPE) + b = torch.randn((K, N), device=DEV, dtype=DTYPE) + + torch_result = torch_matmul(a, b) + naive_result = matmul(a, b) + persistent_result = matmul_persistent(a, b) + + naive_vs_torch = "✅" if torch.allclose(naive_result, torch_result, atol=1.0) else "❌" + persistent_vs_torch = "✅" if torch.allclose(persistent_result, torch_result, atol=1.0) else "❌" + naive_vs_persistent = "✅" if torch.allclose(naive_result, persistent_result, atol=1.0) else "❌" + + print(f"M={M}, N={N}, K={K} verification naive vs torch: {naive_vs_torch} " + f"persistent vs torch: {persistent_vs_torch} naive vs persistent: {naive_vs_persistent}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("-K", type=int, required=False, default=512) + parser.add_argument("--K_range", type=int, nargs=2) + parser.add_argument("--K_step", type=int, default=512) + args = parser.parse_args() + + if args.K and args.K_range is None: + args.K_range = [args.K, args.K] + args.K_step = 1 + + torch.manual_seed(0) + + validate(32, 32, 32) + validate(8192, 8192, 512) + + for K in range(args.K_range[0], args.K_range[1] + 1, args.K_step): + bench(K) diff --git a/third_party/ascend/unittest/CMakeLists.txt b/third_party/ascend/unittest/CMakeLists.txt new file mode 100644 index 0000000000..8ca3dc6877 --- /dev/null +++ b/third_party/ascend/unittest/CMakeLists.txt @@ -0,0 +1,27 @@ +llvm_canonicalize_cmake_booleans( + MLIR_ENABLE_BINDINGS_PYTHON +) + +configure_lit_site_cfg( + ${CMAKE_CURRENT_SOURCE_DIR}/lit.site.cfg.py.in + ${CMAKE_CURRENT_BINARY_DIR}/lit.site.cfg.py + MAIN_CONFIG + ${CMAKE_CURRENT_SOURCE_DIR}/lit.cfg.py +) + +set(TRITON_TEST_DEPENDS + triton-opt +) + +set(FILECHECK_PATH "${LLVM_LIBRARY_DIR}/../bin/FileCheck") +set(LIT_ARGS "-Dfilecheck=${FILECHECK_PATH}") + +add_lit_testsuite(check-triton-ascend-lit-tests "Running the triton regression tests" + ${CMAKE_CURRENT_BINARY_DIR} + ARGS ${LIT_ARGS} + DEPENDS ${TRITON_TEST_DEPENDS} + ) + +set_target_properties(check-triton-ascend-lit-tests PROPERTIES FOLDER "Tests") + +add_lit_testsuites(TRITON-LIT-TESTS ${CMAKE_CURRENT_SOURCE_DIR} DEPENDS ${TRITON_TEST_DEPENDS}) diff --git a/third_party/ascend/unittest/Conversion/General/AddDynamicCVPipeline/AddMultiCache/Inner-scope.mlir b/third_party/ascend/unittest/Conversion/General/AddDynamicCVPipeline/AddMultiCache/Inner-scope.mlir new file mode 100644 index 0000000000..a822fe0b9f --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/AddDynamicCVPipeline/AddMultiCache/Inner-scope.mlir @@ -0,0 +1,573 @@ +// RUN: triton-opt --add_multi_buffer_inner_scope %s | FileCheck %s + +// Consolidated Unit Tests for AddMultiBufferInnerScope Pass +// Tests buffer creation, producer logic, consumer logic with ping-pong select chain + +module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} { + + //===--------------------------------------------------------------------===// + // T1: Double Buffer Creation Test + // Test: cross-block tensor with block_id 5→6 triggers buffer creation + // Key Check: memref.alloc creates ping buffer, scf.if selects buffer + // : producer tag [groupId, 1] on memref.memory_space_cast + // : consumer tag [groupId, 0] on scf.if + //===--------------------------------------------------------------------===// + // CHECK-LABEL: func.func @test_t1_double_buffer + // CHECK-DAG: memref.alloc + // CHECK-DAG: memref.memory_space_cast + // CHECK-DAG: bufferization.to_tensor + // CHECK-DAG: hivm.hir.copy + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 6 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 6 + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 5 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 5 + func.func @test_t1_double_buffer() { + %c0_i32 = arith.constant 0 : i32 + %c100_i32 = arith.constant 100 : i32 + %c1_i32 = arith.constant 1 : i32 + %cst = arith.constant 1.0 : f32 + %empty = tensor.empty() : tensor<128xf32> + scope.scope : () -> () { + %prod = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst : f32) outs(%empty : tensor<128xf32>) -> tensor<128xf32> + %loop_result = scf.for %i = %c0_i32 to %c100_i32 step %c1_i32 iter_args(%arg = %prod) -> (tensor<128xf32>) : i32 { + %consumed = arith.addf %arg, %arg {ssbuffer.block_id = 6 : i32} : tensor<128xf32> + %new_prod = arith.addf %consumed, %consumed {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + scf.yield %new_prod : tensor<128xf32> + } {ssbuffer.main_loop = 1 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return + } + + //===--------------------------------------------------------------------===// + // T2: Multiple Block IDs Test + // Test: block_id 5, 6, 7 all present in cross-block deps + // Key Check: multiple buffers created for different block_ids + //===--------------------------------------------------------------------===// + // CHECK-LABEL: func.func @test_t2_multiple_blocks + // CHECK-DAG: memref.alloc + // CHECK-DAG: memref.alloc + // CHECK-DAG: memref.alloc + // CHECK-DAG: memref.alloc + // CHECK-DAG: linalg.fill {ssbuffer.block_id = 5 : i32} + // CHECK-DAG: linalg.fill {ssbuffer.block_id = 7 : i32} + func.func @test_t2_multiple_blocks() { + %c0_i32 = arith.constant 0 : i32 + %c100_i32 = arith.constant 100 : i32 + %c1_i32 = arith.constant 1 : i32 + %cst = arith.constant 1.0 : f32 + %empty = tensor.empty() : tensor<64xf32> + scope.scope : () -> () { + %prod5 = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst : f32) outs(%empty : tensor<64xf32>) -> tensor<64xf32> + %prod7 = linalg.fill {ssbuffer.block_id = 7 : i32} ins(%cst : f32) outs(%empty : tensor<64xf32>) -> tensor<64xf32> + %loop_result:2 = scf.for %i = %c0_i32 to %c100_i32 step %c1_i32 iter_args(%arg5 = %prod5, %arg7 = %prod7) -> (tensor<64xf32>, tensor<64xf32>) : i32 { + %cons6 = arith.addf %arg5, %arg5 {ssbuffer.block_id = 6 : i32} : tensor<64xf32> + %cons8 = arith.mulf %arg7, %arg7 {ssbuffer.block_id = 8 : i32} : tensor<64xf32> + %new_prod5 = arith.addf %cons6, %cons6 {ssbuffer.block_id = 5 : i32} : tensor<64xf32> + %new_prod7 = arith.mulf %cons8, %cons8 {ssbuffer.block_id = 7 : i32} : tensor<64xf32> + scf.yield %new_prod5, %new_prod7 : tensor<64xf32>, tensor<64xf32> + } {ssbuffer.main_loop = 1 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return + } + + //===--------------------------------------------------------------------===// + // T3: 2D Tensor Double Buffer + // Test: 2D tensor with cross-block dependency + // Key Check: 2D memref buffer created with correct shape + //===--------------------------------------------------------------------===// + // CHECK-LABEL: func.func @test_t3_2d_tensor + // CHECK: memref.alloc + // CHECK: tensor<128x64xf32> + func.func @test_t3_2d_tensor() { + %c0_i32 = arith.constant 0 : i32 + %c100_i32 = arith.constant 100 : i32 + %c1_i32 = arith.constant 1 : i32 + %cst = arith.constant 1.0 : f32 + %empty = tensor.empty() : tensor<128x64xf32> + scope.scope : () -> () { + %prod = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst : f32) outs(%empty : tensor<128x64xf32>) -> tensor<128x64xf32> + %loop_result = scf.for %i = %c0_i32 to %c100_i32 step %c1_i32 iter_args(%arg = %prod) -> (tensor<128x64xf32>) : i32 { + %consumed = arith.mulf %arg, %arg {ssbuffer.block_id = 6 : i32} : tensor<128x64xf32> + %new_prod = arith.addf %consumed, %consumed {ssbuffer.block_id = 5 : i32} : tensor<128x64xf32> + scf.yield %new_prod : tensor<128x64xf32> + } {ssbuffer.main_loop = 1 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return + } + + //===--------------------------------------------------------------------===// + // T4: getIterCount - Loop iteration count calculation + // Test: lb=0, step=1, iv=i32 → iterCount = i + // Key Check: main_loop preserved, no buffer if no cross-block dep + //===--------------------------------------------------------------------===// + // CHECK-LABEL: func.func @test_t4_getitercount + // CHECK: scope.scope + // CHECK: linalg.fill + // CHECK: scf.for + // CHECK: {ssbuffer.main_loop = 1 : i64} + func.func @test_t4_getitercount() { + %c0_i32 = arith.constant 0 : i32 + %c100_i32 = arith.constant 100 : i32 + %c1_i32 = arith.constant 1 : i32 + %cst = arith.constant 1.0 : f32 + %empty = tensor.empty() : tensor<128xf32> + scope.scope : () -> () { + %prod = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst : f32) outs(%empty : tensor<128xf32>) -> tensor<128xf32> + scf.for %i = %c0_i32 to %c100_i32 step %c1_i32 iter_args(%arg = %prod) -> (tensor<128xf32>) : i32 { + %r = arith.addf %arg, %arg {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + scf.yield %r : tensor<128xf32> + } {ssbuffer.main_loop = 1 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return + } + + //===--------------------------------------------------------------------===// + // T5: findMainLoopInScope - Find main loop in scope + // Test: forOp has main_loop attribute → returns that forOp + // Key Check: main_loop attribute preserved + //===--------------------------------------------------------------------===// + // CHECK-LABEL: func.func @test_t5_findmainloop + // CHECK: linalg.fill {ssbuffer.block_id = 5 : i32} + // CHECK: scf.for + // CHECK: {ssbuffer.main_loop = 1 : i64} + func.func @test_t5_findmainloop() { + %c0_i32 = arith.constant 0 : i32 + %c100_i32 = arith.constant 100 : i32 + %c1_i32 = arith.constant 1 : i32 + %cst = arith.constant 1.0 : f32 + %empty = tensor.empty() : tensor<128xf32> + scope.scope : () -> () { + %prod = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst : f32) outs(%empty : tensor<128xf32>) -> tensor<128xf32> + scf.for %i = %c0_i32 to %c100_i32 step %c1_i32 iter_args(%arg = %prod) -> (tensor<128xf32>) : i32 { + %r = arith.addf %arg, %arg {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + scf.yield %r : tensor<128xf32> + } {ssbuffer.main_loop = 1 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return + } + + //===--------------------------------------------------------------------===// + // T6: collectBlockInfo - Group ops by ssbuffer.block_id + // Test: single block → returns 1 block group (no buffer needed) + // Key Check: block_id 5 both producer and consumer + //===--------------------------------------------------------------------===// + // CHECK-LABEL: func.func @test_t6_collectblockinfo + // CHECK: linalg.fill {ssbuffer.block_id = 5 : i32} + // CHECK: arith.addf {{.*}} {ssbuffer.block_id = 5 : i32} + func.func @test_t6_collectblockinfo() { + %c0_i32 = arith.constant 0 : i32 + %c100_i32 = arith.constant 100 : i32 + %c1_i32 = arith.constant 1 : i32 + %cst = arith.constant 1.0 : f32 + %empty = tensor.empty() : tensor<64xf32> + scope.scope : () -> () { + %prod = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst : f32) outs(%empty : tensor<64xf32>) -> tensor<64xf32> + %cons = arith.addf %prod, %prod {ssbuffer.block_id = 5 : i32} : tensor<64xf32> + scf.for %i = %c0_i32 to %c100_i32 step %c1_i32 iter_args(%arg = %cons) -> (tensor<64xf32>) : i32 { + %r = arith.addf %arg, %arg {ssbuffer.block_id = 5 : i32} : tensor<64xf32> + scf.yield %r : tensor<64xf32> + } {ssbuffer.main_loop = 1 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return + } + + //===--------------------------------------------------------------------===// + // T7: dep_mark Array Attribute Test + // Test: scalar deps inside main_loop get dep_mark attribute + // Key Check: dep_mark array attribute on ops + //===--------------------------------------------------------------------===// + // CHECK-LABEL: func.func @test_t7_dep_mark + // CHECK-DAG: memref.load {{.*}} {ssbuffer.block_id = 11 : i32} + // CHECK-DAG: arith.addi {{.*}} ssbuffer.dep_mark + // CHECK-DAG: arith.muli {{.*}} ssbuffer.dep_mark + // CHECK-DAG: arith.divsi {{.*}} ssbuffer.dep_mark + func.func @test_t7_dep_mark() { + %true = arith.constant true + %c0_i64 = arith.constant 0 : i64 + %c100_i64 = arith.constant 100 : i64 + %c1_i64 = arith.constant 1 : i64 + %c0 = arith.constant 0 : index + scope.scope : () -> () { + %alloc = memref.alloc() : memref<1xi64> + %cast_alloc = memref.cast %alloc : memref<1xi64> to memref<1xi64, strided<[1], offset: ?>> + scf.for %i = %c0_i64 to %c100_i64 step %c1_i64 : i64 { + %load = memref.load %cast_alloc[%c0] {ssbuffer.block_id = 11 : i32} : memref<1xi64, strided<[1], offset: ?>> + %compute = arith.addi %load, %c1_i64 {ssbuffer.block_id = 11 : i32} : i64 + %result = scf.if %true -> (i64) { + %r1 = arith.muli %compute, %c1_i64 {ssbuffer.block_id = 8 : i32} : i64 + scf.yield %r1 : i64 + } else { + %r2 = arith.divsi %compute, %c1_i64 {ssbuffer.block_id = 9 : i32} : i64 + scf.yield %r2 : i64 + } + memref.store %compute, %cast_alloc[%c0] {ssbuffer.block_id = 12 : i32} : memref<1xi64, strided<[1], offset: ?>> + } {ssbuffer.block_id = 26 : i32, ssbuffer.main_loop = 1 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return + } + + //===--------------------------------------------------------------------===// + // T8: Memref Operand Test + // Test: memref.alloc with block_id passed to materialize_in_destination + // with different block_id. The tensor operand crosses blocks, but + // memref operand is not tracked as cross-block dep in typical flow. + // Key Check: double buffer is created since cross-block dep is tensor type + //===--------------------------------------------------------------------===// + // CHECK-LABEL: func.func @test_t8_memref_operand + // CHECK: scope.scope + // CHECK: memref.memory_space_cast {{.*}} {ssbuffer.intraDeps + func.func @test_t8_memref_operand() { + %c0_i32 = arith.constant 0 : i32 + %c100_i32 = arith.constant 100 : i32 + %c1_i32 = arith.constant 1 : i32 + %cst = arith.constant 1.0 : f32 + %empty = tensor.empty() : tensor<128xf32> + scope.scope : () -> () { + %mem = memref.alloc() {ssbuffer.block_id = 5 : i32} : memref<128xf32> + %prod = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst : f32) outs(%empty : tensor<128xf32>) -> tensor<128xf32> + %loop_result = scf.for %i = %c0_i32 to %c100_i32 step %c1_i32 iter_args(%arg = %prod) -> (tensor<128xf32>) : i32 { + %consumed = arith.addf %arg, %arg {ssbuffer.block_id = 6 : i32} : tensor<128xf32> + bufferization.materialize_in_destination %consumed in writable %mem {ssbuffer.block_id = 6 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + %new_prod = arith.addf %consumed, %consumed {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + scf.yield %new_prod : tensor<128xf32> + } {ssbuffer.main_loop = 1 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return + } + + //===--------------------------------------------------------------------===// + // T9: Multiple Main Loops - Global groupId Uniqueness + // Test: Two parallel main_loop forOps in same vector scope + // Key Check: intraDeps groupId is globally unique across all main_loops, + // NOT reset to 0 for each main_loop + //===--------------------------------------------------------------------===// + // CHECK-LABEL: func.func @test_t9_multiple_main_loops + // CHECK-DAG: memref.memory_space_cast {{.*}} {ssbuffer.intraDeps = [0 : i32, 1 : i32]} + // CHECK-DAG: memref.memory_space_cast {{.*}} {ssbuffer.intraDeps = [1 : i32, 1 : i32]} + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 6 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 6 + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 5 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 5 + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 8 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 8 + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 7 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 7 + // CHECK-DAG: } {ssbuffer.main_loop = 1 : i64} + // CHECK-DAG: } {ssbuffer.main_loop = 1 : i64} + // CHECK-DAG: ssbuffer.intraDeps = [0 : i32, 0 : i32] + // CHECK-DAG: ssbuffer.intraDeps = [1 : i32, 0 : i32] + func.func @test_t9_multiple_main_loops() { + %c0_i32 = arith.constant 0 : i32 + %c100_i32 = arith.constant 100 : i32 + %c1_i32 = arith.constant 1 : i32 + %cst = arith.constant 1.0 : f32 + %empty_128 = tensor.empty() : tensor<128xf32> + %empty_64 = tensor.empty() : tensor<64xf32> + scope.scope : () -> () { + %prod1 = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst : f32) outs(%empty_128 : tensor<128xf32>) -> tensor<128xf32> + %loop_result1 = scf.for %i = %c0_i32 to %c100_i32 step %c1_i32 iter_args(%arg1 = %prod1) -> (tensor<128xf32>) : i32 { + %cons1 = arith.addf %arg1, %arg1 {ssbuffer.block_id = 6 : i32} : tensor<128xf32> + %new_prod1 = arith.addf %cons1, %cons1 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + scf.yield %new_prod1 : tensor<128xf32> + } {ssbuffer.main_loop = 1 : i64} + %prod2 = linalg.fill {ssbuffer.block_id = 7 : i32} ins(%cst : f32) outs(%empty_64 : tensor<64xf32>) -> tensor<64xf32> + %loop_result2 = scf.for %j = %c0_i32 to %c100_i32 step %c1_i32 iter_args(%arg2 = %prod2) -> (tensor<64xf32>) : i32 { + %cons2 = arith.mulf %arg2, %arg2 {ssbuffer.block_id = 8 : i32} : tensor<64xf32> + %new_prod2 = arith.mulf %cons2, %cons2 {ssbuffer.block_id = 7 : i32} : tensor<64xf32> + scf.yield %new_prod2 : tensor<64xf32> + } {ssbuffer.main_loop = 1 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return + } + + //===--------------------------------------------------------------------===// + // T10: Single Buffer Test (N==1) - Verification of Single Buffer Tags + // Test: When BufferCountManager returns 1, single buffer is created + // Key Check: producer (materialize_in_destination) and consumer (to_tensor) + // : both tagged with block_id, intraDeps=[groupId,0], intra_buffer + // Note: This test shows the expected format when bufNum==1. + // Default bufNum is 2, so this test follows double buffer output. + // For true single buffer test, need BufferCountManager::setBufferCount(IntraCore, 1) + //===--------------------------------------------------------------------===// + // CHECK-LABEL: func.func @test_t10_single_buffer + // CHECK-DAG: memref.memory_space_cast {{.*}} {ssbuffer.intraDeps = [0 : i32, 1 : i32]} + // CHECK-DAG: memref.memory_space_cast {{.*}} {ssbuffer.intraDeps = [0 : i32, 1 : i32]} + // CHECK-DAG: } {ssbuffer.block_id = 5 : i32, ssbuffer.intraDeps = [0 : i32, 0 : i32], ssbuffer.intra_buffer} + func.func @test_t10_single_buffer() { + %c0_i32 = arith.constant 0 : i32 + %c100_i32 = arith.constant 100 : i32 + %c1_i32 = arith.constant 1 : i32 + %cst = arith.constant 1.0 : f32 + %empty = tensor.empty() : tensor<128xf32> + scope.scope : () -> () { + %prod = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst : f32) outs(%empty : tensor<128xf32>) -> tensor<128xf32> + %loop_result = scf.for %i = %c0_i32 to %c100_i32 step %c1_i32 iter_args(%arg = %prod) -> (tensor<128xf32>) : i32 { + %consumed = arith.addf %arg, %arg {ssbuffer.block_id = 6 : i32} : tensor<128xf32> + %new_prod = arith.addf %consumed, %consumed {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + scf.yield %new_prod : tensor<128xf32> + } {ssbuffer.main_loop = 1 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return + } + + //===--------------------------------------------------------------------===// + // T11: DivUIOp/SubIOp Counter Test (lb != 0 or step != 1) + // Test: When lb != 0, SubIOp is created for (iv - lb) / step + // When step != 1, DivUIOp is created for iterIdx / step + // Key Check: Both SubIOp and DivUIOp get block_id tag + // : remsi and cmpi also get block_id tag + //===--------------------------------------------------------------------===// + // CHECK-LABEL: func.func @test_t11_counter_with_divsub + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 6 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 6 + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 5 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 5 + // CHECK-DAG: arith.divui {{.*}} {ssbuffer.block_id = 5 + func.func @test_t11_counter_with_divsub() { + %c10_i32 = arith.constant 10 : i32 + %c100_i32 = arith.constant 100 : i32 + %c2_i32 = arith.constant 2 : i32 + %cst = arith.constant 1.0 : f32 + %empty = tensor.empty() : tensor<128xf32> + scope.scope : () -> () { + %prod = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst : f32) outs(%empty : tensor<128xf32>) -> tensor<128xf32> + %loop_result = scf.for %i = %c10_i32 to %c100_i32 step %c2_i32 iter_args(%arg = %prod) -> (tensor<128xf32>) : i32 { + %consumed = arith.addf %arg, %arg {ssbuffer.block_id = 6 : i32} : tensor<128xf32> + %new_prod = arith.addf %consumed, %consumed {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + scf.yield %new_prod : tensor<128xf32> + } {ssbuffer.main_loop = 1 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return + } + + //===--------------------------------------------------------------------===// + // T12: Index Type Loop Variable Test (IndexCastOp block_id) + // Test: When loop induction variable is 'index' type, IndexCastOp is created + // to convert from index to i32 for buffer index calculation. + // Key Check: IndexCastOp gets block_id tag from the block it belongs to + //===--------------------------------------------------------------------===// + // CHECK-LABEL: func.func @test_t12_index_type_loop_var + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 6 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 6 + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 5 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 5 + // CHECK-DAG: arith.index_cast {{.*}} {ssbuffer.block_id = 5 + func.func @test_t12_index_type_loop_var() { + %c0 = arith.constant 0 : index + %c100 = arith.constant 100 : index + %c1 = arith.constant 1 : index + %cst = arith.constant 1.0 : f32 + %empty = tensor.empty() : tensor<128xf32> + scope.scope : () -> () { + %prod = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst : f32) outs(%empty : tensor<128xf32>) -> tensor<128xf32> + %loop_result = scf.for %i = %c0 to %c100 step %c1 iter_args(%arg = %prod) -> (tensor<128xf32>) : index { + %consumed = arith.addf %arg, %arg {ssbuffer.block_id = 6 : i32} : tensor<128xf32> + %new_prod = arith.addf %consumed, %consumed {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + scf.yield %new_prod : tensor<128xf32> + } {ssbuffer.main_loop = 1 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return + } + + //===--------------------------------------------------------------------===// + // T13: Three Parallel Main Loops - Global groupId Uniqueness + // Test: Three parallel main_loop forOps in same vector scope + // Key Check: groupId is globally unique across all main_loops (0, 1, 2) + // : Each main_loop's intraDeps reference different groupIds + // : remsi and cmpi get correct block_id tags + //===--------------------------------------------------------------------===// + // CHECK-LABEL: func.func @test_t13_three_parallel_main_loops + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 6 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 6 + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 5 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 5 + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 8 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 8 + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 7 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 7 + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 10 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 10 + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 9 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 9 + // CHECK-DAG: } {ssbuffer.main_loop = 1 : i64} + // CHECK-DAG: } {ssbuffer.main_loop = 1 : i64} + // CHECK-DAG: } {ssbuffer.main_loop = 1 : i64} + // CHECK-DAG: ssbuffer.intraDeps = [0 : i32, 1 : i32] + // CHECK-DAG: ssbuffer.intraDeps = [1 : i32, 1 : i32] + // CHECK-DAG: ssbuffer.intraDeps = [2 : i32, 1 : i32] + func.func @test_t13_three_parallel_main_loops() { + %c0_i32 = arith.constant 0 : i32 + %c100_i32 = arith.constant 100 : i32 + %c1_i32 = arith.constant 1 : i32 + %cst = arith.constant 1.0 : f32 + %empty_128 = tensor.empty() : tensor<128xf32> + %empty_64 = tensor.empty() : tensor<64xf32> + %empty_32 = tensor.empty() : tensor<32xf32> + scope.scope : () -> () { + // First main_loop: block_id 5 -> 6 + %prod1 = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst : f32) outs(%empty_128 : tensor<128xf32>) -> tensor<128xf32> + %loop_result1 = scf.for %i = %c0_i32 to %c100_i32 step %c1_i32 iter_args(%arg1 = %prod1) -> (tensor<128xf32>) : i32 { + %cons1 = arith.addf %arg1, %arg1 {ssbuffer.block_id = 6 : i32} : tensor<128xf32> + %new_prod1 = arith.addf %cons1, %cons1 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + scf.yield %new_prod1 : tensor<128xf32> + } {ssbuffer.main_loop = 1 : i64} + // Second main_loop: block_id 7 -> 8 + %prod2 = linalg.fill {ssbuffer.block_id = 7 : i32} ins(%cst : f32) outs(%empty_64 : tensor<64xf32>) -> tensor<64xf32> + %loop_result2 = scf.for %j = %c0_i32 to %c100_i32 step %c1_i32 iter_args(%arg2 = %prod2) -> (tensor<64xf32>) : i32 { + %cons2 = arith.mulf %arg2, %arg2 {ssbuffer.block_id = 8 : i32} : tensor<64xf32> + %new_prod2 = arith.mulf %cons2, %cons2 {ssbuffer.block_id = 7 : i32} : tensor<64xf32> + scf.yield %new_prod2 : tensor<64xf32> + } {ssbuffer.main_loop = 1 : i64} + // Third main_loop: block_id 9 -> 10 + %prod3 = linalg.fill {ssbuffer.block_id = 9 : i32} ins(%cst : f32) outs(%empty_32 : tensor<32xf32>) -> tensor<32xf32> + %loop_result3 = scf.for %k = %c0_i32 to %c100_i32 step %c1_i32 iter_args(%arg3 = %prod3) -> (tensor<32xf32>) : i32 { + %cons3 = arith.subf %arg3, %arg3 {ssbuffer.block_id = 10 : i32} : tensor<32xf32> + %new_prod3 = arith.addf %cons3, %cons3 {ssbuffer.block_id = 9 : i32} : tensor<32xf32> + scf.yield %new_prod3 : tensor<32xf32> + } {ssbuffer.main_loop = 1 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return + } + + //===--------------------------------------------------------------------===// + // T14: i64 Loop Variable Test (TruncIOp block_id) + // Test: When loop induction variable is i64 (width > 32), TruncIOp is created + // to truncate i64 to i32 for buffer index calculation. + // Key Check: TruncIOp gets block_id tag from the block it belongs to + //===--------------------------------------------------------------------===// + // CHECK-LABEL: func.func @test_t14_i64_loop_var + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 6 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 6 + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 5 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 5 + // CHECK-DAG: arith.trunci {{.*}} {ssbuffer.block_id = 5 + func.func @test_t14_i64_loop_var() { + %c0_i64 = arith.constant 0 : i64 + %c100_i64 = arith.constant 100 : i64 + %c1_i64 = arith.constant 1 : i64 + %cst = arith.constant 1.0 : f32 + %empty = tensor.empty() : tensor<128xf32> + scope.scope : () -> () { + %prod = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst : f32) outs(%empty : tensor<128xf32>) -> tensor<128xf32> + %loop_result = scf.for %i = %c0_i64 to %c100_i64 step %c1_i64 iter_args(%arg = %prod) -> (tensor<128xf32>) : i64 { + %consumed = arith.addf %arg, %arg {ssbuffer.block_id = 6 : i32} : tensor<128xf32> + %new_prod = arith.addf %consumed, %consumed {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + scf.yield %new_prod : tensor<128xf32> + } {ssbuffer.main_loop = 1 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return + } + + //===--------------------------------------------------------------------===// + // T15: i16 Loop Variable Test (ExtSIOp block_id) + // Test: When loop induction variable is i16 (width < 32), ExtSIOp is created + // to extend i16 to i32 for buffer index calculation. + // Key Check: ExtSIOp gets block_id tag from the block it belongs to + //===--------------------------------------------------------------------===// + // CHECK-LABEL: func.func @test_t15_i16_loop_var + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 6 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 6 + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 5 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 5 + // CHECK-DAG: arith.extsi {{.*}} {ssbuffer.block_id = 5 + func.func @test_t15_i16_loop_var() { + %c0_i16 = arith.constant 0 : i16 + %c100_i16 = arith.constant 100 : i16 + %c1_i16 = arith.constant 1 : i16 + %cst = arith.constant 1.0 : f32 + %empty = tensor.empty() : tensor<128xf32> + scope.scope : () -> () { + %prod = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst : f32) outs(%empty : tensor<128xf32>) -> tensor<128xf32> + %loop_result = scf.for %i = %c0_i16 to %c100_i16 step %c1_i16 iter_args(%arg = %prod) -> (tensor<128xf32>) : i16 { + %consumed = arith.addf %arg, %arg {ssbuffer.block_id = 6 : i32} : tensor<128xf32> + %new_prod = arith.addf %consumed, %consumed {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + scf.yield %new_prod : tensor<128xf32> + } {ssbuffer.main_loop = 1 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return + } + + //===--------------------------------------------------------------------===// + // T16: Multiple Parallel Main Loops with Complex Loop Variables + // Test: Three parallel main_loop forOps with different loop variable types + // - loop1: i32, lb=10, step=2 -> creates SubIOp, DivUIOp + // - loop2: index type -> creates IndexCastOp + // - loop3: i64 type -> creates TruncIOp + // Key Check: All counter operations get correct block_id tags from each loop + //===--------------------------------------------------------------------===// + // CHECK-LABEL: func.func @test_t16_parallel_complex_loops + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 6 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 6 + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 5 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 5 + // CHECK-DAG: arith.divui {{.*}} {ssbuffer.block_id = 5 + // CHECK-DAG: arith.subi {{.*}} {ssbuffer.block_id = 5 + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 8 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 8 + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 7 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 7 + // CHECK-DAG: arith.index_cast {{.*}} {ssbuffer.block_id = 7 + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 10 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 10 + // CHECK-DAG: arith.remsi {{.*}} {ssbuffer.block_id = 9 + // CHECK-DAG: arith.cmpi eq, {{.*}} {ssbuffer.block_id = 9 + // CHECK-DAG: arith.trunci {{.*}} {ssbuffer.block_id = 9 + func.func @test_t16_parallel_complex_loops() { + %c0_i32 = arith.constant 0 : i32 + %c100_i32 = arith.constant 100 : i32 + %c1_i32 = arith.constant 1 : i32 + %c10_i32 = arith.constant 10 : i32 + %c2_i32 = arith.constant 2 : i32 + %c0 = arith.constant 0 : index + %c100 = arith.constant 100 : index + %c1 = arith.constant 1 : index + %c0_i64 = arith.constant 0 : i64 + %c100_i64 = arith.constant 100 : i64 + %c1_i64 = arith.constant 1 : i64 + %cst = arith.constant 1.0 : f32 + %empty_128 = tensor.empty() : tensor<128xf32> + %empty_64 = tensor.empty() : tensor<64xf32> + %empty_32 = tensor.empty() : tensor<32xf32> + scope.scope : () -> () { + // First main_loop: i32, lb=10, step=2 -> creates SubIOp, DivUIOp + %prod1 = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst : f32) outs(%empty_128 : tensor<128xf32>) -> tensor<128xf32> + %loop_result1 = scf.for %i = %c10_i32 to %c100_i32 step %c2_i32 iter_args(%arg1 = %prod1) -> (tensor<128xf32>) : i32 { + %cons1 = arith.addf %arg1, %arg1 {ssbuffer.block_id = 6 : i32} : tensor<128xf32> + %new_prod1 = arith.addf %cons1, %cons1 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + scf.yield %new_prod1 : tensor<128xf32> + } {ssbuffer.main_loop = 1 : i64} + // Second main_loop: index type -> creates IndexCastOp + %prod2 = linalg.fill {ssbuffer.block_id = 7 : i32} ins(%cst : f32) outs(%empty_64 : tensor<64xf32>) -> tensor<64xf32> + %loop_result2 = scf.for %j = %c0 to %c100 step %c1 iter_args(%arg2 = %prod2) -> (tensor<64xf32>) : index { + %cons2 = arith.mulf %arg2, %arg2 {ssbuffer.block_id = 8 : i32} : tensor<64xf32> + %new_prod2 = arith.mulf %cons2, %cons2 {ssbuffer.block_id = 7 : i32} : tensor<64xf32> + scf.yield %new_prod2 : tensor<64xf32> + } {ssbuffer.main_loop = 1 : i64} + // Third main_loop: i64 type -> creates TruncIOp + %prod3 = linalg.fill {ssbuffer.block_id = 9 : i32} ins(%cst : f32) outs(%empty_32 : tensor<32xf32>) -> tensor<32xf32> + %loop_result3 = scf.for %k = %c0_i64 to %c100_i64 step %c1_i64 iter_args(%arg3 = %prod3) -> (tensor<32xf32>) : i64 { + %cons3 = arith.subf %arg3, %arg3 {ssbuffer.block_id = 10 : i32} : tensor<32xf32> + %new_prod3 = arith.addf %cons3, %cons3 {ssbuffer.block_id = 9 : i32} : tensor<32xf32> + scf.yield %new_prod3 : tensor<32xf32> + } {ssbuffer.main_loop = 1 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return + } + +} diff --git a/third_party/ascend/unittest/Conversion/General/AutoBlockify/auto_blockify.mlir b/third_party/ascend/unittest/Conversion/General/AutoBlockify/auto_blockify.mlir new file mode 100644 index 0000000000..c9a1dcb409 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/AutoBlockify/auto_blockify.mlir @@ -0,0 +1,134 @@ +// RUN: triton-opt --auto-blockify="auto-blockify-size=5" --split-input-file %s | FileCheck %s + +// ----- + +// CHECK-LABEL: tt.func @kernel( +// CHECK-SAME: %[[VAL_0:.*]]: !tt.ptr) attributes {auto_blockify_size = 5 : i32} { +// CHECK: %[[VAL_1:.*]] = arith.constant dense<0.000000e+00> : tensor<5x8xf32> +// CHECK: %[[VAL_2:.*]] = arith.constant dense<8> : tensor<5xi32> +// CHECK: %[[VAL_3:.*]] = arith.constant dense<0> : tensor<5xi32> +// CHECK: %[[VAL_4:.*]] = tt.get_num_programs x : i32 +// CHECK: %[[VAL_5:.*]] = tt.get_num_programs y : i32 +// CHECK: %[[VAL_6:.*]] = tt.get_num_programs z : i32 +// CHECK: %[[VAL_7:.*]] = arith.muli %[[VAL_5]], %[[VAL_6]] : i32 +// CHECK: %[[VAL_8:.*]] = arith.muli %[[VAL_7]], %[[VAL_4]] : i32 +// CHECK: %[[VAL_9:.*]] = tt.get_program_id x {logical_block_id} : i32 +// CHECK: %[[VAL_10:.*]] = tt.get_program_id y {logical_block_id} : i32 +// CHECK: %[[VAL_11:.*]] = tt.get_program_id z {logical_block_id} : i32 +// CHECK: %[[VAL_12:.*]] = arith.muli %[[VAL_9]], %[[VAL_7]] : i32 +// CHECK: %[[VAL_13:.*]] = arith.muli %[[VAL_10]], %[[VAL_6]] : i32 +// CHECK: %[[VAL_14:.*]] = arith.addi %[[VAL_12]], %[[VAL_13]] : i32 +// CHECK: %[[VAL_15:.*]] = arith.addi %[[VAL_14]], %[[VAL_11]] : i32 +// CHECK: %[[VAL_16:.*]] = tt.make_range {end = 5 : i32, start = 0 : i32} : tensor<5xi32> +// CHECK: %[[VAL_17:.*]] = tt.splat %[[VAL_15]] : i32 -> tensor<5xi32> +// CHECK: %[[VAL_18:.*]] = arith.addi %[[VAL_17]], %[[VAL_16]] : tensor<5xi32> +// CHECK: %[[VAL_19:.*]] = tt.splat %[[VAL_8]] : i32 -> tensor<5xi32> +// CHECK: %[[VAL_20:.*]] = arith.cmpi slt, %[[VAL_18]], %[[VAL_19]] : tensor<5xi32> +// CHECK: %[[VAL_21:.*]] = arith.cmpi sge, %[[VAL_18]], %[[VAL_3]] : tensor<5xi32> +// CHECK: %[[VAL_22:.*]] = arith.ori %[[VAL_20]], %[[VAL_21]] : tensor<5xi1> +// CHECK: %[[VAL_23:.*]] = tt.splat %[[VAL_7]] : i32 -> tensor<5xi32> +// CHECK: %[[VAL_24:.*]] = arith.divsi %[[VAL_18]], %[[VAL_23]] : tensor<5xi32> +// CHECK: %[[VAL_25:.*]] = tt.splat %[[VAL_4]] : i32 -> tensor<5xi32> +// CHECK: %[[VAL_26:.*]] = arith.remsi %[[VAL_24]], %[[VAL_25]] : tensor<5xi32> +// CHECK: %[[VAL_27:.*]] = arith.muli %[[VAL_26]], %[[VAL_2]] : tensor<5xi32> +// CHECK: %[[VAL_28:.*]] = tt.expand_dims %[[VAL_27]] {axis = 1 : i32} : tensor<5xi32> -> tensor<5x1xi32> +// CHECK: %[[VAL_29:.*]] = tt.broadcast %[[VAL_28]] : tensor<5x1xi32> -> tensor<5x8xi32> +// CHECK: %[[VAL_30:.*]] = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> +// CHECK: %[[VAL_31:.*]] = tt.expand_dims %[[VAL_30]] {axis = 0 : i32} : tensor<8xi32> -> tensor<1x8xi32> +// CHECK: %[[VAL_32:.*]] = tt.broadcast %[[VAL_31]] : tensor<1x8xi32> -> tensor<5x8xi32> +// CHECK: %[[VAL_33:.*]] = arith.addi %[[VAL_29]], %[[VAL_32]] : tensor<5x8xi32> +// CHECK: %[[VAL_34:.*]] = tt.splat %[[VAL_0]] : !tt.ptr -> tensor<5x8x!tt.ptr> +// CHECK: %[[VAL_35:.*]] = tt.addptr %[[VAL_34]], %[[VAL_33]] : tensor<5x8x!tt.ptr>, tensor<5x8xi32> +// CHECK: %[[VAL_36:.*]] = tt.expand_dims %[[VAL_22]] {axis = 1 : i32} : tensor<5xi1> -> tensor<5x1xi1> +// CHECK: %[[VAL_37:.*]] = tt.broadcast %[[VAL_36]] : tensor<5x1xi1> -> tensor<5x8xi1> +// CHECK: tt.store %[[VAL_35]], %[[VAL_1]], %[[VAL_37]] : tensor<5x8x!tt.ptr> +// CHECK: tt.return +// CHECK: } +tt.func @kernel(%arg0: !tt.ptr) { + %cst = arith.constant dense<0.000000e+00> : tensor<8xf32> + %c8_i32 = arith.constant 8 : i32 + %0 = tt.get_program_id x : i32 + %1 = arith.muli %0, %c8_i32 : i32 + %2 = tt.splat %1 : i32 -> tensor<8xi32> + %3 = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> + %4 = arith.addi %2, %3 : tensor<8xi32> + %5 = tt.splat %arg0 : !tt.ptr -> tensor<8x!tt.ptr> + %6 = tt.addptr %5, %4 : tensor<8x!tt.ptr>, tensor<8xi32> + tt.store %6, %cst : tensor<8x!tt.ptr> + tt.return +} + +// ----- + +// CHECK-LABEL: tt.func @kernel2( +// CHECK-SAME: %[[VAL_0:.*]]: !tt.ptr) attributes {auto_blockify_size = 5 : i32} { +// CHECK: %[[VAL_1:.*]] = arith.constant dense<8> : tensor<5xi32> +// CHECK: %[[VAL_2:.*]] = arith.constant 0 : i32 +// CHECK: %[[VAL_3:.*]] = arith.constant 5 : index +// CHECK: %[[VAL_4:.*]] = arith.constant 1 : index +// CHECK: %[[VAL_5:.*]] = arith.constant 0 : index +// CHECK: %[[VAL_6:.*]] = arith.constant dense<0.000000e+00> : tensor<8xf32> +// CHECK: %[[VAL_7:.*]] = tt.get_num_programs x : i32 +// CHECK: %[[VAL_8:.*]] = tt.get_num_programs y : i32 +// CHECK: %[[VAL_9:.*]] = tt.get_num_programs z : i32 +// CHECK: %[[VAL_10:.*]] = arith.muli %[[VAL_8]], %[[VAL_9]] : i32 +// CHECK: %[[VAL_11:.*]] = arith.muli %[[VAL_10]], %[[VAL_7]] : i32 +// CHECK: %[[VAL_12:.*]] = tt.get_program_id x {logical_block_id} : i32 +// CHECK: %[[VAL_13:.*]] = tt.get_program_id y {logical_block_id} : i32 +// CHECK: %[[VAL_14:.*]] = tt.get_program_id z {logical_block_id} : i32 +// CHECK: %[[VAL_15:.*]] = arith.muli %[[VAL_12]], %[[VAL_10]] : i32 +// CHECK: %[[VAL_16:.*]] = arith.muli %[[VAL_13]], %[[VAL_9]] : i32 +// CHECK: %[[VAL_17:.*]] = arith.addi %[[VAL_15]], %[[VAL_16]] : i32 +// CHECK: %[[VAL_18:.*]] = arith.addi %[[VAL_17]], %[[VAL_14]] : i32 +// CHECK: %[[VAL_19:.*]] = tt.make_range {end = 5 : i32, start = 0 : i32} : tensor<5xi32> +// CHECK: %[[VAL_20:.*]] = tt.splat %[[VAL_18]] : i32 -> tensor<5xi32> +// CHECK: %[[VAL_21:.*]] = arith.addi %[[VAL_20]], %[[VAL_19]] : tensor<5xi32> +// CHECK: %[[VAL_22:.*]] = tt.splat %[[VAL_10]] : i32 -> tensor<5xi32> +// CHECK: %[[VAL_23:.*]] = arith.divsi %[[VAL_21]], %[[VAL_22]] : tensor<5xi32> +// CHECK: %[[VAL_24:.*]] = tt.splat %[[VAL_7]] : i32 -> tensor<5xi32> +// CHECK: %[[VAL_25:.*]] = arith.remsi %[[VAL_23]], %[[VAL_24]] : tensor<5xi32> +// CHECK: %[[VAL_26:.*]] = tt.splat %[[VAL_9]] : i32 -> tensor<5xi32> +// CHECK: %[[VAL_27:.*]] = arith.divsi %[[VAL_21]], %[[VAL_26]] : tensor<5xi32> +// CHECK: %[[VAL_28:.*]] = tt.splat %[[VAL_8]] : i32 -> tensor<5xi32> +// CHECK: %[[VAL_29:.*]] = arith.remsi %[[VAL_27]], %[[VAL_28]] : tensor<5xi32> +// CHECK: %[[VAL_30:.*]] = arith.cmpi slt, %[[VAL_29]], %[[VAL_1]] : tensor<5xi32> +// CHECK: %[[VAL_31:.*]] = arith.muli %[[VAL_25]], %[[VAL_1]] : tensor<5xi32> +// CHECK: %[[VAL_32:.*]] = tt.expand_dims %[[VAL_31]] {axis = 1 : i32} : tensor<5xi32> -> tensor<5x1xi32> +// CHECK: %[[VAL_33:.*]] = tt.broadcast %[[VAL_32]] : tensor<5x1xi32> -> tensor<5x8xi32> +// CHECK: %[[VAL_34:.*]] = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> +// CHECK: %[[VAL_35:.*]] = tt.expand_dims %[[VAL_34]] {axis = 0 : i32} : tensor<8xi32> -> tensor<1x8xi32> +// CHECK: %[[VAL_36:.*]] = tt.broadcast %[[VAL_35]] : tensor<1x8xi32> -> tensor<5x8xi32> +// CHECK: %[[VAL_37:.*]] = arith.addi %[[VAL_33]], %[[VAL_36]] : tensor<5x8xi32> +// CHECK: %[[VAL_38:.*]] = tt.splat %[[VAL_0]] : !tt.ptr -> tensor<5x8x!tt.ptr> +// CHECK: %[[VAL_39:.*]] = tt.addptr %[[VAL_38]], %[[VAL_37]] : tensor<5x8x!tt.ptr>, tensor<5x8xi32> +// CHECK: %[[VAL_40:.*]] = arith.subi %[[VAL_11]], %[[VAL_18]] : i32 +// CHECK: %[[VAL_41:.*]] = arith.maxsi %[[VAL_40]], %[[VAL_2]] : i32 +// CHECK: %[[VAL_42:.*]] = arith.index_cast %[[VAL_41]] : i32 to index +// CHECK: %[[VAL_43:.*]] = arith.minsi %[[VAL_42]], %[[VAL_3]] : index +// CHECK: scf.for %[[VAL_44:.*]] = %[[VAL_5]] to %[[VAL_43]] step %[[VAL_4]] { +// CHECK: %[[VAL_45:.*]] = tensor.extract %[[VAL_30]]{{\[}}%[[VAL_44]]] : tensor<5xi1> +// CHECK: scf.if %[[VAL_45]] { +// CHECK: %[[VAL_46:.*]] = tensor.extract_slice %[[VAL_39]]{{\[}}%[[VAL_44]], 0] [1, 8] [1, 1] : tensor<5x8x!tt.ptr> to tensor<8x!tt.ptr> +// CHECK: tt.store %[[VAL_46]], %[[VAL_6]] : tensor<8x!tt.ptr> +// CHECK: } +// CHECK: } {auto_blockify_loop} +// CHECK: tt.return +// CHECK: } +tt.func @kernel2(%arg0: !tt.ptr) { + %cst = arith.constant dense<0.000000e+00> : tensor<8xf32> + %c8_i32 = arith.constant 8 : i32 + %0 = tt.get_program_id x : i32 + %a = tt.get_program_id y : i32 + %b = arith.cmpi slt, %a, %c8_i32 : i32 + %1 = arith.muli %0, %c8_i32 : i32 + %2 = tt.splat %1 : i32 -> tensor<8xi32> + %3 = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> + %4 = arith.addi %2, %3 : tensor<8xi32> + %5 = tt.splat %arg0 : !tt.ptr -> tensor<8x!tt.ptr> + %6 = tt.addptr %5, %4 : tensor<8x!tt.ptr>, tensor<8xi32> + scf.if %b { + tt.store %6, %cst : tensor<8x!tt.ptr> + scf.yield + } + tt.return +} diff --git a/third_party/ascend/unittest/Conversion/General/DiscreteMaskAccess/atomic.mlir b/third_party/ascend/unittest/Conversion/General/DiscreteMaskAccess/atomic.mlir new file mode 100644 index 0000000000..9d17d7bc6e --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/DiscreteMaskAccess/atomic.mlir @@ -0,0 +1,201 @@ +// RUN: triton-opt --discrete-mask-access-conversion --split-input-file %s | FileCheck %s + +// CHECK-LABEL: tt.func @atomic_add_i32 +// CHECK: %[[default:.*]] = arith.constant dense<0> : tensor<1024xi32> +// CHECK: %[[value:.*]] = arith.select %[[mask:.*]], %[[origin:.*]], %[[default]] +// CHECK: %[[result:.*]] = tt.atomic_rmw add, acq_rel, gpu, %[[ptr:.*]], %[[value]] +tt.func @atomic_add_i32(%arg0: !tt.ptr, %arg1: !tt.ptr) { + %cst = arith.constant dense<200> : tensor<1024xi32> + %cst_0 = arith.constant dense<400> : tensor<1024xi32> + %0 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> + %1 = arith.cmpi slt, %0, %cst : tensor<1024xi32> + %2 = arith.cmpi sgt, %0, %cst_0 : tensor<1024xi32> + %3 = arith.ori %1, %2 : tensor<1024xi1> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<1024x!tt.ptr> + %5 = tt.addptr %4, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<1024x!tt.ptr> + %7 = tt.addptr %6, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %8 = tt.load %7 : tensor<1024x!tt.ptr> + %9 = tt.atomic_rmw add, acq_rel, gpu, %5, %8, %3 : (tensor<1024x!tt.ptr>, tensor<1024xi32>, tensor<1024xi1>) -> tensor<1024xi32> + tt.return +} + +// CHECK-LABEL: tt.func @atomic_fadd_f32 +// CHECK: %[[default:.*]] = arith.constant dense<0.000000e+00> : tensor<1024xf32> +// CHECK: %[[value:.*]] = arith.select %[[mask:.*]], %[[origin:.*]], %[[default]] +// CHECK: %[[result:.*]] = tt.atomic_rmw fadd, acq_rel, gpu, %[[ptr:.*]], %[[value]] +tt.func @atomic_fadd_f32(%arg0: !tt.ptr, %arg1: !tt.ptr) { + %cst = arith.constant dense<200> : tensor<1024xi32> + %cst_0 = arith.constant dense<400> : tensor<1024xi32> + %0 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> + %1 = arith.cmpi slt, %0, %cst : tensor<1024xi32> + %2 = arith.cmpi sgt, %0, %cst_0 : tensor<1024xi32> + %3 = arith.ori %1, %2 : tensor<1024xi1> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<1024x!tt.ptr> + %5 = tt.addptr %4, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<1024x!tt.ptr> + %7 = tt.addptr %6, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %8 = tt.load %7 : tensor<1024x!tt.ptr> + %9 = tt.atomic_rmw fadd, acq_rel, gpu, %5, %8, %3 : (tensor<1024x!tt.ptr>, tensor<1024xf32>, tensor<1024xi1>) -> tensor<1024xf32> + tt.return +} + +// CHECK-LABEL: tt.func @atomic_max_i32 +// CHECK: %[[default:.*]] = arith.constant dense<-2147483648> : tensor<1024xi32> +// CHECK: %[[value:.*]] = arith.select %[[mask:.*]], %[[origin:.*]], %[[default]] +// CHECK: %[[result:.*]] = tt.atomic_rmw max, acq_rel, gpu, %[[ptr:.*]], %[[value]] +tt.func @atomic_max_i32(%arg0: !tt.ptr, %arg1: !tt.ptr) { + %cst = arith.constant dense<200> : tensor<1024xi32> + %cst_0 = arith.constant dense<400> : tensor<1024xi32> + %0 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> + %1 = arith.cmpi slt, %0, %cst : tensor<1024xi32> + %2 = arith.cmpi sgt, %0, %cst_0 : tensor<1024xi32> + %3 = arith.ori %1, %2 : tensor<1024xi1> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<1024x!tt.ptr> + %5 = tt.addptr %4, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<1024x!tt.ptr> + %7 = tt.addptr %6, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %8 = tt.load %7 : tensor<1024x!tt.ptr> + %9 = tt.atomic_rmw max, acq_rel, gpu, %5, %8, %3 : (tensor<1024x!tt.ptr>, tensor<1024xi32>, tensor<1024xi1>) -> tensor<1024xi32> + tt.return +} + +// CHECK-LABEL: tt.func @atomic_umax_i32 +// CHECK: %[[default:.*]] = arith.constant dense<0> : tensor<1024xi32> +// CHECK: %[[value:.*]] = arith.select %[[mask:.*]], %[[origin:.*]], %[[default]] +// CHECK: %[[result:.*]] = tt.atomic_rmw umax, acq_rel, gpu, %[[ptr:.*]], %[[value]] +tt.func @atomic_umax_i32(%arg0: !tt.ptr, %arg1: !tt.ptr) { + %cst = arith.constant dense<200> : tensor<1024xi32> + %cst_0 = arith.constant dense<400> : tensor<1024xi32> + %0 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> + %1 = arith.cmpi slt, %0, %cst : tensor<1024xi32> + %2 = arith.cmpi sgt, %0, %cst_0 : tensor<1024xi32> + %3 = arith.ori %1, %2 : tensor<1024xi1> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<1024x!tt.ptr> + %5 = tt.addptr %4, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<1024x!tt.ptr> + %7 = tt.addptr %6, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %8 = tt.load %7 : tensor<1024x!tt.ptr> + %9 = tt.atomic_rmw umax, acq_rel, gpu, %5, %8, %3 : (tensor<1024x!tt.ptr>, tensor<1024xi32>, tensor<1024xi1>) -> tensor<1024xi32> + tt.return +} + +// CHECK-LABEL: tt.func @atomic_min_i32 +// CHECK: %[[default:.*]] = arith.constant dense<2147483647> : tensor<1024xi32> +// CHECK: %[[value:.*]] = arith.select %[[mask:.*]], %[[origin:.*]], %[[default]] +// CHECK: %[[result:.*]] = tt.atomic_rmw min, acq_rel, gpu, %[[ptr:.*]], %[[value]] +tt.func @atomic_min_i32(%arg0: !tt.ptr, %arg1: !tt.ptr) { + %cst = arith.constant dense<200> : tensor<1024xi32> + %cst_0 = arith.constant dense<400> : tensor<1024xi32> + %0 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> + %1 = arith.cmpi slt, %0, %cst : tensor<1024xi32> + %2 = arith.cmpi sgt, %0, %cst_0 : tensor<1024xi32> + %3 = arith.ori %1, %2 : tensor<1024xi1> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<1024x!tt.ptr> + %5 = tt.addptr %4, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<1024x!tt.ptr> + %7 = tt.addptr %6, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %8 = tt.load %7 : tensor<1024x!tt.ptr> + %9 = tt.atomic_rmw min, acq_rel, gpu, %5, %8, %3 : (tensor<1024x!tt.ptr>, tensor<1024xi32>, tensor<1024xi1>) -> tensor<1024xi32> + tt.return +} + +// CHECK-LABEL: tt.func @atomic_umin_i32 +// CHECK: %[[default:.*]] = arith.constant dense<2147483647> : tensor<1024xi32> +// CHECK: %[[value:.*]] = arith.select %[[mask:.*]], %[[origin:.*]], %[[default]] +// CHECK: %[[result:.*]] = tt.atomic_rmw umin, acq_rel, gpu, %[[ptr:.*]], %[[value]] +tt.func @atomic_umin_i32(%arg0: !tt.ptr, %arg1: !tt.ptr) { + %cst = arith.constant dense<200> : tensor<1024xi32> + %cst_0 = arith.constant dense<400> : tensor<1024xi32> + %0 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> + %1 = arith.cmpi slt, %0, %cst : tensor<1024xi32> + %2 = arith.cmpi sgt, %0, %cst_0 : tensor<1024xi32> + %3 = arith.ori %1, %2 : tensor<1024xi1> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<1024x!tt.ptr> + %5 = tt.addptr %4, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<1024x!tt.ptr> + %7 = tt.addptr %6, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %8 = tt.load %7 : tensor<1024x!tt.ptr> + %9 = tt.atomic_rmw umin, acq_rel, gpu, %5, %8, %3 : (tensor<1024x!tt.ptr>, tensor<1024xi32>, tensor<1024xi1>) -> tensor<1024xi32> + tt.return +} + +// CHECK-LABEL: tt.func @atomic_and_i32 +// CHECK: %[[default:.*]] = arith.constant dense<2147483647> : tensor<1024xi32> +// CHECK: %[[value:.*]] = arith.select %[[mask:.*]], %[[origin:.*]], %[[default]] +// CHECK: %[[result:.*]] = tt.atomic_rmw and, acq_rel, gpu, %[[ptr:.*]], %[[value]] +tt.func @atomic_and_i32(%arg0: !tt.ptr, %arg1: !tt.ptr) { + %cst = arith.constant dense<200> : tensor<1024xi32> + %cst_0 = arith.constant dense<400> : tensor<1024xi32> + %0 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> + %1 = arith.cmpi slt, %0, %cst : tensor<1024xi32> + %2 = arith.cmpi sgt, %0, %cst_0 : tensor<1024xi32> + %3 = arith.ori %1, %2 : tensor<1024xi1> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<1024x!tt.ptr> + %5 = tt.addptr %4, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<1024x!tt.ptr> + %7 = tt.addptr %6, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %8 = tt.load %7 : tensor<1024x!tt.ptr> + %9 = tt.atomic_rmw and, acq_rel, gpu, %5, %8, %3 : (tensor<1024x!tt.ptr>, tensor<1024xi32>, tensor<1024xi1>) -> tensor<1024xi32> + tt.return +} + +// CHECK-LABEL: tt.func @atomic_or_i32 +// CHECK: %[[default:.*]] = arith.constant dense<0> : tensor<1024xi32> +// CHECK: %[[value:.*]] = arith.select %[[mask:.*]], %[[origin:.*]], %[[default]] +// CHECK: %[[result:.*]] = tt.atomic_rmw or, acq_rel, gpu, %[[ptr:.*]], %[[value]] +tt.func @atomic_or_i32(%arg0: !tt.ptr, %arg1: !tt.ptr) { + %cst = arith.constant dense<200> : tensor<1024xi32> + %cst_0 = arith.constant dense<400> : tensor<1024xi32> + %0 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> + %1 = arith.cmpi slt, %0, %cst : tensor<1024xi32> + %2 = arith.cmpi sgt, %0, %cst_0 : tensor<1024xi32> + %3 = arith.ori %1, %2 : tensor<1024xi1> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<1024x!tt.ptr> + %5 = tt.addptr %4, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<1024x!tt.ptr> + %7 = tt.addptr %6, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %8 = tt.load %7 : tensor<1024x!tt.ptr> + %9 = tt.atomic_rmw or, acq_rel, gpu, %5, %8, %3 : (tensor<1024x!tt.ptr>, tensor<1024xi32>, tensor<1024xi1>) -> tensor<1024xi32> + tt.return +} + +// CHECK-LABEL: tt.func @atomic_max_i16 +// CHECK: %[[default:.*]] = arith.constant dense<-32768> : tensor<1024xi16> +// CHECK: %[[value:.*]] = arith.select %[[mask:.*]], %[[origin:.*]], %[[default]] +// CHECK: %[[result:.*]] = tt.atomic_rmw max, acq_rel, gpu, %[[ptr:.*]], %[[value]] +tt.func @atomic_max_i16(%arg0: !tt.ptr, %arg1: !tt.ptr) { + %cst = arith.constant dense<200> : tensor<1024xi32> + %cst_0 = arith.constant dense<400> : tensor<1024xi32> + %0 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> + %1 = arith.cmpi slt, %0, %cst : tensor<1024xi32> + %2 = arith.cmpi sgt, %0, %cst_0 : tensor<1024xi32> + %3 = arith.ori %1, %2 : tensor<1024xi1> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<1024x!tt.ptr> + %5 = tt.addptr %4, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<1024x!tt.ptr> + %7 = tt.addptr %6, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %8 = tt.load %7 : tensor<1024x!tt.ptr> + %9 = tt.atomic_rmw max, acq_rel, gpu, %5, %8, %3 : (tensor<1024x!tt.ptr>, tensor<1024xi16>, tensor<1024xi1>) -> tensor<1024xi16> + tt.return +} + +// CHECK-LABEL: tt.func @atomic_max_f16 +// CHECK: %[[default:.*]] = arith.constant dense<0xFC00> : tensor<1024xf16> +// CHECK: %[[value:.*]] = arith.select %[[mask:.*]], %[[origin:.*]], %[[default]] +// CHECK: %[[result:.*]] = tt.atomic_rmw max, acq_rel, gpu, %[[ptr:.*]], %[[value]] +tt.func @atomic_max_f16(%arg0: !tt.ptr, %arg1: !tt.ptr) { + %cst = arith.constant dense<200> : tensor<1024xi32> + %cst_0 = arith.constant dense<400> : tensor<1024xi32> + %0 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> + %1 = arith.cmpi slt, %0, %cst : tensor<1024xi32> + %2 = arith.cmpi sgt, %0, %cst_0 : tensor<1024xi32> + %3 = arith.ori %1, %2 : tensor<1024xi1> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<1024x!tt.ptr> + %5 = tt.addptr %4, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<1024x!tt.ptr> + %7 = tt.addptr %6, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %8 = tt.load %7 : tensor<1024x!tt.ptr> + %9 = tt.atomic_rmw max, acq_rel, gpu, %5, %8, %3 : (tensor<1024x!tt.ptr>, tensor<1024xf16>, tensor<1024xi1>) -> tensor<1024xf16> + tt.return +} diff --git a/third_party/ascend/unittest/Conversion/General/DiscreteMaskAccess/loadstore.mlir b/third_party/ascend/unittest/Conversion/General/DiscreteMaskAccess/loadstore.mlir new file mode 100644 index 0000000000..88a26cade3 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/DiscreteMaskAccess/loadstore.mlir @@ -0,0 +1,66 @@ +// RUN: triton-opt --discrete-mask-access-conversion --split-input-file %s | FileCheck %s + +// CHECK-LABEL: tt.func @discrete_load +// CHECK: %[[loaded_value:.*]] = tt.load %[[load_ptr:.*]] +// CHECK: %[[value:.*]] = arith.select %[[mask:.*]], %[[loaded_value]], %[[other:.*]] +// CHECK: tt.store %[[store_ptr:.*]], %[[value]] +tt.func @discrete_load(%arg0: !tt.ptr, %arg1: !tt.ptr) { + %cst = arith.constant dense<0> : tensor<1024xi32> + %cst_0 = arith.constant dense<200> : tensor<1024xi32> + %cst_1 = arith.constant dense<400> : tensor<1024xi32> + %0 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> + %1 = arith.cmpi slt, %0, %cst_0 : tensor<1024xi32> + %2 = arith.cmpi sgt, %0, %cst_1 : tensor<1024xi32> + %3 = arith.ori %1, %2 : tensor<1024xi1> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<1024x!tt.ptr> + %5 = tt.addptr %4, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<1024x!tt.ptr> + %7 = tt.addptr %6, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %8 = tt.load %5, %3, %cst : tensor<1024x!tt.ptr> + tt.store %7, %8 : tensor<1024x!tt.ptr> + tt.return +} + +// CHECK-LABEL: tt.func @discrete_load_without_other +// CHECK: %[[other:.*]] = arith.constant dense<0> +// CHECK: %[[loaded_value:.*]] = tt.load %[[load_ptr:.*]] +// CHECK: %[[value:.*]] = arith.select %[[mask:.*]], %[[loaded_value]], %[[other]] +// CHECK: tt.store %[[store_ptr:.*]], %[[value]] +tt.func @discrete_load_without_other(%arg0: !tt.ptr, %arg1: !tt.ptr) { + %cst = arith.constant dense<0> : tensor<1024xi32> + %cst_0 = arith.constant dense<200> : tensor<1024xi32> + %cst_1 = arith.constant dense<400> : tensor<1024xi32> + %0 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> + %1 = arith.cmpi slt, %0, %cst_0 : tensor<1024xi32> + %2 = arith.cmpi sgt, %0, %cst_1 : tensor<1024xi32> + %3 = arith.ori %1, %2 : tensor<1024xi1> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<1024x!tt.ptr> + %5 = tt.addptr %4, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<1024x!tt.ptr> + %7 = tt.addptr %6, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %8 = tt.load %5, %3 : tensor<1024x!tt.ptr> + tt.store %7, %8 : tensor<1024x!tt.ptr> + tt.return +} + +// CHECK-LABEL: tt.func @discrete_store +// CHECK: %[[loaded_value:.*]] = tt.load %[[load_ptr:.*]] : tensor<1024x!tt.ptr> +// CHECK: %[[origin_value:.*]] = tt.load %[[store_ptr:.*]] : tensor<1024x!tt.ptr> +// CHECK: %[[store_value:.*]] = arith.select %[[mask:.*]], %[[loaded_value]], %[[origin_value]] +// CHECK: tt.store %[[store_ptr]], %[[store_value]] +tt.func @discrete_store(%arg0: !tt.ptr, %arg1: !tt.ptr) { + %cst = arith.constant dense<0> : tensor<1024xi32> + %cst_0 = arith.constant dense<200> : tensor<1024xi32> + %cst_1 = arith.constant dense<400> : tensor<1024xi32> + %0 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> + %1 = arith.cmpi slt, %0, %cst_0 : tensor<1024xi32> + %2 = arith.cmpi sgt, %0, %cst_1 : tensor<1024xi32> + %3 = arith.ori %1, %2 : tensor<1024xi1> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<1024x!tt.ptr> + %5 = tt.addptr %4, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<1024x!tt.ptr> + %7 = tt.addptr %6, %0 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %8 = tt.load %5 : tensor<1024x!tt.ptr> + tt.store %7, %8, %3 : tensor<1024x!tt.ptr> + tt.return +} diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/add_controlflow_condition_test.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/add_controlflow_condition_test.mlir new file mode 100644 index 0000000000..5a3857ac0a --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/add_controlflow_condition_test.mlir @@ -0,0 +1,169 @@ +// RUN: triton-opt --add-control-flow-condition %s | FileCheck %s + +// Unit Tests for AddControlFlowCondition Pass +// Test Case: fa_fwd_after_multicache.mlir +// +// This test verifies the following check points: +// 1. ssbuffer initialization (llvm.inttoptr exists) +// 2. mainloop upper bound is extended with new calculation ops +// 3. ops in mainloop are wrapped by ifOp with same block_id + +module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} { + func.func @_attn_fwd(%arg0: memref, %arg1: memref, %arg2: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg3: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg5: memref {tt.divisibility = 16 : i32}, %arg6: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg7: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg8: memref {tt.divisibility = 16 : i32}, %arg9: i32, %arg10: i32, %arg11: i32, %arg12: i32, %arg13: i32, %arg14: i32, %arg15: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} { + %cst = arith.constant {ssbuffer.block_id = 5 : i32} dense<[8, 8, 16, 16]> : tensor<4xi64> + %cst_0 = arith.constant {ssbuffer.block_id = 5 : i32} dense<[128, 8, 16]> : tensor<3xi64> + %cst_1 = arith.constant {ssbuffer.block_id = 8 : i32} 0.000000e+00 : f32 + %c28_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 28 : i32 + %c65536_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 65536 : i32 + %c8192_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 8192 : i32 + %c128_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 128 : i32 + %c1048576_i64 = arith.constant {ssbuffer.block_id = 8 : i32} 1048576 : i64 + %c8388608_i64 = arith.constant {ssbuffer.block_id = 8 : i32} 8388608 : i64 + %c8_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 8 : i32 + %c64_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 64 : i32 + %c0_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 0 : i32 + %cst_2 = arith.constant {ssbuffer.block_id = 8 : i32} 5.000000e-01 : f32 + %cst_3 = arith.constant {ssbuffer.block_id = 8 : i32} 0xFF800000 : f32 + %cst_4 = arith.constant {ssbuffer.block_id = 8 : i32} 1.000000e+00 : f32 + %c128 = arith.constant {ssbuffer.block_id = 8 : i32} 128 : index + // CHECK: %{{.+}} = llvm.mlir.constant + // CHECK: %{{.+}} = llvm.inttoptr + // CHECK: llvm.store + scope.scope : () -> () { + %0 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128x128xf32> + %1 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128x128xf32> + %2 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_1 : f32) outs(%0 : tensor<128x128xf32>) -> tensor<128x128xf32> + %3 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_2 : f32) outs(%1 : tensor<128x128xf32>) -> tensor<128x128xf32> + %4 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128xf32> + %5 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_3 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + %6 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_4 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + scf.for %arg16 = %arg13 to %c65536_i32 step %c28_i32 : i32 { + %alloc = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + annotation.mark %alloc {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + %alloc_5 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32, ssbuffer.crossDeps = [0 : i32, 1 : i32]} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_5 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + hivm.hir.sync_block_set {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + %alloc_6 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32, ssbuffer.crossDeps = [1 : i32, 1 : i32]} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_6 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<2>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + hivm.hir.sync_block_set {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + %alloc_7 = memref.alloc() : memref<128xf32, #hivm.address_space> + %memspacecast = memref.memory_space_cast %alloc_7 {ssbuffer.intraDeps = [0 : i32, 1 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %7 = bufferization.to_tensor %memspacecast writable : memref<128xf32> to tensor<128xf32> + %alloc_8 = memref.alloc() : memref<128xf32, #hivm.address_space> + %memspacecast_9 = memref.memory_space_cast %alloc_8 {ssbuffer.intraDeps = [0 : i32, 1 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %8 = bufferization.to_tensor %memspacecast_9 writable : memref<128xf32> to tensor<128xf32> + // CHECK: %{{.+}} = arith.subi %c8192_i32 + // CHECK: %{{.+}} = arith.ceildivui + // CHECK: %{{.+}} = arith.muli + // CHECK: %{{.+}} = arith.addi + // CHECK: %{{.+}} = arith.muli %c128_i32 + %9:3 = scf.for %arg17 = %c0_i32 to %c8192_i32 step %c128_i32 iter_args(%arg18 = %6, %arg19 = %2, %arg20 = %5) -> (tensor<128xf32>, tensor<128x128xf32>, tensor<128xf32>) : i32 { + hivm.hir.sync_block_wait {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + %memspacecast_11 = memref.memory_space_cast %alloc_5 {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32, ssbuffer.crossDeps = [0 : i32, 0 : i32]} : memref<128x128xf32, #hivm.address_space> to memref<128x128xf32> + %33 = bufferization.to_tensor %memspacecast_11 restrict writable {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32> to tensor<128x128xf32> + %34 = arith.mulf %33, %3 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %reduced = linalg.reduce ins(%34 : tensor<128x128xf32>) outs(%5 : tensor<128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + (%in: f32, %init: f32) { + %55 = arith.maximumf %in, %init : f32 + linalg.yield %55 : f32 + } + %35 = arith.maximumf %arg20, %reduced {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %broadcasted_12 = linalg.broadcast ins(%35 : tensor<128xf32>) outs(%1 : tensor<128x128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + %36 = arith.subf %34, %broadcasted_12 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %37 = math.exp %36 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %38 = arith.truncf %37 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> to tensor<128x128xf16> + %reshape = tensor.reshape %38(%cst_0) {ssbuffer.block_id = 5 : i32} : (tensor<128x128xf16>, tensor<3xi64>) -> tensor<128x8x16xf16> + annotation.mark %reshape {ssbuffer.block_id = 5 : i32, tiling_dim_mapping = {"1" = 1 : index}} : tensor<128x8x16xf16> + %39 = tensor.empty() {ssbuffer.block_id = 5 : i32} : tensor<8x128x16xf16> + %transposed = linalg.transpose ins(%reshape : tensor<128x8x16xf16>) outs(%39 : tensor<8x128x16xf16>) permutation = [1, 0, 2] {ssbuffer.block_id = 5 : i32} + %reshape_13 = tensor.reshape %transposed(%cst) {ssbuffer.block_id = 5 : i32} : (tensor<8x128x16xf16>, tensor<4xi64>) -> tensor<8x8x16x16xf16> + annotation.mark %reshape_13 {ssbuffer.block_id = 5 : i32, tiling_dim_mapping = {"1" = 1 : index}} : tensor<8x8x16x16xf16> + %40 = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst_1 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + %reduced_14 = linalg.reduce ins(%37 : tensor<128x128xf32>) outs(%40 : tensor<128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + (%in: f32, %init: f32) { + %55 = arith.addf %in, %init : f32 + linalg.yield %55 : f32 + } + %41 = arith.subf %arg20, %35 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %42 = math.exp %41 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %43 = arith.divui %arg17, %c128_i32 {ssbuffer.block_id = 5 : i32} : i32 + %c0_i32_15 = arith.constant {ssbuffer.block_id = 5 : i32} 0 : i32 + %c2_i32 = arith.constant {ssbuffer.block_id = 5 : i32} 2 : i32 + %44 = arith.remsi %43, %c2_i32 {ssbuffer.block_id = 5 : i32} : i32 + %45 = arith.cmpi eq, %44, %c0_i32_15 {ssbuffer.block_id = 5 : i32} : i32 + scf.if %45 { + bufferization.materialize_in_destination %42 in restrict writable %memspacecast {ssbuffer.block_id = 5 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + } else { + bufferization.materialize_in_destination %42 in restrict writable %memspacecast_9 {ssbuffer.block_id = 5 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + } {ssbuffer.block_id = 5 : i32} + %46 = arith.mulf %arg18, %42 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %47 = arith.addf %46, %reduced_14 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + hivm.hir.sync_block_set {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + hivm.hir.sync_block_wait {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32}[, , ] flag = 1 + hivm.hir.copy ins(%reshape_13 : tensor<8x8x16x16xf16>) outs(%alloc : memref<8x8x16x16xf16, #hivm.address_space>) {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32} + hivm.hir.sync_block_set {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32}[, , ] flag = 1 + hivm.hir.sync_block_wait {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + %memspacecast_16 = memref.memory_space_cast %alloc_6 {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32, ssbuffer.crossDeps = [1, 0]} : memref<128x128xf32, #hivm.address_space> to memref<128x128xf32> + %48 = bufferization.to_tensor %memspacecast_16 restrict writable {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32> to tensor<128x128xf32> + %49 = arith.divui %arg17, %c128_i32 {ssbuffer.block_id = 6 : i32} : i32 + %c2_i32_17 = arith.constant {ssbuffer.block_id = 6 : i32} 2 : i32 + %50 = arith.remsi %49, %c2_i32_17 {ssbuffer.block_id = 6 : i32} : i32 + %c0_i32_18 = arith.constant {ssbuffer.block_id = 6 : i32} 0 : i32 + %51 = arith.cmpi eq, %50, %c0_i32_18 {ssbuffer.block_id = 6 : i32} : i32 + %52 = scf.if %51 -> (tensor<128xf32>) { + %55 = bufferization.to_tensor %memspacecast restrict writable {ssbuffer.block_id = 6 : i32} : memref<128xf32> to tensor<128xf32> + scf.yield %55 : tensor<128xf32> + } else { + %55 = bufferization.to_tensor %memspacecast_9 restrict writable {ssbuffer.block_id = 6 : i32} : memref<128xf32> to tensor<128xf32> + scf.yield %55 : tensor<128xf32> + } {ssbuffer.block_id = 6 : i32, ssbuffer.intraDeps = [0 : i32, 0 : i32]} + %broadcasted_19 = linalg.broadcast ins(%52 : tensor<128xf32>) outs(%1 : tensor<128x128xf32>) dimensions = [1] {ssbuffer.block_id = 6 : i32} + %53 = arith.mulf %arg19, %broadcasted_19 {ssbuffer.block_id = 6 : i32} : tensor<128x128xf32> + %54 = arith.addf %48, %53 {ssbuffer.block_id = 6 : i32} : tensor<128x128xf32> + hivm.hir.sync_block_set {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + scf.yield %47, %54, %35 : tensor<128xf32>, tensor<128x128xf32>, tensor<128xf32> + } {ssbuffer.block_id = 9 : i32, ssbuffer.main_loop = 0 : i64} + hivm.hir.sync_block_wait {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32}[, , ] flag = 1 + %10 = arith.divsi %arg16, %c64_i32 {ssbuffer.block_id = 7 : i32} : i32 + %11 = arith.remsi %arg16, %c64_i32 {ssbuffer.block_id = 7 : i32} : i32 + %12 = arith.divsi %10, %c8_i32 {ssbuffer.block_id = 7 : i32} : i32 + %13 = arith.remsi %10, %c8_i32 {ssbuffer.block_id = 7 : i32} : i32 + %14 = arith.extsi %12 {ssbuffer.block_id = 7 : i32} : i32 to i64 + %15 = arith.muli %14, %c8388608_i64 {ssbuffer.block_id = 7 : i32} : i64 + %16 = arith.extsi %13 {ssbuffer.block_id = 7 : i32} : i32 to i64 + %17 = arith.muli %16, %c1048576_i64 {ssbuffer.block_id = 7 : i32} : i64 + %18 = arith.addi %15, %17 {ssbuffer.block_id = 7 : i32} : i64 + %19 = arith.index_cast %18 {ssbuffer.block_id = 7 : i32} : i64 to index + %20 = arith.muli %11, %c128_i32 {ssbuffer.block_id = 7 : i32} : i32 + %21 = arith.maxsi %20, %c0_i32 {ssbuffer.block_id = 7 : i32} : i32 + %22 = arith.index_cast %21 {ssbuffer.block_id = 7 : i32} : i32 to index + %23 = arith.muli %22, %c128 {ssbuffer.block_id = 7 : i32} : index + %24 = arith.addi %23, %19 {ssbuffer.block_id = 7 : i32} : index + %reinterpret_cast = memref.reinterpret_cast %arg7 to offset: [%24], sizes: [128, 128], strides: [128, 1] {ssbuffer.block_id = 7 : i32} : memref to memref<128x128xf16, strided<[128, 1], offset: ?>> + %25 = math.log %9#0 {ssbuffer.block_id = 7 : i32} : tensor<128xf32> + %26 = arith.addf %9#2, %25 {ssbuffer.block_id = 7 : i32} : tensor<128xf32> + %broadcasted = linalg.broadcast ins(%9#0 : tensor<128xf32>) outs(%1 : tensor<128x128xf32>) dimensions = [1] {ssbuffer.block_id = 7 : i32} + %27 = arith.divf %9#1, %broadcasted {ssbuffer.block_id = 7 : i32} : tensor<128x128xf32> + %28 = arith.muli %10, %c8192_i32 {ssbuffer.block_id = 7 : i32} : i32 + %29 = arith.index_cast %28 {ssbuffer.block_id = 7 : i32} : i32 to index + %30 = arith.index_cast %20 {ssbuffer.block_id = 7 : i32} : i32 to index + %31 = arith.addi %29, %30 {ssbuffer.block_id = 7 : i32} : index + %reinterpret_cast_10 = memref.reinterpret_cast %arg6 to offset: [%31], sizes: [128], strides: [1] {ssbuffer.block_id = 7 : i32} : memref to memref<128xf32, strided<[1], offset: ?>> + bufferization.materialize_in_destination %26 in writable %reinterpret_cast_10 {ssbuffer.block_id = 7 : i32} : (tensor<128xf32>, memref<128xf32, strided<[1], offset: ?>>) -> () + %32 = arith.truncf %27 {ssbuffer.block_id = 7 : i32} : tensor<128x128xf32> to tensor<128x128xf16> + bufferization.materialize_in_destination %32 in writable %reinterpret_cast {ssbuffer.block_id = 7 : i32} : (tensor<128x128xf16>, memref<128x128xf16, strided<[128, 1], offset: ?>>) -> () + } {ssbuffer.block_id = 10 : i32} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + scope.scope : () -> () { + %alloc = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + annotation.mark %alloc {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + %alloc_5 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_5 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + %alloc_6 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_6 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<2>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return {ssbuffer.core_type = "VECTOR"} + } +} diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/clone-ops.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/clone-ops.mlir new file mode 100644 index 0000000000..292d7d45a3 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/clone-ops.mlir @@ -0,0 +1,129 @@ +// RUN: triton-opt --clone-ops %s --allow-unregistered-dialect | FileCheck %s + +// CHECK: func.func @test_clone_ops +module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} { + func.func @test_clone_ops(%arg0: memref, %arg1: memref, %arg2: memref, %arg3: memref, %arg4: memref, %arg5: memref, %arg6: memref, %arg7: memref, %arg8: memref, %arg9: memref, %arg10: f32, %arg11: f32, %arg12: i32, %arg13: i32, %arg14: i32, %arg15: i32, %arg16: i32, %arg17: i32, %arg18: i32, %arg19: i32, %arg20: i32) attributes {global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} { + %c0_i64 = arith.constant {ssbuffer.block_id = 17 : i32} 0 : i64 + %c1_i64 = arith.constant {ssbuffer.block_id = 17 : i32} 1 : i64 + %c8_i64 = arith.constant {ssbuffer.block_id = 18 : i32} 8 : i64 + %c63_i32 = arith.constant {ssbuffer.block_id = 17 : i32} 63 : i32 + %c64_i32 = arith.constant {ssbuffer.block_id = 17 : i32} 64 : i32 + %c256_i32 = arith.constant {ssbuffer.block_id = 17 : i32} 256 : i32 + %c256_i64 = arith.constant {ssbuffer.block_id = 17 : i32} 256 : i64 + %c0_i32 = arith.constant {ssbuffer.block_id = 17 : i32} 0 : i32 + %c3_i32 = arith.constant {ssbuffer.block_id = 17 : i32} 3 : i32 + %c1_i32 = arith.constant {ssbuffer.block_id = 17 : i32} 1 : i32 + %c2_i32 = arith.constant {ssbuffer.block_id = 17 : i32} 2 : i32 + %c0 = arith.constant {ssbuffer.block_id = 17 : i32} 0 : index + %c1 = arith.constant {ssbuffer.block_id = 17 : i32} 1 : index + %c256 = arith.constant {ssbuffer.block_id = 17 : i32} 256 : index + scope.scope : () -> () { + %0 = tensor.empty() {ssbuffer.block_id = 17 : i32} : tensor<256x64xf32> + %1 = linalg.fill {ssbuffer.block_id = 17 : i32} ins(%arg11 : f32) outs(%0 : tensor<256x64xf32>) -> tensor<256x64xf32> + %2 = linalg.fill {ssbuffer.block_id = 17 : i32} ins(%arg11 : f32) outs(%0 : tensor<256x64xf32>) -> tensor<256x64xf32> + %3 = arith.cmpi sle, %arg12, %c256_i32 {ssbuffer.block_id = 17 : i32} : i32 + %4 = scf.if %3 -> (i64) { + %ext = arith.extsi %c2_i32 {ssbuffer.block_id = 17 : i32} : i32 to i64 + scf.yield %ext : i64 + } else { + %reinterpret_cast = memref.reinterpret_cast %arg7 to offset: [2], sizes: [1], strides: [1] {ssbuffer.block_id = 7 : i32} : memref to memref<1xi32, strided<[1], offset: 2>> + %24 = memref.load %reinterpret_cast[%c0] {ssbuffer.block_id = 7 : i32} : memref<1xi32, strided<[1], offset: 2>> + %25 = arith.extsi %24 {ssbuffer.block_id = 7 : i32} : i32 to i64 + scf.yield %25 : i64 + } {ssbuffer.block_id = 19 : i32} + %5 = arith.muli %4, %c8_i64 {ssbuffer.block_id = 18 : i32} : i64 + %6 = arith.extsi %arg15 {ssbuffer.block_id = 18 : i32} : i32 to i64 + %7 = arith.minsi %6, %5 {ssbuffer.block_id = 18 : i32} : i64 + %8 = arith.divsi %5, %7 {ssbuffer.block_id = 18 : i32} : i64 + %9 = arith.addi %8, %c1_i64 {ssbuffer.block_id = 18 : i32} : i64 + %10 = arith.remsi %5, %7 {ssbuffer.block_id = 18 : i32} : i64 + %11 = arith.extsi %arg18 {ssbuffer.block_id = 18 : i32} : i32 to i64 + %12 = arith.cmpi slt, %11, %7 {ssbuffer.block_id = 18 : i32} : i64 + %13 = arith.cmpi slt, %11, %10 {ssbuffer.block_id = 18 : i32} : i64 + %14 = arith.muli %11, %9 {ssbuffer.block_id = 18 : i32} : i64 + %15 = arith.muli %10, %9 {ssbuffer.block_id = 18 : i32} : i64 + %16 = arith.subi %11, %10 {ssbuffer.block_id = 18 : i32} : i64 + %17 = arith.muli %16, %8 {ssbuffer.block_id = 18 : i32} : i64 + %18 = arith.addi %15, %17 {ssbuffer.block_id = 18 : i32} : i64 + %19 = arith.select %13, %14, %18 {ssbuffer.block_id = 18 : i32} : i64 + %20 = arith.select %12, %19, %c0_i64 {ssbuffer.block_id = 18 : i32} : i64 + %21 = arith.select %13, %9, %8 {ssbuffer.block_id = 18 : i32} : i64 + %22 = arith.select %12, %21, %c0_i64 {ssbuffer.block_id = 18 : i32} : i64 + %23 = arith.cmpi sge, %11, %7 {ssbuffer.block_id = 18 : i32} : i64 + scf.if %23 { + } else { + %25 = arith.addi %arg13, %c63_i32 {ssbuffer.block_id = 15 : i32} : i32 + %26 = arith.divsi %25, %c64_i32 {ssbuffer.block_id = 15 : i32} : i32 + %27 = arith.extsi %26 {ssbuffer.block_id = 15 : i32} : i32 to i64 + %28 = arith.muli %22, %27 {ssbuffer.block_id = 15 : i32} : i64 + %29 = linalg.fill {ssbuffer.block_id = 15 : i32} ins(%arg10 : f32) outs(%0 : tensor<256x64xf32>) -> tensor<256x64xf32> + %30 = linalg.fill {ssbuffer.block_id = 15 : i32} ins(%arg11 : f32) outs(%0 : tensor<256x64xf32>) -> tensor<256x64xf32> + %alloc = memref.alloc() {ssbuffer.block_id = 26 : i32, ssbuffer.transfer_id = 0 : i32} : memref<4x16x16x16xf16, #hivm.address_space> + %alloc_4 = memref.alloc() {ssbuffer.block_id = 26 : i32, ssbuffer.transfer_id = 1 : i32} : memref<256x64xf32, #hivm.address_space> + %alloc_5 = memref.alloc() {ssbuffer.block_id = 26 : i32, ssbuffer.transfer_id = 2 : i32} : memref<256x64xf32, #hivm.address_space> + scf.for %arg21 = %c0_i64 to %28 step %c1_i64 : i64 { + %31 = arith.divsi %arg21, %27 {ssbuffer.block_id = 11 : i32, ssbuffer.dep_mark = [14 : i32]} : i64 + %32 = arith.addi %20, %31 {ssbuffer.block_id = 11 : i32, ssbuffer.dep_mark = [5 : i32]} : i64 + %33 = arith.remsi %32, %4 {ssbuffer.block_id = 11 : i32, ssbuffer.dep_mark = [3 : i32, 4 : i32]} : i64 + // Test: op within the region depends on the preceding op in another block + // CHECK: arith.divsi {{.*}} {ssbuffer.block_id = 22 : i32, ssbuffer.clone = 11 : i32, ssbuffer.dep_mark = [14 : i32]} + // CHECK: arith.addi {{.*}} {ssbuffer.block_id = 22 : i32, ssbuffer.clone = 11 : i32, ssbuffer.dep_mark = [5 : i32]} + // CHECK: arith.remsi {{.*}} {ssbuffer.block_id = 22 : i32, ssbuffer.clone = 11 : i32, ssbuffer.dep_mark = [3 : i32, 4 : i32]} + %34:2 = scf.if %3 -> (i64, i64) { + scf.yield %33, %c0_i64 : i64, i64 + } else { + %66:2 = scf.for %arg22 = %c0_i32 to %c2_i32 step %c1_i32 iter_args(%arg23 = %c0_i32, %arg24 = %c3_i32) -> (i32, i32) : i32 { + %73 = arith.addi %arg23, %arg24 {ssbuffer.block_id = 9 : i32} : i32 + %74 = arith.divsi %73, %c2_i32 {ssbuffer.block_id = 9 : i32} : i32 + %75 = arith.index_cast %74 {ssbuffer.block_id = 9 : i32} : i32 to index + %reinterpret_cast_11 = memref.reinterpret_cast %arg7 to offset: [%75], sizes: [1], strides: [1] {ssbuffer.block_id = 9 : i32} : memref to memref<1xi32, strided<[1], offset: ?>> + %76 = memref.load %reinterpret_cast_11[%c0] {ssbuffer.block_id = 9 : i32} : memref<1xi32, strided<[1], offset: ?>> + %77 = arith.extsi %76 {ssbuffer.block_id = 9 : i32} : i32 to i64 + %78 = arith.cmpi sle, %77, %33 {ssbuffer.block_id = 9 : i32, ssbuffer.dep_mark = [3 : i32, 4 : i32]} : i64 + %79 = arith.select %78, %arg24, %74 {ssbuffer.block_id = 9 : i32} : i32 + %80 = scf.if %78 -> (i32) { + %81 = arith.addi %74, %c1_i32 {ssbuffer.block_id = 8 : i32} : i32 + scf.yield %81 : i32 + } else { + scf.yield %arg23 : i32 + } {ssbuffer.block_id = 20 : i32} + scf.yield %80, %79 : i32, i32 + } {ssbuffer.block_id = 21 : i32} + %67 = arith.subi %66#0, %c1_i32 {ssbuffer.block_id = 10 : i32} : i32 + %68 = arith.extsi %67 {ssbuffer.block_id = 10 : i32} : i32 to i64 + %69 = arith.index_cast %67 {ssbuffer.block_id = 10 : i32} : i32 to index + %reinterpret_cast_10 = memref.reinterpret_cast %arg7 to offset: [%69], sizes: [1], strides: [1] {ssbuffer.block_id = 10 : i32} : memref to memref<1xi32, strided<[1], offset: ?>> + %70 = memref.load %reinterpret_cast_10[%c0] {ssbuffer.block_id = 10 : i32} : memref<1xi32, strided<[1], offset: ?>> + %71 = arith.extsi %70 {ssbuffer.block_id = 10 : i32} : i32 to i64 + %72 = arith.subi %33, %71 {ssbuffer.block_id = 10 : i32, ssbuffer.dep_mark = [3 : i32, 4 : i32]} : i64 + scf.yield %68, %72 : i64, i64 + } {ssbuffer.block_id = 22 : i32, ssbuffer.dep_mark = [1 : i32, 2 : i32]} + // Test: op depends on the preceding region op in another block + // CHECK: arith.divsi {{.*}} {ssbuffer.block_id = 12 : i32, ssbuffer.clone = 11 : i32, ssbuffer.dep_mark = [14 : i32]} + // CHECK: arith.addi {{.*}} {ssbuffer.block_id = 12 : i32, ssbuffer.clone = 11 : i32, ssbuffer.dep_mark = [5 : i32]} + // CHECK: arith.remsi {{.*}} {ssbuffer.block_id = 12 : i32, ssbuffer.clone = 11 : i32, ssbuffer.dep_mark = [3 : i32, 4 : i32]} + // CHECK: } {ssbuffer.block_id = 12 : i32, ssbuffer.clone = 22 : i32, ssbuffer.dep_mark = [1 : i32, 2 : i32]} + %35 = arith.index_cast %34#0 {ssbuffer.block_id = 12 : i32, ssbuffer.dep_mark = [1 : i32]} : i64 to index + %reinterpret_cast = memref.reinterpret_cast %arg5 to offset: [%35], sizes: [1], strides: [1] {ssbuffer.block_id = 12 : i32} : memref to memref<1xi64, strided<[1], offset: ?>> + %36 = memref.load %reinterpret_cast[%c0] {ssbuffer.block_id = 12 : i32, ssbuffer.dep_mark = [6 : i32]} : memref<1xi64, strided<[1], offset: ?>> + %37 = arith.addi %35, %c1 {ssbuffer.block_id = 12 : i32} : index + %reinterpret_cast_6 = memref.reinterpret_cast %arg5 to offset: [%37], sizes: [1], strides: [1] {ssbuffer.block_id = 12 : i32} : memref to memref<1xi64, strided<[1], offset: ?>> + %38 = memref.load %reinterpret_cast_6[%c0] {ssbuffer.block_id = 12 : i32} : memref<1xi64, strided<[1], offset: ?>> + %39 = arith.subi %38, %36 {ssbuffer.block_id = 12 : i32} : i64 + %40 = arith.muli %34#1, %c256_i64 {ssbuffer.block_id = 12 : i32, ssbuffer.dep_mark = [2 : i32]} : i64 + %41 = arith.index_cast %40 {ssbuffer.block_id = 12 : i32, ssbuffer.dep_mark = [7 : i32]} : i64 to index + %42 = arith.addi %41, %c256 {ssbuffer.block_id = 12 : i32} : index + %43 = arith.index_cast %39 {ssbuffer.block_id = 12 : i32} : i64 to index + %44 = arith.maxsi %41, %43 {ssbuffer.block_id = 12 : i32} : index + %45 = arith.minsi %42, %44 {ssbuffer.block_id = 12 : i32} : index + %46 = arith.subi %45, %41 {ssbuffer.block_id = 12 : i32, ssbuffer.dep_mark = [8 : i32]} : index + // Test: op depends on the preceding op in another block + // CHECK: arith.divsi {{.*}} {ssbuffer.block_id = 13 : i32, ssbuffer.clone = 11 : i32, ssbuffer.dep_mark = [14 : i32]} + %47 = arith.addi %20, %31 {ssbuffer.block_id = 13 : i32, ssbuffer.dep_mark = [14 : i32]} : i64 + } {ssbuffer.block_id = 26 : i32, ssbuffer.main_loop = 0 : i32} + } {ssbuffer.block_id = 27 : i32} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return {ssbuffer.core_type = "VECTOR"} + } +} diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/create-if-ops.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/create-if-ops.mlir new file mode 100644 index 0000000000..cfa0e92300 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/create-if-ops.mlir @@ -0,0 +1,146 @@ +// RUN: triton-opt --create-if-ops %s --allow-unregistered-dialect | FileCheck %s + +// Test create-if-ops pass wraps operations by block_id with scf.if %true +// CHECK: func.func @test_create_if_ops +module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} { + func.func @test_create_if_ops(%arg0: memref, %arg1: memref, %arg2: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg3: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg5: memref {tt.divisibility = 16 : i32}, %arg6: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg7: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg8: memref {tt.divisibility = 16 : i32}, %arg9: i32, %arg10: i32, %arg11: i32, %arg12: i32, %arg13: i32, %arg14: i32, %arg15: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} { + %cst = arith.constant {ssbuffer.block_id = 5 : i32} dense<[8, 8, 16, 16]> : tensor<4xi64> + %cst_0 = arith.constant {ssbuffer.block_id = 5 : i32} dense<[128, 8, 16]> : tensor<3xi64> + %cst_1 = arith.constant {ssbuffer.block_id = 8 : i32} 0.000000e+00 : f32 + %c28_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 28 : i32 + %c65536_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 65536 : i32 + %c8192_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 8192 : i32 + %c128_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 128 : i32 + %c1048576_i64 = arith.constant {ssbuffer.block_id = 8 : i32} 1048576 : i64 + %c8388608_i64 = arith.constant {ssbuffer.block_id = 8 : i32} 8388608 : i64 + %c8_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 8 : i32 + %c64_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 64 : i32 + %c0_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 0 : i32 + %cst_2 = arith.constant {ssbuffer.block_id = 8 : i32} 5.000000e-01 : f32 + %cst_3 = arith.constant {ssbuffer.block_id = 8 : i32} 0xFF800000 : f32 + %cst_4 = arith.constant {ssbuffer.block_id = 8 : i32} 1.000000e+00 : f32 + %c128 = arith.constant {ssbuffer.block_id = 8 : i32} 128 : index + scope.scope : () -> () { + %0 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128x128xf32> + %1 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128x128xf32> + %2 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_1 : f32) outs(%0 : tensor<128x128xf32>) -> tensor<128x128xf32> + %3 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_2 : f32) outs(%1 : tensor<128x128xf32>) -> tensor<128x128xf32> + %4 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128xf32> + %5 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_3 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + %6 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_4 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + scf.for %arg16 = %arg13 to %c65536_i32 step %c28_i32 : i32 { + %alloc = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + annotation.mark %alloc {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + %alloc_5 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.crossDeps = [0 : i32, 1 : i32], ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_5 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + hivm.hir.sync_block_set {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + %alloc_6 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.crossDeps = [1 : i32, 1 : i32], ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_6 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<2>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + hivm.hir.sync_block_set {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + %alloc_7 = memref.alloc() : memref<128xf32, #hivm.address_space> + %memspacecast = memref.memory_space_cast %alloc_7 {ssbuffer.intraDeps = [0 : i32, 1 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %7 = bufferization.to_tensor %memspacecast writable : memref<128xf32> to tensor<128xf32> + %alloc_8 = memref.alloc() : memref<128xf32, #hivm.address_space> + %memspacecast_9 = memref.memory_space_cast %alloc_8 {ssbuffer.intraDeps = [0 : i32, 1 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %8 = bufferization.to_tensor %memspacecast_9 writable : memref<128xf32> to tensor<128xf32> + %9:3 = scf.for %arg17 = %c0_i32 to %c8192_i32 step %c128_i32 iter_args(%arg18 = %6, %arg19 = %2, %arg20 = %5) -> (tensor<128xf32>, tensor<128x128xf32>, tensor<128xf32>) : i32 { + // CHECK: %true = arith.constant true + // CHECK: %{{.*}}:2 = scf.if %true -> (tensor<128xf32>, tensor<128xf32>) { + // CHECK: hivm.hir.sync_block_wait {ssbuffer.block_id = 5 : i32 + // CHECK: %13 = arith.mulf %12, %3 {ssbuffer.block_id = 5 : i32} + // CHECK: %16 = math.exp %15 {ssbuffer.block_id = 5 : i32} + // CHECK: %22 = arith.divui %arg17, %c128_i32 {ssbuffer.block_id = 5 : i32} + // CHECK: %c0_i32_14 = arith.constant {ssbuffer.block_id = 5 : i32} 0 : i32 + // CHECK: %c2_i32 = arith.constant {ssbuffer.block_id = 5 : i32} 2 : i32 + // CHECK: %23 = arith.remsi %22, %c2_i32 {ssbuffer.block_id = 5 : i32} + // CHECK: %24 = arith.cmpi eq, %23, %c0_i32_14 {ssbuffer.block_id = 5 : i32} + // CHECK: %25 = arith.mulf %arg18, %21 {ssbuffer.block_id = 5 : i32} + // CHECK: %26 = arith.addf %25, %reduced_13 {ssbuffer.block_id = 5 : i32} + // CHECK: hivm.hir.sync_block_set {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + // CHECK: hivm.hir.sync_block_wait {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32}[, , ] flag = 1 + // CHECK: hivm.hir.copy ins(%reshape_12 : tensor<8x8x16x16xf16>) outs(%alloc : memref<8x8x16x16xf16, #hivm.address_space>) {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32} + // CHECK: } {ssbuffer.if = 5 : i32} + hivm.hir.sync_block_wait {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + %memspacecast_11 = memref.memory_space_cast %alloc_5 {ssbuffer.block_id = 5 : i32, ssbuffer.crossDeps = [0 : i32, 0 : i32], ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> to memref<128x128xf32> + %33 = bufferization.to_tensor %memspacecast_11 restrict writable {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32> to tensor<128x128xf32> + %34 = arith.mulf %33, %3 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %reduced = linalg.reduce ins(%34 : tensor<128x128xf32>) outs(%5 : tensor<128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + (%in: f32, %init: f32) { + %55 = arith.maximumf %in, %init : f32 + linalg.yield %55 : f32 + } + %35 = arith.maximumf %arg20, %reduced {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %broadcasted_12 = linalg.broadcast ins(%35 : tensor<128xf32>) outs(%1 : tensor<128x128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + %36 = arith.subf %34, %broadcasted_12 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %37 = math.exp %36 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %38 = arith.truncf %37 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> to tensor<128x128xf16> + %reshape = tensor.reshape %38(%cst_0) {ssbuffer.block_id = 5 : i32} : (tensor<128x128xf16>, tensor<3xi64>) -> tensor<128x8x16xf16> + annotation.mark %reshape {ssbuffer.block_id = 5 : i32, tiling_dim_mapping = {"1" = 1 : index}} : tensor<128x8x16xf16> + %39 = tensor.empty() {ssbuffer.block_id = 5 : i32} : tensor<8x128x16xf16> + %transposed = linalg.transpose ins(%reshape : tensor<128x8x16xf16>) outs(%39 : tensor<8x128x16xf16>) permutation = [1, 0, 2] {ssbuffer.block_id = 5 : i32} + %reshape_13 = tensor.reshape %transposed(%cst) {ssbuffer.block_id = 5 : i32} : (tensor<8x128x16xf16>, tensor<4xi64>) -> tensor<8x8x16x16xf16> + annotation.mark %reshape_13 {ssbuffer.block_id = 5 : i32, tiling_dim_mapping = {"1" = 1 : index}} : tensor<8x8x16x16xf16> + %40 = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst_1 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + %reduced_14 = linalg.reduce ins(%37 : tensor<128x128xf32>) outs(%40 : tensor<128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + (%in: f32, %init: f32) { + %55 = arith.addf %in, %init : f32 + linalg.yield %55 : f32 + } + %41 = arith.subf %arg20, %35 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %42 = math.exp %41 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %43 = arith.divui %arg17, %c128_i32 {ssbuffer.block_id = 5 : i32} : i32 + %c0_i32_15 = arith.constant {ssbuffer.block_id = 5 : i32} 0 : i32 + %c2_i32 = arith.constant {ssbuffer.block_id = 5 : i32} 2 : i32 + %44 = arith.remsi %43, %c2_i32 {ssbuffer.block_id = 5 : i32} : i32 + %45 = arith.cmpi eq, %44, %c0_i32_15 {ssbuffer.block_id = 5 : i32} : i32 + scf.if %45 { + bufferization.materialize_in_destination %42 in restrict writable %memspacecast {ssbuffer.block_id = 5 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + } else { + bufferization.materialize_in_destination %42 in restrict writable %memspacecast_9 {ssbuffer.block_id = 5 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + } {ssbuffer.block_id = 5 : i32} + %46 = arith.mulf %arg18, %42 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %47 = arith.addf %46, %reduced_14 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + hivm.hir.sync_block_set {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + hivm.hir.sync_block_wait {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32}[, , ] flag = 1 + hivm.hir.copy ins(%reshape_13 : tensor<8x8x16x16xf16>) outs(%alloc : memref<8x8x16x16xf16, #hivm.address_space>) {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32} + hivm.hir.sync_block_set {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32}[, , ] flag = 1 + // CHECK: %true_10 = arith.constant true + // CHECK: %11 = scf.if %true_10 -> (tensor<128x128xf32>) { + // CHECK: hivm.hir.sync_block_wait {ssbuffer.block_id = 6 : i32 + // CHECK: %13 = arith.divui %arg17, %c128_i32 {ssbuffer.block_id = 6 : i32} + // CHECK: %c2_i32 = arith.constant {ssbuffer.block_id = 6 : i32} 2 : i32 + // CHECK: %14 = arith.remsi %13, %c2_i32 {ssbuffer.block_id = 6 : i32} + // CHECK: %c0_i32_12 = arith.constant {ssbuffer.block_id = 6 : i32} 0 : i32 + // CHECK: %15 = arith.cmpi eq, %14, %c0_i32_12 {ssbuffer.block_id = 6 : i32} + // CHECK: linalg.broadcast ins(%16 : tensor<128xf32>) outs(%1 : tensor<128x128xf32>) dimensions = [1] {ssbuffer.block_id = 6 : i32} + // CHECK: %17 = arith.mulf %arg19, %broadcasted {ssbuffer.block_id = 6 : i32} + // CHECK: %18 = arith.addf %12, %17 {ssbuffer.block_id = 6 : i32} + // CHECK: hivm.hir.sync_block_set {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + // CHECK: } {ssbuffer.if = 6 : i32} + hivm.hir.sync_block_wait {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + %memspacecast_16 = memref.memory_space_cast %alloc_6 {ssbuffer.block_id = 6 : i32, ssbuffer.crossDeps = [1, 0], ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> to memref<128x128xf32> + %48 = bufferization.to_tensor %memspacecast_16 restrict writable {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32> to tensor<128x128xf32> + %49 = arith.divui %arg17, %c128_i32 {ssbuffer.block_id = 6 : i32} : i32 + %c2_i32_17 = arith.constant {ssbuffer.block_id = 6 : i32} 2 : i32 + %50 = arith.remsi %49, %c2_i32_17 {ssbuffer.block_id = 6 : i32} : i32 + %c0_i32_18 = arith.constant {ssbuffer.block_id = 6 : i32} 0 : i32 + %51 = arith.cmpi eq, %50, %c0_i32_18 {ssbuffer.block_id = 6 : i32} : i32 + %52 = scf.if %51 -> (tensor<128xf32>) { + %55 = bufferization.to_tensor %memspacecast restrict writable {ssbuffer.block_id = 6 : i32} : memref<128xf32> to tensor<128xf32> + scf.yield %55 : tensor<128xf32> + } else { + %55 = bufferization.to_tensor %memspacecast_9 restrict writable {ssbuffer.block_id = 6 : i32} : memref<128xf32> to tensor<128xf32> + scf.yield %55 : tensor<128xf32> + } {ssbuffer.block_id = 6 : i32, ssbuffer.intraDeps = [0 : i32, 0 : i32]} + %broadcasted_19 = linalg.broadcast ins(%52 : tensor<128xf32>) outs(%1 : tensor<128x128xf32>) dimensions = [1] {ssbuffer.block_id = 6 : i32} + %53 = arith.mulf %arg19, %broadcasted_19 {ssbuffer.block_id = 6 : i32} : tensor<128x128xf32> + %54 = arith.addf %48, %53 {ssbuffer.block_id = 6 : i32} : tensor<128x128xf32> + hivm.hir.sync_block_set {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + scf.yield %47, %54, %35 : tensor<128xf32>, tensor<128x128xf32>, tensor<128xf32> + } {ssbuffer.block_id = 9 : i32, ssbuffer.main_loop = 0 : i64} + } {ssbuffer.block_id = 10 : i32} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return {ssbuffer.core_type = "VECTOR"} + } +} diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/crossCoreDeps_test.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/crossCoreDeps_test.mlir new file mode 100644 index 0000000000..86b68a5731 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/crossCoreDeps_test.mlir @@ -0,0 +1,140 @@ +// RUN: triton-opt --add-control-flow-condition %s | FileCheck %s + +// Unit Tests for AddControlFlowCondition Pass - Cross Core Dependencies Test +// +// This test verifies: +// 1. ssbuffer initialization at scope start (constants 0, 0, 1024, 4, 1028, 8, 1032) +// 2. if condition calculation before scf.if (llvm.load, arith.cmpi) +// 3. update operations at end of wrapped compute ops (llvm.load, arith.subi/arith.addi, llvm.store) + +module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} { + func.func @_attn_fwd(%arg0: memref, %arg1: memref, %arg2: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg3: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg5: memref {tt.divisibility = 16 : i32}, %arg6: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg7: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg8: memref {tt.divisibility = 16 : i32}, %arg9: i32, %arg10: i32, %arg11: i32, %arg12: i32, %arg13: i32, %arg14: i32, %arg15: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} { + %c8192_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 8192 : i32 + %c128_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 128 : i32 + %c0_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 0 : i32 + // ssbuffer init + // CHECK: %[[ZERO:.*]] = llvm.mlir.constant(0 : i32) : i32 + // CHECK: %[[ADDR0:.*]] = llvm.mlir.constant(0 : i64) : i64 + // CHECK: %[[ADDR1:.*]] = llvm.mlir.constant(1024 : i64) : i64 + // CHECK: %[[PTR0:.*]] = llvm.inttoptr %[[ADDR0]] : i64 to !llvm.ptr<11> + // CHECK: %[[PTR1:.*]] = llvm.inttoptr %[[ADDR1]] : i64 to !llvm.ptr<11> + // CHECK: llvm.store %[[ZERO]], %[[PTR0]] : i32, !llvm.ptr<11> + // CHECK: llvm.store %[[ZERO]], %[[PTR1]] : i32, !llvm.ptr<11> + // CHECK: %[[ADDR2:.*]] = llvm.mlir.constant(4 : i64) : i64 + // CHECK: %[[ADDR3:.*]] = llvm.mlir.constant(1028 : i64) : i64 + // CHECK: %[[PTR2:.*]] = llvm.inttoptr %[[ADDR2]] : i64 to !llvm.ptr<11> + // CHECK: %[[PTR3:.*]] = llvm.inttoptr %[[ADDR3]] : i64 to !llvm.ptr<11> + // CHECK: llvm.store %[[ZERO]], %[[PTR2]] : i32, !llvm.ptr<11> + // CHECK: llvm.store %[[ZERO]], %[[PTR3]] : i32, !llvm.ptr<11> + // CHECK: %[[ADDR4:.*]] = llvm.mlir.constant(8 : i64) : i64 + // CHECK: %[[ADDR5:.*]] = llvm.mlir.constant(1032 : i64) : i64 + // CHECK: %[[PTR4:.*]] = llvm.inttoptr %[[ADDR4]] : i64 to !llvm.ptr<11> + // CHECK: %[[PTR5:.*]] = llvm.inttoptr %[[ADDR5]] : i64 to !llvm.ptr<11> + // CHECK: llvm.store %[[ZERO]], %[[PTR4]] : i32, !llvm.ptr<11> + // CHECK: llvm.store %[[ZERO]], %[[PTR5]] : i32, !llvm.ptr<11> + scope.scope : () -> () { + // Vector scope only process each vector core's ssbuffer ptr + // CHECK: arith.constant 1024 + // CHECK: hivm.hir.get_sub_block_idx + // CHECK: arith.muli + // CHECK: arith.addi + // CHECK: llvm.inttoptr + %alloc = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + annotation.mark %alloc {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + %alloc_5 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32, ssbuffer.crossDeps = [0 : i32, 1 : i32]} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_5 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + %alloc_6 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32, ssbuffer.crossDeps = [1 : i32, 1 : i32]} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_6 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<2>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + %30:2 = scf.for %arg16 = %c0_i32 to %c8192_i32 step %c128_i32 iter_args(%arg17 = %c0_i32, %arg18 = %c0_i32) -> (i32, i32) : i32 { + // This block contain an input and a output + // CHECK: llvm.load + // CHECK: arith.cmpi sgt + // CHECK: llvm.load + // CHECK: arith.cmpi slt + // CHECK: scf.if + %memspacecast_11 = memref.memory_space_cast %alloc_5 {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32, ssbuffer.crossDeps = [0 : i32, 0 : i32]} : memref<128x128xf32, #hivm.address_space> to memref<128x128xf32> + %reshape_13 = tensor.empty() {ssbuffer.block_id = 5 : i32} : tensor<8x8x16x16xf16> + hivm.hir.copy ins(%reshape_13 : tensor<8x8x16x16xf16>) outs(%alloc : memref<8x8x16x16xf16, #hivm.address_space>) {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32} + // CHECK: llvm.load + // CHECK: arith.subi + // CHECK: llvm.store + // CHECK: llvm.load + // CHECK: arith.addi + // CHECK: llvm.store + // CHECK: } {ssbuffer.if = 5 : i32} + + // This block only contain an input + // CHECK: llvm.load + // CHECK: arith.cmpi sgt + // CHECK: scf.if + %memspacecast_16 = memref.memory_space_cast %alloc_6 {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32, ssbuffer.crossDeps = [1, 0]} : memref<128x128xf32, #hivm.address_space> to memref<128x128xf32> + %48 = bufferization.to_tensor %memspacecast_16 restrict writable {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32> to tensor<128x128xf32> + // CHECK: llvm.load + // CHECK: arith.subi + // CHECK: llvm.store + // CHECK: } {ssbuffer.if = 6 : i32} + scf.yield %arg17, %arg18 : i32, i32 + } {ssbuffer.block_id = 9 : i32, ssbuffer.main_loop = 0 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + scope.scope : () -> () { + %alloc_5 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32, ssbuffer.crossDeps = [2 : i32, 1 : i32]} : memref<8x8x16x16xf16, #hivm.address_space> + annotation.mark %alloc_5 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + %alloc_6 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_6 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + %alloc_7 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_7 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<2>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + %20:2 = scf.for %arg16 = %c0_i32 to %c8192_i32 step %c128_i32 iter_args(%arg17 = %c0_i32, %arg18 = %c0_i32) -> (i32, i32) : i32 { + // This block contain a output + // in CUBE scope, it need to process two ssbuffer ptr + // CHECK: llvm.load + // CHECK: llvm.load + // CHECK: arith.cmpi slt + // CHECK: arith.cmpi slt + // CHECK: scf.if + %30 = tensor.empty() {ssbuffer.block_id = 0 : i32} : tensor<128x128xf32> + hivm.hir.fixpipe {dma_mode = #hivm.dma_mode, ssbuffer.block_id = 0 : i32, ssbuffer.transfer_id = 1 : i32} ins(%30 : tensor<128x128xf32>) outs(%alloc_6 : memref<128x128xf32, #hivm.address_space>) + // CHECK: llvm.load + // CHECK: llvm.load + // CHECK: arith.addi + // CHECK: arith.addi + // CHECK: llvm.store + // CHECK: llvm.store + // CHECK: } {ssbuffer.if = 0 : i32} + + // This block contain an input and a output + // CHECK: llvm.load + // CHECK: llvm.load + // CHECK: arith.cmpi sgt + // CHECK: arith.cmpi sgt + // CHECK: llvm.load + // CHECK: llvm.load + // CHECK: arith.cmpi slt + // CHECK: arith.cmpi slt + // CHECK: scf.if + %32 = hivm.hir.convert_layout %alloc_5 output_shape [128, 128] {dstLayout = #hivm.data_layout, srcLayout = #hivm.data_layout, ssbuffer.block_id = 1 : i32, ssbuffer.transfer_id = 0 : i32, ssbuffer.crossDeps = [2 : i32, 0 : i32]} : (memref<8x8x16x16xf16, #hivm.address_space>) -> memref<128x128xf16, #hivm.address_space> + %40 = tensor.empty() {ssbuffer.block_id = 1 : i32} : tensor<128x128xf32> + hivm.hir.fixpipe {dma_mode = #hivm.dma_mode, ssbuffer.block_id = 1 : i32, ssbuffer.transfer_id = 2 : i32} ins(%40 : tensor<128x128xf32>) outs(%alloc_7 : memref<128x128xf32, #hivm.address_space>) + // CHECK: llvm.load + // CHECK: llvm.load + // CHECK: arith.subi + // CHECK: arith.subi + // CHECK: llvm.store + // CHECK: llvm.store + // CHECK: llvm.load + // CHECK: llvm.load + // CHECK: arith.addi + // CHECK: arith.addi + // CHECK: llvm.store + // CHECK: llvm.store + // CHECK: } {ssbuffer.if = 1 : i32} + scf.yield %arg17, %arg18 : i32, i32 + } {ssbuffer.block_id = 9 : i32, ssbuffer.main_loop = 0 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + scope.scope : () -> () { + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return {ssbuffer.core_type = "VECTOR"} + } +} diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/multi_deps_add_condition_test.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/multi_deps_add_condition_test.mlir new file mode 100644 index 0000000000..1ba1269192 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/multi_deps_add_condition_test.mlir @@ -0,0 +1,210 @@ +// RUN: triton-opt -add-control-flow-condition %s | FileCheck %s + +// CHECK-LABEL: module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} +// CHECK-LABEL: func.func @_attn_fwd( +// CHECK-DAG: [[C3:%[^ ]+]] = arith.constant 3 : i32 + +// CHECK-DAG: arith.cmpi slt +// CHECK-DAG: arith.cmpi slt +// CHECK-DAG: arith.cmpi slt +// CHECK-DAG: arith.cmpi slt +// CHECK-DAG: arith.cmpi slt +// CHECK-DAG: arith.cmpi slt +// CHECK-DAG: arith.cmpi slt +// CHECK-DAG: arith.cmpi slt + +// CHECK: [[AND1:%[^ ]+]] = arith.andi +// CHECK: [[AND2:%[^ ]+]] = arith.andi [[AND1]] +// CHECK: [[AND3:%[^ ]+]] = arith.andi [[AND2]] +// CHECK: [[AND4:%[^ ]+]] = arith.andi [[AND3]] +// CHECK: [[AND5:%[^ ]+]] = arith.andi [[AND4]] +// CHECK: [[AND6:%[^ ]+]] = arith.andi [[AND5]] +// CHECK: [[AND7:%[^ ]+]] = arith.andi [[AND6]] + +// CHECK: [[FINAL_CMP:%[^ ]+]] = arith.cmpi slt +// CHECK: [[AND8:%[^ ]+]] = arith.andi +// CHECK: [[COND:%[^ ]+]] = arith.andi [[AND8]] + +// CHECK: scf.if [[COND]] + + +// 匹配常量 1 +// CHECK: [[C1:%[^ ]+]] = arith.constant 1 : i32 +// CHECK-DAG: arith.addi %arg24, [[C1]] : i32 +// CHECK-DAG: arith.addi %arg25, [[C1]] : i32 +// CHECK-DAG: arith.addi %arg23, [[C1]] : i32 +// CHECK-DAG: arith.addi %arg27, [[C1]] : i32 +// CHECK-DAG: arith.addi %arg26, [[C1]] : i32 +// CHECK-DAG: arith.addi %arg28, [[C1]] : i32 +// CHECK-DAG: arith.addi %arg29, [[C1]] : i32 +// CHECK-DAG: arith.addi %arg30, [[C1]] : i32 +module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} { + func.func @_attn_fwd(%arg0: memref, %arg1: memref, %arg2: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg3: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg5: memref {tt.divisibility = 16 : i32}, %arg6: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg7: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg8: memref {tt.divisibility = 16 : i32}, %arg9: i32, %arg10: i32, %arg11: i32, %arg12: i32, %arg13: i32, %arg14: i32, %arg15: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} { + %cst = arith.constant {ssbuffer.block_id = 5 : i32} dense<[8, 8, 16, 16]> : tensor<4xi64> + %cst_0 = arith.constant {ssbuffer.block_id = 5 : i32} dense<[128, 8, 16]> : tensor<3xi64> + %cst_1 = arith.constant {ssbuffer.block_id = 8 : i32} 0.000000e+00 : f32 + %c28_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 28 : i32 + %c65536_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 65536 : i32 + %c8192_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 8192 : i32 + %c128_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 128 : i32 + %c1048576_i64 = arith.constant {ssbuffer.block_id = 8 : i32} 1048576 : i64 + %c8388608_i64 = arith.constant {ssbuffer.block_id = 8 : i32} 8388608 : i64 + %c8_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 8 : i32 + %c64_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 64 : i32 + %c0_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 0 : i32 + %cst_2 = arith.constant {ssbuffer.block_id = 8 : i32} 5.000000e-01 : f32 + %cst_3 = arith.constant {ssbuffer.block_id = 8 : i32} 0xFF800000 : f32 + %cst_4 = arith.constant {ssbuffer.block_id = 8 : i32} 1.000000e+00 : f32 + %c128 = arith.constant {ssbuffer.block_id = 8 : i32} 128 : index + scope.scope : () -> () { + %0 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128x128xf32> + %1 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128x128xf32> + %2 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_1 : f32) outs(%0 : tensor<128x128xf32>) -> tensor<128x128xf32> + %3 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_2 : f32) outs(%1 : tensor<128x128xf32>) -> tensor<128x128xf32> + %4 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128xf32> + %5 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_3 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + %6 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_4 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + scf.for %arg16 = %arg13 to %c65536_i32 step %c28_i32 : i32 { + %alloc = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + annotation.mark %alloc {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + %alloc_5 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32, ssbuffer.crossDeps = [0 : i32, 1 : i32]} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_5 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + hivm.hir.sync_block_set {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + %alloc_6 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32, ssbuffer.crossDeps = [1 : i32, 1 : i32]} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_6 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<2>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + hivm.hir.sync_block_set {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + %alloc_7 = memref.alloc() : memref<128xf32, #hivm.address_space> + %memspacecast = memref.memory_space_cast %alloc_7 {ssbuffer.intraDeps = [0 : i32, 1 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %7 = bufferization.to_tensor %memspacecast writable : memref<128xf32> to tensor<128xf32> + %alloc_8 = memref.alloc() : memref<128xf32, #hivm.address_space> + %memspacecast_9 = memref.memory_space_cast %alloc_8 {ssbuffer.intraDeps = [0 : i32, 1 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %8 = bufferization.to_tensor %memspacecast_9 writable : memref<128xf32> to tensor<128xf32> + %9:3 = scf.for %arg17 = %c0_i32 to %c8192_i32 step %c128_i32 iter_args(%arg18 = %6, %arg19 = %2, %arg20 = %5) -> (tensor<128xf32>, tensor<128x128xf32>, tensor<128xf32>) : i32 { + hivm.hir.sync_block_wait {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + %memspacecast_11 = memref.memory_space_cast %alloc_5 {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32, ssbuffer.crossDeps = [0 : i32, 0 : i32]} : memref<128x128xf32, #hivm.address_space> to memref<128x128xf32> + %33 = bufferization.to_tensor %memspacecast_11 restrict writable {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32> to tensor<128x128xf32> + %34 = arith.mulf %33, %3 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %reduced = linalg.reduce ins(%34 : tensor<128x128xf32>) outs(%5 : tensor<128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + (%in: f32, %init: f32) { + %55 = arith.maximumf %in, %init : f32 + linalg.yield %55 : f32 + } + %35 = arith.maximumf %arg20, %reduced {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %broadcasted_12 = linalg.broadcast ins(%35 : tensor<128xf32>) outs(%1 : tensor<128x128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + %36 = arith.subf %34, %broadcasted_12 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %37 = math.exp %36 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %38 = arith.truncf %37 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> to tensor<128x128xf16> + %reshape = tensor.reshape %38(%cst_0) {ssbuffer.block_id = 5 : i32} : (tensor<128x128xf16>, tensor<3xi64>) -> tensor<128x8x16xf16> + annotation.mark %reshape {ssbuffer.block_id = 5 : i32, tiling_dim_mapping = {"1" = 1 : index}} : tensor<128x8x16xf16> + %39 = tensor.empty() {ssbuffer.block_id = 5 : i32} : tensor<8x128x16xf16> + %transposed = linalg.transpose ins(%reshape : tensor<128x8x16xf16>) outs(%39 : tensor<8x128x16xf16>) permutation = [1, 0, 2] {ssbuffer.block_id = 5 : i32} + %reshape_13 = tensor.reshape %transposed(%cst) {ssbuffer.block_id = 5 : i32} : (tensor<8x128x16xf16>, tensor<4xi64>) -> tensor<8x8x16x16xf16> + annotation.mark %reshape_13 {ssbuffer.block_id = 5 : i32, tiling_dim_mapping = {"1" = 1 : index}} : tensor<8x8x16x16xf16> + %40 = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst_1 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + %reduced_14 = linalg.reduce ins(%37 : tensor<128x128xf32>) outs(%40 : tensor<128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + (%in: f32, %init: f32) { + %55 = arith.addf %in, %init : f32 + linalg.yield %55 : f32 + } + %41 = arith.subf %arg20, %35 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %42 = math.exp %41 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %43 = arith.divui %arg17, %c128_i32 {ssbuffer.block_id = 5 : i32} : i32 + %c0_i32_15 = arith.constant {ssbuffer.block_id = 5 : i32} 0 : i32 + %c2_i32 = arith.constant {ssbuffer.block_id = 5 : i32} 2 : i32 + %44 = arith.remsi %43, %c2_i32 {ssbuffer.block_id = 5 : i32} : i32 + %45 = arith.cmpi eq, %44, %c0_i32_15 {ssbuffer.block_id = 5 : i32} : i32 + scf.if %45 { + bufferization.materialize_in_destination %42 in restrict writable %memspacecast {ssbuffer.block_id = 5 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + } else { + bufferization.materialize_in_destination %42 in restrict writable %memspacecast_9 {ssbuffer.block_id = 5 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + } {ssbuffer.block_id = 5 : i32} + %46 = arith.mulf %arg18, %42 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %47 = arith.addf %46, %reduced_14 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + hivm.hir.sync_block_set {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + hivm.hir.sync_block_wait {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32}[, , ] flag = 1 + hivm.hir.copy ins(%reshape_13 : tensor<8x8x16x16xf16>) outs(%alloc : memref<8x8x16x16xf16, #hivm.address_space>) {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32} + hivm.hir.sync_block_set {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32}[, , ] flag = 1 + hivm.hir.sync_block_wait {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + %memspacecast_16 = memref.memory_space_cast %alloc_6 {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32, ssbuffer.crossDeps = [1, 0]} : memref<128x128xf32, #hivm.address_space> to memref<128x128xf32> + %48 = bufferization.to_tensor %memspacecast_16 restrict writable {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32> to tensor<128x128xf32> + %49 = arith.divui %arg17, %c128_i32 {ssbuffer.block_id = 6 : i32} : i32 + %c2_i32_17 = arith.constant {ssbuffer.block_id = 6 : i32} 2 : i32 + %50 = arith.remsi %49, %c2_i32_17 {ssbuffer.block_id = 6 : i32} : i32 + %c0_i32_18 = arith.constant {ssbuffer.block_id = 6 : i32} 0 : i32 + %51 = arith.cmpi eq, %50, %c0_i32_18 {ssbuffer.block_id = 6 : i32} : i32 + %producer_buf = memref.alloc() {ssbuffer.block_id = 6 : i32} : memref<128xf32, #hivm.address_space> + %producer_cast = memref.memory_space_cast %producer_buf {ssbuffer.block_id = 6 : i32, ssbuffer.intraDeps = [0 : i32, 1 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %producer_tensor = bufferization.to_tensor %producer_cast writable {ssbuffer.block_id = 6 : i32} : memref<128xf32> to tensor<128xf32> + + %c1_buf = memref.alloc() {ssbuffer.block_id = 6 : i32} : memref<128xf32, #hivm.address_space> + %c1_cast = memref.memory_space_cast %c1_buf {ssbuffer.block_id = 6 : i32, ssbuffer.intraDeps = [0 : i32, 0 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %c1_tensor = bufferization.to_tensor %c1_cast writable {ssbuffer.block_id = 6 : i32} : memref<128xf32> to tensor<128xf32> + %use1 = arith.mulf %producer_tensor, %producer_tensor {ssbuffer.block_id = 6 : i32, ssbuffer.intraDeps = [0 : i32, 0 : i32]} : tensor<128xf32> + + %c2_buf = memref.alloc() {ssbuffer.block_id = 6 : i32} : memref<128xf32, #hivm.address_space> + %c2_cast = memref.memory_space_cast %c2_buf {ssbuffer.block_id = 6 : i32, ssbuffer.intraDeps = [0 : i32, 0 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %c2_tensor = bufferization.to_tensor %c2_cast writable {ssbuffer.block_id = 6 : i32} : memref<128xf32> to tensor<128xf32> + %use2 = arith.addf %producer_tensor, %producer_tensor {ssbuffer.block_id = 6 : i32, ssbuffer.intraDeps = [0 : i32, 0 : i32]} : tensor<128xf32> + + %c3_buf = memref.alloc() {ssbuffer.block_id = 6 : i32} : memref<128xf32, #hivm.address_space> + %c3_cast = memref.memory_space_cast %c3_buf {ssbuffer.block_id = 6 : i32, ssbuffer.intraDeps = [0 : i32, 0 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %c3_tensor = bufferization.to_tensor %c3_cast writable {ssbuffer.block_id = 6 : i32} : memref<128xf32> to tensor<128xf32> + %use3 = arith.subf %producer_tensor, %producer_tensor {ssbuffer.block_id = 6 : i32, ssbuffer.intraDeps = [0 : i32, 0 : i32]} : tensor<128xf32> + + %52 = scf.if %51 -> (tensor<128xf32>) { + %55 = bufferization.to_tensor %memspacecast restrict writable {ssbuffer.block_id = 6 : i32, ssbuffer.intraDeps = [0 : i32, 0 : i32]} : memref<128xf32> to tensor<128xf32> + scf.yield %55 : tensor<128xf32> + } else { + %55 = bufferization.to_tensor %memspacecast_9 restrict writable {ssbuffer.block_id = 6 : i32, ssbuffer.intraDeps = [0 : i32, 0 : i32]} : memref<128xf32> to tensor<128xf32> + scf.yield %55 : tensor<128xf32> + } {ssbuffer.block_id = 6 : i32} + %broadcasted_19 = linalg.broadcast ins(%52 : tensor<128xf32>) outs(%1 : tensor<128x128xf32>) dimensions = [1] {ssbuffer.block_id = 6 : i32} + %53 = arith.mulf %arg19, %broadcasted_19 {ssbuffer.block_id = 6 : i32} : tensor<128x128xf32> + %54 = arith.addf %48, %53 {ssbuffer.block_id = 6 : i32} : tensor<128x128xf32> + hivm.hir.sync_block_set {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + scf.yield %47, %54, %35 : tensor<128xf32>, tensor<128x128xf32>, tensor<128xf32> + } {ssbuffer.block_id = 9 : i32, ssbuffer.main_loop = 0 : i64} + hivm.hir.sync_block_wait {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32}[, , ] flag = 1 + %10 = arith.divsi %arg16, %c64_i32 {ssbuffer.block_id = 7 : i32} : i32 + %11 = arith.remsi %arg16, %c64_i32 {ssbuffer.block_id = 7 : i32} : i32 + %12 = arith.divsi %10, %c8_i32 {ssbuffer.block_id = 7 : i32} : i32 + %13 = arith.remsi %10, %c8_i32 {ssbuffer.block_id = 7 : i32} : i32 + %14 = arith.extsi %12 {ssbuffer.block_id = 7 : i32} : i32 to i64 + %15 = arith.muli %14, %c8388608_i64 {ssbuffer.block_id = 7 : i32} : i64 + %16 = arith.extsi %13 {ssbuffer.block_id = 7 : i32} : i32 to i64 + %17 = arith.muli %16, %c1048576_i64 {ssbuffer.block_id = 7 : i32} : i64 + %18 = arith.addi %15, %17 {ssbuffer.block_id = 7 : i32} : i64 + %19 = arith.index_cast %18 {ssbuffer.block_id = 7 : i32} : i64 to index + %20 = arith.muli %11, %c128_i32 {ssbuffer.block_id = 7 : i32} : i32 + %21 = arith.maxsi %20, %c0_i32 {ssbuffer.block_id = 7 : i32} : i32 + %22 = arith.index_cast %21 {ssbuffer.block_id = 7 : i32} : i32 to index + %23 = arith.muli %22, %c128 {ssbuffer.block_id = 7 : i32} : index + %24 = arith.addi %23, %19 {ssbuffer.block_id = 7 : i32} : index + %reinterpret_cast = memref.reinterpret_cast %arg7 to offset: [%24], sizes: [128, 128], strides: [128, 1] {ssbuffer.block_id = 7 : i32} : memref to memref<128x128xf16, strided<[128, 1], offset: ?>> + %25 = math.log %9#0 {ssbuffer.block_id = 7 : i32} : tensor<128xf32> + %26 = arith.addf %9#2, %25 {ssbuffer.block_id = 7 : i32} : tensor<128xf32> + %broadcasted = linalg.broadcast ins(%9#0 : tensor<128xf32>) outs(%1 : tensor<128x128xf32>) dimensions = [1] {ssbuffer.block_id = 7 : i32} + %27 = arith.divf %9#1, %broadcasted {ssbuffer.block_id = 7 : i32} : tensor<128x128xf32> + %28 = arith.muli %10, %c8192_i32 {ssbuffer.block_id = 7 : i32} : i32 + %29 = arith.index_cast %28 {ssbuffer.block_id = 7 : i32} : i32 to index + %30 = arith.index_cast %20 {ssbuffer.block_id = 7 : i32} : i32 to index + %31 = arith.addi %29, %30 {ssbuffer.block_id = 7 : i32} : index + %reinterpret_cast_10 = memref.reinterpret_cast %arg6 to offset: [%31], sizes: [128], strides: [1] {ssbuffer.block_id = 7 : i32} : memref to memref<128xf32, strided<[1], offset: ?>> + bufferization.materialize_in_destination %26 in writable %reinterpret_cast_10 {ssbuffer.block_id = 7 : i32} : (tensor<128xf32>, memref<128xf32, strided<[1], offset: ?>>) -> () + %32 = arith.truncf %27 {ssbuffer.block_id = 7 : i32} : tensor<128x128xf32> to tensor<128x128xf16> + bufferization.materialize_in_destination %32 in writable %reinterpret_cast {ssbuffer.block_id = 7 : i32} : (tensor<128x128xf16>, memref<128x128xf16, strided<[128, 1], offset: ?>>) -> () + } {ssbuffer.block_id = 10 : i32} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + scope.scope : () -> () { + %alloc = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + annotation.mark %alloc {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + %alloc_5 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_5 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + %alloc_6 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_6 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<2>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + scope.return + }{hivm.tcore_type = #hivm.tcore_type} + return {ssbuffer.core_type = "VECTOR"} + } +} diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/no_deps_add_condition_test.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/no_deps_add_condition_test.mlir new file mode 100644 index 0000000000..cbba0366f6 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/no_deps_add_condition_test.mlir @@ -0,0 +1,112 @@ +// RUN: triton-opt -add-control-flow-condition --verify-diagnostics %s +module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} { // expected-error {{[AddControlFlowCondition] Pass failed!}} + func.func @_attn_fwd(%arg0: memref, %arg1: memref, %arg2: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg3: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg5: memref {tt.divisibility = 16 : i32}, %arg6: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg7: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg8: memref {tt.divisibility = 16 : i32}, %arg9: i32, %arg10: i32, %arg11: i32, %arg12: i32, %arg13: i32, %arg14: i32, %arg15: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} { + %cst = arith.constant {ssbuffer.block_id = 5 : i32} dense<[8, 8, 16, 16]> : tensor<4xi64> + %cst_0 = arith.constant {ssbuffer.block_id = 5 : i32} dense<[128, 8, 16]> : tensor<3xi64> + %cst_1 = arith.constant {ssbuffer.block_id = 8 : i32} 0.000000e+00 : f32 + %c28_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 28 : i32 + %c65536_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 65536 : i32 + %c8192_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 8192 : i32 + %c128_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 128 : i32 + %c0_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 0 : i32 + %cst_2 = arith.constant {ssbuffer.block_id = 8 : i32} 5.000000e-01 : f32 + %cst_3 = arith.constant {ssbuffer.block_id = 8 : i32} 0xFF800000 : f32 + %cst_4 = arith.constant {ssbuffer.block_id = 8 : i32} 1.000000e+00 : f32 + scope.scope : () -> () { + %0 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128x128xf32> + %1 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128x128xf32> + %2 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_1 : f32) outs(%0 : tensor<128x128xf32>) -> tensor<128x128xf32> + %3 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_2 : f32) outs(%1 : tensor<128x128xf32>) -> tensor<128x128xf32> + %4 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128xf32> + %5 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_3 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + %6 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_4 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + scf.for %arg16 = %arg13 to %c65536_i32 step %c28_i32 : i32 { + %alloc = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + annotation.mark %alloc {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + %alloc_5 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_5 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + hivm.hir.sync_block_set {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + %alloc_6 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_6 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<2>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + hivm.hir.sync_block_set {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + %alloc_7 = memref.alloc() : memref<128xf32, #hivm.address_space> + %memspacecast = memref.memory_space_cast %alloc_7 {} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %7 = bufferization.to_tensor %memspacecast writable : memref<128xf32> to tensor<128xf32> + %alloc_8 = memref.alloc() : memref<128xf32, #hivm.address_space> + %memspacecast_9 = memref.memory_space_cast %alloc_8 {} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %8 = bufferization.to_tensor %memspacecast_9 writable : memref<128xf32> to tensor<128xf32> + %9:3 = scf.for %arg17 = %c0_i32 to %c8192_i32 step %c128_i32 iter_args(%arg18 = %6, %arg19 = %2, %arg20 = %5) -> (tensor<128xf32>, tensor<128x128xf32>, tensor<128xf32>) : i32 { + hivm.hir.sync_block_wait {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + %memspacecast_11 = memref.memory_space_cast %alloc_5 {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> to memref<128x128xf32> + %33 = bufferization.to_tensor %memspacecast_11 restrict writable {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32> to tensor<128x128xf32> + %34 = arith.mulf %33, %3 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %reduced = linalg.reduce ins(%34 : tensor<128x128xf32>) outs(%5 : tensor<128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + (%in: f32, %init: f32) { + %55 = arith.maximumf %in, %init : f32 + linalg.yield %55 : f32 + } + %35 = arith.maximumf %arg20, %reduced {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %broadcasted_12 = linalg.broadcast ins(%35 : tensor<128xf32>) outs(%1 : tensor<128x128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + %36 = arith.subf %34, %broadcasted_12 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %37 = math.exp %36 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %38 = arith.truncf %37 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> to tensor<128x128xf16> + %reshape = tensor.reshape %38(%cst_0) {ssbuffer.block_id = 5 : i32} : (tensor<128x128xf16>, tensor<3xi64>) -> tensor<128x8x16xf16> + annotation.mark %reshape {ssbuffer.block_id = 5 : i32, tiling_dim_mapping = {"1" = 1 : index}} : tensor<128x8x16xf16> + %39 = tensor.empty() {ssbuffer.block_id = 5 : i32} : tensor<8x128x16xf16> + %transposed = linalg.transpose ins(%reshape : tensor<128x8x16xf16>) outs(%39 : tensor<8x128x16xf16>) permutation = [1, 0, 2] {ssbuffer.block_id = 5 : i32} + %reshape_13 = tensor.reshape %transposed(%cst) {ssbuffer.block_id = 5 : i32} : (tensor<8x128x16xf16>, tensor<4xi64>) -> tensor<8x8x16x16xf16> + annotation.mark %reshape_13 {ssbuffer.block_id = 5 : i32, tiling_dim_mapping = {"1" = 1 : index}} : tensor<8x8x16x16xf16> + %40 = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst_1 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + %reduced_14 = linalg.reduce ins(%37 : tensor<128x128xf32>) outs(%40 : tensor<128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + (%in: f32, %init: f32) { + %55 = arith.addf %in, %init : f32 + linalg.yield %55 : f32 + } + %41 = arith.subf %arg20, %35 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %42 = math.exp %41 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %43 = arith.divui %arg17, %c128_i32 {ssbuffer.block_id = 5 : i32} : i32 + %c0_i32_15 = arith.constant {ssbuffer.block_id = 5 : i32} 0 : i32 + %c2_i32 = arith.constant {ssbuffer.block_id = 5 : i32} 2 : i32 + %44 = arith.remsi %43, %c2_i32 {ssbuffer.block_id = 5 : i32} : i32 + %45 = arith.cmpi eq, %44, %c0_i32_15 {ssbuffer.block_id = 5 : i32} : i32 + scf.if %45 { + bufferization.materialize_in_destination %42 in restrict writable %memspacecast {ssbuffer.block_id = 5 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + } else { + bufferization.materialize_in_destination %42 in restrict writable %memspacecast_9 {ssbuffer.block_id = 5 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + } {ssbuffer.block_id = 5 : i32} + %46 = arith.mulf %arg18, %42 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %47 = arith.addf %46, %reduced_14 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + hivm.hir.sync_block_set {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + hivm.hir.sync_block_wait {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32}[, , ] flag = 1 + hivm.hir.copy ins(%reshape_13 : tensor<8x8x16x16xf16>) outs(%alloc : memref<8x8x16x16xf16, #hivm.address_space>) {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32} + hivm.hir.sync_block_set {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32}[, , ] flag = 1 + hivm.hir.sync_block_wait {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + %memspacecast_16 = memref.memory_space_cast %alloc_6 {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> to memref<128x128xf32> + %48 = bufferization.to_tensor %memspacecast_16 restrict writable {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32> to tensor<128x128xf32> + %49 = arith.divui %arg17, %c128_i32 {ssbuffer.block_id = 6 : i32} : i32 + %c2_i32_17 = arith.constant {ssbuffer.block_id = 6 : i32} 2 : i32 + %50 = arith.remsi %49, %c2_i32_17 {ssbuffer.block_id = 6 : i32} : i32 + %c0_i32_18 = arith.constant {ssbuffer.block_id = 6 : i32} 0 : i32 + %51 = arith.cmpi eq, %50, %c0_i32_18 {ssbuffer.block_id = 6 : i32} : i32 + %52 = scf.if %51 -> (tensor<128xf32>) { + %55 = bufferization.to_tensor %memspacecast restrict writable {ssbuffer.block_id = 6 : i32} : memref<128xf32> to tensor<128xf32> + scf.yield %55 : tensor<128xf32> + } else { + %55 = bufferization.to_tensor %memspacecast_9 restrict writable {ssbuffer.block_id = 6 : i32} : memref<128xf32> to tensor<128xf32> + scf.yield %55 : tensor<128xf32> + } {ssbuffer.block_id = 6 : i32} + %broadcasted_19 = linalg.broadcast ins(%52 : tensor<128xf32>) outs(%1 : tensor<128x128xf32>) dimensions = [1] {ssbuffer.block_id = 6 : i32} + %53 = arith.mulf %arg19, %broadcasted_19 {ssbuffer.block_id = 6 : i32} : tensor<128x128xf32> + %54 = arith.addf %48, %53 {ssbuffer.block_id = 6 : i32} : tensor<128x128xf32> + hivm.hir.sync_block_set {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + scf.yield %47, %54, %35 : tensor<128xf32>, tensor<128x128xf32>, tensor<128xf32> + } {ssbuffer.block_id = 9 : i32, ssbuffer.main_loop = 0 : i64} + } {ssbuffer.block_id = 10 : i32} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + scope.scope : () -> () { + scope.return + }{hivm.tcore_type = #hivm.tcore_type} + return {ssbuffer.core_type = "VECTOR"} + } +} diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/single_deps_add_condition_test.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/single_deps_add_condition_test.mlir new file mode 100644 index 0000000000..1a5fdc5b1c --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/single_deps_add_condition_test.mlir @@ -0,0 +1,138 @@ +// RUN: triton-opt -add-control-flow-condition %s | FileCheck %s + +// CHECK-LABEL: module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} +// CHECK-LABEL: func.func @_attn_fwd( +// CHECK: [[C2:%[^ ]+]] = arith.constant 2 +// CHECK: [[CMP1:%[^ ]+]] = arith.cmpi slt, %arg23, [[C2]] +// CHECK: [[CMP2:%[^ ]+]] = arith.cmpi slt, %arg21, %c8192_i32 +// CHECK: [[AND1:%[^ ]+]] = arith.andi {{%[^ ]+}}, [[CMP1]] +// CHECK-NEXT: [[AND2:%[^ ]+]] = arith.andi [[AND1]], [[CMP2]] +// CHECK-NEXT: [[IF_RES:%[^ ]+]]:4 = scf.if [[AND2]] +// CHECK: [[C1:%[^ ]+]] = arith.constant 1 : i32 +// CHECK: [[ADD3:%[^ ]+]] = arith.addi %arg23, [[C1]] : i32 +// CHECK: [[ADD4:%[^ ]+]] = arith.addi %arg21, %c128_i32 : i32 +// CHECK: scf.yield {{%[^ ]+}}, {{%[^ ]+}}, [[ADD3]], [[ADD4]] +// CHECK: scf.yield %arg20, %arg18, %arg23, %arg21 +// CHECK: [[C0:%[^ ]+]] = arith.constant 0 : i32 +// CHECK: [[CMP3:%[^ ]+]] = arith.cmpi sgt, [[IF_RES]]{{.*}}, +// CHECK-NEXT: [[CMP4:%[^ ]+]] = arith.cmpi slt, %arg22, %c8192_i32 : i32 +// CHECK-NEXT: [[AND5:%[^ ]+]] = arith.andi {{%[^ ]+}}, [[CMP3]] : i1 +// CHECK-NEXT: [[FINAL:%[^ ]+]] = arith.andi [[AND5]], [[CMP4]] : i1 +// CHECK-NEXT: {{%[^ ]+}}:3 = scf.if [[FINAL]] +module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} { + func.func @_attn_fwd(%arg0: memref, %arg1: memref, %arg2: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg3: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg5: memref {tt.divisibility = 16 : i32}, %arg6: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg7: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg8: memref {tt.divisibility = 16 : i32}, %arg9: i32, %arg10: i32, %arg11: i32, %arg12: i32, %arg13: i32, %arg14: i32, %arg15: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} { + %cst = arith.constant {ssbuffer.block_id = 5 : i32} dense<[8, 8, 16, 16]> : tensor<4xi64> + %cst_0 = arith.constant {ssbuffer.block_id = 5 : i32} dense<[128, 8, 16]> : tensor<3xi64> + %cst_1 = arith.constant {ssbuffer.block_id = 8 : i32} 0.000000e+00 : f32 + %c28_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 28 : i32 + %c65536_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 65536 : i32 + %c8192_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 8192 : i32 + %c128_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 128 : i32 + %c0_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 0 : i32 + %cst_2 = arith.constant {ssbuffer.block_id = 8 : i32} 5.000000e-01 : f32 + %cst_3 = arith.constant {ssbuffer.block_id = 8 : i32} 0xFF800000 : f32 + %cst_4 = arith.constant {ssbuffer.block_id = 8 : i32} 1.000000e+00 : f32 + scope.scope : () -> () { + %0 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128x128xf32> + %1 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128x128xf32> + %2 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_1 : f32) outs(%0 : tensor<128x128xf32>) -> tensor<128x128xf32> + %3 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_2 : f32) outs(%1 : tensor<128x128xf32>) -> tensor<128x128xf32> + %4 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128xf32> + %5 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_3 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + %6 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_4 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + scf.for %arg16 = %arg13 to %c65536_i32 step %c28_i32 : i32 { + %alloc = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + annotation.mark %alloc {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + %alloc_5 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32, ssbuffer.crossDeps = [0 : i32, 1 : i32]} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_5 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + hivm.hir.sync_block_set {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + %alloc_6 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32, ssbuffer.crossDeps = [1 : i32, 1 : i32]} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_6 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<2>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + hivm.hir.sync_block_set {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + %alloc_7 = memref.alloc() : memref<128xf32, #hivm.address_space> + %memspacecast = memref.memory_space_cast %alloc_7 {ssbuffer.intraDeps = [0 : i32, 1 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %7 = bufferization.to_tensor %memspacecast writable : memref<128xf32> to tensor<128xf32> + %alloc_8 = memref.alloc() : memref<128xf32, #hivm.address_space> + %memspacecast_9 = memref.memory_space_cast %alloc_8 {ssbuffer.intraDeps = [0 : i32, 1 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %8 = bufferization.to_tensor %memspacecast_9 writable : memref<128xf32> to tensor<128xf32> + %9:3 = scf.for %arg17 = %c0_i32 to %c8192_i32 step %c128_i32 iter_args(%arg18 = %6, %arg19 = %2, %arg20 = %5) -> (tensor<128xf32>, tensor<128x128xf32>, tensor<128xf32>) : i32 { + hivm.hir.sync_block_wait {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + %memspacecast_11 = memref.memory_space_cast %alloc_5 {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32, ssbuffer.crossDeps = [0 : i32, 0 : i32]} : memref<128x128xf32, #hivm.address_space> to memref<128x128xf32> + %33 = bufferization.to_tensor %memspacecast_11 restrict writable {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32> to tensor<128x128xf32> + %34 = arith.mulf %33, %3 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %reduced = linalg.reduce ins(%34 : tensor<128x128xf32>) outs(%5 : tensor<128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + (%in: f32, %init: f32) { + %55 = arith.maximumf %in, %init : f32 + linalg.yield %55 : f32 + } + %35 = arith.maximumf %arg20, %reduced {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %broadcasted_12 = linalg.broadcast ins(%35 : tensor<128xf32>) outs(%1 : tensor<128x128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + %36 = arith.subf %34, %broadcasted_12 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %37 = math.exp %36 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %38 = arith.truncf %37 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> to tensor<128x128xf16> + %reshape = tensor.reshape %38(%cst_0) {ssbuffer.block_id = 5 : i32} : (tensor<128x128xf16>, tensor<3xi64>) -> tensor<128x8x16xf16> + annotation.mark %reshape {ssbuffer.block_id = 5 : i32, tiling_dim_mapping = {"1" = 1 : index}} : tensor<128x8x16xf16> + %39 = tensor.empty() {ssbuffer.block_id = 5 : i32} : tensor<8x128x16xf16> + %transposed = linalg.transpose ins(%reshape : tensor<128x8x16xf16>) outs(%39 : tensor<8x128x16xf16>) permutation = [1, 0, 2] {ssbuffer.block_id = 5 : i32} + %reshape_13 = tensor.reshape %transposed(%cst) {ssbuffer.block_id = 5 : i32} : (tensor<8x128x16xf16>, tensor<4xi64>) -> tensor<8x8x16x16xf16> + annotation.mark %reshape_13 {ssbuffer.block_id = 5 : i32, tiling_dim_mapping = {"1" = 1 : index}} : tensor<8x8x16x16xf16> + %40 = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst_1 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + %reduced_14 = linalg.reduce ins(%37 : tensor<128x128xf32>) outs(%40 : tensor<128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + (%in: f32, %init: f32) { + %55 = arith.addf %in, %init : f32 + linalg.yield %55 : f32 + } + %41 = arith.subf %arg20, %35 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %42 = math.exp %41 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %43 = arith.divui %arg17, %c128_i32 {ssbuffer.block_id = 5 : i32} : i32 + %c0_i32_15 = arith.constant {ssbuffer.block_id = 5 : i32} 0 : i32 + %c2_i32 = arith.constant {ssbuffer.block_id = 5 : i32} 2 : i32 + %44 = arith.remsi %43, %c2_i32 {ssbuffer.block_id = 5 : i32} : i32 + %45 = arith.cmpi eq, %44, %c0_i32_15 {ssbuffer.block_id = 5 : i32} : i32 + scf.if %45 { + bufferization.materialize_in_destination %42 in restrict writable %memspacecast {ssbuffer.block_id = 5 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + } else { + bufferization.materialize_in_destination %42 in restrict writable %memspacecast_9 {ssbuffer.block_id = 5 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + } {ssbuffer.block_id = 5 : i32} + %46 = arith.mulf %arg18, %42 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %47 = arith.addf %46, %reduced_14 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + hivm.hir.sync_block_set {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + hivm.hir.sync_block_wait {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32}[, , ] flag = 1 + hivm.hir.copy ins(%reshape_13 : tensor<8x8x16x16xf16>) outs(%alloc : memref<8x8x16x16xf16, #hivm.address_space>) {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32} + hivm.hir.sync_block_set {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32}[, , ] flag = 1 + hivm.hir.sync_block_wait {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + %memspacecast_16 = memref.memory_space_cast %alloc_6 {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32, ssbuffer.crossDeps = [1, 0]} : memref<128x128xf32, #hivm.address_space> to memref<128x128xf32> + %48 = bufferization.to_tensor %memspacecast_16 restrict writable {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32> to tensor<128x128xf32> + %49 = arith.divui %arg17, %c128_i32 {ssbuffer.block_id = 6 : i32} : i32 + %c2_i32_17 = arith.constant {ssbuffer.block_id = 6 : i32} 2 : i32 + %50 = arith.remsi %49, %c2_i32_17 {ssbuffer.block_id = 6 : i32} : i32 + %c0_i32_18 = arith.constant {ssbuffer.block_id = 6 : i32} 0 : i32 + %51 = arith.cmpi eq, %50, %c0_i32_18 {ssbuffer.block_id = 6 : i32} : i32 + %52 = scf.if %51 -> (tensor<128xf32>) { + %55 = bufferization.to_tensor %memspacecast restrict writable {ssbuffer.block_id = 6 : i32} : memref<128xf32> to tensor<128xf32> + scf.yield %55 : tensor<128xf32> + } else { + %55 = bufferization.to_tensor %memspacecast_9 restrict writable {ssbuffer.block_id = 6 : i32} : memref<128xf32> to tensor<128xf32> + scf.yield %55 : tensor<128xf32> + } {ssbuffer.block_id = 6 : i32, ssbuffer.intraDeps = [0 : i32, 0 : i32]} + %broadcasted_19 = linalg.broadcast ins(%52 : tensor<128xf32>) outs(%1 : tensor<128x128xf32>) dimensions = [1] {ssbuffer.block_id = 6 : i32} + %53 = arith.mulf %arg19, %broadcasted_19 {ssbuffer.block_id = 6 : i32} : tensor<128x128xf32> + %54 = arith.addf %48, %53 {ssbuffer.block_id = 6 : i32} : tensor<128x128xf32> + hivm.hir.sync_block_set {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + scf.yield %47, %54, %35 : tensor<128xf32>, tensor<128x128xf32>, tensor<128xf32> + } {ssbuffer.block_id = 9 : i32, ssbuffer.main_loop = 0 : i64} + } {ssbuffer.block_id = 10 : i32} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + scope.scope : () -> () { + %alloc = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + annotation.mark %alloc {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + %alloc_5 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_5 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + %alloc_6 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_6 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<2>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + scope.return + }{hivm.tcore_type = #hivm.tcore_type} + return {ssbuffer.core_type = "VECTOR"} + } +} diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/skip-control-flow.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/skip-control-flow.mlir new file mode 100644 index 0000000000..5cba890399 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/skip-control-flow.mlir @@ -0,0 +1,117 @@ +// RUN: triton-opt --add-control-flow-condition %s --allow-unregistered-dialect | FileCheck %s --implicit-check-not=ssbuffer.clone --implicit-check-not=ssbuffer.if +// CHECK: module attributes +// CHECK: func.func @test_skip_control_flow + +module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} { + func.func @test_skip_control_flow(%arg0: memref, %arg1: memref, %arg2: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg3: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg5: memref {tt.divisibility = 16 : i32}, %arg6: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg7: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg8: memref {tt.divisibility = 16 : i32}, %arg9: i32, %arg10: i32, %arg11: i32, %arg12: i32, %arg13: i32, %arg14: i32, %arg15: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} { + %cst = arith.constant {ssbuffer.block_id = 5 : i32} dense<[8, 8, 16, 16]> : tensor<4xi64> + %cst_0 = arith.constant {ssbuffer.block_id = 5 : i32} dense<[128, 8, 16]> : tensor<3xi64> + %cst_1 = arith.constant {ssbuffer.block_id = 8 : i32} 0.000000e+00 : f32 + %c28_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 28 : i32 + %c65536_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 65536 : i32 + %c8192_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 8192 : i32 + %c128_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 128 : i32 + %c1048576_i64 = arith.constant {ssbuffer.block_id = 8 : i32} 1048576 : i64 + %c8388608_i64 = arith.constant {ssbuffer.block_id = 8 : i32} 8388608 : i64 + %c8_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 8 : i32 + %c64_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 64 : i32 + %c0_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 0 : i32 + %cst_2 = arith.constant {ssbuffer.block_id = 8 : i32} 5.000000e-01 : f32 + %cst_3 = arith.constant {ssbuffer.block_id = 8 : i32} 0xFF800000 : f32 + %cst_4 = arith.constant {ssbuffer.block_id = 8 : i32} 1.000000e+00 : f32 + %c128 = arith.constant {ssbuffer.block_id = 8 : i32} 128 : index + scope.scope : () -> () { + %0 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128x128xf32> + %1 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128x128xf32> + %2 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_1 : f32) outs(%0 : tensor<128x128xf32>) -> tensor<128x128xf32> + %3 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_2 : f32) outs(%1 : tensor<128x128xf32>) -> tensor<128x128xf32> + %4 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128xf32> + %5 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_3 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + %6 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_4 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + scf.for %arg16 = %arg13 to %c65536_i32 step %c28_i32 : i32 { + %alloc = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + annotation.mark %alloc {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + %alloc_5 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32, ssbuffer.crossDeps = [0 : i32, 1 : i32]} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_5 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + hivm.hir.sync_block_set {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + %alloc_6 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32, ssbuffer.crossDeps = [1 : i32, 1 : i32]} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_6 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<2>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + hivm.hir.sync_block_set {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + %alloc_7 = memref.alloc() : memref<128xf32, #hivm.address_space> + %memspacecast = memref.memory_space_cast %alloc_7 {ssbuffer.intraDeps = [0 : i32, 1 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %7 = bufferization.to_tensor %memspacecast writable : memref<128xf32> to tensor<128xf32> + %alloc_8 = memref.alloc() : memref<128xf32, #hivm.address_space> + %memspacecast_9 = memref.memory_space_cast %alloc_8 {ssbuffer.intraDeps = [0 : i32, 1 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %8 = bufferization.to_tensor %memspacecast_9 writable : memref<128xf32> to tensor<128xf32> + %9:3 = scf.for %arg17 = %c0_i32 to %c8192_i32 step %c128_i32 iter_args(%arg18 = %6, %arg19 = %2, %arg20 = %5) -> (tensor<128xf32>, tensor<128x128xf32>, tensor<128xf32>) : i32 { + hivm.hir.sync_block_wait {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + %memspacecast_11 = memref.memory_space_cast %alloc_5 {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32, ssbuffer.crossDeps = [0 : i32, 0 : i32]} : memref<128x128xf32, #hivm.address_space> to memref<128x128xf32> + %33 = bufferization.to_tensor %memspacecast_11 restrict writable {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32> to tensor<128x128xf32> + %34 = arith.mulf %33, %3 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %reduced = linalg.reduce ins(%34 : tensor<128x128xf32>) outs(%5 : tensor<128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + (%in: f32, %init: f32) { + %55 = arith.maximumf %in, %init : f32 + linalg.yield %55 : f32 + } + %35 = arith.maximumf %arg20, %reduced {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %broadcasted_12 = linalg.broadcast ins(%35 : tensor<128xf32>) outs(%1 : tensor<128x128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + %36 = arith.subf %34, %broadcasted_12 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %37 = math.exp %36 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %38 = arith.truncf %37 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> to tensor<128x128xf16> + %reshape = tensor.reshape %38(%cst_0) {ssbuffer.block_id = 5 : i32} : (tensor<128x128xf16>, tensor<3xi64>) -> tensor<128x8x16xf16> + annotation.mark %reshape {ssbuffer.block_id = 5 : i32, tiling_dim_mapping = {"1" = 1 : index}} : tensor<128x8x16xf16> + %39 = tensor.empty() {ssbuffer.block_id = 5 : i32} : tensor<8x128x16xf16> + %transposed = linalg.transpose ins(%reshape : tensor<128x8x16xf16>) outs(%39 : tensor<8x128x16xf16>) permutation = [1, 0, 2] {ssbuffer.block_id = 5 : i32} + %reshape_13 = tensor.reshape %transposed(%cst) {ssbuffer.block_id = 5 : i32} : (tensor<8x128x16xf16>, tensor<4xi64>) -> tensor<8x8x16x16xf16> + annotation.mark %reshape_13 {ssbuffer.block_id = 5 : i32, tiling_dim_mapping = {"1" = 1 : index}} : tensor<8x8x16x16xf16> + %40 = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst_1 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + %reduced_14 = linalg.reduce ins(%37 : tensor<128x128xf32>) outs(%40 : tensor<128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + (%in: f32, %init: f32) { + %55 = arith.addf %in, %init : f32 + linalg.yield %55 : f32 + } + %41 = arith.subf %arg20, %35 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %42 = math.exp %41 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %43 = arith.divui %arg17, %c128_i32 {ssbuffer.block_id = 5 : i32} : i32 + %c0_i32_15 = arith.constant {ssbuffer.block_id = 5 : i32} 0 : i32 + %c2_i32 = arith.constant {ssbuffer.block_id = 5 : i32} 2 : i32 + %44 = arith.remsi %43, %c2_i32 {ssbuffer.block_id = 5 : i32} : i32 + %45 = arith.cmpi eq, %44, %c0_i32_15 {ssbuffer.block_id = 5 : i32} : i32 + scf.if %45 { + bufferization.materialize_in_destination %42 in restrict writable %memspacecast {ssbuffer.block_id = 5 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + } else { + bufferization.materialize_in_destination %42 in restrict writable %memspacecast_9 {ssbuffer.block_id = 5 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + } {ssbuffer.block_id = 5 : i32} + %46 = arith.mulf %arg18, %42 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %47 = arith.addf %46, %reduced_14 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + hivm.hir.sync_block_set {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + hivm.hir.sync_block_wait {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32}[, , ] flag = 1 + hivm.hir.copy ins(%reshape_13 : tensor<8x8x16x16xf16>) outs(%alloc : memref<8x8x16x16xf16, #hivm.address_space>) {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32} + hivm.hir.sync_block_set {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32}[, , ] flag = 1 + hivm.hir.sync_block_wait {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + %memspacecast_16 = memref.memory_space_cast %alloc_6 {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32, ssbuffer.crossDeps = [1, 0]} : memref<128x128xf32, #hivm.address_space> to memref<128x128xf32> + %48 = bufferization.to_tensor %memspacecast_16 restrict writable {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32> to tensor<128x128xf32> + %49 = arith.divui %arg17, %c128_i32 {ssbuffer.block_id = 6 : i32} : i32 + %c2_i32_17 = arith.constant {ssbuffer.block_id = 6 : i32} 2 : i32 + %50 = arith.remsi %49, %c2_i32_17 {ssbuffer.block_id = 6 : i32} : i32 + %c0_i32_18 = arith.constant {ssbuffer.block_id = 6 : i32} 0 : i32 + %51 = arith.cmpi eq, %50, %c0_i32_18 {ssbuffer.block_id = 6 : i32} : i32 + %52 = scf.if %51 -> (tensor<128xf32>) { + %55 = bufferization.to_tensor %memspacecast restrict writable {ssbuffer.block_id = 6 : i32} : memref<128xf32> to tensor<128xf32> + scf.yield %55 : tensor<128xf32> + } else { + %55 = bufferization.to_tensor %memspacecast_9 restrict writable {ssbuffer.block_id = 6 : i32} : memref<128xf32> to tensor<128xf32> + scf.yield %55 : tensor<128xf32> + } {ssbuffer.block_id = 6 : i32, ssbuffer.intraDeps = [0 : i32, 0 : i32]} + %broadcasted_19 = linalg.broadcast ins(%52 : tensor<128xf32>) outs(%1 : tensor<128x128xf32>) dimensions = [1] {ssbuffer.block_id = 6 : i32} + %53 = arith.mulf %arg19, %broadcasted_19 {ssbuffer.block_id = 6 : i32} : tensor<128x128xf32> + %54 = arith.addf %48, %53 {ssbuffer.block_id = 6 : i32} : tensor<128x128xf32> + hivm.hir.sync_block_set {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + scf.yield %47, %54, %35 : tensor<128xf32>, tensor<128x128xf32>, tensor<128xf32> + } {ssbuffer.block_id = 9 : i32, ssbuffer.main_loop = 0 : i64} + } {ssbuffer.block_id = 10 : i32} + scope.return + } {hivm.tcore_type = #hivm.tcore_type, ssbuffer.skip} + return {ssbuffer.core_type = "VECTOR"} + } +} diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/update-for-ops.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/update-for-ops.mlir new file mode 100644 index 0000000000..e6edd1fa37 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/update-for-ops.mlir @@ -0,0 +1,131 @@ +// RUN: triton-opt --update-for-ops %s --allow-unregistered-dialect | FileCheck %s + +// Test update-for-ops pass adds iter_args for block counters, inner dependency conditions and inserts PIPE_S sync +// CHECK: func.func @test_update_for_ops +module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} { + func.func @test_update_for_ops(%arg0: memref, %arg1: memref, %arg2: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg3: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg5: memref {tt.divisibility = 16 : i32}, %arg6: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg7: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg8: memref {tt.divisibility = 16 : i32}, %arg9: i32, %arg10: i32, %arg11: i32, %arg12: i32, %arg13: i32, %arg14: i32, %arg15: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} { + %cst = arith.constant {ssbuffer.block_id = 5 : i32} dense<[8, 8, 16, 16]> : tensor<4xi64> + %cst_0 = arith.constant {ssbuffer.block_id = 5 : i32} dense<[128, 8, 16]> : tensor<3xi64> + %cst_1 = arith.constant {ssbuffer.block_id = 8 : i32} 0.000000e+00 : f32 + %c28_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 28 : i32 + %c65536_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 65536 : i32 + %c8192_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 8192 : i32 + %c128_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 128 : i32 + %c0_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 0 : i32 + %cst_2 = arith.constant {ssbuffer.block_id = 8 : i32} 5.000000e-01 : f32 + %cst_3 = arith.constant {ssbuffer.block_id = 8 : i32} 0xFF800000 : f32 + %cst_4 = arith.constant {ssbuffer.block_id = 8 : i32} 1.000000e+00 : f32 + %c128 = arith.constant {ssbuffer.block_id = 8 : i32} 128 : index + scope.scope : () -> () { + // CHECK: hivm.hir.sync_block_set[, , ] flag = 15 + %0 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128x128xf32> + %1 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128x128xf32> + %2 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_1 : f32) outs(%0 : tensor<128x128xf32>) -> tensor<128x128xf32> + %3 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_2 : f32) outs(%1 : tensor<128x128xf32>) -> tensor<128x128xf32> + %4 = tensor.empty() {ssbuffer.block_id = 8 : i32} : tensor<128xf32> + %5 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_3 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + %6 = linalg.fill {ssbuffer.block_id = 8 : i32} ins(%cst_4 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + scf.for %arg16 = %arg13 to %c65536_i32 step %c28_i32 : i32 { + %alloc = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + annotation.mark %alloc {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<0>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 0 : i32} : memref<8x8x16x16xf16, #hivm.address_space> + %alloc_5 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.crossDeps = [0 : i32, 1 : i32], ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_5 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<1>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> + hivm.hir.sync_block_set {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + %alloc_6 = memref.alloc() {ssbuffer.block_id = 9 : i32, ssbuffer.crossDeps = [1 : i32, 1 : i32], ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + annotation.mark %alloc_6 {effects = ["write", "read"], hivm.tightly_coupled_buffer = #hivm.tightly_coupled_buffer<2>, ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> + hivm.hir.sync_block_set {ssbuffer.block_id = 9 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + %alloc_7 = memref.alloc() : memref<128xf32, #hivm.address_space> + %memspacecast = memref.memory_space_cast %alloc_7 {ssbuffer.intraDeps = [0 : i32, 1 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %7 = bufferization.to_tensor %memspacecast writable : memref<128xf32> to tensor<128xf32> + %alloc_8 = memref.alloc() : memref<128xf32, #hivm.address_space> + %memspacecast_9 = memref.memory_space_cast %alloc_8 {ssbuffer.intraDeps = [0 : i32, 1 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %8 = bufferization.to_tensor %memspacecast_9 writable : memref<128xf32> to tensor<128xf32> + // CHECK: scf.for %arg17 + // CHECK: iter_args({{.*}}, {{.*}}, {{.*}}, %{{.*}}, %{{.*}}) -> (tensor<128xf32>, tensor<128x128xf32>, tensor<128xf32>, i32, i32) + %9:3 = scf.for %arg17 = %c0_i32 to %c8192_i32 step %c128_i32 iter_args(%arg18 = %6, %arg19 = %2, %arg20 = %5) -> (tensor<128xf32>, tensor<128x128xf32>, tensor<128xf32>) : i32 { + // CHECK: hivm.hir.sync_block_wait[, , ] flag = 15 + %true = arith.constant true + %10:2 = scf.if %true -> (tensor<128xf32>, tensor<128xf32>) { + hivm.hir.sync_block_wait {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + %memspacecast_11 = memref.memory_space_cast %alloc_5 {ssbuffer.block_id = 5 : i32, ssbuffer.crossDeps = [0 : i32, 0 : i32], ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32, #hivm.address_space> to memref<128x128xf32> + %12 = bufferization.to_tensor %memspacecast_11 restrict writable {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32} : memref<128x128xf32> to tensor<128x128xf32> + %13 = arith.mulf %12, %3 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %reduced = linalg.reduce ins(%13 : tensor<128x128xf32>) outs(%5 : tensor<128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + (%in: f32, %init: f32) { + %27 = arith.maximumf %in, %init : f32 + linalg.yield %27 : f32 + } + %14 = arith.maximumf %arg20, %reduced {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %broadcasted = linalg.broadcast ins(%14 : tensor<128xf32>) outs(%1 : tensor<128x128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + %15 = arith.subf %13, %broadcasted {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %16 = math.exp %15 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> + %17 = arith.truncf %16 {ssbuffer.block_id = 5 : i32} : tensor<128x128xf32> to tensor<128x128xf16> + %reshape = tensor.reshape %17(%cst_0) {ssbuffer.block_id = 5 : i32} : (tensor<128x128xf16>, tensor<3xi64>) -> tensor<128x8x16xf16> + annotation.mark %reshape {ssbuffer.block_id = 5 : i32, tiling_dim_mapping = {"1" = 1 : index}} : tensor<128x8x16xf16> + %18 = tensor.empty() {ssbuffer.block_id = 5 : i32} : tensor<8x128x16xf16> + %transposed = linalg.transpose ins(%reshape : tensor<128x8x16xf16>) outs(%18 : tensor<8x128x16xf16>) permutation = [1, 0, 2] {ssbuffer.block_id = 5 : i32} + %reshape_12 = tensor.reshape %transposed(%cst) {ssbuffer.block_id = 5 : i32} : (tensor<8x128x16xf16>, tensor<4xi64>) -> tensor<8x8x16x16xf16> + annotation.mark %reshape_12 {ssbuffer.block_id = 5 : i32, tiling_dim_mapping = {"1" = 1 : index}} : tensor<8x8x16x16xf16> + %19 = linalg.fill {ssbuffer.block_id = 5 : i32} ins(%cst_1 : f32) outs(%4 : tensor<128xf32>) -> tensor<128xf32> + %reduced_13 = linalg.reduce ins(%16 : tensor<128x128xf32>) outs(%19 : tensor<128xf32>) dimensions = [1] {ssbuffer.block_id = 5 : i32} + (%in: f32, %init: f32) { + %27 = arith.addf %in, %init : f32 + linalg.yield %27 : f32 + } + %20 = arith.subf %arg20, %14 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %21 = math.exp %20 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %22 = arith.divui %arg17, %c128_i32 {ssbuffer.block_id = 5 : i32} : i32 + %c0_i32_14 = arith.constant {ssbuffer.block_id = 5 : i32} 0 : i32 + %c2_i32 = arith.constant {ssbuffer.block_id = 5 : i32} 2 : i32 + %23 = arith.remsi %22, %c2_i32 {ssbuffer.block_id = 5 : i32} : i32 + %24 = arith.cmpi eq, %23, %c0_i32_14 {ssbuffer.block_id = 5 : i32} : i32 + scf.if %24 { + bufferization.materialize_in_destination %21 in restrict writable %memspacecast {ssbuffer.block_id = 5 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + } else { + bufferization.materialize_in_destination %21 in restrict writable %memspacecast_9 {ssbuffer.block_id = 5 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + } {ssbuffer.block_id = 5 : i32} + %25 = arith.mulf %arg18, %21 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + %26 = arith.addf %25, %reduced_13 {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + hivm.hir.sync_block_set {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 1 : i32}[, , ] flag = 2 + hivm.hir.sync_block_wait {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32}[, , ] flag = 1 + hivm.hir.copy ins(%reshape_12 : tensor<8x8x16x16xf16>) outs(%alloc : memref<8x8x16x16xf16, #hivm.address_space>) {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32} + hivm.hir.sync_block_set {ssbuffer.block_id = 5 : i32, ssbuffer.transfer_id = 0 : i32}[, , ] flag = 1 + scf.yield %14, %26 : tensor<128xf32>, tensor<128xf32> + } else { + scf.yield %arg20, %arg18 : tensor<128xf32>, tensor<128xf32> + } {ssbuffer.if = 5 : i32} + %true_10 = arith.constant true + %11 = scf.if %true_10 -> (tensor<128x128xf32>) { + hivm.hir.sync_block_wait {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + %memspacecast_11 = memref.memory_space_cast %alloc_6 {ssbuffer.block_id = 6 : i32, ssbuffer.crossDeps = [1, 0], ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32, #hivm.address_space> to memref<128x128xf32> + %12 = bufferization.to_tensor %memspacecast_11 restrict writable {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32} : memref<128x128xf32> to tensor<128x128xf32> + %13 = arith.divui %arg17, %c128_i32 {ssbuffer.block_id = 6 : i32} : i32 + %c2_i32 = arith.constant {ssbuffer.block_id = 6 : i32} 2 : i32 + %14 = arith.remsi %13, %c2_i32 {ssbuffer.block_id = 6 : i32} : i32 + %c0_i32_12 = arith.constant {ssbuffer.block_id = 6 : i32} 0 : i32 + %15 = arith.cmpi eq, %14, %c0_i32_12 {ssbuffer.block_id = 6 : i32} : i32 + %16 = scf.if %15 -> (tensor<128xf32>) { + %19 = bufferization.to_tensor %memspacecast restrict writable {ssbuffer.block_id = 6 : i32} : memref<128xf32> to tensor<128xf32> + scf.yield %19 : tensor<128xf32> + } else { + %19 = bufferization.to_tensor %memspacecast_9 restrict writable {ssbuffer.block_id = 6 : i32} : memref<128xf32> to tensor<128xf32> + scf.yield %19 : tensor<128xf32> + } {ssbuffer.block_id = 6 : i32, ssbuffer.intraDeps = [0 : i32, 0 : i32]} + %broadcasted = linalg.broadcast ins(%16 : tensor<128xf32>) outs(%1 : tensor<128x128xf32>) dimensions = [1] {ssbuffer.block_id = 6 : i32} + %17 = arith.mulf %arg19, %broadcasted {ssbuffer.block_id = 6 : i32} : tensor<128x128xf32> + %18 = arith.addf %12, %17 {ssbuffer.block_id = 6 : i32} : tensor<128x128xf32> + hivm.hir.sync_block_set {ssbuffer.block_id = 6 : i32, ssbuffer.transfer_id = 2 : i32}[, , ] flag = 3 + scf.yield %18 : tensor<128x128xf32> + } else { + scf.yield %arg19 : tensor<128x128xf32> + } {ssbuffer.if = 6 : i32} + // CHECK: hivm.hir.sync_block_set[, , ] flag = 15 + // CHECK: scf.yield %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} + scf.yield %10#1, %11, %10#0 : tensor<128xf32>, tensor<128x128xf32>, tensor<128xf32> + } {ssbuffer.block_id = 9 : i32, ssbuffer.main_loop = 0 : i64} + } {ssbuffer.block_id = 10 : i32} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return {ssbuffer.core_type = "VECTOR"} + } +} diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/updateLoopIterTimes2_test.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/updateLoopIterTimes2_test.mlir new file mode 100644 index 0000000000..0f4bc5fc1c --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/updateLoopIterTimes2_test.mlir @@ -0,0 +1,75 @@ +// RUN: triton-opt --add-control-flow-condition %s | FileCheck %s + +// ============================================================================ +// TC-02: insufficient multicache +// ============================================================================ +// If Vm has intra-core dependency on Vn (m>n), to achieve optimal pipeline, at most m-n+1 data flows need to exist simultaneously, which is requireNum +// If multicache count is insufficient (bufferNum < requireNum), iteration count needs to be multiplied by requireNum / bufferNum + +module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} { + func.func @_attn_fwd(%arg0: memref, %arg1: memref, %arg2: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg3: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg5: memref {tt.divisibility = 16 : i32}, %arg6: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg7: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg8: memref {tt.divisibility = 16 : i32}, %arg9: i32, %arg10: i32, %arg11: i32, %arg12: i32, %arg13: i32, %arg14: i32, %arg15: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} { + %c8192_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 8192 : i32 + %c128_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 128 : i32 + %c0_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 0 : i32 + scope.scope : () -> () { + %alloc_7 = memref.alloc() : memref<128xf32, #hivm.address_space> + %memspacecast = memref.memory_space_cast %alloc_7 {ssbuffer.intraDeps = [0 : i32, 1 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %7 = bufferization.to_tensor %memspacecast writable : memref<128xf32> to tensor<128xf32> + %alloc_8 = memref.alloc() : memref<128xf32, #hivm.address_space> + %memspacecast_9 = memref.memory_space_cast %alloc_8 {ssbuffer.intraDeps = [0 : i32, 1 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + + // CHECK: %[[IFCOUNT:.*]] = arith.constant 3 + // CHECK: %[[REQUIRENUM:.*]] = arith.constant 3 + // CHECK: %[[BUFFERNUM:.*]] = arith.constant 2 + // CHECK: arith.subi + // CHECK: arith.ceildivui + // CHECK: arith.muli %{{.*}}, %[[REQUIRENUM]] + // CHECK: arith.ceildivui %{{.*}}, %[[BUFFERNUM]] + // CHECK: arith.addi %{{.*}}, %[[IFCOUNT]] + // CHECK: arith.muli + // CHECK: arith.addi + // CHECK: scf.for + // CHECK: iter_args(%arg17 = %c0_i32, %arg18 = %c0_i32, %[[CNT0:.*]] = %{{.+}}, %[[CNT1:.*]] = %{{.+}}, %[[CNT2:.*]] = %{{.+}}, %{{.+}} = %{{.+}}) + %30:2 = scf.for %arg16 = %c0_i32 to %c8192_i32 step %c128_i32 iter_args(%arg17 = %c0_i32, %arg18 = %c0_i32) -> (i32, i32) : i32 { + %42 = tensor.empty() {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + // CHECK: arith.divui %[[CNT0]], %c128_i32 {ssbuffer.block_id = 5 : i32} : i32 + %43 = arith.divui %arg16, %c128_i32 {ssbuffer.block_id = 5 : i32} : i32 + %c0_i32_15 = arith.constant {ssbuffer.block_id = 5 : i32} 0 : i32 + %c2_i32 = arith.constant {ssbuffer.block_id = 5 : i32} 2 : i32 + %44 = arith.remsi %43, %c2_i32 {ssbuffer.block_id = 5 : i32} : i32 + %45 = arith.cmpi eq, %44, %c0_i32_15 {ssbuffer.block_id = 5 : i32} : i32 + scf.if %45 { + bufferization.materialize_in_destination %42 in restrict writable %memspacecast {ssbuffer.block_id = 5 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + } else { + bufferization.materialize_in_destination %42 in restrict writable %memspacecast_9 {ssbuffer.block_id = 5 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + } {ssbuffer.block_id = 5 : i32} + + // CHECK: arith.divui %[[CNT1]], %c128_i32 {ssbuffer.block_id = 6 : i32} : i32 + %46 = arith.divui %arg16, %c128_i32 {ssbuffer.block_id = 6 : i32} : i32 + %c0_i32_16 = arith.constant {ssbuffer.block_id = 6 : i32} 0 : i32 + %c2_i32_6 = arith.constant {ssbuffer.block_id = 6 : i32} 2 : i32 + + // CHECK: arith.divui %[[CNT2]], %c128_i32 {ssbuffer.block_id = 7 : i32} : i32 + %49 = arith.divui %arg16, %c128_i32 {ssbuffer.block_id = 7 : i32} : i32 + %c2_i32_17 = arith.constant {ssbuffer.block_id = 7 : i32} 2 : i32 + %50 = arith.remsi %49, %c2_i32_17 {ssbuffer.block_id = 7 : i32} : i32 + %c0_i32_18 = arith.constant {ssbuffer.block_id = 7 : i32} 0 : i32 + %51 = arith.cmpi eq, %50, %c0_i32_18 {ssbuffer.block_id = 7 : i32} : i32 + %52 = scf.if %51 -> (tensor<128xf32>) { + %55 = bufferization.to_tensor %memspacecast restrict writable {ssbuffer.block_id = 7 : i32} : memref<128xf32> to tensor<128xf32> + scf.yield %55 : tensor<128xf32> + } else { + %55 = bufferization.to_tensor %memspacecast_9 restrict writable {ssbuffer.block_id = 7 : i32} : memref<128xf32> to tensor<128xf32> + scf.yield %55 : tensor<128xf32> + } {ssbuffer.block_id = 7 : i32, ssbuffer.intraDeps = [0 : i32, 0 : i32]} + scf.yield %arg17, %arg18 : i32, i32 + } {ssbuffer.block_id = 9 : i32, ssbuffer.main_loop = 0 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + + scope.scope : () -> () { + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return {ssbuffer.core_type = "VECTOR"} + } +} diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/updateLoopIterTimes_test.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/updateLoopIterTimes_test.mlir new file mode 100644 index 0000000000..1862037a4a --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/AddControlFlowCondition/updateLoopIterTimes_test.mlir @@ -0,0 +1,73 @@ +// RUN: triton-opt --add-control-flow-condition %s | FileCheck %s + + +// ============================================================================ +// TC-01: sufficient multicache +// ============================================================================ +// This test verifies: +// 1. New loop upper bound calculation (arith.subi, arith.ceildivui, arith.muli, arith.addi) +// 2. iter_args expanded +// 3. MainLoop counter will be replaced by each counter of ifblocks +// If Vm has intra-core dependency on Vn (m>n), to achieve optimal pipeline, at most m-n+1 data flows need to exist simultaneously, which is requireNum +// If multicache count is sufficient (bufferNum >= requireNum), only need to increase iteration times with number of ifblock + +module attributes {hacc.target = #hacc.target<"Ascend950PR_9579">} { + func.func @_attn_fwd(%arg0: memref, %arg1: memref, %arg2: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg3: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg5: memref {tt.divisibility = 16 : i32}, %arg6: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg7: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg8: memref {tt.divisibility = 16 : i32}, %arg9: i32, %arg10: i32, %arg11: i32, %arg12: i32, %arg13: i32, %arg14: i32, %arg15: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "mix", parallel_mode = "simd"} { + %c8192_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 8192 : i32 + %c128_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 128 : i32 + %c0_i32 = arith.constant {ssbuffer.block_id = 8 : i32} 0 : i32 + scope.scope : () -> () { + %alloc_7 = memref.alloc() : memref<128xf32, #hivm.address_space> + %memspacecast = memref.memory_space_cast %alloc_7 {ssbuffer.intraDeps = [0 : i32, 1 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + %7 = bufferization.to_tensor %memspacecast writable : memref<128xf32> to tensor<128xf32> + %alloc_8 = memref.alloc() : memref<128xf32, #hivm.address_space> + %memspacecast_9 = memref.memory_space_cast %alloc_8 {ssbuffer.intraDeps = [0 : i32, 1 : i32]} : memref<128xf32, #hivm.address_space> to memref<128xf32> + + // CHECK: %[[IFCOUNT:.*]] = arith.constant 2 + // CHECK: %[[REQUIRENUM:.*]] = arith.constant 1 + // CHECK: %[[BUFFERNUM:.*]] = arith.constant 1 + // CHECK: arith.subi + // CHECK: arith.ceildivui + // CHECK: arith.muli %{{.*}}, %[[REQUIRENUM]] + // CHECK: arith.ceildivui %{{.*}}, %[[BUFFERNUM]] + // CHECK: arith.addi %{{.*}}, %[[IFCOUNT]] + // CHECK: arith.muli + // CHECK: arith.addi + // CHECK: scf.for + // CHECK: iter_args(%arg17 = %c0_i32, %arg18 = %c0_i32, %[[CNT0:.*]] = %{{.+}}, %[[CNT1:.*]] = %{{.+}}, %arg21 = %{{.+}}) + %30:2 = scf.for %arg16 = %c0_i32 to %c8192_i32 step %c128_i32 iter_args(%arg17 = %c0_i32, %arg18 = %c0_i32) -> (i32, i32) : i32 { + %42 = tensor.empty() {ssbuffer.block_id = 5 : i32} : tensor<128xf32> + // CHECK: arith.divui %[[CNT0]], %c128_i32 {ssbuffer.block_id = 5 : i32} : i32 + %43 = arith.divui %arg16, %c128_i32 {ssbuffer.block_id = 5 : i32} : i32 + %c0_i32_15 = arith.constant {ssbuffer.block_id = 5 : i32} 0 : i32 + %c2_i32 = arith.constant {ssbuffer.block_id = 5 : i32} 2 : i32 + %44 = arith.remsi %43, %c2_i32 {ssbuffer.block_id = 5 : i32} : i32 + %45 = arith.cmpi eq, %44, %c0_i32_15 {ssbuffer.block_id = 5 : i32} : i32 + scf.if %45 { + bufferization.materialize_in_destination %42 in restrict writable %memspacecast {ssbuffer.block_id = 5 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + } else { + bufferization.materialize_in_destination %42 in restrict writable %memspacecast_9 {ssbuffer.block_id = 5 : i32} : (tensor<128xf32>, memref<128xf32>) -> () + } {ssbuffer.block_id = 5 : i32} + // CHECK: arith.divui %[[CNT1]], %c128_i32 {ssbuffer.block_id = 7 : i32} : i32 + %49 = arith.divui %arg16, %c128_i32 {ssbuffer.block_id = 7 : i32} : i32 + %c2_i32_17 = arith.constant {ssbuffer.block_id = 7 : i32} 2 : i32 + %50 = arith.remsi %49, %c2_i32_17 {ssbuffer.block_id = 7 : i32} : i32 + %c0_i32_18 = arith.constant {ssbuffer.block_id = 7 : i32} 0 : i32 + %51 = arith.cmpi eq, %50, %c0_i32_18 {ssbuffer.block_id = 7 : i32} : i32 + %52 = scf.if %51 -> (tensor<128xf32>) { + %55 = bufferization.to_tensor %memspacecast restrict writable {ssbuffer.block_id = 7 : i32} : memref<128xf32> to tensor<128xf32> + scf.yield %55 : tensor<128xf32> + } else { + %55 = bufferization.to_tensor %memspacecast_9 restrict writable {ssbuffer.block_id = 7 : i32} : memref<128xf32> to tensor<128xf32> + scf.yield %55 : tensor<128xf32> + } {ssbuffer.block_id = 7 : i32, ssbuffer.intraDeps = [0 : i32, 0 : i32]} + scf.yield %arg17, %arg18 : i32, i32 + } {ssbuffer.block_id = 9 : i32, ssbuffer.main_loop = 0 : i64} + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + scope.scope : () -> () { + scope.return + } {hivm.tcore_type = #hivm.tcore_type} + return {ssbuffer.core_type = "VECTOR"} + } +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonAscendAllPass/conv1d.mlir b/third_party/ascend/unittest/Conversion/General/TritonAscendAllPass/conv1d.mlir new file mode 100644 index 0000000000..c3f1379a79 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonAscendAllPass/conv1d.mlir @@ -0,0 +1,160 @@ +// RUN: triton-opt --triton-to-structured '--discrete-mask-access-conversion=compile-on-910-95=False force-simt-template=False' '--triton-to-unstructure=compile-on-910-95=False force-simt-template=False' --triton-to-hivm --triton-to-hfusion --triton-to-llvm --bubble-up-operation --triton-to-structured --triton-to-linalg --split-input-file %s | FileCheck %s + +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + tt.func public @triton_conv1d_2d_kernel(%input_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %weight_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %bias_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %output_ptr: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %output_offsets = arith.constant dense<126> : tensor<16x1xi32> + %weight_offsets = arith.constant dense<3> : tensor<1x16x1xi32> + %weight_offsets_0 = arith.constant dense<48> : tensor<16x1x1xi32> + %input_offsets = arith.constant dense<128> : tensor<16x1xi32> + %c_in_offsets = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> + %c_in_offsets_1 = tt.expand_dims %c_in_offsets {axis = 1 : i32} : tensor<16xi32> -> tensor<16x1xi32> + %l_in_offsets = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32> + %l_in_offsets_2 = tt.expand_dims %l_in_offsets {axis = 0 : i32} : tensor<128xi32> -> tensor<1x128xi32> + %input_offsets_3 = arith.muli %c_in_offsets_1, %input_offsets : tensor<16x1xi32> + %input_offsets_4 = tt.broadcast %input_offsets_3 : tensor<16x1xi32> -> tensor<16x128xi32> + %input_offsets_5 = tt.broadcast %l_in_offsets_2 : tensor<1x128xi32> -> tensor<16x128xi32> + %input_offsets_6 = arith.addi %input_offsets_4, %input_offsets_5 : tensor<16x128xi32> + %input_tensor = tt.splat %input_ptr : !tt.ptr -> tensor<16x128x!tt.ptr> + %input_tensor_7 = tt.addptr %input_tensor, %input_offsets_6 : tensor<16x128x!tt.ptr>, tensor<16x128xi32> + %input_tensor_8 = tt.load %input_tensor_7 : tensor<16x128x!tt.ptr> + %c_out_offsets = tt.expand_dims %c_in_offsets_1 {axis = 2 : i32} : tensor<16x1xi32> -> tensor<16x1x1xi32> + %c_in_group_offsets = tt.expand_dims %c_in_offsets {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32> + %c_in_group_offsets_9 = tt.expand_dims %c_in_group_offsets {axis = 2 : i32} : tensor<1x16xi32> -> tensor<1x16x1xi32> + %k_offsets = tt.make_range {end = 3 : i32, start = 0 : i32} : tensor<3xi32> + %k_offsets_10 = tt.expand_dims %k_offsets {axis = 0 : i32} : tensor<3xi32> -> tensor<1x3xi32> + %k_offsets_11 = tt.expand_dims %k_offsets_10 {axis = 1 : i32} : tensor<1x3xi32> -> tensor<1x1x3xi32> + %weight_offsets_12 = arith.muli %c_out_offsets, %weight_offsets_0 : tensor<16x1x1xi32> + %weight_offsets_13 = arith.muli %c_in_group_offsets_9, %weight_offsets : tensor<1x16x1xi32> + %weight_offsets_14 = tt.broadcast %weight_offsets_12 : tensor<16x1x1xi32> -> tensor<16x16x1xi32> + %weight_offsets_15 = tt.broadcast %weight_offsets_13 : tensor<1x16x1xi32> -> tensor<16x16x1xi32> + %weight_offsets_16 = arith.addi %weight_offsets_14, %weight_offsets_15 : tensor<16x16x1xi32> + %weight_offsets_17 = tt.broadcast %weight_offsets_16 : tensor<16x16x1xi32> -> tensor<16x16x3xi32> + %weight_offsets_18 = tt.broadcast %k_offsets_11 : tensor<1x1x3xi32> -> tensor<16x16x3xi32> + %weight_offsets_19 = arith.addi %weight_offsets_17, %weight_offsets_18 : tensor<16x16x3xi32> + %weight_tensor = tt.splat %weight_ptr : !tt.ptr -> tensor<16x16x3x!tt.ptr> + %weight_tensor_20 = tt.addptr %weight_tensor, %weight_offsets_19 : tensor<16x16x3x!tt.ptr>, tensor<16x16x3xi32> + %weight_tensor_21 = tt.load %weight_tensor_20 : tensor<16x16x3x!tt.ptr> + %bias_tensor = tt.splat %bias_ptr : !tt.ptr -> tensor<16x!tt.ptr> + %bias_tensor_22 = tt.addptr %bias_tensor, %c_in_offsets : tensor<16x!tt.ptr>, tensor<16xi32> + %bias_tensor_23 = tt.load %bias_tensor_22 : tensor<16x!tt.ptr> + %output = ascend.conv1d(%input_tensor_8, %weight_tensor_21, %bias_tensor_23) {dilation = 1 : i64, groups = 1 : i64, padding_size = 0 : i64, stride = 1 : i64} : (tensor<16x128xf16>, tensor<16x16x3xf16>, tensor<16xf16>) -> tensor<16x126xf16> + %l_out_off = tt.make_range {end = 126 : i32, start = 0 : i32} : tensor<126xi32> + %l_out_off_24 = tt.expand_dims %l_out_off {axis = 0 : i32} : tensor<126xi32> -> tensor<1x126xi32> + %output_offsets_25 = arith.muli %c_in_offsets_1, %output_offsets : tensor<16x1xi32> + %output_offsets_26 = tt.broadcast %output_offsets_25 : tensor<16x1xi32> -> tensor<16x126xi32> + %output_offsets_27 = tt.broadcast %l_out_off_24 : tensor<1x126xi32> -> tensor<16x126xi32> + %output_offsets_28 = arith.addi %output_offsets_26, %output_offsets_27 : tensor<16x126xi32> + %0 = tt.splat %output_ptr : !tt.ptr -> tensor<16x126x!tt.ptr> + %1 = tt.addptr %0, %output_offsets_28 : tensor<16x126x!tt.ptr>, tensor<16x126xi32> + tt.store %1, %output : tensor<16x126x!tt.ptr> + tt.return + } +} + +// CHECK-LABEL: func.func @triton_conv1d_2d_kernel( +// CHECK: %[[VAL_0:.*]] = memref.reinterpret_cast +// CHECK-SAME: to offset: [0], sizes: [16, 128], strides: [128, 1] +// CHECK: %[[VAL_1:.*]] = memref.alloc() : memref<16x128xf16> +// CHECK: memref.copy %[[VAL_0]], %[[VAL_1]] +// CHECK: %[[VAL_2:.*]] = memref.reinterpret_cast +// CHECK-SAME: to offset: [0], sizes: [16, 16, 3], strides: [48, 3, 1] +// CHECK: %[[VAL_3:.*]] = memref.alloc() : memref<16x16x3xf16> +// CHECK: memref.copy %[[VAL_2]], %[[VAL_3]] +// CHECK: %[[VAL_4:.*]] = memref.reinterpret_cast +// CHECK-SAME: to offset: [0], sizes: [16], strides: [1] +// CHECK: %[[VAL_5:.*]] = memref.alloc() : memref<16xf16> +// CHECK: memref.copy %[[VAL_4]], %[[VAL_5]] +// CHECK: %[[VAL_6:.*]] = hfusion.conv1d +// CHECK-SAME: {dilation = 1 : i32, groups = 1 : i32, padding = 0 : i32, stride = 1 : i32} +// CHECK: %[[VAL_7:.*]] = memref.reinterpret_cast +// CHECK-SAME: to offset: [0], sizes: [16, 126], strides: [126, 1] +// CHECK: bufferization.materialize_in_destination %[[VAL_6]] in writable %[[VAL_7]] + + +// ----- + + +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + tt.func public @triton_conv1d_3d_kernel(%input_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %weight_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %bias_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %output_ptr: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %output_offsets = arith.constant dense<128> : tensor<1x30x1xi32> + %output_offsets_0 = arith.constant dense<3840> : tensor<2x1x1xi32> + %weight_offsets = arith.constant dense<3> : tensor<1x32x1xi32> + %weight_offsets_1 = arith.constant dense<96> : tensor<30x1x1xi32> + %input_offsets = arith.constant dense<128> : tensor<1x32x1xi32> + %input_offsets_2 = arith.constant dense<4096> : tensor<2x1x1xi32> + %n_offsets = tt.make_range {end = 2 : i32, start = 0 : i32} : tensor<2xi32> + %n_offsets_3 = tt.expand_dims %n_offsets {axis = 1 : i32} : tensor<2xi32> -> tensor<2x1xi32> + %n_offsets_4 = tt.expand_dims %n_offsets_3 {axis = 2 : i32} : tensor<2x1xi32> -> tensor<2x1x1xi32> + %c_in_offsets = tt.make_range {end = 32 : i32, start = 0 : i32} : tensor<32xi32> + %c_in_offsets_5 = tt.expand_dims %c_in_offsets {axis = 0 : i32} : tensor<32xi32> -> tensor<1x32xi32> + %c_in_offsets_6 = tt.expand_dims %c_in_offsets_5 {axis = 2 : i32} : tensor<1x32xi32> -> tensor<1x32x1xi32> + %l_in_offsets = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32> + %l_in_offsets_7 = tt.expand_dims %l_in_offsets {axis = 0 : i32} : tensor<128xi32> -> tensor<1x128xi32> + %l_in_offsets_8 = tt.expand_dims %l_in_offsets_7 {axis = 1 : i32} : tensor<1x128xi32> -> tensor<1x1x128xi32> + %input_offsets_9 = arith.muli %n_offsets_4, %input_offsets_2 : tensor<2x1x1xi32> + %input_offsets_10 = arith.muli %c_in_offsets_6, %input_offsets : tensor<1x32x1xi32> + %input_offsets_11 = tt.broadcast %input_offsets_9 : tensor<2x1x1xi32> -> tensor<2x32x1xi32> + %input_offsets_12 = tt.broadcast %input_offsets_10 : tensor<1x32x1xi32> -> tensor<2x32x1xi32> + %input_offsets_13 = arith.addi %input_offsets_11, %input_offsets_12 : tensor<2x32x1xi32> + %input_offsets_14 = tt.broadcast %input_offsets_13 : tensor<2x32x1xi32> -> tensor<2x32x128xi32> + %input_offsets_15 = tt.broadcast %l_in_offsets_8 : tensor<1x1x128xi32> -> tensor<2x32x128xi32> + %input_offsets_16 = arith.addi %input_offsets_14, %input_offsets_15 : tensor<2x32x128xi32> + %input_tensor = tt.splat %input_ptr : !tt.ptr -> tensor<2x32x128x!tt.ptr> + %input_tensor_17 = tt.addptr %input_tensor, %input_offsets_16 : tensor<2x32x128x!tt.ptr>, tensor<2x32x128xi32> + %input_tensor_18 = tt.load %input_tensor_17 : tensor<2x32x128x!tt.ptr> + %c_out_offsets = tt.make_range {end = 30 : i32, start = 0 : i32} : tensor<30xi32> + %c_out_offsets_19 = tt.expand_dims %c_out_offsets {axis = 1 : i32} : tensor<30xi32> -> tensor<30x1xi32> + %c_out_offsets_20 = tt.expand_dims %c_out_offsets_19 {axis = 2 : i32} : tensor<30x1xi32> -> tensor<30x1x1xi32> + %k_offsets = tt.make_range {end = 3 : i32, start = 0 : i32} : tensor<3xi32> + %k_offsets_21 = tt.expand_dims %k_offsets {axis = 0 : i32} : tensor<3xi32> -> tensor<1x3xi32> + %k_offsets_22 = tt.expand_dims %k_offsets_21 {axis = 1 : i32} : tensor<1x3xi32> -> tensor<1x1x3xi32> + %weight_offsets_23 = arith.muli %c_out_offsets_20, %weight_offsets_1 : tensor<30x1x1xi32> + %weight_offsets_24 = arith.muli %c_in_offsets_6, %weight_offsets : tensor<1x32x1xi32> + %weight_offsets_25 = tt.broadcast %weight_offsets_23 : tensor<30x1x1xi32> -> tensor<30x32x1xi32> + %weight_offsets_26 = tt.broadcast %weight_offsets_24 : tensor<1x32x1xi32> -> tensor<30x32x1xi32> + %weight_offsets_27 = arith.addi %weight_offsets_25, %weight_offsets_26 : tensor<30x32x1xi32> + %weight_offsets_28 = tt.broadcast %weight_offsets_27 : tensor<30x32x1xi32> -> tensor<30x32x3xi32> + %weight_offsets_29 = tt.broadcast %k_offsets_22 : tensor<1x1x3xi32> -> tensor<30x32x3xi32> + %weight_offsets_30 = arith.addi %weight_offsets_28, %weight_offsets_29 : tensor<30x32x3xi32> + %weight_tensor = tt.splat %weight_ptr : !tt.ptr -> tensor<30x32x3x!tt.ptr> + %weight_tensor_31 = tt.addptr %weight_tensor, %weight_offsets_30 : tensor<30x32x3x!tt.ptr>, tensor<30x32x3xi32> + %weight_tensor_32 = tt.load %weight_tensor_31 : tensor<30x32x3x!tt.ptr> + %bias_tensor = tt.splat %bias_ptr : !tt.ptr -> tensor<30x!tt.ptr> + %bias_tensor_33 = tt.addptr %bias_tensor, %c_out_offsets : tensor<30x!tt.ptr>, tensor<30xi32> + %bias_tensor_34 = tt.load %bias_tensor_33 : tensor<30x!tt.ptr> + %output = ascend.conv1d(%input_tensor_18, %weight_tensor_32, %bias_tensor_34) {dilation = 1 : i64, groups = 1 : i64, padding_size = 1 : i64, stride = 1 : i64} : (tensor<2x32x128xf16>, tensor<30x32x3xf16>, tensor<30xf16>) -> tensor<2x30x128xf16> + %c_out_offsets_35 = tt.expand_dims %c_out_offsets {axis = 0 : i32} : tensor<30xi32> -> tensor<1x30xi32> + %c_out_offsets_36 = tt.expand_dims %c_out_offsets_35 {axis = 2 : i32} : tensor<1x30xi32> -> tensor<1x30x1xi32> + %output_offsets_37 = arith.muli %n_offsets_4, %output_offsets_0 : tensor<2x1x1xi32> + %output_offsets_38 = arith.muli %c_out_offsets_36, %output_offsets : tensor<1x30x1xi32> + %output_offsets_39 = tt.broadcast %output_offsets_37 : tensor<2x1x1xi32> -> tensor<2x30x1xi32> + %output_offsets_40 = tt.broadcast %output_offsets_38 : tensor<1x30x1xi32> -> tensor<2x30x1xi32> + %output_offsets_41 = arith.addi %output_offsets_39, %output_offsets_40 : tensor<2x30x1xi32> + %output_offsets_42 = tt.broadcast %output_offsets_41 : tensor<2x30x1xi32> -> tensor<2x30x128xi32> + %output_offsets_43 = tt.broadcast %l_in_offsets_8 : tensor<1x1x128xi32> -> tensor<2x30x128xi32> + %output_offsets_44 = arith.addi %output_offsets_42, %output_offsets_43 : tensor<2x30x128xi32> + %0 = tt.splat %output_ptr : !tt.ptr -> tensor<2x30x128x!tt.ptr> + %1 = tt.addptr %0, %output_offsets_44 : tensor<2x30x128x!tt.ptr>, tensor<2x30x128xi32> + tt.store %1, %output : tensor<2x30x128x!tt.ptr> + tt.return + } +} + +// CHECK-LABEL: func.func @triton_conv1d_3d_kernel( +// CHECK: %[[VAL_0:.*]] = memref.reinterpret_cast +// CHECK-SAME: to offset: [0], sizes: [2, 32, 128], strides: [4096, 128, 1] +// CHECK: %[[VAL_1:.*]] = memref.alloc() : memref<2x32x128xf16> +// CHECK: memref.copy %[[VAL_0]], %[[VAL_1]] +// CHECK: %[[VAL_2:.*]] = memref.reinterpret_cast +// CHECK-SAME: to offset: [0], sizes: [30, 32, 3], strides: [96, 3, 1] +// CHECK: %[[VAL_3:.*]] = memref.alloc() : memref<30x32x3xf16> +// CHECK: memref.copy %[[VAL_2]], %[[VAL_3]] +// CHECK: %[[VAL_4:.*]] = memref.reinterpret_cast +// CHECK-SAME: to offset: [0], sizes: [30], strides: [1] +// CHECK: %[[VAL_5:.*]] = memref.alloc() : memref<30xf16> +// CHECK: memref.copy %[[VAL_4]], %[[VAL_5]] +// CHECK: %[[VAL_6:.*]] = hfusion.conv1d +// CHECK-SAME: {dilation = 1 : i32, groups = 1 : i32, padding = 1 : i32, stride = 1 : i32} +// CHECK: %[[VAL_7:.*]] = memref.reinterpret_cast +// CHECK-SAME: to offset: [0], sizes: [2, 30, 128], strides: [3840, 128, 1] +// CHECK: bufferization.materialize_in_destination %[[VAL_6]] in writable %[[VAL_7]] diff --git a/third_party/ascend/unittest/Conversion/General/TritonToHFusion/fp_to_fp_rtz.mlir b/third_party/ascend/unittest/Conversion/General/TritonToHFusion/fp_to_fp_rtz.mlir new file mode 100644 index 0000000000..fbff6cb341 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToHFusion/fp_to_fp_rtz.mlir @@ -0,0 +1,18 @@ +// RUN: triton-opt %s -triton-to-hfusion | FileCheck %s + +// CHECK-LABEL: tt.func @test_fp32_to_fp16_rtz +tt.func @test_fp32_to_fp16_rtz(%arg0: tensor<1024xf32>) -> tensor<1024xf16> { + // CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<1024xf16> + // CHECK: %[[RESULT:.*]] = hfusion.cast {mode = #hfusion.round_mode} ins(%arg0 : tensor<1024xf32>) outs(%[[EMPTY]] : tensor<1024xf16>) -> tensor<1024xf16> + %0 = tt.fp_to_fp %arg0, rounding = rtz : tensor<1024xf32> -> tensor<1024xf16> + // CHECK: return %[[RESULT]] + tt.return %0 : tensor<1024xf16> +} + + +// CHECK-LABEL: tt.func @test_fp32_to_fp16_rtz_fail +tt.func @test_fp32_to_fp16_rtz_fail(%arg0: tensor<1024xf32>) -> tensor<1024xf16> { + %0 = tt.fp_to_fp %arg0, rounding = rtne : tensor<1024xf32> -> tensor<1024xf16> + // CHECK: %{{.*}} = tt.fp_to_fp %arg0, rounding = rtne : tensor<1024xf32> -> tensor<1024xf16> + tt.return %0 : tensor<1024xf16> +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToHFusion/histogram.mlir b/third_party/ascend/unittest/Conversion/General/TritonToHFusion/histogram.mlir new file mode 100644 index 0000000000..c1ce707f91 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToHFusion/histogram.mlir @@ -0,0 +1,9 @@ +// RUN: triton-opt %s -triton-to-hfusion | FileCheck %s + + +// CHECK-LABEL: tt.func @test_histogram +tt.func @test_histogram(%arg0 : tensor<16xi32>) -> tensor<2xi32> { + // CHECK: hfusion.histogram %{{.*}}, 2 + %res = tt.histogram %arg0 : tensor<16xi32> -> tensor<2xi32> + tt.return %res : tensor<2xi32> +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToHFusion/histogram_expect_fail.mlir b/third_party/ascend/unittest/Conversion/General/TritonToHFusion/histogram_expect_fail.mlir new file mode 100644 index 0000000000..62d69153d3 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToHFusion/histogram_expect_fail.mlir @@ -0,0 +1,16 @@ +// RUN: ! triton-opt %s -triton-to-hfusion 2>&1 | FileCheck %s --check-prefix=CHECK-FAIL-ZERO +// Expected failure case, only for branch coverage +// CHECK-FAIL-ZERO: error: 'hfusion.histogram' op output length (0) must equal num_bins +tt.func @test_histogram_fail_with_zero_n(%arg0 : tensor<16xi32>) -> tensor<0xi32> { + %res = tt.histogram %arg0 : tensor<16xi32> -> tensor<0xi32> + tt.return %res : tensor<0xi32> +} + + +// RUN: ! triton-opt %s -triton-to-hfusion 2>&1 | FileCheck %s --check-prefix=CHECK-FAIL-DYN +// Expected failure case, only for branch coverage +// CHECK-FAIL-DYN: error: 'hfusion.histogram' op output must have static shape +tt.func @test_histogram_fail_with_dyn_n(%arg0 : tensor<16xi32>) -> tensor { + %res = tt.histogram %arg0 : tensor<16xi32> -> tensor + tt.return %res : tensor +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToHFusion/mod.mlir b/third_party/ascend/unittest/Conversion/General/TritonToHFusion/mod.mlir new file mode 100644 index 0000000000..d4c264913a --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToHFusion/mod.mlir @@ -0,0 +1,11 @@ +// RUN: triton-opt %s -triton-to-hfusion | FileCheck %s + +// CHECK: tensor.empty() : tensor<1xf32> +// CHECK: hfusion.elemwise_binary {fun = #hfusion.binary_fn} ins(%arg0, %arg1 : tensor<1xf32>, tensor<1xf32>) outs(%0 : tensor<1xf32>) -> tensor<1xf32> + +module { + tt.func @test_mod(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> { + %0 = ascend.mod %arg0, %arg1 : tensor<1xf32> tensor<1xf32> -> tensor<1xf32> + tt.return %0 : tensor<1xf32> + } +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToHIVM/sync_block_op_conversion.mlir b/third_party/ascend/unittest/Conversion/General/TritonToHIVM/sync_block_op_conversion.mlir new file mode 100644 index 0000000000..e116f6d8b7 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToHIVM/sync_block_op_conversion.mlir @@ -0,0 +1,20 @@ +// RUN: triton-opt %s --triton-to-hivm | FileCheck %s + +// CHECK-LABEL: tt.func @triton_func +tt.func @triton_func() { + ascend.custom "sync_block_set" {str_args = ["vector", 1 : i32]} + ascend.custom "sync_block_wait" {str_args = ["vector", 1 : i32]} + ascend.custom "sync_block_set" {str_args = ["cube", 2 : i32]} + ascend.custom "sync_block_wait" {str_args = ["cube", 2 : i32]} + ascend.custom "sync_block_all" {str_args = ["all_cube", 1 : i32]} + ascend.custom "sync_block_all" {str_args = ["all_vector", 1 : i32]} + ascend.custom "sync_block_all" {str_args = ["all", 1 : i32]} + tt.return +} +// CHECK: hivm.hir.sync_block_set[, , ] flag = 1 +// CHECK: hivm.hir.sync_block_wait[, , ] flag = 1 +// CHECK: hivm.hir.sync_block_set[, , ] flag = 2 +// CHECK: hivm.hir.sync_block_wait[, , ] flag = 2 +// CHECK: hivm.hir.sync_block[, 1 : i16] tcube_pipe = +// CHECK: hivm.hir.sync_block[, 1 : i16] tvector_pipe = +// CHECK: hivm.hir.sync_block[, 1 : i16] tcube_pipe = tvector_pipe = diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/atomic_rmw.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/atomic_rmw.mlir new file mode 100644 index 0000000000..f3f36e5dcc --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/atomic_rmw.mlir @@ -0,0 +1,109 @@ +// RUN: triton-opt --triton-to-linalg="named-ops=True" --split-input-file %s | FileCheck %s +// CHECK-LABEL: func.func @matmul_atomic_add +// CHECK-NOT: GenericAtomicRMW +// CHECK: tensor.extract_slice +// CHECK: hivm.hir.store ins(%{{.*}} : tensor) outs(%{{.*}} : memref) atomic = + + tt.func public @matmul_atomic_add(%arg0: !tt.ptr, %arg1: !tt.ptr, %arg2: !tt.ptr, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32, %arg7: i32, %arg8: i32, %arg9: i32, %arg10: i32, %arg11: i32) attributes {noinline = false} { + %0 = tt.get_program_id x : i32 + %1 = tt.get_program_id y : i32 + %2 = tt.get_program_id z : i32 + %c16_i32 = arith.constant 16 : i32 + %c16_i32_0 = arith.constant 16 : i32 + %3 = arith.muli %0, %c16_i32_0 : i32 + %4 = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> + %5 = tt.splat %3 : i32 -> tensor<16xi32> + %6 = arith.addi %5, %4 : tensor<16xi32> + %c16_i32_1 = arith.constant 16 : i32 + %c16_i32_2 = arith.constant 16 : i32 + %7 = arith.muli %1, %c16_i32_2 : i32 + %8 = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> + %9 = tt.splat %7 : i32 -> tensor<16xi32> + %10 = arith.addi %9, %8 : tensor<16xi32> + %11 = tt.expand_dims %6 {axis = 1 : i32} : tensor<16xi32> -> tensor<16x1xi32> + %12 = tt.splat %arg10 : i32 -> tensor<16x1xi32> + %13 = arith.muli %11, %12 : tensor<16x1xi32> + %14 = tt.splat %arg2 : !tt.ptr -> tensor<16x1x!tt.ptr> + %15 = tt.addptr %14, %13 : tensor<16x1x!tt.ptr>, tensor<16x1xi32> + %16 = tt.expand_dims %10 {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32> + %17 = tt.splat %arg11 : i32 -> tensor<1x16xi32> + %18 = arith.muli %16, %17 : tensor<1x16xi32> + %19 = tt.broadcast %15 : tensor<16x1x!tt.ptr> -> tensor<16x16x!tt.ptr> + %20 = tt.broadcast %18 : tensor<1x16xi32> -> tensor<16x16xi32> + %21 = tt.addptr %19, %20 : tensor<16x16x!tt.ptr>, tensor<16x16xi32> + %22 = tt.expand_dims %6 {axis = 1 : i32} : tensor<16xi32> -> tensor<16x1xi32> + %23 = tt.splat %arg3 : i32 -> tensor<16x1xi32> + %24 = arith.cmpi slt, %22, %23 : tensor<16x1xi32> + %25 = tt.expand_dims %10 {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32> + %26 = tt.splat %arg4 : i32 -> tensor<1x16xi32> + %27 = arith.cmpi slt, %25, %26 : tensor<1x16xi32> + %28 = tt.broadcast %24 : tensor<16x1xi1> -> tensor<16x16xi1> + %29 = tt.broadcast %27 : tensor<1x16xi1> -> tensor<16x16xi1> + %30 = arith.andi %28, %29 : tensor<16x16xi1> + %c16_i32_3 = arith.constant 16 : i32 + %c16_i32_4 = arith.constant 16 : i32 + %31 = arith.muli %2, %c16_i32_4 : i32 + %c32_i32 = arith.constant 32 : i32 + %32 = arith.bitcast %31 : i32 to i32 + %33 = arith.bitcast %arg5 : i32 to i32 + %34 = arith.bitcast %c32_i32 : i32 to i32 + %35 = ub.poison : i32 + scf.for %arg12 = %32 to %33 step %34 : i32 { + %36 = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> + %37 = tt.splat %arg12 : i32 -> tensor<16xi32> + %38 = arith.addi %37, %36 : tensor<16xi32> + %39 = tt.expand_dims %6 {axis = 1 : i32} : tensor<16xi32> -> tensor<16x1xi32> + %40 = tt.splat %arg6 : i32 -> tensor<16x1xi32> + %41 = arith.muli %39, %40 : tensor<16x1xi32> + %42 = tt.splat %arg0 : !tt.ptr -> tensor<16x1x!tt.ptr> + %43 = tt.addptr %42, %41 : tensor<16x1x!tt.ptr>, tensor<16x1xi32> + %44 = tt.expand_dims %38 {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32> + %45 = tt.splat %arg7 : i32 -> tensor<1x16xi32> + %46 = arith.muli %44, %45 : tensor<1x16xi32> + %47 = tt.broadcast %43 : tensor<16x1x!tt.ptr> -> tensor<16x16x!tt.ptr> + %48 = tt.broadcast %46 : tensor<1x16xi32> -> tensor<16x16xi32> + %49 = tt.addptr %47, %48 : tensor<16x16x!tt.ptr>, tensor<16x16xi32> + %50 = tt.expand_dims %38 {axis = 1 : i32} : tensor<16xi32> -> tensor<16x1xi32> + %51 = tt.splat %arg8 : i32 -> tensor<16x1xi32> + %52 = arith.muli %50, %51 : tensor<16x1xi32> + %53 = tt.splat %arg1 : !tt.ptr -> tensor<16x1x!tt.ptr> + %54 = tt.addptr %53, %52 : tensor<16x1x!tt.ptr>, tensor<16x1xi32> + %55 = tt.expand_dims %10 {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32> + %56 = tt.splat %arg9 : i32 -> tensor<1x16xi32> + %57 = arith.muli %55, %56 : tensor<1x16xi32> + %58 = tt.broadcast %54 : tensor<16x1x!tt.ptr> -> tensor<16x16x!tt.ptr> + %59 = tt.broadcast %57 : tensor<1x16xi32> -> tensor<16x16xi32> + %60 = tt.addptr %58, %59 : tensor<16x16x!tt.ptr>, tensor<16x16xi32> + %61 = tt.expand_dims %6 {axis = 1 : i32} : tensor<16xi32> -> tensor<16x1xi32> + %62 = tt.splat %arg3 : i32 -> tensor<16x1xi32> + %63 = arith.cmpi slt, %61, %62 : tensor<16x1xi32> + %64 = tt.expand_dims %38 {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32> + %65 = tt.splat %arg5 : i32 -> tensor<1x16xi32> + %66 = arith.cmpi slt, %64, %65 : tensor<1x16xi32> + %67 = tt.broadcast %63 : tensor<16x1xi1> -> tensor<16x16xi1> + %68 = tt.broadcast %66 : tensor<1x16xi1> -> tensor<16x16xi1> + %69 = arith.andi %67, %68 : tensor<16x16xi1> + %70 = tt.expand_dims %38 {axis = 1 : i32} : tensor<16xi32> -> tensor<16x1xi32> + %71 = tt.splat %arg5 : i32 -> tensor<16x1xi32> + %72 = arith.cmpi slt, %70, %71 : tensor<16x1xi32> + %73 = tt.expand_dims %10 {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32> + %74 = tt.splat %arg4 : i32 -> tensor<1x16xi32> + %75 = arith.cmpi slt, %73, %74 : tensor<1x16xi32> + %76 = tt.broadcast %72 : tensor<16x1xi1> -> tensor<16x16xi1> + %77 = tt.broadcast %75 : tensor<1x16xi1> -> tensor<16x16xi1> + %78 = arith.andi %76, %77 : tensor<16x16xi1> + %c0_i32 = arith.constant 0 : i32 + %cst = arith.constant dense<0> : tensor<16x16xi32> + %79 = arith.sitofp %cst : tensor<16x16xi32> to tensor<16x16xf32> + %80 = tt.load %49, %69, %79 : tensor<16x16x!tt.ptr> + %c0_i32_5 = arith.constant 0 : i32 + %cst_6 = arith.constant dense<0> : tensor<16x16xi32> + %81 = arith.sitofp %cst_6 : tensor<16x16xi32> to tensor<16x16xf32> + %82 = tt.load %60, %78, %81 : tensor<16x16x!tt.ptr> + %cst_7 = arith.constant 0.000000e+00 : f32 + %cst_8 = arith.constant dense<0.000000e+00> : tensor<16x16xf32> + %83 = tt.dot %80, %82, %cst_8 : tensor<16x16xf32> * tensor<16x16xf32> -> tensor<16x16xf32> + %84 = tt.atomic_rmw fadd, acq_rel, gpu, %21, %83, %30 : (tensor<16x16x!tt.ptr>, tensor<16x16xf32>, tensor<16x16xi1>) -> tensor<16x16xf32> + } + tt.return + } diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/clamp_fp8.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/clamp_fp8.mlir new file mode 100644 index 0000000000..5f8baf1438 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/clamp_fp8.mlir @@ -0,0 +1,84 @@ +// RUN: triton-opt --discrete-mask-access-conversion="compile-on-910-95=True force-simt-template=False" --triton-to-annotation '--triton-to-unstructure=compile-on-910-95=True force-simt-template=False' --triton-to-hivm --triton-to-hfusion --triton-to-llvm --bubble-up-operation '--triton-to-linalg=global-kernel=false named-ops=True enable-nd2nz-on-vector=False compile-on-910-95=True' --split-input-file %s | FileCheck %s + +module { + tt.func public @tt_clamp_3d(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: !tt.ptr {tt.divisibility = 16 : i32}, %arg3: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %cst = arith.constant dense<4> : tensor<1x29x1xi32> + %cst_0 = arith.constant dense<4> : tensor<2x1x1xi32> + %cst_1 = arith.constant dense<29> : tensor<2x1x1xi32> + %c4_i32 = arith.constant 4 : i32 + %c29_i32 = arith.constant 29 : i32 + %c2_i32 = arith.constant 2 : i32 + %0 = tt.get_program_id x : i32 + %1 = arith.muli %0, %c2_i32 : i32 + %2 = tt.get_program_id y : i32 + %3 = arith.muli %2, %c29_i32 : i32 + %4 = tt.get_program_id z : i32 + %5 = arith.muli %4, %c4_i32 : i32 + %6 = tt.make_range {end = 2 : i32, start = 0 : i32} : tensor<2xi32> + %7 = tt.splat %1 : i32 -> tensor<2xi32> + %8 = arith.addi %6, %7 : tensor<2xi32> + %9 = tt.make_range {end = 29 : i32, start = 0 : i32} : tensor<29xi32> + %10 = tt.splat %3 : i32 -> tensor<29xi32> + %11 = arith.addi %9, %10 : tensor<29xi32> + %12 = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> + %13 = tt.splat %5 : i32 -> tensor<4xi32> + %14 = arith.addi %12, %13 : tensor<4xi32> + %15 = tt.expand_dims %8 {axis = 1 : i32} : tensor<2xi32> -> tensor<2x1xi32> + %16 = tt.expand_dims %15 {axis = 2 : i32} : tensor<2x1xi32> -> tensor<2x1x1xi32> + %17 = arith.muli %16, %cst_1 : tensor<2x1x1xi32> + %18 = arith.muli %17, %cst_0 : tensor<2x1x1xi32> + %19 = tt.expand_dims %11 {axis = 0 : i32} : tensor<29xi32> -> tensor<1x29xi32> + %20 = tt.expand_dims %19 {axis = 2 : i32} : tensor<1x29xi32> -> tensor<1x29x1xi32> + %21 = arith.muli %20, %cst : tensor<1x29x1xi32> + %22 = tt.broadcast %18 : tensor<2x1x1xi32> -> tensor<2x29x1xi32> + %23 = tt.broadcast %21 : tensor<1x29x1xi32> -> tensor<2x29x1xi32> + %24 = arith.addi %22, %23 : tensor<2x29x1xi32> + %25 = tt.expand_dims %14 {axis = 0 : i32} : tensor<4xi32> -> tensor<1x4xi32> + %26 = tt.expand_dims %25 {axis = 1 : i32} : tensor<1x4xi32> -> tensor<1x1x4xi32> + %27 = tt.broadcast %24 : tensor<2x29x1xi32> -> tensor<2x29x4xi32> + %28 = tt.broadcast %26 : tensor<1x1x4xi32> -> tensor<2x29x4xi32> + %29 = arith.addi %27, %28 : tensor<2x29x4xi32> + %30 = tt.splat %arg0 : !tt.ptr -> tensor<2x29x4x!tt.ptr> + %31 = tt.addptr %30, %29 : tensor<2x29x4x!tt.ptr>, tensor<2x29x4xi32> + %32 = tt.load %31 : tensor<2x29x4x!tt.ptr> + %33 = tt.splat %arg2 : !tt.ptr -> tensor<2x29x4x!tt.ptr> + %34 = tt.addptr %33, %29 : tensor<2x29x4x!tt.ptr>, tensor<2x29x4xi32> + %35 = tt.load %34 : tensor<2x29x4x!tt.ptr> + %36 = tt.splat %arg3 : !tt.ptr -> tensor<2x29x4x!tt.ptr> + %37 = tt.addptr %36, %29 : tensor<2x29x4x!tt.ptr>, tensor<2x29x4xi32> + %38 = tt.load %37 : tensor<2x29x4x!tt.ptr> + %39 = tt.fp_to_fp %32 : tensor<2x29x4xf8E4M3FN> -> tensor<2x29x4xf32> + %40 = tt.clampf %39, %35, %38, propagateNan = none : tensor<2x29x4xf32> + %41 = tt.splat %arg1 : !tt.ptr -> tensor<2x29x4x!tt.ptr> + %42 = tt.addptr %41, %29 : tensor<2x29x4x!tt.ptr>, tensor<2x29x4xi32> + %43 = tt.fp_to_fp %40, rounding = rtne : tensor<2x29x4xf32> -> tensor<2x29x4xf8E4M3FN> + tt.store %42, %43 : tensor<2x29x4x!tt.ptr> + tt.return + } +} + +// CHECK: %[[CAST_F8:[A-Za-z0-9_]+]] = memref.reinterpret_cast %arg2 to offset: [%9], sizes: [2, 29, 4], strides: [116, 4, 1] : memref to memref<2x29x4xf8E4M3FN, strided<[116, 4, 1], offset: ?>> +// CHECK: %[[ALLOC_F8:[A-Za-z0-9_]+]] = memref.alloc() : memref<2x29x4xf8E4M3FN> +// CHECK: memref.copy %[[CAST_F8]], %[[ALLOC_F8]] : memref<2x29x4xf8E4M3FN, strided<[116, 4, 1], offset: ?>> to memref<2x29x4xf8E4M3FN> +// CHECK: %[[TENSOR_F8:[A-Za-z0-9_]+]] = bufferization.to_tensor %[[ALLOC_F8]] restrict writable : memref<2x29x4xf8E4M3FN> + +// CHECK: %[[CAST_F32_MIN:[A-Za-z0-9_]+]] = memref.reinterpret_cast %arg4 to offset: [%9], sizes: [2, 29, 4], strides: [116, 4, 1] : memref to memref<2x29x4xf32, strided<[116, 4, 1], offset: ?>> +// CHECK: %[[ALLOC_F32_MIN:[A-Za-z0-9_]+]] = memref.alloc() : memref<2x29x4xf32> +// CHECK: memref.copy %[[CAST_F32_MIN]], %[[ALLOC_F32_MIN]] : memref<2x29x4xf32, strided<[116, 4, 1], offset: ?>> to memref<2x29x4xf32> +// CHECK: %[[TENSOR_F32_MIN:[A-Za-z0-9_]+]] = bufferization.to_tensor %[[ALLOC_F32_MIN]] restrict writable : memref<2x29x4xf32> + +// CHECK: %[[CAST_F32_MAX:[A-Za-z0-9_]+]] = memref.reinterpret_cast %arg5 to offset: [%9], sizes: [2, 29, 4], strides: [116, 4, 1] : memref to memref<2x29x4xf32, strided<[116, 4, 1], offset: ?>> +// CHECK: %[[ALLOC_F32_MAX:[A-Za-z0-9_]+]] = memref.alloc() : memref<2x29x4xf32> +// CHECK: memref.copy %[[CAST_F32_MAX]], %[[ALLOC_F32_MAX]] : memref<2x29x4xf32, strided<[116, 4, 1], offset: ?>> to memref<2x29x4xf32> +// CHECK: %[[TENSOR_F32_MAX:[A-Za-z0-9_]+]] = bufferization.to_tensor %[[ALLOC_F32_MAX]] restrict writable : memref<2x29x4xf32> + +// CHECK: %[[F8_TO_F32:[A-Za-z0-9_]+]] = arith.extf %[[TENSOR_F8]] {round_mode = #hfusion.round_mode} : tensor<2x29x4xf8E4M3FN> to tensor<2x29x4xf32> +// CHECK: %[[CLAMP_MAX:[A-Za-z0-9_]+]] = arith.minnumf %[[F8_TO_F32]], %[[TENSOR_F32_MAX]] : tensor<2x29x4xf32> +// CHECK: %[[CLAMP_RESULT:[A-Za-z0-9_]+]] = arith.maxnumf %[[TENSOR_F32_MIN]], %[[CLAMP_MAX]] : tensor<2x29x4xf32> + +// CHECK: %[[CAST_OUTPUT:[A-Za-z0-9_]+]] = memref.reinterpret_cast %arg3 to offset: [%9], sizes: [2, 29, 4], strides: [116, 4, 1] : memref to memref<2x29x4xf8E4M3FN, strided<[116, 4, 1], offset: ?>> +// CHECK: %[[F32_TO_F8:[A-Za-z0-9_]+]] = arith.truncf %[[CLAMP_RESULT]] {round_mode = #hfusion.round_mode} : tensor<2x29x4xf32> to tensor<2x29x4xf8E4M3FN> +// CHECK: bufferization.materialize_in_destination %[[F32_TO_F8]] in writable %[[CAST_OUTPUT]] : (tensor<2x29x4xf8E4M3FN>, memref<2x29x4xf8E4M3FN, strided<[116, 4, 1], offset: ?>>) -> () + + +// ----- diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/dotScaled_fp8.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/dotScaled_fp8.mlir new file mode 100644 index 0000000000..f43805d667 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/dotScaled_fp8.mlir @@ -0,0 +1,62 @@ +// RUN: triton-opt --pass-pipeline="builtin.module(triton-to-unstructure{compile-on-910-95=true force-simt-template=false},triton-to-linalg{compile-on-910-95=true enable-nd2nz-on-vector=false global-kernel=false named-ops=true})" --split-input-file %s | FileCheck %s + +module attributes {hacc.target = #hacc.target<"Ascend950PR_957c">} { + tt.func public @dot_scale_fp8_kernel(%a_base: !tt.ptr {tt.divisibility = 16 : i32}, %stride_a1: i32 {tt.divisibility = 16 : i32}, %a_scale: !tt.ptr {tt.divisibility = 16 : i32}, %b_base: !tt.ptr {tt.divisibility = 16 : i32}, %stride_b0: i32 {tt.divisibility = 16 : i32}, %out: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %out_ptr = arith.constant dense<32> : tensor<32x1xi32> + %c = arith.constant dense<0.000000e+00> : tensor<32x32xf32> + %scale_a_ptr = arith.constant dense<4> : tensor<32x1xi32> + %a_ptr = tt.make_range {end = 32 : i32, start = 0 : i32} : tensor<32xi32> + %a_ptr_0 = tt.expand_dims %a_ptr {axis = 1 : i32} : tensor<32xi32> -> tensor<32x1xi32> + %a_ptr_1 = tt.splat %a_base : !tt.ptr -> tensor<32x1x!tt.ptr> + %a_ptr_2 = tt.addptr %a_ptr_1, %a_ptr_0 : tensor<32x1x!tt.ptr>, tensor<32x1xi32> + %a_ptr_3 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32> + %a_ptr_4 = tt.expand_dims %a_ptr_3 {axis = 0 : i32} : tensor<128xi32> -> tensor<1x128xi32> + %a_ptr_5 = tt.splat %stride_a1 : i32 -> tensor<1x128xi32> + %a_ptr_6 = arith.muli %a_ptr_4, %a_ptr_5 : tensor<1x128xi32> + %a_ptr_7 = tt.broadcast %a_ptr_2 : tensor<32x1x!tt.ptr> -> tensor<32x128x!tt.ptr> + %a_ptr_8 = tt.broadcast %a_ptr_6 : tensor<1x128xi32> -> tensor<32x128xi32> + %a_ptr_9 = tt.addptr %a_ptr_7, %a_ptr_8 : tensor<32x128x!tt.ptr>, tensor<32x128xi32> + %b_ptr = tt.expand_dims %a_ptr_3 {axis = 1 : i32} : tensor<128xi32> -> tensor<128x1xi32> + %b_ptr_10 = tt.splat %stride_b0 : i32 -> tensor<128x1xi32> + %b_ptr_11 = arith.muli %b_ptr, %b_ptr_10 : tensor<128x1xi32> + %b_ptr_12 = tt.splat %b_base : !tt.ptr -> tensor<128x1x!tt.ptr> + %b_ptr_13 = tt.addptr %b_ptr_12, %b_ptr_11 : tensor<128x1x!tt.ptr>, tensor<128x1xi32> + %b_ptr_14 = tt.expand_dims %a_ptr {axis = 0 : i32} : tensor<32xi32> -> tensor<1x32xi32> + %b_ptr_15 = tt.broadcast %b_ptr_13 : tensor<128x1x!tt.ptr> -> tensor<128x32x!tt.ptr> + %b_ptr_16 = tt.broadcast %b_ptr_14 : tensor<1x32xi32> -> tensor<128x32xi32> + %b_ptr_17 = tt.addptr %b_ptr_15, %b_ptr_16 : tensor<128x32x!tt.ptr>, tensor<128x32xi32> + %scale_a_ptr_18 = arith.muli %a_ptr_0, %scale_a_ptr : tensor<32x1xi32> + %scale_a_ptr_19 = tt.splat %a_scale : !tt.ptr -> tensor<32x1x!tt.ptr> + %scale_a_ptr_20 = tt.addptr %scale_a_ptr_19, %scale_a_ptr_18 : tensor<32x1x!tt.ptr>, tensor<32x1xi32> + %scale_a_ptr_21 = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> + %scale_a_ptr_22 = tt.expand_dims %scale_a_ptr_21 {axis = 0 : i32} : tensor<4xi32> -> tensor<1x4xi32> + %scale_a_ptr_23 = tt.broadcast %scale_a_ptr_20 : tensor<32x1x!tt.ptr> -> tensor<32x4x!tt.ptr> + %scale_a_ptr_24 = tt.broadcast %scale_a_ptr_22 : tensor<1x4xi32> -> tensor<32x4xi32> + %scale_a_ptr_25 = tt.addptr %scale_a_ptr_23, %scale_a_ptr_24 : tensor<32x4x!tt.ptr>, tensor<32x4xi32> + %a = tt.load %a_ptr_9 : tensor<32x128x!tt.ptr> + %b = tt.load %b_ptr_17 : tensor<128x32x!tt.ptr> + %a_scale_26 = tt.load %scale_a_ptr_25 : tensor<32x4x!tt.ptr> + %c_27 = tt.bitcast %a : tensor<32x128xi8> -> tensor<32x128xf8E5M2> + %c_28 = tt.bitcast %b : tensor<128x32xi8> -> tensor<128x32xf8E4M3FN> + %c_29 = tt.dot_scaled %c_27 scale %a_scale_26, %c_28, %c lhs = e5m2 rhs = e4m3 {fastMath = false} : tensor<32x128xf8E5M2>, tensor<32x4xi8> * tensor<128x32xf8E4M3FN> -> tensor<32x32xf32> + %out_ptr_30 = arith.muli %a_ptr_0, %out_ptr : tensor<32x1xi32> + %out_ptr_31 = tt.splat %out : !tt.ptr -> tensor<32x1x!tt.ptr> + %out_ptr_32 = tt.addptr %out_ptr_31, %out_ptr_30 : tensor<32x1x!tt.ptr>, tensor<32x1xi32> + %out_ptr_33 = tt.broadcast %out_ptr_32 : tensor<32x1x!tt.ptr> -> tensor<32x32x!tt.ptr> + %out_ptr_34 = tt.broadcast %b_ptr_14 : tensor<1x32xi32> -> tensor<32x32xi32> + %out_ptr_35 = tt.addptr %out_ptr_33, %out_ptr_34 : tensor<32x32x!tt.ptr>, tensor<32x32xi32> + %0 = arith.truncf %c_29 : tensor<32x32xf32> to tensor<32x32xbf16> + tt.store %out_ptr_35, %0 : tensor<32x32x!tt.ptr> + tt.return + } +} + + +// CHECK-LABEL: func.func @dot_scale_fp8_kernel +// CHECK: hfusion.matmul_mx +// CHECK-DAG: tensor<32x128xf8E5M2> +// CHECK-DAG: tensor<128x32xf8E4M3FN> +// CHECK-DAG: tensor<32x4xi8> +// CHECK: outs([[OUT:.*]] : tensor<32x32xf32>) -> tensor<32x32xf32> +// CHECK: arith.truncf {{.*}} : tensor<32x32xf32> to tensor<32x32xbf16> +// CHECK: bufferization.materialize_in_destination diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/hoist_broadcast.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/hoist_broadcast.mlir new file mode 100644 index 0000000000..efb3abf186 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/hoist_broadcast.mlir @@ -0,0 +1,76 @@ +// RUN: triton-opt --triton-to-linalg="named-ops=True" --split-input-file %s | FileCheck %s + +module { +tt.func public @copy_all_layer_kv_cache2(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: !tt.ptr {tt.divisibility = 16 : i32}, %arg3: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %c0_i32 = arith.constant 0 : i32 + %cst = arith.constant dense<0> : tensor<1x128xi8> + %c1_i64 = arith.constant 1 : i64 + %c0_i64 = arith.constant 0 : i64 + %c127_i64 = arith.constant 127 : i64 + %true = arith.constant true + %c128_i64 = arith.constant 128 : i64 + + %0 = tt.get_program_id x : i32 + %1 = tt.addptr %arg1, %0 : !tt.ptr, i32 + %2 = tt.load %1 : !tt.ptr + %3 = tt.addptr %arg0, %0 : !tt.ptr, i32 + %4 = tt.load %3 : !tt.ptr + %5 = tt.int_to_ptr %4 : i64 -> !tt.ptr + + %6 = tt.addptr %arg2, %c0_i32 : !tt.ptr, i32 + %7 = tt.splat %6 : !tt.ptr -> tensor<1x!tt.ptr> + %8 = tt.load %7 : tensor<1x!tt.ptr> + + %9 = tt.addptr %arg3, %c0_i32 : !tt.ptr, i32 + %10 = tt.splat %9 : !tt.ptr -> tensor<1x!tt.ptr> + %11 = tt.load %10 : tensor<1x!tt.ptr> + + tt.assert %true, "int32 overflow detected for operation sub" : i1 + + %12 = arith.addi %2, %c127_i64 : i64 + %13 = arith.divsi %12, %c128_i64 : i64 + + %14 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32> + %15 = arith.extsi %14 : tensor<128xi32> to tensor<128xi64> + %16 = tt.splat %2 : i64 -> tensor<128xi64> + + %17 = tt.expand_dims %11 {axis = 1 : i32} : tensor<1xi32> -> tensor<1x1xi32> + %18 = arith.extsi %17 : tensor<1x1xi32> to tensor<1x1xi64> + %19 = tt.splat %2 : i64 -> tensor<1x1xi64> + %20 = arith.muli %18, %19 : tensor<1x1xi64> + %21 = tt.broadcast %20 : tensor<1x1xi64> -> tensor<1x128xi64> + %22 = tt.splat %5 : !tt.ptr -> tensor<1x128x!tt.ptr> + + %23 = tt.expand_dims %8 {axis = 1 : i32} : tensor<1xi32> -> tensor<1x1xi32> + %24 = arith.extsi %23 : tensor<1x1xi32> to tensor<1x1xi64> + %25 = arith.muli %24, %19 : tensor<1x1xi64> + %26 = tt.broadcast %25 : tensor<1x1xi64> -> tensor<1x128xi64> + + scf.for %arg4 = %c0_i64 to %13 step %c1_i64 : i64 { + %27 = arith.muli %arg4, %c128_i64 : i64 + %28 = tt.splat %27 : i64 -> tensor<128xi64> + %29 = arith.addi %15, %28 : tensor<128xi64> + %30 = arith.cmpi slt, %29, %16 : tensor<128xi64> + %31 = tt.expand_dims %30 {axis = 0 : i32} : tensor<128xi1> -> tensor<1x128xi1> + %32 = tt.expand_dims %29 {axis = 0 : i32} : tensor<128xi64> -> tensor<1x128xi64> + %33 = arith.addi %21, %32 : tensor<1x128xi64> + %34 = tt.addptr %22, %33 : tensor<1x128x!tt.ptr>, tensor<1x128xi64> + %35 = tt.load %34, %31, %cst : tensor<1x128x!tt.ptr> + %36 = arith.addi %26, %32 : tensor<1x128xi64> + %37 = tt.addptr %22, %36 : tensor<1x128x!tt.ptr>, tensor<1x128xi64> + tt.store %37, %35, %31 : tensor<1x128x!tt.ptr> + } + + tt.return +} +} + +// CHECK-LABEL: func.func @copy_all_layer_kv_cache2 +// CHECK: memref.reinterpret_cast +// CHECK: memref.load +// CHECK: memref.alloc +// CHECK: memref.copy +// CHECK: scf.for +// CHECK: linalg.fill +// CHECK: memref.subview +// CHECK: bufferization.materialize_in_destination diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/hoist_broadcast2.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/hoist_broadcast2.mlir new file mode 100644 index 0000000000..4e6f22ac92 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/hoist_broadcast2.mlir @@ -0,0 +1,145 @@ +// RUN: triton-opt --pass-pipeline="builtin.module(auto-blockify{auto-blockify-size=1},triton-to-structured{enable-mask-fallback-conversion=false optimize-dynamic-offset=false},discrete-mask-access-conversion{compile-on-910-95=true force-simt-template=false},triton-to-annotation,triton-to-unstructure{compile-on-910-95=true force-scalarize-mode=false force-simt-template=false},triton-to-hivm,triton-to-hfusion,triton-to-llvm,bubble-up-operation{enable-aggressive-mode=true},triton-to-structured{enable-mask-fallback-conversion=false optimize-dynamic-offset=false},triton-to-linalg{compile-on-910-95=true enable-nd2nz-on-vector=false enable-select-analysis=true global-kernel=false named-ops=true})" --split-input-file %s | FileCheck %s + +module { + tt.func public @_attn_bwd(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, + %arg1: !tt.ptr {tt.divisibility = 16 : i32}, + %arg2: !tt.ptr {tt.divisibility = 16 : i32}, + %arg3: !tt.ptr {tt.divisibility = 16 : i32}, + %arg4: !tt.ptr {tt.divisibility = 16 : i32}, + %arg5: !tt.ptr {tt.divisibility = 16 : i32}, + %arg6: !tt.ptr {tt.divisibility = 16 : i32}, + %arg7: !tt.ptr {tt.divisibility = 16 : i32}, + %arg8: !tt.ptr {tt.divisibility = 16 : i32}, + %arg9: f32) attributes {noinline = false} { + %cst = arith.constant dense<0.000000e+00> : tensor<128x128xf32> + %cst_0 = arith.constant dense<0.000000e+00> : tensor<128x64xf32> + %c1_i32 = arith.constant 1 : i32 + %c0_i32 = arith.constant 0 : i32 + %c28_i32 = arith.constant 28 : i32 + %c8192_i32 = arith.constant 8192 : i32 + %cst_1 = arith.constant dense : tensor<128x64xi1> + %cst_2 = arith.constant dense<64> : tensor<128x1xi32> + %c128_i32 = arith.constant 128 : i32 + %c524288_i32 = arith.constant 524288 : i32 + %c65536_i32 = arith.constant 65536 : i32 + %c1024_i32 = arith.constant 1024 : i32 + %c8_i32 = arith.constant 8 : i32 + %0 = tt.get_program_id x : i32 + %1 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32> + %2 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32> + %3 = tt.expand_dims %1 {axis = 0 : i32} : tensor<64xi32> -> tensor<1x64xi32> + %4 = tt.broadcast %3 : tensor<1x64xi32> -> tensor<128x64xi32> + %5 = tt.splat %arg9 : f32 -> tensor<128x128xf32> + scf.for %arg10 = %0 to %c8192_i32 step %c28_i32 : i32 { + %6 = arith.divsi %arg10, %c8_i32 : i32 + %7 = arith.muli %6, %c8_i32 : i32 + %8 = arith.subi %arg10, %7 : i32 + %9 = arith.muli %6, %c1024_i32 : i32 + %10 = arith.extsi %9 : i32 to i64 + %11 = arith.remsi %6, %c8_i32 : i32 + %12 = arith.muli %11, %c65536_i32 : i32 + %13 = arith.divsi %6, %c8_i32 : i32 + %14 = arith.muli %13, %c524288_i32 : i32 + %15 = arith.addi %12, %14 : i32 + %16 = arith.extsi %15 : i32 to i64 + %17 = tt.addptr %arg0, %16 : !tt.ptr, i64 + %18 = tt.addptr %arg1, %16 : !tt.ptr, i64 + %19 = tt.addptr %arg2, %16 : !tt.ptr, i64 + %20 = tt.addptr %arg3, %16 : !tt.ptr, i64 + %21 = tt.addptr %arg4, %16 : !tt.ptr, i64 + %22 = tt.addptr %arg5, %16 : !tt.ptr, i64 + %23 = tt.addptr %arg6, %16 : !tt.ptr, i64 + %24 = tt.addptr %arg7, %10 : !tt.ptr, i64 + %25 = tt.addptr %arg8, %10 : !tt.ptr, i64 + %26 = arith.muli %8, %c128_i32 : i32 + %27 = tt.splat %26 : i32 -> tensor<128xi32> + %28 = arith.addi %27, %2 : tensor<128xi32> + %29 = tt.expand_dims %28 {axis = 1 : i32} : tensor<128xi32> -> tensor<128x1xi32> + %30 = arith.muli %29, %cst_2 : tensor<128x1xi32> + %31 = tt.splat %18 : !tt.ptr -> tensor<128x1x!tt.ptr> + %32 = tt.addptr %31, %30 : tensor<128x1x!tt.ptr>, tensor<128x1xi32> + %33 = tt.broadcast %32 : tensor<128x1x!tt.ptr> -> tensor<128x64x!tt.ptr> + %34 = tt.addptr %33, %4 : tensor<128x64x!tt.ptr>, tensor<128x64xi32> + %35 = tt.load %34 : tensor<128x64x!tt.ptr> + %36 = tt.splat %19 : !tt.ptr -> tensor<128x1x!tt.ptr> + %37 = tt.addptr %36, %30 : tensor<128x1x!tt.ptr>, tensor<128x1xi32> + %38 = tt.broadcast %37 : tensor<128x1x!tt.ptr> -> tensor<128x64x!tt.ptr> + %39 = tt.addptr %38, %4 : tensor<128x64x!tt.ptr>, tensor<128x64xi32> + %40 = tt.load %39 : tensor<128x64x!tt.ptr> + %41 = tt.splat %17 : !tt.ptr -> tensor<128x1x!tt.ptr> + %42 = tt.splat %20 : !tt.ptr -> tensor<128x1x!tt.ptr> + %43 = tt.splat %24 : !tt.ptr -> tensor<128x!tt.ptr> + %44 = tt.splat %25 : !tt.ptr -> tensor<128x!tt.ptr> + %45 = tt.trans %35 {order = array} : tensor<128x64xf16> -> tensor<64x128xf16> + %46 = tt.trans %40 {order = array} : tensor<128x64xf16> -> tensor<64x128xf16> + %47 = tt.splat %21 : !tt.ptr -> tensor<128x1x!tt.ptr> + %48:2 = scf.for %arg11 = %c0_i32 to %c8_i32 step %c1_i32 iter_args(%arg12 = %cst_0, %arg13 = %cst_0) -> (tensor<128x64xf32>, tensor<128x64xf32>) : i32 { + %59 = arith.muli %arg11, %c128_i32 : i32 + %60 = tt.splat %59 : i32 -> tensor<128xi32> + %61 = arith.addi %60, %2 : tensor<128xi32> + %62 = tt.expand_dims %61 {axis = 1 : i32} : tensor<128xi32> -> tensor<128x1xi32> + %63 = arith.muli %62, %cst_2 : tensor<128x1xi32> + %64 = tt.addptr %41, %63 : tensor<128x1x!tt.ptr>, tensor<128x1xi32> + %65 = tt.broadcast %64 : tensor<128x1x!tt.ptr> -> tensor<128x64x!tt.ptr> + %66 = tt.addptr %65, %4 : tensor<128x64x!tt.ptr>, tensor<128x64xi32> + %67 = tt.load %66 : tensor<128x64x!tt.ptr> + %68 = tt.addptr %42, %63 : tensor<128x1x!tt.ptr>, tensor<128x1xi32> + %69 = tt.broadcast %68 : tensor<128x1x!tt.ptr> -> tensor<128x64x!tt.ptr> + %70 = tt.addptr %69, %4 : tensor<128x64x!tt.ptr>, tensor<128x64xi32> + %71 = tt.load %70 : tensor<128x64x!tt.ptr> + %72 = tt.addptr %43, %61 : tensor<128x!tt.ptr>, tensor<128xi32> + %73 = tt.load %72 : tensor<128x!tt.ptr> + %74 = tt.addptr %44, %61 : tensor<128x!tt.ptr>, tensor<128xi32> + %75 = tt.load %74 : tensor<128x!tt.ptr> + %76 = tt.dot %67, %45, %cst : tensor<128x64xf16> * tensor<64x128xf16> -> tensor<128x128xf32> + %77 = arith.mulf %76, %5 : tensor<128x128xf32> + %78 = tt.expand_dims %73 {axis = 1 : i32} : tensor<128xf32> -> tensor<128x1xf32> + %79 = tt.broadcast %78 : tensor<128x1xf32> -> tensor<128x128xf32> + %80 = arith.subf %77, %79 : tensor<128x128xf32> + %81 = math.exp %80 : tensor<128x128xf32> + %82 = arith.truncf %81 : tensor<128x128xf32> to tensor<128x128xf16> + %83 = tt.trans %82 {order = array} : tensor<128x128xf16> -> tensor<128x128xf16> + %84 = tt.dot %83, %71, %arg13 : tensor<128x128xf16> * tensor<128x64xf16> -> tensor<128x64xf32> + %85 = tt.dot %71, %46, %cst : tensor<128x64xf16> * tensor<64x128xf16> -> tensor<128x128xf32> + %86 = tt.expand_dims %75 {axis = 1 : i32} : tensor<128xf32> -> tensor<128x1xf32> + %87 = tt.broadcast %86 : tensor<128x1xf32> -> tensor<128x128xf32> + %88 = arith.subf %85, %87 : tensor<128x128xf32> + %89 = arith.extf %82 : tensor<128x128xf16> to tensor<128x128xf32> + %90 = arith.mulf %89, %88 : tensor<128x128xf32> + %91 = arith.mulf %90, %5 : tensor<128x128xf32> + %92 = arith.truncf %91 : tensor<128x128xf32> to tensor<128x128xf16> + %93 = tt.trans %92 {order = array} : tensor<128x128xf16> -> tensor<128x128xf16> + %94 = tt.dot %93, %67, %arg12 : tensor<128x128xf16> * tensor<128x64xf16> -> tensor<128x64xf32> + %95 = tt.dot %92, %35, %cst_0 : tensor<128x128xf16> * tensor<128x64xf16> -> tensor<128x64xf32> + %96 = tt.addptr %47, %63 : tensor<128x1x!tt.ptr>, tensor<128x1xi32> + %97 = tt.broadcast %96 : tensor<128x1x!tt.ptr> -> tensor<128x64x!tt.ptr> + %98 = tt.addptr %97, %4 : tensor<128x64x!tt.ptr>, tensor<128x64xi32> + %99 = arith.truncf %95 : tensor<128x64xf32> to tensor<128x64xf16> + %100 = tt.atomic_rmw fadd, acq_rel, gpu, %98, %99, %cst_1 : (tensor<128x64x!tt.ptr>, tensor<128x64xf16>, tensor<128x64xi1>) -> tensor<128x64xf16> + scf.yield %94, %84 : tensor<128x64xf32>, tensor<128x64xf32> + } + %49 = tt.splat %22 : !tt.ptr -> tensor<128x1x!tt.ptr> + %50 = tt.addptr %49, %30 : tensor<128x1x!tt.ptr>, tensor<128x1xi32> + %51 = tt.broadcast %50 : tensor<128x1x!tt.ptr> -> tensor<128x64x!tt.ptr> + %52 = tt.addptr %51, %4 : tensor<128x64x!tt.ptr>, tensor<128x64xi32> + %53 = arith.truncf %48#0 : tensor<128x64xf32> to tensor<128x64xf16> + tt.store %52, %53 : tensor<128x64x!tt.ptr> + %54 = tt.splat %23 : !tt.ptr -> tensor<128x1x!tt.ptr> + %55 = tt.addptr %54, %30 : tensor<128x1x!tt.ptr>, tensor<128x1xi32> + %56 = tt.broadcast %55 : tensor<128x1x!tt.ptr> -> tensor<128x64x!tt.ptr> + %57 = tt.addptr %56, %4 : tensor<128x64x!tt.ptr>, tensor<128x64xi32> + %58 = arith.truncf %48#1 : tensor<128x64xf32> to tensor<128x64xf16> + tt.store %57, %58 : tensor<128x64x!tt.ptr> + } + tt.return + } +} + +// CHECK-LABEL: func.func @_attn_bwd +// CHECK: linalg.fill +// CHECK: memref.reinterpret_cast +// CHECK: memref.alloc +// CHECK: memref.copy +// CHECK: scf.for +// CHECK: hivm.hir.store +// CHECK: bufferization.materialize_in_destination diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/parse_custom.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/parse_custom.mlir new file mode 100644 index 0000000000..88a7949efe --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/parse_custom.mlir @@ -0,0 +1,69 @@ +// RUN: triton-opt --triton-to-unstructure --triton-to-linalg="named-ops=True" --split-input-file %s | FileCheck %s +// CHECK-LABEL: func.func @parse_custom +// CHECK: %4 = hivm.hir.custom +// CHECK: %reinterpret_cast = memref.reinterpret_cast %4 to offset: [0], sizes: [32], strides: [1] : memref to memref<32xf16, strided<[1]>> +// CHECK: %alloc = memref.alloc() : memref<32xf16> +// CHECK: memref.copy %reinterpret_cast, %alloc : memref<32xf16, strided<[1]>> to memref<32xf16> + +module { + tt.func public @parse_custom(%arg0: !tt.ptr, %arg1: !tt.ptr, %arg2: i32, %arg3: i32) { + %c0_i64 = arith.constant 0 : i64 + %c1_i32 = arith.constant 1 : i32 + %0 = hivm.hir.get_sub_block_idx -> i64 + %1 = arith.cmpi eq, %0, %c0_i64 : i64 + scf.if %1 { + %2 = arith.addi %arg2, %c1_i32 : i32 + %3 = arith.remsi %2, %arg3 : i32 + %4 = hivm.hir.custom {hivm.pipe = #hivm.pipe, hivm.tcore_type = #hivm.tcore_type, hivm.vf_mode = #hivm.vf_mode, symbol = "foo1"} "foo1" ins(%arg0, %3 : !tt.ptr, i32) -> !tt.ptr + %5 = tt.make_range {end = 32 : i32, start = 0 : i32} : tensor<32xi32> + %6 = tt.splat %4 : !tt.ptr -> tensor<32x!tt.ptr> + %7 = tt.addptr %6, %5 : tensor<32x!tt.ptr>, tensor<32xi32> + %8 = tt.load %7 : tensor<32x!tt.ptr> + %9 = tt.splat %arg1 : !tt.ptr -> tensor<32x!tt.ptr> + %10 = tt.addptr %9, %5 : tensor<32x!tt.ptr>, tensor<32xi32> + tt.store %10, %8 : tensor<32x!tt.ptr> + } + hivm.hir.custom {hivm.pipe = #hivm.pipe, hivm.tcore_type = #hivm.tcore_type, hivm.vf_mode = #hivm.vf_mode, symbol = "foo2"} "foo2" + tt.return + } +} + +// ----- + +// CHECK-LABEL: func.func @parse_custom +// CHECK: %4 = hivm.hir.custom {{.*}} "foo1" ins(%reinterpret_cast, %arg7, %c0_i64 : memref<1xi64, strided<[1]>>, i32, i64) -> i32 +// CHECK: %5 = hivm.hir.custom {{.*}} "foo2" ins(%arg2, %3 : memref, i32) -> memref +// CHECK: %reinterpret_cast_0 = memref.reinterpret_cast %5 to offset: [0], sizes: [32], strides: [1] : memref to memref<32xf16, strided<[1]>> +// CHECK: %6 = hivm.hir.custom {{.*}} "foo3" ins(%reinterpret_cast_0, %4 : memref<32xf16, strided<[1]>>, i32) -> memref<32xf16> +// CHECK: %reinterpret_cast_1 = memref.reinterpret_cast %6 to offset: [0], sizes: [32], strides: [1] : memref<32xf16> to memref<32xf16, strided<[1]>> +// CHECK: %alloc = memref.alloc() : memref<32xf16> +// CHECK: memref.copy %reinterpret_cast_1, %alloc : memref<32xf16, strided<[1]>> to memref<32xf16> +module { + tt.func public @parse_custom(%arg0: !tt.ptr, %arg1: !tt.ptr, %arg2: i32, %arg3: i32, %arg4: !tt.ptr) { + %c0_i64 = arith.constant 0 : i64 + %c0_i32 = arith.constant 0 : i32 + %c1_i32 = arith.constant 1 : i32 + %0 = hivm.hir.get_sub_block_idx -> i64 + %1 = arith.cmpi eq, %0, %c0_i64 : i64 + scf.if %1 { + %2 = arith.addi %arg2, %c1_i32 : i32 + %3 = arith.remsi %2, %arg3 : i32 + %4 = tt.get_num_programs x : i32 + %5 = tt.addptr %arg4, %c0_i32 : !tt.ptr, i32 + %6 = hivm.hir.custom {hivm.pipe = #hivm.pipe, hivm.tcore_type = #hivm.tcore_type, hivm.vf_mode = #hivm.vf_mode, symbol = "foo1"} "foo1" ins(%5, %4, %c0_i64 : !tt.ptr, i32, i64) -> i32 + %7 = hivm.hir.custom {hivm.pipe = #hivm.pipe, hivm.tcore_type = #hivm.tcore_type, hivm.vf_mode = #hivm.vf_mode, symbol = "foo2"} "foo2" ins(%arg0, %3 : !tt.ptr, i32) -> !tt.ptr + %8 = tt.make_range {end = 32 : i32, start = 0 : i32} : tensor<32xi32> + %9 = tt.splat %7 : !tt.ptr -> tensor<32x!tt.ptr> + %10 = tt.addptr %9, %8 : tensor<32x!tt.ptr>, tensor<32xi32> + %11 = tensor.empty() : tensor<32x!tt.ptr> + %12 = hivm.hir.custom {SrcPtrIndex = array, hivm.pipe = #hivm.pipe, hivm.tcore_type = #hivm.tcore_type, hivm.vf_mode = #hivm.vf_mode, symbol = "foo3"} "foo3" ins(%10, %6 : tensor<32x!tt.ptr>, i32) outs(%11 : tensor<32x!tt.ptr>) -> tensor<32x!tt.ptr> + annotation.mark %12 {ContinuousMemAccess} : tensor<32x!tt.ptr> + %13 = tt.load %12 : tensor<32x!tt.ptr> + %14 = tt.splat %arg1 : !tt.ptr -> tensor<32x!tt.ptr> + %15 = tt.addptr %14, %8 : tensor<32x!tt.ptr>, tensor<32xi32> + tt.store %15, %13 : tensor<32x!tt.ptr> + } + hivm.hir.custom {hivm.pipe = #hivm.pipe, hivm.tcore_type = #hivm.tcore_type, hivm.vf_mode = #hivm.vf_mode, symbol = "foo4"} "foo4" + tt.return + } +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/parse_select.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/parse_select.mlir new file mode 100644 index 0000000000..f844ca90d5 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/parse_select.mlir @@ -0,0 +1,37 @@ +// RUN: triton-opt --triton-to-linalg --split-input-file %s -verify-each 2>&1 | FileCheck %s --check-prefix=NOERR +// NOERR-NOT: parseSelect currently supports all-ones shape unless cond=i1 with dense constants +// CHECK-LABEL: func.func public @triton_for_if_load + +module { + tt.func public @triton_for_if_load(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %cst = arith.constant dense<0> : tensor<16xi32> + %c0_i32 = arith.constant 0 : i32 + %c16_i32 = arith.constant 16 : i32 + %cst_0 = arith.constant dense<1> : tensor<16xi32> + %cst_1 = arith.constant dense<32> : tensor<16xi32> + %cst_2 = arith.constant dense<0.000000e+00> : tensor<16xf32> + %c2_i32 = arith.constant 2 : i32 + %c1_i32 = arith.constant 1 : i32 + %0 = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> + %1 = tt.get_program_id x : i32 + %2 = arith.cmpi ne, %1, %c0_i32 : i32 + %3 = tt.splat %arg0 : !tt.ptr -> tensor<16x!tt.ptr> + %4 = tt.splat %arg1 : !tt.ptr -> tensor<16x!tt.ptr> + %5:2 = scf.for %arg2 = %c0_i32 to %c2_i32 step %c1_i32 iter_args(%arg3 = %0, %arg4 = %0) -> (tensor<16xi32>, tensor<16xi32>) : i32 { + %6 = arith.muli %arg2, %c16_i32 : i32 + %7 = tt.splat %6 : i32 -> tensor<16xi32> + %8 = arith.addi %arg3, %7 : tensor<16xi32> + %9 = arith.addi %arg4, %7 : tensor<16xi32> + %10 = arith.select %2, %cst_0, %cst : tensor<16xi32> + %11 = arith.addi %8, %10 : tensor<16xi32> + %12 = tt.addptr %3, %11 : tensor<16x!tt.ptr>, tensor<16xi32> + %13 = arith.cmpi slt, %11, %cst_1 : tensor<16xi32> + %14 = tt.load %12 : tensor<16x!tt.ptr> + %15 = arith.select %13, %14, %cst_2 : tensor<16xi1>, tensor<16xf32> + %16 = tt.addptr %4, %9 : tensor<16x!tt.ptr>, tensor<16xi32> + tt.store %16, %15 : tensor<16x!tt.ptr> + scf.yield %11, %9 : tensor<16xi32>, tensor<16xi32> + } + tt.return + } +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_implicit_permute.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_implicit_permute.mlir new file mode 100644 index 0000000000..e7593291d6 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_implicit_permute.mlir @@ -0,0 +1,224 @@ +// RUN: triton-opt '--triton-to-linalg=global-kernel=False named-ops=True enable-nd2nz-on-vector=False enable-select-analysis=True compile-on-910-95=True' --split-input-file %s | FileCheck %s + +// CHECK-LABEL: func.func @addptr_implicit_perm_load_2d( +// CHECK: %[[IN_STRIDE0:[A-Za-z0-9_]+]] = arith.index_cast %arg6 : i32 to index +// CHECK: %[[IN_CAST0:[A-Za-z0-9_]+]] = memref.reinterpret_cast %arg2 to offset: [%[[IN_OFF0:[A-Za-z0-9_]+]]], sizes: [4, 64], strides: [1, %[[IN_STRIDE0]]] : memref to memref<4x64xf32, strided<[1, ?], offset: ?>> +// CHECK: memref.copy %{{.*}}, %{{.*}} : memref> to memref> +// CHECK: %[[TENSOR0:[A-Za-z0-9_]+]] = bufferization.to_tensor %{{.*}} restrict writable : memref<4x64xf32> +// CHECK: %[[OUT_STRIDE0:[A-Za-z0-9_]+]] = arith.index_cast %arg7 : i32 to index +// CHECK: %[[OUT_CAST0:[A-Za-z0-9_]+]] = memref.reinterpret_cast %arg3 to offset: [%[[OUT_OFF0:[A-Za-z0-9_]+]]], sizes: [4, 64], strides: [%[[OUT_STRIDE0]], 1] : memref to memref<4x64xf32, strided<[?, 1], offset: ?>> +// CHECK: %[[SLICE0:[A-Za-z0-9_]+]] = tensor.extract_slice %[[TENSOR0]][0, 0] [%[[ROWS0:[A-Za-z0-9_]+]], %[[COLS0:[A-Za-z0-9_]+]]] [1, 1] : tensor<4x64xf32> to tensor +// CHECK-NOT: linalg.transpose + +// CHECK-LABEL: func.func @addptr_implicit_perm_store_2d( +// CHECK: %[[IN_STRIDE1:[A-Za-z0-9_]+]] = arith.index_cast %arg6 : i32 to index +// CHECK: %[[IN_CAST1:[A-Za-z0-9_]+]] = memref.reinterpret_cast %arg2 to offset: [%[[IN_OFF1:[A-Za-z0-9_]+]]], sizes: [4, 64], strides: [%[[IN_STRIDE1]], 1] : memref to memref<4x64xf32, strided<[?, 1], offset: ?>> +// CHECK: memref.copy %{{.*}}, %{{.*}} : memref> to memref> +// CHECK: %[[TENSOR1:[A-Za-z0-9_]+]] = bufferization.to_tensor %{{.*}} restrict writable : memref<4x64xf32> +// CHECK: %[[OUT_STRIDE1:[A-Za-z0-9_]+]] = arith.index_cast %arg7 : i32 to index +// CHECK: %[[OUT_CAST1:[A-Za-z0-9_]+]] = memref.reinterpret_cast %arg3 to offset: [%[[OUT_OFF1:[A-Za-z0-9_]+]]], sizes: [64, 4], strides: [%[[OUT_STRIDE1]], 1] : memref to memref<64x4xf32, strided<[?, 1], offset: ?>> +// CHECK: %[[EMPTY1:[A-Za-z0-9_]+]] = tensor.empty() : tensor<64x4xf32> +// CHECK-NEXT: %[[TRANS1:[A-Za-z0-9_]+]] = linalg.transpose ins(%[[TENSOR1]] : tensor<4x64xf32>) outs(%[[EMPTY1]] : tensor<64x4xf32>) permutation = [1, 0] +// CHECK: %[[SLICE1:[A-Za-z0-9_]+]] = tensor.extract_slice %[[TRANS1]][0, 0] [%[[COLS1:[A-Za-z0-9_]+]], %[[ROWS1:[A-Za-z0-9_]+]]] [1, 1] : tensor<64x4xf32> to tensor + +// CHECK-LABEL: func.func @make_tensor_ptr_implicit_perm_load_2d( +// CHECK: %[[IN_STRIDE2:[A-Za-z0-9_]+]] = arith.index_cast %arg6 : i32 to index +// CHECK: %[[IN_CAST2:[A-Za-z0-9_]+]] = memref.reinterpret_cast %arg2 to offset: [%[[IN_OFF2:[A-Za-z0-9_]+]]], sizes: [4, 64], strides: [1, %[[IN_STRIDE2]]] : memref to memref<4x64xf32, strided<[1, ?], offset: ?>> +// CHECK: memref.copy %[[IN_CAST2]], %{{.*}} : memref<4x64xf32, strided<[1, ?], offset: ?>> to memref<4x64xf32> +// CHECK: %[[TENSOR2:[A-Za-z0-9_]+]] = bufferization.to_tensor %{{.*}} restrict writable : memref<4x64xf32> +// CHECK: %[[OUT_STRIDE2:[A-Za-z0-9_]+]] = arith.index_cast %arg7 : i32 to index +// CHECK: %[[OUT_CAST2:[A-Za-z0-9_]+]] = memref.reinterpret_cast %arg3 to offset: [%[[OUT_OFF2:[A-Za-z0-9_]+]]], sizes: [4, 64], strides: [%[[OUT_STRIDE2]], 1] : memref to memref<4x64xf32, strided<[?, 1], offset: ?>> +// CHECK: %[[SLICE2:[A-Za-z0-9_]+]] = tensor.extract_slice %[[TENSOR2]][0, 0] [%[[ROWS2:[A-Za-z0-9_]+]], %[[COLS2:[A-Za-z0-9_]+]]] [1, 1] : tensor<4x64xf32> to tensor +// CHECK-NOT: linalg.transpose + +// CHECK-LABEL: func.func @make_tensor_ptr_implicit_perm_store_2d( +// CHECK: %[[IN_STRIDE3:[A-Za-z0-9_]+]] = arith.index_cast %arg6 : i32 to index +// CHECK: %[[IN_CAST3:[A-Za-z0-9_]+]] = memref.reinterpret_cast %arg2 to offset: [%[[IN_OFF3:[A-Za-z0-9_]+]]], sizes: [4, 64], strides: [%[[IN_STRIDE3]], 1] : memref to memref<4x64xf32, strided<[?, 1], offset: ?>> +// CHECK: memref.copy %{{.*}}, %{{.*}} : memref> to memref> +// CHECK: %[[TENSOR3:[A-Za-z0-9_]+]] = bufferization.to_tensor %{{.*}} restrict writable : memref<4x64xf32> +// CHECK: %[[OUT_STRIDE3:[A-Za-z0-9_]+]] = arith.index_cast %arg7 : i32 to index +// CHECK: %[[OUT_CAST3:[A-Za-z0-9_]+]] = memref.reinterpret_cast %arg3 to offset: [%[[OUT_OFF3:[A-Za-z0-9_]+]]], sizes: [64, 4], strides: [%[[OUT_STRIDE3]], 1] : memref to memref<64x4xf32, strided<[?, 1], offset: ?>> +// CHECK: %[[EMPTY3:[A-Za-z0-9_]+]] = tensor.empty() : tensor<64x4xf32> +// CHECK-NEXT: %[[TRANS3:[A-Za-z0-9_]+]] = linalg.transpose ins(%[[TENSOR3]] : tensor<4x64xf32>) outs(%[[EMPTY3]] : tensor<64x4xf32>) permutation = [1, 0] +// CHECK: %[[SLICE3:[A-Za-z0-9_]+]] = tensor.extract_slice %[[TRANS3]]{{.*}} : tensor<64x4xf32> to tensor + +// addptr load implicit permutation +module { + tt.func public @addptr_implicit_perm_load_2d(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: i32 {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}, %arg4: i32 {tt.divisibility = 16 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %cst = arith.constant dense<0.000000e+00> : tensor<4x64xf32> + %c64_i32 = arith.constant 64 : i32 + %c4_i32 = arith.constant 4 : i32 + %0 = tt.get_program_id x : i32 + %1 = arith.muli %0, %c4_i32 : i32 + %2 = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> + %3 = tt.expand_dims %2 {axis = 1 : i32} : tensor<4xi32> -> tensor<4x1xi32> + %4 = tt.splat %1 : i32 -> tensor<4x1xi32> + %5 = arith.addi %4, %3 : tensor<4x1xi32> + %6 = tt.get_program_id y : i32 + %7 = arith.muli %6, %c64_i32 : i32 + %8 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32> + %9 = tt.expand_dims %8 {axis = 0 : i32} : tensor<64xi32> -> tensor<1x64xi32> + %10 = tt.splat %7 : i32 -> tensor<1x64xi32> + %11 = arith.addi %10, %9 : tensor<1x64xi32> + %12 = tt.splat %arg3 : i32 -> tensor<4x1xi32> + %13 = arith.cmpi slt, %5, %12 : tensor<4x1xi32> + %14 = tt.splat %arg2 : i32 -> tensor<1x64xi32> + %15 = arith.cmpi slt, %11, %14 : tensor<1x64xi32> + %16 = tt.broadcast %13 : tensor<4x1xi1> -> tensor<4x64xi1> + %17 = tt.broadcast %15 : tensor<1x64xi1> -> tensor<4x64xi1> + %18 = arith.andi %16, %17 : tensor<4x64xi1> + %19 = tt.splat %arg4 : i32 -> tensor<1x64xi32> + %20 = arith.muli %11, %19 : tensor<1x64xi32> + %21 = tt.broadcast %5 : tensor<4x1xi32> -> tensor<4x64xi32> + %22 = tt.broadcast %20 : tensor<1x64xi32> -> tensor<4x64xi32> + %23 = arith.addi %21, %22 : tensor<4x64xi32> + %24 = tt.splat %arg0 : !tt.ptr -> tensor<4x64x!tt.ptr> + %25 = tt.addptr %24, %23 : tensor<4x64x!tt.ptr>, tensor<4x64xi32> + %26 = tt.load %25, %18, %cst : tensor<4x64x!tt.ptr> + %27 = tt.splat %arg5 : i32 -> tensor<4x1xi32> + %28 = arith.muli %5, %27 : tensor<4x1xi32> + %29 = tt.broadcast %28 : tensor<4x1xi32> -> tensor<4x64xi32> + %30 = tt.broadcast %11 : tensor<1x64xi32> -> tensor<4x64xi32> + %31 = arith.addi %29, %30 : tensor<4x64xi32> + %32 = tt.splat %arg1 : !tt.ptr -> tensor<4x64x!tt.ptr> + %33 = tt.addptr %32, %31 : tensor<4x64x!tt.ptr>, tensor<4x64xi32> + tt.store %33, %26, %18 : tensor<4x64x!tt.ptr> + tt.return + } +} + +// ----- + +// addptr store implicit permutation +module { + tt.func public @addptr_implicit_perm_store_2d(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: i32 {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}, %arg4: i32 {tt.divisibility = 16 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %cst = arith.constant dense<0.000000e+00> : tensor<4x64xf32> + %c64_i32 = arith.constant 64 : i32 + %c4_i32 = arith.constant 4 : i32 + %0 = tt.get_program_id x : i32 + %1 = arith.muli %0, %c4_i32 : i32 + %2 = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> + %3 = tt.expand_dims %2 {axis = 1 : i32} : tensor<4xi32> -> tensor<4x1xi32> + %4 = tt.splat %1 : i32 -> tensor<4x1xi32> + %5 = arith.addi %4, %3 : tensor<4x1xi32> + %6 = tt.get_program_id y : i32 + %7 = arith.muli %6, %c64_i32 : i32 + %8 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32> + %9 = tt.expand_dims %8 {axis = 0 : i32} : tensor<64xi32> -> tensor<1x64xi32> + %10 = tt.splat %7 : i32 -> tensor<1x64xi32> + %11 = arith.addi %10, %9 : tensor<1x64xi32> + %12 = tt.splat %arg3 : i32 -> tensor<4x1xi32> + %13 = arith.cmpi slt, %5, %12 : tensor<4x1xi32> + %14 = tt.splat %arg2 : i32 -> tensor<1x64xi32> + %15 = arith.cmpi slt, %11, %14 : tensor<1x64xi32> + %16 = tt.broadcast %13 : tensor<4x1xi1> -> tensor<4x64xi1> + %17 = tt.broadcast %15 : tensor<1x64xi1> -> tensor<4x64xi1> + %18 = arith.andi %16, %17 : tensor<4x64xi1> + %19 = tt.splat %arg4 : i32 -> tensor<4x1xi32> + %20 = arith.muli %5, %19 : tensor<4x1xi32> + %21 = tt.broadcast %20 : tensor<4x1xi32> -> tensor<4x64xi32> + %22 = tt.broadcast %11 : tensor<1x64xi32> -> tensor<4x64xi32> + %23 = arith.addi %21, %22 : tensor<4x64xi32> + %24 = tt.splat %arg0 : !tt.ptr -> tensor<4x64x!tt.ptr> + %25 = tt.addptr %24, %23 : tensor<4x64x!tt.ptr>, tensor<4x64xi32> + %26 = tt.load %25, %18, %cst : tensor<4x64x!tt.ptr> + %27 = tt.splat %arg5 : i32 -> tensor<1x64xi32> + %28 = arith.muli %11, %27 : tensor<1x64xi32> + %29 = tt.broadcast %5 : tensor<4x1xi32> -> tensor<4x64xi32> + %30 = tt.broadcast %28 : tensor<1x64xi32> -> tensor<4x64xi32> + %31 = arith.addi %29, %30 : tensor<4x64xi32> + %32 = tt.splat %arg1 : !tt.ptr -> tensor<4x64x!tt.ptr> + %33 = tt.addptr %32, %31 : tensor<4x64x!tt.ptr>, tensor<4x64xi32> + tt.store %33, %26, %18 : tensor<4x64x!tt.ptr> + tt.return + } +} + +// ----- + +// make_tensor_ptr load implicit permutation +module { + tt.func public @make_tensor_ptr_implicit_perm_load_2d(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: i32 {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}, %arg4: i32 {tt.divisibility = 16 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %c1_i64 = arith.constant 1 : i64 + %c4_i32 = arith.constant 4 : i32 + %c64_i32 = arith.constant 64 : i32 + %0 = tt.get_program_id y : i32 + %1 = arith.muli %0, %c64_i32 : i32 + %2 = tt.get_program_id x : i32 + %3 = arith.muli %2, %c4_i32 : i32 + %4 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32> + %5 = tt.expand_dims %4 {axis = 0 : i32} : tensor<64xi32> -> tensor<1x64xi32> + %6 = tt.splat %1 : i32 -> tensor<1x64xi32> + %7 = arith.addi %6, %5 : tensor<1x64xi32> + %8 = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> + %9 = tt.expand_dims %8 {axis = 1 : i32} : tensor<4xi32> -> tensor<4x1xi32> + %10 = tt.splat %3 : i32 -> tensor<4x1xi32> + %11 = arith.addi %10, %9 : tensor<4x1xi32> + %12 = tt.splat %arg3 : i32 -> tensor<4x1xi32> + %13 = arith.cmpi slt, %11, %12 : tensor<4x1xi32> + %14 = tt.splat %arg2 : i32 -> tensor<1x64xi32> + %15 = arith.cmpi slt, %7, %14 : tensor<1x64xi32> + %16 = tt.broadcast %13 : tensor<4x1xi1> -> tensor<4x64xi1> + %17 = tt.broadcast %15 : tensor<1x64xi1> -> tensor<4x64xi1> + %18 = arith.andi %16, %17 : tensor<4x64xi1> + %19 = arith.extsi %arg3 : i32 to i64 + %20 = arith.extsi %arg2 : i32 to i64 + %21 = arith.extsi %arg4 : i32 to i64 + %22 = tt.make_tensor_ptr %arg0, [%19, %20], [%c1_i64, %21], [%3, %1] {order = array} : > + %23 = tt.load %22 : !tt.ptr> + %24 = tt.splat %arg5 : i32 -> tensor<4x1xi32> + %25 = arith.muli %11, %24 : tensor<4x1xi32> + %26 = tt.broadcast %25 : tensor<4x1xi32> -> tensor<4x64xi32> + %27 = tt.broadcast %7 : tensor<1x64xi32> -> tensor<4x64xi32> + %28 = arith.addi %26, %27 : tensor<4x64xi32> + %29 = tt.splat %arg1 : !tt.ptr -> tensor<4x64x!tt.ptr> + %30 = tt.addptr %29, %28 : tensor<4x64x!tt.ptr>, tensor<4x64xi32> + tt.store %30, %23, %18 : tensor<4x64x!tt.ptr> + tt.return + } +} + +// ----- + +// make_tensor_ptr store implicit permutation +module { + tt.func public @make_tensor_ptr_implicit_perm_store_2d(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: i32 {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}, %arg4: i32 {tt.divisibility = 16 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %c1_i64 = arith.constant 1 : i64 + %cst = arith.constant dense<0.000000e+00> : tensor<4x64xf32> + %c4_i32 = arith.constant 4 : i32 + %c64_i32 = arith.constant 64 : i32 + %0 = tt.get_program_id y : i32 + %1 = arith.muli %0, %c64_i32 : i32 + %2 = tt.get_program_id x : i32 + %3 = arith.muli %2, %c4_i32 : i32 + %4 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32> + %5 = tt.expand_dims %4 {axis = 0 : i32} : tensor<64xi32> -> tensor<1x64xi32> + %6 = tt.splat %1 : i32 -> tensor<1x64xi32> + %7 = arith.addi %6, %5 : tensor<1x64xi32> + %8 = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> + %9 = tt.expand_dims %8 {axis = 1 : i32} : tensor<4xi32> -> tensor<4x1xi32> + %10 = tt.splat %3 : i32 -> tensor<4x1xi32> + %11 = arith.addi %10, %9 : tensor<4x1xi32> + %12 = tt.splat %arg3 : i32 -> tensor<4x1xi32> + %13 = arith.cmpi slt, %11, %12 : tensor<4x1xi32> + %14 = tt.splat %arg2 : i32 -> tensor<1x64xi32> + %15 = arith.cmpi slt, %7, %14 : tensor<1x64xi32> + %16 = tt.broadcast %13 : tensor<4x1xi1> -> tensor<4x64xi1> + %17 = tt.broadcast %15 : tensor<1x64xi1> -> tensor<4x64xi1> + %18 = arith.andi %16, %17 : tensor<4x64xi1> + %19 = tt.splat %arg4 : i32 -> tensor<4x1xi32> + %20 = arith.muli %11, %19 : tensor<4x1xi32> + %21 = tt.broadcast %20 : tensor<4x1xi32> -> tensor<4x64xi32> + %22 = tt.broadcast %7 : tensor<1x64xi32> -> tensor<4x64xi32> + %23 = arith.addi %21, %22 : tensor<4x64xi32> + %24 = tt.splat %arg0 : !tt.ptr -> tensor<4x64x!tt.ptr> + %25 = tt.addptr %24, %23 : tensor<4x64x!tt.ptr>, tensor<4x64xi32> + %26 = tt.load %25, %18, %cst : tensor<4x64x!tt.ptr> + %27 = arith.extsi %arg3 : i32 to i64 + %28 = arith.extsi %arg2 : i32 to i64 + %29 = arith.extsi %arg5 : i32 to i64 + %30 = tt.make_tensor_ptr %arg1, [%27, %28], [%c1_i64, %29], [%3, %1] {order = array} : > + tt.store %30, %26 {boundaryCheck = array} : !tt.ptr> + tt.return + } +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_load_out_of_bound.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_load_out_of_bound.mlir new file mode 100644 index 0000000000..05b526ed98 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_load_out_of_bound.mlir @@ -0,0 +1,107 @@ +// RUN: triton-opt "--triton-to-linalg=global-kernel=false named-ops=True" --split-input-file %s | FileCheck %s + +// CHECK-LABEL: func.func @_triton_mrope_forward +// CHECK: %[[CST:.*]] = arith.constant 0.000000e+00 : f32 +// CHECK: %[[ALLOC_5:.*]] = memref.alloc() : memref<64xf32> +// CHECK: linalg.fill ins(%[[CST]] : f32) outs(%[[ALLOC_5]] : memref<64xf32>) +// CHECK: bufferization.to_tensor %[[ALLOC_5]] restrict writable : memref<64xf32> to tensor<64xf32> + +module attributes {hacc.target = #hacc.target<"Ascend910_9382">} { + tt.func public @_triton_mrope_forward(%q_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %k_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %cos: !tt.ptr {tt.divisibility = 16 : i32}, %sin: !tt.ptr {tt.divisibility = 16 : i32}, %num_tokens: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %cst = arith.constant dense<0.000000e+00> : tensor<16x64xf32> + %cst_0 = arith.constant dense<0.000000e+00> : tensor<64xf32> + %cst_1 = arith.constant dense<64> : tensor<16x64xi32> + %cst_2 = arith.constant dense<64> : tensor<1x64xi32> + %cst_3 = arith.constant dense<16> : tensor<16x1xi32> + %cst_4 = arith.constant dense<128> : tensor<16x1xi32> + %cst_5 = arith.constant dense<64> : tensor<64xi32> + %cst_6 = arith.constant dense<32> : tensor<64xi32> + %c64_i32 = arith.constant 64 : i32 + %c2048_i32 = arith.constant 2048 : i32 + %pid = tt.get_program_id x : i32 + %q_ptr_7 = arith.muli %pid, %c2048_i32 : i32 + %q_ptr_8 = tt.addptr %q_ptr, %q_ptr_7 : !tt.ptr, i32 + %k_ptr_9 = tt.addptr %k_ptr, %q_ptr_7 : !tt.ptr, i32 + %t_cos = arith.muli %pid, %c64_i32 : i32 + %t_cos_10 = tt.addptr %cos, %t_cos : !tt.ptr, i32 + %h_cos = arith.muli %num_tokens, %c64_i32 : i32 + %h_cos_11 = tt.addptr %t_cos_10, %h_cos : !tt.ptr, i32 + %w_cos = tt.addptr %h_cos_11, %h_cos : !tt.ptr, i32 + %t_sin = tt.addptr %sin, %t_cos : !tt.ptr, i32 + %h_sin = tt.addptr %t_sin, %h_cos : !tt.ptr, i32 + %w_sin = tt.addptr %h_sin, %h_cos : !tt.ptr, i32 + %cos_offsets = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32> + %t_mask = arith.cmpi slt, %cos_offsets, %cst_6 : tensor<64xi32> + %h_mask = arith.cmpi sge, %cos_offsets, %cst_6 : tensor<64xi32> + %h_mask_12 = arith.cmpi slt, %cos_offsets, %cst_5 : tensor<64xi32> + %h_mask_13 = arith.andi %h_mask, %h_mask_12 : tensor<64xi1> + %w_mask = arith.cmpi sge, %cos_offsets, %cst_5 : tensor<64xi32> + %w_mask_14 = arith.andi %w_mask, %h_mask_12 : tensor<64xi1> + %t_cos_row = tt.splat %t_cos_10 : !tt.ptr -> tensor<64x!tt.ptr> + %t_cos_row_15 = tt.addptr %t_cos_row, %cos_offsets : tensor<64x!tt.ptr>, tensor<64xi32> + %t_cos_row_16 = tt.load %t_cos_row_15, %t_mask, %cst_0 : tensor<64x!tt.ptr> + %h_cos_row = tt.splat %h_cos_11 : !tt.ptr -> tensor<64x!tt.ptr> + %h_cos_row_17 = tt.addptr %h_cos_row, %cos_offsets : tensor<64x!tt.ptr>, tensor<64xi32> + %h_cos_row_18 = tt.load %h_cos_row_17, %h_mask_13, %cst_0 : tensor<64x!tt.ptr> + %w_cos_row = tt.splat %w_cos : !tt.ptr -> tensor<64x!tt.ptr> + %w_cos_row_19 = tt.addptr %w_cos_row, %cos_offsets : tensor<64x!tt.ptr>, tensor<64xi32> + %w_cos_row_20 = tt.load %w_cos_row_19, %w_mask_14, %cst_0 : tensor<64x!tt.ptr> + %t_sin_row = tt.splat %t_sin : !tt.ptr -> tensor<64x!tt.ptr> + %t_sin_row_21 = tt.addptr %t_sin_row, %cos_offsets : tensor<64x!tt.ptr>, tensor<64xi32> + %t_sin_row_22 = tt.load %t_sin_row_21, %t_mask, %cst_0 : tensor<64x!tt.ptr> + %h_sin_row = tt.splat %h_sin : !tt.ptr -> tensor<64x!tt.ptr> + %h_sin_row_23 = tt.addptr %h_sin_row, %cos_offsets : tensor<64x!tt.ptr>, tensor<64xi32> + %h_sin_row_24 = tt.load %h_sin_row_23, %h_mask_13, %cst_0 : tensor<64x!tt.ptr> + %w_sin_row = tt.splat %w_sin : !tt.ptr -> tensor<64x!tt.ptr> + %w_sin_row_25 = tt.addptr %w_sin_row, %cos_offsets : tensor<64x!tt.ptr>, tensor<64xi32> + %w_sin_row_26 = tt.load %w_sin_row_25, %w_mask_14, %cst_0 : tensor<64x!tt.ptr> + %cos_row = arith.addf %t_cos_row_16, %h_cos_row_18 : tensor<64xf32> + %cos_row_27 = arith.addf %cos_row, %w_cos_row_20 : tensor<64xf32> + %sin_row = arith.addf %t_sin_row_22, %h_sin_row_24 : tensor<64xf32> + %sin_row_28 = arith.addf %sin_row, %w_sin_row_26 : tensor<64xf32> + %first_half_q_offsets = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> + %first_half_q_offsets_29 = tt.expand_dims %first_half_q_offsets {axis = 1 : i32} : tensor<16xi32> -> tensor<16x1xi32> + %first_half_q_offsets_30 = arith.muli %first_half_q_offsets_29, %cst_4 : tensor<16x1xi32> + %first_half_q_offsets_31 = tt.expand_dims %cos_offsets {axis = 0 : i32} : tensor<64xi32> -> tensor<1x64xi32> + %first_half_q_offsets_32 = tt.broadcast %first_half_q_offsets_30 : tensor<16x1xi32> -> tensor<16x64xi32> + %first_half_q_offsets_33 = tt.broadcast %first_half_q_offsets_31 : tensor<1x64xi32> -> tensor<16x64xi32> + %first_half_q_offsets_34 = arith.addi %first_half_q_offsets_32, %first_half_q_offsets_33 : tensor<16x64xi32> + %first_q_mask = arith.cmpi slt, %first_half_q_offsets_29, %cst_3 : tensor<16x1xi32> + %first_q_mask_35 = arith.cmpi slt, %first_half_q_offsets_31, %cst_2 : tensor<1x64xi32> + %first_q_mask_36 = tt.broadcast %first_q_mask : tensor<16x1xi1> -> tensor<16x64xi1> + %first_q_mask_37 = tt.broadcast %first_q_mask_35 : tensor<1x64xi1> -> tensor<16x64xi1> + %first_q_mask_38 = arith.andi %first_q_mask_36, %first_q_mask_37 : tensor<16x64xi1> + %q_tile_1 = tt.splat %q_ptr_8 : !tt.ptr -> tensor<16x64x!tt.ptr> + %q_tile_1_39 = tt.addptr %q_tile_1, %first_half_q_offsets_34 : tensor<16x64x!tt.ptr>, tensor<16x64xi32> + %q_tile_1_40 = tt.load %q_tile_1_39, %first_q_mask_38, %cst : tensor<16x64x!tt.ptr> + %k_tile_1 = tt.splat %k_ptr_9 : !tt.ptr -> tensor<16x64x!tt.ptr> + %k_tile_1_41 = tt.addptr %k_tile_1, %first_half_q_offsets_34 : tensor<16x64x!tt.ptr>, tensor<16x64xi32> + %k_tile_1_42 = tt.load %k_tile_1_41, %first_q_mask_38, %cst : tensor<16x64x!tt.ptr> + %second_half_q_offsets = arith.addi %first_half_q_offsets_34, %cst_1 : tensor<16x64xi32> + %q_tile_2 = tt.addptr %q_tile_1, %second_half_q_offsets : tensor<16x64x!tt.ptr>, tensor<16x64xi32> + %q_tile_2_43 = tt.load %q_tile_2, %first_q_mask_38, %cst : tensor<16x64x!tt.ptr> + %k_tile_2 = tt.addptr %k_tile_1, %second_half_q_offsets : tensor<16x64x!tt.ptr>, tensor<16x64xi32> + %k_tile_2_44 = tt.load %k_tile_2, %first_q_mask_38, %cst : tensor<16x64x!tt.ptr> + %new_q_tile_1 = tt.expand_dims %cos_row_27 {axis = 0 : i32} : tensor<64xf32> -> tensor<1x64xf32> + %new_q_tile_1_45 = tt.broadcast %new_q_tile_1 : tensor<1x64xf32> -> tensor<16x64xf32> + %new_q_tile_1_46 = arith.mulf %q_tile_1_40, %new_q_tile_1_45 : tensor<16x64xf32> + %new_q_tile_1_47 = tt.expand_dims %sin_row_28 {axis = 0 : i32} : tensor<64xf32> -> tensor<1x64xf32> + %new_q_tile_1_48 = tt.broadcast %new_q_tile_1_47 : tensor<1x64xf32> -> tensor<16x64xf32> + %new_q_tile_1_49 = arith.mulf %q_tile_2_43, %new_q_tile_1_48 : tensor<16x64xf32> + %new_q_tile_1_50 = arith.subf %new_q_tile_1_46, %new_q_tile_1_49 : tensor<16x64xf32> + tt.store %q_tile_1_39, %new_q_tile_1_50, %first_q_mask_38 : tensor<16x64x!tt.ptr> + %new_q_tile_2 = arith.mulf %q_tile_2_43, %new_q_tile_1_45 : tensor<16x64xf32> + %new_q_tile_2_51 = arith.mulf %q_tile_1_40, %new_q_tile_1_48 : tensor<16x64xf32> + %new_q_tile_2_52 = arith.addf %new_q_tile_2, %new_q_tile_2_51 : tensor<16x64xf32> + tt.store %q_tile_2, %new_q_tile_2_52, %first_q_mask_38 : tensor<16x64x!tt.ptr> + %new_k_tile_1 = arith.mulf %k_tile_1_42, %new_q_tile_1_45 : tensor<16x64xf32> + %new_k_tile_1_53 = arith.mulf %k_tile_2_44, %new_q_tile_1_48 : tensor<16x64xf32> + %new_k_tile_1_54 = arith.subf %new_k_tile_1, %new_k_tile_1_53 : tensor<16x64xf32> + tt.store %k_tile_1_41, %new_k_tile_1_54, %first_q_mask_38 : tensor<16x64x!tt.ptr> + %new_k_tile_2 = arith.mulf %k_tile_2_44, %new_q_tile_1_45 : tensor<16x64xf32> + %new_k_tile_2_55 = arith.mulf %k_tile_1_42, %new_q_tile_1_48 : tensor<16x64xf32> + %new_k_tile_2_56 = arith.addf %new_k_tile_2, %new_k_tile_2_55 : tensor<16x64xf32> + tt.store %k_tile_2, %new_k_tile_2_56, %first_q_mask_38 : tensor<16x64x!tt.ptr> + tt.return + } +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_nd2nz_block_ptr.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_nd2nz_block_ptr.mlir new file mode 100644 index 0000000000..8edd3bf152 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_nd2nz_block_ptr.mlir @@ -0,0 +1,20 @@ +// RUN: triton-opt --triton-to-linalg=enable-nd2nz-on-vector=true %s | FileCheck %s + +// CHECK-LABEL: func.func @test_nd2nz_block_ptr + +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + tt.func public @test_nd2nz_block_ptr( + %base_ptr: !tt.ptr {tt.divisibility = 16 : i32}, + %out_ptr: !tt.ptr {tt.divisibility = 16 : i32} + ) attributes {noinline = false} { + %c0_i32 = arith.constant 0 : i32 + %c1_i64 = arith.constant 1 : i64 + %c128_i64 = arith.constant 128 : i64 + %c32_i64 = arith.constant 32 : i64 + %load_ptr = tt.make_tensor_ptr %base_ptr, [%c128_i64, %c32_i64], [%c32_i64, %c1_i64], [%c0_i32, %c0_i32] {order = array} : > + %data = tt.load %load_ptr {boundaryCheck = array} : !tt.ptr> + %store_ptr = tt.make_tensor_ptr %out_ptr, [%c128_i64, %c32_i64], [%c32_i64, %c1_i64], [%c0_i32, %c0_i32] {order = array} : > + tt.store %store_ptr, %data {boundaryCheck = array} : !tt.ptr> + tt.return + } +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_atomic_rmw.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_atomic_rmw.mlir new file mode 100644 index 0000000000..dbe1f2636a --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_atomic_rmw.mlir @@ -0,0 +1,41 @@ +// RUN: triton-opt --triton-to-linalg --split-input-file %s | FileCheck %s + +// CHECK-LABEL: func.func @test_parse_atomic_rmw +// CHECK: %[[C8_I32:.*]] = arith.constant 8 : i32 +// CHECK: %[[FILL:.*]] = linalg.fill ins(%[[C8_I32]] : i32) +// CHECK: hivm.hir.store ins(%[[FILL]] : tensor<1xi32>) outs(%{{.*}} : memref<1xi32, strided<[1]>>) atomic = +// CHECK: %[[EXTRACTED:.*]] = tensor.extract %{{.*}}[%{{.*}}] : tensor<1xi32> +// CHECK: %[[IDX_CAST:.*]] = arith.index_cast %[[EXTRACTED]] : i32 to index +// CHECK: memref.reinterpret_cast %{{.*}} to offset: [%[[IDX_CAST]]], sizes: [8], strides: [1] + + +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + tt.func public @test_parse_atomic_rmw(%counter_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %src_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %dst_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %numel: i32) attributes {noinline = false} { + %zero_i32 = arith.constant dense<0> : tensor<8xi32> + %c8_i32 = arith.constant 8 : i32 + %pid = tt.get_program_id x : i32 + %offs = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> + %base = arith.muli %pid, %c8_i32 : i32 + %base_splat = tt.splat %base : i32 -> tensor<8xi32> + %offs_full = arith.addi %offs, %base_splat : tensor<8xi32> + %numel_splat = tt.splat %numel : i32 -> tensor<8xi32> + %mask = arith.cmpi slt, %offs_full, %numel_splat : tensor<8xi32> + %src_ptrs = tt.splat %src_ptr : !tt.ptr -> tensor<8x!tt.ptr> + %src_addptr = tt.addptr %src_ptrs, %offs_full : tensor<8x!tt.ptr>, tensor<8xi32> + %data = tt.load %src_addptr, %mask, %zero_i32 : tensor<8x!tt.ptr> + + // Tensor-level atomic: tensor<1> + %ctr_range = tt.make_range {start = 0 : i32, end = 1 : i32} : tensor<1xi32> + %ctr_ptrs = tt.splat %counter_ptr : !tt.ptr -> tensor<1x!tt.ptr> + %ctr_addptr = tt.addptr %ctr_ptrs, %ctr_range : tensor<1x!tt.ptr>, tensor<1xi32> + %incr = arith.constant dense<8> : tensor<1xi32> + %old = tt.atomic_rmw add, acq_rel, gpu, %ctr_addptr, %incr : (tensor<1x!tt.ptr>, tensor<1xi32>) -> tensor<1xi32> + %old_broadcast = tt.broadcast %old : tensor<1xi32> -> tensor<8xi32> + %write_offs = arith.addi %old_broadcast, %offs : tensor<8xi32> + + %dst_ptrs = tt.splat %dst_ptr : !tt.ptr -> tensor<8x!tt.ptr> + %dst_addptr = tt.addptr %dst_ptrs, %write_offs : tensor<8x!tt.ptr>, tensor<8xi32> + tt.store %dst_addptr, %data, %mask : tensor<8x!tt.ptr> + tt.return + } +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_div.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_div.mlir new file mode 100644 index 0000000000..398ee5b52a --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_div.mlir @@ -0,0 +1,45 @@ +// RUN: triton-opt --triton-to-linalg --split-input-file %s | FileCheck %s + +// CHECK-LABEL: func.func @test_parse_div +// CHECK: %[[SUB:.*]] = linalg.generic +// CHECK-SAME: iterator_types = ["parallel"] +// CHECK: ^bb0(%[[IN:.*]]: i32, %[[IN_5:.*]]: i32, %[[OUT:.*]]: i32): +// CHECK: %[[SUBI:.*]] = arith.subi %[[IN]], %[[IN_5]] : i32 +// CHECK: linalg.yield %[[SUBI]] : i32 + +// CHECK: %[[DIV:.*]] = linalg.generic +// CHECK-SAME: iterator_types = ["parallel"] +// CHECK: ^bb0(%[[IN:.*]]: i32, %[[IN_5:.*]]: i32, %[[OUT:.*]]: i32): +// CHECK: %[[DIVI:.*]] = arith.divsi %[[IN]], %[[IN_5]] : i32 +// CHECK: linalg.yield %[[DIVI]] : i32 + +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + tt.func public @test_parse_div(%idx_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %src_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %dst_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %numel: i32, %bias: i32, %stride: i32) attributes {noinline = false} { + %zero_f32 = arith.constant dense<0.000000e+00> : tensor<8xf32> + %zero_i32 = arith.constant dense<0> : tensor<8xi32> + %c8_i32 = arith.constant 8 : i32 + %pid = tt.get_program_id x : i32 + %offs = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> + %base = arith.muli %pid, %c8_i32 : i32 + %base_splat = tt.splat %base : i32 -> tensor<8xi32> + %offs_full = arith.addi %offs, %base_splat : tensor<8xi32> + %numel_splat = tt.splat %numel : i32 -> tensor<8xi32> + %mask = arith.cmpi slt, %offs_full, %numel_splat : tensor<8xi32> + %idx_ptrs = tt.splat %idx_ptr : !tt.ptr -> tensor<8x!tt.ptr> + %idx_addptr = tt.addptr %idx_ptrs, %offs_full : tensor<8x!tt.ptr>, tensor<8xi32> + %indices = tt.load %idx_addptr, %mask, %zero_i32 : tensor<8x!tt.ptr> + %bias_splat = tt.splat %bias : i32 -> tensor<8xi32> + %subbed = arith.subi %indices, %bias_splat : tensor<8xi32> + %bias_splat2 = tt.splat %bias : i32 -> tensor<8xi32> + %stride_splat = tt.splat %stride : i32 -> tensor<8xi32> + %scalar_div = arith.divsi %bias_splat2, %stride_splat : tensor<8xi32> + %combined = arith.addi %subbed, %scalar_div : tensor<8xi32> + %src_ptrs = tt.splat %src_ptr : !tt.ptr -> tensor<8x!tt.ptr> + %src_addptr = tt.addptr %src_ptrs, %offs_full : tensor<8x!tt.ptr>, tensor<8xi32> + %data = tt.load %src_addptr, %mask, %zero_f32 : tensor<8x!tt.ptr> + %dst_ptrs = tt.splat %dst_ptr : !tt.ptr -> tensor<8x!tt.ptr> + %dst_addptr = tt.addptr %dst_ptrs, %combined : tensor<8x!tt.ptr>, tensor<8xi32> + tt.store %dst_addptr, %data, %mask : tensor<8x!tt.ptr> + tt.return + } +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_fptosi.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_fptosi.mlir new file mode 100644 index 0000000000..8ea78b2af5 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_fptosi.mlir @@ -0,0 +1,34 @@ +// RUN: triton-opt --triton-to-linalg --split-input-file %s | FileCheck %s + +// CHECK-LABEL: func.func @test_parse_fptosi +// CHECK: %[[GENERIC:.*]] = linalg.generic +// CHECK-SAME: iterator_types = ["parallel"] +// CHECK: ^bb0(%[[IN:.*]]: f32, %[[OUT:.*]]: i32): +// CHECK: %[[FPTOSI:.*]] = arith.fptosi %[[IN]] : f32 to i32 +// CHECK: linalg.yield %[[FPTOSI]] : i32 + +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + tt.func public @test_parse_fptosi(%fidx_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %src_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %dst_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %numel: i32) attributes {noinline = false} { + %zero_f32 = arith.constant dense<0.000000e+00> : tensor<8xf32> + %c8_i32 = arith.constant 8 : i32 + %pid = tt.get_program_id x : i32 + %offs = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> + %base = arith.muli %pid, %c8_i32 : i32 + %base_splat = tt.splat %base : i32 -> tensor<8xi32> + %offs_full = arith.addi %offs, %base_splat : tensor<8xi32> + %numel_splat = tt.splat %numel : i32 -> tensor<8xi32> + %mask = arith.cmpi slt, %offs_full, %numel_splat : tensor<8xi32> + %offs_i64 = arith.extsi %offs_full : tensor<8xi32> to tensor<8xi64> + %fidx_ptrs = tt.splat %fidx_ptr : !tt.ptr -> tensor<8x!tt.ptr> + %fidx_addptr = tt.addptr %fidx_ptrs, %offs_full : tensor<8x!tt.ptr>, tensor<8xi32> + %float_indices = tt.load %fidx_addptr, %mask, %zero_f32 : tensor<8x!tt.ptr> + %int_indices = arith.fptosi %float_indices : tensor<8xf32> to tensor<8xi32> + %src_ptrs = tt.splat %src_ptr : !tt.ptr -> tensor<8x!tt.ptr> + %src_addptr = tt.addptr %src_ptrs, %offs_full : tensor<8x!tt.ptr>, tensor<8xi32> + %data = tt.load %src_addptr, %mask, %zero_f32 : tensor<8x!tt.ptr> + %dst_ptrs = tt.splat %dst_ptr : !tt.ptr -> tensor<8x!tt.ptr> + %dst_addptr = tt.addptr %dst_ptrs, %int_indices : tensor<8x!tt.ptr>, tensor<8xi32> + tt.store %dst_addptr, %data, %mask : tensor<8x!tt.ptr> + tt.return + } +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_select_dense_const.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_select_dense_const.mlir new file mode 100644 index 0000000000..ecb5ba4efb --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_select_dense_const.mlir @@ -0,0 +1,36 @@ +// RUN: triton-opt --triton-to-linalg --split-input-file %s | FileCheck %s + +// CHECK-LABEL: func.func @test_parse_select_dense_const +// CHECK: %[[C32_I32:.*]] = arith.constant 32 : i32 +// CHECK: %[[C0_I32:.*]] = arith.constant 0 : i32 +// CHECK: %[[CMP:.*]] = arith.cmpi sgt, %{{.*}}, %[[C0_I32]] : i32 +// CHECK: %[[SELECT:.*]] = arith.select %[[CMP]], %[[C0_I32]], %[[C32_I32]] : i32 +// CHECK: %[[IDX_CAST:.*]] = arith.index_cast %[[SELECT]] : i32 to index +// CHECK: %[[ADD:.*]] = arith.addi %[[IDX_CAST]], %{{.*}} : index + +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + tt.func public @test_parse_select_dense_const(%src_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %dst_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %numel: i32, %flag: i32) attributes {noinline = false} { + %zero_f32 = arith.constant dense<0.000000e+00> : tensor<8xf32> + %c8_i32 = arith.constant 8 : i32 + %c0_i32 = arith.constant 0 : i32 + %pid = tt.get_program_id x : i32 + %offs = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> + %base = arith.muli %pid, %c8_i32 : i32 + %base_splat = tt.splat %base : i32 -> tensor<8xi32> + %offs_full = arith.addi %offs, %base_splat : tensor<8xi32> + %numel_splat = tt.splat %numel : i32 -> tensor<8xi32> + %mask = arith.cmpi slt, %offs_full, %numel_splat : tensor<8xi32> + %cond = arith.cmpi sgt, %flag, %c0_i32 : i32 + %const_a = arith.constant dense<0> : tensor<8xi32> + %const_b = arith.constant dense<32> : tensor<8xi32> + %chosen = arith.select %cond, %const_a, %const_b : tensor<8xi32> + %write_offs = arith.addi %chosen, %offs_full : tensor<8xi32> + %src_ptrs = tt.splat %src_ptr : !tt.ptr -> tensor<8x!tt.ptr> + %src_addptr = tt.addptr %src_ptrs, %offs_full : tensor<8x!tt.ptr>, tensor<8xi32> + %data = tt.load %src_addptr, %mask, %zero_f32 : tensor<8x!tt.ptr> + %dst_ptrs = tt.splat %dst_ptr : !tt.ptr -> tensor<8x!tt.ptr> + %dst_addptr = tt.addptr %dst_ptrs, %write_offs : tensor<8x!tt.ptr>, tensor<8xi32> + tt.store %dst_addptr, %data, %mask : tensor<8x!tt.ptr> + tt.return + } +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_sub.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_sub.mlir new file mode 100644 index 0000000000..b9d1086358 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_parse_sub.mlir @@ -0,0 +1,34 @@ +// RUN: triton-opt --triton-to-linalg --split-input-file %s | FileCheck %s + +// CHECK-LABEL: func.func @sub_offset_kernel + +// CHECK: memref.subview {{.*}} : memref<8xi32, strided<[1], offset: ?>> to memref> +// CHECK: memref.subview {{.*}} : memref<8xi32> to memref> +// CHECK: memref.copy + +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + tt.func public @sub_offset_kernel(%idx_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %src_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %dst_ptr: !tt.ptr {tt.divisibility = 16 : i32}, %numel: i32, %correction: i32) attributes {noinline = false} { + %val = arith.constant dense<0.000000e+00> : tensor<8xf32> + %base_idx = arith.constant dense<0> : tensor<8xi32> + %c8_i32 = arith.constant 8 : i32 + %pid = tt.get_program_id x : i32 + %offs = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> + %offs_0 = arith.muli %pid, %c8_i32 : i32 + %offs_1 = tt.splat %offs_0 : i32 -> tensor<8xi32> + %offs_2 = arith.addi %offs, %offs_1 : tensor<8xi32> + %mask = tt.splat %numel : i32 -> tensor<8xi32> + %mask_3 = arith.cmpi slt, %offs_2, %mask : tensor<8xi32> + %base_idx_4 = tt.splat %idx_ptr : !tt.ptr -> tensor<8x!tt.ptr> + %base_idx_5 = tt.addptr %base_idx_4, %offs_2 : tensor<8x!tt.ptr>, tensor<8xi32> + %base_idx_6 = tt.load %base_idx_5, %mask_3, %base_idx : tensor<8x!tt.ptr> + %adjusted = tt.splat %correction : i32 -> tensor<8xi32> + %adjusted_7 = arith.subi %base_idx_6, %adjusted : tensor<8xi32> + %val_8 = tt.splat %src_ptr : !tt.ptr -> tensor<8x!tt.ptr> + %val_9 = tt.addptr %val_8, %offs_2 : tensor<8x!tt.ptr>, tensor<8xi32> + %val_10 = tt.load %val_9, %mask_3, %val : tensor<8x!tt.ptr> + %0 = tt.splat %dst_ptr : !tt.ptr -> tensor<8x!tt.ptr> + %1 = tt.addptr %0, %adjusted_7 : tensor<8x!tt.ptr>, tensor<8xi32> + tt.store %1, %val_10, %mask_3 : tensor<8x!tt.ptr> + tt.return + } +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_while_loop_index_tensor.mlir b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_while_loop_index_tensor.mlir new file mode 100644 index 0000000000..097c0cf466 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToLinalg/test_while_loop_index_tensor.mlir @@ -0,0 +1,48 @@ +// RUN: triton-opt --triton-to-linalg %s | FileCheck %s + +// CHECK-LABEL: func.func @test_while_loop_index_tensor +// CHECK: %[[C8:.*]] = arith.constant 8 : index +// CHECK: %[[C0_I32:.*]] = arith.constant 0 : i32 +// CHECK: %[[C1_I32:.*]] = arith.constant 1 : i32 +// CHECK: %[[C4_I32:.*]] = arith.constant 4 : i32 +// CHECK: %[[C0:.*]] = arith.constant 0 : index +// CHECK: scf.while (%[[ARG9:.*]] = %[[C0_I32]], %[[ARG10:.*]] = %[[C0]]) +// CHECK: arith.cmpi slt, %[[ARG9]], %[[C4_I32]] +// CHECK: scf.condition +// CHECK: ^bb0(%[[ARG9:.*]]: i32, %[[ARG10:.*]]: index) +// CHECK: arith.addi %[[ARG9]], %[[C1_I32]] +// CHECK: arith.addi %[[ARG10]], %[[C8]] + +module attributes {hacc.target = #hacc.target<"Ascend910B2">} { + tt.func public @test_while_loop_index_tensor( + %src_ptr: !tt.ptr {tt.divisibility = 16 : i32}, + %dst_ptr: !tt.ptr {tt.divisibility = 16 : i32}, + %numel: i32 + ) attributes {noinline = false} { + %zero_f32 = arith.constant dense<0.000000e+00> : tensor<8xf32> + %c0_i32 = arith.constant 0 : i32 + %c1_i32 = arith.constant 1 : i32 + %c8_i32 = arith.constant 8 : i32 + %c4_i32 = arith.constant 4 : i32 + %init_offs = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> + %step = arith.constant dense<8> : tensor<8xi32> + %result:2 = scf.while (%offs = %init_offs, %iter = %c0_i32) : (tensor<8xi32>, i32) -> (tensor<8xi32>, i32) { + %cond = arith.cmpi slt, %iter, %c4_i32 : i32 + scf.condition(%cond) %offs, %iter : tensor<8xi32>, i32 + } do { + ^bb0(%offs_arg: tensor<8xi32>, %iter_arg: i32): + %ptrs = tt.splat %src_ptr : !tt.ptr -> tensor<8x!tt.ptr> + %addptr = tt.addptr %ptrs, %offs_arg : tensor<8x!tt.ptr>, tensor<8xi32> + %numel_splat = tt.splat %numel : i32 -> tensor<8xi32> + %mask = arith.cmpi slt, %offs_arg, %numel_splat : tensor<8xi32> + %data = tt.load %addptr, %mask, %zero_f32 : tensor<8x!tt.ptr> + %dst_ptrs = tt.splat %dst_ptr : !tt.ptr -> tensor<8x!tt.ptr> + %dst_addptr = tt.addptr %dst_ptrs, %offs_arg : tensor<8x!tt.ptr>, tensor<8xi32> + tt.store %dst_addptr, %data, %mask : tensor<8x!tt.ptr> + %new_offs = arith.addi %offs_arg, %step : tensor<8xi32> + %new_iter = arith.addi %iter_arg, %c1_i32 : i32 + scf.yield %new_offs, %new_iter : tensor<8xi32>, i32 + } + tt.return + } +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToStructured/CmpConverter.mlir b/third_party/ascend/unittest/Conversion/General/TritonToStructured/CmpConverter.mlir new file mode 100644 index 0000000000..909de6174d --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToStructured/CmpConverter.mlir @@ -0,0 +1,24 @@ +// RUN: triton-opt '--triton-to-structured=enable-mask-fallback-conversion=false optimize-dynamic-offset=true' --split-input-file %s | FileCheck %s + +module { + tt.func public @test_cmp(%arg0: tensor<128xi32>) -> tensor<128xi1> { + %cst_12 = arith.constant dense<0> : tensor<128xi32> + %cst_13 = arith.constant dense<1> : tensor<128xi32> + %cst_14 = arith.constant dense<100> : tensor<128xi32> + %39 = arith.cmpi slt, %arg0, %cst_14 : tensor<128xi32> + %40 = arith.select %39, %cst_13, %cst_12 : tensor<128xi1>, tensor<128xi32> + %41 = arith.cmpi ne, %40, %cst_12 : tensor<128xi32> + tt.return %41 : tensor<128xi1> + } +} + +// CHECK-LABEL: tt.func public @test_cmp( +// CHECK-SAME: %[[VAL_0:.*]]: tensor<128xi32>) -> tensor<128xi1> { +// CHECK: %[[VAL_1:.*]] = arith.constant dense<0> : tensor<128xi32> +// CHECK: %[[VAL_2:.*]] = arith.constant dense<1> : tensor<128xi32> +// CHECK: %[[VAL_3:.*]] = arith.constant dense<100> : tensor<128xi32> +// CHECK: %[[VAL_4:.*]] = arith.cmpi slt, %[[VAL_0]], %[[VAL_3]] : tensor<128xi32> +// CHECK: %[[VAL_5:.*]] = arith.select %[[VAL_4]], %[[VAL_2]], %[[VAL_1]] : tensor<128xi1>, tensor<128xi32> +// CHECK: %[[VAL_6:.*]] = arith.cmpi ne, %[[VAL_5]], %[[VAL_1]] : tensor<128xi32> +// CHECK: tt.return %[[VAL_6]] : tensor<128xi1> +// CHECK: } diff --git a/third_party/ascend/unittest/Conversion/General/TritonToStructured/PromotePointerIterArgsPattern.mlir b/third_party/ascend/unittest/Conversion/General/TritonToStructured/PromotePointerIterArgsPattern.mlir new file mode 100644 index 0000000000..7163b89f5f --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToStructured/PromotePointerIterArgsPattern.mlir @@ -0,0 +1,73 @@ +// RUN: triton-opt '--triton-to-structured=enable-mask-fallback-conversion=false optimize-dynamic-offset=true' --split-input-file %s | FileCheck %s + +module { + tt.func public @test_promote_pointer_iter(%base_ptr: !tt.ptr {tt.divisibility = 16 : i32}) -> !tt.ptr { + %c1_i32 = arith.constant 1 : i32 + %c0_f32 = arith.constant 0.000000e+00 : f32 + %true_mask = arith.constant 1 : i1 + %c0_index = arith.constant 0 : index + %c10_index = arith.constant 10 : index + %c1_index = arith.constant 1 : index + %final_ptr = scf.for %iv = %c0_index to %c10_index step %c1_index iter_args(%ptr = %base_ptr) -> (!tt.ptr) { + %data = tt.load %ptr, %true_mask, %c0_f32 : !tt.ptr + %new_ptr = tt.addptr %ptr, %c1_i32 : !tt.ptr, i32 + scf.yield %new_ptr : !tt.ptr + } + tt.return %final_ptr : !tt.ptr + } +} + +// CHECK-LABEL: tt.func public @test_promote_pointer_iter( +// CHECK-SAME: %[[VAL_0:.*]]: !tt.ptr {tt.divisibility = 16 : i32}) -> !tt.ptr { +// CHECK: %[[VAL_1:.*]] = arith.constant 1 : i32 +// CHECK: %[[VAL_2:.*]] = arith.constant 0 : index +// CHECK: %[[VAL_3:.*]] = arith.constant 10 : index +// CHECK: %[[VAL_4:.*]] = arith.constant 1 : index +// CHECK: %[[VAL_5:.*]] = scf.for %[[VAL_6:.*]] = %[[VAL_2]] to %[[VAL_3]] step %[[VAL_4]] iter_args(%[[VAL_7:.*]] = %[[VAL_0]]) -> (!tt.ptr) { +// CHECK: %[[VAL_8:.*]] = tt.addptr %[[VAL_7]], %[[VAL_1]] : !tt.ptr, i32 +// CHECK: scf.yield %[[VAL_8]] : !tt.ptr +// CHECK: } +// CHECK: tt.return %[[VAL_5]] : !tt.ptr +// CHECK: } + + +// ----- + + +module { + tt.func public @test_promote_pointer_iter_advance(%base_ptr: !tt.ptr) -> !tt.ptr>{ + %c0_i32 = arith.constant 0 : i32 + %c1_i64 = arith.constant 1 : i64 + %c32_i64 = arith.constant 32 : i64 + %c1_i32 = arith.constant 1 : i32 // nonZeroConstant 需要 1 + %c0_i32_2 = arith.constant 0 : i32 + %c0_index = arith.constant 0 : index + %c10_index = arith.constant 10 : index + %c1_index = arith.constant 1 : index + %cst = arith.constant dense<0.000000e+00> : tensor<32xf16> + %ptr0 = tt.make_tensor_ptr %base_ptr, [%c32_i64], [%c1_i64], [%c0_i32] {order = array} : !tt.ptr> + %final_ptr = scf.for %iv = %c0_index to %c10_index step %c1_index iter_args(%ptr = %ptr0) -> !tt.ptr> { + %data = tt.load %ptr : !tt.ptr> + %new_ptr = tt.advance %ptr, [%c1_i32, %c0_i32_2] : !tt.ptr> + scf.yield %new_ptr : !tt.ptr> + } + tt.return %final_ptr : !tt.ptr> + } +} + +// CHECK-LABEL: tt.func public @test_promote_pointer_iter_advance( +// CHECK-SAME: %[[VAL_0:.*]]: !tt.ptr) -> !tt.ptr> { +// CHECK: %[[VAL_1:.*]] = arith.constant 0 : i32 +// CHECK: %[[VAL_2:.*]] = arith.constant 1 : i64 +// CHECK: %[[VAL_3:.*]] = arith.constant 32 : i64 +// CHECK: %[[VAL_4:.*]] = arith.constant 1 : i32 +// CHECK: %[[VAL_5:.*]] = arith.constant 0 : index +// CHECK: %[[VAL_6:.*]] = arith.constant 10 : index +// CHECK: %[[VAL_7:.*]] = arith.constant 1 : index +// CHECK: %[[VAL_8:.*]] = tt.make_tensor_ptr %[[VAL_0]], [%[[VAL_3]]], [%[[VAL_2]]], [%[[VAL_1]]] {order = array} : > +// CHECK: %[[VAL_9:.*]] = scf.for %[[VAL_10:.*]] = %[[VAL_5]] to %[[VAL_6]] step %[[VAL_7]] iter_args(%[[VAL_11:.*]] = %[[VAL_8]]) -> (!tt.ptr>) { +// CHECK: %[[VAL_12:.*]] = tt.advance %[[VAL_11]], [%[[VAL_4]], %[[VAL_1]]] : > +// CHECK: scf.yield %[[VAL_12]] : !tt.ptr> +// CHECK: } +// CHECK: tt.return %[[VAL_9]] : !tt.ptr> +// CHECK: } diff --git a/third_party/ascend/unittest/Conversion/General/TritonToStructured/SplatCmpConverter.mlir b/third_party/ascend/unittest/Conversion/General/TritonToStructured/SplatCmpConverter.mlir new file mode 100644 index 0000000000..f046891ec1 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToStructured/SplatCmpConverter.mlir @@ -0,0 +1,16 @@ +// RUN: triton-opt '--triton-to-structured=enable-mask-fallback-conversion=false optimize-dynamic-offset=true' --split-input-file %s | FileCheck %s +module { + tt.func public @test_splat_cmp(%arg0: i32, %arg1: i32) -> tensor<128xi1> { + %0 = tt.splat %arg0 : i32 -> tensor<128xi32> + %1 = tt.splat %arg1 : i32 -> tensor<128xi32> + %2 = arith.cmpi slt, %0, %1 : tensor<128xi32> + tt.return %2 : tensor<128xi1> + } +} + +// CHECK-LABEL: tt.func public @test_splat_cmp( +// CHECK-SAME: %[[VAL_0:.*]]: i32, %[[VAL_1:.*]]: i32) -> tensor<128xi1> { +// CHECK: %[[VAL_2:.*]] = arith.cmpi slt, %[[VAL_0]], %[[VAL_1]] : i32 +// CHECK: %[[VAL_3:.*]] = tt.splat %[[VAL_2]] : i1 -> tensor<128xi1> +// CHECK: tt.return %[[VAL_3]] : tensor<128xi1> +// CHECK: } diff --git a/third_party/ascend/unittest/Conversion/General/TritonToStructured/parseCmp.mlir b/third_party/ascend/unittest/Conversion/General/TritonToStructured/parseCmp.mlir new file mode 100644 index 0000000000..abdae3fa16 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToStructured/parseCmp.mlir @@ -0,0 +1,117 @@ +// RUN: triton-opt '--triton-to-structured=enable-mask-fallback-conversion=false optimize-dynamic-offset=true' --split-input-file %s | FileCheck %s + +tt.func public @test_cmp_ult(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %cst = arith.constant dense<0.000000e+00> : tensor<1024xf32> + %cst_0 = arith.constant dense<512> : tensor<1024xi32> + %c1024_i32 = arith.constant 1024 : i32 + %0 = tt.get_program_id x : i32 + %1 = arith.muli %0, %c1024_i32 {tt.divisibility = dense<512> : tensor<1xi32>} : i32 + %2 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> + %3 = tt.splat %1 : i32 -> tensor<1024xi32> + %4 = arith.addi %3, %2 : tensor<1024xi32> + %5 = arith.divsi %4, %cst_0 : tensor<1024xi32> + %6 = arith.cmpi ult, %5, %cst_0 : tensor<1024xi32> + %7 = arith.muli %5, %cst_0 : tensor<1024xi32> + %8 = tt.splat %arg0 : !tt.ptr -> tensor<1024x!tt.ptr> + %9 = tt.addptr %8, %7 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %10 = tt.load %9, %6, %cst : tensor<1024x!tt.ptr> + %11 = arith.muli %0, %c1024_i32 : i32 + %12 = tt.splat %11 : i32 -> tensor<1024xi32> + %13 = arith.addi %12, %2 : tensor<1024xi32> + %14 = tt.splat %arg1 : !tt.ptr -> tensor<1024x!tt.ptr> + %15 = tt.addptr %14, %13 : tensor<1024x!tt.ptr>, tensor<1024xi32> + tt.store %15, %10 : tensor<1024x!tt.ptr> + tt.return +} + +// CHECK-LABEL: tt.func public @test_cmp_ult( +// CHECK-SAME: %[[VAL_0:.*]]: !tt.ptr {tt.divisibility = 16 : i32}, %[[VAL_1:.*]]: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { +// CHECK: %[[VAL_2:.*]] = arith.constant dense<1024> : tensor<1xi64> +// CHECK: %[[VAL_3:.*]] = arith.constant dense<0.000000e+00> : tensor<2xf32> +// CHECK: %[[VAL_4:.*]] = arith.constant dense<512> : tensor<2xi32> +// CHECK: %[[VAL_5:.*]] = arith.constant 512 : index +// CHECK: %[[VAL_6:.*]] = arith.constant 1024 : i32 +// CHECK: %[[VAL_7:.*]] = tt.get_program_id x : i32 +// CHECK: %[[VAL_8:.*]] = arith.muli %[[VAL_7]], %[[VAL_6]] {tt.divisibility = dense<512> : tensor<1xi32>} : i32 +// CHECK: %[[VAL_9:.*]] = arith.index_cast %[[VAL_8]] : i32 to index +// CHECK: %[[VAL_10:.*]] = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> +// CHECK: %[[VAL_11:.*]] = arith.divsi %[[VAL_9]], %[[VAL_5]] : index +// CHECK: %[[VAL_12:.*]] = arith.muli %[[VAL_11]], %[[VAL_5]] : index +// CHECK: %[[VAL_13:.*]] = tt.make_range {end = 2 : i32, start = 0 : i32} : tensor<2xi32> +// CHECK: %[[VAL_14:.*]] = arith.muli %[[VAL_13]], %[[VAL_4]] : tensor<2xi32> +// CHECK: %[[VAL_15:.*]] = arith.index_cast %[[VAL_12]] : index to i32 +// CHECK: %[[VAL_16:.*]] = tt.splat %[[VAL_15]] : i32 -> tensor<2xi32> +// CHECK: %[[VAL_17:.*]] = arith.addi %[[VAL_14]], %[[VAL_16]] : tensor<2xi32> +// CHECK: %[[VAL_18:.*]] = tt.splat %[[VAL_0]] : !tt.ptr -> tensor<2x!tt.ptr> +// CHECK: %[[VAL_19:.*]] = tt.addptr %[[VAL_18]], %[[VAL_17]] : tensor<2x!tt.ptr>, tensor<2xi32> +// CHECK: %[[VAL_20:.*]] = arith.index_cast %[[VAL_11]] : index to i32 +// CHECK: %[[VAL_21:.*]] = tt.splat %[[VAL_20]] : i32 -> tensor<2xi32> +// CHECK: %[[VAL_22:.*]] = arith.addi %[[VAL_13]], %[[VAL_21]] : tensor<2xi32> +// CHECK: %[[VAL_23:.*]] = arith.cmpi slt, %[[VAL_22]], %[[VAL_4]] : tensor<2xi32> +// CHECK: %[[VAL_24:.*]] = tt.load %[[VAL_19]], %[[VAL_23]], %[[VAL_3]] : tensor<2x!tt.ptr> +// CHECK: %[[VAL_25:.*]] = tensor.empty() : tensor<2x512xf32> +// CHECK: %[[VAL_26:.*]] = linalg.broadcast ins(%[[VAL_24]] : tensor<2xf32>) outs(%[[VAL_25]] : tensor<2x512xf32>) dimensions = [1] +// CHECK: %[[VAL_27:.*]] = tensor.reshape %[[VAL_26]](%[[VAL_2]]) : (tensor<2x512xf32>, tensor<1xi64>) -> tensor<1024xf32> +// CHECK: %[[VAL_28:.*]] = arith.muli %[[VAL_7]], %[[VAL_6]] : i32 +// CHECK: %[[VAL_29:.*]] = tt.splat %[[VAL_28]] : i32 -> tensor<1024xi32> +// CHECK: %[[VAL_30:.*]] = arith.addi %[[VAL_29]], %[[VAL_10]] : tensor<1024xi32> +// CHECK: %[[VAL_31:.*]] = tt.splat %[[VAL_1]] : !tt.ptr -> tensor<1024x!tt.ptr> +// CHECK: %[[VAL_32:.*]] = tt.addptr %[[VAL_31]], %[[VAL_30]] : tensor<1024x!tt.ptr>, tensor<1024xi32> +// CHECK: tt.store %[[VAL_32]], %[[VAL_27]] : tensor<1024x!tt.ptr> +// CHECK: tt.return +// CHECK: } + + +// ----- + + +tt.func public @test_cmp_uge(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %cst = arith.constant dense<0.000000e+00> : tensor<1024xf32> + %cst_0 = arith.constant dense<511> : tensor<1024xi32> + %cst_1 = arith.constant dense<512> : tensor<1024xi32> + %c1024_i32 = arith.constant 1024 : i32 + %0 = tt.get_program_id x : i32 + %1 = arith.muli %0, %c1024_i32 {tt.divisibility = dense<512> : tensor<1xi32>} : i32 + %2 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> + %3 = tt.splat %1 : i32 -> tensor<1024xi32> + %4 = arith.addi %3, %2 : tensor<1024xi32> + %5 = arith.divsi %4, %cst_1 : tensor<1024xi32> + %6 = arith.cmpi uge, %cst_0, %5 : tensor<1024xi32> + %7 = arith.muli %5, %cst_1 : tensor<1024xi32> + %8 = tt.splat %arg0 : !tt.ptr -> tensor<1024x!tt.ptr> + %9 = tt.addptr %8, %7 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %10 = tt.load %9, %6, %cst : tensor<1024x!tt.ptr> + %11 = arith.muli %0, %c1024_i32 : i32 + %12 = tt.splat %11 : i32 -> tensor<1024xi32> + %13 = arith.addi %12, %2 : tensor<1024xi32> + %14 = tt.splat %arg1 : !tt.ptr -> tensor<1024x!tt.ptr> + %15 = tt.addptr %14, %13 : tensor<1024x!tt.ptr>, tensor<1024xi32> + tt.store %15, %10 : tensor<1024x!tt.ptr> + tt.return +} + +// CHECK-LABEL: tt.func public @test_cmp_uge( +// CHECK-SAME: %[[VAL_0:.*]]: !tt.ptr {tt.divisibility = 16 : i32}, %[[VAL_1:.*]]: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { +// CHECK: %[[VAL_2:.*]] = arith.constant dense<0.000000e+00> : tensor<1024xf32> +// CHECK: %[[VAL_3:.*]] = arith.constant dense<511> : tensor<1024xi32> +// CHECK: %[[VAL_4:.*]] = arith.constant dense<512> : tensor<1024xi32> +// CHECK: %[[VAL_5:.*]] = arith.constant 1024 : i32 +// CHECK: %[[VAL_6:.*]] = tt.get_program_id x : i32 +// CHECK: %[[VAL_7:.*]] = arith.muli %[[VAL_6]], %[[VAL_5]] {tt.divisibility = dense<512> : tensor<1xi32>} : i32 +// CHECK: %[[VAL_8:.*]] = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> +// CHECK: %[[VAL_9:.*]] = tt.splat %[[VAL_7]] : i32 -> tensor<1024xi32> +// CHECK: %[[VAL_10:.*]] = arith.addi %[[VAL_9]], %[[VAL_8]] : tensor<1024xi32> +// CHECK: %[[VAL_11:.*]] = arith.divsi %[[VAL_10]], %[[VAL_4]] : tensor<1024xi32> +// CHECK: %[[VAL_12:.*]] = arith.cmpi ule, %[[VAL_11]], %[[VAL_3]] : tensor<1024xi32> +// CHECK: %[[VAL_13:.*]] = arith.muli %[[VAL_11]], %[[VAL_4]] : tensor<1024xi32> +// CHECK: %[[VAL_14:.*]] = tt.splat %[[VAL_0]] : !tt.ptr -> tensor<1024x!tt.ptr> +// CHECK: %[[VAL_15:.*]] = tt.addptr %[[VAL_14]], %[[VAL_13]] : tensor<1024x!tt.ptr>, tensor<1024xi32> +// CHECK: %[[VAL_16:.*]] = tt.load %[[VAL_15]], %[[VAL_12]], %[[VAL_2]] : tensor<1024x!tt.ptr> +// CHECK: %[[VAL_17:.*]] = arith.muli %[[VAL_6]], %[[VAL_5]] : i32 +// CHECK: %[[VAL_18:.*]] = tt.splat %[[VAL_17]] : i32 -> tensor<1024xi32> +// CHECK: %[[VAL_19:.*]] = arith.addi %[[VAL_18]], %[[VAL_8]] : tensor<1024xi32> +// CHECK: %[[VAL_20:.*]] = tt.splat %[[VAL_1]] : !tt.ptr -> tensor<1024x!tt.ptr> +// CHECK: %[[VAL_21:.*]] = tt.addptr %[[VAL_20]], %[[VAL_19]] : tensor<1024x!tt.ptr>, tensor<1024xi32> +// CHECK: tt.store %[[VAL_21]], %[[VAL_16]] : tensor<1024x!tt.ptr> +// CHECK: tt.return +// CHECK: } diff --git a/third_party/ascend/unittest/Conversion/General/TritonToStructured/parseConstant.mlir b/third_party/ascend/unittest/Conversion/General/TritonToStructured/parseConstant.mlir new file mode 100644 index 0000000000..4e1053793d --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToStructured/parseConstant.mlir @@ -0,0 +1,28 @@ +// RUN: triton-opt '--triton-to-structured=enable-mask-fallback-conversion=false optimize-dynamic-offset=true' --split-input-file %s | FileCheck %s + +tt.func public @test_non_splat_mask(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}) { + %non_splat_mask = arith.constant dense<[false, true]> : tensor<2xi1> + %c0_f32 = arith.constant dense<0.000000e+00> : tensor<2xf32> + %0 = tt.make_range {end = 2 : i32, start = 0 : i32} : tensor<2xi32> + %1 = tt.splat %arg0 : !tt.ptr -> tensor<2x!tt.ptr> + %ptr_load = tt.addptr %1, %0 : tensor<2x!tt.ptr>, tensor<2xi32> + %2 = tt.load %ptr_load, %non_splat_mask, %c0_f32 : tensor<2x!tt.ptr> + %3 = tt.splat %arg1 : !tt.ptr -> tensor<2x!tt.ptr> + %ptr_store = tt.addptr %3, %0 : tensor<2x!tt.ptr>, tensor<2xi32> + tt.store %ptr_store, %2 : tensor<2x!tt.ptr> + tt.return +} + +// CHECK-LABEL: tt.func public @test_non_splat_mask( +// CHECK-SAME: %[[VAL_0:.*]]: !tt.ptr {tt.divisibility = 16 : i32}, %[[VAL_1:.*]]: !tt.ptr {tt.divisibility = 16 : i32}) { +// CHECK: %[[VAL_2:.*]] = arith.constant dense<[false, true]> : tensor<2xi1> +// CHECK: %[[VAL_3:.*]] = arith.constant dense<0.000000e+00> : tensor<2xf32> +// CHECK: %[[VAL_4:.*]] = tt.make_range {end = 2 : i32, start = 0 : i32} : tensor<2xi32> +// CHECK: %[[VAL_5:.*]] = tt.splat %[[VAL_0]] : !tt.ptr -> tensor<2x!tt.ptr> +// CHECK: %[[VAL_6:.*]] = tt.addptr %[[VAL_5]], %[[VAL_4]] : tensor<2x!tt.ptr>, tensor<2xi32> +// CHECK: %[[VAL_7:.*]] = tt.load %[[VAL_6]], %[[VAL_2]], %[[VAL_3]] : tensor<2x!tt.ptr> +// CHECK: %[[VAL_8:.*]] = tt.splat %[[VAL_1]] : !tt.ptr -> tensor<2x!tt.ptr> +// CHECK: %[[VAL_9:.*]] = tt.addptr %[[VAL_8]], %[[VAL_4]] : tensor<2x!tt.ptr>, tensor<2xi32> +// CHECK: tt.store %[[VAL_9]], %[[VAL_7]] : tensor<2x!tt.ptr> +// CHECK: tt.return +// CHECK: } diff --git a/third_party/ascend/unittest/Conversion/General/TritonToStructured/parseMakeRange.mlir b/third_party/ascend/unittest/Conversion/General/TritonToStructured/parseMakeRange.mlir new file mode 100644 index 0000000000..efdaad3204 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToStructured/parseMakeRange.mlir @@ -0,0 +1,24 @@ +// RUN: triton-opt '--triton-to-structured=enable-mask-fallback-conversion=false optimize-dynamic-offset=true' --split-input-file %s | FileCheck %s + +module { + tt.func public @test_stride_not_one(%arg0: !tt.ptr {tt.divisibility = 16 : i32}) { + %c0_f32 = arith.constant dense<0.000000e+00> : tensor<4xf32> + %fake_range_mask = arith.constant dense<[false, false, false, false]> : tensor<4xi1> + %0 = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> + %1 = tt.splat %arg0 : !tt.ptr -> tensor<4x!tt.ptr> + %ptr = tt.addptr %1, %0 : tensor<4x!tt.ptr>, tensor<4xi32> + %2 = tt.load %ptr, %fake_range_mask, %c0_f32 : tensor<4x!tt.ptr> + tt.store %ptr, %2 : tensor<4x!tt.ptr> + tt.return + } +} + +// CHECK-LABEL: tt.func public @test_stride_not_one( +// CHECK-SAME: %[[VAL_0:.*]]: !tt.ptr {tt.divisibility = 16 : i32}) { +// CHECK: %[[VAL_1:.*]] = arith.constant dense<0.000000e+00> : tensor<4xf32> +// CHECK: %[[VAL_2:.*]] = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> +// CHECK: %[[VAL_3:.*]] = tt.splat %[[VAL_0]] : !tt.ptr -> tensor<4x!tt.ptr> +// CHECK: %[[VAL_4:.*]] = tt.addptr %[[VAL_3]], %[[VAL_2]] : tensor<4x!tt.ptr>, tensor<4xi32> +// CHECK: tt.store %[[VAL_4]], %[[VAL_1]] : tensor<4x!tt.ptr> +// CHECK: tt.return +// CHECK: } diff --git a/third_party/ascend/unittest/Conversion/General/TritonToStructured/parseRem.mlir b/third_party/ascend/unittest/Conversion/General/TritonToStructured/parseRem.mlir new file mode 100644 index 0000000000..b48d7fc288 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToStructured/parseRem.mlir @@ -0,0 +1,81 @@ +// RUN: triton-opt '--triton-to-structured=enable-mask-fallback-conversion=false optimize-dynamic-offset=true' --split-input-file %s | FileCheck %s + +tt.func public @kernel_with_rem_safe(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %cst = arith.constant dense<0.000000e+00> : tensor<256xf32> + %cst_0 = arith.constant dense<1024> : tensor<256xi32> + %c256_i32 = arith.constant 256 : i32 + %cst_1 = arith.constant dense<64> : tensor<256xi32> + %cst_2 = arith.constant dense<128> : tensor<256xi32> + %0 = tt.get_program_id x : i32 + %1 = tt.make_range {end = 256 : i32, start = 0 : i32} : tensor<256xi32> + %2 = arith.remsi %1, %cst_2 : tensor<256xi32> + %3 = arith.cmpi slt, %2, %cst_1 : tensor<256xi32> + %4 = arith.muli %0, %c256_i32 : i32 + %5 = tt.splat %4 : i32 -> tensor<256xi32> + %6 = arith.addi %5, %1 : tensor<256xi32> + %7 = arith.cmpi slt, %6, %cst_0 : tensor<256xi32> + %8 = arith.andi %3, %7 : tensor<256xi1> + %9 = tt.splat %arg0 : !tt.ptr -> tensor<256x!tt.ptr> + %10 = tt.addptr %9, %2 : tensor<256x!tt.ptr>, tensor<256xi32> + %11 = tt.load %10, %8, %cst : tensor<256x!tt.ptr> + %12 = tt.splat %arg1 : !tt.ptr -> tensor<256x!tt.ptr> + %13 = tt.addptr %12, %1 : tensor<256x!tt.ptr>, tensor<256xi32> + tt.store %13, %11, %8 : tensor<256x!tt.ptr> + tt.return +} + +// CHECK-LABEL: tt.func public @kernel_with_rem_safe( +// CHECK-SAME: %[[VAL_0:.*]]: !tt.ptr {tt.divisibility = 16 : i32}, %[[VAL_1:.*]]: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { +// CHECK: %[[VAL_2:.*]] = arith.constant dense<0.000000e+00> : tensor<256xf32> +// CHECK: %[[VAL_3:.*]] = arith.constant dense<1024> : tensor<256xi32> +// CHECK: %[[VAL_4:.*]] = arith.constant 256 : i32 +// CHECK: %[[VAL_5:.*]] = arith.constant dense<64> : tensor<256xi32> +// CHECK: %[[VAL_6:.*]] = arith.constant dense<128> : tensor<256xi32> +// CHECK: %[[VAL_7:.*]] = tt.get_program_id x : i32 +// CHECK: %[[VAL_8:.*]] = tt.make_range {end = 256 : i32, start = 0 : i32} : tensor<256xi32> +// CHECK: %[[VAL_9:.*]] = arith.remsi %[[VAL_8]], %[[VAL_6]] : tensor<256xi32> +// CHECK: %[[VAL_10:.*]] = arith.cmpi slt, %[[VAL_9]], %[[VAL_5]] : tensor<256xi32> +// CHECK: %[[VAL_11:.*]] = arith.muli %[[VAL_7]], %[[VAL_4]] : i32 +// CHECK: %[[VAL_12:.*]] = tt.splat %[[VAL_11]] : i32 -> tensor<256xi32> +// CHECK: %[[VAL_13:.*]] = arith.addi %[[VAL_12]], %[[VAL_8]] : tensor<256xi32> +// CHECK: %[[VAL_14:.*]] = arith.cmpi slt, %[[VAL_13]], %[[VAL_3]] : tensor<256xi32> +// CHECK: %[[VAL_15:.*]] = arith.andi %[[VAL_10]], %[[VAL_14]] : tensor<256xi1> +// CHECK: %[[VAL_16:.*]] = tt.splat %[[VAL_0]] : !tt.ptr -> tensor<256x!tt.ptr> +// CHECK: %[[VAL_17:.*]] = tt.addptr %[[VAL_16]], %[[VAL_9]] : tensor<256x!tt.ptr>, tensor<256xi32> +// CHECK: %[[VAL_18:.*]] = tt.load %[[VAL_17]], %[[VAL_15]], %[[VAL_2]] : tensor<256x!tt.ptr> +// CHECK: %[[VAL_19:.*]] = tt.splat %[[VAL_1]] : !tt.ptr -> tensor<256x!tt.ptr> +// CHECK: %[[VAL_20:.*]] = tt.addptr %[[VAL_19]], %[[VAL_8]] : tensor<256x!tt.ptr>, tensor<256xi32> +// CHECK: tt.store %[[VAL_20]], %[[VAL_18]], %[[VAL_15]] : tensor<256x!tt.ptr> +// CHECK: tt.return +// CHECK: } + + +// ----- + + +tt.func public @test_remsi_with_broadcast(%arg0: !tt.ptr {tt.divisibility = 16 : i32}) -> tensor<2x2xf32> { + %c0_f32 = arith.constant dense<0.000000e+00> : tensor<2x2xf32> + %c4 = arith.constant dense<4> : tensor<2x2xi32> + %0 = tt.make_range {end = 2 : i32, start = 0 : i32} : tensor<2xi32> + %1 = tt.expand_dims %0 {axis = 0 : i32} : tensor<2xi32> -> tensor<1x2xi32> + %2 = tt.broadcast %1 : tensor<1x2xi32> -> tensor<2x2xi32> + %3 = arith.remsi %2, %c4 : tensor<2x2xi32> + %4 = arith.trunci %3 : tensor<2x2xi32> to tensor<2x2xi1> + %ptrs = tt.splat %arg0 : !tt.ptr -> tensor<2x2x!tt.ptr> + %vals = tt.load %ptrs, %4, %c0_f32 : tensor<2x2x!tt.ptr> + tt.return %vals : tensor<2x2xf32> +} + +// CHECK-LABEL: tt.func public @test_remsi_with_broadcast( +// CHECK-SAME: %[[VAL_0:.*]]: !tt.ptr {tt.divisibility = 16 : i32}) -> tensor<2x2xf32> { +// CHECK: %[[VAL_1:.*]] = arith.constant dense<0.000000e+00> : tensor<2x2xf32> +// CHECK: %[[VAL_2:.*]] = arith.constant dense<4> : tensor<2x2xi32> +// CHECK: %[[VAL_3:.*]] = tt.make_range {end = 2 : i32, start = 0 : i32} : tensor<2xi32> +// CHECK: %[[VAL_4:.*]] = tt.expand_dims %[[VAL_3]] {axis = 0 : i32} : tensor<2xi32> -> tensor<1x2xi32> +// CHECK: %[[VAL_5:.*]] = tt.broadcast %[[VAL_4]] : tensor<1x2xi32> -> tensor<2x2xi32> +// CHECK: %[[VAL_6:.*]] = arith.remsi %[[VAL_5]], %[[VAL_2]] : tensor<2x2xi32> +// CHECK: %[[VAL_7:.*]] = arith.trunci %[[VAL_6]] : tensor<2x2xi32> to tensor<2x2xi1> +// CHECK: %[[VAL_8:.*]] = tt.splat %[[VAL_0]] : !tt.ptr -> tensor<2x2x!tt.ptr> +// CHECK: %[[VAL_9:.*]] = tt.load %[[VAL_8]], %[[VAL_7]], %[[VAL_1]] : tensor<2x2x!tt.ptr> +// CHECK: tt.return %[[VAL_9]] : tensor<2x2xf32> +// CHECK: } diff --git a/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/bubbleupoperation.mlir b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/bubbleupoperation.mlir new file mode 100644 index 0000000000..626f670d1d --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/bubbleupoperation.mlir @@ -0,0 +1,127 @@ +// RUN: triton-opt %s --bubble-up-operation | FileCheck %s + +// CHECK-LABEL: tt.func @test_subi_extract_bubbleup +tt.func @test_subi_extract_bubbleup(%a: tensor<128xi32>, %b: tensor<128xi32>, %i: index, %c: i32) -> i32 { + %0 = arith.subi %a, %b : tensor<128xi32> + %1 = tensor.extract %0[%i] : tensor<128xi32> + %2 = arith.muli %1, %c : i32 + tt.return %2 : i32 +} + + +// CHECK-LABEL: tt.func @test_maxsi_extract_bubbleup +tt.func @test_maxsi_extract_bubbleup(%a: tensor<128xi32>, %b: tensor<128xi32>, %i: index, %c: i32) -> i32 { + %0 = arith.maxsi %a, %b : tensor<128xi32> + %1 = tensor.extract %0[%i] : tensor<128xi32> + %2 = arith.muli %1, %c : i32 + tt.return %2 : i32 +} + + +// CHECK-LABEL: tt.func @test_minsi_extract_bubbleup +tt.func @test_minsi_extract_bubbleup(%a: tensor<128xi32>, %b: tensor<128xi32>, %i: index, %c: i32) -> i32 { + %0 = arith.minsi %a, %b : tensor<128xi32> + %1 = tensor.extract %0[%i] : tensor<128xi32> + %2 = arith.muli %1, %c : i32 + tt.return %2 : i32 +} + + +// CHECK-LABEL: tt.func @test_extf_extract_bubbleup +tt.func @test_extf_extract_bubbleup(%a: tensor<128xf16>, %i: index, %c: f32) -> f32 { + %0 = arith.extf %a : tensor<128xf16> to tensor<128xf32> + %1 = tensor.extract %0[%i] : tensor<128xf32> + %2 = arith.mulf %1, %c : f32 + tt.return %2 : f32 +} + + +// CHECK-LABEL: tt.func @test_minnumf_extract_bubbleup +tt.func @test_minnumf_extract_bubbleup(%a: tensor<128xf32>, %b: tensor<128xf32>, %i: index, %c: f32) -> f32 { + %0 = arith.minnumf %a, %b : tensor<128xf32> + %1 = tensor.extract %0[%i] : tensor<128xf32> + %2 = arith.mulf %1, %c : f32 + tt.return %2 : f32 +} + + +// CHECK-LABEL: tt.func @test_maxnumf_extract_bubbleup +tt.func @test_maxnumf_extract_bubbleup(%a: tensor<128xf32>, %b: tensor<128xf32>, %i: index, %c: f32) -> f32 { + %0 = arith.maxnumf %a, %b : tensor<128xf32> + %1 = tensor.extract %0[%i] : tensor<128xf32> + %2 = arith.mulf %1, %c : f32 + tt.return %2 : f32 +} + + +// CHECK-LABEL: tt.func @test_cmpf_extract_bubbleup +tt.func @test_cmpf_extract_bubbleup(%a: tensor<128xf32>, %b: tensor<128xf32>, %i: index) -> i1 { + %0 = arith.cmpf olt, %a, %b : tensor<128xf32> + %1 = tensor.extract %0[%i] : tensor<128xi1> + tt.return %1 : i1 +} + + +// CHECK-LABEL: tt.func @test_addptr_extract_bubbleup +tt.func @test_addptr_extract_bubbleup(%a: tensor<128x!tt.ptr>, %b: tensor<128xi32>, %i: index) -> !tt.ptr { + %0 = tt.addptr %a, %b : tensor<128x!tt.ptr>, tensor<128xi32> + %1 = tensor.extract %0[%i] : tensor<128x!tt.ptr> + tt.return %1 : !tt.ptr +} + + +// CHECK-LABEL: tt.func @test_ceil_extract_bubbleup +tt.func @test_ceil_extract_bubbleup(%a: tensor<128xf32>, %i: index, %c: f32) -> f32 { + %0 = math.ceil %a : tensor<128xf32> + %1 = tensor.extract %0[%i] : tensor<128xf32> + %2 = arith.mulf %1, %c : f32 + tt.return %2 : f32 +} + + +// CHECK-LABEL: tt.func @test_slice_extract_dropdim_bubbleup +tt.func @test_slice_extract_dropdim_bubbleup(%a: tensor<128x128x128xf32>, %i: index, %j: index) -> f32 { + %0 = tensor.extract_slice %a[0, %i, 0][1, 1, 128][1, 1, 1] : tensor<128x128x128xf32> to tensor<128xf32> + %1 = tensor.extract %0[%j] : tensor<128xf32> + tt.return %1 : f32 +} + + +// CHECK-LABEL: tt.func @test_expand_slice_bubbleup +tt.func @test_expand_slice_bubbleup(%a: tensor<128xf32>, %i: index, %c: f32) -> tensor<1x1xf32> { + %0 = tt.expand_dims %a {axis = 0 : i32} : tensor<128xf32> -> tensor<1x128xf32> + %1 = tensor.extract_slice %0[0, %i][1, 1][1, 1] : tensor<1x128xf32> to tensor<1x1xf32> + tt.return %1 : tensor<1x1xf32> +} + + +// CHECK-LABEL: tt.func @test_expand_slice_dropdim_bubbleup +tt.func @test_expand_slice_dropdim_bubbleup(%a: tensor<128x128xf32>, %i: index, %c: f32) -> tensor<128x1xf32> { + %0 = tt.expand_dims %a {axis = 2 : i32} : tensor<128x128xf32> -> tensor<128x128x1xf32> + %1 = tensor.extract_slice %0[%i, 0, 0][1, 128, 1][1, 1, 1] : tensor<128x128x1xf32> to tensor<128x1xf32> + tt.return %1 : tensor<128x1xf32> +} + + +// CHECK-LABEL: tt.func @test_splat_slice_bubbleup +tt.func @test_splat_slice_bubbleup(%a: f32, %i: index, %c: f32) -> tensor<1xf32> { + %0 = tt.splat %a : f32 -> tensor<128xf32> + %1 = tensor.extract_slice %0[%i][1][1] : tensor<128xf32> to tensor<1xf32> + tt.return %1 : tensor<1xf32> +} + + +// CHECK-LABEL: tt.func @test_makerange_slice_bubbleup +tt.func @test_makerange_slice_bubbleup(%i: index, %c: f32) -> tensor<1xi32> { + %0 = tt.make_range {start = 0 : i32, end = 128 : i32} : tensor<128xi32> + %1 = tensor.extract_slice %0[%i][1][1] : tensor<128xi32> to tensor<1xi32> + tt.return %1 : tensor<1xi32> +} + + +// CHECK-LABEL: tt.func @test_slice_all_bubbleup +tt.func @test_slice_all_bubbleup(%i: index, %c: f32) -> tensor<128xi32> { + %0 = tt.make_range {start = 0 : i32, end = 128 : i32} : tensor<128xi32> + %1 = tensor.extract_slice %0[0][128][1] : tensor<128xi32> to tensor<128xi32> + tt.return %1 : tensor<128xi32> +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/if_simplifier.mlir b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/if_simplifier.mlir new file mode 100644 index 0000000000..7f71ec94a9 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/if_simplifier.mlir @@ -0,0 +1,45 @@ +// RUN: triton-opt --triton-to-unstructure --split-input-file %s | FileCheck %s --implicit-check-not="DiscreteMemAccess" +// CHECK-LABEL: tt.func public @triton_for_if_load +// CHECK: %[[CST:.*]] = arith.constant dense<0> : tensor<16xi32> +// CHECK: %[[CST0:.*]] = arith.constant dense<1> : tensor<16xi32> +// CHECK: %[[SEL:.*]] = arith.select %{{.*}}, %[[CST0]], %[[CST]] : tensor<16xi32> +// CHECK: %[[ADD:.*]] = arith.addi %{{.*}}, %[[SEL]] : tensor<16xi32> +// CHECK: %[[ADDPTR:.*]] = tt.addptr %{{.*}}, %[[ADD]] : tensor<16x!tt.ptr>, tensor<16xi32> + + +module { + tt.func public @triton_for_if_load(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %c1_i32 = arith.constant 1 : i32 + %c2_i32 = arith.constant 2 : i32 + %cst = arith.constant dense<0.000000e+00> : tensor<16xf32> + %cst_0 = arith.constant dense<32> : tensor<16xi32> + %cst_1 = arith.constant dense<1> : tensor<16xi32> + %c16_i32 = arith.constant 16 : i32 + %c0_i32 = arith.constant 0 : i32 + %0 = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> + %1 = tt.get_program_id x : i32 + %2 = arith.cmpi ne, %1, %c0_i32 : i32 + %3 = tt.splat %arg0 : !tt.ptr -> tensor<16x!tt.ptr> + %4 = tt.splat %arg1 : !tt.ptr -> tensor<16x!tt.ptr> + %5:2 = scf.for %arg2 = %c0_i32 to %c2_i32 step %c1_i32 iter_args(%arg3 = %0, %arg4 = %0) -> (tensor<16xi32>, tensor<16xi32>) : i32 { + %6 = arith.muli %arg2, %c16_i32 : i32 + %7 = tt.splat %6 : i32 -> tensor<16xi32> + %8 = arith.addi %arg3, %7 : tensor<16xi32> + %9 = arith.addi %arg4, %7 : tensor<16xi32> + %10 = scf.if %2 -> (tensor<16xi32>) { + %16 = arith.addi %8, %cst_1 : tensor<16xi32> + scf.yield %16 : tensor<16xi32> + } else { + scf.yield %8 : tensor<16xi32> + } + %11 = tt.addptr %3, %10 : tensor<16x!tt.ptr>, tensor<16xi32> + %12 = arith.cmpi slt, %10, %cst_0 : tensor<16xi32> + %13 = tt.load %11 : tensor<16x!tt.ptr> + %14 = arith.select %12, %13, %cst : tensor<16xi1>, tensor<16xf32> + %15 = tt.addptr %4, %9 : tensor<16x!tt.ptr>, tensor<16xi32> + tt.store %15, %14 : tensor<16x!tt.ptr> + scf.yield %10, %9 : tensor<16xi32>, tensor<16xi32> + } + tt.return + } +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/indirect_func.mlir b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/indirect_func.mlir new file mode 100644 index 0000000000..1e4cf2a3d0 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/indirect_func.mlir @@ -0,0 +1,141 @@ +// RUN: triton-opt --triton-to-structured '--discrete-mask-access-conversion=compile-on-910-95=True force-simt-template=True' '--triton-to-unstructure=compile-on-910-95=True force-simt-template=True' %s --split-input-file | FileCheck %s + +// tt.load -> tt.indirect_load +tt.func public @triton_ldst_indirect_05_kernel(%arg0: !tt.ptr, %arg1: !tt.ptr, %arg2: !tt.ptr, %arg3: i32) attributes {noinline = false} { + %cst = arith.constant dense<16> : tensor<8x1xi32> + %c8_i32 = arith.constant 8 : i32 + %0 = tt.get_program_id x : i32 + %1 = arith.muli %0, %c8_i32 : i32 + %2 = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> + %3 = tt.splat %1 : i32 -> tensor<8xi32> + %4 = arith.addi %3, %2 : tensor<8xi32> + %5 = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<16x!tt.ptr> + %7 = tt.addptr %6, %5 : tensor<16x!tt.ptr>, tensor<16xi32> + %8 = tt.load %7 : tensor<16x!tt.ptr> + %9 = tt.expand_dims %2 {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32> + %10 = tt.splat %arg3 : i32 -> tensor<8x1xi32> + %11 = arith.muli %9, %10 : tensor<8x1xi32> + %12 = tt.expand_dims %8 {axis = 0 : i32} : tensor<16xi64> -> tensor<1x16xi64> + %13 = arith.extsi %11 : tensor<8x1xi32> to tensor<8x1xi64> + %14 = tt.broadcast %13 : tensor<8x1xi64> -> tensor<8x16xi64> + %15 = tt.broadcast %12 : tensor<1x16xi64> -> tensor<8x16xi64> + %16 = arith.addi %14, %15 : tensor<8x16xi64> + %17 = tt.splat %arg2 : !tt.ptr -> tensor<8x16x!tt.ptr> + %18 = tt.addptr %17, %16 : tensor<8x16x!tt.ptr>, tensor<8x16xi64> + %19 = tt.load %18 : tensor<8x16x!tt.ptr> + %20 = math.exp %19 : tensor<8x16xf32> + %21 = tt.expand_dims %4 {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32> + %22 = arith.muli %21, %cst : tensor<8x1xi32> + %23 = tt.expand_dims %5 {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32> + %24 = tt.broadcast %22 : tensor<8x1xi32> -> tensor<8x16xi32> + %25 = tt.broadcast %23 : tensor<1x16xi32> -> tensor<8x16xi32> + %26 = arith.addi %24, %25 : tensor<8x16xi32> + %27 = tt.splat %arg0 : !tt.ptr -> tensor<8x16x!tt.ptr> + %28 = tt.addptr %27, %26 : tensor<8x16x!tt.ptr>, tensor<8x16xi32> + tt.store %28, %20 : tensor<8x16x!tt.ptr> + tt.return +} + +// CHECK-LABEL: tt.func public @triton_ldst_indirect_05_kernel( +// CHECK-SAME: %[[VAL_0:.*]]: !tt.ptr, %[[VAL_1:.*]]: !tt.ptr, %[[VAL_2:.*]]: !tt.ptr, %[[VAL_3:.*]]: i32) attributes {noinline = false} { +// CHECK: %[[VAL_4:.*]] = arith.constant dense<16> : tensor<8x1xi32> +// CHECK: %[[VAL_6:.*]] = tt.get_program_id x : i32 +// CHECK: %[[VAL_7:.*]] = arith.muli %[[VAL_6]], %{{.*}} : i32 +// CHECK: %[[VAL_8:.*]] = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> +// CHECK: %[[VAL_9:.*]] = tt.splat %[[VAL_7]] : i32 -> tensor<8xi32> +// CHECK: %[[VAL_10:.*]] = arith.addi %[[VAL_9]], %[[VAL_8]] : tensor<8xi32> +// CHECK: %[[VAL_11:.*]] = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> +// CHECK: %[[VAL_12:.*]] = tt.splat %[[VAL_1:.*]] : !tt.ptr -> tensor<16x!tt.ptr> +// CHECK: %[[VAL_13:.*]] = tt.addptr %[[VAL_12]], %[[VAL_11]] : tensor<16x!tt.ptr>, tensor<16xi32> +// CHECK: %[[VAL_14:.*]] = tt.load %[[VAL_13]] : tensor<16x!tt.ptr> +// CHECK: %[[VAL_15:.*]] = tt.expand_dims %[[VAL_8]] {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32> +// CHECK: %[[VAL_16:.*]] = tt.splat %[[VAL_3:.*]] : i32 -> tensor<8x1xi32> +// CHECK: %[[VAL_17:.*]] = arith.muli %[[VAL_15]], %[[VAL_16]] : tensor<8x1xi32> +// CHECK: %[[VAL_18:.*]] = tt.expand_dims %[[VAL_14]] {axis = 0 : i32} : tensor<16xi64> -> tensor<1x16xi64> +// CHECK: %[[VAL_19:.*]] = arith.extsi %[[VAL_17]] : tensor<8x1xi32> to tensor<8x1xi64> +// CHECK: %[[VAL_20:.*]] = tt.broadcast %[[VAL_19]] : tensor<8x1xi64> -> tensor<8x16xi64> +// CHECK: %[[VAL_21:.*]] = tt.broadcast %[[VAL_18]] : tensor<1x16xi64> -> tensor<8x16xi64> +// CHECK: %[[VAL_22:.*]] = arith.addi %[[VAL_20]], %[[VAL_21]] : tensor<8x16xi64> +// CHECK: %[[VAL_23:.*]] = ascend.indirect_load %[[VAL_2:.*]] : , %[[VAL_22]] : tensor<8x16xi64> -> tensor<8x16xf32> +// CHECK: %[[VAL_24:.*]] = math.exp %[[VAL_23]] : tensor<8x16xf32> +// CHECK: %[[VAL_25:.*]] = tt.expand_dims %[[VAL_10]] {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32> +// CHECK: %[[VAL_26:.*]] = arith.muli %[[VAL_25]], %[[VAL_4]] : tensor<8x1xi32> +// CHECK: %[[VAL_27:.*]] = tt.expand_dims %[[VAL_11]] {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32> +// CHECK: %[[VAL_28:.*]] = tt.broadcast %[[VAL_26]] : tensor<8x1xi32> -> tensor<8x16xi32> +// CHECK: %[[VAL_29:.*]] = tt.broadcast %[[VAL_27]] : tensor<1x16xi32> -> tensor<8x16xi32> +// CHECK: %[[VAL_30:.*]] = arith.addi %[[VAL_28]], %[[VAL_29]] : tensor<8x16xi32> +// CHECK: %[[VAL_31:.*]] = tt.splat %[[VAL_0:.*]] : !tt.ptr -> tensor<8x16x!tt.ptr> +// CHECK: %[[VAL_32:.*]] = tt.addptr %[[VAL_31]], %[[VAL_30]] : tensor<8x16x!tt.ptr>, tensor<8x16xi32> +// CHECK: tt.store %[[VAL_32]], %[[VAL_24]] : tensor<8x16x!tt.ptr> +// CHECK: tt.return +// CHECK: } + +// ----- + +// tt.store -> tt.indirect_store +tt.func public @triton_ldst_indirect_08_kernel(%arg0: !tt.ptr, %arg1: !tt.ptr, %arg2: !tt.ptr, %arg3: i32) attributes {noinline = false} { + %cst = arith.constant dense<32> : tensor<8x1xi32> + %c8_i32 = arith.constant 8 : i32 + %0 = tt.get_program_id x : i32 + %1 = arith.muli %0, %c8_i32 : i32 + %2 = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> + %3 = tt.splat %1 : i32 -> tensor<8xi32> + %4 = arith.addi %3, %2 : tensor<8xi32> + %5 = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<16x!tt.ptr> + %7 = tt.addptr %6, %5 : tensor<16x!tt.ptr>, tensor<16xi32> + %8 = tt.load %7 : tensor<16x!tt.ptr> + %9 = tt.expand_dims %2 {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32> + %10 = tt.splat %arg3 : i32 -> tensor<8x1xi32> + %11 = arith.muli %9, %10 : tensor<8x1xi32> + %12 = tt.expand_dims %8 {axis = 0 : i32} : tensor<16xi64> -> tensor<1x16xi64> + %13 = arith.extsi %11 : tensor<8x1xi32> to tensor<8x1xi64> + %14 = tt.broadcast %13 : tensor<8x1xi64> -> tensor<8x16xi64> + %15 = tt.broadcast %12 : tensor<1x16xi64> -> tensor<8x16xi64> + %16 = arith.addi %14, %15 : tensor<8x16xi64> + %17 = tt.splat %arg2 : !tt.ptr -> tensor<8x16x!tt.ptr> + %18 = tt.addptr %17, %16 : tensor<8x16x!tt.ptr>, tensor<8x16xi64> + %19 = tt.load %18 : tensor<8x16x!tt.ptr> + %20 = math.exp %19 : tensor<8x16xf32> + %21 = tt.expand_dims %4 {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32> + %22 = arith.muli %21, %cst : tensor<8x1xi32> + %23 = arith.extsi %22 : tensor<8x1xi32> to tensor<8x1xi64> + %24 = tt.broadcast %23 : tensor<8x1xi64> -> tensor<8x16xi64> + %25 = arith.addi %24, %15 : tensor<8x16xi64> + %26 = tt.splat %arg0 : !tt.ptr -> tensor<8x16x!tt.ptr> + %27 = tt.addptr %26, %25 : tensor<8x16x!tt.ptr>, tensor<8x16xi64> + tt.store %27, %20 : tensor<8x16x!tt.ptr> + tt.return +} + +// CHECK-LABEL: tt.func public @triton_ldst_indirect_08_kernel( +// CHECK-SAME: %[[VAL_0:.*]]: !tt.ptr, %[[VAL_1:.*]]: !tt.ptr, %[[VAL_2:.*]]: !tt.ptr, %[[VAL_3:.*]]: i32) attributes {noinline = false} { +// CHECK: %[[VAL_4:.*]] = arith.constant dense<32> : tensor<8x1xi32> +// CHECK: %[[VAL_6:.*]] = tt.get_program_id x : i32 +// CHECK: %[[VAL_7:.*]] = arith.muli %[[VAL_6:.*]], %{{.*}} : i32 +// CHECK: %[[VAL_8:.*]] = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> +// CHECK: %[[VAL_9:.*]] = tt.splat %[[VAL_7:.*]] : i32 -> tensor<8xi32> +// CHECK: %[[VAL_10:.*]] = arith.addi %[[VAL_9:.*]], %[[VAL_8:.*]] : tensor<8xi32> +// CHECK: %[[VAL_11:.*]] = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> +// CHECK: %[[VAL_12:.*]] = tt.splat %[[VAL_1:.*]] : !tt.ptr -> tensor<16x!tt.ptr> +// CHECK: %[[VAL_13:.*]] = tt.addptr %[[VAL_12:.*]], %[[VAL_11:.*]] : tensor<16x!tt.ptr>, tensor<16xi32> +// CHECK: %[[VAL_14:.*]] = tt.load %[[VAL_13:.*]] : tensor<16x!tt.ptr> +// CHECK: %[[VAL_15:.*]] = tt.expand_dims %[[VAL_8:.*]] {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32> +// CHECK: %[[VAL_16:.*]] = tt.splat %[[VAL_3:.*]] : i32 -> tensor<8x1xi32> +// CHECK: %[[VAL_17:.*]] = arith.muli %[[VAL_15:.*]], %[[VAL_16:.*]] : tensor<8x1xi32> +// CHECK: %[[VAL_18:.*]] = tt.expand_dims %[[VAL_14:.*]] {axis = 0 : i32} : tensor<16xi64> -> tensor<1x16xi64> +// CHECK: %[[VAL_19:.*]] = arith.extsi %[[VAL_17:.*]] : tensor<8x1xi32> to tensor<8x1xi64> +// CHECK: %[[VAL_20:.*]] = tt.broadcast %[[VAL_19:.*]] : tensor<8x1xi64> -> tensor<8x16xi64> +// CHECK: %[[VAL_21:.*]] = tt.broadcast %[[VAL_18:.*]] : tensor<1x16xi64> -> tensor<8x16xi64> +// CHECK: %[[VAL_22:.*]] = arith.addi %[[VAL_20:.*]], %[[VAL_21:.*]] : tensor<8x16xi64> +// CHECK: %[[VAL_23:.*]] = ascend.indirect_load %[[VAL_2:.*]] : , %[[VAL_22:.*]] : tensor<8x16xi64> -> tensor<8x16xf32> +// CHECK: %[[VAL_24:.*]] = math.exp %[[VAL_23:.*]] : tensor<8x16xf32> +// CHECK: %[[VAL_25:.*]] = tt.expand_dims %[[VAL_10:.*]] {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32> +// CHECK: %[[VAL_26:.*]] = arith.muli %[[VAL_25:.*]], %[[VAL_4:.*]] : tensor<8x1xi32> +// CHECK: %[[VAL_27:.*]] = arith.extsi %[[VAL_26:.*]] : tensor<8x1xi32> to tensor<8x1xi64> +// CHECK: %[[VAL_28:.*]] = tt.broadcast %[[VAL_27:.*]] : tensor<8x1xi64> -> tensor<8x16xi64> +// CHECK: %[[VAL_29:.*]] = arith.addi %[[VAL_28:.*]], %[[VAL_21:.*]] : tensor<8x16xi64> +// CHECK: ascend.indirect_store %[[VAL_0:.*]] : , %[[VAL_29:.*]] : tensor<8x16xi64>, %[[VAL_24:.*]] : tensor<8x16xf32> +// CHECK: tt.return +// CHECK: } diff --git a/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/indirect_load.mlir b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/indirect_load.mlir new file mode 100644 index 0000000000..40f75c27f2 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/indirect_load.mlir @@ -0,0 +1,42 @@ +// RUN: triton-opt --triton-to-structured '--discrete-mask-access-conversion=compile-on-910-95=True force-simt-template=True' '--triton-to-unstructure=compile-on-910-95=True force-simt-template=True' %s --split-input-file | FileCheck %s + +// tt.load -> ascend.indirect_load +tt.func public @triton_indirect_load_kernel(%arg0: !tt.ptr, %arg1: !tt.ptr, %arg2: !tt.ptr, %arg3: i32) attributes {noinline = false} { + %cst = arith.constant dense<32> : tensor<8x1xi32> + %c8_i32 = arith.constant 8 : i32 + %0 = tt.get_program_id x : i32 + %1 = arith.muli %0, %c8_i32 : i32 + %2 = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> + %3 = tt.splat %1 : i32 -> tensor<8xi32> + %4 = arith.addi %3, %2 : tensor<8xi32> + %5 = tt.make_range {end = 32 : i32, start = 0 : i32} : tensor<32xi32> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<32x!tt.ptr> + %7 = tt.addptr %6, %5 : tensor<32x!tt.ptr>, tensor<32xi32> + %8 = tt.load %7 : tensor<32x!tt.ptr> + %9 = tt.expand_dims %2 {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32> + %10 = tt.splat %arg3 : i32 -> tensor<8x1xi32> + %11 = arith.muli %9, %10 : tensor<8x1xi32> + %12 = tt.expand_dims %8 {axis = 0 : i32} : tensor<32xi64> -> tensor<1x32xi64> + %13 = arith.extsi %11 : tensor<8x1xi32> to tensor<8x1xi64> + %14 = tt.broadcast %13 : tensor<8x1xi64> -> tensor<8x32xi64> + %15 = tt.broadcast %12 : tensor<1x32xi64> -> tensor<8x32xi64> + %16 = arith.addi %14, %15 : tensor<8x32xi64> + %17 = tt.splat %arg2 : !tt.ptr -> tensor<8x32x!tt.ptr> + %18 = tt.addptr %17, %16 : tensor<8x32x!tt.ptr>, tensor<8x32xi64> + %19 = tt.load %18 : tensor<8x32x!tt.ptr> + %20 = math.exp %19 : tensor<8x32xf32> + %21 = tt.expand_dims %4 {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32> + %22 = arith.muli %21, %cst : tensor<8x1xi32> + %23 = tt.expand_dims %5 {axis = 0 : i32} : tensor<32xi32> -> tensor<1x32xi32> + %24 = tt.broadcast %22 : tensor<8x1xi32> -> tensor<8x32xi32> + %25 = tt.broadcast %23 : tensor<1x32xi32> -> tensor<8x32xi32> + %26 = arith.addi %24, %25 : tensor<8x32xi32> + %27 = tt.splat %arg0 : !tt.ptr -> tensor<8x32x!tt.ptr> + %28 = tt.addptr %27, %26 : tensor<8x32x!tt.ptr>, tensor<8x32xi32> + tt.store %28, %20 : tensor<8x32x!tt.ptr> + tt.return +} + +// CHECK-LABEL: tt.func public @triton_indirect_load_kernel +// CHECK: ascend.indirect_load{{.*}} : tensor<8x32xi64> -> tensor<8x32xf32> +// CHECK: tt.store diff --git a/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/indirect_mem_access.mlir b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/indirect_mem_access.mlir new file mode 100644 index 0000000000..a80a233fe2 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/indirect_mem_access.mlir @@ -0,0 +1,129 @@ +// RUN: triton-opt %s --triton-to-unstructure | FileCheck %s +// RUN: triton-opt %s '--triton-to-unstructure=force-scalarize-mode=True' | FileCheck %s +// RUN: triton-opt %s '--triton-to-unstructure=compile-on-910-95=True' | FileCheck %s +// RUN: triton-opt %s '--triton-to-unstructure=force-simt-template=True' | FileCheck %s + + +// CHECK-LABEL: tt.func @triton_indirect_load +tt.func @triton_indirect_load(%arg0: !tt.ptr, %arg1: !tt.ptr, %arg2: !tt.ptr) -> tensor<4096xf32> { + %0 = tt.make_range {end = 4096 : i32, start = 0 : i32} : tensor<4096xi32> + %1 = tt.splat %arg2 : !tt.ptr -> tensor<4096x!tt.ptr> + %2 = tt.addptr %1, %0 : tensor<4096x!tt.ptr>, tensor<4096xi32> + %3 = tt.load %2 : tensor<4096x!tt.ptr> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<4096x!tt.ptr> + %5 = tt.addptr %4, %3 : tensor<4096x!tt.ptr>, tensor<4096xi64> + %6 = tt.load %5 : tensor<4096x!tt.ptr> + annotation.mark %6 {mayDiscretememaccess} : tensor<4096xf32> + tt.return %6 : tensor<4096xf32> +} + + +// CHECK-LABEL: tt.func @triton_indirect_load_error_mark +tt.func @triton_indirect_load_error_mark(%arg0: !tt.ptr, %arg1: !tt.ptr, %arg2: !tt.ptr) -> tensor<4096xf32> { + %0 = tt.make_range {end = 4096 : i32, start = 0 : i32} : tensor<4096xi32> + %1 = tt.splat %arg2 : !tt.ptr -> tensor<4096x!tt.ptr> + %2 = tt.addptr %1, %0 : tensor<4096x!tt.ptr>, tensor<4096xi32> + %3 = tt.load %2 : tensor<4096x!tt.ptr> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<4096x!tt.ptr> + %5 = tt.addptr %4, %3 : tensor<4096x!tt.ptr>, tensor<4096xi64> + %6 = tt.load %5 : tensor<4096x!tt.ptr> + annotation.mark %6 {error_mark} : tensor<4096xf32> + tt.return %6 : tensor<4096xf32> +} + + +// CHECK-LABEL: tt.func @triton_indirect_store +tt.func @triton_indirect_store(%arg0: !tt.ptr, %arg1: !tt.ptr, %arg2: !tt.ptr) { + %0 = tt.make_range {end = 4096 : i32, start = 0 : i32} : tensor<4096xi32> + %1 = tt.splat %arg2 : !tt.ptr -> tensor<4096x!tt.ptr> + %2 = tt.addptr %1, %0 : tensor<4096x!tt.ptr>, tensor<4096xi32> + %3 = tt.load %2 : tensor<4096x!tt.ptr> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<4096x!tt.ptr> + %5 = tt.addptr %4, %0 : tensor<4096x!tt.ptr>, tensor<4096xi32> + %6 = tt.load %5 : tensor<4096x!tt.ptr> + annotation.mark %6 {mayDiscretememaccess} : tensor<4096xf32> + %7 = tt.splat %arg1 : !tt.ptr -> tensor<4096x!tt.ptr> + %8 = tt.addptr %7, %3 : tensor<4096x!tt.ptr>, tensor<4096xi64> + tt.store %8, %6 : tensor<4096x!tt.ptr> + tt.return +} + + +// CHECK-LABEL: tt.func @triton_indirect_store_error_mark +tt.func @triton_indirect_store_error_mark(%arg0: !tt.ptr, %arg1: !tt.ptr, %arg2: !tt.ptr) { + %0 = tt.make_range {end = 4096 : i32, start = 0 : i32} : tensor<4096xi32> + %1 = tt.splat %arg2 : !tt.ptr -> tensor<4096x!tt.ptr> + %2 = tt.addptr %1, %0 : tensor<4096x!tt.ptr>, tensor<4096xi32> + %3 = tt.load %2 : tensor<4096x!tt.ptr> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<4096x!tt.ptr> + %5 = tt.addptr %4, %0 : tensor<4096x!tt.ptr>, tensor<4096xi32> + %6 = tt.load %5 : tensor<4096x!tt.ptr> + annotation.mark %6 {error_mark} : tensor<4096xf32> + %7 = tt.splat %arg1 : !tt.ptr -> tensor<4096x!tt.ptr> + %8 = tt.addptr %7, %3 : tensor<4096x!tt.ptr>, tensor<4096xi64> + tt.store %8, %6 : tensor<4096x!tt.ptr> + tt.return +} + + +// CHECK-LABEL: tt.func @triton_indirect_atomic_add +tt.func @triton_indirect_atomic_add( + %arg0: !tt.ptr {tt.divisibility = 16 : i32}, + %arg1: !tt.ptr {tt.divisibility = 16 : i32}, + %arg2: !tt.ptr {tt.divisibility = 16 : i32} +) attributes {noinline = false} { + %cst = arith.constant dense<1.000000e+00> : tensor<4096xf32> + %cst_0 = arith.constant dense : tensor<4096xi1> + %0 = tt.make_range {end = 4096 : i32, start = 0 : i32} : tensor<4096xi32> + %1 = tt.splat %arg2 : !tt.ptr -> tensor<4096x!tt.ptr> + %2 = tt.addptr %1, %0 : tensor<4096x!tt.ptr>, tensor<4096xi32> + %3 = tt.load %2 : tensor<4096x!tt.ptr> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<4096x!tt.ptr> + %5 = tt.addptr %4, %3 : tensor<4096x!tt.ptr>, tensor<4096xi64> + %6 = tt.atomic_rmw fadd, acq_rel, gpu, %5, %cst, %cst_0 : (tensor<4096x!tt.ptr>, tensor<4096xf32>, tensor<4096xi1>) -> tensor<4096xf32> + tt.return +} + + +// CHECK-LABEL: tt.func @triton_indirect_atomic_add_with_discret_mask +tt.func @triton_indirect_atomic_add_with_discret_mask( + %arg0: !tt.ptr {tt.divisibility = 16 : i32}, + %arg1: !tt.ptr {tt.divisibility = 16 : i32}, + %arg2: !tt.ptr {tt.divisibility = 16 : i32} +) attributes {noinline = false} { + %cst = arith.constant dense<1.000000e+00> : tensor<4096xf32> + %cst_0 = arith.constant dense<4096> : tensor<4096xi64> + %c1 = arith.constant 1.0 : f32 + %c0 = arith.constant 0.0 : f32 + %cst_true = tt.splat %c1 : f32 -> tensor<4096xf32> + %cst_false = tt.splat %c0 : f32 -> tensor<4096xf32> + %0 = tt.make_range {end = 4096 : i32, start = 0 : i32} : tensor<4096xi32> + %1 = tt.splat %arg2 : !tt.ptr -> tensor<4096x!tt.ptr> + %2 = tt.addptr %1, %0 : tensor<4096x!tt.ptr>, tensor<4096xi32> + %3 = tt.load %2 : tensor<4096x!tt.ptr> + %mask = arith.cmpi slt, %3, %cst_0 : tensor<4096xi64> + %selected_val = arith.select %mask, %cst_true, %cst_false : tensor<4096xi1>, tensor<4096xf32> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<4096x!tt.ptr> + %5 = tt.addptr %4, %3 : tensor<4096x!tt.ptr>, tensor<4096xi64> + %6 = tt.atomic_rmw fadd, acq_rel, gpu, %5, %selected_val, %mask {DiscreteMask} : (tensor<4096x!tt.ptr>, tensor<4096xf32>, tensor<4096xi1>) -> tensor<4096xf32> + tt.return +} + + +// CHECK-LABEL: tt.func @triton_indirect_atomic_cas +tt.func @triton_indirect_atomic_cas( + %arg0: !tt.ptr {tt.divisibility = 16 : i32}, + %arg1: !tt.ptr {tt.divisibility = 16 : i32}, + %arg2: !tt.ptr {tt.divisibility = 16 : i32} +) attributes {noinline = false} { + %cst = arith.constant dense<2.0> : tensor<4096xf32> + %cst_0 = arith.constant dense<1.0> : tensor<4096xf32> + %0 = tt.make_range {end = 4096 : i32, start = 0 : i32} : tensor<4096xi32> + %1 = tt.splat %arg2 : !tt.ptr -> tensor<4096x!tt.ptr> + %2 = tt.addptr %1, %0 : tensor<4096x!tt.ptr>, tensor<4096xi32> + %3 = tt.load %2 : tensor<4096x!tt.ptr> + %4 = tt.splat %arg0 : !tt.ptr -> tensor<4096x!tt.ptr> + %5 = tt.addptr %4, %3 : tensor<4096x!tt.ptr>, tensor<4096xi64> + %6 = tt.atomic_cas acq_rel, gpu, %5, %cst_0, %cst : (tensor<4096x!tt.ptr>, tensor<4096xf32>, tensor<4096xf32>) -> tensor<4096xf32> + tt.return +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/indirect_store.mlir b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/indirect_store.mlir new file mode 100644 index 0000000000..263e66453f --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/indirect_store.mlir @@ -0,0 +1,42 @@ +// RUN: triton-opt --triton-to-structured '--discrete-mask-access-conversion=compile-on-910-95=True force-simt-template=True' '--triton-to-unstructure=compile-on-910-95=True force-simt-template=True' %s --split-input-file | FileCheck %s + +// tt.store -> ascend.indirect_store +tt.func public @triton_indirect_store_kernel(%arg0: !tt.ptr, %arg1: !tt.ptr, %arg2: !tt.ptr, %arg3: i32) attributes {noinline = false} { + %cst = arith.constant dense<32> : tensor<8x1xi32> + %c8_i32 = arith.constant 8 : i32 + %0 = tt.get_program_id x : i32 + %1 = arith.muli %0, %c8_i32 : i32 + %2 = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> + %3 = tt.splat %1 : i32 -> tensor<8xi32> + %4 = arith.addi %3, %2 : tensor<8xi32> + %5 = tt.make_range {end = 32 : i32, start = 0 : i32} : tensor<32xi32> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<32x!tt.ptr> + %7 = tt.addptr %6, %5 : tensor<32x!tt.ptr>, tensor<32xi32> + %8 = tt.load %7 : tensor<32x!tt.ptr> + %9 = tt.expand_dims %2 {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32> + %10 = tt.splat %arg3 : i32 -> tensor<8x1xi32> + %11 = arith.muli %9, %10 : tensor<8x1xi32> + %12 = tt.expand_dims %8 {axis = 0 : i32} : tensor<32xi64> -> tensor<1x32xi64> + %13 = arith.extsi %11 : tensor<8x1xi32> to tensor<8x1xi64> + %14 = tt.broadcast %13 : tensor<8x1xi64> -> tensor<8x32xi64> + %15 = tt.broadcast %12 : tensor<1x32xi64> -> tensor<8x32xi64> + %16 = arith.addi %14, %15 : tensor<8x32xi64> + %17 = tt.splat %arg2 : !tt.ptr -> tensor<8x32x!tt.ptr> + %18 = tt.addptr %17, %16 : tensor<8x32x!tt.ptr>, tensor<8x32xi64> + %19 = tt.load %18 : tensor<8x32x!tt.ptr> + %20 = math.exp %19 : tensor<8x32xf32> + %21 = tt.expand_dims %4 {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32> + %22 = arith.muli %21, %cst : tensor<8x1xi32> + %23 = arith.extsi %22 : tensor<8x1xi32> to tensor<8x1xi64> + %24 = tt.broadcast %23 : tensor<8x1xi64> -> tensor<8x32xi64> + %25 = arith.addi %24, %15 : tensor<8x32xi64> + %26 = tt.splat %arg0 : !tt.ptr -> tensor<8x32x!tt.ptr> + %27 = tt.addptr %26, %25 : tensor<8x32x!tt.ptr>, tensor<8x32xi64> + tt.store %27, %20 : tensor<8x32x!tt.ptr> + tt.return +} + +// CHECK-LABEL: tt.func public @triton_indirect_store_kernel( +// CHECK: ascend.indirect_store %[[VAL_0:.*]] : , %[[VAL_29:.*]] : tensor<8x32xi64>, %[[VAL_24:.*]] : tensor<8x32xf32> +// CHECK: tt.return +// CHECK: } diff --git a/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/nested_loop.mlir b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/nested_loop.mlir new file mode 100644 index 0000000000..e19caadd25 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/nested_loop.mlir @@ -0,0 +1,204 @@ +// RUN: triton-opt --triton-to-unstructure --split-input-file %s | FileCheck %s + +tt.func public @test_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: !tt.ptr {tt.divisibility = 16 : i32}, %arg3: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %c3_i32 = arith.constant 3 : i32 + %c1_i32 = arith.constant 1 : i32 + %c0_i32 = arith.constant 0 : i32 + %cst = arith.constant dense<128> : tensor<128xi32> + %cst_0 = arith.constant dense<0> : tensor<128xi32> + %cst_1 = arith.constant dense<300> : tensor<128xi32> + %c128_i32 = arith.constant 128 : i32 + %0 = tt.get_program_id x : i32 + %1 = arith.muli %0, %c128_i32 : i32 + %2 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32> + %3 = tt.splat %1 : i32 -> tensor<128xi32> + %4 = arith.addi %3, %2 : tensor<128xi32> + %5 = tt.splat %arg0 : !tt.ptr -> tensor<128x!tt.ptr> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<128x!tt.ptr> + %7 = tt.addptr %6, %4 : tensor<128x!tt.ptr>, tensor<128xi32> + %8 = tt.splat %arg2 : !tt.ptr -> tensor<128x!tt.ptr> + %9 = tt.addptr %8, %4 : tensor<128x!tt.ptr>, tensor<128xi32> + %10 = tt.load %9 : tensor<128x!tt.ptr> + %11 = tt.splat %arg3 : !tt.ptr -> tensor<128x!tt.ptr> + %12 = tt.addptr %11, %10 : tensor<128x!tt.ptr>, tensor<128xi32> + %13:3 = scf.for %arg4 = %c0_i32 to %c3_i32 step %c1_i32 iter_args(%arg5 = %4, %arg6 = %7, %arg7 = %12) -> (tensor<128xi32>, tensor<128x!tt.ptr>, tensor<128x!tt.ptr>) : i32 { + %14:3 = scf.for %arg8 = %c0_i32 to %c3_i32 step %c1_i32 iter_args(%arg9 = %arg5, %arg10 = %arg6, %arg11 = %arg7) -> (tensor<128xi32>, tensor<128x!tt.ptr>, tensor<128x!tt.ptr>) : i32 { + %18 = arith.cmpi slt, %arg9, %cst_1 : tensor<128xi32> + %19 = tt.addptr %5, %arg9 : tensor<128x!tt.ptr>, tensor<128xi32> + %20 = tt.load %19, %18, %cst_0 : tensor<128x!tt.ptr> + %21 = tt.load %arg11 : tensor<128x!tt.ptr> + %22 = arith.addi %20, %21 : tensor<128xi32> + tt.store %arg10, %22, %18 : tensor<128x!tt.ptr> + %23 = arith.addi %arg9, %cst : tensor<128xi32> + %24 = tt.addptr %arg10, %cst : tensor<128x!tt.ptr>, tensor<128xi32> + %25 = tt.addptr %arg11, %cst : tensor<128x!tt.ptr>, tensor<128xi32> + scf.yield %23, %24, %25 : tensor<128xi32>, tensor<128x!tt.ptr>, tensor<128x!tt.ptr> + } + %15 = arith.addi %14#0, %cst : tensor<128xi32> + %16 = tt.addptr %14#1, %cst : tensor<128x!tt.ptr>, tensor<128xi32> + %17 = tt.addptr %14#2, %cst : tensor<128x!tt.ptr>, tensor<128xi32> + scf.yield %15, %16, %17 : tensor<128xi32>, tensor<128x!tt.ptr>, tensor<128x!tt.ptr> + } + tt.return +} + +// CHECK-LABEL: tt.func public @test_kernel( +// CHECK-SAME: %[[VAL_0:.*]]: !tt.ptr {tt.divisibility = 16 : i32}, %[[VAL_1:.*]]: !tt.ptr {tt.divisibility = 16 : i32}, %[[VAL_2:.*]]: !tt.ptr {tt.divisibility = 16 : i32}, %[[VAL_3:.*]]: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { +// CHECK: %[[VAL_4:.*]] = arith.constant dense<128> : tensor<128xi64> +// CHECK: %[[VAL_8:.*]] = arith.constant 3 : i32 +// CHECK: %[[VAL_15:.*]] = tt.get_program_id x : i32 +// CHECK: %[[VAL_16:.*]] = arith.muli %[[VAL_15]], %{{.*}} : i32 +// CHECK: %[[VAL_17:.*]] = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32> +// CHECK: %[[VAL_18:.*]] = tt.splat %[[VAL_16]] : i32 -> tensor<128xi32> +// CHECK: %[[VAL_19:.*]] = arith.addi %[[VAL_18]], %[[VAL_17]] : tensor<128xi32> +// CHECK: %[[VAL_20:.*]] = tt.splat %[[VAL_0]] : !tt.ptr -> tensor<128x!tt.ptr> +// CHECK: %[[VAL_21:.*]] = arith.extsi %[[VAL_19]] : tensor<128xi32> to tensor<128xi64> +// CHECK: %[[VAL_22:.*]] = tt.splat %[[VAL_2]] : !tt.ptr -> tensor<128x!tt.ptr> +// CHECK: %[[VAL_23:.*]] = tt.addptr %[[VAL_22]], %[[VAL_19]] : tensor<128x!tt.ptr>, tensor<128xi32> +// CHECK: %[[VAL_24:.*]] = tt.load %[[VAL_23]] : tensor<128x!tt.ptr> +// CHECK: %[[VAL_25:.*]] = arith.extsi %[[VAL_24]] : tensor<128xi32> to tensor<128xi64> +// CHECK: %[[VAL_26:.*]]:3 = scf.for %[[VAL_27:.*]] = %{{.*}} to %[[VAL_8]] step %{{.*}} iter_args(%[[VAL_28:.*]] = %[[VAL_19]], %[[VAL_29:.*]] = %[[VAL_21]], %[[VAL_30:.*]] = %[[VAL_25]]) -> (tensor<128xi32>, tensor<128xi64>, tensor<128xi64>) : i32 { +// CHECK: %[[VAL_31:.*]]:3 = scf.for %[[VAL_32:.*]] = %{{.*}} to %[[VAL_8]] step %{{.*}} iter_args(%[[VAL_33:.*]] = %[[VAL_28]], %[[VAL_34:.*]] = %[[VAL_29]], %[[VAL_35:.*]] = %[[VAL_30]]) -> (tensor<128xi32>, tensor<128xi64>, tensor<128xi64>) : i32 { +// CHECK: %[[VAL_36:.*]] = tt.splat %[[VAL_1]] : !tt.ptr -> tensor<128x!tt.ptr> +// CHECK: %[[VAL_37:.*]] = tt.addptr %[[VAL_36]], %[[VAL_34]] : tensor<128x!tt.ptr>, tensor<128xi64> +// CHECK: %[[VAL_38:.*]] = arith.cmpi slt, %[[VAL_33]], %{{.*}} : tensor<128xi32> +// CHECK: %[[VAL_39:.*]] = tt.addptr %[[VAL_20]], %[[VAL_33]] : tensor<128x!tt.ptr>, tensor<128xi32> +// CHECK: %[[VAL_40:.*]] = tt.load %[[VAL_39]], %[[VAL_38]], %{{.*}} : tensor<128x!tt.ptr> +// CHECK: %[[VAL_41:.*]] = tensor.empty() : tensor<128xi32> +// CHECK: %[[VAL_42:.*]] = scf.for %[[VAL_43:.*]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[VAL_44:.*]] = %[[VAL_41]]) -> (tensor<128xi32>) { +// CHECK: %[[VAL_45:.*]] = tensor.extract %[[VAL_35]]{{\[}}%[[VAL_43]]] {DiscreteMemAccess} : tensor<128xi64> +// CHECK: %[[VAL_46:.*]] = tt.addptr %[[VAL_3]], %[[VAL_45]] : !tt.ptr, i64 +// CHECK: %[[VAL_47:.*]] = tt.load %[[VAL_46]] {DiscreteMemAccess} : !tt.ptr +// CHECK: scf.yield {DiscreteMemAccess} %{{.*}} : tensor<128xi32> +// CHECK: } {ExtractedLoadOrStore} +// CHECK: %[[VAL_50:.*]] = arith.addi %[[VAL_40]], %[[VAL_42]] : tensor<128xi32> +// CHECK: tt.store %[[VAL_37]], %[[VAL_50]], %[[VAL_38]] : tensor<128x!tt.ptr> +// CHECK: %[[VAL_51:.*]] = arith.addi %[[VAL_33]], %{{.*}} : tensor<128xi32> +// CHECK: %[[VAL_52:.*]] = arith.addi %[[VAL_34]], %[[VAL_4]] : tensor<128xi64> +// CHECK: %[[VAL_53:.*]] = arith.addi %[[VAL_35]], %[[VAL_4]] : tensor<128xi64> +// CHECK: scf.yield %[[VAL_51]], %[[VAL_52]], %[[VAL_53]] : tensor<128xi32>, tensor<128xi64>, tensor<128xi64> +// CHECK: } +// CHECK: %[[VAL_54:.*]] = arith.addi %[[VAL_55:.*]]#0, %{{.*}} : tensor<128xi32> +// CHECK: %[[VAL_56:.*]] = arith.addi %[[VAL_55]]#1, %[[VAL_4]] : tensor<128xi64> +// CHECK: %[[VAL_57:.*]] = arith.addi %[[VAL_55]]#2, %[[VAL_4]] : tensor<128xi64> +// CHECK: scf.yield %[[VAL_54]], %[[VAL_56]], %[[VAL_57]] : tensor<128xi32>, tensor<128xi64>, tensor<128xi64> +// CHECK: } +// CHECK: tt.return +// CHECK: } + +// ----- + +tt.func public @test_kernel2(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: !tt.ptr {tt.divisibility = 16 : i32}, %arg3: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %c300_i32 = arith.constant 300 : i32 + %c3_i32 = arith.constant 3 : i32 + %c1_i32 = arith.constant 1 : i32 + %c0_i32 = arith.constant 0 : i32 + %cst = arith.constant dense<128> : tensor<128xi32> + %cst_0 = arith.constant dense<0> : tensor<128xi32> + %cst_1 = arith.constant dense<300> : tensor<128xi32> + %c128_i32 = arith.constant 128 : i32 + %0 = tt.get_program_id x : i32 + %1 = arith.muli %0, %c128_i32 : i32 + %2 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32> + %3 = tt.splat %1 : i32 -> tensor<128xi32> + %4 = arith.addi %3, %2 : tensor<128xi32> + %5 = tt.splat %arg0 : !tt.ptr -> tensor<128x!tt.ptr> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<128x!tt.ptr> + %7 = tt.addptr %6, %4 : tensor<128x!tt.ptr>, tensor<128xi32> + %8 = tt.splat %arg2 : !tt.ptr -> tensor<128x!tt.ptr> + %9 = tt.addptr %8, %4 : tensor<128x!tt.ptr>, tensor<128xi32> + %10 = tt.load %9 : tensor<128x!tt.ptr> + %11 = tt.splat %arg3 : !tt.ptr -> tensor<128x!tt.ptr> + %12 = tt.addptr %11, %10 : tensor<128x!tt.ptr>, tensor<128xi32> + %13:3 = scf.while (%arg4 = %cst, %arg5 = %4, %arg6 = %7, %arg7 = %12) : (tensor<128xi32>, tensor<128xi32>, tensor<128x!tt.ptr>, tensor<128x!tt.ptr>) -> (tensor<128x!tt.ptr>, tensor<128xi32>, tensor<128x!tt.ptr>) { + %14 = "tt.reduce"(%arg4) <{axis = 0 : i32}> ({ + ^bb0(%arg8: i32, %arg9: i32): + %16 = arith.addi %arg8, %arg9 : i32 + tt.reduce.return %16 : i32 + }) : (tensor<128xi32>) -> i32 + %15 = arith.cmpi slt, %14, %c300_i32 : i32 + scf.condition(%15) %arg7, %arg5, %arg6 : tensor<128x!tt.ptr>, tensor<128xi32>, tensor<128x!tt.ptr> + } do { + ^bb0(%arg4: tensor<128x!tt.ptr>, %arg5: tensor<128xi32>, %arg6: tensor<128x!tt.ptr>): + %14:4 = scf.while (%arg7 = %c0_i32, %arg8 = %arg6, %arg9 = %arg4, %arg10 = %arg5) : (i32, tensor<128x!tt.ptr>, tensor<128x!tt.ptr>, tensor<128xi32>) -> (tensor<128xi32>, i32, tensor<128x!tt.ptr>, tensor<128x!tt.ptr>) { + %18 = arith.cmpi slt, %arg7, %c3_i32 : i32 + scf.condition(%18) %arg10, %arg7, %arg8, %arg9 : tensor<128xi32>, i32, tensor<128x!tt.ptr>, tensor<128x!tt.ptr> + } do { + ^bb0(%arg7: tensor<128xi32>, %arg8: i32, %arg9: tensor<128x!tt.ptr>, %arg10: tensor<128x!tt.ptr>): + %18 = arith.cmpi slt, %arg7, %cst_1 : tensor<128xi32> + %19 = tt.addptr %5, %arg7 : tensor<128x!tt.ptr>, tensor<128xi32> + %20 = tt.load %19, %18, %cst_0 : tensor<128x!tt.ptr> + %21 = tt.load %arg10 : tensor<128x!tt.ptr> + %22 = arith.addi %20, %21 : tensor<128xi32> + tt.store %arg9, %22, %18 : tensor<128x!tt.ptr> + %23 = arith.addi %arg7, %cst : tensor<128xi32> + %24 = arith.addi %arg8, %c1_i32 : i32 + %25 = tt.addptr %arg9, %cst : tensor<128x!tt.ptr>, tensor<128xi32> + %26 = tt.addptr %arg10, %cst : tensor<128x!tt.ptr>, tensor<128xi32> + scf.yield %24, %25, %26, %23 : i32, tensor<128x!tt.ptr>, tensor<128x!tt.ptr>, tensor<128xi32> + } + %15 = arith.addi %14#0, %cst : tensor<128xi32> + %16 = tt.addptr %14#2, %cst : tensor<128x!tt.ptr>, tensor<128xi32> + %17 = tt.addptr %14#3, %cst : tensor<128x!tt.ptr>, tensor<128xi32> + scf.yield %14#0, %15, %16, %17 : tensor<128xi32>, tensor<128xi32>, tensor<128x!tt.ptr>, tensor<128x!tt.ptr> + } + tt.return +} + +// CHECK-LABEL: tt.func public @test_kernel2( +// CHECK-SAME: %[[VAL_0:.*]]: !tt.ptr {tt.divisibility = 16 : i32}, %[[VAL_1:.*]]: !tt.ptr {tt.divisibility = 16 : i32}, %[[VAL_2:.*]]: !tt.ptr {tt.divisibility = 16 : i32}, %[[VAL_3:.*]]: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { +// CHECK: %[[VAL_4:.*]] = arith.constant dense<128> : tensor<128xi64> +// CHECK: %[[VAL_8:.*]] = arith.constant 300 : i32 +// CHECK: %[[VAL_16:.*]] = tt.get_program_id x : i32 +// CHECK: %[[VAL_17:.*]] = arith.muli %[[VAL_16]], %{{.*}} : i32 +// CHECK: %[[VAL_18:.*]] = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32> +// CHECK: %[[VAL_19:.*]] = tt.splat %[[VAL_17]] : i32 -> tensor<128xi32> +// CHECK: %[[VAL_20:.*]] = arith.addi %[[VAL_19]], %[[VAL_18]] : tensor<128xi32> +// CHECK: %[[VAL_21:.*]] = tt.splat %[[VAL_0]] : !tt.ptr -> tensor<128x!tt.ptr> +// CHECK: %[[VAL_22:.*]] = arith.extsi %[[VAL_20]] : tensor<128xi32> to tensor<128xi64> +// CHECK: %[[VAL_23:.*]] = tt.splat %[[VAL_2]] : !tt.ptr -> tensor<128x!tt.ptr> +// CHECK: %[[VAL_24:.*]] = tt.addptr %[[VAL_23]], %[[VAL_20]] : tensor<128x!tt.ptr>, tensor<128xi32> +// CHECK: %[[VAL_25:.*]] = tt.load %[[VAL_24]] : tensor<128x!tt.ptr> +// CHECK: %[[VAL_26:.*]] = arith.extsi %[[VAL_25]] : tensor<128xi32> to tensor<128xi64> +// CHECK: %[[VAL_27:.*]]:3 = scf.while (%[[VAL_28:.*]] = %{{.*}}, %[[VAL_29:.*]] = %[[VAL_20]], %[[VAL_30:.*]] = %[[VAL_22]], %[[VAL_31:.*]] = %[[VAL_26]]) : (tensor<128xi32>, tensor<128xi32>, tensor<128xi64>, tensor<128xi64>) -> (tensor<128xi64>, tensor<128xi32>, tensor<128xi64>) { +// CHECK: %[[VAL_32:.*]] = "tt.reduce"(%[[VAL_28]]) <{axis = 0 : i32}> ({ +// CHECK: ^bb0(%[[VAL_33:.*]]: i32, %[[VAL_34:.*]]: i32): +// CHECK: %[[VAL_35:.*]] = arith.addi %[[VAL_33]], %[[VAL_34]] : i32 +// CHECK: tt.reduce.return %[[VAL_35]] : i32 +// CHECK: }) : (tensor<128xi32>) -> i32 +// CHECK: %[[VAL_36:.*]] = arith.cmpi slt, %[[VAL_32]], %[[VAL_8]] : i32 +// CHECK: scf.condition(%[[VAL_36]]) %[[VAL_31]], %[[VAL_29]], %[[VAL_30]] : tensor<128xi64>, tensor<128xi32>, tensor<128xi64> +// CHECK: } do { +// CHECK: ^bb0(%[[VAL_37:.*]]: tensor<128xi64>, %[[VAL_38:.*]]: tensor<128xi32>, %[[VAL_39:.*]]: tensor<128xi64>): +// CHECK: %[[VAL_40:.*]]:4 = scf.while (%[[VAL_41:.*]] = %{{.*}}, %[[VAL_42:.*]] = %[[VAL_39]], %[[VAL_43:.*]] = %[[VAL_37]], %[[VAL_44:.*]] = %[[VAL_38]]) : (i32, tensor<128xi64>, tensor<128xi64>, tensor<128xi32>) -> (i32, tensor<128xi64>, tensor<128xi64>, tensor<128xi32>) { +// CHECK: %[[VAL_45:.*]] = arith.cmpi slt, %[[VAL_41]], %{{.*}} : i32 +// CHECK: scf.condition(%[[VAL_45]]) %[[VAL_41]], %[[VAL_42]], %[[VAL_43]], %[[VAL_44]] : i32, tensor<128xi64>, tensor<128xi64>, tensor<128xi32> +// CHECK: } do { +// CHECK: ^bb0(%[[VAL_46:.*]]: i32, %[[VAL_47:.*]]: tensor<128xi64>, %[[VAL_48:.*]]: tensor<128xi64>, %[[VAL_49:.*]]: tensor<128xi32>): +// CHECK: %[[VAL_50:.*]] = tt.splat %[[VAL_1]] : !tt.ptr -> tensor<128x!tt.ptr> +// CHECK: %[[VAL_51:.*]] = tt.addptr %[[VAL_50]], %[[VAL_47]] : tensor<128x!tt.ptr>, tensor<128xi64> +// CHECK: %[[VAL_52:.*]] = arith.cmpi slt, %[[VAL_49]], %{{.*}} : tensor<128xi32> +// CHECK: %[[VAL_53:.*]] = tt.addptr %[[VAL_21]], %[[VAL_49]] : tensor<128x!tt.ptr>, tensor<128xi32> +// CHECK: %[[VAL_54:.*]] = tt.load %[[VAL_53]], %[[VAL_52]], %{{.*}} : tensor<128x!tt.ptr> +// CHECK: %[[VAL_55:.*]] = tensor.empty() : tensor<128xi32> +// CHECK: %[[VAL_56:.*]] = scf.for %[[VAL_57:.*]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[VAL_58:.*]] = %[[VAL_55]]) -> (tensor<128xi32>) { +// CHECK: %[[VAL_59:.*]] = tensor.extract %[[VAL_48]]{{\[}}%[[VAL_57]]] {DiscreteMemAccess} : tensor<128xi64> +// CHECK: %[[VAL_60:.*]] = tt.addptr %[[VAL_3]], %[[VAL_59]] : !tt.ptr, i64 +// CHECK: %[[VAL_61:.*]] = tt.load %[[VAL_60]] {DiscreteMemAccess} : !tt.ptr +// CHECK: scf.yield {DiscreteMemAccess} %{{.*}} : tensor<128xi32> +// CHECK: } {ExtractedLoadOrStore} +// CHECK: %[[VAL_64:.*]] = arith.addi %[[VAL_54]], %[[VAL_56]] : tensor<128xi32> +// CHECK: tt.store %[[VAL_51]], %[[VAL_64]], %[[VAL_52]] : tensor<128x!tt.ptr> +// CHECK: %[[VAL_65:.*]] = arith.addi %[[VAL_49]], %{{.*}} : tensor<128xi32> +// CHECK: %[[VAL_66:.*]] = arith.addi %[[VAL_46]], %{{.*}} : i32 +// CHECK: %[[VAL_67:.*]] = arith.addi %[[VAL_47]], %[[VAL_4]] : tensor<128xi64> +// CHECK: %[[VAL_68:.*]] = arith.addi %[[VAL_48]], %[[VAL_4]] : tensor<128xi64> +// CHECK: scf.yield %[[VAL_66]], %[[VAL_67]], %[[VAL_68]], %[[VAL_65]] : i32, tensor<128xi64>, tensor<128xi64>, tensor<128xi32> +// CHECK: } +// CHECK: %[[VAL_69:.*]] = arith.addi %[[VAL_70:.*]]#3, %{{.*}} : tensor<128xi32> +// CHECK: %[[VAL_71:.*]] = arith.addi %[[VAL_70]]#1, %[[VAL_4]] : tensor<128xi64> +// CHECK: %[[VAL_72:.*]] = arith.addi %[[VAL_70]]#2, %[[VAL_4]] : tensor<128xi64> +// CHECK: scf.yield %[[VAL_70]]#3, %[[VAL_69]], %[[VAL_71]], %[[VAL_72]] : tensor<128xi32>, tensor<128xi32>, tensor<128xi64>, tensor<128xi64> +// CHECK: } +// CHECK: tt.return +// CHECK: } diff --git a/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/splat.mlir b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/splat.mlir new file mode 100644 index 0000000000..5ff35c244d --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/splat.mlir @@ -0,0 +1,14 @@ +// RUN: triton-opt %s --triton-to-unstructure | FileCheck %s + +// CHECK-LABEL: tt.func @test_unstructure_splatandloadscenario +// CHECK: %[[EXT:.*]] = tensor.extract %{{.*}}[%{{.*}}] {DiscreteMemAccess} : tensor<128x!tt.ptr> +// CHECK: %[[VAL1:.*]] = tt.load %[[EXT]] : !tt.ptr +// CHECK: %[[VAL2:.*]] = tt.splat %[[VAL1]] : f32 -> tensor<128xf32> +tt.func @test_unstructure_splatandloadscenario(%base: !tt.ptr) -> tensor<128xf32> { + %offset = arith.constant 10 : i64 + %offset_tensor = tt.splat %offset : i64 -> tensor<128xi64> + %base_tensor = tt.splat %base : !tt.ptr -> tensor<128x!tt.ptr> + %ptr = tt.addptr %base_tensor, %offset_tensor : tensor<128x!tt.ptr>, tensor<128xi64> + %val = tt.load %ptr : tensor<128x!tt.ptr> + tt.return %val : tensor<128xf32> +} diff --git a/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/unstructure_mix.mlir b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/unstructure_mix.mlir new file mode 100644 index 0000000000..4516711f08 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/TritonToUnstructure/unstructure_mix.mlir @@ -0,0 +1,82 @@ +// RUN: triton-opt --triton-to-unstructure %s | FileCheck %s + +tt.func public @indirect_mix_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: !tt.ptr {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %cst = arith.constant dense<16> : tensor<1x8xi32> + %c8_i32 = arith.constant 8 : i32 + %0 = tt.get_program_id x : i32 + %1 = arith.muli %0, %c8_i32 : i32 + %2 = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> + %3 = tt.splat %1 : i32 -> tensor<8xi32> + %4 = arith.addi %3, %2 : tensor<8xi32> + %5 = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> + %6 = tt.splat %arg1 : !tt.ptr -> tensor<16x!tt.ptr> + %7 = tt.addptr %6, %5 : tensor<16x!tt.ptr>, tensor<16xi32> + %8 = tt.load %7 : tensor<16x!tt.ptr> + %9 = tt.expand_dims %2 {axis = 0 : i32} : tensor<8xi32> -> tensor<1x8xi32> + %10 = tt.splat %arg3 : i32 -> tensor<1x8xi32> + %11 = arith.muli %9, %10 : tensor<1x8xi32> + %12 = tt.expand_dims %8 {axis = 1 : i32} : tensor<16xi64> -> tensor<16x1xi64> + %13 = arith.extsi %11 : tensor<1x8xi32> to tensor<1x8xi64> + %14 = tt.broadcast %13 : tensor<1x8xi64> -> tensor<16x8xi64> + %15 = tt.broadcast %12 : tensor<16x1xi64> -> tensor<16x8xi64> + %16 = arith.addi %14, %15 : tensor<16x8xi64> + %17 = tt.splat %arg2 : !tt.ptr -> tensor<16x8x!tt.ptr> + %18 = tt.addptr %17, %16 : tensor<16x8x!tt.ptr>, tensor<16x8xi64> + %19 = tt.load %18 : tensor<16x8x!tt.ptr> + %20 = math.exp %19 : tensor<16x8xf32> + %21 = tt.expand_dims %4 {axis = 0 : i32} : tensor<8xi32> -> tensor<1x8xi32> + %22 = arith.muli %21, %cst : tensor<1x8xi32> + %23 = tt.expand_dims %5 {axis = 1 : i32} : tensor<16xi32> -> tensor<16x1xi32> + %24 = tt.broadcast %22 : tensor<1x8xi32> -> tensor<16x8xi32> + %25 = tt.broadcast %23 : tensor<16x1xi32> -> tensor<16x8xi32> + %26 = arith.addi %24, %25 : tensor<16x8xi32> + %27 = tt.splat %arg0 : !tt.ptr -> tensor<16x8x!tt.ptr> + %28 = tt.addptr %27, %26 : tensor<16x8x!tt.ptr>, tensor<16x8xi32> + tt.store %28, %20 : tensor<16x8x!tt.ptr> + tt.return +} + +// CHECK-LABEL: tt.func public @indirect_mix_kernel( +// CHECK-SAME: %[[VAL_0:.*]]: !tt.ptr {tt.divisibility = 16 : i32}, %[[VAL_1:.*]]: !tt.ptr {tt.divisibility = 16 : i32}, %[[VAL_2:.*]]: !tt.ptr {tt.divisibility = 16 : i32}, %[[VAL_3:.*]]: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { +// CHECK: %[[VAL_4:.*]] = arith.constant 16 : index +// CHECK: %[[VAL_5:.*]] = arith.constant 1 : index +// CHECK: %[[VAL_6:.*]] = arith.constant 0 : index +// CHECK: %[[VAL_7:.*]] = arith.constant dense<16> : tensor<1x8xi32> +// CHECK: %[[VAL_9:.*]] = tt.get_program_id x : i32 +// CHECK: %[[VAL_10:.*]] = arith.muli %[[VAL_9]], %{{.*}} : i32 +// CHECK: %[[VAL_11:.*]] = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> +// CHECK: %[[VAL_12:.*]] = tt.splat %[[VAL_10]] : i32 -> tensor<8xi32> +// CHECK: %[[VAL_13:.*]] = arith.addi %[[VAL_12]], %[[VAL_11]] : tensor<8xi32> +// CHECK: %[[VAL_14:.*]] = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32> +// CHECK: %[[VAL_15:.*]] = tt.splat %[[VAL_1]] : !tt.ptr -> tensor<16x!tt.ptr> +// CHECK: %[[VAL_16:.*]] = tt.addptr %[[VAL_15]], %[[VAL_14]] : tensor<16x!tt.ptr>, tensor<16xi32> +// CHECK: %[[VAL_17:.*]] = tt.load %[[VAL_16]] : tensor<16x!tt.ptr> +// CHECK: %[[VAL_18:.*]] = tt.expand_dims %[[VAL_11]] {axis = 0 : i32} : tensor<8xi32> -> tensor<1x8xi32> +// CHECK: %[[VAL_19:.*]] = tt.splat %[[VAL_3]] : i32 -> tensor<1x8xi32> +// CHECK: %[[VAL_20:.*]] = arith.muli %[[VAL_18]], %[[VAL_19]] : tensor<1x8xi32> +// CHECK: %[[VAL_21:.*]] = tt.expand_dims %[[VAL_17]] {axis = 1 : i32} : tensor<16xi64> -> tensor<16x1xi64> +// CHECK: %[[VAL_22:.*]] = arith.extsi %[[VAL_20]] : tensor<1x8xi32> to tensor<1x8xi64> +// CHECK: %[[VAL_23:.*]] = tt.broadcast %[[VAL_22]] : tensor<1x8xi64> -> tensor<16x8xi64> +// CHECK: %[[VAL_24:.*]] = tt.broadcast %[[VAL_21]] : tensor<16x1xi64> -> tensor<16x8xi64> +// CHECK: %[[VAL_25:.*]] = arith.addi %[[VAL_23]], %[[VAL_24]] : tensor<16x8xi64> +// CHECK: %[[VAL_26:.*]] = tensor.empty() : tensor<16x8xf32> +// CHECK: %[[VAL_27:.*]] = scf.for %[[VAL_28:.*]] = %[[VAL_6]] to %[[VAL_4]] step %[[VAL_5]] iter_args(%[[VAL_29:.*]] = %[[VAL_26]]) -> (tensor<16x8xf32>) { +// CHECK: %[[VAL_30:.*]] = tensor.extract_slice %[[VAL_25]]{{\[}}%[[VAL_28]], 0] [1, 8] [1, 1] {DiscreteMemAccess} : tensor<16x8xi64> to tensor<1x8xi64> +// CHECK: %[[VAL_31:.*]] = tt.splat %[[VAL_2]] : !tt.ptr -> tensor<1x8x!tt.ptr> +// CHECK: %[[VAL_32:.*]] = tt.addptr %[[VAL_31]], %[[VAL_30]] : tensor<1x8x!tt.ptr>, tensor<1x8xi64> +// CHECK: %[[VAL_33:.*]] = tt.load %[[VAL_32]] {DiscreteMemAccess} : tensor<1x8x!tt.ptr> +// CHECK: %[[VAL_34:.*]] = tensor.insert_slice %[[VAL_33]] into %[[VAL_29]]{{\[}}%[[VAL_28]], 0] [1, 8] [1, 1] : tensor<1x8xf32> into tensor<16x8xf32> +// CHECK: scf.yield {DiscreteMemAccess} %[[VAL_34]] : tensor<16x8xf32> +// CHECK: } {ExtractedLoadOrStore} +// CHECK: %[[VAL_35:.*]] = math.exp %[[VAL_27]] : tensor<16x8xf32> +// CHECK: %[[VAL_36:.*]] = tt.expand_dims %[[VAL_13]] {axis = 0 : i32} : tensor<8xi32> -> tensor<1x8xi32> +// CHECK: %[[VAL_37:.*]] = arith.muli %[[VAL_36]], %[[VAL_7]] : tensor<1x8xi32> +// CHECK: %[[VAL_38:.*]] = tt.expand_dims %[[VAL_14]] {axis = 1 : i32} : tensor<16xi32> -> tensor<16x1xi32> +// CHECK: %[[VAL_39:.*]] = tt.broadcast %[[VAL_37]] : tensor<1x8xi32> -> tensor<16x8xi32> +// CHECK: %[[VAL_40:.*]] = tt.broadcast %[[VAL_38]] : tensor<16x1xi32> -> tensor<16x8xi32> +// CHECK: %[[VAL_41:.*]] = arith.addi %[[VAL_39]], %[[VAL_40]] : tensor<16x8xi32> +// CHECK: %[[VAL_42:.*]] = tt.splat %[[VAL_0]] : !tt.ptr -> tensor<16x8x!tt.ptr> +// CHECK: %[[VAL_43:.*]] = tt.addptr %[[VAL_42]], %[[VAL_41]] : tensor<16x8x!tt.ptr>, tensor<16x8xi32> +// CHECK: tt.store %[[VAL_43]], %[[VAL_35]] : tensor<16x8x!tt.ptr> +// CHECK: tt.return +// CHECK: } diff --git a/third_party/ascend/unittest/affine_map/affine_map.py b/third_party/ascend/unittest/affine_map/affine_map.py new file mode 100755 index 0000000000..5d8460aed6 --- /dev/null +++ b/third_party/ascend/unittest/affine_map/affine_map.py @@ -0,0 +1,40 @@ +#!/usr/bin/env python3 + +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir + + +def main(): + with ir.context() as ctx: + ir.load_dialects(ctx) + ascend_ir.load_dialects(ctx) + + d0 = ascend_ir.affine_expr.get_dim(0) + d1 = ascend_ir.affine_expr.get_dim(1) + c2 = ascend_ir.affine_expr.get_constant(2) + + expr = (d0 + c2) * d1 + print("expr:", expr) + print("expr pure affine:", expr.is_pure_affine()) + print("expr hashable:", hash(expr)) + + m0 = ascend_ir.affine_map.get_identity(2) + m1 = ascend_ir.affine_map.get(2, 0, [d1, d0]) + m2 = ascend_ir.affine_map.get(2, 0, [d0 + d1, d1]) + m3 = ascend_ir.affine_map.get_constant(7) + minor = ascend_ir.affine_map.get_minor_identity(3, 2) + + print("m0:", m0) + print("m1:", m1) + print("m2:", m2) + print("m1 inverse:", m1.inverse_permutation()) + print("m2 submap[1]:", m2.get_sub_map([1])) + print("m2 compose m1:", m2.compose(m1)) + print("m1 as dict:", m1.to_dict()) + print("m3 constant:", m3, "value=", m3.get_constant_result()) + print("minor identity:", minor) + print("m2 results:", [str(x) for x in m2.get_results()]) + + +if __name__ == "__main__": + main() diff --git a/third_party/ascend/unittest/affine_map/affine_map_buffer_type_demo.py b/third_party/ascend/unittest/affine_map/affine_map_buffer_type_demo.py new file mode 100755 index 0000000000..a3e248a454 --- /dev/null +++ b/third_party/ascend/unittest/affine_map/affine_map_buffer_type_demo.py @@ -0,0 +1,27 @@ +#!/usr/bin/env python3 + +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir + + +def main(): + with ir.context() as ctx: + ir.load_dialects(ctx) + ascend_ir.load_dialects(ctx) + + builder = ascend_ir.ascendnpu_ir_builder(ctx) + f32 = builder.get_float_ty() + ub_space = builder.get_target_attribute(ascend_ir.AddressSpace.UB) + + # Build a memref type using an explicit affine map layout. + transpose_map = ascend_ir.affine_map.get(2, 0, [1, 0]) + memref_ty = builder.get_buffer_ty_with_affine_map([8, 16], f32, transpose_map, ub_space) + map_attr = builder.get_affine_map_attr(transpose_map) + + print("affine map:", transpose_map) + print("affine map attr:", map_attr) + print("memref type:", memref_ty) + + +if __name__ == "__main__": + main() diff --git a/third_party/ascend/unittest/affine_map/affine_map_complex_expr_demo.py b/third_party/ascend/unittest/affine_map/affine_map_complex_expr_demo.py new file mode 100755 index 0000000000..4748e69a02 --- /dev/null +++ b/third_party/ascend/unittest/affine_map/affine_map_complex_expr_demo.py @@ -0,0 +1,36 @@ +#!/usr/bin/env python3 + +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir + + +def main(): + with ir.context() as ctx: + ir.load_dialects(ctx) + ascend_ir.load_dialects(ctx) + + d0 = ascend_ir.affine_expr.get_dim(0) + d1 = ascend_ir.affine_expr.get_dim(1) + s0 = ascend_ir.affine_expr.get_symbol(0) + c3 = ascend_ir.affine_expr.get_constant(3) + c4 = ascend_ir.affine_expr.get_constant(4) + + # Complex expressions with symbols and integer arithmetic. + tiled_row = (d0 + s0).floordiv(c4) + tiled_col = (d1 + c3).ceildiv(c4) + inner = (d0 + d1).mod(c4) + + map_a = ascend_ir.affine_map.get(2, 1, [tiled_row, tiled_col, inner]) + map_b = ascend_ir.affine_map.get(2, 0, [d1, d0]) + map_comp = map_a.compose(map_b) + + print("map_a:", map_a) + print("map_b:", map_b) + print("map_a composed with map_b:", map_comp) + print("map_a results:", [str(r) for r in map_a.get_results()]) + print("map_a submap [0, 2]:", map_a.get_sub_map([0, 2])) + print("map_a metadata:", map_a.to_dict()) + + +if __name__ == "__main__": + main() diff --git a/third_party/ascend/unittest/affine_map/affine_map_indexing_map_demo.py b/third_party/ascend/unittest/affine_map/affine_map_indexing_map_demo.py new file mode 100755 index 0000000000..d8be55d37e --- /dev/null +++ b/third_party/ascend/unittest/affine_map/affine_map_indexing_map_demo.py @@ -0,0 +1,33 @@ +#!/usr/bin/env python3 + +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir + + +def main(): + with ir.context() as ctx: + ir.load_dialects(ctx) + ascend_ir.load_dialects(ctx) + + builder = ascend_ir.ascendnpu_ir_builder(ctx) + + d0 = ascend_ir.affine_expr.get_dim(0) + d1 = ascend_ir.affine_expr.get_dim(1) + c8 = ascend_ir.affine_expr.get_constant(8) + + # Example indexing maps: transpose and a tiled/reduced projection. + map_in0 = ascend_ir.affine_map.get(2, 0, [d1, d0]) + map_in1 = ascend_ir.affine_map.get(2, 0, [d0, d1]) + map_out = ascend_ir.affine_map.get(2, 0, [d0.floordiv(c8), d1.mod(c8)]) + + indexing_map_attr = builder.get_affine_map_array_attr([map_in0, map_in1, map_out]) + print("indexing_map attr:", indexing_map_attr) + + ub_space = builder.get_target_attribute(ascend_ir.AddressSpace.UB) + f32 = builder.get_float_ty() + memref_ty = builder.get_buffer_ty_with_affine_map([16, 32], f32, map_in0, ub_space) + print("buffer type with map_in0:", memref_ty) + + +if __name__ == "__main__": + main() diff --git a/third_party/ascend/unittest/affine_map/affine_map_parse_demo.py b/third_party/ascend/unittest/affine_map/affine_map_parse_demo.py new file mode 100755 index 0000000000..9bb8e09667 --- /dev/null +++ b/third_party/ascend/unittest/affine_map/affine_map_parse_demo.py @@ -0,0 +1,29 @@ +#!/usr/bin/env python3 + +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir + + +def main(): + with ir.context() as ctx: + ir.load_dialects(ctx) + ascend_ir.load_dialects(ctx) + + identity_map = ascend_ir.affine_map.get_identity(2) + transpose_map = ascend_ir.affine_map.get(2, 0, [1, 0]) + + print("identity map:", identity_map) + print(" dims:", identity_map.get_num_dims()) + print(" symbols:", identity_map.get_num_symbols()) + print(" results:", identity_map.get_num_results()) + print(" is_identity:", identity_map.is_identity()) + print(" is_permutation:", identity_map.is_permutation()) + + print("transpose map:", transpose_map) + print(" is_identity:", transpose_map.is_identity()) + print(" is_permutation:", transpose_map.is_permutation()) + print(" as python object:", transpose_map.to_dict()) + + +if __name__ == "__main__": + main() diff --git a/third_party/ascend/unittest/autotune_ut/01-vector-add.py b/third_party/ascend/unittest/autotune_ut/01-vector-add.py new file mode 100644 index 0000000000..1219c1a2e3 --- /dev/null +++ b/third_party/ascend/unittest/autotune_ut/01-vector-add.py @@ -0,0 +1,86 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Vector Add +============= +""" + +import os + +import torch +import torch_npu +import triton +import triton.language as tl +import triton.backends.ascend.runtime +from triton.backends.ascend.testing import do_bench_npu + + +@triton.autotune(configs=[], key=["n_elements"]) +@triton.jit +def add_kernel(x_ptr, # *Pointer* to first input vector. + y_ptr, # *Pointer* to second input vector. + output_ptr, # *Pointer* to output vector. + n_elements, # Size of the vector. + BLOCK_SIZE: tl.constexpr, # Number of elements each program should process. + # NOTE: `constexpr` so it can be used as a shape value. + ): + # There are multiple 'programs' processing different data. We identify which program + # we are here: + pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. + # This program will process inputs that are offset from the initial data. + # For instance, if you had a vector of length 256 and block_size of 64, the programs + # would each access the elements [0:64, 64:128, 128:192, 192:256]. + # Note that offsets is a list of pointers: + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + # Create a mask to guard memory operations against out-of-bounds accesses. + mask = offsets < n_elements + # Load x and y from DRAM, masking out any extra elements in case the input is not a + # multiple of the block size. + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + # Write x + y back to DRAM. + tl.store(output_ptr + offsets, output, mask=mask) + + +def add_torch(x, y): + return x + y + + +def add_autotune(x, y): + output = torch.empty_like(x) + n_elements = output.numel() + add_kernel[lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), )](x, y, output, n_elements) + return output + + +def test_add(size: int): + x = torch.rand(size, device="npu") + y = torch.rand(size, device="npu") + + output_torch = add_torch(x, y) + output_triton = add_autotune(x, y) + assert torch.allclose(output_triton, output_torch) + print(f"Vector Add {size} PASSED!") + + +if __name__ == "__main__": + test_add(98432) diff --git a/third_party/ascend/unittest/autotune_ut/02-fused-softmax.py b/third_party/ascend/unittest/autotune_ut/02-fused-softmax.py new file mode 100644 index 0000000000..66f6c7a371 --- /dev/null +++ b/third_party/ascend/unittest/autotune_ut/02-fused-softmax.py @@ -0,0 +1,99 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Fused Softmax +============= +""" + +import os + +import torch +import torch_npu +import triton +import triton.language as tl +import triton.backends.ascend.runtime +from triton.backends.ascend.testing import do_bench_npu + + +@triton.autotune( + configs=[], + key=["n_rows", "n_cols"], +) +@triton.jit +def softmax_kernel( + output_ptr, + input_ptr, + input_row_stride, + output_row_stride, + n_rows, + n_cols, + BLOCK_SIZE: tl.constexpr, + XBLOCK: tl.constexpr, + XBLOCK_SUB: tl.constexpr, +): + # starting row of the program + row_start = tl.program_id(0) * XBLOCK + for row_idx in tl.range(0, XBLOCK, XBLOCK_SUB): + # The stride represents how much we need to increase the pointer to advance 1 row + row_offsets = row_start + row_idx + tl.arange(0, XBLOCK_SUB)[:, None] + col_offsets = tl.arange(0, BLOCK_SIZE)[None, :] + xmask = row_offsets < n_rows + ymask = col_offsets < n_cols + mask = xmask & ymask + input_ptrs = input_ptr + (row_offsets * input_row_stride + col_offsets) + # Load the row into SRAM, using a mask since BLOCK_SIZE may be > than n_cols + row = tl.load(input_ptrs, mask=mask, other=-float("inf")) + # Subtract maximum for numerical stability + row_minus_max = row - tl.max(row, axis=1).reshape(XBLOCK_SUB, 1).broadcast_to(XBLOCK_SUB, BLOCK_SIZE) + # Note that exponentiation in Triton is fast but approximate (i.e., think __expf in CUDA) + numerator = tl.exp(row_minus_max) + denominator = (tl.sum(numerator, axis=1).reshape(XBLOCK_SUB, 1).broadcast_to(XBLOCK_SUB, BLOCK_SIZE)) + softmax_output = numerator / denominator + # Write back output to DRAM + output_ptrs = output_ptr + (row_offsets * output_row_stride + col_offsets) + tl.store(output_ptrs, softmax_output, mask=mask) + + +def softmax_torch(x): + return torch.softmax(x, axis=-1) + + +def softmax_autotune(x): + n_rows, n_cols = x.shape + BLOCK_SIZE = n_cols + + # Allocate output + y = torch.empty_like(x) + # Create a number of persistent programs. + softmax_kernel[lambda meta: (triton.cdiv(n_rows, meta["XBLOCK"]), 1, 1)](y, x, x.stride(0), y.stride(0), n_rows, + n_cols, BLOCK_SIZE=BLOCK_SIZE) + return y + + +def test_softmax(shape, dtype): + x = torch.randn(shape, dtype=dtype, device="npu") + y_torch = softmax_torch(x) + y_triton = softmax_autotune(x) + assert torch.allclose(y_triton, y_torch) + print(f"Fused Softmax {shape} {dtype} PASSED!") + + +if __name__ == "__main__": + test_softmax((16896, 1024), torch.float32) diff --git a/third_party/ascend/unittest/autotune_ut/03-layer-norm.py b/third_party/ascend/unittest/autotune_ut/03-layer-norm.py new file mode 100644 index 0000000000..ab547cfa36 --- /dev/null +++ b/third_party/ascend/unittest/autotune_ut/03-layer-norm.py @@ -0,0 +1,133 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Layer Normalization +============= +""" + +import os + +import torch +import torch_npu +import triton +import triton.language as tl +import triton.backends.ascend.runtime +from triton.backends.ascend.testing import do_bench_npu + + +@triton.autotune( + configs=[], + key=["M", "N"], +) +@triton.jit +def _layer_norm_fwd_fused( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride, # how much to increase the pointer when moving by 1 row + N, + M, # number of columns in X + eps, # epsilon to avoid division by zero + XBLOCK_SIZE: tl.constexpr, + RBLOCK_SIZE: tl.constexpr, +): + # Map the program id to the row of X and Y it should compute. + row_begin = tl.program_id(0) * XBLOCK_SIZE + row_idx = row_begin + tl.arange(0, XBLOCK_SIZE) + row_mask = row_idx < M + row_offsets = row_idx[:, None] * stride + # Compute mean + _mean = tl.zeros((XBLOCK_SIZE, RBLOCK_SIZE), dtype=tl.float32) + for off in range(0, N, RBLOCK_SIZE): + col_idx = off + tl.arange(0, RBLOCK_SIZE) + col_mask = col_idx < N + mask = row_mask[:, None] & col_mask[None, :] + a = tl.load(X + row_offsets + col_idx[None, :], mask=mask, other=0.0).to(tl.float32) + _mean += a + mean = tl.sum(_mean, axis=1, keep_dims=True) / N + # Compute variance + _var = tl.zeros((XBLOCK_SIZE, RBLOCK_SIZE), dtype=tl.float32) + for off in range(0, N, RBLOCK_SIZE): + col_idx = off + tl.arange(0, RBLOCK_SIZE) + col_mask = col_idx < N + mask = row_mask[:, None] & col_mask[None, :] + x = tl.load(X + row_offsets + col_idx[None, :], mask=mask, other=0.0).to(tl.float32) + x = tl.where(mask, x - mean, 0.0) + _var += x * x + var = tl.sum(_var, axis=1, keep_dims=True) / N + rstd = 1 / tl.sqrt(var + eps) + # Write mean / rstd + tl.store(Mean + row_idx[:, None], mean, mask=row_mask[:, None]) + tl.store(Rstd + row_idx[:, None], rstd, mask=row_mask[:, None]) + # Normalize and apply linear transformation + for off in range(0, N, RBLOCK_SIZE): + col_idx = off + tl.arange(0, RBLOCK_SIZE) + col_mask = col_idx < N + mask = row_mask[:, None] & col_mask[None, :] + w = tl.load(W + col_idx, mask=col_mask).reshape((1, RBLOCK_SIZE)) + b = tl.load(B + col_idx, mask=col_mask).reshape((1, RBLOCK_SIZE)) + x = tl.load(X + row_offsets + col_idx[None, :], mask=mask, other=0.0).to(tl.float32) + x_hat = (x - mean) * rstd + y = x_hat * w + b + # Write output + tl.store(Y + row_offsets + col_idx[None, :], y, mask=mask) + + +def layer_norm_torch(args): + x, w_shape, weight, bias, eps, dtype = args + return torch.nn.functional.layer_norm(x, w_shape, weight, bias, eps).to(dtype) + + +def layer_norm_autotune(args): + x, weight, bias, eps = args + # allocate output + y = torch.empty_like(x) + # reshape input data into 2D tensor + x_arg = x.reshape(-1, x.shape[-1]) + M, N = x_arg.shape + mean = torch.empty((M, ), dtype=torch.float32, device=x.device) + rstd = torch.empty((M, ), dtype=torch.float32, device=x.device) + + # enqueue kernel + _layer_norm_fwd_fused[lambda meta: (triton.cdiv(M, meta["XBLOCK_SIZE"]), 1, 1)]( # + x_arg, y, weight, bias, mean, rstd, x_arg.stride(0), N, M, eps # + ) + return y + + +def test_layer_norm(shape, dtype, eps=1e-5): + M, N = shape + device = "npu" + x_shape = shape + w_shape = (x_shape[-1], ) + weight = torch.rand(w_shape, dtype=dtype, device=device) + bias = torch.rand(w_shape, dtype=dtype, device=device) + x = -2.3 + 0.5 * torch.randn(x_shape, dtype=dtype, device=device) + y_torch = layer_norm_torch((x, w_shape, weight, bias, eps, dtype)) + y_triton = layer_norm_autotune((x, weight, bias, eps)) + assert torch.allclose(y_triton, y_torch, atol=1e-2, rtol=0) + print(f"Layer Normalization {M},{N} {dtype} PASSED!") + + +if __name__ == "__main__": + test_layer_norm((128, 32), torch.float16) diff --git a/third_party/ascend/unittest/autotune_ut/04-libentry.py b/third_party/ascend/unittest/autotune_ut/04-libentry.py new file mode 100644 index 0000000000..e956ba034b --- /dev/null +++ b/third_party/ascend/unittest/autotune_ut/04-libentry.py @@ -0,0 +1,97 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Vector Add with Libentry +============= +""" + +import os + +import torch +import torch_npu +import triton +import triton.language as tl +import triton.backends.ascend.runtime +from triton.runtime.libentry import libentry +from triton.backends.ascend.testing import do_bench_npu + + +# NB: Inserting any other decorator between @triton.autotune and @triton.jit disables +# parallel compilation during autotuning. To enable parallel compilation, apply @triton.autotune +# directly around @triton.jit (i.e., nest autotune as the outermost decorator on the JIT-compiled function) +@triton.autotune( + configs=[ + triton.Config({'BLOCK_SIZE': 1 * 1024, 'multibuffer': True}), + triton.Config({'BLOCK_SIZE': 12 * 1024, 'multibuffer': True}), + triton.Config({'BLOCK_SIZE': 12 * 1024, 'multibuffer': False}), + triton.Config({'BLOCK_SIZE': 8 * 1024, 'multibuffer': True}), + ], key=["n_elements"]) +@libentry() +@triton.jit +def add_kernel(x_ptr, # *Pointer* to first input vector. + y_ptr, # *Pointer* to second input vector. + output_ptr, # *Pointer* to output vector. + n_elements, # Size of the vector. + BLOCK_SIZE: tl.constexpr, # Number of elements each program should process. + # NOTE: `constexpr` so it can be used as a shape value. + ): + # There are multiple 'programs' processing different data. We identify which program + # we are here: + pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. + # This program will process inputs that are offset from the initial data. + # For instance, if you had a vector of length 256 and block_size of 64, the programs + # would each access the elements [0:64, 64:128, 128:192, 192:256]. + # Note that offsets is a list of pointers: + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + # Create a mask to guard memory operations against out-of-bounds accesses. + mask = offsets < n_elements + # Load x and y from DRAM, masking out any extra elements in case the input is not a + # multiple of the block size. + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + # Write x + y back to DRAM. + tl.store(output_ptr + offsets, output, mask=mask) + + +def add_torch(x, y): + return x + y + + +def add_autotune(x, y): + output = torch.empty_like(x) + n_elements = output.numel() + add_kernel[lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), )](x, y, output, n_elements) + return output + + +def test_add(size: int): + x = torch.rand(size, device="npu") + y = torch.rand(size, device="npu") + + output_torch = add_torch(x, y) + output_triton = add_autotune(x, y) + assert torch.allclose(output_triton, output_torch) + print(f"Vector Add {size} with libentry PASSED!") + + +if __name__ == "__main__": + test_add(98432) diff --git a/third_party/ascend/unittest/autotune_ut/test_autotune_param_valid.py b/third_party/ascend/unittest/autotune_ut/test_autotune_param_valid.py new file mode 100644 index 0000000000..4ce48e6c19 --- /dev/null +++ b/third_party/ascend/unittest/autotune_ut/test_autotune_param_valid.py @@ -0,0 +1,163 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import os + +import pytest +import torch +import torch_npu +import triton +import triton.backends.ascend.runtime +import triton.language as tl + + +@triton.autotune( + configs=[], key={"x": "n_elements"}, hints={ + "split_params": {"x": "BLOCK_SIZE"}, + "tiling_params": {"x": "BLOCK_SIZE_SUB"}, + "low_dim_axes": ["x"], + "reduction_axes": [], + }) +@triton.jit +def add_kernel( + x_ptr, + y_ptr, + output_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, + BLOCK_SIZE_SUB: tl.constexpr, +): + offset = tl.program_id(0) * BLOCK_SIZE + loops1 = (BLOCK_SIZE + BLOCK_SIZE_SUB - 1) // BLOCK_SIZE_SUB + for loop in range(0, loops1): + x0 = offset + loop * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE_SUB) + mask = x0 < n_elements + x = tl.load(x_ptr + x0, mask) + y = tl.load(y_ptr + x0, mask) + output = x + y + tl.store(output_ptr + x0, output) + + +def add_torch(x, y): + return x + y + + +def add_autotune(x, y): + output = torch.empty_like(x) + n_elements = output.numel() + add_kernel[lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), )](x, y, output, n_elements) + return output + + +@pytest.mark.autotune +@pytest.mark.parametrize('size', [ + 2048, +]) +def test_add(size: int): + x = torch.rand(size, device="npu") + y = torch.rand(size, device="npu") + + output_torch = add_torch(x, y) + output_triton = add_autotune(x, y) + assert torch.allclose(output_triton, output_torch) + + +@pytest.mark.autotune +def test_add_no_reduction_axes(): + try: + + @triton.autotune( + configs=[], key={"x": "n_elements"}, hints={ + "split_params": {"x": "BLOCK_SIZE"}, + "tiling_params": {"x": "BLOCK_SIZE_SUB"}, + "low_dim_axes": ["x"], + }) + @triton.jit + def add_kernel_exception(): + pass + except ValueError as e: + assert "reduction_axes must be a list" in str(e) + + +@pytest.mark.autotune +def test_add_no_low_dim_axes(): + try: + + @triton.autotune( + configs=[], key={"x": "n_elements"}, hints={ + "split_params": {"x": "BLOCK_SIZE"}, + "tiling_params": {"x": "BLOCK_SIZE_SUB"}, + "reduction_axes": [], + }) + @triton.jit + def add_kernel_exception(): + pass + except ValueError as e: + assert "low_dim_axes must be a list" in str(e) + + +@pytest.mark.autotune +def test_add_no_tiling_params(): + try: + + @triton.autotune(configs=[], key={"x": "n_elements"}, hints={ + "split_params": {"x": "BLOCK_SIZE"}, + "low_dim_axes": ["x"], + "reduction_axes": [], + }) + @triton.jit + def add_kernel_exception(): + pass + except ValueError as e: + assert "tiling_params must be a dict" in str(e) + + +@pytest.mark.autotune +def test_add_no_split_params(): + try: + + @triton.autotune( + configs=[], key={"x": "n_elements"}, hints={ + "tiling_params": {"x": "BLOCK_SIZE_SUB"}, + "low_dim_axes": ["x"], + "reduction_axes": [], + }) + @triton.jit + def add_kernel_exception(): + pass + except ValueError as e: + assert "split_params must be a dict" in str(e) + + +@pytest.mark.autotune +def test_add_no_keyname(): + try: + + @triton.autotune( + configs=[], key={"x0": "n_elements"}, hints={ + "tiling_params": {"x": "BLOCK_SIZE_SUB"}, + "low_dim_axes": ["x"], + "reduction_axes": [], + }) + @triton.jit + def add_kernel_exception(): + pass + except ValueError as e: + assert "All keys in 'key' must be valid axis names" in str(e) diff --git a/third_party/ascend/unittest/autotune_ut/test_common.py b/third_party/ascend/unittest/autotune_ut/test_common.py new file mode 100644 index 0000000000..b02f466334 --- /dev/null +++ b/third_party/ascend/unittest/autotune_ut/test_common.py @@ -0,0 +1,103 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import unittest.mock as mock +import pytest +import torch + + +def MockAutoTilingTunerRun(self, *args, **kwargs): + self.nargs = dict(zip(self.arg_names, args)) + + # generate key + all_args = {**self.nargs, **kwargs} + try: + self._autoparse_axis_params(all_args) + except ValueError as e: + if "Missing required arguments" in str(e): + pass + else: + raise + return { + "keys": self.keys, + "split_params": self.split_params, + "tiling_params": self.tiling_params, + "low_dim_axes": self.low_dim_axes, + "reduction_axes": self.reduction_axes, + "persistent_reduction": self.persistent_reduction, + } + + +def check_axes_parse_res(act: dict, ref: dict): + """ + Compare two axes parse results that may use different symbolic axis names, + but map to the same semantic dimensions via the 'keys' field. + """ + ref_keys = ref["keys"] + act_keys = act["keys"] + + assert set(ref_keys.values()) == set(act_keys.values()), \ + f"Semantic dimensions mismatch: ref={set(ref_keys.values())}, act={set(act_keys.values())}" + + def normalize_param_dict(param_dict: dict, sym_to_sem: dict) -> dict: + """Convert {symbol: value} -> {semantic: value}""" + return {sym_to_sem[sym]: value for sym, value in param_dict.items()} + + ref_split = normalize_param_dict(ref["split_params"], ref_keys) + act_split = normalize_param_dict(act["split_params"], act_keys) + + ref_tiling = normalize_param_dict(ref["tiling_params"], ref_keys) + act_tiling = normalize_param_dict(act["tiling_params"], act_keys) + + def normalize_axis_list(axis_list: list, sym_to_sem: dict) -> list: + return sorted(sym_to_sem[sym] for sym in axis_list) + + ref_low = normalize_axis_list(ref["low_dim_axes"], ref_keys) + act_low = normalize_axis_list(act["low_dim_axes"], act_keys) + + ref_red = normalize_axis_list(ref["reduction_axes"], ref_keys) + act_red = normalize_axis_list(act["reduction_axes"], act_keys) + + # Compare normalized structures + assert ref_split == act_split, f"split_params mismatch: {ref_split} vs {act_split}" + assert ref_tiling == act_tiling, f"tiling_params mismatch: {ref_tiling} vs {act_tiling}" + assert ref_low == act_low, f"low_dim_axes mismatch: {ref_low} vs {act_low}" + assert ref_red == act_red, f"reduction_axes mismatch: {ref_red} vs {act_red}" + + +@pytest.fixture +def mock_autotuner(): + with mock.patch("triton.backends.ascend.runtime.autotuner.AutoTilingTuner.run", new=MockAutoTilingTunerRun): + yield + + +def generate_tensor(shape, dtype): + if dtype == 'float32' or dtype == 'float16' or dtype == 'bfloat16': + return torch.randn(size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16': + return torch.randint(low=0, high=2000, size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'int8': + return torch.randint(low=0, high=127, size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'bool': + return torch.randint(low=0, high=2, size=shape).bool() + elif dtype == 'uint8': + return torch.randint(low=0, high=255, size=shape, dtype=torch.uint8) + else: + raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) diff --git a/third_party/ascend/unittest/autotune_ut/test_customized_config.py b/third_party/ascend/unittest/autotune_ut/test_customized_config.py new file mode 100644 index 0000000000..3431109569 --- /dev/null +++ b/third_party/ascend/unittest/autotune_ut/test_customized_config.py @@ -0,0 +1,91 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import os + +import pytest +import torch +import torch_npu +import triton +import triton.backends.ascend.runtime +import triton.language as tl + +os.environ['TRITON_PRINT_AUTOTUNING'] = '0' + + +@triton.autotune( + configs=[ + triton.Config({'XBLOCK': 128, 'XBLOCK_SUB': 32}), + triton.Config({'XBLOCK': 128, 'XBLOCK_SUB': 64}), + triton.Config({'XBLOCK': 396, 'XBLOCK_SUB': 6}), + ], key=["n_rows", "n_cols"], hints={ + "auto_gen_config": False, + }) +@triton.jit +def softmax_kernel( + output_ptr, + input_ptr, + input_row_stride, + output_row_stride, + n_rows, + n_cols, + BLOCK_SIZE: tl.constexpr, + XBLOCK: tl.constexpr, + XBLOCK_SUB: tl.constexpr, +): + row_start = tl.program_id(0) * XBLOCK + for row_idx in tl.range(0, XBLOCK, XBLOCK_SUB): + row_offsets = row_start + row_idx + tl.arange(0, XBLOCK_SUB)[:, None] + col_offsets = tl.arange(0, BLOCK_SIZE)[None, :] + xmask = row_offsets < n_rows + ymask = col_offsets < n_cols + mask = xmask & ymask + input_ptrs = input_ptr + (row_offsets * input_row_stride + col_offsets) + row = tl.load(input_ptrs, mask=mask, other=-float("inf")) + row_minus_max = row - tl.max(row, axis=1).reshape(XBLOCK_SUB, 1).broadcast_to(XBLOCK_SUB, BLOCK_SIZE) + numerator = tl.exp(row_minus_max) + denominator = (tl.sum(numerator, axis=1).reshape(XBLOCK_SUB, 1).broadcast_to(XBLOCK_SUB, BLOCK_SIZE)) + softmax_output = numerator / denominator + output_ptrs = output_ptr + (row_offsets * output_row_stride + col_offsets) + tl.store(output_ptrs, softmax_output, mask=mask) + + +def softmax_torch(x): + return torch.softmax(x, axis=-1) + + +def softmax_autotune(x): + n_rows, n_cols = x.shape + BLOCK_SIZE = n_cols + y = torch.empty_like(x) + softmax_kernel[lambda meta: (triton.cdiv(n_rows, meta["XBLOCK"]), 1, 1)](y, x, x.stride(0), y.stride(0), n_rows, + n_cols, BLOCK_SIZE=BLOCK_SIZE) + return y + + +@pytest.mark.autotune +@pytest.mark.parametrize('shape,dtype', [ + ((16896, 1024), torch.float32), +]) +def test_softmax(shape, dtype): + x = torch.randn(shape, dtype=dtype, device="npu") + y_torch = softmax_torch(x) + y_triton = softmax_autotune(x) + torch.testing.assert_close(y_torch, y_triton, rtol=1e-03, atol=1e-03, equal_nan=True) diff --git a/third_party/ascend/unittest/autotune_ut/test_do_bench_compat.py b/third_party/ascend/unittest/autotune_ut/test_do_bench_compat.py new file mode 100644 index 0000000000..25012f8daa --- /dev/null +++ b/third_party/ascend/unittest/autotune_ut/test_do_bench_compat.py @@ -0,0 +1,172 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +from types import MethodType + +import pytest +import triton +from triton.runtime.autotuner import Config +from triton.backends.ascend.runtime.autotuner import AutoTilingTuner + + +def _make_tuner(do_bench): + tuner = object.__new__(AutoTilingTuner) + tuner.compile_parallel = False + tuner.do_bench = do_bench + tuner.user_defined_do_bench = True + + def _make_kernel_call(self, *args, config, **meta): + + def kernel_call(warmup): + return None + + return kernel_call + + tuner._make_kernel_call = MethodType(_make_kernel_call, tuner) + return tuner + + +def test_batch_bench_supports_do_bench_with_quantiles(): + record = {} + + def _do_bench(fn, quantiles): + record["quantiles"] = quantiles + fn() + return (1.0, 1.0, 1.0) + + tuner = _make_tuner(_do_bench) + cfg = Config({}) + + result = tuner._batch_bench(configs=[cfg]) + + assert result[cfg] == (1.0, 1.0, 1.0) + assert record["quantiles"] == (0.5, 0.2, 0.8) + + +def test_batch_bench_requires_do_bench_quantiles_parameter(): + + def _do_bench(fn): + fn() + return (2.0, 2.0, 2.0) + + tuner = _make_tuner(_do_bench) + cfg = Config({}) + + with pytest.raises(TypeError): + tuner._batch_bench(configs=[cfg]) + + +def test_batch_bench_npu_env_respects_user_do_bench(monkeypatch): + calls = {"do_bench": 0} + + def _do_bench(fn, quantiles): + calls["do_bench"] += 1 + fn() + return (3.0, 3.0, 3.0) + + def _unexpected_do_bench_npu(*args, **kwargs): + raise AssertionError("do_bench_npu should not be used when user do_bench is provided") + + tuner = _make_tuner(_do_bench) + cfg0 = Config({"ID": 0}) + cfg1 = Config({"ID": 1}) + monkeypatch.setenv("TRITON_BENCH_METHOD", "npu") + monkeypatch.setattr("triton.backends.ascend.testing.do_bench_npu", _unexpected_do_bench_npu) + + result = tuner._batch_bench(configs=[cfg0, cfg1]) + + assert calls["do_bench"] == 2 + assert result[cfg0] == (3.0, 3.0, 3.0) + assert result[cfg1] == (3.0, 3.0, 3.0) + + +def test_batch_bench_npu_env_uses_do_bench_npu_without_user_do_bench(monkeypatch): + + def _do_bench(fn, quantiles): + raise AssertionError("self.do_bench should not be used when no user do_bench is provided") + + calls = {"do_bench_npu": 0} + + def _do_bench_npu(funcs, clear_l2_cache=False): + calls["do_bench_npu"] += 1 + assert len(funcs) == 2 + return [1.0, 2.0] + + tuner = _make_tuner(_do_bench) + tuner.user_defined_do_bench = False + cfg0 = Config({"ID": 0}) + cfg1 = Config({"ID": 1}) + monkeypatch.setenv("TRITON_BENCH_METHOD", "npu") + monkeypatch.setattr("triton.backends.ascend.testing.do_bench_npu", _do_bench_npu) + + result = tuner._batch_bench(configs=[cfg0, cfg1]) + + assert calls["do_bench_npu"] == 1 + assert result[cfg0] == 1.0 + assert result[cfg1] == 2.0 + + +def test_autotilingtuner_marks_user_defined_do_bench(): + marker = {"called": False} + + def _do_bench(fn, quantiles): + marker["called"] = True + return (0.0, 0.0, 0.0) + + def _dummy_kernel(): + return None + + _dummy_kernel.arg_names = [] + + tuner = AutoTilingTuner( + _dummy_kernel, + [], + [Config({})], + [], + None, + None, + do_bench=_do_bench, + ) + + assert tuner.user_defined_do_bench is True + assert marker["called"] is False + + +def test_ascend_autotune_decorator_forwards_do_bench(monkeypatch): + import triton.backends.ascend.runtime.autotuner as ascend_autotuner + + captured = {} + + class DummyAutoTilingTuner: + + def __init__(self, *args, **kwargs): + captured["do_bench"] = kwargs.get("do_bench") + + monkeypatch.setattr(ascend_autotuner, "AutoTilingTuner", DummyAutoTilingTuner) + + def _dummy_kernel(): + return None + + _dummy_kernel.arg_names = [] + my_do_bench = lambda kernel_call, quantiles: (0.0, 0.0, 0.0) + + ascend_autotuner.autotune(configs=[object()], key=[], do_bench=my_do_bench)(_dummy_kernel) + + assert captured["do_bench"] is my_do_bench diff --git a/third_party/ascend/unittest/autotune_ut/test_low_dim_axes_parse.py b/third_party/ascend/unittest/autotune_ut/test_low_dim_axes_parse.py new file mode 100644 index 0000000000..7d95aa87c2 --- /dev/null +++ b/third_party/ascend/unittest/autotune_ut/test_low_dim_axes_parse.py @@ -0,0 +1,173 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import unittest.mock as mock + +import triton +import triton.language as tl +from test_common import check_axes_parse_res, mock_autotuner + + +def _mock_single_reduction_axis(self): + if "x" in self.keys and "rx" not in self.keys: + self.keys["rx"] = self.keys.pop("x") + return ["rx"] + + +def test_low_dim_axis_parse_base_case1(mock_autotuner): + import triton.backends.ascend.runtime + + @triton.autotune(configs=[], key=["n_elements"]) + @triton.jit + def triton_low_dim_axis_parse_base_case1(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE # <- Separate assignment + + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + + tl.store(output_ptr + offsets, output, mask=mask) + + ref_res = { + "keys": {"x": "n_elements"}, + "split_params": {"x": "BLOCK_SIZE"}, + "tiling_params": {}, + "low_dim_axes": ["x"], + "reduction_axes": [], + } + grid = lambda meta: (meta["BLOCK_SIZE"], ) + act_res = triton_low_dim_axis_parse_base_case1[grid]() + + check_axes_parse_res(act_res, ref_res) + + +def test_low_dim_axis_parse_empty_is_non_fatal(mock_autotuner): + import triton.backends.ascend.runtime + + @triton.autotune(configs=[], key=["n_elements"]) + @triton.jit + def triton_low_dim_axis_parse_empty_case(x_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + mask = block_start < n_elements + + _ = tl.load(x_ptr + block_start, mask=mask, other=0) + + ref_res = { + "keys": {"x": "n_elements"}, + "split_params": {"x": "BLOCK_SIZE"}, + "tiling_params": {}, + "low_dim_axes": [], + "reduction_axes": [], + } + grid = lambda meta: (meta["BLOCK_SIZE"], ) + act_res = triton_low_dim_axis_parse_empty_case[grid]() + + check_axes_parse_res(act_res, ref_res) + + +def test_low_dim_axis_parse_empty_no_persistent_reduction_index_error(mock_autotuner): + import triton.backends.ascend.runtime + from triton.backends.ascend.runtime.autotuner import AutoTilingTuner + + @triton.autotune(configs=[], key=["n_elements"]) + @triton.jit + def triton_low_dim_axis_parse_guard_case(x_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + mask = block_start < n_elements + + _ = tl.load(x_ptr + block_start, mask=mask, other=0) + + with mock.patch.object(AutoTilingTuner, "_autoparse_reduction_axes", return_value=["rx"]), \ + mock.patch.object(AutoTilingTuner, "_autoparse_low_dim_axes", return_value=[]), \ + mock.patch.object(AutoTilingTuner, "_autoparse_split_params", return_value={}), \ + mock.patch.object(AutoTilingTuner, "_autoparse_tiling_params", return_value={}): + act_res = triton_low_dim_axis_parse_guard_case[(1, )]() + + assert act_res["low_dim_axes"] == [] + assert act_res["reduction_axes"] == ["rx"] + + +def test_persistent_reduction_inner_axis_threshold(mock_autotuner): + import triton.backends.ascend.runtime + from triton.backends.ascend.runtime.autotuner import AutoTilingTuner + + @triton.autotune(configs=[], key=["n_elements"]) + @triton.jit + def triton_persistent_reduction_inner_axis_case(x_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + mask = block_start < n_elements + _ = tl.load(x_ptr + block_start, mask=mask, other=0) + + with mock.patch.object(AutoTilingTuner, "_autoparse_reduction_axes", new=_mock_single_reduction_axis), \ + mock.patch.object(AutoTilingTuner, "_autoparse_low_dim_axes", return_value=["rx"]), \ + mock.patch.object(AutoTilingTuner, "_autoparse_split_params", return_value={}), \ + mock.patch.object(AutoTilingTuner, "_autoparse_tiling_params", return_value={}): + act_res = triton_persistent_reduction_inner_axis_case[(1, )](None, 1024) + + assert act_res["persistent_reduction"] is True + + +def test_persistent_reduction_outer_axis_threshold(mock_autotuner): + import triton.backends.ascend.runtime + from triton.backends.ascend.runtime.autotuner import AutoTilingTuner + + @triton.autotune(configs=[], key=["n_elements"]) + @triton.jit + def triton_persistent_reduction_outer_axis_case(x_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + mask = block_start < n_elements + _ = tl.load(x_ptr + block_start, mask=mask, other=0) + + with mock.patch.object(AutoTilingTuner, "_autoparse_reduction_axes", new=_mock_single_reduction_axis), \ + mock.patch.object(AutoTilingTuner, "_autoparse_low_dim_axes", return_value=[]), \ + mock.patch.object(AutoTilingTuner, "_autoparse_split_params", return_value={}), \ + mock.patch.object(AutoTilingTuner, "_autoparse_tiling_params", return_value={}): + act_res = triton_persistent_reduction_outer_axis_case[(1, )](None, 64) + + assert act_res["persistent_reduction"] is True + + +def test_persistent_reduction_outer_axis_over_threshold(mock_autotuner): + import triton.backends.ascend.runtime + from triton.backends.ascend.runtime.autotuner import AutoTilingTuner + + @triton.autotune(configs=[], key=["n_elements"]) + @triton.jit + def triton_persistent_reduction_outer_axis_over_threshold_case(x_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + mask = block_start < n_elements + _ = tl.load(x_ptr + block_start, mask=mask, other=0) + + with mock.patch.object(AutoTilingTuner, "_autoparse_reduction_axes", new=_mock_single_reduction_axis), \ + mock.patch.object(AutoTilingTuner, "_autoparse_low_dim_axes", return_value=[]), \ + mock.patch.object(AutoTilingTuner, "_autoparse_split_params", return_value={}), \ + mock.patch.object(AutoTilingTuner, "_autoparse_tiling_params", return_value={}): + act_res = triton_persistent_reduction_outer_axis_over_threshold_case[(1, )](None, 65) + + assert act_res["persistent_reduction"] is False diff --git a/third_party/ascend/unittest/autotune_ut/test_mask_parse.py b/third_party/ascend/unittest/autotune_ut/test_mask_parse.py new file mode 100644 index 0000000000..7a18f9f9f8 --- /dev/null +++ b/third_party/ascend/unittest/autotune_ut/test_mask_parse.py @@ -0,0 +1,162 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton.language as tl +from test_common import check_axes_parse_res, mock_autotuner + + +def test_triton_dot_case1(mock_autotuner): + """ + The current operator is only used for aixs analysis test cases. + CV fused operators do not support autotuning for now. + """ + import triton.backends.ascend.runtime + + @triton.autotune(configs=[], key=["M", "N", "K"]) + @triton.jit + def triton_dot_case1( + A, + B, + C, + M: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, + MBLOCK: tl.constexpr, + NBLOCK: tl.constexpr, + MBLOCK_SUB: tl.constexpr, + NBLOCK_SUB: tl.constexpr, + KBLOCK_SUB: tl.constexpr, + ): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + base_m = pid_m * MBLOCK + base_n = pid_n * NBLOCK + + loops_m = (MBLOCK + MBLOCK_SUB - 1) // MBLOCK_SUB + loops_n = (NBLOCK + NBLOCK_SUB - 1) // NBLOCK_SUB + loops_k = (K + KBLOCK_SUB - 1) // KBLOCK_SUB + + for loop_m in range(loops_m): + for loop_n in range(loops_n): + acc = tl.zeros((MBLOCK_SUB, NBLOCK_SUB), dtype=tl.float32) + + mdx = base_m + loop_m * MBLOCK_SUB + tl.arange(0, MBLOCK_SUB)[:, None] + ndx = base_n + loop_n * NBLOCK_SUB + tl.arange(0, NBLOCK_SUB)[None, :] + + for loop_k in range(loops_k): + kdx = loop_k * KBLOCK_SUB + tl.arange(0, KBLOCK_SUB) + kdx_m = kdx[None, :] # <- + A_ptr = A + mdx * K + kdx_m + a_mask = (mdx < M) & (kdx_m < K) # Use res of Subscript in mask compare + a = tl.load(A_ptr, mask=a_mask, other=0.0) + + kdx_n = kdx[:, None] + B_ptr = B + kdx_n * N + ndx + b_mask = (kdx_n < K) & (ndx < N) + b = tl.load(B_ptr, mask=b_mask, other=0.0) + + acc += tl.dot(a, b) + + C_ptr = C + mdx * N + ndx + c_mask = (mdx < M) & (ndx < N) + tl.store(C_ptr, acc, mask=c_mask) + + ref_res = { + "keys": {"x": "M", "y": "N", "z": "K"}, + "split_params": {"x": "MBLOCK", "y": "NBLOCK"}, + "tiling_params": {"x": "MBLOCK_SUB", "y": "NBLOCK_SUB", "z": "KBLOCK_SUB"}, + "low_dim_axes": ["y", "z"], + "reduction_axes": [], + } + grid = lambda meta: (meta["MBLOCK"], meta["NBLOCK"]) + act_res = triton_dot_case1[grid]() + + check_axes_parse_res(act_res, ref_res) + + +@pytest.mark.skip +def test_triton_dot_case2(mock_autotuner): + """ + The current operator is only used for aixs analysis test cases. + CV fused operators do not support autotuning for now. + """ + import triton.backends.ascend.runtime + + @triton.autotune(configs=[], key=["M", "N", "K"]) + @triton.jit + def triton_dot_case2( + A, + B, + C, + M: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, + MBLOCK: tl.constexpr, + NBLOCK: tl.constexpr, + MBLOCK_SUB: tl.constexpr, + NBLOCK_SUB: tl.constexpr, + KBLOCK_SUB: tl.constexpr, + ): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + base_m = pid_m * MBLOCK + base_n = pid_n * NBLOCK + + loops_m = (MBLOCK + MBLOCK_SUB - 1) // MBLOCK_SUB + loops_n = (NBLOCK + NBLOCK_SUB - 1) // NBLOCK_SUB + loops_k = (K + KBLOCK_SUB - 1) // KBLOCK_SUB + + for loop_m in range(loops_m): + for loop_n in range(loops_n): + acc = tl.zeros((MBLOCK_SUB, NBLOCK_SUB), dtype=tl.float32) + + mdx = base_m + loop_m * MBLOCK_SUB + tl.arange(0, MBLOCK_SUB)[:, None] + ndx = base_n + loop_n * NBLOCK_SUB + tl.arange(0, NBLOCK_SUB)[None, :] + + for loop_k in range(loops_k): + kdx = loop_k * KBLOCK_SUB + tl.arange(0, KBLOCK_SUB) + A_ptr = A + mdx * K + kdx[None, :] # <- + a_mask = (mdx < M) & (kdx[None, :] < K) # Cal subsript directly in mask compare + a = tl.load(A_ptr, mask=a_mask, other=0.0) + + B_ptr = B + kdx[:, None] * N + ndx + b_mask = (kdx[:, None] < K) & (ndx < N) + b = tl.load(B_ptr, mask=b_mask, other=0.0) + + acc += tl.dot(a, b) + + C_ptr = C + mdx * N + ndx + c_mask = (mdx < M) & (ndx < N) + tl.store(C_ptr, acc, mask=c_mask) + + ref_res = { + "keys": {"x": "M", "y": "N", "z": "K"}, + "split_params": {"x": "MBLOCK", "y": "NBLOCK"}, + "tiling_params": {"x": "MBLOCK_SUB", "y": "NBLOCK_SUB", "z": "KBLOCK_SUB"}, + "low_dim_axes": ["y", "z"], + "reduction_axes": [], + } + grid = lambda meta: (meta["MBLOCK"], meta["NBLOCK"]) + act_res = triton_dot_case2[grid]() + + check_axes_parse_res(act_res, ref_res) diff --git a/third_party/ascend/unittest/autotune_ut/test_no_tiling_axis_parse.py b/third_party/ascend/unittest/autotune_ut/test_no_tiling_axis_parse.py new file mode 100644 index 0000000000..f708cef694 --- /dev/null +++ b/third_party/ascend/unittest/autotune_ut/test_no_tiling_axis_parse.py @@ -0,0 +1,99 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +import os +import shutil +import pytest +import torch +import torch_npu +import triton +import triton.backends.ascend.runtime +import triton.language as tl +from triton.tools.get_ascend_devices import is_compile_on_910_95 + +import test_common + +os.environ['TRITON_ALWAYS_COMPILE'] = '1' +os.environ['TRITON_AUTOTUNE_PARALLEL_COMPILE'] = '0' + + +def case_torch(x): + return torch.permute(x, (1, 0)) + + +@triton.autotune(configs=[], key=['xnumel', 'ynumel'], hints={ + "auto_gen_config": True, +}) +@triton.jit +def triton_permute_2d( + output_ptr, + x_ptr, + xnumel: tl.constexpr, + ynumel: tl.constexpr, + XBLOCK: tl.constexpr, + YBLOCK: tl.constexpr, +): + xpid = tl.program_id(0) + ypid = tl.program_id(1) + + x_off = xpid * XBLOCK + tl.arange(0, XBLOCK)[:, None] + y_off = ypid * YBLOCK + tl.arange(0, YBLOCK)[None, :] + mask = (x_off < xnumel) & (y_off < ynumel) + offs = y_off + x_off * ynumel + b = tl.load(x_ptr + offs, mask=mask) + ox_off = ypid * YBLOCK + tl.arange(0, YBLOCK)[:, None] + oy_off = xpid * XBLOCK + tl.arange(0, XBLOCK)[None, :] + o_mask = (ox_off < ynumel) & (oy_off < xnumel) + o_offs = oy_off + ox_off * xnumel + ret = tl.permute(b, (1, 0)) + tl.store(output_ptr + o_offs, ret, mask=o_mask) + + +def case_triton(x_cal, is_simt_only=False): + xnumel = x_cal.shape[0] + ynumel = x_cal.shape[1] + output = torch.randint(1, (ynumel, xnumel), dtype=x_cal.dtype, device=x_cal.device) + if is_simt_only: + (triton_permute_2d[lambda meta: (triton.cdiv(xnumel, meta['XBLOCK']), triton.cdiv(ynumel, meta['YBLOCK']), 1)]( + output, x_cal, xnumel, ynumel, force_simt_only=True)) + else: + (triton_permute_2d[lambda meta: + (triton.cdiv(xnumel, meta['XBLOCK']), triton.cdiv(ynumel, meta['YBLOCK']), 1)](output, x_cal, + xnumel, + ynumel)) + return output + + +@pytest.mark.parametrize('shape', [(1024, 32), (32, 8)]) +@pytest.mark.parametrize('dtype', ['bfloat16']) +def test_permute(shape, dtype): + x_cal = test_common.generate_tensor(shape, dtype).npu() + torch_output = case_torch(x_cal) + triton_output = case_triton(x_cal) + torch.testing.assert_close(torch_output, triton_output, rtol=1e-03, atol=1e-03, equal_nan=True) + + +@pytest.mark.skipif(not is_compile_on_910_95, reason="only support A5") +@pytest.mark.parametrize('shape', [(1024, 32)]) +@pytest.mark.parametrize('dtype', ['bfloat16']) +def test_permute_simt(shape, dtype): + x_cal = test_common.generate_tensor(shape, dtype).npu() + torch_output = case_torch(x_cal) + triton_output = case_triton(x_cal, True) + torch.testing.assert_close(torch_output, triton_output, rtol=1e-03, atol=1e-03, equal_nan=True) diff --git a/third_party/ascend/unittest/autotune_ut/test_reduction_axes_parse.py b/third_party/ascend/unittest/autotune_ut/test_reduction_axes_parse.py new file mode 100644 index 0000000000..e6f3baa442 --- /dev/null +++ b/third_party/ascend/unittest/autotune_ut/test_reduction_axes_parse.py @@ -0,0 +1,157 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton.language as tl +from test_common import check_axes_parse_res, mock_autotuner + + +def test_triton_max_last_dim_case1(mock_autotuner): + import triton.backends.ascend.runtime + + @triton.autotune(configs=[], key=["x0_numel", "r1_numel"]) + @triton.jit + def triton_max_last_dim1( + in_ptr0, + out_ptr0, + x0_numel, + r1_numel, + X0BLOCK: tl.constexpr, + X0BLOCK_SUB: tl.constexpr, + R1BLOCK_SUB: tl.constexpr, + ): + x0_offset = tl.program_id(0) * X0BLOCK + base_x0 = tl.arange(0, X0BLOCK_SUB) + loops_x0 = (X0BLOCK + X0BLOCK_SUB - 1) // X0BLOCK_SUB + base_r1 = tl.arange(0, R1BLOCK_SUB) + loops_r1 = (r1_numel + R1BLOCK_SUB - 1) // R1BLOCK_SUB + for loop_x0 in range(loops_x0): + x0 = x0_offset + (loop_x0 * X0BLOCK_SUB) + base_x0[:, None] + x0_mask = x0 < min(X0BLOCK + x0_offset, x0_numel) + block_val = tl.full([X0BLOCK_SUB, R1BLOCK_SUB], float("-inf"), tl.float32) + for loop_r1 in range(loops_r1): + r1 = (loop_r1 * R1BLOCK_SUB) + base_r1[None, :] + r1_mask = r1 < r1_numel + tmp = tl.load(in_ptr0 + (r1 + r1_numel * x0), r1_mask & x0_mask, other=float("-inf")) + block_val = tl.maximum(block_val, tmp) + # Reduce along axis = 1 (the last dimension in this 2D tensor) + block_res = tl.max(block_val, axis=1)[:, None] # <- explicit positive axis index + tl.store(out_ptr0 + x0, block_res, x0_mask) + + ref_res = { + "keys": {"x": "x0_numel", "ry": "r1_numel"}, + "split_params": {"x": "X0BLOCK"}, + "tiling_params": {"x": "X0BLOCK_SUB", "ry": "R1BLOCK_SUB"}, + "low_dim_axes": ["ry"], + "reduction_axes": ["ry"], + } + grid = lambda meta: (meta["X0BLOCK"], ) + act_res = triton_max_last_dim1[grid]() + + check_axes_parse_res(act_res, ref_res) + + +def test_triton_max_last_dim_case2(mock_autotuner): + import triton.backends.ascend.runtime + + @triton.autotune(configs=[], key=["x0_numel", "r1_numel"]) + @triton.jit + def triton_max_last_dim2( + in_ptr0, + out_ptr0, + x0_numel, + r1_numel, + X0BLOCK: tl.constexpr, + X0BLOCK_SUB: tl.constexpr, + R1BLOCK_SUB: tl.constexpr, + ): + x0_offset = tl.program_id(0) * X0BLOCK + base_x0 = tl.arange(0, X0BLOCK_SUB) + loops_x0 = (X0BLOCK + X0BLOCK_SUB - 1) // X0BLOCK_SUB + base_r1 = tl.arange(0, R1BLOCK_SUB) + loops_r1 = (r1_numel + R1BLOCK_SUB - 1) // R1BLOCK_SUB + for loop_x0 in range(loops_x0): + x0 = x0_offset + (loop_x0 * X0BLOCK_SUB) + base_x0[:, None] + x0_mask = x0 < min(X0BLOCK + x0_offset, x0_numel) + block_val = tl.full([X0BLOCK_SUB, R1BLOCK_SUB], float("-inf"), tl.float32) + for loop_r1 in range(loops_r1): + r1 = (loop_r1 * R1BLOCK_SUB) + base_r1[None, :] + r1_mask = r1 < r1_numel + tmp = tl.load(in_ptr0 + (r1 + r1_numel * x0), r1_mask & x0_mask, other=float("-inf")) + block_val = tl.maximum(block_val, tmp) + # Reduce along axis=-1 (the last dimension, equivalent to axis=1 in 2D) + block_res = tl.max(block_val, axis=-1)[:, None] # <- negative axis index (last dim) + tl.store(out_ptr0 + x0, block_res, x0_mask) + + ref_res = { + "keys": {"x": "x0_numel", "ry": "r1_numel"}, + "split_params": {"x": "X0BLOCK"}, + "tiling_params": {"x": "X0BLOCK_SUB", "ry": "R1BLOCK_SUB"}, + "low_dim_axes": ["ry"], + "reduction_axes": ["ry"], + } + grid = lambda meta: (meta["X0BLOCK"], ) + act_res = triton_max_last_dim2[grid]() + + check_axes_parse_res(act_res, ref_res) + + +def test_triton_max_last_dim_case3(mock_autotuner): + import triton.backends.ascend.runtime + + @triton.autotune(configs=[], key=["x0_numel", "r1_numel"]) + @triton.jit + def triton_max_last_dim3( + in_ptr0, + out_ptr0, + x0_numel, + r1_numel, + X0BLOCK: tl.constexpr, + X0BLOCK_SUB: tl.constexpr, + R1BLOCK_SUB: tl.constexpr, + ): + x0_offset = tl.program_id(0) * X0BLOCK + base_x0 = tl.arange(0, X0BLOCK_SUB) + loops_x0 = (X0BLOCK + X0BLOCK_SUB - 1) // X0BLOCK_SUB + base_r1 = tl.arange(0, R1BLOCK_SUB) + loops_r1 = (r1_numel + R1BLOCK_SUB - 1) // R1BLOCK_SUB + for loop_x0 in range(loops_x0): + x0 = x0_offset + (loop_x0 * X0BLOCK_SUB) + base_x0[:, None] + x0_mask = x0 < min(X0BLOCK + x0_offset, x0_numel) + block_val = tl.full([X0BLOCK_SUB, R1BLOCK_SUB], float("-inf"), tl.float32) + for loop_r1 in range(loops_r1): + r1 = (loop_r1 * R1BLOCK_SUB) + base_r1[None, :] + r1_mask = r1 < r1_numel + tmp = tl.load(in_ptr0 + (r1 + r1_numel * x0), r1_mask & x0_mask, other=float("-inf")) + block_val = tl.maximum(block_val, tmp) + # Reduce along axis=1, passed as a positional argument (not keyword `axis=...`) + block_res = tl.max(block_val, 1)[:, None] # <- explicit positive axis index + tl.store(out_ptr0 + x0, block_res, x0_mask) + + ref_res = { + "keys": {"x": "x0_numel", "ry": "r1_numel"}, + "split_params": {"x": "X0BLOCK"}, + "tiling_params": {"x": "X0BLOCK_SUB", "ry": "R1BLOCK_SUB"}, + "low_dim_axes": ["ry"], + "reduction_axes": ["ry"], + } + grid = lambda meta: (meta["X0BLOCK"], ) + act_res = triton_max_last_dim3[grid]() + + check_axes_parse_res(act_res, ref_res) diff --git a/third_party/ascend/unittest/autotune_ut/test_split_axis_parse.py b/third_party/ascend/unittest/autotune_ut/test_split_axis_parse.py new file mode 100644 index 0000000000..2aef37c4f6 --- /dev/null +++ b/third_party/ascend/unittest/autotune_ut/test_split_axis_parse.py @@ -0,0 +1,180 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import unittest.mock as mock + +import triton +import triton.language as tl + +from test_common import check_axes_parse_res, mock_autotuner + + +def test_split_axis_parse_base_case1(mock_autotuner): + import triton.backends.ascend.runtime + + @triton.autotune(configs=[], key=["n_elements"]) + @triton.jit + def triton_split_axis_parse_base_case1(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE # <- Separate assignment + + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + + tl.store(output_ptr + offsets, output, mask=mask) + + ref_res = { + "keys": {"x": "n_elements"}, + "split_params": {"x": "BLOCK_SIZE"}, + "tiling_params": {}, + "low_dim_axes": ["x"], + "reduction_axes": [], + } + grid = lambda meta: (meta["BLOCK_SIZE"], ) + act_res = triton_split_axis_parse_base_case1[grid]() + + check_axes_parse_res(act_res, ref_res) + + +def test_split_axis_parse_base_case2(mock_autotuner): + import triton.backends.ascend.runtime + + @triton.autotune(configs=[], key=["n_elements"]) + @triton.jit + def triton_split_axis_parse_base_case2(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + block_start = tl.program_id(axis=0) * BLOCK_SIZE # <- Computed inline but still named + + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + + tl.store(output_ptr + offsets, output, mask=mask) + + ref_res = { + "keys": {"x": "n_elements"}, + "split_params": {"x": "BLOCK_SIZE"}, + "tiling_params": {}, + "low_dim_axes": ["x"], + "reduction_axes": [], + } + grid = lambda meta: (meta["BLOCK_SIZE"], ) + act_res = triton_split_axis_parse_base_case2[grid]() + + check_axes_parse_res(act_res, ref_res) + + +def test_split_axis_parse_base_case3(mock_autotuner): + import triton.backends.ascend.runtime + + @triton.autotune(configs=[], key=["n_elements"]) + @triton.jit + def triton_split_axis_parse_base_case3(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + offsets = tl.program_id(axis=0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) # <- Fully fused + mask = offsets < n_elements + + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + + tl.store(output_ptr + offsets, output, mask=mask) + + ref_res = { + "keys": {"x": "n_elements"}, + "split_params": {"x": "BLOCK_SIZE"}, + "tiling_params": {}, + "low_dim_axes": ["x"], + "reduction_axes": [], + } + grid = lambda meta: (meta["BLOCK_SIZE"], ) + act_res = triton_split_axis_parse_base_case3[grid]() + + check_axes_parse_res(act_res, ref_res) + + +def test_split_axis_parse_base_case4(mock_autotuner): + import triton.backends.ascend.runtime + + @triton.autotune(configs=[], key=["n_elements"]) + @triton.jit + def triton_split_axis_parse_base_case4(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + offsets = tl.program_id(axis=0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) # <- Fully fused + mask = offsets < n_elements + + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output_x = x[:, None].to(tl.float32) * 1 # left=Call(func=Attribute(value=Subscipt(...),),)) + output_y = 1 * y[None, :].to(tl.float32) + + output_offsets = tl.program_id(axis=0) * BLOCK_SIZE * BLOCK_SIZE + \ + tl.arange(0, BLOCK_SIZE)[:, None] * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)[None, :] + tl.store(output_ptr + output_offsets, output_x + output_y, mask=mask) + + ref_res = { + "keys": {"x": "n_elements"}, + "split_params": {"x": "BLOCK_SIZE"}, + "tiling_params": {}, + "low_dim_axes": ["x"], + "reduction_axes": [], + } + grid = lambda meta: (meta["BLOCK_SIZE"], ) + act_res = triton_split_axis_parse_base_case4[grid]() + + check_axes_parse_res(act_res, ref_res) + + +def test_grid_stride_loop_block_only_tiling_semantics(mock_autotuner): + import triton.backends.ascend.runtime + + @triton.autotune(configs=[], key=["N", "index_len"]) + @triton.jit + def triton_grid_stride_loop_block_only_tiling_semantics( + input_ptr, + output_ptr, + index_ptr, + N: tl.constexpr, + index_len: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + ): + pid_x = tl.program_id(axis=0) + pid_y = tl.program_id(axis=1) + grid_x = tl.num_programs(axis=0) + grid_y = tl.num_programs(axis=1) + for x in range(pid_x * BLOCK_M, index_len, grid_x * BLOCK_M): + row_offsets = x + tl.arange(0, BLOCK_M) + indices = tl.load(index_ptr + row_offsets, mask=row_offsets < index_len, other=0) + for y in range(pid_y * BLOCK_N, N, grid_y * BLOCK_N): + col_offsets = y + tl.arange(0, BLOCK_N) + col_mask = col_offsets < N + inp_offset = indices[:, None] * N + col_offsets[None, :] + out_offset = row_offsets[:, None] * N + col_offsets[None, :] + selected = tl.load(input_ptr + inp_offset, mask=col_mask[None, :], other=0.0) + tl.store(output_ptr + out_offset, selected, mask=col_mask[None, :]) + + act_res = triton_grid_stride_loop_block_only_tiling_semantics[(1, 1)]() + assert act_res["split_params"] == {} + assert act_res["tiling_params"] == {"y": "BLOCK_M", "x": "BLOCK_N"} diff --git a/third_party/ascend/unittest/autotune_ut/test_tiling_axis_parse.py b/third_party/ascend/unittest/autotune_ut/test_tiling_axis_parse.py new file mode 100644 index 0000000000..1c5e150932 --- /dev/null +++ b/third_party/ascend/unittest/autotune_ut/test_tiling_axis_parse.py @@ -0,0 +1,123 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +from test_common import check_axes_parse_res, mock_autotuner + + +def test_tiling_axis_parse_base_case1(mock_autotuner): + import triton.backends.ascend.runtime + + @triton.autotune(configs=[], key=["n_elements"]) + @triton.jit + def triton_tiling_axis_parse_base_case1(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr, + BLOCK_SUB: tl.constexpr): + offset = tl.program_id(axis=0) * BLOCK_SIZE + base = tl.arange(0, BLOCK_SUB) + loops = (BLOCK_SIZE + BLOCK_SUB - 1) // BLOCK_SUB # <- + for loop in range(loops): + offsets = offset + (loop * BLOCK_SUB) + base + mask = offsets < min(BLOCK_SIZE + offset, n_elements) + + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + + tl.store(output_ptr + offsets, output, mask=mask) + + ref_res = { + "keys": {"x": "n_elements"}, + "split_params": {"x": "BLOCK_SIZE"}, + "tiling_params": {"x": "BLOCK_SUB"}, + "low_dim_axes": ["x"], + "reduction_axes": [], + } + grid = lambda meta: (meta["BLOCK_SIZE"], ) + act_res = triton_tiling_axis_parse_base_case1[grid]() + + check_axes_parse_res(act_res, ref_res) + + +@pytest.mark.skip +def test_tiling_axis_parse_base_case2(mock_autotuner): + import triton.backends.ascend.runtime + + @triton.autotune(configs=[], key=["n_elements"]) + @triton.jit + def triton_tiling_axis_parse_base_case2(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr, + BLOCK_SUB: tl.constexpr): + offset = tl.program_id(axis=0) * BLOCK_SIZE + base = tl.arange(0, BLOCK_SUB) + for offset_sub in range(0, BLOCK_SIZE, BLOCK_SUB): + offsets = offset + offset_sub + base[:] # <- + mask = offsets < min(BLOCK_SIZE + offset, n_elements) + + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + + tl.store(output_ptr + offsets, output, mask=mask) + + ref_res = { + "keys": {"x": "n_elements"}, + "split_params": {"x": "BLOCK_SIZE"}, + "tiling_params": {"x": "BLOCK_SUB"}, + "low_dim_axes": ["x"], + "reduction_axes": [], + } + grid = lambda meta: (meta["BLOCK_SIZE"], ) + act_res = triton_tiling_axis_parse_base_case2[grid]() + + check_axes_parse_res(act_res, ref_res) + + +@pytest.mark.skip +def test_tiling_axis_parse_base_case3(mock_autotuner): + import triton.backends.ascend.runtime + + @triton.autotune(configs=[], key=["n_elements"]) + @triton.jit + def triton_tiling_axis_parse_base_case3(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr, + BLOCK_SUB: tl.constexpr): + offset = tl.program_id(axis=0) * BLOCK_SIZE + base = tl.arange(0, BLOCK_SUB)[:] # <- + for offset_sub in range(0, BLOCK_SIZE, BLOCK_SUB): + offsets = offset + offset_sub + base + mask = offsets < min(BLOCK_SIZE + offset, n_elements) + + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + + tl.store(output_ptr + offsets, output, mask=mask) + + ref_res = { + "keys": {"x": "n_elements"}, + "split_params": {"x": "BLOCK_SIZE"}, + "tiling_params": {"x": "BLOCK_SUB"}, + "low_dim_axes": ["x"], + "reduction_axes": [], + } + grid = lambda meta: (meta["BLOCK_SIZE"], ) + act_res = triton_tiling_axis_parse_base_case3[grid]() + + check_axes_parse_res(act_res, ref_res) diff --git a/third_party/ascend/unittest/conftest.py b/third_party/ascend/unittest/conftest.py new file mode 100644 index 0000000000..a284dcdec0 --- /dev/null +++ b/third_party/ascend/unittest/conftest.py @@ -0,0 +1,38 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import torch + + +@pytest.fixture(scope="session", autouse=True) +def assign_npu(worker_id): + npu_count = torch.npu.device_count() + if worker_id == "master": + npu_id = 0 + else: + idx = int(worker_id.replace("gw", "")) + npu_id = idx % npu_count + torch.npu.set_device(npu_id) + + +def pytest_addoption(parser): + parser.addoption("--kernel", action="append", default=None, + help="run only specified kernel(s); can be supplied multiple times") diff --git a/third_party/ascend/unittest/custom_op/builtin_ops_demo.py b/third_party/ascend/unittest/custom_op/builtin_ops_demo.py new file mode 100755 index 0000000000..e3e9fc6750 --- /dev/null +++ b/third_party/ascend/unittest/custom_op/builtin_ops_demo.py @@ -0,0 +1,52 @@ +#!/usr/bin/env python3 +import subprocess +import triton +import triton.language as tl +import triton.language.extra.cann.extension as al +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir +from triton.backends.ascend.compiler import NPUOptions, ttir_to_linalg + + +@triton.jit +def my_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + x = tl.load(x_ptr + i, mask=i < n) + y = tl.load(y_ptr + i, mask=i < n) + index = tl.full([8], 0, tl.int32) + value = tl.full([8, 64], 0, tl.float32) + tmp = tl.full([8], 0, tl.float32) + x = al.custom("__builtin_index_select", x_ptr, index, dim=0, bound=100, end_offset=(2, 2), start_offset=(0, 0), + src_stride=(4, 1), out=x) + al.custom("__builtin_index_put", x_ptr, index, value, dim=0, bound=12, dst_shape=(1, 2, 3), dst_offset=(4, 5, 6), + dst_stride=(8, 4, 1)) + tmp = al.custom("__builtin_gather_load", y_ptr, index, bound=100, dim=0, src_stride=(1, ), index_shape=(3, ), + offsets=(0, ), out=tmp) + al.custom("__builtin_scatter_store", out_ptr, value, index, 1, 0, (1, ), (2, ), (1, )) + y = al.custom("__builtin_indirect_load", x_ptr, index, mask=i < n, other=y, out=y) + al.custom("__builtin_indirect_store", out_ptr, index, value) + tl.store(out_ptr + i, y, mask=i < n) + + +if __name__ == "__main__": + src = ASTSource(my_kernel, {"x_ptr": "*fp32", "y_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, {"BLOCK": 256}) + context = ir.context() + ir.load_dialects(context) + ascend_ir.load_dialects(context) + options = NPUOptions() + try: + ttir = ast_to_ttir(my_kernel, src, context, options, {}, {}) + print("=== TTIR ===") + print(ttir) + metadata = { + **options.__dict__, + } + linalg = ttir_to_linalg(ttir, metadata, options, named_ops=True) + print("=== MLIR (linalg) ===") + print(linalg) + except subprocess.CalledProcessError as ex: + print(ex.stdout.decode()) + print(ex.stderr.decode()) + print("failed") diff --git a/third_party/ascend/unittest/custom_op/custom_op_demo.py b/third_party/ascend/unittest/custom_op/custom_op_demo.py new file mode 100755 index 0000000000..e3bbd5e412 --- /dev/null +++ b/third_party/ascend/unittest/custom_op/custom_op_demo.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +import subprocess +import os +import triton +import triton.language as tl +import triton.language.extra.cann.extension as al +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir +from triton.backends.ascend.compiler import NPUOptions, ttir_to_linalg + + +@al.register_custom_op +class min_custom_op: + core = al.CORE.VECTOR + pipe = al.PIPE.PIPE_MTE2 + mode = al.MODE.SIMD + + symbol = 'min_custom_op_impl' + bitcode = os.path.abspath(__file__) + + +@al.register_custom_op +class simple_custom_op: + # name is optional, use class name by default. + name = 'simple_custom_op' + + # required attributes. + core = al.CORE.VECTOR + pipe = al.PIPE.PIPE_V + mode = al.MODE.SIMT + + symbol = 'simple_custom_op_impl' + bitcode = os.path.abspath(__file__) + + # __init__ method is optional, but it can be used for better user experience + # when provided. for example, you can validate arguments here. + def __init__(self, x, y, dim=0, out=None): + assert x.shape == y.shape, "x and y should have same shape" + assert isinstance(dim, int), "dim should be const integer" + assert out, "out is required" + + +@al.register_custom_op +class _example_custom_op: + name = 'example_custom_op' + core = al.CORE.VECTOR + pipe = al.PIPE.PIPE_V + mode = al.MODE.SIMT + + symbol = 'example_custom_op_impl' + bitcode = os.path.abspath(__file__) + + def __init__(self, src, index, offset: tl.int64, axis, out=None): + # support validate arguments in __init__ method. + assert isinstance(src, tl.tensor), "src should be tensor" + assert index.dtype.is_int(), "index should be integer tensor" + assert isinstance(offset, int), "offset should be integer" + assert isinstance(axis, int), "axis should be integer" + + # support multi-output by using tuple or list. + assert isinstance(out, tuple) and len(out) == 2, "out should be tuple of 2 items" + + # setup the symbol name of the function that will be called at runtime. + rank = len(index.shape) + self.symbol = f"{self.name}_{rank}d_{src.dtype.cname}_{index.dtype.cname}" + + # setup source and compile command if it is implemented by user source code. + self.source = f"workspace/example_custom_op_impl.cce" + self.compile = "bisheng -O2 -std=c++17 -o $@ -c $<" + + # dynamic set argument type. + self.arg_type['axis'] = index.dtype + + +@al.builtin +def example_op(src, index, offset, axis, _semantic=None): + # you can wrap a custom op as a builtin operation, + # output can be provided here to make it easy to use. + x = _semantic.full(src.shape, 0, tl.float32) + y = _semantic.full(index.shape, 0, tl.float32) + return al.custom_semantic(_example_custom_op.name, src, index, offset, axis, out=(x, y), _semantic=_semantic) + + +@triton.jit +def my_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + x = tl.load(x_ptr + i, mask=i < n) + y = tl.load(y_ptr + i, mask=i < n) + y = al.custom("min_custom_op", x, x_ptr, y_ptr + i, al.int64(0), (1, 2, 3), [4.1, 5.2], out=y) + y = al.custom("simple_custom_op", x, y, dim=1, out=y) + index = tl.full((2, 3), 0, tl.int64) + x, y = al.custom("example_custom_op", x, index, offset=1, axis=0, out=(x, y)) + result, _ = example_op(x, index, offset=2, axis=1) + tl.store(out_ptr + i, result, mask=i < n) + + +if __name__ == "__main__": + src = ASTSource(my_kernel, {"x_ptr": "*fp32", "y_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, {"BLOCK": 256}) + context = ir.context() + ir.load_dialects(context) + ascend_ir.load_dialects(context) + options = NPUOptions() + try: + ttir = ast_to_ttir(my_kernel, src, context, options, {}, {}) + print("=== TTIR ===") + print(ttir) + metadata = { + **options.__dict__, + } + linalg = ttir_to_linalg(ttir, metadata, options, named_ops=True) + print("=== MLIR (linalg) ===") + print(linalg) + except subprocess.CalledProcessError as ex: + print(ex.stdout.decode()) + print(ex.stderr.decode()) + print("failed") diff --git a/third_party/ascend/unittest/custom_op/custom_op_extra_buffer_demo.py b/third_party/ascend/unittest/custom_op/custom_op_extra_buffer_demo.py new file mode 100755 index 0000000000..fe2dbe351b --- /dev/null +++ b/third_party/ascend/unittest/custom_op/custom_op_extra_buffer_demo.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python3 +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Demo: declare scratch/extra buffers on a custom op via `extra_buffers` (dtype, size) +# and read back the sizes from lowered HIVM MLIR (`extra_buffers_sizes` attribute). + +from __future__ import annotations + +import os +import re +import subprocess + +import triton +import triton.language as tl +import triton.language.extra.cann.extension as al +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir +from triton.backends.ascend.compiler import NPUOptions, ttir_to_linalg + +# Scratch buffers requested by the custom kernel (element type + length in elements). +SCRATCH_SPEC = [ + (tl.float32, 1024), + (tl.bfloat16, 512), + (tl.int32, 256), +] + + +@al.register_custom_op +class demo_extra_buffer_op: + """Custom op that advertises extra device buffers for the NPU compiler / runtime.""" + + core = al.CORE.VECTOR + pipe = al.PIPE.PIPE_V + mode = al.MODE.SIMT + symbol = "demo_extra_buffer_op_impl" + bitcode = os.path.abspath(__file__) + + def __init__(self, x, out=None): + self.indexing_map = [al.affine_map.get_identity(1)] + self.extra_buffers = list(SCRATCH_SPEC) + + +@triton.jit +def kernel_extra_buf(x_ptr, out_ptr, n, BLOCK: tl.constexpr): + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + x = tl.load(x_ptr + i, mask=i < n) + y = tl.load(out_ptr + i, mask=i < n) + r = al.custom("demo_extra_buffer_op", x, out=y) + tl.store(out_ptr + i, r, mask=i < n) + + +def compile_to_linalg_mlir(kernel, signature: dict, constants: dict) -> str | None: + src = ASTSource(kernel, signature, constants) + ctx = ir.context() + ir.load_dialects(ctx) + ascend_ir.load_dialects(ctx) + options = NPUOptions() + try: + ttir = ast_to_ttir(kernel, src, ctx, options, {}, {}) + meta = {**options.__dict__} + return str(ttir_to_linalg(ttir, meta, options, named_ops=True)) + except subprocess.CalledProcessError as ex: + print(ex.stdout.decode()) + print(ex.stderr.decode()) + return None + + +def extract_extra_buffer_sizes_from_mlir(mlir: str) -> list[int]: + """ + Parse `extra_buffers_sizes` from HIVM custom op text. + """ + # Parse [1024, 512, 256, ...] + m = re.search(r"extra_buffers_sizes\s*=\s*\[([^\]]+)\]", mlir) + if m: + raw = m.group(1).replace(" ", "") + return [int(x) for x in raw.split(",") if x] + + return [] + + +def main() -> None: + expected_sizes = [size for _, size in SCRATCH_SPEC] + print("Declared extra_buffers (dtype, element_count):") + for dt, sz in SCRATCH_SPEC: + print(f" {dt} -> {sz} elements") + + mlir = compile_to_linalg_mlir( + kernel_extra_buf, + {"x_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, + {"BLOCK": 128}, + ) + if not mlir: + print("Compilation failed.") + return + + parsed = extract_extra_buffer_sizes_from_mlir(mlir) + print("\nParsed extra_buffers_sizes from MLIR:", parsed) + if parsed == expected_sizes: + print("OK: MLIR sizes match the Python extra_buffers specification.") + elif parsed: + print("Note: parsed sizes differ from spec; inspect MLIR spelling below.") + else: + print("Could not parse extra_buffers_sizes automatically; " + "search the dump for 'extra_buffers_sizes'.") + + print("\n--- MLIR excerpt (lines containing hivm.hir.custom) ---") + for line in mlir.splitlines(): + if "hivm.hir.custom" in line and "demo_extra_buffer_op" in line: + print(line) + + +if __name__ == "__main__": + main() diff --git a/third_party/ascend/unittest/custom_op/custom_op_indexing_map_complex_demo.py b/third_party/ascend/unittest/custom_op/custom_op_indexing_map_complex_demo.py new file mode 100755 index 0000000000..11e6dcf3cf --- /dev/null +++ b/third_party/ascend/unittest/custom_op/custom_op_indexing_map_complex_demo.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 +import os +import subprocess +import triton +import triton.language as tl +import triton.language.extra.cann.extension as al +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir +from triton.backends.ascend.compiler import NPUOptions, ttir_to_linalg + + +def _make_indexing_maps(): + d0 = ascend_ir.affine_expr.get_dim(0) + d1 = ascend_ir.affine_expr.get_dim(1) + c8 = ascend_ir.affine_expr.get_constant(8) + + # Input maps use transpose and identity-like projections. + in0 = ascend_ir.affine_map.get(2, 0, [d1, d0]) + in1 = ascend_ir.affine_map.get(2, 0, [d0, d1]) + + # Output map models tiled coordinates. + out = ascend_ir.affine_map.get(2, 0, [d0.floordiv(c8), d1.mod(c8)]) + return [in0, in1, out] + + +@al.register_custom_op +class complex_indexing_map_custom_op: + core = al.CORE.VECTOR + pipe = al.PIPE.PIPE_V + mode = al.MODE.SIMT + symbol = "complex_indexing_map_custom" + # Fake path: this example checks IR lowering only. + bitcode = os.path.abspath(__file__) + + def __init__(self, x, y, out=None): + assert out is not None + self.indexing_map = _make_indexing_maps() + + +@triton.jit +def my_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + x = tl.load(x_ptr + i, mask=i < n) + y = tl.load(y_ptr + i, mask=i < n) + out = al.custom("complex_indexing_map_custom_op", x, y, out=x) + tl.store(out_ptr + i, out, mask=i < n) + + +if __name__ == "__main__": + src = ASTSource( + my_kernel, + {"x_ptr": "*fp32", "y_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, + {"BLOCK": 256}, + ) + context = ir.context() + ir.load_dialects(context) + ascend_ir.load_dialects(context) + options = NPUOptions() + try: + ttir = ast_to_ttir(my_kernel, src, context, options, {}, {}) + print("=== TTIR ===") + print(ttir) + linalg = ttir_to_linalg(ttir, {**options.__dict__}, options, named_ops=True) + print("=== MLIR (linalg) ===") + print(linalg) + except subprocess.CalledProcessError as ex: + print(ex.stdout.decode()) + print(ex.stderr.decode()) + print("failed") diff --git a/third_party/ascend/unittest/custom_op/custom_op_indexing_map_compose_demo.py b/third_party/ascend/unittest/custom_op/custom_op_indexing_map_compose_demo.py new file mode 100755 index 0000000000..6d42c5b832 --- /dev/null +++ b/third_party/ascend/unittest/custom_op/custom_op_indexing_map_compose_demo.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +import os +import subprocess +import triton +import triton.language as tl +import triton.language.extra.cann.extension as al +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir +from triton.backends.ascend.compiler import NPUOptions, ttir_to_linalg + + +def _compose_indexing_maps(): + d0 = ascend_ir.affine_expr.get_dim(0) + d1 = ascend_ir.affine_expr.get_dim(1) + c4 = ascend_ir.affine_expr.get_constant(4) + + # Base permutation map. + perm = ascend_ir.affine_map.get_permutation([1, 0]) + # Tile map (row-major tile decomposition). + tile = ascend_ir.affine_map.get(2, 0, [d0.floordiv(c4), d1.mod(c4)]) + # Compose tile with permutation to build a different output indexing. + out = tile.compose(perm) + + in0 = ascend_ir.affine_map.get_identity(2) + in1 = perm + return [in0, in1, out] + + +@al.register_custom_op +class compose_indexing_map_custom_op: + core = al.CORE.VECTOR + pipe = al.PIPE.PIPE_V + mode = al.MODE.SIMT + symbol = "compose_indexing_map_custom" + bitcode = os.path.abspath(__file__) + + def __init__(self, x, y, out=None): + assert out is not None + self.indexing_map = _compose_indexing_maps() + + +@triton.jit +def my_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + x = tl.load(x_ptr + i, mask=i < n) + y = tl.load(y_ptr + i, mask=i < n) + out = al.custom("compose_indexing_map_custom_op", x, y, out=y) + tl.store(out_ptr + i, out, mask=i < n) + + +if __name__ == "__main__": + src = ASTSource( + my_kernel, + {"x_ptr": "*fp32", "y_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, + {"BLOCK": 256}, + ) + context = ir.context() + ir.load_dialects(context) + ascend_ir.load_dialects(context) + options = NPUOptions() + try: + ttir = ast_to_ttir(my_kernel, src, context, options, {}, {}) + print("=== TTIR ===") + print(ttir) + linalg = ttir_to_linalg(ttir, {**options.__dict__}, options, named_ops=True) + print("=== MLIR (linalg) ===") + print(linalg) + except subprocess.CalledProcessError as ex: + print(ex.stdout.decode()) + print(ex.stderr.decode()) + print("failed") diff --git a/third_party/ascend/unittest/custom_op/test_gather_load.py b/third_party/ascend/unittest/custom_op/test_gather_load.py new file mode 100755 index 0000000000..03e83b9171 --- /dev/null +++ b/third_party/ascend/unittest/custom_op/test_gather_load.py @@ -0,0 +1,32 @@ +#!/usr/bin/env python3 +import torch +import triton +import triton.language as tl +import triton.language.extra.cann.extension as al + + +@triton.jit +def test_gather_load_kernel(src_ptr, index_ptr, out_ptr): + # index tile shape: (2, 2) + cols = tl.arange(0, 2)[None, :] # [[0, 1]] + rows = tl.arange(0, 2)[:, None] # [[0],[1]] + mask = (rows < 2) & (cols < 2) + + # load index tile to UB + index = tl.load(index_ptr + rows * 2 + cols, mask) + + # gather load from GM to UB + dst = tl.full(index.shape, 0, tl.float32) + gathered = al.custom("__builtin_gather_load", src_ptr, index, bound=4, dim=0, src_stride=(2, 1), index_shape=(2, 2), + offsets=(0, 0), out=dst) + + # store result to GM + tl.store(out_ptr + rows * 2 + cols, gathered, mask) + + +if __name__ == "__main__": + src = torch.tensor([[1., 2.], [3., 4.], [5., 6.], [7., 8.]], device='npu') + index = torch.tensor([[0, 1], [2, 3]], device='npu') + out = torch.empty((2, 2), device='npu', dtype=torch.float32) + test_gather_load_kernel[(1, )](src, index, out) + print("result: ", out) # [[1., 4.], [5., 8.]] diff --git a/third_party/ascend/unittest/custom_op/test_index_select.py b/third_party/ascend/unittest/custom_op/test_index_select.py new file mode 100644 index 0000000000..d06174fde1 --- /dev/null +++ b/third_party/ascend/unittest/custom_op/test_index_select.py @@ -0,0 +1,44 @@ +import pytest +import torch +import triton +import triton.language as tl +import triton.language.extra.cann.extension as al + + +@triton.jit +def builtin_index_select_kernel(src_ptr, index_ptr, out_ptr): + # Define 2x2 tile indices for output tensor + r = tl.arange(0, 2)[:, None] # Row indices: shape [2, 1] + c = tl.arange(0, 2)[None, :] # Column indices: shape [1, 2] + + # Load index tensor (shape [2]) from GM to UB + idx = tl.load(index_ptr + tl.arange(0, 2)) + # Initialize empty 2x2 output tile in UB (default value: 0) + dst = tl.full((2, 2), 0, dtype=tl.float32) + + # Invoke __builtin_index_select custom op to gather elements + out_tile = al.custom("__builtin_index_select", src_ptr, # Pointer to source tensor in GM + idx, # Index tensor (in UB) for gathering + dim=0, # Dimension to gather along + bound=4, # Upper bound for valid index values (out-of-bound check) + end_offset=(2, 2), # End offsets of each dimension for the index tensor + start_offset=(0, 0), # Start offsets of each dimension for the source tensor + src_stride=(4, 1), # Stride of each dimension for the source tensor in GM + out=dst # Output tensor (in UB) to store gathered elements + ) + + # Store the gathered tile from UB to output tensor in GM + tl.store(out_ptr + r * 2 + c, out_tile) + + +if __name__ == "__main__": + src = torch.tensor( + [[10., 11., 12., 13.], [20., 21., 22., 23.], [30., 31., 32., 33.], [40., 41., 42., 43.]], + device="npu", + dtype=torch.float32, + ) + index = torch.tensor([2, 0], device="npu", dtype=torch.int32) + out = torch.empty((2, 2), device="npu", dtype=torch.float32) + ref = torch.index_select(src, 0, index.to(torch.int64))[:, :2] + builtin_index_select_kernel[(1, )](src, index, out) + torch.testing.assert_close(out, ref) # ref: [[30., 31.], [10., 11.]] diff --git a/third_party/ascend/unittest/generalization_cases/acc_util.py b/third_party/ascend/unittest/generalization_cases/acc_util.py new file mode 100644 index 0000000000..b1295885a6 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/acc_util.py @@ -0,0 +1,144 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import numpy as np +import torch +import torch_npu + +eval_standard = { + torch.float32: { + "rtol": 1e-6, + "small_value": 1e-6, + "small_value_atol": 1e-9, + "etol": 1e-4, + }, + torch.float16: { + "rtol": 1e-3, + "small_value": 1e-3, + "small_value_atol": 1e-5, + "etol": 1e-3, + }, + torch.bfloat16: { + "rtol": 4e-3, + "small_value": 1e-3, + "small_value_atol": 1e-5, + "etol": 1e-3, + }, +} + + +def assert_close(gold: torch.Tensor, act: torch.Tensor, eval_type: str = 'DEFAULT'): + gold = gold.cpu() + act = act.cpu() + if act.dtype == torch.float16 or act.dtype == torch.float32 or act.dtype == torch.bfloat16: + assert gold.dtype == torch.float32, "golden should be f32" + assert not (torch.isnan(act).any() or torch.isinf(act).any()), "actual tensor can not have 'inf' or 'nan'" + eps = eval_standard[act.dtype]['small_value'] + rtol = eval_standard[act.dtype]['rtol'] + atol = eval_standard[act.dtype]['small_value_atol'] + if eval_type == 'DEFAULT': + ae = torch.abs(act - gold) + re = ae / torch.abs(gold) + mask = torch.abs(gold) < eps + + print(f"count ae > {atol}: {(ae > atol).sum()}") + print(f"count re > {rtol}: {(re > rtol).sum()}") + + not_close = torch.where(mask, ae > atol, re > rtol) + print(f"count not_close = {torch.sum(not_close).item()}") + print(f"not_close.numel = {not_close.numel()}, gold.numel = {gold.numel()}") + print(f"not close ratio = {torch.sum(not_close).item() / not_close.numel()}") + if not torch.any(not_close): + return False + + assert torch.sum( + not_close).item() < not_close.numel() * eps, "actual tensor are not close enough with golden tensor,\ +you can use 'benchmark_compare_close' function to compare again!" + + elif eval_type == 'ABS': + act = act.to(gold.dtype) + assert torch.equal(gold, act), "actual tensor and golden tensor are not binary equal!" + else: + assert 0, "ERROR! invalid eval_type" + return False + + +def benchmark_compare_close(gold: torch.Tensor, act: torch.Tensor, std: torch.tensor): + assert act.dtype == std.dtype, "standard tensor's dtype must equal to actual tensor's dtype!" + if act.dtype == torch.float16 or act.dtype == torch.float32 or act.dtype == torch.bfloat16: + assert gold.dtype == torch.float32, "golden should be f32" + assert not (torch.isnan(act).any() or torch.isinf(act).any()), "actual tensor can not have 'inf' or 'nan'" + + gold = gold.cpu() + act = act.cpu() + std = std.cpu() + + eps = eval_standard[act.dtype]['small_value'] + atol = eval_standard[act.dtype]['small_value_atol'] + + mask = torch.abs(gold) <= eps + small_count = mask.sum().item() + + def calculate_relative_errors_except_small(tensor): + re = torch.abs(gold - tensor) / torch.abs(gold) + return torch.where(mask, 0, re) + + act_re = calculate_relative_errors_except_small(act) + std_re = calculate_relative_errors_except_small(std) + act_ae = torch.abs(gold - std) + std_ae = torch.abs(gold - std) + + # 小值域的定义为golden小于某个阈值 eps + act_small_error_count = (mask & (act_ae > atol)).sum().item() + std_small_error_count = (mask & (std_ae > atol)).sum().item() + act_total = act.numel() + std_total = std.numel() + + act_small_error_ratio = act_small_error_count / act_total + std_small_error_ratio = std_small_error_count / std_total + + def calculate_rmse(tensor): + dlt2 = (tensor - gold)**2 + dlt2_except_small_mean = torch.where(mask, 0, dlt2).sum() / small_count + return torch.sqrt(dlt2_except_small_mean) + + act_rmse = calculate_rmse(act) + std_rmse = calculate_rmse(std) + + print(f"act_re.max = {act_re.max()}, std_re.max = {std_re.max()}, limit ratio = 10") + print(f"act_re.sum = {act_re.sum()}, std_re.sum = {std_re.sum()}, limit_ratio = 2") + print( + f"act_small_error_ratio = {act_small_error_ratio}, std_small_error_ratio = {std_small_error_ratio}, limit_ratio = 2" + ) + print(f"act_rmse = {act_rmse}, std_rmse = {std_rmse}, limit_ratio = 2") + + # 条件 1：actual 与 golden 相对误差最大值超过 10 倍 standard 与 golden 相对误差最大值 + assert act_re.max() <= 10 * std_re.max(), "actual re max > stdandard re max's 10 times" + + # 条件 2：actual 与 golden 相对误差均值超过 2 倍 standard 与 golden 相对误差均值 + assert act_re.sum() <= 2 * std_re.sum(), "actual re sum > stdandard re sum's 2 times" + + # 条件 3：actual 小值域 ERROR 占比超过 standard 的两倍 + assert act_small_error_ratio <= 2 * std_small_error_ratio, "act_small_error_ratio > std_small_error_ratio 's 2 times" + + # 条件 4：actual 均方根误差差于 standard 的两倍 + assert act_rmse <= 2 * std_rmse, "act_rmse > std_rmse 's 2 times" + + return False diff --git a/third_party/ascend/unittest/generalization_cases/conftest.py b/third_party/ascend/unittest/generalization_cases/conftest.py new file mode 100644 index 0000000000..633db8b1f4 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/conftest.py @@ -0,0 +1,33 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import torch + + +@pytest.fixture(scope="session", autouse=True) +def assign_npu(worker_id): + npu_count = torch.npu.device_count() + if worker_id == "master": + npu_id = 0 + else: + idx = int(worker_id.replace("gw", "")) + npu_id = idx % npu_count + torch.npu.set_device(npu_id) diff --git a/third_party/ascend/unittest/generalization_cases/test_abs.py b/third_party/ascend/unittest/generalization_cases/test_abs.py new file mode 100644 index 0000000000..c5d06d0d40 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_abs.py @@ -0,0 +1,157 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils, avoid_not_support +import math +import logging + + +def torch_pointwise(x0): + if x0.dtype != torch.uint32: + return torch.abs(x0) + else: + return torch.abs(x0.to(torch.float32)) + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.abs(X) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_abs_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = tl.abs(x_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16', 'int8', 'int16', 'int32', 'int64', 'bool']) +def test_case2(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.debug(f"output.dtype={output.dtype}") + + ans = torch_pointwise(x) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16', 'int8', 'int16', 'int32', 'int64', 'bool']) +def test_abs_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_pointwise(x) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_abs_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_add_mindspore.py b/third_party/ascend/unittest/generalization_cases/test_add_mindspore.py new file mode 100644 index 0000000000..5d5a1638fe --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_add_mindspore.py @@ -0,0 +1,84 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import os +import triton +import triton.language as tl +import numpy as np +import mindspore +import pytest + +pytestmark = pytest.mark.backend("mindspore") + + +@triton.jit +def add_kernel(x_ptr, # *Pointer* to first input vector. + y_ptr, # *Pointer* to second input vector. + output_ptr, # *Pointer* to output vector. + n_elements, # Size of the vector. + BLOCK_SIZE: tl.constexpr, # Number of elements each program should process. + # NOTE: `constexpr` so it can be used as a shape value. + ): + # There are multiple 'programs' processing different data. We identify which program + # we are here: + pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. + # This program will process inputs that are offset from the initial data. + # For instance, if you had a vector of length 256 and block_size of 64, the programs + # would each access the elements [0:64, 64:128, 128:192, 192:256]. + # Note that offsets is a list of pointers: + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + # Create a mask to guard memory operations against out-of-bounds accesses. + mask = offsets < n_elements + # Load x and y from DRAM, masking out any extra elements in case the input is not a + # multiple of the block size. + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + # Write x + y back to DRAM. + tl.store(output_ptr + offsets, output, mask=mask) + + +def add(x: mindspore.Tensor, y: mindspore.Tensor): + output = mindspore.mint.empty_like(x) + n_elements = output.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024) + return output + + +def add_mindspore(x, y): + return x + y + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8)], + ['float16', (2, 4096, 8)], +]) +def test_add_mindspore(param_list): + os.environ["TRITON_BACKEND"] = "mindspore" + dtype, shape = param_list + mindspore.set_seed(0) + x = mindspore.ops.randn(shape, dtype=eval('mindspore.' + dtype)) + y = mindspore.ops.randn(shape, dtype=eval('mindspore.' + dtype)) + output_triton = add(x, y) + output_mindspore = add_mindspore(x, y) + assert np.allclose(output_triton.asnumpy(), output_mindspore.asnumpy(), rtol=1e-3, atol=1e-3) + del os.environ["TRITON_BACKEND"] diff --git a/third_party/ascend/unittest/generalization_cases/test_advance.py b/third_party/ascend/unittest/generalization_cases/test_advance.py new file mode 100644 index 0000000000..8c7a75dc46 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_advance.py @@ -0,0 +1,221 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common +from test_common import TestUtils + + +@triton.jit +def fn_npu_1d(output_ptr, x_ptr, y_ptr, z_ptr, output_ptr1, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + block_ptr_in = tl.make_block_ptr( + base=x_ptr, + shape=(XB, ), + strides=(1, ), + offsets=(5, ), + block_shape=(XB, ), + order=(0, ), + ) + bbptr = tl.advance(block_ptr_in, (-5, )) + # XB,YB,1 + X = tl.load(bbptr) + + block_ptr_out = tl.make_block_ptr( + base=output_ptr, + shape=(XB, ), + strides=(1, ), + offsets=(0, ), + block_shape=(XB, ), + order=(0, ), + ) + tl.store(block_ptr_out, X) + + +@triton.jit +def fn_npu_2d(output_ptr, x_ptr, y_ptr, z_ptr, output_ptr1, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xoffset = tl.program_id(0) + block_ptr_in = tl.make_block_ptr( + base=x_ptr, + shape=(XB, YB), + strides=(YB, 1), + offsets=(6 + xoffset, 5), + block_shape=(XB, YB), + order=(1, 0), + ) + bbptr = tl.advance(block_ptr_in, (-6, -5)) + # XB,YB,1 + X = tl.load(bbptr) + + block_ptr_out = tl.make_block_ptr( + base=output_ptr, + shape=(XB, YB), + strides=(YB, 1), + offsets=(xoffset, 0), + block_shape=(XB, YB), + order=(1, 0), + ) + tl.store(block_ptr_out, X) + + +@triton.jit +def fn_npu_3d(output_ptr, x_ptr, y_ptr, z_ptr, output_ptr1, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + block_ptr_in = tl.make_block_ptr( + base=x_ptr, + shape=(XB, YB, ZB), + strides=(YB * ZB, ZB, 1), + offsets=(3, 1, 2), + block_shape=(XB, YB, ZB), + order=(2, 1, 0), + ) + bbptr = tl.advance(block_ptr_in, (-3, -1, -2)) + # XB,YB,1 + X = tl.load(bbptr) + + block_ptr_out = tl.make_block_ptr( + base=output_ptr, + shape=(XB, YB, ZB), + strides=(YB * ZB, ZB, 1), + offsets=(0, 0, 0), + block_shape=(XB, YB, ZB), + order=(2, 1, 0), + ) + tl.store(block_ptr_out, X) + + +@triton.jit +def triton_advance_4d( + output_ptr, + x_ptr, + BLOCK_0: tl.constexpr, + BLOCK_1: tl.constexpr, + BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, + SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, + SHAPE_3: tl.constexpr, + STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, + STRIDE_2: tl.constexpr, + STRIDE_3: tl.constexpr, +): + block_ptr_in = tl.make_block_ptr( + base=x_ptr, + shape=(SHAPE_0, SHAPE_1, SHAPE_2, SHAPE_3), + strides=(STRIDE_0, STRIDE_1, STRIDE_2, STRIDE_3), + offsets=(6, 5, 4, 3), + block_shape=(BLOCK_0, BLOCK_1, BLOCK_2, BLOCK_3), + order=(3, 2, 1, 0), + ) + bbptr = tl.advance(block_ptr_in, (-6, -5, -4, -3)) + x = tl.load(bbptr) + + block_ptr_out = tl.make_block_ptr( + base=output_ptr, + shape=(SHAPE_0, SHAPE_1, SHAPE_2, SHAPE_3), + strides=(STRIDE_0, STRIDE_1, STRIDE_2, STRIDE_3), + offsets=(0, 0, 0, 0), + block_shape=(BLOCK_0, BLOCK_1, BLOCK_2, BLOCK_3), + order=(3, 2, 1, 0), + ) + tl.store(block_ptr_out, x) + + +@triton.jit +def triton_advance_5d( + output_ptr, + x_ptr, + BLOCK_0: tl.constexpr, + BLOCK_1: tl.constexpr, + BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, + BLOCK_4: tl.constexpr, + SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, + SHAPE_3: tl.constexpr, + SHAPE_4: tl.constexpr, + STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, + STRIDE_2: tl.constexpr, + STRIDE_3: tl.constexpr, + STRIDE_4: tl.constexpr, +): + block_ptr_in = tl.make_block_ptr( + base=x_ptr, + shape=(SHAPE_0, SHAPE_1, SHAPE_2, SHAPE_3, SHAPE_4), + strides=(STRIDE_0, STRIDE_1, STRIDE_2, STRIDE_3, STRIDE_4), + offsets=(6, 5, 4, 3, 2), + block_shape=(BLOCK_0, BLOCK_1, BLOCK_2, BLOCK_3, BLOCK_4), + order=(4, 3, 2, 1, 0), + ) + bbptr = tl.advance(block_ptr_in, (-6, -5, -4, -3, -2)) + x = tl.load(bbptr) + + block_ptr_out = tl.make_block_ptr( + base=output_ptr, + shape=(SHAPE_0, SHAPE_1, SHAPE_2, SHAPE_3, SHAPE_4), + strides=(STRIDE_0, STRIDE_1, STRIDE_2, STRIDE_3, STRIDE_4), + offsets=(0, 0, 0, 0, 0), + block_shape=(BLOCK_0, BLOCK_1, BLOCK_2, BLOCK_3, BLOCK_4), + order=(4, 3, 2, 1, 0), + ) + tl.store(block_ptr_out, x) + + +temporarily_not_support_dtype = ['bool'] + + +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('shape', TestUtils.full_shape) +def test_npu(dtype, shape): + if dtype in temporarily_not_support_dtype: + return + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + + a = x + blocks = list(x.size()) + strides = list(x.stride()) + grid = (1, ) + if len(shape) == 5: + triton_advance_5d[grid](output, x, *blocks, *blocks, *strides) + elif len(shape) == 4: + triton_advance_4d[grid](output, x, *blocks, *blocks, *strides) + elif len(shape) == 3: + fn_npu_3d[1, 1, 1](output, x, y, z, output1, XB=shape[0], YB=shape[1], ZB=shape[2]) + elif len(shape) == 2: + if x.numel() * x.element_size() > 8192: + fn_npu_2d[shape[0], 1, 1](output, x, y, z, output1, XB=1, YB=shape[1], ZB=1) + else: + fn_npu_2d[1, 1, 1](output, x, y, z, output1, XB=shape[0], YB=shape[1], ZB=1) + else: + fn_npu_1d[1, 1, 1](output, x, y, z, output1, XB=shape[0], YB=1, ZB=1) + + torch.testing.assert_close(output, a) diff --git a/third_party/ascend/unittest/generalization_cases/test_and.py b/third_party/ascend/unittest/generalization_cases/test_and.py new file mode 100644 index 0000000000..4bac287eaf --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_and.py @@ -0,0 +1,161 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils +import math + + +def torch_pointwise(x, y): + res = x & y + return res + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = X & Y + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_and_4d_5d(output_ptr, x_ptr, y_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = x_val & y_val + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'bool']) +def test_case2(dtype, shape): + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.debug(f"output.dtype={output.dtype}") + + ans = torch_pointwise(x, y) + output = torch.zeros_like(ans) + + if len(shape) == 1: + fn_npu_[1, 1, shape[0]](output, x, y, z, 1, 1, 1, 1, 1, shape[0]) + elif len(shape) == 2: + if shape[0] > shape[1]: + fn_npu_[1, shape[0], 1](output, x, y, z, 1, 1, shape[1], 1, shape[0], shape[1]) + else: + fn_npu_[1, 1, shape[1]](output, x, y, z, 1, shape[0], 1, 1, shape[0], shape[1]) + elif len(shape) == 3: + if max(shape[0], shape[1], shape[2]) == shape[0]: + fn_npu_[shape[0], 1, 1](output, x, y, z, 1, shape[1], shape[2], shape[0], shape[1], shape[2]) + elif max(shape[0], shape[1], shape[2]) == shape[1]: + fn_npu_[1, shape[1], 1](output, x, y, z, shape[0], 1, shape[2], shape[0], shape[1], shape[2]) + else: + fn_npu_[1, 1, shape[2]](output, x, y, z, shape[0], shape[1], 1, shape[0], shape[1], shape[2]) + else: + fn_npu_[1, 1, 1](output, x, y, z, 1, 1, 1, 1, 1, 1) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'bool']) +def test_and_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = x & y + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_and_4d_5d[grid](output, x, y, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) + + +invalid_types = [ + 'float16', + 'float32', + 'bfloat16', +] + + +@pytest.mark.parametrize("sigtype", invalid_types) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "unexpected type") +def test_invalid_types(sigtype): + N = 32 + x = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + y = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + z = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + output = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + + fn_npu_[1, 1, 1](output, x, y, z, 32, 1, 1, 32, 1, 1) diff --git a/third_party/ascend/unittest/generalization_cases/test_argmax.py b/third_party/ascend/unittest/generalization_cases/test_argmax.py new file mode 100644 index 0000000000..edbf8b9d8d --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_argmax.py @@ -0,0 +1,362 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import math +import pytest +import torch +import torch_npu +import numpy as np +import triton +import triton.language as tl + +import test_common +from test_common import TestUtils, check_ub_mem_overflow, get_dtype_size + +logger = logging.getLogger(__name__) + + +# <<<<<<< test_argmax_1d +def torch_argmax(x0, dim, keepdim): + x0 = x0 if x0.device == "cpu" else x0.cpu() + return torch.argmax(x0, dim=dim, keepdim=keepdim).npu() + + +@triton.jit +def triton_argmax_1d(in_ptr0, out_ptr1, xnumel, XBLOCK: tl.constexpr): + xoffset = tl.program_id(0) + tl.arange(0, XBLOCK) + tmp0 = tl.load(in_ptr0 + xoffset, None) + tmp4 = tl.argmax(tmp0, 0) + tl.store(out_ptr1, tmp4, None) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32']) +def test_argmax_1d(dtype, shape): + dtype_size = get_dtype_size(dtype) + if dtype_size * math.prod(shape) >= (TestUtils.ub_size / 3): + logger.warning(f"dtype:{dtype} shape:{shape} mem overflow") + return + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty(1, dtype=torch.int32).npu() + numel = shape[0] + triton_argmax_1d[1, 1, 1](x0, triton_res, numel, numel) + torch_res = torch_argmax(x0, dim=0, keepdim=True) + test_common.validate_cmp("int32", triton_res, torch_res) + + +# >>>>>>> test_argmax_1d + + +# <<<<<<< test_argmax_2d +@triton.jit +def triton_argmax_2d(in_ptr0, out_ptr0, dim: tl.constexpr, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, + NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, MNUMEL) + nblk_idx = tl.arange(0, NNUMEL) + mmask = mblk_idx < M + nmask = nblk_idx < N + mask = (mmask[:, None]) & (nmask[None, :]) + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + x = tl.load(in_ptr0 + idx, mask=mask, other=-float('inf')) + tmp4 = tl.argmax(x, dim) + if dim == 0: + tl.store(out_ptr0 + tl.arange(0, N), tmp4, None) + else: + tl.store(out_ptr0 + tl.arange(0, M), tmp4, None) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32']) +@pytest.mark.parametrize('dim', [0, 1]) +def test_argmax_2d(dtype, shape, dim): + dtype_size = get_dtype_size(dtype) + if dtype_size * math.prod(shape) >= (TestUtils.ub_size / 3): + logger.warning(f"dtype:{dtype} shape:{shape} mem overflow") + return + shapex, shapey = shape + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty([ + shape[1 - dim], + ], dtype=torch.int32).npu() + triton_argmax_2d[1, 1, 1](x0, triton_res, dim, shapex, shapey, shapex, shapey) + torch_res = torch_argmax(x0, dim=dim, keepdim=False) + test_common.validate_cmp("int32", triton_res, torch_res) + + +# >>>>>>> test_argmax_2d + + +# <<<<<<< test_argmax_3d +def torch_argmax_3d(x0, no_reduce_dim): + x0 = x0 if x0.device == "cpu" else x0.cpu() + if x0.dtype in (torch.int8, torch.int16, torch.int32): + x0 = x0.to(torch.int64) + if no_reduce_dim == 0: + return torch.argmax(torch.max(x0, 1)[0], 1).npu() + elif no_reduce_dim == 1: + return torch.argmax(torch.max(x0, 0)[0], 1).npu() + elif no_reduce_dim == 2: + return torch.argmax(torch.max(x0, 0)[0], 0).npu() + else: + assert False, f"no reduce dim not right, no_reduce_dim = {no_reduce_dim}" + + +@triton.jit +def triton_argmax_3d_0_1(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + + x = tl.load(in_ptr + idx) + + tmp = tl.max(x, 0) + ret = tl.argmax(tmp, 0) + oidx = zidx + tl.store(out_ptr + oidx, ret) + + +@triton.jit +def triton_argmax_3d_0_2(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + + x = tl.load(in_ptr + idx) + + tmp = tl.max(x, 0) + ret = tl.argmax(tmp, 1) + oidx = yidx + tl.store(out_ptr + oidx, ret) + + +@triton.jit +def triton_argmax_3d_1_2(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + + x = tl.load(in_ptr + idx) + + tmp = tl.max(x, 1) + ret = tl.argmax(tmp, 1) + oidx = xidx + tl.store(out_ptr + oidx, ret) + + +def triton_argmax_3d(in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB, no_reduce_dim): + if no_reduce_dim == 0: + triton_argmax_3d_1_2[1, 1, 1](in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB) + elif no_reduce_dim == 1: + triton_argmax_3d_0_2[1, 1, 1](in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB) + elif no_reduce_dim == 2: + triton_argmax_3d_0_1[1, 1, 1](in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32']) +@pytest.mark.parametrize('no_reduce_dim', [0, 1, 2]) +def test_argmax_3d(dtype, shape, no_reduce_dim): + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty([ + shape[no_reduce_dim], + ], dtype=torch.int32).npu() + triton_argmax_3d(x0, triton_res, shape[0], shape[1], shape[2], shape[0], shape[1], shape[2], no_reduce_dim) + torch_res = torch_argmax_3d(x0, no_reduce_dim) + test_common.validate_cmp("int32", triton_res, torch_res) + + +# >>>>>>> test_argmax_3d + + +# <<<<<<< test_argmax_4d +def torch_argmax_4d(x0, dim): + x0 = x0 if x0.device == "cpu" else x0.cpu() + if x0.dtype in (torch.int8, torch.int16, torch.int32): + x0 = x0.to(torch.int64) + return torch.argmax(x0, dim) + + +@triton.jit +def argmax_4d(out_ptr, x, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, DIM: tl.constexpr): + if DIM == 0: + ret = tl.reshape(tl.argmax(x, DIM), XB * YB * ZB * MB // XB) + o_idx = tl.arange(0, XB * YB * ZB * MB // XB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 1: + ret = tl.reshape(tl.argmax(x, DIM), XB * YB * ZB * MB // YB) + o_idx = tl.arange(0, XB * YB * ZB * MB // YB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 2: + ret = tl.reshape(tl.argmax(x, DIM), XB * YB * ZB * MB // ZB) + o_idx = tl.arange(0, XB * YB * ZB * MB // ZB) + tl.store(out_ptr + o_idx, ret) + else: + ret = tl.reshape(tl.argmax(x, DIM), XB * YB * ZB * MB // MB) + o_idx = tl.arange(0, XB * YB * ZB * MB // MB) + tl.store(out_ptr + o_idx, ret) + + +@triton.jit +def triton_argmax_kernel_4d(in_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + DIM: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + midx = tl.arange(0, MB) + + idx = xidx[:, None, None, None] * YB * ZB * MB + yidx[None, :, None, None] * ZB * MB + zidx[ + None, None, :, None] * MB + midx[None, None, None, :] + + x = tl.load(in_ptr + idx) + + argmax_4d(out_ptr, x, XB, YB, ZB, MB, DIM) + + +def triton_argmax_4d(in_ptr, out_ptr, XB, YB, ZB, MB, dim): + triton_argmax_kernel_4d[(1, )](in_ptr, out_ptr, XB, YB, ZB, MB, dim) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (2, 2, 4, 8), + (2, 3, 4, 8), +]) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32']) +@pytest.mark.parametrize('dim', [0]) +def test_argmax_4d(dtype, shape, dim): + x0 = test_common.generate_tensor(shape, dtype).npu() + torch_res = torch_argmax_4d(x0, dim).to(torch.int32) + triton_res = torch.empty_like(torch_res, dtype=torch.int32).npu() + triton_argmax_4d(x0, triton_res, shape[0], shape[1], shape[2], shape[3], dim) + + test_common.validate_cmp("int32", triton_res, torch_res) + + +# >>>>>>> test_argmax_4d + + +# <<<<<<< test_argmax_5d +def torch_argmax_5d(x0, dim): + x0 = x0 if x0.device == "cpu" else x0.cpu() + if x0.dtype in (torch.int8, torch.int16, torch.int32): + x0 = x0.to(torch.int64) + return torch.argmax(x0, dim) + + +@triton.jit +def argmax_5d(out_ptr, x, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, NB: tl.constexpr, + DIM: tl.constexpr): + if DIM == 0: + ret = tl.reshape(tl.argmax(x, DIM), XB * YB * ZB * MB * NB // XB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // XB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 1: + ret = tl.reshape(tl.argmax(x, DIM), XB * YB * ZB * MB * NB // YB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // YB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 2: + ret = tl.reshape(tl.argmax(x, DIM), XB * YB * ZB * MB * NB // ZB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // ZB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 3: + ret = tl.reshape(tl.argmax(x, DIM), XB * YB * ZB * MB * NB // MB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // MB) + tl.store(out_ptr + o_idx, ret) + else: + ret = tl.reshape(tl.argmax(x, DIM), XB * YB * ZB * MB * NB // NB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // NB) + tl.store(out_ptr + o_idx, ret) + + +@triton.jit +def triton_argmax_kernel_5d(in_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr, DIM: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + midx = tl.arange(0, MB) + nidx = tl.arange(0, NB) + + idx = xidx[:, None, None, None, None] * YB * ZB * MB * NB + yidx[None, :, None, None, None] * ZB * MB * NB + zidx[ + None, None, :, None, None] * MB * NB + midx[None, None, None, :, None] * NB + nidx[None, None, None, None, :] + + x = tl.load(in_ptr + idx) + + argmax_5d(out_ptr, x, XB, YB, ZB, MB, NB, DIM) + + +def triton_argmax_5d(in_ptr, out_ptr, XB, YB, ZB, MB, NB, dim): + triton_argmax_kernel_5d[(1, )](in_ptr, out_ptr, XB, YB, ZB, MB, NB, dim) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (2, 2, 2, 4, 8), + (2, 2, 3, 4, 8), +]) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32']) +@pytest.mark.parametrize('dim', [0]) +def test_argmax_5d(dtype, shape, dim): + x0 = test_common.generate_tensor(shape, dtype).npu() + torch_res = torch_argmax_5d(x0, dim).to(torch.int32) + triton_res = torch.empty_like(torch_res, dtype=torch.int32).npu() + triton_argmax_5d(x0, triton_res, shape[0], shape[1], shape[2], shape[3], shape[4], dim) + + test_common.validate_cmp("int32", triton_res, torch_res) + + +# >>>>>>> test_argmax_5d + + +# <<<<<<< test_argmax_1d_bool +@triton.jit +def triton_argmax_1d_bool(in_ptr0, out_ptr1, xnumel, XBLOCK: tl.constexpr): + xoffset = tl.program_id(0) + tl.arange(0, XBLOCK) + tmp0 = tl.load(in_ptr0 + xoffset, None).to(tl.int1) + tmp4 = tl.argmax(tmp0, 0) + tl.store(out_ptr1, tmp4, None) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', ['bool']) +def test_argmax_1d_bool(dtype, shape): + dtype_size = get_dtype_size(dtype) + if dtype_size * math.prod(shape) >= (TestUtils.ub_size / 3): + logger.warning(f"dtype:{dtype} shape:{shape} mem overflow") + return + x0 = test_common.generate_tensor(shape, dtype) + triton_res = torch.empty(1, dtype=torch.int32).npu() + numel = shape[0] + triton_argmax_1d_bool[1, 1, 1](x0.npu(), triton_res, numel, numel) + np_res = np.argmax(x0.numpy()) + np.equal(triton_res.item(), np_res) + + +# >>>>>>> test_argmax_1d_bool diff --git a/third_party/ascend/unittest/generalization_cases/test_argmin.py b/third_party/ascend/unittest/generalization_cases/test_argmin.py new file mode 100644 index 0000000000..36a671d1ba --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_argmin.py @@ -0,0 +1,360 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import math +import pytest +import torch +import torch_npu +import numpy as np +import triton +import triton.language as tl + +import test_common +from test_common import TestUtils, check_ub_mem_overflow, get_dtype_size + +logger = logging.getLogger(__name__) + + +# <<<<<<< test_argmin_1d +def torch_argmin(input_tensor, dim, keepdim): + return torch.argmin(input_tensor, dim=dim, keepdim=keepdim) + + +@triton.jit +def triton_argmin_1d(in_ptr0, out_ptr1, xnumel, XBLOCK: tl.constexpr): + xoffset = tl.program_id(0) + tl.arange(0, XBLOCK) + tmp0 = tl.load(in_ptr0 + xoffset, None) + tmp4 = tl.argmin(tmp0, 0) + tl.store(out_ptr1, tmp4, None) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32']) +def test_argmin_1d(dtype, shape): + dtype_size = get_dtype_size(dtype) + if dtype_size * math.prod(shape) >= (TestUtils.ub_size / 3): + logger.warning(f"dtype:{dtype} shape:{shape} mem overflow") + return + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty(1, dtype=torch.int32).npu() + numel = shape[0] + triton_argmin_1d[1, 1, 1](x0, triton_res, numel, numel) + torch_res = torch_argmin(x0, dim=0, keepdim=True) + test_common.validate_cmp("int32", triton_res, torch_res) + + +# >>>>>>> test_argmin_1d + + +# <<<<<<< test_argmin_2d +@triton.jit +def triton_argmin_2d(in_ptr0, out_ptr0, dim: tl.constexpr, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, + NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, MNUMEL) + nblk_idx = tl.arange(0, NNUMEL) + mmask = mblk_idx < M + nmask = nblk_idx < N + mask = (mmask[:, None]) & (nmask[None, :]) + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + x = tl.load(in_ptr0 + idx, mask=mask, other=-float('inf')) + tmp4 = tl.argmin(x, dim) + if dim == 0: + tl.store(out_ptr0 + tl.arange(0, N), tmp4, None) + else: + tl.store(out_ptr0 + tl.arange(0, M), tmp4, None) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32']) +@pytest.mark.parametrize('dim', [0, 1]) +def test_argmin_2d(dtype, shape, dim): + dtype_size = get_dtype_size(dtype) + if dtype_size * math.prod(shape) >= (TestUtils.ub_size / 3): + logger.warning(f"dtype:{dtype} shape:{shape} mem overflow") + return + shapex, shapey = shape + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty([ + shape[1 - dim], + ], dtype=torch.int32).npu() + triton_argmin_2d[1, 1, 1](x0, triton_res, dim, shapex, shapey, shapex, shapey) + torch_res = torch_argmin(x0, dim=dim, keepdim=False) + test_common.validate_cmp("int32", triton_res, torch_res) + + +# >>>>>>> test_argmin_2d + + +# <<<<<<< test_argmin_3d +def torch_argmin_3d(x0, no_reduce_dim): + if x0.dtype in (torch.int8, torch.int16, torch.int32): + x0 = x0.to(torch.int64) + if no_reduce_dim == 0: + return torch.argmin(torch.min(x0, 1)[0], 1) + elif no_reduce_dim == 1: + return torch.argmin(torch.min(x0, 0)[0], 1) + elif no_reduce_dim == 2: + return torch.argmin(torch.min(x0, 0)[0], 0) + else: + assert False, f"no reduce dim not right, no_reduce_dim = {no_reduce_dim}" + + +@triton.jit +def triton_argmin_3d_0_1(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + + x = tl.load(in_ptr + idx) + + tmp = tl.min(x, 0) + ret = tl.argmin(tmp, 0) + oidx = zidx + tl.store(out_ptr + oidx, ret) + + +@triton.jit +def triton_argmin_3d_0_2(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + + x = tl.load(in_ptr + idx) + + tmp = tl.min(x, 0) + ret = tl.argmin(tmp, 1) + oidx = yidx + tl.store(out_ptr + oidx, ret) + + +@triton.jit +def triton_argmin_3d_1_2(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + + x = tl.load(in_ptr + idx) + + tmp = tl.min(x, 1) + ret = tl.argmin(tmp, 1) + oidx = xidx + tl.store(out_ptr + oidx, ret) + + +def triton_argmin_3d(in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB, no_reduce_dim): + if no_reduce_dim == 0: + triton_argmin_3d_1_2[1, 1, 1](in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB) + elif no_reduce_dim == 1: + triton_argmin_3d_0_2[1, 1, 1](in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB) + elif no_reduce_dim == 2: + triton_argmin_3d_0_1[1, 1, 1](in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32']) +@pytest.mark.parametrize('no_reduce_dim', [0, 1, 2]) +def test_argmin_3d(dtype, shape, no_reduce_dim): + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty([ + shape[no_reduce_dim], + ], dtype=torch.int32).npu() + triton_argmin_3d(x0, triton_res, shape[0], shape[1], shape[2], shape[0], shape[1], shape[2], no_reduce_dim) + torch_res = torch_argmin_3d(x0, no_reduce_dim) + test_common.validate_cmp("int32", triton_res, torch_res) + + +# >>>>>>> test_argmin_3d + + +# <<<<<<< test_argmin_4d +def torch_argmin_4d(x0, dim): + x0 = x0 if x0.device == "cpu" else x0.cpu() + if x0.dtype in (torch.int8, torch.int16, torch.int32): + x0 = x0.to(torch.int64) + return torch.argmin(x0, dim) + + +@triton.jit +def argmin_4d(out_ptr, x, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, DIM: tl.constexpr): + if DIM == 0: + ret = tl.reshape(tl.argmin(x, DIM), XB * YB * ZB * MB // XB) + o_idx = tl.arange(0, XB * YB * ZB * MB // XB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 1: + ret = tl.reshape(tl.argmin(x, DIM), XB * YB * ZB * MB // YB) + o_idx = tl.arange(0, XB * YB * ZB * MB // YB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 2: + ret = tl.reshape(tl.argmin(x, DIM), XB * YB * ZB * MB // ZB) + o_idx = tl.arange(0, XB * YB * ZB * MB // ZB) + tl.store(out_ptr + o_idx, ret) + else: + ret = tl.reshape(tl.argmin(x, DIM), XB * YB * ZB * MB // MB) + o_idx = tl.arange(0, XB * YB * ZB * MB // MB) + tl.store(out_ptr + o_idx, ret) + + +@triton.jit +def triton_argmin_kernel_4d(in_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + DIM: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + midx = tl.arange(0, MB) + + idx = xidx[:, None, None, None] * YB * ZB * MB + yidx[None, :, None, None] * ZB * MB + zidx[ + None, None, :, None] * MB + midx[None, None, None, :] + + x = tl.load(in_ptr + idx) + + argmin_4d(out_ptr, x, XB, YB, ZB, MB, DIM) + + +def triton_argmin_4d(in_ptr, out_ptr, XB, YB, ZB, MB, dim): + triton_argmin_kernel_4d[(1, )](in_ptr, out_ptr, XB, YB, ZB, MB, dim) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (2, 2, 4, 8), + (2, 3, 4, 8), +]) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32']) +@pytest.mark.parametrize('dim', [0]) +def test_argmin_4d(dtype, shape, dim): + x0 = test_common.generate_tensor(shape, dtype).npu() + torch_res = torch_argmin_4d(x0, dim).to(torch.int32) + triton_res = torch.empty_like(torch_res, dtype=torch.int32).npu() + triton_argmin_4d(x0, triton_res, shape[0], shape[1], shape[2], shape[3], dim) + + test_common.validate_cmp("int32", triton_res, torch_res) + + +# >>>>>>> test_argmin_4d + + +# <<<<<<< test_argmin_5d +def torch_argmin_5d(x0, dim): + x0 = x0 if x0.device == "cpu" else x0.cpu() + if x0.dtype in (torch.int8, torch.int16, torch.int32): + x0 = x0.to(torch.int64) + return torch.argmin(x0, dim) + + +@triton.jit +def argmin_5d(out_ptr, x, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, NB: tl.constexpr, + DIM: tl.constexpr): + if DIM == 0: + ret = tl.reshape(tl.argmin(x, DIM), XB * YB * ZB * MB * NB // XB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // XB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 1: + ret = tl.reshape(tl.argmin(x, DIM), XB * YB * ZB * MB * NB // YB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // YB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 2: + ret = tl.reshape(tl.argmin(x, DIM), XB * YB * ZB * MB * NB // ZB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // ZB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 3: + ret = tl.reshape(tl.argmin(x, DIM), XB * YB * ZB * MB * NB // MB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // MB) + tl.store(out_ptr + o_idx, ret) + else: + ret = tl.reshape(tl.argmin(x, DIM), XB * YB * ZB * MB * NB // NB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // NB) + tl.store(out_ptr + o_idx, ret) + + +@triton.jit +def triton_argmin_kernel_5d(in_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr, DIM: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + midx = tl.arange(0, MB) + nidx = tl.arange(0, NB) + + idx = xidx[:, None, None, None, None] * YB * ZB * MB * NB + yidx[None, :, None, None, None] * ZB * MB * NB + zidx[ + None, None, :, None, None] * MB * NB + midx[None, None, None, :, None] * NB + nidx[None, None, None, None, :] + + x = tl.load(in_ptr + idx) + + argmin_5d(out_ptr, x, XB, YB, ZB, MB, NB, DIM) + + +def triton_argmin_5d(in_ptr, out_ptr, XB, YB, ZB, MB, NB, dim): + triton_argmin_kernel_5d[(1, )](in_ptr, out_ptr, XB, YB, ZB, MB, NB, dim) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (2, 2, 2, 4, 8), + (2, 2, 3, 4, 8), +]) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32']) +@pytest.mark.parametrize('dim', [0]) +def test_argmin_5d(dtype, shape, dim): + x0 = test_common.generate_tensor(shape, dtype).npu() + torch_res = torch_argmin_5d(x0, dim).to(torch.int32) + triton_res = torch.empty_like(torch_res, dtype=torch.int32).npu() + triton_argmin_5d(x0, triton_res, shape[0], shape[1], shape[2], shape[3], shape[4], dim) + + test_common.validate_cmp("int32", triton_res, torch_res) + + +# >>>>>>> test_argmin_5d + + +# <<<<<<< test_argmin_1d_bool +@triton.jit +def triton_argmin_1d_bool(in_ptr0, out_ptr1, xnumel, XBLOCK: tl.constexpr): + xoffset = tl.program_id(0) + tl.arange(0, XBLOCK) + tmp0 = tl.load(in_ptr0 + xoffset, None).to(tl.int1) + tmp4 = tl.argmin(tmp0, 0) + tl.store(out_ptr1, tmp4, None) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', ['bool']) +def test_argmin_1d_bool(dtype, shape): + dtype_size = get_dtype_size(dtype) + if dtype_size * math.prod(shape) >= (TestUtils.ub_size / 3): + logger.warning(f"dtype:{dtype} shape:{shape} mem overflow") + return + x0 = test_common.generate_tensor(shape, dtype) + triton_res = torch.empty(1, dtype=torch.int32).npu() + numel = shape[0] + triton_argmin_1d_bool[1, 1, 1](x0.npu(), triton_res, numel, numel) + np_res = np.argmin(x0.numpy()) + np.equal(triton_res.item(), np_res) + + +# >>>>>>> test_argmin_1d_bool diff --git a/third_party/ascend/unittest/generalization_cases/test_associative_scan.py b/third_party/ascend/unittest/generalization_cases/test_associative_scan.py new file mode 100644 index 0000000000..249abc09fb --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_associative_scan.py @@ -0,0 +1,523 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import pytest +import random +import torch +import torch_npu +import triton +import triton.language as tl + +import test_common +from test_common import TestUtils, get_dtype_size + + +def combine_fn_test_torch(a, b, combine_fn): + if combine_fn == 'maximum_fn': + return torch.maximum(a, b) # 最大值 + elif combine_fn == 'minimum_fn': + return torch.minimum(a, b) # 最小值 + elif combine_fn == 'bitwise_xor_fn': + return a ^ b # 按位异或 + elif combine_fn == 'bitwise_or_fn': + return a | b # 按位异 + elif combine_fn == 'bitwise_and_fn': + return a & b # 按位与 + else: + pytest.skip("The combine_fn is not within the following scope , skipping.") + + +def torch_func_scan(input: torch.Tensor, dim: int, combine_fn='maximum', reverse=False): + """ + PyTorch 实现 associative_scan，语义与 Triton 完全对齐 + 支持任意 combine_fn（如 a|b, a&b, min, max 等） + """ + dim = dim % input.ndim + + if reverse: + input = input.flip(dim) + + N = input.size(dim) + + tensors = torch.unbind(input, dim=dim) + + outputs = [] + + carry = tensors[0] + outputs.append(carry) + + for i in range(1, N): + carry = combine_fn_test_torch(tensors[i], carry, combine_fn) + outputs.append(carry) + + output = torch.stack(outputs, dim=dim) + + if reverse: + output = output.flip(dim) + + return output + + +@triton.jit +def bitwise_and_fn(a, b): + return a & b + + +@triton.jit +def bitwise_or_fn(a, b): + return a | b + + +@triton.jit +def bitwise_xor_fn(a, b): + return a ^ b + + +@triton.jit +def minimum_fn(a, b): + return tl.minimum(a, b) + + +@triton.jit +def maximum_fn(a, b): + return tl.maximum(a, b) + + +@triton.jit +def triton_kernel_1d_scan( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + numel_x: tl.constexpr, + XBLOCK: tl.constexpr, + combine_fn_name: tl.constexpr, +): + tl.static_assert(numel_x == XBLOCK, "numel_x must be equal to XBLOCK in this kernel") + idx = tl.arange(0, XBLOCK) + x = tl.load(in_ptr0 + idx) + if combine_fn_name == "maximum_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=maximum_fn) + elif combine_fn_name == "minimum_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=minimum_fn) + elif combine_fn_name == "bitwise_or_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=bitwise_or_fn) + elif combine_fn_name == "bitwise_xor_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=bitwise_xor_fn) + elif combine_fn_name == "bitwise_and_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=bitwise_and_fn) + + tl.store(out_ptr0 + idx, ret) + + +@triton.jit +def triton_kernel_2d_scan( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + numel_x: tl.constexpr, + numel_r: tl.constexpr, + XBLOCK: tl.constexpr, + RBLOCK: tl.constexpr, + combine_fn_name: tl.constexpr, +): + tl.static_assert(numel_x == XBLOCK, "numel_x must be equal to XBLOCK in this kernel") + tl.static_assert(numel_r == RBLOCK, "numel_r must be equal to RBLOCK in this kernel") + idx_x = tl.arange(0, XBLOCK) + idx_r = tl.arange(0, RBLOCK) + idx = idx_x[:, None] * numel_r + idx_r[None, :] + x = tl.load(in_ptr0 + idx) + + if combine_fn_name == "maximum_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=maximum_fn) + elif combine_fn_name == "minimum_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=minimum_fn) + elif combine_fn_name == "bitwise_or_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=bitwise_or_fn) + elif combine_fn_name == "bitwise_xor_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=bitwise_xor_fn) + elif combine_fn_name == "bitwise_and_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=bitwise_and_fn) + tl.store(out_ptr0 + idx, ret) + + +@triton.jit +def triton_kernel_3d_scan( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + numel_x: tl.constexpr, + numel_r: tl.constexpr, + numel_z: tl.constexpr, + XBLOCK: tl.constexpr, + RBLOCK: tl.constexpr, + ZBLOCK: tl.constexpr, + combine_fn_name: tl.constexpr, +): + tl.static_assert(numel_x == XBLOCK, "numel_x must be equal to XBLOCK in this kernel") + tl.static_assert(numel_r == RBLOCK, "numel_r must be equal to RBLOCK in this kernel") + tl.static_assert(numel_z == ZBLOCK, "numel_z must be equal to ZBLOCK in this kernel") + idx_x = tl.arange(0, XBLOCK) + idx_r = tl.arange(0, RBLOCK) + idx_z = tl.arange(0, ZBLOCK) + idx = idx_x[:, None, None] * numel_r * numel_z + idx_r[None, :, None] * numel_z + idx_z[None, None, :] + x = tl.load(in_ptr0 + idx) + if combine_fn_name == "maximum_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=maximum_fn) + elif combine_fn_name == "minimum_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=minimum_fn) + elif combine_fn_name == "bitwise_or_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=bitwise_or_fn) + elif combine_fn_name == "bitwise_xor_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=bitwise_xor_fn) + elif combine_fn_name == "bitwise_and_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=bitwise_and_fn) + tl.store(out_ptr0 + idx, ret) + + +@triton.jit +def triton_kernel_4d_scan( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + XB: tl.constexpr, + YB: tl.constexpr, + ZB: tl.constexpr, + MB: tl.constexpr, + combine_fn_name: tl.constexpr, +): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + midx = tl.arange(0, MB) + idx = (xidx[:, None, None, None] * YB * ZB * MB + yidx[None, :, None, None] * ZB * MB + + zidx[None, None, :, None] * MB + midx[None, None, None, :]) + x = tl.load(in_ptr0 + idx) + if combine_fn_name == "maximum_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=maximum_fn) + elif combine_fn_name == "minimum_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=minimum_fn) + elif combine_fn_name == "bitwise_or_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=bitwise_or_fn) + elif combine_fn_name == "bitwise_xor_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=bitwise_xor_fn) + elif combine_fn_name == "bitwise_and_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=bitwise_and_fn) + tl.store(out_ptr0 + idx, ret) + + +@triton.jit +def triton_kernel_5d_scan( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + XB: tl.constexpr, + YB: tl.constexpr, + ZB: tl.constexpr, + MB: tl.constexpr, + NB: tl.constexpr, + combine_fn_name: tl.constexpr, +): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + midx = tl.arange(0, MB) + nidx = tl.arange(0, NB) + idx = (xidx[:, None, None, None, None] * YB * ZB * MB * NB + yidx[None, :, None, None, None] * ZB * MB * NB + + zidx[None, None, :, None, None] * MB * NB + midx[None, None, None, :, None] * NB + + nidx[None, None, None, None, :]) + x = tl.load(in_ptr0 + idx) + if combine_fn_name == "maximum_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=maximum_fn) + elif combine_fn_name == "minimum_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=minimum_fn) + elif combine_fn_name == "bitwise_or_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=bitwise_or_fn) + elif combine_fn_name == "bitwise_xor_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=bitwise_xor_fn) + elif combine_fn_name == "bitwise_and_fn": + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=bitwise_and_fn) + tl.store(out_ptr0 + idx, ret) + + +def triton_func_scan(x, dim, combine_fn, reverse): + res = torch.empty_like(x) + shape = x.size() + + if len(shape) == 1: + if dim >= 1: + pytest.skip("dim >= 1 for 1D tensor, skipping.") + triton_kernel_1d_scan[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[0], combine_fn) + elif len(shape) == 2: + if dim >= 2: + pytest.skip("dim >= 2 for 2D tensor, skipping.") + triton_kernel_2d_scan[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[1], x.shape[0], x.shape[1], combine_fn) + elif len(shape) == 3: + if dim >= 3: + pytest.skip("dim >= 3 for 3D tensor, skipping.") + triton_kernel_3d_scan[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[1], x.shape[2], x.shape[0], x.shape[1], + x.shape[2], combine_fn) + elif len(shape) == 4: + if dim >= 4: + pytest.skip("dim >= 4 for 4D tensor, skipping.") + triton_kernel_4d_scan[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[1], x.shape[2], x.shape[3], combine_fn) + elif len(shape) == 5: + if dim >= 5: + pytest.skip("dim >= 5 for 5D tensor, skipping.") + triton_kernel_5d_scan[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[1], x.shape[2], x.shape[3], x.shape[4], + combine_fn) + else: + pytest.skip(f"Unsupported tensor dimension: {len(shape)}") + + return res + + +def should_skip_due_to_mem(dtype, shape): + dtype_size = get_dtype_size(dtype) + total_mem = dtype_size * math.prod(shape) + if dtype in ('int8', 'bool'): + threshold = TestUtils.ub_size / 13 + else: + threshold = TestUtils.ub_size / 6 + + if total_mem >= threshold: + pytest.skip(f"dtype:{dtype} shape:{shape} mem overflow") + + +@pytest.mark.parametrize("dtype", ['int8', 'int16', 'int32', 'int64', 'bool']) +@pytest.mark.parametrize("shape", TestUtils.test_shape1d) +@pytest.mark.parametrize("dim", [0]) +@pytest.mark.parametrize("combine_fn", + ['maximum_fn', 'minimum_fn', 'bitwise_or_fn', 'bitwise_xor_fn', 'bitwise_and_fn']) +@pytest.mark.parametrize("reverse", [False]) +def test_scan_1d(dtype, shape, dim, combine_fn, reverse): + should_skip_due_to_mem(dtype, shape) + torch.manual_seed(0) + x = test_common.generate_tensor(shape=shape, dtype=dtype) + x_gold = x + cpu_res = torch_func_scan(x_gold, dim, combine_fn, reverse) + + x_npu = x.npu() + triton_res = triton_func_scan(x_npu, dim, combine_fn, reverse) + + test_common.validate_cmp(dtype, triton_res, cpu_res) + + +@pytest.mark.parametrize("dtype", ['int8', 'int16', 'int32', 'int64', 'bool']) +@pytest.mark.parametrize("shape", TestUtils.test_shape2d) +@pytest.mark.parametrize("dim", [1]) +@pytest.mark.parametrize("combine_fn", + ['maximum_fn', 'minimum_fn', 'bitwise_or_fn', 'bitwise_xor_fn', 'bitwise_and_fn']) +@pytest.mark.parametrize("reverse", [False]) +def test_scan_2d(dtype, shape, dim, combine_fn, reverse): + should_skip_due_to_mem(dtype, shape) + torch.manual_seed(0) + x = test_common.generate_tensor(shape=shape, dtype=dtype) + x_gold = x + cpu_res = torch_func_scan(x_gold, dim, combine_fn, reverse) + + x_npu = x.npu() + triton_res = triton_func_scan(x_npu, dim, combine_fn, reverse) + + test_common.validate_cmp(dtype, triton_res, cpu_res) + + +@pytest.mark.parametrize("dtype", ['int8', 'int16', 'int32', 'int64', 'bool']) +@pytest.mark.parametrize("shape", TestUtils.test_shape3d) +@pytest.mark.parametrize("dim", [2]) +@pytest.mark.parametrize("combine_fn", + ['maximum_fn', 'minimum_fn', 'bitwise_or_fn', 'bitwise_xor_fn', 'bitwise_and_fn']) +@pytest.mark.parametrize("reverse", [False]) +def test_scan_3d(dtype, shape, dim, combine_fn, reverse): + should_skip_due_to_mem(dtype, shape) + torch.manual_seed(0) + x = test_common.generate_tensor(shape=shape, dtype=dtype) + x_gold = x + cpu_res = torch_func_scan(x_gold, dim, combine_fn, reverse) + + x_npu = x.npu() + triton_res = triton_func_scan(x_npu, dim, combine_fn, reverse) + + test_common.validate_cmp(dtype, triton_res, cpu_res) + + +@pytest.mark.parametrize("dtype", ['int8', 'int16', 'int32', 'int64', 'bool']) +@pytest.mark.parametrize("shape", TestUtils.test_shape4d) +@pytest.mark.parametrize("dim", [3]) +@pytest.mark.parametrize("combine_fn", + ['maximum_fn', 'minimum_fn', 'bitwise_or_fn', 'bitwise_xor_fn', 'bitwise_and_fn']) +@pytest.mark.parametrize("reverse", [False]) +def test_scan_4d(dtype, shape, dim, combine_fn, reverse): + should_skip_due_to_mem(dtype, shape) + torch.manual_seed(0) + x = test_common.generate_tensor(shape=shape, dtype=dtype) + x_gold = x + cpu_res = torch_func_scan(x_gold, dim, combine_fn, reverse) + + x_npu = x.npu() + triton_res = triton_func_scan(x_npu, dim, combine_fn, reverse) + + test_common.validate_cmp(dtype, triton_res, cpu_res) + + +@pytest.mark.parametrize("dtype", ['int8', 'int16', 'int32', 'int64', 'bool']) +@pytest.mark.parametrize("shape", TestUtils.test_shape5d) +@pytest.mark.parametrize("dim", [4]) +@pytest.mark.parametrize("combine_fn", + ['maximum_fn', 'minimum_fn', 'bitwise_or_fn', 'bitwise_xor_fn', 'bitwise_and_fn']) +@pytest.mark.parametrize("reverse", [False]) +def test_scan_5d(dtype, shape, dim, combine_fn, reverse): + should_skip_due_to_mem(dtype, shape) + torch.manual_seed(0) + x = test_common.generate_tensor(shape=shape, dtype=dtype) + x_gold = x + cpu_res = torch_func_scan(x_gold, dim, combine_fn, reverse) + + x_npu = x.npu() + triton_res = triton_func_scan(x_npu, dim, combine_fn, reverse) + + test_common.validate_cmp(dtype, triton_res, cpu_res) + + +@pytest.mark.parametrize("dtype", ['float16', 'float32']) +@pytest.mark.parametrize("shape", TestUtils.test_shape1d) +@pytest.mark.parametrize("dim", [0]) +@pytest.mark.parametrize("combine_fn", ['maximum_fn', 'minimum_fn']) +@pytest.mark.parametrize("reverse", [False]) +def test_scan_float_1d(dtype, shape, dim, combine_fn, reverse): + should_skip_due_to_mem(dtype, shape) + torch.manual_seed(0) + x = test_common.generate_tensor(shape=shape, dtype=dtype) + x_gold = x + cpu_res = torch_func_scan(x_gold, dim, combine_fn, reverse) + + x_npu = x.npu() + triton_res = triton_func_scan(x_npu, dim, combine_fn, reverse) + + test_common.validate_cmp(dtype, triton_res, cpu_res) + + +@pytest.mark.parametrize("dtype", ['float16', 'float32']) +@pytest.mark.parametrize("shape", random.sample(TestUtils.test_shape2d, 5)) +@pytest.mark.parametrize("dim", [1]) +@pytest.mark.parametrize("combine_fn", ['maximum_fn', 'minimum_fn']) +@pytest.mark.parametrize("reverse", [False]) +def test_scan_float_2d(dtype, shape, dim, combine_fn, reverse): + should_skip_due_to_mem(dtype, shape) + torch.manual_seed(0) + x = test_common.generate_tensor(shape=shape, dtype=dtype) + x_gold = x + cpu_res = torch_func_scan(x_gold, dim, combine_fn, reverse) + + x_npu = x.npu() + triton_res = triton_func_scan(x_npu, dim, combine_fn, reverse) + + test_common.validate_cmp(dtype, triton_res, cpu_res) + + +@pytest.mark.parametrize("dtype", ['float16', 'float32']) +@pytest.mark.parametrize("shape", TestUtils.test_shape3d) +@pytest.mark.parametrize("dim", [2]) +@pytest.mark.parametrize("combine_fn", ['maximum_fn', 'minimum_fn']) +@pytest.mark.parametrize("reverse", [False]) +def test_scan_float_1d(dtype, shape, dim, combine_fn, reverse): + should_skip_due_to_mem(dtype, shape) + torch.manual_seed(0) + x = test_common.generate_tensor(shape=shape, dtype=dtype) + x_gold = x + cpu_res = torch_func_scan(x_gold, dim, combine_fn, reverse) + + x_npu = x.npu() + triton_res = triton_func_scan(x_npu, dim, combine_fn, reverse) + + test_common.validate_cmp(dtype, triton_res, cpu_res) + + +@pytest.mark.parametrize("dtype", ['float16', 'float32']) +@pytest.mark.parametrize("shape", TestUtils.test_shape4d) +@pytest.mark.parametrize("dim", [3]) +@pytest.mark.parametrize("combine_fn", ['maximum_fn', 'minimum_fn']) +@pytest.mark.parametrize("reverse", [False]) +def test_scan_float_1d(dtype, shape, dim, combine_fn, reverse): + should_skip_due_to_mem(dtype, shape) + torch.manual_seed(0) + x = test_common.generate_tensor(shape=shape, dtype=dtype) + x_gold = x + cpu_res = torch_func_scan(x_gold, dim, combine_fn, reverse) + + x_npu = x.npu() + triton_res = triton_func_scan(x_npu, dim, combine_fn, reverse) + + test_common.validate_cmp(dtype, triton_res, cpu_res) + + +@pytest.mark.parametrize("dtype", ['float16', 'float32']) +@pytest.mark.parametrize("shape", TestUtils.test_shape5d) +@pytest.mark.parametrize("dim", [4]) +@pytest.mark.parametrize("combine_fn", ['maximum_fn', 'minimum_fn']) +@pytest.mark.parametrize("reverse", [False]) +def test_scan_float_1d(dtype, shape, dim, combine_fn, reverse): + should_skip_due_to_mem(dtype, shape) + torch.manual_seed(0) + x = test_common.generate_tensor(shape=shape, dtype=dtype) + x_gold = x + cpu_res = torch_func_scan(x_gold, dim, combine_fn, reverse) + + x_npu = x.npu() + triton_res = triton_func_scan(x_npu, dim, combine_fn, reverse) + + test_common.validate_cmp(dtype, triton_res, cpu_res) + + +@pytest.mark.parametrize("dtype", ['float16', 'float32']) +@pytest.mark.parametrize("shape", TestUtils.test_shape1d) +@pytest.mark.parametrize("dim", [0]) +@pytest.mark.parametrize("combine_fn", ['bitwise_or_fn', 'bitwise_xor_fn', 'bitwise_and_fn']) +@pytest.mark.parametrize("reverse", [False]) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "unexpected type") +def test_scan_float_invalid(dtype, shape, dim, combine_fn, reverse): + should_skip_due_to_mem(dtype, shape) + torch.manual_seed(0) + x = test_common.generate_tensor(shape=shape, dtype=dtype) + + x_npu = x.npu() + triton_res = triton_func_scan(x_npu, dim, combine_fn, reverse) + + +@pytest.mark.parametrize("dtype", ['int32']) +@pytest.mark.parametrize("shape", TestUtils.test_shape1d) +@pytest.mark.parametrize("dim", [0]) +@pytest.mark.parametrize("combine_fn", + ['maximum_fn', 'minimum_fn', 'bitwise_or_fn', 'bitwise_xor_fn', 'bitwise_and_fn']) +@pytest.mark.parametrize("reverse", [True]) +@test_common.raises_with_match(triton.compiler.errors.MLIRCompilationError, + "reverse=True is not yet supported for scan op") +def test_scan_float_invalid_reverse(dtype, shape, dim, combine_fn, reverse): + should_skip_due_to_mem(dtype, shape) + torch.manual_seed(0) + x = test_common.generate_tensor(shape=shape, dtype=dtype) + + x_npu = x.npu() + triton_res = triton_func_scan(x_npu, dim, combine_fn, reverse) diff --git a/third_party/ascend/unittest/generalization_cases/test_atan.py b/third_party/ascend/unittest/generalization_cases/test_atan.py new file mode 100644 index 0000000000..37fe8dadfe --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_atan.py @@ -0,0 +1,62 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils +import math +import triton.language.extra.ascend.libdevice as libdevice + + +def torch_pointwise(x0): + res = torch.atan(x0) + return res + + +@triton.jit +def triton_atan(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop1 in range(loops1): + x0_prime = offset + (loop1 * XBLOCK_SUB) + base1 + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp2 = libdevice.atan(tmp0) + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', ['float32', 'float16']) +def test_case(dtype, shape): + x0 = test_common.generate_tensor(shape, dtype).npu() + + numel = x0.numel() + ncore = 1 if numel <= 32 else 32 + xblock = math.ceil(numel / ncore) + xblock_sub = numel if numel <= ncore else math.ceil(numel / ncore) + + y_ref = torch_pointwise(x0) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_atan[ncore, 1, 1](x0, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/generalization_cases/test_atomic_add.py b/third_party/ascend/unittest/generalization_cases/test_atomic_add.py new file mode 100644 index 0000000000..d55448ae66 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_atomic_add.py @@ -0,0 +1,576 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import pytest +import torch +import triton + +import triton.language as tl + +import numpy as np +import test_common +from test_common import TestUtils + +filtered_dtype = [dtype for dtype in TestUtils.full_dtype if dtype not in {'int64', 'bool'}] + + +@triton.jit +def atomic_add(in_ptr0, out_ptr0, out_ptr1, n_elements, BLOCK_SIZE: tl.constexpr): + offset = tl.program_id(0) * BLOCK_SIZE + index = offset + tl.arange(0, BLOCK_SIZE)[:] + xmask = index < n_elements + + tmp0 = tl.load(in_ptr0 + (index), xmask) + tmp1 = tl.load(out_ptr0 + (index), xmask) + tl.atomic_add(out_ptr1 + (index), tmp0, xmask) + tl.atomic_add(out_ptr1 + (index), tmp1, xmask) + + +@triton.jit +def atomic_add_broadcast(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(0) + + x = tl.load(x_ptr) # x is scalar or 1D, no mask needed + + # Compute y indices + y_offset = pid * BLOCK_SIZE + y_indices = y_offset + tl.arange(0, BLOCK_SIZE) + y_mask = y_indices < n_elements + + y_value = tl.load(y_ptr + y_indices, y_mask) + # Atomic add: y += x (broadcasted) + tl.atomic_add(out_ptr + y_indices, y_value, mask=y_mask) + tl.atomic_add(out_ptr + y_indices, x, mask=y_mask) + + +# 定义不同测试场景的参数组合 (x_shape, y_shape, BLOCK_SIZE) +test_cases = [ + ((1, 1, 1, 1), (1, 1, 1, 4), 4), + ((1, 1, 1, 3), (1, 5, 1, 3), 5), + ((3, ), (2, 3, 3, 3, 3), 81), + ((3, ), (2, 3, 3, 3), 27), + ((3, ), (2, 3, 3), 9), + ((3, ), (2, 3), 3), +] + + +def promote_dtype(x_dtype, y_dtype): + """ + 如果 y 的精度低于 x, 则提升 y 的精度以匹配 x。 + """ + # 如果两个数据类型一致，直接返回 + if x_dtype == y_dtype: + return y_dtype + + # 构建类型的优先级列表（从低到高） + priority = [torch.int8, torch.int16, torch.int32, torch.float16, torch.bfloat16, torch.float32] + + # 查找两种类型在优先级列表中的位置 + x_priority = priority.index(x_dtype) + y_priority = priority.index(y_dtype) + + # 如果y的优先级比x小，则提升到x的类型 + if y_priority < x_priority: + return x_dtype + else: + return y_dtype + + +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +@pytest.mark.parametrize('y_dtype_str', filtered_dtype) +@pytest.mark.parametrize('x_shape, y_shape, BLOCK_SIZE', test_cases) +def test_atomic_add_broadcast_combined(x_dtype_str, y_dtype_str, x_shape, y_shape, BLOCK_SIZE): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + # 先构造 x0 + x0 = torch.full(x_shape, 83.0000, dtype=x_dtype).npu() + + y_raw_dtype = eval('torch.' + y_dtype_str) + + out_dtype = promote_dtype(x_dtype, y_raw_dtype) + if out_dtype == torch.bfloat16: + out_dtype = torch.float32 + + # 构造y和out + y = torch.full(y_shape, -105, dtype=y_raw_dtype).npu() + out = torch.full(y_shape, 0, dtype=out_dtype).npu() + + # 保存副本用于验证 + x_temp = x0.clone() + y_temp = y.clone() + out_temp = out.clone() + + # 计算网格大小和元素总数 + n_elements = y.numel() + grid = (n_elements // BLOCK_SIZE, ) # 自动计算需要的线程块数量 + + # 调用 Triton 核函数 + atomic_add_broadcast[grid](x_ptr=x0, y_ptr=y, out_ptr=out, n_elements=n_elements, BLOCK_SIZE=BLOCK_SIZE) + + # 验证结果：y += x (广播加法) + expected = out_temp + y_temp + x_temp + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d + TestUtils.test_shape1d) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +@pytest.mark.parametrize('y_dtype_str', filtered_dtype) +def test_atomic_add(x_dtype_str, y_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + y_dtype = eval('torch.' + y_dtype_str) + + x0 = test_common.generate_tensor(shape, x_dtype_str).npu() + x1 = test_common.generate_tensor(shape, y_dtype_str).npu() + out_dtype = promote_dtype(x_dtype, y_dtype) + if out_dtype == torch.bfloat16: + out_dtype = torch.float32 + y = torch.full(x1.shape, 0, dtype=out_dtype).npu() + + # 保存副本用于验证 + x0_temp = x0.clone() + x1_temp = x1.clone() + y_temp = y.clone() + + if len(shape) == 2: + n_elements = shape[0] * shape[1] + atomic_add[shape[0], 1, 1](x0, x1, y, n_elements, BLOCK_SIZE=shape[1]) + elif len(shape) == 1: + n_elements = shape[0] + BLOCK_SIZE = min(1024, shape[0]) # 1024:限制最大线程块大小 + grid_size = (n_elements + BLOCK_SIZE - 1) // BLOCK_SIZE # 向上取整 + atomic_add[grid_size, 1, 1](x0, x1, y, n_elements, BLOCK_SIZE=BLOCK_SIZE) + + expected = y_temp + x1_temp + x0_temp + torch.testing.assert_close(y, expected) + + +# 3d +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +@pytest.mark.parametrize('y_dtype_str', filtered_dtype) +def test_atomic_add_3d(x_dtype_str, y_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + y_dtype = eval('torch.' + y_dtype_str) + + x0 = test_common.generate_tensor(shape, x_dtype_str).npu() + x1 = test_common.generate_tensor(shape, y_dtype_str).npu() + out_dtype = promote_dtype(x_dtype, y_dtype) + if out_dtype == torch.bfloat16: + out_dtype = torch.float32 + y = torch.full(x1.shape, 0, dtype=out_dtype).npu() + + # 保存副本用于验证 + x0_temp = x0.clone() + x1_temp = x1.clone() + y_temp = y.clone() + + n_elements = shape[0] * shape[1] * shape[2] + atomic_add[1, 1, 1](x0, x1, y, n_elements, BLOCK_SIZE=shape[0] * shape[1] * shape[2]) + + expected = y_temp + x1_temp + x0_temp + torch.testing.assert_close(y, expected) + + +@triton.jit +def atomic_add_multi_d(in_ptr0, out_ptr0, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if (YB * ZB * MB * NB) > 1: + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if (ZB * MB * NB) > 1: + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if (MB * NB) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if NB > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + tmp0 = tl.load(in_ptr0 + offsets) + tl.atomic_add(out_ptr0 + offsets, tmp0) + + +# multi_d +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (2, 4, 8, 4), + (8, 4, 2, 4), + (2, 8, 2, 2), + (2, 4, 8, 4, 2), + (8, 4, 2, 4, 4), + (2, 8, 2, 2, 2), +]) +@pytest.mark.parametrize('dtype', filtered_dtype) +def test_atomic_add_4d_5d(dtype, shape): + x0_value = 3 + x0 = torch.full(shape, x0_value, dtype=eval('torch.' + dtype)).npu() + x1 = torch.full(shape, 2, dtype=eval('torch.' + dtype)).npu() + + x1_ref = x1 + x0_value + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + atomic_add_multi_d[(1, )](x0, x1, *triton_shape) + test_common.validate_cmp(dtype, x1, x1_ref) + + +@triton.jit +def atomic_add_5d(x_ptr, y_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr, XB1: tl.constexpr, YB1: tl.constexpr, ZB1: tl.constexpr, MB1: tl.constexpr, + NB1: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + offsets1 = tl.arange(0, XB1) * (YB1 * ZB1 * MB1 * NB1) + offsets1 = offsets1[:, None] + tl.arange(0, YB1)[None, :] * (ZB1 * MB1 * NB1) + offsets1 = offsets1[:, :, None] + tl.arange(0, ZB1)[None, None, :] * (MB1 * NB1) + offsets1 = offsets1[:, :, :, None] + tl.arange(0, MB1)[None, None, None, :] * NB1 + offsets1 = offsets1[:, :, :, :, None] + tl.arange(0, NB1)[None, None, None, None, :] + + tmp0 = tl.load(x_ptr + offsets) + tmp1 = tl.load(y_ptr + offsets1) + tl.atomic_add(out_ptr + offsets1, tmp0) + tl.atomic_add(out_ptr + offsets1, tmp1) + + +@pytest.mark.parametrize('param_list', [ + [(1, 1, 2, 1, 1), (1, 1, 2, 1, 2)], +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +@pytest.mark.parametrize('y_dtype_str', filtered_dtype) +def test_atomic_add_5d(x_dtype_str, y_dtype_str, param_list): + x0_shape, y_shape = param_list + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + y_dtype = eval('torch.' + y_dtype_str) + if x_dtype == torch.int8 or x_dtype == torch.int16 or x_dtype == torch.int32: + x0 = torch.randint(low=0, high=100, size=x0_shape, dtype=x_dtype).npu() + else: + x0 = torch.randn(x0_shape, dtype=eval('torch.' + x_dtype_str)).npu() + + if y_dtype == torch.int8 or y_dtype == torch.int16 or y_dtype == torch.int32: + y = torch.randint(low=0, high=100, size=y_shape, dtype=y_dtype).npu() + else: + y = torch.randn(y_shape, dtype=eval('torch.' + y_dtype_str)).npu() + + out_dtype = promote_dtype(x_dtype, y_dtype) + if out_dtype == torch.bfloat16: + out_dtype = torch.float32 + out = torch.full(y_shape, 0, dtype=out_dtype).npu() + + x0_temp = x0.clone() + y_temp = y.clone() + out_temp = out.clone() + + triton_shape = [*x0_shape] + while len(triton_shape) < 5: + triton_shape.append(1) + XB, YB, ZB, MB, NB = triton_shape + + triton_shape1 = [*y_shape] + while len(triton_shape1) < 5: + triton_shape1.append(1) + XB1, YB1, ZB1, MB1, NB1 = triton_shape1 + + atomic_add_5d[(1, )]( + x_ptr=x0, + y_ptr=y, + out_ptr=out, + XB=XB, + YB=YB, + ZB=ZB, + MB=MB, + NB=NB, + XB1=XB1, + YB1=YB1, + ZB1=ZB1, + MB1=MB1, + NB1=NB1, + ) + + expected = out_temp + y_temp + x0_temp + torch.testing.assert_close(out, expected) + + +@triton.jit +def atomic_add_4d(x_ptr, y_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + XB1: tl.constexpr, YB1: tl.constexpr, ZB1: tl.constexpr, MB1: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB) + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB) + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB) + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] + + offsets1 = tl.arange(0, XB1) * (YB1 * ZB1 * MB1) + offsets1 = offsets1[:, None] + tl.arange(0, YB1)[None, :] * (ZB1 * MB1) + offsets1 = offsets1[:, :, None] + tl.arange(0, ZB1)[None, None, :] * (MB1) + offsets1 = offsets1[:, :, :, None] + tl.arange(0, MB1)[None, None, None, :] + + tmp0 = tl.load(x_ptr + offsets) + tmp1 = tl.load(y_ptr + offsets1) + tl.atomic_add(out_ptr + offsets1, tmp0) + tl.atomic_add(out_ptr + offsets1, tmp1) + + +@pytest.mark.parametrize('param_list', [ + [(1, 1, 2, 1), (1, 1, 2, 2)], + [(1, 1, 1, 1), (1, 1, 2, 2)], +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +@pytest.mark.parametrize('y_dtype_str', filtered_dtype) +def test_atomic_add_4d(x_dtype_str, y_dtype_str, param_list): + x0_shape, y_shape = param_list + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + y_dtype = eval('torch.' + y_dtype_str) + if x_dtype == torch.int8 or x_dtype == torch.int16 or x_dtype == torch.int32: + x0 = torch.randint(low=0, high=100, size=x0_shape, dtype=x_dtype).npu() + else: + x0 = torch.randn(x0_shape, dtype=eval('torch.' + x_dtype_str)).npu() + + if y_dtype == torch.int8 or y_dtype == torch.int16 or y_dtype == torch.int32: + y = torch.randint(low=0, high=100, size=y_shape, dtype=y_dtype).npu() + else: + y = torch.randn(y_shape, dtype=eval('torch.' + y_dtype_str)).npu() + + out_dtype = promote_dtype(x_dtype, y_dtype) + if out_dtype == torch.bfloat16: + out_dtype = torch.float32 + out = torch.full(y_shape, 0, dtype=out_dtype).npu() + + x0_temp = x0.clone() + y_temp = y.clone() + out_temp = out.clone() + + triton_shape = [*x0_shape] + while len(triton_shape) < 4: + triton_shape.append(1) + XB, YB, ZB, MB = triton_shape + + triton_shape1 = [*y_shape] + while len(triton_shape1) < 4: + triton_shape1.append(1) + XB1, YB1, ZB1, MB1 = triton_shape1 + + atomic_add_4d[(1, )]( + x_ptr=x0, + y_ptr=y, + out_ptr=out, + XB=XB, + YB=YB, + ZB=ZB, + MB=MB, + XB1=XB1, + YB1=YB1, + ZB1=ZB1, + MB1=MB1, + ) + + expected = out_temp + y_temp + x0_temp + torch.testing.assert_close(out, expected) + + +@triton.jit +def atomic_add_3d(x_ptr, y_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XB1: tl.constexpr, + YB1: tl.constexpr, ZB1: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB) + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB) + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] + + offsets1 = tl.arange(0, XB1) * (YB1 * ZB1) + offsets1 = offsets1[:, None] + tl.arange(0, YB1)[None, :] * (ZB1) + offsets1 = offsets1[:, :, None] + tl.arange(0, ZB1)[None, None, :] + + tmp0 = tl.load(x_ptr + offsets) + tmp1 = tl.load(y_ptr + offsets1) + tl.atomic_add(out_ptr + offsets1, tmp0) + tl.atomic_add(out_ptr + offsets1, tmp1) + + +@pytest.mark.parametrize('param_list', [ + [(1, 1, 2), (1, 2, 2)], +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +@pytest.mark.parametrize('y_dtype_str', filtered_dtype) +def test_atomic_add_3d_2(x_dtype_str, y_dtype_str, param_list): + x0_shape, y_shape = param_list + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + y_dtype = eval('torch.' + y_dtype_str) + if x_dtype == torch.int8 or x_dtype == torch.int16 or x_dtype == torch.int32: + x0 = torch.randint(low=0, high=100, size=x0_shape, dtype=x_dtype).npu() + else: + x0 = torch.randn(x0_shape, dtype=eval('torch.' + x_dtype_str)).npu() + + if y_dtype == torch.int8 or y_dtype == torch.int16 or y_dtype == torch.int32: + y = torch.randint(low=0, high=100, size=y_shape, dtype=y_dtype).npu() + else: + y = torch.randn(y_shape, dtype=eval('torch.' + y_dtype_str)).npu() + + out_dtype = promote_dtype(x_dtype, y_dtype) + if out_dtype == torch.bfloat16: + out_dtype = torch.float32 + out = torch.full(y_shape, 0, dtype=out_dtype).npu() + + x0_temp = x0.clone() + y_temp = y.clone() + out_temp = out.clone() + + triton_shape = [*x0_shape] + while len(triton_shape) < 3: + triton_shape.append(1) + XB, YB, ZB = triton_shape + + triton_shape1 = [*y_shape] + while len(triton_shape1) < 3: + triton_shape1.append(1) + XB1, YB1, ZB1 = triton_shape1 + + atomic_add_3d[(1, )]( + x_ptr=x0, + y_ptr=y, + out_ptr=out, + XB=XB, + YB=YB, + ZB=ZB, + XB1=XB1, + YB1=YB1, + ZB1=ZB1, + ) + + expected = out_temp + y_temp + x0_temp + torch.testing.assert_close(out, expected) + + +@triton.jit +def atomic_add_2d(x_ptr, y_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, XB1: tl.constexpr, YB1: tl.constexpr): + offsets = tl.arange(0, XB) * (YB) + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] + + offsets1 = tl.arange(0, XB1) * (YB1) + offsets1 = offsets1[:, None] + tl.arange(0, YB1)[None, :] + + tmp0 = tl.load(x_ptr + offsets) + tmp1 = tl.load(y_ptr + offsets1) + tl.atomic_add(out_ptr + offsets1, tmp0) + tl.atomic_add(out_ptr + offsets1, tmp1) + + +@pytest.mark.parametrize('param_list', [ + [(1, 2), (2, 2)], + [(1, 1), (2, 2)], +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +@pytest.mark.parametrize('y_dtype_str', filtered_dtype) +def test_atomic_add_2d(x_dtype_str, y_dtype_str, param_list): + x0_shape, y_shape = param_list + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + y_dtype = eval('torch.' + y_dtype_str) + if x_dtype == torch.int8 or x_dtype == torch.int16 or x_dtype == torch.int32: + x0 = torch.randint(low=0, high=100, size=x0_shape, dtype=x_dtype).npu() + else: + x0 = torch.randn(x0_shape, dtype=eval('torch.' + x_dtype_str)).npu() + + if y_dtype == torch.int8 or y_dtype == torch.int16 or y_dtype == torch.int32: + y = torch.randint(low=0, high=100, size=y_shape, dtype=y_dtype).npu() + else: + y = torch.randn(y_shape, dtype=eval('torch.' + y_dtype_str)).npu() + + out_dtype = promote_dtype(x_dtype, y_dtype) + if out_dtype == torch.bfloat16: + out_dtype = torch.float32 + out = torch.full(y_shape, 0, dtype=out_dtype).npu() + + x0_temp = x0.clone() + y_temp = y.clone() + out_temp = out.clone() + + triton_shape = [*x0_shape] + while len(triton_shape) < 2: + triton_shape.append(1) + XB, YB = triton_shape + + triton_shape1 = [*y_shape] + while len(triton_shape1) < 2: + triton_shape1.append(1) + XB1, YB1 = triton_shape1 + + atomic_add_2d[(1, )]( + x_ptr=x0, + y_ptr=y, + out_ptr=out, + XB=XB, + YB=YB, + XB1=XB1, + YB1=YB1, + ) + + expected = out_temp + y_temp + x0_temp + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('param_list', [ + ['uint8', (32, 32), 2], + ['uint16', (32, 32), 2], + ['uint32', (32, 32), 2], +]) +def test_atomic_add_uint(param_list): + dtype, shape, ncore = param_list + block_size = shape[0] * shape[1] / ncore + split_size = shape[0] // ncore + x0_value = 3 + x0_cpu = torch.full(shape, x0_value, dtype=eval(f'torch.{dtype}')).cpu() + x0 = x0_cpu.to("npu") + x1_cpu = torch.full((split_size, shape[1]), 4, dtype=eval(f'torch.{dtype}')).cpu() + x1 = x1_cpu.to("npu") + y_cpu = torch.full((split_size, shape[1]), -10, dtype=eval(f'torch.{dtype}')).cpu() + y = y_cpu.to("npu") + + x1_np = x1_cpu.numpy() + y_ref_np = x1_np + 0 + x1_ref_np = x1_np + ncore * x0_value + + x1_ref = torch.from_numpy(x1_ref_np).npu() + y_ref = torch.from_numpy(y_ref_np).npu() + + @triton.jit + def atomic_add_uint(in_ptr0, out_ptr0, out_ptr1, n_elements, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + yindex = tl.arange(0, BLOCK_SIZE)[:] + xmask = xindex < n_elements + x0 = xindex + x1 = yindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.atomic_add(out_ptr0 + (x1), tmp0, xmask) + tl.store(out_ptr1 + (x1), tmp1, xmask) + + n_elements = shape[0] * shape[1] + atomic_add_uint[ncore, 1, 1](x0, x1, y, n_elements, BLOCK_SIZE=split_size * shape[1]) + test_common.validate_cmp(dtype, x1, x1_ref) diff --git a/third_party/ascend/unittest/generalization_cases/test_atomic_and.py b/third_party/ascend/unittest/generalization_cases/test_atomic_and.py new file mode 100644 index 0000000000..0ef250741c --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_atomic_and.py @@ -0,0 +1,562 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import pytest +import torch +import triton + +import triton.language as tl + +import test_common +from test_common import TestUtils + +filtered_dtype = [dtype for dtype in TestUtils.full_dtype if dtype not in {'float16', 'float32', 'bfloat16', 'bool'}] + + +@triton.jit +def atomic_and(in_ptr0, out_ptr0, n_elements, BLOCK_SIZE: tl.constexpr, BLOCK_NUM: tl.constexpr): + in_offset = tl.program_id(0) * BLOCK_SIZE + out_offset = (tl.program_id(0) % BLOCK_NUM) * BLOCK_SIZE + in_index = in_offset + tl.arange(0, BLOCK_SIZE) + out_index = out_offset + tl.arange(0, BLOCK_SIZE) + xmask = in_index < n_elements + + tmp0 = tl.load(in_ptr0 + (in_index), xmask) + tl.atomic_and(out_ptr0 + (out_index), tmp0, xmask) + + +@triton.jit +def atomic_and_broadcast(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(0) + + x = tl.load(x_ptr) # x is scalar or 1D, no mask needed + + # Compute y indices + y_offset = pid * BLOCK_SIZE + y_indices = y_offset + tl.arange(0, BLOCK_SIZE) + y_mask = y_indices < n_elements + + y_value = tl.load(y_ptr + y_indices, y_mask) + # Atomic or: y &= x (broadcasted) + tl.atomic_and(out_ptr + y_indices, y_value, mask=y_mask) + tl.atomic_and(out_ptr + y_indices, x, mask=y_mask) + + +# 定义不同测试场景的参数组合 (x_shape, y_shape, BLOCK_SIZE) +test_cases = [ + ((1, 1, 1, 1), (1, 1, 1, 4), 4), + ((1, 1, 1, 3), (1, 5, 1, 3), 5), + ((3, ), (2, 3, 3, 3, 3), 81), + ((3, ), (2, 3, 3, 3), 27), + ((3, ), (2, 3, 3), 9), + ((3, ), (2, 3), 3), +] + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d + TestUtils.test_shape1d) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_and(x_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = test_common.generate_tensor(x_shape, x_dtype_str).npu() + # OR的时候任何位和0做OR都不变 任何位和1做AND也都不变，所以为了保持不变 不能用0 只能用1 + y = torch.full(shape, torch.iinfo(x_dtype).max, dtype=x_dtype).npu() + + # 保存副本用于验证 + x_temp = x.clone() + y_temp = y.clone() + + if len(shape) == 2: + n_elements = shape[0] * shape[1] * 2 + atomic_and[shape[0] * 2, 1, 1](x, y, n_elements, BLOCK_SIZE=shape[1], BLOCK_NUM=shape[0]) + elif len(shape) == 1: + n_elements = shape[0] + BLOCK_SIZE = min(1024, shape[0]) # 1024:限制最大线程块大小 + grid_size = (n_elements + BLOCK_SIZE - 1) // BLOCK_SIZE # 向上取整 + aligned_size = grid_size * BLOCK_SIZE + x_concat = torch.full([aligned_size * 2], 0, dtype=x_dtype).npu() + x_concat[0:n_elements] = x[0:n_elements] + x_concat[aligned_size:(aligned_size + n_elements)] = x[n_elements:(n_elements * 2)] + atomic_and[grid_size * 2, 1, 1](x_concat, y, aligned_size * 2, BLOCK_SIZE=BLOCK_SIZE, BLOCK_NUM=grid_size) + + expected = y_temp & x_temp[0:shape[0]] & x_temp[shape[0]:(shape[0] * 2)] + torch.testing.assert_close(y, expected) + + +# 3d +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_and_3d(x_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = test_common.generate_tensor(x_shape, x_dtype_str).npu() + y = torch.full(shape, 0, dtype=x_dtype).npu() + + # 保存副本用于验证 + x_temp = x.clone() + y_temp = y.clone() + + n_elements = shape[0] * shape[1] * shape[2] + atomic_and[2, 1, 1](x, y, n_elements * 2, BLOCK_SIZE=shape[0] * shape[1] * shape[2], BLOCK_NUM=1) + + expected = y_temp & x_temp[0:shape[0]] & x_temp[shape[0]:(shape[0] * 2)] + torch.testing.assert_close(y, expected) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape_ub_overflow) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +@test_common.raises_with_match(triton.compiler.errors.MLIRCompilationError, "ub overflow") +def test_atomic_and_ub_overflow(x_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = test_common.generate_tensor(x_shape, x_dtype_str).npu() + y = torch.full(shape, 0, dtype=x_dtype).npu() + + # 保存副本用于验证 + x_temp = x.clone() + y_temp = y.clone() + + n_elements = shape[0] * shape[1] * shape[2] + atomic_and[2, 1, 1](x, y, n_elements * 2, BLOCK_SIZE=shape[0] * shape[1] * shape[2], BLOCK_NUM=1) + + +@triton.jit +def atomic_and_multi_d(in_ptr0, out_ptr0, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if (YB * ZB * MB * NB) > 1: + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if (ZB * MB * NB) > 1: + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if (MB * NB) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if NB > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + tmp0 = tl.load(in_ptr0 + offsets) + tl.atomic_and(out_ptr0 + offsets, tmp0) + + +# multi_d +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (2, 4, 8, 4), + (8, 4, 2, 4), + (2, 8, 2, 2), + (2, 4, 8, 4, 2), + (8, 4, 2, 4, 4), + (2, 8, 2, 2, 2), +]) +@pytest.mark.parametrize('dtype', filtered_dtype) +def test_atomic_and_4d_5d(dtype, shape): + x0_value = 3 + x0 = torch.full(shape, x0_value, dtype=eval('torch.' + dtype)).npu() + x1 = torch.full(shape, 2, dtype=eval('torch.' + dtype)).npu() + + x1_ref = x1 & x0_value + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + atomic_and_multi_d[(1, )](x0, x1, *triton_shape) + test_common.validate_cmp(dtype, x1, x1_ref) + + +@triton.jit +def atomic_and_5d(x_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr, XB1: tl.constexpr, YB1: tl.constexpr, ZB1: tl.constexpr, MB1: tl.constexpr, + NB1: tl.constexpr): + base = tl.program_id(0) * (XB * YB * ZB * MB * NB) + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + offsets1 = tl.arange(0, XB1) * (YB1 * ZB1 * MB1 * NB1) + offsets1 = offsets1[:, None] + tl.arange(0, YB1)[None, :] * (ZB1 * MB1 * NB1) + offsets1 = offsets1[:, :, None] + tl.arange(0, ZB1)[None, None, :] * (MB1 * NB1) + offsets1 = offsets1[:, :, :, None] + tl.arange(0, MB1)[None, None, None, :] * NB1 + offsets1 = offsets1[:, :, :, :, None] + tl.arange(0, NB1)[None, None, None, None, :] + + based_offsets = offsets + base + + tmp0 = tl.load(x_ptr + based_offsets) + tl.atomic_and(out_ptr + offsets1, tmp0) + + +@pytest.mark.parametrize('param_list', [ + [(1, 1, 2, 1, 1), (1, 1, 2, 1, 2)], +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_and_5d(x_dtype_str, param_list): + x0_shape, y_shape = param_list + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(x0_shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + + out = torch.full(y_shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*x0_shape] + while len(triton_shape) < 5: + triton_shape.append(1) + XB, YB, ZB, MB, NB = triton_shape + + triton_shape1 = [*y_shape] + while len(triton_shape1) < 5: + triton_shape1.append(1) + XB1, YB1, ZB1, MB1, NB1 = triton_shape1 + + atomic_and_5d[(2, )]( + x_ptr=x, + out_ptr=out, + XB=XB, + YB=YB, + ZB=ZB, + MB=MB, + NB=NB, + XB1=XB1, + YB1=YB1, + ZB1=ZB1, + MB1=MB1, + NB1=NB1, + ) + + expected = out_temp & x_temp[0:x0_shape[0]] & x_temp[x0_shape[0]:x_shape[0]] + torch.testing.assert_close(out, expected) + + +@triton.jit +def atomic_and_4d(x_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + XB1: tl.constexpr, YB1: tl.constexpr, ZB1: tl.constexpr, MB1: tl.constexpr): + base = tl.program_id(0) * (XB * YB * ZB * MB) + offsets = tl.arange(0, XB) * (YB * ZB * MB) + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB) + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB) + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] + + offsets1 = tl.arange(0, XB1) * (YB1 * ZB1 * MB1) + offsets1 = offsets1[:, None] + tl.arange(0, YB1)[None, :] * (ZB1 * MB1) + offsets1 = offsets1[:, :, None] + tl.arange(0, ZB1)[None, None, :] * (MB1) + offsets1 = offsets1[:, :, :, None] + tl.arange(0, MB1)[None, None, None, :] + + based_offsets = offsets + base + + tmp0 = tl.load(x_ptr + based_offsets) + tl.atomic_and(out_ptr + offsets1, tmp0) + + +@pytest.mark.parametrize('param_list', [ + [(1, 1, 2, 1), (1, 1, 2, 2)], + [(1, 1, 1, 1), (1, 1, 2, 2)], +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_and_4d(x_dtype_str, param_list): + x0_shape, y_shape = param_list + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(x0_shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + out = torch.full(y_shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*x0_shape] + while len(triton_shape) < 4: + triton_shape.append(1) + XB, YB, ZB, MB = triton_shape + + triton_shape1 = [*y_shape] + while len(triton_shape1) < 4: + triton_shape1.append(1) + XB1, YB1, ZB1, MB1 = triton_shape1 + + atomic_and_4d[(2, )]( + x_ptr=x, + out_ptr=out, + XB=XB, + YB=YB, + ZB=ZB, + MB=MB, + XB1=XB1, + YB1=YB1, + ZB1=ZB1, + MB1=MB1, + ) + + expected = out_temp & x_temp[0:x0_shape[0]] & x_temp[x0_shape[0]:x_shape[0]] + torch.testing.assert_close(out, expected) + + +@triton.jit +def atomic_and_3d(x_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XB1: tl.constexpr, + YB1: tl.constexpr, ZB1: tl.constexpr): + base = tl.program_id(0) * (XB * YB * ZB) + offsets = tl.arange(0, XB) * (YB * ZB) + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB) + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] + + offsets1 = tl.arange(0, XB1) * (YB1 * ZB1) + offsets1 = offsets1[:, None] + tl.arange(0, YB1)[None, :] * (ZB1) + offsets1 = offsets1[:, :, None] + tl.arange(0, ZB1)[None, None, :] + + based_offsets = offsets + base + + tmp0 = tl.load(x_ptr + based_offsets) + tl.atomic_and(out_ptr + offsets1, tmp0) + + +@pytest.mark.parametrize('param_list', [ + [(1, 1, 1), (1, 1, 2)], + [(1, 1, 2), (1, 2, 2)], +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_and_3d_2(x_dtype_str, param_list): + x0_shape, y_shape = param_list + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(x0_shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + out = torch.full(y_shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*x0_shape] + while len(triton_shape) < 3: + triton_shape.append(1) + XB, YB, ZB = triton_shape + + triton_shape1 = [*y_shape] + while len(triton_shape1) < 3: + triton_shape1.append(1) + XB1, YB1, ZB1 = triton_shape1 + + atomic_and_3d[(2, )]( + x_ptr=x, + out_ptr=out, + XB=XB, + YB=YB, + ZB=ZB, + XB1=XB1, + YB1=YB1, + ZB1=ZB1, + ) + + expected = out_temp & x_temp[0:x0_shape[0]] & x_temp[x0_shape[0]:x_shape[0]] + torch.testing.assert_close(out, expected) + + +@triton.jit +def atomic_and_2d(x_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, XB1: tl.constexpr, YB1: tl.constexpr): + base = tl.program_id(0) * (XB * YB) + offsets = tl.arange(0, XB) * (YB) + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] + + offsets1 = tl.arange(0, XB1) * (YB1) + offsets1 = offsets1[:, None] + tl.arange(0, YB1)[None, :] + + based_offsets = offsets + base + + tmp0 = tl.load(x_ptr + based_offsets) + tl.atomic_and(out_ptr + offsets1, tmp0) + + +@pytest.mark.parametrize('param_list', [ + [(1, 2), (2, 2)], + [(1, 1), (2, 2)], +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_and_2d(x_dtype_str, param_list): + x0_shape, y_shape = param_list + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(x0_shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + out = torch.full(y_shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*x0_shape] + while len(triton_shape) < 2: + triton_shape.append(1) + XB, YB = triton_shape + + triton_shape1 = [*y_shape] + while len(triton_shape1) < 2: + triton_shape1.append(1) + XB1, YB1 = triton_shape1 + + atomic_and_2d[(2, )]( + x_ptr=x, + out_ptr=out, + XB=XB, + YB=YB, + XB1=XB1, + YB1=YB1, + ) + + expected = out_temp & x_temp[0:x0_shape[0]] & x_temp[x0_shape[0]:x_shape[0]] + torch.testing.assert_close(out, expected) + + +@triton.jit +def atomic_and(in_ptr0, out_ptr0, n_elements, BLOCK_SIZE: tl.constexpr, BLOCK_NUM: tl.constexpr, + mode: tl.constexpr = 0): + in_offset = tl.program_id(0) * BLOCK_SIZE + out_offset = (tl.program_id(0) % BLOCK_NUM) * BLOCK_SIZE + in_index = in_offset + tl.arange(0, BLOCK_SIZE) + out_index = out_offset + tl.arange(0, BLOCK_SIZE) + xmask = in_index < n_elements + + tmp0 = tl.load(in_ptr0 + (in_index), xmask) + if mode == 0: + tl.atomic_and(out_ptr0 + (out_index), tmp0, xmask, 'acq_rel', 'cta') + elif mode == 1: + tl.atomic_and(out_ptr0 + (out_index), tmp0, xmask, "test") + elif mode == 2: + tl.atomic_and(out_ptr0 + (out_index), tmp0, xmask, "acq_rel", "test") + + +invalid_types_float = ['float16', 'float32', 'bfloat16'] + + +@pytest.mark.parametrize("sigtype", invalid_types_float) +@test_common.raises_with_match(triton.compiler.errors.MLIRCompilationError, "must be signless-integer-like") +def test_invalid_types_float(sigtype): + N = 32 + x = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + y = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + + atomic_and[1, 1, 1](x, y, 1, 1, 32) + + +default_types = ['int8'] + + +@pytest.mark.parametrize("sigtype", default_types) +@pytest.mark.parametrize("test_type", ["sem", "scope"]) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "Memory semantic test not supported") +def test_invalid_sem_scope(sigtype, test_type): + N = 32 + x = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + y = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + + if test_type == "sem": + atomic_and[1, 1, 1](x, y, 1, 1, 32, 1) + elif test_type == "scope": + atomic_and[1, 1, 1](x, y, 1, 1, 32, 2) + + +@triton.jit +def _atomic_and_ss(in_ptr, out_ptr, n_cols, BLOCK_SIZE: tl.constexpr, SEM: tl.constexpr, SCOPE: tl.constexpr): + pid = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = pid < n_cols + val = tl.load(in_ptr + pid, mask) + tl.atomic_and(out_ptr + pid, val, mask, sem=SEM, scope=SCOPE) + + +SEMS = ("relaxed", "acquire", "release", "acq_rel") +SCOPES = ("cta", "gpu", "sys") + + +@pytest.mark.parametrize("sem", SEMS) +@pytest.mark.parametrize("scope", SCOPES) +def test_atomic_sem_vs_scope(sem: str, scope: str): + n_cols = 1024 + BLOCK = 128 + grid = (triton.cdiv(n_cols, BLOCK), ) + + inp = torch.full((n_cols, ), 0xFF, dtype=torch.int32, device="npu") + + base = torch.full_like(inp, 0xFF) + _atomic_and_ss[grid](inp, base, n_cols, BLOCK_SIZE=BLOCK, SEM="acq_rel", SCOPE="gpu") + + cur = torch.full_like(inp, 0xFF) + _atomic_and_ss[grid](inp, cur, n_cols, BLOCK_SIZE=BLOCK, SEM=sem, SCOPE=scope) + + torch.testing.assert_close(cur, base) + + +@pytest.mark.parametrize('param_list', [ + ['uint8', (32, 32), 2], + ['uint16', (32, 32), 2], + ['uint32', (32, 32), 2], + ['uint64', (32, 32), 2], +]) +def test_atomic_and_uint(param_list): + dtype, shape, ncore = param_list + block_size = shape[0] * shape[1] // ncore + split_size = shape[0] // ncore + + val_cpu = torch.randint(low=0, high=10, size=shape, dtype=eval(f'torch.{dtype}')).cpu() + val = val_cpu.to("npu") + + pointer_cpu = torch.randint(low=0, high=10, size=(split_size, shape[1]), dtype=eval(f'torch.{dtype}')).cpu() + pointer = pointer_cpu.to("npu") + pointer_old_cpu = torch.full_like(pointer_cpu, -10).cpu() + pointer_old = pointer_old_cpu.to("npu") + pointer_ref_cpu = pointer_cpu.clone() + + for i in range(ncore - 1): + pointer_ref_cpu &= val_cpu[(i * split_size):((i + 1) * split_size)] + + pointer_ref_last = pointer_ref_cpu.clone() + pointer_ref_cpu &= val_cpu[((ncore - 1) * split_size):(ncore * split_size)] + pointer_ref = pointer_ref_cpu.to("npu") + + @triton.jit + def atomic_and_uint(in_ptr0, out_ptr0, out_ptr1, n_elements, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + yindex = tl.arange(0, BLOCK_SIZE)[:] + xmask = xindex < n_elements + x0 = xindex + x1 = yindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.atomic_and(out_ptr0 + (x1), tmp0, xmask) + tl.store(out_ptr1 + (x1), tmp1, xmask) + + n_elements = shape[0] * shape[1] + atomic_and_uint[ncore, 1, 1](val, pointer, pointer_old, n_elements, BLOCK_SIZE=split_size * shape[1]) + test_common.validate_cmp(dtype, pointer, pointer_ref) diff --git a/third_party/ascend/unittest/generalization_cases/test_atomic_cas.py b/third_party/ascend/unittest/generalization_cases/test_atomic_cas.py new file mode 100644 index 0000000000..eab2568755 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_atomic_cas.py @@ -0,0 +1,484 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import pytest +import torch +import triton + +import triton.language as tl + +import test_common +from test_common import TestUtils +import numpy as np + +filtered_dtype = [dtype for dtype in TestUtils.full_dtype if dtype not in {'uint32', 'bfloat16', 'int8', 'bool'}] + + +@triton.jit +def atomic_cas(in_ptr0, in_ptr1, out_ptr0, n_elements, BLOCK_SIZE: tl.constexpr, BLOCK_NUM: tl.constexpr): + in_offset = tl.program_id(0) * BLOCK_SIZE + out_offset = (tl.program_id(0) % BLOCK_NUM) * BLOCK_SIZE + in_index = in_offset + tl.arange(0, BLOCK_SIZE) + out_index = out_offset + tl.arange(0, BLOCK_SIZE) + xmask = in_index < n_elements + + tmp0 = tl.load(in_ptr0 + (in_index), xmask) + tmp1 = tl.load(in_ptr1 + (in_index), xmask) + tl.atomic_cas(out_ptr0 + (out_index), tmp1, tmp0) + + +@triton.jit +def atomic_cas_ndim(x_ptr, y_ptr, out_ptr, NCORE: tl.constexpr, BLOCK_SIZE: tl.constexpr, DIM0: tl.constexpr, + DIM1: tl.constexpr, DIM2: tl.constexpr, DIM3: tl.constexpr, DIM4: tl.constexpr): + sub_idx = tl.program_id(1) + base_src = tl.program_id(0) * DIM4 + sub_idx * BLOCK_SIZE + base_dst = (tl.program_id(0) % (DIM0 * DIM1 * DIM2 * DIM3)) * DIM4 + sub_idx * BLOCK_SIZE + offsets_src = tl.arange(0, BLOCK_SIZE) + base_src + offsets_dst = tl.arange(0, BLOCK_SIZE) + base_dst + mask = tl.arange(0, BLOCK_SIZE) + sub_idx * BLOCK_SIZE < DIM4 + tmp = tl.load(x_ptr + offsets_src, mask) + tmp_c = tl.load(y_ptr + offsets_src, mask) + tl.atomic_cas(out_ptr + offsets_dst, tmp_c, tmp) + + +@triton.jit +def atomic_cas_broadcast(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(0) + + x = tl.load(x_ptr) # x is scalar or 1D, no mask needed + + # Compute y indices + y_offset = pid * BLOCK_SIZE + y_indices = y_offset + tl.arange(0, BLOCK_SIZE) + y_mask = y_indices < n_elements + + y_value = tl.load(y_ptr + y_indices, y_mask) + # Atomic or: y |= x (broadcasted) + tl.atomic_cas(out_ptr + y_indices, y_value, mask=y_mask) + tl.atomic_cas(out_ptr + y_indices, x, mask=y_mask) + + +# 定义不同测试场景的参数组合 (x_shape, y_shape, BLOCK_SIZE) +test_cases = [ + ((1, 1, 1, 1), (1, 1, 1, 4), 4), + ((1, 1, 1, 3), (1, 5, 1, 3), 5), + ((3, ), (2, 3, 3, 3, 3), 81), + ((3, ), (2, 3, 3, 3), 27), + ((3, ), (2, 3, 3), 9), + ((3, ), (2, 3), 3), +] + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d + TestUtils.test_shape1d) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_cas(x_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = test_common.generate_tensor(x_shape, x_dtype_str).npu() + c = torch.randint(low=0, high=2, size=x_shape, dtype=x_dtype).npu() + y = torch.randint(low=0, high=2, size=shape, dtype=x_dtype).npu() + + # 保存副本用于验证 + x_temp = x.clone() + c_temp = c.clone() + y_temp = y.clone() + + if len(shape) == 2: + n_elements = shape[0] * shape[1] * 2 + atomic_cas[shape[0] * 2, 1, 1](x, c, y, n_elements, BLOCK_SIZE=shape[1], BLOCK_NUM=shape[0]) + elif len(shape) == 1: + n_elements = shape[0] + BLOCK_SIZE = min(1024, shape[0]) # 1024:限制最大线程块大小 + grid_size = (n_elements + BLOCK_SIZE - 1) // BLOCK_SIZE # 向上取整 + aligned_size = grid_size * BLOCK_SIZE + # value + x_concat = torch.full([aligned_size * 2], 0, dtype=x_dtype).npu() + x_concat[0:n_elements] = x[0:n_elements] + x_concat[aligned_size:(aligned_size + n_elements)] = x[n_elements:(n_elements * 2)] + # compare + c_concat = torch.full([aligned_size * 2], 0, dtype=x_dtype).npu() + c_concat[0:n_elements] = c[0:n_elements] + c_concat[aligned_size:(aligned_size + n_elements)] = c[n_elements:(n_elements * 2)] + atomic_cas[grid_size * 2, 1, 1](x_concat, c_concat, y, aligned_size * 2, BLOCK_SIZE=BLOCK_SIZE, + BLOCK_NUM=grid_size) + + expected = torch.where(y_temp == c_temp[0:shape[0]], x_temp[0:shape[0]], y_temp) + expected = torch.where(expected == c_temp[shape[0]:(shape[0] * 2)], x_temp[shape[0]:(shape[0] * 2)], expected) + torch.testing.assert_close(y, expected) + + +# 3d +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_cas_3d(x_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = test_common.generate_tensor(x_shape, x_dtype_str).npu() + c = torch.randint(low=3, high=5, size=x_shape, dtype=x_dtype).npu() + y = torch.randint(low=3, high=5, size=shape, dtype=x_dtype).npu() + + # 保存副本用于验证 + x_temp = x.clone() + c_temp = c.clone() + y_temp = y.clone() + + n_elements = shape[0] * shape[1] * shape[2] + atomic_cas[2, 1, 1](x, c, y, n_elements * 2, BLOCK_SIZE=shape[0] * shape[1] * shape[2], BLOCK_NUM=1) + + expected = torch.where(y_temp == c_temp[0:shape[0]], x_temp[0:shape[0]], y_temp) + expected = torch.where(expected == c_temp[shape[0]:(shape[0] * 2)], x_temp[shape[0]:(shape[0] * 2)], expected) + torch.testing.assert_close(y, expected) + + +@triton.jit +def atomic_cas_multi_d(in_ptr0, in_ptr1, out_ptr0, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, + MB: tl.constexpr, NB: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if (YB * ZB * MB * NB) > 1: + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if (ZB * MB * NB) > 1: + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if (MB * NB) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if NB > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + tmp0 = tl.load(in_ptr0 + offsets) + tmp1 = tl.load(in_ptr1 + offsets) + tl.atomic_cas(out_ptr0 + offsets, tmp1, tmp0) + + +# multi_d +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (2, 4, 8, 4), + (8, 4, 2, 4), + (2, 8, 2, 2), + (2, 4, 8, 4, 2), + (8, 4, 2, 4, 4), + (2, 8, 2, 2, 2), +]) +@pytest.mark.parametrize('dtype', filtered_dtype) +def test_atomic_cas_4d_5d(dtype, shape): + x0_value = 3 + x0 = torch.full(shape, x0_value, dtype=eval('torch.' + dtype)).npu() + c = torch.randint(low=2, high=4, size=shape, dtype=eval('torch.' + dtype)).npu() + x1 = torch.randint(low=2, high=4, size=shape, dtype=eval('torch.' + dtype)).npu() + + x1_ref = torch.where(x1 == c, 3, x1) + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + + atomic_cas_multi_d[(1, )](x0, c, x1, *triton_shape) + test_common.validate_cmp(dtype, x1, x1_ref) + + +@pytest.mark.parametrize('shape', [ + (1, 1, 1, 1, 2), + (10, 1, 15, 1, 7), + (1, 1, 1, 1, 257), +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_cas_5d(x_dtype_str, shape): + shape = shape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + + c = torch.randint(low=3, high=5, size=x_shape, dtype=x_dtype).npu() + out = torch.randint(low=3, high=5, size=shape, dtype=x_dtype).npu() + + x_temp = x.clone() + c_temp = c.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + XB, YB, ZB, MB, NB = triton_shape + BLOCK_SIZE = 256 + ncore = (NB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_cas_ndim[(2 * XB * YB * ZB * MB, ncore)]( + x_ptr=x, + y_ptr=c, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=XB, + DIM1=YB, + DIM2=ZB, + DIM3=MB, + DIM4=NB, + ) + + expected = torch.where(out_temp == c_temp[0:shape[0]], x_temp[0:shape[0]], out_temp) + expected = torch.where(expected == c_temp[shape[0]:(x_shape[0])], x_temp[shape[0]:(x_shape[0])], expected) + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('shape', [ + (1, 1, 1, 1), + (1, 1, 2, 2), + (1, 3, 2, 7), + (1, 3, 2, 651), +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_cas_4d(x_dtype_str, shape): + shape = shape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + c = torch.randint(low=3, high=5, size=x_shape, dtype=x_dtype).npu() + out = torch.randint(low=3, high=5, size=shape, dtype=x_dtype).npu() + + x_temp = x.clone() + c_temp = c.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 4: + triton_shape.append(1) + XB, YB, ZB, MB = triton_shape + + BLOCK_SIZE = 256 + ncore = (MB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_cas_ndim[(2 * XB * YB * ZB, ncore)]( + x_ptr=x, + y_ptr=c, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=1, + DIM1=XB, + DIM2=YB, + DIM3=ZB, + DIM4=MB, + ) + + expected = torch.where(out_temp == c_temp[0:shape[0]], x_temp[0:shape[0]], out_temp) + expected = torch.where(expected == c_temp[shape[0]:(x_shape[0])], x_temp[shape[0]:(x_shape[0])], expected) + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('shape', [ + (1, 1, 1), + (1, 1, 2), + (1, 31, 275), +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_cas_3d_2(x_dtype_str, shape): + shape = shape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + c = torch.randint(low=3, high=5, size=x_shape, dtype=x_dtype).npu() + out = torch.randint(low=3, high=5, size=shape, dtype=x_dtype).npu() + + x_temp = x.clone() + c_temp = c.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 3: + triton_shape.append(1) + XB, YB, ZB = triton_shape + BLOCK_SIZE = 256 + ncore = (ZB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_cas_ndim[(2 * XB * YB, ncore)]( + x_ptr=x, + y_ptr=c, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=1, + DIM1=1, + DIM2=XB, + DIM3=YB, + DIM4=ZB, + ) + + expected = torch.where(out_temp == c_temp[0:shape[0]], x_temp[0:shape[0]], out_temp) + expected = torch.where(expected == c_temp[shape[0]:(x_shape[0])], x_temp[shape[0]:(x_shape[0])], expected) + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('shape', [ + (1, 2), + (1, 1), + (257, 1), + (257, 2), +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_cas_2d(x_dtype_str, shape): + shape = shape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + c = torch.randint(low=3, high=5, size=x_shape, dtype=x_dtype).npu() + out = torch.randint(low=3, high=5, size=shape, dtype=x_dtype).npu() + + x_temp = x.clone() + c_temp = c.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 2: + triton_shape.append(1) + XB, YB = triton_shape + BLOCK_SIZE = 256 + ncore = (YB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_cas_ndim[(2 * XB, ncore)]( + x_ptr=x, + y_ptr=c, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=1, + DIM1=1, + DIM2=1, + DIM3=XB, + DIM4=YB, + ) + + expected = torch.where(out_temp == c_temp[0:shape[0]], x_temp[0:shape[0]], out_temp) + expected = torch.where(expected == c_temp[shape[0]:(x_shape[0])], x_temp[shape[0]:(x_shape[0])], expected) + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('shape', [(1, ), (9, ), (256, ), (257, ), (65535, ), (65536, )]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_cas_1d(x_dtype_str, shape): + shape = shape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + c = torch.randint(low=3, high=5, size=x_shape, dtype=x_dtype).npu() + out = torch.full(shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + c_temp = c.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 2: + triton_shape.append(1) + XB = triton_shape[0] + BLOCK_SIZE = 256 + ncore = (XB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_cas_ndim[(2, ncore)]( + x_ptr=x, + y_ptr=c, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=1, + DIM1=1, + DIM2=1, + DIM3=1, + DIM4=XB, + ) + + expected = torch.where(out_temp == c_temp[0:shape[0]], x_temp[0:shape[0]], out_temp) + expected = torch.where(expected == c_temp[shape[0]:(x_shape[0])], x_temp[shape[0]:(x_shape[0])], expected) + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('param_list', [ + ['uint16', (32, 32), 2], + ['uint32', (32, 32), 2], + ['uint64', (32, 32), 2], +]) +def test_atomic_cas_uint(param_list): + dtype, shape, ncore = param_list + block_size = shape[0] * shape[1] // ncore + split_size = shape[0] // ncore + + import random + cmp_val = [random.randint(0, 10) for _ in range(ncore)] + + cmp_cpu_parts = [] + for i in range(ncore): + part = torch.ones(split_size, shape[1], dtype=eval(f'torch.{dtype}')) * cmp_val[i] + cmp_cpu_parts.append(part) + cmp_cpu = torch.cat(cmp_cpu_parts, dim=0) + cmp = cmp_cpu.to("npu") + + val_cpu = torch.randint(low=0, high=10, size=shape, dtype=eval(f'torch.{dtype}')).cpu() + val = val_cpu.to("npu") + + pointer_cpu = torch.randint(low=0, high=10, size=(split_size, shape[1]), dtype=eval(f'torch.{dtype}')).cpu() + pointer = pointer_cpu.to("npu") + pointer_old_cpu = torch.full_like(pointer_cpu, -10).cpu() + pointer_old = pointer_old_cpu.to("npu") + pointer_ref_cpu = pointer_cpu.clone() + + pointer_ref_np = pointer_cpu.numpy() + val_np = val_cpu.numpy() + for i in range(ncore): + val_subview_np = val_np[(i * split_size):((i + 1) * split_size)] + pointer_ref_np = np.where(pointer_ref_np == cmp_val[i], val_subview_np, pointer_ref_np) + pointer_ref_cpu = torch.from_numpy(pointer_ref_np) + pointer_ref = pointer_ref_cpu.to("npu") + + @triton.jit + def atomic_cas_uint(in_ptr0, in_ptr1, out_ptr0, out_ptr1, n_elements, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + yindex = tl.arange(0, BLOCK_SIZE)[:] + xmask = xindex < n_elements + x0 = xindex + x1 = yindex + val = tl.load(in_ptr0 + (x0), xmask) + cmp = tl.load(in_ptr1 + (x0), xmask) + tmp1 = tl.atomic_cas(out_ptr0 + (x1), cmp, val) + tl.store(out_ptr1 + (x1), tmp1, xmask) + + n_elements = shape[0] * shape[1] + atomic_cas_uint[ncore, 1, 1](val, cmp, pointer, pointer_old, n_elements, BLOCK_SIZE=split_size * shape[1]) + test_common.validate_cmp(dtype, pointer, pointer_ref) diff --git a/third_party/ascend/unittest/generalization_cases/test_atomic_max.py b/third_party/ascend/unittest/generalization_cases/test_atomic_max.py new file mode 100644 index 0000000000..87347cb2fd --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_atomic_max.py @@ -0,0 +1,258 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import random +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils + + +@triton.jit +def triton_test_fn_atomic_max_dma(in_ptr0, in_ptr1, out_ptr1, n_elements: tl.constexpr, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + index = xoffset + tl.arange(0, BLOCK_SIZE)[:] + mask = index < n_elements + inp0 = tl.load(in_ptr0 + (index), mask) + inp1 = tl.load(in_ptr1 + (index), mask) + tmp1 = tl.atomic_max(out_ptr1 + (index), inp0, mask) + tmp2 = tl.atomic_max(out_ptr1 + (index), inp1, mask) + + +def promote_dtype(x_dtype, y_dtype): + """ + 如果 y 的精度低于 x, 则提升 y 的精度以匹配 x。 + """ + # 如果两个数据类型一致，直接返回 + if x_dtype == y_dtype: + return y_dtype + + # 构建类型的优先级列表（从低到高） + priority = [torch.int8, torch.int16, torch.int32, torch.float16, torch.bfloat16, torch.float32] + + # 查找两种类型在优先级列表中的位置 + x_priority = priority.index(x_dtype) + y_priority = priority.index(y_dtype) + + # 如果y的优先级比x小，则提升到x的类型 + if y_priority < x_priority: + return x_dtype + else: + return y_dtype + + +# torch.max do not support int +@pytest.mark.parametrize('shape', random.sample(TestUtils.test_shape2d + TestUtils.test_shape1d, 5)) +@pytest.mark.parametrize('x_dtype_str', ['float32', 'int32', 'int8', 'int16', 'bfloat16', 'float16']) +@pytest.mark.parametrize('y_dtype_str', ['float32', 'int32', 'int8', 'int16', 'bfloat16', 'float16']) +def test_atomic_max(x_dtype_str, y_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + y_dtype = eval('torch.' + y_dtype_str) + + x0 = test_common.generate_tensor(shape, x_dtype_str) + x1 = test_common.generate_tensor(shape, y_dtype_str) + out_dtype = promote_dtype(x_dtype, y_dtype) + if out_dtype == torch.bfloat16: + out_dtype = torch.float32 + out = torch.full(x1.shape, 0, dtype=out_dtype) + + out_ref = torch.maximum(out, x0) + out_ref = torch.maximum(out_ref, x1) + out_ref = out_ref.npu() + x0 = x0.npu() + x1 = x1.npu() + out = out.npu() + + if len(shape) == 2: + n_elements = shape[0] * shape[1] + triton_test_fn_atomic_max_dma[shape[0], 1, 1](x0, x1, out, n_elements, BLOCK_SIZE=shape[1]) + elif len(shape) == 1: + n_elements = shape[0] + BLOCK_SIZE = min(1024, shape[0]) # 1024:限制最大线程块大小 + grid_size = (n_elements + BLOCK_SIZE - 1) // BLOCK_SIZE # 向上取整 + triton_test_fn_atomic_max_dma[grid_size, 1, 1](x0, x1, out, n_elements, BLOCK_SIZE=BLOCK_SIZE) + + torch.testing.assert_close(out, out_ref) + + +# 3d +testlist = [ + (1, 22, 39), + (27, 1, 39), + (27, 22, 1), + (1, 1, 23), + (23, 1, 1), + (1, 23, 1), + (27, 5, 3), + (2, 29, 4), + (7, 31, 7), + (3, 5, 8), + (7, 17, 15), + (25, 5, 16), + (23, 5, 31), + (7, 11, 32), + (7, 11, 33), + (2, 3, 255), + (3, 3, 256), + (3, 2, 257), +] + + +@pytest.mark.parametrize('shape', random.sample(testlist, 5)) +@pytest.mark.parametrize('x_dtype_str', ['float32', 'int32', 'int8', 'int16', 'bfloat16', 'float16']) +@pytest.mark.parametrize('y_dtype_str', ['float32', 'int32', 'int8', 'int16', 'bfloat16', 'float16']) +def test_atomic_max_3d(x_dtype_str, y_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + y_dtype = eval('torch.' + y_dtype_str) + + ncore = 1 + split_size = shape[0] // ncore + x0 = test_common.generate_tensor(shape, x_dtype_str) + x1 = test_common.generate_tensor(shape, y_dtype_str) + out_dtype = promote_dtype(x_dtype, y_dtype) + if out_dtype == torch.bfloat16: + out_dtype = torch.float32 + y = torch.full(shape, 0, dtype=out_dtype) + + out_ref = torch.full_like(x0, 0, dtype=out_dtype) + out_ref = torch.maximum(out_ref, x0) + out_ref = torch.maximum(out_ref, x1) + x0 = x0.npu() + x1 = x1.npu() + y = y.npu() + + n_elements = shape[0] * shape[1] * shape[2] + triton_test_fn_atomic_max_dma[ncore, 1, 1](x0, x1, y, n_elements, BLOCK_SIZE=split_size * shape[1] * shape[2]) + y = y.cpu() + torch.testing.assert_close(y, out_ref) + + +@triton.jit +def atomic_max_multi_d(in_ptr0, out_ptr0, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if (YB * ZB * MB * NB) > 1: + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if (ZB * MB * NB) > 1: + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if (MB * NB) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if NB > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + tmp0 = tl.load(in_ptr0 + offsets) + tl.atomic_max(out_ptr0 + offsets, tmp0) + + +filtered_dtype = [dtype for dtype in TestUtils.full_dtype if dtype not in {'int64', 'bool'}] + + +# multi_d +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (2, 4, 8, 4), + (8, 4, 2, 4), + (2, 8, 2, 2), + (2, 4, 8, 4, 2), + (8, 4, 2, 4, 4), + (2, 8, 2, 2, 2), +]) +@pytest.mark.parametrize('dtype', filtered_dtype) +def test_atomic_max_4d_5d(dtype, shape): + x0_value = 3 + x0 = torch.full(shape, x0_value, dtype=eval('torch.' + dtype)).npu() + x1 = torch.full(shape, 2, dtype=eval('torch.' + dtype)).npu() + + x1_ref = torch.maximum(x1, x0) + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + atomic_max_multi_d[(1, )](x0, x1, *triton_shape) + test_common.validate_cmp(dtype, x1, x1_ref) + + +@triton.jit +def atomic_max_multi_d_2(in_ptr0, out_ptr0, out_ptr1, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, + MB: tl.constexpr, NB: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if (YB * ZB * MB * NB) > 1: + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if (ZB * MB * NB) > 1: + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if (MB * NB) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if NB > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + tmp0 = tl.load(in_ptr0 + offsets) + tmp1 = tl.load(out_ptr0 + offsets) + tl.atomic_max(out_ptr1 + offsets, tmp0) + tl.atomic_max(out_ptr1 + offsets, tmp1) + + +# multi_d +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (2, 4, 8, 4), + (8, 4, 2, 4), + (2, 8, 2, 2), + (2, 4, 8, 4, 2), + (8, 4, 2, 4, 4), + (2, 8, 2, 2, 2), +]) +@pytest.mark.parametrize('x_dtype_str', ['float32', 'int32', 'int8', 'int16', 'bfloat16', 'float16']) +@pytest.mark.parametrize('y_dtype_str', ['float32', 'int32', 'int8', 'int16', 'bfloat16', 'float16']) +def test_atomic_max_4d_5d_2(x_dtype_str, y_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + y_dtype = eval('torch.' + y_dtype_str) + + if x_dtype == torch.int8 or x_dtype == torch.int16 or x_dtype == torch.int32: + x0 = torch.randint(low=0, high=100, size=shape, dtype=x_dtype).npu() + else: + x0 = torch.randn(shape, dtype=eval('torch.' + x_dtype_str)).npu() + + if y_dtype == torch.int8 or y_dtype == torch.int16 or y_dtype == torch.int32: + x1 = torch.randint(low=0, high=100, size=shape, dtype=y_dtype).npu() + else: + x1 = torch.randn(shape, dtype=eval('torch.' + y_dtype_str)).npu() + + out_dtype = promote_dtype(x_dtype, y_dtype) + if out_dtype == torch.bfloat16: + out_dtype = torch.float32 + if out_dtype == torch.int8 or out_dtype == torch.int16 or out_dtype == torch.int32: + y = torch.full(shape, torch.iinfo(out_dtype).min, dtype=out_dtype).npu() + else: + y = torch.full(shape, float('-inf'), dtype=out_dtype).npu() + + y_tmp = y + x1_ref = torch.maximum(y_tmp, x0) + x1_ref = torch.maximum(x1_ref, x1) + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + atomic_max_multi_d_2[(1, )](x0, x1, y, *triton_shape) + torch.testing.assert_close(y, x1_ref) diff --git a/third_party/ascend/unittest/generalization_cases/test_atomic_min.py b/third_party/ascend/unittest/generalization_cases/test_atomic_min.py new file mode 100644 index 0000000000..a74e99058f --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_atomic_min.py @@ -0,0 +1,264 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import random +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils + + +@triton.jit +def triton_test_fn_atomic_min_dma(in_ptr0, in_ptr1, out_ptr1, n_elements: tl.constexpr, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + index = xoffset + tl.arange(0, BLOCK_SIZE)[:] + mask = index < n_elements + inp0 = tl.load(in_ptr0 + (index), mask) + inp1 = tl.load(in_ptr1 + (index), mask) + tmp1 = tl.atomic_min(out_ptr1 + (index), inp0, mask) + tmp2 = tl.atomic_min(out_ptr1 + (index), inp1, mask) + + +def promote_dtype(x_dtype, y_dtype): + """ + 如果 y 的精度低于 x, 则提升 y 的精度以匹配 x。 + """ + # 如果两个数据类型一致，直接返回 + if x_dtype == y_dtype: + return y_dtype + + # 构建类型的优先级列表（从低到高） + priority = [torch.int8, torch.int16, torch.int32, torch.float16, torch.bfloat16, torch.float32] + + # 查找两种类型在优先级列表中的位置 + x_priority = priority.index(x_dtype) + y_priority = priority.index(y_dtype) + + # 如果y的优先级比x小，则提升到x的类型 + if y_priority < x_priority: + return x_dtype + else: + return y_dtype + + +# torch.min do not support int +@pytest.mark.parametrize('shape', random.sample(TestUtils.test_shape2d + TestUtils.test_shape1d, 5)) +@pytest.mark.parametrize('x_dtype_str', ['float32', 'int32', 'int8', 'int16', 'bfloat16', 'float16']) +@pytest.mark.parametrize('y_dtype_str', ['float32', 'int32', 'int8', 'int16', 'bfloat16', 'float16']) +def test_atomic_min(x_dtype_str, y_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + y_dtype = eval('torch.' + y_dtype_str) + + x0 = test_common.generate_tensor(shape, x_dtype_str) + x1 = test_common.generate_tensor(shape, y_dtype_str) + out_dtype = promote_dtype(x_dtype, y_dtype) + if out_dtype == torch.bfloat16: + out_dtype = torch.float32 + if out_dtype == torch.int8 or out_dtype == torch.int16 or out_dtype == torch.int32: # 判断是否是整数类型 + out = torch.full(x1.shape, torch.iinfo(out_dtype).max, dtype=out_dtype) + else: + out = torch.full(x1.shape, torch.finfo(out_dtype).max, dtype=out_dtype) + + out_ref = torch.minimum(out, x0) + out_ref = torch.minimum(out_ref, x1) + out_ref = out_ref.npu() + x0 = x0.npu() + x1 = x1.npu() + out = out.npu() + + if len(shape) == 2: + n_elements = shape[0] * shape[1] + triton_test_fn_atomic_min_dma[shape[0], 1, 1](x0, x1, out, n_elements, BLOCK_SIZE=shape[1]) + elif len(shape) == 1: + n_elements = shape[0] + BLOCK_SIZE = min(1024, shape[0]) # 1024:限制最大线程块大小 + grid_size = (n_elements + BLOCK_SIZE - 1) // BLOCK_SIZE # 向上取整 + triton_test_fn_atomic_min_dma[grid_size, 1, 1](x0, x1, out, n_elements, BLOCK_SIZE=BLOCK_SIZE) + + torch.testing.assert_close(out, out_ref) + + +# 3d +testlist = [ + (1, 22, 39), + (27, 1, 39), + (27, 22, 1), + (1, 1, 23), + (23, 1, 1), + (1, 23, 1), + (27, 5, 3), + (2, 29, 4), + (7, 31, 7), + (3, 5, 8), + (7, 17, 15), + (25, 5, 16), + (23, 5, 31), + (7, 11, 32), + (7, 11, 33), + (2, 3, 255), + (3, 3, 256), + (3, 2, 257), +] + + +@pytest.mark.parametrize('shape', random.sample(testlist, 5)) +@pytest.mark.parametrize('x_dtype_str', ['float32', 'int32', 'int8', 'int16', 'bfloat16', 'float16']) +@pytest.mark.parametrize('y_dtype_str', ['float32', 'int32', 'int8', 'int16', 'bfloat16', 'float16']) +def test_atomic_min_3d(x_dtype_str, y_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + y_dtype = eval('torch.' + y_dtype_str) + + ncore = 1 + split_size = shape[0] // ncore + x0 = test_common.generate_tensor(shape, x_dtype_str) + x1 = test_common.generate_tensor(shape, y_dtype_str) + out_dtype = promote_dtype(x_dtype, y_dtype) + if out_dtype == torch.bfloat16: + out_dtype = torch.float32 + if out_dtype == torch.int8 or out_dtype == torch.int16 or out_dtype == torch.int32: + y = torch.full(shape, torch.iinfo(out_dtype).max, dtype=out_dtype) + else: + y = torch.full(shape, float('inf'), dtype=out_dtype) + + y_tmp = y + x1_ref = torch.minimum(y_tmp, x0) + x1_ref = torch.minimum(x1_ref, x1) + x0 = x0.npu() + x1 = x1.npu() + y = y.npu() + + n_elements = shape[0] * shape[1] * shape[2] + triton_test_fn_atomic_min_dma[ncore, 1, 1](x0, x1, y, n_elements, BLOCK_SIZE=split_size * shape[1] * shape[2]) + y = y.cpu() + torch.testing.assert_close(y, x1_ref) + + +@triton.jit +def atomic_min_multi_d(in_ptr0, out_ptr0, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if (YB * ZB * MB * NB) > 1: + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if (ZB * MB * NB) > 1: + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if (MB * NB) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if NB > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + tmp0 = tl.load(in_ptr0 + offsets) + tl.atomic_min(out_ptr0 + offsets, tmp0) + + +filtered_dtype = [dtype for dtype in TestUtils.full_dtype if dtype not in {'int64', 'bool'}] + + +# multi_d +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (2, 4, 8, 4), + (8, 4, 2, 4), + (2, 8, 2, 2), + (2, 4, 8, 4, 2), + (8, 4, 2, 4, 4), + (2, 8, 2, 2, 2), +]) +@pytest.mark.parametrize('dtype', filtered_dtype) +def test_atomic_min_4d_5d(dtype, shape): + x0_value = 1 + x0 = torch.full(shape, x0_value, dtype=eval('torch.' + dtype)).npu() + x1 = torch.full(shape, 2, dtype=eval('torch.' + dtype)).npu() + + x1_ref = torch.minimum(x1, x0) + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + atomic_min_multi_d[(1, )](x0, x1, *triton_shape) + test_common.validate_cmp(dtype, x1, x1_ref) + + +@triton.jit +def atomic_min_multi_d_2(in_ptr0, out_ptr0, out_ptr1, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, + MB: tl.constexpr, NB: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if (YB * ZB * MB * NB) > 1: + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if (ZB * MB * NB) > 1: + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if (MB * NB) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if NB > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + tmp0 = tl.load(in_ptr0 + offsets) + tmp1 = tl.load(out_ptr0 + offsets) + tl.atomic_min(out_ptr1 + offsets, tmp0) + tl.atomic_min(out_ptr1 + offsets, tmp1) + + +# multi_d +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (2, 4, 8, 4), + (8, 4, 2, 4), + (2, 8, 2, 2), + (2, 4, 8, 4, 2), + (8, 4, 2, 4, 4), + (2, 8, 2, 2, 2), +]) +@pytest.mark.parametrize('x_dtype_str', ['float32', 'int32', 'int8', 'int16', 'bfloat16', 'float16']) +@pytest.mark.parametrize('y_dtype_str', ['float32', 'int32', 'int8', 'int16', 'bfloat16', 'float16']) +def test_atomic_min_4d_5d_2(x_dtype_str, y_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + y_dtype = eval('torch.' + y_dtype_str) + + if x_dtype == torch.int8 or x_dtype == torch.int16 or x_dtype == torch.int32: + x0 = torch.randint(low=0, high=100, size=shape, dtype=x_dtype).npu() + else: + x0 = torch.randn(shape, dtype=eval('torch.' + x_dtype_str)).npu() + + if y_dtype == torch.int8 or y_dtype == torch.int16 or y_dtype == torch.int32: + x1 = torch.randint(low=0, high=100, size=shape, dtype=y_dtype).npu() + else: + x1 = torch.randn(shape, dtype=eval('torch.' + y_dtype_str)).npu() + + out_dtype = promote_dtype(x_dtype, y_dtype) + if out_dtype == torch.bfloat16: + out_dtype = torch.float32 + if out_dtype == torch.int8 or out_dtype == torch.int16 or out_dtype == torch.int32: + y = torch.full(shape, torch.iinfo(out_dtype).max, dtype=out_dtype).npu() + else: + y = torch.full(shape, float('inf'), dtype=out_dtype).npu() + + y_tmp = y + x1_ref = torch.minimum(y_tmp, x0) + x1_ref = torch.minimum(x1_ref, x1) + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + atomic_min_multi_d_2[(1, )](x0, x1, y, *triton_shape) + torch.testing.assert_close(y, x1_ref) diff --git a/third_party/ascend/unittest/generalization_cases/test_atomic_or.py b/third_party/ascend/unittest/generalization_cases/test_atomic_or.py new file mode 100644 index 0000000000..4e5493b362 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_atomic_or.py @@ -0,0 +1,438 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import pytest +import torch +import triton + +import triton.language as tl + +import test_common +from test_common import TestUtils + +filtered_dtype = [ + dtype for dtype in TestUtils.full_dtype + if dtype not in {'uint32', 'float16', 'float32', 'bfloat16', 'int64', 'bool'} +] + + +@triton.jit +def atomic_or(in_ptr0, out_ptr0, n_elements, BLOCK_SIZE: tl.constexpr, BLOCK_NUM: tl.constexpr): + in_offset = tl.program_id(0) * BLOCK_SIZE + out_offset = (tl.program_id(0) % BLOCK_NUM) * BLOCK_SIZE + in_index = in_offset + tl.arange(0, BLOCK_SIZE) + out_index = out_offset + tl.arange(0, BLOCK_SIZE) + xmask = in_index < n_elements + + tmp0 = tl.load(in_ptr0 + (in_index), xmask) + tl.atomic_or(out_ptr0 + (out_index), tmp0, xmask) + + +@triton.jit +def atomic_or_ndim(x_ptr, out_ptr, NCORE: tl.constexpr, BLOCK_SIZE: tl.constexpr, DIM0: tl.constexpr, + DIM1: tl.constexpr, DIM2: tl.constexpr, DIM3: tl.constexpr, DIM4: tl.constexpr): + sub_idx = tl.program_id(1) + base_src = tl.program_id(0) * DIM4 + sub_idx * BLOCK_SIZE + base_dst = (tl.program_id(0) % (DIM0 * DIM1 * DIM2 * DIM3)) * DIM4 + sub_idx * BLOCK_SIZE + offsets_src = tl.arange(0, BLOCK_SIZE) + base_src + offsets_dst = tl.arange(0, BLOCK_SIZE) + base_dst + mask = tl.arange(0, BLOCK_SIZE) + sub_idx * BLOCK_SIZE < DIM4 + tmp = tl.load(x_ptr + offsets_src, mask) + tl.atomic_or(out_ptr + offsets_dst, tmp, mask) + + +@triton.jit +def atomic_or_broadcast(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(0) + + x = tl.load(x_ptr) # x is scalar or 1D, no mask needed + + # Compute y indices + y_offset = pid * BLOCK_SIZE + y_indices = y_offset + tl.arange(0, BLOCK_SIZE) + y_mask = y_indices < n_elements + + y_value = tl.load(y_ptr + y_indices, y_mask) + # Atomic or: y |= x (broadcasted) + tl.atomic_or(out_ptr + y_indices, y_value, mask=y_mask) + tl.atomic_or(out_ptr + y_indices, x, mask=y_mask) + + +# 定义不同测试场景的参数组合 (x_shape, y_shape, BLOCK_SIZE) +test_cases = [ + ((1, 1, 1, 1), (1, 1, 1, 4), 4), + ((1, 1, 1, 3), (1, 5, 1, 3), 5), + ((3, ), (2, 3, 3, 3, 3), 81), + ((3, ), (2, 3, 3, 3), 27), + ((3, ), (2, 3, 3), 9), + ((3, ), (2, 3), 3), +] + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d + TestUtils.test_shape1d) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_or(x_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = test_common.generate_tensor(x_shape, x_dtype_str).npu() + y = torch.full(shape, 0, dtype=x_dtype).npu() + + # 保存副本用于验证 + x_temp = x.clone() + y_temp = y.clone() + + if len(shape) == 2: + n_elements = shape[0] * shape[1] * 2 + atomic_or[shape[0] * 2, 1, 1](x, y, n_elements, BLOCK_SIZE=shape[1], BLOCK_NUM=shape[0]) + elif len(shape) == 1: + n_elements = shape[0] + BLOCK_SIZE = min(1024, shape[0]) # 1024:限制最大线程块大小 + grid_size = (n_elements + BLOCK_SIZE - 1) // BLOCK_SIZE # 向上取整 + aligned_size = grid_size * BLOCK_SIZE + x_concat = torch.full([aligned_size * 2], 0, dtype=x_dtype).npu() + x_concat[0:n_elements] = x[0:n_elements] + x_concat[aligned_size:(aligned_size + n_elements)] = x[n_elements:(n_elements * 2)] + atomic_or[grid_size * 2, 1, 1](x_concat, y, aligned_size * 2, BLOCK_SIZE=BLOCK_SIZE, BLOCK_NUM=grid_size) + + expected = y_temp | x_temp[0:shape[0]] | x_temp[shape[0]:(shape[0] * 2)] + torch.testing.assert_close(y, expected) + + +# 3d +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_or_3d(x_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = test_common.generate_tensor(x_shape, x_dtype_str).npu() + y = torch.full(shape, 0, dtype=x_dtype).npu() + + # 保存副本用于验证 + x_temp = x.clone() + y_temp = y.clone() + + n_elements = shape[0] * shape[1] * shape[2] + atomic_or[2, 1, 1](x, y, n_elements * 2, BLOCK_SIZE=shape[0] * shape[1] * shape[2], BLOCK_NUM=1) + + expected = y_temp | x_temp[0:shape[0]] | x_temp[shape[0]:(shape[0] * 2)] + torch.testing.assert_close(y, expected) + + +@triton.jit +def atomic_or_multi_d(in_ptr0, out_ptr0, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if (YB * ZB * MB * NB) > 1: + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if (ZB * MB * NB) > 1: + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if (MB * NB) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if NB > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + tmp0 = tl.load(in_ptr0 + offsets) + tl.atomic_or(out_ptr0 + offsets, tmp0) + + +# multi_d +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (2, 4, 8, 4), + (8, 4, 2, 4), + (2, 8, 2, 2), + (2, 4, 8, 4, 2), + (8, 4, 2, 4, 4), + (2, 8, 2, 2, 2), +]) +@pytest.mark.parametrize('dtype', filtered_dtype) +def test_atomic_or_4d_5d(dtype, shape): + x0_value = 3 + x0 = torch.full(shape, x0_value, dtype=eval('torch.' + dtype)).npu() + x1 = torch.full(shape, 2, dtype=eval('torch.' + dtype)).npu() + + x1_ref = x1 | x0_value + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + atomic_or_multi_d[(1, )](x0, x1, *triton_shape) + test_common.validate_cmp(dtype, x1, x1_ref) + + +@pytest.mark.parametrize('shape', [ + (1, 1, 1, 1, 2), + (10, 1, 15, 1, 7), + (1, 1, 1, 1, 257), +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_or_5d(x_dtype_str, shape): + shape = shape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + + out = torch.full(shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + XB, YB, ZB, MB, NB = triton_shape + BLOCK_SIZE = 256 + ncore = (NB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_or_ndim[(2 * XB * YB * ZB * MB, ncore)]( + x_ptr=x, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=XB, + DIM1=YB, + DIM2=ZB, + DIM3=MB, + DIM4=NB, + ) + + expected = out_temp | x_temp[0:shape[0]] | x_temp[shape[0]:x_shape[0]] + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('shape', [ + (1, 1, 1, 1), + (1, 1, 2, 2), + (1, 3, 2, 7), + (1, 3, 2, 651), +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_or_4d(x_dtype_str, shape): + shape = shape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + out = torch.full(shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 4: + triton_shape.append(1) + XB, YB, ZB, MB = triton_shape + + BLOCK_SIZE = 256 + ncore = (MB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_or_ndim[(2 * XB * YB * ZB, ncore)]( + x_ptr=x, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=1, + DIM1=XB, + DIM2=YB, + DIM3=ZB, + DIM4=MB, + ) + + expected = out_temp | x_temp[0:shape[0]] | x_temp[shape[0]:x_shape[0]] + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('shape', [ + (1, 1, 1), + (1, 1, 2), + (1, 31, 275), +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_or_3d_2(x_dtype_str, shape): + shape = shape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + out = torch.full(shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 3: + triton_shape.append(1) + XB, YB, ZB = triton_shape + BLOCK_SIZE = 256 + ncore = (ZB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_or_ndim[(2 * XB * YB, ncore)]( + x_ptr=x, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=1, + DIM1=1, + DIM2=XB, + DIM3=YB, + DIM4=ZB, + ) + + expected = out_temp | x_temp[0:shape[0]] | x_temp[shape[0]:x_shape[0]] + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('shape', [ + (1, 2), + (1, 1), + (257, 1), + (257, 2), +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_or_2d(x_dtype_str, shape): + shape = shape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + out = torch.full(shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 2: + triton_shape.append(1) + XB, YB = triton_shape + BLOCK_SIZE = 256 + ncore = (YB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_or_ndim[(2 * XB, ncore)]( + x_ptr=x, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=1, + DIM1=1, + DIM2=1, + DIM3=XB, + DIM4=YB, + ) + + expected = out_temp | x_temp[0:shape[0]] | x_temp[shape[0]:x_shape[0]] + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('shape', [(1, ), (9, ), (256, ), (257, ), (65535, ), (65536, )]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_or_1d(x_dtype_str, shape): + shape = shape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + out = torch.full(shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 2: + triton_shape.append(1) + XB = triton_shape[0] + BLOCK_SIZE = 256 + ncore = (XB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_or_ndim[(2, ncore)]( + x_ptr=x, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=1, + DIM1=1, + DIM2=1, + DIM3=1, + DIM4=XB, + ) + + expected = out_temp | x_temp[0:shape[0]] | x_temp[shape[0]:x_shape[0]] + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('param_list', [ + ['uint8', (32, 32), 2], + ['uint16', (32, 32), 2], + ['uint32', (32, 32), 2], + ['uint64', (32, 32), 2], +]) +def test_atomic_or_uint(param_list): + dtype, shape, ncore = param_list + block_size = shape[0] * shape[1] // ncore + split_size = shape[0] // ncore + + val_cpu = torch.randint(low=0, high=10, size=shape, dtype=eval(f'torch.{dtype}')).cpu() + val = val_cpu.to("npu") + + pointer_cpu = torch.randint(low=0, high=10, size=(split_size, shape[1]), dtype=eval(f'torch.{dtype}')).cpu() + pointer = pointer_cpu.to("npu") + pointer_old_cpu = torch.full_like(pointer_cpu, -10).cpu() + pointer_old = pointer_old_cpu.to("npu") + pointer_ref_cpu = pointer_cpu.clone() + + for i in range(ncore - 1): + pointer_ref_cpu |= val_cpu[(i * split_size):((i + 1) * split_size)] + + pointer_ref_last = pointer_ref_cpu.clone() + pointer_ref_cpu |= val_cpu[((ncore - 1) * split_size):(ncore * split_size)] + pointer_ref = pointer_ref_cpu.to("npu") + + @triton.jit + def atomic_or_uint(in_ptr0, out_ptr0, out_ptr1, n_elements, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + yindex = tl.arange(0, BLOCK_SIZE)[:] + xmask = xindex < n_elements + x0 = xindex + x1 = yindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.atomic_or(out_ptr0 + (x1), tmp0, xmask) + tl.store(out_ptr1 + (x1), tmp1, xmask) + + n_elements = shape[0] * shape[1] + atomic_or_uint[ncore, 1, 1](val, pointer, pointer_old, n_elements, BLOCK_SIZE=split_size * shape[1]) + test_common.validate_cmp(dtype, pointer, pointer_ref) diff --git a/third_party/ascend/unittest/generalization_cases/test_atomic_xchg.py b/third_party/ascend/unittest/generalization_cases/test_atomic_xchg.py new file mode 100644 index 0000000000..740378e929 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_atomic_xchg.py @@ -0,0 +1,434 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import pytest +import torch +import triton + +import triton.language as tl + +import test_common +from test_common import TestUtils + +filtered_dtype = [dtype for dtype in TestUtils.full_dtype if dtype not in {'uint32', 'bfloat16', 'bool'}] + + +@triton.jit +def atomic_xchg(in_ptr0, out_ptr0, n_elements, BLOCK_SIZE: tl.constexpr, BLOCK_NUM: tl.constexpr): + in_offset = tl.program_id(0) * BLOCK_SIZE + out_offset = (tl.program_id(0) % BLOCK_NUM) * BLOCK_SIZE + in_index = in_offset + tl.arange(0, BLOCK_SIZE) + out_index = out_offset + tl.arange(0, BLOCK_SIZE) + xmask = in_index < n_elements + + tmp0 = tl.load(in_ptr0 + (in_index), xmask) + tl.atomic_xchg(out_ptr0 + (out_index), tmp0, xmask) + + +@triton.jit +def atomic_xchg_ndim(x_ptr, out_ptr, NCORE: tl.constexpr, BLOCK_SIZE: tl.constexpr, DIM0: tl.constexpr, + DIM1: tl.constexpr, DIM2: tl.constexpr, DIM3: tl.constexpr, DIM4: tl.constexpr): + sub_idx = tl.program_id(1) + base_src = tl.program_id(0) * DIM4 + sub_idx * BLOCK_SIZE + base_dst = (tl.program_id(0) % (DIM0 * DIM1 * DIM2 * DIM3)) * DIM4 + sub_idx * BLOCK_SIZE + offsets_src = tl.arange(0, BLOCK_SIZE) + base_src + offsets_dst = tl.arange(0, BLOCK_SIZE) + base_dst + mask = tl.arange(0, BLOCK_SIZE) + sub_idx * BLOCK_SIZE < DIM4 + tmp = tl.load(x_ptr + offsets_src, mask) + tl.atomic_xchg(out_ptr + offsets_dst, tmp, mask) + + +@triton.jit +def atomic_xchg_broadcast(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(0) + + x = tl.load(x_ptr) # x is scalar or 1D, no mask needed + + # Compute y indices + y_offset = pid * BLOCK_SIZE + y_indices = y_offset + tl.arange(0, BLOCK_SIZE) + y_mask = y_indices < n_elements + + y_value = tl.load(y_ptr + y_indices, y_mask) + # Atomic or: y |= x (broadcasted) + tl.atomic_xchg(out_ptr + y_indices, y_value, mask=y_mask) + tl.atomic_xchg(out_ptr + y_indices, x, mask=y_mask) + + +# 定义不同测试场景的参数组合 (x_shape, y_shape, BLOCK_SIZE) +test_cases = [ + ((1, 1, 1, 1), (1, 1, 1, 4), 4), + ((1, 1, 1, 3), (1, 5, 1, 3), 5), + ((3, ), (2, 3, 3, 3, 3), 81), + ((3, ), (2, 3, 3, 3), 27), + ((3, ), (2, 3, 3), 9), + ((3, ), (2, 3), 3), +] + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d + TestUtils.test_shape1d) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_xchg(x_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = test_common.generate_tensor(x_shape, x_dtype_str).npu() + y = torch.full(shape, 0, dtype=x_dtype).npu() + + # 保存副本用于验证 + x_temp = x.clone() + y_temp = y.clone() + + if len(shape) == 2: + n_elements = shape[0] * shape[1] * 2 + atomic_xchg[shape[0] * 2, 1, 1](x, y, n_elements, BLOCK_SIZE=shape[1], BLOCK_NUM=shape[0]) + elif len(shape) == 1: + n_elements = shape[0] + BLOCK_SIZE = min(1024, shape[0]) # 1024:限制最大线程块大小 + grid_size = (n_elements + BLOCK_SIZE - 1) // BLOCK_SIZE # 向上取整 + aligned_size = grid_size * BLOCK_SIZE + x_concat = torch.full([aligned_size * 2], 0, dtype=x_dtype).npu() + x_concat[0:n_elements] = x[0:n_elements] + x_concat[aligned_size:(aligned_size + n_elements)] = x[n_elements:(n_elements * 2)] + atomic_xchg[grid_size * 2, 1, 1](x_concat, y, aligned_size * 2, BLOCK_SIZE=BLOCK_SIZE, BLOCK_NUM=grid_size) + + expected = x_temp[shape[0]:(shape[0] * 2)].expand(y_temp.shape) + torch.testing.assert_close(y, expected) + + +# 3d +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_xchg_3d(x_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = test_common.generate_tensor(x_shape, x_dtype_str).npu() + y = torch.full(shape, 0, dtype=x_dtype).npu() + + # 保存副本用于验证 + x_temp = x.clone() + y_temp = y.clone() + + n_elements = shape[0] * shape[1] * shape[2] + atomic_xchg[2, 1, 1](x, y, n_elements * 2, BLOCK_SIZE=shape[0] * shape[1] * shape[2], BLOCK_NUM=1) + + expected = x_temp[shape[0]:(shape[0] * 2)].expand(y_temp.shape) + torch.testing.assert_close(y, expected) + + +@triton.jit +def atomic_xchg_multi_d(in_ptr0, out_ptr0, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if (YB * ZB * MB * NB) > 1: + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if (ZB * MB * NB) > 1: + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if (MB * NB) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if NB > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + tmp0 = tl.load(in_ptr0 + offsets) + tl.atomic_xchg(out_ptr0 + offsets, tmp0) + + +# multi_d +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (2, 4, 8, 4), + (8, 4, 2, 4), + (2, 8, 2, 2), + (2, 4, 8, 4, 2), + (8, 4, 2, 4, 4), + (2, 8, 2, 2, 2), +]) +@pytest.mark.parametrize('dtype', filtered_dtype) +def test_atomic_xchg_4d_5d(dtype, shape): + x0_value = 3 + x0 = torch.full(shape, x0_value, dtype=eval('torch.' + dtype)).npu() + x1 = torch.full(shape, 2, dtype=eval('torch.' + dtype)).npu() + + x1_ref = x0 + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + atomic_xchg_multi_d[(1, )](x0, x1, *triton_shape) + test_common.validate_cmp(dtype, x1, x1_ref) + + +@pytest.mark.parametrize('shaape', [ + (1, 1, 1, 1, 2), + (10, 1, 15, 1, 7), + (1, 1, 1, 1, 257), +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_xchg_5d(x_dtype_str, shaape): + shape = shaape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + + out = torch.full(shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + XB, YB, ZB, MB, NB = triton_shape + BLOCK_SIZE = 256 + ncore = (NB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_xchg_ndim[(2 * XB * YB * ZB * MB, ncore)]( + x_ptr=x, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=XB, + DIM1=YB, + DIM2=ZB, + DIM3=MB, + DIM4=NB, + ) + + expected = x_temp[shape[0]:x_shape[0]].expand(out_temp.shape) + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('shaape', [ + (1, 1, 1, 1), + (1, 1, 2, 2), + (1, 3, 2, 7), + (1, 3, 2, 651), +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_xchg_4d(x_dtype_str, shaape): + shape = shaape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + out = torch.full(shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 4: + triton_shape.append(1) + XB, YB, ZB, MB = triton_shape + + BLOCK_SIZE = 256 + ncore = (MB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_xchg_ndim[(2 * XB * YB * ZB, ncore)]( + x_ptr=x, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=1, + DIM1=XB, + DIM2=YB, + DIM3=ZB, + DIM4=MB, + ) + + expected = x_temp[shape[0]:x_shape[0]].expand(out_temp.shape) + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('shaape', [ + (1, 1, 1), + (1, 1, 2), + (1, 31, 275), +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_xchg_3d_2(x_dtype_str, shaape): + shape = shaape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + out = torch.full(shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 3: + triton_shape.append(1) + XB, YB, ZB = triton_shape + BLOCK_SIZE = 256 + ncore = (ZB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_xchg_ndim[(2 * XB * YB, ncore)]( + x_ptr=x, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=1, + DIM1=1, + DIM2=XB, + DIM3=YB, + DIM4=ZB, + ) + + expected = x_temp[shape[0]:x_shape[0]].expand(out_temp.shape) + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('shaape', [ + (1, 2), + (1, 1), + (257, 1), + (257, 2), +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_xchg_2d(x_dtype_str, shaape): + shape = shaape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + out = torch.full(shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 2: + triton_shape.append(1) + XB, YB = triton_shape + BLOCK_SIZE = 256 + ncore = (YB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_xchg_ndim[(2 * XB, ncore)]( + x_ptr=x, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=1, + DIM1=1, + DIM2=1, + DIM3=XB, + DIM4=YB, + ) + + expected = x_temp[shape[0]:x_shape[0]].expand(out_temp.shape) + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('shaape', [(1, ), (9, ), (256, ), (257, ), (65535, ), (65536, )]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_xchg_1d(x_dtype_str, shaape): + shape = shaape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + out = torch.full(shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 2: + triton_shape.append(1) + XB = triton_shape[0] + BLOCK_SIZE = 256 + ncore = (XB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_xchg_ndim[(2, ncore)]( + x_ptr=x, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=1, + DIM1=1, + DIM2=1, + DIM3=1, + DIM4=XB, + ) + + expected = x_temp[shape[0]:x_shape[0]].expand(out_temp.shape) + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('param_list', [['uint8', (32, 32), 2], ['uint16', + (32, 32), 2], ['uint32', + (32, 32), 2], ['uint64', (32, 32), 2]]) +def test_atomic_xchg_uint(param_list): + dtype, shape, ncore = param_list + block_size = shape[0] * shape[1] // ncore + split_size = shape[0] // ncore + + val_cpu = torch.randint(low=0, high=10, size=shape, dtype=eval(f'torch.{dtype}')).cpu() + val = val_cpu.to("npu") + + pointer_cpu = torch.randint(low=0, high=10, size=(split_size, shape[1]), dtype=eval(f'torch.{dtype}')).cpu() + pointer = pointer_cpu.to("npu") + + pointer_ref = pointer.clone() + pointer_old_cpu = torch.full_like(val_cpu, -10).cpu() + pointer_old = pointer_old_cpu.to("npu") + pointer_old_ref = pointer_old.clone() + + pointer_ref = val[((ncore - 1) * split_size):(ncore * split_size)].clone() + pointer_old_ref[0:split_size] = pointer + pointer_old_ref[split_size:((ncore - 1) * split_size)] = val[0:(ncore - 2) * split_size] + + @triton.jit + def atomic_xchg_uint(in_ptr0, out_ptr0, out_ptr1, n_elements, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + yindex = tl.arange(0, BLOCK_SIZE)[:] + xmask = xindex < n_elements + x0 = xindex + x1 = yindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.atomic_xchg(out_ptr0 + (x1), tmp0, xmask) + tl.store(out_ptr1 + (x0), tmp1, xmask) + + n_elements = shape[0] * shape[1] + atomic_xchg_uint[ncore, 1, 1](val, pointer, pointer_old, n_elements, BLOCK_SIZE=split_size * shape[1]) + + pointer_cpu = pointer.cpu() + pointer_ref_cpu = pointer_ref.cpu() + assert (pointer_cpu == pointer_ref_cpu).all() diff --git a/third_party/ascend/unittest/generalization_cases/test_atomic_xor.py b/third_party/ascend/unittest/generalization_cases/test_atomic_xor.py new file mode 100644 index 0000000000..4a83697261 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_atomic_xor.py @@ -0,0 +1,441 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import pytest +import torch +import triton + +import triton.language as tl + +import test_common +from test_common import TestUtils + +filtered_dtype = [ + dtype for dtype in TestUtils.full_dtype if dtype not in {'uint32', 'float16', 'float32', 'bfloat16', 'bool'} +] + + +@triton.jit +def atomic_xor(in_ptr0, out_ptr0, n_elements, BLOCK_SIZE: tl.constexpr, BLOCK_NUM: tl.constexpr): + in_offset = tl.program_id(0) * BLOCK_SIZE + out_offset = (tl.program_id(0) % BLOCK_NUM) * BLOCK_SIZE + in_index = in_offset + tl.arange(0, BLOCK_SIZE) + out_index = out_offset + tl.arange(0, BLOCK_SIZE) + xmask = in_index < n_elements + + tmp0 = tl.load(in_ptr0 + (in_index), xmask) + tl.atomic_xor(out_ptr0 + (out_index), tmp0, xmask) + + +@triton.jit +def atomic_xor_ndim(x_ptr, out_ptr, NCORE: tl.constexpr, BLOCK_SIZE: tl.constexpr, DIM0: tl.constexpr, + DIM1: tl.constexpr, DIM2: tl.constexpr, DIM3: tl.constexpr, DIM4: tl.constexpr): + sub_idx = tl.program_id(1) + base_src = tl.program_id(0) * DIM4 + sub_idx * BLOCK_SIZE + base_dst = (tl.program_id(0) % (DIM0 * DIM1 * DIM2 * DIM3)) * DIM4 + sub_idx * BLOCK_SIZE + offsets_src = tl.arange(0, BLOCK_SIZE) + base_src + offsets_dst = tl.arange(0, BLOCK_SIZE) + base_dst + mask = tl.arange(0, BLOCK_SIZE) + sub_idx * BLOCK_SIZE < DIM4 + tmp = tl.load(x_ptr + offsets_src, mask) + tl.atomic_xor(out_ptr + offsets_dst, tmp, mask) + + +@triton.jit +def atomic_xor_broadcast(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(0) + + x = tl.load(x_ptr) # x is scalar or 1D, no mask needed + + # Compute y indices + y_offset = pid * BLOCK_SIZE + y_indices = y_offset + tl.arange(0, BLOCK_SIZE) + y_mask = y_indices < n_elements + + y_value = tl.load(y_ptr + y_indices, y_mask) + # Atomic or: y |= x (broadcasted) + tl.atomic_xor(out_ptr + y_indices, y_value, mask=y_mask) + tl.atomic_xor(out_ptr + y_indices, x, mask=y_mask) + + +# 定义不同测试场景的参数组合 (x_shape, y_shape, BLOCK_SIZE) +test_cases = [ + ((1, 1, 1, 1), (1, 1, 1, 4), 4), + ((1, 1, 1, 3), (1, 5, 1, 3), 5), + ((3, ), (2, 3, 3, 3, 3), 81), + ((3, ), (2, 3, 3, 3), 27), + ((3, ), (2, 3, 3), 9), + ((3, ), (2, 3), 3), +] + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d + TestUtils.test_shape1d) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_xor(x_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + + if len(shape) == 1 and shape[0] == 1: # golden 问题，手动验证 + return + + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = test_common.generate_tensor(x_shape, x_dtype_str).npu() + y = torch.full(shape, 0, dtype=x_dtype).npu() + + # 保存副本用于验证 + x_temp = x.clone() + y_temp = y.clone() + + if len(shape) == 2: + n_elements = shape[0] * shape[1] * 2 + atomic_xor[shape[0] * 2, 1, 1](x, y, n_elements, BLOCK_SIZE=shape[1], BLOCK_NUM=shape[0]) + elif len(shape) == 1: + n_elements = shape[0] + BLOCK_SIZE = min(1024, shape[0]) # 1024:限制最大线程块大小 + grid_size = (n_elements + BLOCK_SIZE - 1) // BLOCK_SIZE # 向上取整 + aligned_size = grid_size * BLOCK_SIZE + x_concat = torch.full([aligned_size * 2], 0, dtype=x_dtype).npu() + x_concat[0:n_elements] = x[0:n_elements] + x_concat[aligned_size:(aligned_size + n_elements)] = x[n_elements:(n_elements * 2)] + atomic_xor[grid_size * 2, 1, 1](x_concat, y, aligned_size * 2, BLOCK_SIZE=BLOCK_SIZE, BLOCK_NUM=grid_size) + + expected = y_temp ^ x_temp[0:shape[0]] ^ x_temp[shape[0]:(shape[0] * 2)] + torch.testing.assert_close(y, expected) + + +# 3d +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_xor_3d(x_dtype_str, shape): + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = test_common.generate_tensor(x_shape, x_dtype_str).npu() + y = torch.full(shape, 0, dtype=x_dtype).npu() + + # 保存副本用于验证 + x_temp = x.clone() + y_temp = y.clone() + + n_elements = shape[0] * shape[1] * shape[2] + atomic_xor[2, 1, 1](x, y, n_elements * 2, BLOCK_SIZE=shape[0] * shape[1] * shape[2], BLOCK_NUM=1) + + expected = y_temp ^ x_temp[0:shape[0]] ^ x_temp[shape[0]:(shape[0] * 2)] + torch.testing.assert_close(y, expected) + + +@triton.jit +def atomic_xor_multi_d(in_ptr0, out_ptr0, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if (YB * ZB * MB * NB) > 1: + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if (ZB * MB * NB) > 1: + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if (MB * NB) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if NB > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + tmp0 = tl.load(in_ptr0 + offsets) + tl.atomic_xor(out_ptr0 + offsets, tmp0) + + +# multi_d +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (2, 4, 8, 4), + (8, 4, 2, 4), + (2, 8, 2, 2), + (2, 4, 8, 4, 2), + (8, 4, 2, 4, 4), + (2, 8, 2, 2, 2), +]) +@pytest.mark.parametrize('dtype', filtered_dtype) +def test_atomic_xor_4d_5d(dtype, shape): + x0_value = 3 + x0 = torch.full(shape, x0_value, dtype=eval('torch.' + dtype)).npu() + x1 = torch.full(shape, 2, dtype=eval('torch.' + dtype)).npu() + + x1_ref = x1 ^ x0_value + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + atomic_xor_multi_d[(1, )](x0, x1, *triton_shape) + test_common.validate_cmp(dtype, x1, x1_ref) + + +@pytest.mark.parametrize('shape', [ + (1, 1, 1, 1, 2), + (10, 1, 15, 1, 7), + (1, 1, 1, 1, 257), +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_xor_5d(x_dtype_str, shape): + shape = shape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + + out = torch.full(shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + XB, YB, ZB, MB, NB = triton_shape + BLOCK_SIZE = 256 + ncore = (NB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_xor_ndim[(2 * XB * YB * ZB * MB, ncore)]( + x_ptr=x, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=XB, + DIM1=YB, + DIM2=ZB, + DIM3=MB, + DIM4=NB, + ) + + expected = out_temp ^ x_temp[0:shape[0]] ^ x_temp[shape[0]:x_shape[0]] + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('shape', [ + (1, 1, 1, 1), + (1, 1, 2, 2), + (1, 3, 2, 7), + (1, 3, 2, 651), +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_xor_4d(x_dtype_str, shape): + shape = shape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + out = torch.full(shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 4: + triton_shape.append(1) + XB, YB, ZB, MB = triton_shape + + BLOCK_SIZE = 256 + ncore = (MB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_xor_ndim[(2 * XB * YB * ZB, ncore)]( + x_ptr=x, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=1, + DIM1=XB, + DIM2=YB, + DIM3=ZB, + DIM4=MB, + ) + + expected = out_temp ^ x_temp[0:shape[0]] ^ x_temp[shape[0]:x_shape[0]] + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('shape', [ + (1, 1, 1), + (1, 1, 2), + (1, 31, 275), +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_xor_3d_2(x_dtype_str, shape): + shape = shape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + out = torch.full(shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 3: + triton_shape.append(1) + XB, YB, ZB = triton_shape + BLOCK_SIZE = 256 + ncore = (ZB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_xor_ndim[(2 * XB * YB, ncore)]( + x_ptr=x, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=1, + DIM1=1, + DIM2=XB, + DIM3=YB, + DIM4=ZB, + ) + + expected = out_temp ^ x_temp[0:shape[0]] ^ x_temp[shape[0]:x_shape[0]] + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('shape', [ + (1, 2), + (1, 1), + (257, 1), + (257, 2), +]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_xor_2d(x_dtype_str, shape): + shape = shape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + out = torch.full(shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 2: + triton_shape.append(1) + XB, YB = triton_shape + BLOCK_SIZE = 256 + ncore = (YB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_xor_ndim[(2 * XB, ncore)]( + x_ptr=x, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=1, + DIM1=1, + DIM2=1, + DIM3=XB, + DIM4=YB, + ) + + expected = out_temp ^ x_temp[0:shape[0]] ^ x_temp[shape[0]:x_shape[0]] + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('shape', [(1, ), (9, ), (256, ), (257, ), (65535, ), (65536, )]) +@pytest.mark.parametrize('x_dtype_str', filtered_dtype) +def test_atomic_xor_1d(x_dtype_str, shape): + shape = shape + + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + x_shape = list(shape[:]) + x_shape[0] *= 2 + x = torch.randint(low=0, high=100, size=x_shape, dtype=x_dtype).npu() + out = torch.full(shape, 0, dtype=x_dtype).npu() + + x_temp = x.clone() + out_temp = out.clone() + + triton_shape = [*shape] + while len(triton_shape) < 2: + triton_shape.append(1) + XB = triton_shape[0] + BLOCK_SIZE = 256 + ncore = (XB + BLOCK_SIZE - 1) // BLOCK_SIZE + + atomic_xor_ndim[(2, ncore)]( + x_ptr=x, + out_ptr=out, + NCORE=ncore, + BLOCK_SIZE=BLOCK_SIZE, + DIM0=1, + DIM1=1, + DIM2=1, + DIM3=1, + DIM4=XB, + ) + + expected = out_temp ^ x_temp[0:shape[0]] ^ x_temp[shape[0]:x_shape[0]] + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('param_list', [ + ['uint8', (32, 32), 2], + ['uint16', (32, 32), 2], + ['uint32', (32, 32), 2], + ['uint64', (32, 32), 2], +]) +def test_atomic_xor_uint(param_list): + dtype, shape, ncore = param_list + block_size = shape[0] * shape[1] // ncore + split_size = shape[0] // ncore + + val_value = 3 + val_cpu = torch.full(shape, val_value, dtype=eval(f'torch.{dtype}')).cpu() + val = val_cpu.to("npu") + + pointer_value = 5 + pointer_cpu = torch.full((split_size, shape[1]), pointer_value, dtype=eval(f'torch.{dtype}')).cpu() + pointer = pointer_cpu.to("npu") + pointer_old_cpu = torch.full_like(pointer_cpu, -10).cpu() + pointer_old = pointer_old_cpu.to("npu") + + pointer_result = pointer_value + for _ in range(ncore): + pointer_result ^= val_value + + pointer_ref_cpu = torch.full_like(pointer_cpu, pointer_result).cpu() + pointer_ref = pointer_ref_cpu.to("npu") + + @triton.jit + def atomic_xor_uint(in_ptr0, out_ptr0, out_ptr1, n_elements, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + yindex = tl.arange(0, BLOCK_SIZE)[:] + xmask = xindex < n_elements + x0 = xindex + x1 = yindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.atomic_xor(out_ptr0 + (x1), tmp0, xmask) + tl.store(out_ptr1 + (x1), tmp1, xmask) + + n_elements = shape[0] * shape[1] + atomic_xor_uint[ncore, 1, 1](val, pointer, pointer_old, n_elements, BLOCK_SIZE=split_size * shape[1]) + test_common.validate_cmp(dtype, pointer, pointer_ref) diff --git a/third_party/ascend/unittest/generalization_cases/test_broadcast.py b/third_party/ascend/unittest/generalization_cases/test_broadcast.py new file mode 100644 index 0000000000..e9f7a46d8b --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_broadcast.py @@ -0,0 +1,324 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils + + +@triton.jit +def fn_broadcast_1d(output_ptr, x_ptr, XS: tl.constexpr, YS: tl.constexpr): + xidx = tl.arange(0, XS)[None, :] + base = tl.load(x_ptr + xidx) + out = base.broadcast_to((YS, XS)) + oidx = tl.arange(0, YS)[:, None] * XS + tl.arange(0, XS)[None, :] + tl.store(output_ptr + oidx, out) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_npu_1d(shape, dtype): + XS = shape[0] + YS = 4 + + x = test_common.generate_tensor((XS, ), dtype=dtype).npu() + std = torch.broadcast_to(x, (YS, XS)) + output = test_common.generate_tensor((YS, XS), dtype=dtype).npu() + fn_broadcast_1d[1, 1, 1](output, x, XS, YS) + test_common.validate_cmp(dtype, std, output) + + +@triton.jit +def fn_broadcast_2d(output_ptr, x_ptr, NUMEL: tl.constexpr, XS: tl.constexpr, YS: tl.constexpr, ZS: tl.constexpr): + zoffset = tl.program_id(0) * ZS + zidx = tl.arange(0, ZS)[None, :] + base = tl.load(x_ptr + zoffset + zidx) + out = base.broadcast_to((YS, ZS)) + oidx = zoffset * YS + tl.arange(0, YS)[:, None] * ZS + zidx + tl.store(output_ptr + oidx, out) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_npu_2d(shape, dtype): + XS = shape[0] + ZS = shape[1] + YS = 4 + NUMEL = XS * ZS + + x = test_common.generate_tensor((XS, 1, ZS), dtype=dtype).npu() # randn not support int type + std = torch.broadcast_to(x, (XS, YS, ZS)) + output = test_common.generate_tensor((XS, YS, ZS), dtype=dtype).npu() + fn_broadcast_2d[XS, 1, 1](output, x, NUMEL, XS, YS, ZS) + test_common.validate_cmp(dtype, std, output) + + +@triton.jit +def triton_broadcast_to_dim0(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = tl.arange(0, 1)[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x = tl.load(in_ptr0 + idx) + x1 = tl.load(out_ptr0 + odx) + ret = tl.broadcast(x, x1) + tl.store(out_ptr0 + odx, ret) + + +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +def test_broadcast_to_dim0(shape, dtype): + L, M, N = shape + x0 = test_common.generate_tensor(shape=(1, M, N), dtype=dtype).npu() + ans = x0.repeat(L, 1, 1) + output = torch.zeros((L, M, N), dtype=eval('torch.' + dtype)).npu() + triton_broadcast_to_dim0[1, 1, 1](x0, output, L, M, N) + test_common.validate_cmp(dtype, output, ans) + + +@triton.jit +def triton_broadcast_to_dim1(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = lblk_idx[:, None, None] * N * 1 + tl.arange(0, 1)[None, :, None] * N + nblk_idx[None, None, :] + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x = tl.load(in_ptr0 + idx) + x1 = tl.load(out_ptr0 + odx) + ret = tl.broadcast(x, x1) + tl.store(out_ptr0 + odx, ret) + + +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +def test_broadcast_to_dim1(shape, dtype): + L, M, N = shape + x0 = test_common.generate_tensor(shape=(L, 1, N), dtype=dtype).npu() + ans = x0.repeat(1, M, 1) + output = torch.zeros((L, M, N), dtype=eval('torch.' + dtype)).npu() + triton_broadcast_to_dim1[1, 1, 1](x0, output, L, M, N) + test_common.validate_cmp(dtype, output, ans) + + +@triton.jit +def triton_broadcast_to_dim2(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = lblk_idx[:, None, None] * 1 * M + mblk_idx[None, :, None] * 1 + tl.arange(0, 1)[None, None, :] + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x = tl.load(in_ptr0 + idx) + x1 = tl.load(out_ptr0 + odx) + ret = tl.broadcast(x, x1) + tl.store(out_ptr0 + odx, ret) + + +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +def test_broadcast_to_dim2(shape, dtype): + L, M, N = shape + x0 = test_common.generate_tensor(shape=(L, M, 1), dtype=dtype).npu() + ans = x0.repeat(1, 1, N) + output = torch.zeros((L, M, N), dtype=eval('torch.' + dtype)).npu() + triton_broadcast_to_dim2[1, 1, 1](x0, output, L, M, N) + test_common.validate_cmp(dtype, output, ans) + + +@triton.jit +def triton_broadcast_to_dim01(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = tl.arange(0, 1)[:, None, None] * N * 1 + tl.arange(0, 1)[None, :, None] * N + nblk_idx[None, None, :] + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x = tl.load(in_ptr0 + idx) + x1 = tl.load(out_ptr0 + odx) + ret = tl.broadcast(x, x1) + tl.store(out_ptr0 + odx, ret) + + +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +def test_broadcast_to_dim01(shape, dtype): + L, M, N = shape + x0 = test_common.generate_tensor(shape=(1, 1, N), dtype=dtype).npu() + ans = x0.repeat(L, M, 1) + output = torch.zeros((L, M, N), dtype=eval('torch.' + dtype)).npu() + triton_broadcast_to_dim01[1, 1, 1](x0, output, L, M, N) + test_common.validate_cmp(dtype, output, ans) + + +@triton.jit +def triton_broadcast_to_dim02(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = tl.arange(0, 1)[:, None, None] * M * 1 + mblk_idx[None, :, None] * 1 + tl.arange(0, 1)[None, None, :] + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x = tl.load(in_ptr0 + idx) + x1 = tl.load(out_ptr0 + odx) + ret = tl.broadcast(x, x1) + tl.store(out_ptr0 + odx, ret) + + +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +def test_broadcast_to_dim02(shape, dtype): + L, M, N = shape + x0 = test_common.generate_tensor(shape=(1, M, 1), dtype=dtype).npu() + ans = x0.repeat(L, 1, N) + output = torch.zeros((L, M, N), dtype=eval('torch.' + dtype)).npu() + triton_broadcast_to_dim02[1, 1, 1](x0, output, L, M, N) + test_common.validate_cmp(dtype, output, ans) + + +@triton.jit +def triton_broadcast_to_dim12(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = lblk_idx[:, None, None] * 1 * 1 + tl.arange(0, 1)[None, :, None] * 1 + tl.arange(0, 1)[None, None, :] + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x = tl.load(in_ptr0 + idx) + x1 = tl.load(out_ptr0 + odx) + ret = tl.broadcast(x, x1) + tl.store(out_ptr0 + odx, ret) + + +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +def test_broadcast_to_dim12(shape, dtype): + L, M, N = shape + x0 = test_common.generate_tensor(shape=(L, 1, 1), dtype=dtype).npu() + ans = x0.repeat(1, M, N) + output = torch.zeros((L, M, N), dtype=eval('torch.' + dtype)).npu() + triton_broadcast_to_dim12[1, 1, 1](x0, output, L, M, N) + test_common.validate_cmp(dtype, output, ans) + + +@triton.jit +def fn_broadcast_multi_d(to_ptr, from_ptr, F_L: tl.constexpr, F_M: tl.constexpr, F_N: tl.constexpr, F_X: tl.constexpr, + F_Y: tl.constexpr, T_L: tl.constexpr, T_M: tl.constexpr, T_N: tl.constexpr, T_X: tl.constexpr, + T_Y: tl.constexpr): + from_offsets = tl.arange(0, F_L) + if F_M is not None: + from_offsets = from_offsets[:, None] * F_M + tl.arange(0, F_M)[None, :] + if F_N is not None: + from_offsets = from_offsets[:, :, None] * F_N + tl.arange(0, F_N)[None, None, :] + if F_X is not None: + from_offsets = from_offsets[:, :, :, None] * F_X + tl.arange(0, F_X)[None, None, None, :] + if F_Y is not None: + from_offsets = from_offsets[:, :, :, :, None] * F_Y + tl.arange(0, F_Y)[None, None, None, None, :] + + to_offsets = tl.arange(0, T_L) + if T_M is not None: + to_offsets = to_offsets[:, None] * T_M + tl.arange(0, T_M)[None, :] + if T_N is not None: + to_offsets = to_offsets[:, :, None] * T_N + tl.arange(0, T_N)[None, None, :] + if T_X is not None: + to_offsets = to_offsets[:, :, :, None] * T_X + tl.arange(0, T_X)[None, None, None, :] + if T_Y is not None: + to_offsets = to_offsets[:, :, :, :, None] * T_Y + tl.arange(0, T_Y)[None, None, None, None, :] + + from_data = tl.load(from_ptr + from_offsets) + to_data = tl.load(to_ptr + to_offsets) + ret_data = tl.broadcast(from_data, to_data) + + tl.store(to_ptr + to_offsets, ret_data) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shapes', [ + [(1, 64, 16, 1), (2, 64, 16, 2)], + [(8, 1, 1, 2), (8, 8, 4, 2)], +]) +@pytest.mark.parametrize('dtype', ["int32", "int64", "float16", "float32", "bfloat16"]) +def test_broadcast_to_4d(shapes, dtype): + from_shape, to_shape = shapes + dtype = eval(f"torch.{dtype}") + + x = torch.randint(0, 8, from_shape, dtype=dtype).npu() + y = torch.randint(0, 8, to_shape, dtype=dtype).npu() + expected = x.expand(to_shape) + + grid = (1, ) + triton_from_shape = [*from_shape] + triton_to_shape = [*to_shape] + while len(triton_from_shape) < 5: + triton_from_shape.append(None) + triton_to_shape.append(None) + fn_broadcast_multi_d[grid](y, x, *triton_from_shape, *triton_to_shape) + assert (torch.equal(y, expected)) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('dtype', ["int32", "int64", "float16", "float32", "bfloat16"]) +@pytest.mark.parametrize('shapes', [ + [(1, 4, 2, 1, 4), (2, 4, 2, 8, 4)], + [(3, 1, 2, 1, 4), (3, 4, 2, 8, 4)], +]) +def test_broadcast_to_5d(shapes, dtype): + from_shape, to_shape = shapes + dtype = eval(f"torch.{dtype}") + + x = torch.randint(0, 8, from_shape, dtype=dtype).npu() + y = torch.randint(0, 8, to_shape, dtype=dtype).npu() + expected = x.expand(to_shape) + + grid = (1, ) + triton_from_shape = [*from_shape] + triton_to_shape = [*to_shape] + while len(triton_from_shape) < 5: + triton_from_shape.append(None) + triton_to_shape.append(None) + fn_broadcast_multi_d[grid](y, x, *triton_from_shape, *triton_to_shape) + assert (torch.equal(y, expected)) + + +@triton.jit +def fn_broadcast(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = tl.arange(0, 1)[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x = tl.load(in_ptr0 + idx) + x1 = tl.load(out_ptr0 + odx) + ret = tl.broadcast(x, x1) + tl.store(out_ptr0 + odx, ret) + + +XS: tl.constexpr = 2 +YS: tl.constexpr = 4 +ZS: tl.constexpr = 8 + + +@pytest.mark.parametrize('dtype', + ["uint8", "int8", "int16", "int32", "int64", "float16", "float32", "bfloat16", "bool"]) +def test_broadcast_alltype(dtype): + input = test_common.generate_tensor((1, YS, ZS), dtype).npu() + ans = input.repeat(XS, 1, 1) + output = torch.zeros((XS, YS, ZS), dtype=eval('torch.' + dtype)).npu() + fn_broadcast[1, 1, 1](input, output, XS, YS, ZS) + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_broadcast_to.py b/third_party/ascend/unittest/generalization_cases/test_broadcast_to.py new file mode 100644 index 0000000000..4ec6173874 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_broadcast_to.py @@ -0,0 +1,327 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time + +import torch +import torch_npu +import test_common +from test_common import TestUtils + + +@triton.jit +def fn_broadcast_1d(output_ptr, x_ptr, XS: tl.constexpr, YS: tl.constexpr): + xidx = tl.arange(0, XS)[None, :] + base = tl.load(x_ptr + xidx) + out = base.broadcast_to((YS, XS)) + oidx = tl.arange(0, YS)[:, None] * XS + tl.arange(0, XS)[None, :] + tl.store(output_ptr + oidx, out) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_npu_1d(shape, dtype): + XS = shape[0] + YS = 4 + + x = test_common.generate_tensor((XS, ), dtype=dtype).npu() + std = torch.broadcast_to(x, (YS, XS)) + output = test_common.generate_tensor((YS, XS), dtype=dtype).npu() + fn_broadcast_1d[1, 1, 1](output, x, XS, YS) + test_common.validate_cmp(dtype, std, output) + + +@triton.jit +def fn_broadcast_2d(output_ptr, x_ptr, NUMEL: tl.constexpr, XS: tl.constexpr, YS: tl.constexpr, ZS: tl.constexpr): + zoffset = tl.program_id(0) * ZS + zidx = tl.arange(0, ZS)[None, :] + base = tl.load(x_ptr + zoffset + zidx) + out = base.broadcast_to((YS, ZS)) + oidx = zoffset * YS + tl.arange(0, YS)[:, None] * ZS + zidx + tl.store(output_ptr + oidx, out) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_npu_2d(shape, dtype): + XS = shape[0] + ZS = shape[1] + YS = 4 + NUMEL = XS * ZS + + x = test_common.generate_tensor((XS, 1, ZS), dtype=dtype).npu() + std = torch.broadcast_to(x, (XS, YS, ZS)) + output = test_common.generate_tensor((XS, YS, ZS), dtype=dtype).npu() + fn_broadcast_2d[XS, 1, 1](output, x, NUMEL, XS, YS, ZS) + test_common.validate_cmp(dtype, std, output) + + +@triton.jit +def triton_broadcast_to_dim0(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = tl.arange(0, 1)[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x = tl.load(in_ptr0 + idx) + ret = x.broadcast_to(L, M, N) + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + tl.store(out_ptr0 + odx, ret) + + +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +def test_broadcast_to_dim0(shape, dtype): + L, M, N = shape + x0 = test_common.generate_tensor(shape=(1, M, N), dtype=dtype).npu() + ans = x0.repeat(L, 1, 1) + output = torch.zeros((L, M, N), dtype=eval('torch.' + dtype)).npu() + triton_broadcast_to_dim0[1, 1, 1](x0, output, L, M, N) + test_common.validate_cmp(dtype, output, ans) + + +@triton.jit +def triton_broadcast_to_dim1(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = lblk_idx[:, None, None] * N * 1 + tl.arange(0, 1)[None, :, None] * N + nblk_idx[None, None, :] + x = tl.load(in_ptr0 + idx) + ret = x.broadcast_to(L, M, N) + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + tl.store(out_ptr0 + odx, ret) + + +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +def test_broadcast_to_dim1(shape, dtype): + L, M, N = shape + x0 = test_common.generate_tensor(shape=(L, 1, N), dtype=dtype).npu() + ans = x0.repeat(1, M, 1) + output = torch.zeros((L, M, N), dtype=eval('torch.' + dtype)).npu() + triton_broadcast_to_dim1[1, 1, 1](x0, output, L, M, N) + test_common.validate_cmp(dtype, output, ans) + + +@triton.jit +def triton_broadcast_to_dim2(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = lblk_idx[:, None, None] * 1 * M + mblk_idx[None, :, None] * 1 + tl.arange(0, 1)[None, None, :] + x = tl.load(in_ptr0 + idx) + ret = x.broadcast_to(L, M, N) + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + tl.store(out_ptr0 + odx, ret) + + +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +def test_broadcast_to_dim2(shape, dtype): + L, M, N = shape + x0 = test_common.generate_tensor(shape=(L, M, 1), dtype=dtype).npu() + ans = x0.repeat(1, 1, N) + output = torch.zeros((L, M, N), dtype=eval('torch.' + dtype)).npu() + triton_broadcast_to_dim2[1, 1, 1](x0, output, L, M, N) + test_common.validate_cmp(dtype, output, ans) + + +@triton.jit +def triton_broadcast_to_dim01(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = tl.arange(0, 1)[:, None, None] * N * 1 + tl.arange(0, 1)[None, :, None] * N + nblk_idx[None, None, :] + x = tl.load(in_ptr0 + idx) + ret = x.broadcast_to(L, M, N) + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + tl.store(out_ptr0 + odx, ret) + + +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +def test_broadcast_to_dim01(shape, dtype): + L, M, N = shape + x0 = test_common.generate_tensor(shape=(1, 1, N), dtype=dtype).npu() + ans = x0.repeat(L, M, 1) + output = torch.zeros((L, M, N), dtype=eval('torch.' + dtype)).npu() + triton_broadcast_to_dim01[1, 1, 1](x0, output, L, M, N) + test_common.validate_cmp(dtype, output, ans) + + +@triton.jit +def triton_broadcast_to_dim02(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = tl.arange(0, 1)[:, None, None] * M * 1 + mblk_idx[None, :, None] * 1 + tl.arange(0, 1)[None, None, :] + x = tl.load(in_ptr0 + idx) + ret = x.broadcast_to(L, M, N) + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + tl.store(out_ptr0 + odx, ret) + + +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +def test_broadcast_to_dim02(shape, dtype): + L, M, N = shape + x0 = test_common.generate_tensor(shape=(1, M, 1), dtype=dtype).npu() + ans = x0.repeat(L, 1, N) + output = torch.zeros((L, M, N), dtype=eval('torch.' + dtype)).npu() + triton_broadcast_to_dim02[1, 1, 1](x0, output, L, M, N) + test_common.validate_cmp(dtype, output, ans) + + +@triton.jit +def triton_broadcast_to_dim12(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = lblk_idx[:, None, None] * 1 * 1 + tl.arange(0, 1)[None, :, None] * 1 + tl.arange(0, 1)[None, None, :] + x = tl.load(in_ptr0 + idx) + ret = x.broadcast_to(L, M, N) + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + tl.store(out_ptr0 + odx, ret) + + +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +def test_broadcast_to_dim12(shape, dtype): + L, M, N = shape + x0 = test_common.generate_tensor(shape=(L, 1, 1), dtype=dtype).npu() + ans = x0.repeat(1, M, N) + output = torch.zeros((L, M, N), dtype=eval('torch.' + dtype)).npu() + triton_broadcast_to_dim12[1, 1, 1](x0, output, L, M, N) + test_common.validate_cmp(dtype, output, ans) + + +@triton.jit +def fn_broadcast_to_multi_d(to_ptr, from_ptr, F_L: tl.constexpr, F_M: tl.constexpr, F_N: tl.constexpr, + F_X: tl.constexpr, F_Y: tl.constexpr, T_L: tl.constexpr, T_M: tl.constexpr, + T_N: tl.constexpr, T_X: tl.constexpr, T_Y: tl.constexpr): + from_offsets = tl.arange(0, F_L) + if F_M is not None: + from_offsets = from_offsets[:, None] * F_M + tl.arange(0, F_M)[None, :] + if F_N is not None: + from_offsets = from_offsets[:, :, None] * F_N + tl.arange(0, F_N)[None, None, :] + if F_X is not None: + from_offsets = from_offsets[:, :, :, None] * F_X + tl.arange(0, F_X)[None, None, None, :] + if F_Y is not None: + from_offsets = from_offsets[:, :, :, :, None] * F_Y + tl.arange(0, F_Y)[None, None, None, None, :] + + to_offsets = tl.arange(0, T_L) + if T_M is not None: + to_offsets = to_offsets[:, None] * T_M + tl.arange(0, T_M)[None, :] + if T_N is not None: + to_offsets = to_offsets[:, :, None] * T_N + tl.arange(0, T_N)[None, None, :] + if T_X is not None: + to_offsets = to_offsets[:, :, :, None] * T_X + tl.arange(0, T_X)[None, None, None, :] + if T_Y is not None: + to_offsets = to_offsets[:, :, :, :, None] * T_Y + tl.arange(0, T_Y)[None, None, None, None, :] + + from_data = tl.load(from_ptr + from_offsets) + if F_Y is not None: + ret_data = from_data.broadcast_to((T_L, T_M, T_N, T_X, T_Y)) + elif F_X is not None: + ret_data = from_data.broadcast_to((T_L, T_M, T_N, T_X)) + elif F_N is not None: + ret_data = from_data.broadcast_to((T_L, T_M, T_N)) + elif F_M is not None: + ret_data = from_data.broadcast_to((T_L, T_M)) + else: + ret_data = from_data.broadcast_to((T_L)) + + tl.store(to_ptr + to_offsets, ret_data) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shapes', [ + [(1, 64, 16, 1), (2, 64, 16, 2)], + [(8, 1, 1, 2), (8, 8, 4, 2)], +]) +@pytest.mark.parametrize('dtype', ["int32", "int64", "float16", "float32", "bfloat16"]) +def test_broadcast_to_4d(shapes, dtype): + from_shape, to_shape = shapes + dtype = eval(f"torch.{dtype}") + + x = torch.randint(0, 8, from_shape, dtype=dtype).npu() + y = torch.randint(0, 8, to_shape, dtype=dtype).npu() + expected = x.expand(to_shape) + + grid = (1, ) + triton_from_shape = [*from_shape] + triton_to_shape = [*to_shape] + while len(triton_from_shape) < 5: + triton_from_shape.append(None) + triton_to_shape.append(None) + fn_broadcast_to_multi_d[grid](y, x, *triton_from_shape, *triton_to_shape) + assert (torch.equal(y, expected)) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('dtype', ["int32", "int64", "float16", "float32", "bfloat16"]) +@pytest.mark.parametrize('shapes', [ + [(1, 4, 2, 1, 4), (2, 4, 2, 8, 4)], + [(3, 1, 2, 1, 4), (3, 4, 2, 8, 4)], +]) +def test_broadcast_to_5d(shapes, dtype): + from_shape, to_shape = shapes + dtype = eval(f"torch.{dtype}") + + x = torch.randint(0, 8, from_shape, dtype=dtype).npu() + y = torch.randint(0, 8, to_shape, dtype=dtype).npu() + expected = x.expand(to_shape) + + grid = (1, ) + triton_from_shape = [*from_shape] + triton_to_shape = [*to_shape] + while len(triton_from_shape) < 5: + triton_from_shape.append(None) + triton_to_shape.append(None) + fn_broadcast_to_multi_d[grid](y, x, *triton_from_shape, *triton_to_shape) + assert (torch.equal(y, expected)) + + +XS: tl.constexpr = 2 +YS: tl.constexpr = 4 +ZS: tl.constexpr = 8 +NUMEL: tl.constexpr = XS * ZS + + +@triton.jit +def fn_broadcast_to(output_ptr, input_ptr, length): + col_offsets = tl.arange(0, NUMEL) + input = tl.load(input_ptr + col_offsets) + result = input.reshape((XS, 1, ZS)).broadcast_to((XS, YS, ZS)).reshape((XS * YS * ZS)) + brc_col_offsets = tl.arange(0, NUMEL * YS) + tl.store(output_ptr + brc_col_offsets, result) + + +@pytest.mark.parametrize('dtype', + ["uint8", "int8", "int16", "int32", "int64", "float16", "float32", "bfloat16", "bool"]) +def test_broadcast_to_alltype(dtype): + length = NUMEL + input = test_common.generate_tensor((XS, 1, ZS), dtype).npu() + output = test_common.generate_tensor((XS, YS, ZS), dtype).npu() + fn_broadcast_to[1, 1, 1](output, input, length, debug=True) + assert (torch.equal(output, input.repeat(1, YS, 1))) diff --git a/third_party/ascend/unittest/generalization_cases/test_cast.py b/third_party/ascend/unittest/generalization_cases/test_cast.py new file mode 100644 index 0000000000..3e1608cb97 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_cast.py @@ -0,0 +1,391 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import random +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common +from test_common import TestUtils, check_ub_mem_overflow, get_dtype_size + + +@triton.jit +def cast_to_bool(output_ptr, x_ptr, x_stride, y_stride, z_stride, DIM: tl.constexpr, XB: tl.constexpr, YB: tl.constexpr, + ZB: tl.constexpr): + if DIM == 1: + xidx = tl.arange(0, XB) + idx = xidx * x_stride + elif DIM == 2: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + idx = xidx[:, None] * x_stride + yidx[None, :] * y_stride + elif DIM == 3: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * x_stride + yidx[None, :, None] * y_stride + zidx[None, None, :] * z_stride + + X = tl.load(x_ptr + idx) + ret = tl.cast(X, dtype=tl.int1) + tl.store(output_ptr + idx, ret) + + +@triton.jit +def cast_to_i8(output_ptr, x_ptr, x_stride, y_stride, z_stride, DIM: tl.constexpr, XB: tl.constexpr, YB: tl.constexpr, + ZB: tl.constexpr): + if DIM == 1: + xidx = tl.arange(0, XB) + idx = xidx * x_stride + elif DIM == 2: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + idx = xidx[:, None] * x_stride + yidx[None, :] * y_stride + elif DIM == 3: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * x_stride + yidx[None, :, None] * y_stride + zidx[None, None, :] * z_stride + + X = tl.load(x_ptr + idx) + ret = tl.cast(X, dtype=tl.int8) + tl.store(output_ptr + idx, ret) + + +@triton.jit +def cast_to_i16(output_ptr, x_ptr, x_stride, y_stride, z_stride, DIM: tl.constexpr, XB: tl.constexpr, YB: tl.constexpr, + ZB: tl.constexpr): + if DIM == 1: + xidx = tl.arange(0, XB) + idx = xidx * x_stride + elif DIM == 2: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + idx = xidx[:, None] * x_stride + yidx[None, :] * y_stride + elif DIM == 3: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * x_stride + yidx[None, :, None] * y_stride + zidx[None, None, :] * z_stride + + X = tl.load(x_ptr + idx) + ret = tl.cast(X, dtype=tl.int16) + tl.store(output_ptr + idx, ret) + + +@triton.jit +def cast_to_i32(output_ptr, x_ptr, x_stride, y_stride, z_stride, DIM: tl.constexpr, XB: tl.constexpr, YB: tl.constexpr, + ZB: tl.constexpr): + if DIM == 1: + xidx = tl.arange(0, XB) + idx = xidx * x_stride + elif DIM == 2: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + idx = xidx[:, None] * x_stride + yidx[None, :] * y_stride + elif DIM == 3: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * x_stride + yidx[None, :, None] * y_stride + zidx[None, None, :] * z_stride + + X = tl.load(x_ptr + idx) + ret = tl.cast(X, dtype=tl.int32) + tl.store(output_ptr + idx, ret) + + +@triton.jit +def cast_to_i64(output_ptr, x_ptr, x_stride, y_stride, z_stride, DIM: tl.constexpr, XB: tl.constexpr, YB: tl.constexpr, + ZB: tl.constexpr): + if DIM == 1: + xidx = tl.arange(0, XB) + idx = xidx * x_stride + elif DIM == 2: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + idx = xidx[:, None] * x_stride + yidx[None, :] * y_stride + elif DIM == 3: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * x_stride + yidx[None, :, None] * y_stride + zidx[None, None, :] * z_stride + + X = tl.load(x_ptr + idx) + ret = tl.cast(X, dtype=tl.int64) + tl.store(output_ptr + idx, ret) + + +@triton.jit +def cast_to_fp32(output_ptr, x_ptr, x_stride, y_stride, z_stride, DIM: tl.constexpr, XB: tl.constexpr, YB: tl.constexpr, + ZB: tl.constexpr): + if DIM == 1: + xidx = tl.arange(0, XB) + idx = xidx * x_stride + elif DIM == 2: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + idx = xidx[:, None] * x_stride + yidx[None, :] * y_stride + elif DIM == 3: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * x_stride + yidx[None, :, None] * y_stride + zidx[None, None, :] * z_stride + + X = tl.load(x_ptr + idx) + ret = tl.cast(X, dtype=tl.float32) + tl.store(output_ptr + idx, ret) + + +@triton.jit +def cast_to_fp16(output_ptr, x_ptr, x_stride, y_stride, z_stride, DIM: tl.constexpr, XB: tl.constexpr, YB: tl.constexpr, + ZB: tl.constexpr): + if DIM == 1: + xidx = tl.arange(0, XB) + idx = xidx * x_stride + elif DIM == 2: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + idx = xidx[:, None] * x_stride + yidx[None, :] * y_stride + elif DIM == 3: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * x_stride + yidx[None, :, None] * y_stride + zidx[None, None, :] * z_stride + + X = tl.load(x_ptr + idx) + ret = tl.cast(X, dtype=tl.float16) + tl.store(output_ptr + idx, ret) + + +@triton.jit +def cast_to_bf16(output_ptr, x_ptr, x_stride, y_stride, z_stride, DIM: tl.constexpr, XB: tl.constexpr, YB: tl.constexpr, + ZB: tl.constexpr): + if DIM == 1: + xidx = tl.arange(0, XB) + idx = xidx * x_stride + elif DIM == 2: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + idx = xidx[:, None] * x_stride + yidx[None, :] * y_stride + elif DIM == 3: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * x_stride + yidx[None, :, None] * y_stride + zidx[None, None, :] * z_stride + + X = tl.load(x_ptr + idx) + ret = tl.cast(X, dtype=tl.bfloat16) + tl.store(output_ptr + idx, ret) + + +@triton.jit +def cast_to_uint32(output_ptr, x_ptr, x_stride, y_stride, z_stride, DIM: tl.constexpr, XB: tl.constexpr, + YB: tl.constexpr, ZB: tl.constexpr): + if DIM == 1: + xidx = tl.arange(0, XB) + idx = xidx * x_stride + elif DIM == 2: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + idx = xidx[:, None] * x_stride + yidx[None, :] * y_stride + elif DIM == 3: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * x_stride + yidx[None, :, None] * y_stride + zidx[None, None, :] * z_stride + + X = tl.load(x_ptr + idx) + ret = tl.cast(X, dtype=tl.uint32) + tl.store(output_ptr + idx, ret) + + +@triton.jit +def cast_to_int64(output_ptr, x_ptr, x_stride, y_stride, z_stride, DIM: tl.constexpr, XB: tl.constexpr, + YB: tl.constexpr, ZB: tl.constexpr): + if DIM == 1: + xidx = tl.arange(0, XB) + idx = xidx * x_stride + elif DIM == 2: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + idx = xidx[:, None] * x_stride + yidx[None, :] * y_stride + elif DIM == 3: + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * x_stride + yidx[None, :, None] * y_stride + zidx[None, None, :] * z_stride + + X = tl.load(x_ptr + idx) + ret = tl.cast(X, dtype=tl.int64) + tl.store(output_ptr + idx, ret) + + +triton_func_map = { + "bool": cast_to_bool, "int8": cast_to_i8, "int16": cast_to_i16, "int32": cast_to_i32, "float16": cast_to_fp16, + "bfloat16": cast_to_bf16, "float32": cast_to_fp32, "uint32": cast_to_uint32, "int64": cast_to_int64 +} + + +def structParam(x0): + dim = x0.dim() + stride0, stride1, stride2 = 0, 0, 0 + shape0, shape1, shape2 = 0, 0, 0 + if dim >= 1: + stride0 = x0.stride(0) + shape0 = x0.shape[0] + if dim >= 2: + stride1 = x0.stride(1) + shape1 = x0.shape[1] + if dim == 3: + stride2 = x0.stride(2) + shape2 = x0.shape[2] + return dim, stride0, stride1, stride2, shape0, shape1, shape2 + + +@pytest.mark.parametrize('shape', random.sample(TestUtils.full_shape, 5)) +@pytest.mark.parametrize('srcDtype', TestUtils.full_dtype) +@pytest.mark.parametrize('dstDtype', TestUtils.full_dtype) +def test_cast(srcDtype, dstDtype, shape): + if srcDtype == dstDtype: + return + srcBytes = get_dtype_size(srcDtype) + dstBytes = get_dtype_size(dstDtype) + dtype_size = max(srcBytes, dstBytes) + if dstDtype == 'int8': + if dtype_size * math.prod(shape) >= (TestUtils.ub_size / 100): + print(f"srcDtype:{srcDtype}, dstDtype:{dstDtype}, shape:{shape} mem overflow") + return + elif dtype_size * math.prod(shape) >= (TestUtils.ub_size / 12): + print(f"srcDtype:{srcDtype}, dstDtype:{dstDtype}, shape:{shape} mem overflow") + return + + x0 = test_common.generate_tensor(shape, srcDtype) + torch_res = x0.to(eval("torch." + dstDtype)) + x0 = x0.npu() + triton_func = triton_func_map.get(dstDtype, None) + assert triton_func is not None, f"triton_func not Found, srcDtype:{srcDtype}, dstDtype:{dstDtype}" + triton_res = torch.empty(shape, dtype=eval("torch." + dstDtype)).npu() + dim, stride0, stride1, stride2, XB, YB, ZB = structParam(x0) + assert 0 <= dim <= 3, f"dim out of range [0, 3], dim:{dim}" + triton_func[1, 1, 1](triton_res, x0, stride0, stride1, stride2, dim, XB, YB, ZB) + test_common.validate_cmp(dstDtype, triton_res, torch_res) + + +@triton.jit +def cast_to_multi_d(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr): + dtype = output_ptr.type.element_ty + + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if (YB * ZB * MB * NB) > 1: + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if (ZB * MB * NB) > 1: + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if (MB * NB) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if NB > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + X = tl.load(x_ptr + offsets) + ret = tl.cast(X, dtype=dtype) + + tl.store(output_ptr + offsets, ret) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (6, 2, 4, 2), + (4, 2, 8, 4), + (4, 3, 8, 4), +]) +@pytest.mark.parametrize('srcDtype', ['int8', 'float16', 'float32']) +@pytest.mark.parametrize('dstDtype', ['int8', 'float16', 'float32']) +def test_cast_4d(srcDtype, dstDtype, shape): + if srcDtype == dstDtype: + return + srcBytes = get_dtype_size(srcDtype) + dstBytes = get_dtype_size(dstDtype) + dtype_size = max(srcBytes, dstBytes) + if dstDtype == 'int8': + if dtype_size * math.prod(shape) >= (TestUtils.ub_size / 100): + print(f"srcDtype:{srcDtype}, dstDtype:{dstDtype}, shape:{shape} mem overflow") + return + elif dtype_size * math.prod(shape) >= (TestUtils.ub_size / 12): + print(f"srcDtype:{srcDtype}, dstDtype:{dstDtype}, shape:{shape} mem overflow") + return + + x0 = test_common.generate_tensor(shape, srcDtype) + torch_res = x0.to(eval("torch." + dstDtype)) + x0 = x0.npu() + + triton_res = torch.empty(shape, dtype=eval("torch." + dstDtype)).npu() + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + grid = (1, ) + cast_to_multi_d[grid](triton_res, x0, *triton_shape) + test_common.validate_cmp(dstDtype, triton_res, torch_res) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (2, 6, 2, 4, 2), + (2, 4, 2, 8, 4), + (3, 4, 2, 8, 4), +]) +@pytest.mark.parametrize('srcDtype', ['int8', 'float16', 'float32']) +@pytest.mark.parametrize('dstDtype', ['int8', 'float16', 'float32']) +def test_cast_5d(srcDtype, dstDtype, shape): + if srcDtype == dstDtype: + return + srcBytes = get_dtype_size(srcDtype) + dstBytes = get_dtype_size(dstDtype) + dtype_size = max(srcBytes, dstBytes) + if dstDtype == 'int8': + if dtype_size * math.prod(shape) >= (TestUtils.ub_size / 100): + print(f"srcDtype:{srcDtype}, dstDtype:{dstDtype}, shape:{shape} mem overflow") + return + elif dtype_size * math.prod(shape) >= (TestUtils.ub_size / 12): + print(f"srcDtype:{srcDtype}, dstDtype:{dstDtype}, shape:{shape} mem overflow") + return + + x0 = test_common.generate_tensor(shape, srcDtype) + torch_res = x0.to(eval("torch." + dstDtype)) + x0 = x0.npu() + + triton_res = torch.empty(shape, dtype=eval("torch." + dstDtype)).npu() + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + grid = (1, ) + cast_to_multi_d[grid](triton_res, x0, *triton_shape) + test_common.validate_cmp(dstDtype, triton_res, torch_res) + + +if __name__ == "__main__": + for shape in [(3, ), (3, 3), (3, 3, 3)]: + for srcDtype in ['int8', 'float32', 'bool']: + for dstDtype in ['int8', 'float32', 'bool']: + test_cast(srcDtype, dstDtype, shape) diff --git a/third_party/ascend/unittest/generalization_cases/test_cdiv.py b/third_party/ascend/unittest/generalization_cases/test_cdiv.py new file mode 100644 index 0000000000..4f9afe73c6 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_cdiv.py @@ -0,0 +1,162 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common +from test_common import TestUtils +import math +import logging + + +def torch_cdiv(x0, x1, dtype): + return (x0 + x1 - 1) // x1 + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.cdiv(X, Y) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_cdiv_4d_5d(output_ptr, x_ptr, y_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = tl.cdiv(x_val, y_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64']) +def test_case2(dtype, shape): + # 生成数据, cdiv int8 溢出的行为triton与torch_cpu不一致 + x = (test_common.generate_tensor(shape, dtype) // 2).abs().npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + y = (y.abs() // 2 + 1) + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.debug(f"output.dtype={output.dtype}") + + ans = torch_cdiv(x.cpu(), y.cpu(), eval('torch.' + dtype)) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if dtype == 'int8': + if x.numel() * x.element_size() >= 512: + grid = (1, 1, ZB) + ZB = 1 + else: + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64']) +def test_cdiv_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = (test_common.generate_tensor(shape, dtype) // 2).abs().npu() + y = test_common.generate_tensor(shape, dtype).npu() + y = (y.abs() // 2 + 1) + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_cdiv(x.cpu(), y.cpu(), eval('torch.' + dtype)) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_cdiv_4d_5d[grid](output, x, y, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_ceil.py b/third_party/ascend/unittest/generalization_cases/test_ceil.py new file mode 100644 index 0000000000..bb0e925658 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_ceil.py @@ -0,0 +1,157 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time + +import torch +import torch_npu +import test_common +from test_common import TestUtils +import logging +import math + + +def torch_ceil(x0): + res = torch.ceil(x0) + return res + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.ceil(X) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_ceil_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = tl.ceil(x_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_case2(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.debug(f"output.dtype={output.dtype}") + + ans = torch_ceil(x) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_ceil_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_ceil(x) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_ceil_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_common.py b/third_party/ascend/unittest/generalization_cases/test_common.py new file mode 100644 index 0000000000..e6cf112f74 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_common.py @@ -0,0 +1,343 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import os +import re +import torch +import torch_npu +import math +import logging +from typing import AnyStr +import pytest +import functools +import numpy as np + +_float_dtypes = ['float32', 'float16', 'bfloat16'] +_int_dtypes = ['int32', 'int64', 'int16', 'int8'] +_uint_dtypes = ['uint8', 'uint16', 'uint32', 'uint64'] + +log_level = os.getenv("LOG_LEVEL", "WARN").upper() +level_mapping = { + "DEBUG": logging.DEBUG, "INFO": logging.INFO, "WARN": logging.WARNING, "ERROR": logging.ERROR, "CRITICAL": + logging.CRITICAL +} + +logging.basicConfig(level=level_mapping.get(log_level, logging.WARNING), + format="[%(asctime)s][%(levelname)s] %(message)s") + +bisheng_not_support_dtypes = { + 'abs': [], 'eq': [], 'ne': [], 'flip': ['int64', + 'bfloat16'], 'load_store': ['int64'], 'permute2d': ['int64'], 'permute3d': + ['int64'], 'trans2d': ['int64'], 'trans3d': ['int64'], 'matmul': ['int16', 'int32', 'uint32', 'int64', 'bool'] +} + +tritonascend_not_support_dtypes = { + 'abs': ['bool'], + 'eq': ['bool'], + 'ne': ['bool'], + 'flip': ['bool'], + 'load_store': ['bool'], + 'permute2d': ['bool'], + 'permute3d': ['bool'], + 'trans2d': ['bool'], + 'trans3d': ['bool'], +} + + +def avoid_not_support(op: AnyStr): + + def decorator(test_func): + + @functools.wraps(test_func) + def wrapper(shape, dtype, *args, **kwargs): + if dtype in bisheng_not_support_dtypes.get(op, []): + logging.warn(f'skiped bisheng not support dtype:{dtype}') + return + if dtype in tritonascend_not_support_dtypes.get(op, []): + logging.warn(f'skiped triton ascend not support dtype:{dtype}') + return + return test_func(shape, dtype, *args, **kwargs) + + return wrapper + + return decorator + + +def get_shape1d(in_shape1d): + result = [] + for i in in_shape1d: + v = tuple((i, )) + result.append(v) + return result + + +def get_shape2d(in_shape1d, custom_shape): + result = [] + for a in in_shape1d: + for b in custom_shape: + t1 = tuple((a, b)) + t2 = tuple((b, a)) + if t1 not in result: + result.append(t1) + if t2 not in result: + result.append(t2) + return result + + +def get_shape3d(): + return [(1, 22, 39), (27, 1, 39), (27, 22, 1), (23, 1, 1), (1, 23, 1), (1, 1, 23), (37, 5, 3), (2, 29, 4), + (7, 31, 7), (3, 5, 8), (7, 17, 15), (23, 5, 16), (23, 5, 31), (7, 11, 32), (7, 11, 33), (2, 3, 255), + (3, 3, 256), (3, 2, 257)] + + +def get_shape1_2_3d(in_shape1d, custom_shape): + return get_shape1d(in_shape1d) + get_shape2d(in_shape1d, custom_shape) + get_shape3d() + + +class TestUtils: + in_shape1d = [1, 2, 3, 4, 8, 16, 32, 64, 128, 256, 37, 741] + custom_shape = [3, 13, 32, 256] + batch = [1, 2, 3, 4, 5, 8] + test_shape1d = get_shape1d(in_shape1d) + test_shape2d = get_shape2d(in_shape1d, custom_shape) + test_shape3d = [ + (1, 22, 39), + (27, 1, 39), + (27, 22, 1), + (1, 1, 23), + (23, 1, 1), + (1, 23, 1), + (37, 5, 3), + (2, 29, 4), + (7, 31, 7), + (3, 5, 8), + (7, 17, 15), + (25, 5, 16), + (23, 5, 31), + (7, 11, 32), + (7, 11, 33), + (2, 3, 255), + (3, 3, 256), + (3, 2, 257), + ] + test_shape4d = [(8, 4, 8, 8), (1, 11, 16, 2)] + test_shape5d = [(2, 3, 4, 5, 6), (1, 3, 4, 5, 6), (3, 6, 2, 4, 4)] + test_shape6d = [(2, 3, 5, 6, 3, 2)] + test_shape7d = [(1, 2, 3, 4, 3, 2, 2)] + test_shape_ub_overflow = [(10, 50, 1000)] + test_shape8d = [(1, 2, 3, 2, 5, 3, 7, 2), (1, 3, 2, 5, 6, 7, 2, 1), (2, 3, 7, 3, 2, 3, 2, 3)] + full_shape_4_8d = test_shape4d + test_shape5d + test_shape6d + test_shape7d + test_shape8d + + full_shape = test_shape1d + test_shape2d + test_shape3d + test_shape1_2_3d = full_shape + full_dtype = ['int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32', 'bool'] + ub_size = 98304 * 2 + dtype_list = full_dtype + + +def get_dtype_size(dtype): + torch_dtype = eval('torch.' + dtype) + bits = 0 + if torch_dtype == torch.bool: + bits = 8 + elif torch.is_floating_point(torch.tensor(0, dtype=torch_dtype)): + bits = torch.finfo(torch_dtype).bits + else: + bits = torch.iinfo(torch_dtype).bits + return bits // 8 + + +def check_ub_mem_overflow(dtype, shape): + bytes = get_dtype_size(dtype) + if bytes * math.prod(shape) > TestUtils.ub_size: + logging.warning(f'dtype:{dtype} shape:{shape} mem overflow') + return True + return False + + +def generate_numpy(shape, dtype, low=None, high=None): + if dtype in _int_dtypes + _uint_dtypes: + iinfo = np.iinfo(getattr(np, dtype)) + low = iinfo.min if low is None else max(low, iinfo.min) + high = iinfo.max if high is None else min(high, iinfo.max) + dty = getattr(np, dtype) + return np.random.randint(low, high, shape, dtype=dty) + elif dtype == 'float16' or dtype == 'float32': + return np.random.normal(0, 1, shape).astype(dtype) + elif dtype == 'bfloat16': + return (np.random.normal(0, 1, shape).astype('float32').view('uint32') & np.uint32(0xffff0000)).view('float32') + elif dtype == 'bool': + return np.random.randint(low=0, high=2, size=shape).astype(bool) + else: + raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + + +def generate_tensor(shape, dtype): + if dtype == 'float32' or dtype == 'float16' or dtype == 'bfloat16': + return torch.randn(size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16' or dtype == 'uint32': + return torch.randint(low=0, high=2000, size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'int8': + return torch.randint(low=0, high=127, size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'bool': + return torch.randint(low=0, high=2, size=shape).bool() + elif dtype == 'uint8': + return torch.randint(low=0, high=255, size=shape, dtype=torch.uint8) + else: + raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + + +def generate_tensor_int_withSigns(shape, dtype): + if dtype == 'int32' or dtype == 'int64' or dtype == 'int16': + return torch.randint(low=-32768, high=32767, size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'int8': + return torch.randint(low=-128, high=127, size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'bool': + return torch.randint(low=0, high=2, size=shape).bool() + else: + raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + + +def get_triton_sig_typename(dtype): + if dtype == 'float32': + tyname = "*fp32" + elif dtype == 'int32': + tyname = "*i32" + elif dtype == 'int64': + tyname = "*i64" + elif dtype == 'float16': + tyname = "*fp16" + elif dtype == 'int16': + tyname = "*i16" + elif dtype == 'int8': + tyname = "*i8" + elif dtype == 'bool': + tyname = "*i1" + else: + raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + return tyname + + +# Relative error: abs(x_ref - x_cal) / abs(x_ref) +# Absolute error: abs(x_ref - x_cal) + + +# calculation type operators require different error range +# It is a stricter verification and not satisfied now, save it here +def validate_cal(dtype, y_cal, y_ref): + if dtype == 'float16': + if torch.mean(y_ref) < 0.001: + assert torch.abs(y_cal - y_ref) < 0.001, "|y_cal - y_ref| < 0.001 is required !" + else: + diff = torch.div(torch.abs(y_cal - y_ref), torch.abs(y_cal)) < 0.001 + # all true + assert diff.all(), "Relative error is less than 0.001 !" + if dtype == 'float32': + if torch.mean(y_ref) < 0.0001: + assert torch.abs(y_cal - y_ref) < 0.0001, "|y_cal - y_ref| < 0.0001 is required !" + else: + diff = torch.div(torch.abs(y_cal - y_ref), torch.abs(y_cal)) < 0.0001 + assert diff.all(), "Relative error is less than 0.001 !" + elif dtype == 'bfloat16': + diff = torch.div(torch.abs(y_cal - y_ref), torch.abs(y_cal)) < 0.001 + assert diff.all(), "Relative error is less than 0.001 !" + elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16' or dtype == 'int8': + assert torch.equal(y_cal, y_ref) + elif dtype == 'uint8': + assert torch.equal(y_cal, y_ref) + elif dtype == 'bool': + assert torch.equal(y_cal, y_ref) + else: + raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + + +# moving and comparison ops require no precision error +def validate_cmp(dtype, y_cal, y_ref): + y_cal = y_cal.npu() + y_ref = y_ref.npu() + if dtype == 'float16': + torch.testing.assert_close(y_ref, y_cal, rtol=1e-03, atol=1e-03, equal_nan=True) + elif dtype == 'bfloat16': + torch.testing.assert_close(y_ref.to(torch.float32), y_cal.to(torch.float32), rtol=1e-03, atol=1e-03, + equal_nan=True) + elif dtype == 'float32': + torch.testing.assert_close(y_ref, y_cal, rtol=1e-04, atol=1e-04, equal_nan=True) + elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16' or dtype == 'int8': + assert torch.equal(y_cal, y_ref) + elif dtype == 'uint8' or dtype == 'uint16' or dtype == 'uint32' or dtype == 'uint64': + assert torch.equal(y_cal, y_ref) + elif dtype == 'bool': + assert torch.equal(y_cal.cpu(), y_ref.cpu()) + else: + raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + + +def validate_cmp_with_expection(dtype, y_cal, y_ref, expect): + if dtype == 'float32' or dtype == 'float16' or dtype == 'bfloat16': + if expect: + assert torch.allclose(y_ref, y_cal, rtol=1e-03, atol=1e-03, equal_nan=True) + else: + assert not torch.allclose(y_ref, y_cal, rtol=1e-03, atol=1e-03, equal_nan=True) + elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16' or dtype == 'int8' \ + or dtype == 'uint8' or dtype == 'uint16' or dtype == 'uint32' or dtype == 'uint64': + if expect: + assert torch.equal(y_cal, y_ref) + else: + assert not torch.equal(y_cal, y_ref) + else: + raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + + +def raises_with_match(expected_exception, match_pattern): + + def decorator(test_func): + + @functools.wraps(test_func) + def wrapper(*args, **kwargs): + with pytest.raises(expected_exception, match=match_pattern): + return test_func(*args, **kwargs) + + return wrapper + + return decorator + + +def capture_output(expected_output): + + def decorator(test_func): + + @functools.wraps(test_func) + def wrapper(*args, **kwargs): + capsys = kwargs.pop('capsys', None) + if capsys is None: + try: + capsys = pytest.fixture(capsys)() + except: + raise RuntimeError("This decorator requires pytest's capsys fixture") + test_func(capsys, *args, **kwargs) + captured = capsys.readouterr() + # pybind11::scoped_ostream_redirect captures std::cout with \x00 inserted + # for now, no idea how to eliminate \x00 from C++ side. + cleaned = re.sub(r"\x00", "", captured.out) + assert expected_output in cleaned + + return wrapper + + return decorator diff --git a/third_party/ascend/unittest/generalization_cases/test_cos.py b/third_party/ascend/unittest/generalization_cases/test_cos.py new file mode 100644 index 0000000000..9f1980a515 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_cos.py @@ -0,0 +1,157 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time + +import torch +import torch_npu +import test_common +from test_common import TestUtils +import logging +import math + + +def torch_cos(x0): + res = torch.cos(x0) + return res + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.cos(X) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_cos_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = tl.cos(x_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_case2(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.debug(f"output.dtype={output.dtype}") + + ans = torch_cos(x) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_cos_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_cos(x) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_cos_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_count_dim0.py b/third_party/ascend/unittest/generalization_cases/test_count_dim0.py new file mode 100644 index 0000000000..826f649909 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_count_dim0.py @@ -0,0 +1,152 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +# Copyright (c) Huawei Technologies Co., Ltd. 2025-2025. All rights reserved. +import pytest +import triton +import triton.language as tl +import time +import torch +import torch_npu +import test_common +from test_common import TestUtils + + +def standard_count(x0, cmp_val, dim, dtype): + res = (x0 == cmp_val).sum(dim=dim) + return res + + +def standard_count_gt(x0, cmp_val, dim, dtype): + res = (x0 > cmp_val).sum(dim=dim) + return res + + +def standard_count_lt(x0, cmp_val, dim, dtype): + res = (x0 < cmp_val).sum(dim=dim) + return res + + +@triton.jit +def count(in_ptr0, out_ptr0, cmp_val, dim: tl.constexpr, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, + NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, MNUMEL) + nblk_idx = tl.arange(0, N) + tl.program_id(2) * N + mmask = mblk_idx < MNUMEL + nmask = nblk_idx < NNUMEL + mask = (mmask[:, None]) & (nmask[None, :]) + idx = mblk_idx[:, None] * NNUMEL + nblk_idx[None, :] + x = tl.load(in_ptr0 + idx, mask=mask, other=0) + tmp1 = (x == cmp_val) + tmp2 = tmp1.to(tl.float32) + ret = tl.sum(tmp2, dim) + tl.store(out_ptr0 + nblk_idx, ret, mask=nmask) + + +@triton.jit +def count_gt(in_ptr0, out_ptr0, cmp_val, dim: tl.constexpr, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, + NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, MNUMEL) + nblk_idx = tl.arange(0, N) + tl.program_id(2) * N + mmask = mblk_idx < MNUMEL + nmask = nblk_idx < NNUMEL + mask = (mmask[:, None]) & (nmask[None, :]) + idx = mblk_idx[:, None] * NNUMEL + nblk_idx[None, :] + x = tl.load(in_ptr0 + idx, mask=mask, other=0) + tmp1 = (x > cmp_val) + tmp2 = tmp1.to(tl.float32) + ret = tl.sum(tmp2, dim) + tl.store(out_ptr0 + nblk_idx, ret, mask=nmask) + + +@triton.jit +def count_lt(in_ptr0, out_ptr0, cmp_val, dim: tl.constexpr, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, + NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, MNUMEL) + nblk_idx = tl.arange(0, N) + tl.program_id(2) * N + mmask = mblk_idx < MNUMEL + nmask = nblk_idx < NNUMEL + mask = (mmask[:, None]) & (nmask[None, :]) + idx = mblk_idx[:, None] * NNUMEL + nblk_idx[None, :] + x = tl.load(in_ptr0 + idx, mask=mask, other=0) + tmp1 = (x < cmp_val) + tmp2 = tmp1.to(tl.float32) + ret = tl.sum(tmp2, dim) + tl.store(out_ptr0 + nblk_idx, ret, mask=nmask) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', ['int8']) +def test_count_dim0_common(shape, dtype): + rblock = shape[1] + xblock = shape[0] + x0 = test_common.generate_tensor(shape, dtype).npu() + + if dtype == torch.int8: + cmp_val = 8 + else: + cmp_val = 0.5 + + ans = standard_count(x0, cmp_val, 0, dtype) + + output = torch.zeros((shape[1], ), dtype=torch.float32).npu() + count[1, 1, rblock](x0, output, cmp_val, 0, xblock, 1, xblock, rblock) + + test_common.validate_cmp("float32", output, ans.to(torch.float32)) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'int8']) +def test_count_gt_dim0_common(shape, dtype): + rblock = shape[1] + xblock = shape[0] + x0 = test_common.generate_tensor(shape, dtype).npu() + + if dtype == torch.int8: + cmp_val = 8 + else: + cmp_val = 0.5 + + ans = standard_count_gt(x0, cmp_val, 0, dtype) + + output = torch.zeros((shape[1], ), dtype=torch.float32).npu() + count_gt[1, 1, rblock](x0, output, cmp_val, 0, xblock, 1, xblock, rblock) + + test_common.validate_cmp("float32", output, ans.to(torch.float32)) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'int8']) +def test_count_lt_dim0_common(shape, dtype): + rblock = shape[1] + xblock = shape[0] + x0 = test_common.generate_tensor(shape, dtype).npu() + + if dtype == torch.int8: + cmp_val = 8 + else: + cmp_val = 0.5 + + ans = standard_count_lt(x0, cmp_val, 0, dtype) + + output = torch.zeros((shape[1], ), dtype=torch.float32).npu() + count_lt[1, 1, rblock](x0, output, cmp_val, 0, xblock, 1, xblock, rblock) + + test_common.validate_cmp("float32", output, ans.to(torch.float32)) diff --git a/third_party/ascend/unittest/generalization_cases/test_count_dim1.py b/third_party/ascend/unittest/generalization_cases/test_count_dim1.py new file mode 100644 index 0000000000..ebd19cf7ab --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_count_dim1.py @@ -0,0 +1,151 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import time +import torch +import torch_npu +import test_common +from test_common import TestUtils + + +def standard_count(x0, cmp_val, dim, dtype): + res = (x0 == cmp_val).sum(dim=dim) + return res + + +def standard_count_gt(x0, cmp_val, dim, dtype): + res = (x0 > cmp_val).sum(dim=dim) + return res + + +def standard_count_lt(x0, cmp_val, dim, dtype): + res = (x0 < cmp_val).sum(dim=dim) + return res + + +@triton.jit +def count(in_ptr0, out_ptr0, cmp_val, dim: tl.constexpr, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, + NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, M) + tl.program_id(1) * M + nblk_idx = tl.arange(0, NNUMEL) + mmask = mblk_idx < MNUMEL + nmask = nblk_idx < NNUMEL + mask = (mmask[:, None]) & (nmask[None, :]) + idx = mblk_idx[:, None] * NNUMEL + nblk_idx[None, :] + x = tl.load(in_ptr0 + idx, mask=mask, other=0) + tmp1 = (x == cmp_val) + tmp2 = tmp1.to(tl.float32) + ret = tl.sum(tmp2, dim) + tl.store(out_ptr0 + mblk_idx, ret, mask=mmask) + + +@triton.jit +def count_gt(in_ptr0, out_ptr0, cmp_val, dim: tl.constexpr, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, + NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, M) + tl.program_id(1) * M + nblk_idx = tl.arange(0, NNUMEL) + mmask = mblk_idx < MNUMEL + nmask = nblk_idx < NNUMEL + mask = (mmask[:, None]) & (nmask[None, :]) + idx = mblk_idx[:, None] * NNUMEL + nblk_idx[None, :] + x = tl.load(in_ptr0 + idx, mask=mask, other=0) + tmp1 = (x > cmp_val) + tmp2 = tmp1.to(tl.float32) + ret = tl.sum(tmp2, dim) + tl.store(out_ptr0 + mblk_idx, ret, mask=mmask) + + +@triton.jit +def count_lt(in_ptr0, out_ptr0, cmp_val, dim: tl.constexpr, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, + NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, M) + tl.program_id(1) * M + nblk_idx = tl.arange(0, NNUMEL) + mmask = mblk_idx < MNUMEL + nmask = nblk_idx < NNUMEL + mask = (mmask[:, None]) & (nmask[None, :]) + idx = mblk_idx[:, None] * NNUMEL + nblk_idx[None, :] + x = tl.load(in_ptr0 + idx, mask=mask, other=0) + tmp1 = (x < cmp_val) + tmp2 = tmp1.to(tl.float32) + ret = tl.sum(tmp2, dim) + tl.store(out_ptr0 + mblk_idx, ret, mask=mmask) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', ['int8']) +def test_count_dim1_common(shape, dtype): + rblock = shape[1] + xblock = shape[0] + x0 = test_common.generate_tensor(shape, dtype).npu() + + if dtype == torch.int8: + cmp_val = 8 + else: + cmp_val = 0.5 + + ans = standard_count(x0, cmp_val, 1, dtype) + + output = torch.zeros((shape[0], ), dtype=torch.float32).npu() + count[1, xblock, 1](x0, output, cmp_val, 1, 1, rblock, xblock, rblock) + + test_common.validate_cmp("float32", output, ans.to(torch.float32)) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'int8']) +def test_count_gt_dim1_common(shape, dtype): + rblock = shape[1] + xblock = shape[0] + x0 = test_common.generate_tensor(shape, dtype).npu() + + if dtype == torch.int8: + cmp_val = 8 + else: + cmp_val = 0.5 + + ans = standard_count_gt(x0, cmp_val, 1, dtype) + + output = torch.zeros((shape[0], ), dtype=torch.float32).npu() + count_gt[1, xblock, 1](x0, output, cmp_val, 1, 1, rblock, xblock, rblock) + + test_common.validate_cmp("float32", output, ans.to(torch.float32)) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'int8']) +def test_count_lt_dim1_common(shape, dtype): + rblock = shape[1] + xblock = shape[0] + x0 = test_common.generate_tensor(shape, dtype).npu() + + if dtype == torch.int8: + cmp_val = 8 + else: + cmp_val = 0.5 + + ans = standard_count_lt(x0, cmp_val, 1, dtype) + + output = torch.zeros((shape[0], ), dtype=torch.float32).npu() + count_lt[1, xblock, 1](x0, output, cmp_val, 1, 1, rblock, xblock, rblock) + + test_common.validate_cmp("float32", output, ans.to(torch.float32)) diff --git a/third_party/ascend/unittest/generalization_cases/test_cumprod.py b/third_party/ascend/unittest/generalization_cases/test_cumprod.py new file mode 100644 index 0000000000..9af5216d9b --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_cumprod.py @@ -0,0 +1,254 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import pytest +import torch +import torch_npu +import triton +import triton.language as tl +from triton.runtime.libentry import libentry + +from test_common import TestUtils, validate_cmp, get_dtype_size + + +def torch_func(x, dim, reverse): + is_bf16 = x.dtype == torch.bfloat16 + if is_bf16: + x = x.to(torch.float32) + if reverse: + x = torch.flip(x, [dim]) + res = torch.cumprod(x, dim=dim) + if is_bf16: + res = res.to(torch.bfloat16) + return res + + +@libentry() +@triton.jit +def triton_kernel_1d( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + numel_x: tl.constexpr, + XBLOCK: tl.constexpr, +): + tl.static_assert(numel_x == XBLOCK, "numel_x must be equal to XBLOCK in this kernel") + idx = tl.arange(0, XBLOCK) + x = tl.load(in_ptr0 + idx) + ret = tl.cumprod(x, axis=dim, reverse=reverse) + tl.store(out_ptr0 + idx, ret) + + +@libentry() +@triton.jit +def triton_kernel_2d( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + numel_x: tl.constexpr, + numel_r: tl.constexpr, + XBLOCK: tl.constexpr, + RBLOCK: tl.constexpr, +): + tl.static_assert(numel_x == XBLOCK, "numel_x must be equal to XBLOCK in this kernel") + tl.static_assert(numel_r == RBLOCK, "numel_r must be equal to RBLOCK in this kernel") + idx_x = tl.arange(0, XBLOCK) + idx_r = tl.arange(0, RBLOCK) + idx = idx_x[:, None] * numel_r + idx_r[None, :] + x = tl.load(in_ptr0 + idx) + ret = tl.cumprod(x, axis=dim, reverse=reverse) + tl.store(out_ptr0 + idx, ret) + + +@libentry() +@triton.jit +def triton_kernel_3d( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + numel_x: tl.constexpr, + numel_r: tl.constexpr, + numel_z: tl.constexpr, + XBLOCK: tl.constexpr, + RBLOCK: tl.constexpr, + ZBLOCK: tl.constexpr, +): + tl.static_assert(numel_x == XBLOCK, "numel_x must be equal to XBLOCK in this kernel") + tl.static_assert(numel_r == RBLOCK, "numel_r must be equal to RBLOCK in this kernel") + tl.static_assert(numel_z == ZBLOCK, "numel_z must be equal to ZBLOCK in this kernel") + idx_x = tl.arange(0, XBLOCK) + idx_r = tl.arange(0, RBLOCK) + idx_z = tl.arange(0, ZBLOCK) + idx = idx_x[:, None, None] * numel_r * numel_z + idx_r[None, :, None] * numel_z + idx_z[None, None, :] + x = tl.load(in_ptr0 + idx) + ret = tl.cumprod(x, axis=dim, reverse=reverse) + tl.store(out_ptr0 + idx, ret) + + +@libentry() +@triton.jit +def triton_kernel_4d( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + XB: tl.constexpr, + YB: tl.constexpr, + ZB: tl.constexpr, + MB: tl.constexpr, +): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + midx = tl.arange(0, MB) + idx = (xidx[:, None, None, None] * YB * ZB * MB + yidx[None, :, None, None] * ZB * MB + + zidx[None, None, :, None] * MB + midx[None, None, None, :]) + x = tl.load(in_ptr0 + idx) + ret = tl.cumprod(x, axis=dim, reverse=reverse) + tl.store(out_ptr0 + idx, ret) + + +@libentry() +@triton.jit +def triton_kernel_5d( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + XB: tl.constexpr, + YB: tl.constexpr, + ZB: tl.constexpr, + MB: tl.constexpr, + NB: tl.constexpr, +): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + midx = tl.arange(0, MB) + nidx = tl.arange(0, NB) + idx = (xidx[:, None, None, None, None] * YB * ZB * MB * NB + yidx[None, :, None, None, None] * ZB * MB * NB + + zidx[None, None, :, None, None] * MB * NB + midx[None, None, None, :, None] * NB + + nidx[None, None, None, None, :]) + x = tl.load(in_ptr0 + idx) + ret = tl.cumprod(x, axis=dim, reverse=reverse) + tl.store(out_ptr0 + idx, ret) + + +def convert_cumprod_dtype(x: torch.Tensor) -> torch.Tensor: + """ + 根据 cumprod 类型转换规则，返回转换后的张量。 + """ + dtype_map = { + torch.int8: torch.int64, + torch.int16: torch.int64, + torch.int32: torch.int64, + torch.int64: torch.int64, + torch.bfloat16: torch.bfloat16, + torch.float16: torch.float16, + torch.float32: torch.float32, + torch.bool: torch.int64, + } + + target_dtype = dtype_map.get(x.dtype, None) + if target_dtype is None: + raise ValueError(f"Unsupported input dtype for cumprod conversion: {x.dtype}") + + return x.to(target_dtype) + + +def triton_func(x, dim, reverse): + x = convert_cumprod_dtype(x) + + res = torch.empty_like(x) + shape = x.size() + if len(shape) == 1: + if dim >= 1: + pytest.skip("dim >= 1 for 1D tensor, skipping.") + triton_kernel_1d[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[0]) + elif len(shape) == 2: + if dim >= 2: + pytest.skip("dim >= 2 for 2D tensor, skipping.") + triton_kernel_2d[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[1], x.shape[0], x.shape[1]) + elif len(shape) == 3: + if dim >= 3: + pytest.skip("dim >= 3 for 3D tensor, skipping.") + triton_kernel_3d[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[1], x.shape[2], x.shape[0], x.shape[1], + x.shape[2]) + elif len(shape) == 4: + if dim >= 4: + pytest.skip("dim >= 4 for 4D tensor, skipping.") + triton_kernel_4d[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[1], x.shape[2], x.shape[3]) + elif len(shape) == 5: + if dim >= 5: + pytest.skip("dim >= 5 for 5D tensor, skipping.") + triton_kernel_5d[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[1], x.shape[2], x.shape[3], x.shape[4]) + else: + pytest.skip(f"Unsupported tensor dimension: {len(shape)}") + + return res + + +def cumprod_generate_tensor(shape, dtype): + if dtype == 'float32' or dtype == 'float16' or dtype == 'bfloat16': + return torch.rand(size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16': + return torch.randint(low=1, high=5, size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'int8': + return torch.randint(low=1, high=5, size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'bool': + return torch.randint(low=0, high=2, size=shape).bool() + else: + raise ValueError(f"Unsupported dtype: {dtype}") + + +def should_skip_due_to_mem(dtype, shape): + dtype_size = get_dtype_size(dtype) + total_mem = dtype_size * math.prod(shape) + + if dtype in ('int8', 'bool'): + threshold = TestUtils.ub_size / 13 + else: + threshold = TestUtils.ub_size / 6 + + if total_mem >= threshold: + pytest.skip(f"dtype:{dtype} shape:{shape} mem overflow") + + +# reverse=True not support; +@pytest.mark.parametrize("dtype", TestUtils.full_dtype) +@pytest.mark.parametrize("shape", TestUtils.full_shape) +@pytest.mark.parametrize("dim", [0, 1, 2, 3, 4]) +@pytest.mark.parametrize("reverse", [False]) +def test_cumprod(dtype, shape, dim, reverse): + should_skip_due_to_mem(dtype, shape) + + x = cumprod_generate_tensor(shape=shape, dtype=dtype) + x_npu = x.npu() + + triton_res = triton_func(x_npu, dim, reverse) + + x_gold = x + cpu_res = torch_func(x_gold, dim, reverse) + + validate_cmp(dtype, triton_res, cpu_res) diff --git a/third_party/ascend/unittest/generalization_cases/test_cumsum.py b/third_party/ascend/unittest/generalization_cases/test_cumsum.py new file mode 100644 index 0000000000..06ef04eceb --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_cumsum.py @@ -0,0 +1,253 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import pytest +import torch +import torch_npu +import triton +import triton.language as tl +from triton.runtime.libentry import libentry + +import acc_util +import test_common +from test_common import TestUtils, get_dtype_size + + +def torch_func(x, dim, reverse): + if reverse: + x = torch.flip(x, [dim]) + res = torch.cumsum(x, dim=dim) + return res + + +@libentry() +@triton.jit +def triton_kernel_1d( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + numel_x: tl.constexpr, + XBLOCK: tl.constexpr, +): + tl.static_assert(numel_x == XBLOCK, "numel_x must be equal to XBLOCK in this kernel") + idx = tl.arange(0, XBLOCK) + x = tl.load(in_ptr0 + idx) + ret = tl.cumsum(x, axis=dim, reverse=reverse) + tl.store(out_ptr0 + idx, ret) + + +@libentry() +@triton.jit +def triton_kernel_2d( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + numel_x: tl.constexpr, + numel_r: tl.constexpr, + XBLOCK: tl.constexpr, + RBLOCK: tl.constexpr, +): + tl.static_assert(numel_x == XBLOCK, "numel_x must be equal to XBLOCK in this kernel") + tl.static_assert(numel_r == RBLOCK, "numel_r must be equal to RBLOCK in this kernel") + idx_x = tl.arange(0, XBLOCK) + idx_r = tl.arange(0, RBLOCK) + idx = idx_x[:, None] * numel_r + idx_r[None, :] + x = tl.load(in_ptr0 + idx) + ret = tl.cumsum(x, axis=dim, reverse=reverse) + tl.store(out_ptr0 + idx, ret) + + +@libentry() +@triton.jit +def triton_kernel_3d( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + numel_x: tl.constexpr, + numel_r: tl.constexpr, + numel_z: tl.constexpr, + XBLOCK: tl.constexpr, + RBLOCK: tl.constexpr, + ZBLOCK: tl.constexpr, +): + tl.static_assert(numel_x == XBLOCK, "numel_x must be equal to XBLOCK in this kernel") + tl.static_assert(numel_r == RBLOCK, "numel_r must be equal to RBLOCK in this kernel") + tl.static_assert(numel_z == ZBLOCK, "numel_z must be equal to ZBLOCK in this kernel") + idx_x = tl.arange(0, XBLOCK) + idx_r = tl.arange(0, RBLOCK) + idx_z = tl.arange(0, ZBLOCK) + idx = idx_x[:, None, None] * numel_r * numel_z + idx_r[None, :, None] * numel_z + idx_z[None, None, :] + x = tl.load(in_ptr0 + idx) + ret = tl.cumsum(x, axis=dim, reverse=reverse) + tl.store(out_ptr0 + idx, ret) + + +@libentry() +@triton.jit +def triton_kernel_4d( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + XB: tl.constexpr, + YB: tl.constexpr, + ZB: tl.constexpr, + MB: tl.constexpr, +): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + midx = tl.arange(0, MB) + idx = (xidx[:, None, None, None] * YB * ZB * MB + yidx[None, :, None, None] * ZB * MB + + zidx[None, None, :, None] * MB + midx[None, None, None, :]) + x = tl.load(in_ptr0 + idx) + ret = tl.cumsum(x, axis=dim, reverse=reverse) + tl.store(out_ptr0 + idx, ret) + + +@libentry() +@triton.jit +def triton_kernel_5d( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + XB: tl.constexpr, + YB: tl.constexpr, + ZB: tl.constexpr, + MB: tl.constexpr, + NB: tl.constexpr, +): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + midx = tl.arange(0, MB) + nidx = tl.arange(0, NB) + idx = (xidx[:, None, None, None, None] * YB * ZB * MB * NB + yidx[None, :, None, None, None] * ZB * MB * NB + + zidx[None, None, :, None, None] * MB * NB + midx[None, None, None, :, None] * NB + + nidx[None, None, None, None, :]) + x = tl.load(in_ptr0 + idx) + ret = tl.cumsum(x, axis=dim, reverse=reverse) + tl.store(out_ptr0 + idx, ret) + + +def convert_cumsum_dtype(x: torch.Tensor) -> torch.Tensor: + """ + 根据 cumsum 类型转换规则，返回转换后的张量。 + """ + dtype_map = { + torch.int8: torch.int64, + torch.int16: torch.int64, + torch.int32: torch.int64, + torch.int64: torch.int64, + torch.bfloat16: torch.bfloat16, + torch.float16: torch.float16, + torch.float32: torch.float32, + torch.bool: torch.int64, + } + + target_dtype = dtype_map.get(x.dtype, None) + if target_dtype is None: + raise ValueError(f"Unsupported input dtype for cumsum conversion: {x.dtype}") + + return x.to(target_dtype) + + +def triton_func(x, dim, reverse): + x = convert_cumsum_dtype(x) + + res = torch.empty_like(x) + shape = x.size() + if len(shape) == 1: + if dim >= 1: + pytest.skip("dim >= 1 for 1D tensor, skipping.") + triton_kernel_1d[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[0]) + elif len(shape) == 2: + if dim >= 2: + pytest.skip("dim >= 2 for 2D tensor, skipping.") + triton_kernel_2d[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[1], x.shape[0], x.shape[1]) + elif len(shape) == 3: + if dim >= 3: + pytest.skip("dim >= 3 for 3D tensor, skipping.") + triton_kernel_3d[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[1], x.shape[2], x.shape[0], x.shape[1], + x.shape[2]) + elif len(shape) == 4: + if dim >= 4: + pytest.skip("dim >= 4 for 4D tensor, skipping.") + triton_kernel_4d[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[1], x.shape[2], x.shape[3]) + elif len(shape) == 5: + if dim >= 5: + pytest.skip("dim >= 5 for 5D tensor, skipping.") + triton_kernel_5d[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[1], x.shape[2], x.shape[3], x.shape[4]) + else: + pytest.skip(f"Unsupported tensor dimension: {len(shape)}") + + return res + + +def cumsum_generate_tensor(shape, dtype): + if dtype == 'float32' or dtype == 'float16' or dtype == 'bfloat16': + return torch.rand(size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16': + return torch.randint(low=0, high=3, size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'int8': + return torch.randint(low=0, high=3, size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'bool': + return torch.randint(low=0, high=2, size=shape).bool() + else: + raise ValueError(f"Unsupported dtype: {dtype}") + + +def should_skip_due_to_mem(dtype, shape): + dtype_size = get_dtype_size(dtype) + total_mem = dtype_size * math.prod(shape) + + if dtype in ('int8', 'bool'): + threshold = TestUtils.ub_size / 13 + else: + threshold = TestUtils.ub_size / 6 + + if total_mem >= threshold: + pytest.skip(f"dtype:{dtype} shape:{shape} mem overflow") + + +# reverse=True not support; + + +@pytest.mark.parametrize("dtype", TestUtils.full_dtype) +@pytest.mark.parametrize("shape", TestUtils.full_shape) +@pytest.mark.parametrize("dim", [0, 1, 2, 3, 4]) +@pytest.mark.parametrize("reverse", [False]) +def test_cumsum(dtype, shape, dim, reverse): + should_skip_due_to_mem(dtype, shape) + + x = cumsum_generate_tensor(shape=shape, dtype=dtype) + x_npu = x.npu() + + triton_res = triton_func(x_npu, dim, reverse) + + x_gold = x + cpu_res = torch_func(x_gold, dim, reverse) + + test_common.validate_cmp(dtype, triton_res, cpu_res) diff --git a/third_party/ascend/unittest/generalization_cases/test_debug_barrier.py b/third_party/ascend/unittest/generalization_cases/test_debug_barrier.py new file mode 100644 index 0000000000..fb17fcb23f --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_debug_barrier.py @@ -0,0 +1,168 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import numpy as np +import torch +import logging +import pytest +import test_common +from test_common import TestUtils + + +def torch_invert(x0, ddtype): + if 'float' in str(ddtype): + x0 = x0.to(torch.int32) + y_ref = ~x0 + y_ref = y_ref.to(ddtype) + else: + y_ref = ~x0 + return y_ref + + +@triton.jit +def triton_sub(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, + XNUMEL: tl.constexpr, YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = X - Y + tl.debug_barrier() + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_invert_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = ~x_val + tl.debug_barrier() + tl.store(output_ptr + offsets, ret, mask=masks) + + +test_shape_1d_2d_3d = [(1, ), (2, ), (1, 1), (3, 13), (1, 1, 1), (4, 3, 8)] +test_shape_4_5d = [(1, 1, 1, 1), (2, 2, 2, 2), (1, 1, 1, 1, 1), (2, 2, 2, 2, 1)] + + +@pytest.mark.parametrize('shape', test_shape_1d_2d_3d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +def test_sub(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = x - y + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if dtype == 'int8': + if x.numel() * x.element_size() >= 512: + grid = (1, 1, ZB) + ZB = 1 + else: + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + triton_sub[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', test_shape_1d_2d_3d + test_shape_4_5d) +@pytest.mark.parametrize('dtype', ['bool']) +def test_invert_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_invert(x, eval('torch.' + dtype)) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_invert_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_device_print.py b/third_party/ascend/unittest/generalization_cases/test_device_print.py new file mode 100644 index 0000000000..5421db96ed --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_device_print.py @@ -0,0 +1,106 @@ +import torch +import torch_npu +import triton +import triton.language as tl +import pytest +import sys +import os +import subprocess +import tempfile +import textwrap + +os.environ["TRITON_DEVICE_PRINT"] = "1" +os.environ["TRITON_ENABLE_TASKQUEUE"] = "0" + +shape = (8, ) +XS = 8 +XVALS_INT = [ + 0, + torch.iinfo(torch.int8).min, + torch.iinfo(torch.int8).max, + torch.iinfo(torch.int16).min, + torch.iinfo(torch.int16).max, + torch.iinfo(torch.int32).min, + torch.iinfo(torch.int32).max, + torch.iinfo(torch.int32).max + 1 +] + + +@pytest.mark.parametrize('sigtype', ['int32', 'int64', 'int16', 'int8', 'float32', 'float16', 'bfloat16']) +def test_device_print_int32(sigtype): + + with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f: + temp_script = f.name + + f.write( + textwrap.dedent(f""" +import torch +import torch_npu +import triton +import triton.language as tl +import os +import sys + +os.environ["TRITON_DEVICE_PRINT"] = "1" +os.environ["TRITON_ENABLE_TASKQUEUE"] = "0" + +@triton.jit +def triton_kernel(out_ptr0, in_ptr0, in_ptr1, XBLOCK: tl.constexpr): + idx = tl.arange(0, XBLOCK) + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tl.load(in_ptr1 + idx) + tmp2 = tmp0 + tmp1 + tl.device_print("OUTPUT = ", tmp2) + tl.store(out_ptr0 + idx, tmp2) + +def main(): + shape = (8,) + XS = 8 + dtype = torch.{sigtype} + + x0 = torch.zeros(shape, dtype=dtype).npu() + x1 = torch.ones(shape, dtype=dtype).npu() + + XVALS_INT = [0, -128, 127, -32768, 32767, -2147483648, 2147483647, 2147483648] + for i in range(x1.numel()): + x1[i] = XVALS_INT[i] + + out = torch.empty_like(x0) + + triton_kernel[1,](out, x0, x1, XS) + + print("Kernel execution completed") + + return out + +if __name__ == "__main__": + result = main() + print(f"Result shape: {{result.shape}}") + """)) + dtype = eval(f"torch.{sigtype}") + x0 = torch.zeros(shape, dtype=dtype).npu() + x1 = torch.ones(shape, dtype=dtype).npu() + for i in range(x1.numel()): + x1[i] = XVALS_INT[i] + + torch_ref = x0 + x1 + if 'int' in sigtype: + torch_ref_str = ','.join([str(int(val)) for val in torch_ref.cpu().numpy()]) + else: + values = torch_ref.cpu() + if values.dtype == torch.bfloat16: + values = values.float() + torch_ref_str = ','.join([f"{float(val):.6f}" for val in values.numpy()]) + + result = subprocess.run([sys.executable, temp_script], capture_output=True, text=True, env=os.environ.copy()) + + captured_output = result.stdout + "\n=== STDERR ===\n" + result.stderr + + ##with open(f"manual_capture_{sigtype}.txt", "w") as f: + ##f.write(captured_output) + ##f.write(f"torch_ref:{torch_ref_str}") + + if os.path.exists(temp_script): + os.remove(temp_script) + + assert torch_ref_str in captured_output diff --git a/third_party/ascend/unittest/generalization_cases/test_div_rn.py b/third_party/ascend/unittest/generalization_cases/test_div_rn.py new file mode 100644 index 0000000000..f0ce253288 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_div_rn.py @@ -0,0 +1,154 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time + +import torch +import torch_npu +import test_common +from test_common import TestUtils +import logging +import math + + +def torch_divRn(x0, x1): + return x0 / x1 + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.div_rn(X, Y) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_div_rn_4d_5d(output_ptr, x_ptr, y_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = tl.div_rn(x_val, y_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_case2(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.debug(f"output.dtype={output.dtype}") + + ans = torch_divRn(x, y) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_div_rn_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + ans = torch_divRn(x, y) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_div_rn_4d_5d[grid](output, x, y, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_dot_scaled.py b/third_party/ascend/unittest/generalization_cases/test_dot_scaled.py new file mode 100644 index 0000000000..428180db43 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_dot_scaled.py @@ -0,0 +1,459 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import contextlib +import itertools +import re +import math +import textwrap +import os +import inspect +import pathlib +import test_common +import numpy as np +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +from numpy.random import RandomState +from triton.language.extra import libdevice +from triton.tools.get_ascend_devices import is_compile_on_910_95 + + +@triton.jit +def dot_scale_kernel(a_base, stride_a0: tl.constexpr, stride_a1: tl.constexpr, a_scale, b_base, stride_b0: tl.constexpr, + stride_b1: tl.constexpr, b_scale, out, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, type_a: tl.constexpr, type_b: tl.constexpr, acc_num: tl.constexpr): + PACKED_BLOCK_K_A: tl.constexpr = BLOCK_K + PACKED_BLOCK_K_B: tl.constexpr = BLOCK_K + str_a0: tl.constexpr = stride_a0 + a_ptr = a_base + tl.arange(0, BLOCK_M)[:, None] * stride_a0 + tl.arange(0, str_a0)[None, :] * stride_a1 + b_ptr = b_base + tl.arange(0, PACKED_BLOCK_K_B)[:, None] * stride_b0 + tl.arange(0, BLOCK_N)[None, :] * stride_b1 + + a = tl.load(a_ptr) + b = tl.load(b_ptr) + SCALE_BLOCK_K: tl.constexpr = BLOCK_K // 32 + accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + if a_scale is not None: + scale_a_ptr = a_scale + tl.arange(0, BLOCK_M)[:, None] * SCALE_BLOCK_K + tl.arange(0, SCALE_BLOCK_K)[None, :] + a_scale = tl.load(scale_a_ptr) + if b_scale is not None: + scale_b_ptr = b_scale + tl.arange(0, BLOCK_N)[:, None] * SCALE_BLOCK_K + tl.arange(0, SCALE_BLOCK_K)[None, :] + b_scale = tl.load(scale_b_ptr) + accumulator = tl.dot_scaled(a, a_scale, type_a, b, b_scale, type_b, acc=accumulator, out_dtype=tl.float32) + if acc_num is not None: + for _ in range(acc_num): + accumulator = tl.dot_scaled(a, a_scale, type_a, b, b_scale, type_b, acc=accumulator, out_dtype=tl.float32) + + out_ptr = out + tl.arange(0, BLOCK_M)[:, None] * BLOCK_N + tl.arange(0, BLOCK_N)[None, :] + tl.store(out_ptr, accumulator.to(a.dtype)) + + +def golden_ref(x, scale_x, y, scale_y): + shape_expand_x = x.shape[-1] // scale_x.shape[-1] + if x.dtype == torch.bfloat16: + upscale_x = scale_x.repeat_interleave(shape_expand_x, dim=1).to(torch.int16) + upscale_x = (upscale_x + 127 << 7).view(torch.bfloat16) + else: + scale_fp32 = scale_x.repeat_interleave(shape_expand_x, dim=1).to(torch.int32) + scale_fp32 = (scale_fp32 + 127 << 23).view(torch.float32) + upscale_x = scale_fp32.to(torch.float16) + upscale_y = None + if scale_y is None: + upscale_y = torch.ones_like(y) + else: + scale_y = scale_y.T + shape_expand_y = y.shape[0] // scale_y.shape[0] + if y.dtype == torch.bfloat16: + upscale_y = scale_y.repeat_interleave(shape_expand_y, dim=0).to(torch.int16) + upscale_y = (upscale_y + 127 << 7).view(torch.bfloat16) + else: + scale_fp32 = scale_y.repeat_interleave(shape_expand_y, dim=0).to(torch.int32) + scale_fp32 = (scale_fp32 + 127 << 23).view(torch.float32) + upscale_y = scale_fp32.to(torch.float16) + ret = torch.matmul(x * upscale_x, y * upscale_y) + return ret + + +@pytest.mark.parametrize("M, N, K, rhs_scale, normal_type, acc_num, num_warps", + [(M, N, K, rhs_scale, normal_type, acc_num, 4) + for M, N, K in itertools.product([16, 32, 64, 128], [16, 32, 64, 128], [32, 64]) + for rhs_scale in [False, True] + for normal_type in ["bf16", "fp16"] + for acc_num in [None, 1, 2]]) +def test_scaled_dot(M, N, K, rhs_scale, normal_type, num_warps, acc_num): + device = "npu" + + # The max exponent we use to initialize data in the x/y and associated scale tensor to avoid + # overflow when scaling. + comp_dtype_max_exp = 6 if normal_type == "fp16" else 15 + + torch.manual_seed(0) + + def make_arg(shape, ty): + if ty == "bf16" or ty == "fp16": + comp_dtype = torch.float16 if ty == "fp16" else torch.bfloat16 + ret = torch.randn(shape, dtype=comp_dtype, device=device) + # Clamp to avoid relative error issues + ret.clamp_(-2**comp_dtype_max_exp, 2**comp_dtype_max_exp - 1) + else: + ret = torch.randint(256, shape, dtype=torch.int8, device=device) + return ret + + type_a = normal_type + type_b = type_a + + x = make_arg((M, K), type_a) + y = make_arg((K, N), type_b) + + min_scale, max_scale = (0, 142) if type_a == torch.bfloat16 else (124, 131) + scale_x = torch.randint(min_scale - 128, max_scale - 127, (M, K // 32), dtype=torch.int8, device=device) + min_scale, max_scale = (0, 142) if type_b == torch.bfloat16 else (124, 131) + scale_y = torch.randint(min_scale - 128, max_scale - 127, (N, K // 32), dtype=torch.int8, device=device) + + if not rhs_scale: + scale_y = None + + kernel_kwargs = {"num_warps": num_warps} + z = x.new_empty((M, N), dtype=x.dtype) + pgm = dot_scale_kernel[(1, )](x, *x.stride(), scale_x, y, *y.stride(), scale_y, z, M, N, K, type_a, type_b, acc_num, + **kernel_kwargs) + z_ref = golden_ref(x, scale_x, y, scale_y) + if acc_num is not None: + z_ref = z_ref * (acc_num + 1) + + atol = 1e-5 + rtol = 1e-2 + torch.testing.assert_close(z, z_ref, atol=atol, rtol=rtol) + + +@pytest.mark.parametrize("B, M, N, K", [(1, 32, 64, 64)]) +def test_4d_dot(B, M, N, K): + device = "npu" + torch.manual_seed(0) + + x4d = torch.randn((B, B, M, N), dtype=torch.float16, device=device) + y4d = torch.randn((B, B, N, K), dtype=torch.float16, device=device) + + x2d = x4d.view(-1, N) # shape (B*B*M, N) + y2d = y4d.view(-1, K) # shape (B*B*N, K) + scale_x = torch.randint(-10, 10, (x2d.shape[0], N // 32), dtype=torch.int8, device=device) + scale_y = torch.randint(-10, 10, (y2d.shape[1], N // 32), dtype=torch.int8, device=device) + + z = torch.empty((x2d.shape[0], y2d.shape[0]), dtype=x2d.dtype, device=device) + acc_num = None + dot_scale_kernel[(1, )](x2d, *x2d.stride(), scale_x, y2d, *y2d.stride(), None, z, x2d.shape[0], y2d.shape[0], K, + "fp16", "fp16", None, num_warps=4) + z_ref = golden_ref(x2d, scale_x, y2d, None) + if acc_num is not None: + z_ref = z_ref * (acc_num + 1) + + atol = 1e-5 + rtol = 1e-2 + torch.testing.assert_close(z, z_ref, atol=atol, rtol=rtol) + + +@pytest.mark.parametrize("B, M, N, K", [(2, 16, 16, 32)]) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, + r"lhs last dimension .* must equal rhs penultimate dimension") +def test_2d_dot_invaild_shape(B, M, N, K): + device = "npu" + torch.manual_seed(0) + + x4d = torch.randn((B, B, M, N), dtype=torch.float16, device=device) + y4d = torch.randn((B, B, N, K), dtype=torch.float16, device=device) + + x2d = x4d.view(-1, N) # shape (B*B*M, N) + y2d = y4d.view(-1, K) # shape (B*B*N, K) + scale_x = torch.randint(-10, 10, (x2d.shape[0], N // 32), dtype=torch.int8, device=device) + scale_y = torch.randint(-10, 10, (y2d.shape[1], N // 32), dtype=torch.int8, device=device) + + z = torch.empty((x2d.shape[0], y2d.shape[0]), dtype=x2d.dtype, device=device) + acc_num = None + dot_scale_kernel[(1, )](x2d, *x2d.stride(), scale_x, y2d, *y2d.stride(), None, z, x2d.shape[0], y2d.shape[0], K, + "fp16", "fp16", None, num_warps=4) + + +VALID_MAIN_DTYPES = { + torch.float16, # fp16 + torch.bfloat16, # bf16 +} + +ALL_DTYPES = { + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.float32, # fp32 + torch.bool, +} +ILLEGAL_MAIN_DTYPES = ALL_DTYPES - VALID_MAIN_DTYPES + +ILLEGAL_SCALE_DTYPES = { + torch.int16, + torch.int32, + torch.int64, + torch.float16, + torch.float32, + torch.bfloat16, + torch.bool, +} + +from itertools import product + + +def is_legal_dtype(lhs_dtype, rhs_dtype, lhs_scale_dtype, rhs_scale_dtype): + return (lhs_dtype in VALID_MAIN_DTYPES and rhs_dtype in VALID_MAIN_DTYPES and lhs_scale_dtype is torch.int8 + and rhs_scale_dtype is torch.int8) + + +illegal_cases = [] +for lhs, rhs, lhs_s, rhs_s in product( + VALID_MAIN_DTYPES | ILLEGAL_MAIN_DTYPES, + VALID_MAIN_DTYPES | ILLEGAL_MAIN_DTYPES, + {torch.int8} | ILLEGAL_SCALE_DTYPES, + {torch.int8} | ILLEGAL_SCALE_DTYPES, +): + + if not is_legal_dtype(lhs, rhs, lhs_s, rhs_s): + illegal_cases.append((lhs, rhs, lhs_s, rhs_s)) + +illegal_cases = sorted(set(illegal_cases), key=lambda t: tuple(str(i) for i in t)) + + +@pytest.mark.parametrize( + "lhs_dtype, rhs_dtype, lhs_scale_dtype, rhs_scale_dtype", + illegal_cases, +) +@test_common.raises_with_match(Exception, r"(?i)invalid|unsupported|dtype") +def test_invalid_dtype_should_fail(lhs_dtype, rhs_dtype, lhs_scale_dtype, rhs_scale_dtype): + device = "npu" + M, N, K = 32, 32, 64 + num_warps = 4 + + def make_tensor(shape, dtype): + return torch.randn(shape, dtype=dtype, device=device) \ + if dtype.is_floating_point else \ + torch.randint(-10, 10, shape, dtype=dtype, device=device) + + def make_scale(shape, dtype): + return torch.randint(-10, 10, shape, dtype=dtype, device=device) + + x = make_tensor((M, K), lhs_dtype) + y = make_tensor((K, N), rhs_dtype) + lhs_scale = make_scale((M, K // 32), lhs_scale_dtype) + rhs_scale = make_scale((N, K // 32), rhs_scale_dtype) + z = torch.empty((M, N), dtype=lhs_dtype, device=device) + + dot_scale_kernel[(1, )]( + x, + *x.stride(), + lhs_scale, + y, + *y.stride(), + rhs_scale, + z, + M, + N, + K, + str(lhs_dtype).split('.')[-1], + str(rhs_dtype).split('.')[-1], + None, + num_warps=num_warps, + ) + + +@pytest.mark.parametrize( + "M, N, K, col_a, col_b, type_a, type_b, num_warps", + list( + itertools.product([32, 64, 128], # M + [32, 64, 128], # N + [64, 128], # K + [True, False], # col_a + [True, False], # col_b + ["e4m3", "e5m2"], # type_a + ["e4m3", "e5m2"], # type_b + [4] # num_warps + ))) +def test_scaled_dot_fp8(M, N, K, col_a, col_b, type_a, type_b, num_warps): + device = "npu" + if not is_compile_on_910_95: + pytest.skip("Skipping dot_scaled on A2/A3 case") + + @triton.jit + def dot_scale_fp8_kernel(a_base, stride_a0, stride_a1, a_scale, b_base, stride_b0, stride_b1, out, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, type_a: tl.constexpr, + type_b: tl.constexpr): + tl.static_assert(type_b == "e4m3" or type_b == "e5m2", "type_b must be fp8") + IS_FP8: tl.constexpr = type_a == "e4m3" or type_a == "e5m2" + DIV_FACTOR: tl.constexpr = 1 if IS_FP8 else 2 + PACKED_BLOCK_K_A: tl.constexpr = BLOCK_K // DIV_FACTOR + PACKED_BLOCK_K_B: tl.constexpr = BLOCK_K + a_ptr = a_base + tl.arange(0, BLOCK_M)[:, None] * stride_a0 + tl.arange(0, + PACKED_BLOCK_K_A)[None, :] * stride_a1 + b_ptr = b_base + tl.arange(0, PACKED_BLOCK_K_B)[:, None] * stride_b0 + tl.arange(0, + BLOCK_N)[None, :] * stride_b1 + + SCALE_BLOCK_K: tl.constexpr = BLOCK_K // 32 + scale_a_ptr = a_scale + tl.arange(0, BLOCK_M)[:, None] * SCALE_BLOCK_K + tl.arange(0, SCALE_BLOCK_K)[None, :] + + a = tl.load(a_ptr) + b = tl.load(b_ptr) + a_scale = tl.load(scale_a_ptr) + c = tl.dot_scaled(a, a_scale, type_a, b, None, type_b) + out_ptr = out + tl.arange(0, BLOCK_M)[:, None] * BLOCK_N + tl.arange(0, BLOCK_N)[None, :] + tl.store(out_ptr, c.to(tl.bfloat16)) + + @triton.jit + def mxfp_to_bf16_kernel( + x_ptr, + scale_ptr, + mxfp_ptr, + N, + e_bits: tl.constexpr, + m_bits: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + ): + is_fp8: tl.constexpr = e_bits + m_bits == 7 + # fp8: BLOCK_SIZE -> BLOCK_SIZE // 32, 32 + # fp4: BLOCK_SIZE // 2 -> BLOCK_SIZE // 32 , 16 + PARALLEL_DIM: tl.constexpr = BLOCK_SIZE // 32 + LAST_DIM: tl.constexpr = 32 if is_fp8 else 16 + LOAD_SIZE: tl.constexpr = LAST_DIM * PARALLEL_DIM + + offsets = (tl.program_id(0) * LOAD_SIZE + tl.arange(0, PARALLEL_DIM)[:, None] * LAST_DIM + + tl.arange(0, LAST_DIM)[None, :]) + x = tl.load(x_ptr + offsets, mask=offsets < N * LAST_DIM) + + offsets = tl.program_id(0) * PARALLEL_DIM + tl.arange(0, PARALLEL_DIM)[:, None] + scale = tl.load(scale_ptr + offsets, mask=offsets < N) + tl.static_assert(scale.dtype == tl.uint8) + tl.static_assert(x.dtype == tl.uint8) + + scale_bf16 = (scale.to(tl.uint16) << 7).to(tl.bfloat16, bitcast=True) + if is_fp8: + if e_bits == 5 and m_bits == 2: + x_f8 = x.to(tl.float8e5, bitcast=True) + x_bf16 = x_f8.to(tl.bfloat16) + # Preserve infs and nans. FIXME Fp8E5M2_to_Bf16 doesn't preserve them! + non_finite_mask: tl.constexpr = ((1 << e_bits) - 1) << m_bits + non_finite_mask_bf16: tl.constexpr = ((1 << 8) - 1) << 7 + x_bf16 = tl.where( + x & non_finite_mask == non_finite_mask, + (x_bf16.to(tl.uint16, bitcast=True) | non_finite_mask_bf16).to(tl.bfloat16, bitcast=True), + x_bf16, + ) + else: + tl.static_assert(e_bits == 4 and m_bits == 3) + x_f8 = x.to(tl.float8e4nv, bitcast=True) + x_bf16 = x_f8.to(tl.bfloat16) + else: + # e2m1 + em0 = x & 0x70 + em1 = x & 0x7 + x0 = (em0.to(tl.uint16) << 2) | ((x & 0x80).to(tl.uint16) << 8) + x1 = (em1.to(tl.uint16) << (2 + 4)) | ((x & 0x8).to(tl.uint16) << (8 + 4)) + # Three cases: + # 1) x is normal and non-zero: Correct bias + x0 = tl.where((em0 & 0x60) != 0, x0 + ((127 - 1) << 7), x0) + x1 = tl.where((em1 & 0x6) != 0, x1 + ((127 - 1) << 7), x1) + # 2) x is subnormal (x == 0bs001 where s is the sign): Map to +-0.5 in bf16 + x0 = tl.where(em0 == 0x10, 16128 | (x0 & 0x8000), x0) + x1 = tl.where(em1 == 0x1, 16128 | (x1 & 0x8000), x1) + # 3) x is zero, do nothing + x_bf16 = tl.interleave(x0, x1).to(tl.bfloat16, bitcast=True) + # Multiplication preserves infs and NaNs in x_bf16 + mxfp = x_bf16 * scale_bf16 + # If scale is NaN, we encode it as an bf16 inf, so we need to correct for that + mxfp = tl.where(scale == 0xFF, float("nan"), mxfp) + + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + tl.store(mxfp_ptr + offsets, tl.ravel(mxfp), mask=offsets < N * 32) + + def dot_scale_ref(x, scale, y, type_x, type_y): + e_bits, m_bits = {"e2m1": (2, 1), "e4m3": (4, 3), "e5m2": (5, 2)}[type_x] + type_fp8_y = {"e4m3": torch.float8_e4m3fn, "e5m2": torch.float8_e5m2}[type_y] + + comp_dtype = torch.bfloat16 + + x = x.contiguous() + x_upcast = x.new_empty(scale.shape[:-1] + (32 * scale.shape[-1], ), dtype=comp_dtype) + + N = x_upcast.numel() + BLOCK_SIZE = 512 + grid = ((N + BLOCK_SIZE - 1) // BLOCK_SIZE, ) + mxfp_to_bf16_kernel[grid](x, scale, x_upcast, scale.numel(), e_bits, m_bits, BLOCK_SIZE, num_warps=num_warps) + assert x_upcast.isfinite().all() + + y_upcast = y.view(type_fp8_y).to(comp_dtype) + + class AccumulateInFp32: + + def __enter__(self): + self.prev_value = torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False + + def __exit__(self, exc_type, exc_val, exc_tb): + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = self.prev_value + + with AccumulateInFp32(): + return torch.matmul(x_upcast.to(comp_dtype), y_upcast.to(comp_dtype)) + + torch.manual_seed(0) + + def create_uint8(shape, col_major=False, max_val=255): + if col_major: + shape = shape[:-2] + (shape[-1], shape[-2]) + ret = torch.randint(max_val + 1, shape, dtype=torch.uint8, device=device) + if col_major: + ret = ret.mT + return ret + + DIV_FACTOR = 2 if type_a == "e2m1" else 1 + x = create_uint8((M, K // DIV_FACTOR), col_major=col_a) + y = create_uint8((K, N), col_major=col_b) + + # sample scales that don't overflow as otherwise it's implementation defined (underflowing is alright) + # We substract a reasonably high number (64) so that the sum of all the mxfp elements does not overflow + m_bytes = int(type_a[1]) + bias_type_a = 1 << (m_bytes - 1) - 1 + max_exponent_type_a = (1 << m_bytes) - 1 - bias_type_a + scale_x = create_uint8((M, K // 32), max_val=255 - max_exponent_type_a - 64) + + def make_finite(x, dtype): + # e5m2 has too many non-finite values when sampled uniformly (1 / 32) and + # Fp8E5M2_to_Bf16 doesn't preserve NaNs (fixme) + if dtype not in ("e5m2", "e4m3"): + return x + mask = 0x7C if dtype == "e5m2" else 0x7F + finite = torch.arange(x.numel(), device=device, dtype=torch.int32).reshape_as(x) % mask + x_finite = torch.where(x & mask == mask, finite | (0x80 & x), x) + x.copy_(x_finite) + return x + + x = make_finite(x, type_a) + y = make_finite(y, type_b) + + z = x.new_empty((M, N), dtype=torch.bfloat16) + pgm = dot_scale_fp8_kernel[(1, )]( + x, *x.stride(), scale_x, y, *y.stride(), z, M, N, K, type_a, type_b, + num_warps=num_warps) # to compare with "dot_scale_ref(x, scale_x, y, type_a, type_b)" diff --git a/third_party/ascend/unittest/generalization_cases/test_eq.py b/third_party/ascend/unittest/generalization_cases/test_eq.py new file mode 100644 index 0000000000..94292ac21c --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_eq.py @@ -0,0 +1,128 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common +from test_common import TestUtils +import math +import logging + + +def torch_eq(x0, x1): + if x0.dtype != torch.uint32: + return x0 == x1 + else: + return x0.to(torch.float32) == x1.to(torch.float32) + + +@triton.jit +def triton_eq(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x_index = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + x_index, mask=x_index < N) + tmp1 = tl.load(in_ptr1 + x_index, mask=x_index < N) + tmp2 = tmp0 == tmp1 + tl.store(out_ptr0 + x_index, tmp2, mask=x_index < N) + + +@triton.jit +def triton_eq_4d_5d(x_ptr, y_ptr, output_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = x_val == y_val + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1_2_3d) +@pytest.mark.parametrize('dtype', ['bool', 'int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32']) +def test_eq(shape, dtype): + logging.debug(f'dtype:{dtype} shape:{shape}') + # 生成数据 + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + + numel = x0.numel() + ncore = 1 if numel <= 32 else 32 + xblock = math.ceil(numel / ncore) + xblock_sub = numel if numel <= ncore else math.ceil(numel / ncore) + + # torch结果 + torch_res = torch_eq(x0, x1).to(eval('torch.' + dtype)) + # triton结果 + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + N = triton_res.numel() + triton_eq[ncore, 1, 1](x0, x1, triton_res, N, xblock, xblock_sub) + # 比较结果 + torch_res = torch_res if dtype != 'uint32' else torch_res.to(torch.float32) + triton_res = triton_res if dtype != 'uint32' else triton_res.to(torch.float32) + cmp_dtype = dtype if dtype != 'uint32' else 'float32' + test_common.validate_cmp(cmp_dtype, triton_res, torch_res) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +def test_eq_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_eq(x, y).to(eval('torch.' + dtype)) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_eq_4d_5d[grid](x, y, output, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_erf.py b/third_party/ascend/unittest/generalization_cases/test_erf.py new file mode 100644 index 0000000000..b82945f0f0 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_erf.py @@ -0,0 +1,172 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common +from test_common import TestUtils +import math + + +def torch_erf(x0): + res = torch.erf(x0) + return res + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.erf(X) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_erf_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr, + BLOCK_TOTAL: tl.constexpr): + + pid = tl.program_id(0) + start_idx = pid * BLOCK_TOTAL + local_idx = tl.arange(0, BLOCK_TOTAL) + global_idx = start_idx + local_idx + total_elements = SHAPE_0 * SHAPE_1 * SHAPE_2 * SHAPE_3 * SHAPE_4 + masks = global_idx < total_elements + + dim1_base = SHAPE_1 * SHAPE_2 * SHAPE_3 * SHAPE_4 + dim2_base = SHAPE_2 * SHAPE_3 * SHAPE_4 + dim3_base = SHAPE_3 * SHAPE_4 + dim4_base = SHAPE_4 + + idx_0 = (global_idx // dim1_base) % SHAPE_0 + idx_1 = (global_idx // dim2_base) % SHAPE_1 + idx_2 = (global_idx // dim3_base) % SHAPE_2 + idx_3 = (global_idx // dim4_base) % SHAPE_3 + idx_4 = global_idx % SHAPE_4 + + offsets = idx_0 * STRIDE_0 + idx_1 * STRIDE_1 + idx_2 * STRIDE_2 + idx_3 * STRIDE_3 + idx_4 * STRIDE_4 + + x_val = tl.load(x_ptr + offsets, mask=masks) + ret = tl.erf(x_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_case2(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.debug(f"output.dtype={output.dtype}") + + ans = torch_erf(x) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_erf_4d_5d(shape, dtype): + logging.debug(f"Testing erf for shape={shape}, dtype={dtype}") + + x = test_common.generate_tensor(shape, dtype).npu() + output = torch.empty_like(x) + + ans = torch_erf(x) + + shape_5d = list(shape) + strides_5d = list(x.stride()) + while len(shape_5d) < 5: + shape_5d.append(1) + strides_5d.append(1) + + MAX_BLOCK_ELEMENTS = 1024 + total_elements = x.numel() + + block_5d = [1] * 5 + for i in reversed(range(5)): + if shape_5d[i] == 0: + continue + max_block_i = min(shape_5d[i], MAX_BLOCK_ELEMENTS // (torch.prod(torch.tensor(block_5d)).item())) + block_5d[i] = max_block_i + if torch.prod(torch.tensor(block_5d)).item() >= MAX_BLOCK_ELEMENTS: + break + block_total = torch.prod(torch.tensor(block_5d)).item() + + grid = (triton.cdiv(total_elements, block_total), ) + logging.debug(f"Grid={grid}, block_5d={block_5d}, block_total={block_total}") + + triton_erf_4d_5d[grid](output, x, *block_5d, *shape_5d, *strides_5d, block_total) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_exp.py b/third_party/ascend/unittest/generalization_cases/test_exp.py new file mode 100644 index 0000000000..52233f89c8 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_exp.py @@ -0,0 +1,155 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging + +import triton +import triton.language as tl +import numpy as np +import torch +import pytest +import test_common +from test_common import TestUtils +import math + + +def torch_pointwise(x0): + res = torch.exp(x0) + return res + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.exp(X) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_exp_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = tl.exp(x_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_case2(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.debug(f"output.dtype={output.dtype}") + + ans = torch_pointwise(x) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_exp_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_pointwise(x) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_exp_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_exp2.py b/third_party/ascend/unittest/generalization_cases/test_exp2.py new file mode 100644 index 0000000000..b8e8aa3122 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_exp2.py @@ -0,0 +1,153 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common +from test_common import TestUtils +import math + + +def torch_exp2(x0): + res = torch.pow(2, x0, out=None) + return res + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.exp2(X) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_exp2_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = tl.exp2(x_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_case2(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.debug(f"output.dtype={output.dtype}") + + ans = torch_exp2(x) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_exp2_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_exp2(x) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_exp2_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_expand_dims.py b/third_party/ascend/unittest/generalization_cases/test_expand_dims.py new file mode 100644 index 0000000000..f9a85f044c --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_expand_dims.py @@ -0,0 +1,223 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common +from test_common import TestUtils +import logging + + +@triton.jit +def fn_npu_1d(output_ptr, x_ptr, YB: tl.constexpr, ZB: tl.constexpr): + yidx = tl.arange(0, YB) + + X = tl.load(x_ptr + yidx) + + ret = tl.expand_dims(X, 1) + + oidx = yidx[:, None] + tl.arange(0, 1)[None, :] + + tl.store(output_ptr + oidx, ret) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_expand_dims_1d(shape, dtype): + x = test_common.generate_tensor(shape, dtype).npu() + a = x.unsqueeze(1) + + output = torch.randint(1, (shape[0], 1), dtype=eval('torch.' + dtype)).npu() + + fn_npu_1d[1, 1, 1](output, x, YB=shape[0], ZB=1, debug=True) + + torch.testing.assert_close(output, a) + + +@triton.jit +def fn_npu_2d(output_ptr, x_ptr, YB: tl.constexpr, ZB: tl.constexpr): + yoffs = tl.program_id(0) + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + + idx = yidx[:, None] * ZB + zidx[None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.expand_dims(X, 1) + + oidx = yidx[:, None, None] * ZB + tl.arange(0, 1)[None, :, None] + zidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_expand_dims_2d(shape, dtype): + x = test_common.generate_tensor(shape, dtype).npu() + a = x.unsqueeze(1) + + output = torch.randint(1, (shape[0], 1, shape[1]), dtype=eval('torch.' + dtype)).npu() + + if x.numel() * x.element_size() > 8192: + fn_npu_2d[shape[0], 1, 1](output, x, YB=1, ZB=shape[1]) + else: + fn_npu_2d[1, 1, 1](output, x, YB=shape[0], ZB=shape[1]) + + torch.testing.assert_close(output, a) + + +@triton.jit +def fn_npu_3d(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.expand_dims(X, 2) + + oidx = xidx[:, None, None, None] * YB * ZB + yidx[None, :, None, None] * ZB + tl.arange( + 0, 1)[None, None, :, None] + zidx[None, None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +def test_expand_dims_3d(dtype, shape): + x = test_common.generate_tensor(shape, dtype).npu() + a = x.unsqueeze(2) + + output = torch.randint(1, (shape[0], shape[1], 1, shape[2]), dtype=eval('torch.' + dtype)).npu() + + fn_npu_3d[1, 1, 1](output, x, XB=shape[0], YB=shape[1], ZB=shape[2]) + + torch.testing.assert_close(output, a) + + +@triton.jit +def fn_npu_multi_d(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr, DIMS: tl.constexpr, DIM: tl.constexpr): + in_offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if DIMS > 1: + in_offsets = in_offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if DIMS > 2: + in_offsets = in_offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if DIMS > 3: + in_offsets = in_offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if DIMS > 4: + in_offsets = in_offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + X = tl.load(x_ptr + in_offsets) + + ret = tl.expand_dims(X, DIM).reshape(XB * YB * ZB * MB * NB) + + out_offstes = tl.arange(0, XB * YB * ZB * MB * NB) + tl.store(output_ptr + out_offstes, ret) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('dtype', ['int8', 'float16', 'float32']) +@pytest.mark.parametrize('shape', [ + (2, 64, 16, 2), + (8, 8, 4, 2), + (8, 8, 4, 1), +]) +@pytest.mark.parametrize('dim', [-1, 0, 1, 2, 3]) +def test_npu_4d(shape, dtype, dim): + x = test_common.generate_tensor(shape, dtype).npu() + expected = x.unsqueeze(dim) + + output = torch.empty_like(expected) + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + grid = (1, ) + fn_npu_multi_d[grid](output, x, *triton_shape, len(shape), dim) + + torch.testing.assert_close(output, expected) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('dtype', ['int8', 'float16', 'float32']) +@pytest.mark.parametrize('shape', [ + (2, 32, 3, 16, 2), + (8, 8, 3, 4, 2), + (8, 8, 3, 4, 1), +]) +@pytest.mark.parametrize('dim', [-1, 0, 1, 2, 3, 4]) +def test_npu_5d(shape, dtype, dim): + x = test_common.generate_tensor(shape, dtype).npu() + expected = x.unsqueeze(dim) + + output = torch.empty_like(expected) + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + grid = (1, ) + fn_npu_multi_d[grid](output, x, *triton_shape, len(shape), dim) + + torch.testing.assert_close(output, expected) + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.expand_dims(X, 2) + + oidx = xidx[:, None, None, None] * YB * ZB + yidx[None, :, None, None] * ZB + tl.arange( + 0, 1)[None, None, :, None] + zidx[None, None, None, :] + + tl.store(output_ptr + oidx, ret) + + +paras = [ + ('bfloat16', eval('torch.bfloat16'), 1, 255, 8, 8, 4), + ('uint8', eval('torch.uint8'), 1, 125, 1, 256, 16), + ('uint16', eval('torch.uint16'), 1, 256, 2, 2, 3), + ('uint32', eval('torch.uint32'), 1, 256, 8, 8, 4), + ('uint64', eval('torch.uint64'), 1, 256, 8, 8, 4), + ('bool', eval('torch.bool'), 0, 2, 1, 1, 2), +] + + +@pytest.mark.parametrize('para_type,data_type,low,top,XB,YB,ZB', paras) +def test_expand_dims(para_type, data_type, low, top, XB, YB, ZB): + x = torch.randint(low=low, high=top, size=(XB, YB, ZB), dtype=data_type).npu() + a = x.unsqueeze(2) + output = torch.randint(1, (XB, YB, 1, ZB), dtype=data_type).npu() + fn_npu_[1, 1, 1](output, x, XB=XB, YB=YB, ZB=ZB, debug=True) + test_common.validate_cmp(para_type, output, a) diff --git a/third_party/ascend/unittest/generalization_cases/test_fdiv.py b/third_party/ascend/unittest/generalization_cases/test_fdiv.py new file mode 100644 index 0000000000..099a82387f --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_fdiv.py @@ -0,0 +1,155 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common +from test_common import TestUtils +import math + + +def torch_fdiv(x0, x1): + res = x0 / x1 + return res + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.fdiv(X, Y) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_fdiv_4d_5d(output_ptr, x_ptr, y_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = tl.fdiv(x_val, y_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_case2(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.debug(f"output.dtype={output.dtype}") + + ans = torch_fdiv(x, y) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_fdiv_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_fdiv(x, y) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_fdiv_4d_5d[grid](output, x, y, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_flip.py b/third_party/ascend/unittest/generalization_cases/test_flip.py new file mode 100644 index 0000000000..f043d0cbc1 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_flip.py @@ -0,0 +1,29 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common +import logging +from test_common import TestUtils, check_ub_mem_overflow diff --git a/third_party/ascend/unittest/generalization_cases/test_full_op.py b/third_party/ascend/unittest/generalization_cases/test_full_op.py new file mode 100644 index 0000000000..e74314d52a --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_full_op.py @@ -0,0 +1,1096 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import test_common + +from test_common import TestUtils +import torch +import torch_npu +import pytest +import math +import random + + +@triton.jit +def fn_npu_int8_3d(output_ptr, X: tl.constexpr, Y: tl.constexpr, Z: tl.constexpr): + xidx = tl.arange(0, X) + yidx = tl.arange(0, Y) + zidx = tl.arange(0, Z) + ret = tl.full((X, Y, Z), value=100, dtype=tl.int8) + oidx = xidx[:, None, None] * Y * Z + yidx[None, :, None] * Z + zidx[None, None, :] + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int16_3d(output_ptr, X: tl.constexpr, Y: tl.constexpr, Z: tl.constexpr): + xidx = tl.arange(0, X) + yidx = tl.arange(0, Y) + zidx = tl.arange(0, Z) + ret = tl.full((X, Y, Z), value=100, dtype=tl.int16) + oidx = xidx[:, None, None] * Y * Z + yidx[None, :, None] * Z + zidx[None, None, :] + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_uint32_3d(output_ptr, X: tl.constexpr, Y: tl.constexpr, Z: tl.constexpr): + xidx = tl.arange(0, X) + yidx = tl.arange(0, Y) + zidx = tl.arange(0, Z) + ret = tl.full((X, Y, Z), value=100, dtype=tl.uint32) + oidx = xidx[:, None, None] * Y * Z + yidx[None, :, None] * Z + zidx[None, None, :] + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int32_3d(output_ptr, X: tl.constexpr, Y: tl.constexpr, Z: tl.constexpr): + xidx = tl.arange(0, X) + yidx = tl.arange(0, Y) + zidx = tl.arange(0, Z) + ret = tl.full((X, Y, Z), value=100, dtype=tl.int32) + oidx = xidx[:, None, None] * Y * Z + yidx[None, :, None] * Z + zidx[None, None, :] + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int64_3d(output_ptr, X: tl.constexpr, Y: tl.constexpr, Z: tl.constexpr): + xidx = tl.arange(0, X) + yidx = tl.arange(0, Y) + zidx = tl.arange(0, Z) + ret = tl.full((X, Y, Z), value=100, dtype=tl.int64) + oidx = xidx[:, None, None] * Y * Z + yidx[None, :, None] * Z + zidx[None, None, :] + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_fp16_3d(output_ptr, X: tl.constexpr, Y: tl.constexpr, Z: tl.constexpr): + xidx = tl.arange(0, X) + yidx = tl.arange(0, Y) + zidx = tl.arange(0, Z) + ret = tl.full((X, Y, Z), value=100, dtype=tl.float16) + oidx = xidx[:, None, None] * Y * Z + yidx[None, :, None] * Z + zidx[None, None, :] + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_fp32_3d(output_ptr, X: tl.constexpr, Y: tl.constexpr, Z: tl.constexpr): + xidx = tl.arange(0, X) + yidx = tl.arange(0, Y) + zidx = tl.arange(0, Z) + ret = tl.full((X, Y, Z), value=100, dtype=tl.float32) + oidx = xidx[:, None, None] * Y * Z + yidx[None, :, None] * Z + zidx[None, None, :] + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_bf16_3d(output_ptr, X: tl.constexpr, Y: tl.constexpr, Z: tl.constexpr): + xidx = tl.arange(0, X) + yidx = tl.arange(0, Y) + zidx = tl.arange(0, Z) + ret = tl.full((X, Y, Z), value=100, dtype=tl.bfloat16) + oidx = xidx[:, None, None] * Y * Z + yidx[None, :, None] * Z + zidx[None, None, :] + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_bool_3d(output_ptr, X: tl.constexpr, Y: tl.constexpr, Z: tl.constexpr): + xidx = tl.arange(0, X) + yidx = tl.arange(0, Y) + zidx = tl.arange(0, Z) + ret = tl.full((X, Y, Z), value=0, dtype=tl.int1) + oidx = xidx[:, None, None] * Y * Z + yidx[None, :, None] * Z + zidx[None, None, :] + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int8_2d(output_ptr, Y: tl.constexpr, Z: tl.constexpr): + yoffs = tl.program_id(0) * Y + yidx = tl.arange(0, Y) + yoffs + zidx = tl.arange(0, Z) + ret = tl.full((Y, Z), value=100, dtype=tl.int8) + oidx = yidx[:, None] * Z + zidx[None, :] + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int16_2d(output_ptr, Y: tl.constexpr, Z: tl.constexpr): + yoffs = tl.program_id(0) * Y + yidx = tl.arange(0, Y) + yoffs + zidx = tl.arange(0, Z) + ret = tl.full((Y, Z), value=100, dtype=tl.int16) + oidx = yidx[:, None] * Z + zidx[None, :] + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_uint32_2d(output_ptr, Y: tl.constexpr, Z: tl.constexpr): + yoffs = tl.program_id(0) * Y + yidx = tl.arange(0, Y) + yoffs + zidx = tl.arange(0, Z) + ret = tl.full((Y, Z), value=100, dtype=tl.uint32) + oidx = yidx[:, None] * Z + zidx[None, :] + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int32_2d(output_ptr, Y: tl.constexpr, Z: tl.constexpr): + yoffs = tl.program_id(0) * Y + yidx = tl.arange(0, Y) + yoffs + zidx = tl.arange(0, Z) + ret = tl.full((Y, Z), value=100, dtype=tl.int32) + oidx = yidx[:, None] * Z + zidx[None, :] + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int64_2d(output_ptr, Y: tl.constexpr, Z: tl.constexpr): + yoffs = tl.program_id(0) * Y + yidx = tl.arange(0, Y) + yoffs + zidx = tl.arange(0, Z) + ret = tl.full((Y, Z), value=100, dtype=tl.int64) + oidx = yidx[:, None] * Z + zidx[None, :] + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_fp16_2d(output_ptr, Y: tl.constexpr, Z: tl.constexpr): + yoffs = tl.program_id(0) * Y + yidx = tl.arange(0, Y) + yoffs + zidx = tl.arange(0, Z) + ret = tl.full((Y, Z), value=100, dtype=tl.float16) + oidx = yidx[:, None] * Z + zidx[None, :] + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_fp32_2d(output_ptr, Y: tl.constexpr, Z: tl.constexpr): + yoffs = tl.program_id(0) * Y + yidx = tl.arange(0, Y) + yoffs + zidx = tl.arange(0, Z) + ret = tl.full((Y, Z), value=100, dtype=tl.float32) + oidx = yidx[:, None] * Z + zidx[None, :] + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_bf16_2d(output_ptr, Y: tl.constexpr, Z: tl.constexpr): + yoffs = tl.program_id(0) * Y + yidx = tl.arange(0, Y) + yoffs + zidx = tl.arange(0, Z) + ret = tl.full((Y, Z), value=100, dtype=tl.bfloat16) + oidx = yidx[:, None] * Z + zidx[None, :] + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_bool_2d(output_ptr, Y: tl.constexpr, Z: tl.constexpr): + yoffs = tl.program_id(0) * Y + yidx = tl.arange(0, Y) + yoffs + zidx = tl.arange(0, Z) + ret = tl.full((Y, Z), value=0, dtype=tl.int1) + oidx = yidx[:, None] * Z + zidx[None, :] + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int8_1d(output_ptr, Z: tl.constexpr): + zidx = tl.arange(0, Z) + ret = tl.full((Z, ), value=100, dtype=tl.int8) + oidx = zidx + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int16_1d(output_ptr, Z: tl.constexpr): + zidx = tl.arange(0, Z) + ret = tl.full((Z, ), value=100, dtype=tl.int16) + oidx = zidx + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_uint32_1d(output_ptr, Z: tl.constexpr): + zidx = tl.arange(0, Z) + ret = tl.full((Z, ), value=100, dtype=tl.uint32) + oidx = zidx + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int32_1d(output_ptr, Z: tl.constexpr): + zidx = tl.arange(0, Z) + ret = tl.full((Z, ), value=100, dtype=tl.int32) + oidx = zidx + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int64_1d(output_ptr, Z: tl.constexpr): + zidx = tl.arange(0, Z) + ret = tl.full((Z, ), value=100, dtype=tl.int64) + oidx = zidx + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_fp16_1d(output_ptr, Z: tl.constexpr): + zidx = tl.arange(0, Z) + ret = tl.full((Z, ), value=100, dtype=tl.float16) + oidx = zidx + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_fp32_1d(output_ptr, Z: tl.constexpr): + zidx = tl.arange(0, Z) + ret = tl.full((Z, ), value=100, dtype=tl.float32) + oidx = zidx + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_bf16_1d(output_ptr, Z: tl.constexpr): + zidx = tl.arange(0, Z) + ret = tl.full((Z, ), value=100, dtype=tl.bfloat16) + oidx = zidx + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_bool_1d(output_ptr, Z: tl.constexpr): + zidx = tl.arange(0, Z) + ret = tl.full((Z, ), value=0, dtype=tl.int1) + oidx = zidx + tl.store(output_ptr + oidx, ret) + + +test_dtype = ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16', 'bool'] +test_shape1d = TestUtils.test_shape1d +test_shape2d = TestUtils.test_shape2d +test_shape3d = TestUtils.test_shape3d + +# 定义 dtype 到 (test_func, test_sigtype) 的映射 +dtype_mapping3d = { + 'int8': (fn_npu_int8_3d, torch.int8), + 'int16': (fn_npu_int16_3d, torch.int16), + 'int32': (fn_npu_int32_3d, torch.int32), + 'uint32': (fn_npu_uint32_3d, torch.uint32), + 'int64': (fn_npu_int64_3d, torch.int64), + 'float16': (fn_npu_fp16_3d, torch.float16), + 'float32': (fn_npu_fp32_3d, torch.float32), + 'bfloat16': (fn_npu_bf16_3d, torch.bfloat16), + 'bool': (fn_npu_bool_3d, torch.bool), +} +dtype_mapping2d = { + 'int8': (fn_npu_int8_2d, torch.int8), + 'int16': (fn_npu_int16_2d, torch.int16), + 'int32': (fn_npu_int32_2d, torch.int32), + 'uint32': (fn_npu_uint32_2d, torch.uint32), + 'int64': (fn_npu_int64_2d, torch.int64), + 'float16': (fn_npu_fp16_2d, torch.float16), + 'float32': (fn_npu_fp32_2d, torch.float32), + 'bfloat16': (fn_npu_bf16_2d, torch.bfloat16), + 'bool': (fn_npu_bool_2d, torch.bool), +} +dtype_mapping1d = { + 'int8': (fn_npu_int8_1d, torch.int8), + 'int16': (fn_npu_int16_1d, torch.int16), + 'int32': (fn_npu_int32_1d, torch.int32), + 'uint32': (fn_npu_uint32_1d, torch.uint32), + 'int64': (fn_npu_int64_1d, torch.int64), + 'float16': (fn_npu_fp16_1d, torch.float16), + 'float32': (fn_npu_fp32_1d, torch.float32), + 'bfloat16': (fn_npu_bf16_1d, torch.bfloat16), + 'bool': (fn_npu_bool_1d, torch.bool), +} + +# 生成测试用例 +testlist = [(func, sigtype, dtype, shape) + for sigtype in test_dtype + for shape in test_shape1d + for func, dtype in [dtype_mapping1d[sigtype]] # 直接解包映射结果 + ] + +testlist += [(func, sigtype, dtype, shape) + for sigtype in test_dtype + for shape in test_shape2d + for func, dtype in [dtype_mapping2d[sigtype]] # 直接解包映射结果 + ] + +testlist += [(func, sigtype, dtype, shape) + for sigtype in test_dtype + for shape in test_shape3d + for func, dtype in [dtype_mapping3d[sigtype]] # 直接解包映射结果 + ] + + +@pytest.mark.parametrize('testfunc, sigtype, dtype, shape', testlist) +def test_npu(testfunc, sigtype, dtype, shape): + x = 0 + output = 0 + if len(shape) == 3: + if dtype == torch.bool: + x = torch.full((shape[0], shape[1], shape[2]), 0, dtype=dtype).npu() + else: + x = torch.full((shape[0], shape[1], shape[2]), 100, dtype=dtype).npu() + output = torch.randint(1, (shape[0], shape[1], shape[2]), dtype=dtype).npu() + testfunc[(1, 1, 1)](output, shape[0], shape[1], shape[2], debug=True) + if len(shape) == 2: + if dtype == torch.bool: + x = torch.full((shape[0], shape[1]), 0, dtype=dtype).npu() + else: + x = torch.full((shape[0], shape[1]), 100, dtype=dtype).npu() + output = torch.randint(1, (shape[0], shape[1]), dtype=dtype).npu() + shape0 = shape[0] + shape1 = shape[1] + if x.numel() * x.element_size() >= 8192: + grid = (shape0, 1, 1) + shape0 = 1 + else: + grid = (1, 1, 1) + testfunc[grid](output, shape0, shape1, debug=True) + if len(shape) == 1: + if dtype == torch.bool: + x = torch.full((shape[0], ), 0, dtype=dtype).npu() + else: + x = torch.full((shape[0], ), 100, dtype=dtype).npu() + output = torch.randint(1, (shape[0], ), dtype=dtype).npu() + testfunc[1, 1, 1](output, shape[0], debug=True) + test_common.validate_cmp(sigtype, output, x) + + +@triton.jit +def fn_npu_multi_d(output_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr): + dtype = output_ptr.type.element_ty + + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if (YB * ZB * MB * NB) > 1: + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if (ZB * MB * NB) > 1: + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if (MB * NB) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if NB > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + if (YB * ZB * MB * NB) == 1: + ret = tl.full((XB, ), value=100, dtype=dtype) + elif (ZB * MB * NB) == 1: + ret = tl.full((XB, YB), value=100, dtype=dtype) + elif (MB * NB) == 1: + ret = tl.full((XB, YB, ZB), value=100, dtype=dtype) + elif NB == 1: + ret = tl.full((XB, YB, ZB, MB), value=100, dtype=dtype) + else: + ret = tl.full((XB, YB, ZB, MB, NB), value=100, dtype=dtype) + + tl.store(output_ptr + offsets, ret) + + +testlist_multi_d = [ + (fn_npu_multi_d, 'float32', torch.float32, (4, 2, 16, 16)), + (fn_npu_multi_d, 'float32', torch.float32, (2, 4, 2, 16, 16)), + (fn_npu_multi_d, 'float32', torch.float16, (4, 2, 16, 16)), + (fn_npu_multi_d, 'float32', torch.float16, (2, 4, 2, 16, 16)), + (fn_npu_multi_d, 'float32', torch.int8, (4, 2, 16, 16)), + (fn_npu_multi_d, 'float32', torch.int8, (2, 4, 2, 16, 16)), +] + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('testfunc, sigtype, dtype, shape', testlist_multi_d) +def test_npu_4d_5d(testfunc, sigtype, dtype, shape): + x = torch.full(shape, 100, dtype=dtype).npu() + + print(f"shape = {x.shape}") + print(x.dtype) + print(torch.flatten(x)[0:16]) + + output = torch.randint(1, shape, dtype=dtype).npu() + + print(f"output.dtype={output.dtype}") + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + testfunc[(1, )](output, *triton_shape) + print(torch.flatten(output)[0:16]) + + test_common.validate_cmp(sigtype, output, x) + + +@triton.jit +def fn_npu_bf16_6d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr): + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + + ret = tl.full((A, B, C, D, E, F), value=10, dtype=tl.bfloat16) + + oidx = (aidx[:, None, None, None, None, None] * B * C * D * E * F + + bidx[None, :, None, None, None, None] * C * D * E * F + cidx[None, None, :, None, None, None] * D * E * F + + didx[None, None, None, :, None, None] * E * F + eidx[None, None, None, None, :, None] * F + + fidx[None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int8_6d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr): + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + + ret = tl.full((A, B, C, D, E, F), value=10, dtype=tl.int8) + + oidx = (aidx[:, None, None, None, None, None] * B * C * D * E * F + + bidx[None, :, None, None, None, None] * C * D * E * F + cidx[None, None, :, None, None, None] * D * E * F + + didx[None, None, None, :, None, None] * E * F + eidx[None, None, None, None, :, None] * F + + fidx[None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int16_6d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr): + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + + ret = tl.full((A, B, C, D, E, F), value=10, dtype=tl.int16) + + oidx = (aidx[:, None, None, None, None, None] * B * C * D * E * F + + bidx[None, :, None, None, None, None] * C * D * E * F + cidx[None, None, :, None, None, None] * D * E * F + + didx[None, None, None, :, None, None] * E * F + eidx[None, None, None, None, :, None] * F + + fidx[None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int32_6d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr): + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + + ret = tl.full((A, B, C, D, E, F), value=10, dtype=tl.int32) + + oidx = (aidx[:, None, None, None, None, None] * B * C * D * E * F + + bidx[None, :, None, None, None, None] * C * D * E * F + cidx[None, None, :, None, None, None] * D * E * F + + didx[None, None, None, :, None, None] * E * F + eidx[None, None, None, None, :, None] * F + + fidx[None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int64_6d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr): + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + + ret = tl.full((A, B, C, D, E, F), value=10, dtype=tl.int64) + + oidx = (aidx[:, None, None, None, None, None] * B * C * D * E * F + + bidx[None, :, None, None, None, None] * C * D * E * F + cidx[None, None, :, None, None, None] * D * E * F + + didx[None, None, None, :, None, None] * E * F + eidx[None, None, None, None, :, None] * F + + fidx[None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_fp16_6d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr): + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + + ret = tl.full((A, B, C, D, E, F), value=10, dtype=tl.float16) + + oidx = (aidx[:, None, None, None, None, None] * B * C * D * E * F + + bidx[None, :, None, None, None, None] * C * D * E * F + cidx[None, None, :, None, None, None] * D * E * F + + didx[None, None, None, :, None, None] * E * F + eidx[None, None, None, None, :, None] * F + + fidx[None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_fp32_6d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr): + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + + ret = tl.full((A, B, C, D, E, F), value=10, dtype=tl.float32) + + oidx = (aidx[:, None, None, None, None, None] * B * C * D * E * F + + bidx[None, :, None, None, None, None] * C * D * E * F + cidx[None, None, :, None, None, None] * D * E * F + + didx[None, None, None, :, None, None] * E * F + eidx[None, None, None, None, :, None] * F + + fidx[None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_bool_6d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr): + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + + ret = tl.full((A, B, C, D, E, F), value=0, dtype=tl.int1) + + oidx = (aidx[:, None, None, None, None, None] * B * C * D * E * F + + bidx[None, :, None, None, None, None] * C * D * E * F + cidx[None, None, :, None, None, None] * D * E * F + + didx[None, None, None, :, None, None] * E * F + eidx[None, None, None, None, :, None] * F + + fidx[None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +test_dtype = ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16', 'bool'] +test_shape6d = TestUtils.test_shape6d +dtype_mapping6d = { + 'int8': (fn_npu_int8_6d, torch.int8), + 'int16': (fn_npu_int16_6d, torch.int16), + 'int32': (fn_npu_int32_6d, torch.int32), + 'int64': (fn_npu_int64_6d, torch.int64), + 'float16': (fn_npu_fp16_6d, torch.float16), + 'float32': (fn_npu_fp32_6d, torch.float32), + 'bfloat16': (fn_npu_bf16_6d, torch.bfloat16), + 'bool': (fn_npu_bool_6d, torch.bool), +} + +testlist6d = [(func, sigtype, dtype, shape) + for sigtype in test_dtype + for shape in test_shape6d + for func, dtype in [dtype_mapping6d[sigtype]]] + + +@pytest.mark.parametrize('testfunc, sigtype, dtype, shape', testlist6d) +def test_npu_6d(testfunc, sigtype, dtype, shape): + x = 0 + output = 0 + if len(shape) == 6: + if dtype == torch.bool: + x = torch.full((shape[0], shape[1], shape[2], shape[3], shape[4], shape[5]), 0, dtype=dtype).npu() + else: + x = torch.full(shape, 10, dtype=dtype).npu() + output = torch.randint(1, shape, dtype=dtype).npu() + testfunc[1, 1, 1](output, *shape, debug=True) + test_common.validate_cmp(sigtype, output, x) + + +@triton.jit +def fn_npu_bf16_7d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr, G: tl.constexpr): + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + gidx = tl.arange(0, G) + + ret = tl.full((A, B, C, D, E, F, G), value=10, dtype=tl.bfloat16) + + oidx = (aidx[:, None, None, None, None, None, None] * B * C * D * E * F * G + + bidx[None, :, None, None, None, None, None] * C * D * E * F * G + + cidx[None, None, :, None, None, None, None] * D * E * F * G + + didx[None, None, None, :, None, None, None] * E * F * G + + eidx[None, None, None, None, :, None, None] * F * G + fidx[None, None, None, None, None, :, None] * G + + gidx[None, None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int8_7d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr, G: tl.constexpr): + + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + gidx = tl.arange(0, G) + + ret = tl.full((A, B, C, D, E, F, G), value=10, dtype=tl.int8) + + oidx = (aidx[:, None, None, None, None, None, None] * B * C * D * E * F * G + + bidx[None, :, None, None, None, None, None] * C * D * E * F * G + + cidx[None, None, :, None, None, None, None] * D * E * F * G + + didx[None, None, None, :, None, None, None] * E * F * G + + eidx[None, None, None, None, :, None, None] * F * G + fidx[None, None, None, None, None, :, None] * G + + gidx[None, None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int16_7d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr, G: tl.constexpr): + + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + gidx = tl.arange(0, G) + + ret = tl.full((A, B, C, D, E, F, G), value=10, dtype=tl.int16) + + oidx = (aidx[:, None, None, None, None, None, None] * B * C * D * E * F * G + + bidx[None, :, None, None, None, None, None] * C * D * E * F * G + + cidx[None, None, :, None, None, None, None] * D * E * F * G + + didx[None, None, None, :, None, None, None] * E * F * G + + eidx[None, None, None, None, :, None, None] * F * G + fidx[None, None, None, None, None, :, None] * G + + gidx[None, None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int32_7d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr, G: tl.constexpr): + + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + gidx = tl.arange(0, G) + + ret = tl.full((A, B, C, D, E, F, G), value=10, dtype=tl.int32) + + oidx = (aidx[:, None, None, None, None, None, None] * B * C * D * E * F * G + + bidx[None, :, None, None, None, None, None] * C * D * E * F * G + + cidx[None, None, :, None, None, None, None] * D * E * F * G + + didx[None, None, None, :, None, None, None] * E * F * G + + eidx[None, None, None, None, :, None, None] * F * G + fidx[None, None, None, None, None, :, None] * G + + gidx[None, None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int64_7d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr, G: tl.constexpr): + + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + gidx = tl.arange(0, G) + + ret = tl.full((A, B, C, D, E, F, G), value=10, dtype=tl.int64) + + oidx = (aidx[:, None, None, None, None, None, None] * B * C * D * E * F * G + + bidx[None, :, None, None, None, None, None] * C * D * E * F * G + + cidx[None, None, :, None, None, None, None] * D * E * F * G + + didx[None, None, None, :, None, None, None] * E * F * G + + eidx[None, None, None, None, :, None, None] * F * G + fidx[None, None, None, None, None, :, None] * G + + gidx[None, None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_fp16_7d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr, G: tl.constexpr): + + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + gidx = tl.arange(0, G) + + ret = tl.full((A, B, C, D, E, F, G), value=10, dtype=tl.float16) + + oidx = (aidx[:, None, None, None, None, None, None] * B * C * D * E * F * G + + bidx[None, :, None, None, None, None, None] * C * D * E * F * G + + cidx[None, None, :, None, None, None, None] * D * E * F * G + + didx[None, None, None, :, None, None, None] * E * F * G + + eidx[None, None, None, None, :, None, None] * F * G + fidx[None, None, None, None, None, :, None] * G + + gidx[None, None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_fp32_7d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr, G: tl.constexpr): + + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + gidx = tl.arange(0, G) + + ret = tl.full((A, B, C, D, E, F, G), value=10, dtype=tl.float32) + + oidx = (aidx[:, None, None, None, None, None, None] * B * C * D * E * F * G + + bidx[None, :, None, None, None, None, None] * C * D * E * F * G + + cidx[None, None, :, None, None, None, None] * D * E * F * G + + didx[None, None, None, :, None, None, None] * E * F * G + + eidx[None, None, None, None, :, None, None] * F * G + fidx[None, None, None, None, None, :, None] * G + + gidx[None, None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_bool_7d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr, G: tl.constexpr): + + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + gidx = tl.arange(0, G) + + ret = tl.full((A, B, C, D, E, F, G), value=0, dtype=tl.int1) + + oidx = (aidx[:, None, None, None, None, None, None] * B * C * D * E * F * G + + bidx[None, :, None, None, None, None, None] * C * D * E * F * G + + cidx[None, None, :, None, None, None, None] * D * E * F * G + + didx[None, None, None, :, None, None, None] * E * F * G + + eidx[None, None, None, None, :, None, None] * F * G + fidx[None, None, None, None, None, :, None] * G + + gidx[None, None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +test_dtype = ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16', 'bool'] +test_shape7d = TestUtils.test_shape7d +dtype_mapping7d = { + 'int8': (fn_npu_int8_7d, torch.int8), + 'int16': (fn_npu_int16_7d, torch.int16), + 'int32': (fn_npu_int32_7d, torch.int32), + 'int64': (fn_npu_int64_7d, torch.int64), + 'float16': (fn_npu_fp16_7d, torch.float16), + 'float32': (fn_npu_fp32_7d, torch.float32), + 'bfloat16': (fn_npu_bf16_7d, torch.bfloat16), + 'bool': (fn_npu_bool_7d, torch.bool), +} + +testlist7d = [(func, sigtype, dtype, shape) + for sigtype in test_dtype + for shape in test_shape7d + for func, dtype in [dtype_mapping7d[sigtype]]] + + +@pytest.mark.parametrize('testfunc, sigtype, dtype, shape', testlist7d) +def test_npu_7d(testfunc, sigtype, dtype, shape): + x = 0 + output = 0 + if len(shape) == 7: + if dtype == torch.bool: + x = torch.full((shape[0], shape[1], shape[2], shape[3], shape[4], shape[5], shape[6]), 0, dtype=dtype).npu() + else: + x = torch.full(shape, 10, dtype=dtype).npu() + output = torch.randint(1, shape, dtype=dtype).npu() + testfunc[1, 1, 1](output, *shape, debug=True) + test_common.validate_cmp(sigtype, output, x) + + +@triton.jit +def fn_npu_bf16_8d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr, G: tl.constexpr, H: tl.constexpr): + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + gidx = tl.arange(0, G) + hidx = tl.arange(0, H) + + ret = tl.full((A, B, C, D, E, F, G, H), value=10, dtype=tl.bfloat16) + + oidx = (aidx[:, None, None, None, None, None, None, None] * B * C * D * E * F * G * H + + bidx[None, :, None, None, None, None, None, None] * C * D * E * F * G * H + + cidx[None, None, :, None, None, None, None, None] * D * E * F * G * H + + didx[None, None, None, :, None, None, None, None] * E * F * G * H + + eidx[None, None, None, None, :, None, None, None] * F * G * H + + fidx[None, None, None, None, None, :, None, None] * G * H + + gidx[None, None, None, None, None, None, :, None] * H + hidx[None, None, None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int8_8d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr, G: tl.constexpr, H: tl.constexpr): + + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + gidx = tl.arange(0, G) + hidx = tl.arange(0, H) + + ret = tl.full((A, B, C, D, E, F, G, H), value=10, dtype=tl.int8) + + oidx = (aidx[:, None, None, None, None, None, None, None] * B * C * D * E * F * G * H + + bidx[None, :, None, None, None, None, None, None] * C * D * E * F * G * H + + cidx[None, None, :, None, None, None, None, None] * D * E * F * G * H + + didx[None, None, None, :, None, None, None, None] * E * F * G * H + + eidx[None, None, None, None, :, None, None, None] * F * G * H + + fidx[None, None, None, None, None, :, None, None] * G * H + + gidx[None, None, None, None, None, None, :, None] * H + hidx[None, None, None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int16_8d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr, G: tl.constexpr, H: tl.constexpr): + + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + gidx = tl.arange(0, G) + hidx = tl.arange(0, H) + + ret = tl.full((A, B, C, D, E, F, G, H), value=10, dtype=tl.int16) + + oidx = (aidx[:, None, None, None, None, None, None, None] * B * C * D * E * F * G * H + + bidx[None, :, None, None, None, None, None, None] * C * D * E * F * G * H + + cidx[None, None, :, None, None, None, None, None] * D * E * F * G * H + + didx[None, None, None, :, None, None, None, None] * E * F * G * H + + eidx[None, None, None, None, :, None, None, None] * F * G * H + + fidx[None, None, None, None, None, :, None, None] * G * H + + gidx[None, None, None, None, None, None, :, None] * H + hidx[None, None, None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int32_8d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr, G: tl.constexpr, H: tl.constexpr): + + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + gidx = tl.arange(0, G) + hidx = tl.arange(0, H) + + ret = tl.full((A, B, C, D, E, F, G, H), value=10, dtype=tl.int32) + + oidx = (aidx[:, None, None, None, None, None, None, None] * B * C * D * E * F * G * H + + bidx[None, :, None, None, None, None, None, None] * C * D * E * F * G * H + + cidx[None, None, :, None, None, None, None, None] * D * E * F * G * H + + didx[None, None, None, :, None, None, None, None] * E * F * G * H + + eidx[None, None, None, None, :, None, None, None] * F * G * H + + fidx[None, None, None, None, None, :, None, None] * G * H + + gidx[None, None, None, None, None, None, :, None] * H + hidx[None, None, None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_int64_8d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr, G: tl.constexpr, H: tl.constexpr): + + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + gidx = tl.arange(0, G) + hidx = tl.arange(0, H) + + ret = tl.full((A, B, C, D, E, F, G, H), value=10, dtype=tl.int64) + + oidx = (aidx[:, None, None, None, None, None, None, None] * B * C * D * E * F * G * H + + bidx[None, :, None, None, None, None, None, None] * C * D * E * F * G * H + + cidx[None, None, :, None, None, None, None, None] * D * E * F * G * H + + didx[None, None, None, :, None, None, None, None] * E * F * G * H + + eidx[None, None, None, None, :, None, None, None] * F * G * H + + fidx[None, None, None, None, None, :, None, None] * G * H + + gidx[None, None, None, None, None, None, :, None] * H + hidx[None, None, None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_fp16_8d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr, G: tl.constexpr, H: tl.constexpr): + + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + gidx = tl.arange(0, G) + hidx = tl.arange(0, H) + + ret = tl.full((A, B, C, D, E, F, G, H), value=10, dtype=tl.float16) + + oidx = (aidx[:, None, None, None, None, None, None, None] * B * C * D * E * F * G * H + + bidx[None, :, None, None, None, None, None, None] * C * D * E * F * G * H + + cidx[None, None, :, None, None, None, None, None] * D * E * F * G * H + + didx[None, None, None, :, None, None, None, None] * E * F * G * H + + eidx[None, None, None, None, :, None, None, None] * F * G * H + + fidx[None, None, None, None, None, :, None, None] * G * H + + gidx[None, None, None, None, None, None, :, None] * H + hidx[None, None, None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_fp32_8d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr, G: tl.constexpr, H: tl.constexpr): + + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + gidx = tl.arange(0, G) + hidx = tl.arange(0, H) + + ret = tl.full((A, B, C, D, E, F, G, H), value=10, dtype=tl.float32) + + oidx = (aidx[:, None, None, None, None, None, None, None] * B * C * D * E * F * G * H + + bidx[None, :, None, None, None, None, None, None] * C * D * E * F * G * H + + cidx[None, None, :, None, None, None, None, None] * D * E * F * G * H + + didx[None, None, None, :, None, None, None, None] * E * F * G * H + + eidx[None, None, None, None, :, None, None, None] * F * G * H + + fidx[None, None, None, None, None, :, None, None] * G * H + + gidx[None, None, None, None, None, None, :, None] * H + hidx[None, None, None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_bool_8d(output_ptr, A: tl.constexpr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr, E: tl.constexpr, + F: tl.constexpr, G: tl.constexpr, H: tl.constexpr): + + aidx = tl.arange(0, A) + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + eidx = tl.arange(0, E) + fidx = tl.arange(0, F) + gidx = tl.arange(0, G) + hidx = tl.arange(0, H) + + ret = tl.full((A, B, C, D, E, F, G, H), value=0, dtype=tl.int1) + + oidx = (aidx[:, None, None, None, None, None, None, None] * B * C * D * E * F * G * H + + bidx[None, :, None, None, None, None, None, None] * C * D * E * F * G * H + + cidx[None, None, :, None, None, None, None, None] * D * E * F * G * H + + didx[None, None, None, :, None, None, None, None] * E * F * G * H + + eidx[None, None, None, None, :, None, None, None] * F * G * H + + fidx[None, None, None, None, None, :, None, None] * G * H + + gidx[None, None, None, None, None, None, :, None] * H + hidx[None, None, None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +test_dtype = ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16', 'bool'] +test_shape8d = TestUtils.test_shape8d +dtype_mapping8d = { + 'int8': (fn_npu_int8_8d, torch.int8), + 'int16': (fn_npu_int16_8d, torch.int16), + 'int32': (fn_npu_int32_8d, torch.int32), + 'int64': (fn_npu_int64_8d, torch.int64), + 'float16': (fn_npu_fp16_8d, torch.float16), + 'float32': (fn_npu_fp32_8d, torch.float32), + 'bfloat16': (fn_npu_bf16_8d, torch.bfloat16), + 'bool': (fn_npu_bool_8d, torch.bool), +} + +testlist8d = [(func, sigtype, dtype, shape) + for sigtype in test_dtype + for shape in test_shape8d + for func, dtype in [dtype_mapping8d[sigtype]]] + + +@pytest.mark.parametrize('testfunc, sigtype, dtype, shape', testlist8d) +def test_npu_8d(testfunc, sigtype, dtype, shape): + x = 0 + output = 0 + if len(shape) == 8: + if dtype == torch.bool: + x = torch.full((shape[0], shape[1], shape[2], shape[3], shape[4], shape[5], shape[6], shape[7]), 0, + dtype=dtype).npu() + else: + x = torch.full(shape, 10, dtype=dtype).npu() + output = torch.randint(1, shape, dtype=dtype).npu() + testfunc[1, 1, 1](output, *shape, debug=True) + test_common.validate_cmp(sigtype, output, x) diff --git a/third_party/ascend/unittest/generalization_cases/test_ge_op.py b/third_party/ascend/unittest/generalization_cases/test_ge_op.py new file mode 100644 index 0000000000..d23da78f38 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_ge_op.py @@ -0,0 +1,158 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import time +import torch +import torch_npu +import test_common +from test_common import TestUtils +import logging + + +@triton.jit +def triton_ge_3d(in_ptr0, in_ptr1, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0 >= x1 + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_ge_2d(in_ptr0, in_ptr1, out_ptr0, M: tl.constexpr, N: tl.constexpr): + moffs = tl.program_id(0) * M + mblk_idx = tl.arange(0, M) + moffs + nblk_idx = tl.arange(0, N) + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0 >= x1 + odx = mblk_idx[:, None] * N + nblk_idx[None, :] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_ge_1d(in_ptr0, in_ptr1, out_ptr0, L: tl.constexpr): + lblk_idx = tl.arange(0, L) + idx = lblk_idx[:] + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0 >= x1 + odx = lblk_idx[:] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_ge_4d_5d(x_ptr, y_ptr, output_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = x_val >= y_val + tl.store(output_ptr + offsets, ret, mask=masks) + + +typelist = ['bool', 'int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32'] + +dtype_mapping = { + 'int8': (torch.int8), + 'int16': (torch.int16), + 'int32': (torch.int32), + 'uint32': (torch.uint32), + 'int64': (torch.int64), + 'float16': (torch.float16), + 'float32': (torch.float32), + 'bfloat16': (torch.bfloat16), + 'bool': (torch.bool), +} + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1_2_3d) +@pytest.mark.parametrize('sigtype', typelist) +def test_ge(sigtype, shape): + dtype = dtype_mapping[sigtype] + x0 = test_common.generate_tensor(shape=shape, dtype=sigtype).npu() + x1 = test_common.generate_tensor(shape=shape, dtype=sigtype).npu() + # ncore, xblock, xblock_sub = 2, 32768, 1024 + y_ref = torch.where(torch.ge(x0, x1), torch.ones_like(x0), torch.zeros_like(x0)).to(eval('torch.' + sigtype)) + output = torch.zeros(shape, dtype=dtype).npu() + if len(shape) == 3: + triton_ge_3d[1, 1, 1](x0, x1, output, shape[0], shape[1], shape[2]) + if len(shape) == 2: + shape0 = shape[0] + shape1 = shape[1] + if x0.numel() * x0.element_size() >= 8192: + grid = (shape0, 1, 1) + shape0 = 1 + else: + grid = (1, 1, 1) + triton_ge_2d[grid](x0, x1, output, shape0, shape1) + if len(shape) == 1: + triton_ge_1d[1, 1, 1](x0, x1, output, shape[0]) + test_common.validate_cmp(sigtype, output, y_ref) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32']) +def test_ge_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.where(torch.ge(x, y), torch.ones_like(x), torch.zeros_like(x)).to(eval('torch.' + dtype)) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_ge_4d_5d[grid](x, y, output, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_add.py b/third_party/ascend/unittest/generalization_cases/test_general_add.py new file mode 100644 index 0000000000..8ddb6adeb0 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_add.py @@ -0,0 +1,438 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import torch +import test_common +from test_common import TestUtils +import logging +import numpy as np + + +@triton.jit +def triton_add(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, + XNUMEL: tl.constexpr, YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = X + Y + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_add_broadcast(in_ptr0, in_ptr1, out_ptr0, X_SHAPE_0: tl.constexpr, X_SHAPE_1: tl.constexpr, + X_SHAPE_2: tl.constexpr, X_SHAPE_3: tl.constexpr, X_SHAPE_4: tl.constexpr, + Y_SHAPE_0: tl.constexpr, Y_SHAPE_1: tl.constexpr, Y_SHAPE_2: tl.constexpr, + Y_SHAPE_3: tl.constexpr, Y_SHAPE_4: tl.constexpr): + x_idx0 = tl.arange(0, X_SHAPE_0) + x_idx1 = tl.arange(0, X_SHAPE_1) + x_idx2 = tl.arange(0, X_SHAPE_2) + x_idx3 = tl.arange(0, X_SHAPE_3) + x_idx4 = tl.arange(0, X_SHAPE_4) + + y_idx0 = tl.arange(0, Y_SHAPE_0) + y_idx1 = tl.arange(0, Y_SHAPE_1) + y_idx2 = tl.arange(0, Y_SHAPE_2) + y_idx3 = tl.arange(0, Y_SHAPE_3) + y_idx4 = tl.arange(0, Y_SHAPE_4) + + xidx = x_idx0[:, None, None, None, None] * X_SHAPE_1 * X_SHAPE_2 * X_SHAPE_3 * X_SHAPE_4 + \ + x_idx1[None, :, None, None, None] * X_SHAPE_2 * X_SHAPE_3 * X_SHAPE_4 + \ + x_idx2[None, None, :, None, None] * X_SHAPE_3 * X_SHAPE_4 + \ + x_idx3[None, None, None, :, None] * X_SHAPE_4 + x_idx4[None, None, None, None, :] + + yidx = y_idx0[:, None, None, None, None] * Y_SHAPE_1 * Y_SHAPE_2 * Y_SHAPE_3 * Y_SHAPE_4 + \ + y_idx1[None, :, None, None, None] * Y_SHAPE_2 * Y_SHAPE_3 * Y_SHAPE_4 + \ + y_idx2[None, None, :, None, None] * Y_SHAPE_3 * Y_SHAPE_4 + \ + y_idx3[None, None, None, :, None] * Y_SHAPE_4 + y_idx4[None, None, None, None, :] + + X = tl.load(in_ptr0 + xidx) + Y = tl.load(in_ptr1 + yidx) + ret = X + Y + + tl.store(out_ptr0 + xidx, ret) + + +@triton.jit +def triton_add_4d_5d(output_ptr, x_ptr, y_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = x_val + y_val + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['bool', 'int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +def test_add(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + + ans = x + y + output = torch.zeros_like(ans) + + if len(shape) == 1: + triton_add[1, 1, shape[0]](output, x, y, z, 1, 1, 1, 1, 1, shape[0]) + elif len(shape) == 2: + if shape[0] > shape[1]: + triton_add[1, shape[0], 1](output, x, y, z, 1, 1, shape[1], 1, shape[0], shape[1]) + else: + triton_add[1, 1, shape[1]](output, x, y, z, 1, shape[0], 1, 1, shape[0], shape[1]) + elif len(shape) == 3: + if max(shape[0], shape[1], shape[2]) == shape[0]: + triton_add[shape[0], 1, 1](output, x, y, z, 1, shape[1], shape[2], shape[0], shape[1], shape[2]) + elif max(shape[0], shape[1], shape[2]) == shape[1]: + triton_add[1, shape[1], 1](output, x, y, z, shape[0], 1, shape[2], shape[0], shape[1], shape[2]) + else: + triton_add[1, 1, shape[2]](output, x, y, z, shape[0], shape[1], 1, shape[0], shape[1], shape[2]) + else: + triton_add[1, 1, 1](output, x, y, z, 1, 1, 1, 1, 1, 1) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +def test_add_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = x + y + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_add_4d_5d[grid](output, x, y, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) + + +def promote_dtype(x_dtype, y_dtype): + """ + 如果 y 的精度低于 x, 则提升 y 的精度以匹配 x。 + """ + # 如果两个数据类型一致，直接返回 + if x_dtype == y_dtype: + return y_dtype + + # 构建类型的优先级列表（从低到高） + priority = [ + torch.bool, torch.int8, torch.int16, torch.int32, torch.int64, torch.float16, torch.bfloat16, torch.float32 + ] + + # 查找两种类型在优先级列表中的位置 + x_priority = priority.index(x_dtype) + y_priority = priority.index(y_dtype) + + # 如果y的优先级比x小，则提升到x的类型 + if y_priority < x_priority: + return x_dtype + else: + return y_dtype + + +@pytest.mark.parametrize('param_list', + [[(5, 1, 1, 1, 1), + (5, 1, 1, 2, 1)], [(2, 1), (2, 4)], [(2, 1, 1), (2, 4, 2)], [(2, 1, 1, 1), (2, 4, 2, 2)], + [(2, 1, 1, 1, 1), + (2, 4, 2, 2, 2)], [(1, ), (4, )], [(1, 2, 1), (1, 2, 3)], [(1, 1, 1, 1), (7, 1, 1, 1)]]) +@pytest.mark.parametrize('x_dtype_str', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16', 'bool']) +@pytest.mark.parametrize('y_dtype_str', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16', 'bool']) +def test_add_broadcast(param_list, x_dtype_str, y_dtype_str): + x_shape, y_shape = param_list + x_dtype = eval('torch.' + x_dtype_str) + y_dtype = eval('torch.' + y_dtype_str) + + x = test_common.generate_tensor(x_shape, x_dtype_str).npu() + y = test_common.generate_tensor(y_shape, y_dtype_str).npu() + if y.numel() > x.numel(): + tmp = y + y = x + x = tmp + ans = x + y + while x.dim() < 5: + x = x.unsqueeze(-1) + while y.dim() < 5: + y = y.unsqueeze(-1) + bf2fpFlag = False + out_dtype = promote_dtype(x_dtype, y_dtype) + if (x_dtype == torch.bfloat16 and y_dtype == torch.float16) or \ + (x_dtype == torch.float16 and y_dtype == torch.bfloat16): + out_dtype = torch.float32 + bf2fpFlag = True + out_dtype = str(out_dtype).split('.')[-1] + out = test_common.generate_tensor(x.shape, out_dtype).npu() + + triton_add_broadcast[1, 1, 1](x, y, out, *x.shape, *y.shape) + while out.dim() > ans.dim(): + out = out.squeeze(-1) + + if bf2fpFlag: + torch.testing.assert_close(out, ans, rtol=1e-3, atol=1e-3) + else: + torch.testing.assert_close(out, ans) + + +@triton.jit +def add_5d(x_ptr, y_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr, XB1: tl.constexpr, YB1: tl.constexpr, ZB1: tl.constexpr, MB1: tl.constexpr, + NB1: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + offsets1 = tl.arange(0, XB1) * (YB1 * ZB1 * MB1 * NB1) + offsets1 = offsets1[:, None] + tl.arange(0, YB1)[None, :] * (ZB1 * MB1 * NB1) + offsets1 = offsets1[:, :, None] + tl.arange(0, ZB1)[None, None, :] * (MB1 * NB1) + offsets1 = offsets1[:, :, :, None] + tl.arange(0, MB1)[None, None, None, :] * NB1 + offsets1 = offsets1[:, :, :, :, None] + tl.arange(0, NB1)[None, None, None, None, :] + + tmp0 = tl.load(x_ptr + offsets) + tmp1 = tl.load(y_ptr + offsets1) + tmp2 = tl.load(out_ptr + offsets1) + out = tmp2 + tmp1 + tmp0 + tl.store(out_ptr + offsets1, out) + + +@triton.jit +def add_4d(x_ptr, y_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + XB1: tl.constexpr, YB1: tl.constexpr, ZB1: tl.constexpr, MB1: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB) + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB) + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB) + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] + + offsets1 = tl.arange(0, XB1) * (YB1 * ZB1 * MB1) + offsets1 = offsets1[:, None] + tl.arange(0, YB1)[None, :] * (ZB1 * MB1) + offsets1 = offsets1[:, :, None] + tl.arange(0, ZB1)[None, None, :] * (MB1) + offsets1 = offsets1[:, :, :, None] + tl.arange(0, MB1)[None, None, None, :] + + tmp0 = tl.load(x_ptr + offsets) + tmp1 = tl.load(y_ptr + offsets1) + tmp2 = tl.load(out_ptr + offsets1) + out = tmp2 + tmp1 + tmp0 + tl.store(out_ptr + offsets1, out) + + +@triton.jit +def add_3d(x_ptr, y_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XB1: tl.constexpr, + YB1: tl.constexpr, ZB1: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB) + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB) + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] + + offsets1 = tl.arange(0, XB1) * (YB1 * ZB1) + offsets1 = offsets1[:, None] + tl.arange(0, YB1)[None, :] * (ZB1) + offsets1 = offsets1[:, :, None] + tl.arange(0, ZB1)[None, None, :] + + tmp0 = tl.load(x_ptr + offsets) + tmp1 = tl.load(y_ptr + offsets1) + tmp2 = tl.load(out_ptr + offsets1) + out = tmp2 + tmp1 + tmp0 + tl.store(out_ptr + offsets1, out) + + +@triton.jit +def add_2d(x_ptr, y_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, XB1: tl.constexpr, YB1: tl.constexpr): + offsets = tl.arange(0, XB) * (YB) + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] + + offsets1 = tl.arange(0, XB1) * (YB1) + offsets1 = offsets1[:, None] + tl.arange(0, YB1)[None, :] + + tmp0 = tl.load(x_ptr + offsets) + tmp1 = tl.load(y_ptr + offsets1) + tmp2 = tl.load(out_ptr + offsets1) + out = tmp2 + tmp1 + tmp0 + tl.store(out_ptr + offsets1, out) + + +@pytest.mark.parametrize('param_list', [ + [(5, 1, 1, 1, 1), (5, 1, 1, 2, 1)], +]) +@pytest.mark.parametrize('x_dtype_str', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16', 'bool']) +@pytest.mark.parametrize('y_dtype_str', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16', 'bool']) +def test_add_2d_to_5d(x_dtype_str, y_dtype_str, param_list): + x0_shape, y_shape = param_list + ndim = max(len(x0_shape), len(y_shape)) + # 获取原始类型 + x_dtype = eval('torch.' + x_dtype_str) + y_dtype = eval('torch.' + y_dtype_str) + + x0 = test_common.generate_tensor(x0_shape, x_dtype_str).npu() + y = test_common.generate_tensor(y_shape, y_dtype_str).npu() + + out_dtype = promote_dtype(x_dtype, y_dtype) + if out_dtype == torch.bfloat16: + out_dtype = torch.float32 + out = torch.full(y_shape, 0, dtype=out_dtype).npu() + + x0_temp = x0.clone() + y_temp = y.clone() + out_temp = out.clone() + + triton_shape = [*x0_shape] + while len(triton_shape) < ndim: + triton_shape.append(1) + + triton_shape1 = [*y_shape] + while len(triton_shape1) < ndim: + triton_shape1.append(1) + + # 按维度分支 + if ndim == 2: + XB, YB = triton_shape + XB1, YB1 = triton_shape1 + + add_2d[(1, )]( + x_ptr=x0, + y_ptr=y, + out_ptr=out, + XB=XB, + YB=YB, + XB1=XB1, + YB1=YB1, + ) + + elif ndim == 3: + XB, YB, ZB = triton_shape + XB1, YB1, ZB1 = triton_shape1 + + add_3d[(1, )]( + x_ptr=x0, + y_ptr=y, + out_ptr=out, + XB=XB, + YB=YB, + ZB=ZB, + XB1=XB1, + YB1=YB1, + ZB1=ZB1, + ) + + elif ndim == 4: + XB, YB, ZB, MB = triton_shape + XB1, YB1, ZB1, MB1 = triton_shape1 + + add_4d[(1, )]( + x_ptr=x0, + y_ptr=y, + out_ptr=out, + XB=XB, + YB=YB, + ZB=ZB, + MB=MB, + XB1=XB1, + YB1=YB1, + ZB1=ZB1, + MB1=MB1, + ) + + elif ndim == 5: + XB, YB, ZB, MB, NB = triton_shape + XB1, YB1, ZB1, MB1, NB1 = triton_shape1 + + add_5d[(1, )]( + x_ptr=x0, + y_ptr=y, + out_ptr=out, + XB=XB, + YB=YB, + ZB=ZB, + MB=MB, + NB=NB, + XB1=XB1, + YB1=YB1, + ZB1=ZB1, + MB1=MB1, + NB1=NB1, + ) + + else: + raise ValueError(f"Unsupported tensor dim: {ndim}") + expected = out_temp + y_temp + x0_temp + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', ['uint16', 'uint32', 'uint64']) +def test_add_uint(shape, dtype): + torch_dtype = eval('torch.' + dtype) + np_x0 = test_common.generate_numpy(shape, dtype) + np_x1 = test_common.generate_numpy(shape, dtype) + np_x2 = test_common.generate_numpy(shape, dtype) + + x0 = torch.from_numpy(np_x0).to(torch_dtype).npu() + x1 = torch.from_numpy(np_x1).to(torch_dtype).npu() + x2 = torch.from_numpy(np_x2).to(torch_dtype).npu() + + #numpy result + ans_numpy = np_x0 + np_x1 + z_ref1 = torch.from_numpy(ans_numpy).npu() + + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_add[1, 1, shape[0]](triton_res, x0, x1, x2, 1, 1, 1, 1, 1, shape[0]) + test_common.validate_cmp(dtype, z_ref1, triton_res) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_arange.py b/third_party/ascend/unittest/generalization_cases/test_general_arange.py new file mode 100644 index 0000000000..4e7f2df67b --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_arange.py @@ -0,0 +1,66 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import pytest +import torch +import torch_npu +import triton +import triton.language as tl +import test_common +from test_common import TestUtils + + +def torch_pointwise(length): + res = (torch.arange(0, length) / 2.7) * torch.arange(0, length) + return res + + +def torch_arange(start, end): + TRITON_MAX_TENSOR_NUMEL = 1048576 + if end < start: + raise ValueError("arange's end argument must be greater than the start argument") + if end - start > TRITON_MAX_TENSOR_NUMEL: + raise ValueError( + f"end - start must be less than or equal to TRITON_MAX_TENSOR_NUMEL = {TRITON_MAX_TENSOR_NUMEL}") + return torch.arange(start, end) + + +@triton.jit +def triton_arange(z, BLOCK: tl.constexpr, START: tl.constexpr, END: tl.constexpr): + off = tl.arange(0, BLOCK) + val = tl.arange(START, END) + tl.store(z + off, val) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +def test_case(shape): + start = 0 + end = shape[0] + shape = [end - start] + block = end - start + dtype = 'int32' + + y_ref = torch_arange(start, end) + y_cal = torch.zeros(shape, dtype=torch.int32).npu() + + triton_arange[(1, )](y_cal, START=start, END=end, BLOCK=block) + + assert torch.equal(y_cal.cpu(), y_ref.cpu()) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_cat.py b/third_party/ascend/unittest/generalization_cases/test_general_cat.py new file mode 100644 index 0000000000..a24232fccb --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_cat.py @@ -0,0 +1,61 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common +from test_common import TestUtils +import math + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, XB: tl.constexpr): + + idx = tl.arange(0, XB) + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.cat(X, Y, can_reorder=True) + + oidx = tl.arange(0, XB * 2) + + tl.store(output_ptr + oidx, ret) + + +# The CAT operator in the Triton community also does not support boolean types. +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) #triton only support 1D cat +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16', 'int32', 'int16', 'int8', 'int64']) +def test_cat(shape, dtype): + m = shape[0] + x = torch.full((m, ), 100, dtype=eval("torch." + dtype)).npu() + y = torch.full((m, ), 30, dtype=eval("torch." + dtype)).npu() + + output = torch.randint(1, (m * 2, ), dtype=eval("torch." + dtype)).npu() + + ans = torch.cat((x, y), dim=0) + + fn_npu_[1, 1, 1](output, x, y, m) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_clamp.py b/third_party/ascend/unittest/generalization_cases/test_general_clamp.py new file mode 100644 index 0000000000..d6d91c20e7 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_clamp.py @@ -0,0 +1,179 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time +import torch +import torch_npu +import test_common +from test_common import TestUtils +import logging + + +def torch_clamp(x0, min_, max_): + res = torch.clamp(x0, min_, max_) + return res + + +@triton.jit +def tt_clamp_1d(in_ptr, out_ptr, min_ptr, max_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + idx = tl.arange(0, XB) + + x = tl.load(in_ptr + idx) + min_ = tl.load(min_ptr + idx) + max_ = tl.load(max_ptr + idx) + ret = tl.clamp(x, min_, max_) + + tl.store(out_ptr + idx, ret) + + +@triton.jit +def tt_clamp_2d(in_ptr, out_ptr, min_ptr, max_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + idx = xidx[:, None] * ynumel + yidx[None, :] + + x = tl.load(in_ptr + idx) + min_ = tl.load(min_ptr + idx) + max_ = tl.load(max_ptr + idx) + ret = tl.clamp(x, min_, max_) + + tl.store(out_ptr + idx, ret) + + +@triton.jit +def tt_clamp_3d(in_ptr, out_ptr, min_ptr, max_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + + x = tl.load(in_ptr + idx) + min_ = tl.load(min_ptr + idx) + max_ = tl.load(max_ptr + idx) + ret = tl.clamp(x, min_, max_) + + tl.store(out_ptr + idx, ret) + + +@triton.jit +def triton_clamp_4d_5d(x_ptr, output_ptr, min_ptr, max_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, + BLOCK_2: tl.constexpr, BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, + STRIDE_0: tl.constexpr, STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, + STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + min_ = tl.load(min_ptr + offsets) + max_ = tl.load(max_ptr + offsets) + ret = tl.clamp(x_val, min_, max_) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_clamp(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + torch.manual_seed(0) + x = test_common.generate_tensor(shape, dtype).npu() + a = test_common.generate_tensor(shape, dtype) + b = test_common.generate_tensor(shape, dtype) + min_ = torch.min(a, b).npu() + max_ = torch.max(a, b).npu() + + grid = (1, 1, 1) + + y_cal = torch.empty(shape, dtype=eval('torch.' + dtype), device="npu") + + y_ref = torch_clamp(x, min_, max_) + if len(shape) == 1: + tt_clamp_1d[grid](x, y_cal, min_, max_, x.numel(), 1, 1, x.numel(), 1, 1) + elif len(shape) == 2: + xnumel, ynumel, znumel = shape + (1, ) + XB, YB, ZB = xnumel, ynumel, znumel + if x.numel() * x.element_size() > 8192: + grid = (1, ynumel, 1) + YB = 1 + tt_clamp_2d[grid](x, y_cal, min_, max_, xnumel, ynumel, znumel, XB, YB, ZB) + + elif len(shape) == 3: + xnumel, ynumel, znumel = shape + XB, YB, ZB = xnumel, ynumel, znumel + tt_clamp_3d[grid](x, y_cal, min_, max_, xnumel, ynumel, znumel, XB, YB, ZB) + + test_common.validate_cmp(dtype, y_cal, y_ref) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_clamp_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + torch.manual_seed(0) + x = test_common.generate_tensor(shape, dtype).npu() + a = test_common.generate_tensor(shape, dtype) + b = test_common.generate_tensor(shape, dtype) + min_ = torch.min(a, b).npu() + max_ = torch.max(a, b).npu() + + output = torch.empty(shape, dtype=eval('torch.' + dtype), device="npu") + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_clamp(x, min_, max_) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_clamp_4d_5d[grid](x, output, min_, max_, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_div.py b/third_party/ascend/unittest/generalization_cases/test_general_div.py new file mode 100644 index 0000000000..fd4e252177 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_div.py @@ -0,0 +1,139 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import torch +import test_common +from test_common import TestUtils +import logging + + +@triton.jit +def triton_div(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, + XNUMEL: tl.constexpr, YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = X / Y + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_div_4d_5d(output_ptr, x_ptr, y_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = x_val / y_val + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) # some shape with int8 over ub +@pytest.mark.parametrize('dtype', ['bool', 'int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +def test_div(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + y[y == 0] = 1 + + ans = x / y + output = torch.zeros_like(ans) + if len(shape) == 1: + triton_div[1, 1, shape[0]](output, x, y, z, 1, 1, 1, 1, 1, shape[0]) + elif len(shape) == 2: + if shape[0] > shape[1]: + triton_div[1, shape[0], 1](output, x, y, z, 1, 1, shape[1], 1, shape[0], shape[1]) + else: + triton_div[1, 1, shape[1]](output, x, y, z, 1, shape[0], 1, 1, shape[0], shape[1]) + elif len(shape) == 3: + if max(shape[0], shape[1], shape[2]) == shape[0]: + triton_div[shape[0], 1, 1](output, x, y, z, 1, shape[1], shape[2], shape[0], shape[1], shape[2]) + elif max(shape[0], shape[1], shape[2]) == shape[1]: + triton_div[1, shape[1], 1](output, x, y, z, shape[0], 1, shape[2], shape[0], shape[1], shape[2]) + else: + triton_div[1, 1, shape[2]](output, x, y, z, shape[0], shape[1], 1, shape[0], shape[1], shape[2]) + else: + triton_div[1, 1, 1](output, x, y, z, 1, 1, 1, 1, 1, 1) + # change dtype beacuse of triton processing, triton div op will change from int to float + if dtype in ['int8', 'int16', 'int32', 'int64']: + dtype = 'float32' + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +def test_div_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + y[y == 0] = 1 + + new_shape = shape + if dtype == 'int8' or dtype == 'int16' or dtype == 'int32' or dtype == 'int64': + output = torch.randint(1, new_shape, dtype=eval('torch.float32')).npu() + dtype = 'float32' + else: + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + + ans = x / y + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_div_4d_5d[grid](output, x, y, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_floor.py b/third_party/ascend/unittest/generalization_cases/test_general_floor.py new file mode 100644 index 0000000000..38bc1621ac --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_floor.py @@ -0,0 +1,148 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import logging +import pytest +import test_common +from test_common import TestUtils +import math + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = X + tl.floor(Y) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_floor_4d_5d(output_ptr, x_ptr, y_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = x_val + tl.floor(y_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_floor(dtype, shape): + logging.log(logging.DEBUG, f"shape = {shape}") + x = torch.rand(shape, dtype=eval('torch.' + dtype)).npu() + y = torch.rand(shape, dtype=eval('torch.' + dtype)).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = x + torch.floor(y) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_floor_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = x + torch.floor(y) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_floor_4d_5d[grid](output, x, y, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_floordiv.py b/third_party/ascend/unittest/generalization_cases/test_general_floordiv.py new file mode 100644 index 0000000000..9e5cf1c6ef --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_floordiv.py @@ -0,0 +1,163 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import torch +import test_common +from test_common import TestUtils +import logging + + +@triton.jit +def triton_floordiv(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, + XNUMEL: tl.constexpr, YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = X // Y + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_floordiv_4d_5d(output_ptr, x_ptr, y_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, + STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = x_val // y_val + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) # some shape with int8 over ub +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64']) +def test_floordiv(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor_int_withSigns(shape, dtype).npu() + y = test_common.generate_tensor_int_withSigns(shape, dtype).npu() + z = test_common.generate_tensor_int_withSigns(shape, dtype).npu() + + new_shape = shape + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + y[y == 0] = 1 + ans = x // y + ans_mask = (x.to(torch.int64) % y.to(torch.int64) != 0) & (~((x ^ y) > 0)).to(ans.dtype) + ans = ans + ans_mask + + if len(shape) == 1: + triton_floordiv[1, 1, shape[0]](output, x, y, z, 1, 1, 1, 1, 1, shape[0]) + elif len(shape) == 2: + if shape[0] > shape[1]: + triton_floordiv[1, shape[0], 1](output, x, y, z, 1, 1, shape[1], 1, shape[0], shape[1]) + else: + triton_floordiv[1, 1, shape[1]](output, x, y, z, 1, shape[0], 1, 1, shape[0], shape[1]) + elif len(shape) == 3: + if max(shape[0], shape[1], shape[2]) == shape[0]: + triton_floordiv[shape[0], 1, 1](output, x, y, z, 1, shape[1], shape[2], shape[0], shape[1], shape[2]) + elif max(shape[0], shape[1], shape[2]) == shape[1]: + triton_floordiv[1, shape[1], 1](output, x, y, z, shape[0], 1, shape[2], shape[0], shape[1], shape[2]) + else: + triton_floordiv[1, 1, shape[2]](output, x, y, z, shape[0], shape[1], 1, shape[0], shape[1], shape[2]) + else: + triton_floordiv[1, 1, 1](output, x, y, z, 1, 1, 1, 1, 1, 1) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64']) +def test_floordiv_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + new_shape = shape + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + y[y == 0] = 1 + ans = x // y + ans_mask = (x.to(torch.int64) % y.to(torch.int64) != 0) & (~((x ^ y) > 0)).to(ans.dtype) + ans = ans + ans_mask + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_floordiv_4d_5d[grid](output, x, y, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) + + +invalid_types = [ + 'float16', + 'float32', + 'bfloat16', +] + + +@pytest.mark.parametrize("sigtype", invalid_types) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "unexpected type") +def test_invalid_types(sigtype): + N = 32 + x = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + y = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + y = y.masked_fill(y == 0, 1) + z = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + output = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + + triton_floordiv[1, 1, 1](output, x, y, z, 32, 1, 1, 32, 1, 1) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_fma.py b/third_party/ascend/unittest/generalization_cases/test_general_fma.py new file mode 100644 index 0000000000..eb255558ca --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_fma.py @@ -0,0 +1,162 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import logging +import pytest +import test_common +from test_common import TestUtils +import math + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + Z = tl.load(z_ptr + idx) + + ret = tl.fma(X, Y, Z) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_fma_4d_5d(output_ptr, x_ptr, y_ptr, z_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, + BLOCK_2: tl.constexpr, BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, + STRIDE_0: tl.constexpr, STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, + STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + z_val = tl.load(z_ptr + offsets, masks) + ret = tl.fma(x_val, y_val, z_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) # math.fma do not support int dtype +def test_fma(dtype, shape): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + if (x.dtype == torch.bfloat16): + ans = x.to(torch.float32) * y.to(torch.float32) + z.to(torch.float32) + ans = ans.to(torch.bfloat16) + else: + ans = x * y + z + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_fma_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + if (x.dtype == torch.bfloat16): + ans = x.to(torch.float32) * y.to(torch.float32) + z.to(torch.float32) + ans = ans.to(torch.bfloat16) + else: + ans = x * y + z + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_fma_4d_5d[grid](output, x, y, z, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_gather.py b/third_party/ascend/unittest/generalization_cases/test_general_gather.py new file mode 100644 index 0000000000..ee2b8bc437 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_gather.py @@ -0,0 +1,150 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import numpy as np +import torch +import torch_npu +import triton +import triton.language as tl +import test_common +import pytest +from test_common import TestUtils, check_ub_mem_overflow, get_dtype_size + + +@pytest.mark.parametrize("src_shape, indices_shape, axis", [ + ([2, 2], [4, 2], 0), + ([3, 3], [1, 3], 0), + ([3, 4], [4, 4], 0), + ([4, 4], [8, 4], 0), + ([4, 32], [4, 16], 1), + ([4, 64], [4, 32], 1), + ([128, 64], [128, 128], 1), +]) +def test_gather(src_shape, indices_shape, axis): + + @triton.jit + def gather_kernel(src_ptr, idx_ptr, out_ptr, axis: tl.constexpr, src_dim0: tl.constexpr, src_dim1: tl.constexpr, + src_stride0: tl.constexpr, src_stride1: tl.constexpr, idx_dim0: tl.constexpr, + idx_dim1: tl.constexpr, idx_stride0: tl.constexpr, idx_stride1: tl.constexpr, + out_dim0: tl.constexpr, out_dim1: tl.constexpr, out_stride0: tl.constexpr, + out_stride1: tl.constexpr): + src_offs = (tl.arange(0, src_dim0)[:, None] * src_stride0 + tl.arange(0, src_dim1)[None, :] * src_stride1) + src = tl.load(src_ptr + src_offs) + + idx_offs = (tl.arange(0, idx_dim0)[:, None] * idx_stride0 + tl.arange(0, idx_dim1)[None, :] * idx_stride1) + idx = tl.load(idx_ptr + idx_offs) + + out = tl.gather(src, idx, axis) + + out_offs = (tl.arange(0, out_dim0)[:, None] * out_stride0 + tl.arange(0, out_dim1)[None, :] * out_stride1) + tl.store(out_ptr + out_offs, out) + + def triton_gather(src: torch.Tensor, axis: int, indices: torch.Tensor): + output = torch.empty(indices.shape, dtype=src.dtype, device=src.device) + gather_kernel[(1, )](src, indices, output, axis, src.shape[0], src.shape[1], + src.stride(0), src.stride(1), indices.shape[0], indices.shape[1], indices.stride(0), + indices.stride(1), output.shape[0], output.shape[1], output.stride(0), output.stride(1)) + return output + + DEV = "npu" + src = torch.randn(src_shape, device=DEV) + indices = torch.randint(0, src.shape[axis], indices_shape, device=DEV) + + dtype_size = get_dtype_size('int32') + if dtype_size * math.prod(src.shape) >= (TestUtils.ub_size / 8): + print(f"dtype:int32 shape:{src.shape} mem overflow") + return + + ref = torch.gather(src, axis, indices) + result = triton_gather(src, axis, indices) + torch.testing.assert_close(result, ref, rtol=0, atol=0) + + +@triton.jit +def gather_kernel_multi_d(src_ptr, idx_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, + MB: tl.constexpr, NB: tl.constexpr, I_XB: tl.constexpr, I_YB: tl.constexpr, + I_ZB: tl.constexpr, I_MB: tl.constexpr, I_NB: tl.constexpr, DIMS: tl.constexpr, + AXIS: tl.constexpr): + in_offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if DIMS > 1: + in_offsets = in_offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if DIMS > 2: + in_offsets = in_offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if DIMS > 3: + in_offsets = in_offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if DIMS > 4: + in_offsets = in_offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + idx_offsets = tl.arange(0, I_XB) * (I_YB * I_ZB * I_MB * I_NB) + if DIMS > 1: + idx_offsets = idx_offsets[:, None] + tl.arange(0, I_YB)[None, :] * (I_ZB * I_MB * I_NB) + if DIMS > 2: + idx_offsets = idx_offsets[:, :, None] + tl.arange(0, I_ZB)[None, None, :] * (I_MB * I_NB) + if DIMS > 3: + idx_offsets = idx_offsets[:, :, :, None] + tl.arange(0, I_MB)[None, None, None, :] * I_NB + if DIMS > 4: + idx_offsets = idx_offsets[:, :, :, :, None] + tl.arange(0, I_NB)[None, None, None, None, :] + + src = tl.load(src_ptr + in_offsets) + idx = tl.load(idx_ptr + idx_offsets) + + out = tl.gather(src, idx, AXIS) + + tl.store(out_ptr + idx_offsets, out) + + +def triton_gather_multi_d(src: torch.Tensor, axis: int, indices: torch.Tensor): + output = torch.empty(indices.shape, dtype=src.dtype, device=src.device) + + s_shape = [*(src.shape)] + while len(s_shape) < 5: + s_shape.append(1) + i_shape = [*(indices.shape)] + while len(i_shape) < 5: + i_shape.append(1) + gather_kernel_multi_d[(1, )](src, indices, output, *s_shape, *i_shape, len(src.shape), axis) + return output + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize("src_shape, indices_shape, axis", [ + ((2, 2, 4, 8), (2, 2, 4, 8), 0), + ((2, 2, 4, 1), (2, 2, 4, 1), 3), + ((2, 3, 4, 8), (2, 3, 4, 8), 1), + ((2, 3, 4, 8), (2, 3, 4, 8), 2), + ((2, 2, 2, 4, 1), (2, 2, 2, 4, 1), 4), + ((2, 2, 2, 4, 8), (2, 2, 2, 4, 8), 1), + ((2, 2, 3, 4, 8), (2, 2, 3, 4, 8), 2), + ((2, 2, 3, 4, 8), (2, 2, 3, 4, 8), 0), +]) +def test_gather_4d_5d(src_shape, indices_shape, axis): + DEV = "npu" + src = torch.randn(src_shape, device=DEV) + indices = torch.randint(0, src.shape[axis], indices_shape, device=DEV) + + ref = torch.gather(src, axis, indices) + result = triton_gather_multi_d(src, axis, indices) + torch.testing.assert_close(result, ref, rtol=0, atol=0) + + +if __name__ == "__main__": + test_gather([4, 64], [4, 32], 1) + print("success: test_gather") diff --git a/third_party/ascend/unittest/generalization_cases/test_general_interleave.py b/third_party/ascend/unittest/generalization_cases/test_general_interleave.py new file mode 100644 index 0000000000..cce95fef86 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_interleave.py @@ -0,0 +1,216 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import logging +import torch +import torch_npu +import pytest +import test_common +from test_common import TestUtils + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + zoffs2 = tl.program_id(2) * ZB * 2 + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + zidx2 = tl.arange(0, 2 * ZB) + zoffs2 + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.interleave(X, Y) + + oidx = xidx[:, None, None] * YNUMEL * ZNUMEL * 2 + yidx[None, :, None] * ZNUMEL * 2 + zidx2[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def triton_interleave_4d( + output_ptr, + x_ptr, + y_ptr, + BLOCK_0: tl.constexpr, + BLOCK_1: tl.constexpr, + BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, + SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, + SHAPE_3: tl.constexpr, + STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, + STRIDE_2: tl.constexpr, + STRIDE_3: tl.constexpr, +): + pid = tl.program_id(0) + tmp0 = tl.arange(0, BLOCK_0)[:, None, None, None] + tmp1 = tl.arange(0, BLOCK_1)[None, :, None, None] + tmp2 = tl.arange(0, BLOCK_2)[None, None, :, None] + tmp3 = tl.arange(0, BLOCK_3)[None, None, None, :] + tmp4 = tl.arange(0, 2 * BLOCK_3)[None, None, None, :] + offsets = pid + tmp0 * STRIDE_0 + tmp1 * STRIDE_1 + tmp2 * STRIDE_2 + tmp3 * STRIDE_3 + masks = (tmp0 < SHAPE_0) & (tmp1 < SHAPE_1) & (tmp2 < SHAPE_2) & (tmp3 < SHAPE_3) + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + + ret = tl.interleave(x_val, y_val) + + out_offsets = pid + tmp0 * STRIDE_0 * 2 + tmp1 * STRIDE_1 * 2 + tmp2 * STRIDE_2 * 2 + tmp4 * STRIDE_3 + out_masks = (tmp0 < SHAPE_0) & (tmp1 < SHAPE_1) & (tmp2 < SHAPE_2) & (tmp4 < 2 * SHAPE_3) + tl.store(output_ptr + out_offsets, ret, mask=out_masks) + + +@triton.jit +def triton_interleave_5d(output_ptr, x_ptr, y_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, + STRIDE_4: tl.constexpr): + pid = tl.program_id(0) + tmp0 = tl.arange(0, BLOCK_0)[:, None, None, None, None] + tmp1 = tl.arange(0, BLOCK_1)[None, :, None, None, None] + tmp2 = tl.arange(0, BLOCK_2)[None, None, :, None, None] + tmp3 = tl.arange(0, BLOCK_3)[None, None, None, :, None] + tmp4 = tl.arange(0, BLOCK_4)[None, None, None, None, :] + tmp5 = tl.arange(0, 2 * BLOCK_4)[None, None, None, None, :] + offsets = pid + tmp0 * STRIDE_0 + tmp1 * STRIDE_1 + tmp2 * STRIDE_2 + tmp3 * STRIDE_3 + tmp4 * STRIDE_4 + masks = (tmp0 < SHAPE_0) & (tmp1 < SHAPE_1) & (tmp2 < SHAPE_2) & (tmp3 < SHAPE_3) & (tmp4 < SHAPE_4) + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + + ret = tl.interleave(x_val, y_val) + + out_offsets = pid + tmp0 * STRIDE_0 * 2 + tmp1 * STRIDE_1 * 2 + tmp2 * STRIDE_2 * 2 + tmp3 * STRIDE_3 * 2 + tmp5 * STRIDE_4 + out_masks = (tmp0 < SHAPE_0) & (tmp1 < SHAPE_1) & (tmp2 < SHAPE_2) & (tmp3 < SHAPE_3) & (tmp5 < 2 * SHAPE_4) + tl.store(output_ptr + out_offsets, ret, mask=out_masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_interleave(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = torch.full(shape, 100, dtype=eval('torch.' + dtype)).npu() + y = torch.full(shape, 30, dtype=eval('torch.' + dtype)).npu() + new_shape = shape[:-1] + (2 * shape[-1], ) + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.stack((x, y), dim=-1).reshape(new_shape) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_interleave_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape[:-1] + (2 * shape[-1], ) + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.stack((x, y), dim=-1).reshape(new_shape) + + blocks = list(x.size()) + strides = list(x.stride()) + + grid = (1, ) + if len(shape) == 4: + triton_interleave_4d[grid](output, x, y, *blocks, *blocks, *strides) + else: + triton_interleave_5d[grid](output, x, y, *blocks, *blocks, *strides) + test_common.validate_cmp(dtype, ans, output) + + +@triton.jit +def fn_npu_dtype(output_ptr, x_ptr, y_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.interleave(X, Y) + + oidx = xidx[:, None, None] * YB * ZB * 2 + yidx[None, :, None] * ZB * 2 + tl.arange(0, 2 * ZB)[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@pytest.mark.parametrize('para_type,data_type,XB,YB,ZB', [ + ('bfloat16', eval('torch.bfloat16'), 8, 8, 4), + ('uint8', eval('torch.uint8'), 1, 256, 16), + ('bool', eval('torch.bool'), 1, 1, 2), +]) +def test_interleave_u(para_type, data_type, XB, YB, ZB): + x = torch.full((XB, YB, ZB), 100, dtype=data_type).npu() + y = torch.full((XB, YB, ZB), 30, dtype=data_type).npu() + output = torch.randint(1, (XB, YB, ZB * 2), dtype=data_type).npu() + ans = torch.stack((x, y), dim=-1).reshape(XB, YB, ZB * 2) + fn_npu_dtype[1, 1, 1](output, x, y, XB, YB, ZB) + test_common.validate_cmp(para_type, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_join.py b/third_party/ascend/unittest/generalization_cases/test_general_join.py new file mode 100644 index 0000000000..a1d8cd3cd0 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_join.py @@ -0,0 +1,229 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common +from test_common import TestUtils +import logging + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + ret = tl.join(X, Y) + + oidx = xidx[:, None, None, None] * YNUMEL * ZNUMEL * 2 + yidx[None, :, None, None] * ZNUMEL * 2 + \ + zidx[None, None, :, None] * 2 + tl.arange(0, 2)[None, None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def triton_join_4d( + output_ptr, + x_ptr, + y_ptr, + BLOCK_0: tl.constexpr, + BLOCK_1: tl.constexpr, + BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, + SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, + SHAPE_3: tl.constexpr, + STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, + STRIDE_2: tl.constexpr, + STRIDE_3: tl.constexpr, +): + pid = tl.program_id(0) + tmp0 = tl.arange(0, BLOCK_0)[:, None, None, None] + tmp1 = tl.arange(0, BLOCK_1)[None, :, None, None] + tmp2 = tl.arange(0, BLOCK_2)[None, None, :, None] + tmp3 = tl.arange(0, BLOCK_3)[None, None, None, :] + + offsets = pid + tmp0 * STRIDE_0 + tmp1 * STRIDE_1 + tmp2 * STRIDE_2 + tmp3 * STRIDE_3 + masks = (tmp0 < SHAPE_0) & (tmp1 < SHAPE_1) & (tmp2 < SHAPE_2) & (tmp3 < SHAPE_3) + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + + ret = tl.join(x_val, y_val) + + out_tmp0 = tl.arange(0, BLOCK_0)[:, None, None, None, None] + out_tmp1 = tl.arange(0, BLOCK_1)[None, :, None, None, None] + out_tmp2 = tl.arange(0, BLOCK_2)[None, None, :, None, None] + out_tmp3 = tl.arange(0, BLOCK_3)[None, None, None, :, None] + out_tmp4 = tl.arange(0, 2)[None, None, None, None, :] + out_offsets = pid + out_tmp0 * STRIDE_0 * 2 + out_tmp1 * STRIDE_1 * 2 + out_tmp2 * STRIDE_2 * 2 \ + + out_tmp3 * STRIDE_3 * 2 + out_tmp4 + out_masks = (out_tmp0 < SHAPE_0) & (out_tmp1 < SHAPE_1) & (out_tmp2 < SHAPE_2) \ + & (out_tmp3 < SHAPE_3) & (out_tmp4 < 2) + tl.store(output_ptr + out_offsets, ret, mask=out_masks) + + +@triton.jit +def triton_join_5d(output_ptr, x_ptr, y_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + pid = tl.program_id(0) + tmp0 = tl.arange(0, BLOCK_0)[:, None, None, None, None] + tmp1 = tl.arange(0, BLOCK_1)[None, :, None, None, None] + tmp2 = tl.arange(0, BLOCK_2)[None, None, :, None, None] + tmp3 = tl.arange(0, BLOCK_3)[None, None, None, :, None] + tmp4 = tl.arange(0, BLOCK_4)[None, None, None, None, :] + + offsets = pid + tmp0 * STRIDE_0 + tmp1 * STRIDE_1 + tmp2 * STRIDE_2 + tmp3 * STRIDE_3 + tmp4 * STRIDE_4 + masks = (tmp0 < SHAPE_0) & (tmp1 < SHAPE_1) & (tmp2 < SHAPE_2) & (tmp3 < SHAPE_3) & (tmp4 < SHAPE_4) + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + + ret = tl.join(x_val, y_val) + + out_tmp0 = tl.arange(0, BLOCK_0)[:, None, None, None, None, None] + out_tmp1 = tl.arange(0, BLOCK_1)[None, :, None, None, None, None] + out_tmp2 = tl.arange(0, BLOCK_2)[None, None, :, None, None, None] + out_tmp3 = tl.arange(0, BLOCK_3)[None, None, None, :, None, None] + out_tmp4 = tl.arange(0, BLOCK_4)[None, None, None, None, :, None] + out_tmp5 = tl.arange(0, 2)[None, None, None, None, None, :] + out_offsets = pid + out_tmp0 * STRIDE_0 * 2 + out_tmp1 * STRIDE_1 * 2 + out_tmp2 * STRIDE_2 * 2 \ + + out_tmp3 * STRIDE_3 * 2 + out_tmp4 * STRIDE_4 * 2 + out_tmp5 + out_masks = (out_tmp0 < SHAPE_0) & (out_tmp1 < SHAPE_1) & (out_tmp2 < SHAPE_2) \ + & (out_tmp3 < SHAPE_3) & (out_tmp4 < SHAPE_4) & (out_tmp5 < 2) + tl.store(output_ptr + out_offsets, ret, mask=out_masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_join(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = torch.full(shape, 100, dtype=eval('torch.' + dtype)).npu() + y = torch.full(shape, 30, dtype=eval('torch.' + dtype)).npu() + new_shape = shape + (2, ) + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.stack((x, y), dim=-1) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +def test_join_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape + (2, ), dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.stack((x, y), dim=-1) + + blocks = list(x.size()) + strides = list(x.stride()) + + grid = (1, ) + if len(shape) == 4: + triton_join_4d[grid](output, x, y, *blocks, *blocks, *strides) + else: + triton_join_5d[grid](output, x, y, *blocks, *blocks, *strides) + test_common.validate_cmp(dtype, ans, output) + + +@triton.jit +def fn_npu_dtype(output_ptr, x_ptr, y_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + + idx = xidx[:, None] * YB + yidx[None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.join(X, Y) + + oidx = xidx[:, None, None] * YB * 2 + yidx[None, :, None] * 2 + tl.arange(0, 2)[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@pytest.mark.parametrize('para_type,data_type,XB,YB,ZB', [ + ('bfloat16', eval('torch.bfloat16'), 8, 8, 4), + ('uint8', eval('torch.uint8'), 1, 256, 16), + ('bool', eval('torch.bool'), 1, 1, 2), +]) +def test_join_u(para_type, data_type, XB, YB, ZB): + x = torch.full((XB, YB), 100, dtype=data_type).npu() + y = torch.full((XB, YB), 30, dtype=data_type).npu() + + ans = torch.stack((x, y), dim=-1) + output = torch.randint(1, (XB, YB, 2), dtype=data_type).npu() + fn_npu_dtype[1, 1, 1](output, x, y, XB, YB, ZB, debug=True) + test_common.validate_cmp(para_type, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_log.py b/third_party/ascend/unittest/generalization_cases/test_general_log.py new file mode 100644 index 0000000000..8bcafa2274 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_log.py @@ -0,0 +1,148 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import triton.language.extra.ascend.libdevice as libdevice +import test_common +from test_common import TestUtils +import math + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.log(X) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_log_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = tl.log(x_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_log(dtype, shape): + logging.log(logging.DEBUG, f"shape = {shape}") + x = torch.rand(shape, dtype=eval('torch.' + dtype)).npu() + y = torch.rand(shape, dtype=eval('torch.' + dtype)).npu() + z = torch.rand(shape, dtype=eval('torch.' + dtype)).npu() + + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.log(x).to(eval('torch.' + dtype)) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_log_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.log(x).to(eval('torch.' + dtype)) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_log_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_log2.py b/third_party/ascend/unittest/generalization_cases/test_general_log2.py new file mode 100644 index 0000000000..0a0321466d --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_log2.py @@ -0,0 +1,153 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import triton.language.extra.ascend.libdevice as libdevice +import test_common +from test_common import TestUtils +import math + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.log2(X) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_log2_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = tl.log2(x_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_log2(dtype, shape): + logging.log(logging.DEBUG, f"shape = {shape}") + x = torch.rand(shape, dtype=eval('torch.' + dtype)).npu() + y = torch.rand(shape, dtype=eval('torch.' + dtype)).npu() + z = torch.rand(shape, dtype=eval('torch.' + dtype)).npu() + + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.log2(x).to(eval('torch.' + dtype)) + + if len(shape) == 1: + fn_npu_[1, 1, 1](output, x, y, z, 1, 1, shape[0], 1, 1, shape[0]) + elif len(shape) == 2: + fn_npu_[1, shape[0], 1](output, x, y, z, 1, 1, shape[1], 1, shape[0], shape[1]) + else: + mx = max(shape[0], shape[1], shape[2]) + if mx == shape[0]: + fn_npu_[shape[0], 1, 1](output, x, y, z, 1, shape[1], shape[2], shape[0], shape[1], shape[2]) + elif mx == shape[1]: + fn_npu_[1, shape[1], 1](output, x, y, z, shape[0], 1, shape[2], shape[0], shape[1], shape[2]) + else: + fn_npu_[1, 1, shape[2]](output, x, y, z, shape[0], shape[1], 1, shape[0], shape[1], shape[2]) + + test_common.validate_cmp(dtype, ans, output) + + +invalid_dtypes = [ + 'int8', + 'int16', + 'int32', + 'uint32', + 'int64', + 'bool', +] + + +@pytest.mark.parametrize("dtype", invalid_dtypes) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "Expected dtype") +def test_log2_invalid_dtype_case(dtype): + x = test_common.generate_tensor((1, ), dtype).npu() + y = test_common.generate_tensor((1, ), dtype).npu() + z = test_common.generate_tensor((1, ), dtype).npu() + + output = torch.randint(1, (1, ), dtype=eval('torch.' + dtype)).npu() + fn_npu_[1, 1, 1](output, x, y, z, 1, 1, 1, 1, 1, 1) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_log2_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.log2(x).to(eval('torch.' + dtype)) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_log2_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_maximum.py b/third_party/ascend/unittest/generalization_cases/test_general_maximum.py new file mode 100644 index 0000000000..e7de75d0d3 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_maximum.py @@ -0,0 +1,153 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils +import math + + +def torch_maximum(x, y): + return torch.maximum(x, y) + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.maximum(X, Y) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_maximum_4d_5d(output_ptr, x_ptr, y_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, + STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = tl.maximum(x_val, y_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16', 'int8', 'int16', 'int32', 'int64', 'bool']) +def test_maximum(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_maximum(x, y) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16', 'bool']) +def test_maximum_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_maximum(x, y) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_maximum_4d_5d[grid](output, x, y, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_minimum.py b/third_party/ascend/unittest/generalization_cases/test_general_minimum.py new file mode 100644 index 0000000000..1b419c0e5e --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_minimum.py @@ -0,0 +1,153 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils +import math + + +def torch_minimum(x, y): + return torch.minimum(x, y) + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.minimum(X, Y) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_minimum_4d_5d(output_ptr, x_ptr, y_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, + STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = tl.minimum(x_val, y_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16', 'int8', 'int16', 'int32', 'int64', 'bool']) +def test_minimum(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_minimum(x, y) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16', 'bool']) +def test_minimum_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_minimum(x, y) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_minimum_4d_5d[grid](output, x, y, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_mul.py b/third_party/ascend/unittest/generalization_cases/test_general_mul.py new file mode 100644 index 0000000000..29666d72e4 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_mul.py @@ -0,0 +1,155 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import torch +import test_common +from test_common import TestUtils +import logging +import numpy as np + + +@triton.jit +def triton_mul(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, + XNUMEL: tl.constexpr, YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = X * Y + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_mul_4d_5d(output_ptr, x_ptr, y_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = x_val * y_val + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) # some shape with int8 over ub +@pytest.mark.parametrize('dtype', ['bool', 'int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +def test_mul(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + + ans = x * y + output = torch.zeros_like(ans) + + if len(shape) == 1: + triton_mul[1, 1, shape[0]](output, x, y, z, 1, 1, 1, 1, 1, shape[0]) + elif len(shape) == 2: + if shape[0] > shape[1]: + triton_mul[1, shape[0], 1](output, x, y, z, 1, 1, shape[1], 1, shape[0], shape[1]) + else: + triton_mul[1, 1, shape[1]](output, x, y, z, 1, shape[0], 1, 1, shape[0], shape[1]) + elif len(shape) == 3: + if max(shape[0], shape[1], shape[2]) == shape[0]: + triton_mul[shape[0], 1, 1](output, x, y, z, 1, shape[1], shape[2], shape[0], shape[1], shape[2]) + elif max(shape[0], shape[1], shape[2]) == shape[1]: + triton_mul[1, shape[1], 1](output, x, y, z, shape[0], 1, shape[2], shape[0], shape[1], shape[2]) + else: + triton_mul[1, 1, shape[2]](output, x, y, z, shape[0], shape[1], 1, shape[0], shape[1], shape[2]) + else: + triton_mul[1, 1, 1](output, x, y, z, 1, 1, 1, 1, 1, 1) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +def test_mul_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = x * y + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_mul_4d_5d[grid](output, x, y, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', ['uint16', 'uint32', 'uint64']) +def test_mul_uint(shape, dtype): + torch_dtype = eval('torch.' + dtype) + np_x0 = test_common.generate_numpy(shape, dtype) + np_x1 = test_common.generate_numpy(shape, dtype) + np_x2 = test_common.generate_numpy(shape, dtype) + + x0 = torch.from_numpy(np_x0).to(torch_dtype).npu() + x1 = torch.from_numpy(np_x1).to(torch_dtype).npu() + x2 = torch.from_numpy(np_x2).to(torch_dtype).npu() + + #numpy result + ans_numpy = np_x0 * np_x1 + z_ref1 = torch.from_numpy(ans_numpy).npu() + + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_mul[1, 1, shape[0]](triton_res, x0, x1, x2, 1, 1, 1, 1, 1, shape[0]) + test_common.validate_cmp(dtype, z_ref1, triton_res) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_ravel.py b/third_party/ascend/unittest/generalization_cases/test_general_ravel.py new file mode 100644 index 0000000000..2e6735709d --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_ravel.py @@ -0,0 +1,189 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import logging +import pytest +import test_common +from test_common import TestUtils +import math + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.ravel(X) + + oidx = tl.arange(0, XB * YB * ZB) + xoffs * YNUMEL * ZNUMEL + yoffs * ZNUMEL + zoffs + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def triton_ravel_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = tl.ravel(x_val) + + pid0 = tl.program_id(0) + + flat_idx = tl.arange(0, BLOCK_0 * BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) + out_offsets = pid0 * BLOCK_0 * BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4 + flat_idx + out_masks = out_offsets < SHAPE_0 * SHAPE_1 * SHAPE_2 * SHAPE_3 * SHAPE_4 + tl.store(output_ptr + out_offsets, ret, mask=out_masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_ravel(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = torch.full(shape, 100, dtype=eval('torch.' + dtype)).npu() + y = torch.full(shape, 30, dtype=eval('torch.' + dtype)).npu() + new_shape = (x.numel(), ) + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.ravel(x) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + if xnumel > 1: + grid = (XB, 1, 1) + XB = 1 + elif ynumel > 1: + grid = (1, YB, 1) + YB = 1 + else: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_ravel_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = torch.full(shape, 100, dtype=eval('torch.' + dtype)).npu() + + output = torch.randint(1, (x.numel(), ), dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.ravel(x) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_ravel_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) + + +@triton.jit +def fn_npu_dtype(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.ravel(X) + + oidx = tl.arange(0, XB * YB * ZB) + tl.store(output_ptr + oidx, ret) + + +@pytest.mark.parametrize('sigtype, dtype, XB, YB, ZB', [ + ('bfloat16', torch.bfloat16, 2, 8, 4), + ('uint8', torch.uint8, 1, 256, 16), + ('bool', torch.bool, 1, 1, 2), +]) +def test_ravel_u(sigtype, dtype, XB, YB, ZB): + x = test_common.generate_tensor((XB, YB, ZB), sigtype).npu() + ans = torch.ravel(x) + output = test_common.generate_tensor((1, XB * YB * ZB), sigtype).npu() + output = output.reshape(-1) + fn_npu_dtype[1, 1, 1](output, x, XB, YB, ZB) + test_common.validate_cmp(sigtype, output, ans) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_reshape.py b/third_party/ascend/unittest/generalization_cases/test_general_reshape.py new file mode 100644 index 0000000000..25b661b0ad --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_reshape.py @@ -0,0 +1,160 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils +import math +import logging + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.reshape(X, (ZB * YB * XB, )) + + oidx = tl.arange(0, XB * YB * ZB) + xoffs * YNUMEL * ZNUMEL + yoffs * ZNUMEL + zoffs + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def triton_reshape_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, + STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = tl.reshape(x_val, (SHAPE_0 * SHAPE_1 * SHAPE_2 * SHAPE_3 * SHAPE_4, )) + + pid0 = tl.program_id(0) + + flat_idx = tl.arange(0, BLOCK_0 * BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) + out_offsets = pid0 * BLOCK_0 * BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4 + flat_idx + out_masks = out_offsets < SHAPE_0 * SHAPE_1 * SHAPE_2 * SHAPE_3 * SHAPE_4 + tl.store(output_ptr + out_offsets, ret, mask=out_masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_reshape(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = torch.full(shape, 100, dtype=eval('torch.' + dtype)).npu() + y = torch.full(shape, 30, dtype=eval('torch.' + dtype)).npu() + new_shape = (x.numel(), ) + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = x.reshape(-1) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + if xnumel > 1: + grid = (XB, 1, 1) + XB = 1 + elif ynumel > 1: + grid = (1, YB, 1) + YB = 1 + else: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_reshape_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = torch.full(shape, 100, dtype=eval('torch.' + dtype)).npu() + + output = torch.randint(1, (x.numel(), ), dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = x.reshape(-1) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_reshape_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_rsqrt.py b/third_party/ascend/unittest/generalization_cases/test_general_rsqrt.py new file mode 100644 index 0000000000..54a0984fec --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_rsqrt.py @@ -0,0 +1,151 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import logging +import pytest +import test_common +from test_common import TestUtils + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.rsqrt(X) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_rsqrt_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = tl.rsqrt(x_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_rsqrt( + dtype, + shape, +): + x = test_common.generate_tensor(shape, dtype).abs().npu() + y = test_common.generate_tensor(shape, dtype).abs().npu() + z = test_common.generate_tensor(shape, dtype).abs().npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.rsqrt(x) + + if len(shape) == 1: + fn_npu_[1, 1, 1](output, x, y, z, 1, 1, shape[0], 1, 1, shape[0]) + elif len(shape) == 2: + fn_npu_[1, shape[0], 1](output, x, y, z, 1, 1, shape[1], 1, shape[0], shape[1]) + else: + mx = max(shape[0], shape[1], shape[2]) + if mx == shape[0]: + fn_npu_[shape[0], 1, 1](output, x, y, z, 1, shape[1], shape[2], shape[0], shape[1], shape[2]) + elif mx == shape[1]: + fn_npu_[1, shape[1], 1](output, x, y, z, shape[0], 1, shape[2], shape[0], shape[1], shape[2]) + else: + fn_npu_[1, 1, shape[2]](output, x, y, z, shape[0], shape[1], 1, shape[0], shape[1], shape[2]) + + test_common.validate_cmp(dtype, ans, output) + + +invalid_dtypes = [ + 'int8', + 'int16', + 'int32', + 'uint32', + 'int64', + 'bool', +] + + +@pytest.mark.parametrize("dtype", invalid_dtypes) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "Expected dtype") +def test_rsqrt_invalid_dtype_case(dtype): + x = test_common.generate_tensor((1, ), dtype).npu() + y = test_common.generate_tensor((1, ), dtype).npu() + z = test_common.generate_tensor((1, ), dtype).npu() + + output = torch.randint(1, (1, ), dtype=eval('torch.' + dtype)).npu() + fn_npu_[1, 1, 1](output, x, y, z, 1, 1, 1, 1, 1, 1) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_rsqrt_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.rsqrt(x) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_rsqrt_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_sigmoid.py b/third_party/ascend/unittest/generalization_cases/test_general_sigmoid.py new file mode 100644 index 0000000000..91737647e6 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_sigmoid.py @@ -0,0 +1,175 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import logging +import pytest +import test_common +from test_common import TestUtils +import math + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.sigmoid(X) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_sigmoid_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, + STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = tl.sigmoid(x_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_sigmoid( + dtype, + shape, +): + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + if (x.dtype == torch.bfloat16): + ans = torch.sigmoid(x.to(torch.float32)).to(torch.bfloat16) + else: + ans = torch.sigmoid(x) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +invalid_dtypes = [ + 'int8', + 'int16', + 'int32', + 'uint32', + 'int64', + 'bool', +] + + +@pytest.mark.parametrize("dtype", invalid_dtypes) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "Expected dtype") +def test_sigmoid_invalid_dtype_case(dtype): + x = test_common.generate_tensor((1, ), dtype).npu() + y = test_common.generate_tensor((1, ), dtype).npu() + z = test_common.generate_tensor((1, ), dtype).npu() + + output = torch.randint(1, (1, ), dtype=eval('torch.' + dtype)).npu() + fn_npu_[1, 1, 1](output, x, y, z, 1, 1, 1, 1, 1, 1) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_sigmoid_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + if (x.dtype == torch.bfloat16): + ans = torch.sigmoid(x.to(torch.float32)).to(torch.bfloat16) + else: + ans = torch.sigmoid(x) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_sigmoid_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_sin.py b/third_party/ascend/unittest/generalization_cases/test_general_sin.py new file mode 100644 index 0000000000..f52d0405de --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_sin.py @@ -0,0 +1,173 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import numpy as np +import pytest +import test_common +from test_common import TestUtils +import math + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.sin(X) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_sin_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = tl.sin(x_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +import logging + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_sin( + dtype, + shape, +): + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.sin(x) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +invalid_dtypes = [ + 'int8', + 'int16', + 'int32', + 'uint32', + 'int64', + 'bool', +] + + +@pytest.mark.parametrize("dtype", invalid_dtypes) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "Expected dtype") +def test_sin_invalid_dtype_case(dtype): + x = test_common.generate_tensor((1, ), dtype).npu() + y = test_common.generate_tensor((1, ), dtype).npu() + z = test_common.generate_tensor((1, ), dtype).npu() + + output = torch.randint(1, (1, ), dtype=eval('torch.' + dtype)).npu() + fn_npu_[1, 1, 1](output, x, y, z, 1, 1, 1, 1, 1, 1) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_sin_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.sin(x) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_sin_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_softmax.py b/third_party/ascend/unittest/generalization_cases/test_general_softmax.py new file mode 100644 index 0000000000..ce4b34d4a1 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_softmax.py @@ -0,0 +1,179 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +# 实际实现与官网定义不符，可能和triton submodule版本有关， 当前的submodule 不接受指定dim，都是按第0维做softmax +# arith.maximum 不支持类似 1x3 -> 3 和 1 -> 1 的reduce +import triton +import triton.language as tl +import torch +import logging +import pytest +import test_common +from test_common import TestUtils +import math + + +def torch_softmax_d0(x1): + res = torch.softmax(x1, axis=0).to(x1.dtype) + return res + + +@triton.jit +def tt_softmax_1d(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, XB: tl.constexpr, + YB: tl.constexpr, ZB: tl.constexpr): + idx = tl.arange(0, XB) + x = tl.load(in_ptr + idx) + ret = tl.softmax(x) + tl.store(out_ptr + idx, ret) + + +@triton.jit +def tt_softmax_2d(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, XB: tl.constexpr, + YB: tl.constexpr, ZB: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + idx = xidx[:, None] * ynumel + yidx[None, :] + + a = tl.load(in_ptr + idx) + ret = tl.softmax(a) + + tl.store(out_ptr + idx, ret) + + +@triton.jit +def tt_softmax_3d(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, XB: tl.constexpr, + YB: tl.constexpr, ZB: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + + a = tl.load(in_ptr + idx) + ret = tl.softmax(a) + + tl.store(out_ptr + idx, ret) + + +@triton.jit +def triton_softmax_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, + STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = tl.softmax(x_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_softmax(dtype, shape): + logging.log(logging.DEBUG, f"shape = {shape}", flush=True) + torch.manual_seed(0) + x = torch.rand(shape, dtype=eval('torch.' + dtype), device="npu") * 10 + grid = (1, 1, 1) + + y_cal = torch.rand(shape, dtype=eval('torch.' + dtype), device="npu") + + y_ref = torch_softmax_d0(x) + if len(shape) == 1: + tt_softmax_1d[grid](x, y_cal, x.numel(), 1, 1, x.numel(), 1, 1) + elif len(shape) == 2: + xnumel, ynumel, znumel = shape + (1, ) + XB, YB, ZB = xnumel, ynumel, znumel + if x.numel() * x.element_size() > 8192: + grid = (1, ynumel, 1) + YB = 1 + tt_softmax_2d[grid](x, y_cal, xnumel, ynumel, znumel, XB, YB, ZB) + + elif len(shape) == 3: + mx = max(shape[1], shape[2]) + if mx == shape[1]: + tt_softmax_3d[1, shape[1], 1](x, y_cal, shape[0], shape[1], shape[2], shape[0], 1, shape[2]) + else: + tt_softmax_3d[1, 1, shape[2]](x, y_cal, shape[0], shape[1], shape[2], shape[0], shape[1], 1) + + test_common.validate_cmp(dtype, y_cal, y_ref) + + +invalid_types = [ + 'int8', + 'int16', + 'int32', + 'uint32', + 'int64', + 'bool', +] + + +@pytest.mark.parametrize("dtype", invalid_types) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "Expected dtype") +def test_softmax_invalid_dtype_case(dtype): + x0 = test_common.generate_tensor((1, ), dtype).npu() + + y_cal = torch.zeros((1, ), dtype=eval('torch.' + dtype)).npu() + tt_softmax_1d[1, 1, 1](x0, y_cal, 0, 0, 0, 1, 0, 0) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_softmax_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_softmax_d0(x) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_softmax_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_split.py b/third_party/ascend/unittest/generalization_cases/test_general_split.py new file mode 100644 index 0000000000..9efedcf71a --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_split.py @@ -0,0 +1,224 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils +import math + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, output_ptr1, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx=xidx[:,None,None,None]*YNUMEL*ZNUMEL*2+yidx[None,:,None,None]*ZNUMEL*2+ \ + zidx[None,None,:,None]*2 + tl.arange(0,2)[None,None,None,:] + + X = tl.load(x_ptr + idx) + + xx, yy = tl.split(X) + + oidx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + tl.store(output_ptr + oidx, xx) + tl.store(output_ptr1 + oidx, yy) + + +import logging + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_split(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = torch.full(shape, 100, dtype=eval('torch.' + dtype)).npu() + y = torch.full(shape, 30, dtype=eval('torch.' + dtype)).npu() + xx = torch.stack((x, y), dim=-1) + + a, b = torch.split(xx, 1, dim=-1) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + a = a.reshape(XB, YB, ZB) + b = b.reshape(XB, YB, ZB) + output = torch.randint(1, (XB, YB, ZB), dtype=eval('torch.' + dtype)).npu() + output1 = torch.randint(1, (XB, YB, ZB), dtype=eval('torch.' + dtype)).npu() + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + if xnumel > 1: + grid = (XB, 1, 1) + XB = 1 + elif ynumel > 1: + grid = (1, YB, 1) + YB = 1 + else: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, xx, output1, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, a, output) + test_common.validate_cmp(dtype, b, output1) + + +@triton.jit +def fn_npu_4_8d(output_ptr, x_ptr, output_ptr1, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, WB: tl.constexpr, + VB: tl.constexpr, UB: tl.constexpr, TB: tl.constexpr, SB: tl.constexpr): + + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + widx = tl.arange(0, WB) + vidx = tl.arange(0, VB) + uidx = tl.arange(0, UB) + tidx = tl.arange(0, TB) + sidx = tl.arange(0, SB) + + idx = (xidx[:, None, None, None, None, None, None, None, None] * YB * ZB * WB * VB * UB * TB * SB * 2 + + yidx[None, :, None, None, None, None, None, None, None] * ZB * WB * VB * UB * TB * SB * 2 + + zidx[None, None, :, None, None, None, None, None, None] * WB * VB * UB * TB * SB * 2 + + widx[None, None, None, :, None, None, None, None, None] * VB * UB * TB * SB * 2 + + vidx[None, None, None, None, :, None, None, None, None] * UB * TB * SB * 2 + + uidx[None, None, None, None, None, :, None, None, None] * TB * SB * 2 + + tidx[None, None, None, None, None, None, :, None, None] * SB * 2 + + sidx[None, None, None, None, None, None, None, :, None] * 2 + + tl.arange(0, 2)[None, None, None, None, None, None, None, None, :]) + + X = tl.load(x_ptr + idx) + xx, yy = tl.split(X) + + oidx = (xidx[:, None, None, None, None, None, None, None] * YB * ZB * WB * VB * UB * TB * SB + + yidx[None, :, None, None, None, None, None, None] * ZB * WB * VB * UB * TB * SB + + zidx[None, None, :, None, None, None, None, None] * WB * VB * UB * TB * SB + + widx[None, None, None, :, None, None, None, None] * VB * UB * TB * SB + + vidx[None, None, None, None, :, None, None, None] * UB * TB * SB + + uidx[None, None, None, None, None, :, None, None] * TB * SB + + tidx[None, None, None, None, None, None, :, None] * SB + sidx[None, None, None, None, None, None, None, :]) + + tl.store(output_ptr + oidx, xx) + tl.store(output_ptr1 + oidx, yy) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape_4_8d) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_split_4_8d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = torch.full(shape, 100, dtype=eval('torch.' + dtype)).npu() + y = torch.full(shape, 30, dtype=eval('torch.' + dtype)).npu() + xx = torch.stack((x, y), dim=-1) + + a, b = torch.split(xx, 1, dim=-1) + + if len(shape) == 1: + XB, YB, ZB, WB, VB, UB, TB, SB = 1, 1, 1, 1, 1, 1, 1, shape[0] + elif len(shape) == 2: + XB, YB, ZB, WB, VB, UB, TB, SB = 1, 1, 1, 1, 1, 1, shape[0], shape[1] + elif len(shape) == 3: + XB, YB, ZB, WB, VB, UB, TB, SB = 1, 1, 1, 1, 1, shape[0], shape[1], shape[2] + elif len(shape) == 4: + XB, YB, ZB, WB, VB, UB, TB, SB = 1, 1, 1, 1, shape[0], shape[1], shape[2], shape[3] + elif len(shape) == 5: + XB, YB, ZB, WB, VB, UB, TB, SB = 1, 1, 1, shape[0], shape[1], shape[2], shape[3], shape[4] + elif len(shape) == 6: + XB, YB, ZB, WB, VB, UB, TB, SB = 1, 1, shape[0], shape[1], shape[2], shape[3], shape[4], shape[5] + elif len(shape) == 7: + XB, YB, ZB, WB, VB, UB, TB, SB = 1, shape[0], shape[1], shape[2], shape[3], shape[4], shape[5], shape[6] + else: + XB, YB, ZB, WB, VB, UB, TB, SB = shape + + a = a.reshape(XB, YB, ZB, WB, VB, UB, TB, SB) + b = b.reshape(XB, YB, ZB, WB, VB, UB, TB, SB) + + output = torch.randint(1, (XB, YB, ZB, WB, VB, UB, TB, SB), dtype=eval('torch.' + dtype)).npu() + output1 = torch.randint(1, (XB, YB, ZB, WB, VB, UB, TB, SB), dtype=eval('torch.' + dtype)).npu() + + grid = (1, 1, 1) + fn_npu_4_8d[grid](output, xx, output1, XB, YB, ZB, WB, VB, UB, TB, SB) + + test_common.validate_cmp(dtype, a, output) + test_common.validate_cmp(dtype, b, output1) + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, output_ptr1, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + + xx, yy = tl.split(X) + + oidx = xidx[:, None] * YB + yidx[None, :] + + tl.store(output_ptr + oidx, xx) + tl.store(output_ptr1 + oidx, yy) + + +@pytest.mark.parametrize('para_type, data_type, XB, YB, ZB', [ + ('bfloat16', torch.bfloat16, 2, 8, 2), + ('uint8', torch.uint8, 1, 256, 2), + ('bool', torch.bool, 1, 1, 2), +]) +def test_split_u(para_type, data_type, XB, YB, ZB): + x = test_common.generate_tensor((XB, YB, ZB), para_type).npu() + a, b = torch.split(x, 1, dim=-1) + a = a.reshape(XB, YB) + b = b.reshape(XB, YB) + + output = test_common.generate_tensor((XB, YB), para_type).npu() + output1 = test_common.generate_tensor((XB, YB), para_type).npu() + fn_npu_[1, 1, 1](output, x, output1, XB, YB, ZB, debug=True) + + test_common.validate_cmp(para_type, a, output) + test_common.validate_cmp(para_type, b, output1) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_sub.py b/third_party/ascend/unittest/generalization_cases/test_general_sub.py new file mode 100644 index 0000000000..a831d6c421 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_sub.py @@ -0,0 +1,155 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import torch +import test_common +from test_common import TestUtils +import logging +import numpy as np + + +@triton.jit +def triton_sub(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, + XNUMEL: tl.constexpr, YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = X - Y + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_sub_4d_5d(output_ptr, x_ptr, y_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = x_val - y_val + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['bool', 'int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +def test_sub(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + + ans = x - y + output = torch.zeros_like(ans) + + if len(shape) == 1: + triton_sub[1, 1, shape[0]](output, x, y, z, 1, 1, 1, 1, 1, shape[0]) + elif len(shape) == 2: + if shape[0] > shape[1]: + triton_sub[1, shape[0], 1](output, x, y, z, 1, 1, shape[1], 1, shape[0], shape[1]) + else: + triton_sub[1, 1, shape[1]](output, x, y, z, 1, shape[0], 1, 1, shape[0], shape[1]) + elif len(shape) == 3: + if max(shape[0], shape[1], shape[2]) == shape[0]: + triton_sub[shape[0], 1, 1](output, x, y, z, 1, shape[1], shape[2], shape[0], shape[1], shape[2]) + elif max(shape[0], shape[1], shape[2]) == shape[1]: + triton_sub[1, shape[1], 1](output, x, y, z, shape[0], 1, shape[2], shape[0], shape[1], shape[2]) + else: + triton_sub[1, 1, shape[2]](output, x, y, z, shape[0], shape[1], 1, shape[0], shape[1], shape[2]) + else: + triton_sub[1, 1, 1](output, x, y, z, 1, 1, 1, 1, 1, 1) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +def test_sub_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = x - y + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_sub_4d_5d[grid](output, x, y, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', ['uint16', 'uint32', 'uint64']) +def test_sub_uint(shape, dtype): + torch_dtype = eval('torch.' + dtype) + np_x0 = test_common.generate_numpy(shape, dtype) + np_x1 = test_common.generate_numpy(shape, dtype) + np_x2 = test_common.generate_numpy(shape, dtype) + + x0 = torch.from_numpy(np_x0).to(torch_dtype).npu() + x1 = torch.from_numpy(np_x1).to(torch_dtype).npu() + x2 = torch.from_numpy(np_x2).to(torch_dtype).npu() + + #numpy result + ans_numpy = np_x0 - np_x1 + z_ref1 = torch.from_numpy(ans_numpy).npu() + + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_sub[1, 1, shape[0]](triton_res, x0, x1, x2, 1, 1, 1, 1, 1, shape[0]) + test_common.validate_cmp(dtype, z_ref1, triton_res) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_tensor_descriptor.py b/third_party/ascend/unittest/generalization_cases/test_general_tensor_descriptor.py new file mode 100644 index 0000000000..ec4a7c8bc0 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_tensor_descriptor.py @@ -0,0 +1,268 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common +from test_common import TestUtils + +full_dtype = test_common._float_dtypes + test_common._int_dtypes + test_common._uint_dtypes +temporarily_not_support_dtype = ['bool'] + + +@triton.jit +def triton_tensor_descriptor_2d( + out_ptr, + x_ptr, + M: tl.constexpr, + N: tl.constexpr, + M_BLOCK: tl.constexpr, + N_BLOCK: tl.constexpr, +): + in_desc = tl.make_tensor_descriptor( + x_ptr, + shape=[M, N], + strides=[N, 1], + block_shape=[M_BLOCK, N_BLOCK], + ) + out_desc = tl.make_tensor_descriptor( + out_ptr, + shape=[M, N], + strides=[N, 1], + block_shape=[M_BLOCK, N_BLOCK], + ) + moffset = tl.program_id(0) * M_BLOCK + noffset = tl.program_id(1) * N_BLOCK + block = in_desc.load([moffset, noffset]) + out_desc.store([moffset, noffset], block) + + +@triton.jit +def triton_tensor_descriptor_3d( + out_ptr, + x_ptr, + M: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, + stride_m: tl.constexpr, + stride_n: tl.constexpr, + stride_k: tl.constexpr, + M_BLOCK: tl.constexpr, + N_BLOCK: tl.constexpr, + K_BLOCK: tl.constexpr, +): + in_desc = tl.make_tensor_descriptor( + x_ptr, + shape=[M, N, K], + strides=[stride_m, stride_n, stride_k], + block_shape=[M_BLOCK, N_BLOCK, K_BLOCK], + ) + out_desc = tl.make_tensor_descriptor( + out_ptr, + shape=[M, N, K], + strides=[stride_m, stride_n, stride_k], + block_shape=[M_BLOCK, N_BLOCK, K_BLOCK], + ) + moffset = tl.program_id(0) * M_BLOCK + noffset = tl.program_id(1) * N_BLOCK + koffset = tl.program_id(2) * K_BLOCK + block = in_desc.load([moffset, noffset, koffset]) + out_desc.store([moffset, noffset, koffset], block) + + +@triton.jit +def triton_tensor_descriptor_4d( + out_ptr, + x_ptr, + SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, + SHAPE_3: tl.constexpr, + STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, + STRIDE_2: tl.constexpr, + STRIDE_3: tl.constexpr, + BLOCK_0: tl.constexpr, + BLOCK_1: tl.constexpr, + BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, +): + pid0 = tl.program_id(0) + pid1 = tl.program_id(1) + pid2 = tl.program_id(2) + idx2 = pid2 // BLOCK_3 + idx3 = pid2 % BLOCK_3 + o1 = pid0 * BLOCK_0 + o2 = pid1 * BLOCK_1 + o3 = idx2 * BLOCK_2 + o4 = idx3 * BLOCK_3 + in_desc = tl.make_tensor_descriptor( + x_ptr, + shape=[SHAPE_0, SHAPE_1, SHAPE_2, SHAPE_3], + strides=[STRIDE_0, STRIDE_1, STRIDE_2, STRIDE_3], + block_shape=[BLOCK_0, BLOCK_1, BLOCK_2, BLOCK_3], + ) + out_desc = tl.make_tensor_descriptor( + out_ptr, + shape=[SHAPE_0, SHAPE_1, SHAPE_2, SHAPE_3], + strides=[STRIDE_0, STRIDE_1, STRIDE_2, STRIDE_3], + block_shape=[BLOCK_0, BLOCK_1, BLOCK_2, BLOCK_3], + ) + block = in_desc.load([o1, o2, o3, o4]) + out_desc.store([o1, o2, o3, o4], block) + + +@triton.jit +def triton_tensor_descriptor_5d( + out_ptr, + x_ptr, + SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, + SHAPE_3: tl.constexpr, + SHAPE_4: tl.constexpr, + STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, + STRIDE_2: tl.constexpr, + STRIDE_3: tl.constexpr, + STRIDE_4: tl.constexpr, + BLOCK_0: tl.constexpr, + BLOCK_1: tl.constexpr, + BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, + BLOCK_4: tl.constexpr, +): + pid0 = tl.program_id(0) + pid1 = tl.program_id(1) + pid2 = tl.program_id(2) + idx3 = pid2 // (BLOCK_3 * BLOCK_4) + idx4 = (pid2 // BLOCK_4) % BLOCK_3 + idx5 = pid2 % BLOCK_4 + o1 = pid0 * BLOCK_0 + o2 = pid1 * BLOCK_1 + o3 = idx3 * BLOCK_2 + o4 = idx4 * BLOCK_3 + o5 = idx5 * BLOCK_4 + in_desc = tl.make_tensor_descriptor( + x_ptr, + shape=[SHAPE_0, SHAPE_1, SHAPE_2, SHAPE_3, SHAPE_4], + strides=[STRIDE_0, STRIDE_1, STRIDE_2, STRIDE_3, STRIDE_4], + block_shape=[BLOCK_0, BLOCK_1, BLOCK_2, BLOCK_3, BLOCK_4], + ) + out_desc = tl.make_tensor_descriptor( + out_ptr, + shape=[SHAPE_0, SHAPE_1, SHAPE_2, SHAPE_3, SHAPE_4], + strides=[STRIDE_0, STRIDE_1, STRIDE_2, STRIDE_3, STRIDE_4], + block_shape=[BLOCK_0, BLOCK_1, BLOCK_2, BLOCK_3, BLOCK_4], + ) + block = in_desc.load([o1, o2, o3, o4, o5]) + out_desc.store([o1, o2, o3, o4, o5], block) + + +@triton.jit +def triton_tensor_descriptor_function_2d( + out_ptr, + x_ptr, + M: tl.constexpr, + N: tl.constexpr, + M_BLOCK: tl.constexpr, + N_BLOCK: tl.constexpr, +): + in_desc = tl.make_tensor_descriptor( + x_ptr, + shape=[M, N], + strides=[N, 1], + block_shape=[M_BLOCK, N_BLOCK], + ) + out_desc = tl.make_tensor_descriptor( + out_ptr, + shape=[M, N], + strides=[N, 1], + block_shape=[M_BLOCK, N_BLOCK], + ) + moffset = tl.program_id(0) * M_BLOCK + noffset = tl.program_id(1) * N_BLOCK + block = tl.load_tensor_descriptor(in_desc, [moffset, noffset]) + tl.store_tensor_descriptor(out_desc, [moffset, noffset], block) + + +@pytest.mark.parametrize('dtype', full_dtype) +@pytest.mark.parametrize('shape', TestUtils.full_shape) +def test_tensor_descriptor_load_store_nd(dtype, shape): + """test tensor_descriptor load/store for nd tensor""" + + if dtype in temporarily_not_support_dtype: + pytest.skip(f"{dtype} not supported") + + inp = test_common.generate_numpy(shape, dtype) + inp = torch.from_numpy(inp).npu() + out = inp.new_empty(shape) + blocks = list(inp.size()) + strides = list(inp.stride()) + grid = (1, ) + dims = len(shape) + + # 如果最后一维小于16字节，则跳过 + itemsize = torch.tensor([], dtype=inp.dtype).element_size() + if blocks[-1] * itemsize < 16: + pytest.skip(f"last dimension must be at least 16 bytes, but got {blocks[-1] * itemsize} bytes") + + if dims == 2: + if inp.numel() * inp.element_size() > 8192: + triton_tensor_descriptor_2d[shape[0], 1, 1](out, inp, 1, shape[1], 1, shape[1]) + else: + triton_tensor_descriptor_2d[grid](out, inp, *shape, *blocks) + test_common.validate_cmp(dtype, inp, out) + elif dims == 3: + triton_tensor_descriptor_3d[grid](out, inp, *shape, *strides, *blocks) + test_common.validate_cmp(dtype, inp, out) + elif dims == 4: + triton_tensor_descriptor_4d[grid](out, inp, *shape, *strides, *blocks) + test_common.validate_cmp(dtype, inp, out) + elif dims == 5: + triton_tensor_descriptor_5d[grid](out, inp, *shape, *strides, *blocks) + test_common.validate_cmp(dtype, inp, out) + else: + pytest.skip(f"{dims}d not supported") + + +@pytest.mark.parametrize("dtype", test_common._uint_dtypes) +def test_tensor_descriptor_in_function(dtype): + """test tensor_descriptor load/store in function""" + + if dtype in temporarily_not_support_dtype: + pytest.skip(f"{dtype} not supported") + + M, N = 32, 128 + inp = test_common.generate_numpy((M, N), dtype) + inp = torch.from_numpy(inp).npu() + out = inp.new_empty((M, N)) + + M_BLOCK = 8 + N_BLOCK = 32 + grid_m = M // M_BLOCK + grid_n = N // N_BLOCK + + triton_tensor_descriptor_function_2d[(grid_m, grid_n)](out, inp, M, N, M_BLOCK, N_BLOCK) + test_common.validate_cmp(dtype, inp, out) diff --git a/third_party/ascend/unittest/generalization_cases/test_general_view.py b/third_party/ascend/unittest/generalization_cases/test_general_view.py new file mode 100644 index 0000000000..7f0f9b1532 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_general_view.py @@ -0,0 +1,159 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import logging +import torch +import pytest +import test_common +from test_common import TestUtils +import math + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.view(X, (ZB * YB * XB, )) + + oidx = tl.arange(0, XB * YB * ZB) + xoffs * YNUMEL * ZNUMEL + yoffs * ZNUMEL + zoffs + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def triton_view_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = tl.view(x_val, (SHAPE_0 * SHAPE_1 * SHAPE_2 * SHAPE_3 * SHAPE_4, )) + + pid0 = tl.program_id(0) + + flat_idx = tl.arange(0, BLOCK_0 * BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) + out_offsets = pid0 * BLOCK_0 * BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4 + flat_idx + out_masks = out_offsets < SHAPE_0 * SHAPE_1 * SHAPE_2 * SHAPE_3 * SHAPE_4 + tl.store(output_ptr + out_offsets, ret, mask=out_masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_view(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = torch.full(shape, 100, dtype=eval('torch.' + dtype)).npu() + y = torch.full(shape, 30, dtype=eval('torch.' + dtype)).npu() + new_shape = (x.numel(), ) + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = x.view(new_shape) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + if xnumel > 1: + grid = (XB, 1, 1) + XB = 1 + elif ynumel > 1: + grid = (1, YB, 1) + YB = 1 + else: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_view_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = torch.full(shape, 100, dtype=eval('torch.' + dtype)).npu() + + output = torch.randint(1, (x.numel(), ), dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = x.view(x.numel(), ) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_view_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_gt_op.py b/third_party/ascend/unittest/generalization_cases/test_gt_op.py new file mode 100644 index 0000000000..7079457b2e --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_gt_op.py @@ -0,0 +1,158 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import pytest +import triton +import triton.language as tl +import time +import torch +import torch_npu +import test_common +from test_common import TestUtils + + +@triton.jit +def triton_gt_3d(in_ptr0, in_ptr1, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0 > x1 + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_gt_2d(in_ptr0, in_ptr1, out_ptr0, M: tl.constexpr, N: tl.constexpr): + moffs = tl.program_id(0) * M + mblk_idx = tl.arange(0, M) + moffs + nblk_idx = tl.arange(0, N) + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0 > x1 + odx = mblk_idx[:, None] * N + nblk_idx[None, :] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_gt_1d(in_ptr0, in_ptr1, out_ptr0, L: tl.constexpr): + lblk_idx = tl.arange(0, L) + idx = lblk_idx[:] + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0 > x1 + odx = lblk_idx[:] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_gt_4d_5d(x_ptr, y_ptr, output_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = x_val > y_val + tl.store(output_ptr + offsets, ret, mask=masks) + + +typelist = ['bool', 'int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32'] + +dtype_mapping = { + 'int8': (torch.int8), + 'int16': (torch.int16), + 'int32': (torch.int32), + 'uint32': (torch.uint32), + 'int64': (torch.int64), + 'float16': (torch.float16), + 'float32': (torch.float32), + 'bfloat16': (torch.bfloat16), + 'bool': (torch.bool), +} + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1_2_3d) +@pytest.mark.parametrize('sigtype', typelist) +def test_gt(sigtype, shape): + dtype = dtype_mapping[sigtype] + x0 = test_common.generate_tensor(shape=shape, dtype=sigtype).npu() + x1 = test_common.generate_tensor(shape=shape, dtype=sigtype).npu() + # ncore, xblock, xblock_sub = 2, 32768, 1024 + y_ref = torch.where(torch.gt(x0, x1), torch.ones_like(x0), torch.zeros_like(x0)).to(eval('torch.' + sigtype)) + output = torch.zeros(shape, dtype=dtype).npu() + if len(shape) == 3: + triton_gt_3d[1, 1, 1](x0, x1, output, shape[0], shape[1], shape[2]) + if len(shape) == 2: + shape0 = shape[0] + shape1 = shape[1] + if x0.numel() * x0.element_size() >= 8192: + grid = (shape0, 1, 1) + shape0 = 1 + else: + grid = (1, 1, 1) + triton_gt_2d[grid](x0, x1, output, shape0, shape1) + if len(shape) == 1: + triton_gt_1d[1, 1, 1](x0, x1, output, shape[0]) + test_common.validate_cmp(sigtype, output, y_ref) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32']) +def test_gt_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.where(torch.gt(x, y), torch.ones_like(x), torch.zeros_like(x)).to(eval('torch.' + dtype)) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_gt_4d_5d[grid](x, y, output, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_invalid_fp8.py b/third_party/ascend/unittest/generalization_cases/test_invalid_fp8.py new file mode 100644 index 0000000000..7101136517 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_invalid_fp8.py @@ -0,0 +1,55 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import numpy as np +import torch +import pytest +import test_common + + +@triton.jit +def triton_test_fp8(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop1 in range(loops1): + x0_prime = offset + (loop1 * XBLOCK_SUB) + base1 + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.load(in_ptr1 + (x0), None) + tmp0 = tmp0.to(tl.float8e5) + tmp1 = tmp1.to(tl.float8e5) + tmp2 = tmp0 + tmp1 + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], +]) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "unexpected type fp8") +def test_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_test_fp8[ncore, 1, 1](x0, x1, y_cal, xblock, xblock_sub) diff --git a/third_party/ascend/unittest/generalization_cases/test_invert.py b/third_party/ascend/unittest/generalization_cases/test_invert.py new file mode 100644 index 0000000000..698cb8f11c --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_invert.py @@ -0,0 +1,166 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils +import math + + +def torch_invert(x0, ddtype): + if 'float' in str(ddtype): + x0 = x0.to(torch.int32) + y_ref = ~x0 + y_ref = y_ref.to(ddtype) + else: + y_ref = ~x0 + return y_ref + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = ~X + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_invert_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = ~x_val + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'bool']) +def test_case2(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.debug(f"output.dtype={output.dtype}") + + ddtype = eval('torch.' + dtype) + ans = torch_invert(x, ddtype) + + if len(shape) == 1: + fn_npu_[1, 1, shape[0]](output, x, y, z, 1, 1, 1, 1, 1, shape[0]) + elif len(shape) == 2: + if shape[0] > shape[1]: + fn_npu_[1, shape[0], 1](output, x, y, z, 1, 1, shape[1], 1, shape[0], shape[1]) + else: + fn_npu_[1, 1, shape[1]](output, x, y, z, 1, shape[0], 1, 1, shape[0], shape[1]) + elif len(shape) == 3: + if max(shape[0], shape[1], shape[2]) == shape[0]: + fn_npu_[shape[0], 1, 1](output, x, y, z, 1, shape[1], shape[2], shape[0], shape[1], shape[2]) + elif max(shape[0], shape[1], shape[2]) == shape[1]: + fn_npu_[1, shape[1], 1](output, x, y, z, shape[0], 1, shape[2], shape[0], shape[1], shape[2]) + else: + fn_npu_[1, 1, shape[2]](output, x, y, z, shape[0], shape[1], 1, shape[0], shape[1], shape[2]) + else: + fn_npu_[1, 1, 1](output, x, y, z, 1, 1, 1, 1, 1, 1) + + test_common.validate_cmp(dtype, ans, output) + + +invalid_types = [ + 'float16', + 'float32', + 'bfloat16', +] + + +@pytest.mark.parametrize("sigtype", invalid_types) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "unexpected type") +def test_invalid_types(sigtype): + N = 32 + x = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + y = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + z = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + output = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + + fn_npu_[1, 1, 1](output, x, y, z, 32, 1, 1, 32, 1, 1) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'bool']) +def test_invert_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_invert(x, eval('torch.' + dtype)) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_invert_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_le_op.py b/third_party/ascend/unittest/generalization_cases/test_le_op.py new file mode 100644 index 0000000000..d305395417 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_le_op.py @@ -0,0 +1,158 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import pytest +import triton +import triton.language as tl +import time +import torch +import torch_npu +import test_common +from test_common import TestUtils + + +@triton.jit +def triton_le_3d(in_ptr0, in_ptr1, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0 <= x1 + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_le_2d(in_ptr0, in_ptr1, out_ptr0, M: tl.constexpr, N: tl.constexpr): + moffs = tl.program_id(0) * M + mblk_idx = tl.arange(0, M) + moffs + nblk_idx = tl.arange(0, N) + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0 <= x1 + odx = mblk_idx[:, None] * N + nblk_idx[None, :] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_le_1d(in_ptr0, in_ptr1, out_ptr0, L: tl.constexpr): + lblk_idx = tl.arange(0, L) + idx = lblk_idx[:] + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0 <= x1 + odx = lblk_idx[:] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_le_4d_5d(x_ptr, y_ptr, output_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = x_val <= y_val + tl.store(output_ptr + offsets, ret, mask=masks) + + +typelist = ['bool', 'int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32'] + +dtype_mapping = { + 'int8': (torch.int8), + 'int16': (torch.int16), + 'int32': (torch.int32), + 'uint32': (torch.uint32), + 'int64': (torch.int64), + 'float16': (torch.float16), + 'float32': (torch.float32), + 'bfloat16': (torch.bfloat16), + 'bool': (torch.bool), +} + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1_2_3d) +@pytest.mark.parametrize('sigtype', typelist) +def test_le(sigtype, shape): + dtype = dtype_mapping[sigtype] + x0 = test_common.generate_tensor(shape=shape, dtype=sigtype).npu() + x1 = test_common.generate_tensor(shape=shape, dtype=sigtype).npu() + # ncore, xblock, xblock_sub = 2, 32768, 1024 + y_ref = torch.where(torch.le(x0, x1), torch.ones_like(x0), torch.zeros_like(x0)).to(eval('torch.' + sigtype)) + output = torch.zeros(shape, dtype=dtype).npu() + if len(shape) == 3: + triton_le_3d[1, 1, 1](x0, x1, output, shape[0], shape[1], shape[2]) + if len(shape) == 2: + shape0 = shape[0] + shape1 = shape[1] + if x0.numel() * x0.element_size() >= 8192: + grid = (shape0, 1, 1) + shape0 = 1 + else: + grid = (1, 1, 1) + triton_le_2d[grid](x0, x1, output, shape0, shape1) + if len(shape) == 1: + triton_le_1d[1, 1, 1](x0, x1, output, shape[0]) + test_common.validate_cmp(sigtype, output, y_ref) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32']) +def test_le_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.where(torch.le(x, y), torch.ones_like(x), torch.zeros_like(x)).to(eval('torch.' + dtype)) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_le_4d_5d[grid](x, y, output, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_load_store.py b/third_party/ascend/unittest/generalization_cases/test_load_store.py new file mode 100644 index 0000000000..82013d9f38 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_load_store.py @@ -0,0 +1,168 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common +from test_common import TestUtils +import logging + + +@triton.jit +def fn_npu_1d(output_ptr, x_ptr, YB: tl.constexpr): + idx = tl.arange(0, YB) + X = tl.load(x_ptr + idx) + tl.store(output_ptr + idx, X) + + +def torch_fn_npu_1d(x): + return x + + +@triton.jit +def fn_npu_2d(output_ptr, x_ptr, YB: tl.constexpr, ZB: tl.constexpr): + pid = tl.program_id(0) + y_idx = tl.arange(0, YB)[:, None] + pid * YB + z_idx = tl.arange(0, ZB)[None, :] + idx = y_idx * ZB + z_idx + + X = tl.load(x_ptr + idx) + + tl.store(output_ptr + idx, X) + + +def torch_fn_npu_2d(x): + return x + + +@triton.jit +def fn_npu_3d(output_ptr, x_ptr, YB: tl.constexpr, ZB: tl.constexpr, KB: tl.constexpr): + y = tl.arange(0, YB)[:, None, None] + z = tl.arange(0, ZB)[None, :, None] + k = tl.arange(0, KB)[None, None, :] + + idx = y * ZB * KB + z * KB + k + + X = tl.load(x_ptr + idx) + + tl.store(output_ptr + idx, X) + + +def torch_fn_npu_3d(x): + return x + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1_2_3d) +@pytest.mark.parametrize('dtype', TestUtils.dtype_list) +def test_npu(shape, dtype): + logging.debug(f'dtype:{dtype} shape:{shape}') + data_type = eval('torch.' + dtype) + x = torch.randint(low=0, high=2, size=shape, dtype=data_type).npu() + triton_res = torch.empty(shape, dtype=data_type).npu() + torch_res = x + if len(shape) == 1: + torch_res = torch_fn_npu_1d(x) + fn_npu_1d[1, 1, 1](triton_res, x, shape[0]) + # uint32 转成 float32算精度，因为torch_npu不支持uint32类型张量的slice + torch_res = torch_res if dtype != 'uint32' else torch_res.to(torch.float32) + triton_res = triton_res if dtype != 'uint32' else triton_res.to(torch.float32) + cmp_type = dtype if dtype != 'uint32' else 'float32' + test_common.validate_cmp(cmp_type, triton_res[:2 * shape[0] // 3], torch_res[:2 * shape[0] // 3]) + elif len(shape) == 2: + torch_res = torch_fn_npu_2d(x) + fn_npu_2d[shape[0], 1, 1](triton_res, x, 1, shape[1]) + torch_res = torch_res if dtype != 'uint32' else torch_res.to(torch.float32) + triton_res = triton_res if dtype != 'uint32' else triton_res.to(torch.float32) + cmp_type = dtype if dtype != 'uint32' else 'float32' + test_common.validate_cmp(cmp_type, triton_res[:2 * shape[0] // 3, :2 * shape[1] // 3], + torch_res[:2 * shape[0] // 3, :2 * shape[1] // 3]) + elif len(shape) == 3: + torch_res = torch_fn_npu_3d(x) + fn_npu_3d[1, 1, 1](triton_res, x, shape[0], shape[1], shape[2]) + torch_res = torch_res if dtype != 'uint32' else torch_res.to(torch.float32) + triton_res = triton_res if dtype != 'uint32' else triton_res.to(torch.float32) + cmp_type = dtype if dtype != 'uint32' else 'float32' + test_common.validate_cmp(cmp_type, triton_res[:2 * shape[0] // 3, :2 * shape[1] // 3, :2 * shape[2] // 3], + torch_res[:2 * shape[0] // 3, :2 * shape[1] // 3, :2 * shape[2] // 3]) + + +# require: all data (4d and 5d) can be placed into but without ub overflow +@triton.jit +def triton_load_store_multi_d(in_ptr0, out_ptr0, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, + SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, STRIDE_1: tl.constexpr, + STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + tmp_in = tl.load(in_ptr0 + offsets, masks) + tmp_out = tmp_in + tl.store(out_ptr0 + offsets, tmp_out, masks) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('param_list', [ + ['float32', (8, 4, 16, 16)], + ['float16', (8, 4, 16, 16)], + ['int8', (8, 4, 16, 16)], + ['float32', (8, 8, 4, 4)], + ['float16', (8, 8, 4, 4)], + ['int8', (8, 8, 4, 4)], + ['float32', (3, 8, 2, 16, 16)], + ['float16', (3, 8, 2, 16, 16)], + ['int8', (9, 8, 8, 16, 16)], + ['float32', (11, 8, 8, 4, 4)], + ['float16', (11, 8, 8, 4, 4)], + ['int8', (11, 8, 8, 4, 4)], +]) +def test_load_store_4d_5d(param_list): + # 生成数据 + dtype, shape = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + y_expect = x0 + y_actual = test_common.generate_tensor(shape, dtype).npu() + # triton结果 + blocks = list(x0.size()) + shapes = list(x0.stride()) + while len(blocks) < 5: + blocks.append(1) + shapes.append(1) + triton_load_store_multi_d[(1, )](x0, y_actual, *blocks, *blocks, *shapes) + # 比较结果 + test_common.validate_cmp(dtype, y_actual, y_expect) diff --git a/third_party/ascend/unittest/generalization_cases/test_log1p.py b/third_party/ascend/unittest/generalization_cases/test_log1p.py new file mode 100644 index 0000000000..fa37cbd298 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_log1p.py @@ -0,0 +1,62 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils +import math +import triton.language.extra.ascend.libdevice as libdevice + + +def torch_pointwise(x0): + res = torch.log1p(x0) + return res + + +@triton.jit +def triton_log1p(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop1 in range(loops1): + x0_prime = offset + (loop1 * XBLOCK_SUB) + base1 + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp2 = libdevice.log1p(tmp0) + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', ['float32', 'float16']) +def test_case(dtype, shape): + x0 = test_common.generate_tensor(shape, dtype).npu() + + numel = x0.numel() + ncore = 1 if numel <= 32 else 32 + xblock = math.ceil(numel / ncore) + xblock_sub = numel if numel <= ncore else math.ceil(numel / ncore) + + y_ref = torch_pointwise(x0) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_log1p[ncore, 1, 1](x0, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/generalization_cases/test_logical_and_op.py b/third_party/ascend/unittest/generalization_cases/test_logical_and_op.py new file mode 100644 index 0000000000..b89ca7f08f --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_logical_and_op.py @@ -0,0 +1,157 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import time +import torch +import torch_npu +import test_common +from test_common import TestUtils, generate_tensor +import logging + + +@triton.jit +def triton_logical_and_1d(in_ptr0, in_ptr1, out_ptr0, L: tl.constexpr): + lblk_idx = tl.arange(0, L) + idx = lblk_idx + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0.logical_and(x1) + odx = lblk_idx + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_logical_and_2d(in_ptr0, in_ptr1, out_ptr0, L: tl.constexpr, M: tl.constexpr): + loffs = tl.program_id(0) * L + lblk_idx = tl.arange(0, L) + loffs + mblk_idx = tl.arange(0, M) + idx = lblk_idx[:, None] * M + mblk_idx[None, :] + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0.logical_and(x1) + odx = lblk_idx[:, None] * M + mblk_idx[None, :] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_logical_and_3d(in_ptr0, in_ptr1, out_ptr0, XB, YB, ZB, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + tl.program_id(0) * XB + mblk_idx = tl.arange(0, M) + tl.program_id(1) * YB + nblk_idx = tl.arange(0, N) + tl.program_id(2) * ZB + idx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0.logical_and(x1) + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_logical_and_4d_5d(x_ptr, y_ptr, output_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, + BLOCK_2: tl.constexpr, BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, + STRIDE_0: tl.constexpr, STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, + STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = x_val.logical_and(y_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +support_typelist = [ + 'bool', +] + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1_2_3d) +@pytest.mark.parametrize('sigtype', support_typelist) +def test_logical_and(shape, sigtype): + logging.debug(f"dtype:{sigtype} shape:{shape}") + dtype = eval('torch.' + sigtype) + x0 = generate_tensor(shape=shape, dtype=sigtype).npu() + x1 = generate_tensor(shape=shape, dtype=sigtype).npu() + # ncore, xblock, xblock_sub = 2, 32768, 1024 + y_ref = torch.logical_and(x0, x1) + output = torch.zeros(shape, dtype=dtype).npu() + if len(shape) == 1: + triton_logical_and_1d[1, 1, 1](x0, x1, output, shape[0]) + elif len(shape) == 2: + shape0 = shape[0] + shape1 = shape[1] + if x0.numel() * x0.element_size() >= 8192: + grid = (shape0, 1, 1) + shape0 = 1 + else: + grid = (1, 1, 1) + triton_logical_and_2d[grid](x0, x1, output, shape0, shape1) + elif len(shape) == 3: + mx = max(shape[0], shape[1], shape[2]) + if mx == shape[0]: + triton_logical_and_3d[shape[0], 1, 1](x0, x1, output, 1, shape[1], shape[2], shape[0], shape[1], shape[2]) + elif mx == shape[1]: + triton_logical_and_3d[1, shape[1], 1](x0, x1, output, shape[0], 1, shape[2], shape[0], shape[1], shape[2]) + else: + triton_logical_and_3d[1, 1, shape[2]](x0, x1, output, shape[0], shape[1], 1, shape[0], shape[1], shape[2]) + + test_common.validate_cmp(sigtype, output, y_ref) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['bool']) +def test_logical_and_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.logical_and(x, y) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_logical_and_4d_5d[grid](x, y, output, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_logical_or_op.py b/third_party/ascend/unittest/generalization_cases/test_logical_or_op.py new file mode 100644 index 0000000000..f470de056f --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_logical_or_op.py @@ -0,0 +1,156 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import time +import torch +import torch_npu +import test_common +from test_common import TestUtils, generate_tensor +import logging + + +@triton.jit +def triton_logical_or_1d(in_ptr0, in_ptr1, out_ptr0, L: tl.constexpr): + lblk_idx = tl.arange(0, L) + idx = lblk_idx + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0.logical_or(x1) + odx = lblk_idx + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_logical_or_2d(in_ptr0, in_ptr1, out_ptr0, L: tl.constexpr, M: tl.constexpr): + pid = tl.program_id(0) + lblk_idx = tl.arange(0, L) + pid * L + mblk_idx = tl.arange(0, M) + idx = lblk_idx[:, None] * M + mblk_idx[None, :] + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0.logical_or(x1) + odx = lblk_idx[:, None] * M + mblk_idx[None, :] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_logical_or_3d(in_ptr0, in_ptr1, out_ptr0, XB, YB, ZB, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + tl.program_id(0) * XB + mblk_idx = tl.arange(0, M) + tl.program_id(1) * YB + nblk_idx = tl.arange(0, N) + tl.program_id(2) * ZB + idx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0.logical_or(x1) + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_logical_or_4d_5d(x_ptr, y_ptr, output_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, + BLOCK_2: tl.constexpr, BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, + STRIDE_0: tl.constexpr, STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, + STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = x_val.logical_or(y_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +support_typelist = [ + 'bool', +] + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('sigtype', support_typelist) +def test_logical_or(shape, sigtype): + logging.debug(f"dtype:{sigtype} shape:{shape}") + dtype = eval('torch.' + sigtype) + x0 = generate_tensor(shape=shape, dtype=sigtype).npu() + x1 = generate_tensor(shape=shape, dtype=sigtype).npu() + # ncore, xblock, xblock_sub = 2, 32768, 1024 + y_ref = torch.logical_or(x0, x1) + output = torch.zeros(shape, dtype=dtype).npu() + if len(shape) == 1: + triton_logical_or_1d[1, 1, 1](x0, x1, output, shape[0]) + elif len(shape) == 2: + shape0 = shape[0] + shape1 = shape[1] + if x0.numel() * x0.element_size() >= 8192: + grid = (shape0, 1, 1) + shape0 = 1 + else: + grid = (1, 1, 1) + triton_logical_or_2d[grid](x0, x1, output, shape0, shape1) + elif len(shape) == 3: + mx = max(shape[0], shape[1], shape[2]) + if mx == shape[0]: + triton_logical_or_3d[shape[0], 1, 1](x0, x1, output, 1, shape[1], shape[2], shape[0], shape[1], shape[2]) + elif mx == shape[1]: + triton_logical_or_3d[1, shape[1], 1](x0, x1, output, shape[0], 1, shape[2], shape[0], shape[1], shape[2]) + else: + triton_logical_or_3d[1, 1, shape[2]](x0, x1, output, shape[0], shape[1], 1, shape[0], shape[1], shape[2]) + test_common.validate_cmp(sigtype, output, y_ref) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['bool']) +def test_logical_or_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.logical_or(x, y) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_logical_or_4d_5d[grid](x, y, output, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_lshift_op.py b/third_party/ascend/unittest/generalization_cases/test_lshift_op.py new file mode 100644 index 0000000000..b70020aca5 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_lshift_op.py @@ -0,0 +1,183 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import pytest +import triton +import triton.language as tl +import time +import torch +import torch_npu +import test_common +from test_common import TestUtils + + +@triton.jit +def triton_lshift_1d(in_ptr0, out_ptr0, L: tl.constexpr): + lblk_idx = tl.arange(0, L) + idx = lblk_idx[:] + x0 = tl.load(in_ptr0 + idx) + ret = x0 << 2 + odx = lblk_idx[:] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_lshift_2d(in_ptr0, out_ptr0, M: tl.constexpr, N: tl.constexpr): + moffs = tl.program_id(0) * M + mblk_idx = tl.arange(0, M) + moffs + nblk_idx = tl.arange(0, N) + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + x0 = tl.load(in_ptr0 + idx) + ret = x0 << 2 + odx = mblk_idx[:, None] * N + nblk_idx[None, :] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_lshift_3d(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + loffs = tl.program_id(0) * L + lblk_idx = tl.arange(0, L) + loffs + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x0 = tl.load(in_ptr0 + idx) + ret = x0 << 2 + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_lshift_4d_5d(x_ptr, output_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = x_val << 2 + tl.store(output_ptr + offsets, ret, mask=masks) + + +dtype_mapping = { + 'int8': (torch.int8), + 'int16': (torch.int16), + 'int32': (torch.int32), + 'uint32': (torch.uint32), + 'int64': (torch.int64), + 'float16': (torch.float16), + 'float32': (torch.float32), + 'bfloat16': (torch.bfloat16), + 'bool': (torch.bool), +} + +typelist = [ + 'int8', + 'int16', + 'int32', + 'int64', +] + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1_2_3d) +@pytest.mark.parametrize('sigtype', typelist) +def test_lshift(sigtype, shape): + dtype = dtype_mapping[sigtype] + x0 = test_common.generate_tensor(shape=shape, dtype=sigtype).npu() + # ncore, xblock, xblock_sub = 2, 32768, 1024 + y_ref = x0 << 2 + output = torch.zeros(shape, dtype=dtype).npu() + if len(shape) == 3: + shape0 = shape[0] + shape1 = shape[1] + shape2 = shape[2] + if x0.numel() * x0.element_size() >= 1024: + grid = (shape0, 1, 1) + shape0 = 1 + else: + grid = (1, 1, 1) + triton_lshift_3d[grid](x0, output, shape0, shape1, shape2) + if len(shape) == 2: + shape0 = shape[0] + shape1 = shape[1] + if x0.numel() * x0.element_size() >= 1024: + grid = (shape0, 1, 1) + shape0 = 1 + else: + grid = (1, 1, 1) + triton_lshift_2d[grid](x0, output, shape0, shape1) + if len(shape) == 1: + triton_lshift_1d[1, 1, 1](x0, output, shape[0]) + test_common.validate_cmp(sigtype, output, y_ref) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64']) +def test_lshift_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + + output = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = x << 2 + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_lshift_4d_5d[grid](x, output, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) + + +invalid_types = [ + 'float16', + 'float32', + 'bfloat16', +] + + +@pytest.mark.parametrize("sigtype", invalid_types) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "unexpected type") +def test_invalid_types(sigtype): + N = 32 + x = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + output = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + + triton_lshift_1d[1, 1, 1](x, output, N) diff --git a/third_party/ascend/unittest/generalization_cases/test_lt_op.py b/third_party/ascend/unittest/generalization_cases/test_lt_op.py new file mode 100644 index 0000000000..8f013d7c9f --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_lt_op.py @@ -0,0 +1,158 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import time +import torch +import torch_npu +import test_common +from test_common import TestUtils +import logging + + +@triton.jit +def triton_lt_3d(in_ptr0, in_ptr1, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0 < x1 + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_lt_2d(in_ptr0, in_ptr1, out_ptr0, M: tl.constexpr, N: tl.constexpr): + moffs = tl.program_id(0) * M + mblk_idx = tl.arange(0, M) + moffs + nblk_idx = tl.arange(0, N) + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0 < x1 + odx = mblk_idx[:, None] * N + nblk_idx[None, :] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_lt_1d(in_ptr0, in_ptr1, out_ptr0, L: tl.constexpr): + lblk_idx = tl.arange(0, L) + idx = lblk_idx[:] + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0 < x1 + odx = lblk_idx[:] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_lt_4d_5d(x_ptr, y_ptr, output_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = x_val < y_val + tl.store(output_ptr + offsets, ret, mask=masks) + + +typelist = ['bool', 'int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32'] + +dtype_mapping = { + 'int8': (torch.int8), + 'int16': (torch.int16), + 'int32': (torch.int32), + 'uint32': (torch.uint32), + 'int64': (torch.int64), + 'float16': (torch.float16), + 'float32': (torch.float32), + 'bfloat16': (torch.bfloat16), + 'bool': (torch.bool), +} + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1_2_3d) +@pytest.mark.parametrize('sigtype', typelist) +def test_lt(sigtype, shape): + dtype = dtype_mapping[sigtype] + x0 = test_common.generate_tensor(shape=shape, dtype=sigtype).npu() + x1 = test_common.generate_tensor(shape=shape, dtype=sigtype).npu() + # ncore, xblock, xblock_sub = 2, 32768, 1024 + y_ref = torch.where(torch.lt(x0, x1), torch.ones_like(x0), torch.zeros_like(x0)).to(eval('torch.' + sigtype)) + output = torch.zeros(shape, dtype=dtype).npu() + if len(shape) == 3: + triton_lt_3d[1, 1, 1](x0, x1, output, shape[0], shape[1], shape[2]) + if len(shape) == 2: + shape0 = shape[0] + shape1 = shape[1] + if x0.numel() * x0.element_size() >= 8192: + grid = (shape0, 1, 1) + shape0 = 1 + else: + grid = (1, 1, 1) + triton_lt_2d[grid](x0, x1, output, shape0, shape1) + if len(shape) == 1: + triton_lt_1d[1, 1, 1](x0, x1, output, shape[0]) + test_common.validate_cmp(sigtype, output, y_ref) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32']) +def test_lt_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.where(torch.lt(x, y), torch.ones_like(x), torch.zeros_like(x)).to(eval('torch.' + dtype)) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_lt_4d_5d[grid](x, y, output, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_make_blkptr_matmul.py b/third_party/ascend/unittest/generalization_cases/test_make_blkptr_matmul.py new file mode 100644 index 0000000000..0fdc244e79 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_make_blkptr_matmul.py @@ -0,0 +1,74 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +import test_common +from test_common import TestUtils, avoid_not_support, get_dtype_size + + +@triton.jit +def matmul_kernel( + a_ptr, + b_ptr, + c_ptr, + M: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, + acc_dtype: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + matxa_ptr_in = tl.make_block_ptr(a_ptr, (M, K), (K, 1), (0, 0), (M, K), order=(1, 0)) + matxb_ptr_in = tl.make_block_ptr(b_ptr, (K, N), (N, 1), (0, 0), (K, N), order=(1, 0)) + matxc_ptr_in = tl.make_block_ptr(c_ptr, (M, N), (N, 1), (0, 0), (M, N), order=(1, 0)) + + accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=acc_dtype) + a = tl.load(matxa_ptr_in) + b = tl.load(matxb_ptr_in) + accumulator = tl.dot(a, b, accumulator, out_dtype=acc_dtype) + c = accumulator.to(c_ptr.dtype.element_ty) + tl.store(matxc_ptr_in, c) + + +@avoid_not_support('matmul') +@pytest.mark.parametrize('shape', [(16, 32)]) +@pytest.mark.parametrize('dtype', ['float32']) +def test_matmul(shape, dtype): + M, N, K = shape[0], shape[0], shape[1] + + BLOCK_M, BLOCK_N, BLOCK_K = M, N, K + a = test_common.generate_tensor((M, K), dtype) + b = test_common.generate_tensor((K, N), dtype) + + triton_res = torch.zeros((M, N), dtype=eval('torch.' + dtype)).npu() + accumulator_type = tl.float32 + + matmul_kernel[ + 1, + ](a.npu(), b.npu(), triton_res, M, N, K, accumulator_type, BLOCK_M, BLOCK_N, BLOCK_K, enable_nd2nz_on_vector=False) + + print("PASSED") diff --git a/third_party/ascend/unittest/generalization_cases/test_make_block_ptr.py b/third_party/ascend/unittest/generalization_cases/test_make_block_ptr.py new file mode 100644 index 0000000000..4c95d5623b --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_make_block_ptr.py @@ -0,0 +1,212 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common +from test_common import TestUtils + + +@triton.jit +def fn_npu_1d(output_ptr, x_ptr, y_ptr, z_ptr, output_ptr1, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + block_ptr_in = tl.make_block_ptr( + base=x_ptr, + shape=(XB, ), + strides=(1, ), + offsets=(0, ), + block_shape=(XB, ), + order=(0, ), + ) + X = tl.load(block_ptr_in) + + block_ptr_out = tl.make_block_ptr( + base=output_ptr, + shape=(XB, ), + strides=(1, ), + offsets=(0, ), + block_shape=(XB, ), + order=(0, ), + ) + tl.store(block_ptr_out, X) + + +@triton.jit +def fn_npu_2d(output_ptr, x_ptr, y_ptr, z_ptr, output_ptr1, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xoffset = tl.program_id(0) + block_ptr_in = tl.make_block_ptr( + base=x_ptr, + shape=(XB, YB), + strides=(YB, 1), + offsets=(xoffset, 0), + block_shape=(XB, YB), + order=(1, 0), + ) + X = tl.load(block_ptr_in) + + block_ptr_out = tl.make_block_ptr( + base=output_ptr, + shape=(XB, YB), + strides=(YB, 1), + offsets=(xoffset, 0), + block_shape=(XB, YB), + order=(1, 0), + ) + tl.store(block_ptr_out, X) + + +@triton.jit +def fn_npu_3d(output_ptr, x_ptr, y_ptr, z_ptr, output_ptr1, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + block_ptr_in = tl.make_block_ptr( + base=x_ptr, + shape=(XB, YB, ZB), + strides=(YB * ZB, ZB, 1), + offsets=(0, 0, 0), + block_shape=(XB, YB, ZB), + order=(2, 1, 0), + ) + X = tl.load(block_ptr_in) + + block_ptr_out = tl.make_block_ptr( + base=output_ptr, + shape=(XB, YB, ZB), + strides=(YB * ZB, ZB, 1), + offsets=(0, 0, 0), + block_shape=(XB, YB, ZB), + order=(2, 1, 0), + ) + tl.store(block_ptr_out, X) + + +@triton.jit +def triton_make_block_ptr_4d( + output_ptr, + x_ptr, + BLOCK_0: tl.constexpr, + BLOCK_1: tl.constexpr, + BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, + SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, + SHAPE_3: tl.constexpr, + STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, + STRIDE_2: tl.constexpr, + STRIDE_3: tl.constexpr, +): + block_ptr_in = tl.make_block_ptr( + base=x_ptr, + shape=(SHAPE_0, SHAPE_1, SHAPE_2, SHAPE_3), + strides=(STRIDE_0, STRIDE_1, STRIDE_2, STRIDE_3), + offsets=(0, 0, 0, 0), + block_shape=(BLOCK_0, BLOCK_1, BLOCK_2, BLOCK_3), + order=(3, 2, 1, 0), + ) + x = tl.load(block_ptr_in) + + block_ptr_out = tl.make_block_ptr( + base=output_ptr, + shape=(SHAPE_0, SHAPE_1, SHAPE_2, SHAPE_3), + strides=(STRIDE_0, STRIDE_1, STRIDE_2, STRIDE_3), + offsets=(0, 0, 0, 0), + block_shape=(BLOCK_0, BLOCK_1, BLOCK_2, BLOCK_3), + order=(3, 2, 1, 0), + ) + tl.store(block_ptr_out, x) + + +@triton.jit +def triton_make_block_ptr_5d( + output_ptr, + x_ptr, + BLOCK_0: tl.constexpr, + BLOCK_1: tl.constexpr, + BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, + BLOCK_4: tl.constexpr, + SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, + SHAPE_3: tl.constexpr, + SHAPE_4: tl.constexpr, + STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, + STRIDE_2: tl.constexpr, + STRIDE_3: tl.constexpr, + STRIDE_4: tl.constexpr, +): + block_ptr_in = tl.make_block_ptr( + base=x_ptr, + shape=(SHAPE_0, SHAPE_1, SHAPE_2, SHAPE_3, SHAPE_4), + strides=(STRIDE_0, STRIDE_1, STRIDE_2, STRIDE_3, STRIDE_4), + offsets=(0, 0, 0, 0, 0), + block_shape=(BLOCK_0, BLOCK_1, BLOCK_2, BLOCK_3, BLOCK_4), + order=(4, 3, 2, 1, 0), + ) + x = tl.load(block_ptr_in) + + block_ptr_out = tl.make_block_ptr( + base=output_ptr, + shape=(SHAPE_0, SHAPE_1, SHAPE_2, SHAPE_3, SHAPE_4), + strides=(STRIDE_0, STRIDE_1, STRIDE_2, STRIDE_3, STRIDE_4), + offsets=(0, 0, 0, 0, 0), + block_shape=(BLOCK_0, BLOCK_1, BLOCK_2, BLOCK_3, BLOCK_4), + order=(4, 3, 2, 1, 0), + ) + tl.store(block_ptr_out, x) + + +temporarily_not_support_dtype = ['bool'] + + +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('shape', TestUtils.full_shape) +def test_npu(dtype, shape): + if dtype in temporarily_not_support_dtype: + return + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + + a = x + blocks = list(x.size()) + strides = list(x.stride()) + grid = (1, ) + if len(shape) == 5: + triton_make_block_ptr_5d[grid](output, x, *blocks, *blocks, *strides) + elif len(shape) == 4: + triton_make_block_ptr_4d[grid](output, x, *blocks, *blocks, *strides) + elif len(shape) == 3: + fn_npu_3d[1, 1, 1](output, x, y, z, output1, XB=shape[0], YB=shape[1], ZB=shape[2]) + elif len(shape) == 2: + if x.numel() * x.element_size() > 8192: + fn_npu_2d[shape[0], 1, 1](output, x, y, z, output1, XB=1, YB=shape[1], ZB=1) + else: + fn_npu_2d[1, 1, 1](output, x, y, z, output1, XB=shape[0], YB=shape[1], ZB=1) + else: + fn_npu_1d[1, 1, 1](output, x, y, z, output1, XB=shape[0], YB=1, ZB=1) + torch.testing.assert_close(output, a) diff --git a/third_party/ascend/unittest/generalization_cases/test_matmul.py b/third_party/ascend/unittest/generalization_cases/test_matmul.py new file mode 100644 index 0000000000..edeca4f170 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_matmul.py @@ -0,0 +1,186 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +import acc_util +import test_common +from test_common import TestUtils, avoid_not_support, get_dtype_size + + +@triton.jit +def matmul_kernel( + a_ptr, + b_ptr, + c_ptr, + M: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, + acc_dtype: tl.constexpr, + stride_am: tl.constexpr, + stride_ak: tl.constexpr, + stride_bk: tl.constexpr, + stride_bn: tl.constexpr, + stride_cm: tl.constexpr, + stride_cn: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + pid = tl.program_id(axis=0) + num_pid_n = tl.cdiv(N, BLOCK_N) + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + + offs_am = (pid_m * BLOCK_M + tl.arange(0, BLOCK_M)) + offs_bn = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)) + offs_k = tl.arange(0, BLOCK_K) + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=acc_dtype) + for k in range(0, tl.cdiv(K, BLOCK_K)): + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0) + accumulator = tl.dot(a, b, accumulator, out_dtype=acc_dtype) + a_ptrs += BLOCK_K * stride_ak + b_ptrs += BLOCK_K * stride_bk + c = accumulator.to(c_ptr.dtype.element_ty) + + offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +@avoid_not_support('matmul') +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', TestUtils.dtype_list) +def test_matmul(shape, dtype): + M, N, K = shape[0], shape[0], shape[1] + # 32byte/Dtype_bytes + kalign = 32 // get_dtype_size(dtype) + BLOCK_M, BLOCK_N, BLOCK_K = min(max(M, 16), 32), min(max(N, 16), 32), min(max(K, kalign), 32) + a = test_common.generate_tensor((M, K), dtype) + b = test_common.generate_tensor((K, N), dtype) + + if dtype == "int8": + triton_res = torch.zeros((M, N), dtype=torch.int32).npu() + accumulator_type = tl.int32 + else: + triton_res = torch.zeros((M, N), dtype=eval('torch.' + dtype)).npu() + accumulator_type = tl.float32 + grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N), ) + + matmul_kernel[grid](a.npu(), b.npu(), triton_res, M, N, K, accumulator_type, a.stride(0), a.stride(1), b.stride(0), + b.stride(1), triton_res.stride(0), triton_res.stride(1), BLOCK_M, BLOCK_N, BLOCK_K) + + a_gold = a.to(torch.float32) + b_gold = b.to(torch.float32) + cpu_res = torch.mm(a_gold, b_gold) + + if dtype == "int8": + # torch_npu do not support int8 matmul + a_npu = a.npu().to(torch.float32) + b_npu = b.npu().to(torch.float32) + torch_res = torch.mm(a_npu, b_npu) + triton_res = triton_res.to(torch.float32) + else: + a_npu = a.npu() + b_npu = b.npu() + torch_res = torch.mm(a_npu, b_npu) + + try: + print("starting compare of cpu vs triton:") + acc_util.assert_close(cpu_res, triton_res) + except Exception as e: + print(e) + print("starting compare of cpu vs triton vs torch_npu:") + acc_util.benchmark_compare_close(cpu_res, triton_res, torch_res) + print("PASSED") + + +@avoid_not_support('matmul') +@pytest.mark.parametrize('batch', TestUtils.batch) +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', TestUtils.dtype_list) +def test_batch_matmul(shape, dtype, batch): + M, N, K = shape[0], shape[0], shape[1] + # 32byte/Dtype_bytes + kalign = 32 // get_dtype_size(dtype) + BLOCK_M, BLOCK_N, BLOCK_K = min(max(M, 16), 32), min(max(N, 16), 32), min(max(K, kalign), 32) + + aa = test_common.generate_tensor((batch, M, K), dtype) + bb = test_common.generate_tensor((batch, K, N), dtype) + + if dtype == "int8": + final_triton_res = torch.zeros((batch, M, N), dtype=torch.int32).npu() + accumulator_type = tl.int32 + else: + final_triton_res = torch.zeros((batch, M, N), dtype=eval('torch.' + dtype)).npu() + accumulator_type = tl.float32 + + for i in range(0, batch): + if dtype == "int8": + triton_res = torch.zeros((M, N), dtype=torch.int32).npu() + else: + triton_res = torch.zeros((M, N), dtype=eval('torch.' + dtype)).npu() + grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N), ) + a = aa[i] + b = bb[i] + matmul_kernel[grid](a.npu(), b.npu(), triton_res, M, N, K, accumulator_type, a.stride(0), a.stride(1), + b.stride(0), b.stride(1), triton_res.stride(0), triton_res.stride(1), BLOCK_M, BLOCK_N, + BLOCK_K) + final_triton_res[i] = triton_res + + a_gold = aa.to(torch.float32) + b_gold = bb.to(torch.float32) + cpu_res = torch.bmm(a_gold, b_gold) + + if dtype == "int8": + a_npu = aa.npu().to(torch.float32) + b_npu = bb.npu().to(torch.float32) + final_triton_res = final_triton_res.to(torch.float32) + else: + a_npu = aa.npu() + b_npu = bb.npu() + torch_res = torch.bmm(a_npu, b_npu) + + try: + print("starting compare of cpu vs triton:") + acc_util.assert_close(cpu_res, final_triton_res) + except Exception as e: + print(e) + print("starting compare of cpu vs triton vs torch_npu:") + acc_util.benchmark_compare_close(cpu_res, final_triton_res, torch_res) + print("PASSED") + + +if __name__ == "__main__": + test_matmul((16, 32), 'float32') + test_matmul((16, 32), 'int8') + test_batch_matmul(2, (16, 32), 'float32') + test_batch_matmul(2, (16, 32), 'int8') diff --git a/third_party/ascend/unittest/generalization_cases/test_matmul_mindspore.py b/third_party/ascend/unittest/generalization_cases/test_matmul_mindspore.py new file mode 100644 index 0000000000..dc43a2ab0a --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_matmul_mindspore.py @@ -0,0 +1,189 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import os +import triton +import mindspore +import pytest +import numpy as np +import triton.language as tl + +pytestmark = pytest.mark.backend("mindspore") +activation = "leaky_relu_custom" + + +def get_autotune_config(): + return [ + triton.Config({"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 128}), + triton.Config({"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64}), + ] + + +@triton.autotune( + configs=get_autotune_config(), + key=["M", "N", "K"], +) +@triton.jit +def matmul_kernel( + # Pointers to matrices + a_ptr, + b_ptr, + c_ptr, + # Matrix dimensions + M, + N, + K, + # The stride variables represent how much to increase the ptr by when moving by 1 + # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr` + # by to get the element one row down (A has M rows). + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + ACTIVATION: tl.constexpr, +): + """Kernel for computing the matmul C = A x B. + A has shape (M, K), B has shape (K, N) and C has shape (M, N) + """ + GROUP_SIZE_M: tl.constexpr = 1 + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + # See above `L2 Cache Optimizations` section for details. + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # ---------------------------------------------------------- + # Create pointers for the first blocks of A and B. + # We will advance this pointer as we move in the K direction + # and accumulate + # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers + # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers + # See above `Pointer Arithmetic` section for details + offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs_base = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs_base = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + msk_m = offs_am < M + msk_n = offs_bn < N + # ----------------------------------------------------------- + # Iterate to compute a block of the C matrix. + # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block + # of fp32 values for higher accuracy. + # `accumulator` will be converted back to fp16 after the loop. + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + # Load the next block of A and B, generate a mask by checking the K dimension. + # If it is out of bounds, set it to 0. + a_ptrs = a_ptrs_base + k * BLOCK_SIZE_K * stride_ak + b_ptrs = b_ptrs_base + k * BLOCK_SIZE_K * stride_bk + a = tl.load( + a_ptrs, + mask=msk_m[:, None] and (offs_k[None, :] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + b = tl.load( + b_ptrs, + mask=msk_n[None, :] and (offs_k[:, None] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + # We accumulate along the K dimension. + accumulator = tl.dot(a, b, accumulator) + # You can fuse arbitrary activation functions here + # while the accumulator is still in FP32! + # Original vector operations + if ACTIVATION == "leaky_relu_custom": + accumulator = leaky_relu_custom(accumulator) + c = accumulator.to(tl.float16) + # ----------------------------------------------------------- + # Write back the block of the output matrix C with masks. + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +# We can fuse `leaky_relu_custom` by providing it as an `ACTIVATION` meta-parameter in `matmul_kernel`. +@triton.jit +def leaky_relu_custom(x): + return tl.where(x >= 0, x, 0.01 * x) + 1.0 + + +def mindspore_matmul(a, b, activation_func=""): + c = mindspore.mint.matmul(a, b) + if activation_func == "leaky_relu_custom": + c = mindspore.mint.where(c >= 0, c, 0.01 * c) + 1.0 + return c + + +def matmul(a, b, activation_func=""): + # Check constraints. + assert a.shape[1] == b.shape[0], "Incompatible dimensions" + assert a.is_contiguous(), "Matrix A must be contiguous" + M, K = a.shape + K, N = b.shape + # Allocates output. + c = mindspore.mint.empty((M, N), dtype=mindspore.float16) + # 1D launch kernel where each block gets its own program. + grid = lambda META: (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), ) + matmul_kernel[grid]( + a, + b, + c, + M, + N, + K, + a.stride(0), + a.stride(1), + b.stride(0), + b.stride(1), + c.stride(0), + c.stride(1), + ACTIVATION=activation_func, + ) + return c + + +@pytest.mark.parametrize('param_list', [(512, 512), (1024, 1024)]) +def test_matmul_mindspore(param_list): + os.environ["TRITON_BACKEND"] = "mindspore" + mindspore.set_seed(0) + shape = param_list + a = mindspore.ops.randn(shape, dtype=mindspore.float16) + b = mindspore.ops.randn(shape, dtype=mindspore.float16) + triton_output = matmul(a, b, activation) + mindspore_output = mindspore_matmul(a, b, activation) + assert np.allclose(triton_output.asnumpy(), mindspore_output.asnumpy(), rtol=1e-3, atol=1e-3) + del os.environ["TRITON_BACKEND"] diff --git a/third_party/ascend/unittest/generalization_cases/test_max.py b/third_party/ascend/unittest/generalization_cases/test_max.py new file mode 100644 index 0000000000..6029b77d9b --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_max.py @@ -0,0 +1,313 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common +from test_common import TestUtils, check_ub_mem_overflow, get_dtype_size + + +# <<<<<<< test_max_1d +def torch_max(x0, dim, keepdim): + inp = x0 if x0.device == "cpu" else x0.cpu() + return torch.max(inp, dim=dim, keepdim=keepdim)[0].npu() + + +@triton.jit +def triton_max_1d(in_ptr0, out_ptr1, xnumel, XBLOCK: tl.constexpr): + xoffset = tl.program_id(0) + tl.arange(0, XBLOCK) + tmp0 = tl.load(in_ptr0 + xoffset, None) + tmp4 = tl.max(tmp0, 0) + tl.store(out_ptr1, tmp4, None) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_max_1d(dtype, shape): + if check_ub_mem_overflow(dtype, shape): + return + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty(1, dtype=eval("torch." + dtype)).npu() + numel = shape[0] + triton_max_1d[1, 1, 1](x0, triton_res, numel, numel) + torch_res = torch_max(x0, dim=0, keepdim=True) + test_common.validate_cmp(dtype, triton_res, torch_res) + + +# >>>>>>> test_max_1d + + +# <<<<<<< test_max_2d +@triton.jit +def triton_max_2d(in_ptr0, out_ptr0, dim: tl.constexpr, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, + NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, MNUMEL) + nblk_idx = tl.arange(0, NNUMEL) + mmask = mblk_idx < M + nmask = nblk_idx < N + mask = (mmask[:, None]) & (nmask[None, :]) + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + x = tl.load(in_ptr0 + idx, mask=mask, other=-float('inf')) + tmp4 = tl.max(x, dim) + if dim == 0: + tl.store(out_ptr0 + tl.arange(0, N), tmp4, None) + else: + tl.store(out_ptr0 + tl.arange(0, M), tmp4, None) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('dim', [0, 1]) +def test_max_2d(dtype, shape, dim): + dtype_size = get_dtype_size(dtype) + if dtype == 'int8' or dtype == 'bool': + if dtype_size * math.prod(shape) >= (TestUtils.ub_size / 20): + pytest.skip(f"dtype:{dtype} shape:{shape} mem overflow") + elif dtype_size * math.prod(shape) >= (TestUtils.ub_size / 5): + pytest.skip(f"dtype:{dtype} shape:{shape} mem overflow") + shapex, shapey = shape + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty([ + shape[1 - dim], + ], dtype=eval("torch." + dtype)).npu() + triton_max_2d[1, 1, 1](x0, triton_res, dim, shapex, shapey, shapex, shapey) + torch_res = torch_max(x0, dim=dim, keepdim=False) + test_common.validate_cmp(dtype, triton_res, torch_res) + + +# >>>>>>> test_max_2d + + +# <<<<<<< test_max_3d +def torch_max_3d(x0, no_reduce_dim): + inp = x0 if x0.device == "cpu" else x0.cpu() + if no_reduce_dim == 0: + return torch.max(torch.max(inp, 1)[0], 1)[0].npu() + elif no_reduce_dim == 1: + return torch.max(torch.max(inp, 0)[0], 1)[0].npu() + elif no_reduce_dim == 2: + return torch.max(torch.max(inp, 0)[0], 0)[0].npu() + else: + assert False, f"no reduce dim not right, no_reduce_dim = {no_reduce_dim}" + + +@triton.jit +def triton_max_3d_0_1(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + x = tl.load(in_ptr + idx) + tmp = tl.max(x, 0) + ret = tl.max(tmp, 0) + oidx = zidx + tl.store(out_ptr + oidx, ret) + + +@triton.jit +def triton_max_3d_0_2(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + x = tl.load(in_ptr + idx) + tmp = tl.max(x, 0) + ret = tl.max(tmp, 1) + oidx = yidx + tl.store(out_ptr + oidx, ret) + + +@triton.jit +def triton_max_3d_1_2(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + x = tl.load(in_ptr + idx) + tmp = tl.max(x, 1) + ret = tl.max(tmp, 1) + oidx = xidx + tl.store(out_ptr + oidx, ret) + + +def triton_max_3d(in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB, no_reduce_dim): + if no_reduce_dim == 0: + triton_max_3d_1_2[1, 1, 1](in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB) + elif no_reduce_dim == 1: + triton_max_3d_0_2[1, 1, 1](in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB) + elif no_reduce_dim == 2: + triton_max_3d_0_1[1, 1, 1](in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('no_reduce_dim', [0, 1, 2]) +def test_max_3d(dtype, shape, no_reduce_dim): + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty([ + shape[no_reduce_dim], + ], dtype=eval("torch." + dtype)).npu() + triton_max_3d(x0, triton_res, shape[0], shape[1], shape[2], shape[0], shape[1], shape[2], no_reduce_dim) + torch_res = torch_max_3d(x0, no_reduce_dim) + test_common.validate_cmp(dtype, triton_res, torch_res) + + +# >>>>>>> test_max_3d + + +# <<<<<<< test_max_4d +def torch_max_4d(x0, dim): + x0 = x0 if x0.device == "cpu" else x0.cpu() + if x0.dtype in (torch.int8, torch.int16, torch.int32): + x0 = x0.to(torch.int64) + return torch.max(x0, dim=dim)[0] + + +@triton.jit +def max_4d(out_ptr, x, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, DIM: tl.constexpr): + if DIM == 0: + ret = tl.reshape(tl.max(x, DIM), XB * YB * ZB * MB // XB) + o_idx = tl.arange(0, XB * YB * ZB * MB // XB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 1: + ret = tl.reshape(tl.max(x, DIM), XB * YB * ZB * MB // YB) + o_idx = tl.arange(0, XB * YB * ZB * MB // YB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 2: + ret = tl.reshape(tl.max(x, DIM), XB * YB * ZB * MB // ZB) + o_idx = tl.arange(0, XB * YB * ZB * MB // ZB) + tl.store(out_ptr + o_idx, ret) + else: + ret = tl.reshape(tl.max(x, DIM), XB * YB * ZB * MB // MB) + o_idx = tl.arange(0, XB * YB * ZB * MB // MB) + tl.store(out_ptr + o_idx, ret) + + +@triton.jit +def triton_max_kernel_4d(in_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + DIM: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + midx = tl.arange(0, MB) + + idx = xidx[:, None, None, None] * YB * ZB * MB + yidx[None, :, None, None] * ZB * MB + zidx[ + None, None, :, None] * MB + midx[None, None, None, :] + + x = tl.load(in_ptr + idx) + + max_4d(out_ptr, x, XB, YB, ZB, MB, DIM) + + +def triton_max_4d(in_ptr, out_ptr, XB, YB, ZB, MB, dim): + triton_max_kernel_4d[(1, )](in_ptr, out_ptr, XB, YB, ZB, MB, dim) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [(2, 2, 4, 8)]) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('dim', [0]) +def test_max_4d(dtype, shape, dim): + x0 = test_common.generate_tensor(shape, dtype).npu() + torch_res = torch_max_4d(x0, dim) + triton_res = torch.empty_like(torch_res).npu() + triton_max_4d(x0, triton_res, shape[0], shape[1], shape[2], shape[3], dim) + + test_common.validate_cmp(dtype, triton_res, torch_res) + + +# >>>>>>> test_max_4d + + +# <<<<<<< test_max_5d +def torch_max_5d(x0, dim): + x0 = x0 if x0.device == "cpu" else x0.cpu() + if x0.dtype in (torch.int8, torch.int16, torch.int32): + x0 = x0.to(torch.int64) + return torch.max(x0, dim=dim)[0] + + +@triton.jit +def max_5d(out_ptr, x, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, NB: tl.constexpr, + DIM: tl.constexpr): + if DIM == 0: + ret = tl.reshape(tl.max(x, DIM), XB * YB * ZB * MB * NB // XB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // XB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 1: + ret = tl.reshape(tl.max(x, DIM), XB * YB * ZB * MB * NB // YB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // YB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 2: + ret = tl.reshape(tl.max(x, DIM), XB * YB * ZB * MB * NB // ZB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // ZB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 3: + ret = tl.reshape(tl.max(x, DIM), XB * YB * ZB * MB * NB // MB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // MB) + tl.store(out_ptr + o_idx, ret) + else: + ret = tl.reshape(tl.max(x, DIM), XB * YB * ZB * MB * NB // NB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // NB) + tl.store(out_ptr + o_idx, ret) + + +@triton.jit +def triton_max_kernel_5d(in_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr, DIM: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + midx = tl.arange(0, MB) + nidx = tl.arange(0, NB) + + idx = xidx[:, None, None, None, None] * YB * ZB * MB * NB + yidx[None, :, None, None, None] * ZB * MB * NB + zidx[ + None, None, :, None, None] * MB * NB + midx[None, None, None, :, None] * NB + nidx[None, None, None, None, :] + + x = tl.load(in_ptr + idx) + + max_5d(out_ptr, x, XB, YB, ZB, MB, NB, DIM) + + +def triton_max_5d(in_ptr, out_ptr, XB, YB, ZB, MB, NB, dim): + triton_max_kernel_5d[(1, )](in_ptr, out_ptr, XB, YB, ZB, MB, NB, dim) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [(2, 2, 2, 4, 8)]) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('dim', [0]) +def test_max_5d(dtype, shape, dim): + x0 = test_common.generate_tensor(shape, dtype).npu() + torch_res = torch_max_5d(x0, dim) + triton_res = torch.empty_like(torch_res).npu() + triton_max_5d(x0, triton_res, shape[0], shape[1], shape[2], shape[3], shape[4], dim) + + test_common.validate_cmp(dtype, triton_res, torch_res) + + +# >>>>>>> test_max_5d diff --git a/third_party/ascend/unittest/generalization_cases/test_min.py b/third_party/ascend/unittest/generalization_cases/test_min.py new file mode 100644 index 0000000000..2c080e3b18 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_min.py @@ -0,0 +1,313 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common +from test_common import TestUtils, check_ub_mem_overflow, get_dtype_size + + +# <<<<<<< test_min_1d +def torch_min(x0, dim, keepdim): + inp = x0 if x0.device == "cpu" else x0.cpu() + return torch.min(inp, dim=dim, keepdim=keepdim)[0].npu() + + +@triton.jit +def triton_min_1d(in_ptr0, out_ptr1, xnumel, XBLOCK: tl.constexpr): + xoffset = tl.program_id(0) + tl.arange(0, XBLOCK) + tmp0 = tl.load(in_ptr0 + xoffset, None) + tmp4 = tl.min(tmp0, 0) + tl.store(out_ptr1, tmp4, None) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +def test_min_1d(dtype, shape): + if check_ub_mem_overflow(dtype, shape): + return + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty(1, dtype=eval("torch." + dtype)).npu() + numel = shape[0] + triton_min_1d[1, 1, 1](x0, triton_res, numel, numel) + torch_res = torch_min(x0, dim=0, keepdim=True) + test_common.validate_cmp(dtype, triton_res, torch_res) + + +# >>>>>>> test_min_1d + + +# <<<<<<< test_min_2d +@triton.jit +def triton_min_2d(in_ptr0, out_ptr0, dim: tl.constexpr, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, + NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, MNUMEL) + nblk_idx = tl.arange(0, NNUMEL) + mmask = mblk_idx < M + nmask = nblk_idx < N + mask = (mmask[:, None]) & (nmask[None, :]) + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + x = tl.load(in_ptr0 + idx, mask=mask, other=-float('inf')) + tmp4 = tl.min(x, dim) + if dim == 0: + tl.store(out_ptr0 + tl.arange(0, N), tmp4, None) + else: + tl.store(out_ptr0 + tl.arange(0, M), tmp4, None) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('dim', [0, 1]) +def test_min_2d(dtype, shape, dim): + dtype_size = get_dtype_size(dtype) + if dtype == 'int8' or dtype == 'bool': + if dtype_size * math.prod(shape) >= (TestUtils.ub_size / 20): + pytest.skip(f"dtype:{dtype} shape:{shape} mem overflow") + elif dtype_size * math.prod(shape) >= (TestUtils.ub_size / 5): + pytest.skip(f"dtype:{dtype} shape:{shape} mem overflow") + shapex, shapey = shape + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty([ + shape[1 - dim], + ], dtype=eval("torch." + dtype)).npu() + triton_min_2d[1, 1, 1](x0, triton_res, dim, shapex, shapey, shapex, shapey) + torch_res = torch_min(x0, dim=dim, keepdim=False) + test_common.validate_cmp(dtype, triton_res, torch_res) + + +# >>>>>>> test_min_2d + + +# <<<<<<< test_min_3d +def torch_min_3d(x0, no_reduce_dim): + inp = x0 if x0.device == "cpu" else x0.cpu() + if no_reduce_dim == 0: + return torch.min(torch.min(inp, 1)[0], 1)[0].npu() + elif no_reduce_dim == 1: + return torch.min(torch.min(inp, 0)[0], 1)[0].npu() + elif no_reduce_dim == 2: + return torch.min(torch.min(inp, 0)[0], 0)[0].npu() + else: + assert False, f"no reduce dim not right, no_reduce_dim = {no_reduce_dim}" + + +@triton.jit +def triton_min_3d_0_1(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + x = tl.load(in_ptr + idx) + tmp = tl.min(x, 0) + ret = tl.min(tmp, 0) + oidx = zidx + tl.store(out_ptr + oidx, ret) + + +@triton.jit +def triton_min_3d_0_2(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + x = tl.load(in_ptr + idx) + tmp = tl.min(x, 0) + ret = tl.min(tmp, 1) + oidx = yidx + tl.store(out_ptr + oidx, ret) + + +@triton.jit +def triton_min_3d_1_2(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + x = tl.load(in_ptr + idx) + tmp = tl.min(x, 1) + ret = tl.min(tmp, 1) + oidx = xidx + tl.store(out_ptr + oidx, ret) + + +def triton_min_3d(in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB, no_reduce_dim): + if no_reduce_dim == 0: + triton_min_3d_1_2[1, 1, 1](in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB) + elif no_reduce_dim == 1: + triton_min_3d_0_2[1, 1, 1](in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB) + elif no_reduce_dim == 2: + triton_min_3d_0_1[1, 1, 1](in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('no_reduce_dim', [0, 1, 2]) +def test_min_3d(dtype, shape, no_reduce_dim): + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty([ + shape[no_reduce_dim], + ], dtype=eval("torch." + dtype)).npu() + triton_min_3d(x0, triton_res, shape[0], shape[1], shape[2], shape[0], shape[1], shape[2], no_reduce_dim) + torch_res = torch_min_3d(x0, no_reduce_dim) + test_common.validate_cmp(dtype, triton_res, torch_res) + + +# >>>>>>> test_min_3d + + +# <<<<<<< test_min_4d +def torch_min_4d(x0, dim): + x0 = x0 if x0.device == "cpu" else x0.cpu() + if x0.dtype in (torch.int8, torch.int16, torch.int32): + x0 = x0.to(torch.int64) + return torch.min(x0, dim=dim)[0] + + +@triton.jit +def min_4d(out_ptr, x, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, DIM: tl.constexpr): + if DIM == 0: + ret = tl.reshape(tl.min(x, DIM), XB * YB * ZB * MB // XB) + o_idx = tl.arange(0, XB * YB * ZB * MB // XB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 1: + ret = tl.reshape(tl.min(x, DIM), XB * YB * ZB * MB // YB) + o_idx = tl.arange(0, XB * YB * ZB * MB // YB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 2: + ret = tl.reshape(tl.min(x, DIM), XB * YB * ZB * MB // ZB) + o_idx = tl.arange(0, XB * YB * ZB * MB // ZB) + tl.store(out_ptr + o_idx, ret) + else: + ret = tl.reshape(tl.min(x, DIM), XB * YB * ZB * MB // MB) + o_idx = tl.arange(0, XB * YB * ZB * MB // MB) + tl.store(out_ptr + o_idx, ret) + + +@triton.jit +def triton_min_kernel_4d(in_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + DIM: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + midx = tl.arange(0, MB) + + idx = xidx[:, None, None, None] * YB * ZB * MB + yidx[None, :, None, None] * ZB * MB + zidx[ + None, None, :, None] * MB + midx[None, None, None, :] + + x = tl.load(in_ptr + idx) + + min_4d(out_ptr, x, XB, YB, ZB, MB, DIM) + + +def triton_min_4d(in_ptr, out_ptr, XB, YB, ZB, MB, dim): + triton_min_kernel_4d[(1, )](in_ptr, out_ptr, XB, YB, ZB, MB, dim) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [(2, 2, 4, 8)]) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('dim', [0]) +def test_min_4d(dtype, shape, dim): + x0 = test_common.generate_tensor(shape, dtype).npu() + torch_res = torch_min_4d(x0, dim) + triton_res = torch.empty_like(torch_res).npu() + triton_min_4d(x0, triton_res, shape[0], shape[1], shape[2], shape[3], dim) + + test_common.validate_cmp(dtype, triton_res, torch_res) + + +# >>>>>>> test_min_4d + + +# <<<<<<< test_min_5d +def torch_min_5d(x0, dim): + x0 = x0 if x0.device == "cpu" else x0.cpu() + if x0.dtype in (torch.int8, torch.int16, torch.int32): + x0 = x0.to(torch.int64) + return torch.min(x0, dim=dim)[0] + + +@triton.jit +def min_5d(out_ptr, x, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, NB: tl.constexpr, + DIM: tl.constexpr): + if DIM == 0: + ret = tl.reshape(tl.min(x, DIM), XB * YB * ZB * MB * NB // XB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // XB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 1: + ret = tl.reshape(tl.min(x, DIM), XB * YB * ZB * MB * NB // YB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // YB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 2: + ret = tl.reshape(tl.min(x, DIM), XB * YB * ZB * MB * NB // ZB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // ZB) + tl.store(out_ptr + o_idx, ret) + elif DIM == 3: + ret = tl.reshape(tl.min(x, DIM), XB * YB * ZB * MB * NB // MB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // MB) + tl.store(out_ptr + o_idx, ret) + else: + ret = tl.reshape(tl.min(x, DIM), XB * YB * ZB * MB * NB // NB) + o_idx = tl.arange(0, XB * YB * ZB * MB * NB // NB) + tl.store(out_ptr + o_idx, ret) + + +@triton.jit +def triton_min_kernel_5d(in_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr, DIM: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + midx = tl.arange(0, MB) + nidx = tl.arange(0, NB) + + idx = xidx[:, None, None, None, None] * YB * ZB * MB * NB + yidx[None, :, None, None, None] * ZB * MB * NB + zidx[ + None, None, :, None, None] * MB * NB + midx[None, None, None, :, None] * NB + nidx[None, None, None, None, :] + + x = tl.load(in_ptr + idx) + + min_5d(out_ptr, x, XB, YB, ZB, MB, NB, DIM) + + +def triton_min_5d(in_ptr, out_ptr, XB, YB, ZB, MB, NB, dim): + triton_min_kernel_5d[(1, )](in_ptr, out_ptr, XB, YB, ZB, MB, NB, dim) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [(2, 2, 2, 4, 8)]) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('dim', [0]) +def test_min_5d(dtype, shape, dim): + x0 = test_common.generate_tensor(shape, dtype).npu() + torch_res = torch_min_5d(x0, dim) + triton_res = torch.empty_like(torch_res).npu() + triton_min_5d(x0, triton_res, shape[0], shape[1], shape[2], shape[3], shape[4], dim) + + test_common.validate_cmp(dtype, triton_res, torch_res) + + +# >>>>>>> test_min_5d diff --git a/third_party/ascend/unittest/generalization_cases/test_mod.py b/third_party/ascend/unittest/generalization_cases/test_mod.py new file mode 100644 index 0000000000..ce15ea3d84 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_mod.py @@ -0,0 +1,227 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils +import math + + +def torch_pointwise(x, y): + res = x % y + return res + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = X % Y + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_mod_4d( + output_ptr, + x_ptr, + y_ptr, + BLOCK_SIZE: tl.constexpr, + SUB_BLOCK: tl.constexpr, + SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, + SHAPE_3: tl.constexpr, + STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, + STRIDE_2: tl.constexpr, + STRIDE_3: tl.constexpr, +): + pid = tl.program_id(0) + for loop in range(0, tl.cdiv(BLOCK_SIZE, SUB_BLOCK)): + base_idx = tl.arange(0, SUB_BLOCK) + pid_tensor = tl.full((SUB_BLOCK, ), pid * BLOCK_SIZE + loop * SUB_BLOCK, dtype=tl.int32) + tmp0 = (pid_tensor + base_idx)[:, None, None, None] + tmp1 = tl.arange(0, SHAPE_1)[None, :, None, None] + tmp2 = tl.arange(0, SHAPE_2)[None, None, :, None] + tmp3 = tl.arange(0, SHAPE_3)[None, None, None, :] + offsets = tmp0 * STRIDE_0 + tmp1 * STRIDE_1 + tmp2 * STRIDE_2 + tmp3 * STRIDE_3 + masks = tmp0 < SHAPE_0 + x = tl.load(x_ptr + offsets, mask=masks) + y = tl.load(y_ptr + offsets, mask=masks) + ret = x % y + tl.store(output_ptr + offsets, ret, mask=masks) + + +@triton.jit +def triton_mod_5d( + output_ptr, + x_ptr, + y_ptr, + BLOCK_SIZE: tl.constexpr, + SUB_BLOCK: tl.constexpr, + SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, + SHAPE_3: tl.constexpr, + SHAPE_4: tl.constexpr, + STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, + STRIDE_2: tl.constexpr, + STRIDE_3: tl.constexpr, + STRIDE_4: tl.constexpr, +): + pid = tl.program_id(0) + for loop in range(0, tl.cdiv(BLOCK_SIZE, SUB_BLOCK)): + base_idx = tl.arange(0, SUB_BLOCK) + pid_tensor = tl.full((SUB_BLOCK, ), pid * BLOCK_SIZE + loop * SUB_BLOCK, dtype=tl.int32) + tmp0 = (pid_tensor + base_idx)[:, None, None, None, None] + tmp1 = tl.arange(0, SHAPE_1)[None, :, None, None, None] + tmp2 = tl.arange(0, SHAPE_2)[None, None, :, None, None] + tmp3 = tl.arange(0, SHAPE_3)[None, None, None, :, None] + tmp4 = tl.arange(0, SHAPE_4)[None, None, None, None, :] + offsets = tmp0 * STRIDE_0 + tmp1 * STRIDE_1 + tmp2 * STRIDE_2 + tmp3 * STRIDE_3 + tmp4 * STRIDE_4 + masks = tmp0 < SHAPE_0 + x = tl.load(x_ptr + offsets, mask=masks) + y = tl.load(y_ptr + offsets, mask=masks) + ret = x % y + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1_2_3d) +@pytest.mark.parametrize('dtype', ['float16', 'float32', 'bfloat16', 'int8', 'int16', 'int32', 'int64']) +def test_case2(dtype, shape): + if dtype in ['int8', 'int16', 'int32', 'int64']: + x = test_common.generate_tensor_int_withSigns(shape, dtype).npu() + y = test_common.generate_tensor_int_withSigns(shape, dtype).npu() + z = test_common.generate_tensor_int_withSigns(shape, dtype).npu() + else: + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + + x[x <= 0] = 1 + y[y <= 0] = 1 + z[z <= 0] = 1 + + ans = torch_pointwise(x.cpu(), y.cpu()) + ans = ans.npu() + output = torch.zeros_like(ans) + + if len(shape) == 1: + fn_npu_[1, 1, shape[0]](output, x, y, z, 1, 1, 1, 1, 1, shape[0]) + elif len(shape) == 2: + if shape[0] > shape[1]: + fn_npu_[1, shape[0], 1](output, x, y, z, 1, 1, shape[1], 1, shape[0], shape[1]) + else: + fn_npu_[1, 1, shape[1]](output, x, y, z, 1, shape[0], 1, 1, shape[0], shape[1]) + elif len(shape) == 3: + if max(shape[0], shape[1], shape[2]) == shape[0]: + fn_npu_[shape[0], 1, 1](output, x, y, z, 1, shape[1], shape[2], shape[0], shape[1], shape[2]) + elif max(shape[0], shape[1], shape[2]) == shape[1]: + fn_npu_[1, shape[1], 1](output, x, y, z, shape[0], 1, shape[2], shape[0], shape[1], shape[2]) + else: + fn_npu_[1, 1, shape[2]](output, x, y, z, shape[0], shape[1], 1, shape[0], shape[1], shape[2]) + else: + fn_npu_[1, 1, 1](output, x, y, z, 1, 1, 1, 1, 1, 1) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + [(25, 2, 3, 31), (2, 2, 39, 23), (17, 27, 3, 3), + (3, 2, 27, 37)]) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +def test_mod_4d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + if dtype in ['int8', 'int16', 'int32', 'int64']: + x = test_common.generate_tensor_int_withSigns(shape, dtype).npu() + y = test_common.generate_tensor_int_withSigns(shape, dtype).npu() + else: + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + x[x <= 0] = 1 + y[y <= 0] = 1 + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_pointwise(x.cpu(), y.cpu()) + ans = ans.npu() + + n = x.numel() + block_size = min(triton.next_power_of_2(n), 64) + sub_block_size = 1 + grid = (triton.cdiv(n, block_size), ) + print(" ") + print(f"=== loops: {triton.cdiv(block_size, sub_block_size)}") + print(f"=== grid : {grid}") + triton_mod_4d[grid](output, x, y, block_size, sub_block_size, *list(shape), *list(x.stride())) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape5d + [(32, 5, 3, 1, 8)]) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +def test_mod_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + if dtype in ['int8', 'int16', 'int32', 'int64']: + x = test_common.generate_tensor_int_withSigns(shape, dtype).npu() + y = test_common.generate_tensor_int_withSigns(shape, dtype).npu() + else: + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + x[x <= 0] = 1 + y[y <= 0] = 1 + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_pointwise(x.cpu(), y.cpu()) + ans = ans.npu() + + n = x.numel() + block_size = min(triton.next_power_of_2(n), 32) + sub_block_size = 1 + grid = (triton.cdiv(n, block_size), ) + print(" ") + print(f"=== loops: {triton.cdiv(block_size, sub_block_size)}") + print(f"=== grid : {grid}") + triton_mod_5d[grid](output, x, y, block_size, sub_block_size, *list(shape), *list(x.stride())) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_multi-process.py b/third_party/ascend/unittest/generalization_cases/test_multi-process.py new file mode 100644 index 0000000000..8963097b48 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_multi-process.py @@ -0,0 +1,106 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +testing the compilation of multi-process operators +============= +""" +import multiprocessing +import os +import shutil + +import psutil +import torch +import torch_npu +import triton +import triton.backends.ascend.runtime +import triton.language as tl + + +@triton.jit +def add_kernel( + x_ptr, + y_ptr, + output_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + tl.store(output_ptr + offsets, output, mask=mask) + + +def add_torch(x, y): + return x + y + + +def add_triton(x, y): + output = torch.empty_like(x) + n_elements = output.numel() + BLOCK_SIZE = 1024 + add_kernel[ + triton.cdiv(n_elements, BLOCK_SIZE), + ]( + x, + y, + output, + n_elements, + BLOCK_SIZE, + ) + return output + + +def case_add(size: int): + x = torch.rand(size, device="npu") + y = torch.rand(size, device="npu") + try: + output_torch = add_torch(x, y) + output_triton = add_triton(x, y) + assert torch.allclose(output_triton, output_torch) + except Exception: + import traceback + traceback.print_exc() + return f"{multiprocessing.current_process().name}:failed" + return f"{multiprocessing.current_process().name}:success" + + +def test_multi_process(): + # init env + TA_CACHE_DEFAULT_PATH = os.environ.get("TRITON_CACHE_DIR", "~/.triton/cache") + TA_CACHE_PATH = os.path.join(TA_CACHE_DEFAULT_PATH, "multi_process") + if os.path.exists(TA_CACHE_PATH): + shutil.rmtree(TA_CACHE_PATH) + os.environ["TRITON_CACHE_DIR"] = TA_CACHE_PATH + # multi process + process_num = min(len(psutil.Process().cpu_affinity()), 64) + + multiprocessing.set_start_method("spawn", force=True) + results = [] + with multiprocessing.Pool(processes=process_num) as pool: + results = pool.map(case_add, [(98432, )] * process_num) + success_num = 0 + for result in results: + if result and "success" in result: + success_num += 1 + assert success_num == len(results), f"multi-process failed,failed {len(results) - success_num}" diff --git a/third_party/ascend/unittest/generalization_cases/test_ne.py b/third_party/ascend/unittest/generalization_cases/test_ne.py new file mode 100644 index 0000000000..a05220da45 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_ne.py @@ -0,0 +1,128 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common +from test_common import TestUtils +import math +import logging + + +def torch_ne(x0, x1): + if x0.dtype != torch.uint32: + return x0 != x1 + else: + return x0.to(torch.float32) != x1.to(torch.float32) + + +@triton.jit +def triton_ne(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x_index = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + x_index, mask=x_index < N) + tmp1 = tl.load(in_ptr1 + x_index, mask=x_index < N) + tmp2 = tmp0 != tmp1 + tl.store(out_ptr0 + x_index, tmp2, mask=x_index < N) + + +@triton.jit +def triton_ne_4d_5d(x_ptr, y_ptr, output_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = x_val != y_val + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1_2_3d) +@pytest.mark.parametrize('dtype', ['bool', 'int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32']) +def test_ne(shape, dtype): + logging.debug(f'dtype:{dtype} shape:{shape}') + # 生成数据 + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + + numel = x0.numel() + ncore = 1 if numel <= 32 else 32 + xblock = math.ceil(numel / ncore) + xblock_sub = numel if numel <= ncore else math.ceil(numel / ncore) + + # torch结果 + torch_res = torch_ne(x0, x1).to(eval('torch.' + dtype)) + # triton结果 + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + N = triton_res.numel() + triton_ne[ncore, 1, 1](x0, x1, triton_res, N, xblock, xblock_sub) + # 比较结果 + torch_res = torch_res if dtype != 'uint32' else torch_res.to(torch.float32) + triton_res = triton_res if dtype != 'uint32' else triton_res.to(torch.float32) + cmp_dtype = dtype if dtype != 'uint32' else 'float32' + test_common.validate_cmp(cmp_dtype, triton_res, torch_res) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +def test_ne_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_ne(x, y).to(eval('torch.' + dtype)) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_ne_4d_5d[grid](x, y, output, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_neg.py b/third_party/ascend/unittest/generalization_cases/test_neg.py new file mode 100644 index 0000000000..07c738d28f --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_neg.py @@ -0,0 +1,158 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils +import math + + +def torch_pointwise(x): + res = -x + return res + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = -X + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_neg_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = -x_val + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16', 'int8', 'int16', 'int32', 'int64']) +def test_case2(dtype, shape): + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.debug(f"output.dtype={output.dtype}") + + ans = torch_pointwise(x.cpu()) + ans = ans.npu() + + if len(shape) == 1: + fn_npu_[1, 1, shape[0]](output, x, y, z, 1, 1, 1, 1, 1, shape[0]) + elif len(shape) == 2: + if shape[0] > shape[1]: + fn_npu_[1, shape[0], 1](output, x, y, z, 1, 1, shape[1], 1, shape[0], shape[1]) + else: + fn_npu_[1, 1, shape[1]](output, x, y, z, 1, shape[0], 1, 1, shape[0], shape[1]) + elif len(shape) == 3: + if max(shape[0], shape[1], shape[2]) == shape[0]: + fn_npu_[shape[0], 1, 1](output, x, y, z, 1, shape[1], shape[2], shape[0], shape[1], shape[2]) + elif max(shape[0], shape[1], shape[2]) == shape[1]: + fn_npu_[1, shape[1], 1](output, x, y, z, shape[0], 1, shape[2], shape[0], shape[1], shape[2]) + else: + fn_npu_[1, 1, shape[2]](output, x, y, z, shape[0], shape[1], 1, shape[0], shape[1], shape[2]) + else: + fn_npu_[1, 1, 1](output, x, y, z, 1, 1, 1, 1, 1, 1) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16', 'int8', 'int16', 'int32', 'int64']) +def test_neg_4d_5d(shape, dtype): + x = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + logging.debug(f"output.dtype={output.dtype}") + + ans = torch_pointwise(x.cpu()) + ans = ans.npu() + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_neg_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) + + +invalid_types = [ + 'bool', +] + + +@pytest.mark.parametrize("sigtype", invalid_types) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "unexpected type") +def test_invalid_types(sigtype): + N = 32 + x = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + y = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + z = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + output = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + + fn_npu_[1, 1, 1](output, x, y, z, 32, 1, 1, 32, 1, 1) diff --git a/third_party/ascend/unittest/generalization_cases/test_not.py b/third_party/ascend/unittest/generalization_cases/test_not.py new file mode 100644 index 0000000000..21397985d6 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_not.py @@ -0,0 +1,173 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common +from test_common import TestUtils +import math + + +def torch_not(x0): + res = torch.bitwise_not(x0) + return res + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = not (X) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_not_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = not (x_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'bool']) +def test_case2(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.debug(f"output.dtype={output.dtype}") + + ans = torch_not(x) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +invalid_types = [ + 'float16', + 'float32', + 'bfloat16', +] + + +@pytest.mark.parametrize("sigtype", invalid_types) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "unexpected type") +def test_invalid_types(sigtype): + N = 32 + x = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + y = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + z = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + output = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + + fn_npu_[1, 1, 1](output, x, y, z, 32, 1, 1, 32, 1, 1) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'bool']) +def test_not_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_not(x) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_not_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_or.py b/third_party/ascend/unittest/generalization_cases/test_or.py new file mode 100644 index 0000000000..9861b8daf7 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_or.py @@ -0,0 +1,159 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common +from test_common import TestUtils +import math + + +def torch_or(x0, x1): + return x0 | x1 + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = X | Y + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_or_4d_5d(output_ptr, x_ptr, y_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = x_val | y_val + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'bool']) +def test_case2(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + ans = torch_or(x, y) + output = torch.zeros_like(ans) + + if len(shape) == 1: + fn_npu_[1, 1, shape[0]](output, x, y, z, 1, 1, 1, 1, 1, shape[0]) + elif len(shape) == 2: + if shape[0] > shape[1]: + fn_npu_[1, shape[0], 1](output, x, y, z, 1, 1, shape[1], 1, shape[0], shape[1]) + else: + fn_npu_[1, 1, shape[1]](output, x, y, z, 1, shape[0], 1, 1, shape[0], shape[1]) + elif len(shape) == 3: + if max(shape[0], shape[1], shape[2]) == shape[0]: + fn_npu_[shape[0], 1, 1](output, x, y, z, 1, shape[1], shape[2], shape[0], shape[1], shape[2]) + elif max(shape[0], shape[1], shape[2]) == shape[1]: + fn_npu_[1, shape[1], 1](output, x, y, z, shape[0], 1, shape[2], shape[0], shape[1], shape[2]) + else: + fn_npu_[1, 1, shape[2]](output, x, y, z, shape[0], shape[1], 1, shape[0], shape[1], shape[2]) + else: + fn_npu_[1, 1, 1](output, x, y, z, 1, 1, 1, 1, 1, 1) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'bool']) +def test_or_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = x | y + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_or_4d_5d[grid](output, x, y, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) + + +invalid_types = [ + 'float16', + 'float32', + 'bfloat16', +] + + +@pytest.mark.parametrize("sigtype", invalid_types) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "unexpected type") +def test_invalid_types(sigtype): + N = 32 + x = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + y = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + z = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + output = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + + fn_npu_[1, 1, 1](output, x, y, z, 32, 1, 1, 32, 1, 1) diff --git a/third_party/ascend/unittest/generalization_cases/test_permute_1d_2d.py b/third_party/ascend/unittest/generalization_cases/test_permute_1d_2d.py new file mode 100644 index 0000000000..70d41abc22 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_permute_1d_2d.py @@ -0,0 +1,95 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils, check_ub_mem_overflow +import math +import logging + + +@triton.jit +def fn_npu_1d(output_ptr, x_ptr, xnumel: tl.constexpr): + idx = tl.arange(0, xnumel) + + X = tl.load(x_ptr + idx) + + ret = tl.permute(X, (0, )) + + tl.store(output_ptr + idx, ret) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', TestUtils.dtype_list) +def test_permute_1d(shape, dtype): + logging.debug(f'dtype:{dtype} shape:{shape}') + + data_type = eval('torch.' + dtype) + x = torch.randint(low=0, high=2, size=shape, dtype=data_type).npu() + + triton_res = torch.randint(1, shape, dtype=data_type).npu() + torch_res = torch.permute(x, (0, )) + fn_npu_1d[1, 1, 1](triton_res, x, shape[0]) + test_common.validate_cmp(dtype, triton_res, torch_res) + + +@triton.jit +def fn_npu_021(output_ptr, x_ptr, YB: tl.constexpr, ZB: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr): + pid = tl.program_id(0) + yidx = tl.arange(0, YB) + pid * YB + zidx = tl.arange(0, ZB) + idx = yidx[:, None] * znumel + zidx[None, :] + + # XB,YB,1 + X = tl.load(x_ptr + idx) + + ret = tl.permute(X, (1, 0)) + + oidx = zidx[:, None] * ynumel + yidx[None, :] + + tl.store(output_ptr + oidx, ret) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', TestUtils.dtype_list) +def test_permute(shape, dtype): + logging.debug(f'dtype:{dtype} shape:{shape}') + + ynumel = shape[0] + YB = 1 + znumel = shape[1] + ZB = shape[1] + + data_type = eval('torch.' + dtype) + x = torch.randint(low=0, high=2, size=(shape[0], shape[1]), dtype=data_type).npu() + + triton_res = torch.randint(1, (shape[1], shape[0]), dtype=data_type).npu() + torch_res = torch.permute(x, (1, 0)) + fn_npu_021[shape[0], 1, 1](triton_res, x, YB, ZB, ynumel, znumel) + test_common.validate_cmp(dtype, triton_res, torch_res) + + +if __name__ == "__main__": + for shape in [(37, 3)]: + for dtype in TestUtils.dtype_list: + test_permute(shape, dtype) diff --git a/third_party/ascend/unittest/generalization_cases/test_permute_3d.py b/third_party/ascend/unittest/generalization_cases/test_permute_3d.py new file mode 100644 index 0000000000..9696e09f32 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_permute_3d.py @@ -0,0 +1,101 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils, check_ub_mem_overflow +import math +import logging + + +@triton.jit +def fn_npu_102(output_ptr, x_ptr, YB: tl.constexpr, ZB: tl.constexpr, KB: tl.constexpr): + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + kidx = tl.arange(0, KB) + idx = yidx[:, None, None] * ZB * KB + zidx[None, :, None] * KB + kidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.permute(X, (1, 0, 2)) + + oidx = zidx[:, None, None] * YB * KB + yidx[None, :, None] * KB + kidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_210(output_ptr, x_ptr, YB: tl.constexpr, ZB: tl.constexpr, KB: tl.constexpr): + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + kidx = tl.arange(0, KB) + idx = yidx[:, None, None] * ZB * KB + zidx[None, :, None] * KB + kidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.permute(X, (2, 1, 0)) + + oidx = kidx[:, None, None] * ZB * YB + zidx[None, :, None] * YB + yidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_021(output_ptr, x_ptr, YB: tl.constexpr, ZB: tl.constexpr, KB: tl.constexpr): + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + kidx = tl.arange(0, KB) + idx = yidx[:, None, None] * ZB * KB + zidx[None, :, None] * KB + kidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.permute(X, (0, 2, 1)) + + oidx = yidx[:, None, None] * ZB * KB + kidx[None, :, None] * ZB + zidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +@pytest.mark.parametrize('dtype', ["int8", 'int16', 'int32', 'float16', 'float32', 'bfloat16', 'int64']) +def test_permute_3d(shape, dtype): + logging.debug(f'dtype:{dtype} shape:{shape}') + + data_type = eval('torch.' + dtype) + x = torch.randint(low=0, high=2, size=shape, dtype=data_type).npu() + + triton_res = torch.empty((shape[1], shape[0], shape[2]), dtype=data_type).npu() + torch_res = torch.permute(x, (1, 0, 2)) + fn_npu_102[1, 1, 1](triton_res, x, shape[0], shape[1], shape[2]) + test_common.validate_cmp(dtype, triton_res, torch_res) + + # not support yet: need bisheng support later + # triton_res = torch.empty((shape[2], shape[1], shape[0]), dtype=data_type).npu() + # torch_res = torch.permute(x, (2, 1, 0)) + # fn_npu_210[1, 1, 1](triton_res, x, shape[0], shape[1], shape[2]) + # test_common.validate_cmp(dtype, triton_res, torch_res) + + triton_res = torch.empty((shape[0], shape[2], shape[1]), dtype=data_type).npu() + torch_res = torch.permute(x, (0, 2, 1)) + fn_npu_021[1, 1, 1](triton_res, x, shape[0], shape[1], shape[2]) + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/generalization_cases/test_permute_4d_5d.py b/third_party/ascend/unittest/generalization_cases/test_permute_4d_5d.py new file mode 100644 index 0000000000..615ff3cd6e --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_permute_4d_5d.py @@ -0,0 +1,223 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils, check_ub_mem_overflow +import math +import logging + + +@triton.jit +def triton_permute_4d( + output_ptr, + x_ptr, + PERM: tl.constexpr, + BLOCK_0: tl.constexpr, + BLOCK_1: tl.constexpr, + BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, + SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, + SHAPE_3: tl.constexpr, + STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, + STRIDE_2: tl.constexpr, + STRIDE_3: tl.constexpr, +): + pid = tl.program_id(0) + tmp0 = tl.arange(0, BLOCK_0)[:, None, None, None] + tmp0_1 = tl.arange(0, BLOCK_0)[None, :, None, None] + tmp1 = tl.arange(0, BLOCK_1)[None, :, None, None] + tmp1_0 = tl.arange(0, BLOCK_1)[:, None, None, None] + tmp1_2 = tl.arange(0, BLOCK_1)[None, None, :, None] + tmp2 = tl.arange(0, BLOCK_2)[None, None, :, None] + tmp2_1 = tl.arange(0, BLOCK_2)[None, :, None, None] + tmp2_3 = tl.arange(0, BLOCK_2)[None, None, None, :] + tmp3 = tl.arange(0, BLOCK_3)[None, None, None, :] + tmp3_2 = tl.arange(0, BLOCK_3)[None, None, :, None] + offsets = pid + tmp0 * STRIDE_0 + tmp1 * STRIDE_1 + tmp2 * STRIDE_2 + tmp3 * STRIDE_3 + masks = (tmp0 < SHAPE_0) & (tmp1 < SHAPE_1) & (tmp2 < SHAPE_2) & (tmp3 < SHAPE_3) + x_val = tl.load(x_ptr + offsets, masks) + + if PERM == 0: # 1, 0, 2, 3 + ret = tl.permute(x_val, (1, 0, 2, 3)) + shape0 = SHAPE_1 + shape1 = SHAPE_0 + shape2 = SHAPE_2 + shape3 = SHAPE_3 + elif PERM == 1: # 0, 2, 1, 3 + ret = tl.permute(x_val, (0, 2, 1, 3)) + shape0 = SHAPE_0 + shape1 = SHAPE_2 + shape2 = SHAPE_1 + shape3 = SHAPE_3 + else: # 0, 1, 3, 2 + ret = tl.permute(x_val, (0, 1, 3, 2)) + shape0 = SHAPE_0 + shape1 = SHAPE_1 + shape2 = SHAPE_3 + shape3 = SHAPE_2 + + s3 = 1 + s2 = s3 * shape3 + s1 = s2 * shape2 + s0 = s1 * shape1 + + if PERM == 0: # 1, 0, 2, 3 + out_offsets = pid + tmp1_0 * s0 + tmp0_1 * s1 + tmp2 * s2 + tmp3 * s3 + out_masks = (tmp1_0 < shape0) & (tmp0_1 < shape1) & (tmp2 < shape2) & (tmp3 < shape3) + elif PERM == 1: # 0, 2, 1, 3 + out_offsets = pid + tmp0 * s0 + tmp2_1 * s1 + tmp1_2 * s2 + tmp3 * s3 + out_masks = (tmp0 < shape0) & (tmp1_2 < shape2) & (tmp2_1 < shape1) & (tmp3 < shape3) + else: # 0, 1, 3, 2 + out_offsets = pid + tmp0 * s0 + tmp1 * s1 + tmp3_2 * s2 + tmp2_3 * s3 + out_masks = (tmp0 < shape0) & (tmp1 < shape1) & (tmp3_2 < shape2) & (tmp2_3 < shape3) + tl.store(output_ptr + out_offsets, ret, mask=out_masks) + + +@triton.jit +def triton_permute_5d(output_ptr, x_ptr, PERM: tl.constexpr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, + BLOCK_2: tl.constexpr, BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, + STRIDE_0: tl.constexpr, STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, + STRIDE_4: tl.constexpr): + pid = tl.program_id(0) + tmp0 = tl.arange(0, BLOCK_0)[:, None, None, None, None] + tmp1 = tl.arange(0, BLOCK_1)[None, :, None, None, None] + tmp2 = tl.arange(0, BLOCK_2)[None, None, :, None, None] + tmp3 = tl.arange(0, BLOCK_3)[None, None, None, :, None] + tmp4 = tl.arange(0, BLOCK_4)[None, None, None, None, :] + + tmp0_1 = tl.arange(0, BLOCK_0)[None, :, None, None, None] + tmp1_0 = tl.arange(0, BLOCK_1)[:, None, None, None, None] + + tmp1_2 = tl.arange(0, BLOCK_1)[None, None, :, None, None] + tmp2_1 = tl.arange(0, BLOCK_2)[None, :, None, None, None] + + tmp2_3 = tl.arange(0, BLOCK_2)[None, None, None, :, None] + tmp3_2 = tl.arange(0, BLOCK_3)[None, None, :, None, None] + + tmp3_4 = tl.arange(0, BLOCK_3)[None, None, None, None, :] + tmp4_3 = tl.arange(0, BLOCK_4)[None, None, None, :, None] + + offsets = pid + tmp0 * STRIDE_0 + tmp1 * STRIDE_1 + tmp2 * STRIDE_2 + tmp3 * STRIDE_3 + tmp4 * STRIDE_4 + masks = (tmp0 < SHAPE_0) & (tmp1 < SHAPE_1) & (tmp2 < SHAPE_2) & (tmp3 < SHAPE_3) & (tmp4 < SHAPE_4) + x_val = tl.load(x_ptr + offsets, masks) + + if PERM == 0: # 1, 0, 2, 3, 4 + ret = tl.permute(x_val, 1, 0, 2, 3, 4) + shape0 = SHAPE_1 + shape1 = SHAPE_0 + shape2 = SHAPE_2 + shape3 = SHAPE_3 + shape4 = SHAPE_4 + elif PERM == 1: # 0, 2, 1, 3, 4 + ret = tl.permute(x_val, 0, 2, 1, 3, 4) + shape0 = SHAPE_0 + shape1 = SHAPE_2 + shape2 = SHAPE_1 + shape3 = SHAPE_3 + shape4 = SHAPE_4 + elif PERM == 2: # 0, 1, 3, 2, 4 + ret = tl.permute(x_val, 0, 1, 3, 2, 4) + shape0 = SHAPE_0 + shape1 = SHAPE_1 + shape2 = SHAPE_3 + shape3 = SHAPE_2 + shape4 = SHAPE_4 + else: # 0, 1, 2, 4, 3 + ret = tl.permute(x_val, 0, 1, 2, 4, 3) + shape0 = SHAPE_0 + shape1 = SHAPE_1 + shape2 = SHAPE_2 + shape3 = SHAPE_4 + shape4 = SHAPE_3 + + s4 = 1 + s3 = s4 * shape4 + s2 = s3 * shape3 + s1 = s2 * shape2 + s0 = s1 * shape1 + + if PERM == 0: # 1, 0, 2, 3, 4 + out_offsets = pid + tmp1_0 * s0 + tmp0_1 * s1 + tmp2 * s2 + tmp3 * s3 + tmp4 * s4 + out_masks = (tmp1_0 < shape0) & (tmp0_1 < shape1) & (tmp2 < shape2) & (tmp3 < shape3) & (tmp4 < shape4) + elif PERM == 1: # 0, 2, 1, 3, 4 + out_offsets = pid + tmp0 * s0 + tmp2_1 * s1 + tmp1_2 * s2 + tmp3 * s3 + tmp4 * s4 + out_masks = (tmp0 < shape0) & (tmp1_2 < shape2) & (tmp2_1 < shape1) & (tmp3 < shape3) & (tmp4 < shape4) + elif PERM == 2: # 0, 1, 3, 2, 4 + out_offsets = pid + tmp0 * s0 + tmp1 * s1 + tmp3_2 * s2 + tmp2_3 * s3 + tmp4 * s4 + out_masks = (tmp0 < shape0) & (tmp1 < shape1) & (tmp3_2 < shape2) & (tmp2_3 < shape3) & (tmp4 < shape4) + else: # 0, 1, 2, 4, 3 + out_offsets = pid + tmp0 * s0 + tmp1 * s1 + tmp2 * s2 + tmp4_3 * s3 + tmp3_4 * s4 + out_masks = (tmp0 < shape0) & (tmp1 < shape1) & (tmp2 < shape2) & (tmp4_3 < shape3) & (tmp3_4 < shape4) + tl.store(output_ptr + out_offsets, ret, mask=out_masks) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +@pytest.mark.parametrize('perm', [0, 1, 2, 3]) # 4d: support 3 mode; 5d: support 4 mode +def test_permute_4d_5d(shape, dtype, perm): + logging.log(logging.DEBUG, f"shape = {shape}") + x = torch.randint(low=0, high=2, size=shape, dtype=eval('torch.' + dtype)).npu() + grid = (1, ) + if len(shape) == 4: + blocks = list(x.size()) + strides = list(x.stride()) + if perm == 0: # 1, 0, 2, 3; exchange axis 0, 1 + output = torch.empty((shape[1], shape[0], shape[2], shape[3]), dtype=eval('torch.' + dtype)).npu() + ans_4d = torch.permute(x, (1, 0, 2, 3)) + triton_permute_4d[grid](output, x, perm, *blocks, *blocks, *strides) + test_common.validate_cmp(dtype, ans_4d, output) + elif perm == 1: # 0, 2, 1, 3; exchange axis 1, 2 + output = torch.empty((shape[0], shape[2], shape[1], shape[3]), dtype=eval('torch.' + dtype)).npu() + ans_4d = torch.permute(x, (0, 2, 1, 3)) + triton_permute_4d[grid](output, x, perm, *blocks, *blocks, *strides) + test_common.validate_cmp(dtype, ans_4d, output) + elif perm == 2: # 0, 1, 3, 2; exchange axis 2, 3 + output = torch.empty((shape[0], shape[1], shape[3], shape[2]), dtype=eval('torch.' + dtype)).npu() + ans_4d = torch.permute(x, (0, 1, 3, 2)) + triton_permute_4d[grid](output, x, perm, *blocks, *blocks, *strides) + test_common.validate_cmp(dtype, ans_4d, output) + else: + pass + else: + blocks = list(x.size()) + strides = list(x.stride()) + + if perm == 0: # 1, 0, 2, 3, 4; exchange axis 0, 1 + output = torch.empty((shape[1], shape[0], shape[2], shape[3], shape[4]), dtype=eval('torch.' + dtype)).npu() + ans_5d = torch.permute(x, (1, 0, 2, 3, 4)) + elif perm == 1: # 0, 2, 1, 3, 4; exchange axis 1, 2 + output = torch.empty((shape[0], shape[2], shape[1], shape[3], shape[4]), dtype=eval('torch.' + dtype)).npu() + ans_5d = torch.permute(x, (0, 2, 1, 3, 4)) + elif perm == 2: # 0, 1, 3, 2, 4; exchange axis 2, 3 + output = torch.empty((shape[0], shape[1], shape[3], shape[2], shape[4]), dtype=eval('torch.' + dtype)).npu() + ans_5d = torch.permute(x, (0, 1, 3, 2, 4)) + else: # 0, 1, 2, 4, 3; exchange axis 3, 4 + output = torch.empty((shape[0], shape[1], shape[2], shape[4], shape[3]), dtype=eval('torch.' + dtype)).npu() + ans_5d = torch.permute(x, (0, 1, 2, 4, 3)) + triton_permute_5d[grid](output, x, perm, *blocks, *blocks, *strides) + test_common.validate_cmp(dtype, ans_5d, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_rand.py b/third_party/ascend/unittest/generalization_cases/test_rand.py new file mode 100644 index 0000000000..8e66e48bfa --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_rand.py @@ -0,0 +1,304 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils +import math +import numpy as np +import scipy + + +@triton.jit +def kernel_rand(x_ptr, n_rounds: tl.constexpr, N: tl.constexpr, XBLOCK: tl.constexpr): + block_offset = tl.program_id(0) * XBLOCK + block_size = XBLOCK if block_offset + XBLOCK <= N else N - block_offset + for inner_idx in range(block_size): + global_offset = block_offset + inner_idx + rand_vals = tl.rand(5, 10 + global_offset, n_rounds) # 对每个索引生成一个随机数 + tl.store(x_ptr + global_offset, rand_vals) # 存储随机数 + + +@triton.jit +def triton_rand_4d_5d(output_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + # 1D program_id for flatten multi-d offset + pid = tl.program_id(0) + # base offset for dimension 0 + offsets = pid + tl.arange(0, BLOCK_0) * STRIDE_0 + mask = tl.arange(0, BLOCK_0) < SHAPE_0 + # nested offset expansion + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + mask = mask[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + mask = mask[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + mask = mask[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + mask = mask[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + ret = tl.rand(5, offsets, 10) + tl.store(output_ptr + offsets, ret, mask=mask) + + +@triton.jit +def kernel_randn(x_ptr, n_rounds: tl.constexpr, N: tl.constexpr, XBLOCK: tl.constexpr): + block_offset = tl.program_id(0) * XBLOCK + block_size = XBLOCK if block_offset + XBLOCK <= N else N - block_offset + for inner_idx in range(block_size): + global_offset = block_offset + inner_idx + rand_vals = tl.randn(5, 10 + global_offset, n_rounds) # 对每个索引生成一个随机数 + tl.store(x_ptr + global_offset, rand_vals) # 存储随机数 + + +@triton.jit +def triton_randn_4d_5d(output_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + # 1D program_id for flatten multi-d offset + pid = tl.program_id(0) + # base offset for dimension 0 + offsets = pid + tl.arange(0, BLOCK_0) * STRIDE_0 + mask = tl.arange(0, BLOCK_0) < SHAPE_0 + # nested offset expansion + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + mask = mask[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + mask = mask[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + mask = mask[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + mask = mask[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + ret = tl.randn(5, offsets, 10) + tl.store(output_ptr + offsets, ret, mask=mask) + + +@triton.jit +def kernel_randint(x_ptr, n_rounds: tl.constexpr, N: tl.constexpr, XBLOCK: tl.constexpr): + block_offset = tl.program_id(0) * XBLOCK + block_size = XBLOCK if block_offset + XBLOCK <= N else N - block_offset + for inner_idx in range(block_size): + global_offset = block_offset + inner_idx + rand_vals = tl.randint(5, 10 + global_offset, n_rounds) # 对每个索引生成一个随机数 + tl.store(x_ptr + global_offset, rand_vals) # 存储随机数 + + +@triton.jit +def triton_randint_4d_5d(output_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, + STRIDE_4: tl.constexpr): + # 1D program_id for flatten multi-d offset + pid = tl.program_id(0) + # base offset for dimension 0 + offsets = pid + tl.arange(0, BLOCK_0) * STRIDE_0 + mask = tl.arange(0, BLOCK_0) < SHAPE_0 + # nested offset expansion + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + mask = mask[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + mask = mask[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + mask = mask[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + mask = mask[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + ret = tl.randint(5, offsets, 10) + tl.store(output_ptr + offsets, ret, mask=mask) + + +@triton.jit +def kernel_randint4x(x_ptr, n_rounds: tl.constexpr, N: tl.constexpr, XBLOCK: tl.constexpr): + block_offset = tl.program_id(0) * XBLOCK + indices = tl.arange(0, 4) + block_size = XBLOCK if block_offset + XBLOCK <= N else N - block_offset + for inner_idx in range(0, block_size + 4, step=4): + global_offset = block_offset + inner_idx + rand_vals = tl.randint4x(5, 10 + global_offset, n_rounds) # 对每个索引生成一个随机数 + mask = (global_offset + indices) < (block_offset + block_size) + tl.store(x_ptr + global_offset + indices, rand_vals, mask) # 存储随机数 + + +@triton.jit +def triton_randint4x_4d_5d(output_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, + STRIDE_4: tl.constexpr): + # 1D program_id for flatten multi-d offset + pid = tl.program_id(0) + # base offset for dimension 0 + offsets = pid + tl.arange(0, BLOCK_0) * STRIDE_0 + mask = tl.arange(0, BLOCK_0) < SHAPE_0 + # nested offset expansion + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + mask = mask[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + mask = mask[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + mask = mask[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + mask = mask[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + ret = tl.randint4x(5, offsets, 10) + tl.store(output_ptr + offsets, ret, mask=mask) + + +# With alpha=0.01, z=-3.0902, N=100, we have (1-0.01)+(-3.0902)*sqrt(0.01*(1-0.01)/100)=0.9593, +# so there must be 96 cases for each shape to have pvalue larger than 0.01. +# There is higher possibility to fail with small shapes, so we will use large shape. +@pytest.mark.parametrize('shape', [ + (256, 256), + (512, 512), + (1024, 1024), +]) +def test_rand_case(shape): + y_calf = torch.zeros(shape, dtype=eval('torch.float32')).npu() + + numel = y_calf.numel() + ncore = 1 if numel < 32 else 32 + xblock = math.ceil(numel / ncore) + + correctness = 0 + for _ in range(100): + ref = np.random.random_sample(shape).flatten() + kernel_rand[ncore, 1, 1](y_calf, 10, numel, xblock) + + pvalue = scipy.stats.kstest(ref, y_calf.cpu().numpy().flatten()).pvalue + if pvalue > 0.01: + correctness += 1 + + assert correctness > 95 + + +@pytest.mark.parametrize('shape', [ + (256, 256), + (512, 512), + (1024, 1024), +]) +def test_randn_case(shape): + y_calf = torch.zeros(shape, dtype=eval('torch.float32')).npu() + + numel = y_calf.numel() + ncore = 1 if numel < 32 else 32 + xblock = math.ceil(numel / ncore) + + correctness = 0 + for _ in range(100): + ref = np.random.standard_normal(shape).flatten() + kernel_randn[ncore, 1, 1](y_calf, 10, numel, xblock) + + pvalue = scipy.stats.kstest(ref, y_calf.cpu().numpy().flatten()).pvalue + if pvalue > 0.01: + correctness += 1 + + assert correctness > 95 + + +@pytest.mark.parametrize('shape', [ + (256, 256), + (512, 512), + (1024, 1024), +]) +def test_randint_case(shape): + y_cali = torch.zeros(shape, dtype=eval('torch.int32')).npu() + + numel = y_cali.numel() + ncore = 1 if numel < 32 else 32 + xblock = math.ceil(numel / ncore) + + correctness = 0 + ii32 = np.iinfo(np.int32) + for _ in range(100): + ref = np.random.randint(low=ii32.min, high=ii32.max, size=shape).flatten() + kernel_randint[ncore, 1, 1](y_cali, 10, numel, xblock) + + pvalue = scipy.stats.kstest(ref, y_cali.cpu().numpy().flatten()).pvalue + if pvalue > 0.01: + correctness += 1 + + assert correctness > 95 + + +@pytest.mark.parametrize('shape', [ + (256, 256), + (512, 512), + (1024, 1024), +]) +def test_randint4x_case(shape): + y_cali = torch.zeros(shape, dtype=eval('torch.int32')).npu() + + numel = y_cali.numel() + ncore = 1 if numel < 32 else 32 + xblock = math.ceil(numel / ncore) + + correctness = 0 + ii32 = np.iinfo(np.int32) + for _ in range(100): + ref = np.random.randint(low=ii32.min, high=ii32.max, size=shape).flatten() + kernel_randint4x[ncore, 1, 1](y_cali, 10, numel, xblock) + + pvalue = scipy.stats.kstest(ref, y_cali.cpu().numpy().flatten()).pvalue + if pvalue > 0.01: + correctness += 1 + + assert correctness > 95 + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +def test_rand_4d_5d(shape): + x = torch.zeros(shape, dtype=eval('torch.float32')).npu() + y = torch.zeros(shape, dtype=eval('torch.int32')).npu() + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_rand_4d_5d[grid](x, *blocks, *blocks, *strides) + triton_randn_4d_5d[grid](x, *blocks, *blocks, *strides) + triton_randint_4d_5d[grid](y, *blocks, *blocks, *strides) + triton_randint4x_4d_5d[grid](y, *blocks, *blocks, *strides) diff --git a/third_party/ascend/unittest/generalization_cases/test_range.py b/third_party/ascend/unittest/generalization_cases/test_range.py new file mode 100644 index 0000000000..992076fc6d --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_range.py @@ -0,0 +1,187 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import torch +import test_common +import logging + +import triton.language as tl +from test_common import TestUtils + + +@triton.jit +def triton_range(output_ptr, x_ptr, y_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = X + Y + for _ in tl.range(2, 5, 2): + ret = ret + X + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_static_range(output_ptr, x_ptr, y_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, + XNUMEL: tl.constexpr, YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = X + Y + for _ in tl.static_range(2, 5, 2): + ret = ret + X + + tl.store(output_ptr + idx, ret) + + +test_shape = [(1, ), (2, ), (1, 1), (3, 4), (1, 1, 1), (2, 4, 8)] + + +@pytest.mark.parametrize('shape', test_shape) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +def test_range(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + if dtype == 'bfloat16': + ans = (x.to(torch.float32) + y.to(torch.float32) + x.to(torch.float32) + x.to(torch.float32)).to(torch.bfloat16) + else: + ans = x + y + x + x + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if dtype == 'int8': + if x.numel() * x.element_size() >= 512: + grid = (1, 1, ZB) + ZB = 1 + else: + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + triton_range[grid](output, x, y, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', test_shape) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +def test_static_range(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + if dtype == 'bfloat16': + ans = (x.to(torch.float32) + y.to(torch.float32) + x.to(torch.float32) + x.to(torch.float32)).to(torch.bfloat16) + else: + ans = x + y + x + x + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if dtype == 'int8': + if x.numel() * x.element_size() >= 512: + grid = (1, 1, ZB) + ZB = 1 + else: + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + triton_static_range[grid](output, x, y, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_reduce.py b/third_party/ascend/unittest/generalization_cases/test_reduce.py new file mode 100644 index 0000000000..14ab4696cb --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_reduce.py @@ -0,0 +1,338 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import random +import pytest +import torch +import triton +import triton.language as tl + +import test_common +from test_common import TestUtils, get_dtype_size + + +def torch_reduce(x1, dim): + if x1.dtype == torch.float16 or x1.dtype == torch.float32: + res = torch.sum(x1.to(torch.float32), dim=dim).to(x1.dtype) + else: + res = torch.sum(x1, dim=dim).to(x1.dtype) + return res + + +@triton.jit +def _reduce_combine(a, b): + return a + b + + +@triton.jit +def tt_reduce_1d(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, XB: tl.constexpr, + YB: tl.constexpr, ZB: tl.constexpr, dim: tl.constexpr): + idx = tl.arange(0, XB) + x = tl.load(in_ptr + idx) + ret = tl.reduce(x, dim, _reduce_combine) + tl.store(out_ptr + tl.arange(0, 1), ret) + + +@triton.jit +def tt_reduce_2d(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, XB: tl.constexpr, + YB: tl.constexpr, ZB: tl.constexpr, dim: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + idx = xidx[:, None] * ynumel + yidx[None, :] + + x = tl.load(in_ptr + idx) + ret = tl.reduce(x, dim, _reduce_combine) + + if dim == 0: + oidx = yidx + else: + oidx = xidx + tl.store(out_ptr + oidx, ret) + + +@triton.jit +def tt_reduce_1d_dim_none(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, dim: tl.constexpr): + idx = tl.arange(0, XB) + x = tl.load(in_ptr + idx) + ret = tl.reduce(x, dim, _reduce_combine) + tl.store(out_ptr + tl.arange(0, 1), ret) + + +@triton.jit +def tt_reduce_2d_dim_none(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, dim: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + idx = xidx[:, None] * ynumel + yidx[None, :] + + x = tl.load(in_ptr + idx) + ret = tl.reduce(x, dim, _reduce_combine) + + tl.store(out_ptr + tl.arange(0, 1), ret) + + +@triton.jit +def tt_reduce_3d_dim_none(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, dim: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + + x = tl.load(in_ptr + idx) + ret = tl.reduce(x, dim, _reduce_combine) + + tl.store(out_ptr, ret) + + +@triton.jit +def tt_reduce_3d(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, XB: tl.constexpr, + YB: tl.constexpr, ZB: tl.constexpr, dim: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + + x = tl.load(in_ptr + idx) + ret = tl.reduce(x, dim, _reduce_combine) + + if dim == 0: + oidx = yidx[:, None] * znumel + zidx[None, :] + elif dim == 1: + oidx = xidx[:, None] * znumel + zidx[None, :] + else: + oidx = xidx[:, None] * ynumel + yidx[None, :] + + tl.store(out_ptr + oidx, ret) + + +@triton.jit +def tt_reduce_3d_0_1(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, dim: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + + x = tl.load(in_ptr + idx) + + tmp = tl.reduce(x, 0, _reduce_combine) + ret = tl.reduce(tmp, 0, _reduce_combine) + oidx = zidx + + tl.store(out_ptr + oidx, ret) + + +@triton.jit +def tt_reduce_3d_0_2(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, dim: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + + x = tl.load(in_ptr + idx) + + tmp = tl.reduce(x, 0, _reduce_combine) + ret = tl.reduce(tmp, 1, _reduce_combine) + oidx = yidx + + tl.store(out_ptr + oidx, ret) + + +@triton.jit +def tt_reduce_3d_1_2(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, dim: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + + x = tl.load(in_ptr + idx) + + tmp = tl.reduce(x, 1, _reduce_combine) + ret = tl.reduce(tmp, 1, _reduce_combine) + oidx = xidx + + tl.store(out_ptr + oidx, ret) + + +def is_legal_combine(shape, dims): + return dims is None or (len(shape) == 3) or \ + (len(dims) == 1 and dims[0] < len(shape)) + + +dims_map = {(0, 1): tt_reduce_3d_0_1, (1, 2): tt_reduce_3d_1_2, (0, 2): tt_reduce_3d_0_2} + +shape_map = { + 1: {"append_shape": (1, 1), "func": tt_reduce_1d}, 2: {"append_shape": (1, ), "func": tt_reduce_2d}, 3: + {"append_shape": (), "func": tt_reduce_3d} +} + + +def reduce_check_ub_mem_overflow(dtype, shape): + dtype_size = get_dtype_size(dtype) + if (dtype == "int8" or dtype == "bool") and dtype_size * math.prod(shape) >= (TestUtils.ub_size / 20): + pytest.skip("dtype:{dtype} shape:{shape} mem overflow, skipping.") + elif dtype_size * math.prod(shape) >= (TestUtils.ub_size / 6): + pytest.skip("dtype:{dtype} shape:{shape} mem overflow, skipping.") + + +@pytest.mark.parametrize('shape', random.sample(TestUtils.full_shape, 5)) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('dims', [None, (0, ), (1, ), (2, ), (0, 1), (1, 2), (0, 2)]) +def test_reduce(dtype, shape, dims): + if not is_legal_combine(shape, dims): + return + + torch.manual_seed(0) + x = test_common.generate_tensor(shape, dtype).npu() + grid = (1, 1, 1) + + y_ref = torch_reduce(x, dims) + y_cal = torch.empty(y_ref.shape, dtype=eval('torch.' + dtype), device="npu") + + if dims is None: + reduce_check_ub_mem_overflow(dtype, shape) + append_shape, tt_kernel = shape_map[len(shape)]["append_shape"], shape_map[len(shape)]["func"] + xnumel, ynumel, znumel = shape + append_shape + XB, YB, ZB = xnumel, ynumel, znumel + if len(shape) == 1: + tt_reduce_1d_dim_none[1, 1, 1](x, y_cal, xnumel, ynumel, znumel, XB, YB, ZB, dims) + if len(shape) == 2: + tt_reduce_2d_dim_none[1, 1, 1](x, y_cal, xnumel, ynumel, znumel, XB, YB, ZB, dims) + if len(shape) == 3: + tt_reduce_3d_dim_none[1, 1, 1](x, y_cal, xnumel, ynumel, znumel, XB, YB, ZB, dims) + + test_common.validate_cmp(dtype, y_cal, y_ref) + + elif len(dims) == 1: # 1d reduce, 1-3d shape + append_shape, tt_kernel = shape_map[len(shape)]["append_shape"], shape_map[len(shape)]["func"] + xnumel, ynumel, znumel = shape + append_shape + XB, YB, ZB = xnumel, ynumel, znumel + if (len(shape) == 2) and (x.numel() * x.element_size() > 8192): + if dims[0] == 0: + grid = (1, ynumel, 1) + YB = 1 + else: + grid = (xnumel, 1, 1) + XB = 1 + tt_kernel[grid](x, y_cal, xnumel, ynumel, znumel, XB, YB, ZB, dims[0]) + test_common.validate_cmp(dtype, y_cal, y_ref) + else: # 3d shape, 2d reduce + tt_kernel = dims_map[dims] + xnumel, ynumel, znumel = shape + XB, YB, ZB = xnumel, ynumel, znumel + + tt_kernel[grid](x, y_cal, xnumel, ynumel, znumel, XB, YB, ZB, dims[0]) + test_common.validate_cmp(dtype, y_cal, y_ref) + + +@triton.jit +def triton_reduce_multi_d(in_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr, DIMS: tl.constexpr, DIM: tl.constexpr, REDUCE_NUMEL: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if DIMS > 1: + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if DIMS > 2: + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if DIMS > 3: + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if DIMS > 4: + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + x = tl.load(in_ptr + offsets) + + if DIM is not None: + ret = tl.reshape(tl.reduce(x, DIM, _reduce_combine), REDUCE_NUMEL) + o_offsets = tl.arange(0, REDUCE_NUMEL) + tl.store(out_ptr + o_offsets, ret) + else: + ret = tl.reduce(x, DIM, _reduce_combine) + tl.store(out_ptr, ret) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (4, 2, 8, 4), + (4, 3, 8, 1), +]) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('dims', [None, (0, ), (1, ), (2, ), (3, )]) +def test_reduce_4d(dtype, shape, dims): + torch.manual_seed(0) + + x = test_common.generate_tensor(shape, dtype).npu() + dim = dims[0] if dims is not None else None + + y_ref = torch_reduce(x, dim) + y_cal = torch.empty(y_ref.shape, dtype=eval('torch.' + dtype), device="npu") + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + reduce_numel = math.prod(triton_shape) // triton_shape[dim] if dim is not None else None + grid = (1, ) + triton_reduce_multi_d[grid](x, y_cal, *triton_shape, len(shape), dim, reduce_numel) + test_common.validate_cmp(dtype, y_cal, y_ref) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (2, 4, 2, 8, 4), + (3, 4, 2, 8, 1), +]) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('dims', [None, (0, ), (1, ), (2, ), (3, ), (4, )]) +def test_reduce_5d(dtype, shape, dims): + torch.manual_seed(0) + + x = test_common.generate_tensor(shape, dtype).npu() + dim = dims[0] if dims is not None else None + + y_ref = torch_reduce(x, dim) + y_cal = torch.empty(y_ref.shape, dtype=eval('torch.' + dtype), device="npu") + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + reduce_numel = math.prod(triton_shape) // triton_shape[dim] if dim is not None else None + grid = (1, ) + triton_reduce_multi_d[grid](x, y_cal, *triton_shape, len(shape), dim, reduce_numel) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/generalization_cases/test_relu.py b/third_party/ascend/unittest/generalization_cases/test_relu.py new file mode 100644 index 0000000000..21880163bf --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_relu.py @@ -0,0 +1,66 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common +import triton.language.extra.ascend.libdevice as libdevice +from test_common import TestUtils +import math + + +def torch_relu(x0, x1): + res = x0 + torch.relu(x1) + return res + + +@triton.jit +def triton_relu(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + x_index = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = x_index < xnumel + tmp0 = tl.load(in_ptr0 + x_index, xmask) + tmp1 = tl.load(in_ptr1 + x_index, xmask) + tmp2 = tmp0 + libdevice.relu(tmp1) + tl.store(out_ptr0 + x_index, tmp2, xmask) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', ['float32', 'float16']) +def test_relu(dtype, shape): + # 生成数据 + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + + numel = x0.numel() + ncore = 1 if numel <= 32 else 32 + xblock = math.ceil(numel / ncore) + xblock_sub = numel if numel <= ncore else math.ceil(numel / ncore) + + # torch结果 + torch_res = torch_relu(x0, x1) + # triton结果 + triton_res = test_common.generate_tensor(shape, dtype).npu() + triton_relu[ncore, 1, 1](x0, x1, triton_res, x0.numel(), xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/generalization_cases/test_rshift_op.py b/third_party/ascend/unittest/generalization_cases/test_rshift_op.py new file mode 100644 index 0000000000..33b6fffecd --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_rshift_op.py @@ -0,0 +1,184 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import pytest +import triton +import triton.language as tl +import time +import torch +import torch_npu +import test_common +from test_common import TestUtils + + +@triton.jit +def triton_rshift_1d(in_ptr0, out_ptr0, L: tl.constexpr): + lblk_idx = tl.arange(0, L) + idx = lblk_idx[:] + x0 = tl.load(in_ptr0 + idx) + ret = x0 >> 2 + odx = lblk_idx[:] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_rshift_2d(in_ptr0, out_ptr0, M: tl.constexpr, N: tl.constexpr): + moffs = tl.program_id(0) * M + mblk_idx = tl.arange(0, M) + moffs + nblk_idx = tl.arange(0, N) + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + x0 = tl.load(in_ptr0 + idx) + ret = x0 >> 2 + odx = mblk_idx[:, None] * N + nblk_idx[None, :] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_rshift_3d(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + loffs = tl.program_id(0) * L + lblk_idx = tl.arange(0, L) + loffs + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x0 = tl.load(in_ptr0 + idx) + ret = x0 >> 2 + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_rshift_4d_5d(x_ptr, output_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = x_val >> 2 + tl.store(output_ptr + offsets, ret, mask=masks) + + +dtype_mapping = { + 'int8': (torch.int8), + 'int16': (torch.int16), + 'int32': (torch.int32), + 'uint32': (torch.uint32), + 'int64': (torch.int64), + 'float16': (torch.float16), + 'float32': (torch.float32), + 'bfloat16': (torch.bfloat16), + 'bool': (torch.bool), +} + +typelist = [ + 'int8', + 'int16', + 'int32', + 'int64', +] + + +# @pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('shape', TestUtils.test_shape1_2_3d) +@pytest.mark.parametrize('sigtype', typelist) +def test_lshift(sigtype, shape): + dtype = dtype_mapping[sigtype] + x0 = test_common.generate_tensor(shape=shape, dtype=sigtype).npu() + # ncore, xblock, xblock_sub = 2, 32768, 1024 + y_ref = x0 >> 2 + output = torch.zeros(shape, dtype=dtype).npu() + if len(shape) == 3: + shape0 = shape[0] + shape1 = shape[1] + shape2 = shape[2] + if x0.numel() * x0.element_size() >= 1024: + grid = (shape0, 1, 1) + shape0 = 1 + else: + grid = (1, 1, 1) + triton_rshift_3d[grid](x0, output, shape0, shape1, shape2) + if len(shape) == 2: + shape0 = shape[0] + shape1 = shape[1] + if x0.numel() * x0.element_size() >= 1024: + grid = (shape0, 1, 1) + shape0 = 1 + else: + grid = (1, 1, 1) + triton_rshift_2d[grid](x0, output, shape0, shape1) + if len(shape) == 1: + triton_rshift_1d[1, 1, 1](x0, output, shape[0]) + test_common.validate_cmp(sigtype, output, y_ref) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64']) +def test_rshift_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + + output = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = x >> 2 + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_rshift_4d_5d[grid](x, output, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) + + +invalid_types = [ + 'float16', + 'float32', + 'bfloat16', +] + + +@pytest.mark.parametrize("sigtype", invalid_types) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "unexpected type") +def test_invalid_types(sigtype): + N = 32 + x = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + output = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + + triton_rshift_1d[1, 1, 1](x, output, N) diff --git a/third_party/ascend/unittest/generalization_cases/test_scalar_tensor.py b/third_party/ascend/unittest/generalization_cases/test_scalar_tensor.py new file mode 100644 index 0000000000..defb936aa7 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_scalar_tensor.py @@ -0,0 +1,82 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import numpy as np +import torch +import pytest +import test_common + + +def torch_(x0, x1, op_type): + if op_type == 'mul': + return torch.tensor(x0 * x1) + elif op_type == 'lshift': + return torch.tensor(x0 << x1) + elif op_type == 'eq': + return torch.tensor(x0 == x1) + else: + raise TypeError('Invalid op_type') + + +@triton.jit +def scalar_mul(out_ptr0, val0: tl.constexpr, val1: tl.constexpr): + scalar0 = tl.core.tensor(val0, tl.core.block_type(tl.float32, [])) + scalar1 = tl.core.tensor(val1, tl.core.block_type(tl.float32, [])) + ret = scalar0 * scalar1 + tl.store(out_ptr0, ret) + + +@triton.jit +def scalar_lshift(out_ptr0, val0: tl.constexpr, val1: tl.constexpr): + scalar0 = tl.core.tensor(val0, tl.core.block_type(tl.int32, [])) + scalar1 = tl.core.tensor(val1, tl.core.block_type(tl.int32, [])) + ret = scalar0 << scalar1 + tl.store(out_ptr0, ret) + + +@triton.jit +def scalar_eq(out_ptr0, val0: tl.constexpr, val1: tl.constexpr): + scalar0 = tl.core.tensor(val0, tl.core.block_type(tl.int16, [])) + scalar1 = tl.core.tensor(val1, tl.core.block_type(tl.int16, [])) + ret = scalar0 == scalar1 + tl.store(out_ptr0, ret) + + +@pytest.mark.parametrize('param_list', [ + ['float32', 'mul', (1, ), 3.14, 6.66], + ['int32', 'lshift', (1, ), 6, 7], + ['bool', 'eq', (1, ), 5, 5], +]) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "0d block_type is forbidden") +def test_case(param_list): + dtype, op_type, shape, lval, rval = param_list + ans = torch_(lval, rval, op_type) + ret = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + + if op_type == 'mul': + scalar_mul[1, 1, 1](ret, lval, rval) + elif op_type == 'lshift': + scalar_lshift[1, 1, 1](ret, lval, rval) + elif op_type == 'eq': + scalar_eq[1, 1, 1](ret, lval, rval) + + test_common.validate_cmp(dtype, ans, ret) diff --git a/third_party/ascend/unittest/generalization_cases/test_softmax_mindspore.py b/third_party/ascend/unittest/generalization_cases/test_softmax_mindspore.py new file mode 100644 index 0000000000..9efd0fbcd0 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_softmax_mindspore.py @@ -0,0 +1,99 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import os +import mindspore +import triton +import triton.language as tl +import numpy as np +from triton.runtime import driver +import pytest + +pytestmark = pytest.mark.backend("mindspore") + + +@triton.jit +def softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride, n_rows, n_cols, + BLOCK_SIZE: tl.constexpr): + # starting row of the program + row_start = tl.program_id(0) + row_step = tl.num_programs(0) + for row_idx in tl.range(row_start, n_rows, row_step): + # The stride represents how much we need to increase the pointer to advance 1 row + row_start_ptr = input_ptr + row_idx * input_row_stride + # The block size is the next power of two greater than n_cols, so we can fit each + # row in a single block + col_offsets = tl.arange(0, BLOCK_SIZE) + input_ptrs = row_start_ptr + col_offsets + # Load the row into SRAM, using a mask since BLOCK_SIZE may be > than n_cols + mask = col_offsets < n_cols + row = tl.load(input_ptrs, mask=mask, other=-float('inf')) + # Subtract maximum for numerical stability + row_minus_max = row - tl.max(row, axis=0) + + numerator = tl.exp(row_minus_max) + denominator = tl.sum(numerator, axis=0) + softmax_output = numerator / denominator + # Write back output to DRAM + output_row_start_ptr = output_ptr + row_idx * output_row_stride + output_ptrs = output_row_start_ptr + col_offsets + tl.store(output_ptrs, softmax_output, mask=mask) + + +target = triton.runtime.driver.active.get_current_target() +kernels = {} + + +def softmax(x): + n_rows, n_cols = x.shape + + # The block size of each loop iteration is the smallest power of two greater than the number of columns in `x` + BLOCK_SIZE = triton.next_power_of_2(n_cols) + + # Allocate output + y = mindspore.mint.empty_like(x) + + # pre-compile kernel to get register usage and compute thread occupancy. + kernel, num_programs = kernels.get(BLOCK_SIZE, (None, 0)) + if kernel is None: + num_programs = 32 + kernel = softmax_kernel + kernels[BLOCK_SIZE] = (kernel, num_programs) + + num_programs = min(num_programs, n_rows) + + # Create a number of persistent programs. + kernel[(num_programs, 1, 1)](y, x, x.stride(0), y.stride(0), n_rows, n_cols, BLOCK_SIZE) + return y + + +@pytest.mark.parametrize('param_list', [ + ['float32', (1823, 781)], + ['float16', (1823, 781)], +]) +def test_softmax_mindspore(param_list): + os.environ["TRITON_BACKEND"] = "mindspore" + dtype, shape = param_list + mindspore.set_seed(0) + x = mindspore.ops.randn(shape, dtype=eval('mindspore.' + dtype)) + output_triton = softmax(x) + output_mindspore = mindspore.ops.softmax(x, axis=1) + assert np.allclose(output_triton.asnumpy(), output_mindspore.asnumpy(), rtol=1e-3, atol=1e-3) + del os.environ["TRITON_BACKEND"] diff --git a/third_party/ascend/unittest/generalization_cases/test_sort.py b/third_party/ascend/unittest/generalization_cases/test_sort.py new file mode 100644 index 0000000000..543915ec2b --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_sort.py @@ -0,0 +1,229 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import pytest +import torch +import numpy as np +import triton.language as tl +import test_common +from test_common import TestUtils + + +# ---------------------- +# 1D sort kernel +# ---------------------- +@triton.jit +def sort_kernel_1d(X, Z, M: tl.constexpr, descending: tl.constexpr): + off = tl.arange(0, M) + x = tl.load(X + off) + x = tl.sort(x, descending=descending, dim=0) + tl.store(Z + off, x) + + +@pytest.mark.interpreter +@pytest.mark.parametrize("shape", TestUtils.test_shape1d) +@pytest.mark.parametrize("descending", [False, True]) +@pytest.mark.parametrize("dtype", ["int8", "int16", "float16", "float32", "bfloat16", "bool"]) +def test_sort_1d(shape, descending, dtype): + if dtype == "bool": + x = test_common.generate_tensor(shape, dtype) + np_sorted = np.sort(x) + if descending: + np_sorted = np_sorted[::-1].copy() + torch_res = torch.from_numpy(np_sorted).npu() + else: + x = test_common.generate_tensor(shape, dtype).npu() + torch_res = torch.sort(x, descending=descending)[0] + + x = x.npu() + triton_res = torch.zeros_like(x) + M = x.shape[0] + sort_kernel_1d[(1, )](x, triton_res, M, descending) + assert torch.equal(torch_res, triton_res) + + +# ---------------------- +# 2D sort kernel (split by rows, not cutting M axis) +# ---------------------- +@triton.jit +def sort_kernel_2d(X, Z, N: tl.constexpr, M: tl.constexpr, descending: tl.constexpr): + pid = tl.program_id(0) + offx = tl.arange(0, M) + offy = pid * M + off2d = offx + offy + x = tl.load(X + off2d) + x = tl.sort(x, descending=descending, dim=0) + tl.store(Z + off2d, x) + + +@pytest.mark.interpreter +@pytest.mark.parametrize("shape", TestUtils.test_shape2d) +@pytest.mark.parametrize("descending", [False, True]) +@pytest.mark.parametrize("dtype", ["int8", "int16", "float16", "float32", "bfloat16", "bool"]) +def test_sort_2d(shape, descending, dtype): + if dtype == "bool": + x = test_common.generate_tensor(shape, dtype) + np_sorted = np.sort(x) + if descending: + np_sorted = np_sorted[:, ::-1].copy() + torch_res = torch.from_numpy(np_sorted).npu() + else: + x = test_common.generate_tensor(shape, dtype).npu() + torch_res = torch.sort(x, descending=descending)[0] + + x = x.npu() + triton_res = torch.zeros_like(x) + N, M = x.shape + # 每行一个 block + sort_kernel_2d[(N, )](x, triton_res, N, M, descending) + assert torch.equal(torch_res, triton_res), (torch_res, triton_res) + + +# ---------------------- +# 3D sort kernel (split by D0, D1, not cutting D2) +# ---------------------- +@triton.jit +def sort_kernel_3d(X, Z, D0: tl.constexpr, D1: tl.constexpr, D2: tl.constexpr, descending: tl.constexpr): + pid = tl.program_id(0) + row_id = pid % D1 + batch_id = pid // D1 + + off2 = tl.arange(0, D2) + off1 = row_id * D2 + off0 = batch_id * D1 * D2 + off = off2 + off1 + off0 + + x = tl.load(X + off) + x = tl.sort(x, descending=descending, dim=0) # 一整行排序 + tl.store(Z + off, x) + + +@pytest.mark.interpreter +@pytest.mark.parametrize("shape", TestUtils.test_shape3d) +@pytest.mark.parametrize("descending", [False, True]) +@pytest.mark.parametrize("dtype", ["int8", "int16", "float16", "float32", "bfloat16", "bool"]) +def test_sort_3d(shape, descending, dtype): + if dtype == "bool": + x = test_common.generate_tensor(shape, dtype) + np_sorted = np.sort(x) + if descending: + np_sorted = np_sorted[:, :, ::-1].copy() + torch_res = torch.from_numpy(np_sorted).npu() + else: + x = test_common.generate_tensor(shape, dtype).npu() + torch_res = torch.sort(x, descending=descending)[0] + + x = x.npu() + triton_res = torch.zeros_like(x) + D0, D1, D2 = x.shape + # 每个 (D0,D1) 对应一个 block + sort_kernel_3d[(D0 * D1, )](x, triton_res, D0, D1, D2, descending) + assert torch.equal(torch_res, triton_res), (torch_res, triton_res) + + +# ---------------------- +# 4D sort kernel +# ---------------------- +@triton.jit +def sort_kernel_4d(X, Z, D0: tl.constexpr, D1: tl.constexpr, D2: tl.constexpr, D3: tl.constexpr, + descending: tl.constexpr): + pid = tl.program_id(0) + row_id = pid % D2 + col_id = (pid // D2) % D1 + batch_id = pid // (D1 * D2) + + off3 = tl.arange(0, D3) + off2 = row_id * D3 + off1 = col_id * D2 * D3 + off0 = batch_id * D1 * D2 * D3 + off = off3 + off2 + off1 + off0 + + x = tl.load(X + off) + x = tl.sort(x, descending=descending, dim=0) + tl.store(Z + off, x) + + +@pytest.mark.interpreter +@pytest.mark.parametrize("shape", TestUtils.test_shape4d) +@pytest.mark.parametrize("descending", [False, True]) +@pytest.mark.parametrize("dtype", ["int8", "int16", "float16", "float32", "bfloat16", "bool"]) +def test_sort_4d(shape, descending, dtype): + if dtype == "bool": + x = test_common.generate_tensor(shape, dtype) + np_sorted = np.sort(x) + if descending: + np_sorted = np_sorted[:, :, :, ::-1].copy() + torch_res = torch.from_numpy(np_sorted).npu() + else: + x = test_common.generate_tensor(shape, dtype).npu() + torch_res = torch.sort(x, descending=descending)[0] + + x = x.npu() + triton_res = torch.zeros_like(x) + D0, D1, D2, D3 = x.shape + sort_kernel_4d[(D0 * D1 * D2, )](x, triton_res, D0, D1, D2, D3, descending) + assert torch.equal(torch_res, triton_res) + + +# ---------------------- +# 5D sort kernel +# ---------------------- +@triton.jit +def sort_kernel_5d(X, Z, D0: tl.constexpr, D1: tl.constexpr, D2: tl.constexpr, D3: tl.constexpr, D4: tl.constexpr, + descending: tl.constexpr): + pid = tl.program_id(0) + row_id = pid % D3 + col_id = (pid // D3) % D2 + depth_id = (pid // (D2 * D3)) % D1 + batch_id = pid // (D1 * D2 * D3) + + off4 = tl.arange(0, D4) + off3 = row_id * D4 + off2 = col_id * D3 * D4 + off1 = depth_id * D2 * D3 * D4 + off0 = batch_id * D1 * D2 * D3 * D4 + off = off4 + off3 + off2 + off1 + off0 + + x = tl.load(X + off) + x = tl.sort(x, descending=descending, dim=0) + tl.store(Z + off, x) + + +@pytest.mark.interpreter +@pytest.mark.parametrize("shape", TestUtils.test_shape5d) +@pytest.mark.parametrize("descending", [False, True]) +@pytest.mark.parametrize("dtype", ["int8", "int16", "float16", "float32", "bfloat16", "bool"]) +def test_sort_5d(shape, descending, dtype): + if dtype == "bool": + x = test_common.generate_tensor(shape, dtype) + np_sorted = np.sort(x) + if descending: + np_sorted = np_sorted[:, :, :, :, ::-1].copy() + torch_res = torch.from_numpy(np_sorted).npu() + else: + x = test_common.generate_tensor(shape, dtype).npu() + torch_res = torch.sort(x, descending=descending)[0] + + x = x.npu() + triton_res = torch.zeros_like(x) + D0, D1, D2, D3, D4 = x.shape + sort_kernel_5d[(D0 * D1 * D2 * D3, )](x, triton_res, D0, D1, D2, D3, D4, descending) + assert torch.equal(torch_res, triton_res) diff --git a/third_party/ascend/unittest/generalization_cases/test_sqrt.py b/third_party/ascend/unittest/generalization_cases/test_sqrt.py new file mode 100644 index 0000000000..49055ff811 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_sqrt.py @@ -0,0 +1,174 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging + +import triton +import triton.language as tl +import torch +import numpy as np +import pytest +import test_common +from test_common import TestUtils +import math + + +def torch_sqrt(x0): + res = torch.sqrt(x0) + return res + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.sqrt(X) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_sqrt_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = tl.sqrt(x_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_case2(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.debug(f"output.dtype={output.dtype}") + + ans = torch_sqrt(x) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +invalid_types = [ + 'int8', + 'int16', + 'int32', + 'uint32', + 'int64', + 'bool', +] + + +@pytest.mark.parametrize("dtype", invalid_types) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "Expected dtype") +def test_sqrt_invalid_dtype_case(dtype): + x = test_common.generate_tensor((1, ), dtype).npu() + y = test_common.generate_tensor((1, ), dtype).npu() + z = test_common.generate_tensor((1, ), dtype).npu() + + output = torch.randint(1, (1, ), dtype=eval('torch.' + dtype)).npu() + fn_npu_[1, 1, 1](output, x, y, z, 1, 1, 1, 1, 1, 1) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_sqrt_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch.sqrt(x) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_sqrt_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_sqrt_rn.py b/third_party/ascend/unittest/generalization_cases/test_sqrt_rn.py new file mode 100644 index 0000000000..a1add886b4 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_sqrt_rn.py @@ -0,0 +1,174 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils +import math + + +def torch_sqrt_rn(x0): + tmp = torch.sqrt(x0) + return tmp + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.sqrt_rn(X) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_sqrt_rn_4d_5d(output_ptr, x_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, + STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + ret = tl.sqrt_rn(x_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_case2(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.debug(f"output.dtype={output.dtype}") + + ans = torch_sqrt_rn(x) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +invalid_types = [ + 'int8', + 'int16', + 'int32', + 'uint32', + 'int64', + 'bool', +] + + +@pytest.mark.parametrize("dtype", invalid_types) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "Expected dtype") +def test_sqrt_rn_invalid_dtype_case(dtype): + x = test_common.generate_tensor((1, ), dtype).npu() + y = test_common.generate_tensor((1, ), dtype).npu() + z = test_common.generate_tensor((1, ), dtype).npu() + + output = torch.randint(1, (1, ), dtype=eval('torch.' + dtype)).npu() + fn_npu_[1, 1, 1](output, x, y, z, 1, 1, 1, 1, 1, 1) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16']) +def test_sqrt_rn_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = torch_sqrt_rn(x) + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_sqrt_rn_4d_5d[grid](output, x, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_static_print_and_assert_op.py b/third_party/ascend/unittest/generalization_cases/test_static_print_and_assert_op.py new file mode 100644 index 0000000000..d383648c4c --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_static_print_and_assert_op.py @@ -0,0 +1,73 @@ +import torch +import torch_npu +import triton +import triton.language as tl +import pytest +import test_common + +import os + +os.environ["TRITON_ALWAYS_COMPILE"] = "1" +os.environ["PYTEST_ADDOPTS"] = "-sv" + +shape = (8, ) +XS = 8 +XVALS_INT = [ + 0, + torch.iinfo(torch.int8).min, + torch.iinfo(torch.int8).max, + torch.iinfo(torch.int16).min, + torch.iinfo(torch.int16).max, + torch.iinfo(torch.int32).min, + torch.iinfo(torch.int32).max, + torch.iinfo(torch.int32).max + 1 +] + + +def torch_func(x0, x1): + res = x0 + x1 + return res + + +@triton.jit +def triton_kernel(out_ptr0, in_ptr0, in_ptr1, XBLOCK: tl.constexpr): + idx = tl.arange(0, XBLOCK) + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tl.load(in_ptr1 + idx) + tmp2 = tmp0 + tmp1 + tl.static_print(XBLOCK) + tl.static_print(tmp2) + tl.static_assert(XBLOCK == 8) + tl.store(out_ptr0 + idx, tmp2) + + +def triton_func(x0, x1, XS): + out = torch.empty_like(x0) + triton_kernel[ + 1, + ](out, x0, x1, XS) + return out + + +@pytest.mark.parametrize('sigtype', ['int32', 'int64', 'int16', 'int8', 'float32', 'float16', 'bfloat16']) +def test_static_print_and_assert(capsys, sigtype): + dtype = eval(f"torch.{sigtype}") + x0 = torch.zeros(shape, dtype=dtype).npu() + x1 = torch.ones(shape, dtype=dtype).npu() + for i in range(x1.numel()): + x1[i] = XVALS_INT[i] + torch_ref = torch_func(x0, x1) + triton_cal = triton_func(x0, x1, XS) + captured = capsys.readouterr() + + if sigtype == "float32": + assert "fp32" in captured.out + if sigtype == "float16": + assert "fp16" in captured.out + if sigtype == "bfloat16": + assert "bf16" in captured.out + if "int" in sigtype: + assert sigtype in captured.out + assert "8" in captured.out + + test_common.validate_cmp(sigtype, triton_cal, torch_ref) diff --git a/third_party/ascend/unittest/generalization_cases/test_sum.py b/third_party/ascend/unittest/generalization_cases/test_sum.py new file mode 100644 index 0000000000..2879e5fe78 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_sum.py @@ -0,0 +1,355 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import random +import pytest +import torch +import triton +import triton.language as tl + +import test_common +from test_common import TestUtils, get_dtype_size + + +def torch_sum(x1, dim): + if x1.dtype == torch.float16 or x1.dtype == torch.bfloat16: + res = torch.sum(x1.to(torch.float32), dim=dim, keepdim=False).to(x1.dtype) + else: + res = torch.sum(x1, dim=dim, keepdim=False).to(x1.dtype) + return res + + +@triton.jit +def tt_sum_1d(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, XB: tl.constexpr, + YB: tl.constexpr, ZB: tl.constexpr, dim: tl.constexpr): + idx = tl.arange(0, XB) + x = tl.load(in_ptr + idx) + ret = tl.sum(x, dim) + tl.store(out_ptr + tl.arange(0, 1), ret) + + +@triton.jit +def tt_sum_2d(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, XB: tl.constexpr, + YB: tl.constexpr, ZB: tl.constexpr, dim: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + idx = xidx[:, None] * ynumel + yidx[None, :] + + x = tl.load(in_ptr + idx) + ret = tl.sum(x, dim) + + if dim == 0: + oidx = yidx + else: + oidx = xidx + tl.store(out_ptr + oidx, ret) + + +@triton.jit +def tt_sum_1d_dim_none(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, dim: tl.constexpr): + idx = tl.arange(0, XB) + x = tl.load(in_ptr + idx) + ret = tl.sum(x, dim) + tl.store(out_ptr + tl.arange(0, 1), ret) + + +@triton.jit +def tt_sum_2d_dim_none(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, dim: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + idx = xidx[:, None] * ynumel + yidx[None, :] + + x = tl.load(in_ptr + idx) + ret = tl.sum(x, dim) + + tl.store(out_ptr + tl.arange(0, 1), ret) + + +@triton.jit +def tt_sum_3d_dim_none(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, dim: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + + x = tl.load(in_ptr + idx) + ret = tl.sum(x, dim) + + tl.store(out_ptr, ret) + + +@triton.jit +def tt_sum_3d(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, XB: tl.constexpr, + YB: tl.constexpr, ZB: tl.constexpr, dim: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + + x = tl.load(in_ptr + idx) + ret = tl.sum(x, dim) + + if dim == 0: + oidx = yidx[:, None] * znumel + zidx[None, :] + elif dim == 1: + oidx = xidx[:, None] * znumel + zidx[None, :] + else: + oidx = xidx[:, None] * ynumel + yidx[None, :] + + tl.store(out_ptr + oidx, ret) + + +@triton.jit +def tt_sum_3d_0_1(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, XB: tl.constexpr, + YB: tl.constexpr, ZB: tl.constexpr, dim: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + + x = tl.load(in_ptr + idx) + + tmp = tl.sum(x, 0) + ret = tl.sum(tmp, 0) + oidx = zidx + + tl.store(out_ptr + oidx, ret) + + +@triton.jit +def tt_sum_3d_0_2(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, XB: tl.constexpr, + YB: tl.constexpr, ZB: tl.constexpr, dim: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + + x = tl.load(in_ptr + idx) + + tmp = tl.sum(x, 0) + ret = tl.sum(tmp, 1) + oidx = yidx + + tl.store(out_ptr + oidx, ret) + + +@triton.jit +def tt_sum_3d_1_2(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, XB: tl.constexpr, + YB: tl.constexpr, ZB: tl.constexpr, dim: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + + x = tl.load(in_ptr + idx) + + tmp = tl.sum(x, 1) + ret = tl.sum(tmp, 1) + oidx = xidx + + tl.store(out_ptr + oidx, ret) + + +def is_legal_combine(shape, dims): + return dims is None or (len(shape) == 3) or \ + (len(dims) == 1 and dims[0] < len(shape)) + + +dims_map = {(0, 1): tt_sum_3d_0_1, (1, 2): tt_sum_3d_1_2, (0, 2): tt_sum_3d_0_2} + +shape_map = { + 1: {"append_shape": (1, 1), "func": tt_sum_1d}, 2: {"append_shape": (1, ), "func": tt_sum_2d}, 3: + {"append_shape": (), "func": tt_sum_3d} +} + + +def reduce_check_ub_mem_overflow(dtype, shape): + dtype_size = get_dtype_size(dtype) + if (dtype == "int8" or dtype == "bool") and dtype_size * math.prod(shape) >= (TestUtils.ub_size / 20): + pytest.skip("dtype:{dtype} shape:{shape} mem overflow, skipping.") + elif dtype_size * math.prod(shape) >= (TestUtils.ub_size / 6): + pytest.skip("dtype:{dtype} shape:{shape} mem overflow, skipping.") + + +@pytest.mark.parametrize('shape', random.sample(TestUtils.full_shape, 5)) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('dims', [None, (0, ), (1, ), (2, ), (0, 1), (1, 2), (0, 2)]) +def test_sum(dtype, shape, dims): + if not is_legal_combine(shape, dims): + return + + torch.manual_seed(0) + x = test_common.generate_tensor(shape, dtype).npu() + grid = (1, 1, 1) + + y_ref = torch_sum(x, dims) + y_cal = torch.empty(y_ref.shape, dtype=eval('torch.' + dtype), device="npu") + + if dims is None: + reduce_check_ub_mem_overflow(dtype, shape) + append_shape, tt_kernel = shape_map[len(shape)]["append_shape"], shape_map[len(shape)]["func"] + xnumel, ynumel, znumel = shape + append_shape + XB, YB, ZB = xnumel, ynumel, znumel + if len(shape) == 1: + tt_sum_1d_dim_none[1, 1, 1](x, y_cal, xnumel, ynumel, znumel, XB, YB, ZB, dims) + if len(shape) == 2: + tt_sum_2d_dim_none[1, 1, 1](x, y_cal, xnumel, ynumel, znumel, XB, YB, ZB, dims) + if len(shape) == 3: + tt_sum_3d_dim_none[1, 1, 1](x, y_cal, xnumel, ynumel, znumel, XB, YB, ZB, dims) + + test_common.validate_cmp(dtype, y_cal, y_ref) + + elif len(dims) == 1: # 1d sum, 1-3d shape + append_shape, tt_kernel = shape_map[len(shape)]["append_shape"], shape_map[len(shape)]["func"] + xnumel, ynumel, znumel = shape + append_shape + XB, YB, ZB = xnumel, ynumel, znumel + if (len(shape) == 2) and (x.numel() * x.element_size() > 8192): + if dims[0] == 0: + grid = (1, ynumel, 1) + YB = 1 + else: + grid = (xnumel, 1, 1) + XB = 1 + tt_kernel[grid](x, y_cal, xnumel, ynumel, znumel, XB, YB, ZB, dims[0]) + test_common.validate_cmp(dtype, y_cal, y_ref) + else: # 3d shape, 2d sum + tt_kernel = dims_map[dims] + xnumel, ynumel, znumel = shape + XB, YB, ZB = xnumel, ynumel, znumel + tt_kernel[grid](x, y_cal, xnumel, ynumel, znumel, XB, YB, ZB, dims[0]) + test_common.validate_cmp(dtype, y_cal, y_ref) + + +@triton.jit +def triton_sum_multi_d(in_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr, DIMS: tl.constexpr, DIM: tl.constexpr, REDUCE_NUMEL: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if DIMS > 1: + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if DIMS > 2: + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if DIMS > 3: + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if DIMS > 4: + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + x = tl.load(in_ptr + offsets) + + if DIM is not None: + ret = tl.reshape(tl.sum(x, DIM), REDUCE_NUMEL) + o_offsets = tl.arange(0, REDUCE_NUMEL) + tl.store(out_ptr + o_offsets, ret) + else: + ret = tl.sum(x, DIM) + tl.store(out_ptr, ret) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (4, 2, 8, 4), + (4, 3, 8, 1), +]) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('dims', [None, (0, ), (1, ), (2, ), (3, )]) +def test_sum_4d(dtype, shape, dims): + torch.manual_seed(0) + + x = test_common.generate_tensor(shape, dtype).npu() + dim = dims[0] if dims is not None else None + + y_ref = torch_sum(x, dim) + y_cal = torch.empty(y_ref.shape, dtype=eval('torch.' + dtype), device="npu") + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + reduce_numel = math.prod(triton_shape) // triton_shape[dim] if dim is not None else None + grid = (1, ) + triton_sum_multi_d[grid](x, y_cal, *triton_shape, len(shape), dim, reduce_numel) + test_common.validate_cmp(dtype, y_cal, y_ref) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (2, 4, 2, 8, 4), + (3, 4, 2, 8, 1), +]) +@pytest.mark.parametrize('dtype', TestUtils.full_dtype) +@pytest.mark.parametrize('dims', [None, (0, ), (1, ), (2, ), (3, ), (4, )]) +def test_sum_5d(dtype, shape, dims): + torch.manual_seed(0) + + x = test_common.generate_tensor(shape, dtype).npu() + dim = dims[0] if dims is not None else None + + y_ref = torch_sum(x, dim) + y_cal = torch.empty(y_ref.shape, dtype=eval('torch.' + dtype), device="npu") + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + reduce_numel = math.prod(triton_shape) // triton_shape[dim] if dim is not None else None + grid = (1, ) + triton_sum_multi_d[grid](x, y_cal, *triton_shape, len(shape), dim, reduce_numel) + test_common.validate_cmp(dtype, y_cal, y_ref) + + +@triton.jit +def tt_sum_mask(in_ptr, out_ptr, XB: tl.constexpr): + idx = tl.arange(0, XB) + x = tl.load(in_ptr + idx) + ret = tl.sum(x > 0) + tl.store(out_ptr + tl.arange(0, 1), ret) + + +@pytest.mark.parametrize('xb', [32, 256, 1024]) +def test_sum_i1_counts(xb): + # Summing an i1 mask must count the set lanes. Reducing at the native 1-bit + # width turns `a + b` into a logical or, which saturates to 1. + torch.manual_seed(0) + x = torch.randn(xb).npu() + + y_ref = (x > 0).sum().to(torch.int32).reshape(1) + y_cal = torch.empty(1, dtype=torch.int32, device="npu") + + tt_sum_mask[1, 1, 1](x, y_cal, xb) + assert y_ref.item() > 1, "shape does not exercise the saturation bug" + test_common.validate_cmp('int32', y_cal, y_ref) diff --git a/third_party/ascend/unittest/generalization_cases/test_sum_dim0.py b/third_party/ascend/unittest/generalization_cases/test_sum_dim0.py new file mode 100644 index 0000000000..9ef39548d7 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_sum_dim0.py @@ -0,0 +1,74 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils, get_dtype_size +import math + + +def torch_sum(x0): + res = torch.sum(x0, 0) + return res + + +@triton.jit +def triton_sum(in_ptr0, out_ptr1, xnumel, rnumel, XBLOCK: tl.constexpr, RBLOCK: tl.constexpr, RBLOCK_SUB: tl.constexpr): + xindex = tl.arange(0, XBLOCK) + xmask = xindex[:, None] < xnumel + for roffset_sub in range(0, RBLOCK, RBLOCK_SUB): + rindex = roffset_sub + tl.arange(0, RBLOCK_SUB) + x0 = xindex + r1 = rindex + rmask = rindex < rnumel + tmp0 = tl.load(in_ptr0 + (r1 + (RBLOCK * x0[:, None])), xmask & rmask) + tmp2 = tl.reshape(tmp0, [XBLOCK, RBLOCK_SUB]) + tmp4 = tl.sum(tmp2, 0) + tl.store(out_ptr1 + (rindex), tmp4, rmask) + + +def should_skip_due_to_mem(dtype, shape): + dtype_size = get_dtype_size(dtype) + total_mem = dtype_size * math.prod(shape) + threshold = TestUtils.ub_size / 1.5 + + if total_mem >= threshold: + pytest.skip(f"dtype:{dtype} shape:{shape} mem overflow") + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', ['float32', 'int32']) +def test_case(dtype, shape): + should_skip_due_to_mem(dtype, shape) + x0 = test_common.generate_tensor(shape, dtype).npu() + + rblock = shape[1] + xblock = shape[0] + ncore = 1 #if numel <= 32 else 32 + rblock_sub = rblock #if xblock <= 16 else 16 + RBLOCK_tl = 256 if rblock > 1 else 1 + + y_ref = torch_sum(x0) + y_cal = torch.zeros(shape[1], dtype=eval('torch.' + dtype)).npu() + triton_sum[ncore, 1, 1](x0, y_cal, xblock, rblock, xblock, rblock, rblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/generalization_cases/test_sum_dim1.py b/third_party/ascend/unittest/generalization_cases/test_sum_dim1.py new file mode 100644 index 0000000000..dd304da524 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_sum_dim1.py @@ -0,0 +1,66 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils +import math + + +def torch_sum(x0): + res = torch.sum(x0, 1) + return res + + +@triton.jit +def triton_sum(in_ptr0, out_ptr1, xnumel, rnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr, RBLOCK: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + rindex = tl.arange(0, RBLOCK)[None, :] + rmask = rindex < rnumel + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + xindex = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB) + x0 = xindex + r1 = rindex + xmask = xindex[:, None] < xnumel + xmask_prime = xindex < xnumel + tmp0 = tl.load(in_ptr0 + (r1 + (RBLOCK * x0[:, None])), rmask & xmask) + tmp2 = tl.reshape(tmp0, [XBLOCK_SUB, RBLOCK]) + tmp4 = tl.sum(tmp2, 1) + tl.store(out_ptr1 + (xindex), tmp4, xmask_prime) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'int32']) +def test_case(dtype, shape): + x0 = test_common.generate_tensor(shape, dtype).npu() + + rblock = shape[1] + xblock = shape[0] + ncore = 1 #if numel <= 32 else 32 + xblock_sub = xblock if xblock <= 16 else 16 + RBLOCK_tl = 256 if rblock > 1 else 1 + + y_ref = torch_sum(x0) + y_cal = torch.zeros(shape[:-1], dtype=eval('torch.' + dtype)).npu() + triton_sum[ncore, 1, 1](x0, y_cal, xblock, rblock, xblock, xblock_sub, rblock) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/generalization_cases/test_swizzle2d.py b/third_party/ascend/unittest/generalization_cases/test_swizzle2d.py new file mode 100644 index 0000000000..46bd2ebdf8 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_swizzle2d.py @@ -0,0 +1,72 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import random +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common +from test_common import TestUtils + + +def swizzle2d(size_i, size_j, size_g): + i = torch.arange(0, size_i)[:, None] + j = torch.arange(0, size_j)[None, :] + ij = i * size_j + j + size_gj = size_g * size_j + group_id = ij // size_gj + off_i = group_id * size_g + size_g = torch.min(size_i - off_i, torch.tensor(size_g).expand_as(off_i)) + ij = ij % size_gj + new_i = off_i + ij % size_g + new_j = ij // size_g + ret = new_i * size_i + new_j + return ret + + +@triton.jit +def fn_npu_(out0, out1, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + i = tl.arange(0, XB)[:, None] + j = tl.arange(0, YB)[None, :] + ij = i * YB + j + xx, yy = tl.swizzle2d(i, j, size_i=XB, size_j=YB, size_g=ZB) + + ptr = tl.load(out0) + xx = tl.cast(xx, dtype=ptr.dtype) + yy = tl.cast(yy, dtype=ptr.dtype) + tl.store(out0 + ij, xx) + tl.store(out1 + ij, yy) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64']) +def test_swizzle2d(shape, dtype): + if (shape[0] > 255) or (shape[1] > 255): + return + size_g = random.randint(1, min(shape[0], shape[1])) + ans = swizzle2d(shape[0], shape[1], size_g).to(eval('torch.' + dtype)).npu() + + out0 = test_common.generate_tensor(shape, dtype).npu() + out1 = test_common.generate_tensor(shape, dtype).npu() + fn_npu_[1, 1, 1](out0, out1, shape[0], shape[1], size_g) + triton_ret = out0 * shape[0] + out1 + torch.testing.assert_close(triton_ret, ans) diff --git a/third_party/ascend/unittest/generalization_cases/test_tan.py b/third_party/ascend/unittest/generalization_cases/test_tan.py new file mode 100644 index 0000000000..4d6b6454cb --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_tan.py @@ -0,0 +1,62 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils +import math +import triton.language.extra.ascend.libdevice as libdevice + + +def torch_pointwise(x0): + res = torch.tan(x0) + return res + + +@triton.jit +def triton_tan(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop1 in range(loops1): + x0_prime = offset + (loop1 * XBLOCK_SUB) + base1 + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp2 = libdevice.tan(tmp0) + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', ['float32', 'float16']) +def test_case(dtype, shape): + x0 = test_common.generate_tensor(shape, dtype).npu() + + numel = x0.numel() + ncore = 1 if numel <= 32 else 32 + xblock = math.ceil(numel / ncore) + xblock_sub = numel if numel <= ncore else math.ceil(numel / ncore) + + y_ref = torch_pointwise(x0) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_tan[ncore, 1, 1](x0, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/generalization_cases/test_trans_1d_2d.py b/third_party/ascend/unittest/generalization_cases/test_trans_1d_2d.py new file mode 100644 index 0000000000..d56ec1bbb4 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_trans_1d_2d.py @@ -0,0 +1,114 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils, check_ub_mem_overflow +import math +import logging + + +@triton.jit +def fn_npu_1d(output_ptr, x_ptr, xnumel: tl.constexpr): + idx = tl.arange(0, xnumel) + + X = tl.load(x_ptr + idx) + + ret = tl.trans(X, 0) + + tl.store(output_ptr + idx, ret) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', TestUtils.dtype_list) +def test_trans_1d(shape, dtype): + logging.debug(f'dtype:{dtype} shape:{shape}') + + data_type = eval('torch.' + dtype) + x = torch.randint(low=0, high=2, size=shape, dtype=data_type).npu() + + triton_res = torch.randint(1, shape, dtype=data_type).npu() + torch_res = torch.permute(x, (0, )) + fn_npu_1d[1, 1, 1](triton_res, x, shape[0]) + test_common.validate_cmp(dtype, triton_res, torch_res) + + +@triton.jit +def fn_npu_021(output_ptr, x_ptr, YB: tl.constexpr, ZB: tl.constexpr): + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = yidx[:, None] * ZB + zidx[None, :] + + # XB,YB,1 + X = tl.load(x_ptr + idx) + + ret = tl.trans(X, 1, 0) + + oidx = zidx[:, None] * YB + yidx[None, :] + + tl.store(output_ptr + oidx, ret) + + +bisheng_notsupport_dtype = ['int64'] +tritonascend_notsupport_dtype = ['bool'] +# check_ub_mem_overflow没拦住，在kernel中最大ub占用超过ubsize +mem_overflow_scene = [ + ('bfloat16', (128, 256)), + ('bfloat16', (256, 128)), + ('int8', (741, 256)), + ('int8', (256, 741)), + ('int16', (256, 256)), + ('float16', (256, 256)), + ('bfloat16', (256, 256)), + ('int32', (128, 256)), + ('int32', (256, 128)), + ('float32', (128, 256)), + ('float32', (256, 128)), +] + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', TestUtils.dtype_list) +def test_permute(shape, dtype): + logging.debug(f'dtype:{dtype} shape:{shape}') + if dtype in bisheng_notsupport_dtype or dtype in tritonascend_notsupport_dtype: + return + if (dtype, shape) in mem_overflow_scene: + return + if check_ub_mem_overflow(dtype, shape): + return + YB = shape[0] + ZB = shape[1] + data_type = eval('torch.' + dtype) + x = torch.randint(low=0, high=2, size=(YB, ZB), dtype=data_type).npu() + + triton_res = torch.randint(1, (ZB, YB), dtype=data_type).npu() + torch_res = torch.permute(x, (1, 0)) + fn_npu_021[1, 1, 1](triton_res, x, YB, ZB) + test_common.validate_cmp(dtype, triton_res, torch_res) + + +if __name__ == "__main__": + for shape in [(37, 3)]: + for dtype in TestUtils.dtype_list: + test_permute(shape, dtype) diff --git a/third_party/ascend/unittest/generalization_cases/test_trans_3d.py b/third_party/ascend/unittest/generalization_cases/test_trans_3d.py new file mode 100644 index 0000000000..6f8428e575 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_trans_3d.py @@ -0,0 +1,144 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils, check_ub_mem_overflow +import math +import logging + + +@triton.jit +def fn_npu_102(output_ptr, x_ptr, YB: tl.constexpr, ZB: tl.constexpr, KB: tl.constexpr): + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + kidx = tl.arange(0, KB) + idx = yidx[:, None, None] * ZB * KB + zidx[None, :, None] * KB + kidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.trans(X, 1, 0, 2) + + oidx = zidx[:, None, None] * YB * KB + yidx[None, :, None] * KB + kidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_210(output_ptr, x_ptr, YB: tl.constexpr, ZB: tl.constexpr, KB: tl.constexpr): + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + kidx = tl.arange(0, KB) + idx = yidx[:, None, None] * ZB * KB + zidx[None, :, None] * KB + kidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.trans(X, 2, 1, 0) + + oidx = kidx[:, None, None] * ZB * YB + zidx[None, :, None] * YB + yidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_021(output_ptr, x_ptr, YB: tl.constexpr, ZB: tl.constexpr, KB: tl.constexpr): + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + kidx = tl.arange(0, KB) + idx = yidx[:, None, None] * ZB * KB + zidx[None, :, None] * KB + kidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.trans(X, 0, 2, 1) + + oidx = yidx[:, None, None] * ZB * KB + kidx[None, :, None] * ZB + zidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +bisheng_notsupport_dtype = [] +tritonascend_notsupport_dtype = ['bool'] + + +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +@pytest.mark.parametrize('dtype', TestUtils.dtype_list) +def test_permute_3d(shape, dtype): + logging.debug(f'dtype:{dtype} shape:{shape}') + if dtype in bisheng_notsupport_dtype or dtype in tritonascend_notsupport_dtype: + return + if check_ub_mem_overflow(dtype, shape): + return + + data_type = eval('torch.' + dtype) + x = torch.randint(low=0, high=2, size=shape, dtype=data_type).npu() + + triton_res = torch.empty((shape[1], shape[0], shape[2]), dtype=data_type).npu() + torch_res = torch.permute(x, (1, 0, 2)) + fn_npu_102[1, 1, 1](triton_res, x, shape[0], shape[1], shape[2]) + test_common.validate_cmp(dtype, triton_res, torch_res) + + # not support yet: need bisheng support later + # triton_res = torch.empty((shape[2], shape[1], shape[0]), dtype=data_type).npu() + # torch_res = torch.permute(x, (2, 1, 0)) + # fn_npu_210[1, 1, 1](triton_res, x, shape[0], shape[1], shape[2]) + # test_common.validate_cmp(dtype, triton_res, torch_res) + + triton_res = torch.empty((shape[0], shape[2], shape[1]), dtype=data_type).npu() + torch_res = torch.permute(x, (0, 2, 1)) + fn_npu_021[1, 1, 1](triton_res, x, shape[0], shape[1], shape[2]) + test_common.validate_cmp(dtype, triton_res, torch_res) + + +if __name__ == "__main__": + for shape in [(1, 22, 39)]: + for dtype in TestUtils.dtype_list: + test_permute_3d(shape, dtype) + + +@triton.jit +def fn_npu_102(output_ptr, x_ptr, YB: tl.constexpr, ZB: tl.constexpr, KB: tl.constexpr): + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + kidx = tl.arange(0, KB) + idx = yidx[:, None, None] * ZB * KB + zidx[None, :, None] * KB + kidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.trans(X, 1, 0, 2) + + oidx = (zidx[:, None, None] * YB * KB + yidx[None, :, None] * KB + kidx[None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@pytest.mark.parametrize('sigtype, dtype, XB, YB, ZB', [ + ('bfloat16', torch.bfloat16, 2, 8, 4), + ('uint8', torch.uint8, 1, 256, 16), + ('bool', torch.bool, 1, 1, 2), +]) +def test_permute_3d_u(sigtype, dtype, XB, YB, ZB): + x = test_common.generate_tensor((XB, YB, ZB), sigtype).npu() + triton_res = torch.empty((YB, XB, ZB), dtype=dtype).npu() + torch_res = torch.permute(x, (1, 0, 2)) + fn_npu_102[1, 1, 1](triton_res, x, XB, YB, ZB) + test_common.validate_cmp(sigtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/generalization_cases/test_trans_4d_5d.py b/third_party/ascend/unittest/generalization_cases/test_trans_4d_5d.py new file mode 100644 index 0000000000..8505e974f3 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_trans_4d_5d.py @@ -0,0 +1,223 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils, check_ub_mem_overflow +import math +import logging + + +@triton.jit +def triton_trans_4d( + output_ptr, + x_ptr, + PERM: tl.constexpr, + BLOCK_0: tl.constexpr, + BLOCK_1: tl.constexpr, + BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, + SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, + SHAPE_3: tl.constexpr, + STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, + STRIDE_2: tl.constexpr, + STRIDE_3: tl.constexpr, +): + pid = tl.program_id(0) + tmp0 = tl.arange(0, BLOCK_0)[:, None, None, None] + tmp0_1 = tl.arange(0, BLOCK_0)[None, :, None, None] + tmp1 = tl.arange(0, BLOCK_1)[None, :, None, None] + tmp1_0 = tl.arange(0, BLOCK_1)[:, None, None, None] + tmp1_2 = tl.arange(0, BLOCK_1)[None, None, :, None] + tmp2 = tl.arange(0, BLOCK_2)[None, None, :, None] + tmp2_1 = tl.arange(0, BLOCK_2)[None, :, None, None] + tmp2_3 = tl.arange(0, BLOCK_2)[None, None, None, :] + tmp3 = tl.arange(0, BLOCK_3)[None, None, None, :] + tmp3_2 = tl.arange(0, BLOCK_3)[None, None, :, None] + offsets = pid + tmp0 * STRIDE_0 + tmp1 * STRIDE_1 + tmp2 * STRIDE_2 + tmp3 * STRIDE_3 + masks = (tmp0 < SHAPE_0) & (tmp1 < SHAPE_1) & (tmp2 < SHAPE_2) & (tmp3 < SHAPE_3) + x_val = tl.load(x_ptr + offsets, masks) + + if PERM == 0: # 1, 0, 2, 3 + ret = tl.trans(x_val, (1, 0, 2, 3)) + shape0 = SHAPE_1 + shape1 = SHAPE_0 + shape2 = SHAPE_2 + shape3 = SHAPE_3 + elif PERM == 1: # 0, 2, 1, 3 + ret = tl.trans(x_val, (0, 2, 1, 3)) + shape0 = SHAPE_0 + shape1 = SHAPE_2 + shape2 = SHAPE_1 + shape3 = SHAPE_3 + else: # 0, 1, 3, 2 + ret = tl.trans(x_val, (0, 1, 3, 2)) + shape0 = SHAPE_0 + shape1 = SHAPE_1 + shape2 = SHAPE_3 + shape3 = SHAPE_2 + + s3 = 1 + s2 = s3 * shape3 + s1 = s2 * shape2 + s0 = s1 * shape1 + + if PERM == 0: # 1, 0, 2, 3 + out_offsets = pid + tmp1_0 * s0 + tmp0_1 * s1 + tmp2 * s2 + tmp3 * s3 + out_masks = (tmp1_0 < shape0) & (tmp0_1 < shape1) & (tmp2 < shape2) & (tmp3 < shape3) + elif PERM == 1: # 0, 2, 1, 3 + out_offsets = pid + tmp0 * s0 + tmp2_1 * s1 + tmp1_2 * s2 + tmp3 * s3 + out_masks = (tmp0 < shape0) & (tmp1_2 < shape2) & (tmp2_1 < shape1) & (tmp3 < shape3) + else: # 0, 1, 3, 2 + out_offsets = pid + tmp0 * s0 + tmp1 * s1 + tmp3_2 * s2 + tmp2_3 * s3 + out_masks = (tmp0 < shape0) & (tmp1 < shape1) & (tmp3_2 < shape2) & (tmp2_3 < shape3) + tl.store(output_ptr + out_offsets, ret, mask=out_masks) + + +@triton.jit +def triton_trans_5d(output_ptr, x_ptr, PERM: tl.constexpr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, + BLOCK_2: tl.constexpr, BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, + STRIDE_0: tl.constexpr, STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, + STRIDE_4: tl.constexpr): + pid = tl.program_id(0) + tmp0 = tl.arange(0, BLOCK_0)[:, None, None, None, None] + tmp1 = tl.arange(0, BLOCK_1)[None, :, None, None, None] + tmp2 = tl.arange(0, BLOCK_2)[None, None, :, None, None] + tmp3 = tl.arange(0, BLOCK_3)[None, None, None, :, None] + tmp4 = tl.arange(0, BLOCK_4)[None, None, None, None, :] + + tmp0_1 = tl.arange(0, BLOCK_0)[None, :, None, None, None] + tmp1_0 = tl.arange(0, BLOCK_1)[:, None, None, None, None] + + tmp1_2 = tl.arange(0, BLOCK_1)[None, None, :, None, None] + tmp2_1 = tl.arange(0, BLOCK_2)[None, :, None, None, None] + + tmp2_3 = tl.arange(0, BLOCK_2)[None, None, None, :, None] + tmp3_2 = tl.arange(0, BLOCK_3)[None, None, :, None, None] + + tmp3_4 = tl.arange(0, BLOCK_3)[None, None, None, None, :] + tmp4_3 = tl.arange(0, BLOCK_4)[None, None, None, :, None] + + offsets = pid + tmp0 * STRIDE_0 + tmp1 * STRIDE_1 + tmp2 * STRIDE_2 + tmp3 * STRIDE_3 + tmp4 * STRIDE_4 + masks = (tmp0 < SHAPE_0) & (tmp1 < SHAPE_1) & (tmp2 < SHAPE_2) & (tmp3 < SHAPE_3) & (tmp4 < SHAPE_4) + x_val = tl.load(x_ptr + offsets, masks) + + if PERM == 0: # 1, 0, 2, 3, 4 + ret = tl.trans(x_val, 1, 0, 2, 3, 4) + shape0 = SHAPE_1 + shape1 = SHAPE_0 + shape2 = SHAPE_2 + shape3 = SHAPE_3 + shape4 = SHAPE_4 + elif PERM == 1: # 0, 2, 1, 3, 4 + ret = tl.trans(x_val, 0, 2, 1, 3, 4) + shape0 = SHAPE_0 + shape1 = SHAPE_2 + shape2 = SHAPE_1 + shape3 = SHAPE_3 + shape4 = SHAPE_4 + elif PERM == 2: # 0, 1, 3, 2, 4 + ret = tl.trans(x_val, 0, 1, 3, 2, 4) + shape0 = SHAPE_0 + shape1 = SHAPE_1 + shape2 = SHAPE_3 + shape3 = SHAPE_2 + shape4 = SHAPE_4 + else: # 0, 1, 2, 4, 3 + ret = tl.trans(x_val, 0, 1, 2, 4, 3) + shape0 = SHAPE_0 + shape1 = SHAPE_1 + shape2 = SHAPE_2 + shape3 = SHAPE_4 + shape4 = SHAPE_3 + + s4 = 1 + s3 = s4 * shape4 + s2 = s3 * shape3 + s1 = s2 * shape2 + s0 = s1 * shape1 + + if PERM == 0: # 1, 0, 2, 3, 4 + out_offsets = pid + tmp1_0 * s0 + tmp0_1 * s1 + tmp2 * s2 + tmp3 * s3 + tmp4 * s4 + out_masks = (tmp1_0 < shape0) & (tmp0_1 < shape1) & (tmp2 < shape2) & (tmp3 < shape3) & (tmp4 < shape4) + elif PERM == 1: # 0, 2, 1, 3, 4 + out_offsets = pid + tmp0 * s0 + tmp2_1 * s1 + tmp1_2 * s2 + tmp3 * s3 + tmp4 * s4 + out_masks = (tmp0 < shape0) & (tmp1_2 < shape2) & (tmp2_1 < shape1) & (tmp3 < shape3) & (tmp4 < shape4) + elif PERM == 2: # 0, 1, 3, 2, 4 + out_offsets = pid + tmp0 * s0 + tmp1 * s1 + tmp3_2 * s2 + tmp2_3 * s3 + tmp4 * s4 + out_masks = (tmp0 < shape0) & (tmp1 < shape1) & (tmp3_2 < shape2) & (tmp2_3 < shape3) & (tmp4 < shape4) + else: # 0, 1, 2, 4, 3 + out_offsets = pid + tmp0 * s0 + tmp1 * s1 + tmp2 * s2 + tmp4_3 * s3 + tmp3_4 * s4 + out_masks = (tmp0 < shape0) & (tmp1 < shape1) & (tmp2 < shape2) & (tmp4_3 < shape3) & (tmp3_4 < shape4) + tl.store(output_ptr + out_offsets, ret, mask=out_masks) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +@pytest.mark.parametrize('perm', [0, 1, 2, 3]) # 4d: support 3 mode; 5d: support 4 mode +def test_trans_4d_5d(shape, dtype, perm): + logging.log(logging.DEBUG, f"shape = {shape}") + x = torch.randint(low=0, high=2, size=shape, dtype=eval('torch.' + dtype)).npu() + grid = (1, ) + if len(shape) == 4: + blocks = list(x.size()) + strides = list(x.stride()) + if perm == 0: # 1, 0, 2, 3; exchange axis 0, 1 + output = torch.empty((shape[1], shape[0], shape[2], shape[3]), dtype=eval('torch.' + dtype)).npu() + ans_4d = torch.permute(x, (1, 0, 2, 3)) + triton_trans_4d[grid](output, x, perm, *blocks, *blocks, *strides) + test_common.validate_cmp(dtype, ans_4d, output) + elif perm == 1: # 0, 2, 1, 3; exchange axis 1, 2 + output = torch.empty((shape[0], shape[2], shape[1], shape[3]), dtype=eval('torch.' + dtype)).npu() + ans_4d = torch.permute(x, (0, 2, 1, 3)) + triton_trans_4d[grid](output, x, perm, *blocks, *blocks, *strides) + test_common.validate_cmp(dtype, ans_4d, output) + elif perm == 2: # 0, 1, 3, 2; exchange axis 2, 3 + output = torch.empty((shape[0], shape[1], shape[3], shape[2]), dtype=eval('torch.' + dtype)).npu() + ans_4d = torch.permute(x, (0, 1, 3, 2)) + triton_trans_4d[grid](output, x, perm, *blocks, *blocks, *strides) + test_common.validate_cmp(dtype, ans_4d, output) + else: + pass + else: + blocks = list(x.size()) + strides = list(x.stride()) + + if perm == 0: # 1, 0, 2, 3, 4; exchange axis 0, 1 + output = torch.empty((shape[1], shape[0], shape[2], shape[3], shape[4]), dtype=eval('torch.' + dtype)).npu() + ans_5d = torch.permute(x, (1, 0, 2, 3, 4)) + elif perm == 1: # 0, 2, 1, 3, 4; exchange axis 1, 2 + output = torch.empty((shape[0], shape[2], shape[1], shape[3], shape[4]), dtype=eval('torch.' + dtype)).npu() + ans_5d = torch.permute(x, (0, 2, 1, 3, 4)) + elif perm == 2: # 0, 1, 3, 2, 4; exchange axis 2, 3 + output = torch.empty((shape[0], shape[1], shape[3], shape[2], shape[4]), dtype=eval('torch.' + dtype)).npu() + ans_5d = torch.permute(x, (0, 1, 3, 2, 4)) + else: # 0, 1, 2, 4, 3; exchange axis 3, 4 + output = torch.empty((shape[0], shape[1], shape[2], shape[4], shape[3]), dtype=eval('torch.' + dtype)).npu() + ans_5d = torch.permute(x, (0, 1, 2, 4, 3)) + triton_trans_5d[grid](output, x, perm, *blocks, *blocks, *strides) + test_common.validate_cmp(dtype, ans_5d, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_umulhi.py b/third_party/ascend/unittest/generalization_cases/test_umulhi.py new file mode 100644 index 0000000000..421fc77322 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_umulhi.py @@ -0,0 +1,147 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import triton +import torch +import pytest +import test_common + +import numpy as np +import triton.language as tl +from test_common import TestUtils + + +# inp the two 32 bit signed integers. +@triton.jit +def umulhi_kernel(X, Y, Z, N: tl.constexpr): + offs = tl.arange(0, N) + x = tl.load(X + offs) + y = tl.load(Y + offs) + z = tl.umulhi(x, y) + tl.store(Z + tl.arange(0, N), z) + + +@triton.jit +def triton_umulhi_4d_5d(output_ptr, x_ptr, y_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = tl.umulhi(x_val, y_val) + tl.store(output_ptr + offsets, ret, mask=masks) + + +# accuracy reference +def umulhi32(a, b): + a_64 = a.astype(np.int64) + b_64 = b.astype(np.int64) + product_64 = a_64 * b_64 + # get the high part + result_high_32 = product_64 >> 32 + return result_high_32.astype(np.int32) + + +@pytest.mark.parametrize('dtype', ['int32']) +@pytest.mark.parametrize('shape', TestUtils.full_shape) +def test_case2(dtype, shape): + N = shape[0] + dtypes = eval('torch.' + dtype) + x = torch.randint(low=0, high=2000, size=shape, dtype=dtypes) + y = torch.randint(low=0, high=2000, size=shape, dtype=dtypes) + xx = x.npu() + yy = y.npu() + z_tri = torch.zeros(size=shape, dtype=dtypes).npu() + umulhi_kernel[(1, )](xx, yy, z_tri, N=N) + + xxx = x.numpy() + yyy = y.numpy() + z_ref = umulhi32(xxx, yyy) + z_ref1 = torch.from_numpy(z_ref).npu() + torch.equal(z_tri, z_ref1) + + +invalid_types = [ + 'int8', + 'int16', + 'int64', + 'float16', + 'float32', + 'bfloat16', + 'bool', +] + + +@pytest.mark.parametrize("dtype", invalid_types) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "Expected dtype") +def test_umulhi_invalid_dtype_case(dtype): + x0 = test_common.generate_tensor((1, ), dtype).npu() + x1 = test_common.generate_tensor((1, ), dtype).npu() + + y_cal = torch.zeros((1, ), dtype=eval('torch.' + dtype)).npu() + umulhi_kernel[(1, )](x0, x1, y_cal, 1) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int32']) +def test_umulhi_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = torch.randint(low=0, high=2000, size=shape, dtype=eval('torch.' + dtype)) + y = torch.randint(low=0, high=2000, size=shape, dtype=eval('torch.' + dtype)) + xx = x.npu() + yy = y.npu() + + output = torch.zeros(size=shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + xxx = x.numpy() + yyy = y.numpy() + z = umulhi32(xxx, yyy) + ans = torch.from_numpy(z).npu() + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_umulhi_4d_5d[grid](output, xx, yy, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_where.py b/third_party/ascend/unittest/generalization_cases/test_where.py new file mode 100644 index 0000000000..79345b6f21 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_where.py @@ -0,0 +1,135 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging + +import triton +import triton.language as tl +import torch +import pytest +import test_common +from test_common import TestUtils +import math + + +def torch_pointwise(x0, x1): + res = torch.where(x0 < x1, x0, 1) + return res + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + tmp2 = X < Y + ret = tl.where(tmp2, X, 1) + + tl.store(output_ptr + idx, ret) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['bool', 'float32', 'float16', 'bfloat16', 'int8', 'int16', 'int32', 'int64']) +def test_case2(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + ans = torch_pointwise(x, y) + output = torch.zeros_like(ans) + + if len(shape) == 1: + fn_npu_[1, 1, shape[0]](output, x, y, z, 1, 1, 1, 1, 1, shape[0]) + elif len(shape) == 2: + if shape[0] > shape[1]: + fn_npu_[1, shape[0], 1](output, x, y, z, 1, 1, shape[1], 1, shape[0], shape[1]) + else: + fn_npu_[1, 1, shape[1]](output, x, y, z, 1, shape[0], 1, 1, shape[0], shape[1]) + elif len(shape) == 3: + if max(shape[0], shape[1], shape[2]) == shape[0]: + fn_npu_[shape[0], 1, 1](output, x, y, z, 1, shape[1], shape[2], shape[0], shape[1], shape[2]) + else: + fn_npu_[1, 1, shape[2]](output, x, y, z, shape[0], shape[1], 1, shape[0], shape[1], shape[2]) + else: + fn_npu_[1, 1, 1](output, x, y, z, 1, 1, 1, 1, 1, 1) + + test_common.validate_cmp(dtype, ans, output) + + +@triton.jit +def fn_npu_multi_d(output_ptr, x_ptr, y_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr, DIMS: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if DIMS > 1: + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if DIMS > 2: + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if DIMS > 3: + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if DIMS > 4: + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + X = tl.load(x_ptr + offsets) + Y = tl.load(y_ptr + offsets) + + tmp2 = X < Y + ret = tl.where(tmp2, X, 1) + + tl.store(output_ptr + offsets, ret) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (4, 2, 8, 4), + (2, 4, 2, 8, 1), + (4, 3, 8, 1), + (3, 4, 2, 8, 4), +]) +@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16', 'int8', 'int16', 'int32', 'int64']) +def test_case_4d_5d(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + + ans = torch_pointwise(x, y) + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + grid = (1, ) + fn_npu_multi_d[grid](output, x, y, *triton_shape, len(shape)) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/generalization_cases/test_xor.py b/third_party/ascend/unittest/generalization_cases/test_xor.py new file mode 100644 index 0000000000..fe696552af --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_xor.py @@ -0,0 +1,174 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common +from test_common import TestUtils +import math + + +def torch_xor(x0, x1): + return x0 ^ x1 + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = X ^ Y + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def triton_xor_4d_5d(output_ptr, x_ptr, y_ptr, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, SHAPE_1: tl.constexpr, + SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, + STRIDE_1: tl.constexpr, STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + x_val = tl.load(x_ptr + offsets, masks) + y_val = tl.load(y_ptr + offsets, masks) + ret = x_val ^ y_val + tl.store(output_ptr + offsets, ret, mask=masks) + + +@pytest.mark.parametrize('shape', TestUtils.full_shape) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'bool']) +def test_case2(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + logging.debug(f"output.dtype={output.dtype}") + + ans = torch_xor(x, y) + + if len(shape) == 1: + XB = 1 + xnumel = 1 + YB = 1 + ynumel = 1 + ZB = shape[0] + znumel = shape[0] + elif len(shape) == 2: + XB = 1 + xnumel = 1 + YB = shape[0] + ynumel = shape[0] + ZB = shape[1] + znumel = shape[1] + else: + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + if x.numel() * x.element_size() >= 8192: + grid = (1, 1, ZB) + ZB = 1 + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape4d + TestUtils.test_shape5d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'bool']) +def test_xor_4d_5d(shape, dtype): + logging.log(logging.DEBUG, f"shape = {shape}") + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + + logging.log(logging.DEBUG, f"output.dtype={output.dtype}") + + ans = x ^ y + + blocks = list(x.size()) + strides = list(x.stride()) + while len(blocks) < 5: + blocks.append(1) + strides.append(1) + + grid = (1, ) + triton_xor_4d_5d[grid](output, x, y, *blocks, *blocks, *strides) + + test_common.validate_cmp(dtype, ans, output) + + +invalid_types = [ + 'float16', + 'float32', + 'bfloat16', +] + + +@pytest.mark.parametrize("sigtype", invalid_types) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "unexpected type") +def test_invalid_types(sigtype): + N = 32 + x = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + y = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + z = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + output = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + + fn_npu_[1, 1, 1](output, x, y, z, 32, 1, 1, 32, 1, 1) diff --git a/third_party/ascend/unittest/generalization_cases/test_xorsum.py b/third_party/ascend/unittest/generalization_cases/test_xorsum.py new file mode 100644 index 0000000000..633db01c15 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_xorsum.py @@ -0,0 +1,281 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common +import functools +from test_common import TestUtils, check_ub_mem_overflow, get_dtype_size + + +# <<<<<<< test_xorsum_1d +def torch_xorsum(tensor, dim=None, keepdim=False): + if dim is None: + result = tensor.flatten()[0] + for x in tensor.flatten()[1:]: + result = result ^ x + return result + else: + assert dim < tensor.dim(), f"Invalid dim {dim} for tensor shape {tensor.shape}" + result = tensor.select(dim, 0) + for i in range(1, tensor.size(dim)): + result = result ^ tensor.select(dim, i) + if keepdim: + result = result.unsqueeze(dim) + return result + + +@triton.jit +def triton_xorsum_1d(in_ptr0, out_ptr1, xnumel, XBLOCK: tl.constexpr): + xoffset = tl.program_id(0) + tl.arange(0, XBLOCK) + tmp0 = tl.load(in_ptr0 + xoffset, None) + tmp4 = tl.xor_sum(tmp0, 0) + tl.store(out_ptr1, tmp4, None) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape1d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'bool']) +def test_xorsum_1d(dtype, shape): + if check_ub_mem_overflow(dtype, shape): + return + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty(1, dtype=eval("torch." + dtype)).npu() + numel = shape[0] + triton_xorsum_1d[1, 1, 1](x0, triton_res, numel, numel) + torch_res = torch_xorsum(x0, dim=0, keepdim=True) + test_common.validate_cmp(dtype, triton_res, torch_res) + + +# >>>>>>> test_xorsum_1d + + +# <<<<<<< test_xorsum_2d +@triton.jit +def triton_xorsum_2d(in_ptr0, out_ptr0, dim: tl.constexpr, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, + NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, MNUMEL) + nblk_idx = tl.arange(0, NNUMEL) + mmask = mblk_idx < M + nmask = nblk_idx < N + mask = (mmask[:, None]) & (nmask[None, :]) + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + x = tl.load(in_ptr0 + idx, mask=mask, other=-float('inf')) + tmp4 = tl.xor_sum(x, dim) + if dim == 0: + tl.store(out_ptr0 + tl.arange(0, N), tmp4, None) + else: + tl.store(out_ptr0 + tl.arange(0, M), tmp4, None) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape2d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'bool']) +@pytest.mark.parametrize('dim', [0, 1]) +def test_xorsum_2d(dtype, shape, dim): + dtype_size = get_dtype_size(dtype) + if dtype in ['int8', 'int16', 'int32', 'int64']: + if dtype_size * math.prod(shape) >= (TestUtils.ub_size / 3): + pytest.skip(f"dtype:{dtype} shape:{shape} mem overflow") + elif dtype in ['bool']: + if dtype_size * math.prod(shape) >= (TestUtils.ub_size / 5): + pytest.skip(f"dtype:{dtype} shape:{shape} mem overflow") + shapex, shapey = shape + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty([ + shape[1 - dim], + ], dtype=eval("torch." + dtype)).npu() + triton_xorsum_2d[1, 1, 1](x0, triton_res, dim, shapex, shapey, shapex, shapey) + torch_res = torch_xorsum(x0, dim=dim, keepdim=False) + test_common.validate_cmp(dtype, triton_res, torch_res) + + +# >>>>>>> test_xorsum_2d + + +# <<<<<<< test_xorsum_3d +def torch_xorsum_3d(x0, no_reduce_dim): + inp = x0 if x0.device == "cpu" else x0.cpu() + if no_reduce_dim == 0: + return torch_xorsum(torch_xorsum(inp, 1), 1).npu() + elif no_reduce_dim == 1: + return torch_xorsum(torch_xorsum(inp, 0), 1).npu() + elif no_reduce_dim == 2: + return torch_xorsum(torch_xorsum(inp, 0), 0).npu() + else: + assert False, f"no reduce dim not right, no_reduce_dim = {no_reduce_dim}" + + +@triton.jit +def triton_xorsum_3d_0_1(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + x = tl.load(in_ptr + idx) + tmp = tl.xor_sum(x, 0) + ret = tl.xor_sum(tmp, 0) + oidx = zidx + tl.store(out_ptr + oidx, ret) + + +@triton.jit +def triton_xorsum_3d_0_2(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + x = tl.load(in_ptr + idx) + tmp = tl.xor_sum(x, 0) + ret = tl.xor_sum(tmp, 1) + oidx = yidx + tl.store(out_ptr + oidx, ret) + + +@triton.jit +def triton_xorsum_3d_1_2(in_ptr, out_ptr, xnumel: tl.constexpr, ynumel: tl.constexpr, znumel: tl.constexpr, + XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * ynumel * znumel + yidx[None, :, None] * znumel + zidx[None, None, :] + x = tl.load(in_ptr + idx) + tmp = tl.xor_sum(x, 1) + ret = tl.xor_sum(tmp, 1) + oidx = xidx + tl.store(out_ptr + oidx, ret) + + +def triton_xorsum_3d(in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB, no_reduce_dim): + if no_reduce_dim == 0: + triton_xorsum_3d_1_2[1, 1, 1](in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB) + elif no_reduce_dim == 1: + triton_xorsum_3d_0_2[1, 1, 1](in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB) + elif no_reduce_dim == 2: + triton_xorsum_3d_0_1[1, 1, 1](in_ptr, out_ptr, xnumel, ynumel, znumel, XB, YB, ZB) + + +@pytest.mark.parametrize('shape', TestUtils.test_shape3d) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'bool']) +@pytest.mark.parametrize('no_reduce_dim', [0, 1, 2]) +def test_xorsum_3d(dtype, shape, no_reduce_dim): + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty([ + shape[no_reduce_dim], + ], dtype=eval("torch." + dtype)).npu() + triton_xorsum_3d(x0, triton_res, shape[0], shape[1], shape[2], shape[0], shape[1], shape[2], no_reduce_dim) + torch_res = torch_xorsum_3d(x0, no_reduce_dim) + test_common.validate_cmp(dtype, triton_res, torch_res) + + +# >>>>>>> test_xorsum_3d + + +# <<<<<<< test_xorsum_4d +@triton.jit +def triton_xorsum_multi_d(in_ptr, out_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr, DIMS: tl.constexpr, DIM: tl.constexpr, REDUCE_NUMEL: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if DIMS > 1: + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if DIMS > 2: + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if DIMS > 3: + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if DIMS > 4: + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + x = tl.load(in_ptr + offsets) + + if DIM is not None: + ret = tl.reshape(tl.xor_sum(x, DIM), REDUCE_NUMEL) + o_offsets = tl.arange(0, REDUCE_NUMEL) + tl.store(out_ptr + o_offsets, ret) + else: + ret = tl.xor_sum(x, DIM) + tl.store(out_ptr, ret) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (4, 2, 8, 4), + (4, 3, 8, 1), +]) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'bool']) +@pytest.mark.parametrize('dim', [0, 1, 2, 3]) +def test_xorsum_4d(dtype, shape, dim): + dtype_size = get_dtype_size(dtype) + if dtype in ['int8', 'int16', 'int32', 'int64']: + if dtype_size * math.prod(shape) >= (TestUtils.ub_size / 3): + print(f"dtype:{dtype} shape:{shape} mem overflow") + return + + x0 = test_common.generate_tensor(shape, dtype).npu() + torch_res = torch_xorsum(x0, dim=dim, keepdim=False) + triton_res = torch.empty_like(torch_res, dtype=eval("torch." + dtype)).npu() + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + reduce_numel = math.prod(triton_shape) // triton_shape[dim] if dim is not None else None + grid = (1, ) + triton_xorsum_multi_d[grid](x0, triton_res, *triton_shape, len(shape), dim, reduce_numel) + test_common.validate_cmp(dtype, triton_res, torch_res) + + +# >>>>>>> test_xorsum_4d + + +# <<<<<<< test_xorsum_5d +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('shape', [ + (2, 4, 2, 8, 4), + (3, 4, 2, 8, 1), +]) +@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'bool']) +@pytest.mark.parametrize('dim', [0, 1, 2, 3, 4]) +def test_xorsum_5d(dtype, shape, dim): + dtype_size = get_dtype_size(dtype) + if dtype in ['int8', 'int16', 'int32', 'int64']: + if dtype_size * math.prod(shape) >= (TestUtils.ub_size / 3): + print(f"dtype:{dtype} shape:{shape} mem overflow") + return + + x0 = test_common.generate_tensor(shape, dtype).npu() + torch_res = torch_xorsum(x0, dim=dim, keepdim=False) + triton_res = torch.empty_like(torch_res, dtype=eval("torch." + dtype)).npu() + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + reduce_numel = math.prod(triton_shape) // triton_shape[dim] if dim is not None else None + grid = (1, ) + triton_xorsum_multi_d[grid](x0, triton_res, *triton_shape, len(shape), dim, reduce_numel) + test_common.validate_cmp(dtype, triton_res, torch_res) + + +# >>>>>>> test_xorsum_5d + +if __name__ == "__main__": + test_xorsum_3d('int8', (3, 3, 3), 0) diff --git a/third_party/ascend/unittest/generalization_cases/test_zeros_op.py b/third_party/ascend/unittest/generalization_cases/test_zeros_op.py new file mode 100644 index 0000000000..7e5304d153 --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_zeros_op.py @@ -0,0 +1,534 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import pytest +import random +import torch +import torch_npu +import triton +import triton.language as tl + +import test_common +from test_common import TestUtils, check_ub_mem_overflow + + +@triton.jit +def fn_npu_int8_3d(output_ptr, X: tl.constexpr, Y: tl.constexpr, Z: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xidx = tl.arange(0, XNUMEL) + yidx = tl.arange(0, YNUMEL) + zidx = tl.arange(0, ZNUMEL) + Xmask = xidx < X + Ymask = yidx < Y + Zmask = zidx < Z + mask = (Xmask[:, None, None]) & (Ymask[None, :, None]) & (Zmask[None, None, :]) + ret = tl.zeros((XNUMEL, YNUMEL, ZNUMEL), dtype=tl.int8) + oidx = xidx[:, None, None] * Y * Z + yidx[None, :, None] * Z + zidx[None, None, :] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_int16_3d(output_ptr, X: tl.constexpr, Y: tl.constexpr, Z: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xidx = tl.arange(0, XNUMEL) + yidx = tl.arange(0, YNUMEL) + zidx = tl.arange(0, ZNUMEL) + Xmask = xidx < X + Ymask = yidx < Y + Zmask = zidx < Z + mask = (Xmask[:, None, None]) & (Ymask[None, :, None]) & (Zmask[None, None, :]) + ret = tl.zeros((XNUMEL, YNUMEL, ZNUMEL), dtype=tl.int16) + oidx = xidx[:, None, None] * Y * Z + yidx[None, :, None] * Z + zidx[None, None, :] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_int32_3d(output_ptr, X: tl.constexpr, Y: tl.constexpr, Z: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xidx = tl.arange(0, XNUMEL) + yidx = tl.arange(0, YNUMEL) + zidx = tl.arange(0, ZNUMEL) + Xmask = xidx < X + Ymask = yidx < Y + Zmask = zidx < Z + mask = (Xmask[:, None, None]) & (Ymask[None, :, None]) & (Zmask[None, None, :]) + ret = tl.zeros((XNUMEL, YNUMEL, ZNUMEL), dtype=tl.int32) + oidx = xidx[:, None, None] * Y * Z + yidx[None, :, None] * Z + zidx[None, None, :] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_int64_3d(output_ptr, X: tl.constexpr, Y: tl.constexpr, Z: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xidx = tl.arange(0, XNUMEL) + yidx = tl.arange(0, YNUMEL) + zidx = tl.arange(0, ZNUMEL) + Xmask = xidx < X + Ymask = yidx < Y + Zmask = zidx < Z + mask = (Xmask[:, None, None]) & (Ymask[None, :, None]) & (Zmask[None, None, :]) + ret = tl.zeros((XNUMEL, YNUMEL, ZNUMEL), dtype=tl.int64) + oidx = xidx[:, None, None] * Y * Z + yidx[None, :, None] * Z + zidx[None, None, :] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_fp16_3d(output_ptr, X: tl.constexpr, Y: tl.constexpr, Z: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xidx = tl.arange(0, XNUMEL) + yidx = tl.arange(0, YNUMEL) + zidx = tl.arange(0, ZNUMEL) + Xmask = xidx < X + Ymask = yidx < Y + Zmask = zidx < Z + mask = (Xmask[:, None, None]) & (Ymask[None, :, None]) & (Zmask[None, None, :]) + ret = tl.zeros((XNUMEL, YNUMEL, ZNUMEL), dtype=tl.float16) + oidx = xidx[:, None, None] * Y * Z + yidx[None, :, None] * Z + zidx[None, None, :] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_fp32_3d(output_ptr, X: tl.constexpr, Y: tl.constexpr, Z: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xidx = tl.arange(0, XNUMEL) + yidx = tl.arange(0, YNUMEL) + zidx = tl.arange(0, ZNUMEL) + Xmask = xidx < X + Ymask = yidx < Y + Zmask = zidx < Z + mask = (Xmask[:, None, None]) & (Ymask[None, :, None]) & (Zmask[None, None, :]) + ret = tl.zeros((XNUMEL, YNUMEL, ZNUMEL), dtype=tl.float32) + oidx = xidx[:, None, None] * Y * Z + yidx[None, :, None] * Z + zidx[None, None, :] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_bf16_3d(output_ptr, X: tl.constexpr, Y: tl.constexpr, Z: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xidx = tl.arange(0, XNUMEL) + yidx = tl.arange(0, YNUMEL) + zidx = tl.arange(0, ZNUMEL) + Xmask = xidx < X + Ymask = yidx < Y + Zmask = zidx < Z + mask = (Xmask[:, None, None]) & (Ymask[None, :, None]) & (Zmask[None, None, :]) + ret = tl.zeros((XNUMEL, YNUMEL, ZNUMEL), dtype=tl.bfloat16) + oidx = xidx[:, None, None] * Y * Z + yidx[None, :, None] * Z + zidx[None, None, :] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_bool_3d(output_ptr, X: tl.constexpr, Y: tl.constexpr, Z: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xidx = tl.arange(0, XNUMEL) + yidx = tl.arange(0, YNUMEL) + zidx = tl.arange(0, ZNUMEL) + Xmask = xidx < X + Ymask = yidx < Y + Zmask = zidx < Z + mask = (Xmask[:, None, None]) & (Ymask[None, :, None]) & (Zmask[None, None, :]) + ret = tl.zeros((XNUMEL, YNUMEL, ZNUMEL), dtype=tl.int1) + oidx = xidx[:, None, None] * Y * Z + yidx[None, :, None] * Z + zidx[None, None, :] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_int8_2d(output_ptr, Y: tl.constexpr, Z: tl.constexpr, YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + yidx = tl.arange(0, YNUMEL) + zidx = tl.arange(0, ZNUMEL) + Ymask = yidx < Y + Zmask = zidx < Z + mask = (Ymask[:, None]) & (Zmask[None, :]) + ret = tl.zeros((YNUMEL, ZNUMEL), dtype=tl.int8) + oidx = yidx[:, None] * Z + zidx[None, :] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_int16_2d(output_ptr, Y: tl.constexpr, Z: tl.constexpr, YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + yidx = tl.arange(0, YNUMEL) + zidx = tl.arange(0, ZNUMEL) + Ymask = yidx < Y + Zmask = zidx < Z + mask = (Ymask[:, None]) & (Zmask[None, :]) + ret = tl.zeros((YNUMEL, ZNUMEL), dtype=tl.int16) + oidx = yidx[:, None] * Z + zidx[None, :] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_int32_2d(output_ptr, Y: tl.constexpr, Z: tl.constexpr, YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + yidx = tl.arange(0, YNUMEL) + zidx = tl.arange(0, ZNUMEL) + Ymask = yidx < Y + Zmask = zidx < Z + mask = (Ymask[:, None]) & (Zmask[None, :]) + ret = tl.zeros((YNUMEL, ZNUMEL), dtype=tl.int32) + oidx = yidx[:, None] * Z + zidx[None, :] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_int64_2d(output_ptr, Y: tl.constexpr, Z: tl.constexpr, YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + yidx = tl.arange(0, YNUMEL) + zidx = tl.arange(0, ZNUMEL) + Ymask = yidx < Y + Zmask = zidx < Z + mask = (Ymask[:, None]) & (Zmask[None, :]) + ret = tl.zeros((YNUMEL, ZNUMEL), dtype=tl.int64) + oidx = yidx[:, None] * Z + zidx[None, :] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_fp16_2d(output_ptr, Y: tl.constexpr, Z: tl.constexpr, YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + yidx = tl.arange(0, YNUMEL) + zidx = tl.arange(0, ZNUMEL) + Ymask = yidx < Y + Zmask = zidx < Z + mask = (Ymask[:, None]) & (Zmask[None, :]) + ret = tl.zeros((YNUMEL, ZNUMEL), dtype=tl.float16) + oidx = yidx[:, None] * Z + zidx[None, :] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_fp32_2d(output_ptr, Y: tl.constexpr, Z: tl.constexpr, YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + yidx = tl.arange(0, YNUMEL) + zidx = tl.arange(0, ZNUMEL) + Ymask = yidx < Y + Zmask = zidx < Z + mask = (Ymask[:, None]) & (Zmask[None, :]) + ret = tl.zeros((YNUMEL, ZNUMEL), dtype=tl.float32) + oidx = yidx[:, None] * Z + zidx[None, :] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_bf16_2d(output_ptr, Y: tl.constexpr, Z: tl.constexpr, YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + yidx = tl.arange(0, YNUMEL) + zidx = tl.arange(0, ZNUMEL) + Ymask = yidx < Y + Zmask = zidx < Z + mask = (Ymask[:, None]) & (Zmask[None, :]) + ret = tl.zeros((YNUMEL, ZNUMEL), dtype=tl.bfloat16) + oidx = yidx[:, None] * Z + zidx[None, :] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_bool_2d(output_ptr, Y: tl.constexpr, Z: tl.constexpr, YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + yidx = tl.arange(0, YNUMEL) + zidx = tl.arange(0, ZNUMEL) + Ymask = yidx < Y + Zmask = zidx < Z + mask = (Ymask[:, None]) & (Zmask[None, :]) + ret = tl.zeros((YNUMEL, ZNUMEL), dtype=tl.int1) + oidx = yidx[:, None] * Z + zidx[None, :] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_int8_1d(output_ptr, Z: tl.constexpr, ZNUMEL: tl.constexpr): + zidx = tl.arange(0, ZNUMEL) + Zmask = zidx < Z + mask = (Zmask[:]) + ret = tl.zeros((ZNUMEL, ), dtype=tl.int8) + oidx = zidx[:] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_int16_1d(output_ptr, Z: tl.constexpr, ZNUMEL: tl.constexpr): + zidx = tl.arange(0, ZNUMEL) + Zmask = zidx < Z + mask = (Zmask[:]) + ret = tl.zeros((ZNUMEL, ), dtype=tl.int16) + oidx = zidx[:] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_int32_1d(output_ptr, Z: tl.constexpr, ZNUMEL: tl.constexpr): + zidx = tl.arange(0, ZNUMEL) + Zmask = zidx < Z + mask = (Zmask[:]) + ret = tl.zeros((ZNUMEL, ), dtype=tl.int32) + oidx = zidx[:] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_int64_1d(output_ptr, Z: tl.constexpr, ZNUMEL: tl.constexpr): + zidx = tl.arange(0, ZNUMEL) + Zmask = zidx < Z + mask = (Zmask[:]) + ret = tl.zeros((ZNUMEL, ), dtype=tl.int64) + oidx = zidx[:] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_fp16_1d(output_ptr, Z: tl.constexpr, ZNUMEL: tl.constexpr): + zidx = tl.arange(0, ZNUMEL) + Zmask = zidx < Z + mask = (Zmask[:]) + ret = tl.zeros((ZNUMEL, ), dtype=tl.float16) + oidx = zidx[:] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_fp32_1d(output_ptr, Z: tl.constexpr, ZNUMEL: tl.constexpr): + zidx = tl.arange(0, ZNUMEL) + Zmask = zidx < Z + mask = (Zmask[:]) + ret = tl.zeros((ZNUMEL, ), dtype=tl.float32) + oidx = zidx[:] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_bf16_1d(output_ptr, Z: tl.constexpr, ZNUMEL: tl.constexpr): + zidx = tl.arange(0, ZNUMEL) + Zmask = zidx < Z + mask = (Zmask[:]) + ret = tl.zeros((ZNUMEL, ), dtype=tl.bfloat16) + oidx = zidx[:] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_bool_1d(output_ptr, Z: tl.constexpr, ZNUMEL: tl.constexpr): + zidx = tl.arange(0, ZNUMEL) + Zmask = zidx < Z + mask = (Zmask[:]) + ret = tl.zeros((ZNUMEL, ), dtype=tl.int1) + oidx = zidx[:] + tl.store(output_ptr + oidx, ret, mask=mask) + + +@triton.jit +def fn_npu_int8_0d(output_ptr, N: tl.constexpr): + zero = tl.zeros((), dtype=tl.int8) + tl.store(output_ptr, zero) + + +@triton.jit +def fn_npu_int16_0d(output_ptr, N: tl.constexpr): + zero = tl.zeros((), dtype=tl.int16) + tl.store(output_ptr, zero) + + +@triton.jit +def fn_npu_int32_0d(output_ptr, N: tl.constexpr): + zero = tl.zeros((), dtype=tl.int32) + tl.store(output_ptr, zero) + + +@triton.jit +def fn_npu_int64_0d(output_ptr, N: tl.constexpr): + zero = tl.zeros((), dtype=tl.int64) + tl.store(output_ptr, zero) + + +@triton.jit +def fn_npu_fp16_0d(output_ptr, N: tl.constexpr): + zero = tl.zeros((), dtype=tl.float16) + tl.store(output_ptr, zero) + + +@triton.jit +def fn_npu_fp32_0d(output_ptr, N: tl.constexpr): + zero = tl.zeros((), dtype=tl.float32) + tl.store(output_ptr, zero) + + +@triton.jit +def fn_npu_bf16_0d(output_ptr, N: tl.constexpr): + zero = tl.zeros((), dtype=tl.bfloat16) + tl.store(output_ptr, zero) + + +@triton.jit +def fn_npu_bool_0d(output_ptr, N: tl.constexpr): + zero = tl.zeros((), dtype=tl.int1) + tl.store(output_ptr, zero) + + +test_dtype = ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16', 'bool'] +test_shape0d = [()] +test_shape1d = TestUtils.test_shape1d +test_shape2d = TestUtils.test_shape2d +test_shape3d = TestUtils.test_shape3d + +# 定义 dtype 到 (test_func, test_sigtype) 的映射 +dtype_mapping3d = { + 'int8': (fn_npu_int8_3d, torch.int8), + 'int16': (fn_npu_int16_3d, torch.int16), + 'int32': (fn_npu_int32_3d, torch.int32), + 'int64': (fn_npu_int64_3d, torch.int64), + 'float16': (fn_npu_fp16_3d, torch.float16), + 'float32': (fn_npu_fp32_3d, torch.float32), + 'bfloat16': (fn_npu_bf16_3d, torch.bfloat16), + 'bool': (fn_npu_bool_3d, torch.bool), +} +dtype_mapping2d = { + 'int8': (fn_npu_int8_2d, torch.int8), + 'int16': (fn_npu_int16_2d, torch.int16), + 'int32': (fn_npu_int32_2d, torch.int32), + 'int64': (fn_npu_int64_2d, torch.int64), + 'float16': (fn_npu_fp16_2d, torch.float16), + 'float32': (fn_npu_fp32_2d, torch.float32), + 'bfloat16': (fn_npu_bf16_2d, torch.bfloat16), + 'bool': (fn_npu_bool_2d, torch.bool), +} +dtype_mapping1d = { + 'int8': (fn_npu_int8_1d, torch.int8), + 'int16': (fn_npu_int16_1d, torch.int16), + 'int32': (fn_npu_int32_1d, torch.int32), + 'int64': (fn_npu_int64_1d, torch.int64), + 'float16': (fn_npu_fp16_1d, torch.float16), + 'float32': (fn_npu_fp32_1d, torch.float32), + 'bfloat16': (fn_npu_bf16_1d, torch.bfloat16), + 'bool': (fn_npu_bool_1d, torch.bool), +} +dtype_mapping0d = { + 'int8': (fn_npu_int8_0d, torch.int8), + 'int16': (fn_npu_int16_0d, torch.int16), + 'int32': (fn_npu_int32_0d, torch.int32), + 'int64': (fn_npu_int64_0d, torch.int64), + 'float16': (fn_npu_fp16_0d, torch.float16), + 'float32': (fn_npu_fp32_0d, torch.float32), + 'bfloat16': (fn_npu_bf16_0d, torch.bfloat16), + 'bool': (fn_npu_bool_0d, torch.bool), +} + +# 生成测试用例 +testlist = [(func, sigtype, dtype, shape) + for sigtype in test_dtype + for shape in test_shape0d + for func, dtype in [dtype_mapping0d[sigtype]] # 直接解包映射结果 + ] + +testlist += [(func, sigtype, dtype, shape) + for sigtype in test_dtype + for shape in test_shape1d + for func, dtype in [dtype_mapping1d[sigtype]] # 直接解包映射结果 + ] + +testlist += [(func, sigtype, dtype, shape) + for sigtype in test_dtype + for shape in test_shape2d + for func, dtype in [dtype_mapping2d[sigtype]] # 直接解包映射结果 + ] + +testlist += [(func, sigtype, dtype, shape) + for sigtype in test_dtype + for shape in test_shape3d + for func, dtype in [dtype_mapping3d[sigtype]] # 直接解包映射结果 + ] + + +@pytest.mark.parametrize('testfunc, sigtype, dtype, shape', testlist) +def test_npu(testfunc, sigtype, dtype, shape): + if check_ub_mem_overflow(sigtype, shape): + pytest.skip(f"dtype:{sigtype} shape:{shape} mem overflow") + x = 0 + output = 0 + if len(shape) == 3: + x = torch.full((shape[0], shape[1], shape[2]), 0, dtype=dtype).npu() + output = torch.randint(1, (shape[0], shape[1], shape[2]), dtype=dtype).npu() + testfunc[(1, 1, 1)](output, shape[0], shape[1], shape[2], shape[0], shape[1], shape[2]) + if len(shape) == 2: + x = torch.full((shape[0], shape[1]), 0, dtype=dtype).npu() + output = torch.randint(1, (shape[0], shape[1]), dtype=dtype).npu() + shape0 = shape[0] + shape1 = shape[1] + if x.numel() * x.element_size() >= 8192: + grid = (shape0, 1, 1) + shape0 = 1 + else: + grid = (1, 1, 1) + testfunc[grid](output, shape0, shape1, shape0, shape1) + if len(shape) == 1: + x = torch.full((shape[0], ), 0, dtype=dtype).npu() + output = torch.randint(1, (shape[0], ), dtype=dtype).npu() + testfunc[1, 1, 1](output, shape[0], shape[0]) + if len(shape) == 0: + output = torch.randint(1, size=shape, dtype=dtype).npu() + x = torch.zeros_like(output) + testfunc[(1, )](output_ptr=output, N=1) + test_common.validate_cmp(sigtype, output, x) + + +@triton.jit +def fn_npu_multi_d(output_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr): + dtype = output_ptr.type.element_ty + + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if (YB * ZB * MB * NB) > 1: + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if (ZB * MB * NB) > 1: + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if (MB * NB) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if NB > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + if (YB * ZB * MB * NB) == 1: + ret = tl.zeros((XB, ), dtype=dtype) + elif (ZB * MB * NB) == 1: + ret = tl.zeros((XB, YB), dtype=dtype) + elif (MB * NB) == 1: + ret = tl.zeros((XB, YB, ZB), dtype=dtype) + elif NB == 1: + ret = tl.zeros((XB, YB, ZB, MB), dtype=dtype) + else: + ret = tl.zeros((XB, YB, ZB, MB, NB), dtype=dtype) + + tl.store(output_ptr + offsets, ret) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('param_list', [ + ('float32', (4, 2, 16, 16)), + ('float32', (2, 4, 2, 16, 16)), + ('float32', (4, 2, 16, 16)), + ('float32', (2, 4, 2, 16, 16)), + ('float32', (4, 2, 16, 16)), + ('float32', (2, 4, 2, 16, 16)), +]) +def test_case_4d_5d(param_list): + dtype, shape = param_list + if check_ub_mem_overflow(dtype, shape): + pytest.skip(f"dtype:{dtype} shape:{shape} mem overflow") + y_ref = torch.full(shape, 0, dtype=eval('torch.' + dtype)).npu() + print(f"y_ref = {torch.flatten(y_ref)[0:4]}") + + y_cal = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + fn_npu_multi_d[(1, )](y_cal, *triton_shape) + print(f"y_cal = {torch.flatten(y_cal)[0:4]}") + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/generalization_cases/test_zeroslike.py b/third_party/ascend/unittest/generalization_cases/test_zeroslike.py new file mode 100644 index 0000000000..84763df73d --- /dev/null +++ b/third_party/ascend/unittest/generalization_cases/test_zeroslike.py @@ -0,0 +1,170 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +import test_common +from test_common import TestUtils, check_ub_mem_overflow + + +@triton.jit +def fn_npu_0d(output_ptr, x_ptr, YB: tl.constexpr): + yidx = tl.arange(0, YB) + + idx = yidx + + X = tl.load(x_ptr + idx) + + ret = tl.zeros_like(X) + + oidx = yidx + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_1d(output_ptr, x_ptr, YB: tl.constexpr): + yidx = tl.arange(0, YB) + + idx = yidx + + X = tl.load(x_ptr + idx) + + ret = tl.zeros_like(X) + + oidx = yidx + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_2d(output_ptr, x_ptr, YB: tl.constexpr, ZB: tl.constexpr): + pid = tl.program_id(0) + yidx = tl.arange(0, YB)[:, None] + pid * YB + zidx = tl.arange(0, ZB)[None, :] + + idx = yidx * ZB + zidx + + X = tl.load(x_ptr + idx) + + ret = tl.zeros_like(X) + + oidx = yidx * ZB + zidx + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_3d(output_ptr, x_ptr, YB: tl.constexpr, ZB: tl.constexpr, KB: tl.constexpr): + yidx = tl.arange(0, YB)[:, None, None] * ZB * KB + zidx = tl.arange(0, ZB)[None, :, None] * KB + kidx = tl.arange(0, KB)[None, None, :] + + idx = yidx + zidx + kidx + + X = tl.load(x_ptr + idx) + + ret = tl.zeros_like(X) + + oidx = yidx + zidx + kidx + + tl.store(output_ptr + oidx, ret) + + +test_shape0d = [()] +testlist = test_shape0d + TestUtils.test_shape1_2_3d + + +@pytest.mark.parametrize('shape', testlist) +@pytest.mark.parametrize('dtype', TestUtils.dtype_list) +def test_npu(shape, dtype): + logging.debug(f'dtype:{dtype} shape:{shape}') + if check_ub_mem_overflow(dtype, shape): + return + x = torch.full(shape, 0, dtype=eval('torch.' + dtype)).npu() + triton_res = torch.empty(shape, dtype=eval('torch.' + dtype)).npu() + torch_res = x + + if len(shape) == 0: + fn_npu_0d[1, 1, 1](triton_res, x, 1) + elif len(shape) == 1: + fn_npu_1d[1, 1, 1](triton_res, x, shape[0]) + elif len(shape) == 2: + fn_npu_2d[shape[0], 1, 1](triton_res, x, 1, shape[1]) + elif len(shape) == 3: + fn_npu_3d[1, 1, 1](triton_res, x, shape[0], shape[1], shape[2]) + + test_common.validate_cmp(dtype, triton_res, torch_res) + + +@triton.jit +def fn_npu_multi_d(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, MB: tl.constexpr, + NB: tl.constexpr): + offsets = tl.arange(0, XB) * (YB * ZB * MB * NB) + if (YB * ZB * MB * NB) > 1: + offsets = offsets[:, None] + tl.arange(0, YB)[None, :] * (ZB * MB * NB) + if (ZB * MB * NB) > 1: + offsets = offsets[:, :, None] + tl.arange(0, ZB)[None, None, :] * (MB * NB) + if (MB * NB) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, MB)[None, None, None, :] * NB + if NB > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, NB)[None, None, None, None, :] + + X = tl.load(x_ptr + offsets) + ret = tl.zeros_like(X) + + tl.store(output_ptr + offsets, ret) + + +@pytest.mark.shape_4d_5d +@pytest.mark.parametrize('param_list', [ + ('float32', (4, 2, 16, 16)), + ('float32', (2, 4, 2, 16, 16)), + ('float32', (4, 2, 16, 16)), + ('float32', (2, 4, 2, 16, 16)), + ('float32', (4, 2, 16, 16)), + ('float32', (2, 4, 2, 16, 16)), +]) +def test_case_4d_5d(param_list): + dtype, shape = param_list + if check_ub_mem_overflow(dtype, shape): + return + x0 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch.zeros_like(x0, dtype=eval('torch.' + dtype)).npu() + print(f"y_ref = {torch.flatten(y_ref)[0:4]}") + y_cal = torch.ones(shape, dtype=eval('torch.' + dtype)).npu() + + triton_shape = [*shape] + while len(triton_shape) < 5: + triton_shape.append(1) + fn_npu_multi_d[(1, )](y_cal, x0, *triton_shape) + print(f"y_cal = {torch.flatten(y_cal)[0:4]}") + test_common.validate_cmp(dtype, y_cal, y_ref) + + +if __name__ == "__main__": + for dtype in TestUtils.dtype_list: + for shape in [(37, ), (37, 3), (1, 22, 39)]: + test_npu(shape, dtype) diff --git a/third_party/ascend/unittest/kernels/README.md b/third_party/ascend/unittest/kernels/README.md new file mode 100644 index 0000000000..ae6c02be87 --- /dev/null +++ b/third_party/ascend/unittest/kernels/README.md @@ -0,0 +1,65 @@ +# 指导：如何新增kernel测试用例 + +新增kernel测试用例可以分为三大步: +1、准备pt文件 +2、在triton-ascend仓中添加kernel算子，完成本地kernel测试 +3、将pt文件上传到obs桶中 + +## 1、准备pt文件 + +pt 文件用于把 GPU（或参考实现）上的输入与输出作为 golden 数据，后续测试会在 NPU 上运行 Triton kernel 并与之比对。 + +**三步生成流程** + +- **步骤 1 — 构造GPU输入并保存副本预处理成NPU kernel的输入**：根据GPU上kernel或pytorch算子的参数构造 `input_data`（键名须与 kernel 参数一致），把所有 Tensor 克隆到 CPU，形成 `input_data_before`，若GPU上算子的输入和NPU上算子有出入，需要提前预处理使`input_data_before`符合NPU上算子入参的要求。 +- **步骤 2 — 运行GPU Kernel获取输出**：在GPU上运行GPU kernel，得到 `gpu_output`，并将 Tensor 转为 CPU。 +- **步骤 3 — 打包并保存**：把 `input_data_before`、`grid`、`gpu_output` 封装为字典，通过 `torch.save` 保存为 `{kernel_name}.pt`。如果有多组用例，保存为 list-of-dicts（`[case0, case1]`）。 + +**精简示例** + +```python +import copy +import torch + +DEVICE = torch.device("cuda:0") +batch_size = 2 +grid = (batch_size,) + +input_data = { +"output_token_ids_ptr": torch.zeros((batch_size, 4), dtype=torch.int32, device=DEVICE), +"cu_num_draft_tokens_ptr": torch.tensor([2, 1], dtype=torch.int32, device=DEVICE), +# ... 其它字段 +} + +# 保存输入副本到 CPU +input_data_before = { + k: (v.clone().cpu() if isinstance(v, torch.Tensor) else copy.deepcopy(v)) + for k, v in input_data.items() +} +# 预处理 input_data_before 符合 NPU kernel 输入 +input_data_before["npu_need_param_key"] = NPU_NEED_PARAMS_VALUE +# 运行 kernel（在 GPU / 参考实现上）并收集输出 +triton_kernel[grid](**input_data) +# 这里用 input_data 作为示例，实际应调用对应的 triton/pytorch 函数 +gpu_output = {k: (v.cpu() if isinstance(v, torch.Tensor) else v) for k, v in input_data.items()} + +save_obj = {"input_data": input_data_before, "grid": grid, "gpu_output": gpu_output} +torch.save(save_obj, ".pt") +# 多组用例场景：torch.save([save_obj1, save_obj2], ".pt") +``` + +## 2、在triton-ascend新增三方kernel测试用例 + +- **步骤 1 — 在triton-ascend仓中新增kernel算子** ：本地验证阶段，在 kernels/xxx(例如vllm、sglang) 下新增与算子同名的 Python 文件，内容为Triton kernel函数。 +- **步骤 2 — 本地测试** ：将pt文件放在kernels目录下，在项目根目录运行 +python -m pytest -v third_party/ascend/unittest/kernels/test_triton_kernel.py + +**说明** + +- 指定单个 kernel：在项目根目录下执行 python -m pytest -v third_party/ascend/unittest/kernels/test_triton_kernel.py --kernel={kernel_name} +- pt文件查找策略：优先使用仓库内匹配的本地 pt，若本地不存在则按需从远端 OBS 下载 {kernel_name}.pt文件。 +- 本地已存在的pt文件，在执行完测试后不会删除，从obs桶取的文件在跑完测试后会被测试程序直接删除。 + +## 3、将pt文件上传至obs桶 + +本地验证通过后，将pt文件统一上传到OBS桶当中，OBS桶链接：`https://triton-ascend-artifacts.obs.cn-southwest-2.myhuaweicloud.com/test/kernels/{xxx}_pt/{kernel_name}.pt，xxx为vllm或sglang` diff --git a/third_party/ascend/unittest/kernels/common_kernel.py b/third_party/ascend/unittest/kernels/common_kernel.py new file mode 100644 index 0000000000..fbbce42dd9 --- /dev/null +++ b/third_party/ascend/unittest/kernels/common_kernel.py @@ -0,0 +1,7 @@ +import triton +import triton.language as tl + + +@triton.jit +def safe_exp(x): + return tl.exp(tl.where(x <= 0, x, float("-inf"))) diff --git a/third_party/ascend/unittest/kernels/test_common.py b/third_party/ascend/unittest/kernels/test_common.py new file mode 100644 index 0000000000..4632fca352 --- /dev/null +++ b/third_party/ascend/unittest/kernels/test_common.py @@ -0,0 +1,119 @@ +from typing import Optional +import torch +import pytest + +DEVICE_TYPE_NPU = 'npu' + + +def validate_cmp(dtype, y_cal, y_ref, overflow_mode: Optional[str] = None, device_type: Optional[str] = None): + if device_type is not None: + target_device = torch.device(device_type) + y_cal = y_cal.to(target_device) + y_ref = y_ref.to(target_device) + else: + y_cal = y_cal.npu() + y_ref = y_ref.npu() + if overflow_mode == "saturate": + if dtype in ['float32', 'float16']: + min_value = -torch.finfo(dtype).min + max_value = torch.finfo(dtype).max + elif dtype in ['int32', 'int16', 'int8']: + min_value = torch.iinfo(dtype).min + max_value = torch.iinfo(dtype).max + elif dtype == 'bool': + min_value = 0 + max_value = 1 + else: + raise ValueError('Invalid parameter "dtype" is found : {}'.format(dtype)) + y_ref = torch.clamp(y_ref, min=min_value, max=max_value) + if dtype == 'float16': + torch.testing.assert_close(y_ref, y_cal, rtol=1e-03, atol=1e-03, equal_nan=True) + elif dtype == 'bfloat16': + torch.testing.assert_close(y_ref.to(torch.float32), y_cal.to(torch.float32), rtol=1e-03, atol=1e-03, + equal_nan=True) + elif dtype == 'float32': + torch.testing.assert_close(y_ref, y_cal, rtol=1e-03, atol=1e-03, equal_nan=True) + elif dtype in ['int64', 'int32', 'int16', 'int8']: + assert torch.equal(y_cal, y_ref) + elif dtype == 'bool': + assert torch.equal(y_cal, y_ref) + else: + raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + + +def convert_tensor_with_device_type(indata: dict, device_type: str): + target_device = torch.device(device_type) + outdata = {} + + for key, value in indata.items(): + if isinstance(value, torch.Tensor): + if value.device.type != target_device.type: + outdata[key] = value.to(target_device) + else: + outdata[key] = value + else: + outdata[key] = value + + return outdata + + +def compare_data_precision(dict_ref: dict, dict_cal: dict, device_type: str): + keys_ref, keys_cal = set(dict_ref.keys()), set(dict_cal.keys()) + if not keys_ref.issubset(keys_cal): + raise ValueError("The keys of dict_ref is not subset of dict_cal") + + for key in dict_ref.keys(): + val_a, val_b = dict_ref[key], dict_cal[key] + if not isinstance(val_b, type(val_a)): + raise ValueError("The data type of two dicts are different") + + if isinstance(val_a, torch.Tensor): + try: + promoted_dtype = torch.promote_types(val_a.dtype, val_b.dtype) + except Exception: + # Fallback: if promote_types fails for some exotic dtypes, cast both to float32 + promoted_dtype = torch.float32 + val_a_cmp = val_a.to(promoted_dtype) if val_a.dtype != promoted_dtype else val_a + val_b_cmp = val_b.to(promoted_dtype) if val_b.dtype != promoted_dtype else val_b + validate_cmp(dtype=str(promoted_dtype).split('.')[-1], y_ref=val_a_cmp, y_cal=val_b_cmp, + device_type=device_type) + + +def run_and_compare_ptfile(ptfile_path: str, kernel_runner, device_type: str = DEVICE_TYPE_NPU): + try: + datas = torch.load(ptfile_path, map_location=torch.device('cpu')) + except Exception as e: + pytest.fail(f"load file {ptfile_path} failed: {e}") + + def _run_single_case(data): + if not isinstance(data, dict): + pytest.fail("Each case loaded from pt file must be a dict") + + input_data = convert_tensor_with_device_type(data.get("input_data", {}), device_type=device_type) + grid = data.get("grid") + try: + kernel_runner(input_data, grid) + except Exception as e: + pytest.fail(f"kernel_runner execution failed: {e}") + + output_data_cpu = convert_tensor_with_device_type(input_data, device_type='cpu') + expected = data.get("gpu_output", {}) + expected_filtered = {k: expected[k] for k in output_data_cpu.keys() if k in expected} + if not expected_filtered: + pytest.fail("No matching expected outputs found in pt file for comparison") + try: + compare_data_precision(expected_filtered, output_data_cpu, device_type='cpu') + except Exception as e: + pytest.fail(f"The testcase failed: {e}") + + # Supports three scenarios: + # 1) The file stores a single dict (existing behavior) + # 2) The file stores a list, where each element is a case dict + # 3) The file stores a dict, but some tensors represent multiple cases in batch on the 0th dimension (no automatic splitting; it is recommended to use a list) + if isinstance(datas, list): + for _, data in enumerate(datas): + _run_single_case(data) + elif isinstance(datas, dict): + _run_single_case(datas) + else: + pytest.fail("Unsupported pt file format: must be a dict or a list of dicts") diff --git a/third_party/ascend/unittest/kernels/test_triton_kernel.py b/third_party/ascend/unittest/kernels/test_triton_kernel.py new file mode 100644 index 0000000000..528c8eb088 --- /dev/null +++ b/third_party/ascend/unittest/kernels/test_triton_kernel.py @@ -0,0 +1,73 @@ +import importlib +import os +import urllib.request +from pathlib import Path + +import pytest + +import test_common + + +def discover_kernels(): + kernels = [] + kernels_root_path = Path(__file__).parents[0] + for p in kernels_root_path.rglob("*.py"): + if not p.is_file(): + continue + if p.parent == kernels_root_path: + continue + rel = p.relative_to(kernels_root_path) + if len(rel.parts) == 1 or p.name == "__init__.py": + continue + module_path = ".".join(rel.with_suffix("").parts) + kernels.append((module_path, p.stem)) + return sorted(kernels, key=lambda x: x[1]) + + +KERNEL_ITEMS = discover_kernels() + + +@pytest.mark.parametrize("module_path, kernel_name", KERNEL_ITEMS) +def test_triton_kernel(module_path, kernel_name, pytestconfig): + selected = pytestconfig.getoption("kernel") + if selected: + if kernel_name not in selected: + pytest.skip(f"skip {kernel_name} due to --kernel filter") + base_url = "https://triton-ascend-artifacts.obs.cn-southwest-2.myhuaweicloud.com" + rel = module_path + parts = rel.split(".") if rel else [] + pt_url = f"{base_url}/test/kernels/{parts[0]}_pt/{kernel_name}.pt" + local_pt = Path(__file__).parent / f"{kernel_name}.pt" + downloaded = False + if not local_pt.exists(): + try: + urllib.request.urlretrieve(pt_url, local_pt) + downloaded = True + except Exception as e: + pytest.fail( + f"Failed to download the {kernel_name}.pt file. Please check whether the {kernel_name}.pt file has been uploaded to the OBS bucket: {e}" + ) + try: + mod = importlib.import_module(module_path) + except Exception as e: + pytest.fail(f"import {module_path} failed: {e}") + + if hasattr(mod, kernel_name): + kernel_attr = kernel_name + else: + candidates = [a for a in dir(mod) if a.endswith("_kernel")] + kernel_attr = candidates[0] if candidates else None + + if not kernel_attr: + pytest.fail(f"No kernel callable found in {module_path}") + + kernel_callable = getattr(mod, kernel_attr) + + def runner(input_data, grid): + kernel_callable[grid](**input_data) + + try: + test_common.run_and_compare_ptfile(str(local_pt), runner, device_type='npu') + finally: + if downloaded and local_pt.exists(): + local_pt.unlink() diff --git a/third_party/ascend/unittest/kernels/vllm/_triton_rope.py b/third_party/ascend/unittest/kernels/vllm/_triton_rope.py new file mode 100644 index 0000000000..d8ae6003d9 --- /dev/null +++ b/third_party/ascend/unittest/kernels/vllm/_triton_rope.py @@ -0,0 +1,112 @@ +import triton +import triton.language as tl + + +@triton.jit +def _triton_rope( + q_ptr, + q_row_stride, + k_ptr, + k_row_stride, + cos, + cos_row_stride, + sin, + sin_row_stride, + num_tokens, + n_qh: tl.constexpr, + n_kh: tl.constexpr, + hd: tl.constexpr, + rope_dim: tl.constexpr, + pad_n_qh: tl.constexpr, + pad_n_kh: tl.constexpr, + pad_rope_dim: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + IS_NEOX_STYLE: tl.constexpr, +): + """ + This triton kernel applies rotary embedding on q and k. + It supports rope_dim != head_dim scenario. + It supports both neox style and non-neox style rope computation. + + Input tensor layout assumptions: + + q size: (num_tokens, num_q_heads, head_dim) + q stride: (num_q_heads * head_dim, head_dim, 1) + k size: (num_tokens, num_kv_heads, head_dim) + k stride: (num_kv_heads * head_dim, head_dim, 1) + cos/sin size: (num_tokens, rope_dim/2) + cos/sin stride: (rope_dim/2, 1) + + Different compute pattern of IS_NEOX_STYLE: + + if IS_NEOX_STYLE: + x1, x2 = torch.chunk(x, 2, dim=-1) + else: + x1 = x[..., ::2] + x2 = x[..., 1::2] + o1 = x1 * cos - x2 * sin + o2 = x2 * cos + x1 * sin + if IS_NEOX_STYLE: + return torch.cat((o1, o2), dim=-1) + else: + return torch.stack((o1, o2), dim=-1).flatten(-2) + """ + pid = tl.program_id(0).to(tl.int64) + row_block_size = tl.num_programs(0) + + for row_idx in tl.range(pid, num_tokens, row_block_size): + q_start_ptr = q_ptr + row_idx * q_row_stride + k_start_ptr = k_ptr + row_idx * k_row_stride + + # #################################################################### + # get the cos(mθ_{i...d/2}) and sin(mθ_{i...d/2}) for token position + # m of this program instance + # #################################################################### + cos_start_ptr = cos + row_idx * cos_row_stride + sin_start_ptr = sin + row_idx * sin_row_stride + + cos_offsets = tl.arange(0, pad_rope_dim // 2) + cos_mask = cos_offsets < (rope_dim // 2) + cos_row = tl.load(cos_start_ptr + cos_offsets, mask=cos_mask, other=0).to(tl.float32) + sin_row = tl.load(sin_start_ptr + cos_offsets, mask=cos_mask, other=0).to(tl.float32) + + # #################################################################### + # Load the left and right half of q and k for the current + # program instance (i.e. for the current token) separately + # #################################################################### + # left half of the head + if IS_NEOX_STYLE: + first_half_q_offsets = tl.arange(0, pad_n_qh)[:, None] * hd + tl.arange(0, pad_rope_dim // 2)[None, :] + first_half_k_offsets = tl.arange(0, pad_n_kh)[:, None] * hd + tl.arange(0, pad_rope_dim // 2)[None, :] + else: + first_half_q_offsets = tl.arange(0, pad_n_qh)[:, None] * hd + (2 * tl.arange(0, pad_rope_dim // 2)[None, :]) + first_half_k_offsets = tl.arange(0, pad_n_kh)[:, None] * hd + (2 * tl.arange(0, pad_rope_dim // 2)[None, :]) + + first_q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & (tl.arange(0, pad_rope_dim // 2)[None, :] + < (rope_dim // 2)) + first_k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & (tl.arange(0, pad_rope_dim // 2)[None, :] + < (rope_dim // 2)) + q_tile_1 = tl.load(q_start_ptr + first_half_q_offsets, mask=first_q_mask, other=0).to(sin_row.dtype) + k_tile_1 = tl.load(k_start_ptr + first_half_k_offsets, mask=first_k_mask, other=0).to(sin_row.dtype) + + # right half of the head + if IS_NEOX_STYLE: + second_half_q_offsets = first_half_q_offsets + (rope_dim // 2) + second_half_k_offsets = first_half_k_offsets + (rope_dim // 2) + else: + second_half_q_offsets = first_half_q_offsets + 1 + second_half_k_offsets = first_half_k_offsets + 1 + second_q_mask = first_q_mask + second_k_mask = first_k_mask + q_tile_2 = tl.load(q_start_ptr + second_half_q_offsets, mask=second_q_mask, other=0).to(sin_row.dtype) + k_tile_2 = tl.load(k_start_ptr + second_half_k_offsets, mask=second_k_mask, other=0).to(sin_row.dtype) + + new_q_tile_1 = q_tile_1 * cos_row - q_tile_2 * sin_row + tl.store(q_start_ptr + first_half_q_offsets, new_q_tile_1, mask=first_q_mask) + new_q_tile_2 = q_tile_2 * cos_row + q_tile_1 * sin_row + tl.store(q_start_ptr + second_half_q_offsets, new_q_tile_2, mask=second_q_mask) + + new_k_tile_1 = k_tile_1 * cos_row - k_tile_2 * sin_row + tl.store(k_start_ptr + first_half_k_offsets, new_k_tile_1, mask=first_k_mask) + new_k_tile_2 = k_tile_2 * cos_row + k_tile_1 * sin_row + tl.store(k_start_ptr + second_half_k_offsets, new_k_tile_2, mask=second_k_mask) diff --git a/third_party/ascend/unittest/kernels/vllm/chunk_gated_delta_rule_fwd_kernel_h_blockdim64.py b/third_party/ascend/unittest/kernels/vllm/chunk_gated_delta_rule_fwd_kernel_h_blockdim64.py new file mode 100644 index 0000000000..b4d6468328 --- /dev/null +++ b/third_party/ascend/unittest/kernels/vllm/chunk_gated_delta_rule_fwd_kernel_h_blockdim64.py @@ -0,0 +1,162 @@ +import triton +import triton.language as tl + + +@triton.jit +def safe_exp(x): + return tl.exp(tl.where(x <= 0, x, float("-inf"))) + + +@triton.heuristics({ + "USE_G": lambda args: args["g"] is not None, + "USE_INITIAL_STATE": lambda args: args["h0"] is not None, + "STORE_FINAL_STATE": lambda args: args["ht"] is not None, + "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, +}) +@triton.jit(do_not_specialize=["T"]) +def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( + k, + v, + w, + v_new, + g, + h, + h0, + ht, + cu_seqlens, + chunk_offsets, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + USE_G: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + SAVE_NEW_VALUE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_nh = tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + T_max = 1 * T + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + stride_v = H * V + stride_k = Hg * K + stride_w = H * K + + b_h1_bv1 = tl.zeros([128, 64], dtype=tl.float32) + b_h1_bv2 = tl.zeros([128, 64], dtype=tl.float32) + + v_start1 = 0 + v_start2 = 64 + + offs_k = tl.arange(0, 128)[:, None] + offs_v1 = v_start1 + tl.arange(0, 64)[None, :] + offs_v2 = v_start2 + tl.arange(0, 64)[None, :] + mask_kv1 = (offs_k < K) & (offs_v1 < V) + mask_kv2 = (offs_k < K) & (offs_v2 < V) + + # load initial state + if USE_INITIAL_STATE: + h0_ptr = h0 + i_nh * K * V + ptr_h0_bv1 = h0_ptr + offs_k * V + offs_v1 * 1 + b_h1_bv1 += tl.load(ptr_h0_bv1, mask=mask_kv1, other=0.0).to(tl.float32) + + ptr_h0_bv2 = h0_ptr + offs_k * V + offs_v2 * 1 + b_h1_bv2 += tl.load(ptr_h0_bv2, mask=mask_kv2, other=0.0).to(tl.float32) + + # main recurrence + for i_t in range(NT): + h_base = h + (boh + i_t) * H * K * V + i_h * K * V + + p_h1_bv1 = tl.make_block_ptr(h_base, (K, V), (V, 1), (0, v_start1), (128, 64), (1, 0)) + tl.store(p_h1_bv1, b_h1_bv1.to(p_h1_bv1.dtype.element_ty), boundary_check=(0, 1)) + + p_h1_bv2 = tl.make_block_ptr(h_base, (K, V), (V, 1), (0, v_start2), (128, 64), (1, 0)) + tl.store(p_h1_bv2, b_h1_bv2.to(p_h1_bv2.dtype.element_ty), boundary_check=(0, 1)) + + offs_t_wv = (i_t * BT + tl.arange(0, BT))[:, None] + offs_k_wv = tl.arange(0, 128)[None, :] + mask_w = (offs_t_wv < T) & (offs_k_wv < K) + + w_base = w + bos * H * K + i_h * K + ptr_w = w_base + offs_t_wv * stride_w + offs_k_wv * 1 + b_w = tl.load(ptr_w, mask=mask_w, other=0.0) + + k_base = k + bos * Hg * K + (i_h // (H // Hg)) * K + p_k = tl.make_block_ptr(k_base, (K, T), (1, stride_k), (0, i_t * BT), (128, BT), (0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + + v_new_base = v_new + bos * H * V + i_h * V + + last_idx = min((i_t + 1) * BT, T) - 1 + b_g_last = tl.load(g + bos + i_h * T_max + last_idx) + + offs_t = i_t * BT + tl.arange(0, BT) + mask_t = offs_t < T + g_ptr = g + bos + i_h * T_max + b_g = tl.load(g_ptr + offs_t, mask=mask_t, other=0.0) + + b_g = safe_exp(b_g_last - b_g) + b_g_last = tl.exp(b_g_last) + + offs_t_v = (i_t * BT + tl.arange(0, BT))[:, None] + mask_v1 = (offs_t_v < T) & (offs_v1 < V) + + v_base = v + bos * H * V + i_h * V + ptr_v1 = v_base + offs_t_v * stride_v + offs_v1 * 1 + b_v1 = tl.load(ptr_v1, mask=mask_v1, other=0.0) + b_v_new1 = b_v1.to(tl.float32) + b_v_new1 -= tl.dot(b_w, b_h1_bv1.to(b_w.dtype)) + + if SAVE_NEW_VALUE: + p_v_new1 = tl.make_block_ptr(v_new_base, (T, V), (stride_v, 1), (i_t * BT, v_start1), (BT, 64), (1, 0)) + tl.store(p_v_new1, b_v_new1.to(p_v_new1.dtype.element_ty), boundary_check=(0, 1)) + + if USE_G: + b_v_new1 = b_v_new1 * b_g[:, None] + b_h1_bv1 = b_h1_bv1 * b_g_last + + b_v_new1 = b_v_new1.to(k.dtype.element_ty) + b_h1_bv1 += tl.dot(b_k, b_v_new1) + + mask_v2 = (offs_t_v < T) & (offs_v2 < V) + ptr_v2 = v_base + offs_t_v * stride_v + offs_v2 * 1 + b_v2 = tl.load(ptr_v2, mask=mask_v2, other=0.0) + b_v_new2 = b_v2.to(tl.float32) + b_v_new2 -= tl.dot(b_w, b_h1_bv2.to(b_w.dtype)) + + if SAVE_NEW_VALUE: + p_v_new2 = tl.make_block_ptr(v_new_base, (T, V), (stride_v, 1), (i_t * BT, v_start2), (BT, 64), (1, 0)) + tl.store(p_v_new2, b_v_new2.to(p_v_new2.dtype.element_ty), boundary_check=(0, 1)) + + if USE_G: + b_v_new2 = b_v_new2 * b_g[:, None] + b_h1_bv2 = b_h1_bv2 * b_g_last + + b_v_new2 = b_v_new2.to(k.dtype.element_ty) + b_h1_bv2 += tl.dot(b_k, b_v_new2) + + # epilogue + if STORE_FINAL_STATE: + ht_ptr = ht + i_nh * K * V + + p_ht1_bv1 = tl.make_block_ptr(ht_ptr, (K, V), (V, 1), (0, v_start1), (128, 64), (1, 0)) + tl.store(p_ht1_bv1, b_h1_bv1.to(p_ht1_bv1.dtype.element_ty), boundary_check=(0, 1)) + + p_ht1_bv2 = tl.make_block_ptr(ht_ptr, (K, V), (V, 1), (0, v_start2), (128, 64), (1, 0)) + tl.store(p_ht1_bv2, b_h1_bv2.to(p_ht1_bv2.dtype.element_ty), boundary_check=(0, 1)) diff --git a/third_party/ascend/unittest/kernels/vllm/chunk_local_cumsum_scalar_kernel.py b/third_party/ascend/unittest/kernels/vllm/chunk_local_cumsum_scalar_kernel.py new file mode 100644 index 0000000000..244cc22569 --- /dev/null +++ b/third_party/ascend/unittest/kernels/vllm/chunk_local_cumsum_scalar_kernel.py @@ -0,0 +1,59 @@ +import triton +import triton.language as tl + + +@triton.heuristics( + {'HAS_SCALE': lambda args: args['scale'] is not None, 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) +@triton.jit(do_not_specialize=['T']) +def chunk_local_cumsum_scalar_kernel( + s, + o, + scale, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + BLOCK_T: tl.constexpr, + REVERSE: tl.constexpr, + HAS_SCALE: tl.constexpr, + IS_VARLEN: tl.constexpr, + HEAD_FIRST: tl.constexpr, + CHUNK_SIZE: tl.constexpr = 64, +): + i_block, i_b = tl.program_id(0), tl.program_id(1) + N_CHUNKS: tl.constexpr = BLOCK_T // CHUNK_SIZE + + if IS_VARLEN: + i_s, i_block = tl.load(chunk_indices + i_block * 2).to(tl.int32), tl.load(chunk_indices + i_block * 2 + 1).to( + tl.int32) + bos, eos = tl.load(cu_seqlens + i_s).to(tl.int32), tl.load(cu_seqlens + i_s + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if HEAD_FIRST: + ptr_s = tl.make_block_ptr(s + bos * H, (H, T), (T, 1), (0, i_block * BLOCK_T), (H, BLOCK_T), (1, 0)) + ptr_o = tl.make_block_ptr(o + bos * H, (H, T), (T, 1), (0, i_block * BLOCK_T), (H, BLOCK_T), (1, 0)) + b_s = tl.load(ptr_s, boundary_check=(0, )).to(tl.float32) + b_s = tl.reshape(b_s, (H, N_CHUNKS, CHUNK_SIZE)) + b_s = tl.trans(b_s, (2, 0, 1)) + b_o = tl.cumsum(b_s, axis=0, reverse=REVERSE) + if HAS_SCALE: + b_o *= scale + b_o = tl.trans(b_o, (2, 0, 1)) + b_o = tl.reshape(b_o, (H, BLOCK_T)) + else: + ptr_s = tl.make_block_ptr(s + bos * H, (T, H), (H, 1), (i_block * BLOCK_T, 0), (BLOCK_T, H), (1, 0)) + ptr_o = tl.make_block_ptr(o + bos * H, (T, H), (H, 1), (i_block * BLOCK_T, 0), (BLOCK_T, H), (1, 0)) + b_s = tl.load(ptr_s, boundary_check=(0, )).to(tl.float32) + b_s = tl.reshape(b_s, (N_CHUNKS, CHUNK_SIZE, H)) + b_s = tl.trans(b_s, (1, 0, 2)) + b_o = tl.cumsum(b_s, axis=0, reverse=REVERSE) + if HAS_SCALE: + b_o *= scale + b_o = tl.trans(b_o, (1, 0, 2)) + b_o = tl.reshape(b_o, (BLOCK_T, H)) + + tl.store(ptr_o, b_o.to(s.dtype.element_ty), boundary_check=(0, )) + return diff --git a/third_party/ascend/unittest/kernels/vllm/chunk_scaled_dot_kkt_fwd_kernel.py b/third_party/ascend/unittest/kernels/vllm/chunk_scaled_dot_kkt_fwd_kernel.py new file mode 100644 index 0000000000..43cf91bf7d --- /dev/null +++ b/third_party/ascend/unittest/kernels/vllm/chunk_scaled_dot_kkt_fwd_kernel.py @@ -0,0 +1,67 @@ +import triton +import triton.language as tl + + +@triton.jit +def safe_exp(x): + return tl.exp(tl.where(x <= 0, x, float("-inf"))) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + 'USE_G': lambda args: args['g_cumsum'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def chunk_scaled_dot_kkt_fwd_kernel( + k, + beta, # [H, B, T] + g_cumsum, # [H, B, T] + A, + cu_seqlens, + chunk_indices, + T, + B, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_G: tl.constexpr, +): + bt_stride = B * T + i_t_i, _ = tl.program_id(0), tl.program_id(1) + + for i_bh in range(B * H): + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t_i * 2).to(tl.int32), tl.load(chunk_indices + i_t_i * 2 + 1).to( + tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + i_t = i_t_i + o_t = tl.arange(0, BT) + o_t_fp32 = o_t.to(tl.float32) + + p_beta = tl.make_block_ptr(beta + i_h * bt_stride + bos, (T, ), (1, ), (i_t * BT, ), (BT, ), (0, )) + b_beta = tl.load(p_beta, boundary_check=(0, )) + + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k + (bos * Hg + i_h // (H // Hg)) * K, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), + (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_A += tl.dot(b_k, tl.trans(b_k)) + + if USE_G: + p_g = tl.make_block_ptr(g_cumsum + i_h * bt_stride + bos, (T, ), (1, ), (i_t * BT, ), (BT, ), (0, )) + b_g = tl.load(p_g, boundary_check=(0, )) + b_g_diff = b_g[:, None] - b_g[None, :] + b_A *= safe_exp(b_g_diff) + + b_A *= b_beta[:, None] + b_A = tl.where(o_t_fp32[:, None] > o_t_fp32[None, :], b_A, 0) + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) diff --git a/third_party/ascend/unittest/kernels/vllm/expand_kernel.py b/third_party/ascend/unittest/kernels/vllm/expand_kernel.py new file mode 100644 index 0000000000..8c87b6d0b2 --- /dev/null +++ b/third_party/ascend/unittest/kernels/vllm/expand_kernel.py @@ -0,0 +1,33 @@ +import triton +import triton.language as tl +import triton.language.extra.cann.extension as extension + + +@triton.jit(do_not_specialize=["replace_from", "replace_to"]) +def expand_kernel( + output_ptr, # [num_tokens] + input_ptr, # [batch_size] + cu_num_tokens_ptr, # [batch_size] + replace_from, + replace_to, + vec_len, + MAX_NUM_TOKENS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + req_idx = tl.program_id(0) + offset = req_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + len_mask = offset < vec_len + + start_idx = tl.where(offset == 0, 0, tl.load(cu_num_tokens_ptr + offset - 1, len_mask)) + end_idx = tl.load(cu_num_tokens_ptr + offset, len_mask) + num_tokens = end_idx - start_idx + + src_val = tl.load(input_ptr + offset, len_mask) + src_val = tl.where(src_val == replace_from, replace_to, src_val) + + for i in tl.range(0, BLOCK_SIZE): + num_tokens1 = extension.get_element(num_tokens, (i, )) + start_idx1 = extension.get_element(start_idx, (i, )) + src_val1 = extension.get_element(src_val, (i, )) + offset1 = tl.arange(0, MAX_NUM_TOKENS) + tl.store(output_ptr + start_idx1 + offset1, src_val1, mask=offset1 < num_tokens1) diff --git a/third_party/ascend/unittest/kernels/vllm/fused_gdn_gating_kernel.py b/third_party/ascend/unittest/kernels/vllm/fused_gdn_gating_kernel.py new file mode 100644 index 0000000000..674b34f3b7 --- /dev/null +++ b/third_party/ascend/unittest/kernels/vllm/fused_gdn_gating_kernel.py @@ -0,0 +1,47 @@ +import triton +import triton.language as tl + + +@triton.jit +def fused_gdn_gating_kernel( + g, + beta_output, + A_log, + a, + b, + dt_bias, + seq_len, + NUM_HEADS: tl.constexpr, + NUM_BATCHES: tl.constexpr, + beta: tl.constexpr, + threshold: tl.constexpr, + BLK_HEADS: tl.constexpr, + COL_ITER: tl.constexpr, + BLK_BATCHES: tl.constexpr, + ROW_ITER: tl.constexpr, +): + i_b, i_s = tl.program_id(0), tl.program_id(1) + for row_idx in range(0, ROW_ITER): + batch_off = i_b * ROW_ITER * BLK_BATCHES + row_idx * BLK_BATCHES + tl.arange(0, BLK_BATCHES) + + for col_idx in range(0, COL_ITER): + head_off = col_idx * BLK_HEADS + tl.arange(0, BLK_HEADS) + + off = batch_off[:, None] * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off[None, :] + head_mask = head_off < NUM_HEADS + mask = head_mask[None, :] & (batch_off[:, None] < NUM_BATCHES) + + blk_A_log = tl.load(A_log + head_off, mask=head_mask) + blk_a = tl.load(a + off, mask=mask) + blk_b = tl.load(b + off, mask=mask) + blk_bias = tl.load(dt_bias + head_off, mask=head_mask) + + x = blk_a.to(tl.float32) + blk_bias.to(tl.float32)[None, :] + softplus_x = tl.where(beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x) + + blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x + tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask) + + # compute beta_output = sigmoid(b) + blk_beta_output = tl.sigmoid(blk_b.to(tl.float32)) + tl.store(beta_output + off, blk_beta_output.to(beta_output.dtype.element_ty), mask=mask) diff --git a/third_party/ascend/unittest/kernels/vllm/fused_qkvzba_split_reshape_cat_kernel.py b/third_party/ascend/unittest/kernels/vllm/fused_qkvzba_split_reshape_cat_kernel.py new file mode 100644 index 0000000000..291e8331ae --- /dev/null +++ b/third_party/ascend/unittest/kernels/vllm/fused_qkvzba_split_reshape_cat_kernel.py @@ -0,0 +1,50 @@ +import triton +import triton.language as tl + + +@triton.jit +def fused_qkvzba_split_reshape_cat_kernel( + mixed_qkv, + z, + b, + a, + mixed_qkvz, + mixed_ba, + NUM_HEADS_QK: tl.constexpr, + NUM_HEADS_V: tl.constexpr, + HEAD_QK: tl.constexpr, + HEAD_V: tl.constexpr, +): + i_bs, i_qk = tl.program_id(0), tl.program_id(1) + QKVZ_DIM_T: tl.constexpr = HEAD_QK * 2 + NUM_HEADS_V // NUM_HEADS_QK * HEAD_V * 2 + BA_DIM_T: tl.constexpr = NUM_HEADS_V // NUM_HEADS_QK * 2 + QKV_DIM_T: tl.constexpr = HEAD_QK * 2 + NUM_HEADS_V // NUM_HEADS_QK * HEAD_V + q_end: tl.constexpr = HEAD_QK + blk_q_ptr = (mixed_qkvz + i_bs * NUM_HEADS_QK * QKVZ_DIM_T + i_qk * QKVZ_DIM_T + tl.arange(0, q_end)) + k_end: tl.constexpr = q_end + HEAD_QK + blk_k_ptr = (mixed_qkvz + i_bs * NUM_HEADS_QK * QKVZ_DIM_T + i_qk * QKVZ_DIM_T + tl.arange(q_end, k_end)) + v_end: tl.constexpr = k_end + NUM_HEADS_V // NUM_HEADS_QK * HEAD_V + blk_v_ptr = (mixed_qkvz + i_bs * NUM_HEADS_QK * QKVZ_DIM_T + i_qk * QKVZ_DIM_T + tl.arange(k_end, v_end)) + z_end: tl.constexpr = v_end + NUM_HEADS_V // NUM_HEADS_QK * HEAD_V + blk_z_ptr = (mixed_qkvz + i_bs * NUM_HEADS_QK * QKVZ_DIM_T + i_qk * QKVZ_DIM_T + tl.arange(v_end, z_end)) + blk_q_st_ptr = (mixed_qkv + i_bs * NUM_HEADS_QK * QKV_DIM_T + i_qk * HEAD_QK + tl.arange(0, HEAD_QK)) + blk_k_st_ptr = (mixed_qkv + i_bs * NUM_HEADS_QK * QKV_DIM_T + NUM_HEADS_QK * HEAD_QK + i_qk * HEAD_QK + + tl.arange(0, HEAD_QK)) + blk_v_st_ptr = (mixed_qkv + i_bs * NUM_HEADS_QK * QKV_DIM_T + NUM_HEADS_QK * HEAD_QK * 2 + + i_qk * HEAD_V * NUM_HEADS_V // NUM_HEADS_QK + tl.arange(0, HEAD_V * NUM_HEADS_V // NUM_HEADS_QK)) + blk_z_st_ptr = (z + i_bs * NUM_HEADS_V * HEAD_V + i_qk * HEAD_V * NUM_HEADS_V // NUM_HEADS_QK + + tl.arange(0, HEAD_V * NUM_HEADS_V // NUM_HEADS_QK)) + tl.store(blk_q_st_ptr, tl.load(blk_q_ptr)) + tl.store(blk_k_st_ptr, tl.load(blk_k_ptr)) + tl.store(blk_v_st_ptr, tl.load(blk_v_ptr)) + tl.store(blk_z_st_ptr, tl.load(blk_z_ptr)) + b_end: tl.constexpr = NUM_HEADS_V // NUM_HEADS_QK + a_end: tl.constexpr = b_end + NUM_HEADS_V // NUM_HEADS_QK + for i in tl.static_range(b_end): + blk_b_ptr = mixed_ba + i_bs * NUM_HEADS_QK * BA_DIM_T + i_qk * BA_DIM_T + i + blk_b_st_ptr = b + i_bs * NUM_HEADS_V + i_qk * NUM_HEADS_V // NUM_HEADS_QK + i + tl.store(blk_b_st_ptr, tl.load(blk_b_ptr)) + for i in tl.static_range(b_end, a_end): + blk_a_ptr = mixed_ba + i_bs * NUM_HEADS_QK * BA_DIM_T + i_qk * BA_DIM_T + i + blk_a_st_ptr = (a + i_bs * NUM_HEADS_V + i_qk * NUM_HEADS_V // NUM_HEADS_QK + (i - b_end)) + tl.store(blk_a_st_ptr, tl.load(blk_a_ptr)) diff --git a/third_party/ascend/unittest/kernels/vllm/l2norm_fwd_kernel2_loop.py b/third_party/ascend/unittest/kernels/vllm/l2norm_fwd_kernel2_loop.py new file mode 100644 index 0000000000..012a76e569 --- /dev/null +++ b/third_party/ascend/unittest/kernels/vllm/l2norm_fwd_kernel2_loop.py @@ -0,0 +1,19 @@ +import triton +import triton.language as tl + + +@triton.jit +def l2norm_fwd_kernel2_loop(X, Y, eps, M, N: tl.constexpr, MBLOCK: tl.constexpr, NUM_CHUNKS: tl.constexpr): + base_row = tl.program_id(0) * (NUM_CHUNKS * MBLOCK) + rindex = tl.arange(0, N)[None, :] + + for chunk in range(NUM_CHUNKS): + row_idx = base_row + chunk * MBLOCK + tl.arange(0, MBLOCK)[:, None] + xmask = row_idx < M + + xs = tl.load(X + (rindex + N * row_idx), mask=xmask, other=0.0).to(tl.float32) + square = xs * xs + square_sum = tl.sum(square, 1)[:, None] + rsqrt = tl.rsqrt(square_sum + eps) + + tl.store(Y + (rindex + N * row_idx), xs * rsqrt, xmask) diff --git a/third_party/ascend/unittest/kernels/vllm/layer_norm_fwd_kernel.py b/third_party/ascend/unittest/kernels/vllm/layer_norm_fwd_kernel.py new file mode 100644 index 0000000000..77fa44b477 --- /dev/null +++ b/third_party/ascend/unittest/kernels/vllm/layer_norm_fwd_kernel.py @@ -0,0 +1,79 @@ +import triton +import triton.language as tl + + +@triton.heuristics({ + "HAS_BIAS": lambda args: args["B"] is not None, + "HAS_Z": lambda args: args["Z"] is not None, +}) +@triton.jit +def layer_norm_fwd_kernel( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + Z, # pointer to the other branch + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride_x_row, # how much to increase the pointer when moving by 1 row + stride_y_row, + stride_z_row, + M, # number of rows in X_base + N, # number of columns in X_base + eps, # epsilon to avoid division by zero + BLOCK_N: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_Z: tl.constexpr, + NORM_BEFORE_GATE: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + N_CORES: tl.constexpr, +): + # Map the program id to the row of X_base and Y_base it should compute. + row = tl.program_id(0) + group = tl.program_id(1) + + BLOCK_ROWS = M if M < N_CORES else N_CORES + n_iters = M // BLOCK_ROWS + remain = M % BLOCK_ROWS + if row < remain: + n_iters = n_iters + 1 + + for i in tl.range(n_iters): + X_base = X + (i * BLOCK_ROWS * stride_x_row) + row * stride_x_row + group * N + Y_base = Y + (i * BLOCK_ROWS * stride_y_row) + row * stride_y_row + group * N + if HAS_Z: + Z_base = Z + (i * BLOCK_ROWS * stride_z_row) + row * stride_z_row + group * N + if not IS_RMS_NORM: + Mean_base = Mean + (i * BLOCK_ROWS) + group * M + Rstd_base = Rstd + (i * BLOCK_ROWS) + group * M + W_base = W + group * N + if HAS_BIAS: + B_base = B + group * N + # Compute mean and variance + cols = tl.arange(0, BLOCK_N) + x = tl.load(X_base + cols, mask=cols < N, other=0.).to(tl.float32) + if HAS_Z and not NORM_BEFORE_GATE: + z = tl.load(Z_base + cols, mask=cols < N).to(tl.float32) + x *= z * tl.sigmoid(z) + if not IS_RMS_NORM: + mean = tl.sum(x, axis=0) / N + tl.store(Mean_base + row, mean) + xbar = tl.where(cols < N, x - mean, 0.) + var = tl.sum(xbar * xbar, axis=0) / N + else: + xbar = tl.where(cols < N, x, 0.) + var = tl.sum(xbar * xbar, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + tl.store(Rstd_base + row, rstd) + # Normalize and apply linear transformation + mask = cols < N + w = tl.load(W_base + cols, mask=mask).to(tl.float32) + if HAS_BIAS: + b = tl.load(B_base + cols, mask=mask).to(tl.float32) + x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd + y = x_hat * w + b if HAS_BIAS else x_hat * w + if HAS_Z and NORM_BEFORE_GATE: + z = tl.load(Z_base + cols, mask=mask).to(tl.float32) + y *= z * tl.sigmoid(z) + # Write output + tl.store(Y_base + cols, y, mask=mask) diff --git a/third_party/ascend/unittest/kernels/vllm/merge_16x16_to_32x32_inverse_kernel.py b/third_party/ascend/unittest/kernels/vllm/merge_16x16_to_32x32_inverse_kernel.py new file mode 100644 index 0000000000..6b48646a9a --- /dev/null +++ b/third_party/ascend/unittest/kernels/vllm/merge_16x16_to_32x32_inverse_kernel.py @@ -0,0 +1,66 @@ +import triton +import triton.language as tl + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.jit(do_not_specialize=["T"]) +def merge_16x16_to_32x32_inverse_kernel( + A, + Ad, + Ai, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + A += (bos * H + i_h) * 32 + Ad += (bos * H + i_h) * 16 + Ai += (bos * H + i_h) * 32 + + p_A_21 = tl.make_block_ptr(A, (T, 32), (H * 32, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0)) + p_Ad_11 = tl.make_block_ptr(Ad, (T, 16), (H * 16, 1), (i_t * 32, 0), (16, 16), (1, 0)) + p_Ad_22 = tl.make_block_ptr(Ad, (T, 16), (H * 16, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0)) + p_Ai_11 = tl.make_block_ptr(Ai, (T, 32), (H * 32, 1), (i_t * 32, 0), (16, 16), (1, 0)) + p_Ai_22 = tl.make_block_ptr(Ai, (T, 32), (H * 32, 1), (i_t * 32 + 16, 16), (16, 16), (1, 0)) + p_Ai_21 = tl.make_block_ptr(Ai, (T, 32), (H * 32, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0)) + + A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) + Ai_11 = tl.load(p_Ad_11, boundary_check=(0, 1)).to(tl.float32) + Ai_22 = tl.load(p_Ad_22, boundary_check=(0, 1)).to(tl.float32) + Ai_21 = -tl.dot( + tl.dot(Ai_22, A_21, input_precision="ieee"), + Ai_11, + input_precision="ieee", + ) + tl.store( + p_Ai_11, + Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_22, + Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_21, + Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) diff --git a/third_party/ascend/unittest/kernels/vllm/merge_16x16_to_64x64_inverse_kernel.py b/third_party/ascend/unittest/kernels/vllm/merge_16x16_to_64x64_inverse_kernel.py new file mode 100644 index 0000000000..63a001123f --- /dev/null +++ b/third_party/ascend/unittest/kernels/vllm/merge_16x16_to_64x64_inverse_kernel.py @@ -0,0 +1,133 @@ +import triton +import triton.language as tl +import triton.language.extra.cann.extension as extension + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.jit(do_not_specialize=["T"]) +def merge_16x16_to_64x64_inverse_kernel( + A, + Ad, + Ai, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t_val = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + i_t = i_t_val + else: + bos, eos = i_b * T, i_b * T + T + + # Base pointers (already offset by batch and head) + A += (bos * H + i_h) * 64 + Ad += (bos * H + i_h) * 16 + Ai += (bos * H + i_h) * 64 + + # load Ai_22 (Ad block at row i_t * 64 + 16, col 0, 16 * 16) + offs_m = i_t * 64 + 16 + tl.arange(0, 16) + offs_n = tl.arange(0, 16) + mask_Ad = (offs_m[:, None] < T) & (offs_n[None, :] < 16) + ptr_Ad = Ad + offs_m[:, None] * (H * 16) + offs_n[None, :] + Ai_22 = tl.load(ptr_Ad, mask=mask_Ad, other=0.0).to(tl.float32) + + # load A_21 (A block at row i_t * 64 + 16, col 0, 16 * 16) + mask_A = (offs_m[:, None] < T) & (offs_n[None, :] < 64) + ptr_A = A + offs_m[:, None] * (H * 64) + offs_n[None, :] + A_21 = tl.load(ptr_A, mask=mask_A, other=0.0).to(tl.float32) + tmp = tl.dot(Ai_22, A_21, input_precision="ieee") + + # load Ai_11 (Ad block at row i_t * 64, col 0, 16 * 16) + offs_m = i_t * 64 + tl.arange(0, 16) + offs_n = tl.arange(0, 16) + mask_Ad = (offs_m[:, None] < T) & (offs_n[None, :] < 16) + ptr_Ad = Ad + offs_m[:, None] * (H * 16) + offs_n[None, :] + Ai_11 = tl.load(ptr_Ad, mask=mask_Ad, other=0.0).to(tl.float32) + + Ai_21 = -tl.dot(tmp, Ai_11, input_precision="ieee") + + # load Ai_44 (Ad block at row i_t * 64 + 48, col 0, 16 * 16) + offs_m = i_t * 64 + 48 + tl.arange(0, 16) + offs_n = tl.arange(0, 16) + mask_Ad = (offs_m[:, None] < T) & (offs_n[None, :] < 16) + ptr_Ad = Ad + offs_m[:, None] * (H * 16) + offs_n[None, :] + Ai_44 = tl.load(ptr_Ad, mask=mask_Ad, other=0.0).to(tl.float32) + + # load A_43 (Ad block at row i_t * 64 + 48, col 32, 16 * 16) + offs_n = 32 + tl.arange(0, 16) + mask_A = (offs_m[:, None] < T) & (offs_n[None, :] < 64) + ptr_A = A + offs_m[:, None] * (H * 64) + offs_n[None, :] + A_43 = tl.load(ptr_A, mask=mask_A, other=0.0).to(tl.float32) + tmp = tl.dot(Ai_44, A_43, input_precision="ieee") + + # load Ai_33 (Ad block at row i_t * 64 + 32, col 0, 16 * 16) + offs_m = i_t * 64 + 32 + tl.arange(0, 16) + offs_n = tl.arange(0, 16) + mask_Ad = (offs_m[:, None] < T) & (offs_n[None, :] < 16) + ptr_Ad = Ad + offs_m[:, None] * (H * 16) + offs_n[None, :] + Ai_33 = tl.load(ptr_Ad, mask=mask_Ad, other=0.0).to(tl.float32) + + Ai_43 = -tl.dot(tmp, Ai_33, input_precision="ieee") + + # build Ai_22_32 (32 * 32) + Ai_22_32 = tl.zeros((32, 32), tl.float32) + Ai_22_32 = extension.insert_slice(Ai_22_32, Ai_33, (0, 0), (16, 16), (1, 1)) + Ai_22_32 = extension.insert_slice(Ai_22_32, Ai_44, (16, 16), (16, 16), (1, 1)) + Ai_22_32 = extension.insert_slice(Ai_22_32, Ai_43, (16, 0), (16, 16), (1, 1)) + + # load A_21_32 (A block at row i_t * 64 + 32, col 0, 32 * 32) + offs_m = i_t * 64 + 32 + tl.arange(0, 32) + offs_n = tl.arange(0, 32) + mask_A = (offs_m[:, None] < T) & (offs_n[None, :] < 64) + ptr_A = A + offs_m[:, None] * (H * 64) + offs_n[None, :] + A_21_32 = tl.load(ptr_A, mask=mask_A, other=0.0).to(tl.float32) + tmp = tl.dot(Ai_22_32, A_21_32, input_precision="ieee") + + # build Ai_11_32 (32 * 32) + Ai_11_32 = tl.zeros((32, 32), tl.float32) + Ai_11_32 = extension.insert_slice(Ai_11_32, Ai_11, (0, 0), (16, 16), (1, 1)) + Ai_11_32 = extension.insert_slice(Ai_11_32, Ai_22, (16, 16), (16, 16), (1, 1)) + Ai_11_32 = extension.insert_slice(Ai_11_32, Ai_21, (16, 0), (16, 16), (1, 1)) + + Ai_21_32 = -tl.dot(tmp, Ai_11_32, input_precision="ieee") + + # store Ai_11_32 to (i_t * 64, 0) + offs_m = i_t * 64 + tl.arange(0, 32) + offs_n = tl.arange(0, 32) + mask_store = (offs_m[:, None] < T) & (offs_n[None, :] < 64) + ptr_Ai = Ai + offs_m[:, None] * (H * 64) + offs_n[None, :] + tl.store(ptr_Ai, Ai_11_32.to(ptr_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), mask=mask_store) + + # store Ai_22_32 to (i_t * 64 + 32, 32) + offs_m = i_t * 64 + 32 + tl.arange(0, 32) + offs_n = 32 + tl.arange(0, 32) + mask_store = (offs_m[:, None] < T) & (offs_n[None, :] < 64) + ptr_Ai = Ai + offs_m[:, None] * (H * 64) + offs_n[None, :] + tl.store(ptr_Ai, Ai_22_32.to(ptr_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), mask=mask_store) + + # store Ai_21_32 to (i_t * 64 + 32, 32) + offs_n = tl.arange(0, 32) + mask_store = (offs_m[:, None] < T) & (offs_n[None, :] < 64) + ptr_Ai = Ai + offs_m[:, None] * (H * 64) + offs_n[None, :] + tl.store(ptr_Ai, Ai_21_32.to(ptr_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), mask=mask_store) + + # zero out the upper-right 32 * 32 block (rows 0 ~ 31, cols 32 ~ 63) + offs_m = i_t * 64 + tl.arange(0, 32) + offs_n = 32 + tl.arange(0, 32) + mask_store = (offs_m[:, None] < T) & (offs_n[None, :] < BT) + ptr_Ai = Ai + offs_m[:, None] * (H * BT) + offs_n[None, :] + zero_block = tl.zeros((32, 32), dtype=ptr_Ai.dtype.element_ty) + tl.store(ptr_Ai, zero_block, mask=mask_store) diff --git a/third_party/ascend/unittest/kernels/vllm/recompute_w_u_fwd_kernel.py b/third_party/ascend/unittest/kernels/vllm/recompute_w_u_fwd_kernel.py new file mode 100644 index 0000000000..d603b83473 --- /dev/null +++ b/third_party/ascend/unittest/kernels/vllm/recompute_w_u_fwd_kernel.py @@ -0,0 +1,62 @@ +import triton +import triton.language as tl + + +@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) +@triton.jit(do_not_specialize=['T']) +def recompute_w_u_fwd_kernel(k, v, beta, w, u, A, g, cu_seqlens, chunk_indices, T, H: tl.constexpr, Hg: tl.constexpr, + K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, + IS_VARLEN: tl.constexpr): + T_max = T + i_t_o = tl.program_id(0) + + for i_bh in range(H): + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t_o * 2).to(tl.int32), tl.load(chunk_indices + i_t_o * 2 + 1).to( + tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + offs_t = tl.arange(0, BT) + global_offs_t = i_t * BT + offs_t + mask_t = global_offs_t < T + + offs_t_2d = global_offs_t[:, None] + offs_bt = tl.arange(0, BT)[None, :] + ptr_A = (A + (bos * H + i_h) * BT + offs_t_2d * (H * BT) + offs_bt * 1) + mask_A = mask_t[:, None] + b_A = tl.load(ptr_A, mask=mask_A, other=0.0).to(tl.float32) + + ptr_g = g + bos + i_h * T_max + global_offs_t + b_g = tl.exp(tl.load(ptr_g, mask=mask_t, other=0.0)).to(tl.float32) + + ptr_beta = beta + bos + i_h * T_max + global_offs_t + b_beta = tl.load(ptr_beta, mask=mask_t, other=0.0).to(tl.float32) + + for i_v in range(tl.cdiv(V, BV)): + offs_v = i_v * BV + tl.arange(0, BV)[None, :] + mask_v = (mask_t[:, None]) & (offs_v < V) + + ptr_v = (v + (bos * H + i_h) * V + offs_t_2d * (H * V) + offs_v * 1) + b_v = tl.load(ptr_v, mask=mask_v, other=0.0).to(tl.float32) + + b_vb = (b_v * b_beta[:, None]) + b_u = tl.dot(b_A, b_vb, allow_tf32=False) + + ptr_u = (u + (bos * H + i_h) * V + offs_t_2d * (H * V) + offs_v * 1) + tl.store(ptr_u, b_u.to(ptr_u.dtype.element_ty), mask=mask_v) + + for i_k in range(tl.cdiv(K, BK)): + offs_k = i_k * BK + tl.arange(0, BK)[None, :] + mask_k = (mask_t[:, None]) & (offs_k < K) + ptr_k = (k + (bos * Hg + i_h // (H // Hg)) * K + offs_t_2d * (Hg * K) + offs_k * 1) + b_k = tl.load(ptr_k, mask=mask_k, other=0.0).to(tl.float32) + + b_kb = (b_k * b_beta[:, None] * b_g[:, None]) + b_w = tl.dot(b_A, b_kb) + + ptr_w = (w + (bos * H + i_h) * K + offs_t_2d * (H * K) + offs_k * 1) + tl.store(ptr_w, b_w.to(ptr_w.dtype.element_ty), mask=mask_k) diff --git a/third_party/ascend/unittest/kernels/vllm/rejection_greedy_sample_spec_len_1_triton.py b/third_party/ascend/unittest/kernels/vllm/rejection_greedy_sample_spec_len_1_triton.py new file mode 100644 index 0000000000..efda3cbf7d --- /dev/null +++ b/third_party/ascend/unittest/kernels/vllm/rejection_greedy_sample_spec_len_1_triton.py @@ -0,0 +1,42 @@ +import triton +import triton.language as tl +import triton.language.extra.cann.extension as extension + + +@triton.jit(do_not_specialize=["max_spec_len"]) +def bonus_renew_1( + bonus_token_ids_ptr, + position, + output_token_ids_ptr, +): + bonus_token_id = tl.load(bonus_token_ids_ptr + position) + tl.store(output_token_ids_ptr + position * 2 + 1, bonus_token_id) + + +@triton.jit(do_not_specialize=["max_spec_len"]) +def rejection_greedy_sample_spec_len_1_triton( + output_token_ids_ptr, # [batch_size, 2] + draft_token_ids_ptr, # [num_tokens] + target_argmax_ptr, # [num_tokens] + bonus_token_ids_ptr, + vec_len, + BLOCK_SIZE: tl.constexpr, +): + block_idx = tl.program_id(0) + offset = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offset < vec_len + + draft_token_id = tl.load(draft_token_ids_ptr + offset, mask) + target_argmax_id = tl.load(target_argmax_ptr + offset, mask) + tl.store(output_token_ids_ptr + offset * 2, target_argmax_id, mask) + + for pos in tl.range(0, BLOCK_SIZE): + draft_token_id1 = extension.get_element(draft_token_id, (pos, )) + target_argmax1 = extension.get_element(target_argmax_id, (pos, )) + position = block_idx * BLOCK_SIZE + pos + if draft_token_id1 == target_argmax1: + bonus_renew_1( + bonus_token_ids_ptr, + position, + output_token_ids_ptr, + ) diff --git a/third_party/ascend/unittest/kernels/vllm/rejection_random_sample_kernel.py b/third_party/ascend/unittest/kernels/vllm/rejection_random_sample_kernel.py new file mode 100644 index 0000000000..b5d124a912 --- /dev/null +++ b/third_party/ascend/unittest/kernels/vllm/rejection_random_sample_kernel.py @@ -0,0 +1,55 @@ +import triton +import triton.language as tl + + +@triton.jit(do_not_specialize=["max_spec_len"]) +def rejection_random_sample_kernel( + output_token_ids_ptr, # [batch_size, max_spec_len + 1] + cu_num_draft_tokens_ptr, # [batch_size] + draft_token_ids_ptr, # [num_tokens] + draft_probs_ptr, # [num_tokens, vocab_size] or None + target_probs_ptr, # [num_tokens, vocab_size] + bonus_token_ids_ptr, # [batch_size] + recovered_token_ids_ptr, # [num_tokens] + uniform_probs_ptr, # [num_tokens] + is_greedy_ptr, # [batch_size] + max_spec_len, + vocab_size, + NO_DRAFT_PROBS: tl.constexpr, +): + req_idx = tl.program_id(0) + is_greedy = tl.load(is_greedy_ptr + req_idx) + if is_greedy: + # Early exost for greedy sampling requests + return + + start_idx = 0 if req_idx == 0 else tl.load(cu_num_draft_tokens_ptr + req_idx - 1) + end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx) + num_draft_tokens = end_idx - start_idx + + rejected = False + for pos in range(num_draft_tokens): + if not rejected: + draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos) + if NO_DRAFT_PROBS: + draft_prob = 1 + else: + draft_prob = tl.load(draft_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id) + target_prob = tl.load(target_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id) + uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos) + if draft_prob > 0 and target_prob / draft_prob >= uniform_prob: + # Accept + token_id = draft_token_id + else: + # Reject. Use recovered token + rejected = True + token_id = tl.load(recovered_token_ids_ptr + start_idx + pos) + tl.store(output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos, token_id) + + if not rejected: + # If all tokens are accepted, append the bonus token + bonus_token_id = tl.load(bonus_token_ids_ptr + req_idx) + tl.store( + output_token_ids_ptr + req_idx * (max_spec_len + 1) + num_draft_tokens, + bonus_token_id, + ) diff --git a/third_party/ascend/unittest/kernels/vllm/sample_recovered_tokens_kernel.py b/third_party/ascend/unittest/kernels/vllm/sample_recovered_tokens_kernel.py new file mode 100644 index 0000000000..24aa9c7b7c --- /dev/null +++ b/third_party/ascend/unittest/kernels/vllm/sample_recovered_tokens_kernel.py @@ -0,0 +1,77 @@ +import triton +import triton.language as tl +import triton.language.extra.cann.extension as extension + + +@triton.jit +def sample_recovered_tokens_kernel( + output_token_ids_ptr, # [num_tokens] + cu_num_draft_tokens_ptr, # [batch_size] + draft_token_ids_ptr, # [num_tokens] + draft_probs_ptr, # [num_tokens, vocab_size] or None + target_probs_ptr, # [num_tokens, vocab_size] + q_ptr, # [batch_size, vocab_size] + vocab_size, + PADDED_VOCAB_SIZE: tl.constexpr, + NO_DRAFT_PROBS: tl.constexpr, + SUB_BLOCK: tl.constexpr, +): + req_idx = tl.program_id(0) + start_idx = 0 if req_idx == 0 else tl.load(cu_num_draft_tokens_ptr + req_idx - 1) + end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx) + num_draft_tokens = end_idx - start_idx + + # Early exit for out-of-range positions. + pos = tl.program_id(1) + if pos >= num_draft_tokens: + return + + loop = (vocab_size + SUB_BLOCK - 1) // SUB_BLOCK + global_recovered_id = -1 + global_max_p = -1.0 + if NO_DRAFT_PROBS: + draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos) + orig_prob = tl.load(target_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id) + # Temporarily zero out the probability of the draft token. + # This is essentially the same as target_prob - draft_prob, except that + # n-gram does not have draft_prob. We regard it as 1. + tl.store(target_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id, 0) + for loop_i in range(loop): + vocab_start = loop_i * SUB_BLOCK + vocab_offset = vocab_start + tl.arange(0, SUB_BLOCK) + prob = tl.load(target_probs_ptr + (start_idx + pos) * vocab_size + vocab_offset, mask=vocab_offset + < vocab_size, other=0) + q = tl.load(q_ptr + req_idx * vocab_size + vocab_offset, mask=vocab_offset < vocab_size, + other=float("-inf")) + new_p = prob / q + recovered_id = tl.argmax(new_p, axis=-1) + max_p = extension.get_element(new_p, (recovered_id, )) + if max_p > global_max_p: + global_max_p = max_p + global_recovered_id = vocab_start + recovered_id + else: + for loop_i in range(loop): + vocab_start = loop_i * SUB_BLOCK + vocab_offset = vocab_start + tl.arange(0, SUB_BLOCK) + draft_prob = tl.load(draft_probs_ptr + (start_idx + pos) * vocab_size + vocab_offset, mask=vocab_offset + < vocab_size, other=0) + target_prob = tl.load(target_probs_ptr + (start_idx + pos) * vocab_size + vocab_offset, mask=vocab_offset + < vocab_size, other=0) + prob = tl.maximum(target_prob - draft_prob, 0) + # NOTE(woosuk): We don't need `prob = prob / tl.sum(prob)` here because + # `tl.argmax` will select the maximum value. + + q = tl.load(q_ptr + req_idx * vocab_size + vocab_offset, mask=vocab_offset < vocab_size, + other=float("-inf")) + new_p = prob / q + recovered_id = tl.argmax(new_p, axis=-1) + max_p = extension.get_element(new_p, (recovered_id, )) + if max_p > global_max_p: + global_max_p = max_p + global_recovered_id = vocab_start + recovered_id + + tl.store(output_token_ids_ptr + start_idx + pos, global_recovered_id) + + if NO_DRAFT_PROBS: + # Restore the original probability. + tl.store(target_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id, orig_prob) diff --git a/third_party/ascend/unittest/kernels/vllm/solve_tril_16x16_kernel.py b/third_party/ascend/unittest/kernels/vllm/solve_tril_16x16_kernel.py new file mode 100644 index 0000000000..c3742f0a6d --- /dev/null +++ b/third_party/ascend/unittest/kernels/vllm/solve_tril_16x16_kernel.py @@ -0,0 +1,108 @@ +import triton +import triton.language as tl +import triton.language.extra.cann.extension as extension + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.jit(do_not_specialize=["T"]) +def solve_tril_16x16_kernel( + A, + Ad, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + IS_VARLEN: tl.constexpr, + LARGE_BLOCK_T: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + A = A + (bos * H + i_h) * BT + Ad = Ad + (bos * H + i_h) * 16 + + base_t = i_t * LARGE_BLOCK_T + + NTASKS: tl.constexpr = 2 + N_BLOCKS: tl.constexpr = LARGE_BLOCK_T // 16 // NTASKS + + for taskid in range(0, NTASKS): + base_t += taskid * (LARGE_BLOCK_T // NTASKS) + + # use make_block_ptr to reduce vector computation + b_A = tl.zeros((N_BLOCKS, 16, 16), dtype=tl.float32) + for blkid in range(0, N_BLOCKS): + row_start_o = base_t + blkid * 16 + col_start_o = row_start_o % BT + + # 1 Create in-block offset + offs_rows_in_block = tl.arange(0, 16) + offs_cols_in_block = tl.arange(0, 16) + + # 2 Calculate the pointer of each element + ptr_A_subrec16 = (A + row_start_o * H * BT + col_start_o + offs_rows_in_block[:, None] * H * BT + + offs_cols_in_block[None, :]) + + # 3 Create a mask to prevent out-of-bounds access + global_rows = row_start_o + offs_rows_in_block[:, None] + global_cols = col_start_o + offs_cols_in_block[None, :] + load_mask = (global_rows < T) & (global_cols < BT) + + # 4 Use mask to safely load data + b_A_subrec16 = tl.load(ptr_A_subrec16, mask=load_mask, other=0.0).to(tl.float32) + b_A = extension.insert_slice(ful=b_A, sub=b_A_subrec16[None, :, :], # (1, 16, 16) + offsets=[blkid, 0, 0], sizes=[1, 16, 16], strides=[1, 1, 1]) + + local_ori_A = tl.trans(b_A, (1, 0, 2)) + local_ori_A = tl.reshape(local_ori_A, (16, 16 * N_BLOCKS)) + + # Convert mask into matrix multiplication to avoid for loops ub oom + tmp = tl.arange(0, 16).to(tl.float32) + rows = tmp[:, None] + cols = tmp[None, :] + is_lower = (rows > cols).to(b_A.dtype) + b_A = -b_A * is_lower + + # for loop to update N_BLOCKS row vector + for i in range(1, 16): + nblks_vec16 = -extension.extract_slice(local_ori_A, (i, 0), (1, 16 * N_BLOCKS), (16 * N_BLOCKS, 1)) + b_a = tl.reshape(nblks_vec16, (N_BLOCKS, 16)) + + dot_tmp = tl.trans(b_a[:, :, None] * b_A, (1, 0, 2)) + dot_product = tl.sum(dot_tmp, 0) + b_a = b_a + dot_product + + b_a_new_expanded = b_a[:, None, :] + b_A = extension.insert_slice(ful=b_A, sub=b_a_new_expanded, offsets=[0, i, 0], sizes=[N_BLOCKS, 1, 16], + strides=[1, 1, 1]) + + on_diagonal = (rows == cols) + b_A = tl.where(on_diagonal, b_A + 1.0, b_A) + + b_A = tl.reshape(b_A, (N_BLOCKS * 16, 16)) + p_Ai = tl.make_block_ptr(Ad, (T, 16), (H * 16, 1), (base_t, 0), (N_BLOCKS * 16, 16), (1, 0)) + + # 1 Create in-block offset + offs_rows_to_store = tl.arange(0, N_BLOCKS * 16) + offs_cols_to_store = tl.arange(0, 16) + + # 2 Calculate the pointer of each element + p_Ai = (Ad + base_t * H * 16 + 0 + offs_rows_to_store[:, None] * H * 16 + offs_cols_to_store[None, :]) + # 3 Create a mask to prevent out-of-bounds access, only check rows + global_store_rows = base_t + offs_rows_to_store[:, None] + store_mask = global_store_rows < T + # 4 use mask to save data safely + tl.store(p_Ai, b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), mask=store_mask) diff --git a/third_party/ascend/unittest/lit.cfg.py b/third_party/ascend/unittest/lit.cfg.py new file mode 100644 index 0000000000..aa693e4dae --- /dev/null +++ b/third_party/ascend/unittest/lit.cfg.py @@ -0,0 +1,64 @@ +# -*- Python -*- + +import os +import platform +import re +import subprocess +import tempfile + +import lit.formats +import lit.util +from lit.llvm import llvm_config +from lit.llvm.subst import FindTool, ToolSubst + +# Configuration file for the 'lit' test runner + +# name: The name of this test suite +config.name = 'TRITON-ASCEND-MLIR' + +config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell) + +# suffixes: A list of file extensions to treat as test files. +config.suffixes = ['.mlir'] + +# test_source_root: The root path where tests are located. +config.test_source_root = os.path.dirname(os.path.abspath(__file__)) + +# test_exec_root: The root path where tests should be run. +config.test_exec_root = os.path.join(config.triton_obj_root, 'test') +config.substitutions.append(('%PATH%', config.environment['PATH'])) + +llvm_config.with_system_environment(['HOME', 'INCLUDE', 'LIB', 'TMP', 'TEMP']) + +# excludes: A list of directories to exclude from the testsuite. The 'Inputs' +# subdirectories contain auxiliary inputs for various tests in their parent +# directories. +config.excludes = ['Inputs', 'Examples', 'CMakeLists.txt', 'README.txt', 'LICENSE.txt'] + +# test_source_root: The root path where tests are located. +config.test_source_root = os.path.dirname(__file__) + +# test_exec_root: The root path where tests should be run. +config.test_exec_root = os.path.join(config.triton_obj_root, 'test') +config.triton_tools_dir = os.path.join(config.triton_obj_root, 'bin') +config.filecheck_dir = config.llvm_tools_dir +# FileCheck -enable-var-scope is enabled by default in MLIR test +# This option avoids to accidentally reuse variable across -LABEL match, +# it can be explicitly opted-in by prefixing the variable name with $ +config.environment["FILECHECK_OPTS"] = "--enable-var-scope" + +tool_dirs = [config.triton_tools_dir, config.llvm_tools_dir, config.filecheck_dir] + +# Tweak the PATH to include the tools dir. +for d in tool_dirs: + llvm_config.with_environment('PATH', d, append_path=True) +tools = [ + 'triton-opt', + ToolSubst('%PYTHON', config.python_executable, unresolved='ignore'), +] + +llvm_config.add_tool_substitutions(tools, tool_dirs) + +llvm_config.with_environment('PYTHONPATH', [ + os.path.join(config.mlir_binary_dir, 'python_packages', 'triton'), +], append_path=True) diff --git a/third_party/ascend/unittest/lit.site.cfg.py.in b/third_party/ascend/unittest/lit.site.cfg.py.in new file mode 100644 index 0000000000..8bbb8e0224 --- /dev/null +++ b/third_party/ascend/unittest/lit.site.cfg.py.in @@ -0,0 +1,23 @@ +@LIT_SITE_CFG_IN_HEADER@ + +import sys + +config.triton_obj_root = "@triton_BINARY_DIR@" +config.llvm_src_root = "@LLVM_SOURCE_DIR@" +config.llvm_obj_root = "@LLVM_BINARY_DIR@" +config.llvm_tools_dir = "@LLVM_TOOLS_DIR@" +config.llvm_lib_dir = "@LLVM_LIBS_DIR@" +config.llvm_shlib_dir = "@CMAKE_LIBRARY_OUTPUT_DIRECTORY@" +config.llvm_shlib_ext = "@CMAKE_SHARED_LIBRARY_SUFFIX@" +config.llvm_exe_ext = "@EXEEXT@" +config.lit_tools_dir = "@LLVM_LIT_TOOLS_DIR@" +config.mlir_binary_dir = "@MLIR_BINARY_DIR@" +config.python_executable = "@Python3_EXECUTABLE@" +config.enable_bindings_python = @MLIR_ENABLE_BINDINGS_PYTHON@ + + +import lit.llvm +lit.llvm.initialize(lit_config, config) + +# Let the main config do the real work +lit_config.load_config(config, "@CMAKE_CURRENT_SOURCE_DIR@/lit.cfg.py") diff --git a/third_party/ascend/unittest/pytest_ut/attn_cp_triton_kernel_3d.py b/third_party/ascend/unittest/pytest_ut/attn_cp_triton_kernel_3d.py new file mode 100644 index 0000000000..195458a936 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/attn_cp_triton_kernel_3d.py @@ -0,0 +1,182 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu + +#number of block per aiv core +NBLOCKS = 32 +#constants +TL_DTYPE_ATTN = tl.bfloat16 + + +@triton.jit +def tl_fn_forward_update(output_ptr0, output_ptr1, output_ptr2, prev_attn_out, prev_softmax_max, prev_softmax_sum, + cur_attn_out, cur_softmax_max, cur_softmax_sum, S: tl.constexpr, B: tl.constexpr, + H: tl.constexpr, N: tl.constexpr, E: tl.constexpr, CUR_SM_N_STRIDE: tl.constexpr, + PREV_SM_N_STRIDE: tl.constexpr, S_NBLOCKS: tl.constexpr, S_SUB: tl.constexpr, + BRC_SIZE: tl.constexpr): + + S_E_SUB_SIZE: tl.constexpr = S_SUB * E + D: tl.constexpr = H // N + S_BLOCK: tl.constexpr = (S + S_NBLOCKS - 1) // S_NBLOCKS + S_NSUB: tl.constexpr = (S_BLOCK + S_SUB - 1) // S_SUB + + D_SUB: tl.constexpr = E * BRC_SIZE + D_NLOOP: tl.constexpr = (D + D_SUB - 1) // D_SUB + + LOOP_COUNT: tl.constexpr = (S_BLOCK * B * N + S_SUB - 1) // S_SUB + block_idx = tl.program_id(0) + s_block_offset = block_idx % S_NBLOCKS * S_BLOCK + + for loop_index in range(LOOP_COUNT): + b = loop_index // (N * S_NSUB) + n = (loop_index // S_NSUB) % N + s_loop_offset = loop_index % S_NSUB * S_SUB + s = s_block_offset + s_loop_offset + tl.arange(0, S_SUB) #index on S axis + se_offset = E * (s_block_offset + s_loop_offset) + tl.arange(0, S_E_SUB_SIZE) + #assume no slice on N axis + offsets_prev = PREV_SM_N_STRIDE * (b * N + n) + se_offset + offsets_cur = CUR_SM_N_STRIDE * (b * N + n) + se_offset + mask0 = None + prev_softmax_max_local = tl.load(prev_softmax_max + offsets_prev, mask0) + cur_softmax_max_local = tl.load(cur_softmax_max + offsets_cur, mask0) + softmax_max = tl.maximum(prev_softmax_max_local, cur_softmax_max_local) + prev_scale = tl.exp(prev_softmax_max_local - softmax_max) + cur_scale = tl.exp(cur_softmax_max_local - softmax_max) + prev_softmax_sum_local = tl.load(prev_softmax_sum + offsets_prev, mask0) + cur_softmax_sum_local = tl.load(cur_softmax_sum + offsets_cur, mask0) + prev_softmax_sum_scaled = prev_softmax_sum_local * prev_scale + cur_softmax_sum_scaled = cur_softmax_sum_local * cur_scale + softmax_sum = prev_softmax_sum_scaled + cur_softmax_sum_scaled + prev_out_scale_local = prev_softmax_sum_scaled / softmax_sum + cur_out_scale_local = cur_softmax_sum_scaled / softmax_sum + prev_out_scale_out = prev_out_scale_local + cur_out_scale_out = cur_out_scale_local + mask1 = None + tl.store(output_ptr1 + offsets_cur, softmax_max, mask0) + tl.store(output_ptr2 + offsets_cur, softmax_sum, mask0) + prev_out_scale = prev_out_scale_local + prev_out_scale = prev_out_scale.reshape(S_SUB, 1, E) #(s_sub,8)->(1,1,s_sub,1,8) + cur_out_scale = cur_out_scale_local + cur_out_scale = cur_out_scale.reshape(S_SUB, 1, E) + for d_index in range(D_NLOOP): + # (s,b,h) -> (s,b,n*d) -> (s,b,n,d) + d = d_index * D_SUB + tl.arange(0, D_SUB) + offsets2 = s[:, None] * B * H + b * H + n * D + d[None, :] + mask2 = None + prev_attn_out_local = tl.load(prev_attn_out + offsets2, mask2) + cur_attn_out_local = tl.load(cur_attn_out + offsets2, mask2) + prev_attn_out_local = prev_attn_out_local.to(tl.float32) + cur_attn_out_local = cur_attn_out_local.to(tl.float32) + prev_attn_out_local = prev_attn_out_local.reshape(S_SUB, BRC_SIZE, E) + cur_attn_out_local = cur_attn_out_local.reshape(S_SUB, BRC_SIZE, E) + prev_out_scale_brc = prev_out_scale + prev_out_scale_brc = prev_out_scale_brc.broadcast_to(S_SUB, BRC_SIZE, E) + cur_out_scale_brc = cur_out_scale + cur_out_scale_brc = cur_out_scale_brc.broadcast_to(S_SUB, BRC_SIZE, E) + attn_out_local = prev_attn_out_local * prev_out_scale_brc + cur_attn_out_local * cur_out_scale_brc + attn_out_local = attn_out_local.reshape(S_SUB, D_SUB) + attn_out_local = attn_out_local.to(TL_DTYPE_ATTN) + tl.store(output_ptr0 + offsets2, attn_out_local, mask2) + + +#target_name = torch.npu.get_device_name(torch.npu.current_device()) +guards = {"dummy": None} + + +def forward_update_triton(prev_attn_out, prev_softmax_max, prev_softmax_sum, cur_attn_out, cur_softmax_max, + cur_softmax_sum): + + # size of sub_block in one block + S_SUB_SIZE = 64 + BROARDCAST_SIZE = 8 + + (S, B, H) = cur_attn_out.shape + N = prev_softmax_max.shape[1] + E = prev_softmax_max.shape[3] + D = S // H + D_SUB = E * BROARDCAST_SIZE + #(b,n,s,8) + PREV_SM_N_STRIDE = prev_softmax_max.stride()[1] + CUR_SM_N_STRIDE = cur_softmax_max.stride()[1] + + GUARD = (S % NBLOCKS == 0 and H % N == 0 and (H // N) % (D_SUB) == 0) + + if (not GUARD): + print(f"parameter does not meet compiling GUARD , fallback to eager foward_update \ + (S,H,N,D,D_SUB,NBLOCKS):{S},{H},{N},{D},{D_SUB},{NBLOCKS}") + + org_dtype = cur_attn_out.dtype + device_id = cur_attn_out.device.index + device = "npu:" + str(device_id) + + softmax_max = torch.empty_strided(cur_softmax_max.shape, cur_softmax_max.stride(), dtype=cur_softmax_max.dtype, + device=device) + softmax_sum = torch.empty_strided(cur_softmax_sum.shape, cur_softmax_sum.stride(), dtype=cur_softmax_max.dtype, + device=device) + attn_out = torch.empty_strided(cur_attn_out.shape, cur_attn_out.stride(), dtype=cur_attn_out.dtype, device=device) + + tl_fn_forward_update[NBLOCKS, 1, 1](attn_out, softmax_max, softmax_sum, prev_attn_out, prev_softmax_max, + prev_softmax_sum, cur_attn_out, cur_softmax_max, cur_softmax_sum, S=S, B=B, H=H, + N=N, E=E, CUR_SM_N_STRIDE=CUR_SM_N_STRIDE, PREV_SM_N_STRIDE=PREV_SM_N_STRIDE, + S_NBLOCKS=NBLOCKS, S_SUB=S_SUB_SIZE, BRC_SIZE=BROARDCAST_SIZE, debug=True) + + return attn_out, softmax_max, softmax_sum + + +@triton.jit +def tl_fn_backward_update(dq, dk, dv, cur_dq, cur_dk, cur_dv, qnumel: tl.constexpr, knumel: tl.constexpr, + XBLOCK: tl.constexpr, SIMD_SIZE: tl.constexpr): + + block_idx = tl.program_id(0) + block_offset = block_idx * XBLOCK + LOOP_COUNT: tl.constexpr = (XBLOCK + SIMD_SIZE - 1) // SIMD_SIZE + for loop_index in range(LOOP_COUNT): + loop_offset = block_offset + loop_index * SIMD_SIZE + tl.arange(0, SIMD_SIZE) + mask0 = loop_offset < qnumel + tmp0 = tl.load(dq + loop_offset, mask0).to(tl.float32) + tmp1 = tl.load(cur_dq + loop_offset, mask0).to(tl.float32) + tmp0 = tmp1 + tmp0 + tl.store(dq + loop_offset, tmp0.to(tl.bfloat16), mask0) + + mask1 = loop_offset < knumel + tmp2 = tl.load(dk + loop_offset, mask1).to(tl.float32) + tmp3 = tl.load(cur_dk + loop_offset, mask1).to(tl.float32) + tmp2 = tmp2 + tmp3 + + tmp4 = tl.load(dv + loop_offset, mask1).to(tl.float32) + tmp5 = tl.load(cur_dv + loop_offset, mask1).to(tl.float32) + tmp4 = tmp4 + tmp5 + + tl.store(dk + loop_offset, tmp2.to(tl.bfloat16), mask1) + tl.store(dv + loop_offset, tmp4.to(tl.bfloat16), mask1) + + +def backward_update_triton(dq, dk, dv, cur_dq, cur_dk, cur_dv): + # parameters need auto-tune. + SIMD_SIZE = 4 * 1024 + + xblock = max(dq.numel(), dk.numel()) // NBLOCKS + + tl_fn_backward_update[40, 1, 1](dq, dk, dv, cur_dq, cur_dk, cur_dv, dq.numel(), dk.numel(), xblock, SIMD_SIZE, + debug=True) diff --git a/third_party/ascend/unittest/pytest_ut/attn_cp_triton_kernel_3d_la.py b/third_party/ascend/unittest/pytest_ut/attn_cp_triton_kernel_3d_la.py new file mode 100644 index 0000000000..6617584ae5 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/attn_cp_triton_kernel_3d_la.py @@ -0,0 +1,100 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu + +#number of block per aiv core +NBLOCKS = 32 +# size of sub_block in one block +# reduce the size from 64 to 32 due to UB overflow +S_SUB_SIZE = 32 + + +@triton.jit +def tl_fn_forward_update_la(prev_attn_out, prev_softmax_log_max_sum, cur_attn_out, cur_softmax_log_max_sum, + B: tl.constexpr, N: tl.constexpr, S: tl.constexpr, D: tl.constexpr, + PREV_ATTN_NSTRIDE: tl.constexpr, PREV_SOFTMAX_NSTRIDE: tl.constexpr, + CUR_ATTN_NSTRIDE: tl.constexpr, CUR_SOFTMAX_NSTRIDE: tl.constexpr, S_NBLOCKS: tl.constexpr, + S_SUB: tl.constexpr): + + S_BLOCK: tl.constexpr = (S + S_NBLOCKS - 1) // S_NBLOCKS + S_NSUB: tl.constexpr = (S_BLOCK + S_SUB - 1) // S_SUB + LOOP_COUNT: tl.constexpr = (S_BLOCK * B * N + S_SUB - 1) // S_SUB + block_idx = tl.program_id(0) + + s_block_start = block_idx * S_BLOCK + # assuming S stride is D, if D is not contiguous, need to use 2-d offset + SIMD_SIZE: tl.constexpr = S_SUB * D + + for loop_index in range(LOOP_COUNT): + b = loop_index // (N * S_NSUB) + n = (loop_index // S_NSUB) % N + s_loop_start = (loop_index % S_NSUB) * S_SUB + s = s_block_start + s_loop_start + + sd_offsets = D * s + tl.arange(0, SIMD_SIZE) + s1_offsets = s + tl.arange(0, S_SUB) + + mask0 = None + softmax_offsets = PREV_SOFTMAX_NSTRIDE * (b * N + n) + s1_offsets + prev_softmax_local = tl.load(prev_softmax_log_max_sum + softmax_offsets, mask0) + offsets = CUR_SOFTMAX_NSTRIDE * (b * N + n) + s1_offsets + cur_softmax_local = tl.load(cur_softmax_log_max_sum + offsets, mask0) + + attn_offsets = PREV_ATTN_NSTRIDE * (b * N + n) + sd_offsets + prev_attn_local = tl.load(prev_attn_out + attn_offsets, mask0) + offsets = CUR_ATTN_NSTRIDE * (b * N + n) + sd_offsets + cur_attn_local = tl.load(cur_attn_out + offsets, mask0) + + tmp0 = tl.exp(cur_softmax_local) + tmp1 = tl.exp(prev_softmax_local) + softmax_log_max_sum = tl.log(tmp0 + tmp1) + tmp2 = (prev_softmax_local - softmax_log_max_sum).reshape(S_SUB, 1).broadcast_to(S_SUB, D) + tmp3 = (cur_softmax_local - softmax_log_max_sum).reshape(S_SUB, 1).broadcast_to(S_SUB, D) + + attn_out = tl.exp(tmp2) * prev_attn_local.reshape(S_SUB, D) + (tl.exp(tmp3) * cur_attn_local.reshape(S_SUB, D)) + mask1 = None + tl.store(prev_softmax_log_max_sum + softmax_offsets, softmax_log_max_sum, mask1) + tl.store(prev_attn_out + attn_offsets, attn_out.reshape(SIMD_SIZE, ), mask1) + + +#target_name = torch.npu.get_device_name(torch.npu.current_device()) +guards = {"dummy": None} + + +def forward_update_triton(prev_attn_out, prev_softmax_log_max_sum, cur_attn_out, cur_softmax_log_max_sum): + (B, N, S, D) = cur_attn_out.shape + #shape is (b,n,s,d) + PREV_ATTN_NSTRIDE = prev_attn_out.stride()[1] + PREV_SOFTMAX_NSTRIDE = prev_softmax_log_max_sum.stride()[1] + CUR_ATTN_NSTRIDE = cur_attn_out.stride()[1] + CUR_SOFTMAX_NSTRIDE = cur_softmax_log_max_sum.stride()[1] + + device_id = cur_attn_out.device.index + device = "npu:" + str(device_id) + + tl_fn_forward_update_la[NBLOCKS, 1, + 1](prev_attn_out, prev_softmax_log_max_sum, cur_attn_out, cur_softmax_log_max_sum, B=B, N=N, + S=S, D=D, PREV_ATTN_NSTRIDE=PREV_ATTN_NSTRIDE, PREV_SOFTMAX_NSTRIDE=PREV_SOFTMAX_NSTRIDE, + CUR_ATTN_NSTRIDE=CUR_ATTN_NSTRIDE, CUR_SOFTMAX_NSTRIDE=CUR_SOFTMAX_NSTRIDE, + S_NBLOCKS=NBLOCKS, S_SUB=S_SUB_SIZE, debug=True) diff --git a/third_party/ascend/unittest/pytest_ut/conftest.py b/third_party/ascend/unittest/pytest_ut/conftest.py new file mode 100644 index 0000000000..7c76d57ceb --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/conftest.py @@ -0,0 +1,48 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + + +@pytest.fixture(scope="module", autouse=True) +def assign_npu(request, worker_id): + marker = request.node.get_closest_marker("backend") + if marker: + backend = marker.args[0] + else: + backend = "torch_npu" + if backend == "torch_npu": + import torch + npu_count = torch.npu.device_count() + if worker_id == "master": + npu_id = 0 + else: + idx = int(worker_id.replace("gw", "")) + npu_id = idx % npu_count + torch.npu.set_device(npu_id) + elif backend == "mindspore": + import mindspore + npu_count = mindspore.device_context.ascend.device_count() + if worker_id == "master": + npu_id = 0 + else: + idx = int(worker_id.replace("gw", "")) + npu_id = idx % npu_count + mindspore.set_device("Ascend", npu_id) diff --git a/third_party/ascend/unittest/pytest_ut/test_01_vector_add.py b/third_party/ascend/unittest/pytest_ut/test_01_vector_add.py new file mode 100644 index 0000000000..d305609863 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_01_vector_add.py @@ -0,0 +1,70 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +""" +Vector Addition - Pytest Version +""" + +import torch +import torch_npu + +import triton +import triton.language as tl + + +@triton.jit +def add_kernel(x_ptr, # *Pointer* to first input vector. + y_ptr, # *Pointer* to second input vector. + output_ptr, # *Pointer* to output vector. + n_elements, # Size of the vector. + BLOCK_SIZE: tl.constexpr, # Number of elements each program should process. + # NOTE: `constexpr` so it can be used as a shape value. + ): + pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + tl.store(output_ptr + offsets, output, mask=mask) + + +def add(x: torch.Tensor, y: torch.Tensor): + output = torch.empty_like(x) + n_elements = output.numel() + + def grid(meta): + return (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + + add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024) + return output + + +def test_vector_addition(): + torch.manual_seed(0) + DEVICE = 'npu' + size = 98432 + x = torch.rand(size, device=DEVICE) + y = torch.rand(size, device=DEVICE) + output_torch = x + y + output_triton = add(x, y) + torch.testing.assert_close(output_triton, output_torch) diff --git a/third_party/ascend/unittest/pytest_ut/test_02_fused_softmax.py b/third_party/ascend/unittest/pytest_ut/test_02_fused_softmax.py new file mode 100644 index 0000000000..91fde18a95 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_02_fused_softmax.py @@ -0,0 +1,121 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Fused Softmax +============= +""" + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + + +def naive_softmax(x): + """Compute row-wise softmax of X using native pytorch + + We subtract the maximum element in order to avoid overflows. Softmax is invariant to + this shift. + """ + # read MN elements ; write M elements + x_max = x.max(dim=1)[0] + # read MN + M elements ; write MN elements + z = x - x_max[:, None] + # read MN elements ; write MN elements + numerator = torch.exp(z) + # read MN elements ; write M elements + denominator = numerator.sum(dim=1) + # read MN + M elements ; write MN elements + ret = numerator / denominator[:, None] + # in total: read 5MN + 2M elements ; wrote 3MN + 2M elements + return ret + + +@triton.jit +def softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride, n_rows, n_cols, + BLOCK_SIZE: tl.constexpr): + # starting row of the program + row_start = tl.program_id(0) + row_step = tl.num_programs(0) + for row_idx in tl.range(row_start, n_rows, row_step): + # The stride represents how much we need to increase the pointer to advance 1 row + row_start_ptr = input_ptr + row_idx * input_row_stride + # The block size is the next power of two greater than n_cols, so we can fit each + # row in a single block + col_offsets = tl.arange(0, BLOCK_SIZE) + input_ptrs = row_start_ptr + col_offsets + # Load the row into SRAM, using a mask since BLOCK_SIZE may be > than n_cols + mask = col_offsets < n_cols + row = tl.load(input_ptrs, mask=mask, other=-float('inf')) + # Subtract maximum for numerical stability + row_minus_max = row - tl.max(row, axis=0) + + numerator = tl.exp(row_minus_max) + denominator = tl.sum(numerator, axis=0) + softmax_output = numerator / denominator + # Write back output to DRAM + output_row_start_ptr = output_ptr + row_idx * output_row_stride + output_ptrs = output_row_start_ptr + col_offsets + tl.store(output_ptrs, softmax_output, mask=mask) + + +kernels = {} + + +def softmax(x): + n_rows, n_cols = x.shape + + # The block size of each loop iteration is the smallest power of two greater than the number of columns in `x` + BLOCK_SIZE = triton.next_power_of_2(n_cols) + + # Allocate output + y = torch.empty_like(x) + + # pre-compile kernel to get register usage and compute thread occupancy. + kernel, num_programs = kernels.get(BLOCK_SIZE, (None, 0)) + if kernel is None: + num_programs = 32 + kernel = softmax_kernel + kernels[BLOCK_SIZE] = (kernel, num_programs) + + num_programs = min(num_programs, n_rows) + + # Create a number of persistent programs. + kernel[(num_programs, 1, 1)](y, x, x.stride(0), y.stride(0), n_rows, n_cols, BLOCK_SIZE) + return y + + +@pytest.mark.parametrize( + "shape", + [ + (1823, 781), + (128, 257), + ], +) +@pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) +def test_fused_softmax(shape, dtype): + torch.manual_seed(0) + x = torch.randn(shape, dtype=dtype, device="npu") + + y_triton = softmax(x) + y_torch = torch.softmax(x, axis=1) + + torch.testing.assert_close(y_triton, y_torch, atol=1e-4, rtol=1e-4) diff --git a/third_party/ascend/unittest/pytest_ut/test_03_matrix_multiplication.py b/third_party/ascend/unittest/pytest_ut/test_03_matrix_multiplication.py new file mode 100644 index 0000000000..ec5f8e8655 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_03_matrix_multiplication.py @@ -0,0 +1,211 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Matrix Multiplication +=============== +""" + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import triton.language.extra.cann.extension as extension + +DEV = "npu" + + +def get_autotune_config(): + return [ + triton.Config({"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 128}), + triton.Config({"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64}), + ] + + +@triton.autotune( + configs=get_autotune_config(), + key=["M", "N", "K"], +) +@triton.jit +def matmul_kernel( + # Pointers to matrices + a_ptr, b_ptr, c_ptr, + # Matrix dimensions + M, N, K, + # The stride variables represent how much to increase the ptr by when moving by 1 + # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr` + # by to get the element one row down (A has M rows). + stride_am, stride_ak, # + stride_bk, stride_bn, # + stride_cm, stride_cn, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, # + ACTIVATION: tl.constexpr, # +): + """Kernel for computing the matmul C = A x B. + A has shape (M, K), B has shape (K, N) and C has shape (M, N) + """ + GROUP_SIZE_M: tl.constexpr = 1 + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + # See above `L2 Cache Optimizations` section for details. + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # ---------------------------------------------------------- + # Create pointers for the first blocks of A and B. + # We will advance this pointer as we move in the K direction + # and accumulate + # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers + # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers + # See above `Pointer Arithmetic` section for details + offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs_base = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs_base = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + msk_m = offs_am < M + msk_n = offs_bn < N + # ----------------------------------------------------------- + # Iterate to compute a block of the C matrix. + # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block + # of fp32 values for higher accuracy. + # `accumulator` will be converted back to fp16 after the loop. + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + # Load the next block of A and B, generate a mask by checking the K dimension. + # If it is out of bounds, set it to 0. + a_ptrs = a_ptrs_base + k * BLOCK_SIZE_K * stride_ak + b_ptrs = b_ptrs_base + k * BLOCK_SIZE_K * stride_bk + a = tl.load( + a_ptrs, + mask=msk_m[:, None] and (offs_k[None, :] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + b = tl.load( + b_ptrs, + mask=msk_n[None, :] and (offs_k[:, None] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + # We accumulate along the K dimension. + accumulator = tl.dot(a, b, accumulator) + # You can fuse arbitrary activation functions here + # while the accumulator is still in FP32! + # Original vector operations + # # ----------------------------------------------------------- + # # Write back the block of the output matrix C with masks. + # Comment out the following lines to enable split the workload to two vector cores + SUB_BLK_M: tl.constexpr = BLOCK_SIZE_M // 2 + for s in extension.parallel(0, 2, bind_sub_block=True): + vec_sub_blk = extension.extract_slice(accumulator, (s * SUB_BLK_M, 0), (SUB_BLK_M, BLOCK_SIZE_N), (1, 1)) + if ACTIVATION == "leaky_relu_custom": + vec_sub_blk = leaky_relu_custom(vec_sub_blk) + c_sub_blk = vec_sub_blk.to(tl.float16) + # ----------------------------------------------------------- + # Write back the block of the output matrix C with masks. + offs_cm = pid_m * BLOCK_SIZE_M + s * SUB_BLK_M + tl.arange(0, SUB_BLK_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c_sub_blk, mask=c_mask) + + +# We can fuse `leaky_relu_custom` by providing it as an `ACTIVATION` meta-parameter in `matmul_kernel`. +@triton.jit +def leaky_relu_custom(x): + return tl.where(x >= 0, x, 0.01 * x) + 1.0 + + +def torch_matmul(a, b, activation=""): + c = torch.matmul(a, b) + if activation == "leaky_relu_custom": + c = torch.where(c >= 0, c, 0.01 * c) + 1.0 + return c + + +# %% +# We can now create a convenience wrapper function that only takes two input tensors, +# and (1) checks any shape constraint; (2) allocates the output; (3) launches the above kernel. + + +def matmul(a, b, activation=""): + # Check constraints. + assert a.shape[1] == b.shape[0], "Incompatible dimensions" + assert a.is_contiguous(), "Matrix A must be contiguous" + M, K = a.shape + K, N = b.shape + # Allocates output. + c = torch.empty((M, N), device=a.device, dtype=torch.float16) + + # 1D launch kernel where each block gets its own program. + + def grid(META): + return (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), ) + + matmul_kernel[grid]( + a, b, c, # + M, N, K, # + a.stride(0), a.stride(1), # + b.stride(0), b.stride(1), # + c.stride(0), c.stride(1), # + ACTIVATION=activation, # + ) + return c + + +# %% +# Unit Test +# --------- +# +# We can test our custom matrix multiplication operation against a native torch implementation. +@pytest.mark.parametrize( + "shape", + [ + (512, 512, 512), + (256, 384, 128), + ], +) +@pytest.mark.parametrize( + "activation", + [ + "", + pytest.param("leaky_relu_custom", + marks=pytest.mark.skip(reason="temporarily skip leaky_relu_custom ub overflow case")), + ], +) +def test_matrix_multiplication(shape, activation): + m, k, n = shape + torch.manual_seed(0) + + a = torch.randn((m, k), device=DEV, dtype=torch.float16) + b = torch.randn((k, n), device=DEV, dtype=torch.float16) + + triton_output = matmul(a, b, activation) + torch_output = torch_matmul(a, b, activation) + + torch.testing.assert_close(triton_output, torch_output, atol=1e-3, rtol=1e-3) diff --git a/third_party/ascend/unittest/pytest_ut/test_04_low_memory_dropout.py b/third_party/ascend/unittest/pytest_ut/test_04_low_memory_dropout.py new file mode 100644 index 0000000000..d5f13ef462 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_04_low_memory_dropout.py @@ -0,0 +1,127 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Low-Memory Dropout +================== +""" + +import pytest +import torch +import torch_npu + +import triton +import triton.language as tl + +DEV = "npu" + + +@triton.jit +def _dropout( + x_ptr, # pointer to the input + x_keep_ptr, # pointer to a mask of 0s and 1s + output_ptr, # pointer to the output + n_elements, # number of elements in the `x` tensor + p, # probability that an element of `x` is changed to zero + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + x_keep = tl.load(x_keep_ptr + offsets, mask=mask) + output = tl.where(x_keep != 0, x / (1 - p), 0.0) + tl.store(output_ptr + offsets, output, mask=mask) + + +def dropout(x, x_keep, p): + output = torch.empty_like(x) + assert x.is_contiguous() + n_elements = x.numel() + + def grid(meta): + return (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + + _dropout[grid](x, x_keep, output, n_elements, p, BLOCK_SIZE=1024) + return output + + +@triton.jit +def _seeded_dropout( + x_ptr, + output_ptr, + n_elements, + p, + seed, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + random_seed = tl.rand(seed, offsets) + x_keep = random_seed > p + output = tl.where(x_keep, x / (1 - p), 0.0) + tl.store(output_ptr + offsets, output, mask=mask) + + +def seeded_dropout(x, p, seed): + output = torch.empty_like(x) + assert x.is_contiguous() + n_elements = x.numel() + + def grid(meta): + return (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + + _seeded_dropout[grid](x, output, n_elements, p, seed, BLOCK_SIZE=1024) + return output + + +@pytest.mark.parametrize("shape,p", [((10, ), 0.5), ((256, ), 0.5), ((513, ), 0.2), ((32, 64), 0.35)]) +def test_dropout_matches_reference(shape, p): + torch.manual_seed(0) + x = torch.randn(size=shape, device=DEV, dtype=torch.float32) + x_keep = (torch.rand(size=shape, device=DEV) > p).to(torch.int32) + + output = dropout(x, x_keep=x_keep, p=p) + expected = torch.where(x_keep != 0, x / (1 - p), torch.zeros_like(x)) + + torch.testing.assert_close(output, expected, atol=1e-6, rtol=0) + + +@pytest.mark.parametrize("shape,p,seed", [((10, ), 0.5, 123), ((256, ), 0.5, 123), ((513, ), 0.2, 7), + ((32, 64), 0.35, 999)]) +def test_seeded_dropout_is_deterministic(shape, p, seed): + torch.manual_seed(0) + x = torch.randn(size=shape, device=DEV, dtype=torch.float32) + + output = seeded_dropout(x, p=p, seed=seed) + output_same_seed = seeded_dropout(x, p=p, seed=seed) + output_different_seed = seeded_dropout(x, p=p, seed=512) + + torch.testing.assert_close(output, output_same_seed, atol=1e-6, rtol=0) + + assert output.shape == x.shape + assert output.dtype == x.dtype + assert torch.count_nonzero(output != output_different_seed).item() > 0 + assert torch.count_nonzero(output).item() <= x.numel() diff --git a/third_party/ascend/unittest/pytest_ut/test_05_layer_norm.py b/third_party/ascend/unittest/pytest_ut/test_05_layer_norm.py new file mode 100644 index 0000000000..84e3268b9b --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_05_layer_norm.py @@ -0,0 +1,119 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Layer Normalization +============= +""" + +import pytest +import torch +import triton +import triton.language as tl +import torch_npu + + +@triton.jit +def _layer_norm_fwd_fused( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride, # how much to increase the pointer when moving by 1 row + N, # number of columns in X + eps, # epsilon to avoid division by zero + BLOCK_SIZE: tl.constexpr, +): + # Map the program id to the row of X and Y it should compute. + row = tl.program_id(0) + Y += row * stride + X += row * stride + # Compute mean + mean = 0 + _mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + a = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32) + _mean += a + mean = tl.sum(_mean, axis=0) / N + # Compute variance + _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + x = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32) + x = tl.where(cols < N, x - mean, 0.) + _var += x * x + var = tl.sum(_var, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + # Write mean / rstd + tl.store(Mean + row, mean) + tl.store(Rstd + row, rstd) + # Normalize and apply linear transformation + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + mask = cols < N + w = tl.load(W + cols, mask=mask) + b = tl.load(B + cols, mask=mask) + x = tl.load(X + cols, mask=mask, other=0.).to(tl.float32) + x_hat = (x - mean) * rstd + y = x_hat * w + b + # Write output + tl.store(Y + cols, y, mask=mask) + + +@torch.inference_mode() +def layer_norm(x, normalized_shape, weight, bias, eps=1e-5): + # allocate output + y = torch.empty_like(x) + # reshape input data into 2D tensor + x_arg = x.reshape(-1, x.shape[-1]) + M, N = x_arg.shape + mean = torch.empty((M, ), dtype=torch.float32, device=x.device) + rstd = torch.empty((M, ), dtype=torch.float32, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + # heuristics for number of warps + num_warps = min(max(BLOCK_SIZE // 256, 1), 8) + # enqueue kernel + kernel = _layer_norm_fwd_fused[(M, )]( # + x_arg, y, weight, bias, mean, rstd, # + x_arg.stride(0), N, eps, # + BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_ctas=1) + return y + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_layer_norm(dtype): + M, N = 128, 128 + eps = 1e-5 + device = 'npu' + + x_shape = (M, N) + w_shape = (x_shape[-1], ) + weight = torch.rand(w_shape, dtype=dtype, device=device) + bias = torch.rand(w_shape, dtype=dtype, device=device) + x = -2.3 + 0.5 * torch.randn(x_shape, dtype=dtype, device=device) + + y_tri = layer_norm(x, w_shape, weight, bias, eps) + y_ref = torch.nn.functional.layer_norm(x, w_shape, weight, bias, eps).to(dtype) + + assert torch.allclose(y_tri, y_ref, atol=1e-2, rtol=0) diff --git a/third_party/ascend/unittest/pytest_ut/test_06_fused_attention.py b/third_party/ascend/unittest/pytest_ut/test_06_fused_attention.py new file mode 100644 index 0000000000..8f83d6a730 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_06_fused_attention.py @@ -0,0 +1,340 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Fused Attention +=============== + +This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao (https://tridao.me/publications/flash2/flash2.pdf) + +Credits: OpenAI kernel team + +Extra Credits: + +* Original flash attention paper (https://arxiv.org/abs/2205.14135) +* Rabe and Staats (https://arxiv.org/pdf/2112.05682v2.pdf) + +""" + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl +import triton.language.extra.cann.extension as extension + +DEVICE = "npu" + + +@triton.jit +def _attn_fwd_inner(acc_ptr, l_i, m_i, q, # Accumulator, local l, local m, query vector + K_block_ptr, V_block_ptr, # Key and value block pointers for current stage + start_m, qk_scale, # Starting position of current query block, qk scale factor + BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr, BLOCK_N: tl.constexpr, # Block size constants + STAGE: tl.constexpr, offs_m: tl.constexpr, + offs_n: tl.constexpr, # Current stage flag, m and n offset indices + N_CTX: tl.constexpr, + fp8_v: tl.constexpr): # Total context length, whether to enable FP8 for value precision + # Set the processing range [lo, hi) for the current stage (in column block units) + # Causal attention, as the name implies, restricts the flow of information during computation, + # only allowing the model to see the current and previous positions. + # In other words, the output at the current position can only depend on the input at or before this position, + # and cannot access information from future positions. + # Causal attention ensures sequential order and prevents "leakage of future information." + # But the following logic will also be triggered + if STAGE == 1: + # Stage 1: process all tokens before the query block + tl.static_assert(BLOCK_M >= BLOCK_N) + lo, hi = 0, start_m * BLOCK_M + elif STAGE == 2: + # Stage 2: process the current query block + tl.static_assert(BLOCK_M >= BLOCK_N) + lo, hi = start_m * BLOCK_M, (start_m + 1) * BLOCK_M + lo = tl.multiple_of(lo, BLOCK_M) # Align starting position + # causal = False (no need for masking) + else: + lo, hi = 0, N_CTX # Process the entire context + + # Adjust K and V block pointers to the starting position `lo` + K_block_ptr = tl.advance(K_block_ptr, (lo, 0)) # K is [HEAD_DIM, N_CTX], shift along the second dim by lo + V_block_ptr = tl.advance(V_block_ptr, (lo, 0)) # V is [N_CTX, HEAD_DIM], shift along the first dim by lo + + # Index mapping for the accumulator , used for slicing when HEAD_DIM >= 256 + row = tl.arange(0, BLOCK_M)[:, None] + col_head_dim = tl.arange(0, HEAD_DIM)[None, :] + block2d_acc = row * HEAD_DIM + col_head_dim + + # Iterate over all k, v blocks in the current stage and accumulate the output + for start_n in range(lo, hi, BLOCK_N): # Process BLOCK_N columns at a time + start_n = tl.multiple_of(start_n, BLOCK_N) # Align column start position + # -- Compute qk ---- + k = tl.load(K_block_ptr) + # Modify K + trans_k = tl.trans(k) + qk = tl.dot(q, trans_k) + # Apply causal mask for STAGE 2 + if STAGE == 2: + mask = offs_m[:, None] >= (start_n + offs_n[None, :]) # Construct upper triangular mask + qk = qk * qk_scale + tl.where(mask, 0, -1.0e6) # Set invalid positions to -∞ + m_ij = tl.maximum(m_i, tl.max(qk, 1)) # Update m_ij = max(m_i, max(qk)) + qk -= m_ij[:, None] # Subtract max for softmax stability + else: + qk = qk * qk_scale + m_ij = tl.maximum(m_i, tl.max(qk, 1)) # Scaled max + qk = qk - m_ij[:, None] # Stabilize + + # Softmax weights p = exp(qk) + p = tl.math.exp(qk) + + # Convert softmax weight type depending on FP8 usage + if fp8_v: + p_cast = p.to(tl.float8e5) # Convert to FP8 format (save memory) + else: + p_cast = p.to(k.dtype) + + v = tl.load(V_block_ptr) # Load corresponding V block + pv = tl.dot(p_cast, v) + l_ij = tl.sum(p, 1) # Softmax denominator (sum of each row) + # -- Update m_i and l_i + alpha = tl.math.exp(m_i - m_ij) # Update factor: exp difference between old and new max + l_i = l_i * alpha + l_ij # Update softmax denominator + # -- Update output accumulator -- + if HEAD_DIM < 256: + acc_ptr = acc_ptr * alpha[:, None] + acc_ptr = tl.dot(p_cast, v, acc_ptr) + else: + # 1. Load current slice of accumulator + acc = tl.load(acc_ptr + block2d_acc) + # 2. Update in slices (split by 1/4 of BLOCK_M to avoid ub overflow) + for i in range(4): + # Calculate start/end rows for current slice + offset = i * (BLOCK_M // 4) + # Extract slice data + acc_i = extension.extract_slice(acc, (offset, 0), (BLOCK_M // 4, HEAD_DIM), (1, 1)) + alpha_i = extension.extract_slice(alpha, [offset], [BLOCK_M // 4], [1]) + pv_i = extension.extract_slice(pv, (offset, 0), (BLOCK_M // 4, HEAD_DIM), (1, 1)) + # Incrementally update slice: acc = acc * alpha + pv + acc_i = acc_i * alpha_i[:, None] + pv_i + # Write updated slice back to accumulator + acc = extension.insert_slice(acc, acc_i, (offset, 0), (BLOCK_M // 4, HEAD_DIM), (1, 1)) + # 3. updated accumulator + tl.store(acc_ptr + block2d_acc, acc) + + m_i = m_ij # Update current block max + # Advance V and K block pointers to next BLOCK_N range + V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0)) + K_block_ptr = tl.advance(K_block_ptr, (BLOCK_N, 0)) + # Return accumulated output acc_ptr, softmax denominator l_i, and max value m_i + return acc_ptr, l_i, m_i + + +@triton.jit +def _attn_fwd(Q, K, V, M, Out, acc, sm_scale, stride_qz: tl.constexpr, stride_qh: tl.constexpr, stride_qm: tl.constexpr, + stride_qk: tl.constexpr, stride_kz: tl.constexpr, stride_kh: tl.constexpr, stride_kn: tl.constexpr, + stride_kk: tl.constexpr, stride_vz: tl.constexpr, stride_vh: tl.constexpr, stride_vn: tl.constexpr, + stride_vk: tl.constexpr, stride_oz: tl.constexpr, stride_oh: tl.constexpr, stride_om: tl.constexpr, + stride_on: tl.constexpr, Z: tl.constexpr, H: tl.constexpr, N_CTX: tl.constexpr, HEAD_DIM: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, STAGE: tl.constexpr): + # Total number of blocks in sequence dimension (M) + NUM_BLOCKS_M = N_CTX // BLOCK_M + # Total tasks = number of sequence blocks × batch size (Z) × number of attention heads (H) + NUM_BLOCKS = NUM_BLOCKS_M * Z * H + + # Current M-dimension block index + pid = tl.program_id(0) + + for block_idx in range(pid, NUM_BLOCKS, 20): + task_hz_idx = block_idx // NUM_BLOCKS_M + task_m_idx = block_idx % NUM_BLOCKS_M + off_z = task_hz_idx // H + off_h = task_hz_idx % H + qvk_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh + # Create block pointers for Q, K, V, Output + Q_block_ptr = tl.make_block_ptr( + base=Q + qvk_offset, + shape=(N_CTX, HEAD_DIM), + strides=(stride_qm, stride_qk), + offsets=(task_m_idx * BLOCK_M, 0), + block_shape=(BLOCK_M, HEAD_DIM), + order=(1, 0), + ) + V_block_ptr = tl.make_block_ptr( + base=V + qvk_offset, + shape=(N_CTX, HEAD_DIM), + strides=(stride_vn, stride_vk), + offsets=(0, 0), + block_shape=(BLOCK_N, HEAD_DIM), + order=(1, 0), + ) + K_block_ptr = tl.make_block_ptr( + base=K + qvk_offset, + shape=(N_CTX, HEAD_DIM), + strides=(stride_kn, stride_kk), + offsets=(0, 0), + block_shape=(BLOCK_N, HEAD_DIM), + order=(1, 0), + ) + O_block_ptr = tl.make_block_ptr( + base=Out + qvk_offset, + shape=(N_CTX, HEAD_DIM), + strides=(stride_om, stride_on), + offsets=(task_m_idx * BLOCK_M, 0), + block_shape=(BLOCK_M, HEAD_DIM), + order=(1, 0), + ) + # Initialize offsets + offs_m = task_m_idx * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0 + + # Initialize accumulator + if HEAD_DIM < 256: + acc_ptr = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) + else: + acc_offset = (off_z.to(tl.int64) * stride_qz // stride_qm * HEAD_DIM + + off_h.to(tl.int64) * stride_qh // stride_qm * HEAD_DIM + task_m_idx * BLOCK_M * HEAD_DIM) + acc_ptr = acc + acc_offset + + # load q: it will stay in SRAM throughout + q = tl.load(Q_block_ptr) + + # stage 1: off-band + # For causal = True, STAGE = 3 and _attn_fwd_inner gets 1 as its STAGE + # For causal = False, STAGE = 1, and _attn_fwd_inner gets 3 as its STAGE + if STAGE & 1: + acc_ptr, l_i, m_i = _attn_fwd_inner(acc_ptr, l_i, m_i, q, K_block_ptr, V_block_ptr, # + task_m_idx, sm_scale, # + BLOCK_M, HEAD_DIM, BLOCK_N, # + 4 - STAGE, offs_m, offs_n, N_CTX, V.dtype.element_ty == tl.float8e5 # + ) + # stage 2: on-band + if STAGE & 2: + # barrier makes it easier for compielr to schedule the + # two loops independently + acc_ptr, l_i, m_i = _attn_fwd_inner(acc_ptr, l_i, m_i, q, K_block_ptr, V_block_ptr, # + task_m_idx, sm_scale, # + BLOCK_M, HEAD_DIM, BLOCK_N, # + 2, offs_m, offs_n, N_CTX, V.dtype.element_ty == tl.float8e5 # + ) + + m_i += tl.math.log(l_i) + if HEAD_DIM < 256: + accumulator = acc_ptr / l_i[:, None] + else: + row = tl.arange(0, BLOCK_M)[:, None] + col_head_dim = tl.arange(0, HEAD_DIM)[None, :] + block2d_acc = row * HEAD_DIM + col_head_dim + accumulator = tl.load(acc_ptr + block2d_acc) + accumulator = accumulator / l_i[:, None] + + m_ptrs = M + task_hz_idx * N_CTX + offs_m + + tl.store(m_ptrs, m_i) + tl.store(O_block_ptr, accumulator.to(Out.type.element_ty)) + + +class _attention(torch.autograd.Function): + + @staticmethod + def forward(ctx, q, k, v, causal, sm_scale, BM, BN): + """ + Forward computation interface: + Args: + ctx: Context object + q: Query tensor (Q), shape [Z, H, N_CTX, HEAD_DIM] + k: Key tensor (K), shape [Z, H, N_CTX, HEAD_DIM] + v: Value tensor (V), shape [Z, H, N_CTX, HEAD_DIM] + causal: Whether to enable causal attention + sm_scale: Scaling factor for QK product + BM: Q block size (BLOCK_M) + BN: K/V block size (BLOCK_N) + Returns: + o: Attention output tensor, shape [Z, H, N_CTX, HEAD_DIM] + """ + # shape constraints + HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1] + # when v is in float8_e5m2 it is transposed. + HEAD_DIM_V = v.shape[-1] + assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V + assert HEAD_DIM_K in {16, 32, 64, 128, 256} + + out = torch.empty_like(q) + stage = 3 if causal else 1 + extra_kern_args = {} + + # Number of NPU cores (adjust based on hardware) + num_cores = 20 + acc = torch.zeros((q.shape[0], q.shape[1], q.shape[2], HEAD_DIM_K), dtype=torch.float32, device=q.device) + M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32) + + _attn_fwd[(num_cores, )](q, k, v, M, out, acc, sm_scale, q.stride(0), q.stride(1), q.stride(2), q.stride(3), + k.stride(0), k.stride(1), k.stride(2), k.stride(3), v.stride(0), v.stride(1), + v.stride(2), v.stride(3), out.stride(0), out.stride(1), out.stride(2), out.stride(3), + q.shape[0], q.shape[1], N_CTX=q.shape[2], HEAD_DIM=HEAD_DIM_K, BLOCK_M=BM, BLOCK_N=BN, + STAGE=stage, **extra_kern_args) + + ctx.save_for_backward(q, k, v, out, M) + ctx.sm_scale = sm_scale + ctx.HEAD_DIM = HEAD_DIM_K + ctx.causal = causal + return out + + +attention = _attention.apply + + +# ==================== Pytest Test ==================== +@pytest.mark.parametrize("Z, H, N_CTX, HEAD_DIM, causal, dtype, BM, BN", [ + (1, 1, 128, 128, False, torch.float16, 32, 128), + (1, 1, 128, 128, False, torch.bfloat16, 64, 128), + (1, 2, 256, 256, False, torch.bfloat16, 32, 256), + (2, 2, 128, 256, False, torch.float16, 64, 128), + (4, 32, 64, 64, False, torch.float16, 32, 64), + (4, 32, 1024, 64, False, torch.bfloat16, 64, 128), + (4, 32, 4096, 64, False, torch.float16, 128, 128), +]) +def test_attention_fused(Z, H, N_CTX, HEAD_DIM, causal, dtype, BM, BN): + if N_CTX % BM != 0 or N_CTX % BN != 0 or HEAD_DIM % 16 != 0: + pytest.skip("Skipping non-divisible case") + + torch.manual_seed(20) + q = torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device=DEVICE).normal_(mean=0.0, std=0.5).requires_grad_() + k = torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device=DEVICE).normal_(mean=0.0, std=0.5).requires_grad_() + v = torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device=DEVICE).normal_(mean=0.0, std=0.5).requires_grad_() + + sm_scale = 0.5 + tri_out = attention(q, k, v, causal, sm_scale, BM, BN) + ref_out = torch_npu.npu_fusion_attention( + q, + k, + v, + H, + padding_mask=None, + atten_mask=None, + scale=sm_scale, + keep_prob=1.0, + input_layout="BNSD", + pre_tockens=65535, + next_tockens=65535, + sparse_mode=0, + )[0] + + torch.testing.assert_close(ref_out, tri_out, atol=1e-2, rtol=1e-2, equal_nan=True) diff --git a/third_party/ascend/unittest/pytest_ut/test_07_extern_functions.py b/third_party/ascend/unittest/pytest_ut/test_07_extern_functions.py new file mode 100644 index 0000000000..8ffc40242e --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_07_extern_functions.py @@ -0,0 +1,85 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Libdevice (`tl.extra.libdevice`) function +============================== +""" + +import pytest +import torch +import torch_npu + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +from triton.backends.ascend.compiler import get_libdevice + +DEV = "npu" + + +@triton.jit +def asin_kernel( + x_ptr, + y_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + x = libdevice.asin(x) + tl.store(y_ptr + offsets, x, mask=mask) + + +extern_libs = {'libdevice': get_libdevice()} + + +def run_asin_case(size, use_extern_libs): + torch.manual_seed(0) + x = torch.rand(size, device=DEV) + output_triton = torch.empty_like(x) + output_torch = torch.asin(x) + assert x.device.type == DEV and output_triton.device.type == DEV + + n_elements = output_torch.numel() + + def grid(meta): + return (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + + launch_kwargs = {"BLOCK_SIZE": 1024} + if use_extern_libs: + launch_kwargs["extern_libs"] = extern_libs + + asin_kernel[grid](x, output_triton, n_elements, **launch_kwargs) + torch.testing.assert_close(output_torch, output_triton, rtol=1e-4, atol=1e-4) + + +@pytest.mark.skip(reason="Precision issue with libdevice.asin, it will be fixed later.") +@pytest.mark.parametrize("size", [98432, 1024]) +def test_asin_kernel_matches_torch(size): + run_asin_case(size=size, use_extern_libs=False) + + +@pytest.mark.skip(reason="Precision issue with libdevice.asin, it will be fixed later.") +@pytest.mark.parametrize("size", [98432, 1024]) +def test_asin_kernel_matches_torch_with_extern_libs(size): + run_asin_case(size=size, use_extern_libs=True) diff --git a/third_party/ascend/unittest/pytest_ut/test_08_grouped_gemm.py b/third_party/ascend/unittest/pytest_ut/test_08_grouped_gemm.py new file mode 100644 index 0000000000..830f9e57d6 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_08_grouped_gemm.py @@ -0,0 +1,287 @@ +# Copyright (c) 2023 NVIDIA Corporation & Affiliates. All rights reserved. +# Copyright (c) Huawei Technologies Co., Ltd. 2025. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Group GEMM +============================ +""" + +import pytest +import torch +import torch_npu + +import triton +import triton.language as tl +import triton.runtime.driver as driver + +DEV = "npu" + + +def get_npu_properties(): + device = torch.npu.current_device() + return driver.active.utils.get_device_properties(device) + + +NUM_CORES = get_npu_properties()["num_aicore"] + + +@triton.autotune( + configs=[ + triton.Config({ + 'BLOCK_SIZE_M': 128, + 'BLOCK_SIZE_N': 128, + 'BLOCK_SIZE_K': 32, + 'NUM_SM': NUM_CORES, + }), + triton.Config({ + 'BLOCK_SIZE_M': 128, + 'BLOCK_SIZE_N': 128, + 'BLOCK_SIZE_K': 32, + 'NUM_SM': NUM_CORES, + }), + triton.Config({ + 'BLOCK_SIZE_M': 64, + 'BLOCK_SIZE_N': 64, + 'BLOCK_SIZE_K': 32, + 'NUM_SM': NUM_CORES, + }), + triton.Config({ + 'BLOCK_SIZE_M': 64, + 'BLOCK_SIZE_N': 64, + 'BLOCK_SIZE_K': 32, + 'NUM_SM': NUM_CORES, + }), + ], + key=['group_size'], +) +@triton.jit +def grouped_matmul_kernel( + group_a_ptrs, + group_b_ptrs, + group_c_ptrs, + group_gemm_sizes, + g_lds, + group_size, + NUM_SM: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, +): + tile_idx = tl.program_id(0) + last_problem_end = 0 + for g in range(group_size): + gm = tl.load(group_gemm_sizes + g * 3) + gn = tl.load(group_gemm_sizes + g * 3 + 1) + gk = tl.load(group_gemm_sizes + g * 3 + 2) + num_m_tiles = tl.cdiv(gm, BLOCK_SIZE_M) + num_n_tiles = tl.cdiv(gn, BLOCK_SIZE_N) + num_tiles = num_m_tiles * num_n_tiles + while (tile_idx >= last_problem_end and tile_idx < last_problem_end + num_tiles): + k = gk + lda = tl.load(g_lds + g * 3) + ldb = tl.load(g_lds + g * 3 + 1) + ldc = tl.load(g_lds + g * 3 + 2) + a_ptr = tl.load(group_a_ptrs + g).to(tl.pointer_type(tl.float16)) + b_ptr = tl.load(group_b_ptrs + g).to(tl.pointer_type(tl.float16)) + c_ptr = tl.load(group_c_ptrs + g).to(tl.pointer_type(tl.float16)) + tile_idx_in_gemm = tile_idx - last_problem_end + tile_m_idx = tile_idx_in_gemm // num_n_tiles + tile_n_idx = tile_idx_in_gemm % num_n_tiles + + offs_am = tile_m_idx * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_bn = tile_n_idx * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + offs_am[:, None] * lda + offs_k[None, :] + b_ptrs = b_ptr + offs_k[:, None] * ldb + offs_bn[None, :] + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for _ in range(0, tl.cdiv(k, BLOCK_SIZE_K)): + tl.multiple_of(a_ptrs, [16, 16]) + tl.multiple_of(b_ptrs, [16, 16]) + a = tl.load(a_ptrs) + b = tl.load(b_ptrs) + accumulator += tl.dot(a, b) + a_ptrs += BLOCK_SIZE_K + b_ptrs += BLOCK_SIZE_K * ldb + c = accumulator.to(tl.float16) + + offs_cm = tile_m_idx * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = tile_n_idx * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + ldc * offs_cm[:, None] + offs_cn[None, :] + + tl.store(c_ptrs, c) + tile_idx += NUM_SM + + last_problem_end = last_problem_end + num_tiles + + +def group_gemm_fn(group_A, group_B): + device = torch.device(DEV) + assert len(group_A) == len(group_B) + group_size = len(group_A) + + A_addrs = [] + B_addrs = [] + C_addrs = [] + g_sizes = [] + g_lds = [] + group_C = [] + for i in range(group_size): + A = group_A[i] + B = group_B[i] + assert A.shape[1] == B.shape[0] + M, K = A.shape + K, N = B.shape + C = torch.empty((M, N), device=device, dtype=A.dtype) + group_C.append(C) + A_addrs.append(A.data_ptr()) + B_addrs.append(B.data_ptr()) + C_addrs.append(C.data_ptr()) + g_sizes += [M, N, K] + g_lds += [A.stride(0), B.stride(0), C.stride(0)] + + d_a_ptrs = torch.tensor(A_addrs, device=device) + d_b_ptrs = torch.tensor(B_addrs, device=device) + d_c_ptrs = torch.tensor(C_addrs, device=device) + d_g_sizes = torch.tensor(g_sizes, dtype=torch.int32, device=device) + d_g_lds = torch.tensor(g_lds, dtype=torch.int32, device=device) + + def grid(meta): + return (meta['NUM_SM'], ) + + grouped_matmul_kernel[grid]( + d_a_ptrs, + d_b_ptrs, + d_c_ptrs, + d_g_sizes, + d_g_lds, + group_size, + ) + + return group_C + + +def build_group_inputs(group_m, group_n, group_k): + assert len(group_m) == len(group_n) + assert len(group_n) == len(group_k) + + group_A = [] + group_B = [] + for m, n, k in zip(group_m, group_n, group_k): + group_A.append(torch.rand((m, k), device=DEV, dtype=torch.float16)) + group_B.append(torch.rand((k, n), device=DEV, dtype=torch.float16)) + return group_A, group_B + + +def run_group_gemm_case(group_m, group_n, group_k): + group_A, group_B = build_group_inputs(group_m, group_n, group_k) + + tri_out = group_gemm_fn(group_A, group_B) + ref_out = [torch.matmul(a, b) for a, b in zip(group_A, group_B)] + + assert len(tri_out) == len(ref_out) + for tri_tensor, ref_tensor, m, n in zip(tri_out, ref_out, group_m, group_n): + assert tri_tensor.shape == (m, n) + assert tri_tensor.dtype == torch.float16 + torch.testing.assert_close(ref_tensor, tri_tensor, atol=1e-2, rtol=1e-3) + + +@pytest.mark.parametrize( + "group_m,group_n,group_k", + [([1024, 512, 256, 128], [1024, 512, 256, 128], [1024, 512, 256, 128])], +) +def test_grouped_gemm_tutorial_example(group_m, group_n, group_k): + run_group_gemm_case( + group_m=group_m, + group_n=group_n, + group_k=group_k, + ) + + +def triton_perf_fn(a_ptrs, b_ptrs, c_ptrs, sizes, lds, group_size): + + def grid(meta): + return (meta['NUM_SM'], ) + + grouped_matmul_kernel[grid]( + a_ptrs, + b_ptrs, + c_ptrs, + sizes, + lds, + group_size, + ) + + +def torch_perf_fn(group_A, group_B): + for a, b in zip(group_A, group_B): + torch.matmul(a, b) + + +def run_benchmark_case(N, provider): + group_size = 4 + group_A = [] + group_B = [] + group_C = [] + A_addrs = [] + B_addrs = [] + C_addrs = [] + g_sizes = [] + g_lds = [] + for _ in range(group_size): + A = torch.rand((N, N), device=DEV, dtype=torch.float16) + B = torch.rand((N, N), device=DEV, dtype=torch.float16) + C = torch.empty((N, N), device=DEV, dtype=torch.float16) + group_A.append(A) + group_B.append(B) + group_C.append(C) + A_addrs.append(A.data_ptr()) + B_addrs.append(B.data_ptr()) + C_addrs.append(C.data_ptr()) + g_sizes += [N, N, N] + g_lds += [N, N, N] + + d_a_ptrs = torch.tensor(A_addrs, device=DEV) + d_b_ptrs = torch.tensor(B_addrs, device=DEV) + d_c_ptrs = torch.tensor(C_addrs, device=DEV) + d_g_sizes = torch.tensor(g_sizes, dtype=torch.int32, device=DEV) + d_g_lds = torch.tensor(g_lds, dtype=torch.int32, device=DEV) + + quantiles = [0.5, 0.2, 0.8] + + def bench_torch(): + torch_perf_fn(group_A, group_B) + + def bench_triton(): + triton_perf_fn(d_a_ptrs, d_b_ptrs, d_c_ptrs, d_g_sizes, d_g_lds, group_size) + + if provider == 'torch': + ms, min_ms, max_ms = triton.testing.do_bench(bench_torch, quantiles=quantiles) + if provider == 'triton': + ms, min_ms, max_ms = triton.testing.do_bench(bench_triton, quantiles=quantiles) + + assert ms >= 0 + assert min_ms >= 0 + assert max_ms >= 0 + + +@pytest.mark.parametrize("N", [2**i for i in range(7, 11)]) +@pytest.mark.parametrize("provider", ["torch", "triton"]) +def test_grouped_gemm_benchmark_cases(N, provider): + run_benchmark_case(N=N, provider=provider) diff --git a/third_party/ascend/unittest/pytest_ut/test_09_persistent_matmul.py b/third_party/ascend/unittest/pytest_ut/test_09_persistent_matmul.py new file mode 100644 index 0000000000..f73a3fc46f --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_09_persistent_matmul.py @@ -0,0 +1,325 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Persistent Matmul +===================== +""" + +import time + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl +import triton.runtime.driver as driver + +DEV = "npu" +DTYPE = torch.float16 + + +def get_npu_properties(): + device = torch.npu.current_device() + return driver.active.utils.get_device_properties(device) + + +def get_num_compute_cores(): + return get_npu_properties()["num_aicore"] + + +def _matmul_launch_metadata(grid, kernel, args): + ret = {} + M, N, K = args["M"], args["N"], args["K"] + bytes_per_elem = args["c_ptr"].element_size() + ret["name"] = f"{kernel.name} [M={M}, N={N}, K={K}]" + ret[f"flops{bytes_per_elem * 8}"] = 2.0 * M * N * K + ret["bytes"] = bytes_per_elem * (M * K + N * K + M * N) + return ret + + +@triton.jit(launch_metadata=_matmul_launch_metadata) +def matmul_kernel( + a_ptr, + b_ptr, + c_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + start_m = pid_m * BLOCK_SIZE_M + start_n = pid_n * BLOCK_SIZE_N + + offs_am = start_m + tl.arange(0, BLOCK_SIZE_M) + offs_bn = start_n + tl.arange(0, BLOCK_SIZE_N) + offs_am = tl.where(offs_am < M, offs_am, 0) + offs_bn = tl.where(offs_bn < N, offs_bn, 0) + + offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M) + offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + accumulator = tl.dot(a, b, accumulator) + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + c = accumulator.to(tl.float16) + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +@triton.jit(launch_metadata=_matmul_launch_metadata) +def matmul_kernel_persistent( + a_ptr, + b_ptr, + c_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_SMS: tl.constexpr, +): + start_pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + k_tiles = tl.cdiv(K, BLOCK_SIZE_K) + num_tiles = num_pid_m * num_pid_n + + tiles_per_sm = num_tiles // NUM_SMS + if start_pid < num_tiles % NUM_SMS: + tiles_per_sm += 1 + + tile_id = start_pid - NUM_SMS + ki = -1 + offs_k_for_mask = tl.arange(0, BLOCK_SIZE_K) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + + pid_m = 0 + pid_n = 0 + offs_am = tl.arange(0, BLOCK_SIZE_M) + offs_bn = tl.arange(0, BLOCK_SIZE_N) + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + for _ in range(0, k_tiles * tiles_per_sm): + ki = tl.where(ki == k_tiles - 1, 0, ki + 1) + if ki == 0: + tile_id += NUM_SMS + group_id = tile_id // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (tile_id % group_size_m) + pid_n = (tile_id % num_pid_in_group) // group_size_m + + start_m = pid_m * BLOCK_SIZE_M + start_n = pid_n * BLOCK_SIZE_N + offs_am = start_m + tl.arange(0, BLOCK_SIZE_M) + offs_bn = start_n + tl.arange(0, BLOCK_SIZE_N) + offs_am = tl.where(offs_am < M, offs_am, 0) + offs_bn = tl.where(offs_bn < N, offs_bn, 0) + offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M) + offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N) + + offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + a = tl.load(a_ptrs, mask=offs_k_for_mask[None, :] < K - ki * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k_for_mask[:, None] < K - ki * BLOCK_SIZE_K, other=0.0) + accumulator = tl.dot(a, b, accumulator) + + if ki == k_tiles - 1: + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + c = accumulator.to(tl.float16) + tl.store(c_ptrs, c, mask=c_mask) + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + +def get_configs(dtype): + return { + torch.float16: { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_stages": 3, + "num_warps": 8, + } + }[dtype] + + +def matmul(a, b): + assert a.shape[1] == b.shape[0], "Incompatible dimensions" + assert a.dtype == b.dtype, "Incompatible dtypes" + M, K = a.shape + _, N = b.shape + configs = get_configs(a.dtype) + c = torch.empty((M, N), device=a.device, dtype=a.dtype) + + def grid(meta): + return (triton.cdiv(M, meta["BLOCK_SIZE_M"]) * triton.cdiv(N, meta["BLOCK_SIZE_N"]), ) + + matmul_kernel[grid]( + a, + b, + c, + M, + N, + K, + a.stride(0), + a.stride(1), + b.stride(0), + b.stride(1), + c.stride(0), + c.stride(1), + BLOCK_SIZE_M=configs["BLOCK_SIZE_M"], + BLOCK_SIZE_N=configs["BLOCK_SIZE_N"], + BLOCK_SIZE_K=configs["BLOCK_SIZE_K"], + GROUP_SIZE_M=configs["GROUP_SIZE_M"], + num_stages=configs["num_stages"], + num_warps=configs["num_warps"], + ) + return c + + +def matmul_persistent(a, b): + assert a.shape[1] == b.shape[0], "Incompatible dimensions" + assert a.dtype == b.dtype, "Incompatible dtypes" + num_sms = get_num_compute_cores() + M, K = a.shape + _, N = b.shape + configs = get_configs(a.dtype) + c = torch.empty((M, N), device=a.device, dtype=a.dtype) + + def grid(meta): + return (min(num_sms, triton.cdiv(M, meta["BLOCK_SIZE_M"]) * triton.cdiv(N, meta["BLOCK_SIZE_N"])), ) + + matmul_kernel_persistent[grid]( + a, + b, + c, + M, + N, + K, + a.stride(0), + a.stride(1), + b.stride(0), + b.stride(1), + c.stride(0), + c.stride(1), + BLOCK_SIZE_M=configs["BLOCK_SIZE_M"], + BLOCK_SIZE_N=configs["BLOCK_SIZE_N"], + BLOCK_SIZE_K=configs["BLOCK_SIZE_K"], + GROUP_SIZE_M=configs["GROUP_SIZE_M"], + NUM_SMS=num_sms, + num_stages=configs["num_stages"], + num_warps=configs["num_warps"], + ) + return c + + +def torch_matmul(a, b): + return torch.matmul(a, b) + + +def bench(K, reps=10): + M = 8192 + N = 8192 + a = torch.randn((M, K), device=DEV, dtype=DTYPE) + b = torch.randn((K, N), device=DEV, dtype=DTYPE) + + for _ in range(reps): + _ = torch_matmul(a, b) + time.sleep(0.01) + for _ in range(reps): + _ = matmul(a, b) + time.sleep(0.01) + for _ in range(reps): + _ = matmul_persistent(a, b) + time.sleep(0.01) + + +def validate(M, N, K): + a = torch.randn((M, K), device=DEV, dtype=DTYPE) + b = torch.randn((K, N), device=DEV, dtype=DTYPE) + + torch_result = torch_matmul(a, b) + naive_result = matmul(a, b) + persistent_result = matmul_persistent(a, b) + return torch_result, naive_result, persistent_result + + +@pytest.mark.skip(reason="temporarily skip persistent matmul validate cases until UB overflow issue is fixed") +@pytest.mark.parametrize( + "M,N,K", + [ + (32, 32, 32), + (8192, 8192, 512), + ], +) +def test_persistent_matmul_validate_cases(M, N, K): + torch.manual_seed(0) + torch_result, naive_result, persistent_result = validate(M, N, K) + + torch.testing.assert_close(naive_result, torch_result, atol=1.0, rtol=0) + torch.testing.assert_close(persistent_result, torch_result, atol=1.0, rtol=0) + torch.testing.assert_close(naive_result, persistent_result, atol=1.0, rtol=0) diff --git a/third_party/ascend/unittest/pytest_ut/test_10_gather_sorted.py b/third_party/ascend/unittest/pytest_ut/test_10_gather_sorted.py new file mode 100644 index 0000000000..8fcc3f6b9e --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_10_gather_sorted.py @@ -0,0 +1,201 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Gather sorted +=============== +This is an example only for npu. +""" + +import pytest +import torch +import torch_npu +import triton +import triton.runtime.driver as driver +import triton.language as tl + + +# get device properties of npu +def get_npu_properties(): + device = torch.npu.current_device() + return driver.active.utils.get_device_properties(device) + + +# a torch-version gather_sorted benchmark +def torch_gather_sorted(embeddings, sorted_idxes, aux_idxes): + # make the result tensor + res = torch.empty((aux_idxes.shape[0], embeddings.shape[-1]), dtype=embeddings.dtype, device=embeddings.device) + + # scatter embeddings + res[aux_idxes] = embeddings[sorted_idxes] + + return res + + +# triton-version gather_sorted's kernel +@triton.jit +def gather_sorted_kernel(embeddings_ptr, sorted_indices_ptr, aux_indices_ptr, res_ptr, rows, cols, + DEFAULT_VALUE: tl.constexpr, BIG_CORE_NUM: tl.constexpr, BIG_ROW_BLOCK_SIZE: tl.constexpr, + COL_BLOCK_SIZE: tl.constexpr, COL_BLOCK_SIZE_SUB: tl.constexpr): + SMALL_ROW_BLOCK_SIZE = BIG_ROW_BLOCK_SIZE - 1 + + emb_dtype = embeddings_ptr.type.element_ty + default_value = tl.cast(DEFAULT_VALUE, dtype=emb_dtype) + + core_idx = tl.program_id(0) + # compute the the size and start index of block + row_block_size = BIG_ROW_BLOCK_SIZE if (core_idx < BIG_CORE_NUM) else SMALL_ROW_BLOCK_SIZE + row_start_idx = (core_idx * BIG_ROW_BLOCK_SIZE) if (core_idx < BIG_CORE_NUM) else ( + BIG_CORE_NUM * BIG_ROW_BLOCK_SIZE + (core_idx - BIG_CORE_NUM) * SMALL_ROW_BLOCK_SIZE) + + # this version has 3-buffers, initilize for buffers + row_block_size_0 = tl.cdiv(row_block_size, 3) + remain_row_block_size = row_block_size - row_block_size_0 + row_block_size_1 = tl.cdiv(remain_row_block_size, 2) + row_block_size_2 = remain_row_block_size - row_block_size_1 + + row_start_idx_0 = row_start_idx + row_start_idx_1 = row_start_idx + row_block_size_0 + row_start_idx_2 = row_start_idx + row_block_size_0 + row_block_size_1 + + # process blocks witn shape (row_block_size, COL_BLOCK_SIZE_SUB) one by one + for col_idx in tl.range(0, COL_BLOCK_SIZE, COL_BLOCK_SIZE_SUB): + + embedding_0 = tl.full((COL_BLOCK_SIZE_SUB, ), default_value, dtype=emb_dtype) + embedding_1 = embedding_0 + 0 + embedding_2 = embedding_0 + 0 + + emb_offsets = col_idx + tl.arange(0, COL_BLOCK_SIZE_SUB) + emb_mask = emb_offsets < cols + + prev_embedding_idx_0 = tl.cast(-1, dtype=tl.int32) + prev_embedding_idx_1 = tl.cast(-1, dtype=tl.int32) + prev_embedding_idx_2 = tl.cast(-1, dtype=tl.int32) + for row_idx in tl.range(row_start_idx_0, row_start_idx_1): + # process the first buffer + embedding_idx_0 = tl.load(sorted_indices_ptr + row_idx) + res_idx_0 = tl.load(aux_indices_ptr + row_idx) + + if (embedding_idx_0 != 0) and (embedding_idx_0 != prev_embedding_idx_0): + embedding_0 = tl.load(embeddings_ptr + embedding_idx_0 * cols + emb_offsets, emb_mask) + tl.store(res_ptr + res_idx_0 * cols + emb_offsets, embedding_0, emb_mask) + else: + tl.store(res_ptr + res_idx_0 * cols + emb_offsets, embedding_0, emb_mask) + + prev_embedding_idx_0 = embedding_idx_0 + + # process the second buffer + if (row_idx + row_block_size_0) < (row_start_idx_1 + row_block_size_1): + embedding_idx_1 = tl.load(sorted_indices_ptr + row_idx + row_block_size_0) + res_idx_1 = tl.load(aux_indices_ptr + row_idx + row_block_size_0) + + if (embedding_idx_1 != 0) and (embedding_idx_1 != prev_embedding_idx_1): + embedding_1 = tl.load(embeddings_ptr + embedding_idx_1 * cols + emb_offsets, emb_mask) + tl.store(res_ptr + res_idx_1 * cols + emb_offsets, embedding_1, emb_mask) + else: + tl.store(res_ptr + res_idx_1 * cols + emb_offsets, embedding_1, emb_mask) + + prev_embedding_idx_1 = embedding_idx_1 + + # process the third buffer + if (row_idx + row_block_size_0 + row_block_size_1) < (row_start_idx_2 + row_block_size_2): + embedding_idx_2 = tl.load(sorted_indices_ptr + row_idx + row_block_size_0 + row_block_size_1) + res_idx_2 = tl.load(aux_indices_ptr + row_idx + row_block_size_0 + row_block_size_1) + + if (embedding_idx_2 != 0) and (embedding_idx_2 != prev_embedding_idx_2): + embedding_2 = tl.load(embeddings_ptr + embedding_idx_2 * cols + emb_offsets, emb_mask) + tl.store(res_ptr + res_idx_2 * cols + emb_offsets, embedding_2, emb_mask) + else: + tl.store(res_ptr + res_idx_2 * cols + emb_offsets, embedding_2, emb_mask) + + prev_embedding_idx_2 = embedding_idx_2 + + +# triton-version gather_sorted's host +def triton_gather_sorted(embeddings: torch.Tensor, sorted_indices: torch.Tensor, aux_indices: torch.Tensor, + default_value=1.0): + # constant settings for npu + ALIGNED = 32 + USE_SIZE = 96 * 1024 + CORE_NUM = get_npu_properties()["num_vectorcore"] + + n_rows = sorted_indices.shape[0] + n_cols = embeddings.shape[1] + # make the result tensor + output = torch.empty(n_rows, n_cols, dtype=embeddings.dtype, device=embeddings.device) + + # when writing an npu kernel using triton, + # you should note that the difference between BLOCK_SIZE and BLOCK_SIZE_SUB + # BLOCK_SIZE specifies the size of data that are processed in one program + col_size_aligned = triton.cdiv(embeddings.shape[-1] * embeddings.element_size(), + ALIGNED) * ALIGNED // embeddings.element_size() + # the data are scattered to multiple programs, which can not be even + # some process more data, some process less + big_row_block_size = triton.cdiv(n_rows, CORE_NUM) + big_core_num = CORE_NUM - ((big_row_block_size * CORE_NUM) - n_rows) + col_block_size = col_size_aligned + # BLOCK_SIZE_SUB specifies the size of data that are processed in one loop of a program + col_block_size_sub = min(1024, col_size_aligned) + + grid = (min(n_rows, CORE_NUM), triton.cdiv(n_cols, col_block_size)) + # launch the kernel + gather_sorted_kernel[grid](embeddings, sorted_indices, aux_indices, output, n_rows, n_cols, default_value, + BIG_CORE_NUM=big_core_num, BIG_ROW_BLOCK_SIZE=big_row_block_size, + COL_BLOCK_SIZE=col_block_size, COL_BLOCK_SIZE_SUB=col_block_size_sub) + + return output + + +# genreate the desired inputs +def generate_inputs(index_shape, table_shape, dtype): + sorted_indices = torch.randint(1, table_shape[0], index_shape, dtype=torch.int32).npu() + mask = torch.rand_like(sorted_indices, dtype=torch.float).npu() < 0.2 + + # make sorted_indices + sorted_indices[mask] = 0 + sorted_indices, _ = torch.sort(sorted_indices) + counts = torch.bincount(sorted_indices) + _, _indices = torch.sort(counts[sorted_indices], descending=True, stable=True) + sorted_indices = sorted_indices[_indices] + + # make aux_indicess + aux_indices = torch.arange(0, index_shape[0], dtype=torch.int32).npu() + _indices = torch.randperm(aux_indices.size(0)) + aux_indices = aux_indices[_indices] + + # make table, the first contains only 1.0 + table = torch.randn(table_shape, dtype=dtype).npu() + table[0] = 1.0 + + return table, sorted_indices, aux_indices + + +# ==================== Pytest Test ==================== +@pytest.mark.parametrize("table_rows", [500, 1000]) +@pytest.mark.parametrize("table_cols", [16, 17, 31, 32, 63, 64, 128, 256, 819, 512, 1024, 8192, 1001, 2003, 17000]) +@pytest.mark.parametrize("index_num", [19, 123, 4321, 54321, 100, 200, 819, 500, 700, 1000]) +def test_gather_sorted(table_rows, table_cols, index_num): + table, sorted_indices, aux_indices = generate_inputs((index_num, ), (table_rows, table_cols), torch.float) + + expect = torch_gather_sorted(table, sorted_indices, aux_indices).cpu() + torch.npu.synchronize() + actual = triton_gather_sorted(table, sorted_indices, aux_indices).cpu() + torch.npu.synchronize() + + torch.testing.assert_close(actual, expect) diff --git a/third_party/ascend/unittest/pytest_ut/test_11_rab_time.py b/third_party/ascend/unittest/pytest_ut/test_11_rab_time.py new file mode 100644 index 0000000000..6603a2a457 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_11_rab_time.py @@ -0,0 +1,351 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Relative Attention Bias Timestamps +=============== +""" + +import math +import pytest +import torch +import torch_npu +import triton +import triton.language as tl +import triton.runtime.driver as driver + +NUM_BUCKETS = 128 +BUCKET_DIVISOR = 0.301 + + +# get device properties of npu +def get_npu_properties(): + device = torch.npu.current_device() + return driver.active.utils.get_device_properties(device) + + +def create_pos_w(train_len: int, num_layers: int) -> torch.Tensor: + return torch.arange(0, 2 * train_len + 1).unsqueeze(1).repeat(1, num_layers) + + +def create_past_valid_lens(bs: int, past_len: int) -> torch.Tensor: + return torch.randint(0, past_len, (bs, )) + + +def create_timestamps(train_len: int, candidate_len: int, past_valid_lens: torch.Tensor) -> torch.Tensor: + bs = past_valid_lens.size(0) + timestamps = torch.zeros(bs, train_len + candidate_len // 2) + for i, valid_len in enumerate(past_valid_lens): + if valid_len > 0: + timestamps[i, :valid_len] = torch.arange(1, valid_len.int() + 1) + + if candidate_len <= 0: + return timestamps + timestamps[:, -candidate_len // 2:] = train_len + 1 + + return timestamps + + +def create_timestamps_weights(num_layers: int): + return (torch.arange(0, NUM_BUCKETS + 1).repeat(num_layers).reshape(NUM_BUCKETS + 1, num_layers)) + + +def create_rab_time_grad(num_layers: int, batchsize: int, s: int): + return torch.rand(num_layers, batchsize, s, s) * 1e-4 + + +def create_bucket_timestamps(batchsize: int, s: int): + result = torch.arange(batchsize * s) % NUM_BUCKETS + result = result.unsqueeze(-1).repeat(1, 1, s) + return result + + +@triton.jit +def rab_time_forward_kernel( + inp, + out, + index, + index_len: tl.constexpr, + inp_row_stride: tl.constexpr, + clamp_max: tl.constexpr, + bucketization_divisor: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + COL_BLOCK_SIZE: tl.constexpr, +): + pid0 = tl.program_id(axis=0) + pid1 = tl.program_id(axis=1) + + col_iter_num = tl.cdiv(BLOCK_SIZE, COL_BLOCK_SIZE) + + for col_idx in tl.range(0, col_iter_num): + cols_offsets = (pid0 * BLOCK_SIZE + col_idx * COL_BLOCK_SIZE + tl.arange(0, COL_BLOCK_SIZE)) + cols_mask = cols_offsets < index_len + + out_mask = cols_offsets < index_len + + index_val = tl.load(index + cols_offsets, mask=cols_mask, other=0.0) + index_val = tl.abs(index_val) + index_val = tl.minimum(tl.maximum(index_val, 1.0), clamp_max) + index_val = tl.log(index_val) + index_val = index_val / bucketization_divisor + index_val = tl.cast(index_val, tl.int64) + + inp_val = tl.load(inp + pid1 * inp_row_stride + tl.arange(0, inp_row_stride)) + out_val = tl.gather(inp_val, index_val, 0) + + tl.store(out + pid1 * index_len + cols_offsets, out_val, mask=out_mask) + + +def get_outer_loop_num(num_layers, index_len): + sub_num_layers = num_layers + while sub_num_layers * index_len >= 2**31 - 1: + sub_num_layers = sub_num_layers // 2 + outer_loop_num = (num_layers + sub_num_layers - 1) // sub_num_layers + remain_layers = num_layers % sub_num_layers + return outer_loop_num, sub_num_layers, remain_layers + + +def rab_time_forward_triton(ts_w, timestamps, bucketization_divisor): + ts_w_trans = ts_w.t().contiguous() + + bs, seq_len = timestamps.shape + infer_len = 2 * seq_len + num_layers = ts_w.shape[1] + num_buckets = ts_w.shape[0] - 1 + + timestamps_expanded = timestamps.unsqueeze(-1).repeat(1, 1, 2) + timestamps_expanded = timestamps_expanded.reshape(bs, infer_len, 1) - timestamps_expanded.reshape(bs, 1, infer_len) + + timestamps_expanded = timestamps_expanded.view(-1) + timestamps_expanded = timestamps_expanded.contiguous() + + clamp_max = torch.exp(torch.tensor(num_buckets * bucketization_divisor)).item() + index_len = bs * infer_len * infer_len + + out = torch.empty((num_layers, index_len), dtype=ts_w.dtype, device=ts_w.device) + outer_loop_num, sub_num_layers, remain_layers = get_outer_loop_num(num_layers, index_len) + + CORE_NUM = get_npu_properties()["num_vectorcore"] + BLOCK_SIZE = math.ceil(index_len / CORE_NUM) + COL_BLOCK_SIZE = 8 * 1024 + + curr_layers = sub_num_layers + for i in range(outer_loop_num): + if i == outer_loop_num - 1 and remain_layers != 0: + curr_layers = remain_layers + + def grid(meta): + return (triton.cdiv(index_len, meta["BLOCK_SIZE"]), curr_layers) + + rab_time_forward_kernel[grid]( + ts_w_trans[i * sub_num_layers], + out[i * sub_num_layers], + timestamps_expanded, + index_len, + num_buckets + 1, + clamp_max, + bucketization_divisor, + BLOCK_SIZE, + COL_BLOCK_SIZE, + ) + + out = out.view(num_layers, bs, infer_len, infer_len) + + return out + + +@triton.jit +def rab_time_backward_kernel(inp, src, index, index_len, BLOCK_SIZE: tl.constexpr, COL_BLOCK_SIZE: tl.constexpr): + pid0 = tl.program_id(axis=0) + total_col_num = (BLOCK_SIZE if pid0 * BLOCK_SIZE + BLOCK_SIZE < index_len else index_len - pid0 * BLOCK_SIZE) + COL_BLOCK_SIZE = min(COL_BLOCK_SIZE, total_col_num) + col_iter_num = (total_col_num + COL_BLOCK_SIZE - 1) // COL_BLOCK_SIZE + + for col_idx in tl.range(0, col_iter_num): + base_idx = 0 + base_idx = base_idx.to(index.dtype.element_ty) + + col_start_offset = col_idx * COL_BLOCK_SIZE + + acc_result = 0.0 + acc_result = acc_result.to(inp.dtype.element_ty) + cur_col_num = (COL_BLOCK_SIZE if col_start_offset + COL_BLOCK_SIZE < total_col_num else total_col_num - + col_start_offset) + + for cur_idx in range(0, cur_col_num): + cur_offset = pid0 * BLOCK_SIZE + col_start_offset + cur_idx + + src_val = tl.load(src + cur_offset) + new_idx = tl.load(index + cur_offset) + + if base_idx == new_idx: + acc_result += src_val + else: + tl.atomic_add(inp + base_idx, acc_result) + + base_idx = new_idx + acc_result = 0.0 + acc_result = acc_result.to(inp.dtype.element_ty) + acc_result += src_val + + tl.atomic_add(inp + base_idx, acc_result) + + +def rab_time_backward_triton(rab_time_grad: torch.Tensor, bucket_timestamps: torch.Tensor): + num_layers, b, s, _ = rab_time_grad.shape + tsw_grad = torch.zeros(num_layers, NUM_BUCKETS, dtype=torch.float32).to(rab_time_grad.device) + + bucket_timestamps_expand = (bucket_timestamps.reshape(b, s // 2, 1, s // 2, + 1).repeat(1, 1, 2, 1, 2).reshape(b, s, + s).to(torch.int64)).view(-1) + + index_len = bucket_timestamps_expand.numel() + + rab_time_grad_f32 = rab_time_grad.to(torch.float32) + sorted_bucket_timestamps_expand, sorted_idx = torch.sort(bucket_timestamps_expand.view(-1)) + + torch.npu.synchronize() + + def grid(meta): + return (triton.cdiv(index_len, meta["BLOCK_SIZE"]), ) + + CORE_NUM = get_npu_properties()["num_vectorcore"] + BLOCK_SIZE = math.ceil(index_len / CORE_NUM) + + COL_BLOCK_SIZE = 8 * 1024 + + for layer_idx in range(num_layers): + curr_sorted_grad_f32 = rab_time_grad_f32[layer_idx].view(-1)[sorted_idx] + rab_time_backward_kernel[grid]( + tsw_grad[layer_idx], + curr_sorted_grad_f32, + sorted_bucket_timestamps_expand, + index_len, + BLOCK_SIZE, + COL_BLOCK_SIZE, + ) + + return tsw_grad + + +def rab_time_forward_golden(ts_w: torch.Tensor, timestamps: torch.Tensor, bucketization_divisor: float) -> torch.Tensor: + """ + torch realization of rab time forward for reference. + """ + infer_len = timestamps.shape[1] * 2 + bs = timestamps.shape[0] + num_layers = ts_w.shape[1] + + timestamps = timestamps.unsqueeze(-1).repeat(1, 1, 2) + diff_timestamps = timestamps.reshape(bs, infer_len, 1) - timestamps.reshape(bs, 1, infer_len) + + clamp_max = torch.exp(torch.tensor(NUM_BUCKETS * BUCKET_DIVISOR)) + diff_timestamps = (torch.log(torch.abs(diff_timestamps).clamp(1, clamp_max)) / bucketization_divisor) + bucket_timestamps = diff_timestamps.long() + bucket_timestamps = bucket_timestamps.view(-1) + result = torch.index_select(ts_w, dim=0, index=bucket_timestamps) + + result = result.t() + + result = result.view(num_layers, bs, infer_len, infer_len) + return result + + +def rab_time_backward_golden(rab_time_grad: torch.Tensor, bucket_timestamps: torch.Tensor): + """ + torch realization of rab time backward for reference. + """ + num_layers, b, s, _ = rab_time_grad.shape + tsw_grad = torch.zeros(num_layers, NUM_BUCKETS, dtype=torch.float32).to(rab_time_grad.device) + + bucket_timestamps_expand = (bucket_timestamps.reshape(b, s // 2, 1, s // 2, + 1).repeat(1, 1, 2, 1, 2).reshape(b, s, s).to(torch.int64)) + for n, grad in enumerate(rab_time_grad.to(torch.float32)): + tsw_grad[n] = tsw_grad[n].scatter_add(src=grad.view(-1), index=bucket_timestamps_expand.view(-1), dim=0) + return tsw_grad + + +def run_rab_time_forward_case(num_layers, train_len, candidate_len, bs, dtype): + past_valid_lens = create_past_valid_lens(bs, train_len).to(torch.int32) + timestamps = create_timestamps(train_len, candidate_len, past_valid_lens).to(torch.int32) + timestamps_weights = create_timestamps_weights(num_layers).to(dtype) + timestamps = timestamps.npu() + timestamps_weights = timestamps_weights.npu() + + torch_npu.npu.synchronize() + + # triton output + rab_time_out_triton = rab_time_forward_triton( + ts_w=timestamps_weights, + timestamps=timestamps, + bucketization_divisor=BUCKET_DIVISOR, + ) + torch_npu.npu.synchronize() + + # pytorch output + rab_time_out_golden = rab_time_forward_golden( + ts_w=timestamps_weights, + timestamps=timestamps, + bucketization_divisor=BUCKET_DIVISOR, + ) + torch_npu.npu.synchronize() + + torch.testing.assert_close(rab_time_out_triton, rab_time_out_golden) + + +def run_rab_time_backward_case(num_layers: int, batchsize: int, s: int, dtype: torch.dtype): + grad = create_rab_time_grad(num_layers, batchsize, s).to(dtype).npu() + bucket_timestamps = (create_bucket_timestamps(batchsize, s // 2).to(torch.int32).npu()) + + torch_npu.npu.synchronize() + + golden_result = (rab_time_backward_golden(grad, bucket_timestamps).to(torch.float32).cpu()) + op_result = (rab_time_backward_triton(grad, bucket_timestamps).to(torch.float32).cpu()) + + loss = 1e-4 if dtype == torch.float32 else 1e-3 + torch.testing.assert_close(op_result, golden_result, rtol=loss, atol=loss) + + +@pytest.mark.parametrize( + "num_layers, train_len, candidate_len, batch_size, dtype", + [ + pytest.param( + 8, + 500, + 500, + 4, + torch.float32, + marks=pytest.mark.skip(reason="temporarily skip UB overflow case"), + ), + ], +) +def test_rab_time_forward(num_layers, train_len, candidate_len, batch_size, dtype): + torch.manual_seed(0) + run_rab_time_forward_case(num_layers, train_len, candidate_len, batch_size, dtype) + + +@pytest.mark.parametrize( + "num_layers, batch_size, seq_len, dtype", + [ + (8, 4, 1500, torch.float32), + ], +) +def test_rab_time_backward(num_layers, batch_size, seq_len, dtype): + torch.manual_seed(0) + run_rab_time_backward_case(num_layers, batch_size, seq_len, dtype) diff --git a/third_party/ascend/unittest/pytest_ut/test_12_hstu_attention.py b/third_party/ascend/unittest/pytest_ut/test_12_hstu_attention.py new file mode 100644 index 0000000000..f4fe9b523d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_12_hstu_attention.py @@ -0,0 +1,933 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +HSTU Attention +=============== +""" + +from dataclasses import dataclass +from typing import List, Optional, Tuple +import pytest +import torch +import torch_npu +import triton +import triton.language as tl +import triton.runtime.driver as driver +import numpy as np +import torch.nn.functional as F + +DEVICE = "npu" +BLOCK_FWD = 64 +BLOCK_BWD = 32 + + +@dataclass +class JaggedData: + grad: torch.Tensor + q: torch.Tensor + k: torch.Tensor + v: torch.Tensor + bias: torch.Tensor + mask: torch.Tensor + max_seq_len: int + seq_offset: np.ndarray + + +def get_npu_properties(coreType): + device = torch.npu.current_device() + return driver.active.utils.get_device_properties(device)[coreType] + + +@triton.jit +def _hstu_attn_fwd_one_block( + q, + k_block_ptr, + v_block_ptr, + bias_block_ptr, + alpha, + MAX_SEQ_LEN, + CAUSAL: tl.constexpr, + HAS_BIAS: tl.constexpr, + mask_block, +): + k = tl.load(k_block_ptr) + qk = tl.dot(q, tl.trans(k)) * alpha + if HAS_BIAS: + rel_attn_bias = tl.load(bias_block_ptr) + qk = qk + rel_attn_bias + silu = qk / (1.0 + tl.exp(-qk)) * (1.0 / MAX_SEQ_LEN) + if CAUSAL: + silu = tl.where(mask_block != 0, silu, 0) + v = tl.load(v_block_ptr) + silu = silu.to(v.dtype) + return tl.dot(silu, v) + + +@triton.jit +def _hstu_attn_fwd_compute( # noqa C901 + Q, + K, + V, + seq_offsets, + Out, + stride_qm: tl.constexpr, + stride_qh: tl.constexpr, + stride_kn: tl.constexpr, + stride_kh: tl.constexpr, + stride_vn: tl.constexpr, + stride_vh: tl.constexpr, + stride_om: tl.constexpr, + stride_oh: tl.constexpr, + alpha, + head_num, + MAX_SEQ_LEN, + off_batch, + off_head, + start_m, + seq_start, + seq_len, + CAUSAL: tl.constexpr, + HAS_BIAS: tl.constexpr, + BLOCK_D_Q: tl.constexpr, + BLOCK_D_V: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + mask_block, + bias, +): + off_head = off_head.to(tl.int64) + off_seq = seq_start.to(tl.int64) + start_m = start_m.to(tl.int32) + + # initialize offsets + q_offset = off_seq * stride_qm + off_head * stride_qh + k_offset = off_seq * stride_kn + off_head * stride_kh + v_offset = off_seq * stride_vn + off_head * stride_kh + + Q_block_ptr = tl.make_block_ptr( + base=Q + q_offset, + shape=(seq_len, BLOCK_D_Q), + strides=(stride_qm, 1), + offsets=(start_m, 0), + block_shape=(BLOCK_M, BLOCK_D_Q), + order=(1, 0), + ) + k_block_ptr = tl.make_block_ptr( + base=K + k_offset, + shape=(seq_len, BLOCK_D_Q), + strides=(stride_kn, 1), + offsets=(0, 0), + block_shape=(BLOCK_N, BLOCK_D_Q), + order=(1, 0), + ) + v_block_ptr = tl.make_block_ptr( + base=V + v_offset, + shape=(seq_len, BLOCK_D_V), + strides=(stride_vn, 1), + offsets=(0, 0), + block_shape=(BLOCK_N, BLOCK_D_V), + order=(1, 0), + ) + q = tl.load(Q_block_ptr) + + acc = tl.zeros([BLOCK_M, BLOCK_D_V], dtype=tl.float32) + if CAUSAL: + low = 0 + high = start_m + BLOCK_M + else: + low = 0 + high = seq_len + + bias_block_ptr = None + if HAS_BIAS: + bias_block_ptr = tl.make_block_ptr( + base=bias + off_batch * head_num * MAX_SEQ_LEN * MAX_SEQ_LEN + off_head * MAX_SEQ_LEN * MAX_SEQ_LEN, + shape=(MAX_SEQ_LEN, MAX_SEQ_LEN), + strides=(MAX_SEQ_LEN, 1), + offsets=(start_m, 0), + block_shape=(BLOCK_M, BLOCK_N), + order=(1, 0), + ) + + for start_n in range(low, high, BLOCK_N): + acc += _hstu_attn_fwd_one_block( + q=q, + k_block_ptr=k_block_ptr, + v_block_ptr=v_block_ptr, + bias_block_ptr=bias_block_ptr, + alpha=alpha, + MAX_SEQ_LEN=MAX_SEQ_LEN, + CAUSAL=CAUSAL and start_m == start_n, + HAS_BIAS=HAS_BIAS, + mask_block=mask_block, + ) + k_block_ptr = tl.advance(k_block_ptr, (BLOCK_N, 0)) + v_block_ptr = tl.advance(v_block_ptr, (BLOCK_N, 0)) + if HAS_BIAS: + bias_block_ptr = tl.advance(bias_block_ptr, (0, BLOCK_N)) + + # rematerialize offsets to save registers + offs_v_d = tl.arange(0, BLOCK_D_V) + off_o = Out + off_seq * stride_om + off_head * stride_oh + offs_m = start_m + tl.arange(0, BLOCK_M) + out_ptrs = off_o + offs_m[:, None] * stride_om + offs_v_d[None, :] + tl.store(out_ptrs, acc, mask=(offs_m < seq_len)[:, None]) + + +@triton.jit +def _hstu_attn_fwd( # noqa C901 + Q, + K, + V, + seq_offsets, + Out, + stride_qm: tl.constexpr, + stride_qh: tl.constexpr, + stride_kn: tl.constexpr, + stride_kh: tl.constexpr, + stride_vn: tl.constexpr, + stride_vh: tl.constexpr, + stride_om: tl.constexpr, + stride_oh: tl.constexpr, + alpha: tl.constexpr, + batch: tl.constexpr, + head_num: tl.constexpr, + MAX_SEQ_LEN: tl.constexpr, + head_dim: tl.constexpr, + CAUSAL: tl.constexpr, + HAS_BIAS: tl.constexpr, + CORE_NUM: tl.constexpr, + tasks: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + mask, + bias, +): + core_id = tl.program_id(0) + cur_batch = 0 + mask_block = None + if CAUSAL and mask is not None: + mask_ptr = tl.make_block_ptr( + base=mask, + shape=(MAX_SEQ_LEN, MAX_SEQ_LEN), + strides=(MAX_SEQ_LEN, 1), + offsets=(0, 0), + block_shape=(BLOCK_M, BLOCK_M), + order=(1, 0), + ) + mask_block = tl.load(mask_ptr) + for col in range(core_id, tasks, CORE_NUM): + seq_end = tl.load(seq_offsets + cur_batch + 1) + start_m = col * BLOCK_M + while start_m >= seq_end * head_num // 2: + cur_batch += 1 + seq_end = tl.load(seq_offsets + cur_batch + 1) + seq_start = tl.load(seq_offsets + cur_batch) + seq_len = seq_end - seq_start + off_batch = cur_batch + off_head = (start_m - seq_start * head_num // 2) // (seq_len // 2) + start_m_1 = (start_m - seq_start * head_num // 2) % (seq_len // 2) + start_m_2 = seq_len - start_m_1 - BLOCK_M + _hstu_attn_fwd_compute( + Q, + K, + V, + seq_offsets, + Out, + stride_qm, + stride_qh, + stride_kn, + stride_kh, + stride_vn, + stride_vh, + stride_om, + stride_oh, + alpha, + head_num, + MAX_SEQ_LEN, + off_batch, + off_head, + start_m_1, + seq_start, + seq_len, + CAUSAL, + HAS_BIAS, + head_dim, + head_dim, + BLOCK_M, + BLOCK_N, + mask_block=mask_block, + bias=bias, + ) + _hstu_attn_fwd_compute( + Q, + K, + V, + seq_offsets, + Out, + stride_qm, + stride_qh, + stride_kn, + stride_kh, + stride_vn, + stride_vh, + stride_om, + stride_oh, + alpha, + head_num, + MAX_SEQ_LEN, + off_batch, + off_head, + start_m_2, + seq_start, + seq_len, + CAUSAL, + HAS_BIAS, + head_dim, + head_dim, + BLOCK_M, + BLOCK_N, + mask_block=mask_block, + bias=bias, + ) + + +@triton.jit +def _hstu_attn_bwd_one_block( # noqa C901 + start_m, + offs_n, + offs_m, + q_ptrs, + dq_ptrs, + mask_n, + do_ptrs, + dk, + dv, + k, + v, + pos_offs_n, + seq_len, + max_ids, + stride_qm, + stride_dom, + stride_dqm, + alpha, + MAX_SEQ_LEN, + CAUSAL: tl.constexpr, + HAS_BIAS: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + bias_block_ptr, +): + pos_offs_m = offs_m + start_m + mask_m = pos_offs_m < seq_len + # recompute qk and silu + q = tl.load( + q_ptrs + start_m * stride_qm, + mask=mask_m[:, None], + other=0.0, + ) + q_trans = tl.trans(q) + qk_trans = tl.dot(k, q_trans) * alpha + if HAS_BIAS: + rel_attn_bias = tl.load(bias_block_ptr) + qk_trans = qk_trans + tl.trans(rel_attn_bias) + sig_trans = 1.0 / (1.0 + tl.exp(-qk_trans)) + silu_trans = qk_trans * sig_trans * (1.0 / MAX_SEQ_LEN) + if CAUSAL: + invalid_mask_trans = pos_offs_m[None, :] == offs_n[:, None] + pos_offs_m_minus_n = pos_offs_m[None, :] - pos_offs_n[:, None] + invalid_mask_trans = invalid_mask_trans | (pos_offs_m_minus_n > 0) + silu_trans = tl.where(invalid_mask_trans, silu_trans, 0) + silu_trans = silu_trans.to(k.dtype) + # compute dv + do = tl.load( + do_ptrs + start_m * stride_dom, + mask=mask_m[:, None], + other=0.0, + ) + dv += tl.dot(silu_trans, do) + # compute dk and dq (dqk = do * v^T dk = dqk^T * q dq = dqk * k) + dqk_trans = tl.dot(v, tl.trans(do)) + dqk_trans = dqk_trans * sig_trans * (1 + qk_trans * (1 - sig_trans)) * (1.0 / MAX_SEQ_LEN) + if CAUSAL: + dqk_trans = tl.where(invalid_mask_trans, dqk_trans, 0) + dqk_trans = dqk_trans.to(k.dtype) + dq = tl.load( + dq_ptrs + start_m * stride_dqm, + mask=mask_m[:, None], + other=0.0, + ) + dq += tl.dot(tl.trans(dqk_trans), k) * alpha + tl.store( + dq_ptrs + start_m * stride_dqm, + dq, + mask=mask_m[:, None], + ) + # Note: the factor `alpha` is delayed until the end of the function to reduce the cost + dk += tl.dot(dqk_trans, q) + return dk, dv + + +@triton.jit +def _hstu_attn_bwd_one_col_block( # noqa C901 + start_n, + seq_len, + Q, + K, + V, + DOut, + DQ, + DK, + DV, + stride_qm, + stride_kn, + stride_vn, + stride_dom, + stride_dqm, + stride_dkn, + stride_dvn, + alpha, + MAX_SEQ_LEN, + CAUSAL: tl.constexpr, + HAS_BIAS: tl.constexpr, + BLOCK_D_Q: tl.constexpr, + BLOCK_D_V: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + bias, +): + # Work on the subsequence dv[start_n, start_n + BLOCK_N, :] + if CAUSAL: + low = start_n + high = seq_len + else: + low = 0 + high = seq_len + + # initialize row/col offsets + offs_m = tl.arange(0, BLOCK_M) + offs_qk_d = tl.arange(0, BLOCK_D_Q) + offs_v_d = tl.arange(0, BLOCK_D_V) + offs_n = start_n + tl.arange(0, BLOCK_N) + + dq_ptrs = DQ + (offs_m[:, None] * stride_dqm + offs_qk_d[None, :]) + dk = tl.zeros([BLOCK_N, BLOCK_D_Q], dtype=tl.float32) + dv = tl.zeros([BLOCK_N, BLOCK_D_V], dtype=tl.float32) + + mask_n = offs_n < seq_len + q_ptrs = Q + (offs_m[:, None] * stride_qm + offs_qk_d[None, :]) + do_ptrs = DOut + (offs_m[:, None] * stride_dom + offs_v_d[None, :]) + k_ptrs = K + (offs_n[:, None] * stride_kn + offs_qk_d[None, :]) + v_ptrs = V + (offs_n[:, None] * stride_vn + offs_v_d[None, :]) + k = tl.load(k_ptrs, mask=mask_n[:, None], other=0.0) + v = tl.load(v_ptrs, mask=mask_n[:, None], other=0.0) + + max_ids = seq_len + pos_offs_n = offs_n + # loop over rows + for start_m in tl.range(low, high, BLOCK_M): + bias_block_ptr = None + if HAS_BIAS: + bias_block_ptr = tl.make_block_ptr( + base=bias, + shape=(MAX_SEQ_LEN, MAX_SEQ_LEN), + strides=(MAX_SEQ_LEN, 1), + offsets=(start_m, start_n), + block_shape=(BLOCK_M, BLOCK_N), + order=(1, 0), + ) + start_m = tl.multiple_of(start_m, BLOCK_M) + dk, dv = _hstu_attn_bwd_one_block( + start_m=start_m, + offs_n=offs_n, + offs_m=offs_m, + q_ptrs=q_ptrs, + dq_ptrs=dq_ptrs, + mask_n=mask_n, + do_ptrs=do_ptrs, + dk=dk, + dv=dv, + k=k, + v=v, + pos_offs_n=pos_offs_n, + seq_len=seq_len, + max_ids=max_ids, + stride_qm=stride_qm, + stride_dom=stride_dom, + stride_dqm=stride_dqm, + alpha=alpha, + MAX_SEQ_LEN=MAX_SEQ_LEN, + CAUSAL=CAUSAL, + HAS_BIAS=HAS_BIAS, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + bias_block_ptr=bias_block_ptr, + ) + # write-back + dk = dk * alpha + dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_v_d[None, :]) + dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_qk_d[None, :]) + tl.store(dv_ptrs, dv.to(k.dtype), mask=mask_n[:, None]) + tl.store(dk_ptrs, dk.to(k.dtype), mask=mask_n[:, None]) + + +@triton.jit +def _hstu_attn_bwd( # noqa C901 + Q, + K, + V, + Grad, + DQ, + DK, + DV, + stride_qm: tl.constexpr, + stride_qh: tl.constexpr, + stride_kn: tl.constexpr, + stride_kh: tl.constexpr, + stride_vn: tl.constexpr, + stride_vh: tl.constexpr, + stride_dom: tl.constexpr, + stride_doh: tl.constexpr, + seq_offsets, + alpha: tl.constexpr, + batch: tl.constexpr, + head_num: tl.constexpr, + MAX_SEQ_LEN: tl.constexpr, + head_dim: tl.constexpr, + CAUSAL: tl.constexpr, + HAS_BIAS: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + bias, +): + off = tl.program_id(0) + off_batch = off // head_num + off_head = off % head_num + off_head = off_head.to(tl.int64) + seq_start = tl.load(seq_offsets + off_batch).to(tl.int64) + seq_end = tl.load(seq_offsets + off_batch + 1) + seq_len = (seq_end - seq_start).to(tl.int32) + # offset pointers for batch/head + q_offset = seq_start * stride_qm + off_head * stride_qh + k_offset = seq_start * stride_kn + off_head * stride_kh + v_offset = seq_start * stride_vn + off_head * stride_vh + grad_offset = seq_start * stride_dom + off_head * stride_doh + bias_offset = off_batch * head_num * MAX_SEQ_LEN * MAX_SEQ_LEN + off_head * MAX_SEQ_LEN * MAX_SEQ_LEN + for start_n in range(0, seq_len, BLOCK_N): + _hstu_attn_bwd_one_col_block( + start_n=start_n, + seq_len=seq_len, + Q=Q + q_offset, + K=K + k_offset, + V=V + v_offset, + DOut=Grad + grad_offset, + DQ=DQ + q_offset, + DK=DK + k_offset, + DV=DV + v_offset, + stride_qm=stride_qm, + stride_kn=stride_kn, + stride_vn=stride_vn, + stride_dom=stride_dom, + stride_dqm=stride_qm, + stride_dkn=stride_kn, + stride_dvn=stride_vn, + alpha=alpha, + MAX_SEQ_LEN=MAX_SEQ_LEN, + CAUSAL=CAUSAL, + HAS_BIAS=HAS_BIAS, + BLOCK_D_Q=head_dim, + BLOCK_D_V=head_dim, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + bias=bias + bias_offset if HAS_BIAS else bias, + ) + + +def triton_hstu_attention_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + seq_offsets: torch.Tensor, + max_seq_len: int, + alpha: float, + causal: bool, + mask: torch.Tensor, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + batch = seq_offsets.numel() - 1 + total_seq, head_num, head_dim = q.shape + out = torch.empty_like(v) + BLOCK_M = BLOCK_FWD + BLOCK_N = BLOCK_FWD + if total_seq == 0: + print("error") + return out + has_bias = bias is not None + core_num = get_npu_properties('num_aicore') + tasks = total_seq * head_num // BLOCK_M // 2 + grid = (core_num, 1, 1) + _hstu_attn_fwd[grid]( + q, + k, + v, + seq_offsets, + out, + q.stride(0), + q.stride(1), + k.stride(0), + k.stride(1), + v.stride(0), + v.stride(1), + out.stride(0), + out.stride(1), + alpha, + batch, + head_num, + max_seq_len, + head_dim, + causal, + has_bias, + core_num, + tasks, + BLOCK_M, + BLOCK_N, + mask, + bias, + ) + return out + + +def triton_hstu_attention_bwd( + grad: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + seq_offsets: torch.Tensor, + max_seq_len: int, + alpha: float, + causal: bool, + bias: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + dq = torch.zeros_like(q) + dk = torch.zeros_like(k) + dv = torch.zeros_like(v) + if grad.shape[0] == 0: + return dq, dk, dv + batch = seq_offsets.numel() - 1 + _, head_num, head_dim = q.shape + has_bias = bias is not None + grid = ( + batch * head_num, + 1, + ) + _hstu_attn_bwd[grid]( + q, + k, + v, + grad, + dq, + dk, + dv, + q.stride(0), + q.stride(1), + k.stride(0), + k.stride(1), + v.stride(0), + v.stride(1), + grad.stride(0), + grad.stride(1), + seq_offsets, + alpha, + batch, + head_num, + max_seq_len, + head_dim, + causal, + has_bias, + BLOCK_BWD, + BLOCK_BWD, + bias, + ) + return dq, dk, dv + + +def jagged_data_gen(batch_size, max_seq_len, num_heads, attention_dim, dataType) -> JaggedData: + seq_array = np.arange(256, max_seq_len + 1, 256) + seq_lens = np.random.choice(seq_array, size=batch_size) + if not np.isin(max_seq_len, seq_lens): + seq_lens[np.random.randint(0, batch_size)] = max_seq_len + seq_lens_tensor = torch.from_numpy(seq_lens).to(dtype=torch.int64) + seq_offset = torch.concat(( + torch.zeros((1, ), dtype=torch.int64, device=seq_lens_tensor.device), + torch.cumsum(seq_lens_tensor, axis=0), + )).numpy() + max_seq_len = np.max(seq_lens) + total_seqs = np.sum(seq_lens) + grad = torch.rand((int(total_seqs), num_heads, attention_dim), dtype=dataType) + q = torch.rand((int(total_seqs), num_heads, attention_dim), dtype=dataType) + k = torch.rand((int(total_seqs), num_heads, attention_dim), dtype=dataType) + v = torch.rand((int(total_seqs), num_heads, attention_dim), dtype=dataType) + + bias = torch.empty(batch_size, num_heads, max_seq_len, max_seq_len, dtype=dataType).uniform_(-1, 1) + mask = 1 - torch.triu(torch.ones(batch_size, num_heads, max_seq_len, max_seq_len), diagonal=1).to(torch.float32) + return JaggedData( + grad=grad, + q=q, + k=k, + v=v, + bias=bias, + mask=mask, + max_seq_len=max_seq_len, + seq_offset=seq_offset, + ) + + +def dense_to_jagged(q, dense_tensor, seq_lens): + tensor = torch.zeros_like(q) + offset = 0 + for batch_id, seq_len in enumerate(seq_lens): + tensor[offset:offset + seq_len, :, :] = dense_tensor[batch_id, 0:seq_len, :, :] + offset = offset + seq_len + return tensor + + +def jagged_to_dense(jagged_tensor, seq_lens, head_nums, atten_dim): + need_pad_seq = [] + offset = 0 + for _, seq_len in enumerate(seq_lens): + src_tensor = jagged_tensor[offset:offset + seq_len, :, :].reshape(seq_len, head_nums, atten_dim) + need_pad_seq.append(src_tensor) + offset = offset + seq_len + + dense_tensor = torch.nn.utils.rnn.pad_sequence(need_pad_seq, batch_first=True) + return dense_tensor + + +def gloden_fwd(q, k, v, mask, alpha, seq_offset, attnBias, max_seq_len, enable_mask, enableBias, dataType): + head_nums = q.shape[1] + head_dim = q.shape[2] + batch_size = attnBias.shape[0] + seq_lens = np.zeros((batch_size, )).astype(np.int64) + for batch_id in range(batch_size): + seq_lens[batch_id] = seq_offset[batch_id + 1] - seq_offset[batch_id] + q_dens = jagged_to_dense(q, seq_lens, head_nums, head_dim).to(dataType) + k_dens = jagged_to_dense(k, seq_lens, head_nums, head_dim).to(dataType) + v_dens = jagged_to_dense(v, seq_lens, head_nums, head_dim).to(dataType) + q_dens = q_dens.permute(0, 2, 1, 3) + k_dens = k_dens.permute(0, 2, 3, 1) + v_dens = v_dens.permute(0, 2, 1, 3) + + qk_attn = torch.matmul(q_dens, k_dens) * alpha + qk_attn = qk_attn.to(torch.float32) + attnBias = attnBias.to(torch.float32) + mask = mask.to(torch.float32) + if enableBias: + qk_attn = qk_attn + attnBias + silu = F.silu(qk_attn) * (1 / max_seq_len) + if enable_mask: + silu = silu * mask + silu = silu.to(dataType) + atten_output = torch.matmul(silu, v_dens) + + atten_output = atten_output.permute(0, 2, 1, 3) + atten_output = dense_to_jagged(q, atten_output, seq_lens) + return atten_output.to(dataType) + + +def run_fwd_case(batch_size, max_seq_len, num_heads, attention_dim, data_type): + alpha = 1 + jagged_data = jagged_data_gen(batch_size, max_seq_len, num_heads, attention_dim, data_type) + # golden 输出 + golden_output = gloden_fwd( + jagged_data.q, + jagged_data.k, + jagged_data.v, + jagged_data.mask, + alpha, + jagged_data.seq_offset, + jagged_data.bias, + jagged_data.max_seq_len, + True, + False, + data_type, + ) + # triton 输出 + seq_offsets = torch.tensor(jagged_data.seq_offset, dtype=torch.int64, device=DEVICE) + triton_output = triton_hstu_attention_fwd( + q=jagged_data.q.npu(), + k=jagged_data.k.npu(), + v=jagged_data.v.npu(), + seq_offsets=seq_offsets, + max_seq_len=int(jagged_data.max_seq_len), + alpha=alpha, + causal=True, + mask=jagged_data.mask.npu(), + ) + loss = 1e-4 + if data_type == torch.float16: + loss = 1e-3 + elif data_type == torch.bfloat16: + loss = 1e-2 + torch.testing.assert_close(triton_output.cpu(), golden_output.cpu(), atol=loss, rtol=loss) + + +def golden_bwd(grad, q, k, v, bias, mask, max_seq_len, seq_offset, enable_mask, silu_scale, enable_bias, data_type): + + def jagged_to_dense_bwd(jagged_tensor, seq_lens, max_seq_len, head_num, head_dim): + batch_size = len(seq_lens) + dense_tensor = torch.zeros(batch_size, max_seq_len, head_num, head_dim, dtype=jagged_tensor.dtype) + + offset = 0 + for batch_id, seq_len in enumerate(seq_lens): + dense_tensor[batch_id, :seq_len, :, :] = jagged_tensor[offset:offset + seq_len, :, :] + offset = offset + seq_len + + return dense_tensor + + def dense_to_jagged_bwd(jagged_tensor, dense_tensor, seq_lens): + tensor = torch.zeros_like(jagged_tensor) + + offset = 0 + for batch_id, seq_len in enumerate(seq_lens): + tensor[offset:offset + seq_len, :, :] = dense_tensor[batch_id, 0:seq_len, :, :] + offset = offset + seq_len + + return tensor + + q = q.cpu() + k = k.cpu() + v = v.cpu() + grad = grad.cpu() + head_nums = grad.shape[1] + head_dim = grad.shape[2] + batch_size = bias.shape[0] + seq_lens = np.zeros((batch_size, )).astype(np.int64) + for batch_id in range(batch_size): + seq_lens[batch_id] = seq_offset[batch_id + 1] - seq_offset[batch_id] + grad_dens = jagged_to_dense_bwd(grad, seq_lens, max_seq_len, head_nums, head_dim).to(data_type) + q_dens = jagged_to_dense_bwd(q, seq_lens, max_seq_len, head_nums, head_dim).to(data_type) + k_dens = jagged_to_dense_bwd(k, seq_lens, max_seq_len, head_nums, head_dim).to(data_type) + v_dens = jagged_to_dense_bwd(v, seq_lens, max_seq_len, head_nums, head_dim).to(data_type) + actual_seq_lens = torch.from_numpy(seq_lens).reshape(batch_size, 1, 1, 1).to(data_type) + actual_seq_lens = torch.broadcast_to(actual_seq_lens, bias.shape) + qk = torch.matmul(q_dens.permute(0, 2, 1, 3), k_dens.permute(0, 2, 3, 1)) + gv = torch.matmul(grad_dens.permute(0, 2, 1, 3), v_dens.permute(0, 2, 3, 1)) + qk = qk.float() + gv = gv.float() + bias = bias.float() + if enable_mask: + mask = mask.to(data_type) + mask = mask.float() + if enable_bias: + bias = bias.to(data_type) + bias = bias.float() + qkb = qk + bias + else: + qkb = qk + real_silu_scale = 1 / max_seq_len if silu_scale == 0.0 else silu_scale + + if enable_mask: + score = F.silu(qkb) * real_silu_scale * mask + else: + score = F.silu(qkb) * real_silu_scale + score = score.to(data_type) + v_grad_dens = torch.matmul(score.permute(0, 1, 3, 2), grad_dens.permute(0, 2, 1, 3)).permute(0, 2, 1, 3) + if enable_mask: + bias_grad = gv * real_silu_scale * mask * F.sigmoid(qkb) * (1 + qkb * (1 - F.sigmoid(qkb))) + else: + bias_grad = gv * real_silu_scale * F.sigmoid(qkb) * (1 + qkb * (1 - F.sigmoid(qkb))) + bias_grad = bias_grad.to(data_type) + k_grad_dens = torch.matmul(bias_grad.permute(0, 1, 3, 2), q_dens.permute(0, 2, 1, 3)).permute(0, 2, 1, 3) + q_grad_dens = torch.matmul(bias_grad, k_dens.permute(0, 2, 1, 3)).permute(0, 2, 1, 3) + bias_grad = bias_grad.cpu() + q_grad_dens = q_grad_dens.cpu() + q_grad = dense_to_jagged_bwd(q, q_grad_dens, seq_lens) + k_grad_dens = k_grad_dens.cpu() + k_grad = dense_to_jagged_bwd(k, k_grad_dens, seq_lens) + v_grad_dens = v_grad_dens.cpu() + v_grad = dense_to_jagged_bwd(v, v_grad_dens, seq_lens) + torch.npu.synchronize() + return q_grad, k_grad, v_grad, bias_grad + + +def run_bwd_case(batch_size, max_seq_len, num_heads, attention_dim, data_type): + alpha = 1 + jagged_data = jagged_data_gen(batch_size, max_seq_len, num_heads, attention_dim, data_type) + # golden 输出 + q_grad_golden, k_grad_golden, v_grad_golden, _ = golden_bwd( + jagged_data.grad, + jagged_data.q, + jagged_data.k, + jagged_data.v, + jagged_data.bias, + jagged_data.mask, + jagged_data.max_seq_len, + jagged_data.seq_offset, + True, + 0, + False, + data_type, + ) + + # triton 输出 + seq_offsets = torch.tensor(jagged_data.seq_offset, dtype=torch.int64, device=DEVICE) + dq, dk, dv = triton_hstu_attention_bwd( + grad=jagged_data.grad.npu(), + q=jagged_data.q.npu(), + k=jagged_data.k.npu(), + v=jagged_data.v.npu(), + seq_offsets=seq_offsets, + max_seq_len=int(jagged_data.max_seq_len), + alpha=alpha, + causal=True, + ) + loss = 1e-4 + if data_type == torch.float16: + loss = 1e-3 + elif data_type == torch.bfloat16: + loss = 1e-2 + torch.testing.assert_close(dq.cpu(), q_grad_golden.cpu(), atol=loss, rtol=loss) + torch.testing.assert_close(dk.cpu(), k_grad_golden.cpu(), atol=loss, rtol=loss) + torch.testing.assert_close(dv.cpu(), v_grad_golden.cpu(), atol=loss, rtol=loss) + + +@pytest.mark.parametrize( + "batch_size, max_seq_len, num_heads, attention_dim, data_type", + [ + (2, 1024, 2, 32, torch.float32), + ], +) +def test_hstu_attention_fwd(batch_size, max_seq_len, num_heads, attention_dim, data_type): + np.random.seed(0) + torch.manual_seed(0) + run_fwd_case(batch_size, max_seq_len, num_heads, attention_dim, data_type) + + +@pytest.mark.parametrize( + "batch_size, max_seq_len, num_heads, attention_dim, data_type", + [ + (2, 1024, 2, 32, torch.float32), + ], +) +def test_hstu_attention_bwd(batch_size, max_seq_len, num_heads, attention_dim, data_type): + np.random.seed(0) + torch.manual_seed(0) + run_bwd_case(batch_size, max_seq_len, num_heads, attention_dim, data_type) diff --git a/third_party/ascend/unittest/pytest_ut/test_13_matrix_multiplication_optimized.py b/third_party/ascend/unittest/pytest_ut/test_13_matrix_multiplication_optimized.py new file mode 100644 index 0000000000..27df8cf162 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_13_matrix_multiplication_optimized.py @@ -0,0 +1,216 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl +import triton.runtime.driver as driver +import triton.language.extra.cann.extension as extension + + +# get device properties of npu +def get_npu_properties(): + device = torch.npu.current_device() + return driver.active.utils.get_device_properties(device) + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 256, "BLOCK_TRESHHOLD": 4}), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 256, "BLOCK_TRESHHOLD": 5}), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 256, "BLOCK_TRESHHOLD": 6}), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 256, "BLOCK_TRESHHOLD": 7}), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 256, "BLOCK_TRESHHOLD": 8}), + triton.Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 256, "BLOCK_TRESHHOLD": 4}), + triton.Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 256, "BLOCK_TRESHHOLD": 5}), + triton.Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 256, "BLOCK_TRESHHOLD": 6}), + triton.Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 256, "BLOCK_TRESHHOLD": 7}), + triton.Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 256, "BLOCK_TRESHHOLD": 8}), + ], key=["M", "N", "K"]) +@triton.jit +def matmul_kernel( + mat_a, + mat_b, + mat_c, + M: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, + num_cores: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_TRESHHOLD: tl.constexpr, +): + pid = tl.program_id(axis=0) + task_m_idx = 0 + task_n_idx = 0 + ''' + 水平分核方式每个任务块编号如下 + [0, 1, 2, 3, 4, 5, 6, 7] + [8, 9, 10, 11, 12, 13, 14, 15] + [16, 17, 18, 19, 20, 21, 22, 23] + [24, 25, 26, 27, 28, 29, 30, 31] + [32, 33, 34, 35, 36, 37, 38, 39] + [40, 41, 42, 43, 44, 45, 46, 47] + [48, 49, 50, 51, 52, 53, 54, 55] + [56, 57, 58, 59, 60, 61, 62, 63] + 0核处理 0 20 40 60 4块任务 + 1核处理 1 21 41 61 4块任务 + 2核处理 2 22 42 62 4块任务 + ... + 19核处理 19 39 59 3块任务 + + 大shape下如果使用传统水平分核方式,会有如下问题 + 1:同一时间大量核心需要访问同一块左矩阵内存,产生Bank冲突,导致硬件访问效率降低 + 2:当完成一整行mat_c运算时,已经将所有右矩阵数据全部使用上,右矩阵较大时会超过L2Cache的容量上限, + 从而导致L2Cache的搬入及换出,此后每行运算都会或多或少产生CacheMiss,导致L2Cche命中率较低,影响 + 算子执行效率 + 此处使用8 * 8对角线分核方式可以按8 * 8的方块沿对角线方向分核计算,可以很大程度优化上面两点。 + + 此处以8*8对角线分核为例,实际以BLOCK_TRESHHOLD为tune参数选择最优的阈值 + 8 * 8 对角线分核方式中,每8 * 8分格内任务块编号如下 + [0, 8, 16, 24, 32, 40, 48, 56] + [57, 1, 9, 17, 25, 33, 41, 49] + [50, 58, 2, 10, 18, 26, 34, 42] + [43, 51, 59, 3, 11, 19, 27, 35] + [36, 44, 52, 60, 4, 12, 20, 28] + [29, 37, 45, 53, 61, 5, 13, 21] + [22, 30, 38, 46, 54, 62, 6, 14] + [15, 23, 31, 39, 47, 55, 63, 7] + + M轴方向超过8个基本块时,使用对角线分核可以明显减小Bank冲突 + 当右矩阵大小超过L2Cache大小时,采取对角线分核可以提升L2Cache利用率 + 所以当矩阵在M和N方向均超过8块时使能对角线分核即可有优化,当右矩阵大小超过L2Cache大小时优化效果尤为明显 + ''' + NUM_BLOCKS_M = triton.cdiv(M, BLOCK_M) + NUM_BLOCKS_N = triton.cdiv(N, BLOCK_N) + NUM_BLOCKS = NUM_BLOCKS_M * NUM_BLOCKS_N + # 当任务量较多时，可以使能对角线分核策略进行优化 + if NUM_BLOCKS_M >= BLOCK_TRESHHOLD and NUM_BLOCKS_N >= BLOCK_TRESHHOLD: + for block_idx in range(pid, NUM_BLOCKS, num_cores): + # 8 * 8 对角线分核代码实现 + curThresholdM = BLOCK_TRESHHOLD if block_idx < ( + NUM_BLOCKS_M // BLOCK_TRESHHOLD * BLOCK_TRESHHOLD) * NUM_BLOCKS_N else NUM_BLOCKS_M % BLOCK_TRESHHOLD + curThresholdM_thresholdN = curThresholdM * BLOCK_TRESHHOLD + curThresholdN = BLOCK_TRESHHOLD if block_idx % (NUM_BLOCKS_N * BLOCK_TRESHHOLD) < ( + curThresholdM * + NUM_BLOCKS_N) // curThresholdM_thresholdN * curThresholdM_thresholdN else NUM_BLOCKS_N % BLOCK_TRESHHOLD + localRelativeBlock = block_idx % (BLOCK_TRESHHOLD * NUM_BLOCKS_N) % (BLOCK_TRESHHOLD * curThresholdM) + task_m_idx = localRelativeBlock % curThresholdM + block_idx // (BLOCK_TRESHHOLD * + NUM_BLOCKS_N) * BLOCK_TRESHHOLD + # 求最小公倍数，方便求基本块的坐标 + x, y = curThresholdM, curThresholdN if curThresholdM > curThresholdN else curThresholdN, curThresholdM + while y != 0: + x, y = y, x % y + lcm = curThresholdM * curThresholdN // x + task_n_idx = (localRelativeBlock + (localRelativeBlock // lcm)) % curThresholdN + block_idx % ( + BLOCK_TRESHHOLD * NUM_BLOCKS_N) // curThresholdM_thresholdN * BLOCK_TRESHHOLD + + m_start = task_m_idx * BLOCK_M + n_start = task_n_idx * BLOCK_N + + mat_c_block = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k_start in range(0, K, BLOCK_K): + mat_a_offset = ( + (m_start + tl.arange(0, BLOCK_M)) * K)[:, None] + (k_start + tl.arange(0, BLOCK_K))[None, :] + mat_a_mask = ((m_start + tl.arange(0, BLOCK_M)) < M)[:, None] & ( + (k_start + tl.arange(0, BLOCK_K)) < K)[None, :] + mat_a_block = tl.load(mat_a + mat_a_offset, mask=mat_a_mask, other=0.0) + extension.compile_hint(mat_a_block, "dot_pad_only_k") + mat_b_offset = ( + (k_start + tl.arange(0, BLOCK_K)) * N)[:, None] + (n_start + tl.arange(0, BLOCK_N))[None, :] + mat_b_mask = ((k_start + tl.arange(0, BLOCK_K)) < K)[:, None] & ( + (n_start + tl.arange(0, BLOCK_N)) < N)[None, :] + mat_b_block = tl.load(mat_b + mat_b_offset, mask=mat_b_mask, other=0.0) + extension.compile_hint(mat_b_block, "dot_pad_only_k") + mat_c_block = tl.dot(mat_a_block, mat_b_block, mat_c_block) + mat_c_offset = ((m_start + tl.arange(0, BLOCK_M)) * N)[:, None] + (n_start + tl.arange(0, BLOCK_N))[None, :] + mat_c_mask = ((m_start + tl.arange(0, BLOCK_M)) < M)[:, None] & ( + (n_start + tl.arange(0, BLOCK_N)) < N)[None, :] + tl.store(mat_c + mat_c_offset, mat_c_block.to(tl.bfloat16), mask=mat_c_mask) + else: + # 传统顺序分核 + for block_idx in range(pid, NUM_BLOCKS, num_cores): + task_m_idx = block_idx // NUM_BLOCKS_N + task_n_idx = block_idx % NUM_BLOCKS_N + m_start = task_m_idx * BLOCK_M + n_start = task_n_idx * BLOCK_N + + mat_c_block = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k_start in range(0, K, BLOCK_K): + mat_a_offset = ( + (m_start + tl.arange(0, BLOCK_M)) * K)[:, None] + (k_start + tl.arange(0, BLOCK_K))[None, :] + mat_a_mask = ((m_start + tl.arange(0, BLOCK_M)) < M)[:, None] & ( + (k_start + tl.arange(0, BLOCK_K)) < K)[None, :] + mat_a_block = tl.load(mat_a + mat_a_offset, mask=mat_a_mask, other=0.0) + extension.compile_hint(mat_a_block, "dot_pad_only_k") + mat_b_offset = ( + (k_start + tl.arange(0, BLOCK_K)) * N)[:, None] + (n_start + tl.arange(0, BLOCK_N))[None, :] + mat_b_mask = ((k_start + tl.arange(0, BLOCK_K)) < K)[:, None] & ( + (n_start + tl.arange(0, BLOCK_N)) < N)[None, :] + mat_b_block = tl.load(mat_b + mat_b_offset, mask=mat_b_mask, other=0.0) + extension.compile_hint(mat_b_block, "dot_pad_only_k") + mat_c_block = tl.dot(mat_a_block, mat_b_block, mat_c_block) + mat_c_offset = ((m_start + tl.arange(0, BLOCK_M)) * N)[:, None] + (n_start + tl.arange(0, BLOCK_N))[None, :] + mat_c_mask = ((m_start + tl.arange(0, BLOCK_M)) < M)[:, None] & ( + (n_start + tl.arange(0, BLOCK_N)) < N)[None, :] + tl.store(mat_c + mat_c_offset, mat_c_block.to(tl.bfloat16), mask=mat_c_mask) + + +def triton_matmul( + mat_a, + mat_b, +): + m = mat_a.shape[0] + k = mat_a.shape[1] + n = mat_b.shape[1] + mat_c = torch.empty(m, n, dtype=mat_a.dtype, device=mat_a.device) + ''' + NPU芯片更加亲和512B对齐场景,如下分块通用性能较好,可以使用autotune选取最优 + BLOCK_M = 128 + BLOCK_N = 256 + BLOCK_K = 256 + ''' + + num_cores = get_npu_properties()["num_aicore"] + + matmul_kernel[(num_cores, )](mat_a, mat_b, mat_c, m, n, k, num_cores) + return mat_c + + +# ==================== Pytest Test ==================== +def test_matmul_extension(): + M = 2048 + K = 7168 + N = 16384 + + mat_a = torch.randn([M, K], dtype=torch.bfloat16, device="npu") + mat_b = torch.randn([K, N], dtype=torch.bfloat16, device="npu") + + result = triton_matmul(mat_a, mat_b) + golden = torch.matmul(mat_a, mat_b) + + mask = golden.abs() < 1.0 + tmpatol = tmprtol = 2**-6 + + torch.testing.assert_close(result[mask], golden[mask], atol=tmpatol, rtol=0) + torch.testing.assert_close(result[~mask], golden[~mask], atol=0, rtol=tmprtol) diff --git a/third_party/ascend/unittest/pytest_ut/test_14_accuracy_comparison.py b/third_party/ascend/unittest/pytest_ut/test_14_accuracy_comparison.py new file mode 100644 index 0000000000..bb7720ec4a --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_14_accuracy_comparison.py @@ -0,0 +1,144 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + + +def run_add(x0, x1): + """ + 测试 Triton 实现的向量加法与 PyTorch 的结果，精度比对是否一致。 + + 步骤： + 1. 使用 PyTorch 计算参考结果（torch_ref） + 2. 使用 Triton 编写 kernel 并计算结果（triton_cal） + 3. 调用 accuracy_comparison 进行精度比对 + """ + + # 1. 使用 PyTorch 作为参考实现（golden truth） + def torch_func(x0, x1): + res = x0 + x1 + return res + + # 2. 定义 Triton kernel（在 NPU/GPU 上执行） + @triton.jit + def triton_kernel_add(out_ptr0, # 输出指针：结果存储位置 + in_ptr0, # 输入指针0：x0 的起始地址 + in_ptr1, # 输入指针1：x1 的起始地址 + XS: tl.constexpr # constexpr 参数：向量长度，在编译时确定 + ): + # 生成 [0, 1, 2, ..., XS-1] 的索引数组 + idx = tl.arange(0, XS) + # 从 in_ptr0 + idx 处加载 x0 的值 + tmp0 = tl.load(in_ptr0 + idx) + # 从 in_ptr1 + idx 处加载 x1 的值 + tmp1 = tl.load(in_ptr1 + idx) + # 执行加法 + tmp2 = tmp0 + tmp1 + # 将结果写入 out_ptr0 + idx + tl.store(out_ptr0 + idx, tmp2) + + # 3. Triton 封装函数：调用 kernel 并返回结果 + def triton_func(x0, x1): + y0 = torch.empty_like(x0) # 创建与输入形状、dtype 相同的输出张量 + # 启动 kernel：grid = [1, 1, 1] 表示仅使用一个 block + # 注意：XS 必须作为参数传入，因为它是 tl.constexpr 类型 + triton_kernel_add[1, 1, 1](y0, x0, x1, XS=x0.numel()) + return y0 + + # 4. 获取参考结果和 Triton 计算结果 + torch_ref = torch_func(x0, x1) + triton_cal = triton_func(x0, x1) + + # 5. 精度比对 + accuracy_comparison(triton_cal, torch_ref) + + # 6. 打印成功信息 + print( + f"== dtype:{triton_cal.dtype} == The accuracy comparison between triton_result and torch_result was successful." + ) + + +def accuracy_comparison(y_cal, y_ref): + """ + 精度比对函数：根据数据类型选择合适的比对策略。 + + 不同数据类型的处理策略： + - 浮点类型（float16/32, bfloat16）：使用 torch.testing.assert_close，设置相对/绝对误差容限 + - 整数类型（int8/16/32/64）：要求完全相等（torch.equal） + - 布尔类型（bool）：CPU 上严格比较（避免设备差异） + """ + # 检查输出数据类型是否一致 + assert y_cal.dtype == y_ref.dtype, f"dtype mismatch: {y_cal.dtype} vs {y_ref.dtype}" + tensor_dtype = y_cal.dtype + + # 将张量移动到 NPU（假设测试在 NPU 上进行） + y_cal = y_cal.npu() + y_ref = y_ref.npu() + + # 根据数据类型选择不同的比对方式 + if tensor_dtype == torch.float16: + # float16 精度较低，允许稍大误差 + torch.testing.assert_close(y_ref, y_cal, rtol=1e-3, atol=1e-3, equal_nan=True) + elif tensor_dtype == torch.bfloat16: + # bfloat16 精度更低，建议转为 float32 再比较 + torch.testing.assert_close(y_ref.to(torch.float32), y_cal.to(torch.float32), rtol=1e-3, atol=1e-3, + equal_nan=True) + elif tensor_dtype == torch.float32: + # float32 精度较高，使用更严格的容差 + torch.testing.assert_close(y_ref, y_cal, rtol=1e-4, atol=1e-4, equal_nan=True) + elif tensor_dtype in [torch.int64, torch.int32, torch.int16, torch.int8, torch.uint32]: + # 整数类型应完全相等 + assert torch.equal(y_cal, y_ref), f"Integer tensors are not equal for dtype {tensor_dtype}" + elif tensor_dtype == torch.bool: + # 布尔类型建议在 CPU 上比较，避免设备间布尔表示差异 + assert torch.equal(y_cal.cpu(), y_ref.cpu()), "Boolean tensors are not equal" + else: + raise ValueError(f'Invalid or unsupported tensor dtype: {tensor_dtype}') + + +# ==================== Pytest Test ==================== +@pytest.mark.parametrize("dtype_name, dtype, low, high", [ + ("fp32", torch.float32, 0, 1), + ("fp16", torch.float16, 0, 1), + ("bf16", torch.bfloat16, 0, 1), + ("i64", torch.int64, 1, 100), + ("i32", torch.int32, 1, 100), + ("i16", torch.int16, 1, 100), + ("i8", torch.int8, 1, 100), + ("i1", torch.bool, 0, 2), +]) +def test_all_dtypes(dtype_name, dtype, low, high): + N = 1024 + if dtype == torch.bool: + x0 = torch.randint(low=low, high=high, size=(N, )).bool().npu() + x1 = torch.randint(low=low, high=high, size=(N, )).bool().npu() + elif dtype.is_floating_point: + x0 = torch.rand((N, ), dtype=dtype).npu() + x1 = torch.rand((N, ), dtype=dtype).npu() + else: + x0 = torch.randint(low=low, high=high, size=(N, ), dtype=dtype).npu() + x1 = torch.randint(low=low, high=high, size=(N, ), dtype=dtype).npu() + + print(f"Running test for {dtype_name}...") + run_add(x0, x1) diff --git a/third_party/ascend/unittest/pytest_ut/test_15_demo_autotune.py b/third_party/ascend/unittest/pytest_ut/test_15_demo_autotune.py new file mode 100644 index 0000000000..a63a0b7ec3 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_15_demo_autotune.py @@ -0,0 +1,88 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Autotune +============= +""" +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + + +# Return a set of different kernel configurations for autotune +def get_autotune_config(): + return [ + triton.Config({'XS': 1 * 128, 'multibuffer': True}), + triton.Config({'XS': 12 * 1024, 'multibuffer': True}), + triton.Config({'XS': 12 * 1024, 'multibuffer': False}), + triton.Config({'XS': 8 * 1024, 'multibuffer': True}), + ] + + +# Use @autotune decorator to automatically select the best kernel configuration +@triton.autotune( + configs=get_autotune_config(), + key=["numel"], +) +@triton.jit +def triton_calc_kernel(out_ptr0, in_ptr0, in_ptr1, numel, + XS: tl.constexpr # Block size controlling how many elements each thread block processes + ): + pid = tl.program_id(0) + idx = pid * XS + tl.arange(0, XS) + msk = idx < numel + for i in range(10000): + tmp0 = tl.load(in_ptr0 + idx, mask=msk, other=0.0) + tmp1 = tl.load(in_ptr1 + idx, mask=msk, other=0.0) + tmp2 = tl.math.exp(tmp0) + tmp1 + i + tl.store(out_ptr0 + idx, tmp2, mask=msk) + + +# Function to call the Triton kernel with autotuned configuration +def triton_calc_func(x0, x1): + n = x0.numel() + y0 = torch.empty_like(x0) + + def grid(meta): + return (triton.cdiv(n, meta["XS"]), 1, 1) + + triton_calc_kernel[grid](y0, x0, x1, n) + return y0 + + +# Reference implementation using PyTorch for correctness check +def torch_calc_func(x0, x1): + return torch.exp(x0) + x1 + 10000 - 1 + + +# ==================== Pytest Test ==================== +def test_triton_autotune(): + DEV = "npu" + DTYPE = torch.float32 + N = 192 * 1024 + x0 = torch.randn((N, ), dtype=DTYPE, device=DEV) + x1 = torch.randn((N, ), dtype=DTYPE, device=DEV) + + torch_ref = torch_calc_func(x0, x1) + triton_cal = triton_calc_func(x0, x1) + + torch.testing.assert_close(triton_cal, torch_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_16_profiler.py b/third_party/ascend/unittest/pytest_ut/test_16_profiler.py new file mode 100644 index 0000000000..393969cf63 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_16_profiler.py @@ -0,0 +1,116 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + + +def profiler_wrapper(fn, *args): + result_path = "./result_profiling" + skip_first = 10 + wait = 0 + warmup = 3 + active = 30 + repeat = 1 + stream = torch.npu.current_stream() + experimental_config = torch_npu.profiler._ExperimentalConfig( + aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, l2_cache=False, data_simplification=False) + with torch_npu.profiler.profile( + activities=[torch_npu.profiler.ProfilerActivity.CPU, torch_npu.profiler.ProfilerActivity.NPU], + schedule=torch_npu.profiler.schedule(wait=wait, warmup=warmup, active=active, repeat=repeat, + skip_first=skip_first), + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(result_path), record_shapes=True, + profile_memory=False, with_stack=False, with_flops=False, with_modules=False, + experimental_config=experimental_config) as prof: + stream.synchronize() + for _ in range(skip_first + (wait + warmup + active) * repeat): + fn(*args) + prof.step() + stream.synchronize() + + +@triton.jit +def triton_kernel_add(out_ptr0, in_ptr0, in_ptr1, XS: tl.constexpr): + idx = tl.arange(0, XS) + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tl.load(in_ptr1 + idx) + tmp2 = tmp0 + tmp1 + tl.store(out_ptr0 + idx, tmp2) + + +@triton.jit +def triton_kernel_or(out_ptr0, in_ptr0, in_ptr1, XS: tl.constexpr): + idx = tl.arange(0, XS) + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tl.load(in_ptr1 + idx) + tmp2 = tmp0 | tmp1 + tl.store(out_ptr0 + idx, tmp2) + + +def triton_add_func(x0, x1, N): + y0 = torch.empty_like(x0) + triton_kernel_add[1, 1, 1](y0, x0, x1, N) + return y0 + + +def triton_or_func(x0, x1, N): + y0 = torch.empty_like(x0) + triton_kernel_or[1, 1, 1](y0, x0, x1, N) + return y0 + + +# ==================== Pytest Test ==================== +@pytest.mark.parametrize("dtype, low, high", [ + (torch.float32, 0, 1), + (torch.float16, 0, 1), + (torch.bfloat16, 0, 1), + (torch.int64, 1, 100), + (torch.int32, 1, 100), + (torch.int16, 1, 100), + (torch.int8, 1, 100), + (torch.bool, 0, 2), +]) +def test_elementwise_ops(dtype, low, high): + N = 1024 + test_case_is_inductor = False + + if dtype == torch.bool: + x0 = torch.randint(low=low, high=high, size=(N, )).bool().npu() + x1 = torch.randint(low=low, high=high, size=(N, )).bool().npu() + triton_cal = triton_or_func(x0, x1, N) + ref = x0 | x1 + else: + if dtype.is_floating_point: + x0 = torch.rand((N, ), dtype=dtype).npu() + x1 = torch.rand((N, ), dtype=dtype).npu() + else: + x0 = torch.randint(low=low, high=high, size=(N, ), dtype=dtype).npu() + x1 = torch.randint(low=low, high=high, size=(N, ), dtype=dtype).npu() + + triton_cal = triton_add_func(x0, x1, N) + ref = x0 + x1 + + torch.testing.assert_close(triton_cal, ref) + + def wrapper(): + _ = triton_add_func(x0, x1, N) if dtype != torch.bool else triton_or_func(x0, x1, N) + + profiler_wrapper(wrapper) diff --git a/third_party/ascend/unittest/pytest_ut/test_17_demo_libentry.py b/third_party/ascend/unittest/pytest_ut/test_17_demo_libentry.py new file mode 100644 index 0000000000..a6172552f4 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_17_demo_libentry.py @@ -0,0 +1,131 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import time + +import pytest +import torch +import torch_npu + +import triton +import triton.language as tl +from triton.runtime.libentry import libentry + +DEV = "npu" +DTYPE = torch.float32 +SEQ_LEN = 2 * 1024 +device = torch.npu.current_device() +stream = torch.npu.current_stream(device) + + +def benchmark(func): + warmup = 10 + repeat = 100 + + def wrapper(*args, **kwargs): + # + for _ in range(warmup): + result = func(*args, **kwargs) + stream.synchronize() + # + start_time = time.perf_counter_ns() + for _ in range(repeat): + result = func(*args, **kwargs) + stream.synchronize() + end_time = time.perf_counter_ns() + # + start_time = start_time * 1e-3 + end_time = end_time * 1e-3 + elapsed_time = (end_time - start_time) / repeat + return (result, elapsed_time) + + return wrapper + + +@libentry() +@triton.jit +def softmax_kernel( + output_ptr, + input_ptr, + input_row_stride, + output_row_stride, + n_rows, + n_cols, + XBLOCK: tl.constexpr, + XBLOCK_SUB: tl.constexpr, + RBLOCK: tl.constexpr, +): + pid = tl.program_id(0) + row_start = pid * XBLOCK + rblk_idx = tl.arange(0, XBLOCK_SUB) + col_idx = tl.arange(0, RBLOCK) + for row_idx in tl.range(0, XBLOCK, XBLOCK_SUB): + row_offsets = row_start + row_idx + rblk_idx[:, None] + col_offsets = col_idx[None, :] + xmask = row_offsets < n_rows + ymask = col_offsets < n_cols + mask = xmask & ymask + input_idx = row_offsets * input_row_stride + col_offsets + input_ptrs = input_ptr + input_idx + row = tl.load(input_ptrs, mask=mask, other=-float('inf')) + row_minus_max = row - tl.max(row, axis=1).reshape(XBLOCK_SUB, 1) + numerator = tl.exp(row_minus_max) + denominator = tl.sum(numerator, axis=1).reshape(XBLOCK_SUB, 1) + softmax_output = numerator / denominator + output_ptrs = output_ptr + (row_offsets * output_row_stride + col_offsets) + tl.store(output_ptrs, softmax_output, mask=mask) + + +@benchmark +def torch_func(x0: torch.Tensor): + m = torch.nn.Softmax(dim=1) + return m(x0) + + +@benchmark +def triton_func(y0: torch.Tensor, x0: torch.Tensor): + n_rows, n_cols = x0.shape + ncore = 40 + xs = (n_rows + ncore - 1) // ncore + xss = min(xs, 5) + softmax_kernel[(ncore, 1, 1)]( + y0, + x0, + x0.stride(0), + y0.stride(0), + n_rows, + n_cols, + XBLOCK=xs, + XBLOCK_SUB=xss, + RBLOCK=n_cols, + ) + return y0 + + +@pytest.mark.parametrize("batch", [1000 * x for x in range(1, 16 + 1)]) +def test_demo_libentry_softmax(batch): + torch.manual_seed(0) + x = torch.rand((batch, SEQ_LEN), dtype=DTYPE, device=DEV) + y = torch.empty_like(x) + + torch_out, _ = torch_func(x) + triton_out, _ = triton_func(y, x) + + torch.testing.assert_close(triton_out, torch_out) diff --git a/third_party/ascend/unittest/pytest_ut/test_18_gather.py b/third_party/ascend/unittest/pytest_ut/test_18_gather.py new file mode 100644 index 0000000000..f8a4c70898 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_18_gather.py @@ -0,0 +1,139 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Gather +=============== +This is an example only for npu. +""" + +import pytest +import torch +import torch_npu +import triton +import triton.runtime.driver as driver +import triton.language as tl + + +# get device properties of npu +def get_npu_properties(): + device = torch.npu.current_device() + return driver.active.utils.get_device_properties(device) + + +# a torch-version gather benchmark +def torch_gather(embeddings, idxes, default_value=0.0): + # make the result tensor + res = torch.empty((idxes.shape[0], embeddings.shape[-1]), dtype=embeddings.dtype, device=embeddings.device) + + # scatter embeddings + res[idxes >= 0] = embeddings[idxes[idxes >= 0]] + # set default values + res[idxes < 0] = default_value + + return res + + +# triton-version gather's kernel +@triton.jit +def gather_kernel(embeddings_ptr, idxes_ptr, res_ptr, rows, cols, DEFAULT_VALUE: tl.constexpr, + BIG_CORE_NUM: tl.constexpr, BIG_ROW_BLOCK_SIZE: tl.constexpr, COL_BLOCK_SIZE: tl.constexpr, + COL_BLOCK_SIZE_SUB: tl.constexpr): + SMALL_ROW_BLOCK_SIZE = BIG_ROW_BLOCK_SIZE - 1 + + embedding_dtype = embeddings_ptr.type.element_ty + default_value = tl.cast(DEFAULT_VALUE, dtype=embedding_dtype) + default_embedding = tl.full((COL_BLOCK_SIZE_SUB, ), default_value, dtype=embedding_dtype) + + core_idx = tl.program_id(0) + # compute the the size and start index of block + row_block_size = BIG_ROW_BLOCK_SIZE if (core_idx < BIG_CORE_NUM) else SMALL_ROW_BLOCK_SIZE + row_start_idx = (core_idx * BIG_ROW_BLOCK_SIZE) if (core_idx < BIG_CORE_NUM) else ( + BIG_CORE_NUM * BIG_ROW_BLOCK_SIZE + (core_idx - BIG_CORE_NUM) * SMALL_ROW_BLOCK_SIZE) + + # process blocks witn shape (row_block_size, COL_BLOCK_SIZE_SUB) one by one + for col_idx in tl.range(0, COL_BLOCK_SIZE, COL_BLOCK_SIZE_SUB): + emb_col_offsets = col_idx + tl.arange(0, COL_BLOCK_SIZE_SUB) + emb_col_mask = emb_col_offsets < cols + + for row_idx in tl.range(row_start_idx, min(row_start_idx + row_block_size, rows)): + idx_val = tl.load(idxes_ptr + row_idx) + + write_row_offset = row_idx * cols + write_emb_mask = emb_col_mask + + if idx_val >= 0: + read_row_offset = idx_val * cols + read_emb_mask = emb_col_mask + # read embedding + embedding = tl.load(embeddings_ptr + read_row_offset + emb_col_offsets, mask=read_emb_mask) + tl.store(res_ptr + write_row_offset + emb_col_offsets, embedding, write_emb_mask) + else: + # set default values + tl.store(res_ptr + write_row_offset + emb_col_offsets, default_embedding, write_emb_mask) + + +# triton-version gather's host +def triton_gather(embeddings: torch.Tensor, indices: torch.Tensor, default_value=0.0): + # constant settings for npu + USE_SIZE = 96 * 1024 + CORE_NUM = get_npu_properties()["num_vectorcore"] + + n_rows = indices.shape[0] + n_cols = embeddings.shape[1] + # make the result tensor + output = torch.empty(n_rows, n_cols, dtype=embeddings.dtype, device=embeddings.device) + + # when writing an npu kernel using triton, + # you should note that the difference between BLOCK_SIZE and BLOCK_SIZE_SUB + # BLOCK_SIZE specifies the size of data that are processed in one program + col_size_aligned = triton.cdiv(embeddings.shape[-1] * embeddings.element_size(), + 32) * 32 // embeddings.element_size() + # the data are scattered to multiple programs, which can not be even + # some process more data, some process less + big_row_block_size = triton.cdiv(n_rows, CORE_NUM) + big_core_num = CORE_NUM - ((big_row_block_size * CORE_NUM) - n_rows) + col_block_size = col_size_aligned + + # BLOCK_SIZE_SUB specifies the size of data that are processed in one loop of a program + max_col_block_size_sub = USE_SIZE // embeddings.element_size() // 2 + col_block_size_sub = min(col_size_aligned, max_col_block_size_sub) + + grid = (min(n_rows, CORE_NUM), triton.cdiv(n_cols, col_block_size)) + # launch the kernel + gather_kernel[grid](embeddings, indices, output, n_rows, n_cols, default_value, BIG_CORE_NUM=big_core_num, + BIG_ROW_BLOCK_SIZE=big_row_block_size, COL_BLOCK_SIZE=col_block_size, + COL_BLOCK_SIZE_SUB=col_block_size_sub) + + return output + + +# ==================== Pytest Test ==================== +@pytest.mark.parametrize("n_rows", [500, 1000]) +@pytest.mark.parametrize("n_cols", [16, 17, 31, 32, 63, 64, 128, 256, 819, 512, 1024, 8192, 1001, 2003, 17000]) +@pytest.mark.parametrize("index_num", [19, 123, 4321, 54321, 100, 200, 819, 500, 700, 1000]) +def test_gather(n_rows, n_cols, index_num): + indices = torch.randint(0, n_rows, (index_num, ), dtype=torch.int32).npu() + embeddings = torch.randn(n_rows, n_cols, dtype=torch.float).npu() + + expect = torch_gather(embeddings, indices).cpu() + actual = triton_gather(embeddings, indices).cpu() + torch.npu.synchronize() + + torch.testing.assert_close(actual, expect) diff --git a/third_party/ascend/unittest/pytest_ut/test_2d_permute.py b/third_party/ascend/unittest/pytest_ut/test_2d_permute.py new file mode 100644 index 0000000000..24fcee34d2 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_2d_permute.py @@ -0,0 +1,54 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl + +import torch +import torch_npu + + +def fn(x): + return x.t() + + +@triton.jit +def triton_2d_permute(output_ptr, input_ptr, X: tl.constexpr, Y: tl.constexpr): + xindex = tl.arange(0, X * Y) + input_local = tl.load(input_ptr + xindex) + output_local = input_local.reshape(X, Y).trans().reshape(X * Y) + tl.store(output_ptr + xindex, output_local) + + +@pytest.mark.parametrize('X', [32, 64, 256]) +@pytest.mark.parametrize('Y', [16, 32]) +def test_cases(X, Y): + + x = torch.randn((X, Y)).npu() + output1 = fn(x) + output2 = torch.randn(output1.shape, dtype=output1.dtype).npu() + + triton_2d_permute[1, 1, 1](output2, x, X, Y, debug=True) + print(output1) + print(output2) + + torch.testing.assert_close(output1, output2, rtol=1e-3, atol=1e-3) diff --git a/third_party/ascend/unittest/pytest_ut/test_3Dgrid.py b/third_party/ascend/unittest/pytest_ut/test_3Dgrid.py new file mode 100644 index 0000000000..b779cee73b --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_3Dgrid.py @@ -0,0 +1,91 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import pytest + +BLOCK: tl.constexpr = 32 + + +@triton.jit +def triton_(in_ptr0, out_ptr0, x0_numel, r1_numel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr, + block_id_threshold: tl.constexpr, XBLOCK1: tl.constexpr, num_core: tl.constexpr): + RBLOCK: tl.constexpr = 64 + + block_idx = tl.program_id(0) * tl.num_programs(1) * tl.num_programs(2) + tl.program_id(1) * tl.num_programs( + 2) + tl.program_id(2) + if (block_idx < block_id_threshold): + offset = block_idx * XBLOCK + loops1 = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB # 32+23 / 24 = 2 + upper = offset + XBLOCK + else: + offset = block_id_threshold * XBLOCK + (block_idx - + block_id_threshold) * XBLOCK1 #pid=34 offset = 9*32 + (34-9)*24 = 888 + loops1 = (XBLOCK1 + XBLOCK_SUB - 1) // XBLOCK_SUB #1 + if (block_idx == num_core - 1): + upper = x0_numel + else: + upper = offset + XBLOCK1 # 912 + + base1 = tl.arange(0, XBLOCK_SUB) + base2 = tl.arange(0, RBLOCK) + loops2: tl.constexpr = (r1_numel + RBLOCK - 1) // RBLOCK + for loop1 in range(loops1): + x = offset + (loop1 * XBLOCK_SUB) + base1 + x0_prime = offset + (loop1 * XBLOCK_SUB) + base1[None, :] + x0 = offset + (loop1 * XBLOCK_SUB) + base1[:, None] + xmask = x0 < upper + r1_prime = base2[:, None] + rindex = base2 + r1 = base2[None, :] + rmask = r1 < r1_numel + tmp0 = tl.load(in_ptr0 + (r1 + (64 * x0)), rmask & xmask, other=0.0) + + tmp1 = tl.reshape(tmp0, [XBLOCK_SUB, RBLOCK]) + tmp2_tmp = tl.sum(tmp1, 1) + tmp2 = tmp2_tmp.reshape(XBLOCK_SUB, 1) + + tl.store(out_ptr0 + (x0), tmp2, xmask) + + +guards = {"dummy": None} + + +# @pytest.mark.skip(reason="multi-process error, to be fixed.") +@pytest.mark.parametrize("size", [(1025, 64)]) +def test_3dgrid(size): + b = torch.randn((size), dtype=torch.float32).npu() + c = torch.sum(b, dim=1) + + ret = torch.randn((size[0]), dtype=torch.float32).npu() + + triton_[5, 2, 4](b, ret, size[0], size[1], XBLOCK=32, XBLOCK_SUB=16, block_id_threshold=9, XBLOCK1=24, num_core=40, + debug=True) + print(c[0:8]) + print(ret[0:8]) + torch.testing.assert_close(c, ret) + print("test 3D launch passed") + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/third_party/ascend/unittest/pytest_ut/test_abs.py b/third_party/ascend/unittest/pytest_ut/test_abs.py new file mode 100644 index 0000000000..bb5964bc9c --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_abs.py @@ -0,0 +1,69 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import numpy as np +import torch +import pytest +import test_common + + +def torch_pointwise(x0): + if x0.dtype in [torch.uint8, torch.uint16, torch.uint32, torch.uint64]: + return x0 + res = torch.abs(x0) + return res + + +@triton.jit +def triton_abs(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop1 in range(loops1): + x0_prime = offset + (loop1 * XBLOCK_SUB) + base1 + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp2 = tl.abs(tmp0) + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float16', (2, 4096, 8), 32, 2048, 64], + # ['bfloat16', (2, 4096, 8), 32, 2048, 64], + ['float32', (2, 4096, 8), 32, 2048, 64], + ['int8', (2, 4096, 8), 32, 2048, 64], + # ['int16', (2, 4096, 8), 32, 2048, 64], + # ['int32', (2, 4096, 8), 32, 2048, 64], + # ['int64', (2, 4096, 8), 32, 2048, 64], + ['uint8', (2, 4096, 8), 32, 2048, 64], + # ['uint16', (2, 4096, 8), 32, 2048, 64], + # ['uint32', (2, 4096, 8), 32, 2048, 64], + # ['uint64', (2, 4096, 8), 32, 2048, 64], +]) +def test_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + np_x0 = test_common.generate_numpy(shape, dtype) + x0 = torch.from_numpy(np_x0).to(eval('torch.' + dtype)).npu() + y_ref = torch_pointwise(x0) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_abs[ncore, 1, 1](x0, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_abs_2.py b/third_party/ascend/unittest/pytest_ut/test_abs_2.py new file mode 100644 index 0000000000..4f4df6c757 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_abs_2.py @@ -0,0 +1,62 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time + +import torch +import torch_npu +import test_common + + +def torch_abs(x0): + res = torch.abs(x0) + return res + + +@triton.jit +def triton_abs(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.abs(tmp0) + tl.store(out_ptr0 + (x0), tmp1, None) + + +@pytest.mark.parametrize('param_list', [ + ['float16', (4, 4), 4, 4, 4], + ['float32', (4, 4), 4, 4, 4], +]) +def test_abs(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype) + y_ref = torch_abs(x0) + tyname = test_common.get_triton_sig_typename(dtype) + + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + x0 = x0.npu() + triton_abs[ncore, 1, 1](x0, y_cal, xblock, xblock_sub, debug=True) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_acos.py b/third_party/ascend/unittest/pytest_ut/test_acos.py new file mode 100644 index 0000000000..ed0c5bedfd --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_acos.py @@ -0,0 +1,87 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + +import torch +import torch_npu + + +@triton.jit +def triton_acos(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = libdevice.acos(tmp0) + tl.store(out_ptr0 + (x0), tmp1, None) + + +@pytest.mark.parametrize('param_list', ['float32', 'float16', 'bfloat16']) +def test_asinh_special(param_list): + dtype = param_list + x0 = torch.linspace(-1.0 + 1e-6, 1.0 - 1e-6, 256, dtype=eval("torch." + dtype)).npu() + + y_ref = torch.acos(x0) + y_cal = torch.zeros_like(y_ref) + triton_acos[1, 1, 1](x0, y_cal, x0.shape[0], x0.shape[0]) + bf16_tolerance = 1.0 / 128 + if dtype == 'bfloat16': + torch.testing.assert_close(y_ref, y_cal, rtol=bf16_tolerance, atol=bf16_tolerance) + else: + torch.testing.assert_close(y_ref, y_cal, rtol=1e-3, atol=1e-3) + + +@triton.jit +def acos_kernel(in_ptr0, out_ptr0, n_elements, BLOCK_SIZE: tl.constexpr): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(in_ptr0 + offsets, mask=mask) + x = x.to(tl.float32) + y = libdevice.acos(x) + tl.store(out_ptr0 + offsets, y, mask=mask) + + +@pytest.mark.acos +@pytest.mark.parametrize('shape', [(), (1, ), (1024, 1024), (20, 320, 15), (16, 128, 64, 60), (16, 7, 57, 32, 29)]) +@pytest.mark.parametrize('dtype', [torch.float16, torch.float32, torch.bfloat16]) +def test_accuracy_acos(shape, dtype): + inp = torch.randn(shape, dtype=dtype, device='npu') + ref_out = torch.acos(inp) + res_out = torch.zeros_like(ref_out) + n_elements = inp.numel() + grid = (triton.cdiv(n_elements, 1024), ) + acos_kernel[grid](inp, res_out, n_elements, BLOCK_SIZE=1024) + print("ref_out:", ref_out) + print("res_out:", res_out) + if dtype == torch.bfloat16: + bf16_tolerance = 1.0 / 128 + torch.testing.assert_close(res_out, ref_out, rtol=bf16_tolerance, atol=bf16_tolerance, equal_nan=True) + elif dtype == torch.float16: + torch.testing.assert_close(res_out, ref_out, rtol=1e-3, atol=1e-4, equal_nan=True) + else: + torch.testing.assert_close(res_out, ref_out, rtol=1e-6, atol=1e-4, equal_nan=True) diff --git a/third_party/ascend/unittest/pytest_ut/test_acosh.py b/third_party/ascend/unittest/pytest_ut/test_acosh.py new file mode 100644 index 0000000000..78705b33b0 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_acosh.py @@ -0,0 +1,73 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + +import torch +import torch_npu + + +def torch_acosh(x0): + return torch.acosh(x0) + + +@triton.jit +def triton_acosh(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = libdevice.acosh(tmp0) + tl.store(out_ptr0 + (x0), tmp1, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['bfloat16', (2, 4096, 8), 2, 32768, 1024], +]) +def test_acosh_common(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu().abs() + 1.0 + y_ref = torch_acosh(x0) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_acosh[ncore, 1, 1](x0, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['bfloat16', (2, 4096, 8), 2, 32768, 1024], +]) +def test_acosh_special(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = -test_common.generate_tensor(shape, dtype).npu().abs() + 1.0 + y_ref = torch_acosh(x0) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_acosh[ncore, 1, 1](x0, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_add.py b/third_party/ascend/unittest/pytest_ut/test_add.py new file mode 100644 index 0000000000..be88ce1a5d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_add.py @@ -0,0 +1,91 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import numpy as np +import torch +import pytest +import test_common + + +def torch_pointwise(x0, x1): + res = x0 + x1 + return res + + +@triton.jit +def triton_add(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop1 in range(loops1): + x0_prime = offset + (loop1 * XBLOCK_SUB) + base1 + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.load(in_ptr1 + (x0), None) + tmp2 = tmp0 + tmp1 + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], + ['uint8', (2, 4096, 8), 2, 32768, 1024], +]) +def test_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_pointwise(x0, x1) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_add[ncore, 1, 1](x0, x1, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_all_blocks_parallel(param_list, monkeypatch): + monkeypatch.setenv("TRITON_ALL_BLOCKS_PARALLEL", "1") + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_pointwise(x0, x1) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_add[ncore, 1, 1](x0, x1, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) + monkeypatch.delenv("TRITON_ALL_BLOCKS_PARALLEL") + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_auto_blockify(param_list, monkeypatch): + monkeypatch.setenv("TRITON_ALL_BLOCKS_PARALLEL", "1") + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_pointwise(x0, x1) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_add[ncore, 1, 1](x0, x1, y_cal, xblock, xblock_sub, auto_blockify_size=ncore) + test_common.validate_cmp(dtype, y_cal, y_ref) + monkeypatch.delenv("TRITON_ALL_BLOCKS_PARALLEL") diff --git a/third_party/ascend/unittest/pytest_ut/test_address_check.py b/third_party/ascend/unittest/pytest_ut/test_address_check.py new file mode 100644 index 0000000000..3a8429c08f --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_address_check.py @@ -0,0 +1,68 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import pytest + + +@triton.jit +def simple_kernel(x_ptr, y_ptr, output_ptr, n_elements): + pid = tl.program_id(axis=0) + offsets = pid * 1024 + tl.arange(0, 1024) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + ret = x + y + tl.store(output_ptr + offsets, ret, mask=mask) + + +def test_npu_tensor_should_success(): + print("Test the NPU tensor. The NPU tensor should be passed and executed properly.") + + size = 1024 + x_npu = torch.rand(size, device='npu') + y_npu = torch.rand(size, device='npu') + output = torch.empty(size, device='npu') + + simple_kernel[(1, )](x_npu, y_npu, output, size) + + expected = x_npu + y_npu + actual = output + + torch.testing.assert_close(expected, actual, rtol=1e-03, atol=1e-03) + + +def test_cpu_tensor_should_fail(): + print("Test the CPU tensor. An address check exception should be raised.") + + size = 1024 + x_cpu = torch.rand(size, device='cpu') + y_cpu = torch.rand(size, device='cpu') + output = torch.empty(size, device='npu') + + with pytest.raises(ValueError) as exc_info: + simple_kernel[(1, )](x_cpu, y_cpu, output, size) + + error_msg = str(exc_info.value) + assert "cannot be accessed from Triton (cpu tensor?)" in error_msg, \ + f"Expected error message to contain CPU tensor rejection hint, but got: {error_msg}" diff --git a/third_party/ascend/unittest/pytest_ut/test_advance.py b/third_party/ascend/unittest/pytest_ut/test_advance.py new file mode 100644 index 0000000000..d76f0f3278 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_advance.py @@ -0,0 +1,182 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, output_ptr1, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + # idx = tl.arange(0,XB*YB*ZB) + block_ptr_in = tl.make_block_ptr( + base=x_ptr, + shape=(XB, YB, ZB), + strides=(YB * ZB, ZB, 1), + offsets=(9, 6, 5), + block_shape=(XB, YB, ZB), + order=(2, 1, 0), + ) + bbptr = tl.advance(block_ptr_in, (-9, -6, -5)) + # XB,YB,1 + X = tl.load(bbptr) + # X = tl.load(x_ptr + idx) + # Y = tl.load(y_ptr + idx) + + # xx=tl.view(X,(ZB*YB,XB)) + + oidx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + block_ptr_out = tl.make_block_ptr( + base=output_ptr, + shape=(XB, YB, ZB), + strides=(YB * ZB, ZB, 1), + offsets=(0, 0, 0), + block_shape=(XB, YB, ZB), + order=(2, 1, 0), + ) + tl.store(block_ptr_out, X) + # tl.store(output_ptr + tl.arange(0,ZB*YB)[:,None]*XB+xidx[None,:], xx) + # tl.store(output_ptr + xidx[:,None]*YB+yidx[None,:], yy) + + +@triton.jit +def fn_npu_2d(output_ptr, x_ptr, y_ptr, z_ptr, output_ptr1, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xoffset = tl.program_id(0) + block_ptr_in = tl.make_block_ptr( + base=x_ptr, + shape=(XB, YB), + strides=(YB, 1), + offsets=(6 + xoffset, 5), + block_shape=(XB, YB), + order=(1, 0), + ) + bbptr = tl.advance(block_ptr_in, (-6, -5)) + # XB,YB,1 + X = tl.load(bbptr) + + block_ptr_out = tl.make_block_ptr( + base=output_ptr, + shape=(XB, YB), + strides=(YB, 1), + offsets=(xoffset, 0), + block_shape=(XB, YB), + order=(1, 0), + ) + tl.store(block_ptr_out, X) + + +@triton.jit +def fn_npu_3d(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + block_ptr_in = tl.make_block_ptr( + base=x_ptr, + shape=(XB, YB, ZB), + strides=(YB * ZB, ZB, 1), + offsets=(0, 0, 0), + block_shape=(XB, YB, 2), + order=(2, 1, 0), + ) + + block_ptr_out = tl.make_block_ptr( + base=output_ptr, + shape=(XB, YB, ZB), + strides=(YB * ZB, ZB, 1), + offsets=(0, 0, 0), + block_shape=(XB, YB, 2), + order=(2, 1, 0), + ) + + for _ in range(ZB // 2): + X = tl.load(block_ptr_in, boundary_check=(0, 1, 2)) + tl.store(block_ptr_out, X, boundary_check=(0, 1, 2)) + block_ptr_in = tl.advance(block_ptr_in, (0, 0, 2)) + block_ptr_out = tl.advance(block_ptr_out, (0, 0, 2)) + + +@pytest.mark.parametrize('dtype', ["int32", "float32", "int16"]) +@pytest.mark.parametrize('shape', [(33, 9, 6), (8, 8, 4)]) +def test_advance_with_boundary_check(dtype, shape): + x = torch.randint(low=-128, high=128, size=shape, dtype=eval('torch.' + dtype)).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + a = x + + fn_npu_3d[1, 1, 1](output, x, XB=shape[0], YB=shape[1], ZB=shape[2]) + + torch.testing.assert_close(output, a) + + +@pytest.mark.parametrize('dtype', ["int32", "float32", "int16"]) +@pytest.mark.parametrize('shape', [(1, 3), (3, 1), (1, 13), (13, 1)]) +def test_advance_supplement(dtype, shape): + x = torch.randint(low=-128, high=128, size=shape, dtype=eval('torch.' + dtype)).npu() + y = torch.randint(low=-128, high=128, size=shape, dtype=eval('torch.' + dtype)).npu() + z = torch.randint(low=-128, high=128, size=shape, dtype=eval('torch.' + dtype)).npu() + + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + + a = x + + fn_npu_2d[1, 1, 1](output, x, y, z, output1, XB=shape[0], YB=shape[1], ZB=1) + + torch.testing.assert_close(output, a) + + +paras = [ + ('*fp32', eval('torch.float32'), 2, 256, 16), + ('*fp32', eval('torch.float32'), 8, 8, 4), + ('*fp16', eval('torch.float16'), 2, 256, 16), + ('*fp16', eval('torch.float16'), 8, 8, 4), + ('*i8', eval('torch.int8'), 2, 256, 16), + ('*i8', eval('torch.int8'), 8, 8, 4), +] + + +@pytest.mark.parametrize('para_type,data_type,XB,YB,ZB', paras) +def test_npu(para_type, data_type, XB, YB, ZB): + + x = torch.randint(low=-128, high=128, size=(XB, YB, ZB), dtype=data_type).npu() + y = torch.randint(low=-128, high=128, size=(XB, YB, ZB), dtype=data_type).npu() + z = torch.randint(low=-128, high=128, size=(XB, YB, ZB), dtype=data_type).npu() + + print(f"shape = {x.shape}") + print(x.dtype) + + output = torch.randint(1, (XB, YB, ZB), dtype=data_type).npu() + output1 = output + print(f"output.dtype={output.dtype}") + + a = x + print(a) + fn_npu_[1, 1, 1](output, x, y, z, output1, XB=XB, YB=YB, ZB=ZB, debug=True) + print(output) + torch.testing.assert_close(output, a) + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/third_party/ascend/unittest/pytest_ut/test_advance_ptr.py b/third_party/ascend/unittest/pytest_ut/test_advance_ptr.py new file mode 100644 index 0000000000..f9052cd28e --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_advance_ptr.py @@ -0,0 +1,51 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest + + +@triton.jit +def fn_npu_3d(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + block_ptr_in = tl.make_block_ptr(base=x_ptr, shape=(XB, YB, ZB), strides=(YB * ZB, ZB, 1), offsets=(0, 0, 0), + block_shape=(XB, YB, 2), order=(2, 1, 0)) + block_ptr_out = tl.make_block_ptr(base=output_ptr, shape=(XB, YB, ZB), strides=(YB * ZB, ZB, 1), offsets=(0, 0, 0), + block_shape=(XB, YB, 2), order=(2, 1, 0)) + pid = tl.program_id(axis=0) # pid=0,1 BLOCK_SIZE_N=8 + for _ in range(ZB // 2): + X = tl.load(block_ptr_in, boundary_check=(0, 1, 2)) + tl.store(block_ptr_out, X, boundary_check=(0, 1, 2)) + block_ptr_in = tl.advance(block_ptr_in, (0, 0, 2)) + block_ptr_out = tl.advance(block_ptr_out, (0, 0, 2)) + + +@pytest.mark.parametrize('dtype', ["int32", "float32", "int16"]) +@pytest.mark.parametrize('shape', [(33, 9, 6), (8, 8, 4)]) +def test_advance_with_boundary_check(dtype, shape): + x = torch.randint(low=-128, high=128, size=shape, dtype=eval('torch.' + dtype)).npu() + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + expected = x + output = torch.randint(1, shape, dtype=eval('torch.' + dtype)).npu() + fn_npu_3d[1, 1, 1](output, x, XB=shape[0], YB=shape[1], ZB=shape[2]) + torch.testing.assert_close(output, expected) diff --git a/third_party/ascend/unittest/pytest_ut/test_affine_map_binding.py b/third_party/ascend/unittest/pytest_ut/test_affine_map_binding.py new file mode 100644 index 0000000000..b46f389eb9 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_affine_map_binding.py @@ -0,0 +1,114 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir +import triton.language.extra.cann.extension as al +import pytest + + +def test_extension_reexports_affine_bindings(): + assert al.affine_map is ascend_ir.affine_map + assert al.affine_expr is ascend_ir.affine_expr + assert al.affine_constant_expr is ascend_ir.affine_constant_expr + assert al.affine_dim_expr is ascend_ir.affine_dim_expr + assert al.affine_symbol_expr is ascend_ir.affine_symbol_expr + assert al.affine_binary_op_expr is ascend_ir.affine_binary_op_expr + assert al.AffineMap is al.affine_map + assert al.AffineExpr is al.affine_expr + assert al.AffineConstantExpr is al.affine_constant_expr + assert al.AffineDimExpr is al.affine_dim_expr + assert al.AffineSymbolExpr is al.affine_symbol_expr + assert al.AffineBinaryOpExpr is al.affine_binary_op_expr + + +def test_make_affine_map(): + with ir.context() as ctx: + ir.load_dialects(ctx) + ascend_ir.load_dialects(ctx) + + d0 = ascend_ir.affine_expr.get_dim(0) + d1 = ascend_ir.affine_expr.get_dim(1) + c2 = ascend_ir.affine_expr.get_constant(2) + + expr = (d0 + c2) * d1 + assert "d0" in str(expr) and "d1" in str(expr) + assert not expr.is_pure_affine() + assert hash(expr) == hash(expr) + assert d0 == ascend_ir.affine_expr.get_dim(0) + assert c2 == ascend_ir.affine_expr.get_constant(2) + assert isinstance(c2, ascend_ir.affine_expr) + assert isinstance(d0, ascend_ir.affine_expr) + + identity_map = ascend_ir.affine_map.get_identity(2) + transpose_map = ascend_ir.affine_map.get(2, 0, [1, 0]) + transpose_map_by_expr = ascend_ir.affine_map.get(2, 0, [d1, d0]) + sum_map = ascend_ir.affine_map.get(2, 0, [d0 + d1, d1]) + const_map = ascend_ir.affine_map.get_constant(7) + minor_identity_map = ascend_ir.affine_map.get_minor_identity(3, 2) + + assert identity_map.is_identity() + assert identity_map.is_permutation() + assert identity_map.get_num_dims() == 2 + assert identity_map.get_num_symbols() == 0 + assert identity_map.get_num_results() == 2 + assert str(identity_map) == "(d0, d1) -> (d0, d1)" + + assert not transpose_map.is_identity() + assert transpose_map.is_permutation() + assert str(transpose_map) == "(d0, d1) -> (d1, d0)" + assert str(transpose_map_by_expr) == "(d0, d1) -> (d1, d0)" + assert str(sum_map) == "(d0, d1) -> (d0 + d1, d1)" + assert transpose_map.to_dict() == { + "num_dims": 2, + "num_symbols": 0, + "results": [1, 0], + } + assert str(sum_map.get_sub_map([1])) == "(d0, d1) -> (d1)" + assert str(sum_map.compose(transpose_map)) == "(d0, d1) -> (d0 + d1, d0)" + assert str(transpose_map.inverse_permutation()) == "(d0, d1) -> (d1, d0)" + assert transpose_map == transpose_map_by_expr + assert hash(transpose_map) == hash(transpose_map) + assert [str(r) for r in sum_map.get_results()] == ["d0 + d1", "d1"] + assert const_map.is_single_constant() + assert const_map.get_constant_result() == 7 + assert str(minor_identity_map) == "(d0, d1, d2) -> (d1, d2)" + + +def test_build_buffer_type_with_affine_map(): + with ir.context() as ctx: + ir.load_dialects(ctx) + ascend_ir.load_dialects(ctx) + builder = ascend_ir.ascendnpu_ir_builder(ctx) + + transpose_map = ascend_ir.affine_map.get(2, 0, [1, 0]) + ub_attr = builder.get_target_attribute(ascend_ir.AddressSpace.UB) + + buffer_ty = builder.get_buffer_ty_with_affine_map([8, 16], builder.get_float_ty(), transpose_map, ub_attr) + + assert "memref<8x16xf32" in str(buffer_ty) + assert "affine_map<(d0, d1) -> (d1, d0)>" in str(buffer_ty) + assert "ub" in str(buffer_ty) + + +if __name__ == '__main__': + test_build_buffer_type_with_affine_map() + test_extension_reexports_affine_bindings() + test_make_affine_map() diff --git a/third_party/ascend/unittest/pytest_ut/test_alloc.py b/third_party/ascend/unittest/pytest_ut/test_alloc.py new file mode 100644 index 0000000000..1b250ea35f --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_alloc.py @@ -0,0 +1,96 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import os + +import pytest +import triton +import triton.language as tl +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +import triton.extension.buffer.language as bl +import triton.language.extra.cann.extension as al +from triton._C.libtriton import ir, buffer_ir +from triton._C.libtriton.ascend import ir as ascend_ir + +os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" + + +class Options: + num_warps = 4 + num_stages = 3 + num_ctas = 1 + cluster_dims = (1, 1, 1) + enable_fp_fusion = True + debug = False + + +def compile_kernel(kernel, signature, constants): + """Helper to compile a kernel to MLIR.""" + src = ASTSource(kernel, signature, constants) + context = ir.context() + ir.load_dialects(context) + buffer_ir.load_dialects(context) + ascend_ir.load_dialects(context) + module = ast_to_ttir(kernel, src, context, Options(), {"create_address_space": al.semantic.create_address_space}, + {}) + return str(module) + + +# ============== Kernel definitions ============== + + +@triton.jit +def allocate_local_buffer(XBLOCK: tl.constexpr): + # this statement has no effect, just to test the builder + bl.alloc(tl.float32, [XBLOCK]) + bl.alloc(tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.UB) + bl.alloc(tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.L1) + bl.alloc(tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.L0A) + bl.alloc(tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.L0B) + bl.alloc(tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.L0C) + bl.alloc(tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.UB, is_mem_unique=True) + + +@triton.jit +def allocate_to_smem_buffer(x_ptr, XBLOCK: tl.constexpr): + x_l1_keep = bl.alloc(tl.float32, [XBLOCK], al.ascend_address_space.L1) + for i in range(XBLOCK): + offsets = tl.arange(0, XBLOCK) + x = tl.load(x_ptr + offsets, mask=offsets < XBLOCK) + bl.to_buffer(tensor=x, bind_buffer=x_l1_keep) + tl.store(x_ptr + offsets, x, mask=offsets < XBLOCK) + + +def test_allocate_local_buffer(): + """Test allocating local buffers in different address spaces.""" + mlir = compile_kernel(allocate_to_smem_buffer, {"x_ptr": "*fp32"}, {"XBLOCK": 256}) + print(f"✅ Generated MLIR ({len(mlir)} chars):\n") + + +# ============== Main for manual testing ============== + +if __name__ == "__main__": + print("=" * 60) + print("Test 1: Nested Scopes") + print("=" * 60) + mlir = compile_kernel(allocate_local_buffer, {}, {"XBLOCK": 256}) + print(f"✅ Generated MLIR ({len(mlir)} chars):\n") + print(mlir) diff --git a/third_party/ascend/unittest/pytest_ut/test_and.py b/third_party/ascend/unittest/pytest_ut/test_and.py new file mode 100644 index 0000000000..c9fdb44898 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_and.py @@ -0,0 +1,61 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common + + +def torch_and(x0, x1): + res = x0 & x1 + return res + + +@triton.jit +def triton_and(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x_index = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + x_index) + tmp1 = tl.load(in_ptr1 + x_index) + tmp2 = tmp0 & tmp1 + tl.store(out_ptr0 + x_index, tmp2) + + +@pytest.mark.parametrize('param_list', [ + ['int32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_and(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + torch_res = torch_and(x0, x1) + # triton结果 + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_and[ncore, 1, 1](x0, x1, triton_res, xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_annotations.py b/third_party/ascend/unittest/pytest_ut/test_annotations.py new file mode 100644 index 0000000000..8067dc8298 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_annotations.py @@ -0,0 +1,71 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +from __future__ import annotations +import torch +import triton +import triton.language as tl +import pytest + + +def annotated_function(return_type=None, **arg_types): + """A decorator to add annotations to a function.""" + + def decorator(func): + func.__annotations__ = {**arg_types, 'return': return_type} + return func + + return decorator + + +# Test integer annotations +@pytest.mark.parametrize(("signed", "width"), [ + (signed, width) for signed in [False, True]\ + for width in [8, 16, 32, 64] +] + [(False, 1)] + ) +def test_int_annotation(signed, width, device="npu"): + + @triton.jit + @annotated_function(X=torch.tensor, v=f"tl.{'' if signed else 'u'}int{width}") + def _kernel(X, v): + tl.store(X, v) + + h = _kernel[(1, )](torch.empty(1, device=device), 3) + pfx = 'si' if signed else 'ui' + assert f'%v: i{width}' in h.asm["ttir"] + assert f'arith.{pfx}tofp' in h.asm["ttir"] + + +# Test that unknown annotations do not emit an error +def test_unknown_annotation(device="npu"): + + @triton.jit + def _kernel(X: torch.Tensor, N: int, BLOCK_SIZE: tl.constexpr): + pass + + x = torch.empty(1, device=device) + _kernel[(1, )](x, x.shape[0], 32) + try: + _kernel[(1, )](x.shape[0], x.shape[0], 32) + except AttributeError: + pass diff --git a/third_party/ascend/unittest/pytest_ut/test_arange.py b/third_party/ascend/unittest/pytest_ut/test_arange.py new file mode 100644 index 0000000000..5c049bec80 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_arange.py @@ -0,0 +1,113 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import pytest +import torch +import torch_npu +import triton +import triton.language as tl +import test_common + + +def torch_arange(start, end): + TRITON_MAX_TENSOR_NUMEL = 1048576 + if end < start: + raise ValueError("arange's end argument must be greater than the start argument") + if end - start > TRITON_MAX_TENSOR_NUMEL: + raise ValueError( + f"end - start must be less than or equal to TRITON_MAX_TENSOR_NUMEL = {TRITON_MAX_TENSOR_NUMEL}") + return torch.arange(start, end) + + +def torch_arange_access(start, end): + z = torch.zeros([end], dtype=torch.int32).npu() + v = torch.arange(start, end).npu() + z[start:end] = v + return z + + +@triton.jit +def triton_arange(z, BLOCK: tl.constexpr, START: tl.constexpr, END: tl.constexpr): + off = tl.arange(0, BLOCK) + val = tl.arange(START, END) + tl.store(z + off, val) + + +@triton.jit +def triton_arange_access(z, BLOCK: tl.constexpr, START: tl.constexpr, END: tl.constexpr): + off = tl.arange(START, END) + val = tl.arange(START, END) + tl.store(z + off, val) + + +@pytest.mark.parametrize('param_list', [ + [0, 128], + [7, 128], + [128, 1024], +]) +def test_case(param_list): + start, end = param_list + shape = [end - start] + block = end - start + dtype = 'int32' + + y_ref = torch_arange(start, end) + y_cal = torch.zeros(shape, dtype=torch.int32).npu() + + triton_arange[(1, )](y_cal, START=start, END=end, BLOCK=block) + + test_common.validate_cmp(dtype, y_cal, y_ref) + + +@pytest.mark.parametrize('param_list', [ + [0, 128], + [7, 128], + [128, 1024], +]) +def test_case_access(param_list): + start, end = param_list + shape = [end] + block = end - start + dtype = 'int32' + + y_ref = torch_arange_access(start, end) + y_cal = torch.zeros(shape, dtype=torch.int32).npu() + + triton_arange_access[(1, )](y_cal, START=start, END=end, BLOCK=block) + + test_common.validate_cmp(dtype, y_cal, y_ref) + + +@pytest.mark.parametrize('invalid_param_list', [ + [0, 10000000], + [1024, 128], +]) +def test_arange_invalid_range(invalid_param_list): + start, end = invalid_param_list + shape = [end - start] + block = end - start + flag = False + try: + y_cal = torch.zeros(shape, dtype=torch.int32).npu() + triton_arange[(1, )](y_cal, START=start, END=end, BLOCK=block) + except Exception as e: + flag = True + assert flag diff --git a/third_party/ascend/unittest/pytest_ut/test_arch.py b/third_party/ascend/unittest/pytest_ut/test_arch.py new file mode 100644 index 0000000000..4946ca3538 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_arch.py @@ -0,0 +1,94 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import os +import pytest +import triton +import triton.language as tl +import triton.extension.buffer.language as bl +import triton.language.extra.cann.extension as al +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +from triton._C.libtriton import ir, buffer_ir +from triton._C.libtriton.ascend import ir as ascend_ir + +os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" + + +class Options: + num_warps = 4 + num_stages = 3 + num_ctas = 1 + cluster_dims = (1, 1, 1) + enable_fp_fusion = True + debug = False + arch = "Ascend950" + + +def compile_kernel(kernel, signature, constants): + """Helper to compile a kernel to MLIR.""" + src = ASTSource(kernel, signature, constants) + context = ir.context() + ir.load_dialects(context) + buffer_ir.load_dialects(context) + ascend_ir.load_dialects(context) + module = ast_to_ttir(kernel, src, context, Options(), {}, {}) + return str(module) + + +@triton.jit +def copy( + A_ptr, + A1_ptr, + M: tl.constexpr, + N: tl.constexpr, +): + offs_a = tl.arange(0, M)[:, None] + offs_b = tl.arange(0, N)[None, :] + + offs_c = (offs_a) * M + (offs_b) + a_ptr = A_ptr + offs_c + a_val = tl.load(a_ptr) + a1_ptr = A1_ptr + offs_c + a1_val = tl.load(a1_ptr) + + add = tl.add(a_val, a1_val) + + add_ub = bl.to_buffer(add, al.ascend_address_space.UB) + A_l1 = bl.alloc(tl.float32, [M, N], al.ascend_address_space.L1) + al.copy(add_ub, A_l1) + + +def test_arch(): + print("=" * 60) + print("Test 1: copy ") + print("=" * 60) + mlir = compile_kernel( + copy, + {"A_ptr": "*fp32", "A1_ptr": "*fp32"}, + {"M": 16, "N": 16}, + ) + print(f"✅ Generated MLIR ({len(mlir)} chars):\n") + print(mlir) + + +# ============== Main for manual testing ============== +if __name__ == "__main__": + test_arch() diff --git a/third_party/ascend/unittest/pytest_ut/test_argmax.py b/third_party/ascend/unittest/pytest_ut/test_argmax.py new file mode 100644 index 0000000000..3d84f2829d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_argmax.py @@ -0,0 +1,65 @@ +import logging +import math +import pytest +import torch +import torch_npu +import numpy as np +import triton +import triton.language as tl + +import test_common + + +def torch_argmax(x0, dim, keepdim): + x0 = x0 if x0.device == "cpu" else x0.cpu() + return torch.argmax(x0, dim=dim, keepdim=keepdim).npu() + + +@triton.jit +def triton_argmax_1d(in_ptr0, out_ptr1, xnumel, XBLOCK: tl.constexpr): + xoffset = tl.program_id(0) + tl.arange(0, XBLOCK) + tmp0 = tl.load(in_ptr0 + xoffset, None) + tmp4 = tl.argmax(tmp0, 0) + tl.store(out_ptr1, tmp4, None) + + +@pytest.mark.parametrize('shape', [(128, ), (256, ), (37, ), (741, )]) +@pytest.mark.parametrize('dtype', ['int32', 'float32', 'uint8', 'int8']) +def test_argmax_1d(dtype, shape): + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty(1, dtype=torch.int32).npu() + numel = shape[0] + triton_argmax_1d[(1, )](x0, triton_res, numel, numel) + torch_res = torch_argmax(x0, dim=0, keepdim=True) + test_common.validate_cmp("int32", triton_res, torch_res) + + +@triton.jit +def triton_argmax_2d(in_ptr0, out_ptr0, dim: tl.constexpr, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, + NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, MNUMEL) + nblk_idx = tl.arange(0, NNUMEL) + mmask = mblk_idx < M + nmask = nblk_idx < N + mask = (mmask[:, None]) & (nmask[None, :]) + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + x = tl.load(in_ptr0 + idx, mask=mask, other=float('-inf')) + tmp4 = tl.argmax(x, dim) + if dim == 0: + tl.store(out_ptr0 + tl.arange(0, N), tmp4, None) + else: + tl.store(out_ptr0 + tl.arange(0, M), tmp4, None) + + +@pytest.mark.parametrize('shape', [(37, 125), (29, 4), (7, 31)]) +@pytest.mark.parametrize('dtype', ['int32', 'float32', 'uint8', 'int8']) +@pytest.mark.parametrize('dim', [0, 1]) +def test_argmax_2d(dtype, shape, dim): + shapex, shapey = shape + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty([ + shape[1 - dim], + ], dtype=torch.int32).npu() + triton_argmax_2d[(1, 1)](x0, triton_res, dim, shapex, shapey, shapex, shapey) + torch_res = torch_argmax(x0, dim=dim, keepdim=False) + test_common.validate_cmp("int32", triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_argmin.py b/third_party/ascend/unittest/pytest_ut/test_argmin.py new file mode 100644 index 0000000000..16be10a5f4 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_argmin.py @@ -0,0 +1,65 @@ +import logging +import math +import pytest +import torch +import torch_npu +import numpy as np +import triton +import triton.language as tl + +import test_common + + +def torch_argmin(x0, dim, keepdim): + x0 = x0 if x0.device == "cpu" else x0.cpu() + return torch.argmin(x0, dim=dim, keepdim=keepdim).npu() + + +@triton.jit +def triton_argmin_1d(in_ptr0, out_ptr1, xnumel, XBLOCK: tl.constexpr): + xoffset = tl.program_id(0) + tl.arange(0, XBLOCK) + tmp0 = tl.load(in_ptr0 + xoffset, None) + tmp4 = tl.argmin(tmp0, 0) + tl.store(out_ptr1, tmp4, None) + + +@pytest.mark.parametrize('shape', [(128, ), (256, ), (37, ), (741, )]) +@pytest.mark.parametrize('dtype', ['int32', 'float32', 'uint8', 'int8']) +def test_argmin_1d(dtype, shape): + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty(1, dtype=torch.int32).npu() + numel = shape[0] + triton_argmin_1d[(1, )](x0, triton_res, numel, numel) + torch_res = torch_argmin(x0, dim=0, keepdim=True) + test_common.validate_cmp("int32", triton_res, torch_res) + + +@triton.jit +def triton_argmin_2d(in_ptr0, out_ptr0, dim: tl.constexpr, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, + NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, MNUMEL) + nblk_idx = tl.arange(0, NNUMEL) + mmask = mblk_idx < M + nmask = nblk_idx < N + mask = (mmask[:, None]) & (nmask[None, :]) + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + x = tl.load(in_ptr0 + idx, mask=mask, other=float('inf')) + tmp4 = tl.argmin(x, dim) + if dim == 0: + tl.store(out_ptr0 + tl.arange(0, N), tmp4, None) + else: + tl.store(out_ptr0 + tl.arange(0, M), tmp4, None) + + +@pytest.mark.parametrize('shape', [(37, 125), (29, 4), (7, 31)]) +@pytest.mark.parametrize('dtype', ['int32', 'float32', 'uint8', 'int8']) +@pytest.mark.parametrize('dim', [0, 1]) +def test_argmin_2d(dtype, shape, dim): + shapex, shapey = shape + x0 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.empty([ + shape[1 - dim], + ], dtype=torch.int32).npu() + triton_argmin_2d[(1, 1)](x0, triton_res, dim, shapex, shapey, shapex, shapey) + torch_res = torch_argmin(x0, dim=dim, keepdim=False) + test_common.validate_cmp("int32", triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_ascend_barrier.py b/third_party/ascend/unittest/pytest_ut/test_ascend_barrier.py new file mode 100755 index 0000000000..78573ac163 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_ascend_barrier.py @@ -0,0 +1,75 @@ +#!/usr/bin/env python3 +import os + +import pytest +import triton +import triton.language as tl +import triton.language.extra.cann.extension as al +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir + +os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" + + +class Options: + num_warps = 4 + num_stages = 3 + num_ctas = 1 + cluster_dims = (1, 1, 1) + enable_fp_fusion = True + debug = False + + +def compile_kernel(kernel, signature, constants): + """Helper to compile a kernel to MLIR.""" + src = ASTSource(kernel, signature, constants) + context = ir.context() + ir.load_dialects(context) + ascend_ir.load_dialects(context) + module = ast_to_ttir(kernel, src, context, Options(), {}, {}) + return str(module) + + +# ============== Kernel definitions ============== + + +@triton.jit +def kernel_debug_barrier(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): + """Test debug barrier.""" + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + with al.scope(core_mode="vector"): + al.debug_barrier(al.SYNC_IN_VF.VV_ALL) + x = tl.load(x_ptr + i, mask=i < n) + y = tl.load(y_ptr + i, mask=i < n) + result = x + y + tl.store(out_ptr + i, result, mask=i < n) + + +# ============== Pytest tests ============== + + +def test_debug_barrier(): + """Test debug barrier generates.""" + mlir = compile_kernel( + kernel_debug_barrier, + {"x_ptr": "*fp32", "y_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, + {"BLOCK": 256}, + ) + assert "annotation.mark" in mlir and "VV_ALL" in mlir + + +# ============== Main for manual testing ============== + +if __name__ == "__main__": + print("=" * 60) + print("Test: debug barrier") + print("=" * 60) + mlir = compile_kernel( + kernel_debug_barrier, + {"x_ptr": "*fp32", "y_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, + {"BLOCK": 256}, + ) + print(f"✅ Generated MLIR ({len(mlir)} chars):\n") + print(mlir) diff --git a/third_party/ascend/unittest/pytest_ut/test_asin.py b/third_party/ascend/unittest/pytest_ut/test_asin.py new file mode 100644 index 0000000000..8d9d7249ed --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_asin.py @@ -0,0 +1,90 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import torch +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + + +@triton.jit +def asin_kernel(x_ptr, y_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + + mask = offsets < n_elements + + x = tl.load(x_ptr + offsets, mask=mask) + + y = libdevice.asin(x) + + tl.store(y_ptr + offsets, y, mask=mask) + + +@pytest.mark.parametrize('shape', [ + (12, 16), +]) +@pytest.mark.parametrize('dtype', ['float32']) +def test_asin(shape, dtype): + n_elements = shape[0] * shape[1] + + x = test_common.generate_tensor(shape, dtype).npu() + + # Ensure to include some boundary cases + x[0, 0] = 0.0 + x[0, 1] = 0.5 + x[0, 2] = -0.5 + x[0, 3] = 1.0 + x[0, 4] = -1.0 + x[0, 5] = 0.707 # sin(π/4) + x[0, 6] = 0.866 # sin(π/3) + + # Add some out-of-range values + x[0, 7] = 1.1 + x[0, 8] = -1.1 + + y = torch.empty_like(x) + + BLOCK_SIZE = 192 + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + + asin_kernel[grid](x, y, n_elements, BLOCK_SIZE=BLOCK_SIZE) + + expected = torch.asin(x) + + # Check the accuracy for values within the effective range. + valid_mask = (x >= -1) & (x <= 1) + + if torch.any(valid_mask): + valid_y = y[valid_mask] + valid_expected = expected[valid_mask] + + torch.testing.assert_close(valid_y, valid_expected, rtol=1e-3, atol=1e-3) + + # Check if values outside the range return NaN + invalid_mask = (x < -1) | (x > 1) + if torch.any(invalid_mask): + invalid_y = y[invalid_mask] + assert torch.all(torch.isnan(invalid_y)), "Invalid inputs should return NaN" + + print("✓ ASIN test PASSED!") diff --git a/third_party/ascend/unittest/pytest_ut/test_asinh.py b/third_party/ascend/unittest/pytest_ut/test_asinh.py new file mode 100644 index 0000000000..38865e92be --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_asinh.py @@ -0,0 +1,58 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + +import torch +import torch_npu + + +@triton.jit +def triton_asinh(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = libdevice.asinh(tmp0) + tl.store(out_ptr0 + (x0), tmp1, None) + + +@pytest.mark.parametrize('param_list', ['float32', 'float16', 'bfloat16']) +def test_asinh_special(param_list): + dtype = param_list + x_near = torch.linspace(1.0 + 1e-6, 2.0, 500, dtype=eval("torch." + dtype)).npu() + x_far = torch.linspace(2.0, 1e4, 300, dtype=eval("torch." + dtype)).npu() + x0 = torch.cat([x_near, x_far], dim=0) + x_special = torch.tensor([1.0, 1.0 + 1e-8, 2.0, 1000.0, 1e4], dtype=eval("torch." + dtype)).npu() + x0 = torch.cat([x0, x_special], dim=0) + + if dtype == 'float16': + x0 = torch.clamp(x0, max=240) + y_ref = torch.asinh(x0) + y_cal = torch.zeros_like(y_ref) + triton_asinh[1, 1, 1](x0, y_cal, x0.shape[0], x0.shape[0]) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_asm.py b/third_party/ascend/unittest/pytest_ut/test_asm.py new file mode 100644 index 0000000000..ce4c7dca25 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_asm.py @@ -0,0 +1,96 @@ +import triton +import triton.language as tl +import numpy as np +import torch +import pytest +import test_common + + +def torch_add(x, y): + res = x + y + return res + + +@triton.jit +def triton_asm_add( + x_ptr, + y_ptr, + output_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = tl.inline_asm_elementwise( + asm=""" + ADD.s64 $0, $1, $2 + """, + constraints=("=l,l,l"), + args=[x, y], + dtype=tl.int64, + is_pure=True, + pack=1, + ) + tl.store(output_ptr + offsets, output, mask=mask) + + +@pytest.mark.parametrize('param_list', [ + ['int64', 4096, 1024], +]) +def test_case(param_list): + dtype, length, block_size = param_list + ncore = length // block_size + x = test_common.generate_tensor((length, ), dtype).npu() + y = test_common.generate_tensor((length, ), dtype).npu() + res_ref = torch_add(x, y) + res_cal = torch.zeros((length, ), dtype=eval('torch.' + dtype)).npu() + triton_asm_add[(ncore, )](x, y, res_cal, length, BLOCK_SIZE=block_size) + test_common.validate_cmp(dtype, res_cal, res_ref) + + +@triton.jit +def triton_asm_add_2d( + x_ptr, + y_ptr, + output_ptr, + M, + N, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid = tl.program_id(axis=0) + row_offsets = pid * BLOCK_M + tl.arange(0, BLOCK_M) + col_offsets = tl.arange(0, BLOCK_N) + offsets = row_offsets[:, None] * N + col_offsets[None, :] + mask = (row_offsets[:, None] < M) & (col_offsets[None, :] < N) + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = tl.inline_asm_elementwise( + asm=""" + ADD.s64 $0, $1, $2 + """, + constraints=("=l,l,l"), + args=[x, y], + dtype=tl.int64, + is_pure=True, + pack=1, + ) + tl.store(output_ptr + offsets, output, mask=mask) + + +@pytest.mark.parametrize('param_list', [ + ['int64', 64, 32, 16, 32], +]) +def test_case_2d(param_list): + dtype, M, N, block_m, block_n = param_list + ncore = M // block_m + x = test_common.generate_tensor((M, N), dtype).npu() + y = test_common.generate_tensor((M, N), dtype).npu() + res_ref = torch_add(x, y) + res_cal = torch.zeros((M, N), dtype=eval('torch.' + dtype)).npu() + triton_asm_add_2d[(ncore, )](x, y, res_cal, M, N, BLOCK_M=block_m, BLOCK_N=block_n) + test_common.validate_cmp(dtype, res_cal, res_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_asm_scalar.py b/third_party/ascend/unittest/pytest_ut/test_asm_scalar.py new file mode 100644 index 0000000000..23b8f6553c --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_asm_scalar.py @@ -0,0 +1,32 @@ +import torch +import torch_npu +import triton +import triton.language as tl +import pytest + + +@triton.jit +def triton_asm_time(output_ptr, ): + y = tl.inline_asm_elementwise( + asm=""" + MOV $0, SYS_CNT + """, + constraints="=l", + args=[], + dtype=(tl.int64), + is_pure=False, + pack=1, + ) + tl.store(output_ptr, y) + + +@pytest.mark.parametrize( + "param_list", + [[ + "int64", + ]], +) +def test_case(param_list): + (dtype, ) = param_list + res_cal = torch.zeros((1, ), dtype=eval("torch." + dtype)).npu() + triton_asm_time[(1, )](res_cal, ) diff --git a/third_party/ascend/unittest/pytest_ut/test_associative_scan.py b/third_party/ascend/unittest/pytest_ut/test_associative_scan.py new file mode 100644 index 0000000000..248d08cba4 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_associative_scan.py @@ -0,0 +1,179 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +import test_common + + +def torch_func(x, dim, reverse): + if reverse: + x = torch.flip(x, [dim]) + res = torch.cumsum(x, dim=dim) + return res + + +def combine_fn_test_torch(a, b, combine_fn): + return torch.maximum(a, b) + + +def torch_func_scan(x: torch.Tensor, dim: int, combine_fn='maximum', reverse=False): + """ + PyTorch implements associative_scan, with semantics fully aligned with Triton. + """ + dim = dim % x.ndim + + if reverse: + x = x.flip(dim) + + N = x.size(dim) + tensors = torch.unbind(x, dim=dim) + + outputs = [] + carry = tensors[0] + outputs.append(carry) + + for i in range(1, N): + carry = combine_fn_test_torch(tensors[i], carry, combine_fn) + outputs.append(carry) + + output = torch.stack(outputs, dim=dim) + + if reverse: + output = output.flip(dim) + + return output + + +@triton.jit +def combine_fn_test(a, b): + return tl.maximum(a, b) + + +@triton.jit +def triton_kernel_1d_scan( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + numel_x: tl.constexpr, + XBLOCK: tl.constexpr, +): + tl.static_assert(numel_x == XBLOCK, "numel_x must be equal to XBLOCK in this kernel") + idx = tl.arange(0, XBLOCK) + x = tl.load(in_ptr0 + idx) + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=combine_fn_test) + tl.store(out_ptr0 + idx, ret) + + +@triton.jit +def triton_kernel_2d_scan( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + numel_x: tl.constexpr, + numel_r: tl.constexpr, + XBLOCK: tl.constexpr, + RBLOCK: tl.constexpr, +): + tl.static_assert(numel_x == XBLOCK, "numel_x must be equal to XBLOCK in this kernel") + tl.static_assert(numel_r == RBLOCK, "numel_r must be equal to RBLOCK in this kernel") + idx_x = tl.arange(0, XBLOCK) + idx_r = tl.arange(0, RBLOCK) + idx = idx_x[:, None] * numel_r + idx_r[None, :] + x = tl.load(in_ptr0 + idx) + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=combine_fn_test) + tl.store(out_ptr0 + idx, ret) + + +@triton.jit +def triton_kernel_3d_scan( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + numel_x: tl.constexpr, + numel_r: tl.constexpr, + numel_z: tl.constexpr, + XBLOCK: tl.constexpr, + RBLOCK: tl.constexpr, + ZBLOCK: tl.constexpr, +): + tl.static_assert(numel_x == XBLOCK, "numel_x must be equal to XBLOCK in this kernel") + tl.static_assert(numel_r == RBLOCK, "numel_r must be equal to RBLOCK in this kernel") + tl.static_assert(numel_z == ZBLOCK, "numel_z must be equal to ZBLOCK in this kernel") + idx_x = tl.arange(0, XBLOCK) + idx_r = tl.arange(0, RBLOCK) + idx_z = tl.arange(0, ZBLOCK) + idx = idx_x[:, None, None] * numel_r * numel_z + idx_r[None, :, None] * numel_z + idx_z[None, None, :] + x = tl.load(in_ptr0 + idx) + ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=combine_fn_test) + tl.store(out_ptr0 + idx, ret) + + +def triton_func_scan(x, dim, reverse): + res = torch.empty_like(x) + print(f"res.dtype = {res.dtype}") + shape = x.size() + if len(shape) == 1: + if dim >= 1: + pytest.skip("dim >= 1 for 1D tensor, skipping.") + triton_kernel_1d_scan[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[0]) + elif len(shape) == 2: + if dim >= 2: + pytest.skip("dim >= 2 for 2D tensor, skipping.") + triton_kernel_2d_scan[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[1], x.shape[0], x.shape[1]) + elif len(shape) == 3: + if dim >= 3: + pytest.skip("dim >= 3 for 3D tensor, skipping.") + triton_kernel_3d_scan[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[1], x.shape[2], x.shape[0], x.shape[1], + x.shape[2]) + else: + pytest.skip(f"This testcase unsupported tensor dimension: {len(shape)}") + + return res + + +@pytest.mark.parametrize("dtype", ['int32', 'float32']) +@pytest.mark.parametrize("shape", [(128, ), (8, 4), (128, 4, 16)]) +@pytest.mark.parametrize("dim", [0, 1, 2]) +@pytest.mark.parametrize("combine_fn", [ + 'maximum', +]) +@pytest.mark.parametrize("reverse", [False, True]) +def test_scan(dtype, shape, dim, combine_fn, reverse): + torch.manual_seed(0) + x = test_common.generate_tensor(shape=shape, dtype=dtype) + x_gold = x + cpu_res = torch_func_scan(x_gold, dim, combine_fn, reverse) + print(f"cpu_res: {cpu_res}") + + x_npu = x.npu() + triton_res = triton_func_scan(x_npu, dim, reverse) + print(f"triton_res: {triton_res}") + + test_common.validate_cmp(dtype, triton_res, cpu_res) + print(f"Validate PASS") diff --git a/third_party/ascend/unittest/pytest_ut/test_associative_scan_multi_input.py b/third_party/ascend/unittest/pytest_ut/test_associative_scan_multi_input.py new file mode 100644 index 0000000000..6fb8813dc4 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_associative_scan_multi_input.py @@ -0,0 +1,157 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import triton +import triton.language as tl +import pytest + + +def combine_fn_test_torch_add(a, b): + return a + b + + +def torch_func_scan(x: torch.Tensor, dim: int, reverse: bool): + """ + PyTorch implements associative_scan, with semantics fully aligned with Triton. + """ + dim = dim % x.ndim + + if reverse: + x = x.flip(dim) + + N = x.size(dim) + tensors = torch.unbind(x, dim=dim) + + outputs = [] + carry = tensors[0] + outputs.append(carry) + + for i in range(1, N): + carry = combine_fn_test_torch_add(tensors[i], carry) + outputs.append(carry) + + output = torch.stack(outputs, dim=dim) + + if reverse: + output = output.flip(dim) + + return output + + +@triton.jit +def add_combine_fn(a_value, a_index, b_value, b_index): + new_val = a_value + b_value + new_idx = a_index + b_index + return (new_val, new_idx) + + +@triton.jit +def prefix_scan_last_dim_kernel( + vals_ptr, + idx_ptr, + out_vals_ptr, + out_idxs_ptr, + axis_size, + total_slices, + BLOCK_SIZE: tl.constexpr, + reverse: tl.constexpr, +): + slice_id = tl.program_id(0) + if slice_id >= total_slices: + return + + offs = tl.arange(0, BLOCK_SIZE) + mask = offs < axis_size + base = slice_id * axis_size + + vals = tl.load(vals_ptr + base + offs, mask=mask, other=0.0) + idxs = tl.load(idx_ptr + base + offs, mask=mask, other=0) + + pre_vals, pre_idxs = tl.associative_scan((vals, idxs), axis=0, combine_fn=add_combine_fn, reverse=reverse) + + tl.store(out_vals_ptr + base + offs, pre_vals, mask=mask) + tl.store(out_idxs_ptr + base + offs, pre_idxs, mask=mask) + + +def multi_input_prefix_sum(values: torch.Tensor, index: torch.Tensor, axis=0, reverse=False): + assert values.shape == index.shape + assert values.device == index.device + rank = values.ndim + if axis < 0: + axis += rank + assert 0 <= axis < rank + + # 1. change to make axis the last dim + order = list(range(rank)) + if axis != rank - 1: + order[axis], order[-1] = order[-1], order[axis] + inv_order = [0] * rank + for i, o in enumerate(order): + inv_order[o] = i + + vals_p = values.permute(order).contiguous() + idxs_p = index.permute(order).contiguous() + + shape_p = vals_p.shape + axis_size = shape_p[-1] + total_slices = 1 + for d in shape_p[:-1]: + total_slices *= d + + out_vals_p = torch.empty_like(vals_p) + out_idxs_p = torch.empty_like(idxs_p) + + BLOCK_SIZE = 1 << (axis_size - 1).bit_length() + + prefix_scan_last_dim_kernel[(total_slices, )](vals_p, idxs_p, out_vals_p, out_idxs_p, axis_size, total_slices, + BLOCK_SIZE=BLOCK_SIZE, reverse=reverse) + + # 2. permute back + out_vals = out_vals_p.permute(inv_order) + out_idxs = out_idxs_p.permute(inv_order) + return out_vals, out_idxs + + +@pytest.mark.parametrize("shape, axis", [ + ((10, ), 0), + ((4, 4), 0), + ((4, 4), 1), + ((2, 10, 5), 0), + ((2, 10, 5), 1), + ((2, 10, 5), 2), +]) +@pytest.mark.parametrize("reverse", [False, True]) +def test_multi_input_prefix_sum(shape, axis, reverse): + torch.manual_seed(0) + device = "npu" + + values = torch.randn(shape, device=device, dtype=torch.float32) + index = torch.arange(values.numel(), device=device, dtype=torch.int32).reshape(shape) + + triton_vals, triton_idxs = multi_input_prefix_sum(values, index, axis=axis, reverse=reverse) + + torch_vals = torch_func_scan(values, axis, reverse) + torch_idxs = torch_func_scan(index, axis, reverse) + + assert torch.allclose(triton_vals, torch_vals, rtol=1e-3, atol=1e-3), \ + f"数值不匹配！shape={shape}, axis={axis}\nTriton: {triton_vals}\nPyTorch: {torch_vals}" + assert torch.equal(triton_idxs, torch_idxs), \ + f"索引不匹配！shape={shape}, axis={axis}\nTriton: {triton_idxs}\nPyTorch: {torch_idxs}" diff --git a/third_party/ascend/unittest/pytest_ut/test_assume.py b/third_party/ascend/unittest/pytest_ut/test_assume.py new file mode 100644 index 0000000000..b18c0f2db5 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_assume.py @@ -0,0 +1,51 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import pytest +import test_common + + +@triton.jit +def triton_assume(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + tl.assume((XBLOCK & (XBLOCK - 1)) == 0) + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + xindex = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = xindex < xnumel + x0 = xindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + num = tl.sum(tmp0) + tl.assume(num > 0) + tmp2 = tmp0 + tl.store(out_ptr0 + (xindex), tmp2, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_assume(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + y_ref = x0 + y_cal = test_common.generate_tensor(shape, dtype).npu() + triton_assume[(ncore, )](x0, y_cal, x0.numel(), xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_assume1.py b/third_party/ascend/unittest/pytest_ut/test_assume1.py new file mode 100644 index 0000000000..05d5630092 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_assume1.py @@ -0,0 +1,32 @@ +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + +from triton._internal_testing import (is_interpreter) + + +@triton.jit +def assume(out_ptr, N: tl.constexpr, BLOCK_N: tl.constexpr): + current_size = N - tl.program_id(0) * BLOCK_N + tl.assume(current_size >= BLOCK_N) + if current_size >= BLOCK_N: + tl.store(out_ptr + tl.program_id(0), current_size) + else: + tl.store(out_ptr + tl.program_id(0), current_size + 101024) + + +@pytest.mark.parametrize('dtype', ["float32"]) +def test_assume(dtype): + NBLOCKS = 1024 // 128 + BLOCK_N = 128 + N = 1024 + output = torch.zeros(NBLOCKS, device='npu') + pgm = assume[(NBLOCKS, )](output, N=N, BLOCK_N=BLOCK_N) + + if is_interpreter(): + return + + assert 'llvm.intr.assume' in pgm.asm['ttadapter'] diff --git a/third_party/ascend/unittest/pytest_ut/test_atan.py b/third_party/ascend/unittest/pytest_ut/test_atan.py new file mode 100644 index 0000000000..eca6b31f07 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_atan.py @@ -0,0 +1,101 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + +import triton.language.extra.cann.libdevice as libdevice + + +def standard_unary(x0, dtype): + res = torch.atan(x0) + return res + + +def standard_binary(x0, y0, dtype): + res = x0 + y0 + return res + + +@triton.jit +def triton_elementwise_unary(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + ret = libdevice.atan(x) + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +@triton.jit +def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + y = tl.load(in_ptr1 + idx_block, mask=idx_block < N) + ret = x + y + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + # (torch.bfloat16, 'bfloat16'), + # (torch.int8, 'int8'), + # (torch.int16, 'int16'), + # (torch.int32, 'int32'), + # (torch.int64, 'int64'), +] + +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('N,NUMEL', shapes) +def test_elementwsie_common(dtype, sigtype, N, NUMEL): + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == "int64": + N = map_for_64_t[N] if N in map_for_64_t else N + + print(f"elementwise : ({N},) {dtype} {sigtype}") + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standard_unary(x0, dtype) + x0 = x0.npu() + print(ans) + + out = torch.zeros((N, ), dtype=dtype).npu() + triton_elementwise_unary[1, 1, 1](x0, out, N=N, NUMEL=NUMEL, debug=True) + print(out) + + test_common.validate_cmp(sigtype, out, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_atan2.py b/third_party/ascend/unittest/pytest_ut/test_atan2.py new file mode 100644 index 0000000000..121866037d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_atan2.py @@ -0,0 +1,129 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + +import triton.language.extra.cann.libdevice as libdevice + + +def standard_unary(x0, y0, dtype): + res = torch.atan2(y0, x0) + return res + + +def standard_binary(x0, y0, dtype): + res = x0 + y0 + return res + + +@triton.jit +def triton_elementwise_unary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + y = tl.load(in_ptr1 + idx_block, mask=idx_block < N) + ret = libdevice.atan2(y, x) + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +@triton.jit +def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + y = tl.load(in_ptr1 + idx_block, mask=idx_block < N) + ret = x + y + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), +] + +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('N,NUMEL', shapes) +def test_atan2_common(dtype, sigtype, N, NUMEL): + if N == 0: + return + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == "int64": + N = map_for_64_t[N] if N in map_for_64_t else N + + print(f"elementwise : ({N},) {dtype} {sigtype}") + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + y0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standard_unary(x0, y0, dtype) + x0 = x0.npu() + y0 = y0.npu() + + out = torch.zeros((N, ), dtype=dtype).npu() + triton_elementwise_unary[1, 1, 1](x0, y0, out, N=N, NUMEL=NUMEL, debug=True) + + test_common.validate_cmp(sigtype, out, ans) + + +input_vals = [ + (0.0, 1.0), + (0.0, -1.0), +] + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('N,NUMEL', shapes[0:2]) +@pytest.mark.parametrize('X,Y', input_vals) +def test_atan2_special(dtype, sigtype, N, NUMEL, X, Y): + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == "int64": + N = map_for_64_t[N] if N in map_for_64_t else N + + print(f"elementwise : ({N},) {dtype} {sigtype}") + + x0 = torch.full((N, ), X, dtype=eval(f'torch.{sigtype}')) + y0 = torch.full((N, ), Y, dtype=eval(f'torch.{sigtype}')) + + ans = standard_unary(x0, y0, dtype) + x0 = x0.npu() + y0 = y0.npu() + + out = torch.zeros((N, ), dtype=dtype).npu() + triton_elementwise_unary[1, 1, 1](x0, y0, out, N=N, NUMEL=NUMEL, debug=True) + + test_common.validate_cmp(sigtype, out, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_atanh.py b/third_party/ascend/unittest/pytest_ut/test_atanh.py new file mode 100644 index 0000000000..a8d8ad3fe1 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_atanh.py @@ -0,0 +1,77 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + +import torch +import torch_npu + + +def torch_atanh(x0): + return torch.atanh(x0) + + +@triton.jit +def triton_atanh(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = libdevice.atanh(tmp0) + tl.store(out_ptr0 + (x0), tmp1, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['bfloat16', (2, 4096, 8), 2, 32768, 1024], +]) +def test_atanh_common(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu().clamp(-0.99, 0.99) + y_ref = torch_atanh(x0) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_atanh[ncore, 1, 1](x0, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) + + +input_vals = [-1, 1] + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['bfloat16', (2, 4096, 8), 2, 32768, 1024], +]) +@pytest.mark.parametrize('input_val', input_vals) +def test_atanh_special(param_list, input_val): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = torch.full(shape, input_val, dtype=eval('torch.' + dtype)).npu() + y_ref = torch_atanh(x0) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_atanh[ncore, 1, 1](x0, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_atomic_add.py b/third_party/ascend/unittest/pytest_ut/test_atomic_add.py new file mode 100644 index 0000000000..967717ad33 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_atomic_add.py @@ -0,0 +1,207 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import pytest +import test_common +import torch +import torch_npu + + +@triton.jit +def atomic_add(in_ptr0, out_ptr0, out_ptr1, n_elements, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + yindex = tl.arange(0, BLOCK_SIZE)[:] + xmask = xindex < n_elements + x0 = xindex + x1 = yindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.atomic_add(out_ptr0 + (x1), tmp0, xmask) + tl.store(out_ptr1 + (x1), tmp1, xmask) + + +@triton.jit +def atomic_add_supply(in_ptr0, out_ptr0, n_elements, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + yindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + xmask = xindex < n_elements + x0 = xindex + x1 = yindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.atomic_add(out_ptr0 + (x1), tmp0, xmask) + + +@triton.jit +def atomic_add_for_load_offset(index_ptr, in_ptr0, out_ptr0): + index = tl.atomic_add(index_ptr, 1) + val = tl.load(in_ptr0 + index) + tl.store(out_ptr0, val) + + +@triton.jit +def atomic_add_for_store_offset(index_ptr, out_ptr0): + index = tl.atomic_add(index_ptr, 1) + tl.store(out_ptr0 + index, 1) + + +@pytest.mark.parametrize('param_list', [ + ['int64', (256, 32), 2], + ['int32', (32, 32), 2], + ['int16', (32, 32), 2], + ['int8', (32, 32), 2], + ['uint8', (32, 32), 2], + ['float32', (32, 32), 2], + ['float16', (64, 64), 4], + ['bfloat16', (64, 64), 4], + ['float32', (128, 128), 8], + ['float16', (128, 128), 16], + ['float32', (32768, 16), 32], +]) +def test_atomic_add(param_list): + dtype, shape, ncore = param_list + block_size = shape[0] * shape[1] / ncore + split_size = shape[0] // ncore + x0_value = 3 + x0 = torch.full(shape, x0_value, dtype=eval(f'torch.{dtype}')).npu() + if dtype == 'int64': + x1 = torch.randint(-10**15, 10**15, (split_size, shape[1]), dtype=eval(f'torch.{dtype}')).npu() + else: + x1 = torch.full((split_size, shape[1]), 2, dtype=eval(f'torch.{dtype}')).npu() + y = torch.full((split_size, shape[1]), -10, dtype=eval(f'torch.{dtype}')).npu() + + y_ref = x1 + 0 + x1_ref = x1 + ncore * x0_value + + n_elements = shape[0] * shape[1] + atomic_add[ncore, 1, 1](x0, x1, y, n_elements, BLOCK_SIZE=split_size * shape[1]) + test_common.validate_cmp(dtype, x1, x1_ref) + + +@pytest.mark.parametrize('param_list', [ + ['int16', (32, 32), 1], + ['int32', (32, 32), 1], + ['float32', (32, 32), 1], + ['float16', (64, 64), 1], +]) +def test_atomic_add_return_value(param_list): + dtype, shape, ncore = param_list + block_size = shape[0] * shape[1] / ncore + split_size = shape[0] // ncore + x0_value = 3 + x0 = torch.full(shape, x0_value, dtype=eval(f'torch.{dtype}')).npu() + x1 = torch.full((split_size, shape[1]), 2, dtype=eval(f'torch.{dtype}')).npu() + y = torch.full((split_size, shape[1]), -10, dtype=eval(f'torch.{dtype}')).npu() + + y_ref = x1 + 0 + x1_ref = x1 + ncore * x0_value + + n_elements = shape[0] * shape[1] + atomic_add[ncore, 1, 1](x0, x1, y, n_elements, BLOCK_SIZE=split_size * shape[1]) + test_common.validate_cmp(dtype, x1, x1_ref) + test_common.validate_cmp(dtype, y, y_ref) + + +@triton.jit +def atomic_add_2d(in_ptr0, out_ptr0, out_ptr1, numel_0, numel_1, BLOCK_SIZE_0: tl.constexpr, + BLOCK_SIZE_1: tl.constexpr): + pid = tl.program_id(0) + idx0_in = pid * BLOCK_SIZE_0 + tl.arange(0, BLOCK_SIZE_0)[:, None] + idx0_out = tl.arange(0, BLOCK_SIZE_0)[:, None] + idx1 = tl.arange(0, BLOCK_SIZE_1)[None, :] + idx_in = idx0_in * BLOCK_SIZE_1 + idx1 + idx_out = idx0_out * BLOCK_SIZE_1 + idx1 + msk_in = (idx0_in < numel_0) & (idx1 < numel_1) + msk_out = (idx0_out < numel_0) & (idx1 < numel_1) + tmp0 = tl.load(in_ptr0 + idx_in, msk_in) + tmp1 = tl.atomic_add(out_ptr0 + idx_out, tmp0, msk_out) + tl.store(out_ptr1 + idx_out, tmp1, msk_out) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (32, 32), 2], +]) +def test_atomic_add_2d(param_list): + dtype, shape, ncore = param_list + split_size = shape[0] // ncore + block_size_0 = split_size + block_size_1 = shape[1] + x0_value = 3 + x0 = torch.full(shape, x0_value, dtype=eval('torch.float32')).npu() + x1 = torch.full((split_size, shape[1]), 2, dtype=eval('torch.float32')).npu() + y = torch.full((split_size, shape[1]), -10, dtype=eval('torch.float32')).npu() + + y_ref = x1 + 0 + x1_ref = x1 + ncore * x0_value + + atomic_add_2d[ncore, 1, 1](x0, x1, y, shape[0], shape[1], BLOCK_SIZE_0=block_size_0, BLOCK_SIZE_1=block_size_1) + test_common.validate_cmp(dtype, x1, x1_ref) + + +@pytest.mark.parametrize('shape', [(3, 1), (13, 1), (32, 1), (256, 1)]) +@pytest.mark.parametrize('dtype', ['float32']) +def test_atomic_add_2d_supply(dtype, shape): + ncore = 1 + block_size = shape[0] * shape[1] / ncore + split_size = shape[0] // ncore + x0_value = 3 + x0 = torch.full(shape, x0_value, dtype=eval('torch.' + dtype)).npu() + x1 = torch.full((split_size, shape[1]), 2, dtype=eval('torch.' + dtype)).npu() + + y_ref = x1 + 0 + x1_ref = x1 + ncore * x0_value + + n_elements = shape[0] * shape[1] + atomic_add_supply[shape[0], 1, 1](x0, x1, n_elements, BLOCK_SIZE=shape[1]) + test_common.validate_cmp(dtype, x1, x1_ref) + + +def test_atomic_add_for_load_offset(): + index = torch.tensor([1]).npu() + input_tensor = torch.zeros(5).npu() + output = torch.tensor([1]).npu() + index_ref = index.clone() + index_ref += 1 + output_ref = output.clone() + output_ref = input_tensor[index] + + atomic_add_for_load_offset[(1, )](index, input_tensor, output) + torch.equal(index, index_ref) + torch.equal(output, output_ref) + + +def test_atomic_add_for_store_offset(): + index = torch.tensor([1]).npu() + output = torch.zeros(5).npu() + index_ref = index.clone() + index_ref += 1 + output_ref = output.clone() + output_ref[index] = 1 + + atomic_add_for_store_offset[(1, )](index, output) + torch.equal(index, index_ref) + torch.equal(output, output_ref) + + +if __name__ == "__main__": + param_list = ['float32', (32, 32), 2] + test_atomic_add_2d(param_list) diff --git a/third_party/ascend/unittest/pytest_ut/test_atomic_and.py b/third_party/ascend/unittest/pytest_ut/test_atomic_and.py new file mode 100644 index 0000000000..675f65f773 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_atomic_and.py @@ -0,0 +1,68 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import pytest +import test_common +import torch +import torch_npu + + +@triton.jit +def atomic_and(in_ptr0, out_ptr0, out_ptr1, n_elements, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + yindex = tl.arange(0, BLOCK_SIZE)[:] + xmask = xindex < n_elements + x0 = xindex + x1 = yindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.atomic_and(out_ptr0 + (x1), tmp0, xmask) + tl.store(out_ptr1 + (x1), tmp1, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['int64', (32, 32), 2], + ['int32', (32, 32), 2], + ['int16', (32, 32), 2], + ['int8', (16, 16), 4], + ['uint8', (16, 16), 4], +]) +def test_atomic_and(param_list): + dtype, shape, ncore = param_list + block_size = shape[0] * shape[1] // ncore + split_size = shape[0] // ncore + + val = torch.randint(low=0, high=10, size=shape, dtype=eval(f'torch.{dtype}')).npu() + + pointer = torch.randint(low=0, high=10, size=(split_size, shape[1]), dtype=eval(f'torch.{dtype}')).npu() + pointer_old = torch.full_like(pointer, -10).npu() + pointer_ref = pointer.clone() + + for i in range(ncore - 1): + pointer_ref &= val[(i * split_size):((i + 1) * split_size)] + + pointer_ref_last = pointer_ref.clone() + pointer_ref &= val[((ncore - 1) * split_size):(ncore * split_size)] + + n_elements = shape[0] * shape[1] + atomic_and[ncore, 1, 1](val, pointer, pointer_old, n_elements, BLOCK_SIZE=split_size * shape[1]) + test_common.validate_cmp(dtype, pointer, pointer_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_atomic_cas.py b/third_party/ascend/unittest/pytest_ut/test_atomic_cas.py new file mode 100644 index 0000000000..e834aa8d79 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_atomic_cas.py @@ -0,0 +1,190 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import pytest +import test_common +import torch +import torch_npu + +types_all = [ + (torch.float32, 'float32'), +] + + +def ceil_div(a, b): + return (a + b - 1) // b + + +@triton.jit +def atomic_cas(in_ptr0, in_ptr1, out_ptr0, out_ptr1, n_elements, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + yindex = tl.arange(0, BLOCK_SIZE)[:] + xmask = xindex < n_elements + x0 = xindex + x1 = yindex + val = tl.load(in_ptr0 + (x0), xmask) + cmp = tl.load(in_ptr1 + (x0), xmask) + tmp1 = tl.atomic_cas(out_ptr0 + (x1), cmp, val) + tl.store(out_ptr1 + (x1), tmp1, xmask) + + +@triton.jit +def atomic_cas_with_full( + ptr, + out, + n_elements, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + x = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = x < n_elements + + cmp = tl.full((BLOCK_SIZE, ), 2.0, tl.float32) + val = tl.full((BLOCK_SIZE, ), 1.0, tl.float32) + + old = tl.atomic_cas(ptr + x, cmp, val) # in_ptr(origin 2) -> ref: 1 X + tl.store(out + x, old, mask=mask) # out(origin 1) -> ref: old in_ptr(2) √ + + +@triton.jit +def atomic_cas_without_full( + ptr, + cmp_ptr, + val_ptr, + out_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + x = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = x < n_elements + + cmp = tl.load(cmp_ptr + x, mask) # 2 + val = tl.load(val_ptr + x, mask) # 1 + + old = tl.atomic_cas(ptr + x, cmp, val) # old : 2 + tl.store(out_ptr + x, old, mask=mask) + + +@pytest.mark.parametrize('param_list', [ + ['int16', (8, 8), 2], + ['int32', (32, 32), 6], + ['int64', (32, 32), 2], + ['float32', (32, 32), 2], + ['float16', (64, 64), 4], + ['float32', (128, 128), 8], + ['float16', (128, 128), 16], +]) +def test_atomic_cas(param_list): + dtype, shape, ncore = param_list + block_size = shape[0] * shape[1] // ncore + split_size = shape[0] // ncore + + import random + cmp_val = [random.randint(0, 10) for _ in range(ncore)] + + cmp = torch.ones(split_size, shape[1], dtype=eval(f'torch.{dtype}')).to().npu() * cmp_val[0] + for i in range(1, ncore): + append = torch.ones(split_size, shape[1], dtype=eval(f'torch.{dtype}')).to().npu() * cmp_val[i] + cmp = torch.cat([cmp, append], dim=0) + + val = torch.randint(low=0, high=10, size=shape, dtype=eval(f'torch.{dtype}')).npu() + + pointer = torch.randint(low=0, high=10, size=(split_size, shape[1]), dtype=eval(f'torch.{dtype}')).npu() + pointer_old = torch.full_like(pointer, -10).npu() + pointer_ref = pointer.clone() + + for i in range(ncore): + val_subview = val[(i * split_size):((i + 1) * split_size)] + pointer_ref = torch.where(pointer_ref == cmp_val[i], val_subview, pointer_ref) + + n_elements = shape[0] * shape[1] + atomic_cas[ncore, 1, 1](val, cmp, pointer, pointer_old, n_elements, BLOCK_SIZE=split_size * shape[1]) + test_common.validate_cmp(dtype, pointer, pointer_ref) + + +@pytest.mark.parametrize('param_list', [ + ['int16', (8, 8), 1], + ['int32', (32, 32), 1], + ['float32', (32, 32), 1], +]) +def test_atomic_cas_return_value(param_list): + dtype, shape, ncore = param_list + block_size = shape[0] * shape[1] // ncore + split_size = shape[0] // ncore + + import random + cmp_val = [random.randint(0, 10) for _ in range(ncore)] + + cmp = torch.ones(split_size, shape[1], dtype=eval(f'torch.{dtype}')).to().npu() * cmp_val[0] + for i in range(1, ncore): + append = torch.ones(split_size, shape[1], dtype=eval(f'torch.{dtype}')).to().npu() * cmp_val[i] + cmp = torch.cat([cmp, append], dim=0) + + val = torch.randint(low=0, high=10, size=shape, dtype=eval(f'torch.{dtype}')).npu() + + pointer = torch.randint(low=0, high=10, size=(split_size, shape[1]), dtype=eval(f'torch.{dtype}')).npu() + pointer_old_ref = pointer.clone() + pointer_old = torch.full_like(pointer, -10).npu() + pointer_ref = pointer.clone() + + for i in range(ncore): + val_subview = val[(i * split_size):((i + 1) * split_size)] + pointer_ref = torch.where(pointer_ref == cmp_val[i], val_subview, pointer_ref) + + n_elements = shape[0] * shape[1] + atomic_cas[ncore, 1, 1](val, cmp, pointer, pointer_old, n_elements, BLOCK_SIZE=split_size * shape[1]) + test_common.validate_cmp(dtype, pointer, pointer_ref) + test_common.validate_cmp(dtype, pointer_old, pointer_old_ref) + + +@pytest.mark.parametrize('dtype,sigtype', types_all) +@pytest.mark.parametrize('n_elements, BLOCK_SIZE', [(4096, 256)]) +@pytest.mark.skip(reason="full tensor has problem, skipped") +def test_atomic_cas_with_full(n_elements, BLOCK_SIZE, dtype, sigtype): + in_ptr = torch.full((n_elements, ), 2, dtype=dtype).npu() + out_ptr = torch.empty_like(in_ptr) + + grid = (ceil_div(n_elements, BLOCK_SIZE), 1, 1) + atomic_cas_with_full[grid](in_ptr, out_ptr, n_elements, BLOCK_SIZE=BLOCK_SIZE) + + # old should be all 2 (for in-range) + torch.testing.assert_close(out_ptr, torch.full_like(out_ptr, 2.0)) + + # final ptr should be all 1 + torch.testing.assert_close(in_ptr, torch.ones_like(in_ptr)) + + +@pytest.mark.parametrize('dtype,sigtype', types_all) +@pytest.mark.parametrize('n_elements, BLOCK_SIZE', [(4096, 256)]) +def test_atomic_cas_without_full(n_elements, BLOCK_SIZE, dtype, sigtype): + in_ptr = torch.full((n_elements, ), 2, dtype=dtype).npu() + cmp_ptr = torch.full((n_elements, ), 2, dtype=dtype).npu() + val_ptr = torch.full((n_elements, ), 1, dtype=dtype).npu() + out_ptr = torch.full((n_elements, ), 1, dtype=dtype).npu() # ref: in_ptr + + grid = (ceil_div(n_elements, BLOCK_SIZE), 1, 1) + atomic_cas_without_full[grid](in_ptr, cmp_ptr, val_ptr, out_ptr, n_elements, BLOCK_SIZE=BLOCK_SIZE) + + torch.testing.assert_close(in_ptr, torch.full_like(in_ptr, 1.0)) + torch.testing.assert_close(out_ptr, torch.full_like(out_ptr, 2.0)) diff --git a/third_party/ascend/unittest/pytest_ut/test_atomic_max.py b/third_party/ascend/unittest/pytest_ut/test_atomic_max.py new file mode 100644 index 0000000000..7a7a200e24 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_atomic_max.py @@ -0,0 +1,110 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import pytest +import test_common +import torch +import torch_npu +import numpy as np + + +@triton.jit +def triton_test_fn_atomic_max_dma(in_ptr0, out_ptr0, out_ptr1, n_elements, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + yindex = tl.arange(0, BLOCK_SIZE)[:] + xmask = xindex < n_elements + x0 = xindex + x1 = yindex + tmp0 = tl.load(in_ptr0 + (x0)) + # only set mask of atomic_max + tl.atomic_max(out_ptr0 + (x1), tmp0, xmask) + + +@triton.jit +def triton_test_fn_atomic_max_dma_supply(in_ptr0, out_ptr0, n_elements: tl.constexpr, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + yindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + xmask = xindex < n_elements + x0 = xindex + x1 = yindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.atomic_max(out_ptr0 + (x1), tmp0, xmask) + + +# torch.max do not support int +@pytest.mark.parametrize('param_list', [ + ['uint8', (32, 32), 2], + ['int16', (32, 32), 2], + ['bfloat16', (32, 32), 2], + ['float16', (32, 32), 2], + ['float32', (128, 128), 8], + ['float32', (32768, 16), 32], + ['int32', (32, 32), 2], + ['int32', (128, 128), 8], + ['int32', (32768, 16), 32], + ['int64', (32, 32), 2], + ['int64', (128, 128), 8], + ['int64', (8192, 16), 32], +]) +def test_atomic_max(param_list): + dtype, shape, ncore = param_list + block_size = shape[0] * shape[1] / ncore + split_size = shape[0] // ncore + # old size: (32768, 256) + # tensor of (1024, 256) is too big, and it will lead to failure in the backend + # so here we make it smaller + x0 = test_common.generate_tensor(shape, dtype) + x1 = test_common.generate_tensor((split_size, shape[1]), dtype) + y = test_common.generate_tensor((split_size, shape[1]), dtype) + + merged_tensor = torch.cat((x0, x1), dim=0) + chunks = torch.stack(torch.chunk(merged_tensor, ncore + 1, dim=0)) + x1_ref = torch.max(chunks, dim=0)[0] + x0 = x0.npu() + x1 = x1.npu() + y = y.npu() + + n_elements = shape[0] * shape[1] + triton_test_fn_atomic_max_dma[ncore, 1, 1](x0, x1, y, n_elements, BLOCK_SIZE=split_size * shape[1]) + test_common.validate_cmp(dtype, x1, x1_ref) + + +@pytest.mark.parametrize('shape', [(3, 1), (13, 1), (32, 1), (256, 1)]) +@pytest.mark.parametrize('dtype', ['float32']) +def test_atomic_max_2d_supply(dtype, shape): + # old size: (32768, 256) + # tensor of (1024, 256) is too big, and it will lead to failure in the backend + # so here we make it smaller + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + + x1_ref = torch.maximum(x0, x1) + + n_elements = shape[0] * shape[1] + triton_test_fn_atomic_max_dma_supply[shape[0], 1, 1](x0, x1, n_elements, BLOCK_SIZE=shape[1]) + test_common.validate_cmp(dtype, x1, x1_ref) + + +# if __name__ == "__main__": +# test_atomic_max(['int32', (8, 8), 2]) diff --git a/third_party/ascend/unittest/pytest_ut/test_atomic_min.py b/third_party/ascend/unittest/pytest_ut/test_atomic_min.py new file mode 100644 index 0000000000..461a3baa9c --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_atomic_min.py @@ -0,0 +1,94 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import pytest +import test_common +import torch +import torch_npu +import numpy as np + + +@triton.jit +def triton_test_fn_atomic_min_dma(in_ptr0, out_ptr0, out_ptr1, n_elements, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + yindex = tl.arange(0, BLOCK_SIZE)[:] + xmask = xindex < n_elements + x0 = xindex + x1 = yindex + tmp0 = tl.load(in_ptr0 + (x0)) + tl.atomic_min(out_ptr0 + (x1), tmp0, xmask) + + +@triton.jit +def triton_test_fn_atomic_min_dma_supply(in_ptr0, out_ptr0, n_elements: tl.constexpr, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + yindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + xmask = xindex < n_elements + x0 = xindex + x1 = yindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.atomic_min(out_ptr0 + (x1), tmp0, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['uint8', (32, 32), 2], + ['int8', (32, 32), 2], + ['int16', (32, 32), 2], + ['int32', (32, 32), 2], + ['int64', (32, 32), 2], + ['bfloat16', (64, 64), 4], + ['float16', (64, 64), 4], + ['float32', (32, 32), 2], +]) +def test_atomic_min(param_list): + dtype, shape, ncore = param_list + block_size = shape[0] * shape[1] / ncore + split_size = shape[0] // ncore + x0 = test_common.generate_tensor(shape, dtype) + x1 = test_common.generate_tensor((split_size, shape[1]), dtype) + y = test_common.generate_tensor((split_size, shape[1]), dtype) + + merged_tensor = torch.cat((x0, x1), dim=0) + chunks = torch.stack(torch.chunk(merged_tensor, ncore + 1, dim=0)) + x1_ref = torch.min(chunks, dim=0)[0] + x0 = x0.npu() + x1 = x1.npu() + y = y.npu() + + n_elements = shape[0] * shape[1] + triton_test_fn_atomic_min_dma[ncore, 1, 1](x0, x1, y, n_elements, BLOCK_SIZE=split_size * shape[1]) + test_common.validate_cmp(dtype, x1, x1_ref) + + +@pytest.mark.parametrize('shape', [(3, 1), (13, 1), (32, 1), (256, 1)]) +@pytest.mark.parametrize('dtype', ['float32']) +def test_atomic_min_2d_supply(dtype, shape): + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + + x1_ref = torch.minimum(x0, x1) + + n_elements = shape[0] * shape[1] + triton_test_fn_atomic_min_dma_supply[shape[0], 1, 1](x0, x1, n_elements, BLOCK_SIZE=shape[1]) + test_common.validate_cmp(dtype, x1, x1_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_atomic_or.py b/third_party/ascend/unittest/pytest_ut/test_atomic_or.py new file mode 100644 index 0000000000..0b627ce69e --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_atomic_or.py @@ -0,0 +1,67 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import pytest +import test_common +import torch +import torch_npu + + +@triton.jit +def atomic_or(in_ptr0, out_ptr0, out_ptr1, n_elements, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + yindex = tl.arange(0, BLOCK_SIZE)[:] + xmask = xindex < n_elements + x0 = xindex + x1 = yindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.atomic_or(out_ptr0 + (x1), tmp0, xmask) + tl.store(out_ptr1 + (x1), tmp1, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['int64', (32, 32), 2], + ['int32', (32, 32), 2], + ['int16', (32, 32), 2], + ['int8', (16, 16), 4], +]) +def test_atomic_or(param_list): + dtype, shape, ncore = param_list + block_size = shape[0] * shape[1] // ncore + split_size = shape[0] // ncore + + val = torch.randint(low=0, high=10, size=shape, dtype=eval(f'torch.{dtype}')).npu() + + pointer = torch.randint(low=0, high=10, size=(split_size, shape[1]), dtype=eval(f'torch.{dtype}')).npu() + pointer_old = torch.full_like(pointer, -10).npu() + pointer_ref = pointer.clone() + + for i in range(ncore - 1): + pointer_ref |= val[(i * split_size):((i + 1) * split_size)] + + pointer_ref_last = pointer_ref.clone() + pointer_ref |= val[((ncore - 1) * split_size):(ncore * split_size)] + + n_elements = shape[0] * shape[1] + atomic_or[ncore, 1, 1](val, pointer, pointer_old, n_elements, BLOCK_SIZE=split_size * shape[1]) + test_common.validate_cmp(dtype, pointer, pointer_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_atomic_rmw_useanalysis.py b/third_party/ascend/unittest/pytest_ut/test_atomic_rmw_useanalysis.py new file mode 100644 index 0000000000..09a1fc2498 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_atomic_rmw_useanalysis.py @@ -0,0 +1,70 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def atomic_rmw_useanalysis_kernel( + input_ptr, + output_ptr, + m_ptr, + d_ptr, + N: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + base_idx = pid * 8 + + term1 = 15.0 * 15.0 + term2 = 8.0 * (7.0 - base_idx) + + delta = term1 + term2 + sqrt_delta = tl.sqrt(delta) + + task_idx = tl.ceil((15.0 - sqrt_delta) / 2.0) + task_idx_i32 = task_idx.to(tl.int32) + + block_start = task_idx_i32 * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < N + + data = tl.load(input_ptr + offsets, mask=mask, other=0.0) + m_val = tl.load(m_ptr + offsets, mask=mask, other=0.0) + d_val = tl.load(d_ptr + offsets, mask=mask, other=0.0) + + scaled = data - m_val + p = tl.exp(scaled) + + result = p * (data * 2.0 - d_val) + + output_offsets = offsets + tl.atomic_add(output_ptr + output_offsets, result, mask=mask) + + +def test_atomic_rmw_useanalysis(): + DEVICE = "npu" + N = 1024 + BLOCK_SIZE = 128 + + torch.manual_seed(42) + input_data = torch.randn(N, dtype=torch.float32, device=DEVICE) + m_data = torch.randn(N, dtype=torch.float32, device=DEVICE) + d_data = torch.randn(N, dtype=torch.float32, device=DEVICE) + output_data = torch.zeros(N, dtype=torch.float32, device=DEVICE) + + grid = (8, ) + + atomic_rmw_useanalysis_kernel[grid]( + input_data, + output_data, + m_data, + d_data, + N=N, + BLOCK_SIZE=BLOCK_SIZE, + ) + output_sum = output_data.abs().sum().item() + + if output_sum == 0: + raise AssertionError("UseAnalysis bug detected: atomic_rmw dependencies were erased") + else: + print(" AtomicRMW UseAnalysis is working correctly.") diff --git a/third_party/ascend/unittest/pytest_ut/test_atomic_xchg.py b/third_party/ascend/unittest/pytest_ut/test_atomic_xchg.py new file mode 100644 index 0000000000..0ce72bb2e6 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_atomic_xchg.py @@ -0,0 +1,68 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import pytest +import test_common +import torch +import torch_npu + + +@triton.jit +def atomic_xchg(in_ptr0, out_ptr0, out_ptr1, n_elements, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + yindex = tl.arange(0, BLOCK_SIZE)[:] + xmask = xindex < n_elements + x0 = xindex + x1 = yindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.atomic_xchg(out_ptr0 + (x1), tmp0, xmask) + tl.store(out_ptr1 + (x0), tmp1, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['int32', (32, 32), 2], + ['int16', (16, 16), 2], + ['int8', (32, 32), 2], + ['float32', (16, 16), 2], + ['float16', (128, 128), 16], + ['float32', (32768, 16), 32], +]) +def test_atomic_xchg(param_list): + dtype, shape, ncore = param_list + block_size = shape[0] * shape[1] // ncore + split_size = shape[0] // ncore + + val = torch.randint(low=0, high=10, size=shape, dtype=eval(f'torch.{dtype}')).npu() + + pointer = torch.randint(low=0, high=10, size=(split_size, shape[1]), dtype=eval(f'torch.{dtype}')).npu() + pointer_ref = pointer.clone() + pointer_old = torch.full_like(val, -10).npu() + pointer_old_ref = pointer_old.clone() + + pointer_ref = val[((ncore - 1) * split_size):(ncore * split_size)].clone() + pointer_old_ref[0:split_size] = pointer + pointer_old_ref[split_size:((ncore - 1) * split_size)] = val[0:(ncore - 2) * split_size] + + n_elements = shape[0] * shape[1] + atomic_xchg[ncore, 1, 1](val, pointer, pointer_old, n_elements, BLOCK_SIZE=split_size * shape[1]) + test_common.validate_cmp(dtype, pointer, pointer_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_atomic_xor.py b/third_party/ascend/unittest/pytest_ut/test_atomic_xor.py new file mode 100644 index 0000000000..7616e9e0d9 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_atomic_xor.py @@ -0,0 +1,73 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import pytest +import test_common +import torch +import torch_npu + + +@triton.jit +def atomic_xor(in_ptr0, out_ptr0, out_ptr1, n_elements, BLOCK_SIZE: tl.constexpr): + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE)[:] + yindex = tl.arange(0, BLOCK_SIZE)[:] + xmask = xindex < n_elements + x0 = xindex + x1 = yindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.atomic_xor(out_ptr0 + (x1), tmp0, xmask) + tl.store(out_ptr1 + (x1), tmp1, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['int64', (32, 32), 2], + ['int32', (32, 32), 2], + ['int16', (32, 32), 7], + ['int8', (32, 32), 10], +]) +def test_atomic_xor(param_list): + dtype, shape, ncore = param_list + block_size = shape[0] * shape[1] // ncore + split_size = shape[0] // ncore + + # 初始化原始值 val，全为 0b0011（十进制 3） + val_value = 3 + val = torch.full(shape, val_value, dtype=eval(f'torch.{dtype}')).npu() + + # 每个线程使用不同输入值 x1，全为 0b0101（十进制 5） + pointer_value = 5 + pointer = torch.full((split_size, shape[1]), pointer_value, dtype=eval(f'torch.{dtype}')).npu() + pointer_old = torch.full_like(pointer, -10) + + # 原子异或后：val ^= pointer 每个线程执行一次 + # 因为异或操作具有可逆性和对称性，参考更新次数 + # 所以参考值为 val_value ^ pointer_value ^ pointer_value ^ ...（ncore 次） + pointer_result = pointer_value + for _ in range(ncore): + pointer_result ^= val_value + + pointer_ref = torch.full_like(pointer, pointer_result) + + n_elements = shape[0] * shape[1] + atomic_xor[ncore, 1, 1](val, pointer, pointer_old, n_elements, BLOCK_SIZE=split_size * shape[1]) + test_common.validate_cmp(dtype, pointer, pointer_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_attn_cp.py b/third_party/ascend/unittest/pytest_ut/test_attn_cp.py new file mode 100644 index 0000000000..6a83bd5e50 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_attn_cp.py @@ -0,0 +1,500 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import time +import torch +import torch_npu +import pytest +import sys + +sys.path.append("../../../") + +from attn_cp_triton_kernel_3d import forward_update_triton, backward_update_triton +from attn_cp_triton_kernel_3d_la import forward_update_triton as forward_update_triton_la + + +def collect_time(model, example_inputs, times: int = 1): + stream = torch.npu.current_stream() + warmup = 1 + stream.synchronize() + for i in range(times + warmup): + out = model(*example_inputs) + if i < warmup: + stream.synchronize() + t0 = time.perf_counter() + else: + t1 = time.perf_counter() + stream.synchronize() + t1 = time.perf_counter() + # GC the result after timing + assert out is not None + return (t1 - t0) / times + + +def print_performance(fn, args=(), times=10, repeat=10, baseline=1.0): + stream = torch.npu.current_stream() + + start = time.perf_counter() + + stream.synchronize() + + for _ in range(repeat * times): + fn(*args) + + stream.synchronize() + + end = time.perf_counter() + took = (end - start) / (times * repeat) + print(f"{took:.6f}") + return took + + +def profile_test(fn, fn_triton, args=(), name="gen_fn", times=10, repeat=10): + print(f"--------------------profiling {name} for {times * repeat} times--------------------") + stream = torch.npu.current_stream() + + experimental_config = torch_npu.profiler._ExperimentalConfig( + aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, + ) + prof = torch_npu.profiler.profile( + activities=[ + # torch_npu.profiler.ProfilerActivity.CPU, + torch_npu.profiler.ProfilerActivity.NPU + ], record_shapes=False, profile_memory=False, with_stack=False, + schedule=torch_npu.profiler.schedule(wait=0, warmup=1, active=100, repeat=1, skip_first=10), + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./result_dir"), + experimental_config=experimental_config) + + stream.synchronize() + prof.start() + + for _ in range(times * repeat): + fn_triton(*args) + prof.step() + + prof.stop() + + +def benchmark_test(fn, fn_triton, args=(), name="gen_fn", times=10, repeat=10, profile=False): + print(f"--------------------benchmark_{name} for {times * repeat} times--------------------") + stream = torch.npu.current_stream() + + experimental_config = torch_npu.profiler._ExperimentalConfig( + aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, + ) + prof = torch_npu.profiler.profile( + activities=[ + # torch_npu.profiler.ProfilerActivity.CPU, + torch_npu.profiler.ProfilerActivity.NPU + ], record_shapes=False, profile_memory=False, with_stack=False, + schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=100, repeat=1, skip_first=10), + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./result_dir"), + experimental_config=experimental_config) + + stream.synchronize() + prof.start() + start = time.perf_counter() + for _ in range(times * repeat): + fn_triton(*args) + if profile: + prof.step() + stream.synchronize() + end = time.perf_counter() + time_compiled = (end - start) / (times * repeat) + print(f"{time_compiled:.6f}") + + print(f"Runing eager {name} for {times * repeat} times") + start = time.perf_counter() + for _ in range(times * repeat): + fn(*args) + if profile: + prof.step() + stream.synchronize() + end = time.perf_counter() + time_eager = (end - start) / (times * repeat) + print(f"{time_eager:.6f}") + + time_eager *= 1000000 + time_compiled *= 1000000 + print( + f"Accelerated: {(time_eager - time_compiled) / time_compiled * 100:.4f}% eager takes {time_eager:.3f} us, triton takes {time_compiled:.3f} us" + ) + + return time_eager, time_compiled + + +from einops import rearrange + + +def trans_BNSD2SBH(x): + """Trans data layout from BNSD to SBH""" + return rearrange(x, 'b n s d -> s b (n d)').contiguous() + + +def broadcast_and_trans_BNSD2SBH(x, h): + """broadcast and trans a tensor from [b, n, s, 8] to [s, b, h]""" + n = x.shape[1] + d = h // n + # [b, n, s, 8] -> [b, n, s, d] + new_x = x[..., 0].unsqueeze(3) + new_x = new_x.repeat(1, 1, 1, d) + return trans_BNSD2SBH(new_x) + + +def forward_update(prev_attn_out, prev_softmax_max, prev_softmax_sum, cur_attn_out, cur_softmax_max, cur_softmax_sum): + # update softmax_max + org_dtype = prev_attn_out.dtype + softmax_max = torch.maximum(prev_softmax_max, cur_softmax_max) + prev_scale = torch.exp(prev_softmax_max - softmax_max) + cur_scale = torch.exp(cur_softmax_max - softmax_max) + # update softmax_sum + prev_softmax_sum_scaled = prev_softmax_sum * prev_scale + cur_softmax_sum_scaled = cur_softmax_sum * cur_scale + softmax_sum = prev_softmax_sum_scaled + cur_softmax_sum_scaled + # out updating scale + prev_out_scale = prev_softmax_sum_scaled / softmax_sum + cur_out_scale = cur_softmax_sum_scaled / softmax_sum + # [b, n, s, 8] -> [s, b, h] + prev_out_scale_sbh = broadcast_and_trans_BNSD2SBH(prev_out_scale, prev_attn_out.shape[-1]) + cur_out_scale_sbh = broadcast_and_trans_BNSD2SBH(cur_out_scale, prev_attn_out.shape[-1]) + # update output + attn_out = prev_attn_out * prev_out_scale_sbh + cur_attn_out * cur_out_scale_sbh + attn_out = attn_out.to(org_dtype) + return attn_out, softmax_max, softmax_sum + + +def prove_forward_update(): + + def data_validation(prev_softmax_max, cur_softmax_max, prev_softmax_sum, cur_softmax_sum, prev_attn_out, + cur_attn_out): + + (tt_attn_out, tt_softmax_max, tt_softmax_sum) = forward_update_triton(prev_attn_out, prev_softmax_max, + prev_softmax_sum, cur_attn_out, + cur_softmax_max, cur_softmax_sum) + + (attn_out, softmax_max, softmax_sum) = forward_update(prev_attn_out, prev_softmax_max, prev_softmax_sum, + cur_attn_out, cur_softmax_max, cur_softmax_sum) + + try: + assert torch.equal(softmax_max, tt_softmax_max) + print("max comparition passed.") + assert torch.equal(softmax_sum, tt_softmax_sum) + print("sum comparition passed.") + torch.testing.assert_close(attn_out, tt_attn_out) + print("atten comparition passed.") + + except Exception as e: + print(e) + print("comparison not passed") + print( + f"proving finished, attn shape:{prev_attn_out.shape}, stride:{prev_attn_out.stride(), cur_attn_out.stride()}, softmax shape:{prev_softmax_sum.shape}, stride:{prev_softmax_sum.stride(), cur_softmax_sum.stride()}" + ) + + (S, B, H, N) = (4096, 1, 6144, 48) + DS = 2 * S + DTYPE_ATTN = torch.bfloat16 + DTYPE = torch.float32 + F32_BLK_SIZE = 8 + for i in range(1): + print("prove_forward_update round:", i) + prev_attn_out = torch.randn((DS, B, H), dtype=DTYPE_ATTN).npu() + prev_softmax_max = torch.rand((B, N, DS), dtype=DTYPE).npu().unsqueeze(3).repeat(1, 1, 1, F32_BLK_SIZE) + prev_softmax_sum = torch.rand((B, N, DS), dtype=DTYPE).npu().unsqueeze(3).repeat(1, 1, 1, F32_BLK_SIZE) + cur_attn_out = torch.randn((DS, B, H), dtype=DTYPE_ATTN).npu() + cur_softmax_max = torch.rand((B, N, DS), dtype=DTYPE).npu().unsqueeze(3).repeat(1, 1, 1, F32_BLK_SIZE) + cur_softmax_sum = torch.rand((B, N, DS), dtype=DTYPE).npu().unsqueeze(3).repeat(1, 1, 1, F32_BLK_SIZE) + + prev_attn_out_s = prev_attn_out.view(2, S, B, H)[1] + prev_softmax_max_s = prev_softmax_max.view(B, N, 2, S, 8)[:, :, 1, :, :] + prev_softmax_sum_s = prev_softmax_sum.view(B, N, 2, S, 8)[:, :, 1, :, :] + cur_attn_out_s = torch.randn((S, B, H), dtype=DTYPE_ATTN).npu() + cur_softmax_max_s = torch.rand((B, N, S), dtype=DTYPE).npu().unsqueeze(3).repeat(1, 1, 1, F32_BLK_SIZE) + cur_softmax_sum_s = torch.rand((B, N, S), dtype=DTYPE).npu().unsqueeze(3).repeat(1, 1, 1, F32_BLK_SIZE) + + print("--------------------prove_forward_update:2S----------------------") + data_validation(prev_softmax_max, cur_softmax_max, prev_softmax_sum, cur_softmax_sum, prev_attn_out, + cur_attn_out) + + print("--------------------prove_forward_update:1S---------------------- ") + data_validation(prev_softmax_max_s, cur_softmax_max_s, prev_softmax_sum_s, cur_softmax_sum_s, prev_attn_out_s, + cur_attn_out_s) + + +def benchmark_forward_update(): + (S, B, H, N) = (4096, 1, 6144, 48) + DS = 2 * S + DTYPE_ATTN = torch.bfloat16 + DTYPE = torch.float32 + F32_BLK_SIZE = 8 + prev_attn_out = torch.randn((DS, B, H), dtype=DTYPE_ATTN).npu() + prev_softmax_max = torch.rand((B, N, DS), dtype=DTYPE).npu().unsqueeze(3).repeat(1, 1, 1, F32_BLK_SIZE) + prev_softmax_sum = torch.rand((B, N, DS), dtype=DTYPE).npu().unsqueeze(3).repeat(1, 1, 1, F32_BLK_SIZE) + cur_attn_out = torch.randn((DS, B, H), dtype=DTYPE_ATTN).npu() + cur_softmax_max = torch.rand((B, N, DS), dtype=DTYPE).npu().unsqueeze(3).repeat(1, 1, 1, F32_BLK_SIZE) + cur_softmax_sum = torch.rand((B, N, DS), dtype=DTYPE).npu().unsqueeze(3).repeat(1, 1, 1, F32_BLK_SIZE) + + prev_attn_out_s = prev_attn_out.view(2, S, B, H)[1] + prev_softmax_max_s = prev_softmax_max.view(B, N, 2, S, 8)[:, :, 1, :, :] + prev_softmax_sum_s = prev_softmax_sum.view(B, N, 2, S, 8)[:, :, 1, :, :] + cur_attn_out_s = torch.randn((S, B, H), dtype=DTYPE_ATTN).npu() + cur_softmax_max_s = torch.rand((B, N, S), dtype=DTYPE).npu().unsqueeze(3).repeat(1, 1, 1, F32_BLK_SIZE) + cur_softmax_sum_s = torch.rand((B, N, S), dtype=DTYPE).npu().unsqueeze(3).repeat(1, 1, 1, F32_BLK_SIZE) + + benchmark_test( + forward_update, forward_update_triton, + args=(prev_attn_out, prev_softmax_max, prev_softmax_sum, cur_attn_out, cur_softmax_max, cur_softmax_sum), + name="forward_update_2s", times=10, repeat=10) + + benchmark_test( + forward_update, forward_update_triton, + args=(prev_attn_out, prev_softmax_max, prev_softmax_sum, cur_attn_out, cur_softmax_max, cur_softmax_sum), + name="forward_update_s", times=10, repeat=10) + + +def profile_forward_update(): + (S, B, H, N) = (4096, 1, 6144, 48) + # (S, B, H, N) = (2048,1,1536,12) + DS = 2 * S + DTYPE_ATTN = torch.bfloat16 + DTYPE = torch.float32 + F32_BLK_SIZE = 8 + prev_attn_out = torch.randn((DS, B, H), dtype=DTYPE_ATTN).npu() + prev_softmax_max = torch.rand((B, N, DS), dtype=DTYPE).npu().unsqueeze(3).repeat(1, 1, 1, F32_BLK_SIZE) + prev_softmax_sum = torch.rand((B, N, DS), dtype=DTYPE).npu().unsqueeze(3).repeat(1, 1, 1, F32_BLK_SIZE) + cur_attn_out = torch.randn((DS, B, H), dtype=DTYPE_ATTN).npu() + cur_softmax_max = torch.rand((B, N, DS), dtype=DTYPE).npu().unsqueeze(3).repeat(1, 1, 1, F32_BLK_SIZE) + cur_softmax_sum = torch.rand((B, N, DS), dtype=DTYPE).npu().unsqueeze(3).repeat(1, 1, 1, F32_BLK_SIZE) + + prev_attn_out_s = prev_attn_out.view(2, S, B, H)[1] + prev_softmax_max_s = prev_softmax_max.view(B, N, 2, S, 8)[:, :, 1, :, :] + prev_softmax_sum_s = prev_softmax_sum.view(B, N, 2, S, 8)[:, :, 1, :, :] + cur_attn_out_s = torch.randn((S, B, H), dtype=DTYPE_ATTN).npu() + cur_softmax_max_s = torch.rand((B, N, S), dtype=DTYPE).npu().unsqueeze(3).repeat(1, 1, 1, F32_BLK_SIZE) + cur_softmax_sum_s = torch.rand((B, N, S), dtype=DTYPE).npu().unsqueeze(3).repeat(1, 1, 1, F32_BLK_SIZE) + + profile_test( + forward_update, forward_update_triton, + args=(prev_attn_out, prev_softmax_max, prev_softmax_sum, cur_attn_out, cur_softmax_max, cur_softmax_sum), + name="forward_update_2s", times=10, repeat=10) + + profile_test( + forward_update, forward_update_triton, + args=(prev_attn_out, prev_softmax_max, prev_softmax_sum, cur_attn_out, cur_softmax_max, cur_softmax_sum), + name="forward_update_s", times=10, repeat=10) + + +def backward_update(dq, dk, dv, cur_dq, cur_dk, cur_dv, i=7, rank=1): + cp_size = 8 + if i >= cp_size - rank - 1: + if i == cp_size - 1: + cur_dq = cur_dq.view(dq.shape) + cur_dk = cur_dk.view(dk.shape) + cur_dv = cur_dv.view(dv.shape) + dq.add_(cur_dq) + dk.add_(cur_dk) + dv.add_(cur_dv) + else: + cur_dq = cur_dq.view(dq.shape) + dq.add_(cur_dq) + dk[0].add_(cur_dk) + dv[0].add_(cur_dv) + else: + dq[1].add_(cur_dq) + cur_dk = cur_dk.view(dk.shape) # [2s, b, h] -> [2, s, b, h] + cur_dv = cur_dv.view(dv.shape) + dk.add_(cur_dk) + dv.add_(cur_dv) + + +def prove_backward_update(): + (S, B, H) = (16384, 1, 1536) + DS = 2 * S + DTYPE_ATTN = torch.bfloat16 + DTYPE = torch.float32 + + d_dq = torch.randn((2, S, B, H), dtype=DTYPE_ATTN).npu() + cur_dq = torch.randn((S, B, H), dtype=DTYPE_ATTN).npu() + d_cur_dq = torch.randn((DS, B, H), dtype=DTYPE_ATTN).npu() + + d_dk = torch.randn((2, S, B, H), dtype=DTYPE_ATTN).npu() + cur_dk = torch.randn((S, B, H), dtype=DTYPE_ATTN).npu() + d_cur_dk = torch.randn((DS, B, H), dtype=DTYPE_ATTN).npu() + + d_dv = torch.randn((2, S, B, H), dtype=DTYPE_ATTN).npu() + cur_dv = torch.randn((S, B, H), dtype=DTYPE_ATTN).npu() + d_cur_dv = torch.randn((DS, B, H), dtype=DTYPE_ATTN).npu() + + def data_validate(dq, dk, dv, cur_dq, cur_dk, cur_dv, i, rank): + dq_c = dq.detach().clone() + dk_c = dk.detach().clone() + dv_c = dv.detach().clone() + + if (i == 7): + backward_update(dq, dk, dv, cur_dq, cur_dk, cur_dv, i, rank) + backward_update_triton(dq_c, dk_c, dv_c, cur_dq, cur_dk, cur_dv) + elif (i == 6): + backward_update(dq, dk, dv, cur_dq, cur_dk, cur_dv, i, rank) + backward_update_triton(dq_c, dk_c[0], dv_c[0], cur_dq, cur_dk, cur_dv) + else: + backward_update(dq, dk, dv, cur_dq, cur_dk, cur_dv, i, rank) + backward_update_triton(dq_c[1], dk_c, dv_c, cur_dq, cur_dk, cur_dv) + + torch.testing.assert_close(dq, dq_c) + print("dq comparison passed") + torch.testing.assert_close(dk, dk_c) + print("dk comparison passed") + torch.testing.assert_close(dv, dv_c) + print("passed comparison ") + + print("--------------------prove_backward_update case 0----------------------") + data_validate(d_dq, d_dk, d_dv, d_cur_dq, d_cur_dk, d_cur_dv, 7, 1) + print("--------------------prove_backward_update case 1----------------------") + data_validate(d_dq, d_dk, d_dv, d_cur_dq, cur_dk, cur_dv, 6, 1) + print("--------------------prove_backward_update case 2----------------------") + data_validate(d_dq, d_dk, d_dv, cur_dq, d_cur_dk, d_cur_dv, 5, 1) + + +def benchmark_backward_update(): + (S, B, H) = (16384, 1, 1536) + DS = 2 * S + DTYPE_ATTN = torch.bfloat16 + DTYPE = torch.float32 + + d_dq = torch.randn((2, S, B, H), dtype=DTYPE_ATTN).npu() + cur_dq = torch.randn((S, B, H), dtype=DTYPE_ATTN).npu() + d_cur_dq = torch.randn((DS, B, H), dtype=DTYPE_ATTN).npu() + + d_dk = torch.randn((2, S, B, H), dtype=DTYPE_ATTN).npu() + cur_dk = torch.randn((S, B, H), dtype=DTYPE_ATTN).npu() + d_cur_dk = torch.randn((DS, B, H), dtype=DTYPE_ATTN).npu() + + d_dv = torch.randn((2, S, B, H), dtype=DTYPE_ATTN).npu() + cur_dv = torch.randn((S, B, H), dtype=DTYPE_ATTN).npu() + d_cur_dv = torch.randn((DS, B, H), dtype=DTYPE_ATTN).npu() + # benchmark case 0 + benchmark_test(backward_update, backward_update_triton, args=(d_dq, d_dk, d_dv, d_cur_dq, d_cur_dk, d_cur_dv), + name="backward_update_0") + + +def forward_update_la(prev_attn_out, prev_softmax_log_max_sum, cur_attn_out, cur_softmax_log_max_sum): + if prev_attn_out is None: + return cur_attn_out, cur_softmax_log_max_sum + softmax_log_max_sum = torch.log(torch.exp(cur_softmax_log_max_sum) + torch.exp(prev_softmax_log_max_sum)) + attn_out = torch.exp(prev_softmax_log_max_sum - softmax_log_max_sum) * prev_attn_out + torch.exp( + cur_softmax_log_max_sum - softmax_log_max_sum) * cur_attn_out + return attn_out, softmax_log_max_sum + + +# simuldate origin code :call foward_update then call copy +def forward_update_copy(prev_attn_out, prev_softmax_log_max_sum, cur_attn_out, cur_softmax_log_max_sum): + attn_out, softmax = forward_update_la(prev_attn_out, prev_softmax_log_max_sum, cur_attn_out, + cur_softmax_log_max_sum) + prev_attn_out.copy_(attn_out) + prev_softmax_log_max_sum.copy_(softmax) + + +def prove_forward_update_la(): + + def data_validation(prev_attn_out, prev_softmax_sum, cur_attn_out, cur_softmax_sum): + + (attn_out, softmax_sum) = forward_update_la(prev_attn_out, prev_softmax_sum, cur_attn_out, cur_softmax_sum) + forward_update_triton_la(prev_attn_out, prev_softmax_sum, cur_attn_out, cur_softmax_sum) + + try: + torch.testing.assert_close(softmax_sum, prev_softmax_sum) + print("softmax comparition passed.") + torch.testing.assert_close(attn_out, prev_attn_out) + print("attn comparition passed.") + + except Exception as e: + print(e) + print("comparison not passed") + + print( + f"proving finished, attn shape:{prev_attn_out.shape}, stride:{prev_attn_out.stride(), cur_attn_out.stride()}, softmax shape:{prev_softmax_sum.shape}, stride:{prev_softmax_sum.stride(), cur_softmax_sum.stride()}" + ) + + (S, B, N, D) = (4096, 1, 48, 128) + DS = 2 * S + DTYPE_ATTN = torch.float32 + DTYPE = torch.float32 + + for i in range(1): + print("round:", i) + attn_out = torch.randn((B, N, DS, D), dtype=DTYPE_ATTN).npu() + softmax_sum = torch.rand((B, N, DS, 1), dtype=DTYPE).npu() + + cur_attn_out = torch.randn((B, N, DS, D), dtype=DTYPE_ATTN).npu() + cur_softmax_sum = torch.rand((B, N, DS, 1), dtype=DTYPE).npu() + + cur_attn_out_s = torch.randn((B, N, S, D), dtype=DTYPE_ATTN).npu() + cur_softmax_sum_s = torch.rand((B, N, S, 1), dtype=DTYPE).npu() + + print("---------------------prove_forward_updat_la 2S-------------------------------") + data_validation(attn_out, softmax_sum, cur_attn_out, cur_softmax_sum) + + print("--------------------prove_forward_update_la 1S------------------------------") + # [b, n, 2s, d] -> [b, n, 2, s, d] + attn_out_s = attn_out.view(*attn_out.shape[:2], 2, attn_out.shape[2] // 2, attn_out.shape[-1]) + # [b, n, 2s, 1] -> [b, n, 2, s, 1] + softmax_sum_s = softmax_sum.view(*softmax_sum.shape[:2], 2, softmax_sum.shape[2] // 2, softmax_sum.shape[-1]) + + data_validation(attn_out_s[:, :, 1], softmax_sum_s[:, :, 1], cur_attn_out_s, cur_softmax_sum_s) + + +def benchmark_forward_update_la(): + (S, B, N, D) = (4096, 1, 48, 128) + DS = 2 * S + DTYPE_ATTN = torch.float32 + DTYPE = torch.float32 + + attn_out = torch.randn((B, N, DS, D), dtype=DTYPE_ATTN).npu() + softmax_sum = torch.rand((B, N, DS, 1), dtype=DTYPE).npu() + cur_attn_out = torch.randn((B, N, DS, D), dtype=DTYPE_ATTN).npu() + cur_softmax_sum = torch.rand((B, N, DS, 1), dtype=DTYPE).npu() + cur_attn_out_s = torch.randn((B, N, S, D), dtype=DTYPE_ATTN).npu() + cur_softmax_sum_s = torch.rand((B, N, S, 1), dtype=DTYPE).npu() + attn_out_s = attn_out.view(*attn_out.shape[:2], 2, attn_out.shape[2] // 2, attn_out.shape[-1]) + softmax_sum_s = softmax_sum.view(*softmax_sum.shape[:2], 2, softmax_sum.shape[2] // 2, softmax_sum.shape[-1]) + + benchmark_test(forward_update_copy, forward_update_triton_la, + args=(attn_out_s[:, :, 1], softmax_sum_s[:, :, 1], cur_attn_out_s, cur_softmax_sum_s), + name="forward_update_la", profile=False, repeat=1000) + + +@pytest.mark.skip(reason="attn_cp") +def test_prove_forward_update(): + prove_forward_update() + + +@pytest.mark.skip(reason="attn_cp") +def test_prove_forward_update_la(): + prove_forward_update_la() + + +@pytest.mark.skip(reason="attn_cp") +def test_prove_backward_update(): + prove_backward_update() + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/third_party/ascend/unittest/pytest_ut/test_bind_buffer.py b/third_party/ascend/unittest/pytest_ut/test_bind_buffer.py new file mode 100644 index 0000000000..a8e59d0767 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_bind_buffer.py @@ -0,0 +1,67 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +#!/usr/bin/env python3 +import os + +import triton +import triton.language as tl +import triton.extension.buffer.language as bl +import triton.language.extra.cann.extension as al +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir + +os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" + + +class Options: + num_warps = 4 + num_stages = 3 + num_ctas = 1 + cluster_dims = (1, 1, 1) + enable_fp_fusion = True + debug = False + + +def compile_kernel(kernel, signature, constants): + """Helper to compile a kernel to MLIR.""" + src = ASTSource(kernel, signature, constants) + context = ir.context() + ir.load_dialects(context) + ascend_ir.load_dialects(context) + module = ast_to_ttir(kernel, src, context, Options(), {}, {}) + return str(module) + + +@triton.jit +def bind_buffer(): + alloc = bl.alloc(tl.float32, [32, 32], al.ascend_address_space.UB) + tensor = tl.full((32, 32), 0, dtype=tl.float32) + bl.to_buffer(tensor, bind_buffer=alloc) + + +# ============== Main for manual testing ============== + +if __name__ == "__main__": + mlir = compile_kernel(bind_buffer, {}, {}) + assert len(mlir) > 0 + print(mlir) diff --git a/third_party/ascend/unittest/pytest_ut/test_block_ptr.py b/third_party/ascend/unittest/pytest_ut/test_block_ptr.py new file mode 100644 index 0000000000..e8132a4508 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_block_ptr.py @@ -0,0 +1,195 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, output_ptr1, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + block_ptr_in = tl.make_block_ptr( + base=x_ptr, + shape=(XB, YB, ZB), + strides=(YB * ZB, ZB, 1), + offsets=(0, 0, 0), + block_shape=(XB, YB, ZB), + order=(2, 1, 0), + ) + X = tl.load(block_ptr_in) + + oidx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + block_ptr_out = tl.make_block_ptr( + base=output_ptr, + shape=(XB, YB, ZB), + strides=(YB * ZB, ZB, 1), + offsets=(0, 0, 0), + block_shape=(XB, YB, ZB), + order=(2, 1, 0), + ) + tl.store(block_ptr_out, X) + + +paras = [ + ('*fp32', eval('torch.float32'), 2, 256, 16), + ('*fp32', eval('torch.float32'), 8, 8, 4), + ('*fp16', eval('torch.float16'), 2, 256, 16), + ('*fp16', eval('torch.float16'), 8, 8, 4), + ('*i8', eval('torch.int8'), 2, 256, 16), + ('*i8', eval('torch.int8'), 8, 8, 4), +] + + +@pytest.mark.parametrize('para_type,data_type,XB,YB,ZB', paras) +def test_npu(para_type, data_type, XB, YB, ZB): + + x = torch.randint(low=-128, high=128, size=(XB, YB, ZB), dtype=data_type).npu() + y = torch.randint(low=-128, high=128, size=(XB, YB, ZB), dtype=data_type).npu() + z = torch.randint(low=-128, high=128, size=(XB, YB, ZB), dtype=data_type).npu() + + print(f"shape = {x.shape}") + print(x.dtype) + + output = torch.randint(1, (XB, YB, ZB), dtype=data_type).npu() + output1 = output + print(f"output.dtype={output.dtype}") + + a = x + print(a) + fn_npu_[1, 1, 1](output, x, y, z, output1, XB=XB, YB=YB, ZB=ZB, debug=True) + print(output) + torch.testing.assert_close(output, a) + + +@triton.jit +def dma_block_ptr( + input_ptr, + output_ptr, + scale_ptr, + batch_size, + cu_seqlens_ptr, + stride_i_m, + stride_i_n, + stride_o_m, + stride_o_n, + stride_s_b, + HEAD_DIM, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, +): + n_progs = tl.num_programs(0) + pid = tl.program_id(0) + + cu_num_blocks = 0 + for bid in range(batch_size): + start_loc = tl.load(cu_seqlens_ptr + bid) + end_loc = tl.load(cu_seqlens_ptr + bid + 1) + scale = tl.load(scale_ptr + bid * stride_s_b) + + len_loc = end_loc - start_loc + prev_num_blocks = cu_num_blocks + new_num_blocks = tl.cdiv(len_loc, BLOCK_SIZE_M).to(tl.int32) + i_block_ptr_bbase = tl.make_block_ptr( + input_ptr + start_loc * stride_i_m, + shape=(len_loc, HEAD_DIM), + strides=(stride_i_m, stride_i_n), + offsets=(0, 0), + block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N), + order=(1, 0), + ) + o_block_ptr_bbase = tl.make_block_ptr( + output_ptr + start_loc * stride_o_m, + shape=(len_loc, HEAD_DIM), + strides=(stride_o_m, stride_o_n), + offsets=(0, 0), + block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N), + order=(1, 0), + ) + cu_num_blocks += new_num_blocks + for m_id in range((prev_num_blocks + pid) % n_progs, new_num_blocks, n_progs): + i_block_ptr = tl.advance(i_block_ptr_bbase, (m_id * BLOCK_SIZE_M, 0)) + o_block_ptr = tl.advance(o_block_ptr_bbase, (m_id * BLOCK_SIZE_M, 0)) + i_tile = tl.load(i_block_ptr, boundary_check=[0, 1], padding_option="zero") + o_tile = i_tile.to(tl.float32) * scale + tl.store(o_block_ptr, o_tile.to(i_tile.dtype), boundary_check=[0, 1]) + + +def ref_func(inputs, scale, cu_lens): + outputs = torch.zeros_like(inputs) + bsz = cu_lens.size(0) - 1 + for bid in range(bsz): + tmp = inputs[cu_lens[bid]:cu_lens[bid + 1]].to(torch.float32) * scale[bid] + outputs[cu_lens[bid]:cu_lens[bid + 1]] = tmp.to(outputs.dtype) + return outputs + + +def tt_func(inputs, scale, cu_lens): + bsz = cu_lens.size(0) - 1 + outputs = torch.zeros_like(inputs) + head_dim = inputs.size(-1) + assert head_dim <= 1024 + BLOCK_SIZE_N = 1024 + BLOCK_SIZE_M = 4 + dma_block_ptr[ + 20, + ]( + inputs, + outputs, + scale, + bsz, + cu_lens, + inputs.stride(0), + inputs.stride(1), + outputs.stride(0), + outputs.stride(1), + scale.stride(0), + head_dim, + BLOCK_SIZE_M=BLOCK_SIZE_M, + BLOCK_SIZE_N=BLOCK_SIZE_N, + ) + return outputs + + +@pytest.mark.parametrize('param_list', [ + [8, 1024, 1024, True], + [8, 1024, 1024, False], +]) +def test_func(param_list): + bsz, max_len, max_n, test_align = param_list + lens = torch.randint(max_len // 2, max_len, (bsz, ), dtype=torch.int32, device="npu") + n = torch.randint(max_n // 2, max_n, (1, ), dtype=torch.int32, device="npu")[0].item() + if test_align: + lens = (lens + 1023) // 1024 * 1024 + n = (n + 1023) // 1024 * 1024 + cu_lens = torch.cumsum(lens, dim=0) + cu_lens = torch.cat([torch.zeros(1, dtype=torch.int32, device="npu"), cu_lens], dim=0) + inputs = torch.randn(cu_lens[-1], n, dtype=torch.float16, device="npu") + scale = torch.randn(bsz, dtype=torch.float32, device="npu") + ref_output = ref_func(inputs, scale, cu_lens) + tt_output = tt_func(inputs, scale, cu_lens) + torch.testing.assert_close(ref_output, tt_output) diff --git a/third_party/ascend/unittest/pytest_ut/test_boundary_check.py b/third_party/ascend/unittest/pytest_ut/test_boundary_check.py new file mode 100644 index 0000000000..c7dabda498 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_boundary_check.py @@ -0,0 +1,275 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import triton +import triton.language as tl +import pytest + + +# ========== Test 1: Static base address + boundary_check ========== +@triton.jit +def static_base_boundary_check_kernel( + out_ptr, + in_ptr, + BLOCK_SIZE: tl.constexpr, +): + ptr = tl.make_block_ptr(base=in_ptr, shape=(BLOCK_SIZE * 2, ), strides=(1, ), offsets=(0, ), + block_shape=(BLOCK_SIZE, ), order=(0, )) + data = tl.load(ptr, boundary_check=(0, ), padding_option="zero") + result = tl.sum(data) + tl.store(out_ptr, result) + + +def ref_static_base(in_tensor, BLOCK_SIZE): + return in_tensor[:BLOCK_SIZE].sum().item() + + +def test_static_base(): + BLOCK_SIZE = 64 + in_tensor = torch.randn(BLOCK_SIZE * 2, dtype=torch.float32).npu() + out_tensor = torch.zeros(1, dtype=torch.float32).npu() + static_base_boundary_check_kernel[(1, )]( + out_ptr=out_tensor, + in_ptr=in_tensor, + BLOCK_SIZE=BLOCK_SIZE, + ) + expected = ref_static_base(in_tensor.cpu(), BLOCK_SIZE) + assert torch.allclose(out_tensor.cpu(), torch.tensor(expected, device='cpu'), atol=1e-4) + + +# ========== Test 2: Simple dynamic base address + boundary_check ========== +@triton.jit +def simple_dynamic_base_boundary_check_kernel( + out_ptr, + in_ptr, + offset: tl.int32, + BLOCK_SIZE: tl.constexpr, +): + base = in_ptr + offset + ptr = tl.make_block_ptr(base=base, shape=(BLOCK_SIZE * 2, ), strides=(1, ), offsets=(0, ), + block_shape=(BLOCK_SIZE, ), order=(0, )) + data = tl.load(ptr, boundary_check=(0, ), padding_option="zero") + result = tl.sum(data) + tl.store(out_ptr, result) + + +def test_simple_dynamic_base(): + BLOCK_SIZE = 64 + offset = 32 + in_tensor = torch.randn(BLOCK_SIZE * 4, dtype=torch.float32).npu() + out_tensor = torch.zeros(1, dtype=torch.float32).npu() + simple_dynamic_base_boundary_check_kernel[(1, )]( + out_ptr=out_tensor, + in_ptr=in_tensor, + offset=offset, + BLOCK_SIZE=BLOCK_SIZE, + ) + expected = in_tensor.cpu()[offset:offset + BLOCK_SIZE].sum().item() + assert torch.allclose(out_tensor.cpu(), torch.tensor(expected, device='cpu'), atol=1e-4) + + +# ========== Test 3: Nested loop + dynamic base address + advance + boundary_check ========== +@triton.jit +def nested_dynamic_advance_boundary_kernel( + out_ptr, + in_ptr, + stride_in: tl.int32, + OUTER_LOOP: tl.constexpr, + INNER_LOOP: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """ + Smallest reproducible code: The dynamic base address is in the outer loop, + and tl.advance is in the inner loop, where there is a boundary_check. + """ + for i in range(OUTER_LOOP): + base = in_ptr + i * stride_in + ptr = tl.make_block_ptr(base=base, shape=(INNER_LOOP * BLOCK_SIZE, ), strides=(1, ), offsets=(0, ), + block_shape=(BLOCK_SIZE, ), order=(0, )) + for j in range(INNER_LOOP): + cur_ptr = tl.advance(ptr, (j * BLOCK_SIZE, )) + data = tl.load(cur_ptr, boundary_check=(0, ), padding_option="zero") + result = tl.sum(data) + tl.store(out_ptr + i * INNER_LOOP + j, result) + + +def ref_nested_dynamic(in_tensor, OUTER_LOOP, INNER_LOOP, BLOCK_SIZE): + """ + PyTorch equivalent implementation: + - Treat in_tensor as a tensor of shape [OUTER_LOOP, INNER_LOOP * BLOCK_SIZE] + - For each (i, j) block: take the BLOCK_SIZE elements starting from j*BLOCK_SIZE in the i-th row and sum them up. + - Note: There is boundary_check + zero padding, but there is no out-of-bound access in this case, so no special handling is needed. + """ + reshaped = in_tensor[:OUTER_LOOP * INNER_LOOP * BLOCK_SIZE].view(OUTER_LOOP, INNER_LOOP * BLOCK_SIZE) + blocks = reshaped.unfold(1, BLOCK_SIZE, BLOCK_SIZE) + return blocks.sum(dim=-1).flatten() + + +def test_nested_dynamic(): + BLOCK_SIZE = 8 + OUTER_LOOP = 2 + INNER_LOOP = 2 + in_tensor = torch.randn(OUTER_LOOP * INNER_LOOP * BLOCK_SIZE * 2, dtype=torch.float32).npu() + out_tensor = torch.zeros(OUTER_LOOP * INNER_LOOP, dtype=torch.float32).npu() + nested_dynamic_advance_boundary_kernel[(1, )]( + out_ptr=out_tensor, + in_ptr=in_tensor, + stride_in=INNER_LOOP * BLOCK_SIZE, + OUTER_LOOP=OUTER_LOOP, + INNER_LOOP=INNER_LOOP, + BLOCK_SIZE=BLOCK_SIZE, + ) + ref = ref_nested_dynamic(in_tensor.cpu(), OUTER_LOOP, INNER_LOOP, BLOCK_SIZE) + assert torch.allclose(out_tensor.cpu(), ref, atol=1e-4) + + +# ========== Test 4: Explicit out-of-bounds access + zero padding + boundary_check ========== +@triton.jit +def out_of_bound_zero_padding_kernel( + out_ptr, + in_ptr, + BLOCK_SIZE: tl.constexpr, +): + ptr = tl.make_block_ptr(base=in_ptr, shape=(BLOCK_SIZE, ), strides=(1, ), offsets=(0, ), + block_shape=(BLOCK_SIZE * 2, ), order=(0, )) + data = tl.load(ptr, boundary_check=(0, ), padding_option="zero") + result = tl.sum(data) + tl.store(out_ptr, result) + + +def test_out_of_bound(): + BLOCK_SIZE = 64 + in_tensor = torch.randn(BLOCK_SIZE, dtype=torch.float32).npu() + out_tensor = torch.zeros(1, dtype=torch.float32).npu() + out_of_bound_zero_padding_kernel[(1, )]( + out_ptr=out_tensor, + in_ptr=in_tensor, + BLOCK_SIZE=BLOCK_SIZE, + ) + expected = in_tensor.cpu().sum().item() + assert torch.allclose(out_tensor.cpu(), torch.tensor(expected, device='cpu'), atol=1e-4) + + +# ========== Test 5：padding_option = NAN + boundary_check========== +@triton.jit +def nan_padding_kernel( + out_ptr, + in_ptr, + BLOCK_SIZE: tl.constexpr, +): + ptr = tl.make_block_ptr(base=in_ptr, shape=(BLOCK_SIZE, ), strides=(1, ), offsets=(0, ), + block_shape=(BLOCK_SIZE * 2, ), order=(0, )) + data = tl.load(ptr, boundary_check=(0, ), padding_option="nan") + result = tl.sum(data) + tl.store(out_ptr, result) + + +def test_nan_padding(): + BLOCK_SIZE = 64 + in_tensor = torch.randn(BLOCK_SIZE, dtype=torch.float32).npu() + out_tensor = torch.zeros(1, dtype=torch.float32).npu() + try: + nan_padding_kernel[(1, )]( + out_ptr=out_tensor, + in_ptr=in_tensor, + BLOCK_SIZE=BLOCK_SIZE, + ) + assert torch.isnan(out_tensor.cpu()).any() + except Exception as e: + print(f"Warning: NAN padding test may not be supported: {e}") + + +# ========== Test 6：Multi-layer advance + boundary_check ========== +@triton.jit +def multi_advance_kernel( + out_ptr, + in_ptr, + BLOCK_SIZE: tl.constexpr, +): + base = in_ptr + ptr0 = tl.make_block_ptr(base=base, shape=(BLOCK_SIZE * 4, ), strides=(1, ), offsets=(0, ), + block_shape=(BLOCK_SIZE, ), order=(0, )) + ptr1 = tl.advance(ptr0, (BLOCK_SIZE, )) + ptr2 = tl.advance(ptr1, (BLOCK_SIZE, )) + data = tl.load(ptr2, boundary_check=(0, ), padding_option="zero") + result = tl.sum(data) + tl.store(out_ptr, result) + + +def test_multi_advance(): + BLOCK_SIZE = 64 + in_tensor = torch.randn(BLOCK_SIZE * 4, dtype=torch.float32).npu() + out_tensor = torch.zeros(1, dtype=torch.float32).npu() + multi_advance_kernel[(1, )]( + out_ptr=out_tensor, + in_ptr=in_tensor, + BLOCK_SIZE=BLOCK_SIZE, + ) + expected = in_tensor.cpu()[2 * BLOCK_SIZE:3 * BLOCK_SIZE].sum().item() + assert torch.allclose(out_tensor.cpu(), torch.tensor(expected, device='cpu'), atol=1e-4) + + +# ========== Test 7：Complex base address calculation + boundary_check ========== +@triton.jit +def complex_base_calculation_kernel( + out_ptr, + in_ptr, + offset1: tl.int32, + offset2: tl.int32, + scale: tl.int32, + BLOCK_SIZE: tl.constexpr, +): + base = in_ptr + offset1 * scale + offset2 + ptr = tl.make_block_ptr(base=base, shape=(BLOCK_SIZE * 2, ), strides=(1, ), offsets=(0, ), + block_shape=(BLOCK_SIZE, ), order=(0, )) + data = tl.load(ptr, boundary_check=(0, ), padding_option="zero") + result = tl.sum(data) + tl.store(out_ptr, result) + + +def test_complex_base(): + BLOCK_SIZE = 64 + offset1, offset2, scale = 2, 16, 32 + total_offset = offset1 * scale + offset2 + in_tensor = torch.randn(total_offset + BLOCK_SIZE * 2, dtype=torch.float32).npu() + out_tensor = torch.zeros(1, dtype=torch.float32).npu() + complex_base_calculation_kernel[(1, )]( + out_ptr=out_tensor, + in_ptr=in_tensor, + offset1=offset1, + offset2=offset2, + scale=scale, + BLOCK_SIZE=BLOCK_SIZE, + ) + expected = in_tensor.cpu()[total_offset:total_offset + BLOCK_SIZE].sum().item() + assert torch.allclose(out_tensor.cpu(), torch.tensor(expected, device='cpu'), atol=1e-4) + + +if __name__ == "__main__": + print("Running all boundary_check tests...") + test_static_base() + test_simple_dynamic_base() + test_nested_dynamic() + test_out_of_bound() + test_nan_padding() + test_multi_advance() + test_complex_base() + print("All tests completed successfully!") diff --git a/third_party/ascend/unittest/pytest_ut/test_broadcast.py b/third_party/ascend/unittest/pytest_ut/test_broadcast.py new file mode 100644 index 0000000000..10b17adc52 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_broadcast.py @@ -0,0 +1,57 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu + +import pytest + +import test_common + +NBLOCKS = 1 +XS = tl.constexpr(2) +YS = tl.constexpr(4) +ZS = tl.constexpr(8) +NUMEL = tl.constexpr(XS * ZS) + + +@triton.jit +def fn_broadcast(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = tl.arange(0, 1)[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x = tl.load(in_ptr0 + idx) + x1 = tl.load(out_ptr0 + odx) + ret, _ = tl.broadcast(x, x1) + tl.store(out_ptr0 + odx, ret) + + +@pytest.mark.parametrize('dtype', ["bfloat16"]) +def test_broadcast_alltype(dtype): + input = test_common.generate_tensor((1, YS, ZS), dtype).npu() + ans = input.repeat(XS, 1, 1) + output = torch.zeros((XS, YS, ZS), dtype=eval('torch.' + dtype)).npu() + fn_broadcast[1, 1, 1](input, output, XS, YS, ZS) + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/pytest_ut/test_broadcast_op.py b/third_party/ascend/unittest/pytest_ut/test_broadcast_op.py new file mode 100644 index 0000000000..924964e5f7 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_broadcast_op.py @@ -0,0 +1,54 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu + +import pytest + +import test_common + +NBLOCKS = 1 +XS = tl.constexpr(2) +YS = tl.constexpr(4) +ZS = tl.constexpr(8) +NUMEL = tl.constexpr(XS * ZS) + + +@triton.jit +def fn_broadcast_to(output_ptr, input_ptr, length): + col_offsets = tl.arange(0, NUMEL) + input = tl.load(input_ptr + col_offsets) + result = input.reshape((XS, 1, ZS)).broadcast_to((XS, YS, ZS)).reshape((XS * YS * ZS)) + brc_col_offsets = tl.arange(0, NUMEL * YS) + tl.store(output_ptr + brc_col_offsets, result) + + +@pytest.mark.parametrize('dtype', ["float32"]) +def test_broadcast_to(dtype): + length = NUMEL + dtype = eval(f"torch.{dtype}") + x = torch.randn((XS, 1, ZS), dtype=dtype).npu() + output = torch.randn((XS, YS, ZS), dtype=dtype).npu() + fn_broadcast_to[NBLOCKS, 1, 1](output, x, length, debug=True) + assert (torch.equal(output, x.repeat(1, YS, 1))) diff --git a/third_party/ascend/unittest/pytest_ut/test_cannonicalize_tl_where.py b/third_party/ascend/unittest/pytest_ut/test_cannonicalize_tl_where.py new file mode 100644 index 0000000000..80a259d766 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_cannonicalize_tl_where.py @@ -0,0 +1,124 @@ +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + +types_all = [ + (torch.float32, 'float32'), +] + +shapes_common = [(128, 256), (127, 256), (127, 16), (129, 256), (77, 1024), (69, 512), (512, 512)] + +block_size = [128, 256, 1024] + + +def ceil_div(a, b): + return (a + b - 1) // b + + +def profiler_wrapper(fn, *args): + result_path = "./result_profiling_tl_where" + skip_first = 10 + wait = 0 + warmup = 3 + active = 30 + repeat = 1 + stream = torch.npu.current_stream() + experimental_config = torch_npu.profiler._ExperimentalConfig( + aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, l2_cache=False, data_simplification=False) + with torch_npu.profiler.profile( + activities=[torch_npu.profiler.ProfilerActivity.CPU, torch_npu.profiler.ProfilerActivity.NPU], + schedule=torch_npu.profiler.schedule(wait=wait, warmup=warmup, active=active, repeat=repeat, + skip_first=skip_first), + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(result_path), record_shapes=True, + profile_memory=False, with_stack=False, with_flops=False, with_modules=False, + experimental_config=experimental_config) as prof: + stream.synchronize() + for _ in range(skip_first + (wait + warmup + active) * repeat): + fn(*args) + prof.step() + stream.synchronize() + + +@triton.jit +def tl_where_kernel( + in_ptr, + output_ptr, + N: tl.constexpr, + M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, +): + pid = tl.program_id(axis=0) + offset = tl.multiple_of(pid * BLOCK_SIZE_N, N) + x1 = (offset + tl.arange(0, BLOCK_SIZE_N)) // N + mask1 = tl.where(x1 < M, 1, 0).to(tl.int1) + data = tl.load(in_ptr + x1 * N, mask=mask1, other=0) + x2 = pid * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + tl.store(output_ptr + x2, data) + + +def torch_tl_where(in_tensor): + M = in_tensor.shape[0] // 2 + N = in_tensor.shape[1] + + output = torch.zeros_like(in_tensor) + + first_elements = in_tensor[:M, 0:1] + output[:M] = first_elements.expand(-1, N) + + return output + + +@pytest.mark.parametrize('dtype, sigtype', types_all) +@pytest.mark.parametrize('M, N', shapes_common) +@pytest.mark.parametrize('BLOCK_SIZE_N', block_size) +def test_tl_where(M, N, BLOCK_SIZE_N, dtype, sigtype): + + in_tensor = torch.randn(2 * M, N, dtype=dtype).npu() + + triton_output = torch.zeros_like(in_tensor) + + grid = (ceil_div(2 * M * N, BLOCK_SIZE_N), ) + + tl_where_kernel[grid](in_tensor, triton_output, N=N, M=M, BLOCK_SIZE_N=BLOCK_SIZE_N, optimize_dynamic_offset=False) + + torch_output = torch_tl_where(in_tensor.clone()) + assert torch.allclose(triton_output, torch_output, rtol=1e-5, atol=1e-8) + + +def triton_tl_where(in_tensor, BLOCK_SIZE): + M = in_tensor.shape[0] // 2 + N = in_tensor.shape[1] + + triton_output = torch.zeros_like(in_tensor) + grid = (ceil_div(2 * M * N, BLOCK_SIZE), ) + + tl_where_kernel[grid](in_tensor, triton_output, N=N, M=M, BLOCK_SIZE_N=BLOCK_SIZE, optimize_dynamic_offset=True) + + +def profile_performance_test(M, N, dtype, BLOCK_SIZE): + print(f"\nDetailed performance analysis: M={M}, N={N}, dtype={dtype}, block_size={BLOCK_SIZE}") + + in_tensor = torch.randn(2 * M, N, dtype=dtype).npu() + + def wrapper_func(x): + triton_tl_where(x, BLOCK_SIZE=BLOCK_SIZE) + + # Run performance analysis + profiler_wrapper(wrapper_func, in_tensor) + + +if __name__ == "__main__": + + # Optional: Run detailed profiler test (specific configuration) + profile_performance_test(512, 512, torch.float32, BLOCK_SIZE=1024) + + print("\n" + "=" * 80) + print("Test completed!") + print(f"Detailed performance analysis results saved in: ./result_profiling_tl_where/") + print("=" * 80) diff --git a/third_party/ascend/unittest/pytest_ut/test_cast_full.py b/third_party/ascend/unittest/pytest_ut/test_cast_full.py new file mode 100644 index 0000000000..d3463f913b --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_cast_full.py @@ -0,0 +1,229 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common + +Dimensions = tuple[int, int, int] + + +@triton.jit +def cast_to_bool(output_ptr, x_ptr, dims: Dimensions, overflow_mode: tl.constexpr): + xidx = tl.arange(0, dims.XB) + yidx = tl.arange(0, dims.YB) + zidx = tl.arange(0, dims.ZB) + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + overflow_mode = "trunc" if overflow_mode == 0 else "saturate" + ret = tl.cast(X, dtype=tl.int1, overflow_mode=overflow_mode) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def cast_to_i8(output_ptr, x_ptr, dims: Dimensions, overflow_mode: tl.constexpr): + xidx = tl.arange(0, dims.XB) + yidx = tl.arange(0, dims.YB) + zidx = tl.arange(0, dims.ZB) + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + overflow_mode = "trunc" if overflow_mode == 0 else "saturate" + ret = tl.cast(X, dtype=tl.int8, overflow_mode=overflow_mode) + + tl.store(output_ptr + idx, ret) + + +def cast_to_i16(output_ptr, x_ptr, dims: Dimensions, overflow_mode: tl.constexpr): + xidx = tl.arange(0, dims.XB) + yidx = tl.arange(0, dims.YB) + zidx = tl.arange(0, dims.ZB) + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + overflow_mode = "trunc" if overflow_mode == 0 else "saturate" + ret = tl.cast(X, dtype=tl.int16, overflow_mode=overflow_mode) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def cast_to_i32(output_ptr, x_ptr, dims: Dimensions, overflow_mode: tl.constexpr): + xidx = tl.arange(0, dims.XB) + yidx = tl.arange(0, dims.YB) + zidx = tl.arange(0, dims.ZB) + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + overflow_mode = "trunc" if overflow_mode == 0 else "saturate" + ret = tl.cast(X, dtype=tl.int32, overflow_mode=overflow_mode) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def cast_to_i64(output_ptr, x_ptr, dims: Dimensions, overflow_mode: tl.constexpr): + xidx = tl.arange(0, dims.XB) + yidx = tl.arange(0, dims.YB) + zidx = tl.arange(0, dims.ZB) + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + overflow_mode = "trunc" if overflow_mode == 0 else "saturate" + ret = tl.cast(X, dtype=tl.int64, overflow_mode=overflow_mode) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def cast_to_fp32(output_ptr, x_ptr, dims: Dimensions, overflow_mode: tl.constexpr): + xidx = tl.arange(0, dims.XB) + yidx = tl.arange(0, dims.YB) + zidx = tl.arange(0, dims.ZB) + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + overflow_mode = "trunc" if overflow_mode == 0 else "saturate" + ret = tl.cast(X, dtype=tl.float32, overflow_mode=overflow_mode) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def cast_to_fp16(output_ptr, x_ptr, dims: Dimensions, overflow_mode: tl.constexpr): + xidx = tl.arange(0, dims.XB) + yidx = tl.arange(0, dims.YB) + zidx = tl.arange(0, dims.ZB) + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + overflow_mode = "trunc" if overflow_mode == 0 else "saturate" + ret = tl.cast(X, dtype=tl.float16, overflow_mode=overflow_mode) + + tl.store(output_ptr + idx, ret) + + +@triton.jit +def cast_to_bf16(output_ptr, x_ptr, dims: Dimensions, overflow_mode: tl.constexpr): + xidx = tl.arange(0, dims.XB) + yidx = tl.arange(0, dims.YB) + zidx = tl.arange(0, dims.ZB) + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + overflow_mode = "trunc" if overflow_mode == 0 else "saturate" + ret = tl.cast(X, dtype=tl.bfloat16, overflow_mode=overflow_mode) + + tl.store(output_ptr + idx, ret) + + +import numpy as np + + +def cast_npu(para_type, data_type, to_para, to_dtype, XB, YB, ZB, overflow_mode): + + print(f"TESTING: cast from {para_type} to {to_para} in shape ({XB}, {YB}, {ZB})") + + if para_type == "*i1": + x = torch.randint(low=0, high=2, size=(XB, YB, ZB), dtype=data_type).npu() + elif (para_type == "*i8" or para_type == "*i16" or para_type == "*i32" or para_type == "*64"): + x = torch.randint(low=-128, high=128, size=(XB, YB, ZB), dtype=data_type).npu() + elif para_type == "*i16": + x = torch.randint(low=-32768, high=32768, size=(XB, YB, ZB), dtype=data_type).npu() + elif para_type == "*i32": + x = torch.randint(low=-65536, high=65536, size=(XB, YB, ZB), dtype=data_type).npu() + elif para_type == "*i64": + x = torch.randint(low=-65536, high=65536, size=(XB, YB, ZB), dtype=data_type).npu() + else: # float + x = torch.randn((XB, YB, ZB), dtype=data_type).npu() + + if to_para == "*i1": + triton_func = cast_to_bool + cmp_type = "bool" + elif to_para == "*i8": + triton_func = cast_to_i8 + cmp_type = "int8" + elif to_para == "*i16": + triton_func = cast_to_i16 + cmp_type = "int16" + elif to_para == "*i32": + triton_func = cast_to_i32 + cmp_type = "int32" + elif to_para == "*i64": + triton_func = cast_to_i64 + cmp_type = "int64" + elif to_para == "*fp16": + triton_func = cast_to_fp16 + cmp_type = "float16" + elif to_para == "*fp32": + triton_func = cast_to_fp32 + cmp_type = "float32" + elif to_para == "*bf16": + triton_func = cast_to_bf16 + cmp_type = "bfloat16" + + output = torch.randint(1, (XB, YB, ZB), dtype=to_dtype).npu() + + a = x.to(to_dtype) + dims = Dimensions(XB=XB, YB=YB, ZB=ZB) + + triton_func[1, 1, 1](output, x, dims, overflow_mode) + + test_common.validate_cmp(cmp_type, a, output, overflow_mode) + + +def test_cast_high_priority_dtype(): + + typelist = [ + (torch.int8, "*i8"), + (torch.float32, "*fp32"), + (torch.float16, "*fp16"), + ] + + overflow_mode = [ + 0, # "trunc", + 1, # "saturate", + ] + + shapes = [(8, 32, 32)] + ContinueList = [] + for src in typelist: + for dst in typelist: + if src != dst and (src[1], dst[1]) not in ContinueList: + for shape in shapes: + for mode in overflow_mode: + ( + src[1], + src[0], + dst[1], + dst[0], + shape[0], + shape[1], + shape[2], + mode, + ) + + print("test_cast_full passed") diff --git a/third_party/ascend/unittest/pytest_ut/test_cat.py b/third_party/ascend/unittest/pytest_ut/test_cat.py new file mode 100644 index 0000000000..c1fb9ecbf2 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_cat.py @@ -0,0 +1,61 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common +import math + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, XB: tl.constexpr): + + idx = tl.arange(0, XB) + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.cat(X, Y, can_reorder=True) + + oidx = tl.arange(0, XB * 2) + + tl.store(output_ptr + oidx, ret) + + +# The CAT operator in the Triton community also does not support boolean types. +@pytest.mark.parametrize('shape', [(32, ), (741, )]) #triton only support 1D cat +@pytest.mark.parametrize('dtype', [ + 'float32', +]) +def test_cat(shape, dtype): + m = shape[0] + x = torch.full((m, ), 100, dtype=eval("torch." + dtype)).npu() + y = torch.full((m, ), 30, dtype=eval("torch." + dtype)).npu() + + output = torch.randint(1, (m * 2, ), dtype=eval("torch." + dtype)).npu() + + ans = torch.cat((x, y), dim=0) + + fn_npu_[1, 1, 1](output, x, y, m) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/pytest_ut/test_cat_dim.py b/third_party/ascend/unittest/pytest_ut/test_cat_dim.py new file mode 100644 index 0000000000..c0d6690e95 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_cat_dim.py @@ -0,0 +1,119 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest + + +@triton.jit +def fn3_dim0(output_ptr, x1_ptr, x2_ptr, x3_ptr, x1_shape: tl.constexpr, x2_shape: tl.constexpr, + x3_shape: tl.constexpr): + idx_start = 0 + x1_idx = tl.arange(0, x1_shape) + X1 = tl.load(x1_ptr + x1_idx) + tl.store(output_ptr + x1_idx, X1) + + idx_start += x1_shape + x2_idx = tl.arange(0, x2_shape) + X2 = tl.load(x2_ptr + x2_idx) + tl.store(output_ptr + idx_start + x2_idx, X2) + + idx_start += x2_shape + x3_idx = tl.arange(0, x3_shape) + X3 = tl.load(x3_ptr + x3_idx) + tl.store(output_ptr + idx_start + x3_idx, X3) + + +@triton.jit +def fn4_dim1(output_ptr, x0_ptr, x1_ptr, x2_ptr, x3_ptr, dim0_len: tl.constexpr, x0_len: tl.constexpr, + x1_len: tl.constexpr, x2_len: tl.constexpr, x3_len: tl.constexpr): + + total_dim1_len = x0_len + x1_len + x2_len + x3_len + x0 = tl.load(x0_ptr + tl.arange(0, dim0_len * x0_len)) + x0 = x0.reshape(dim0_len, x0_len) + x1 = tl.load(x1_ptr + tl.arange(0, dim0_len * x1_len)) + x1 = x1.reshape(dim0_len, x1_len) + x2 = tl.load(x2_ptr + tl.arange(0, dim0_len * x2_len)) + x2 = x2.reshape(dim0_len, x2_len) + x3 = tl.load(x3_ptr + tl.arange(0, dim0_len * x3_len)) + x3 = x3.reshape(dim0_len, x3_len) + # (86,5), (86,36), (86,5) (86,4) + #torch.arange(0,86)[:,None] * 50 + torch.arange(0,36) + idx_start = 0 + nidx0 = (tl.arange(0, dim0_len)[:, None] * total_dim1_len + idx_start) + tl.arange(0, x0_len) + tl.store(output_ptr + nidx0, x0) + + idx_start += x0_len + nidx1 = (tl.arange(0, dim0_len)[:, None] * total_dim1_len + idx_start) + tl.arange(0, x1_len) + tl.store(output_ptr + nidx1, x1) + + idx_start += x1_len + nidx2 = (tl.arange(0, dim0_len)[:, None] * total_dim1_len + idx_start) + tl.arange(0, x2_len) + tl.store(output_ptr + nidx2, x2) + + idx_start += x2_len + nidx3 = (tl.arange(0, dim0_len)[:, None] * total_dim1_len + idx_start) + tl.arange(0, x3_len) + tl.store(output_ptr + nidx3, x3) + + +def cat_dim0(data_type): + x1_shape = (86, 48) + x2_shape = (8, 48) + x3_shape = (16, 48) + + x1 = torch.rand(x1_shape, dtype=data_type).npu() + x2 = torch.rand(x2_shape, dtype=data_type).npu() + x3 = torch.rand(x3_shape, dtype=data_type).npu() + res = torch.zeros((x1_shape[0] + x2_shape[0] + x3_shape[0], 48), dtype=data_type).npu() + fn3_dim0[(1, 1, 1)](res, x1, x2, x3, x1_shape[0] * x1_shape[1], x2_shape[0] * x2_shape[1], + x3_shape[0] * x3_shape[1]) + + res_ref = torch.cat((x1, x2, x3), dim=0) + assert torch.allclose(res_ref, res, rtol=1e-03, atol=1e-03, equal_nan=True) + + +def cat_dim1(data_type): + #data_type = torch.float16 + x0_shape = (86, 5) + x1_shape = (86, 36) + x2_shape = (86, 5) + x3_shape = (86, 4) + + dim = 1 + x0 = torch.rand(x0_shape, dtype=data_type).npu() + x1 = torch.rand(x1_shape, dtype=data_type).npu() + x2 = torch.rand(x2_shape, dtype=data_type).npu() + x3 = torch.rand(x3_shape, dtype=data_type).npu() + res = torch.zeros((86, x0_shape[dim] + x1_shape[dim] + x2_shape[dim] + x3_shape[dim]), dtype=data_type).npu() + fn4_dim1[(1, 1, 1)](res, x0, x1, x2, x3, 86, x0_shape[1], x1_shape[1], x2_shape[1], x3_shape[1]) + #print("res_tri=", res) + + res_ref = torch.cat((x0, x1, x2, x3), dim=1) + #print("res_ref=", res_ref) + assert torch.allclose(res_ref, res, rtol=1e-03, atol=1e-03, equal_nan=True) + + +def test_cat(): + cat_dim0(torch.float16) + cat_dim1(torch.float16) diff --git a/third_party/ascend/unittest/pytest_ut/test_cat_help_func.py b/third_party/ascend/unittest/pytest_ut/test_cat_help_func.py new file mode 100644 index 0000000000..4b102b9335 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_cat_help_func.py @@ -0,0 +1,728 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import random +import pytest + +import triton +import triton.language as tl +import torch +import test_common +import numpy as np +import triton.language.extra.cann.extension as extension + + +def gen_1d_cat_shapes(min_val=1, max_val=4096): + shape1 = random.randint(min_val, max_val) + shape2 = random.randint(min_val, max_val) + return (shape1, ), (shape2, ), 0 + + +def gen_2d_cat_shapes(dim=0, min_val=1, max_val=4096): + if dim == 0: + common_col = random.randint(min_val, max_val) + row1 = random.randint(min_val, max_val) + row2 = random.randint(min_val, max_val) + shape1 = (row1, common_col) + shape2 = (row2, common_col) + elif dim == 1: + common_row = random.randint(min_val, max_val) + col1 = random.randint(min_val, max_val) + col2 = random.randint(min_val, max_val) + shape1 = (common_row, col1) + shape2 = (common_row, col2) + else: + raise ValueError("2d shape only support dim=0 or dim=1") + return shape1, shape2, dim + + +def gen_3d_cat_shapes(dim=0, min_val=1, max_val=4096): + if dim not in [0, 1, 2]: + raise ValueError("3d shape only support dim=0/1/2") + + if dim == 0: + common_d1 = random.randint(min_val, max_val) + common_d2 = random.randint(min_val, max_val) + d0_1 = random.randint(min_val, max_val) + d0_2 = random.randint(min_val, max_val) + shape1 = (d0_1, common_d1, common_d2) + shape2 = (d0_2, common_d1, common_d2) + + elif dim == 1: + common_d0 = random.randint(min_val, max_val) + common_d2 = random.randint(min_val, max_val) + d1_1 = random.randint(min_val, max_val) + d1_2 = random.randint(min_val, max_val) + shape1 = (common_d0, d1_1, common_d2) + shape2 = (common_d0, d1_2, common_d2) + + else: # dim == 2 + common_d0 = random.randint(min_val, max_val) + common_d1 = random.randint(min_val, max_val) + d2_1 = random.randint(min_val, max_val) + d2_2 = random.randint(min_val, max_val) + shape1 = (common_d0, common_d1, d2_1) + shape2 = (common_d0, common_d1, d2_2) + + return shape1, shape2, dim + + +def gen_100_cat_shapes(num_groups=100, mix_ratio=(0.3, 0.3, 0.4), min_val=1, max_val=4096): + + shape_list = [] + num_1d = int(num_groups * mix_ratio[0]) + num_2d = int(num_groups * mix_ratio[1]) + num_3d = num_groups - num_1d - num_2d + + for _ in range(num_1d): + shape_list.append(gen_1d_cat_shapes(min_val, max_val)) + + for _ in range(num_2d): + dim = random.choice([0, 1]) + shape_list.append(gen_2d_cat_shapes(dim, min_val, max_val)) + + for _ in range(num_3d): + dim = random.choice([0, 1, 2]) + shape_list.append(gen_3d_cat_shapes(dim, min_val, max_val)) + + random.shuffle(shape_list) + return shape_list + + +full_shape = gen_100_cat_shapes(num_groups=100, mix_ratio=(0.3, 0.4, 0.3), min_val=1, max_val=4096) + + +@triton.jit +def _cat_helper_func_2D_1( + in_ptr0, + in_ptr1, + out_ptr0, + in0_x: tl.constexpr, + in1_x: tl.constexpr, + y0_numel, + x1_numel, + Y0BLOCK: tl.constexpr, + Y0BLOCK_SUB: tl.constexpr, +): + y0_offset = tl.program_id(0) * Y0BLOCK_SUB + base_y0 = tl.arange(0, Y0BLOCK_SUB) + loops_y0 = (Y0BLOCK + Y0BLOCK_SUB - 1) // Y0BLOCK_SUB + base_input0_x1 = tl.arange(0, in0_x)[None, :] + base_input1_x1 = tl.arange(0, in1_x)[None, :] + x1 = tl.arange(0, in0_x + in1_x)[None, :] + + for loop in range(loops_y0): + y0 = y0_offset + (loop * Y0BLOCK_SUB) + base_y0[:, None] + y0_mask = y0 < min(Y0BLOCK + y0_offset, y0_numel) + x1_mask = x1 < x1_numel + tmp0 = tl.load(in_ptr0 + (base_input0_x1 + in0_x * y0), y0_mask) + tmp1 = tl.load(in_ptr1 + (base_input1_x1 + in1_x * y0), y0_mask) + tmp2 = tl.zeros((Y0BLOCK_SUB, in0_x + in1_x), dtype=tmp0.dtype) + tmp3 = extension.insert_slice(tmp2, tmp0, [0, 0], [Y0BLOCK_SUB, in0_x], [1, 1]) + tmp4 = extension.insert_slice(tmp3, tmp1, [0, in0_x], [Y0BLOCK_SUB, in1_x], [1, 1]) + tl.store(out_ptr0 + (x1 + (in0_x + in1_x) * y0), tmp4, x1_mask & y0_mask) + + +@triton.jit +def triton_unk_fused_cat_dim0_sameshape(output_ptr, x_ptr, y_ptr, y0_numel, x1_numel, Y0BLOCK: tl.constexpr, + Y0BLOCK_SUB: tl.constexpr, X1BLOCK_SUB: tl.constexpr): + y0_offset = tl.program_id(0) * Y0BLOCK + base_y0 = tl.arange(0, Y0BLOCK_SUB) + loops_y0 = (Y0BLOCK + Y0BLOCK_SUB - 1) // Y0BLOCK_SUB + base_x1 = tl.arange(0, X1BLOCK_SUB) + loops_x1 = (x1_numel + X1BLOCK_SUB - 1) // X1BLOCK_SUB + for loop_y0 in range(loops_y0): + y0 = y0_offset + (loop_y0 * Y0BLOCK_SUB) + base_y0[:, None] + y0_mask = y0 < min(Y0BLOCK + y0_offset, y0_numel) + for loop_x1 in range(loops_x1): + x1 = (loop_x1 * X1BLOCK_SUB) + base_x1[None, :] + x1_mask = x1 < x1_numel + + tmp0 = tl.load(x_ptr + (x1 + x1_numel * y0), x1_mask & y0_mask) + tmp8 = tl.load(y_ptr + (x1 + x1_numel * y0), x1_mask & y0_mask) + tmp10 = tl.zeros((2 * Y0BLOCK_SUB, X1BLOCK_SUB), dtype=tmp0.dtype) + tmp11 = extension.insert_slice(tmp10, tmp0, [0, 0], [Y0BLOCK_SUB, X1BLOCK_SUB], [1, 1]) + tmp12 = extension.insert_slice(tmp11, tmp8, [Y0BLOCK_SUB, 0], [Y0BLOCK_SUB, X1BLOCK_SUB], [1, 1]) + tmp13 = tl.reshape(tmp12, (2, Y0BLOCK_SUB, X1BLOCK_SUB)) + + new_base_x2 = tl.arange(0, X1BLOCK_SUB) + new_x2 = (loop_x1 * X1BLOCK_SUB) + new_base_x2[None, None, :] + new_base_y1 = tl.arange(0, Y0BLOCK_SUB) + new_y1 = y0_offset + (loop_y0 * Y0BLOCK_SUB) + new_base_y1[None, :, None] + new_z0 = tl.arange(0, 2)[:, None, None] + new_x2_mask = new_x2 < x1_numel + new_y1_mask = new_y1 < y0_numel + tl.store(output_ptr + (new_x2 + x1_numel * (new_y1 + y0_numel * new_z0)), tmp13, new_x2_mask & new_y1_mask) + + +@triton.jit +def triton_unk_fused_cat_dim0_diffshape(output_ptr, x_ptr, y_ptr, y0_numel, y1_numel, x1_numel, YBLOCK: tl.constexpr, + YBLOCK_2: tl.constexpr, YBLOCK_SUB: tl.constexpr, X1BLOCK_SUB: tl.constexpr): + y0_offset = tl.program_id(0) * YBLOCK + base_y0 = tl.arange(0, YBLOCK_SUB) + loops_y0 = (YBLOCK + YBLOCK_SUB - 1) // YBLOCK_SUB + base_x1 = tl.arange(0, X1BLOCK_SUB) + loops_x1 = (x1_numel + X1BLOCK_SUB - 1) // X1BLOCK_SUB + min_numel = 0 + max_numel = 0 + clone_numel = 0 + if y0_numel < y1_numel: + min_numel = y0_numel + max_numel = y1_numel + clone_numel = y1_numel - y0_numel + else: + min_numel = y1_numel + max_numel = y0_numel + clone_numel = y0_numel - y1_numel + + for loops_y in range(loops_y0): + y0 = y0_offset + (loops_y * YBLOCK_SUB) + base_y0[:, None] + y0_mask = y0 < min(YBLOCK + y0_offset, min_numel) + for loop_x1 in range(loops_x1): + x1 = (loop_x1 * X1BLOCK_SUB) + base_x1[None, :] + x1_mask = x1 < x1_numel + + tmp0 = tl.load(x_ptr + (x1 + x1_numel * y0), x1_mask & y0_mask) + tmp8 = tl.load(y_ptr + (x1 + x1_numel * y0), x1_mask & y0_mask) + tmp10 = tl.zeros((2 * YBLOCK_SUB, X1BLOCK_SUB), dtype=tmp0.dtype) + tmp11 = extension.insert_slice(tmp10, tmp0, [0, 0], [YBLOCK_SUB, X1BLOCK_SUB], [1, 1]) + tmp12 = extension.insert_slice(tmp11, tmp8, [YBLOCK_SUB, 0], [YBLOCK_SUB, X1BLOCK_SUB], [1, 1]) + tmp13 = tl.reshape(tmp12, (2, YBLOCK_SUB, X1BLOCK_SUB)) + + new_base_x2 = tl.arange(0, X1BLOCK_SUB) + new_x2 = (loop_x1 * X1BLOCK_SUB) + new_base_x2[None, None, :] + new_base_y1 = tl.arange(0, YBLOCK_SUB) + new_y1 = y0_offset + (loops_y * YBLOCK_SUB) + new_base_y1[None, :, None] + new_z0 = tl.arange(0, 2)[:, None, None] + new_x2_mask = new_x2 < x1_numel + new_y1_mask = new_y1 < min_numel + tl.store(output_ptr + (new_x2 + x1_numel * new_y1 + x1_numel * y0_numel * new_z0), tmp13, + new_x2_mask & new_y1_mask) + + loops_y1 = (YBLOCK_2 + YBLOCK_SUB - 1) // YBLOCK_SUB + y2_offset = tl.program_id(0) * YBLOCK_2 + min_numel + if y0_numel < y1_numel: + for loops_y1 in range(loops_y1): + y0 = y2_offset + (loops_y1 * YBLOCK_SUB) + base_y0[:, None] + y0_mask = y0 < min(YBLOCK_2 + y2_offset, y1_numel) + for loop_x1 in range(loops_x1): + x1 = (loop_x1 * X1BLOCK_SUB) + base_x1[None, :] + x1_mask = x1 < x1_numel + + tmp8 = tl.load(y_ptr + (x1 + x1_numel * y0), x1_mask & y0_mask) + new_base_x2 = tl.arange(0, X1BLOCK_SUB) + new_x2 = (loop_x1 * X1BLOCK_SUB) + new_base_x2[None, :] + new_base_y1 = tl.arange(0, YBLOCK_SUB) + new_y1 = y2_offset + y0_numel + (loops_y1 * YBLOCK_SUB) + new_base_y1[:, None] + sum_numel = y0_numel + y1_numel + new_x2_mask = new_x2 < x1_numel + new_y1_mask = new_y1 < sum_numel + tl.store(output_ptr + (new_x2 + x1_numel * new_y1), tmp8, new_x2_mask & new_y1_mask) + else: + for loops_y1 in range(loops_y1): + y0 = y2_offset + (loops_y1 * YBLOCK_SUB) + base_y0[:, None] + y0_mask = y0 < min(YBLOCK_2 + y2_offset, y0_numel) + for loop_x1 in range(loops_x1): + x1 = (loop_x1 * X1BLOCK_SUB) + base_x1[None, :] + x1_mask = x1 < x1_numel + + tmp8 = tl.load(x_ptr + (x1 + x1_numel * y0), x1_mask & y0_mask) + new_base_x2 = tl.arange(0, X1BLOCK_SUB) + new_x2 = (loop_x1 * X1BLOCK_SUB) + new_base_x2[None, :] + new_base_y1 = tl.arange(0, YBLOCK_SUB) + new_y1 = y2_offset + (loops_y1 * YBLOCK_SUB) + new_base_y1[:, None] + new_x2_mask = new_x2 < x1_numel + new_y1_mask = new_y1 < y0_numel + tl.store(output_ptr + (new_x2 + x1_numel * new_y1), tmp8, new_x2_mask & new_y1_mask) + + +@triton.jit +def triton_unk_fused_cat_dim1_sameshape(output_ptr, x_ptr, y_ptr, y0_numel, x1_numel, Y0BLOCK: tl.constexpr, + Y0BLOCK_SUB: tl.constexpr, X1BLOCK_SUB: tl.constexpr): + y0_offset = tl.program_id(0) * Y0BLOCK + base_y0 = tl.arange(0, Y0BLOCK_SUB) + loops_y0 = (Y0BLOCK + Y0BLOCK_SUB - 1) // Y0BLOCK_SUB + base_x1 = tl.arange(0, X1BLOCK_SUB) + loops_x1 = (x1_numel + X1BLOCK_SUB - 1) // X1BLOCK_SUB + for loop_y0 in range(loops_y0): + y0 = y0_offset + (loop_y0 * Y0BLOCK_SUB) + base_y0[:, None] + y0_mask = y0 < min(Y0BLOCK + y0_offset, y0_numel) + for loop_x1 in range(loops_x1): + x1 = (loop_x1 * X1BLOCK_SUB) + base_x1[None, :] + x1_mask = x1 < x1_numel + + tmp0 = tl.load(x_ptr + (x1 + x1_numel * y0), x1_mask & y0_mask) + tmp8 = tl.load(y_ptr + (x1 + x1_numel * y0), x1_mask & y0_mask) + tmp10 = tl.zeros((Y0BLOCK_SUB, 2 * X1BLOCK_SUB), dtype=tmp0.dtype) + tmp11 = extension.insert_slice(tmp10, tmp0, [0, 0], [Y0BLOCK_SUB, X1BLOCK_SUB], [1, 1]) + tmp12 = extension.insert_slice(tmp11, tmp8, [0, X1BLOCK_SUB], [Y0BLOCK_SUB, X1BLOCK_SUB], [1, 1]) + tmp13 = tl.reshape(tmp12, (Y0BLOCK_SUB, 2, X1BLOCK_SUB)) + + new_base_x2 = tl.arange(0, X1BLOCK_SUB) + new_x2 = (loop_x1 * X1BLOCK_SUB) + new_base_x2[None, None, :] + new_base_y1 = tl.arange(0, Y0BLOCK_SUB) + new_y1 = y0_offset + (loop_y0 * Y0BLOCK_SUB) + new_base_y1[:, None, None] + new_z0 = tl.arange(0, 2)[None, :, None] + new_x2_mask = new_x2 < x1_numel + new_y1_mask = new_y1 < y0_numel + tl.store(output_ptr + (new_x2 + 2 * x1_numel * new_y1 + x1_numel * new_z0), tmp13, + new_x2_mask & new_y1_mask) + + +@triton.jit +def triton_unk_fused_cat_dim1_diffshape(output_ptr, x_ptr, y_ptr, y0_numel, x0_numel, x1_numel, Y0BLOCK: tl.constexpr, + Y0BLOCK_SUB: tl.constexpr, XBLOCK_SUB: tl.constexpr): + y0_offset = tl.program_id(0) * Y0BLOCK + base_y0 = tl.arange(0, Y0BLOCK_SUB) + loops_y0 = (Y0BLOCK + Y0BLOCK_SUB - 1) // Y0BLOCK_SUB + base_x = tl.arange(0, XBLOCK_SUB) + min_numel = 0 + max_numel = 0 + clone_numel = 0 + if x0_numel < x1_numel: + min_numel = x0_numel + max_numel = x1_numel + clone_numel = x1_numel - x0_numel + else: + min_numel = x1_numel + max_numel = x0_numel + clone_numel = x0_numel - x1_numel + loops_x = (min_numel + XBLOCK_SUB - 1) // XBLOCK_SUB + loops_x2 = (clone_numel + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop_y0 in range(loops_y0): + y0 = y0_offset + (loop_y0 * Y0BLOCK_SUB) + base_y0[:, None] + y0_mask = y0 < min(Y0BLOCK + y0_offset, y0_numel) + for loop_x in range(loops_x): + x = (loop_x * XBLOCK_SUB) + base_x[None, :] + x_mask = x < min_numel + + tmp0 = tl.load(x_ptr + (x + x0_numel * y0), x_mask & y0_mask) + tmp8 = tl.load(y_ptr + (x + x1_numel * y0), x_mask & y0_mask) + tmp10 = tl.zeros((Y0BLOCK_SUB, 2 * XBLOCK_SUB), dtype=tmp0.dtype) + tmp11 = extension.insert_slice(tmp10, tmp0, [0, 0], [Y0BLOCK_SUB, XBLOCK_SUB], [1, 1]) + tmp12 = extension.insert_slice(tmp11, tmp8, [0, XBLOCK_SUB], [Y0BLOCK_SUB, XBLOCK_SUB], [1, 1]) + tmp13 = tl.reshape(tmp12, (Y0BLOCK_SUB, 2, XBLOCK_SUB)) + + new_base_x2 = tl.arange(0, XBLOCK_SUB) + new_x2 = (loop_x * XBLOCK_SUB) + new_base_x2[None, None, :] + new_base_y1 = tl.arange(0, Y0BLOCK_SUB) + new_y1 = y0_offset + (loop_y0 * Y0BLOCK_SUB) + new_base_y1[:, None, None] + new_z0 = tl.arange(0, 2)[None, :, None] + new_x2_mask = new_x2 < min_numel + new_y1_mask = new_y1 < y0_numel + sum_numel = x0_numel + x1_numel + tl.store(output_ptr + (new_x2 + sum_numel * new_y1 + x0_numel * new_z0), tmp13, new_x2_mask & new_y1_mask) + + if x0_numel < x1_numel: + for loop_y0 in range(loops_y0): + y0 = y0_offset + (loop_y0 * Y0BLOCK_SUB) + base_y0[:, None] + y0_mask = y0 < min(Y0BLOCK + y0_offset, y0_numel) + for loop_x2 in range(loops_x2): + x = (loop_x2 * XBLOCK_SUB) + base_x[None, :] + min_numel + x_mask = x < x1_numel + + tmp8 = tl.load(y_ptr + (x + x1_numel * y0), x_mask & y0_mask) + new_base_x2 = tl.arange(0, XBLOCK_SUB) + new_x2 = x0_numel + min_numel + (loop_x2 * XBLOCK_SUB) + new_base_x2[None, :] + new_base_y1 = tl.arange(0, Y0BLOCK_SUB) + new_y1 = y0_offset + (loop_y0 * Y0BLOCK_SUB) + new_base_y1[:, None] + sum_numel = x0_numel + x1_numel + new_x2_mask = new_x2 < sum_numel + new_y1_mask = new_y1 < y0_numel + tl.store(output_ptr + (new_x2 + sum_numel * new_y1), tmp8, new_x2_mask & new_y1_mask) + else: + for loop_y0 in range(loops_y0): + y0 = y0_offset + (loop_y0 * Y0BLOCK_SUB) + base_y0[:, None] + y0_mask = y0 < min(Y0BLOCK + y0_offset, y0_numel) + for loop_x2 in range(loops_x2): + x = (loop_x2 * XBLOCK_SUB) + base_x[None, :] + min_numel + x_mask = x < x0_numel + + tmp8 = tl.load(x_ptr + (x + x0_numel * y0), x_mask & y0_mask) + new_base_x2 = tl.arange(0, XBLOCK_SUB) + new_x2 = min_numel + (loop_x2 * XBLOCK_SUB) + new_base_x2[None, :] + new_base_y1 = tl.arange(0, Y0BLOCK_SUB) + new_y1 = y0_offset + (loop_y0 * Y0BLOCK_SUB) + new_base_y1[:, None] + sum_numel = x0_numel + x1_numel + new_x2_mask = new_x2 < x0_numel + new_y1_mask = new_y1 < y0_numel + tl.store(output_ptr + (new_x2 + sum_numel * new_y1), tmp8, new_x2_mask & new_y1_mask) + + +@triton.jit +def triton_unk_fused_cat_3d_dim0(output_ptr, x_ptr, y_ptr, z0_numel, z1_numel, y1_numel, x1_numel, ZBLOCK: tl.constexpr, + ZBLOCK_2: tl.constexpr, ZBLOCK_SUB: tl.constexpr, X1BLOCK_SUB: tl.constexpr): + z0_offset = tl.program_id(0) * ZBLOCK + base_z0 = tl.arange(0, ZBLOCK_SUB) + loops_z0 = (ZBLOCK + ZBLOCK_SUB - 1) // ZBLOCK_SUB + xy_numel = x1_numel * y1_numel + base_x1 = tl.arange(0, X1BLOCK_SUB) + loops_x1 = (xy_numel + X1BLOCK_SUB - 1) // X1BLOCK_SUB + min_numel = 0 + max_numel = 0 + clone_numel = 0 + if z0_numel < z1_numel: + min_numel = z0_numel + max_numel = z1_numel + clone_numel = z1_numel - z0_numel + else: + min_numel = z1_numel + max_numel = z0_numel + clone_numel = z0_numel - z1_numel + + for loops_z in range(loops_z0): + z0 = z0_offset + (loops_z * ZBLOCK_SUB) + base_z0[:, None] + z0_mask = z0 < min(ZBLOCK + z0_offset, min_numel) + for loop_x1 in range(loops_x1): + x1 = (loop_x1 * X1BLOCK_SUB) + base_x1[None, :] + x1_mask = x1 < xy_numel + + tmp0 = tl.load(x_ptr + (x1 + xy_numel * z0), x1_mask & z0_mask) + tmp8 = tl.load(y_ptr + (x1 + xy_numel * z0), x1_mask & z0_mask) + tmp10 = tl.zeros((2 * ZBLOCK_SUB, X1BLOCK_SUB), dtype=tmp0.dtype) + tmp11 = extension.insert_slice(tmp10, tmp0, [0, 0], [ZBLOCK_SUB, X1BLOCK_SUB], [1, 1]) + tmp12 = extension.insert_slice(tmp11, tmp8, [ZBLOCK_SUB, 0], [ZBLOCK_SUB, X1BLOCK_SUB], [1, 1]) + tmp13 = tl.reshape(tmp12, (2, ZBLOCK_SUB, X1BLOCK_SUB)) + + new_base_x2 = tl.arange(0, X1BLOCK_SUB) + new_x2 = (loop_x1 * X1BLOCK_SUB) + new_base_x2[None, None, :] + new_base_z1 = tl.arange(0, ZBLOCK_SUB) + new_z1 = z0_offset + (loops_z * ZBLOCK_SUB) + new_base_z1[None, :, None] + new_z0 = tl.arange(0, 2)[:, None, None] + new_x2_mask = new_x2 < xy_numel + new_z1_mask = new_z1 < min_numel + tl.store(output_ptr + (new_x2 + xy_numel * new_z1 + xy_numel * z0_numel * new_z0), tmp13, + new_x2_mask & new_z1_mask) + + loops_z1 = (ZBLOCK_2 + ZBLOCK_SUB - 1) // ZBLOCK_SUB + z2_offset = tl.program_id(0) * ZBLOCK_2 + min_numel + if z0_numel < z1_numel: + for loops_z1 in range(loops_z1): + z0 = z2_offset + (loops_z1 * ZBLOCK_SUB) + base_z0[:, None] + z0_mask = z0 < min(ZBLOCK_2 + z2_offset, z1_numel) + for loop_x1 in range(loops_x1): + x1 = (loop_x1 * X1BLOCK_SUB) + base_x1[None, :] + x1_mask = x1 < xy_numel + + tmp8 = tl.load(y_ptr + (x1 + xy_numel * z0), x1_mask & z0_mask) + new_base_x2 = tl.arange(0, X1BLOCK_SUB) + new_x2 = (loop_x1 * X1BLOCK_SUB) + new_base_x2[None, :] + new_base_z1 = tl.arange(0, ZBLOCK_SUB) + new_z1 = z2_offset + z0_numel + (loops_z1 * ZBLOCK_SUB) + new_base_z1[:, None] + sum_numel = z0_numel + z1_numel + new_x2_mask = new_x2 < xy_numel + new_z1_mask = new_z1 < sum_numel + tl.store(output_ptr + (new_x2 + xy_numel * new_z1), tmp8, new_x2_mask & new_z1_mask) + else: + for loops_z1 in range(loops_z1): + z0 = z2_offset + (loops_z1 * ZBLOCK_SUB) + base_z0[:, None] + z0_mask = z0 < min(ZBLOCK_2 + z2_offset, z0_numel) + for loop_x1 in range(loops_x1): + x1 = (loop_x1 * X1BLOCK_SUB) + base_x1[None, :] + x1_mask = x1 < xy_numel + + tmp8 = tl.load(x_ptr + (x1 + xy_numel * z0), x1_mask & z0_mask) + new_base_x2 = tl.arange(0, X1BLOCK_SUB) + new_x2 = (loop_x1 * X1BLOCK_SUB) + new_base_x2[None, :] + new_base_z1 = tl.arange(0, ZBLOCK_SUB) + new_z1 = z2_offset + (loops_z1 * ZBLOCK_SUB) + new_base_z1[:, None] + new_x2_mask = new_x2 < xy_numel + new_z1_mask = new_z1 < z0_numel + tl.store(output_ptr + (new_x2 + xy_numel * new_z1), tmp8, new_x2_mask & new_z1_mask) + + +@triton.jit +def triton_unk_fused_cat_3d_dim1(output_ptr, x_ptr, y_ptr, z0_numel, y0_numel, y1_numel, x0_numel, + Z0BLOCK: tl.constexpr, Z0BLOCK_SUB: tl.constexpr, XBLOCK_SUB: tl.constexpr): + z0_offset = tl.program_id(0) * Z0BLOCK + base_z0 = tl.arange(0, Z0BLOCK_SUB) + loops_z0 = (Z0BLOCK + Z0BLOCK_SUB - 1) // Z0BLOCK_SUB + base_x = tl.arange(0, XBLOCK_SUB) + min_numel = 0 + max_numel = 0 + clone_numel = 0 + if y0_numel < y1_numel: + min_numel = y0_numel * x0_numel + max_numel = y1_numel * x0_numel + clone_numel = (y1_numel - y0_numel) * x0_numel + else: + min_numel = y1_numel * x0_numel + max_numel = y0_numel * x0_numel + clone_numel = (y0_numel - y1_numel) * x0_numel + loops_x = (min_numel + XBLOCK_SUB - 1) // XBLOCK_SUB + loops_x2 = (clone_numel + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop_z0 in range(loops_z0): + z0 = z0_offset + (loop_z0 * Z0BLOCK_SUB) + base_z0[:, None] + z0_mask = z0 < min(Z0BLOCK + z0_offset, z0_numel) + for loop_x in range(loops_x): + x = (loop_x * XBLOCK_SUB) + base_x[None, :] + x_mask = x < min_numel + + tmp0 = tl.load(x_ptr + (x + x0_numel * y0_numel * z0), x_mask & z0_mask) + tmp8 = tl.load(y_ptr + (x + x0_numel * y1_numel * z0), x_mask & z0_mask) + tmp10 = tl.zeros((Z0BLOCK_SUB, 2 * XBLOCK_SUB), dtype=tmp0.dtype) + tmp11 = extension.insert_slice(tmp10, tmp0, [0, 0], [Z0BLOCK_SUB, XBLOCK_SUB], [1, 1]) + tmp12 = extension.insert_slice(tmp11, tmp8, [0, XBLOCK_SUB], [Z0BLOCK_SUB, XBLOCK_SUB], [1, 1]) + tmp13 = tl.reshape(tmp12, (Z0BLOCK_SUB, 2, XBLOCK_SUB)) + + new_base_x2 = tl.arange(0, XBLOCK_SUB) + new_x2 = (loop_x * XBLOCK_SUB) + new_base_x2[None, None, :] + new_base_z1 = tl.arange(0, Z0BLOCK_SUB) + new_z1 = z0_offset + (loop_z0 * Z0BLOCK_SUB) + new_base_z1[:, None, None] + new_z0 = tl.arange(0, 2)[None, :, None] + new_x2_mask = new_x2 < min_numel + new_z1_mask = new_z1 < z0_numel + sum_numel = min_numel + max_numel + tl.store(output_ptr + (new_x2 + sum_numel * new_z1 + x0_numel * y0_numel * new_z0), tmp13, + new_x2_mask & new_z1_mask) + + if y0_numel == y1_numel: + return + + if y0_numel < y1_numel: + for loop_z0 in range(loops_z0): + z0 = z0_offset + (loop_z0 * Z0BLOCK_SUB) + base_z0[:, None] + z0_mask = z0 < min(Z0BLOCK + z0_offset, z0_numel) + for loop_x2 in range(loops_x2): + x = (loop_x2 * XBLOCK_SUB) + base_x[None, :] + min_numel + x_mask = x < y1_numel * x0_numel + + tmp8 = tl.load(y_ptr + (x + x0_numel * y1_numel * z0), x_mask & z0_mask) + new_base_x2 = tl.arange(0, XBLOCK_SUB) + new_x2 = x0_numel * y0_numel + min_numel + (loop_x2 * XBLOCK_SUB) + new_base_x2[None, :] + new_base_z1 = tl.arange(0, Z0BLOCK_SUB) + new_z1 = z0_offset + (loop_z0 * Z0BLOCK_SUB) + new_base_z1[:, None] + sum_numel = min_numel + max_numel + new_x2_mask = new_x2 < sum_numel + new_z1_mask = new_z1 < z0_numel + tl.store(output_ptr + (new_x2 + sum_numel * new_z1), tmp8, new_x2_mask & new_z1_mask) + else: + for loop_z0 in range(loops_z0): + z0 = z0_offset + (loop_z0 * Z0BLOCK_SUB) + base_z0[:, None] + z0_mask = z0 < min(Z0BLOCK + z0_offset, z0_numel) + for loop_x2 in range(loops_x2): + x = (loop_x2 * XBLOCK_SUB) + base_x[None, :] + min_numel + x_mask = x < x0_numel * y0_numel + + tmp8 = tl.load(x_ptr + (x + x0_numel * y0_numel * z0), x_mask & z0_mask) + new_base_x2 = tl.arange(0, XBLOCK_SUB) + new_x2 = min_numel + (loop_x2 * XBLOCK_SUB) + new_base_x2[None, :] + new_base_z1 = tl.arange(0, Z0BLOCK_SUB) + new_z1 = z0_offset + (loop_z0 * Z0BLOCK_SUB) + new_base_z1[:, None] + sum_numel = min_numel + max_numel + new_x2_mask = new_x2 < x0_numel * y0_numel + new_z1_mask = new_z1 < z0_numel + tl.store(output_ptr + (new_x2 + sum_numel * new_z1), tmp8, new_x2_mask & new_z1_mask) + + +@triton.jit +def triton_unk_fused_cat_3d_dim2(output_ptr, x_ptr, y_ptr, z0_numel, y0_numel, x0_numel, x1_numel, + Y0BLOCK: tl.constexpr, Y0BLOCK_SUB: tl.constexpr, XBLOCK_SUB: tl.constexpr): + y0_offset = tl.program_id(0) * Y0BLOCK + base_y0 = tl.arange(0, Y0BLOCK_SUB) + loops_y0 = (Y0BLOCK + Y0BLOCK_SUB - 1) // Y0BLOCK_SUB + base_x = tl.arange(0, XBLOCK_SUB) + min_numel = 0 + max_numel = 0 + clone_numel = 0 + zy_numel = z0_numel * y0_numel + if x0_numel < x1_numel: + min_numel = x0_numel + max_numel = x1_numel + clone_numel = x1_numel - x0_numel + else: + min_numel = x1_numel + max_numel = x0_numel + clone_numel = x0_numel - x1_numel + loops_x = (min_numel + XBLOCK_SUB - 1) // XBLOCK_SUB + loops_x2 = (clone_numel + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop_y0 in range(loops_y0): + y0 = y0_offset + (loop_y0 * Y0BLOCK_SUB) + base_y0[:, None] + y0_mask = y0 < min(Y0BLOCK + y0_offset, zy_numel) + for loop_x in range(loops_x): + x = (loop_x * XBLOCK_SUB) + base_x[None, :] + x_mask = x < min_numel + + tmp0 = tl.load(x_ptr + (x + x0_numel * y0), x_mask & y0_mask) + tmp8 = tl.load(y_ptr + (x + x1_numel * y0), x_mask & y0_mask) + tmp10 = tl.zeros((Y0BLOCK_SUB, 2 * XBLOCK_SUB), dtype=tmp0.dtype) + tmp11 = extension.insert_slice(tmp10, tmp0, [0, 0], [Y0BLOCK_SUB, XBLOCK_SUB], [1, 1]) + tmp12 = extension.insert_slice(tmp11, tmp8, [0, XBLOCK_SUB], [Y0BLOCK_SUB, XBLOCK_SUB], [1, 1]) + tmp13 = tl.reshape(tmp12, (Y0BLOCK_SUB, 2, XBLOCK_SUB)) + + new_base_x2 = tl.arange(0, XBLOCK_SUB) + new_x2 = (loop_x * XBLOCK_SUB) + new_base_x2[None, None, :] + new_base_y1 = tl.arange(0, Y0BLOCK_SUB) + new_y1 = y0_offset + (loop_y0 * Y0BLOCK_SUB) + new_base_y1[:, None, None] + new_z0 = tl.arange(0, 2)[None, :, None] + new_x2_mask = new_x2 < min_numel + new_y1_mask = new_y1 < zy_numel + sum_numel = x0_numel + x1_numel + tl.store(output_ptr + (new_x2 + sum_numel * new_y1 + x0_numel * new_z0), tmp13, new_x2_mask & new_y1_mask) + + if x0_numel == x1_numel: + return + + if x0_numel < x1_numel: + for loop_y0 in range(loops_y0): + y0 = y0_offset + (loop_y0 * Y0BLOCK_SUB) + base_y0[:, None] + y0_mask = y0 < min(Y0BLOCK + y0_offset, zy_numel) + for loop_x2 in range(loops_x2): + x = (loop_x2 * XBLOCK_SUB) + base_x[None, :] + min_numel + x_mask = x < x1_numel + + tmp8 = tl.load(y_ptr + (x + x1_numel * y0), x_mask & y0_mask) + new_base_x2 = tl.arange(0, XBLOCK_SUB) + new_x2 = x0_numel + min_numel + (loop_x2 * XBLOCK_SUB) + new_base_x2[None, :] + new_base_y1 = tl.arange(0, Y0BLOCK_SUB) + new_y1 = y0_offset + (loop_y0 * Y0BLOCK_SUB) + new_base_y1[:, None] + sum_numel = x0_numel + x1_numel + new_x2_mask = new_x2 < sum_numel + new_y1_mask = new_y1 < zy_numel + tl.store(output_ptr + (new_x2 + sum_numel * new_y1), tmp8, new_x2_mask & new_y1_mask) + else: + for loop_y0 in range(loops_y0): + y0 = y0_offset + (loop_y0 * Y0BLOCK_SUB) + base_y0[:, None] + y0_mask = y0 < min(Y0BLOCK + y0_offset, zy_numel) + for loop_x2 in range(loops_x2): + x = (loop_x2 * XBLOCK_SUB) + base_x[None, :] + min_numel + x_mask = x < x0_numel + + tmp8 = tl.load(x_ptr + (x + x0_numel * y0), x_mask & y0_mask) + new_base_x2 = tl.arange(0, XBLOCK_SUB) + new_x2 = min_numel + (loop_x2 * XBLOCK_SUB) + new_base_x2[None, :] + new_base_y1 = tl.arange(0, Y0BLOCK_SUB) + new_y1 = y0_offset + (loop_y0 * Y0BLOCK_SUB) + new_base_y1[:, None] + sum_numel = x0_numel + x1_numel + new_x2_mask = new_x2 < x0_numel + new_y1_mask = new_y1 < zy_numel + tl.store(output_ptr + (new_x2 + sum_numel * new_y1), tmp8, new_x2_mask & new_y1_mask) + + +testlist = [ + # ===================== 1D场景（15组，dim=0） ===================== + ((3, ), (3, ), 0), + ((7, ), (9, ), 0), + ((13, ), (11, ), 0), + ((2047, ), (2047, ), 0), + ((2701, ), (3003, ), 0), + ((4093, ), (3095, ), 0), + + # ===================== 2D场景（20组，dim0/dim1） ===================== + # dim0（行拼接，列维度一致） + ((3, 5), (3, 5), 0), + ((1005, 300), (2007, 300), 0), + ((1307, 400), (309, 400), 0), + ((303, 500), (303, 500), 0), + # dim1（列拼接，行维度一致） + ((7, 9), (7, 9), 1), + ((100, 1001), (100, 2003), 1), + ((200, 2005), (200, 207), 1), + ((300, 707), (300, 707), 1), + + # ===================== 3D场景（15组，dim0/dim1/dim2） ===================== + # dim0（第0维拼接，d1/d2一致） + ((378, 200, 300), (101, 200, 300), 0), + ((378, 70, 50), (601, 70, 50), 0), + # dim1（第1维拼接，d0/d2一致） + ((100, 452, 300), (100, 201, 300), 1), + ((65, 1735, 57), (65, 2001, 57), 1), + # dim2（第2维拼接，d0/d1一致） + ((87, 200, 387), (87, 200, 501), 2), + ((20, 337, 543), (20, 337, 401), 2), +] + + +@pytest.mark.parametrize('testlists', testlist) +@pytest.mark.parametrize('dtype', ['bool', 'int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16']) +def test_cat_bigshape(testlists, dtype): + torch_dtype = eval('torch.' + dtype) + np_x0 = test_common.generate_numpy(testlists[0], dtype) + np_x1 = test_common.generate_numpy(testlists[1], dtype) + cat_dim = testlists[2] + + x0 = torch.from_numpy(np_x0).to(torch_dtype).npu() + x1 = torch.from_numpy(np_x1).to(torch_dtype).npu() + + if len(x0.shape) > 3: + pytest.skip("dim > 3 for 3D+ tensor, skipping.") + + torch_res = torch.cat([x0, x1], dim=cat_dim) + triton_res = torch.zeros_like(torch_res) + num_core = 32 + if len(x0.shape) == 3: + if cat_dim == 0: + ZBLOCK = (min(x0.shape[0], x1.shape[0]) + num_core - 1) // num_core + ZBLOCK_2 = (max(x0.shape[0], x1.shape[0]) - min(x0.shape[0], x1.shape[0]) + num_core - 1) // num_core + triton_unk_fused_cat_3d_dim0[num_core, 1, 1](triton_res, x0, x1, x0.shape[0], x1.shape[0], x0.shape[1], + x0.shape[2], ZBLOCK, ZBLOCK_2, 1, 256) + elif cat_dim == 1: + Z0BLOCK = (x0.shape[0] + num_core - 1) // num_core + triton_unk_fused_cat_3d_dim1[num_core, 1, 1](triton_res, x0, x1, x0.shape[0], x0.shape[1], x1.shape[1], + x1.shape[2], Z0BLOCK, 1, 256) + else: + Y0BLOCK = (x0.shape[0] * x0.shape[1] + num_core - 1) // num_core + triton_unk_fused_cat_3d_dim2[num_core, 1, 1](triton_res, x0, x1, x0.shape[0], x0.shape[1], x0.shape[2], + x1.shape[2], Y0BLOCK, 1, 256) + test_common.validate_cmp(dtype, torch_res, triton_res) + return + numel_large = torch_res.numel() > 512 and len(x0.shape) < 3 + if numel_large or (cat_dim == 0 and len(x0.shape) == 2): + squeeze_flag = False + if len(x0.shape) == 1: + squeeze_flag = True + x0 = torch.unsqueeze(x0, dim=0) + x1 = torch.unsqueeze(x1, dim=0) + triton_res = torch.unsqueeze(triton_res, dim=0) + cat_dim = 1 + if cat_dim == 1: + Y0BLOCK = (x0.shape[0] + num_core - 1) // num_core + if x0.shape[1] == x1.shape[1]: + triton_unk_fused_cat_dim1_sameshape[num_core, 1, 1](triton_res, x0, x1, x0.shape[0], x0.shape[1], + Y0BLOCK, 1, 256) + else: + triton_unk_fused_cat_dim1_diffshape[num_core, 1, 1](triton_res, x0, x1, x0.shape[0], x0.shape[1], + x1.shape[1], Y0BLOCK, 1, 256) + else: + if x0.shape[0] == x1.shape[0]: + Y0BLOCK = (x0.shape[0] + num_core - 1) // num_core + triton_unk_fused_cat_dim0_sameshape[num_core, 1, 1](triton_res, x0, x1, x0.shape[0], x0.shape[1], + Y0BLOCK, 1, 256) + else: + YBLOCK = (min(x0.shape[0], x1.shape[0]) + num_core - 1) // num_core + YBLOCK_2 = (max(x0.shape[0], x1.shape[0]) - min(x0.shape[0], x1.shape[0]) + num_core - 1) // num_core + triton_unk_fused_cat_dim0_diffshape[num_core, 1, 1](triton_res, x0, x1, x0.shape[0], x1.shape[0], + x1.shape[1], YBLOCK, YBLOCK_2, 1, 256) + if squeeze_flag: + triton_res = triton_res.squeeze() + else: + squeeze_flag = False + if len(x0.shape) == 1: + squeeze_flag = True + x0 = torch.unsqueeze(x0, dim=0) + x1 = torch.unsqueeze(x1, dim=0) + triton_res = torch.unsqueeze(triton_res, dim=0) + _cat_helper_func_2D_1[num_core, 1, 1](x0, x1, triton_res, x0.shape[1], x1.shape[1], x0.shape[0], + x0.shape[1] + x1.shape[1], 256, 16) + if squeeze_flag: + triton_res = triton_res.squeeze() + + test_common.validate_cmp(dtype, torch_res, triton_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_cdiv.py b/third_party/ascend/unittest/pytest_ut/test_cdiv.py new file mode 100644 index 0000000000..f47a64eeea --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_cdiv.py @@ -0,0 +1,113 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common + + +def torch_cdiv(x0, x1, dtype_x): + if dtype_x in [torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64]: + return torch.div(x0, x1, rounding_mode='floor') + (x0 % x1 != 0).to(torch.int) + else: + if dtype_x in ["float16", "bfloat16"]: + x0 = x0.to(torch.float32) + x1 = x1.to(torch.float32) + return torch.ceil(x0 / x1).to(eval("torch." + dtype_x)) + + +def torch_cdiv_special(x0, x1, dtype_x): + if dtype_x in [torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64]: + return torch.div(x0, x1, rounding_mode='floor') + (x0 % x1 != 0).to(torch.int) + else: + if dtype_x in ["float16", "bfloat16"]: + x0 = x0.to(torch.float32) + return torch.ceil(x0 / x1).to(eval("torch." + dtype_x)) + + +@triton.jit +def triton_cdiv(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = tl.cdiv(XBLOCK, XBLOCK_SUB) + for loop1 in range(loops1): + x_index = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + x_index, None) + tmp1 = tl.load(in_ptr1 + x_index, None) + tmp2 = tl.cdiv(tmp0, tmp1) + tl.store(out_ptr0 + x_index, tmp2, None) + + +@triton.jit +def triton_cdiv_special(in_ptr0, div, out_ptr0, XBLOCK, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = tl.cdiv(XBLOCK, XBLOCK_SUB) + for loop1 in range(loops1): + x_index = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + x_index, None) + tmp2 = tl.cdiv(tmp0, div) + tl.store(out_ptr0 + x_index, tmp2, None) + + +param_lists = [ + # ['int8', (4096,), 1, 4096, 4096], + # ['int16', (4096,), 1, 4096, 4096], + ['int32', (4096, ), 1, 4096, 4096], + # ['int64', (4096,), 1, 4096, 4096], + # ['float16', (4096,), 1, 4096, 4096], + ['float32', (4096, ), 1, 4096, 4096], + # ['bfloat16', (4096,), 1, 4096, 4096], +] + + +@pytest.mark.parametrize('param_list', param_lists) +def test_cdiv(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + torch_dtype = eval('torch.' + dtype) + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + torch_res = torch_cdiv(x0, x1, dtype) + # triton结果 + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_cdiv[ncore, 1, 1](x0, x1, triton_res, xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) + + +@pytest.mark.parametrize('param_list', param_lists) +def test_cdiv_special(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + torch_dtype = eval('torch.' + dtype) + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = 2 + # torch结果 + torch_res = torch_cdiv_special(x0, x1, dtype) + # triton结果 + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_cdiv_special[ncore, 1, 1](x0, x1, triton_res, xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_ceil.py b/third_party/ascend/unittest/pytest_ut/test_ceil.py new file mode 100644 index 0000000000..faecd4fa8d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_ceil.py @@ -0,0 +1,76 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time + +import torch +import torch_npu +import test_common + + +def torch_ceil(x0): + if x0.dtype in [ + torch.int8, torch.int16, torch.int32, torch.int64, torch.uint8, torch.uint16, torch.uint32, torch.uint64 + ]: + return x0 + res = torch.ceil(x0) + return res + + +@triton.jit +def triton_ceil(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.ceil(tmp0) + tl.store(out_ptr0 + (x0), tmp1, None) + + +@pytest.mark.parametrize('param_list', [ + ['float16', (2, 4096, 8), 32, 2048, 64], + # ['bfloat16', (2, 4096, 8), 32, 2048, 64], + ['float32', (2, 4096, 8), 32, 2048, 64], + ['int8', (2, 4096, 8), 32, 2048, 64], + # ['int16', (2, 4096, 8), 32, 2048, 64], + # ['int32', (2, 4096, 8), 32, 2048, 64], + # ['int64', (2, 4096, 8), 32, 2048, 64], + ['uint8', (2, 4096, 8), 32, 2048, 64], + # ['uint16', (2, 4096, 8), 32, 2048, 64], + # ['uint32', (2, 4096, 8), 32, 2048, 64], + # ['uint64', (2, 4096, 8), 32, 2048, 64], +]) +def test_ceil(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + np_x0 = test_common.generate_numpy(shape, dtype) + x0 = torch.from_numpy(np_x0).to(eval('torch.' + dtype)).npu() + y_ref = torch_ceil(x0) + + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + x0 = x0.npu() + triton_ceil[ncore, 1, 1](x0, y_cal, xblock, xblock_sub, debug=True) + y_ref = y_ref.npu() + test_common.validate_cmp_with_expection(dtype, y_cal, y_ref, True) diff --git a/third_party/ascend/unittest/pytest_ut/test_celoss_indices.py b/third_party/ascend/unittest/pytest_ut/test_celoss_indices.py new file mode 100644 index 0000000000..da9f76e796 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_celoss_indices.py @@ -0,0 +1,112 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import triton +import triton.language as tl + + +@triton.jit +def celoss_indices_kernel( + inp_ptr, + tgt_ptr, + w_ptr, + out_ptr, + w_tgt_ptr, + ignore_index, + C, + D, + BLOCK_C: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_d = tl.program_id(0).to(tl.int64) + pid_n = tl.program_id(1).to(tl.int64) + offset_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D).to(tl.int64) + + tgt_ptrs = tgt_ptr + pid_n * D + offset_d + tgt_mask = offset_d < D + tgt = tl.load(tgt_ptrs, mask=tgt_mask, other=0) + + ignore_mask = not (tgt == ignore_index) and tgt_mask + + tmp_max = tl.zeros([BLOCK_C, BLOCK_D], dtype=tl.float32) + tmp_sum = tl.zeros([BLOCK_C, BLOCK_D], dtype=tl.float32) + + for off in range(0, C, BLOCK_C): + offset_c = off + tl.arange(0, BLOCK_C) + inp_ptrs = inp_ptr + pid_n * C * D + offset_c[:, None] * D + offset_d[None, :] + inp_mask = offset_c[:, None] < C and offset_d[None, :] < D + inp = tl.load(inp_ptrs, inp_mask, other=-float("inf")).to(tl.float32) + cur_max = tl.maximum(tmp_max, inp) + cur_exp = tl.exp(inp - cur_max) + tmp_sum = tmp_sum * tl.exp(tmp_max - cur_max) + cur_exp + tmp_max = cur_max + + final_max = tl.max(tmp_max, axis=0) + tmp_sum = tmp_sum * tl.exp(tmp_max - final_max[None, :]) + final_sum = tl.log(tl.sum(tmp_sum, axis=0)) + inp_tgt_ptrs = inp_ptr + pid_n * C * D + tgt * D + offset_d + inp_tgt = tl.load(inp_tgt_ptrs, mask=tgt_mask, other=-float("inf")).to(tl.float32) + + out = final_sum + final_max - inp_tgt + w_tgt_ptrs = w_tgt_ptr + pid_n * D + offset_d + + if w_ptr is None: + w_tgt = ignore_mask + else: + w_tgt = tl.load(w_ptr + tgt, mask=tgt_mask, other=0).to(tl.float32) + w_tgt = tl.where(ignore_mask, w_tgt, 0) + + tl.store(w_tgt_ptrs, w_tgt, mask=tgt_mask) + out *= w_tgt + out_ptrs = out_ptr + pid_n * D + offset_d + tl.store(out_ptrs, out, mask=tgt_mask) + + +def test_celoss_indices_kernel(shape=(1, 2)): + device = "npu" + dtype = torch.float16 + ignore_index = -100 + BLOCK_C = 256 + BLOCK_D = 1 + + N, C = shape + D = 1 + + inp = torch.randn(shape, dtype=dtype, device=device) + tgt = torch.randint(0, C, (N, ), dtype=torch.int64, device=device) + wgt = torch.randn(C, dtype=dtype, device=device) + + out_triton = torch.empty((N * D, ), dtype=torch.float32, device=device) + w_tgt_triton = torch.empty((N * D, ), dtype=torch.float32, device=device) + + grid = (triton.cdiv(D, BLOCK_D), N) + celoss_indices_kernel[grid]( + inp, + tgt, + wgt, + out_triton, + w_tgt_triton, + ignore_index, + C, + D, + BLOCK_C=BLOCK_C, + BLOCK_D=BLOCK_D, + ) diff --git a/third_party/ascend/unittest/pytest_ut/test_clamp.py b/third_party/ascend/unittest/pytest_ut/test_clamp.py new file mode 100644 index 0000000000..426bbd01db --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_clamp.py @@ -0,0 +1,62 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time +import torch +import torch_npu +import test_common + + +def torch_clamp_float(x0): + res = torch.clamp(x0, 0.0, 100.0) + return res + + +@triton.jit +def triton_clamp_float(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.clamp(tmp0, 0.0, 100.0) + tl.store(out_ptr0 + (x0), tmp1, None) + + +@pytest.mark.parametrize('param_list', [ + # int原生不支持 + ['float16', (4, 4), 4, 4, 4], + ['float32', (4, 4), 4, 4, 4], +]) +def test_clamp(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype) + y_ref = torch_clamp_float(x0) + tyname = test_common.get_triton_sig_typename(dtype) + + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + x0 = x0.npu() + triton_clamp_float[ncore, 1, 1](x0, y_cal, xblock, xblock_sub, debug=True) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_common.py b/third_party/ascend/unittest/pytest_ut/test_common.py new file mode 100644 index 0000000000..3a38ae9b66 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_common.py @@ -0,0 +1,223 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +from typing import Optional +import torch +import torch_npu +import pytest +import functools +import re +import numpy as np + +_float_dtypes = ['float32', 'float16', 'bfloat16'] +_int_dtypes = ['int32', 'int64', 'int16', 'int8'] +_uint_dtypes = ['uint8', 'uint16', 'uint32', 'uint64'] +_all_dtypes_no_bool = _float_dtypes + _int_dtypes +_all_dtypes = _all_dtypes_no_bool + ['bool'] +_32bit_dtypes = ['float32', 'int32'] +_16bit_dtypes = ['float16', 'bfloat16', 'int16'] + + +def generate_numpy(shape, dtype, low=None, high=None): + if dtype in _int_dtypes + _uint_dtypes: + iinfo = np.iinfo(getattr(np, dtype)) + low = iinfo.min if low is None else max(low, iinfo.min) + high = iinfo.max if high is None else min(high, iinfo.max) + dty = getattr(np, dtype) + return np.random.randint(low, high, shape, dtype=dty) + elif dtype == 'float16' or dtype == 'float32': + return np.random.normal(0, 1, shape).astype(dtype) + elif dtype == 'bfloat16': + return (np.random.normal(0, 1, shape).astype('float32').view('uint32') & np.uint32(0xffff0000)).view('float32') + elif dtype == 'bool': + return np.random.randint(low=0, high=2, size=shape).astype(bool) + else: + raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + + +def generate_tensor(shape, dtype): + if dtype == 'float32' or dtype == 'float16' or dtype == 'bfloat16': + return torch.randn(size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16': + return torch.randint(low=0, high=2000, size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'int8': + return torch.randint(low=0, high=127, size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'bool': + return torch.randint(low=0, high=2, size=shape).bool() + elif dtype == 'uint8': + return torch.randint(low=0, high=255, size=shape, dtype=torch.uint8) + else: + raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + + +def get_triton_sig_typename(dtype): + if dtype == 'float32': + tyname = "*fp32" + elif dtype == 'int32': + tyname = "*i32" + elif dtype == 'int64': + tyname = "*i64" + elif dtype == 'float16': + tyname = "*fp16" + elif dtype == 'int16': + tyname = "*i16" + elif dtype == 'int8': + tyname = "*i8" + elif dtype == 'bool': + tyname = "*i1" + else: + raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + return tyname + + +# Relative error: abs(x_ref - x_cal) / abs(x_ref) +# Absolute error: abs(x_ref - x_cal) + + +# calculation type operators require different error range +# It is a stricter verification and not satisfied now, save it here +def validate_cal(dtype, y_cal, y_ref): + if dtype == 'float16': + if torch.mean(y_ref) < 0.001: + assert torch.abs(y_cal - y_ref) < 0.001, "|y_cal - y_ref| < 0.001 is required !" + else: + diff = torch.div(torch.abs(y_cal - y_ref), torch.abs(y_cal)) < 0.001 + # all true + assert diff.all(), "Relative error is less than 0.001 !" + if dtype == 'float32': + if torch.mean(y_ref) < 0.0001: + assert torch.abs(y_cal - y_ref) < 0.0001, "|y_cal - y_ref| < 0.0001 is required !" + else: + diff = torch.div(torch.abs(y_cal - y_ref), torch.abs(y_cal)) < 0.0001 + assert diff.all(), "Relative error is less than 0.001 !" + elif dtype == 'bfloat16': + diff = torch.div(torch.abs(y_cal - y_ref), torch.abs(y_cal)) < 0.001 + assert diff.all(), "Relative error is less than 0.001 !" + elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16' or dtype == 'int8': + assert torch.equal(y_cal, y_ref) + elif dtype == 'uint8' or dtype == 'uint16' or dtype == 'uint32' or dtype == 'uint64': + assert torch.equal(y_cal, y_ref) + elif dtype == 'bool': + assert torch.equal(y_cal, y_ref) + else: + raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + + +# moving and comparison ops require no precision error +def validate_cmp(dtype, y_cal, y_ref, overflow_mode: Optional[str] = None): + y_cal = y_cal.npu() + y_ref = y_ref.npu() + if overflow_mode == "saturate": + if dtype in ['float32', 'float16']: + min_value = -torch.finfo(dtype).min + max_value = torch.finfo(dtype).max + elif dtype in ['int32', 'int16', 'int8']: + min_value = torch.iinfo(dtype).min + max_value = torch.iinfo(dtype).max + elif dtype == 'bool': + min_value = 0 + max_value = 1 + else: + raise ValueError('Invalid parameter "dtype" is found : {}'.format(dtype)) + y_ref = torch.clamp(y_ref, min=min_value, max=max_value) + if dtype == 'float16': + torch.testing.assert_close(y_ref, y_cal, rtol=1e-03, atol=1e-03, equal_nan=True) + elif dtype == 'bfloat16': + torch.testing.assert_close(y_ref.to(torch.float32), y_cal.to(torch.float32), rtol=1e-03, atol=1e-03, + equal_nan=True) + elif dtype == 'float32': + torch.testing.assert_close(y_ref, y_cal, rtol=1e-04, atol=1e-04, equal_nan=True) + elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16' or dtype == 'int8': + assert torch.equal(y_cal, y_ref) + elif dtype == 'uint8' or dtype == 'uint16' or dtype == 'uint32' or dtype == 'uint64': + assert torch.equal(y_cal, y_ref) + elif dtype == 'bool': + assert torch.equal(y_cal, y_ref) + else: + raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + + +def validate_cmp_with_expection(dtype, y_cal, y_ref, expect): + if dtype == 'float32' or dtype == 'float16' or dtype == 'bfloat16': + if expect: + assert torch.allclose(y_ref, y_cal, rtol=1e-03, atol=1e-03, equal_nan=True) + else: + assert not torch.allclose(y_ref, y_cal, rtol=1e-03, atol=1e-03, equal_nan=True) + elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16' or dtype == 'int8' \ + or dtype == 'uint8' or dtype == 'uint16' or dtype == 'uint32' or dtype == 'uint64': + if expect: + assert torch.equal(y_cal, y_ref) + else: + assert not torch.equal(y_cal, y_ref) + else: + raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + + +# Use the following pytest fixture to run one test case by only single worker. +# Refer to https://pytest-xdist.readthedocs.io/en/stable/how-to.html#making-session-scoped-fixtures-execute-only-once +@pytest.fixture(scope="function") +def pytest_runonce(worker_id, request, cache): + if (cache.get(request.node.nodeid, "none")) == "none": + cache.set(request.node.nodeid, worker_id) + else: + file_name = f"pytest_{worker_id}.txt" + with open(file_name, 'a') as file: + file.write(f"{request.node.nodeid} is already processed by {worker_id}") + return True + yield True + cache.set(request.node.nodeid, "none") + + +def raises_with_match(expected_exception, match_pattern): + + def decorator(test_func): + + @functools.wraps(test_func) + def wrapper(*args, **kwargs): + with pytest.raises(expected_exception, match=match_pattern): + return test_func(*args, **kwargs) + + return wrapper + + return decorator + + +def capture_output(expected_output): + + def decorator(test_func): + + @functools.wraps(test_func) + def wrapper(*args, **kwargs): + capsys = kwargs.pop('capsys', None) + if capsys is None: + try: + capsys = pytest.fixture(capsys)() + except: + raise RuntimeError("This decorator requires pytest's capsys fixture") + test_func(capsys, *args, **kwargs) + captured = capsys.readouterr() + # pybind11::scoped_ostream_redirect captures std::cout with \x00 inserted + # for now, no idea how to eliminate \x00 from C++ side. + cleaned = re.sub(r"\x00", "", captured.out) + assert expected_output in cleaned + + return wrapper + + return decorator diff --git a/third_party/ascend/unittest/pytest_ut/test_compile_hint.py b/third_party/ascend/unittest/pytest_ut/test_compile_hint.py new file mode 100644 index 0000000000..4582201936 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_compile_hint.py @@ -0,0 +1,58 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import triton.language.extra.cann.extension as extension +import pytest +import test_common + +# eg: pytest -v test.py::test_compile_hint +############################# + + +@triton.jit +def triton_compile_hint(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + xindex = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = xindex < xnumel + x0 = xindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + extension.compile_hint(tmp0, "hint_a") + extension.multibuffer(tmp0, 2) + tmp2 = tmp0 + extension.compile_hint(tmp2, "hint_b", 42) + extension.compile_hint(tmp2, "hint_c", True) + extension.compile_hint(tmp2, "hint_d", [XBLOCK, XBLOCK_SUB]) + tl.store(out_ptr0 + (xindex), tmp2, xmask) + + +@pytest.mark.skip(reason="not supported after the NPUIR is updated in April, and will be fixed later") +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_compile_hint(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + y_ref = x0 + y_cal = test_common.generate_tensor(shape, dtype).npu() + triton_compile_hint[(ncore, )](x0, y_cal, x0.numel(), xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_complex_mask.py b/third_party/ascend/unittest/pytest_ut/test_complex_mask.py new file mode 100644 index 0000000000..8a89039e4d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_complex_mask.py @@ -0,0 +1,69 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + + +def copy(x): + return x.clone() + + +@triton.jit +def copy_kernel(in_ptr, out_ptr, N: tl.constexpr, NUMEL): + idx_block = tl.arange(0, N) + is_valid = N <= NUMEL + x = tl.load(in_ptr + idx_block, mask=idx_block < N) + mask_i1 = is_valid[None, None] & (idx_block < N)[None, :] + tl.store(out_ptr + idx_block[None, :], x[None, :], mask=mask_i1) + + +@triton.jit +def permute_copy_kernel(in_ptr, out_ptr, M: tl.constexpr, N: tl.constexpr, NUMEL): + idx_block_n = tl.arange(0, N) + idx_block_m = tl.arange(0, M) + idx_block = idx_block_m[:, None] + idx_block_n[None, :] * M + is_valid = N <= NUMEL + x = tl.load(in_ptr + idx_block, mask=(idx_block_m[:, None] < M) & (idx_block_n[None, :] < N)) + mask_i1 = (is_valid[None, None, None]) & (idx_block_m[None, :, None] < M) & (idx_block_n[None, None, :] < N) + tl.store(out_ptr + idx_block[None, :], x[None, :], mask=mask_i1) + + +def test_complex_mask_copy(): + N = 1024 + x = torch.randn(N, dtype=torch.float32).npu() + y = torch.empty_like(x).npu() + copy_kernel[(1, )](x, y, N=N, NUMEL=N) + torch.testing.assert_close(x, y) + + +def test_complex_mask_permute_copy(): + M = 4 + N = 32 + x = torch.randn(M * N, dtype=torch.float32).npu() + y = torch.empty_like(x).npu() + permute_copy_kernel[(1, )](x, y, M=M, N=N, NUMEL=M * N) + torch.testing.assert_close(x, y) diff --git a/third_party/ascend/unittest/pytest_ut/test_copy.py b/third_party/ascend/unittest/pytest_ut/test_copy.py new file mode 100644 index 0000000000..4cf9daa322 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_copy.py @@ -0,0 +1,97 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import os +import pytest +import triton +import triton.language as tl +import triton.extension.buffer.language as bl +import triton.language.extra.cann.extension as al +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +from triton._C.libtriton import ir, buffer_ir +from triton._C.libtriton.ascend import ir as ascend_ir + +os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" + + +class Options: + num_warps = 4 + num_stages = 3 + num_ctas = 1 + cluster_dims = (1, 1, 1) + enable_fp_fusion = True + debug = False + arch = "Ascend910_95" + + +def compile_kernel(kernel, signature, constants): + """Helper to compile a kernel to MLIR.""" + src = ASTSource(kernel, signature, constants) + context = ir.context() + ir.load_dialects(context) + buffer_ir.load_dialects(context) + ascend_ir.load_dialects(context) + module = ast_to_ttir(kernel, src, context, Options(), {}, {}) + return str(module) + + +@triton.jit +def copy( + A_ptr, + A1_ptr, + M: tl.constexpr, + N: tl.constexpr, +): + offs_a = tl.arange(0, M)[:, None] + offs_b = tl.arange(0, N)[None, :] + + offs_c = (offs_a) * M + (offs_b) + a_ptr = A_ptr + offs_c + a_val = tl.load(a_ptr) + a1_ptr = A1_ptr + offs_c + a1_val = tl.load(a1_ptr) + + add = tl.add(a_val, a1_val) + add_ub = bl.to_buffer(add, al.ascend_address_space.UB) + + A_l1 = bl.alloc(tl.float32, [M, N], al.ascend_address_space.L1) + al.copy_from_ub_to_l1(add_ub, A_l1) + + A_ub = bl.alloc(tl.float32, [M, N], al.ascend_address_space.UB) + al.copy(add_ub, A_ub) + + +def test_copy(): + print("=" * 60) + print("Test 1: copy ") + print("=" * 60) + mlir = compile_kernel( + copy, + {"A_ptr": "*fp32", "A1_ptr": "*fp32"}, + {"M": 16, "N": 16}, + ) + print(f"✅ Generated MLIR ({len(mlir)} chars):\n") + print(mlir) + + +# ============== Main for manual testing ============== +if __name__ == "__main__": + test_copy() diff --git a/third_party/ascend/unittest/pytest_ut/test_copysign.py b/third_party/ascend/unittest/pytest_ut/test_copysign.py new file mode 100644 index 0000000000..40aa235f2a --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_copysign.py @@ -0,0 +1,108 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import torch +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + + +@triton.jit +def copysign_kernel(x_ptr, y_ptr, z_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + + mask = offsets < n_elements + + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + + z = libdevice.copysign(x, y) + + tl.store(z_ptr + offsets, z, mask=mask) + + +@pytest.mark.parametrize('shape', [ + (12, 16), +]) +@pytest.mark.parametrize('dtype', ['float32']) +def test_copysign(shape, dtype): + n_elements = shape[0] * shape[1] + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + + # Ensure to include some boundary cases + x[0, 0] = 3.14 + y[0, 0] = 1.0 # The result should be 3.14 + + x[0, 1] = 3.14 + y[0, 1] = -1.0 # The result should be -3.14 + + x[0, 2] = -3.14 + y[0, 2] = 1.0 # The result should be 3.14 + + x[0, 3] = -3.14 + y[0, 3] = -1.0 # The result should be -3.14 + + x[0, 4] = 0.0 + y[0, 4] = -1.0 # The result should be -0.0 + + x[0, 5] = 0.0 + y[0, 5] = 1.0 # The result should be 0.0 + + x[0, 6] = 3.14 + y[0, 6] = 0.0 # The result should be 3.14 + + x[0, 7] = 3.14 + y[0, 7] = -0.0 # The result should be -3.14 + + x[0, 8] = -3.14 + y[0, 8] = 0.0 # The result should be 3.14 + + x[0, 9] = -3.14 + y[0, 9] = -0.0 # The result should be -3.14 + + x[0, 10] = 0.0 + y[0, 10] = 0.0 # The result should be 0.0 + + x[0, 11] = 0.0 + y[0, 11] = -0.0 # The result should be -0.0 + + x[0, 12] = -0.0 + y[0, 12] = 0.0 # The result should be 0.0 + + x[0, 13] = -0.0 + y[0, 13] = -0.0 # The result should be -0.0 + + z = torch.empty_like(x) + + BLOCK_SIZE = 192 + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + + copysign_kernel[grid](x, y, z, n_elements, BLOCK_SIZE=BLOCK_SIZE) + + expected = torch.copysign(x, y) + + torch.testing.assert_close(z, expected, rtol=1e-3, atol=1e-3) + + print("✓ COPYSIGN test PASSED!") diff --git a/third_party/ascend/unittest/pytest_ut/test_cos.py b/third_party/ascend/unittest/pytest_ut/test_cos.py new file mode 100644 index 0000000000..7c88852df5 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_cos.py @@ -0,0 +1,99 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + + +def standard_unary(x0, dtype): + res = torch.cos(x0) + return res + + +def standard_binary(x0, y0, dtype): + res = x0 + y0 + return res + + +@triton.jit +def triton_elementwise_unary(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + ret = tl.cos(x) + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +@triton.jit +def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + y = tl.load(in_ptr1 + idx_block, mask=idx_block < N) + ret = x + y + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +types = [ + (torch.float32, 'float32'), + # (torch.float16, 'float16'), + # (torch.bfloat16, 'bfloat16'), + # (torch.int8, 'int8'), + # (torch.int16, 'int16'), + # (torch.int32, 'int32'), + # (torch.int64, 'int64'), +] + +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('N,NUMEL', shapes) +def test_elementwsie_common(dtype, sigtype, N, NUMEL): + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == "int64": + N = map_for_64_t[N] if N in map_for_64_t else N + + print(f"elementwise : ({N},) {dtype} {sigtype}") + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standard_unary(x0, dtype) + x0 = x0.npu() + print(ans) + + out = torch.zeros((N, ), dtype=dtype).npu() + triton_elementwise_unary[1, 1, 1](x0, out, N=N, NUMEL=NUMEL, debug=True) + print(out) + + test_common.validate_cmp(sigtype, out, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_cos_2.py b/third_party/ascend/unittest/pytest_ut/test_cos_2.py new file mode 100644 index 0000000000..b7de878120 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_cos_2.py @@ -0,0 +1,60 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time + +import torch +import torch_npu +import test_common + + +def torch_cos(x0): + res = torch.cos(x0) + return res + + +@triton.jit +def triton_cos(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.cos(tmp0) + tl.store(out_ptr0 + (x0), tmp1, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 32, 2048, 64], +]) +def test_cos(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype) + y_ref = torch_cos(x0) + tyname = test_common.get_triton_sig_typename(dtype) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + x0 = x0.npu() + triton_cos[ncore, 1, 1](x0, y_cal, xblock, xblock_sub, debug=True) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_cosh.py b/third_party/ascend/unittest/pytest_ut/test_cosh.py new file mode 100644 index 0000000000..ee35268b6a --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_cosh.py @@ -0,0 +1,59 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + +import torch +import torch_npu + + +@triton.jit +def triton_cosh(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = libdevice.cosh(tmp0) + tl.store(out_ptr0 + (x0), tmp1, None) + + +@pytest.mark.parametrize('param_list', ['float32', 'float16', 'bfloat16']) +def test_cosh_special(param_list): + dtype = param_list + x_near_zero = torch.linspace(-1.0, 1.0, 256, dtype=eval("torch." + dtype)).npu() + x_medium_neg = torch.linspace(-10.0, -1.0, 128, dtype=eval("torch." + dtype)).npu() + x_medium_pos = torch.linspace(1.0, 10.0, 128, dtype=eval("torch." + dtype)).npu() + + x_large_neg = torch.linspace(-1e4, -10.0, 64, dtype=eval("torch." + dtype)).npu() + x_large_pos = torch.linspace(10.0, 1e4, 64, dtype=eval("torch." + dtype)).npu() + + x0 = torch.cat([x_near_zero, x_medium_neg, x_medium_pos, x_large_neg, x_large_pos], dim=0) + + y_ref = torch.cosh(x0) + y_cal = torch.zeros_like(y_ref) + triton_cosh[1, 1, 1](x0, y_cal, x0.shape[0], x0.shape[0]) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_cumprod.py b/third_party/ascend/unittest/pytest_ut/test_cumprod.py new file mode 100644 index 0000000000..feed59b7c4 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_cumprod.py @@ -0,0 +1,97 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl +from triton.runtime.libentry import libentry + +from test_common import _all_dtypes_no_bool, validate_cmp + + +def torch_func(x, dim, reverse): + is_bf16 = x.dtype == torch.bfloat16 + if is_bf16: + x = x.to(torch.float32) + if reverse: + x = torch.flip(x, [dim]) + res = torch.cumprod(x, dim=dim) + if is_bf16: + res = res.to(torch.bfloat16) + if reverse: + res = torch.flip(res, [dim]) + return res + + +@libentry() +@triton.jit +def triton_kernel( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + numel_x: tl.constexpr, + numel_r: tl.constexpr, + XBLOCK: tl.constexpr, + RBLOCK: tl.constexpr, +): + tl.static_assert(numel_x == XBLOCK, "numel_x must be equal to XBLOCK in this kernel") + tl.static_assert(numel_r == RBLOCK, "numel_r must be equal to RBLOCK in this kernel") + idx_x = tl.arange(0, XBLOCK) + idx_r = tl.arange(0, RBLOCK) + idx = idx_x[:, None] * numel_r + idx_r[None, :] + x = tl.load(in_ptr0 + idx) + ret = tl.cumprod(x, axis=dim, reverse=reverse) + tl.store(out_ptr0 + idx, ret) + + +def triton_func(x, dim, reverse): + res = torch.empty_like(x) + triton_kernel[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[1], x.shape[0], x.shape[1]) + return res + + +def cumprod_generate_tensor(shape, dtype): + if dtype == 'float32' or dtype == 'float16' or dtype == 'bfloat16': + return torch.rand(size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16': + return torch.randint(low=0, high=3, size=shape, dtype=eval('torch.' + dtype)) + elif dtype == 'int8': + return torch.randint(low=0, high=3, size=shape, dtype=eval('torch.' + dtype)) + else: + raise ValueError(f"Unsupported dtype: {dtype}") + + +# dtype=int8, reverse=True not support; +not_support_dtype = {'int8', 'bool'} +support_dtypes = [dtype for dtype in _all_dtypes_no_bool if dtype not in not_support_dtype] + + +@pytest.mark.parametrize("dtype", support_dtypes) +@pytest.mark.parametrize("shape", [(7, 23)]) +@pytest.mark.parametrize("dim", [0, 1]) +@pytest.mark.parametrize("reverse", [False, True]) +def test_cumprod(dtype, shape, dim, reverse): + x0 = cumprod_generate_tensor(shape=shape, dtype=dtype).npu() + triton_cal = triton_func(x0, dim, reverse) + torch_ref = torch_func(x0, dim, reverse) + validate_cmp(dtype, torch_ref, triton_cal) diff --git a/third_party/ascend/unittest/pytest_ut/test_cumsum.py b/third_party/ascend/unittest/pytest_ut/test_cumsum.py new file mode 100644 index 0000000000..c9196b6bf7 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_cumsum.py @@ -0,0 +1,84 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl +from triton.runtime.libentry import libentry + +from test_common import _all_dtypes_no_bool, generate_tensor, validate_cmp + + +def torch_func(x, dim, reverse): + if reverse: + x = torch.flip(x, [dim]) + res = torch.cumsum(x, dim=dim) + if reverse: + res = torch.flip(res, [dim]) + return res + + +@libentry() +@triton.jit +def triton_kernel( + out_ptr0, + in_ptr0, + dim: tl.constexpr, + reverse: tl.constexpr, + numel_x: tl.constexpr, + numel_r: tl.constexpr, + XBLOCK: tl.constexpr, + RBLOCK: tl.constexpr, +): + tl.static_assert(numel_x == XBLOCK, "numel_x must be equal to XBLOCK in this kernel") + tl.static_assert(numel_r == RBLOCK, "numel_r must be equal to RBLOCK in this kernel") + idx_x = tl.arange(0, XBLOCK) + idx_r = tl.arange(0, RBLOCK) + idx = idx_x[:, None] * numel_r + idx_r[None, :] + x = tl.load(in_ptr0 + idx) + ret = tl.cumsum(x, axis=dim, reverse=reverse) + tl.store(out_ptr0 + idx, ret) + + +def triton_func(x, dim, reverse): + res = torch.empty_like(x) + triton_kernel[1, 1, 1](res, x, dim, reverse, x.shape[0], x.shape[1], x.shape[0], x.shape[1]) + return res + + +# dtype=int8, reverse=True not support; +not_support_dtype = {'int8', 'bool'} +support_dtypes = [dtype for dtype in _all_dtypes_no_bool if dtype not in not_support_dtype] + + +@pytest.mark.parametrize("dtype", support_dtypes) +@pytest.mark.parametrize("shape", [(7, 23)]) +@pytest.mark.parametrize("dim", [0, 1]) +@pytest.mark.parametrize("reverse", [False, True]) +def test_cumsum(dtype, shape, dim, reverse): + x0 = generate_tensor(shape=shape, dtype=dtype).npu() + triton_cal = triton_func(x0, dim, reverse) + torch_dtype = eval('torch.' + dtype) + if torch_dtype == torch.float16 or torch_dtype == torch.float32: + x0 = x0.to(torch.float32) + torch_ref = torch_func(x0, dim, reverse).to(torch_dtype) + validate_cmp(dtype, torch_ref, triton_cal) diff --git a/third_party/ascend/unittest/pytest_ut/test_custom.py b/third_party/ascend/unittest/pytest_ut/test_custom.py new file mode 100755 index 0000000000..97c2a0ffe9 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_custom.py @@ -0,0 +1,299 @@ +#!/usr/bin/env python3 +import subprocess +import os +import triton +import triton.language as tl +import triton.language.extra.cann.extension as al +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir +from triton.backends.ascend.compiler import NPUOptions, ttir_to_linalg +import pytest + + +def compile_kernel(kernel, signature, constants): + """Helper to compile a kernel function to MLIR in linalg dialect.""" + src = ASTSource(kernel, signature, constants) + context = ir.context() + ir.load_dialects(context) + ascend_ir.load_dialects(context) + try: + options = NPUOptions() + ttir = ast_to_ttir(kernel, src, context, options, {}, {}) + metadata = { + **options.__dict__, + } + linalg = ttir_to_linalg(ttir, metadata, options, named_ops=True) + return str(linalg) + except subprocess.CalledProcessError as ex: + print(ex.stdout.decode()) + print(ex.stderr.decode()) + print("failed") + return None + + +# ============== Kernel definitions ============== + + +@al.register_custom_op +class my_custom_op: + core = al.CORE.VECTOR + pipe = al.PIPE.PIPE_V + mode = al.MODE.SIMT + symbol = "my_custom_func" + # fake path, this test only check Triton successfully lowered to MLIR + bitcode = os.path.abspath(__file__) + iterator_types = [ + al.IteratorType.Parallel, + al.IteratorType.Broadcast, + al.IteratorType.Transpose, + al.IteratorType.Reduction, + al.IteratorType.Interleave, + al.IteratorType.Deinterleave, + al.IteratorType.Inverse, + al.IteratorType.Pad, + al.IteratorType.Concat, + al.IteratorType.Gather, + al.IteratorType.Cumulative, + al.IteratorType.Opaque, + ] + + def __init__(self, x, ptr1, ptr2, offset: tl.int64, other, out=None): + # Add optional custom-op attribute: ArrayAttr. + self.indexing_map = [al.affine_map.get_identity(1)] + + # Tag ptr2 as an argument that should be aligned at dimension 1. + # Tag 2nd argument that should be aligned at dimension 0. + self.align_dim = {"ptr2": 1, 1: 0} + + +@triton.jit +def my_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + x = tl.load(x_ptr + i, mask=i < n) + y = tl.load(y_ptr + i, mask=i < n) + result = al.custom("my_custom_op", x, x_ptr, y_ptr + i, (1, 2, 3), [4.1, 5.2], out=y) + a = 123 + result = al.custom("my_custom_op", x, x_ptr, y_ptr, (a, n), (1.2, 3.4), out=result) + tl.store(out_ptr + i, result, mask=i < n) + + +@al.register_custom_op +class my_custom_op_extra_buf: + """Custom op declaring extra_buffers with several scalar Triton dtypes.""" + + core = al.CORE.VECTOR + pipe = al.PIPE.PIPE_V + mode = al.MODE.SIMT + symbol = "my_extra_buf_func" + bitcode = os.path.abspath(__file__) + + def __init__(self, x, out=None): + self.indexing_map = [al.affine_map.get_identity(1)] + self.extra_buffers = [ + (tl.bfloat16, 256), + (tl.float64, 424242), + (tl.int8, 11), + (tl.float16, 22), + (tl.int32, 33), + ] + + +@al.register_custom_op +class my_custom_op_extra_buf_single_buf: + """Custom op declaring extra_buffers with single buf.""" + + core = al.CORE.VECTOR + pipe = al.PIPE.PIPE_V + mode = al.MODE.SIMT + symbol = "my_extra_buf_func_single_buf" + bitcode = os.path.abspath(__file__) + + def __init__(self, x, out=None): + self.indexing_map = [al.affine_map.get_identity(1)] + self.extra_buffers = (tl.bfloat16, 256) + + +@triton.jit +def kernel_extra_buf(x_ptr, out_ptr, n, BLOCK: tl.constexpr): + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + x = tl.load(x_ptr + i, mask=i < n) + y = tl.load(out_ptr + i, mask=i < n) + r = al.custom("my_custom_op_extra_buf", x, out=y) + tl.store(out_ptr + i, r, mask=i < n) + + +@triton.jit +def kernel_extra_buf_single_buf(x_ptr, out_ptr, n, BLOCK: tl.constexpr): + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + x = tl.load(x_ptr + i, mask=i < n) + y = tl.load(out_ptr + i, mask=i < n) + r = al.custom("my_custom_op_extra_buf_single_buf", x, out=y) + tl.store(out_ptr + i, r, mask=i < n) + + +@al.register_custom_op +class my_custom_op_extra_buf_wide: + """Cover more integer widths and unsigned dtypes in extra_buffers_types.""" + + core = al.CORE.VECTOR + pipe = al.PIPE.PIPE_V + mode = al.MODE.SIMT + symbol = "my_extra_buf_wide_func" + bitcode = os.path.abspath(__file__) + + def __init__(self, x, out=None): + self.indexing_map = [al.affine_map.get_identity(1)] + self.extra_buffers = [ + (tl.int16, 1001), + (tl.uint16, 1002), + (tl.int64, 1003), + (tl.uint32, 1004), + (tl.uint8, 1005), + ] + + +@triton.jit +def kernel_extra_buf_wide(x_ptr, out_ptr, n, BLOCK: tl.constexpr): + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + x = tl.load(x_ptr + i, mask=i < n) + y = tl.load(out_ptr + i, mask=i < n) + r = al.custom("my_custom_op_extra_buf_wide", x, out=y) + tl.store(out_ptr + i, r, mask=i < n) + + +# ============== Pytest tests ============== + + +def test_custom_op(): + """Test custom op compile to linalg MLIR.""" + mlir = compile_kernel(my_kernel, {"x_ptr": "*fp32", "y_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, + {"BLOCK": 256}) + assert mlir and len(mlir) > 0 + assert "func.func @my_kernel(" in mlir + assert "hivm.hir.custom" in mlir + for line in mlir.splitlines(): + if "hivm.hir.custom" in line: + # custom op name + assert '"my_custom_op"' in line + # All tt.ptr converted to memref. + assert "tt.ptr" not in line + # Required attributes are set. + assert "hivm.pipe = #hivm.pipe" in line + assert "hivm.tcore_type = #hivm.tcore_type" in line + assert "hivm.vf_mode = #hivm.vf_mode" in line + # Optional indexing map attribute should be attached. + assert "indexing_map = [" in line + # Tagged argument alignment info is attached as integer operand attr. + assert "align_dim = 1" in line + assert "align_dim = 0" in line + # All offset converted to int64. + assert 'i64, ' in line + assert 'i32, ' not in line + assert "iterator_types" in line + for iterator_name in ( + "parallel", + "broadcast", + "transpose", + "reduction", + "interleave", + "deinterleave", + "inverse", + "pad", + "concat", + "gather", + "cumulative", + "opaque", + ): + assert iterator_name in line + + +def _custom_lines(mlir: str, op_name: str): + # Match the MLIR string attribute exactly (avoid `my_custom_op` matching + # `my_custom_op_extra_buf`). + quoted = f'"{op_name}"' + return [line for line in mlir.splitlines() if "hivm.hir.custom" in line and quoted in line] + + +def test_custom_op_extra_buffers_mixed_scalar_types(): + """extra_buffers_types must preserve bf16/f64/i8/f16/i32 (not all lowered to f32).""" + mlir = compile_kernel( + kernel_extra_buf, + {"x_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, + {"BLOCK": 256}, + ) + assert mlir and len(mlir) > 0 + lines = _custom_lines(mlir, "my_custom_op_extra_buf") + assert lines, "expected at least one hivm.hir.custom line for my_custom_op_extra_buf" + line = lines[0] + assert "extra_buffers_types" in line + assert "extra_buffers_sizes" in line + assert "bf16" in line + assert "f64" in line + assert "i8" in line + assert "f16" in line + assert "i32" in line + assert "424242" in line + + +def test_custom_op_extra_buffers_single_buffer(): + mlir = compile_kernel( + kernel_extra_buf_single_buf, + {"x_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, + {"BLOCK": 256}, + ) + assert mlir and len(mlir) > 0 + lines = _custom_lines(mlir, "my_custom_op_extra_buf_single_buf") + assert lines, "expected at least one hivm.hir.custom line for my_custom_op_extra_buf_single_buf" + line = lines[0] + assert "extra_buffers_types" in line + assert "extra_buffers_sizes" in line + assert "f32" in line + + +def test_custom_op_extra_buffers_integer_variants(): + """extra_buffers accept int16/uint16/int64/uint32/uint8 (IR uses i* storage types).""" + mlir = compile_kernel( + kernel_extra_buf_wide, + {"x_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, + {"BLOCK": 256}, + ) + assert mlir and len(mlir) > 0 + lines = _custom_lines(mlir, "my_custom_op_extra_buf_wide") + assert lines + line = lines[0] + assert "extra_buffers_types" in line + assert "extra_buffers_sizes" in line + assert "i16" in line + assert "i64" in line + assert "i32" in line + assert "i8" in line + assert "1001" in line and "1005" in line + + +def test_custom_op_without_extra_buffers_has_no_extra_buffer_attrs(): + """Ops that do not set extra_buffers should not emit extra_buffers_* attributes.""" + mlir = compile_kernel( + my_kernel, + {"x_ptr": "*fp32", "y_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, + {"BLOCK": 256}, + ) + assert mlir + for line in _custom_lines(mlir, "my_custom_op"): + assert "extra_buffers_types" not in line + assert "extra_buffers_sizes" not in line + + +# ============== Main for manual testing ============== + +if __name__ == "__main__": + test_custom_op() + test_custom_op_without_extra_buffers_has_no_extra_buffer_attrs() + test_custom_op_extra_buffers_integer_variants() + test_custom_op_extra_buffers_mixed_scalar_types() + test_custom_op_extra_buffers_single_buffer() + mlir = compile_kernel(my_kernel, {"x_ptr": "*fp32", "y_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, + {"BLOCK": 256}) + print(f"✅ Generated MLIR ({len(mlir)} chars):\n") + print(mlir) diff --git a/third_party/ascend/unittest/pytest_ut/test_cyl_bessel_i0.py b/third_party/ascend/unittest/pytest_ut/test_cyl_bessel_i0.py new file mode 100644 index 0000000000..5eba4e9690 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_cyl_bessel_i0.py @@ -0,0 +1,61 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + +import torch +import torch_npu +import numpy as np + + +def torch_modified_bessel_i0(x0): + return torch.special.modified_bessel_i0(x0) + + +@triton.jit +def triton_modified_bessel_i0(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = libdevice.cyl_bessel_i0(tmp0) + tl.store(out_ptr0 + (x0), tmp1, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], +]) +def test_modified_bessel_i0(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + np_x0 = test_common.generate_numpy(shape, dtype) + y_ref = np.i0(np_x0) + y_ref = torch.from_numpy(y_ref).to(eval('torch.' + dtype)) + x0 = torch.from_numpy(np_x0).to(eval('torch.' + dtype)).npu() + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_modified_bessel_i0[ncore, 1, 1](x0, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_debug_barrier.py b/third_party/ascend/unittest/pytest_ut/test_debug_barrier.py new file mode 100644 index 0000000000..98e9266bd8 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_debug_barrier.py @@ -0,0 +1,59 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import numpy as np +import torch +import pytest +import test_common + + +def torch_pointwise(x0, x1): + res = x0 - x1 + return res + + +@triton.jit +def triton_sub(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop1 in range(loops1): + x0_prime = offset + (loop1 * XBLOCK_SUB) + base1 + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.load(in_ptr1 + (x0), None) + tmp2 = tmp0 - tmp1 + tl.debug_barrier() + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_pointwise(x0, x1) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_sub[ncore, 1, 1](x0, x1, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_debug_triton_opt_cmd.py b/third_party/ascend/unittest/pytest_ut/test_debug_triton_opt_cmd.py new file mode 100644 index 0000000000..2962fe3eed --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_debug_triton_opt_cmd.py @@ -0,0 +1,115 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Test that ttir_to_linalg prints the triton-opt command line when opt.debug=True. +Does not require NPU hardware. +""" + +from unittest.mock import patch, MagicMock +import pytest + +MOD = "triton.backends.ascend.compiler" + +FAKE_TOOL = "/fake/triton-opt" +FAKE_PIPELINE = "builtin.module(some-pass,another-pass)" + + +def _make_metadata(): + return { + "enable_nd2nz_on_vector": False, + "enable_select_analysis": False, + "compile_on_910_95": False, + "force_simt_template": False, + "enable_sync_block_lock": False, + "enable_mask_fallback_conversion": False, + "optimize_dynamic_offset": False, + "auto_blockify_size": 1, + "add_auto_scheduling": False, + "enable_dynamic_cv_pipeline": False, + "enable_reduce_sum_strength": False, + "hash": "deadbeef", + } + + +def _make_opt(debug): + opt = MagicMock() + opt.debug = debug + return opt + + +def _run_ttir_to_linalg(debug, capsys): + """ + Patch all MLIR machinery so ttir_to_linalg can run without hardware, + then return the captured stdout. + """ + from triton.backends.ascend.compiler import ttir_to_linalg + + pm_mock = MagicMock() + pm_mock.get_pipeline_str.return_value = FAKE_PIPELINE + + ir_mock = MagicMock() + ir_mock.pass_manager.return_value = pm_mock + + dump_manager_mock = MagicMock() + + with patch(f"{MOD}._get_triton_opt_path", return_value=FAKE_TOOL), \ + patch(f"{MOD}._get_triton_adapter_opt_path", return_value="/fake/triton-adapter-opt"), \ + patch(f"{MOD}._is_auto_map_parallel_blocks_enabled", return_value=False), \ + patch(f"{MOD}.ir", ir_mock), \ + patch(f"{MOD}.ascend", MagicMock()), \ + patch(f"{MOD}.passes", MagicMock()), \ + patch(f"{MOD}.get_dump_manager", return_value=dump_manager_mock): + ttir_to_linalg(MagicMock(), _make_metadata(), _make_opt(debug)) + + return capsys.readouterr().out + + +def test_debug_print_contains_tool_path(capsys): + out = _run_ttir_to_linalg(debug=True, capsys=capsys) + assert FAKE_TOOL in out + + +def test_debug_print_contains_pass_pipeline(capsys): + out = _run_ttir_to_linalg(debug=True, capsys=capsys) + assert f"--pass-pipeline={FAKE_PIPELINE}" in out + + +def test_debug_print_contains_debug_info_flag(capsys): + out = _run_ttir_to_linalg(debug=True, capsys=capsys) + assert "--mlir-print-debuginfo" in out + + +def test_debug_print_contains_output_flag(capsys): + out = _run_ttir_to_linalg(debug=True, capsys=capsys) + assert " -o " in out + + +def test_debug_print_has_prefix(capsys): + out = _run_ttir_to_linalg(debug=True, capsys=capsys) + assert "[DEBUG] cmd list:" in out + + +def test_no_debug_print_when_debug_is_false(capsys): + out = _run_ttir_to_linalg(debug=False, capsys=capsys) + assert "[DEBUG] cmd list:" not in out + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/third_party/ascend/unittest/pytest_ut/test_device_print.py b/third_party/ascend/unittest/pytest_ut/test_device_print.py new file mode 100644 index 0000000000..29c5d2b17c --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_device_print.py @@ -0,0 +1,158 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import pytest +import test_common +import os + +os.environ["TRITON_DEVICE_PRINT"] = "1" +os.environ["TRITON_ENABLE_TASKQUEUE"] = "0" +shape = (8, ) +XS = 8 +XVALS_INT = [ + 0, + torch.iinfo(torch.int8).min, + torch.iinfo(torch.int8).max, + torch.iinfo(torch.int16).min, + torch.iinfo(torch.int16).max, + torch.iinfo(torch.int32).min, + torch.iinfo(torch.int32).max, + torch.iinfo(torch.int32).max + 1 +] +XVALS_FP = [ + 0, + torch.finfo(torch.float32).eps, + torch.finfo(torch.float16).eps, + torch.finfo(torch.bfloat16).eps, + torch.finfo(torch.float32).max, + torch.finfo(torch.float16).max, + torch.finfo(torch.bfloat16).max, 1 +] + + +def torch_func(x0, x1): + res = x0 + x1 + return res + + +@triton.jit +def triton_kernel(out_ptr0, in_ptr0, in_ptr1, XBLOCK: tl.constexpr): + idx = tl.arange(0, XBLOCK) + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tl.load(in_ptr1 + idx) + tmp2 = tmp0 + tmp1 + tl.device_print("OUTPUT = ", tmp2) + tl.store(out_ptr0 + idx, tmp2) + + +def triton_func(x0, x1, XS): + out = torch.empty_like(x0) + triton_kernel[1, 1, 1](out, x0, x1, XS) + return out + + +@pytest.mark.skip(reason="waiting for bishengir-compile to support") +@pytest.mark.parametrize('sigtype', ['int64']) +def test_device_print_int64(capsys, sigtype): + dtype = eval(f"torch.{sigtype}") + x0 = torch.zeros(shape, dtype=dtype).npu() + x1 = torch.ones(shape, dtype=dtype).npu() + for i in range(x1.numel()): + x1[i] = XVALS_INT[i] + torch_ref = torch_func(x0, x1) + triton_cal = triton_func(x0, x1, XS) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) + + +@pytest.mark.parametrize('sigtype', ['int32']) +def test_device_print_int32(capsys, sigtype): + dtype = eval(f"torch.{sigtype}") + x0 = torch.zeros(shape, dtype=dtype).npu() + x1 = torch.ones(shape, dtype=dtype).npu() + for i in range(x1.numel()): + x1[i] = XVALS_INT[i] + torch_ref = torch_func(x0, x1) + triton_cal = triton_func(x0, x1, XS) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) + + +@pytest.mark.parametrize('sigtype', ['int16']) +def test_device_print_int16(capsys, sigtype): + dtype = eval(f"torch.{sigtype}") + x0 = torch.zeros(shape, dtype=dtype).npu() + x1 = torch.ones(shape, dtype=dtype).npu() + for i in range(x1.numel()): + x1[i] = XVALS_INT[i] + torch_ref = torch_func(x0, x1) + triton_cal = triton_func(x0, x1, XS) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) + + +@pytest.mark.parametrize('sigtype', ['int8']) +def test_device_print_int8(capsys, sigtype): + dtype = eval(f"torch.{sigtype}") + x0 = torch.zeros(shape, dtype=dtype).npu() + x1 = torch.ones(shape, dtype=dtype).npu() + for i in range(x1.numel()): + x1[i] = XVALS_INT[i] + torch_ref = torch_func(x0, x1) + triton_cal = triton_func(x0, x1, XS) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) + + +@pytest.mark.parametrize('sigtype', ['float32']) +def test_device_print_fp32(capsys, sigtype): + dtype = eval(f"torch.{sigtype}") + x0 = torch.zeros(shape, dtype=dtype).npu() + x1 = torch.ones(shape, dtype=dtype).npu() + for i in range(x1.numel()): + x1[i] = XVALS_FP[i] + torch_ref = torch_func(x0, x1) + triton_cal = triton_func(x0, x1, XS) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) + + +@pytest.mark.parametrize('sigtype', ['float16']) +def test_device_print_fp16(capsys, sigtype): + dtype = eval(f"torch.{sigtype}") + x0 = torch.zeros(shape, dtype=dtype).npu() + x1 = torch.ones(shape, dtype=dtype).npu() + for i in range(x1.numel()): + x1[i] = XVALS_FP[i] + torch_ref = torch_func(x0, x1) + triton_cal = triton_func(x0, x1, XS) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) + + +@pytest.mark.skip(reason="waiting for bishengir-compile to support") +@pytest.mark.parametrize('sigtype', ['bfloat16']) +def test_device_print_bf16(capsys, sigtype): + dtype = eval(f"torch.{sigtype}") + x0 = torch.zeros(shape, dtype=dtype).npu() + x1 = torch.ones(shape, dtype=dtype).npu() + for i in range(x1.numel()): + x1[i] = XVALS_FP[i] + torch_ref = torch_func(x0, x1) + triton_cal = triton_func(x0, x1, XS) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_device_print_comprehensive.py b/third_party/ascend/unittest/pytest_ut/test_device_print_comprehensive.py new file mode 100644 index 0000000000..d623fc51d2 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_device_print_comprehensive.py @@ -0,0 +1,122 @@ +import os +import sys +import subprocess +import tempfile +import textwrap + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +expected_prints = [ + "Offsets:", + "Mask:", + "Pointer offsets:", + "Loaded x:", + "Scalar factor:", + "Temp result (x * 2):", + "Final y (x * 2 + 1):", + "Positive mask:", + "Block ID:", + "Block start:", + "Valid elements in this block:", +] + + +def test_comprehensive_print(): + + with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f: + temp_script = f.name + + f.write( + textwrap.dedent(f""" +import os +import sys +import subprocess +import tempfile +import textwrap + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +os.environ["TRITON_DEVICE_PRINT"] = "1" + +@triton.jit +def comprehensive_print_kernel( + x_ptr, + y_ptr, + mask_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + block_start = pid * BLOCK_SIZE + + offsets = block_start + tl.arange(0, BLOCK_SIZE) + tl.device_print("Offsets: ", offsets) + + mask = offsets < n_elements + tl.device_print("Mask: ", mask) + + x_ptrs = x_ptr + offsets + tl.device_print("Pointer offsets: ", offsets) + + x = tl.load(x_ptrs, mask=mask, other=0.0) + tl.device_print("Loaded x: ", x) + + scalar_factor = 2.0 + tl.device_print("Scalar factor: ", scalar_factor) + + y_temp = x * scalar_factor + tl.device_print("Temp result (x * 2): ", y_temp) + + y = y_temp + 1.0 + tl.device_print("Final y (x * 2 + 1): ", y) + + positive_mask = y > 0.0 + tl.device_print("Positive mask: ", positive_mask) + + tl.device_print("Block ID: ", pid) + tl.device_print("Block start: ", block_start) + + y_ptrs = y_ptr + offsets + tl.store(y_ptrs, y, mask=mask) + mask_count = tl.sum(mask.to(tl.int32)) + tl.device_print("Valid elements in this block: ", mask_count) + + +def test_comprehensive_print(): + size = 16 + x = torch.randn(size).npu() + y = torch.zeros(size).npu() + mask = torch.ones(size, dtype=torch.bool).npu() + BLOCK_SIZE = 32 + + h = comprehensive_print_kernel[1,](x, y, mask, size, BLOCK_SIZE=BLOCK_SIZE) + + expected = x * 2.0 + 1.0 + torch.testing.assert_close(y, expected, rtol=1e-5, atol=1e-5) + + for i in range(11): + opStr = "call @triton_print_" + str(i) + assert opStr in h.asm["ttadapter"] + + print("passed!") + + +if __name__ == "__main__": + test_comprehensive_print() + """)) + + result = subprocess.run([sys.executable, temp_script], capture_output=True, text=True, env=os.environ.copy()) + + captured_output = result.stdout + "\n=== STDERR ===\n" + result.stderr + + assert "passed!" in captured_output + for prefix in expected_prints: + assert prefix in captured_output diff --git a/third_party/ascend/unittest/pytest_ut/test_device_print_script.py b/third_party/ascend/unittest/pytest_ut/test_device_print_script.py new file mode 100644 index 0000000000..15136f67cf --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_device_print_script.py @@ -0,0 +1,117 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import os +import subprocess +import re +import pytest + + +# 执行test_device_print.py -> 生成 test_device_print.log, 以便执行后续验证 +def run_test_device_print_py(test_name, log_name): + testfile_path = os.path.join(os.getcwd(), test_name) + logfile_path = os.path.join(os.getcwd(), log_name) + + with open(logfile_path, 'w') as f: + try: + subprocess.run(["pytest", testfile_path], stdout=f, stderr=subprocess.STDOUT, check=True) + print(f"Run 【{test_name}】 successfully!") + except Exception as e: + print(f"Run 【{test_name}】 unsuccessfully: ", e) + + +def assert_close(expected_output, logfile): + # 读回日志内容 + with open(logfile, "r", encoding="utf-8") as f: + raw = f.read() + cleaned = re.sub(r"\x00", "", raw) + + # 在文件内容里做检查 + assert expected_output in cleaned, f"Expected '{expected_output}' not found in log file." + + +@pytest.mark.skip(reason="waiting for TA to support") +def test_device_print_int8(): + expected_output = "0,-128,127,0,-1,0,-1,0" + test_name = "test_device_print.py::test_device_print_int8[int8]" + log_name = "test_device_print_int8.log" + logfile = os.path.join(os.getcwd(), log_name) + run_test_device_print_py(test_name, logfile) + assert_close(expected_output, logfile) + + +@pytest.mark.skip(reason="waiting for TA to support") +def test_device_print_int16(): + expected_output = "0,-128,127,-32768,32767,0,-1,0" + test_name = "test_device_print.py::test_device_print_int16[int16]" + log_name = "test_device_print_int16.log" + logfile = os.path.join(os.getcwd(), log_name) + run_test_device_print_py(test_name, logfile) + assert_close(expected_output, logfile) + + +@pytest.mark.skip(reason="waiting for TA to support") +def test_device_print_int32(): + expected_output = "0,-128,127,-32768,32767,-2147483648,2147483647,-2147483648" + test_name = "test_device_print.py::test_device_print_int32[int32]" + log_name = "test_device_print_int32.log" + logfile = os.path.join(os.getcwd(), log_name) + run_test_device_print_py(test_name, logfile) + assert_close(expected_output, logfile) + + +@pytest.mark.skip(reason="waiting for compiler to support") +def test_device_print_int64(): + expected_output = "???" + test_name = "test_device_print.py::test_device_print_int64[int64]" + log_name = "test_device_print_int64.log" + logfile = os.path.join(os.getcwd(), log_name) + run_test_device_print_py(test_name, logfile) + assert_close(expected_output, logfile) + + +@pytest.mark.skip(reason="waiting for TA to support") +def test_device_print_fp16(): + expected_output = "0.000000,0.000000,0.000977,0.007812,inf,65504.000000,inf,1.000000" + test_name = "test_device_print.py::test_device_print_fp16[float16]" + log_name = "test_device_print_fp16.log" + logfile = os.path.join(os.getcwd(), log_name) + run_test_device_print_py(test_name, logfile) + assert_close(expected_output, logfile) + + +@pytest.mark.skip(reason="waiting for TA to support") +def test_device_print_fp32(): + expected_output = "0.000000,0.000000,0.000977,0.007812,340282346638528859811704183484516925440.000000,65504.000000,338953138925153547590470800371487866880.000000,1.000000" + test_name = "test_device_print.py::test_device_print_fp32[float32]" + log_name = "test_device_print_fp16.log" + logfile = os.path.join(os.getcwd(), log_name) + run_test_device_print_py(test_name, logfile) + assert_close(expected_output, logfile) + + +@pytest.mark.skip(reason="waiting for compiler to support") +def test_device_print_bf16(): + expected_output = "???" + test_name = "test_device_print.py::test_device_print_bf16[bfloat16]" + log_name = "test_device_print_bf16.log" + logfile = os.path.join(os.getcwd(), log_name) + run_test_device_print_py(test_name, logfile) + assert_close(expected_output, logfile) diff --git a/third_party/ascend/unittest/pytest_ut/test_discrete_mask_atomic.py b/third_party/ascend/unittest/pytest_ut/test_discrete_mask_atomic.py new file mode 100644 index 0000000000..c0527c7dc4 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_discrete_mask_atomic.py @@ -0,0 +1,48 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import triton +import triton.language as tl +import torch_npu +import pytest + + +@triton.jit +def single_disc_mask_atomic_add_kernel( + in_ptr, + BLOCK_N: tl.constexpr, +): + col_offs = tl.arange(0, BLOCK_N) + disc_mask = (col_offs * 2) < BLOCK_N + ptr_in = in_ptr + col_offs + tl.atomic_add(ptr_in, 1, mask=disc_mask) + + +@pytest.mark.parametrize("BLOCK_N", [8]) +def test_single_discrete_mask_atomic_add(BLOCK_N): + in_tensor = torch.arange(BLOCK_N, dtype=torch.float16, device='npu') + expected = in_tensor.clone() + single_disc_mask_atomic_add_kernel[(1, )](in_tensor, BLOCK_N=BLOCK_N) + + half = BLOCK_N // 2 + expected[:half] += 1 + assert torch.allclose(in_tensor, expected), \ + f"Expected:\n{expected.cpu()}\nGot:\n{in_tensor.cpu()}" diff --git a/third_party/ascend/unittest/pytest_ut/test_discrete_mask_loadstore.py b/third_party/ascend/unittest/pytest_ut/test_discrete_mask_loadstore.py new file mode 100644 index 0000000000..cb03ff1bec --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_discrete_mask_loadstore.py @@ -0,0 +1,408 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +# ============================================================================= +# Discrete mask access conversion test suite +# +# Test matrix (mask type x operation type): +# +# | mask type | load only | store only | load + store | +# |---------------------------------|-----------|------------|--------------| +# | single discrete mask | (A) | (B) | - | +# | single continuous mask | (C) | (D) | - | +# | continuous & discrete 2-way | (E) | (F) | (G) | +# | continuous & discrete 4-way | - | - | (H) | +# | broadcast(cont & disc) 2-D AND | (I) | - | (J) | +# +# ============================================================================= + +import torch +import triton +import triton.language as tl +import torch_npu +import pytest + + +# ============================================================================= +# (A) Single discrete mask -- load only +# ============================================================================= +@triton.jit +def single_disc_mask_load_kernel( + in_ptr, + out_ptr, + BLOCK_N: tl.constexpr, +): + col_offs = tl.arange(0, BLOCK_N) + disc_mask = (col_offs * 2) < BLOCK_N + ptr_in = in_ptr + col_offs + ptr_out = out_ptr + col_offs + data = tl.load(ptr_in, mask=disc_mask, other=0.0) + tl.store(ptr_out, data) + + +@pytest.mark.parametrize("BLOCK_N", [8]) +def test_single_discrete_mask_load(BLOCK_N): + in_tensor = torch.arange(BLOCK_N, dtype=torch.float16, device='npu') + out_tensor = torch.empty(BLOCK_N, dtype=torch.float16, device='npu') + + single_disc_mask_load_kernel[(1, )](in_tensor, out_tensor, BLOCK_N=BLOCK_N) + + half = BLOCK_N // 2 + expected = torch.zeros(BLOCK_N, dtype=torch.float16, device='npu') + expected[:half] = in_tensor[:half] + assert torch.allclose(out_tensor, expected), \ + f"Expected:\n{expected.cpu()}\nGot:\n{out_tensor.cpu()}" + + +# ============================================================================= +# (B) Single discrete mask -- store only +# ============================================================================= +@triton.jit +def single_disc_mask_store_kernel( + in_ptr, + out_ptr, + BLOCK_N: tl.constexpr, +): + col_offs = tl.arange(0, BLOCK_N) + disc_mask = (col_offs * 2) < BLOCK_N + ptr_in = in_ptr + col_offs + ptr_out = out_ptr + col_offs + data = tl.load(ptr_in) + tl.store(ptr_out, data, mask=disc_mask) + + +@pytest.mark.parametrize("BLOCK_N", [8]) +def test_single_discrete_mask_store(BLOCK_N): + in_tensor = torch.arange(BLOCK_N, dtype=torch.float16, device='npu') + out_tensor = torch.full((BLOCK_N, ), -1.0, dtype=torch.float16, device='npu') + + single_disc_mask_store_kernel[(1, )](in_tensor, out_tensor, BLOCK_N=BLOCK_N) + + half = BLOCK_N // 2 + expected = torch.full((BLOCK_N, ), -1.0, dtype=torch.float16, device='npu') + expected[:half] = in_tensor[:half] + assert torch.allclose(out_tensor, expected), \ + f"Expected:\n{expected.cpu()}\nGot:\n{out_tensor.cpu()}" + + +# ============================================================================= +# (C) Single continuous mask -- load only +# ============================================================================= +@triton.jit +def single_cont_mask_load_kernel( + in_ptr, + out_ptr, + M: tl.constexpr, + BLOCK_M: tl.constexpr, +): + row_offs = tl.arange(0, BLOCK_M) + cont_mask = row_offs < M # Continuous mask + ptr_in = in_ptr + row_offs + ptr_out = out_ptr + row_offs + data = tl.load(ptr_in, mask=cont_mask, other=0.0) + tl.store(ptr_out, data, mask=cont_mask) + + +@pytest.mark.parametrize("M,BLOCK_M", [(6, 8)]) +def test_single_continuous_mask_load(M, BLOCK_M): + in_tensor = torch.arange(BLOCK_M, dtype=torch.float16, device='npu') + out_tensor = torch.full((BLOCK_M, ), -1.0, dtype=torch.float16, device='npu') + + single_cont_mask_load_kernel[(1, )](in_tensor, out_tensor, M=M, BLOCK_M=BLOCK_M) + + expected = torch.full((BLOCK_M, ), -1.0, dtype=torch.float16, device='npu') + expected[:M] = in_tensor[:M] + assert torch.allclose(out_tensor, expected), \ + f"Expected:\n{expected.cpu()}\nGot:\n{out_tensor.cpu()}" + + +# ============================================================================= +# (D) Single continuous mask -- store only +# ============================================================================= +@triton.jit +def single_cont_mask_store_kernel( + in_ptr, + out_ptr, + M: tl.constexpr, + BLOCK_M: tl.constexpr, +): + row_offs = tl.arange(0, BLOCK_M) + cont_mask = row_offs < M + ptr_in = in_ptr + row_offs + ptr_out = out_ptr + row_offs + data = tl.load(ptr_in) + tl.store(ptr_out, data, mask=cont_mask) + + +@pytest.mark.parametrize("M,BLOCK_M", [(6, 8)]) +def test_single_continuous_mask_store(M, BLOCK_M): + in_tensor = torch.arange(BLOCK_M, dtype=torch.float16, device='npu') + out_tensor = torch.full((BLOCK_M, ), -1.0, dtype=torch.float16, device='npu') + single_cont_mask_store_kernel[(1, )](in_tensor, out_tensor, M=M, BLOCK_M=BLOCK_M) + expected = torch.full((BLOCK_M, ), -1.0, dtype=torch.float16, device='npu') + expected[:M] = in_tensor[:M] + assert torch.allclose(out_tensor, expected), \ + f"Expected:\n{expected.cpu()}\nGot:\n{out_tensor.cpu()}" + + +# ============================================================================= +# (E) Continuous & discrete 2-way AND -- load only +# ============================================================================= +@triton.jit +def cont_disc_combined_mask_load_kernel( + in_ptr, + out_ptr, + M: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + row_offs = tl.arange(0, BLOCK_M) + col_offs = tl.arange(0, BLOCK_N) + # Continuous mask + row_boundary = row_offs < M + # Discrete mask + col_stride = (col_offs * 2) < BLOCK_N + mask = row_boundary[:, None] & col_stride[None, :] + ptr_in = in_ptr + row_offs[:, None] * BLOCK_N + col_offs[None, :] + ptr_out = out_ptr + row_offs[:, None] * BLOCK_N + col_offs[None, :] + data = tl.load(ptr_in, mask=mask, other=0.0) + tl.store(ptr_out, data) + + +@pytest.mark.parametrize("M,BLOCK_M,BLOCK_N", [(6, 8, 8)]) +def test_cont_disc_combined_mask_load(M, BLOCK_M, BLOCK_N): + in_tensor = torch.ones((BLOCK_M, BLOCK_N), dtype=torch.float16, device='npu') + out_tensor = torch.empty((BLOCK_M, BLOCK_N), dtype=torch.float16, device='npu') + + cont_disc_combined_mask_load_kernel[(1, )](in_tensor, out_tensor, M=M, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N) + + half_n = BLOCK_N // 2 + expected = torch.zeros((BLOCK_M, BLOCK_N), dtype=torch.float16, device='npu') + expected[:M, :half_n] = 1.0 + assert torch.allclose(out_tensor, expected), \ + f"Expected:\n{expected.cpu()}\nGot:\n{out_tensor.cpu()}" + + +# ============================================================================= +# (F) Continuous & discrete 2-way AND -- store only +# ============================================================================= +@triton.jit +def cont_disc_combined_mask_store_kernel( + in_ptr, + out_ptr, + M: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + row_offs = tl.arange(0, BLOCK_M) + col_offs = tl.arange(0, BLOCK_N) + row_boundary = row_offs < M # continuous -> contLeaf + col_stride = (col_offs * 2) < BLOCK_N # discrete -> discLeaf + mask = row_boundary[:, None] & col_stride[None, :] + ptr_in = in_ptr + row_offs[:, None] * BLOCK_N + col_offs[None, :] + ptr_out = out_ptr + row_offs[:, None] * BLOCK_N + col_offs[None, :] + data = tl.load(ptr_in) + tl.store(ptr_out, data, mask=mask) + + +@pytest.mark.parametrize("M,BLOCK_M,BLOCK_N", [(6, 8, 8)]) +def test_cont_disc_combined_mask_store(M, BLOCK_M, BLOCK_N): + in_tensor = torch.ones((BLOCK_M, BLOCK_N), dtype=torch.float16, device='npu') + out_tensor = torch.full((BLOCK_M, BLOCK_N), -1.0, dtype=torch.float16, device='npu') + cont_disc_combined_mask_store_kernel[(1, )](in_tensor, out_tensor, M=M, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N) + half_n = BLOCK_N // 2 + expected = torch.full((BLOCK_M, BLOCK_N), -1.0, dtype=torch.float16, device='npu') + expected[:M, :half_n] = 1.0 + assert torch.allclose(out_tensor, expected), \ + f"Expected:\n{expected.cpu()}\nGot:\n{out_tensor.cpu()}" + + +# ============================================================================= +# (G) Continuous & discrete 2-way AND -- load + store (complex interleave, original) +# ============================================================================= +@triton.jit +def interleave_cont_disc_mask_kernel( + in_ptr, + out_ptr, + M: tl.constexpr, + N: tl.constexpr, +): + pid = tl.program_id(0) + col_offs = tl.arange(0, N) + even_col_offs = tl.arange(0, N // 2) * 2 + even_col_mask = even_col_offs < N # discrete: cmpi(muli(range,2), N) + row_offs = tl.arange(0, M) + row_mask = row_offs < M # continuous: cmpi(range_M, M) + in_even_ptr = in_ptr + row_offs[:, None] * N + even_col_offs[None, :] + in_odd_ptr = in_ptr + row_offs[:, None] * N + even_col_offs[None, :] + 1 + even_data = tl.load(in_even_ptr, mask=row_mask[:, None] & even_col_mask[None, :], other=0.0) + odd_data = tl.load(in_odd_ptr) + rotated_data = tl.interleave(-odd_data, even_data) + out_ptr = out_ptr + row_offs[:, None] * N + col_offs[None, :] + tl.store(out_ptr, rotated_data) + + +@pytest.mark.skip(reason="not supported after the NPUIR is updated in April, and will be fixed later") +@pytest.mark.parametrize("M", [4]) +@pytest.mark.parametrize("N", [8]) +def test_discrete_mask_load_store(M, N): + """Regression test: mask=row_mask & even_col_mask (continuous & discrete 2-way)""" + input_tensor = torch.arange(M * N, dtype=torch.float16, device='npu').reshape(M, N) + output_tensor = torch.empty_like(input_tensor) + interleave_cont_disc_mask_kernel[(1, )](input_tensor, output_tensor, M=M, N=N) + even_cols = input_tensor[:, 0::2] + odd_cols = input_tensor[:, 1::2] + ref_output = torch.empty_like(input_tensor) + ref_output[:, 0::2] = -odd_cols + ref_output[:, 1::2] = even_cols + assert torch.allclose(output_tensor.float(), ref_output.float()) + + +# ============================================================================= +# (H) Continuous & discrete 4-way AND -- load + store +# ============================================================================= +@triton.jit +def multi_cont_disc_mask_load_store_kernel( + in_ptr, + out_ptr, + M: tl.constexpr, + N: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + row_offs = tl.arange(0, BLOCK_M) + col_offs = tl.arange(0, BLOCK_N) + + row_boundary = row_offs < M # continuous mask + col_boundary = col_offs < N # continuous mask + row_stride = (row_offs * 2) < BLOCK_M # discrete mask + col_stride = (col_offs * 2) < BLOCK_N # discrete mask + + mask = (row_boundary[:, None] & col_boundary[None, :] & row_stride[:, None] & col_stride[None, :]) + + ptr_in = in_ptr + row_offs[:, None] * BLOCK_N + col_offs[None, :] + ptr_out = out_ptr + row_offs[:, None] * BLOCK_N + col_offs[None, :] + + data = tl.load(ptr_in, mask=mask, other=0.0) + result = data + 1.0 + tl.store(ptr_out, result, mask=mask) + + +@pytest.mark.parametrize("M,N,BLOCK_M,BLOCK_N", [ + (6, 6, 8, 8), +]) +def test_multi_cont_disc_mask_load_store(M, N, BLOCK_M, BLOCK_N): + in_tensor = torch.ones((BLOCK_M, BLOCK_N), dtype=torch.float16, device='npu') + out_tensor = torch.zeros((BLOCK_M, BLOCK_N), dtype=torch.float16, device='npu') + + multi_cont_disc_mask_load_store_kernel[(1, )](in_tensor, out_tensor, M=M, N=N, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N) + + half_m = BLOCK_M // 2 # = 4 + half_n = BLOCK_N // 2 # = 4 + expected = torch.zeros((BLOCK_M, BLOCK_N), dtype=torch.float16, device='npu') + expected[:half_m, :half_n] = 2.0 + + assert torch.allclose(out_tensor, expected), (f"BLOCK=({BLOCK_M},{BLOCK_N}), valid=({M},{N})\n" + f"Expected:\n{expected.cpu()}\nGot:\n{out_tensor.cpu()}") + + +# ============================================================================= +# (I) broadcast(continuous & discrete) 2-D AND -- load only +# ============================================================================= +@triton.jit +def broadcast_cont_disc_2d_load_kernel( + in_ptr, + out_ptr, + M: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + row_offs = tl.arange(0, BLOCK_M) + col_offs = tl.arange(0, BLOCK_N) + + row_boundary = row_offs < M + row_disc = (row_offs * 2) < BLOCK_M + mask = row_boundary[:, None] & row_disc[:, None] & (col_offs < BLOCK_N)[None, :] + + ptr_in = in_ptr + row_offs[:, None] * BLOCK_N + col_offs[None, :] + ptr_out = out_ptr + row_offs[:, None] * BLOCK_N + col_offs[None, :] + + data = tl.load(ptr_in, mask=mask, other=0.0) + tl.store(ptr_out, data) + + +@pytest.mark.parametrize("M,BLOCK_M,BLOCK_N", [(3, 4, 8)]) +def test_broadcast_cont_disc_2d_load(M, BLOCK_M, BLOCK_N): + in_tensor = torch.ones((BLOCK_M, BLOCK_N), dtype=torch.float16, device='npu') + out_tensor = torch.empty((BLOCK_M, BLOCK_N), dtype=torch.float16, device='npu') + + broadcast_cont_disc_2d_load_kernel[(1, )](in_tensor, out_tensor, M=M, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N) + + disc_true_rows = BLOCK_M // 2 + both_true_rows = min(M, disc_true_rows) + + expected = torch.zeros((BLOCK_M, BLOCK_N), dtype=torch.float16, device='npu') + expected[:both_true_rows, :] = 1.0 + + assert torch.allclose(out_tensor, expected), (f"M={M}, BLOCK_M={BLOCK_M}, BLOCK_N={BLOCK_N}\n" + f"Expected:\n{expected.cpu()}\nGot:\n{out_tensor.cpu()}") + + +# ============================================================================= +# (J) broadcast(continuous & discrete) 2-D AND -- load + store +# ============================================================================= +@triton.jit +def broadcast_cont_disc_2d_load_store_kernel( + in_ptr, + out_ptr, + M: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + row_offs = tl.arange(0, BLOCK_M) + col_offs = tl.arange(0, BLOCK_N) + + row_boundary = row_offs < M + row_disc = (row_offs * 2) < BLOCK_M + + combined = row_boundary[:, None] & row_disc[:, None] & (col_offs < BLOCK_N)[None, :] + + ptr_in = in_ptr + row_offs[:, None] * BLOCK_N + col_offs[None, :] + ptr_out = out_ptr + row_offs[:, None] * BLOCK_N + col_offs[None, :] + + data = tl.load(ptr_in, mask=combined, other=0.0) + tl.store(ptr_out, data, mask=combined) + + +@pytest.mark.parametrize("M,BLOCK_M,BLOCK_N", [(3, 4, 8)]) +def test_broadcast_cont_disc_2d_load_store(M, BLOCK_M, BLOCK_N): + in_tensor = torch.ones((BLOCK_M, BLOCK_N), dtype=torch.float16, device='npu') + out_tensor = torch.full((BLOCK_M, BLOCK_N), -1.0, dtype=torch.float16, device='npu') + + broadcast_cont_disc_2d_load_store_kernel[(1, )](in_tensor, out_tensor, M=M, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N) + + disc_true_rows = BLOCK_M // 2 + both_true_rows = min(M, disc_true_rows) + + expected = torch.full((BLOCK_M, BLOCK_N), -1.0, dtype=torch.float16, device='npu') + expected[:both_true_rows, :] = 1.0 + + assert torch.allclose(out_tensor, expected), (f"M={M}, BLOCK_M={BLOCK_M}, BLOCK_N={BLOCK_N}\n" + f"Expected:\n{expected.cpu()}\nGot:\n{out_tensor.cpu()}") diff --git a/third_party/ascend/unittest/pytest_ut/test_discrete_mask_tail_block_mte_oob.py b/third_party/ascend/unittest/pytest_ut/test_discrete_mask_tail_block_mte_oob.py new file mode 100644 index 0000000000..d8129d2677 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_discrete_mask_tail_block_mte_oob.py @@ -0,0 +1,238 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +# ============================================================================= +# MTE (Memory Tag Extension) OOB regression test for DiscreteMaskAccessConversionPass +# +# This file verifies that DiscreteMaskAccessConversionPass correctly bounds +# global-memory accesses when the load/store mask is a combined discrete mask +# +# Test strategy +# ------------- +# The test engineers this condition in four steps: +# Step 1 — probe: Trigger a fresh 2 MB NPU segment; measure its size. +# Step 2 — pre_fill: Fill the segment with small tensors until the remaining +# free space is in [IN_BYTES, TARGET_FREE]. +# Step 3 — in_tensor: Allocate the test tensor; it lands at the segment tail +# with only ~7680 bytes gap to the boundary. +# Step 4 — kernel: Run the kernel + synchronize. Before the fix the +# OOB read (24576 bytes) crosses the boundary → MTE. +# After the fix the copy is bounded to IN_BYTES → no MTE. +# +# Memory layout at the time of the kernel call (before fix): +# +# ┌──────────────────────────── 2 MB segment ────────────────────────────────┐ +# │ probe(512 B) │←────── pre_fill (~2025 KB) ──────→│ in_tensor(8192 B) │gap│ +# └──────────────────────────────────────────────────────────────────────────┘ +# ↑ segment end +# ├──── OOB_BYTES (24576 B) ─────→ +# crosses boundary → MTE ✓ +# +# ============================================================================= + +import math +import torch +import triton +import triton.language as tl +import torch_npu +import pytest + + +@triton.jit +def cont_disc_oob_inplace_2d_kernel( + ptr, + M, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_m = tl.program_id(0) + row_offs = tl.arange(0, BLOCK_M) + col_offs = tl.arange(0, BLOCK_N) + + # Continuous bound (contMask): marks the M valid rows in this tile. + row_boundary = row_offs < (M - pid_m * BLOCK_M) + row_disc = (row_offs * 2) < BLOCK_M + combined = row_boundary[:, None] & row_disc[:, None] & (col_offs < BLOCK_N)[None, :] + + row_start = pid_m * BLOCK_M + ptr_2d = ptr + (row_start + row_offs[:, None]) * BLOCK_N + col_offs[None, :] + + # load triggers DiscreteMaskAccessConversionPass. + # Before fix: copy size = BLOCK_M × BLOCK_N × 2 bytes = 32768 bytes (OOB). + # After fix: copy size = M × BLOCK_N × 2 bytes = 8192 bytes (safe). + data = tl.load(ptr_2d, mask=combined, other=0.0) + tl.store(ptr_2d, data, mask=row_boundary[:, None]) + + +# ============================================================================= +# Memory setup helper +# ============================================================================= +def _fill_segment_to_boundary(dtype, device, in_bytes, target_free, chunk_max_bytes): + """Allocate a fresh NPU segment and fill it so that only ~target_free bytes remain. + + Returns + ------- + pre_fillers : list of torch.Tensor + All tensors allocated (probe + fill chunks). The caller is responsible + for deleting them in `finally`. + pool_free_after_fill : int + Segment free space after filling, in bytes. + seg_size : int + Total size of the triggered segment, in bytes. + """ + elem_size = torch.finfo(dtype).bits // 8 + + # --- Step 1: probe — trigger a fresh 2 MB small-alloc segment ---------- + pool0 = torch.npu.memory_reserved(0) + alloc0 = torch.npu.memory_allocated(0) + + probe = torch.empty(1, dtype=dtype, device=device) + + pool1 = torch.npu.memory_reserved(0) + alloc1 = torch.npu.memory_allocated(0) + + seg_size = pool1 - pool0 # should be 2 MB = 2097152 bytes + probe_actual = alloc1 - alloc0 # NPU 512-byte aligned → 512 bytes + + print(f"\n[mte] Step 1: probe") + print(f"[mte] segment_size = {seg_size} bytes ({seg_size // 1024} KB)") + print(f"[mte] probe_actual = {probe_actual} bytes") + print(f"[mte] pool_free = {seg_size - probe_actual} bytes") + + # --- Step 2: pre_fill — leave only [in_bytes, target_free] bytes free --- + # Chunks are kept ≤ chunk_max_bytes to stay in the small-alloc pool and + # avoid opening a new segment via the large-alloc path. + pre_fillers = [probe] + + for chunk in [ + chunk_max_bytes, chunk_max_bytes // 2, chunk_max_bytes // 4, chunk_max_bytes // 8, 32 * 1024, 16 * 1024, + 8 * 1024, 4 * 1024, 2 * 1024, 1024, 512 + ]: + while True: + free = torch.npu.memory_reserved(0) - torch.npu.memory_allocated(0) + if free <= in_bytes: + break # not enough room even for in_tensor; stop + if free <= target_free: + break # already in target range; try smaller chunk + if free <= target_free + chunk: + break # this chunk would overshoot; try smaller chunk + try: + t = torch.empty(chunk // elem_size, dtype=dtype, device=device) + pre_fillers.append(t) + except RuntimeError: + break # segment exhausted; try smaller chunk + + pool_free_after_fill = torch.npu.memory_reserved(0) - torch.npu.memory_allocated(0) + pre_bytes = sum(t.numel() * elem_size for t in pre_fillers) + print(f"\n[mte] Step 2: pre_fill") + print(f"[mte] tensors = {len(pre_fillers)}, total = {pre_bytes} bytes ({pre_bytes // 1024} KB)") + print(f"[mte] pool_free = {pool_free_after_fill} bytes (target [{in_bytes}, {target_free}] bytes)") + + return pre_fillers, pool_free_after_fill, seg_size + + +# ============================================================================= +# Test: MTE OOB via segment-boundary placement +# ============================================================================= +@pytest.mark.parametrize("BLOCK_M,BLOCK_N,M", [ + (4, 4096, 1), +]) +def test_mte_segment_boundary_oob(BLOCK_M, BLOCK_N, M): + """Regression: combined discrete mask load causes OOB on tail blocks. + + Verifies that DiscreteMaskAccessConversionPass correctly bounds + the memory copy to M rows (the contiguous range), not BLOCK_M rows (the full tile). + + Test outcome: + - Before fix: RuntimeError (MTE OOB) — the test would fail. + - After fix: no exception — the test passes. + """ + dtype = torch.float16 + device = 'npu' + elem_size = 2 # float16 + + in_bytes = M * BLOCK_N * elem_size # 8192 bytes + oob_bytes = (BLOCK_M - M) * BLOCK_N * elem_size # 24576 bytes + # TARGET_FREE: midpoint between in_bytes and oob_bytes. + # Ensures in_tensor fits AND gap < oob_bytes so OOB crosses segment boundary. + target_free = (in_bytes + oob_bytes) // 2 # 16384 bytes + chunk_max_bytes = 512 * 1024 # 512 KB + + print(f"\n[mte] BLOCK_M={BLOCK_M} BLOCK_N={BLOCK_N} M={M}") + print(f"[mte] in_bytes = {in_bytes} bytes (in_tensor: {M}×{BLOCK_N}×{elem_size})") + print(f"[mte] oob_bytes = {oob_bytes} bytes (unfixed copy: {BLOCK_M}×{BLOCK_N}×{elem_size} - in_bytes)") + print(f"[mte] target_free = {target_free} bytes (must satisfy in_bytes < target_free < oob_bytes)") + + torch.npu.empty_cache() + + pre_fillers = [] + in_tensor = None + + try: + pre_fillers, pool_free_after_fill, _ = _fill_segment_to_boundary(dtype, device, in_bytes, target_free, + chunk_max_bytes) + except Exception as exc: + torch.npu.empty_cache() + pytest.skip(f"Memory layout setup failed (allocator behaviour may differ): {exc}") + + # Verify pre_fill achieved the required free-space window. + if not (in_bytes <= pool_free_after_fill <= target_free): + for t in reversed(pre_fillers): + del t + torch.npu.empty_cache() + pytest.skip(f"pre_fill did not reach target range [{in_bytes}, {target_free}] bytes; " + f"got {pool_free_after_fill} bytes. " + f"Skipping MTE check (NPU allocator behaviour may differ).") + + try: + # Step 3: allocate in_tensor — lands at the very end of the segment. + # NPU 512-byte alignment means the allocator consumes + # in_bytes + 512 = 8704 bytes, leaving gap ≈ target_free - 8704 = 7680 bytes. + in_tensor = torch.ones(M * BLOCK_N, dtype=dtype, device=device).view(M, BLOCK_N) + + gap = torch.npu.memory_reserved(0) - torch.npu.memory_allocated(0) + print(f"\n[mte] Step 3: in_tensor") + print(f"[mte] address = [{in_tensor.data_ptr():#x}, {in_tensor.data_ptr() + in_bytes:#x})") + print(f"[mte] gap = {gap} bytes (in_tensor end → segment end)") + + if oob_bytes <= gap: + pytest.skip(f"gap ({gap} bytes) >= oob_bytes ({oob_bytes} bytes): " + f"OOB would not cross the segment boundary. " + f"Skipping MTE check.") + print(f"[mte] oob_bytes({oob_bytes} B) > gap({gap} B) → MTE expected if unfixed ✓") + + # Step 4: run kernel + num_pids_m = math.ceil(M / BLOCK_M) + print(f"\n[mte] Step 4: kernel (grid=({num_pids_m},))") + cont_disc_oob_inplace_2d_kernel[(num_pids_m, )](in_tensor, M=M, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N) + torch.npu.synchronize() + print("[mte] PASSED: fix is effective, no OOB.") + + except RuntimeError as exc: + pytest.fail(f"MTE OOB triggered — DiscreteMaskAccessConversionPass fix " + f"may not be applied or is incomplete.\nError: {exc}") + + finally: + if in_tensor is not None: + del in_tensor + for t in reversed(pre_fillers): + del t + torch.npu.empty_cache() + print("[mte] Memory released.") diff --git a/third_party/ascend/unittest/pytest_ut/test_discrete_overlap_mask.py b/third_party/ascend/unittest/pytest_ut/test_discrete_overlap_mask.py new file mode 100644 index 0000000000..2d4f6437dc --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_discrete_overlap_mask.py @@ -0,0 +1,235 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +# --------------------------------------------------------------------------- +# Fixed Constants +# --------------------------------------------------------------------------- +_M_ROWS = 16 # Rows per program +_OFFS = 8 # Write window offset for two programs +_HALF = 12 # Mask threshold +_NUM_C = 24 # Rows of matrix C (= OFFS + M_ROWS, ensure pid=1 write window not out of bounds) + +assert _OFFS < _HALF < _M_ROWS, "OFFS < HALF < M_ROWS to ensure True/False on both sides" +assert _NUM_C >= _OFFS + _M_ROWS, "NUM_C must accommodate upper bound of pid=1 write window" + + +# --------------------------------------------------------------------------- +# Triton Kernel +# --------------------------------------------------------------------------- +@triton.jit +def _copy_matrix_kernel( + A_ptr, + idx_ptr, + C_ptr, + idx_stride, + A_row_stride, + A_col_stride, + C_row_stride, + C_col_stride, + BLOCK_N: tl.constexpr, + HALF: tl.constexpr, +): + """ + Discrete memory access + overlapping write window + runtime mask. + + pid=0 write window: rows [0, 15], mask=True when idx < HALF + pid=1 write window: rows [8, 23], mask=True when idx >= HALF + Overlap region : rows [8, 15] -> triggers load-select-store RMW + """ + program_id = tl.program_id(axis=0).to(tl.int64) + N_id = tl.program_id(axis=1).to(tl.int64) + + OFFS: tl.constexpr = 8 + M_ROWS: tl.constexpr = 16 + + N_BLOCK = N_id * BLOCK_N + tl.arange(0, BLOCK_N) # shape: (BLOCK_N,) + M_BLOCK = tl.arange(0, M_ROWS) # shape: (M_ROWS,) + + # Discrete row indices (loaded at runtime -> mask cannot be statically analyzed) + idx = tl.load(idx_ptr + program_id * idx_stride + M_BLOCK) + + # Runtime mask (generates scf.if -> compiler converts to load-select-store) + if program_id == 0: + mask = idx < HALF + else: + mask = idx >= HALF + + val = tl.load( + A_ptr + idx[:, None] * A_row_stride + tl.arange(0, BLOCK_N)[None, :] * A_col_stride, + mask=mask[:, None], + ) + + # Write to C (mask=False rows rely on load-select-store to preserve original values) + tl.store( + C_ptr + (OFFS * program_id + M_BLOCK[:, None]) * C_row_stride + N_BLOCK[None, :] * C_col_stride, + val, + mask=mask[:, None], + ) + # C (24 × N) Program 0 Program 1 + # Row 0~7 ──────────── write value (mask=True) ── Not involved + # Row 8~11 ──────────── write value (mask=True) ── Not written (mask=False, overwritten by P0 to 0) + # Row 12~15 ──────────── Not written (mask=False) ── write value (mask=True) + # Row 16~23 ──────────── Not involved ── write value (mask=True) + + +# --------------------------------------------------------------------------- +# Helper: Construct discrete index vector +# --------------------------------------------------------------------------- +def _make_idx(device: str) -> torch.Tensor: + """ + Construct 2x16 index matrix that meets mask distribution requirements. + + pid=0 row (idx0): + First HALF=12 values ∈ [0, HALF) -> mask=True + Last 4 values ∈ [HALF, M_ROWS) -> mask=False + pid=1 row (idx1): + First 4 values ∈ [OFFS, OFFS+4) -> mask=False + Last HALF=12 values ∈ [HALF, HALF*2) -> mask=True + """ + + def shuffle_quads(lst: list) -> list: + """Reverse each group of 4 elements (ignore if less than 4).""" + out = lst[:] + for i in range(0, len(out) - 3, 4): + out[i], out[i + 1], out[i + 2], out[i + 3] = \ + out[i + 3], out[i + 2], out[i + 1], out[i] + return out + + num_false = _M_ROWS - _HALF # = 4 + + seg0_true = shuffle_quads(list(range(0, _HALF))) # 12 values, < 12 + seg0_false = shuffle_quads(list(range(_HALF, _HALF + num_false))) # 4 values, >= 12 + idx0 = seg0_true + seg0_false # Total length 16 + + seg1_false = shuffle_quads(list(range(_OFFS, _OFFS + num_false))) # 4 values, < 12 + seg1_true = shuffle_quads(list(range(_HALF, _HALF + _HALF))) # 12 values, >= 12 + idx1 = seg1_false + seg1_true # Total length 16 + + assert len(idx0) == _M_ROWS, f"idx0 length error: {len(idx0)}" + assert len(idx1) == _M_ROWS, f"idx1 length error: {len(idx1)}" + assert all(v < _HALF for v in seg0_true), "pid=0 True segment should all be < HALF" + assert all(v >= _HALF for v in seg0_false), "pid=0 False segment should all be >= HALF" + assert all(v < _HALF for v in seg1_false), "pid=1 False segment should all be < HALF" + assert all(v >= _HALF for v in seg1_true), "pid=1 True segment should all be >= HALF" + + return torch.tensor([idx0, idx1], dtype=torch.int32, device=device) + + +# --------------------------------------------------------------------------- +# Dtype Mapping +# --------------------------------------------------------------------------- +_DTYPE_MAP = { + 'int32': torch.int32, + 'float32': torch.float32, + 'float16': torch.float16, + 'int16': torch.int16, +} + + +# --------------------------------------------------------------------------- +# Single Execution + Verification +# --------------------------------------------------------------------------- +def _run_once(BLOCK_N: int, dtype_str: str) -> None: + """ + Execute kernel once and verify results. + + Expectations: + C[0:HALF, :] all 0 -- pid=0 writes rows [0,HALF) of A (all 0) + C[HALF:NUM_C, :] all 1 -- pid=1 writes rows [HALF, NUM_C) of A (all 1) + """ + dev = 'npu' + td = _DTYPE_MAP[dtype_str] + zero_val = 0.0 if dtype_str.startswith('float') else 0 + one_val = 1.0 if dtype_str.startswith('float') else 1 + + # A: First HALF rows all 0, last HALF rows all 1 + A = torch.zeros((_NUM_C, BLOCK_N), dtype=td, device=dev) + A[_HALF:, :] = one_val + + idx = _make_idx(dev) + + # C: Fill all with 2 + C = torch.full((_NUM_C, BLOCK_N), 2, dtype=td, device=dev) + + grid = (2, 1) + _copy_matrix_kernel[grid]( + A_ptr=A, + idx_ptr=idx, + C_ptr=C, + idx_stride=idx.stride(0), + A_row_stride=A.stride(0), + A_col_stride=A.stride(1), + C_row_stride=C.stride(0), + C_col_stride=C.stride(1), + BLOCK_N=BLOCK_N, + HALF=_HALF, + enable_sync_block_lock=True, + ) + + # Verification + assert torch.all(C[:_HALF] == zero_val), ( + f"[dtype={dtype_str}, BLOCK_N={BLOCK_N}] " + f"C[:HALF] should all be {zero_val}, actual unique values: {C[:_HALF].unique().tolist()}") + assert torch.all(C[_HALF:] == one_val), ( + f"[dtype={dtype_str}, BLOCK_N={BLOCK_N}] " + f"C[HALF:] should all be {one_val}, actual unique values: {C[_HALF:].unique().tolist()}") + + +@pytest.mark.parametrize("param_list", [ + # --- int32 --- + (16, 'int32'), + (32, 'int32'), + (64, 'int32'), + # --- float32 --- + (16, 'float32'), + (32, 'float32'), + (64, 'float32'), +]) +def test_discrete_overlap_mask(param_list): + """ + Verify no precision issues in discrete access + overlapping write window + runtime mask scenario. + + Race condition errors are probabilistic. Each parameter combination is executed 10 times + to fully cover concurrent timing scenarios. + If sync_block_lock fix is effective, all 10 runs pass; if race condition exists, assertion failure + occurs with high probability. + """ + BLOCK_N, dtype_str = param_list + for _ in range(10): + _run_once(BLOCK_N, dtype_str) + + +if __name__ == "__main__": + configs = [ + (32, 'int32'), + (32, 'float32'), + ] + for BLOCK_N, dtype_str in configs: + print(f"Testing BLOCK_N={BLOCK_N}, dtype={dtype_str} ...", end=" ", flush=True) + for _ in range(10): + _run_once(BLOCK_N, dtype_str) + print("PASS (10 rounds)") + print("All tests passed.") diff --git a/third_party/ascend/unittest/pytest_ut/test_div.py b/third_party/ascend/unittest/pytest_ut/test_div.py new file mode 100644 index 0000000000..a145b32fbf --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_div.py @@ -0,0 +1,63 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import numpy as np +import torch +import pytest +import test_common + + +def torch_pointwise(x0, x1): + res = x0 / x1 + return res + + +@triton.jit +def triton_div(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop1 in range(loops1): + x0_prime = offset + (loop1 * XBLOCK_SUB) + base1 + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.load(in_ptr1 + (x0), None) + tmp2 = tmp0 / tmp1 + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], +]) +def test_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + if dtype == 'int8': + dtype = 'float32' + x1 = x1.masked_fill(x1 == 0, 1) + y_ref = torch_pointwise(x0, x1) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_div[ncore, 1, 1](x0, x1, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_dot.py b/third_party/ascend/unittest/pytest_ut/test_dot.py new file mode 100644 index 0000000000..635df74c5e --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_dot.py @@ -0,0 +1,181 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl +import test_common + + +@pytest.fixture(scope="function") +def restore_npu_hf32_setting(): + original_allow_hf32 = torch_npu.npu.matmul.allow_hf32 + try: + torch_npu.npu.matmul.allow_hf32 = True + yield + finally: + torch_npu.npu.matmul.allow_hf32 = original_allow_hf32 + + +def torch_dot_None(x0, x1): + res = torch.matmul(x0, x1) + return res + + +@triton.jit +def triton_dot_2_None(output_ptr, x_ptr, y_ptr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr): + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + + x_mask = (bidx[:, None] < B) & (cidx[None, :] < C) + y_mask = (cidx[:, None] < C) & (didx[None, :] < D) + out_mask = (bidx[:, None] < B) & (didx[None, :] < D) + Xidx = bidx[:, None] * C + cidx[None, :] + Yidx = cidx[:, None] * D + didx[None, :] + X = tl.load(x_ptr + Xidx, mask=x_mask, other=0.0) + Y = tl.load(y_ptr + Yidx, mask=y_mask, other=0.0) + ret = tl.dot(X, Y, input_precision="hf32") + oidx = bidx[:, None] * D + didx[None, :] + tl.store(output_ptr + oidx, ret, mask=out_mask) + + +@triton.jit +def triton_dot_2_allow_tf32(output_ptr, x_ptr, y_ptr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr): + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + + x_mask = (bidx[:, None] < B) & (cidx[None, :] < C) + y_mask = (cidx[:, None] < C) & (didx[None, :] < D) + out_mask = (bidx[:, None] < B) & (didx[None, :] < D) + Xidx = bidx[:, None] * C + cidx[None, :] + Yidx = cidx[:, None] * D + didx[None, :] + X = tl.load(x_ptr + Xidx, mask=x_mask, other=0.0) + Y = tl.load(y_ptr + Yidx, mask=y_mask, other=0.0) + ret = tl.dot(X, Y, allow_tf32=True) + oidx = bidx[:, None] * D + didx[None, :] + tl.store(output_ptr + oidx, ret, mask=out_mask) + + +@triton.jit +def triton_dot_2_input_tf32(output_ptr, x_ptr, y_ptr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr): + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + + x_mask = (bidx[:, None] < B) & (cidx[None, :] < C) + y_mask = (cidx[:, None] < C) & (didx[None, :] < D) + out_mask = (bidx[:, None] < B) & (didx[None, :] < D) + Xidx = bidx[:, None] * C + cidx[None, :] + Yidx = cidx[:, None] * D + didx[None, :] + X = tl.load(x_ptr + Xidx, mask=x_mask, other=0.0) + Y = tl.load(y_ptr + Yidx, mask=y_mask, other=0.0) + ret = tl.dot(X, Y, input_precision="tf32") + oidx = bidx[:, None] * D + didx[None, :] + tl.store(output_ptr + oidx, ret, mask=out_mask) + + +@triton.jit +def triton_dot_2_ignore_tf32(output_ptr, x_ptr, y_ptr, B: tl.constexpr, C: tl.constexpr, D: tl.constexpr): + bidx = tl.arange(0, B) + cidx = tl.arange(0, C) + didx = tl.arange(0, D) + + x_mask = (bidx[:, None] < B) & (cidx[None, :] < C) + y_mask = (cidx[:, None] < C) & (didx[None, :] < D) + out_mask = (bidx[:, None] < B) & (didx[None, :] < D) + Xidx = bidx[:, None] * C + cidx[None, :] + Yidx = cidx[:, None] * D + didx[None, :] + X = tl.load(x_ptr + Xidx, mask=x_mask, other=0.0) + Y = tl.load(y_ptr + Yidx, mask=y_mask, other=0.0) + ret = tl.dot(X, Y, input_precision="hf32") + oidx = bidx[:, None] * D + didx[None, :] + tl.store(output_ptr + oidx, ret, mask=out_mask) + + +testlist1 = [ + (10, 13, 35, 39), +] + +testlist2 = [(16, 32, 16)] + +typelist = [ + 'float32', +] + + +@pytest.mark.skip(reason="not supported after the NPUIR is updated in April, and will be fixed later") +@pytest.mark.parametrize("B, C, D", testlist2) +@pytest.mark.parametrize("sigtype", typelist) +def test_dot_2(restore_npu_hf32_setting, sigtype, B, C, D): + x = test_common.generate_tensor((B, C), sigtype).npu() + y = test_common.generate_tensor((C, D), sigtype).npu() + z_ref = torch_dot_None(x, y).to(torch.float32) + z = torch.zeros((B, D), dtype=torch.float32).npu() + triton_dot_2_None[1, 1, 1](z, x, y, B, C, D) + test_common.validate_cmp(sigtype, z, z_ref) + + +@pytest.mark.xfail( + reason="Temporarily disabled: TA backend does not support allow_tf32 yet. Will be fixed in follow-up.") +@pytest.mark.parametrize("B, C, D", testlist2) +@pytest.mark.parametrize("sigtype", typelist) +def test_dot_2_allow_tf32(restore_npu_hf32_setting, sigtype, B, C, D): + x = test_common.generate_tensor((B, C), sigtype).npu() + y = test_common.generate_tensor((C, D), sigtype).npu() + z_ref = torch_dot_None(x, y).to(torch.float32) + z = torch.zeros((B, D), dtype=torch.float32).npu() + triton_dot_2_allow_tf32[1, 1, 1](z, x, y, B, C, D) + test_common.validate_cmp(sigtype, z, z_ref) + + +@pytest.mark.skip(reason="not supported after the NPUIR is updated in April, and will be fixed later") +@pytest.mark.parametrize("B, C, D", testlist2) +@pytest.mark.parametrize("sigtype", typelist) +def test_dot_2_input_tf32(restore_npu_hf32_setting, sigtype, B, C, D): + x = test_common.generate_tensor((B, C), sigtype).npu() + y = test_common.generate_tensor((C, D), sigtype).npu() + z_ref = torch_dot_None(x, y).to(torch.float32) + z = torch.zeros((B, D), dtype=torch.float32).npu() + triton_dot_2_input_tf32[1, 1, 1](z, x, y, B, C, D) + test_common.validate_cmp(sigtype, z, z_ref) + + +@pytest.mark.parametrize("B, C, D", testlist2) +@pytest.mark.parametrize("sigtype", typelist) +def test_dot_2_ignore_tf32(sigtype, B, C, D): + input_type = "bfloat16" + x = test_common.generate_tensor((B, C), input_type).npu() + y = test_common.generate_tensor((C, D), input_type).npu() + z = torch.zeros((B, D), dtype=torch.float32).npu() + + original_allow_hf32 = torch_npu.npu.matmul.allow_hf32 + try: + torch_npu.npu.matmul.allow_hf32 = False + z_ref = torch_dot_None(x.to(torch.float32), y.to(torch.float32)).to(torch.float32) + + finally: + torch_npu.npu.matmul.allow_hf32 = original_allow_hf32 + + triton_dot_2_ignore_tf32[1, 1, 1](z, x, y, B, C, D) + test_common.validate_cmp(sigtype, z, z_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_dot_scaled.py b/third_party/ascend/unittest/pytest_ut/test_dot_scaled.py new file mode 100644 index 0000000000..9b0f0b6b62 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_dot_scaled.py @@ -0,0 +1,224 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import contextlib +import itertools +import re +import math +import textwrap +import os +import inspect +import pathlib + +import numpy as np +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +from numpy.random import RandomState +from triton.language.extra import libdevice + + +@pytest.mark.parametrize("M, N, K, rhs_scale, normal_type, acc_num, num_warps", + [(M, N, K, rhs_scale, normal_type, acc_num, 4) + for M, N, K in itertools.product([32, 64], [32, 64], [32]) + for rhs_scale in [False, True] + for normal_type in ["bf16", "fp16"] + for acc_num in [None, 1, 2]]) +def test_scaled_dot(M, N, K, rhs_scale, normal_type, num_warps, acc_num): + device = "npu" + + @triton.jit + def dot_scale_kernel(a_base, stride_a0, stride_a1, a_scale, b_base, stride_b0, stride_b1, b_scale, out, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, type_a: tl.constexpr, + type_b: tl.constexpr, acc_num: tl.constexpr): + + PACKED_BLOCK_K_A: tl.constexpr = BLOCK_K + PACKED_BLOCK_K_B: tl.constexpr = BLOCK_K + a_ptr = a_base + tl.arange(0, BLOCK_M)[:, None] * stride_a0 + tl.arange(0, + PACKED_BLOCK_K_A)[None, :] * stride_a1 + b_ptr = b_base + tl.arange(0, PACKED_BLOCK_K_B)[:, None] * stride_b0 + tl.arange(0, + BLOCK_N)[None, :] * stride_b1 + + a = tl.load(a_ptr) + b = tl.load(b_ptr) + SCALE_BLOCK_K: tl.constexpr = BLOCK_K // 32 + accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + if a_scale is not None: + scale_a_ptr = a_scale + tl.arange(0, BLOCK_M)[:, None] * SCALE_BLOCK_K + tl.arange(0, + SCALE_BLOCK_K)[None, :] + a_scale = tl.load(scale_a_ptr) + if b_scale is not None: + scale_b_ptr = b_scale + tl.arange(0, BLOCK_N)[:, None] * SCALE_BLOCK_K + tl.arange(0, + SCALE_BLOCK_K)[None, :] + b_scale = tl.load(scale_b_ptr) + accumulator = tl.dot_scaled(a, a_scale, type_a, b, b_scale, type_b, acc=accumulator, out_dtype=tl.float32) + if acc_num is not None: + for _ in range(acc_num): + accumulator = tl.dot_scaled(a, a_scale, type_a, b, b_scale, type_b, acc=accumulator, + out_dtype=tl.float32) + + out_ptr = out + tl.arange(0, BLOCK_M)[:, None] * BLOCK_N + tl.arange(0, BLOCK_N)[None, :] + tl.store(out_ptr, accumulator.to(a.dtype)) + + # The max exponent we use to initialize data in the x/y and associated scale tensor to avoid + # overflow when scaling. + comp_dtype_max_exp = 6 if normal_type == "fp16" else 15 + + torch.manual_seed(0) + + def make_arg(shape, ty): + if ty == "bf16" or ty == "fp16": + comp_dtype = torch.float16 if ty == "fp16" else torch.bfloat16 + ret = torch.randn(shape, dtype=comp_dtype, device=device) + # Clamp to avoid relative error issues + ret.clamp_(-2**comp_dtype_max_exp, 2**comp_dtype_max_exp - 1) + else: + ret = torch.randint(256, shape, dtype=torch.int8, device=device) + return ret + + type_a = normal_type + type_b = type_a + + x = make_arg((M, K), type_a) + y = make_arg((K, N), type_b) + + min_scale, max_scale = (0, 142) if type_a == torch.bfloat16 else (124, 131) + scale_x = torch.randint(min_scale - 128, max_scale - 127, (M, K // 32), dtype=torch.int8, device=device) + min_scale, max_scale = (0, 142) if type_b == torch.bfloat16 else (124, 131) + scale_y = torch.randint(min_scale - 128, max_scale - 127, (N, K // 32), dtype=torch.int8, device=device) + + if not rhs_scale: + scale_y = None + + def golden_ref(x, scale_x, y, scale_y): + shape_expand_x = x.shape[-1] // scale_x.shape[-1] + if x.dtype == torch.bfloat16: + upscale_x = scale_x.repeat_interleave(shape_expand_x, dim=1).to(torch.int16) + upscale_x = (upscale_x + 127 << 7).view(torch.bfloat16) + else: + scale_fp32 = scale_x.repeat_interleave(shape_expand_x, dim=1).to(torch.int32) + scale_fp32 = (scale_fp32 + 127 << 23).view(torch.float32) + upscale_x = scale_fp32.to(torch.float16) + upscale_y = None + if scale_y is None: + upscale_y = torch.ones_like(y) + else: + scale_y = scale_y.T + shape_expand_y = y.shape[0] // scale_y.shape[0] + if y.dtype == torch.bfloat16: + upscale_y = scale_y.repeat_interleave(shape_expand_y, dim=0).to(torch.int16) + upscale_y = (upscale_y + 127 << 7).view(torch.bfloat16) + else: + scale_fp32 = scale_y.repeat_interleave(shape_expand_y, dim=0).to(torch.int32) + scale_fp32 = (scale_fp32 + 127 << 23).view(torch.float32) + upscale_y = scale_fp32.to(torch.float16) + ret = torch.matmul(x * upscale_x, y * upscale_y) + return ret + + kernel_kwargs = {"num_warps": num_warps} + z = x.new_empty((M, N), dtype=x.dtype) + pgm = dot_scale_kernel[(1, )](x, *x.stride(), scale_x, y, *y.stride(), scale_y, z, M, N, K, type_a, type_b, acc_num, + **kernel_kwargs) + z_ref = golden_ref(x, scale_x, y, scale_y) + if acc_num is not None: + z_ref = z_ref * (acc_num + 1) + + atol = 1e-5 + rtol = 1e-2 + torch.testing.assert_close(z, z_ref, atol=atol, rtol=rtol) + + +@pytest.mark.parametrize("normal_type", ["bf16", "fp16"]) +def test_scaled_dot_fast_math(normal_type): + device = "npu" + m = n = k = 32 + + @triton.jit + def dot_scale_kernel(a_base, stride_a0, stride_a1, a_scale, b_base, stride_b0, stride_b1, b_scale, out, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, type_a: tl.constexpr, + type_b: tl.constexpr, fast_math: tl.constexpr): + + a_ptr = a_base + tl.arange(0, BLOCK_M)[:, None] * stride_a0 + tl.arange(0, BLOCK_K)[None, :] * stride_a1 + b_ptr = b_base + tl.arange(0, BLOCK_K)[:, None] * stride_b0 + tl.arange(0, BLOCK_N)[None, :] * stride_b1 + a = tl.load(a_ptr) + b = tl.load(b_ptr) + + scale_block_k: tl.constexpr = BLOCK_K // 32 + scale_a_ptr = a_scale + tl.arange(0, BLOCK_M)[:, None] * scale_block_k + tl.arange(0, scale_block_k)[None, :] + scale_b_ptr = b_scale + tl.arange(0, BLOCK_N)[:, None] * scale_block_k + tl.arange(0, scale_block_k)[None, :] + a_scale = tl.load(scale_a_ptr) + b_scale = tl.load(scale_b_ptr) + + out_tensor = tl.dot_scaled(a, a_scale, type_a, b, b_scale, type_b, fast_math=fast_math, out_dtype=tl.float32) + out_ptr = out + tl.arange(0, BLOCK_M)[:, None] * BLOCK_N + tl.arange(0, BLOCK_N)[None, :] + tl.store(out_ptr, out_tensor.to(a.dtype)) + + def make_scale_tensor(scale, data_dtype): + if data_dtype == torch.bfloat16: + scale_i16 = scale.to(torch.int16) + return ((scale_i16 + 127) << 7).view(torch.bfloat16) + scale_fp32 = ((scale.to(torch.int32) + 127) << 23).view(torch.float32) + return scale_fp32.to(torch.float16) + + def golden_ref(x, scale_x, y, scale_y, fast_math): + upscale_x = make_scale_tensor(scale_x.repeat_interleave(x.shape[1] // scale_x.shape[1], dim=1), x.dtype) + scale_y_t = scale_y.T + upscale_y = make_scale_tensor(scale_y_t.repeat_interleave(y.shape[0] // scale_y_t.shape[0], dim=0), y.dtype) + scaled_x = x * upscale_x + scaled_y = y * upscale_y + + if not fast_math: + lhs_nan_mask = scale_x.eq(-1).repeat_interleave(x.shape[1] // scale_x.shape[1], dim=1) + rhs_nan_mask = scale_y_t.eq(-1).repeat_interleave(y.shape[0] // scale_y_t.shape[0], dim=0) + scaled_x = torch.where(lhs_nan_mask, torch.full_like(scaled_x, float("nan")), scaled_x) + scaled_y = torch.where(rhs_nan_mask, torch.full_like(scaled_y, float("nan")), scaled_y) + + return torch.matmul(scaled_x, scaled_y) + + dtype = torch.float16 if normal_type == "fp16" else torch.bfloat16 + x = torch.full((m, k), 2.0, dtype=dtype, device=device) + y = torch.full((k, n), 3.0, dtype=dtype, device=device) + scale_x = torch.zeros((m, k // 32), dtype=torch.int8, device=device) + scale_y = torch.zeros((n, k // 32), dtype=torch.int8, device=device) + scale_x[3, 0] = -1 + scale_y[5, 0] = -1 + + def run_kernel(fast_math): + out = torch.empty((m, n), dtype=dtype, device=device) + dot_scale_kernel[(1, )](x, *x.stride(), scale_x, y, *y.stride(), scale_y, out, m, n, k, normal_type, + normal_type, fast_math, num_warps=4) + return out + + out_fast_math = run_kernel(True) + out_precise = run_kernel(False) + + ref_fast_math = golden_ref(x, scale_x, y, scale_y, fast_math=True) + ref_precise = golden_ref(x, scale_x, y, scale_y, fast_math=False) + + assert not torch.isnan(out_fast_math).any() + assert torch.isnan(out_precise[3, :]).all() + assert torch.isnan(out_precise[:, 5]).all() + assert not torch.isnan(out_precise[:3, :5]).any() + + torch.testing.assert_close(out_fast_math, ref_fast_math, atol=1e-3, rtol=1e-2) + torch.testing.assert_close(out_precise, ref_precise, atol=1e-3, rtol=1e-2, equal_nan=True) diff --git a/third_party/ascend/unittest/pytest_ut/test_downgrade.py b/third_party/ascend/unittest/pytest_ut/test_downgrade.py new file mode 100644 index 0000000000..810fdc7b92 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_downgrade.py @@ -0,0 +1,49 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +from triton.backends.ascend.utils import downgrade_llir, _downgrade_mem_attrs, _downgrade_stacksaverestore_intrinsics + + +@pytest.mark.parametrize("new_attr,legacy_attrs", [ + ("memory(none)", ["readnone"]), + ("memory(read)", ["readonly"]), + ("memory(write)", ["writeonly"]), + ("memory(readwrite)", []), + ("memory(argmem: read)", ["readonly", "argmemonly"]), + ("memory(argmem: read, inaccessiblemem: write)", ["inaccessiblemem_or_argmemonly"]), + ("memory(read, argmem: readwrite)", []), + ("memory(readwrite, argmem: none)", []), +]) +def test_mem_attrs(new_attr, legacy_attrs): + assert _downgrade_mem_attrs(new_attr).strip().split() == legacy_attrs + + +@pytest.mark.parametrize("new_intr,legacy_intr", [ + ("declare ptr @llvm.stacksave.p0()", "declare ptr @llvm.stacksave()"), + ("declare ptr addrspace(5) @llvm.stacksave.p5()", "declare ptr addrspace(5) @llvm.stacksave()"), + ("declare void @llvm.stackrestore.p0(ptr %ptr)", "declare void @llvm.stackrestore(ptr %ptr)"), + ("declare void @llvm.stackrestore.p5(ptr addrspace(5) %ptr)", + "declare void @llvm.stackrestore(ptr addrspace(5) %ptr)"), + ("%53 = call ptr @llvm.stacksave.p0()", "%53 = call ptr @llvm.stacksave()"), + ("call void @llvm.stackrestore.p0(ptr %53)", "call void @llvm.stackrestore(ptr %53)"), +]) +def test_stacksaverestore_intrinsics(new_intr, legacy_intr): + assert _downgrade_stacksaverestore_intrinsics(new_intr).strip() == legacy_intr diff --git a/third_party/ascend/unittest/pytest_ut/test_elementwise_ceil.py b/third_party/ascend/unittest/pytest_ut/test_elementwise_ceil.py new file mode 100644 index 0000000000..690cc0bc81 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_elementwise_ceil.py @@ -0,0 +1,86 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time +import test_common +import os +import shutil + +import torch +import torch_npu + + +def standard_ceil(x0): + res = torch.ceil(x0) + return res + + +@triton.jit +def triton_ceil(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + mask = idx_block < N + x = tl.load(in_ptr0 + idx_block, mask=mask) + res = tl.ceil(x) + tl.store(out_ptr0 + idx_block, res, mask=mask) + + +types = [ + (torch.float32, 'float32'), +] + +# if shape axis = 32/256 , then actual shape = axis/element_size() +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + +ops = [ + ('ceil', triton_ceil, standard_ceil), +] + + +@pytest.mark.parametrize('opName, tritonOp, standOp', ops) +@pytest.mark.parametrize('dtype, sigtype', types) +@pytest.mark.parametrize('N, NUMEL', shapes) +def test_elementwise_common(opName, tritonOp, standOp, dtype, sigtype, N, NUMEL): + torch.manual_seed(0) + torch_npu.npu.utils.set_device(0) + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + N = map_for_64_t[N] if N in map_for_64_t else N + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standOp(x0) + x0 = x0.npu() + + output = torch.zeros((N, ), dtype=dtype).npu() + tritonOp[1, 1, 1](x0, output, N=N, NUMEL=NUMEL, debug=True) + test_common.validate_cmp(sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_elementwise_clip.py b/third_party/ascend/unittest/pytest_ut/test_elementwise_clip.py new file mode 100644 index 0000000000..5dce754aaf --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_elementwise_clip.py @@ -0,0 +1,88 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time +import test_common +import os +import shutil + +import torch +import torch_npu + + +def standard_clamp(x0): + res = torch.clamp(x0, min=-10, max=10) + return res + + +@triton.jit +def triton_clamp(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + mask = idx_block < N + x = tl.load(in_ptr0 + idx_block, mask=mask) + res = tl.clamp(x, -10, 10) + tl.store(out_ptr0 + idx_block, res, mask=mask) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + (torch.bfloat16, 'bfloat16'), +] + +# if shape axis = 32/256 , then actual shape = axis/element_size() +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + +ops = [ + ('clamp', triton_clamp, standard_clamp), +] + + +@pytest.mark.parametrize('opName, tritonOp, standOp', ops) +@pytest.mark.parametrize('dtype, sigtype', types) +@pytest.mark.parametrize('N, NUMEL', shapes) +def test_elementwise_common(opName, tritonOp, standOp, dtype, sigtype, N, NUMEL): + torch.manual_seed(0) + torch_npu.npu.utils.set_device(0) + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + N = map_for_64_t[N] if N in map_for_64_t else N + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standOp(x0) + x0 = x0.npu() + + output = torch.zeros((N, ), dtype=dtype).npu() + tritonOp[1, 1, 1](x0, output, N=N, NUMEL=NUMEL, debug=True) + test_common.validate_cmp(sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_elementwise_f2i.py b/third_party/ascend/unittest/pytest_ut/test_elementwise_f2i.py new file mode 100644 index 0000000000..a5bf1b21e7 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_elementwise_f2i.py @@ -0,0 +1,137 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time +import test_common +import os +import shutil + +import torch +import torch_npu + + +def standard_f2i32(x0): + res = x0.to(torch.int32) + return res + + +def standard_f2i8(x0): + res = x0.to(torch.int8) + return res + + +def standard_f2i16(x0): + res = x0.to(torch.int16) + return res + + +def standard_f2i64(x0): + res = x0.to(torch.int64) + return res + + +@triton.jit +def triton_f2i8(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + mask = idx_block < N + x = tl.load(in_ptr0 + idx_block, mask=mask) + res = tl.cast(x, tl.int8) + tl.store(out_ptr0 + idx_block, res, mask=mask) + + +@triton.jit +def triton_f2i16(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + mask = idx_block < N + x = tl.load(in_ptr0 + idx_block, mask=mask) + res = tl.cast(x, tl.int16) + tl.store(out_ptr0 + idx_block, res, mask=mask) + + +@triton.jit +def triton_f2i32(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + mask = idx_block < N + x = tl.load(in_ptr0 + idx_block, mask=mask) + res = tl.cast(x, tl.int32) + tl.store(out_ptr0 + idx_block, res, mask=mask) + + +@triton.jit +def triton_f2i64(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + mask = idx_block < N + x = tl.load(in_ptr0 + idx_block, mask=mask) + res = tl.cast(x, tl.int64) + tl.store(out_ptr0 + idx_block, res, mask=mask) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + (torch.bfloat16, 'bfloat16'), +] + +# if shape axis = 32/256 , then actual shape = axis/element_size() +shapes = [ + (3, 32), +] + +map_for_64_t = {37: 31} + +ops = [ + ('f2i8', triton_f2i8, standard_f2i8, 'int8'), + ('f2i16', triton_f2i16, standard_f2i16, 'int16'), + ('f2i32', triton_f2i32, standard_f2i32, 'int32'), + ('f2i64', triton_f2i64, standard_f2i64, 'int64'), +] + + +def continue_func(opName, d_type): + if 'f2i' in opName and 'int' in d_type: + return True + + +@pytest.mark.parametrize('opName, tritonOp, standOp, dst_sigtype', ops) +@pytest.mark.parametrize('dtype, sigtype', types) +@pytest.mark.parametrize('N, NUMEL', shapes) +def test_elementwise_common(opName, tritonOp, standOp, dst_sigtype, dtype, sigtype, N, NUMEL): + if continue_func(opName, sigtype): + return + + torch_npu.npu.utils.set_device(0) + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + N = map_for_64_t[N] if N in map_for_64_t else N + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standOp(x0) + x0 = x0.npu() + + output = test_common.generate_tensor(shape=(N, ), dtype=dst_sigtype).npu() + tritonOp[1, 1, 1](x0, output, N=N, NUMEL=NUMEL, debug=True) + + test_common.validate_cmp(dst_sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_elementwise_floor.py b/third_party/ascend/unittest/pytest_ut/test_elementwise_floor.py new file mode 100644 index 0000000000..c78b908710 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_elementwise_floor.py @@ -0,0 +1,86 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time +import test_common +import os +import shutil + +import torch +import torch_npu + + +def standard_floor(x0): + res = torch.floor(x0) + return res + + +@triton.jit +def triton_floor(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + mask = idx_block < N + x = tl.load(in_ptr0 + idx_block, mask=mask) + res = tl.floor(x) + tl.store(out_ptr0 + idx_block, res, mask=mask) + + +types = [ + (torch.float32, 'float32'), +] + +# if shape axis = 32/256 , then actual shape = axis/element_size() +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + +ops = [ + ('floor', triton_floor, standard_floor), +] + + +@pytest.mark.parametrize('opName, tritonOp, standOp', ops) +@pytest.mark.parametrize('dtype, sigtype', types) +@pytest.mark.parametrize('N, NUMEL', shapes) +def test_elementwise_common(opName, tritonOp, standOp, dtype, sigtype, N, NUMEL): + torch.manual_seed(0) + torch_npu.npu.utils.set_device(0) + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + N = map_for_64_t[N] if N in map_for_64_t else N + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standOp(x0) + x0 = x0.npu() + + output = torch.zeros((N, ), dtype=dtype).npu() + tritonOp[1, 1, 1](x0, output, N=N, NUMEL=NUMEL, debug=True) + test_common.validate_cmp(sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_elementwise_i2f.py b/third_party/ascend/unittest/pytest_ut/test_elementwise_i2f.py new file mode 100644 index 0000000000..d12738101a --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_elementwise_i2f.py @@ -0,0 +1,123 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time +import test_common +import os +import shutil + +import torch +import torch_npu + + +def standard_i2f_float32(x0): + res = x0.to(torch.float32) + return res + + +def standard_i2f_float16(x0): + res = x0.to(torch.float16) + return res + + +def standard_i2f_bfloat16(x0): + res = x0.to(torch.bfloat16) + return res + + +@triton.jit +def triton_i2f_float32(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + mask = idx_block < N + x = tl.load(in_ptr0 + idx_block, mask=mask) + res = tl.cast(x, tl.float32) + tl.store(out_ptr0 + idx_block, res, mask=mask) + + +@triton.jit +def triton_i2f_float16(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + mask = idx_block < N + x = tl.load(in_ptr0 + idx_block, mask=mask) + res = tl.cast(x, tl.float16) + tl.store(out_ptr0 + idx_block, res, mask=mask) + + +@triton.jit +def triton_i2f_bfloat16(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + mask = idx_block < N + x = tl.load(in_ptr0 + idx_block, mask=mask) + res = tl.cast(x, tl.bfloat16) + tl.store(out_ptr0 + idx_block, res, mask=mask) + + +types = [ + # (torch.int8, 'int8'), # TO BE FIXED i8 -> f16、bf16 + # (torch.int16, 'int16'), # TO BE FIXED i16 -> f32、bf16 + (torch.int32, 'int32'), # TO BE FIXED i32 -> f16、bf16 + # (torch.int64, 'int64'), # TO BE FIXED i64 -> bf16 +] + +# if shape axis = 32/256 , then actual shape = axis/element_size() +shapes = [ + (3, 32), +] + +map_for_64_t = {37: 31} + +ops = [ + # ('i2f16', triton_i2f_float16, standard_i2f_float16, 'float16'), + ('i2f32', triton_i2f_float32, standard_i2f_float32, 'float32'), + # ('i2fbf16', triton_i2f_bfloat16, standard_i2f_bfloat16, 'bfloat16'), +] + + +def continue_func(opName, d_type): + if 'i2f' in opName and 'float' in d_type: + return True + + +@pytest.mark.parametrize('opName, tritonOp, standOp, dst_sigtype', ops) +@pytest.mark.parametrize('dtype, sigtype', types) +@pytest.mark.parametrize('N, NUMEL', shapes) +def test_elementwise_common(opName, tritonOp, standOp, dst_sigtype, dtype, sigtype, N, NUMEL): + if continue_func(opName, sigtype): + return + + torch_npu.npu.utils.set_device(0) + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + N = map_for_64_t[N] if N in map_for_64_t else N + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standOp(x0) + x0 = x0.npu() + + output = test_common.generate_tensor(shape=(N, ), dtype=dst_sigtype).npu() + tritonOp[1, 1, 1](x0, output, N=N, NUMEL=NUMEL, debug=True) + + test_common.validate_cmp(dst_sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_elementwise_round.py b/third_party/ascend/unittest/pytest_ut/test_elementwise_round.py new file mode 100644 index 0000000000..2cd3acb6bc --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_elementwise_round.py @@ -0,0 +1,92 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common +import os +import shutil + +import torch +import torch_npu + + +def standard_round_to_nearest_neighbor_even(x0): # TO BE FIXED round to nearest neighbor even + res = torch.round(x0) + return res + + +def standard_round(x0): + res = torch.where(x0 >= 0, torch.floor(x0 + 0.5), torch.ceil(x0 - 0.5)) + return res + + +@triton.jit +def triton_round(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + mask = idx_block < N + x = tl.load(in_ptr0 + idx_block, mask=mask) + res = tl.where(x >= 0, tl.floor(x.to(tl.float32) + 0.5), tl.ceil(x.to(tl.float32) - 0.5)) + tl.store(out_ptr0 + idx_block, res, mask=mask) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + (torch.bfloat16, 'bfloat16'), +] + +# if shape axis = 32/256 , then actual shape = axis/element_size() +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + +ops = [ + ('round', triton_round, standard_round), +] + + +@pytest.mark.parametrize('opName, tritonOp, standOp', ops) +@pytest.mark.parametrize('dtype, sigtype', types) +@pytest.mark.parametrize('N, NUMEL', shapes) +def test_elementwise_common(opName, tritonOp, standOp, dtype, sigtype, N, NUMEL): + torch.manual_seed(0) + torch_npu.npu.utils.set_device(0) + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + N = map_for_64_t[N] if N in map_for_64_t else N + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standOp(x0) + x0 = x0.npu() + + output = torch.zeros((N, ), dtype=dtype).npu() + tritonOp[1, 1, 1](x0, output, N=N, NUMEL=NUMEL, debug=True) + test_common.validate_cmp(sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_eq.py b/third_party/ascend/unittest/pytest_ut/test_eq.py new file mode 100644 index 0000000000..35ad5458c3 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_eq.py @@ -0,0 +1,79 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + + +def standard_binary(x0, y0): + res = x0 == y0 + return res + + +@triton.jit +def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + y = tl.load(in_ptr1 + idx_block, mask=idx_block < N) + ret = x == y + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + # (torch.bfloat16, 'bfloat16'), + (torch.int8, 'int8'), + (torch.int16, 'int16'), + (torch.int32, 'int32'), + (torch.int64, 'int64'), +] + +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('N,NUMEL', shapes) +def test_elementwsie_common(dtype, sigtype, N, NUMEL): + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == "int64": + N = map_for_64_t[N] if N in map_for_64_t else N + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + y0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + ans = standard_binary(x0, y0) + out = torch.zeros((N, ), dtype=torch.bool).npu() + triton_elementwise_binary[1, 1, 1](x0, y0, out, N, NUMEL) + test_common.validate_cmp(sigtype, out, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_eq_2.py b/third_party/ascend/unittest/pytest_ut/test_eq_2.py new file mode 100644 index 0000000000..f99d537001 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_eq_2.py @@ -0,0 +1,76 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common + + +def torch_eq(x0, x1): + return x0 == x1 + + +def torch_eq_from_np(x0, x1, dtype): + return torch.from_numpy(x0 == x1).to(eval('torch.' + dtype)) + + +@triton.jit +def triton_eq(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x_index = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + x_index, None) + tmp1 = tl.load(in_ptr1 + x_index, None) + tmp2 = tmp0 == tmp1 + tl.store(out_ptr0 + x_index, tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + # ['float16', (2, 4096, 8), 2, 32768, 1024], + # ['bfloat16', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], + # ['int16', (2, 4096, 8), 2, 32768, 1024], + # ['int32', (2, 4096, 8), 2, 32768, 1024], + # ['int64', (2, 4096, 8), 2, 32768, 1024], + ['uint8', (2, 4096, 8), 2, 32768, 1024], + # ['uint16', (2, 4096, 8), 2, 32768, 1024], + # ['uint32', (2, 4096, 8), 2, 32768, 1024], + # ['uint64', (2, 4096, 8), 2, 32768, 1024], +]) +def test_eq(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + np_x0 = test_common.generate_numpy(shape, dtype) + x0 = torch.from_numpy(np_x0).to(eval('torch.' + dtype)).npu() + np_x1 = test_common.generate_numpy(shape, dtype) + x1 = torch.from_numpy(np_x1).to(eval('torch.' + dtype)).npu() + # torch结果 + torch_res = torch_eq_from_np(np_x0, np_x1, dtype).npu() + # triton结果 + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_eq[ncore, 1, 1](x0, x1, triton_res, xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_erfinv.py b/third_party/ascend/unittest/pytest_ut/test_erfinv.py new file mode 100644 index 0000000000..0c9a1ffe67 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_erfinv.py @@ -0,0 +1,156 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common +import triton.language.extra.cann.libdevice as libdevice + + +@triton.jit +def triton_erfinv(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + xindex = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = xindex < xnumel + x0 = tl.load(in_ptr0 + xindex, xmask) + y = libdevice.erfinv(x0) + tl.store(out_ptr0 + xindex, y, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_erfinv_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x = test_common.generate_tensor(shape, dtype).npu() + x[0][0][0] = 1 # erfinv(1) -> ∞ + x[0][0][1] = -1 # erfinv(-1) -> -∞ + + # Avoid numerical instability near ±1 + # Move values in (threshold, 1) to threshold and (-1, -threshold) to -threshold + threshold = 1 - 1.1e-4 + too_close_pos = (x > threshold) & (x < 1) + too_close_neg = (x < -threshold) & (x > -1) + x[too_close_pos] = threshold + x[too_close_neg] = -threshold + y_ref = torch.erfinv(x).npu() + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_erfinv[ncore, 1, 1](x, y_cal, x.numel(), xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_all_blocks_parallel(param_list, monkeypatch): + monkeypatch.setenv("TRITON_ALL_BLOCKS_PARALLEL", "1") + dtype, shape, ncore, xblock, xblock_sub = param_list + x = test_common.generate_tensor(shape, dtype).npu() + x[0][0][0] = 1 # erfinv(1) -> ∞ + x[0][0][1] = -1 # erfinv(-1) -> -∞ + + # Avoid numerical instability near ±1 + # Move values in (threshold, 1) to threshold and (-1, -threshold) to -threshold + threshold = 1 - 1.1e-4 + too_close_pos = (x > threshold) & (x < 1) + too_close_neg = (x < -threshold) & (x > -1) + x[too_close_pos] = threshold + x[too_close_neg] = -threshold + y_ref = torch.erfinv(x).npu() + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_erfinv[ncore, 1, 1](x, y_cal, x.numel(), xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) + monkeypatch.delenv("TRITON_ALL_BLOCKS_PARALLEL") + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_auto_blockify(param_list, monkeypatch): + monkeypatch.setenv("TRITON_ALL_BLOCKS_PARALLEL", "1") + dtype, shape, ncore, xblock, xblock_sub = param_list + x = test_common.generate_tensor(shape, dtype).npu() + x[0][0][0] = 1 # erfinv(1) -> ∞ + x[0][0][1] = -1 # erfinv(-1) -> -∞ + + # Avoid numerical instability near ±1 + # Move values in (threshold, 1) to threshold and (-1, -threshold) to -threshold + threshold = 1 - 1.1e-4 + too_close_pos = (x > threshold) & (x < 1) + too_close_neg = (x < -threshold) & (x > -1) + x[too_close_pos] = threshold + x[too_close_neg] = -threshold + y_ref = torch.erfinv(x).npu() + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_erfinv[ncore, 1, 1](x, y_cal, x.numel(), xblock, xblock_sub, auto_blockify_size=ncore) + test_common.validate_cmp(dtype, y_cal, y_ref) + monkeypatch.delenv("TRITON_ALL_BLOCKS_PARALLEL") + + +def prepare_erfinv_input(x: torch.Tensor, dtype: str) -> torch.Tensor: + """Inject special values and clip regions near ±1 to avoid numerical instability.""" + # Inject ±1 at known positions + if x.numel() >= 2: + with torch.no_grad(): + x_flat = x.view(-1) + x_flat[0] = 1.0 # erfinv(1) → +inf + x_flat[1] = -1.0 # erfinv(-1) → -inf + + # Threshold to avoid instability near ±1 + threshold = 1.0 - 1.1e-4 + too_close_pos = (x > threshold) & (x < 1.0) + too_close_neg = (x < -threshold) & (x > -1.0) + + with torch.no_grad(): + x[too_close_pos] = threshold + x[too_close_neg] = -threshold + + return x + + +@pytest.mark.parametrize("dtype", ["float32"]) +@pytest.mark.parametrize("shape", [ + (2, 4096, 8), +]) +def test_erfinv_all_blocks_parallel(dtype, shape, monkeypatch): + monkeypatch.setenv("TRITON_ALL_BLOCKS_PARALLEL", "1") + + XBLOCK = 1024 + XBLOCK_SUB = 128 + + x = test_common.generate_tensor(shape, dtype).npu() + x = prepare_erfinv_input(x, dtype) + + y_ref = torch.erfinv(x).npu() + y_cal = torch.zeros_like(x).npu() + numel = x.numel() + + def grid(meta): + return (triton.cdiv(numel, meta['XBLOCK']), ) + + triton_erfinv[grid](x, y_cal, numel, XBLOCK=XBLOCK, XBLOCK_SUB=XBLOCK_SUB) + + test_common.validate_cmp(dtype, y_cal, y_ref) + + monkeypatch.delenv("TRITON_ALL_BLOCKS_PARALLEL") diff --git a/third_party/ascend/unittest/pytest_ut/test_exp.py b/third_party/ascend/unittest/pytest_ut/test_exp.py new file mode 100644 index 0000000000..651e9d5eb3 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_exp.py @@ -0,0 +1,56 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import numpy as np +import torch +import pytest +import test_common + + +def torch_pointwise(x0): + res = torch.exp(x0) + return res + + +@triton.jit +def triton_exp(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop1 in range(loops1): + x0_prime = offset + (loop1 * XBLOCK_SUB) + base1 + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp2 = tl.exp(tmp0) + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_pointwise(x0) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_exp[ncore, 1, 1](x0, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_exp2.py b/third_party/ascend/unittest/pytest_ut/test_exp2.py new file mode 100644 index 0000000000..f45bdd13bd --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_exp2.py @@ -0,0 +1,59 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common + + +def torch_exp2(x0): + res = torch.pow(2, x0, out=None) + return res + + +@triton.jit +def triton_exp2(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x_index = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + x_index, None) + tmp1 = tl.exp2(tmp0) + tl.store(out_ptr0 + x_index, tmp1, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_exp2(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + torch_res = torch_exp2(x0) + # triton结果 + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_exp2[ncore, 1, 1](x0, triton_res, xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_exp_.py b/third_party/ascend/unittest/pytest_ut/test_exp_.py new file mode 100644 index 0000000000..ba1647185a --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_exp_.py @@ -0,0 +1,99 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + + +def standard_unary(x0, dtype): + res = torch.exp(x0) + return res + + +def standard_binary(x0, y0, dtype): + res = x0 + y0 + return res + + +@triton.jit +def triton_elementwise_unary(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + ret = tl.math.exp(x) + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +@triton.jit +def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + y = tl.load(in_ptr1 + idx_block, mask=idx_block < N) + ret = x + y + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + # (torch.bfloat16, 'bfloat16'), + # (torch.int8, 'int8'), + # (torch.int16, 'int16'), + # (torch.int32, 'int32'), + # (torch.int64, 'int64'), +] + +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('N,NUMEL', shapes) +def test_elementwsie_common(dtype, sigtype, N, NUMEL): + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == "int64": + N = map_for_64_t[N] if N in map_for_64_t else N + + print(f"elementwise : ({N},) {dtype} {sigtype}") + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standard_unary(x0, dtype) + x0 = x0.npu() + print(ans) + + out = torch.zeros((N, ), dtype=dtype).npu() + triton_elementwise_unary[1, 1, 1](x0, out, N=N, NUMEL=NUMEL, debug=True) + print(out) + + test_common.validate_cmp(sigtype, out, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_expand_dims.py b/third_party/ascend/unittest/pytest_ut/test_expand_dims.py new file mode 100644 index 0000000000..304ec9ff8d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_expand_dims.py @@ -0,0 +1,74 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.expand_dims(X, 2) + + oidx = xidx[:, None, None, None] * YB * ZB + yidx[None, :, None, None] * ZB + tl.arange( + 0, 1)[None, None, :, None] + zidx[None, None, None, :] + + tl.store(output_ptr + oidx, ret) + + +paras = [ + ('*fp32', eval('torch.float32'), 2, 256, 16), + ('*fp32', eval('torch.float32'), 8, 8, 4), + ('*fp16', eval('torch.float16'), 2, 256, 16), + ('*fp16', eval('torch.float16'), 8, 8, 4), + ('*i8', eval('torch.int8'), 2, 256, 16), + ('*i8', eval('torch.int8'), 8, 8, 4), +] + + +@pytest.mark.parametrize('para_type,data_type,XB,YB,ZB', paras) +def test_npu(para_type, data_type, XB, YB, ZB): + + x = torch.randint(low=-128, high=128, size=(XB, YB, ZB), dtype=data_type).npu() + a = x.unsqueeze(2) + + print(f"shape = {x.shape}") + print(x.dtype) + print(a[0, 0:16, 0, 0]) + + output = torch.randint(1, (XB, YB, 1, ZB), dtype=data_type).npu() + + print(f"output.dtype={output.dtype}") + + fn_npu_[1, 1, 1](output, x, XB=XB, YB=YB, ZB=ZB, debug=True) + print(output[0, 0:16, 0, 0]) + + torch.testing.assert_close(output, a) diff --git a/third_party/ascend/unittest/pytest_ut/test_expm1.py b/third_party/ascend/unittest/pytest_ut/test_expm1.py new file mode 100644 index 0000000000..7a06320d61 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_expm1.py @@ -0,0 +1,62 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + +import torch +import torch_npu + + +def torch_expm1(x0): + return torch.expm1(x0) + + +@triton.jit +def triton_expm1(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = libdevice.expm1(tmp0) + tl.store(out_ptr0 + (x0), tmp1, None) + + +@pytest.mark.skip(reason="expm1 failed sometimes, wait for fix") +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_expm1(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0_ref = test_common.generate_tensor(shape, dtype) + # error gains when input is large; expm1 used to handle small input when exp(x) close to 1 + x0_ref[x0_ref > 1] = 1 + y_ref = torch_expm1(x0_ref) + + x0 = x0_ref.npu() + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_expm1[ncore, 1, 1](x0, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_extract_slice.py b/third_party/ascend/unittest/pytest_ut/test_extract_slice.py new file mode 100644 index 0000000000..7014e73e1d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_extract_slice.py @@ -0,0 +1,60 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu + +import triton +import triton.language as tl +import triton.language.extra.cann.extension as extension + +import pytest + + +@triton.jit +def triton_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + out_sub = extension.extract_slice(output, [block_start], [32], [1]) + out_idx = block_start + tl.arange(0, 32) + out_msk = out_idx < n_elements + tl.store(output_ptr + out_idx, out_sub, mask=out_msk) + + +def triton_func(x: torch.Tensor, y: torch.Tensor): + output = torch.empty_like(x) + n_elements = output.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + triton_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024) + return output + + +def test_extract_slice(): + size = 1024 + x = torch.rand(size, device='npu') + y = torch.rand(size, device='npu') + torch_ref = x + y + triton_cal = triton_func(x, y) + torch.testing.assert_close(triton_cal[:32], torch_ref[:32]) diff --git a/third_party/ascend/unittest/pytest_ut/test_fast_dividef.py b/third_party/ascend/unittest/pytest_ut/test_fast_dividef.py new file mode 100644 index 0000000000..19cded059b --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_fast_dividef.py @@ -0,0 +1,59 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import numpy as np +import torch +import pytest +import test_common +import triton.language.extra.cann.libdevice as libdevice + + +def torch_pointwise(x0, x1): + res = x0 / x1 + return res + + +@triton.jit +def triton_fast_dividef(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.load(in_ptr1 + (x0), None) + tmp2 = libdevice.fast_dividef(tmp0, tmp1) + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], +]) +def test_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_pointwise(x0, x1) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_fast_dividef[ncore, 1, 1](x0, x1, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_fast_expf.py b/third_party/ascend/unittest/pytest_ut/test_fast_expf.py new file mode 100644 index 0000000000..fc8217500e --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_fast_expf.py @@ -0,0 +1,56 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import numpy as np +import torch +import pytest +import test_common +import triton.language.extra.cann.libdevice as libdevice + + +def torch_pointwise(x0): + res = torch.exp(x0) + return res + + +@triton.jit +def triton_fast_expf(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp2 = libdevice.fast_expf(tmp0) + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_pointwise(x0) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_fast_expf[ncore, 1, 1](x0, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_fdiv.py b/third_party/ascend/unittest/pytest_ut/test_fdiv.py new file mode 100644 index 0000000000..0714798a24 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_fdiv.py @@ -0,0 +1,63 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common + + +def torch_fdiv(x0, x1): + res = x0 / x1 + return res + + +@triton.jit +def triton_fdiv(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + x_index = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + tmp0 = tl.load(in_ptr0 + x_index) + tmp1 = tl.load(in_ptr1 + x_index) + tmp2 = tl.fdiv(tmp0, tmp1) + tl.store(out_ptr0 + x_index, tmp2) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], +]) +def test_fdiv(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + y_tmp = test_common.generate_tensor(shape, dtype) + y0 = y_tmp.masked_fill(y_tmp == 0, 1) + y0 = y0.npu() + + # torch结果 + y_ref = torch_fdiv(x0, y0).to(eval('torch.' + dtype)) + # triton结果 + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_fdiv[ncore, 1, 1](x0, y0, y_cal, xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_fixpipe.py b/third_party/ascend/unittest/pytest_ut/test_fixpipe.py new file mode 100644 index 0000000000..22b9aa98fe --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_fixpipe.py @@ -0,0 +1,84 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import os +import pytest +import triton +import triton.language as tl +import triton.extension.buffer.language as bl +import triton.language.extra.cann.extension as al +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir + +os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" + + +class Options: + num_warps = 4 + num_stages = 3 + num_ctas = 1 + cluster_dims = (1, 1, 1) + enable_fp_fusion = True + debug = False + arch = "Ascend910_95" + + +def compile_kernel(kernel, signature, constants): + """Helper to compile a kernel to MLIR.""" + src = ASTSource(kernel, signature, constants) + context = ir.context() + ir.load_dialects(context) + ascend_ir.load_dialects(context) + module = ast_to_ttir(kernel, src, context, Options(), {}, {}) + return str(module) + + +@triton.jit +def fixpipe( + A_ptr, + M: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, +): + + row_matmul = tl.program_id(0) + + offs_i = tl.arange(0, tl.constexpr(M))[:, None] # [M,1] (row axis) + offs_k = tl.arange(0, K) # [K] + + a_ptrs = A_ptr + (row_matmul + offs_i) * K + offs_k[None, :] + a_vals = tl.load(a_ptrs) # [M, K] + + ub = bl.alloc(tl.float32, [M, N], al.ascend_address_space.UB) + al.fixpipe(a_vals, ub, dual_dst_mode=al.FixpipeDualDstMode.NO_DUAL) + + +@pytest.mark.parametrize("M, K, N", [(16, 16, 16)]) +def test_fixpipe(M, K, N): + mlir = compile_kernel( + fixpipe, + { + "A_ptr": "*fp32", + }, + {"M": M, "K": K, "N": N}, + ) + assert len(mlir) > 0 diff --git a/third_party/ascend/unittest/pytest_ut/test_fixpipe_tensor_dst.py b/third_party/ascend/unittest/pytest_ut/test_fixpipe_tensor_dst.py new file mode 100644 index 0000000000..797e30fe60 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_fixpipe_tensor_dst.py @@ -0,0 +1,169 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import os +import pytest +import triton +import triton.language as tl +import triton.extension.buffer.language as bl +import triton.language.extra.cann.extension as al +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir +from triton.tools.get_ascend_devices import is_compile_on_910_95 + +os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" + + +class Options: + num_warps = 4 + num_stages = 3 + num_ctas = 1 + cluster_dims = (1, 1, 1) + enable_fp_fusion = True + debug = False + arch = "Ascend910_95" + + +def compile_kernel(kernel, signature, constants): + """Helper to compile a kernel to MLIR.""" + src = ASTSource(kernel, signature, constants) + context = ir.context() + ir.load_dialects(context) + ascend_ir.load_dialects(context) + module = ast_to_ttir(kernel, src, context, Options(), {}, {}) + return str(module) + + +@triton.jit +def fixpipe_without_dst( + A_ptr, + M: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, +): + row_matmul = tl.program_id(0) + + offs_i = tl.arange(0, tl.constexpr(M))[:, None] + offs_k = tl.arange(0, K) + + a_ptrs = A_ptr + (row_matmul + offs_i) * K + offs_k[None, :] + a_vals = tl.load(a_ptrs) + + result = al.fixpipe(a_vals, dual_dst_mode=al.FixpipeDualDstMode.NO_DUAL) + tl.store(a_ptrs, result) + + +@triton.jit +def fixpipe_with_buffer_dst( + A_ptr, + M: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, +): + row_matmul = tl.program_id(0) + + offs_i = tl.arange(0, tl.constexpr(M))[:, None] + offs_k = tl.arange(0, K) + + a_ptrs = A_ptr + (row_matmul + offs_i) * K + offs_k[None, :] + a_vals = tl.load(a_ptrs) + + ub = bl.alloc(tl.float32, [M, N], al.ascend_address_space.UB) + al.fixpipe(a_vals, ub, dual_dst_mode=al.FixpipeDualDstMode.NO_DUAL) + + +@pytest.mark.parametrize("M, K, N", [(16, 16, 16)]) +def test_fixpipe_without_dst(M, K, N): + mlir = compile_kernel( + fixpipe_without_dst, + {"A_ptr": "*fp32"}, + {"M": M, "K": K, "N": N}, + ) + assert len(mlir) > 0 + + +@triton.jit +def fixpipe_row_split( + A_ptr, + M: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, +): + row_matmul = tl.program_id(0) + + offs_i = tl.arange(0, tl.constexpr(M))[:, None] + offs_k = tl.arange(0, K) + + a_ptrs = A_ptr + (row_matmul + offs_i) * K + offs_k[None, :] + a_vals = tl.load(a_ptrs) + + result = al.fixpipe(a_vals, dual_dst_mode=al.FixpipeDualDstMode.ROW_SPLIT) + + +@triton.jit +def fixpipe_column_split( + A_ptr, + M: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, +): + row_matmul = tl.program_id(0) + + offs_i = tl.arange(0, tl.constexpr(M))[:, None] + offs_k = tl.arange(0, K) + + a_ptrs = A_ptr + (row_matmul + offs_i) * K + offs_k[None, :] + a_vals = tl.load(a_ptrs) + + result = al.fixpipe(a_vals, dual_dst_mode=al.FixpipeDualDstMode.COLUMN_SPLIT) + + +@pytest.mark.skipif(not is_compile_on_910_95, reason="only support A5") +@pytest.mark.parametrize("M, K, N", [(32, 16, 16)]) +def test_fixpipe_row_split(M, K, N): + mlir = compile_kernel( + fixpipe_row_split, + {"A_ptr": "*fp32"}, + {"M": M, "K": K, "N": N}, + ) + assert len(mlir) > 0 + + +@pytest.mark.skipif(not is_compile_on_910_95, reason="only support A5") +@pytest.mark.parametrize("M, K, N", [(16, 16, 32)]) +def test_fixpipe_column_split(M, K, N): + mlir = compile_kernel( + fixpipe_column_split, + {"A_ptr": "*fp32"}, + {"M": M, "K": K, "N": N}, + ) + assert len(mlir) > 0 + + +@pytest.mark.parametrize("M, K, N", [(16, 16, 16)]) +def test_fixpipe_with_buffer_dst(M, K, N): + mlir = compile_kernel( + fixpipe_with_buffer_dst, + {"A_ptr": "*fp32"}, + {"M": M, "K": K, "N": N}, + ) + assert len(mlir) > 0 diff --git a/third_party/ascend/unittest/pytest_ut/test_flip.py b/third_party/ascend/unittest/pytest_ut/test_flip.py new file mode 100644 index 0000000000..b1120234d4 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_flip.py @@ -0,0 +1,28 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common +from triton.runtime.libentry import libentry diff --git a/third_party/ascend/unittest/pytest_ut/test_floor.py b/third_party/ascend/unittest/pytest_ut/test_floor.py new file mode 100644 index 0000000000..13a83dbd91 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_floor.py @@ -0,0 +1,59 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common + + +def torch_floor(x0, x1): + res = x0 + torch.floor(x1) + return res + + +@triton.jit +def triton_floor(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + x_index = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = x_index < xnumel + tmp0 = tl.load(in_ptr0 + x_index, xmask) + tmp1 = tl.load(in_ptr1 + x_index, xmask) + tmp2 = tmp0 + tl.floor(tmp1) + tl.store(out_ptr0 + x_index, tmp2, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_floor(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + y_ref = torch_floor(x0, x1) + # triton结果 + y_cal = test_common.generate_tensor(shape, dtype).npu() + triton_floor[ncore, 1, 1](x0, x1, y_cal, x0.numel(), xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_floordiv.py b/third_party/ascend/unittest/pytest_ut/test_floordiv.py new file mode 100644 index 0000000000..1edc15866a --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_floordiv.py @@ -0,0 +1,72 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +def torch_func(x0, x1): + res = x0 // x1 + return res + + +@triton.jit +def triton_kernel(out_ptr0, in_ptr0, in_ptr1, N: tl.constexpr): + idx = tl.arange(0, N) + x = tl.load(in_ptr0 + idx) + y = tl.load(in_ptr1 + idx) + ret = x // y + tl.store(out_ptr0 + idx, ret) + + +def triton_func(x0, x1, N): + out = torch.empty_like(x0) + triton_kernel[1, 1, 1](out, x0, x1, N) + return out + + +types = [ + "int32", + "bool", +] + +shapes = [ + 3, + 32, + 37, + 256, + 781, +] + + +@pytest.mark.parametrize("sigtype", types) +@pytest.mark.parametrize("N", shapes) +def test_floordiv(sigtype, N): + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + x1 = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + x1 = x1.masked_fill(x1 == 0, 1) + + torch_ref = torch_func(x0, x1) + triton_cal = triton_func(x0, x1, N) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_for_ptr.py b/third_party/ascend/unittest/pytest_ut/test_for_ptr.py new file mode 100644 index 0000000000..47b8d28ba2 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_for_ptr.py @@ -0,0 +1,116 @@ +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + +types_all = [ + (torch.float32, 'float32'), +] + +shapes_common = [(128, 256), (127, 256), (127, 16), (129, 256), (77, 1024), (69, 512)] + +block_size = [128, 256, 1024] + + +def ceil_div(a, b): + return (a + b - 1) // b + + +def profiler_wrapper(fn, *args): + result_path = "./result_profiling_for" + skip_first = 10 + wait = 0 + warmup = 3 + active = 30 + repeat = 1 + stream = torch.npu.current_stream() + experimental_config = torch_npu.profiler._ExperimentalConfig( + aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, l2_cache=False, data_simplification=False) + with torch_npu.profiler.profile( + activities=[torch_npu.profiler.ProfilerActivity.CPU, torch_npu.profiler.ProfilerActivity.NPU], + schedule=torch_npu.profiler.schedule(wait=wait, warmup=warmup, active=active, repeat=repeat, + skip_first=skip_first), + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(result_path), record_shapes=True, + profile_memory=False, with_stack=False, with_flops=False, with_modules=False, + experimental_config=experimental_config) as prof: + stream.synchronize() + for _ in range(skip_first + (wait + warmup + active) * repeat): + fn(*args) + prof.step() + stream.synchronize() + + +@triton.jit +def for_ptr_kernel( + in_ptr, + output_ptr, + N: tl.constexpr, + M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, +): + pid = tl.program_id(axis=0) + offset = 2 * pid * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + x_in = in_ptr + offset + x_out = output_ptr + offset + for _ in range(0, 2): + mask = (offset < M * N) + data = tl.load(x_in, mask=mask, other=0) + tl.store(x_out, data, mask) + x_in += BLOCK_SIZE_N + x_out += BLOCK_SIZE_N + + +@pytest.mark.parametrize('dtype, sigtype', types_all) +@pytest.mark.parametrize('M, N', shapes_common) +@pytest.mark.parametrize('BLOCK_SIZE_N', block_size) +def test_for_ptr(M, N, BLOCK_SIZE_N, dtype, sigtype): + + in_tensor = torch.randn(M, N, dtype=dtype).npu() + + triton_output = torch.zeros_like(in_tensor) + + grid = (ceil_div(2 * M * N, BLOCK_SIZE_N), ) + + for_ptr_kernel[grid](in_tensor, triton_output, N=N, M=M, BLOCK_SIZE_N=BLOCK_SIZE_N, optimize_dynamic_offset=False) + + assert torch.allclose(triton_output, in_tensor, rtol=1e-5, atol=1e-8) + + +def triton_lfor_ptr(in_tensor, BLOCK_SIZE): + M = in_tensor.shape[0] + N = in_tensor.shape[1] + + triton_output = torch.zeros_like(in_tensor) + grid = (ceil_div(2 * M * N, BLOCK_SIZE), ) + + for_ptr_kernel[grid](in_tensor, triton_output, N=N, M=M, BLOCK_SIZE_N=BLOCK_SIZE, optimize_dynamic_offset=True) + + +def profile_performance_test(M, N, dtype, BLOCK_SIZE): + print(f"\nDetailed performance analysis: M={M}, N={N}, dtype={dtype}, block_size={BLOCK_SIZE}") + + in_tensor = torch.randn(2 * M, N, dtype=dtype).npu() + + def wrapper_func(x): + triton_lfor_ptr(x, BLOCK_SIZE=BLOCK_SIZE) + + # Run performance analysis + profiler_wrapper(wrapper_func, in_tensor) + + +if __name__ == "__main__": + print("For Kernel Performance Test Suite") + print("Function: Broadcast first element") + + # Optional: Run detailed profiler test (specific configuration) + profile_performance_test(512, 512, torch.float32, BLOCK_SIZE=1024) + + print("\n" + "=" * 80) + print("Test completed!") + print(f"Detailed performance analysis results saved in: ./result_profiling_for/") + print("=" * 80) diff --git a/third_party/ascend/unittest/pytest_ut/test_full.py b/third_party/ascend/unittest/pytest_ut/test_full.py new file mode 100644 index 0000000000..a366ace2c1 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_full.py @@ -0,0 +1,95 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu +import pytest + + +@triton.jit +def fn_npu_f32(output_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + ret = tl.full((XB, YB, ZB), value=100, dtype=tl.float32) + + oidx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_f16(output_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + ret = tl.full((XB, YB, ZB), value=100, dtype=tl.float16) + + oidx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_i8(output_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + ret = tl.full((XB, YB, ZB), value=100, dtype=tl.int8) + + oidx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +testlist = [ + (fn_npu_f32, 'float32', torch.float32, 2, 256, 16), + (fn_npu_f32, 'float32', torch.float32, 8, 8, 4), + (fn_npu_f16, 'float16', torch.float16, 2, 256, 16), + (fn_npu_f16, 'float16', torch.float16, 8, 8, 4), + (fn_npu_i8, 'int8', torch.int8, 2, 256, 16), + (fn_npu_i8, 'int8', torch.int8, 8, 8, 4), +] + + +@pytest.mark.parametrize('testfunc, sigtype, dtype, XB, YB, ZB', testlist) +def test_npu(testfunc, sigtype, dtype, XB, YB, ZB): + + x = torch.full((XB, YB, ZB), 100, dtype=dtype).npu() + + print(f"shape = {x.shape}") + print(x.dtype) + print(x[0, 0:16, 0]) + + output = torch.randint(1, (XB, YB, ZB), dtype=dtype).npu() + + print(f"output.dtype={output.dtype}") + + testfunc[1, 1, 1](output, XB, YB, ZB, debug=True) + print(output[0, 0:16, 0]) + + test_common.validate_cmp(sigtype, output, x) diff --git a/third_party/ascend/unittest/pytest_ut/test_fusedattention.py b/third_party/ascend/unittest/pytest_ut/test_fusedattention.py new file mode 100644 index 0000000000..f884dd3dbe --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_fusedattention.py @@ -0,0 +1,354 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Fused Attention +=============== + +This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao (https://tridao.me/publications/flash2/flash2.pdf) + +Credits: OpenAI kernel team + +Extra Credits: + +* Original flash attention paper (https://arxiv.org/abs/2205.14135) +* Rabe and Staats (https://arxiv.org/pdf/2112.05682v2.pdf) + +""" + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl +import triton.language.extra.cann.extension as extension + +DEVICE = "npu" + + +@triton.jit +def _attn_fwd_inner(acc_ptr, l_i, m_i, q, # Accumulator, local l, local m, query vector + K_block_ptr, V_block_ptr, # Key and value block pointers for current stage + start_m, qk_scale, # Starting position of current query block, qk scale factor + BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr, BLOCK_N: tl.constexpr, # Block size constants + STAGE: tl.constexpr, offs_m: tl.constexpr, + offs_n: tl.constexpr, # Current stage flag, m and n offset indices + N_CTX: tl.constexpr, + fp8_v: tl.constexpr): # Total context length, whether to enable FP8 for value precision + # Set the processing range [lo, hi) for the current stage (in column block units) + # causal = true + # stage = 1 + # Causal attention, as the name implies, restricts the flow of information during computation, + # only allowing the model to see the current and previous positions. + # In other words, the output at the current position can only depend on the input at or before this position, + # and cannot access information from future positions. + # Causal attention ensures sequential order and prevents "leakage of future information." + # But the following logic will also be triggered + if STAGE == 1: + # Stage 1: process all tokens before the query block + tl.static_assert(BLOCK_M >= BLOCK_N) + lo, hi = 0, start_m * BLOCK_M + elif STAGE == 2: + # Stage 2: process the current query block + tl.static_assert(BLOCK_M >= BLOCK_N) + lo, hi = start_m * BLOCK_M, (start_m + 1) * BLOCK_M + lo = tl.multiple_of(lo, BLOCK_M) # Align starting position + # causal = False (no need for masking) + else: + lo, hi = 0, N_CTX # Process the entire context + + # Adjust K and V block pointers to the starting position `lo` + K_block_ptr = tl.advance(K_block_ptr, (lo, 0)) # K is [HEAD_DIM, N_CTX], shift along the second dim by lo + V_block_ptr = tl.advance(V_block_ptr, (lo, 0)) # V is [N_CTX, HEAD_DIM], shift along the first dim by lo + + # Index mapping for the accumulator , used for slicing when HEAD_DIM >= 256 + row = tl.arange(0, BLOCK_M)[:, None] + col_head_dim = tl.arange(0, HEAD_DIM)[None, :] + block2d_acc = row * HEAD_DIM + col_head_dim + + # Iterate over all k, v blocks in the current stage and accumulate the output + for start_n in range(lo, hi, BLOCK_N): # Process BLOCK_N columns at a time + start_n = tl.multiple_of(start_n, BLOCK_N) # Align column start position + # -- Compute qk ---- + k = tl.load(K_block_ptr) + # Modify K + trans_k = tl.trans(k) + qk = tl.dot(q, trans_k) + # Apply causal mask for STAGE 2 + if STAGE == 2: + mask = offs_m[:, None] >= (start_n + offs_n[None, :]) # Construct upper triangular mask + qk = qk * qk_scale + tl.where(mask, 0, -1.0e6) # Set invalid positions to -∞ + m_ij = tl.maximum(m_i, tl.max(qk, 1)) # Update m_ij = max(m_i, max(qk)) + qk -= m_ij[:, None] # Subtract max for softmax stability + else: + qk = qk * qk_scale + m_ij = tl.maximum(m_i, tl.max(qk, 1)) # Scaled max + qk = qk - m_ij[:, None] # Stabilize + + # Softmax weights p = exp(qk) + p = tl.math.exp(qk) + + # Convert softmax weight type depending on FP8 usage + if fp8_v: + p_cast = p.to(tl.float8e5) # Convert to FP8 format (save memory) + else: + p_cast = p.to(k.dtype) + + v = tl.load(V_block_ptr) # Load corresponding V block + pv = tl.dot(p_cast, v) + l_ij = tl.sum(p, 1) # Softmax denominator (sum of each row) + # -- Update m_i and l_i + alpha = tl.math.exp(m_i - m_ij) # Update factor: exp difference between old and new max + l_i = l_i * alpha + l_ij # Update softmax denominator + # -- Update output accumulator -- + if HEAD_DIM < 256: + acc_ptr = acc_ptr * alpha[:, None] + acc_ptr = tl.dot(p_cast, v, acc_ptr) + else: + # 1. Load current slice of accumulator + acc = tl.load(acc_ptr + block2d_acc) + # 2. Update in slices (split by 1/4 of BLOCK_M to avoid ub overflow) + for i in range(4): + # Calculate start/end rows for current slice + offset = i * (BLOCK_M // 4) + # Extract slice data + acc_i = extension.extract_slice(acc, (offset, 0), (BLOCK_M // 4, HEAD_DIM), (1, 1)) + alpha_i = extension.extract_slice(alpha, [offset], [BLOCK_M // 4], [1]) + pv_i = extension.extract_slice(pv, (offset, 0), (BLOCK_M // 4, HEAD_DIM), (1, 1)) + # Incrementally update slice: acc = acc * alpha + pv + acc_i = acc_i * alpha_i[:, None] + pv_i + # Write updated slice back to accumulator + acc = extension.insert_slice(acc, acc_i, (offset, 0), (BLOCK_M // 4, HEAD_DIM), (1, 1)) + # 3. updated accumulator + tl.store(acc_ptr + block2d_acc, acc) + + m_i = m_ij # Update current block max + # Advance V and K block pointers to next BLOCK_N range + V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0)) + K_block_ptr = tl.advance(K_block_ptr, (BLOCK_N, 0)) + # Return accumulated output acc_ptr, softmax denominator l_i, and max value m_i + return acc_ptr, l_i, m_i + + +@triton.jit +def _attn_fwd(Q, K, V, M, Out, acc, sm_scale, stride_qz: tl.constexpr, stride_qh: tl.constexpr, stride_qm: tl.constexpr, + stride_qk: tl.constexpr, stride_kz: tl.constexpr, stride_kh: tl.constexpr, stride_kn: tl.constexpr, + stride_kk: tl.constexpr, stride_vz: tl.constexpr, stride_vh: tl.constexpr, stride_vn: tl.constexpr, + stride_vk: tl.constexpr, stride_oz: tl.constexpr, stride_oh: tl.constexpr, stride_om: tl.constexpr, + stride_on: tl.constexpr, Z: tl.constexpr, H: tl.constexpr, N_CTX: tl.constexpr, HEAD_DIM: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, STAGE: tl.constexpr): + # Total number of blocks in sequence dimension (M) + NUM_BLOCKS_M = N_CTX // BLOCK_M + # Total tasks = number of sequence blocks × batch size (Z) × number of attention heads (H) + NUM_BLOCKS = NUM_BLOCKS_M * Z * H + + # Current M-dimension block index + pid = tl.program_id(0) + + for block_idx in range(pid, NUM_BLOCKS, 20): + task_hz_idx = block_idx // NUM_BLOCKS_M + task_m_idx = block_idx % NUM_BLOCKS_M + off_z = task_hz_idx // H + off_h = task_hz_idx % H + qvk_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh + # Create block pointers for Q, K, V, Output + Q_block_ptr = tl.make_block_ptr( + base=Q + qvk_offset, + shape=(N_CTX, HEAD_DIM), + strides=(stride_qm, stride_qk), + offsets=(task_m_idx * BLOCK_M, 0), + block_shape=(BLOCK_M, HEAD_DIM), + order=(1, 0), + ) + V_block_ptr = tl.make_block_ptr( + base=V + qvk_offset, + shape=(N_CTX, HEAD_DIM), + strides=(stride_vn, stride_vk), + offsets=(0, 0), + block_shape=(BLOCK_N, HEAD_DIM), + order=(1, 0), + ) + K_block_ptr = tl.make_block_ptr( + base=K + qvk_offset, + shape=(N_CTX, HEAD_DIM), + strides=(stride_kn, stride_kk), + offsets=(0, 0), + block_shape=(BLOCK_N, HEAD_DIM), + order=(1, 0), + ) + O_block_ptr = tl.make_block_ptr( + base=Out + qvk_offset, + shape=(N_CTX, HEAD_DIM), + strides=(stride_om, stride_on), + offsets=(task_m_idx * BLOCK_M, 0), + block_shape=(BLOCK_M, HEAD_DIM), + order=(1, 0), + ) + # Initialize offsets + offs_m = task_m_idx * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0 + + # Initialize accumulator + if HEAD_DIM < 256: + acc_ptr = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) + else: + acc_offset = (off_z.to(tl.int64) * stride_qz // stride_qm * HEAD_DIM + + off_h.to(tl.int64) * stride_qh // stride_qm * HEAD_DIM + task_m_idx * BLOCK_M * HEAD_DIM) + acc_ptr = acc + acc_offset + + # load q: it will stay in SRAM throughout + q = tl.load(Q_block_ptr) + + # stage 1: off-band + # For causal = True, STAGE = 3 and _attn_fwd_inner gets 1 as its STAGE + # For causal = False, STAGE = 1, and _attn_fwd_inner gets 3 as its STAGE + if STAGE & 1: + acc_ptr, l_i, m_i = _attn_fwd_inner(acc_ptr, l_i, m_i, q, K_block_ptr, V_block_ptr, # + task_m_idx, sm_scale, # + BLOCK_M, HEAD_DIM, BLOCK_N, # + 4 - STAGE, offs_m, offs_n, N_CTX, V.dtype.element_ty == tl.float8e5 # + ) + # stage 2: on-band + if STAGE & 2: + # barrier makes it easier for compielr to schedule the + # two loops independently + acc_ptr, l_i, m_i = _attn_fwd_inner(acc_ptr, l_i, m_i, q, K_block_ptr, V_block_ptr, # + task_m_idx, sm_scale, # + BLOCK_M, HEAD_DIM, BLOCK_N, # + 2, offs_m, offs_n, N_CTX, V.dtype.element_ty == tl.float8e5 # + ) + + m_i += tl.math.log(l_i) + if HEAD_DIM < 256: + accumulator = acc_ptr / l_i[:, None] + else: + row = tl.arange(0, BLOCK_M)[:, None] + col_head_dim = tl.arange(0, HEAD_DIM)[None, :] + block2d_acc = row * HEAD_DIM + col_head_dim + accumulator = tl.load(acc_ptr + block2d_acc) + accumulator = accumulator / l_i[:, None] + + m_ptrs = M + task_hz_idx * N_CTX + offs_m + + tl.store(m_ptrs, m_i) + tl.store(O_block_ptr, accumulator.to(Out.type.element_ty)) + + +class _attention(torch.autograd.Function): + + @staticmethod + def forward(ctx, q, k, v, causal, sm_scale, BM, BN): + """ + Forward computation interface: + Args: + ctx: Context object + q: Query tensor (Q), shape [Z, H, N_CTX, HEAD_DIM] + k: Key tensor (K), shape [Z, H, N_CTX, HEAD_DIM] + v: Value tensor (V), shape [Z, H, N_CTX, HEAD_DIM] + causal: Whether to enable causal attention + sm_scale: Scaling factor for QK product + BM: Q block size (BLOCK_M) + BN: K/V block size (BLOCK_N) + Returns: + o: Attention output tensor, shape [Z, H, N_CTX, HEAD_DIM] + """ + # shape constraints + HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1] + # when v is in float8_e5m2 it is transposed. + HEAD_DIM_V = v.shape[-1] + assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V + assert HEAD_DIM_K in {16, 32, 64, 128, 256} + + o = torch.empty_like(q) + stage = 3 if causal else 1 + extra_kern_args = {} + + # Number of NPU cores (adjust based on hardware) + num_cores = 20 + acc = torch.zeros((q.shape[0], q.shape[1], q.shape[2], HEAD_DIM_K), dtype=torch.float32, device=q.device) + M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32) + + _attn_fwd[(num_cores, )](q, k, v, M, o, acc, sm_scale, q.stride(0), q.stride(1), q.stride(2), q.stride(3), + k.stride(0), k.stride(1), k.stride(2), k.stride(3), v.stride(0), v.stride(1), + v.stride(2), v.stride(3), o.stride(0), o.stride(1), o.stride(2), o.stride(3), + q.shape[0], q.shape[1], N_CTX=q.shape[2], HEAD_DIM=HEAD_DIM_K, BLOCK_M=BM, BLOCK_N=BN, + STAGE=stage, **extra_kern_args) + + ctx.save_for_backward(q, k, v, o, M) + ctx.sm_scale = sm_scale + ctx.HEAD_DIM = HEAD_DIM_K + ctx.causal = causal + return o + + +attention = _attention.apply + + +@pytest.mark.parametrize("Z, H, N_CTX, HEAD_DIM, causal, dtype, BM, BN", [ + (1, 1, 128, 128, False, torch.float16, 32, 128), + (1, 1, 128, 128, False, torch.bfloat16, 64, 128), + (1, 2, 256, 256, False, torch.bfloat16, 32, 256), + (2, 2, 128, 256, False, torch.float16, 64, 128), + (4, 32, 64, 64, False, torch.float16, 32, 64), + (4, 32, 1024, 64, False, torch.bfloat16, 64, 128), + (4, 32, 4096, 64, False, torch.float16, 128, 128), +]) +def test_op(Z, H, N_CTX, HEAD_DIM, causal, dtype, BM, BN): + # Filter out non-integer cases; N_CTX must be divisible by BM and BN, and HEAD_DIM must be divisible by 16. + if N_CTX % BM != 0 or N_CTX % BN != 0 or HEAD_DIM % 16 != 0: + pytest.skip("Skipping non-divisible case") + + torch.manual_seed(20) + q = (torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device=DEVICE).normal_(mean=0.0, std=0.5).requires_grad_()) + k = (torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device=DEVICE).normal_(mean=0.0, std=0.5).requires_grad_()) + v = (torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device=DEVICE).normal_(mean=0.0, std=0.5).requires_grad_()) + + sm_scale = 0.5 + + tri_out = attention(q, k, v, causal, sm_scale, BM, BN) + ref_out = torch_npu.npu_fusion_attention( + q, + k, + v, + H, + padding_mask=None, + atten_mask=None, + scale=sm_scale, + keep_prob=1.0, + input_layout="BNSD", + pre_tockens=65535, + next_tockens=65535, + sparse_mode=0, + )[0] + + torch.testing.assert_close(ref_out, tri_out, atol=1e-2, rtol=1e-2, equal_nan=True) + print(f"[PASSED] Attention shape:({Z}, {H}, {N_CTX}, {HEAD_DIM}), BM: {BM}, BN: {BN}, dtype: {dtype}") + + +if __name__ == "__main__": + test_op(1, 1, 128, 128, causal=False, dtype=torch.float16, BM=32, BN=128) + test_op(1, 1, 128, 128, causal=False, dtype=torch.bfloat16, BM=64, BN=128) + test_op(1, 2, 256, 256, causal=False, dtype=torch.bfloat16, BM=32, BN=256) + test_op(2, 2, 128, 256, causal=False, dtype=torch.float16, BM=64, BN=128) + test_op(4, 32, 64, 64, causal=False, dtype=torch.float16, BM=32, BN=64) + test_op(4, 32, 1024, 64, causal=False, dtype=torch.bfloat16, BM=64, BN=128) + test_op(4, 32, 4096, 64, causal=False, dtype=torch.float16, BM=128, BN=128) diff --git a/third_party/ascend/unittest/pytest_ut/test_gamma.py b/third_party/ascend/unittest/pytest_ut/test_gamma.py new file mode 100644 index 0000000000..b3cdf3b9af --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_gamma.py @@ -0,0 +1,86 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common +import triton.language.extra.cann.libdevice as libdevice +import numpy as np +from scipy.special import gamma + + +@triton.jit +def triton_gamma(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + xindex = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = xindex < xnumel + x0 = tl.load(in_ptr0 + xindex, xmask) + y = libdevice.gamma(x0) + tl.store(out_ptr0 + xindex, y, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_gamma_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x = torch.abs(test_common.generate_tensor(shape, dtype)) + x_np = x.cpu().numpy() + x = x.npu() + y_ref = torch.from_numpy(gamma(x_np)).to(x.device).to(x.dtype).npu() + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_gamma[ncore, 1, 1](x, y_cal, x.numel(), xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_all_blocks_parallel(param_list, monkeypatch): + monkeypatch.setenv("TRITON_ALL_BLOCKS_PARALLEL", "1") + dtype, shape, ncore, xblock, xblock_sub = param_list + x = torch.abs(test_common.generate_tensor(shape, dtype)) + x_np = x.cpu().numpy() + x = x.npu() + y_ref = torch.from_numpy(gamma(x_np)).to(x.device).to(x.dtype).npu() + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_gamma[ncore, 1, 1](x, y_cal, x.numel(), xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) + monkeypatch.delenv("TRITON_ALL_BLOCKS_PARALLEL") + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 2048, 8), 2, 32768, 512], +]) +def test_auto_blockify(param_list, monkeypatch): + monkeypatch.setenv("TRITON_ALL_BLOCKS_PARALLEL", "1") + dtype, shape, ncore, xblock, xblock_sub = param_list + x = torch.abs(test_common.generate_tensor(shape, dtype)) + x_np = x.cpu().numpy() + x = x.npu() + y_ref = torch.from_numpy(gamma(x_np)).to(x.device).to(x.dtype).npu() + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_gamma[ncore, 1, 1](x, y_cal, x.numel(), xblock, xblock_sub, auto_blockify_size=ncore) + test_common.validate_cmp(dtype, y_cal, y_ref) + monkeypatch.delenv("TRITON_ALL_BLOCKS_PARALLEL") diff --git a/third_party/ascend/unittest/pytest_ut/test_gather.py b/third_party/ascend/unittest/pytest_ut/test_gather.py new file mode 100644 index 0000000000..08946f8944 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_gather.py @@ -0,0 +1,73 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import triton.language.extra.cann.extension as extension +import numpy as np +import test_common +import pytest + + +@pytest.mark.skip(reason="waiting for the compiler to support.") +@pytest.mark.parametrize("src_shape, indices_shape, axis", [ + ([4, 4], [8, 4], 0), + ([4, 64], [4, 32], 1), + ([128, 64], [128, 128], 1), +]) +def test_gather(src_shape, indices_shape, axis): + + @triton.jit + def gather_kernel(src_ptr, idx_ptr, out_ptr, axis: tl.constexpr, src_dim0: tl.constexpr, src_dim1: tl.constexpr, + src_stride0: tl.constexpr, src_stride1: tl.constexpr, idx_dim0: tl.constexpr, + idx_dim1: tl.constexpr, idx_stride0: tl.constexpr, idx_stride1: tl.constexpr, + out_dim0: tl.constexpr, out_dim1: tl.constexpr, out_stride0: tl.constexpr, + out_stride1: tl.constexpr): + src_offs = (tl.arange(0, src_dim0)[:, None] * src_stride0 + tl.arange(0, src_dim1)[None, :] * src_stride1) + src = tl.load(src_ptr + src_offs) + + idx_offs = (tl.arange(0, idx_dim0)[:, None] * idx_stride0 + tl.arange(0, idx_dim1)[None, :] * idx_stride1) + idx = tl.load(idx_ptr + idx_offs) + + out = extension.gather(src, idx, axis) + + out_offs = (tl.arange(0, out_dim0)[:, None] * out_stride0 + tl.arange(0, out_dim1)[None, :] * out_stride1) + tl.store(out_ptr + out_offs, out) + + def triton_gather(src: torch.Tensor, axis: int, indices: torch.Tensor): + output = torch.empty(indices.shape, dtype=src.dtype, device=src.device) + gather_kernel[(1, )](src, indices, output, axis, src.shape[0], src.shape[1], + src.stride(0), src.stride(1), indices.shape[0], indices.shape[1], indices.stride(0), + indices.stride(1), output.shape[0], output.shape[1], output.stride(0), output.stride(1)) + return output + + DEV = "npu" + src = torch.randn(src_shape, device=DEV) + indices = torch.randint(0, src.shape[axis], indices_shape, device=DEV) + ref = torch.gather(src, axis, indices) + result = triton_gather(src, axis, indices) + torch.testing.assert_close(result, ref, rtol=0, atol=0) + + +if __name__ == "__main__": + test_gather([4, 64], [4, 32], 1) + print("success: test_gather") diff --git a/third_party/ascend/unittest/pytest_ut/test_gather_simd.py b/third_party/ascend/unittest/pytest_ut/test_gather_simd.py new file mode 100644 index 0000000000..6225e4ef2c --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_gather_simd.py @@ -0,0 +1,46 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Unit test for gather_2d_simd kernel. +""" + +import torch +import torch_npu +import triton +import triton.language as tl +import pytest + +from triton.language.extra.kernels import gather_2d_simd + + +@pytest.mark.parametrize("M,N,K", [ + (32, 128, 64), +]) +def test_gather_2d_simd(M, N, K): + """Test gather_2d_simd with various tensor sizes.""" + src = torch.randn(M, N, dtype=torch.float32, device='npu') + indices = torch.randint(0, N, (M, K), dtype=torch.int32, device='npu') + output = torch.empty((M, K), dtype=src.dtype, device='npu') + + grid = (triton.cdiv(M, 32), ) + gather_2d_simd[grid](src, indices, output, M, N, K, XBLOCK=32, XBLOCK_SUB=4) + + ref = torch.gather(src, 1, indices.long()) + assert torch.allclose(output, ref, rtol=1e-5, atol=1e-5) diff --git a/third_party/ascend/unittest/pytest_ut/test_ge.py b/third_party/ascend/unittest/pytest_ut/test_ge.py new file mode 100644 index 0000000000..e86e6dc632 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_ge.py @@ -0,0 +1,79 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + + +def standard_binary(x0, y0): + res = x0 >= y0 + return res + + +@triton.jit +def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + y = tl.load(in_ptr1 + idx_block, mask=idx_block < N) + ret = x >= y + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + # (torch.bfloat16, 'bfloat16'), + (torch.int8, 'int8'), + (torch.int16, 'int16'), + (torch.int32, 'int32'), + (torch.int64, 'int64'), +] + +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('N,NUMEL', shapes) +def test_elementwsie_common(dtype, sigtype, N, NUMEL): + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == "int64": + N = map_for_64_t[N] if N in map_for_64_t else N + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + y0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + ans = standard_binary(x0, y0) + out = torch.zeros((N, ), dtype=torch.bool).npu() + triton_elementwise_binary[1, 1, 1](x0, y0, out, N, NUMEL) + test_common.validate_cmp(sigtype, out, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_ge_2.py b/third_party/ascend/unittest/pytest_ut/test_ge_2.py new file mode 100644 index 0000000000..c5e9b6ebc2 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_ge_2.py @@ -0,0 +1,63 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common + + +def torch_ge(x0, x1, dtype): + res = torch.where(torch.ge(x0, x1), torch.ones_like(x0), torch.zeros_like(x0)).to(eval('torch.' + dtype)) + return res + + +@triton.jit +def triton_ge(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.load(in_ptr1 + (x0), None) + tmp2 = tmp0 >= tmp1 + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], +]) +def test_ge(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + torch_res = torch_ge(x0, x1, dtype) + # triton结果 + triton_res = torch.empty_like(x0) + triton_ge[ncore, 1, 1](x0, x1, triton_res, xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_gelu.py b/third_party/ascend/unittest/pytest_ut/test_gelu.py new file mode 100644 index 0000000000..ff6cfe2286 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_gelu.py @@ -0,0 +1,99 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + + +def standard_unary(x0, dtype): + res = x0 * 0.5 * (1.0 + torch.erf(x0 / torch.sqrt(torch.tensor(2.0)))) + return res + + +def standard_binary(x0, y0, dtype): + res = x0 + y0 + return res + + +@triton.jit +def triton_elementwise_unary(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + ret = x * 0.5 * (1.0 + tl.erf(x / tl.sqrt(2.0))) + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +@triton.jit +def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + y = tl.load(in_ptr1 + idx_block, mask=idx_block < N) + ret = x + y + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +types = [ + (torch.float32, 'float32'), + # (torch.float16, 'float16'), + # (torch.bfloat16, 'bfloat16'), + # (torch.int8, 'int8'), + # (torch.int16, 'int16'), + # (torch.int32, 'int32'), + # (torch.int64, 'int64'), +] + +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('N,NUMEL', shapes) +def test_elementwsie_common(dtype, sigtype, N, NUMEL): + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == "int64": + N = map_for_64_t[N] if N in map_for_64_t else N + + print(f"elementwise : ({N},) {dtype} {sigtype}") + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standard_unary(x0, dtype) + x0 = x0.npu() + print(ans) + + out = torch.zeros((N, ), dtype=dtype).npu() + triton_elementwise_unary[1, 1, 1](x0, out, N=N, NUMEL=NUMEL, debug=True) + print(out) + + test_common.validate_cmp(sigtype, out, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_gt.py b/third_party/ascend/unittest/pytest_ut/test_gt.py new file mode 100644 index 0000000000..e5d81e0584 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_gt.py @@ -0,0 +1,61 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common + + +def torch_gt(x0, x1, dtype): + res = torch.where(torch.gt(x0, x1), torch.ones_like(x0), torch.zeros_like(x0)).to(eval('torch.' + dtype)) + return res + + +@triton.jit +def triton_gt(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.load(in_ptr1 + (x0), None) + tmp2 = tmp0 > tmp1 + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_gt(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + torch_res = torch_gt(x0, x1, dtype) + # triton结果 + triton_res = torch.empty_like(x0) + triton_gt[ncore, 1, 1](x0, x1, triton_res, xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_hd_permute.py b/third_party/ascend/unittest/pytest_ut/test_hd_permute.py new file mode 100644 index 0000000000..c58e985b5d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_hd_permute.py @@ -0,0 +1,53 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu + +X_SIZE = tl.constexpr(4) +Y_SIZE = tl.constexpr(64) +Z_SIZE = tl.constexpr(32) +NUMEL = X_SIZE * Y_SIZE * Z_SIZE + + +def torch_permute(x): + return x.reshape((X_SIZE, Y_SIZE, Z_SIZE)).permute(1, 0, 2).reshape((X_SIZE * Y_SIZE * Z_SIZE)) + + +@triton.jit +def triton_permute(output_ptr, input_ptr): + x_index = tl.arange(0, X_SIZE * Y_SIZE * Z_SIZE) + input_local = tl.load(input_ptr + x_index) + output_local = input_local.reshape((X_SIZE, Y_SIZE, Z_SIZE)).permute(1, 0, 2).reshape((X_SIZE * Y_SIZE * Z_SIZE)) + tl.store(output_ptr + x_index, output_local) + + +def test_hd_permute(): + # 生成数据 + x = torch.randn(NUMEL).npu() + # torch结果 + torch_res = torch_permute(x) + # triton结果 + triton_res = torch.randn(torch_res.shape, dtype=torch_res.dtype).npu() + triton_permute[1, 1, 1](triton_res, x) + # 比较结果 + torch.testing.assert_close(triton_res, torch_res, rtol=1e-3, atol=1e-3) diff --git a/third_party/ascend/unittest/pytest_ut/test_histogram.py b/third_party/ascend/unittest/pytest_ut/test_histogram.py new file mode 100644 index 0000000000..4a54934e81 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_histogram.py @@ -0,0 +1,69 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import math +import test_common + + +@triton.jit +def histogram_kernel(x_ptr, z_ptr, M: tl.constexpr, N: tl.constexpr): + offset1 = tl.arange(0, M) + offset2 = tl.arange(0, N) + x = tl.load(x_ptr + offset1) + z = tl.histogram(x, N) + tl.store(z_ptr + offset2, z) + + +@pytest.mark.parametrize("M", [2048]) +@pytest.mark.parametrize("N", [2]) +@pytest.mark.parametrize("ncore", [2]) +@pytest.mark.parametrize("dtype", ["int32", "int64"]) +def test_histogram(M, N, ncore, dtype): + torch.manual_seed(17) + x = torch.randint(low=0, high=N, size=(M, ), dtype=eval(f'torch.{dtype}')).npu() + # torch结果 + y_cal = torch.histc(x.float(), bins=N, min=0, max=N - 1) + # triton结果 + y_ref = torch.empty(N, dtype=eval(f'torch.{dtype}'), device="npu") + histogram_kernel[(ncore, )](x, y_ref, M=M, N=N) + print(y_cal) + print(y_ref) + test_common.validate_cmp(dtype, y_cal, y_ref) + + +@pytest.mark.parametrize("M", [2048]) +@pytest.mark.parametrize("N", [2]) +@pytest.mark.parametrize("ncore", [2]) +@pytest.mark.parametrize("dtype", ["uint32", "uint64"]) +def test_histogram_uint(M, N, ncore, dtype): + torch.manual_seed(17) + x_cpu = torch.randint(low=0, high=N, size=(M, ), dtype=eval(f'torch.{dtype}'), device="cpu") + x = x_cpu.to("npu") + # torch结果 + y_cal = torch.histc(x.float(), bins=N, min=0, max=N - 1) + y_cal = y_cal.to(eval(f'torch.{dtype}')) + # triton结果 + y_ref = torch.empty(N, dtype=eval(f'torch.{dtype}'), device="npu") + histogram_kernel[(ncore, )](x, y_ref, M=M, N=N) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_hoistbroadcast.py b/third_party/ascend/unittest/pytest_ut/test_hoistbroadcast.py new file mode 100644 index 0000000000..79ea5b3ca7 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_hoistbroadcast.py @@ -0,0 +1,220 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu +import numpy as np + + +@triton.jit +def copy( + data_ptrs, + tgt_loc_ptr, + src_loc_ptr, + num_locs, + stride: tl.constexpr, + num_locs_upper: tl.constexpr, +): + bid = tl.program_id(0) + data_ptr = tl.load(data_ptrs + bid) + data_ptr = tl.cast(data_ptr, tl.pointer_type(tl.float16)) + num_locs_offset = tl.arange(0, num_locs_upper) + tgt_locs = tl.load(tgt_loc_ptr + num_locs_offset) + src_locs = tl.load(src_loc_ptr + num_locs_offset) + copy_offset = tl.arange(0, stride) + value = tl.load(data_ptr + src_locs[:, None] * stride + copy_offset[None, :]) + value += 1 + tl.store(data_ptr + tgt_locs[:, None] * stride + copy_offset[None, :], value) + + +@triton.jit +def copy_all_layer_kv_cache( + data_ptrs, + strides, + tgt_loc_ptr, + src_loc_ptr, + num_locs, + num_locs_upper: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 128 + + bid = tl.program_id(0) + stride = tl.load(strides + bid) + + data_ptr = tl.load(data_ptrs + bid) + data_ptr = tl.cast(data_ptr, tl.pointer_type(tl.uint8)) + + num_locs_offset = tl.arange(0, num_locs_upper) + tgt_locs = tl.load(tgt_loc_ptr + num_locs_offset) + src_locs = tl.load(src_loc_ptr + num_locs_offset) + + num_loop = tl.cdiv(stride, BLOCK_SIZE) + for i in range(num_loop): + copy_offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + mask = (num_locs_offset < num_locs)[:, None] & (copy_offset < stride)[None, :] + value = tl.load(data_ptr + src_locs[:, None] * stride + copy_offset[None, :], mask=mask) + value *= 1 + tl.store( + data_ptr + tgt_locs[:, None] * stride + copy_offset[None, :], + value, + mask=mask, + ) + + +@triton.jit +def copy_all_layer_kv_cache2( + data_ptrs, + strides, + tgt_loc_ptr, + src_loc_ptr, + num_locs, + num_locs_upper: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 128 + + bid = tl.program_id(0) + stride = tl.load(strides + bid) + + data_ptr = tl.load(data_ptrs + bid) + data_ptr = tl.cast(data_ptr, tl.pointer_type(tl.uint8)) + + num_locs_offset = tl.arange(0, num_locs_upper) + tgt_locs = tl.load(tgt_loc_ptr + num_locs_offset) + src_locs = tl.load(src_loc_ptr + num_locs_offset) + + num_loop = tl.cdiv(stride, BLOCK_SIZE) + for i in range(num_loop): + copy_offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + mask = (num_locs_offset < num_locs)[:, None] & (copy_offset < stride)[None, :] + value = tl.load(data_ptr + (src_locs[:, None] * stride + copy_offset[None, :]), mask=mask) + value *= 1 + tl.store( + data_ptr + (tgt_locs[:, None] * stride + copy_offset[None, :]), + value, + mask=mask, + ) + + +@pytest.mark.parametrize('param_list', [ + ['float16', (10, 10), 'npu'], +]) +def test_copy(param_list): + dtype, shape, device = param_list + data = torch.zeros(shape, dtype=eval('torch.' + dtype), device=device) + data_ref = torch.zeros(shape, dtype=eval('torch.' + dtype), device=device) + + src_loc = torch.tensor([0], dtype=torch.int32, device=device) + tgt_loc = torch.tensor([0], dtype=torch.int32, device=device) + + data_ptr = torch.tensor([data.data_ptr()], dtype=torch.uint64, device=device) + stride = shape[1] + copy[(1, )](data_ptr, tgt_loc, src_loc, 1, stride, 1) + data_ref[0, :] += 1 + test_common.validate_cmp(dtype, data, data_ref) + + +@pytest.mark.parametrize('param_list', [ + ['float16', 3, 20, 16, 4, 16, 'npu'], +]) +def test_hoistbroadcast_compare(param_list): + dtype, layer_num, page_num, page_size, head_num, head_dim, device = param_list + kv_buffer = torch.randn((2, layer_num, page_num, page_size, head_num, head_dim), dtype=eval('torch.' + dtype), + device=device) + kv_buffer_ref = kv_buffer.clone() + k_buffer = kv_buffer[0] + v_buffer = kv_buffer[1] + k_buffer_ref = kv_buffer_ref[0] + v_buffer_ref = kv_buffer_ref[1] + + data_ptrs = torch.tensor([x.data_ptr() for x in [k_buffer]] + [x.data_ptr() for x in [v_buffer]], + dtype=torch.uint64, device=device) + data_ptrs_ref = torch.tensor([x.data_ptr() for x in [k_buffer_ref]] + [x.data_ptr() for x in [v_buffer_ref]], + dtype=torch.uint64, device=device) + + data_strides = torch.cat([ + torch.tensor([np.prod(x.shape[1:]) * x.dtype.itemsize for x in k_buffer], device=device), + torch.tensor([np.prod(x.shape[1:]) * x.dtype.itemsize for x in v_buffer], device=device) + ], dim=0) + data_strides_ref = data_strides.clone() + + src_loc = torch.tensor([0], dtype=torch.int32, device=device) + tgt_loc = torch.tensor([0], dtype=torch.int32, device=device) + + copy_all_layer_kv_cache[(len(data_ptrs), )](data_ptrs, data_strides, tgt_loc, src_loc, len(tgt_loc), 1) + copy_all_layer_kv_cache2[(len(data_ptrs_ref), )](data_ptrs_ref, data_strides_ref, tgt_loc, src_loc, len(tgt_loc), 1) + test_common.validate_cmp(dtype, kv_buffer, kv_buffer_ref) + + +def torch_pointwise(x0): + if x0.dtype != torch.uint32: + return torch.abs(x0) + else: + return torch.abs(x0.to(torch.float32)) + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, z_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr, + YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xoffs = tl.program_id(0) * XB + yoffs = tl.program_id(1) * YB + zoffs = tl.program_id(2) * ZB + + xidx = tl.arange(0, XB) + xoffs + yidx = tl.arange(0, YB) + yoffs + zidx = tl.arange(0, ZB) + zoffs + + X = tl.load(x_ptr + xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :]) + ret = tl.abs(X) + tl.store(output_ptr + xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :], + ret) + + +@pytest.mark.parametrize('shape', [(8, 16, 16)]) +@pytest.mark.parametrize('dtype', ['float32']) +def test_case2(dtype, shape): + # 生成数据 + x = test_common.generate_tensor(shape, dtype).npu() + y = test_common.generate_tensor(shape, dtype).npu() + z = test_common.generate_tensor(shape, dtype).npu() + new_shape = shape + + output = torch.randint(1, new_shape, dtype=eval('torch.' + dtype)).npu() + output1 = output + + ans = torch_pointwise(x) + + XB = shape[0] + xnumel = shape[0] + YB = shape[1] + ynumel = shape[1] + ZB = shape[2] + znumel = shape[2] + + grid = (1, 1, 1) + + fn_npu_[grid](output, x, y, z, XB, YB, ZB, xnumel, ynumel, znumel) + + test_common.validate_cmp(dtype, ans, output) diff --git a/third_party/ascend/unittest/pytest_ut/test_hypot.py b/third_party/ascend/unittest/pytest_ut/test_hypot.py new file mode 100644 index 0000000000..9bb1e6b31c --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_hypot.py @@ -0,0 +1,61 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + +import torch +import torch_npu + + +def torch_hypot(x0, x1): + return torch.hypot(x0, x1) + + +@triton.jit +def triton_hypot(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.load(in_ptr1 + (x0), None) + tmp2 = libdevice.hypot(tmp0, tmp1) + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['bfloat16', (2, 4096, 8), 2, 32768, 1024], +]) +def test_hypot(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_hypot(x0, x1) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_hypot[ncore, 1, 1](x0, x1, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_if_advance.py b/third_party/ascend/unittest/pytest_ut/test_if_advance.py new file mode 100644 index 0000000000..a51cc81374 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_if_advance.py @@ -0,0 +1,39 @@ +import torch +import torch_npu +import triton +import triton.language as tl +import triton.language.extra.cann.extension as al + + +@triton.jit +def triton_if_advance_kernel(in_ptr0, in_ptr1, out_ptr, xnumel, ynumel, k_loops, XBLOCK: tl.constexpr, + YBLOCK: tl.constexpr): + + K_block_ptr = tl.make_block_ptr(base=in_ptr0, shape=(xnumel, ynumel), strides=(ynumel, 1), offsets=(0, 0), + block_shape=(XBLOCK, YBLOCK), order=(1, 0)) + V_block_ptr = tl.make_block_ptr(base=in_ptr1, shape=(ynumel, xnumel), strides=(xnumel, 1), offsets=(0, 0), + block_shape=(YBLOCK, XBLOCK), order=(1, 0)) + O_block_ptr = tl.make_block_ptr(base=out_ptr, shape=(xnumel, xnumel), strides=(xnumel, 1), offsets=(0, 0), + block_shape=(XBLOCK, XBLOCK), order=(1, 0)) + res = tl.zeros([XBLOCK, XBLOCK], tl.float32) + for i in range(0, k_loops): + if i > 0: + K_block_ptr = tl.advance(K_block_ptr, (0, YBLOCK)) + V_block_ptr = tl.advance(V_block_ptr, (YBLOCK, 0)) + a = tl.load(K_block_ptr) + b = tl.load(V_block_ptr) + res = tl.dot(a, b, acc=res) + tl.store(O_block_ptr, res) + + +def test_if_advance(): + x = torch.randn((64, 256), dtype=torch.float32, device="npu") + y = torch.randn((256, 64), dtype=torch.float32, device="npu") + out_tri = torch.empty((64, 64), dtype=torch.float32, device="npu") + out_std = torch.empty((64, 64), dtype=torch.float32, device="npu") + torch.matmul(x, y, out=out_std) + triton_if_advance_kernel[1, 1, 1](x, y, out_tri, 64, 256, 4, 64, 64) + torch.testing.assert_close(out_std, out_tri, atol=1e-2, rtol=1e-2) + + +test_if_advance() diff --git a/third_party/ascend/unittest/pytest_ut/test_if_load.py b/third_party/ascend/unittest/pytest_ut/test_if_load.py new file mode 100644 index 0000000000..7932366cb4 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_if_load.py @@ -0,0 +1,82 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + +import torch +import torch_npu + + +@triton.jit +def triton_if_load(in_ptr0, out_ptr0, XBLOCK: tl.constexpr): + base1 = tl.arange(0, XBLOCK) + index = base1 + if tl.program_id(0) == 0: + base1 = base1 * 1 + else: + base1 = base1 * 2 + tmp0 = tl.load(in_ptr0 + base1, base1 < XBLOCK, other=0.0) + tl.store(out_ptr0 + index, tmp0, None) + + +@triton.jit +def triton_for_if_load(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + base1 = tl.arange(0, XBLOCK_SUB) + index = base1 + loops = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for i in range(loops): + base1 = base1 + i * XBLOCK_SUB + index = index + i * XBLOCK_SUB + if tl.program_id(0) != 0: + base1 = base1 + 1 + + tmp0 = tl.load(in_ptr0 + base1, base1 < XBLOCK, other=0.0) + tl.store(out_ptr0 + index, tmp0, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (32, ), 32], +]) +def test_if_load(param_list): + dtype, shape, xblock = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + y_ref = x0.clone() + + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_if_load[(1, )](x0, y_cal, xblock) + test_common.validate_cmp(dtype, y_cal, y_ref) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (32, ), 32, 16], +]) +def test_if_load(param_list): + dtype, shape, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + y_ref = x0.clone() + + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_for_if_load[(1, )](x0, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_if_tensor.py b/third_party/ascend/unittest/pytest_ut/test_if_tensor.py new file mode 100644 index 0000000000..7ceadb105e --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_if_tensor.py @@ -0,0 +1,134 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import torch +import triton +import triton.language as tl + +import pytest + + +@triton.jit +def if_tensor_kernel( + kv_start_idx, + output_ptr, +): + pid = tl.program_id(0) + if kv_start_idx: + value = tl.load(kv_start_idx + pid) + tl.store(output_ptr + pid, value) + + +def test_kernel(): + n = 8 + device = 'npu' + + kv_start_idx = torch.arange(n, dtype=torch.float32, device=device) + output1 = torch.zeros(n, dtype=torch.float32, device=device) + if_tensor_kernel[(n, )]( + kv_start_idx, + output1, + ) + + expected = torch.arange(n, dtype=torch.float32, device=device) + assert torch.allclose(output1, expected), f"Output {output1} != Expected {expected}" + print(f"RESULT: output1 = {output1}") + print("✅ Test passed!") + + +@triton.jit +def mul_if_block_kernel(value, value_stride0, value_stride1, output, output_stride0, output_stride1, lengths, bs, dim, + max_seq_len, DIM_SIZE: tl.constexpr, BLOCK_SIZE: tl.constexpr): + batch_idx = tl.program_id(0) + if batch_idx >= bs: + return + + need_reverse = tl.program_id(1) == 0 + + block_idx = tl.program_id(2) + block_start = block_idx * BLOCK_SIZE + + reverse_len = tl.load(lengths + batch_idx) + copy_len = max_seq_len - reverse_len + + value = value + batch_idx * value_stride0 + output = output + batch_idx * output_stride0 + + if need_reverse: + if block_start >= reverse_len: + return + value_block_ptr = tl.make_block_ptr(base=value, shape=(reverse_len, dim), strides=(value_stride0, 1), + offsets=(block_start, 0), block_shape=(BLOCK_SIZE, DIM_SIZE), order=(0, 1)) + + block_values = tl.load(value_block_ptr, boundary_check=(0, 1), padding_option="zero") + + output_block_ptr = tl.make_block_ptr(base=output + (reverse_len - 1) * output_stride1, shape=(reverse_len, dim), + strides=(output_stride0, 1), offsets=(block_start, 0), + block_shape=(BLOCK_SIZE, DIM_SIZE), order=(0, 1)) + tl.store(output_block_ptr, block_values, boundary_check=(0, 1)) + else: + if block_start >= copy_len: + return + value_block_ptr = tl.make_block_ptr(base=value + reverse_len * value_stride1, shape=(max_seq_len, dim), + strides=(value_stride0, 1), offsets=(block_start, 0), + block_shape=(BLOCK_SIZE, DIM_SIZE), order=(0, 1)) + + block_values = tl.load(value_block_ptr, boundary_check=(0, 1), padding_option="zero") + + output_block_ptr = tl.make_block_ptr(base=output + reverse_len * output_stride1, shape=(max_seq_len, dim), + strides=(output_stride0, 1), offsets=(block_start, 0), + block_shape=(BLOCK_SIZE, DIM_SIZE), order=(0, 1)) + tl.store(output_block_ptr, block_values, boundary_check=(0, 1)) + + +def ref_reverse(value: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor: + bs, max_seq_len, dim = value.shape + out = value.clone() + for b in range(bs): + r = int(lengths[b]) + if r > 0: + out[b, :r] = value[b, :r].flip(0) + return out + + +@pytest.mark.parametrize("bs,max_seq_len,dim,BLOCK_SIZE", [ + (2, 32, 16, 8), +]) +def test_reverse_sequence_kernel(bs, max_seq_len, dim, BLOCK_SIZE): + device = "npu" + value = torch.randn(bs, max_seq_len, dim, device=device, dtype=torch.float32) + lengths = torch.tensor([0, max_seq_len // 2, max_seq_len] + [max_seq_len // 3] * max(0, bs - 3), device=device, + dtype=torch.int32)[:bs] + output = torch.empty_like(value) + + value_stride0 = value.stride(0) + value_stride1 = value.stride(1) + output_stride0 = output.stride(0) + output_stride1 = output.stride(1) + + DIM_SIZE = dim + grid = (bs, 2, math.ceil(max_seq_len / BLOCK_SIZE)) + mul_if_block_kernel[grid](value, value_stride0, value_stride1, output, output_stride0, output_stride1, lengths, bs, + dim, max_seq_len, DIM_SIZE=DIM_SIZE, BLOCK_SIZE=BLOCK_SIZE) + + +if __name__ == "__main__": + test_kernel() diff --git a/third_party/ascend/unittest/pytest_ut/test_implicit_atomic.py b/third_party/ascend/unittest/pytest_ut/test_implicit_atomic.py new file mode 100644 index 0000000000..b7a539f392 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_implicit_atomic.py @@ -0,0 +1,140 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import pytest +import triton +import test_common +import triton.language as tl + +types_all = [ + (torch.float32, 'float32'), +] + + +def ceil_div(a, b): + return (a + b - 1) // b + + +@triton.jit +def addptr_implicit_perm_atomic_add_2d( + ptr, + out, + ynumel, + xnumel, + YBLOCK: tl.constexpr, + XBLOCK: tl.constexpr, +): + y = tl.program_id(1) * YBLOCK + tl.arange(0, YBLOCK)[None, :] # [1, YB] + x = tl.program_id(0) * XBLOCK + tl.arange(0, XBLOCK)[:, None] # [XB, 1] + + val = 1.0 + (x.to(tl.float32) * 0.01) + (y.to(tl.float32) * 0.001) # [XB, YB] + xmask = x < xnumel + ymask = y < ynumel + old = tl.atomic_add(ptr + (x + 4 * y), val, xmask & ymask) + + tl.store(out + (x + 4 * y), old) + + +@triton.jit +def addptr_implicit_perm_atomic_cas_2d( + ptr, + out, + cmp_ptr, + val_ptr, + ynumel, + xnumel, + YBLOCK: tl.constexpr, + XBLOCK: tl.constexpr, +): + y = tl.program_id(1) * YBLOCK + tl.arange(0, YBLOCK)[None, :] + x = tl.program_id(0) * XBLOCK + tl.arange(0, XBLOCK)[:, None] + + xmask = x < xnumel + ymask = y < ynumel + mask = xmask & ymask + + offset = x + 4 * y + + cmp = tl.load(cmp_ptr + offset, mask=mask, other=0.0).to(tl.float32) + val = tl.load(val_ptr + offset, mask=mask, other=0.0).to(tl.float32) + + old = tl.atomic_cas(ptr + offset, cmp, val) + + tl.store(out + offset, old, mask=mask) + + +@pytest.mark.parametrize('dtype,sigtype', types_all) +@pytest.mark.parametrize('xnumel, ynumel, XBLOCK, YBLOCK', [(4, 512, 4, 64)]) +def test_addptr_implicit_perm_atomic_add_2d( + dtype, + sigtype, + xnumel, + ynumel, + XBLOCK, + YBLOCK, +): + in_ptr = torch.zeros((ynumel * 4, ), dtype=dtype).npu() + out_ptr = torch.ones_like(in_ptr) + + grid = (ceil_div(xnumel, XBLOCK), ceil_div(ynumel, YBLOCK), 1) + addptr_implicit_perm_atomic_add_2d[grid](in_ptr, out_ptr, ynumel, xnumel, YBLOCK=YBLOCK, XBLOCK=XBLOCK) + + y_idx = torch.arange(ynumel).unsqueeze(1).npu() + x_idx = torch.arange(xnumel).unsqueeze(0).npu() + idx = (x_idx + 4 * y_idx).reshape(-1) + torch.testing.assert_close(out_ptr[idx], torch.zeros_like(out_ptr[idx])) + + val_ref = (1.0 + 0.01 * x_idx.to(torch.float32) + 0.001 * y_idx.to(torch.float32)).reshape(-1) + torch.testing.assert_close(in_ptr[idx], val_ref, rtol=1e-5, atol=1e-5) + + +@pytest.mark.parametrize('dtype,sigtype', types_all) +@pytest.mark.parametrize('xnumel, ynumel, XBLOCK, YBLOCK', [(4, 512, 4, 64)]) +def test_addptr_implicit_perm_atomic_cas_2d( + dtype, + sigtype, + xnumel, + ynumel, + XBLOCK, + YBLOCK, +): + in_ptr = torch.full((ynumel * 4, ), 2, dtype=dtype).npu() + out_ptr = torch.full((ynumel * 4, ), 1, dtype=dtype).npu() + cmp_ptr = torch.full((ynumel * 4, ), 2, dtype=dtype).npu() + val_ptr = torch.full((ynumel * 4, ), 1, dtype=dtype).npu() + + grid = (ceil_div(xnumel, XBLOCK), ceil_div(ynumel, YBLOCK), 1) + addptr_implicit_perm_atomic_cas_2d[grid](in_ptr, out_ptr, cmp_ptr, val_ptr, ynumel, xnumel, YBLOCK=YBLOCK, + XBLOCK=XBLOCK) + + y_idx = torch.arange(ynumel).unsqueeze(1).npu() + x_idx = torch.arange(xnumel).unsqueeze(0).npu() + idx = (x_idx + 4 * y_idx).reshape(-1) + + torch.testing.assert_close(out_ptr[idx], torch.full_like(out_ptr[idx], 2.0)) + + torch.testing.assert_close(in_ptr[idx], torch.ones_like(in_ptr[idx])) + + +if __name__ == '__main__': + case_2d = (4, 512, 4, 64) + test_addptr_implicit_perm_atomic_add_2d(*types_all[0], *case_2d) + test_addptr_implicit_perm_atomic_cas_2d(*types_all[0], *case_2d) diff --git a/third_party/ascend/unittest/pytest_ut/test_implicit_permute.py b/third_party/ascend/unittest/pytest_ut/test_implicit_permute.py new file mode 100644 index 0000000000..c5016fc665 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_implicit_permute.py @@ -0,0 +1,1338 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import pytest +import triton +import triton.language as tl +import test_common + +types_all = [ + (torch.float32, 'float32'), +] + +case_2d = [ + # X, Y, XBLOCK, YBLOCK + (512, 32, 4, 64), +] + +case_3d = [ + # X, Y, Z, XBLOCK, YBLOCK, ZBLOCK + (100, 40, 32, 10, 4, 4), +] + +case_4d = [ + # X, Y, Z, W, XBLOCK, YBLOCK, ZBLOCK, WBLOCK + (100, 80, 20, 16, 20, 4, 4, 4), +] + + +# ---------------------------------------------------------- +# Triton kernel: addptr with implicit permute kernels +# ---------------------------------------------------------- +@triton.jit +def addptr_implicit_perm_load_2d( + ptr, + out, + ynumel, + xnumel, + stride_y, + stride_x, + out_stride_x, + out_stride_y, + YBLOCK: tl.constexpr, + XBLOCK: tl.constexpr, +): + x = tl.program_id(0) * XBLOCK + tl.arange(0, XBLOCK)[:, None] + y = tl.program_id(1) * YBLOCK + tl.arange(0, YBLOCK)[None, :] + mask = (x < xnumel) & (y < ynumel) + + offset = x * stride_x + y * stride_y + val = tl.load(ptr + offset, mask=mask) + + out_offset = x * out_stride_x + y * out_stride_y + tl.store(out + out_offset, val, mask=mask) + + +@triton.jit +def addptr_implicit_perm_store_2d( + ptr, + out, + ynumel, + xnumel, + in_stride_x, + in_stride_y, + stride_y, + stride_x, + YBLOCK: tl.constexpr, + XBLOCK: tl.constexpr, +): + x = tl.program_id(0) * XBLOCK + tl.arange(0, XBLOCK)[:, None] + y = tl.program_id(1) * YBLOCK + tl.arange(0, YBLOCK)[None, :] + mask = (x < xnumel) & (y < ynumel) + + in_offset = x * in_stride_x + y * in_stride_y + val = tl.load(ptr + in_offset, mask=mask) + + out_offset = x * stride_x + y * stride_y + tl.store(out + out_offset, val, mask=mask) + + +@triton.jit +def addptr_implicit_perm_load_store_2d_static_stride(ptr, out, ynumel, xnumel, stride_y: tl.constexpr, + stride_x: tl.constexpr, YBLOCK: tl.constexpr, + XBLOCK: tl.constexpr): + # logical indices (A^T view) + x = tl.program_id(0) * XBLOCK + tl.arange(0, XBLOCK)[:, None] # [XBLOCK, 1] + y = tl.program_id(1) * YBLOCK + tl.arange(0, YBLOCK)[None, :] # [1, YBLOCK] + mask = (x < xnumel) & (y < ynumel) + + # IMPORTANT: + # ptr is a row-major A, but we interpret it as A^T via stride + offset = x * stride_x + y * stride_y + + val = tl.load(ptr + offset, mask) + tl.store(out + offset, val, mask) + + +@triton.jit +def addptr_implicit_perm_load_store_2d( + ptr, + out, + ynumel, + xnumel, + stride_y, + stride_x, + YBLOCK: tl.constexpr, + XBLOCK: tl.constexpr, +): + # logical indices (A^T view) + x = tl.program_id(0) * XBLOCK + tl.arange(0, XBLOCK)[:, None] # [XBLOCK, 1] + y = tl.program_id(1) * YBLOCK + tl.arange(0, YBLOCK)[None, :] # [1, YBLOCK] + + mask = (x < xnumel) & (y < ynumel) + + # IMPORTANT: + # ptr is a row-major A, but we interpret it as A^T via stride + offset = x * stride_x + y * stride_y + + val = tl.load(ptr + offset, mask=mask) + tl.store(out + offset, val, mask=mask) + + +@triton.jit +def addptr_implicit_perm_load_store_3d_static_stride( + ptr, + out, + znumel, + ynumel, + xnumel, + stride_z: tl.constexpr, + stride_y: tl.constexpr, + stride_x: tl.constexpr, + ZBLOCK: tl.constexpr, + YBLOCK: tl.constexpr, + XBLOCK: tl.constexpr, +): + pid_x = tl.program_id(0) + pid_y = tl.program_id(1) + pid_z = tl.program_id(2) + + x = pid_x * XBLOCK + tl.arange(0, XBLOCK)[:, None, None] + y = pid_y * YBLOCK + tl.arange(0, YBLOCK)[None, :, None] + z = pid_z * ZBLOCK + tl.arange(0, ZBLOCK)[None, None, :] + + mask = (x < xnumel) & (y < ynumel) & (z < znumel) + + offset = x * stride_x + y * stride_y + z * stride_z + val = tl.load(ptr + offset, mask=mask) + tl.store(out + offset, val, mask=mask) + + +@triton.jit +def addptr_implicit_perm_load_store_3d( + ptr, + out, + znumel, + ynumel, + xnumel, + stride_z, + stride_y, + stride_x, + ZBLOCK: tl.constexpr, + YBLOCK: tl.constexpr, + XBLOCK: tl.constexpr, +): + pid_x = tl.program_id(0) + pid_y = tl.program_id(1) + pid_z = tl.program_id(2) + + x = pid_x * XBLOCK + tl.arange(0, XBLOCK)[:, None, None] + y = pid_y * YBLOCK + tl.arange(0, YBLOCK)[None, :, None] + z = pid_z * ZBLOCK + tl.arange(0, ZBLOCK)[None, None, :] + + mask = (x < xnumel) & (y < ynumel) & (z < znumel) + + offset = x * stride_x + y * stride_y + z * stride_z + val = tl.load(ptr + offset, mask=mask) + tl.store(out + offset, val, mask=mask) + + +@triton.jit +def addptr_implicit_perm_load_store_4d_static_stride( + ptr, + out, + wnumel, + znumel, + ynumel, + xnumel, + stride_w: tl.constexpr, + stride_z: tl.constexpr, + stride_y: tl.constexpr, + stride_x: tl.constexpr, + WBLOCK: tl.constexpr, + ZBLOCK: tl.constexpr, + YBLOCK: tl.constexpr, + XBLOCK: tl.constexpr, +): + pid0 = tl.program_id(0) # covers (w, x) + pid1 = tl.program_id(1) # y + pid2 = tl.program_id(2) # z + + xblocks_per_w = (xnumel + XBLOCK - 1) // XBLOCK + + w_pid = pid0 // xblocks_per_w + x_pid = pid0 - w_pid * xblocks_per_w + + x0 = x_pid * XBLOCK + y0 = pid1 * YBLOCK + z0 = pid2 * ZBLOCK + w0 = w_pid * WBLOCK + + x = x0 + tl.arange(0, XBLOCK)[:, None, None, None] + y = y0 + tl.arange(0, YBLOCK)[None, :, None, None] + z = z0 + tl.arange(0, ZBLOCK)[None, None, :, None] + w = w0 + tl.arange(0, WBLOCK)[None, None, None, :] + + mask = (x < xnumel) & (y < ynumel) & (z < znumel) & (w < wnumel) + + offset = x * stride_x + y * stride_y + z * stride_z + w * stride_w + val = tl.load(ptr + offset, mask=mask, other=0.0) + tl.store(out + offset, val, mask=mask) + + +@triton.jit +def addptr_implicit_perm_load_store_4d( + ptr, + out, + wnumel, + znumel, + ynumel, + xnumel, + stride_w, + stride_z, + stride_y, + stride_x, + WBLOCK: tl.constexpr, + ZBLOCK: tl.constexpr, + YBLOCK: tl.constexpr, + XBLOCK: tl.constexpr, +): + pid0 = tl.program_id(0) # covers (w, x) + pid1 = tl.program_id(1) # y + pid2 = tl.program_id(2) # z + + xblocks_per_w = (xnumel + XBLOCK - 1) // XBLOCK + + w_pid = pid0 // xblocks_per_w + x_pid = pid0 - w_pid * xblocks_per_w + + x0 = x_pid * XBLOCK + y0 = pid1 * YBLOCK + z0 = pid2 * ZBLOCK + w0 = w_pid * WBLOCK + + x = x0 + tl.arange(0, XBLOCK)[:, None, None, None] + y = y0 + tl.arange(0, YBLOCK)[None, :, None, None] + z = z0 + tl.arange(0, ZBLOCK)[None, None, :, None] + w = w0 + tl.arange(0, WBLOCK)[None, None, None, :] + + mask = (x < xnumel) & (y < ynumel) & (z < znumel) & (w < wnumel) + + offset = x * stride_x + y * stride_y + z * stride_z + w * stride_w + val = tl.load(ptr + offset, mask=mask, other=0.0) + tl.store(out + offset, val, mask=mask) + + +# ---------------------------------------------------------- +# Triton kernel: make_tensor_ptr with implicit permute kernels +# ---------------------------------------------------------- +@triton.jit +def make_tensor_ptr_implicit_perm_load_2d( + ptr, + out, + ynumel, + xnumel, + STRIDE_Y, + STRIDE_X, + OUT_STRIDE_X, + OUT_STRIDE_Y, + YBLOCK: tl.constexpr, + XBLOCK: tl.constexpr, +): + y0 = tl.program_id(1) * YBLOCK + x0 = tl.program_id(0) * XBLOCK + y = y0 + tl.arange(0, YBLOCK)[None, :] + x = x0 + tl.arange(0, XBLOCK)[:, None] + xmask = x < xnumel + ymask = y < ynumel + mask = xmask & ymask + + tptr = tl.make_block_ptr( + base=ptr, + shape=(xnumel, ynumel), + strides=(STRIDE_X, STRIDE_Y), + offsets=(x0, y0), + block_shape=(XBLOCK, YBLOCK), + order=(0, 1), + ) + + val = tl.load(tptr) + out_offset = x * OUT_STRIDE_X + y * OUT_STRIDE_Y + tl.store(out + out_offset, val, mask=mask) + + +@triton.jit +def make_tensor_ptr_implicit_perm_store_2d( + ptr, + out, + ynumel, + xnumel, + IN_STRIDE_X, + IN_STRIDE_Y, + STRIDE_Y, + STRIDE_X, + YBLOCK: tl.constexpr, + XBLOCK: tl.constexpr, +): + y0 = tl.program_id(1) * YBLOCK + x0 = tl.program_id(0) * XBLOCK + y = y0 + tl.arange(0, YBLOCK)[None, :] + x = x0 + tl.arange(0, XBLOCK)[:, None] + xmask = x < xnumel + ymask = y < ynumel + mask = xmask & ymask + + in_offset = x * IN_STRIDE_X + y * IN_STRIDE_Y + val = tl.load(ptr + in_offset, mask=mask) + + tptr = tl.make_block_ptr( + base=out, + shape=(xnumel, ynumel), + strides=(STRIDE_X, STRIDE_Y), + offsets=(x0, y0), + block_shape=(XBLOCK, YBLOCK), + order=(0, 1), + ) + tl.store(tptr, val, boundary_check=(0, 1)) + + +@triton.jit +def make_tensor_ptr_implicit_perm_load_store_2d_static_stride(ptr, out, ynumel, xnumel, STRIDE_Y: tl.constexpr, + STRIDE_X: tl.constexpr, YBLOCK: tl.constexpr, + XBLOCK: tl.constexpr): + y0 = tl.program_id(1) * YBLOCK + x0 = tl.program_id(0) * XBLOCK + y = y0 + tl.arange(0, YBLOCK)[None, :] # [1, YBLOCK] + x = x0 + tl.arange(0, XBLOCK)[:, None] # [XBLOCK, 1] + xmask = x < xnumel + ymask = y < ynumel + mask = xmask & ymask + + tptr = tl.make_block_ptr( + base=ptr, + shape=(xnumel, ynumel), + strides=(STRIDE_X, STRIDE_Y), + offsets=(x0, y0), + block_shape=(XBLOCK, YBLOCK), + order=(0, 1), + ) + + val = tl.load(tptr) + tl.store(out + (x * STRIDE_X + STRIDE_Y * y), val, mask=mask) + + +@triton.jit +def make_tensor_ptr_implicit_perm_load_store_3d_static_stride( + ptr, + out, + znumel, + ynumel, + xnumel, + STRIDE_Z: tl.constexpr, + STRIDE_Y: tl.constexpr, + STRIDE_X: tl.constexpr, + ZBLOCK: tl.constexpr, + YBLOCK: tl.constexpr, + XBLOCK: tl.constexpr, +): + pid_x = tl.program_id(0) + pid_y = tl.program_id(1) + pid_z = tl.program_id(2) + + x0 = pid_x * XBLOCK + y0 = pid_y * YBLOCK + z0 = pid_z * ZBLOCK + + x = x0 + tl.arange(0, XBLOCK)[:, None, None] + y = y0 + tl.arange(0, YBLOCK)[None, :, None] + z = z0 + tl.arange(0, ZBLOCK)[None, None, :] + + mask = (x < xnumel) & (y < ynumel) & (z < znumel) + + tptr = tl.make_block_ptr( + base=ptr, + shape=(xnumel, ynumel, znumel), + strides=(STRIDE_X, STRIDE_Y, STRIDE_Z), + offsets=(x0, y0, z0), + block_shape=(XBLOCK, YBLOCK, ZBLOCK), + order=(0, 1, 2), + ) + + val = tl.load(tptr) + tl.store(out + (x * STRIDE_X + y * STRIDE_Y + z * STRIDE_Z), val, mask=mask) + + +@triton.jit +def make_tensor_ptr_implicit_perm_load_store_3d( + ptr, + out, + znumel, + ynumel, + xnumel, + STRIDE_Z, + STRIDE_Y, + STRIDE_X, + ZBLOCK: tl.constexpr, + YBLOCK: tl.constexpr, + XBLOCK: tl.constexpr, +): + pid_x = tl.program_id(0) + pid_y = tl.program_id(1) + pid_z = tl.program_id(2) + + x0 = pid_x * XBLOCK + y0 = pid_y * YBLOCK + z0 = pid_z * ZBLOCK + + x = x0 + tl.arange(0, XBLOCK)[:, None, None] + y = y0 + tl.arange(0, YBLOCK)[None, :, None] + z = z0 + tl.arange(0, ZBLOCK)[None, None, :] + + mask = (x < xnumel) & (y < ynumel) & (z < znumel) + + tptr = tl.make_block_ptr( + base=ptr, + shape=(xnumel, ynumel, znumel), + strides=(STRIDE_X, STRIDE_Y, STRIDE_Z), + offsets=(x0, y0, z0), + block_shape=(XBLOCK, YBLOCK, ZBLOCK), + order=(0, 1, 2), + ) + + val = tl.load(tptr) + tl.store(out + (x * STRIDE_X + y * STRIDE_Y + z * STRIDE_Z), val, mask=mask) + + +@triton.jit +def make_tensor_ptr_implicit_perm_load_3d_static_stride( + ptr, + out, + znumel, # logical z (== X) + ynumel, # logical y (== Y) + xnumel, # logical x (== Z) + STRIDE_Z: tl.constexpr, + STRIDE_Y: tl.constexpr, + STRIDE_X: tl.constexpr, + # out is row-major with shape (xnumel, ynumel, znumel) + OUT_STRIDE0: tl.constexpr, # = ynumel*znumel + OUT_STRIDE1: tl.constexpr, # = znumel + OUT_STRIDE2: tl.constexpr, # = 1 + ZBLOCK: tl.constexpr, + YBLOCK: tl.constexpr, + XBLOCK: tl.constexpr, +): + pid_x = tl.program_id(0) + pid_y = tl.program_id(1) + pid_z = tl.program_id(2) + + x0 = pid_x * XBLOCK + y0 = pid_y * YBLOCK + z0 = pid_z * ZBLOCK + + x = x0 + tl.arange(0, XBLOCK)[:, None, None] + y = y0 + tl.arange(0, YBLOCK)[None, :, None] + z = z0 + tl.arange(0, ZBLOCK)[None, None, :] + + mask = (x < xnumel) & (y < ynumel) & (z < znumel) + + # load: implicit permute view + tptr = tl.make_block_ptr( + base=ptr, + shape=(xnumel, ynumel, znumel), + strides=(STRIDE_X, STRIDE_Y, STRIDE_Z), + offsets=(x0, y0, z0), + block_shape=(XBLOCK, YBLOCK, ZBLOCK), + order=(0, 1, 2), + ) + val = tl.load(tptr) + + # store: row-major output (no implicit permute) + out_offset = x * OUT_STRIDE0 + y * OUT_STRIDE1 + z * OUT_STRIDE2 + tl.store(out + out_offset, val, mask=mask) + + +# ---------------------------------------------------------- +# Triton kernel: advance with implicit permute kernels +# ---------------------------------------------------------- +@triton.jit +def advance_implicit_perm_load_store_2d_static_stride(ptr, out, ynumel, xnumel, STRIDE_Y: tl.constexpr, + STRIDE_X: tl.constexpr, YBLOCK: tl.constexpr, + XBLOCK: tl.constexpr): + y0 = tl.program_id(1) * YBLOCK + x0 = tl.program_id(0) * XBLOCK + y = y0 + tl.arange(0, YBLOCK)[None, :] + x = x0 + tl.arange(0, XBLOCK)[:, None] + mask = (x < xnumel) & (y < ynumel) + + tptr = tl.make_block_ptr( + base=ptr, + shape=(xnumel, ynumel), + strides=(STRIDE_X, STRIDE_Y), + offsets=(0, 0), + block_shape=(XBLOCK, YBLOCK), + order=(0, 1), + ) + tptr2 = tl.advance(tptr, (x0, y0)) + val = tl.load(tptr2) + tl.store(out + (x * STRIDE_X + y * STRIDE_Y), val, mask=mask) + + +@triton.jit +def advance_implicit_perm_load_store_3d_static_stride( + ptr, + out, + znumel, + ynumel, + xnumel, + STRIDE_Z: tl.constexpr, + STRIDE_Y: tl.constexpr, + STRIDE_X: tl.constexpr, + ZBLOCK: tl.constexpr, + YBLOCK: tl.constexpr, + XBLOCK: tl.constexpr, +): + pid_x = tl.program_id(0) + pid_y = tl.program_id(1) + pid_z = tl.program_id(2) + + x0 = pid_x * XBLOCK + y0 = pid_y * YBLOCK + z0 = pid_z * ZBLOCK + + x = x0 + tl.arange(0, XBLOCK)[:, None, None] + y = y0 + tl.arange(0, YBLOCK)[None, :, None] + z = z0 + tl.arange(0, ZBLOCK)[None, None, :] + + mask = (x < xnumel) & (y < ynumel) & (z < znumel) + + tptr0 = tl.make_block_ptr( + base=ptr, + shape=(xnumel, ynumel, znumel), + strides=(STRIDE_X, STRIDE_Y, STRIDE_Z), + offsets=(0, 0, 0), + block_shape=(XBLOCK, YBLOCK, ZBLOCK), + order=(0, 1, 2), + ) + tptr = tl.advance(tptr0, (x0, y0, z0)) + val = tl.load(tptr) + tl.store(out + (x * STRIDE_X + y * STRIDE_Y + z * STRIDE_Z), val, mask=mask) + + +# ---------------------------------------------------------- +# pytest case +# ---------------------------------------------------------- +def ceil_div(a, b): + return (a + b - 1) // b + + +def _assert_row_major_2d(A, X, Y): + assert tuple(A.shape) == (X, Y) + assert A.is_contiguous() + assert A.stride() == (Y, 1) + + +def _assert_row_major_3d(A, X, Y, Z): + # [X, Y, Z] contiguous -> stride = (Y*Z, Z, 1) + assert tuple(A.shape) == (X, Y, Z) + assert A.is_contiguous() + assert A.stride() == (Y * Z, Z, 1) + + +def _assert_row_major_4d(A, X, Y, Z, W): + # [X, Y, Z, W] contiguous -> stride = (Y*Z*W, Z*W, W, 1) + assert tuple(A.shape) == (X, Y, Z, W) + assert A.is_contiguous() + assert A.stride() == (Y * Z * W, Z * W, W, 1) + + +# ---------------------------------------------------------- +# pytest case: addptr kernels +# ---------------------------------------------------------- +@pytest.mark.parametrize("X, Y, XBLOCK, YBLOCK", case_2d) +@pytest.mark.parametrize("dtype, sigtype", types_all) +def test_addptr_implicit_perm_load_2d(X, Y, XBLOCK, YBLOCK, dtype, sigtype): + A = test_common.generate_tensor(shape=(X, Y), dtype=sigtype).npu() + _assert_row_major_2d(A, X, Y) + + xnumel = Y + ynumel = X + + STRIDE_X = 1 + STRIDE_Y = Y + + out = torch.empty((xnumel, ynumel), device="npu", dtype=A.dtype) + assert out.is_contiguous() + OUT_STRIDE_X = ynumel + OUT_STRIDE_Y = 1 + + grid = (ceil_div(xnumel, XBLOCK), ceil_div(ynumel, YBLOCK), 1) + + addptr_implicit_perm_load_2d[grid]( + A, + out, + ynumel, + xnumel, + STRIDE_Y, + STRIDE_X, + OUT_STRIDE_X, + OUT_STRIDE_Y, + YBLOCK=YBLOCK, + XBLOCK=XBLOCK, + ) + + ref = A.T.contiguous() + torch.testing.assert_close(out, ref) + + +@pytest.mark.parametrize("X, Y, XBLOCK, YBLOCK", case_2d) +@pytest.mark.parametrize("dtype, sigtype", types_all) +def test_addptr_implicit_perm_store_2d(X, Y, XBLOCK, YBLOCK, dtype, sigtype): + """ + Test goal: + - Real memory layout: A[X, Y], row-major (stride = (Y, 1)) + - Load uses row-major semantics only + - Store uses implicit transpose semantics only + - Output shape: [Y, X] row-major + - Result: out == A.T + """ + A = test_common.generate_tensor(shape=(X, Y), dtype=sigtype).npu() + _assert_row_major_2d(A, X, Y) + + xnumel = X + ynumel = Y + + IN_STRIDE_X = Y + IN_STRIDE_Y = 1 + + STRIDE_X = 1 + STRIDE_Y = X + + out = torch.empty((Y, X), device="npu", dtype=A.dtype) + assert out.is_contiguous() + + grid = (ceil_div(xnumel, XBLOCK), ceil_div(ynumel, YBLOCK), 1) + + addptr_implicit_perm_store_2d[grid]( + A, + out, + ynumel, + xnumel, + IN_STRIDE_X, + IN_STRIDE_Y, + STRIDE_Y, + STRIDE_X, + YBLOCK=YBLOCK, + XBLOCK=XBLOCK, + ) + + ref = A.T.contiguous() + torch.testing.assert_close(out, ref) + + +@pytest.mark.parametrize("X, Y, XBLOCK, YBLOCK", case_2d) +@pytest.mark.parametrize("dtype, sigtype", types_all) +def test_addptr_implicit_perm_load_store_2d_static_stride( + X, + Y, + XBLOCK, + YBLOCK, + dtype, + sigtype, +): + """ + Test goal: + - Real memory layout: A[X, Y], row-major (stride = (Y, 1)) + - Kernel view: A^T[Y, X], stride = (1, Y) + - Kernel does load+store with identical offsets + - Result must satisfy: out == in + """ + A = test_common.generate_tensor( + shape=(X, Y), + dtype=sigtype, + ).npu() + + _assert_row_major_2d(A, X, Y) + + out = torch.zeros_like(A) + + # A^T logical shape + xnumel = Y # cols of A + ynumel = X # rows of A + + # A^T logical stride + stride_x = 1 + stride_y = Y + + grid = ( + ceil_div(xnumel, XBLOCK), + ceil_div(ynumel, YBLOCK), + 1, + ) + + addptr_implicit_perm_load_store_2d_static_stride[grid]( + A, + out, + ynumel, + xnumel, + stride_y, + stride_x, + XBLOCK=XBLOCK, + YBLOCK=YBLOCK, + ) + + torch.testing.assert_close(out, A) + + +@pytest.mark.parametrize("X, Y, XBLOCK, YBLOCK", case_2d) +@pytest.mark.parametrize("dtype, sigtype", types_all) +def test_addptr_implicit_perm_load_store_2d( + X, + Y, + XBLOCK, + YBLOCK, + dtype, + sigtype, +): + """ + Same as static-stride version, but stride passed as runtime values. + """ + A = test_common.generate_tensor( + shape=(X, Y), + dtype=sigtype, + ).npu() + + _assert_row_major_2d(A, X, Y) + + out = torch.zeros_like(A) + + # A^T logical shape + xnumel = Y # cols of A + ynumel = X # rows of A + + # A^T logical stride + stride_x = 1 + stride_y = Y + + grid = ( + ceil_div(xnumel, XBLOCK), + ceil_div(ynumel, YBLOCK), + 1, + ) + + addptr_implicit_perm_load_store_2d[grid]( + A, + out, + ynumel, + xnumel, + stride_y, + stride_x, + XBLOCK=XBLOCK, + YBLOCK=YBLOCK, + ) + + torch.testing.assert_close(out, A) + + +@pytest.mark.parametrize("X, Y, Z, XBLOCK, YBLOCK, ZBLOCK", case_3d) +@pytest.mark.parametrize("dtype, sigtype", types_all) +def test_addptr_implicit_perm_load_store_3d_static_stride(X, Y, Z, XBLOCK, YBLOCK, ZBLOCK, dtype, sigtype): + """ + Test goal: + - Real memory layout: A[X, Y, Z], row-major (stride = (Y*Z, Z, 1)) + - Kernel view: treat as permuted logical coords via stride: + offset = x*1 + y*Z + z*(Y*Z) + i.e. (x,y,z) mapped to base index (z, y, x) + - Kernel does load+store with identical offsets + - Result must satisfy: out == in + """ + A = test_common.generate_tensor(shape=(X, Y, Z), dtype=sigtype).npu() + _assert_row_major_3d(A, X, Y, Z) + out = torch.zeros_like(A) + + # Logical shape for "A^(perm)" (x fastest) + xnumel = Z + ynumel = Y + znumel = X + + # Logical strides (in elements): (1, Z, Y*Z) + stride_x = 1 + stride_y = Z + stride_z = Y * Z + + grid = ( + ceil_div(xnumel, XBLOCK), + ceil_div(ynumel, YBLOCK), + ceil_div(znumel, ZBLOCK), + ) + + addptr_implicit_perm_load_store_3d_static_stride[grid]( + A, + out, + znumel, + ynumel, + xnumel, + stride_z, + stride_y, + stride_x, + ZBLOCK=ZBLOCK, + YBLOCK=YBLOCK, + XBLOCK=XBLOCK, + ) + + torch.testing.assert_close(out, A) + + +@pytest.mark.parametrize("X, Y, Z, XBLOCK, YBLOCK, ZBLOCK", case_3d) +@pytest.mark.parametrize("dtype, sigtype", types_all) +def test_addptr_implicit_perm_load_store_3d(X, Y, Z, XBLOCK, YBLOCK, ZBLOCK, dtype, sigtype): + """ + Same as static-stride version, but stride passed as runtime values. + """ + A = test_common.generate_tensor(shape=(X, Y, Z), dtype=sigtype).npu() + _assert_row_major_3d(A, X, Y, Z) + out = torch.zeros_like(A) + + xnumel = Z + ynumel = Y + znumel = X + + stride_x = 1 + stride_y = Z + stride_z = Y * Z + + grid = ( + ceil_div(xnumel, XBLOCK), + ceil_div(ynumel, YBLOCK), + ceil_div(znumel, ZBLOCK), + ) + + addptr_implicit_perm_load_store_3d[grid]( + A, + out, + znumel, + ynumel, + xnumel, + stride_z, + stride_y, + stride_x, + ZBLOCK=ZBLOCK, + YBLOCK=YBLOCK, + XBLOCK=XBLOCK, + ) + + torch.testing.assert_close(out, A) + + +@pytest.mark.parametrize("X, Y, Z, W, XBLOCK, YBLOCK, ZBLOCK, WBLOCK", case_4d) +@pytest.mark.parametrize("dtype, sigtype", types_all) +def test_addptr_implicit_perm_load_store_4d_static_stride(X, Y, Z, W, XBLOCK, YBLOCK, ZBLOCK, WBLOCK, dtype, sigtype): + """ + Test goal: + - Real memory layout: A[X, Y, Z, W], row-major (stride = (Y*Z*W, Z*W, W, 1)) + - Kernel view: treat as permuted logical coords via stride: + offset = x*1 + y*W + z*(Z*W) + w*(Y*Z*W) + i.e. (x,y,z,w) mapped to base index (w, z, y, x) + - Kernel does load+store with identical offsets + - Result must satisfy: out == in + """ + A = test_common.generate_tensor(shape=(X, Y, Z, W), dtype=sigtype).npu() + _assert_row_major_4d(A, X, Y, Z, W) + out = torch.zeros_like(A) + + # Logical shape (x fastest) + xnumel = W + ynumel = Z + znumel = Y + wnumel = X + + # Logical strides (in elements): (1, W, Z*W, Y*Z*W) + stride_x = 1 + stride_y = W + stride_z = Z * W + stride_w = Y * Z * W + + # Kernel maps pid0 over (w, x). It uses xblocks_per_w computed from xnumel. + xblocks_per_w = ceil_div(xnumel, XBLOCK) + grid0 = wnumel * xblocks_per_w + grid = ( + grid0, + ceil_div(ynumel, YBLOCK), + ceil_div(znumel, ZBLOCK), + ) + + addptr_implicit_perm_load_store_4d_static_stride[grid]( + A, + out, + wnumel, + znumel, + ynumel, + xnumel, + stride_w, + stride_z, + stride_y, + stride_x, + WBLOCK=WBLOCK, + ZBLOCK=ZBLOCK, + YBLOCK=YBLOCK, + XBLOCK=XBLOCK, + ) + + torch.testing.assert_close(out, A) + + +@pytest.mark.parametrize("X, Y, Z, W, XBLOCK, YBLOCK, ZBLOCK, WBLOCK", case_4d) +@pytest.mark.parametrize("dtype, sigtype", types_all) +def test_addptr_implicit_perm_load_store_4d(X, Y, Z, W, XBLOCK, YBLOCK, ZBLOCK, WBLOCK, dtype, sigtype): + """ + Same as static-stride version, but stride passed as runtime values. + """ + A = test_common.generate_tensor(shape=(X, Y, Z, W), dtype=sigtype).npu() + _assert_row_major_4d(A, X, Y, Z, W) + out = torch.zeros_like(A) + + xnumel = W + ynumel = Z + znumel = Y + wnumel = X + + stride_x = 1 + stride_y = W + stride_z = Z * W + stride_w = Y * Z * W + + xblocks_per_w = ceil_div(xnumel, XBLOCK) + grid0 = wnumel * xblocks_per_w + grid = ( + grid0, + ceil_div(ynumel, YBLOCK), + ceil_div(znumel, ZBLOCK), + ) + + addptr_implicit_perm_load_store_4d[grid]( + A, + out, + wnumel, + znumel, + ynumel, + xnumel, + stride_w, + stride_z, + stride_y, + stride_x, + WBLOCK=WBLOCK, + ZBLOCK=ZBLOCK, + YBLOCK=YBLOCK, + XBLOCK=XBLOCK, + ) + + torch.testing.assert_close(out, A) + + +# ---------------------------------------------------------- +# pytest case: make_tensor_ptr kernels +# ---------------------------------------------------------- +@pytest.mark.parametrize("X, Y, XBLOCK, YBLOCK", case_2d) +@pytest.mark.parametrize("dtype, sigtype", types_all) +def test_make_tensor_ptr_implicit_perm_load_2d(X, Y, XBLOCK, YBLOCK, dtype, sigtype): + """ + Test goal: + - Real memory layout: A[X, Y], row-major (stride = (Y, 1)) + - Kernel view (logical): shape = (Y, X), stride = (1, Y) (implicit transpose) + - Load via make_block_ptr with implicit permute stride + - Store with row-major output stride + - Output shape: [Y, X] row-major + - Result: out == A.T + """ + A = test_common.generate_tensor(shape=(X, Y), dtype=sigtype).npu() + _assert_row_major_2d(A, X, Y) + + xnumel = Y + ynumel = X + + STRIDE_X = 1 + STRIDE_Y = Y + + out = torch.empty((xnumel, ynumel), device="npu", dtype=A.dtype) + assert out.is_contiguous() + OUT_STRIDE_X = ynumel + OUT_STRIDE_Y = 1 + + grid = (ceil_div(xnumel, XBLOCK), ceil_div(ynumel, YBLOCK), 1) + + make_tensor_ptr_implicit_perm_load_2d[grid]( + A, + out, + ynumel, + xnumel, + STRIDE_Y, + STRIDE_X, + OUT_STRIDE_X, + OUT_STRIDE_Y, + YBLOCK=YBLOCK, + XBLOCK=XBLOCK, + ) + + ref = A.T.contiguous() + torch.testing.assert_close(out, ref) + + +@pytest.mark.parametrize("X, Y, XBLOCK, YBLOCK", case_2d) +@pytest.mark.parametrize("dtype, sigtype", types_all) +def test_make_tensor_ptr_implicit_perm_store_2d(X, Y, XBLOCK, YBLOCK, dtype, sigtype): + """ + Test goal: + - Real memory layout: A[X, Y], row-major (stride = (Y, 1)) + - Load uses row-major semantics only + - Store via make_block_ptr with implicit transpose semantics only + - Output shape: [Y, X] row-major + - Result: out == A.T + """ + A = test_common.generate_tensor(shape=(X, Y), dtype=sigtype).npu() + _assert_row_major_2d(A, X, Y) + + xnumel = X + ynumel = Y + + IN_STRIDE_X = Y + IN_STRIDE_Y = 1 + + STRIDE_X = 1 + STRIDE_Y = X + + out = torch.empty((Y, X), device="npu", dtype=A.dtype) + assert out.is_contiguous() + + grid = (ceil_div(xnumel, XBLOCK), ceil_div(ynumel, YBLOCK), 1) + + make_tensor_ptr_implicit_perm_store_2d[grid]( + A, + out, + ynumel, + xnumel, + IN_STRIDE_X, + IN_STRIDE_Y, + STRIDE_Y, + STRIDE_X, + YBLOCK=YBLOCK, + XBLOCK=XBLOCK, + ) + + ref = A.T.contiguous() + torch.testing.assert_close(out, ref) + + +@pytest.mark.parametrize("X, Y, XBLOCK, YBLOCK", case_2d) +@pytest.mark.parametrize("dtype, sigtype", types_all) +def test_make_tensor_ptr_implicit_perm_load_store_2d_static_stride(X, Y, XBLOCK, YBLOCK, dtype, sigtype): + """ + Test goal matches addptr_2d_static_stride, but uses tl.make_block_ptr + tl.load(tptr). + Real layout: A[X,Y] row-major stride=(Y,1) + Kernel view: A^T[Y,X] stride=(1,Y) + Store is by explicit linear offset with same logical stride. + """ + A = test_common.generate_tensor(shape=(X, Y), dtype=sigtype).npu() + _assert_row_major_2d(A, X, Y) + out = torch.zeros_like(A) + + xnumel = Y + ynumel = X + STRIDE_X = 1 + STRIDE_Y = Y + + grid = (ceil_div(xnumel, XBLOCK), ceil_div(ynumel, YBLOCK), 1) + make_tensor_ptr_implicit_perm_load_store_2d_static_stride[grid]( + A, + out, + ynumel, + xnumel, + STRIDE_Y=STRIDE_Y, + STRIDE_X=STRIDE_X, + YBLOCK=YBLOCK, + XBLOCK=XBLOCK, + ) + + torch.testing.assert_close(out, A) + + +@pytest.mark.parametrize("X, Y, Z, XBLOCK, YBLOCK, ZBLOCK", case_3d) +@pytest.mark.parametrize("dtype, sigtype", types_all) +def test_make_tensor_ptr_implicit_perm_load_store_3d_static_stride(X, Y, Z, XBLOCK, YBLOCK, ZBLOCK, dtype, sigtype): + """ + Real layout: A[X,Y,Z] row-major stride=(Y*Z, Z, 1) + Kernel view (logical): shape=(Z,Y,X), strides=(1, Z, Y*Z) + """ + A = test_common.generate_tensor(shape=(X, Y, Z), dtype=sigtype).npu() + _assert_row_major_3d(A, X, Y, Z) + out = torch.zeros_like(A) + + xnumel = Z + ynumel = Y + znumel = X + STRIDE_X = 1 + STRIDE_Y = Z + STRIDE_Z = Y * Z + + grid = ( + ceil_div(xnumel, XBLOCK), + ceil_div(ynumel, YBLOCK), + ceil_div(znumel, ZBLOCK), + ) + + make_tensor_ptr_implicit_perm_load_store_3d_static_stride[grid]( + A, + out, + znumel, + ynumel, + xnumel, + STRIDE_Z=STRIDE_Z, + STRIDE_Y=STRIDE_Y, + STRIDE_X=STRIDE_X, + ZBLOCK=ZBLOCK, + YBLOCK=YBLOCK, + XBLOCK=XBLOCK, + ) + + torch.testing.assert_close(out, A) + + +@pytest.mark.parametrize("X, Y, Z, XBLOCK, YBLOCK, ZBLOCK", case_3d) +@pytest.mark.parametrize("dtype, sigtype", types_all) +def test_make_tensor_ptr_implicit_perm_load_store_3d(X, Y, Z, XBLOCK, YBLOCK, ZBLOCK, dtype, sigtype): + """ + Same as static stride but STRIDE_* passed at runtime. + """ + A = test_common.generate_tensor(shape=(X, Y, Z), dtype=sigtype).npu() + _assert_row_major_3d(A, X, Y, Z) + out = torch.zeros_like(A) + + xnumel = Z + ynumel = Y + znumel = X + STRIDE_X = 1 + STRIDE_Y = Z + STRIDE_Z = Y * Z + + grid = ( + ceil_div(xnumel, XBLOCK), + ceil_div(ynumel, YBLOCK), + ceil_div(znumel, ZBLOCK), + ) + + make_tensor_ptr_implicit_perm_load_store_3d[grid]( + A, + out, + znumel, + ynumel, + xnumel, + STRIDE_Z, + STRIDE_Y, + STRIDE_X, + ZBLOCK=ZBLOCK, + YBLOCK=YBLOCK, + XBLOCK=XBLOCK, + ) + + torch.testing.assert_close(out, A) + + +@pytest.mark.parametrize("X, Y, Z, XBLOCK, YBLOCK, ZBLOCK", case_3d) +@pytest.mark.parametrize("dtype, sigtype", types_all) +def test_make_tensor_ptr_implicit_perm_load_3d_static_stride(X, Y, Z, XBLOCK, YBLOCK, ZBLOCK, dtype, sigtype): + A = test_common.generate_tensor(shape=(X, Y, Z), dtype=sigtype).npu() + _assert_row_major_3d(A, X, Y, Z) + + # logical/permuted shape + xnumel = Z + ynumel = Y + znumel = X + + # implicit-permute strides (elements) + STRIDE_X = 1 + STRIDE_Y = Z + STRIDE_Z = Y * Z + + # output is row-major of shape (Z, Y, X) + out = torch.empty((xnumel, ynumel, znumel), device="npu", dtype=A.dtype) + assert out.is_contiguous() + OUT_STRIDE0 = ynumel * znumel # Y*X + OUT_STRIDE1 = znumel # X + OUT_STRIDE2 = 1 + + grid = ( + ceil_div(xnumel, XBLOCK), + ceil_div(ynumel, YBLOCK), + ceil_div(znumel, ZBLOCK), + ) + + make_tensor_ptr_implicit_perm_load_3d_static_stride[grid]( + A, + out, + znumel, + ynumel, + xnumel, + STRIDE_Z=STRIDE_Z, + STRIDE_Y=STRIDE_Y, + STRIDE_X=STRIDE_X, + OUT_STRIDE0=OUT_STRIDE0, + OUT_STRIDE1=OUT_STRIDE1, + OUT_STRIDE2=OUT_STRIDE2, + ZBLOCK=ZBLOCK, + YBLOCK=YBLOCK, + XBLOCK=XBLOCK, + ) + + # expected: out[x,y,z] == A[z,y,x] => out == A.permute(2,1,0) + ref = A.permute(2, 1, 0).contiguous() + torch.testing.assert_close(out, ref) + + +# ---------------------------------------------------------- +# pytest case: advance kernels +# ---------------------------------------------------------- +@pytest.mark.parametrize("X, Y, XBLOCK, YBLOCK", case_2d) +@pytest.mark.parametrize("dtype, sigtype", types_all) +def test_advance_implicit_perm_load_store_2d_static_stride(X, Y, XBLOCK, YBLOCK, dtype, sigtype): + """ + Same goal as addptr_2d_static_stride, but uses tl.make_block_ptr + tl.advance. + """ + A = test_common.generate_tensor(shape=(X, Y), dtype=sigtype).npu() + _assert_row_major_2d(A, X, Y) + out = torch.zeros_like(A) + + xnumel = Y + ynumel = X + STRIDE_X = 1 + STRIDE_Y = Y + + grid = (ceil_div(xnumel, XBLOCK), ceil_div(ynumel, YBLOCK), 1) + advance_implicit_perm_load_store_2d_static_stride[grid]( + A, + out, + ynumel, + xnumel, + STRIDE_Y=STRIDE_Y, + STRIDE_X=STRIDE_X, + YBLOCK=YBLOCK, + XBLOCK=XBLOCK, + ) + + torch.testing.assert_close(out, A) + + +@pytest.mark.parametrize("X, Y, Z, XBLOCK, YBLOCK, ZBLOCK", case_3d) +@pytest.mark.parametrize("dtype, sigtype", types_all) +def test_advance_implicit_perm_load_store_3d_static_stride(X, Y, Z, XBLOCK, YBLOCK, ZBLOCK, dtype, sigtype): + """ + Real layout: A[X,Y,Z] row-major stride=(Y*Z, Z, 1) + Kernel view (logical): shape=(Z,Y,X), strides=(1, Z, Y*Z) + """ + A = test_common.generate_tensor(shape=(X, Y, Z), dtype=sigtype).npu() + _assert_row_major_3d(A, X, Y, Z) + out = torch.zeros_like(A) + + xnumel = Z + ynumel = Y + znumel = X + STRIDE_X = 1 + STRIDE_Y = Z + STRIDE_Z = Y * Z + + grid = ( + ceil_div(xnumel, XBLOCK), + ceil_div(ynumel, YBLOCK), + ceil_div(znumel, ZBLOCK), + ) + + advance_implicit_perm_load_store_3d_static_stride[grid]( + A, + out, + znumel, + ynumel, + xnumel, + STRIDE_Z=STRIDE_Z, + STRIDE_Y=STRIDE_Y, + STRIDE_X=STRIDE_X, + ZBLOCK=ZBLOCK, + YBLOCK=YBLOCK, + XBLOCK=XBLOCK, + ) + + torch.testing.assert_close(out, A) + + +if __name__ == "__main__": + test_addptr_implicit_perm_load_2d(*case_2d[0], *types_all[0]) + test_addptr_implicit_perm_store_2d(*case_2d[0], *types_all[0]) + test_addptr_implicit_perm_load_store_2d_static_stride(*case_2d[0], *types_all[0]) + test_addptr_implicit_perm_load_store_2d(*case_2d[0], *types_all[0]) + test_addptr_implicit_perm_load_store_3d_static_stride(*case_3d[0], *types_all[0]) + test_addptr_implicit_perm_load_store_3d(*case_3d[0], *types_all[0]) + test_addptr_implicit_perm_load_store_4d_static_stride(*case_4d[0], *types_all[0]) + test_addptr_implicit_perm_load_store_4d(*case_4d[0], *types_all[0]) + + test_make_tensor_ptr_implicit_perm_load_2d(*case_2d[0], *types_all[0]) + test_make_tensor_ptr_implicit_perm_store_2d(*case_2d[0], *types_all[0]) + test_make_tensor_ptr_implicit_perm_load_3d_static_stride(*case_3d[0], *types_all[0]) + test_make_tensor_ptr_implicit_perm_load_store_2d_static_stride(*case_2d[0], *types_all[0]) + test_make_tensor_ptr_implicit_perm_load_store_3d_static_stride(*case_3d[0], *types_all[0]) + test_make_tensor_ptr_implicit_perm_load_store_3d(*case_3d[0], *types_all[0]) + test_advance_implicit_perm_load_store_2d_static_stride(*case_2d[0], *types_all[0]) + test_advance_implicit_perm_load_store_3d_static_stride(*case_3d[0], *types_all[0]) diff --git a/third_party/ascend/unittest/pytest_ut/test_index_select.py b/third_party/ascend/unittest/pytest_ut/test_index_select.py new file mode 100644 index 0000000000..a6d4737e4c --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_index_select.py @@ -0,0 +1,268 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import triton.language.extra.cann.extension as extension + +import torch +import torch_npu +import pytest +import test_common + +# ============================================================================ +# PyTorch Reference Implementation +# ============================================================================ + + +def torch_index_select(x0, dim, indices): + """PyTorch reference implementation using torch.index_select.""" + return torch.index_select(x0, dim, indices) + + +# ============================================================================ +# Triton Kernel Implementations +# ============================================================================ + + +@triton.jit +def index_select_manual_kernel(in_ptr, indices_ptr, out_ptr, dim, g_stride: tl.constexpr, indice_length: tl.constexpr, + g_block: tl.constexpr, g_block_sub: tl.constexpr, other_block: tl.constexpr): + """ + Manual implementation using extension.get_element and extension.insert_slice. + + This is a baseline implementation without using the index_select_simd intrinsic. + """ + g_begin = tl.program_id(0) * g_block + for goffs in range(0, g_block, g_block_sub): + g_idx = tl.arange(0, g_block_sub) + g_begin + goffs + g_mask = g_idx < indice_length + indices = tl.load(indices_ptr + g_idx, g_mask, other=0) + + for other_offset in range(0, g_stride, other_block): + tmp_buf = tl.zeros((g_block_sub, other_block), in_ptr.dtype.element_ty) + other_idx = tl.arange(0, other_block) + other_offset + other_mask = other_idx < g_stride + + # Manual gather: iterate over each index + for i in range(0, g_block_sub): + gather_offset = extension.get_element(indices, (i, )) * g_stride + val = tl.load(in_ptr + gather_offset + other_idx, other_mask) + tmp_buf = extension.insert_slice(tmp_buf, val[None, :], offsets=(i, 0), sizes=(1, other_block), + strides=(1, 1)) + + tl.store(out_ptr + g_idx[:, None] * g_stride + other_idx[None, :], tmp_buf, + g_mask[:, None] & other_mask[None, :]) + + +@triton.jit +def index_select_extension_kernel(in_ptr, indices_ptr, out_ptr, dim: tl.constexpr, other_numel: tl.constexpr, + g_stride: tl.constexpr, indice_length: tl.constexpr, g_block: tl.constexpr, + g_block_sub: tl.constexpr, other_block: tl.constexpr): + """ + Implementation using extension.index_select_simd intrinsic. + + This uses the hardware-accelerated SIMD index_select operation. + """ + g_begin = tl.program_id(0) * g_block + for goffs in range(0, g_block, g_block_sub): + g_idx = tl.arange(0, g_block_sub) + g_begin + goffs + g_mask = g_idx < indice_length + indices = tl.load(indices_ptr + g_idx, g_mask, other=0) + + for other_offset in range(0, g_stride, other_block): + other_idx = tl.arange(0, other_block) + other_offset + other_mask = other_idx < g_stride + + # Use extension index_select_simd + tmp_buf = extension.index_select_simd(src=in_ptr, dim=dim, index=indices, src_shape=(other_numel, g_stride), + src_offset=(-1, 0), read_shape=(-1, other_block)) + + tl.store(out_ptr + g_idx[:, None] * g_stride + other_idx[None, :], tmp_buf, + g_mask[:, None] & other_mask[None, :]) + + +@triton.jit +def index_select_auto_kernel(in_ptr, indices_ptr, out_ptr, dim: tl.constexpr, other_numel: tl.constexpr, + g_stride: tl.constexpr, indice_length: tl.constexpr, g_block: tl.constexpr, + g_block_sub: tl.constexpr, other_block: tl.constexpr): + """ + Auto-lowering implementation using standard tl.load with computed offsets. + + This lets the compiler automatically handle the index_select pattern. + """ + g_begin = tl.program_id(0) * g_block + for goffs in range(0, g_block, g_block_sub): + g_idx = tl.arange(0, g_block_sub) + g_begin + goffs + g_mask = g_idx < indice_length + indices = tl.load(indices_ptr + g_idx, g_mask, other=0) + + for other_offset in range(0, g_stride, other_block): + other_idx = tl.arange(0, other_block) + other_offset + other_mask = other_idx < g_stride + + # Auto-lowering: compute offsets and use standard load + src_offsets = indices[:, None] * g_stride + other_idx[None, :] + tmp_buf = tl.load(in_ptr + src_offsets, g_mask[:, None] & other_mask[None, :]) + + tl.store(out_ptr + g_idx[:, None] * g_stride + other_idx[None, :], tmp_buf, + g_mask[:, None] & other_mask[None, :]) + + +# ============================================================================ +# Host Functions +# ============================================================================ + + +def triton_index_select(x0, dim, indices, impl='extension', num_vec_core=48): + """ + Triton implementation of index_select. + + Args: + x0: Source tensor + dim: Dimension to select from + indices: Indices to select + impl: Implementation to use ('manual', 'extension', or 'auto') + num_vec_core: Number of vector cores to use + + Returns: + Output tensor with selected indices + """ + sz = list(x0.shape) + sz[dim] = len(indices) + out = torch.empty(tuple(sz), dtype=x0.dtype).npu() + + g_stride = x0.stride(dim) + indice_length = indices.numel() + g_block = (indice_length - 1) // num_vec_core + 1 + + # Calculate UB space allocation + enable_multi_buffer = True + ub_size = 125 * 1024 // (2 if enable_multi_buffer else 1) + other_block = g_stride + g_block_sub = ub_size // ( + # max memory consumption: arith.select + other (mask handling in auto) + x0.element_size() * g_stride * 3 + indices.element_size()) + if g_block_sub < 1: + other_block = (ub_size - indices.element_size()) // x0.element_size() + g_block_sub = 1 + + # Select kernel based on implementation + if impl == 'manual': + kernel = index_select_manual_kernel + kernel[num_vec_core, 1, 1](x0, indices, out, dim, g_stride=g_stride, indice_length=indice_length, + g_block=g_block, g_block_sub=g_block_sub, other_block=other_block, multibuffer=False) + elif impl == 'extension': + kernel = index_select_extension_kernel + kernel[num_vec_core, 1, + 1](x0, indices, out, dim, other_numel=sz[0], g_stride=g_stride, indice_length=indice_length, + g_block=g_block, g_block_sub=g_block_sub, other_block=other_block) + elif impl == 'auto': + kernel = index_select_auto_kernel + kernel[num_vec_core, 1, + 1](x0, indices, out, dim, other_numel=sz[0], g_stride=g_stride, indice_length=indice_length, + g_block=g_block, g_block_sub=g_block_sub, other_block=other_block) + else: + raise ValueError(f"Unknown implementation: {impl}") + + return out + + +# ============================================================================ +# Test Parameters +# ============================================================================ + +# Test cases with various shapes and sizes +# Format: (src_shape, dim, indice_shape, dtype) +INDEX_SELECT_TEST_CASES = [ + # Small scale tests (< 100us) + ((10, 16), 0, (1024, ), "float32"), # 27us fp32 + ((20669, 8), 0, (2048, ), "float32"), # 31us fp32 + ((1034, 16), 0, (48000, ), "float32"), # 38us fp32 + ((270098, 32), 0, (512, ), "float32"), # 40us fp32 + ((1000000, 8), 0, (21329, ), "float32"), # 44us fp32 + + # Medium scale tests (100us - 2000us) + ((22717, 8), 0, (278400, ), "float32"), # 76us fp32 + ((2000000, 32), 0, (270244, ), "float32"), # 155us fp32 + ((500000, 37), 0, (324344, ), "float32"), # 266us +- 10us fp32 + ((270610, 32), 0, (2928000, ), "float32"), # 553us fp32 + ((3200, 16), 0, (1971940, ), "float32"), # 1104us fp32 + ((3200, 1, 37), 0, (1022226, ), "float32"), # 1965us +- 65us fp32 + + # Large scale tests (> 2000us) + ((500000, 240), 0, (375144, ), "float32"), # 1984us +- 5us fp32 48core + ((323357, 37), 0, (1022226, ), "float32"), # 3500us +- 50us fp32 + ((480000, 16), 0, (3943880, ), "float32"), # 3636 +- 42us fp32 + ((480000, 32), 0, (3943880, ), "float32"), # 3678 +- 5us fp32 + ((480000, 64), 0, (3943880, ), "float32"), # 5392 +- 20us fp32 + ((378103, 240), 0, (1992337, ), "float32"), # 10000us fp32 + ((374035, 240), 0, (1971940, ), "float32"), # 10100us fp32 +] + +# ============================================================================ +# Test Cases +# ============================================================================ + + +@pytest.mark.parametrize("src_shape, dim, indice_shape, dtype", INDEX_SELECT_TEST_CASES) +def test_index_select_manual(src_shape, dim, indice_shape, dtype): + """Test manual implementation using extension.get_element and extension.insert_slice.""" + x0 = test_common.generate_tensor(shape=src_shape, dtype=dtype).npu() + indices = torch.randint(0, src_shape[dim], size=indice_shape, dtype=torch.int32).npu() + + torch_ref = torch_index_select(x0, dim, indices) + triton_cal = triton_index_select(x0, dim, indices, impl='manual', num_vec_core=40) + + assert torch.equal(torch_ref, triton_cal), \ + f"Manual implementation failed for shape={src_shape}, dim={dim}, indices={indice_shape}" + + +@pytest.mark.parametrize("src_shape, dim, indice_shape, dtype", INDEX_SELECT_TEST_CASES) +def test_index_select_extension(src_shape, dim, indice_shape, dtype): + """Test extension.index_select_simd implementation.""" + x0 = test_common.generate_tensor(shape=src_shape, dtype=dtype).npu() + indices = torch.randint(0, src_shape[dim], size=indice_shape, dtype=torch.int32).npu() + + torch_ref = torch_index_select(x0, dim, indices) + triton_cal = triton_index_select(x0, dim, indices, impl='extension', num_vec_core=48) + + test_common.validate_cmp(dtype, triton_cal, torch_ref) + + +@pytest.mark.parametrize("src_shape, dim, indice_shape, dtype", INDEX_SELECT_TEST_CASES) +def test_index_select_auto(src_shape, dim, indice_shape, dtype): + """Test auto-lowering implementation using standard tl.load.""" + x0 = test_common.generate_tensor(shape=src_shape, dtype=dtype).npu() + indices = torch.randint(0, src_shape[dim], size=indice_shape, dtype=torch.int32).npu() + + torch_ref = torch_index_select(x0, dim, indices) + triton_cal = triton_index_select(x0, dim, indices, impl='auto', num_vec_core=48) + + test_common.validate_cmp(dtype, triton_cal, torch_ref) + + +# Quick smoke test +if __name__ == "__main__": + test_index_select_extension((500000, 37), 0, (324344, ), "float32") + print("extension implementation passed") + + test_index_select_auto((500000, 37), 0, (324344, ), "float32") + print("auto-lowering implementation passed") diff --git a/third_party/ascend/unittest/pytest_ut/test_index_select_inductor.py b/third_party/ascend/unittest/pytest_ut/test_index_select_inductor.py new file mode 100644 index 0000000000..b66367b6eb --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_index_select_inductor.py @@ -0,0 +1,91 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu + +import test_common + + +@triton.jit +def triton_unk_fused_embedding_0(in_ptr0, in_ptr1, out_ptr0, y0_numel, x1_numel, Y0BLOCK: tl.constexpr, + Y0BLOCK_SUB: tl.constexpr): + x1_numel = 64 + X1BLOCK_SUB: tl.constexpr = 64 + y0_offset = tl.program_id(0) * Y0BLOCK + base_y0 = tl.arange(0, Y0BLOCK_SUB) + loops_y0 = (Y0BLOCK + Y0BLOCK_SUB - 1) // Y0BLOCK_SUB + base_x1 = tl.arange(0, X1BLOCK_SUB) + for loop_y0 in range(loops_y0): + y0 = y0_offset + (loop_y0 * Y0BLOCK_SUB) + base_y0[:, None] + y0_mask = y0 < min(Y0BLOCK + y0_offset, y0_numel) + x1 = base_x1[None, :] + tmp0 = tl.load(in_ptr0 + (y0), y0_mask) + tmp6 = tl.load(in_ptr1 + (x1 + 64 * tmp0)) + tl.store(out_ptr0 + (x1 + 64 * y0), tmp6, y0_mask) + + +def torch_embedding_impl(in_ptr0, in_ptr1): + indices_flat = in_ptr0.flatten() # [6400] + vocab_size = 1353406 + out_flat = torch.embedding(in_ptr1, indices_flat) # [6400, 16] + out = out_flat.view(128, 50, 64) + return out + + +def test_kernel(): + arg0_1 = torch.randint(0, 1353406, (128, 50), device='npu', dtype=torch.int64) + arg1_1 = torch.randn((1353406, 64), device='npu', dtype=torch.float32) + buf0 = torch.empty((128, 50, 64), device='npu', dtype=torch.float32) + + y0_numel = 128 * 50 + x1_numel = 64 + Y0BLOCK = 256 + Y0BLOCK_SUB = 16 + + grid = (triton.cdiv(y0_numel, Y0BLOCK), ) + + print(f"Grid size: {grid}") + print(f"Total programs: {grid[0]}") + print(f"y0_numel: {y0_numel}, Y0BLOCK: {Y0BLOCK}") + + triton_unk_fused_embedding_0[grid]( + arg0_1, + arg1_1, + buf0, + y0_numel, + x1_numel, + Y0BLOCK=Y0BLOCK, + Y0BLOCK_SUB=Y0BLOCK_SUB, + ) + + print(f"Output shape: {buf0.shape}") + print(f"Output sample: {buf0[0, 0, :]}") + + expected = torch_embedding_impl(arg0_1, arg1_1) + + test_common.validate_cmp("float32", buf0, expected) + + +if __name__ == "__main__": + test_kernel() diff --git a/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_add.py b/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_add.py new file mode 100644 index 0000000000..fddfdbd285 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_add.py @@ -0,0 +1,520 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +# ============================================================================= +# Indirect atomic_add test suite +# +# Test matrix (access pattern x tensor rank): +# +# | access pattern | 1-D | 2-D | 3-D | 4-D | 5-D | +# |-------------------------------------------------|------|------|------|------|------| +# | structured pointer + discrete mask | (A1) | (A2) | (A3) | (A4) | (A5) | +# | partial structured: high-dim disc + low-dim cont| - | (B2) | (B3) | (B4) | (B5) | +# | fully unstructured indirect offsets | (C1) | (C2) | (C3) | (C4) | (C5) | +# +# Notes: +# 1. Case A exercises the structured-pointer discrete-mask atomic_add path. +# 2. Case B exercises a single high-dimension discrete remap with the +# remaining lower dimensions kept contiguous. +# 3. Case C exercises fully unstructured indirect offsets. +# 4. All cases validate both the final destination tensor and the atomic_add +# return value, which must be the old value observed at each access. +# ============================================================================= + +import math + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +SUPPORTED_DTYPES = [ + ("int32", torch.int32), + ("int64", torch.int64), + ("float16", torch.float16), + ("float32", torch.float32), + ("bfloat16", torch.bfloat16), +] + +RANK_SHAPES = { + 1: (8, ), + 2: (4, 4), + 3: (2, 3, 4), + 4: (2, 2, 3, 4), + 5: (2, 2, 2, 3, 4), +} + +PARTIAL_STRUCTURED_SHAPES = { + 2: (2, 16), + 3: (2, 2, 8), + 4: (2, 2, 2, 8), + 5: (2, 2, 2, 2, 8), +} + +TEST_RANKS = [3] +TEST_DTYPE = [("int32", torch.int32), ("bfloat16", torch.bfloat16)] + + +@triton.jit +def structured_disc_mask_atomic_add_1d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, +): + i0 = tl.arange(0, D0) + linear = i0 + mask = (linear * 2) < D0 + value = tl.load(val_ptr + linear) + old = tl.atomic_add(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_add_2d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + i1 = tl.arange(0, D1)[None, :] + linear = i0 * D1 + i1 + mask = (linear * 2) < (D0 * D1) + value = tl.load(val_ptr + linear) + old = tl.atomic_add(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_add_3d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1) * D2 + i2 + mask = (linear * 2) < (D0 * D1 * D2) + value = tl.load(val_ptr + linear) + old = tl.atomic_add(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_add_4d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = ((i0 * D1 + i1) * D2 + i2) * D3 + i3 + mask = (linear * 2) < (D0 * D1 * D2 * D3) + value = tl.load(val_ptr + linear) + old = tl.atomic_add(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_add_5d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None, None] + i2 = tl.arange(0, D2)[None, None, :, None, None] + i3 = tl.arange(0, D3)[None, None, None, :, None] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2) * D3 + i3) * D4 + i4 + mask = (linear * 2) < (D0 * D1 * D2 * D3 * D4) + value = tl.load(val_ptr + linear) + old = tl.atomic_add(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def fully_unstructured_atomic_add_1d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, +): + i0 = tl.arange(0, D0) + linear = i0 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_add(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_add_2d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + i1 = tl.arange(0, D1)[None, :] + linear = i0 * D1 + i1 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_add(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_add_3d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1) * D2 + i2 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_add(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_add_4d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = ((i0 * D1 + i1) * D2 + i2) * D3 + i3 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_add(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_add_5d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None, None] + i2 = tl.arange(0, D2)[None, None, :, None, None] + i3 = tl.arange(0, D3)[None, None, None, :, None] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2) * D3 + i3) * D4 + i4 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_add(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_add_2d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + remapped_i0 = (i0 + 1) % D0 + inner_idx = tl.arange(0, D1)[None, :] + linear = i0 * D1 + inner_idx + value = tl.load(val_ptr + linear) + ptr = out_ptr + remapped_i0 * D1 + inner_idx + old = tl.atomic_add(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_add_3d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1)[:, :, None] * D2 + i2 + value = tl.load(val_ptr + linear) + ptr = out_ptr + ((remapped_i0 * D1 + i1)[:, :, None] * D2 + i2) + old = tl.atomic_add(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_add_4d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2)[:, :, :, None] * D3 + i3) + value = tl.load(val_ptr + linear) + ptr = out_ptr + ((((remapped_i0 * D1 + i1) * D2 + i2)[:, :, :, None] * D3) + i3) + old = tl.atomic_add(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_add_5d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = ((((i0 * D1 + i1) * D2 + i2) * D3 + i3)[:, :, :, :, None] * D4 + i4) + value = tl.load(val_ptr + linear) + ptr = out_ptr + (((((remapped_i0 * D1 + i1) * D2 + i2) * D3 + i3)[:, :, :, :, None] * D4) + i4) + old = tl.atomic_add(ptr, value) + tl.store(old_ptr + linear, old) + + +STRUCTURED_DISC_MASK_KERNELS = { + 1: structured_disc_mask_atomic_add_1d, + 2: structured_disc_mask_atomic_add_2d, + 3: structured_disc_mask_atomic_add_3d, + 4: structured_disc_mask_atomic_add_4d, + 5: structured_disc_mask_atomic_add_5d, +} + +PARTIAL_STRUCTURED_KERNELS = { + 2: partial_structured_atomic_add_2d, + 3: partial_structured_atomic_add_3d, + 4: partial_structured_atomic_add_4d, + 5: partial_structured_atomic_add_5d, +} + +FULLY_UNSTRUCTURED_KERNELS = { + 1: fully_unstructured_atomic_add_1d, + 2: fully_unstructured_atomic_add_2d, + 3: fully_unstructured_atomic_add_3d, + 4: fully_unstructured_atomic_add_4d, + 5: fully_unstructured_atomic_add_5d, +} + + +def _build_partial_structured_offsets(shape): + assert len(shape) >= 2 + first_dim = shape[0] + tail_shape = shape[1:] + tail_numel = math.prod(tail_shape) + remapped_first = ((torch.arange(first_dim, dtype=torch.int64) + 1) % first_dim) + first_offsets = remapped_first.reshape((first_dim, ) + (1, ) * len(tail_shape)) + tail_offsets = torch.arange(tail_numel, dtype=torch.int64).reshape((1, ) + tail_shape) + return first_offsets * tail_numel + tail_offsets, math.prod(shape) + + +def _build_fully_unstructured_offsets(shape): + numel = math.prod(shape) + output_numel = numel * 2 + base = torch.arange(numel, dtype=torch.int64) + offsets = (base % 2) * numel + (base // 2) + offsets = offsets.reshape(shape) + return offsets, output_numel + + +def _build_discrete_mask(shape): + numel = math.prod(shape) + return ((torch.arange(numel, dtype=torch.int64).reshape(shape) * 2) < numel) + + +def _build_value_tensor(shape, dtype): + return torch.ones(shape, dtype=dtype) + + +def _build_output_baseline(output_numel, dtype): + return torch.arange(output_numel, dtype=dtype) + + +def _simulate_atomic_add(base_output, offsets, values, mask=None): + expected_output = base_output.reshape(-1).clone().cpu() + flat_offsets = offsets.reshape(-1).to(torch.int64).cpu() + flat_values = values.reshape(-1).cpu() + flat_old = torch.zeros(flat_values.shape, dtype=base_output.dtype) + if mask is None: + flat_mask = torch.ones(flat_offsets.shape, dtype=torch.bool) + else: + flat_mask = mask.reshape(-1).cpu() + + for idx in range(flat_offsets.numel()): + if not bool(flat_mask[idx]): + continue + offset = int(flat_offsets[idx].item()) + flat_old[idx] = expected_output[offset] + expected_output[offset] = expected_output[offset] + flat_values[idx].to(expected_output.dtype) + + return expected_output, flat_old.reshape(values.shape) + + +def _assert_equal(actual, expected, dtype_name, rank, scenario): + actual_cpu = actual.cpu() + expected_cpu = expected.cpu() + assert torch.equal(actual_cpu, expected_cpu), (f"scenario={scenario}, dtype={dtype_name}, rank={rank}\n" + f"Expected:\n{expected_cpu}\nGot:\n{actual_cpu}") + + +def _launch_structured_discrete_mask(rank, values, output, old, shape): + kernel = STRUCTURED_DISC_MASK_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](values, output, old, **kwargs) + + +def _launch_partial_structured(rank, values, output, old, shape): + kernel = PARTIAL_STRUCTURED_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](values, output, old, **kwargs) + + +def _launch_fully_unstructured(rank, offsets, values, output, old, shape): + kernel = FULLY_UNSTRUCTURED_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](offsets, values, output, old, **kwargs) + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_add_structured_pointer_with_discrete_mask(dtype_name, torch_dtype, rank): + shape = RANK_SHAPES[rank] + values = _build_value_tensor(shape, torch_dtype).npu() + output_numel = math.prod(shape) + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_structured_discrete_mask(rank, values, output, old, shape) + + identity_offsets = torch.arange(output_numel, dtype=torch.int64).reshape(shape) + mask = _build_discrete_mask(shape) + expected_output, expected_old = _simulate_atomic_add( + baseline, + identity_offsets, + _build_value_tensor(shape, torch_dtype), + mask=mask, + ) + _assert_equal(output, expected_output, dtype_name, rank, "structured+discrete-mask/output") + _assert_equal(old, expected_old, dtype_name, rank, "structured+discrete-mask/old") + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_add_partially_structured_indirect_offsets(dtype_name, torch_dtype, rank): + shape = PARTIAL_STRUCTURED_SHAPES[rank] + offsets, output_numel = _build_partial_structured_offsets(shape) + values = _build_value_tensor(shape, torch_dtype).npu() + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_partial_structured(rank, values, output, old, shape) + + expected_output, expected_old = _simulate_atomic_add( + baseline, + offsets, + _build_value_tensor(shape, torch_dtype), + ) + _assert_equal(output, expected_output, dtype_name, rank, "partially-structured-indirect/output") + _assert_equal(old, expected_old, dtype_name, rank, "partially-structured-indirect/old") + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_add_fully_unstructured_indirect_offsets(dtype_name, torch_dtype, rank): + shape = RANK_SHAPES[rank] + offsets, output_numel = _build_fully_unstructured_offsets(shape) + values = _build_value_tensor(shape, torch_dtype).npu() + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_fully_unstructured(rank, offsets.npu(), values, output, old, shape) + + expected_output, expected_old = _simulate_atomic_add( + baseline, + offsets, + _build_value_tensor(shape, torch_dtype), + ) + _assert_equal(output, expected_output, dtype_name, rank, "fully-unstructured-indirect/output") + _assert_equal(old, expected_old, dtype_name, rank, "fully-unstructured-indirect/old") diff --git a/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_and.py b/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_and.py new file mode 100644 index 0000000000..51b2713eb7 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_and.py @@ -0,0 +1,517 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +# ============================================================================= +# Indirect atomic_and test suite +# +# Test matrix (access pattern x tensor rank): +# +# | access pattern | 1-D | 2-D | 3-D | 4-D | 5-D | +# |-------------------------------------------------|------|------|------|------|------| +# | structured pointer + discrete mask | (A1) | (A2) | (A3) | (A4) | (A5) | +# | partial structured: high-dim disc + low-dim cont| - | (B2) | (B3) | (B4) | (B5) | +# | fully unstructured indirect offsets | (C1) | (C2) | (C3) | (C4) | (C5) | +# +# Notes: +# 1. Case A exercises the structured-pointer discrete-mask atomic_and path. +# 2. Case B exercises a single high-dimension discrete remap with the +# remaining lower dimensions kept contiguous. +# 3. Case C exercises fully unstructured indirect offsets. +# 4. All cases validate both the final destination tensor and the atomic_and +# return value, which must be the old value observed at each access. +# ============================================================================= + +import math + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +SUPPORTED_DTYPES = [ + ("int32", torch.int32), + ("int64", torch.int64), +] + +RANK_SHAPES = { + 1: (8, ), + 2: (4, 4), + 3: (2, 3, 4), + 4: (2, 2, 3, 4), + 5: (2, 2, 2, 3, 4), +} + +PARTIAL_STRUCTURED_SHAPES = { + 2: (2, 16), + 3: (2, 2, 8), + 4: (2, 2, 2, 8), + 5: (2, 2, 2, 2, 8), +} + +TEST_RANKS = [3] +TEST_DTYPE = [("int32", torch.int32)] + + +@triton.jit +def structured_disc_mask_atomic_and_1d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, +): + i0 = tl.arange(0, D0) + linear = i0 + mask = (linear * 2) < D0 + value = tl.load(val_ptr + linear) + old = tl.atomic_and(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_and_2d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + i1 = tl.arange(0, D1)[None, :] + linear = i0 * D1 + i1 + mask = (linear * 2) < (D0 * D1) + value = tl.load(val_ptr + linear) + old = tl.atomic_and(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_and_3d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1) * D2 + i2 + mask = (linear * 2) < (D0 * D1 * D2) + value = tl.load(val_ptr + linear) + old = tl.atomic_and(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_and_4d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = ((i0 * D1 + i1) * D2 + i2) * D3 + i3 + mask = (linear * 2) < (D0 * D1 * D2 * D3) + value = tl.load(val_ptr + linear) + old = tl.atomic_and(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_and_5d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None, None] + i2 = tl.arange(0, D2)[None, None, :, None, None] + i3 = tl.arange(0, D3)[None, None, None, :, None] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2) * D3 + i3) * D4 + i4 + mask = (linear * 2) < (D0 * D1 * D2 * D3 * D4) + value = tl.load(val_ptr + linear) + old = tl.atomic_and(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def fully_unstructured_atomic_and_1d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, +): + i0 = tl.arange(0, D0) + linear = i0 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_and(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_and_2d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + i1 = tl.arange(0, D1)[None, :] + linear = i0 * D1 + i1 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_and(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_and_3d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1) * D2 + i2 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_and(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_and_4d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = ((i0 * D1 + i1) * D2 + i2) * D3 + i3 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_and(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_and_5d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None, None] + i2 = tl.arange(0, D2)[None, None, :, None, None] + i3 = tl.arange(0, D3)[None, None, None, :, None] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2) * D3 + i3) * D4 + i4 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_and(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_and_2d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + remapped_i0 = (i0 + 1) % D0 + inner_idx = tl.arange(0, D1)[None, :] + linear = i0 * D1 + inner_idx + value = tl.load(val_ptr + linear) + ptr = out_ptr + remapped_i0 * D1 + inner_idx + old = tl.atomic_and(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_and_3d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1)[:, :, None] * D2 + i2 + value = tl.load(val_ptr + linear) + ptr = out_ptr + ((remapped_i0 * D1 + i1)[:, :, None] * D2 + i2) + old = tl.atomic_and(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_and_4d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2)[:, :, :, None] * D3 + i3) + value = tl.load(val_ptr + linear) + ptr = out_ptr + ((((remapped_i0 * D1 + i1) * D2 + i2)[:, :, :, None] * D3) + i3) + old = tl.atomic_and(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_and_5d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = ((((i0 * D1 + i1) * D2 + i2) * D3 + i3)[:, :, :, :, None] * D4 + i4) + value = tl.load(val_ptr + linear) + ptr = out_ptr + (((((remapped_i0 * D1 + i1) * D2 + i2) * D3 + i3)[:, :, :, :, None] * D4) + i4) + old = tl.atomic_and(ptr, value) + tl.store(old_ptr + linear, old) + + +STRUCTURED_DISC_MASK_KERNELS = { + 1: structured_disc_mask_atomic_and_1d, + 2: structured_disc_mask_atomic_and_2d, + 3: structured_disc_mask_atomic_and_3d, + 4: structured_disc_mask_atomic_and_4d, + 5: structured_disc_mask_atomic_and_5d, +} + +PARTIAL_STRUCTURED_KERNELS = { + 2: partial_structured_atomic_and_2d, + 3: partial_structured_atomic_and_3d, + 4: partial_structured_atomic_and_4d, + 5: partial_structured_atomic_and_5d, +} + +FULLY_UNSTRUCTURED_KERNELS = { + 1: fully_unstructured_atomic_and_1d, + 2: fully_unstructured_atomic_and_2d, + 3: fully_unstructured_atomic_and_3d, + 4: fully_unstructured_atomic_and_4d, + 5: fully_unstructured_atomic_and_5d, +} + + +def _build_partial_structured_offsets(shape): + assert len(shape) >= 2 + first_dim = shape[0] + tail_shape = shape[1:] + tail_numel = math.prod(tail_shape) + remapped_first = ((torch.arange(first_dim, dtype=torch.int64) + 1) % first_dim) + first_offsets = remapped_first.reshape((first_dim, ) + (1, ) * len(tail_shape)) + tail_offsets = torch.arange(tail_numel, dtype=torch.int64).reshape((1, ) + tail_shape) + return first_offsets * tail_numel + tail_offsets, math.prod(shape) + + +def _build_fully_unstructured_offsets(shape): + numel = math.prod(shape) + output_numel = numel * 2 + base = torch.arange(numel, dtype=torch.int64) + offsets = (base % 2) * numel + (base // 2) + offsets = offsets.reshape(shape) + return offsets, output_numel + + +def _build_discrete_mask(shape): + numel = math.prod(shape) + return ((torch.arange(numel, dtype=torch.int64).reshape(shape) * 2) < numel) + + +def _build_value_tensor(shape, dtype): + return torch.ones(shape, dtype=dtype) + + +def _build_output_baseline(output_numel, dtype): + return torch.arange(output_numel, dtype=dtype) + + +def _simulate_atomic_and(base_output, offsets, values, mask=None): + expected_output = base_output.reshape(-1).clone().cpu() + flat_offsets = offsets.reshape(-1).to(torch.int64).cpu() + flat_values = values.reshape(-1).cpu() + flat_old = torch.zeros(flat_values.shape, dtype=base_output.dtype) + if mask is None: + flat_mask = torch.ones(flat_offsets.shape, dtype=torch.bool) + else: + flat_mask = mask.reshape(-1).cpu() + + for idx in range(flat_offsets.numel()): + if not bool(flat_mask[idx]): + continue + offset = int(flat_offsets[idx].item()) + flat_old[idx] = expected_output[offset] + expected_output[offset] = expected_output[offset] & flat_values[idx].to(expected_output.dtype) + + return expected_output, flat_old.reshape(values.shape) + + +def _assert_equal(actual, expected, dtype_name, rank, scenario): + actual_cpu = actual.cpu() + expected_cpu = expected.cpu() + assert torch.equal(actual_cpu, expected_cpu), (f"scenario={scenario}, dtype={dtype_name}, rank={rank}\n" + f"Expected:\n{expected_cpu}\nGot:\n{actual_cpu}") + + +def _launch_structured_discrete_mask(rank, values, output, old, shape): + kernel = STRUCTURED_DISC_MASK_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](values, output, old, **kwargs) + + +def _launch_partial_structured(rank, values, output, old, shape): + kernel = PARTIAL_STRUCTURED_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](values, output, old, **kwargs) + + +def _launch_fully_unstructured(rank, offsets, values, output, old, shape): + kernel = FULLY_UNSTRUCTURED_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](offsets, values, output, old, **kwargs) + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_and_structured_pointer_with_discrete_mask(dtype_name, torch_dtype, rank): + shape = RANK_SHAPES[rank] + values = _build_value_tensor(shape, torch_dtype).npu() + output_numel = math.prod(shape) + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_structured_discrete_mask(rank, values, output, old, shape) + + identity_offsets = torch.arange(output_numel, dtype=torch.int64).reshape(shape) + mask = _build_discrete_mask(shape) + expected_output, expected_old = _simulate_atomic_and( + baseline, + identity_offsets, + _build_value_tensor(shape, torch_dtype), + mask=mask, + ) + _assert_equal(output, expected_output, dtype_name, rank, "structured+discrete-mask/output") + _assert_equal(old, expected_old, dtype_name, rank, "structured+discrete-mask/old") + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_and_partially_structured_indirect_offsets(dtype_name, torch_dtype, rank): + shape = PARTIAL_STRUCTURED_SHAPES[rank] + offsets, output_numel = _build_partial_structured_offsets(shape) + values = _build_value_tensor(shape, torch_dtype).npu() + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_partial_structured(rank, values, output, old, shape) + + expected_output, expected_old = _simulate_atomic_and( + baseline, + offsets, + _build_value_tensor(shape, torch_dtype), + ) + _assert_equal(output, expected_output, dtype_name, rank, "partially-structured-indirect/output") + _assert_equal(old, expected_old, dtype_name, rank, "partially-structured-indirect/old") + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_and_fully_unstructured_indirect_offsets(dtype_name, torch_dtype, rank): + shape = RANK_SHAPES[rank] + offsets, output_numel = _build_fully_unstructured_offsets(shape) + values = _build_value_tensor(shape, torch_dtype).npu() + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_fully_unstructured(rank, offsets.npu(), values, output, old, shape) + + expected_output, expected_old = _simulate_atomic_and( + baseline, + offsets, + _build_value_tensor(shape, torch_dtype), + ) + _assert_equal(output, expected_output, dtype_name, rank, "fully-unstructured-indirect/output") + _assert_equal(old, expected_old, dtype_name, rank, "fully-unstructured-indirect/old") diff --git a/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_cas.py b/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_cas.py new file mode 100644 index 0000000000..3d7812471b --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_cas.py @@ -0,0 +1,407 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +# ============================================================================= +# Indirect atomic_cas test suite +# +# Test matrix (access pattern x tensor rank): +# +# | access pattern | 1-D | 2-D | 3-D | 4-D | 5-D | +# |-------------------------------------------------|------|------|------|------|------| +# | partial structured: high-dim disc + low-dim cont| - | (B2) | (B3) | (B4) | (B5) | +# | fully unstructured indirect offsets | (C1) | (C2) | (C3) | (C4) | (C5) | +# +# Notes: +# 1. Case B exercises a single high-dimension discrete remap with the +# remaining lower dimensions kept contiguous. +# 2. Case C exercises fully unstructured indirect offsets. +# 3. All cases validate both the final destination tensor and the atomic_cas +# return value, which must be the old value observed at each access. +# ============================================================================= + +import math + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +SUPPORTED_DTYPES = [ + ("int32", torch.int32), + ("int64", torch.int64), + ("float16", torch.float16), + ("float32", torch.float32), + ("bfloat16", torch.bfloat16), +] + +RANK_SHAPES = { + 1: (8, ), + 2: (4, 4), + 3: (2, 3, 4), + 4: (2, 2, 3, 4), + 5: (2, 2, 2, 3, 4), +} + +PARTIAL_STRUCTURED_SHAPES = { + 2: (2, 16), + 3: (2, 2, 8), + 4: (2, 2, 2, 8), + 5: (2, 2, 2, 2, 8), +} + +TEST_RANKS = [3] +TEST_DTYPE = [("int32", torch.int32), ("bfloat16", torch.bfloat16)] + + +@triton.jit +def fully_unstructured_atomic_cas_1d( + idx_ptr, + cmp_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, +): + i0 = tl.arange(0, D0) + linear = i0 + offset = tl.load(idx_ptr + linear) + compare = tl.load(cmp_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_cas(out_ptr + offset, compare, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_cas_2d( + idx_ptr, + cmp_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + i1 = tl.arange(0, D1)[None, :] + linear = i0 * D1 + i1 + offset = tl.load(idx_ptr + linear) + compare = tl.load(cmp_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_cas(out_ptr + offset, compare, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_cas_3d( + idx_ptr, + cmp_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1) * D2 + i2 + offset = tl.load(idx_ptr + linear) + compare = tl.load(cmp_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_cas(out_ptr + offset, compare, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_cas_4d( + idx_ptr, + cmp_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = ((i0 * D1 + i1) * D2 + i2) * D3 + i3 + offset = tl.load(idx_ptr + linear) + compare = tl.load(cmp_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_cas(out_ptr + offset, compare, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_cas_5d( + idx_ptr, + cmp_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None, None] + i2 = tl.arange(0, D2)[None, None, :, None, None] + i3 = tl.arange(0, D3)[None, None, None, :, None] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2) * D3 + i3) * D4 + i4 + offset = tl.load(idx_ptr + linear) + compare = tl.load(cmp_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_cas(out_ptr + offset, compare, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_cas_2d( + cmp_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :] + linear = i0 * D1 + i1 + compare = tl.load(cmp_ptr + linear) + value = tl.load(val_ptr + linear) + ptr = out_ptr + remapped_i0 * D1 + i1 + old = tl.atomic_cas(ptr, compare, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_cas_3d( + cmp_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1)[:, :, None] * D2 + i2 + compare = tl.load(cmp_ptr + linear) + value = tl.load(val_ptr + linear) + ptr = out_ptr + ((remapped_i0 * D1 + i1)[:, :, None] * D2 + i2) + old = tl.atomic_cas(ptr, compare, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_cas_4d( + cmp_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2)[:, :, :, None] * D3 + i3) + compare = tl.load(cmp_ptr + linear) + value = tl.load(val_ptr + linear) + ptr = out_ptr + ((((remapped_i0 * D1 + i1) * D2 + i2)[:, :, :, None] * D3) + i3) + old = tl.atomic_cas(ptr, compare, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_cas_5d( + cmp_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = ((((i0 * D1 + i1) * D2 + i2) * D3 + i3)[:, :, :, :, None] * D4 + i4) + compare = tl.load(cmp_ptr + linear) + value = tl.load(val_ptr + linear) + ptr = out_ptr + (((((remapped_i0 * D1 + i1) * D2 + i2) * D3 + i3)[:, :, :, :, None] * D4) + i4) + old = tl.atomic_cas(ptr, compare, value) + tl.store(old_ptr + linear, old) + + +PARTIAL_STRUCTURED_KERNELS = { + 2: partial_structured_atomic_cas_2d, + 3: partial_structured_atomic_cas_3d, + 4: partial_structured_atomic_cas_4d, + 5: partial_structured_atomic_cas_5d, +} + +FULLY_UNSTRUCTURED_KERNELS = { + 1: fully_unstructured_atomic_cas_1d, + 2: fully_unstructured_atomic_cas_2d, + 3: fully_unstructured_atomic_cas_3d, + 4: fully_unstructured_atomic_cas_4d, + 5: fully_unstructured_atomic_cas_5d, +} + + +def _build_partial_structured_offsets(shape): + assert len(shape) >= 2 + first_dim = shape[0] + tail_shape = shape[1:] + tail_numel = math.prod(tail_shape) + remapped_first = ((torch.arange(first_dim, dtype=torch.int64) + 1) % first_dim) + first_offsets = remapped_first.reshape((first_dim, ) + (1, ) * len(tail_shape)) + tail_offsets = torch.arange(tail_numel, dtype=torch.int64).reshape((1, ) + tail_shape) + return first_offsets * tail_numel + tail_offsets, math.prod(shape) + + +def _build_fully_unstructured_offsets(shape): + numel = math.prod(shape) + output_numel = numel * 2 + base = torch.arange(numel, dtype=torch.int64) + offsets = (base % 2) * numel + (base // 2) + offsets = offsets.reshape(shape) + return offsets, output_numel + + +def _build_value_tensor(shape, dtype): + return torch.full(shape, 7, dtype=dtype) + + +def _build_output_baseline(output_numel, dtype): + return torch.arange(output_numel, dtype=dtype) + + +def _build_compare_tensor(offsets, baseline, dtype): + flat_offsets = offsets.reshape(-1).to(torch.int64) + base_flat = baseline.reshape(-1) + compare = torch.empty(flat_offsets.numel(), dtype=dtype) + for idx, offset in enumerate(flat_offsets.tolist()): + compare[idx] = base_flat[offset] if idx % 2 == 0 else base_flat[offset] + 1 + return compare.reshape(offsets.shape) + + +def _simulate_atomic_cas(base_output, offsets, compare, values): + expected_output = base_output.reshape(-1).clone().cpu() + flat_offsets = offsets.reshape(-1).to(torch.int64).cpu() + flat_compare = compare.reshape(-1).cpu() + flat_values = values.reshape(-1).cpu() + flat_old = torch.zeros(flat_values.shape, dtype=base_output.dtype) + + for idx in range(flat_offsets.numel()): + offset = int(flat_offsets[idx].item()) + flat_old[idx] = expected_output[offset] + if expected_output[offset] == flat_compare[idx].to(expected_output.dtype): + expected_output[offset] = flat_values[idx].to(expected_output.dtype) + + return expected_output, flat_old.reshape(values.shape) + + +def _assert_equal(actual, expected, dtype_name, rank, scenario): + actual_cpu = actual.cpu() + expected_cpu = expected.cpu() + assert torch.equal(actual_cpu, expected_cpu), (f"scenario={scenario}, dtype={dtype_name}, rank={rank}\n" + f"Expected:\n{expected_cpu}\nGot:\n{actual_cpu}") + + +def _launch_partial_structured(rank, compare, values, output, old, shape): + kernel = PARTIAL_STRUCTURED_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](compare, values, output, old, **kwargs) + + +def _launch_fully_unstructured(rank, offsets, compare, values, output, old, shape): + kernel = FULLY_UNSTRUCTURED_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](offsets, compare, values, output, old, **kwargs) + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_cas_partially_structured_indirect_offsets(dtype_name, torch_dtype, rank): + shape = PARTIAL_STRUCTURED_SHAPES[rank] + offsets, output_numel = _build_partial_structured_offsets(shape) + baseline = _build_output_baseline(output_numel, torch_dtype) + compare = _build_compare_tensor(offsets, baseline, torch_dtype) + values = _build_value_tensor(shape, torch_dtype).npu() + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_partial_structured(rank, compare.npu(), values, output, old, shape) + + expected_output, expected_old = _simulate_atomic_cas( + baseline, + offsets, + compare, + _build_value_tensor(shape, torch_dtype), + ) + _assert_equal(output, expected_output, dtype_name, rank, "partially-structured-indirect/output") + _assert_equal(old, expected_old, dtype_name, rank, "partially-structured-indirect/old") + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_cas_fully_unstructured_indirect_offsets(dtype_name, torch_dtype, rank): + shape = RANK_SHAPES[rank] + offsets, output_numel = _build_fully_unstructured_offsets(shape) + baseline = _build_output_baseline(output_numel, torch_dtype) + compare = _build_compare_tensor(offsets, baseline, torch_dtype) + values = _build_value_tensor(shape, torch_dtype).npu() + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_fully_unstructured(rank, offsets.npu(), compare.npu(), values, output, old, shape) + + expected_output, expected_old = _simulate_atomic_cas( + baseline, + offsets, + compare, + _build_value_tensor(shape, torch_dtype), + ) + _assert_equal(output, expected_output, dtype_name, rank, "fully-unstructured-indirect/output") + _assert_equal(old, expected_old, dtype_name, rank, "fully-unstructured-indirect/old") diff --git a/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_max.py b/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_max.py new file mode 100644 index 0000000000..62dbd10a09 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_max.py @@ -0,0 +1,520 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +# ============================================================================= +# Indirect atomic_max test suite +# +# Test matrix (access pattern x tensor rank): +# +# | access pattern | 1-D | 2-D | 3-D | 4-D | 5-D | +# |-------------------------------------------------|------|------|------|------|------| +# | structured pointer + discrete mask | (A1) | (A2) | (A3) | (A4) | (A5) | +# | partial structured: high-dim disc + low-dim cont| - | (B2) | (B3) | (B4) | (B5) | +# | fully unstructured indirect offsets | (C1) | (C2) | (C3) | (C4) | (C5) | +# +# Notes: +# 1. Case A exercises the structured-pointer discrete-mask atomic_max path. +# 2. Case B exercises a single high-dimension discrete remap with the +# remaining lower dimensions kept contiguous. +# 3. Case C exercises fully unstructured indirect offsets. +# 4. All cases validate both the final destination tensor and the atomic_max +# return value, which must be the old value observed at each access. +# ============================================================================= + +import math + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +SUPPORTED_DTYPES = [ + ("int32", torch.int32), + ("int64", torch.int64), + ("float16", torch.float16), + ("float32", torch.float32), + ("bfloat16", torch.bfloat16), +] + +RANK_SHAPES = { + 1: (8, ), + 2: (4, 4), + 3: (2, 3, 4), + 4: (2, 2, 3, 4), + 5: (2, 2, 2, 3, 4), +} + +PARTIAL_STRUCTURED_SHAPES = { + 2: (2, 16), + 3: (2, 2, 8), + 4: (2, 2, 2, 8), + 5: (2, 2, 2, 2, 8), +} + +TEST_RANKS = [3] +TEST_DTYPE = [("int32", torch.int32), ("bfloat16", torch.bfloat16)] + + +@triton.jit +def structured_disc_mask_atomic_max_1d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, +): + i0 = tl.arange(0, D0) + linear = i0 + mask = (linear * 2) < D0 + value = tl.load(val_ptr + linear) + old = tl.atomic_max(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_max_2d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + i1 = tl.arange(0, D1)[None, :] + linear = i0 * D1 + i1 + mask = (linear * 2) < (D0 * D1) + value = tl.load(val_ptr + linear) + old = tl.atomic_max(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_max_3d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1) * D2 + i2 + mask = (linear * 2) < (D0 * D1 * D2) + value = tl.load(val_ptr + linear) + old = tl.atomic_max(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_max_4d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = ((i0 * D1 + i1) * D2 + i2) * D3 + i3 + mask = (linear * 2) < (D0 * D1 * D2 * D3) + value = tl.load(val_ptr + linear) + old = tl.atomic_max(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_max_5d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None, None] + i2 = tl.arange(0, D2)[None, None, :, None, None] + i3 = tl.arange(0, D3)[None, None, None, :, None] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2) * D3 + i3) * D4 + i4 + mask = (linear * 2) < (D0 * D1 * D2 * D3 * D4) + value = tl.load(val_ptr + linear) + old = tl.atomic_max(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def fully_unstructured_atomic_max_1d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, +): + i0 = tl.arange(0, D0) + linear = i0 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_max(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_max_2d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + i1 = tl.arange(0, D1)[None, :] + linear = i0 * D1 + i1 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_max(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_max_3d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1) * D2 + i2 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_max(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_max_4d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = ((i0 * D1 + i1) * D2 + i2) * D3 + i3 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_max(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_max_5d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None, None] + i2 = tl.arange(0, D2)[None, None, :, None, None] + i3 = tl.arange(0, D3)[None, None, None, :, None] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2) * D3 + i3) * D4 + i4 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_max(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_max_2d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + remapped_i0 = (i0 + 1) % D0 + inner_idx = tl.arange(0, D1)[None, :] + linear = i0 * D1 + inner_idx + value = tl.load(val_ptr + linear) + ptr = out_ptr + remapped_i0 * D1 + inner_idx + old = tl.atomic_max(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_max_3d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1)[:, :, None] * D2 + i2 + value = tl.load(val_ptr + linear) + ptr = out_ptr + ((remapped_i0 * D1 + i1)[:, :, None] * D2 + i2) + old = tl.atomic_max(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_max_4d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2)[:, :, :, None] * D3 + i3) + value = tl.load(val_ptr + linear) + ptr = out_ptr + ((((remapped_i0 * D1 + i1) * D2 + i2)[:, :, :, None] * D3) + i3) + old = tl.atomic_max(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_max_5d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = ((((i0 * D1 + i1) * D2 + i2) * D3 + i3)[:, :, :, :, None] * D4 + i4) + value = tl.load(val_ptr + linear) + ptr = out_ptr + (((((remapped_i0 * D1 + i1) * D2 + i2) * D3 + i3)[:, :, :, :, None] * D4) + i4) + old = tl.atomic_max(ptr, value) + tl.store(old_ptr + linear, old) + + +STRUCTURED_DISC_MASK_KERNELS = { + 1: structured_disc_mask_atomic_max_1d, + 2: structured_disc_mask_atomic_max_2d, + 3: structured_disc_mask_atomic_max_3d, + 4: structured_disc_mask_atomic_max_4d, + 5: structured_disc_mask_atomic_max_5d, +} + +PARTIAL_STRUCTURED_KERNELS = { + 2: partial_structured_atomic_max_2d, + 3: partial_structured_atomic_max_3d, + 4: partial_structured_atomic_max_4d, + 5: partial_structured_atomic_max_5d, +} + +FULLY_UNSTRUCTURED_KERNELS = { + 1: fully_unstructured_atomic_max_1d, + 2: fully_unstructured_atomic_max_2d, + 3: fully_unstructured_atomic_max_3d, + 4: fully_unstructured_atomic_max_4d, + 5: fully_unstructured_atomic_max_5d, +} + + +def _build_partial_structured_offsets(shape): + assert len(shape) >= 2 + first_dim = shape[0] + tail_shape = shape[1:] + tail_numel = math.prod(tail_shape) + remapped_first = ((torch.arange(first_dim, dtype=torch.int64) + 1) % first_dim) + first_offsets = remapped_first.reshape((first_dim, ) + (1, ) * len(tail_shape)) + tail_offsets = torch.arange(tail_numel, dtype=torch.int64).reshape((1, ) + tail_shape) + return first_offsets * tail_numel + tail_offsets, math.prod(shape) + + +def _build_fully_unstructured_offsets(shape): + numel = math.prod(shape) + output_numel = numel * 2 + base = torch.arange(numel, dtype=torch.int64) + offsets = (base % 2) * numel + (base // 2) + offsets = offsets.reshape(shape) + return offsets, output_numel + + +def _build_discrete_mask(shape): + numel = math.prod(shape) + return ((torch.arange(numel, dtype=torch.int64).reshape(shape) * 2) < numel) + + +def _build_value_tensor(shape, dtype): + return torch.full(shape, math.prod(shape) + 7, dtype=dtype) + + +def _build_output_baseline(output_numel, dtype): + return torch.arange(output_numel, dtype=dtype) + + +def _simulate_atomic_max(base_output, offsets, values, mask=None): + expected_output = base_output.reshape(-1).clone().cpu() + flat_offsets = offsets.reshape(-1).to(torch.int64).cpu() + flat_values = values.reshape(-1).cpu() + flat_old = torch.zeros(flat_values.shape, dtype=base_output.dtype) + if mask is None: + flat_mask = torch.ones(flat_offsets.shape, dtype=torch.bool) + else: + flat_mask = mask.reshape(-1).cpu() + + for idx in range(flat_offsets.numel()): + if not bool(flat_mask[idx]): + continue + offset = int(flat_offsets[idx].item()) + flat_old[idx] = expected_output[offset] + expected_output[offset] = torch.maximum(expected_output[offset], flat_values[idx].to(expected_output.dtype)) + + return expected_output, flat_old.reshape(values.shape) + + +def _assert_equal(actual, expected, dtype_name, rank, scenario): + actual_cpu = actual.cpu() + expected_cpu = expected.cpu() + assert torch.equal(actual_cpu, expected_cpu), (f"scenario={scenario}, dtype={dtype_name}, rank={rank}\n" + f"Expected:\n{expected_cpu}\nGot:\n{actual_cpu}") + + +def _launch_structured_discrete_mask(rank, values, output, old, shape): + kernel = STRUCTURED_DISC_MASK_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](values, output, old, **kwargs) + + +def _launch_partial_structured(rank, values, output, old, shape): + kernel = PARTIAL_STRUCTURED_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](values, output, old, **kwargs) + + +def _launch_fully_unstructured(rank, offsets, values, output, old, shape): + kernel = FULLY_UNSTRUCTURED_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](offsets, values, output, old, **kwargs) + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_max_structured_pointer_with_discrete_mask(dtype_name, torch_dtype, rank): + shape = RANK_SHAPES[rank] + values = _build_value_tensor(shape, torch_dtype).npu() + output_numel = math.prod(shape) + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_structured_discrete_mask(rank, values, output, old, shape) + + identity_offsets = torch.arange(output_numel, dtype=torch.int64).reshape(shape) + mask = _build_discrete_mask(shape) + expected_output, expected_old = _simulate_atomic_max( + baseline, + identity_offsets, + _build_value_tensor(shape, torch_dtype), + mask=mask, + ) + _assert_equal(output, expected_output, dtype_name, rank, "structured+discrete-mask/output") + _assert_equal(old, expected_old, dtype_name, rank, "structured+discrete-mask/old") + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_max_partially_structured_indirect_offsets(dtype_name, torch_dtype, rank): + shape = PARTIAL_STRUCTURED_SHAPES[rank] + offsets, output_numel = _build_partial_structured_offsets(shape) + values = _build_value_tensor(shape, torch_dtype).npu() + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_partial_structured(rank, values, output, old, shape) + + expected_output, expected_old = _simulate_atomic_max( + baseline, + offsets, + _build_value_tensor(shape, torch_dtype), + ) + _assert_equal(output, expected_output, dtype_name, rank, "partially-structured-indirect/output") + _assert_equal(old, expected_old, dtype_name, rank, "partially-structured-indirect/old") + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_max_fully_unstructured_indirect_offsets(dtype_name, torch_dtype, rank): + shape = RANK_SHAPES[rank] + offsets, output_numel = _build_fully_unstructured_offsets(shape) + values = _build_value_tensor(shape, torch_dtype).npu() + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_fully_unstructured(rank, offsets.npu(), values, output, old, shape) + + expected_output, expected_old = _simulate_atomic_max( + baseline, + offsets, + _build_value_tensor(shape, torch_dtype), + ) + _assert_equal(output, expected_output, dtype_name, rank, "fully-unstructured-indirect/output") + _assert_equal(old, expected_old, dtype_name, rank, "fully-unstructured-indirect/old") diff --git a/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_min.py b/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_min.py new file mode 100644 index 0000000000..2a2b98bccb --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_min.py @@ -0,0 +1,520 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +# ============================================================================= +# Indirect atomic_min test suite +# +# Test matrix (access pattern x tensor rank): +# +# | access pattern | 1-D | 2-D | 3-D | 4-D | 5-D | +# |-------------------------------------------------|------|------|------|------|------| +# | structured pointer + discrete mask | (A1) | (A2) | (A3) | (A4) | (A5) | +# | partial structured: high-dim disc + low-dim cont| - | (B2) | (B3) | (B4) | (B5) | +# | fully unstructured indirect offsets | (C1) | (C2) | (C3) | (C4) | (C5) | +# +# Notes: +# 1. Case A exercises the structured-pointer discrete-mask atomic_min path. +# 2. Case B exercises a single high-dimension discrete remap with the +# remaining lower dimensions kept contiguous. +# 3. Case C exercises fully unstructured indirect offsets. +# 4. All cases validate both the final destination tensor and the atomic_min +# return value, which must be the old value observed at each access. +# ============================================================================= + +import math + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +SUPPORTED_DTYPES = [ + ("int32", torch.int32), + ("int64", torch.int64), + ("float16", torch.float16), + ("float32", torch.float32), + ("bfloat16", torch.bfloat16), +] + +RANK_SHAPES = { + 1: (8, ), + 2: (4, 4), + 3: (2, 3, 4), + 4: (2, 2, 3, 4), + 5: (2, 2, 2, 3, 4), +} + +PARTIAL_STRUCTURED_SHAPES = { + 2: (2, 16), + 3: (2, 2, 8), + 4: (2, 2, 2, 8), + 5: (2, 2, 2, 2, 8), +} + +TEST_RANKS = [3] +TEST_DTYPE = [("int32", torch.int32), ("bfloat16", torch.bfloat16)] + + +@triton.jit +def structured_disc_mask_atomic_min_1d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, +): + i0 = tl.arange(0, D0) + linear = i0 + mask = (linear * 2) < D0 + value = tl.load(val_ptr + linear) + old = tl.atomic_min(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_min_2d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + i1 = tl.arange(0, D1)[None, :] + linear = i0 * D1 + i1 + mask = (linear * 2) < (D0 * D1) + value = tl.load(val_ptr + linear) + old = tl.atomic_min(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_min_3d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1) * D2 + i2 + mask = (linear * 2) < (D0 * D1 * D2) + value = tl.load(val_ptr + linear) + old = tl.atomic_min(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_min_4d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = ((i0 * D1 + i1) * D2 + i2) * D3 + i3 + mask = (linear * 2) < (D0 * D1 * D2 * D3) + value = tl.load(val_ptr + linear) + old = tl.atomic_min(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_min_5d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None, None] + i2 = tl.arange(0, D2)[None, None, :, None, None] + i3 = tl.arange(0, D3)[None, None, None, :, None] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2) * D3 + i3) * D4 + i4 + mask = (linear * 2) < (D0 * D1 * D2 * D3 * D4) + value = tl.load(val_ptr + linear) + old = tl.atomic_min(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def fully_unstructured_atomic_min_1d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, +): + i0 = tl.arange(0, D0) + linear = i0 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_min(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_min_2d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + i1 = tl.arange(0, D1)[None, :] + linear = i0 * D1 + i1 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_min(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_min_3d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1) * D2 + i2 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_min(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_min_4d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = ((i0 * D1 + i1) * D2 + i2) * D3 + i3 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_min(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_min_5d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None, None] + i2 = tl.arange(0, D2)[None, None, :, None, None] + i3 = tl.arange(0, D3)[None, None, None, :, None] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2) * D3 + i3) * D4 + i4 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_min(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_min_2d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + remapped_i0 = (i0 + 1) % D0 + inner_idx = tl.arange(0, D1)[None, :] + linear = i0 * D1 + inner_idx + value = tl.load(val_ptr + linear) + ptr = out_ptr + remapped_i0 * D1 + inner_idx + old = tl.atomic_min(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_min_3d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1)[:, :, None] * D2 + i2 + value = tl.load(val_ptr + linear) + ptr = out_ptr + ((remapped_i0 * D1 + i1)[:, :, None] * D2 + i2) + old = tl.atomic_min(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_min_4d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2)[:, :, :, None] * D3 + i3) + value = tl.load(val_ptr + linear) + ptr = out_ptr + ((((remapped_i0 * D1 + i1) * D2 + i2)[:, :, :, None] * D3) + i3) + old = tl.atomic_min(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_min_5d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = ((((i0 * D1 + i1) * D2 + i2) * D3 + i3)[:, :, :, :, None] * D4 + i4) + value = tl.load(val_ptr + linear) + ptr = out_ptr + (((((remapped_i0 * D1 + i1) * D2 + i2) * D3 + i3)[:, :, :, :, None] * D4) + i4) + old = tl.atomic_min(ptr, value) + tl.store(old_ptr + linear, old) + + +STRUCTURED_DISC_MASK_KERNELS = { + 1: structured_disc_mask_atomic_min_1d, + 2: structured_disc_mask_atomic_min_2d, + 3: structured_disc_mask_atomic_min_3d, + 4: structured_disc_mask_atomic_min_4d, + 5: structured_disc_mask_atomic_min_5d, +} + +PARTIAL_STRUCTURED_KERNELS = { + 2: partial_structured_atomic_min_2d, + 3: partial_structured_atomic_min_3d, + 4: partial_structured_atomic_min_4d, + 5: partial_structured_atomic_min_5d, +} + +FULLY_UNSTRUCTURED_KERNELS = { + 1: fully_unstructured_atomic_min_1d, + 2: fully_unstructured_atomic_min_2d, + 3: fully_unstructured_atomic_min_3d, + 4: fully_unstructured_atomic_min_4d, + 5: fully_unstructured_atomic_min_5d, +} + + +def _build_partial_structured_offsets(shape): + assert len(shape) >= 2 + first_dim = shape[0] + tail_shape = shape[1:] + tail_numel = math.prod(tail_shape) + remapped_first = ((torch.arange(first_dim, dtype=torch.int64) + 1) % first_dim) + first_offsets = remapped_first.reshape((first_dim, ) + (1, ) * len(tail_shape)) + tail_offsets = torch.arange(tail_numel, dtype=torch.int64).reshape((1, ) + tail_shape) + return first_offsets * tail_numel + tail_offsets, math.prod(shape) + + +def _build_fully_unstructured_offsets(shape): + numel = math.prod(shape) + output_numel = numel * 2 + base = torch.arange(numel, dtype=torch.int64) + offsets = (base % 2) * numel + (base // 2) + offsets = offsets.reshape(shape) + return offsets, output_numel + + +def _build_discrete_mask(shape): + numel = math.prod(shape) + return ((torch.arange(numel, dtype=torch.int64).reshape(shape) * 2) < numel) + + +def _build_value_tensor(shape, dtype): + return torch.ones(shape, dtype=dtype) + + +def _build_output_baseline(output_numel, dtype): + return torch.arange(output_numel, dtype=dtype) + + +def _simulate_atomic_min(base_output, offsets, values, mask=None): + expected_output = base_output.reshape(-1).clone().cpu() + flat_offsets = offsets.reshape(-1).to(torch.int64).cpu() + flat_values = values.reshape(-1).cpu() + flat_old = torch.zeros(flat_values.shape, dtype=base_output.dtype) + if mask is None: + flat_mask = torch.ones(flat_offsets.shape, dtype=torch.bool) + else: + flat_mask = mask.reshape(-1).cpu() + + for idx in range(flat_offsets.numel()): + if not bool(flat_mask[idx]): + continue + offset = int(flat_offsets[idx].item()) + flat_old[idx] = expected_output[offset] + expected_output[offset] = torch.minimum(expected_output[offset], flat_values[idx].to(expected_output.dtype)) + + return expected_output, flat_old.reshape(values.shape) + + +def _assert_equal(actual, expected, dtype_name, rank, scenario): + actual_cpu = actual.cpu() + expected_cpu = expected.cpu() + assert torch.equal(actual_cpu, expected_cpu), (f"scenario={scenario}, dtype={dtype_name}, rank={rank}\n" + f"Expected:\n{expected_cpu}\nGot:\n{actual_cpu}") + + +def _launch_structured_discrete_mask(rank, values, output, old, shape): + kernel = STRUCTURED_DISC_MASK_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](values, output, old, **kwargs) + + +def _launch_partial_structured(rank, values, output, old, shape): + kernel = PARTIAL_STRUCTURED_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](values, output, old, **kwargs) + + +def _launch_fully_unstructured(rank, offsets, values, output, old, shape): + kernel = FULLY_UNSTRUCTURED_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](offsets, values, output, old, **kwargs) + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_min_structured_pointer_with_discrete_mask(dtype_name, torch_dtype, rank): + shape = RANK_SHAPES[rank] + values = _build_value_tensor(shape, torch_dtype).npu() + output_numel = math.prod(shape) + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_structured_discrete_mask(rank, values, output, old, shape) + + identity_offsets = torch.arange(output_numel, dtype=torch.int64).reshape(shape) + mask = _build_discrete_mask(shape) + expected_output, expected_old = _simulate_atomic_min( + baseline, + identity_offsets, + _build_value_tensor(shape, torch_dtype), + mask=mask, + ) + _assert_equal(output, expected_output, dtype_name, rank, "structured+discrete-mask/output") + _assert_equal(old, expected_old, dtype_name, rank, "structured+discrete-mask/old") + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_min_partially_structured_indirect_offsets(dtype_name, torch_dtype, rank): + shape = PARTIAL_STRUCTURED_SHAPES[rank] + offsets, output_numel = _build_partial_structured_offsets(shape) + values = _build_value_tensor(shape, torch_dtype).npu() + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_partial_structured(rank, values, output, old, shape) + + expected_output, expected_old = _simulate_atomic_min( + baseline, + offsets, + _build_value_tensor(shape, torch_dtype), + ) + _assert_equal(output, expected_output, dtype_name, rank, "partially-structured-indirect/output") + _assert_equal(old, expected_old, dtype_name, rank, "partially-structured-indirect/old") + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_min_fully_unstructured_indirect_offsets(dtype_name, torch_dtype, rank): + shape = RANK_SHAPES[rank] + offsets, output_numel = _build_fully_unstructured_offsets(shape) + values = _build_value_tensor(shape, torch_dtype).npu() + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_fully_unstructured(rank, offsets.npu(), values, output, old, shape) + + expected_output, expected_old = _simulate_atomic_min( + baseline, + offsets, + _build_value_tensor(shape, torch_dtype), + ) + _assert_equal(output, expected_output, dtype_name, rank, "fully-unstructured-indirect/output") + _assert_equal(old, expected_old, dtype_name, rank, "fully-unstructured-indirect/old") diff --git a/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_or.py b/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_or.py new file mode 100644 index 0000000000..f9e09b085a --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_or.py @@ -0,0 +1,517 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +# ============================================================================= +# Indirect atomic_or test suite +# +# Test matrix (access pattern x tensor rank): +# +# | access pattern | 1-D | 2-D | 3-D | 4-D | 5-D | +# |-------------------------------------------------|------|------|------|------|------| +# | structured pointer + discrete mask | (A1) | (A2) | (A3) | (A4) | (A5) | +# | partial structured: high-dim disc + low-dim cont| - | (B2) | (B3) | (B4) | (B5) | +# | fully unstructured indirect offsets | (C1) | (C2) | (C3) | (C4) | (C5) | +# +# Notes: +# 1. Case A exercises the structured-pointer discrete-mask atomic_or path. +# 2. Case B exercises a single high-dimension discrete remap with the +# remaining lower dimensions kept contiguous. +# 3. Case C exercises fully unstructured indirect offsets. +# 4. All cases validate both the final destination tensor and the atomic_or +# return value, which must be the old value observed at each access. +# ============================================================================= + +import math + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +SUPPORTED_DTYPES = [ + ("int32", torch.int32), + ("int64", torch.int64), +] + +RANK_SHAPES = { + 1: (8, ), + 2: (4, 4), + 3: (2, 3, 4), + 4: (2, 2, 3, 4), + 5: (2, 2, 2, 3, 4), +} + +PARTIAL_STRUCTURED_SHAPES = { + 2: (2, 16), + 3: (2, 2, 8), + 4: (2, 2, 2, 8), + 5: (2, 2, 2, 2, 8), +} + +TEST_RANKS = [3] +TEST_DTYPE = [("int32", torch.int32)] + + +@triton.jit +def structured_disc_mask_atomic_or_1d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, +): + i0 = tl.arange(0, D0) + linear = i0 + mask = (linear * 2) < D0 + value = tl.load(val_ptr + linear) + old = tl.atomic_or(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_or_2d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + i1 = tl.arange(0, D1)[None, :] + linear = i0 * D1 + i1 + mask = (linear * 2) < (D0 * D1) + value = tl.load(val_ptr + linear) + old = tl.atomic_or(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_or_3d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1) * D2 + i2 + mask = (linear * 2) < (D0 * D1 * D2) + value = tl.load(val_ptr + linear) + old = tl.atomic_or(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_or_4d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = ((i0 * D1 + i1) * D2 + i2) * D3 + i3 + mask = (linear * 2) < (D0 * D1 * D2 * D3) + value = tl.load(val_ptr + linear) + old = tl.atomic_or(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_or_5d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None, None] + i2 = tl.arange(0, D2)[None, None, :, None, None] + i3 = tl.arange(0, D3)[None, None, None, :, None] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2) * D3 + i3) * D4 + i4 + mask = (linear * 2) < (D0 * D1 * D2 * D3 * D4) + value = tl.load(val_ptr + linear) + old = tl.atomic_or(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def fully_unstructured_atomic_or_1d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, +): + i0 = tl.arange(0, D0) + linear = i0 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_or(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_or_2d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + i1 = tl.arange(0, D1)[None, :] + linear = i0 * D1 + i1 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_or(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_or_3d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1) * D2 + i2 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_or(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_or_4d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = ((i0 * D1 + i1) * D2 + i2) * D3 + i3 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_or(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_or_5d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None, None] + i2 = tl.arange(0, D2)[None, None, :, None, None] + i3 = tl.arange(0, D3)[None, None, None, :, None] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2) * D3 + i3) * D4 + i4 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_or(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_or_2d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + remapped_i0 = (i0 + 1) % D0 + inner_idx = tl.arange(0, D1)[None, :] + linear = i0 * D1 + inner_idx + value = tl.load(val_ptr + linear) + ptr = out_ptr + remapped_i0 * D1 + inner_idx + old = tl.atomic_or(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_or_3d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1)[:, :, None] * D2 + i2 + value = tl.load(val_ptr + linear) + ptr = out_ptr + ((remapped_i0 * D1 + i1)[:, :, None] * D2 + i2) + old = tl.atomic_or(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_or_4d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2)[:, :, :, None] * D3 + i3) + value = tl.load(val_ptr + linear) + ptr = out_ptr + ((((remapped_i0 * D1 + i1) * D2 + i2)[:, :, :, None] * D3) + i3) + old = tl.atomic_or(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_or_5d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = ((((i0 * D1 + i1) * D2 + i2) * D3 + i3)[:, :, :, :, None] * D4 + i4) + value = tl.load(val_ptr + linear) + ptr = out_ptr + (((((remapped_i0 * D1 + i1) * D2 + i2) * D3 + i3)[:, :, :, :, None] * D4) + i4) + old = tl.atomic_or(ptr, value) + tl.store(old_ptr + linear, old) + + +STRUCTURED_DISC_MASK_KERNELS = { + 1: structured_disc_mask_atomic_or_1d, + 2: structured_disc_mask_atomic_or_2d, + 3: structured_disc_mask_atomic_or_3d, + 4: structured_disc_mask_atomic_or_4d, + 5: structured_disc_mask_atomic_or_5d, +} + +PARTIAL_STRUCTURED_KERNELS = { + 2: partial_structured_atomic_or_2d, + 3: partial_structured_atomic_or_3d, + 4: partial_structured_atomic_or_4d, + 5: partial_structured_atomic_or_5d, +} + +FULLY_UNSTRUCTURED_KERNELS = { + 1: fully_unstructured_atomic_or_1d, + 2: fully_unstructured_atomic_or_2d, + 3: fully_unstructured_atomic_or_3d, + 4: fully_unstructured_atomic_or_4d, + 5: fully_unstructured_atomic_or_5d, +} + + +def _build_partial_structured_offsets(shape): + assert len(shape) >= 2 + first_dim = shape[0] + tail_shape = shape[1:] + tail_numel = math.prod(tail_shape) + remapped_first = ((torch.arange(first_dim, dtype=torch.int64) + 1) % first_dim) + first_offsets = remapped_first.reshape((first_dim, ) + (1, ) * len(tail_shape)) + tail_offsets = torch.arange(tail_numel, dtype=torch.int64).reshape((1, ) + tail_shape) + return first_offsets * tail_numel + tail_offsets, math.prod(shape) + + +def _build_fully_unstructured_offsets(shape): + numel = math.prod(shape) + output_numel = numel * 2 + base = torch.arange(numel, dtype=torch.int64) + offsets = (base % 2) * numel + (base // 2) + offsets = offsets.reshape(shape) + return offsets, output_numel + + +def _build_discrete_mask(shape): + numel = math.prod(shape) + return ((torch.arange(numel, dtype=torch.int64).reshape(shape) * 2) < numel) + + +def _build_value_tensor(shape, dtype): + return torch.ones(shape, dtype=dtype) + + +def _build_output_baseline(output_numel, dtype): + return torch.arange(output_numel, dtype=dtype) + + +def _simulate_atomic_or(base_output, offsets, values, mask=None): + expected_output = base_output.reshape(-1).clone().cpu() + flat_offsets = offsets.reshape(-1).to(torch.int64).cpu() + flat_values = values.reshape(-1).cpu() + flat_old = torch.zeros(flat_values.shape, dtype=base_output.dtype) + if mask is None: + flat_mask = torch.ones(flat_offsets.shape, dtype=torch.bool) + else: + flat_mask = mask.reshape(-1).cpu() + + for idx in range(flat_offsets.numel()): + if not bool(flat_mask[idx]): + continue + offset = int(flat_offsets[idx].item()) + flat_old[idx] = expected_output[offset] + expected_output[offset] = expected_output[offset] | flat_values[idx].to(expected_output.dtype) + + return expected_output, flat_old.reshape(values.shape) + + +def _assert_equal(actual, expected, dtype_name, rank, scenario): + actual_cpu = actual.cpu() + expected_cpu = expected.cpu() + assert torch.equal(actual_cpu, expected_cpu), (f"scenario={scenario}, dtype={dtype_name}, rank={rank}\n" + f"Expected:\n{expected_cpu}\nGot:\n{actual_cpu}") + + +def _launch_structured_discrete_mask(rank, values, output, old, shape): + kernel = STRUCTURED_DISC_MASK_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](values, output, old, **kwargs) + + +def _launch_partial_structured(rank, values, output, old, shape): + kernel = PARTIAL_STRUCTURED_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](values, output, old, **kwargs) + + +def _launch_fully_unstructured(rank, offsets, values, output, old, shape): + kernel = FULLY_UNSTRUCTURED_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](offsets, values, output, old, **kwargs) + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_or_structured_pointer_with_discrete_mask(dtype_name, torch_dtype, rank): + shape = RANK_SHAPES[rank] + values = _build_value_tensor(shape, torch_dtype).npu() + output_numel = math.prod(shape) + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_structured_discrete_mask(rank, values, output, old, shape) + + identity_offsets = torch.arange(output_numel, dtype=torch.int64).reshape(shape) + mask = _build_discrete_mask(shape) + expected_output, expected_old = _simulate_atomic_or( + baseline, + identity_offsets, + _build_value_tensor(shape, torch_dtype), + mask=mask, + ) + _assert_equal(output, expected_output, dtype_name, rank, "structured+discrete-mask/output") + _assert_equal(old, expected_old, dtype_name, rank, "structured+discrete-mask/old") + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_or_partially_structured_indirect_offsets(dtype_name, torch_dtype, rank): + shape = PARTIAL_STRUCTURED_SHAPES[rank] + offsets, output_numel = _build_partial_structured_offsets(shape) + values = _build_value_tensor(shape, torch_dtype).npu() + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_partial_structured(rank, values, output, old, shape) + + expected_output, expected_old = _simulate_atomic_or( + baseline, + offsets, + _build_value_tensor(shape, torch_dtype), + ) + _assert_equal(output, expected_output, dtype_name, rank, "partially-structured-indirect/output") + _assert_equal(old, expected_old, dtype_name, rank, "partially-structured-indirect/old") + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_or_fully_unstructured_indirect_offsets(dtype_name, torch_dtype, rank): + shape = RANK_SHAPES[rank] + offsets, output_numel = _build_fully_unstructured_offsets(shape) + values = _build_value_tensor(shape, torch_dtype).npu() + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_fully_unstructured(rank, offsets.npu(), values, output, old, shape) + + expected_output, expected_old = _simulate_atomic_or( + baseline, + offsets, + _build_value_tensor(shape, torch_dtype), + ) + _assert_equal(output, expected_output, dtype_name, rank, "fully-unstructured-indirect/output") + _assert_equal(old, expected_old, dtype_name, rank, "fully-unstructured-indirect/old") diff --git a/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_xchg.py b/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_xchg.py new file mode 100644 index 0000000000..b1e6882232 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_xchg.py @@ -0,0 +1,520 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +# ============================================================================= +# Indirect atomic_xchg test suite +# +# Test matrix (access pattern x tensor rank): +# +# | access pattern | 1-D | 2-D | 3-D | 4-D | 5-D | +# |-------------------------------------------------|------|------|------|------|------| +# | structured pointer + discrete mask | (A1) | (A2) | (A3) | (A4) | (A5) | +# | partial structured: high-dim disc + low-dim cont| - | (B2) | (B3) | (B4) | (B5) | +# | fully unstructured indirect offsets | (C1) | (C2) | (C3) | (C4) | (C5) | +# +# Notes: +# 1. Case A exercises the structured-pointer discrete-mask atomic_xchg path. +# 2. Case B exercises a single high-dimension discrete remap with the +# remaining lower dimensions kept contiguous. +# 3. Case C exercises fully unstructured indirect offsets. +# 4. All cases validate both the final destination tensor and the atomic_xchg +# return value, which must be the old value observed at each access. +# ============================================================================= + +import math + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +SUPPORTED_DTYPES = [ + ("int32", torch.int32), + ("int64", torch.int64), + ("float16", torch.float16), + ("float32", torch.float32), + ("bfloat16", torch.bfloat16), +] + +RANK_SHAPES = { + 1: (8, ), + 2: (4, 4), + 3: (2, 3, 4), + 4: (2, 2, 3, 4), + 5: (2, 2, 2, 3, 4), +} + +PARTIAL_STRUCTURED_SHAPES = { + 2: (2, 16), + 3: (2, 2, 8), + 4: (2, 2, 2, 8), + 5: (2, 2, 2, 2, 8), +} + +TEST_RANKS = [3] +TEST_DTYPE = [("int32", torch.int32), ("bfloat16", torch.bfloat16)] + + +@triton.jit +def structured_disc_mask_atomic_xchg_1d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, +): + i0 = tl.arange(0, D0) + linear = i0 + mask = (linear * 2) < D0 + value = tl.load(val_ptr + linear) + old = tl.atomic_xchg(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_xchg_2d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + i1 = tl.arange(0, D1)[None, :] + linear = i0 * D1 + i1 + mask = (linear * 2) < (D0 * D1) + value = tl.load(val_ptr + linear) + old = tl.atomic_xchg(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_xchg_3d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1) * D2 + i2 + mask = (linear * 2) < (D0 * D1 * D2) + value = tl.load(val_ptr + linear) + old = tl.atomic_xchg(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_xchg_4d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = ((i0 * D1 + i1) * D2 + i2) * D3 + i3 + mask = (linear * 2) < (D0 * D1 * D2 * D3) + value = tl.load(val_ptr + linear) + old = tl.atomic_xchg(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_xchg_5d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None, None] + i2 = tl.arange(0, D2)[None, None, :, None, None] + i3 = tl.arange(0, D3)[None, None, None, :, None] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2) * D3 + i3) * D4 + i4 + mask = (linear * 2) < (D0 * D1 * D2 * D3 * D4) + value = tl.load(val_ptr + linear) + old = tl.atomic_xchg(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def fully_unstructured_atomic_xchg_1d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, +): + i0 = tl.arange(0, D0) + linear = i0 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_xchg(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_xchg_2d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + i1 = tl.arange(0, D1)[None, :] + linear = i0 * D1 + i1 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_xchg(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_xchg_3d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1) * D2 + i2 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_xchg(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_xchg_4d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = ((i0 * D1 + i1) * D2 + i2) * D3 + i3 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_xchg(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_xchg_5d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None, None] + i2 = tl.arange(0, D2)[None, None, :, None, None] + i3 = tl.arange(0, D3)[None, None, None, :, None] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2) * D3 + i3) * D4 + i4 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_xchg(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_xchg_2d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + remapped_i0 = (i0 + 1) % D0 + inner_idx = tl.arange(0, D1)[None, :] + linear = i0 * D1 + inner_idx + value = tl.load(val_ptr + linear) + ptr = out_ptr + remapped_i0 * D1 + inner_idx + old = tl.atomic_xchg(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_xchg_3d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1)[:, :, None] * D2 + i2 + value = tl.load(val_ptr + linear) + ptr = out_ptr + ((remapped_i0 * D1 + i1)[:, :, None] * D2 + i2) + old = tl.atomic_xchg(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_xchg_4d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2)[:, :, :, None] * D3 + i3) + value = tl.load(val_ptr + linear) + ptr = out_ptr + ((((remapped_i0 * D1 + i1) * D2 + i2)[:, :, :, None] * D3) + i3) + old = tl.atomic_xchg(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_xchg_5d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = ((((i0 * D1 + i1) * D2 + i2) * D3 + i3)[:, :, :, :, None] * D4 + i4) + value = tl.load(val_ptr + linear) + ptr = out_ptr + (((((remapped_i0 * D1 + i1) * D2 + i2) * D3 + i3)[:, :, :, :, None] * D4) + i4) + old = tl.atomic_xchg(ptr, value) + tl.store(old_ptr + linear, old) + + +STRUCTURED_DISC_MASK_KERNELS = { + 1: structured_disc_mask_atomic_xchg_1d, + 2: structured_disc_mask_atomic_xchg_2d, + 3: structured_disc_mask_atomic_xchg_3d, + 4: structured_disc_mask_atomic_xchg_4d, + 5: structured_disc_mask_atomic_xchg_5d, +} + +PARTIAL_STRUCTURED_KERNELS = { + 2: partial_structured_atomic_xchg_2d, + 3: partial_structured_atomic_xchg_3d, + 4: partial_structured_atomic_xchg_4d, + 5: partial_structured_atomic_xchg_5d, +} + +FULLY_UNSTRUCTURED_KERNELS = { + 1: fully_unstructured_atomic_xchg_1d, + 2: fully_unstructured_atomic_xchg_2d, + 3: fully_unstructured_atomic_xchg_3d, + 4: fully_unstructured_atomic_xchg_4d, + 5: fully_unstructured_atomic_xchg_5d, +} + + +def _build_partial_structured_offsets(shape): + assert len(shape) >= 2 + first_dim = shape[0] + tail_shape = shape[1:] + tail_numel = math.prod(tail_shape) + remapped_first = ((torch.arange(first_dim, dtype=torch.int64) + 1) % first_dim) + first_offsets = remapped_first.reshape((first_dim, ) + (1, ) * len(tail_shape)) + tail_offsets = torch.arange(tail_numel, dtype=torch.int64).reshape((1, ) + tail_shape) + return first_offsets * tail_numel + tail_offsets, math.prod(shape) + + +def _build_fully_unstructured_offsets(shape): + numel = math.prod(shape) + output_numel = numel * 2 + base = torch.arange(numel, dtype=torch.int64) + offsets = (base % 2) * numel + (base // 2) + offsets = offsets.reshape(shape) + return offsets, output_numel + + +def _build_discrete_mask(shape): + numel = math.prod(shape) + return ((torch.arange(numel, dtype=torch.int64).reshape(shape) * 2) < numel) + + +def _build_value_tensor(shape, dtype): + return torch.ones(shape, dtype=dtype) + + +def _build_output_baseline(output_numel, dtype): + return torch.arange(output_numel, dtype=dtype) + + +def _simulate_atomic_xchg(base_output, offsets, values, mask=None): + expected_output = base_output.reshape(-1).clone().cpu() + flat_offsets = offsets.reshape(-1).to(torch.int64).cpu() + flat_values = values.reshape(-1).cpu() + flat_old = torch.zeros(flat_values.shape, dtype=base_output.dtype) + if mask is None: + flat_mask = torch.ones(flat_offsets.shape, dtype=torch.bool) + else: + flat_mask = mask.reshape(-1).cpu() + + for idx in range(flat_offsets.numel()): + if not bool(flat_mask[idx]): + continue + offset = int(flat_offsets[idx].item()) + flat_old[idx] = expected_output[offset] + expected_output[offset] = flat_values[idx].to(expected_output.dtype) + + return expected_output, flat_old.reshape(values.shape) + + +def _assert_equal(actual, expected, dtype_name, rank, scenario): + actual_cpu = actual.cpu() + expected_cpu = expected.cpu() + assert torch.equal(actual_cpu, expected_cpu), (f"scenario={scenario}, dtype={dtype_name}, rank={rank}\n" + f"Expected:\n{expected_cpu}\nGot:\n{actual_cpu}") + + +def _launch_structured_discrete_mask(rank, values, output, old, shape): + kernel = STRUCTURED_DISC_MASK_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](values, output, old, **kwargs) + + +def _launch_partial_structured(rank, values, output, old, shape): + kernel = PARTIAL_STRUCTURED_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](values, output, old, **kwargs) + + +def _launch_fully_unstructured(rank, offsets, values, output, old, shape): + kernel = FULLY_UNSTRUCTURED_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](offsets, values, output, old, **kwargs) + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_xchg_structured_pointer_with_discrete_mask(dtype_name, torch_dtype, rank): + shape = RANK_SHAPES[rank] + values = _build_value_tensor(shape, torch_dtype).npu() + output_numel = math.prod(shape) + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_structured_discrete_mask(rank, values, output, old, shape) + + identity_offsets = torch.arange(output_numel, dtype=torch.int64).reshape(shape) + mask = _build_discrete_mask(shape) + expected_output, expected_old = _simulate_atomic_xchg( + baseline, + identity_offsets, + _build_value_tensor(shape, torch_dtype), + mask=mask, + ) + _assert_equal(output, expected_output, dtype_name, rank, "structured+discrete-mask/output") + _assert_equal(old, expected_old, dtype_name, rank, "structured+discrete-mask/old") + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_xchg_partially_structured_indirect_offsets(dtype_name, torch_dtype, rank): + shape = PARTIAL_STRUCTURED_SHAPES[rank] + offsets, output_numel = _build_partial_structured_offsets(shape) + values = _build_value_tensor(shape, torch_dtype).npu() + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_partial_structured(rank, values, output, old, shape) + + expected_output, expected_old = _simulate_atomic_xchg( + baseline, + offsets, + _build_value_tensor(shape, torch_dtype), + ) + _assert_equal(output, expected_output, dtype_name, rank, "partially-structured-indirect/output") + _assert_equal(old, expected_old, dtype_name, rank, "partially-structured-indirect/old") + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_xchg_fully_unstructured_indirect_offsets(dtype_name, torch_dtype, rank): + shape = RANK_SHAPES[rank] + offsets, output_numel = _build_fully_unstructured_offsets(shape) + values = _build_value_tensor(shape, torch_dtype).npu() + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_fully_unstructured(rank, offsets.npu(), values, output, old, shape) + + expected_output, expected_old = _simulate_atomic_xchg( + baseline, + offsets, + _build_value_tensor(shape, torch_dtype), + ) + _assert_equal(output, expected_output, dtype_name, rank, "fully-unstructured-indirect/output") + _assert_equal(old, expected_old, dtype_name, rank, "fully-unstructured-indirect/old") diff --git a/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_xor.py b/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_xor.py new file mode 100644 index 0000000000..72855dbebe --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_indirect_atomic_xor.py @@ -0,0 +1,517 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +# ============================================================================= +# Indirect atomic_xor test suite +# +# Test matrix (access pattern x tensor rank): +# +# | access pattern | 1-D | 2-D | 3-D | 4-D | 5-D | +# |-------------------------------------------------|------|------|------|------|------| +# | structured pointer + discrete mask | (A1) | (A2) | (A3) | (A4) | (A5) | +# | partial structured: high-dim disc + low-dim cont| - | (B2) | (B3) | (B4) | (B5) | +# | fully unstructured indirect offsets | (C1) | (C2) | (C3) | (C4) | (C5) | +# +# Notes: +# 1. Case A exercises the structured-pointer discrete-mask atomic_xor path. +# 2. Case B exercises a single high-dimension discrete remap with the +# remaining lower dimensions kept contiguous. +# 3. Case C exercises fully unstructured indirect offsets. +# 4. All cases validate both the final destination tensor and the atomic_xor +# return value, which must be the old value observed at each access. +# ============================================================================= + +import math + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +SUPPORTED_DTYPES = [ + ("int32", torch.int32), + ("int64", torch.int64), +] + +RANK_SHAPES = { + 1: (8, ), + 2: (4, 4), + 3: (2, 3, 4), + 4: (2, 2, 3, 4), + 5: (2, 2, 2, 3, 4), +} + +PARTIAL_STRUCTURED_SHAPES = { + 2: (2, 16), + 3: (2, 2, 8), + 4: (2, 2, 2, 8), + 5: (2, 2, 2, 2, 8), +} + +TEST_RANKS = [3] +TEST_DTYPE = [("int32", torch.int32)] + + +@triton.jit +def structured_disc_mask_atomic_xor_1d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, +): + i0 = tl.arange(0, D0) + linear = i0 + mask = (linear * 2) < D0 + value = tl.load(val_ptr + linear) + old = tl.atomic_xor(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_xor_2d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + i1 = tl.arange(0, D1)[None, :] + linear = i0 * D1 + i1 + mask = (linear * 2) < (D0 * D1) + value = tl.load(val_ptr + linear) + old = tl.atomic_xor(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_xor_3d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1) * D2 + i2 + mask = (linear * 2) < (D0 * D1 * D2) + value = tl.load(val_ptr + linear) + old = tl.atomic_xor(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_xor_4d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = ((i0 * D1 + i1) * D2 + i2) * D3 + i3 + mask = (linear * 2) < (D0 * D1 * D2 * D3) + value = tl.load(val_ptr + linear) + old = tl.atomic_xor(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def structured_disc_mask_atomic_xor_5d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None, None] + i2 = tl.arange(0, D2)[None, None, :, None, None] + i3 = tl.arange(0, D3)[None, None, None, :, None] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2) * D3 + i3) * D4 + i4 + mask = (linear * 2) < (D0 * D1 * D2 * D3 * D4) + value = tl.load(val_ptr + linear) + old = tl.atomic_xor(out_ptr + linear, value, mask=mask) + tl.store(old_ptr + linear, old, mask=mask) + + +@triton.jit +def fully_unstructured_atomic_xor_1d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, +): + i0 = tl.arange(0, D0) + linear = i0 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_xor(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_xor_2d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + i1 = tl.arange(0, D1)[None, :] + linear = i0 * D1 + i1 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_xor(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_xor_3d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1) * D2 + i2 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_xor(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_xor_4d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = ((i0 * D1 + i1) * D2 + i2) * D3 + i3 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_xor(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def fully_unstructured_atomic_xor_5d( + idx_ptr, + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None, None] + i1 = tl.arange(0, D1)[None, :, None, None, None] + i2 = tl.arange(0, D2)[None, None, :, None, None] + i3 = tl.arange(0, D3)[None, None, None, :, None] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2) * D3 + i3) * D4 + i4 + offset = tl.load(idx_ptr + linear) + value = tl.load(val_ptr + linear) + old = tl.atomic_xor(out_ptr + offset, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_xor_2d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + remapped_i0 = (i0 + 1) % D0 + inner_idx = tl.arange(0, D1)[None, :] + linear = i0 * D1 + inner_idx + value = tl.load(val_ptr + linear) + ptr = out_ptr + remapped_i0 * D1 + inner_idx + old = tl.atomic_xor(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_xor_3d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :] + i2 = tl.arange(0, D2)[None, None, :] + linear = (i0 * D1 + i1)[:, :, None] * D2 + i2 + value = tl.load(val_ptr + linear) + ptr = out_ptr + ((remapped_i0 * D1 + i1)[:, :, None] * D2 + i2) + old = tl.atomic_xor(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_xor_4d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :, None] + i2 = tl.arange(0, D2)[None, None, :] + i3 = tl.arange(0, D3)[None, None, None, :] + linear = (((i0 * D1 + i1) * D2 + i2)[:, :, :, None] * D3 + i3) + value = tl.load(val_ptr + linear) + ptr = out_ptr + ((((remapped_i0 * D1 + i1) * D2 + i2)[:, :, :, None] * D3) + i3) + old = tl.atomic_xor(ptr, value) + tl.store(old_ptr + linear, old) + + +@triton.jit +def partial_structured_atomic_xor_5d( + val_ptr, + out_ptr, + old_ptr, + D0: tl.constexpr, + D1: tl.constexpr, + D2: tl.constexpr, + D3: tl.constexpr, + D4: tl.constexpr, +): + i0 = tl.arange(0, D0)[:, None, None, None] + remapped_i0 = (i0 + 1) % D0 + i1 = tl.arange(0, D1)[None, :, None, None] + i2 = tl.arange(0, D2)[None, None, :, None] + i3 = tl.arange(0, D3)[None, None, None, :] + i4 = tl.arange(0, D4)[None, None, None, None, :] + linear = ((((i0 * D1 + i1) * D2 + i2) * D3 + i3)[:, :, :, :, None] * D4 + i4) + value = tl.load(val_ptr + linear) + ptr = out_ptr + (((((remapped_i0 * D1 + i1) * D2 + i2) * D3 + i3)[:, :, :, :, None] * D4) + i4) + old = tl.atomic_xor(ptr, value) + tl.store(old_ptr + linear, old) + + +STRUCTURED_DISC_MASK_KERNELS = { + 1: structured_disc_mask_atomic_xor_1d, + 2: structured_disc_mask_atomic_xor_2d, + 3: structured_disc_mask_atomic_xor_3d, + 4: structured_disc_mask_atomic_xor_4d, + 5: structured_disc_mask_atomic_xor_5d, +} + +PARTIAL_STRUCTURED_KERNELS = { + 2: partial_structured_atomic_xor_2d, + 3: partial_structured_atomic_xor_3d, + 4: partial_structured_atomic_xor_4d, + 5: partial_structured_atomic_xor_5d, +} + +FULLY_UNSTRUCTURED_KERNELS = { + 1: fully_unstructured_atomic_xor_1d, + 2: fully_unstructured_atomic_xor_2d, + 3: fully_unstructured_atomic_xor_3d, + 4: fully_unstructured_atomic_xor_4d, + 5: fully_unstructured_atomic_xor_5d, +} + + +def _build_partial_structured_offsets(shape): + assert len(shape) >= 2 + first_dim = shape[0] + tail_shape = shape[1:] + tail_numel = math.prod(tail_shape) + remapped_first = ((torch.arange(first_dim, dtype=torch.int64) + 1) % first_dim) + first_offsets = remapped_first.reshape((first_dim, ) + (1, ) * len(tail_shape)) + tail_offsets = torch.arange(tail_numel, dtype=torch.int64).reshape((1, ) + tail_shape) + return first_offsets * tail_numel + tail_offsets, math.prod(shape) + + +def _build_fully_unstructured_offsets(shape): + numel = math.prod(shape) + output_numel = numel * 2 + base = torch.arange(numel, dtype=torch.int64) + offsets = (base % 2) * numel + (base // 2) + offsets = offsets.reshape(shape) + return offsets, output_numel + + +def _build_discrete_mask(shape): + numel = math.prod(shape) + return ((torch.arange(numel, dtype=torch.int64).reshape(shape) * 2) < numel) + + +def _build_value_tensor(shape, dtype): + return torch.ones(shape, dtype=dtype) + + +def _build_output_baseline(output_numel, dtype): + return torch.arange(output_numel, dtype=dtype) + + +def _simulate_atomic_xor(base_output, offsets, values, mask=None): + expected_output = base_output.reshape(-1).clone().cpu() + flat_offsets = offsets.reshape(-1).to(torch.int64).cpu() + flat_values = values.reshape(-1).cpu() + flat_old = torch.zeros(flat_values.shape, dtype=base_output.dtype) + if mask is None: + flat_mask = torch.ones(flat_offsets.shape, dtype=torch.bool) + else: + flat_mask = mask.reshape(-1).cpu() + + for idx in range(flat_offsets.numel()): + if not bool(flat_mask[idx]): + continue + offset = int(flat_offsets[idx].item()) + flat_old[idx] = expected_output[offset] + expected_output[offset] = expected_output[offset] ^ flat_values[idx].to(expected_output.dtype) + + return expected_output, flat_old.reshape(values.shape) + + +def _assert_equal(actual, expected, dtype_name, rank, scenario): + actual_cpu = actual.cpu() + expected_cpu = expected.cpu() + assert torch.equal(actual_cpu, expected_cpu), (f"scenario={scenario}, dtype={dtype_name}, rank={rank}\n" + f"Expected:\n{expected_cpu}\nGot:\n{actual_cpu}") + + +def _launch_structured_discrete_mask(rank, values, output, old, shape): + kernel = STRUCTURED_DISC_MASK_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](values, output, old, **kwargs) + + +def _launch_partial_structured(rank, values, output, old, shape): + kernel = PARTIAL_STRUCTURED_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](values, output, old, **kwargs) + + +def _launch_fully_unstructured(rank, offsets, values, output, old, shape): + kernel = FULLY_UNSTRUCTURED_KERNELS[rank] + kwargs = {f"D{dim}": size for dim, size in enumerate(shape)} + kernel[(1, )](offsets, values, output, old, **kwargs) + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_xor_structured_pointer_with_discrete_mask(dtype_name, torch_dtype, rank): + shape = RANK_SHAPES[rank] + values = _build_value_tensor(shape, torch_dtype).npu() + output_numel = math.prod(shape) + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_structured_discrete_mask(rank, values, output, old, shape) + + identity_offsets = torch.arange(output_numel, dtype=torch.int64).reshape(shape) + mask = _build_discrete_mask(shape) + expected_output, expected_old = _simulate_atomic_xor( + baseline, + identity_offsets, + _build_value_tensor(shape, torch_dtype), + mask=mask, + ) + _assert_equal(output, expected_output, dtype_name, rank, "structured+discrete-mask/output") + _assert_equal(old, expected_old, dtype_name, rank, "structured+discrete-mask/old") + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_xor_partially_structured_indirect_offsets(dtype_name, torch_dtype, rank): + shape = PARTIAL_STRUCTURED_SHAPES[rank] + offsets, output_numel = _build_partial_structured_offsets(shape) + values = _build_value_tensor(shape, torch_dtype).npu() + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_partial_structured(rank, values, output, old, shape) + + expected_output, expected_old = _simulate_atomic_xor( + baseline, + offsets, + _build_value_tensor(shape, torch_dtype), + ) + _assert_equal(output, expected_output, dtype_name, rank, "partially-structured-indirect/output") + _assert_equal(old, expected_old, dtype_name, rank, "partially-structured-indirect/old") + + +@pytest.mark.parametrize("dtype_name, torch_dtype", TEST_DTYPE) +@pytest.mark.parametrize("rank", TEST_RANKS) +def test_atomic_xor_fully_unstructured_indirect_offsets(dtype_name, torch_dtype, rank): + shape = RANK_SHAPES[rank] + offsets, output_numel = _build_fully_unstructured_offsets(shape) + values = _build_value_tensor(shape, torch_dtype).npu() + baseline = _build_output_baseline(output_numel, torch_dtype) + output = baseline.clone().npu() + old = torch.zeros(shape, dtype=torch_dtype).npu() + + _launch_fully_unstructured(rank, offsets.npu(), values, output, old, shape) + + expected_output, expected_old = _simulate_atomic_xor( + baseline, + offsets, + _build_value_tensor(shape, torch_dtype), + ) + _assert_equal(output, expected_output, dtype_name, rank, "fully-unstructured-indirect/output") + _assert_equal(old, expected_old, dtype_name, rank, "fully-unstructured-indirect/old") diff --git a/third_party/ascend/unittest/pytest_ut/test_indirect_scalar_load_offset.py b/third_party/ascend/unittest/pytest_ut/test_indirect_scalar_load_offset.py new file mode 100644 index 0000000000..b9a34efccd --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_indirect_scalar_load_offset.py @@ -0,0 +1,98 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + + +@triton.jit +def gather_after_reduce_kernel( + logits_ptr, + topk_ids_ptr, + output_ptr, + logits_stride, + vocab_size, + BLOCK: tl.constexpr, +): + req_idx = tl.program_id(0) + + max_val = -float('inf') + for start in range(0, vocab_size, BLOCK): + offsets = start + tl.arange(0, BLOCK) + mask = offsets < vocab_size + vals = tl.load( + logits_ptr + req_idx * logits_stride + offsets, + mask=mask, + other=-float('inf'), + ) + block_max = tl.max(vals) + max_val = tl.maximum(max_val, block_max) + + topk_id = tl.load(topk_ids_ptr + req_idx + tl.arange(0, 1)) + val = tl.load(logits_ptr + req_idx * logits_stride + topk_id) + tl.store(output_ptr + req_idx + tl.arange(0, 1), val - max_val) + + +def torch_reference(logits, topk_ids): + num_rows = logits.shape[0] + output = torch.empty(num_rows, dtype=logits.dtype) + for i in range(num_rows): + max_val = logits[i].max() + output[i] = logits[i, topk_ids[i]] - max_val + return output + + +shapes = [ + (4, 128), + (8, 256), + (16, 1024), +] + + +@pytest.mark.parametrize('num_rows,vocab_size', shapes) +def test_gather_after_reduce(num_rows, vocab_size): + BLOCK = 128 + + logits_ref = test_common.generate_tensor(shape=(num_rows, vocab_size), dtype='float32') + logits = logits_ref.npu() + logits_flat = logits.reshape(-1) + + topk_ids_ref = torch.randint(0, vocab_size, (num_rows, ), dtype=torch.int64) + topk_ids = topk_ids_ref.npu() + + output = torch.empty(num_rows, dtype=torch.float32).npu() + + gather_after_reduce_kernel[(num_rows, )]( + logits_flat, + topk_ids, + output, + vocab_size, + vocab_size, + BLOCK=BLOCK, + ) + + output_ref = torch_reference(logits_ref, topk_ids_ref) + test_common.validate_cmp('float32', output, output_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_insert_slice.py b/third_party/ascend/unittest/pytest_ut/test_insert_slice.py new file mode 100644 index 0000000000..9f968971d7 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_insert_slice.py @@ -0,0 +1,64 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu + +import triton +import triton.language as tl +import triton.language.extra.cann.extension as extension +import pytest + + +@triton.jit +def triton_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr, SLICE_OFFSET: tl.constexpr, + SLICE_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + x_sub = extension.extract_slice(x, [block_start + SLICE_OFFSET], [SLICE_SIZE], [1]) + y_sub = extension.extract_slice(y, [block_start + SLICE_OFFSET], [SLICE_SIZE], [1]) + output_sub = x_sub + y_sub + output = tl.load(output_ptr + offsets, mask=mask) + output = extension.insert_slice(output, output_sub, [block_start + SLICE_OFFSET], [SLICE_SIZE], [1]) + tl.store(output_ptr + offsets, output, mask=mask) + + +def triton_func(x: torch.Tensor, y: torch.Tensor, slice_offset: int, slice_size: int): + output = torch.empty_like(x) + n_elements = output.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + triton_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024, SLICE_OFFSET=0, SLICE_SIZE=32) + return output + + +def test_extract_slice(): + size = 1024 + slice_offset = 0 + slice_size = 32 + x = torch.rand(size, device='npu') + y = torch.rand(size, device='npu') + torch_ref = x + y + triton_cal = triton_func(x, y, slice_offset, slice_size) + torch.testing.assert_close(triton_cal[slice_offset:slice_offset + slice_size], + torch_ref[slice_offset:slice_offset + slice_size]) diff --git a/third_party/ascend/unittest/pytest_ut/test_interleave.py b/third_party/ascend/unittest/pytest_ut/test_interleave.py new file mode 100644 index 0000000000..f78b397393 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_interleave.py @@ -0,0 +1,75 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.interleave(X, Y) + + oidx = xidx[:, None, None] * YB * ZB * 2 + yidx[None, :, None] * ZB * 2 + tl.arange(0, 2 * ZB)[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@pytest.mark.parametrize('para_type,data_type,XB,YB,ZB', [ + ['float32', torch.float32, 2, 64, 16], + ['float32', torch.float32, 8, 8, 4], + ['float16', torch.float16, 2, 64, 16], + ['float16', torch.float16, 8, 8, 4], + ['int8', torch.int8, 2, 64, 32], + ['int8', torch.int8, 8, 8, 4], +]) +def test_interleave(para_type, data_type, XB, YB, ZB): + + x = torch.full((XB, YB, ZB), 100, dtype=data_type).npu() + y = torch.full((XB, YB, ZB), 30, dtype=data_type).npu() + + print(f"shape = {x.shape}") + print(x.dtype) + + output = torch.randint(1, (XB, YB, ZB * 2), dtype=data_type).npu() + output1 = output + print(f"output.dtype={output.dtype}") + + ans = torch.stack((x, y), dim=-1).reshape(XB, YB, ZB * 2) + print(ans) + print(ans.shape) + + fn_npu_[1, 1, 1](output, x, y, XB, YB, ZB) + print(output) + + test_common.validate_cmp(para_type, ans, output) diff --git a/third_party/ascend/unittest/pytest_ut/test_interleave_optimizaiton.py b/third_party/ascend/unittest/pytest_ut/test_interleave_optimizaiton.py new file mode 100644 index 0000000000..7f491e0494 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_interleave_optimizaiton.py @@ -0,0 +1,125 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common + + +def torch_interleave_load(q, k, head_dim_half, bias): + d_indices = torch.arange(0, head_dim_half) + k[d_indices * 2 + bias] = q[d_indices * 2 + bias] + k[d_indices * 2 + 1 + bias] = -q[d_indices * 2 + 1 + bias] + return k + + +def torch_interleave_load_with_mask(q, k, head_dim_half, bias, numel): + d_indices = torch.arange(0, min(head_dim_half, numel)) + k[d_indices * 2 + bias] = q[d_indices * 2 + bias] + k[d_indices * 2 + 1 + bias] = -q[d_indices * 2 + 1 + bias] + return k + + +def torch_interleave_loadstore_with_mask(q, head_dim_half, bias, numel): + d_indices = torch.arange(0, min(head_dim_half, numel)) + # it's unneccessary since we store it back without edit: q[d_indices * 2 + bias] = q[d_indices * 2 + bias] + q[d_indices * 2 + 1 + bias] = -q[d_indices * 2 + 1 + bias] + return q + + +@triton.jit +def triton_interleave_load(q_ptr, k_ptr, head_dim_half: tl.constexpr, bias: tl.constexpr): + d_indices = tl.program_id(0) + tl.arange(0, head_dim_half) + q_real = tl.load(q_ptr + d_indices * 2 + bias) + q_imag = tl.load(q_ptr + d_indices * 2 + 1 + bias) + new_q_real = q_real + new_q_imag = -q_imag + tl.store(k_ptr + d_indices * 2 + bias, new_q_real) + tl.store(k_ptr + d_indices * 2 + 1 + bias, new_q_imag) + + +@triton.jit +def triton_interleave_load_with_mask(q_ptr, k_ptr, head_dim_half: tl.constexpr, bias: tl.constexpr, + numel: tl.constexpr): + d_indices = tl.program_id(0) + tl.arange(0, head_dim_half) + mask = d_indices < numel + q_real = tl.load(q_ptr + d_indices * 2 + bias, mask) + q_imag = tl.load(q_ptr + d_indices * 2 + 1 + bias, mask) + new_q_real = q_real + new_q_imag = -q_imag + tl.store(k_ptr + d_indices * 2 + bias, new_q_real, mask) + tl.store(k_ptr + d_indices * 2 + 1 + bias, new_q_imag, mask) + + +# when load and store are on the same pointer, sometimes we can only optimize the store with mask +@triton.jit +def triton_interleave_loadstore_with_mask(q_ptr, head_dim_half: tl.constexpr, bias: tl.constexpr, numel: tl.constexpr): + d_indices = tl.arange(0, head_dim_half) + mask = d_indices < numel + q_real = tl.load(q_ptr + d_indices * 2 + bias, mask) + q_imag = tl.load(q_ptr + d_indices * 2 + 1 + bias, mask) + new_q_real = q_real + new_q_imag = -q_imag + tl.store(q_ptr + d_indices * 2 + bias, new_q_real, mask) + tl.store(q_ptr + d_indices * 2 + 1 + bias, new_q_imag, mask) + + +@pytest.mark.parametrize('para_type,data_type,head_dim_half,bias', [ + ['float32', torch.float32, 16, 4], +]) +def test_interleave(para_type, data_type, head_dim_half, bias): + length = bias + head_dim_half * 2 + q = torch.randn((length, ), dtype=data_type).npu() + k = torch.zeros_like(q, dtype=data_type).npu() + k_ref = torch.zeros_like(q, dtype=data_type).npu() + + triton_interleave_load[(1, )](q, k, head_dim_half, bias) + k_ref = torch_interleave_load(q, k_ref, head_dim_half, bias) + assert torch.allclose(k, k_ref) + + +@pytest.mark.parametrize('para_type,data_type,head_dim_half,bias,numel', [ + ['float32', torch.float32, 16, 0, 8], +]) +def test_interleave_with_mask(para_type, data_type, head_dim_half, bias, numel): + length = bias + head_dim_half * 2 + q = torch.randn((length, ), dtype=data_type).npu() + k = torch.zeros_like(q, dtype=data_type).npu() + k_ref = torch.zeros_like(q, dtype=data_type).npu() + + triton_interleave_load_with_mask[(1, )](q, k, head_dim_half, bias, numel) + k_ref = torch_interleave_load_with_mask(q, k_ref, head_dim_half, bias, numel) + assert torch.allclose(k, k_ref) + + +@pytest.mark.parametrize('para_type,data_type,head_dim_half,bias,numel', [ + ['float32', torch.float32, 16, 0, 8], +]) +def test_interleave_loadstore_with_mask(para_type, data_type, head_dim_half, bias, numel): + length = bias + head_dim_half * 2 + q = torch.randn((length, ), dtype=data_type).npu() + q_ref = q.clone() + + triton_interleave_loadstore_with_mask[(1, )](q, head_dim_half, bias, numel) + q_ref = torch_interleave_loadstore_with_mask(q_ref, head_dim_half, bias, numel) + assert torch.allclose(q, q_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_invert.py b/third_party/ascend/unittest/pytest_ut/test_invert.py new file mode 100644 index 0000000000..1612315b43 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_invert.py @@ -0,0 +1,63 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common + + +def torch_invert(x0): + res = ~(x0) + return res + + +@triton.jit +def triton_invert(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + xindex = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = xindex < xnumel + x0 = xindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp2 = ~tmp0 + tl.store(out_ptr0 + (xindex), tmp2, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['int8', (2, 4096, 8), 2, 32768, 1024], + ['int16', (2, 4096, 8), 2, 32768, 1024], + ['int32', (2, 4096, 8), 2, 32768, 1024], + ['int64', (2, 4096, 8), 2, 32768, 1024], + ['bool', (2, 4096, 8), 2, 32768, 1024], +]) +def test_invert(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + torch_res = torch_invert(x0) + # triton结果 + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_invert[ncore, 1, 1](x0, triton_res, x0.numel(), xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_isfinited.py b/third_party/ascend/unittest/pytest_ut/test_isfinited.py new file mode 100644 index 0000000000..aa87ecb84f --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_isfinited.py @@ -0,0 +1,134 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import torch +import torch_npu +import pytest +import test_common + +types = [ + "float32", + "float16", + "bfloat16", +] + +shapes = [ + # 3, + # 32, + 37, + # 256, + # 781, +] + + +@pytest.mark.parametrize("sigtype", types) +@pytest.mark.parametrize("N", shapes) +@pytest.mark.parametrize("val", ['1.0', 'nan', 'inf', '-inf']) +def test_isfinited(sigtype, N, val): + + def torch_func(x0): + res = torch.isfinite(x0) + return res + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = tl.arange(0, N) + x0 = tl.load(in_ptr0 + idx) + ret = libdevice.isfinited(x0) + tl.store(out_ptr0 + idx, ret) + + def triton_func(x0, N): + out = torch.zeros(x0.shape, dtype=torch.bool).npu() + triton_kernel[1, 1, 1](out, x0, N) + return out + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + x0[1] = float(val) + + torch_ref = torch_func(x0) + triton_cal = triton_func(x0, N) + + test_common.validate_cmp("bool", triton_cal, torch_ref) + + +@pytest.mark.parametrize("N", shapes) +@pytest.mark.parametrize("val", ['1.0', 'nan', 'inf', '-inf']) +def test_finitef(N, val): + + def torch_func(x0): + res = torch.isfinite(x0) + return res + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = tl.arange(0, N) + x0 = tl.load(in_ptr0 + idx) + ret = libdevice.finitef(x0) + tl.store(out_ptr0 + idx, ret) + + def triton_func(x0, N): + out = torch.zeros(x0.shape, dtype=torch.bool).npu() + triton_kernel[1, 1, 1](out, x0, N) + return out + + x0 = test_common.generate_tensor(shape=(N, ), dtype='float32').npu() + x0[1] = float(val) + + torch_ref = torch_func(x0) + triton_cal = triton_func(x0, N) + test_common.validate_cmp("bool", triton_cal, torch_ref) + + +invalid_types = [ + "int32", +] + + +@pytest.mark.parametrize("sigtype", invalid_types) +@pytest.mark.parametrize("N", shapes) +@test_common.raises_with_match(triton.compiler.errors.CompilationError, "Expected dtype fp16/fp32/bf16, but got") +def test_isfinited_invalid_dtype(sigtype, N): + + def torch_func(x0): + res = torch.isfinite(x0) + return res + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = tl.arange(0, N) + x0 = tl.load(in_ptr0 + idx) + ret = libdevice.isfinited(x0) + tl.store(out_ptr0 + idx, ret) + + def triton_func(x0, N): + out = torch.zeros(x0.shape, dtype=torch.bool).npu() + triton_kernel[1, 1, 1](out, x0, N) + return out + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + x0[1] = float('nan') + + torch_ref = torch_func(x0) + triton_cal = triton_func(x0, N) + test_common.validate_cmp("bool", triton_cal, torch_ref) + assert triton_cal[1] == True diff --git a/third_party/ascend/unittest/pytest_ut/test_isnan.py b/third_party/ascend/unittest/pytest_ut/test_isnan.py new file mode 100644 index 0000000000..dedf3d903b --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_isnan.py @@ -0,0 +1,107 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + +types = [ + "float32", + "float16", + "bfloat16", +] + +shapes = [ + # 3, + # 32, + 37, + # 256, + # 781, +] + + +@pytest.mark.parametrize("sigtype", types) +@pytest.mark.parametrize("N", shapes) +def test_isnan(sigtype, N): + + def torch_func(x0): + res = torch.isnan(x0) + return res + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = tl.arange(0, N) + x0 = tl.load(in_ptr0 + idx) + ret = tl.extra.cann.libdevice.isnan(x0) + tl.store(out_ptr0 + idx, ret) + + def triton_func(x0, N): + out = torch.zeros(x0.shape, dtype=torch.bool).npu() + triton_kernel[1, 1, 1](out, x0, N) + return out + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + x0[1] = float('nan') + + torch_ref = torch_func(x0) + triton_cal = triton_func(x0, N) + test_common.validate_cmp("bool", triton_cal, torch_ref) + assert triton_cal[1] == True + + +invalid_types = [ + "int32", +] + + +@pytest.mark.parametrize("sigtype", invalid_types) +@pytest.mark.parametrize("N", shapes) +def test_isnan_invalid_dtype(sigtype, N): + + def torch_func(x0): + res = torch.isnan(x0) + return res + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = tl.arange(0, N) + x0 = tl.load(in_ptr0 + idx) + ret = tl.extra.cann.libdevice.isnan(x0) + tl.store(out_ptr0 + idx, ret) + + def triton_func(x0, N): + out = torch.zeros(x0.shape, dtype=torch.bool).npu() + triton_kernel[1, 1, 1](out, x0, N) + return out + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + x0[1] = float('nan') + flag = False + try: + torch_ref = torch_func(x0) + triton_cal = triton_func(x0, N) + test_common.validate_cmp("bool", triton_cal, torch_ref) + assert triton_cal[1] == True + except Exception as e: + flag = True + assert flag diff --git a/third_party/ascend/unittest/pytest_ut/test_join.py b/third_party/ascend/unittest/pytest_ut/test_join.py new file mode 100644 index 0000000000..293cf927ae --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_join.py @@ -0,0 +1,66 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, y_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + + idx = xidx[:, None] * YB + yidx[None, :] + + X = tl.load(x_ptr + idx) + Y = tl.load(y_ptr + idx) + + ret = tl.join(X, Y) + + oidx = xidx[:, None, None] * YB * 2 + yidx[None, :, None] * 2 + tl.arange(0, 2)[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@pytest.mark.parametrize('para_type,data_type,XB,YB,ZB', [ + ['float32', torch.float32, 4, 64, 4], + ['float32', torch.float32, 8, 8, 4], + ['float16', torch.float16, 4, 64, 4], + ['float16', torch.float16, 8, 8, 4], + ['int8', torch.int8, 4, 128, 4], + ['int8', torch.int8, 8, 8, 4], +]) +def test_join(para_type, data_type, XB, YB, ZB): + x = torch.full((XB, YB), 100, dtype=data_type).npu() + y = torch.full((XB, YB), 30, dtype=data_type).npu() + + ans = torch.stack((x, y), dim=-1) + print(ans) + + output = torch.randint(1, (XB, YB, 2), dtype=data_type).npu() + fn_npu_[1, 1, 1](output, x, y, XB, YB, ZB, debug=True) + + print(output) + test_common.validate_cmp(para_type, ans, output) diff --git a/third_party/ascend/unittest/pytest_ut/test_lanzcos.py b/third_party/ascend/unittest/pytest_ut/test_lanzcos.py new file mode 100644 index 0000000000..6e20d3c083 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_lanzcos.py @@ -0,0 +1,246 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch, torch_npu +import triton +import triton.language as tl +import numpy as np +import math +import pytest + + +@triton.jit +def lanczos_resize_kernel( + img_src_ptr, + img_dst_ptr, + img_coeffs_ptr, + src_rows, + src_cols, + dst_rows, + dst_cols, + R_H, + R_W, + C, + stride_in_h, + stride_in_w, + stride_in_c, + stride_out_h, + stride_out_w, + stride_out_c, + BLOCK_SIZE: tl.constexpr, +): + block_id_c = tl.program_id(0) + block_id_h = tl.program_id(1) + block_id_w = tl.program_id(2) + dest_h_offs = block_id_h * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + dest_w_offs = block_id_w * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + dest_offs = (block_id_c[None, None] * stride_out_c + dest_h_offs[:, None] * stride_out_h + + dest_w_offs[None, :] * stride_out_w) + + RR_H = 1.0 / R_H + RR_W = 1.0 / R_W + + fy = (dest_h_offs + 0.5) * RR_H - 0.5 + sy = tl.floor(fy) + fx = (dest_w_offs + 0.5) * RR_W - 0.5 + sx = tl.floor(fx) + + idxY = tl.floor((fy - sy) * 24.999999).to(tl.int32) + idxX = tl.floor((fx - sx) * 24.999999).to(tl.int32) + tableIndex = idxY[:, None] * 25 + idxX[None, :] + res = tl.zeros((BLOCK_SIZE, BLOCK_SIZE), tl.float32) + + for ii in range(4): + for jj in range(4): + src_offsets = (block_id_c[None, None] * stride_in_c + (tl.clamp( + (sy + ii - 1), 0, src_rows - 1)).to(tl.int32)[:, None] * stride_in_h + (tl.clamp( + (sx + jj - 1), 0, src_cols - 1)).to(tl.int32)[None, :] * stride_in_w) + src_val = tl.load(img_src_ptr + src_offsets) + coeffs_offs = tableIndex[:, :] * 16 + (ii * 4 + jj)[None, None] + coeffs = tl.load(img_coeffs_ptr + coeffs_offs) + res = res + src_val * coeffs + dst_mask = (dest_h_offs[:, None] < dst_rows) & (dest_w_offs[None, :] < dst_cols) + res = tl.clamp(res, 0.0, 1.0) + tl.store(img_dst_ptr + dest_offs, res, mask=dst_mask) + + +def lanczos_resize_triton(img_src, img_dst, c_lanczosCoeffs, dst_rows, dst_cols): + N, C, src_rows, src_cols = img_src.shape + R_H = float(dst_rows) / src_rows + R_W = float(dst_cols) / src_cols + + stride_in_n, stride_in_c, stride_in_h, stride_in_w = img_src.stride() + stride_out_n, stride_out_c, stride_out_h, stride_out_w = img_dst.stride() + BLOCK_SIZE = 16 + grid = lambda meta: ( + C, + triton.cdiv(dst_rows, meta["BLOCK_SIZE"]), + triton.cdiv(dst_cols, meta["BLOCK_SIZE"]), + ) + lanczos_resize_kernel[grid]( + img_src, + img_dst, + c_lanczosCoeffs, + src_rows, + src_cols, + dst_rows, + dst_cols, + R_H, + R_W, + C, + stride_in_h, + stride_in_w, + stride_in_c, + stride_out_h, + stride_out_w, + stride_out_c, + BLOCK_SIZE=BLOCK_SIZE, + ) + return img_dst + + +def lanczos_resize_cpu(img_src, img_dst, img_coeffs, dst_rows, dst_cols): + N, C, src_rows, src_cols = img_src.shape + R_H = float(dst_rows) / src_rows + R_W = float(dst_cols) / src_cols + for i in range(dst_rows): + for j in range(dst_cols): + RR_H = 1.0 / R_H + RR_W = 1.0 / R_W + fy = (i + 0.5) * RR_H - 0.5 + sy = math.floor(fy) + fx = (j + 0.5) * RR_W - 0.5 + sx = math.floor(fx) + idxY = math.floor((fy - np.floor(fy)) * 24.999999) + idxX = math.floor((fx - np.floor(fx)) * 24.999999) + tableIndex = idxY * 25 + idxX + res = (0.0, 0.0, 0.0, 0.0) + for ii in range(4): + for jj in range(4): + idx_y = np.clip(sy + ii - 1, 0, src_rows - 1) + idx_x = np.clip(sx + jj - 1, 0, src_cols - 1) + src_val = img_src[0, :, idx_y, idx_x] + coeffs_offs = tableIndex * 16 + (ii * 4 + jj) + coeffs = img_coeffs[coeffs_offs] + res = res + src_val * coeffs + + img_dst[0, :, i, j] = np.clip(res, 0.0, 1.0) + + +@pytest.mark.parametrize("shapes", [ + [360, 640, 140, 280], +]) +def test_lanzcos(shapes): + c_lanczosCoeffs = torch.randn(10000, dtype=torch.float32, device="npu") / 4.0 + src_rows, src_cols, dst_rows, dst_cols = shapes + img_src = torch.randn(1, 4, src_rows, src_cols, dtype=torch.float32, device="npu") + img_dst = torch.zeros( + (1, img_src.shape[1], dst_rows, dst_cols), + dtype=img_src.dtype, + device=img_src.device, + ) + resized_image = lanczos_resize_triton(img_src, img_dst, c_lanczosCoeffs, dst_rows, dst_cols) + img_src_cpu = img_src.cpu().numpy() + img_dst_cpu = torch.zeros((1, img_src_cpu.shape[1], dst_rows, dst_cols), dtype=img_src.dtype, device="cpu").numpy() + lanczos_resize_cpu(img_src_cpu, img_dst_cpu, c_lanczosCoeffs.cpu().numpy(), dst_rows, dst_cols) + torch.testing.assert_close(resized_image.cpu(), torch.from_numpy(img_dst_cpu), atol=1.0 / 255, rtol=0) + + +def benchmark_test(fn_ref, fn_triton, ref_args=(), triton_args=(), name="gen_fn", times=10, repeat=10): + import time + + print(f"--------------------benchmark_{name} for {times * repeat} times--------------------") + stream = torch.npu.current_stream() + # warm_up + stream.synchronize() + for _ in range(10): + fn_triton(*triton_args) + stream.synchronize() + + start = time.perf_counter() + for _ in range(times * repeat): + fn_triton(*triton_args) + stream.synchronize() + end = time.perf_counter() + + time_compiled = (end - start) / (times * repeat) + time_compiled *= 1000000 + + # warm_up + stream.synchronize() + for _ in range(10): + std = fn_ref(*ref_args) + stream.synchronize() + + start = time.perf_counter() + for _ in range(times * repeat): + std = fn_ref(*ref_args) + stream.synchronize() + end = time.perf_counter() + time_eager = (end - start) / (times * repeat) + time_eager *= 1000000 + + accelerated = (time_eager - time_compiled) / time_compiled * 100 + print(f"Accelerated: {accelerated:.4f}% eager takes {time_eager:.3f} us, triton takes {time_compiled:.3f} us") + + return accelerated, time_eager, time_compiled + + +if __name__ == "__main__": + c_lanczosCoeffs = torch.randn(10000, dtype=torch.float32, device="npu") / 4.0 + + src_rows, src_cols = 360, 640 + dst_rows, dst_cols = 140, 280 + img_src = torch.randn(1, 4, src_rows, src_cols, dtype=torch.float32, device="npu") + + print("==========run npu===============") + img_dst = torch.zeros( + (1, img_src.shape[1], dst_rows, dst_cols), + dtype=img_src.dtype, + device=img_src.device, + ) + resized_image = lanczos_resize_triton(img_src, img_dst, c_lanczosCoeffs, dst_rows, dst_cols) + resized_cpu = resized_image.cpu().numpy() + print("==========run cpu===============") + img_src_cpu = img_src.cpu().numpy() + img_dst_cpu = torch.zeros((1, img_src_cpu.shape[1], dst_rows, dst_cols), dtype=img_src.dtype, device="cpu").numpy() + lanczos_resize_cpu(img_src_cpu, img_dst_cpu, c_lanczosCoeffs.cpu().numpy(), dst_rows, dst_cols) + + print("==========compare result===============") + diff = np.abs(resized_cpu - img_dst_cpu) + max_diff_value = np.max(diff) + print("max diff float = ", max_diff_value) + print("max diff * 255 int = ", int(max_diff_value * 255)) + torch.testing.assert_close(resized_image.cpu(), torch.from_numpy(img_dst_cpu), atol=1.0 / 255, rtol=0) + + print("==========profiling===============") + accelerate, eager_time, triton_time = benchmark_test( + lanczos_resize_cpu, + lanczos_resize_triton, + ref_args=( + img_src_cpu, + img_dst_cpu, + c_lanczosCoeffs.cpu().numpy(), + dst_rows, + dst_cols, + ), + triton_args=(img_src, img_dst, c_lanczosCoeffs, dst_rows, dst_cols), + name="lanzcos", + ) diff --git a/third_party/ascend/unittest/pytest_ut/test_launcher_empty_signature.py b/third_party/ascend/unittest/pytest_ut/test_launcher_empty_signature.py new file mode 100644 index 0000000000..a9414ca661 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_launcher_empty_signature.py @@ -0,0 +1,37 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import os +import pytest + +import triton +import triton.language as tl + + +@triton.jit +def _empty_kernel(): + return + + +@pytest.mark.interpreter +def test_launcher_empty_signature(): + grid = (1, ) + _empty_kernel[grid]() + assert True diff --git a/third_party/ascend/unittest/pytest_ut/test_layernorm.py b/third_party/ascend/unittest/pytest_ut/test_layernorm.py new file mode 100644 index 0000000000..328bae5a0d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_layernorm.py @@ -0,0 +1,126 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import torch +import triton +import triton.language as tl +import torch_npu + + +@triton.jit +def _layer_norm_fwd_fused( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride, # how much to increase the pointer when moving by 1 row + N, # number of columns in X + eps, # epsilon to avoid division by zero + BLOCK_SIZE: tl.constexpr, +): + # Map the program id to the row of X and Y it should compute. + row = tl.program_id(0) + Y += row * stride + X += row * stride + # Compute mean + mean = 0 + _mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + a = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32) + _mean += a + mean = tl.sum(_mean, axis=0) / N + # Compute variance + _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + x = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32) + x = tl.where(cols < N, x - mean, 0.) + _var += x * x + var = tl.sum(_var, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + # Write mean / rstd + tl.store(Mean + row, mean) + tl.store(Rstd + row, rstd) + # Normalize and apply linear transformation + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + mask = cols < N + w = tl.load(W + cols, mask=mask) + b = tl.load(B + cols, mask=mask) + x = tl.load(X + cols, mask=mask, other=0.).to(tl.float32) + x_hat = (x - mean) * rstd + y = x_hat * w + b + # Write output + tl.store(Y + cols, y, mask=mask) + + +@torch.inference_mode() +def layer_norm(x, normalized_shape, weight, bias, eps=1e-5): + # allocate output + y = torch.empty_like(x) + # reshape input data into 2D tensor + x_arg = x.reshape(-1, x.shape[-1]) + M, N = x_arg.shape + mean = torch.empty((M, ), dtype=torch.float32, device=x.device) + rstd = torch.empty((M, ), dtype=torch.float32, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + # heuristics for number of warps + num_warps = min(max(BLOCK_SIZE // 256, 1), 8) + # enqueue kernel + kernel = _layer_norm_fwd_fused[(M, )]( # + x_arg, y, weight, bias, mean, rstd, # + x_arg.stride(0), N, eps, # + BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_ctas=1) + # print(kernel.asm['ttir']) + return y + + +def _layer_norm(M, N, dtype, eps=1e-5, device='npu'): + # create data + x_shape = (M, N) + w_shape = (x_shape[-1], ) + weight = torch.rand(w_shape, dtype=dtype, device=device, requires_grad=True) + bias = torch.rand(w_shape, dtype=dtype, device=device, requires_grad=True) + x = -2.3 + 0.5 * torch.randn(x_shape, dtype=dtype, device=device) + dy = .1 * torch.randn_like(x) + x.requires_grad_(True) + # forward pass + y_tri = layer_norm(x, w_shape, weight, bias, eps) + y_ref = torch.nn.functional.layer_norm(x, w_shape, weight, bias, eps).to(dtype) + # compare + assert torch.allclose(y_tri, y_ref, atol=1e-2, rtol=0) + print(f"layernorm {M},{N} {dtype} passed") + + +def test_layernorm(): + _layer_norm(128, 128, torch.float16) + _layer_norm(128, 128, torch.bfloat16) + _layer_norm(128, 128, torch.float32) + + # _layer_norm(128, 3, torch.bfloat16) + # _layer_norm(128, 16, torch.bfloat16) + # _layer_norm(128, 37, torch.bfloat16) + # _layer_norm(128, 781, torch.bfloat16) diff --git a/third_party/ascend/unittest/pytest_ut/test_ldst.py b/third_party/ascend/unittest/pytest_ut/test_ldst.py new file mode 100644 index 0000000000..c5b989fa70 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_ldst.py @@ -0,0 +1,449 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch, torch_npu +import triton +import triton.language as tl +import triton.language.math as tl_math +import pytest + + +def test_ldst_indirect_00(): + + @triton.jit + def triton_ldst_indirect_00_kernel(out_ptr0, in_ptr0, in_ptr1, OFFSET0: tl.constexpr, XS: tl.constexpr): + pid = tl.program_id(0) + offset1 = tl.load(in_ptr0 + OFFSET0) + idx_in1 = offset1 + pid * XS + tl.arange(0, XS) + tmp0 = tl.load(in_ptr1 + idx_in1, cache_modifier=".ca", eviction_policy="evict_first", volatile=True) + tmp1 = tl_math.exp(tmp0) + idx_out0 = pid * XS + tl.arange(0, XS) + tl.store(out_ptr0 + idx_out0, tmp1) + + def triton_ldst_indirect_00_func(x0, x1, s, xs): + n = x1.numel() + ns = n - s + assert ns == xs, "test only single core" + y0 = torch.empty((ns, ), dtype=x1.dtype, device=x1.device) + triton_ldst_indirect_00_kernel[ns // xs, 1, 1](y0, x0, x1, OFFSET0=s, XS=xs) + return y0 + + def torch_ldst_indirect_00_func(x0, x1, s): + offset = x0[s] + return torch.exp(x1[offset:]) + + DEV = "npu" + DTYPE = torch.float32 + offset = 0 + N0, N1 = 16, 16 + blocksize = 16 + assert N0 > offset, "offset must be < N0" + N1 = N1 + offset + x0 = torch.arange(0, N0, dtype=torch.int32, device=DEV) + x1 = torch.randn((N1, ), dtype=DTYPE, device=DEV) + torch_ref = torch_ldst_indirect_00_func(x0, x1, offset) + triton_cal = triton_ldst_indirect_00_func(x0, x1, offset, blocksize) + torch.testing.assert_close(triton_cal, torch_ref) + + +def test_ldst_indirect_01(): + + @triton.jit + def triton_ldst_indirect_01_kernel(out_ptr0, in_ptr0, in_ptr1, OFFSET0: tl.constexpr, XS: tl.constexpr): + pid = tl.program_id(0) + offset1 = tl.load(in_ptr0 + OFFSET0) + idx_in1 = offset1 + pid * XS + tl.arange(0, XS) + tmp0 = tl.load(in_ptr1 + idx_in1) + tmp1 = tl_math.exp(tmp0) + idx_out0 = pid * XS + tl.arange(0, XS) + tl.store(out_ptr0 + idx_out0, tmp1) + + def triton_ldst_indirect_01_func(x0, x1, s, xs): + n = x1.numel() + ns = n - s + assert ns == xs, "test only single core" + y0 = torch.empty((ns, ), dtype=x1.dtype, device=x1.device) + triton_ldst_indirect_01_kernel[ns // xs, 1, 1](y0, x0, x1, OFFSET0=s, XS=xs) + return y0 + + def torch_ldst_indirect_01_func(x0, x1, s): + offset = x0[s] + return torch.exp(x1[offset:]) + + DEV = "npu" + DTYPE = torch.float32 + offset = 0 + N0, N1 = 16, 16 + blocksize = 16 + assert N0 > offset, "offset must be < N0" + N1 = N1 + offset + x0 = torch.arange(0, N0, device=DEV) # int64 + x1 = torch.randn((N1, ), dtype=DTYPE, device=DEV) + torch_ref = torch_ldst_indirect_01_func(x0, x1, offset) + triton_cal = triton_ldst_indirect_01_func(x0, x1, offset, blocksize) + torch.testing.assert_close(triton_cal, torch_ref) + + +def test_ldst_indirect_02(): + + @triton.jit + def triton_ldst_indirect_02_kernel(out_ptr0, in_ptr0, in_ptr1, XS: tl.constexpr): + pid = tl.program_id(0) + for i in tl.range(0, XS): + tmp0 = tl.load(in_ptr0 + i) + tmp1 = tl.load(in_ptr1 + tmp0) + tmp2 = tl_math.exp(tmp1) + tl.store(out_ptr0 + i, tmp2) + + def triton_ldst_indirect_02_func(x0, x1, xs): + n0 = x0.numel() + assert n0 == xs, "test only single core" + y0 = torch.empty((n0, ), dtype=x1.dtype, device=x1.device) + triton_ldst_indirect_02_kernel[n0 // xs, 1, 1](y0, x0, x1, XS=xs) + return y0 + + def torch_ldst_indirect_02_func(x0, x1): + return torch.exp(x1[x0]) + + DEV = "npu" + DTYPE = torch.float32 + offset = 8 + N0, N1 = 16, 32 + blocksize = 16 + assert N1 >= N0 + offset, "N1 must be >= N0+offset" + assert N0 == blocksize, "N0 must be == blocksize" + x0 = offset + torch.arange(0, N0, device=DEV) # int64 + x1 = torch.randn((N1, ), dtype=DTYPE, device=DEV) + torch_ref = torch_ldst_indirect_02_func(x0, x1) + triton_cal = triton_ldst_indirect_02_func(x0, x1, blocksize) + torch.testing.assert_close(triton_cal, torch_ref) + + +def test_ldst_indirect_03(): + + @triton.jit + def triton_ldst_indirect_03_kernel(out_ptr0, in_ptr0, in_ptr1, XS: tl.constexpr): + pid = tl.program_id(0) + in_idx0 = pid * XS + tl.arange(0, XS) + tmp0 = tl.load(in_ptr0 + in_idx0) + tmp1 = tl.load(in_ptr1 + tmp0) + tmp2 = tl_math.exp(tmp1) + out0_idx = pid * XS + tl.arange(0, XS) + tl.store(out_ptr0 + out0_idx, tmp2) + + def triton_ldst_indirect_03_func(x0, x1, xs): + n0 = x0.numel() + assert n0 == xs, "test only single core" + y0 = torch.empty((n0, ), dtype=x1.dtype, device=x1.device) + triton_ldst_indirect_03_kernel[n0 // xs, 1, 1](y0, x0, x1, XS=xs) + return y0 + + def torch_ldst_indirect_03_func(x0, x1): + return torch.exp(x1[x0]) + + DEV = "npu" + DTYPE = torch.float32 + offset = 8 + N0, N1 = 16, 32 + blocksize = 16 + assert N1 >= N0 + offset, "N1 must be >= N0+offset" + assert N0 == blocksize, "N0 must be == blocksize" + x0 = offset + torch.arange(0, N0, device=DEV) # int64 + x1 = torch.randn((N1, ), dtype=DTYPE, device=DEV) + torch_ref = torch_ldst_indirect_03_func(x0, x1) + triton_cal = triton_ldst_indirect_03_func(x0, x1, blocksize) + torch.testing.assert_close(triton_cal, torch_ref) + + +def test_ldst_indirect_04(): + + @triton.jit + def triton_ldst_indirect_04_kernel(out_ptr0, in_ptr0, in_ptr1, XS: tl.constexpr): + pid = tl.program_id(0) + in_idx0 = pid * XS + tl.arange(0, XS) + tmp0 = tl.load(in_ptr0 + in_idx0) + tmp0min = tl.min(tmp0, axis=0) + tmp0max = tl.max(tmp0, axis=0) + tmp0 = tmp0 * 2.0 + tmp0 = tl.clamp(tmp0, tmp0min, tmp0max) + tmp0 = tmp0.to(tl.int32) + tmp1 = tl.load(in_ptr1 + tmp0) + tmp2 = tl_math.exp(tmp1) + out0_idx = pid * XS + tl.arange(0, XS) + tl.store(out_ptr0 + out0_idx, tmp2) + + def triton_ldst_indirect_04_func(x0, x1, xs): + n0 = x0.numel() + assert n0 == xs, "test only single core" + y0 = torch.empty((n0, ), dtype=x1.dtype, device=x1.device) + triton_ldst_indirect_04_kernel[n0 // xs, 1, 1](y0, x0, x1, XS=xs) + return y0 + + def torch_ldst_indirect_04_func(x0, x1): + x0min = torch.min(x0) + x0max = torch.max(x0) + idx = torch.clamp(x0 * 2, x0min, x0max) + return torch.exp(x1[idx.to(torch.int32)]) + + DEV = "npu" + DTYPE = torch.float32 + offset = 8 + N0, N1 = 16, 32 + blocksize = 16 + assert N1 >= N0 + offset, "N1 must be >= N0+offset" + assert N0 == blocksize, "N0 must be == blocksize" + x0 = offset + torch.arange(0, N0, dtype=torch.float32, device=DEV) + x1 = torch.randn((N1, ), dtype=DTYPE, device=DEV) + torch_ref = torch_ldst_indirect_04_func(x0, x1) + triton_cal = triton_ldst_indirect_04_func(x0, x1, blocksize) + torch.testing.assert_close(triton_cal, torch_ref) + + +def test_ldst_indirect_05(): + + @triton.jit + def triton_ldst_indirect_05_kernel(out_ptr0, in_ptr1, in_ptr2, stride_in_r, XS: tl.constexpr, RS: tl.constexpr): + pid = tl.program_id(0) + in_idx0 = pid * XS + tl.arange(0, XS) + in_idx1 = tl.arange(0, RS) + tmp0 = tl.arange(0, XS) + tmp1 = tl.load(in_ptr1 + in_idx1) + in_idx2 = tmp0[:, None] * stride_in_r + tmp1[None, :] + tmp2 = tl.load(in_ptr2 + in_idx2) + tmp2 = tl_math.exp(tmp2) + out0_idx = in_idx0[:, None] * RS + in_idx1[None, :] + tl.store(out_ptr0 + out0_idx, tmp2) + + def triton_ldst_indirect_05_func(xc, x2, xs, rs): + nr = x2.size()[0] + nc = xc.numel() + stride_in_r = x2.stride()[0] + assert nr == xs, "test only single core" + y0 = torch.empty((nr, nc), dtype=x2.dtype, device=x2.device) + triton_ldst_indirect_05_kernel[nr // xs, 1, 1](y0, xc, x2, stride_in_r, XS=xs, RS=rs) + return y0 + + def torch_ldst_indirect_05_func(xr, xc, x2): + flatten_idx = (xr[:, None] * x2.stride()[0] + xc[None, :]).flatten() + extracted = x2.flatten()[flatten_idx].reshape([xr.numel(), xc.numel()]) + return torch.exp(extracted) + + DEV = "npu" + DTYPE = torch.float32 + offset = 8 + N0, N1 = 16, 32 + blocksize = 8 + lowdimsize = N0 + assert N1 >= N0 + offset, "N1 must be >= N0+offset" + assert N0 == lowdimsize, "N0 must be == lowdimsize" + xc = offset + torch.arange(0, N0, device=DEV) + xr = torch.arange(0, blocksize, device=DEV) + x2 = torch.randn((blocksize, N1), dtype=DTYPE, device=DEV) + torch_ref = torch_ldst_indirect_05_func(xr, xc, x2) + triton_cal = triton_ldst_indirect_05_func(xc, x2, blocksize, lowdimsize) + torch.testing.assert_close(triton_cal, torch_ref) + + +def test_ldst_indirect_06(): + + @triton.jit + def triton_ldst_indirect_06_kernel(out_ptr0, in_ptr0, in_ptr1, in_ptr2, stride_in_r, XS: tl.constexpr, + RS: tl.constexpr): + pid = tl.program_id(0) + in_idx0 = pid * XS + tl.arange(0, XS) + in_idx1 = tl.arange(0, RS) + tmp0 = tl.load(in_ptr0 + in_idx0) + tmp1 = tl.load(in_ptr1 + in_idx1) + in_idx2 = tmp0[:, None] * stride_in_r + tmp1[None, :] + tmp2 = tl.load(in_ptr2 + in_idx2) + tmp2 = tl_math.exp(tmp2) + out0_idx = in_idx0[:, None] * RS + in_idx1[None, :] + tl.store(out_ptr0 + out0_idx, tmp2) + + def triton_ldst_indirect_06_func(xr, xc, x2, xs, rs): + nr = x2.size()[0] + nc = xc.numel() + stride_in_r = x2.stride()[0] + assert nr == xs, "test only single core" + y0 = torch.empty((nr, nc), dtype=x2.dtype, device=x2.device) + triton_ldst_indirect_06_kernel[nr // xs, 1, 1](y0, xr, xc, x2, stride_in_r, XS=xs, RS=rs) + return y0 + + def torch_ldst_indirect_06_func(xr, xc, x2): + flatten_idx = (xr[:, None] * x2.stride()[0] + xc[None, :]).flatten() + extracted = x2.flatten()[flatten_idx].reshape([xr.numel(), xc.numel()]) + return torch.exp(extracted) + + DEV = "npu" + DTYPE = torch.float32 + offset = 8 + N0, N1 = 16, 32 + blocksize = 4 + lowdimsize = N0 + assert N1 >= N0 + offset, "N1 must be >= N0+offset" + assert N0 == lowdimsize, "N0 must be == lowdimsize" + xc = offset + torch.arange(0, N0, device=DEV) + xr = torch.arange(0, blocksize, device=DEV) + x2 = torch.randn((blocksize, N1), dtype=DTYPE, device=DEV) + torch_ref = torch_ldst_indirect_06_func(xr, xc, x2) + triton_cal = triton_ldst_indirect_06_func(xr, xc, x2, blocksize, lowdimsize) + torch.testing.assert_close(triton_cal, torch_ref) + + +def test_ldst_indirect_07(): + + @triton.jit + def triton_ldst_indirect_07_kernel(out_ptr0, in_ptr0, in_ptr1, in_ptr2, stride_in_r, XS: tl.constexpr, + RS: tl.constexpr): + pid = tl.program_id(0) + in_idx0 = pid * XS + tl.arange(0, XS) + in_idx1 = tl.arange(0, RS) + tmp0 = tl.load(in_ptr0 + in_idx0) + tmp1 = tl.load(in_ptr1 + in_idx1) + in_idx2 = tmp0[:, None] * stride_in_r + tmp1[None, :] + tmp2 = tl.load(in_ptr2 + in_idx2) + out0_idx = in_idx0[:, None] * RS + in_idx1[None, :] + tl.store(out_ptr0 + out0_idx, tmp2) + + def triton_ldst_indirect_07_func(xr, xc, x2, xs, rs): + nr = x2.size()[0] + nc = xc.numel() + stride_in_r = x2.stride()[0] + assert nr == xs, "test only single core" + y0 = torch.empty((nr, nc), dtype=x2.dtype, device=x2.device) + triton_ldst_indirect_07_kernel[nr // xs, 1, 1](y0, xr, xc, x2, stride_in_r, XS=xs, RS=rs) + return y0 + + def torch_ldst_indirect_07_func(xr, xc, x2): + flatten_idx = (xr[:, None] * x2.stride()[0] + xc[None, :]).flatten() + extracted = x2.flatten()[flatten_idx].reshape([xr.numel(), xc.numel()]) + return extracted + + DEV = "npu" + DTYPE = torch.float32 + offset = 8 + N0, N1 = 16, 32 + blocksize = 4 + lowdimsize = N0 + assert N1 >= N0 + offset, "N1 must be >= N0+offset" + assert N0 == lowdimsize, "N0 must be == lowdimsize" + xc = offset + torch.arange(0, N0, device=DEV) + xr = torch.arange(0, blocksize, device=DEV) + x2 = torch.randn((blocksize, N1), dtype=DTYPE, device=DEV) + torch_ref = torch_ldst_indirect_07_func(xr, xc, x2) + triton_cal = triton_ldst_indirect_07_func(xr, xc, x2, blocksize, lowdimsize) + torch.testing.assert_close(triton_cal, torch_ref) + + +def test_ldst_indirect_08(): + + @triton.jit + def triton_ldst_indirect_08_kernel(out_ptr0, in_ptr_xc, in_ptr_x2, stride_in_r, OUT_COLS: tl.constexpr, + XS: tl.constexpr, RS: tl.constexpr): + pid = tl.program_id(0) + row_idx_full = pid * XS + tl.arange(0, XS) + col_pos = tl.arange(0, RS) + xc_vals = tl.load(in_ptr_xc + col_pos) + row_arange = tl.arange(0, XS) + gather_flat = row_arange[:, None] * stride_in_r + xc_vals[None, :] + vals = tl.load(in_ptr_x2 + gather_flat) + vals = tl_math.exp(vals) + out_flat = row_idx_full[:, None] * OUT_COLS + xc_vals[None, :] + tl.store(out_ptr0 + out_flat, vals) + + def triton_ldst_indirect_08_func(xc, x2, xs, rs): + nr = x2.size(0) + out_cols = x2.size(1) + stride_in_r = x2.stride(0) + assert nr == xs, "test only single core" + y0 = torch.zeros((nr, out_cols), dtype=x2.dtype, device=x2.device) + triton_ldst_indirect_08_kernel[nr // xs, 1, 1](y0, xc, x2, stride_in_r, OUT_COLS=out_cols, XS=xs, RS=rs) + return y0 + + def torch_ldst_indirect_08_func(xr, xc, x2): + out = torch.zeros((xr.numel(), x2.size(1)), dtype=x2.dtype, device=x2.device) + gathered = torch.exp(x2[xr[:, None], xc[None, :]]) + out.scatter_(1, xc.expand(xr.numel(), -1), gathered) + return out + + DEV = "npu" + DTYPE = torch.float32 + offset = 8 + N0, N1 = 16, 32 + blocksize = 8 + lowdimsize = N0 + assert N1 >= N0 + offset, "N1 must be >= N0+offset" + assert N0 == lowdimsize, "N0 must be == lowdimsize" + xc = offset + torch.arange(0, N0, device=DEV) + xr = torch.arange(0, blocksize, device=DEV) + x2 = torch.randn((blocksize, N1), dtype=DTYPE, device=DEV) + torch_ref = torch_ldst_indirect_08_func(xr, xc, x2) + triton_cal = triton_ldst_indirect_08_func(xc, x2, blocksize, lowdimsize) + torch.testing.assert_close(triton_cal, torch_ref) + + +def test_ldst_indirect_09(): + + @triton.jit + def triton_ldst_indirect_09_kernel(out_ptr0, in_ptr1, in_ptr2, stride_in_r, offset: tl.constexpr, XS: tl.constexpr, + RS: tl.constexpr): + pid = tl.program_id(0) + in_idx0 = tl.arange(0, XS) + in_idx1 = tl.arange(0, RS) + tmp0 = pid * XS + tl.load(in_ptr1 + in_idx0) + tmp1 = tl.arange(0, RS) + offset + in_idx2 = tmp0[:, None] * stride_in_r + tmp1[None, :] + tmp2 = tl.load(in_ptr2 + in_idx2) + tmp2 = tl_math.exp(tmp2) + out0_idx = pid * XS * RS + in_idx0[:, None] * RS + in_idx1[None, :] + tl.store(out_ptr0 + out0_idx, tmp2) + + def triton_ldst_indirect_09_func(xr, x2, offset, xs, rs): + nr = xr.numel() + nc = rs + stride_in_r = x2.stride()[0] + y0 = torch.empty((nr, nc), dtype=x2.dtype, device=x2.device) + triton_ldst_indirect_09_kernel[nr // xs, 1, 1](y0, xr, x2, stride_in_r, offset=offset, XS=xs, RS=rs) + return y0 + + def torch_ldst_indirect_09_func(xr, xc, x2): + flatten_idx = (xr[:, None] * x2.stride()[0] + xc[None, :]).flatten() + extracted = x2.flatten()[flatten_idx].reshape([xr.numel(), xc.numel()]) + return torch.exp(extracted) + + DEV = "npu" + DTYPE = torch.float32 + offset = 8 + N0, N1 = 16, 32 + blocksize = 8 + lowdimsize = N0 + assert N1 >= N0 + offset, "N1 must be >= N0+offset" + assert N0 == lowdimsize, "N0 must be == lowdimsize" + xc = offset + torch.arange(0, N0, device=DEV) + xr = torch.arange(0, blocksize, device=DEV) + x2 = torch.randn((blocksize, N1), dtype=DTYPE, device=DEV) + torch_ref = torch_ldst_indirect_09_func(xr, xc, x2) + triton_cal = triton_ldst_indirect_09_func(xr, x2, offset, blocksize, lowdimsize) + torch.testing.assert_close(triton_cal, torch_ref) + + +if __name__ == "__main__": + test_ldst_indirect_05() + print("success: test_ldst_indirect_05") diff --git a/third_party/ascend/unittest/pytest_ut/test_le.py b/third_party/ascend/unittest/pytest_ut/test_le.py new file mode 100644 index 0000000000..b03e6fa7e5 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_le.py @@ -0,0 +1,79 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + + +def standard_binary(x0, y0): + res = x0 <= y0 + return res + + +@triton.jit +def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + y = tl.load(in_ptr1 + idx_block, mask=idx_block < N) + ret = x <= y + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + # (torch.bfloat16, 'bfloat16'), + (torch.int8, 'int8'), + (torch.int16, 'int16'), + (torch.int32, 'int32'), + (torch.int64, 'int64'), +] + +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('N,NUMEL', shapes) +def test_elementwsie_common(dtype, sigtype, N, NUMEL): + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == "int64": + N = map_for_64_t[N] if N in map_for_64_t else N + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + y0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + ans = standard_binary(x0, y0) + out = torch.zeros((N, ), dtype=torch.bool).npu() + triton_elementwise_binary[1, 1, 1](x0, y0, out, N, NUMEL) + test_common.validate_cmp(sigtype, out, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_lgamma.py b/third_party/ascend/unittest/pytest_ut/test_lgamma.py new file mode 100644 index 0000000000..fc56fffa24 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_lgamma.py @@ -0,0 +1,114 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common +import triton.language.extra.cann.libdevice as libdevice + + +@triton.jit +def triton_lgamma(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + xindex = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = xindex < xnumel + x0 = tl.load(in_ptr0 + xindex, xmask) + y = libdevice.lgamma(x0) + tl.store(out_ptr0 + xindex, y, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_lgamma_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x = test_common.generate_tensor(shape, dtype).npu() + + # Avoid numerical instability near negative integer + nearest_int = torch.round(x) + neg_mask = nearest_int <= -1 + threshold = torch.zeros_like(x) + if neg_mask.any(): + neg_ints = nearest_int[neg_mask] + threshold[neg_mask] = 5.75e-5 * (2.42**(-1 - neg_ints)) + mask = (torch.abs(x - nearest_int) < threshold) & (nearest_int <= -1) + if mask.any(): + x = torch.where(mask, nearest_int, x) + + y_ref = torch.lgamma(x).npu() + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_lgamma[ncore, 1, 1](x, y_cal, x.numel(), xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_all_blocks_parallel(param_list, monkeypatch): + monkeypatch.setenv("TRITON_ALL_BLOCKS_PARALLEL", "1") + dtype, shape, ncore, xblock, xblock_sub = param_list + x = test_common.generate_tensor(shape, dtype).npu() + + # Avoid numerical instability near negative integer + nearest_int = torch.round(x) + neg_mask = nearest_int <= -1 + threshold = torch.zeros_like(x) + if neg_mask.any(): + neg_ints = nearest_int[neg_mask] + threshold[neg_mask] = 5.75e-5 * (2.42**(-1 - neg_ints)) + mask = (torch.abs(x - nearest_int) < threshold) & (nearest_int <= -1) + if mask.any(): + x = torch.where(mask, nearest_int, x) + + y_ref = torch.lgamma(x).npu() + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_lgamma[ncore, 1, 1](x, y_cal, x.numel(), xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) + monkeypatch.delenv("TRITON_ALL_BLOCKS_PARALLEL") + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 2048, 8), 2, 32768, 512], +]) +def test_auto_blockify(param_list, monkeypatch): + monkeypatch.setenv("TRITON_ALL_BLOCKS_PARALLEL", "1") + dtype, shape, ncore, xblock, xblock_sub = param_list + x = test_common.generate_tensor(shape, dtype).npu() + + # Avoid numerical instability near negative integer + nearest_int = torch.round(x) + neg_mask = nearest_int <= -1 + threshold = torch.zeros_like(x) + if neg_mask.any(): + neg_ints = nearest_int[neg_mask] + threshold[neg_mask] = 5.75e-5 * (2.42**(-1 - neg_ints)) + mask = (torch.abs(x - nearest_int) < threshold) & (nearest_int <= -1) + if mask.any(): + x = torch.where(mask, nearest_int, x) + + y_ref = torch.lgamma(x).npu() + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_lgamma[ncore, 1, 1](x, y_cal, x.numel(), xblock, xblock_sub, auto_blockify_size=ncore) + test_common.validate_cmp(dtype, y_cal, y_ref) + monkeypatch.delenv("TRITON_ALL_BLOCKS_PARALLEL") diff --git a/third_party/ascend/unittest/pytest_ut/test_linearize.py b/third_party/ascend/unittest/pytest_ut/test_linearize.py new file mode 100644 index 0000000000..4a33ef367b --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_linearize.py @@ -0,0 +1,602 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + +num_cube_core = 20 +num_vector_core = 20 + + +def foo(a, d, shape): + y = a.reshape(shape) + y = y.permute(0, 2, 1) + d + return y + + +@triton.jit +def triton_gpu_revised(in_ptr0, in_ptr1, out_ptr0, ynumel, xnumel, YBLOCK: tl.constexpr, SHAPE0: tl.constexpr, + SHAPE1: tl.constexpr, SHAPE2: tl.constexpr, XBLOCK: tl.constexpr): + yoffset = tl.program_id(1) * (tl.program_id(2) + 1) * YBLOCK + yindex = yoffset + tl.arange(0, YBLOCK) + ymask = yindex < ynumel + xoffset = tl.program_id(0) * XBLOCK + xindex = xoffset + tl.arange(0, XBLOCK) + xmask = xindex < xnumel + x2 = xindex[:, None] + x2_1 = xindex[None, :] + y3 = yindex[:, None] + y0 = (yindex % SHAPE1)[None, :] + y1 = (yindex // SHAPE1)[None, :] + tmp0 = tl.load(in_ptr0 + (x2_1 + (SHAPE2 * y3)), (x2_1 < xnumel) & (y3 < ynumel)) + tmp1 = tl.load(in_ptr1 + (y0 + (SHAPE1*x2) + (SHAPE1*SHAPE2*y1)), \ + (xindex[:,None] < xnumel) & (yindex[None,:] < ynumel)) + # (x2 < xnumel) & (y0 < SHAPE1)) + # (xindex[:,None] < xnumel) & (yindex[None,:] < ynumel)) + tmp10 = tmp1.permute(1, 0) + tmp2 = tmp0 + tmp10 + tl.store(out_ptr0 + (x2_1 + (SHAPE2 * y3)), tmp2, (x2_1 < xnumel) & (y3 < ynumel)) + + +def biggest_divisor(num): + for i in range(2, num): + if num % i == 0: + return num // i + return num + + +def find_good_yblock(ynumel, xnumel, y_upper, dtype): + y = ynumel + x = xnumel + + align_numel = 4 if dtype == torch.int64 else 8 + ub_upper = 3900 if dtype == torch.int64 else 8000 + + # optimize block_dim + def get_block_dim(y, x): + return ((xnumel + x - 1) // x) * ((y_upper + y - 1) // y) + + count = 0 + while (get_block_dim(y, x) < num_vector_core and y > 8 and count < 20): + y_1 = biggest_divisor(y) + if get_block_dim(y_1, x) > num_vector_core: + break + y = y_1 + if get_block_dim(y, x) < num_vector_core and x > align_numel: + x = x // 2 + count = count + 1 + + # optimize block_size to avoid ub-overflow + while (y * x > ub_upper): + y_1 = biggest_divisor(y) + if y_1 == y or y_1 <= align_numel: + break + y = y_1 + + while (y * x > ub_upper and x > align_numel): + x_1 = x // 2 + if x_1 <= align_numel: + break + x = x_1 + + return (x, y) + + +def triton_foo(a, d, shape, dtype): + z, y, x = shape + out = torch.empty_strided((z, x, y), (x * y, 1, x), device='npu', dtype=dtype) + XBLOCK, YBLOCK = find_good_yblock(y, x, y * z, dtype=dtype) + print(f"XBLOCK={XBLOCK},YBLOCK={YBLOCK}, block_dim={((x + XBLOCK -1 )//XBLOCK) * (((y*z) + YBLOCK -1 ) // YBLOCK)}") + grid = ((x + XBLOCK - 1) // XBLOCK, ((y * z) + YBLOCK - 1) // YBLOCK, 1) + + triton_gpu_revised[grid](a, d, out, y * z, x, SHAPE0=z, SHAPE1=y, SHAPE2=x, YBLOCK=YBLOCK, XBLOCK=XBLOCK) + return out + + +types = [ + (torch.float32, 'float32'), +] + +shapes = [(8, 2048, 4)] + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('Z,Y,X', shapes) +def test_linearize(Z, Y, X, dtype, sigtype): + shape = (Z, Y, X) + print(f"start data validation on shape:{shape}") + a = test_common.generate_tensor(shape=(Z, Y * X), dtype=sigtype).npu() + d = test_common.generate_tensor(shape=(Z, X, Y), dtype=sigtype).npu() + r = triton_foo(a, d, shape, dtype) + r1 = foo(a, d, shape) + test_common.validate_cmp(sigtype, r1, r) + print(f"data validation passed") + + +# Test linearize offset handling with expert routing pattern +@triton.jit +def linearize_offset_kernel( + bias_ptr, + output_ptr, + experts_ids_ptr, + N: tl.constexpr, + EM: tl.constexpr, + stride_bias_e, + stride_bias_n, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + off_experts = tl.load(experts_ids_ptr + pid_m).to(tl.int64) + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N + + if bias_ptr is not None: + bias = tl.load(bias_ptr + off_experts * stride_bias_e + offs_bn[None, :] * stride_bias_n) + tl.store(output_ptr + pid * 16 + tl.arange(0, 16), bias.reshape(16)) + + +def torch_linearize_offset(bias_ptr, experts_ids_ptr, N, EM, stride_bias_e, stride_bias_n, BLOCK_SIZE_M, BLOCK_SIZE_N, + GROUP_SIZE_M): + """PyTorch reference implementation for offset handling""" + output = torch.empty([16, 16], dtype=bias_ptr.dtype, device=bias_ptr.device) + + num_pid_m = (EM + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + num_pid_n = (N + BLOCK_SIZE_N - 1) // BLOCK_SIZE_N + num_pid_in_group = GROUP_SIZE_M * num_pid_n + + for pid in range(16): + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + off_experts = experts_ids_ptr[pid_m].to(torch.int64).item() + offs_bn = torch.arange(pid_n * BLOCK_SIZE_N, + (pid_n + 1) * BLOCK_SIZE_N, dtype=torch.int64, device=bias_ptr.device) % N + + bias = bias_ptr[off_experts, offs_bn] + output[pid] = bias + + return output + + +@pytest.mark.parametrize('dtype,sigtype', [(torch.float32, 'float32')]) +def test_linearize_offset_handling(dtype, sigtype): + """ + Test linearization's handling of complex offset patterns. + This test simulates expert routing scenarios where offsets are computed + dynamically based on expert IDs, validating correct pointer arithmetic + and linearization in the compiler. + """ + print(f"Testing linearize offset handling with dtype={sigtype}") + + # Setup test data + num_experts = 4 + hidden_dim = 64 + bias_ptr = torch.arange(0, num_experts * hidden_dim, dtype=dtype).npu().reshape(num_experts, hidden_dim) + output_ptr = torch.empty([16, 16], dtype=dtype).npu() + experts_ids_ptr = torch.tensor([1, 2, 3, 1], dtype=torch.int32).npu() + + # Kernel parameters + N = 64 + EM = 64 + stride_bias_e = 64 + stride_bias_n = 1 + BLOCK_SIZE_M = 16 + BLOCK_SIZE_N = 16 + GROUP_SIZE_M = 4 + + # Run triton kernel + linearize_offset_kernel[(16, )](bias_ptr=bias_ptr, output_ptr=output_ptr, experts_ids_ptr=experts_ids_ptr, N=N, + EM=EM, stride_bias_e=stride_bias_e, stride_bias_n=stride_bias_n, + BLOCK_SIZE_M=BLOCK_SIZE_M, BLOCK_SIZE_N=BLOCK_SIZE_N, GROUP_SIZE_M=GROUP_SIZE_M) + + # Compute reference result + expected = torch_linearize_offset(bias_ptr, experts_ids_ptr, N, EM, stride_bias_e, stride_bias_n, BLOCK_SIZE_M, + BLOCK_SIZE_N, GROUP_SIZE_M) + + # Validate results + test_common.validate_cmp(sigtype, expected, output_ptr) + print(f"Linearize offset handling test passed") + + +def torch_expand_dims_and_add(buffer, cache, buffer_stride, BLOCK, NUMEL): + for i in range(buffer.shape[0]): + tmp = buffer[i, 0, :buffer_stride] + accumulator = torch.zeros(2, buffer_stride).npu() + accumulator += tmp[None, :] + cache[i, 0, :] = accumulator.reshape(buffer_stride * 2) + + +@triton.jit +def expand_dims_and_add( + buffer_ptr, + cache_ptr, + buffer_stride: tl.constexpr, + BLOCK: tl.constexpr, + NUMEL: tl.constexpr, +): + pid = tl.program_id(0) + buffer_offset = pid * buffer_stride + buffer_index = (buffer_offset + tl.arange(0, buffer_stride)) % NUMEL + + cache_offset = pid * buffer_stride * 2 + cache_index = cache_offset + tl.arange(0, buffer_stride * 2) + + tmp = tl.load(buffer_ptr + buffer_index[None, :]) + accumulator = tl.zeros((2, buffer_stride), dtype=tl.float32) + accumulator += tmp + tl.store(cache_ptr + cache_index, accumulator.reshape(buffer_stride * 2)) + + +cache_shapes = [ + (5, 15), +] + + +@pytest.mark.parametrize("dtype, sigtype", types) +@pytest.mark.parametrize("batch_size, buffer_len", cache_shapes) +def test_expand_dims_and_add(batch_size, buffer_len, dtype, sigtype): + block = biggest_divisor(buffer_len) + cache_len = buffer_len * 2 + buffer = test_common.generate_tensor(shape=(batch_size, 1, buffer_len), dtype=sigtype).npu() + cache = torch.zeros(batch_size, 1, cache_len, dtype=dtype).npu() + cache_ref = torch.zeros(batch_size, 1, cache_len, dtype=dtype).npu() + numel = batch_size * buffer_len * 2 + torch_expand_dims_and_add(buffer, cache_ref, buffer_len, block, numel) + expand_dims_and_add[block, cache_len](buffer, cache, buffer_len, block, numel) + + test_common.validate_cmp(sigtype, cache, cache_ref) + + +def torch_save_cache_to_buffer(buffer, cache1, cache2, buffer_stride, cache_stride, BLOCK): + idx = torch.arange(0, cache_stride) + mask = ((idx // BLOCK) % 2 == 0) + max_len = min(buffer_stride, mask.shape[0]) + for i in range(buffer.shape[0]): + if i % 2 == 0: + tmp = cache1[i, 0, mask] + buffer[i, 0, :max_len] = tmp[:max_len] + else: + tmp = cache2[i, 0, mask] + buffer[i, 0, :max_len] = tmp[:max_len] + + +def torch_save_cache_to_buffer_with_offset(buffer, cache1, cache2, buffer_stride, cache_stride, BLOCK): + idx = torch.arange(0, cache_stride) + mask = ((idx // BLOCK) % 2 == 0) + max_len = min(buffer_stride, mask.shape[0]) + for i in range(buffer.shape[0]): + if i % 2 == 0: + tmp = cache1[i, 0, mask] + buffer[i, 0, :max_len] = tmp[:max_len] + else: + tmp = cache2[i, 0, ~mask] + buffer[i, 0, :max_len] = tmp[:max_len] + + +def torch_rearrange_and_combine_two_buffer( + buffer1, + buffer2, + cache, + buffer_stride, + NUM_BLOCK, + BLOCK, +): + for i in range(buffer1.shape[0]): + tmp1 = buffer1[i, 0, :] + tmp1 = tmp1.reshape(NUM_BLOCK, BLOCK) + tmp1 = tmp1.permute(1, 0) + tmp1 = tmp1.reshape(1, -1) + + tmp2 = buffer2[i, 0, :] + tmp2 = tmp2.reshape(NUM_BLOCK, BLOCK) + tmp2 = tmp2.permute(1, 0) + tmp2 = tmp2.reshape(1, -1) + + cache[i, 0, :buffer_stride] = tmp1 + cache[i, 0, buffer_stride:] = tmp2 + + +def torch_save_cache_to_buffer_with_mask(buffer, cache1, cache2, mask_int, buffer_stride, cache_stride, BLOCK, + MASK_NUM): + idx = torch.arange(0, cache_stride) + mask_idx = torch.arange(0, buffer_stride) + mask = ((idx // BLOCK) % 2 == 0) + max_len = min(buffer_stride, mask.shape[0]) + for i in range(buffer.shape[0]): + if i % 2 == 0: + tmp = cache1[i, 0, :] + tmp[~(idx < MASK_NUM)] = 0 + tmp = tmp[mask] + buffer[i, 0, :max_len] = tmp[:max_len] + else: + tmp = cache2[i, 0, mask] + tmp[~((mask_idx < MASK_NUM) & (mask_idx < mask_int[i]))] = 0 + buffer[i, 0, :max_len] = tmp[:max_len] + + +def torch_rearrange_cache_with_mask( + cache1, + cache2, + half_buffer_num, + buffer_len, + NUM_BLOCK, + BLOCK, +): + buffer_num = half_buffer_num * 2 + idx1 = torch.arange(0, half_buffer_num) + idx2 = torch.arange(0, buffer_len) + # all true + mask1 = idx1 < buffer_num + mask2 = idx2 < buffer_len + + for i in range(cache1.shape[0]): + for j in range(2): + mask = (mask1[j] & mask2) + tmp1 = cache1[i, j, :] + tmp1[~mask] = 0 + tmp1 = tmp1.reshape(NUM_BLOCK, BLOCK) + tmp1 = tmp1.permute(1, 0) + tmp1 = tmp1.reshape(1, -1) + + tmp2 = cache1[i, j + half_buffer_num, :] + tmp2[~mask] = 0 + tmp2 = tmp2.reshape(NUM_BLOCK, BLOCK) + tmp2 = tmp2.permute(1, 0) + tmp2 = tmp2.reshape(1, -1) + + cache2[i, j, :] = -tmp1 + cache2[i, j + half_buffer_num, :] = tmp2 + + +@triton.jit +def save_cache_to_buffer(buffer_ptr, cache_ptr1, cache_ptr2, buffer_stride: tl.constexpr, BLOCK: tl.constexpr): + pid_loc = tl.program_id(0) + buffer_offset = pid_loc * buffer_stride + buffer_index = buffer_offset + tl.arange(0, buffer_stride) + + cache_offset = pid_loc * buffer_stride + cache_index = cache_offset + tl.arange(0, buffer_stride) + cache_index_0 = cache_index // BLOCK + cache_index_1 = cache_index % BLOCK + + if pid_loc % 2 == 0: + tmp = tl.load(cache_ptr1 + (2 * BLOCK * cache_index_0 + cache_index_1)) + tl.store(buffer_ptr + buffer_index, tmp) + if pid_loc % 2 == 1: + tmp = tl.load(cache_ptr2 + (2 * BLOCK * cache_index_0 + cache_index_1)) + tl.store(buffer_ptr + buffer_index, tmp) + + +@triton.jit +def save_cache_to_buffer_with_offset(buffer_ptr, cache_ptr1, cache_ptr2, buffer_stride: tl.constexpr, + BLOCK: tl.constexpr): + pid_loc = tl.program_id(0) + buffer_offset = pid_loc * buffer_stride + buffer_index = buffer_offset + tl.arange(0, buffer_stride) + + cache_offset = pid_loc * buffer_stride + cache_index = cache_offset + tl.arange(0, buffer_stride) + cache_index_0 = cache_index // BLOCK + cache_index_1 = cache_index % BLOCK + + if pid_loc % 2 == 0: + tmp = tl.load(cache_ptr1 + (2 * BLOCK * cache_index_0 + cache_index_1)) + tl.store(buffer_ptr + buffer_index, tmp) + if pid_loc % 2 == 1: + tmp = tl.load(cache_ptr2 + BLOCK + (2 * BLOCK * cache_index_0 + cache_index_1)) + tl.store(buffer_ptr + buffer_index, tmp) + + +@triton.jit +def rearrange_and_combine_two_buffer( + buffer_ptr1, + buffer_ptr2, + cache_ptr, + buffer_stride: tl.constexpr, + NUM_BLOCK: tl.constexpr, + BLOCK: tl.constexpr, +): + pid_loc = tl.program_id(0) + buffer_offset = pid_loc * buffer_stride + buffer_index = buffer_offset + tl.arange(0, buffer_stride) + + tmp1 = tl.load(buffer_ptr1 + buffer_index) + tmp2 = tl.load(buffer_ptr2 + buffer_index) + + x_offset = pid_loc * 2 * BLOCK * buffer_stride + x_index = x_offset + tl.arange(0, buffer_stride) + + for i in range(2): + n_index = (x_index // BLOCK) + b_index = (x_index % BLOCK) + + if i % 2 == 0: + tl.store(cache_ptr + NUM_BLOCK * b_index + n_index, tmp1) + if i % 2 == 1: + tl.store(cache_ptr + NUM_BLOCK * b_index + n_index, tmp2) + + x_index = x_index + buffer_stride * BLOCK + + +@triton.jit +def save_cache_to_buffer_with_mask(buffer_ptr, cache_ptr1, cache_ptr2, mask_int_ptr, buffer_stride: tl.constexpr, + BLOCK: tl.constexpr, MASK_NUM: tl.constexpr): + pid_loc = tl.program_id(0) + + buffer_offset = pid_loc * buffer_stride + + buffer_index = tl.arange(0, buffer_stride) + index = buffer_offset + buffer_index + cache_index_0 = index // BLOCK + cache_index_1 = index % BLOCK + mask_int = tl.load(mask_int_ptr + pid_loc) + if pid_loc % 2 == 0: + tmp = tl.load(cache_ptr1 + (2*BLOCK*cache_index_0 + cache_index_1), \ + ((2*BLOCK*cache_index_0 + cache_index_1) < buffer_offset * 2 + MASK_NUM)) + tl.store(buffer_ptr + index, tmp) + if pid_loc % 2 == 1: + tmp = tl.load(cache_ptr2 + (2*BLOCK*cache_index_0 + cache_index_1), \ + (buffer_index < MASK_NUM) & (buffer_index < mask_int)) + tl.store(buffer_ptr + index, tmp) + + +@triton.jit +def rearrange_cache_with_mask( + cache_ptr1, + cache_ptr2, + half_buffer_num: tl.constexpr, + buffer_stride: tl.constexpr, + NUM_BLOCK: tl.constexpr, + BLOCK: tl.constexpr, +): + pid_loc = tl.program_id(0) + buffer_num = half_buffer_num * 2 + buffer_offset = pid_loc * buffer_stride * buffer_num + + index = tl.arange(0, buffer_stride) + buffer_index = buffer_offset + index + mask_index = index[None, :] + load_index = buffer_index[None, :] + + buffer_index2 = buffer_offset * BLOCK + index + n_index = (buffer_index2 // BLOCK) + b_index = (buffer_index2 % BLOCK) + store_index = NUM_BLOCK * b_index + n_index + store_index2 = store_index[None, :] + + buffer_num_offset = 0 + for i in range(2): + buffer_num_index = buffer_num_offset + tl.arange(0, half_buffer_num) + buffer_num_index2 = buffer_num_index[:, None] + + tmp1 = tl.load(cache_ptr1 + buffer_stride * buffer_num_index2 + load_index) + tmp1 = tmp1 * (2 * i - 1) + mask1 = buffer_num_index2 < buffer_num + mask2 = mask_index < buffer_stride + tl.store(cache_ptr2 + buffer_stride * buffer_num_index2 + store_index2, tmp1, mask1 & mask2) + # tl.store(cache_ptr2 + buffer_stride * buffer_num_index2 + store_index2, tmp1, True) + + buffer_num_offset = buffer_num_offset + half_buffer_num + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('batch_size,buffer_len', cache_shapes) +def test_linearize_jump_load(batch_size, buffer_len, dtype, sigtype): + block = biggest_divisor(buffer_len) + cache_len = buffer_len * 2 + buffer_ref = torch.zeros(batch_size, 1, buffer_len, dtype=dtype) + buffer = buffer_ref.npu() + cache1_ref = test_common.generate_tensor(shape=(batch_size, 1, cache_len), dtype=sigtype) + cache1 = cache1_ref.npu() + cache2_ref = test_common.generate_tensor(shape=(batch_size, 1, cache_len), dtype=sigtype) + cache2 = cache2_ref.npu() + + torch_save_cache_to_buffer(buffer_ref, cache1_ref, cache2_ref, buffer_len, cache_len, block) + save_cache_to_buffer[(batch_size, 1, 1)](buffer, cache1, cache2, buffer_len, block) + test_common.validate_cmp(sigtype, buffer, buffer_ref) + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('batch_size,buffer_len', cache_shapes) +def test_linearize_jump_load_with_offset(batch_size, buffer_len, dtype, sigtype): + block = biggest_divisor(buffer_len) + cache_len = buffer_len * 2 + buffer_ref = torch.zeros(batch_size, 1, buffer_len, dtype=dtype) + buffer = buffer_ref.npu() + cache1_ref = test_common.generate_tensor(shape=(batch_size, 1, cache_len), dtype=sigtype) + cache1 = cache1_ref.npu() + cache2_ref = test_common.generate_tensor(shape=(batch_size, 1, cache_len), dtype=sigtype) + cache2 = cache2_ref.npu() + + torch_save_cache_to_buffer_with_offset(buffer_ref, cache1_ref, cache2_ref, buffer_len, cache_len, block) + save_cache_to_buffer_with_offset[(batch_size, 1, 1)](buffer, cache1, cache2, buffer_len, block) + test_common.validate_cmp(sigtype, buffer, buffer_ref) + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('batch_size,buffer_len', cache_shapes) +def test_linearize_rearrange(batch_size, buffer_len, dtype, sigtype): + block = biggest_divisor(buffer_len) + num_block = int(buffer_len / block) + cache_len = buffer_len * 2 + buffer1_ref = test_common.generate_tensor(shape=(batch_size, 1, buffer_len), dtype=sigtype) + buffer1 = buffer1_ref.npu() + buffer2_ref = test_common.generate_tensor(shape=(batch_size, 1, buffer_len), dtype=sigtype) + buffer2 = buffer2_ref.npu() + cache_ref = torch.zeros(batch_size, 1, cache_len, dtype=dtype) + cache = cache_ref.npu() + + torch_rearrange_and_combine_two_buffer(buffer1_ref, buffer2_ref, cache_ref, buffer_len, num_block, block) + rearrange_and_combine_two_buffer[(batch_size, 1, 1)](buffer1, buffer2, cache, buffer_len, num_block, block) + test_common.validate_cmp(sigtype, cache, cache_ref) + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('batch_size,buffer_len', cache_shapes) +def test_linearize_jump_load_with_mask(batch_size, buffer_len, dtype, sigtype): + block = biggest_divisor(buffer_len) + cache_len = buffer_len * 2 + buffer_ref = torch.zeros(batch_size, 1, buffer_len, dtype=dtype) + buffer = buffer_ref.npu() + cache1_ref = test_common.generate_tensor(shape=(batch_size, 1, cache_len), dtype=sigtype) + cache1 = cache1_ref.npu() + cache2_ref = test_common.generate_tensor(shape=(batch_size, 1, cache_len), dtype=sigtype) + cache2 = cache2_ref.npu() + mask_ref = torch.arange(0, batch_size, dtype=torch.int64) * 2 + mask = mask_ref.npu() + mask_num = 16 + torch_save_cache_to_buffer_with_mask(buffer_ref, cache1_ref, cache2_ref, mask_ref, buffer_len, cache_len, block, + mask_num) + save_cache_to_buffer_with_mask[(batch_size, 1, 1)](buffer, cache1, cache2, mask, buffer_len, block, mask_num) + test_common.validate_cmp(sigtype, buffer, buffer_ref) + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('batch_size,buffer_len', cache_shapes) +def test_linearize_rearrange_with_mask(batch_size, buffer_len, dtype, sigtype): + block = biggest_divisor(buffer_len) + num_block = int(buffer_len / block) + cache1_ref = test_common.generate_tensor(shape=(batch_size, 4, buffer_len), dtype=sigtype) + cache1 = cache1_ref.npu() + cache2_ref = torch.zeros(batch_size, 4, buffer_len, dtype=dtype) + cache2 = cache2_ref.npu() + + torch_rearrange_cache_with_mask(cache1_ref, cache2_ref, 2, buffer_len, num_block, block) + rearrange_cache_with_mask[(batch_size, 1, 1)](cache1, cache2, 2, buffer_len, num_block, block) + test_common.validate_cmp(sigtype, cache2, cache2_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_linearize_fallback.py b/third_party/ascend/unittest/pytest_ut/test_linearize_fallback.py new file mode 100644 index 0000000000..841c2562d2 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_linearize_fallback.py @@ -0,0 +1,86 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + + +def torch_save_cache_to_buffer_indirect(buffer, cache, loc, buffer_stride, cache_stride, BLOCK): + max_len = min(buffer_stride, loc.shape[0]) + for i in range(buffer.shape[0]): + index = (loc // BLOCK) * BLOCK + loc % BLOCK + tmp = cache[i, 0, index] + buffer[i, 0, :max_len] = tmp[:max_len] + + +@triton.jit +def save_cache_to_buffer_indirect(buffer_ptr, cache_ptr, loc_ptr, buffer_stride: tl.constexpr, BLOCK: tl.constexpr): + pid_loc = tl.program_id(0) + index = tl.arange(0, buffer_stride) + loc = tl.load(loc_ptr + index) + buffer_offset = pid_loc * buffer_stride + buffer_index = buffer_offset + index + + cache_offset = pid_loc * buffer_stride * 2 + cache_index = cache_offset + loc + cache_index_0 = cache_index // BLOCK + cache_index_1 = cache_index % BLOCK + + tmp = tl.load(cache_ptr + (BLOCK * cache_index_0 + cache_index_1)) + tl.store(buffer_ptr + buffer_index, tmp) + + +def biggest_divisor(num): + for i in range(2, num): + if num % i == 0: + return num // i + return num + + +types = [ + (torch.float32, 'float32'), +] + +cache_shapes = [ + (5, 15), +] + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('batch_size,buffer_len', cache_shapes) +def test_linearize_indirect_load(batch_size, buffer_len, dtype, sigtype): + block = biggest_divisor(buffer_len) + cache_len = buffer_len * 2 + buffer_ref = torch.zeros(batch_size, 1, buffer_len, dtype=dtype) + buffer = buffer_ref.npu() + cache_ref = test_common.generate_tensor(shape=(batch_size, 1, cache_len), dtype=sigtype) + cache = cache_ref.npu() + + loc_ref = torch.arange(0, buffer_len) + loc = loc_ref.npu() + torch_save_cache_to_buffer_indirect(buffer_ref, cache_ref, loc_ref, buffer_len, cache_len, block) + save_cache_to_buffer_indirect[(batch_size, 1, 1)](buffer, cache, loc, buffer_len, block) + test_common.validate_cmp(sigtype, buffer, buffer_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_linearize_mask.py b/third_party/ascend/unittest/pytest_ut/test_linearize_mask.py new file mode 100644 index 0000000000..b3993de102 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_linearize_mask.py @@ -0,0 +1,168 @@ +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + +types_all = [ + (torch.float32, 'float32'), +] + +shapes_common = [(128, 256), (127, 256), (127, 16), (129, 256), (77, 1024), (69, 512)] + +block_size = [128, 256, 1024] + + +def ceil_div(a, b): + return (a + b - 1) // b + + +def profiler_wrapper(fn, *args): + result_path = "./result_profiling_broadcast" + skip_first = 10 + wait = 0 + warmup = 3 + active = 30 + repeat = 1 + stream = torch.npu.current_stream() + experimental_config = torch_npu.profiler._ExperimentalConfig( + aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, l2_cache=False, data_simplification=False) + with torch_npu.profiler.profile( + activities=[torch_npu.profiler.ProfilerActivity.CPU, torch_npu.profiler.ProfilerActivity.NPU], + schedule=torch_npu.profiler.schedule(wait=wait, warmup=warmup, active=active, repeat=repeat, + skip_first=skip_first), + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(result_path), record_shapes=True, + profile_memory=False, with_stack=False, with_flops=False, with_modules=False, + experimental_config=experimental_config) as prof: + stream.synchronize() + for _ in range(skip_first + (wait + warmup + active) * repeat): + fn(*args) + prof.step() + stream.synchronize() + + +@triton.jit +def linearize_mask_broadcast_kernel( + in_ptr, + output_ptr, + N: tl.constexpr, + M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, +): + pid = tl.program_id(axis=0) + offset = tl.multiple_of(pid * BLOCK_SIZE_N, N) + x1 = (offset + tl.arange(0, BLOCK_SIZE_N)) // N + mask1 = (x1 < M) + data = tl.load(in_ptr + x1 * N, mask=mask1, other=0) + x2 = pid * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + tl.store(output_ptr + x2, data) + + +def torch_linearize_mask_broadcast(in_tensor): + M = in_tensor.shape[0] // 2 + N = in_tensor.shape[1] + + output = torch.zeros_like(in_tensor) + + first_elements = in_tensor[:M, 0:1] + output[:M] = first_elements.expand(-1, N) + + return output + + +@pytest.mark.parametrize('dtype, sigtype', types_all) +@pytest.mark.parametrize('M, N', shapes_common) +@pytest.mark.parametrize('BLOCK_SIZE_N', block_size) +def test_linearize_mask_broadcast(M, N, BLOCK_SIZE_N, dtype, sigtype): + + in_tensor = torch.randn(2 * M, N, dtype=dtype).npu() + + triton_output = torch.zeros_like(in_tensor) + + grid = (ceil_div(2 * M * N, BLOCK_SIZE_N), ) + + linearize_mask_broadcast_kernel[grid](in_tensor, triton_output, N=N, M=M, BLOCK_SIZE_N=BLOCK_SIZE_N, + optimize_dynamic_offset=True) + + torch_output = torch_linearize_mask_broadcast(in_tensor.clone()) + assert torch.allclose(triton_output, torch_output, rtol=1e-5, atol=1e-8) + + +def triton_linearize_mask_broadcast(in_tensor, BLOCK_SIZE): + M = in_tensor.shape[0] // 2 + N = in_tensor.shape[1] + + triton_output = torch.zeros_like(in_tensor) + grid = (ceil_div(2 * M * N, BLOCK_SIZE), ) + + linearize_mask_broadcast_kernel[grid](in_tensor, triton_output, N=N, M=M, BLOCK_SIZE_N=BLOCK_SIZE, + optimize_dynamic_offset=True) + + +@triton.jit +def rem_kernel(in_ptr0, in_ptr1, out_ptr, N: tl.constexpr, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(0) + x = tl.arange(0, BLOCK_SIZE) + + base_offset = pid * BLOCK_SIZE + x + + rem_result = base_offset % 128 + mask = rem_result < 64 + + tmp0 = tl.load(in_ptr0 + base_offset, mask=mask, other=0.0) + tmp1 = tl.load(in_ptr1 + base_offset, mask=mask, other=0.0) + tmp2 = tmp0 + tmp1 + + tl.store(out_ptr + base_offset, tmp2, mask=mask) + + +def test_linearize_mask_rem(): + N = 1024 + BLOCK_SIZE = 256 + dtype = 'float32' + shape = (N, ) + + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + triton_res = torch.zeros(shape).npu() + + grid = (ceil_div(N, BLOCK_SIZE), ) + rem_kernel[grid](x0, x1, triton_res, N, BLOCK_SIZE=BLOCK_SIZE) + + base_offsets = torch.arange(N).npu() + rem_results = base_offsets % 128 + mask_bool = rem_results < 64 + + torch_res = torch.zeros((N, )).npu() + torch_res[mask_bool] = x0[mask_bool] + x1[mask_bool] + + test_common.validate_cmp(dtype, triton_res, torch_res) + + +def profile_performance_test(M, N, dtype, BLOCK_SIZE): + print(f"\nDetailed performance analysis: M={M}, N={N}, dtype={dtype}, block_size={BLOCK_SIZE}") + + in_tensor = torch.randn(2 * M, N, dtype=dtype).npu() + + def wrapper_func(x): + triton_linearize_mask_broadcast(x, BLOCK_SIZE=BLOCK_SIZE) + + # Run performance analysis + profiler_wrapper(wrapper_func, in_tensor) + + +if __name__ == "__main__": + print("Broadcast Kernel Performance Test Suite") + print("Function: Broadcast first element of first M rows, set remaining M rows to zero") + + # Optional: Run detailed profiler test (specific configuration) + profile_performance_test(512, 512, torch.float32, BLOCK_SIZE=1024) + + print("\n" + "=" * 80) + print("Test completed!") + print(f"Detailed performance analysis results saved in: ./result_profiling_broadcast/") + print("=" * 80) diff --git a/third_party/ascend/unittest/pytest_ut/test_linearize_mask_fallback.py b/third_party/ascend/unittest/pytest_ut/test_linearize_mask_fallback.py new file mode 100644 index 0000000000..17b6a0987a --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_linearize_mask_fallback.py @@ -0,0 +1,129 @@ +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + +types_all = [ + (torch.float32, 'float32'), +] + +shapes_common = [(128, 256), (127, 256), (127, 16), (129, 256), (77, 1024), (69, 512)] + +block_size = [128, 256, 1024] + + +def ceil_div(a, b): + return (a + b - 1) // b + + +def profiler_wrapper(fn, *args): + result_path = "./result_profiling_mask_fallback" + skip_first = 10 + wait = 0 + warmup = 3 + active = 30 + repeat = 1 + stream = torch.npu.current_stream() + experimental_config = torch_npu.profiler._ExperimentalConfig( + aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, l2_cache=False, data_simplification=False) + with torch_npu.profiler.profile( + activities=[torch_npu.profiler.ProfilerActivity.CPU, torch_npu.profiler.ProfilerActivity.NPU], + schedule=torch_npu.profiler.schedule(wait=wait, warmup=warmup, active=active, repeat=repeat, + skip_first=skip_first), + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(result_path), record_shapes=True, + profile_memory=False, with_stack=False, with_flops=False, with_modules=False, + experimental_config=experimental_config) as prof: + stream.synchronize() + for _ in range(skip_first + (wait + warmup + active) * repeat): + fn(*args) + prof.step() + stream.synchronize() + + +@triton.jit +def linearize_mask_broadcast_kernel( + in_ptr, + output_ptr, + N: tl.constexpr, + M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, +): + pid = tl.program_id(axis=0) + offset = tl.multiple_of(pid * BLOCK_SIZE_N, N) + x1 = (offset + tl.arange(0, BLOCK_SIZE_N)) + x2 = x1 // N + mask1 = (x1 < M * N) + data = tl.load(in_ptr + x2 * N, mask=mask1, other=0) + x3 = pid * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + tl.store(output_ptr + x3, data) + + +def torch_linearize_mask_broadcast(in_tensor): + M = in_tensor.shape[0] // 2 + N = in_tensor.shape[1] + + output = torch.zeros_like(in_tensor) + + first_elements = in_tensor[:M, 0:1] + output[:M] = first_elements.expand(-1, N) + + return output + + +@pytest.mark.parametrize('dtype, sigtype', types_all) +@pytest.mark.parametrize('M, N', shapes_common) +@pytest.mark.parametrize('BLOCK_SIZE_N', block_size) +def test_linearize_mask_broadcast(M, N, BLOCK_SIZE_N, dtype, sigtype): + + in_tensor = torch.randn(2 * M, N, dtype=dtype).npu() + + triton_output = torch.zeros_like(in_tensor) + + grid = (ceil_div(2 * M * N, BLOCK_SIZE_N), ) + + linearize_mask_broadcast_kernel[grid](in_tensor, triton_output, N=N, M=M, BLOCK_SIZE_N=BLOCK_SIZE_N, + optimize_dynamic_offset=True, enable_mask_fallback_conversion=True) + + torch_output = torch_linearize_mask_broadcast(in_tensor.clone()) + assert torch.allclose(triton_output, torch_output, rtol=1e-5, atol=1e-8) + + +def triton_linearize_mask_broadcast(in_tensor, BLOCK_SIZE): + M = in_tensor.shape[0] // 2 + N = in_tensor.shape[1] + + triton_output = torch.zeros_like(in_tensor) + grid = (ceil_div(2 * M * N, BLOCK_SIZE), ) + + linearize_mask_broadcast_kernel[grid](in_tensor, triton_output, N=N, M=M, BLOCK_SIZE_N=BLOCK_SIZE, + optimize_dynamic_offset=True, enable_mask_fallback_conversion=False) + + +def profile_performance_test(M, N, dtype, BLOCK_SIZE): + print(f"\nDetailed performance analysis: M={M}, N={N}, dtype={dtype}, block_size={BLOCK_SIZE}") + + in_tensor = torch.randn(2 * M, N, dtype=dtype).npu() + + def wrapper_func(x): + triton_linearize_mask_broadcast(x, BLOCK_SIZE=BLOCK_SIZE) + + # Run performance analysis + profiler_wrapper(wrapper_func, in_tensor) + + +if __name__ == "__main__": + print("mask fallback Kernel Performance Test Suite") + print("Function: Broadcast first element of first M rows, set remaining M rows to zero") + + # Optional: Run detailed profiler test (specific configuration) + profile_performance_test(512, 512, torch.float32, BLOCK_SIZE=1024) + + print("\n" + "=" * 80) + print("Test completed!") + print(f"Detailed performance analysis results saved in: ./result_profiling_mask_fallback/") + print("=" * 80) diff --git a/third_party/ascend/unittest/pytest_ut/test_linearize_permute.py b/third_party/ascend/unittest/pytest_ut/test_linearize_permute.py new file mode 100644 index 0000000000..8aaa295f52 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_linearize_permute.py @@ -0,0 +1,517 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + +types_all = [ + (torch.float32, 'float32'), +] + +shapes_common = [ + (8, 2048, 4), +] + + +def foo(a, b, shape): + y = a.reshape(shape) + y = y.permute(0, 2, 1) + b + return y + + +def ceil_div(a, b): + return (a + b - 1) // b + + +@triton.jit +def triton_gpu(in_ptr0, in_ptr1, out_ptr, ynumel, xnumel, YBLOCK: tl.constexpr, XBLOCK: tl.constexpr, + SHAPE0: tl.constexpr, SHAPE1: tl.constexpr, SHAPE2: tl.constexpr): + yoffset = tl.program_id(1) * (tl.program_id(2) + 1) * YBLOCK + yindex = yoffset + tl.arange(0, YBLOCK)[None, :] # (1, YBLOCK) + ymask = yindex < ynumel + xoffset = tl.program_id(0) * XBLOCK + xindex = xoffset + tl.arange(0, XBLOCK)[:, None] # (XBLOCK, 1) + xmask = xindex < xnumel + x2 = xindex # (XBLOCK, 1) + y3 = yindex # (1, YBLOCK) + y0 = yindex % SHAPE1 # (1, YBLOCK) + y1 = (yindex // SHAPE1) # (1, YBLOCK) + tmp0 = tl.load(in_ptr0 + (x2 + (SHAPE2 * y3)), xmask) # (XBLOCK, YBLOCK) + tmp1 = tl.load(in_ptr1 + (y0 + (SHAPE1 * x2) + (SHAPE1 * SHAPE2 * y1)), xmask) # (XBLOCK, YBLOCK) + tmp2 = tmp0 + tmp1 + tl.store(out_ptr + (x2 + (SHAPE2 * y3)), tmp2, xmask) + + +@triton.jit +def k_load_perm_select(ptr, out, ynumel, xnumel, YBLOCK: tl.constexpr, XBLOCK: tl.constexpr): + y = tl.program_id(1) * YBLOCK + tl.arange(0, YBLOCK)[None, :] + x = tl.program_id(0) * XBLOCK + tl.arange(0, XBLOCK)[:, None] + xmask = x < xnumel + ymask = y < ynumel + bad_mask = xmask | ymask + val = tl.load(ptr + (x + 4 * y), bad_mask) + tl.store(out + (x + 4 * y), val, xmask & ymask) + + +@triton.jit +def k_store_perm_select(in_ptr, out_ptr, ynumel, xnumel, YBLOCK: tl.constexpr, XBLOCK: tl.constexpr): + y = tl.program_id(1) * YBLOCK + tl.arange(0, YBLOCK)[None, :] + x = tl.program_id(0) * XBLOCK + tl.arange(0, XBLOCK)[:, None] + xmask = x < xnumel + ymask = y < ynumel + bad_mask = xmask | ymask + val = tl.load(in_ptr + (x + 4 * y), xmask & ymask) + tl.store(out_ptr + (x + 4 * y), val, bad_mask) + + +@triton.jit +def k_load_moddiv_noperm(in_ptr, out_ptr, ynumel, xnumel, YBLOCK: tl.constexpr, XBLOCK: tl.constexpr, + SHAPE0: tl.constexpr, SHAPE1: tl.constexpr, SHAPE2: tl.constexpr): + pid_y = tl.program_id(1) + pid_x = tl.program_id(0) + + yindex = pid_y * YBLOCK + tl.arange(0, YBLOCK) # (YBLOCK,) + xindex = pid_x * XBLOCK + tl.arange(0, XBLOCK) # (XBLOCK,) + + # 2D tile：[YBLOCK, XBLOCK] + y = yindex[:, None] # (YBLOCK, 1) + x = xindex[None, :] # (1, XBLOCK) + + mask = y < ynumel + + z = (yindex // SHAPE1)[:, None] # (YBLOCK, 1) + y_ = (yindex % SHAPE1)[:, None] # (YBLOCK, 1) + y_linear = z * SHAPE1 + y_ # (YBLOCK, 1) + offset_load = x + SHAPE2 * y_linear # (YBLOCK, XBLOCK) + val = tl.load(in_ptr + offset_load, mask=mask) + + offset_store = x + SHAPE2 * y # (YBLOCK, XBLOCK) + tl.store(out_ptr + offset_store, val, mask=mask) + + +@triton.jit +def k_store_moddiv_noperm(in_ptr, out_ptr, ynumel, xnumel, YBLOCK: tl.constexpr, XBLOCK: tl.constexpr, + SHAPE0: tl.constexpr, SHAPE1: tl.constexpr, SHAPE2: tl.constexpr): + pid_y = tl.program_id(1) + pid_x = tl.program_id(0) + + yindex = pid_y * YBLOCK + tl.arange(0, YBLOCK) # (YBLOCK,) + xindex = pid_x * XBLOCK + tl.arange(0, XBLOCK) # (XBLOCK,) + + y = yindex[:, None] # (YBLOCK, 1) + x = xindex[None, :] # (1, XBLOCK) + + mask = (y < ynumel) & (x < xnumel) + + offset_load = x + SHAPE2 * y # (YBLOCK, XBLOCK) + val = tl.load(in_ptr + offset_load, mask=mask) + + z = (yindex // SHAPE1)[:, None] # (YBLOCK, 1) + y_ = (yindex % SHAPE1)[:, None] # (YBLOCK, 1) + y_linear = z * SHAPE1 + y_ # (YBLOCK, 1) + offset_store = x + SHAPE2 * y_linear # (YBLOCK, XBLOCK) + tl.store(out_ptr + offset_store, val, mask=mask) + + +@triton.jit +def k_load_moddiv_perm(in_ptr, out_ptr, ynumel, xnumel, YBLOCK: tl.constexpr, XBLOCK: tl.constexpr, + SHAPE0: tl.constexpr, SHAPE1: tl.constexpr, SHAPE2: tl.constexpr): + pid_y = tl.program_id(1) + pid_x = tl.program_id(0) + + yindex = pid_y * YBLOCK + tl.arange(0, YBLOCK) # (YBLOCK,) + xindex = pid_x * XBLOCK + tl.arange(0, XBLOCK) # (XBLOCK,) + + X = xindex[:, None] # (XBLOCK, 1) + Y = yindex[None, :] # (1, YBLOCK) + mask_load = (X < xnumel) & (Y < ynumel) # (XBLOCK, YBLOCK) + + z = (yindex // SHAPE1)[None, :] # (1, YBLOCK) + y_ = (yindex % SHAPE1)[None, :] # (1, YBLOCK) + offset_load = X + SHAPE2 * y_ + SHAPE2 * SHAPE1 * z # (XBLOCK, YBLOCK) + val = tl.load(in_ptr + offset_load, mask=mask_load) # (XBLOCK, YBLOCK) + + y2 = yindex[:, None] # (YBLOCK, 1) + x2 = xindex[None, :] # (1, XBLOCK) + mask_store = (y2 < ynumel) & (x2 < xnumel) # (YBLOCK, XBLOCK) + offset_store = x2 + SHAPE2 * y2 # (YBLOCK, XBLOCK) + + tl.store(out_ptr + offset_store, val.permute(1, 0), mask=mask_store) + + +@triton.jit +def k_store_moddiv_perm(in_ptr, out_ptr, ynumel, xnumel, YBLOCK: tl.constexpr, XBLOCK: tl.constexpr, + SHAPE0: tl.constexpr, SHAPE1: tl.constexpr, SHAPE2: tl.constexpr): + pid_y = tl.program_id(1) + pid_x = tl.program_id(0) + + yindex = pid_y * YBLOCK + tl.arange(0, YBLOCK) # (YBLOCK,) + xindex = pid_x * XBLOCK + tl.arange(0, XBLOCK) # (XBLOCK,) + + y = yindex[:, None] # (YBLOCK, 1) + x = xindex[None, :] # (1, XBLOCK) + mask_load = (y < ynumel) & (x < xnumel) + + offset_load = x + SHAPE2 * y # (YBLOCK, XBLOCK) + val = tl.load(in_ptr + offset_load, mask=mask_load) # (YBLOCK, XBLOCK) + + X = xindex[:, None] # (XBLOCK, 1) + Y = yindex[None, :] # (1, YBLOCK) + mask_store = (X < xnumel) & (Y < ynumel) # (XBLOCK, YBLOCK) + + z = Y // SHAPE1 # (1, YBLOCK) + y_ = Y % SHAPE1 # (1, YBLOCK) + offset_store = X + SHAPE2 * y_ + SHAPE2 * SHAPE1 * z # (XBLOCK, YBLOCK) + + tl.store(out_ptr + offset_store, val.permute(1, 0), mask=mask_store) + + +@triton.jit +def k_load_store_moddiv_noperm(in_ptr, out_ptr, ynumel, xnumel, YBLOCK: tl.constexpr, XBLOCK: tl.constexpr, + SHAPE0: tl.constexpr, SHAPE1: tl.constexpr, SHAPE2: tl.constexpr): + pid_y = tl.program_id(1) + pid_x = tl.program_id(0) + + yindex = pid_y * YBLOCK + tl.arange(0, YBLOCK) # (YBLOCK,) + xindex = pid_x * XBLOCK + tl.arange(0, XBLOCK) # (XBLOCK,) + + y = yindex[:, None] # (YBLOCK, 1) + x = xindex[None, :] # (1, XBLOCK) + + mask = (y < ynumel) & (x < xnumel) + + z = (yindex // SHAPE1)[:, None] # (YBLOCK, 1) + y_ = (yindex % SHAPE1)[:, None] # (YBLOCK, 1) + y_linear = z * SHAPE1 + y_ # (YBLOCK, 1) + + offset = x + SHAPE2 * y_linear # (YBLOCK, XBLOCK) + + val = tl.load(in_ptr + offset, mask=mask) + val = val + 2 + tl.store(out_ptr + offset, val, mask=mask) + + +@triton.jit +def k_load_store_moddiv_perm(in_ptr, out_ptr, ynumel, xnumel, YBLOCK: tl.constexpr, XBLOCK: tl.constexpr, + SHAPE0: tl.constexpr, SHAPE1: tl.constexpr, SHAPE2: tl.constexpr): + # Program ID + pid_y = tl.program_id(1) + pid_x = tl.program_id(0) + + # Block index calculation + yindex = pid_y * YBLOCK + tl.arange(0, YBLOCK) # (YBLOCK,) + xindex = pid_x * XBLOCK + tl.arange(0, XBLOCK) # (XBLOCK,) + + # Broadcasting to 2D + x = xindex[:, None] # (XBLOCK, 1) + y = yindex[None, :] # (1, YBLOCK) + + # Valid mask + mask = (x < xnumel) & (y < ynumel) + + # Simulate linear index back to 3D + z = y // SHAPE1 # (1, YBLOCK) + y_ = y % SHAPE1 # (1, YBLOCK) + + # compute input offset (simulate [z, y_, x] access in strided format) + offset = x + SHAPE2 * y_ + SHAPE2 * SHAPE1 * z # (XBLOCK, YBLOCK) + + # load with implicit transpose (will be interpreted as [Y, X] then transposed) + val = tl.load(in_ptr + offset, mask=mask) + + # apply some dummy operation + val = val + 1 + + # store it back to out_ptr with same offset + tl.store(out_ptr + offset, val, mask=mask) + + +@triton.jit +def k_load_perm_scalar_ref(in_ptr, out_ptr, y1_numel, y0_numel, x2_numel, Y1BLOCK: tl.constexpr, Y0BLOCK: tl.constexpr, + Y0BLOCK_SUB: tl.constexpr, X2BLOCK_SUB: tl.constexpr): + pid_y1 = tl.program_id(0) + pid_y0 = tl.program_id(1) + + y1_scalar = pid_y1 * Y1BLOCK + + base_y0 = tl.arange(0, Y0BLOCK_SUB) + base_x2 = tl.arange(0, X2BLOCK_SUB) + + y0_offset = pid_y0 * Y0BLOCK + loops_y0 = (Y0BLOCK + Y0BLOCK_SUB - 1) // Y0BLOCK_SUB + + for loop_y0 in range(loops_y0): + y0_0 = y0_offset + loop_y0 * Y0BLOCK_SUB + base_y0[:, None] # (Y0BLOCK_SUB, 1) + y0 = y0_offset + (loop_y0 * Y0BLOCK_SUB) + base_y0[None, :] # (1, Y0BLOCK_SUB) + x2_1 = base_x2[None, :] # (1, X2BLOCK_SUB) + x2 = base_x2[:, None] # (X2BLOCK_SUB, 1) + y0_0_mask = y0_0 < min(Y0BLOCK + y0_offset, y0_numel) + y0_mask = y0 < min(Y0BLOCK + y0_offset, y0_numel) + val0 = tl.load(in_ptr + (y0 + 8 * x2 + 8 * 16 * y1_scalar), mask=y0_mask) + tl.store(out_ptr + (x2_1 + 16 * y0_0 + 8 * 16 * y1_scalar), val0.permute([1, 0]), mask=y0_0_mask) + + +@triton.jit +def k_load_perm_scalar(in_ptr, out_ptr, y1_numel, y0_numel, x2_numel, Y1BLOCK: tl.constexpr, Y0BLOCK: tl.constexpr, + Y0BLOCK_SUB: tl.constexpr, X2BLOCK_SUB: tl.constexpr): + pid_y1 = tl.program_id(0) + pid_y0 = tl.program_id(1) + + y1_scalar = pid_y1 * Y1BLOCK + + base_y0 = tl.arange(0, Y0BLOCK_SUB) + base_x2 = tl.arange(0, X2BLOCK_SUB) + + y0_offset = pid_y0 * Y0BLOCK + loops_y0 = (Y0BLOCK + Y0BLOCK_SUB - 1) // Y0BLOCK_SUB + + for loop_y0 in range(loops_y0): + y0_0 = y0_offset + loop_y0 * Y0BLOCK_SUB + base_y0[:, None] # (Y0BLOCK_SUB, 1) + x2_1 = base_x2[None, :] # (1, X2BLOCK_SUB) + y0_0_mask = y0_0 < min(Y0BLOCK + y0_offset, y0_numel) + val0 = tl.load(in_ptr + (y0_0 + 8 * x2_1 + 8 * 16 * y1_scalar), mask=y0_0_mask) + tl.store(out_ptr + (x2_1 + 16 * y0_0 + 8 * 16 * y1_scalar), val0, mask=y0_0_mask) + + +@pytest.mark.parametrize('dtype, sigtype', types_all) +@pytest.mark.parametrize('Z, Y, X', shapes_common) +def test_triton_gpu_kernel(Z, Y, X, dtype, sigtype): + shape = (Z, Y, X) + + a = test_common.generate_tensor(shape=(Z, Y * X), dtype=sigtype).npu() + b = test_common.generate_tensor(shape=(Z, X, Y), dtype=sigtype).npu() + + # must set device='npu' for empty_strided, do not use out.npu() later + out = torch.empty_strided((Z, X, Y), (X * Y, 1, X), device='npu', dtype=dtype) + + out_ref = foo(a, b, shape) + + ynumel = Z * Y + xnumel = X + XBLOCK = X + YBLOCK = 64 + + if ynumel % YBLOCK != 0: + pytest.skip(f"ynumel:{ynumel} not divisible by YBLOCK:{YBLOCK}") + + grid = (ceil_div(xnumel, XBLOCK), ceil_div(ynumel, YBLOCK), 1) + triton_gpu[grid](a, b, out, ynumel, xnumel, YBLOCK=YBLOCK, XBLOCK=XBLOCK, SHAPE0=Z, SHAPE1=Y, SHAPE2=X) + + test_common.validate_cmp(sigtype, out_ref, out) + + +@pytest.mark.parametrize('dtype,sigtype', types_all) +@pytest.mark.parametrize('xnumel, ynumel, XBLOCK, YBLOCK', [(3, 513, 4, 64)]) +def test_k_load_perm_select(xnumel, ynumel, XBLOCK, YBLOCK, dtype, sigtype): + + in_ptr = test_common.generate_tensor(shape=(ynumel * 4, ), dtype=sigtype).npu() + out_ptr = torch.zeros_like(in_ptr) + + grid = (ceil_div(xnumel, XBLOCK), ceil_div(ynumel, YBLOCK), 1) + k_load_perm_select[grid](in_ptr, out_ptr, ynumel, xnumel, YBLOCK=YBLOCK, XBLOCK=XBLOCK) + + out_ref = torch.zeros_like(out_ptr) + y_idx = torch.arange(ynumel).unsqueeze(1).npu() # [ynumel, 1] + x_idx = torch.arange(xnumel).unsqueeze(0).npu() # [1, xnumel] + idx = (x_idx + 4 * y_idx).reshape(-1) # [ynumel * xnumel] + out_ref[idx] = in_ptr[idx] + torch.testing.assert_close(out_ptr[idx], out_ref[idx]) + + +@pytest.mark.parametrize('dtype,sigtype', types_all) +@pytest.mark.parametrize('xnumel, ynumel, XBLOCK, YBLOCK', [(3, 513, 4, 64)]) +def test_k_store_perm_select(xnumel, ynumel, XBLOCK, YBLOCK, dtype, sigtype): + in_ptr = test_common.generate_tensor(shape=(ynumel * 4, ), dtype=sigtype).npu() + out_ptr = torch.zeros_like(in_ptr) + + grid = (ceil_div(xnumel, XBLOCK), ceil_div(ynumel, YBLOCK), 1) + k_store_perm_select[grid]( + in_ptr, + out_ptr, + ynumel, + xnumel, + YBLOCK=YBLOCK, + XBLOCK=XBLOCK, + ) + + out_ref = torch.zeros_like(out_ptr) + y_idx = torch.arange(ynumel).unsqueeze(1).npu() + x_idx = torch.arange(xnumel).unsqueeze(0).npu() + idx = (x_idx + 4 * y_idx).reshape(-1) + out_ref[idx] = in_ptr[idx] + torch.testing.assert_close(out_ptr[idx], out_ref[idx]) + + +@pytest.mark.parametrize('dtype, sigtype', types_all) +@pytest.mark.parametrize('Z, Y, X', shapes_common) +def test_k_load_moddiv_noperm(Z, Y, X, dtype, sigtype): + a = test_common.generate_tensor(shape=(Z, Y, X), dtype=sigtype).npu() + in_flat = a.contiguous().view(-1) + out = torch.zeros_like(in_flat) + + ynumel = Z * Y + xnumel = X + XBLOCK = X + YBLOCK = 64 + grid = (ceil_div(xnumel, XBLOCK), ceil_div(ynumel, YBLOCK), 1) + + k_load_moddiv_noperm[grid](in_flat, out, ynumel, xnumel, YBLOCK=YBLOCK, XBLOCK=XBLOCK, SHAPE0=Z, SHAPE1=Y, SHAPE2=X) + + torch.testing.assert_close(out, in_flat) + + +@pytest.mark.parametrize('dtype, sigtype', types_all) +@pytest.mark.parametrize('Z, Y, X', shapes_common) +def test_k_store_moddiv_noperm(Z, Y, X, dtype, sigtype): + a = test_common.generate_tensor(shape=(Z, Y, X), dtype=sigtype).npu() + in_flat = a.contiguous().view(-1) + out = torch.zeros_like(in_flat) + + ynumel = Z * Y + xnumel = X + XBLOCK = X + YBLOCK = 64 + grid = (ceil_div(xnumel, XBLOCK), ceil_div(ynumel, YBLOCK), 1) + + k_store_moddiv_noperm[grid](in_flat, out, ynumel, xnumel, YBLOCK=YBLOCK, XBLOCK=XBLOCK, SHAPE0=Z, SHAPE1=Y, + SHAPE2=X) + + torch.testing.assert_close(out, in_flat) + + +@pytest.mark.parametrize('dtype, sigtype', types_all) +@pytest.mark.parametrize('Z, Y, X', shapes_common) +def test_k_load_moddiv_perm(Z, Y, X, dtype, sigtype): + a = test_common.generate_tensor(shape=(Z, Y, X), dtype=sigtype).npu() + in_flat = a.contiguous().view(-1) + out = torch.zeros_like(in_flat) + + ynumel = Z * Y + xnumel = X + XBLOCK = X + YBLOCK = 64 + + grid = (ceil_div(xnumel, XBLOCK), ceil_div(ynumel, YBLOCK), 1) + + k_load_moddiv_perm[grid](in_flat, out, ynumel, xnumel, YBLOCK=YBLOCK, XBLOCK=XBLOCK, SHAPE0=Z, SHAPE1=Y, SHAPE2=X) + + torch.testing.assert_close(out, in_flat) + + +@pytest.mark.parametrize('dtype, sigtype', types_all) +@pytest.mark.parametrize('Z, Y, X', shapes_common) +def test_k_store_moddiv_perm(Z, Y, X, dtype, sigtype): + a = test_common.generate_tensor(shape=(Z, Y, X), dtype=sigtype).npu() + in_flat = a.contiguous().view(-1) + out = torch.zeros_like(in_flat) + + ynumel = Z * Y + xnumel = X + XBLOCK = X + YBLOCK = 64 + + grid = (ceil_div(xnumel, XBLOCK), ceil_div(ynumel, YBLOCK), 1) + + k_store_moddiv_perm[grid](in_flat, out, ynumel, xnumel, YBLOCK=YBLOCK, XBLOCK=XBLOCK, SHAPE0=Z, SHAPE1=Y, SHAPE2=X) + + torch.testing.assert_close(out, in_flat) + + +@pytest.mark.parametrize('dtype, sigtype', types_all) +@pytest.mark.parametrize('Z, Y, X', shapes_common) +def test_k_load_store_moddiv_noperm(Z, Y, X, dtype, sigtype): + a = test_common.generate_tensor(shape=(Z, Y, X), dtype=sigtype).npu() + in_flat = a.contiguous().view(-1) + out = torch.empty_like(in_flat) + + ynumel = Z * Y + xnumel = X + XBLOCK = 4 + YBLOCK = 256 + grid = (ceil_div(xnumel, XBLOCK), ceil_div(ynumel, YBLOCK), 1) + + k_load_store_moddiv_noperm[grid](in_flat, out, ynumel, xnumel, YBLOCK=YBLOCK, XBLOCK=XBLOCK, SHAPE0=Z, SHAPE1=Y, + SHAPE2=X) + + ref = (a + 2).contiguous().view(-1) + torch.testing.assert_close(out, ref) + + +@pytest.mark.parametrize('dtype, sigtype', types_all) +@pytest.mark.parametrize('Z, Y, X', shapes_common) +def test_k_load_store_moddiv_perm(Z, Y, X, dtype, sigtype): + shape = (Z, Y, X) + numel = Z * Y * X + + a = test_common.generate_tensor(shape, dtype=sigtype).npu() + a_flat = a.contiguous().view(-1) + + out = torch.zeros_like(a_flat) + + ynumel = Z * Y + xnumel = X + XBLOCK = 4 + YBLOCK = 256 + grid = (ceil_div(xnumel, XBLOCK), ceil_div(ynumel, YBLOCK), 1) + + k_load_store_moddiv_perm[grid](a_flat, out, ynumel, xnumel, YBLOCK=YBLOCK, XBLOCK=XBLOCK, SHAPE0=Z, SHAPE1=Y, + SHAPE2=X) + + a_reshaped = a + 1 + out_ref = a_reshaped.contiguous().view(-1) + + torch.testing.assert_close(out, out_ref) + + +@pytest.mark.parametrize('dtype', [torch.float32]) +@pytest.mark.parametrize('y1_numel, y0_numel, x2_numel, Y1BLOCK, Y0BLOCK, Y0BLOCK_SUB, X2BLOCK_SUB', [ + (2, 8, 16, 1, 8, 4, 16), +]) +def test_k_load_perm_scalar(y1_numel, y0_numel, x2_numel, Y1BLOCK, Y0BLOCK, Y0BLOCK_SUB, X2BLOCK_SUB, dtype): + in_ptr = torch.randn((y0_numel * 8, x2_numel), dtype=dtype, device='npu') + + out_ptr_triton_ref = torch.zeros((y1_numel, y0_numel, x2_numel), dtype=dtype, device='npu') + out_ptr_triton = torch.zeros((y1_numel, y0_numel, x2_numel), dtype=dtype, device='npu') + + grid = (ceil_div(y1_numel, Y1BLOCK), ceil_div(y0_numel, Y0BLOCK)) + + k_load_perm_scalar_ref[grid](in_ptr=in_ptr, out_ptr=out_ptr_triton_ref, y1_numel=y1_numel, y0_numel=y0_numel, + x2_numel=x2_numel, Y1BLOCK=Y1BLOCK, Y0BLOCK=Y0BLOCK, Y0BLOCK_SUB=Y0BLOCK_SUB, + X2BLOCK_SUB=X2BLOCK_SUB) + + k_load_perm_scalar[grid]( + in_ptr=in_ptr, + out_ptr=out_ptr_triton, + y1_numel=y1_numel, + y0_numel=y0_numel, + x2_numel=x2_numel, + Y1BLOCK=Y1BLOCK, + Y0BLOCK=Y0BLOCK, + Y0BLOCK_SUB=Y0BLOCK_SUB, + X2BLOCK_SUB=X2BLOCK_SUB, + ) + + torch.testing.assert_close(out_ptr_triton, out_ptr_triton_ref, rtol=1e-5, atol=1e-6) diff --git a/third_party/ascend/unittest/pytest_ut/test_load.py b/third_party/ascend/unittest/pytest_ut/test_load.py new file mode 100644 index 0000000000..440467ece3 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_load.py @@ -0,0 +1,117 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import numpy as np +import torch +import pytest +import test_common + +# eg: pytest -v test.py::test_add +############################# + + +@triton.jit +def triton_load_store(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + xindex = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = xindex < xnumel + x0 = xindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp2 = tmp0 + tl.store(out_ptr0 + (xindex), tmp2, xmask) + + +# require: all data (4d and 5d) can be placed into but without ub overflow +@triton.jit +def triton_load_store_multi_d(in_ptr0, out_ptr0, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, + SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, STRIDE_1: tl.constexpr, + STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + tmp_in = tl.load(in_ptr0 + offsets, masks) + tmp_out = tmp_in + tl.store(out_ptr0 + offsets, tmp_out, masks) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], + ['float32', (8, 8, 4), 2, 128, 64], + ['float16', (8, 8, 4), 2, 128, 64], + ['int8', (8, 8, 4), 2, 128, 64], + ['int8', (8, 7, 4), 2, 128, 64], +]) +def test_load_store(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + y_ref = x0 + y_cal = test_common.generate_tensor(shape, dtype).npu() + triton_load_store[(ncore, )](x0, y_cal, x0.numel(), xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (8, 4, 16, 16)], + ['float16', (8, 4, 16, 16)], + ['int8', (8, 4, 16, 16)], + ['float32', (8, 8, 4, 4)], + ['float16', (8, 8, 4, 4)], + ['int8', (8, 8, 4, 4)], + ['float32', (3, 8, 2, 16, 16)], + ['float16', (3, 8, 2, 16, 16)], + ['int8', (9, 8, 8, 16, 16)], + ['float32', (11, 8, 8, 4, 4)], + ['float16', (11, 8, 8, 4, 4)], + ['int8', (11, 8, 8, 4, 4)], +]) +def test_load_store_multi_d(param_list): + dtype, shape = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + y_expect = x0 + y_actual = test_common.generate_tensor(shape, dtype).npu() + + blocks = list(x0.size()) + shapes = list(x0.stride()) + while len(blocks) < 5: + blocks.append(1) + shapes.append(1) + triton_load_store_multi_d[(1, )](x0, y_actual, *blocks, *blocks, *shapes) + test_common.validate_cmp(dtype, y_actual, y_expect) diff --git a/third_party/ascend/unittest/pytest_ut/test_load_store.py b/third_party/ascend/unittest/pytest_ut/test_load_store.py new file mode 100644 index 0000000000..8986a5f65d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_load_store.py @@ -0,0 +1,151 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common + + +@triton.jit +def triton_load_store(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + x_index = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = x_index < xnumel + tmp0 = tl.load(in_ptr0 + x_index, xmask) + tmp2 = tmp0 + tl.store(out_ptr0 + x_index, tmp2, xmask) + + +# require: all data (4d and 5d) can be placed into but without ub overflow +@triton.jit +def triton_load_store_multi_d(in_ptr0, out_ptr0, BLOCK_0: tl.constexpr, BLOCK_1: tl.constexpr, BLOCK_2: tl.constexpr, + BLOCK_3: tl.constexpr, BLOCK_4: tl.constexpr, SHAPE_0: tl.constexpr, + SHAPE_1: tl.constexpr, SHAPE_2: tl.constexpr, SHAPE_3: tl.constexpr, + SHAPE_4: tl.constexpr, STRIDE_0: tl.constexpr, STRIDE_1: tl.constexpr, + STRIDE_2: tl.constexpr, STRIDE_3: tl.constexpr, STRIDE_4: tl.constexpr): + offsets = tl.program_id(0) + + offsets = offsets + tl.arange(0, BLOCK_0) * STRIDE_0 + masks = tl.arange(0, BLOCK_0) < SHAPE_0 + if (BLOCK_1 * BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, None] + tl.arange(0, BLOCK_1)[None, :] * STRIDE_1 + masks = masks[:, None] & (tl.arange(0, BLOCK_1)[None, :] < SHAPE_1) + if (BLOCK_2 * BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, None] + tl.arange(0, BLOCK_2)[None, None, :] * STRIDE_2 + masks = masks[:, :, None] & (tl.arange(0, BLOCK_2)[None, None, :] < SHAPE_2) + if (BLOCK_3 * BLOCK_4) > 1: + offsets = offsets[:, :, :, None] + tl.arange(0, BLOCK_3)[None, None, None, :] * STRIDE_3 + masks = masks[:, :, :, None] & (tl.arange(0, BLOCK_3)[None, None, None, :] < SHAPE_3) + if BLOCK_4 > 1: + offsets = offsets[:, :, :, :, None] + tl.arange(0, BLOCK_4)[None, None, None, None, :] * STRIDE_4 + masks = masks[:, :, :, :, None] & (tl.arange(0, BLOCK_4)[None, None, None, None, :] < SHAPE_4) + + tmp_in = tl.load(in_ptr0 + offsets, masks) + tmp_out = tmp_in + tl.store(out_ptr0 + offsets, tmp_out, masks) + + +@triton.jit +def triton_load_store_sle_mask(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + x_index = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = x_index <= xnumel + tmp0 = tl.load(in_ptr0 + x_index, xmask) + tmp2 = tmp0 + tl.store(out_ptr0 + x_index, tmp2, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], + ['float32', (8, 8, 4), 2, 128, 64], + ['float16', (8, 8, 4), 2, 128, 64], + ['int8', (8, 8, 4), 2, 128, 64], +]) +def test_load_store(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + y_ref = x0 + # triton结果 + y_cal = test_common.generate_tensor(shape, dtype).npu() + triton_load_store[(ncore, )](x0, y_cal, x0.numel(), xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, y_cal, y_ref) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (8, 4, 16, 16)], + ['float16', (8, 4, 16, 16)], + ['int8', (8, 4, 16, 16)], + ['float32', (8, 8, 4, 4)], + ['float16', (8, 8, 4, 4)], + ['int8', (8, 8, 4, 4)], + ['float32', (3, 8, 2, 16, 16)], + ['float16', (3, 8, 2, 16, 16)], + ['int8', (9, 8, 8, 16, 16)], + ['float32', (11, 8, 8, 4, 4)], + ['float16', (11, 8, 8, 4, 4)], + ['int8', (11, 8, 8, 4, 4)], +]) +def test_load_store_multi_d(param_list): + # 生成数据 + dtype, shape = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + y_expect = x0 + y_actual = test_common.generate_tensor(shape, dtype).npu() + # triton结果 + blocks = list(x0.size()) + shapes = list(x0.stride()) + while len(blocks) < 5: + blocks.append(1) + shapes.append(1) + triton_load_store_multi_d[(1, )](x0, y_actual, *blocks, *blocks, *shapes) + # 比较结果 + test_common.validate_cmp(dtype, y_actual, y_expect) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], + ['float32', (8, 8, 4), 2, 128, 64], + ['float16', (8, 8, 4), 2, 128, 64], + ['int8', (8, 8, 4), 2, 128, 64], +]) +def test_load_store_sle_mask(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + y_ref = x0 + # triton结果 + y_cal = test_common.generate_tensor(shape, dtype).npu() + triton_load_store_sle_mask[(ncore, )](x0, y_cal, x0.numel() - 1, xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_log.py b/third_party/ascend/unittest/pytest_ut/test_log.py new file mode 100644 index 0000000000..8263a9cdc9 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_log.py @@ -0,0 +1,99 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + + +def standard_unary(x0, dtype): + res = torch.log(x0) + return res + + +def standard_binary(x0, y0, dtype): + res = x0 + y0 + return res + + +@triton.jit +def triton_elementwise_unary(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + ret = tl.math.log(x) + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +@triton.jit +def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + y = tl.load(in_ptr1 + idx_block, mask=idx_block < N) + ret = x + y + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + # (torch.bfloat16, 'bfloat16'), + # (torch.int8, 'int8'), + # (torch.int16, 'int16'), + # (torch.int32, 'int32'), + # (torch.int64, 'int64'), +] + +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('N,NUMEL', shapes) +def test_elementwsie_common(dtype, sigtype, N, NUMEL): + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == "int64": + N = map_for_64_t[N] if N in map_for_64_t else N + + print(f"elementwise : ({N},) {dtype} {sigtype}") + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standard_unary(x0, dtype) + x0 = x0.npu() + print(ans) + + out = torch.zeros((N, ), dtype=dtype).npu() + triton_elementwise_unary[1, 1, 1](x0, out, N=N, NUMEL=NUMEL, debug=True) + print(out) + + test_common.validate_cmp(sigtype, out, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_log10.py b/third_party/ascend/unittest/pytest_ut/test_log10.py new file mode 100644 index 0000000000..caa8bbcce6 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_log10.py @@ -0,0 +1,90 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import torch +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + + +@triton.jit +def log10_kernel(x_ptr, y_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + + mask = offsets < n_elements + + x = tl.load(x_ptr + offsets, mask=mask) + + y = libdevice.log10(x) + + tl.store(y_ptr + offsets, y, mask=mask) + + +@pytest.mark.parametrize('shape', [ + (12, 16), +]) +@pytest.mark.parametrize('dtype', ['float32']) +def test_log10(shape, dtype): + x = test_common.generate_tensor(shape, dtype).npu() + + x[0, 0] = 1.0 # log10(1) = 0 + x[0, 1] = 10.0 # log10(10) = 1 + x[0, 2] = 100.0 # log10(100) = 2 + x[0, 3] = 0.1 # log10(0.1) = -1 + x[0, 4] = 0.0 # log10(0) = -inf + x[0, 5] = -1.0 # log10(-1) = NaN + x[0, 6] = 2.0 # log10(2) ≈ 0.3010 + + y = torch.empty_like(x) + + BLOCK_SIZE = 192 + grid = lambda meta: (triton.cdiv(192, meta['BLOCK_SIZE']), ) + + log10_kernel[grid](x, y, 192, BLOCK_SIZE=BLOCK_SIZE) + + expected = torch.log10(x) + print(f"triton_ret = {y}") + print(f"triton_ret = {expected}") + + valid_mask = (x > 0) + + if torch.any(valid_mask): + valid_y = y[valid_mask] + valid_expected = expected[valid_mask] + + torch.testing.assert_close(valid_y, valid_expected, rtol=1e-3, atol=1e-3) + + # Check if negative values return NaN + negative_mask = (x < 0) + if torch.any(negative_mask): + negative_y = y[negative_mask] + assert torch.all(torch.isnan(negative_y)), "Negative inputs should return NaN" + + # Check if zero value returns -inf + zero_mask = (x == 0) + if torch.any(zero_mask): + zero_y = y[zero_mask] + assert torch.all(torch.isinf(zero_y) & (zero_y < 0)), "Zero inputs should return -inf" + + print("✓ LOG10 test PASSED!") diff --git a/third_party/ascend/unittest/pytest_ut/test_log1p.py b/third_party/ascend/unittest/pytest_ut/test_log1p.py new file mode 100644 index 0000000000..c066ea35ca --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_log1p.py @@ -0,0 +1,60 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import torch +import test_common +import triton.language.extra.cann.libdevice as libdevice + + +def torch_log1p(x0, x1): + res = x0 + torch.log1p(x1) + return res + + +@triton.jit +def triton_log1p(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + x_index = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = x_index < xnumel + tmp0 = tl.load(in_ptr0 + x_index, xmask) + tmp1 = tl.load(in_ptr1 + x_index, xmask) + tmp2 = tmp0 + libdevice.log1p(tmp1) + tl.store(out_ptr0 + x_index, tmp2, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_log1p(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + torch_res = torch_log1p(x0, x1) + # triton结果 + triton_res = test_common.generate_tensor(shape, dtype).npu() + triton_log1p[ncore, 1, 1](x0, x1, triton_res, x0.numel(), xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_log2.py b/third_party/ascend/unittest/pytest_ut/test_log2.py new file mode 100644 index 0000000000..37c912d623 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_log2.py @@ -0,0 +1,58 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common + + +def torch_log2(x0): + res = torch.log2(x0) + return res + + +@triton.jit +def triton_log2(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x_inedx = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + x_inedx, None) + tmp2 = tl.log2(tmp0) + tl.store(out_ptr0 + x_inedx, tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_log2(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + torch_res = torch_log2(x0) + # triton结果 + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_log2[ncore, 1, 1](x0, triton_res, xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_log_2.py b/third_party/ascend/unittest/pytest_ut/test_log_2.py new file mode 100644 index 0000000000..f45bdd13bd --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_log_2.py @@ -0,0 +1,59 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common + + +def torch_exp2(x0): + res = torch.pow(2, x0, out=None) + return res + + +@triton.jit +def triton_exp2(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x_index = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + x_index, None) + tmp1 = tl.exp2(tmp0) + tl.store(out_ptr0 + x_index, tmp1, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_exp2(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + torch_res = torch_exp2(x0) + # triton结果 + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_exp2[ncore, 1, 1](x0, triton_res, xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_logical_and.py b/third_party/ascend/unittest/pytest_ut/test_logical_and.py new file mode 100644 index 0000000000..b5aefc368b --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_logical_and.py @@ -0,0 +1,77 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import numpy as np +import test_common + + +def torch_logical_and(x0, x1): + res = torch.logical_and(x0, x1) + return res + + +@triton.jit +def triton_logical_and(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x_index = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + x_index) + tmp1 = tl.load(in_ptr1 + x_index) + tmp2 = tmp0.logical_and(tmp1) + tl.store(out_ptr0 + x_index, tmp2) + + +@pytest.mark.parametrize('param_list', [ + ['bool', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], + # ['int16', (2, 4096, 8), 2, 32768, 1024], + # ['int32', (2, 4096, 8), 2, 32768, 1024], + # ['int64', (2, 4096, 8), 2, 32768, 1024], + ['uint8', (2, 4096, 8), 2, 32768, 1024], + # ['uint16', (2, 4096, 8), 2, 32768, 1024], + # ['uint32', (2, 4096, 8), 2, 32768, 1024], + # ['uint64', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + # ['float32', (2, 4096, 8), 2, 32768, 1024], + # ['bfloat16', (2, 4096, 8), 2, 32768, 1024], +]) +def test_and(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + torch_dtype = eval('torch.' + dtype) + np_x0 = test_common.generate_numpy(shape, dtype) + np_x1 = test_common.generate_numpy(shape, dtype) + x0 = torch.from_numpy(np_x0).to(torch_dtype).npu() + x1 = torch.from_numpy(np_x1).to(torch_dtype).npu() + # torch结果 + np_res = np.logical_and(np_x0, np_x1) + torch_res = torch.from_numpy(np_res).to(torch.bool) + # triton结果 + triton_res = torch.zeros(shape, dtype=torch.bool).npu() + triton_logical_and[ncore, 1, 1](x0, x1, triton_res, xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_logical_or.py b/third_party/ascend/unittest/pytest_ut/test_logical_or.py new file mode 100644 index 0000000000..fd40ca5232 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_logical_or.py @@ -0,0 +1,77 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import numpy as np +import test_common + + +def torch_logical_or(x0, x1): + res = torch.logical_or(x0, x1) + return res + + +@triton.jit +def triton_logical_or(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x_index = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + x_index) + tmp1 = tl.load(in_ptr1 + x_index) + tmp2 = tmp0.logical_or(tmp1) + tl.store(out_ptr0 + x_index, tmp2) + + +@pytest.mark.parametrize('param_list', [ + ['bool', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], + # ['int16', (2, 4096, 8), 2, 32768, 1024], + # ['int32', (2, 4096, 8), 2, 32768, 1024], + # ['int64', (2, 4096, 8), 2, 32768, 1024], + ['uint8', (2, 4096, 8), 2, 32768, 1024], + # ['uint16', (2, 4096, 8), 2, 32768, 1024], + # ['uint32', (2, 4096, 8), 2, 32768, 1024], + # ['uint64', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + # ['float32', (2, 4096, 8), 2, 32768, 1024], + # ['bfloat16', (2, 4096, 8), 2, 32768, 1024], +]) +def test_logical_or(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + torch_dtype = eval('torch.' + dtype) + np_x0 = test_common.generate_numpy(shape, dtype) + np_x1 = test_common.generate_numpy(shape, dtype) + x0 = torch.from_numpy(np_x0).to(torch_dtype).npu() + x1 = torch.from_numpy(np_x1).to(torch_dtype).npu() + # torch结果 + np_res = np.logical_or(np_x0, np_x1) + torch_res = torch.from_numpy(np_res).to(torch.bool) + # triton结果 + triton_res = torch.zeros(shape, dtype=torch.bool).npu() + triton_logical_or[ncore, 1, 1](x0, x1, triton_res, xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_lshift.py b/third_party/ascend/unittest/pytest_ut/test_lshift.py new file mode 100644 index 0000000000..6df7e5ed08 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_lshift.py @@ -0,0 +1,99 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import time +import test_common +import torch +import torch_npu + + +def standard_unary(x0, dtype): + res = x0 << 2 + return res + + +def standard_binary(x0, y0, dtype): + res = x0 + y0 + return res + + +@triton.jit +def triton_elementwise_unary(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + tmp = tl.cast(2, tl.int8) + ret = x << tmp + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +@triton.jit +def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + y = tl.load(in_ptr1 + idx_block, mask=idx_block < N) + ret = x + y + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +types = [ + # (torch.float32, 'float32'), + # (torch.float16, 'float16'), + # (torch.bfloat16, 'bfloat16'), + (torch.int8, 'int8'), + # (torch.int16, 'int16'), + # (torch.int32, 'int32'), + # (torch.int64, 'int64'), +] + +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('N,NUMEL', shapes) +def test_elementwsie_common(dtype, sigtype, N, NUMEL): + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == "int64": + N = map_for_64_t[N] if N in map_for_64_t else N + + print(f"elementwise : ({N},) {dtype} {sigtype}") + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standard_unary(x0, dtype) + x0 = x0.npu() + # print(ans) + + out = torch.zeros((N, ), dtype=dtype).npu() + triton_elementwise_unary[1, 1, 1](x0, out, N=N, NUMEL=NUMEL, debug=True) + # print(out) + + test_common.validate_cmp(sigtype, out, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_lt.py b/third_party/ascend/unittest/pytest_ut/test_lt.py new file mode 100644 index 0000000000..ab790b05e5 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_lt.py @@ -0,0 +1,60 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common + + +def torch_lt(x0, x1): + return x0 < x1 + + +@triton.jit +def triton_lt(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x_index = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + x_index, None) + tmp1 = tl.load(in_ptr1 + x_index, None) + tmp2 = tmp0 < tmp1 + tl.store(out_ptr0 + x_index, tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (32, ), 1, 32, 32], +]) +def test_lt(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + torch_res = torch_lt(x0, x1).to(eval('torch.' + dtype)) + # triton结果 + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_lt[ncore, 1, 1](x0, x1, triton_res, xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_makeblockptr_negative_padding.py b/third_party/ascend/unittest/pytest_ut/test_makeblockptr_negative_padding.py new file mode 100644 index 0000000000..4d339997ab --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_makeblockptr_negative_padding.py @@ -0,0 +1,136 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is herey_size granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import pytest +import test_common + + +@triton.jit +def negative_padding_with_load_kernel( + input_ptr, + output_ptr, + x_offset: tl.constexpr, + y_offset: tl.constexpr, + x_size: tl.constexpr, + y_size: tl.constexpr, +): + in_ptr = tl.make_block_ptr( + base=input_ptr, + shape=(x_size, y_size), + strides=(y_size, 1), + offsets=(x_offset, y_offset), + block_shape=(x_size, y_size), + order=(1, 0), + ) + out_ptr = tl.make_block_ptr( + base=output_ptr, + shape=(x_size, y_size), + strides=(y_size, 1), + offsets=(0, 0), + block_shape=(x_size, y_size), + order=(1, 0), + ) + in_val = tl.load(in_ptr, boundary_check=(0, 1), padding_option="zero") + tl.store(out_ptr, in_val) + + +@triton.jit +def negative_padding_with_store_kernel( + input_ptr, + output_ptr, + x_offset: tl.constexpr, + y_offset: tl.constexpr, + x_size: tl.constexpr, + y_size: tl.constexpr, +): + in_ptr = tl.make_block_ptr( + base=input_ptr, + shape=(x_size, y_size), + strides=(y_size, 1), + offsets=(0, 0), + block_shape=(x_size, y_size), + order=(1, 0), + ) + out_ptr = tl.make_block_ptr( + base=output_ptr, + shape=(x_size, y_size), + strides=(y_size, 1), + offsets=(x_offset, y_offset), + block_shape=(x_size, y_size), + order=(1, 0), + ) + in_val = tl.load(in_ptr) + tl.store(out_ptr, in_val, boundary_check=(0, 1)) + + +@pytest.mark.parametrize('param_list', [(8, 8), (16, 16), (32, 32), (64, 64)]) +def test_makeblockptr_load_with_negative_padding(param_list): + shape = param_list + torch.manual_seed(1) + x_offset = torch.randint(shape[0], size=()).item() + # y_offset = torch.randint(shape[1], size=()).item() + y_offset = 0 + input_tensor = torch.arange(start=1, end=shape[0] * shape[1] + 1, dtype=torch.int32).view(shape).npu() + output = torch.zeros(shape, dtype=torch.int32).npu() + negative_padding_with_load_kernel[(1, )]( + input_tensor, + output, + -x_offset, + -y_offset, + shape[0], + shape[1], + ) + output_ref = torch.zeros((shape[0] + x_offset, shape[1] + y_offset), dtype=torch.int32).cpu() + output_subview = torch.narrow(output_ref, 0, x_offset, shape[0]) + output_subview = torch.narrow(output_subview, 1, y_offset, shape[1]) + output_subview.copy_(input_tensor) + output_ref = torch.narrow(output_ref, 0, 0, shape[0]) + output_ref = torch.narrow(output_ref, 1, 0, shape[1]) + test_common.validate_cmp("int32", output, output_ref) + + +@pytest.mark.parametrize('param_list', [(8, 8), (16, 16), (32, 32), (64, 64)]) +def test_makeblockptr_store_with_negative_padding(param_list): + shape = param_list + torch.manual_seed(1) + x_offset = torch.randint(shape[0], size=()).item() + # y_offset = torch.randint(shape[1], size=()).item() + y_offset = 0 + input_tensor = torch.arange(start=1, end=shape[0] * shape[1] + 1, dtype=torch.int32).view(shape).npu() + output = torch.zeros(shape, dtype=torch.int32).npu() + negative_padding_with_store_kernel[(1, )]( + input_tensor, + output, + -x_offset, + -y_offset, + shape[0], + shape[1], + ) + output_ref = torch.zeros(shape, dtype=torch.int32).cpu() + input_subview = torch.narrow(input_tensor, 0, x_offset, shape[0] - x_offset) + input_subview = torch.narrow(input_subview, 1, y_offset, shape[1] - y_offset) + output_subview = torch.narrow(output_ref, 0, 0, shape[0] - x_offset) + output_subview = torch.narrow(output_subview, 1, 0, shape[1] - y_offset) + output_subview.copy_(input_subview) + test_common.validate_cmp("int32", output, output_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_makeblockptr_permute.py b/third_party/ascend/unittest/pytest_ut/test_makeblockptr_permute.py new file mode 100644 index 0000000000..5b5ffe1de7 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_makeblockptr_permute.py @@ -0,0 +1,115 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import pytest +import test_common + + +@pytest.mark.parametrize('shape', [(1, 4, 2)]) +@pytest.mark.parametrize('permute_order', [(2, 0, 1)]) +def test_makeblockptr_order(shape, permute_order): + + @triton.jit + def triton_kernel( + in0_ptr: tl.tensor, # of tl.pointer_type + out0_ptr: tl.tensor, # of tl.pointer_type + in0_stride0: int, + in0_stride1: int, + in0_stride2: int, # strides for in0 + in0_stride_order0: tl.constexpr, + in0_stride_order1: tl.constexpr, + in0_stride_order2: tl.constexpr, # stride order for in0 + out0_stride0: int, + out0_stride1: int, + out0_stride2: int, # strides for out0 + out0_stride_order0: tl.constexpr, + out0_stride_order1: tl.constexpr, + out0_stride_order2: tl.constexpr, # stride order for out0 + s0: int, + s1: int, + s2: int, + tile_size0: tl.constexpr, + tile_size1: tl.constexpr, + tile_size2: tl.constexpr, + ): + tile_id0 = tl.program_id(axis=0) + tile_id1 = tl.program_id(axis=1) + tile_id2 = tl.program_id(axis=2) + offset0 = (tile_id0 * tile_size0).to(tl.int32) + offset1 = (tile_id1 * tile_size1).to(tl.int32) + offset2 = (tile_id2 * tile_size2).to(tl.int32) + in0_bptr = tl.make_block_ptr(in0_ptr, (s0, s1, s2), (in0_stride0, in0_stride1, in0_stride2), + (offset0, offset1, offset2), (tile_size0, tile_size1, tile_size2), + order=(in0_stride_order0, in0_stride_order1, in0_stride_order2)) + in0 = tl.load(in0_bptr, boundary_check=(in0_stride_order0, in0_stride_order1, + in0_stride_order2)).to(in0_ptr.type.element_ty) + + out0 = in0 + + out0_bptr = tl.make_block_ptr(out0_ptr, (s0, s1, s2), (out0_stride0, out0_stride1, out0_stride2), + (offset0, offset1, offset2), (tile_size0, tile_size1, tile_size2), + order=(out0_stride_order0, out0_stride_order1, out0_stride_order2)) + tl.store(out0_bptr, out0.to(out0_bptr.type.element_ty), + boundary_check=(out0_stride_order0, out0_stride_order1, out0_stride_order2)) + + def triton_func(in0: torch.Tensor, permute_order): + # in fact, it adjusts the layout metadata instead of doing a real permutation. + in0_permuted_tmp = in0.permute(permute_order) + in0_permuted_shape = in0_permuted_tmp.size() + in0_permuted_strides = in0_permuted_tmp.stride() + in0_stride_order = [len(permute_order) - 1 - i for i in permute_order] + shape = (in0_permuted_shape[0], in0_permuted_shape[1], in0_permuted_shape[2]) + tile_sizes = (shape[0], shape[1], shape[2]) + out0 = torch.empty(shape, dtype=in0.dtype, device=in0.device) + out0_strides = out0.stride() + out0_stride_order = [len(permute_order) - 1 - i for i in range(len(permute_order))] + grid = (shape[0] // tile_sizes[0], shape[1] // tile_sizes[1], shape[2] // tile_sizes[2]) + triton_kernel[grid]( + in0, + out0, + in0_permuted_strides[0], + in0_permuted_strides[1], + in0_permuted_strides[2], # stride for in0 + in0_stride_order[0], + in0_stride_order[1], + in0_stride_order[2], # stride order for in0 + out0_strides[0], + out0_strides[1], + out0_strides[2], # stride for out0 + out0_stride_order[0], + out0_stride_order[1], + out0_stride_order[2], # stride orderfor out0 + shape[0], + shape[1], + shape[2], # task indexing space + tile_size0=tile_sizes[0], + tile_size1=tile_sizes[1], + tile_size2=tile_sizes[2], + ) + return out0 + + x0 = torch.randint(0, 9, shape, dtype=torch.int32).npu() + torch_ref = torch.permute(x0, permute_order) + triton_cal = triton_func(x0, permute_order) + test_common.validate_cmp("int32", triton_cal, torch_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_max_constancy.py b/third_party/ascend/unittest/pytest_ut/test_max_constancy.py new file mode 100644 index 0000000000..6798790144 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_max_constancy.py @@ -0,0 +1,27 @@ +import torch +import triton +import triton.language as tl +import triton.language.extra.cann.extension as extension +import pytest + + +@triton.jit +def compile_hint_kernel(input_ptr, output_ptr, n_elements: tl.constexpr, SIZE: tl.constexpr): + offsets = tl.arange(0, n_elements) + val = tl.load(input_ptr + offsets) + val = tl.max_constancy(val, SIZE) + tl.store(output_ptr + offsets, val) + + +@pytest.mark.parametrize('sigtype', [ + 'int32', + #'int64', 'int16', 'int8', + #'uint8', 'uint16', 'uint32', 'uint64', + #'float32', 'float16', 'bfloat16', 'bool' +]) +def test_compile_hint(sigtype): + n_elements = 10 + dtype = eval(f"torch.{sigtype}") + x = torch.ones((n_elements, ), dtype=dtype).npu() + y = torch.zeros((n_elements, ), dtype=dtype).npu() + compile_hint_kernel[(1, )](x, y, n_elements, 1) diff --git a/third_party/ascend/unittest/pytest_ut/test_max_contiguous.py b/third_party/ascend/unittest/pytest_ut/test_max_contiguous.py new file mode 100644 index 0000000000..c5c2268250 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_max_contiguous.py @@ -0,0 +1,26 @@ +import torch +import triton +import triton.language as tl +import pytest + + +@triton.jit +def compile_hint_kernel(input_ptr, output_ptr, n_elements: tl.constexpr, SIZE: tl.constexpr): + offsets = tl.arange(0, n_elements) + val = tl.load(input_ptr + offsets) + val = tl.max_contiguous(val, SIZE) + tl.store(output_ptr + offsets, val) + + +@pytest.mark.parametrize('sigtype', [ + 'int32', + #'int64', 'int16', 'int8', + #'uint8', 'uint16', 'uint32', 'uint64', + #'float32', 'float16', 'bfloat16', 'bool' +]) +def test_compile_hint(sigtype): + n_elements = 10 + dtype = eval(f"torch.{sigtype}") + x = torch.ones((n_elements, ), dtype=dtype).npu() + y = torch.zeros((n_elements, ), dtype=dtype).npu() + compile_hint_kernel[(1, )](x, y, n_elements, 1) diff --git a/third_party/ascend/unittest/pytest_ut/test_max_dim0.py b/third_party/ascend/unittest/pytest_ut/test_max_dim0.py new file mode 100644 index 0000000000..3ebcd9e225 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_max_dim0.py @@ -0,0 +1,130 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time + +import torch +import torch_npu +import test_common + + +def standard_max(x0, dim, dtype): + # fix with aclnnMaxDim support list:[DT_FLOAT,DT_FLOAT16,DT_INT64,DT_BOOL,DT_BFLOAT16,]. + if x0.dtype == torch.int8: + x0 = x0.to(torch.int64) + res, index = torch.max(x0, dim) + return res.to(dtype) + + +@triton.jit +def triton_max_dim0(in_ptr0, out_ptr0, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, MNUMEL) + nblk_idx = tl.arange(0, NNUMEL) + + mmask = mblk_idx < M + nmask = nblk_idx < N + + mask = (mmask[:, None]) & (nmask[None, :]) + + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + + x = tl.load(in_ptr0 + idx, mask=mask, other=-float('inf')) + + ret = tl.max(x, 0) + + tl.store(out_ptr0 + nblk_idx, ret, mask=nmask) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + # (torch.bfloat16,'bfloat16'), TODO: waiting for supporting or testing + (torch.int8, 'int8'), + # (torch.int16,'int16'), TODO: waiting for supporting or testing + # (torch.int32,'int32'), TODO: waiting for supporting or testing + # (torch.int64,'int64'), TODO: waiting for supporting or testing +] + +# if shape axis = 32/256 , then actual shape = axis/element_size() +shapes = [ + (57, 3, 64, 16), + (57, -32, 64, 32), + (57, 37, 64, 64), + (57, -256, 64, 256), + (57, 263, 64, 512), + (64, 3, 64, 16), + (64, -32, 64, 32), + (64, 37, 64, 64), + (64, -256, 64, 256), + (64, 263, 64, 512), + (3, 3, 8, 8), + (-32, 3, 32, 8), + (37, 3, 64, 8), + (-256, 3, 256, 8), + (263, 3, 512, 8), + (3, 1, 8, 8), + (-32, 1, 32, 8), + (37, 1, 64, 8), + (-256, 1, 256, 8), + (263, 1, 512, 8), +] + +map_for_64_t = {37: (31, 32), 263: (107, 128)} +map_for_32_t = {263: (137, 256)} + + +# @pytest.mark.parametrize('dtype, sigtype',[(torch.float32,'float32'),]) +@pytest.mark.parametrize('M, N, MNUMEL, NNUMEL', [(64, -32, 64, 32)]) +@pytest.mark.parametrize('dtype, sigtype', types) +# @pytest.mark.parametrize('M, N, MNUMEL, NNUMEL',shapes) +def test_max_dim0(dtype, sigtype, M, N, MNUMEL, NNUMEL): + + M = (-M) // torch.tensor(0, dtype=dtype).element_size() if M < 0 else M + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + M = map_for_64_t[M][0] if M in map_for_64_t else M + MNUMEL = map_for_64_t[M][1] if M in map_for_64_t else MNUMEL + N = map_for_64_t[N][0] if N in map_for_64_t else N + NNUMEL = map_for_64_t[N][1] if N in map_for_64_t else NNUMEL + + elif sigtype == 'float32' or sigtype == 'bfloat16' or sigtype == 'int32': + M = map_for_32_t[M][0] if M in map_for_32_t else M + MNUMEL = map_for_32_t[M][1] if M in map_for_32_t else MNUMEL + N = map_for_32_t[N][0] if N in map_for_32_t else N + NNUMEL = map_for_32_t[N][1] if N in map_for_32_t else NNUMEL + + print(f"max : ({M}, {N}) {dtype} {sigtype}") + x0 = test_common.generate_tensor(shape=(M, N), dtype=sigtype) + + ans = standard_max(x0, 0, dtype) + + x0 = x0.npu() + print(ans) + + output = torch.zeros((N, ), dtype=dtype).npu() + triton_max_dim0[1, 1, 1](x0, output, M=M, N=N, MNUMEL=MNUMEL, NNUMEL=NNUMEL) + print(output) + + test_common.validate_cmp(sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_max_dim1.py b/third_party/ascend/unittest/pytest_ut/test_max_dim1.py new file mode 100644 index 0000000000..77999b76fe --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_max_dim1.py @@ -0,0 +1,137 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time + +import torch +import torch_npu +import test_common + + +def standard_max(x0, dim, dtype): + # fix with aclnnMaxDim support list:[DT_FLOAT,DT_FLOAT16,DT_INT64,DT_BOOL,DT_BFLOAT16,]. + if x0.dtype == torch.int8: + x0 = x0.to(torch.int64) + res, index = torch.max(x0, dim) + return res.to(dtype) + + +@triton.jit +def triton_max_dim1(in_ptr0, out_ptr0, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, MNUMEL) + nblk_idx = tl.arange(0, NNUMEL) + + mmask = mblk_idx < M + nmask = nblk_idx < N + + mask = (mmask[:, None]) & (nmask[None, :]) + + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + + if in_ptr0.dtype == tl.int8: + padding = -128 + else: + padding = -float('inf') + + x = tl.load(in_ptr0 + idx, mask=mask, other=padding) + + ret = tl.max(x, 1) + + tl.store(out_ptr0 + mblk_idx, ret, mask=mmask) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + # (torch.bfloat16,'bfloat16'), waiting for supporting or testing + (torch.int8, 'int8'), + # (torch.int16,'int16'), waiting for supporting or testing + # (torch.int32,'int32'), waiting for supporting or testing + # (torch.int64,'int64'), waiting for supporting or testing +] + +# if shape axis = 32/256 , then actual shape = axis/element_size() +shapes = [ + (57, 3, 64, 16), + (57, -32, 64, 32), + (57, 37, 64, 64), + (57, -256, 64, 256), + (57, 263, 64, 512), + (64, 3, 64, 16), + (64, -32, 64, 32), + (64, 37, 64, 64), + (64, -256, 64, 256), + (64, 263, 64, 512), + (3, 3, 8, 8), + (-32, 3, 32, 8), + (37, 3, 64, 8), + (-256, 3, 256, 8), + (263, 3, 512, 8), + (3, 1, 8, 8), + (-32, 1, 32, 8), + (37, 1, 64, 8), + (-256, 1, 256, 8), + (263, 1, 512, 8), +] + +map_for_64_t = {37: (31, 32), 263: (107, 128)} +map_for_32_t = {263: (137, 256)} + + +@pytest.mark.parametrize('M, N, MNUMEL, NNUMEL', [ + (64, -32, 64, 32), +]) +# @pytest.mark.parametrize('M, N',[(263,3),(-256,3)]) +@pytest.mark.parametrize('dtype, sigtype', types) +# @pytest.mark.parametrize('M, N, MNUMEL, NNUMEL',shapes) +def test_max_dim1(dtype, sigtype, M, N, MNUMEL, NNUMEL): + + M = (-M) // torch.tensor(0, dtype=dtype).element_size() if M < 0 else M + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + M = map_for_64_t[M][0] if M in map_for_64_t else M + MNUMEL = map_for_64_t[M][1] if M in map_for_64_t else MNUMEL + N = map_for_64_t[N][0] if N in map_for_64_t else N + NNUMEL = map_for_64_t[N][1] if N in map_for_64_t else NNUMEL + + elif sigtype == 'float32' or sigtype == 'bfloat16' or sigtype == 'int32': + M = map_for_32_t[M][0] if M in map_for_32_t else M + MNUMEL = map_for_32_t[M][1] if M in map_for_32_t else MNUMEL + N = map_for_32_t[N][0] if N in map_for_32_t else N + NNUMEL = map_for_32_t[N][1] if N in map_for_32_t else NNUMEL + + print(f"max : ({M}, {N}) {dtype} {sigtype}") + x0 = test_common.generate_tensor(shape=(M, N), dtype=sigtype) + + ans = standard_max(x0, 1, dtype) + + x0 = x0.npu() + print(ans) + + output = torch.zeros((M, ), dtype=dtype).npu() + triton_max_dim1[1, 1, 1](x0, output, M=M, N=N, MNUMEL=MNUMEL, NNUMEL=NNUMEL) + print(output) + + test_common.validate_cmp(sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_max_propagate_nan.py b/third_party/ascend/unittest/pytest_ut/test_max_propagate_nan.py new file mode 100644 index 0000000000..198a37180d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_max_propagate_nan.py @@ -0,0 +1,20 @@ +import torch +import torch_npu +import triton +import triton.language as tl +import pytest + + +def test_max_propagate_nan(): + + @triton.jit + def func(in_ptr, out_ptr): + a = tl.load(in_ptr + tl.arange(0, 8)[:, None] * 8 + tl.arange(0, 8)[None, :]) + a = tl.max(a, 0, propagate_nan=True) + tl.store(out_ptr + tl.arange(0, 8), a) + + a = torch.randn((8, 8), device="npu") + std = a.max(0)[0] + ans = torch.zeros((8, ), dtype=torch.float32, device="npu") + func[1, 1, 1](a, ans) + torch.testing.assert_close(std, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_max_vector.py b/third_party/ascend/unittest/pytest_ut/test_max_vector.py new file mode 100644 index 0000000000..5e3d8f8fa9 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_max_vector.py @@ -0,0 +1,89 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + + +def standard_(x0, dtype): + res, index = torch.max(x0, 0, keepdim=True) + return res + + +@triton.jit +def triton_max_vector(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + + if in_ptr0.dtype == tl.int8: + padding = -128 + else: + padding = -float('inf') + + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N, other=padding) + ret = tl.max(x, 0) + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < 1) + + +types = [ + (torch.float32, 'float32'), + # (torch.float16,'float16'), TODO : fix reduceConverter bug + # (torch.bfloat16,'bfloat16'), waiting for supporting or testing + # (torch.int8,'int8'), TODO : fix compiler bug + # (torch.int16,'int16'), waiting for supporting or testing + # (torch.int32,'int32'), waiting for supporting or testing + # (torch.int64,'int64'), waiting for supporting or testing +] + +# if shape axis = 32/256 , then actual shape = axis/element_size() +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.skip(reason="randomly failed") +@pytest.mark.parametrize('dtype, sigtype', types) +@pytest.mark.parametrize('N, NUMEL', shapes) +def test_reduce_dim0_common(dtype, sigtype, N, NUMEL): + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + N = map_for_64_t[N] if N in map_for_64_t else N + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standard_(x0, dtype) + x0 = x0.npu() + + output = torch.zeros((1, ), dtype=dtype).npu() + triton_max_vector[1, 1, 1](x0, output, N=N, NUMEL=NUMEL, debug=True) + + test_common.validate_cmp(sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_maximum.py b/third_party/ascend/unittest/pytest_ut/test_maximum.py new file mode 100644 index 0000000000..abdd398e14 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_maximum.py @@ -0,0 +1,61 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common + + +def torch_maximum(x0, x1): + res = torch.maximum(x0, x1) + return res + + +@triton.jit +def triton_maximum(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + x_index = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = x_index < xnumel + tmp0 = tl.load(in_ptr0 + x_index, xmask) + tmp1 = tl.load(in_ptr1 + x_index, xmask) + tmp2 = tl.maximum(tmp0, tmp1) + tl.store(out_ptr0 + x_index, tmp2, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], +]) +def test_maximum(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + torch_res = torch_maximum(x0, x1) + # triton结果 + triton_res = test_common.generate_tensor(shape, dtype).npu() + triton_maximum[ncore, 1, 1](x0, x1, triton_res, x0.numel(), xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_mean_dim0.py b/third_party/ascend/unittest/pytest_ut/test_mean_dim0.py new file mode 100644 index 0000000000..f317ca7f74 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_mean_dim0.py @@ -0,0 +1,145 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time + +import torch +import torch_npu +import test_common + + +def standard_mean(x0, dim, dtype): + res = torch.mean(x0, dim, dtype=dtype) + return res + + +@triton.jit +def triton_mean_dim0(in_ptr0, out_ptr0, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, MNUMEL) + nblk_idx = tl.arange(0, NNUMEL) + + mmask = mblk_idx < M + nmask = nblk_idx < N + + mask = (mmask[:, None]) & (nmask[None, :]) + + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + + x = tl.load(in_ptr0 + idx, mask=mask, other=0) + + if x.dtype == tl.bfloat16: + ret = (tl.sum(x.to(tl.float32), 0) / M).to(tl.bfloat16) + elif x.dtype == tl.float16: + ret = tl.sum(x, 0) / M + else: + ret = tl.sum(x.to(tl.float32), 0) / M + + tl.store(out_ptr0 + nblk_idx, ret, mask=nmask) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + # (torch.bfloat16,'bfloat16'), TODO: waiting for supporting or testing + (torch.int8, 'int8'), + # (torch.int16,'int16'), TODO: waiting for supporting or testing + # (torch.int32,'int32'), TODO: waiting for supporting or testing + # (torch.int64,'int64'), TODO: waiting for supporting or testing +] + +# if shape axis = 32/256 , then actual shape = axis/element_size() +shapes = [ + (57, 3, 64, 16), + (57, -32, 64, 32), + (57, 37, 64, 64), + (57, -256, 64, 256), + (57, 263, 64, 512), + (64, 3, 64, 16), + (64, -32, 64, 32), + (64, 37, 64, 64), + (64, -256, 64, 256), + (64, 263, 64, 512), + (3, 3, 8, 8), + (-32, 3, 32, 8), + (37, 3, 64, 8), + (-256, 3, 256, 8), + (263, 3, 512, 8), + (3, 1, 8, 8), + (-32, 1, 32, 8), + (37, 1, 64, 8), + (-256, 1, 256, 8), + (263, 1, 512, 8), +] + +map_for_64_t = {37: (31, 32)} +map_for_32_t = {263: (137, 256)} + + +# @pytest.mark.parametrize('dtype, sigtype',[(torch.float32,'float32'),]) +@pytest.mark.parametrize('M, N, MNUMEL, NNUMEL', [ + (57, 3, 64, 16), + (64, -32, 64, 32), + (37, 3, 64, 8), + (263, 1, 512, 8), + (-256, 3, 256, 8), +]) +@pytest.mark.parametrize('dtype, sigtype', types) +# @pytest.mark.parametrize('M, N, MNUMEL, NNUMEL',shapes) +def test_mean_dim0(dtype, sigtype, M, N, MNUMEL, NNUMEL): + + M = (-M) // torch.tensor(0, dtype=dtype).element_size() if M < 0 else M + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + M = map_for_64_t[M][0] if M in map_for_64_t else M + MNUMEL = map_for_64_t[M][1] if M in map_for_64_t else MNUMEL + N = map_for_64_t[N][0] if N in map_for_64_t else N + NNUMEL = map_for_64_t[N][1] if N in map_for_64_t else NNUMEL + + res_dtype = dtype + res_sigtype = sigtype + should_cast_to_fp32 = ['int8', 'int16', 'int32', 'int64', 'float32', 'bfloat16'] + + if sigtype in should_cast_to_fp32: + M = map_for_32_t[M][0] if M in map_for_32_t else M + MNUMEL = map_for_32_t[M][1] if M in map_for_32_t else MNUMEL + N = map_for_32_t[N][0] if N in map_for_32_t else N + NNUMEL = map_for_32_t[N][1] if N in map_for_32_t else NNUMEL + if sigtype != 'bfloat16': + res_dtype = torch.float32 + res_sigtype = 'float32' + + print(f"sum : ({M}, {N}) {dtype} {sigtype}") + x0 = test_common.generate_tensor(shape=(M, N), dtype=sigtype) + + ans = standard_mean(x0, 0, res_dtype) + + x0 = x0.npu() + print(ans) + + output = torch.zeros((N, ), dtype=res_dtype).npu() + triton_mean_dim0[1, 1, 1](x0, output, M=M, N=N, MNUMEL=MNUMEL, NNUMEL=NNUMEL, debug=True) + print(output) + + test_common.validate_cmp(res_sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_mean_dim1.py b/third_party/ascend/unittest/pytest_ut/test_mean_dim1.py new file mode 100644 index 0000000000..e79f09c4ee --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_mean_dim1.py @@ -0,0 +1,145 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time + +import torch +import torch_npu +import test_common + + +def standard_mean(x0, dim, dtype): + res = torch.mean(x0, dim, dtype=dtype) + return res + + +@triton.jit +def triton_mean_dim1(in_ptr0, out_ptr0, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, MNUMEL) + nblk_idx = tl.arange(0, NNUMEL) + + mmask = mblk_idx < M + nmask = nblk_idx < N + + mask = (mmask[:, None]) & (nmask[None, :]) + + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + + x = tl.load(in_ptr0 + idx, mask=mask, other=0) + + if x.dtype == tl.bfloat16: + ret = (tl.sum(x.to(tl.float32), 1) / N).to(tl.bfloat16) + elif x.dtype == tl.float16: + ret = tl.sum(x, 1) / N + else: + ret = tl.sum(x.to(tl.float32), 1) / N + + tl.store(out_ptr0 + mblk_idx, ret, mask=mmask) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + # (torch.bfloat16,'bfloat16'), TODO: waiting for supporting or testing + (torch.int8, 'int8'), + # (torch.int16,'int16'), TODO: waiting for supporting or testing + # (torch.int32,'int32'), TODO: waiting for supporting or testing + # (torch.int64,'int64'), TODO: waiting for supporting or testing +] + +# if shape axis = 32/256 , then actual shape = axis/element_size() +shapes = [ + (57, 3, 64, 16), + (57, -32, 64, 32), + (57, 37, 64, 64), + (57, -256, 64, 256), + (57, 263, 64, 512), + (64, 3, 64, 16), + (64, -32, 64, 32), + (64, 37, 64, 64), + (64, -256, 64, 256), + (64, 263, 64, 512), + (3, 3, 8, 8), + (-32, 3, 32, 8), + (37, 3, 64, 8), + (-256, 3, 256, 8), + (263, 3, 512, 8), + (3, 1, 8, 8), + (-32, 1, 32, 8), + (37, 1, 64, 8), + (-256, 1, 256, 8), + (263, 1, 512, 8), +] + +map_for_64_t = {37: (31, 32)} +map_for_32_t = {263: (137, 256)} + + +# @pytest.mark.parametrize('dtype, sigtype',[(torch.float32,'float32'),]) +@pytest.mark.parametrize('M, N, MNUMEL, NNUMEL', [ + (57, 3, 64, 16), + (64, -32, 64, 32), + (37, 3, 64, 8), + (263, 1, 512, 8), + (-256, 3, 256, 8), +]) +@pytest.mark.parametrize('dtype, sigtype', types) +# @pytest.mark.parametrize('M, N, MNUMEL, NNUMEL',shapes) +def test_mean_dim1(dtype, sigtype, M, N, MNUMEL, NNUMEL): + + M = (-M) // torch.tensor(0, dtype=dtype).element_size() if M < 0 else M + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + M = map_for_64_t[M][0] if M in map_for_64_t else M + MNUMEL = map_for_64_t[M][1] if M in map_for_64_t else MNUMEL + N = map_for_64_t[N][0] if N in map_for_64_t else N + NNUMEL = map_for_64_t[N][1] if N in map_for_64_t else NNUMEL + + res_dtype = dtype + res_sigtype = sigtype + should_cast_to_fp32 = ['int8', 'int16', 'int32', 'int64', 'float32', 'bfloat16'] + + if sigtype in should_cast_to_fp32: + M = map_for_32_t[M][0] if M in map_for_32_t else M + MNUMEL = map_for_32_t[M][1] if M in map_for_32_t else MNUMEL + N = map_for_32_t[N][0] if N in map_for_32_t else N + NNUMEL = map_for_32_t[N][1] if N in map_for_32_t else NNUMEL + if sigtype != 'bfloat16': + res_dtype = torch.float32 + res_sigtype = 'float32' + + print(f"sum : ({M}, {N}) {dtype} {sigtype}") + x0 = test_common.generate_tensor(shape=(M, N), dtype=sigtype) + + ans = standard_mean(x0, 1, res_dtype) + + x0 = x0.npu() + print(ans) + + output = torch.zeros((M, ), dtype=res_dtype).npu() + triton_mean_dim1[1, 1, 1](x0, output, M=M, N=N, MNUMEL=MNUMEL, NNUMEL=NNUMEL, debug=True) + print(output) + + test_common.validate_cmp(res_sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_mean_vector.py b/third_party/ascend/unittest/pytest_ut/test_mean_vector.py new file mode 100644 index 0000000000..7c40a2b5cf --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_mean_vector.py @@ -0,0 +1,105 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time + +import torch +import torch_npu +import test_common + + +def standard_mean(x0, dim, dtype): + res = torch.mean(x0, dim, keepdim=True, dtype=dtype) + return res + + +@triton.jit +def triton_mean_dim0(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N, other=0) + + if x.dtype == tl.bfloat16: + ret = (tl.sum(x.to(tl.float32), 0) / N).to(tl.bfloat16) + elif x.dtype == tl.float16: + ret = tl.sum(x, 0) / N + else: + ret = tl.sum(x.to(tl.float32), 0) / N + + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < 1) + + +types = [ + (torch.float32, 'float32'), + # (torch.float16,'float16'), TODO: should fix reduceConverter's bug + # (torch.bfloat16,'bfloat16'), TODO: waiting for supporting or testing + (torch.int8, 'int8'), + # (torch.int16,'int16'), TODO: waiting for supporting or testing + # (torch.int32,'int32'), TODO: waiting for supporting or testing + # (torch.int64,'int64'), TODO: waiting for supporting or testing +] + +# if shape axis = 32/256 , then actual shape = axis/element_size() +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: (31, 32)} + + +@pytest.mark.parametrize('dtype, sigtype', types) +@pytest.mark.parametrize('N, NUMEL', shapes) +def test_mean_dim0(dtype, sigtype, N, NUMEL): + + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + N = map_for_64_t[N][0] if N in map_for_64_t else N + NUMEL = map_for_64_t[N][1] if N in map_for_64_t else NUMEL + + res_dtype = dtype + res_sigtype = sigtype + should_cast_to_fp32 = ['int8', 'int16', 'int32', 'int64', 'float32'] + + if sigtype in should_cast_to_fp32: + res_dtype = torch.float32 + res_sigtype = 'float32' + + print(f"sum : ({N},) {dtype} {sigtype}") + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standard_mean(x0, 0, res_dtype) + + x0 = x0.npu() + print(ans) + + output = torch.zeros((1, ), dtype=res_dtype).npu() + triton_mean_dim0[1, 1, 1](x0, output, N=N, NUMEL=NUMEL, debug=True) + print(output) + + test_common.validate_cmp(res_sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_min_dim0.py b/third_party/ascend/unittest/pytest_ut/test_min_dim0.py new file mode 100644 index 0000000000..41125de925 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_min_dim0.py @@ -0,0 +1,152 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time + +import torch +import torch_npu +import test_common + + +@pytest.mark.skip(reason="to be supported by bishengir-compile") +def test_min_dim0_3d(): + + def torch_func(x, dim): + res = torch.min(x, dim) + return res + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N0: tl.constexpr, N1: tl.constexpr, N2: tl.constexpr): + idx0 = tl.arange(0, N0) + idx1 = tl.arange(0, N1) + idx2 = tl.arange(0, N2) + in_idx = idx2[None, None, :] + idx1[None, :, None] * N2 + idx0[:, None, None] * N2 * N1 + tmp0 = tl.load(in_ptr0 + in_idx) + tmp1 = tl.min(tmp0, 0) + out_idx = idx2[None, :] + idx1[:, None] * N2 + tl.store(out_ptr0 + out_idx, tmp1) + + def triton_func(x0, dim): + N0, N1, N2 = x0.size() + y0 = test_common.generate_tensor(shape=(N1, N2), dtype="float32").npu() + triton_kernel[1, 1, 1](y0, x0, N0, N1, N2) + return y0 + + dim = 0 + N0, N1, N2 = 1, 22, 13 + x0 = test_common.generate_tensor(shape=(N0, N1, N2), dtype="float32").npu() + torch_ref = torch_func(x0, dim) + triton_cal = triton_func(x0, dim) + test_common.validate_cmp("float32", triton_cal, torch_ref) + + +def standard_min(x0, dim, dtype): + # fix with aclnnMinDim support list:[DT_FLOAT,DT_FLOAT16,DT_INT64,DT_BOOL,DT_BFLOAT16,]. + if x0.dtype == torch.int8: + x0 = x0.to(torch.int64) + res, index = torch.min(x0, dim) + return res.to(dtype) + + +@triton.jit +def triton_min_dim0(in_ptr0, out_ptr0, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, MNUMEL) + nblk_idx = tl.arange(0, NNUMEL) + + mmask = mblk_idx < M + nmask = nblk_idx < N + + mask = (mmask[:, None]) & (nmask[None, :]) + + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + if in_ptr0.dtype == tl.int8: + padding = 127 + else: + padding = float('inf') + x = tl.load(in_ptr0 + idx, mask=mask, other=padding) + + ret = tl.min(x, 0) + + tl.store(out_ptr0 + nblk_idx, ret, mask=nmask) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + # (torch.bfloat16,'bfloat16'), TODO: waiting for supporting or testing + (torch.int8, 'int8'), + # (torch.int16,'int16'), TODO: waiting for supporting or testing + # (torch.int32,'int32'), TODO: waiting for supporting or testing + # (torch.int64,'int64'), TODO: waiting for supporting or testing +] + +# if shape axis = 32/256 , then actual shape = axis/element_size() +#shapes=[ +# (57,3,64,16), (57,-32,64,32), (57,37,64,64), (57,-256,64,256), (57,263,64,512), +# (64,3,64,16), (64,-32,64,32), (64,37,64,64), (64,-256,64,256), (64,263,64,512), +# (3,3,8,8), (-32,3,32,8), (37,3,64,8), (-256,3,256,8), (263,3,512,8), +# (3,1,8,8), (-32,1,32,8), (37,1,64,8), (-256,1,256,8), (263,1,512,8), +#] +shapes = [ + (64, -32, 64, 32), +] + +map_for_64_t = {37: (31, 32), 263: (107, 128)} +map_for_32_t = {263: (137, 256)} + + +# @pytest.mark.parametrize('dtype, sigtype',[(torch.float32,'float32'),]) +@pytest.mark.parametrize('M, N, MNUMEL, NNUMEL', [(64, -32, 64, 32)]) +@pytest.mark.parametrize('dtype, sigtype', types) +# @pytest.mark.parametrize('M, N, MNUMEL, NNUMEL',shapes) +def test_min_dim0(dtype, sigtype, M, N, MNUMEL, NNUMEL): + + M = (-M) // torch.tensor(0, dtype=dtype).element_size() if M < 0 else M + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + M = map_for_64_t[M][0] if M in map_for_64_t else M + MNUMEL = map_for_64_t[M][1] if M in map_for_64_t else MNUMEL + N = map_for_64_t[N][0] if N in map_for_64_t else N + NNUMEL = map_for_64_t[N][1] if N in map_for_64_t else NNUMEL + + elif sigtype == 'float32' or sigtype == 'bfloat16' or sigtype == 'int32': + M = map_for_32_t[M][0] if M in map_for_32_t else M + MNUMEL = map_for_32_t[M][1] if M in map_for_32_t else MNUMEL + N = map_for_32_t[N][0] if N in map_for_32_t else N + NNUMEL = map_for_32_t[N][1] if N in map_for_32_t else NNUMEL + + print(f"min : ({M}, {N}) {dtype} {sigtype}") + x0 = test_common.generate_tensor(shape=(M, N), dtype=sigtype) + + ans = standard_min(x0, 0, dtype) + + x0 = x0.npu() + print(ans) + + output = torch.zeros((N, ), dtype=dtype).npu() + triton_min_dim0[1, 1, 1](x0, output, M=M, N=N, MNUMEL=MNUMEL, NNUMEL=NNUMEL) + print(output) + + test_common.validate_cmp(sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_min_dim1.py b/third_party/ascend/unittest/pytest_ut/test_min_dim1.py new file mode 100644 index 0000000000..b2567e18d4 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_min_dim1.py @@ -0,0 +1,135 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time + +import torch +import torch_npu +import test_common + + +def standard_min(x0, dim, dtype): + # fix with aclnnMinDim support list:[DT_FLOAT,DT_FLOAT16,DT_INT64,DT_BOOL,DT_BFLOAT16,]. + if x0.dtype == torch.int8: + x0 = x0.to(torch.int64) + res, index = torch.min(x0, dim) + return res.to(dtype) + + +@triton.jit +def triton_min_dim1(in_ptr0, out_ptr0, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, MNUMEL) + nblk_idx = tl.arange(0, NNUMEL) + + mmask = mblk_idx < M + nmask = nblk_idx < N + + mask = (mmask[:, None]) & (nmask[None, :]) + + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + if in_ptr0.dtype == tl.int8: + padding = 127 + else: + padding = float('inf') + x = tl.load(in_ptr0 + idx, mask=mask, other=padding) + + ret = tl.min(x, 1) + + tl.store(out_ptr0 + mblk_idx, ret, mask=mmask) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + # (torch.bfloat16,'bfloat16'), waiting for supporting or testing + (torch.int8, 'int8'), + # (torch.int16,'int16'), waiting for supporting or testing + # (torch.int32,'int32'), waiting for supporting or testing + # (torch.int64,'int64'), waiting for supporting or testing +] + +# if shape axis = 32/256 , then actual shape = axis/element_size() +shapes = [ + (57, 3, 64, 16), + (57, -32, 64, 32), + (57, 37, 64, 64), + (57, -256, 64, 256), + (57, 263, 64, 512), + (64, 3, 64, 16), + (64, -32, 64, 32), + (64, 37, 64, 64), + (64, -256, 64, 256), + (64, 263, 64, 512), + (3, 3, 8, 8), + (-32, 3, 32, 8), + (37, 3, 64, 8), + (-256, 3, 256, 8), + (263, 3, 512, 8), + (3, 1, 8, 8), + (-32, 1, 32, 8), + (37, 1, 64, 8), + (-256, 1, 256, 8), + (263, 1, 512, 8), +] + +map_for_64_t = {37: (31, 32), 263: (107, 128)} +map_for_32_t = {263: (137, 256)} + + +@pytest.mark.parametrize('M, N, MNUMEL, NNUMEL', [ + (64, -32, 64, 32), +]) +# @pytest.mark.parametrize('M, N',[(263,3),(-256,3)]) +@pytest.mark.parametrize('dtype, sigtype', types) +# @pytest.mark.parametrize('M, N, MNUMEL, NNUMEL',shapes) +def test_min_dim1(dtype, sigtype, M, N, MNUMEL, NNUMEL): + + M = (-M) // torch.tensor(0, dtype=dtype).element_size() if M < 0 else M + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + M = map_for_64_t[M][0] if M in map_for_64_t else M + MNUMEL = map_for_64_t[M][1] if M in map_for_64_t else MNUMEL + N = map_for_64_t[N][0] if N in map_for_64_t else N + NNUMEL = map_for_64_t[N][1] if N in map_for_64_t else NNUMEL + + elif sigtype == 'float32' or sigtype == 'bfloat16' or sigtype == 'int32': + M = map_for_32_t[M][0] if M in map_for_32_t else M + MNUMEL = map_for_32_t[M][1] if M in map_for_32_t else MNUMEL + N = map_for_32_t[N][0] if N in map_for_32_t else N + NNUMEL = map_for_32_t[N][1] if N in map_for_32_t else NNUMEL + + print(f"min : ({M}, {N}) {dtype} {sigtype}") + x0 = test_common.generate_tensor(shape=(M, N), dtype=sigtype) + + ans = standard_min(x0, 1, dtype) + + x0 = x0.npu() + print(ans) + + output = torch.zeros((M, ), dtype=dtype).npu() + triton_min_dim1[1, 1, 1](x0, output, M=M, N=N, MNUMEL=MNUMEL, NNUMEL=NNUMEL) + print(output) + + test_common.validate_cmp(sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_min_vector.py b/third_party/ascend/unittest/pytest_ut/test_min_vector.py new file mode 100644 index 0000000000..7a0a6c9255 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_min_vector.py @@ -0,0 +1,88 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import test_common +import torch +import torch_npu + + +def standard_(x0, dtype): + res, index = torch.min(x0, 0, keepdim=True) + return res + + +@triton.jit +def triton_min_vector(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + if in_ptr0.dtype == tl.int8: + padding = 127 + else: + padding = float('inf') + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N, other=padding) + + ret = tl.min(x, 0) + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < 1) + + +types = [ + (torch.float32, 'float32'), + # (torch.float16,'float16'), TODO : fix reduceConverter bug + # (torch.bfloat16,'bfloat16'), waiting for supporting or testing + # (torch.int8,'int8'), TODO : fix compiler bug + # (torch.int16,'int16'), waiting for supporting or testing + # (torch.int32,'int32'), waiting for supporting or testing + # (torch.int64,'int64'), waiting for supporting or testing +] + +# if shape axis = 32/256 , then actual shape = axis/element_size() +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.skip(reason="randomly failed") +@pytest.mark.parametrize('dtype, sigtype', types) +@pytest.mark.parametrize('N, NUMEL', shapes) +def test_reduce_dim0_common(dtype, sigtype, N, NUMEL): + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + N = map_for_64_t[N] if N in map_for_64_t else N + + print(f"elementwise : ({N},) {dtype} {sigtype}") + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standard_(x0, dtype) + x0 = x0.npu() + + output = torch.zeros((1, ), dtype=dtype).npu() + triton_min_vector[1, 1, 1](x0, output, N=N, NUMEL=NUMEL, debug=True) + + test_common.validate_cmp(sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_minimum.py b/third_party/ascend/unittest/pytest_ut/test_minimum.py new file mode 100644 index 0000000000..8029cbfa6d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_minimum.py @@ -0,0 +1,61 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common + + +def torch_minimum(x0, x1): + res = torch.minimum(x0, x1) + return res + + +@triton.jit +def triton_minimum(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + x_index = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = x_index < xnumel + tmp0 = tl.load(in_ptr0 + x_index, xmask) + tmp1 = tl.load(in_ptr1 + x_index, xmask) + tmp2 = tl.minimum(tmp0, tmp1) + tl.store(out_ptr0 + x_index, tmp2, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], +]) +def test_minimum(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + y_ref = torch_minimum(x0, x1) + # triton结果 + y_cal = test_common.generate_tensor(shape, dtype).npu() + triton_minimum[ncore, 1, 1](x0, x1, y_cal, x0.numel(), xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_mod.py b/third_party/ascend/unittest/pytest_ut/test_mod.py new file mode 100644 index 0000000000..fc91c60afb --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_mod.py @@ -0,0 +1,80 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +def torch_pointwise(x0, x1, dtype): + if dtype == 'float16': + x0 = x0.to(torch.float32) + x1 = x1.to(torch.float32) + elif dtype == 'float32': + x0 = x0.to(torch.float64) + x1 = x1.to(torch.float64) + res = torch.div(x0, x1, rounding_mode="trunc") + res = x0 - x1 * res + return res + + +@triton.jit +def triton_mod(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop1 in range(loops1): + x0_prime = offset + (loop1 * XBLOCK_SUB) + base1 + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.load(in_ptr1 + (x0), None) + tmp2 = tmp0 % tmp1 + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], +]) +def test_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + if dtype == 'int8': + x0 = torch.randint(low=1, high=127, size=shape, dtype=eval('torch.' + dtype)).npu() + x1 = torch.randint(low=1, high=127, size=shape, dtype=eval('torch.' + dtype)).npu() + else: + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_pointwise(x0, x1, dtype) + if dtype == "float16": + y_ref = y_ref.to(torch.float16) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_mod[ncore, 1, 1](x0, x1, y_cal, xblock, xblock_sub) + #test_common.validate_cmp(dtype, y_cal, y_ref.npu()) + if dtype == 'int8': + torch.equal(y_cal, y_ref) + else: + res = torch.isclose(y_cal, y_ref, rtol=1e-3, atol=1e-3, equal_nan=True) + if not res.all(): + max_diff = torch.max((y_ref - y_cal)).item() + raise ValueError(f"Tensors are not close, diff is {max_diff}") diff --git a/third_party/ascend/unittest/pytest_ut/test_mul.py b/third_party/ascend/unittest/pytest_ut/test_mul.py new file mode 100644 index 0000000000..feb98bb668 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_mul.py @@ -0,0 +1,61 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import numpy as np +import torch +import pytest +import test_common + + +def torch_pointwise(x0, x1): + res = x0 * x1 + return res + + +@triton.jit +def triton_mul(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop1 in range(loops1): + x0_prime = offset + (loop1 * XBLOCK_SUB) + base1 + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.load(in_ptr1 + (x0), None) + tmp2 = tmp0 * tmp1 + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], + ['uint8', (2, 4096, 8), 2, 32768, 1024], +]) +def test_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_pointwise(x0, x1) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_mul[ncore, 1, 1](x0, x1, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_mul_reduce.py b/third_party/ascend/unittest/pytest_ut/test_mul_reduce.py new file mode 100644 index 0000000000..5ca4246331 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_mul_reduce.py @@ -0,0 +1,55 @@ +import pytest +import torch +import torch_npu +import triton +import triton.language as tl +import numpy as np + + +@triton.jit +def minimum(a, b): + ret = tl.minimum(a, b, tl.PropagateNan.ALL) + if a.dtype == tl.bfloat16: + ret = ret.to(tl.bfloat16) + return ret + + +@triton.jit +def triton_pw_rdc5d(in_ptr0, in_ptr1, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, + Z: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + kblk_idx = tl.arange(0, K) + zblk_idx = tl.arange(0, Z) + idx = (lblk_idx[:, None, None, None, None] * Z * K * N * M + mblk_idx[None, :, None, None, None] * Z * K * N + + nblk_idx[None, None, :, None, None] * Z * K + kblk_idx[None, None, None, :, None] * Z + + zblk_idx[None, None, None, None, :]) + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret0 = x0 * x1 + ret = tl.reduce(ret0, 4, minimum, keep_dims=True) + zblk_idx = tl.arange(0, 1) + odx = (lblk_idx[:, None, None, None, None] * K * N * M + mblk_idx[None, :, None, None, None] * K * N + + nblk_idx[None, None, :, None, None] * K + kblk_idx[None, None, None, :, None] + + zblk_idx[None, None, None, None, :]) + tl.store(out_ptr0 + odx, ret) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("shape", [(16, 1, 1, 1, 1)]) # L=16, others=1 +def test_pw_rdc5d(dtype, shape): + L, M, N, K, Z = shape + a = torch.randn(*shape, dtype=dtype, device='npu') + b = torch.randn(*shape, dtype=dtype, device='npu') + out = torch.empty(*shape, dtype=dtype, device='npu') + + expected = (a * b).to(dtype) + + triton_pw_rdc5d[(1, )](a, b, out, L=L, M=M, N=N, K=K, Z=Z) + + torch.testing.assert_close(out.cpu(), expected.cpu(), rtol=1e-3, atol=1e-3) + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/third_party/ascend/unittest/pytest_ut/test_multi_return.py b/third_party/ascend/unittest/pytest_ut/test_multi_return.py new file mode 100644 index 0000000000..1935e5225b --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_multi_return.py @@ -0,0 +1,599 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +from typing import Optional +import pytest +import torch +import torch_npu +import triton +from triton.language.math import tanh +import triton.language as tl + +device = 'npu' + + +@triton.jit +def triton_multi_return_cross_entropy_case( + X_ptr, + Y_ptr, + X_stride, + Y_stride, + n_cols, + ignore_index, + BLOCK_SIZE: tl.constexpr, +): + program_id = tl.program_id(0).to(tl.int64) + + Y_ptr += program_id * Y_stride + y = tl.load(Y_ptr) + X_ptr += program_id * X_stride + + if y == ignore_index: + for i in range(0, n_cols, BLOCK_SIZE): + X_offsets = i + tl.arange(0, BLOCK_SIZE) + tl.store(X_ptr + X_offsets, 0.0, mask=X_offsets < n_cols) + return + else: + for i in range(0, n_cols, BLOCK_SIZE): + X_offsets = i + tl.arange(0, BLOCK_SIZE) + tl.store(X_ptr + X_offsets, 0.0, mask=X_offsets < n_cols) + + +@pytest.mark.parametrize('param_list', [ + [1024, 152064, 11480, 'bfloat16', 'int32', 'float32'], +]) +def test_case(param_list): + N, M, V, dtype_1, dtype_2, dtype_3 = param_list + logits_chunk = torch.rand([N, M], device=device, dtype=eval('torch.' + dtype_1)) + target_chunk = torch.rand(N, device=device, dtype=eval('torch.' + dtype_2)) + loss_1d_slice = torch.rand(N, device=device, dtype=eval('torch.' + dtype_3)) + ce_weight = None + z_loss_1d_slice = None + softcap = None + return_z_loss = False + BLOCK_SIZE = N + triton_multi_return_cross_entropy_case[(N, )](X_ptr=logits_chunk, X_stride=logits_chunk.stride(-2), + Y_ptr=target_chunk, Y_stride=target_chunk.stride(-1), + n_cols=logits_chunk.stride(-2), ignore_index=-100, + BLOCK_SIZE=BLOCK_SIZE, num_warps=32) + + +@triton.jit +def liger_cross_entropy_kernel( + X_ptr, + X_stride, + Y_ptr, + Y_stride, + weight_ptr, + loss_ptr, + z_loss_ptr, + loss_stride, + n_cols, + n_non_ignore, + sum_non_ignore_weight, + weight_sum, + ignore_index, + lse_square_scale: tl.constexpr, + label_smoothing: tl.constexpr, + reduction: tl.constexpr, # set it as constexpr since reduction is always known at compile time + softcap, + RETURN_Z_LOSS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_SOFTCAPPING: tl.constexpr, +): + """ + This kernel computes both cross entropy loss and the gradient of the input. + + Parameters: + X_ptr: Pointer to input tensor. + X_stride (int): The stride of the input tensor. + Y_ptr: Pointer to target tensor. + Y_stride (int): The stride of the target tensor. + weight_ptr: Pointer to weight tensor. + loss_ptr: Pointer to tensor to store the loss. + z_loss_ptr: Pointer to tensor to store the z loss. No operation if RETURN_Z_LOSS is 0. + loss_stride (int): The stride of the loss tensor. + n_cols (int): The number of columns in the input tensor. + n_non_ignore (float): The number of non-ignored elements in the batch. + sum_non_ignore_weight (float): The sum of non-ignored target's weights in the batch. + weight_sum (float): The sum of weight tensor. + ignore_index (int): The index to ignore in the target. + label_smoothing (float): The amount of smoothing when computing the loss, where 0.0 means no smoothing. + lse_square_scale (float): The scaler of (logsumexp(_input)) ^ 2 adding to the loss for the stability of training. + reduction (str): The string for the reduction to apply + softcap (float): The upper threshold for scaling logits to the range (-softcap, +softcap). + RETURN_Z_LOSS (int): The boolean value to decide whether storing z loss to z_loss_ptr or not. It must be 0 or 1. + BLOCK_SIZE (int): The block size for Triton operations. + HAS_WEIGHT (bool): The boolean value to determine whether assigning weight to each of the classes. + HAS_SOFTCAPPING (bool): The boolean value to determine whether applying soft-capping or not. + """ + + # If B*T*V is too large, program_id * stride will overflow out of int32, so we convert to int64 + program_id = tl.program_id(0).to(tl.int64) + + # 1. Load Y_ptr first because if the target is ignore_index, we can return right away + Y_ptr += program_id * Y_stride + y = tl.load(Y_ptr) + + # 2. locate the start index + X_ptr += program_id * X_stride + + if y == ignore_index: + # set all X_ptr as 0 + for i in range(0, n_cols, BLOCK_SIZE): + X_offsets = i + tl.arange(0, BLOCK_SIZE) + tl.store(X_ptr + X_offsets, 0.0, mask=X_offsets < n_cols) + return + + loss_ptr += program_id * loss_stride + if RETURN_Z_LOSS: + z_loss_ptr += program_id * loss_stride + + if HAS_WEIGHT: + weight_y = tl.load(weight_ptr + y).cast(tl.float32) + + # Online softmax: 2 loads + 1 store (compared with 3 loads + 1 store for the safe softmax) + + # 3. [Online softmax] first pass: find max + sum + m = float("-inf") # m is the max value. use the notation from the paper + d = 0.0 # d is the sum. use the notation from the paper + ori_X_y = tl.load(X_ptr + y).cast(tl.float32) # we need to store the original value of X_y for the loss calculation + if HAS_SOFTCAPPING: + ori_X_y = softcap * tanh(ori_X_y / softcap) + + # Label smoothing is a general case of normal cross entropy + scaled_x_sum = 0.0 + eps = label_smoothing / n_cols + + for i in range(0, n_cols, BLOCK_SIZE): + X_offsets = i + tl.arange(0, BLOCK_SIZE) + X_block = tl.load( + X_ptr + X_offsets, + mask=X_offsets < n_cols, + other=float("-inf"), + # Ensure float32 precision for softmax calculation + ).cast(tl.float32) + if HAS_SOFTCAPPING: + X_block = softcap * tanh(X_block / softcap) + block_max = tl.max(X_block) + if label_smoothing > 0: + # scale X beforehand to avoid overflow + if HAS_WEIGHT: + weight_block = tl.load(weight_ptr + X_offsets, mask=X_offsets < n_cols) + scaled_x_sum += tl.sum(tl.where(X_offsets < n_cols, -eps * X_block * weight_block, 0.0)) + else: + scaled_x_sum += tl.sum(tl.where(X_offsets < n_cols, -eps * X_block, 0.0)) + m_new = tl.maximum(m, block_max) + d = d * tl.exp(m - m_new) + tl.sum(tl.exp(X_block - m_new)) + m = m_new + + # log (sum(e^(X_i))) = log (sum(e ^ (max(X) * e ^ (X_i - max(X))))) + # = log (e^(max(X)) * sum(e ^ (X_i - max(X)))) + # = max(X) + log (sum(e ^ (X_i - max(X)))) = m + log d + lse = m + tl.log(d) + + # 4. [Online Softmax] Second pass: compute gradients + for i in range(0, n_cols, BLOCK_SIZE): + X_offsets = i + tl.arange(0, BLOCK_SIZE) + X_block = tl.load( + X_ptr + X_offsets, + mask=X_offsets < n_cols, + other=float("-inf"), + # Ensure float32 precision for softmax calculation + ).cast(tl.float32) + if HAS_SOFTCAPPING: + intermediate = tanh(X_block / softcap) + X_block = softcap * intermediate + + if not HAS_WEIGHT: + X_block = tl.exp(X_block - m) / d + # derivative of z-loss: 2 * lse_square_scale * lse * softmax(x_i) + X_block += 2 * lse_square_scale * lse * X_block + # smoothing term + X_block += -eps + # special handle dx_y + X_block = tl.where(X_offsets != y, X_block, X_block - (1 - label_smoothing)) + # reduction scale + if reduction == "mean": + X_block = X_block / n_non_ignore + else: + weight_block = tl.load(weight_ptr + X_offsets, mask=X_offsets < n_cols) + softmax_X = tl.exp(X_block - m) / d + # derivative of original_loss + dloss_ori = (1 - label_smoothing) * softmax_X + # specially handle dx_y + dloss_ori = tl.where(X_offsets != y, dloss_ori, dloss_ori - (1 - label_smoothing)) + dloss_ori = dloss_ori * weight_y + # derivative of smooth_loss + dloss_smooth = eps * (-weight_block + softmax_X * weight_sum) + # derivative of z-loss + dz_loss = 2 * lse_square_scale * lse * softmax_X + # reduction scale + if reduction == "mean": + dloss_ori = dloss_ori / sum_non_ignore_weight + dloss_smooth = dloss_smooth / sum_non_ignore_weight + dz_loss = dz_loss / n_non_ignore + # derivative of total_loss + X_block = dloss_ori + dloss_smooth + dz_loss + + # chain rule softcapping + # d(softcap * tanh(x / softcap)) = (1 - tanh^2(x / softcap)) + if HAS_SOFTCAPPING: + X_block = X_block * (1 - intermediate * intermediate) + + tl.store(X_ptr + X_offsets, X_block, mask=X_offsets < n_cols) + + # We need tl.debug_barrier() to ensure the new result of X_ptr is written as mentioned in + tl.debug_barrier() + + # 5. Calculate the loss + + # loss = log (softmax(X_y)) = log ((e ^ (X_y - max(X)) / sum(e ^ (X - max(X)))) + # = (X_y - max(X)) - log(sum(e ^ (X - max(X)))) + # = X_y - m - log d = X_y - lse + # sum(e ^ (X - max(X))) must >= 1 because the max term is e ^ 0 = 1 + # So we can safely calculate log (softmax(X_y)) without overflow + loss = lse - ori_X_y + if HAS_WEIGHT: + loss = weight_y * loss + + # Original loss = H(q, p), with label smoothing regularization = H(q', p) and (label_smoothing / V) = eps + # H(q', p) = (1 - label_smoothing) * H(q, p) + label_smoothing * H(u, p) + # = (1 - label_smoothing) * H(q, p) + eps * sum(logsoftmax(x_i)) + # By using m (global max of xi) and d (sum of e^(xi-m)), we can simplify as: + # = (1 - label_smoothing) * H(q, p) + (sum(-eps * x_i) + label_smoothing * (m + logd)) + if label_smoothing > 0: + if HAS_WEIGHT: + smooth_loss = scaled_x_sum + eps * lse * weight_sum + else: + smooth_loss = scaled_x_sum + label_smoothing * lse + loss = loss * (1 - label_smoothing) + smooth_loss + + # An auxiliary loss, z_loss + z_loss = lse_square_scale * lse * lse + # Normalize the loss by the number of non-ignored elements if reduction is "mean" + if reduction == "mean": + if HAS_WEIGHT: + loss = loss / sum_non_ignore_weight + else: + loss = loss / n_non_ignore + z_loss = z_loss / n_non_ignore + loss += z_loss + + tl.store(loss_ptr, loss) + if RETURN_Z_LOSS: + tl.store(z_loss_ptr, z_loss) + + +@triton.jit +def element_mul_kernel( + X_ptr, + X_stride, + grad_output_ptr, + n_cols, + BLOCK_SIZE: tl.constexpr, +): + """ + This function multiplies each element of the tensor pointed by X_ptr with the value pointed by grad_output_ptr. + The multiplication is performed in-place on the tensor pointed by X_ptr. + + Parameters: + X_ptr: Pointer to the input tensor. + X_stride (int): The stride of the input tensor. + grad_output_ptr: Pointer to the gradient output value. + n_cols (int): The number of columns in the input tensor. + BLOCK_SIZE (int): The block size for Triton operations. + """ + + # Get the program ID and convert it to int64 to avoid overflow + program_id = tl.program_id(0).to(tl.int64) + + # Locate the start index + X_ptr += program_id * X_stride + + # Load the gradient output value + grad_output = tl.load(grad_output_ptr) + + # Perform the element-wise multiplication + for i in range(0, n_cols, BLOCK_SIZE): + X_offsets = i + tl.arange(0, BLOCK_SIZE) + X_block = tl.load(X_ptr + X_offsets, mask=X_offsets < n_cols) + tl.store(X_ptr + X_offsets, X_block * grad_output, mask=X_offsets < n_cols) + + +MAX_FUSED_SIZE = 65536 // 2 + + +def cross_entropy_forward( + _input, + target, + weight, + ignore_index, + lse_square_scale, + label_smoothing, + reduction, + softcap, + return_z_loss, +): + assert isinstance(return_z_loss, bool), f"return_z_loss must be True or False. Got: {return_z_loss}" + + BT, V = _input.shape + n_rows = BT + + BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V)) + # unreduced loss + loss_1d = torch.zeros(n_rows, dtype=_input.dtype, device=_input.device) + z_loss_1d = torch.zeros(n_rows, dtype=_input.dtype, device=_input.device) if return_z_loss else None + + target_mask = target != ignore_index + n_non_ignore = target_mask.sum().item() + assert (target * target_mask).max() < _input.shape[-1], ( + f"Target {target.max()} is out of bounds. Expected < {_input.shape[-1]}") + assert (target * target_mask).min() >= 0, f"Target {target.min()} is out of bounds. Expected >= 0" + sum_non_ignore_weight = n_non_ignore + weight_sum = 0.0 + if weight is not None: + assert weight.shape[0] == V, f"If given, weight has to be a Tensor of size V. Got: {weight.shape}" + assert torch.is_floating_point(weight), ( + f"If given, weight has to be a Tensor of floating point dtype. Got: {weight.dtype}") + sum_non_ignore_weight = torch.gather(weight, dim=0, index=target.masked_select(target_mask)).sum().item() + weight_sum = weight.sum().item() + # ensure weight is contiguous + if weight.stride(-1) != 1: + weight = weight.contiguous() + + # ensure _input and target are contiguous in the last dimension + if _input.stride(-1) != 1: + _input = _input.contiguous() + if target.stride(-1) != 1: + target = target.contiguous() + + # Here we use a trick to store X_ptr gradient in X_ptr so we can save memory + liger_cross_entropy_kernel[(n_rows, )]( + X_ptr=_input, + X_stride=_input.stride(-2), + Y_ptr=target, + Y_stride=target.stride(-1), # always 1 + weight_ptr=weight, # dummy if None + loss_ptr=loss_1d, + z_loss_ptr=z_loss_1d, + loss_stride=loss_1d.stride(-1), # always 1 + n_cols=V, + n_non_ignore=n_non_ignore, + sum_non_ignore_weight=sum_non_ignore_weight, + ignore_index=ignore_index, + weight_sum=weight_sum, + lse_square_scale=lse_square_scale, + label_smoothing=label_smoothing, + reduction=reduction, + softcap=softcap, + RETURN_Z_LOSS=return_z_loss, + BLOCK_SIZE=BLOCK_SIZE, + HAS_WEIGHT=True if weight is not None else False, + HAS_SOFTCAPPING=True if softcap is not None else False, + ) + + if reduction == "none": + loss = loss_1d + z_loss = z_loss_1d if return_z_loss else None + else: + loss = torch.sum(loss_1d) + z_loss = torch.sum(z_loss_1d) if return_z_loss else None + + return loss, z_loss, _input + + +def cross_entropy_backward(_input, grad_output): + # If cross entropy is the last layer, grad_output is 1.0. Skip the mul to save time + if torch.equal(grad_output, torch.tensor(1.0, device=grad_output.device)): + pass + # If reduction is 'none' + elif grad_output.ndim > 0: + _input = _input * grad_output.unsqueeze(dim=1) + # If reduction is ['mean', 'sum'], grad_output is just a scalar + # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place + # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton. + else: + BT, V = _input.shape + n_rows = BT + BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V)) + + element_mul_kernel[(n_rows, )]( + _input, + _input.stride(-2), + grad_output, + V, + BLOCK_SIZE=BLOCK_SIZE, + ) + + return _input + + +class LigerCrossEntropyFunction(torch.autograd.Function): + """ + This class implements a custom autograd function for the Liger Cross Entropy loss. + It overrides the forward and backward methods of the torch.autograd.Function class. + """ + + @staticmethod + def forward( + ctx, + _input: torch.Tensor, + target: torch.Tensor, + weight: Optional[torch.FloatTensor], + ignore_index: int = -100, + lse_square_scale: float = 0.0, + label_smoothing: float = 0.0, + reduction: str = "mean", + softcap: Optional[float] = None, + return_z_loss: bool = False, + ): + """ + The forward pass of the Liger Cross Entropy loss. + + Parameters: + ctx : The context object. + _input (tensor): The input tensor of shape (BT, V) where B is batch size, T is sequence length, V is vocab size. + target (tensor): The target tensor of shape (BT) where each value is in [0, V-1]. + weight(Tensor, optional): a manual rescaling weight given to each class. If given, has to be a Tensor of size V and floating point dtype + ignore_index (int): The index to ignore in the target. + lse_square_scale (float): The scaler of (logsumexp(_input)) ^ 2 adding to the loss for the stability of training. + label_smoothing (float): The amount of smoothing when computing the loss, where 0.0 means no smoothing. + reduction (str): The reduction to apply to the output: "none" | "mean | "sum". + softcap (Optional[float]): The upper threshold for scaling logits to the range (-softcap, +softcap). + return_z_loss (bool): When `return_z_loss` is `True`, returns (loss, z_loss) instead of (loss, None). Default: `False` + + Returns: + tuple: A tuple with the compouted losses with respect to loss and z loss. The elements are tensors or None. + """ + loss, z_loss, _input = cross_entropy_forward( + _input, + target, + weight, + ignore_index, + lse_square_scale, + label_smoothing, + reduction, + softcap, + return_z_loss, + ) + # If we don't detach the _input tensor, the memory will double + # Not sure why but seems that there will be a time both grad and value exist but in different location + ctx.save_for_backward(_input.detach()) + ctx.return_z_loss = return_z_loss + + return loss, z_loss + + @staticmethod + def backward(ctx, grad_output, grad_ouput2): + """ + The backward pass of the Liger Cross Entropy loss. + + Parameters: + ctx : The context object with saved tensors. + grad_output (tensor): The tensor containing the gradient of the loss with respect to the output. + grad_output2 (tenosr): No use. + Returns: + tuple: A tuple with the gradients with respect to the inputs. The elements are tensors or None. + """ + if ctx.return_z_loss: + del grad_ouput2 # z_loss is only for logging + + (_input, ) = ctx.saved_tensors + _input = cross_entropy_backward(_input, grad_output) + return ( + _input, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +def liger_cross_entropy( + input, + target, + weight=None, + size_average=None, + ignore_index: int = -100, + reduce=None, + reduction: str = "mean", + label_smoothing: float = 0.0, + lse_square_scale: float = 0.0, + softcap: Optional[float] = None, + return_z_loss: bool = False, +): + loss, z_loss = LigerCrossEntropyFunction.apply( + input, + target, + weight, + ignore_index, + lse_square_scale, + label_smoothing, + reduction, + softcap, + return_z_loss, + ) + if not return_z_loss: + return loss + return loss, z_loss + + +def _test_correctness_functional( + B, + T, + V, + scalar, + dtype, + atol, + rtol, +): + _input = torch.randn(B * T, V, device=device, dtype=dtype) * scalar + + x1 = _input.clone().requires_grad_(True) + x2 = _input.clone().requires_grad_(True) + + target = torch.randint(0, V, (B * T, ), device=device, dtype=torch.long) + + y1, y1_z = liger_cross_entropy( + x1, + target, + None, + ignore_index=0, + lse_square_scale=1e-4, + label_smoothing=0.1, + reduction="mean", + softcap=30.0, + return_z_loss=True, + ) + y2, y2_z = LigerCrossEntropyFunction.apply(x2, target, None, 0, 1e-4, 0.1, "mean", 30.0, True) + + assert torch.allclose(y1, y2, atol=atol, rtol=rtol) + assert torch.allclose(y1_z, y2_z, atol=atol, rtol=rtol) + + grad = torch.randn_like(y2) + + y1.backward(grad) + y2.backward(grad) + + assert torch.allclose(x1.grad, x2.grad, atol=atol, rtol=rtol) + + +@pytest.mark.parametrize( + "B, T, V", + [ + (2, 512, 4096), + ], +) +@pytest.mark.parametrize( + "scalar, dtype, atol, rtol", + [ + (1.0, torch.bfloat16, 1e-8, 5e-2), + (1.0, torch.float32, 1e-8, 1e-6), + (1.0, torch.int32, 0, 0), + ], +) +def test_correctness_functional(B, T, V, scalar, dtype, atol, rtol): + _test_correctness_functional(B, T, V, scalar, dtype, atol, rtol) diff --git a/third_party/ascend/unittest/pytest_ut/test_multibuffer.py b/third_party/ascend/unittest/pytest_ut/test_multibuffer.py new file mode 100644 index 0000000000..04d2471e22 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_multibuffer.py @@ -0,0 +1,73 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import os + +import pytest +import triton +import triton.language as tl +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +import triton.extension.buffer.language as bl +import triton.language.extra.cann.extension as al +from triton._C.libtriton import ir, buffer_ir +from triton._C.libtriton.ascend import ir as ascend_ir + +os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" + + +class Options: + num_warps = 4 + num_stages = 3 + num_ctas = 1 + cluster_dims = (1, 1, 1) + enable_fp_fusion = True + debug = False + + +def compile_kernel(kernel, signature, constants): + """Helper to compile a kernel to MLIR.""" + src = ASTSource(kernel, signature, constants) + context = ir.context() + ir.load_dialects(context) + buffer_ir.load_dialects(context) + ascend_ir.load_dialects(context) + module = ast_to_ttir(kernel, src, context, Options(), {}, {}) + return str(module) + + +@triton.jit +def multibuffer(XBLOCK: tl.constexpr): + buf = bl.alloc(tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.UB) + al.multibuffer(buf, 2) + + +def test_multibuffer(): + print("=" * 60) + print("Test 1: test_alloc_ub_multibuffer") + print("=" * 60) + mlir = compile_kernel(multibuffer, {}, {"XBLOCK": 256}) + print(f"Generated MLIR ({len(mlir)} chars):\n") + print(mlir) + + +# ============== Main for manual testing ============== +if __name__ == "__main__": + test_multibuffer() diff --git a/third_party/ascend/unittest/pytest_ut/test_multiple_of.py b/third_party/ascend/unittest/pytest_ut/test_multiple_of.py new file mode 100644 index 0000000000..6813ee8bc4 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_multiple_of.py @@ -0,0 +1,26 @@ +import torch +import triton +import triton.language as tl +import pytest + + +@triton.jit +def compile_hint_kernel(input_ptr, output_ptr, n_elements: tl.constexpr, SIZE: tl.constexpr): + offsets = tl.arange(0, n_elements) + val = tl.load(input_ptr + offsets) + val = tl.multiple_of(val, SIZE) + tl.store(output_ptr + offsets, val) + + +@pytest.mark.parametrize('sigtype', [ + 'int32', + #'int64', 'int16', 'int8', + #'uint8', 'uint16', 'uint32', 'uint64', + #'float32', 'float16', 'bfloat16', 'bool' +]) +def test_compile_hint(sigtype): + n_elements = 10 + dtype = eval(f"torch.{sigtype}") + x = torch.ones((n_elements, ), dtype=dtype).npu() + y = torch.zeros((n_elements, ), dtype=dtype).npu() + compile_hint_kernel[(1, )](x, y, n_elements, 1) diff --git a/third_party/ascend/unittest/pytest_ut/test_ne.py b/third_party/ascend/unittest/pytest_ut/test_ne.py new file mode 100644 index 0000000000..00ecc5b137 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_ne.py @@ -0,0 +1,76 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common + + +def torch_ne(x0, x1): + return x0 != x1 + + +def torch_ne_from_np(x0, x1, dtype): + return torch.from_numpy(x0 != x1).to(eval('torch.' + dtype)) + + +@triton.jit +def triton_ne(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x_index = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + x_index, None) + tmp1 = tl.load(in_ptr1 + x_index, None) + tmp2 = tmp0 != tmp1 + tl.store(out_ptr0 + x_index, tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + # ['float16', (2, 4096, 8), 2, 32768, 1024], + # ['bfloat16', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], + # ['int16', (2, 4096, 8), 2, 32768, 1024], + # ['int32', (2, 4096, 8), 2, 32768, 1024], + # ['int64', (2, 4096, 8), 2, 32768, 1024], + ['uint8', (2, 4096, 8), 2, 32768, 1024], + # ['uint16', (2, 4096, 8), 2, 32768, 1024], + # ['uint32', (2, 4096, 8), 2, 32768, 1024], + # ['uint64', (2, 4096, 8), 2, 32768, 1024], +]) +def test_ne(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + np_x0 = test_common.generate_numpy(shape, dtype) + x0 = torch.from_numpy(np_x0).to(eval('torch.' + dtype)).npu() + np_x1 = test_common.generate_numpy(shape, dtype) + x1 = torch.from_numpy(np_x1).to(eval('torch.' + dtype)).npu() + # torch结果 + torch_res = torch_ne_from_np(np_x0, np_x1, dtype).npu() + # triton结果 + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_ne[ncore, 1, 1](x0, x1, triton_res, xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_nearbyint.py b/third_party/ascend/unittest/pytest_ut/test_nearbyint.py new file mode 100644 index 0000000000..f02b896d7c --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_nearbyint.py @@ -0,0 +1,69 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import torch +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + + +@triton.jit +def nearbyint_kernel(x_ptr, y_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + + mask = offsets < n_elements + + x = tl.load(x_ptr + offsets, mask=mask) + + y = libdevice.nearbyint(x) + + tl.store(y_ptr + offsets, y, mask=mask) + + +@pytest.mark.parametrize('shape', [ + (12, 16), +]) +@pytest.mark.parametrize('dtype', ['float32']) +def test_nearbyint(shape, dtype): + n_elements = shape[0] * shape[1] + x = test_common.generate_tensor(shape, dtype).npu() + + # Ensure some boundary cases are included + x[0, 0] = 0.0 + x[0, 1] = 3.14 # Should be rounded to 3.0 + x[0, 2] = -2.71 # Should be rounded to -3.0 + x[0, 3] = 5.0 # Integer, should remain unchanged + x[0, 4] = -3.0 # Negative integer, should remain unchanged + x[0, 5] = 2.5 # Should be rounded to 3.0 + x[0, 6] = -1.5 # Should be rounded to -2.0 + + y = torch.empty_like(x) + + BLOCK_SIZE = 192 + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + + nearbyint_kernel[grid](x, y, n_elements, BLOCK_SIZE=BLOCK_SIZE) + + expected = torch.round(x) + test_common.validate_cmp('float32', y, expected) diff --git a/third_party/ascend/unittest/pytest_ut/test_nearest.py b/third_party/ascend/unittest/pytest_ut/test_nearest.py new file mode 100644 index 0000000000..f6df8d3ded --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_nearest.py @@ -0,0 +1,116 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import math +import numpy as np +import pytest + + +@triton.jit +def nearest_resize_kernel(img_src_ptr, img_dst_ptr, src_rows, src_cols, dst_rows, dst_cols, RR_H, RR_W, C, stride_in_h, + stride_in_w, stride_in_c, stride_out_h, stride_out_w, stride_out_c, BLOCK_SIZE: tl.constexpr): + #RR_H和RR_W分别为高和宽的缩放比例 + block_id_c = tl.program_id(0) + block_id_h = tl.program_id(1) + block_id_w = tl.program_id(2) + dest_h_offs = (block_id_h * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)) + dest_w_offs = (block_id_w * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)) + dest_offs = (block_id_c[None, None] * stride_out_c + dest_h_offs[:, None] * stride_out_h + + dest_w_offs[None, :] * stride_out_w) + #根据output image的坐标值(dest_h_offs, dest_w_offs)计算input image的坐标值(sy, sx) + fy = dest_h_offs * RR_H + sy = tl.floor(fy) + fx = dest_w_offs * RR_W + sx = tl.floor(fx) + + src_offsets = (block_id_c[None, None] * stride_in_c + + tl.clamp(sy, 0, src_rows - 1)[:, None].to(tl.int32) * stride_in_h + + tl.clamp(sx, 0, src_cols - 1)[None, :].to(tl.int32) * stride_in_w) + src_val = tl.load(img_src_ptr + src_offsets) + dst_mask = (dest_h_offs[:, None] < dst_rows) & (dest_w_offs[None, :] < dst_cols) + tl.store(img_dst_ptr + dest_offs, src_val, mask=dst_mask) + + +def triton_kernel(img_src, img_dst): + N, C, src_rows, src_cols = img_src.shape + _, _, dst_rows, dst_cols = img_dst.shape + R_H = float(dst_rows) / src_rows + R_W = float(dst_cols) / src_cols + RR_H = 1.0 / R_H + RR_W = 1.0 / R_W + stride_in_n, stride_in_c, stride_in_h, stride_in_w = img_src.stride() + stride_out_n, stride_out_c, stride_out_h, stride_out_w = img_dst.stride() + bs = 16 + grid = lambda meta: ( + C, + triton.cdiv(dst_rows, meta["BLOCK_SIZE"]), + triton.cdiv(dst_cols, meta["BLOCK_SIZE"]), + ) + nearest_resize_kernel[grid](img_src, img_dst, src_rows, src_cols, dst_rows, dst_cols, RR_H, RR_W, C, stride_in_h, + stride_in_w, stride_in_c, stride_out_h, stride_out_w, stride_out_c, bs) + return img_dst + + +def nearest_resize_cpu(img_src, img_dst): + N, C, src_rows, src_cols = img_src.shape + _, _, dst_rows, dst_cols = img_dst.shape + #RR_H和RR_W分别为高和宽的缩放比例 + RR_H = src_rows / float(dst_rows) + RR_W = src_cols / float(dst_cols) + #根据output image的坐标值(i,j)计算input image的坐标值(sy, sx) + for i in range(dst_rows): + for j in range(dst_cols): + fy = (i * RR_H) + sy = math.floor(fy) + fx = (j * RR_W) + sx = math.floor(fx) + src_val = img_src[0, :, np.clip(sy, 0, src_rows - 1), np.clip(sx, 0, src_cols - 1)] + img_dst[0, :, i, j] = src_val + return img_dst + + +@pytest.mark.parametrize("shapes", [ + [360, 640, 140, 280], +]) +def test_nearest(shapes): + src_rows, src_cols, dst_rows, dst_cols = shapes + img_src = torch.rand(1, 4, src_rows, src_cols, dtype=torch.float32, device='npu') + img_dst = torch.zeros((1, img_src.shape[1], dst_rows, dst_cols), dtype=img_src.dtype, device=img_src.device) + torch_ref = nearest_resize_cpu(img_src.cpu(), img_dst.cpu()) + triton_cal = triton_kernel(img_src, img_dst) + torch.testing.assert_close(torch_ref.npu(), triton_cal) + + +if __name__ == "__main__": + src_rows, src_cols = 360, 640 + dst_rows, dst_cols = 140, 280 + img_src = torch.rand(1, 4, src_rows, src_cols, dtype=torch.float32, device='npu') + img_dst = torch.zeros((1, img_src.shape[1], dst_rows, dst_cols), dtype=img_src.dtype, device=img_src.device) + + assert img_src.shape[ + 0] == 1, "currently supports only shape[0] == 1 which does not change the functionality of thie case" + torch_ref = nearest_resize_cpu(img_src.cpu(), img_dst.cpu()) + triton_cal = triton_kernel(img_src, img_dst) + torch.testing.assert_close(torch_ref.npu(), triton_cal) + print("success") diff --git a/third_party/ascend/unittest/pytest_ut/test_neg.py b/third_party/ascend/unittest/pytest_ut/test_neg.py new file mode 100644 index 0000000000..a32a5039d8 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_neg.py @@ -0,0 +1,58 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import time +import torch +import torch_npu +import test_common + + +def torch_neg(x0): + res = -x0 + return res + + +@triton.jit +def triton_neg(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1 = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = -tmp0 + tl.store(out_ptr0 + (x0), tmp1, None) + + +@pytest.mark.parametrize('param_list', [ + ['float16', (8, 8), 8, 8, 8], + ['float32', (8, 8), 8, 8, 8], + ['int8', (2, 4096, 8), 32, 2048, 64], +]) +def test_neg(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_neg(x0) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_neg[ncore, 1, 1](x0, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_negative_mask_dim.py b/third_party/ascend/unittest/pytest_ut/test_negative_mask_dim.py new file mode 100644 index 0000000000..55ea0cb752 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_negative_mask_dim.py @@ -0,0 +1,48 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl + +import torch +import torch_npu + + +@triton.jit +def triton_negative_mask_dim(in_ptr0, out_ptr0, XBLOCK: tl.constexpr): + index = tl.arange(0, XBLOCK) + mask = (index < 1) & (index + 1 >= XBLOCK) + tmp0 = tl.load(in_ptr0 + index, mask, other=0.0) + tl.store(out_ptr0 + index, tmp0, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (32, ), 32], +]) +def test_negative_mask_dim(param_list): + dtype, shape, xblock = param_list + x0 = torch.ones(shape, dtype=eval('torch.' + dtype)).npu() + y_ref = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + + y_cal = torch.ones(shape, dtype=eval('torch.' + dtype)).npu() + triton_negative_mask_dim[(1, )](x0, y_cal, xblock) + assert torch.allclose(y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_nextafter.py b/third_party/ascend/unittest/pytest_ut/test_nextafter.py new file mode 100644 index 0000000000..007f74d9dc --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_nextafter.py @@ -0,0 +1,63 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + +import torch +import torch_npu +import numpy as np + + +def torch_nextafter(x0, x1): + return torch.nextafter(x0, x1) + + +@triton.jit +def triton_nextafter(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.load(in_ptr1 + (x0), None) + tmp2 = libdevice.nextafter(tmp0, tmp1) + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], +]) +def test_nextafter(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0_ref = test_common.generate_tensor(shape, dtype) + x1_ref = test_common.generate_tensor(shape, dtype) + x0 = x0_ref.npu() + x1 = x1_ref.npu() + y_ref = torch_nextafter(x0_ref, x1_ref) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_nextafter[ncore, 1, 1](x0, x1, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_not.py b/third_party/ascend/unittest/pytest_ut/test_not.py new file mode 100644 index 0000000000..51a4871993 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_not.py @@ -0,0 +1,63 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import torch +import torch_npu +import test_common + + +def torch_not(x0): + res = torch.bitwise_not(x0) + return res + + +@triton.jit +def triton_not(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + xindex = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = xindex < xnumel + x0 = xindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp2 = not (tmp0) + tl.store(out_ptr0 + (xindex), tmp2, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['int8', (2, 4096, 8), 2, 32768, 1024], + ['int16', (2, 4096, 8), 2, 32768, 1024], + ['int32', (2, 4096, 8), 2, 32768, 1024], + ['int64', (2, 4096, 8), 2, 32768, 1024], + ['bool', (2, 4096, 8), 2, 32768, 1024], +]) +def test_not(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + torch_res = torch_not(x0) + # triton结果 + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_not[ncore, 1, 1](x0, triton_res, x0.numel(), xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_npu_indexing.py b/third_party/ascend/unittest/pytest_ut/test_npu_indexing.py new file mode 100644 index 0000000000..d1abbc3aab --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_npu_indexing.py @@ -0,0 +1,143 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import time + + +def foo(a, b, c): + Z, Y, X, R = (1, 1, 64, 64) + y = a + b + y = y.sum(-1) + y = y.unsqueeze(3) + y = y.broadcast_to(Z, Y, X, R) + b + y = c + y.permute(0, 1, 3, 2) + return y + + +@triton.jit +def triton_foo( + in_ptr0, + in_ptr1, + in_ptr2, + out_ptr0, + BLOCK1: tl.constexpr, + BLOCK1_SUB: tl.constexpr, + BLOCK2: tl.constexpr, + Z: tl.constexpr, + Y: tl.constexpr, + X: tl.constexpr, + R: tl.constexpr, + Z_STRIDE: tl.constexpr, + Y_STRIDE: tl.constexpr, + X_STRIDE: tl.constexpr, + R_STRIDE: tl.constexpr, + Z_STRIDE1: tl.constexpr, + Y_STRIDE1: tl.constexpr, + X_STRIDE1: tl.constexpr, + R_STRIDE1: tl.constexpr, +): + offset: tl.constexpr = tl.program_id(0) * BLOCK1 + base1 = tl.arange(0, BLOCK1_SUB) + base2 = tl.arange(0, BLOCK2) + nsub: tl.constexpr = BLOCK1 // BLOCK1_SUB + # loops1 : tl.constexpr = nsub * Y * Z + loops1: tl.constexpr = nsub + loops2: tl.constexpr = R // BLOCK2 + + for z in range(Z): + for y in range(Y): + for loop1 in range(loops1): + # y = (loop1 // nsub) % Y + # z = loop1 // nsub // Y + # off1 = (loop1 % nsub) + off1 = loop1 + x = offset + (off1 * BLOCK1_SUB) + base1[:, None] + x1 = offset + (off1 * BLOCK1_SUB) + base1[None, :] + _tmp4 = tl.full([BLOCK1_SUB, BLOCK2], 0, tl.float32) + for loop2 in range(loops2): + r = loop2 * BLOCK2 + base2[None, :] + tmp0 = tl.load(in_ptr0 + (R_STRIDE * r + (X_STRIDE * x) + (Y_STRIDE * y) + (Z_STRIDE * z)), None) + tmp1 = tl.load(in_ptr1 + (R_STRIDE * r + (X_STRIDE * x) + (Y_STRIDE * y) + (Z_STRIDE * z)), None) + tmp2 = tmp0 + tmp1 + _tmp4 = _tmp4 + tmp2 + tmp4 = tl.sum(_tmp4, 1)[:, None] + tmp5 = tmp4.reshape(BLOCK1_SUB, 1).broadcast_to(BLOCK1_SUB, BLOCK2) + + for loop2 in range(loops2): + r = loop2 * BLOCK2 + base2[None, :] + tmp6 = tl.load(in_ptr1 + (R_STRIDE * r + (X_STRIDE * x) + (Y_STRIDE * y) + (Z_STRIDE * z)), None) + tmp7 = tmp6 + tmp5 + r1 = loop2 * BLOCK2 + base2[:, None] + tmp8 = tl.load(in_ptr2 + (R_STRIDE1 * x1 + (X_STRIDE1 * r1) + (Y_STRIDE1 * y) + (Z_STRIDE1 * z)), + None) + + tmp9 = tmp8.reshape(BLOCK2, BLOCK1_SUB) + tmp7.reshape(BLOCK1_SUB, BLOCK2).permute(1, 0) + tl.store(out_ptr0 + (R_STRIDE1 * x1 + (X_STRIDE1 * r1) + (Y_STRIDE1 * y) + (Z_STRIDE1 * z)), tmp9, + None) + + +def foo_triton_wrapper(a, b, c): + NBLOCKS = 1 + BLOCK1 = a.shape[2] // NBLOCKS + BLOCK1_SUB = 64 + BLOCK2 = 64 + + value = torch.empty_strided((c.shape[0], c.shape[1], c.shape[2], c.shape[3]), + (c.stride()[0], c.stride()[1], c.stride()[2], c.stride()[3]), + dtype=torch.float32).npu() + + triton_foo[NBLOCKS, 1, 1]( + a, + b, + c, + value, + BLOCK1, + BLOCK1_SUB, + BLOCK2, + a.shape[0], + a.shape[1], + a.shape[2], + a.shape[3], + a.stride()[0], + a.stride()[1], + a.stride()[2], + a.stride()[3], + c.stride()[0], + c.stride()[1], + c.stride()[2], + c.stride()[3], + ) + return value + + +def test_npu_indexing(): + Z, Y, X, R = (1, 1, 64, 64) + a = torch.randn((Z, Y, X, R), dtype=torch.float32).npu() + b = torch.randn((Z, Y, X, R), dtype=torch.float32).npu() + c = torch.randn((Z, Y, R, X), dtype=torch.float32).npu() + r = foo_triton_wrapper(a, b, c) + r1 = foo(a, b, c) + print(r[0, 0, 0:8, 0:8]) + print(r1[0, 0, 0:8, 0:8]) + torch.testing.assert_close(r, r1) diff --git a/third_party/ascend/unittest/pytest_ut/test_npu_indexing2.py b/third_party/ascend/unittest/pytest_ut/test_npu_indexing2.py new file mode 100644 index 0000000000..e36cb8403f --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_npu_indexing2.py @@ -0,0 +1,92 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import time + + +def foo(a, b, c): + y = a + b + c + y = y.sum(dim=1) + return y + + +@triton.jit +def triton_codegen2(in_ptr0, in_ptr1, in_ptr2, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr, + RBLOCK: tl.constexpr): + ynumel = 8 + rnumel = 2048 + xnumel = 1024 + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + base2 = tl.arange(0, RBLOCK) + loops2: tl.constexpr = rnumel // RBLOCK + for y in range(ynumel): + y0 = y + for loop1 in range(loops1): + x = offset + (loop1 * XBLOCK_SUB) + base1 + x1 = offset + (loop1 * XBLOCK_SUB) + base1[None, :] + _tmp6 = tl.full([XBLOCK_SUB, RBLOCK], 0, tl.float32) + for loop2 in range(loops2): + r2 = loop2 * RBLOCK + base2[:, None] + tmp0 = tl.load(in_ptr0 + (x1 + (1024 * r2) + (2097152 * y0)), None, eviction_policy='evict_last') + tmp1 = tl.load(in_ptr1 + (x1 + (1024 * r2) + (2097152 * y0)), None, eviction_policy='evict_last') + tmp3 = tl.load(in_ptr2 + (x1 + (1024 * r2) + (2097152 * y0)), None, eviction_policy='evict_last') + tmp2 = tmp0 + tmp1 + tmp4 = tmp2 + tmp3 + tmp5 = tl.reshape(tmp4, [RBLOCK, XBLOCK_SUB]) + tmp7 = _tmp6 + tmp5 + _tmp6 = tmp7 + tmp6 = tl.sum(_tmp6, 0).reshape(XBLOCK_SUB) + + tl.store(out_ptr0 + (x + (1024 * y0)), tmp6, None) + + +def foo_triton_wrapper(a, b, c): + NBLOCKS = 8 + BLOCK1 = a.shape[2] // NBLOCKS + BLOCK1_SUB = 64 + BLOCK2 = 64 + + value = torch.empty_strided((c.shape[0], c.shape[2]), (c.shape[2], 1), dtype=torch.float32).npu() + + triton_codegen2[NBLOCKS, 1, 1](a, b, c, value, BLOCK1, BLOCK1_SUB, BLOCK2) + + return value + + +def test_npu_indexing2(): + + Y, X, R = (8, 2048, 1024) + a = torch.randn((Y, X, R), dtype=torch.float32).npu() + b = torch.randn((Y, X, R), dtype=torch.float32).npu() + c = torch.randn((Y, X, R), dtype=torch.float32).npu() + r = foo_triton_wrapper(a, b, c) + r1 = foo(a, b, c) + print(r[ + 0:8, + 0:8, + ]) + print(r1[0:8, 0:8]) + torch.testing.assert_close(r, r1, rtol=1e-3, atol=1e-3) diff --git a/third_party/ascend/unittest/pytest_ut/test_or.py b/third_party/ascend/unittest/pytest_ut/test_or.py new file mode 100644 index 0000000000..5a03d301c5 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_or.py @@ -0,0 +1,60 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common + + +def torch_or(x0, x1): + res = x0 | x1 + return res + + +@triton.jit +def triton_or(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x_index = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + x_index, None) + tmp1 = tl.load(in_ptr1 + x_index, None) + tmp2 = tmp0 | tmp1 + tl.store(out_ptr0 + x_index, tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['int32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_or(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + torch_res = torch_or(x0, x1) + # triton结果 + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_or[ncore, 1, 1](x0, x1, triton_res, xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_paged_kvcache_krope.py b/third_party/ascend/unittest/pytest_ut/test_paged_kvcache_krope.py new file mode 100644 index 0000000000..ebc419b4fa --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_paged_kvcache_krope.py @@ -0,0 +1,72 @@ +import torch +import triton +import triton.language as tl +import pytest + + +@triton.jit +def rope_like_load_kernel( + Kv_cache, + Req_to_tokens, + output_ptr, + stride_kv: tl.constexpr, + HEAD_DIM_V: tl.constexpr, + HEAD_DIM: tl.constexpr, + PAGE_SIZE: tl.constexpr, + BLOCK_N: tl.constexpr, + ROPE_DIM: tl.constexpr, +): + + offs_d_kpe = tl.arange(HEAD_DIM_V, HEAD_DIM) + offs_n = tl.arange(0, BLOCK_N) + + kv_page_number = tl.load(Req_to_tokens + offs_n // PAGE_SIZE) + kv_loc = kv_page_number * PAGE_SIZE + offs_n % PAGE_SIZE + + offs_k_pe = kv_loc[None, :] * stride_kv + offs_d_kpe[:, None] + + k_pe = tl.load(Kv_cache + offs_k_pe) + + offs_out = offs_n[:, None] * ROPE_DIM + tl.arange(0, ROPE_DIM)[None, :] + tl.store(output_ptr + offs_out, tl.trans(k_pe)) + + +def test_bubbleup_extract_nonzero_offset(): + device = "npu" + + PAGE_SIZE = 2 + BLOCK_N = 4 + head_dim = 32 + head_dim_v = 24 + rope_dim = head_dim - head_dim_v + num_pages = BLOCK_N // PAGE_SIZE + + req_to_tokens = torch.arange(num_pages, dtype=torch.int32, device=device) + total_tokens = num_pages * PAGE_SIZE + kv_cache = torch.zeros(total_tokens, head_dim, dtype=torch.float32, device=device) + for token_id in range(total_tokens): + kv_cache[token_id, :head_dim_v] = (torch.arange(head_dim_v, dtype=torch.float32) + token_id * 100) + kv_cache[token_id, head_dim_v:] = (torch.arange(head_dim_v, head_dim, dtype=torch.float32) + token_id * 1000) + output = torch.zeros(BLOCK_N, rope_dim, dtype=torch.float32, device=device) + + rope_like_load_kernel[(1, )]( + kv_cache.flatten(), + req_to_tokens, + output.flatten(), + stride_kv=head_dim, + HEAD_DIM_V=head_dim_v, + HEAD_DIM=head_dim, + PAGE_SIZE=PAGE_SIZE, + BLOCK_N=BLOCK_N, + ROPE_DIM=rope_dim, + ) + + expected = torch.zeros(BLOCK_N, rope_dim, dtype=torch.float32, device=device) + for token_id in range(BLOCK_N): + expected[token_id] = (torch.arange(head_dim_v, head_dim, dtype=torch.float32) + token_id * 1000) + + buggy = torch.zeros(BLOCK_N, rope_dim, dtype=torch.float32, device=device) + for token_id in range(BLOCK_N): + buggy[token_id] = (torch.arange(rope_dim, dtype=torch.float32) + token_id * 100) + + assert torch.allclose(output, expected, atol=1e-5) diff --git a/third_party/ascend/unittest/pytest_ut/test_parallel.py b/third_party/ascend/unittest/pytest_ut/test_parallel.py new file mode 100644 index 0000000000..f81c060811 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_parallel.py @@ -0,0 +1,100 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl + +import torch +import torch_npu + +import test_common +import triton.language.extra.cann.extension as extension + + +@triton.jit +def triton_add(in_ptr0, in_ptr1, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0 + x1 + + for _ in extension.parallel(2, 5, 2, bind_sub_block=False): + ret = ret + x1 + + for _ in extension.parallel(2, 10, 3, bind_sub_block=False): + ret = ret + x0 + + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + tl.store(out_ptr0 + odx, ret) + + +testlist = [ + (3, 5, 8), +] + + +def get_torch_typename(dtype): + if dtype == 'float32': + tyname = torch.float32 + elif dtype == 'int32': + tyname = torch.int32 + elif dtype == 'int64': + tyname = torch.int64 + elif dtype == 'float16': + tyname = torch.float16 + elif dtype == 'bfloat16': + tyname = torch.bfloat16 + elif dtype == 'int16': + tyname = torch.int16 + elif dtype == 'int8': + tyname = torch.int8 + elif dtype == 'bool': + tyname = torch.bool + else: + raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + + return tyname + + +typelist = ['int8', 'int16', 'int32', 'int64'] + + +@pytest.mark.skip(reason="not supported after the NPUIR is updated in April, and will be fixed later") +@pytest.mark.parametrize('L, M, N', testlist) +@pytest.mark.parametrize('sigtype', typelist) +def test_add_bind_false(sigtype, L, M, N): + dtype = get_torch_typename(sigtype) + shape = (L, M, N) + x0 = test_common.generate_tensor(shape=(L, M, N), dtype=sigtype).npu() + x1 = test_common.generate_tensor(shape=(L, M, N), dtype=sigtype).npu() + y_ref = x0 + x1 + x1 + x1 + x0 + x0 + x0 + + output = torch.zeros(shape, dtype=dtype).npu() + h = triton_add[1, 1, 1](x0, x1, output, L, M, N) + + test_common.validate_cmp(sigtype, output, y_ref) + code_str = h.asm["ttadapter"] + count = code_str.count("hivm.parallel_loop") + assert count == 2 diff --git a/third_party/ascend/unittest/pytest_ut/test_permute.py b/third_party/ascend/unittest/pytest_ut/test_permute.py new file mode 100644 index 0000000000..9850b26579 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_permute.py @@ -0,0 +1,127 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import time + + +@triton.jit +def triton_foo( + in_ptr0, + in_ptr1, + in_ptr2, + out_ptr0, + BLOCK1: tl.constexpr, + BLOCK1_SUB: tl.constexpr, + BLOCK2: tl.constexpr, + X: tl.constexpr, + Y: tl.constexpr, + Z: tl.constexpr, + R: tl.constexpr, + Z_STRIDE: tl.constexpr, + Y_STRIDE: tl.constexpr, + X_STRIDE: tl.constexpr, + R_STRIDE: tl.constexpr, + X_STRIDE1: tl.constexpr, + Y_STRIDE1: tl.constexpr, + Z_STRIDE1: tl.constexpr, + R_STRIDE1: tl.constexpr, +): + offset: tl.constexpr = tl.program_id(0) * BLOCK1 + base1 = tl.arange(0, BLOCK1_SUB) + base2 = tl.arange(0, BLOCK2) + nsub: tl.constexpr = BLOCK1 // BLOCK1_SUB + # loops1 : tl.constexpr = nsub * Y * Z + loops1: tl.constexpr = nsub + loops2: tl.constexpr = R // BLOCK2 + + for z in range(Z): + for y in range(Y): + for loop1 in range(loops1): + off1 = loop1 + x = offset + (off1 * BLOCK1_SUB) + base1[:, None] + x1 = offset + (off1 * BLOCK1_SUB) + base1[None, :] + + for loop2 in range(loops2): + r = loop2 * BLOCK2 + base2[None, :] + r1 = loop2 * BLOCK2 + base2[:, None] + tmp0 = tl.load(in_ptr0 + ((R_STRIDE * r) + (X_STRIDE * x) + (Y_STRIDE * y) + (Z_STRIDE * z)), None) + tmp1 = tl.load(in_ptr1 + ((R_STRIDE * r) + (X_STRIDE * x) + (Y_STRIDE * y) + (Z_STRIDE * z)), None) + tmp2 = tmp0 + tmp1 + + tmp8 = tl.load(in_ptr2 + (R_STRIDE1 * r + X_STRIDE1 * x + (Y_STRIDE1 * y) + (Z_STRIDE1 * z)), None) + tmp9 = tmp8 + tmp2 + tl.store(out_ptr0 + (R_STRIDE1 * r + X_STRIDE1 * x + (Y_STRIDE1 * y) + (Z_STRIDE1 * z)), tmp9, None) + + +def foo_triton_wrapper(a, b, c): + NBLOCKS = 32 if c.shape[0] >= 256 else 1 + BLOCK1 = c.shape[0] // NBLOCKS + BLOCK1_SUB = BLOCK1 if BLOCK1 < 64 else 64 + BLOCK2 = c.shape[3] if c.shape[3] < 64 else 64 + + value = torch.empty_strided((c.shape[0], c.shape[1], c.shape[2], c.shape[3]), + (c.stride()[0], c.stride()[1], c.stride()[2], c.stride()[3]), + dtype=torch.float32).npu() + + triton_foo[NBLOCKS, 1, 1]( + a, + b, + c, + value, + BLOCK1, + BLOCK1_SUB, + BLOCK2, + c.shape[0], + c.shape[1], + c.shape[2], + c.shape[3], + a.stride()[0], + a.stride()[1], + a.stride()[2], + a.stride()[3], + c.stride()[0], + c.stride()[1], + c.stride()[2], + c.stride()[3], + ) + return value + + +def foo(a, b, c): + y = a + b + y = c + y.permute(2, 1, 0, 3) + return y + + +def test_permute_handwritten(): + + Z, Y, X, R = (1, 12, 4096, 8) + a = torch.randn((Z, Y, X, R), dtype=torch.float32).npu() + b = torch.randn((Z, Y, X, R), dtype=torch.float32).npu() + c = torch.randn((X, Y, Z, R), dtype=torch.float32).npu() + r = foo_triton_wrapper(a, b, c) + r1 = foo(a, b, c) + print(r[0, 0, 0:8, 0:8]) + print(r1[0, 0, 0:8, 0:8]) + torch.testing.assert_close(r, r1, rtol=1e-3, atol=1e-3) diff --git a/third_party/ascend/unittest/pytest_ut/test_permute_full.py b/third_party/ascend/unittest/pytest_ut/test_permute_full.py new file mode 100644 index 0000000000..e717b8a1c3 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_permute_full.py @@ -0,0 +1,171 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common + + +@triton.jit +def fn_npu_021(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + # XB,YB,1 + X = tl.load(x_ptr + idx) + + ret = tl.permute(X, (0, 2, 1)) + + oidx = xidx[:, None, None] * YB * ZB + zidx[None, :, None] * YB + yidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_102(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + # XB,YB,1 + X = tl.load(x_ptr + idx) + + ret = tl.permute(X, (1, 0, 2)) + + oidx = yidx[:, None, None] * XB * ZB + xidx[None, :, None] * ZB + zidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_210(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + # XB,YB,1 + X = tl.load(x_ptr + idx) + + ret = tl.permute(X, (2, 1, 0)) + + oidx = zidx[:, None, None] * YB * XB + yidx[None, :, None] * XB + xidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_201(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + # XB,YB,1 + X = tl.load(x_ptr + idx) + + ret = tl.permute(X, (2, 0, 1)) + + oidx = zidx[:, None, None] * YB * XB + xidx[None, :, None] * YB + yidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_120(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + # XB,YB,1 + X = tl.load(x_ptr + idx) + + ret = tl.permute(X, (1, 2, 0)) + + oidx = yidx[:, None, None] * ZB * XB + zidx[None, :, None] * XB + xidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@pytest.mark.parametrize('para_type,data_type,XB,YB,ZB', [ + # ['float32',eval('torch.float32'),2,4,3], + ['float32', eval('torch.float32'), 2, 4, 8], + # ['float32',eval('torch.float32'),2,4,37], + ['float32', eval('torch.float32'), 2, 4, 64], + # ['float32',eval('torch.float32'),2,4,781], + + # ['float16',eval('torch.float16'),2,4,3], + ['float16', eval('torch.float16'), 2, 4, 8], + # ['float16',eval('torch.float16'),2,4,37], + ['float16', eval('torch.float16'), 2, 4, 64], + # ['float16',eval('torch.float16'),2,4,781], + + # ['int8',eval('torch.int8'),2,4,3], + ['int8', eval('torch.int8'), 2, 4, 8], + # ['int8',eval('torch.int8'),2,4,37], + ['int8', eval('torch.int8'), 2, 4, 64], + # ['int8',eval('torch.int8'),2,4,781], +]) +def test_permute(para_type, data_type, XB, YB, ZB): + + x = torch.randint(low=0, high=2, size=(XB, YB, ZB), dtype=data_type).npu() + + output = torch.randint(1, (XB, ZB, YB), dtype=data_type).npu() + torch_021 = torch.permute(x, (0, 2, 1)) + fn_npu_021[1, 1, 1](output, x, XB, YB, ZB) + torch.testing.assert_close(output, torch_021) + + print(" test permute 021 passed") + + output = torch.randint(1, (YB, XB, ZB), dtype=data_type).npu() + torch_102 = torch.permute(x, (1, 0, 2)) + fn_npu_102[1, 1, 1](output, x, XB, YB, ZB) + torch.testing.assert_close(output, torch_102) + + print(" test permute 102 passed") + + output = torch.randint(1, (ZB, XB, YB), dtype=data_type).npu() + torch_201 = torch.permute(x, (2, 0, 1)) + fn_npu_201[1, 1, 1](output, x, XB, YB, ZB) + torch.testing.assert_close(output, torch_201) + + print(" test permute 201 passed") + + output = torch.randint(1, (ZB, YB, XB), dtype=data_type).npu() + torch_210 = torch.permute(x, (2, 1, 0)) + fn_npu_210[1, 1, 1](output, x, XB, YB, ZB) + torch.testing.assert_close(output, torch_210) + + print(" test permute 210 passed") + + output = torch.randint(1, (YB, ZB, XB), dtype=data_type).npu() + torch_120 = torch.permute(x, (1, 2, 0)) + fn_npu_120[1, 1, 1](output, x, XB, YB, ZB) + torch.testing.assert_close(output, torch_120) + + print(" test permute 120 passed") diff --git a/third_party/ascend/unittest/pytest_ut/test_permute_reshape.py b/third_party/ascend/unittest/pytest_ut/test_permute_reshape.py new file mode 100644 index 0000000000..65b5775b31 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_permute_reshape.py @@ -0,0 +1,91 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import time + + +@triton.jit +def triton_foo( + in_ptr0, + in_ptr1, + in_ptr2, + out_ptr0, + BLOCK1: tl.constexpr, + BLOCK1_SUB: tl.constexpr, + BLOCK2: tl.constexpr, + S: tl.constexpr, + N: tl.constexpr, + D: tl.constexpr, +): + offset: tl.constexpr = tl.program_id(0) * BLOCK1 + base1 = tl.arange(0, BLOCK1_SUB) + base2 = tl.arange(0, BLOCK2) + loops1: tl.constexpr = BLOCK1 // BLOCK1_SUB + loops2: tl.constexpr = D // BLOCK2 + + for loop1 in range(loops1): + off1 = loop1 + s = offset + (off1 * BLOCK1_SUB) + base1[:, None] + for n in range(N): + for loop2 in range(loops2): + d = loop2 * BLOCK2 + base2[None, :] + tmp0 = tl.load(in_ptr0 + ((32768 * n) + (8 * s) + d), None) + tmp1 = tl.load(in_ptr1 + ((32768 * n) + (8 * s) + d), None) + tmp2 = tmp0 + tmp1 + + tmp3 = tl.load(in_ptr2 + ((8 * n) + d + (96 * s)), None) + tmp9 = tmp3 + tmp2 + tl.store(out_ptr0 + ((8 * n) + d + (96 * s)), tmp9, None) + + +def foo_triton_wrapper(a, b, c): + NBLOCKS = 32 if a.shape[2] >= 256 else 1 + BLOCK1 = a.shape[2] // NBLOCKS + BLOCK1_SUB = BLOCK1 if BLOCK1 < 64 else 64 + BLOCK2 = a.shape[3] if a.shape[3] < 64 else 64 + + value = torch.empty_strided((c.shape[0], c.shape[1], c.shape[2]), (c.stride()[0], c.stride()[1], c.stride()[2]), + dtype=torch.float32).npu() + triton_foo[NBLOCKS, 1, 1](a, b, c, value, BLOCK1, BLOCK1_SUB, BLOCK2, a.shape[2], a.shape[1], a.shape[3]) + + return value + + +def foo(a, b, c): + B, N, S, D = (1, 12, 4096, 8) + y = a + b + y = c + y.permute(2, 0, 1, 3).reshape(S, B, N * D) + return y + + +def test_permute_reshape(): + B, N, S, D = (1, 12, 4096, 8) + a = torch.randn((B, N, S, D), dtype=torch.float32).npu() + b = torch.randn((B, N, S, D), dtype=torch.float32).npu() + c = torch.randn((S, B, N * D), dtype=torch.float32).npu() + r = foo_triton_wrapper(a, b, c) + r1 = foo(a, b, c) + print(r[0:8, 0, 0:8]) + print(r1[0:8, 0, 0:8]) + torch.testing.assert_close(r, r1, rtol=1e-3, atol=1e-3) diff --git a/third_party/ascend/unittest/pytest_ut/test_permuted_boundary_check.py b/third_party/ascend/unittest/pytest_ut/test_permuted_boundary_check.py new file mode 100644 index 0000000000..c232f62de0 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_permuted_boundary_check.py @@ -0,0 +1,38 @@ +import torch +import torch_npu +import triton +import triton.language as tl +import pytest + + +@triton.jit +def zj_fa_fwd_pattern(in_ptr0, in_ptr1, out_ptr, M, K, N, MBLOCK: tl.constexpr, NBLOCK: tl.constexpr, + KBLOCK: tl.constexpr): + a_ptr = tl.make_block_ptr(base=in_ptr0, shape=(M, K), # 8, 3 + strides=(K, 1), offsets=(0, 0), block_shape=(MBLOCK, KBLOCK), order=(1, 0)) + + b_ptr = tl.make_block_ptr(base=in_ptr1, shape=(K, N), # 3, 8 + strides=(1, K), offsets=(0, 0), block_shape=(KBLOCK, NBLOCK), order=(0, 1)) + + c_ptr = tl.make_block_ptr(base=out_ptr, shape=(M, N), strides=(1, M), offsets=(0, 0), block_shape=(MBLOCK, NBLOCK), + order=(0, 1)) + + a = tl.load(a_ptr, boundary_check=(0, ), padding_option="zero") + b = tl.load(b_ptr, boundary_check=(0, ), padding_option="zero") + c = tl.dot(a, b) + tl.store(c_ptr, c, boundary_check=(0, 1)) + + +def test_permute_boundary_check(): + M = 8 + K = 3 + N = 8 + MBLOCK = 8 + NBLOCK = 8 + KBLOCK = 4 + a = torch.randn((M, K), device="npu") # 8, 3 + b = torch.randn((N, K), device="npu") # 8, 3 + c = torch.empty((N, M), device="npu") + zj_fa_fwd_pattern[(1, 1, 1)](a, b, c, M, K, N, MBLOCK, NBLOCK, KBLOCK) + std = a @ b.T + torch.testing.assert_close(std, c.T, atol=1e-2, rtol=1e-2) diff --git a/third_party/ascend/unittest/pytest_ut/test_pointer_type.py b/third_party/ascend/unittest/pytest_ut/test_pointer_type.py new file mode 100644 index 0000000000..7641b25d66 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_pointer_type.py @@ -0,0 +1,49 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common + + +@triton.jit +def kernel(ans_ptr, x_ptr): + val = tl.load(x_ptr) + output_ptr = tl.load(ans_ptr) + output_ptr = output_ptr.to(tl.pointer_type(val.dtype)) + tl.store(output_ptr, val) + + +@pytest.mark.parametrize("literal, dtype_str", + [[0, eval('torch.int8')], [0, eval('torch.int16')], [0, eval('torch.int32')], + [0, eval('torch.int64')], [0, eval('torch.float16')], [0, eval('torch.float32')]]) +def test_pointer_type(literal, dtype_str): + x = torch.randint(low=0, high=5, size=(1, ), dtype=dtype_str).npu() + output = torch.zeros((1, ), dtype=dtype_str).npu() + ans = [] + ans.append(output.data_ptr()) + ans_tensor = torch.tensor(ans).npu() + kernel[(1, )](ans_tensor, x) + assert torch.isclose(x, output) + print("Pointer type convert successful") diff --git a/third_party/ascend/unittest/pytest_ut/test_pow.py b/third_party/ascend/unittest/pytest_ut/test_pow.py new file mode 100644 index 0000000000..310aa9afa7 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_pow.py @@ -0,0 +1,121 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +from triton.language.extra.cann.libdevice import pow +import torch +import torch_npu +import pytest +import test_common + +shapes = [ + # 3, + # 32, + 37, + # 256, + # 781, +] + + +@pytest.mark.parametrize("sigtype", ["float32"]) +@pytest.mark.parametrize("N", shapes) +def test_pow_vv(sigtype, N): + + def torch_func(x0, x1): + res = torch.pow(x0, x1) + return res + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, in_ptr1, N: tl.constexpr): + idx = tl.arange(0, N) + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = pow(x0, x1) + tl.store(out_ptr0 + idx, ret) + + def triton_func(x0, x1, N): + out = torch.empty_like(x0) + triton_kernel[1, 1, 1](out, x0, x1, N) + return out + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + x1 = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + + triton_cal = triton_func(x0, x1, N) + torch_ref = torch_func(x0, x1) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) + + +@pytest.mark.parametrize("sigtype", ["float32"]) +@pytest.mark.parametrize("N", shapes) +def test_pow_vs_dynamic(sigtype, N): + + def torch_func(x0, x1): + res = torch.pow(x0, x1) + return res + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, in_ptr1, N: tl.constexpr): + idx = tl.arange(0, N) + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1) + ret = pow(x0, x1) + tl.store(out_ptr0 + idx, ret) + + def triton_func(x0, x1, N): + out = torch.empty_like(x0) + triton_kernel[1, 1, 1](out, x0, x1, N) + return out + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + x1 = test_common.generate_tensor(shape=(1, ), dtype=sigtype).npu() + + triton_cal = triton_func(x0, x1, N) + torch_ref = torch_func(x0, x1) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) + + +@pytest.mark.parametrize("sigtype", ["float32"]) +@pytest.mark.parametrize("N", shapes) +def test_pow_vs_const(sigtype, N): + + def torch_func(x0, x1): + res = torch.pow(x0, x1) + return res + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, x1: tl.constexpr, N: tl.constexpr): + idx = tl.arange(0, N) + x0 = tl.load(in_ptr0 + idx) + ret = pow(x0, x1) + tl.store(out_ptr0 + idx, ret) + + def triton_func(x0, x1, N): + out = torch.empty_like(x0) + triton_kernel[1, 1, 1](out, x0, x1, N) + return out + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + x1 = test_common.generate_tensor(shape=(1, ), dtype=sigtype).item() + + triton_cal = triton_func(x0, x1, N) + torch_ref = torch_func(x0, x1) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_precise_div.py b/third_party/ascend/unittest/pytest_ut/test_precise_div.py new file mode 100644 index 0000000000..c40afb552e --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_precise_div.py @@ -0,0 +1,60 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl + +import torch +import torch_npu +import test_common + + +def torch_divRn(x0, x1): + return x0 / x1 + + +@triton.jit +def triton_divRn(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.load(in_ptr1 + (x0), None) + tmp2 = tl.div_rn(tmp0, tmp1) + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 32, 2048, 64], +]) +def test_divRn(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype) + x2 = x1.masked_fill(x1 == 0, 1) + x2 = x2.npu() + y_ref = torch_divRn(x0, x2) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_divRn[ncore, 1, 1](x0, x2, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_precise_sqrt.py b/third_party/ascend/unittest/pytest_ut/test_precise_sqrt.py new file mode 100644 index 0000000000..1aaa97cceb --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_precise_sqrt.py @@ -0,0 +1,55 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl + +torch.set_printoptions(precision=10) + + +@triton.jit +def sqrtrn_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + id = tl.program_id(axis=0) + start = id * BLOCK_SIZE + offsets = start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + + output = x + tl.sqrt_rn(y) + tl.store(output_ptr + offsets, output, mask=mask) + + +def sqrtrn(x: torch.Tensor, y: torch.Tensor): + output = torch.empty_like(y) + grid = lambda meta: (triton.cdiv(output.numel(), meta['BLOCK_SIZE']), ) + sqrtrn_kernel[grid](x, y, output, output.numel(), BLOCK_SIZE=512) + return output + + +def test_sqrtrn_fp32(): + size = 10240 + x = torch.abs(torch.randn(size, device='npu', dtype=torch.float32)) + y = torch.abs(torch.randn(size, device='npu', dtype=torch.float32)) + ref = x + torch.sqrt(y) + cal = sqrtrn(x, y) + torch.testing.assert_close(cal, ref, rtol=1e-06, atol=1e-06, equal_nan=True) diff --git a/third_party/ascend/unittest/pytest_ut/test_ptr_add.py b/third_party/ascend/unittest/pytest_ut/test_ptr_add.py new file mode 100644 index 0000000000..bb9f685640 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_ptr_add.py @@ -0,0 +1,64 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import pytest + + +def ptr_add(device): + + @triton.jit + def wrap_stacked(a_ptr, c_ptr, M: tl.constexpr, N: tl.constexpr, stride_am: tl.constexpr, stride_an: tl.constexpr, + stride_cm: tl.constexpr, stride_cn: tl.constexpr, BLOCK_SIZE_K: tl.constexpr): + offs_am = (2 + tl.arange(0, 4)) % M + offs_an = tl.arange(0, 4) + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_an[None, :] * stride_an) + + offs_cm = tl.arange(0, 4) + offs_cn = tl.arange(0, 4) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + for k in range(0, 2): + a = tl.load(a_ptrs) + tl.store(c_ptrs, a) + a_ptrs += BLOCK_SIZE_K * stride_an + c_ptrs += BLOCK_SIZE_K * stride_an + + M = 4 + N = 8 + AA = torch.arange(0, M * N, dtype=torch.float32).npu() + A = AA.reshape((M, N)) + out = torch.full((M, N), 88888, dtype=torch.float32).npu() + grid = lambda meta: (1, ) + + wrap_stacked[grid](A, out, M, N, A.stride(0), A.stride(1), out.stride(0), out.stride(1), BLOCK_SIZE_K=4) + + # Expected output copied from running triton on NPU + expected_out = torch.tensor([[16, 17, 18, 19, 20, 21, 22, 23], [24, 25, 26, 27, 28, 29, 30, 31], + [0, 1, 2, 3, 4, 5, 6, 7], [8, 9, 10, 11, 12, 13, 14, 15]], ).npu() + + assert torch.equal(expected_out.int(), out.int()) + + +def test_ptr_add(): + device = 'npu' + ptr_add(device) diff --git a/third_party/ascend/unittest/pytest_ut/test_rand.py b/third_party/ascend/unittest/pytest_ut/test_rand.py new file mode 100644 index 0000000000..cce83ba2ac --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_rand.py @@ -0,0 +1,88 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common +import math + + +@triton.jit +def kernel_rand(x_ptr, n_rounds: tl.constexpr, N: tl.constexpr, XBLOCK: tl.constexpr): + block_offset = tl.program_id(0) * XBLOCK + block_size = XBLOCK if block_offset + XBLOCK <= N else N - block_offset + for inner_idx in range(block_size): + global_offset = block_offset + inner_idx + rand_vals = tl.rand(5, 10 + global_offset, n_rounds) # 对每个索引生成一个随机数 + tl.store(x_ptr + global_offset, rand_vals) # 存储随机数 + + +@triton.jit +def kernel_randn(x_ptr, n_rounds: tl.constexpr, N: tl.constexpr, XBLOCK: tl.constexpr): + block_offset = tl.program_id(0) * XBLOCK + block_size = XBLOCK if block_offset + XBLOCK <= N else N - block_offset + for inner_idx in range(block_size): + global_offset = block_offset + inner_idx + rand_vals = tl.randn(5, 10 + global_offset, n_rounds) # 对每个索引生成一个随机数 + tl.store(x_ptr + global_offset, rand_vals) # 存储随机数 + + +@triton.jit +def kernel_randint(x_ptr, n_rounds: tl.constexpr, N: tl.constexpr, XBLOCK: tl.constexpr): + block_offset = tl.program_id(0) * XBLOCK + block_size = XBLOCK if block_offset + XBLOCK <= N else N - block_offset + for inner_idx in range(block_size): + global_offset = block_offset + inner_idx + rand_vals = tl.randint(5, 10 + global_offset, n_rounds) # 对每个索引生成一个随机数 + tl.store(x_ptr + global_offset, rand_vals) # 存储随机数 + + +@triton.jit +def kernel_randint4x(x_ptr, n_rounds: tl.constexpr, N: tl.constexpr, XBLOCK: tl.constexpr): + block_offset = tl.program_id(0) * XBLOCK + indices = tl.arange(0, 4) + block_size = XBLOCK if block_offset + XBLOCK <= N else N - block_offset + for inner_idx in range(0, block_size, step=4): + global_offset = block_offset + inner_idx + rand_vals, _, _, _ = tl.randint4x(5, 10 + global_offset, n_rounds) # 对每个索引生成一个随机数 + mask = (global_offset + indices) < N + tl.store(x_ptr + global_offset + indices, rand_vals, mask) # 存储随机数 + + +shapes = [(1, 3)] + + +@pytest.mark.parametrize('shape', shapes) +def test_case(shape): + y_calf = torch.zeros(shape, dtype=eval('torch.float32')).npu() + + numel = y_calf.numel() + ncore = 1 if numel < 32 else 32 + xblock = math.ceil(numel / ncore) + + kernel_rand[ncore, 1, 1](y_calf, 10, numel, xblock) + kernel_randn[ncore, 1, 1](y_calf, 10, numel, xblock) + + y_cali = torch.zeros(shape, dtype=eval('torch.int32')).npu() + + kernel_randint[ncore, 1, 1](y_cali, 10, numel, xblock) + kernel_randint4x[ncore, 1, 1](y_cali, 10, numel, xblock) diff --git a/third_party/ascend/unittest/pytest_ut/test_range.py b/third_party/ascend/unittest/pytest_ut/test_range.py new file mode 100644 index 0000000000..2544f77a39 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_range.py @@ -0,0 +1,89 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time + +import torch +import torch_npu +import test_common + + +@triton.jit +def triton_add(in_ptr0, in_ptr1, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + idx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + x0 = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret = x0 + x1 + for i in tl.range(2, 5, 2): + ret = ret + x1 + for i in tl.static_range(2, 10, 3): + ret = ret + x0 + odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :] + tl.store(out_ptr0 + odx, ret) + + +testlist = [ + (3, 5, 8), +] + + +def get_torch_typename(dtype): + if dtype == 'float32': + tyname = torch.float32 + elif dtype == 'int32': + tyname = torch.int32 + elif dtype == 'int64': + tyname = torch.int64 + elif dtype == 'float16': + tyname = torch.float16 + elif dtype == 'bfloat16': + tyname = torch.bfloat16 + elif dtype == 'int16': + tyname = torch.int16 + elif dtype == 'int8': + tyname = torch.int8 + elif dtype == 'bool': + tyname = torch.bool + else: + raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + return tyname + + +typelist = ['int8', 'int16', 'int32', 'int64'] + + +@pytest.mark.parametrize('L, M, N', testlist) +@pytest.mark.parametrize('sigtype', typelist) +def test_add(sigtype, L, M, N): + dtype = get_torch_typename(sigtype) + shape = (L, M, N) + x0 = test_common.generate_tensor(shape=(L, M, N), dtype=sigtype).npu() + x1 = test_common.generate_tensor(shape=(L, M, N), dtype=sigtype).npu() + y_ref = x0 + x1 + x1 + x1 + x0 + x0 + x0 + output = torch.zeros(shape, dtype=dtype).npu() + triton_add[1, 1, 1](x0, x1, output, L, M, N) + test_common.validate_cmp(sigtype, output, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_ravel.py b/third_party/ascend/unittest/pytest_ut/test_ravel.py new file mode 100644 index 0000000000..4d41dd7019 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_ravel.py @@ -0,0 +1,71 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, can_reorder: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.ravel(X, can_reorder=can_reorder) + + oidx = tl.arange(0, XB * YB * ZB) + tl.store(output_ptr + oidx, ret) + + +testlist = [ + ('float32', torch.float32, 2, 256, 16), + ('float32', torch.float32, 8, 8, 4), + ('float16', torch.float16, 2, 256, 16), + ('float16', torch.float16, 8, 8, 4), + ('int8', torch.int8, 2, 256, 16), + ('int8', torch.int8, 8, 8, 4), +] + + +@pytest.mark.parametrize('sigtype, dtype, XB, YB, ZB', testlist) +@pytest.mark.parametrize('can_reorder', [False, True]) +def test_ravel(sigtype, dtype, XB, YB, ZB, can_reorder): + + x = torch.randint(low=-128, high=128, size=(XB, YB, ZB), dtype=dtype).npu() + ans = torch.ravel(x) + + print(ans[0:16]) + + output = torch.randint(1, (XB * YB * ZB, ), dtype=dtype).npu() + + fn_npu_[1, 1, 1](output, x, XB, YB, ZB, can_reorder) + + print(output[0:16]) + + test_common.validate_cmp(sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_reduce_count_vector.py b/third_party/ascend/unittest/pytest_ut/test_reduce_count_vector.py new file mode 100644 index 0000000000..b4ebd2df46 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_reduce_count_vector.py @@ -0,0 +1,151 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time +import test_common + +import torch +import torch_npu + + +def standard_count(x0, cmp_val, dim): + res = (x0 == cmp_val).sum(dim=dim) + return res + + +def standard_gt(x0, cmp_val, dim): + res = (x0 > cmp_val).sum(dim=dim) + return res + + +def standard_lt(x0, cmp_val, dim): + res = (x0 < cmp_val).sum(dim=dim) + return res + + +@triton.jit +def triton_count(in_ptr0, out_ptr0, cmp_val, dim: tl.constexpr, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, N) + x = tl.load(in_ptr0 + idx_block) + + tmp3 = (x == cmp_val) + # tmp3 bool -> tl.float32 + tmp4 = tmp3.to(tl.float32) + res = tl.sum(tmp4, dim) + + tl.store(out_ptr0 + idx_block, res) + + +@triton.jit +def triton_gt(in_ptr0, out_ptr0, cmp_val, dim: tl.constexpr, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, N) + x = tl.load(in_ptr0 + idx_block) + + tmp3 = (x > cmp_val) + # tmp3 bool -> tl.float32 + tmp4 = tmp3.to(tl.float32) + res = tl.sum(tmp4, dim) + + tl.store(out_ptr0 + idx_block, res) + + +@triton.jit +def triton_lt(in_ptr0, out_ptr0, cmp_val, dim: tl.constexpr, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, N) + x = tl.load(in_ptr0 + idx_block) + + tmp3 = (x < cmp_val) + # tmp3 bool -> tl.float32 + tmp4 = tmp3.to(tl.float32) + res = tl.sum(tmp4, dim) + + tl.store(out_ptr0 + idx_block, res) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + (torch.bfloat16, 'bfloat16'), + (torch.int8, 'int8'), + (torch.int16, 'int16'), + (torch.int32, 'int32'), + (torch.int64, 'int64'), +] + +# if shape axis = 32/256 , then actual shape = axis/element_size() +shapes = [ + (32, 32), +] + +map_for_64_t = {37: 31} + +CPM_VAL_INT = 8 +CPM_VAL_FLOAT = 0.5 + +# TO BE FIXED with mask +ops = [ + ('counti', triton_count, standard_count, CPM_VAL_INT), + ('countf', triton_gt, standard_gt, CPM_VAL_FLOAT), + ('countf', triton_lt, standard_lt, CPM_VAL_FLOAT), +] + + +def judge_continue(opName, sigtype): + if opName == 'counti' and 'int' in sigtype: + return False + if opName == 'countf' and 'float' in sigtype: + return False + return True + + +@pytest.mark.parametrize('opName, tritonOp, standOp, cmp_val', ops) +@pytest.mark.parametrize('dtype, sigtype', types) +@pytest.mark.parametrize('N, NUMEL', shapes) +def test_reduce_count_vector(opName, tritonOp, standOp, cmp_val, dtype, sigtype, N, NUMEL): + if judge_continue(opName, sigtype): + return + torch.manual_seed(0) + torch_npu.npu.utils.set_device(0) + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + N = map_for_64_t[N] if N in map_for_64_t else N + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + ans = standOp(x0, cmp_val, 0) + x0 = x0.npu() + + output = torch.tensor(0, dtype=torch.float32).npu() + tritonOp[1, 1, 1](x0, output, cmp_val, dim=0, N=N, NUMEL=NUMEL, debug=True) + output = output.cpu().to(torch.int32) + # print(f'x0:{x0}\ntriton:{output}\ntorch:{ans}') + assert torch.equal(output, ans) + + +if __name__ == "__main__": + dtype = torch.float32 + sigtype = 'float32' + allshape = [(3, 32)] + for shape in allshape: + test_reduce_count_vector('countf', triton_lt, standard_lt, CPM_VAL_FLOAT, dtype, sigtype, shape[0], shape[1]) diff --git a/third_party/ascend/unittest/pytest_ut/test_reduce_maximum.py b/third_party/ascend/unittest/pytest_ut/test_reduce_maximum.py new file mode 100644 index 0000000000..ec29ccbbfe --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_reduce_maximum.py @@ -0,0 +1,308 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +@triton.jit +def maximum(a, b): + ret = tl.maximum(a, b, tl.PropagateNan.ALL) + # 经过测试发现，tl.maximum仅在输入类型为bfloat16时，输出的结果会转变为float32，从而导致编译报错。在GPU上测试发现，和NPU上错误的现象一致。 + # 因此此处针对输入类型为bfloat16的情况，对输出进行了类型转换来规避该错误引起的编译报错。 + if a.dtype == tl.bfloat16: + ret = ret.to(tl.bfloat16) + return ret + + +@triton.jit +def triton_max_5d_dim024(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, + Z: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + kblk_idx = tl.arange(0, K) + zblk_idx = tl.arange(0, Z) + idx = lblk_idx[:, None, None, None, None] * Z * K * N * M + mblk_idx[None, :, None, None, None] * Z * K * N + \ + nblk_idx[None, None, :, None, None] * Z * K + kblk_idx[None, None, None, :, None] * Z + zblk_idx[ + None, None, None, None, :] + odx = mblk_idx[:, None] * K + kblk_idx[None, :] + x = tl.load(in_ptr0 + idx) + ret = tl.reduce(x, 4, maximum) + ret1 = tl.reduce(ret, 2, maximum) + ret2 = tl.reduce(ret1, 0, maximum) + tl.store(out_ptr0 + odx, ret2) + + +@triton.jit +def triton_max_5d_dim13(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, + Z: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + kblk_idx = tl.arange(0, K) + zblk_idx = tl.arange(0, Z) + + idx = (lblk_idx[:, None, None, None, None] * Z * K * N * M + mblk_idx[None, :, None, None, None] * Z * K * N + + nblk_idx[None, None, :, None, None] * Z * K + kblk_idx[None, None, None, :, None] * Z + + zblk_idx[None, None, None, None, :]) + + odx = (lblk_idx[:, None, None] * N * Z + nblk_idx[None, :, None] * Z + zblk_idx[None, None, :]) + + x = tl.load(in_ptr0 + idx) + ret_k = tl.reduce(x, 3, maximum) # [L, M, N, Z] + ret_m = tl.reduce(ret_k, 1, maximum) # [L, N, Z] + tl.store(out_ptr0 + odx, ret_m) + + +@triton.jit +def triton_max_5d_dim0(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, + Z: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + kblk_idx = tl.arange(0, K) + zblk_idx = tl.arange(0, Z) + + idx = (lblk_idx[:, None, None, None, None] * Z * K * N * M + mblk_idx[None, :, None, None, None] * Z * K * N + + nblk_idx[None, None, :, None, None] * Z * K + kblk_idx[None, None, None, :, None] * Z + + zblk_idx[None, None, None, None, :]) + + odx = (mblk_idx[:, None, None, None] * N * K * Z + nblk_idx[None, :, None, None] * K * Z + + kblk_idx[None, None, :, None] * Z + zblk_idx[None, None, None, :]) + + x = tl.load(in_ptr0 + idx) + ret = tl.reduce(x, 0, maximum) + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_max_5d_dim1(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, + Z: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + kblk_idx = tl.arange(0, K) + zblk_idx = tl.arange(0, Z) + + idx = (lblk_idx[:, None, None, None, None] * Z * K * N * M + mblk_idx[None, :, None, None, None] * Z * K * N + + nblk_idx[None, None, :, None, None] * Z * K + kblk_idx[None, None, None, :, None] * Z + + zblk_idx[None, None, None, None, :]) + + odx = (lblk_idx[:, None, None, None] * N * K * Z + nblk_idx[None, :, None, None] * K * Z + + kblk_idx[None, None, :, None] * Z + zblk_idx[None, None, None, :]) + + x = tl.load(in_ptr0 + idx) + ret = tl.reduce(x, 1, maximum) + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_max_5d_dim2(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, + Z: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + kblk_idx = tl.arange(0, K) + zblk_idx = tl.arange(0, Z) + + idx = (lblk_idx[:, None, None, None, None] * Z * K * N * M + mblk_idx[None, :, None, None, None] * Z * K * N + + nblk_idx[None, None, :, None, None] * Z * K + kblk_idx[None, None, None, :, None] * Z + + zblk_idx[None, None, None, None, :]) + + odx = (lblk_idx[:, None, None, None] * M * K * Z + mblk_idx[None, :, None, None] * K * Z + + kblk_idx[None, None, :, None] * Z + zblk_idx[None, None, None, :]) + + x = tl.load(in_ptr0 + idx) + ret = tl.reduce(x, 2, maximum) + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_max_5d_dim3(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, + Z: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + kblk_idx = tl.arange(0, K) + zblk_idx = tl.arange(0, Z) + + idx = (lblk_idx[:, None, None, None, None] * Z * K * N * M + mblk_idx[None, :, None, None, None] * Z * K * N + + nblk_idx[None, None, :, None, None] * Z * K + kblk_idx[None, None, None, :, None] * Z + + zblk_idx[None, None, None, None, :]) + + odx = (lblk_idx[:, None, None, None] * M * N * Z + mblk_idx[None, :, None, None] * N * Z + + nblk_idx[None, None, :, None] * Z + zblk_idx[None, None, None, :]) + + x = tl.load(in_ptr0 + idx) + ret = tl.reduce(x, 3, maximum) + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_max_5d_dim4(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, + Z: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + kblk_idx = tl.arange(0, K) + zblk_idx = tl.arange(0, Z) + + idx = (lblk_idx[:, None, None, None, None] * Z * K * N * M + mblk_idx[None, :, None, None, None] * Z * K * N + + nblk_idx[None, None, :, None, None] * Z * K + kblk_idx[None, None, None, :, None] * Z + + zblk_idx[None, None, None, None, :]) + + odx = (lblk_idx[:, None, None, None] * M * N * K + mblk_idx[None, :, None, None] * N * K + + nblk_idx[None, None, :, None] * K + kblk_idx[None, None, None, :]) + + x = tl.load(in_ptr0 + idx) + ret = tl.reduce(x, 4, maximum) + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_max_5d_all(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, + Z: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + kblk_idx = tl.arange(0, K) + zblk_idx = tl.arange(0, Z) + + idx = (lblk_idx[:, None, None, None, None] * Z * K * N * M + mblk_idx[None, :, None, None, None] * Z * K * N + + nblk_idx[None, None, :, None, None] * Z * K + kblk_idx[None, None, None, :, None] * Z + + zblk_idx[None, None, None, None, :]) + + x = tl.load(in_ptr0 + idx) + ret1 = tl.reduce(x, 4, maximum) + ret2 = tl.reduce(ret1, 3, maximum) + ret3 = tl.reduce(ret2, 2, maximum) + ret4 = tl.reduce(ret3, 1, maximum) + ret5 = tl.reduce(ret4, 0, maximum) + tl.store(out_ptr0, ret5) + + +testlist = [ + (triton_max_5d_dim024, (1, 1, 1, 1, 1), "dim024"), + (triton_max_5d_dim024, (2, 2, 2, 2, 2), "dim024"), + (triton_max_5d_dim024, (3, 11, 1, 3, 42), "dim024"), + (triton_max_5d_dim13, (1, 1, 1, 1, 1024), "dim13"), + (triton_max_5d_dim0, (2, 2, 2, 2, 2), "dim0"), + (triton_max_5d_dim1, (2, 2, 2, 2, 2), "dim1"), + (triton_max_5d_dim2, (2, 2, 2, 2, 2), "dim2"), + (triton_max_5d_dim3, (2, 2, 2, 2, 2), "dim3"), + (triton_max_5d_dim4, (2, 2, 2, 2, 2), "dim4"), + (triton_max_5d_all, (3, 11, 1, 3, 42), "all"), +] + +typelist = ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16', 'bool'] + +ids = [ + "{}-{}-{}".format(testfunc.__name__, "-".join(map(str, shape)), dim_name) for testfunc, shape, dim_name in testlist +] + + +@pytest.mark.parametrize('testfunc, shape, dim_name', testlist, ids=ids) +@pytest.mark.parametrize('dtype', typelist) +def test_max(testfunc, dtype, shape, dim_name): + x0 = test_common.generate_tensor(shape=shape, dtype=dtype).npu() + + if dim_name == "dim024": + if 'int' in dtype: + ans, _ = torch.max(x0.to(torch.int64), 4) + ans, _ = torch.max(ans.to(torch.int64), 2) + ans, _ = torch.max(ans.to(torch.int64), 0) + ans = ans.to(dtype=eval('torch.' + dtype)) + else: + ans, _ = torch.max(x0, 4) + ans, _ = torch.max(ans, 2) + ans, _ = torch.max(ans, 0) + output = torch.zeros((shape[1], ) + (shape[3], ), dtype=eval('torch.' + dtype)).npu() + + elif dim_name == "dim13": + if 'int' in dtype: + ans, _ = torch.max(x0.to(torch.int64), 3) + ans, _ = torch.max(ans.to(torch.int64), 1) + ans = ans.to(dtype=eval('torch.' + dtype)) + else: + ans, _ = torch.max(x0, 3) + ans, _ = torch.max(ans, 1) + output = torch.zeros((shape[0], ) + (shape[2], ) + (shape[4], ), dtype=eval('torch.' + dtype)).npu() + + elif dim_name == "dim0": + if 'int' in dtype: + ans, _ = torch.max(x0.to(torch.int64), 0) + ans = ans.to(dtype=eval('torch.' + dtype)) + else: + ans, _ = torch.max(x0, 0) + output = torch.zeros((shape[1], ) + (shape[2], ) + (shape[3], ) + (shape[4], ), + dtype=eval('torch.' + dtype)).npu() + + elif dim_name == "dim1": + if 'int' in dtype: + ans, _ = torch.max(x0.to(torch.int64), 1) + ans = ans.to(dtype=eval('torch.' + dtype)) + else: + ans, _ = torch.max(x0, 1) + output = torch.zeros((shape[0], ) + (shape[2], ) + (shape[3], ) + (shape[4], ), + dtype=eval('torch.' + dtype)).npu() + + elif dim_name == "dim2": + if 'int' in dtype: + ans, _ = torch.max(x0.to(torch.int64), 2) + ans = ans.to(dtype=eval('torch.' + dtype)) + else: + ans, _ = torch.max(x0, 2) + output = torch.zeros((shape[0], ) + (shape[1], ) + (shape[3], ) + (shape[4], ), + dtype=eval('torch.' + dtype)).npu() + + elif dim_name == "dim3": + if 'int' in dtype: + ans, _ = torch.max(x0.to(torch.int64), 3) + ans = ans.to(dtype=eval('torch.' + dtype)) + else: + ans, _ = torch.max(x0, 3) + output = torch.zeros((shape[0], ) + (shape[1], ) + (shape[2], ) + (shape[4], ), + dtype=eval('torch.' + dtype)).npu() + + elif dim_name == "dim4": + if 'int' in dtype: + ans, _ = torch.max(x0.to(torch.int64), 4) + ans = ans.to(dtype=eval('torch.' + dtype)) + else: + ans, _ = torch.max(x0, 4) + output = torch.zeros((shape[0], ) + (shape[1], ) + (shape[2], ) + (shape[3], ), + dtype=eval('torch.' + dtype)).npu() + + elif dim_name == "all": + if 'int' in dtype: + ans = torch.max(x0.to(torch.int64)) + ans = torch.tensor([ans], dtype=eval('torch.' + dtype)) + else: + ans = torch.tensor([torch.max(x0)], dtype=eval('torch.' + dtype)) + output = torch.zeros((1, ), dtype=eval('torch.' + dtype)).npu() + + testfunc[(1, )](x0, output, *shape) + + test_common.validate_cmp(dtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_reduce_mean.py b/third_party/ascend/unittest/pytest_ut/test_reduce_mean.py new file mode 100644 index 0000000000..f22cd47c39 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_reduce_mean.py @@ -0,0 +1,89 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common +import numpy as np + + +def numpy_mean_pr(x0, x1): + res = np.mean(x0, axis=-1) + x1 + return res + + +@triton.jit +def triton_mean_pr(out_ptr0, in_ptr0, in_ptr1, xnumel, rnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr, + RBLOCK: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + rindex = tl.arange(0, RBLOCK)[None, :] + rmask = rindex < rnumel + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + xindex = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB) + xmask = xindex[:, None] < xnumel + x0 = xindex + r1 = rindex + tmp0 = tl.load(in_ptr0 + (r1 + (RBLOCK * x0[:, None])), xmask & rmask) + tmp4 = tl.load(in_ptr1 + (x0), xindex < xnumel) + tmp1 = tl.reshape(tmp0, [XBLOCK_SUB, RBLOCK]) + tmp3 = tl.sum(tmp1, 1) / RBLOCK + tmp5 = tmp3 + tmp4 + tl.store(out_ptr0 + (xindex), tmp5, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (8, 8, 4), 8, 2], + ['float32', (8, 8, 64), 8, 2], + ['float32', (8, 8, 1024), 8, 2], + ['float16', (8, 8, 4), 8, 2], + ['float16', (8, 8, 64), 8, 2], + ['float16', (8, 8, 1024), 8, 2], + ['int8', (8, 8, 4), 8, 2], + ['int8', (8, 8, 64), 8, 2], + pytest.param( + ['int8', (8, 8, 1024), 8, 2], + marks=pytest.mark.skip( + reason="tl.sum does not support int8 reduction with RBLOCK=1024, and will be fixed later"), + ), +]) +def test_mean_pr(param_list): + dtype, shape, ncore, xblock_sub = param_list + import math + numel = math.prod(shape) + xblock = numel // shape[-1] // ncore + rblock = shape[-1] + assert (ncore * xblock * shape[-1] == numel) + xn1 = np.random.randn(shape[0], shape[1], shape[2]).astype(eval('np.' + dtype)) + xn2 = np.random.randn(shape[0], shape[1]).astype(eval('np.' + dtype)) + x0 = torch.tensor(xn1).npu() + x1 = torch.tensor(xn2).npu() + y_ref = numpy_mean_pr(xn1, xn2) + if dtype == 'int8': + y_cal = test_common.generate_tensor(shape[:-1], 'float32').npu() + else: + y_cal = test_common.generate_tensor(shape[:-1], dtype).npu() + triton_mean_pr[ncore, 1, 1](y_cal, x0, x1, x1.numel(), rblock, xblock, xblock_sub, rblock) + if dtype == 'int8': + torch.allclose(torch.tensor(y_ref.astype(np.float32)).npu(), y_cal, rtol=1e-03, atol=1e-03, equal_nan=True) + else: + torch.allclose(torch.tensor(y_ref).npu(), y_cal, rtol=1e-03, atol=1e-03, equal_nan=True) diff --git a/third_party/ascend/unittest/pytest_ut/test_reduce_min_4_keepdim_True_with_index_op.py b/third_party/ascend/unittest/pytest_ut/test_reduce_min_4_keepdim_True_with_index_op.py new file mode 100644 index 0000000000..26ffbe0c51 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_reduce_min_4_keepdim_True_with_index_op.py @@ -0,0 +1,96 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Huawei Technologies Co., Ltd. 2023-2024. All rights reserved. +import time + +import pytest +import torch +import torch_npu + +import triton +import triton.language as tl +import test_common + + +@triton.jit +def promote_to_tensor(x): + # Addition promotes to tensor for us + return x + tl.zeros((1, ), tl.int1) + + +@triton.jit +def minimum_with_index(a_value, a_index, b_value, b_index): + mask = a_value < b_value + equal = a_value == b_value + if promote_to_tensor(a_value).dtype.is_floating(): + a_isnan = a_value != a_value + b_isnan = b_value != b_value + mask |= a_isnan and not b_isnan + # Consider NaNs as equal + equal |= a_isnan and b_isnan + # Prefer lowest index if values are equal + mask |= equal & (a_index < b_index) + return tl.where(mask, a_value, b_value), tl.where(mask, a_index, b_index) + + +@triton.jit +def triton_min_5d_dim4_keepdim(in_ptr0, in_ptr1, out_ptr0, out_ptr1, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr, + K: tl.constexpr, Z: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + kblk_idx = tl.arange(0, K) + zblk_idx = tl.arange(0, Z) + idx = lblk_idx[:, None, None, None, None] * Z * K * N * M + mblk_idx[None, :, None, None, None] * Z * K * N + \ + nblk_idx[None, None, :, None, None] * Z * K + kblk_idx[None, None, None, :, None] * Z + zblk_idx[None, None, None, None, :] + x = tl.load(in_ptr0 + idx) + x1 = tl.load(in_ptr1 + idx) + ret, ret1 = tl.reduce((x, x1), 4, minimum_with_index, keep_dims=True) + zblk_idx = tl.arange(0, 1) + odx = lblk_idx[:, None, None, None, None] * K * N * M + mblk_idx[None, :, None, None, None] * K * N + \ + nblk_idx[None, None, :, None, None] * K + kblk_idx[None, None, None, :, None] \ + + zblk_idx[None, None, None, None, :] + tl.store(out_ptr0 + odx, ret) + tl.store(out_ptr1 + odx, ret1) + + +testlist = [ + # 5D + (triton_min_5d_dim4_keepdim, (1, 1, 1, 1, 1)), + (triton_min_5d_dim4_keepdim, (2, 2, 2, 2, 2)), + (triton_min_5d_dim4_keepdim, (9, 3, 2, 4, 17)), + (triton_min_5d_dim4_keepdim, (3, 11, 1, 3, 42)), + (triton_min_5d_dim4_keepdim, (2, 51, 3, 13, 1)), + (triton_min_5d_dim4_keepdim, (129, 1, 5, 1, 4)), + (triton_min_5d_dim4_keepdim, (203, 1, 2, 2, 3)), + (triton_min_5d_dim4_keepdim, (512, 1, 1, 1, 1)), + (triton_min_5d_dim4_keepdim, (3, 1, 1, 2, 600)), + (triton_min_5d_dim4_keepdim, (1, 1, 1, 1, 1024)), + (triton_min_5d_dim4_keepdim, (15, 2, 2, 2, 54)), + (triton_min_5d_dim4_keepdim, (2, 91, 4, 2, 4)), + (triton_min_5d_dim4_keepdim, (1, 1, 3, 2, 600)), + (triton_min_5d_dim4_keepdim, (5, 2, 4, 1, 26)), + (triton_min_5d_dim4_keepdim, (2, 2, 2, 4, 8)), +] + +typelist = ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16', 'bool'] + +ids = ["{}-{}".format(testfunc.__name__, "-".join(map(str, shape))) for testfunc, shape in testlist] + + +@pytest.mark.parametrize('testfunc, shape', testlist, ids=ids) +@pytest.mark.parametrize('sigtype', typelist) +def test_min_dim4_keepdim(testfunc, sigtype, shape): + dtype = eval('torch.' + sigtype) + x0 = torch.randn(shape).to(dtype).npu() + + x1 = torch.arange(x0.numel()).view(x0.shape).npu().to(torch.int32) + if 'int' in sigtype: + ans, ans1 = torch.min(x0.to(torch.int64), 4) + ans = ans.to(dtype) + else: + ans, ans1 = torch.min(x0, 4) + output = torch.zeros(shape[0:4], dtype=dtype).npu() + output1 = torch.zeros(shape[0:4], dtype=torch.int32).npu() + testfunc[(1, )](x0, x1, output, output1, *shape, debug=True) + test_common.validate_cmp(sigtype, output, ans) + test_common.validate_cmp('int32', output1, ans1) diff --git a/third_party/ascend/unittest/pytest_ut/test_reduce_minimum.py b/third_party/ascend/unittest/pytest_ut/test_reduce_minimum.py new file mode 100644 index 0000000000..f472eef67e --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_reduce_minimum.py @@ -0,0 +1,308 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +@triton.jit +def minimum(a, b): + ret = tl.minimum(a, b, tl.PropagateNan.ALL) + # 经过测试发现，tl.minimum仅在输入类型为bfloat16时，输出的结果会转变为float32，从而导致编译报错。在GPU上测试发现，和NPU上错误的现象一致。 + # 因此此处针对输入类型为bfloat16的情况，对输出进行了类型转换来规避该错误引起的编译报错。 + if a.dtype == tl.bfloat16: + ret = ret.to(tl.bfloat16) + return ret + + +@triton.jit +def triton_min_5d_dim024(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, + Z: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + kblk_idx = tl.arange(0, K) + zblk_idx = tl.arange(0, Z) + idx = lblk_idx[:, None, None, None, None] * Z * K * N * M + mblk_idx[None, :, None, None, None] * Z * K * N + \ + nblk_idx[None, None, :, None, None] * Z * K + kblk_idx[None, None, None, :, None] * Z + zblk_idx[ + None, None, None, None, :] + odx = mblk_idx[:, None] * K + kblk_idx[None, :] + x = tl.load(in_ptr0 + idx) + ret = tl.reduce(x, 4, minimum) + ret1 = tl.reduce(ret, 2, minimum) + ret2 = tl.reduce(ret1, 0, minimum) + tl.store(out_ptr0 + odx, ret2) + + +@triton.jit +def triton_min_5d_dim13(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, + Z: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + kblk_idx = tl.arange(0, K) + zblk_idx = tl.arange(0, Z) + + idx = (lblk_idx[:, None, None, None, None] * Z * K * N * M + mblk_idx[None, :, None, None, None] * Z * K * N + + nblk_idx[None, None, :, None, None] * Z * K + kblk_idx[None, None, None, :, None] * Z + + zblk_idx[None, None, None, None, :]) + + odx = (lblk_idx[:, None, None] * N * Z + nblk_idx[None, :, None] * Z + zblk_idx[None, None, :]) + + x = tl.load(in_ptr0 + idx) + ret_k = tl.reduce(x, 3, minimum) # [L, M, N, Z] + ret_m = tl.reduce(ret_k, 1, minimum) # [L, N, Z] + tl.store(out_ptr0 + odx, ret_m) + + +@triton.jit +def triton_min_5d_dim0(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, + Z: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + kblk_idx = tl.arange(0, K) + zblk_idx = tl.arange(0, Z) + + idx = (lblk_idx[:, None, None, None, None] * Z * K * N * M + mblk_idx[None, :, None, None, None] * Z * K * N + + nblk_idx[None, None, :, None, None] * Z * K + kblk_idx[None, None, None, :, None] * Z + + zblk_idx[None, None, None, None, :]) + + odx = (mblk_idx[:, None, None, None] * N * K * Z + nblk_idx[None, :, None, None] * K * Z + + kblk_idx[None, None, :, None] * Z + zblk_idx[None, None, None, :]) + + x = tl.load(in_ptr0 + idx) + ret = tl.reduce(x, 0, minimum) + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_min_5d_dim1(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, + Z: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + kblk_idx = tl.arange(0, K) + zblk_idx = tl.arange(0, Z) + + idx = (lblk_idx[:, None, None, None, None] * Z * K * N * M + mblk_idx[None, :, None, None, None] * Z * K * N + + nblk_idx[None, None, :, None, None] * Z * K + kblk_idx[None, None, None, :, None] * Z + + zblk_idx[None, None, None, None, :]) + + odx = (lblk_idx[:, None, None, None] * N * K * Z + nblk_idx[None, :, None, None] * K * Z + + kblk_idx[None, None, :, None] * Z + zblk_idx[None, None, None, :]) + + x = tl.load(in_ptr0 + idx) + ret = tl.reduce(x, 1, minimum) + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_min_5d_dim2(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, + Z: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + kblk_idx = tl.arange(0, K) + zblk_idx = tl.arange(0, Z) + + idx = (lblk_idx[:, None, None, None, None] * Z * K * N * M + mblk_idx[None, :, None, None, None] * Z * K * N + + nblk_idx[None, None, :, None, None] * Z * K + kblk_idx[None, None, None, :, None] * Z + + zblk_idx[None, None, None, None, :]) + + odx = (lblk_idx[:, None, None, None] * M * K * Z + mblk_idx[None, :, None, None] * K * Z + + kblk_idx[None, None, :, None] * Z + zblk_idx[None, None, None, :]) + + x = tl.load(in_ptr0 + idx) + ret = tl.reduce(x, 2, minimum) + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_min_5d_dim3(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, + Z: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + kblk_idx = tl.arange(0, K) + zblk_idx = tl.arange(0, Z) + + idx = (lblk_idx[:, None, None, None, None] * Z * K * N * M + mblk_idx[None, :, None, None, None] * Z * K * N + + nblk_idx[None, None, :, None, None] * Z * K + kblk_idx[None, None, None, :, None] * Z + + zblk_idx[None, None, None, None, :]) + + odx = (lblk_idx[:, None, None, None] * M * N * Z + mblk_idx[None, :, None, None] * N * Z + + nblk_idx[None, None, :, None] * Z + zblk_idx[None, None, None, :]) + + x = tl.load(in_ptr0 + idx) + ret = tl.reduce(x, 3, minimum) + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_min_5d_dim4(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, + Z: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + kblk_idx = tl.arange(0, K) + zblk_idx = tl.arange(0, Z) + + idx = (lblk_idx[:, None, None, None, None] * Z * K * N * M + mblk_idx[None, :, None, None, None] * Z * K * N + + nblk_idx[None, None, :, None, None] * Z * K + kblk_idx[None, None, None, :, None] * Z + + zblk_idx[None, None, None, None, :]) + + odx = (lblk_idx[:, None, None, None] * M * N * K + mblk_idx[None, :, None, None] * N * K + + nblk_idx[None, None, :, None] * K + kblk_idx[None, None, None, :]) + + x = tl.load(in_ptr0 + idx) + ret = tl.reduce(x, 4, minimum) + tl.store(out_ptr0 + odx, ret) + + +@triton.jit +def triton_min_5d_all(in_ptr0, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, + Z: tl.constexpr): + lblk_idx = tl.arange(0, L) + mblk_idx = tl.arange(0, M) + nblk_idx = tl.arange(0, N) + kblk_idx = tl.arange(0, K) + zblk_idx = tl.arange(0, Z) + + idx = (lblk_idx[:, None, None, None, None] * Z * K * N * M + mblk_idx[None, :, None, None, None] * Z * K * N + + nblk_idx[None, None, :, None, None] * Z * K + kblk_idx[None, None, None, :, None] * Z + + zblk_idx[None, None, None, None, :]) + + x = tl.load(in_ptr0 + idx) + ret1 = tl.reduce(x, 4, minimum) + ret2 = tl.reduce(ret1, 3, minimum) + ret3 = tl.reduce(ret2, 2, minimum) + ret4 = tl.reduce(ret3, 1, minimum) + ret5 = tl.reduce(ret4, 0, minimum) + tl.store(out_ptr0, ret5) + + +testlist = [ + (triton_min_5d_dim024, (1, 1, 1, 1, 1), "dim024"), + (triton_min_5d_dim024, (2, 2, 2, 2, 2), "dim024"), + (triton_min_5d_dim024, (3, 11, 1, 3, 42), "dim024"), + (triton_min_5d_dim13, (1, 1, 1, 1, 1024), "dim13"), + (triton_min_5d_dim0, (2, 2, 2, 2, 2), "dim0"), + (triton_min_5d_dim1, (2, 2, 2, 2, 2), "dim1"), + (triton_min_5d_dim2, (2, 2, 2, 2, 2), "dim2"), + (triton_min_5d_dim3, (2, 2, 2, 2, 2), "dim3"), + (triton_min_5d_dim4, (2, 2, 2, 2, 2), "dim4"), + (triton_min_5d_all, (3, 11, 1, 3, 42), "all"), +] + +typelist = ['int8', 'int16', 'int32', 'int64', 'float16', 'float32', 'bfloat16', 'bool'] + +ids = [ + "{}-{}-{}".format(testfunc.__name__, "-".join(map(str, shape)), dim_name) for testfunc, shape, dim_name in testlist +] + + +@pytest.mark.parametrize('testfunc, shape, dim_name', testlist, ids=ids) +@pytest.mark.parametrize('dtype', typelist) +def test_min(testfunc, dtype, shape, dim_name): + x0 = test_common.generate_tensor(shape=shape, dtype=dtype).npu() + + if dim_name == "dim024": + if 'int' in dtype: + ans, _ = torch.min(x0.to(torch.int64), 4) + ans, _ = torch.min(ans.to(torch.int64), 2) + ans, _ = torch.min(ans.to(torch.int64), 0) + ans = ans.to(dtype=eval('torch.' + dtype)) + else: + ans, _ = torch.min(x0, 4) + ans, _ = torch.min(ans, 2) + ans, _ = torch.min(ans, 0) + output = torch.zeros((shape[1], ) + (shape[3], ), dtype=eval('torch.' + dtype)).npu() + + elif dim_name == "dim13": + if 'int' in dtype: + ans, _ = torch.min(x0.to(torch.int64), 3) + ans, _ = torch.min(ans.to(torch.int64), 1) + ans = ans.to(dtype=eval('torch.' + dtype)) + else: + ans, _ = torch.min(x0, 3) + ans, _ = torch.min(ans, 1) + output = torch.zeros((shape[0], ) + (shape[2], ) + (shape[4], ), dtype=eval('torch.' + dtype)).npu() + + elif dim_name == "dim0": + if 'int' in dtype: + ans, _ = torch.min(x0.to(torch.int64), 0) + ans = ans.to(dtype=eval('torch.' + dtype)) + else: + ans, _ = torch.min(x0, 0) + output = torch.zeros((shape[1], ) + (shape[2], ) + (shape[3], ) + (shape[4], ), + dtype=eval('torch.' + dtype)).npu() + + elif dim_name == "dim1": + if 'int' in dtype: + ans, _ = torch.min(x0.to(torch.int64), 1) + ans = ans.to(dtype=eval('torch.' + dtype)) + else: + ans, _ = torch.min(x0, 1) + output = torch.zeros((shape[0], ) + (shape[2], ) + (shape[3], ) + (shape[4], ), + dtype=eval('torch.' + dtype)).npu() + + elif dim_name == "dim2": + if 'int' in dtype: + ans, _ = torch.min(x0.to(torch.int64), 2) + ans = ans.to(dtype=eval('torch.' + dtype)) + else: + ans, _ = torch.min(x0, 2) + output = torch.zeros((shape[0], ) + (shape[1], ) + (shape[3], ) + (shape[4], ), + dtype=eval('torch.' + dtype)).npu() + + elif dim_name == "dim3": + if 'int' in dtype: + ans, _ = torch.min(x0.to(torch.int64), 3) + ans = ans.to(dtype=eval('torch.' + dtype)) + else: + ans, _ = torch.min(x0, 3) + output = torch.zeros((shape[0], ) + (shape[1], ) + (shape[2], ) + (shape[4], ), + dtype=eval('torch.' + dtype)).npu() + + elif dim_name == "dim4": + if 'int' in dtype: + ans, _ = torch.min(x0.to(torch.int64), 4) + ans = ans.to(dtype=eval('torch.' + dtype)) + else: + ans, _ = torch.min(x0, 4) + output = torch.zeros((shape[0], ) + (shape[1], ) + (shape[2], ) + (shape[3], ), + dtype=eval('torch.' + dtype)).npu() + + elif dim_name == "all": + if 'int' in dtype: + ans = torch.min(x0.to(torch.int64)) + ans = torch.tensor([ans], dtype=eval('torch.' + dtype)) + else: + ans = torch.tensor([torch.min(x0)], dtype=eval('torch.' + dtype)) + output = torch.zeros((1, ), dtype=eval('torch.' + dtype)).npu() + + testfunc[(1, )](x0, output, *shape) + + test_common.validate_cmp(dtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_reduce_sum.py b/third_party/ascend/unittest/pytest_ut/test_reduce_sum.py new file mode 100644 index 0000000000..a435124417 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_reduce_sum.py @@ -0,0 +1,90 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import numpy as np + +import os +import sys + +parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..')) +sys.path.append(parent_dir) +import test_common + + +# PR: Pointiwise-Reduction pattern, reduction in last axis +def numpy_sum_pr(x0, x1): + res = np.sum(x0, axis=-1) + x1 + return res + + +@triton.jit +def triton_sum_pr(out_ptr0, in_ptr0, in_ptr1, xnumel, rnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr, + RBLOCK: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + rindex = tl.arange(0, RBLOCK)[None, :] + rmask = rindex < rnumel + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + xindex = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB) + xmask = xindex[:, None] < xnumel + x0 = xindex + r1 = rindex + tmp0 = tl.load(in_ptr0 + (r1 + (RBLOCK * x0[:, None])), xmask & rmask) + tmp4 = tl.load(in_ptr1 + (x0), xindex < xnumel) + tmp1 = tl.reshape(tmp0, [XBLOCK_SUB, RBLOCK]) + tmp3 = tl.sum(tmp1, 1) + tmp5 = tmp3 + tmp4 + tl.store(out_ptr0 + (xindex), tmp5, None) + + +# fp16 use numpy +@pytest.mark.parametrize('param_list', [ + # ['float32', (8, 8, 4), 8, 2], + ['float32', (8, 8, 64), 8, 2], + ['float32', (8, 8, 512), 8, 2], + ['float16', (8, 8, 4), 8, 2], + ['float16', (8, 8, 64), 8, 2], + ['float16', (8, 8, 512), 8, 2], + ['int8', (8, 8, 4), 8, 2], + ['int8', (8, 8, 64), 8, 2], + ['int8', (8, 8, 512), 8, 2], +]) +def test_sum_pr(param_list): + dtype, shape, ncore, xblock_sub = param_list + import math + numel = math.prod(shape) + xblock = numel // shape[-1] // ncore + rblock = shape[-1] + assert (ncore * xblock * shape[-1] == numel) + xn1 = np.random.randn(shape[0], shape[1], shape[2]).astype(eval('np.' + dtype)) + xn2 = np.random.randn(shape[0], shape[1]).astype(eval('np.' + dtype)) + x0 = torch.tensor(xn1).npu() + x1 = torch.tensor(xn2).npu() + if dtype == 'int8': + y_ref = numpy_sum_pr(xn1, xn2).astype(np.int8) + else: + y_ref = numpy_sum_pr(xn1, xn2) + y_cal = test_common.generate_tensor(shape[:-1], dtype).npu() + triton_sum_pr[ncore, 1, 1](y_cal, x0, x1, x1.numel(), rblock, xblock, xblock_sub, rblock) + test_common.validate_cmp(dtype, y_cal, torch.tensor(y_ref).npu()) diff --git a/third_party/ascend/unittest/pytest_ut/test_relu.py b/third_party/ascend/unittest/pytest_ut/test_relu.py new file mode 100644 index 0000000000..63e9bd9714 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_relu.py @@ -0,0 +1,61 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common +import triton.language.extra.cann.libdevice as libdevice + + +def torch_relu(x0, x1): + res = x0 + torch.relu(x1) + return res + + +@triton.jit +def triton_relu(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + x_index = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = x_index < xnumel + tmp0 = tl.load(in_ptr0 + x_index, xmask) + tmp1 = tl.load(in_ptr1 + x_index, xmask) + tmp2 = tmp0 + libdevice.relu(tmp1) + tl.store(out_ptr0 + x_index, tmp2, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], +]) +def test_relu(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + torch_res = torch_relu(x0, x1) + # triton结果 + triton_res = test_common.generate_tensor(shape, dtype).npu() + triton_relu[ncore, 1, 1](x0, x1, triton_res, x0.numel(), xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_reshape.py b/third_party/ascend/unittest/pytest_ut/test_reshape.py new file mode 100644 index 0000000000..11778d3c7d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_reshape.py @@ -0,0 +1,70 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.reshape(X, (ZB, XB * YB)) + + oidx = tl.arange(0, ZB)[:, None] * XB * YB + tl.arange(0, XB * YB)[None, :] + + tl.store(output_ptr + oidx, ret) + + +testlist = [ + ('float32', torch.float32, 2, 256, 16), + ('float32', torch.float32, 8, 8, 4), + ('float16', torch.float16, 2, 256, 16), + ('float16', torch.float16, 8, 8, 4), + ('int8', torch.int8, 2, 256, 16), + ('int8', torch.int8, 8, 8, 4), +] + + +@pytest.mark.parametrize('sigtype, dtype, XB, YB, ZB', testlist) +def test_ravel(sigtype, dtype, XB, YB, ZB): + + x = torch.randint(low=-128, high=128, size=(XB, YB, ZB), dtype=dtype).npu() + ans = torch.reshape(x, (ZB, XB * YB)) + + print(ans[0, 0:16]) + + output = torch.randint(1, (ZB, XB * YB), dtype=dtype).npu() + + fn_npu_[1, 1, 1](output, x, XB, YB, ZB) + print(output[0, 0:16]) + + test_common.validate_cmp(sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_resize_performance.py b/third_party/ascend/unittest/pytest_ut/test_resize_performance.py new file mode 100644 index 0000000000..6478e2f6b5 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_resize_performance.py @@ -0,0 +1,102 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch, torch_npu +import triton +import triton.language as tl +import numpy as np +import math +import pytest + + +@triton.jit +def nearest_resize_kernel_col_tile(img_src_ptr, img_dst_ptr, src_rows: tl.constexpr, src_cols: tl.constexpr, + dst_rows: tl.constexpr, dst_cols: tl.constexpr, RR_H: tl.constexpr, + RR_W: tl.constexpr, stride_in_h: tl.constexpr, stride_in_w: tl.constexpr, + stride_in_c: tl.constexpr, stride_out_h: tl.constexpr, stride_out_w: tl.constexpr, + stride_out_c: tl.constexpr, BLOCK_SIZE: tl.constexpr): + #RR_H和RR_W分别为高和宽的缩放比例 + block_id_c = tl.program_id(0) + block_id_h = tl.program_id(1) + block_id_w = tl.program_id(2) + + dest_w_offs = (block_id_w * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)) + + dest_offs = (block_id_c[None, None] * stride_out_c + block_id_h[None, None] * stride_out_h + + dest_w_offs[None, :] * stride_out_w) + + fx = dest_w_offs * RR_W + sx = tl.floor(fx) + + new_col = block_id_h * RR_H + src_offsets = (block_id_c[None, None] * stride_in_c + new_col[None, None].to(tl.int32) * stride_in_h + + tl.clamp(sx, 0, src_cols - 1)[None, :].to(tl.int32) * stride_in_w) + src_val = tl.load(img_src_ptr + src_offsets) + dst_mask = dest_w_offs[None, :] < dst_cols + tl.store(img_dst_ptr + dest_offs, src_val, mask=dst_mask) + + +def nearest_resize_cpu(img_src, img_dst, dst_rows, dst_cols): + N, C, src_rows, src_cols = img_src.shape + # 2,4 64, 32 + #RR_H和RR_W分别为高和宽的缩放比例 + RR_H = src_rows / float(dst_rows) + RR_W = src_cols / float(dst_cols) + print("RR_H RR_W", RR_H, RR_W) + # 2, 2 + #根据output image的坐标值(i,j)计算input image的坐标值(sy, sx) + for i in range(dst_rows): #32 + for j in range(dst_cols): #16 + # fy = i * 2 = 0/1/..31 * 2 = 0/2/4...62 + fy = (i * RR_H) + sy = math.floor(fy) + # fx = j * 2 = 0/1/2..14 * 2 = 0/2/4...28 + fx = (j * RR_W) + sx = math.floor(fx) + src_val = img_src[0, :, np.clip(sy, 0, src_rows - 1), np.clip(sx, 0, src_cols - 1)] + # img_dst[0, :, i, j] 表示取批量中第 0 张图像、第 i 行第 j 列位置上的所有通道像素值 + img_dst[0, :, i, j] = src_val + + +def test_nearest_resize(): + n, c, h, w = 1, 4, 64, 64 + img_src = torch.randint(0, 255, size=(n, c, h, w)) + dst_rows = h // 2 + dst_cols = w // 2 + img_dst_cpu = torch.randint(0, 255, size=(n, c, dst_rows, dst_cols)) + nearest_resize_cpu(img_src, img_dst_cpu, dst_rows, dst_cols) + + # call triton kernel + img_src = img_src.npu() + RR_H = h / float(dst_rows) + RR_W = w / float(dst_cols) + img_dst_npu = torch.randint(0, 255, size=(n, c, dst_rows, dst_cols)).npu() + stride_in_h = h + stride_in_w = 1 + stride_in_c = h * w + stride_out_h = dst_rows + stride_out_w = 1 + stride_out_c = dst_rows * dst_cols + BLOCK_SIZE = 32 + # best performance + nearest_resize_kernel_col_tile[(4, 32, 1)](img_src, img_dst_npu, h, w, dst_rows, dst_cols, RR_H, RR_W, stride_in_h, + stride_in_w, stride_in_c, stride_out_h, stride_out_w, stride_out_c, + BLOCK_SIZE) + assert torch.equal(img_dst_cpu.cpu(), img_dst_npu.cpu()) diff --git a/third_party/ascend/unittest/pytest_ut/test_rint.py b/third_party/ascend/unittest/pytest_ut/test_rint.py new file mode 100644 index 0000000000..23b1bb3f13 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_rint.py @@ -0,0 +1,65 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice + +import time + +import torch +import torch_npu +import test_common + + +@triton.jit +def triton_rint(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + for loop1 in range(0, XBLOCK, XBLOCK_SUB): + x0 = offset + loop1 + tl.arange(0, XBLOCK_SUB) + xmask = x0 < xnumel + tmp0 = tl.load(in_ptr0 + x0, mask=xmask) + tmp1 = libdevice.rint(tmp0) + tl.store(out_ptr0 + x0, tmp1, mask=xmask) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 32, 2048, 64], +]) +def test_rint(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch.round(x) + y_cal = test_common.generate_tensor(x.shape, dtype).npu() + triton_rint[ncore, 1, 1](x, y_cal, x.numel(), xblock, xblock_sub, debug=True) + test_common.validate_cmp_with_expection(dtype, y_cal, y_ref, True) + + +@pytest.mark.parametrize('dtype', [ + 'float32', +]) +def test_rint_half(dtype): + x0 = torch.tensor([-2.5, -1.5, -0.5, 0, 0.5, 1.5, 2.5], dtype=eval('torch.' + dtype)).npu() + y_ref = torch.round(x0) + y_cal = test_common.generate_tensor(x0.shape, dtype).npu() + triton_rint[32, 1, 1](x0, y_cal, x0.numel(), 2048, 64, debug=True) + test_common.validate_cmp_with_expection(dtype, y_cal, y_ref, True) diff --git a/third_party/ascend/unittest/pytest_ut/test_rms_norm.py b/third_party/ascend/unittest/pytest_ut/test_rms_norm.py new file mode 100644 index 0000000000..95e1baf02e --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_rms_norm.py @@ -0,0 +1,142 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import triton +import triton.language as tl +import torch_npu + + +@triton.jit +def _rms_norm_fwd_fused( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + stride, # how much to increase the pointer when moving by 1 row + N, # number of columns in X + eps, # epsilon to avoid division by zero + BLOCK_SIZE: tl.constexpr, +): + # Map the program id to the row of X and Y it should compute. + row = tl.program_id(0) + Y += row * stride + X += row * stride + # Compute variance + _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) + _var += x * x + var = tl.sum(_var, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + # Normalize and apply linear transformation + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + mask = cols < N + w = tl.load(W + cols, mask=mask).to(tl.float32) + x = tl.load(X + cols, mask=mask, other=0.0).to(tl.float32) + x_hat = x * rstd + y = x_hat * w + # Write output + tl.store(Y + cols, y.to(tl.float16), mask=mask) + + +# have to change the block_size +@torch.inference_mode() +def rms_norm(x, weight, eps, out=None): + # allocate output, tl.store save y in tl.float16 + y = torch.empty_like(x, dtype=torch.float16) if out is None else out + # reshape input data into 2D tensor + x_arg = x.view(-1, x.shape[-1]) + M, N = x_arg.shape + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + if N > BLOCK_SIZE: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + # heuristics for number of warps + num_warps = min(max(BLOCK_SIZE // 256, 1), 8) + BLOCK_SIZE = 128 * 2 * 2 * 2 * 2 * 2 * 2 + num_warps = 8 + # enqueue kernel + kernel = _rms_norm_fwd_fused[(M, )](x_arg, y, weight, x_arg.stride(0), N, eps, BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps) + return y, kernel + + +def _rms_norm(shape, datatype): + x = torch.randn(shape[0], shape[1], dtype=datatype, device="npu") + weight = torch.randn(shape[1], dtype=datatype, device="npu") + y, kernel = rms_norm(x, weight, eps=1e-5) + eps1 = 1e-5 + if datatype == torch.bfloat16 or datatype == torch.float16: + x = x.to(torch.float32) + rms = torch.sqrt(x.pow(2).mean(-1, keepdim=True) + eps1) # 计算均方根 + x_norm = x / rms # 标准化 + y_ref = weight * x_norm + y_ref = y_ref.to(torch.float16) + torch.testing.assert_close(y_ref, y, rtol=1e-3, atol=1e-3) + + +def test_cases(): + _rms_norm((16, 1024), torch.float16) + _rms_norm((16, 1024), torch.float32) + _rms_norm((16, 1024), torch.bfloat16) + _rms_norm((128, 3), torch.bfloat16) + _rms_norm((128, 16), torch.bfloat16) + _rms_norm((128, 37), torch.bfloat16) + _rms_norm((128, 64), torch.bfloat16) + _rms_norm((128, 781), torch.bfloat16) + _rms_norm((128, 781), torch.bfloat16) + _rms_norm((16, 1024), torch.float16) + _rms_norm((16, 1024), torch.float32) + _rms_norm((16, 1024), torch.bfloat16) + + _rms_norm((128, 128), torch.float16) + _rms_norm((128, 128), torch.float32) + _rms_norm((128, 128), torch.bfloat16) + + _rms_norm((1, 128), torch.float16) + _rms_norm((1, 128), torch.float32) + _rms_norm((1, 128), torch.bfloat16) + + _rms_norm((65535, 128), torch.float16) + _rms_norm((65535, 128), torch.float32) + _rms_norm((65535, 128), torch.bfloat16) + + _rms_norm((128, 3), torch.float16) + _rms_norm((128, 3), torch.float32) + _rms_norm((128, 3), torch.bfloat16) + + _rms_norm((128, 16), torch.float16) + _rms_norm((128, 16), torch.float32) + _rms_norm((128, 16), torch.bfloat16) + + _rms_norm((128, 37), torch.float16) + _rms_norm((128, 37), torch.float32) + _rms_norm((128, 37), torch.bfloat16) + + _rms_norm((128, 64), torch.float16) + _rms_norm((128, 64), torch.float32) + _rms_norm((128, 64), torch.bfloat16) + + _rms_norm((128, 781), torch.float16) + _rms_norm((128, 781), torch.float32) + _rms_norm((128, 781), torch.bfloat16) diff --git a/third_party/ascend/unittest/pytest_ut/test_rotary_embedding.py b/third_party/ascend/unittest/pytest_ut/test_rotary_embedding.py new file mode 100644 index 0000000000..7a5f37c225 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_rotary_embedding.py @@ -0,0 +1,175 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +"""Rotary embedding kernel implemented by Triton. + +GPT-NeoX style +""" + +import torch +import torch_npu +import triton +import triton.language as tl + + +@triton.jit +def rotary_embedding_kernel( + state, # [num_tokens, head_num, head_dim] + cos, # [num_tokens, 1, head_dim // 2] + sin, # [num_tokens, 1, head_dim // 2] + stride_state_n, + stride_state_h, + stride_state_d, + stride_cos_n, + stride_cos_d, + # stride_sin_n, + # stride_sin_d, + num_tokens, + num_heads, + BLOCK_N: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, +): + token_index = tl.program_id(0) + token_range = token_index * BLOCK_N + tl.arange(0, BLOCK_N) + head_index = tl.program_id(1) + head_range = head_index * BLOCK_H + tl.arange(0, BLOCK_H) + + dim_range_x = tl.arange(0, BLOCK_D // 2) + dim_range_y = tl.arange(BLOCK_D // 2, BLOCK_D) + + state_x_offset = (token_range[:, None, None] * stride_state_n + head_range[None, :, None] * stride_state_h + + dim_range_x[None, None, :] * stride_state_d) + state_y_offset = (token_range[:, None, None] * stride_state_n + head_range[None, :, None] * stride_state_h + + dim_range_y[None, None, :] * stride_state_d) + + cos_sim_offset = (token_range[:, None, None] * stride_cos_n + dim_range_x[None, None, :] * stride_cos_d) + + state_x = tl.load( + state + state_x_offset, + mask=(token_range[:, None, None] < num_tokens) + & (head_range[None, :, None] < num_heads), + other=0.0, + ) + state_y = tl.load( + state + state_y_offset, + mask=(token_range[:, None, None] < num_tokens) + & (head_range[None, :, None] < num_heads), + other=0.0, + ) + + cos_loaded = tl.load( + cos + cos_sim_offset, + mask=token_range[:, None, None] < num_tokens, + other=0.0, + ) + sin_loaded = tl.load( + sin + cos_sim_offset, + mask=token_range[:, None, None] < num_tokens, + other=0.0, + ) + + out_x = state_x * cos_loaded - state_y * sin_loaded + out_y = state_x * sin_loaded + state_y * cos_loaded + + tl.store( + state + state_x_offset, + out_x, + mask=(token_range[:, None, None] < num_tokens) + & (head_range[None, :, None] < num_heads), + ) + tl.store( + state + state_y_offset, + out_y, + mask=(token_range[:, None, None] < num_tokens) + & (head_range[None, :, None] < num_heads), + ) + + +@torch.inference_mode() +def rotary_embedding(state, cos, sin): + num_tokens = state.shape[0] + num_heads = state.shape[1] + head_dim = state.shape[2] + + #BLOCK_N = 32 + BLOCK_N = 16 + BLOCK_H = 4 + grid = ( + triton.cdiv(num_tokens, BLOCK_N), + triton.cdiv(num_heads, BLOCK_H), + ) + if head_dim >= 128: + num_warps = 8 + else: + num_warps = 4 + + kernel = rotary_embedding_kernel[grid]( + state, + cos, + sin, + state.stride(0), + state.stride(1), + state.stride(2), + cos.stride(0), + cos.stride(2), + # sin.stride(0), + # sin.stride(2), + num_tokens, + num_heads, + BLOCK_N=BLOCK_N, + BLOCK_H=BLOCK_H, + BLOCK_D=head_dim, + num_warps=num_warps, + num_stages=1, + ) + return + + +def torch_rotary_embedding(state, cos, sin): + _, _, dim = state.shape + state_x = state[:, :, 0:dim // 2] + state_y = state[:, :, dim // 2:dim] + out_x = state_x * cos - state_y * sin + out_y = state_x * sin + state_y * cos + return torch.cat((out_x, out_y), dim=-1) + + +def rotary_emb(tokens, heads, headdim, dtype): + tokens_num = tokens + num_heads = heads + head_dim = headdim + max_positions = 1024 + + # torch.float16 has floating point problem in Triton 2.0.0 + # But it works fine in Triton 2.1.0 + state = torch.randn((tokens_num, num_heads, head_dim), dtype=dtype, device="npu") + cos_shape = (tokens_num, 1, head_dim // 2) + cos = -1.2 + 0.5 * torch.randn(cos_shape, dtype=dtype, device="npu") + sin = -2.0 + 0.5 * torch.randn(cos_shape, dtype=dtype, device="npu") + # forward pass + torch_result = torch_rotary_embedding(state, cos, sin) + rotary_embedding(state, cos, sin) + triton_result = state # state is modified in-place + torch.testing.assert_close(torch_result, triton_result, rtol=1e-3, atol=1e-3) + + +def test_cases(): + rotary_emb(256, 96, 128, torch.float16) + rotary_emb(256, 96, 128, torch.float32) diff --git a/third_party/ascend/unittest/pytest_ut/test_rotatry_gpt.py b/third_party/ascend/unittest/pytest_ut/test_rotatry_gpt.py new file mode 100644 index 0000000000..5a8bdac7f3 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_rotatry_gpt.py @@ -0,0 +1,186 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +Rotary embedding kernel implemented by Triton. +GPT-J style +""" + +import torch +import torch_npu +import triton +import triton.language as tl + + +@triton.jit +def rotary_embedding_kernel( + state, # [num_tokens, head_num, head_dim] + cos, # [num_tokens, 1, head_dim // 2] + sin, # [num_tokens, 1, head_dim // 2] + stride_state_n, + stride_state_h, + stride_state_d, + stride_cos_n, + stride_cos_d, + # stride_sin_n, + # stride_sin_d, + num_tokens, + num_heads, + BLOCK_N: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, +): + token_index = tl.program_id(0) + token_range = token_index * BLOCK_N + tl.arange(0, BLOCK_N) + head_index = tl.program_id(1) + head_range = head_index * BLOCK_H + tl.arange(0, BLOCK_H) + + dim_range = tl.arange(0, BLOCK_D // 2) + dim_range_x = dim_range * 2 + dim_range_y = dim_range * 2 + 1 + + # tl.device_print("dim x", dim_range_x) + # tl.device_print("dim y", dim_range_y) + + state_x_offset = (token_range[:, None, None] * stride_state_n + head_range[None, :, None] * stride_state_h + + dim_range_x[None, None, :] * stride_state_d) + + state_y_offset = (token_range[:, None, None] * stride_state_n + head_range[None, :, None] * stride_state_h + + dim_range_y[None, None, :] * stride_state_d) + + cos_sim_offset = (token_range[:, None, None] * stride_cos_n + dim_range[None, None, :] * stride_cos_d) + + state_x = tl.load( + state + state_x_offset, + mask=(token_range[:, None, None] < num_tokens) + & (head_range[None, :, None] < num_heads), + other=0.0, + ) + state_y = tl.load( + state + state_y_offset, + mask=(token_range[:, None, None] < num_tokens) + & (head_range[None, :, None] < num_heads), + other=0.0, + ) + + cos_loaded = tl.load( + cos + cos_sim_offset, + mask=token_range[:, None, None] < num_tokens, + other=0.0, + ) + sin_loaded = tl.load( + sin + cos_sim_offset, + mask=token_range[:, None, None] < num_tokens, + other=0.0, + ) + + out_x = state_x * cos_loaded - state_y * sin_loaded + out_y = state_x * sin_loaded + state_y * cos_loaded + + tl.store( + state + state_x_offset, + out_x, + mask=(token_range[:, None, None] < num_tokens) + & (head_range[None, :, None] < num_heads), + ) + tl.store( + state + state_y_offset, + out_y, + mask=(token_range[:, None, None] < num_tokens) + & (head_range[None, :, None] < num_heads), + ) + + +@torch.inference_mode() +def rotary_embedding(state, cos, sin): + num_tokens = state.shape[0] + num_heads = state.shape[1] + head_dim = state.shape[2] + + BLOCK_N = 8 + BLOCK_H = 4 + grid = ( + triton.cdiv(num_tokens, BLOCK_N), + triton.cdiv(num_heads, BLOCK_H), + ) + if head_dim >= 128: + num_warps = 8 + else: + num_warps = 4 + + kernel = rotary_embedding_kernel[grid]( + state, + cos, + sin, + state.stride(0), + state.stride(1), + state.stride(2), + cos.stride(0), + cos.stride(2), + # sin.stride(0), + # sin.stride(2), + num_tokens, + num_heads, + BLOCK_N=BLOCK_N, + BLOCK_H=BLOCK_H, + BLOCK_D=head_dim, + num_warps=num_warps, + num_stages=1, + ) + # print(kernel.asm['ttir']) + return + + +def torch_rotary_embedding(state, cos, sin): + _, _, dim = state.shape + state_x = state[:, :, 0:dim:2] + state_y = state[:, :, 1:dim:2] + out_x = state_x * cos - state_y * sin + out_y = state_x * sin + state_y * cos + out = torch.empty_like(state).npu() + out[:, :, 0:dim:2] = out_x + out[:, :, 1:dim:2] = out_y + return out + + +def _rotary_emb(dtype): + tokens_num = 256 + num_heads = 96 + head_dim = 128 + max_positions = 1024 + + # torch.float16 has floating point problem in Triton 2.0.0 + # But it works fine in Triton 2.1.0 + state = torch.randn((tokens_num, num_heads, head_dim), dtype=dtype, device="npu") + cos_shape = (tokens_num, 1, head_dim // 2) + cos = -1.2 + 0.5 * torch.randn(cos_shape, dtype=dtype, device="npu") + sin = -2.0 + 0.5 * torch.randn(cos_shape, dtype=dtype, device="npu") + # forward pass + torch_result = torch_rotary_embedding(state, cos, sin) + rotary_embedding(state, cos, sin) + triton_result = state # state is modified in-place + # print(torch_result[1][0]) + # print(triton_result[1][0]) + # Note: This test is not accurate enough. + assert torch.allclose(torch_result, triton_result, atol=1e-2, rtol=1e-7) + + +def test_rotary_emb(): + _rotary_emb(torch.float16) + _rotary_emb(torch.float32) diff --git a/third_party/ascend/unittest/pytest_ut/test_rotaty_embedding_gpt.py b/third_party/ascend/unittest/pytest_ut/test_rotaty_embedding_gpt.py new file mode 100644 index 0000000000..5f73042f78 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_rotaty_embedding_gpt.py @@ -0,0 +1,174 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +"""Rotary embedding kernel implemented by Triton. + +GPT-NeoX style +""" + +import torch +import torch_npu +import triton +import triton.language as tl + + +@triton.jit +def rotary_embedding_kernel( + state, # [num_tokens, head_num, head_dim] + cos, # [num_tokens, 1, head_dim // 2] + sin, # [num_tokens, 1, head_dim // 2] + stride_state_n, + stride_state_h, + stride_state_d, + stride_cos_n, + stride_cos_d, + # stride_sin_n, + # stride_sin_d, + num_tokens, + num_heads, + BLOCK_N: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, +): + token_index = tl.program_id(0) + token_range = token_index * BLOCK_N + tl.arange(0, BLOCK_N) + head_index = tl.program_id(1) + head_range = head_index * BLOCK_H + tl.arange(0, BLOCK_H) + + dim_range_x = tl.arange(0, BLOCK_D // 2) + dim_range_y = tl.arange(BLOCK_D // 2, BLOCK_D) + + state_x_offset = (token_range[:, None, None] * stride_state_n + head_range[None, :, None] * stride_state_h + + dim_range_x[None, None, :] * stride_state_d) + state_y_offset = (token_range[:, None, None] * stride_state_n + head_range[None, :, None] * stride_state_h + + dim_range_y[None, None, :] * stride_state_d) + + cos_sim_offset = (token_range[:, None, None] * stride_cos_n + dim_range_x[None, None, :] * stride_cos_d) + + state_x = tl.load( + state + state_x_offset, + mask=(token_range[:, None, None] < num_tokens) + & (head_range[None, :, None] < num_heads), + other=0.0, + ) + state_y = tl.load( + state + state_y_offset, + mask=(token_range[:, None, None] < num_tokens) + & (head_range[None, :, None] < num_heads), + other=0.0, + ) + + cos_loaded = tl.load( + cos + cos_sim_offset, + mask=token_range[:, None, None] < num_tokens, + other=0.0, + ) + sin_loaded = tl.load( + sin + cos_sim_offset, + mask=token_range[:, None, None] < num_tokens, + other=0.0, + ) + + out_x = state_x * cos_loaded - state_y * sin_loaded + out_y = state_x * sin_loaded + state_y * cos_loaded + + tl.store( + state + state_x_offset, + out_x, + mask=(token_range[:, None, None] < num_tokens) + & (head_range[None, :, None] < num_heads), + ) + tl.store( + state + state_y_offset, + out_y, + mask=(token_range[:, None, None] < num_tokens) + & (head_range[None, :, None] < num_heads), + ) + + +@torch.inference_mode() +def rotary_embedding(state, cos, sin): + num_tokens = state.shape[0] + num_heads = state.shape[1] + head_dim = state.shape[2] + + BLOCK_N = 16 + BLOCK_H = 4 + grid = ( + triton.cdiv(num_tokens, BLOCK_N), + triton.cdiv(num_heads, BLOCK_H), + ) + if head_dim >= 128: + num_warps = 8 + else: + num_warps = 4 + + kernel = rotary_embedding_kernel[grid]( + state, + cos, + sin, + state.stride(0), + state.stride(1), + state.stride(2), + cos.stride(0), + cos.stride(2), + # sin.stride(0), + # sin.stride(2), + num_tokens, + num_heads, + BLOCK_N=BLOCK_N, + BLOCK_H=BLOCK_H, + BLOCK_D=head_dim, + num_warps=num_warps, + num_stages=1, + ) + return + + +def torch_rotary_embedding(state, cos, sin): + _, _, dim = state.shape + state_x = state[:, :, 0:dim // 2] + state_y = state[:, :, dim // 2:dim] + out_x = state_x * cos - state_y * sin + out_y = state_x * sin + state_y * cos + return torch.cat((out_x, out_y), dim=-1) + + +def rotary_emb(tokens, heads, headdim, dtype): + tokens_num = tokens + num_heads = heads + head_dim = headdim + #max_positions = 1024 + + # torch.float16 has floating point problem in Triton 2.0.0 + # But it works fine in Triton 2.1.0 + state = torch.randn((tokens_num, num_heads, head_dim), dtype=dtype, device="npu") + cos_shape = (tokens_num, 1, head_dim // 2) + cos = -1.2 + 0.5 * torch.randn(cos_shape, dtype=dtype, device="npu") + sin = -2.0 + 0.5 * torch.randn(cos_shape, dtype=dtype, device="npu") + # forward pass + torch_result = torch_rotary_embedding(state, cos, sin) + rotary_embedding(state, cos, sin) + triton_result = state # state is modified in-place + torch.testing.assert_close(torch_result, triton_result, rtol=1e-3, atol=1e-3) + + +def test_cases(): + rotary_emb(256, 96, 128, torch.float16) + rotary_emb(256, 96, 128, torch.float32) diff --git a/third_party/ascend/unittest/pytest_ut/test_rshift.py b/third_party/ascend/unittest/pytest_ut/test_rshift.py new file mode 100644 index 0000000000..f24676572b --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_rshift.py @@ -0,0 +1,99 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import time +import test_common +import torch +import torch_npu + + +def standard_unary(x0, dtype): + res = x0 >> 2 + return res + + +def standard_binary(x0, y0, dtype): + res = x0 + y0 + return res + + +@triton.jit +def triton_elementwise_unary(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + tmp = tl.cast(2, tl.int8) + ret = x >> tmp + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +@triton.jit +def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + y = tl.load(in_ptr1 + idx_block, mask=idx_block < N) + ret = x + y + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +types = [ + # (torch.float32, 'float32'), + # (torch.float16, 'float16'), + # (torch.bfloat16, 'bfloat16'), + (torch.int8, 'int8'), + # (torch.int16, 'int16'), + # (torch.int32, 'int32'), + # (torch.int64, 'int64'), +] + +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('N,NUMEL', shapes) +def test_elementwsie_common(dtype, sigtype, N, NUMEL): + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == "int64": + N = map_for_64_t[N] if N in map_for_64_t else N + + print(f"elementwise : ({N},) {dtype} {sigtype}") + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standard_unary(x0, dtype) + x0 = x0.npu() + # print(ans) + + out = torch.zeros((N, ), dtype=dtype).npu() + triton_elementwise_unary[1, 1, 1](x0, out, N=N, NUMEL=NUMEL, debug=True) + # print(out) + + test_common.validate_cmp(sigtype, out, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_rsqrt.py b/third_party/ascend/unittest/pytest_ut/test_rsqrt.py new file mode 100644 index 0000000000..f6be226a50 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_rsqrt.py @@ -0,0 +1,63 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import numpy as np +import pytest +import test_common + + +def numpy_rsqrt(x0, x1): + res = x0 + 1.0 / (np.sqrt(x1)) + return res + + +@triton.jit +def triton_rsqrt(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + xindex = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = xindex < xnumel + x0 = xindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.load(in_ptr1 + (x0), xmask) + tmp2 = tmp0 + tl.rsqrt(tmp1) + tl.store(out_ptr0 + (xindex), tmp2, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_rsqrt(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = np.abs(np.random.randn(shape[0], shape[1], shape[2])).astype(eval('np.' + dtype)) + x1 = np.abs(np.random.randn(shape[0], shape[1], shape[2])).astype(eval('np.' + dtype)) + x0_npu = torch.tensor(x0).npu() + x1_npu = torch.tensor(x1).npu() + # numpy结果 + numpy_res = numpy_rsqrt(x0, x1) + # triton结果 + triton_res = test_common.generate_tensor(shape, dtype).npu() + triton_rsqrt[ncore, 1, 1](x0_npu, x1_npu, triton_res, x0_npu.numel(), xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, triton_res, torch.tensor(numpy_res).npu()) diff --git a/third_party/ascend/unittest/pytest_ut/test_runtime_utils.py b/third_party/ascend/unittest/pytest_ut/test_runtime_utils.py new file mode 100644 index 0000000000..6a70866cfe --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_runtime_utils.py @@ -0,0 +1,33 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +import logging +import os +from triton.backends.ascend import utils + + +def test_get_logger(): + logger = utils.get_logger("test_utils", "INFO") + assert logger.level == logging.INFO + + +def test_get_ascend_arch_from_env(): + os.environ["TRITON_ASCEND_ARCH"] = "Ascend910_9599" + result = utils.get_ascend_arch_from_env() + assert result == "Ascend910_9599" diff --git a/third_party/ascend/unittest/pytest_ut/test_scalarPointer.py b/third_party/ascend/unittest/pytest_ut/test_scalarPointer.py new file mode 100644 index 0000000000..bbe7f65f74 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_scalarPointer.py @@ -0,0 +1,178 @@ +import pytest +import os +import sys +import math + +import pytest +import os +import sys +import math + +current_file = os.path.abspath(__file__) +vllm_root_path = os.path.dirname(os.path.dirname(os.path.dirname(current_file))) +sys.path.insert(0, vllm_root_path) + +#from vllm.model_executor.layers.mamba.ops.ssd_state_passing import _state_passing_fwd_kernel +from einops import rearrange +import torch +import torch_npu +import triton +import triton.language as tl + + +@triton.autotune( + configs=[ + triton.Config({'BLOCK_SIZE': 64}), + ], + key=['dim'], +) +@triton.jit +def _state_passing_fwd_kernel( + # Pointers to matrices + states_ptr, + out_ptr, + final_states_ptr, + dA_cs_ptr, + initstates_ptr, + seq_idx_ptr, + # Matrix dimensions + dim, + nchunks, + seqlen, + chunk_size, + # Strides + stride_states_batch, + stride_states_chunk, + stride_states_head, + stride_states_dim, + stride_out_batch, + stride_out_chunk, + stride_out_head, + stride_out_dim, + stride_final_states_batch, + stride_final_states_head, + stride_final_states_dim, + stride_dA_cs_batch, + stride_dA_cs_chunk, + stride_dA_cs_head, + stride_initstates_batch, + stride_initstates_head, + stride_initstates_dim, + stride_seq_idx_batch, + stride_seq_idx_seqlen, + # Meta-parameters + HAS_INITSTATES: tl.constexpr, + HAS_SEQ_IDX: tl.constexpr, + IS_CONT_BATCHED: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + pid_m = tl.program_id(axis=0) + states_ptr += pid_b * stride_states_batch + pid_h * stride_states_head + dA_cs_ptr += pid_b * stride_dA_cs_batch + pid_h * stride_dA_cs_head + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + final_states_ptr += pid_b * stride_final_states_batch + pid_h * stride_final_states_head + + offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + states_ptrs = states_ptr + offs_m * stride_states_dim + out_ptrs = out_ptr + offs_m * stride_out_dim + final_states_ptrs = final_states_ptr + offs_m * stride_final_states_dim + + # - states will be the past state of the sequence that continues on the current check + if not HAS_INITSTATES: + states = tl.zeros((BLOCK_SIZE, ), dtype=tl.float32) + + tl.store(out_ptrs, states, mask=offs_m < dim) + out_ptrs += stride_out_chunk + for c in range(nchunks): + new_states = tl.load(states_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + #dA_cs = tl.load(dA_cs_ptr + tl.arange(0,1)).to(tl.float32) + dA_cs = tl.load(dA_cs_ptr).to(tl.float32) + scale = tl.exp(dA_cs) + states = scale * states + new_states + if c < nchunks - 1: + tl.store(out_ptrs, states, mask=offs_m < dim) + else: + tl.store(final_states_ptrs, states, mask=offs_m < dim) + states_ptrs += stride_states_chunk + dA_cs_ptr += stride_dA_cs_chunk + out_ptrs += stride_out_chunk + + +def _state_passing_fwd( + states, + dA_chunk_cumsum, + initial_states=None, + seq_idx=None, + chunk_size=None, + out_dtype=None, + is_cont_batched=False, +): + batch, nchunks, nheads, dim = states.shape + assert dA_chunk_cumsum.shape == (batch, nheads, nchunks) + if initial_states is not None: + if is_cont_batched: + # - if cu_seqlens is provided, then the initial states + # are used for continuous batching. In which case we + # require seq_idx to be provided + assert seq_idx is not None, "" + else: + # - this is the regular batching case, where initial + # states are used are for each example of the batch. + assert initial_states.shape == (batch, nheads, dim) + + if seq_idx is not None: + assert chunk_size is not None + seqlen = seq_idx.shape[-1] + assert seq_idx.shape == (batch, seqlen) + out_dtype = states.dtype if out_dtype is None else out_dtype + out = torch.empty((batch, nchunks, nheads, dim), device=states.device, dtype=out_dtype) + final_states = torch.empty((batch, nheads, dim), device=states.device, dtype=torch.float32) + grid = lambda META: (triton.cdiv(dim, META['BLOCK_SIZE']), batch, nheads) + # with torch.cuda.device(states.device.index): + _state_passing_fwd_kernel[grid]( + states, + out, + final_states, + dA_chunk_cumsum, + initial_states, + seq_idx, + dim, + nchunks, + seqlen if seq_idx is not None else 0, + chunk_size if seq_idx is not None else 0, + states.stride(0), + states.stride(1), + states.stride(2), + states.stride(3), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + final_states.stride(0), + final_states.stride(1), + final_states.stride(2), + dA_chunk_cumsum.stride(0), + dA_chunk_cumsum.stride(2), + dA_chunk_cumsum.stride(1), + *((initial_states.stride(0), initial_states.stride(1), + initial_states.stride(2)) if initial_states is not None else (0, 0, 0)), + *((seq_idx.stride(0), seq_idx.stride(1)) if seq_idx is not None else (0, 0)), + HAS_INITSTATES=initial_states is not None, + HAS_SEQ_IDX=seq_idx is not None, + IS_CONT_BATCHED=is_cont_batched, + ) + return out, final_states + + +@pytest.mark.perf(repeat=17) +def test_state_passing_fwd(): + states = torch.randn((1, 7, 3, 5, 5), dtype=torch.float32).npu() + dA_cumsum = torch.randn((1, 3, 7, 17), dtype=torch.float32).npu() + initial_states = None + cu_seqlens = None + states, final_states = _state_passing_fwd( + rearrange(states, "... p n -> ... (p n)"), dA_cumsum[:, :, :, -1], + initial_states=rearrange(initial_states, "... p n -> ... (p n)") if initial_states is not None else None, + seq_idx=None, chunk_size=17, out_dtype=torch.float32, is_cont_batched=cu_seqlens is not None) diff --git a/third_party/ascend/unittest/pytest_ut/test_scalar_calc.py b/third_party/ascend/unittest/pytest_ut/test_scalar_calc.py new file mode 100644 index 0000000000..53517eb4eb --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_scalar_calc.py @@ -0,0 +1,694 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice + +import pytest +import test_common + + +### add +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_add_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tmp0 + 2.0 + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = y + 2.0 + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### sub +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_sub_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tmp0 - 2.0 + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = y - 2.0 + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### mul +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_mul_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tmp0 * 2.0 + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = y * 2.0 + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### div +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_div_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tmp0 / 2.0 + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = y / 2.0 + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### remf +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_remf_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tmp0 % 2.0 + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = y - 2.0 * torch.div(y, 2.0, rounding_mode="trunc") + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### negf +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_negf_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = -tmp0 + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = -y + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### cmpf +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_cmpf_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = (tmp0 > 0.5).to(tmp0.dtype) + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = (y > 0.5).to(y.dtype) + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### ceil +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_ceil_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tl.math.ceil(tmp0) + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = torch.ceil(y) + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### floor +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_floor_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tl.math.floor(tmp0) + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = torch.floor(y) + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### maximum(propagate_nan == tl.PropagateNan.ALL) +# setting propagate_nan=tl.PropagateNan.ALL to generate arith::MaximumFOp +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_maximum_nanall_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + tl.static_assert(N > 1) + tmp0 = tl.load(in_ptr0 + 0) + tmp1 = tl.load(in_ptr0 + 1) + tmp1 = tl.maximum(tmp0, tmp1, propagate_nan=tl.PropagateNan.ALL) + tl.store(out_ptr0 + 0, tmp1) + + def torch_func(x0): + y0 = x0[0] + y1 = x0[1] + y = torch.maximum(y0, y1) + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### maximum(propagate_nan == tl.PropagateNan.NONE) +# setting propagate_nan=tl.PropagateNan.NONE to generate arith::MaxNumFOp +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_maximum_nannone_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + tl.static_assert(N > 1) + tmp0 = tl.load(in_ptr0 + 0) + tmp1 = tl.load(in_ptr0 + 1) + tmp1 = tl.maximum(tmp0, tmp1, propagate_nan=tl.PropagateNan.ALL) + tl.store(out_ptr0 + 0, tmp1) + + def torch_func(x0): + y0 = x0[0] + y1 = x0[1] + y = torch.fmax(y0, y1) + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### minimum(propagate_nan == tl.PropagateNan.ALL) +# setting propagate_nan=tl.PropagateNan.ALL to generate arith::MinimumFOp +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_minimum_nanall_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + tl.static_assert(N > 1) + tmp0 = tl.load(in_ptr0 + 0) + tmp1 = tl.load(in_ptr0 + 1) + tmp1 = tl.minimum(tmp0, tmp1, propagate_nan=tl.PropagateNan.ALL) + tl.store(out_ptr0 + 0, tmp1) + + def torch_func(x0): + y0 = x0[0] + y1 = x0[1] + y = torch.minimum(y0, y1) + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### minimum(propagate_nan == tl.PropagateNan.NONE) +# setting propagate_nan=tl.PropagateNan.NONE to generate arith::MinNumFOp +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_minimum_nannone_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + tl.static_assert(N > 1) + tmp0 = tl.load(in_ptr0 + 0) + tmp1 = tl.load(in_ptr0 + 1) + tmp1 = tl.minimum(tmp0, tmp1, propagate_nan=tl.PropagateNan.NONE) + tl.store(out_ptr0 + 0, tmp1) + + def torch_func(x0): + y0 = x0[0] + y1 = x0[1] + y = torch.fmin(y0, y1) + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### extf +@pytest.mark.parametrize('param_list', [['float16', 'float32', 16]]) +def test_scalar_extf_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tmp0.to(tl.float32) + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = y.to(torch.float32) + return torch.tensor(y) + + src_dtype, dst_dtype, N = param_list + x0 = test_common.generate_tensor((N, ), src_dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dst_dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dst_dtype, y_cal[0], y_ref) + + +### truncf +@pytest.mark.parametrize('param_list', [['float32', 'float16', 16]]) +def test_scalar_truncf_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tmp0.to(tl.float16) + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = y.to(torch.float16) + return torch.tensor(y) + + src_dtype, dst_dtype, N = param_list + x0 = test_common.generate_tensor((N, ), src_dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dst_dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dst_dtype, y_cal[0], y_ref) + + +### exp +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_exp_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tl.math.exp(tmp0) + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = torch.exp(y) + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### exp2 +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_exp_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tl.math.exp2(tmp0) + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = torch.exp2(y) + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### log +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_log_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp0 = tl.abs(tmp0) + tmp1 = tl.log(tmp0) + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = torch.abs(y) + y = torch.log(y) + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### log2 +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_log2_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp0 = tl.abs(tmp0) + tmp1 = tl.log2(tmp0) + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = torch.abs(y) + y = torch.log2(y) + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### sin +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_sin_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tl.sin(tmp0) + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = torch.sin(y) + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### cos +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_cos_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tl.cos(tmp0) + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = torch.cos(y) + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### abs +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_abs_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tl.abs(tmp0) + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = torch.abs(y) + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### erf +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_erf_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tl.erf(tmp0) + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = torch.erf(y) + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### sqrt +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_sqrt_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp0 = tl.abs(tmp0) + tmp1 = tl.math.sqrt(tmp0) + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = torch.abs(y) + y = torch.sqrt(y) + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### rsqrt +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_rsqrt_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp0 = tl.abs(tmp0) + tmp1 = tl.math.rsqrt(tmp0) + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = torch.abs(y) + y = torch.rsqrt(y) + return y.clone().detach() + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### tanh +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_tanh_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + idx = 0 + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = libdevice.tanh(tmp0) + tl.store(out_ptr0 + idx, tmp1) + + def torch_func(x0): + y = x0[0] + y = torch.tanh(y) + return y.clone().detach() + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) + + +### sum +@pytest.mark.parametrize('param_list', [['float32', 16]]) +def test_scalar_sum_calc(param_list): + + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): + tmp0 = tl.load(in_ptr0 + tl.arange(0, N)) + tmp1 = tl.sum(tmp0, 0) + tl.store(out_ptr0 + 0, tmp1) + + def torch_func(x0): + y = torch.sum(x0, 0) + return torch.tensor(y) + + dtype, N = param_list + x0 = test_common.generate_tensor((N, ), dtype).npu() + y_ref = torch_func(x0) + y_cal = test_common.generate_tensor((1, ), dtype).npu() + triton_kernel[1, 1, 1](y_cal, x0, N=N) + test_common.validate_cmp(dtype, y_cal[0], y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_scope.py b/third_party/ascend/unittest/pytest_ut/test_scope.py new file mode 100755 index 0000000000..9973348bd1 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_scope.py @@ -0,0 +1,159 @@ +#!/usr/bin/env python3 +import os + +os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" + +import pytest +import triton +import triton.language as tl +import triton.language.extra.cann.extension as al +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir + + +class Options: + num_warps = 4 + num_stages = 3 + num_ctas = 1 + cluster_dims = (1, 1, 1) + enable_fp_fusion = True + debug = False + + +def compile_kernel(kernel, signature, constants): + """Helper to compile a kernel to MLIR.""" + src = ASTSource(kernel, signature, constants) + context = ir.context() + ir.load_dialects(context) + ascend_ir.load_dialects(context) + module = ast_to_ttir(kernel, src, context, Options(), {}, {}) + return str(module) + + +# ============== Kernel definitions ============== + + +@triton.jit +def kernel_nested_scope(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): + """Test nested scopes.""" + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + with al.scope(core_mode="vector"): + with al.scope(core_mode="vector"): + with al.scope(core_mode="cube"): + x = tl.load(x_ptr + i, mask=i < n) + y = tl.load(y_ptr + i, mask=i < n) + result = x + y + tl.store(out_ptr + i, result, mask=i < n) + + +@triton.jit +def kernel_scope_escape(x_ptr, out_ptr, n, BLOCK: tl.constexpr): + """Test variable defined inside scope, used outside.""" + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + with al.scope(core_mode="vector"): + x = tl.load(x_ptr + i, mask=i < n) + # Use x outside of the scope + a = x + 1.0 + tl.store(out_ptr + i, a, mask=i < n) + + +@triton.jit +def kernel_scope_cube(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): + """Test cube core mode.""" + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + with al.scope(core_mode="cube"): + x = tl.load(x_ptr + i, mask=i < n) + y = tl.load(y_ptr + i, mask=i < n) + result = x + y + tl.store(out_ptr + i, result, mask=i < n) + + +@triton.jit +def kernel_scope_vector(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): + """Test vector core mode.""" + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + with al.scope(core_mode="vector"): + x = tl.load(x_ptr + i, mask=i < n) + y = tl.load(y_ptr + i, mask=i < n) + result = x + y + tl.store(out_ptr + i, result, mask=i < n) + + +@triton.jit +def kernel_scope_disable_auto_sync(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): + """Test disable auto sync.""" + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + with al.scope(core_mode="vector", disable_auto_sync=True): + x = tl.load(x_ptr + i, mask=i < n) + y = tl.load(y_ptr + i, mask=i < n) + result = x + y + tl.store(out_ptr + i, result, mask=i < n) + + +# ============== Pytest tests ============== + + +def test_nested_scope(): + """Test nested scopes compile successfully.""" + mlir = compile_kernel(kernel_nested_scope, {"x_ptr": "*fp32", "y_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, + {"BLOCK": 256}) + assert "scope.scope" in mlir + assert len(mlir) > 0 + + +def test_scope_escape(): + """Test variable escaping from scope.""" + mlir = compile_kernel(kernel_scope_escape, {"x_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, {"BLOCK": 256}) + assert "scope.scope" in mlir + assert len(mlir) > 0 + + +def test_scope_cube_mode(): + """Test cube core mode generates correct attributes.""" + mlir = compile_kernel(kernel_scope_cube, {"x_ptr": "*fp32", "y_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, + {"BLOCK": 256}) + assert "scope.scope" in mlir + # Check for cube core type attribute + assert "hivm.tcore_type" in mlir or "CUBE" in mlir.upper() + + +def test_scope_vector_mode(): + """Test vector core mode generates correct attributes.""" + mlir = compile_kernel(kernel_scope_vector, {"x_ptr": "*fp32", "y_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, + {"BLOCK": 256}) + assert "scope.scope" in mlir + # Check for vector core type attribute + assert "hivm.tcore_type" in mlir or "VECTOR" in mlir.upper() + + +def test_scope_disable_auto_sync(): + """Test disable auto sync generates correct attributes.""" + mlir = compile_kernel( + kernel_scope_disable_auto_sync, + {"x_ptr": "*fp32", "y_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, + {"BLOCK": 256}, + ) + assert "scope.scope" in mlir + # Check for disable auto sync attribute + assert "hivm.disable_auto_sync" in mlir + + +# ============== Main for manual testing ============== + +if __name__ == "__main__": + print("=" * 60) + print("Test 1: Nested Scopes") + print("=" * 60) + mlir = compile_kernel(kernel_nested_scope, {"x_ptr": "*fp32", "y_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, + {"BLOCK": 256}) + print(f"✅ Generated MLIR ({len(mlir)} chars):\n") + print(mlir) + + print("\n" + "=" * 60) + print("Test 2: Scope Escape") + print("=" * 60) + mlir = compile_kernel(kernel_scope_escape, {"x_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32"}, {"BLOCK": 256}) + print(f"✅ Generated MLIR ({len(mlir)} chars):\n") + print(mlir) diff --git a/third_party/ascend/unittest/pytest_ut/test_scope_buffer.py b/third_party/ascend/unittest/pytest_ut/test_scope_buffer.py new file mode 100644 index 0000000000..a5c4d9f094 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_scope_buffer.py @@ -0,0 +1,97 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import triton.extension.buffer.language as bl +import triton.language.extra.cann.extension as al +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +from triton._C.libtriton import ir, buffer_ir +from triton._C.libtriton.ascend import ir as ascend_ir + + +class Options: + num_warps = 4 + num_stages = 3 + num_ctas = 1 + cluster_dims = (1, 1, 1) + enable_fp_fusion = True + debug = False + + +def compile_kernel(kernel, signature, constants): + src = ASTSource(kernel, signature, constants) + context = ir.context() + ir.load_dialects(context) + buffer_ir.load_dialects(context) + ascend_ir.load_dialects(context) + module = ast_to_ttir( + kernel, + src, + context, + Options(), + {"create_address_space": al.semantic.create_address_space}, + {}, + ) + return str(module) + + +@triton.jit +def helper_use_q(q, XBLOCK: tl.constexpr): + tmp = q + 1.0 + tl.device_assert(tmp.shape[0] == XBLOCK if isinstance(tmp.shape[0], tl.constexpr) else True) + + +@triton.jit +def helper_use_buffer(q_l1, XBLOCK: tl.constexpr): + q = bl.to_tensor(q_l1) + helper_use_q(q, XBLOCK) + + +@triton.jit +def kernel_buffer_arg(XBLOCK: tl.constexpr): + q_l1 = bl.alloc(tl.float32, [XBLOCK], al.ascend_address_space.L1) + helper_use_buffer(q_l1, XBLOCK) + + +@triton.jit +def kernel_scope_for_if_call_buffer(x_ptr, XBLOCK: tl.constexpr): + q_l1 = bl.alloc(tl.float32, [XBLOCK], al.ascend_address_space.L1) + offsets = tl.arange(0, XBLOCK) + with al.scope(core_mode="vector"): + for skv_loop_idx in range(2): + if skv_loop_idx == 0: + q = tl.load(x_ptr + offsets, mask=offsets < XBLOCK) + bl.to_buffer(tensor=q, bind_buffer=q_l1) + else: + q = bl.to_tensor(q_l1) + helper_use_buffer(q_l1, XBLOCK) + + +def test_buffer_argument_to_jit_function(): + mlir = compile_kernel(kernel_buffer_arg, {}, {"XBLOCK": 16}) + assert len(mlir) > 0 + + +def test_buffer_helper_call_inside_scope_loop_if(): + mlir = compile_kernel(kernel_scope_for_if_call_buffer, {"x_ptr": "*fp32"}, {"XBLOCK": 16}) + assert "scf.for" in mlir + assert "scope.scope" in mlir diff --git a/third_party/ascend/unittest/pytest_ut/test_select_analysis.py b/third_party/ascend/unittest/pytest_ut/test_select_analysis.py new file mode 100644 index 0000000000..57d8c3e24b --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_select_analysis.py @@ -0,0 +1,438 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +@triton.jit +def kernel_cal_select_mask_bool( + Output_ptr, + Indices_ptr, + numel: tl.constexpr, + BLOCK: tl.constexpr, +): + offs = tl.arange(0, BLOCK) + index = tl.load(Indices_ptr) + + true_tensor = tl.arange(0, BLOCK) < numel + false_tensor = tl.arange(0, BLOCK) >= numel + mask = offs < index + res = tl.where(mask, true_tensor, false_tensor) + tl.store(Output_ptr + offs, res) + + +@triton.jit +def kernel_cal_select_mask_one_dim_static_offset(QK_ptr, Other_ptr, Output_ptr, BLOCK: tl.constexpr, + start_zero: tl.constexpr): + offs = tl.arange(0, BLOCK) + + qk_ub = tl.load(QK_ptr + offs) + other = tl.load(Other_ptr + offs) + + if start_zero: + mask = offs < 1 + else: + mask = offs >= 1 + + res = tl.where(mask, qk_ub, other) + tl.store(Output_ptr + offs, res) + + +@triton.jit +def kernel_cal_select_mask_one_dim_dyn_offset(QK_ptr, Other_ptr, Output_ptr, Indices_ptr, BLOCK: tl.constexpr, + start_zero: tl.constexpr): + offs = tl.arange(0, BLOCK) + index = tl.load(Indices_ptr) + + qk_ub = tl.load(QK_ptr + offs) + other = tl.load(Other_ptr + offs) + + if start_zero: + mask = offs < index + else: + mask = offs >= index + + res = tl.where(mask, qk_ub, other) + tl.store(Output_ptr + offs, res) + + +@triton.jit +def kernel_cal_select_mask_one_dim_two_dyn_offset( + QK_ptr, + Other_ptr, + Output_ptr, + Indices_ptr, + BLOCK: tl.constexpr, +): + offs = tl.arange(0, BLOCK) + minIndex = tl.load(Indices_ptr) + maxIndex = tl.load(Indices_ptr + 1) + + qk_ub = tl.load(QK_ptr + offs) + other = tl.load(Other_ptr + offs) + + mask = (offs < maxIndex) and (offs >= minIndex) + res = tl.where(mask, qk_ub, other) + tl.store(Output_ptr + offs, res) + + +@triton.jit +def kernel_cal_select_mask_two_dim_static_offset( + QK_ptr, + Other_ptr, + Output_ptr, + BLOCK: tl.constexpr, + start_zero: tl.constexpr, +): + rows = tl.arange(0, BLOCK) + cols = tl.arange(0, BLOCK) + offs = rows[:, None] * BLOCK + cols[None, :] + + qk_ub = tl.load(QK_ptr + offs) + other = tl.load(Other_ptr + offs) + if start_zero: + mask_rows = rows < 1 + mask_cols = cols < 1 + else: + mask_rows = rows >= 1 + mask_cols = cols >= 1 + + res = tl.where(mask_rows[:, None] & mask_cols[None, :], qk_ub, other) + tl.store(Output_ptr + offs, res) + + +@triton.jit +def kernel_cal_select_mask_two_dim_one_dyn_offset( + QK_ptr, + Other_ptr, + Output_ptr, + Indices_ptr, + BLOCK: tl.constexpr, + start_zero: tl.constexpr, +): + rows = tl.arange(0, BLOCK) + cols = tl.arange(0, BLOCK) + offs = rows[:, None] * BLOCK + cols[None, :] + minIndex = tl.load(Indices_ptr) + maxIndex = tl.load(Indices_ptr + 1) + + qk_ub = tl.load(QK_ptr + offs) + other = tl.load(Other_ptr + offs) + if start_zero: + mask_rows = rows < 1 + else: + mask_rows = rows >= 1 + + mask_cols = (cols < maxIndex) & (cols >= minIndex) + + res = tl.where(mask_rows[:, None] & mask_cols[None, :], qk_ub, other) + tl.store(Output_ptr + offs, res) + + +@triton.jit +def kernel_cal_select_mask_two_dim_all_dyn_offset( + QK_ptr, + Other_ptr, + Output_ptr, + Indices_ptr, + BLOCK: tl.constexpr, +): + rows = tl.arange(0, BLOCK) + cols = tl.arange(0, BLOCK) + offs = rows[:, None] * BLOCK + cols[None, :] + min_index1 = tl.load(Indices_ptr) + max_index1 = tl.load(Indices_ptr + 1) + min_index2 = tl.load(Indices_ptr + 2) + max_index2 = tl.load(Indices_ptr + 3) + + qk_ub = tl.load(QK_ptr + offs) + other = tl.load(Other_ptr + offs) + + mask_rows = (rows < max_index1) & (rows >= min_index1) + mask_cols = (cols < max_index2) & (cols >= min_index2) + + res = tl.where(mask_rows[:, None] & mask_cols[None, :], qk_ub, other) + tl.store(Output_ptr + offs, res) + + +@triton.jit +def kernel_cal_select_mask( + QK_ptr, + Other_ptr, + Output_ptr, + Indices_ptr, + stride_qk: tl.constexpr, + numel: tl.constexpr, + BLOCK: tl.constexpr, +): + rows = tl.arange(0, BLOCK) * stride_qk + cols = tl.arange(0, BLOCK) + offs = rows[:, None] + cols[None, :] + row_indices = tl.load(Indices_ptr) + col_indices = tl.load(Indices_ptr + 1) + + qk_ub = tl.load(QK_ptr + offs) + other = tl.load(Other_ptr + offs) + mask_rows = rows < row_indices * stride_qk + mask_cols = cols < col_indices + + res = tl.where(mask_rows[:, None] & mask_cols[None, :], qk_ub, other) + tl.store(Output_ptr + offs, res) + + +def torch_cal_select_mask_bool( + index: torch.Tensor, + numel, + BLOCK, +): + offs = torch.arange(0, BLOCK) + true_tensor = torch.arange(0, BLOCK) < numel + false_tensor = torch.arange(0, BLOCK) >= numel + mask = offs < index + + res = torch.where(mask, true_tensor, false_tensor) + return res + + +def torch_cal_select_mask_one_dim_static_offset( + QK: torch.Tensor, + Other: torch.Tensor, + start_zero, +): + Output = Other.clone() + if start_zero: + Output[:1] = QK[:1] + else: + Output[1:] = QK[1:] + return Output + + +def torch_cal_select_mask_one_dim_dyn_offset( + QK: torch.Tensor, + Other: torch.Tensor, + Indices: torch.Tensor, + start_zero, +): + limit_idx = Indices[0].item() + Output = Other.clone() + if start_zero: + Output[:limit_idx] = QK[:limit_idx] + else: + Output[limit_idx:] = QK[limit_idx:] + return Output + + +def torch_cal_select_mask_one_dim_two_dyn_offset( + QK: torch.Tensor, + Other: torch.Tensor, + Indices: torch.Tensor, +): + min_idx = Indices[0].item() + max_idx = Indices[1].item() + Output = Other.clone() + Output[min_idx:max_idx] = QK[min_idx:max_idx] + + return Output + + +def torch_cal_select_mask_two_dim_static_offset( + QK: torch.Tensor, + Other: torch.Tensor, + start_zero, +): + Output = Other.clone() + if start_zero: + Output[:1, :1] = QK[:1, :1] + else: + Output[1:, 1:] = QK[1:, 1:] + + return Output + + +def torch_cal_select_mask_two_dim_one_dyn_offset( + QK: torch.Tensor, + Other: torch.Tensor, + Indices: torch.Tensor, + start_zero, +): + min_idx = Indices[0].item() + max_idx = Indices[1].item() + Output = Other.clone() + if start_zero: + Output[:1, min_idx:max_idx] = QK[:1, min_idx:max_idx] + else: + Output[1:, min_idx:max_idx] = QK[1:, min_idx:max_idx] + + return Output + + +def torch_cal_select_mask_two_dim_all_dyn_offset( + QK: torch.Tensor, + Other: torch.Tensor, + Indices: torch.Tensor, +): + min_idx1 = Indices[0].item() + max_idx1 = Indices[1].item() + min_idx2 = Indices[2].item() + max_idx2 = Indices[3].item() + Output = Other.clone() + Output[min_idx1:max_idx1, min_idx2:max_idx2] = QK[min_idx1:max_idx1, min_idx2:max_idx2] + + return Output + + +def torch_cal_select_mask( + QK: torch.Tensor, + Other: torch.Tensor, + Indices: torch.Tensor, +): + row_limit_idx = Indices[0].item() + col_limit_idx = Indices[1].item() + Output = Other.clone() + Output[:row_limit_idx, :col_limit_idx] = QK[:row_limit_idx, :col_limit_idx] + return Output + + +@pytest.mark.parametrize('param_list', [['bool', 64, 63]]) +def test_select_analysis_bool(param_list): + dtype, SEQ_LEN, index = param_list + assert dtype == 'bool' + qk_cal = torch.empty(SEQ_LEN).npu() + indices = torch.tensor([index]).npu() + qk_ref = torch_cal_select_mask_bool(index, SEQ_LEN, SEQ_LEN) + kernel_cal_select_mask_bool[(1, )](qk_cal, indices, SEQ_LEN, SEQ_LEN) + test_common.validate_cmp(dtype, qk_cal, qk_ref) + + +@pytest.mark.parametrize('param_list', [ + ['float16', 64, True], + ['float32', 64, False], +]) +def test_select_analysis_one_dim_static_offset(param_list): + dtype, SEQ_LEN, start_zero = param_list + assert dtype != 'bool' + qk = torch.rand([SEQ_LEN], dtype=eval('torch.' + dtype), device='npu') + qk_cal = torch.empty_like(qk).npu() + other = torch.zeros_like(qk).npu() + qk_ref = torch_cal_select_mask_one_dim_static_offset(qk, other, start_zero) + kernel_cal_select_mask_one_dim_static_offset[(1, )](qk, other, qk_cal, SEQ_LEN, start_zero) + test_common.validate_cmp(dtype, qk_cal, qk_ref) + + +@pytest.mark.parametrize('param_list', [ + ['float16', 64, 62, True], + ['float32', 64, 62, False], +]) +def test_select_analysis_one_dim_dyn_offset(param_list): + dtype, SEQ_LEN, index, start_zero = param_list + assert dtype != 'bool' + qk = torch.rand([SEQ_LEN], dtype=eval('torch.' + dtype), device='npu') + qk_cal = torch.empty_like(qk).npu() + other = torch.zeros_like(qk).npu() + indices_tensor = torch.tensor([index]).npu() + qk_ref = torch_cal_select_mask_one_dim_dyn_offset(qk, other, indices_tensor, start_zero) + kernel_cal_select_mask_one_dim_dyn_offset[(1, )](qk, other, qk_cal, indices_tensor, SEQ_LEN, start_zero) + test_common.validate_cmp(dtype, qk_cal, qk_ref) + + +@pytest.mark.parametrize('param_list', [ + ['float16', 64, 1, 62], + ['float32', 64, 1, 62], +]) +def test_select_analysis_one_dim_two_dyn_offset(param_list): + dtype, SEQ_LEN, minIndex, maxIndex = param_list + assert dtype != 'bool' + qk = torch.rand([SEQ_LEN], dtype=eval('torch.' + dtype), device='npu') + qk_cal = torch.empty_like(qk).npu() + other = torch.zeros_like(qk).npu() + indices_tensor = torch.tensor([minIndex, maxIndex]).npu() + qk_ref = torch_cal_select_mask_one_dim_two_dyn_offset(qk, other, indices_tensor) + kernel_cal_select_mask_one_dim_two_dyn_offset[(1, )](qk, other, qk_cal, indices_tensor, SEQ_LEN) + test_common.validate_cmp(dtype, qk_cal, qk_ref) + + +@pytest.mark.parametrize('param_list', [ + ['float16', 64, True], + ['float32', 64, False], +]) +def test_select_analysis_two_dim_static_offset(param_list): + dtype, SEQ_LEN, start_zero = param_list + assert dtype != 'bool' + qk = torch.rand([SEQ_LEN, SEQ_LEN], dtype=eval('torch.' + dtype), device='npu') + qk_cal = torch.empty_like(qk).npu() + other = torch.zeros_like(qk).npu() + qk_ref = torch_cal_select_mask_two_dim_static_offset(qk, other, start_zero) + kernel_cal_select_mask_two_dim_static_offset[(1, )](qk, other, qk_cal, SEQ_LEN, start_zero) + test_common.validate_cmp(dtype, qk_cal, qk_ref) + + +@pytest.mark.parametrize('param_list', [ + ['float16', 64, 1, 62, True], + ['float32', 64, 1, 62, False], +]) +def test_select_analysis_two_dim_one_dyn_offset(param_list): + dtype, SEQ_LEN, minIndex, maxIndex, start_zero = param_list + assert dtype != 'bool' + + qk = torch.rand([SEQ_LEN, SEQ_LEN], dtype=eval('torch.' + dtype), device='npu') + qk_cal = torch.empty_like(qk).npu() + other = torch.zeros_like(qk).npu() + indices_tensor = torch.tensor([minIndex, maxIndex]).npu() + qk_ref = torch_cal_select_mask_two_dim_one_dyn_offset(qk, other, indices_tensor, start_zero) + kernel_cal_select_mask_two_dim_one_dyn_offset[(1, )](qk, other, qk_cal, indices_tensor, SEQ_LEN, start_zero) + test_common.validate_cmp(dtype, qk_cal, qk_ref) + + +@pytest.mark.parametrize('param_list', [ + ['float16', 64, 2, 62, 1, 63], + ['float32', 64, 2, 62, 1, 63], +]) +def test_select_analysis_two_dim_all_dyn_offset(param_list): + dtype, SEQ_LEN, minIndex1, maxIndex1, minIndex2, maxIndex2 = param_list + assert dtype != 'bool' + + qk = torch.rand([SEQ_LEN, SEQ_LEN], dtype=eval('torch.' + dtype), device='npu') + qk_cal = torch.empty_like(qk).npu() + other = torch.zeros_like(qk).npu() + indices_tensor = torch.tensor([minIndex1, maxIndex1, minIndex2, maxIndex2]).npu() + qk_ref = torch_cal_select_mask_two_dim_all_dyn_offset(qk, other, indices_tensor) + kernel_cal_select_mask_two_dim_all_dyn_offset[(1, )](qk, other, qk_cal, indices_tensor, SEQ_LEN) + test_common.validate_cmp(dtype, qk_cal, qk_ref) + + +@pytest.mark.parametrize('param_list', [ + ['float16', 64, 63, 62], + ['float32', 64, 63, 62], +]) +def test_select_analysis(param_list): + dtype, SEQ_LEN, indice_x, indice_y = param_list + assert dtype != 'bool' + qk = torch.rand([SEQ_LEN, SEQ_LEN], dtype=eval('torch.' + dtype), device='npu') + qk_cal = torch.empty_like(qk).npu() + other = torch.zeros_like(qk).npu() + indices_tensor = torch.tensor([indice_x, indice_y]).npu() + qk_ref = torch_cal_select_mask(qk, other, indices_tensor) + kernel_cal_select_mask[(1, )](qk, other, qk_cal, indices_tensor, qk.stride(0), SEQ_LEN * SEQ_LEN, SEQ_LEN) + test_common.validate_cmp(dtype, qk_cal, qk_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_select_analysis_for_divsiop.py b/third_party/ascend/unittest/pytest_ut/test_select_analysis_for_divsiop.py new file mode 100644 index 0000000000..12e2166723 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_select_analysis_for_divsiop.py @@ -0,0 +1,84 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +def torch_divsiop_select_analysis(offs_num, divnum, maxindex, index, query): + idx = query[index] + offs_m = torch.arange(offs_num, dtype=torch.int32).npu() + query_pos = offs_m // divnum - idx + mask = query_pos < maxindex + + tensor0 = torch.tensor(0, dtype=torch.int32).npu() + tensor1 = torch.tensor(1, dtype=torch.int32).npu() + result = torch.where(mask, tensor1, tensor0).npu() + return result + + +@triton.jit +def divsiop_select_analysis_kernel1(index, out_ptr, query, offs_num: tl.constexpr, divnum: tl.constexpr, + maxindex: tl.constexpr): + idx = tl.load(query + index) + offs_m = tl.arange(0, offs_num) + + query_pos = offs_m // divnum - idx + + mask = query_pos < maxindex + query_mask = tl.where(mask, 1, 0).to(tl.int1) + tl.store(out_ptr + tl.arange(0, offs_num), query_mask) + + +@triton.jit +def divsiop_select_analysis_kernel2(index, out_ptr, query, offs_num: tl.constexpr, divnum: tl.constexpr, + maxindex: tl.constexpr): + idx = tl.load(query + index) + offs_m = tl.arange(0, offs_num) + + query_pos = -idx + offs_m // divnum + mask = query_pos < maxindex + + query_mask = tl.where(mask, 1, 0).to(tl.int1) + tl.store(out_ptr + tl.arange(0, offs_num), query_mask) + + +@pytest.mark.parametrize('param_list', [[16, 4, 2, index] for index in range(0, 4)]) +def test_divsiop_select_analysis1(param_list): + offs_num, divnum, maxindex, index = param_list + query = torch.tensor(range(0, divnum)).npu() + y_ref = torch_divsiop_select_analysis(offs_num, divnum, maxindex, index, query).npu() + y_cal = torch.full((offs_num, ), 2, dtype=torch.int32).npu() + divsiop_select_analysis_kernel1[(1, )](index, y_cal, query, offs_num, divnum, maxindex) + test_common.validate_cmp('int32', y_cal, y_ref) + + +@pytest.mark.parametrize('param_list', [[16, 4, 2, index] for index in range(0, 4)]) +def test_divsiop_select_analysis2(param_list): + offs_num, divnum, maxindex, index = param_list + query = torch.tensor(range(0, divnum)).npu() + y_ref = torch_divsiop_select_analysis(offs_num, divnum, maxindex, index, query).npu() + y_cal = torch.full((offs_num, ), 2, dtype=torch.int32).npu() + divsiop_select_analysis_kernel2[(1, )](index, y_cal, query, offs_num, divnum, maxindex) + test_common.validate_cmp('int32', y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_select_analysis_for_invert.py b/third_party/ascend/unittest/pytest_ut/test_select_analysis_for_invert.py new file mode 100644 index 0000000000..64abd9b3cb --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_select_analysis_for_invert.py @@ -0,0 +1,153 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +@triton.jit +def cal_atten_mask_kernel( + QK_ptr, + Indices_ptr, + stride_qk_m, + stride_qk_n, + stride_ik, + SEQ_LEN: tl.constexpr, + sparse_block_size: tl.constexpr, + BLOCK_SBS: tl.constexpr, + TOPK_BASE: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_n = tl.program_id(axis=1) + + idx_sub_sbs = pid_n + cur_s1 = pid_m * BLOCK_SBS + cur_s2 = cur_s1 + BLOCK_SBS + + if cur_s1 >= SEQ_LEN: + return + + beg_sbs = idx_sub_sbs * BLOCK_SBS // sparse_block_size + end_sbs = ((idx_sub_sbs + 1) * BLOCK_SBS) // sparse_block_size + + valid_col_end = cur_s1 + (cur_s2 - cur_s1) + + offs_m = cur_s1 + tl.arange(0, BLOCK_SBS) + offs_n_base = idx_sub_sbs * BLOCK_SBS + offs_n = offs_n_base + tl.arange(0, BLOCK_SBS) + + mask_m = offs_m < SEQ_LEN + mask_n = offs_n < SEQ_LEN + mask_load = mask_m[:, None] & mask_n[None, :] + + qk_ub = tl.load(QK_ptr + offs_m[:, None] * stride_qk_m + offs_n[None, :] * stride_qk_n, mask=mask_load, other=0.0) + + for idx_k in range(beg_sbs, end_sbs): + idx_s2 = tl.load(Indices_ptr + TOPK_BASE + idx_k * stride_ik) + if idx_s2 != -1 and idx_s2 * sparse_block_size > valid_col_end: + idx_lower_sbs = idx_k * sparse_block_size - \ + idx_sub_sbs * BLOCK_SBS + idx_higher_sbs = (idx_k + 1) * sparse_block_size - \ + idx_sub_sbs * BLOCK_SBS + mask_lower_sbs = tl.arange(0, BLOCK_SBS) >= idx_lower_sbs + mask_higher_sbs = tl.arange(0, BLOCK_SBS) < idx_higher_sbs + qk_ub = tl.where((mask_lower_sbs & mask_higher_sbs)[None, :], float("-inf"), qk_ub) + + tl.store(QK_ptr + offs_m[:, None] * stride_qk_m + offs_n[None, :] * stride_qk_n, qk_ub, mask=mask_load) + + +def launch_cal_atten_mask(qk_tensor, indices_tensor, sparse_block_size=64, block_sbs=128): + """ + qk_tensor: (SEQ_LEN, SEQ_LEN) + indices_tensor: (K,) / (BATCH, K, ...) + """ + assert qk_tensor.is_contiguous() + M, N = qk_tensor.shape + + stride_qk_m = qk_tensor.stride(0) + stride_qk_n = qk_tensor.stride(1) + + stride_ik = 1 + topk_base = 0 + + grid = (triton.cdiv(M, block_sbs), triton.cdiv(N, block_sbs)) + cal_atten_mask_kernel[grid]( + qk_tensor, + indices_tensor, + stride_qk_m, + stride_qk_n, + stride_ik, + SEQ_LEN=M, + sparse_block_size=sparse_block_size, + BLOCK_SBS=block_sbs, + TOPK_BASE=topk_base, + ) + return qk_tensor + + +def torch_cal_atten_mask( + qk, + indices, + sparse_block_size, + block_sbs, + topk_base=0, +): + device = qk.device + dtype = qk.dtype + M, N = qk.shape + + row_ids = torch.arange(M, device=device).unsqueeze(1) + col_ids = torch.arange(N, device=device).unsqueeze(0) + + k_idx_global = col_ids // sparse_block_size + lookup_idx = k_idx_global + topk_base + max_valid_idx = indices.numel() - 1 + + valid_lookup = (lookup_idx >= 0) & (lookup_idx <= max_valid_idx) + safe_lookup_idx = lookup_idx.clamp(0, max_valid_idx) + idx_s2_map = indices.gather(0, safe_lookup_idx.squeeze(0)).unsqueeze(0) + idx_s2_map = torch.where(valid_lookup, idx_s2_map, torch.tensor(-1, device=device)) + + row_block_ends = ((row_ids // block_sbs) + 1) * block_sbs + row_block_ends = torch.min(row_block_ends, torch.tensor(N, device=device)) + + start_pos_k_map = idx_s2_map * sparse_block_size + cond_valid = (idx_s2_map != -1) + cond_exceed = (start_pos_k_map > row_block_ends) + final_mask = cond_valid & cond_exceed + + qk_out = torch.where(final_mask, torch.tensor(float("-inf"), dtype=dtype, device=device), qk) + return qk_out + + +@pytest.mark.parametrize('param_list', [['float32', 1024, 128, 64]]) +def test_divsiop_select_analysis1(param_list): + dtype, SEQ_LEN, BLOCK_SBS, SPARSE_BLOCK = param_list + qk = torch.zeros((SEQ_LEN, SEQ_LEN), dtype=eval('torch.' + dtype), device='npu') + K_SIZE = 20 + indices = torch.full((K_SIZE, ), -1, dtype=torch.int32, device='npu') + indices[10] = 20 + qk_ref = torch_cal_atten_mask(qk.clone(), indices, sparse_block_size=SPARSE_BLOCK, block_sbs=BLOCK_SBS) + qk_cal = launch_cal_atten_mask(qk, indices, sparse_block_size=SPARSE_BLOCK, block_sbs=BLOCK_SBS) + test_common.validate_cmp(dtype, qk_cal, qk_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_sigmoid.py b/third_party/ascend/unittest/pytest_ut/test_sigmoid.py new file mode 100644 index 0000000000..f2d141dc80 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_sigmoid.py @@ -0,0 +1,60 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import pytest +import test_common + + +def torch_sigmoid(x0, x1): + res = x0 + torch.sigmoid(x1) + return res + + +@triton.jit +def triton_sigmoid(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + xindex = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = xindex < xnumel + x0 = xindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.load(in_ptr1 + (x0), xmask) + tmp2 = tmp0 + tl.sigmoid(tmp1) + tl.store(out_ptr0 + (xindex), tmp2, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_sigmoid(param_list): + # 生成数据 + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + # torch结果 + y_ref = torch_sigmoid(x0, x1) + # triton结果 + y_cal = test_common.generate_tensor(shape, dtype).npu() + triton_sigmoid[ncore, 1, 1](x0, x1, y_cal, x0.numel(), xblock, xblock_sub) + # 比较结果 + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_signbit.py b/third_party/ascend/unittest/pytest_ut/test_signbit.py new file mode 100644 index 0000000000..8350094340 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_signbit.py @@ -0,0 +1,81 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common +import triton.language.extra.cann.libdevice as libdevice + + +@triton.jit +def triton_signbit(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + xindex = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = xindex < xnumel + x0 = tl.load(in_ptr0 + xindex, xmask) + y = libdevice.signbit(x0) + tl.store(out_ptr0 + xindex, y, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_signbit_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch.signbit(x).npu() + y_cal = torch.zeros(shape).bool().npu() + triton_signbit[ncore, 1, 1](x, y_cal, x.numel(), xblock, xblock_sub) + test_common.validate_cmp('bool', y_cal, y_ref) + + +@pytest.mark.parametrize('param_list', [ + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_all_blocks_parallel(param_list, monkeypatch): + monkeypatch.setenv("TRITON_ALL_BLOCKS_PARALLEL", "1") + dtype, shape, ncore, xblock, xblock_sub = param_list + x = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch.signbit(x).npu() + y_cal = torch.zeros(shape).bool().npu() + triton_signbit[ncore, 1, 1](x, y_cal, x.numel(), xblock, xblock_sub) + test_common.validate_cmp('bool', y_cal, y_ref) + monkeypatch.delenv("TRITON_ALL_BLOCKS_PARALLEL") + + +@pytest.mark.parametrize('param_list', [ + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_auto_blockify(param_list, monkeypatch): + monkeypatch.setenv("TRITON_ALL_BLOCKS_PARALLEL", "1") + dtype, shape, ncore, xblock, xblock_sub = param_list + x = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch.signbit(x).npu() + y_cal = torch.zeros(shape).bool().npu() + triton_signbit[ncore, 1, 1](x, y_cal, x.numel(), xblock, xblock_sub, auto_blockify_size=ncore) + test_common.validate_cmp('bool', y_cal, y_ref) + monkeypatch.delenv("TRITON_ALL_BLOCKS_PARALLEL") diff --git a/third_party/ascend/unittest/pytest_ut/test_silu.py b/third_party/ascend/unittest/pytest_ut/test_silu.py new file mode 100644 index 0000000000..038dcccec7 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_silu.py @@ -0,0 +1,99 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + + +def standard_unary(x0, dtype): + res = x0 * (1 / (1 + torch.exp(-x0))) + return res + + +def standard_binary(x0, y0, dtype): + res = x0 + y0 + return res + + +@triton.jit +def triton_elementwise_unary(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + ret = x * (1 / (1 + tl.math.exp(-x))) + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +@triton.jit +def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + y = tl.load(in_ptr1 + idx_block, mask=idx_block < N) + ret = x + y + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + # (torch.bfloat16, 'bfloat16'), + # (torch.int8, 'int8'), + # (torch.int16, 'int16'), + # (torch.int32, 'int32'), + # (torch.int64, 'int64'), +] + +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('N,NUMEL', shapes) +def test_elementwsie_common(dtype, sigtype, N, NUMEL): + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == "int64": + N = map_for_64_t[N] if N in map_for_64_t else N + + print(f"elementwise : ({N},) {dtype} {sigtype}") + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standard_unary(x0, dtype) + x0 = x0.npu() + print(ans) + + out = torch.zeros((N, ), dtype=dtype).npu() + triton_elementwise_unary[1, 1, 1](x0, out, N=N, NUMEL=NUMEL, debug=True) + print(out) + + test_common.validate_cmp(sigtype, out, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_simplify_iterargs.py b/third_party/ascend/unittest/pytest_ut/test_simplify_iterargs.py new file mode 100644 index 0000000000..d1dabaf669 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_simplify_iterargs.py @@ -0,0 +1,78 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + +import torch +import torch_npu + + +@triton.jit +def triton_fn_expanddims(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, YBLOCK: tl.constexpr, YBLOCK_SUB: tl.constexpr): + base1 = tl.arange(0, XBLOCK)[:, None] + base2 = tl.arange(0, YBLOCK_SUB)[None, :] + loops1: tl.constexpr = YBLOCK // YBLOCK_SUB # assume it's divisible + for _ in range(loops1): + x0 = base1 * YBLOCK + base2 + base2 = base2 + YBLOCK_SUB + tmp0 = tl.load(in_ptr0 + (x0), None) + tl.store(out_ptr0 + (x0), tmp0, None) + + +@triton.jit +def triton_fn_broadcast(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, YBLOCK: tl.constexpr, YBLOCK_SUB: tl.constexpr): + base1 = tl.arange(0, XBLOCK)[:, None] + base2 = tl.arange(0, YBLOCK_SUB)[None, :] + base2 = base2.broadcast_to((XBLOCK, YBLOCK_SUB)) + loops1: tl.constexpr = YBLOCK // YBLOCK_SUB # assume it's divisible + for _ in range(loops1): + x0 = base1 * YBLOCK + base2 + base2 = base2 + YBLOCK_SUB + tmp0 = tl.load(in_ptr0 + (x0), None) + tl.store(out_ptr0 + (x0), tmp0, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (128, 128), 128, 128, 32], +]) +def test_expanddims(param_list): + dtype, shape, xblock, yblock, yblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_fn_expanddims[(1, )](x0, y_cal, xblock, yblock, yblock_sub) + test_common.validate_cmp(dtype, y_cal, x0) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (128, 128), 128, 128, 32], +]) +def test_broadcast(param_list): + dtype, shape, xblock, yblock, yblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_fn_broadcast[(1, )](x0, y_cal, xblock, yblock, yblock_sub) + test_common.validate_cmp(dtype, y_cal, x0) diff --git a/third_party/ascend/unittest/pytest_ut/test_sin.py b/third_party/ascend/unittest/pytest_ut/test_sin.py new file mode 100644 index 0000000000..6258d35d6b --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_sin.py @@ -0,0 +1,99 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + + +def standard_unary(x0, dtype): + res = torch.sin(x0) + return res + + +def standard_binary(x0, y0, dtype): + res = x0 + y0 + return res + + +@triton.jit +def triton_elementwise_unary(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + ret = tl.sin(x) + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +@triton.jit +def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + y = tl.load(in_ptr1 + idx_block, mask=idx_block < N) + ret = x + y + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +types = [ + (torch.float32, 'float32'), + # (torch.float16, 'float16'), + # (torch.bfloat16, 'bfloat16'), + # (torch.int8, 'int8'), + # (torch.int16, 'int16'), + # (torch.int32, 'int32'), + # (torch.int64, 'int64'), +] + +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('N,NUMEL', shapes) +def test_elementwsie_common(dtype, sigtype, N, NUMEL): + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == "int64": + N = map_for_64_t[N] if N in map_for_64_t else N + + print(f"elementwise : ({N},) {dtype} {sigtype}") + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standard_unary(x0, dtype) + x0 = x0.npu() + print(ans) + + out = torch.zeros((N, ), dtype=dtype).npu() + triton_elementwise_unary[1, 1, 1](x0, out, N=N, NUMEL=NUMEL, debug=True) + print(out) + + test_common.validate_cmp(sigtype, out, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_sinh.py b/third_party/ascend/unittest/pytest_ut/test_sinh.py new file mode 100644 index 0000000000..81227bd868 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_sinh.py @@ -0,0 +1,59 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + +import torch +import torch_npu + + +@triton.jit +def triton_sinh(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = libdevice.sinh(tmp0) + tl.store(out_ptr0 + (x0), tmp1, None) + + +@pytest.mark.parametrize('param_list', ['float32', 'float16', 'bfloat16']) +def test_sinh_special(param_list): + dtype = param_list + x_near_zero = torch.linspace(-1.0, 1.0, 256, dtype=eval("torch." + dtype)).npu() + x_medium_neg = torch.linspace(-10.0, -1.0, 128, dtype=eval("torch." + dtype)).npu() + x_medium_pos = torch.linspace(1.0, 10.0, 128, dtype=eval("torch." + dtype)).npu() + + x_large_neg = torch.linspace(-1e4, -10.0, 64, dtype=eval("torch." + dtype)).npu() + x_large_pos = torch.linspace(10.0, 1e4, 64, dtype=eval("torch." + dtype)).npu() + + x0 = torch.cat([x_near_zero, x_medium_neg, x_medium_pos, x_large_neg, x_large_pos], dim=0) + + y_ref = torch.sinh(x0) + y_cal = torch.zeros_like(y_ref) + triton_sinh[1, 1, 1](x0, y_cal, x0.shape[0], x0.shape[0]) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_sink_broadcast.py b/third_party/ascend/unittest/pytest_ut/test_sink_broadcast.py new file mode 100644 index 0000000000..99b347246d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_sink_broadcast.py @@ -0,0 +1,89 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + +import torch +import torch_npu + + +@triton.jit +def triton_sink_broadcast1(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, YBLOCK: tl.constexpr): + base1 = tl.arange(0, XBLOCK)[:, None] * YBLOCK + base2 = tl.arange(0, YBLOCK)[None, :] + base1 = base1.broadcast_to((XBLOCK, YBLOCK)) + tmp0 = tl.load(in_ptr0 + base1, None) + index = base1 + base2 + tl.store(out_ptr0 + index, tmp0, None) + + +@triton.jit +def triton_sink_broadcast2(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, YBLOCK: tl.constexpr): + base1 = tl.arange(0, XBLOCK)[:, None] * YBLOCK + base2 = tl.arange(0, YBLOCK)[None, :] + base1 = base1.broadcast_to((XBLOCK, YBLOCK)) + tmp0 = tl.load(in_ptr0 + base1, base1 < XBLOCK * YBLOCK, other=0.0) + index = base1 + base2 + tl.store(out_ptr0 + index, tmp0, index < XBLOCK * YBLOCK) + + +@triton.jit +def triton_sink_broadcast3(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, YBLOCK: tl.constexpr): + base1 = (tl.arange(0, XBLOCK) * YBLOCK)[:, None] + base2 = tl.arange(0, YBLOCK)[None, :] + base1 = base1.broadcast_to((XBLOCK, YBLOCK)) + tmp0 = tl.load(in_ptr0 + base1 + base2, (base1 + base2) < XBLOCK * YBLOCK, other=0.0) + index = base1 + base2 + tl.store(out_ptr0 + index, tmp0, index < XBLOCK * YBLOCK) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (32, 32), 32, 32], +]) +def test_sink_broadcast(param_list): + dtype, shape, xblock, yblock = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + y_ref = x0.clone() + y_ref = y_ref[:, 0].unsqueeze(1).expand(-1, x0.size(1)) + + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + y_cal2 = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_sink_broadcast1[(1, )](x0, y_cal, xblock, yblock) + triton_sink_broadcast2[(1, )](x0, y_cal2, xblock, yblock) + test_common.validate_cmp(dtype, y_cal, y_ref) + test_common.validate_cmp(dtype, y_cal2, y_ref) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (32, 32), 32, 32], +]) +def test_sink_broadcast3(param_list): + dtype, shape, xblock, yblock = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + y_ref = x0.clone() + + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_sink_broadcast3[(1, )](x0, y_cal, xblock, yblock) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_softmax.py b/third_party/ascend/unittest/pytest_ut/test_softmax.py new file mode 100644 index 0000000000..e9d25c475f --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_softmax.py @@ -0,0 +1,130 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import test_common +import pytest + + +def naive_softmax(x): + # read MN elements ; write M elements + x_max = x.max(dim=1)[0] + # read MN + M elements ; write MN elements + z = x - x_max[:, None] + # read MN elements ; write MN elements + numerator = torch.exp(z) + # read MN elements ; write M elements + denominator = numerator.sum(dim=1) + # read MN + M elements ; write MN elements + ret = numerator / denominator[:, None] + # in total: read 5MN + 2M elements ; wrote 3MN + 2M elements + return ret + + +@triton.jit +def softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride, n_rows, n_cols, + BLOCK_SIZE: tl.constexpr): + # starting row of the program + row_start = tl.program_id(0) + row_step = tl.num_programs(0) + for row_idx in tl.range(row_start, n_rows, row_step): + # The stride represents how much we need to increase the pointer to advance 1 row + row_start_ptr = input_ptr + row_idx * input_row_stride + # The block size is the next power of two greater than n_cols, so we can fit each + # row in a single block + col_offsets = tl.arange(0, BLOCK_SIZE) + input_ptrs = row_start_ptr + col_offsets + # Load the row into SRAM, using a mask since BLOCK_SIZE may be > than n_cols + mask = col_offsets < n_cols + row = tl.load(input_ptrs, mask=mask, other=-float('inf')) + # Subtract maximum for numerical stability + row_minus_max = row - tl.max(row, axis=0) + numerator = tl.exp(row_minus_max) + denominator = tl.sum(numerator, axis=0) + softmax_output = numerator / denominator + # Write back output to DRAM + output_row_start_ptr = output_ptr + row_idx * output_row_stride + output_ptrs = output_row_start_ptr + col_offsets + tl.store(output_ptrs, softmax_output, mask=mask) + + +kernels = {} + + +def softmax(x, stream): + n_rows, n_cols = x.shape + + BLOCK_SIZE = triton.next_power_of_2(n_cols) + + y = torch.empty_like(x) + + kernel, num_programs = kernels.get(BLOCK_SIZE, (None, 0)) + if kernel is None: + num_programs = 32 + kernel = softmax_kernel + kernels[BLOCK_SIZE] = (kernel, num_programs) + + num_programs = min(num_programs, n_rows) + + # Create a number of persistent programs. + kernel[(num_programs, 1, 1)](y, x, x.stride(0), y.stride(0), n_rows, n_cols, BLOCK_SIZE) + return y + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + (torch.bfloat16, 'bfloat16'), +] + +shapes = [ + (1823, 781), + (1823, 2), + (1823, 4), + (1823, -32), + (1823, -100), + (1823, -256), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.skip(reason="randomly failed") +@pytest.mark.parametrize('dtype, sigtype', types) +@pytest.mark.parametrize('M, N', shapes) +def test_softmax(dtype, sigtype, M, N): + torch_npu.npu.utils.set_device(0) + M = (-M) // torch.tensor(0, dtype=dtype).element_size() if M < 0 else M + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + M = map_for_64_t[M] if M in map_for_64_t else M + N = map_for_64_t[N] if N in map_for_64_t else N + + device = torch.npu.current_device() + stream = torch.npu.current_stream(device).npu_stream + torch.manual_seed(0) + x = torch.randn(M, N, dtype=dtype, device='npu') + y_triton = softmax(x, stream) + y_torch = torch.softmax(x, axis=1) + test_common.validate_cmp(sigtype, y_triton, y_torch) diff --git a/third_party/ascend/unittest/pytest_ut/test_sort.py b/third_party/ascend/unittest/pytest_ut/test_sort.py new file mode 100644 index 0000000000..62a103447e --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_sort.py @@ -0,0 +1,102 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import pytest +import numpy as np +import torch +import triton.language as tl +import triton.language.extra.cann.extension as extension +import test_common + +# --------------- +# test sort op +# --------------- + + +@triton.jit +def sort_kernel_2d(X, Z, N: tl.constexpr, M: tl.constexpr, descending: tl.constexpr): + offx = tl.arange(0, M) + offy = tl.arange(0, N) * M + off2d = offx[None, :] + offy[:, None] + x = tl.load(X + off2d) + x = extension.sort(x, descending=descending, dim=1) + tl.store(Z + off2d, x) + + +@pytest.mark.interpreter +@pytest.mark.parametrize("shape", [(256, 16)]) +@pytest.mark.parametrize("descending", [False, True]) +@pytest.mark.parametrize("dtype", ['int8', 'int16', 'float16', 'float32', 'bfloat16', 'bool']) +def test_sort_2d(shape, descending, dtype): + if dtype == "bool": + x = test_common.generate_tensor(shape, dtype) + np_sorted = np.sort(x) + if descending: + np_sorted = np_sorted[:, ::-1].copy() + torch_res = torch.from_numpy(np_sorted).npu() + else: + x = test_common.generate_tensor(shape, dtype).npu() + torch_res = torch.sort(x, descending=descending)[0] + + x = x.npu() + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + N = x.shape[0] + M = x.shape[1] + sort_kernel_2d[(1, )](x, triton_res, N, M, descending) + assert (torch_res == triton_res).all(), (torch_res, triton_res) + + +@triton.jit +def sort_kernel_3d(X, Z, D0: tl.constexpr, D1: tl.constexpr, D2: tl.constexpr, descending: tl.constexpr): + off2 = tl.arange(0, D2) + off1 = tl.arange(0, D1) * D2 + off0 = tl.arange(0, D0) * D1 * D2 + + off = off2[None, None, :] + off1[None, :, None] + off0[:, None, None] + x = tl.load(X + off) + + x = extension.sort(x, descending=descending, dim=2) + + tl.store(Z + off, x) + + +@pytest.mark.interpreter +@pytest.mark.parametrize("shape", [(8, 4, 16)]) +@pytest.mark.parametrize("descending", [False, True]) +@pytest.mark.parametrize("dtype", ['int8', 'int16', 'float16', 'float32', 'bfloat16', 'bool']) +def test_sort_3d(shape, descending, dtype): + if dtype == "bool": + x = test_common.generate_tensor(shape, dtype) + np_sorted = np.sort(x) + if descending: + np_sorted = np_sorted[:, :, ::-1].copy() + torch_res = torch.from_numpy(np_sorted).npu() + else: + x = test_common.generate_tensor(shape, dtype).npu() + torch_res = torch.sort(x, descending=descending)[0] + + x = x.npu() + triton_res = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + D0 = x.shape[0] + D1 = x.shape[1] + D2 = x.shape[2] + sort_kernel_3d[(1, )](x, triton_res, D0, D1, D2, descending) + assert (torch_res == triton_res).all(), (torch_res, triton_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_split.py b/third_party/ascend/unittest/pytest_ut/test_split.py new file mode 100644 index 0000000000..96af5366b5 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_split.py @@ -0,0 +1,72 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl + +import torch +import torch_npu +import pytest +import test_common + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, output_ptr1, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + + xx, yy = tl.split(X) + + oidx = xidx[:, None] * YB + yidx[None, :] + + tl.store(output_ptr + oidx, xx) + tl.store(output_ptr1 + oidx, yy) + + +@pytest.mark.parametrize('para_type,data_type,XB,YB,ZB', [ + ['float32', torch.float32, 8, 8, 2], + ['float16', torch.float16, 8, 8, 2], + ['int8', torch.int8, 8, 128, 2], + ['int8', torch.int8, 8, 8, 2], +]) +def test_split(para_type, data_type, XB, YB, ZB): + + x = torch.randint(low=-128, high=128, size=(XB, YB, ZB), dtype=data_type).npu() + + a, b = torch.split(x, 1, dim=-1) + a = a.reshape(XB, YB) + b = b.reshape(XB, YB) + print(a) + print(b) + + output = torch.randint(1, (XB, YB), dtype=data_type).npu() + output1 = torch.randint(1, (XB, YB), dtype=data_type).npu() + fn_npu_[1, 1, 1](output, x, output1, XB, YB, ZB, debug=True) + + print(output) + print(output1) + + test_common.validate_cmp(para_type, a, output) + test_common.validate_cmp(para_type, b, output1) diff --git a/third_party/ascend/unittest/pytest_ut/test_sqrt.py b/third_party/ascend/unittest/pytest_ut/test_sqrt.py new file mode 100644 index 0000000000..7a79ed9946 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_sqrt.py @@ -0,0 +1,99 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + + +def standard_unary(x0, dtype): + res = torch.sqrt(x0) + return res + + +def standard_binary(x0, y0, dtype): + res = x0 + y0 + return res + + +@triton.jit +def triton_elementwise_unary(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + ret = tl.math.sqrt(x) + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +@triton.jit +def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + y = tl.load(in_ptr1 + idx_block, mask=idx_block < N) + ret = x + y + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + # (torch.bfloat16, 'bfloat16'), + # (torch.int8, 'int8'), + # (torch.int16, 'int16'), + # (torch.int32, 'int32'), + # (torch.int64, 'int64'), +] + +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('N,NUMEL', shapes) +def test_elementwsie_common(dtype, sigtype, N, NUMEL): + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == "int64": + N = map_for_64_t[N] if N in map_for_64_t else N + + print(f"elementwise : ({N},) {dtype} {sigtype}") + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standard_unary(x0, dtype) + x0 = x0.npu() + print(ans) + + out = torch.zeros((N, ), dtype=dtype).npu() + triton_elementwise_unary[1, 1, 1](x0, out, N=N, NUMEL=NUMEL, debug=True) + print(out) + + test_common.validate_cmp(sigtype, out, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_static_print_and_assert.py b/third_party/ascend/unittest/pytest_ut/test_static_print_and_assert.py new file mode 100644 index 0000000000..132245122c --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_static_print_and_assert.py @@ -0,0 +1,187 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import pytest +import test_common +import functools +import os +import re + +shape = (8, ) +XS = 8 +XVALS_INT = [ + 0, -128, # torch.iinfo(torch.int8).min + 127, # torch.iinfo(torch.int8).max + -32768, # torch.iinfo(torch.int16).min + 32767, # torch.iinfo(torch.int16).max + -2147483648, # torch.iinfo(torch.int32).min + 2147483647, # torch.iinfo(torch.int32).max + 9223372036854775807 +] # torch.iinfo(torch.int64).max + +XVALS_FP = [ + 0.0000000000e+00, # 0 + 1.1921000009e-07, # torch.finfo(torch.float32).eps + 9.7655999707e-04, # torch.finfo(torch.float16).eps + 7.8125000000e-03, # torch.finfo(torch.bfloat16).eps + 3.4027999388e+38, # torch.finfo(torch.float32).max + 6.5504000000e+04, # torch.finfo(torch.float16).max + 3.3894999515e+38, # torch.finfo(torch.bfloat16).max + 1.0000000000e+00 +] # 1 + + +def torch_func(x0, x1): + res = x0 + x1 + return res + + +@triton.jit +def triton_kernel(out_ptr0, in_ptr0, in_ptr1, XBLOCK: tl.constexpr, print_data_ptr: tl.constexpr, + assert_data_ptr: tl.constexpr): + idx = tl.arange(0, XBLOCK) + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tl.load(in_ptr1 + idx) + tmp2 = tmp0 + tmp1 + tl.static_print(print_data_ptr) + tl.static_assert(assert_data_ptr == assert_data_ptr, "assert_data should equal assert_data") + tl.store(out_ptr0 + idx, tmp2) + + +def triton_func(x0, x1, XS, print_data_ptr, assert_data_ptr): + out = torch.empty_like(x0) + triton_kernel[1, 1, 1](out, x0, x1, XS, print_data_ptr, assert_data_ptr) + return out + + +@pytest.mark.skip(reason="waiting for TA to support") +@pytest.mark.parametrize('sigtype', ['int8']) +@test_common.capture_output("-128") +def test_static_print_int8(capsys, sigtype): + dtype = eval(f"torch.{sigtype}") + x0 = torch.zeros(shape, dtype=dtype).npu() + x1 = torch.ones(shape, dtype=dtype).npu() + for i in range(x1.numel()): + x1[i] = XVALS_INT[i] + torch_ref = torch_func(x0, x1) + triton_cal = triton_func(x0, x1, XS, -128, XVALS_INT[0]) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) + + +@pytest.mark.skip(reason="waiting for TA to support") +@pytest.mark.parametrize('sigtype', ['int16']) +@test_common.capture_output("-32768") +def test_static_print_int16(capsys, sigtype): + dtype = eval(f"torch.{sigtype}") + x0 = torch.zeros(shape, dtype=dtype).npu() + x1 = torch.ones(shape, dtype=dtype).npu() + for i in range(x1.numel()): + x1[i] = XVALS_INT[i] + torch_ref = torch_func(x0, x1) + triton_cal = triton_func(x0, x1, XS, -32768, XVALS_INT[2]) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) + + +@pytest.mark.skip(reason="waiting for TA to support") +@pytest.mark.parametrize('sigtype', ['int32']) +@test_common.capture_output("-2147483648") +def test_static_print_int32(capsys, sigtype): + dtype = eval(f"torch.{sigtype}") + x0 = torch.zeros(shape, dtype=dtype).npu() + x1 = torch.ones(shape, dtype=dtype).npu() + for i in range(x1.numel()): + x1[i] = XVALS_INT[i] + torch_ref = torch_func(x0, x1) + triton_cal = triton_func(x0, x1, XS, -2147483648, XVALS_INT[4]) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) + + +@pytest.mark.skip(reason="waiting for TA to support") +@pytest.mark.parametrize('sigtype', ['int64']) +@test_common.capture_output("9223372036854775807") +def test_static_print_int64(capsys, sigtype): + dtype = eval(f"torch.{sigtype}") + x0 = torch.zeros(shape, dtype=dtype).npu() + x1 = torch.ones(shape, dtype=dtype).npu() + for i in range(x1.numel()): + x1[i] = XVALS_INT[i] + torch_ref = torch_func(x0, x1) + triton_cal = triton_func(x0, x1, XS, 9223372036854775807, XVALS_INT[-1]) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) + + +@pytest.mark.skip(reason="waiting for TA to support") +@pytest.mark.parametrize('sigtype', ['float16']) +@test_common.capture_output("1.1921000009e-07") +def test_static_print_float16(capsys, sigtype): + dtype = eval(f"torch.{sigtype}") + x0 = torch.zeros(shape, dtype=dtype).npu() + x1 = torch.ones(shape, dtype=dtype).npu() + for i in range(x1.numel()): + x1[i] = XVALS_FP[i] + torch_ref = torch_func(x0, x1) + triton_cal = triton_func(x0, x1, XS, 1.1921000009e-07, XVALS_FP[1]) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) + + +@pytest.mark.skip(reason="waiting for TA to support") +@pytest.mark.parametrize('sigtype', ['float32']) +@test_common.capture_output("0.0078125") +def test_static_print_float32(capsys, sigtype): + dtype = eval(f"torch.{sigtype}") + x0 = torch.zeros(shape, dtype=dtype).npu() + x1 = torch.ones(shape, dtype=dtype).npu() + for i in range(x1.numel()): + x1[i] = XVALS_FP[i] + torch_ref = torch_func(x0, x1) + triton_cal = triton_func(x0, x1, XS, 7.8125000000e-03, XVALS_FP[0]) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) + + +@pytest.mark.skip(reason="waiting for TA to support") +@pytest.mark.parametrize('sigtype', ['bfloat16']) +@test_common.capture_output("0.00097655999707") +def test_static_print_bfloat16(capsys, sigtype): + dtype = eval(f"torch.{sigtype}") + x0 = torch.zeros(shape, dtype=dtype).npu() + x1 = torch.ones(shape, dtype=dtype).npu() + for i in range(x1.numel()): + x1[i] = XVALS_FP[i] + torch_ref = torch_func(x0, x1) + triton_cal = triton_func(x0, x1, XS, 9.7655999707e-04, XVALS_FP[2]) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) + + +@pytest.mark.skip(reason="waiting for TA to support") +@pytest.mark.parametrize('sigtype', ['int8']) +@test_common.capture_output("True") +def test_static_print_bool(capsys, sigtype): + dtype = eval(f"torch.{sigtype}") + x0 = torch.zeros(shape, dtype=dtype).npu() + x1 = torch.ones(shape, dtype=dtype).npu() + for i in range(x1.numel()): + x1[i] = XVALS_INT[i] + torch_ref = torch_func(x0, x1) + triton_cal = triton_func(x0, x1, XS, True, True) + test_common.validate_cmp(sigtype, triton_cal, torch_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_store_scalar.py b/third_party/ascend/unittest/pytest_ut/test_store_scalar.py new file mode 100644 index 0000000000..52be31e3e8 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_store_scalar.py @@ -0,0 +1,46 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl + + +# load with mask, store with scalar +@triton.jit +def sum_kernel_1(inp, mid, M, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(0) + # 0 / 1 / 2 * 4 + (0,1,2,3) + offset = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + inp_ptrs = inp + offset + mask = offset < M + inp_val = tl.load(inp_ptrs, mask=mask).to(tl.float32) + sum_val = tl.sum(inp_val) + mid_ptr = mid + pid + tl.store(mid_ptr, sum_val) + + +def test_case(): + inp = torch.ones(16, device="npu", dtype=torch.float32) + mid = torch.empty(4, device="npu", dtype=torch.float32) + sum_kernel_1[(4, 1, 1)](inp, mid, 16, 4) + ref = torch.tensor([4.0, 4.0, 4.0, 4.0], device="npu", dtype=torch.float32) + assert torch.allclose(mid, ref, rtol=1e-03, atol=1e-03, equal_nan=True) diff --git a/third_party/ascend/unittest/pytest_ut/test_stride0_load.py b/third_party/ascend/unittest/pytest_ut/test_stride0_load.py new file mode 100644 index 0000000000..9c9295381d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_stride0_load.py @@ -0,0 +1,66 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl + + +def torch_expanddims_load(in0, in1, out0, YBLOCK, XBLOCK): + for i in range(0, YBLOCK): + tmp = in0[i] + 10 + res = in1[tmp] + out0[i, :] = res + + +@triton.jit +def triton_expanddims_load(in0, in1, out0, YBLOCK: tl.constexpr, XBLOCK: tl.constexpr): + base_y = tl.arange(0, YBLOCK) + y_idx = base_y[:, None] + y_mask = y_idx < YBLOCK + base_x = tl.arange(0, XBLOCK) + x_idx = base_x[None, :] + x_mask = x_idx < XBLOCK + + y = tl.load(in0 + y_idx, mask=y_mask) + tmp0 = tl.full([YBLOCK, XBLOCK], 10, tl.int32) + tmp1 = y + tmp0 + res = tl.load(in1 + tmp1, mask=y_mask) + tl.store(out0 + (y_idx * XBLOCK + x_idx), res, mask=y_mask & x_mask) + + +def test_case(): + YBLOCK = 4 + XBLOCK = 8 + in0 = torch.arange(0, YBLOCK, device="npu", dtype=torch.int32) + in1 = torch.arange(0, YBLOCK + 10, device="npu", dtype=torch.int32) + out0 = torch.empty((YBLOCK, XBLOCK), device="npu", dtype=torch.int32) + + torch_expanddims_load(in0, in1, out0, YBLOCK, XBLOCK) + + in0_triton = in0 + in1_triton = in1 + out0_triton = torch.empty((YBLOCK * XBLOCK, ), device="npu", dtype=torch.int32) + + triton_expanddims_load[(1, 1, 1)](in0_triton, in1_triton, out0_triton, YBLOCK, XBLOCK) + out0_triton = out0_triton.view(YBLOCK, XBLOCK) + + assert torch.allclose(out0, out0_triton, rtol=1e-03, atol=1e-03, equal_nan=True) diff --git a/third_party/ascend/unittest/pytest_ut/test_strides.py b/third_party/ascend/unittest/pytest_ut/test_strides.py new file mode 100644 index 0000000000..b2522ebe4d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_strides.py @@ -0,0 +1,140 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import pytest +import torch +import triton +import triton.language as tl + + +@triton.jit +def mv_kernel( + A, + B, + C, + N, + M, + stride_an, + stride_am, + stride_bm, + stride_cn, + BLOCK_N: tl.constexpr, + BLOCK_M: tl.constexpr, +): + pid = tl.program_id(0).to(tl.int64) + offset_n = pid * BLOCK_N + tl.arange(0, BLOCK_N)[:, None] + offset_m = tl.arange(0, BLOCK_M)[None, :] + n_mask = offset_n < N + A_ptrs = A + offset_n * stride_an + offset_m * stride_am + B_ptrs = B + offset_m * stride_bm + acc = tl.zeros((BLOCK_N, BLOCK_M), dtype=tl.float32) + for m in range(0, M, BLOCK_M): + m_mask = m + offset_m < M + a = tl.load(A_ptrs, mask=n_mask & m_mask, other=0.0).to(tl.float32) + b = tl.load(B_ptrs, mask=m_mask, other=0.0).to(tl.float32) + acc += a * b + A_ptrs += BLOCK_M * stride_am + B_ptrs += BLOCK_M * stride_bm + + acc = tl.sum(acc, axis=1) + C_ptrs = C + offset_n * stride_cn + tl.store(C_ptrs, acc[:, None], mask=n_mask) + + +def mv(inp, vec): + assert inp.shape[1] == vec.shape[0], "incompatible dimensions" + N, M = inp.shape + out = torch.empty((N, ), device=inp.device, dtype=inp.dtype) + + def grid(META): + return (triton.cdiv(N, META["BLOCK_N"]), ) + + mv_kernel[grid]( + inp, + vec, + out, + N, + M, + inp.stride(0), + inp.stride(1), + vec.stride(0), + out.stride(0), + 32, + 4, + ) + return out + + +def to_reference(inp, upcast=False): + if inp is None: + return None + ref_inp = inp + ref_inp = ref_inp.to("cpu") + if upcast: + if ref_inp.is_complex(): + ref_inp = ref_inp.to(torch.complex128) + else: + ref_inp = ref_inp.to(torch.float64) + return ref_inp + + +RESOLUTION = { + torch.bool: 0, + torch.int16: 0, + torch.int32: 0, + torch.int64: 0, + torch.float16: 1e-3, + torch.float32: 1.3e-6, + torch.bfloat16: 0.016, + torch.float64: 1e-7, + torch.complex32: 1e-3, + torch.complex64: 1.3e-6, +} + + +def assert_close(res, ref, dtype, equal_nan=False, reduce_dim=1): + assert res.dtype == dtype + ref = ref.to(dtype) + atol = 1e-4 * reduce_dim + rtol = RESOLUTION[dtype] + torch.testing.assert_close(res, ref, atol=atol, rtol=rtol, equal_nan=equal_nan) + + +def gems_assert_close(res, ref, dtype, equal_nan=False, reduce_dim=1): + res = res.to("cpu") + assert ref.device == torch.device("cpu") + assert_close(res, ref, dtype, equal_nan=equal_nan, reduce_dim=reduce_dim) + + +@pytest.mark.mv +@pytest.mark.parametrize("M, N", [(1, 32), (160, 1024), (5333, 497)]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_accuracy_mv(M, N, dtype): + matrix = torch.randn((N, M), dtype=dtype, device="npu") + vector = torch.randn((M, ), dtype=dtype, device="npu") + ref_matrix = to_reference(matrix, True) + ref_vector = to_reference(vector, True) + + ref_out = torch.mv(ref_matrix, ref_vector) + + res_out = mv(matrix, vector) + + gems_assert_close(res_out, ref_out, dtype, reduce_dim=M) diff --git a/third_party/ascend/unittest/pytest_ut/test_sub.py b/third_party/ascend/unittest/pytest_ut/test_sub.py new file mode 100644 index 0000000000..c799a48c5b --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_sub.py @@ -0,0 +1,62 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import numpy as np +import torch +import pytest +import test_common + + +def torch_pointwise(x0, x1): + res = x0 - x1 + return res + + +@triton.jit +def triton_sub(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop1 in range(loops1): + x0_prime = offset + (loop1 * XBLOCK_SUB) + base1 + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.load(in_ptr1 + (x0), None) + tmp2 = tmp0 - tmp1 + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['int32', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], + ['uint8', (2, 4096, 8), 2, 32768, 1024], +]) +def test_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_pointwise(x0, x1) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_sub[ncore, 1, 1](x0, x1, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_sub_vec_id.py b/third_party/ascend/unittest/pytest_ut/test_sub_vec_id.py new file mode 100644 index 0000000000..3fa24e76ed --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_sub_vec_id.py @@ -0,0 +1,132 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import pytest +import test_common +import triton.language.extra.cann.extension as al + + +@triton.jit +def triton_matmul_exp( + A_ptr, + B_ptr, + C_ptr, + TBuff_ptr, + M: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, + sub_M: tl.constexpr, +): + """function: + 1) The matrix matmul + 2) The matrix exp computation + For example, + 1) [2, 3] @ [3, 5] -> [2, 5] + + [[-1.4310, 0.3144, 0.1952], [[-0.1099, 0.7062, 0.6576, 1.3056, 0.3783], + [ 1.6719, -0.2581, -1.0243]] @ [ 0.9769, -0.6924, 0.4765, 1.1012, 0.3814] + [-1.4598, -0.5444, 0.5582, -2.0959, -0.0568]] + -> + [[ 0.1795, -1.3346, -0.6822, -1.9311, -0.4324], + [ 1.0593, 1.9171, 0.4047, 4.0454, 0.5921]] + + 2) exp([2, 5]) + exp([[ 0.1795, -1.3346, -0.6822, -1.9311, -0.4324], -> [[ 1.1966, 0.2633, 0.5055, 0.1450, 0.6489], + [ 1.0593, 1.9171, 0.4047, 4.0454, 0.5921]]) [ 2.8845, 6.8013, 1.4988, 57.1358, 1.8078]] + """ + # Each program computes one element C[row, col] using 2D tl.dot + row_matmul = tl.program_id(0) + col = tl.program_id(1) + + # Build small 2D grids so tl.dot sees [M,K] x [K,N] + offs_i = tl.arange(0, tl.constexpr(M))[:, None] # [M,1] (row axis) + offs_j = tl.arange(0, N)[None, :] # [1,N] (col axis) + offs_k = tl.arange(0, K) # [K] + + # A row: [M, K] + a_ptrs = A_ptr + (row_matmul + offs_i) * K + offs_k[None, :] + a_vals = tl.load(a_ptrs) # [M, K] + + # B column: [K, N] + b_ptrs = B_ptr + offs_k[:, None] * N + (col + offs_j) + b_vals = tl.load(b_ptrs) # [K, N] + + tbuff_ptrs = TBuff_ptr + (row_matmul + offs_i) * N + (col + offs_j) + + # Dot: [M, K] @ [K, N] -> [M, N] + acc_11 = tl.dot(a_vals, b_vals) # [M, N] + tl.store(tbuff_ptrs, acc_11) + + # Load Matrix [M/2, N] + sub_vec_id = al.sub_vec_id() + row_exp = row_matmul + sub_M * sub_vec_id + offs_exp_i = tl.arange(0, tl.constexpr(sub_M))[:, None] # [M/2, 1] (row axis) + tbuff_exp_ptrs = TBuff_ptr + (row_exp + offs_exp_i) * N + (col + offs_j) + acc_11_reload = tl.load(tbuff_exp_ptrs) + # Pointer grid for the single output element: shape [M/2, N] + c_ptrs = C_ptr + (row_exp + offs_exp_i) * N + (col + offs_j) + + # Store exp(acc) without scalar indexing + tl.store(c_ptrs, tl.exp(acc_11_reload)) + + +@pytest.mark.parametrize( + "dtype, ashape, bshape", + [ + # dtype, A-shape, B-shape + ["float32", (2, 3), (3, 5)], + ["float32", (2, 1), (1, 5)], + ], +) +def test_sub_vec_id_1to2(dtype, ashape, bshape): + """function: + A 1:2 demo using sub_vec_id. + 1. The matrix computation and the vector computation unit each have their own independent Scalar scheduler units, + deploying separately on cube core and vector core. + 2. Combine cube core and vector core in a certain ratio (1:2) + + For example, [2, 3] @ [3, 5] -> [2, 5] matrix matmul computation and matrix exp([2, 5]) computation + using sub_vec_id was used during the matrix exp. + """ + M, K = ashape + K2, N = bshape + assert K == K2, "Inner dimensions must match" + assert M % 2 == 0, "M dimensions must be divisible by 2" + sub_M = int(M / 2) + + # Generate input tensors + A = test_common.generate_tensor(ashape, dtype).npu() + B = test_common.generate_tensor(bshape, dtype).npu() + C = test_common.generate_tensor((M, N), dtype).npu() + TBuff = test_common.generate_tensor((M, N), dtype).npu() + + # Run + grid_matmul_exp = (1, ) # grid + triton_matmul_exp[grid_matmul_exp](A, B, C, TBuff, M, N, K, sub_M) + + # Reference result + C_ref = (A @ B).exp() + test_common.validate_cmp(dtype, C, C_ref) + + +if __name__ == "__main__": + test_sub_vec_id_1to2("float32", (2, 3), (3, 5)) diff --git a/third_party/ascend/unittest/pytest_ut/test_sub_vec_num.py b/third_party/ascend/unittest/pytest_ut/test_sub_vec_num.py new file mode 100644 index 0000000000..b46302836b --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_sub_vec_num.py @@ -0,0 +1,130 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import pytest +import test_common +import triton.language.extra.cann.extension as al + + +@triton.jit +def triton_matmul_exp( + A_ptr, + B_ptr, + C_ptr, + TBuff_ptr, + M: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, +): + """function: + 1) The matrix matmul + 2) The matrix exp computation + For example, + 1) [2, 3] @ [3, 5] -> [2, 5] + + [[-1.4310, 0.3144, 0.1952], [[-0.1099, 0.7062, 0.6576, 1.3056, 0.3783], + [ 1.6719, -0.2581, -1.0243]] @ [ 0.9769, -0.6924, 0.4765, 1.1012, 0.3814] + [-1.4598, -0.5444, 0.5582, -2.0959, -0.0568]] + -> + [[ 0.1795, -1.3346, -0.6822, -1.9311, -0.4324], + [ 1.0593, 1.9171, 0.4047, 4.0454, 0.5921]] + + 2) exp([2, 5]) + exp([[ 0.1795, -1.3346, -0.6822, -1.9311, -0.4324], -> [[ 1.1966, 0.2633, 0.5055, 0.1450, 0.6489], + [ 1.0593, 1.9171, 0.4047, 4.0454, 0.5921]]) [ 2.8845, 6.8013, 1.4988, 57.1358, 1.8078]] + """ + # Each program computes one element C[row, col] using 2D tl.dot + row_matmul = tl.program_id(0) + col = tl.program_id(1) + + # Build small 2D grids so tl.dot sees [M,K] x [K,N] + offs_i = tl.arange(0, tl.constexpr(M))[:, None] # [M,1] (row axis) + offs_j = tl.arange(0, N)[None, :] # [1,N] (col axis) + offs_k = tl.arange(0, K) # [K] + + # A row: [M, K] + a_ptrs = A_ptr + (row_matmul + offs_i) * K + offs_k[None, :] + a_vals = tl.load(a_ptrs) # [M, K] + + # B column: [K, N] + b_ptrs = B_ptr + offs_k[:, None] * N + (col + offs_j) + b_vals = tl.load(b_ptrs) # [K, N] + + tbuff_ptrs = TBuff_ptr + (row_matmul + offs_i) * N + (col + offs_j) + + # Dot: [M, K] @ [K, N] -> [M, N] + acc_11 = tl.dot(a_vals, b_vals) # [M, N] + tl.store(tbuff_ptrs, acc_11) + + # Load Matrix [M/2, N] + sub_vec_id = al.sub_vec_id() + row_exp = row_matmul + (M // al.sub_vec_num()) * sub_vec_id + offs_exp_i = tl.arange(0, M // al.sub_vec_num())[:, None] # [M/2, 1] (row axis) + tbuff_exp_ptrs = TBuff_ptr + (row_exp + offs_exp_i) * N + (col + offs_j) + acc_11_reload = tl.load(tbuff_exp_ptrs) + # Pointer grid for the single output element: shape [M/2, N] + c_ptrs = C_ptr + (row_exp + offs_exp_i) * N + (col + offs_j) + + # Store exp(acc) without scalar indexing + tl.store(c_ptrs, tl.exp(acc_11_reload)) + + +@pytest.mark.parametrize( + "dtype, ashape, bshape", + [ + # dtype, A-shape, B-shape + ["float32", (2, 3), (3, 5)], + ["float32", (2, 1), (1, 5)], + ], +) +def test_sub_vec_num_1to2(dtype, ashape, bshape): + """function: + A 1:2 demo using sub_vec_id. + 1. The matrix computation and the vector computation unit each have their own independent Scalar scheduler units, + deploying separately on cube core and vector core. + 2. Combine cube core and vector core in a certain ratio (1:2) + + For example, [2, 3] @ [3, 5] -> [2, 5] matrix matmul computation and matrix exp([2, 5]) computation + using sub_vec_id was used during the matrix exp. + """ + M, K = ashape + K2, N = bshape + assert K == K2, "Inner dimensions must match" + assert M % 2 == 0, "M dimensions must be divisible by 2" + + # Generate input tensors + A = test_common.generate_tensor(ashape, dtype).npu() + B = test_common.generate_tensor(bshape, dtype).npu() + C = test_common.generate_tensor((M, N), dtype).npu() + TBuff = test_common.generate_tensor((M, N), dtype).npu() + + # Run + grid_matmul_exp = (1, ) # grid + triton_matmul_exp[grid_matmul_exp](A, B, C, TBuff, M, N, K) + + # Reference result + C_ref = (A @ B).exp() + test_common.validate_cmp(dtype, C, C_ref) + + +if __name__ == "__main__": + test_sub_vec_num_1to2("float32", (2, 3), (3, 5)) diff --git a/third_party/ascend/unittest/pytest_ut/test_subview.py b/third_party/ascend/unittest/pytest_ut/test_subview.py new file mode 100644 index 0000000000..d5f6179f34 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_subview.py @@ -0,0 +1,89 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import os + +import pytest +import triton +import triton.language as tl +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +import triton.extension.buffer.language as bl +import triton.language.extra.cann.extension as al +from triton._C.libtriton import ir, buffer_ir +from triton._C.libtriton.ascend import ir as ascend_ir + +os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" + + +class Options: + num_warps = 4 + num_stages = 3 + num_ctas = 1 + cluster_dims = (1, 1, 1) + enable_fp_fusion = True + debug = False + + +def compile_kernel(kernel, signature, constants): + """Helper to compile a kernel to MLIR.""" + src = ASTSource(kernel, signature, constants) + context = ir.context() + ir.load_dialects(context) + buffer_ir.load_dialects(context) + ascend_ir.load_dialects(context) + module = ast_to_ttir(kernel, src, context, Options(), {"create_address_space": al.semantic.create_address_space}, + {}) + return str(module) + + +# ============== Kernel definitions ============== + + +@triton.jit +def test_subview_kernel1(XBLOCK: tl.constexpr): + # 1. Allocate a local buffer + src_buffer = bl.alloc(tl.float32, [XBLOCK, XBLOCK]) + + result_buffer = bl.subview(src_buffer, offsets=[1, 1], sizes=[XBLOCK - 2, XBLOCK - 2], strides=[1, 1]) + + +@triton.jit +def test_subview_kernel2(XBLOCK: tl.constexpr, offsets: tl.constexpr, sizes: tl.constexpr, strides: tl.constexpr): + # this statement has no effect, just to test the builder + bl.alloc(tl.float32, [XBLOCK]).subview([offsets], [sizes], [strides]) + + +# ============== Main for manual testing ============== + +if __name__ == "__main__": + print("=" * 60) + print("Test 1: test_subview_function") + print("=" * 60) + mlir = compile_kernel(test_subview_kernel1, {}, {"XBLOCK": 8}) + print(f"Generated MLIR ({len(mlir)} chars):\n") + print(mlir) + + print("\n" + "=" * 60) + print("Test 2: test_subview_constructor") + print("=" * 60) + mlir = compile_kernel(test_subview_kernel2, {}, {"XBLOCK": 32, "offsets": 1, "sizes": 30, "strides": 1}) + print(f"Generated MLIR ({len(mlir)} chars):\n") + print(mlir) diff --git a/third_party/ascend/unittest/pytest_ut/test_sum.py b/third_party/ascend/unittest/pytest_ut/test_sum.py new file mode 100644 index 0000000000..9e3ffeb76a --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_sum.py @@ -0,0 +1,124 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import pytest +import test_common +import time + + +@triton.jit +def sum_loop_high(in_ptr0, in_ptr1, in_ptr2, out_ptr0, rnumel, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr, + RBLOCK: tl.constexpr): + R = rnumel + X = xnumel + xoffset = tl.program_id(0) * XBLOCK + xbase = tl.arange(0, XBLOCK_SUB) + rbase = tl.arange(0, RBLOCK) + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + xindex = xoffset + xoffset_sub + xbase + x0 = xindex[None, :] + _tmp6 = tl.full([RBLOCK, XBLOCK_SUB], 0, tl.float32) + for roffset in range(0, rnumel, RBLOCK): + rindex = roffset + rbase + rmask = None + r1 = rindex[:, None] + tmp0 = tl.load(in_ptr0 + (X * r1 + (x0)), rmask) + tmp1 = tl.load(in_ptr1 + (X * r1 + (x0)), rmask) + tmp3 = tl.load(in_ptr2 + (X * r1 + (x0)), rmask) + tmp2 = tmp0 + tmp1 + tmp4 = tmp2 + tmp3 + _tmp6 = _tmp6 + tmp4 + tmp6 = tl.sum(_tmp6, 0) + tl.store(out_ptr0 + (xindex), tmp6, None) + + +@triton.jit +def sum_loop_low(in_ptr0, in_ptr1, in_ptr2, out_ptr0, xnumel, ynumel, XBLOCK: tl.constexpr, RBLOCK: tl.constexpr): + X = xnumel + Y = ynumel + xoffset = tl.program_id(0) * XBLOCK + xindex = xoffset + tl.arange(0, XBLOCK) + + x0 = xindex[:, None] + rbase = tl.arange(0, RBLOCK) + _tmp6 = tl.full([XBLOCK, RBLOCK], 0, tl.float32) + for roffset in range(0, ynumel, RBLOCK): + rindex = roffset + rbase + rmask = None + r1 = rindex[None, :] + tmp0 = tl.load(in_ptr0 + (r1 + (Y * x0)), rmask) + tmp1 = tl.load(in_ptr1 + (r1 + (Y * x0)), rmask) + tmp3 = tl.load(in_ptr2 + (r1 + (Y * x0)), rmask) + tmp2 = tmp0 + tmp1 + tmp4 = tmp2 + tmp3 + _tmp6 = _tmp6 + tmp4 + tmp6 = tl.sum(_tmp6, 1) + + tl.store(out_ptr0 + (xindex), tmp6, None) + + +def foo(a, b, c): + y = a + b + c + y = y.sum(0) + return y + + +def bar(a, b, c): + y = a + b + c + y = y.sum(1) + return y + + +@pytest.mark.parametrize('param_list', [ + ['float32', (64, 8192), 1, 2, 256, 16], +]) +def test_case_1(param_list): + dtype, shape, ncore, XB, YB, ZB = param_list + a = test_common.generate_tensor(shape, dtype).npu() + b = test_common.generate_tensor(shape, dtype).npu() + c = test_common.generate_tensor(shape, dtype).npu() + value = torch.empty_strided((a.shape[0], ), (1, )).npu() + + std_low_ret = bar(a, b, c) + print(f"std_low_ret = {std_low_ret[0:8]}") + XBLOCK = 64 + RBLOCK = 32 + NBLOCKS = a.shape[0] // XBLOCK + sum_loop_low[NBLOCKS, 1, 1](a, b, c, value, a.shape[0], a.shape[1], XBLOCK, RBLOCK) + triton_low_ret = value + print(f"triton_low_ret = {triton_low_ret[0:8]}") + torch.testing.assert_close(std_low_ret, triton_low_ret, rtol=1e-3, atol=1e-3) + + std_ret2 = foo(a, b, c) + print(f"std_ret2 = {std_ret2[0:8]}") + NBLOCKS = 32 + XBLOCK = a.shape[1] // NBLOCKS + XBLOCK_SUB = min(64, max(XBLOCK // 2, 32)) + RBLOCK = 64 + + value2 = torch.empty_strided((a.shape[1], ), (1, )).npu() + sum_loop_high[NBLOCKS, 1, 1](a, b, c, value2, a.shape[0], a.shape[1], XBLOCK, XBLOCK_SUB, RBLOCK) + triton_ret2 = value2 + print(f"triton_ret2 = {triton_ret2[0:8]}") + torch.testing.assert_close(std_ret2, triton_ret2, rtol=1e-3, atol=1e-3) diff --git a/third_party/ascend/unittest/pytest_ut/test_sum_dim0.py b/third_party/ascend/unittest/pytest_ut/test_sum_dim0.py new file mode 100644 index 0000000000..93b107f839 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_sum_dim0.py @@ -0,0 +1,133 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time + +import torch +import torch_npu +import test_common + + +def standard_sum(x0, dim, dtype): + res = torch.sum(x0, dim, dtype=dtype) + return res + + +@triton.jit +def triton_sum_dim0(in_ptr0, out_ptr0, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, MNUMEL) + nblk_idx = tl.arange(0, NNUMEL) + + mmask = mblk_idx < M + nmask = nblk_idx < N + + mask = (mmask[:, None]) & (nmask[None, :]) + + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + + x = tl.load(in_ptr0 + idx, mask=mask, other=0) + + ret = tl.sum(x, 0) + + tl.store(out_ptr0 + nblk_idx, ret, mask=nmask) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + # (torch.bfloat16,'bfloat16'), TODO: waiting for supporting or testing + (torch.int8, 'int8'), + # (torch.int16,'int16'), TODO: waiting for supporting or testing + # (torch.int32,'int32'), TODO: waiting for supporting or testing + # (torch.int64,'int64'), TODO: waiting for supporting or testing +] + +# if shape axis = 32/256 , then actual shape = axis/element_size() +shapes = [ + (57, 3, 64, 16), + (57, -32, 64, 32), + (57, 37, 64, 64), + (57, -256, 64, 256), + (57, 263, 64, 512), + (64, 3, 64, 16), + (64, -32, 64, 32), + (64, 37, 64, 64), + (64, -256, 64, 256), + (64, 263, 64, 512), + (3, 3, 8, 8), + (-32, 3, 32, 8), + (37, 3, 64, 8), + (-256, 3, 256, 8), + (263, 3, 512, 8), + (3, 1, 8, 8), + (-32, 1, 32, 8), + (37, 1, 64, 8), + (-256, 1, 256, 8), + (263, 1, 512, 8), +] + +map_for_64_t = {37: (31, 32), 263: (107, 128)} +map_for_32_t = {263: (137, 256)} + + +# @pytest.mark.parametrize('dtype, sigtype',[(torch.float32,'float32'),]) +@pytest.mark.parametrize('M, N, MNUMEL, NNUMEL', [ + (57, 3, 64, 16), + (64, -32, 64, 32), + (37, 3, 64, 8), + (263, 1, 512, 8), + (-256, 3, 256, 8), +]) +@pytest.mark.parametrize('dtype, sigtype', types) +# @pytest.mark.parametrize('M, N, MNUMEL, NNUMEL',shapes) +def test_sum_dim0(dtype, sigtype, M, N, MNUMEL, NNUMEL): + + M = (-M) // torch.tensor(0, dtype=dtype).element_size() if M < 0 else M + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + M = map_for_64_t[M][0] if M in map_for_64_t else M + MNUMEL = map_for_64_t[M][1] if M in map_for_64_t else MNUMEL + N = map_for_64_t[N][0] if N in map_for_64_t else N + NNUMEL = map_for_64_t[N][1] if N in map_for_64_t else NNUMEL + + elif sigtype == 'float32' or sigtype == 'bfloat16' or sigtype == 'int32': + M = map_for_32_t[M][0] if M in map_for_32_t else M + MNUMEL = map_for_32_t[M][1] if M in map_for_32_t else MNUMEL + N = map_for_32_t[N][0] if N in map_for_32_t else N + NNUMEL = map_for_32_t[N][1] if N in map_for_32_t else NNUMEL + + print(f"sum : ({M}, {N}) {dtype} {sigtype}") + x0 = test_common.generate_tensor(shape=(M, N), dtype=sigtype) + + ans = standard_sum(x0, 0, dtype) + + x0 = x0.npu() + print(ans) + + output = torch.zeros((N, ), dtype=dtype).npu() + triton_sum_dim0[1, 1, 1](x0, output, M=M, N=N, MNUMEL=MNUMEL, NNUMEL=NNUMEL, debug=True) + print(output) + + test_common.validate_cmp(sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_sum_dim1.py b/third_party/ascend/unittest/pytest_ut/test_sum_dim1.py new file mode 100644 index 0000000000..f7a3e27d5f --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_sum_dim1.py @@ -0,0 +1,133 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import time + +import torch +import torch_npu +import test_common + + +def standard_sum(x0, dim, dtype): + res = torch.sum(x0, dim, dtype=dtype) + return res + + +@triton.jit +def triton_sum_dim1(in_ptr0, out_ptr0, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr, NNUMEL: tl.constexpr): + mblk_idx = tl.arange(0, MNUMEL) + nblk_idx = tl.arange(0, NNUMEL) + + mmask = mblk_idx < M + nmask = nblk_idx < N + + mask = (mmask[:, None]) & (nmask[None, :]) + + idx = mblk_idx[:, None] * N + nblk_idx[None, :] + + x = tl.load(in_ptr0 + idx, mask=mask, other=0) + + ret = tl.sum(x, 1) + + tl.store(out_ptr0 + mblk_idx, ret, mask=mmask) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + # (torch.bfloat16,'bfloat16'), waiting for supporting or testing + (torch.int8, 'int8'), + # (torch.int16,'int16'), waiting for supporting or testing + # (torch.int32,'int32'), waiting for supporting or testing + # (torch.int64,'int64'), waiting for supporting or testing +] + +# if shape axis = 32/256 , then actual shape = axis/element_size() +shapes = [ + (57, 3, 64, 16), + (57, -32, 64, 32), + (57, 37, 64, 64), + (57, -256, 64, 256), + (57, 263, 64, 512), + (64, 3, 64, 16), + (64, -32, 64, 32), + (64, 37, 64, 64), + (64, -256, 64, 256), + (64, 263, 64, 512), + (3, 3, 8, 8), + (-32, 3, 32, 8), + (37, 3, 64, 8), + (-256, 3, 256, 8), + (263, 3, 512, 8), + (3, 1, 8, 8), + (-32, 1, 32, 8), + (37, 1, 64, 8), + (-256, 1, 256, 8), + (263, 1, 512, 8), +] + +map_for_64_t = {37: (31, 32), 263: (107, 128)} +map_for_32_t = {263: (137, 256)} + + +@pytest.mark.parametrize('M, N, MNUMEL, NNUMEL', [ + (57, 3, 64, 16), + (64, -32, 64, 32), + (37, 3, 64, 8), + (263, 1, 512, 8), + (-256, 3, 256, 8), +]) +# @pytest.mark.parametrize('M, N',[(263,3),(-256,3)]) +@pytest.mark.parametrize('dtype, sigtype', types) +# @pytest.mark.parametrize('M, N, MNUMEL, NNUMEL',shapes) +def test_sum_dim1(dtype, sigtype, M, N, MNUMEL, NNUMEL): + + M = (-M) // torch.tensor(0, dtype=dtype).element_size() if M < 0 else M + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == 'int64': + M = map_for_64_t[M][0] if M in map_for_64_t else M + MNUMEL = map_for_64_t[M][1] if M in map_for_64_t else MNUMEL + N = map_for_64_t[N][0] if N in map_for_64_t else N + NNUMEL = map_for_64_t[N][1] if N in map_for_64_t else NNUMEL + + elif sigtype == 'float32' or sigtype == 'bfloat16' or sigtype == 'int32': + M = map_for_32_t[M][0] if M in map_for_32_t else M + MNUMEL = map_for_32_t[M][1] if M in map_for_32_t else MNUMEL + N = map_for_32_t[N][0] if N in map_for_32_t else N + NNUMEL = map_for_32_t[N][1] if N in map_for_32_t else NNUMEL + + print(f"sum : ({M}, {N}) {dtype} {sigtype}") + x0 = test_common.generate_tensor(shape=(M, N), dtype=sigtype) + + ans = standard_sum(x0, 1, dtype) + + x0 = x0.npu() + print(ans) + + output = torch.zeros((M, ), dtype=dtype).npu() + triton_sum_dim1[1, 1, 1](x0, output, M=M, N=N, MNUMEL=MNUMEL, NNUMEL=NNUMEL, debug=True) + print(output) + + test_common.validate_cmp(sigtype, output, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_sum_vector.py b/third_party/ascend/unittest/pytest_ut/test_sum_vector.py new file mode 100644 index 0000000000..efa3dcb70c --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_sum_vector.py @@ -0,0 +1,82 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +from triton.runtime.libentry import libentry +import pytest +from test_common import generate_tensor, validate_cmp, _32bit_dtypes, _16bit_dtypes + + +def torch_func(x0): + return torch.sum(x0) + + +@pytest.mark.parametrize("dtype", _32bit_dtypes) +@pytest.mark.parametrize("shape", [(1, ), (3, ), (8, ), (37, ), (64, ), (781, )]) +def test_sum(dtype, shape): + + @libentry() + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, XBLOCK: tl.constexpr): + idx = tl.arange(0, XBLOCK) + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tl.sum(tmp0) + tl.store(out_ptr0 + idx, tmp1) + + def triton_func(x0): + out = x0[0] + triton_kernel[1, 1, 1](out, x0, x0.numel()) + return out + + x0 = generate_tensor(shape=shape, dtype=dtype).npu() + torch_ref = torch_func(x0) + triton_cal = triton_func(x0) + validate_cmp(dtype, torch_ref, triton_cal) + + +@triton.jit +def _reduce_combine(a, b): + return a + b + + +@pytest.mark.parametrize("dtype", _32bit_dtypes) +@pytest.mark.parametrize("shape", [(1, ), (3, ), (8, ), (37, ), (64, ), (781, )]) +def test_reduce_sum(dtype, shape): + + @libentry() + @triton.jit + def triton_kernel(out_ptr0, in_ptr0, XBLOCK: tl.constexpr): + idx = tl.arange(0, XBLOCK) + tmp0 = tl.load(in_ptr0 + idx) + tmp1 = tl.reduce(tmp0, 0, _reduce_combine) + tl.store(out_ptr0 + idx, tmp1) + + def triton_func(x0): + out = x0[0] + triton_kernel[1, 1, 1](out, x0, x0.numel()) + return out + + x0 = generate_tensor(shape=shape, dtype=dtype).npu() + torch_ref = torch_func(x0) + triton_cal = triton_func(x0) + validate_cmp(dtype, torch_ref, triton_cal) diff --git a/third_party/ascend/unittest/pytest_ut/test_swap.py b/third_party/ascend/unittest/pytest_ut/test_swap.py new file mode 100644 index 0000000000..589def92e0 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_swap.py @@ -0,0 +1,58 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import triton +import triton.language as tl +import torch_npu +import pytest + + +@triton.jit +def swap_kernel(x_ptr, # *Pointer* to first inout vector. + y_ptr, # *Pointer* to second inout vector. + BLOCK_SIZE: tl.constexpr, # Number of elements each program should process. + # NOTE: `constexpr` so it can be used as a shape value. + ): + pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + x = tl.load(x_ptr + offsets) + y = tl.load(y_ptr + offsets) + tl.store(x_ptr + offsets, y) + tl.store(y_ptr + offsets, x) + + +def swap(x: torch.Tensor, y: torch.Tensor, size): + n_elements = x.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), ) + swap_kernel[grid](x, y, BLOCK_SIZE=size) + + +@pytest.mark.parametrize('shape', [(1, ), (3, ), (4, ), (7, ), (8, ), (11, ), (16, ), (512, ), (1024, )]) +def test(shape): + x = torch.rand(shape).npu() + y = torch.rand(shape).npu() + assert not torch.equal(x, y) + x_ = x.clone() + y_ = y.clone() + swap(x, y, shape[0]) + torch.testing.assert_close(x, y_, rtol=1e-04, atol=1e-04, equal_nan=True) + torch.testing.assert_close(y, x_, rtol=1e-04, atol=1e-04, equal_nan=True) diff --git a/third_party/ascend/unittest/pytest_ut/test_swiglu.py b/third_party/ascend/unittest/pytest_ut/test_swiglu.py new file mode 100644 index 0000000000..42d18ddb9d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_swiglu.py @@ -0,0 +1,100 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +# from .utils import calculate_settings +def standard_binary(e, g): + ee = e.to(torch.float32) + f = ee * torch.sigmoid(ee) + h = (f * g).to(g.dtype) + return h + + +@triton.jit +def _fg_kernel( + e, + g, + h, + n_elements, + BLOCK_SIZE: tl.constexpr, +): + block_idx = tl.program_id(0) + offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + + e_row = tl.load(e + offsets, mask=mask, other=0).to(tl.float32) + g_row = tl.load(g + offsets, mask=mask, other=0) #.to(tl.float32) + + # f = e * sigmoid(e) + f_row = e_row * tl.sigmoid(e_row) # e_row / (1 + tl.exp(-e_row)) + # f_row = f_row.to(g_row.dtype) # bf16 should always cast to fp32 when calculating + # h = f * g + h_row = (f_row * g_row).to(g_row.dtype) + # Store h + tl.store(h + offsets, h_row, mask=mask) + + +pass + + +def swiglu_fg_kernel(e, g): + batch, seq_len, hd = e.shape + n_elements = e.numel() + h = torch.empty((batch, seq_len, hd), dtype=e.dtype, device="npu") + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + kk = _fg_kernel[grid]( + e, + g, + h, + n_elements, + BLOCK_SIZE=1024, + ) + print(kk.asm['ttir']) + return h + + +pass + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 128, 128)], + ['float16', (2, 128, 128)], + ['bfloat16', (2, 128, 128)], +]) +def test_case(param_list): + dtype, size = param_list + torch.manual_seed(0) + x = torch.rand(size, device='npu', dtype=eval('torch.' + dtype)) + y = torch.rand(size, device='npu', dtype=eval('torch.' + dtype)) + std_ret = standard_binary(x, y) + print(f"std_ret= {std_ret}") + ret = swiglu_fg_kernel(x, y) + print(f"ret= {ret}") + test_common.validate_cmp(dtype, std_ret, ret) + + +pass diff --git a/third_party/ascend/unittest/pytest_ut/test_swizzle2d.py b/third_party/ascend/unittest/pytest_ut/test_swizzle2d.py new file mode 100644 index 0000000000..49f2dbab4b --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_swizzle2d.py @@ -0,0 +1,55 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + i = tl.arange(0, 4) + j = tl.arange(0, 4) + xx, yy = tl.swizzle2d(i, j, size_i=4, size_j=4, size_g=2) + + tl.store(output_ptr + tl.arange(0, 4), xx) + tl.store(x_ptr + tl.arange(0, 4), yy) + + +@pytest.mark.parametrize('param_list', [ + ['int32', (2, 256, 16), 1, 2, 256, 16], +]) +def test_case(param_list): + dtype, shape, ncore, XB, YB, ZB = param_list + x = test_common.generate_tensor((4, ), dtype).npu() + a = torch.tensor([[0, 2, 1, 3], [4, 6, 5, 7], [8, 10, 9, 11], [12, 14, 13, 15]], dtype=eval('torch.' + dtype)).npu() + output = torch.randint(1, (4, ), dtype=eval('torch.' + dtype)).npu() + fn_npu_[ncore, 1, 1](output, x, XB, YB, ZB) + print(f"output={output}") + triton_ret = output[:, None] * 4 + x[None, :] + print(f"triton_ret={triton_ret}") + torch.testing.assert_close(triton_ret, a) diff --git a/third_party/ascend/unittest/pytest_ut/test_sync_block.py b/third_party/ascend/unittest/pytest_ut/test_sync_block.py new file mode 100644 index 0000000000..8d1df6d243 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_sync_block.py @@ -0,0 +1,99 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +import torch + +import triton +import triton.language as tl +import pytest +import test_common + +import triton.language.extra.cann.extension as extension + +pipe = extension.PIPE + +# eg: pytest -v test_matmul_exp.py::test_matmul_exp +############################# + + +@triton.jit +def triton_matmul_exp(A_ptr, B_ptr, C_ptr, TBuff_ptr, M, N, K: tl.constexpr): + # Each program computes one element C[row, col] using 2D tl.dot + row = tl.program_id(0) + col = tl.program_id(1) + + # Build small 2D grids so tl.dot sees [1,K] x [K,1] + offs_i = tl.arange(0, 1)[:, None] # [1,1] (row axis) + offs_j = tl.arange(0, 1)[None, :] # [1,1] (col axis) + offs_k = tl.arange(0, K) # [K] + + # A row: [1, K] + a_ptrs = A_ptr + (row + offs_i) * K + offs_k[None, :] + a_vals = tl.load(a_ptrs) # [1, K] + + # B column: [K, 1] + b_ptrs = B_ptr + offs_k[:, None] * N + (col + offs_j) + b_vals = tl.load(b_ptrs) # [K, 1] + + tbuff_ptrs = TBuff_ptr + (row + offs_i) * N + (col + offs_j) + + # Dot: [1, K] @ [K, 1] -> [1, 1] + acc_11 = tl.dot(a_vals, b_vals) # [1, 1] + tl.store(tbuff_ptrs, acc_11) + + extension.sync_block_set("cube", "vector", 5, pipe.PIPE_MTE1, pipe.PIPE_MTE3) + extension.sync_block_wait("cube", "vector", 5, pipe.PIPE_MTE1, pipe.PIPE_MTE3) + + acc_11_reload = tl.load(tbuff_ptrs) + # Pointer grid for the single output element: shape [1,1] + c_ptrs = C_ptr + (row + offs_i) * N + (col + offs_j) + + # Store exp(acc) without scalar indexing + tl.store(c_ptrs, tl.exp(acc_11_reload)) + + +@pytest.mark.parametrize('dtype, ashape, bshape', [ + # dtype, A-shape, B-shape + ['float32', (4, 4), (4, 4)], + ['float32', (2, 3), (3, 5)], +]) +def test_matmul_exp(dtype, ashape, bshape): + M, K = ashape + K2, N = bshape + assert K == K2, "Inner dimensions must match" + + # generate input tensors + A = test_common.generate_tensor(ashape, dtype).npu() + B = test_common.generate_tensor(bshape, dtype).npu() + C = test_common.generate_tensor((M, N), dtype).npu() + TBuff = test_common.generate_tensor((M, N), dtype).npu() + + # run kernel + grid = (M, N) # one program per output element + triton_matmul_exp[grid](A, B, C, TBuff, M, N, K) + + # reference result + C_ref = (A @ B).exp() + + # compare + torch.testing.assert_close(C_ref, C, rtol=3e-2, atol=3e-2, equal_nan=True) + + +if __name__ == "__main__": + test_matmul_exp('float32', (4, 4), (4, 4)) diff --git a/third_party/ascend/unittest/pytest_ut/test_sync_block_all.py b/third_party/ascend/unittest/pytest_ut/test_sync_block_all.py new file mode 100755 index 0000000000..a418854321 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_sync_block_all.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python3 + +import os +import pytest +import triton +import triton.language as tl +import triton.language.extra.cann.extension as al +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +from triton._C.libtriton import ir +from triton._C.libtriton.ascend import ir as ascend_ir + +os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" + + +class Options: + num_warps = 4 + num_stages = 3 + num_ctas = 1 + cluster_dims = (1, 1, 1) + enable_fp_fusion = True + debug = False + + +def compile_kernel(kernel, signature, constants): + """Helper to compile a kernel to MLIR.""" + src = ASTSource(kernel, signature, constants) + context = ir.context() + ir.load_dialects(context) + ascend_ir.load_dialects(context) + module = ast_to_ttir(kernel, src, context, Options(), {}, {}) + return str(module) + + +# ============== Kernel definitions ============== + + +@triton.jit +def test_sync_block_all(): + al.sync_block_all("all_cube", 8) + al.sync_block_all("all_vector", 9) + al.sync_block_all("all", 10) + al.sync_block_all("all_sub_vector", 11) + + +# ============== Main for manual testing ============== + +if __name__ == "__main__": + print("=" * 60) + print("Test 1: test_sync_block_all") + print("=" * 60) + mlir = compile_kernel(test_sync_block_all, {}, {}) + print(f"✅ Generated MLIR ({len(mlir)} chars):\n") + print(mlir) diff --git a/third_party/ascend/unittest/pytest_ut/test_tan.py b/third_party/ascend/unittest/pytest_ut/test_tan.py new file mode 100644 index 0000000000..13f63f51ab --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_tan.py @@ -0,0 +1,102 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import test_common + +import torch +import torch_npu + +import triton.language.extra.cann.libdevice as libdevice + + +def standard_unary(x0, dtype): + res = torch.tan(x0) + return res + + +def standard_binary(x0, y0, dtype): + res = x0 + y0 + return res + + +@triton.jit +def triton_elementwise_unary(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + ret = libdevice.tan(x) + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +@triton.jit +def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + y = tl.load(in_ptr1 + idx_block, mask=idx_block < N) + ret = x + y + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +types = [ + (torch.float32, 'float32'), + (torch.float16, 'float16'), + # (torch.bfloat16, 'bfloat16'), + # (torch.int8, 'int8'), + # (torch.int16, 'int16'), + # (torch.int32, 'int32'), + # (torch.int64, 'int64'), +] + +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.skip(reason="randomly failed accuracy test") +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('N,NUMEL', shapes) +def test_elementwsie_common(dtype, sigtype, N, NUMEL): + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == "int64": + N = map_for_64_t[N] if N in map_for_64_t else N + + print(f"elementwise : ({N},) {dtype} {sigtype}") + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype) + + ans = standard_unary(x0, dtype) + x0 = x0.npu() + print(ans) + + out = torch.zeros((N, ), dtype=dtype).npu() + triton_elementwise_unary[1, 1, 1](x0, out, N=N, NUMEL=NUMEL, debug=True) + print(out) + + test_common.validate_cmp(sigtype, out, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_tanh.py b/third_party/ascend/unittest/pytest_ut/test_tanh.py new file mode 100644 index 0000000000..b80eac61a3 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_tanh.py @@ -0,0 +1,59 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + +import torch +import torch_npu + + +def torch_tanh(x0): + return torch.tanh(x0) + + +@triton.jit +def triton_tanh(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = libdevice.tanh(tmp0) + tl.store(out_ptr0 + (x0), tmp1, None) + + +@pytest.mark.parametrize('param_list', [ + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['bfloat16', (2, 4096, 8), 2, 32768, 1024], + ['float32', (2, 4096, 8), 2, 32768, 1024], +]) +def test_tanh(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_tanh(x0) + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + triton_tanh[ncore, 1, 1](x0, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_template.py b/third_party/ascend/unittest/pytest_ut/test_template.py new file mode 100644 index 0000000000..2b1c5a2262 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_template.py @@ -0,0 +1,58 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl + +import time +import torch +import torch_npu +import test_common + +NBLOCKS = 1 +X_SIZE = tl.constexpr(4) +Y_SIZE = tl.constexpr(64) +Z_SIZE = tl.constexpr(32) +NUMEL = X_SIZE * Y_SIZE * Z_SIZE + + +def fn(input): + output = input.reshape((X_SIZE, Y_SIZE, Z_SIZE)).permute((1, 0, 2)).reshape((X_SIZE * Y_SIZE * Z_SIZE)) + return output + + +@triton.jit +def fn_kernel(output_ptr, input_ptr): + col_offsets = tl.arange(0, X_SIZE * Y_SIZE * Z_SIZE) + input_local = tl.load(input_ptr + col_offsets) + input_local = input_local.reshape((X_SIZE, Y_SIZE, Z_SIZE)).permute((1, 0, 2)).reshape((X_SIZE * Y_SIZE * Z_SIZE)) + tl.store(output_ptr + col_offsets, input_local) + + +def test_cases(): + input = torch.randn(NUMEL, dtype=torch.float16).npu() + output = torch.randn(NUMEL, dtype=torch.float16).npu() + output2 = torch.randn(NUMEL, dtype=torch.float16).npu() + fn_kernel[1, 1, 1](output, input) + output2 = fn(input) + test_common.validate_cmp('float16', output, output2) + print("data validation passed") diff --git a/third_party/ascend/unittest/pytest_ut/test_tensor_descriptor.py b/third_party/ascend/unittest/pytest_ut/test_tensor_descriptor.py new file mode 100644 index 0000000000..5a62c96a46 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_tensor_descriptor.py @@ -0,0 +1,225 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import math +import pytest +import torch +import triton +import triton.language as tl +import test_common + + +@pytest.mark.parametrize("dtype", ['float32', 'float16', 'bfloat16', 'int32', 'int64', 'int16', 'int8']) +@pytest.mark.parametrize("M_BLOCK,N_BLOCK", [(2, 16), (8, 16)]) +def test_tensor_descriptor_load_store(dtype, M_BLOCK, N_BLOCK): + + @triton.jit + def kernel(out_ptr, a_ptr, M, N, M_BLOCK: tl.constexpr, N_BLOCK: tl.constexpr): + in_desc = tl.make_tensor_descriptor( + a_ptr, + shape=[M, N], + strides=[N, 1], + block_shape=[M_BLOCK, N_BLOCK], + ) + out_desc = tl.make_tensor_descriptor( + out_ptr, + shape=[M, N], + strides=[N, 1], + block_shape=[M_BLOCK, N_BLOCK], + ) + moffset = tl.program_id(0) * M_BLOCK + noffset = tl.program_id(1) * N_BLOCK + block = in_desc.load([moffset, noffset]) + out_desc.store([moffset, noffset], block) + + M, N = M_BLOCK * 2, N_BLOCK * 2 + inp = test_common.generate_tensor((M, N), dtype).npu() + out = inp.new_empty((M, N)) + + grid_m = M // M_BLOCK + grid_n = N // N_BLOCK + + kernel[(grid_m, grid_n)](out, inp, M, N, M_BLOCK, N_BLOCK) + torch.testing.assert_close(inp, out) + + +@pytest.mark.parametrize("dtype", ['float32', 'float16', 'bfloat16', 'int32', 'int64', 'int16', 'int8']) +def test_tensor_descriptor_load_store3d(dtype): + + @triton.jit + def kernel(out_ptr, a_ptr, M, N, K, stride_m, stride_n, stride_k, M_BLOCK: tl.constexpr, N_BLOCK: tl.constexpr, + K_BLOCK: tl.constexpr): + in_desc = tl.make_tensor_descriptor( + a_ptr, + shape=[M, N, K], + strides=[stride_m, stride_n, stride_k], + block_shape=[M_BLOCK, N_BLOCK, K_BLOCK], + ) + out_desc = tl.make_tensor_descriptor( + out_ptr, + shape=[M, N, K], + strides=[stride_m, stride_n, stride_k], + block_shape=[M_BLOCK, N_BLOCK, K_BLOCK], + ) + moffset = tl.program_id(0) * M_BLOCK + noffset = tl.program_id(1) * N_BLOCK + koffset = tl.program_id(2) * K_BLOCK + block = in_desc.load([moffset, noffset, koffset]) + out_desc.store([moffset, noffset, koffset], block) + + M, N, K = 8, 16, 32 + inp = test_common.generate_tensor((M, N, K), dtype).npu() + out = inp.new_empty((M, N, K)) + + M_BLOCK = 2 + N_BLOCK = 4 + + # 自动调整 K_BLOCK，保证最后一维 block 至少 16 字节 + dtype = getattr(inp, "dtype", None) + itemsize = torch.tensor([], dtype=inp.dtype).element_size() + min_k_block = max(16 // itemsize, 1) + K_BLOCK = max(8, min_k_block) + + grid_m = M // M_BLOCK + grid_n = N // N_BLOCK + grid_k = K // K_BLOCK + + kernel[(grid_m, grid_n, grid_k)](out, inp, *inp.shape, *out.stride(), M_BLOCK, N_BLOCK, K_BLOCK) + torch.testing.assert_close(inp.reshape(M * N * K), out.reshape(M * N * K)) + + +# Exercise the functional load/store builtins once to ensure they map through. +@pytest.mark.parametrize("dtype", ["float32", "uint8"]) +def test_tensor_descriptor_functional_interface(dtype): + """Copies an entire tensor blockwise using the descriptor builtins.""" + + @triton.jit + def kernel(out_ptr, a_ptr, M, N, M_BLOCK: tl.constexpr, N_BLOCK: tl.constexpr): + in_desc = tl.make_tensor_descriptor( + a_ptr, + shape=[M, N], + strides=[N, 1], + block_shape=[M_BLOCK, N_BLOCK], + ) + out_desc = tl.make_tensor_descriptor( + out_ptr, + shape=[M, N], + strides=[N, 1], + block_shape=[M_BLOCK, N_BLOCK], + ) + moffset = tl.program_id(0) * M_BLOCK + noffset = tl.program_id(1) * N_BLOCK + block = tl.load_tensor_descriptor(in_desc, [moffset, noffset]) + tl.store_tensor_descriptor(out_desc, [moffset, noffset], block) + + M, N = 32, 128 + inp = test_common.generate_tensor((M, N), dtype).npu() + + M_BLOCK = 8 + N_BLOCK = 32 + out = inp.new_empty((M, N)) + + grid_m = M // M_BLOCK + grid_n = N // N_BLOCK + + kernel[(grid_m, grid_n)](out, inp, M, N, M_BLOCK, N_BLOCK) + torch.testing.assert_close(inp, out) + + +@pytest.mark.parametrize("dtype_str", ["int32"]) +@pytest.mark.parametrize("shape", [(128, 2, 4), (64, 2, 4), (32, 2, 4), (2, 4, 32), (2, 4, 2)]) +@pytest.mark.parametrize("axis", [0, 1, 2]) +@pytest.mark.parametrize("device", ["npu"]) +def test_reduce_max(dtype_str, shape, axis, device): + + @triton.jit + def kernel( + In, + Out, + in_shape1: tl.constexpr, + in_shape2: tl.constexpr, + in_shape3: tl.constexpr, + ou_shape1: tl.constexpr, + ou_shape2: tl.constexpr, + axis: tl.constexpr, + ): + in_desc = tl.make_tensor_descriptor( + base=In, + shape=[in_shape1 * in_shape2 * in_shape3], + strides=[1], + block_shape=[in_shape1 * in_shape2 * in_shape3], + ) + out_desc = tl.make_tensor_descriptor( + base=Out, + shape=[ou_shape1 * ou_shape2], + strides=[1], + block_shape=[ou_shape1 * ou_shape2], + ) + val = in_desc.load([0]).reshape(in_shape1, in_shape2, in_shape3) + output = tl.max(val, axis=axis) + out_desc.store([0], output.reshape(out_desc.block_shape)) + + inp = torch.arange(math.prod(shape), dtype=getattr(torch, dtype_str), device=device).reshape(shape) + expected, indices = torch.max(inp.to(torch.int64), dim=axis) + expected = expected.to(torch.int32) + actual = torch.zeros(expected.shape, dtype=getattr(torch, dtype_str), device=device) + kernel[(1, )](inp, actual, *shape, *expected.shape, axis=axis) + assert torch.equal(expected, actual) + + +@pytest.mark.parametrize("dtype", ['float32', 'float16', 'bfloat16']) +@pytest.mark.parametrize("padding", ["zero", "nan"]) +def test_tensor_descriptor_padding(dtype, padding): + + @triton.jit + def device_tma_load(in_ptr, out_ptr, IM, IN, YM, YN, M_BLOCK: tl.constexpr, N_BLOCK: tl.constexpr, + padding: tl.constexpr): + x_desc = tl.make_tensor_descriptor(in_ptr, shape=[IM, IN], strides=[IN, 1], block_shape=[M_BLOCK, N_BLOCK], + padding_option=padding) + out_desc = tl.make_tensor_descriptor(out_ptr, shape=[YM, YN], strides=[YN, 1], block_shape=[M_BLOCK, N_BLOCK]) + + moffset = tl.program_id(0) * M_BLOCK + noffset = tl.program_id(1) * N_BLOCK + + value = x_desc.load([moffset, noffset]) + + out_desc.store([moffset, noffset], value) + + def alloc_fn(size: int, alignment: float, stream: float): + return torch.ones(size, device="npu", dtype=torch.float32) + + triton.set_allocator(alloc_fn) + + IM, IN = 48, 48 + OM, ON = 64, 64 + M_BLOCK = 32 + N_BLOCK = 32 + torch_dtype = getattr(torch, dtype) + input_tensor = torch.arange(IM * IN, device="npu", dtype=torch_dtype).reshape(IM, IN) + out_device_tma = torch.zeros((OM, ON), device="npu", dtype=torch_dtype) + grid = (triton.cdiv(OM, M_BLOCK), triton.cdiv(ON, N_BLOCK)) + device_tma_load[grid](input_tensor, out_device_tma, IM, IN, OM, ON, M_BLOCK, N_BLOCK, padding) + expected = torch.zeros((OM, ON), device="npu", dtype=torch_dtype) + expected[0:IN, 0:IM] = input_tensor + if padding == "nan": + expected[:, IN:ON] = float('nan') + expected[IM:OM, :] = float('nan') + + torch.testing.assert_close(expected, out_device_tma, equal_nan=True) diff --git a/third_party/ascend/unittest/pytest_ut/test_to_buffer.py b/third_party/ascend/unittest/pytest_ut/test_to_buffer.py new file mode 100644 index 0000000000..2053bc7944 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_to_buffer.py @@ -0,0 +1,54 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch + +import triton +import triton.language as tl +from triton.compiler import ASTSource +import triton.extension.buffer.language as bl +import triton.language.extra.cann.extension as al + +target = triton.runtime.driver.active.get_current_target() + + +@triton.jit +def to_buffer(): + a = tl.full((32, 2, 4), 0, dtype=tl.int64) + a_buf = bl.to_buffer(a) + b = tl.full((32, 2, 4), 0, dtype=tl.int64) + b_buf = bl.to_buffer(b, al.ascend_address_space.UB) + c = tl.full((32, 2, 4), 0, dtype=tl.int64) + c_buf = bl.to_buffer(c, al.ascend_address_space.L1) + d = tl.full((32, 2, 4), 0, dtype=tl.int64) + d_buf = bl.to_buffer(d, al.ascend_address_space.L0A) + e = tl.full((32, 2, 4), 0, dtype=tl.int64) + e_buf = bl.to_buffer(e, al.ascend_address_space.L0B) + f = tl.full((32, 2, 4), 0, dtype=tl.int64) + f_buf = bl.to_buffer(f, al.ascend_address_space.L0C) + + +def test_to_buffer(): + src = ASTSource( + fn=to_buffer, + constexprs={}, + signature={}, + ) + triton.compile(src=src, target=target) diff --git a/third_party/ascend/unittest/pytest_ut/test_to_tensor.py b/third_party/ascend/unittest/pytest_ut/test_to_tensor.py new file mode 100644 index 0000000000..d75bed7603 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_to_tensor.py @@ -0,0 +1,75 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import os + +import triton +import triton.language as tl +from triton.compiler.compiler import ASTSource +from triton.compiler.code_generator import ast_to_ttir +import triton.extension.buffer.language as bl +import triton.language.extra.cann.extension as al +from triton._C.libtriton import ir, buffer_ir +from triton._C.libtriton.ascend import ir as ascend_ir + +os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" + + +class Options: + num_warps = 4 + num_stages = 3 + num_ctas = 1 + cluster_dims = (1, 1, 1) + enable_fp_fusion = True + debug = False + + +def compile_kernel(kernel, signature, constants): + """Helper to compile a kernel to MLIR.""" + src = ASTSource(kernel, signature, constants) + context = ir.context() + ir.load_dialects(context) + buffer_ir.load_dialects(context) + ascend_ir.load_dialects(context) + module = ast_to_ttir(kernel, src, context, Options(), {"create_address_space": al.semantic.create_address_space}, + {}) + return str(module) + + +# ============== Kernel definitions ============== + + +@triton.jit +def kernel_func(XBLOCK: tl.constexpr): + buffer1 = bl.alloc(tl.float32, [XBLOCK]) + buffer1.to_tensor(writable=True) + buffer2 = bl.alloc(tl.float32, [XBLOCK]) + bl.to_tensor(buffer2, writable=True) + + +# ============== Main for manual testing ============== + +if __name__ == "__main__": + print("=" * 60) + print("Test 1: Nested Scopes") + print("=" * 60) + mlir = compile_kernel(kernel_func, {}, {"XBLOCK": 256}) + print(f"✅ Generated MLIR ({len(mlir)} chars):\n") + print(mlir) diff --git a/third_party/ascend/unittest/pytest_ut/test_top2gating_argmax.py b/third_party/ascend/unittest/pytest_ut/test_top2gating_argmax.py new file mode 100644 index 0000000000..38f31b00bd --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_top2gating_argmax.py @@ -0,0 +1,97 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import pytest +import test_common + + +@triton.jit +def promote_to_tensor(x): + # Addition promotes to tensor for us + return x + tl.zeros((1, ), tl.int1) + + +@triton.jit +def is_floating(x): + return promote_to_tensor(x).dtype.is_floating() + + +@triton.jit +def maximum_with_index(a_value, a_index, b_value, b_index): + mask = a_value > b_value + equal = a_value == b_value + if is_floating(a_value): + a_isnan = a_value != a_value + b_isnan = b_value != b_value + mask |= a_isnan and not b_isnan + # Consider NaNs as equal + equal |= a_isnan and b_isnan + + # Prefer lowest index if values are equal + mask |= equal & (a_index < b_index) + return tl.where(mask, a_value, b_value), tl.where(mask, a_index, b_index) + + +@triton.jit +def max_with_index(value, index, dim): + return tl.reduce((value, index), dim, maximum_with_index) + + +@triton.jit +def triton_4(in_ptr2, in_ptr4, out_ptr10, x0_numel, r1_numel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + RBLOCK: tl.constexpr = 4 + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = XBLOCK // XBLOCK_SUB + base2 = tl.arange(0, RBLOCK) + loops2: tl.constexpr = r1_numel // RBLOCK + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1[:, None] + r1 = base2[None, :] + tmp15 = tl.load(in_ptr2 + (r1 + (8 * x0)), None) + tmp25 = tl.load(in_ptr4 + (r1 + (4 * x0)), None) + tmp21 = tmp15.to(tl.float32) + tmp26 = tmp25 * tmp21 + tmp27 = tl.reshape(tmp26, [XBLOCK_SUB, RBLOCK]) + tmp31 = tl.broadcast_to(r1.reshape(1, RBLOCK), tmp27.shape) + _, tmp30_tmp = max_with_index(tmp27, tmp31, 1) + tmp30 = tmp30_tmp.reshape(XBLOCK_SUB, 1) + tmp43 = tmp30.to(tl.int32) + tmp46 = tmp43 + tl.store(out_ptr10 + (x0), tmp46, None) + + +def test_max_with_index_dim0(): + mask = torch.randint(low=0, high=2, size=(512, 2, 4), dtype=torch.int32).npu() + weights = torch.randn((512, 4), device='npu', dtype=torch.float32) + buf32 = torch.randint(low=0, high=8, size=(512, ), dtype=torch.int32).npu() + XBLOCK = 32 + XBLOCK_SUB = 32 + triton_4[16, 1, 1](mask, weights, buf32, 512, 4, XBLOCK, XBLOCK_SUB) + + _, first_idx = torch.max(weights * mask[:, 0, :], dim=1) + print(f"first_idx: {first_idx[0:8]}") + print(f"triton_idx: {buf32[0:8]}") + + assert torch.all(torch.eq(first_idx, buf32.to(first_idx.dtype))) diff --git a/third_party/ascend/unittest/pytest_ut/test_topk.py b/third_party/ascend/unittest/pytest_ut/test_topk.py new file mode 100644 index 0000000000..b532418e61 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_topk.py @@ -0,0 +1,74 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import pytest + +import triton +import triton.language as tl + +import torch +import torch_npu +import test_common + + +@pytest.mark.interpreter +@pytest.mark.parametrize("shape", [ + (16, 64), + (16, 8, 8), + (2, 8, 8, 8), +]) +@pytest.mark.parametrize("k", [4]) +@pytest.mark.parametrize("dtype", ["float32"]) +def test_topk_nd(shape, k, dtype): + + numel = 1 + for d in shape: + numel *= d + base = torch.arange(numel, dtype=torch.float32).view(shape) / 10.0 + if dtype == "float16": + base = base.half() + elif dtype == "int32": + base = torch.arange(numel, dtype=torch.int32).view(shape) + x = base.npu() + + y = torch.topk(x.cpu(), k=k, dim=-1).values + + M = int(torch.tensor(shape[:-1]).prod().item()) if len(shape) > 1 else 1 + N = shape[-1] + + x_2d = x.view(M, N) + z = torch.empty((M, k), dtype=x_2d.dtype, device=x_2d.device) + + @triton.jit + def topk_kernel_nd(X, stride_xm, Z, stride_zm, M: tl.constexpr, N: tl.constexpr, k: tl.constexpr): + offs_m = tl.arange(0, M) + offs_x_n = tl.arange(0, N) + offs_z_n = tl.arange(0, k) + offs_x = offs_m[:, None] * stride_xm + offs_x_n[None, :] + x_val = tl.load(X + offs_x) + z_val = tl.topk(x_val, k) + offs_z = offs_m[:, None] * stride_zm + offs_z_n[None, :] + tl.store(Z + offs_z, z_val) + + topk_kernel_nd[(1, )](x_2d, x_2d.stride(0), z, z.stride(0), M, N, k, num_warps=8) + + z_view = z.view(*shape[:-1], k) if len(shape) > 1 else z.view(k) + test_common.validate_cmp(dtype, z_view.cpu(), y) diff --git a/third_party/ascend/unittest/pytest_ut/test_trans_3d.py b/third_party/ascend/unittest/pytest_ut/test_trans_3d.py new file mode 100644 index 0000000000..fde7857774 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_trans_3d.py @@ -0,0 +1,80 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import logging +import math + +import pytest +import test_common +import torch + +import triton +import triton.language as tl + + +@triton.jit +def fn_npu_102(output_ptr, x_ptr, YB: tl.constexpr, ZB: tl.constexpr, KB: tl.constexpr): + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + kidx = tl.arange(0, KB) + idx = yidx[:, None, None] * ZB * KB + zidx[None, :, None] * KB + kidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.trans(X, 1, 0, 2) + + oidx = (zidx[:, None, None] * YB * KB + yidx[None, :, None] * KB + kidx[None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_021(output_ptr, x_ptr, YB: tl.constexpr, ZB: tl.constexpr, KB: tl.constexpr): + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + kidx = tl.arange(0, KB) + idx = yidx[:, None, None] * ZB * KB + zidx[None, :, None] * KB + kidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.trans(X, (0, 2, 1)) + + oidx = (yidx[:, None, None] * ZB * KB + kidx[None, :, None] * ZB + zidx[None, None, :]) + + tl.store(output_ptr + oidx, ret) + + +@pytest.mark.parametrize("shape", [(23, 5, 31)]) +@pytest.mark.parametrize("dtype", ["float32"]) +def test_permute_3d(shape, dtype): + logging.debug(f"dtype:{dtype} shape:{shape}") + + data_type = eval("torch." + dtype) + x = torch.randint(low=0, high=2, size=shape, dtype=data_type).npu() + + triton_res = torch.empty((shape[1], shape[0], shape[2]), dtype=data_type).npu() + torch_res = torch.permute(x, (1, 0, 2)) + fn_npu_102[1, 1, 1](triton_res, x, shape[0], shape[1], shape[2]) + test_common.validate_cmp(dtype, triton_res, torch_res) + + triton_res = torch.empty((shape[0], shape[2], shape[1]), dtype=data_type).npu() + torch_res = torch.permute(x, (0, 2, 1)) + fn_npu_021[1, 1, 1](triton_res, x, shape[0], shape[1], shape[2]) + test_common.validate_cmp(dtype, triton_res, torch_res) diff --git a/third_party/ascend/unittest/pytest_ut/test_triton_eq.py b/third_party/ascend/unittest/pytest_ut/test_triton_eq.py new file mode 100644 index 0000000000..6bb1e09096 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_triton_eq.py @@ -0,0 +1,60 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +def torch_pointwise(x0, x1): + res = x0 == x1 + return res + + +@triton.jit +def triton_test(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop1 in range(loops1): + x0_prime = offset + (loop1 * XBLOCK_SUB) + base1 + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.load(in_ptr1 + (x0), None) + tmp2 = tmp0 == tmp1 + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], +]) +def test_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_pointwise(x0, x1) + y_cal = torch.zeros(shape, dtype=eval('torch.' + 'bool')).npu() + triton_test[ncore, 1, 1](x0, x1, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_triton_le.py b/third_party/ascend/unittest/pytest_ut/test_triton_le.py new file mode 100644 index 0000000000..21e7f72dcd --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_triton_le.py @@ -0,0 +1,60 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +def torch_pointwise(x0, x1): + res = x0 <= x1 + return res + + +@triton.jit +def triton_le(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop1 in range(loops1): + x0_prime = offset + (loop1 * XBLOCK_SUB) + base1 + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.load(in_ptr1 + (x0), None) + tmp2 = tmp0 <= tmp1 + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], +]) +def test_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_pointwise(x0, x1) + y_cal = torch.zeros(shape, dtype=eval('torch.' + 'bool')).npu() + triton_le[ncore, 1, 1](x0, x1, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_triton_lt.py b/third_party/ascend/unittest/pytest_ut/test_triton_lt.py new file mode 100644 index 0000000000..1789bbf9b8 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_triton_lt.py @@ -0,0 +1,60 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +def torch_pointwise(x0, x1): + res = x0 < x1 + return res + + +@triton.jit +def triton_lt(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop1 in range(loops1): + x0_prime = offset + (loop1 * XBLOCK_SUB) + base1 + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.load(in_ptr1 + (x0), None) + tmp2 = tmp0 < tmp1 + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], +]) +def test_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_pointwise(x0, x1) + y_cal = torch.zeros(shape, dtype=eval('torch.' + 'bool')).npu() + triton_lt[ncore, 1, 1](x0, x1, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_triton_neq.py b/third_party/ascend/unittest/pytest_ut/test_triton_neq.py new file mode 100644 index 0000000000..a2fbde4242 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_triton_neq.py @@ -0,0 +1,60 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +def torch_pointwise(x0, x1): + res = x0 != x1 + return res + + +@triton.jit +def triton_neq(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + for loop1 in range(loops1): + x0_prime = offset + (loop1 * XBLOCK_SUB) + base1 + x0 = offset + (loop1 * XBLOCK_SUB) + base1 + tmp0 = tl.load(in_ptr0 + (x0), None) + tmp1 = tl.load(in_ptr1 + (x0), None) + tmp2 = tmp0 != tmp1 + tl.store(out_ptr0 + (x0), tmp2, None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 4096, 8), 2, 32768, 1024], + ['float16', (2, 4096, 8), 2, 32768, 1024], + ['int8', (2, 4096, 8), 2, 32768, 1024], +]) +def test_case(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_pointwise(x0, x1) + y_cal = torch.zeros(shape, dtype=eval('torch.' + 'bool')).npu() + triton_neq[ncore, 1, 1](x0, x1, y_cal, xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_triton_unified_attention.py b/third_party/ascend/unittest/pytest_ut/test_triton_unified_attention.py new file mode 100644 index 0000000000..4e3b8a0ad3 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_triton_unified_attention.py @@ -0,0 +1,469 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Optional + +import pytest +import torch +import triton +import triton.language as tl + +NUM_HEADS = [(8, 2), (16, 2)] +HEAD_SIZES = [128] +# only support 32 currently +BLOCK_SIZES = [32] + +DTYPES = [torch.float16] +QDTYPES = [None] +# one value large enough to test overflow in index calculation. +# one value small enough to test the schema op check +NUM_BLOCKS = [32768, 2048] + + +@triton.jit +def cdiv_fn(x, y): + return (x + y - 1) // y + + +@triton.jit +def apply_softcap(S, x): + Sdiv = S / x + p1 = tl.exp(Sdiv) + p2 = tl.exp(-Sdiv) + return x * (p1 - p2) / (p1 + p2) + + +@triton.jit +def kernel_unified_attention_2d(output_ptr, # [num_tokens, num_query_heads, head_size] + query_ptr, # [num_tokens, num_query_heads, head_size] + key_cache_ptr, # [num_blks, blk_size, num_kv_heads, head_size] + value_cache_ptr, # [num_blks, blk_size, num_kv_heads, head_size] + block_tables_ptr, # [num_seqs, max_num_blocks_per_seq] + seq_lens_ptr, # [num_seqs] + alibi_slopes_ptr, # [num_query_heads] + scale, # float32 + k_scale, # float32 + v_scale, # float32 + softcap, # float32 + num_query_heads: tl.constexpr, # int + num_queries_per_kv: tl.constexpr, # int + block_table_stride: tl.int64, # int + query_stride_0: tl.int64, # int + query_stride_1: tl.int64, # int, should be equal to head_size + output_stride_0: tl.int64, # int + output_stride_1: tl.int64, # int, should be equal to head_size + BLOCK_SIZE: tl.constexpr, # int + HEAD_SIZE: tl.constexpr, # int + HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2 + USE_ALIBI_SLOPES: tl.constexpr, # bool + USE_SOFTCAP: tl.constexpr, # bool + SLIDING_WINDOW: tl.constexpr, # int + stride_k_cache_0: tl.int64, # int + stride_k_cache_1: tl.int64, # int + stride_k_cache_2: tl.int64, # int + stride_k_cache_3: tl.constexpr, # int + stride_v_cache_0: tl.int64, # int + stride_v_cache_1: tl.int64, # int + stride_v_cache_2: tl.int64, # int + stride_v_cache_3: tl.constexpr, # int + query_start_len_ptr, # [num_seqs+1] + BLOCK_Q: tl.constexpr, # int + num_seqs: tl.int32, BLOCK_M: tl.constexpr, # int + ): + + q_block_global_idx = tl.program_id(0) + kv_head_idx = tl.program_id(1) + + left: tl.int32 = 0 + right = num_seqs + while left < right: + mid = (left + right) // 2 + mid_val = tl.load(query_start_len_ptr + mid) // BLOCK_Q + mid + if mid_val <= q_block_global_idx: + left = mid + 1 + else: + right = mid + + seq_idx = left - 1 + q_block_start_idx = tl.load(query_start_len_ptr + seq_idx) // BLOCK_Q + seq_idx + + q_block_local_idx = q_block_global_idx - q_block_start_idx + + cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx) + cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx + 1) + + cur_batch_query_len = cur_batch_in_all_stop_index \ + - cur_batch_in_all_start_index + + if q_block_local_idx * BLOCK_Q >= cur_batch_query_len: + return + + offs_m = tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, HEAD_SIZE_PADDED) + query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv + + query_offset_0 = cur_batch_in_all_start_index + query_pos + query_offset_1 = kv_head_idx * num_queries_per_kv + \ + offs_m % num_queries_per_kv + query_offset = (query_offset_0[:, None] * query_stride_0 + query_offset_1[:, None] * query_stride_1 + + offs_d[None, :]) + + dim_mask = tl.where(offs_d < HEAD_SIZE, 1, 0).to(tl.int1) + query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1) + query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1) + + Q = tl.load( + query_ptr + query_offset, + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + other=0.0, + ) + + block_table_offset = seq_idx * block_table_stride + + M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) + L = tl.full([BLOCK_M], 1.0, dtype=tl.float32) + acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32) + + # sequence len for this particular sequence + seq_len = tl.load(seq_lens_ptr + seq_idx) + + # context length for this particular sequences + context_len = seq_len - cur_batch_query_len + + # alibi slope for this head + if USE_ALIBI_SLOPES: + alibi_slope = tl.load(alibi_slopes_ptr + query_offset_1, mask=query_mask_1, other=0.0) + + num_blocks = cdiv_fn(seq_len, BLOCK_SIZE) + + # iterate through tiles + for j in range(0, num_blocks): + + physical_block_idx = tl.load(block_tables_ptr + block_table_offset + j) + + offs_n = tl.arange(0, BLOCK_SIZE) + + v_offset = (physical_block_idx * stride_v_cache_0 + kv_head_idx * stride_v_cache_2 + + offs_d[None, :] * stride_v_cache_3 + offs_n[:, None] * stride_v_cache_1) + + k_offset = (physical_block_idx * stride_k_cache_0 + kv_head_idx * stride_k_cache_2 + + offs_d[:, None] * stride_k_cache_3 + offs_n[None, :] * stride_k_cache_1) + + K_load = tl.load(key_cache_ptr + k_offset, mask=dim_mask[:, None], other=0.0) + + if K_load.dtype.is_fp8(): + if Q.dtype.is_fp8(): + K = K_load + else: + K = (K_load.to(tl.float32) * tl.load(k_scale)).to(Q.dtype) + else: + K = K_load + + V_load = tl.load(value_cache_ptr + v_offset, mask=dim_mask[None, :], other=0.0) + + if V_load.dtype.is_fp8(): + if Q.dtype.is_fp8(): + V = V_load + else: + V = (V_load.to(tl.float32) * tl.load(v_scale)).to(Q.dtype) + else: + V = V_load + + seq_offset = j * BLOCK_SIZE + offs_n + + seq_mask = seq_offset[None, :] < context_len + query_pos[:, None] + 1 + + S = tl.zeros(shape=(BLOCK_M, BLOCK_SIZE), dtype=tl.float32) + + S += scale * tl.dot(Q, K) + + if USE_SOFTCAP: + S = apply_softcap(S, softcap) + + S = tl.where(query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, S, float("-inf")) + + if SLIDING_WINDOW > 0: + S = tl.where((context_len + query_pos[:, None] - seq_offset) < SLIDING_WINDOW, S, float("-inf")) + + if USE_ALIBI_SLOPES: + S += alibi_slope[:, None] * (seq_offset - context_len) + + # compute running maximum + m_j = tl.maximum(M, tl.max(S, axis=1)) + # For sliding window there's a chance the max is -inf due to masking of + # the entire row. In this case we need to set m_j 0 to avoid NaN + m_j = tl.where(m_j > float("-inf"), m_j, 0.0) + + P = tl.exp(S - m_j[:, None]) + + l_j = tl.sum(P, axis=1) + + alpha = tl.exp(M - m_j) + + acc = acc * alpha[:, None] + + # update constants + L = L * alpha + l_j + M = m_j + + acc += tl.dot(P.to(V.dtype), V) + + # epilogue + acc = acc / L[:, None] + + output_offset = (query_offset_0[:, None] * output_stride_0 + query_offset_1[:, None] * output_stride_1 + + offs_d[None, :]) + + tl.store( + output_ptr + output_offset, + acc, + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + ) + + +def unified_attention( + q, + k, + v, + out, + cu_seqlens_q, + max_seqlen_q, + seqused_k, + max_seqlen_k, + softmax_scale, + causal, + window_size, + block_table, + softcap, + q_descale, + k_descale, + v_descale, + alibi_slopes=None, +): + assert causal, "Only causal attention is supported" + assert q_descale is None, "Q scales not supported" + + block_size = v.shape[1] + assert q.element_size() >= 2 or block_size >= 32, \ + "Block size must be at least 32 for fp8" + + use_alibi_slopes = alibi_slopes is not None + + block_size = v.shape[1] + num_seqs = len(seqused_k) + num_query_heads = q.shape[1] + num_kv_heads = k.shape[2] + num_queries_per_kv = num_query_heads // num_kv_heads + head_size = q.shape[2] + + BLOCK_M = 16 + BLOCK_Q = BLOCK_M // num_queries_per_kv + + # Ideally we would launch with kernel with: + # \sum_i[ceil(query_len[i] / BLOCK_Q)] blocks. + # However, it is slow to realize the query_lens on cpu. + # Instead we use upper-bound: + # \sum_i[ceil(query_len[i] / BLOCK_Q)] + # <= \sum_i[floor(query_len[i] / BLOCK_Q) + 1] + # = \sum_i[floor(query_len[i] / BLOCK_Q)] + num_seqs + # <= floor(\sum_i(query_len[i]) / BLOCK_Q) + num_seqs + # = floor(q.shape[0] / BLOCK_Q) + num_seqs + total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs + + kernel_unified_attention_2d[( + total_num_q_blocks, + num_kv_heads, + )]( + output_ptr=out, + query_ptr=q, + key_cache_ptr=k, + value_cache_ptr=v, + block_tables_ptr=block_table, + seq_lens_ptr=seqused_k, + alibi_slopes_ptr=alibi_slopes, + scale=softmax_scale, + k_scale=k_descale, + v_scale=v_descale, + softcap=softcap, + num_query_heads=num_query_heads, + num_queries_per_kv=num_queries_per_kv, + block_table_stride=block_table.stride(0), + query_stride_0=q.stride(0), + query_stride_1=q.stride(1), + output_stride_0=out.stride(0), + output_stride_1=out.stride(1), + BLOCK_SIZE=block_size, + HEAD_SIZE=head_size, + HEAD_SIZE_PADDED=triton.next_power_of_2(head_size), + USE_ALIBI_SLOPES=use_alibi_slopes, + USE_SOFTCAP=(softcap > 0), + SLIDING_WINDOW=(1 + window_size[0]), + stride_k_cache_0=k.stride(0), + stride_k_cache_1=k.stride(1), + stride_k_cache_2=k.stride(2), + stride_k_cache_3=k.stride(3), + stride_v_cache_0=v.stride(0), + stride_v_cache_1=v.stride(1), + stride_v_cache_2=v.stride(2), + stride_v_cache_3=v.stride(3), + query_start_len_ptr=cu_seqlens_q, + BLOCK_Q=BLOCK_Q, + num_seqs=num_seqs, + BLOCK_M=BLOCK_M, + ) + + +def ref_paged_attn( + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + query_lens: list[int], + kv_lens: list[int], + block_tables: torch.Tensor, + scale: float, + sliding_window: Optional[int] = None, + soft_cap: Optional[float] = None, +) -> torch.Tensor: + num_seqs = len(query_lens) + block_tables = block_tables.cpu().numpy() + _, block_size, num_kv_heads, head_size = key_cache.shape + + outputs: list[torch.Tensor] = [] + start_idx = 0 + for i in range(num_seqs): + query_len = query_lens[i] + kv_len = kv_lens[i] + q = query[start_idx:start_idx + query_len] + q *= scale + + num_kv_blocks = (kv_len + block_size - 1) // block_size + block_indices = block_tables[i, :num_kv_blocks] + + k = key_cache[block_indices].view(-1, num_kv_heads, head_size) + k = k[:kv_len] + v = value_cache[block_indices].view(-1, num_kv_heads, head_size) + v = v[:kv_len] + + if q.shape[1] != k.shape[1]: + k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1) + v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1) + attn = torch.einsum("qhd,khd->hqk", q, k).float() + empty_mask = torch.ones(query_len, kv_len) + mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool() + if sliding_window is not None: + sliding_window_mask = torch.triu(empty_mask, + diagonal=kv_len - (query_len + sliding_window) + 1).bool().logical_not() + mask |= sliding_window_mask + if soft_cap is not None and soft_cap > 0: + attn = soft_cap * torch.tanh(attn / soft_cap) + attn.masked_fill_(mask, float("-inf")) + attn = torch.softmax(attn, dim=-1).to(v.dtype) + out = torch.einsum("hqk,khd->qhd", attn, v) + + outputs.append(out) + start_idx += query_len + + return torch.cat(outputs, dim=0) + + +@pytest.mark.perf(repeat=17) +@pytest.mark.parametrize("seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]]) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize("head_size", HEAD_SIZES) +@pytest.mark.parametrize("block_size", BLOCK_SIZES) +@pytest.mark.parametrize("sliding_window", [None]) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("soft_cap", [None]) +@pytest.mark.parametrize("num_blocks", NUM_BLOCKS) +@pytest.mark.parametrize("q_dtype", QDTYPES) +@torch.inference_mode() +def test_triton_unified_attn( + seq_lens: list[tuple[int, int]], + num_heads: tuple[int, int], + head_size: int, + sliding_window: Optional[int], + dtype: torch.dtype, + block_size: int, + soft_cap: Optional[float], + num_blocks: int, + q_dtype: Optional[torch.dtype], +) -> None: + torch.set_default_device("npu") + + if q_dtype is not None and q_dtype.itemsize < 2 and block_size < 32: + pytest.skip("block size must be at least 32 for fp8") + + num_seqs = len(seq_lens) + query_lens = [x[0] for x in seq_lens] + kv_lens = [x[1] for x in seq_lens] + num_query_heads = num_heads[0] + num_kv_heads = num_heads[1] + assert num_query_heads % num_kv_heads == 0 + max_query_len = max(query_lens) + max_kv_len = max(kv_lens) + window_size = ((sliding_window - 1, 0) if sliding_window is not None else (-1, -1)) + scale = head_size**-0.5 + + query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype) + key_cache = torch.randn(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype) + value_cache = torch.randn_like(key_cache) + cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(dim=0, dtype=torch.int32) + kv_lens = torch.tensor(kv_lens, dtype=torch.int32) + + max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size + block_tables = torch.randint(0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32) + + output = torch.empty_like(query) + + maybe_quantized_query = query + maybe_quantized_key_cache = key_cache + maybe_quantized_value_cache = value_cache + q_descale = None + k_descale = None + v_descale = None + if q_dtype is not None: + # QKV are drawn from N(0, 1): no need for a fp8 scaling factor + maybe_quantized_query = query.to(q_dtype) + maybe_quantized_key_cache = key_cache.to(q_dtype) + maybe_quantized_value_cache = value_cache.to(q_dtype) + + scale_shape = (num_seqs, num_kv_heads) + q_descale = None # Not yet supported + k_descale = torch.rand(scale_shape, dtype=torch.float32) + v_descale = torch.rand(scale_shape, dtype=torch.float32) + + unified_attention( + q=maybe_quantized_query, + k=maybe_quantized_key_cache, + v=maybe_quantized_value_cache, + out=output, + cu_seqlens_q=cu_query_lens, + seqused_k=kv_lens, + max_seqlen_q=max_query_len, + max_seqlen_k=max_kv_len, + softmax_scale=scale, + causal=True, + window_size=window_size, + block_table=block_tables, + softcap=soft_cap if soft_cap is not None else 0, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + ) + + ref_output = ref_paged_attn( + query=query, + key_cache=key_cache, + value_cache=value_cache, + query_lens=query_lens, + kv_lens=kv_lens, + block_tables=block_tables, + scale=scale, + sliding_window=sliding_window, + soft_cap=soft_cap, + ) + atol, rtol = 1.5e-2, 1e-2 + if q_dtype is not None: + atol, rtol = 1.5e-1, 1.5e-1 + torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol), \ + f"{torch.max(torch.abs(output - ref_output))}" diff --git a/third_party/ascend/unittest/pytest_ut/test_trunc.py b/third_party/ascend/unittest/pytest_ut/test_trunc.py new file mode 100644 index 0000000000..ec02d678b1 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_trunc.py @@ -0,0 +1,68 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import torch +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + + +@triton.jit +def trunc_kernel(x_ptr, y_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + + mask = offsets < n_elements + + x = tl.load(x_ptr + offsets, mask=mask) + + y = libdevice.trunc(x) + + tl.store(y_ptr + offsets, y, mask=mask) + + +@pytest.mark.parametrize('shape', [ + (12, 16), +]) +@pytest.mark.parametrize('dtype', ['float32']) +def test_cases(shape, dtype): + n_elements = shape[0] * shape[1] + x = test_common.generate_tensor(shape, dtype).npu() + + # Make sure to include some edge cases. + x[0, 0] = 0.0 + x[0, 1] = 3.14 + x[0, 2] = -2.71 + x[0, 3] = 5.0 + x[0, 4] = -3.0 + + y = torch.empty_like(x) + + BLOCK_SIZE = 192 + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + + trunc_kernel[grid](x, y, n_elements, BLOCK_SIZE=BLOCK_SIZE) + + expected = torch.trunc(x) + + torch.testing.assert_close(y, expected, rtol=1e-3, atol=1e-3) diff --git a/third_party/ascend/unittest/pytest_ut/test_umulhi.py b/third_party/ascend/unittest/pytest_ut/test_umulhi.py new file mode 100644 index 0000000000..3f4617bdeb --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_umulhi.py @@ -0,0 +1,62 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import numpy as np +import triton +import triton.language as tl +from numpy.random import RandomState + + +# inp the two 32 bit signed integers. +@triton.jit +def umulhi_kernel(X, Y, Z, N: tl.constexpr): + offs = tl.arange(0, N) + x = tl.load(X + offs) + y = tl.load(Y + offs) + z = tl.umulhi(x, y) + tl.store(Z + tl.arange(0, N), z) + + +# accuracy reference +def umulhi32(a, b): + a_64 = a.astype(np.int64) + b_64 = b.astype(np.int64) + product_64 = a_64 * b_64 + # get the high part + result_high_32 = product_64 >> 32 + return result_high_32.astype(np.int32) + + +def test_umulhi(): + N = 128 + x = torch.randint(low=0, high=2000, size=(N, ), dtype=torch.int32) + y = torch.randint(low=0, high=2000, size=(N, ), dtype=torch.int32) + xx = x.npu() + yy = y.npu() + z_tri = torch.zeros(size=(N, ), dtype=torch.int32).npu() + umulhi_kernel[(1, )](xx, yy, z_tri, N=N) + + xxx = x.cpu().numpy() + yyy = y.cpu().numpy() + z_ref = umulhi32(xxx, yyy) + z_ref1 = torch.from_numpy(z_ref).npu() + torch.equal(z_tri, z_ref1) diff --git a/third_party/ascend/unittest/pytest_ut/test_unlign_max_with_index.py b/third_party/ascend/unittest/pytest_ut/test_unlign_max_with_index.py new file mode 100644 index 0000000000..74e96e5e45 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_unlign_max_with_index.py @@ -0,0 +1,115 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import pytest +import test_common + + +@triton.jit +def promote_to_tensor(x): + # Addition promotes to tensor for us + return x + tl.zeros((1, ), tl.int1) + + +@triton.jit +def is_floating(x): + return promote_to_tensor(x).dtype.is_floating() + + +@triton.jit +def maximum_with_index(a_value, a_index, b_value, b_index): + mask = a_value > b_value + equal = a_value == b_value + if is_floating(a_value): + a_isnan = a_value != a_value + b_isnan = b_value != b_value + mask |= a_isnan and not b_isnan + # Consider NaNs as equal + equal |= a_isnan and b_isnan + + # Prefer lowest index if values are equal + mask |= equal & (a_index < b_index) + return tl.where(mask, a_value, b_value), tl.where(mask, a_index, b_index) + + +@triton.jit +def max_with_index(value, index, dim): + return tl.reduce((value, index), dim, maximum_with_index) + + +def torch_max_pr_index(x0): + res = torch.max(x0, 0) + return res + + +# x0_numel 128, r1_numel 65, 128, 128, 32 +@triton.jit +def triton_kernel(in_ptr0, in_ptr1, out_ptr0, out_ptr1, x0_numel, r1_numel, XBLOCK: tl.constexpr, + XBLOCK_SUB: tl.constexpr, RBLOCK: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + base2 = tl.arange(0, RBLOCK) + loops2: tl.constexpr = (r1_numel + RBLOCK - 1) // RBLOCK + for loop1 in range(loops1): + x0 = offset + (loop1 * XBLOCK_SUB) + base1[:, None] + xmask = x0 < x0_numel + for loop2 in range(loops2): + r1 = loop2 * RBLOCK + base2[None, :] + rmask = r1 < r1_numel + tmp0 = tl.load(in_ptr0 + (r1 + (65 * x0)), rmask & xmask, other=0) + tmp2 = tl.load(in_ptr1 + (r1 + (65 * x0)), rmask & xmask, other=0) + tmp11 = tl.reshape(tmp0, [XBLOCK_SUB, RBLOCK]) + tmp22 = tl.reshape(tmp2, [XBLOCK_SUB, RBLOCK]) + tmp4, tmp5 = max_with_index(tmp11, tmp22, 0) + tl.store(out_ptr0 + r1, tmp4[None, :], None) + tl.store(out_ptr1 + r1, tmp5[None, :], None) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (128, 65), 1, 128, 65, 32], +]) +def test_max_with_index_dim0(param_list): + dtype, shape, ncore, XB, YB, RBLOCK = param_list + import math + numel = math.prod(shape) + xblock = numel // shape[-1] // ncore + assert (ncore * xblock * shape[-1] == numel) + + b = test_common.generate_tensor(shape, dtype).npu() + idx = torch.zeros((shape), device='npu', dtype=torch.int32) + c, d = torch.max(b, dim=0) + + print(f"std_ret = {c[0:4]}") + + ret = torch.empty((YB), device='npu', dtype=torch.float32) + ret1 = torch.zeros((YB), device='npu', dtype=torch.int32) + triton_kernel[ncore, 1, 1](b, idx, ret, ret1, shape[0], shape[1], shape[0], shape[0], RBLOCK) + print(f"triton_ret = {ret[0:4]}") + d = d.to(torch.int32) + ret1 = ret1.to(torch.int32) + print(f"d = {d[0:4]}") + print(f"triton_ret1 = {ret1[0:4]}") + assert torch.allclose(c, ret, rtol=1e-03, atol=1e-03, equal_nan=True) + assert torch.allclose(d, ret1, rtol=1e-03, atol=1e-03, equal_nan=True) diff --git a/third_party/ascend/unittest/pytest_ut/test_unlign_sum.py b/third_party/ascend/unittest/pytest_ut/test_unlign_sum.py new file mode 100644 index 0000000000..31a34f74df --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_unlign_sum.py @@ -0,0 +1,64 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import torch +import torch_npu +import triton +import triton.language as tl +import triton +import triton.language as tl + +# [128,65] -> [128,128] +# [128,128] -> [128,65] mask 作用 +# [128,65] -> [128,64] [128,1] mask 作用 + + +@triton.jit +def triton_unlign(in_ptr0, out_ptr0, x0_numel, r1_numel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr, + RBLOCK: tl.constexpr): + offset = tl.program_id(0) * XBLOCK + base1 = tl.arange(0, XBLOCK_SUB) + loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB + base2 = tl.arange(0, RBLOCK) + loops2: tl.constexpr = (r1_numel + RBLOCK - 1) // RBLOCK + for loop1 in range(loops1): + x = offset + (loop1 * XBLOCK_SUB) + base1 + x0 = offset + (loop1 * XBLOCK_SUB) + base1[:, None] + xmask = x0 < x0_numel + _tmp2 = tl.full([XBLOCK_SUB, RBLOCK], 0, tl.float32) + for loop2 in range(loops2): + r1_prime = loop2 * RBLOCK + base2[:, None] + r1 = loop2 * RBLOCK + base2[None, :] + rmask = r1 < r1_numel + tmp0 = tl.load(in_ptr0 + (r1 + (65 * x0)), rmask & xmask, eviction_policy='evict_first', other=0.0) + tmp1 = tl.reshape(tmp0, [XBLOCK_SUB, RBLOCK]) + tmp3 = _tmp2 + tmp1 + _tmp2 = tmp3 + tmp2 = tl.sum(_tmp2, 1).reshape(XBLOCK_SUB, 1) + tl.store(out_ptr0 + (x0), tmp2, xmask) + + +def test_cases(): + size = (128, 65) + b = weights = torch.randn((size), dtype=torch.float32).npu() + c = torch.sum(b, dim=1) + ret = torch.randn((size[0]), device='npu', dtype=torch.float32).npu().reshape(size[0]) + triton_unlign[1, 1, 1](b, ret, size[0], size[1], size[0], size[0], 32) + assert torch.allclose(c, ret, rtol=1e-03, atol=1e-03, equal_nan=True) diff --git a/third_party/ascend/unittest/pytest_ut/test_unused_func_arg.py b/third_party/ascend/unittest/pytest_ut/test_unused_func_arg.py new file mode 100644 index 0000000000..3e7812de86 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_unused_func_arg.py @@ -0,0 +1,82 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import test_common +import torch +import torch_npu +import pytest +import math + + +def expand_to_next_power_of_two(a): + if a <= 0: + raise ValueError("must >0") + if (math.log2(a)).is_integer(): + return a + return 2**math.ceil(math.log2(a)) + + +@triton.jit +def triton_unused_func_arg_kernel(output_ptr, x_ptr, X: tl.constexpr, Y: tl.constexpr, Z: tl.constexpr, + XNUMEL: tl.constexpr, YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr): + xidx = tl.arange(0, XNUMEL) + yidx = tl.arange(0, YNUMEL) + zidx = tl.arange(0, ZNUMEL) + Xmask = xidx < X + Ymask = yidx < Y + Zmask = zidx < Z + oidx = xidx[:, None, None] * Y * Z + yidx[None, :, None] * Z + zidx[None, None, :] + mask = (Xmask[:, None, None]) & (Ymask[None, :, None]) & (Zmask[None, None, :]) + abc = tl.load(x_ptr + oidx, mask=mask) + ret = tl.zeros_like(abc) + tl.store(output_ptr + oidx, ret, mask=mask) + + +testlist = [ + (triton_unused_func_arg_kernel, 'int8', torch.int8, 2, 255, 9), + (triton_unused_func_arg_kernel, 'int16', torch.int16, 3, 5, 3), + (triton_unused_func_arg_kernel, 'int32', torch.int32, 2, 255, 9), + (triton_unused_func_arg_kernel, 'int64', torch.int64, 2, 5, 3), + (triton_unused_func_arg_kernel, 'float16', torch.float16, 55, 5, 16), + (triton_unused_func_arg_kernel, 'float16', torch.float16, 4, 5, 17), + (triton_unused_func_arg_kernel, 'float16', torch.float16, 6, 5, 15), + (triton_unused_func_arg_kernel, 'float16', torch.float16, 2, 1928, 3), + (triton_unused_func_arg_kernel, 'float32', torch.float32, 2, 255, 9), + (triton_unused_func_arg_kernel, 'bfloat16', torch.bfloat16, 3, 5, 3), + (triton_unused_func_arg_kernel, 'bool', torch.bool, 3, 5, 3), +] + + +@pytest.mark.parametrize('testfunc, sigtype, dtype, X, Y, Z', testlist) +def test_npu(testfunc, sigtype, dtype, X, Y, Z): + XNUMEL = expand_to_next_power_of_two(X) + YNUMEL = expand_to_next_power_of_two(Y) + ZNUMEL = expand_to_next_power_of_two(Z) + x = torch.full((X, Y, Z), 10, dtype=dtype).npu() + y = torch.full((X, Y, Z), 0, dtype=dtype).npu() + output = torch.full((X, Y, Z), 5, dtype=dtype).npu() + testfunc[1, 1, 1](output, x, X, Y, Z, XNUMEL, YNUMEL, ZNUMEL) + test_common.validate_cmp(sigtype, output, y) + + +if __name__ == "__main__": + test_npu(triton_unused_func_arg_kernel, "bool", torch.bool, 3, 5, 3) diff --git a/third_party/ascend/unittest/pytest_ut/test_use_analysis.py b/third_party/ascend/unittest/pytest_ut/test_use_analysis.py new file mode 100644 index 0000000000..06ce0a2470 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_use_analysis.py @@ -0,0 +1,74 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest + +import triton +import triton.language as tl +import triton.language.extra.cann.libdevice as libdevice +import test_common + +import torch +import torch_npu + + +@triton.jit +def triton_reduce_deadcode(v_ptr, in_ptr0, in_ptr1, out_ptr0, VBLOCK: tl.constexpr, XBLOCK: tl.constexpr, + YBLOCK: tl.constexpr): + v_idx = tl.arange(0, VBLOCK) + v = tl.load(v_ptr + v_idx) + v_ret = tl.argmax(v, 0) + if v_ret < v_ret + 1: + for _ in range(v_ret, v_ret + 1): + cube_idx = tl.arange(0, XBLOCK)[:, None] * YBLOCK + tl.arange(0, YBLOCK)[None, :] + c0 = tl.load(in_ptr0 + cube_idx) + c1 = tl.load(in_ptr1 + cube_idx) + ret = tl.dot(c0, c1) + 1 + tl.store(out_ptr0 + cube_idx, ret) + else: + for _ in range(v_ret - 1, v_ret): + cube_idx = tl.arange(0, XBLOCK)[:, None] * YBLOCK + tl.arange(0, YBLOCK)[None, :] + c0 = tl.load(in_ptr0 + cube_idx) + c1 = tl.load(in_ptr1 + cube_idx) + ret = tl.dot(c0, c1) + 1 + tl.store(out_ptr0 + cube_idx, ret) + + +def torch_reduce_deadcode(in0, in1, v): + v_ret = torch.argmax(v) + if v_ret < v_ret + 1: + ret = torch.matmul(in0, in1) + 1 + else: + ret = torch.matmul(in0, in1) + 1 + return ret + + +def test_reduce_deadcode(): + VBLOCK, XBLOCK, YBLOCK = 16, 16, 16 + sigtype = 'float32' + dtype = torch.float32 + in0 = torch.randn((XBLOCK, YBLOCK), dtype=dtype, device='npu') + in1 = torch.randn((XBLOCK, YBLOCK), dtype=dtype, device='npu') + v = torch.randn((VBLOCK, ), dtype=dtype, device='npu') + out = torch.zeros((XBLOCK, YBLOCK), dtype=dtype, device='npu') + + triton_reduce_deadcode[(1, )](v, in0, in1, out, VBLOCK=VBLOCK, XBLOCK=XBLOCK, YBLOCK=YBLOCK) + expected = torch_reduce_deadcode(in0, in1, v) + test_common.validate_cmp(sigtype, out, expected) diff --git a/third_party/ascend/unittest/pytest_ut/test_utils.py b/third_party/ascend/unittest/pytest_ut/test_utils.py new file mode 100644 index 0000000000..ee8ec1b28f --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_utils.py @@ -0,0 +1,48 @@ +import os +import platform +import hashlib +from unittest.mock import patch, mock_open +import unittest +import pytest +from triton.backends.ascend.utils import ( + get_cann_version_file_hash, + get_machine_arch, +) + + +@patch.object(platform, "machine") +def test_get_machine_arch_normal_case(mock_machine): + mock_machine.return_value = "x86_64" + result = get_machine_arch() + assert result == "x86_64" + + +@patch.object(platform, "machine") +def test_get_machine_arch_arm_case(mock_machine): + mock_machine.return_value = "arm64" + result = get_machine_arch() + assert result == "aarch64" + + +@patch.object(platform, "machine") +def test_get_machine_arch_unknown_case(mock_machine): + mock_machine.return_value = "unknown" + with pytest.raises(KeyError): + get_machine_arch() + + +def test_get_cann_version_file_hash_normal_case(): + test_version = "9.0.0" + test_inner_version = "v100rc00" + test_file_context = f"version={test_version}\ninnerversion={test_inner_version}".encode("utf-8") + with patch("builtins.open", mock_open(read_data=test_file_context)): + result = get_cann_version_file_hash() + excepted_hash = hashlib.sha256(test_file_context).hexdigest() + assert result == excepted_hash + + +def test_get_cann_version_file_hash_file_not_find_case(): + with patch("builtins.open") as mock_open_file: + mock_open_file.side_effect = FileNotFoundError + with pytest.raises(FileNotFoundError): + get_cann_version_file_hash() diff --git a/third_party/ascend/unittest/pytest_ut/test_view.py b/third_party/ascend/unittest/pytest_ut/test_view.py new file mode 100644 index 0000000000..988773b2c3 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_view.py @@ -0,0 +1,63 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.view(X, (ZB, XB * YB)) + + oidx = tl.arange(0, ZB)[:, None] * XB * YB + tl.arange(0, XB * YB)[None, :] + + tl.store(output_ptr + oidx, ret) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 256, 16), 1, 2, 256, 16], + ['float32', (8, 8, 4), 1, 8, 8, 4], + ['float16', (2, 256, 16), 1, 2, 256, 16], + ['float16', (8, 8, 4), 1, 8, 8, 4], + ['int8', (2, 256, 16), 1, 2, 256, 16], + ['int8', (8, 8, 4), 1, 8, 8, 4], +]) +def test_case(param_list): + dtype, shape, ncore, XB, YB, ZB = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + y_ref = x0.view(ZB, XB * YB).npu() + print(f"y_ref = {y_ref[0, 0:4]}") + y_cal = torch.empty((ZB, XB * YB), dtype=eval('torch.' + dtype)).npu() + + fn_npu_[ncore, 1, 1](y_cal, x0, XB, YB, ZB) + print(f"y_cal = {y_cal[0, 0:4]}") + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_where_lt.py b/third_party/ascend/unittest/pytest_ut/test_where_lt.py new file mode 100644 index 0000000000..26158b97de --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_where_lt.py @@ -0,0 +1,60 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +def torch_where_lt_case1(x0, x1): + res = torch.where(x0 < x1, x0, 1) + return res + + +@triton.jit +def triton_where_lt_case1(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + xindex = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = xindex < xnumel + x0 = xindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.load(in_ptr1 + (x0), xmask) + tmp2 = tmp0 < tmp1 + tmp3 = tl.where(tmp2, tmp0, 1) + tl.store(out_ptr0 + (xindex), tmp3, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 1024, 8), 2, 8192, 1024], + ['float16', (2, 1024, 8), 2, 8192, 1024], + ['int8', (2, 1024, 8), 2, 8192, 1024], +]) +def test_where_lt_case1(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_where_lt_case1(x0, x1) + y_cal = test_common.generate_tensor(shape, dtype).npu() + triton_where_lt_case1[ncore, 1, 1](x0, x1, y_cal, x0.numel(), xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_where_mask.py b/third_party/ascend/unittest/pytest_ut/test_where_mask.py new file mode 100644 index 0000000000..fb1474792a --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_where_mask.py @@ -0,0 +1,60 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +def torch_where_lt_case2(x0, x1): + res = torch.where(x0 < x1, x0, x1) + return res + + +@triton.jit +def triton_where_lt_case2(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + xindex = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = xindex < xnumel + x0 = xindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.load(in_ptr1 + (x0), xmask) + tmp2 = tmp0 < tmp1 + tmp3 = tl.where(tmp2, tmp0, tmp1) + tl.store(out_ptr0 + (xindex), tmp3, xmask) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 1024, 8), 2, 8192, 1024], + ['float16', (2, 1024, 8), 2, 8192, 1024], + ['int8', (2, 1024, 8), 2, 8192, 1024], +]) +def test_where_lt_case2(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_where_lt_case2(x0, x1) + y_cal = test_common.generate_tensor(shape, dtype).npu() + triton_where_lt_case2[ncore, 1, 1](x0, x1, y_cal, x0.numel(), xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_where_var.py b/third_party/ascend/unittest/pytest_ut/test_where_var.py new file mode 100644 index 0000000000..82adaf50ce --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_where_var.py @@ -0,0 +1,62 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common +# from test_common import pytest_runonce + + +def torch_where_lt_case2(x0, x1): + res = torch.where(x0 < x1, x0, x1) + return res + + +@triton.jit +def triton_where_lt_case2(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): + xoffset = tl.program_id(0) * XBLOCK + for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): + xindex = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] + xmask = xindex < xnumel + x0 = xindex + tmp0 = tl.load(in_ptr0 + (x0), xmask) + tmp1 = tl.load(in_ptr1 + (x0), xmask) + tmp2 = tmp0 < tmp1 + tmp3 = tl.where(tmp2, tmp0, tmp1) + tl.store(out_ptr0 + (xindex), tmp3, xmask) + + +# @pytest.mark.usefixtures("pytest_runonce") +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 1024, 8), 2, 8192, 1024], + ['float16', (2, 1024, 8), 2, 8192, 1024], + ['int8', (2, 1024, 8), 2, 8192, 1024], +]) +def test_where_lt_case2(param_list): + dtype, shape, ncore, xblock, xblock_sub = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + x1 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch_where_lt_case2(x0, x1) + y_cal = test_common.generate_tensor(shape, dtype).npu() + triton_where_lt_case2[ncore, 1, 1](x0, x1, y_cal, x0.numel(), xblock, xblock_sub) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_xor.py b/third_party/ascend/unittest/pytest_ut/test_xor.py new file mode 100644 index 0000000000..f10caf4536 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_xor.py @@ -0,0 +1,97 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import pytest +import triton +import triton.language as tl +import time +import test_common +import torch +import torch_npu + + +def standard_unary(x0, dtype): + res = x0 + return res + + +def standard_binary(x0, y0, dtype): + res = x0 ^ y0 + return res + + +@triton.jit +def triton_elementwise_unary(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + ret = x + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +@triton.jit +def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): + idx_block = tl.arange(0, NUMEL) + x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) + y = tl.load(in_ptr1 + idx_block, mask=idx_block < N) + ret = x ^ y + tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N) + + +types = [ + # (torch.float32, 'float32'), + # (torch.float16, 'float16'), + # (torch.bfloat16, 'bfloat16'), + (torch.int8, 'int8'), + # (torch.int16, 'int16'), + # (torch.int32, 'int32'), + # (torch.int64, 'int64'), +] + +shapes = [ + (3, 32), + (-32, 32), + (37, 64), + (-256, 256), + (781, 1024), +] + +map_for_64_t = {37: 31} + + +@pytest.mark.parametrize('dtype,sigtype', types) +@pytest.mark.parametrize('N,NUMEL', shapes) +def test_elementwsie_common(dtype, sigtype, N, NUMEL): + N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N + + if sigtype == "int64": + N = map_for_64_t[N] if N in map_for_64_t else N + + print(f"elementwise : ({N},) {dtype} {sigtype}") + + x0 = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + x1 = test_common.generate_tensor(shape=(N, ), dtype=sigtype).npu() + ans = standard_binary(x0, x1, dtype) + print(ans) + + out = torch.zeros((N, ), dtype=dtype).npu() + triton_elementwise_binary[1, 1, 1](x0, x1, out, N=N, NUMEL=NUMEL, debug=True) + print(out) + + test_common.validate_cmp(sigtype, out, ans) diff --git a/third_party/ascend/unittest/pytest_ut/test_xor_sum.py b/third_party/ascend/unittest/pytest_ut/test_xor_sum.py new file mode 100644 index 0000000000..218ef6a79d --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_xor_sum.py @@ -0,0 +1,78 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +@triton.jit +def fn_npu_(in_ptr0, out_ptr0, xnumel, ynumel, XBLOCK: tl.constexpr, RBLOCK: tl.constexpr): + X = xnumel + Y = ynumel + xoffset = tl.program_id(0) * XBLOCK + xindex = xoffset + tl.arange(0, XBLOCK) + + x0 = xindex[:, None] + rbase = tl.arange(0, RBLOCK) + _tmp6 = tl.full([XBLOCK, RBLOCK], 0, tl.int8) + for roffset in range(0, ynumel, RBLOCK): + rindex = roffset + rbase + rmask = None + r1 = rindex[None, :] + tmp0 = tl.load(in_ptr0 + (r1 + (Y * x0)), rmask) + _tmp6 = _tmp6 ^ tmp0 + tmp6 = tl.xor_sum(_tmp6, 1) + + tl.store(out_ptr0 + (xindex), tmp6, None) + + +def bar(tensor): + N, M = tensor.shape + result = torch.zeros(N, dtype=tensor.dtype, device=tensor.device) + for i in range(N): + row_xor_sum = 0 + for j in range(M): + row_xor_sum ^= tensor[i, j].item() + result[i] = row_xor_sum + return result + + +@pytest.mark.parametrize('param_list', [ + ['int8', (64, 32), 64, 32], +]) +def test_case(param_list): + dtype, shape, xblock, rblock = param_list + a = test_common.generate_tensor(shape, dtype).npu() + + std_ret = bar(a) + print(f"std_ret={std_ret}") + + value = torch.empty_strided((a.shape[0], ), (1, ), dtype=eval('torch.' + dtype)).npu() + XBLOCK = xblock + RBLOCK = rblock + NBLOCK = a.shape[0] // XBLOCK + fn_npu_[NBLOCK, 1, 1](a, value, a.shape[0], a.shape[1], XBLOCK, RBLOCK) + print(f"triton_ret={value}") + + torch.testing.assert_close(value, std_ret) diff --git a/third_party/ascend/unittest/pytest_ut/test_zeros.py b/third_party/ascend/unittest/pytest_ut/test_zeros.py new file mode 100644 index 0000000000..3dab18af39 --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_zeros.py @@ -0,0 +1,88 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +@triton.jit +def fn_npu_f32(output_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + ret = tl.zeros((XB, YB, ZB), dtype=tl.float32) + + oidx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_f16(output_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + ret = tl.zeros((XB, YB, ZB), dtype=tl.float16) + + oidx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@triton.jit +def fn_npu_i8(output_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + ret = tl.zeros((XB, YB, ZB), dtype=tl.int8) + + oidx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 256, 16), 1, 2, 256, 16], + ['float32', (8, 8, 4), 1, 8, 8, 4], + ['float16', (2, 256, 16), 1, 2, 256, 16], + ['float16', (8, 8, 4), 1, 8, 8, 4], + ['int8', (2, 256, 16), 1, 2, 256, 16], + ['int8', (8, 8, 4), 1, 8, 8, 4], +]) +def test_case(param_list): + dtype, shape, ncore, XB, YB, ZB = param_list + + y_ref = torch.full((XB, YB, ZB), 0, dtype=eval('torch.' + dtype)).npu() + + y_cal = torch.randint(1, (XB, YB, ZB), dtype=eval('torch.' + dtype)).npu() + if dtype == "float32": + fn_npu_f32[ncore, 1, 1](y_cal, XB, YB, ZB) + elif dtype == "float16": + fn_npu_f16[ncore, 1, 1](y_cal, XB, YB, ZB) + else: + fn_npu_i8[ncore, 1, 1](y_cal, XB, YB, ZB) + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/pytest_ut/test_zeroslike.py b/third_party/ascend/unittest/pytest_ut/test_zeroslike.py new file mode 100644 index 0000000000..6d6b4a822f --- /dev/null +++ b/third_party/ascend/unittest/pytest_ut/test_zeroslike.py @@ -0,0 +1,63 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + +import triton +import triton.language as tl +import torch +import torch_npu +import pytest +import test_common + + +@triton.jit +def fn_npu_(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr): + xidx = tl.arange(0, XB) + yidx = tl.arange(0, YB) + zidx = tl.arange(0, ZB) + + idx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + X = tl.load(x_ptr + idx) + + ret = tl.zeros_like(X) + + oidx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + + tl.store(output_ptr + oidx, ret) + + +@pytest.mark.parametrize('param_list', [ + ['float32', (2, 256, 16), 1, 2, 256, 16], + ['float32', (8, 8, 4), 1, 8, 8, 4], + ['float16', (2, 256, 16), 1, 2, 256, 16], + ['float16', (8, 8, 4), 1, 8, 8, 4], + ['int8', (2, 256, 16), 1, 2, 256, 16], + ['int8', (8, 8, 4), 1, 8, 8, 4], +]) +def test_case(param_list): + dtype, shape, ncore, XB, YB, ZB = param_list + x0 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch.zeros_like(x0, dtype=eval('torch.' + dtype)).npu() + print(f"y_ref = {y_ref[0, 0, 0:4]}") + y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + + fn_npu_[ncore, 1, 1](y_cal, x0, XB, YB, ZB) + print(f"y_cal = {y_cal[0, 0, 0:4]}") + test_common.validate_cmp(dtype, y_cal, y_ref) diff --git a/third_party/ascend/unittest/run_prtest.sh b/third_party/ascend/unittest/run_prtest.sh new file mode 100755 index 0000000000..cbff6889be --- /dev/null +++ b/third_party/ascend/unittest/run_prtest.sh @@ -0,0 +1,170 @@ +#!/bin/bash +# notice: this script supports python3.11.x + +set -ex + +script=$(readlink -f "$0") +script_dir=$(dirname "$script") + +# clean old logs +mkdir -p /home/pr_test_log +UNITTEST_DIR="triton-ascend/third_party/ascend/unittest" + +# define summary file path +SUMMARY_FILE="${WORKSPACE}/${UNITTEST_DIR}/summary.txt" + +function clean_cache() { + if [ -d /tmp/torchinductor_* ];then + rm -rf /tmp/torchinductor_* + fi + if [ -d ${HOME}/.triton/dump ];then + rm -rf ${HOME}/.triton/dump + fi + + if [ -d ${HOME}/.triton/cache ];then + rm -rf ${HOME}/.triton/cache + fi +} + + +function run_case_by_multi_card() { + NPU_DEVICES=$(ls /dev/davinci[0-9]* 2>/dev/null | wc -l) + [ $NPU_DEVICES -eq 0 ] && { + echo "No Ascend devices found!" + exit 1 + } + + echo "Detected $NPU_DEVICES Ascend devices" + + test_dir=$1 + cd ${test_dir} + + # clean logs + rm -rf logs && mkdir logs + + # record start time + start_time=$(date +"%Y-%m-%d %H:%M:%S") + echo "===== Test Start Time: ${start_time} =====" + + # run tests and capture exit status + set +e + python -m pytest ${test_dir} -n auto --dist=loadfile -v --junitxml=logs/results.xml | tee logs/raw_output.log + pytest_exit=$? + set -e + + # process logs (add device tags) + awk ' + />> Worker gw[0-9]+ using NPU device/ { + split($0, parts, / /) + dev_id = parts[NF] + worker = parts[3] + print "[" strftime("%Y-%m-%d %H:%M:%S") "| DEV-" dev_id "] " $0 + next + } + { print "[" strftime("%Y-%m-%d %H:%M:%S") "| DEV-" dev_id "] " $0 } + ' logs/raw_output.log > logs/combined.log + + # parse test result statistics + total_tests=0 + passed_tests=0 + failed_tests=0 + skipped_tests=0 + error_tests=0 + + # use Python to parse JUnit XML report + python3 -c " +import xml.etree.ElementTree as ET +import os + +xml_file = os.path.join('logs', 'results.xml') +if not os.path.exists(xml_file): + print('JUnitXML report not found:', xml_file) + exit(1) + +tree = ET.parse(xml_file) +root = tree.getroot() + +total = 0 +passed = 0 +failed = 0 +skipped = 0 +errors = 0 + +# 遍历所有testsuite +for testsuite in root.findall('testsuite'): + total += int(testsuite.get('tests', 0)) + passed += int(testsuite.get('tests', 0)) - int(testsuite.get('errors', 0)) - int(testsuite.get('failures', 0)) - int(testsuite.get('skipped', 0)) + failed += int(testsuite.get('failures', 0)) + skipped += int(testsuite.get('skipped', 0)) + errors += int(testsuite.get('errors', 0)) + +print(f'total_tests={total}') +print(f'passed_tests={passed}') +print(f'failed_tests={failed}') +print(f'skipped_tests={skipped}') +print(f'error_tests={errors}') +" > logs/stats.tmp + + # load stats + source logs/stats.tmp + rm logs/stats.tmp + + # record end time + end_time=$(date +"%Y-%m-%d %H:%M:%S") + duration=$(( $(date -d "$end_time" +%s) - $(date -d "$start_time" +%s) )) + duration_str=$(printf "%02dh %02dm %02ds" $((duration/3600)) $(((duration%3600)/60)) $((duration%60))) + + # generate summary + stats_summary=" +===== Test Summary - [generalization_cases] ===== +Test Directory: $(basename ${test_dir}) +Test Start Time: ${start_time} +Test End Time: ${end_time} +Total Duration: ${duration_str} +------------------------ +Total Tests: ${total_tests} +Passed Tests: ${passed_tests} +Failed Tests: ${failed_tests} +Skipped Tests: ${skipped_tests} +Error Tests: ${error_tests} +Success Rate: $(( passed_tests * 100 / total_tests ))% (Passed/Total) +NPU Devices: ${NPU_DEVICES} +======================== +" + + # output stats summary to console + echo "${stats_summary}" + + # append stats summary to summary.txt + echo "${stats_summary}" >> ${SUMMARY_FILE} + + echo "========================================" + echo "All tests completed!" + echo "JUnit Report: logs/results.xml" + echo "Combined Log: logs/combined.log" + echo "Stats Summary has been appended to: ${SUMMARY_FILE}" + echo "========================================" + + zip_file=$2 + cd ${test_dir}/logs + zip ${zip_file} combined.log + cp ${zip_file} "/home/pr_test_log" + + # return pytest exit status + return $pytest_exit +} + +# initialize stats file +echo "Generate Time: $(date +"%Y-%m-%d %H:%M:%S")" >> ${SUMMARY_FILE} +echo "========================================" >> ${SUMMARY_FILE} + +# run gene case +zip_file="test_generalization_case_$(date +%Y%m%d).zip" +TEST_generalization="${WORKSPACE}/${UNITTEST_DIR}/generalization_cases" +clean_cache +run_case_by_multi_card ${TEST_generalization} ${zip_file} + +echo "========================================" >> ${SUMMARY_FILE} + +# copy summary.txt to /home/pr_test_log +cp ${SUMMARY_FILE} /home/pr_test_log