From 7cd295cececc3144ab0e1001bdd302c7ff21507f Mon Sep 17 00:00:00 2001 From: Junkai-Wu Date: Tue, 7 Jul 2026 10:05:33 +0800 Subject: [PATCH] v4.6 tag release. (#3362) --- CHANGELOG.md | 41 +- CMakeLists.txt | 2 +- CUDA.cmake | 18 + README.md | 37 + ...113_hopper_grouped_gemm_fused_gated_act.cu | 9 +- .../ampere/kernel/dense_gemm/tensorop_gemm.py | 662 +++-- .../blackwell/kernel/attention/fmha/fmha.py | 4 +- .../export/run_with_dynamic_loading.sh | 22 +- .../export/run_with_static_linking.sh | 19 +- include/cute/arch/mma_sm100_desc.hpp | 280 +++ include/cute/atom/mma_traits_sm100.hpp | 876 +------ include/cute/atom/mma_traits_sm100_frag.hpp | 640 +++++ include/cute/config.hpp | 14 + include/cute/container/array.hpp | 6 +- include/cute/container/array_subbyte.hpp | 6 +- include/cute/container/tuple.hpp | 6 +- include/cute/container/type_list.hpp | 6 +- include/cute/numeric/arithmetic_tuple.hpp | 6 +- .../cutlass/arch/grid_dependency_control.h | 2 +- .../gemm/collective/sm120_mma_array_tma.hpp | 10 +- .../sm120_mma_array_tma_blockwise_scaling.hpp | 4 + .../operators/api_reference/arguments.rst | 55 + .../operators/api_reference/discovery.rst | 22 + media/docs/operators/api_reference/index.rst | 20 + .../docs/operators/api_reference/metadata.rst | 45 + media/docs/operators/api_reference/misc.rst | 20 + .../docs/operators/api_reference/operator.rst | 40 + media/docs/operators/overview.rst | 157 ++ .../docs/operators/tutorials/000_gemm.nblink | 1 + .../001_gemm_with_fused_epilogue.nblink | 1 + .../002_bring_your_own_kernel.nblink | 1 + .../003_host_latency_best_practices.nblink | 1 + .../tutorials/004_fake_tensors.nblink | 1 + .../005_grouped_gemm_contiguous_offset.nblink | 1 + .../tutorials/006_block_scaled_gemm.nblink | 1 + media/docs/operators/tutorials/index.rst | 16 + media/docs/pythonDSL/cute_dsl.rst | 1 + .../pythonDSL/cute_dsl_general/debugging.rst | 61 + .../framework_integration.rst | 7 +- .../cute_dsl_general/iket_profiling.rst | 910 +++++++ .../docs/pythonDSL/cute_dsl_general/types.rst | 16 - media/docs/pythonDSL/deprecation.rst | 7 + media/images/fp16_gemm_4_iket_instrument.png | Bin 0 -> 284361 bytes media/images/fp16_gemm_4_iket_quickstart.png | Bin 0 -> 191957 bytes media/images/iket_concept_timeline.svg | 93 + operators/README.md | 149 ++ operators/cutlass/operators/__init__.py | 136 + operators/cutlass/operators/arch.py | 219 ++ .../cutlass/operators/arguments/__init__.py | 60 + operators/cutlass/operators/arguments/base.py | 134 + .../cutlass/operators/arguments/epilogue.py | 196 ++ operators/cutlass/operators/arguments/gemm.py | 212 ++ .../operators/arguments/grouped_gemm.py | 107 + .../cutlass/operators/arguments/operand.py | 396 +++ operators/cutlass/operators/artifact.py | 55 + operators/cutlass/operators/base.py | 318 +++ operators/cutlass/operators/config.py | 210 ++ .../cutlass/operators/fusion/__init__.py | 126 + .../cutlass/operators/fusion/activation.py | 233 ++ .../operators/fusion/backend/__init__.py | 45 + .../operators/fusion/backend/emitter_base.py | 165 ++ .../operators/fusion/backend/sm100_emitter.py | 151 ++ .../operators/fusion/backend/sm100_nodes.py | 140 ++ .../operators/fusion/backend/sm80_emitter.py | 46 + .../operators/fusion/backend/sm80_nodes.py | 247 ++ .../operators/fusion/backend/sm90_emitter.py | 97 + .../operators/fusion/backend/sm90_nodes.py | 327 +++ .../cutlass/operators/fusion/epilogue.py | 88 + operators/cutlass/operators/fusion/evt_ops.py | 116 + .../operators/fusion/frontend/__init__.py | 41 + .../fusion/frontend/frontend_base.py | 303 +++ .../operators/fusion/frontend/python_ast.py | 261 ++ .../cutlass/operators/fusion/ir/__init__.py | 75 + .../cutlass/operators/fusion/ir/c_types.py | 247 ++ .../operators/fusion/ir/compute_nodes.py | 99 + .../cutlass/operators/fusion/ir/dag_ir.py | 258 ++ .../operators/fusion/ir/layout_algorithm.py | 362 +++ .../operators/fusion/ir/layout_nodes.py | 351 +++ .../cutlass/operators/fusion/ir/load_nodes.py | 312 +++ operators/cutlass/operators/fusion/ir/node.py | 330 +++ .../operators/fusion/ir/store_nodes.py | 276 ++ .../cutlass/operators/fusion/ir/tensor.py | 155 ++ operators/cutlass/operators/fusion/library.py | 430 ++++ .../operators/fusion/passes/__init__.py | 59 + .../operators/fusion/passes/graph_drawer.py | 133 + .../fusion/passes/pass_argument_type.py | 136 + .../fusion/passes/pass_dag_2_tree.py | 176 ++ .../fusion/passes/pass_fix_element_d.py | 86 + .../operators/fusion/passes/pass_get_impl.py | 93 + .../fusion/passes/pass_layout_elimination.py | 230 ++ .../operators/fusion/passes/pass_manager.py | 185 ++ .../fusion/passes/pass_no_op_elimination.py | 59 + .../fusion/passes/pass_preprocess_red.py | 96 + .../passes/pass_shape_type_propagation.py | 60 + .../fusion/passes/smem_size_calculator.py | 363 +++ .../cutlass/operators/fusion/passes/util.py | 46 + .../operators/fusion/pycute/__init__.py | 36 + .../operators/fusion/pycute/int_tuple.py | 229 ++ .../cutlass/operators/fusion/pycute/layout.py | 409 +++ .../operators/fusion/pycute/swizzle.py | 133 + .../cutlass/operators/fusion/pycute/typing.py | 42 + operators/cutlass/operators/manifest.py | 153 ++ .../cutlass/operators/metadata/__init__.py | 73 + operators/cutlass/operators/metadata/base.py | 170 ++ .../operators/metadata/design/__init__.py | 45 + .../cutlass/operators/metadata/design/base.py | 109 + .../operators/metadata/design/sm100.py | 111 + .../metadata/design/tile_scheduler.py | 91 + .../cutlass/operators/metadata/epilogue.py | 71 + .../operators/metadata/operand_constraints.py | 255 ++ .../operators/metadata/operands/__init__.py | 37 + .../operators/metadata/operands/base.py | 70 + .../operators/metadata/operands/gemm.py | 90 + .../metadata/operands/grouped_gemm.py | 100 + operators/cutlass/operators/mma.py | 166 ++ .../cutlass/operators/providers/__init__.py | 99 + .../operators/providers/cutedsl/__init__.py | 66 + .../providers/cutedsl/evt/common_efc.py | 2189 ++++++++++++++++ .../providers/cutedsl/evt/converter.py | 470 ++++ .../providers/cutedsl/gemm/__init__.py | 43 + .../gemm/implementations/operator_helpers.py | 73 + .../cutedsl/gemm/implementations/scheduler.py | 589 +++++ ..._contiguous_offset_2d3d_dense_gemm_impl.py | 1976 +++++++++++++++ ...ense_blockscaled_static_persistent_impl.py | 2170 ++++++++++++++++ .../implementations/sm100_mixed_input_impl.py | 2058 +++++++++++++++ .../implementations/sm100_persistent_impl.py | 1277 ++++++++++ ...sm100_persistent_preferred_cluster_impl.py | 1176 +++++++++ .../sm100_static_persistent_efc_impl.py | 1686 +++++++++++++ .../implementations/sm100_tgv_gemm_impl.py | 1085 ++++++++ .../sm80_tensorop_gemm_impl.py | 814 ++++++ .../sm90_static_persistent_impl.py | 1236 +++++++++ ...sm100_contiguous_offset_2d3d_dense_gemm.py | 451 ++++ ...100_dense_blockscaled_static_persistent.py | 606 +++++ .../cutedsl/gemm/sm100_mixed_input.py | 910 +++++++ .../cutedsl/gemm/sm100_persistent.py | 494 ++++ .../sm100_persistent_preferred_cluster.py | 671 +++++ .../gemm/sm100_static_persistent_efc.py | 506 ++++ .../providers/cutedsl/gemm/sm100_tgv_gemm.py | 420 ++++ .../cutedsl/gemm/sm80_tensorop_gemm.py | 316 +++ .../cutedsl/gemm/sm90_static_persistent.py | 441 ++++ .../cutedsl/integration_utils/__init__.py | 40 + .../cutedsl/integration_utils/builders.py | 97 + .../cutedsl/integration_utils/mma.py | 49 + .../operators/providers/cutedsl/operator.py | 90 + .../cutlass/operators/providers/provider.py | 153 ++ operators/cutlass/operators/status.py | 68 + operators/cutlass/operators/typing.py | 73 + operators/cutlass/operators/utils/__init__.py | 39 + operators/cutlass/operators/utils/common.py | 84 + operators/cutlass/operators/utils/device.py | 247 ++ operators/cutlass/operators/utils/dtype.py | 213 ++ operators/cutlass/operators/utils/gemm.py | 96 + operators/cutlass/operators/utils/generate.py | 100 + operators/cutlass/operators/utils/layout.py | 80 + operators/cutlass/operators/utils/tensor.py | 575 +++++ operators/cutlass/operators/workspace.py | 67 + operators/examples/000_gemm.ipynb | 568 +++++ .../001_gemm_with_fused_epilogue.ipynb | 679 +++++ .../examples/002_bring_your_own_kernel.ipynb | 663 +++++ .../003_host_latency_best_practices.ipynb | 619 +++++ operators/examples/004_fake_tensors.ipynb | 167 ++ .../005_grouped_gemm_contiguous_offset.ipynb | 246 ++ .../examples/006_block_scaled_gemm.ipynb | 352 +++ operators/pyproject.toml | 175 ++ operators/test/__init__.py | 27 + operators/test/conftest.py | 104 + .../test/integration/test_blockscaled_gemm.py | 1152 +++++++++ .../test_contiguous_offset_dense_gemm.py | 128 + operators/test/integration/test_cuda_graph.py | 164 ++ .../test/integration/test_cutlass_tensor.py | 302 +++ operators/test/integration/test_gemm.py | 206 ++ .../test/integration/test_gemm_dynamic.py | 128 + .../integration/test_gemm_epilogue_fusion.py | 1158 +++++++++ operators/test/integration/test_gemm_jax.py | 74 + .../test/integration/test_mixed_input_gemm.py | 274 ++ operators/test/integration/test_notebooks.py | 93 + operators/test/integration/test_overhead.py | 311 +++ .../test_preferred_cluster_gemm.py | 415 ++++ operators/test/integration/test_tgv_gemm.py | 89 + operators/test/integration/test_tvm_ffi.py | 230 ++ operators/test/test_utils/__init__.py | 40 + operators/test/test_utils/allocation.py | 511 ++++ operators/test/test_utils/common.py | 214 ++ operators/test/test_utils/gemm.py | 403 +++ operators/test/test_utils/reference_check.py | 454 ++++ operators/test/unit/test_arguments.py | 153 ++ operators/test/unit/test_config.py | 183 ++ .../unit/test_fusion_arguments_parsing.py | 94 + operators/test/unit/test_manifest.py | 167 ++ operators/test/unit/test_metadata.py | 204 ++ operators/test/unit/test_provider.py | 271 ++ operators/test/unit/test_reference_check.py | 477 ++++ operators/test/unit/test_tensor_wrapper.py | 145 ++ python/CuTeDSL/_mlir_helpers/__init__.py | 29 + python/CuTeDSL/_mlir_helpers/arith.py | 1068 ++++++++ python/CuTeDSL/_mlir_helpers/dialect_proxy.py | 118 + python/CuTeDSL/_mlir_helpers/gpu.py | 63 + python/CuTeDSL/_mlir_helpers/lru_cache_ir.py | 79 + python/CuTeDSL/_mlir_helpers/math.py | 2213 +++++++++++++++++ .../CuTeDSL/_mlir_helpers/nanobind_compat.py | 71 + python/CuTeDSL/_mlir_helpers/op.py | 389 +++ python/CuTeDSL/_mlir_helpers/vector.py | 1358 ++++++++++ python/CuTeDSL/cutlass/__init__.py | 55 +- python/CuTeDSL/cutlass/_pth_hook.py | 2 +- python/CuTeDSL/cutlass/address_space.py | 14 + python/CuTeDSL/cutlass/base_dsl/__init__.py | 2 + .../CuTeDSL/cutlass/base_dsl/address_space.py | 31 + python/CuTeDSL/cutlass/base_dsl/arch.py | 102 +- .../CuTeDSL/cutlass/base_dsl/ast_helpers.py | 60 +- .../cutlass/base_dsl/ast_preprocessor.py | 143 +- python/CuTeDSL/cutlass/base_dsl/common.py | 405 +-- python/CuTeDSL/cutlass/base_dsl/compiler.py | 426 +++- .../CuTeDSL/cutlass/base_dsl/diagnostics.py | 1857 ++++++++++++++ python/CuTeDSL/cutlass/base_dsl/dsl.py | 529 +++- .../CuTeDSL/cutlass/base_dsl/env_manager.py | 129 +- .../base_dsl/export/c_header_generator.py | 50 +- .../CuTeDSL/cutlass/base_dsl/export/export.py | 7 +- .../base_dsl/export/external_binary_module.py | 9 +- python/CuTeDSL/cutlass/base_dsl/ffi.py | 18 +- .../CuTeDSL/cutlass/base_dsl/jit_executor.py | 487 +++- python/CuTeDSL/cutlass/base_dsl/leaf_utils.py | 2 +- python/CuTeDSL/cutlass/base_dsl/pointer.py | 55 + .../CuTeDSL/cutlass/base_dsl/runtime/cuda.py | 61 +- .../cutlass/base_dsl/runtime/device_tensor.py | 23 +- .../base_dsl/runtime/jit_arg_adapters.py | 4 +- python/CuTeDSL/cutlass/base_dsl/swizzle.py | 374 +++ python/CuTeDSL/cutlass/base_dsl/torch.py | 2 + python/CuTeDSL/cutlass/base_dsl/typing.py | 979 +++++++- .../cutlass/base_dsl/utils/tree_utils.py | 90 +- python/CuTeDSL/cutlass/base_dsl/vector.py | 42 + python/CuTeDSL/cutlass/compiler/__init__.py | 37 +- python/CuTeDSL/cutlass/cute/__init__.py | 18 +- .../cute/_tvm_ffi_args_spec_converter.py | 26 +- python/CuTeDSL/cutlass/cute/arch/__init__.py | 16 + .../cutlass/cute/arch/numeric_conversion.py | 68 +- .../cutlass/cute/arch/nvvm_wrappers.py | 340 ++- python/CuTeDSL/cutlass/cute/arch/smem.py | 73 +- python/CuTeDSL/cutlass/cute/arch/tmem.py | 5 +- python/CuTeDSL/cutlass/cute/atom.py | 7 +- python/CuTeDSL/cutlass/cute/core.py | 139 +- .../cutlass/cute/experimental/algorithm.py | 18 +- .../CuTeDSL/cutlass/cute/experimental/core.py | 29 +- .../CuTeDSL/cutlass/cute/experimental/math.py | 35 + .../cutlass/cute/experimental/memory.py | 18 +- .../cutlass/cute/experimental/pipeline.py | 16 +- .../cutlass/cute/experimental/utils.py | 12 +- .../CuTeDSL/cutlass/cute/export/aot_config.py | 6 +- .../cutlass/cute/export/c_header_generator.py | 63 +- python/CuTeDSL/cutlass/cute/math.py | 101 +- python/CuTeDSL/cutlass/cute/metadata.py | 814 ++++++ python/CuTeDSL/cutlass/cute/nvgpu/helpers.py | 12 +- .../CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py | 48 +- .../cutlass/cute/nvgpu/warp/__init__.py | 1 + python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py | 65 + .../cutlass/cute/nvgpu/warpgroup/mma.py | 2 +- python/CuTeDSL/cutlass/cute/runtime.py | 151 +- python/CuTeDSL/cutlass/cute/tensor.py | 48 +- python/CuTeDSL/cutlass/cute/typing.py | 89 +- .../CuTeDSL/cutlass/cutlass_dsl/__init__.py | 2 + .../cutlass/cutlass_dsl/cuda_jit_executor.py | 25 +- python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py | 328 ++- .../cutlass_dsl/cutlass_ast_decorators.py | 61 +- .../cutlass/cutlass_dsl/tvm_ffi_provider.py | 36 +- python/CuTeDSL/cutlass/jax/types.py | 24 +- python/CuTeDSL/cutlass/pipeline/__init__.py | 1 + python/CuTeDSL/cutlass/pipeline/helpers.py | 66 +- python/CuTeDSL/cutlass/pipeline/profiling.py | 2 +- python/CuTeDSL/cutlass/pipeline/sm100.py | 106 +- python/CuTeDSL/cutlass/pipeline/sm90.py | 17 +- python/CuTeDSL/cutlass/runtime.py | 131 + python/CuTeDSL/cutlass/testing.py | 51 +- python/CuTeDSL/cutlass/utils/README.md | 637 +++++ python/CuTeDSL/cutlass/utils/__init__.py | 11 +- .../cutlass/utils/blackwell_helpers.py | 10 +- .../cutlass/utils/gemm/tensor_utils.py | 127 +- .../grouped_gemm_tile_scheduler_helper.py | 481 ---- .../cutlass/utils/mixed_input_helpers.py | 159 ++ python/CuTeDSL/cutlass/utils/print_latex.py | 409 ++- .../CuTeDSL/cutlass/utils/smem_allocator.py | 119 +- .../cutlass/utils/tensormap_manager.py | 10 +- .../CuTeDSL/cutlass/utils/tmem_allocator.py | 61 +- python/CuTeDSL/requirements-cu13.txt | 2 +- python/CuTeDSL/requirements.txt | 2 +- python/cutlass_library/generator.py | 67 +- .../CMakeLists.txt | 1 + .../sm120_bs_gemm_mxf4_mxf4_f32_f32.cu | 128 + ...0_bs_gemm_nvf4_nvf4_f32_bf16_group_gemm.cu | 351 +++ ...mm_nvf4_nvf4_f32_nvf4_group_gemm_fusion.cu | 3 - .../device/sm120_tensorop_gemm/CMakeLists.txt | 2 +- ...120_gemm_f8_f8_f32_tensor_op_group_gemm.cu | 7 +- test/unit/nvrtc/thread/nvrtc_config.in | 1 + test/unit/nvrtc/thread/nvrtc_contraction.cu | 2 +- 292 files changed, 65178 insertions(+), 3130 deletions(-) create mode 100644 include/cute/atom/mma_traits_sm100_frag.hpp create mode 100644 media/docs/operators/api_reference/arguments.rst create mode 100644 media/docs/operators/api_reference/discovery.rst create mode 100644 media/docs/operators/api_reference/index.rst create mode 100644 media/docs/operators/api_reference/metadata.rst create mode 100644 media/docs/operators/api_reference/misc.rst create mode 100644 media/docs/operators/api_reference/operator.rst create mode 100644 media/docs/operators/overview.rst create mode 100644 media/docs/operators/tutorials/000_gemm.nblink create mode 100644 media/docs/operators/tutorials/001_gemm_with_fused_epilogue.nblink create mode 100644 media/docs/operators/tutorials/002_bring_your_own_kernel.nblink create mode 100644 media/docs/operators/tutorials/003_host_latency_best_practices.nblink create mode 100644 media/docs/operators/tutorials/004_fake_tensors.nblink create mode 100644 media/docs/operators/tutorials/005_grouped_gemm_contiguous_offset.nblink create mode 100644 media/docs/operators/tutorials/006_block_scaled_gemm.nblink create mode 100644 media/docs/operators/tutorials/index.rst create mode 100644 media/docs/pythonDSL/cute_dsl_general/iket_profiling.rst create mode 100644 media/images/fp16_gemm_4_iket_instrument.png create mode 100644 media/images/fp16_gemm_4_iket_quickstart.png create mode 100644 media/images/iket_concept_timeline.svg create mode 100644 operators/README.md create mode 100644 operators/cutlass/operators/__init__.py create mode 100644 operators/cutlass/operators/arch.py create mode 100644 operators/cutlass/operators/arguments/__init__.py create mode 100644 operators/cutlass/operators/arguments/base.py create mode 100644 operators/cutlass/operators/arguments/epilogue.py create mode 100644 operators/cutlass/operators/arguments/gemm.py create mode 100644 operators/cutlass/operators/arguments/grouped_gemm.py create mode 100644 operators/cutlass/operators/arguments/operand.py create mode 100644 operators/cutlass/operators/artifact.py create mode 100644 operators/cutlass/operators/base.py create mode 100644 operators/cutlass/operators/config.py create mode 100644 operators/cutlass/operators/fusion/__init__.py create mode 100644 operators/cutlass/operators/fusion/activation.py create mode 100644 operators/cutlass/operators/fusion/backend/__init__.py create mode 100644 operators/cutlass/operators/fusion/backend/emitter_base.py create mode 100644 operators/cutlass/operators/fusion/backend/sm100_emitter.py create mode 100644 operators/cutlass/operators/fusion/backend/sm100_nodes.py create mode 100644 operators/cutlass/operators/fusion/backend/sm80_emitter.py create mode 100644 operators/cutlass/operators/fusion/backend/sm80_nodes.py create mode 100644 operators/cutlass/operators/fusion/backend/sm90_emitter.py create mode 100644 operators/cutlass/operators/fusion/backend/sm90_nodes.py create mode 100644 operators/cutlass/operators/fusion/epilogue.py create mode 100644 operators/cutlass/operators/fusion/evt_ops.py create mode 100644 operators/cutlass/operators/fusion/frontend/__init__.py create mode 100644 operators/cutlass/operators/fusion/frontend/frontend_base.py create mode 100644 operators/cutlass/operators/fusion/frontend/python_ast.py create mode 100644 operators/cutlass/operators/fusion/ir/__init__.py create mode 100644 operators/cutlass/operators/fusion/ir/c_types.py create mode 100644 operators/cutlass/operators/fusion/ir/compute_nodes.py create mode 100644 operators/cutlass/operators/fusion/ir/dag_ir.py create mode 100644 operators/cutlass/operators/fusion/ir/layout_algorithm.py create mode 100644 operators/cutlass/operators/fusion/ir/layout_nodes.py create mode 100644 operators/cutlass/operators/fusion/ir/load_nodes.py create mode 100644 operators/cutlass/operators/fusion/ir/node.py create mode 100644 operators/cutlass/operators/fusion/ir/store_nodes.py create mode 100644 operators/cutlass/operators/fusion/ir/tensor.py create mode 100644 operators/cutlass/operators/fusion/library.py create mode 100644 operators/cutlass/operators/fusion/passes/__init__.py create mode 100644 operators/cutlass/operators/fusion/passes/graph_drawer.py create mode 100644 operators/cutlass/operators/fusion/passes/pass_argument_type.py create mode 100644 operators/cutlass/operators/fusion/passes/pass_dag_2_tree.py create mode 100644 operators/cutlass/operators/fusion/passes/pass_fix_element_d.py create mode 100644 operators/cutlass/operators/fusion/passes/pass_get_impl.py create mode 100644 operators/cutlass/operators/fusion/passes/pass_layout_elimination.py create mode 100644 operators/cutlass/operators/fusion/passes/pass_manager.py create mode 100644 operators/cutlass/operators/fusion/passes/pass_no_op_elimination.py create mode 100644 operators/cutlass/operators/fusion/passes/pass_preprocess_red.py create mode 100644 operators/cutlass/operators/fusion/passes/pass_shape_type_propagation.py create mode 100644 operators/cutlass/operators/fusion/passes/smem_size_calculator.py create mode 100644 operators/cutlass/operators/fusion/passes/util.py create mode 100644 operators/cutlass/operators/fusion/pycute/__init__.py create mode 100644 operators/cutlass/operators/fusion/pycute/int_tuple.py create mode 100644 operators/cutlass/operators/fusion/pycute/layout.py create mode 100644 operators/cutlass/operators/fusion/pycute/swizzle.py create mode 100644 operators/cutlass/operators/fusion/pycute/typing.py create mode 100644 operators/cutlass/operators/manifest.py create mode 100644 operators/cutlass/operators/metadata/__init__.py create mode 100644 operators/cutlass/operators/metadata/base.py create mode 100644 operators/cutlass/operators/metadata/design/__init__.py create mode 100644 operators/cutlass/operators/metadata/design/base.py create mode 100644 operators/cutlass/operators/metadata/design/sm100.py create mode 100644 operators/cutlass/operators/metadata/design/tile_scheduler.py create mode 100644 operators/cutlass/operators/metadata/epilogue.py create mode 100644 operators/cutlass/operators/metadata/operand_constraints.py create mode 100644 operators/cutlass/operators/metadata/operands/__init__.py create mode 100644 operators/cutlass/operators/metadata/operands/base.py create mode 100644 operators/cutlass/operators/metadata/operands/gemm.py create mode 100644 operators/cutlass/operators/metadata/operands/grouped_gemm.py create mode 100644 operators/cutlass/operators/mma.py create mode 100644 operators/cutlass/operators/providers/__init__.py create mode 100644 operators/cutlass/operators/providers/cutedsl/__init__.py create mode 100644 operators/cutlass/operators/providers/cutedsl/evt/common_efc.py create mode 100644 operators/cutlass/operators/providers/cutedsl/evt/converter.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/__init__.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/implementations/operator_helpers.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/implementations/scheduler.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_contiguous_offset_2d3d_dense_gemm_impl.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_dense_blockscaled_static_persistent_impl.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_mixed_input_impl.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_persistent_impl.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_persistent_preferred_cluster_impl.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_static_persistent_efc_impl.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_tgv_gemm_impl.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm80_tensorop_gemm_impl.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm90_static_persistent_impl.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/sm100_contiguous_offset_2d3d_dense_gemm.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/sm100_dense_blockscaled_static_persistent.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/sm100_mixed_input.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/sm100_persistent.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/sm100_persistent_preferred_cluster.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/sm100_static_persistent_efc.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/sm100_tgv_gemm.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/sm80_tensorop_gemm.py create mode 100644 operators/cutlass/operators/providers/cutedsl/gemm/sm90_static_persistent.py create mode 100644 operators/cutlass/operators/providers/cutedsl/integration_utils/__init__.py create mode 100644 operators/cutlass/operators/providers/cutedsl/integration_utils/builders.py create mode 100644 operators/cutlass/operators/providers/cutedsl/integration_utils/mma.py create mode 100644 operators/cutlass/operators/providers/cutedsl/operator.py create mode 100644 operators/cutlass/operators/providers/provider.py create mode 100644 operators/cutlass/operators/status.py create mode 100644 operators/cutlass/operators/typing.py create mode 100644 operators/cutlass/operators/utils/__init__.py create mode 100644 operators/cutlass/operators/utils/common.py create mode 100644 operators/cutlass/operators/utils/device.py create mode 100644 operators/cutlass/operators/utils/dtype.py create mode 100644 operators/cutlass/operators/utils/gemm.py create mode 100644 operators/cutlass/operators/utils/generate.py create mode 100644 operators/cutlass/operators/utils/layout.py create mode 100644 operators/cutlass/operators/utils/tensor.py create mode 100644 operators/cutlass/operators/workspace.py create mode 100644 operators/examples/000_gemm.ipynb create mode 100644 operators/examples/001_gemm_with_fused_epilogue.ipynb create mode 100644 operators/examples/002_bring_your_own_kernel.ipynb create mode 100644 operators/examples/003_host_latency_best_practices.ipynb create mode 100644 operators/examples/004_fake_tensors.ipynb create mode 100644 operators/examples/005_grouped_gemm_contiguous_offset.ipynb create mode 100644 operators/examples/006_block_scaled_gemm.ipynb create mode 100644 operators/pyproject.toml create mode 100644 operators/test/__init__.py create mode 100644 operators/test/conftest.py create mode 100644 operators/test/integration/test_blockscaled_gemm.py create mode 100644 operators/test/integration/test_contiguous_offset_dense_gemm.py create mode 100644 operators/test/integration/test_cuda_graph.py create mode 100644 operators/test/integration/test_cutlass_tensor.py create mode 100644 operators/test/integration/test_gemm.py create mode 100644 operators/test/integration/test_gemm_dynamic.py create mode 100644 operators/test/integration/test_gemm_epilogue_fusion.py create mode 100644 operators/test/integration/test_gemm_jax.py create mode 100644 operators/test/integration/test_mixed_input_gemm.py create mode 100644 operators/test/integration/test_notebooks.py create mode 100644 operators/test/integration/test_overhead.py create mode 100644 operators/test/integration/test_preferred_cluster_gemm.py create mode 100644 operators/test/integration/test_tgv_gemm.py create mode 100644 operators/test/integration/test_tvm_ffi.py create mode 100644 operators/test/test_utils/__init__.py create mode 100644 operators/test/test_utils/allocation.py create mode 100644 operators/test/test_utils/common.py create mode 100644 operators/test/test_utils/gemm.py create mode 100644 operators/test/test_utils/reference_check.py create mode 100644 operators/test/unit/test_arguments.py create mode 100644 operators/test/unit/test_config.py create mode 100644 operators/test/unit/test_fusion_arguments_parsing.py create mode 100644 operators/test/unit/test_manifest.py create mode 100644 operators/test/unit/test_metadata.py create mode 100644 operators/test/unit/test_provider.py create mode 100644 operators/test/unit/test_reference_check.py create mode 100644 operators/test/unit/test_tensor_wrapper.py create mode 100644 python/CuTeDSL/_mlir_helpers/__init__.py create mode 100644 python/CuTeDSL/_mlir_helpers/arith.py create mode 100644 python/CuTeDSL/_mlir_helpers/dialect_proxy.py create mode 100644 python/CuTeDSL/_mlir_helpers/gpu.py create mode 100644 python/CuTeDSL/_mlir_helpers/lru_cache_ir.py create mode 100644 python/CuTeDSL/_mlir_helpers/math.py create mode 100644 python/CuTeDSL/_mlir_helpers/nanobind_compat.py create mode 100644 python/CuTeDSL/_mlir_helpers/op.py create mode 100644 python/CuTeDSL/_mlir_helpers/vector.py create mode 100644 python/CuTeDSL/cutlass/address_space.py create mode 100644 python/CuTeDSL/cutlass/base_dsl/address_space.py create mode 100644 python/CuTeDSL/cutlass/base_dsl/diagnostics.py create mode 100644 python/CuTeDSL/cutlass/base_dsl/pointer.py create mode 100644 python/CuTeDSL/cutlass/base_dsl/swizzle.py create mode 100644 python/CuTeDSL/cutlass/base_dsl/vector.py create mode 100644 python/CuTeDSL/cutlass/cute/metadata.py create mode 100644 python/CuTeDSL/cutlass/runtime.py create mode 100644 python/CuTeDSL/cutlass/utils/README.md delete mode 100644 python/CuTeDSL/cutlass/utils/grouped_gemm_tile_scheduler_helper.py create mode 100644 test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/sm120_bs_gemm_nvf4_nvf4_f32_bf16_group_gemm.cu diff --git a/CHANGELOG.md b/CHANGELOG.md index 7a837307f..6446475b4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,20 +2,30 @@ # CUTLASS 4.x -## [4.6.0](https://github.com/NVIDIA/cutlass/tree/main) (2026-06-11) +## [4.6.0](https://github.com/NVIDIA/cutlass/releases/tag/v4.6.0) (2026-07-01) ### CuTe DSL * New features + - New fine-grained compilation API: cute.compile_to that gives control over the what stage the compiler outputs. This feature allows customization of the path from compilation to runtime execution. cute.compile_to is considered experimental in 4.6. + - Experimental Feature: Added the IKET (In-Kernel-Event-Tracing) profiler for instrumentation-based intra-kernel activities tracing. This enables fine-grained profiling and makes it easier to understand persistent, warp-specialized kernels' performance. This is a beta feature provided by CUTLASS Python until a NVIDIA DevTools product is released, there is no guarantee that this interface will remain stable! + - Distribute compiler binaries to accomany cute.compile_to allowing users to build customized compile-exececute pipelines outside of Python. Both static and shared compiler and executor/runtime libraries will be provided. Compiler binaries will be uploaded to GitHub with each release. - Supported AoT cross-compilation for aarch64‑linux‑gnu - Support for two launch attributes: launch completion events (cudaLaunchAttributeLaunchCompletionEvent), for recording an event once all thread blocks have begun executing, and launch programatic events (cudaLaunchAttributeProgrammaticEvent), for PDL event-based synchronization - - Supported auto calculating per-kernel shared memory carveout preference, or use new launch option `preferred_smem_carveout` to set manually. + - Supported auto calculating per-kernel shared memory carveout preference, or use new laucnch option `preferred_smem_carveout` to set manually. - Auto-deduced smem size for launching kernels - Launch config `smem` now defaults to `None` for auto-calculating kernel shared memory usage, which is recommended unless manual control is required. - Warnings will be raised when the manually set shared memory size is insufficient or exceeds the GPU maximum. - The default shared memory usage calculation aligns with CUDA C++ static shared memory behavior, i.e. summing all allocations additively. - An additional launch option `smem_merge_branch_allocs` is provided to merge shared memory allocations across mutually exclusive code branches, which is recommended for inlined mega-kernels to reduce total footprint. + - SASS dumping in DSL is now supported in a self-contained manner - no CUDA toolkit installation required to get nvdisasm * Bug fixing and improvements + - Add the missing elect_one in cute.copy for bulk copy. + - The elect_one required for async bulk copp was missing in cute.copy. It's now generated in cute.copy automatically. + - Nesting elect_one will cause funtionality issues. Please remove elect_one around cute.copy with async bulk copy. + - Elect_one around direct async bulk copy instruction should be kept as it bypasses the cute.copy layer and will not be affected by this fix. + - Affected copy atoms are CopyBulkG2SOp, CopyBulkG2SMulticastOp, CopyBulkS2GOp, CopyBulkS2GByteMaskOp, and CopyBulkS2SOp. + - Example of changes to avoid nesting elect_ones could be found here: https://github.com/Dao-AILab/quack/pull/164 - Improvements on linter support with more type ignores cleaned up - Improvements on tvm-ffi CUDA runtime error diagnostics - Improvements on dataclass support for TVM-FFI @@ -25,6 +35,28 @@ - cute.core.ThrMma, please use cute.ThrMma instead - cute.core.ThrCopy, please use cute.ThrCopy instead - cute.make_fragment, please use cute.make_rmem_tensor instead + - Fixed following issues + - https://github.com/NVIDIA/cutlass/issues/3132 + - https://github.com/NVIDIA/cutlass/issues/3170 + - https://github.com/NVIDIA/cutlass/issues/3259 + - https://github.com/NVIDIA/cutlass/issues/3290 + - https://github.com/NVIDIA/cutlass/issues/3298 + - https://github.com/NVIDIA/cutlass/issues/3312 + - https://github.com/NVIDIA/cutlass/issues/3327 + - https://github.com/NVIDIA/cutlass/issues/3329 + - https://github.com/NVIDIA/cutlass/issues/3342 + - https://github.com/NVIDIA/cutlass/issues/3348 + +### CUTLASS Operator API +* CUTLASS Operator API is a new addition to the CUTLASS Python stack, providing easy interfaces +to discover CUTLASS Python DSL kernels & integrate them in your code. + - `pip install nvidia-cutlass-operators` to get started +* [Operator API Overview](https://docs.nvidia.com/cutlass/latest/media/docs/operators/overview.html) + - See [currently supported kernels](https://docs.nvidia.com/cutlass/latest/master/media/docs/operators/overview.html#supported-and-upcoming-features) +* [Basic GEMM tutorial](https://docs.nvidia.com/cutlass/latest/media/docs/operators/tutorials/000_gemm.html) + - More tutorials [here](https://docs.nvidia.com/cutlass/latest/media/docs/operators/tutorials/index.html) +* [GitHub source](https://github.com/NVIDIA/cutlass/tree/main/operators) +* [API Reference](https://docs.nvidia.com/cutlass/latest/master/media/docs/operators/api_reference/index.html) ### CUTLASS C++ * Add [example 113](https://github.com/NVIDIA/cutlass/tree/main/examples/113_hopper_gemm_activation_fusion) for Hopper GEMM with activation fusion. @@ -32,7 +64,12 @@ - Covers both regular GEMM and grouped GEMM variants. * Improve SM90 grouped/ptr-array GEMM with EVT support. - Adds the EVT (Epilogue Visitor Tree) plumbing required to do activation, bias, and auxiliary-tensor fusion inside SM90 grouped and ptr-array GEMM kernels. +* Add ptr-array TMA collective for tensor/token-scaled FP8 grouped GEMM Blackwell SM120/SM121 kernels. + - Implement `CollectiveMma` and `CollectiveBuilder` specializations for `MainloopSm120ArrayTmaWarpSpecialized`, enabling ptr-array grouped GEMM (MoE expert dispatch) with tensor- and token-level FP8 scaling. + - Corresponding [unit test](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/sm120_tensorop_gemm/sm120_gemm_f8_f8_f32_tensor_op_group_gemm.cu) +* Add tileN = 8,16 for Blackwell SM120 blockscale GEMM kernels. * Fix `DescriptorIterator::operator+` in `mma_traits_sm100.hpp` to use 32-bit arithmetic on CUDA toolkit version <= 13.3, preserving the high half of the smem descriptor. +* Fix a CUDA structured bindings header issue. * Various improvements and fixes from the community and CUTLASS team. Thanks to everyone who submitted PRs! * Optimal code generation with CUDA toolkit versions 13.3. diff --git a/CMakeLists.txt b/CMakeLists.txt index ad2be654e..6b2e290b2 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -771,7 +771,7 @@ if(CUDA_VERSION VERSION_GREATER_EQUAL 13.0) target_include_directories( CUTLASS SYSTEM INTERFACE - $ + $ ) endif() diff --git a/CUDA.cmake b/CUDA.cmake index 5d96279ff..b39e81b74 100644 --- a/CUDA.cmake +++ b/CUDA.cmake @@ -52,6 +52,24 @@ if(NOT CUDA_TOOLKIT_ROOT_DIR) get_filename_component(CUDA_TOOLKIT_ROOT_DIR "${CUDAToolkit_NVCC_EXECUTABLE}/../.." ABSOLUTE) endif() +set(CUTLASS_CUDA_CCCL_INCLUDE_DIR "") +foreach(__CUTLASS_CUDA_INCLUDE_DIR IN LISTS CUDAToolkit_INCLUDE_DIRS CUDA_INCLUDE_DIRS CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES) + if(NOT __CUTLASS_CUDA_INCLUDE_DIR) + continue() + endif() + if(EXISTS "${__CUTLASS_CUDA_INCLUDE_DIR}/cccl/cuda/std/cstdint") + set(CUTLASS_CUDA_CCCL_INCLUDE_DIR "${__CUTLASS_CUDA_INCLUDE_DIR}/cccl") + break() + elseif(EXISTS "${__CUTLASS_CUDA_INCLUDE_DIR}/cuda/std/cstdint") + set(CUTLASS_CUDA_CCCL_INCLUDE_DIR "${__CUTLASS_CUDA_INCLUDE_DIR}") + break() + endif() +endforeach() + +if(NOT CUTLASS_CUDA_CCCL_INCLUDE_DIR) + set(CUTLASS_CUDA_CCCL_INCLUDE_DIR "${CUDA_TOOLKIT_ROOT_DIR}/include/cccl") +endif() + if (CMAKE_CUDA_COMPILER_ID MATCHES "(nvcc|[Nn][Vv][Ii][Dd][Ii][Aa])") set(CUTLASS_NVCC_DEVICE_COMPILE ON CACHE BOOL "Using nvcc tools for device compilation") elseif (CMAKE_CUDA_COMPILER_ID MATCHES "[Cc]lang") diff --git a/README.md b/README.md index 75e590e7f..65bbc3a00 100644 --- a/README.md +++ b/README.md @@ -47,6 +47,9 @@ To get started quickly - please refer : ## CuTe DSL * New features + - New fine-grained compilation API: cute.compile_to that gives control over the what stage the compiler outputs. This feature allows customization of the path from compilation to runtime execution. cute.compile_to is considered experimental in 4.6. + - Experimental Feature: Added the IKET (In-Kernel-Event-Tracing) profiler for instrumentation-based intra-kernel activities tracing. This enables fine-grained profiling and makes it easier to understand persistent, warp-specialized kernels' performance. This is a beta feature provided by CUTLASS Python until a NVIDIA DevTools product is released, there is no guarantee that this interface will remain stable! + - Distribute compiler binaries to accomany cute.compile_to allowing users to build customized compile-exececute pipelines outside of Python. Both static and shared compiler and executor/runtime libraries will be provided. Compiler binaries will be uploaded to GitHub with each release. - Supported AoT cross-compilation for aarch64?~@~Qlinux?~@~Qgnu - Support for two launch attributes: launch completion events (cudaLaunchAttributeLaunchCompletionEvent), for recording an event once all thread blocks have begun executing, and launch programatic events (cudaLaunchAttributeProgrammaticEvent), for PDL event-based synchronization - Supported auto calculating per-kernel shared memory carveout preference, or use new laucnch option `preferred_smem_carveout` to set manually. @@ -55,8 +58,15 @@ To get started quickly - please refer : - Warnings will be raised when the manually set shared memory size is insufficient or exceeds the GPU maximum. - The default shared memory usage calculation aligns with CUDA C++ static shared memory behavior, i.e. summing all allocations additively. - An additional launch option `smem_merge_branch_allocs` is provided to merge shared memory allocations across mutually exclusive code branches, which is recommended for inlined mega-kernels to reduce total footprint. + - SASS dumping in DSL is now supported in a self-contained manner - no CUDA toolkit installation required to get nvdisasm * Bug fixing and improvements + - Add the missing elect_one in cute.copy for bulk copy. + - The elect_one required for async bulk copp was missing in cute.copy. It's now generated in cute.copy automatically. + - Nesting elect_one will cause funtionality issues. Please remove elect_one around cute.copy with async bulk copy. + - Elect_one around direct async bulk copy instruction should be kept as it bypasses the cute.copy layer and will not be affected by this fix. + - Affected copy atoms are CopyBulkG2SOp, CopyBulkG2SMulticastOp, CopyBulkS2GOp, CopyBulkS2GByteMaskOp, and CopyBulkS2SOp. + - Example of changes to avoid nesting elect_ones could be found here: https://github.com/Dao-AILab/quack/pull/164 - Improvements on linter support with more type ignores cleaned up - Improvements on tvm-ffi CUDA runtime error diagnostics - Improvements on dataclass support for TVM-FFI @@ -66,6 +76,28 @@ To get started quickly - please refer : - cute.core.ThrMma, please use cute.ThrMma instead - cute.core.ThrCopy, please use cute.ThrCopy instead - cute.make_fragment, please use cute.make_rmem_tensor instead + - Fixed following issues + - https://github.com/NVIDIA/cutlass/issues/3132 + - https://github.com/NVIDIA/cutlass/issues/3170 + - https://github.com/NVIDIA/cutlass/issues/3259 + - https://github.com/NVIDIA/cutlass/issues/3290 + - https://github.com/NVIDIA/cutlass/issues/3298 + - https://github.com/NVIDIA/cutlass/issues/3312 + - https://github.com/NVIDIA/cutlass/issues/3327 + - https://github.com/NVIDIA/cutlass/issues/3329 + - https://github.com/NVIDIA/cutlass/issues/3342 + - https://github.com/NVIDIA/cutlass/issues/3348 + +### CUTLASS Operator API +* CUTLASS Operator API is a new addition to the CUTLASS Python stack, providing easy interfaces +to discover CUTLASS Python DSL kernels & integrate them in your code. + - `pip install nvidia-cutlass-operators` to get started +* [Operator API Overview](https://docs.nvidia.com/cutlass/latest/media/docs/operators/overview.html) + - See [currently supported kernels](https://docs.nvidia.com/cutlass/latest/master/media/docs/operators/overview.html#supported-and-upcoming-features) +* [Basic GEMM tutorial](https://docs.nvidia.com/cutlass/latest/media/docs/operators/tutorials/000_gemm.html) + - More tutorials [here](https://docs.nvidia.com/cutlass/latest/media/docs/operators/tutorials/index.html) +* [GitHub source](https://github.com/NVIDIA/cutlass/tree/main/operators) +* [API Reference](https://docs.nvidia.com/cutlass/latest/master/media/docs/operators/api_reference/index.html) ## CUTLASS C++ * Add [example 113](https://github.com/NVIDIA/cutlass/tree/main/examples/113_hopper_gemm_activation_fusion) for Hopper GEMM with activation fusion. @@ -73,7 +105,12 @@ To get started quickly - please refer : - Covers both regular GEMM and grouped GEMM variants. * Improve SM90 grouped/ptr-array GEMM with EVT support. - Adds the EVT (Epilogue Visitor Tree) plumbing required to do activation, bias, and auxiliary-tensor fusion inside SM90 grouped and ptr-array GEMM kernels. +* Add ptr-array TMA collective for tensor/token-scaled FP8 grouped GEMM Blackwell SM120/SM121 kernels. + - Implement `CollectiveMma` and `CollectiveBuilder` specializations for `MainloopSm120ArrayTmaWarpSpecialized`, enabling ptr-array grouped GEMM (MoE expert dispatch) with tensor- and token-level FP8 scaling. + - Corresponding [unit test](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/sm120_tensorop_gemm/sm120_gemm_f8_f8_f32_tensor_op_group_gemm.cu) +* Add tileN = 8,16 for Blackwell SM120 blockscale GEMM kernels. * Fix `DescriptorIterator::operator+` in `mma_traits_sm100.hpp` to use 32-bit arithmetic on CUDA toolkit version <= 13.3, preserving the high half of the smem descriptor. +* Fix a CUDA structured bindings header issue. Note: CUTLASS 4.x builds are known to be down on Windows platforms for all CUDA toolkits. CUTLASS team is working on a fix. diff --git a/examples/113_hopper_gemm_activation_fusion/113_hopper_grouped_gemm_fused_gated_act.cu b/examples/113_hopper_gemm_activation_fusion/113_hopper_grouped_gemm_fused_gated_act.cu index f111acbd7..8286c5686 100644 --- a/examples/113_hopper_gemm_activation_fusion/113_hopper_grouped_gemm_fused_gated_act.cu +++ b/examples/113_hopper_gemm_activation_fusion/113_hopper_grouped_gemm_fused_gated_act.cu @@ -87,7 +87,14 @@ using ActivationFn = cutlass::epilogue::thread::Identity; bool constexpr IsFp8 = true; // whether to run with fp8 or fp16 input/output bool constexpr Quantize = true; // whether to quantize output with a per-tensor scale factor -bool constexpr ExactMode = false; // whether to reproduce unfused dual gemm+activation exactly +// Enabled by default here (unlike the other example 113 binaries): in the grouped + gated case +// the fused epilogue accumulates the gated product alpha*acc in f32, while the unfused reference +// rounds it to fp16 (ElementIntermediate) before the activation. After SiLU, gating, and fp8 +// (e4m3) quantization, near-zero outputs can land 1-2 ULPs apart on rare elements (data-dependent +// across the per-group random alphas), which trips the relative-tolerance check. ExactMode makes +// the fused path also use the fp16 intermediate, so fused and reference match bit-exactly and +// verification is deterministic. +bool constexpr ExactMode = true; // whether to reproduce unfused dual gemm+activation exactly bool constexpr BiasBroadcast = true; // whether bias is broadcast along columns in each group bool constexpr Pingpong = true; // whether to use pingpong schedule diff --git a/examples/python/CuTeDSL/cute/ampere/kernel/dense_gemm/tensorop_gemm.py b/examples/python/CuTeDSL/cute/ampere/kernel/dense_gemm/tensorop_gemm.py index dc004b17d..fb1143834 100644 --- a/examples/python/CuTeDSL/cute/ampere/kernel/dense_gemm/tensorop_gemm.py +++ b/examples/python/CuTeDSL/cute/ampere/kernel/dense_gemm/tensorop_gemm.py @@ -28,13 +28,17 @@ import argparse import math +from functools import lru_cache from typing import Tuple, Type +import cuda.bindings.driver as cuda + import cutlass import cutlass.cute as cute -import cutlass.cute.testing as testing +from cutlass import testing import cutlass.utils as utils -from cutlass.cute.runtime import from_dlpack +from cutlass.utils.tensor_helpers import create_cute_tensor_for_fp8 +from cutlass.utils.tensor_helpers import is_fp8_dtype """ A dense GEMM (C = A * B) example for the NVIDIA Ampere architecture using CUTE DSL. @@ -62,7 +66,7 @@ To run this example: .. code-block:: bash - python examples/ampere/tensorop_gemm.py \ + python examples/cute/ampere/kernel/dense_gemm/tensorop_gemm.py \ --mnkl 8192,8192,8192,1 --atom_layout_mnk 2,2,1 \ --ab_dtype Float16 \ --c_dtype Float16 --acc_dtype Float32 \ @@ -77,7 +81,7 @@ To collect performance with NCU profiler: .. code-block:: bash - ncu python examples/ampere/tensorop_gemm.py \ + ncu python examples/cute/ampere/kernel/dense_gemm/tensorop_gemm.py \ --mnkl 8192,8192,8192,1 --atom_layout_mnk 2,2,1 \ --ab_dtype Float16 \ --c_dtype Float16 --acc_dtype Float32 \ @@ -85,9 +89,10 @@ To collect performance with NCU profiler: --skip_ref_check --iterations 2 Constraints: -* Supported input and output data types: fp16 -* Support accumulator data types: f32 -* Default tile shape is set to be 128x128x32 +* Supported input data types: fp16/bf16/fp8 +* Supported output data types: fp16 +* Support accumulator data types: f32/f16 +* Default tile shape is 128x128x32 for fp16/bf16 and 128x128x64 for fp8 * Atom layout's MNK shape is set so that tile shape can be divided by MMA instruction shape * The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, @@ -96,6 +101,13 @@ Constraints: class TensorOpGemm: + _FP16_BF16_DTYPES = (cutlass.Float16, cutlass.BFloat16) + _FP8_DTYPES = (cutlass.Float8E4M3FN, cutlass.Float8E5M2) + _MMA_SHAPE_FP16_BF16 = (16, 8, 16) + _MMA_SHAPE_FP8 = (16, 8, 32) + _CTA_TILER_FP16_BF16 = (128, 128, 32) + _CTA_TILER_FP8 = (128, 128, 64) + def __init__( self, ab_dtype: Type[cutlass.Numeric], @@ -107,14 +119,22 @@ class TensorOpGemm: self.ab_dtype = ab_dtype self.c_dtype = c_dtype self.acc_dtype = acc_dtype - self.cta_tiler = (128, 128, 32) - self.num_stages = 3 + self.is_fp8 = self.ab_dtype in self._FP8_DTYPES + assert self.ab_dtype in self._FP16_BF16_DTYPES + self._FP8_DTYPES, ( + "ab_dtype must be one of Float16, BFloat16, Float8E4M3FN, Float8E5M2" + ) + self.cta_tiler = ( + self._CTA_TILER_FP8 if self.is_fp8 else self._CTA_TILER_FP16_BF16 + ) + self.num_stages = 4 self.atom_layout_mnk = atom_layout_mnk atom_lay_M, atom_lay_N, atom_lay_K = self.atom_layout_mnk self.num_threads = atom_lay_M * atom_lay_N * atom_lay_K * 32 self.bM, self.bN, self.bK = self.cta_tiler - self.mma_inst_shape = (16, 8, 16) + self.mma_inst_shape = ( + self._MMA_SHAPE_FP8 if self.is_fp8 else self._MMA_SHAPE_FP16_BF16 + ) mmaM, mmaN, mmaK = self.mma_inst_shape # M-major C uses C^T = B^T * A^T, swapping atom layout M/N roles. @@ -142,6 +162,7 @@ class TensorOpGemm: mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor, + stream: cuda.CUstream, epilogue_op: cutlass.Constexpr = lambda x: x, ): # The grid divides the problems's M, N, and L dimensions by the @@ -166,7 +187,6 @@ class TensorOpGemm: else: atom_layout_mnk = self.atom_layout_mnk - self.a_major_mode = utils.LayoutEnum.from_tensor(mA) self.b_major_mode = utils.LayoutEnum.from_tensor(mB) self.c_major_mode = utils.LayoutEnum.from_tensor(mC) @@ -183,13 +203,13 @@ class TensorOpGemm: # assume the input is 16B align ab_copy_bits = 128 - sA_layout = self._make_smem_layout_AB( + sA_layout, sA_swizzle = self._make_smem_layout_AB( mA.element_type, self.a_major_mode, ab_copy_bits, (self.cta_tiler[0], self.cta_tiler[2], self.num_stages), ) - sB_layout = self._make_smem_layout_AB( + sB_layout, sB_swizzle = self._make_smem_layout_AB( mB.element_type, self.b_major_mode, ab_copy_bits, @@ -213,9 +233,7 @@ class TensorOpGemm: # Create a copy atom for a global to shared memory asynchronous copy atom_async_copy = cute.make_copy_atom( - cute.nvgpu.cpasync.CopyG2SOp( - cache_mode=cute.nvgpu.cpasync.LoadCacheMode.GLOBAL - ), + cute.nvgpu.cpasync.CopyG2SOp(cache_mode=cute.nvgpu.LoadCacheMode.GLOBAL), mA.element_type, num_bits_per_copy=ab_copy_bits, ) @@ -244,10 +262,14 @@ class TensorOpGemm: # Tiled MMA # /////////////////////////////////////////////////////////////////////////////// - # Creates a mma atom with 16x8x16 shape for MNK - op = cute.nvgpu.warp.MmaF16BF16Op( - self.ab_dtype, self.acc_dtype, self.mma_inst_shape - ) + if cutlass.const_expr(self.is_fp8): + op = cute.nvgpu.warp.MmaFP8Op( + self.ab_dtype, self.acc_dtype, self.mma_inst_shape + ) + else: + op = cute.nvgpu.warp.MmaF16BF16Op( + self.ab_dtype, self.acc_dtype, self.mma_inst_shape + ) permutation_mnk = ( atom_layout_mnk[0] * self.mma_inst_shape[0], @@ -291,7 +313,9 @@ class TensorOpGemm: mB, mC, sA_layout, + sA_swizzle, sB_layout, + sB_swizzle, sC_layout, tiled_copy_A, tiled_copy_B, @@ -302,6 +326,7 @@ class TensorOpGemm: ).launch( grid=rasterization_remap_grid_dim, block=[self.num_threads, 1, 1], + stream=stream, ) @cute.kernel @@ -310,8 +335,10 @@ class TensorOpGemm: mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor, - sA_layout: cute.ComposedLayout, - sB_layout: cute.ComposedLayout, + sA_layout: cute.Layout, + sA_swizzle: cute.Swizzle, + sB_layout: cute.Layout, + sB_swizzle: cute.Swizzle, sC_layout: cute.ComposedLayout, tiled_copy_A: cute.TiledCopy, tiled_copy_B: cute.TiledCopy, @@ -428,8 +455,8 @@ class TensorOpGemm: max(SharedStorageAB.size_in_bytes(), SharedStorageC.size_in_bytes()), byte_alignment=16, ) - sA = SharedStorageAB(storage).a.get_tensor(sA_layout) - sB = SharedStorageAB(storage).b.get_tensor(sB_layout) + sA = SharedStorageAB(storage).a.get_tensor(sA_layout, swizzle=sA_swizzle) + sB = SharedStorageAB(storage).b.get_tensor(sB_layout, swizzle=sB_swizzle) sC = SharedStorageC(storage).c.get_tensor(sC_layout) thr_copy_A = tiled_copy_A.get_slice(tidx) @@ -571,18 +598,42 @@ class TensorOpGemm: # /////////////////////////////////////////////////////////////////////////////// # Create the copy atoms for the copy from shared memory to register - atom_copy_s2r_A = cute.make_copy_atom( - cute.nvgpu.warp.LdMatrix8x8x16bOp( - self.a_major_mode != utils.LayoutEnum.ROW_MAJOR, 4 - ), - mA.element_type, - ) - atom_copy_s2r_B = cute.make_copy_atom( - cute.nvgpu.warp.LdMatrix8x8x16bOp( - self.b_major_mode != utils.LayoutEnum.ROW_MAJOR, 4 - ), - mB.element_type, - ) + if cutlass.const_expr(self.is_fp8): + if cutlass.const_expr(self.a_major_mode == utils.LayoutEnum.ROW_MAJOR): + atom_copy_s2r_A = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x16x8bOp(False, 4), + mA.element_type, + ) + else: + atom_copy_s2r_A = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + mA.element_type, + num_bits_per_copy=8, + ) + if cutlass.const_expr(self.b_major_mode == utils.LayoutEnum.ROW_MAJOR): + atom_copy_s2r_B = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x16x8bOp(False, 4), + mB.element_type, + ) + else: + atom_copy_s2r_B = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + mB.element_type, + num_bits_per_copy=8, + ) + else: + atom_copy_s2r_A = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp( + self.a_major_mode != utils.LayoutEnum.ROW_MAJOR, 4 + ), + mA.element_type, + ) + atom_copy_s2r_B = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp( + self.b_major_mode != utils.LayoutEnum.ROW_MAJOR, 4 + ), + mB.element_type, + ) # Creates the tiled copy so that it matches the thread-value layout # expected by the tiled mma @@ -669,9 +720,7 @@ class TensorOpGemm: tCrB_copy_view[None, None, k_block_next], ) - # Fetch next A: To better interleave global memory access and compute - # instructions, we intentionally use the sequence: copy A, perform GEMM, - # then copy B. + # Fetch next A and B and update smem pipeline read/write if k_block == 0: if k_tile + num_smem_stages - 1 < k_tile_count: cute.copy( @@ -680,18 +729,6 @@ class TensorOpGemm: tAsA[None, None, None, smem_pipe_write], pred=tApA, ) - - # Thread-level register gemm for k_block - cute.gemm( - tiled_mma, - tCrC, - tCrA[None, None, k_block], - tCrB[None, None, k_block], - tCrC, - ) - - # Fetch next B and update smem pipeline read/write - if k_block == 0: if k_tile + num_smem_stages - 1 < k_tile_count: cute.copy( tiled_copy_B, @@ -706,6 +743,15 @@ class TensorOpGemm: if smem_pipe_read == num_smem_stages: smem_pipe_read = 0 + # Thread-level register gemm for k_block + cute.gemm( + tiled_mma, + tCrC, + tCrA[None, None, k_block], + tCrB[None, None, k_block], + tCrC, + ) + # Sync before epilogue cute.arch.cp_async_wait_group(0) cute.arch.sync_threads() @@ -776,23 +822,25 @@ class TensorOpGemm: major_mode_size = ( smem_tiler[1] if major_mode == utils.LayoutEnum.ROW_MAJOR else smem_tiler[0] ) - major_mode_size = 64 if major_mode_size >= 64 else major_mode_size + # Cap at 128 bytes (fp16: 64 elems, fp8: 128 elems) + max_elems = 128 * 8 // dtype.width + major_mode_size = min(major_mode_size, max_elems) swizzle_bits = int(math.log2(major_mode_size * dtype.width // copy_bits)) swizzle_bits = min(swizzle_bits, 3) + # PDSL: base_bits is in bytes (copy_bits / 8), not in elements + base_bits = int(math.log2(copy_bits // 8)) + + shift_bits = int(math.log2(copy_bits // dtype.width)) + swizzle = cute.make_swizzle(swizzle_bits, base_bits, shift_bits) layout_atom_outer = ( cute.make_layout((8, major_mode_size), stride=(major_mode_size, 1)) if major_mode == utils.LayoutEnum.ROW_MAJOR else cute.make_layout((major_mode_size, 8), stride=(1, major_mode_size)) ) - layout_atom = cute.make_composed_layout( - cute.make_swizzle(swizzle_bits, 3, 3), - 0, - layout_atom_outer, - ) - layout = cute.tile_to_shape(layout_atom, smem_tiler, (0, 1, 2)) - return layout + layout = cute.tile_to_shape(layout_atom_outer, smem_tiler, (0, 1, 2)) + return layout, swizzle def _make_smem_layout_C(self, dtype, major_mode, copy_bits, smem_tiler): major_mode_size = ( @@ -872,164 +920,428 @@ class TensorOpGemm: return (new_i, new_j) -def run( +@cute.jit +def bmm( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, # (l, m, k) + b: cute.Tensor, # (l, k, n) + c: cute.Tensor, # (l, m, n) + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, +): + """ + Wrapper API for GEMM kernel to follow the convention of PyTorch's batch matrix-multiply (bmm). + + Internally, the tensors are permuted to match CuTe's convention: + - a: (m, k, l) + - b: (n, k, l) + - c: (m, n, l) + + :param gemm_op: Kernel operation, expects (a, b, c, stream, epilogue_op) + :type gemm_op: cutlass.Constexpr + :param a: Input tensor of shape (l, m, k) + :type a: cute.Tensor + :param b: Input tensor of shape (l, k, n) + :type b: cute.Tensor + :param c: Output tensor of shape (l, m, n) + :type c: cute.Tensor + :param stream: CUDA stream for asynchronous execution + :type stream: cuda.CUstream + :param epilogue_op: Optional elementwise lambda function to apply per output element, defaults to identity + :type epilogue_op: cutlass.Constexpr, optional + """ + # (l,m,k) -> (m,k,l) + a = cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])) + # (l,k,n) -> (n,k,l) + b = cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])) + # (l,m,n) -> (m,n,l) + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + + gemm_op(a, b, c, stream, epilogue_op) + + +@lru_cache(maxsize=1) +def prepare_tensors( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], a_major: str, b_major: str, c_major: str, + init_random: bool = True, +): + """ + Prepare input and output tensors for the GEMM operation. + + :param mnkl: Problem size as a tuple (M, N, K, L). + :type mnkl: Tuple[int, int, int, int] + :param ab_dtype: Data type for input tensors A and B. + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: Data type for output tensor C. + :type c_dtype: Type[cutlass.Numeric] + :param a_major: Major dimension of the A tensor layout ("m" or "k"). + :type a_major: str + :param b_major: Major dimension of the B tensor layout ("n" or "k"). + :type b_major: str + :param c_major: Major dimension of the C tensor layout ("m" or "n"). + :type c_major: str + :param init_random: Whether to initialize tensors with random values, defaults to True. + :type init_random: bool, optional + + :return: Tuple of (a, b, c) torch tensors. + :rtype: Tuple[torch.Tensor, torch.Tensor, torch.Tensor] + """ + import torch + from cutlass.torch import dtype as torch_dtype + + m, n, k, l = mnkl + + if a_major == "k": + a = torch.empty((l, m, k), dtype=torch.float32, device="cuda") + elif a_major == "m": + a = torch.empty((l, k, m), dtype=torch.float32, device="cuda").permute(0, 2, 1) + + if b_major == "n": + b = torch.empty((l, k, n), dtype=torch.float32, device="cuda") + elif b_major == "k": + b = torch.empty((l, n, k), dtype=torch.float32, device="cuda").permute(0, 2, 1) + + if c_major == "n": + c = torch.empty((l, m, n), dtype=torch.float32, device="cuda") + elif c_major == "m": + c = torch.empty((l, n, m), dtype=torch.float32, device="cuda").permute(0, 2, 1) + + if init_random: + a.random_(-2, 3) + b.random_(-2, 3) + c.random_(-2, 3) + + # For fp8 types, use uint8 as storage to avoid dlpack limitation + a_storage_dtype = torch.uint8 if is_fp8_dtype(ab_dtype) else torch_dtype(ab_dtype) + b_storage_dtype = torch.uint8 if is_fp8_dtype(ab_dtype) else torch_dtype(ab_dtype) + c_storage_dtype = torch.uint8 if is_fp8_dtype(c_dtype) else torch_dtype(c_dtype) + + return ( + a.to(dtype=a_storage_dtype), + b.to(dtype=b_storage_dtype), + c.to(dtype=c_storage_dtype), + a.clone(), # fp32 source for fp8 conversion + b.clone(), # fp32 source for fp8 conversion + c.clone(), # fp32 source for fp8 conversion + ) + + +def mark_dynamic_layout( + a, + b, + c, + leading_dim_a: int, + leading_dim_b: int, + leading_dim_c: int, + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_f32=None, + b_f32=None, + c_f32=None, +): + a_ = create_cute_tensor_for_fp8(a, ab_dtype, leading_dim_a, source_f32_tensor=a_f32) + b_ = create_cute_tensor_for_fp8(b, ab_dtype, leading_dim_b, source_f32_tensor=b_f32) + c_ = create_cute_tensor_for_fp8(c, c_dtype, leading_dim_c, source_f32_tensor=c_f32) + + a_.mark_compact_shape_dynamic( + mode=leading_dim_a, + stride_order=(0, 1, 2) if leading_dim_a == 2 else (0, 2, 1), + divisibility=128 // ab_dtype.width, + ) + b_.mark_compact_shape_dynamic( + mode=leading_dim_b, + stride_order=(0, 1, 2) if leading_dim_b == 2 else (0, 2, 1), + divisibility=128 // ab_dtype.width, + ) + c_.mark_compact_shape_dynamic( + mode=leading_dim_c, + stride_order=(0, 1, 2) if leading_dim_c == 2 else (0, 2, 1), + divisibility=128 // c_dtype.width, + ) + return a_, b_, c_ + + +@lru_cache(maxsize=1) +def compile_bmm( + mnkl: Tuple[int, int, int, int], + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], acc_dtype: Type[cutlass.Numeric], - mnkl: Tuple[int, int, int, int], atom_layout_mnk: Tuple[int, int, int], + epilogue_op: cutlass.Constexpr = lambda x: x, +): + """ + Compile the BMM kernel with caching. + + :param mnkl: Problem size as a tuple (M, N, K, L). + :type mnkl: Tuple[int, int, int, int] + :param a: Input tensor A. + :type a: cute.Tensor + :param b: Input tensor B. + :type b: cute.Tensor + :param c: Output tensor C. + :type c: cute.Tensor + :param ab_dtype: Data type for input tensors A and B. + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: Data type for output tensor C. + :type c_dtype: Type[cutlass.Numeric] + :param acc_dtype: Accumulator data type. + :type acc_dtype: Type[cutlass.Numeric] + :param atom_layout_mnk: Atom layout shape (M, N, K). + :type atom_layout_mnk: Tuple[int, int, int] + :param epilogue_op: Optional elementwise lambda function to apply to the output tensor. + :type epilogue_op: cutlass.Constexpr, optional + + :return: Compiled kernel function. + """ + from cutlass.cute.runtime import make_fake_stream + + stream = make_fake_stream() + + is_m_major_c = c.leading_dim == 0 + gemm = TensorOpGemm(ab_dtype, c_dtype, acc_dtype, atom_layout_mnk, is_m_major_c) + return cute.compile(bmm, gemm, a, b, c, stream, epilogue_op) + + +def run( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + atom_layout_mnk: Tuple[int, int, int], + tolerance: float = 1e-03, warmup_iterations: int = 2, iterations: int = 100, skip_ref_check: bool = False, use_cold_l2: bool = False, + benchmark: bool = False, **kwargs, ): + """ + Execute an Ampere tensor core GEMM operation with performance benchmarking. + + Prepares input tensors, configures and launches the GEMM kernel, + optionally performs reference validation, and benchmarks execution. + + :param mnkl: Problem size as a tuple (M, N, K, L). + :type mnkl: Tuple[int, int, int, int] + :param ab_dtype: Data type for input tensors A and B. + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: Data type for output tensor C. + :type c_dtype: Type[cutlass.Numeric] + :param acc_dtype: Accumulator data type for the matrix multiplication. + :type acc_dtype: Type[cutlass.Numeric] + :param a_major: Major dimension of the A tensor layout ("m" or "k"). + :type a_major: str + :param b_major: Major dimension of the B tensor layout ("n" or "k"). + :type b_major: str + :param c_major: Major dimension of the C tensor layout ("m" or "n"). + :type c_major: str + :param atom_layout_mnk: Atom layout shape (M, N, K). + :type atom_layout_mnk: Tuple[int, int, int] + :param tolerance: Tolerance for reference validation, defaults to 1e-03. + :type tolerance: float, optional + :param warmup_iterations: Number of warmup iterations before benchmarking, defaults to 2. + :type warmup_iterations: int, optional + :param iterations: Number of benchmark iterations to run, defaults to 100. + :type iterations: int, optional + :param skip_ref_check: Whether to skip reference result validation, defaults to False. + :type skip_ref_check: bool, optional + :param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False. + :type use_cold_l2: bool, optional + :param benchmark: Whether to only benchmark the kernel, defaults to False. + :type benchmark: bool, optional + :raises RuntimeError: If CUDA GPU is not available. + :return: Execution time of the GEMM kernel. + :rtype: float + """ import torch - import cutlass.torch as cutlass_torch + from cutlass.torch import dtype as torch_dtype - print("Running Ampere tensor core GEMM example:") - print(f"mnkl: {mnkl}") - print( - f"A dtype: {ab_dtype}, B dtype: {ab_dtype}, C dtype: {c_dtype}, Acc dtype: {acc_dtype}" + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + current_stream = cuda.CUstream(torch_stream.cuda_stream) + + # Run and verify BMM with torch + a, b, c, a_f32, b_f32, c_f32 = prepare_tensors( + mnkl, ab_dtype, c_dtype, a_major, b_major, c_major ) - print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") - print(f"Atoms layout: {atom_layout_mnk}") - print(f"Warmup iterations: {warmup_iterations}") - print(f"Iterations: {iterations}") - print(f"Skip reference checking: {skip_ref_check}") - print(f"Use cold L2: {use_cold_l2}") - M, N, K, L = mnkl - # Create and permute tensor A/B/C - def create_and_permute_tensor(l, mode0, mode1, is_mode0_major, dtype): - # is_mode0_major: (l, mode1, mode0) -> (mode0, mode1, l) - # else: (l, mode0, mode1) -> (mode0, mode1, l) - shape = (l, mode1, mode0) if is_mode0_major else (l, mode0, mode1) - permute_order = (2, 1, 0) if is_mode0_major else (1, 2, 0) - torch_tensor = ( - torch.empty(*shape, dtype=torch.int32) - .random_(-2, 2) - .to(dtype=cutlass_torch.dtype(dtype)) - .permute(permute_order) - .cuda() - ) - # assume input is 16B aligned - cute_tensor = ( - from_dlpack(torch_tensor, assumed_align=16) - .mark_layout_dynamic(leading_dim=(1 if not is_mode0_major else 0)) - .mark_compact_shape_dynamic( - mode=(1 if not is_mode0_major else 0), - stride_order=(2, 0, 1) if not is_mode0_major else (2, 1, 0), - divisibility=(128 // dtype.width), - ) - ) - return cute_tensor, torch_tensor + leading_dim_a = 2 if a_major == "k" else 1 + leading_dim_b = 1 if b_major == "k" else 2 + leading_dim_c = 2 if c_major == "n" else 1 - mA, a_torch = create_and_permute_tensor(L, M, K, a_major == "m", ab_dtype) - mB, b_torch = create_and_permute_tensor(L, N, K, b_major == "n", ab_dtype) - mC, c_torch = create_and_permute_tensor(L, M, N, c_major == "m", c_dtype) - - is_m_major_c = c_major == "m" + a_, b_, c_ = mark_dynamic_layout( + a, + b, + c, + leading_dim_a, + leading_dim_b, + leading_dim_c, + ab_dtype, + c_dtype, + a_f32, + b_f32, + c_f32, + ) - tensor_op_gemm = TensorOpGemm( + compiled_fn = compile_bmm( + mnkl, + a_, + b_, + c_, ab_dtype, c_dtype, acc_dtype, atom_layout_mnk, - is_m_major_c + epilogue_op=lambda x: x, ) - print("Compiling kernel with cute.compile ...") - compiled_gemm = cute.compile(tensor_op_gemm, mA, mB, mC) - - print("Executing GEMM kernel...") + print("Running Ampere tensor core GEMM test with:") + print(f"mnkl: {mnkl}") + print(f"Tolerance: {tolerance}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Skip reference checking: {skip_ref_check}") + print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") if not skip_ref_check: - ref = torch.einsum( - "mkl,nkl->mnl", - a_torch.to(dtype=torch.float32), - b_torch.to(dtype=torch.float32), - ).to(cutlass_torch.dtype(c_dtype)) - compiled_gemm(mA, mB, mC) - print("Verifying results...") - torch.testing.assert_close(c_torch.cpu(), ref.cpu(), atol=1e-03, rtol=1e-05) - print("Results verified successfully!") + # Use small random number for deterministic result for reference check + compiled_fn(a_, b_, c_, current_stream) + + # Manually quantize to be comparable + # For fp8 types, use f32 source tensors for reference computation + # since a/b/c may be stored as uint8 + a_ref = a_f32 if is_fp8_dtype(ab_dtype) else a + b_ref = b_f32 if is_fp8_dtype(ab_dtype) else b + ref = ( + torch.bmm(a_ref.to(dtype=torch.float32), b_ref.to(dtype=torch.float32)) + .to(dtype=torch_dtype(c_dtype)) + .to(dtype=torch.float32) + ) + torch.testing.assert_close( + c.to(dtype=torch.float32), ref, atol=tolerance, rtol=1e-05 + ) + + if not benchmark: + return 0 def generate_tensors(): - a_workspace, _ = create_and_permute_tensor(L, M, K, a_major == "m", ab_dtype) - b_workspace, _ = create_and_permute_tensor(L, N, K, b_major == "n", ab_dtype) - c_workspace, _ = create_and_permute_tensor(L, M, N, c_major == "m", c_dtype) - return testing.JitArguments(a_workspace, b_workspace, c_workspace) + a, b, c, a_f32, b_f32, c_f32 = prepare_tensors( + mnkl, + ab_dtype, + c_dtype, + a_major, + b_major, + c_major, + init_random=True, + ) + a_, b_, c_ = mark_dynamic_layout( + a, + b, + c, + leading_dim_a, + leading_dim_b, + leading_dim_c, + ab_dtype, + c_dtype, + a_f32, + b_f32, + c_f32, + ) + return testing.JitArguments(a_, b_, c_, current_stream) workspace_count = 1 if use_cold_l2: one_workspace_bytes = ( - a_torch.numel() * a_torch.element_size() - + b_torch.numel() * b_torch.element_size() - + c_torch.numel() * c_torch.element_size() + a.numel() * a.element_size() + + b.numel() * b.element_size() + + c.numel() * c.element_size() ) workspace_count = testing.get_workspace_count( one_workspace_bytes, warmup_iterations, iterations ) - avg_time_us = testing.benchmark( - compiled_gemm, + # Return execution time in microseconds + exec_time = testing.benchmark( + compiled_fn, workspace_generator=generate_tensors, workspace_count=workspace_count, + stream=current_stream, warmup_iterations=warmup_iterations, iterations=iterations, - use_cuda_graphs=False, ) - - return avg_time_us # Return execution time in microseconds + print(f"[DSL INFO] Execution time: {exec_time} microseconds per iteration") + return exec_time -if __name__ == "__main__": +def _parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) - def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: - try: - return tuple(int(x.strip()) for x in s.split(",")) - except ValueError: - raise argparse.ArgumentTypeError( - "Invalid format. Expected comma-separated integers." - ) - parser = argparse.ArgumentParser( - description="example of multistage block matmul with CuTe on GPU" +def prepare_parser(): + parser = argparse.ArgumentParser(description="Example of Dense GEMM on Ampere.") + + parser.add_argument( + "--mnkl", + type=_parse_comma_separated_ints, + default=(256, 256, 512, 1), + help="mnkl dimensions (comma-separated)", ) parser.add_argument( - "--mnkl", type=parse_comma_separated_ints, default=(112, 136, 40, 1) + "--atom_layout_mnk", + type=_parse_comma_separated_ints, + default=(2, 2, 1), + help="Atom layout (comma-separated)", + ) + parser.add_argument("--ab_dtype", type=cutlass.dtype, default=cutlass.Float16) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.Float16) + parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument("--a_major", choices=["k", "m"], type=str, default="m") + parser.add_argument("--b_major", choices=["k", "n"], type=str, default="n") + parser.add_argument("--c_major", choices=["n", "m"], type=str, default="n") + parser.add_argument( + "--tolerance", type=float, default=1e-03, help="Tolerance for validation" ) parser.add_argument( - "--atom_layout_mnk", type=parse_comma_separated_ints, default=(2, 2, 1) + "--benchmark", + type=str, + default="default", + choices=["default", "none"], + help="Benchmark the kernel with default (cutlass.testing.benchmark) or none", ) parser.add_argument( - "--ab_dtype", - type=cutlass.dtype, - choices=[cutlass.Float16], - default=cutlass.Float16, + "--warmup_iterations", type=int, default=2, help="Warmup iterations" ) parser.add_argument( - "--acc_dtype", - type=cutlass.dtype, - choices=[cutlass.Float32], - default=cutlass.Float32, + "--iterations", + type=int, + default=100, + help="Number of iterations to run the kernel", ) parser.add_argument( - "--c_dtype", - type=cutlass.dtype, - choices=[cutlass.Float16], - default=cutlass.Float16, + "--skip_ref_check", action="store_true", help="Skip reference checking" ) - parser.add_argument("--a_major", choices=["k", "m"], default="m") - parser.add_argument("--b_major", choices=["k", "n"], default="n") - parser.add_argument("--c_major", choices=["n", "m"], default="n") - parser.add_argument("--warmup_iterations", default=2, type=int) - parser.add_argument("--iterations", default=100, type=int) - parser.add_argument("--skip_ref_check", action="store_true") parser.add_argument( "--use_cold_l2", action="store_true", @@ -1037,19 +1349,43 @@ if __name__ == "__main__": help="Use circular buffer tensor sets to ensure L2 cold cache", ) + return parser + + +if __name__ == "__main__": + parser = prepare_parser() + args = parser.parse_args() + + if len(args.mnkl) != 4: + parser.error("--mnkl must contain exactly 4 values") + + if len(args.atom_layout_mnk) != 3: + parser.error("--atom_layout_mnk must contain exactly 3 values") + + print("[DSL INFO] Compiling Ampere Dense GEMM with:") + print( + f"[DSL INFO] A dtype: {args.ab_dtype}, B dtype: {args.ab_dtype}, C dtype: {args.c_dtype}, Acc dtype: {args.acc_dtype}" + ) + print( + f"[DSL INFO] Matrix majors - A: {args.a_major}, B: {args.b_major}, C: {args.c_major}" + ) + print(f"[DSL INFO] Atom layout (M, N, K): {args.atom_layout_mnk}") + run( - args.a_major, - args.b_major, - args.c_major, + args.mnkl, args.ab_dtype, args.c_dtype, args.acc_dtype, - args.mnkl, + args.a_major, + args.b_major, + args.c_major, args.atom_layout_mnk, + args.tolerance, args.warmup_iterations, args.iterations, args.skip_ref_check, args.use_cold_l2, + args.benchmark == "default", ) print("PASS") diff --git a/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py b/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py index d7887f009..a717df31b 100644 --- a/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py +++ b/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py @@ -3255,7 +3255,7 @@ class BlackwellFusedMultiHeadAttentionForward: if cutlass.const_expr(mLSE is not None): scaled_tmp = scale_softmax * tTMEM_LOAD_VECrS[1] - # Convert LSE from natural log to log2 space, consistent with flashinfer trtllm-gen backend + # Convert LSE from natural log to log2 space, consistent with flashinfer backend lse = (cute.math.log(row_sum, fastmath=True) + scaled_tmp) * Float32( 1.4426950408889634 ) @@ -3968,7 +3968,7 @@ def run( cur_p = torch.softmax(cur_s_with_sink, dim=1)[:, :cur_s_k, :, :] else: if lse_calculation: - # Use log2 space for LSE, consistent with flashinfer trtllm-gen backend + # Use log2 space for LSE, consistent with flashinfer backend cur_lse = ( torch.logsumexp(cur_s, dim=1) * 1.4426950408889634 ) # reduce over s_k diff --git a/examples/python/CuTeDSL/dsl_tutorials/export/run_with_dynamic_loading.sh b/examples/python/CuTeDSL/dsl_tutorials/export/run_with_dynamic_loading.sh index 571d06b37..98305415e 100755 --- a/examples/python/CuTeDSL/dsl_tutorials/export/run_with_dynamic_loading.sh +++ b/examples/python/CuTeDSL/dsl_tutorials/export/run_with_dynamic_loading.sh @@ -39,7 +39,25 @@ else echo "nvidia-cutlass-dsl wheel path found at: $WHEEL_PATH" fi -CUTE_DSL_LIB_PATH="${WHEEL_PATH}/lib/" +# The wheel hierarchy ships the runtime under ${WHEEL_PATH}/cuN/{lib,include}. +# Pick the highest cuN flavor not newer than the system CUDA major version. +# nvidia-smi reports either "CUDA Version" or "CUDA UMD Version" (casing/spacing varies). +CUDA_MAJOR_VERSION=$(nvidia-smi --version | grep -oP 'CUDA(?: UMD)? (V|v)ersion\s*:\s*\K[0-9]+' | head -1) +TARGET_CU="" +for CU_VERSION in $(ls "$WHEEL_PATH" | grep -oP 'cu\K[0-9]+' | sort -nr); do + if [ "$CU_VERSION" -le "$CUDA_MAJOR_VERSION" ]; then + TARGET_CU="cu${CU_VERSION}" + break + fi +done +if [ -z "$TARGET_CU" ]; then + echo "No compatible cuN runtime found in $WHEEL_PATH for CUDA major version ${CUDA_MAJOR_VERSION}" + exit 1 +fi +echo "Using DSL runtime flavor: $TARGET_CU" + +CUTE_DSL_LIB_PATH="${WHEEL_PATH}/${TARGET_CU}/lib" +CUTE_DSL_INCLUDE_PATH="${WHEEL_PATH}/${TARGET_CU}/include" export LD_LIBRARY_PATH=${CUTE_DSL_LIB_PATH}:${CUDA_HOME}/lib64:./build if [ -z "$CUDA_HOME" ]; then @@ -67,7 +85,7 @@ fi $CXX -o build/run_with_dynamic_loading \ -I${CUDA_HOME}/include \ - -I${WHEEL_PATH}/include \ + -I${CUTE_DSL_INCLUDE_PATH} \ ${SOURCE_FILE} \ -L${CUTE_DSL_LIB_PATH} \ -L${CUDA_HOME}/lib64 \ diff --git a/examples/python/CuTeDSL/dsl_tutorials/export/run_with_static_linking.sh b/examples/python/CuTeDSL/dsl_tutorials/export/run_with_static_linking.sh index c3540d73c..ec867c34c 100755 --- a/examples/python/CuTeDSL/dsl_tutorials/export/run_with_static_linking.sh +++ b/examples/python/CuTeDSL/dsl_tutorials/export/run_with_static_linking.sh @@ -38,7 +38,24 @@ if [[ -z "$WHEEL_PATH" ]]; then else echo "nvidia-cutlass-dsl wheel path found at: $WHEEL_PATH" fi -CUTE_DSL_LIB_PATH="${WHEEL_PATH}/lib/" +# The wheel hierarchy ships the runtime under ${WHEEL_PATH}/cuN/{lib,include}. +# Pick the highest cuN flavor not newer than the system CUDA major version. +# nvidia-smi reports either "CUDA Version" or "CUDA UMD Version" (casing/spacing varies). +CUDA_MAJOR_VERSION=$(nvidia-smi --version | grep -oP 'CUDA(?: UMD)? (V|v)ersion\s*:\s*\K[0-9]+' | head -1) +TARGET_CU="" +for CU_VERSION in $(ls "$WHEEL_PATH" | grep -oP 'cu\K[0-9]+' | sort -nr); do + if [ "$CU_VERSION" -le "$CUDA_MAJOR_VERSION" ]; then + TARGET_CU="cu${CU_VERSION}" + break + fi +done +if [ -z "$TARGET_CU" ]; then + echo "No compatible cuN runtime found in $WHEEL_PATH for CUDA major version ${CUDA_MAJOR_VERSION}" + exit 1 +fi +echo "Using DSL runtime flavor: $TARGET_CU" + +CUTE_DSL_LIB_PATH="${WHEEL_PATH}/${TARGET_CU}/lib" export LD_LIBRARY_PATH=${CUTE_DSL_LIB_PATH} if [ -z "$CUDA_HOME" ]; then diff --git a/include/cute/arch/mma_sm100_desc.hpp b/include/cute/arch/mma_sm100_desc.hpp index 887b58967..71ed59470 100644 --- a/include/cute/arch/mma_sm100_desc.hpp +++ b/include/cute/arch/mma_sm100_desc.hpp @@ -45,6 +45,9 @@ #include #include // cute::array +#include // Tensor, make_tensor, recast, Layout, logical_divide, etc. +#include // get_swizzle_t + //////////////////////////////////////////////////////////////////////////////////////////////////// namespace cute { @@ -263,6 +266,7 @@ enum class BMatrixBufferId : uint8_t { }; enum class BMatrixBufferReuse : uint8_t { + None = 0u, Keep = 1u, Reuse = 2u, ReuseAndKeep = 3u @@ -647,5 +651,281 @@ make_runtime_instr_desc_block_scaled(UMMA::InstrDescriptorBlockScaled desc_i, return idescE; } +/////////////////////////////////////////////////////////////////////////////// +// Tensor (position-dependent swizzle) to LayoutType utility +/////////////////////////////////////////////////////////////////////////////// + +template +CUTE_HOST_DEVICE constexpr +LayoutType +layout_type(Tensor> const&) +{ + static_assert(is_same::value, + "Expected uint128_t type in LayoutType conversion."); + + using Swizzle = get_swizzle_t; + constexpr int B = Swizzle::num_bits; + constexpr int M = Swizzle::num_base; + constexpr int S = Swizzle::num_shft; + + if constexpr (M == 4) { + static_assert(S == 3, "Expected S = 3 when M == 4. Unsupported layout swizzle."); + switch (B) { + default: static_assert(0 <= B && B <= 3, "Expected B = 0,1,2, or 3 when M == 4. Unsupported layout swizzle."); + case 0: return LayoutType::SWIZZLE_NONE; + case 1: return LayoutType::SWIZZLE_32B; + case 2: return LayoutType::SWIZZLE_64B; + case 3: return LayoutType::SWIZZLE_128B; + } + } else + if constexpr (M == 5) { + static_assert(B == 2, "Expected B = 2 when M == 5. Unsupported layout swizzle."); + static_assert(S == 2, "Expected S = 2 when M == 5. Unsupported layout swizzle."); + return LayoutType::SWIZZLE_128B_BASE32B; + } else { + static_assert(M==5, "Only 16B and 32B Atoms are supported for UMMA. Unsupported layout swizzle."); + return LayoutType::SWIZZLE_NONE; // ERROR + } +} + +/////////////////////////////////////////////////////////////////////////////// +// Construction method for UMMA Descriptors +/////////////////////////////////////////////////////////////////////////////// + +/** +* /////////////////////////////// +* // make_umma_desc // +* /////////////////////////////// +* Each UmmaDescriptor Major-MN describes a canonical layout of the form +* +* LayoutType::INTERLEAVE : Swizzle<0,4,3> o smem_ptr o ((T,1,m),(8,k)):((1,T,SBO),(1T,LBO)) +* LayoutType::B32 : Swizzle<1,4,3> o smem_ptr o ((T,2,m),(8,k)):((1,T,LBO),(2T,SBO)) +* LayoutType::B64 : Swizzle<2,4,3> o smem_ptr o ((T,4,m),(8,k)):((1,T,LBO),(4T,SBO)) +* LayoutType::B128 : Swizzle<3,4,3> o smem_ptr o ((T,8,m),(8,k)):((1,T,LBO),(8T,SBO)) +* LayoutType::128B_BASE32B : Swizzle<2,5,2> o smem_ptr o ((T,8,m),(4,k)):((1,T,LBO),(?T,SBO)) +* +* where +* T : sizeof(uint128_t) / sizeof(value_type) +* m : integer in [1,16] corresponding to UMMA shape +* k : integer in [1,32] corresponding to UMMA shape +* SBO: stride byte offset +* LBO: leading byte offset +* +* See UMMA::Layout_MN_XXX_Atom for building canonical UmmaDescriptor Major-MN layouts. +* For example, +* auto smem_layout = tile_to_shape(Layout_MN_SW128_Atom{}, Shape<_128,_64>{}); +* is guaranteed to be accepted by make_umma_desc for appropriate value_type. +* +* ////////////////////////////// +* // make_umma_desc // +* ////////////////////////////// +* Each UmmaDescriptor Major-K describes a canonical layout of the form +* +* LayoutType::INTERLEAVE : Swizzle<0,4,3> o smem_ptr o ((8,m),(T,2)):((1T,SBO),(1,LBO)) +* LayoutType::B32 : Swizzle<1,4,3> o smem_ptr o ((8,m),(T,2)):((2T,SBO),(1, T )) +* LayoutType::B64 : Swizzle<2,4,3> o smem_ptr o ((8,m),(T,2)):((4T,SBO),(1, T )) +* LayoutType::B128 : Swizzle<3,4,3> o smem_ptr o ((8,m),(T,2)):((8T,SBO),(1, T )) +* +* See UMMA::Layout_K_XXX_Atom for building canonical UmmaDescriptor Major-K layouts. +* For example, +* auto smem_layout = tile_to_shape(Layout_K_SW128_Atom{}, Shape<_128,_64>{}); +* is guaranteed to be accepted by make_umma_desc for appropriate value_type. +*/ +template +CUTE_HOST_DEVICE constexpr +SmemDescriptor +make_umma_desc(Tensor const& tensor) +{ + static_assert(is_smem::value, "UMMA Descriptors can only be constructed on smem."); + static_assert(TLayout::rank == 2, "UMMA Descriptors can only be constructed on rank-2 tensors."); + + Tensor u128_tensor = recast(tensor); + + // Result + SmemDescriptor desc; + desc.version_ = 1; // Set the version for blackwell + desc.lbo_mode_ = 0; // set to legacy mode by default + + // Layout type + constexpr UMMA::LayoutType LAYOUT_TYPE = UMMA::layout_type(u128_tensor); + desc.layout_type_ = uint8_t(LAYOUT_TYPE); + + // Start address (4LSB not included) + uint32_t start_address = cast_smem_ptr_to_uint(raw_pointer_cast(u128_tensor.data())); + desc.start_address_ = static_cast(start_address >> 4); + + constexpr uint8_t base_offset = 0; + desc.base_offset_ = base_offset; + + // LayoutType meta + constexpr int SwizzleAtomMNSize = LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_NONE ? 1 : + LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_32B ? 2 : + LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_64B ? 4 : + LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_128B ? 8 : + LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_128B_BASE32B ? 8 : -1; + + if constexpr (MajorMode == UMMA::Major::MN) + { + /* In units of uint128_t, each UmmaDescriptor Major-MN describes a canonical layout of the form + * + * LayoutType::INTERLEAVE : Swizzle<0,4,3> o smem_ptr o ((1,n),(8,k)):((X,SBO),(1,LBO)) + * LayoutType::B32 : Swizzle<1,4,3> o smem_ptr o ((2,n),(8,k)):((1,LBO),(2,SBO)) + * LayoutType::B64 : Swizzle<2,4,3> o smem_ptr o ((4,n),(8,k)):((1,LBO),(4,SBO)) + * LayoutType::B128 : Swizzle<3,4,3> o smem_ptr o ((8,n),(8,k)):((1,LBO),(8,SBO)) + * LayoutType::B128_BASE32B : Swizzle<2,5,2> o smem_ptr o ((8,n),(4,k)):((1,LBO),(4,SBO)) + */ + + constexpr int SwizzleAtomKSize = LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_128B_BASE32B ? 4 : 8; + + // Construct the canonical UMMA T Layout with shape ((SwizzleAtomMNSize,n),(SwizzleAtomKSize,2)) + Layout canonical_layout = logical_divide(layout(u128_tensor), Tile>,Layout>>{}); + + // Check profile of canonical + CUTE_STATIC_ASSERT_V(congruent(canonical_layout, Shape,Shape<_1,_1>>{}), "Not a canonical UMMA_MN Layout: Expected profile failure."); + // Check canonical mode strides + constexpr uint32_t stride_00 = stride<0,0>(canonical_layout); + constexpr uint32_t expected_stride_00 = LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_NONE ? stride<0,0>(canonical_layout) : 1; + static_assert(stride_00 == expected_stride_00, "Not a canonical UMMA_MN Layout: Expected stride failure."); + constexpr uint32_t stride_10 = stride<1,0>(canonical_layout); + constexpr uint32_t expected_stride_10 = SwizzleAtomMNSize; + static_assert(stride_10 == expected_stride_10, "Not a canonical UMMA_MN Layout: Expected stride failure."); + + // stride dimension byte offset and leading dimension byte offset (4LSB not included == uint128_t units) + constexpr uint32_t stride_01 = stride<0,1>(canonical_layout); + constexpr uint32_t stride_11 = stride<1,1>(canonical_layout); + + desc.stride_byte_offset_ = (LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_NONE) ? stride_01 : stride_11; + desc.leading_byte_offset_ = (LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_NONE) ? stride_11 : stride_01; + } else + if constexpr (MajorMode == UMMA::Major::K) + { + /* In units of uint128_t, each UmmaDescriptor Major-K describes a canonical layout of the form + * + * LayoutType::INTERLEAVE : Swizzle<0,4,3> o smem_ptr o ((8,n),2):((1,SBO),LBO) + * LayoutType::B32 : Swizzle<1,4,3> o smem_ptr o ((8,n),2):((2,SBO),1) + * LayoutType::B64 : Swizzle<2,4,3> o smem_ptr o ((8,n),2):((4,SBO),1) + * LayoutType::B128 : Swizzle<3,4,3> o smem_ptr o ((8,n),2):((8,SBO),1) + * LayoutType::B128_BASE32B : Not applicable for Major-K + */ + + static_assert(LAYOUT_TYPE != UMMA::LayoutType::SWIZZLE_128B_BASE32B, "SWIZZLE_128B_BASE32B is invalid for Major-K"); + CUTE_STATIC_ASSERT_V(size<0>(u128_tensor) % Int<8>{} == Int<0>{}, // N|M size + "Not a canonical UMMA_K Layout: Expected MN-size multiple of 8."); + + // Construct the canonical UMMA N Layout with shape ((8,n),(2,1)) + Layout canonical_layout = logical_divide(layout(u128_tensor), Tile,Layout<_2,_1>>{}); + + // Check profile of canonical + CUTE_STATIC_ASSERT_V(congruent(canonical_layout, Shape,Shape<_1,_1>>{}), "Not a canonical UMMA_K Layout: Expected profile failure."); + // Check canonical mode strides + constexpr uint32_t stride_00 = stride<0,0>(canonical_layout); + constexpr uint32_t expected_stride_00 = SwizzleAtomMNSize; + static_assert(stride_00 == expected_stride_00, "Not a canonical UMMA_K Layout: Expected stride failure."); + constexpr uint32_t stride_10 = stride<1,0>(canonical_layout); + constexpr uint32_t expected_stride_10 = (LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_NONE) ? stride<1,0>(canonical_layout) : 1; + static_assert(stride_10 == expected_stride_10, "Not a canonical UMMA_K Layout: Expected stride failure."); + + // stride dimension byte offset and leading dimension byte offset (4LSB not included == uint128_t units) + constexpr uint32_t stride_01 = stride<0,1>(canonical_layout); + + desc.stride_byte_offset_ = stride_01; + desc.leading_byte_offset_ = stride_10; + } else { + static_assert(MajorMode != UMMA::Major::MN && MajorMode != UMMA::Major::K, "Unrecognized MajorMode!"); + } + return desc; +} + +/////////////////////////////////////////////////////////////////////////////// +// Higher level UMMA Descriptor utilities +/////////////////////////////////////////////////////////////////////////////// + +struct DescriptorIterator +{ + using reference = SmemDescriptor; + using element_type = SmemDescriptor; + using value_type = SmemDescriptor; + + SmemDescriptor desc_; + + // Dereference returns the UmmaDescriptor + CUTE_HOST_DEVICE constexpr + reference operator*() const { return desc_; } + + // Advance and return a new UmmaDescriptor + template + CUTE_HOST_DEVICE constexpr + reference operator[](Index const& i) const { return *(*this + i); } + + // Return an advanced iterator + template + CUTE_HOST_DEVICE constexpr + DescriptorIterator operator+(Index const& offset) const + { +#if (__CUDACC_VER_MAJOR__ > 13) || ((__CUDACC_VER_MAJOR__ == 13) && (__CUDACC_VER_MINOR__ > 3)) + return { desc_ + uint64_t(offset) }; +#else + // Use 32-bit calculation rather than 64-bit calculation as we only update part of desc. + SmemDescriptor ret; + ret.lo = desc_.lo + uint32_t(offset); + ret.hi = desc_.hi; + return { ret }; +#endif + } +}; + +template +CUTE_HOST_DEVICE constexpr +SmemDescriptor +raw_pointer_cast(DescriptorIterator const& ptr) { + return ptr.desc_; +} + +CUTE_HOST_DEVICE void +print(DescriptorIterator const&) { + printf("UMMA::DescriptorIterator"); +} + +// Flag for smem descriptor allocation/creation +template +struct smem_desc : DescriptorIterator {}; + +template +struct sparse_smem_desc : DescriptorIterator {}; + } // end namespace UMMA + +// Customization point for creating a UMMA::smem_desc Tensor +template +struct MakeTensor> +{ + template + CUTE_HOST_DEVICE constexpr auto + operator()(Tensor const& smem_tensor) + { + static_assert(is_smem::value, "Expected SMEM Tensor to construct a UMMA Desc Tensor"); + return make_tensor(UMMA::DescriptorIterator{UMMA::make_umma_desc(tensor<0>(smem_tensor))}, + replace<0>(recast(smem_tensor).layout(), Layout<_1,_0>{})); + } +}; + +// Customization point for creating a UMMA::sparse_smem_desc Tensor +template +struct MakeTensor> +{ + // Note that this is the exact same as UMMA::smem_desc above. + // Only the interface validates that we are passed a sparse_ptr, which is recast away to construct + // the smem desc tensor + template + CUTE_HOST_DEVICE constexpr auto + operator()(Tensor const& smem_tensor) + { + static_assert(is_smem::value, "Expected SMEM Tensor to construct a UMMA Desc Tensor"); + static_assert(is_sparse::value, "Expected sparse value_type."); + static_assert(is_sparse_ptr::value, "Expected sparse iter."); + return make_tensor(UMMA::DescriptorIterator{UMMA::make_umma_desc(tensor<0>(smem_tensor))}, + replace<0>(recast(smem_tensor).layout(), Layout<_1,_0>{})); + } +}; + } // namespace cute diff --git a/include/cute/atom/mma_traits_sm100.hpp b/include/cute/atom/mma_traits_sm100.hpp index a93047a66..d70b3a635 100644 --- a/include/cute/atom/mma_traits_sm100.hpp +++ b/include/cute/atom/mma_traits_sm100.hpp @@ -45,6 +45,7 @@ #include // UTCCP smem desc #include +#include // UMMA::TmemAllocMode, tmem_frg, tmem_e_frg, tmem_sf_frg, tmem_frg_ws // Check that aggregate initialization in .with() initializes all fields #if defined(__GNUG__) @@ -121,284 +122,8 @@ tile_to_mma_shape(LayoutAtom const& atom, MMATileShape const& mma_tile_shape, Mo return tiled_divide(mn_tiled, product_each(shape<0>(mma_tile_shape))); // ((MMA_M,MMA_N),M_MMAs,N_MMAs,...) } -// -// Tensor (position-dependent swizzle) to LayoutType utility -// - -template -CUTE_HOST_DEVICE constexpr -LayoutType -layout_type(Tensor> const&) -{ - static_assert(is_same::value, - "Expected uint128_t type in LayoutType conversion."); - - using Swizzle = get_swizzle_t; - constexpr int B = Swizzle::num_bits; - constexpr int M = Swizzle::num_base; - constexpr int S = Swizzle::num_shft; - - if constexpr (M == 4) { - static_assert(S == 3, "Expected S = 3 when M == 4. Unsupported layout swizzle."); - switch (B) { - default: static_assert(0 <= B && B <= 3, "Expected B = 0,1,2, or 3 when M == 4. Unsupported layout swizzle."); - case 0: return LayoutType::SWIZZLE_NONE; - case 1: return LayoutType::SWIZZLE_32B; - case 2: return LayoutType::SWIZZLE_64B; - case 3: return LayoutType::SWIZZLE_128B; - } - } else - if constexpr (M == 5) { - static_assert(B == 2, "Expected B = 2 when M == 5. Unsupported layout swizzle."); - static_assert(S == 2, "Expected S = 2 when M == 5. Unsupported layout swizzle."); - return LayoutType::SWIZZLE_128B_BASE32B; - } else { - static_assert(M==5, "Only 16B and 32B Atoms are supported for UMMA. Unsupported layout swizzle."); - return LayoutType::SWIZZLE_NONE; // ERROR - } -} - -/////////////////////////////////////////////////////////////////////////////// -// Construction method for UMMA Descriptors -/////////////////////////////////////////////////////////////////////////////// - -/** -* /////////////////////////////// -* // make_umma_desc // -* /////////////////////////////// -* Each UmmaDescriptor Major-MN describes a canonical layout of the form -* -* LayoutType::INTERLEAVE : Swizzle<0,4,3> o smem_ptr o ((T,1,m),(8,k)):((1,T,SBO),(1T,LBO)) -* LayoutType::B32 : Swizzle<1,4,3> o smem_ptr o ((T,2,m),(8,k)):((1,T,LBO),(2T,SBO)) -* LayoutType::B64 : Swizzle<2,4,3> o smem_ptr o ((T,4,m),(8,k)):((1,T,LBO),(4T,SBO)) -* LayoutType::B128 : Swizzle<3,4,3> o smem_ptr o ((T,8,m),(8,k)):((1,T,LBO),(8T,SBO)) -* LayoutType::128B_BASE32B : Swizzle<2,5,2> o smem_ptr o ((T,8,m),(4,k)):((1,T,LBO),(?T,SBO)) -* -* where -* T : sizeof(uint128_t) / sizeof(value_type) -* m : integer in [1,16] corresponding to UMMA shape -* k : integer in [1,32] corresponding to UMMA shape -* SBO: stride byte offset -* LBO: leading byte offset -* -* See UMMA::Layout_MN_XXX_Atom for building canonical UmmaDescriptor Major-MN layouts. -* For example, -* auto smem_layout = tile_to_shape(Layout_MN_SW128_Atom{}, Shape<_128,_64>{}); -* is guaranteed to be accepted by make_umma_desc for appropriate value_type. -* -* ////////////////////////////// -* // make_umma_desc // -* ////////////////////////////// -* Each UmmaDescriptor Major-K describes a canonical layout of the form -* -* LayoutType::INTERLEAVE : Swizzle<0,4,3> o smem_ptr o ((8,m),(T,2)):((1T,SBO),(1,LBO)) -* LayoutType::B32 : Swizzle<1,4,3> o smem_ptr o ((8,m),(T,2)):((2T,SBO),(1, T )) -* LayoutType::B64 : Swizzle<2,4,3> o smem_ptr o ((8,m),(T,2)):((4T,SBO),(1, T )) -* LayoutType::B128 : Swizzle<3,4,3> o smem_ptr o ((8,m),(T,2)):((8T,SBO),(1, T )) -* -* See UMMA::Layout_K_XXX_Atom for building canonical UmmaDescriptor Major-K layouts. -* For example, -* auto smem_layout = tile_to_shape(Layout_K_SW128_Atom{}, Shape<_128,_64>{}); -* is guaranteed to be accepted by make_umma_desc for appropriate value_type. -*/ -template -CUTE_HOST_DEVICE constexpr -SmemDescriptor -make_umma_desc(Tensor const& tensor) -{ - static_assert(is_smem::value, "UMMA Descriptors can only be constructed on smem."); - static_assert(TLayout::rank == 2, "UMMA Descriptors can only be constructed on rank-2 tensors."); - using value_type = typename TEngine::value_type; - - Tensor u128_tensor = recast(tensor); - - // Result - SmemDescriptor desc; - desc.version_ = 1; // Set the version for blackwell - desc.lbo_mode_ = 0; // set to legacy mode by default - - // Layout type - constexpr UMMA::LayoutType LAYOUT_TYPE = UMMA::layout_type(u128_tensor); - desc.layout_type_ = uint8_t(LAYOUT_TYPE); - - // Start address (4LSB not included) - uint32_t start_address = cast_smem_ptr_to_uint(raw_pointer_cast(u128_tensor.data())); - desc.start_address_ = static_cast(start_address >> 4); - - constexpr uint8_t base_offset = 0; - desc.base_offset_ = base_offset; - - // LayoutType meta - constexpr int SwizzleAtomMNSize = LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_NONE ? 1 : - LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_32B ? 2 : - LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_64B ? 4 : - LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_128B ? 8 : - LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_128B_BASE32B ? 8 : -1; - - if constexpr (MajorMode == UMMA::Major::MN) - { - /* In units of uint128_t, each UmmaDescriptor Major-MN describes a canonical layout of the form - * - * LayoutType::INTERLEAVE : Swizzle<0,4,3> o smem_ptr o ((1,n),(8,k)):((X,SBO),(1,LBO)) - * LayoutType::B32 : Swizzle<1,4,3> o smem_ptr o ((2,n),(8,k)):((1,LBO),(2,SBO)) - * LayoutType::B64 : Swizzle<2,4,3> o smem_ptr o ((4,n),(8,k)):((1,LBO),(4,SBO)) - * LayoutType::B128 : Swizzle<3,4,3> o smem_ptr o ((8,n),(8,k)):((1,LBO),(8,SBO)) - * LayoutType::B128_BASE32B : Swizzle<2,5,2> o smem_ptr o ((8,n),(4,k)):((1,LBO),(4,SBO)) - */ - - constexpr int SwizzleAtomKSize = LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_128B_BASE32B ? 4 : 8; - - // Construct the canonical UMMA T Layout with shape ((SwizzleAtomMNSize,n),(SwizzleAtomKSize,2)) - Layout canonical_layout = logical_divide(layout(u128_tensor), Tile>,Layout>>{}); - - // Check profile of canonical - CUTE_STATIC_ASSERT_V(congruent(canonical_layout, Shape,Shape<_1,_1>>{}), "Not a canonical UMMA_MN Layout: Expected profile failure."); - // Check canonical mode strides - constexpr uint32_t stride_00 = stride<0,0>(canonical_layout); - constexpr uint32_t expected_stride_00 = LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_NONE ? stride<0,0>(canonical_layout) : 1; - static_assert(stride_00 == expected_stride_00, "Not a canonical UMMA_MN Layout: Expected stride failure."); - constexpr uint32_t stride_10 = stride<1,0>(canonical_layout); - constexpr uint32_t expected_stride_10 = SwizzleAtomMNSize; - static_assert(stride_10 == expected_stride_10, "Not a canonical UMMA_MN Layout: Expected stride failure."); - - // stride dimension byte offset and leading dimension byte offset (4LSB not included == uint128_t units) - constexpr uint32_t stride_01 = stride<0,1>(canonical_layout); - constexpr uint32_t stride_11 = stride<1,1>(canonical_layout); - - desc.stride_byte_offset_ = (LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_NONE) ? stride_01 : stride_11; - desc.leading_byte_offset_ = (LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_NONE) ? stride_11 : stride_01; - } else - if constexpr (MajorMode == UMMA::Major::K) - { - /* In units of uint128_t, each UmmaDescriptor Major-K describes a canonical layout of the form - * - * LayoutType::INTERLEAVE : Swizzle<0,4,3> o smem_ptr o ((8,n),2):((1,SBO),LBO) - * LayoutType::B32 : Swizzle<1,4,3> o smem_ptr o ((8,n),2):((2,SBO),1) - * LayoutType::B64 : Swizzle<2,4,3> o smem_ptr o ((8,n),2):((4,SBO),1) - * LayoutType::B128 : Swizzle<3,4,3> o smem_ptr o ((8,n),2):((8,SBO),1) - * LayoutType::B128_BASE32B : Not applicable for Major-K - */ - - static_assert(LAYOUT_TYPE != UMMA::LayoutType::SWIZZLE_128B_BASE32B, "SWIZZLE_128B_BASE32B is invalid for Major-K"); - CUTE_STATIC_ASSERT_V(size<0>(u128_tensor) % Int<8>{} == Int<0>{}, // N|M size - "Not a canonical UMMA_K Layout: Expected MN-size multiple of 8."); - - // Construct the canonical UMMA N Layout with shape ((8,n),(2,1)) - Layout canonical_layout = logical_divide(layout(u128_tensor), Tile,Layout<_2,_1>>{}); - - // Check profile of canonical - CUTE_STATIC_ASSERT_V(congruent(canonical_layout, Shape,Shape<_1,_1>>{}), "Not a canonical UMMA_K Layout: Expected profile failure."); - // Check canonical mode strides - constexpr uint32_t stride_00 = stride<0,0>(canonical_layout); - constexpr uint32_t expected_stride_00 = SwizzleAtomMNSize; - static_assert(stride_00 == expected_stride_00, "Not a canonical UMMA_K Layout: Expected stride failure."); - constexpr uint32_t stride_10 = stride<1,0>(canonical_layout); - constexpr uint32_t expected_stride_10 = (LAYOUT_TYPE == UMMA::LayoutType::SWIZZLE_NONE) ? stride<1,0>(canonical_layout) : 1; - static_assert(stride_10 == expected_stride_10, "Not a canonical UMMA_K Layout: Expected stride failure."); - - // stride dimension byte offset and leading dimension byte offset (4LSB not included == uint128_t units) - constexpr uint32_t stride_01 = stride<0,1>(canonical_layout); - - desc.stride_byte_offset_ = stride_01; - desc.leading_byte_offset_ = stride_10; - } else { - static_assert(MajorMode != UMMA::Major::MN && MajorMode != UMMA::Major::K, "Unrecognized MajorMode!"); - } - return desc; -} - -/////////////////////////////////////////////////////////////////////////////// -// Higher level UMMA Descriptor utilities -/////////////////////////////////////////////////////////////////////////////// - -struct DescriptorIterator -{ - using reference = SmemDescriptor; - using element_type = SmemDescriptor; - using value_type = SmemDescriptor; - - SmemDescriptor desc_; - - // Dereference returns the UmmaDescriptor - CUTE_HOST_DEVICE constexpr - reference operator*() const { return desc_; } - - // Advance and return a new UmmaDescriptor - template - CUTE_HOST_DEVICE constexpr - reference operator[](Index const& i) const { return *(*this + i); } - - // Return an advanced iterator - template - CUTE_HOST_DEVICE constexpr - DescriptorIterator operator+(Index const& offset) const - { -#if (__CUDACC_VER_MAJOR__ > 13) || ((__CUDACC_VER_MAJOR__ == 13) && (__CUDACC_VER_MINOR__ > 3)) - return { desc_ + uint64_t(offset) }; -#else - // Use 32-bit calculation rather than 64-bit calculation as we only update part of desc. - SmemDescriptor ret; - ret.lo = desc_.lo + uint32_t(offset); - ret.hi = desc_.hi; - return { ret }; -#endif - } -}; - -template -CUTE_HOST_DEVICE constexpr -SmemDescriptor -raw_pointer_cast(DescriptorIterator const& ptr) { - return ptr.desc_; -} - -CUTE_HOST_DEVICE void -print(DescriptorIterator const&) { - printf("UMMA::DescriptorIterator"); -} - -// Flag for smem descriptor allocation/creation -template -struct smem_desc : DescriptorIterator {}; - -template -struct sparse_smem_desc : DescriptorIterator {}; - } // end namespace UMMA -// Customization point for creating a UMMA::smem_desc Tensor -template -struct MakeTensor> -{ - template - CUTE_HOST_DEVICE constexpr auto - operator()(Tensor const& smem_tensor) - { - static_assert(is_smem::value, "Expected SMEM Tensor to construct a UMMA Desc Tensor"); - return make_tensor(UMMA::DescriptorIterator{UMMA::make_umma_desc(tensor<0>(smem_tensor))}, - replace<0>(recast(smem_tensor).layout(), Layout<_1,_0>{})); - } -}; - -// Customization point for creating a UMMA::sparse_smem_desc Tensor -template -struct MakeTensor> -{ - // Note that this is the exact same as UMMA::smem_desc above. - // Only the interface validates that we are passed a sparse_ptr, which is recast away to construct - // the smem desc tensor - template - CUTE_HOST_DEVICE constexpr auto - operator()(Tensor const& smem_tensor) - { - static_assert(is_smem::value, "Expected SMEM Tensor to construct a UMMA Desc Tensor"); - static_assert(is_sparse::value, "Expected sparse value_type."); - static_assert(is_sparse_ptr::value, "Expected sparse iter."); - return make_tensor(UMMA::DescriptorIterator{UMMA::make_umma_desc(tensor<0>(smem_tensor))}, - replace<0>(recast(smem_tensor).layout(), Layout<_1,_0>{})); - } -}; - // Special smem_desc_iter tensor entry for UTCCP copy. template constexpr auto get_utccp_smem_desc_tensor(Tensor const& smem_utccp_partitioned_tensor) { @@ -422,605 +147,6 @@ constexpr auto get_utccp_smem_desc_tensor(Tensor const& smem_u return make_tensor(core_matrix_desc_tensor.data(), recast_layout(smem_utccp_partitioned_tensor.layout())); } -namespace UMMA { - -// Import TMEM constants -namespace TMEM = cute::TMEM; - -enum class TmemAllocMode { - // Default allocation mode. - // If a TMEM Atom uses a half-subpartition (16DPs), then multiple atoms can be - // interleaved by using the top-half-subpartition and the bottom-half-subpartition. - // Full utilization of TMEM capacity. - Interleaved = 0, - // Prevents interleaving. - // If a TMEM Atom uses a half-subpartition (16DPs), then multiple atoms will not be - // interleaved. - // Required for DP-address equivalence in TMEM-A and TMEM-C allocations in UMMA_TS. - NonInterleaved = 1, - // Duplicates the TMEM allocation across subpartitions. - // E.g. UMMA_2SM_128xNx16_TS uses a "2x2 DP" TMEM Layout, but the TMEM allocation is - // actually doubled and the input data must be duplicated between the - // subpartitions [0,1]<->[2,3], i.e., each subpartition holds all columns - // of the A matrix needed for a single UMMA operation. - // For UMMA_2SM_128xNx16_TS, the distribution of the data is as follows. - // SM0: - // Subpart0 = A[0:32, 0:16], Subpart1 = A[32:64, 0:16], - // Subpart2 = A[A:32, 0:16], Subpart3 = A[32:64, 0:16] - // SM1: - // Subpart0 = A[64:96, 0:16], Subpart1 = A[96:128, 0:16], - // Subpart2 = A[64:96, 0:16], Subpart3 = A[96:128, 0:16] - Duplicated = 2, - // Duplicates the TMEM allocation across subpartitions for scale factor. - // Scale factor TMEM allocation for 4x1 data path - ScaleFactorDuplicated4by1 = 3, - // Scale factor TMEM allocation for 2x2 data path - ScaleFactorDuplicated2by2 = 4 -}; - -struct tmem_frg_base {}; - -// The UMMA Traits below have custom fragment type flags for their tmem tensors. -// These flags specialize a MakeTensor customization point to correctly make the fragment that is desired. -template -struct tmem_frg : tmem_frg_base -{ - static_assert(sizeof_bits_v <= sizeof_bits_v, "TMEM MMA allocations require StorageType big enough for ValueType."); - - // UMMA TMEM Allocator - // Each UMMA expects a specific MxN layout of TMEM for accumulators - // and sometimes a specific MxK layout of TMEM for A-values. - // @tparam ValueType The value type of the TMEM Tensor to allocate. - // @tparam StorageType The storage type of the TMEM Tensor to allocate. - // "Sparse" allocations often allocate ValueType=half_t within StorageType=uint32_t. - // "Dense" allocations often allocate ValueType=half_t within StorageType=half_t. - // @tparam N_SM The number of SMs in this UMMA_XSM instruction. - // @tparam TmemAlloc UMMA-specific allocation modifier for special cases. - // Some UMMA instructions expect strange atoms or tilings of atoms. - // @param tmem_shape ((M_MMA_SM,N_MMA_SM),MMA_M,MMA_N,...) - // The post-MMA-partitioned shape of TMEM to allocate. - // Note for UMMA_2SM_128xNx16, that M_MMA_SM will be 64, for example. - template - CUTE_HOST_DEVICE constexpr static auto - make(TmemShape const& tmem_shape) - { - CUTE_STATIC_ASSERT_V(size(tmem_shape)*Int)>{} <= TMEM::MAX_CAPACITY_BITS{}, - "Requesting more TMEM than is available."); - CUTE_STATIC_ASSERT_V(rank<0>(tmem_shape) == Int<2>{}, "Expected post-partitioned shape ((M_MMA,N_MMA),...)."); - constexpr int R = decltype(rank(tmem_shape))::value; - constexpr int M_MMA = decltype(size<0,0>(tmem_shape))::value; - constexpr int N_MMA = decltype(size<0,1>(tmem_shape))::value; - - // It's convenient to use "virtual tensor memory addressing" - // with DP_STRIDE=1, COL_STRIDE=128 to define the tmem_atom, - // then convert to "logical tensor memory addressing" on return. - using COL_ADDR = C::value / sizeof_bits::value>; - Layout tmem_restride = Layout, - Stride, COL_ADDR>>{}; - - static_assert(N_SM == 1 || N_SM == 2, "UMMA expects N_SM == 1 or N_SM == 2"); - if constexpr (N_SM == 1) - { - static_assert(TmemAlloc == UMMA::TmemAllocMode::Interleaved || TmemAlloc == UMMA::TmemAllocMode::NonInterleaved, - "UMMA_1SM only accepts Interleaved or NonInterleaved"); - static_assert(M_MMA == 64 || M_MMA == 128, "UMMA_1SM M-mode size should be 64 or 128."); - - if constexpr (M_MMA == 64) - { - // Half subpartitions layout atom: (M,N) -> tmem_addr - Layout tmem_atom = Layout, Int>, - Stride, _128>>{}; - // tile_stride = 2 causes the tiling to "skip" the first tile in DPs - constexpr int tile_stride = TmemAlloc == UMMA::TmemAllocMode::Interleaved ? 1 : 2; - // This will tile in DPs first, then COLs - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape), - compact_col_major(take<1,R>(tmem_shape),Int{}))); - // Restride for the DP/COL addressing and return - return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); - } else - if constexpr (M_MMA == 128) - { - // For M_MMA = 128, all datapaths are occupied. TmemAllocMode doesn't change the allocation. - // Full subpartitions layout atom: (M,N) -> tmem_addr - Layout tmem_atom = Layout>, - Stride< _1, _128>>{}; - // This will tile in DPs first, then COLs - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Restride for the DP/COL addressing and return - return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); - } - - } else - if constexpr (N_SM == 2) - { - static_assert(TmemAlloc == UMMA::TmemAllocMode::Interleaved || TmemAlloc == UMMA::TmemAllocMode::Duplicated, - "UMMA_2SM only accepts Interleaved or Duplicated"); - static_assert(M_MMA == 32 || M_MMA == 64 || M_MMA == 128, "UMMA_2SM M-mode size should be 32 or 64 or 128."); - - if constexpr (M_MMA == 32) - { - static_assert(TmemAlloc == UMMA::TmemAllocMode::Interleaved, "Only TmemAllocMode::Interleaved is supported for UMMA_2SM M_MMA=32"); - // The "1x4" layout atom: (M,N) -> tmem_addr - Layout tmem_atom = Layout, _4>>, - Stride< _1,Stride< _128,_32>>>{}; - // This will tile in DPs first, then COLs - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Restride for the DP/COL addressing and return - return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); - } else - if constexpr (M_MMA == 64 && TmemAlloc == UMMA::TmemAllocMode::Interleaved) - { - // The "2x2" layout atom: (M,N) -> tmem_addr - Layout tmem_atom = Layout, _2>>, - Stride< _1,Stride< _128,_64>>>{}; - // This will tile in DPs first, then COLs - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Restride for the DP/COL addressing and return - return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); - - } else - if constexpr (M_MMA == 64 && TmemAlloc == UMMA::TmemAllocMode::Duplicated) - { - // The "2x2" duplicated layout atom: (M,N) -> tmem_addr - Layout tmem_atom = Layout>, - Stride< _1, _128>>{}; - // This will tile in DPs first, then COLs - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Restride for the DP/COL addressing and return - return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); - } else - if constexpr (M_MMA == 128) - { - // For M_MMA = 128, all datapaths are occupied. TmemAllocMode doesn't change the allocation. - // The "4x1" layout atom: (M,N) -> tmem_addr - Layout tmem_atom = Layout>, - Stride< _1, _128>>{}; - // This will tile in DPs first, then COLs - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Restride for the DP/COL addressing and return - return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); - } - } - - CUTE_GCC_UNREACHABLE; - } -}; - -// Convenient aliases for common cases in the UMMA::ElementXFrg below -template -using tmem_frg_1sm = tmem_frg; -template -using tmem_frg_2sm = tmem_frg; - -// Make metadata TMEM fragments for sparse MMAs. -// Also note that the TMEM fragment addresses are assumed to be COL-4 aligned -- working with arch to remove this condition -template -struct tmem_e_frg : tmem_frg_base -{ - template - CUTE_HOST_DEVICE constexpr static auto - make(TmemShape const& tmem_shape) - { - CUTE_STATIC_ASSERT_V(rank<0>(tmem_shape) == Int<2>{}, "Expected post-partitioned shape ((M_MMA,N_MMA),...)."); - constexpr int R = decltype(rank(tmem_shape))::value; - constexpr int M_MMA = decltype(size<0,0>(tmem_shape))::value; - constexpr int N_MMA = decltype(size<0,1>(tmem_shape))::value; - - static_assert(M_MMA == 128, "Only 128 implemented right now."); - - // It's convenient to use "virtual tensor memory addressing" - // with DP_STRIDE=1, COL_STRIDE=128 to define the tmem_atom, - // then convert to "logical tensor memory addressing" on return. - [[maybe_unused]] Layout tmem_restride = Layout, - Stride>{}; - - if constexpr (sizeof_bits::value == 32) // TF32: 128x16 atom - { - static_assert(N_MMA == 16); - Layout tmem_atom = Layout, Shape < _8,_2>>, - Stride, Stride<_128,_8>>>{}; - // Tile to MMA tiling - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Address transformations with upcast<2> for 2-bit base types - Layout tmem_layout = composition(upcast<2>(tmem_restride), tmem_logical_layout); - // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing - return make_tensor(make_tmem_ptr>(), tmem_layout); - } else - if constexpr (sizeof_bits::value == 16) // FP16: 128x32 atom - { - static_assert(N_MMA == 32); - Layout tmem_atom = Layout, Shape < _16,_2>>, - Stride, Stride<_128,_8>>>{}; - // Tile to MMA tiling - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Address transformations - Layout tmem_layout = composition(tmem_restride, tmem_logical_layout); - // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing - return make_tensor(make_tmem_ptr>(), tmem_layout); - } else - if constexpr (sizeof_bits::value == 8) // S8|Mix.F4/F6/F8: 128x64 atom - { - // For Mix 8bit f4/f6/f8, will pass in ValueType = uint8_t - static_assert(N_MMA == 64); - Layout tmem_atom = Layout, - Stride< _1,_128>>{}; - // Tile to MMA tiling - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Address transformations - Layout tmem_layout = composition(tmem_restride, tmem_logical_layout); - // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing - return make_tensor(make_tmem_ptr>(), tmem_layout); - } - if constexpr (sizeof_bits::value == 4) // F4: 128x128 atom - { - // For F4, will pass in ValueType = fp4 - Layout tmem_restride1 = Layout>, - Stride, _1>>{}; - // F4 has roughly same TMEM layout as Mix8bit.F4/F6/F8, the only difference is that K is multiplied by two - static_assert(N_MMA == 128); - Layout tmem_atom = Layout, - Stride< _1, _128>>{}; - // Tile to MMA tiling - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Address transformations - Layout tmem_layout = composition(tmem_restride1, tmem_logical_layout); - // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing - return make_tensor(make_tmem_ptr>(), tmem_layout); - } - - CUTE_GCC_UNREACHABLE; - } -}; - -template -struct tmem_e_frg_ws : tmem_frg_base -{ - template - CUTE_HOST_DEVICE constexpr static auto - make(TmemShape const& tmem_shape) - { - CUTE_STATIC_ASSERT_V(rank<0>(tmem_shape) == Int<2>{}, "Expected post-partitioned shape ((M_MMA,N_MMA),...)."); - constexpr int R = decltype(rank(tmem_shape))::value; - constexpr int M_MMA = decltype(size<0,0>(tmem_shape))::value; - constexpr int N_MMA = decltype(size<0,1>(tmem_shape))::value; - - static_assert(M_MMA == 128 || M_MMA == 64 || M_MMA == 32, "Weight stationary UMMA_1SM M-mode size should be 32 or 64 or 128."); - - // It's convenient to use "virtual tensor memory addressing" - // with DP_STRIDE=1, COL_STRIDE=128 to define the tmem_atom, - // then convert to "logical tensor memory addressing" on return. - Layout tmem_restride = Layout, - Stride>{}; - - if constexpr (sizeof_bits::value == 32) // TF32 - { - // MMA_M x MMA_K: 128x16 atom / 64x16 atom / 32x16 atom - static_assert(N_MMA == 16); - if constexpr (M_MMA == 128) { - Layout tmem_atom = Layout, Shape < _8,_2>>, - Stride, Stride<_128,_8>>>{}; - // Tile to MMA tiling - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Address transformations with upcast<2> for 2-bit base types - Layout tmem_layout = composition(upcast<2>(tmem_restride), tmem_logical_layout); - // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing - return make_tensor(make_tmem_ptr>(), tmem_layout); - } - else if constexpr (M_MMA == 64) { - Layout tmem_atom = Layout, Shape < _8,_2>, _2>, - Stride, Stride<_128,_8>,_64>>{}; - // Tile to MMA tiling - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Address transformations with upcast<2> for 2-bit base types - Layout tmem_layout = composition(upcast<2>(tmem_restride), tmem_logical_layout); - // Sparsity wrap, no sparse_ptr because tmem_ptr handles its own subword addressing - return make_tensor(make_tmem_ptr>(), tmem_layout); - } - else if constexpr (M_MMA == 32) { - Layout tmem_atom = Layout, Shape < _8,_2>, _4>, - Stride, Stride<_128,_8>,_32>>{}; - // Tile to MMA tiling - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Address transformations with upcast<2> for 2-bit base types - Layout tmem_layout = composition(upcast<2>(tmem_restride), tmem_logical_layout); - // Sparsity wrap, no sparse_ptr because tmem_ptr handles its own subword addressing - return make_tensor(make_tmem_ptr>(), tmem_layout); - } - else { - static_assert(dependent_false, "Invalid M_MMA value"); - } - } - else if constexpr (sizeof_bits::value == 16) // FP16 - { - // MMA_M x MMA_K: 128x32 atom / 64x32 atom / 32x32 atom - static_assert(N_MMA == 32); - if constexpr (M_MMA == 128) { - Layout tmem_atom = Layout, Shape < _16,_2>>, - Stride, Stride<_128,_8>>>{}; - // Tile to MMA tiling - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Address transformations - Layout tmem_layout = composition(tmem_restride, tmem_logical_layout); - // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing - return make_tensor(make_tmem_ptr>(), tmem_layout); - } - else if constexpr (M_MMA == 64) { - Layout tmem_atom = Layout, Shape < _16,_2>, _2>, - Stride, Stride<_128,_8>,_64>>{}; - // Tile to MMA tiling - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Address transformations - Layout tmem_layout = composition(tmem_restride, tmem_logical_layout); - // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing - return make_tensor(make_tmem_ptr>(), tmem_layout); - } - else if constexpr (M_MMA == 32) { - Layout tmem_atom = Layout, Shape < _16,_2>, _4>, - Stride, Stride<_128,_8>,_32>>{}; - // Tile to MMA tiling - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Address transformations - Layout tmem_layout = composition(tmem_restride, tmem_logical_layout); - // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing - return make_tensor(make_tmem_ptr>(), tmem_layout); - } - else { - static_assert(dependent_false, "Invalid M_MMA value"); - } - } - else if constexpr (sizeof_bits::value == 8) // I8|F8 - { - // MMA_M x MMA_K: 128x64 atom / 64x64 atom / 32x64 atom - static_assert(N_MMA == 64); - if constexpr (M_MMA == 128) { - Layout tmem_atom = Layout, - Stride< _1,_128>>{}; - // Tile to MMA tiling - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Address transformations - Layout tmem_layout = composition(tmem_restride, tmem_logical_layout); - // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing - return make_tensor(make_tmem_ptr>(), tmem_layout); - } - else if constexpr (M_MMA == 64) { - Layout tmem_atom = Layout>, - Stride< _1, Stride<_128, _64>>>{}; - // Tile to MMA tiling - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Address transformations - Layout tmem_layout = composition(tmem_restride, tmem_logical_layout); - // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing - return make_tensor(make_tmem_ptr>(), tmem_layout); - } - else if constexpr (M_MMA == 32) { - Layout tmem_atom = Layout>, - Stride< _1, Stride<_128, _32>>>{}; - // Tile to MMA tiling - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Address transformations - Layout tmem_layout = composition(tmem_restride, tmem_logical_layout); - // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing - return make_tensor(make_tmem_ptr>(), tmem_layout); - } - else { - static_assert(dependent_false, "Invalid M_MMA value"); - } - } - else { - static_assert(dependent_false, "Invalid ValueType"); - } - - CUTE_GCC_UNREACHABLE; - } -}; - -template -struct tmem_sf_frg: tmem_frg_base -{ - // UMMA TMEM Allocator for Scale Factor A for Mxf4Nvf4 and Mxf8f6f4 instructions - // We expect a tensor that has the same layout as A matrix - // @tparam ValueType: data type of scaling factor - // Note that the StorageType is the same as ValueType, i.e., we always use a compact allocation - // @tparam SFVecSize: The number of values that is scaled by a single scaling factor. - // Valid values are (16, 32) - // @tparam N_SM: Number of SMs in UMMA instruction - // @param tmem_shape: An MMA partitioned shape where first mode encodes, A layout of the MMA instruction. - // Note that the shape doesn't match the actual allocation. size<0,1>(tmem_shape) will give us the number of - // elements in K-mode of MMA rather than the number of scaling factors. - template - CUTE_HOST_DEVICE constexpr static auto - make(TmemShape const& tmem_shape) - { - CUTE_STATIC_ASSERT_V(rank<0>(tmem_shape) == Int<2>{}, "Expected post-partitioned shape ((M_MMA,N_MMA),...)."); - constexpr int MMA_MN = decltype(size<0,0>(tmem_shape))::value; - constexpr int MMA_VS = decltype(size<0,1,0>(tmem_shape))::value; - constexpr int MMA_NSF = decltype(size<0,1,1>(tmem_shape))::value; - constexpr int R_MMA_K = decltype(rank(get<0,1>(tmem_shape)))::value; - constexpr int R = decltype(rank(tmem_shape))::value; - - // We expect an MMA-SF partitioned tensor - // ((MMA_MN, (VecSize, NSF)), num_MMA_MN, num_MMA_K, ...) - // where VecSize*NSF = MMA_K - static_assert(R >= 3, "Expected an MMA partitioned tensor"); // ((MMA), num_MMA_MN, num_MMA_K, ...) - static_assert(R_MMA_K == 2, "Expected an MMA-SF partitioned tensor"); // (VecSize, NSF) - using REP = _4; // Replication factor. Data is always replicated across subpartitions - constexpr int SUBPART_DPs = 32; // Number of DPs in a subpartition - - using COL_ADDR = C::value / sizeof_bits::value>; - Layout tmem_restride = Layout, - Stride, COL_ADDR>>{}; - - if constexpr (Is_SFA || (!Is_SFA && TmemAlloc == UMMA::TmemAllocMode::ScaleFactorDuplicated4by1)) { - // SFA, 2x2 and 4x1 data path - // SFB, 4x1 data path - auto tmem_atom = Layout < Shape< Shape< Shape, Int>, REP>, Shape, Int>>, - Stride, _32>, Stride< _0, _128>>>{}; - - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - auto final_tmem_layout = composition(tmem_restride, tmem_logical_layout); - return make_tensor(make_tmem_ptr(), final_tmem_layout); - } - else { - // SFB, 2x2 datapath - static_assert(!Is_SFA and TmemAlloc == UMMA::TmemAllocMode::ScaleFactorDuplicated2by2); - static_assert(N_SM == 2, "Should be 2x2 Datapath"); - // 2x2 Datapth - auto tmem_atom = Layout < Shape< Shape< Shape, Int>, _2, _2>, Shape, Int>>, - Stride, _64, _32>, Stride< _0, _128>>>{}; - - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - auto final_tmem_layout = composition(tmem_restride, tmem_logical_layout); - return make_tensor(make_tmem_ptr(), final_tmem_layout); - } - } -}; - -// Make C/D Tmem fragment for weight-stationary MMAs -template -struct tmem_frg_ws : tmem_frg_base -{ - static_assert(sizeof_bits_v <= sizeof_bits_v, "TMEM MMA allocations require StorageType big enough for ValueType."); - - // UMMA TMEM Allocator - // Each UMMA expects a specific MxN layout of TMEM for accumulators - // and sometimes a specific MxK layout of TMEM for A-values. - // @tparam ValueType The value type of the TMEM Tensor to allocate. - // @tparam StorageType The storage type of the TMEM Tensor to allocate. - // "Sparse" allocations often allocate ValueType=half_t within StorageType=uint32_t. - // "Dense" allocations often allocate ValueType=half_t within StorageType=half_t. - // @tparam N_SM The number of SMs in this UMMA_XSM instruction. - // @tparam TmemAlloc UMMA-specific allocation modifier for special cases. - // Some UMMA instructions expect strange atoms or tilings of atoms. - // @param tmem_shape ((M_MMA_SM,N_MMA_SM),MMA_M,MMA_N,...) - // The post-MMA-partitioned shape of TMEM to allocate. - // Note for UMMA_2SM_128xNx16, that M_MMA_SM will be 64, for example. - template - CUTE_HOST_DEVICE constexpr static auto - make(TmemShape const& tmem_shape) - { - CUTE_STATIC_ASSERT_V(size(tmem_shape)*Int)>{} <= TMEM::MAX_CAPACITY_BITS{}, - "Requesting more TMEM than is available."); - CUTE_STATIC_ASSERT_V(rank<0>(tmem_shape) == Int<2>{}, "Expected post-partitioned shape ((M_MMA,N_MMA),...)."); - constexpr int R = decltype(rank(tmem_shape))::value; - constexpr int M_MMA = decltype(size<0,0>(tmem_shape))::value; - constexpr int N_MMA = decltype(size<0,1>(tmem_shape))::value; - - // It's convenient to use "virtual tensor memory addressing" - // with DP_STRIDE=1, COL_STRIDE=128 to define the tmem_atom, - // then convert to "logical tensor memory addressing" on return. - using COL_ADDR = C::value / sizeof_bits::value>; - Layout tmem_restride = Layout, - Stride, COL_ADDR>>{}; - - static_assert(N_SM == 1, "UMMA.WS expects N_SM == 1"); - - static_assert(M_MMA == 32 || M_MMA == 64 || M_MMA == 128, - "Weight stationary UMMA_1SM M-mode size should be 32 or 64 or 128."); - static_assert(N_MMA == 64 || N_MMA == 128 || N_MMA == 256, - "Dense weight stationary UMMA_1SM N-mode size should be 64 or 128 or 256."); - // Weight Stationary MMA config - if constexpr (M_MMA == 32) - { - // 1x4 datapath - Layout tmem_atom = Layout, _4>>, - Stride< _1, Stride< _128,_32>> - >{}; - constexpr int tile_stride = 1; - // This will tile in DPs first, then COLs - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape), - compact_col_major(take<1,R>(tmem_shape), Int{}))); - // Restride for the DP/COL addressing and return - return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); - } else - if constexpr (M_MMA == 64) - { - // 2x2 datapath - Layout tmem_atom = Layout, _2>>, - Stride< _1, Stride< _128,_64>> - >{}; - constexpr int tile_stride = 1; - // This will tile in DPs first, then COLs - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape), - compact_col_major(take<1,R>(tmem_shape), Int{}))); - // Restride for the DP/COL addressing and return - return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); - } else - if constexpr (M_MMA == 128) - { - // For M_MMA = 128, all datapaths are occupied. TmemAllocMode doesn't change the allocation. - // Full subpartitions layout atom: (M,N) -> tmem_addr - Layout tmem_atom = Layout>, - Stride< _1, _128>>{}; - // This will tile in DPs first, then COLs - Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); - // Restride for the DP/COL addressing and return - return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); - } - - CUTE_GCC_UNREACHABLE; - } -}; - -// Convenient aliases for common cases in the UMMA::ElementXFrg below -template -using tmem_frg_ws_1sm = tmem_frg_ws; - -} // end namespace UMMA - -// Customization point for creating a UMMA::tmem_frg Tensor -template -struct MakeTensor> -{ - template - CUTE_HOST_DEVICE constexpr auto - operator()(Shape const& tmem_shape) { - return UMMA::tmem_frg::make(shape(tmem_shape)); - } -}; - -template -struct MakeTensor> -{ - template - CUTE_HOST_DEVICE constexpr auto - operator()(Shape const& tmem_shape) { - return UMMA::tmem_frg_ws::make(shape(tmem_shape)); - } -}; - - -// Customization point for creating a UMMA::tmem_frg Tensor -template -struct MakeTensor> -{ - template - CUTE_HOST_DEVICE constexpr auto - operator()(Shape const& tmem_shape) { - return UMMA::tmem_e_frg::make(shape(tmem_shape)); - } -}; - -template -struct MakeTensor> -{ - template - CUTE_HOST_DEVICE constexpr auto - operator()(Shape const& tmem_shape) { - return UMMA::tmem_e_frg_ws::make(shape(tmem_shape)); - } -}; - -template -struct MakeTensor> -{ - template - CUTE_HOST_DEVICE constexpr auto - operator()(Shape const& tmem_shape) { - return UMMA::tmem_sf_frg::make(shape(tmem_shape)); - } -}; - /////////////////////////////////////////////////////////////////////////////// //////////////////////////// MMA_TRAITS /////////////////////////////////////// /////////////////////////////////////////////////////////////////////////////// diff --git a/include/cute/atom/mma_traits_sm100_frag.hpp b/include/cute/atom/mma_traits_sm100_frag.hpp new file mode 100644 index 000000000..0388d47d6 --- /dev/null +++ b/include/cute/atom/mma_traits_sm100_frag.hpp @@ -0,0 +1,640 @@ +/*************************************************************************************************** + * Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#include +#include +#include // UMMA::Major, SmemDescriptor, DescriptorIterator +#include // cute::TMEM::, make_tmem_ptr +#include // sizeof_bits_v + +namespace cute { + +namespace UMMA { + +// Import TMEM constants +namespace TMEM = cute::TMEM; + +enum class TmemAllocMode { + // Default allocation mode. + // If a TMEM Atom uses a half-subpartition (16DPs), then multiple atoms can be + // interleaved by using the top-half-subpartition and the bottom-half-subpartition. + // Full utilization of TMEM capacity. + Interleaved = 0, + // Prevents interleaving. + // If a TMEM Atom uses a half-subpartition (16DPs), then multiple atoms will not be + // interleaved. + // Required for DP-address equivalence in TMEM-A and TMEM-C allocations in UMMA_TS. + NonInterleaved = 1, + // Duplicates the TMEM allocation across subpartitions. + // E.g. UMMA_2SM_128xNx16_TS uses a "2x2 DP" TMEM Layout, but the TMEM allocation is + // actually doubled and the input data must be duplicated between the + // subpartitions [0,1]<->[2,3], i.e., each subpartition holds all columns + // of the A matrix needed for a single UMMA operation. + // For UMMA_2SM_128xNx16_TS, the distribution of the data is as follows. + // SM0: + // Subpart0 = A[0:32, 0:16], Subpart1 = A[32:64, 0:16], + // Subpart2 = A[0:32, 0:16], Subpart3 = A[32:64, 0:16] + // SM1: + // Subpart0 = A[64:96, 0:16], Subpart1 = A[96:128, 0:16], + // Subpart2 = A[64:96, 0:16], Subpart3 = A[96:128, 0:16] + Duplicated = 2, + // Duplicates the TMEM allocation across subpartitions for scale factor. + // Scale factor TMEM allocation for 4x1 data path + ScaleFactorDuplicated4by1 = 3, + // Scale factor TMEM allocation for 2x2 data path + ScaleFactorDuplicated2by2 = 4 +}; + +struct tmem_frg_base {}; + +// The UMMA Traits below have custom fragment type flags for their tmem tensors. +// These flags specialize a MakeTensor customization point to correctly make the fragment that is desired. +template +struct tmem_frg : tmem_frg_base +{ + static_assert(sizeof_bits_v <= sizeof_bits_v, "TMEM MMA allocations require StorageType big enough for ValueType."); + + // UMMA TMEM Allocator + // Each UMMA expects a specific MxN layout of TMEM for accumulators + // and sometimes a specific MxK layout of TMEM for A-values. + // @tparam ValueType The value type of the TMEM Tensor to allocate. + // @tparam StorageType The storage type of the TMEM Tensor to allocate. + // "Sparse" allocations often allocate ValueType=half_t within StorageType=uint32_t. + // "Dense" allocations often allocate ValueType=half_t within StorageType=half_t. + // @tparam N_SM The number of SMs in this UMMA_XSM instruction. + // @tparam TmemAlloc UMMA-specific allocation modifier for special cases. + // Some UMMA instructions expect strange atoms or tilings of atoms. + // @param tmem_shape ((M_MMA_SM,N_MMA_SM),MMA_M,MMA_N,...) + // The post-MMA-partitioned shape of TMEM to allocate. + // Note for UMMA_2SM_128xNx16, that M_MMA_SM will be 64, for example. + template + CUTE_HOST_DEVICE constexpr static auto + make(TmemShape const& tmem_shape) + { + CUTE_STATIC_ASSERT_V(size(tmem_shape)*Int)>{} <= TMEM::MAX_CAPACITY_BITS{}, + "Requesting more TMEM than is available."); + CUTE_STATIC_ASSERT_V(rank<0>(tmem_shape) == Int<2>{}, "Expected post-partitioned shape ((M_MMA,N_MMA),...)."); + constexpr int R = decltype(rank(tmem_shape))::value; + constexpr int M_MMA = decltype(size<0,0>(tmem_shape))::value; + constexpr int N_MMA = decltype(size<0,1>(tmem_shape))::value; + + // It's convenient to use "virtual tensor memory addressing" + // with DP_STRIDE=1, COL_STRIDE=128 to define the tmem_atom, + // then convert to "logical tensor memory addressing" on return. + using COL_ADDR = C::value / sizeof_bits::value>; + Layout tmem_restride = Layout, + Stride, COL_ADDR>>{}; + + static_assert(N_SM == 1 || N_SM == 2, "UMMA expects N_SM == 1 or N_SM == 2"); + if constexpr (N_SM == 1) + { + static_assert(TmemAlloc == UMMA::TmemAllocMode::Interleaved || TmemAlloc == UMMA::TmemAllocMode::NonInterleaved, + "UMMA_1SM only accepts Interleaved or NonInterleaved"); + static_assert(M_MMA == 64 || M_MMA == 128, "UMMA_1SM M-mode size should be 64 or 128."); + + if constexpr (M_MMA == 64) + { + // Half subpartitions layout atom: (M,N) -> tmem_addr + Layout tmem_atom = Layout, Int>, + Stride, _128>>{}; + // tile_stride = 2 causes the tiling to "skip" the first tile in DPs + constexpr int tile_stride = TmemAlloc == UMMA::TmemAllocMode::Interleaved ? 1 : 2; + // This will tile in DPs first, then COLs + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape), + compact_col_major(take<1,R>(tmem_shape),Int{}))); + // Restride for the DP/COL addressing and return + return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); + } else + if constexpr (M_MMA == 128) + { + // For M_MMA = 128, all datapaths are occupied. TmemAllocMode doesn't change the allocation. + // Full subpartitions layout atom: (M,N) -> tmem_addr + Layout tmem_atom = Layout>, + Stride< _1, _128>>{}; + // This will tile in DPs first, then COLs + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Restride for the DP/COL addressing and return + return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); + } + + } else + if constexpr (N_SM == 2) + { + static_assert(TmemAlloc == UMMA::TmemAllocMode::Interleaved || TmemAlloc == UMMA::TmemAllocMode::Duplicated, + "UMMA_2SM only accepts Interleaved or Duplicated"); + static_assert(M_MMA == 32 || M_MMA == 64 || M_MMA == 128, "UMMA_2SM M-mode size should be 32 or 64 or 128."); + + if constexpr (M_MMA == 32) + { + static_assert(TmemAlloc == UMMA::TmemAllocMode::Interleaved, "Only TmemAllocMode::Interleaved is supported for UMMA_2SM M_MMA=32"); + // The "1x4" layout atom: (M,N) -> tmem_addr + Layout tmem_atom = Layout, _4>>, + Stride< _1,Stride< _128,_32>>>{}; + // This will tile in DPs first, then COLs + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Restride for the DP/COL addressing and return + return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); + } else + if constexpr (M_MMA == 64 && TmemAlloc == UMMA::TmemAllocMode::Interleaved) + { + // The "2x2" layout atom: (M,N) -> tmem_addr + Layout tmem_atom = Layout, _2>>, + Stride< _1,Stride< _128,_64>>>{}; + // This will tile in DPs first, then COLs + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Restride for the DP/COL addressing and return + return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); + + } else + if constexpr (M_MMA == 64 && TmemAlloc == UMMA::TmemAllocMode::Duplicated) + { + // The "2x2" duplicated layout atom: (M,N) -> tmem_addr + Layout tmem_atom = Layout>, + Stride< _1, _128>>{}; + // This will tile in DPs first, then COLs + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Restride for the DP/COL addressing and return + return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); + } else + if constexpr (M_MMA == 128) + { + // For M_MMA = 128, all datapaths are occupied. TmemAllocMode doesn't change the allocation. + // The "4x1" layout atom: (M,N) -> tmem_addr + Layout tmem_atom = Layout>, + Stride< _1, _128>>{}; + // This will tile in DPs first, then COLs + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Restride for the DP/COL addressing and return + return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); + } + } + + CUTE_GCC_UNREACHABLE; + } +}; + +// Convenient aliases for common cases in the UMMA::ElementXFrg below +template +using tmem_frg_1sm = tmem_frg; +template +using tmem_frg_2sm = tmem_frg; + +// Make metadata TMEM fragments for sparse MMAs. +// Also note that the TMEM fragment addresses are assumed to be COL-4 aligned -- working with arch to remove this condition +template +struct tmem_e_frg : tmem_frg_base +{ + template + CUTE_HOST_DEVICE constexpr static auto + make(TmemShape const& tmem_shape) + { + CUTE_STATIC_ASSERT_V(rank<0>(tmem_shape) == Int<2>{}, "Expected post-partitioned shape ((M_MMA,N_MMA),...)."); + constexpr int R = decltype(rank(tmem_shape))::value; + constexpr int M_MMA = decltype(size<0,0>(tmem_shape))::value; + constexpr int N_MMA = decltype(size<0,1>(tmem_shape))::value; + + static_assert(M_MMA == 128, "Only 128 implemented right now."); + + // It's convenient to use "virtual tensor memory addressing" + // with DP_STRIDE=1, COL_STRIDE=128 to define the tmem_atom, + // then convert to "logical tensor memory addressing" on return. + [[maybe_unused]] Layout tmem_restride = Layout, + Stride>{}; + + if constexpr (sizeof_bits::value == 32) // TF32: 128x16 atom + { + static_assert(N_MMA == 16); + Layout tmem_atom = Layout, Shape < _8,_2>>, + Stride, Stride<_128,_8>>>{}; + // Tile to MMA tiling + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Address transformations with upcast<2> for 2-bit base types + Layout tmem_layout = composition(upcast<2>(tmem_restride), tmem_logical_layout); + // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing + return make_tensor(make_tmem_ptr>(), tmem_layout); + } else + if constexpr (sizeof_bits::value == 16) // FP16: 128x32 atom + { + static_assert(N_MMA == 32); + Layout tmem_atom = Layout, Shape < _16,_2>>, + Stride, Stride<_128,_8>>>{}; + // Tile to MMA tiling + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Address transformations + Layout tmem_layout = composition(tmem_restride, tmem_logical_layout); + // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing + return make_tensor(make_tmem_ptr>(), tmem_layout); + } else + if constexpr (sizeof_bits::value == 8) // S8|Mix.F4/F6/F8: 128x64 atom + { + // For Mix 8bit f4/f6/f8, will pass in ValueType = uint8_t + static_assert(N_MMA == 64); + Layout tmem_atom = Layout, + Stride< _1,_128>>{}; + // Tile to MMA tiling + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Address transformations + Layout tmem_layout = composition(tmem_restride, tmem_logical_layout); + // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing + return make_tensor(make_tmem_ptr>(), tmem_layout); + } + if constexpr (sizeof_bits::value == 4) // F4: 128x128 atom + { + // For F4, will pass in ValueType = fp4 + Layout tmem_restride1 = Layout>, + Stride, _1>>{}; + // F4 has roughly same TMEM layout as Mix8bit.F4/F6/F8, the only difference is that K is multiplied by two + static_assert(N_MMA == 128); + Layout tmem_atom = Layout, + Stride< _1, _128>>{}; + // Tile to MMA tiling + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Address transformations + Layout tmem_layout = composition(tmem_restride1, tmem_logical_layout); + // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing + return make_tensor(make_tmem_ptr>(), tmem_layout); + } + + CUTE_GCC_UNREACHABLE; + } +}; + +template +struct tmem_e_frg_ws : tmem_frg_base +{ + template + CUTE_HOST_DEVICE constexpr static auto + make(TmemShape const& tmem_shape) + { + CUTE_STATIC_ASSERT_V(rank<0>(tmem_shape) == Int<2>{}, "Expected post-partitioned shape ((M_MMA,N_MMA),...)."); + constexpr int R = decltype(rank(tmem_shape))::value; + constexpr int M_MMA = decltype(size<0,0>(tmem_shape))::value; + constexpr int N_MMA = decltype(size<0,1>(tmem_shape))::value; + + static_assert(M_MMA == 128 || M_MMA == 64 || M_MMA == 32, "Weight stationary UMMA_1SM M-mode size should be 32 or 64 or 128."); + + // It's convenient to use "virtual tensor memory addressing" + // with DP_STRIDE=1, COL_STRIDE=128 to define the tmem_atom, + // then convert to "logical tensor memory addressing" on return. + Layout tmem_restride = Layout, + Stride>{}; + + if constexpr (sizeof_bits::value == 32) // TF32 + { + // MMA_M x MMA_K: 128x16 atom / 64x16 atom / 32x16 atom + static_assert(N_MMA == 16); + if constexpr (M_MMA == 128) { + Layout tmem_atom = Layout, Shape < _8,_2>>, + Stride, Stride<_128,_8>>>{}; + // Tile to MMA tiling + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Address transformations with upcast<2> for 2-bit base types + Layout tmem_layout = composition(upcast<2>(tmem_restride), tmem_logical_layout); + // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing + return make_tensor(make_tmem_ptr>(), tmem_layout); + } + else if constexpr (M_MMA == 64) { + Layout tmem_atom = Layout, Shape < _8,_2>, _2>, + Stride, Stride<_128,_8>,_64>>{}; + // Tile to MMA tiling + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Address transformations with upcast<2> for 2-bit base types + Layout tmem_layout = composition(upcast<2>(tmem_restride), tmem_logical_layout); + // Sparsity wrap, no sparse_ptr because tmem_ptr handles its own subword addressing + return make_tensor(make_tmem_ptr>(), tmem_layout); + } + else if constexpr (M_MMA == 32) { + Layout tmem_atom = Layout, Shape < _8,_2>, _4>, + Stride, Stride<_128,_8>,_32>>{}; + // Tile to MMA tiling + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Address transformations with upcast<2> for 2-bit base types + Layout tmem_layout = composition(upcast<2>(tmem_restride), tmem_logical_layout); + // Sparsity wrap, no sparse_ptr because tmem_ptr handles its own subword addressing + return make_tensor(make_tmem_ptr>(), tmem_layout); + } + else { + static_assert(dependent_false, "Invalid M_MMA value"); + } + } + else if constexpr (sizeof_bits::value == 16) // FP16 + { + // MMA_M x MMA_K: 128x32 atom / 64x32 atom / 32x32 atom + static_assert(N_MMA == 32); + if constexpr (M_MMA == 128) { + Layout tmem_atom = Layout, Shape < _16,_2>>, + Stride, Stride<_128,_8>>>{}; + // Tile to MMA tiling + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Address transformations + Layout tmem_layout = composition(tmem_restride, tmem_logical_layout); + // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing + return make_tensor(make_tmem_ptr>(), tmem_layout); + } + else if constexpr (M_MMA == 64) { + Layout tmem_atom = Layout, Shape < _16,_2>, _2>, + Stride, Stride<_128,_8>,_64>>{}; + // Tile to MMA tiling + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Address transformations + Layout tmem_layout = composition(tmem_restride, tmem_logical_layout); + // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing + return make_tensor(make_tmem_ptr>(), tmem_layout); + } + else if constexpr (M_MMA == 32) { + Layout tmem_atom = Layout, Shape < _16,_2>, _4>, + Stride, Stride<_128,_8>,_32>>{}; + // Tile to MMA tiling + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Address transformations + Layout tmem_layout = composition(tmem_restride, tmem_logical_layout); + // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing + return make_tensor(make_tmem_ptr>(), tmem_layout); + } + else { + static_assert(dependent_false, "Invalid M_MMA value"); + } + } + else if constexpr (sizeof_bits::value == 8) // I8|F8 + { + // MMA_M x MMA_K: 128x64 atom / 64x64 atom / 32x64 atom + static_assert(N_MMA == 64); + if constexpr (M_MMA == 128) { + Layout tmem_atom = Layout, + Stride< _1,_128>>{}; + // Tile to MMA tiling + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Address transformations + Layout tmem_layout = composition(tmem_restride, tmem_logical_layout); + // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing + return make_tensor(make_tmem_ptr>(), tmem_layout); + } + else if constexpr (M_MMA == 64) { + Layout tmem_atom = Layout>, + Stride< _1, Stride<_128, _64>>>{}; + // Tile to MMA tiling + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Address transformations + Layout tmem_layout = composition(tmem_restride, tmem_logical_layout); + // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing + return make_tensor(make_tmem_ptr>(), tmem_layout); + } + else if constexpr (M_MMA == 32) { + Layout tmem_atom = Layout>, + Stride< _1, Stride<_128, _32>>>{}; + // Tile to MMA tiling + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Address transformations + Layout tmem_layout = composition(tmem_restride, tmem_logical_layout); + // Sparsity wrap, no sparse_ptr because tmem_ptr handles it's own subword addressing + return make_tensor(make_tmem_ptr>(), tmem_layout); + } + else { + static_assert(dependent_false, "Invalid M_MMA value"); + } + } + else { + static_assert(dependent_false, "Invalid ValueType"); + } + + CUTE_GCC_UNREACHABLE; + } +}; + +template +struct tmem_sf_frg: tmem_frg_base +{ + // UMMA TMEM Allocator for Scale Factor A for Mxf4Nvf4 and Mxf8f6f4 instructions + // We expect a tensor that has the same layout as A matrix + // @tparam ValueType: data type of scaling factor + // Note that the StorageType is the same as ValueType, i.e., we always use a compact allocation + // @tparam SFVecSize: The number of values that is scaled by a single scaling factor. + // Valid values are (16, 32) + // @tparam N_SM: Number of SMs in UMMA instruction + // @param tmem_shape: An MMA partitioned shape where first mode encodes, A layout of the MMA instruction. + // Note that the shape doesn't match the actual allocation. size<0,1>(tmem_shape) will give us the number of + // elements in K-mode of MMA rather than the number of scaling factors. + template + CUTE_HOST_DEVICE constexpr static auto + make(TmemShape const& tmem_shape) + { + CUTE_STATIC_ASSERT_V(rank<0>(tmem_shape) == Int<2>{}, "Expected post-partitioned shape ((M_MMA,N_MMA),...)."); + constexpr int MMA_MN = decltype(size<0,0>(tmem_shape))::value; + constexpr int MMA_VS = decltype(size<0,1,0>(tmem_shape))::value; + constexpr int MMA_NSF = decltype(size<0,1,1>(tmem_shape))::value; + constexpr int R_MMA_K = decltype(rank(get<0,1>(tmem_shape)))::value; + constexpr int R = decltype(rank(tmem_shape))::value; + + // We expect an MMA-SF partitioned tensor + // ((MMA_MN, (VecSize, NSF)), num_MMA_MN, num_MMA_K, ...) + // where VecSize*NSF = MMA_K + static_assert(R >= 3, "Expected an MMA partitioned tensor"); // ((MMA), num_MMA_MN, num_MMA_K, ...) + static_assert(R_MMA_K == 2, "Expected an MMA-SF partitioned tensor"); // (VecSize, NSF) + using REP = _4; // Replication factor. Data is always replicated across subpartitions + constexpr int SUBPART_DPs = 32; // Number of DPs in a subpartition + + using COL_ADDR = C::value / sizeof_bits::value>; + Layout tmem_restride = Layout, + Stride, COL_ADDR>>{}; + + if constexpr (Is_SFA || (!Is_SFA && TmemAlloc == UMMA::TmemAllocMode::ScaleFactorDuplicated4by1)) { + // SFA, 2x2 and 4x1 data path + // SFB, 4x1 data path + auto tmem_atom = Layout < Shape< Shape< Shape, Int>, REP>, Shape, Int>>, + Stride, _32>, Stride< _0, _128>>>{}; + + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + auto final_tmem_layout = composition(tmem_restride, tmem_logical_layout); + return make_tensor(make_tmem_ptr(), final_tmem_layout); + } + else { + // SFB, 2x2 datapath + static_assert(!Is_SFA and TmemAlloc == UMMA::TmemAllocMode::ScaleFactorDuplicated2by2); + static_assert(N_SM == 2, "Should be 2x2 Datapath"); + // 2x2 Datapth + auto tmem_atom = Layout < Shape< Shape< Shape, Int>, _2, _2>, Shape, Int>>, + Stride, _64, _32>, Stride< _0, _128>>>{}; + + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + auto final_tmem_layout = composition(tmem_restride, tmem_logical_layout); + return make_tensor(make_tmem_ptr(), final_tmem_layout); + } + } +}; + +// Make C/D Tmem fragment for weight-stationary MMAs +template +struct tmem_frg_ws : tmem_frg_base +{ + static_assert(sizeof_bits_v <= sizeof_bits_v, "TMEM MMA allocations require StorageType big enough for ValueType."); + + // UMMA TMEM Allocator + // Each UMMA expects a specific MxN layout of TMEM for accumulators + // and sometimes a specific MxK layout of TMEM for A-values. + // @tparam ValueType The value type of the TMEM Tensor to allocate. + // @tparam StorageType The storage type of the TMEM Tensor to allocate. + // "Sparse" allocations often allocate ValueType=half_t within StorageType=uint32_t. + // "Dense" allocations often allocate ValueType=half_t within StorageType=half_t. + // @tparam N_SM The number of SMs in this UMMA_XSM instruction. + // @tparam TmemAlloc UMMA-specific allocation modifier for special cases. + // Some UMMA instructions expect strange atoms or tilings of atoms. + // @param tmem_shape ((M_MMA_SM,N_MMA_SM),MMA_M,MMA_N,...) + // The post-MMA-partitioned shape of TMEM to allocate. + // Note for UMMA_2SM_128xNx16, that M_MMA_SM will be 64, for example. + template + CUTE_HOST_DEVICE constexpr static auto + make(TmemShape const& tmem_shape) + { + CUTE_STATIC_ASSERT_V(size(tmem_shape)*Int)>{} <= TMEM::MAX_CAPACITY_BITS{}, + "Requesting more TMEM than is available."); + CUTE_STATIC_ASSERT_V(rank<0>(tmem_shape) == Int<2>{}, "Expected post-partitioned shape ((M_MMA,N_MMA),...)."); + constexpr int R = decltype(rank(tmem_shape))::value; + constexpr int M_MMA = decltype(size<0,0>(tmem_shape))::value; + constexpr int N_MMA = decltype(size<0,1>(tmem_shape))::value; + + // It's convenient to use "virtual tensor memory addressing" + // with DP_STRIDE=1, COL_STRIDE=128 to define the tmem_atom, + // then convert to "logical tensor memory addressing" on return. + using COL_ADDR = C::value / sizeof_bits::value>; + Layout tmem_restride = Layout, + Stride, COL_ADDR>>{}; + + static_assert(N_SM == 1, "UMMA.WS expects N_SM == 1"); + + static_assert(M_MMA == 32 || M_MMA == 64 || M_MMA == 128, + "Weight stationary UMMA_1SM M-mode size should be 32 or 64 or 128."); + static_assert(N_MMA == 64 || N_MMA == 128 || N_MMA == 256, + "Dense weight stationary UMMA_1SM N-mode size should be 64 or 128 or 256."); + // Weight Stationary MMA config + if constexpr (M_MMA == 32) + { + // 1x4 datapath + Layout tmem_atom = Layout, _4>>, + Stride< _1, Stride< _128,_32>> + >{}; + constexpr int tile_stride = 1; + // This will tile in DPs first, then COLs + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape), + compact_col_major(take<1,R>(tmem_shape), Int{}))); + // Restride for the DP/COL addressing and return + return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); + } else + if constexpr (M_MMA == 64) + { + // 2x2 datapath + Layout tmem_atom = Layout, _2>>, + Stride< _1, Stride< _128,_64>> + >{}; + constexpr int tile_stride = 1; + // This will tile in DPs first, then COLs + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape), + compact_col_major(take<1,R>(tmem_shape), Int{}))); + // Restride for the DP/COL addressing and return + return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); + } else + if constexpr (M_MMA == 128) + { + // For M_MMA = 128, all datapaths are occupied. TmemAllocMode doesn't change the allocation. + // Full subpartitions layout atom: (M,N) -> tmem_addr + Layout tmem_atom = Layout>, + Stride< _1, _128>>{}; + // This will tile in DPs first, then COLs + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1,R>(tmem_shape))); + // Restride for the DP/COL addressing and return + return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); + } + + CUTE_GCC_UNREACHABLE; + } +}; + +// Convenient aliases for common cases in the UMMA::ElementXFrg below +template +using tmem_frg_ws_1sm = tmem_frg_ws; + +} // end namespace UMMA + +// Customization point for creating a UMMA::tmem_frg Tensor +template +struct MakeTensor> +{ + template + CUTE_HOST_DEVICE constexpr auto + operator()(Shape const& tmem_shape) { + return UMMA::tmem_frg::make(shape(tmem_shape)); + } +}; + +template +struct MakeTensor> +{ + template + CUTE_HOST_DEVICE constexpr auto + operator()(Shape const& tmem_shape) { + return UMMA::tmem_frg_ws::make(shape(tmem_shape)); + } +}; + +// Customization point for creating a UMMA::tmem_e_frg Tensor +template +struct MakeTensor> +{ + template + CUTE_HOST_DEVICE constexpr auto + operator()(Shape const& tmem_shape) { + return UMMA::tmem_e_frg::make(shape(tmem_shape)); + } +}; + +template +struct MakeTensor> +{ + template + CUTE_HOST_DEVICE constexpr auto + operator()(Shape const& tmem_shape) { + return UMMA::tmem_e_frg_ws::make(shape(tmem_shape)); + } +}; + +template +struct MakeTensor> +{ + template + CUTE_HOST_DEVICE constexpr auto + operator()(Shape const& tmem_shape) { + return UMMA::tmem_sf_frg::make(shape(tmem_shape)); + } +}; + +} // namespace cute diff --git a/include/cute/config.hpp b/include/cute/config.hpp index fc5b11670..bbd8c431c 100644 --- a/include/cute/config.hpp +++ b/include/cute/config.hpp @@ -110,6 +110,20 @@ # define CUTE_STL_NAMESPACE std #endif +#if !defined(CUTE_CUDA_STD_STRUCTURED_BINDINGS_HEADER_AVAILABLE) +# if defined(__has_include) +# if __has_include() +# define CUTE_CUDA_STD_STRUCTURED_BINDINGS_HEADER_AVAILABLE 1 +# else +# define CUTE_CUDA_STD_STRUCTURED_BINDINGS_HEADER_AVAILABLE 0 +# endif +# elif (__CUDACC_VER_MAJOR__ >= 13) +# define CUTE_CUDA_STD_STRUCTURED_BINDINGS_HEADER_AVAILABLE 1 +# else +# define CUTE_CUDA_STD_STRUCTURED_BINDINGS_HEADER_AVAILABLE 0 +# endif +#endif + // // Assertion helpers // diff --git a/include/cute/container/array.hpp b/include/cute/container/array.hpp index 0e57859fc..7a58a9886 100644 --- a/include/cute/container/array.hpp +++ b/include/cute/container/array.hpp @@ -447,16 +447,16 @@ struct tuple_element> namespace std { -#if (__CUDACC_VER_MAJOR__ >= 13) +#if CUTE_CUDA_STD_STRUCTURED_BINDINGS_HEADER_AVAILABLE #include #else #if defined(__CUDACC_RTC__) - template + template struct tuple_size; - template + template struct tuple_element; #endif #endif diff --git a/include/cute/container/array_subbyte.hpp b/include/cute/container/array_subbyte.hpp index 2e6347140..03ae1bb87 100644 --- a/include/cute/container/array_subbyte.hpp +++ b/include/cute/container/array_subbyte.hpp @@ -617,16 +617,16 @@ struct tuple_element> namespace std { -#if (__CUDACC_VER_MAJOR__ >= 13) +#if CUTE_CUDA_STD_STRUCTURED_BINDINGS_HEADER_AVAILABLE #include #else #if defined(__CUDACC_RTC__) - template + template struct tuple_size; - template + template struct tuple_element; #endif #endif diff --git a/include/cute/container/tuple.hpp b/include/cute/container/tuple.hpp index f7cfca946..d89b6e409 100644 --- a/include/cute/container/tuple.hpp +++ b/include/cute/container/tuple.hpp @@ -701,16 +701,16 @@ struct tuple_element> namespace std { -#if (__CUDACC_VER_MAJOR__ >= 13) +#if CUTE_CUDA_STD_STRUCTURED_BINDINGS_HEADER_AVAILABLE #include #else #if defined(__CUDACC_RTC__) - template + template struct tuple_size; - template + template struct tuple_element; #endif #endif diff --git a/include/cute/container/type_list.hpp b/include/cute/container/type_list.hpp index 62ebf9dfd..149d90c6d 100644 --- a/include/cute/container/type_list.hpp +++ b/include/cute/container/type_list.hpp @@ -103,16 +103,16 @@ struct tuple_element> namespace std { -#if (__CUDACC_VER_MAJOR__ >= 13) +#if CUTE_CUDA_STD_STRUCTURED_BINDINGS_HEADER_AVAILABLE #include #else #if defined(__CUDACC_RTC__) - template + template struct tuple_size; - template + template struct tuple_element; #endif #endif diff --git a/include/cute/numeric/arithmetic_tuple.hpp b/include/cute/numeric/arithmetic_tuple.hpp index 01dde5f12..3c78eec00 100644 --- a/include/cute/numeric/arithmetic_tuple.hpp +++ b/include/cute/numeric/arithmetic_tuple.hpp @@ -513,16 +513,16 @@ struct tuple_element> namespace std { -#if (__CUDACC_VER_MAJOR__ >= 13) +#if CUTE_CUDA_STD_STRUCTURED_BINDINGS_HEADER_AVAILABLE #include #else #if defined(__CUDACC_RTC__) - template + template struct tuple_size; - template + template struct tuple_element; #endif #endif diff --git a/include/cutlass/arch/grid_dependency_control.h b/include/cutlass/arch/grid_dependency_control.h index 63a7714ee..6e926cdcf 100644 --- a/include/cutlass/arch/grid_dependency_control.h +++ b/include/cutlass/arch/grid_dependency_control.h @@ -94,7 +94,7 @@ void launch_dependent_grids() { CUTLASS_DEVICE void wait_on_dependent_grids() { #if (defined(CUTLASS_GDC_ENABLED)) - asm volatile("griddepcontrol.wait;"); + asm volatile("griddepcontrol.wait;" ::: "memory"); #endif } diff --git a/include/cutlass/gemm/collective/sm120_mma_array_tma.hpp b/include/cutlass/gemm/collective/sm120_mma_array_tma.hpp index 3d5c86875..dc668748b 100644 --- a/include/cutlass/gemm/collective/sm120_mma_array_tma.hpp +++ b/include/cutlass/gemm/collective/sm120_mma_array_tma.hpp @@ -29,10 +29,6 @@ * **************************************************************************************************/ -// CollectiveMma specialization for MainloopSm120ArrayTmaWarpSpecialized. -// Implements ptr-array (grouped GEMM) without blockwise scaling, targeting SM120 (GB10/DGX Spark). -// Fills the gap identified in https://github.com/NVIDIA/cutlass/issues/3263. - #pragma once #include "cutlass/cutlass.h" @@ -437,6 +433,7 @@ struct CollectiveMma< ++smem_pipe_write; } } + __syncwarp(); } CUTLASS_DEVICE void @@ -664,6 +661,11 @@ struct CollectiveMma< tensormaps_cp_fence_release( TensorMapStorage& shared_tensormaps, cute::tuple const& input_tensormaps) { + if (cute::elect_one_sync()) { + cute::tma_desc_commit_group(); + cute::tma_desc_wait_group(); + } + // Entire warp must do this (i.e. it's aligned) tma_descriptor_cp_fence_release(get<0>(input_tensormaps), shared_tensormaps.smem_tensormap_A); tma_descriptor_cp_fence_release(get<1>(input_tensormaps), shared_tensormaps.smem_tensormap_B); } diff --git a/include/cutlass/gemm/collective/sm120_mma_array_tma_blockwise_scaling.hpp b/include/cutlass/gemm/collective/sm120_mma_array_tma_blockwise_scaling.hpp index b10d803c1..c183bcd6c 100644 --- a/include/cutlass/gemm/collective/sm120_mma_array_tma_blockwise_scaling.hpp +++ b/include/cutlass/gemm/collective/sm120_mma_array_tma_blockwise_scaling.hpp @@ -972,6 +972,10 @@ struct CollectiveMma< tensormaps_cp_fence_release ( TensorMapStorage& shared_tensormaps, cute::tuple const& input_tensormaps) { + if (cute::elect_one_sync()) { + cute::tma_desc_commit_group(); + cute::tma_desc_wait_group(); + } // Entire warp must do this (i.e. it's aligned) tma_descriptor_cp_fence_release(get<0>(input_tensormaps), shared_tensormaps.smem_tensormap_A); tma_descriptor_cp_fence_release(get<1>(input_tensormaps), shared_tensormaps.smem_tensormap_B); diff --git a/media/docs/operators/api_reference/arguments.rst b/media/docs/operators/api_reference/arguments.rst new file mode 100644 index 000000000..9d14cad3b --- /dev/null +++ b/media/docs/operators/api_reference/arguments.rst @@ -0,0 +1,55 @@ +.. _operators_api_reference_arguments: + +Arguments and Operands +====================== + +An operation is expressed as a :class:`~cutlass.operators.RuntimeArguments` +object that describes both the operation kind (e.g. GEMM) and the operands it +runs on. Each operation has its own ``RuntimeArguments`` subclass. Operands are +described by :class:`~cutlass.operators.Operand` subclasses such as +:class:`~cutlass.operators.DenseTensor` and +:class:`~cutlass.operators.ScaledOperand`. Tensor- and numeric-typed fields +within Arguments and Operands accept any object matching the +:class:`~cutlass.operators.TensorLike` / :class:`~cutlass.operators.NumericLike` +protocols (e.g. ``torch.Tensor``, ``cute.Tensor``, ``cutlass.Numeric``, +``torch.dtype``). + +RuntimeArguments +---------------- + +.. automodule:: cutlass.operators.arguments.base + :members: + :exclude-members: Operand + :show-inheritance: + +.. automodule:: cutlass.operators.arguments.gemm + :members: + :show-inheritance: + +.. automodule:: cutlass.operators.arguments.grouped_gemm + :members: + :show-inheritance: + +.. automodule:: cutlass.operators.arguments.epilogue + :members: + :show-inheritance: + + +Operands +-------- + +.. autoclass:: cutlass.operators.arguments.base.Operand + :members: + :show-inheritance: + +.. automodule:: cutlass.operators.arguments.operand + :members: + :show-inheritance: + + +Type markers +------------ + +.. automodule:: cutlass.operators.typing + :members: + :show-inheritance: diff --git a/media/docs/operators/api_reference/discovery.rst b/media/docs/operators/api_reference/discovery.rst new file mode 100644 index 000000000..7e463a531 --- /dev/null +++ b/media/docs/operators/api_reference/discovery.rst @@ -0,0 +1,22 @@ +.. _operators_api_reference_discovery: + +Kernel discovery +================ + +:func:`~cutlass.operators.get_operators` is the entry point for finding +Operators that support a given operation, operands, and target. It returns +:class:`Operators ` drawn from the +:class:`~cutlass.operators.Manifest` of registered kernels. + +Operators can be additionally discovered or filtered by target compute +capability, which is described by :class:`~cutlass.operators.TargetSm`. + +.. autofunction:: cutlass.operators.get_operators + +.. automodule:: cutlass.operators.manifest + :members: + :show-inheritance: + +.. automodule:: cutlass.operators.arch + :members: + :show-inheritance: diff --git a/media/docs/operators/api_reference/index.rst b/media/docs/operators/api_reference/index.rst new file mode 100644 index 000000000..fc34722f6 --- /dev/null +++ b/media/docs/operators/api_reference/index.rst @@ -0,0 +1,20 @@ +.. _operators_api_reference: + +API reference +============= + +API Reference for the public interface of ``cutlass.operators``. + +For a more step-by-step introduction to core concepts and worked examples, +see :doc:`/media/docs/operators/tutorials/index`. + +The API reference is grouped into the following pages: + +.. toctree:: + :maxdepth: 1 + + operator + arguments + discovery + metadata + misc diff --git a/media/docs/operators/api_reference/metadata.rst b/media/docs/operators/api_reference/metadata.rst new file mode 100644 index 000000000..3a4c2bb6c --- /dev/null +++ b/media/docs/operators/api_reference/metadata.rst @@ -0,0 +1,45 @@ +.. _operators_api_reference_metadata: + +Metadata +======== + +.. automodule:: cutlass.operators.metadata.base + :members: + :show-inheritance: + + +Operands metadata +----------------- + +.. automodule:: cutlass.operators.metadata.operands + :members: + :show-inheritance: + +.. automodule:: cutlass.operators.metadata.operand_constraints + :members: + :show-inheritance: + + +Design metadata +--------------- + +.. automodule:: cutlass.operators.metadata.design + :members: + :show-inheritance: + + + +Epilogue metadata +----------------- + +.. automodule:: cutlass.operators.metadata.epilogue + :members: + :show-inheritance: + + +MMA instructions +---------------- + +.. automodule:: cutlass.operators.mma + :members: + :show-inheritance: diff --git a/media/docs/operators/api_reference/misc.rst b/media/docs/operators/api_reference/misc.rst new file mode 100644 index 000000000..2dcf8146f --- /dev/null +++ b/media/docs/operators/api_reference/misc.rst @@ -0,0 +1,20 @@ +.. _operators_api_reference_misc: + +Miscellaneous +============= + +Status +------ + +.. automodule:: cutlass.operators.status + :members: + :special-members: __bool__ + :show-inheritance: + + +Configuration +------------- + +.. automodule:: cutlass.operators.config + :members: + :show-inheritance: diff --git a/media/docs/operators/api_reference/operator.rst b/media/docs/operators/api_reference/operator.rst new file mode 100644 index 000000000..3b40b9183 --- /dev/null +++ b/media/docs/operators/api_reference/operator.rst @@ -0,0 +1,40 @@ +.. _operators_api_reference_operator: + +Operator interface +================== + +The :class:`~cutlass.operators.Operator` class is the central interface of Operator API. +Each ``Operator`` wraps a single "kernel" written in DSL and exposes a uniform interface for +executing it, generating its instances, and metadata that describes the kernel's capabilities. + +.. automodule:: cutlass.operators.base + :members: + :show-inheritance: + + +Compiled artifacts +------------------ + +:meth:`Operator.compile ` returns a +:class:`~cutlass.operators.CompiledArtifact` that bundles the compiled kernel +with the originating Operator and the target it was compiled for. It can be +passed back to :meth:`Operator.run ` to skip +recompilation. + +.. automodule:: cutlass.operators.artifact + :members: + :show-inheritance: + + +Workspace +--------- + +Some Operators need a scratch buffer to run. +:meth:`Operator.get_workspace_size ` +returns an :class:`~cutlass.operators.workspace.AllocationRequirement` that +the caller satisfies by passing a :class:`~cutlass.operators.workspace.Workspace` +to :meth:`Operator.run `. + +.. automodule:: cutlass.operators.workspace + :members: + :show-inheritance: diff --git a/media/docs/operators/overview.rst b/media/docs/operators/overview.rst new file mode 100644 index 000000000..ffc2f0bef --- /dev/null +++ b/media/docs/operators/overview.rst @@ -0,0 +1,157 @@ +.. _operators_overview: + +CUTLASS Operator API +==================== + +.. note:: + + CUTLASS Operator API is currently in beta. All interfaces, names, + and paths are subject to change. + +CUTLASS Operator API is a Python interface for integrating kernels written in CUTLASS +Python DSLs (like :doc:`CuTe DSL `) into your code. + +While DSLs focus on kernel authoring, CUTLASS Operator API focuses on ease of +managing and integrating those kernels into libraries that use CUTLASS. + +It views kernels as end-to-end "Operators" that execute an operation (like a GEMM), and provides two things: + +1. a **kernel-agnostic interface** for finding operators that support an + operation/operands, inspecting their properties, and executing them — the same + way regardless of which kernel each operator wraps. +2. a **registry** of officially maintained CUTLASS kernels exposed through that interface. + + +Example +------- + +.. code-block:: python + + import cutlass.operators as ops + import torch + + A, B, out = (torch.randn(128, 128, device="cuda", dtype=torch.float16) for _ in range(3)) + + # Arguments express an operation, and the operands to it + args = ops.GemmArguments(A, B, out, accumulator_type=torch.float32) + + # Find operators that support our provided arguments, and can run on SM100. + # Returns a list of ``Operator``s that wrap ready-to-compile CuTe DSL kernels. + operators = ops.get_operators(args, target_sm="100") + + # JIT compile and execute one of the returned operators using our arguments + operators[0].run(args) + + +Why use it? +=========== + +Any software that uses kernels requires finding kernels that do what you want, wiring up glue code to call them, +and maintaining the integration as kernels evolve. +Without an integration layer, these tasks are manual, error-prone, and repeated for every kernel. +CUTLASS Operator API eases each of these tasks — as kernels rapidly evolve, you don't have +to choose between integration churn and adoption inertia. + +**"Which kernel does what I want?"** + +Without Operator API, finding the right kernel means reading kernel source code and manually +deducing support — which dtypes, layouts, tile sizes, arch features, etc. each kernel supports. Operator API +provides a simple ``get_operators(args)``: express the operation you want to run, and get all the +operators that support it. Each operator also exposes uniform metadata describing +its constraints and design features for more advanced inspection, instead of requiring you to deduce +from source code. + +**"How do I get newer, faster, or fixed kernels?"** + +Adopting kernels directly carries a maintenance burden to mirror bug fixes/optimizations into local copies and +monitor release notes for new relevant kernels. With Operator API, you **integrate once, +not perpetually**. New kernels, fixes, and optimizations land in the registry on each release, and you get +them automatically without changing your integration code — just upgrade ``nvidia-cutlass-operators``. + +**"How do I call this kernel with my torch tensors?"** + +Different kernels have different usage conventions — direct usage requires kernel-specific +glue code to convert your framework tensors to the kernel's expected inputs, set performance +options, and run preparation steps. Operator API wraps every kernel with a consistent interface: pass PyTorch (or any +DLPack-compatible) tensors directly into ``GemmArguments`` and call ``operator.run(args)``. That lets you swap operators without touching your call site. + +It also supports: + +- **Custom epilogue fusions with ease** — pass a plain Python function; Operator API lowers it onto + CuTe DSL's Epilogue Fusion Configuration (EFC) framework and fuses it into supported kernels. +- **Bring-your-own-kernel** — register your own CuTe DSL kernels so you can call them through the same + interfaces as pre-bundled ones from CUTLASS, with no separate integration path for in-house kernels. +- **Negligible runtime overhead** on top of invoking the underlying kernel directly. + + +How to use it? +============== + +Installation +------------ + +To use with PyTorch, install the ``nvidia-cutlass-operators[torch]`` package: + +.. code-block:: bash + + pip install nvidia-cutlass-operators[torch] + +Alternatively, choose which dependencies to install: + +.. code-block:: bash + + # Install only nvidia-cutlass-operators core + pip install nvidia-cutlass-operators + + # Install all dependencies to develop, run tests, etc. + pip install nvidia-cutlass-operators[dev] + + +Next steps +---------- + +- `Core concepts and basic GEMM `__ — a guided tour of the API on a minimal GEMM, covering + ``RuntimeArguments``, operator discovery, compile/run, and metadata filtering. +- More guided notebooks in :doc:`tutorials/index`. +- :doc:`api_reference/index` — full API reference for ``cutlass.operators``. + + +Supported and Upcoming Features +=============================== + +CUTLASS Operator API will support a wide range of functionality, configurations, and optimizations, all robustly tested. + +**Current support:** + +- Kernel coverage + + - Dense GEMMs (F32, F16, BF16, INT8) for Blackwell, Hopper, Ampere + - Preferred and fallback cluster shapes + - Static and dynamic scheduling + - Block-scaled GEMMs (NVFP4, MXFP4, MXFP8, mixed input precision) for Blackwell + - Grouped GEMM (Contiguous offset) for Blackwell + - Low-latency TGV GEMM for Blackwell + +- Custom epilogue fusions (e.g. activations, elementwise ops, aux load/store) +- CUDA Graph support +- Native support for PyTorch and other DLPack tensors +- Bring-your-own-kernel + +**Upcoming support:** + +- Additional GEMM kernel coverage: Sparsity, performance optimizations, grouped GEMM variants, and more +- Ahead-of-time compilation +- JAX Graph support +- nvMatmulHeuristics support + +Community & Feedback +==================== + +We welcome contributions and feedback from the developer community. You can: + +- File bug reports, feature requests, documentation gaps, and upvote roadmap items on + the `CUTLASS GitHub Issues page `__. +- Contribute examples, tutorials, or improvements via pull request to the + `CUTLASS repository `__. +- Ask questions and share ideas in the + `#cutlass channel `__ on `NVIDIA Developer Discord `__. diff --git a/media/docs/operators/tutorials/000_gemm.nblink b/media/docs/operators/tutorials/000_gemm.nblink new file mode 100644 index 000000000..fd83a615f --- /dev/null +++ b/media/docs/operators/tutorials/000_gemm.nblink @@ -0,0 +1 @@ +{"path": "../../../../operators/examples/000_gemm.ipynb"} diff --git a/media/docs/operators/tutorials/001_gemm_with_fused_epilogue.nblink b/media/docs/operators/tutorials/001_gemm_with_fused_epilogue.nblink new file mode 100644 index 000000000..f87a61d9b --- /dev/null +++ b/media/docs/operators/tutorials/001_gemm_with_fused_epilogue.nblink @@ -0,0 +1 @@ +{"path": "../../../../operators/examples/001_gemm_with_fused_epilogue.ipynb"} diff --git a/media/docs/operators/tutorials/002_bring_your_own_kernel.nblink b/media/docs/operators/tutorials/002_bring_your_own_kernel.nblink new file mode 100644 index 000000000..10fde29b2 --- /dev/null +++ b/media/docs/operators/tutorials/002_bring_your_own_kernel.nblink @@ -0,0 +1 @@ +{"path": "../../../../operators/examples/002_bring_your_own_kernel.ipynb"} diff --git a/media/docs/operators/tutorials/003_host_latency_best_practices.nblink b/media/docs/operators/tutorials/003_host_latency_best_practices.nblink new file mode 100644 index 000000000..917bab091 --- /dev/null +++ b/media/docs/operators/tutorials/003_host_latency_best_practices.nblink @@ -0,0 +1 @@ +{"path": "../../../../operators/examples/003_host_latency_best_practices.ipynb"} diff --git a/media/docs/operators/tutorials/004_fake_tensors.nblink b/media/docs/operators/tutorials/004_fake_tensors.nblink new file mode 100644 index 000000000..da233a918 --- /dev/null +++ b/media/docs/operators/tutorials/004_fake_tensors.nblink @@ -0,0 +1 @@ +{"path": "../../../../operators/examples/004_fake_tensors.ipynb"} diff --git a/media/docs/operators/tutorials/005_grouped_gemm_contiguous_offset.nblink b/media/docs/operators/tutorials/005_grouped_gemm_contiguous_offset.nblink new file mode 100644 index 000000000..f7614bab3 --- /dev/null +++ b/media/docs/operators/tutorials/005_grouped_gemm_contiguous_offset.nblink @@ -0,0 +1 @@ +{"path": "../../../../operators/examples/005_grouped_gemm_contiguous_offset.ipynb"} diff --git a/media/docs/operators/tutorials/006_block_scaled_gemm.nblink b/media/docs/operators/tutorials/006_block_scaled_gemm.nblink new file mode 100644 index 000000000..d8365e1e7 --- /dev/null +++ b/media/docs/operators/tutorials/006_block_scaled_gemm.nblink @@ -0,0 +1 @@ +{"path": "../../../../operators/examples/006_block_scaled_gemm.ipynb"} diff --git a/media/docs/operators/tutorials/index.rst b/media/docs/operators/tutorials/index.rst new file mode 100644 index 000000000..c6d1e4a0c --- /dev/null +++ b/media/docs/operators/tutorials/index.rst @@ -0,0 +1,16 @@ +.. _operators_tutorials: + +Tutorials +========= + +Guided notebooks that walk through using the CUTLASS Operator API on concrete +examples. The notebooks are available in +`operators/examples/ `__ on GitHub: + +- `Core concepts and basic GEMM `__ +- `GEMM with fused epilogue `__ +- `Bringing your own kernel `__ +- `Host latency best practices `__ +- `Fake tensors for compilation `__ +- `Grouped GEMM with contiguous offset `__ +- `Block-scaled GEMM (MXFP8) `__ diff --git a/media/docs/pythonDSL/cute_dsl.rst b/media/docs/pythonDSL/cute_dsl.rst index f5675a85e..4a6740505 100644 --- a/media/docs/pythonDSL/cute_dsl.rst +++ b/media/docs/pythonDSL/cute_dsl.rst @@ -17,6 +17,7 @@ CuTe DSL JIT Types Integration with Frameworks Debugging with the DSL + IKET Profiling Autotuning with the DSL Educational Notebooks Deprecation Policy diff --git a/media/docs/pythonDSL/cute_dsl_general/debugging.rst b/media/docs/pythonDSL/cute_dsl_general/debugging.rst index 8b46dfd50..d35125c06 100644 --- a/media/docs/pythonDSL/cute_dsl_general/debugging.rst +++ b/media/docs/pythonDSL/cute_dsl_general/debugging.rst @@ -23,6 +23,64 @@ CuTe DSL provides Python code to PTX/SASS correlation to enable the profiling/de You can enable that globally via the environment variable CUTE_DSL_LINEINFO=1. Alternative, you can use compilation options to enable that per kernel. Please refer to :doc:`./dsl_jit_compilation_options` for more details. +Debug Mode +---------- + +To turn on a broad set of debugging aids at once, set the ``CUTE_DSL_DEBUG`` +environment variable. It is a convenience switch for diagnosing problems and for +reporting issues to the CUTLASS team: + +.. code:: bash + + # Enable debug mode (default: False) + export CUTE_DSL_DEBUG=1 + +When debug mode is enabled, CuTe DSL raises the defaults of several individual +debugging settings so you get more diagnostics from a single switch: + +- Line info is generated for Python-to-PTX/SASS correlation (same effect as + ``CUTE_DSL_LINEINFO=1``). +- Full, unfiltered Python stack traces are shown on failure (internal DSL + frames are no longer hidden). +- Optimization warnings that are normally suppressed are surfaced. +- Trace-time operation verification runs as operations are built, so malformed + operations are reported earlier instead of late in compilation. +- Full per-launch argument validation is performed, so a mismatched or + unsupported argument is reported with a clear error instead of failing later + inside the compiled kernel. + +Each of these behaviors is also controlled by its own environment variable, so +debug mode only changes their *defaults*, and setting a variable explicitly +takes precedence -- except trace-time operation verification, which stays on +while debug mode is enabled. For example, to enable debug mode but keep line +info off: + +.. code:: bash + + export CUTE_DSL_DEBUG=1 + export CUTE_DSL_LINEINFO=0 + +.. note:: + + Debug mode adds extra checks and diagnostics that increase compile time and + may affect the generated code (for example, by embedding line info). Enable + it while debugging, not for production runs. + +.. note:: + + The settings debug mode raises -- line info in particular -- change the + emitted IR/PTX, and every one of these settings is folded into the JIT + kernel cache key. A kernel compiled with debug mode on is therefore cached + separately from the same kernel compiled with it off: toggling + ``CUTE_DSL_DEBUG`` forces a recompile instead of reusing a cached kernel, + and the kernel you inspect or profile under debug mode is not identical to + the one produced for a normal (debug-off) run. Validate performance and + generated-code conclusions with debug mode disabled. Because these settings + are part of the cache key, a debug-built kernel is never silently reused for + a production run. See :doc:`JIT caching <./dsl_jit_caching>` for how the + cache key is formed. + + DSL Debugging ------------- @@ -89,6 +147,9 @@ Use ``CUTE_DSL_KEEP`` with a comma-separated list of artifact tokens: # Save CUBIN binary to a .cubin file export CUTE_DSL_KEEP=cubin + # Save SASS disassembly to a file (requires nvdisasm in PATH) + export CUTE_DSL_KEEP=sass + # Save LLVM IR to a file export CUTE_DSL_KEEP=llvm diff --git a/media/docs/pythonDSL/cute_dsl_general/framework_integration.rst b/media/docs/pythonDSL/cute_dsl_general/framework_integration.rst index 974777b48..11741ad20 100644 --- a/media/docs/pythonDSL/cute_dsl_general/framework_integration.rst +++ b/media/docs/pythonDSL/cute_dsl_general/framework_integration.rst @@ -515,14 +515,13 @@ issues with shape-1 dimension handling. # from cutlass.cute.runtime import make_ptr a_ptr = make_ptr( - cutlass.Float16, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + cutlass.Float16, a.data_ptr(), cutlass.AddressSpace.gmem, assumed_align=32 ) b_ptr = make_ptr( - cutlass.Float16, b.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + cutlass.Float16, b.data_ptr(), cutlass.AddressSpace.gmem, assumed_align=32 ) c_ptr = make_ptr( - cutlass.Float16, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + cutlass.Float16, c.data_ptr(), cutlass.AddressSpace.gmem, assumed_align=32 ) tensor_op_gemm_wrapper(a_ptr, b_ptr, c_ptr, m, n, k, l) - diff --git a/media/docs/pythonDSL/cute_dsl_general/iket_profiling.rst b/media/docs/pythonDSL/cute_dsl_general/iket_profiling.rst new file mode 100644 index 000000000..f9736d9c0 --- /dev/null +++ b/media/docs/pythonDSL/cute_dsl_general/iket_profiling.rst @@ -0,0 +1,910 @@ +.. _iket_profiling: + +IKET Profiling +============== + +.. warning:: + + IKET is an experimental profiling feature for CuTe DSL kernels. The API, + output format, profiler workflow, and overhead characteristics may change + in future releases. Users should understand the intended use cases and + limitations before interpreting profiling results. IKET dialect tool support + may also move to official NVIDIA Nsight tools in the future. + +IKET, short for In-Kernel Event Tracing, lets CuTe DSL kernels emit named +markers and ranges from inside the kernel. The ``run-iket`` profiler collects +those events and generates timeline output that can be inspected in the +Perfetto UI at https://ui.perfetto.dev/, along with machine-readable JSON +output. ``run-iket`` is a purpose-built standalone profiler for collecting IKET +traces in this experimental workflow. Conceptually, IKET is similar to CPU-side +NVTX ranges and markers, but the events are emitted from device code inside the +kernel. IKET supports Hopper and newer GPU architectures, including SM90, +SM100, SM103, SM110, and SM120. The ``run-iket`` profiler is released with the +``nvidia-cutlass-dsl`` package. + +IKET records device-side events at instrumentation points. The figure below is +a conceptual illustration, not a ``run-iket`` output format. It shows the kind +of producer/consumer activity that user-defined IKET ranges and markers can +make visible inside a kernel. Actual timeline viewing uses the Perfetto trace +described later in this guide. + +.. figure:: ../../../images/iket_concept_timeline.svg + :alt: Conceptual IKET timeline showing producer and consumer warp ranges + :width: 90% + + Conceptual view of IKET ranges and markers inside a kernel. ``run-iket`` + emits Perfetto and JSON traces; it does not emit this simplified diagram. + +Requirements +------------ + +Use a ``nvidia-cutlass-dsl`` installation that includes ``run-iket``. A quick +first check is: + +.. code-block:: bash + + run-iket --help + +The profiled workload must run on a supported GPU architecture and must +JIT-compile the instrumented CuTe DSL kernel during the ``run-iket`` profiling +process. Kernels that are already compiled and reused without recompilation do +not gain IKET instrumentation during that run. + +End-to-End Quick Start +---------------------- + +This section shows the minimal end-to-end flow: add IKET calls in kernel code, +run the workload under ``run-iket``, and open the generated trace. + +**Step 1: Add IKET instrumentation.** Place IKET calls inside a +``@cute.kernel`` function. IKET calls in host-side Python wrappers do not emit +in-kernel events. + +The following example shows a small kernel fragment with a marker, a +token-based range, and a stack-based range. + +.. code-block:: python + + import cutlass + import cutlass.cute as cute + + + @cute.kernel + def kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor): + bidx, _, _ = cute.arch.block_idx() + + cute.experimental.iket.mark("kernel_start", bidx) + + load_token = cute.experimental.iket.range_start("load") + # Load data from gA and gB. + cute.experimental.iket.range_end(load_token) + + cute.experimental.iket.range_push("compute") + # Compute and store results. + cute.experimental.iket.range_pop() + +A complete CuTe DSL GEMM example with IKET instrumentation is available at +``examples/python/CuTeDSL/dsl_tutorials/fp16_gemm_4_iket.py``. + +**Step 2: Run the application under** ``run-iket``. The profiler automatically +requests IKET lowering for kernels that are JIT-compiled during the profiled +run. + +.. code-block:: bash + + run-iket profile --postprocess perfetto -- \ + python fp16_gemm_4_iket.py \ + --mnk 512,1024,64 + +**Step 3: Open the trace.** Open the generated ``*.pftrace`` file in the +Perfetto UI at https://ui.perfetto.dev/ and inspect the in-kernel markers and +ranges. + +.. figure:: ../../../images/fp16_gemm_4_iket_quickstart.png + :alt: Cropped IKET timeline in Perfetto UI + :width: 95% + + Cropped Perfetto view from the GEMM tutorial, showing nested IKET ranges + across several warp roles. + +The rest of this guide explains the API, instrumentation patterns, trace output, +limitations, and overhead guidance in more detail. + + +IKET API calls are stripped by default. If neither ``run-iket`` is used to +profile the target kernel nor an explicit compile option enables IKET lowering, +the ``iket.*`` operations do not add instrumentation code to the final kernel. + +To build IKET instrumentation outside a ``run-iket`` profiling run, enable IKET +for all JIT compilations in the process: + +.. code-block:: bash + + export CUTE_DSL_COMPILER_OPT=iket + python my_kernel.py + +Alternatively, enable IKET for one explicit compilation: + +.. code-block:: python + + compiled = cute.compile(host_function, *args, options="iket") + +API Reference +------------- + +IKET APIs are available under ``cutlass.cute.experimental.iket``. The table +below uses ``iket`` as shorthand for that module. IKET API calls should be +placed inside ``@cute.kernel`` code. Instrumentation in host-side Python code +does not create in-kernel events. + +IKET has three basic concepts: + +* An event is one warp-level runtime record emitted by a kernel. Each event + records a timestamp and the metadata needed to identify it in the trace. +* A marker is a point annotation and emits one event. +* A range represents a duration and is usually built from two events: one at + the start and one at the end. Each event can optionally carry one payload + value to record a runtime variable. + +.. list-table:: + :header-rows: 1 + :widths: 25 45 30 + + * - API + - Purpose + - Notes + * - ``iket.mark(name)`` + - Emit a single timestamped marker. + - Use for point events. + * - ``iket.mark(name, payload)`` + - Emit a marker with a numeric payload. + - The payload is stored with the event. + * - ``iket.range_push(name)`` + - Start a stack-based range. + - Closed by the next matching ``iket.range_pop()`` in LIFO order. + * - ``iket.range_push(name, payload)`` + - Start a stack-based range with a numeric payload. + - The payload is attached to the push event. + * - ``iket.range_pop()`` + - End the most recent stack-based range. + - Does not take a range name. + * - ``iket.range_start(name)`` + - Start a token-based range and return a token. + - Closed by ``iket.range_end(token)``. + * - ``iket.range_start(name, payload)`` + - Start a token-based range with a numeric payload. + - Closed by ``iket.range_end(token, payload)`` with matching payload type. + * - ``iket.range_end(token)`` + - End a token-based range. + - The token must come from ``iket.range_start`` or ``iket.sentinel_token``. + * - ``iket.range_end(token, payload)`` + - End a token-based range with a numeric payload. + - The corresponding ``iket.range_start`` call must also have a payload, + and the payload types must match. + * - ``iket.sentinel_token(name)`` + - Create a token without a real runtime event for cross-iteration ranges. + - Use it when ``range_end`` appears before the later ``range_start`` in + source order. + +Choosing a Range API +~~~~~~~~~~~~~~~~~~~~ + +CuTe DSL IKET provides two valid range-pairing models. Choose the one that makes +the pairing easiest to see in the kernel source. + +Use ``range_push`` / ``range_pop`` when the range is naturally nested and the +push and pop calls can stay in the same structured scope. This is often the +clearest shape for phase-style instrumentation such as setup, mainloop, wait, +issue, and epilogue ranges. + +Use ``range_start`` / ``range_end`` when an explicit handle makes the pairing +clearer. This can be useful when a range ends at a later synchronization point, +crosses an iteration boundary, or has multiple mutually exclusive close sites. + +Payloads +~~~~~~~~ + +Payloads attach a runtime value to an event. They are useful for recording +values such as loop indices, block coordinates, or small computed metrics. +Supported payloads include Python boolean, integer, and floating-point literals, +plus CuTe DSL numeric and index scalar values. Do not use tensors, tuples, or +other aggregate values as payloads. Prefer warp-uniform payload values, such as +loop indices or block coordinates, when they describe the event clearly. For +example: + +.. code-block:: python + + for k_tile in cutlass.range(k_tile_count): + cute.experimental.iket.range_push("k_tile", k_tile) + # Work for this K tile. + cute.experimental.iket.range_pop() + +IKET events are warp-level events. If active threads in the participating warp +evaluate a payload expression to different values, the dumped payload value is +from the first active thread. To record the payload value from a specific +thread, guard the IKET call with a predicate such as ``if tidx == 0:``. For +range endpoints, guard paired endpoints consistently. + +Plain Python integer literals are emitted as 32-bit integer payloads, and plain +Python floating-point literals are emitted as 32-bit floating-point payloads. +Use explicit CuTe DSL scalar types for 64-bit literal payloads: + +.. code-block:: python + + cute.experimental.iket.mark("large_count", cutlass.Int64(0x100000000)) + cute.experimental.iket.mark("scale", cutlass.Float64(3.141592653589793)) + +For token-based ranges, the start and end payload forms must match. It is not +allowed to start a range with a payload and end it without one, or to use +different payload types between ``range_start`` and ``range_end``. + +Example Instrumentation Patterns +-------------------------------- + +The examples below use ``cute.experimental.iket`` inside ``@cute.kernel`` code. +IKET calls in host-side Python wrappers do not emit in-kernel events. + +Before Adding Events +~~~~~~~~~~~~~~~~~~~~ + +Start by identifying the kernel body and the work you want to measure. + +1. Find the ``@cute.kernel`` function. Host-side ``@cute.jit`` functions and + launch wrappers are useful context, but IKET instrumentation should be placed + in device kernel code. +2. Split the kernel into natural phases. For a GEMM-shaped kernel this may be + setup, TMA or copy issue, mainloop, MMA, waits, and epilogue. Other kernels + should use names that match their own algorithmic phases. +3. Note warp-specialized regions such as ``if warp_idx == 0:`` or + ``if is_leader_cta:``. Put both ends of a range inside the same role or + guard when the work is role-specific. +4. Identify asynchronous work. For example, TMA copies, ``cp.async``-style + copies, WGMMA or MMA issue, and pipeline or mbarrier operations may have + separate issue and completion points. + +Coarse Phase Timing +~~~~~~~~~~~~~~~~~~~ + +Begin with a small number of coarse ranges. This provides orientation in the +trace and keeps overhead manageable while you decide where more detail is +needed. + +.. code-block:: python + + @cute.kernel + def kernel(...): + user_warp_lifetime = cute.experimental.iket.range_start( + "user_warp_lifetime" + ) + + cute.experimental.iket.range_push("setup") + # Allocate/register fragments, partition tensors, initialize pipelines. + cute.experimental.iket.range_pop() # setup + + cute.experimental.iket.range_push("mainloop") + for k_tile in cutlass.range(k_tile_count): + # Main loop body. + cute.experimental.iket.range_pop() # mainloop + + cute.experimental.iket.range_push("epilogue") + # Convert accumulators and store results. + cute.experimental.iket.range_pop() # epilogue + + cute.experimental.iket.range_end(user_warp_lifetime) + +For warp-specialized code, place the range inside the guard for the warp that +does the work: + +.. code-block:: python + + if warp_idx == tma_warp_id: + cute.experimental.iket.range_push("tma_main") + # TMA producer work. + cute.experimental.iket.range_pop() # tma_main + + if warp_idx == mma_warp_id: + cute.experimental.iket.range_push("mma_main") + # MMA consumer work. + cute.experimental.iket.range_pop() # mma_main + +Timing Waits and Async Work +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For asynchronous operations, decide whether the range measures issue time or +completion/wait time. An event immediately after an async issue point measures +issue-side timing. Completion is usually observed at a pipeline or mbarrier wait. + +To measure issue time: + +.. code-block:: python + + issue_token = cute.experimental.iket.range_start("tma_issue", k_tile) + cute.copy(tma_atom, src_tensor, dst_tensor, tma_bar_ptr=barrier) + cute.experimental.iket.range_end(issue_token, k_tile) + +To measure wait time: + +.. code-block:: python + + cute.experimental.iket.range_push("ab_wait") + ab_full = ab_consumer.wait_and_advance() + cute.experimental.iket.range_pop() + +The same wait pattern can be used around pipeline acquire calls, mbarrier waits, +allocator waits, or other synchronization points whose return marks completion +of the waited-for work. + +Cross-Iteration Wait-Boundary Timing +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Some pipelined loops start work for iteration ``N`` and observe the next useful +boundary for that work in iteration ``N + 1``. For example, the next iteration +may reach a pipeline wait, mbarrier wait, or other ``wait_and_advance``-style +call before the previous tile's range should close. Use ``sentinel_token`` to +initialize the token before the loop. Creating the sentinel token emits no +runtime event. Calling ``range_end`` on the initial sentinel token is valid and +emits no runtime event; after the variable is replaced by a token from +``range_start``, ``range_end`` emits the runtime end event for that real range. + +.. code-block:: python + + iter_token = cute.experimental.iket.sentinel_token("mma_k_tile") + + for k_tile in cutlass.range(k_tile_count): + ... # some setup codes for k_tile + ab_full = ab_consumer.wait_and_advance() + + # Close the previous tile only after this wait boundary is reached. + if k_tile > 0: + cute.experimental.iket.range_end(iter_token) + + iter_token = cute.experimental.iket.range_start("mma_k_tile") + # Work for this tile. + cute.gemm(tiled_mma, tCtAcc, tCrA, tCrB, tCtAcc) + ab_full.release() + + if k_tile_count > 0: + ... # final drain or synchronization boundary for the last tile + cute.experimental.iket.range_end(iter_token) + +Use this pattern only when the cross-iteration boundary is meaningful. For a +simple per-iteration range whose start and end are both inside the same loop +iteration, a push/pop pair inside the loop is simpler: + +.. code-block:: python + + for k_tile in cutlass.range(k_tile_count): + cute.experimental.iket.range_push("k_tile", k_tile) + # Work for this tile. + cute.experimental.iket.range_pop() + +Warp-Specialized Mainloop Example +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The following skeleton shows how to layer ranges by role and by loop level in a +warp-specialized kernel. Adapt the role names and phase names to the actual +kernel. + +.. code-block:: python + + @cute.kernel + def kernel(...): + user_warp_lifetime = cute.experimental.iket.range_start( + "user_warp_lifetime" + ) + + # Work shared by all participating warps. + cute.experimental.iket.range_push("prologue") + # Tensor partitioning, pipeline setup, scheduler setup. + cute.experimental.iket.range_pop() # prologue + + if warp_idx == tma_warp_id: + cute.experimental.iket.range_push("tma_main") + while work_tile.is_valid_tile: + cute.experimental.iket.range_push("tma_tile") + + for k_tile in cutlass.range(k_tile_count): + cute.experimental.iket.range_push("tma_k_tile", k_tile) + ... # some setup codes for k_tile + + cute.experimental.iket.range_push("tma_acquire") + ab_empty = ab_producer.acquire_and_advance() + cute.experimental.iket.range_pop() # tma_acquire + + issue_token = cute.experimental.iket.range_start( + "tma_issue", k_tile + ) + cute.copy(tma_a_atom, tAgA, tAsA, tma_bar_ptr=ab_empty.barrier) + cute.copy(tma_b_atom, tBgB, tBsB, tma_bar_ptr=ab_empty.barrier) + cute.experimental.iket.range_end(issue_token, k_tile) + + cute.experimental.iket.range_pop() # tma_k_tile + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + cute.experimental.iket.range_pop() # tma_tile + + ab_producer.tail() + cute.experimental.iket.range_pop() # tma_main + + if warp_idx == mma_warp_id: + cute.experimental.iket.range_push("mma_main") + while work_tile.is_valid_tile: + cute.experimental.iket.range_push("mma_tile") + + for k_tile in cutlass.range(k_tile_count): + cute.experimental.iket.range_push("mma_k_tile", k_tile) + + cute.experimental.iket.range_push("ab_wait") + ab_full = ab_consumer.wait_and_advance() + cute.experimental.iket.range_pop() # ab_wait + + cute.experimental.iket.range_push("mma_issue") + cute.gemm(tiled_mma, tCtAcc, tCrA, tCrB, tCtAcc) + cute.experimental.iket.range_pop() # mma_issue + + ab_full.release() + cute.experimental.iket.range_pop() # mma_k_tile + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + cute.experimental.iket.range_pop() # mma_tile + + cute.experimental.iket.range_pop() # mma_main + + cute.experimental.iket.range_end(user_warp_lifetime) + +This style makes each role visible in the trace. Prefixing names with the role +such as ``tma_`` and ``mma_`` also makes JSON output easier to filter. + +Instrumentation Guidelines +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Consider the following when placing IKET calls. These choices affect what the +profiler can reconstruct from the runtime events, so unclear or mismatched +instrumentation may produce results that do not match the intended measurement +or even cause the profiler to fail when postprocessing the data. + +* Every dynamic ``range_push`` has exactly one matching ``range_pop`` on each + participating warp execution path, in LIFO order. +* Every dynamic non-sentinel ``range_start`` is closed by ``range_end`` on each + participating warp execution path. It is valid to close the same token in + multiple mutually exclusive branches, as long as each executed path closes it + once. +* Start and end points are in the same warp role and compatible control-flow + path. Do not start a range in one thread-divergent branch and close it in + another. Violating this may cause undefined profiling results. +* Payload-bearing token ranges use matching payload types at start and end. +* Event names are at most 32 characters. +* Reuse the same descriptive name for the same recurring phase inside a loop. + This creates many runtime events but only one unique marker or range name. +* When more than 30 unique marker or range names are used, instrumentation + overhead increases. +* Avoid high-frequency events in innermost unrolled loops unless that detail is + necessary. IKET events can affect compiler scheduling and can create large + traces. +* Do not put IKET range operations inside + ``cutlass.range(..., prefetch_stages=...)``. That loop form is not currently + supported for IKET range instrumentation. + +Profiling with the ``run-iket`` Tool +------------------------------------ + +The ``run-iket`` tool is installed with the ``nvidia-cutlass-dsl`` package. +During profiling, ``run-iket`` automatically enables IKET lowering for +JIT-compiled kernels so that the final kernel contains instrumentation code. +The application command must appear after ``--`` so that workload arguments are +not parsed as profiler arguments. + +.. code-block:: bash + + run-iket \ + --output-dir ./iket_output \ + --clobber \ + profile \ + --postprocess all \ + -- \ + python my_kernel.py + +Important options: + +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Option + - Description + * - ``--output-dir `` + - Directory for profiler output and intermediate files. + * - ``--clobber`` + - Remove an existing output directory and create a new one without + prompting. + * - ``profile`` + - Start a profiling run. + * - ``--postprocess perfetto`` + - Generate a Perfetto timeline trace. + * - ``--postprocess json`` + - Generate JSON output for script-based analysis. + * - ``--postprocess all`` + - Generate both Perfetto and JSON output. + +Output Files +~~~~~~~~~~~~ + +With ``--postprocess perfetto``, ``run-iket`` writes one or more ``*.pftrace`` +files under the output directory. With ``--postprocess json``, it also writes +JSON traces containing the collected ranges, markers, timestamps, and payloads. + +The exact filenames may include the profiled process ID. If the workload uses +multiple processes or GPUs, ``run-iket`` may produce separate traces that are +**NOT** aligned to a single global timeline. + +For a single-process run with ``--postprocess all``, the output directory +contains files shaped like this: + +.. code-block:: text + + iket_output/ + *.pftrace # Open this in Perfetto UI. + *.json # Use this for script-based analysis. + +There may also be profiler intermediate files in the same directory. Start with +the ``*.pftrace`` file for visual inspection, then use the JSON file when you +need scripted filtering or aggregation. + +JSON Trace Shape +~~~~~~~~~~~~~~~~ + +The JSON output is intended for script-based analysis. Its schema is also +experimental and may change in future releases. A trace is organized around +profiled kernel launches, with ranges, markers, warp locations, payload fields, +and warp lifetimes. A simplified trace looks like this: + +.. code-block:: json + + { + "launches": [ + { + "gridId": 0, + "kernelName": "my_kernel", + "ranges": [ + { + "rangeName": "mainloop", + "rangeScope": 0, + "startTs": 1000, + "endTs": 2500, + "warpLocs": [ + { + "smId": 0, + "tpcId": 0, + "gpcId": 0, + "ctaId": [0, 0, 0], + "warpId": 0 + } + ], + "internalEvents": [] + } + ], + "markers": [ + { + "markerName": "checkpoint", + "timestamp": 1500, + "location": { + "smId": 0, + "tpcId": 0, + "gpcId": 0, + "ctaId": [0, 0, 0], + "warpId": 0 + }, + "payloadType": 0, + "payloadVal": 0 + } + ], + "warpLifetimes": [ + { + "startTs": 900, + "endTs": 3800, + "warpLocation": { + "smId": 0, + "tpcId": 0, + "gpcId": 0, + "ctaId": [0, 0, 0], + "warpId": 0 + } + } + ] + } + ] + } + +Important fields include: + +* ``launches[]``: profiled kernel launches. +* ``gridId`` and ``kernelName``: launch identity and kernel name. +* ``ranges[]``: duration ranges with ``rangeName``, ``startTs``, ``endTs``, + ``rangeScope``, and one or more ``warpLocs``. Use ``endTs - startTs`` as the + range duration. +* ``markers[]``: point events with ``markerName``, ``timestamp``, + ``location``, and optional payload fields such as ``payloadType`` and + ``payloadVal``. +* ``warpLifetimes[]``: active spans for warps observed during the profiled + launch. + +Timestamp values are in a trace-local timebase. Compare timestamp differences +within one trace, but do not compare absolute timestamp values across separate +traces. ``rangeScope`` and ``payloadType`` are profiler metadata fields whose +exact numeric values are experimental. A ``warpLocs`` entry identifies the GPU +location for a warp that emitted or participated in the range. + +Viewing a Trace in Perfetto +--------------------------- + +Open the generated ``*.pftrace`` file in the Perfetto UI at +https://ui.perfetto.dev/. + +The basic workflow is: + +1. Open the trace file. +2. Search for the profiled kernel or zoom into the kernel region. Use the + ``W`` / ``A`` / ``S`` / ``D`` keys on your keyboard to pan and zoom. +3. Expand the relevant tracks. +4. Click a marker or range to inspect its name, timing, and payload values. + +The trace is organized around profiled kernel launches and the warp-level IKET +records collected inside those launches. After expanding a kernel region, look +for tracks grouped by the recorded GPU location, such as CTA and warp identity. +The visible marker and range names come from the strings passed to the IKET API +in the kernel. + +The image below shows a more complete expanded trace from the GEMM tutorial. +Use it as a guide to the track structure: + +* The left track hierarchy is expanded by GPU location, then by CTA and warp. +* ``WarpLifeTime`` tracks are generated automatically for kernels that contain + IKET instrumentation. +* Long ranges under each warp show user-instrumented phases such as + ``kernel_e2e``, ``epi_main``, ``mma_tile``, and ``tma_tile``. Token-based + ``range_start`` / ``range_end`` ranges use separate tracks for each range + name. +* Stack-based ``range_push`` / ``range_pop`` ranges are shown on + ``StackedRanges`` tracks, and markers are shown on ``Marker`` tracks. +* Very short colored blocks and marker glyphs represent fine-grained events. + Payload values may appear in event labels and are also available by clicking + the event and inspecting the details panel. + +.. figure:: ../../../images/fp16_gemm_4_iket_instrument.png + :alt: Example IKET timeline in Perfetto UI + :width: 100% + + Example IKET timeline in Perfetto UI. The exact tracks and event names + depend on the kernel instrumentation and workload. This trace was generated + from + ``examples/python/CuTeDSL/dsl_tutorials/fp16_gemm_4_iket.py``. + +Trace viewing is powered by the Perfetto UI (https://ui.perfetto.dev/), part of +the Perfetto project (https://perfetto.dev/) licensed under the Apache License +2.0. Perfetto UI is provided by a third-party site. This product does not modify +or redistribute Perfetto UI code. Perfetto is only a viewer for the generated +trace. The trace content comes from the IKET instrumentation emitted by the CuTe +DSL kernel and collected by ``run-iket``. + +Assumptions, Limitations, and Impact +------------------------------------ + +The ``run-iket`` profiler assumes well-formed instrumentation and sufficiently +convergent execution within participating warps. IKET events are warp-level +records, so placement is easiest to interpret when all participating threads in +a warp follow the same instrumentation path. + +Range Pairing and Warp Divergence +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For token-based ranges, every dynamic non-sentinel ``range_start`` must be +closed by ``range_end`` on each participating warp execution path. A token may +be closed in multiple mutually exclusive branches, but each executed path should +close that dynamic range once. For stack-based ranges, every ``range_push`` must +be balanced by ``range_pop`` using LIFO stack semantics. + +Avoid these patterns: + +* Starting a range in one thread-divergent branch and ending it in another + branch. +* Returning or otherwise exiting early between paired range endpoints. +* Emitting different push/pop nesting on different warp execution paths. + +If a range pair is inside a branch where different threads in a warp diverge, +the trace may be incomplete or may reflect serialized divergent execution. +Incorrectly paired ``range_start`` / ``range_end`` or ``range_push`` / +``range_pop`` calls may cause profiling to fail, may produce incorrect timeline +visualization, or may cause undefined profiling results. Warp-uniform placement +gives cleaner and easier-to-interpret results. + +Event Name Count and Overhead +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Unique user names include marker names and range names from ``mark``, +``range_start``, and ``range_push``. Repeatedly emitting the same name is +supported and is the expected way to record a loop phase or recurring marker. + +Kernels with more than 30 unique marker or range names may use a wider event +encoding during IKET lowering, which can increase instrumentation overhead. +Keep the number of unique names as small and stable as practical, especially in +performance-sensitive kernels. + +Event names may use arbitrary characters and must be at most 32 characters. +Longer names are not supported; use a shorter stable name instead. + +Timestamp Semantics +~~~~~~~~~~~~~~~~~~~ + +An IKET timestamp records the instrumentation point itself. The timer +granularity is 32 ns. If a range starts and ends very close together, its start +and end timestamps may be identical. In the final Perfetto visualization, such +a very short range may look similar to a marker. For asynchronous operations +such as TMA copies, placing an end event immediately after the issue point +usually measures issue-side timing, not completion. To measure completion or +wait time, place the corresponding range endpoint around the synchronization or +wait point where completion is observed. + +Workload Size +~~~~~~~~~~~~~ + +``run-iket`` profiles the whole workload it launches. It does not currently +support selecting a smaller profiling or capture window. Prefer small workloads +with a limited number of instrumented kernel launches while collecting IKET +traces. A kernel with many IKET events can generate a large amount of data +because records are collected per warp, and often across many CTAs. Large +workloads, such as workloads with many instrumented kernel launches, may run +much more slowly under the profiler or may run out of memory. + +Kernel Launch and Overlap Timing +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +``run-iket`` can collect traces from workloads that launch multiple +instrumented kernels. However, IKET is intended for in-kernel timing. Do not use +the IKET trace to measure host-side launch latency, inter-kernel launch gaps, +kernel overlap, or CPU/GPU scheduling behavior. These workload-level analyses +are not IKET's target use case, and the corresponding timing views include +additional ``run-iket`` overhead and are not currently optimized for launch or +overlap analysis. When a trace contains multiple kernels, interpret the IKET +markers and ranges within each individual kernel launch. + +Use NVIDIA Nsight Systems in a separate run when you need accurate kernel launch +latency, kernel overlap, CPU/GPU scheduling, or whole-application timeline +analysis. + +Profiler Compatibility +~~~~~~~~~~~~~~~~~~~~~~ + +``run-iket`` cannot run at the same time as NVIDIA Nsight Compute, NVIDIA +Nsight Systems, or other CUPTI-based profiling and tracing tools due to +conflicts over driver profiling resources. Do not run them together on the same +workload. Use separate runs when collecting IKET traces and other profiler +outputs. + +Buffer Sizing +~~~~~~~~~~~~~ + +``run-iket`` uses multiple profiling passes to allocate device-side trace +buffers. An initial pass estimates how much buffer space the workload needs, +and a later pass allocates memory with some margin and collects the timestamp +and payload records. + +This assumes the number of emitted IKET records per warp is reasonably stable +between those passes. If a kernel emits a very different number of records per +warp between the sizing pass and the collection pass, the trace may contain +incorrect data, and the workload may fail with an illegal memory access. Prefer +deterministic profiled workloads and avoid data-dependent instrumentation rates +that vary substantially between runs. + +Unsupported Prefetch-Stage Loops +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +IKET range operations are not currently supported inside a ``cutlass.range`` +loop that uses ``prefetch_stages``, such as +``cutlass.range(..., prefetch_stages=...)``. This is a known limitation of the +prefetch-stage loop form. Place IKET range instrumentation outside that loop, or +use a loop form without prefetch stages when profiling that region. + +Compiler and Runtime Impact +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +IKET events do slightly change the generated kernel code. Avoid placing too many +events in innermost hot loops, especially loops the compiler may unroll. IKET +markers and ranges can act partly like code-motion barriers and may reduce loop +interleaving optimizations after unrolling. + +High-frequency instrumentation, many unique range or marker names (especially +more than 30), payloads, many kernel launches, or large workloads can increase +overhead, overflow buffers, alter compiler behavior, or create large traces. +Payloads also increase the amount of stored trace data. + +IKET profiling may also add fixed per-kernel entry and exit overhead. This +overhead can affect the trace-reported kernel duration, but it is not shown as +a separate IKET marker or range on the Perfetto timeline. + +IKET profiling adds CPU-side overhead in addition to in-kernel overhead. Host +wall-clock measurements outside CUDA Driver/API timing are therefore not a +clean measure of kernel event overhead. + +Performance Overhead Guidance +----------------------------- + +The primary way to compare profiled overhead is to use kernel durations from +the IKET Perfetto trace. Do not use application wall-clock timing or timing +taken outside CUDA Driver/API boundaries as the measure of event overhead. + +For a reasonable comparison: + +1. Instrument the target kernel. +2. Run ``run-iket`` with a minimal instrumentation variant that emits one + instrumented event site, so the profiler recognizes and profiles the kernel. +3. Run ``run-iket`` again with many instrumented event sites or with the + intended full event set in the instrumented kernel. +4. Compare the trace-reported kernel execution times in Perfetto. + +The difference between those trace-reported kernel times estimates incremental +event overhead. To report amortized overhead, divide the runtime delta by the +number of executed instrumentation points, preferably on a per-warp basis. + +As a separate compiler and binary sanity check, you can compare an +uninstrumented kernel against the same kernel compiled with IKET instrumentation +enabled but not profiled. To build that IKET-enabled binary outside a +``run-iket`` profiling run, use ``CUTE_DSL_COMPILER_OPT=iket`` or +``cute.compile(..., options="iket")``. This is not a replacement for the +Perfetto-based profiled-overhead comparison above. + +Payloads can significantly increase overhead and trace volume. 64-bit payloads +are more expensive than no-payload or 32-bit payload events. + +How It Works +------------ + +At compile time, CuTe DSL emits ``iket.*`` IR operations and event metadata for +the IKET API calls in the kernel. By default, the compiler strips the ``iket.*`` +operations before lowering, so IKET calls do not contribute instrumentation code +to the final kernel. + +When IKET lowering is enabled, either explicitly with +``CUTE_DSL_COMPILER_OPT=iket`` or ``cute.compile(..., options="iket")``, or +automatically by ``run-iket`` during a profiling run, lowering emits placeholder +instrumentation and metadata into the compiled kernel. + +During profiling, ``run-iket`` prepares the kernel for collection, patches the +placeholder instrumentation at runtime, collects timestamp and payload records, +and postprocesses the records into Perfetto and JSON output. + +Troubleshooting +--------------- + +Empty or missing trace + Confirm that the workload is launched under ``run-iket`` and that the + instrumented kernels are JIT-compiled in that profiled process. If running + without ``run-iket``, confirm that ``CUTE_DSL_COMPILER_OPT=iket`` is set or + that the kernel is compiled with ``options="iket"``. + +Trace is very large or profiling fails + Reduce the number of instrumented kernel launches, reduce event frequency, + or reduce payload use. If the workload can legitimately emit more records + per warp than the profiler detects automatically, increase + ``--max-ts-cnt-per-warp `` to reserve space for up to ``N`` events per + warp. Choose ``N`` above the largest expected number of marker, range-start, + range-end, range-push, and range-pop events emitted by one warp in one + kernel launch. + +Expected events do not appear + Confirm the IKET calls are inside ``@cute.kernel`` code and are reachable + on the execution path being profiled. Also check that paired range + endpoints are not split across divergent branches. If the workload caches a + compiled kernel or executor, make sure that compilation happens inside the + ``run-iket`` profiling process or clear the application-level compiled + kernel cache before profiling. + +Unexpected timing around asynchronous work + Move range endpoints to the wait or synchronization point that observes + completion. + +Inspect generated IR + Use existing CuTe DSL debugging options such as ``CUTE_DSL_KEEP=ir`` or + ``CUTE_DSL_PRINT_IR=1`` to inspect generated IR when diagnosing whether + IKET operations were emitted. diff --git a/media/docs/pythonDSL/cute_dsl_general/types.rst b/media/docs/pythonDSL/cute_dsl_general/types.rst index 42452c4fd..cf5f2de97 100644 --- a/media/docs/pythonDSL/cute_dsl_general/types.rst +++ b/media/docs/pythonDSL/cute_dsl_general/types.rst @@ -561,22 +561,6 @@ Same as ``struct``: - ``__alignof__()`` - Returns the alignment of the union in bytes -Deprecated Types ----------------- - -ThrMma -~~~~~~ - -.. deprecated:: - ``cute.core.ThrMma`` is deprecated, use ``cute.ThrMma`` instead - -ThrCopy -~~~~~~~ - -.. deprecated:: - ``cute.core.ThrCopy`` is deprecated, use ``cute.ThrCopy`` instead - - Type Hierarchies and Relationships ----------------------------------- diff --git a/media/docs/pythonDSL/deprecation.rst b/media/docs/pythonDSL/deprecation.rst index dd6ad881d..142accb99 100644 --- a/media/docs/pythonDSL/deprecation.rst +++ b/media/docs/pythonDSL/deprecation.rst @@ -50,6 +50,13 @@ Soft Deprecations * ``alignment`` argument in ``CooperativeGroup`` constructor → Scheduled for deprecation. It was unused; no replacement is suggested. +* ``cute.AddressSpace`` + → Supported today as a quiet compatibility alias. It is scheduled for + deprecation in a later release, where accesses will start emitting a + deprecation warning. Use ``cutlass.AddressSpace`` for new code. The + lowercase address-space members such as ``gmem``, ``smem``, ``rmem``, + ``tmem``, and ``dsmem`` keep the same spelling. + Deprecated Features ------------------- diff --git a/media/images/fp16_gemm_4_iket_instrument.png b/media/images/fp16_gemm_4_iket_instrument.png new file mode 100644 index 0000000000000000000000000000000000000000..ca79f455e20eee72c210c7b57c48ebbf60e624c6 GIT binary patch literal 284361 zcmc$`cT`hbw?2$|kYnWp1r!A|3P=$Y=~X54Dpk6I5IWKc-5x-iG^vqZL+_9Tq)L_E z2?C*o-aEg=DIWEG?>Fwa|9oqp2`f8m&o$R9&nz2nB?ZZIr)f@;kdT~{mVzmhkeoLo zAt8PC`*ASyjf@8ZUQRui(zGKXIdkFY-!YQN=u0Fd*GQyc4^*9_=Y|pXep4Hj%kFO* z*jAC(q~6~R{TNzz>2|uRjcI6+7Po4Kma3XslOkp9ZL?ci_b7i0O-H=Xeffuh6pSL` z!sXVb-SK&W&7qqeQt{t*Ji;6n5u4?x{khE+t*NZXyCSL*36DrFia9$wcP7g>m7gIA z&~SDx2@el1E_Ooj9w$+}ioJep^y${y-@INL7#c2)V4C!kkKKRZCU*RO+!ld`gzl}i z6v?gBOpj9a@JFIxH$RR_6|ZH@A1A0cIlG=uvrC|^NldaIO@T%G&~OX)idI)ub$v7OR&+GW zNvBik^ZKj|-1B|jOJb4=Laj_EekQ516Gou^x?{8WsDIF|szOu2|NLZ?*t-A^0i(Ao zrBC+-R&l$G`1Iz?=It|K3q!QVK3l3Ak%z<0ZVnstLR^ShxlW^S>H*Kab`$i5sA1Kk zUhDUl)}J~$>S>Bn#qx`CL5A>Fa`y z^9PptR(c9*C>AyM^pY|@T7E}mn9%TWnbUkoc*^d0W!05--U}29+FDuy0(=Fg!_cvg zq>p`k5F1CwJwgft4^NNu7iK}ng%OOCzLnL{Lfz@=4t5SAZfl9izJUR@XkOWjZoT&< zd(~rvAdQ6MG+lNM4r!C)tGo2?=!k}6nqO1k#qziXmPPg-=ml@@rC-wd!=5!2ebwV|+<&rGUNtP>iJPv^Br*VV>YU`4rALJzd~r=du?Ziu!rH^0zVk%pApaVz5#%;0kr+FKar3zeZ< za~phz9?sd;fY>8G0k^`fO)+_HGJr%*Bg zbrFJAH%t!>bQay2jBzt%Qiy4vu^$6UxSj7)8A7pp+;SUa%X7V?)(pVRFT#m z3tbGj!~$KVzRe0vQC~NMeh3VFywuY={8Z5C9c6Bjd5S;uLsa-CO2GR~SHvCD$eDaC zG$bf9lWg+TNsI!FEAoRV^Dl%Sz(by=uI)p_b~k2{+}I6X7Ud}#L8RTbAo=78?Wy|n zyLrf(x*_gvX8W_$zkK%=_QU%&mi+v?cke>u&$f$LyzJ>2kb%)XgGWL2<@al_(xuKzoM?~Uc@WjooYUh2ziUk--L4iQk)+0*MTuCC=g=tz12+s_w^ zopLhNTO?r$(L%SOr@S{Y6+(kWrlm&O)A|AMI0qX?svzI+tPp3XU--J$=KQ)7E`D=z z0ppgR?*sPzv3&D%U4H)R0pD{eShC#J#rYX4kEZS~X-^KSCnhK7`^yz!j0%^u`v*HO z((|s+hRTG7r<@UWeE>VCaF&M6cBi4+_tteV`g2*zY4!rT7CL(gLbrao?E&x6h8gY= z2T9RdzYX;Fqcyjy39K5{!`SiJF`Yqt+sKFl%nF8QF^=J1r7kq@?C9y)NHIFYq3d5k zjZ!Ku!W9>%AopCCZ80O)pt51v9;F)O=H?ml<{DRp+-kb}OVyh0L;A}srwl!0ez`eN z^!c@ZiUya!;d1%@}|7K`)qX`M5Z96DuBWsx5edap?MM+8QOP0mg9qVkNN>Zx z5^1|&A>zU99>rt&@nfL#Ue`>)T5|{Y`8~pvksaAtTHO)dpI7!$-kz$1zeu4S&;R*@ zJ#wXRr(awoNu}t!|2u&AU;6w1`(-fgk#kamJuk$E_V0};+uCa7Gagg#>5-{<@>Spe zeVM0d=~p*gB-T=!H(RrM6ZjJE;1tDJ5fsORghymfr_3%XDJH3-aL;XAr|Du^l}yqY z{fL7DeRRSec~4&UJRHnJbcTd<0j9HtyX(y^XkOms45+Fgr zOx1#<4$}ktsiRvKDWA?mSL$Dg0s9>+y04Vc&B(++D+q~=4QkI76V#v%e{f<)T*qI? ze&V&Mw?x9szH-31sMWU9yXu>F>8I41Bi2KM^X-(CSJJtOLqj^R%AFijAKTQeec`|g zEj@)3esLC`_!>Rk9B0@1liO%f*6=%*nr|>lB(w`9hxCUt49R~8`4znAK;s)IX>zR;(D~j$f5$; zd>VYKbWi)Yoxi+ghLw8fh#u0>wY9a?yS((4ldXI%MSsHr|HRyIFw-2|g`s7rbvRWE>K;IlL5%XpVK?`h5uFe6OG^<>1_t6axG-yJ zeGuLKG$xqoP~9_UXD(lcj%lOQk$eKGT$j$Ws%JNaZxD-=Vie$TI0rP872zvY6Phu3 z%2Ti6CR9(4?#+TN=UygqZMMHf=+Y267MTndt915AKSu726oZlXPa>s-#@;1@a?DAGLGaOI< z&L1UYjy2nD9fo!!sYtGl1|HjoJ*)RrY_uB8o!Ym#Kz>!k*rzLF9BytN1C zx1ab);dK@@>Fu)%4p@O%P_TvZ(JF@4yFn&~%)sC<&DYZH)7B z<14a$Iy$S?R{!KD)o+iTeD!ob*;0J%Q+N8|A~6k`XX$DZA$sj+!YcD8ZT$YmsQ{`P zaX!dUP77m?#@K(Gqm-ST?_05JqZ<}~R@=x89UC3J1})`B7^1CloAS+f@d>f5W zvBI#4uk9C-Uk850zjVd9Y)#uFv`#4(>rE+Vt-Ch^?7BQ-5;f`X7$GEz!OnQEBf+A)1$3EACLbF->vY7H}x*`XqcSYHl zfqD}z4FWo4N!J7ho8>{ZGs;@&e-m8BciGnt-s!Na8oIKBJGthP8zqsj zy=qWX&9jbPxM%a347#!(Pp=q*RhlUOB}}<*6S}-*;bKpIHN?9L>I>aUeT9VS`BDT;Pe(z=7B!lLu7A?gg_I(Qj*wbji+*+}S_{NU zrZ*s(`BuR*J!fc9o>*Pt`T4<*Q}wmAQ-m2<&6hX>tTL-Qn(|2dTsNi}q2@hSZ1$4c z+EJu-xONk+^3B(leHG}w>|dxP@l1GM|10KibQ7$NUZnPxs-GHkH0f=;IQKOxG(k)p z(qFpD3*q;e&9vRgR5apXsOs-;|I)7s^GE&85x>?9-Ch=bY*+ig*EoveMIA7*sUEY#EySVKcXYG${d*DJ=}wy#e*0|!#2 zhoRwipB(ytrTAGENG?7EFfhb#oT*K3>s~yF8c~V|&P#sN#NGq|$5m5`P{?Qz*M+$^ zzeBuze6Y5JNStR2L6&MS_wdt*$@JGOKO{^NFR(e46i<|d>6rE9yg){uJ$keB_5K2@y>KD>Gx{^ot0}h#ws(CZ8JGs)xJ<{if$ORz&McS7-4WpD!U< zdVEJ@S05x9E1a40D$0NPv~oP2eevWklU|byX;RQxoG^#OcPGt z*5-0@vh%B;a#&i*CGuJ^7e{iR%E>G8==;1@um%_<7w+!jSX@*zKv%{KlAcdx-O`;k zHMJ!5%CmHMOm%RU8DVq_S-sXn16X2lNy(Pvyikd(Qni;0OX1Gm8{a{(&1w6$lIb8T zc^AxE9%?bCnbs@3KFS`bma6$01Jvo62r}$HjW)(m{rkj~g}4$o7T1SBU3_ zRbn6lxft|kI1~8}-O}2+FqD8=>CaAH?C|ttf&#fc7~ZRWGl}C*Jj9*c(BEHuo5b2< zQaD?BxTywGZ1PU$PF8mqiK6w&OG{#Qojpb2$-vmC+`Mtau=!oETz;u%M$?&j(F2nL z^<9+BR4S2I;{DDYLBY1V0W*z;olbb0`_>m#-Po*;q>qmOe`aD+DqL?-!Qu7#9$Fv` z9}^OtXE?P0AUI)iQguB|8EvMQSXVcKfw-?qtjUv+xn41k>gy}Hkn%ZH zYHP+WMFFKE<43oxUYITd9a~t)v#QP}-Q4Ll&pzw!9;lan_;jZ6bxmR$lOMT74;P2x z=IX#Xvh#Fk)i0e@RdLeeZzcU>OHd&?(Ua61TRkn^49xsSRk(0E5!<1*Y`|5+LG_ki zt3Iu0LQxUVnW>fzvqI_{?1wqdh>7l0uaoOKy1MDxTYCrf*I%`NO>j-2!FFr&Sm#;R z$4a_oJO2hmV{na_^9t$&djeH^pZ>9t_-19 zH+?;#?=6z4tZ1l+ns$~E(mLS>i{;{^vckOuama~@$uhq3l5&B1CdLT7W!I;LoVddd z;`*z1SXU+<1=v8bB?I(Ey+6)=I}&NpN0GN2ZUn-g{oRcngSukRK}5ny2aJ#~?Aa)B z@tLu)x775A`%^()etyAXOIqiprDYs?X+b;NoGY@cpa7Ck#~LJfH&;b8>T8+WC*<_{ zTaOd?JtG2zC{8P3gqFwpTgMpY(N983;sp;XtGIHD3V1%_Hg2jQX_RMDb1O^H+qsSz zdzR32%Zy5cu$9f#!u=sMU)14YPXzE9TJ*UR+SY=9%tj}w|(d(gVFNy-c%e$OjE!@sz3^3?;hdT;6@d6>OGbx8oz<+E=`gkKd>DYi&) z{Y-47=;F*Jl%*2RUnWULs)F8qRROF0J&|0?#+~D-58-#7*s8ldD4wVrMK$T9s}s$l zkscBpRX!W;Cj+zoiV}3X{%luawZKd*9Vt~)+f>`1Ys+8_@h-9L7TTJsm4%KG2#=yq z`}?1V+Ab~KT^+@2{RTNRmB_ z@WoId_CdVS!m673cJG9`v-nCqe0Y64KcEKR z*4x^{&d*e-YcN$Ad$fQA1)|Xhrshg(BK>U zNnT!-8saKgu57iV)A<&fXFPwBU<+b9dx4nX{QUBae2Uz{LW-cizP^Xyyyg^3txLV) z6HS*7!#_l-HyYEu8yXS{l_89csgC&lvLT_k0zzKK_M8PcEaZDSUYmuv8Nf%1pl#nU zLKDt`=#MNZFQZ~h>P}~UtlIFUrJHz8+^d&s#B1_i zh9YZ*W!XxGs)yH!Gbj)u>i1Xgc0 zv~vT6rj{attk7j|>Bh?iAopJ6)Zp1T*G=pduj%Yecv47D-7-v>=KmgIhrxXzg3{wyUH;?t1hR0q>5w>oXc zHH-Pv*twRJkN7G!H@krh3~`=d!^G3-Hvf4MYvX8J;mW)4>J&p;n;)~ONo^^Oc zsOiy}xM05RteZeCbB$6$q~+x=K&-7fR-*&zu-YKq1S0rqJgv5g4K;O&=lbG>zRX8) zP5~%^jA#*}%64`hi0wg00JYtNl~9-}WpChiHkX%|*$Y0Ss5!i)Fam|#kvI>%A4&uj zh-8pLo3R(i9o74P7+!Eq{%8?xPP<7Cv)<$+GD7^rHO9Pxez&r$MZ-z1ypUeO5)k zOFrxO_Wb++QVOJO?L75%GBCY6MP}{Its_?LA3{l4*)*{`1U=zQG@1vkYqPWj(ST2> z-##J!j93&iX-Yd#sGO#1I#~G^(EW_R3$Ej98y4nlKp_NfsET3^#g}`_E%M-6LS`ePm1SKxS`58J>AY> znOINS%f#8Su|J~#twOU~#X8Yf#cv7=Tk=y9yjjm{s47so8Jr_@)-xN9c8i!U0x5fmNTt%S2+QFGTL~iKcEZOvW7YfpN zZB7eGZVHMmpny94ZDD{|o7H%%;CdYH&Fceod&BkFUdkX*4n3F04~>m`a-8ZBCFLi1 z%;Vcl{9}I)@b%?YOw}p9Q2uxcI) z^4P;_aijV^maAZ+)$*XfuKRl({LC^y;Bn<1`2_`MXTM_kwn1DsTq>AZGYE1w7*Q(& z3mvyFVcP1cO6c@9zbm|Wn(YVPuD>kgEMd}A{(`CN{b3Z&PVwUd>syNw2YOa|3wxxp zJQk<^i%UEHzD=a7bQMAvewaILf6^HKOXTMxtvhn=m%ET0Ml+4A^|MZ$4^?Bv=+R%g zbPF067iVQqfK?BG2?o&<#lmKbrN)DC<}e_nqd5d!YPmEo#=4I)ZVHvQ<`AHg5h1 zp`tSOaa-@!h!J#R_c$Ia_wC!ab3I(_|Jfjsx_4YF+=<+XUxXEEYHE7&RaJX?yAY`6 zsHmuL)A(U9?%2J|%*@2Zn;aY*eiEz2wV})XX(}lv$;jYFMoEc@pe~%={pRXe)Q1nY zQ?0aug5@%yESirWhg>_R7`e2xbUBQ*5p)s#G?Z8rYi9C|G7M^cNMEspvf%Hdz5R}c z{3#w5m$-z4@oGml{K0i;}_p5A+NS;xjVL|1)<$@aA(bT9rL<$TQVgDt0Kv zQ7Oj2$f%CiDg^5eIh5pgmr^C6fC>P)JC~`}b!%btNSwfmn1bzT|ZJHVQHw$12A|qwU)o z>skjIK|bT>&5-Ba%g^WRutJ0Y@VZp<4II$vT8`d-_G+}=w5}o~;<`H70;Dij)N|yb zjh)@@m@h}ji!M%3wURs3bOAoMVqk^P)SIbk7&|&S$zxE1936e)C!xY97%vw+Seye6 z;;#{VzPJ&jJDm-dfiQt_pE>j(oS5R^#>30&w70eV=+Ps<^!b6p{M_9Be7K^NR09%n zqsR)h(AXM5`0b~^TI2S4Ec$X-)Y3zoaJ8IfUAGh!6<<|7-&_C)ve$FDmU?Dp<^NHdM0Qm8KFZ1#E!3n;R9|<0WD?4!b$x+PGXG=r}js zo2`%F)z$$lM(bh9R<@JE^=3Pgnxgsb^gOnW%*}}_R=?PvQP0;$;bYv#kCz~(-RYr- z9AHMEwuDB?S5==nXKd~4LhZ~P1Y)Y5+(|MrGSV%zB>)PFt??iYz0GpPug-OIAx6Mq zc4wU^^XAeIFrAyumH^iySOg)HnrRH+%XNbmIG*0q-0U$11CX7@Am1p zcjWZ!&;-eVgI?SL1N3MMBJM6OE@ig7ckVpjS)1|q_irQ&mph3bY>gw4zY{uIS}g0{ zU-qNnX6dVZMMgg2u{)1MA{Vi{S@Kal#3~BO?cc);Pcb;ltNL^Iqv-Mompk=J4FyTzZeKr_Y`xc8qK-kArQ211YrV&wGf%SPhrmgSW1& zt*Nq4HAmzdx6r6qi@iOISj23~S^v`IIBL3F;NP~hy?tPdU>*(PwLV8yN07{SoHc;MEWV+4zV2Rm`{Y(qmsUy)^$smc1&_Y&h+_gz>` zQNWV@3tzhtvv{6 zUKl9M(8IW7V>TjJOn|n@DJs(QJ%0ieH-KJL1hMiAa_aZrm5N-KMytNYiCv(e5ZM}i z1%z7%>xKYgcO)=`aQ!Z`^<2aHo=nY7a3kPdfx*iHoFX)a9vmD1U+Q5Y`t(PsBq8yt zA2X#z73Stf@LG?AFspq2bG&@{a%Xc9pYv2KR3<69d3Jt2CM=8x-Wq+dzuTLkp{NOX z#0)`f4m*`Q&ZkEoAI#D%2N0flJpcOj>%d$j0R9L$=h)t{Op*!3Zcj#m9jz@454vq8 z$%ZfV=cjlc9)PKfK~3Qti^Ju-e0gHdk8Wy!M}Ze1I>H(<%16; ztB(|b6?o_IbDjKoZw9dLt`8r+gvn?H2p4>a4+6%ug4L->BEOG zj$8C2JA}J~bkA$d%*GpyraGo*JdEirRK20WB2=V4MQ8eFEE9~Jy!AO%yB)~LqJOHXOw|00k;b6*EGWS z>bcH0ap0PJM{2w;%+JrOC@Y72(Q(D6AjI5=2eQ*$f)>C>msjxq=H&K02P5dm%BDt%aU{nIPlW;dZhK|!~*3Y5;{ zg`Jn4tj~6|9pz!~B@&Q2I1gPIfSykqHdPg0{At34D)NwHujBw@nJQ(CQfzz722B zT>;79(6+F!FmS)YVr!hs`t1Jh=DN#1P;VG4yO&x^C852fqy!{eS;7HIZIqj@9+_Ak zuNR-Fr}yL%2hbOGlbu}wriUPe0@vGy67>YKaD+&}3$X{in*+M-{~D~?;g0Sp>D`f; z1+Np!H8AEpPWIU&afgR*&>px7?Z5HTQb<4>3mY4o0B|nAI)$=ous_au0?dVM7#lO! zDlkj*Zc8sXZwwZgs+>9F3(VQtnmxhl;lqc+We$Y{@gv@xAmsoYh2=;^=?AyN19wP& zxl?vjR1|Y(OH0eZzyQ~~+41q!V)7%SIoqA#+wlJKJs!(}i*$4&NwU-gbuBF|z|79h zP9R$9kl~q&@z~lg0|~ zf(bgl=br$yzI)~Gy?eoBDp#%3x+BSEq@!>|W~iDO?kg70wV8g5F#Jy=EyAZB-Uh%0 z;M{pyfm=7yq4DY&`tIx9q-UtWv_wiLW%Nt#OJ`M*`xyR21$P78BDdbj65 zSObAPn%1%dj-n4(4-VX39%6OUV}HjR>Feu@6q@VO2B?s**Bf?Qn`(9D9#j zeSH$?;85_|8w9W-cw~B&s5sySC-e8c@;JLOADv(R<#fF zaxks#yupO~XDZo7I$lmeKZ$uxe7Ilq${?Fgl$HnA($q;|br=z03qn}yg!boF_>LMu zxIDLTIybqAXxGRW)_b?VeXc@*%bgR$_M&i;gE<+YeuicSYq|6IXGwKF$l&s9wksUYOu&Mx#9Jx3+Wo!A{KpjfCU29W~AIn z5e87{P3Tl>3~Pa!lamvK-(fc7s5Oad*F0Ou%o!DEV_Gp#b&wbqudz3= zvQf=Ymxlr6-^lrm(0@k$D{&nd#(IBRm67M~&eKkU^`N-J1J~v6BGqBuZEiL{i5YgM zxml-^LjK+xy_ydtsXq8u1Xu`q0MbO+?>TEhTAp|uM-x($-S%|VexnB?Ndl>&vJhx zx4BQlWJ^@bloE$N_9~Fl1c`In+S-uZoSZ*E#B!04m6Vjk2{YHCfI#A|j;8k#LA9bB z^jb7~vv?ftZvvA9^4=;i{<)s2s;3|V0{!J5PM$pJ(?IXJ_q6e-QA;E1EIB!K5I`4w zgM*;Q9?z8?6%>>U>{&eha{xe~dCtoQNxV!Ba|hyOfx(Pnx1ey-{T-}81>GRDriaor`;TORcJ50$93dW@4LOFogK~L3k$GR z?g1)c5h}EE-(_w3 z8-Y*=AWZc;a^HrfZ`Gk!L2?1Iy*-6mU0DG_tXF8(?d9bKG+r@Qs1z6uCdJsZ=gz5d z1_uWV3k$ohO)rlJh_<3c^;0^xcX#Op>`lh&{D9kNd~7KK!Zn~0Lm*ZI_V2`r6Cj!5 zG}pKMg$eZ=$XUSQ@7^4Zbr5O*ivzq3z*><2>H#u#*NNewp@jg5O;1n9(|_&lEe5vR z71;Y?Ykei9ud8jM^p`L5KmGigPS|A#;Gm;?&8>ML5x%QsU{}l4V6zRk77@4?wtpx(a}sX#6BJ2f_6Q zhb4a-u1~5utleq-sZq!|&E&f-y(1WYKM%#1B_4blGyUW9Du&siiEOv6o8 z9-#}lrcZBcw63IG*@{Ho#FvM;a45t4-@5I0AW=FP;)E=b%lUwhm4DRno zHM}}pwd&#M%uJwA#lYvx4;BNZ(m7RDTI#;_?HxzO(o5URmoD`I-wBvR53KWSR~i5q z`zuWxDRL2c=DkM%1H>wRw1PdKKO0OoM}V2twY7r0Jgr=V+C=GKs@qzLs!JOiv3CvM z=I7@lQ&m|Yk1`GgvNaOg-S(F2K_nCbPOk?IYK_s;V`D3wilX3ny^$EqKERoH*F|B5 zQm4fc1OfqYAMjks$cm-UXR?7TPN6)l(vusU_(7fxWY^d<^KRa{Wdii`)TvWPaMizQ z9)QQ^&zAr+X4fsdc=2LLHt&VbySa{zjv4A`8c~luU`;|c7=e&l4sfX?CnXWUgvjCE zDxoQybrHDud0!6fbPqZTS+ToN4niRyOZ*^6{D24X2JeV-!kkC!-{-7!(a4c!>cP(b z$5p6u)xd)!WnZbyZ=x)xs3osMOs> z2hc;^CS1_5iEr5!s$8ycYlzMA2JdZZh$X4{YBcbC6LXsPHXqb2iUvG_CJp`F)O~*@ z?XT=#tu0|F?&0|5%a?(8f*w~o(8p+H)u_)I38E55fOtV{52~J@U>|~m?QCs@goWvg zG?bM)z42=zL%UewB2Es8arqQ)?iW6y%K+`yOq=`?9$>w*I@)o7 z+RG};(b^YhWiT-!t~8beMrRhBJM2~SKKs#LDM=X}(Ubh>Yw8YqhdOh0w~d~m_JDsn z|He+moW8NW+hgl0>ujK4EyX(FQ5Xk3YeK$ai~{gg+?Un{5BFWP3_>W_TsoC2SpGi7 z+Vz515px4LTmuB^Pzul({efd_VnWqk9uA-S;P(#TIS`!-&t-wSDPQo8l*OMzrf6fV z))(;Yp71U(kg!BxKDIkp)H5Gv>tPootMqEPbuS(aBO2AszkhHJ z0=a9o!_#MSC#1we>$5xr5g1!wst`cI)+XcYw6cUj%v4sLyKs7>&>?eZlF$a5Xvo%& zGKS0ktHvte35qSIIe@ww0%nD?NKQ0Wv?HTi?haGXP8?ivMXF@?3PB_?G_^ zKlLn0EeBN2iINft55mK5$EVNrK9GLCv}8VF7%#JK;yL5KXk3DPJSdpPMJ1CcF+Ot8 zqR%?HvNcF8UV~QFRY^1rx3N;A@kPp@7kDgkwGzCAFgJES`*Z#)pu+1ZZSj1-YFq&7 zM{_jNVM9Yh$?{Q?v$I+w=)c|IeH+??^$h=P_j!CKyr!Ky+~3mTtD{nsh)R4bSH=?+ z5MBCdCCFFUE${TBYdJnf*iIb`XRzxP?!51}-QUxLTYI*ppYxNbM$7KfCD!lHuOO?N zP7IvOOPd~8Z!uafZq-Qd|6J9Q;Ia2a;BzE7g^@C@eW{oh~Qr+7#~(Xrie;>90Qm466S#eMv6 z`Za}&f!y5%QYY(0c`XN1W2~q|4|_SQW=3dc9&U4a^>K&)!cb6ajO$c$_u0P%frR8e zeR7+^@3pnHAlP>W!`OjZKxXK=;g^7-NNH=wwbO~N~@9PFfQZr*$Jv&OH*ALqGa&d~?J zlfDyt$3t#7Iy53n?BRwbK5&Mm5>^a6yP_VJx$__K2I-Us(VtjjDugbs714Wm3K2}s z`hSpaav2b@aiEnOn`e|nW9*Has?Hr#Y_rk+_kqEi=Q`oXDBjHc`)JUGz;<=J1q~F< zwoAjJuIH3#MC!jO#rM>xnI8+X&p4#dlJ3U_Ws?c*N2H$rPI4ki@{fQk<9wfFF(CL-+Qvd_S4MYzsOzt_-ilWU!Pp~{ihv%|8C`PM>PDGap%QP zw~+gN(nlZ$eG{erZ_E9cZ})V6sE&CRbW&Z^J!lYRc;q&ne{aQ5{QI(jTyp;;fxw*C z9`#zu@e>ND3{mSJmXoet`=Mi|!PB5}b4y%y#q@t%MyJ{~1Gns($M1V7%p8wdpU3Ze zCH?_v;oD62*=FTIms@X4WwB*O^$$94RZA||I&raB3#`{~ZtuGC4moLOD)Sd^bS)lY z_k%4L7W^aN`ZBQc!}X!K;sb&GIgYLvTmS5A&{?-=TXfZ*-I^s^_j&)ikzKg`0>034 zKS6i5$X_Q&Icbv~V}6j-%Wnqn^4wbJfiF;>z^-}z+pYP>1r=VnQ8F9Q2)b zkR(8Kl4m#A-~P@ZFYHqIA6t)Xso1s3 zma@?SHT}|17MahXuq-DaC0qUP%pFC;m1ZoYdtv-kQtS=Z=( z7WzxC#5M2+In;<<7djE^gPm->X{U|LbK;$CrJQA67tiDNg0n^37Amq%xbi)47`u(5 zQ`u&#r*+;9W#SNxvuxvmvP2a-Z?d+D+VszEU<8pS2Mqe^RNBwK6vEcNp%xE2FLyfO zc8)_@_!n1m8$+tca2{^eU=wubw-ALcb?{{|_kW10a%qS-|k>h?tY5`=keSU~K{eOXP_(3@x2WD{P4xSr{l!vfNaM z=Bdlk)c?@_KVAzj`>`%>_=#KF+vK2DKiYl42v)Ol95OTSKYjcGN)qwa?s>{hAGZ2s zEwR2S1$)vDiF=FlW;dYH6gL!I%aKSUeu62nXmd(`BeYF)n;wH!LcDqU@v->H{UDao zazP}r*QBoEd**{AP~a!!FMWOCH|S}yy}Vl#LAQJiDi>!>XdS*izxJtPhAa8F(^x0` zLfR`W*BPSF~fxn5sS)X);1JKlqEN9r)^~ zs8fi+ZU)Q2_SbiduuIN-mu^KBm*iL(PaOzfM}-sk#2yhEl)aJuVf^+#7(V9z=2GQF zE1*BPR1zrEp!5TOlRfDhNuCN8()~vBDg3TMdx7TS$Yzs!wy-{z;nVuqmUhu3Z^PL= zY~6{cC_%LO_luM^5XfE!U;NTJ-zrS@7eFm-aq-1Tbc0o#|#!yWn>iIZRI>+xgMZj zvc;$VuQSX$|D)TwAKSN8*~Tlfv?|y8uUO^LM*ZaEzu$}Gs%Xg$TjixK1?EP3O4~ti zHI0;tW~VE$a)^k3FCuv!{q5%8t6fzDo49lGg$iv%)&24P`yXAHHxE z{zgST>`F>fQL{VQ0L{sw(1{ygk?$<51Y=p z`hp5(;ljVt-BJ{?a&J5)ALHMQfB(Y=UO^u|YJK{4_Po5`%2r}WhIQnAU*^lU#1u!c z@yl-kMks~nj#p$PAI!>jE-38XcB!)6Z+68B5RVbkbhYECn~x_*Yt#NEOdo8vL-gG| zX4@3AORg^$k2-%hk53YBCR6>X;Ysj%@o zPLGuD&yFh#{<*N%1ZH@rW)Q7dqkppguc@Q?`jCdM<514D`8+mpRMy~OHCHbHh}X1% zNZ)^Z`M1H#l)QZQl*ii3Q>HNF0{W_YqJ?`+#}h`{S0_8hR$FF2C%+IrG)3ryi9_?u z(-m%EGCiL3hGm6I-nIk0E!}AQBS>7=zAABXe$Hdpi<|4sEs^B4QgW}VT8 zg<1L&TNd6=ZHu!BF<~h>ajmdqs;*1a%}$w5kxwrzyp+b9uV3ByH2&s-3>NGJGqm*1 zYx+u&P%^;7_hYNpG`-|llr(e%_u2gkMjtL4oNmhSvFG;8#`z*A2NbCI*00vXFBu?H zrL@(|^|nSAu6U=yG_>KG2CDsrem3wy#R7BIka~wO2JF|$W4)4{=q54c=Uts&c$;^E zrHyWU@V?|*mfFE*KA%PGteE+84}nvY{R*N9vjR{F#{SeiUDe(FaANPm+o=@v@>8* zgjUj4gKMr#$6BwGY6tKaej4CQY0ZNnbl{p}UsXN6E`JuPXBHS)Lm96&%;bO60;?e2 z9(g>*0=slvRqkHZ!EGmt(&@!_o*ZAs)?ON&tfJDbygcpVyjk3{P;vK~%kA0ENpGf@ zl;;dlZo3gDM493b?$^(2v%7&lEhuDl@vZ1^DPfF!;7|d^q@lxILQ*|XJe(e(?X7sebUtDtqNv+g%?8R4 zFylP+@rT3jU;Jpw{)h@Cx)cR&Qqki25x$+Dj*XPcwJsL1*$q2m^npRM!oBH@=fG+q z-x)(Qpzjm(0NXbt8n~Z>VA}U?LEJlVaf5idnl8;SacXCSMJLzeepyOWt+7ttCb9cn zA-zw-=6{0@^sF08op#&~n43j}Ba^kmQr!M>cXw=5FHfDxWVuONyk5Z`w5LEsi8E&s z&I)BHQ~mCP#nEA_sU2zU-Mhxk8RH4!r}lN8%lRXkzvGI)egBBbfv;r;Un}dcGjEwy zp8yXervg8J=P3AGAwLQT<#e=0YW@&|^~&MI5aLg`Vb3(eCBv6}tq4^kSU=y>ZC$78 zLNvtMqL3z<%TrX8;DM8?r(n-iDHFdFdC}AyDfL7RAHAEna=P9Yrka>GL+pJ$h2521K#dzjf!_(dE6^G* z6!=UN@y8$Tu!-~Q0Xijvcqtg$jl>EXw(1Ly9&fOA2j04Lnr!z(472`UQLfxoEs>Q$ z-1M?P?;ImEVj}1aE7`T<;bX54KK>xbd7Y0tb)GV(jup_fh76cJ`S;vDV($Y~^~G}t zfoAhvrj9k%j2jx^B6R;QrAKS-g8ueRafdgt;;*NLKYnyM-f~3PaPewW^+6&vTeV&q zXY0-Q7f%$;h!lgq}hh z;;;QF?q0~ZFLYtg{$RU$Vf)*H(4grVTDT;&QodHOQmPR9-7bnpos|o}ol#tJW(xX& zG`4W5Nb?V-u`5rK*YCkD{vqrg&e241WU(r1UT3r&f7X%Qs-Yed?P=B0B3L$MLOOo2 z3F`k%8$(vKSrNzf+^U6_YCthqDvPW2T6}L5opjCwRT~rM7sk;2%*LpV4UDJB==ft< zZ2MLNbqfCOaM}B~#hvN}Q_?1`=m#xkdVG7ti=2xin&o zmK{0|zg{pwb8}pa4-n8(sCe#_+LoB7e`jV_{)8)!!~TsPgFmg(Hkrne#vT@-2Zwj* z>Q`3J^vdP48zQ$$R=eY-Y+_!t+K$^3xK@dN0rtCBkVq*J<1cw^Jm=g#|8Z~7qvP4? z=OFNmT*OVN*(R+UX6BJ}Sh8k#i(Pd3x3$}`Cpc&igA2 z8<(9w)nDu`29hrqhC$^5|D42xsw+J#91 zA$XABI=H)AaCdi?;O;OX*bsEk;K5;V2=2i>xVyW%eUrVP=RNN^*Y)v(?uMD}s$Np} zx@*;{3pU497n8xEf5?!p>Z~8?@HdyZPP22{-&3?2H-5Rv6&m+Hy(uX>P))ka{A`9n z2^AXq0qeiB?&fLC^7k5z&bl^3emfn^1ph&hU*+*QN8Gok_y@0QUmg6sv`+Q%n~KX; z{AJS5XD?@HcY4AM7NLXiog8si@n4762mcSf1xNWRwA-rtL<*M&0vP1O5~oWg9hq6kMzjicb9NKvDiS zBLBr@tuN(25$ETFdx_hiX2;?GMYCLVrj5k3%s)vN{0)0z$rB;Ic@c( z*LWsvh|M5m{FA;;+CH7Ryt@_8q=q!eqz@vm^^bb?%JElx=5HpgwRrHjd8oVoZy?Zb zfXTa_h02*hW=RS^O0GWoY&tk}UnkYyUhyzjvUaFf7;n^Kxs9iImThtgni$21a{pJ`0Zv1 zDfW`I89Lfq+<%-)Sz!CuesX^K@8nEc7S;`auBbU4%_wGMN4;q%`zHnV|LQXIf^_r$ zm&p9z^xgmY+E)ZvC>{HMuK(xE_$?Is{m<3^I{ESWuh`jn-c@9lj=3?esUhE_Fcc2D zF_PZq74t5(8|XC#ggX?wm%9I9tWz&HGHuW2I(I)$nvx2@nC9_FDX~jDOrC3Fr3HuJ zWsWmeV<6OL3WQD$_Fo{BvWlx8o#bf_L!gday5su_ji=U%r={Hlk~fa3S~SziVza24JE zcDFg3Dmy999j;FEeRP81s-NrMFwMR;)7=Uo9b> zZXRHtL1V4V0>%&oBr6Wd#AxzWic#yWBs)+(LfrPyqxN83q&{(8XI3fw=XUl|V z3+I8LOHNn$l=2VaSW4~IJ35}DLc1^WTCgrksi`#($uV7xrPyL)gqc9sBApUh&Vksq z@54BWNrKH@eu+dq+9lu*+1j+1bA=)DR-ZIpY$etmE$7<|wAvhG_Mb1cKgo$FF>7cF zG;7mL=T&O5IG@2ms{62$%H;fgWtv)%IDKXy7=yJ4$D|%sx=cKeA)3QI}t;&CnL~p|+8d02pF0O+u^>z&%23X?Dx;GeqNubpHhawYJc@~fEwW;*3 z%^QweiB^Zd{URhEf=`5ATu4AS!#rNWAZ<}*oHKaOj?e49#a{qDP!pJb*)YGq;7_V% zB~NazN)Ye$0t}T*4nX5M^;2-GX+u}m{`b$pWAilM0|QUqK_$PH8OgDG;T|!RrOWZc z_k)?PHoyX7zy(78L&gu>w6;Tw+y@eT^Ewj{C!=BI^i&iL`W7abu3+%+=Np!x1}YWm z08!HW-=>e}C(Yw4n}0^RwVrK5o@a_h6QRphU~p|NOLnzoXGcf0wo;sBOEf{zlz)P`f>t{(3uY7fgm&nW*VR%owe-6{@t^+kMCxLJg zC7v}@WH<9B0E}7uq|8^pV}|bzSv~k(yrlO+M~Uw+Y_Y$e{E!e|FNZlV-F&s&!`5Ll zkm~nt43H#ej`4FcY3RHX$T90n%^yU2{L`ZT zmm;c2g_B+k^iQjTX`D!HXuiJ00&xd?qI8tcpJL!`*L289T78r!{Oh-owo~u1i<##< zHtW=9S?RG>xcukn8D|}6Ux!R5;kE;bn8tit@gTG3L3hKVe&Nta0BFDO+u}Tzb-^aY z9w+mW&aD(|rI3EX(fC&k^Y8zPLB`bhvL@NQ2^+>Nsrerv{6ZQj1t&}B+{@!86+hg{ zjzT8n+3#+f_<;>2XWTV}?LWr-=BdWR5%^+T(uH!t6;f@ybmUdR46=rzlBm_#P(ms1 z>;7-m%~a7Dis00iWls`hPgoF0qbizY8tw);*f*$`{kycWprd=^%S=H_kB@yKM0L@< z@67%kswmzILV4qM6^CQSk)gX88@p#W(p zuQQQqJGg7KZ?iaFb?TT>E@VDYG@`@4Xx&z1w>FJ>?Uw7|9JqzZo8@?8??LA%lgG=> zV8+wj@8!dqYa)2+fzmQX)zl@vb$5c5V$pOuSHYm^WS^*9eBAW>bt?j^C&jA|?JyYM zzoB#T(Mha<{HLBdYo?k!GAp6lRn3<`|C5)(w9Q$QsVD&8aFz;Ftb8|BbmYbCq!uJi zGpo&1p*&Utd)uf#R(p5b7~X&0#fLax;3-+YG>lVZ$E}2Ttw?+Jg7@OWpkHB^e@%;;2rq$;3moMm)~c5nOQwCnC0%4{=%dhPl@ z6R*2eo4c~=lab>Eg)PNrb;YK@qFuDGO|0P**hz`}2h3wrVE^-}$E~e_?10Yn%I4{e zr2VOAO2OKLO+M-%Ta$(6W{SkgEWeZ{xuP$LyAlSiQ&o?f3Z5<2)8c~#Rd!mq*XrSo z{~UqMInImuLFImp_|$o;!L}0fc!qBGlB+bT)v|?W?d3>P3H?1lp@DwSuW9^IWR=yk zCsd8-VVY~#NN+QX1&T{>nQ1WT++G-4N3R(Nlgo`9IMY?O0?H<9s;-tlL&+8)ARLKw{WYW_HBB9OD3aPu25*A9>k zJoLzFM3q0bb(v}@>kOV}=%$$Ej@bMj9Max6rdzDHflg0sq4=&I903a?ybWk3iwe`W zzE0gwDkW%5;LD{4)-S7;KsaM*;5~SoPld4B&9|A6$FTfrB?X5lJm3tNe^H!gSY~Ul zN~ez*Y2-#+ev*n>aCAXUd_Ig1_23JsSNP-`BzTK5nldxYp2AGKYX{kr9l=S)Q)H$9 zw2A-KO4)?o!)7z&_s~9mp+t9o@`{=8?m zEk_0zvamSV@6D{um49SvsM2lzGiZIU+*%0LhtMR#!c`6fZHVx=NZeL{z1^&Ejelz> zClbZvAEf$-QRT5#_@jEo+7%Ys-2&F*&)o5YKmp7}XOtyuY$H+sMMT?a#(4UaKFHnv+?r zWweBgq~>4>>?33Ry(tR%CqH)=r}pN<(Rb&4?+~I7qf=yXuf2pdsj^?0N4g-{nI;Sk z1>cz`zmLY=W|Y6|hPstjJ#J92A?R%ISml3|bC%WIa$pbv;D(mz`ouSSVY8-}T*f~{ zdD2@UA|8&iCXYJK-z4Yk@8T1>ue|gYjKshT9uGF#C_S0<@7mq%K|vb{Re6j2TXPxF z{RYZ6Q|t=j6m=3l*gA1tL;-EuxtaDir|!POC$|BAa`i0f8`}65`rd5q;zb3=_+9Qq zN^16pXv-R5Q!ZB}=qE8tnbB%WGAS}BoJF2HgYxs(dwI#iO`JBMsNJ{)(J>98cT5V1 z3wJFV%G92_=owj$hOu3%r9Vo|=L(IuayrRH`mD{V^z^wQnDf`f@ufuC87my3HA)rbrn|2JjbA5b24tXEqQM_wh9lnP#h9 ziQccl)G6I5O>xAVN#{f=UG_YI76wZ+c4gA52YD9BPywp-bu8NQa`xlUC@YHi^JDQ` zI*fU0esCg1uCovc8dvHQH45fKbSv~s92*@XECVH)f0wru8JOTeuAh~-`-;U%qc)IR zd{V*Y2f%6>9kv(4AvcjPol)@#E_b3~4z;-V3JIE zVShlNi}HATfLw@;0cIGqX>K}8-g%)uNQ}oTO`->y&~~v>U>ZTDhuP9bj~$yWl3kjb zwgT^w$LO*;{=AJi{V{OX6+{id+}YI2zUlraD> zQRdPAhF0<0SJ$G!*A~#a4rS&vyLiHaB0d`EQ1%o*NPQD8-8WKr~m=8pDaoS?o(Ytm3dPGVNS z5@da1vrMtzYc5(NPy><0h1ZlI0axc9DpIuzdg!$kXm>uw|cxYAJVOIwcW z{S%5+67DF&euI0MV~Jez?*ixxT{7LW>*EOW2+61^$?c0R1v58z4^}4PD>4EyDFLBY zP$m1ls(UXQ1^tApwWKAtsXZ1t-WVe#U{?)gJOrgMnFW8w$;Ww1F3bE6Q!T9)3U$r} zHTDH1etw8G-#t2DB+UIzs8}cXxLsC;zYEbn;HZ9hqN*1e&wT$Xx-#@ z7ikzlRHpIyfK@Rd#Byl)nA6^qz4W`Q9t@)#CSNbIT+2iHZV_Cd(-I0L;C3N(~gK1_{WINECVb*WYnA1$OfQ3 zWfZSY^%*}*`4b+ijnEaGZCfumDY)S%GC>B1oEQ za(7UEH?dZ~Ic2(X;gd(3J~xlxIfQOQ=RCPf=}*$DH@Z?|>R4T6May~E*3z!&YB^aZ@ z5X^bJRy+;O++&+3%-SFN>rSD9KBL56C$S$i)&{y{dw@B;x(PBaigb*}87h$tzMTD} z_Aw_@KfFgITJ|O#^&&lYHh3bFD&2h6P6;+awk5ndg*pdwl^sGGjWUtu{)L-mgFz#+ z>mFyQ2R*s_>p^9#@=1k7Qb6Z~S{*B{e#8vb6h3tgoFScw|P zF5ZK$Nzz#DyP6ffmo)E2&G%u_YQ3mMAQyZGc67)X?|0!vm)jr8$>U6l8@kxGX6PB{ z?+MRG$YjO#{q`j)UmLISDRzm4$Jx#Yg2JLwM1AO9LR3vtzg$_npM&-;3zEyd!JWyheX|VJGi``pIA5;DD6qDqz01W9IN3+?a$ov<`lOj8t~Vl^I74he zx^^;mUQbZpc1Fp%r$P|%-D<-_?lS_XRnk6iOS&uNv@z3>D%8%XB1&(CG+T6te+nxO zf6TXiJ@U4@jWpL1g?bUwi+;eQ_^{9ye{2C)V2vuU7U(H^Dt3Ks%M~kMO*30D|Htoo zXNyY~m=lGISbrC~B1$=GIOM0pi+(MvV+34ye@kyKZZ6QC0BSk)!~~!JR^2{sO5_5l zhFYsYY|7TBMB|E|lZa5PT4tXY{l(4Os=rS=i<>{{qI9taD(PQC|L%N<3Q!%3PEV^^ zRqZ>oTZ<3ut>N-muG2$B%Y=cd31nA0^m#IaMV8AZuYW2~hUb#iS*N0Si_oTudCu#m zw~nArcJ$+}GQk)lu|9yo^Tq4-PeHQrz|N%{UM32s&!N)dAHQ=cXy8|76&xJ8glB!q z>S3E(=Q%(hy6A^Ic95)>*b_|G$}pJeDdet<)85**5dRT;9Xik*dzn%Y%D#FCyilu; zCs6XeZCREOUr4MI#}*3EmYiQ%@kO1=I?uHAXGhm*XQPEyi9se+z)8YXSsYm;81))^ z6hT8zVJ?$CYi)~pOl5kQl06h|9>4T~NB1_gtuU(TtmnLuklnDu8gkN1gg*cFF|?I- zY2C5a!Sxb{OU4k)JfqFTNDf#AAQ|O=1=b6YgsfkFOSybDy>1C8?zz{miCY8N?Z-M; zZ$JSb0IDW*T{Uen|8)7fq_vvnP?O)QnJ)DT?|yRKtvhZLUJZ|vU~Us4Y%kWOL(oO^ zW`<#|&48fLd>dOaU5v4@OF_$T%I6uoivM#X^C))??)9Zox7UM075^WPu9<9l`_=%x zzB2xBnx!3R^RpjMf}B@~tyq2x?Ju;_HIN4C* zr+1ES133)pd$>s|<#r-(2^`KPn3eM?{o2x*$E8o4CPPSOvg$f76{)*eZ%(HrgtBm- z3zV~`%RFM`GlZ9>-PXg}V&HGfAu>(ii~^r%3C~~j3=h-trq*lu*yFr|V`y(5;X(ag zmMQ=vktAs&zlAp55YYRV6%3*Y+!-wDGojQE=XfPs=iSs+SfG9<$tvhQ>X$U{fRji- zeePdE^LT$$N8l4!iU z>WFx6VAERWG{`#+XRfpG*gM8m#=ko0AG;kaL*JG%{_ytWwk-y!Eru+u(p&4~MM4rXtzi|(MC#iUoptib zfNbp1KU)$FU5ANz%sTNhYlSNXbRg^#3wpp`Ip<9a$F<#t?7eM%xkkyVA=~?MDsYF5 z{%JZRY5YO7j>pm!^~m`~coq=0y%Kh?UD>_IGh%191R$qqmlp4?5e?By+B)ICX*9G! zB|f9LFx*v(%DnZ-|A|$s^fo*ZmRq)Q$h*PyE^62Gxb}fWaN@SG>shJ(qsr^1!gx*E z^BB}~djx9WAzkF_&Qrg?K}Ker^mFD#55X6BA!6EY5{YB)c8Xjd1^`@M3rkMH(L%(11{^tg|Nr*;v`bm zSrAhjuBdS^z>m`&v@~$2-U5yu7+{m@wS6%t*jPb(jI;Jx&;044g+59V{20^<_LnjAcLJgxGuty*Y)j6wylrR&LxKYW5i?t{HWL#AgB@Qu-^b^wQrS7OCI5ZlIf zxx6;*ti5#VMMW02!lOd#7+VdO`#S7l|G?vKrsB?{s8yzA38v5Gina@!_{>i2v`F9oJSMkWgRl8Pay|q~C@I<-|;F6}E! zO`_scuDc?DjR?-tq=-q!)QTd~oHg`%Pqp?vME57-dP0!=-g(P1d-bH<Jd3$ z@}APKQS$)pM!+jr64T)v9o;BhPWVB=$Zzn41-;njs+;YEyGehHtI-fQfc{MTLRiQw zQ6Tyb3tp6BvSX;sp4;*Er6R(7B15|i6Z`IhNI=Lbi$;BA)%1e=NEfL0Gp92b~|%)yiC94d`e>K2WZHpiFonV>_SF- zs`1rb2HEQo&7H_E0)%LzDT!mf@bIhVT1*C+aOF-4z=o|jV(tHi#wZ?UKK?IW`MG-% zlp7rY1@KB+@!)ha3r&R$5~{7LztE%`>{y;Y-fsz_l=y#Kf*sT)u>a!{KI?{cem~*^ zBgeW1V{oDuA0Q_J4Y2*e<1XWU)jVD)^&Jk$%K12hXL6k!8`D&#xPVmm8^J;P^q*va zGK({->UCrD;c>m!_Bg!V52=354p4q|Zp6F)gV@oBGYegf3|1fEW}fbHW2bg#-V9A6 zFmnp$2v`S&kK#Z+D9JS#)bS>UVGy~mQdjY!r-j65gCIZM4eL>?YX!r$Vta9i!^qI^ zWePCraph@-425Hiq}OmdCEuAxS)r;=9Pp<9?pRQB0K)eUJw*5YX;+E7whGj`!xrC)cOQi0 z$w|XcV+o?O+HgHs&|*Gls9l-*8mGVO4{Wc8s_-vEBeTZpiAzPRF_IYM<~V%@7cc$X z86JWj{wZ+|8y2aql}Cq|eOHT89pzE0h*psy&gla_^7BX+6M=2qx%yhA`X$RJ{i?$r zuF>+OhoJRnzlL8{6BlerJn=t|@St>jj9o$~B_LatVE%K;-sJ3rq7(z|bZnREnzFe| z_~Y@(H%CPicfY5&M5f+LO|3lBR+B;J^LF%}qAVbwciNFpY@kdX9XjAY?TH={w)uv0 z3Vb*sW--v#J-Z%=jE;bP*Y|O}Ls)NmvhYBO*`swSJHv`HMV}evl$-p`r6q-N-cJgY zU>)*`Ov^KS_ud2U5LKuI!>;OFv`wca{vuAg4qd)f%f(Z>*1XO z+aDLK$`=Q-zGB?KeP4qX$1g7~D&Y|oK&Eza5h;lw)z;`DNuIW?%#d%M8*RlvnY=wT zt}RPL&aKNpsp=JA6*-9`=gj8fOTXk1;4bYa zw5;!yMT7clsusO8i|;p7P2!*x$jQqG^3=sry0MEDwZ4M1-${u_lg#XE&Tm&@@g1Qz zze}Zr19uiBl6b(dt6vk8NVd^5c17(1WB*8z4~los!^&r6{Dch!<$cW)|Hkm$de51e5x~_W=II_WX)x-{wC3uFaK|+~!)IV}pxu=zND!`LsM| zkyVO?3bB{jo&D&>uUj6|(b3s;+gWDb zn6GX|b#@sZYT0;&EjU?*2PmlYaeLm~M%g2}VEkilsosE`25JM*>lJIZ-%izC0f*{w zv`Ys_oL*Dx@}-n53S7I$|C}6Hjv!3bFv* zs}q~&srSUGqRb29*2ih7Ul^UsX*T^ZHhW3ejJb`9WVyOj`g8XnS7kY2X!EhDplBIOFJa_dN{ZML{8-Q zkz1%Wl9@G-jjOcCX!0DVU7AhpF8r|P85O1|?74AfIEzp#3za&5e++b5G0!f%9*C~k z;|PD=PpsBT%#x9m78~77cBlSHk1kCpX{_Ncc`YD!uNLu0B0M47TapQ#IkT}$F-=U@ z_p4j3mYbiHG`ebCrB~m-pXl6w3&qY~Mkw;~n28T%`-IC079RWl9RbzF@Gr5b`aa8M`bWh-}#1Z|6$NVCX z$BX|so7#uBdDWwu4xIv*K4j4jB*W=w02#c2u5#yg>udCrOYkYD|mm zdbzCjiv*|F9*6@lg1VTiP%$Teq}C{xljkU_{YFreIq}V@It`bElIELIliCr+0?pQI zBnZp7Qk{^w{-0J1$@394Dubt)o3DFxO|ui_zgTq@Z^EN?TPvQ*n+_7mxG%g4=AlO9 z(U`%G^SHIg#kA_e35O1U~_qC3=6~EcyPq z90i%hDaqyRho3m@GT*Qs1x)tj&)C{>T$Y)&jyg*&_`20&x3@MMIoA6X_g>`NX!gOO zHApqWm>=B`7j7xn#DIbkmy6sZW>;cDm{qOTX~&e~wHm+aPUe%KJYpLC#J`cncqcv9 zyqxLT9D-=Su2hJW$TyS>S!}P5NI0ddB3PoRn&2VGr}+A#L6~7C%bFZ8Pe7IPjlc#} z^cwyEov#Wnx~s3@xFjc$Grs+Lgc*XE@nQ~-&YP97SFbcOq>ac0ttDb^)3lWfeg_jL zb6@7cQ2d=TAHOuvN~4smOqqSU!zt2bg-LOsTFQL5^hG@Faj&Awb(ysdtVHl|Pi;LF z5TCHjm8|?@UlI3Gb`$MAh35~>vlcO;pSa1+L=$R@4y9zFZ+-Rcxlc&}le@BrZI&P= zyro0cF^ZiUf1Yg4%Smk=PP2z6#)LykbisxYONo&0M)Bs`BfR>w4?fw_hS@bC`uR+!qr$+F zteB#Sc)`{7J_Rmr2^3bcW|FTm2}~lsGZKj@>AE3C|g9tzw-x zMWLu4cxp%!Hg5FI~S#4o9wRg zenhSkoPFSeHrJiH*)%S_Y86rJn?AcyQ3Z<0PCPHliHSy$rXwFSt*tk$R304{WI z;iJozSuybk4QQ#4N_pgBEBtU!eF<*YSyt_ymSjLisqnOjcpNobe>^4c%M#ekLA_yw z8@lt0%B^m?2Ch$g(6_oAG+mhYM+@8uJ45+^_)ZOO*aSVvb9T5haZ z^7f;tVYVtTS|FHG=O(;6_GPQTQmMX$RM&DSxsJo=`n(c&yt=9&o-N6Y{smh*CLk1U zF!3PNa8SDOx+d;J5CYBA4mIEz0esT~(e8^m@ETe>)2(O05DuLbn!buIzIySVuBfQi zS7DN?zHCf#3d_Qs$S}_>wL8r{nc1lKoQ<%jh_?(30WDP^w`WN|LG9Um6}pnml#yyAkHA*o#Hi?m|H%W-XF1#$0f2x1vl!^B=7jxk|2AiIm^<`O#tkOc;dfM}2aG~~*7wA-Eah26mLVXyv@?Kg{XYAm8 z76eke1bY7?7m8b@<1E_6k@qzv9{jJ(Tb~zeyV=0=XUI`4BWed`GI(Kx%hZ_Qmtj!= z;w|iEGNH{41~|ao{K{0I-b=evR_d&{P|V}v058iY_0vWG`S)?*SJ>_YZs^n)U?bYH z?TQ{rvTfGQW4u36@;5^ugH3wpR~L7_ZJ}(jC_Vw@-tzUGetp0uFQwbn-=!)KzWh*0 z>4x`TKhcB_DMuuO+^8jlK3^k~chyj+6%9c1eUXp$NEY@>Y}^KkB*OdKU5(8P=b2 zJX2uWuHgkr_=TqWIsb*X7Fl_jR;&iP=*wNUyE-5mw>iu*C68XK2Q~!f2h#yY46w%+ ziW_x}E~T1cDuM*f4m__J4)gpt!d6pRLYc@zQ&e74!oTE|GNe;Cc<-}fbLCf%D7*kx z^IyEHjGxftaHbu!jyEzzrlcdueyrBc%HzNmmnh26L(A)FhJbhAgqt1{OH~oW|o2OElu4|-><3CvWZH3GJjJ6BKA2GJ!?@%r$N{b@% z^~HSz8nm{OT)}xXfk5;#G|voOY<}QS&|7}^c4o~Kp~bIAFmJXG@uDcA{am*rCAr3) z`?IxWZLul;9oeGMZl5+ad%CuGbW$FIOm5oF*{u%4_Lsz^Bg~)`1R)T=fMQ^>8vAO_lbd$q&rMcn zpqs{2@@JR>=f$sSiew_ooEq@TM2#N8|0J4Ue#c$&HVe7PD)k5YCfCSY+xc9i-y?=2 zss!kbhToUy)Q?auP~F$Z3gJC?GMKxc5LD8m5Kmk&!8Y9$^ODQ9#gJO@Zk#W&C+R(V zJkwa@=;*@QdZ$p_v73kvKI~C(uv~c6FS-9tH^Y!1B(U0EtmGLNv_21InCjZYy9`7|#*YY$iV!_v-@nXTHsvkdWbi%Pq$%g`4*d%ETk{7x{HdUYayFmF7!kO= zr(GGuw|v{Fo4)!KfPW}0<6=>zt~{b9E(t5u52|<($edI<)CO^$l_s6uofBKUIY!i* zH$(cZAT%=ZEnDIAhp5yNgfD0tDbX)Hl>I_>MUD4i4HEl$wpb7^c$udShGaJlMk?H%^LB@(A19`$HUp-+DckJX|++Isg(Z? zJ>5N<{rKHrFDw%U@cNpT>*52+a<$ymyHrjKv>aDN3)eTvn$K%MWGemUO7U4Jxjt;m&MTLEU|l$b&jk zQ2%v3_NB$A8xr|xrafbY7LUfL_4$sKz(6fNaV^kn(KYq_d&B1`CK{ruEx+t}7YmI} zZkuT}l*mwMlmn%IXNbQWx%`d4Z^tIfD-)M^Bh*WM9tss_p_yWO?2Cj;rw-2x&;{NP zBl5=p+oF|8mt{==U5h=$p_Y7U$Fc2TQ@pKUlL|Viz`$~_n+4K-wuFwMVU(d(L$B35 zOZ&6KgaO}@K!|Ou5oOjH~aSIuzg@|1*`4(Me!!1b6K;sA; z{^wA3Sy2C_rq$M#Lu5eBYoq4nZi#jYq_iY+!>mL7%Fo-VG;CCtAfT z0z?aY29a+E!hXgxyti`@EsSRxLZ?ND)`iVG>b;xxYvABpI1Rbo^jqqznrYoC5NN^U zmHhlSj(&Bt+P+?$;B3pY3?LPFuJ)~Ws_yYLisUsMv}14;ENgYR^tuDSJyl%15eXO^ zqWFA6kuU*<1@+WV)AEXy>J&ZnxxC&4jOa(W0IgmPJ>Bpue3aKdLrIeL*LnWaGf7Cv zdKO+32mbZO%24_2jZbSU4xgM1*%h|5=m(h(w4X(N?&G-0yuBwerD+I0d?cL;hH1`b zXdd!#yeyyfUpD8|YmMUAeiYN4PqNTSl1egPUES;((kSGV7# z@>66;;CYo@*CXcbpTa?ff}zp$w9aYM`6K~-MV79B$)Liqtnm1ZzD-AGA5T4)ck`01 z8-swzjvxc#KK39OmWL9>URkpQ=_c7OcEVlcH!!{LK6||p{Hes#z0GE$MojhF&$n#d zD5OjL4O(@S&-Qdr)t9xt=YWH-3!>`FH>k6quJM=LZpPYg{sFVXVRxib<$t6P5;_&R zPpw$du`^C~{5oM^K5qC)5$?h%NQdsf#HS~N)dw`V6VUG|%(4}XViLm0Ul#{>zir6! z)=Ccf&KYOnz!Y)n!c+>b?DDgG4`+ckLq&9)nh2&t35PX_6K>T+$dvu)cLG}Tx9W3{|Pk{M;#5k0Ztf>$d; z5`SG_u~jb^*@Ttefm5?!qn^CM^o5`AU}i7bK^(%|Kia+dIM%vqF{W!R z%xgjmc@JY2wQbFxd)!@|`lVk2%Fj zAj%|u6FS7l>=&6U9f0GpX^p1Ia8f&c!@q07k|G`8om~Wy|6?wDrre3jZTz#TfM4U9ClM8WX;1wm+Khfy3eoP~(k1~RTw zpN>n_SsE`z(iZUZ*3m_ z@j*>rgP|Q-kd6>|dP`hMoh0w+V;6J%7Hta=Wxfc}qbZQYcS@Rjqa{i50YxPr<w=*DqF-Tk|lFHNHBANdz&}Y!5!xP#Sl_))WY0j zR0%gG$9qB`w8Z`wwNf}Phs%yOv#&zyFvz)F+vp(^PK>-O&op!UZ(vr5ZI{ep$L=M8 zlRY&pcDaW41FDhT&r-+a?FXp^D_dTdL4{tan8`y$wR|gGF6$v@mZpo{UJN_bxPj=D zSYsFEvo2pkDz5emuc@u@kcgO!DI-1!bj`oaVG1P@KL=ai$3H)$NV`BDzW7uo66HIm zfEM!lwqccAXGnfYJp0sW2~S$eNGszk5+5RY_^=R?r*Sq&Via*uoITk^WuNiz?pu0VEypMfJ>T~lF_Hz2GGwbN?V=YQ$isO z!G@bUEoV`mz5SAIRmNh5Yi*^zXZyDl7jKctt(DAwNs4Zp1Sfb%@hzGbYf`AZ*MAoe z)V@3NwDMsVFZhw3!ReV`D`bGjW63S(7OD?1M{&DSzxc?9FkkvxhQVopsD>-HK^`B~ z6zet&HyW%U`%oH&reIBri63rPGvhfWG>wKSj0LkL+#uc73KH7n2x#RTP?c*EM<+aU zeVJ=0tL9AoX2bz?X1(y37WCDcOgcFq(w#y>&9tO2l4To7aqKrD;gi?mC@(tj(3}S5 zG62RrR0ib$MCn!5KCh}^N!5qny2yQ@4g%@FJ;jIVo#YovDctfrPpvtCiBBetd4{|X zVGd)9*e=E2gfP>t<3v;?z;X%tUy;(CdVlcLP_$F4l`15zUJ^sp^mUDWf6ZRo09Hd_ z@BQfp*1t(7P-pWC#w+E2efr(Mu$F` zLF0XMt6Q4yZv?`4F1md#ZefXVS*&=TaHY0lzj#!bmdDM$UFj!fn9pVEI4?GMr)Ka% z?ks+XgFwgMd0wj)kK#5}4e5v#KyG62H&R5ICTgzU_n~H=9PKxr?Y@o4&d`dAy;vJ# zRBH+em1pMZ_C|+lYkluRTWF@6=1wlWSuA#HWBI{@fB0)y&fqsKdaY*HqxR}JEBA-U z4!V`4;PoHpa(!I~U`>PX6lSg1)-#^n=WK|A7A3hL|SRfv|!<%)EMzeRX>VTN@O|r5mFywZO9JZg6$us-*0ff zhxvJw!YN)gbw0)0CyYG6fyP1@&q1?S?_))Q_1WGb_nr) zu>R(wrpmbKXUB84azuOAi7{XI)WI6IpN)5$z^?dk75_)N$ez!GXH9tj`~8WuFs8Jk zjxePL7SzGs+?yU`9L^#I^R5~MMIrRU3GMe^dhszr8SC2xg1Ja3w%DM+v# zH<_mrs$Jj4|DyQljbT$pHFp|gA4O3bQCLl}hM}*@h~y?tOoF*-UdHgi#qu_uK}CMF z(RE-+FF%$k-!f6Qp6S!82Ob`CItGAq^yDFX2-hkE9=Yoi+=!{{F78ra6*`=@&$J`f zz1-+A5UPNM0u|%@6WD@@6b+y}j_H;0o6BYlSvAb%9;YgOtzA4`_WLlDMG}=yoTE?o z7$z@@C9g&7;YjF>-7Oq%&4WZ}i6p;v$3|?MXMA(G?FKQv_UWn+|JCK&c6oLT(2w2W zA@ry%pEh3Z99Fs2m|yYyO%x5ED36;GAc@RU?zH6+Ov_CiXbdi2rt^!TzlnHqXNr)3 zRSdPOiV|+hz5=uyLGgZ0CfE`Zq#rw7+s*}>a7D$cMoMu_@ezmnyA7m1Zj7Z%$s-Ki z*}wBG;3Ghs6AQdknIB48)oFU_9q2RIsj>1D!o5=bkVSkILqUrT7NZy)ML=HtZF>N$ypF3<~4{JcCG%PG+OR zMHlEkaxtA#3vucjroivznx^v6pvEmC(Tp_fM;4wv3d_rkx$C%t3E6os)Qxw_ z-^rf7GYUxvNqG9=z0IH2qv6tv1zM~cpmV#`M3+!~h;+S**=5VX7I&y&x%BO^C7?50 z1%Vq*1?4;>B?*-}WAH^H`LjsQh*HDY_2$O!@ca3yJw>=DxyM*EQ}rR5hB~pSu*=C1 z$jU-qQ$d0i&!dsmr;<CE_kzNHOu{f{>rQxv2Isb*fG5k_Au zoxu{G-crEDkag`&G*RqUPK`BpqaItnqwrB zf0U_WVPuCo9xx`UfoKy2mUMC~lLZ0I1H)S|(;aEwkGI*`!xGC<^1s<9$ie%VMrIDc zk)@Xf{rnGLg5ZZ|C_JxM!sIaB8%=7I1sey4z75rj#(sjZZmUMPvu=>?+Ey(FB){t_ zFJ;_%*o{4|@J~_v0Pscl)$A1YgH=8r-9@mSTBo$9W0IjSXqQM$(LCqR*V(pQ=FVTU zM-&4&a~JV9-8klC0)1%90l(JZ<#?EpVBC_^`pK;*0$qc>P(_bV#-)5Q#UDQh%cv&n zSD;cG)-mu`J^d2%4+^76AT=1|8E66aN>dvWvhzE)TcD!cbJlB6tP0?womzQ^T1mYd zwR&uYcQFP?74?}tdrN(r^KF!7dq~OcA)OdYCmo>jEU!Pr{zgf*Ic!#fH9;xvu~Fyo zJwNqpU+fzG0{ha1XgcN5N059VvMfSBEBR+ot`e{Os)L+KhHdw6iWH`Gfy!ZqVDzZS zAUE1K(#e+HZ^!eo82;kRK;>v}ZG%h8z_d8;pX+O;aOw`)_1H3k zBb}jW4|>d_@w;hGQGELrs@Weuz?#-spF-4E^_fcLcaef_sA~kD5dQP8w}C&{zXO{i z^&z8fbUSEPp3i?Q1XB0a<(0Z_q|_y%eh+d03#WltJ>Ag3;sfkI-`fACo6Gdn!)9VJ zVU^xnQIAN^{^l-y8R6fD{m+g6dG!B#)z2edfcf?F76uf%SkVSMBo9IxJX$~#3pDLuxw6^jj}pBED>8wgWGR`w{E|7!2!@-Yh7VZ_vl2g9txY)+`_`1^i8il)Yj^nSQNZd4?rA> zFj)RJq<@v&5sCz({p`JN@!)5fQmU+-!i<7Q5_;Kt!TD`S9=!GF^NAs#=N@H~7f_&0Ov zp3;ns$rwrrA2c8hp+|kv=WBFoWr?Xq#Wmy@q>}YAk4Q)cu4TvxU~k#_R$xZEhims4 ze=mG)?X=Gbc565YsAQVyDio@&a?#0GV&rv4>?>CSF$U$kxnoV|#@dF+1uqLtsZ2oS zWNYoassDiMv;tk|;9rgn(4=S&UDWhWO-RmJ)R7h84O%^(@NaluK1(ym3{r~Qwi|Yw zarv*&#Q4BTa~%^~<5!idZSu7wkq2UY0wD5ESFnF}52Kf?5+w?&r$UJ)VtX+UBQ2G&u2J=_>CAJQ2vt)pdu*;`qBT8fA@d*M%FXN^9>sX8b%NSB zM}J&Wc%md|?U`cUT-rl+3mcq99LuX=n;Zixd__-+RoXW~^cQzMJ?Q-@hhc6-MPY&3 z*tF<*tw)`_L0vjJm7z4i3*KU=r!GySy6ZOq>Y;;5Pm8L$-7TgvTS#Rw>O{o^(WLyH z4#Nwz@KPNqH99H+j{2rMGRz~rv;OaYC$xK5gdkwsFpHZ>H*@L$cW46Bt?krRuAR12 z0BH1dM|9!>^cON}evb6=Auo1<1Gcsduyi-!YM(zMW9VR|xtt^~9fk~f%jFY>3cU>a z=!|uwV0e8^0x%?QracTu!LhyvUsvczO69{g^m60W&J>l)99xpvv7owx&%W>b;&sX{%4h)I~MZyFJJs5LgfRKZ+fAbPA=SUv(eF_r% zHy5o@T1OY|UA}OEQA+Z0$S54~@Mz-ygP7idtvju?_lGFgPZTqI2E4lG>eLp$2$4Va zj@of}V)dKha;6^zfbd8JvPxqph#`k_r9d=wj23?9#%&6YDrAYtiU$j>83N2RAMNLV zVQ7#GzNckMnHmy2Rz;mA7<9v&4kt3+bJIHJV_}Er4BIq8gzOkR)92y+r>mF2YF>Qk zfB%GpSsn-LCacIpG`sUpinp=yvaKbBa_yK@xdhK8a)XA2>NXYS+toiWzg-Ob`q( zvUyT>;&3kfXKkKyIFli-A08u^a?$=5ewPNd#K*w}D(23S;O2F?+*XkQvd?ISljvms zXaeiW_XhOqh>~$ZBc%l5p!W#bgpeGz2rmT1NtSCb)EdIXa5}XUcG4qQwo3a5yv0 z1eyM;7xwWo4v?Hhxr7ERY0q8a^!Bd~dYTVq8rGbFgnK%8T}h0xMwJq}*}KN=l=685ZL26a6e zrQ9fNw)4#2+(iEXEo#&K@F?8P3Wq;_4n$r^AOzx~qPqb| z@&@4;aNvFq$rKpUsdLDM4T^dMiDN~5w(W4t)j>}E)gyTRDsW@)U!wb&9C(^}t+*1V z{=+8Ffq$h@pmdS(MJ+|nkI`z|nmQ8}+K^RKbGM79@9XD>Lso|WSV2g$m#Zl^$dnEj zz+g_*SIwgUf=En~r(DIdOFD|YlcE*9qQ&Pj9xxM|Ryp4nYK%F|A*ck-#o@7jzCHz&Obl>q zTLqy72oL~1%AR*fnBLdnc_3o|S!BQId7&@!(RwHe<~CJUVp-!qBj)2+`nvXBS*C`Z zN=ptgbrx?9hXxW~lE^NGFRi`$LjE$Kb+wOzRFb_C9Bpz@TWDZ3yN3cju!@fo^3|vc zFTy6;bwupn;w#U=P>7!0xDD>tVopa60OYw|gK$k6UUPV6U?sUF`F_kq9=0P?Sk%h2B^sjL z$K}A)e0GCWV@0sF*X_3oN+xh_ABMteG~bPC@7jT+J4H8>w6|)_M;f%s2YKdQ(y}nxSKRSpU`6*Cf6uo1EZJZPwUTilk5+r{& zF!=b_m0;9o!%jlT;KA`7lk)el5Eu~K2Y^K z;`)^-p&u%stI>1&e7*mfn<*Iw1YtnVQ_fY(ub!5X5&B_m2V~CC5R;z}6PK46+uf{3 z?UW9gq6GWknVIUW#8P%9xS$_i!2AfR!rL=# zZF0QX^55Yw=={$gTCFR@@*nA&wRPK6e4bg8uhg@}>2&01TV(SnF_G#Xwe*}gVT4C) zl5?tu5Cx$E{or6U4-|@eqD}YLvxn z!%;j2uD8Jw80${IitmL2C;Bto*PPEGE+R5mxSR5eU%RUn6GfPI?_Zyseh<-i5_@TH z)QajApeh9Uon2+M?b}azBRh^>m1m_8+v4(aOFt~#pDzlGz-aOO08Pj+T{oLBA{*c^{AkYoChv2I1IKd8I#!(sSAF)Hp%(hMokKooO$) zXLipl{YHXtoI*@4`$E5-ow9u$!N_uJFp5&Mr>?SR@^izS@T5SO=xKVppQ&kjfC)b} zvI72YOgEoKuJ+EKM<+4h-k$fuU6wENNrlWvo2YwyiL?(Js$)v|zBaHEe2Ns3=wL8j z>fG9*Sf0`=Yf|^uRh}-`wW(%mv0}U8!elO$XP=Xkr>;_2=_GO`zby}X-Px+1dnos= zuNM|nL_|y!6DL&u@jHul4vp6Ly4SI)TP~G%)afaFuGdCd3ytQ-0Vpt-nvaIipQ`#C zsvod@anuQrp$}e%KRoB4u&g{8KTh zv0@@RTl^S%p(o!ZO}lncMdIttcvehwG&eeG@^syFzg%JS`rWq;SCyZ-BQMG6NLwulg_Y`%P$ z-=w8}*PDZYCfBv&fDxX8Z0I>yOZKkgPQq0WVyQqCxoJ<2C=Y&{4TG9VmbL84G*o{nuyNh$Ivy%@{QR$XZ>+9& zhyD1b^G+A?mIPqdSa`YKL&ki*`319}*L~JVo))Y>ikq-Ss6w08ufj|H_3|I2 z)ytH<3;f52mY!k@MJT*uStcOaM{u6AG?${6J7aB7s-9?FpoE!Vn9xEaG`qH6~S&y#bK2 z*3OU%@I)>)6g`4Fx?x9CUVp#Ohs2{{_GtFGS`wd@xiKWwp5geCj}YV?@n9uo~{?^dBNLn ziOHo4vtmxVL?6yQ-ybr)?FHQ_OffESc(EOKXZ->VIE4~CZUcn}ycosD5dt`(2JHPk zPbK&RNpbI{`uOPz*T`UiO~DMg)mx{fa92c=ESR%+OzSx%<;7xUl4RDsBmRf?3=*H@ z5PsIx;`qTHrq2?0ilRr!<9y`u9+*VYWh>dl#2bY%akO)X?RMt%;gl1Fj127X*hUx= z=1Uzo=P0jYpx2rBADU^)Kc&8nD)RRwsmm?(3{+2f&vSXO><|9E`US!uI`y~Gzs4Qf zZITf?JeTf%Dzg+N7{lVrT4uybN|f#ASBif+-SE({4*4r)=*c!dLPZ;uC$KUO47G^> z{lQ8tyN4yn`&##j|AQEcl&;9_l$|0!WHFWi@M|FftuwS-(o5E-|gjMs*Q>& zl-+g!MS)YvoC=cj>kP6JSf|H#>ky;Z7poMfgrOovm4||J2T{Q`7Ouv>6-mn0No`)q zU99Nfe#_PlLHnpJ6YH;Dso8!29-u#lyNtnJwBgmMUdOq8T38ZiKHTT&!&zEae#V6~;Tx|_n>P>LN2q;TGP44j)ZX*M zE%LGSnyC6OnAXVmVk^#fnY0^0V3(aIz1J|7<&P(wFc(WL0d3MQ4x$Q2I?mF>6S7hr zbyjmmxMdeut(BcVfMg={~ z6i}l)sA|jTXQPe3|CZx>4{rZ(;AxQ6Dk1OR{;c`qfqg~q*<$6!hb={}xO~WR+}xuk z_FI!72uwz9HY}EW7yOZ$br8V^E;MFSj-1Be;e#Xfroo^+eVMO!O%7z+4?#fm5J8&| z8eZx6(mLT)?0bMT%7YY6uR;+XpZnd{t}}+CRUQI0LOn)a%XSHAj{X9Ptx)Zdye@_A7|KEr~cPGWdwj4OehO4%6>C*zWVaCB+ znRVhCzvdx7jr>yWuxRh8|hFZ{Z}H~-LU;?2h`9nx#v#C#OI0JQ^pA+LBgNG{G)n)WE9cvGhI92U{K@#e)kq(9euM5L6;VyihHT2LlNE zoa#`4!X+{ozxFmhK;6Cs8i+ctaivuAmWlzcuhx@kiZ0O9JT+aa+tTxdwY7cZ~ z2}5pq5$b|Y{-*om!&{G>?)*a+JuOiAj98`~L@N8yPk#jVYiEKxV|CrtNu0>_fW~YQ zYv9Jm?hwM@;wOZQ=MfqF=1>=sE$>sW?eZf{|Jq{50N`+^8X0^be-57WNEac2t3F-%1)Xu=>f{v6Axnu8FEntB7DJ_h z2b?mT!uWZQI|lhdKrAxY@Xfe3>*D^!-H$Q5FoyxO|1=PYC7SjQj8K97-Pk$4^U%(k zG)lJy&_b%aq!#wWF3I@Wk#fA>}sbVGI{8*JiN4R;*9}f*K{5Ul!;0CXUuGsC3<> zxGcbJ>j~$h+R6)m%{2~IG?b7J9z%?UW8>Dt&@mNo=-A}Bs%sYWbW=0~R7#YjMvj(T zCqs1l`PQ=QyH@*m!25gx>aKs1_=F+j?G(9b4`L zDm`~;#u(ecUih@g0g$otDfu|CFX^Wu;I_db;b>K{P&2>+4`5%1n-V+bNBVQ>&$^JE z+lTL^RI>kj>1PrCou9)Kr)Ejdl4*u8r#(}-&(OcF=q9klZnlPC+%sG`j1EG_;sQ%E z=cF(Io~0Q?;6-c{IuizBu`pA`2WbFr<>N&GIP~5y!~X;Zt+K+TF^}+@FU>1v$2?kX zlS6T4d1jySvq&;>$e$dsGH}0+`gABiUb676vwnaL!?c3ecvAI{!mgL4w_5!*dY>Z8w7uW+=!v>fE12$>>R}Tg98&~ z7gjw`;I&YJ5=?*ymDlyDZ;Kc*3PPYvYH8ojgv;4pK0FsoAGo|BT?mSf;Q#|na@jm0 zgCpUKaqaGPbt`2g<&{2ZHKNBbfu2#Ex$dNJ`&T)^=sVOK5iEdA;xmAqxJA76SlTRV z7##{rzlYYg=*0E6ETK?3KE2Op-}#_;r53xrP|LfZ{srQ}e`94)|Ic>%bO`YGr6blE z0V}Ow4X3F023zu_N#8B*d9lv|1t@KloVk@`1wwkUfWeyrud^Onu&6y&3`2mzCO-w> zcN*2Jlhk#rDy@`+1KL{J0DYUh&)7pFzN&?^m=IswVToW73G60wICJFbVEi;Axism( z{FGTr&gHX2V1c>n*MhY32s{$|t5d}=iCEsUHT&oe&+HTb6MD(-`=QM5jRRVm5bQef zX@idhAg$^&D_)EW4qTrMnx6rH$bDUpN-iY8r=sL&qCyIU*y8|B(;BRr%8*ci8fIxD zO$RLSNd^+oWc=1Ghf*g34&c*QEa_^(0+2xveo22+LLO4tp!fXHiaiK*!#e%jNW1)zeB8y)l5lA~FupA*@Ri1z$JsIxdi+L0s|0)5Qp-6W%5Q zC=w|ZlO&7n;UR#>S)*9H!L6o>Ll74PDNHzD|L8ahKj2kdq7q(DZH@_pRoi~7>0XoXRYrsJPP}RS%LWe! ziiT2MEpC5mx~>9_7jGvDDL})x&^L;zxm~o?Arn&~8!sUqD>y7VcX&2gNSine`%M9u zH-dOClpuD@g#zz}FFL@(h~mhkcIQy@_fXX&S42k0(1ST? z2~rzAZ1a}}Y#>Ryh)P!$3`vKZ7qFXpxLNR0Euh7WzK4WCd~S5B;c)o;6~*0C;mrx^ zmy3sLO%EmLNG;RcLpx?4af!{w+z=A?7xPz&8qvkR#Ss z#D%$MDQjPlA=f5H{~jOa!|Y{cCz#roQXh*@34ESq`P-TZZgi&rMO$jImgA`)!^sh- zKcy2%RrN$@*U;cs;_usZA*oEv)dgeBUqa{zkvW4^QCWjT9O9x4Ol^%MshNutgmJJl z+87RoMlJqnw*H7_ei{Toz#pijg6wvLdqwXKsK3KoMc=3@XdlMyc%H z$O3Q(n0RZ97%3b8%alWrqf@HvJRi^3_j2e}>>vJOuna zm$ze^$Pmg%PtyK6Dc+p$*32Ru^c^Oj%j9cl$FPP4&A-&JMxivc@Si-lZffo&T>I>4 zIEneda8QC5KREjPQXPlaXhZYYz)lfkTwz^$fCJp0d*g*+=7(e$;W8}1)5&W+A*+>X ztEdLy7d@gIu9aU21&-mcXHfE_W+n$pQ5G5D!-M-x^z*sCfzrBCn3hpg8iHVc15Ce~ z-VgXLGhRANiqQ3*cZ96-+}75b8}JzVw1qu6q=<{#Mb__BdrNZ$F#`*7<-TCGyKP*Q zwq1~027lu@olz~O_bLPckZp(%{EE|q$tK14=V3{MZ4iw_wM_#NpkL$MWn%*DKsqj$ z>EtKPcnAu=PS(#~Vl(tnu=$zwe+&Lh9>)vc0YX7$p(gQG4!d9sOUz^C7B-o;8_9wb=T61_h~<+UoGIk7GS z&-gR>AG73lIAfbCcz8V52r_YZMq#yyR@Z%^K5v1SNih<8F7p`q|I-$Zx1PlAH_R93 zeYex>*M(+6$>^g^K?a~L_ z=gyf6Wi49dv0KDCEc`Kqm1x~j zGAu|prJ$gf!GNfkhm87zfVfm@_?I?d^V4)DKgFVeJDj_Vrh|z9{*mJS$u!MZwtl@~ zcmTc!{|jb0%JKgNYuJwCNZ%EC0DXqjnI9)Jt|z6Hv_!}^SV4OT`IIlRW?Zol_(a&8 zCN&U#np6n%pp6%5Ny)lO8hGnuqK}4ii-wYrQ41vlE}~ef6{=(QKRGvlM0El4)Fs{G z7P-sxiPn7vqk49~^T}GPWOX2b9v|`R>l;}2twsg9vJk4I2;Lkq659A)J#nUnLnlr^ zzsllNgHNxkJ+dkYF1xvR8%M+FeUMoB0`-Ch$oX6&~AuzPe05R70_h7Ig1DR614u+>w~Z69BkGLSnjo&oX& zQzFr;GmW{D*a{0?lo}$Df!33;af(Zer|4xO+GOF(+v`WnV4Xzn2DgJceb~tv8#)AE z`K8TI*vV3`pEHTCAZEz&8~7EdPxQlM_yqn1fbG2##Nz^|?JtywD8DwaHE7aMCic5$ zCNYBn6y{i3tL{7A&)xfjHunhurkO|p+ZW70S4oujF^mx>lUEc+SAA)=AsKj>m%SrMN(pU-3q9e2I{?65tf*}W%E{z^^OUVN0 zf3;zynXC>_q~EXyOJk{bvl(@Yb9bePqUcEH<#ftOT50YU&l>}ASyT)aS;d0}diTuA z6lFvE6%|Vo#EAzW1QLbH(Fa25W7PRAb2TaX?n}NEsnUJ(^vB2*qa-2^)t(#`<1cB0 zVzL-s@+;u3HW^AF6wXGjAy(oR$?yC*H(}@Y175$4j^8U+th~|HUYMiR-ogZT0t{^v z?ALHvtc9i;_0siNt2vU#`ABnC782^8@0*3ZO8IpKggRlmnK{h?0=qb`PpwT=1oU3H zs+1qc)ncb!>sGTehu|z0vr|gIWT8UjMmjdC+AuI#O7q!PyHEq&rQy3SM<`l2g|}03 z=VhO=u9I4Tf%A@wWAcxYUFf}li}k9K!XQobX za{W#BZoo$k9sv1_2lwt^e?3ZYfNr)aJ1>o$D*%xHrcnEnLQ5K~Q$h`g`=C~s(7x&p zs*g#Aq(4aW6|mc9HFGNkb31(N#JuxG2o8nvbJm-cB@NpFS@xjD&*fW_j~p%vPOPzP z{6*b)AZG+C>KScgZM;`9<(+ERbaU{Po2zbZ*+M?~tqq4rd@262QN zfWpTPh#{LHS&kb$BS3y8q2oDKZ=)PTv=eBWi8KGRK^tVKbA~*@2kWhHRZxKW z;RYGQynpTs+hAXZ12!yti!xozMW)4n^F@NaHlEzo@P(t3<5@nlKrepWw|dVqr;Lls zaML5*d4o1>N=y0H8Lnv*7zOfGDCaM#%X=+!tpS(4{hujNt}J{XZr<8}%f}#VsjA(w zuJ(d)XYfv=lX?ruBb%K`{FY^=7jzq9=4sIbS~6Sjtkx4=xp%{Ax2A-HtyMgFd>01{ z19n~|Q}e5mk-l71GxPJ&k#b{_70+Kd1JdY^tkH>Ba!_CNvH`#LH(h`s6YiR4@NBR019D;?RBVoeJ>c8Ev2H5wO~K|vgW;i@tJ(b43=Aq zjAPJeKYNbe*J&6nLj4)h1@V%Hf95Khcz~Q^P}j%*x(5P}3(J1mIvNv!%j92$=FjIs zd8)UWW^3SbbL!2TFIto4X(Z?~qaQvFztj9#u1GdHo30<&{m6V0@%_{llN`^1wtb0M zIeEEzZSUf@LUr;Pv#mxM$Xf8S?|7wHdI;hgsE&JuPr-ZGXp5mFzFBsOEf+7yh zP>SR^lbIZ}p#d_|KjQFUvX_D^+`{IV^Fh%m?$=aph5eAV(AsBET7YD7!cv;rwoC`r zcu9-j0|_Xnp=-0ennoU2tG>_y8VNcM+k;G#4s2H*@R%7Z<-p&p+C|&mZE%GhxywOu z)*#CSt@gA5#@G1d2QY!NyzB4|tF=$-q&x>e{;=q&V*(n~z129kv}RdWqyaXzw))zL zw~byr8Xoa%Xoeo`SQqy6heYzeKgj)R8uayeDc8|{xl^G|LPDhM+pZH0gp6|K!nuG) zB?vqYEvrz&UBhUfzuce0_}Nxit*a|qT|b{MrJDRLj)g zx6EHMgY)5^lBh{2;(eD}^96HGoVh!kmB^gG{;w7Q6z+Vj{9#f<=i7PG-mwP&UN!CR z-h#197cCDVfD&&(-JF_r(#B;2(|~@Z;V*bVQ_r4^(*3c~w{I+Mo9xJhpZb!!NfdJk zi~#3z&oZ8TG6kELf?o5o?5=1>WClWwsC`Zte|66+y zb)pLP-U{A!G>=V<%QWye6@5h|A7ky^BJjxRuQ5mk&Lb|{?x1c2X1TuzP~ZsU?u>Cj zP#5Uhl{f29tAv+(SI^q@>N#fNkTW!J8grmtYAzL#ct(tx+sf+^H|+(6ZN}W=_d5UMWOU8Jn3Z3(H)P)UyIPR3eSgom4w|a=e zBE1rtWhP#|=iVFLg~d7Yoiw!QLik;%=q;5};oq*{oqBx*@CsXv)AQ$rxGyYy8AG*9 zaQg-RPMno07JpRxobc6DX%<<`9=dHJMQDndbJaOe+Gn(#CBlWm^WXnz+mnO=R__&? zlxex{!fqxrtpK}`a1#gKsrH$+hh#MPjwYtRnd-^=t|rOq$%@!sa=~XU2K_BDxd;C-=;0^1D=m(yK@P-=We^ zz3^@yA_p?Q4v&i|Ct@Z|yN>LZ7tPU zTylxpH^SaEbhW+J61UR;`ntd zJb)B1A;x$6Gyu;}|E1-3(`@uknQhwX8=(}nMZ2y}tM>5ZYB_#t4m5Gpmirefw!zuF z?U>Au1+~c7%q!$5yBL9MabGM1kH%N7>L(A@Jd7!u^UX|gFuw!0nIa&lEADJ~7=g*t z)@}mSw{1ktI*wf@hfKojvhboOSUkI#%?Gw)Dr!^Q7iW*7-Sl_s78OQO=h&IQ&PJvw z@66OX0q|!9vkJ$#^#~e|BPsfujmw8cudP{KVq91q({fK_#%HmD;sNjJ6x-nMe~2;= z^pXyq@Nqa!jZ#Fv-3?fjNMpc}f9J2f=zO&7-jidga~l=@?g5AR__6R7l!-7u2v<&F zP8VyD9&yovTM^Y^x^DJLv$YnlGj(>tPSf^)rseDK;&Sbkn~6Mfm%6hUg63H2OoO;I ze^+f;uDI@=!}g6Q#K1j=bE$4FAL^ezD39(IS;juVXWTcalJrM#Lwxt{U_9KpPL3H{ zC;KthJVIYj_mnV-@y;ufb>0kdhaLWe zBWQ*(_gIHTjH4G`OWJjIDpD&zS-xi8b6g#3oQVp0issy0ZLyFRohZH9?^y4DihRgs zuu_=ccHQT%iHSO&-a5!T}7iGy~L6*tp0xbN9Xj!pk;(#SHK$NT7$@QD5(YK zfX)MfHwgv5(T#p1I&TMcnr0yC;ek7$Lfs#M^Zo@%>Y@`fbD`J+9ShCs*cG3tl(NYhJxJUhy49n9C6Oru$iHn}g_7Rv@|n`b{gsKL4Csky3kb#S z*I=@KSmhozr|MPtz1R5ZWNnA{y}=lhlLZIX$#)YBzL;$YyIgr2l5JWv*w^P-g?!*T zjp=Gyjsk464CZ^z!OvV&+jc>nRT2|(n}df@LVP5k^NKO(trz5^z|va}Os-0Jk~DT+ zQt|^>vLYVOQSFmM$$w;GsC6Xw=m}|Y{bMo*L>C_`OAVm;uBU8dX?l{Z)TqKjvLJlO zI+eb~()bENWVkpg))jh$|4G;xH`wQYz>aT{e(=zOCeNrz zr`ak~!+nZQeOmCIk@mU4?x8~yrn)+&bzlq(9ME_SrKiEB34q07i&~N}Q64(RKq}RI zXH9IT^RTm%sD#=E+ce{~{rbwrO4yS9w@Ft4b$fB+_Pmy9eWc#k z1K`@4v;&ZBTacB_`(sD(bOgBMf!g@U{u2I4NlAKVR^C@FnW%seKU3_&L4R(4j5a(qzWB zR6&PnL}YiHR~;5vuBzA=fpe5wW{sAucAP9xtFryx)?Uf@xhKmlPdSPF8_#Be&J3hs zmmP&BVd9a9x2Kntd}WoMN@_}0Q|~>KJE8y1CLeyHY7#k7Zlfi|wP~)Vb}=gFbG(Ow z$w{snpZ1^LKxTah&>roWg&(uxrK;GxNyXq&RNV|Y^oBbl@lq1n#vC0b#C;AzJcp5^ zYACT__U{1>f8%Cpk}zvIV9s`4{Y?|D97#JS;@;nV>NLxy%mXk^E}!Y&&vUlU&PFib zcB4v213m%keAD+m2QJk-NPu3uVb4LuYC@|_a`4-}9MKJm8v6*JYj_X=1F+~j`!XsW zRate1mP}pIZlxNI_OF8UYt6XLbTV9|C+DPd$zeF zF*3QREOwV&0{Z?1{JtvAU5s=qt)-6I7KH<|b$;%qVf_4%F34a^GcqXwhe7Mmo|(=x zSK2U9Dfb(8Ul*7b75#Kr(Xa9M%|-+NuSe5mqLzFbjZ%wd6BxSqhwk+ihV%# zi4sUrnuqpVRA#xE;Vvf@uFD&W&4~kt9^v-f zAwo1uLXdw5>222MJRskK^Zqtmm@ z>>YVsSJT(^``z3YCoH+<2bXaOWM+%`N+!m2+|nxP(6yuP@xg=~F7JM&s4Hcxd7hYD zjm${nFJ{`4`S&uJ1@`*J`jl>|g&{0{b|<%Pk=U<^&iEhRk_u>?zsb*z-OJi;egK^e ztwmSVDj1j$z(` zjraT_eafS6h>`3Ar0pY68$%T>GF2aPIcB<&wQhgQRpT}Y7&AUK1{p$-l1T99^_V>M znv0*FpxTuhQ=R*Cc(C!q_^o+%?BCx)0FuY^<LXr~T{(k6!0w_?Blxm_VP@lES5n3>P7h^1YEHRtDDAQ&Cq)-!7 zfJDg^{sbvOK!~hKCCRo_yJq$D#^|#Uhz>5fm@x6u7rro>%6!06llo^Z)!RX=ckVm; z^Ku@->E8SsK4&>}v-!*6GCu8unkiXz_zfJV^8G^Dsk-&x)-}(`$Cx&H2?D2Ltl7t! z_WH{OLd5ZPqx%h~?=0@&{`>$vL(XiBhR?XjASNqaq(@fEPCM)ScynXPE0NBb``I`4 zL|-+vmz2LQDo~ zRm1k7W95mA1XSOj&M0M+PX8bs(})Exasn6_Wq%vsI2gV zH{OU%%Zl#X`!Lwp)RNRpg<16kHH&8Hee^PsVeRrMew}EqRUMblZ2Krx*}gUDfaOYY zgYxkRwdJalr$on97OR&|f^OQ4RsV5ou_Ij=+~|SKHyz_53X4|C#2 zb&LG%;UQBNyxWx)^%}VR?d&jnX0%6^mTsld*^kO)4;=mU@RkcbIUerYzTt66hoZR* zJ<2HxJ>G>siOB}A>DWq)N$0)VRm^In?$;fu)86PwyW~z<~Y5K=Kx08ywqXTr1U&M!FSyoALWN- z$D?d5@7{#Hr<=@ynGL@6$6wZb-_oqs>$wXb%-Bq{Xl>=^)2>;5xlmnRp8U0Iv-8pr z`1bCS?C$#QU_+Fpx{7T7qnbH;AGY#u8_^DtVgqyFGn^G>u|vuv&fY5WPNhK?sA($_ z<~2+Am|s-rZ3ml-F;@JmZ}?~H(jAMYG_31)x4-SD{1|-FcVvi?Lu?Ijlvm!=Ls=mHqrC6)M&xYAD#3KM=d%r{pMN~-xhby#V6yz zdB+_;4*LOg5L}`4s|-Eo!6GyKL92$E?}dx<%9DHT_8*xjqr`(-t+zJpFli3{HMy4z zN``O=lZKkawfEO?;j$zx@<+8aV>8315R!9xBqZzI31`d`nfXj+W2Xy`mN@R0gV`)z zGxgF$sd*#S`TN3ic2>gz>!;V!O7(h-*%;c{n5y#P?xd!oAh0c_Xy&n)( zQBupWNn2}iH@~#68#&xv!-XH6|B{YSt>J|d`9VHrFb6tnDusK+VUCTAd%rMLLMpuk zL2YF(y4$dNuBz_5!Ftw?$7(y(awTx9YtQ& zXIyhoNSPyI-60Z|bRFmud7y~xUQ=?VZLyhr@!g(Ir!*`KoYZxGhPl=@b0874#;+bu z)*BvjJn!2r*pe}maq2Y3ES!x@$?NwVizr!w9f$Ne@^;5T8^rFU-B#blzF;9QkiW;a zy?jwGQwbD>5wVKoJvz&DBZ0TGzc=|)?7O#%o)dD|I>%f9?Zoj=^o6(CeU-UyV@`#X z=xFsD#GkTWZ4Dbnl@B%0amsD*wKjWDUzCdKo;!xO!+&G{^`{E6KxeV>ZCLmse;D*Z!Cg!+Yxh0i8vcnda=Dv_#wU>5OAJXFHRCtJw~aB9~dS7WjXH zzl2Jtj_)t8fYM8)sx@(_Q)MI8I>fe4tGQEdw^6K|$*E{MvR!DuSx+F}NCrq3#@n|Gx&77k$X zb;t+4cllp=d+?UdH&V;-IO;~2e;qfrMMIH`b^(dl)}H(L%}S5Z^S}V#|HIi^Kt=U^ zVZ)fHsECxJh>jxCNH-`$cXxM(bSo+Y2uL?5-8D1}s7RM|cQZ5$J@6eA{nh{d*1O)b z6o$EV&ffd%{p@F-d+#Ti40C2t;U@kk%j$Rr372Tz_^F^<;%hG?a`eNG0}%Qa7MU)o z7i|}p72Hf+vm3UX+!NeV#^dB3bjOiBKPNux=%nG9r8sKWhj(8^)@s3q|KfbU+@Od} z>wzrcp5HrE2a{h7r>gb5v7tAH+D+z|CkC;rI0SoBMfZ*nrjGeTqe++LBQpss2SOt4 zZG_Y!oS9>+M|vK$&W>dqk7+O+l}`)=aSJ4659RRLySvI(HbZ~;ooHx`$>#bi za<*5m(-2JFl?UF{)NI@4!dw*0K2x&tEf!o(wMcD9B9F3sux>kt+PL91dFBU&81aB? zrsX3^?T~kaGI34Ed?{Vqt0*a@M*QB=gG2t^71Tm}vZACRdaPVlmpcOk!{ik^(t(iB z{AFIl5xPQtoQQy(Zp`)M*lR788%|31uYz5G^Mo4!F_B4=_=jy2fLjLp~LQDj;{?RF^pjI z41KH|7a6ghdYSL9h@&utF=)gUUUT>KAI?_Rzq`EGr}QF4W|YMu{qx8ka;z35me9sy zdUd_GHInGggyL3RfCyB^nOeQZ&A zCyDm*ek?%zOtm(Dqp{RD54AXdrX4R}}2j3Hqa6DpelO_D=q7FOdLDkxq-u1h& zi>#KVj~yHmz%$T?5xcYBzJ*Rqj74zq7Tqx{EVaK46Q6uK~N)Vu|@%q?Gm07WH>=*s2Sb_v#D%M$s** zd$-dR46$}lAEJ7>k1~fSE@JMk67w8tHIqfK&|sK{=M>V{NalHadAal|B|) z>SobcY8h!$g=bQ{QXby}no2au>SmjAhL%1P?zwke7e1j*D?bdb0c)NCst=2vf>D`ely z$JZFjcfTr0fS0UM51wm^F`-8w3ho}#WfxS_bSPf`LFr-}qoS;ATRwIEhrPTUUcGGs z0Remds!0$|5{F$`yH18;F#b^L?b?mKvO}6Z%=ys@36uWSYcV;whi@aq@bP()$%o8s z1$Qj1?XMUiZ{y+~qL5?|{`WQ>>TMw#tX+hvN3O_<{>4TqgWle4EA0t-)v)-^1#^j| zvEt6*BW?%Fq|%JWek6jEt9aF)!MN@1$7=sB3E9K(ruvg|$>p9sp55fUZU3>SM2g?= zY6L9d@L`k9WYdedOU!|Vt{BXDEu$f}jLw8@LdJRc=yH8N~0 zXYR=qXlRcjYKcnQ!lz)>&x(_C6Qtv-2Zp<+6S?^gV%N8M+$NFDeqCR-y7lPG$`_Yp zyVd09rk3oQh*KgA(n9q$H&>aaj?6}wV3%$qAAX05yW(K%ovAH>Xo|qJbV4+K7=QeOc}z_YbEhM-$GUNKGAW&9eQDtK zQRN)NAEXZ7%X`YqHNKV0DJv*TD=d!Z_92m}j^h`j;i-FSgt&KpJf7$pw>RyLg+ zD%ag4RZ~kyBOe=oUElSQR_a6jZ7>C$6>`(5wzsd8%~VHEi7Hu0v8YMMk9D`y*0%6L z>GO<^_VZPn2@hgL{j(lTckz!-(M5@jj#SZ>=asn9FtLXW+APG>d+<8fj!pZHB4c}M zgX&~*JC9J>l6n1vgoJ`n@QT9Sp>|T>w+;%^I?fx_MMKy?g$+|c2qXQ*<0lw$4h6e; z=RU1yy}js3WNm!|qgeax#UIbyha>tbMGtusL=%XUl#R3fObse80_*)-Gm7h5l^Aqh zh=eB@*$X9H`r9V((I~Fu0e6ADgYW;|`c7hwKp~pCGSE=foY}CUtLh{0asOyPu9=S{ z&_EQjn9#QIa9v}K|6bx<-zrx&RkkGi)Ur}EVVB{#f1ChAs*~W*xn-;W zwMUHG-L&PFG%V3(iVqSu(wrW-5&ZKUpNjYOJG2{mM~%s12}R+22UP#u&TG1*YV&@2 zsl!Dk?tgC23eID}QzRl5Jr>g68Lav5Ff+}#Y?rHoeJQpwsvu@lY^5JS!qkJo`~vi89RP&DYkOGn}fqldLhO6ixnLfnVk0*AK5QD#^JQ z+)ZoA+J^VvYkXgQOh&AjhwNpdZ|ta6QL>^d|`zBBt;eYhje2HoP| ze`9N_gwTS^wpI8u)8ph0@Nl!gx1!ss*a=bu&h|zu9tTP!v>gipk7&g!zm9jSlRLEL znZy5UKR13EA*73G;3Q>H*UR-J*0w~l|0eAiTYYygKh&P2qh!deJ)o!Qe;muaDoGf$ z<6fb-`k=ocu4X#*k+r%;8Jbuc zEORHL^e;&Zxffv!iKYVoz0^i4fvL3d{t<^@gM+cAYTQ-DdLC}=h5DJK2$mv)hPuIj z**(Um&o+M7PpbBZD0M4()K_+J&^PewDg5f>lgiaTGjX*3N3l~*W#u0haOD%CJxE-_ zI^}w|ahNp!D`Cgi+#0CtjN2PK+eTAX(MksKpX{!Wx6nsK@zgQ~QQKm78W$($n>TNkNtt6-096L*`pk#f5bfmTWB_O5wmCZ~9j(0f`Sw-zPFJeR! z|8d8PD#{uL*;kr)rj~gzJ`DnA{Fht)Ig$!gatDaG&nYODT4-yEBJ>OmC1Y8oq3-*$ z5nsOWy6k_r%W(a2ZN2}!m}kM$4S^j|3}lZV*A_HHF&LUK61fOlSganPkP)<+@nK=8 zSF1BpMwZ4GRtp#XJ1=<1=5_Y*$xTc6N5qV?k!=I|lIm zCbERq)>hYr&))B`Npky<Z!h9BwYC;T+-J4oRUW9ZtU*X++#(?vhiV^fG>}9^MfGdAtZXlKb2zMyq)8_1 zZ_WEPCi44*UL)|+-_H%OPEGL&=5#pr~q_+2(KGUTSJ;hKb{7c@NB+T2vBuwAScAd=+Q`JC8FlNg z$VcW@uJWAhcSJtd9T^#kUe-VBjb|S(DFse}&q|_5H2dYsI~|R>_-Xz!DRPyh3kRkb zy!R>=2;x-^1e?a3OVl364A`2O=VzTfpN_gra^~VaFB`8e!bMiiNa`CJ z@>EM{7#JAX*`0S+N3UMJssc=~*M4=xsXtep!LXeZI3(Ei@}O+SIKpL@#C`8;M@Prl zH(+|f!NJ$AU0Ys+PuKf%&;l;eqGIvsDW`aZMq^?>s4^ce+1}nRD=T~Ic;jpgpHO@B zJK~DPEyUjO=|y+>zpJk7J1w8bablis>+ziyiOZWFQ^=Aly&pRYeeR0Co9fp*JLiQR zr^x+W>pSseh|LEl?9>jzAnW7Wr``d$LAPHSm#1vLY2&`dJsa)$mC1r$Im<%QY1wOm zkdrBXi++qv%DO1%<{7&Dl48`i2;UhQ#OB`Pg|60|99J+Hxoj_dZY3WEmh=7lcV)N% z`Rs+u*Odx%_??{g3v?TxPuSU2$Q>LU`WJ0Y`?7^4C25~M%U7am2qYoI@vGCZlugOB z%fzTWMju)Zmr#?Dd2cJ94K`i%g}ijS%Q~-mhd8MZ#3oEqGsP8VsY&Wqf-wCkkV5l3 zLgRQq;_fRkh3?Kx+~%d;b6(IHI75pE)sRfWsB8;8yOX_CI0twJemWyteHFcF3Jfo5 z8)8#IUzsF$IFk7$l@mM>@Mk0JbL}<+v}M>;9iF`jFEFl3JLFS{H#W)YNI4v~z0@m8 zzBXLC=8I1WJS%8W=pY4zl#~?rsTkdz>+EqB-(fgoV`G=S9#K*vb@=>0e2`NyT?aQ( z|7^UJ5z^L>9xm3gO7CFqVL`&I&<1CIaxBFKXWz!C=H(U=RUPH#74pDg#o7@*=uAW+ zF=_-|I;CENza|q-xa&CB&tJN*9pOzuzlm&>--4sU4IX>2%?*Gl>QFzOcD|j)dNOh!scPD>gwuYd&HE_qKbFi&nCmCwjos; zE7&-`IwuEtR?tuYbAh~aoA&BTXz&A#QRocsB96X%T3-Nm#hl2vd)?er0v7X$bZ~NL z#}z~q-->(k^4?qf{uP#FswVVg!OiZjW(L$;u;~$Id8T>^9%(N5kgaq9e}(t>(CT}> zl*QrehX(IQr&8_lgUHy&?1mx_`NvGc(pRx6nffX#ReA)Q0}&C$ttM~{xkHDN$^wK- zV`W;+vr_p0R83}tL!f@S>C?-BAGQ}9tEOh=*|Io$S7l#XEKlH6WLV7y^2eC#$l9AZ z_Wu$p#(k#lZRH$=N~2DQqN1Y9YWB;fPdS;HneiN13*$8>{hA&*-HVBdiLS1$q2KcJ z^Cu=IR)$N-NJ-nGnOZcQfMbKuJ8v7EPUh>E`P{z22h9SqKw!iR`uOn{&Q)w|!XWhF zuArbGrD~1>a;@fs2C_0#T-g}$8%(Oith~Q~?Njs$TSno68rnD0SF}w(j&s<%?DC8P zzWt^VEo4|^#~0?^j1Gy6V50CBH~YN$wNqm>msZ?^c}pFq{Ha!c<FVOZ5aAkB zdcf1UbzU6}MB|eg=OUHuu&XQ>Uzx631I``jS1Q?_xtZClee6|X^7vrydwIGVG&uy^ z+~B}M;@_@5O#=y9;Rv{L=ST@W={w-Y8qqJ=dOj|OUr3j?vRZd4hhO92sR6PZ+VS=C zligItyg{T#E0-bXg^4|ea4CfWk227LMx%*k4M5~p%+mmOrSuUGgGu;7AgR!GXL&G= z!{IU}ro-BpIyCfz#bU%rvVhBVFaN970x)@V@|rV!{;-cUOT8xgu5zf2Z}!3bAU|p5 zhfQ2^{Y41k}DsS5XN_9yt{72vQD*Y(eLz30cqEX+P- z)r_gpqkiB@rqr0^hhDBY7@y!NdY$*a*T8LKpt8fT${bfUU-ct@c85XL{#1d3xejvS z?i@`nqzek}wEEf1mAAwb5R6>T+foDv+l$jm7Mk*1s`!eUZpE7P!h+9b2w34F7~+hY0iXJ#5q{ciE}p{$@YWn1#~C${(wrABGehFQ$&HyQf#U% zA0)GdjgF^|a)zf6BfP|Qq_1kdpRdsieI3*fEQZb zt{kUg%~{Tg_#&5lyB8xNUc4XN{c_rVgqO2&cnb7o$dc)QwGLOtSsa9?(I=S;3BeFn zLq*2>`}tggT>zhOapELc z0D1*siy}sjqSrJJvC~YMcXV_L)sEmYdz&q^*mvoZojF9NMFf9!0po|jeF}~t{(CrU zqe&r#E;xfKg?U{DMFZKyT;C3nQ_v?#Q^_)0R{I4;=JOu*(fF&K5~L>nBGi!F%$ij* zLRArBRc61zYS-jtCCtElx?Ow!+wedEG6>6M|NHo?Vo)_X8p!R5rSUHYS!M+KpGQ<4 zL7tk16{=`w;NcXDbld)AB&MPXy+({R=$)8e6V%q$E2_1T@-2{TFcJCde)jfm7%6{6 z;D;Q*GkLv6d}I?XFNm0+Z2{X_>sx@b6V~i5Y9= zE39kJaLDQKnM7J=!$n94*CgAiPLlL2lNbje6_o-ORz^>@TBRQ z)%ZAbombftWziZ=n_CQ&In2Yn<8N_Z^A!P*z5jaO&v(73Bw~J_k0iM+XK!xE2A4~-$`4k8;zLeqWqTM zy{4&FDK@l_|J-C?lYM`cz^qUqjtPl0qT`+3l5Q1E2`{XanYXpVQpi%1AcEiHG7C$0 zHm=;SJXn1;W*M3(U|Hj2|JzQvY1(i8v5>jQI@jjVR~p){%gcCuvlQFuXEo)f(`XCb z(SzT^)#AsX{w%|}zaqY8pWq6}^8eX9Tf7JV(5lq&O$uR_>@o6D^MIT9)rv6wqKQ2F;DPQ|p88-U=B0_oz9i1GVU5w)XN$2Di6;@>zLcY`f09xjS z*Iv1b!8A2Rgkf^}sh6p*j%gA8g+KC^D+11k@uNs0G-17WANUM5j$0T;_IDng`^0=o z4aI5X(UuShUKC{g^SywY+67+o(>*HcE6NmDz+KD#Y`j>!|K?aRhHlBr9o;f)0I%^= znp?i*iN)cvV2RYcEft_;ho=YIVx3+%H(v9(n{ge*EA?!b6wsH$-a(n&pu>WwQ&|AR6O)p-xVU1{^FT=A1o_6+d?egRyajTNXg|rKJwFE-#We)*A?as*Zx974^f}#o6&+GFRu(Ek5%yEFy6t@T!3Cs@sDwmG#6?WZn!PDMk(VY(aNt!> z?zIC-tz>7=Fo9aNq_yBBfNE~Lqb&BzO8x@f&$tq$m&f;B1GWEZD(^#1Z9bYBNbR%{ zrB8w@FuC>V(Q~1CWbfX+tJ@Gj%xo5^kfmN~F}l>7wFdCY33|;4;Tz6l0l9AJmuZcGlUWqfs2nuo}`FK zb7Nz=wIF;jHY!S4MMWjy!u9Ldm&a>rV0!xcgILrAEgRDfrwJf~td+#M@87ZO>&W?_*2a9P;4_wU5mfu{$ddbHZ)>}H}Y&uw|zcY|3nQ;8& z0``EW(?12}8J?Ju+9ANte+d&4aERRK?QU$IBP&P8GMNWBIGv?;2k^+|V4A42H|BlVES(z4L)}8oAB{2?Bkc zYy^4$2I>_!DWyVoB*;TIXIlU&q>ot!=gaI~1l*wpEp1*HrQ$Idp0%~LY^s>_O>;2k ze$70Pc?7L0y8Y(KD0N$Ol25=~V+cny? z?{C~={&1gdvt-x;Ia*mNk|JYZa045gmYm!PnNw;}xl&SA#;TtOW&o@ay*thW@@;*B z0dSo8NSU;7f;z&v_^VbV*j!}X9qMBYSl5#w(>(LBDyObw;Q)lkW)m4G6M)kUr%J?8 zoKig+K;nSloo|iO#kK@71LO|TOs3(TMn@A~*wC5jK%SkA4V=?dViY*U9Vb20B7?T* zqa#;RQKeJn`v1UG#s_O=?E$%lU7R|1;qr)e191AwQeB1~2a9B0U+GNysPt*Zo&$dK zKC4vIC;4m*1Un#R0h@e}$Gr;n#~p%4db%b*;2rOsuNvP8Amw z6MOQz;>as|kIup(ZGp1X{X}G^Gbf@PS;Mu-$8S72B)M7vs_qnaR%z80o&io|!PnnK z`4bWUlD?;Czf89wf3nN&s%Qj_I^3Xb7|{FfVDcmT8yXsTwptlPd=33A$|m1`=JS64 z_U*GQl^o_rj~yx$Mmvo9AphHDu zvJnBn`7yDvR1lDc0pkL80~Qo^t%Z`y9;5O+;t0n;mSV>aaDyVfsU7<$FoisYtS~P$ zu&kq_qvbF_5diNZCncpNnLx~95gW{r)dQ$S!}eHd1s1b`hY-H^*9+FT4O;IsY!8aZ zF#EuC+0B4&fg`uOaxd7L0uR@4I%Zl4(#{d z8r9mU^x?)L?iZ`^5$=zR__kr6S-7&S_qT4sx!3k26Gn8I?%kv1Jv%HS{Zd+oFH9gQ zicb664VIC7{g!axg`Sd;Nyc;T`_Wb}bS4SOdHrI_=g(=rjT$_=d*Xp%F2*URKKR>sB`-aCWth* zgQ}`H_1Tpg#G9$g=fT7UO?t_%k!9Q!7Y{F)s-}L=_HdQhu;H zos0r;e*xS(ca#(qQnL*}v4pL-7!mxFp#n_cu;{ik0yZ-<18W2x=(I4fJLa+$q3K?} z0sIiKf`Ng7q3GWwU`NMGH!JnyRdUE_n`m1lu-BNsHNUhJ9Ug8#cPBIlsC;zpr@CCYB+C&T6%r`l)TdCZDW+1n`@`23ODedHnFnG z4~C2Y<(*b>09H{gGU~i%S#>(}R^ShSS5(O<0WbFc)p%oWd!WX9O~l`QUau@Ag9$Y} z=)al3S)&fEf={1`=pjH{%)q^Ka2!UP=1R-R@k4VI#+{KlKmuYSBJ?qBy>Kl+1OjWP zg|viI`vn9jT$4=TMi_TL)n{X6WhEs|sdpoG+XC7Eo(B*=s8NS;@pw*U==lp5_>)VN zEkQ*#bYBw$;%4Uuqdgh2fCdQ#;fz!w7b0@xR(Y;YP>fp(pl)9**u z=@fKCsv=$jUk;+QDTlp>$BD}t!l5~o9PEFz2fW_VdTn80At)Fn|268QGM()Zm4t~} z4#ko#FxGg=u$-2Iv=Ri zOu?f#`l=K(Vo4##D;2QDjK?0_g;bQfm9X@gJVO{$zU-MkZu|mmfjq~KtL5_`49zK} z_nPSH>VhTvj0)qVJUruCblYp=m^95TpNn%C!wCX_Dd+f) zRn^{zXU%Ky6kcIRvAfK9($Lv-)sEZ4D!N)&26&Q1Yl1v#mQo`^O-2s`R^qK{C$64s zUo|2bsP>MjbVz%kNctg?Z8_ZObKK!u{vBes4otVa!&mmNzI6orUJKx9;dWv_T2%TH zGR0;0`YW-HpqLXAy}K)#kI1>Z1ue*fcc*W(h2>e5>4^`vdk*RbVUT77cbV(gg;7(@ zLwLm7s}LF29la#mYrh2l>9vZ1jiUbXT89@a`kL=JxcqFRB`~NDJ&+2cUlsTlPCpK#3j=LKp;X*SLxkz5;4r;P=KNylQMgFcr=6# zPEAeqW+_q<_s-7}al4km%!W(aG&X@%E9B+79ooAHD=UM$da7dpupi)`~v3^KX zc=p}zG`fY{`3JH3R32O@)BR}OB5TmGH+i0sD5BeVIHIk_WUHAlTS0WbEQ_nkYXv2W zDz9-p{_&pnfr3IWk+?zg8s&J|#=x;1q^r2&4vut24_c7xS>jMNz||*A%M*RHVBt7& zEzeM>`jIq_wW2C@crVFV@f7Obres<=6OC6rS{9mZpewKM(`JB=qlm3O9^^q76uQ!K zZ;bqx;62~JvIc325hy?b$?O9zx72)?*b3z_d$?#4)anpsdVB&uHj6eTJ2lIm2KXg=Q_yCuESyzUr}*T~w}x z3eHuZ`{~xGd^vf8Cpou+-0qD9-enwAM<&@@mqwd9HV-=KU`pTwC5|N3l^a6}-B4N6 zYloInj;&6;xiKYD(^K(ooQsZ|OYa?iPp(NYxv~L=3oI2Br4qq_T|t5_q78EM)1?C| z#r+gn7cg($=XTvo6^{YcXEspKhrwVTXr$&Td42jHL>Ejw(=VDffR*gG-F0)rpZ^3bP-V93OZ%cK}p_u@*UxF z4tcw39Jkv6(DiB52dOleuhcv|9=)0J&9>zsPsNFVR|J`rjI1nJ%;ThCX1`;sgFqkv zDn00tl>}L;*wr;)uyH&d)ri!}%ELna77fSAcOBe^VoqED2518E1Yqf(DM4O8vOh<) z)MeKq<+#jxdKIJ=^@mPpF_~9AbVOPA-RB9hLRgv8L8hDuPPz7+PRg6Cm^9elWxJBK z#F_S?!E6HkuVxf*6>5rN5$F+L5Ni2Y8XFr6V(Z6`A15ZXW}hY{ zfjDya`g@Df3a}7_oD{?-fu!e*j725Y*RZf8l1ZfZE`h{( zi9>}d94vMKE07!m_)G5NZ*sEPg9Tg{$kjM!aH}VZZ56&yXMxoF_BTa z!6l{oH+E_9gPJ#W8k)Tw++EfKB`sVISO^sZX*KR=<2gHC8avuE(RFHKIeCy1I&$1f zLJ?*%c$9VUA$IiDx_R&L{z?lO7X0Se#t5^h{H$W$hji+a0(kX0_M*+$s)6k6$yi)6 zm4k1&#Q^nIwn9tv$%woSGaq?&W%uPOxwCk`l$FxXZ92cQG7419PGP$|uu*jzV3>;G z>sPLn6>UoaiKM5e50wBLUSi82sCqvc>Uq3|nuh^S)i#l4XZaW67Fk@q9aT)gPDwp3RFS(Y2>18#&RJfh6WS@ZdUa7>#?ZYG|z9jrniW8&Uz zCan>*u%#&r=lOUcJRHfm8jkgCnC6z`D!L$1C(SSg)$qdveDq`Rx)Qs(kAdQZU&UT8QQ&K zGR_C3Xu0NbWb_{9c=(RO?tEZ z^Xr5~fZS@Zi;_x|tgRQ0|c z77uC z&9#dQYf=~8p^QFeI_vE*kIOFcypZ|{_7N9pl$bMf(~h%#-tO4NATJI{URGY-OZfcx zuXXxqHIb6H_@#WiaFV?e&!4B*L$NknBdpLnCA)2@rBUFf_wwb!N|54&&+k?eYk z)9<$$2Jg1Ia<6reF?_rw%i|%Qi~JbOS6Mn!agA=!9Dupu-MmIS>Obb-$&E@oVPF># zl7?x+8iz+W2FQ4zP{OedwWikc!Y2jy(`d-~ze##*oVRO77U$jfegASdfrm%BKu5Ys zUh-&sq7+xUCD$AuE254BJkohq9n;}f=ZIgf#$lZ$Zct+5_1fE=eTN>1`6w6fY@wXs zC;V(Q_FR5O&$hqJrHDL?W23kdqutXWNf2WOLN{Oa8w*x$1tDahd@uP zSdJZL#Yv9vLx&%fd%-S-kV;1lKR{GE4I_bIi&n;kp(UB*(|f;rg@K_}wWp@|*Pj?# z*P(v?Faq_>D5rT%0Vq}8BZyb+UCoF0*E2k}?~hW7-fpq3GWu}2DkGOSy!Yg=atJ0A zo-$B7Mg!4NmT()WK<4mAxNU?49<^xEDkNy%c`6+=sg*4i0}J1zE5z1o(Ri4TUa$Bww^S4mfVZ9n}b*`bsNfNcs+EV8mlpZx;3tu z4_7{3-tEgdwz}Mc`TLxn+5eLAqXfYsJoLN+57%+MW(aN5)kK9nGKi(1t^rc2@<7c$ zJ@DcYwtCrHm{QJabj?cpGR5yxVLWsh_?Kix%&DP_`YUWpd=j^KHW9;iB~yFndqw*) zJ$mxSvn5K~9EqVZaj^4f#8R!i<{k|W&hxvPXI>1PR_{XY-|sirpQ&*`Z%>VLahb0S zww{dK-7ZVZQKh}pMueCoIoY~o_qsyf(0%IIqbypok4jSzT6qGNZy|qF>0))3wvpn` zk8VZ^()<49$$X#_<<8bmad^W?hzhkgWnJ$w3<|?}3MB|?3 zm*H;XtjsJ`yQeMI4JB12)1AYzRLatL5E;Gw9zK*ufoXHI%v%_q46L!&%4i;r&0qW3 z_g#xGa^Cg<_pPgY$^i`Fa@e$Cum!m&FU2!g40=-jt*1jX1j(K}^JA5|JomhqJ)r!@ zchHnJ>c7Dp;65ObaLSk(8BIA?ot73G)^u)%N#c&M*4V=Y4iEcS zRzs>KQ_kBr9N`sz-ImlyH3DAw@uzKmq$*xu{fowG0( zEBug_Gx@o#MM5{_&)!81>Sal1E=e~4Tt=e1q5V5}$ z5J9i;^lzt*HK^%{`LFN4JbCkcP8N@iHzg3ow1BZP5rt?B6)a%=5P^Ob0jjoQG6c zw6EpGy^=iAR{TUxb!nurt>B-0lR`IxJT& z74?8>KWGo#@=CdnoRPIn5idOcI72PD3^D455cq@?9^3JHLG5^l4Fc@^^qN@0rN-w; z_0^x3RMM&`p6p|9SMA%+zgpTdRCL#$|FV)r1au+E$vZf)mEAH_ft+IuvjWG0KOIsQ;>AoMX z9fQEsG^-L-|0U}$jy;hC*_4Ry5WLFA2+LwFhaQAKj4^wsYACPN5&JMPcYg=Y`hI+T zcbr(AHp4PXmqYix`K(JEEtB(ZfYO}OA&1X_+Rss+{&0FC9c@~mtDEC=?4g>Vv2mb8x&LtEj@ z?qS~JM0#|QGtwrcFDTugqr>_=U$)FiZYd2mloQCKy)ehOaT3QnBrgs;si99I`POrV z&&flKyQ^;fdQROxr}U>U)I|xEpXackA71(jo7OlEIcS=dxmd2%6vbALRk%j$Uv05k z3dXNLxaNi)ml%_P!p%(6lnEU$qjAL(h-E#HcW6h}b9?YmhT~&l>IAj}Lm0%Cq#`bOnp}9(}TmBXK<0M;uW#dVC%yPTp?z%lYdLji^;(7w*I5PE|aaxq*J=qLhqK&eY@fbT^vc z%J=lK>|u;Z7V5aaUzg{gJ~qeG4nInbrcIGqM&@9TSLwoE`eqBUtPLD_U(l2&q)P>% zgLe{mXiqf?rXHM5(>Cl|ork&9qhAdmsbdxN)n95cYkF}_Yz5hOcaFSa1nlkYj>L0A z0MAxd9N!J9RsV{6=zPovg=c5cS3`bwBa+iZ%#Sgf8RG4Ls-g(l6hM-YvQgLw#P?=m zy-oR&5IBAtRuNb+zH*fjw=*+1)SU*r0e+|a_Yg3iR!euQsj4J!&pFXTyfwLvey5KO@6|6MiaZTu?=ZU!&)EocDEIQ{@k^gv z5k3dm6i}-rT)oMB8>hvHUgJjiq|dX=i@hd|cmdGbTq(o)Gi(sz-PquMNVvAPm!tua z5xQHQKXk6s>aN5>U`EGpP5^`Mxj$1Ieps!FM341ctHOn!rn!C^EZVotaX~)K2Q%JOQB{r1BOe*_z9EG3oTXpSkdtnI2q|rx45JwRS`T$!72iUE)nK zf8(1tPT#nqa%v<0=0tG2BfHdy6tZTa5F0^y&I0^jSNu5{W=KPB2fc)fA)MYFe2`Tv zEm#tzdcL7=Fvy*JEN!VS4CVS7j)xV2=L5TAuJ!FmI-uiM_Y>B(e@!+jQdQx9dS%IkKkU-~K|2jtMbUT&7eTU8N8P5ne` zC#~7427?Op+c{c29B`wQ0ETPXtAo=gw0x>pf7)w^ghv@-)FV+&a~Fis)wNNX%)Bh- zHU%NJWoOmYXE+7d9ok~YnaUi7?eHOn9~vjk)jdQD)_V>){$uf|K7Gwd)Ns5MlCebx z;#@ji*MqeyaL}sKz%N@pJj0fsgUs_GsOOPmi*3bjnSJnjuFIhxI2oIOIv;+0SG`=H1>$It4JKsf1tmh=bt;`=oPO zZTF^P>hZ@LtCs(8DhXgH+6&ODAamLn4H{;irVLVRLZyK0lt+oQ%oKMjbq*mWj<x-7qjY)(kYZRvM-38aqL0^! zkc#tMw?0+@IMMtx6#jkpEm8L~Hq>!D?93NL8gVnB=fhH?qNulg9`kT44n_^0!ciOj z?~MIvh9ffCkvoC)hdZd_N;RIOUXN$~I35=)V=R>IM-VN1dKl8RFv=0G@{8E| z(0*@maQJWynz!>#1s4(e^;lx_FRAlt2*uC&Hcvvmj`?5q63!hN>}@@2@&K#oeHsZ@ zhI_Y6DD1xcG^Qs2y?$&sl6DZ!`DqB2h(_*%EaJZ9OlaNdF3zh46J`q&^+#atbx>@&Q{MUEjEO+E%)?lw$ zDqpS$hc9B6j#nRMb}NB8AxO{w(!lLW2sonam;ZJ3#;Q*uweYmX@pwu{sh9~qi(B;BQjJAXB#*DzR1 z5@h(95+n<$DLrYYY&MH}jX=CsM?_KMKIQ)YYq^cyu%5)8A)mV8+YLSU`;VsS=za~w zDGlRtJII2xuFqAMQ1r{+QC=zS2UlW~?Z;_yjG}C&X?2Xck!KD*gWi0E6m-bgv{$u; z-(6a&yB! z9fL?wL&K8RfNI3+-2(>IUzrTA+PUlB3e(KR{~Gt-KCyZK|D+gy+wafE|Hmw*{OesO z?aW90#TcH~8@pD)^L>hz(jt(G1NADHXnIprcMCC#gs7W(m3)i^6Rt=SQBp*Cp9Xr{ z=CD6Cxsa}w;UGn!yI!YSPYYRCkRs*ymhGuP$i9Hhmn$Ivem%qxhQ5&7SI5eVSBV!O zZg%G?Jy@KLRSr3*D8yMZmebv)UyIa4ZEu>n4qfuox7`xEOB=?p*(*s)i!j&lZNU!c z$Rvvh8%NNg?4A>aGi(~mGei|AT)6fyd~;T*Az!SwyQRgQgGct#T!fF6NUxkv)nft* z{{9{X9>hr5Sk&od48UdNjN@gLpB~RYF(Z>K+KYWCUVXoTtmQgdP%eEVDn;zghM6H4 z=}LcJO40VG9bO{!Jy+v?fAWxzfwx3`9)>>@wVgA`SvVyL>z|#3VC>rz-Z8z+%TjA} ztmH%IHX#?ELj6wI)GCj`tLwp6s)^PQXi7UU2$opseE+vTw?sppv?#6@Fm}@0T)HoM z=Pfa=s~ASP^hTP#evM{wjAgqWKVYp2i#SY?*R#c2!mjA%T$smGic#RGi=6VMB8*Vv zEH`rvBmSzC6?9EJLuGEE_T`*7X|IQ-`jKDxj%AM`cS+Sz?z`-8QV4QBGb>&#mAqd+ zUqMQXDnc%FwEELjHUC{&W!ap0UzsJT8yQ}y^kMaJt|flARc7{ZAx_SNY^j{Yc)3$ypCQayx6$7nEXZswvS{xjZr3wJXeUNLk_)s z4Xqs2arhyFND8AMRI-6RVA4#*yiMiy0xuNnnaJSptItZ?cnBxzU0* z&X{ym-748eqRJ#g^=dh-Cpe)a*G=Dl`d-1JIDQ~C^}$Q?%UO(H#lP>}@7_qeH7#4y zr*YUE{pduAOYx`D&THoCnyZMbi=_2@U3pz#Ke?5z9^w4-k>sQB+b^hIFdKbOyhM2H zK3Ru7R{_D$4OtWRWtna=Z?ow6AwO#7IqSMjF&d&Cv~en~@Cfu*aL94dq!nIk&) zy__M}NOrC5|39wYGODdEXdA9mC|+FK;O=fM1h*30-Q69EySo>6cP|db-J!TUAvoXZ z{XFkyzj9VqR@OeV_t-Vp%+7vX>q}|(CC7*r^OJ=i%J$2gZiXCTCd`gk)@f-mjd3&b z%6q+lq!ak9~3^i#pNBkgUGWCTZ-!ZI?&-3~&~D;h~J zc6Com_uLu^>iWXz9}MlTkzN@rhoJ^7k)M67I#e~Zz+eYz&ysg!SNRFT^O#GBePGnA z&(RGMqF{+{p-%pKFH3$BV5*th5pnYD?U&ICFE7K#KeDEca1#}2KTsuK6~qyzDocuK#=CG!V>2mjY0Hh}fu)ZDurn_62v?~DoV9ENAkm{u{tuBvj3 z8#Fg{U)cpQ4wKT;#ID9X8>hZqewb~{tKngx)qhk8Iaq;+4BS6Y@3?kIUy6hb-jyF9 z3_Gmbsc4q%AL2>ZFF0PIaiP+LvRsXw&6^)0PR4btJ;;dg2>x?5uKB&FJxpW*nYi5B zj(ltBZO9~67-)+fya_mV0<##V$dPU{Adqzan44=t)dU{@wv)gtI`1|aXb*NomE*AL z2G$H@m|CjLLgzkC0B#d=l)?{_PW!{z%D#kl?47gtBBT9CHa3{XK*f=M5PlQL=WTlR zROqggzr(#e%jQwfZ#R;(86`ivYpvsMp)d zCYZfRsP-^lXljO3BXIcQTzC|M1$BK;KE(04YFp8L0+T6I``_Ha)>eO!JJdP4EgP4- zZM|RhwvvR@G$_?+K&WWz0;62{yseOqk10T7#f~0^p)JZKVVLWzLWfAO6)wgu$;Jtb z=DOHg^DP=j5B-v7S_-rS0f>g)ru!-X*x?DnbV1aGC%17;$?gzdm=TLTB&c zy?)%mmf=!L!L2T-R(w;hUPlu|p6Y)6)D{rd2{Qhe4bHfoi}t~_WyHz%xnXL5)q>2{ zpfS0v@V2LZ(}L_VfKPc!*tcy+RNVD_S+BlhDy`In5opBHlo%q({zX5*Q)*)t^7cLu?r~v z092)YGM@aRu6BWI$HwT%6GJK;$}YGzMbKh)LJ>s91EO5`82UP3q^)1k(1w#au<$t& znDF`^#${kvmJ;9w&jK~|!n2O6rWwpI$~}%kDgu0g2pP5axRS-hM>TkoLy?Xo31uDK z=TehNhyKk*#WcrW4P+5#4sZUfQ6;dLk6}>xNd8&W<+zd8$sZ1K7s>#U90x!J!?jY#s~Cjz=-uCq5XWn1hVJi%j>?m1me=j=f2PBstq^@ufbo};TlUd(U|vs`T? zy6l)s*gVy^Sb(fiMGJxro*W`}rb-GV*JGYypGo>gej)qW>Qk9dGWKsdl=SGdQ#-j! zSeAe9T$;QLAZ+PVvaJ{LgaGTA!PAK1pePo~iP#KpW1I#fG!vOZlBB$5n6-f$Cs#4* z2|5;vwfBx&Vc-`Qx=j$i`K-(HQbEJWNZN4UEV6LH?}m=9SsRbSH5Vr)G^$t-r~hV_ zss;*2ksof^pdx>(?*>eJGzUn%W3sIo%1=SSE z>E)jCKm`=qfW!()s2cf5tN(Y90WD^42YGl$v0;4nL&KR274sJX(~{mxcWfqJNdzFwqIsM4l1(kWXTJC-#q zbtxah?LstEg?MS6j^8dCbVwC7S;Lc59qL9@qV^@+6bv^Ot8Sr5MzX%=w*iCDcUAgd z==K<{1|&i|*fVfI0h~cU9;dnZI=PQs^I7h)7tVDHo7yNnn4=BqdTtiYhB0T&Jn%RY zaXAtlJ=|q=3e`0%v6XoYDlza+l1(v`c{qbm&xC5oJNEAngp{}oZl&tvzy{o%>B0BJ ze7M@2Lsnwo1QdpiVW6#^etGYRHc6RH>pIe@MNUJ@-~^%gA-+M0HC>$=8fa=qGpK?$ zT}|&^9sj7pp*KWI)61OTtx+QLCYKvM3Da=XeXRO81Hx6L=EISUR5ro`7gfPPk;mq?8uBU!Sq>An4dVj*2+cYzGQ|@_K z8Zy==&J6!~RPwca;~Pgih-Z+=T=68uneD{&Ok<71^D}Lf^oy(Dbf?Wf6h@B9)W0VGq+9a6bF5fQV!56t zjoypmx+7>qb%3Lw4_k+m{M^WSde^zdYDpFWmme~c1kY7&r>lq4;MTVnlDwsU<_4R* zQoG?71+xuTj5)S{$qMh_sSQ7R{!GEqao|;uf>Yyds5W)^!K4JdM)d=1)%LeM4u38F zzB~Lk-nd=fr4|{r^V9yft%C9m60NDKNNee{_Qq`8CoG6qOne|zJ+*Qvb4>(ytm?DS zWa96?pG3&H%v9-1NkQ`@NJx7Y7acQzo_`t+8GKRVKucjyDz%=kto|Ic1P+^?3+=BB z6VJZEho!Ha8^IV*fQVw6tt^uS^pJFf6A9mbX0>o`G|&h`jbUV9XWwe7&fHj>c*yLR zR&TXhsszJ*_~sC1+3j%@?R}3D5gma?&+e#$Ai|H3e$~o?lZn)}ylgX;Sdm4nKq}lh zHpEkR@ji#JYmULf`gw+Zk>Mhr`-p-XLc8eI;{dvBn?R2;-nyT{Sf_CeTB=y(a>+*` zUOjZL_Y3x@tm7zKM_))e2BC_xOPqX+IJa*pn!4+G@j`Z3c;MYKiZ$7!u$wVhU zblCNyb4(n;5l9+*_+C*J#c{PY@dIf;i8~ zg0MBK=b2UsF(afl-GKU+Vk5RLwB@Pqh1o-;U-YTU^)=b`sD`TP<%SOF(m4yixaUeM zw=KJkA{X<3?-&4kG90jjI4*+s1njP9$F?R!Ox|3Fl`6j^Fbkqyxj1`q(hkP3{KQ_i zKIOSQNow8b*^p|F0MTx`&{tosSfzrAau(`vcg6=J9)*Qao?DGizzjCx&7KU%faSGM z7L1;;rq}ktL@0WgBF|<9__N;8os$+SoM1eii_|c#mo!X;6m9R>2Y(&EqJzAA6cG4G zc!L^asxMBh-SsYpLAv8j*!Xp-$IwCh;6f#dvS{_^}~+ zzbnA;aBx0X--eq6t+5;mcUL z)b7-tQOT*Qn9NpIkqKZ^43SD1d}g6x!_}i|FT=q28Twus{)79%dgDoT*iFaztk?t{ z&T-%Lco@Sk%FC{KueOuwPoArysxmBwP$aH$hF5agR;UEvVMpUO1DuhE`hobmKhRbNHTrJmz+T+?fqm~J5e_;d;j1%Xu<89Q0>Fp z$?rpi7wxmyrq^gO(91vRO3jQZ!-+QOU&OO3NEoO?SK78-tA<|s#_ZHlwc-DSEM^G7 z-D}F$Q_AJT);^F?v$)A{jK|jv;&4%-p}NGKTOKO_pB?xppJdK9-czD$Y$X}O(B zr5)>hcM#naC_O%Bm&PgdAjBEp7Sfawi8p zS`d!cZy3nc#KAFf9e;@$-B!Oe3v}Cr+v>xkxxUQKI=_Sf<^ZG$f9yc>Is|94Tn#ie z#)u{d*q$NXMVM075;~;b+fjKLsWs@)ywgjWnmBBRSdKJcV*-KN=Fpne{ghV(+`-X8 zBJHIo);UI`{p+Mfb0hFu9m%$5eMWFNv9u$7BYBBffpjySwg8g1Wu$W38wVuK3d<6U5SZ2 zJ5d&lbGVl4T~Y&h3D@bI!_`mIF2rG?7~bl85?^l)bcwIj^! zWcuRhVcj{CWmL*GBw||XcqoG8Js;evo5!~Rh;8g;k--(fBy4VzlY&|0k-GS-$SAL5 zd{{2Q)km(^0x?t47YK^q`yp-vEfzqAqOX|~8#36`;UC<_P%o*u0l5F=`A-`V*dkB9 zhcANM zr^J_u7U_@(e2I|32ea1gNL0D;a?y}rhS_RRQr`Q_?f66Yc4HTk=99ox%}!1*c=yxr zt%RIYfV-EsuL!y&JSv)Ys;t00-5|L(QLnL@bjS7ul`AMOdarM;;mG;nWmy>EV4M;Q z8%cv&U$5Wzh8cc~3p)t8Nz>ax!aDz;{s^6o3mpoRA0iD7qYWPf>_Gdmd>$9h%B(Cn zQIaM$D=f|coxCO_B_S2DI$lc&0pX|pVG4N0^;S_8;ZAd)CGt5fIP`F?7u#g%&c*wP zrv`(&1L_Pmpw5h$zJK%MpDB>}X9|91E=NW#+(X+IGXGGJz(i+M*^0fBfBcShKvyJj zhL>l|92wyhWt`r>BR4~*_><+Lee?yib+-dHP|@tQk-Dui7E2itW&E5_rfp{EKr6i6 z0Vhi>*9TPr`t~Orc6$9)AciD~Q^S`#*0d3+F~j1t(%0b{?Egg_wM&Rf7mLMQzYVgf zJ_n1bTbDgw{ke?`AMATk%JUiK4!mh)WBVZe-is>P)oob;w8^PEWGGuYUPT}Zc3Wn( zhWTjBBVYFHw_g3i3Y8Thq3hXRly8@hz^n%a#Zb#GSXI}+0HG`NQ zu=SAiBEYvBrCfz29U1^z^UOJhh56sjr6wr_!ET(N4#nz7X2U*7JjGo9r`ZIO2$OJ? z6w=_SVv_#o9x;>7fwtzheHwXOzCtY*snrKwZ*xqeQO0!jL$7Kf4HJ@pmGMTT%0To4 zp#is27{#q5Z zwqCG3#DyG77%wsS4bQ`M>2X0+QcITk7J=_ELSY9CGtW0XpEb%cB{F||poNso9t#QW z?s*lnRK+^}QbS3$_;gRrGOOIMJhjAF8Jzmq?o4F=;<^8>1Ke*f5#cSLC`mn==2HSQ z+$DH}W=ozKoHv#oloXx?fqS!?RsA!jr{*}MhFmMz6!G|2cIJ(&okT*9oyp4}ABNKS zVns1Y&KNPV)NS!iU`Vk2e*d-Z?vL;6cb8=R(A+jK|b|dW4y_r^5r|0b-eNS%E z5nYcLm%qnU;OzfbYq_!E;FdJ?HBH@4TENng;HM`AX83L?eUGKOKQBsAO*8$NN3rQ) zamN2YZ563QmAh*TlK^Pbpv0{Zu~K53v!tHt*~+%{#;B>{uPzax$aj!m; zX2(a6eyigkIFa_g{(d78A|}sPt=1V$BhDRC0@m!7lO$3$GV&lvHof-v%og%_M|Ner zUbV7pcC;Qf##WvR%~hLV=}Ma4H>JqHvOOjkAg)`cY6d3rpT{msCTw3I~gY2b_Xa0nGd+; zMhe&Nq-x~KzUMx(@(wL;iP9+zE3Gv8qj4=V1i)^FpG_W>M8Vt3OO{ItR?T@DEPZ0& zr=E(TJSCDP_K9EdErQ@)JT8!(SEp9zTsC-T&^Bv^aEgv0lug9)F&cS9r@peB@Hy`| zf?ey@NrE^c4;9ULUgzE-dowifL+553Mr__%wtRmG<<#FOm@fp4d)K_xsqpD&x(-UN zn_Yi7X;={?)70+*KT=)c5Feq3tH(dVA@&2hpbO@i{o%Ikp(l2T6x7hl6{N&N=0pzt zm%>(S@vQz}xeh)XrKy_JewuhusP#%%J>SqH%R?1A1WL)*o?=k6xrVX_Iq%R(tMLEja~c`?-12&)7ru9&Br8> zPbaoRx&a%i!jNO_6hW}r^ePY+^NcQdJrQxy4lo?|hPuWn1CEh+N64dKep?bf>mSAx zRd+%0)*5PNhUjtk+rcLUH%oOCC|79k&Lu>4YI58jXBBm3Nsxap59Hk!?VKd<ik8KR2^RrYdVtygOj90r@+F_$uHAD|r1s-{R}yIDMIr$s+|=;`${_SNPUlj#$N? z!N_8hiLbP&4vlaEnz4i-yp@h9R22V`?O0V3Vgaq4l*f>KrZz*pUSJs=yb<|KhTA5` z+UBo~!)vi61Nc6s%^Fo9NYnqcjcE~7Fh93Gkg>Z<L1;<~p*%iYy&EaC z)5IQB)?dnCOK<~)vgW5VCM0wO;Xe>r%~V!#c{9eJDb#S#*>s^xtfOg>hLnY zmaqj&mYwdY&l=a&j({;t;cL{CI{Uwa>9Uq9Ze&Vd+G3e%fUgnJ(MtDl6XgQgJVDaW z8`5hg=0AfeDF?ju^7ZPHW>o+;vF1<#e-)9Whslh7 z+0NEACpQ0-aTc%l4hv>qUTp37J(TrsmZpf!b9418!-eqD*YUi`pVv=?%djjf=>(O> zT-8*wcBp;*$@-$G*M@C z0Db99?oDiggf-`|IQU6c=4QDN%YZiGN+qdSMZ&8LDCBI$;e0aZ>^bFw9-Vt6oyeDL zM>suSqxKN7YeTGBTVxss`_=0g`Sa0ZHWP4YmhA`R*71ct$JcQD!E;@s@~7-F)L?Ot zh*0i1*A;#2C`y-?1ZR`1k^W0iXH49GjG`c;asNeAUY&(kAGN_{sFa&K{ zEImprVjILUFl^TDAEvUVdWdhz%V*oKFw!|pOh5b{T4<37)UvPuT$mFJJp6cH=)Gnk zL=caL(72V}>qd5JU6w2kFe}7%J$g@gt=x4f|Ep;`U z=T793OjYJDD+n#XQ-5&hje*erqi5&3Q4Mq7dcf*$%6v#bYwYLc#UyN(zrKbq#CN>O zc=uL6Auo;-jOQ0ZVr10saFB8>!Ztr4UiZ zxOf6nWYm}fiLCn@8B|vcNxuQtxU?@o5S+*$t8XfW+YFexE`#kFjQ{lR#E211x@`UE z?RLh*YYw2!FLhyjtX(`kitjSG&(%TN(2Bq!1;)T}dP@a$8Aw9Zz|3NpLbgJ+!HOtA zZ`>2>pY8|>u5-ih>EC_BU1&|6;+lQWF$^D=bAQ{7Sjx0ujLd7$)P#2SkK08^ay-_9si`UdaJV$JK0*QlGOrE}_u#>*x$j^M#o({Hks-k@L5WPrPpVq#3i{~oRBfqb&0N15fDLUCf3I}ihHtZ&n}m--@4${$(+$@AK7VFLx;mtBJj3Xc1Ey+ zqy_L`qyqH&a>x=U0kW!24$k4**d`0QY5G7M0~AW3H^#Qz3O#@CKglDNdHv09P}&Y8 zB)wA5QSRh^@4ul>+~O)9a+h;?`-Y1BOVIS(fw{J3=~^r46vlr0gUaKn07N@E*#CZE zkAb!h(Urk`n+>!2Lm}fTj<#v&MT$Xg_rMdlcV7hO1jm(QrV5O)vPTXIH1v1f?XreH z-6sjGn+~s`z%#|YGrkDUzkPFuBY8l0#7rzH9_1)55y|=xGFU}F=f;$w&7G%HzX9_J z@*9)xK1Sh9N3C&{pzAc(1q4^FtUz5e^`anIMig_Po${_84~TTE2SX%7gDLHLMMK8>hFIVak;!II!qu2~bJLjizYgyk zdHv=1P$(@^0}0M2CEK&6{T?cK>IABpuTUbcNoRCib!O#O_joPYrw2tDh_4M1LwG=s zXMGbc`YjQ|09_}CG`=SKoekd{uinx8i<7oVWm<3(B*QJ!P$iII5>3wJ>^H2hWZGD^ zV9DF3J25Y=2(MZ)ITW0N`{(&T1Pd*#7sydp$LnJRfkGPSvV-*+`2hF^uq6wpAkbWi z;zJeUut|Z(uM-a4J+Ey!T0^$5*ognChLzRvf+F-=)e4-r&By;$4I>M2(>DWUS0dmB zKUEB<(7GQXU!(C_9Av6-8oR7nph2QVPi@mq%1^P@bY}<*(2n~ z-2(H!yh8NWKdaV+I!i+av3|9Rr*6F`RAZKfCOCVH_!QPTofK9+^7gBBY2#kST^2(d_<+<%70~#~BK^?tC+- z-pze)!#mUEs6oRTZhNYzmBn~kNuUC zkcAL9_4Bge$;e*?YTEZ<&Dk)%ZUPpTveg20i&7UI+q%Wd|5!Fx)#?Ef?Ye7J7%v7& zCbHfywea48@qh{iwSn^pZ_FFeqSThufZLkUH5be6fs8c63Iy7OGdp^*PCD>K0NDv} zbdO2_V4Ye;kuF{8R18S|1Gfenhzo?0c@$S17ViLoA>+CP*r(BL^V3Irsa}u2{l9)h zuXvkQYbfp9t|8jm!kO>(DLMrj3&HVaCh0kE+Iu($9y~%~G1$c1PNfbSRKcXpAP-#S zo3y*Nc^=2Bvfdwcd_~}G6H%3rKZxT>&!(z%EJbN7tZ7nUVeK?4&Q_N}2_SOC;Yd7G zUdmWcU2{J4Ro7po93tFs^xcqZ-P=LY;FZ#_taP!lt)u#-E{7oy4zL9R-YY$+j-itj zrs!=;I)d}GZhivbzcW_{pXRFU)|KgpBkzXA`-9Y+O6ku=IJktyw>C^@&{a=b*P)6+ zHM*=aT}yynhb_A7*CMQh=P*+1mD7EH$4y3>^|+wKHgEBA$Vx_?wbI1q6jpaO`hzqd zfW#9_*0gUamaq_rp=O1+r9|4hp6&9p=u&Sbc1_KY$lCKlRPt)Tr;pCSn?S>NHhcH@LRACZ6kBA0;`X$qg~=21{DT ze9#`oc0!l?pW^*wgU(gLS^gzk2KztUPsXJlDsogq<)yg7T({o6@-sc>e3ABC!y z=lxz|!NYSUk2p#TY)s$KnQo@i5o;Y>!%?EWXAr?lb4e(`i)spQVw~y_M*B4G@yjqR zLf&kgaNy6GCM|tE=8Hk%3@)n`B?)M~lSkjL6<+P{kZkc*HgYIRW)&n z|BGZksUT3#+W;}r$d>IQ-^5srU0Z9FAysq4`wwX#eR`5(fuD^0ar;Y@XIcTU zD(j_$ODQ595Cziwy$ZRpOTZ8(jF;&^HD0zNJW<%tYp>t^UPg@hMzL|KS&+7U&d1S*S4ti2(2|U7v|@ z*~Z%KsT>;)j2qbilZo;OZft3tVCe5`4g0ogw7gexny9M{;f0Yz=~0!c#sh#bkb*r% zz0gXb9m4FK5Oi=NTLboGL&Z;1=F#%}tz8-*&p1pb;$1Nmpm^&hUt658a0odTDv zg9%Yut@&tMyHoM}nq~@I@@@R0f_wAX0T92k>qQBrhY47SHhrK|Yp@vZNoX~I)sI^v z)5LaTngVuBkJ4je8v24ffL&!lPXQ1xGV}VrtsS}l%{z;b2{evK`vY)odKphU>v~%) z4_EaAxUlM7SFc)gH$ETA-ie-JIy~JznHUGxn~q{Dnib*$8soml1kGP_j~w<2VZd^* zr;8zfiSEV_hbO*5)-FCP=U>D(yiVUpN3T%8QG;_djul;`Jfk@Q8>}#1GY3Ir#pLs` zm6zd&Med2r5KrvTdV)&{JwnD{l`1-+2%IfmUrVFaZ-BHglwXd8jmTavkXg42) zET&m~apsRC2Cap}3Y)HHh#Fces9UpYD-12w)le+CkHbUh9XphH;)}{MjArKWgw9&w z5aPw9&HsE?)sX8|_Rz`C?KL*T$K>Ds8aP7CfPa~Y=G-%u)kbUyUnK%yU0ol`Udo9; zQ!-&s2J9{qrCP{&FU8p$OJUz>cf4uyZOevLX!D(bEYxnlsu6%V#J}ohD<)ea6W>HL zR~1*&4;E8KIDX&!Udo^8@)QfuKw7N@uq@%D1KF%772%lq80w0g17yzPO9mxOD#LJ? z)8#=GHAtFSIlYw}sNBaV?MKNVntMlJE+O;W@1Ql9!^$;3P$jb@1+A`R<1UP3K;TNy z_&3h(M$v>J|3KeQ)Ii44Hb4I~Op;soJ1z?%-$OYP-N?w>lN1)3*YuJu>wNnsmwesi zLi}QDyfK2Pw1@-uKw?Az3#ZyE-7| z|1*mP{||^;Q+etBU|=1Oref=gMp9KnGj);CexZy+T*E@}*%GCnfdMW~BUUEItQRQ4 z#u^yaaQLP;85#AGUlfrfn?=G)-A{haroihtl~zf>9{FS~_K5no zK_f7)@WJ_Qh}P`fK|o6XZ}o~Ing(xx!2hQ(aX8atOW4|++7WTwA0NFVgJBYVxX3vq zRI^@7XUiF7qg<`Ox46(BZMLJqi9kJZea&?{*N@mtIH=j9y7EEnm-pX4aI_S`bKKhF zUGgaPIgS#KUwh0Dg`sPwPZvAN8c<08(P9-5Ty-0hhjW9Ls^OD^^#1^npX{X|l9bPo z!&^`}O2~~>sIh{3S??36`10_ zJ{1EfvjmA-(&=3)bfp=(A5USi=@M&Hfx5V>Z&8i(ZxImq)0$o}X-rd9HT4D6C0KFZfRZ`22JVFB=V1?(__W*G%bDa>{ zu|hS%6nKNe7tTnVn_-C9F3>%uc()JfP?&eP;1P$^l?{f_SYX6HE^7#Rg(v43K!Zsi z$e`5U{otYfVpqQbP8}O!2Uco>3YPfYZ(E$djeRxcZVkHFGG1~CQqs_aC%nl> z{ln7m!2tcAWFo0NX0W->yTtTyS1siCe~OCRfOm%9Dd_}NvO;eOI^@8l^fmF1c55%j zXB}HhEvL^u-j$?vuaMYXenv!C{&q2RA(h!Yp5@F9VOe68USX9@tCOBukK=!Ge44sB znBF#h2xlk6UA4T-lb+@e`E1Npy%kZ7)Q{-b_46J27x-+?fio1{ul}TSn0T-&FX(L2 z=`6Gz&W>|$#lB;YpldXv6{>SSDgTW7#~ve8Pmz%&NZ(*&lY`?E(dkM#oc{A|$cSV= z+%S{}YxH!zfAba?N-(QZ7$4(mnrf&icm;W^p`Xj}o2lpaE;P5({H2v~Wa$v7y=&YqI%#bmP0Fxb(y1MtLOZGjTk9v>N7i1( zFH*0IsUlK6QsM=S=eNUTRkP(3j~+EogWV{gSH`|>=$kWrmjnShQYYs45m1y_8oPzp z-n~M&@8j$ZzyDhP+|i&)J8e2IT{whxLCm=`mRVl%=b}DbWW{MRi#{H^DDe3x{Oat9 zrRA$ZlL6y5~ znGmgC=_`qu`p2T(`LbH*?bu`+uBHv`9;T^DIX8o3+zf`_qBMXFHbUth%faOxCD!=< z>dS{0Rp(oLP5*D!1}rp!U?CsEbQcqt{B_C|>d+Z3!SMp2ZS0;$_81jtlgs-y`!`(9 zjiQf1--~_Z^hNtERGK-$`>)~I$%0)4RVzVPDb z{bP^{{0}SG35i9uCibO7e>?mN%HKOC@$PIxgSbz_G}->0ULOJuu}WE|`^ENE$8L2$ zw>fO!up?Z+V=$sdqb}Bf`)TmIoHwNp<}rgKF0mTqapbULd?u)K;T>MD}R5^`as+4JLo_*&I%Lu4t`cn3_z>xw? zuEhuA7f)v!EAx@dMC?sClmItfkd%P8$Wt~i67sruYsuJ;%iN{l`BfF~z!ge-M-tO5 z63p;)+mnbpcOk7*S{bd>nlec}tyELhq~JD;&ayDaysQG77b1Nj#z=0Ojp^93{uJ3G7G|KP4+$o@!> zn8d62TCMhX6U(G|9?O)04#T4j`cSU7YR{i97sn8Hxp$oEaUo2HItB6`YoXAcJK?s9 z-`U*FFUk-4jbOWkh-S1@5jMLW4DApY&3xxHakheK$kB7_Xj_)~Zf!j1#mN5^@pB_> z)rB4OG5JM{YCC@sDv|G{<$HkoZjH-^cdn6lFA7?&*61>Nvqzp~{I=p=U^V38qZq3L zWxgE6H;)ddq)NU^Y^_@c#mhH}f79`+Yf-8CCG~!kDR*<-jF5V~vk$cHuubiiQ!i&O z=Y5!ZBXoQi>OUZ}T#0POn=<{K*MQyChS{N$Ta6bt9jnrwbMOKi};sO55d_-)HiK+lUr|w@lPw0q0)c6 zS_&S^=~OWnEjl?4!=xpK9s~n&z$ZLAIQM6@(@g}V5%yqUtaJ6skHZgIJX)vr-UbOf z)Cies{H;7i=8l-wR`ho;Zr~ud$sfP-HEGJ;XG1oYT}4{G?{o{TzgJ+pUpWdhat>+} zrKVn6DO>LNrsKNIuve!<<1bcQ1{?|FH^ZXVq=hyyJLb+vPk9+JOv39bY^^+bnhB|} zgq=4Gh#uu~X6}9k`b@!(DY|)!6_S;i{thut?T|sq=(%^9hEEkkFP~hh#%S_bP1$*; zNhrV6pIEUfcu^KT=?+9uo1 zSMURRjnReAuK+1bq_ru+UVkI?u|iosgVsw`n}E-rpjQc?$#-t;p3cN{_qxzBNjfY4 z(aWn-x8-y1ysj!54g#FJthEhoCo-O|(FPq(g*o3aUs7W_Q#Y-5;!nwc${95WYp@eOoD#f>-u#7sGz$|Z=XT|r_f-8yS*L*C~7U-g=9s~>%`|3|PSn0^y z`JB;7MqrSuDHfmIAkZeOq(daH*eIMwN~Uxa;yEPqMoiPw)AL!L3h$Hp#TeO;jNlZ( z{U}n-`qk#cI>u9`rgAPS&HnFfBNmT{ zNdLbdVr62E-@VFqs@NueauU^o57%& zoyO%%n;dpbk7NsRs(d>l1IN&?VIu6&-n`HQTd{)xE>Y0>tP;btlMuSYWjb+#mGNal-Mew!ad>3O$9MJAp7VOg^n5-E*MV%dCHByu5X7ax=Jwy= zPELhR{GROkYJ28-SCF3I3Tlk||7^vEUD|3zSGQmcY}gCdOK)o}2IiU16Rj=WN-Eoj zVbE}rwmmSReSeeHDYyM1;#qhsQQuYt@rNEv$x-j+``4)8+u326u?1|dDgD`r-9MCd zNu*Oew3qk9C9!|W?p~fNj+c-B%Hq~ckE7VOY|DG73>pXj&*X&UF#U+k=hqkRSg#;+ zKcQ6vG|vVOs9|3|Ql#-&abz9)llpwIc{|rbyC?1>_@&?#O}fQRRax}0|10pu1;m(x z_EYtKbU5#^N1K`x^PzSLFlOG0iVOfX530;6_;X2jC1BL$6P&5ogAiw^09g4qma| zOkY&^C6r=yNao{jUfbpo+9RUm0j*cb#ab*!Q20?x(}`SC?$%sC!_~R20t1S7dS++W z?w;K#^fAgpaHxF;ENuec3cRRf!tij4j8$PplG_Q;@4iyIey=;HkLm$~k2*cMrIqFv zx2y1P>YP{w|1^I;Njg0s3^2e+Mg0P(*S*Ei87RuWAfJ|gh#)B!3`c8Mjsm7S- z3cdS(F5vs3$Xf3bvTyTNiThnKBaE0dM~lTR9x^|9g{?g^f1cU)_l2*jL%ObYm*8{c{xfdaMss@#E9utN#brwJeu z_}#X7_tU-(U%vbco2%$!E>_l4`tn)KX1v;vo%;vXxf2|&o~zrl%$$2V+F=AjO&^SP zx<6b?Jy*w71=#^r6}XMNBa3KnPT`oc8{;E5Go2AyDda ze$#zfu{SZAYWc@SaZQR;Pw>oye>9;Om+!X|&DS7qI}N+(oOj==7G=?5iisbJy3^kY zfr5+XAHUa3znOGx@BS4*k|P9#r2jX!Qk>O*hwxE{h%@!L?iKiuR|m<9Q&}MXxbFAk z0Wl}~Bcp}6w3FJGmcTe7(57qlK>So|3N+>BbIcc;$e7GKSJ3~FFxkg9^%=OIewd$p zB<10H1wUZx6Ej(%lOEAC8FRe-4I?}HV&J!(((2yDIFpm*UYb=Vbr6zg*wXCToz86ey*iaTa)STV1O}tK#9*RH*(x&k3Z@HuZdcGM|X9HQeaT{epmMFPI|=GJ@q>Rw}+lZ zqI;;UGY2ec)|UM0`G1%HvjoX4Aza*vx}ZTqJk-#EBTY>RzEGQuJUN3g3s zKN@7hxJ~|pFKj|Rg+f^m_j|?IuZ;4@Xwf%a^d(@FA!E!7sP`6!0`W6--yEbzx5CEQ z*-6acAj8%D7XdixOY&~tKTW3e9A0m1?B*4`pHid5x$QFI2`E;IHfupa{vhDsWdQ7kPUgoqdp33N1z_q-3o)T8 zN{rAb>8lEXm=-@mf<61@2j48|UDu^W_>J0EuOQCLv3)NshFh`?R)EbYu#xOKv}Ra; zmYO{Ql=jtc08d@J9W`w9@ah6W2(evGCp%um{cY1@A*7Vb@ivC*MT{Ff51P zI=#Sv8JeUdkk}Jt-~X;NIuJlZ7BW=9Qn&Z~L|L z=7w+g#ih75zM>WQeESle0~iBFP7mH|*WSslE0rH~_2+F9lWc2SuDfO0(LZkd6|g^k zeCxM`)Z9OfjzFca+Bm@T?2D>~hM-fxo4HXH{qc~61cPTs<-_86fA&m)+UwsMO##0t z_v-IOVlz-*i7D^;kS6d4CbWvJ#{6IGy=7P&&C)kcgd|9ULvVrwX9@0-EN+3|n&9qkn?QiY zgDeCO5L^~_c0&jd+}#pjS=?R!lXKtqd7tY%=ga@=d%gW-W@lz+y1J`=Rb5@x-IlF} zINqsdUd^3IiM$wIP~%6OWr_R!dGY%f$QQH-_&L=qO^da#BfQOEGXOS!#-SaQm^zvT)gt2m z)v+7AyD8pKfOB7_E@t9O%vSX#wg8~=0?(f+aLVlEYq=hq&u35k7`412DQc*mNb;pf zLS!%mA-nS@aBtlV0N8nC07^qqpfXTZvFsZ6a*}C_{boDZCN2(BH~dvUCshR z+#)Xx4EYrGd609eimgDole?_^1KL+|r9u{Ur2J$_a_(2LUbdo!>ng0`H zFnSH|_G>VEL z?t)s1F=$@mT_K;)xMCcP^CBT9CQTkZ{owfGGP;SdV~uWh&(Cl0hB%tnAkAq`y|z?0 zm;C6$*kM=5oLbq@_pb56gz&p16UgYJ!2_NU_eoFX>5)>~>0X*`7r7K6*>GJ`5}9%+1RzEnStCMMr`B!4<%4m%l0qZ_5eq5q8=mq2UH zgN&GL_v#k|y}sj+iKd3(7wtiVvn_hq%|rH`*W9ijKKl1GDpgm$d~=rTwo5Qb6OY~)Xa%758<*$X);F% zeBPB-4n^~4laiX`5+A0wfJ9kmFBOMcgtG+>D^rAac4 z5V1e<2Lt4JO&SD1m_?1~ zWK8r7l9>{yDJvWyZl}meLFb@~qS6{Eam3B1j^*=_uuGc2+gF0Uq&ECs9IR@JBFGBw zw4gXzFZW%c{gT@a-!QVX$S4?KiLBi%89X0P(xI`t$qhRsD|Vm7*@%x$-OSFK<^w$&A^Ii{B)hZ!ol3M&gG5 zI{yZYCkFbPBL!DjtG7!pZ#v2T{6TdB;7bIT6Rbh4Tb_#f=<^i;bE`y~Tg*>FZ(>}&rn0hI7~^NIPaUs!k_|SM>G7PS$&8Z5 zj`CT;$96GHU!>%M7fz34RmZY&S-Vtto-*XB#I6v$@`H$-eIK_{Yq4j}QNcq4oDlak z)aQ#O$*ZHo1eN+7U$m2z2y|_afOfI={ zZ5}hD8e^_K<;mtJ1AbFJ11psT!zVfXIbswqLqqkDl8Qgn#wHsz6QEZwhf<1nH9|VI7zGRqm;Gt3jWUXD6i=tOz@GaT5RmH zXnHxG-v!j#O>f_DC;Up-j@k66#mmzw@P-eMD$`#B|9*4so4JR7bI#c9{z`(_E>S@7 zcw&G%es7%hZZTmBz&Wq~<{Xqxzg$er$(i?w4p8^usx9GKbzT^F!waMG^W~|E1nY%n z$seYf3c}&ThD9CkbKW{*i&5W%L|b8Ml`EGkmqi~4te{y@u998b8|y^(f9Fl`-pL%) zvMe(1R<=TCvXl(gAq780)i9bgt>iUE>fcZ4_v;9VNbJoMUh>&ki{V*tF(0|{>+U$* zf->~ze-}?UaV0jgEo;O1&nXBOG&Keq^JAHmMx!HQ$?ZA^U@3Pm87>VA>Kl%>HqTAf4uxpTUx zfN#S-D#IIwA(l6ztkO+9m;@x+MZk7hXH`G2Rr2{4pQ(^oSrMt{N#BOq)rF`&P-<*6 z$mwRL&r{wv-M(()b(K{GEAjb|E&_|v@*y;|(@G#8p2NXRsw-R{f298ME**9|Ei1a2 zRR!O0Xs;3u6FR(OF(8and>&4k58hyS6BR!Bd^D_=-HBvhS|lVC>d2w*=kbn{fc2>` z0eLm|eg4|QX9QsPz0Y==tGKtACnqwC2Z+fUtNk2yu>ZeQUvX@lP=C!iH^D_t*(URZ~zo?@a`Meck<% zdW=W-71*b@sEqKCPenpu7%N(D6 z-@6$taX{wrs!uOXotF9|rOQEX`5Nc<;4Z_%Ose;NFLGbR5WmP(6#nMcHjQ(b*_ZhI zkKxUcl-;Ro65d=g9@VMB6Qxh4$K@j@1P@lKr4<*nv;V@T=_GP}Lx9XuGZqn**oi&X znv`ah70g#MPGxv=`$5-4??LXatAOagI&>}kb4g#OaN0}GZ3KeFVE?|CXVr9WsUkal zUl8VJqOL{^f;cH8vQPmK!*Pmqq57+v2f73f5Z9d zn|!zGKrDjBbK>G^-p{=~t048vd9?3on)+(lUB<~TUG=(h5BslgeYZZY3GNxm?jTrrdI%Jc{53W$5$^Id4` z;q;DIN0U|4rL4ALu2 z2VO!OYsLLd1}~N1{=y4aRXBU>#Wn|EKd2&M(Z|op-eUIB8EDcq?nrpOA1_Hn3QYs~ z{MH^1fw@!nj7UVPOog>ajBBGXxu~S))CA!sJKSlA_NkU73VO@|$aVP9k4~&`=cXnf z2|6Mhoe1*dVd;kaiO+|3M4IlkNGZ&N$mpW`t)m!Zl#v?R-NbEJ=lAX$v8&Skt5+-Q zS{|v^61o%w0m4i^II#>Y;mCa&COW_*%DI2kWC%7t_1%T+wwo7N^c4xJR3dDBbK$EO zA%|({D@uTAAXhno{VYdfG%T-!40=q*nylpJegy1ynHfeo$DWke(=+s3>8Q*bS_pSC z$#k(;qKO9P{8qug!jG|kygtt}Q>XTU^(Z0qzx{H;S<3S__m!G^s>>M4>QNE&`MQ3OAa4S2AwX-*avEl-Y0)yU&=xBgT@#&$f4O#7;< z>JwgQJ??#IJ(_#o;b9hrO_G2#6;lQ$br6?QNt&B0w-rY{rvIQhWLr(34r{{~Idndb z&|jX)I;6rWd|r?OJ<^+wI6OY(_hEPJfnyncfw4vokZ~wgF*^qsJsizK?yGGhD*R|X-sY-E6IvS_ zlyKfM=~9(kYaF0cG<+|KYF8Fr*B8K#OWiHOJuPD@FtiTw{Y?a(p4Xh z81iMt)PtAemWzzJ`el$P>vX3}(=8@(Q2;00#C8ZXAEzL`S!v#L6eMt9WFV z?@8wqXgX&1>L9D1=^-1eZ9B6IuXZ^sG=c-h^;pEg&rcG1=#Vb8oswD=%N2Y$Le&9L z+TQ6-chJ4FI-O`nH6Dru#~6PQaJDB`D@_(CB)O#11J+CyC1rDOe5c1f=XoKh(@osV zcb@$Ulfj^z^$~Rsc|O|~cYul>l5AYL&@sYR6Yz|Mj4MWU2YuCjKxrC4l5Wd15u8sY zsF%lyjSPg03|byIHzLPr_~FamXpMbWON52M|Z0`E+)*w5!_|EjP+EQ z@7E|y+)A5gol%oG27-lJB0e3KWy)>3YJLW{WQTI^5Us=09oz&_IukBfee39Q9k%SqtzH3hy+7+JhXX2Qv1+ ze05YCizhHkP0_fOS5p4=>8fAY1|HrMRiFC#i3CA1`0dM%y75<>!q~^Y0U<$gSuO?( zmzt7wrLvHjov1TY!Q&X3f*-Na|d&hOW(SLqxah z3N0eCCIcg$t)|2n$;tWU`>3r8-8}4?@xz@ykJuiBrBJp;f<`Dvt+b3YuP13l?$YUv6~=o#@5bs)U}DdQ z;uz^PJ;lS0t z@%T0wIa?N;Yb630l25zB^plQ*%jZzU;lE7;jZ>O8?=eFV=@OnJZfIuQ;aiJ~ z&5VYK?y5X5kMsa0zV5a;B;2>7)l?kk+|Bbz=y)3LD?$6-fWrVQYnqZV_XKQN~1O6|O7j^b%@=7j7`omKp9^F7UfJguOp z|NYqDG)#gws>QXN+DgRaNyhs2^_Nof1l_@S*PEzOA_K4bo23k&(pbp)g(JOwt@vqt z7q+*y_wQ6L`r;NgzgZ;le)`_X<9UC&yfBBd_O4f1B(5Gl=PJ}0@X&#yqd>+bRZN$Q zH6EaeQt2L9iFD;IjT4Q&~1q9hoM`spuXXzg2^wE~U9=J(IK9)_=p zl9ihnXnH&TIA2@11|N+zHK|zSU2~nb`z@(h6Bwl{3VXLbXpg5Yk+eCPnOyG;w@$Wr zAWC+fb0HCJc2@~@?Fh^8lHl!(%B2mOFQt`@VlQ(4oc=tJ88SV|CNU2m+ zv(m2SbLdyY5{aE$qtcJPFM`7Yv-;6q#fjl;ued33-j>JQ^mwiYtEU9%c_cWlPuQCl z?$;9h7*wd8+pvM89j5?wy|Xu_L(ccYow>BJ}W7<=18 z#)I9h=dRf4@<4sr*3=pqyR#V9<`Jrctu+?Mzxmq9tVkqJ2{HLkEI?L+kk=Cl!BU}A zx_|if|HR510C({LcL?PzH_Y6JGX^h zq8t;|1Ky#5qt$BoXOXxt9b!#8RPmF^Z+Oj~VI~PhJ~k8dFXyO-w9Gp3*dyNhzljU&R5rGuq#Is}c2Tzfi5ViI8w6z5v>@q}&o#_&8UQ`= zIFyM;)^loMKf>aHN>JrXp)qFKD{|(fK4n?rpZgnL51J94M4Kyf#?YO<^C?(BDTHpH~mR#5EyzgXzH@`AzlI_1Z zn_r}GHV{^6%&7h%l5#?# zNCOq3s?XU)c

C?5MP~!<92ZWV(a*h#Tkl__zy^O13%6@woaIJ?G|cKRmS;i4BL> zOjz>N^23}?E!4*?vRt#TKkcHUHG7FS-AH}BA?<74QhiQ35N)P6 za~?~{4;q`4^PD{8ZtM3)QeX@W`I6iMwlMe*VEs$n%I^)aOs=pbiX(o~LW~B23fP<& zb9QzdfoLmk^DS4?gk!<;oBhwMS5YoSRnD-3xxo5K0UENL*Dpeu^ny9fvKd{O%u4$x z#Kb*WdQGv56O-PUsg#6}6~S+%S?~1e^;T?W;Q>>o)r`!|JzpeEYoAx57MIKyfe9W zo-xy-^+0sf!#OV3E*prdV=Vkg<*~2k?^il%9#FR0<`<2PNo1J5H3^yxxa_`lRTGkiPSF>UfqJ;3mu1o*8$)AZt6U9XHX{2s}2-Jh=E) zDS1Nsn7RjA)Xnw6(&mXs)Ae045y{*ZceQTJdSZB%Qt^mPZ1?RvtwC@Wd@4O)((_ft z;C2~aI=of*q<7i}z8_S(EKNv08d@_wdqRCC**2)e?pUICT=aI*)Zo%YfAEUS`zGkv z1jLDO=2DZsNsk=eMr2*N+V_W3!uzg&oPV31T%_~Gfb zX1@x3w-NOqSJCI~x25H-+_j`l*ULFaLD8gUU~I){4Cm$cY+z*eo7#0d>#|Ma$aqyw z>9pI?lf`TMOMw|eW9y7O>#9_la6~G24O0yKcIMEpIfEgEC4{&Ic(*BI$+$TBzsZm{ z5428tZyZVtLu^{yTy+9YL>r$~BfWf!A(zjhIQ9drMrZC;bb50zmo58C1nM8KZSy5vj?Ssd_f zo-T=ue^W3y%g*a2)4ynGP_?i*z(LR67U=;qdlrKWAYr-YaQVERX6Nh4=uky^*XD;$ zsIPYlOTwU4iH(y+=gqgcf8wH`dr~8RD=uHPYD=bf6*@Q9h6*tidCLY z@-@CUX~n!=s8b|TlgCoQ(J|7j)J)}Bg4Vig+V*no#-^aLXcw`BSNmlyDJ=mxsh!j1 zKIS=U8f?Dqj5vYiGCX(MyBa_m3v?BHja*;JBY%u2u>Ue#?zQtyPmln$nrPzQcKeu% z&#ty;{cJY(n zh1K+wPL5CFY^8SpcV%5sw|iifht2v#(@aL+d97QrrXlU41^{s~p%;CAt4}2;KQLzx zlUTSS!3QJC)Z!8+R?VoF+su;h9+lC!%#$=4!z(rethj$jiunte9i>P-3Dj-)HM0Rb z4J8%QMQqh!Y;J5_uoD3BczrMY1@0c8qo^bVw1@3y;{y(Bmto7y<>QUwSDDvb^dcj_ zlH||q?@QyW;?6$La-ww!rS%kI)|^P#+;OfvV5aa^rxp=7@$y}UHBK5P+11a6%Re{U zX^`37KWwYzipkisG(WS?vat_|0_?K;W>#HEPs#$I%dKsNatZ5F;|H@c{)LH$R;8I! z*%RT>+c8*S3A0J8EQ!h5W!`j8MU}v=dJ}-9v{!2N&G$3(;^78+51ep6yYgb~D-h+L zee**X$GnR(^z9Z00c-W=!7Y|E)b2yI`g7=-X4>g@WqbC!eQCO8R0O?20Hdi@vZ3gZ z-#RWl;}mcEdPJGCPMvQarUouAPrrU7zbE1FF0AyMTy%ztN#(}r-qTuEygNKk2OkDi zR_o~rZ%&HQL}f71m(_V+&oJXnLMAtBA6%)zCk48Sv`_$5UT%(kgwLIW5sNI z*mz>3enO>M>uNK!2`Klxg z;)*LN<4Q73uvXuap>(qUocVT;6l!=$it#pZ>h4Y%R7hXsQxt zjT(OG(^u(lCZx(1;c(g2y9BSEq#gebNX31Xb@lNk=UAWK*7_yR{dg(-!n}DoaI#Y- z{;UcZTi9*aA(VgI76^b9yLWF9^Irk5!|@T@>?IQ9H!#DbTE}BF&Ve9)T+war)G4r) z_>zynQ^HF7^_U(e!7vC@a>?YFc&b$?Q;zZB!)~SAF`w2TkIzgw+moxtb;6=o>}8=6 zeqIYB9g6fKnz6=;7@w0c2IaOj-S6{Mbl=Bt;}ddEg|6_7r;eL&{e(ZhRfMx6-_9o} z!ahJRIM?bwFberOe}ax(l%^q&b?&H-)J>$N0t;vJ{5@2F?NdS<5r8Kzx%#0U{8fWo zurxp;xmoC{fkhis~92Jos_%_p^#Zk)Qj74Lu-?T`z`?gVj{e;8rnKnL{%q=wrYCx3trrLpU8z}q9@dDX`J>p z>arfi_wNFs!mp)oLTgKC<9e&<5#q?uoG$yDpQuW6yNLKiWrgZkXD&fxoov9agKUTx z4}@){FQ9(&TH@x66>j!%baUl}SOPg{PQIlmC!fT=@=7~PhjZwfTnIo6Sxrjto#@Y|zVh9LRIM&ceQQQwow z$sn&?>EXXdc-vhbKTOYZxab#qj*T&*o8<#)|2bJ2SW|Nr)V)aYn0*3%(=!}Odz5;j z$wBH?aQFSXP;D=xPPe0XljCqxOZOsIW5ZsP(`4&Dh&gqUhqIZn$2jIwzORZOG4?h9 zvU1eDES!xEG1mn|D!0S4^&m;O2sze$5&!QAWiWAfC*!gHFI-%>rC3X@9+EwE~U_(p66Ba-w-CZVvMWgf`j-C)TrZ?_- zwx`;9w1+DjXxe|9dCNCZYH;}ZB|!lET_cB`!!|k+n`>=Xl5e-@9(_*d8w9KucDiVz z$(9?YT$ES^Fju*5xvl*`A7=bB7`5oSu}&`HQlNchYUj&atAKNQ3njmN8CXgmZE}*K zzD_f*|NM?K8_cmiwd=TvslC1bFjGXHIbb!GEXl~I`T(JuT}h$RhPingtNn(D!Y|TU zzL4SZPOqSiQzMt+mi=pKgef}UyooBMe_+6A1cMIM@>XK6BT16hE9Fm&?oa&SS6y~G z;w@jc;&sW0Nz3)#oPM5-4y9*odprBnW3jYsj>-*j9nq^)e(oE6zCv6wE~QV$g98FaFT%Za+8wG-Vu)nr8Uc9L(IepHI~j?dnROh?uvT29w6HG1#fL&C%de*%etBIR0>bsVrGVP>Pyf}*tGnQJ-2gw2-b8AXoNNYq^@S1lj(hWIo8Q3_Xtqd`=XB9xYX_&#tZ6maqr#;< z6<%oe!}B34@l35p>^o$cxy|lDfTx^s>C_HgnfA#vb!T1mqpXng2{UYek(s9X-?$>@ zuTGz)8MGdLE)n$Lnoc`tk0eg`JXM>E*`I(7KN@ zSyTf1XHlmY%Eo4*b3ue11D6Lg1C}$If6`~<>0=JHIS`1PtJJe(7|dS3amyG`>n48! z(hjDsDQSyK5uValy|zw<4ZjiCi`MzaMT$vXye+op1xiX>F_7u_gH2LUOX~8FPRe}& z(R#CHgl8Vxjm4y4X7jt~B-F=qMV+q%=;@|ZIc%KCA4q#WRGEwn*`vEzP(E%g4+BCM zli}LTspKuj_haLelf5%+=)f!wFUlYCPcT|17UH%Tl_QnOoZNE8T$J{GfvMx2iOu5& z!8JFklsAkwnv#hgbV=Yuk994syclreCr3ORW+xNc$Gup%FAj)jLq&`Ztcp>2x?f+@ z`m_qv&Io0&3lA_Ss|^TU6N7gsAKqZ1?i28sEWR@@s$+3Z@3rORk=GJZx0 zd|lktWOFNUW1@6st*WxRrUQc2+?vnNDyyGy31*46vXANmo`s{KiiT5jO=DRm*DwV5 zpkR=>P?P9x&z5I%`)?KyD7L}XJaH0bu6eT}1UD?t9DvclRd`#Iq%<|SBrntaS7qap zU$PpNu@+?nxUmxmanU7Psh}0>t+;QZO3%KgC$B=6@B2jY*D>Hec02cZVF0UCl*i~+ z1UEZs1upcoOzSkOGv!<qhN(DV38QLG6|d+HfF;EA+$FLYCPJeVT+$7=Xcxtu@d}4Ej6G z=s&K1!v7l<{^R=lk90t8#eXS`_x~RW|1``0)xGe!xCzLjX|LJ3xulKHxc{lCt*tF6 zCjm7ZI zsKxZe#6)>H7fi9%e!40mBH{w`c2#hG(6S8ChAhBDX`Y z@Pso4@FdRYq*;aiPPX2tiH{Et6WsjB5cTRo9sGLd<#j5E*pFUaT^+*onsez&P?IQn zZ9V{rQ9f?|0AUox!uaGJNH1443hcC<`UkLn1Q8Nq&=r3$X9gyREV6;xe#Qj>6GfH( zs<`a{Xp+ZeB<`+&{nXIJ#OmSsORJarMZZ53=jP@OO}v3;mz{HRax#<>{1DG^x+<@@*!K3rb#{W)K>D*BtrG2!0cR?&ldWk6xfpkM z_p|GpJ-;fBUKc~|wUe{`WiaETyS-IngG6z&+dCT@*v9yTxP$Xm7D{W%K0ZG0-em(U zCq$J^>^NJ;>E8A|1c#lCtyJ6YY^3pjw*t^jk0Q*OrTDb8w7_E{Bd((=f&_8%j=O7X z4f~yx8LDU{P%9q~PtNFi5E=gP0jf9nwz#(TIBoIDQcg~8ucEYcE5#~PuhJBgAUZ0_ zeztCVKy-I&ZLJUvXU@r$i=`2A{#DsUb2(9uIR0~IQIZsPxssr^GgZlO=JWRLd#J6Q zUHho#v_*0&(6jCB?Y^yl)Srxu3=S#BqkzrR%jTx0;^JaWwsHUKy8lK2Bb)elD--YE znDoATZ4Z|X-xwOEr=|HFtaSf+$E{pfUOuqyPB^?iS^z?vnbM2+9GC;5r>gpd;M^Yp zf}HnimmD8v-TDJ!C`TJ`)WZ(hv>B z9wmiH-v34bl>YI*2na!E`f;5JtKR2q=S~xFm}&K^A4IR5~`JvAYEovJ2_>5%`! z`kxat@v5vX%R1yt`l7#Yp)%r^tQs$Hthch=S{#1ei0pUV6I{G{9Ui~n-42y&0y_FfLD1fahF@$P3lTePTse`2 zbX@FfciDp?z7>;Lkl=&yrKLVK%M14eB<3C|!@sH=T$)YUL!_SNeKuL785q0+2ts?rR0GB!3=Q&STk z9}kFZguJO~&Pe&_I_`sZK;U_KdEH!HOSLzFXA_IH<&oEH3d<{CGW z1_!bvA@fZhfZT45m*~{0ViW<4{@1zyp1=wLf699iwgNGbq1BeTbtA$#4C;Pp2e0uu zc7i7k%G$i*WAk_?Gh8;fop_$!VJU3r`ckh|ys#ST@4p&J;zOA+9RS&nd%7#z<7x(H ziB@zQ@XwyoeopU`iG3qu}==L5&am*6u+DysoCG( zr+5AJ?L%iYWu$zGc3I*Hpa=`{@~B1J*M_oR(dXEc#C`w%y)F1*n)37DwxMP(CleEs zgziKbOlw|qd1)!pdVI#NKb7ANJ)okXKrZNL=ppxS(l+GR!<{oKYn|Nk{OiiHLq-<9 zGfV5OSzEKxTlEYqT+DeATZw(%4>IF_N%ow65$)jOdP<8EcwPBR$HFJ@u>MLM*G=`- z&t#oIE-X2=b#o)}@9vv}>g!7wXlPWqqI!C2tLX{x@xN#el^Zs?0+Qm}dQ;Wv9}oaw zMQUc|LZLbXfz9v|Qq9Pq20!lT=&052)N^A@A$OQ%`w%@Sn4(PdT)E+N!4IZqDS3N! z2t2W2S)y!PvQji|3%g?Xdd+4>z0AgsSni6AjfzU=v52VpM|JYu7sY{WQVdmwdqq`I z6vWm=w$5o2WaDV*r(XVu{JD@5+*2%FMoT=hDRBD)lLWKGRp837LCJ=eT4GG8OjOt3 z-gzy^0V)otrT+#sCA7TmAR2P+6jj(2pS@;qvC@U!hym)yywms88EgPX}JC zGZSDvZavx)6csf{aXGm@T|{nS{N3GO0xvsiPHo~|8$qimzCQlR4q_>0iP?|A9(|Aj zM9u9-5XSX~h|2s>%1cW=N|Md1j1)JtQa%dbHAocRKl7s>VK0nS-S+Vta-i2Rk3QD& z^Eh%E9e~4R^R(%^HkX%ak96MGQTwav_v?mAR0#vxx69WJ%ZT zE*cdALtLJ)!j~S;IQ`S2#m{VOV_J1ylS$&4Al;OddU}2DdpZ`!I&3CyWgp-JV)>I4 zu^&hgteQ&{*HE~&AZ?T>$s;pex%%qQq8~shRNmgXUu@1t0oT{3z{h{fM^Q4vE-b89 zWVJolSe>h$r{S?aVli9iKp&r&xMGcmi(6{QWAOTQh5a;7-ml}C#=N||lzvN#^QDeR zgL+4PifBL(i3m`;fa(xdw8x_oddriND&SC@c$EZ${ZYEp)6>l+&Q9 z#03J%cNz_+;_QkSUH{oP8QZU+&ChmR7bG5Q9tA<488%NfW9EOf3wEOGzRv5~x!qoO zH3)&u-Oi5Y;wgJ`yGvCFa3gb^e;_20pM}TGIk6`mEKpCihFY~-*0%;8wNA$LRYcCO zk(&Ym=M1;r*K2;Bm2!S^YhZSIK^yD8TJX!WlDz7{g^JbwGC&YQD z`a*HT^~o$AwFqC})!B!84>KhKeAh?v6j>cqRR?P2x_|zB`oi>w`03m-3RUa9Z%W{^ zQB+W)Uk3e}pYORg^bTltB4*XC*?K2q7+=- z!K}l2tCT?pwmA#3%gp#!mL+-x!s?*Db5UUU+LP^pdNNbyJ*j{Z;}iHzMDrniC*ty9a^g3R9o~P9;z` zah$!IY9ZIkpk>UFV6Zx)W%{MYQi*x-F=DFQaw`Qrx*4$LZrBrb{5 zfQ@%1G$cfrLN10XB_pHNeSPHU$en*x>LqX<>D}B| zi6)>;P8T62U_eKH#3Fit`}B*wMv>itd;E!*?gLNOQ^V}6$rc!I-Sbteuq%hXKtY1m%L)BW1boHOh^9iKSpJ%!DQNZ-UHO05$)7rLH)NEhBSNd0+Su?<)kBkbHTu(-yf>$y> zNqAwN;uGM(!l3tU?%6%9n{(76RkY8eH?4%Rm*h?G0EF-fI{hEsazgzvV7q7`HPiNY z?%aXRQVKe@SoQ-|ebYGH+iUmzBi3?fbZ>uuG!*)}&T$c8G`)dFqjeF=%F2Dqv699b z8u7pYl7=QwFWf;&>f2ufO31~Ljfu%(VIhl!seZ>;(HkS*AAhsX>#i&P1~fpkN{bl) zGaVE0Jk5IGgWkhn@!%Qqco}oA-+_7P>S2aA*R$GF?5c?tD7}N33nst9^nlWUHdl3U z3!gDpw_ff96x!if8LpyNq02d)w0Xr=-6X;H*GFEPjMpFN07}VZj$t$5 zfDiAl(x74NldZk|jNVT<`FZNR`dJ#(yI8zm+bz6R*+w(Ghs!h1QviUbSXsDo@@p+! zr&_Ns&aE;FG_V>zGcTgn>994*C^urIWdg<}&^<(Z=4b%UsHQ zozu?l<6x(^WChhFL4UR_F~m5!se;%_S!`EGhCAUhfT?S4lLW2d=YzA^hqQiwUgsIe z+3q1*JubGt);#(1*??%%L6+r0^tLH{F8jdNa=Ke>fTMDHJL78h_{mQ3q&F>>@XbVb zqZa1vF5;wEfE6;kB~Uar`1}?!?%4d{Z&<)EOlHgW(D6dkJv_B{_rm^mybajJ5j#X) z{}$DUG)n}($xZ80K$YMY7vPK zVhB#2Us2J5oJ>Hi%|v(J>xb#cV-#Y;KVdlM2Lh*p|vgu#%36Cebd2svwJ;uKue@$Zi?9MO6^Ye4ex*%P%(tm%T zfz7e}&KI_pmv_Ix3_aY+Jk|d#Re#r3idyRYALjpu01AAqmz)0|Yi}J;RrjrpZcs4+ z5h>}GR#F-fkdQ8ErMtTog$+oTba!``ba%4}>F$pGP2LwhzjN;S?)~TPKM*%-uQlhG zV?5&-&zO6yyzSn%3U1ZoFe=EH-QMq(z`&<}m-WwqhRsc|hfwGtKJ)J?7nJP@dofFd zENgJ^`1nU~@Z2(o#qV$R^=^E8?th1LJ7+))zkWp|NJD4{yP0M7@8uVetY0)_5Y-NO zwlPDjrCS#C=TeYigwTK6Vt?O(Bp@XIKcsJDs8=<523mAa8Aw1(dcB~g7xH^SPC)@w z#o3g?GUJK;VOdG>IIc?Ycb0%-Pzu*wpKn+#cZ-OKpap?pB_)xxjRNZe5RM8!jF`GO zAhAn_%q@3zaqpy_93Agl*pWb_Hg+F9gvifhJcH!PDJ!Td$jI39ftGO`R@i7HM6NJc3z+l?s_JPlo~)D*k)=rFr{8S3nT zHh^#k>2P)6UH(|UdadI`kxl^b^{LzSK|C;ocRQ%%+)n4c29kji^}d=@Bc1D@n`wT2 zo;7+4Fm3|_gZ%pIst)ZJuQk6_AN-dj@sHIrc-IZjjCFOFgGk>4D$iyxNS3-(Hf8GT z?H$QtG|D#Xd;nA`>Y<^$kA<4!1)sMiMwMXIb^*S=vf=lUUs#W++0|ZbSD(P_;(_8& zE&D&#ID>L{TVf6=mfcE7pvrnprn1m(XNJUbckXR;kAk%H=faPn4u(LNGiHrWRuK{D z0Ci`%*{tjOFjY0oRY@ECxV}6C+AJZrqnx)Ch7M@WuIqR7_05(Vje!y8RcQ$dw(Kpm z`TF?8us6)j&4q?ON9*Y9v}9J!9o1^{2NZFulNc{jPEHPw)0UQjVJGXem6cU-QBi1e znR0Geh@w9?4UN&o@fJ|4!uo!aP*Uoh09D3r2MEzN%TNEyGynO`FtLi`!dPEF1DXTO zF`OZ7c$7*qH$FZNw4~0VAuB)*fEHUb_(r!I8_nT(LJz*{#S3kZVlkR=J`fM-NHN|T z&K!z1tHS~HYfc=u({XBqdS`pPl!Bko0w7G2WhPEKh=ApTq4%`6|9q__RBj3r6%$Jr zj|P6usFU6J2`JXr19@tk@8RDcWA(c3lj54`js4vx6{{tMel#Ann%F)P#=YF zZh@5I@W}qew{Lfo2kb9T_imKv3B7nAjLJ=FX=#Ni$!R!km5PK=cnia?uViIq!x)Ev z#`}YimxqyYKQle<5i+t`k#;=&JsEY@APUk#;6b}9hM6C0r2n|d+WoF02 z)t!1c&zmU+cCP3@N!NCEm00ix5=!Ym~15h9Ve(Dy_fma>ivbAD6D0;i_X(W~GBQ>HTeEqdtI%Rvp083aAwWVxBG7Ei%6l?x1iDS=^U$!7-GJBdUTiXzD%a|C zg}3-(*=~) z8BefZeFBT_e|Yo3RIvzue9lp0*1kAy6#<^h-FLjay#Evxgz(e%ZCpi)rOs#jo14a| zq8|kqnV5i-?DYhP(eY3pZEd>J3IJ1{S~b0DMZ)kHHP6{1daC;y7M2^4cd|&wSBB>T zxMJjoG$6@QTkRe1%+|}MONidf)dUa;*6c>%O?c}ekOC|LS^_}T9sp1{2?{wdENZ4x zR)awS_FCY1(}i01mx^!AX<9b?1zq!*+0aa=d zp_!-8UFMtVWw$LNbZ@)ObFCvI`$}jAqd60x*AT&=!K2ms1DGR6sW=TRe16TokIZie zdAnH`JAg1vzkY?v*Q(Eg*@3|+0mQM*Tt_~!V43-$kzK0hS6lBV#EBF@}45a}r*0&viBoIsne};vTlf=7SpQVYz z&GtQjzj(I;APrsnL}MnS+fn<&>z~)4hCtoreixufI!ft7I8Qu*0^USPQ|e&;4nVFdwOd$$8|D?U8@`1)c8 zw9rlfp!5O1g(V64X7S&000=Oa)1_&!axN8BcQ1D|4o>`wr!U4fzSQTd9hUsft5R)U zdKk*NWs3Rgz$a7-J~iIBmt_3)yceN*>53;ZZk*e-L#j|2SSRJ&V=!lm5Rr+f*jR9} z$EX2DvFNF`3Ie*lF_BSGS=PNB9cV$oMX~*a-b*GAtZr>>WxVtA_O9Ii0co?b#A z?L*}Lh@u9l0v7VBw`Eu{FfmaGxW=^^fk|d+oi5JLqdrii8js<0CK{AZ8riATIXf)L zCl6$VJ^)-|(TmXa^n4>9Xo^e#hd}Qs0H4zqV1;At4A{Q;?DD4Ycp$axR=ZU;6~k49Mk?fdG@=N^vvn|NnFbEGYv+&-W`O$uT@@pBjzQB^ zQc}{6F~L$_A$#Pt)^-?vQE@RqOofrc`$$MU=c~z0UZ}HGHXAoKXuMDO91H6vl}J$B zP6U4EaGd zW!evz%;kD^J^{C5iDj~SzS%(1cPQwT)7W+N_OeEPprWGkpHxEY=<2efWiuN6AmES3 z{#*a1TmG1#W*3lf4&~MWLP9JoHUQ=xjjsTC0P4H1u$)}zjS#$&mFel}>C+C{4>W9( zTgOItYCuB)IS_EO(LA;6l9$CZ#~>eN#$uqgrBwFngRU18kH@(p3#GGAqts)>P_5E( z3CLEVOw&Mnj%3pFX$P!$)%g#>m9kRXjveN*5!U~Z?W=eE;_HGZ*M)rYw<+gKzb~J~ zIS$Wc%M;D0VNs=o{W&n08dCit)^iNJk!jf#+@Ttk>KQG^OH g)%gYtqX3(>Z@1) z^-wdn*;vw9WM7jllLxGRb#$Z9N22=6;JrJbCo5sS`YJ*tfHD4b&MpW!Q!kkZ=Qof(8JB5fT!j z4H%Y^06i(+cx2656fIXP;>V9CI*5wdBh)0^j)zatQYGZ}fa;W|Q9q+w%&r{TVS06S z)#iswNg_&YkS-A`u$Uh2YIz96EYL#+b-WjpjGP<Mvy?Rw(*p&G@&p3I9Q;uR4jLv*eul2KE+2QVI2`yNh_cQoX+D#QKnz2Tg zTI_GB*9pCcCtV)Qy6qoM3ieTF8~QEVw@Qko&01Oa8KN0XmeoTu%w9oeE#nU`Fb2GIKSWy1D=2HF{blD4_h>w8EDae5)f!UKiXh3nIt0t z*<1QtKtKT6UJqKY;EPlh6%~N#gT`@sdb(uHo7XHXHh^r?DCWsIE-x*8_woXZXAa60 zKuBFd+g?-iHz>>^Xq7ZyVTTq?cHPjTj$m@&s{~y3=3sSzSX!Xj;sXH5p}`#i7!z!) z08q9y*XY4%y_ztzcEg`Q0cO?x^EK!w(oj(;)j6MmVkh6lT|A0;ZK`}{eI$D9!C-g>ZLoOau);9NTRKt~6$ zn*$o2AP(cf+n}YxWxXl|aF2wP^x|aaCT#(tjQWCE7?9nI8JjUR3Ao*D(&F>`CP3a& zzQsV2&FfG#Y%L-1l_bi$y-a?`r@iCH$UWG5Bu^G^k#!gSDXdOsy$nd233_C#fP6y@$fF!(&ND@E8Ga?pgN0p-WN$d2qb>j+0_MVJh0h;8Le`` zpOe7#biV#-%gkTjr%U0z--}W}vr&033qSz_=xTRwuX>erCZpOm_#|+nSD*|6E&r3X zp_QU(JiMXi^Ye4SnRi>TG*m&I7Ch4oh>~R7=4$598YYAKWjQ!Czzv1obf8NB-~sE* zKytHl$kp{4xKz2xRNCQ<8{JU!v8zA37#d1nmKHC|$o58BqwO6YD`pOpSf2bm6TP3C zUJgC=Vz%8(8t&y7H&M_@{3T;ttR!z32EFS^=>rwoII6C91}LV_3{p>c#uIB3&jew# zhyEX|H|a?T;%34wtT#_05_avy3K+ru zJV3ls&IPo!+VXw&h$9Hrxw%hZbB}X(Cnzc{t*o-6<#j-%aI?WBF%gtSxGaWppxByq z+{o?32UHW>Ob7!T8yg45VWO_GvJ!{}hNT=_T&h%5ruw#^59-!3wB37~f%HF~b2IX$ z48WJIuJ(#4B7M&P>;>3Uip0I-k3wDIxAE_zf|oRu~=I62!!xtBW3n2zDXSac`L z68Dv@^rw#@UAFlJp;s%l#GRHO24lHh**bKK>EUBd-?0K-`i&8kAf?|XAot?+_03|G zy1F{BF<@-b z#$J!`B8{7r;i9Pgbm12e6|dB&<=jQGWuU5giaom6e2vWO#(v$I%Y)!o1v|-t8ei*X zCEJ!Qd)!(!gB@tI4Y>t?YJ&h1k9h+E)DOrOprpQ-^FS5Q2J*+Y#I5uW$=exD6t5n; z9#WNAUE7)hq-)bwMbIG}2IP${f>2Utd?nLNkWfuaj4@Sfbu>D&oyn|N8!ST$LQW7ql@!p`pV7 zlEXl1P_2|Pc>TgJgi;da;X?~M!&}45<%rMLi~#P?9NLIePdihuBa*$1YJ+<|8Kg!P z=UD%<+0@Xlt!3!KfBQP%0KxY_awVqciN8MUVP)$3p|QJy6rOEmg;OS!C;2oivDMq# zIMF<&um$2;4dMFtyjA9m+|+f&^{*-Nbau#x?pxQpWkvAt34A4eb{?QKH?9M2W55tI zt{(FC7RQJ2mWq-2?d?xNDD{DcC#c?mEb*_%x_;znxAJFy^w?*tYAA?sAhAschSFSyv__(5wIZ%kbMx2Z5H-7Z&@6@!lSgf%d<4F8(%K z|D31aUF!CcB7fe2xFdnH{{7#-zc3U1zpC{8M5pg5S3OAAhxl$kX!_$n>@$~kfbs~J z$1JkPFxJV=_x$y(r0R|<^tVZoa9#WyD@y!c>i4rFU838wu>JTu622=vmIzm8H$#^f zw_US;7VGvvsj&0KT z)jZKUED_h3}DUcOVo;lGkFJ$n6!c1jBvMzF%IXq-oQD+Sd4jxfX98itw zOZsIT$6JzIJ8j9-_h(^@r2oS?vl16Fu)wfH_bqwwrs=o56Ls!6)DeFcIk{{cXjG6c z)clQ3h+-HW8<{xZcptI!)%`)i6j;H4 zX!&sy0Y*@6eD8UJ#$a8Tm_W+W9GZMfiQ>nI`u%Y*`utZB4y)F5+spTlZo|Yb>2Ewq zMUWj55Ru@})E7mGSp{6$e(GUF^WL$INNy?krTzQPv^Pa6r&Z<|Z3$1Krj`<1?oQ34|Bw9pNm&}~+4+aN#S%Yb&kHv>?S#163 zk(d(v2ka!gc+?k%7DNF%r_O#6dxFo`~ z!2OgQ{>*SYNmL&;0%L!(G#BRwI)TVAAd{C}9libAk>wK*7>|R z?HRcxAu%O+m8)cQg45k-kk@>Dd0$KxRw-!{PRqm8Za-Lo|0e~4Y|yOmx$wD$PSN<_ zrK(-ZZEvpu`9E9`eFb%M59I-EH;02CPg-#sFF};$xjAZ85q zCk?t!O8j5C_FoKN){=I^sD@3rwcb~9EuW)5RwekGa;G(&q5llFA`A>9t8uywAnJPA z@fM+tL-z=0y_Lg<+qpo@eZRAAOt$w=m%c{TjpP12xitCn=5OxjZ~xaz|4li6PVoPc z{Qew<{NS(A{=B8n2L`vEA{@U=VAyB4^(@8rcm#>23zaqc(mFe!hWpY~!uIA^bC2O9 zpJS1{ZT0D$+s`UieHlb#@XH^dkx;7Rwj3VHg4o+wvZFH+>ODvLq%s$J74(h5l73b$)a>>CMTeq4PZy0m9B#d8=(Pv)?c9|GmPT zD4Ucr5lKLK%z?gdJ=?ws6^YP}K4Gq#yWXjI7a9a|1GgDYmYot|Q?z=%s2a(V8FnFh zQ?wQJ9U5WszNIO{%Asa;=Gm4v`OO_SNZz?(T9;kr*VS!&Ns;!3VnW*&I^d494*ps! z)rX-i%F69StPs9CbMs)UhTN3&Q14o`mqv9!xmTmg0`6~^uo7LllI}K+jabo*qq*aj z({F^>D8;K*qn>d7JMIP>c3CiP*Q%%0u;q_x)~oq!>ObWNs@P(8M2>_OnTRXahB`hWF$v$vF|z)c@YQijKbjSkq`eKQ;KP1SSd zC@v_t{K}Mi&%*O~>xW&IK&p(vw$Z+f!d?*FO2i@5{$(?#VT<=k_vm2Pxh$x_GW380 zPE!eEUFi?bvKXRPkIN@LGI@@h7I3kAIxCnWZf=$jP@*NkHL#Qgi ze%Ccof3-$I$Z@lBj(H}#YFCViNXWYp$(jb?Rl?)LZ%ewMSk>t5)9P+*GpjZ+yuq{#q?`M41)6 zA6>=TYUL#fg^4!&d|5>?PXi4PeI_hjMN}5>9Z~B`)35Sr<5No@5UIWlSS+MJ0XB4s z<2Sire%)JcHv{h(urskWl9AHWFeylW_l5nU$z?%z$7R}>?qUCT+ei4nKa2+)^TL&u zH%NlbczA7A`oq$7HIl8$cjsN>8U>fmhU|FX@GkU`V^Hw+&6N>fsaiPStlx{5r5W2$ z*N+0F>&^FfHZNh}R`W zw#FrPsT+@Urr z4vu(gj;Fuz@fR`Hz>_2@S8h|0A-f)dRisI3=dbUB$*k4#xr6puwW`rbsrNuK@Y+{u zFgc8Z=rrn<(>&HS%7i`-ThW*})0_|iYF-JwWU!|=?Gi?~^TE9c)Enqa znBP6FI4pMazlaU?Aps=dsJ{hYiuZHvZ5+{m2x=A@(=`IA}!fnYs{dSA|?r>^2uD1BpEfh5m%qH zy1+9;PMU1x(iLvM_>;W{IoAj`qiGq%X7MoPF4H ztscMPdWY1^=yHCs#w%n-0(`F`xIR?=IScCgX6oKZW^Fj0OPApLhc_$N zNL3KSsc1>7N_ZmBRe52_=DD$;`Gjs$lQ*=!g#v>lNx|Sc^zwS|7k8t&^vsJ{%V-^y zV@663PPck1*k5=$uHM$l`{o|{T4ZPGXYuc=Dc|<|6~#A%H5pueeN2mW6bq7(jAc4* zD)MZiK%O{41-y2&?>Ui%A+><8KTTEbbI4*V<&cc-2(*93? zfj0lryu9Dp$o@?)Rv^}NV0c_!V{ReBgoF?*(&sHLznF$3@pI&h(zmsh>=|KjX0 zD=LdqaZqBVk(O)tS?j|g`mxr#kOsB9^Pd~2-e&L5n)5l%=FKX3h1SuFMNwuND6e@4 zzVuwm`eSa>)7G1B#g*Rxue(YU*eMBl>YKaysc|{0EL`q+sN8c^8`%@+{c3?}JDm%h zp)Z1w`7e2|ZSoEanApdgU5oV1Y>IC7D7$aaQ1?wlUuemz9CRzmn^?|Y^_Kok103^{ z0hM67Sv-fR;6{rgt|y(g#iiP` zC(t71U2iv~m_N*}z03-j6zzQck2qC&Z>gA%&L(d$qmgNK(#2MrVK%7tQbFe*or3;T zsr+Iz!(wP`r7hdG9T?|7L*-D3SzDNoFypcFbnR6M2Zb$o?#hZCJhU_ zf7IkTN1f3mw=Jh>*dGb$s4;5jU^FIfWgF~A%?Pgm|8vd%-*m0){pMPYnnWb`pg)$c|(k=)QT z8c29L`@{mPQ$CGjm!1DZw3iwIlimF1h_-*Gm6E6VC|Z(Z?N>Fa@47B0W!5w$e0NR!lXi$7gu)(}cm zu->%q-K%Ic`DoON>Z^-0crSgoY<^pPBa=_kuzx(($_4g8T04C)(z*#k_jlZ;fOma6-bf_J%lkM|)0_D!wr15nAbPtdd-T+CS}JNx2F5IfnwMI&ubeKEoirpvI%Bkz zZK7J><{whQ9O11q6tQjo3Z1PP1A%$v!&m_^>N$ z0j9&ly|ESjZPJ{MB{1Fj=vq0;L~?O`ue8~#VC0tp;xwtDg=lmAq(kYkl^sb18Q0qN zBEn8dm0ox?`dbr=kUC**lG@=fag%RcY^?gRZPxiD`zY!pP!`?Xw1-!ySj3L0v!Xqj zO6EnoH{8Y-*obI#9U}Uv=se;j^-Er|SSqNLrYj2|5_o%Ga|z|mBq&7V_3Q}{_5JqcqS^0p_+%Kk;CMN?Ch{Wb9NZ9t_Y=z*%FccTAq<*Y|z&rFeW?WCfK;f7A|_uCS8 zGvcbQpC+KV6q{cYyVEza=^Cu;x|Jb2mMJxF+P)U!aA;i<6o^>FjSXKPVrdUkN=&TU z1Z&N;;0%@|>QyXLEKR*eZkwcTs8tfO-FnP0d+EHg;Ty@b^^;OL-1zp9StD=M&)MGr zS>)~vz?u@v!C2`6U*w>tqI+Xd8YPa!1dDxO__*9Th*ddsQ2%pV)XE%K*tu~E&GKrF zw|4svhVF4seMps7vg+2U-fy9gdsx@{@F9sNyK>IMQNdo5ok<6PI$?wiWziRYz z)`i>WX46dX)5(u@Q!a1xel{zSQZNz3V3qxTH0@WDiW(o4vtydm&R}QbU;0(ih<9h0 z2rUS%wNulCo7sB>mk1K7?SdJZ9UOQ5DypV^`b=3^nBR-c9NqN6^mAc;{`T64XiKy6 z0A^7{;(HIspYmIp5%D2Drk(OFcd!E+4aLKCCe}@an5 zw_e6yi?y0L;ND?h`QhxN`E}kCvrBCxJn<*;X1orxNd6w=CG{>n>z*02=kE;Vt75XK zcUN{1eNJnai=C1VZFk<}(a;nye12u+-SFLkwHZk*mG!PiBn}20qMVVIH2f+o#a}1~ z8@BZ}NB9NC5%0%`kecXwQ6=y(%N``zOp|6guHrSDUQd2M8wH)da(*grhR+zO)QvuP zX~vDdaIIIVjSeUCGSTkrkKyfcE|P9J!wR1|p}|}!*7zRH9Wi5<#!_?xkvkvV?=tU4 z-nmzr*Zk8g_Ws@7-TpWw9bR)ts~cds{huWes#Ny@Lq zlMmoM&wU{y8GQv?*T8eWMjYH|;WNsC(XZpLUUzl(p+@Q^Uy zpdEjpz|?HHly0hFCcP~Q^RZvD&8&(6ymx;-`uLG+OoB&8e2r6VaSD-E?8xLnu9WC+ zv5$h8q6H*gR4V4@v`Dg$Cu6K@47sGpKBrRJmyNPDzV{$9B9ri1kF7`h1QKq7+%MaA zqqiHnB-Timr=~&=hPMe1p=gQ2%ZkLJr6$*D42In%&RmQ$7i9H%v)OLMF zel+=dhc;gZlHXkryc32%N-bQwsw9zOywoT)=u6Lj{tDS61|=YFz7a~@x@P#=$cl0^ z-!%M}xplq9xiCE{)!S!Ps3RK;W<+w@rzUp0-l>xGOm9W<=%0xag$&amMop+dar~AKhI<#6Hv( zFRh_bcw@-3KYIZLAqo1@k*20Pqfr$hUm6bl$2*=v#_I#&r=0`Y;&>a zNzMGv;%tQ5a^vGAU&-ytN36ACW_|e_06u;w_SN6AuJw}PASw2`AwHIooMvjIbwlN6 zDMcRJ^UDL_!vsRe%j&`7vwMRFG_eXN4bK{z#JQLw=Teo8Ud{CLD}{CVPS({(S$X3_ zn0NBC$A-`Y4uf7Ow|W}0rRpoZ?s@{nXsdHjgYEKMH zq$tdvJ70w<=*v-LM6z4sLAX+aHDMhw&Y$o1-Ht4dWk|fIoAIidbdB9G)W6JCC&)4g z)A5c7y5d9BRBy4D9go>~uyw7Ow6Ml-oKO1 zG?uAWj9LlT(Wf5<6wOvYkk&40%vH)7jIhk#@L3nLv{6e7POX2+{Q0c>tXtG$ZW)n* zV;|_VfxlTYL*IU`$P;l265nlk=SGNJyDUV~nJ>bkKATfrg92FZuoPup-gsteh&Vm>W5`!tp1op%sj={%JE~F^Yu*@W`|_!7n5>J*n@=Fk zUdgg>x$A2svL0d;72~SPNTsZJ`;-7X3rIrsa9(Z0y4wI6XRvA=xtdkNlSfmEn1jgF zkb=u8{ya|L>jBPOfm*6;xyymlIPGcdr`NAybdicx&CcX-VL6zWd+n^{J7=MuEKVvh zUdvJx-G#2-T}0Bfy*rP{KhQtDJ5jdxxYaX^PwGqI!WAb$sL-Qz4jyA;<(1XTcUMK( zHpl)@Gjn~oS>ZU`*@UoU+MAmbJP64SXCNy`Y&k#ssq`p!h)L!w2 zsnqN2ABE_N?K~RVW8apYvzPpjwQ;=~wKvE*?H(>U#_F|abT2ZFis0y9uyH&;DDCj= zwa;N%iOY8uj*;e-X3nOyD@nJVVOfbjQXKVCCyYKPZ~tM_`I0HhO5AfJN~D%#bnTtM zl5v3g!#X)GjMVYE(D+Wp#JW)Nppd=P7m7bt8X~fsp0WnKdSiD3m`p0`1Rs+o4DUnu z_$DHkXJZ#mBd%(E$!$s;i^WR^M5Vkj{gre-g|EYAcaR`g7|1CiZ(m3*98P>qX6#n6 zIeJGXc!yu~>(BCeg`YVo?`G&6eRfl_pfhbTzq? zFOCT!5F&k<+KDvqT+lpOo#MC4X~2Q*@M*VI44d46T%{g$rk*b@rpW;~TXFBjNMo{f zvAKRgJIyDjThery>CmdOEOMP$T_kO%le}T%m3;rsdXK`9ZZx<(NHV^d4Or*SRQ%C2@Ou#8}t>TWp4ti9C+ATeRSHsC50 zY&6`mowlXN`$&yVZ%t&w>X?(O^6cx#G{$sjaRly?4eT!>Q9XRJ>~XFtM(n!sOtyw1 ztcXS>Zkj0U16yX~5>gXOkAELtcfQDox}jnTxuzHx+T2PLCzm&>?z49)@iTlPja&#- zLSc@67OUa40G*m8GK6P(ieL|y3hl+IAthXuUrgoRk`U=dxtZC*oXBjMiO{pVm@~UI zbC-J0WX)J>n$S;&prJ@9^;I}ezs=$e)z#A<#q{HvZoqukkiwGW=~$*Q-ll$Y9$ATZ z_~yKSEuw@0Nl-Q!){tiS_Pu{YD`*S7!>V(F*EgMA+sM~Ok#*3Ixkx zL=x-hF>yAMJg8Az+QJz1E|Tfw9OF2(=7{ElO~LZyER!+~1d&M26G>}fy%&r3A>>#v zVs%rjI~VssFcQSCvWIc!vdQfX4;K#!VsvpYsbT(5=|g6QC|bgYyQyD&A)eXmr4i%l zt|K$E$br<4AojB<$Vt2?$kZ~?A17cJ*+WqWxlz8G?n8R7_1Z^8LU*2Z9&(TuQ8CG! zS6_ZiUn{d+6WmYAI!eRwvss=w2t)N|TBmHR_se3NkXum-j&!y6`Hgk%mt59*9pG$ zFg)qSFM6Y^^{N!HSSowM@I9jCjDz`V6>M{l=xsIAEHu;oVvO=&UKz3GCaA1FxtluOlS5?FgvP@=_%HI-Xg0xpOd=`?_Z^CKT|;= zya(CkW(jnZ2HgHC!u0-KVLtE=VV#MaD_d%9`=h(N{`%{Tfs8gmduf=E-K)BKc+pm( zVhkh!{t+D_VC#R!{d$sOZ6S(I0TPIqPX^*+`;LwKr$ZXA`+4?emtK?z4W}KfGVc)F zkE{+J`g{twf1YpOcU0O0X~+@%Q88XO__m+SiR+F12+GNwwx_)8If=?tdrcOMNNkF>~)q$}$o-RMNus!<-tZk{19)9TsJR-C{tqi zAUH5>{_EF8{DDQoM-Q2{HTYkVO<}Co8(%cYbJ}Bxy}%e{U!~N4H?t8C1 z7(Q%8oUPQB=;1-^<1dgtNhe?nhTXQ$kXLTDZ})eK9}*#tl|kGOw%^@RtihUKL4dU4 zs#skbTgTh2Ru$88+<}KyQUxb;i;^Ym_8t;)FKA^c|F6JdCspn~;%JGkB3qEYU#C#Z|4MU#y-Ql}}=Rgp@GO*uu^IW@cZ1)s$s5rW^j~@Ck?AgA2q}wvq ze#Zr6sv^XF+XqRvRHV2=`9z}&>~)+ z-r{6Q80VlGL3aa8i%IMd_oK(Y?hufYfphllWx7DIB1S2Xp`APnZlAZsOgZLC6t>QR`}ilHt!|8b8#wW+>mQlp?j*g?#r+Q*93Khm1B4tv<7M zKKU(p!d5NAz>!?o%jH-%!TbNy6qI;>D*)LmmQ!nM}I*(EfFe-yjWg=-T0qjr{exuXHj;Y{B z%W3^&_~6Q>2l8|Wxw!T-@zeNI5md-a&chCLDFJhrgKtlvUf@|D9>5J@(=sw4*@U?^ z#E`?CGjfRi{W8I;CdfC3nR^}R{@1fu$=3+8 zu6?f20Q;bRsVsTN2@-VZYXpcp1xBtRS<)V1%##}C?xSy3XS7ooyWH%dZG!2wR8>mI zkjszfww4p?B(v!0we{-<+zSxMteID^R(PGB7Xn}H46UHOb;t|Bdu22~mDBdlF(FXY zZ(3e=1`YXu*GYZb?sSAaxEk8JIx;wc$J|8|L`9O4gCOdIIyO}75c2A}DTIX5)6OyV zR7msiacr*>O6nDEZ8cLwzTj?Kko$v2ZLB}<4&wEAdpRNy)tGScJV!ueR>|!^2Wo}o zmGn%Dr#gU0*nO)WFTU3VkkErtzGsPS(_dCR5gF2U9!L5GRdYd9;ZE^y9e3y zmYRHjDTyp#>m(@gVHhKrN9^l~_t3&DZ^yBxwCWSM)f31n5(^Z%fr5Z=PO|~wbF(3Z ze7g>+YxE-9cXXY340%b96o-Eo(nH~^+5u@*QATj*#P``edibSc<5~-vJ23FR$C0!# z5Q=ff&FTx3E8v7wo$ScE2+kV~k$EvBp9}6ky^>hhAl5W!i9h7_in?N(m*Vu`hq8{T zD7z#%YAo8Z`YATR}7h!iV1#M~vQB;a|n$cWki1d=%R z{yQHYgu>PO5wH(NW}3F9O(X@&@p#{*SkR#{tr6CEp z)CPlQ@}ABZbMX$}0h^(^yWmcAT8_2?^CH=)LinRPbVF#m)9B(-mjmDb z7a(Me;XGvdZZEBTOvQueYjcj!VE^t97|3U8$bb763?LW3d=Wn7T-Zo|BCS(D$z(%l zV007P|CZ>{clzGe7jde50Sl6l5AFl#^cp)u0WHD#A7X-z8rO+^g#-tWoM$6FUyxA{ z-`1c=^9fob6V={+w%lMGiUXJRm_OliP_>?n?i@Yn!F;ASB*|}90pYW=L5JwGsXXwR ztMV|==?t*ZrBNzw6vz;U`p(QgAtE@7BUr8)TI$@tdl4+~+L?+sEZ@?_Kv^)4tWIv* zNA?7B~sm=>HL^r{j<=0uyIKSUk1_o0!5 zp9zWR0p#OZV!Y`4m){>jhSNF>int=6>|9qntn}~#etDCPa4nh0YhrP4+`Cv}Y}sM> z=Zo|l^j=BclQDIS5_+7a_(Ql6rpd4 z*{z01H&RHy`P_vFPHEH7etTAj<>2%S6eNQg?S4RVL{syn=Zd}ErLE#Ur4M)Z|Ln1x*E-totNbCZyzgU1h%$?3Q# zxYHZq_@Be+2GnPso<@t+jy`#S$iSW+^~FADOZ-WN=gf=p>EfTm{SbHi&|Gc7kVgKjgh+_#hpje<;-=S)O_4k!z1FK=D4waO5Frpb_GeFMkGPc(?^06g zXbBKTTns6O$)z$ZDUD%FMl5b{L^!)K$FZ zdA`RrjC)#tc(=~DJ55CVS_TMdDLs6tE|FBYcOap0NRZF8wT@KwAslP3;KFZyw0mm? zQXYYF8H;nk*iH@+?v8c1VtE}?RKHGn}U38#%)ewAHjZyWWFz(TDokTjYx z(093kc7%K=jhuv&au~N;E><*+E_@8HGcs#9pZZtC{!soad3S!3yY6Z7W6ZSGk08Q6 zNRU)r%NmWvXFTY8$3#Jt!cj>hh=e`rWJt2sSu%HDqQA|=CV}kC+e=<(HX68r1FLJ;40w`|=qKODXL-`>Cpr@=aSR?Ca!U9tDEP9_0ThC^QeY zBE43-hhCm9Xol0e&fwgE@cYGAlO#!zkFvjixkt0BdKXfm@7dC!b?piv%7`qD`q`-| zgqk_Gu&T?V(|$yZfA)*`AUD_3C`zPw&~glMs3+HhZ96imkGSsx1GSrMiUP6tY519a z5j{zm_QAr+Q$x`o52n%TvFAf@vHuTUZy6Tl^S%!+Aq~>qDc#cD-QC@Sbc2L6NH<8c zz=E{2NJ@8yv~+hpw|>6A|GVea9?Nm>GCOnCoaZ@nL5X?~e1EB+-MUtae!~wsQ_-`N z;|TRq{u~vO4Pf6gP{GecBw5Igk$77HW@b5Bxi)fP8P@y|gCzM!B0NQ*cc1UN6OM=W zD>2YPs@iXq^ZQiOt_EPk-@u0amQ!JY5B3hCUPK;x)5LSse);uUAtr^y*BS!e0lNis z2_0u|U}kW3{Ndt}!W7M!}>z&zDV+GL@3e5sV@*qao~DgkOxt-Dq$y>ZO0A@ zASe;cxVK@z?tLmn73yjHt;2i0#?wj(CTM8NHbKtTnu8*S0J}5~qXK`q&gx`(0&G}7 zIy=i*WjsTR1N!W|Wbr&DevO-Lo>%%3QE3j}UUW;6ofz(8h&}eCPRS^L8x9LqEAGz) z2Z9CmNaXpH8j-ZzcGCW)Q)a#BDhGB$RVh)kR#}>>C7#Rk987Wjq&HE-l(v|NgAey_ zK;&=kTNw#cV%&?o$Z+}CTfSj|tMQE*86Kyi7+ zsK_lk0V2Ps;)*9f%lYBv$aeXo{ar78QYk(au+}MQ{~?tOpAl4VkFBzl=D-6&T}Q(^ z*fpAc}x@nr4JO+2dwNkrXdn@76zg5*!;yDg9;*qr_mQf`;F^1OxHi zUB9zQcbr*0{5-pTshsRuHcgg;EhV)JF(kL&kEe?0g9za;nj3vy^C zPDpmRr<3L5_mK-`xeV@7%fWDZ%-SD-8un{w&=artw5029Qw&J-_f)#>fF>5FuPhET zAy7~5Foi}BKDLRSp2Bexjot*A$Y)_cyq$%A&6*w)xHr)RQrP0ig$6juBhfb|RaN61 z14^J}8J9upacasjUmnSv$!zhOWl=0rZ*=XLXe)sQQ-j0oUIw6)mX%2vMgakEm@Ex= z(HsUePjFy<2^oIjBWLMMv_D*6J%BvB3MEP!XSCm&9b1o|MDyKbhy7EqVC_+Mz! zuurpMX+T^;*+GDP@Em0x*xA{#)fp+xH4=8j`Z?8 z9*pQ6e@N(wayqYf)Z zgG^#L_Y&&IZCl_vXAfxNtS7a?BdR7`2v<@~Y_ej~*c=}8>ekp9V6XKs#dO#zexKxT zM5n=lgdCPQV{D0pLIMN`ecLeyHsw}zUcJH`y zcykGIig_fR2bYtL?qFT4!>6p7VuGdH#iYkVbx#`5HzA5lU87RD9T|%x;S9KC98KpY z`9tD&xCiR5#fHUHDby~&RgsbB=xP(i!-#k_lCdk4+{GDa-VfYpYXRGnQY&G>qhx=5) z=TIkuK{PT#ZfK$T665s8`+H8g=mKne(C?p|8UgGXAsVTaqL~Qtedti%PrlX}A!`6d zuet&h$SLn>wwRrS-Rc`v;*?H_99u>v<+Pk4Rii?N#RfrYCIPd0@_>}h;-ylY^3DLi zL`_1Lu+aL|zNfmxN1*}~1AQy9!xVZxVd6I+sY=@0M`kRB!E|+*upN3VlOdIzN>!7r zL`t1^CKZHI)w2KcR3BO>)YtaiLkl(>_+=ar$U0Nlr}yFhzCtV684*=#o{m(-nnZ+LP-oi*S69?ReuPeNq_JFb0he#S1+*E80aAsi zRIx!;%@gPH6^y>hE=}kqHM2FwN2)5qLUZ9!6UXY6HwO4%6bqWSDjZEAk7^rmT#SV` zb$zf8dWM#&mq$21IrV~pR+qs-O7iLhE@j6mb!gTt4V*hGcmU%mRTYX)dO}Tuwe25p7&?wst}0$EL{zxnri_U`~z9TS8qxK1_oXCyE))$aEBPAuUGGmHunTb3U#h>2v z58o@+T$D5}!iZr8!SD7QmQ1=tm)fX>v9ATEvR9sE&Rju1+^IW*tRd{j^~LepTy)UN z*0n4?lf=&&yQSYd5kTk_C`BZhWD;)>LA6CtAadCdJuBH~*`r`Q1XY%)C^6xaK|yBH zNA-&^`KywrgVvvqNsfTeyZZA&gO(ssyfQ%^68hchk=}F5DNRv&fpg1sUo({Kr+XZh zO}c)#DIQFIok1Q2RZw_8h)e(w;dU3WlLnzb0G(z_ zmXrOq6LH^wdKqnln=mChE6&EB#Dxi7CDju395_}MvZ^*_9ROoKMW#EZ|5v_}Pvgc= zK!Ik6uz+^F3S`#9rF^my98=Ns9NQT4R;YY+{{lo3L9tvS%J$5jE{J>$>+_LQZTl{{ zkgDRqb~%n1Rs{{cvz&P2RTUqco3z{&bx(aYneoslo>* zcu@cD^%U}Q$XkGU39nB{fl4xONQ>V#2+_2KDA28AY91J%(0cu?dhx`-2$aF6`x~*7 z&X(!vLICeyAx6_!4=Xou%>Xh@Pz+wqllwm{z>9Hlim*U19_VMVKb;z-MLHWE0H1&9 z(Jk{-_({6oEoaO%@DIZxkTAcm9s7}=6D-;YxMYFOl$UJ!zgW|{4i74pqrWyoOj2dy zWxQcrUP(Z@C@HM?A9^K!T|$iA67r~i{m!-T{4Z#ptAwAk6_`g(X?LaN?!1s^qXUkL24sY>2(O3YB(F){)9QC$maEGKk8H__wmm zyQ0FKmdjw!o>0Vsg&07^b$-J@z<*4z0E7zaD^}qff;T~<4%G(U8y%lD9PFf>Gn*61 z_b2WUZ^1exShyy4sXvn5LPvi7oC;$5c9_nWRyFg!s_lda= zIa^?1&XQy6wvFF#j^5^pQ>i^fiY=hS|cX*ii=tWK?7SG$}7GL63I`1^OfBKy})B zO=br~f_x4~wFAGvUT*~=J54U3PL`(m%9aJ?3F=exqQ_@-Pvq_|qg}=WZG5u?N^#X1 z7cB4Zx|DlRndkfj8Vn%O@0nMt#&(czh4}!aH}8XLqi_(Z2)3DaHi*2{jnW6YUi@vZ z-g1Dsea??zv@In_(_?!1e48u`WX+PsHc1Err88+EYW(J!=An{ER5VFINib!)Ne@JR zx9LtbdzU;5P%N62ilQJ9LWz%wULPaz5VITsmC%!}Vk4!I5*%n456?kFHQjUa&luur zDk{M1+|;n#+fQ}^s7CX&KL|J_0VX6MZ?gSV z*2UYW$n!0yV!g$u?CU36#R^bz#s|mu&xiiK(J*976n3L?tW4tVye>QOy|{L_tYXlF zPZMapI!`}E+LL#b!+`=FGM+=^>kz9~Yyr%c*xl+!bn*cZ2i3|x0wxByD4W!nNZ??) zqXvM_1!zI!bpsCVR-~&{AAH@JRlnwlSiMUR-W`~i0c@PveyUv7m`mm_6|go61+A`> zNXP(58@t25pklX*s6NRawutbr@F?ZaOx2?T5LL&aj$(QTZPm zkgNDBPez$XxLvzrap|ro^(%Fzl@01Vx8QgA#-Txv2x;MYM9}rRLBX17{A~nDJ|6x60uUoz?=yE}0m;YtG=){f*j1E|dr|b%^pv-mdl%~|ZK@CYk0l+Z*wj%mdK%# zuVaUAXv@j+a9C}|NwhRT*zl#{uZ&_=4)cKP#R#!glysSb7~XAHGQ5uPY!AboZiIbI zt7B&#&mHi5$l#V%SD-S!dhxC6c=vyCxPIVy2rc&}8V+E>K4@#qU3F(WX2JMz$nt=h z^Nqgp4{Z6|xi=~70dX0b&obqc9XI@1pDSGmKy9a$jiYv*sI$yYGTT#sjfyC<+W~e* zX?X9G2!~!rzTP&6f+ikTe2UdDyzLe)ZAAc{w*FHc$Hu($s)%sv@VKU^e4&Y%2__%A z^K*q`0Ae!La~qk_{YJk2ax>^jVEl8{v4O3$3_aJX^GmvWNQu+dLBbQm0U3Mq2M}nl zwrkh;4u~(MDf%*dmpu`IZ$XFELGlG0ya~@%F6%WQ(_@b=4f`sK)mPq*%kIKevQmC9 zK#l?!qCd>&N^e67$9&a=6ExFVezKgIu| zdP}TF#awoFzV7&28BHPx7;p4r6=y=8bN?4zX%Q$;BG#4!MxP)mIYMWW**Z>Ut+H9j zur|l3c}wRP*H6_unQj)j!n1@}nMt17V=ZJj+X_-zY5@5JFr%r53eya(Z*EH})V8Cy z6=sDHvc2$g9^}=G zSxMwy-hh4|;b4|z>cS?L_b0KAJ+&@~aWz6(tRF1b(q^QS{* z{E4PVB1_)AOQhWi1!2Om?f?*u*d`q!-scfR2S<^4UVcbZ)c_PEy)gjM41$$kfB_k* zL{6w_uo6Bzz9awpi5h#wW~hQ+ul4Jimd$v5l2gTB=FU`-YFRcPiaz0e#-QVHP#QuL znkSJ|`}8$(oBI8$8yYD1JO1LKfX0NLb!_1M7yRASWBs0M=7!vjfrUvw%*K_wgjOlC zGTZ42o(whF-w!>YuWP1opz*Ee#9S;`sdAp>KeXI_gYY2RRvZf-Jf8F957gvCg_d&} z!Jre#iRaamZNz>DAX?X1h&SOYqHYA31EfdnSGiTKxGab#dE}4bc*|y`C80Duvy`8Q zQ6@Q?$_g6yWzJ|T>-=L0wZ(dN-%F)hbR)bcFI8)<=*Eiz%q71b;?K%k_`gfi#>YgT zC(81jJX5a0-$X<$>8Q%HH=T0&LMrevD;uffJrbyuaVnh;n^sq^C6fV(4V+>Fg%U3gJekxjes8I?H5=?__ViU{(Myc{ikDNJ1uLuPHy;ir?@T2xD6Jb2YYQK=)5s)#eYa0>rh|iHh#*xKdBx1QTm-q{ zPj~q$>$a>=Hzh;2C4vX7A@}MZJb-ZP=rwpbEG+t?ikp2H)zAxlVjDu+ieb~vHR&wN z54q}m;>YT9P#*|d8Wt%0O75Y2FI<8L~_idq;FeEdC@?fWpDrx#$oh{*)W;k4$h08<2;#^ntYVD@`Y*? zRzmvLXCP>SQ>lt&y?7VyQ)EK|I+qv^I58MHsJfV>69}u>Vg{UksCD(ugutly_#-o~8 zr>16V4Sr`z*^o|cCl`b=y)XTnIZEF}i9K1^C@0F;eScK^tU`bA(Ca%n^6UIAdP(f~6}lJrM3ba_{PGNs z!@ybOjMkE1{Zxe_??fUt;rNbSzesZ#LB9_UE7vJ}NRj@!_Q&s;w{s%Tahb-*8oH&z z#el|CDc>=M3OGFN4D7jp!6NwYu7n3 zGsCq{RdBCW4)UiC7^ebu$lJY^zmt%_)Ts&$wH6fUZfePx=+p)J7GDoD2BA#sYa8Si zpH^C67XrRv!#R5f))&fUzUJg%jmU_N}I0i+Xvnw{_IEV*}p$-rXq^K3x zYAFyok76$sYns?%NSjYeKw*niQdR0#Y{3;R&WV@0bFG$IC%l=Cci=Gqk+)M_KEVGa ziLfqToH@lX$J)^9I?O<}4}F9?>d(Eb%HhT&7TSLnEDE>;fDI4JbnO7A;AdtPr>{hs z67$DtiMDdwC2O{D+L1@sB}&pMfKH^XH}`5jwVz-M4McOU&WJ38CKYs3VXH|BKuEKK5PFTpm01XX{@+ zI=+7zNBvdP&(91{zkG8ZQr9&1f2m*0id492R>4MQRmjvpJgkOFjK`bg#uEDyv0VdZ zW*~5-9@|?-8OKHetEJzTY?3_W{s-IXV9+mb|CWXvan@eQ$?WH;a&c1^6qKhe$xohB70DDUeRFT!@<7LO8%z= zDRR`|1mb59HDr}O=#?pH7uaTN@MnEcuU##0BpWH`L+2{+>tYs0S29{=B5LXNXbB@; zO!-924Zl>XneVU(SGYeJ>1dFnOccGi`a`3*(>baWUU(Qx)!lCu_ZylrFG2pb z+pM`JLBX|D;9%w#?4tLhZ-F~GFSWt%nPod^WW(pCFbGB7UxyB}yw(<&-H)8Apv^za z6h1!jVQhzv9!T=2=Y@7TfJCOD9xKiNPlh^nX>0PEre!hXFP=Vx;-sejerHD4 zH?4(Cs?60N{F_03GJxJBwMIOEzKBxC(WC-es=E9i;O7(8&9}yjijPu??ovht7HZ|% z?AbTdCbatmEfRQRKn9;_GiFgod8Ri0TVU~pFu`Z5cBHbdxvB-?Gm5H_0B<;u(B;XB zC+u+V&cdH(QJaxO_4;LKaPjh`E=p5FBSGNB%EX_#-;XzN*}HUH>Gd}q&W)xw8Qxts za-RY`dP`h-MR3*f4u zIqFdA?gbdWzgG{uQ$;uzFXgAhBhjPWQlcQe4_tPBnZLgl(cb}t%wQp6Y&?rxf(xlH zksNb@6|-=yVtomcJLx>=MdG{)Z!#d3QI3rtvsk`p*AC8Xq#}!qy)kjKd8||vn7X^1 z9hpMO?gr@kFS||FkW4A+e8Ch=v#G&V5zaNutFQS&N5GDdkkl}9ZgunqABjSEOL}Uz zZ3!_Np)32^ucwm>hrlGOshY8Xb+n=%)EWeXKWMO&-y%ao0~kLLR8~ z(i!&aO<3>OO3q~ysZYZwd53+%aoPMEo&oGm|jnhJT zuU^9F+IRC$uStb6cHaLHP_~D5=SbzCo*91r%}YAkn4AFSkR*dIr0hA*R7uVvQZ;lB zEU4c$x(BIP=B&oelU=gC4Jm`pZNOn3ZBj`C8H?10IvB_pvitX}hZQ znmNGv1w@TW^2Bk$Nc5h=cJrYu`L}}hVyfn8hXC`5CDDF~9O24K68bA@Nz)o+x|Tmq z*T{vA9=NKN=O?5^ypaBl6$j}^VIh48HAgt68ymSd3aF`s4O_Kt*)@7)aP zUXJ+B_gd-@G~_p7)B?Hs^D?aIBR7_B##5!_EEZPEpk%Xl!bI_$+$ap(rDtF z5_8q^Uco|(>on+0O+#hUC1ZOmw#KoS=SPSM>$KqlQ1rcPM5M8}Jr&Wl$8s%v*>_G^ zf<=J~rDr)n?SD9hFgRu&u?HLB0lPN^8a2R$ zW^xh+wMOfiF*V}M_WVnJz?s45eLc_n1XR>v;3J1GRFeVTN`}C*|9wP*slpUJEvKH* z?B5ss?|c6LUT7j%8lmtNX(l7C>RGpQf*P@Z;wpxuHRiK4BJvejEXh0lydw5cAtqov z>x`>B7-2k~`*?OSE@_-`FiI82Ng5I=1(*##ZK-agZvY2UB)W&Iuz`-L zmN|>++>-o77L6G=Yiocf&({0jRo8NBqOfrx2Gg==NtVku9}wW8QSgsEOnwUFxm+x zv$@qj)ApOr4c_A}>^V_9kagsSGuPuQ{z=Kk&C{`GETwBr0@{Vhnhjef+V3mjuswYS zBNz+iGdh04?x8~`z8Hy%t86M8o7+LQ>JY+<@d+apu8b@|;Zms-gh5O%h z@7$D|+#ft-pU^vITCm)TNFy80&vG|7Rn$y&j>L?6J9qtJa{HMx`^8Pgc`tvirtd{> zPVq2bYee_$t-aT$$8SgyUpVfN46H{V=q~P?gvj*0G!($7^r+uon3x>tgxbq~fRBfD z=D4dYfrr@t2wcs3I0)-3aM1a_B}2jC#d+KezcM#^)4TQtp5pG&B!DN-()uI$#l3FV z`hqr;iiwA*t?ma~U1NJjc_0DzT)lHc z5=xJ-5y^>3+PsY1ZgL}n)}aqr({D}(2CAd|fLK=V%y4-$xl}TfX9FFvA;q1sN~9+u zzWDnThM~@Hpu=e`%AF$xDXWnH)iQ_4^5bykroU@h|ElUf30s0Bqz}qf((rm_=<3a?0<;JRh3$8Vy|0P{D8Bz zrGyAFIo0PzB~8#&1K*Mbyo-OdnEVYBfP1&FZLV+&UMPbhKTJ8?Hm?)0UR~0bkVXtS zxh~?}+`WLm57s&aZOh^r0QO@uSnD*>7&OKb0i4SC-#6eZwxA=$eERbJ)unjmvm(s! z!t{of$uV=OF*^6bOu4O-Xnh0Y0;^C5iw{hk8SO{q&rvBX9~N7XPJfIp!kP$b-* zO%Oj|0`iytLM1Qp7%K&A^(I3*WWCN|h0 z&tq0EDKI{-^jb2Uh5X-n_j~&p%Jj6`HQzG?ti+$t3uXFRpno;Kk>~_s7dpbBYg&jol2;WUm*Oj{*go?HQCby-Cj0MTzr! z16t_l!YxDkZx-n(u-3MwknIZKQA(@+EbpfoqnzZwYh=!7v6yMp@!dGjV*=ic5DNm@ z1vNYN?F|u-I*OcS;HS$Iz0%#-h z&s&6&_wb>7$Yn8jQqZ7&TJNUzGASJ5{4SL&94r)61o8X+hKvgv|3U@djcd=*#Sq2} zf{&XcD3wb&_Tm4=R2qgNu$F4fdxmD0aZz-5I>yN=E$@VEeS$AsAC=iqMo@~f==P3% zxMQTOO+W_AX(iHaV8sxlY#iv+fKuv7@kzG;!rF-QceUhxNIVMSn&<#249~4F@7!;4v3RVx!6+Nos`}|3j8}Zj!dzWlwj3J-GD8VZ;89* z&ru>vPGcxND;ja8VdS}_iXN-2noAD7Q|rLA7VuE2lDr9wPk?uY7GK?=Xkm3__g^=$ zPHW(hW`xlH9Y|@&$GGDZWABcEGfhUt7WQurhQ5=x@ohl0yNeJ8+tp>@D|FBwO0t76 zUX(c5%XRu3G35140)4@ym!gC8T^I5P%tQkvx%zyg`Q~UJJL0v;^WvKzc z=!@tpz7ODYqEmX)sa*|@2v@WqZS{QM0B=2$oKP*Q33H{?dkK?!p zEx%*)NLT*7)A^780t(nD@O0%eic00FXWHoRe4@a?Fww;;<)pGM`!o>yeF#AZ+B<0^ zd^dblumOH^RA2i?mJDA@$0wKlk#0+JW*SPTKu-@W1hkT!Vc&K-;2~NVn4a?)f4wl% z2{p!-DEc7fo4lPkrigwzdzv>Q5bnKu7p~|seX|QQ_oFIfxIBBU&&Aie4o(#~KqZS> zHMkyl^{Y+g1;~1)Y#H)>8!MhS27?_(1mGnJi++e6R2$cW$aIatMF zX<+y_bX;$yS%<3g5u*sXHaKwU2agg8I~w_XYU~jOldz9(^_;eQK7Xad3BkwB3f>Gz~|8gy`Rl>)R$wsRIFWp%w!EsL{-$Hi=yM zvrz(C7xQm!1~&_P%-6io4IG6G%uLR*SYEFZl8>>D&3Zc6PvQ$Am)7eAgDF(CO}JK3 zBlJSBV*Y`Bbvy5o5qsql_H1NYtxpyHJ~aQy{Z()V$M1EnF6J=0`|NDlPR7szercJ7 z`7YdkQN%$qN=YH=i!WtgBu@RlU54Fa09!}(;r>?IymVD$H#SbT{C!I*;~s%~qW6SZ@paHQBZ~CpHU*yzUy+KT5apqbe(uQ_|iiZ1yX{9jDS8?5(iHhoOn-#vCb;sk&MSejkH-EAepSr0~|7j%@4l zU+y-aVP0hn?S?`|y{e3X)7N24_SZaP0(%=0GR9;q>%Yt(vFDEp)kTO%P<0_E(%1LL z_K%^Vudu-}j>Gs#orA3@AOI-?Lk?q|@{E`VNk6z3sdK+M7wvbIk3k$e>_R zhcf4u2=hcHzr8rP9sLKAj2haI=rOC@CQ|u9#J!dZFl%a-qjc-}n z&W8u`YrrUdb>{Y2Qm<@u2$6+I>uq_cJw5w93xTh4VUi*5zCCCRzvxjNPo51L*vYJ5 zNiB~5*|2KXP0^PFS6IX{x{rF1yh^aMXa*edWvmUklQRDpiHkPv%@i(K=4-M*^>6x( z_wyx_2Oh1=Z|jppqC*Ww%3#jaW$$B?{;d74bi#(L)c1#F>p90muz*SJ%IdTRove{T z60h`bdpjGbB*vPf6NL@xficK4iJR-oOD2};|N3Kx?*j{knoxYJ>Pd#Ob*1+=m)UcR zw7A5CP-BR{Z+zH->+vtzfQKLEuNN^W@-Jw#_*5F0TM$Pn3$oQR7-Mf*_quokj|U2l zD-Hjs1_U}x?x{`-w(y^ahf*55V=Z*?T<>#C7LbiT=skevkP$(!gM|}oZNH-j;MO0? zn#;cF(5s{wL)LwpQ4vHD8s0xdzkXP_vG5vaw4xhJ-w1x0=tjA3qml_=c3>`HE4$K+ zW*x=6^|o2U0Fr6Y4_impqj38c<18AP_>ve8L*kYd|BjA!Oc@!Y{a~8Ae2&7L;tSa;;8P`M3b)9E0Nb_8YJ6|aQ(PAhDIZN#TRtgb1=>@La+&`@I$T6 zy7#>%D`=(F(GGl^wB>fmQh-=7E#fj=Fc38P`Z==qas(ODKcwReM6^xWtWd(j zNbD<2rRCJpBgP{{2?sfw)u;p3#d(Zz;XAECDalTXgr~dG-ts~mqkgzW)l70f`b=kn zaG@5X?iX&Q2RxmM08L1^F@@HzyM3)?KzpkR`w3a*Gw_}-D3PE=?=njggl1#JzVI=*$u(SmeVG+Mp4GJhQ*N>GWcHDBt6f~PBf zbi~0#6OdvT^1ptJr%HJVA4viNDxsQrz)PnhJ;wf4V!e&N5NdLq#`FFsbEUs&H+mp_ zu3G!cHdOj!QBC{H>kSz;`qIpdtw!o8Df0CDAxd$PgMjNfsxCf~kBqVhV{$h(F5Q97=#;H1O9!}oRlS|2W;?eO}-sIap9ua%KTPz9_KfLvm6y|YxA z=N0O04sPf{6FyZ(bK<0gzkVxdxY;A4{oU+X$2hUWfCK`Cp+Eu7((R!i$-h4T{OKGY z+&x0Dr^_rKy$QQ9O93zbrx6X#uTFa9K;L_{v5Mq5&1X@iK^cE`1UrZi@z$TuVL>w++?;GX-7>q)Frq=^ z3q&Z7d@qSTGEk<8i(J_pBDN+lpwP*v&Ww51-aziF_aLU3Ih-;35tNLxGvaiLEjkOY z>aZ}XqQWm`Jjt<1JC5a6uT_aRgNKJOgTF7wCRf8dCU7s!K@0B#b8zlEGrZc$h(PbN zaItR~wB-6z5J0?BTJZ*iIBAIEwXp*O{YdSi95>_be1Q!0WO*a5?FTL8Ed-!7d_wa# zwLJKNeuGe;?Cb*@4cfmCrAYal5|szO0{cNlp-A+L`$r_dg8%G8MtVKlR~36PM5QrH zTcyX+!nw55K2GtkeGEYZI$6>~el*Z&Y&>6x7__Nd!2JOQU>dCawBZDLdS%#(6Lg0m~d!!poN`I`oe8LVq(FodwRed#Axd<&0|w zmx3gQFFZ%A*HLcWtFax1l<@EIZ)s-}DILVf%07_pm~Ktt*hFdsP;kO?ogROoDi+A} zgGR~%2Z}W=6WbpsQ8~4s??;&X2uewyZxdfnFW$&I8NtP-L<5APNrT2k&-hvXlgG=g}B`nt`OWU2SwXRTF7hqd6CXYa2rPLYvJk;9!#o*|L8T7(!;D4V!5w})Ib z;*w2Aj#-=zk*`n+@&fr6oyA9$iBv2%Tc67F_Z%In6CN^`Q#jSu_cpG&T#UX&G6Bsn zU&<{p5Y;QCsO8s-$O2}@-I}P<`aO2RC7B0b+~H1~R#{yoe>1~YmhV4Kew^Z4NhUo^s81<*O}e)Jwv6iMF^MD;WF znIubP-A}zKmoT&QhxXyccW#D8V@P9*T%$7<@1EBW)*NP~(WasJouMLavjVvnepPwd zhDg}T#x*qiIC&~cdI~k4Ifft-E7NFFEGA0*+@)nP$XvN3wY{S1Z1-4x^C=$8+2=rs z@X@t}U$G%`gfberLR_wqw}|dsx(DiSxw&%Ay!kU-2bT;_T)7>%0*8kXwmrcS@#*~h zBQ8{wu7x8Tu5#li&5y|~FScb{W4eU#zJe@T&bxA)4@?tO6p4^^m4xM@r#)NO#aXef^kB}-IblgPK8il9d%23cKOGMwI;oi}9>BUFC+(_7S zbai;K&vTW~&Y{+oUV4>wqip2(4Uly`=eC*OE;EaMT!C7V(q8uY5GCM)jvokgEpp-F zodg8Bt`|%{^5{Ja&O523qz0|pwTx;XyKvo~i+(Z`f!X%g)L!%PHJ7kRT(0q&oJO{PfM=jfSR?14ws+amfKQFkx`Gx!2!W*Mr28RO1lkj^f(XmrbKo z6-bcA?j&rYTv~iZ8X?F$+)m!@$$Q$L zUbI#n$p}~sh8Xi!TLw%Xt30@Xf@A4}@J;q6>dO+{Sllp8UrDZDVioTfo9r};Kg*H) zd!7P^ME=U{23qwtx|>WvB#d7UZOC&Ha#$*BQV=k}5CLo~p-Gd0x>n+ir7!a;jw_n$ zKH_irXQkxBBC*qX)urKXog@RwT3TBBW__k&2ezu6c}X$yR`!?NM4ttDaI}=m61X05 z%gXAO2=K66Rs-g6ZNgaYf^JUMbvK&b;IA)Nut`q+xsXfj8o63y%_MI0=9XG}I<*&% zX5-WPe}82I#`NF1dM09;*)e#w8O#J`rEXR|Oi!%62gF|Al=7`1b)YaJ)vK?75%P~H-+Kgz5+wU1wmx!hjPI#J7BQPWyyBs%+%FcPQg zYOuD(FXul&w*ZSy(iR}R*W#4U?KC1vK+sxMcM0~EEN*bu(A>UT6W}TuhZfrW7&E<# zkxFecku{5vO6dQhVy@}m%$9KsfnX&fbQouDib?J!;5WVV2?MU~RZBp1&SsLU)J{mS zHJB;a+@kvV$ZOi{ee|JB-m>*$iaLirna}MR^32qnj?e&?z7IRIc|&9hcXnIXtFDu+ ze_NYjV%C@5-%aDTQ-31hW$w1zuIh(i-eW{5IJXIA`Ixx-YabSBmqjZ(C}630GOj4HB||=u*HkM@_Gc=tB?)j&(MmQ2UYlRHhJ9%=BPw+7?m=znkf#dPn5lU|=D=soIDJTaNo)%a2tenvk?aElf-0hD;srIDv zvRLhG=S%Vyyvi#nFq=&I-HAWAP z^uk)Hm8rc_`&+tGS#RCJwK zkky1WA)EWFoZ#K_Sb*blb+!KBvybpVD?YTmd>f&8(~_6?HXp*m!+nQPts1y{$bH-? z=5v$ObUC@v9r>RIX8QeubSb>3C7VOP$LY^5fN4|*mEh*%^G|<+4j9=wUoc`x9US?N zUqU^L@B7+|l3LMdG{lyBgGJthy`7M1%oUlRuf&zXA<^~&yo!&WY?51nm>H?$5>HH}(Y%g^3b z)ypNI78iba?aH?5wLM)5m@>}Ajc>e19-VV(TAspekS^F27wq8m^s_)ap^}Ar%K-m`^S8$ofpRJ~#D3cbiCnV z?O6Ca+nFD9bxie!@}i#x?=w7nkbl)#&|4CWL><;VES$pEP48c)xA$)EUQQc3?N7DN zlTQh_z05~9*dT$fz3(9MZZrD#r%QrXp=2~Bu6pjJ27+5XNs`-fQzvV832%`eCHmti zFaB(B_j%UK6$SbH9A!v+9eUlPa40c`Jf9-+R7zeShIj}Ns$A{L9C@^@L&iHNu?ZnY zSFY;!!pnnpaMOcdXJW+H^;``l(_RBCzln8RuC}hsk5Jb6k0JzwS)Z;SHot(iB|msO zd)qeLoOLm+3+k+{=fg-$`wc{=#VW-iO}={G#P)DE@}1ZA{QPN6N3&M9Q5lB*{5sO1 zDoOg}t6^rLVd*k)m~LE!dB50ffR=SA!F9mH(xF?T{%r1W(sVF=uqVQB%M_iDNr3%u zv>+uz$XNGNhI*gMMiOY3_51k8{pMTrWzRg%!}BS=0-|PJlK8pml=dRD1eu4Y8yPy~ zoV=Ngth(a6J33_@2V7aL(GUOmFWz3!H_2nv(|Gj6Tb-0S61P88vpmntdYTs-0amH- zM`9Ee*=s7=V9V8m=%i7Y9r&O!gz;4ij*mSC2oCDAv{a#te46@XsBVWCyZ; z_{F7hkRX+nb7PkXE+EV`Cb_CC3_`tcoxw6tR(aKTiTb#jj(ZDcDb4lEPB(Tsc*+em z`L$n^5lavdAhtnxDuK3xlg%-e_IRiZ)Q$Bet3DbfNsN1hBF{fJ1JJCNl@D?72m(5N zB7Lo1yVnib9IVwh84^oz_hO5~zJ2|Q4!Hr7v{fs}b@ExX_v}1$+Ym`wc1Lc!CZ;ZK zdYUw!tW&zGnLkM+*Bg$6;CbNt+M-0rh}NXfy<@Z5xf$17C5 z<7c8ly4p^6uiiWmaM?MZOPU`rEipUx?1WfH*r+(_>Abk;x$L4)?XjCO5r_nt*U}8w z3xKHGscTdV_#O#4%gOypmqX-y_NWrA&NsU1wO&PM53ov^J_)o4tc@&Y<;V*Pie29K zIliF8kn7-bKMCq@G0NkqX7X%C`Lo=l)#%j4)Ib>Edhs*F;?z$;ImPj6OYx1-RT(g+ z?qhAm3%Wf#_lX$I3b9(z~ zLLAIT_5QeQhlF^J-q^Pa7m+p&0cxrB`@d$wflz=1(W<-<0s5fOU)KMW$+4{t*dBbO zyx7`SD9u!7Jifjq39KQCAar5c89VUz8_`{j3nBr|Z(1Db8f{W{3{Zs=Wmjocf0|}N zZ6JEsi8x%Q8$`>--2ZT%a_U4fF?4*TNmAci3a6c~HSb|(yW(F{&9@N9PU~COW4t_} zTg=37^xnZOm4bSXb>N`Y2q_1-tls0@YeU&0ozm1S$2(QG%bj%qf7 zWJ%~E^_>OZT3W>Q+!qdy&|+A=H!q>*^+WsG54_HiJlCl~v`Sb}$b{%LHS118;aqLl z7lb}0Jq*P3%a^OnUn$P_e-UuY zG~X6GmoI0Kx9wa)`KBL3VQ=5Q-E%k&;;^uU z2gs4tA)8s!#nuisA;Py}BIC<>&3?MC%BL_@P0d@su`oI=L^*nf>eQDKV-n|%J(gR?) znwcJ54aiSincopYA*@&wPy<{%>?Qv;W_;F31=@4r^(pt=_N%mZPe-a%f#xsH44+W=@_o75x=!cs9h~}Wy;=1 z+_KS(Lk`=xCn^in@a$Mgl7g$@SL_kVriC0b38$&a`IQvUvKzI<;fh=jxqwHHil<>zP4!PQ&u9&k(Xs~%?C6Pd+6Y?{v8#S=^;U7AXPT$lRAXdq8c{>Gu)ei)^8PzlTX*-=QRs7R zGq2YQjp>hKyU`&QNLyO#1w!h&J0vsLe9TZD)z0e2ArIZ(N!ji8W4wQ~l;hvGBeEA$ zS)Rp=AYOBq(KH+LY9JHoOovykXIs1$N3tIE`1^X8W1LR|Sl~q`OIOTzG}DyLsC)v9 z9e5nr$|=TAkWC1C)Ac=9Ehdv)(Q262NT`u(Q)S|dvP>E^H36H_kwj%PagsbTF1dSm z({cIV3fQvy#TLHJkt2*sO-V8&Q%bHh$B*28hLX?|oU=hm50!OZM5hC)BJ&o5mfB#xOBI4u~F3T?ml8m&T_HJlyR&_3-}ASGH@xyn5EqZ7*; z+&5l~K#h7@=WKH`&3**4D|6Q9Y`O zpA18m=JvXUkGef_Z~i%1{N&M+AIY>(SthgKJ&4!q@*zSVUIwcp)!{Vadp$Lj9yZd+ zM08SNT=C9HV$SYcxqT?Ys4)JD-uhCvO%oKE(ZY7yPB9|nnM>B^HCOvZ@)&l3>{?R3 z)rqd%nDG2&@*~e)9!$BM$MpQL@0Pr)&HYs)>Bk)v@Ix`ww|h4kn6%t!toIvxiTJU# zdscgWaZGcqe>Ry0ZHNP zF==V`OcZpbBX0{FOXF2P{y?$K8g96{5BD-X?OzD=rc_2fGAiXSE>6qZ8d9ij=MEh= z$q`_Kx&txh)#MhORoRD>QQU0IkhB>d9SwTOWYaBxf=^CP9@?x@v9cVlCyt6*TUjbvYZz&Eii{8&29Y+BSab4p|H=&*r(KQf4n!N-?_=p={8S~ zFPk-Vj~PwEW|I%AoTRVWKZS2^kLbfX89c!AN!HqBJrp-ZkMm2Ujt6_n8Tdk7hQ|so zIj&C;^BU zHk>k2oWe38T_7dh(HmH2So-ud0BmF`>-enuPNY$uxo184HV7c)j5@3=>rU~MOoUg+ zgJRVb(x36`=dB&KGz?HW?wpV&r14aZPJ0yDz^8b2o!uhxZQzZzs~At6c;qF|qt%#e zbOFtf!^RP@gREeG=386N-Qm;hM|GlS{f(KJma`subGR|A|5gl(otkW(v<<;@B0MZ@i+s`{GAr!AWw_00SgF*xO8-laZTC0ZMa| zj*Pc7C8;Yttl*^eGaqki3+k=x_iANSsomDo2Ed7~8Z67}O|Bg$-h}+K56q%^(56b0P|x&T*P$WOv4~^wj1tlItT;9-yvJ!y6R1?bn#_rga!SW^RrCyNEWw{vUj10DXLNe| zK7Pj3T%8*seq`g{FOL2h*aeDO&7cxcLFdY2FI zgI^y1JV?R%sGl7%57<>`A)|R{AuQ<+bG!2o_|5MGThv4!^SPHN7qiJfWiT|A8VRcP z>IrU?z_$$M;@EdifUC>?E`pA6Qs~`*7p# zIE#!xZ%27t9^q?SBOfh$Aw>JuD(#QFZgfGt+1=pvn%qkv{H|AnT1zs@^YxrPywb3` z9~Mge+XF`<){n9iL}H{n=Q#~|)+J-8jz zIXZz~4bu+No+zpvU#PS=Q+KEsjQJyZZRt&oP)6&N33_d20d+FXFS71R_~B{Sx}~g* zk=Mr`l+rFL?QswAy`Fm5?2{vG5{kyN)gBYXW)m_HeFh4=?{boIErTS6~Qp(7@$-68Q>L}AWfB~<^!d1+<(ZX%aKr1fl7EqBU^ z17~K^V%p?E@DSJ`*9ta(l0NQcIdGdBjwBYE*DotcOG}`_&{Jc@pb&DZc?p-7Dc4Ql zG;9mmOevARu(XxR&MC!w_9qYn18x5|61($jzK#j~My!@=A^p_;LQUrU%Uid0d1s!Q zw_4`F^o>>{>Df(KGxDkpwGM??1cl~xPG4@+@NyVDhzVmVFJV@8vaJb!#mCaZNXEZ* z13xRfV}11D2v|k8&ZdkxO?m#M# zci3F?3t{HnX(IS9D@rS;;-m4I2h>9q@b^NtwT%5uH*Vz!D*O0OqE*rXu>DvtU+#WISuCzrUEjYJ2E%&!ew72 zXnyS=Lj#d)tp8uXaYkSrom(^@`^sWepNIKP&4a&a^FEPHEH_PEK;-LvVr znAoSPQ<+L8gYDF7US5sdmS5u?>JWMYN{6rP_dXm}uARd0y}UYybxmm4qXwv?U0%An zke0(>nal&0Hsxc1k-03RWd5uHE0Iqw#{wbe+olB!PgS!rb1ru+eN&`Sb&-UYOqJMY zl+_MIJ>*E`R#a0XTHlZUwoE;KZ)3#xV~>if>8<9)?d<@`?(n3i*4X=oFIhM$qjS+~ z64-4s^0H_aREAPYwTbAbM)b7%cO~SE-(RDLW4L{^9n##I3{m#+v)7bG>>d|DI91$1 z(834@cdmB5b*!{+5HywW!Faz8YUR@%dk{5eh335KEqssJdEhd_5u2Tl{@OV_Q1szS zZMIn$?^^SU$*7}7&P{fp$}{DY+TXf;T&=!8uv?d@*Mu(ivh~1AL$AyIXdH5!e0X#= zoyn!~;xBLRt9^(-b|oUVrw{h?TlRB|ize&O^7wDv`2Ivm?KD6vEsAxkzUgp>J3oLm zO)?(r{G|z}p*}VpdgIxnfMeJC!QOx<4%3lQB0LYZ*N9=}bC~utaNkaFz3rjIzIbxj zIQ%P#-=NlT%gb&l1r{rdKi5H3^0ss} zo$ZMvcq*RYYu%ohlgy{jlKKzs&n_z>990>kO)3qwsH#f*z8SA~#3VO14%r(TA-%8a ze0R=s+o|(8tFFn)gXa;e_>13^fChuKUJ7cTI9X2Q zNFtz4$9s0-%B<`D!IGp(WKLT1I4c)w_RS!tHsXs|z1zu{x>Wq+q&U-8ez3xArocwuiOAAiegnFOo!6e#uXYkh*rpts0InIM-`6 z5>{#yTDWtI|I1B_+Y1PEB~v!)EsvL%Vuco&s>Z~CkM-W?>{WMpO1{*CU_@S@Dek!H z0Si*evz@URSx?uYY+3O`M+Usxe)DvtWZ8_WL(Ou^J^kB*_C>q@;T!qHRs>bzQ*;HxABmZx6QZho=R=SGuRSXToOK9r&+P5MI+nis% zWg@KY)0M`IN!E^O&LND#WhGl%3(T1M-FxYdb*f_9wZV;ywpGhHYuD@zM~j8<@3o++ zr`NXqzsOTQie-@2ZnjUqD*je!ajK{&w0hcHV)bWbtQ#Wxf*?_95L<*el#Uqtk0SfkK@f? zmerIV#oGIW2nTHQSrDdwq(R#yMN`hlkJc6CApSW6Lm&+HBI@L5y`0@YX>NBf z&bpeHM%DylLS$+r9?Jf_lt;aKN(z64+d0{+&N`5oZja?oj|jVZoEF)59K9w(jj%16 z@Yzwo!Zj=(rK+&;&#^-28QK=dw6O9=z?7+(IeQQQ(A@pR^!i9!PP&GM7PR>|R;frH zWG4*Es%v`U_DqYb1#2o65*;eb(K3>Go(EY2;|kmYj?35HU`r8v+*2mk4Gc@F>g)Nr zxQ@Qi6BA9%z7@>&sMcT)5Jtw)$5_}uovOLJtnj2F;lc6x!tn|3eX*ZSUq4+xT{+t# zjLZ>AxxU&HF|xRtpU}3d_GC1T)TWA;h)mBsAt`<20m;Hn_JH!u>wo4>SLTZc1rm|j zqcf(Y_JoheiLZ*sOm3ZMDVyBBK8zBqu{Rjby5t1X=07a0`c~WmJw3e*)h%HBtFb{s z*UCJPa=fbD59VIhBq3D{@ch>;5eNl6yZ}-R?%vCM@lzWm&BwPz^xnN2<07N|ib#vk ztf+)kJC-z(4@|Dd8^KEuD0j)RvEWXFVdDwIULFM5X#MKlm6Df^3yX_VQOs%Z6r8>q zZ{gL;IJ;NX-r#J0Cb)~k4dJ?lz$Wx60ZFBsh2oU-3k8t6wakj|NoEu7RmqjWd>zqFmv#1?W0Et zxKzLeUnb_{=Lh=wu53L2J;=Y`ybLlv=E1=auH}ot@ccP*tf;I^MMZVscpt|_`GQl_ zU)^|#qw^iCNgQqCrl6ZC%uF{FueC>k$C--G)5xW?qVlib_aGxVX3|C@5@Pxh#MwGg;@!EaVTJ*FM82+u4;?xHbDdAn;%4WQ4D!y}5DorgEN| zqKr)6DU4P*f5@1%l^RV!5VS8v?e$NOx+?PdqR;hGM6?jRfBaB)k@Np4heLBSJ8eaU? zYbX&hac4(Ix^%*{_*>4no`jQmsw-En985VPk2m0Ft2qUR`ueq))tm&kZQtS#&AQIZ zNj@DN9ZI<>TO;aLeA{j0!YS6n*@~@!B&f8n|1(!JYrqsOauq2OAn4%}7sY(5URIX&B0q7fP+HtRy_8 z`JXOv29(1uZ^l^9wfMGy`3&z;04g9OQ>fGMhkx^S3$Ya|E9)m8A6ig&GghP*?zph9 zz;}8$2d1DyS59uAz!QrqEoH?wF*g3ja!5^f*Fzo+RYyaNKbfqIXeWR+a5w89zkaUSGq3B z^*bz9<;Z0j_QY0sU|nja@&{M;H;@Il zL=lsdx%%Bw3H(!*6=YmC`YQuz9qQXVJ2~0e1dLj>y3IH6qnPxBJv>hHHERHpzjyCm zovX)AHwR*=heoxE5q?q}a-YYjC$^KJW(|IJ3}5O^Abb3HvKw!~p(X(=TA)xgmyz{{ zKm9wFt1gzoi_$D81%z}fR#DWdbtarN#3r|@iGjYuPykkS60$swuQrz;Rg0Bu?6E{w z$&+q99QGbdY}4BxAocuZogFOO;**0Zy+w$gjSvfcaw zcq-LmecyAUC^$H{1l(wkuGs(*VPRpZIjREn5cTfIetv$?CS%qKfZbYCed z%{9BZPX)v_Lh~n^?P5o3uy!%bT4CbD0I{hV^R?qKsWIBemaVd8Ma53CL$~mqQ+Dg5 zSJ~~C)_tOH2z%dB)(RTH1Do7G*zcgMkc6&2HpHhc3J>bT;~AG5j3K)i_V)wg%9WCJ zTchwxy4o%)Y56S1tzRCILwxYaNZCxoU=p#MveXH@&f7rO0v~hP%(vm=%Z~WH5fp4> z8`aP&QUMA$S!D}!mX1JNRdsyaJ31sJd2QVsnK`&(*!3+qF0RgL>nEPT$=BbQPEhNW zt@TF+c3Sh=5dIj8E}0IN*R9lyJuFgpg5}hmk*yI3yUi9afqe$b)t^H1+V<+zNBfMc z74EWdNoVZDp|-Gb0Au9I(M#J8j$R25>Ss2-UTAQv|D|Ytvof{OJoQn2K!!?ChK9VV zv7D+VwyXqq3de5c7Zw%{4Gq<~9^5A5*MKrJOYKz~>ch@0&%ru^!yRK6Ae<5g-=8oe z@Hg079R~hZPB;pS(mtzp+?Xy2P`5Rw2_Ntg+p3P)!EpMGx0k}a%cPVTx=JH{-1N&Nn1-F%5+ZzL><`+N6;|VHQWC=6EzTlCNLyn1d!h}*L z-r09Gvyi6`@t_29#t$VM8PeQ@e%!57iyT#OlAJVfrMrgYKI(65wBpjGXXYVh2#m!~ z9kI;r1St1Cd_~_?bQ6l@M@zO) zoDtuCYq8@L>4EC{*EJom`?>}6&`XG{p6$z0ek|i{|xT^{#9d4Zi9}~(L$ux$B(sDB{wL%%G1))^w7Dw zt3z49*JuX_&sAE_{b+0q|IyUcG&VNY9mBCPQHJW_+1RzL1r`-Qio+`Y+++~=r}4}> z^TzM{t5;Y)2nL=4UsO!&ZZMzgerj48wb_*AfwHzc1q47r$o7mEo?lzFJ!)d3!Lq_1 zDBJ3E!nI%OlhYmHe&E}*q8rwn`9xk*+&!hlM8B=@FflWfzKndY5jtPv9S2EPUeS^~ zOA}7pf>FI9j3MA&+-^B(^*en6Ar$Ez$ns+Ji>8JIr`N#Di5Cgcwsarm2CTH=Fwa$ZD96crWKPlrcEWzCoLd}C@<&jfC&JG#vM z*a2TuTKXVDL|C|7Ix4TAK)c2vFJyDV-W-@N*}GaosnyjksF6JIeZAJe2lxwT+8#$~ z9L@WivyBA)6l`pYqDOJe24aee;dj3zCnxIy_W(dy6AXXqbmyWk-xKFIIioeZx$O(J z(mHQF%Db!hq3lpcz8B5G&!g+`g68%6!Ki&? z#m~)E;yD87HKFJ@+)g~+ZR67dy8M>q#5P)T#^qV<#&v4}`!zi2`?bsX(&q2Jh))JR zI6-IWd2Ygtpl_0;K0;F?Z{z;c|38aA#PGXXwgizC)VB5XSZbp+6$aEF5Ng^5!*bQi zR>{r|Tggs0eT4i8!o$OT-$+VIj*4(OZO&{jbRq$=1R!K%A7Mop?7C6M0qg-ZT{4E< zpd)M@Q6aEDrbo$V+R&kDf}pX369L&hSkK02yTN;t1g~XjniI2 zsoE9%UQ!f(O=tdH#nhJNso|6>@aa*vq17U4fA&ah)x84t%9wI?G;Nt8tgVbG@`z&b zro@tO8(|rei#HIUf`BpvR3DRtZ`wQV`oQc!Z3mB)P+{B*I$!X>43CWvUA8#NxvG#H zO~9XY@LP1u7M|B;jL1a-6|_pp43wah65kTP7!#9`ZP3>)M5t$HXY&K>1}v&y+s|fy zLOK*^%$3Vh z=;b?@Ea~Mx9ydyG-yb&;n`QJo>KrdIqAdsw4MjT51sKIU8#a3560@7j>~9VxoZ3y- zI@j$hU$}$!m*!s%_V;_w=N4{}3wryuRp8cGkk)k0ktRt;n|9p~{^r+E8;BrCeWKpKpZD{-Y z^^Zr?<&>c-RZHDHB}(<>+zT z#zzX5vd6@QeZG7dY-X*>We?JXcF`{`&tyK5$UKQ}oie=2ElY6Q}>9!MClkKtYD)GbB97iKKHytR5s zVrksuNh6)-PgwD`<&+t%B#+tXbwgloqlN0wkBSZQnP-Sd$Psuc zHsIl%y#)#v8HYC}c=?}M7976w_&1aF;yZ~IgYON$JHY{aJYgXltBy#Pp9J{SwhA;9 zw}TpI9sbPm(tf;?z!q37F~^z4D@^QMoK`+KWvZ}>403hzoSTJex_zS@nssnfIR;Mz z_K3q6A2&D^$03i(H>O3<}|f%#&;H(K%^97`g_r%m3m z-yWrsy3t!$ZQdHw34Mm(oF27|Bh65s1a9!IU6Dg1wn&fB4qs5L4Yq$4fct0xU-}q# zf^iWM2v^0_Bj%z=gIPFS%3s~H8B*KroGUn-ILbP3am$lC-Po^8n{u$^a_^tx_xUBe zL#=FU)skbLtL=!AEVv*}y|$$#x}=)pnCHen-PL*C=f}>(*@7-t=0??i?w|O}i+f&Q zJg;%i7;Y{1F1K+?Zk^S!sPVbvA?*s>*p3am*_S=rP-GguUs*S~@iqs_*8hT%`MIt* zmb9ky(EaT4&Mz7DrL1pU2n7e{QY&@Oy+A*f8)mV4Y;s8JIcAYz)%DtWx{3|{Ag6t` z9aod3%lrR+75(p5FMi{FedF>khUxvQ-dU3@8->Vmn&aP+{|^T0{i}lhDYt`Qz_DD6 z6UWP|8h){$es>B7a~s^(WHh^Ym;-n*oH|r$6{P~%-E@v!$0s;1X8KS>jI|_&?j1qCZ(S`AabJyJUNe5nU%(n6Kj3_SF9=Ap ze=Y=c__SuiHST6?HM?P_S`z3UKJ15d} zgZ(@oAxigNdU$xay}ccz^xjV^YG^#h$A8Nl^_v;{P6#;aA=nEoKlmQ=G)uuWR#BMR)lsL+Un`0IOetM0}c*#Bs=4HT<@Im zx5!sbrB`v^KE?E%-u@Hk>f_%;yJ)iP=`M+|1O)DkM&CSyMNAn<5siOeM z4YU{*@^#hGvI22@77t3K)(`dp{Zo64u36`R1q9TAH|Kt|-}uA^de|J#sJ*&zq5bb0 zW}Q3%TE0vGh{yATZ)k^X)Qio1WcJ-T9sy+0O#-@aEJlG1Ix27v>_nN_20$*Jr#lIN z03{3tPzQUAJku)DZHD2KV|b=5&oK^QS|%+W0VxSW#dw!hYX9b9x1;UT!1J8me;hi? z4|7*aj*gB%D_mS$U`VNQ%(jbN!Q>*RE7G2r3)p1fGyy(3CsnY(NNc`pxB}SN4pdG- zJ;skV-)Z(UVSe5I1n?!Nzi`Q~HgsSz046VGmX<0&&nct&cj(C~NJwPll0!|LR7T^h zOiPyiIU$Fri5fQT?(9SW6zqGC3h)z;T=$)Lsifm4!5Q!`jPEe{ar1(Wvv`-`4j^%E z+_=H*o-(*XM?+(3ZobxY?Mq-_Jhy!|EG#_S`KBg=^4qtw`|AiE(jUaW;3_)0+?*Vg zT}4@08K8_iG90Qj*1eSZT|t-(JUnpJ6KBo$+1S`>KxkR#>S$qM;mCgxmj4>dcN|s~ zk%Rz$|9ipba}Nko#T@m}_W=EzogToiUAtyAQMxkgL#A{?IbVYsg2%i0y_*9`2LagP zu$j^0ZIMb?*d4xD$} zY9{OzI6GNfTwF@<^jM!L>rhMSKdeDw8K|gQYp@6__}_Ey%m>h}Xeg!ohLMD*=y<(* z4aZPzZEY)=XKiBQOI5&*WA+L><02v`A-K4>wo5&A*nNZ|&9GTPUAp-TtLa+6VK2GV zX{plD(t_=>0Ex_%{OMphC8!=;=TdjNH>$m{TVkZI&wF;VbN%}Dj-NLfYc~Or;-Xix zH==%ytT(htf63{eSDnK1r8U1qJm2FeV+(v!!{rngC5m z4cP%?8@okxh@l$bPECiiQ&UsVkvszf8hd~|1IS+6c_Cbc%w6G8__lw(g|P%`~m|lJ?$zG*xf;CGG3?6Bp*_4&|jhXcB+KQ!ncj~qneH@X_3+0 zi^$Xc3c$mkfD6M}k>}J8{0zGR=uWHJtrl}|`=Z=6O5qieNBnw@0IQGKh zS5Z_v{~5_|FW!Qofd~%l?63zqo6&C^)~~PaiM3DZcWNbZ{N(NZH+$>8uWD$xh|4jV z=vKPW8L3Ve0sapn6Yu~LUl8B^H<-SG0kw~rrRBM#0a;1vaX7oXy9+e?Y`+5T4%jFF zBjK6DtMuyS=hmdes8^J}sIPDH$4y2+EYq>D$UBb$eSpr30M@P6 zc_&O(XQ?Y{2XL=6D9&vOnyhLt-#JnIg}e{*|)X+0w*C1rqsAE;k1N7T-> zIjqhF3FK#I%N2hTdIng$>}=DmpUt#Nxp#seu^7%TEdjE1cD);5+@+-K@c1p8(UpnJ*w^L z4ro=7ECN9iUKqX`NF(H(I9pXML`+U<{7#wfmYdvE4w8sDTKDOVAj6gi;kmDSk+E@% zsYogfv$8 zl$I8?Byd(0@^d_kwxD+V<4p!wP>?bK?o~l4sUBe9G#DaSv0FlDoF>63{^pN$zwfZWipsvf%I?%mRdZL!F?&t z5l{!nmGi#VpMqzMAg$%X4lcLCSr7*f1si)joE`na47&b*-MA#w?9CMMFaR_7P*0at z_MYAsddi~2qUPy4Ouj_-g@i5%rQRIis+^-4MU5O`nU-w0&=GYlTHwhdwOVUE{5xHa z`|)wk@kaGL>U%&S8%Q|3x>^nDurW;!iB9)DXD)^GH)k3ewt*|8gd|t)0 zn1sX-nK~Uo#^4cPIDqxf$<6KBIM_f`NJCFNJag=cZh3$D^r^a<8*5@|nIA8Y7~!ch z>?Xo@MprF;6bESzA&{DT1MLB-54aOPz=e`ugG-#w*wx#s zU8MW@bZo)K0#TkfY#-x_%BzT=Q{Bd(3K&(3=YjRTAN=XlpZF%I{0aakF<5PB&e=~d z4$$Xj_3`@C`3Uh0{kt*^MdVn<#VdP7c8$ zRY>x)Y_0qyiRSG`*Qd6(9CNh$Iwg2_2vWwCTy#2@LmBDE$_%mFty0_@@dLG1D4LtL zlnFT|A0!j_U7vgE752=3kZ*~{+%wWJ)wExUVAP@++jFjYO~h`{kT_?kx?ZH{V*_{D z_<+g&13a36=o!SBC)%SpSEn};lH4|nyn~t-zVK#t-b5=n+{^xq^_8gnntk*X!Yyga z5@h&NL9H&DT$!4lJ})aPir3lZ1xPV~yw1lUpW7rf5ZC>+>kQO%bmtrxK+0!-Mou~) zEXK#j1+Xj9aZ*d?u3STd4zjYc0x1UOL`hKcoZhyoKT1D(EjjplO;br}1msdcLN4G=qZZ_eFt^Je|W$>0XWun-wQvzo!~HG62!X= zgSlM1ghR@rlm0xq!ba_V$6KaYOfo%7_O=nn4?;83lAhXh7VR#q0G%Nzww1-{pF zShr|dw6sC7+yaq~{64&K;NM1_5etF1#Al*`qNx_?3O%{=7el|51M$6kwjj^y4~u`} za^T{aStHQRpJ{<>)d-o@Jqoi60jLe4?#?$J+yKmR!W zkm8|ZU-@ytjMiVkZMJ2hU@}!e5!ZXtB7xrHFt;uy!nbY4x42&8hNHw2>QjlctjpN~9W3vYZC4jBI8w{c@HOA2B zXc|aC-G1~as=KF$^o7oMP+-7eHSGqRAt)g;0L}uq#$jlRGvtu6DrQX)jj_RJ ziq3`nG|5%3I;?>TMt~^_ z;0a)mL7pdbpyo$t9=Oed++2-v^Krn2>}{V1wA&rDHRz3BI5TOvss!yFUF-$20OtO( z_Z6WYPJ6a4mqN!WAxyrPy>h_kuhhO52y8m9{&5MOOar+a$t(fAMTxosKb7$u)_DPY zqp|(MT`-AjZV>wHsfEfrd$2Mrx;rJJ_fUsnPvgBOq_vJe*QbUOUm&PHXR)5Kd*md7 zDh}pqjNkv&CDJd8L2EU&jO5V$#^bfR%YhAnkm4eSXktMyfq^u zYWYx6UI9=MSr^t04RbB}+g}2v8A`k%af}XJ(%l17zGk-~g8&qp`%~ zphExeqYvJNX}9%#AaX-8E&Si|sgTW^k!!s#FZdy3sRt+aRcE3|oCx7HILN{%R> za!P6&2Z(W=ML16F4O?LiBI@TucDFATd%5pm2drDSNj>GLJj?xVIpS~1p{l_3#_ z823ShpG8lvmajPJkgk0~#+G(ae2zd55z|iCP4w+QPd2b~eyf76yPuOd!2SFe4D8EF z7u1xI*W=r*U!MxTzTei>WnT1*fq}s;Zfk3cL>v!9qky%ts55=>FtAJ1g0Kcu{~&nE z(I7$+2_Skfc>t`Zf7?uhm->DGz(7z#p^(F7cWCh|BPGRMX(lPu9^X%;Wi!C#NawQr zCA!!h0!gmJj&i9z>6O zRqV#GS5?V>sJ1*g2~OZ=X`mwK+67`!zz9`=WqOXnf$xGa1J31*&(#Y#qN3%TjHo<@j z(rH%P0n~!I8`)_ObnY1?<-a)puhDGJGiS`lV;{-*3a}9e{MdBsTPP!+$?3*{`8V@& z5hGk4QMQM83szG#TiUbAD#}sXN|m%v?UzrqYfBt=TDMFZ<*q)?F!^v9={_MNZ#Oht zk>}#EMjbhka?eO}uQ%j!%%G1%fVF`^P<;L2EEiRRCw!Qc7IHm{tE9k#sj{QmASit8 zES=?F{YvfKzX|l^Jgx#*#M06d7&T$x7U0f6;b-dH1%i2>(g(eVPd`H#{82L@^1<4h zH^dIBS3I>aYYzg)hSr!emL#~(#!+oTg`n zk^kM<|1buRR_^p6&(q8`2UcI~?q5v$ZI*9qCSKB;F020g9^+9ixwn;96SE9HBNZdj z=aKQf{NQR`O|i41WvZEhvwjh5SlwE&Bi%X{3F&9o$(Lu;)+8O@0^uVsS6i z>81V(ZEzBoZXDLdw^Wly)}Uy}@o)drh4&2qLBU&O=X80Zl30FQh4=ZX!rT{q{rxgC z_J0FzIz~#o*dOuXZ@%ishaJYp6?EdEz&i5k0|05@W z!T-&|{JOj^=%0I2e*Fzz-+TO@RqDIs|CSPCK7;KiTwj}NPyAxt1DXC~bk%Y6{LT#} z8&^+GGdwwqt;{c&hnGIVf8Paj-qP|hul+DOI@^1-gTLh0I7x(mcQ|olr-nFQz*Fqs zHg1#unkbw;M5cGBr-w9VOw`Zx`9$^XppRALdu6C_yu0e+p^2Hcvs)LJfKm`*6Vb z?EG<8o*rTOy^>Zu2>2_GtByih3o>|+zosE6Gx%angQei`&F+XkuxTFSS8TKxplyde zY017JqPq<%@pw#kXY*(pzJoF9Ayuw6i)qm3--{5a$ngNSa&N^`TjQPBY7?n1t9R}7*Dda%zu1V_HB!fB;aQ=YZ^)O7t&v=B`Z_RSlod-ww8hZe_yg> zRih*KY$M#e^tXiGzxl1GIDIkm5((rGL`40%BS+ZAk;8H8hF{y)(4JcaD#so7oV6Jb zYKZxdk2N77%(dMJ@Gw?Q<&u-r7~&3pxW|?(zV)br`&JuCzyk_o?e-_wGhDJ1!sFEh zPwl;3ZT8xF8R$s;Aq=Jx%VaTJwZA~MJ?f!W?e0E1f(w8RSLQvdiz$2G!~7r1exiZk z9u*i}rEb?ap+C%&)%T4;GOfz@L75DJ=ZHFrnS-3GUj2)kY(w5PA6~vZgV&eJ0L$q) zIJhl`FRSQvTB$e~0?SgJO*>S{Ivnhs|yDWt>1U>b{iUq5ugbwz) z%sdF4if6S;g=Gq-PhPC3=?^u(Z;x;Z73AS?-)eII@KR?}ltrXF-pbd8zzVtDuu*_Z zkV#ML+u7B_KYJ~^q0neS1f_lI`i;#jni8V+WD_fog{QlATs_gk?}y$C?wBMYCJ@-z zsT^#TJU2nbwlNh+T5Jy=x2uYb&_K4`%7kETx(s+`8+P*f4OdNAwIz4 z{U4uR9FMHWvp-kP)@4FQn=^!#PdF=Fy6;JwTa^C$rc;FTf84VZJK$GpgM5Ntj?|h|sXx0j?{~)A*iCmKpUUXSY4w#J?TOj;r%D8@Kle3&GyeRMX z{8!}?P{UArR*D`nkXj~w|891JKk9mG`)ZreWJ!UF>0$$%%tW*$K zgtb~uE>+Op<=*p-C6Zf&akEZ+4FAUi37K__`uOwiMO#@KQQosfc>JewfedcX>)wLP zDeu-b86IL+$T;XZ^)2IBnWoam9k@|Jc(TCAwJPhHXYwzzmd;UUE5tgCD)zQ|-v3xa zD2v`B(E}Xvn%8x@|MBqDsw(w!lt)BUJd6L1rx0WMBr1O6$4T*W>hljrHZ~HOy-O-X z_LQI>w6Fd`=RUOF@~)GT^ed-1kB|epbT)@X00P^SX2$iJJM=x&5(9LvE)EcN^k!;9R@g1qCi}xe#v6aW?x1%!)UAe}=RLr+$%#D>l?_Ep z@wb<6_Y<{t?GHT*%D91c{QZoE*YXYV(K%nVZ&w);v8S#hO^eNseg+3+QTsN|>Jup0 zzJt2_@|;$XB$DJ#=ilak?B)YwGxC)~%C;z=AJKsML2^CiH#HgV z%~T>&t;1Mk$XC8*{!XZ>iuNFR3M!03MaqnnVD4w?W1Elg6~Uo(FRTQm#gF&-7Fu@y zHUZh-jA=DY;Zz674&CHsP8VN?Ubf8SK3@@Y|LRGr5O5pmy*}Vxgp}rFFE-n zSguh+p&8lY!=V|g)s}vmEuQJSuFc%uUjBnIJ2kn$5lkf6Ux>Mz)kO>Vs1n)ya$3Cd z{x;z7PXFy`G$af0it*7ijYj#b8Shg<{c@L^U2_iEUgiu+rMidPgQSW>WDZ25Dg4k_ zJium=knU+#eD(9$bd}vQU(Yqh5r-Si)(WQC^jDnDTH1TGK+6OD<%l=V248~XaWBNr zN0HB9Y{+_u!br=W6H0ew&dILy8{F2Hb8Ol6;AMh*HN+`~A$eL^P14Cu?0^5PfhB6*QF zFb;h>dImG)?Gx<3mDm67!NMwz+BxF$4}%^FFPmHbIdO5Po+4}vwrwSjcD%L;Ps&!P;z-1DiXw=iP}ZI9 z87==G?BJ=hOC88f9YP0E0r@lwR3%KJg5ECW0C@Z3;VESEyje0neS7Vkk<;#CZ*W&r z^^1j67chPLLNnOtEg~^GOXn-D0YN_Mp1i_|G+ah$UgNGvmM@F%C`q-a;o?6#t$1gQ zPs<)KXl)19 zo3fHUrQX)EZ^_{6CIZq4A8qN45+TQIx&ZTT&pxs7LD6-7Su8EJ=S3JkZUiK)xUYiJ zudZ{gKo1)y+}5z=a>TGU39udzGJ{Q7RrL-= zHL<26yxHr$)xYp9SUP#4W!B{#2J3oVt!3my{m-x-7ETxlR>f7&99yR}o5&14!oP01 zx;uJVajpeUBOIZ0)i)Xf^2R;bl#u<$$N;_>bLO@3Qhz|3Tx0DjH>GnepDo*0y}I_= zZJy0THrVB>e~Q=4XDeL315$ka{X&3aidSq3OU84Dijz_C9QRF2g5v@qM0|TjiI)a% zpJ7F(2@Ugg+CqeM>h1a$Z<;*+*&%+$-;?BVdvJ4pHYQi7s{DVqj&Jegk-oRm-Um%?`-&@D&9O9k(>xML zU%o^Z#2RRNZ?h%(^jPiNN>nvq^7|#k@?3Pu^x`E_r)fzEJv#IFtB%U-BswalXxc$t zzA%qAR9N~st#I_VTE5mfFcx3>7!g{ZyEW~w`*w&GN*4+pG}JV1yuqQUu2;Y?*?To8e zI-?@7dCJXO3|pLlmWlsB(_3Yiolzg2^v-7J@U+LBO~6oac^M%nP{zhBb+@;ZMkbo# zq~o&{jN{?OkXY#te`$>JD{1nz;2woV{njeL?KZ7f$#Q@|%YXhMFRPfm2ZeeIQzhrH zVsB&360VGN@BZ4!(9<&n*0lt;^fCtrq&f#cg)s zPd(o0bGfZ}GN1hqauUB+cDY8)gmNN#^H@i^=Mr&meV1d7OJ!(Pr-+jA=zbd@+Gi*yPwWR5(Kp^j@<5r((iVOWK zCa7qe2Ajq&rxIWurSyNnG>~%{_lj+S5t#iFq~k%m-Kz#OQ{QqS5 zrf0o?y8C{2H9+LWI30MfoxnI(;z_vqscZ}m3MlY?;0@>E8zz|DuKgdi#mcG{yIt^z zwEjO1U;%mhHNc~f^O%j#(3hHSx8Y?I^gTXnfR{(o^gk;X3&pqw0)S{_9Wq9nfwzOF z$A2_+pO%_pxEJaA#8md6C=vrGpojXyLJ0>do5K=VJbSw1hiG_~q?>^wg1~d6ieoaGt=D!`{S|a9z zuO({vPeuPe`0s~4>VM+w|33PE3v2&A`zK`h@9+P9-~HdadAFs37>&Mj`IR2Jw4~6} zTU_Z?7Pv{2E77)J)6QD@)c~Dtc>bU0Xgf!&i~cI$?-z@&U)(y7n>WKYJf277Ona%l zpUP&%6sILpGVLz>!`v-o#Y`TDduZsf$BLvGU@0&|AeBDN{&G%nRN8rqRfHQotjF7p z=kPQ@e_?D=7_Vn+*Y|_{W;W}>@yAQuVl(Qsy%PXv&4GW)8bD2GdA+&Ma~j!iqa5)p zKtD{DPauN&DNp=G@Mp2nMDKwrCDnvMNp&?^>N!o~3r^9Mc-!B#JjeoPDp4I(m72|I^P=@Nun&wIlT=Bs0yP-23X$f}+n; zD$P?>D^kP}@t~LpPn?s4J5YVq0w#CU3{g2{Yx%nmi4`ODd6{n*rbFfkNKNOy)nh^; z&a{~e?#)}=>|Jxw-o?e^Z$+n`jwaWTND1Mci7lKtzp;`oa<-JH&THi+Ga%s`wx?Nt z;v&JlRgl-l-RRN*&2BaQ+8Lic?R7Tre#|AdW-%?%b!q#n@#cq?*NhG;0`Sr>nLXv| z*GtLvay9Oz8r$Y(oHELHH_kgF5r-$Kin{I+zB*{dB;3j9V4$Gxd@2Xgv0nR~hxG5B zH>1Ok=fd2-;k#$JOQh>^yIExT>HjvCD!4sw=t{%zaHtE@)m$6g{A_CI>NqBbjY*=O zC_0F(+bd8*LhG2dcz!Cj{}j#d2giB7dB!{Lu9P7({>hFPgguvfoIYwijr*L|dMZjZXUjNjz(=e50Ny+?dDBvZqk zCY-ZsjuQ8)@|#tbBfDA&<+v|$moRq7THuW{QK%_)wp4f5DVW^6`+R}3;-d->gp{vF zVbDpES%N^ZihA6ApE1_ww?A7Mdt=UrC*ym2yO(%*&qI%O!=w4F<|Mmb|B7j}mYWPc zHY6u@;^mOaFf8oUesyy+7;&GOtoF-W**hb1G!Z$%ZoKqSUesBt{1D2>`3w3C0Miei zF1xkfZ50jN3JE)W0X00O=dQ%9DXzNe5|${=T&S( zf8teIFdKTe@#e>Q71q>b7pn554G6@_F@vrn_Tubod20|8lECp}#wJcf{WR4;WRFp7 zOx*|Dy)S1?0tk=%8^1%8>ys;5&Bw+8r@tFxqEK7a zz+%HaFTzO#(obir$Vxk@oTstIZ}mFhyUn~!Pr{qo<%y1o`8FS}DDgvuxEIeuVy4yz z!MRA|Z~0gyE=JynwcZWm3&mUpNo9>T=%FQi4pdNWtv1LVcSW120vgfd3CwaS2%ii$f6 zg(5p{upWziM%(g(fA%)x8mf$n#@6tBS~2$r&VBJj)+5roYJFLAm&YiB%Fig+$Hi@- z_~c2q+uo>{&vWx)MR4(wpw98_XA-XS{tNP$lt=3MhHeyaYe3>Pze7Q~iDElof{JBioBWg3 z7PZ!1UJ6(-hRc@SaO~(!X(vop)%&oQ3!Z7oX)u)3*-NjR7EM>ZC+ZE@bklkqs zK25>1cBNLGTTL_tUFi6Y$771r;@4nG!?q!A3+d~N=>+kbCG9<^a+q6Y0(}?@I^*@h z$^msRt z6QcdnuW)C0sK?XgmtxSou-zRo2XkP0OdekPZu_7Sr$dMU&x- z3w(pd2dqA~PG6GT3|t4lkq19=amop8qq5IBAi%%bk?mJ9pCR23#o_uZjqFi*Z9w@* zuH7Gx_KM4{9VeDVyLwfCxGmd?S#UF#jVbS9Ef&G}EtMb7Lo!it&VT#3^?p2`WCIDu z8~TB+oO7p@WiP)%yW3R9@l9x|POIiV%uHjoVgwoyb&=#rEDg-pe`eV@TG~CfcHunHNkZKO3jEVhV+D!&7}Pomvx^{*K5FhmWh{%Q>xtE@Q0f zs7PW>n|VFMi_!AsCH~E9c6?TtdD`Yg(+!aM`gj_h!!VuEF$* zy6rFQ8QTci+Rpk+V*G%F`n?##(0ZKkW_OhedF4G~Lzvi%STl6M%1Ia8mA2%xE-%{V z=RTqEUVm%#&*PoM6<~lakYS?qtI0cILaaY+k2`ueI_BzG{+1UNiLJrsFp1|N*J4r0 z-Tk)`xEFh+9Ddf?+xKsi@0xOkU(^aMI@&S)Jb~YmnCeIgtyNH&DAZqlw$7 z>oCelR;_{L?^miaV{-60JqW-f{D0tgtg=b983pT^9s!HO&6Up;3iTBf6?I@YnXYQKKXGA=6m>^I@%V9Gfz{Qu)7q)OB z11g>^yAkf?-9IW_wtYIjT0xmRZLaVMU7aG~zBs3Xit@ipP**?MYn2aD2D+GsGiYlNjH4Wko4>6O)BRGxW7gDp5GpXrU1$Z>8ilcbc&r0!6pNH&U zJ-T!6PNWI}>e*p|&HbqcD^kwStYtBCUl#L%kf1neJxM&pPOcv4j#P9(99$nIcbkK^ zr!o4!gUt2qC%GAPBZ?s`Qt&q%V`mkp=CO5UWbgMhfG9&U*Z-W(?++9Ga7zDQ6bupa zBs-?DpA-jAVLH>(r-HNMXhga009UV3Vhb3@U?uHvUW0ugndU>*z18AOvaM`@Wwx2> zd1os3f@NCj_j;2uv9@{Jp~V+BR)YWC3(yk$AV?poBh~GCUor0^IieW#aIK;W=H-5& z7kfaid6AUCN1OOnj7U0Tvsy9qa?kOJv*C+_-=)d&)n^ww$r*^`3_ai7&0uUe)CLo3 z6Z%;%V8ShC1!t}nERoG95Ui2W1CUVYAM55>D}7-``r(=ha4urwuHzS-STKF;@h`A> zv8M8V*o}?*;Fm$Ptuk(0w?Fw^y@^fdBk_QfK@1WFG)g1u#by#`-AZ?1P$HzxwCxbV zC|L8gQED#pXIbbUre2TkCBa`!hw-*fDLlZ_X?uklT`4K)tR(*<1U9vBp459xW6Y@Vo8SlCqR*fBOZlOY8T#>GSDDVC39)kz;qnp?H++%Nw@27jNW4Q-$v4s@BSg%D*X_U(dMozd-uEA zBqeiRtl5$2XPHZtxqC4;+!m?)GxLOuPqhCH>u^nlSn`Q=2A*98-kfhv=7L*C%A<*P z7IQ$6oh8l|$MbLFbWO@BJNDP-!Q>FIY$6B{RUjI>P3^GY#&!O~)J=Sk)W3s6hr8)4 z3A`K3#a~M5BaTXmyesggaRQ$p;v{ZpUMo&+NZf1-_Nz`9O2gp9TzlGO==jH)G*NUa z8wo@u6mZwJi?D)=nF6Li%5H$lW8DV*2O#OBGMb>4mS^DW1K1bx`#yfxqz3qPTcd2bwUJvZ7;uKpOGeZs9wYwrW(I&ahF)Uwwqsq zvqKw^4D#LIUR&<&MylJbf(pP`Tfp4$m~KGgBiB5Usu4nV(O2 zDM>g~u`z6r5^&uVwXZQ*gqaOFG?6%^lnNHG?KViKqfg-4+{~=1zGwq>q>1nYnWNf7 zUeFI;+aM7agoI%5VfNz`j9DsxW5i|azd6jRLyQIq11zH#O{4i8nBbmEpPS|2u-L__ zj1A^i1JUW>HB6IjVU*@BJ~*_+8PqlOp4J-jxI%>h^u@y7K0M$;T`#3xNKGBl{=D4x z%(;9nm{qE7W%p#BM=No$+3?g3=-nQ$OJ!h1L)9|E?K-I(d-5idb?Q&}cxS2L3aNx- z?&cCk(EbcJ596TQUC8{m>5ta1v3c6^a(yj!dN*1Upt&=Zbd&q^@Au-x@N74BUB$}Q zqsoSG8IUMLLSjn1tS#=H_kty@7zoLjKeSd;z8b~Xd_X*+N_=P6V5a7WLHXs<(E91x zg=v_in`6nasQCu)x?lH4zMlA_0XjO{DcoS7mf5oXj%`C}ZTTx{VTnI4Xp0Fv_DxiC zZx5S4S(OzV#dFa+i&GN$w(?N-ZC{hqt9ef3-_>Zr12lmh!lV|#jUPMtiM_TSMx^G( z#3c6Oi^*08Q;kb*syAp`;n6?+)JnPG6osc78JC*(dB%1JffN=ILSj(FkgwmD92@1d zj9_F108c`$rMpp3qE#KBpOMM}E1s3>9*VCgE>_$q6>)Fnvr55zeu!O0 zBwUmFa-=4sOR?rx<`+ka6D5#-(4^W_a|0Wd2<&VS#Iz8R9(b3QR zIB~e+M&r&0b&Kvt6NTQq{I+UD+1CJp1E2`Yjp`N{>_hTD^_;9V=zRjvWLD)X4S`*_K4z+o?Z$bZ!>{9`xC=m zgRB2?ulSkOLU9uAWEmT-C#vG{SMJAzn@Z#p4uBEaTz9%SP}i|#&1oM`F8-dz36t44 z_39`|OQ_fihd8`XUnJ$tpQ%M}^U+s46*A`p&H)|XH}wp`k$yp_nnv9}1T2}uW%f_T z!>le77Mmqd^*V!X%R0Ei$jpc6&1g1>VkB@6o_M0V8Eb8e(ix6LQa%f)D(Y^_PJ0u4 z3cV_6Ni|l-I{ShYxD!sbh(5%RX&S-tIz>q0-}Ga4pnjWnI9q>n6}C_}0OTb+?sPh( ztCh-vB-iACHEqS&%KRBByqyP@zS+xnxKS?nAcQK%at`7NEb#CzFJ?(fG=N+h{)&0M}@U z_2S#}T?}xS|9juUR7o?=ej@x4*etRx;ojEqcwe8-L%bDforRa9qcKEX=-;M;s&~!P zn9B775MGW0tc1QJ`3!dptpP3H*YP~o_pXOV+PJF2*!M;&Koi5KjL!Xcx&_<=(O0K; zN9G+)@iKEk`_qcy`N9Djtu#|3Tj!d7)f{en zI@8n+tbnFgmMYoB&7%vpeb8U$I`_PQlP=Daeso*%E6^}bk$&Lgl&J03-eC;@wt&R? zb`wo;#k-xvQ}Bt(f^CyX9*dw{W;QiN%Y(a|KV?J62!`laXY`lM5)K6Nge>zQGW*r^ z0d?Bo-!!uvK!26k@_gO?93qf6<)`QVfw}RBG)?Ysj*fn8E1!bL{!_U2NfV9aEY1fi1J9P6H5YmbtvNd?CvQj@#23%>8i`MwLFDbmY z8ty)DfG2QmxY_4W?@rgjH>~YiXonTPIQT}G@mRwG{pXk+{ms{U*ziLeUx}VI3Un#F zSDxar%0o)%BmVp)C^<&Lpr}?c!P-N?39zA9E!WNNf!0ke-y!L@)8R+OE#Bn{SB^}0 z-ro!D{A9?2hX)5;Po7$0NWmQ3Lg!M&0Uo9Urm`vsm6@#d?Ec~DK2=-Db4EM}wFOig z&C+5|G~x$fRk`-J_KB#lzLNHZLM@}AM{Qoz(8>8C)OC`)5mK9%DaJ6J@gE8@g;O*9 zM>GcQ?i(wdo_>=`dP-P5=oxVu-IdNMEvziZ$F9!iZ9HqnA(1;n4x=8CBv@+Xt{K*4 zQXK~2xHuxU^~=l`llD3wEn*mXt$Q=;H!#L6{UE@PSiHfV}!{Hek+E zywU1r@9XpQiltk*IgDhbr-!}w`0*>JCnd6Wn%- zDy_(g80jOIrgWqPIu+J36mBBp>Czi9l1)-AM zJ*jCA@dQW$iCS65Y+xm2I*CO~6Rd!Sll3hqs{%=EF}9Uk!kz!>C0Br%F+#jYzo!T1 z>X3INv3R^~?}RwS$K@TMA(&H&B9vIsHF;wmTC-{)S8d{EwE__h^;6&x-2?as+)qpG z;t!0(Bi1f_b62B+57-H0NP=&3z79w9j;1Ybb6s{B3693D5#@|ehA9+7cVhJ3ncHn)&uePk~Tm;!0F7wx+{*yF|JgT4tb+-7W(W{3E29NgphCu(-I{tapHq zFOgkm=xF)qyr0w9(I$&}`{bK-7QLF&kQDsmXuwd-dbIF(R$Ur8STZYas`95z;LN`c z9$_%=*{}p(jC3zGp4R#&vJl|G1d247tu4PFcs-XZAQ8*g)hKZGf6Uh+r*djguTd7& z?B#cx>;tmdo8O|alPF{LCzD8USNu#3vUdCWJr;qR12$%~o=btHrI5$t8Z?F_Xj7NN zP?koUA(_w2+udz4eCI(~{FvEG{U~1vz;l;axO>;yS7(xMmFIPyzL}D{Fwh^ zSTj2hvc;8em?~i52cY6&a3Ya?gV8xo7+oJ#aw1o7|MEu)i}5xvjoJCKsytuG&WM7B zJs`d=)#S0M96q~eiG@s-iLIx{k+{4|wfQaqE{kW-0Fr+!grvj z0)EYRd^4b*9}y1WuNM<|`ZUBx*>I(flHena*|4)-Pq$@@NaM-1U&4ONV@n@?@AwQg zV%2*q`%$TF#Sr{cU~Q)@idgb?gXXfyRw8>wr0*x`KT&KV)#)gUd8uwB-}O^UeJYo; zSG7HzSwnAIjcEzPAa1U%`g||lgCJLOSEHebTi7AKh=8__B-m!ZWpokeL-#3$h|lht zf1-n{E>(?sjwKW6gPF@()P>*1546@PJmI_^GD+}zqfDFoP$vZs#zi3o-eY23&&z2b zh%c#9tO02{Zq6yoK!iACUF@=D?zA`dXDeU@%FsL|n%$;ij&EX9fYhfD2JytKyj#qc!dGPn z!rTz0n4K_n+gDhJ_np{GD!b|#QWEKc{|WWEClWk&E45Y^49JhH!&Kq7dcUxPfS{~N z*fI|gn!`TMQXM_|QXi=1&}wN^2dsGI^4E?;xm~}n9r-IpMb~zG!!}L?W*PYe4l~Kw z;KfWC52oQK&X0#u!T?$Jzi|AIQ>2mlF_Yx)=XH+kj>X$r9R;LJ-I1&WM9JTnSZzCR z)@>goV2NBGg@q^(y7i8V?L;%|F~6W{Y`LZN>GZPVjWz~rFl&d6DLLR0=#K%{pJrsG z?6gbeO)xB(9P)PE%Kqp$6#io}Sg9K#$ev>kFyqs)pCspB8@`ZdTf$ATXlA@q7AF#C ztK>%hLx%GliVS%^jo*vg(PnZVz@U@$T)*7yLUg6Ku(pW;iuBQ?;1wc7s}~&a64%Uv z3^`rndJrP`_+xL_kumwydCa7j4|}a1o>bb6R}aH`_J;^~EW37p;yQsa)STHFb?5sc zUD8(&-Q){{Ck9tzE2vA)44$C8Cy?u7{QBZ60W*$ldRLcb^#|e-o$6V4*|~v1B_LhL zF5I@RIaaG?c4pz!>-K7f)n{0^#ADcN#Aop|vnd=U@aL*}DZ|q^#@69+3S!V#W)f<$ zj!(l&QNjVj<+C>IcfZl+pTmvx5ANo&-A`9KEIhlmb>&i9iEz6szALqEv(4JFUJ{wh zVUm965%>IR{MqnQZcvKmfxoCVU#_tPoSpRP1tsFFE8LX%r%~Cwv~)dbks)N_LGtey`0X8d^lv5Wy=L7tkQz^ zM82_7;DV1G2}T1o_zNSv-3dT!{*|&-j3+hUTa7sTWb8B%cX^=Yd-b1=S*b})wS!P4 z7DKs(sXUgHQ=ViRZDWHy;Wi3-Lw-&fW*_2{pBRXpSqj$!ECBn+nCEt=#XgNA!e z(0!2wR(Bz*)bXX!J$#Q;*_i2MDYQQ~O%{?N1!AR_}tG zzxx?&9GHoXe6Cd`^pS;11@$}q&;flXvuNC?9+8jZVm_b-VGSXqi<%HEhoJ?L}gf8}K}%r_+6zf+I6Bwdqz68(4;`bKM?b+h1& z`dN38D5e>cuu;*M>Ac!jH3*f3tP94IuCSOS37?pT(AV6x1?a8S-h6{WZyr0>Jpw&h z+cS>L3NSTiVxs@-?L{&BrH_ggJ=8guJ0*FPrzUrjXJ^QS(`ud4E+ap)icYAi8f93<|m!ZYxlZslg~gaz;z}pP>PX< zU`gNnrau}|w^Ji616VOxl~kl=C}?m={JI$aVtE=pG;Xda%K+AgcXjnIY`wUGo7(n^ z#XVzc^*XMsdoM9QqZvq;sHH={f!=;vRM0+TO?@z2sM|Pw8Tmmip7TRWfiFwc2b7>| z{j5T$=+@d&wMT2Ewe_UAC3m0~*k>*UVE;n1dAPNJQK*zPl}LYmO15VyuY{ zI^#`f`fyVF5-J(XS9ldUp4~?YuR0l^gg4k;sAyHh58AzOi8m=o!$s+U^hK99=NovK zYn}y50-|*cUm*h5sO!*>AoI!at(id64enYl=q-%%^%Am}W^(u{4s1xa#E%zmoFVW1<9l8+ZNbvGCl{p2A1h$Mz$>%;%v?eG3;G@L^1&2r<8?@}h z=pS-6$&(BW4YU~q;_tiSW)C(Whr8H|kM;Pkf9JHA43EV$lp7gHDru^H5}m%>6faKG zIJm9mul$R%*TiT&DEx8-AClgFUMd@K!zNqRf~cBOiPrae7gMX12WxP>O_UiFX6mrb zT%_4^0xq9#NL-e2_pcOJ<*e}9L1c;M{HX1!rM77_u@Ju@bg`p zH(Ku~=lf^oJ9pAkRa- ze|@Q_>0$wreGLOuuj+2+dc8^5z$sHXk=NX+HCD=Xvne%2tYxf*3i?UI4;qDjcTFf> z^-1gFG4iUxNwcbFg;JDc&C4mH6A6c-w4a)NC{eKtrkC5WSelo~bZ>r$b;VQ-3&}?J z1wBuzPvtNzVJy)`1*w)7-mWkyFvgU@ffG$N#l4$J-HW1@8OYU5xPQ2-q$eAhjkZOub~IPcTPF22j#pa^xX3ocXvfwG>XfNmWH z207iE3Apu_Zw(@w5l%U!$fB}ZzT7`% zfyr-1ob>BfaL1)@DPtGBn)WA}Yk*+QinSfR*9R6$jud7G2u%_5l z{=A{Vb*JL<&JD^fg+0yN5zB5%QNC>+@jsZLjU_v+kNm;XclkHKUEDL%QzaLP+-Ch_Osj}x8%Ts!AZ9q)xC(K0xhTiu_ zBK74w(bw4NU8u?nd7K#jNRc-ts<-Rk`;)YhDdqi1Sl+>A`7*MhEt|^1eM(6! zn$abgQI;8PDf?VqMOi7QFLJlT`*}Xj2t^`)KeZ`~E!2_=TVr-~YITLX2-7E+)g;8C z8D-1qo*l!(9C&!5Gvu}0DBKNS$1!w!2E0>7bzlJ06#F3PX3Dc$42NoLifpNV?)fwG zXJ;n>BjX!6?9wE-tXUf_=tcQ)p)BmwI(fRfKDd?xQaRmom#vzANQEEg)P!I2V9IE< zsL3ShgelIwWYjZ!pUc!e-_i5Zzp%vmLOEmqddmuAa%ZvhX3`#8pttGa0t6~;yn5Mv z?(?3;Zd(4*@`euYTFMA#zk7}8ch{&Gd58MPDe9i{aex*)b!#LlL-(KFRB7Y##-%-z zzzH45mxN>y1%dB_ze@b-{yxZ4gUXX^6X8pk+QoXy3~E=i_oRPrn1_nfQt{%C2V*xD z6OST5*r73#5M`*I`sv=|pGryD@W9@$8KHs~3?Z(pIY#jc+iKoMuz7giW6@7vcGVU= z*I(z#UyEX9 z`h>;eL4vO-$JwVNzC_##z&8IOT^qr(PWmC8v(D>*6#rxq)9TZrqL_MTOjX$Z?^UEI zC^_%&aO{sR=!us@Qy3F~&dUuh(loeDUd)FCO`oqcH(egojspb1<&~qmgWfgUGC++ObaHOG|vU};gA)&6Wf&+w?9dzIBE13M*&TUljTys$d3njnLw(= zf+?t70JO|MidH?rSfvH_Q1k?z-tuH*O_+A>FY0>5L3Ux@bigk z@4ACTifj#~IoaJU`u&wjTnt@O#xnJn^^}6N5|K!hw~Y%=Ks?G3=uryh1PYDKUklY@ z4LL!}B)`ez=L&adaX(=(v&Q!81HcJCbo2FiHa#IOET6GPddD?SEdX;4hQ>9h<^!Vy)#@DE8ypF_ zmgk9%u0#oDiF^5Gr<3q=!!Lp0oM6Axm*=$3)nXLtvXh5gNg91?7R$y=?w>35{(crFK-+oJd~p<=)^s}9tQx-mOcWr{ z6_SAOk*AUwiwkNh4~b z?<1+n5C4G1Pf6S`v8eD(%I6q;<#>#f`-T+Jx)m75jEpe>=$;q+%a-+5U7o zZ#e8-#xsQnPl&kyyxeRbql8SuTGDgDZ30Wepfr4>)_^kN$S)h2)xdcS#ocKgen2FDQqNd_@&k zPv#U9u@uNytxrFI0^(cmWQA&V+hwXG=mBH;7<;vsL&tzRLEXzi{m^J}9{ z+aJJ+ojHEC7U?aPZplm?8@KwVD~7=Y>X_*eZjn@{Mb1Z!elA)fJ~KyhbcT9S+Hx&T zUu9%}HwJ+xxVM;5U$uYVqKy$e`DD#Sz^}>_d{J~CAAM0NMiIC6n_4oN%F57mjEmbTV`EW;DrGes|Zg@T8F|MPu%DWBPaAY+A6gi*bK*}d6x`# zP)gH>xlYkqJ|Y{5cmclEH1QzTUw@*h#`7Y)uV4R^xADIIq?JT+*2g%;xXvA{u%iZ3 z{*`a@iVN9M<;LOjng!;BI8eQ=S67DX!sGKdMbvlnT zB_F-ysiBwC{v?7z5$Tp~fG*H;NI=}L$JLv26M=_Lo}X}J`2Ib@xXqeeIdP8aW&w^l zBL5^0hCmI*X&PeAX69QJ~$b<27xnAots3)3y{ctCh4SZbm0;cR@BF6_ez2{mW$ zd(p)P!>_q~PsKu2e{q=!FfHXWj$AbyQOT#g>eKLv_#8LV)eAePuodY;Ups{sLn(j5>O=QF>C(ua>-AH@tItza3WV*WmZW9%%1l__*sgYc z8E~5Mb}O95*fSH+z<))hUW2|vbVSuXSh6^+yc zmD$Q_3~>*mRA0_H6lNcj&#RO34=>MqPxj%rpJap}(cVc{Yi1ZWP>*Ex&ML43o||0t zG^K}8Yx5J(?IGiEeC4u)=_4qDwwar%M#L$M>H_U8d_>QJvK9coU15_Xe!ITKE6Ycn z+%r zGy@H-Y8?H5&1NYy=sE7n`Lcd@psil8^^l2S);}Lxm0tbhL|oAQ=}!@Pp6BY%9W@40 z>n3%@XZW=zoD{}#8E;w)Vo4GTm{7_82B~w_ms8$4`DBu9!RA)AZb{NIu+t_jR#+ z(OlC9##5FzJB9cUx3ITWx|E2Eo<$PGsg%BaHdvYzZ@X^%O`!>e{l%mGF8kZU~dZG2xm zDOq8aNef?WHJ8$+Y`aB6={sTyh=>?d_eNy#cExA>Wcyaj`?TiR(iZKfI{-l`QrK17 zP<}H&n5X_ew>Kro_*BQz!&628$OMsEJtTIs_sUqksb3u4n}KUI8vpS`Ib)?k1{Zc5 zM!Un7`ZZzV!rJ2~*9V`U{v_^ifZU-NANteJcHIts;zm)uM5=AS|6ur>HAeYiuk_M! zijf%d7A(UnW8fsyK&Iwt`C2Tfu_5NrOkwo{unC^4e8Rx!&FK|nt6QUlTYbNPP@bAk+QFx2UEY-r%A(vw^4p;?_oroDKVXO#>YJD zNxj*!SRhTXVtAU+PTXOckb)8r(xsUvB$}wl_$)(Oqka65*>0v#cQVjit5VILFw!xDyY4UEgDoXqfnyi??FznRB3i3UN~{Du{TxXwxIViMg^N=o$x!9z66jc+ z)kG<3r&@j;zn^dExygYz&LGtip9`|$Yv+;_?~LOx?4v(n;=B8jJf8jcCX{j#RoIcN zGMz!2CZjjkg%>#=cKDVp9wTq(QIS3xeNBf;L5r41IORn zznFv($2Cam{WWv9)QxG0ma^Llc&CrT*?L%~^8`t|*u+^7H0Iw+xeF(oS<`?e@fi}V z(yn;SS4wT(17kNzZ8J1>SkA0WQ+Kq<+9PbDk=*G#|56wcdm4*CP=eVU1uJJv_}#|? zv0nw$XTI;6aPm*XVqMWxLc?DS>zVPM80{TsOSN#muhIoL(D0PU-8;fX|1hwJsjRF9 zP!c`K!VU0YyJPsRvhi^j3j*;kW4~VTxf~?UiHUE6HLlQf;ar&*u=}TqWP^o=6(_Sp zW3!(lKrC(CtVap$W)&@BquIMA7mK_H(FlGFk>~3KE1VHt58+cZcViqUlX$)ZjV69ASvzv1WxUy|_Bs z@|X(MaO#=O9~`Jessr5*rxWdpRajH3{$C8`8zNQzR)$%y=o%>Aa736KR&P>I{D~#pq~=W z>TRlR;}YDWuGL!ZN~&%Y7(BAtZ-Z9g?JDs{xV&3RG)jZ10~Th%x{59+MaqIatbTwr z=y{3{r_-k7xEODWj%2TKI(w{qM`aX@B3UP-DPbeTlVe|Um`@s)FUhyzIrriP8~y*X zvs}`ykkyBi9VxY&Jj0U6T8~C^)H*U0NQI;f96;q}GNAai5g~TVm*aSI`G{tZztP`8 zj5>;uL=z2<`iUMgng$e_SDJNha9edsw!4vvoA5hSVNL(Ydtml$Qe4cx|3jy~z#=Z` z(??O0CTH%p|Dj_=zx~3;d37{lf$*cB zWjiX7I=Y~>?<$+~<~#3Z_iHJ{3j;4$cLB1zkgcL7^C_CAHKI05B9f6BznIBCbaiq++5tf98%*kiJAYoKjK8FQat{Tlbny=Z7NF>HqLp*-OqN72Np zAw4heMZx^lZwhHH+LZxRet}KjF+~$`32!-DxuZz$v>0RhD2XP=(sZIT{59&1;5%ZY z953g=^b6<9Pb}}ZEJV{^ndO;8A=6Z%eJG)SXZDViEer>3*Vou6n$i{@f1m>YUaKbC zw4=wzfd*PMF+s?8Umlkbypoa3U)cV|LUV}FS7?wPN;9aM0Wy2s)Y_WqHAYll_ei^q zdWqB^T)(Y}Yl)OT^wPLB%@P#_7EGel8gWZEJ8Mhy7Rd(X+*W>^LL@rP<8M~xrf-1OC!(Pyndol{O6QDJNbDQ4y_9^P2t!I!$raqqzuW% zuXK>${`~8q;HWQ2{yoYil;HpR&$WO5^nZVm9-%B?0lxhd-#{mq2n)9pq$CCcLWQPH z9N<>Nhc`2Vza%5|8m9(rp-3(>d|)o_ z>V0~5y2iTL7m}uN5Q&lM*og%T3?BlM6w$0~!ijOYmzmp@9}e!+r_FrdPDS|w3 z^`lCu?S8!e>X9tdGViY?UVx~{?3{c}#pxIt#-y?&)=~I97{G8fh7MV<@y{G2glN(s|-M5Me(cqgyA8LsZPwFRO>Qs5Ojin zCQ>XRummoC?zW?*Rv2=L@AGg&&px0LU6gM~>kg&O_B*ki0>B)0@%YhAAf%lZ3`g`C zAuED94wmKT5J#>oyxUaIyG{$w-w34(L_@$O4h+BJp zG4NVpn!Yg8^`~br7c+H5DfNsibmS_v%FDJ!+1Ig$+GVX+OMxmVGroX;V{*z!1pLL- zKfg=#Vm0PjSUW`L)@uLoQd=ZXbrr}wyhG|~&5c$m9)MU{6TIZC90)-@O> zPiE(K+LscE_GwcbFr^j}^mMmJ+`uS0&qB~=K~<0+z`%;$FW&NlNq6ua^wO+hu;A$$ zDy+ac41LGCgV+V7ZrF?R9E2{yJD(?Y`>tc0S2hNhWI^bl?nWZy=W3ZK7X62AHpsVd z$ri?BLHR{$ey>Iiq2b0jx#I{0T&ZwzRa%S6fMrHFX?JCt5wm11y4%m+u#PGx@#^{w z5&n=!;A#{k`D40J!2chvq7RBV{WU|`Y%E|?5^?9Dl2qCH7J66 zuO*l)eeWyCYALX8#?TMG%0 z(a85uSWD9r781UqJ^smu`)DqBAO6M6!96-GEgYX7q-*o|iP5IK$CuxjB*4cdpV4=n z>1Xx-J{_8K;(N-=vLsJG?0*ooGTW**lXb58UtmCj<>58ErN1N-=6M@UzM?y5UYnh4 z(QS9}#PFG*qV(5Cpi+^N!6A72@x6U}4mpVNj#&x@Ux8-<>s@qQ7Amcf9r$g1Fv#oYPhIPe)`vN7@$g3znQ|Ddoc+n;PHfFj1ZfSF>;8qhEe{|_1o zIek)JtCS)bsKtO4(tEF6ciR@&X^eo>^pZ5-LA@kTEI+{^M0j#;tR4!z4zlWXb<=}( zgbNrW{VznU9G@bB^w}!t&Nvv?#~1O<9Rp(&Afh&GB_|F} za@QLQxO53kMlW~qDkMgRd+NKJslz_xjwKiAIj zjaF<}4E^Z0l-#X@Xx};90?8eDkaPtJU1J)E`Y*U)1NcZ7Rb*O?kD>r*27R!W7*^;D zKFrq~T>vyeb89H~qsumJYG!Ef$6zge@N%GGm)06KAgQ+c_q5wzaQRQ{MqS%*E&X5l z!%qw!mYmG{>Nn%WvQ$Nka`6-2_ME79si2wxL2svGQTm+7C=F^M0XjNKp}-GluUht3 zoo-4_9zKHr`aW5Wbq`JlcN(Mqt11*&f9iD*_7^-jSv+hkYZbH4&ZYa1I?1GdQPrn*iQ&$#qO=d!3yC@j-`cov8yWg zBlxJGfrB_v(pR!?so^3#4;8Pg7u@KJtH4}2#gXbz+k`B9LfkDZa%Vx>`${55^#=uC zLNb$Z1U7&_S72Dh&%SY7p6-)f2@i`^o@v>5f*J{Ylv+)Z5Xn$227&_J+(+~xF#m() zH%^}sU#EAH?X4Z@nm?YL%TI#zO}?MtWL4m)#r&~_UFRP4X5sj20O^z9RN4Zr>h8)a z#p(c3Ajt~^`rCqqYF2u>?EmN-L4k3Xk8~#L8SnINQWR{{i^+FGfo*omp^bfG`(a~$ z7trD-yjMq!WrQl|J_9%R+v=N>xI5CSt_cO-q_Is+rytak5q-6p1@Of$pVf^Lm8E{= zeGfL&aUPCfqc`~-Pkn$xzYM-2m?O&g^g!Y6?Rn6%fMo#j6KyAwefhRJpSv22d%pYz z39=@E`8i&|=n3IK*%w$LdId6Gc6_)4BpB>OxQ7yW$O3cfMss-KM$yT1%Mg?c#E9WC z>r5vW-(RAMTA+Nl>`W*D)~=Fnk5BNKe^Y+^3*xfZg(R_2#OzoZqffsQ;O%1PzjTZ{ zeD4w1W@hh;-rtk$7YK@`l5@C59BwokuC)Qt%SXynCZ!zw^t!DNQpKWBME^)X^dWN8 z?aNq3j%*o@zYA;-=hAcu$((1tkKmvQ7YgD+PSC_pS=erQxcA1Ly4u&{1_D|0*T6jZvW`2v}JQjv~CgFM~FXcAzUjzkDC`(w66U*)RPwN`Ez|b2gB{64KM7(uG(bjuN$e8fd^VKKV59CmHI$fnf0QSNRf8H|b^z z<49}~D@1g}_obD_7PYs>A%6Y3f+WUCm%SO4GU=lkpx}4&1^e|b3-dP%58-{>s2`cQ zxA6P_UjHc-sJ~^%iPS3 zv}SJKY5atoR8934VCG*5+(Ej$anD*pP9Qet*&BZq=dhG2qDY?4scT|h{IWbA!SU!6Ue~PyZL9^Gq$PWvh zJ0KTVFxd?maJKxE%!y(;?on;E^28?I(#=Web?2S$ksz47CA9qxoGnyKUZFI0fU>1k z`}$hd57J`RXA$Ls8H69fZQo2*|3)g6&@DFj4D`WWsoyxBjR|Q6lw<~s6ZHfL0qwa5 zvPwIQ9dw-#PbxPV;VHEEUGJZ3`>588WFrQ9KJf*=-P3qqxw$5EtI7>0u#=b6kQVI7B-=3xqBAAOhXS zIPbHe81~iZsdj4|{m2xj8h=z^qBcmgbR8@zv6bdgc;n(X*0y?u3vg`Qym;wQA9_D+ zCSV~8^QDf^p9hcEA){7}${kYl2q)5e8+R2Z)QXQ%+9NW%n9BSTJM>mQK{?tbu z^-AE2@-9MurHr-G_=f7of%%NSq^QsDx3K@n~+{l_7(as1<`M!U^qPJJPmg&+Si#!!zeRm(w$Sz0>?=`%Kl z_n$(S(&rKDUl`ZH%e#k*hmkcFod76;PM)@{S$|1ia|MPMZ*xPqcfTP3kZdAa-AmbSgFEFFiymBbn1`*z&}otQt0&WUTB0U{ zGryo!O`~@VW?zMRo|(fY#6opJAwu^PH+D&BKQnusnqWW?Jd|(2Co+Oe=SqtUf zN*e9=`W9)~TrHZv%cEpy=UJWq=*#LoEF}EUccF06B5526&~ey7R#SM1)%Cn;(RpZ` zhemnjOZ6SNSRU`gFFS-YN^BZrQZcQlmndE~DrSARq@`)S8`&|Udg zA>OEtIEb}1rlhwvB=xA(xvA#=Qi` z6-`AkdP(t!nNqV0Cv;kMYD%rb%2XXTK%o+r<It|i=2Im&W1hGx^#`%CS4+Cce((+jj;R;S#kOE&AqxcD9FE5IL+eT$LY*i z0)UEua)h(psiQ-B>Hl;cx}SDMnl_#ZiYNA(Tiz-%r5g+cEl?>tm?nerxwb}a@VA^I z4{mNczbp2iOtKuDhC9#7>5lIH8j|S21ew+Jzzj`FGNU*Y#^vFZ*4aQkqrzs@Z%@dl zwf9Csd)Y5EaI{+L`uQ^p{?mT@717zlOU2qZotP){TO;){jl4p+*?O4~mSu#yR4?zXq*6YA^!e2w z@hXOokPyIlq0(7dk%7u>mNU}Y)gzvt^he*V%;o7nX)JPhr0mjnJ>~T~Anep&8L?Ey z<10vf&Irebpn%f6sGDk?N3zj_pU5FP3l(Eb74}yVmzKgTHDO50Q!mpwP-E@~I?Cqt? zIEOYE5b!^~`O(J;&W=G|bWLKYVabLEz0P`NZ-w$Xv#cWTstP{J7eR8u}0O= zzk#@Mt*IgT=b9fx-wFEDn}9yy-=*uHEI$-|`%f%DUF+jx@Q{?~xlEnaU}zeGlM{X4 zi_Lv|P*Wz?M;lSw0V*@5(aPg_~6G zMJ|Z-QL_z>uP7wgNKJ&)kcg`<7t^t4+;cZ+2!NAkGib@Khinw?rF9PNzU~Prlu>wY zFp946kB4ObeKRYdQw86g_B(1RZ^!;v(F2JL*avSk!nh4*52=|#q(BGytLv*UQkd@m z=c{khQ0@Jx5PL4z%!PYXh9oEr1GSr^$wGpkhJE4C`)p;9(>_#}e5Z`x+8lgyb~L2} ze!>)p2ugFHKhAV7fk}gLs8*w4a98~XVCIn=K~%_1^}G|A0%X5qFQJolkHO#_Lg-9y z5!^nmV7!D)@c1tem}sy3k>u)-8yn(*hvvP_Lo>j|Q}2htyl(m12^=~AmLJ;E$)bEy`E((Bo(A-9 zH&c4tyl>=vwv2bd6_}V+URc3YvlcIR&$09p@U7EDT1t?BOzJ9>PuHn^eiQ;!z3l!o zzw5w->rw&0&#Z_dQHs(q`Xg}49_njfLB439=dlAaP{N8^O}h6ZeMOs|O8SvfG{H-6 z(-Kh0(MmW8>Su(WY$HPv)U9phe#`Z2)SOpSV^d;D`=#=@J`5l%pOyBRMofreu+YDK zT`I1lgX5HrwRl@=!AI!x-Z3zhF6*ISv6YGovlsNcd>bBv_LV{I-T{0GctC%Na*+}Ik zACXhXB;}y=4u~s@^}gR}+!qr}w%+9?6`@fshWIt{E3#Gs_!UKO4)R|_k&CwQ)faI- zkHAa7q+*~S5X2jqbx!42jXjrl!vK%EsJDT1)%l`@m%HWG2gvWaWO0Cvll_)3295`n)=F7Mja5S2@)@Hl9b;aXI{E!1a42}}Q=e^+sKd~?HAasF?1EAP_ zIsDa2fem1kMsq!)mDi2IgQd%z%GX?`j))(=rh{kXGreE#;D-N^FfL!v@e_51viQe4 znFT@&RFN80;2Sj6D7PrhR*^UW0Np55SI710L~zqCQ!2^1_lR@I7BOJ}WE$rQh6gNh z{q2`+m54WyhlQwAxM7U;2#4TjVY+Y^iqRvsy*z4U;Y!jaSVfu;>;?M1%eQ8SaZD`b zj81&PqEHp+iT4cKMvw}HY+r*RF;Nd}a!z7`veUM!E|s&`g%t76M8DkZUsS!ei;s=x zwp+MzZ#>AFgl0cOg5}_&D)CG7LJa83#1@_*boxo`ItY-(AE86*2?XaS6mW$uA=T!b ziLKYfB}AwuHA`Umv6A8PyZvW1ch()PQ90?z}_v1 z$AhYBi%Q_|4rfpTCfz~cCwGEh-_YZ82!mkd$*&*$!^UEP3Q0QvKT~XJeUik{aVLFO zIlPMPCsasT;umg65)6f@l>x?N1`+5IQSPzt*AoIv1lwrebap<%YkD+cTr!hzkkyKO z-0*)rZz~;O3@2ph-tX0z$TQTdB@}ZS;Ei`ltiGn+b}~&)I~GqdygdeQ2OcYx|4#kl zoSU0d6Rv6!9##!HjGrFi&IKFr@iVqcS?|VD3+B%q%3VVS%4GL5IM86=Li}+Cm5b@Q zHgF#%=G&?mBzV`I$9q0h*$2J6J*M#d<7SbB9#G-}nN@m~=SUR(mF$Gr?|UWTzLQ^F zBl$a7$d#4|h3Kj-o`^MGl<7AXH`>+7I{pPqL2n^eUjq_UfO+eSf_$@TVIw4B9$B1> zRI~kjI_MWnbQH$lAKAs&N(MWhF9uN%r2~???f1ZX9?zDWre^>xZiB3}X8Av~JTl08z}MwknpVx)PHBIW22K;jyH zD4)VO(zekw1Vn*$+~Lmrvaqw`ZBSN}C`3EfZE2hXED6m&Rz&DPNG_gAYN>U(WwTac zsAMhzprQNL_w-7wnB4Hg{s*K_JK|MeYF{;j1g#J#YwanNb0IIoJ0+#ONLAIgN~Uao zl0__m4*ZVOZXMlF9fFb^Nu9*dwt?~FJN`Zy=-F~q2_utgo{R5a038+rMi#4njTsY^ zNbsaaKy!6NymxdblOqrguL^6G{gLI25IvNZVGHS_l|?-I3s=M@lBJE$6kAb&L!_r>wu_$ussbB%0#jN^WDp z)y5oCq!WT=d=9H2{FuT+#yTrWEk6J_DSc!A9Vc9qMkeM&DW&v^x$tOPF2R^8Us63C-4PV5Wesq6dh2l}ICz&Fj zmn0QUM*9K&JKx9aZY;LJ`;!b~WTon$K)2r-iBQr@0|Sa=VHm!P5&_T=2BwHQ1+^J` z)oQvPzl@lWRIZt|9Zi$*Pei$H`fZtYM>8!|P+i5$y5NAtK~k%$p(0v{b|S6L_cV~B z1i7cjz%OAuxQ>SIia_K#%^ztdr~H!i^KFT2V+U4)DLa+16?sxDpQ*yAsqtA0Vl*#Z zMn6eT(ql^G^_hb+kSR)~dH1bq|9M^9+};0xr({@pXRR17JR%CHl+%Xu3^ULZkN~nc z6!o7kVe!hQ*w;QU1%%{kqAmIJ4%sA3of;@p$#6lhOw)3RmP$N)dKhKSEyhXI`!0qKj9cAB`Xe#)$I8$ao5 z3Ev+hGR)%4Hh?E5L9VVn_9bt9*+(QPfy z@|Xh8hWp1v?}H%L4~7_L2Mpe@Iq$lY1%UgJyrwz-??na!+E@Q(le5NFk5sv$ z*VV-%7|&TOP0RRB7u|`JC-*l1Xftb6r-^Q{u=JjtMFQF^8YLvCwQDlH7VR(q!X=ji zxgsz3#Q+o;Kcdg^X&D|Ml`dwnl*F)WWqa-B+XD4R0ZJ^Gs}~~1 zPqB>Pv#}cZmhEeVQ~k&hpq*fBrZ6nzb*}mB^-i%i6Dm6M3kQqa)M77b-sam$Lxs1S zk%X&h&9~};qRXx~F2yXxXS25+;Jl#@IrHivoaI0vQ9Bgk%i~#&Cu5!b4?fR_JnTTv zJ(Z`6-<*%h*&k>U&0gycq=EafSH8$-Mn9p8`+Q_wK%40CDA`BT7rkxyngp?e2#E-E zklMdm9>a%!mV|C(`@nm{EdqssX`6?bmIp#&My+3@+4USH6&Arx&563WTF+76-WlPRiEY)1oj3&Jb!NJnlDV~$E?48$|>Jz z0W)uoeOFsqgQWX91C#q|&$KO<%MlD>NL%ugWXj$#e2dYIW!K&-zUH=0C83p?ccQE# z$%|xU&Zig3z>b+Ga7yHZ@UpL*tElyxXqJ;qdo*pQ2Tv2LD@{eTZSU5Rr#^mHGRc*v z5b(I#{+)_<$UFVVS1(Cj>*XV$p`d^6)ephxzt>f@;1U+m> zcG>CTMs%Cdn51MY*}60|3_(FT-P)th%j?M03$9?Lk~Y|wt%i6yVuf-omoV+JW&Xm2 zmz1rz-c(Do_Izi-e{!La)tG{xwD%2zo(7TXHIM*vr>)*4a$s-myViT)4OHhZx%U@( z_K)xJYH@JAO6q&i>?|kdAlf0G(40yH?P2uDgx|bQWN7KsF(L#m%QKmw+T;1C%U8-$ zaPqY3=iUK{o-GzPyfuA-e9p=p9}vNdk89St3#!~g1(s3bj#I4FAAU`-$;vm2MKzY_ zo|%lxmD9fnPo;|(jm{Q}vzruw)*~832>!g`h0j@yE2blWT0dW-0Umi9MOd0ALN>e6 z9c@5+mt*1y^XEmeyMm`ob?}yZ9l-^E zL9I1!jR22Wy=A07kaNiNqj&#gntr$XqKXP{@pNMxQS6z*lxJlSMUOg-UGEm&4zejF z1}vB2puE7L_zEKl6cdhd;BJYLwWYj@h3?bmE}}=Mw56f_nuRrnL0&dvPf_6q#=1QX zW|JLcnn^Kvtp<7pOIWq=0Dk(!SQ*e*>rNM(#a2z@%+s`wz?}pk_ywqb$5yB3?cBM7 z#G!Ia3lmz#eC*ve87Gh+IdC2ZU3n$({sDOc(;xJnZnmkS-Sg^rNFY$Hl(wm=Bu`w| zVak^Yxfz>IOXl8mj!*F59k@&&?MTgGzizRUk%T-qBX+4;Nc20zaO!;}H7hP(u752{7|HA% zGuGDwdwwz4@XwktM_8Z&p!^BcMt-N{ELn( zY#yD_gE%8rbWNK7lh@QyX$I+d1)EUw5XqBqT;iLNu_|4bbb>{f;r>-DT<)q6DzD{-yPrEMgzRh!M z*exBO(m#tR;)1Q*{L5~=-_vx;my-`rSP3K#W7w3Y3(p7PN2evl4P0)1mWPEtGd%h{ zuHCK)?c`4qDbCoqw(UrZP1>0_)bBLKQ{h;9GS<{JXnVjzzcn(4;{@Js{W?E~F4LKo ze;WeK$r;Hyg=#qqa&n9v6=X(3)nCcNm^9iO{GeeFMl0IHb(f7y+CM$lo8OP!QaIFr z^$3gH!!_Haqr(n8WU?3TM|BeT_pz?C0;Tdw*5;lD5=X{2kB$?u=syAjJl*dx6>PrA z6qeDOFZ}uqUp*WtqGdn+`Y=oE?K4KG*RUGZdI?=@Phav`T+M1$xt7)>hUflm0TSuG z&@5w!nPA)7`th^q!nc!}HIt@Wbq3+^YITX5w;@uKfCd40&B4flJFy{)>KqP>`NhWt zhEn9;aq$zeSjAEL^B5;+Eax${e>`V1+nj2SGnRhdt!WzkTu3WuiX4tHu}?}$8#$8G zn65RvDyXgDAx-=!EW>q7kX%flK37F$pNG)TpBVK7MG9VaSzTSd2*@^2&J!&9Fa|eGs%lc;#mn1^l8tt1Z?2sPl-`Jk-=2;!>y3`z1o7_8bSQDB0LX_;f zWSOf7B?83`F?&0%f;<{uSs}Q+>xC6w@_Nky9SaM|@;N+2*Hmc>-dV!Sv&>}lGeZpq zx)o|XZzdGG(cS}8{tM+=he*(N%BXX-N@>@fd-r4Nz*rjjv7P@uD=?hL~qU}ox za;duel%1KX0a0S&V?+O|sHXHf$b}(#{F73MzwCO(fT-Uu=1utYZZ<4V3}f%@JM~qM zJ&0F(%NxwnRK7r6Gyufm_G@pr{QfF6h0vnVea+6ybz%0R8jgcUq^qjs9Lycn z6o^rNkuA;?DlL5_u?E@A|B@S%%}ouxrk77w&$btCOO@oCqYh(;>+|+f9|qeYJjZwG zYb^EBRyGqPcMljh0Xdm{($E=h>{SDcQgeq+e^{*QnrfY9GT#opjzqhQ&u4%ASj*q&4TM^Ck7Qb@%v8CUKQ!W+#vYv`+9h9zFUoKp;KUo1G# z41B5VIWo)$u>ql&LWJTP^13HoenSc`HuN=3`O;>%;rOPmR4L|Ax-)(*@#5D84wmTR zVDl{Un(+d3=Ffjm@ztr2Ht9B>O%`CnZ`qkWyq@#PJQ&Rt(Aek|=d<_AS5hSwNvv5l zTa|uusCA}vxxgGVw|0F<@{3&7U9J>{1~!=Qx>;E@J)^a^`ql~+f6EJ#)oU=;Eomw% z*U8J1!6#Z1aJ$->xxGI42Tn!a9&4C0KR-I0wMpiGc>lF+2-94FA&_E(pqrGgX5(`F z`3jD^>8$i+-Aw6W3jM700AHh=77un_|EWWGpdqfoPX)f9R2>Ivp)Ikl-?Cs}{;NAO zj+Cd?Drw!F?6RvSrtx1VT~C%X-_cDtgpY>SA8`074q`FbFsmTszOeMVa`-LjzGLXE zV#HAA_`-EtktADVT0Hh1Uq6@b)}y+yd5({JW-^!d40(RxYC$AbaR*9(+ujRe$G((E zJztTo3;k(tTSGZsw1`Po`h)CfQv9uiWuiYkQivm;*Lh}nkCxi2KQGYiF@iZ1cYaa4 zcD*15Fqg=$+;5%qIGOk{agg16Lq%kpr)8Jr{oPKJ>#F+H2h3hu<;0FB!leh^nl;kmhw_ zkKuBd62W2b7U4-55%I6VL{R06b=50?*+qIEw1+=h#~>`Kn8;rC21~?^0HA3!8eT!s z6a3+ffbyA{)zWru@QDvpdE}A#F|vRlkT|(+V4Pdp^x_6Z#EmEQ^{Lp{=rKqX*vCWM z$tw-+r;5Rm{z9I5ZC0vwQys>X;?2uH+t$}x$lwypzx@djeOKQ&&)-dbsmtS-9=yhjxwJ^gxq6W<+mD*xOfs1l zh0xIMo`Tf8fe)^o-3beLg#3+UAJ>k->OkC)(WPge*v*_pbSdj`T#~Qob>gi40Hzme z^J%T8^x8y1Gi0u=vI_t+*Kknqr@Ta=ge&s{A)@qEWwp3r&uneb-o(Ce% zrv4Rq=ineMg@RW5*H02WL%8s}Bs;GDjFIf`c$wpF6E&>d7xNAKcP*crByCkSm*>oJ zcl3>#%@o&Kx`G;w>gx*|N_XrD()n~2&VN?p`3%X*@8U01!TycE`X{Vtn@ofY5*SyZ zYPyRc4bU4+GJX?YL~aJh|JTxR;g4Q^Ym3eQ3g~>cTAX}T4H3d||6RG!2sYJWfO79N zU4#`Cmh)1F`S9K zE_tUhnr7f>Lu)!?;TfLo@`yM|*p2+A-}JDoD1~=g^t%-jz>l~s8C2PKxQPaAHrT=T z!hybL3ikUt5B@g&Vt7FMhFo+xfZqT&0*cqJ{#Kck_Xb&0vg6Oda z8a+6K_6dAE;wWjNxQ%(2XtnSh5(8dK)2W4Th@>TD(p61C4pFIoyzLD<(qj@M51mbr z`G~^-(Vn|gdOV=A|4(=!gwIj=q-L`^fxB2<=UGeFyeu7*k3NCx{45E|m(P6W0`O~GqWr~5eJ79;VtFPte9oguDL)!$cX9ubdmp^qqsNT|*frRyFOhMb6wOQI z^)Ohb@`XukKSvlZ#8!@TIR!tV!Oh*%O5N0Xx#o8@fBO5caP(^s|Xj5E|;C*1L!AfgR^8UaZ$U2x>8X08&I z%A)hlP9sVW!I$kLJ_)PU!v4UY?Lk8I%f=3ifu*IHu>+s4tmo@r*%Qb_Mq56!&8Tb5 zCVap_ZGhf%+|OaDEnO zQd3Q!_=JF)7<2(M4<;^64qP~~W-m0f6?|u7WvG7KiE8Lq7 z6Oi}!_KKqh&fTq8d0(3jvQR^FRT}8a?YQcB_COJ(rLEbnT|v8h?NOp}>gs&%o6~pc zqpXPN8EfsND02Dv51?Z__Odq1SEUBYrkWf3zNNM;yBlN14R}&E<3LfM{*mQG>8a+b z{IDh_a@Fv8K|$fA>j4}e9DRp4Qh&|TK!Gi48%crEe6*fQN394s`gp7sGAZAZT=kz= z0P+06WaAI+XD1IU@bCeDVNx-HX zG1k{Jh-#q-?U&MoMF!!ka`t2TmdhiCw`V)$wc=g2kMUHrKy<8R3Nos{+0bqf+Nh!8 z_hXAJzJ-~qDY2=I&Fp*&*kGS6O@SKtw|z}vI1azYUs9Ox9YoqT8T;^MMFYE5a#%-A zEXiI=2_g=SN>}Bxvs#@Miet%rQVDlxs}&xXg9;)40Tg1p-~*h+UsR`U#r6GBpl+D1 z^}yn5cjuz2)Ixfc3Tlj_S$bcYrm3sYZr1#u0(y_tg&1JxyUqeSd3k6Uh)Nl9_BcA_+ssV3AV-5ag_^zSv~kE;`3-Ot{% z^_%GVw1cRKE~ZRp7C3iP)p2DJc%?v`rc0xIELF#v{n{#awrkVqbdeaj3r#3M{co1) zh+0#g23wu-L9Tm_!|~RYKR*P$Wc>*Y$ag=umRhQeiF>c{ZbtYBur6x+9BOEA%#v0O z3!CEptM%ccl2?}`{+s@j!`v#z@YcSqJFh=m zs)U3Tbb7Yd`7L~TiHzv^ukQ?f+O5sy5oW}^2;lUF?YuFg*qI*P+|0jw1;08yXpN9-P9avYa%fBTHF=7)Z0fbJs?u#8`ve(T_U4DabMT zE-GYZhMc%Ke^Y+hPj}0@rkr(T%{9h9T7Gp+6t7x=hk>ySDm3+2OK2g_!zS|XFQ~Gp z6D(HZdyR75#M!K_Ul`4ck0>`;Ijw5`@)@1zHiA#p1O2EG2QhFO=_J2!M(IWJ=iO(x+tjqJCYvgR(M?4Y;2j15 z@Rxo{4iijSA;A4ocxWZy%cE9{GweJq}l|HxL@nM3<>6Q@2nkkfRc35nY{Kd<7br!vxa z0GcgkAA~sUUd-KABeGW}-n-&dtwv0~c4elDAC4*VLrzZ4hK7cYj*hwO;h{S4?z`W;`9^m2 z!h7?$c)~)fHH+_J`2%n-Z>}b7w|ewznlZ$6xaWLoIbOHEYci)wu;KO{UenVL7Cv%R z70pJyG0sFS(9bk8jP!qgc=+z7ud;xNR_N$>aV4)=7?#;_C#G`zl8MzuoT(VBKnA!ul(XD@JX6Y6kLNu5YG@#X&FyZf)qOe=>u z<-0}oC{u^4mKUb!y34h^MLSI*+qk%GZqA0abBzxVjf&_E@2ib(2%S`fLyKMp>`Pg3 z?dhI6doPD~1;cmS)-s%1IOgJCG)Z|86qPM3EcQmzPc)%@8Xet*h=5;=HUaN^2LFCs z!UjR`8`cX#@(0v6y#Mz#v!mU|+DBpU|CNsPmsz**-)y%3t&IDiWJ3S9bMPTQl%NMV z{rdkXA7C$3|HlOY_)Pr&SN2~+^8bM`D>^UKx-H96)0-`=aE3G4XJKJ^ettgI{9n{j zhzjHWk(ku(a^K255qPfo`Sa)cx{kXt#4oG7e?yMzKcZ^q|B$F1xY*l2Fn<6Bx8UIb z%*|a~fVK-04HzsgO*MC9nE7|^;{O`XX{i5qw6&>o25C0(aJ8qYp#g5J1;h8^o6Z0P z(#c%0*eafKjRrlI zOFXjD84fs{HPp{f?U-p%G(Qsr))W3V@?R|-Vq&Cc#(F5o{Nqdl*bHs_U;Q{*qT=PY z+wMOq@R4F+FUQ>B;_of2@do?7?rx!+2@~^>$jCOMuR@nsS5uRd**Q5! zg@~~O*PCHXlvGrFSHrC4vn8s{F17{++nYUM99GLJt<^JT>3rT^1mJkc|51D{tupT| z`XUH~$tfs$X6=sWKN)WAN9WXj%loRRs7MfBX+BpL7Z(Qs^Mnx6TtM@qLrDEk&*iJh9y5XIOz6tQs?*MhXa?UWgE`_YD?k#33}!SWaN+0 z=v30wO<>70FffqL?R+{@tgNS}r=^uKS>|A}P?cBZ4koA|BCmU27{VkL!$Ch}UXa`M zek_#?Ez{1;gA7!yHhOq`oV{@WUrgdtjDM`jUqi22B_x#L zvKd1Ac>nZtD8^8TckFjDz?dRZLVgI3fZoy1u@?qJr7k*%|c{Zul6_oras+ z2ibp&EX8I(XdL5S^T%0O0`nbDz_W6ba05Yek3#u%+hWOkGqqL%RK^9 z^ZM>?t;fTq_sfxXSy>s`x1X}JRhXzEB7IdgoZQ@D!Bs}3dQZ0}pFVxk@BMiIu9=`& z_kWMD3QzDDqiU7GM7Hpt6DFKmy!cFxNMuqHm`V4Kk3F96Z4Jd1v$IJl*zgDlEI;F- zqoeoz#L2;r0h3ln2CvyPwV~m{=}Hq71qA`G$NgYDov*KNZFC}|rW!kUajxA*2DM!O zzpZee({%Q{9H0VIp`)E$(I^ghUPdx_)XV4iZl)DD9?SmSGX9oawD@3^9yd(}M@>H| z^SSXU;k|JAA_w<&>v3RBD+8%9U@zfp@F^#SOu3)-d%9Y&^EIz!v+Ne@ziW^st8?@C zJ}?>>6ciSWmR(9JfQN~RiH;r{6N83|YJ0xk5kbIhR-;z>9S;rd^69KiCgn2?8=I2Q zVy$KIwe>>PPNU;db91u=nZw!Y%lXFt!`fE|M76!`qMjoba1a3{b(AhaLfRNWLAqP% zZWww{6b29wNoncsu2H(XVFu}D=o%RA!Z`Q)e)oR&ue<+XMrUU2wchpC^Stlgy8%48 zr|fVIyvc9t_HR8sn94UW7HgcCs5b72cRSdSmXtIbNRxJySxF8S;ktdpx@OC#=nDJq z*>F=$L2(oE>zl0!XRKuuVvXACRor%Uh|X6y*t5PnsXgtym~7=N{2J(pDRrP#f`sxrZCJ+YK> z35L|#NMi;%k%`eh*<^*qFWc6hmWv-FC0XgxuPKJ31l_9wS|iC-EIZNKaSbHbKF3t@ z%9lOYqA*vAmTG4CJ3^1RjlGRHWT%Ndri!~xD_T90;ddVOBOg_@HC8mHZuoXISQJ;W z@|R@*85?vC82r6^_aOaw+JdRIis>@hs`)an&W~Hf^UB0+9UYXMdhfr#?p<~`#U22& z+tXBSVrFXwU0lg(&JV|c^6y8RK6 zk&&MhQNVYN7V4L3tn)}zR#*Fe_|VkU#ADhIg+lJ1VH8Jz7xxPY5Oi8ob8sjTPnI8# zWL4LDc-yEmiqmfHg<<#@aCgm3P1@F_?MY>laX~>rlBb4(jBOf$;`S$pTf4iv*JhFT zC4z2Ty$U>E*IqZ_Z>`#ErFCg>b8$4?7=JZJFvMHxUplCF&A}?Aq{Pnr9;9VkmbqG> zbYBN4=;rLy%`1D`k};DYJ~AGV=%M6*dl=SADk)y1{petw zNw(KkKR}zBTcE}^!h@TlKn;wJoi*&}l4}I`x*S+7qRpD-3GVM*dafNE&m>U2+pnQo zx}$dH7G3ezTB|?-r(3#mPWxM)1d3mad!k-$;;xcC#I8}&+Sk_?5vR;}b-cg7e|t<_ z45++no>u0-a(zRCu#iw#^S5tLYITs-Q*{l#M6&ThW5fLP^x{yT8+Veub`1&I%H}RB^;)J5r*`;vp=oo8elZEloslG)!QAJt>viajF zu_>ymG#3S&u7v?FA0G?^X2~C=E8{BpIv^^AuY=V05TMXa)*0*tZn##9{~UI#wP+H>%g(@79GOGz!@z&K$wp*)l8FSt?h9GWJZfNLNUbw8Z;1%~`e zN{S?;v$HeH7$l;!Aipd!>;OO8JDd#;XdTQ{q^718OVKR10y&n5^X8AvFOOgwlSc!x zqQa-U0}w9D@e+Z@+P`I%0cWSkN4zM>vn-L2Am$H)S;%hCLi%+^G_vh%6$FjKjaY0-D)I(gWxGa>+t|05NFl=@lJye zfufE=rV^%Rr(4syw$ud}Yz!YuAb=W1`8u`C+3LxCgZ-e$`SW-pylMPBb7A}kt=cXf6iEz9rAQ7>jvOuI!!wlFvM*I$2WkreK2Z@VtW!9eogogj+Y z*g#j<>KncpZ_feuE4Xa@J8I$Wbz)p3FnvTO8{;`4#ky}R3r5Xug~aw#=`th)Rq;aX z9&kKeKC&3tLJ0{*<)jVHl&pPTBc!pUQW`o{12>oQK~$J%nWY*9YTSp=4c`n<3k$qy z|I)Qh*F9X_}=_MInp`Dh133vFNiKp*UC?oi5g$R_yj$*Krc|K*A%XS&Ja4s3UCg~V0*hDFp? zO2k_U=s^1=^BCFahCknw@(UebJ|#N~RsW0cZrj#637Kk_#+e32y*0y&B zn&o1P$nHn;e709DFvRL6>>`?B^bl%?TxQB)10MCp{=Q^OuVS;adN7b5<4e65|$JU2O@)h~ft(>-f{0qVH!IHFHA_rn6EtHifhV zG2B_+klm#}LFZ8F7rpXSh6ZJcx#@)@qol+}>?i=<*7m(v2q=8uw?*u2scP_#jkbmLL%4spPG}I z`gb!BIEStq`9-{Ps=QBla`&wA{(L%KRY@;$UCWpF_~}0Tg(oOUwpV+DF0L;BS3-)d z^84}Y|7jSN5e4CY7yRqZ6X7TSCoPI#7RrjOHUAXl$g-w=adYJVO`35np>5$~i*tP_ zzErYAg*$id07`p5?;@C~wDcZ`4BOj9|FTQr=M>ZY{M_8!!ov0&3c&K5b3cuLsGcxY zO+fW%Wa0nEG&ArKJ^w-5Bwo1pvt;vMzUg(j@jLqOK}3@tAP6dKrk?>iC6rzo0)c30 zYpgN*up&oCM`fzDH8qncl)1=tGgP6r z?nL!pc#bTBt3?6BsG+7V^asFC`JS^<=hRe~s*;wRR#XCPck*w|RD>9_QuOkfaEQKSIF((39& zsktigODU;ho?mF3O31n3c<2CQNV zepVYAIsj5;UFYe)Q_ne4$o5)s4{qHWP7=bVS=JfD*Yo-zW4^eA1b{I;9UW3*wa!~J z9UMq1TH5@zmA&I~v}-od-udf{%*@e3E_UwjCxCdFs)k{bl9B*toUe=2aZyemIK@t# zZHyKs=bmlPcZeSTxMWxjSjzQ@^0N;68w3Qx&YOw3gMi1Y+aEUvTwuGEoLp};rjC-B z_-JDB-(G-w5EmC0U^eG@n=M$VAu*a+U0LL5xX&xRQN5rA@3OnNj z`9wr;M~ev#^MPLC)H=M)C87Jv`Ea{am+ z00BEYJ0rQ8#@BOIn$1eBCd#Ct7cX5BYc$~&IheFxUS5WuVW%d_EYGD5NYFMdn&HC%Pn&l!O#|9IO_t+l;nW##x3T6rkZ}KOtb=|8v!(o9_ z07I05v=o6rJoUajG?wu2KZG8ak-w#$tAd^OrR@gu0R9j_7I4xZ&xwQw>pa}TB~x=c zDQ?{gi;O%#yJ3WPd!Ga3A|E8+u=EKb1lUHP3-T&&fmUGv8 z2%w>5B_##i4>tiD6Xu?)2M9GuuL`iFV4>|++$lEvaIHu5%;Mo~dfE7#oF+ksrQW3_ z6S}9_YDiB2zK_lGYn_0X@_zw5Tyb$R5~%UH#R6s$Q~;D!HZ9RLmD~G4w4x%syb*`2 zzl7;kHEenK(TG@h|73T#Nzu(}d+elV9?KbbAaN{sraJe_8R?`AdlHVzC0#t0L292*clikUJkD~Q~^`V95S0X0P4Ls zKX14?2_yhm8Dk@(@tI=qMzyF|5*{Z5o38)c5C(Jx(x0u$F)p};si}UgfQvs^JAwz$#tOn@pm|!ANv!yq`*eok zlje_q9*n1AbHI55m!Mb{7Z>L!fC3Qs^XJc8j~;ay?SPO1ZM5K3eVitqytTK-#>gn= zbY9DOWod|+0|*&v(Ds?GEFw0Rin!8Y`7`1psI36*m)2vq4^-LgKwV=J#LbY9kZ^)O zk;cY$XC3mK!=5aL5lZtWsKkFIMNUzk;6n_ zM3Rf0ov(<2eg^2|m+uX(?txNC8I>9E%eZZ*QBOQI1qCz2x^{2)U7lbyaTi$nEcHUQWg~`6v`L~c<3phy+yer ztFyDSEJ9^D=N|vO7`8uQeLgKv{Xrn|^Yal(Wh=RWH!i}N2h#{=xj4->2V!wp3^}_N z@SRWSzEv5t-2_e?6pKKulsUT-s0C_M>O=iiiCRodjOT{))9$9tyyP5}jy}tPN8PlA zfy=c0Y*|yY0+?^_YIgOj>m;SdU9{=M94IGhPX8g#Rkadt#}Qy|Eu{G&V95&*yN$V_ zd{!wbGy{bB2m?b5*(m(BU7oLHunT-9fX~gZ-5&>Tyo1-Qbo{*R2}*MFL@#A!K@mW^ z+NspgSnPqI;6xd!$WV&-rHo8jsaB2i7VX(#3vdY*Y#?(ZA|QwZ)nu9U7*RNGz9UjB z1q2C*9Y`~;Nm&kNgdnH{95O7Dhs<(k*3P3QaP%Pm03nN+tv##!JQNibL-aJc%LSf2 zdj_2YEL^ywTz73fp;565C@WzPFskTjLBKVEApTDXu=m{KKlb#bYmbV*TDe?4FuJs( z!sgW?rNGJE7tl*xML| zD@@@V;lFpda$}u)ZYx+2<>=_>MHaVFZGPg0`!1g+pz9IZiqoGW46G_YrrWDL>&kl` zzeQlNys%7uYuIsaFkGYOF7c7$(I9xt1=Zwk_mH|k&_5xQAB#)5n0^Nuimd#tF)}s{ zL@SZ#i5#!5>rA8H;o%{ufb#M25!*UA07591YF^pc*q9D#F8Dx1L`VJGkOgnR3s9pK zgSPoG0~c3LE$F1iN~Rf5G@6^9z6dfK4v6F@ZUF%eXqokt)v9Ts8X$fFkY|Aa$qp|& z0KF^_a|EvK96Mym+JJHfWP8jwfvMER%`Gypb75g259FL*G6bwuk)Eee0}Q&hECRrQ z2uy+Xl!&I()2H8o4Mjynr1gQv<|c?bt|(0iX4Vw`Rv?>4Vh->WmqBeK`{ia?+Ty5> zwGy>yeUD>hC(HG$FmdKOJI+acI>I*Q{tdmt$Z2^2!E1-u$!Oz99@c&QnAu z*LF$d#E`f!>SJn?i%PZ_luFwn?y?9IWDBb@UMFM@UqE4ee=X(XHGK45UyL7s8A zISs0_Ks2Mp#=(eJ(cCYWNp9a}giN{bNP~)(iHV7gjSXlX(9)U&)Ni%psuKV}AUHD# zgS_I>#fv_`Hvz97EOOjxz$FHSuAyt)K>jKaC*wDyqd!{k2DC4Lr1}o09u!o@3f?q; zwD7Q>G}`g0QkAeVYwfK+sLS>Q4mo;0{->6Z=O!=F-Y1gL7K=z5oove|zk^wys!?CX`QfWF`C4e4 zyf1tx1*M|WhNh%cygrCNBMzD3s|B9vt^`9a;?P&zoJzYhwW-Z_#J@Ep_skUK@v(UL z`T7qFPUTHD;}H-*Ss-JPvhzAJw?aHtr|Mwss-91C^Ov1Q4nI)fF9sMBCirRU?!5dR zQ3ka<0((G*!8x@L0?#-QcciXoCO1G^=iuhj}Ad0E!`u81(F*U02y?X=#b& zGnGKssIGRtYI0xA$#w^TMNwc7P-}gx-+V(?cMAN->^!wf8dwIY^x5u!EK8h3AceAq zh5#IZs-4AvR<8?}Xh4bpA|{GAdL|kLaAsp;qdUm^z!a%OJ>5Y@3$p3(c{}h!9*|>$ zY@7bUgBw(jH>1ZbEIJyis;dEPK1q$#(J0jWpnwLN3e?Gbb$M;AMz6^q415Zpk`h|> zCub20cG0XxtU$$>pS77!Fvg?AQ#CY<7;qP+5~BES(y zNl8)inC&hmz-6KyzUfVTdS0~{BLsjMNO5O>pYYnXs)bQd1Sl!tF4ds^=buEVKV|gj z=;-^vp3&-SzviC!=FaxvO(gx9cF$h9b$UA8`tIiS>VwRAh=Nv>X4#J4r77B}(!;N+ zp7u5lN3JrmseMgduwBt%%IiCI*pYuy=ciR+nMzrt_HMA-k8>@1eA|O}CipmvQmNr{ zrL_zFGuBfpy6Tn<^zildsgvjiX}Ts+&&QHuR>Q}6bwr(IYMqoKF5DdV)lG9;`zdPv z5Sm%pb$KG@v!EOSRF_$R+j4w2n#|f7lhv}jwN+W@oKakC@9gYsm~X}%z5QiyH3>YT zU1L3v8Vkfx2H>C520%@uAq52m=2b4v&Hxdn)e~DCTK+aUFN&&94tKGVb`v^82Q)`3 zD=_>!Ts;7%?qT6X#Tu-9LKZr%S%_PAit6CjhXhnQQLv^7q>EX7xAv4Y$5~d4YAE`k zmLB2;cb$lc%RNf@9Zbw*KpE8Lkp%$W_3y6&ZVzA=n4aDY=%f%m-Blz84)uMl)bDzd z7t)E{wg)&6PO=65$|~5MZLkW?aI%1m==W2Z4x76~7+!ZDNk?0@ zlJ|~Oq zd1??+0kiM}fPDaUQ%XRU6?C*6A03TVbN^mJMxL-}?HGJ*f>>Lk5u+JW0GF;73Er)P zt87o_$!G){5!Vh;_W8p?eb@`GUH5Hl zc6XrLlar?Z?Yiac8;vR7dwF@4gM6t!RSHBEz->rJbNe9>y=JBRdFNg8=Lhbcat=hV zG)n9No2jwU_{*t8geOr=rRi4khR)ls>zCidDpjGZWcKwq9#SLYT*Kxg>*T$ajfnZ4 z8=h;593D0X!JH61-aD|oFWOzV0n|MLasCpmOHaNn3$i&Eno@DseAbvmtP#cXWOt^=TgI4-a7o=l;yO(bS_d&NZ`5*Q_@k0NR z4;nkTy}h2=^dWnS+925fmk96%!@O(%$OE)B8gfQP+#(ON<|CUMO{V(&ol>tq!~q^& zICCri8LHYFv_Sv*)N4EiWYfRD`RC zRbi^+<fHucg22XU*nwE;iLgUO1dn~I?lqU%S&sqO6l4*2%NDp6T%jq@z#B}9+$ zVcxjg@dG3W0UiMn*!>FI9yrBLy)km!30r!drkA&Opld8o?B{k_X^HDhSHOnZPxlZ- zt8}+e=ze`kBljNMfEw#nN`;JAqF-Ky)M;{K|&f%(;F#jlc40lG$@P?F= zqrw@sExWoW@S~@dK}nqXYRk@s6pf1df&Q@~I)c{s<}CG7Viz`YzB6S@wU$Y;A+;n3` zW;ht+t8|?R<&~!+=Ex|@X{C?WO(Y$%n_-SuhTJNZRi2o+ZglG#Kf3l&;cf?9aG0w% z+if*%VpgHs@HWlLNVB}+_i=P#dcU5YY`P)EIL`t6FsN*UUB+<-eZE6aY7NSwkYhBo z=+Ti5)KTQ~HV9XRlp0ijs^U?kr(kr}fkQPzRjTjsQI#um7CayJaCye()DKc*va^*T zcZRl%_BnR*sGgPO;9LAk%J#&(w7q$X6j?oHx z@Uzdi3-_^=(Mk(PMAMM=o0yCgD1^>Oj?Sf&lcH5!b=XX<$-r?rk?C4C3;YVv!G*-P z;Buivo#P)Z-qkW%`}X$Tj1T=0t>4RW!{!@Vmp;(Z6CWkyAK(r}4+z&rPpVvTA?}5k zU1iA3;zEsq`8$M#+s^kqwo5FmKL$VKFa?w`PbiXE#<{|(H*#YlM;{iPwhn8O zKpwfF;KTs0neSW--PN%g!no8PxFV?pN?)Vi7`nA5% z#neqYDXh-Psp)By93ARQD2lk6nNZh&yKSBy;5MYRW7hn3uEXW{B8OWCB;*9`g|g@& z`mJXR`-f5xF4*4pNX)8`y#*~Xr+OzU9y>EQ|IBbyH3H*e?4F;FS=M-S;cxNK^=aG* z?U~7RJ~a?+-Fy2U_cD)$&XD=pVxdWVFCXu{hj@ef!IO7F$}HN`)RUoe;qz|s$nu8N zdD&AT^Du;aIm@)G8~gY_ED&UOrKsU@m@)?;NyxnB40DX@6W`twO#F_$z17Nk^cl+D zT9h}@U!eGgSSW7C{5pJnjDtQwQ^2KrbO}s<%0g7f?N@*0Bz?Cbxd*?c+_8V=PV?Ta z5FkH|y_tfK*75OJ8zN#qe)k1eH=U7MbhbYB0jsMPH0_3I3^>7o6H$KBU|fqLa?;7p z0z*a~Df?2J%SFkdHF83RM%bfL1>DE-{P3E1x<|X=wyGuLxee`4Cn=R7OgKGGFAk+Y zDt^fz>hf%LZnMuP*(g9A%8T<{8Jw@LBszG!Gap8hAVd!-R;9_XFr$TyfYo8KKS007 zOlV8c$zmJx%G|UF9yjB;x$7674fi}>olLj1QM6Vib~$>h1B}sQ;xe|9KVR4eEI>pw zYhs^|b`71CiMX+J8v=Q0s$QzFxLTtx`@uNLgw_#K*I@pawY|i@9xalB;F44%< zOG!`z>-r;&JHH~|4=06FU#zvLld-<8c49F0JuVkIS6b+CT8pB}5Oc=lj!YD9Z0v?CkaOAn{uS8G zmtZ#qPjqEzVdQ{%sTd!aK#4(F0jd1gQAM0=kB9a8+k|~OvRg4gv9wPgw-%;3Zrjk} zQQVp-#=8sAD~qCaU-ezzirnzO7A;g#T%A1-CZ=g_zF7lFE}H&fruWPxJArFsc*>W#B>K}_=KCXao z&SDtYDa;HFXf`fA@jk!&2?~|b?i{0mF9V^LK5Ai^hGx^;%&I&Dmu&5A!=$w@N+c&+A6q(+5L?*EM<^TZ}L$j z$D}5H`cKhRtX$sMyXZa%P7fc4A9rVv0FC7}pHo-w*}622++TldEIMISsv)cur(z*k zC$hf=|40{%9{Ot4efj6)@si}tUEi-B(53Zg)FppGwj>e28P|R4Z^Fpa+k;m6x#W1M zHEWK7MJ&DXVT9Q^l;;8G>6<-uyz)yDfeX=Oxr$9l8aZ`YYglB0(MJ>7nbVnmZF?1J z=dC&!81!;%2YWjQd!`;5A6q&sz%R7dHcWd!k3fji0eq4Ph0;MF6s3}YD^ut97}S=3 z3VwNeZ=vddkIJRMc;|yGQSR8wbAO9XWbv z&vjg6EtW7iX=Mh_-N?x8b648IvlsWlRnRwG@lD@lB)pd!JL!doBRvMtN)x|Zcwa-qQy2j~vs@0$f)*OHOeN6pH3=3=aes7k6#i{66 zj2?H4?0kgy;i^`2T~Jq6O1n*7MV5V{3|-{e;;Zi=26cVyZ;6QIjf~Q&tE*+wu^7&y zB!}6wVCoR+pCUO?W%I~aRmv09{Th62H_X0E%bndnOjSYMFvsskudP$dpETU-EaRE= zXWJZcQ;aX>-QDGf{ffPRosw57a7$ZY zRveL#2VI%f8_7oj&O+#ru` z{OwYiKsI{i&wT;-*%g1z82l*+P)W`yUp!WqI!G`zAQ*q1fKDD_RFCkx-sFqL!?Phy z1Km2K^@SD*YDgcfApYamj=w>g%{A51q`m{!Iww})pXPKvsF1?eWm7e*lu!8+46ZR56&9SI}Ros(&W;Uthl||26QR z7l;1*{0B3{-svJ9wO`Gu8rWVu>dTAN)0=9wGCJHx-)Ym_DlG(d~nI4V8QzZ?Rx8`0!IikB| z{=!Q{eZO*5F4%NY-8lVqq_rdwzaW5mUfJW`$$9q9ojqqb)HLVUQ}ck4e!o54$3jMH z1^84nJCZXQd?QgJ*LyGevp2-WUBq*EKuF8&d~j`ND>*2TNu`pdsu4^wAZy@VU<)-R z37Fhkx-?>Kb_w~s-|pyd{=-v3!d|lcp6z$rlpDvZ9|RQ0zm8_fBY$=Uc!X;|GMsc*}g&&5*d zAXob?0Y=)NlKhW9UW9o6lf>UGpAY$yU70Ef21Y$@K>*1 zVqEm6up--qwZ&#s@5+@Bk-7^Sxx0HKM8b8p3l_5Es#u}}Tl1&-g3^?NZ)Q8kMi|l` zKU_En7O9;c8e^{PjBoxLkLF9=U#7g%RY*QXP%0*S;hP_opl6xhj}^w*#p@*WR z!O0P6Y{yQ%nzM+@F$r{#|E}EPOjZbspqp5#w-qW`VJ@AC171-K?1B>m?F`Ze01y%O-Tz{F&BsqB;hlQuEk+IYJtfcZhnc ze!+``-}%I%?b}dcb5DlXU0FWMPwA1De#L(ZJ$P-l^J z7(`)l=`p+ljJ@6Z7@Ehsg|;N0(}dA0#@0$w0|gR^zLDuM!x~*>(-CNobd*i;$JvxF zlCdu^Pk|9NeNT(QLr?9QjhKV!v6*N1^q29v-zD)qq5Qt`OEXjvhu;#}iqW`TG=56g zwa~B6@_SgA%)~N|4nzf5IWsGY>ihTw2HX@k0&ZNH`TLFW_|}a7wL2}fTrU3kF%gIx zJZXYHK~bD_dn0QXpZNN0k^EMiM33XOI^VKRS3&kfC%cuxoL8Y^?%XWNskR$n4Q;Dd z|3C`~>>Q&1Aqrk(kcqI>l6M)44>f(q*zEJZ)?-y0>cGs=GNFAiP;!T@ijEpj>m|-} zc~6<6j@}~>EI_f~m!f_neU@$8XU6@e4%eSn*evv8emNJh)4LfWoUAX&%t8;}EK9a` zNDEt1gShSn$d<*4yo8$TdN!!0yS;AAP)D-YZmpl?x-ai(UM4`rc)E6lJn!zlg~!hK zG2!<@fUk_oPf$s=h1rJiPT|atCg5z=rk@%=^)fW;_FS)*wnp|{<`?LCKJxL+ZTj*G zlg_m-PzP!Ea@a=vI0t(?qHgL~HYymCa={x3KlxHdY(&+u9vv~ z{%OAgpS{OSCG2FhCpH;FX1!+7eu?lWg*9NNrdmc>4>nK3QXQrovAzWG|` zc}B%wVlz8#Ps&!IuoHyyR)8eK5=4$pkuj7au07wr&bl4kOL}dc5WG=!7K{=crWQWI zg}F&)XBnMVmAr8aGSsQuib&}WVQruhh{5&ui`CqHL4@MBct9+Sn>XeWl7*Jzs=BqxiO|0H^lKR=e5$}7N=rWq{%QaF_=h4+ByPRKuGxNeLnU+3JBE#K}& zzgVtZJbxtL;s~ANeFv#Gm)gj{^+lQlrB5|je5E&U+cHKoykF0BV+P3(BNBaivOOoIEaFt}pZoc30 zOvt9ogf`AjNOM)38&?;_ShH6zt8rigQygUDHuu_r_QdS- zhT|qDwgFXmxoOZ4$YuS#zWSF0Yqk5OcgV=$+GYvJ{WXiIm#fV^PdT-%D7lr= zZO1%T=E8Sj`t?i3k5iA=@^tllC$syH=MqLE%C@d!HseE1t;3mkOev8}udL^&eui8x z`*qM#Hc~~Z|ICd?bIVHJvZh%N{l(LH=sWzP@QFCxZouw!M=QPy<5#g#oC?y1$)w1s zM-I?eFP8k)a$vZVA?2^04!fgQkY1uNj9QJyb&f%qwtmw=7$KEFlAE2aVbf=}7*(2a zYq+F9;PH@OdC?7Q|7Q$yV47>NYgI*eEH{5^M+_oOI?4eexbuN7fc`=1jrHNv7PgqO z1pc9iKfoRYzucIgUU4e>F_~i+kUR8wF%RkYCdzu`R6D;wS4mw&q%{8-JJt=efhor+ z@)i2Txnj!%*~mJsOOd^_bVrqNpqzKM);edv71X8=KKO`o8EIt%>%LTOsyfBiO_p>! z8N6iRNN8#PcztI_FfKAiRl%a?>Wzczcax}1TT;7tBnWHi{JDbc6`5E9)A%CJglZ)0 z3d&{AmZGZYnD9(Og4A|tZj%TW_TWXlLA(*ETJ_BQF?r7)9jf3q)$8UJTnD;bgs(uJ3**Of1c_8VtuUbdh2`Og!^;*r(eItHyC6*l!tyTtP8)OEIpdIZ<0?1?5Duf~9-iSimsC1v%4e38uBrWvc+TbFyBwGow7P;wCcIv`0P&`++T$!ze{;7z zqdPtbJ@4nC_mfEk9gS}LAQpr-M`xrd(Z@Tg6PKP9$3FH<6Bj3vd{GjV01cy{#*PP1 zotD91ig6lc=BwWhtxjiVa>E;a#dRr3Z`ZhT3r`zULD_sI@Pr@;R zM>cyhX(lp~t!Pn5P}{kuF!mW~1h@YiHOTgrkK2Ar8`S7P#61Y6TJ&+LV_P3m)WIqG zXLPXf``KuEV3z|^D<-wZn$MEOV}#AOzw`@P2bR6i-I}4L?Wx?@^7Tc{3AnXcPt`l^L{|s&T=xj@ksc1=gYI@QB0J$mgW~YCh#= z>Q$2a7cZ_ls@!VI+*xEIwS%;?)IA<4Yt~a~BejwYbO3jfDqPQ-zE4XnjGGF@jP6KZ z@bVGR>#K$-Z{lc^ONm^#h?(aoQ@{^VtKa z@pMsQ?TJU>2%NA}3i@8nU7@3)^jagK@Q66N$TpqYBPYtZV@nMC^&a2j$udl6bI*yl}j|+=1)1L?=HnX&%c+{qk1|vJJ#Yk zf#V=XahY5p7CSy9lzMd&y0vw*<5RPJ5=27IAryG^hD^{$t3zuW_ax!PBQ=%F#Hilv zd|~itfy>wySx7sRkm#Fz^XtPo50W;BLm2Gs^E_tf^GEU^N#>tFr}9b@T`wNPt$*ttt-UDlT=naq@)A=hl%Ff)^C6nAowRR}rrqppfmlh5;JU+9Rm z!EK}Isv)L+ZSCKu!E9S9&lbQ5ts_%17YORTUMo+!zFKk3eh^k>`2kj`S^bWz+7W7?bY z!QCNFMqKcfWE@TWnAnE{k1g9|0-3o4DfflG$r%X5b+ec#G+p$~eGRQAKMEQjThlx6?bF62)T4WRm7=96hLDcdMn=2QTZ04FManm8cz$fsv=IdGg2*ygbWHw?Z}&7p8o2mdlNM5V1-bg9wIF$Q znRyg^EvfbTGw`Z|U0CJ4SQeTD`cs*Le=6$HSH8%?BBHrRPOQANmw#H5Qcz@u zTM38WqNPCtgKKc#9>H*M9Pf`-3MG#fnf0(kuI+D|@)D$KZJ`u<3UBN*Wr?1NIDEUO zO^E^q6oDBv%i*LIj4Zc1V24mdjCvz#P$<782k(OkU0ki+?3R8r#m=e}IkEgRX9mi+ znZG&N7-9^gUS`(%Ev406@-s2?jw$%KIuwb(jaSyDRQV%P{wBO6T&op!4H8{Gb5Myq z6M#2RfcM87kbW~Q#1i_=eZQs7%E>O%7OAXkn!|0ygbC0+20CKV<&u4K;y>POAha#0ZRq=hGOzvoLD8g!m?-RlYCJl7VFZP&k-3)ce z|BzlKS1yVzf=K05y7ZMu!C6 zYhX^>W9_0KA1uXMRB^v`MEi^G8yQ)M+A||_!rnxE)xL=^nXg09o>iV4VUP=d+hUg$ zFNgZQxlPMAVl?3FIY|Jtl$pRMW-4L!lP8MP4z|cB9i@r`B?@8`r*p;nzz`kxM}&^XNp4|t4Q@2VdNrrad5VFSTcPA=Q(Hf< za%@0BNxAL%yN^wScfHQ8iD9b7la_=HHy=U-V&GQB$9I;~LJWtL!4pI|OM{DEM*ATanQE$6t{-L;cPI%$r_3Ckz?b|8DiAzE4 z^7hHAe4C}KNZ%82%Pi9@#f~zo4GJ=lnPX%H%Tl2rm>r^!F)K~G4vzPeko&_j**wEh zEW56`;-R;1)~)Zh#ewwWkws%Wxz5gTl@wp;8&JDiHkx;njVoy`RAYUh`hpm2xtzY9 z60R#gBB`AVvOVy|lk^+h5IimVbZ(f2?g~}Mu$G#JH~C*jA0|Y3eyO#rVqA~Oj->xH zkCQ5&_ZuQHXr=VJiK4By@+06HfSY{;wOHc5{0Uqxd9mZHv6dVou^~CB#_W_?*N<;~|Eaa(Ojs85e#tKE0Ni8c%x3QZPBW@x+P|nm zCK+I&vu|~wL~%b%?*R|F6}^fj_f7f0)cq&5eJ`L?EyR6wox6>jGWHmchm#`-Z*DI; z8A!_XSx|)H6NL|JW$^MFIk|l*=Ob?Nd4Gy+o*0fZ`%qt{_e@>cu=Y9YNmG*vQA)z{ z9Qi*p>-x!x>U3gs;iQ0C0?dX^NFQs%nxK7=PUqpLhai#)Sfh%MTEC(TZ9zogg}7B$ z72eRFJ_3*~M$-E`#k9k;%P|qFP<~j4$E4PEmTa696}Uf>KiaWWtWU79B@*G3AX&We zvZHN-bbyZeZmYa-?Ts8#(*Su{Xie%q^h=;CL~nFKmB6=yDB)KXP>Yahnpjul_*TP1 zG17To*%a)GgZx{zYTf?TP%97HHi++o`=M=!R+GFCGxgxtep_cwKFV8xiRH-Jg13dx7;sT z3Z#c9u-Go%^4f1|_IjPZ|Hv~!N#(=wVWqj2g+tc3QmD4G$lwSE=YtZr^zNFwtW)LL zecT!Xo@P805N69w^^mzAZVB9z>ZT8O;ieiLJ->J;uK`xn-of*LDa0OOOqi>*S9Ons zIe8po)lGKtzf2O)(e5KZ8aK$|MZ;t4>uPdzwQ4?^f1gaG#1cl8yNbEQWhUh=Bq_Yq zRa-0xS601aZiZ-;&b^nkLJ=C+U>^|R*Q%UZgvt!`B|?wFzEOMoh~Zx*?mLL#9Slge zx4oAHsIt}g;tzkH!yOCgzMbtoh(5_3KyXg%^nZ9uR~rxSkAKu~T_|L-ML7U)wd?r) zXo3KA|3h}+{5f~^FOE0xO#}feeXig({F5Rkw;LY%cBf5kwPB$S63DkbXq!-B{0J3g zX{-@WHCR}fH}&A%o_YUj&i`WVt;5=Ef`(tJJS_#9;?mLvhXTbNS_tm$?(S|?2wuE+ zaff2Tp)JKB6bbI`?s}G|PtW^(?{)q>_th(*AtZb6&d$!v&d!f`#8Tz{2VmJJnFtnf zRhLW&$K3yrs5xJ1=oRx|euj0=d*+SvUj-yDYEgy1O8tS81wI9+JpHLM@~gQ%?9je`HYynxf}UwDis zpE(VOi4`|&l`;?4`7bNyCO zzdIYBI!&4G((ALB?ffD=CIA6Ukk=#%vx)86U4zdJ%=AkR@R^Ngme`NYgNKD~1`(9v8QP1it}8Qi}Y}dry%z z-@2@$`KFIMvao0!_$SeD8xYL;+1#KPMS*EFo{LI>)+Ggd_dKL}6SZLb!i!JdW zCV2Sbi8$w&6j1+!CP-Ju#Kt7iD(^rT>m%0V6m*-41+d0%)Km+WRKwx^l&*0cDi2S2 zt0oucmn30XF)^@(<>CT2@$D~+{RG3*Sx*#V-b`EP#9p0aLk@2O5z|*n*L2?N@0#nV zB@!jlW?8~;V7ov=F>!K2>1tj@9gddb6h<5;3)uMdO2bO8Zh@LSXP(ES2O!32b+UaK zTsE1l-UZF{6ds;VBC<5r{C1ZgS&6LI>1cV;w2I!XY$rfAINoKEQVEV zDkp5p`c%0@(FD=2!KsTOeqO50yH0Bt z9{m{(x8C(D49Re!S}JReE-qY%UCPZ$E_eQsD6G;rKC*oFV{&cM=1;i?p2t3+eOehY zOUNZP+)Ec{3`n4|UyvyODPGGVM+=V_t}=%s83o80@He=t5KZ+V`=gk6)ElDg_*5?F znME5Sr|DD@w*}bZKF|#73S3wOoGr-Ime`dBghb|=n~S`3zr143Qo0LS#3gtVQCA?< zAgryNub-TA0Q_(+$flH_Nu3h62_?|SU=@(8xp;lH>Tj{Z=MxF%)7B}9i4tJefv>=1 zV?*2PpA$+``5Uyl$i;>c|7M>e-rxPC$Dw(=bZ^%B2e%a@1zsGu-O3X3X_-EhfS@d< zwD5V*kXFZ?*IxESK~pe_AZRliS$Q&U16ZM6q!$}hHgmPH{FoAeRa_1_Sjai5yH6f< z7|fD-mAXl!#aSbqG@H!hg?u`)_vV;rIa862D&)ao!vF0SsAW%@VPjv_20 z9*$+2gXz7z<=zA@V!C<_quys#GL^$}+~;om7V*Fa?KeJFip(~MLD6xeWQqbGKEyfz zo_)`!a|w2rpQ5omP~g?#I^6V*kBgOiJl(Mtc$$R9NB;#k}Gw($?7)b z@}#%bl=pD>T|>zvk4US-vZPe!LFk09)lsh?L;R6p>PRu!G=JsS3S;f`SAwt{O@a+P>1lr*>C?=4xZRU>zru z^KqWei-O#`m13*;23Jt8u7B*SlsF9}xwEU1nix%+gX8{PFD{>pgMLfShNGn)tT|qy zVmN)fh28rxcnDsTLT0lRrhkt231AMZ#hd@idc2ONCmjIgUUpJ9c~ ztZ{zCYj`^?ty%q7bQx1JqYJ!WJ=Xm^XOF{miCv~En2#U#y=Q0=XxNxE3y|~hxcC4& zm=BA=X7nY`7n3{utQ%KKkM{I!X{$BcaPC6T;(%T^6N5xxP|VT!I$znQdgk+jO!LL9 zr$$QW-8J)5KS~6qmCKDZh+a$y$c5Dju+y0$x0~2d)>r=QZrEgto5km=ffnffy$759 z?A^WFAma5wFWKQGbud^Nq1{YO;k zB8eoYa0cLLnVX6Zf7^*+LY`R-hgJXnCa))``*`Z&(ri2jWKUdAFgyS1s+tg}&^ngA zx=Iu9?MZnJJ+nT{XQpic-r!KvSyRaSi_CU&a8S10!0N-oH+z4B1Ao!AeY6=^T-u5Z z49sd@p78MPl-lj?RXbv>CDDG_peYQFlDg%DJ|2Q$y_DqA{a?GY{E=C5gxCD>qA(0H z-*mC%T(D=y+3OPkZ>+wa1%hRUTs*Y@Jo7gfU~*^68hWPR*&vMLwSbEkmb{;ky!<x=!vG#G9QOX>MZ&b_I#-63s_@p3S|LJbiZKNPzD zR99ISe`%SM=0$|jC?@)`UzuqPe=<*Oj09}`c-!i;KhQIKy_a?!AVhe-)h&@t|C?m- zh{N1SuoON-LG#TOHpx)@%O#+lb%dJ3%uI-+slrTz`S{%5c2#h;wd_s^*29d z>)ctnQZ}?P!p8}c88)qQX~0MMq`SQ}$HG-wSD+cCqwEfhi?M+`6)#ta!dVU9Iwm-! zdLjWk_B}Jul_*WrWCd>&&Et!SGrDXX`T5Fa_<_D9r@*hM>q3w#_{u%-iXpJqKQLbJB)yu3hP?G8C#CQi<L(zW4ntYov8S}i$5qOCr|&bFrM`Yetr-v_0R4uaLJE-4^k>XI1=AlT^ME> zY~nPwDoO+gM7Aff^B=2smu^r+3I$qYjIfFfl3~M6jgPo{CzW?HQ%lp{^fTr1oQ7t#9DY zSXQ}hPb~m0H%NG~p=V_sn4HwfDeyaf=lZ_4Aopr9_i_`XK+?ss=;9}yP77*LGb>nI zweH9}uTh9xtp;{tP6cvGdWFq;Cw4o~!$Kd}x6G&`yu?@DZAj~y7ierec=ICS>I&@d zE#Q$-pm@``v`O%6Zf)@RcyJXh;>XCE0io6>9Qovx0B;|pF{1In1BW&+0OSX7|xUJ9{U6=6A_43G!qfJ>d5dXc|)<*!ix0 z{nj$7FU921+QT7_qp+Y(28PiX-|VfB(&O!FGxg(4umiDmJ~xkl@2+`33Nil{2jW5RU*Qrvr$DUIl-yxsga^* zL(|2wT@_o+eanJ2IX%vz)S*jzTg0aAP|Rc8NCodHPt&>Ex1G8qyO;{KpRAMFH4lK(+_UH1X<##mlc7)%~yL)(ktO`$>?|-oWP}8wjh%9E9xjbls781iHW`$+1}4Tfb(AH$Rx$4j~c{Kj)LHi4t7Rd_lzgM0jIk_ zFTRJR8$N8xRw#MPnG1dT{MRfqPO8V?EVEYZYZ8J65I6{M**w|~!~rHILJP!qtCYRN zupf;3EFFx_{j5D#;t|;pNc44JU(oZJOn2PKyhhq>ai87^p@`VSKnVt=&Z$eUl4Y!7ET4r=7gCIH&72 zAe_;#L39gk_x(4{{Yb-6piZXPy_K6|Z~Rv7l8E?~9)w_0A}E_nqx-lD>?qq`J5Bo@ zUJmpMa+yy2=srkZH7h(1HI>TPBypAOr%!kKnxp1A>p|jVmc<6rH%13D(909Gi5Fm3 z+$08xkxHt2`&-F}uCmgr>};h09oD|SeLtzf?w~Cfa6X5foiKWBE&BK5hqhiZ$yil4 zjUL&6od<=Cd1IxFq(n*XLa)eA%YGwg;hSuloW3a~g z3X?VEFy%;>@2ylf+8E6lrWs?^z|ZNsH~21T9#d;8^|V}dQPXYfCFIm<-HzqV)ngg+ z?a_;d+;3AUHg@$Hs(HkSdIRM8VYcw-K%@V%->6bk?szoL=`yh@-PDOq$Bb=YU^0&& zoA-j*y6U*>)X`N=>$Es*f4HzBxwR%J$fVA-f~L&bX}>nAdE*y`P&09aB;)zkt_5iq z<9OKvC$n2)MF>Vu0#((^isoY4FPoVw40x&Qs8q^B|F%gY1jl&IUcJ;6O%YAC@uLzh zzq2wkXO-sva=^^!3W|$4hV)oI!-MFF^;|NtHnnEOb+?kTg7vCbhpM*f{dNj^tDr)X z?!U~i`5779E`N|U@z9kbSdIFN-qQ%|wX`A3&ThZ1Q*6cv4GYMM3B7(^RpCQ_aSS~D zE1&tN1TjG-Uz%KPs`fQK{_R0DbJY;NFOtRP)u~3@W|?>FKhfvq53;+`J4LnhtMB7i z+0kZi%&a&ZX7`S5=TK*h+`HGs#aLCHJ4IUNYJPS?0!%UxV)8IL?L{K* zp>2Fh=ls4H@L$p!@?bMo2C=QMM>bxC><8Oy!N_Sr!$**(NWS&irC8KRB$`Yf2N46G zp1Q&nYAt3x{HaBk^wW7|TO>v>{#f-28QNG{VPrZG9J{)sA|GSw@+GxhOyxZLM`_WtMT&j)4>@e6AA$s&46}xl)4?`)yD{@} z46a#9ubsk2N1w`3!_(vTSx(>cyqSeg)}jrw0Ee;;|W;oy8~X zc}v5be6LM_O0y<$V$c@-g4i2e@2ldzfE$3ZH9vo_rhDvrMf2 z)75Bb#g?C9RKT`m0K5~^zwg?sGAlrEXEDbnN-gnA)sfvzGNu|9Zo}-uRvzBJnVx_D z_TT5XUFZcj#o(L?wZE#Ocz=P!MJCyjB*m{z0EUN#7S?v}3CYaN7su;Ym2$a8MpQ0g z&epGKmy|NvxcLg3K_&&uR+qCYevcC6ry6Ico#RRQFD6Jg`)S_J^#R^`tvG{lF1~D| zJOB3K>x}5UA7#&r%FE5Cnp>IMq(#C)HR|ny|NV8vd%n?K=E?_mNXaTND@6M0GA?|N zYrA%*GlH7Kw!OqW7`W?SW5!|LFCJD6e-`vGwSebt0y+Dx686K|o>(Ifv z8}YkHY1uXzoHPzzA_{U{QH#0x$>KK7V6zu5hiOmgrfG3^Df3M$>G^(uf1`hrCSP0G zztd#=c1Jo@CTm7ntgTR6PoX^Lj`nSE!dQ(XCL?$Xw!q?@*=&UPR5$j#cdiPHm(Nsd z$MQp?S{$jPdw1BLW)P8gV!6^=Rl`49)&Ey--~8wM;)nm=Q#b$jmlQ+)i!9Ed70c<( z*Bls#l*q>IOVm7;Lyl*Snik5=@8qave)MAW899uPzzSJCTY21@g-wqxGCuIdMTIWl z7ae-rx8lxV-tn(MhW=jqch>U{w9mespL6O>nVBp+c-Mn4T+p5*F3WJgI{12?a({ph8JVe>mq$p`*_dq^jo$jR=r$mLns1gi{v2*EzFQ_M5E9_OKcI^+Gv;Tn zE!l>6q!|yMb0p+ZQ(0}1C|j4ySfxz1#q9oWhkBC|My)S&xtptBFZHI4J7wTD&s%ZA zSc>@M+1;K(YI>7j;;O0eU6-3LJ^e!Xx4GcFjl>V;XT-9t+Q?;Rzhf#5A-Kijk9_PQ z4wPuhV>ZlgS zIZ{YRx;MX{5A3}4(1oO(=iGccu{#$pw$N}nYt?;>lbWTNjzy!JxA$|$&)R}YJ5Zic z;F52CF+b6dpulCJUKiv?)06x=XZ~$tXx4(iN0uAIXO7Y5*sU@>_dPU8+Kn#Ql>#0! zMNc>V_ORedu~d8ri60_m^jkSEJ8xpQ6Lu3nf^=@SnQe`IC+C1A32e}yk)5$hj zo>xnAR?Ea{PhUV=RXVEkM&(O_vCBh_E<1j`LOES?khMFQrekG^W0_10{=#S(Mi#lR zR?Y2_MbfZ5W872sYu>w)J$)RZn0>i$ljZ@!zy|F}|7Loiz1{~>wiZa+8WXVKaI0(( zX{eBOb-o!pu&lF1LbDFBxuD;pF8|c&c73ubxFboeU50BYkDVmDY$z zG=!<8Wolld1iiOfv`g@ z>IOd9cDnp4)5K7UNnIl9*wq37&F!;;e_DQXK@}eb57g%63rs%8Rj%b@i6&%W9PL<+ z*5wN*!0M{I+pY)&XA;yAb*6T)?39cn!u;GSQ~ceLuE=WRDsF~UjQ$$wr^XVzIePU4 z9Gg>+{Jsa->|!KG;_N)S>>Lrhj(86{F$K2@M?Ho_7lM~xP-^p+QrF}$ptE!|YI&9N z5I|XqxF~GEAVp#anbd1%^#m=)uHg&jl1^26ZfW*GNAMtDDX*ryx{VPB>rYA@wXNnNt*r>iX7a=1^0> zG~AK17Ii3Opkg!D8)nRvy;szRLN<}@ z+yk}bt$N47rj17z=xf?ro4x*?yJpDo6Sg!)jp5g;0@jt|?`d$O%ddDegub1^d{NXS zxE~B_(2}9t6T1SeqRVXw#6*x1aa3zn zBsHIf969=!8ha5IC(q4{CJ7VO%q!N}@Hp?;i|SU%zRB=%W!E}EorfiN^7_dr#%(3P zCu(gBzjKR`;NRp$w6|EVr^uD3I3;6=f3$N!)5Qgk7i5LraFc9Ez{NJPSv^Ueqwd!F!@dJv7eIL5h^PU zcdt>w`m1b$KtUt72xwFaE~8AUO6EuUGcm0&OE_XwT#@;?zjSPV_z-j~&cAOU7%hYzq|k5sq9+OM}7Eb9Yd5_kXP{ z+p!dLQEcN-N`-!tO+`&irUvh@+wdur@Y?oT*)o=dEa3;c=7-X%E z2!Ib7;jqbtM4R?>reCbxclNVNAIzvdsV~bvRt;u}eNmRgkUtr(Aq{(^HvWYvACV-$ z6udVuF;(X|i^y(_K`T0%-6pBZH1Vv<4bI1e!76-_n?|M#IyohIFVYF{i?jviQ)VZP zl9QKE!I-;_$=p9`nqDUTQzG|gw!Ec}bTsBv?1&+s4>ZW63wX&48)Api71uU+MD4~+ zqEhW_6{JZ1%#>HQmOVJb;(AuKcRLO5&`oQjlxaL9@2+oBojc@EW_*DwR~;L1sk&sp zzMjN_D$1?4omsYE50B+`&<8k2F_^t!!*D zE#J7XTCZtT(+GN4n#mZL0P%&3<8Nh=#y^i8Y%0`bPY>>@DE;+#_xAccRq#R7aY*hL zgtDh{)s8OXYcVK9IN|($0rfF$%`DCnm_~KnQ8}$5UNFXS9?9>+zBFy>RL@?abaeqA zq720+5%R9g=%o{{uXu(;Ufif7MY*k_8Z-T#DKcws`de#~y-eOK`+A0V_~eobE*)ea zPV;=T8l1qpmP~x`LtuP5acPj0_sL@J-5l&;WZBmDKs7(jeJ~D1C0TjRUdi3`BhD{B zz)cynjwe38V3K>bxxsde@m=u0WlFv_>vPFoub!3B>IKEzk94yBEOuB&-?eoiSqRE9 z%W`uq87hSpjk*^eEXgs5Il;xJ_y;TIUaw}#oz8o}GHbbG;BXUT{+|+=d+1FJMCb25 zt4Kex*jlk6zZECu&qCTo)=U5F(6{A)5=J5QPu}?JiX9c!xVVsHD{ht7j!2iW+-#wE zMDh!Th4*|}GGP25z}Z{h5QYW4>w_tmA66tw;D?`nuUL8oiB~0Ke70Nc;fHXPXJe*V zh>d)%>xiYHdX37C?As=Z+l2bFg@=7UDzIA>f%)52;p*+DSS-GQD(a+CtXX|@$_fcl zItu)pP--v*W%|OAg$kv`IKG()3We@VZF36SYE2FA;mU3YMr=XZP$_iFXQYykhEXG|D5&FlZ27nunHt&g2NjdAGXfJU|i?5r+J*`?ZP*C)fH zh37QK`b8HJ4z9QoFTIelvPo)(oh1}OK?XWCXvf*_HLy$}5svDl5L!={NPHjpN4Z7g z-*k)9rzoOV5{o);6h;TPGXqYXOj^jNiah-=OCQb8LigTVc%4zoa3f6D^ZH2y^WqGY zdJWafg5e83b$Xtp?`^cu8c)ck5Lt(MF}7UMyC~=?&80l9fz!23NaKR51N#)j#!Lv`5TV zU0&4eB^2cZwwDPGRG$&1s+mrCwDK$}Lsb*$%t~-256(VaKZ#Jc$W5q|r5RXnnihra zF}{QS`WE8*BRik-x-5wnx2N9XQUN-5;llCcRgqx!wWU(a^%e9D+jcKLrl}zuWl@M3 zGHu6h!_P3Jo-e2_11lT(PG$5UU0z<+F$}}rY=Lw=gEHz9T$z&s5>an;Bezd&MIxVY zv`4`|?jLR+8*|4Rsh5dW{Da!RhA{ADD753;4is*HT@ro6JSSf$e;vs>omC6f+Csu^H|mF1VPFsyfy?Wal> zDIwYFX8M)QF!TSHv??Ecck|ZOIQ(hz3j?NnwSq15M&yuFO9Gb&J1!Pso$ojvZ5m9G z{sU-tg0ufpBSj#$Q0_cHjLfjjizZvfK-e z{N(urvMO_U6nvJ@8Q+9~*Xl{carkmcxW<^RTtIV3#X{$v@b!FgE%PG35MiZu#ZLd^ z$Vrmo)oXl6o*g*MEzl$c6hLR3A0%H|=uOt5^eu=(vJV^{yi1+lS)~@pjz4hZJkXy! z^kq@*dQjcNJ3`&4bkYI(q!aTO#?W}>tr5GWp5X)dKMqRtcEq2hVN4TjP)3T^1qw22 zsnRenSrZ=KA7Q_0nkxqyJdQs;XF5ic+e{XYgs;SCiJf~?gHx>29Iq7gl%jl#M@Apq z))xBf&uzHQ?pc*eGcg>}**shXxk}YM!~PE{)fgXsQh~t;7{~E5$LUHRK?3oAII_vP z7|F$n1uzp+mcE2Q%+_zSc17@c48Pta87}{k@~C-`Lshp>I8deXkgF?iAIkKKv@3Tp z{l|nk+kT_;B4M&6EJ@zrx&Nc4O@-+Rtv=TGa6-W%dE3YN(y+mP3d+K!SolX4pG$OU zf^{pyH}`NBUIc?m4afrS&qeVQUGwMXz}pk`{9*ua-{}%KjJ{QNq7PLJ+Z2Vh;-Rm7 z0?HZQ?)bc`R*XyOJJlyAhyZL**;u2346<93MX=AT%cszfIiM9FsP+aH5?qFYpe3|0 zR|@(IY}t-xc^>^VeH*YoncCiB}A&a=>(Rs~Ge6 z84~j+I8HZu@FiXn1oy^f)M$UHIy|z%j6XIDu78ek2!(1?>9ZvCoU<sM_8~z#-Qo`qCuV} zKRQcvq*f!cLx}iZh6*RU&rrqQH%IAoFq1%TA`9bq25JANL#MblbQTt*BVUz^grqtF z9wgA8n8fB|nI%X22UI@MC+Xuw{t!*ZHJR9E?o3W#MoQSgK4`Qv7<}MgWWEAJJa=z? zZ2oXg97pBUX>o|(+aN0h?xHdZ@2<`d6fb#)fvD)NJ)>H?Tcx+>$Yd)H7oGjwx z@9wpq#ar3iG_|9e-nyvaK;mtB-V`OKg_<$H9AyA+A$>xf;BJ^}+7|EWK*$+-NM{>J zS@#i6n3V*UMkAf?Bi*pSn&qm9GJ zT5`h%Z4e5ZdH6eMmKaU*zjN9%bXf^sgdS;xEyfD&l0)*Sorn2YCPckNUnm^}Q>E{D zHO6KKd0rjS)NttL8{x8uW;{>#R7)Q=k?kU2;3L>-Zs#slW>O~X;gzK?k#hC1q=3Y8 zB$vQtqawgtkII zheU~-d{p|iUs*RA=0uKM zKm?<3GVU4sc9$Yn7$vl{UemLE@!cC2Z03tHJ)h&}8j6GvR3YEUPCSE{<3e6?&6n?E zHn9e`U4W)Z{PD?y6R!ruBrqTu1LAh)fhN`PV7%6TuSSc#%or8m@h*P3YjIX+seBs} zNU1__zET3dqHNd2<`!el`+iA5?5f(tMjo}LpW!>FWs;!kPIkDzYZ`mgOSn3jnxV!z z9&TZdzXENhoNeZhco(Hghq<2rF&*zo^dZRylp(7cr1o>~_ z;>afvr{att#MA!dy%jk|y7?;nt9^qct)~_UW-WC3w%WXWoS99I^|c zwUB_|JLeRIwb z(`ZIxCDVHGkdH@>hnv5l{Uh@EkMH*q8Z^r0*0P}(_wc;WJQ#lt8|~g%x}&l8>obN( zhbf`oSe}Pt=3c^GjZCuAM+95Kuro4*7D>C8$m0loCxbh0yJn&UIR&qqgVd-nQw4U3 zp=XvoL(~Kp0++QvN+x*x;rMn_{D%W3*r6Nz$CtOj*>7u$RVECm;bp}$L21hLy*vyC zvtV_%$~HM{9AP{;Mh>zbqd~`T6Ko|X9Q$!%-M`1g)mX>bXGPidC=3WESIz= zWf*2MFEV8>pwXoeA@QY}-(+AI4yW!-mkMHTVjHZP|D2{&FIs{Fk;#>SZAcAVh&kJ^ z&8ix08|!#~L^zm-`r2Cr@U<+C&Ai>A4Ci%QAm(We7LESNvYeu%;^6P1irBcAEgy#%_X;uG}jG19ONKbbHr_By_lp?N)S4b_z0-as4g$RoIpv?rcok(I@YcK^~lY)>BQiX;JU!{>+*uz2>=zqQ~7m z91p%ae<~WZmy=yB0*h8M8l1mQ$-P)5TigCxPy{}Qr;2RQoX!7PQAOX6FjO53Msn5U za{|V1U981$phaw@#O6g7mYZxN{6r`5>kO9rmsVW^_%FnjnxQ+Y54OS@Bg=0adeFHz0H>S>A*zRcH~PkylDv@JM$`5@&m zd$Z!YLJ8)~knY?=E&@64XFF%|s=sqlTZJ33SHhV}{#c&i(b#|=J!-*iZTk^lkmBm9 z(^}2K)xJ~rrod&vs64v&C0X!wpCxR;FQ$y3!7%ptPpmhNS}vw{1LXrRF z4h^i0R#eYc=6W>wq%u@E#?fpyN-uSvXOP{X?M%8-qx7W?E~Kihbn)$&(-6A{1k&R- z?5E9(9V|NUwy)wlOc@0C+Gg$(|Ngt++Ibah1@rsc298w=nz$o0U`bwKw2QJ3+SI1Q z5H*F12_Y~vt)w%8^}&rgIHh%kOx`*8-(u_Qe>r7};}Fga(@TVSE@!`l^i_sHBFcJ) zxH(P0zzK(|j#bMKhkY=+@=!7G@l8|Id?DJ9W|dBO7?f4NP;+le?J~d(AF|w9=+t_v zz5s`0ri2{$L7YiGowTQ<<64l9+5 zA_N<>DtJPVpDyxq`8ZBV?Xb%z2Bb71(?pN9)VJ4h?X#q0WHu*i%1g@8PmRV=mPvfp$L0_vDQF6DsnRW;b_z<>GBzhz7BJjZx6Vnn6%TWG`cgjb zHa)8#=l2pnukjGQ&L+We+goS9#C;qqoX!ol38GV3%g&OQ#{MZkCNZ`YCoyMp4kHk8 zuTN-i>B8MpS}u7c<{aJ>_Vswp{-~3j($=@)wXv- zJD8$&L_btvyOTDznWDFp>RRtC)7;1Vi;pnE(20|Mmwo7T=~xpwjhF7C`5RVNef6-Nemxe9&M+i@Iz0k*UkjxpI$pux}V3WaPb8h zV!g`maSRw(bwrPs^{LbUG@vqT@6YzooNX@tOagI>rh0Yr>>sUyw#rty5;^$J)?^0r z%0XbnlubUnXORFm)yievI$uZGJv)@)t8Gf{FSTd(K4%;hK(~F5Y02PnTU6OHch-}; zj+iXKayB7VSeI`<bN@7F66fUyd_&8T_fCGKyYs;fd(4vNEjANwnFD?QG5uUb=j) z6W_Y9JX}{fcNBfiySY(bSGM0*-jpX8Zo2a_xzlzRw|Cjd~;m!4|5`}Q?Ej?c3AJ85Y@~0eVCelO{>|Ei<;MT zsHO{7OR}}wEQHfA%3)rQmFn*5{-k_=Oa%IPhdO1>F{@j*QjeCatpuU>(Y!L9YqBLA ziE*_t^_ICojy|=h=37g+IWdt)hC}8?;buvbp7#>1@UO20U#j}fCw0-}<-AYa(f{=r zbD0>AQK$!_A;gUQzeEzEl&8{ zlv9OuD-#%l=sX=ozE4?r7*!nb7;iL!V&wnPILD>Zk(bxUG4PRL)3IUGp`(MJbyq5a z9;o0*TEP>Z|H?F@eb4%lp`OF2fE|2wOj%}=idd24M6^*k0 zHuC>?sCFN@sarM3>{CaTjqZ8MgMsGK9jS+_|L_&Z9#9)^nJhB!WSpP24T`WA&BN6= znpFf&u3s$geq_@|Wg-WDi?GKWxp;5Z3*MLvU$5>?I3+idlhw%@oI=p@m8Rvyb-c1t zpnCDe3r#efgNL`-?IH3v=e5>!zw1o;B;S6^;|k^Ee9p!7oZ{DTAlmoQm2vjIs)F;* z!SXs-%I8`*RV_T&XY;1XZ}*p?dcVaLN-Z?AqiMf)7vgxSwx<_;UfF;grbVmFK6Lpn zMD+jO-}e^b)BjzrtN3|BCKkU>vT$hT%lPG+rG(dG&lm>#2$qun`yzmCuiu5i{;W7} zoo?Rd?Op%BSDb&)v>n}rZohW9c~HR*&Z@5d(!z$88$?ZkmINUKzA6tNAZq3V$f3}B{-AVHZbuj_fz3kE zV*zj!v>Da3wY4uPjjF~0bt*A2(TXEc3t{b{c^kB;qT--J3BcgM8o9)`U_jfz-;}y139tc|X!p_ILZ`W=f@+Il5O^)_fE*gM%qcL)WJD3SMvByM1fkgyO z1%(P&@5=E8K&>>uvjk9v%Bp2RUIe%;yoglaGg>~riE%+drY}I%n)W1kb}j%!Iu5mL zHxgOQHxB@8&OFIuNdR30kR5Jg8swrO?tl=sPD?o$!`&8l24woM_l4rdIw{&r|1X>; zX!*_H0l89Lf`j%}3SUkfbT0NVn*7t0&uI3d1Q)Ysd$g&*7;pk5!zXPyi9;>rL`D6N zzk=N3HBh0`LRVK8pipa7us=003QX07O$@6!w&t z3JTB{x!#80;?Me%X7fLzQvz-qp_C(=mzM|FoPbhj*^5l(MSXZW?&jWQ=H%!I%LKh} z7~Kk3UkOoBP5@>G*ps509)KSR2rB}Bztp){4m!@dtT!_=Gc=LIllORYoEi#nce_W`lUo3(VLe&P4gf~8_R3+ajx`gu zRhrIqKIb5_9=B!D(Sz?9q;%98G;05Ydi0Gy}2z5U|iVs1`}BoAPDIKjvH z+Ww6M{cR6&X>ZhrgO*kT1|XrV0GkB?N18`RM?c)hEYmEy0hfVWkeC0=^BSgEX_=K$ z3Hq8q&zqT-SIY(ufj@;eaBdJu&{jTWRrXu^a#ua3j#4 zxzD^bO3%pXF?I6vtQ}MCG6p0){Yp)d8^Y*htqV>BfWCqc0q|=`>G&q#9Kw8kucoiB zj*m9SO8`mihU*2eih_QPvN=k)A~5`y>iQgs+kl{Ua^k+N4Mto>QqprRFNc`jlFba6 z1lR+-&!6kJe}1v1i+~!wBBWb3Wg`}WO_mQCmcro!s~3Qa2mqZGnl^ykcOzw=KY#YD z^ZVPu|2=bM-HPr-CUH9f+}+x)2Aw%Tc)JYi#6?CLR*i>>&I8&_nyrMEu5PhvkL&tS zk8Kq|-h~k|3-MC|7^tA2AK>x;01RFP;LMqsm}F&T0bCW7gJX!qW2jIg#>xzpKS9K5 znk!NYxRf=k*@;(&C4QKsT)%sJx0B+Ol$A487xNteJ}@OAq1Jx(=V-YZr`6EL?XCZI z9nsj(i2LAhsAQ4n8ZUYY$E^qFUD)F-d02Il~VAR0B+*CIhMR11Z2G1)CpjwsHML=CM3+I=>%X(0J|x=!A1bw)^GEC z9nJq%&$nIQhi#_W%@2<^66DiHL@4A*UKMPzLE@7xx3>1~tbMT?Zzw^TZ(YondUnnx zx^AzrpYdL~A5>hW8C-4!lDKyJ&P5ce!Hn(AQz6VBL^6M=SURaCv2HXwklX#U?$8=X z`fPcUNX{qeLh&u5zU2`*Z`U!ZFYB)qki|Kl7FM|)&4&x(M+5*CfgSbmV*jtda%f-D3otTpE*Y&fI!8@#HfIdajnB9<$4aatr{mO zr#dXusdXL&+7>`OEfodE^HmC>qj4K_OGJmv*?`Vfue3z9hLT`Bd<};r@+?B+}_~-zRTV= z&i2%LaZJPI!w3Kkoq*d0#97=3z}scBo6?3A*RIi3UME-;YkbpxV|sIC!*ff5#3dv; zKEDtKyu?xKt>F?qef_T}%bu>T0-O8?4<1-n5QhQ4TTCKxXf|Lxf&~EZZ=!D!UT?1I zdXu>2Rzh;#y9m`Tvl5OfmGf~9Nk!~p_9}u1ay&66U)wwrA~!R>O`BHY43UML^d5G; z%ec8|U)n1f{*=+PO8o|p*4Yxb7&9+)KEIiUTG<<87hK;@t6^PcHXP$Yc=U9`$RKfX zak%E2I)E(##JvRR>;dG%M~|ezB`0UspBV_RusNSr8AjNMQ1J8n06q^;Qv<8cme+uJ zRHV9|DG>#@viV#0z=W92&dC|(&~j1IVwbCA8!|0ddG=LxD1U#V}F6-Xqa8wlzs<~Y7io?Su5l9f%2U@s5T5$u#{UDWe_noTh*Wsp!hZ)<&YJWPBX`;_pB%K(;D>&CoPiX{urK&+8N%*5ld$*-sqCP^LD z3WRD6CE@%PvCDdKe_`m*=Dgwd!2mJLSZKftdpx|Czf_q#0o~|&jN@C3YK520xnO7F zvI1Xa&~;VrI5A0vARpgww(_Hgn=UH)>mF@!n4ue~Gt&pJPk(5%;iH>}l4sai&5@NO z?7Q~TU4m%V0tp1z|8(=)PR71Hji?SbK&ATdW2PusB@eGcHV3Gf!hESXMliXHMy+Kf zB_BR~sK{$%WJFF*4%Rq8N#^jqaNnK*QvgErO8>1UFFjPXppt5pHb}9i22dy=fnj0z z5jRNE(NZHZN5`_W{I>biUcW1HsDgq5bPre#HzZU5MqP7{(R>T|rdsp7f6W8TeyVk( zm`Hj^PGcZ__hZF{my9kk>=@>xHvjN$y|iAMR)=xKUNAzhi07AWbO$V)gJ_sr+ZRL` z>==6zlJ?W!6OV_Qa}VnTHh3VkE?}?kvIKK$gm#54S+qdZ)IGI!FABvz zgF&{FI2#7sie!XE+>B!_1RD9f+M?;(GdYu{;uQDUzfAjF z>Gk!y`FKP`9-5nlbQBa6EG%WU<$&RpGgR)nJvBL=*JthtG_T9LDg*$PLaPMX*eV=5 ztQ{P70h$N(ez<%Yq5tD{-DL5zRDgyCBx3*#1v>o+4J}?E6e&i&4Cm74Pge_|;xGw?&1;a`>ZJ`P8>s zsbu+$Zx1BkB6qOYA2gPn?w-{m2c1rU?Bo35{#u>}E3()mm8h2*0!$u&;FZAn00x>Txwv@yk4udPJeROBIk-5d)4A=NWSFs52NJWTDOg<0 zi6p9ZTF#2b89$HQuhRPbY^7aq2OybPO%s%re@BJRYl=E2xI9;`3WnMk6%}UVW`E7lDeb;BhO}ajZs7kPg=+vo( z?o4@Yd;IOV?~68SqmB64Cs^-_67R9$M_^Kp)8p1@?Xn$)uN&G@Lq`|UQ4wC&=TR~} zsa6~UNI^+*&16w}=>48+RYdS)Rt$OFuSNbTEkskGpLHQ~zN{lwf4s;70)AEX@fZSW z5M_cI`2WUuw5~j{_O3^X5sDSAWD{FQ4}%r-p{57_2-f$Q4M+M^)T(pS#a?jRJXrhk zCZ_8>_5ou!8PQPC)N~E*HN|_c0|13bM$`{c0JxY}a{t~vk8rtxLXC2CGVkA}Ef;;)RUrmT5gbbsy-{`B36YTRag|4;AaX zQ9lU@j8>`kNw^TdWWVHv!hto|uQV>cIR$J|Td~G@QazKQ-#P!m2^iJlv`z~(k2EWl zGWwEK<=0Jati=HjY3Cyx(B}L^g~>qcS%iv~&}g$RbHww6fQi$Oe~`($lN_n<6wJ&B$=GI5pT_oxEJ-(}>=IwET5 z-ygh7C5ykd>ezO9^6icEU!g3nU?=?RjJ(#NbBrad@Pv#S@K^R(TXv8^=W-6q>o&o} z*c^n5LPUy^Z3up+M<6q8$=L;C=6Vf#%+H*MP;Y9Zs$=Q z`>%DoJh_NWCjN6L)!I7#b$^W)G)4#XF6;oJ_aeqRyBw8fM!QD=iD1>@YkN@Hmc~W; z|FHJnQB8JD`zY!oRzMU*sTx!S1O(|_NdTpZ^sXS%d#^ze2_PV%AiYX2(n}~JQj{*8 z&^v?>dPzw5?I`-Z@Av!8S?jEI?mr-e<*D0_C{S@lyUca)O@pJn-Y3O$oasQzPbKBe`hwXFgKCo zrC~fW!{De5{J5w*<*6fz90vvJQjW0qKaQ~r7PC`!h!nM?-Wq)=x^)I?{_{AnJF8t* zdzM}v?C5=}it3dV50YgM*=R|O4Vo8ACnr;_`xO2s7vNXZ?38|!K7Pxig=2s$nvaZ$ z!ByEu{^fkQPIz=T<6wj0+ZXa7wx$&lgJ~2(E4Q6)h5qoH-iI(U%|D2my?vy?>0V-) zdugw}3w9+?aUW7jN4SUY7E#t^iOEsml%}YU4R#%D)Sh4CZqNFyxO`i~#s>aw^T-$C zvS)N)5X~$8qONzfQT@cc@Ve#2eLK$PfjG#+D$bmQmxe`ubZ-=M9$-BF$Hm3u@|PUq zbE9mG#tKDi+J6Pq+3(AY>|gST&+Svf`iZ+4IKb-tz9{_q|830Sdt|g!K)^~@%NJ0G zoY3R07Rjq!;~*c9Y~me>d-+AcWV7de{-FL|HsBThFIjYV3Q%!s0f6dKQ`cn9XG=V8 z@lLE*F|3x-=RgfPNWJ^m8rDdV(ht|p8meP^^_@Bic?2ae@U&Zn?&(#!M zsJO}$W(xQT*-!?1AqG&u5R{;%Jb5xRqY}_Q`zq;YUTpj~@X2 z5{Sq(tRs?mIZ+oYKp1()+lGj34;u6E@Gymun8q`gZo7f;0De-YN~%{p?ceeL_2SEH zYI0=fa}yI2keRyFYrfQ*Wzh*}l>wmh7JD)dRvV~5x$4hV;+XZ>R@pa~C(+L=vjKwz z6zI|#aJPI~Im@eo#RT7LH9UIssMjAX{_~VCCij6=+|dL4I%t|~7}Ez)P>7cm0?LbS zxruQ_wNBb);(M#~xWbyJ8J`yqZ3A)!XkNezvIRNlJcFqKzNHqtLBQ|iP-mxhOt~Oy z#4iQk;N;{K6~(_{tn=jQH2E&g)vKRcuQvsK;h+E?m< zn!2~x*+q;8fqt*97TW^sr|axYpJiujDfguX4TK?S1r)j9q|6+Z89ngEG z&z@Bu5b-r0TkPE2iJ;yn$w$xx#wy_1P+6+1*j-`AIj`0F6HE5MAqKUq-8w848w(tN_w+xuJeFC8l+~%%SeReI^_Xh1GD0 z`N~KoFklIb8HryT8aO&v*4zvLaSS{G)n}muk%tbul+(@ z{H!at9rpna1XxZV`OJTUBF@et@l~e6m;Zk~GhcplK zRGARK0fV0wqg7j=F#ApAgxe5C5vNBP9NgT%lfi?+?1_LX1$_ec!F;^y7)fmE_U7pG zG?1i0-k9539&!XWCS`yz@NEFU4XD7{y}4L4hV|ERF8>bwBW^m2z1;85aA3)1C`1=| zY=N@5nmgF>gM8%5(qLh#Osb0O^N`?R4mLJ-Y|YL}6*?b`ewHa%UW!maq$LAk$V7hG zUhFXi#mp`CmVv|p)G!Vc7x36Lb2c6ZMlNVtEGTRZ90;x1A<2XNjrZ^0BUh>{doxua zMSwL215QE5X9h~!k8Lu6($bB=v|-GWV>abL0GfXloJj-BG5XgA=b zz+Gi?`a}R*;mkcok|uP&VqTIur+yWmq~t*(tzk7`8wE=P!UluE487JGZR-gAU~zzE z<%sTu7wcpD&G7Ja#{UWpRkk5yGeWQ=F8a%kj| z98-Wu_?ZahZ1Fn6cB~E!nlbg#GT^y_TL4<-X;=c{wh!=vV*q1`)654#oHd}7e02Ns z;KmebCaz^XI1+K~?3Ql9G~_FOS;pkxIRjzFa+ICE+lG7}l%t zC`D{85A}HBKp=`Es*~iCsokJ9^juCsE#T{N z>y`f?CYPL++L;5Hry(zmE$@_tM!$r-IWk{&ZuokoA$Bst8K=t87J1=>&;GtoSgV2S zsYAX|NSUI-O9Sd8G71<=8xR=TLJxFC-(>}y-!}xc<9vrTSVdQwW83GZrlz>-_IU*b z%hS?8)R7eG2}(h<5D9qLbsnB>Q3(+d9H{S}$?&w{H3cYo{EPgKSITgio1G180a^?B zv>{PIQpjoHEh0|H{{ARs*}D7^S7k;wdds86fRP(=hL-CM0`Vv;L?(5}%Am|j8;Ih) zdoNHJ3A+hj1PCeJZR%Zgnw*>jn$zMgE30f?R^Ez|%}@~x`n)4<5d+K)h{N(A51>av zML~pQwOf{%0H(5h)CC3Lc(}|OwRHcFMYK=vzfPL3r2 z%}aIqGy%i_^v)Sb#3W>bFu;2K31%k%x{P{{9$kgxXygHu-V6n}dG_SV)&67w5YmXF zC1K~QSN~AE6SfBG@GHpP-dUM*@)~{gWP&BRr$q?Y=!~lCnaKz3}CxsoNleb*}C!jYMdqC_6 zM0ESuwZfY3n*5-ljux~v0HoxwJHSmqNn|$z^?K#X6;rYDBL6eABz|KYkQsm{k~I2u zlKK=C6>a#EOAW=TQz|>~8UTwHZq5L*aR8yxl@rJ6;P6skXP_-La-4srJpiVOx@$Cm z72BSk9?Q#*YWT`gm{FecC0`Il`6M5u;fhI4PNouD{n8T7Qf@W0K4hBcz5bmQ$O$Qw zM)Ai>@fESh_?jx*3->I$EfkJym%LMip9s9B!Wp++L=e@)j0wpz_UuJwZ2277JK8w!2gDA;x*x4N znmO;^Aw#XY?S1Rm(srM9%$e8@2W>#F^JV)Z7=}HUI9k_>1OpK=#2rQ>S|PSWAUl%` zvIFK9No3DjBP~!MhLt_fcuTEOz)f}W>{A684L}E&TGc?D9;>O%h}-mKe?Wkft2{kJ znZ$8G%DJ_RzM<{{y6)AhS0JPUFbEKyjFLWjP&z@I+Yn%_2pMOSsNy1^e160SrNqrH z#qe7SxUEi9S62hIl_S2CUw(}I>?2E4*f*rPjpDBu2UzHuqhwpdokCwQE+n=NM8TFA zYbRgROYBY@qXt&4p+&-9v8rH_&L~?Ext%?v&KkDN7r=Vfb02y0!2tj1(6b>n-IWIi zTXe*1k=#ChKON_)(&o4GFTc_S2-;vw&sz0YX=eCL<3@z})A#xcPYs-{IGDDy5h)dw zYPBRf^epV807Cm?150cPMsPjk6=nH=*r%}uLp1M1_v)ZV$V--~^T@N_GrFpFafp4f z`S86Sax49{xh!oF3hMBwh!g0)WzT6df&C4%#banKSSdf8Gd;JLzqho=DJWR=W`-j! zpCu=`FiSZ=dvm62CAP_Zn6?6&`0*`X6x$ch`6qd6dO)}g$0}GQI=G9 zJVQZ|r^l+AG7wyn(rzmAe##Ne+an$s4dVHU@7qoS) zfj>!e^~##vRW4@I7vHfj&FWip;3XC>#G|&ca@679a!v#xWKuB>tt^ljMfq_Yj@g?y zxLS(P^W2c;S6~woVXGb8Vk-AsWZmI%L-+4DT+xZI#ZSybzfBj3a)7f{Ss>)_4SvFb zGL%=Ky>9ivA?WG}0k={}8QAQxAcRKpaX>`0-tRWX}AU1b8-Z^H5f0mid)) zgN4RG@3pnHSpl;J7~{YKparUHf1?wGiZCMaO5!__x;9yBVShLhE;;D50A~TyEeKmN z6K6Toh49aV8jvFASx<<=9hOfOI}2S4m)eInTSr6;#Z4OLsRph6At{2p!HCopD?-Wv zE2J*ZEeF;dOW|{qI#6jL{q8O@{^F2%FzgwA30}4+@zw0~&!rHifh)DM%Y4?z5Ss9m zX?#$aX;4f0C)*L?geJAUwGzEc&%&Ze&Dz3?l4FbV++K&pDnW7>SXwwd9=tEMYxS2~c$ZV(hp^jd%5 zSAPeje*ovO@!PjuJf1}J0+a3|kMV1|Yrw_Y8Ze|fbLIkdhztt^xP=4{XOU50rV2i9 z`v@GqEmHafAee!xQU~S}09-I)$cp3v0tYIWl1aeoW6xq7@B+Z%gwBA5kz8s+kHs6d ze0}EaAo~YowRnNNAK=I^i+ebMe2l5a;6984$yNqJ1R!>0tOf}`AZE4yp&HKlAv$k#VoL#gIc7i#Ge&v0e+#( zW;8pL8A!7C&ZwrVxoAU;Rx=|VK1HjIL*j89h2oC}{o%#Eq` z4r_BaPHRxwGu0p-$unqKSJCywx*$0_Dt{ibT4Hq8p~5*TuMoJgX23 zgrn?Xl`dcIKIP$Loc&biMR+yr_#)yEMJnw6B%qO~s5b<_cu}!oolV};o`-M!touMnVA3tLayAMn~Sye!ITcLQ{7@? z6Lwm7WM*atR)+cfrAwE9F=<{p{PE-E(zM-&5g;qXb1xNeWWNj58%| zj}X2V05=_gW&X;67M{sfJ!GSGpp#~fVBZ-`KiUywYk<|%xAAN0qzJLGA712=ZLE`I zNM{nXbW}|g&F~OgF*&FzpLtd`?7sbgdqU$7_cF5piIJ=M{g0*X0f>(%Sm z3F6PSv7P-SP-r%CV6x40EW3wt(AW$R{X`X6gq3D8J(JirG8D~bnzuR)!U z_w(m-AW=nPZ{`B^72aovf+ z`F^bCe5JaZ5(%;FkRx-CLOPPRue*Sqbf**3=y>3|i#^xIeLze}|Ft%GpUyd2Z34I)^Fxi=Cd439m z=kzdrgcTiI74%2Qxmo@EVqaX7k(Tb|1aU?AD}aN=NY}(y6$wumA~4XMly}D{8W|B0 zk(f9RKnff}8u5Jw{pZNH2?iP!D0Hm?cAeW&|K8ZxYQ;>gYi-%ohUYP9qYGf?_~I}M z_HEi8J665$rnLCsC!qthni9WeyZdG&UgA{X8(YuOlhpReW%MP*yr)-kQrvuEw#*JZ z$AZm9L{jW6sk3fz=n)sSs85{@RB3+IT$b(i&<^@FaKxIH_peEs@nET|XF#~OWsDIG zGGX(bKvE$xN^%MPCuMa%tm8?Le09g-XGT=r;{B=J6EaO@3kG&ht`CT{8$#l@icFw- zNEUoQs`WJ2aN6%NAnTx?EtXPL`CSwNFt9J$CQpH|e1=6`@z*jSy58Z`zdz}F=F`r< zDYJT2Vm2A6DY+JPHEmPDPKP(^J@PM~ov9H$yZW`;IoWmPxB}MA2{m zE7&&j`^LUKizNQ~@QS)tY&cjPpMku%<4FH{4`h)4zTo%&MkSJ-essbpz*G$68PRB{ zMcRn^bH@*n!t!3=VAWrjZ_nR3Av3#tf1ib>aTT=ZUB#Y_;h(U<#MyrXIC@L( z{6r0DMWtBiLS$TR(GLP1?&L{`!R5ZXrFOulyJrOj>90IHJFZ``mEpH>-cQtE+%2=N zHybiJ5(V21*O&E6yKx7VSgu0j9-ia@fMeiFN>JlN={}8+hhR9hTo~Hy%&@HA$ zaqt_oa;Mg*NDLcdO-1ywt3{NYH219Fn=c)2$`h%Of{MVlG@3UpMf9K=b6y*5xNEVY z=?FP&Mg4r7T2A+;crC%La!)E6EOmI~HBD;~T#vkN(@f;@#c&@*VoK-(O3ba4>J9vb z3xZ?rBQA1KN7bG>4cg*N2wZ2W!4T_h2sNT_M3?$FI|Up{VDHhttId~uIlhEdNrPc; zaY8I?>QZ=^i^44@219qde8|JMh6s_l<~U*bolYI#!%>qyL(OxgZT^kkBztVsURj?? z#X5Mau<~JDYl6Yj?Z@)YvC$DwD9f-3e!dO0K~zEcL^ZGB1Dzq~SD01x+jhz%?Nj!< zuzPibdY>TYEiF_s1P_DaQX!=E{^y$SoGr{b!PJ=Aw=fcTDkZfk-D07f;DN(YPI9v9GEqn~X5y*P*uO=wLvv)P@9{wG{3!q)l?kh9{`_ zqtcn)i^?nFq~P%fdxpZlALP3M`@0w?9MT8q5;6PAlJRhVluMFO>Q@s>=I*@wmw;aL zOq;fZ#GIGyPIxDN&4=3lMLU{Z=ZGF|pY~q$T{mwUfj`{)m{Z6U-Hlhrp@E3HF*Vlz z{l;~ShBs15&R~CvZ>MJ$um?RZRfdK}%@gOlj51K&yffM9bCChLc!rhJc`~$fwECk^ zi(9ZYHBg0!sgV;BbPH0nJ~84Q8BZ5d61H%O?p(BoQ6AzSy#K&zkl15>)5E#nQi|wJ z&M0uq46hOG$BP`-TLdg{9b4jG9p}rTlXR`bTWd!gmUXE)a5w?QcV`c}cMplI?kN6g z;iPyhf+})FLiH5);RMNt{zo|dk?N<8l=f!LxbaGWj2bpF$N9X|Hzac?bL>7Sj(NG} zK|E=sSTR+9asgcP`KR}b4vks6wSV!r;(VoezJsQAACcFzko)K1RM@Tk+N$#F$ra}> z4u;igwdy}V{HGq%|GJ&#U+_L0g5O&=d|!W)OdSN72br{|@;t6zH{gJ<1}lHur;mOU z!X-BMV%hpckw8DplRx7-e6V$phoNEnHNo>&KaSstAsB0-RyOPR=ISx5gm zSkA`A#?H=8POinjS!~9F-}8G=RSV_eO?+EbbBK8{WebnWdBQyU83vKHHpZmRSL}mr z27jLP=l6eh_RkGCPW}a#KOd`j`R|zg{aDL^DYj!eHM^rZOwx5b`Wf%gjz8I{r)Th7 zDZU!VFD^X>hc7qkpe%HCtvS`}baYF%Epy6-ZF4{*=Dxbm%FEM1KWhCM@5>#%cJ0OT zcZKIQb&_16*PedoeOlx-#%<`)8SMzRZVP)WI(CXV1{@yC{HH}^WM{9Ybsm)hN63W; zUcNoNwgnETG^NV_JZSWl5$t;c+mdv3#?hnp2eN|WdlnTezvusp_W4*LJ_!2*(VHNx z2gmc%^hgr=W@pt}dLnTn5*caM_^tOE^|n^(_e(1kzn-N(FlWHcE^_Uf`J)QI!F?K+evlF=)%eMtbKSu@hcyF}$qXLDZ7vpmG!E44_&OIU`Mi$S=NM zYnXoB^Na6?_X6wlb-K=5DRl{iLNQK7E>mg!t;ym$d}LRv&rX8Ltl|*3)?*PIq}U>c zH@=9~-dfEcvNdu~1H+u4=Xpliew|8g{mLy=jSDk;L@3o?RnU?IF3#maU zuG%GH_XY0r;rkD0wXIcJBZ>9Y5xiT~SSxKLAE%{L2f~qVDM!L^dnbEcf4&{B?b4C< zGo?Mzd$(j@^z!o4M=QcT9d#01*Br5-otg!^rjF7k&CR^JkWS=ozzw{h6(EYYF)S(X z=HTesSMRSeX>M#QGUscrg)zncvIrJ5tK17RZsB-T^Dj#hV#!mpo{S=6Yu$$CEf;sS zGU^6ee?&{5Q*={}RY5~)wTL4`eGX5Do4m{0*jbo@^OH+6YB|GE6f&s3Y)| zp15F@t-IYMjX`k)5mPwecg+_5iC2Puu9&FN#8UjTFD?M~8707Qw|#4^muL1oYVJWd zJ&X;Axm$^KF~V+o;8k4s-4fjUH&3NE();_>zOg>HXnb(Dz-yN_qNJyQP=7#iMF0kc zs%Ny}`rkVi_~62GEV8#VbJHU=oBaLd(YRjb8yWMQ zgq==LFBwv&ld$NynlUOk!9~F3j+vy-Mz{oWk870Z#+AAy) zZ;;c;AohH@_xModh0>{2o_WcgU5bkVr$u7PyL@uLGA6!Z5FMCzPUQ36Zr<^pKq&QY z#~V%ckk2Z3ExkD7`3(30;>xruq!q?gYx%R^85`6RS&9|gEG|5;Q_#J%B|-?WC9Y+S z8_YQfEKoF6yxJ@`@R{JI(q938Z#7LCgTp}Ll_amb7G*AtAK10CxGIZ!q-oKZ%=wo) zx#ULZ*p1q5c^FUTbHiQdcLfs;oUcwW3%j^L)gDu|{Q}}m!O_ZrRwjYFHzo2K`!)xs zHWJ*PBy^(7v{(Y(eC!)499E?v_@D8)8MMA)S$9C_WcA+r9t?4qe3`i2cBrz?K00t) z{z02aOh)uFiSBoGHr3Y6_~aR6^w^G)F9vNMfCu zp`%@AnKEqVODzRT@OEAf7w5>{C&XE1QCTPrY^NKg%;~*0mK`K6Vc5DI!faO8a_$}3 zW8T++EdjP->=5LcF-ey+D2>VP$20eT{xW^eDlk?Kd6rR`+3;E90eZ(Xuk4mz{rRBZCpwaSSuClD8S_|->S#*k zklf8dRK!s*6L;>J_7tkO60bqbwNpJrD6Yb%8X-mQ?T-@&)YS_!o)GXL2H7_9)`^l# z2>%PJJYGpo^iamcTEMPPNE$c^3qw8BT3%SKOQL#H&~|+11o%j-nU0npVOR z;v!igG7+-?H$U|)IFnO9P@p$3sQQE_{5C)*_|*8582S3ugAjD9mW^{bc1&flcntYf zNLh;+bu`aBGUB(HwXC)zP@)v|QDyP>tH1 z^1j4_A%c=%?8Xy0hsW31)>uu(YbBB#k(s>4((RutlO)_R@}rUdk=u)Hw_OKsNU% zJ%aD~l*Cu0A@KDnXx;G;sZoJaUbOA=Ddx6p5dsy2?F%$@D(*p-x3y}qDProTJ7BJnUc7L3N;W@g52lrA674D zkMo=stiyLjI;&a@0GcXo3>tlz+e+$?sI1P3}VL)vE?iix>byM(Vz*h+feH z^ zSN691wrqSIluT-O4VX0EW5jq4y$J$h>=F(bJpD*H{qet@6#2?F9 z7B4kq1npR>(cIvMN3saGqnfB-*V#o1r;P(MS*vrihf_^qGsQc^YU2%unlT}JoS^KR z_|#iJsyYo=Tz|1PBAH;m3MAp0&VBcC`Xs#C_GBX5vD@m9ckg(rr;ST%1}tYknbyKL zZ*VyrQWW<&&aE=v9?rg1Rj4{Q|0Zs)&vdPM=VkA7}Xf0$(^Cs zGwqud@r|C6SgR*Ou|)y#j-ymW2n2BHlmRRvB2{>AARrk3j9U7j1=2jd$N}Acp3^`b z{}><_E69U-NGE3H(U`2?9J-9LdP15^9pvOT-1KY*m38959x5z&XMK{FZ1cqT#nkDD z6WGz(uV|a)-avi1V+*@lK`EV=dAn_t2ayEh#DY(%FOHCTT$HXjM&{eK@HHSSM(PW^ z!pRXGL+`V0ne@yg&@3BG?+-b(&Wb0oTQT(P z<+fVGrC&(@0|&U@1LMcch8pv~fXX=1R<#@I$W>A3-3rj*fY3`3PbaTV0E!lUUyaOn z>@^S~vz#oVR&Y+1cTI7n0K??3o80&H6#wTLz9E~of$=z5l54mB@5zW6#;{JlRnHvd$M)Ru20MX|KArztj0^DCbgF<)){9 zwFoKE!Aw&AL*~Ia2DZwJN9AN50rwNgkE*Os=ZD>3wP8)fL*2%PZ%|Hg8;m_aJhzJ$EA7lQE80<&}a6 z^~;6Hm=y7($ICbMdsqB6 zq*3#-xIB6e%ljzSuGsd9CoUDzWOfXwFtRR#Oq+b?YeB!^;NQVmI0nG{-#jJfCj(BH$ zP&^TrsPyBEzVDNe!!IT=bIwzsR zZYuc9bEiEgn2JoJta-EkppxLCbGVFmlq~3u6z#8eRJEiH2oyvP2(D$CPXKSpB$?~7WrxH4o7UaO&m!B>UFKV8ZKWOleAUbqE%e!n@nkG89 z`@u;T`lfGWwuf05zRW_aQ0+bSnae68EGP`uDQKDy72P}A4ThuZIB;3Wy2i!aY=|7= zal#dvHqL~NqQy>q5l$^Jt$~^KJ&su2uXJjcOQHh|{P}WqKFvp~f!;(1j8bRq5? zNqotC&&fMO2is)=Dg}vkUN74j%+&WAn`0X;D_k!9@8`({mJGEw(I7cTUc^CT2L%|i zDviEpOj2e`du+%uSsQtlO?)XnenJElXrRJWWA`%U7F-{ROv=tb@f8f&7q+dYI^H8> zQUgM5uSIXp-jLOS%|xE6#OKg68P?60tx30gG~n^A`HUYJE$`ai7Tuoec*hgitvO?& z6Y^TqmBuzZmty;g;Fa3wSbHUn$9b&;o2NtB(z2owxR2i!S{-^U^;YU^*sTthzNLK( zcKk^L>kc0A?8y8Y1HR2zZN53-6o7#bU|c*17-j8`u`WI?58pScX^ihjQFw;(A(0|H zH7yV3U;G_l2~22U!*dhX;1+kkcbJ73Y&3?yX zi#hO!u5wfXQex&uT8y|btzh?uGAS-`R~Vx&U0mD6i8}_QMN{B}r61G4sKc)ns-#93 zy=HB^1%a%8S+=*uauM+T5b8oYCvMKTtO2&{>~!dm%=N~$oLs_e5GUgm5#g`A_JTov6&+xS9v2q#jBF2o5MKP6TOd5z6OnJFaxhs-xoa62Wr;DM!pZ$Ad#SEROc zW~m5RGw;3|d2`%Q{OC6GM*2-uUN!>a&`=S_A^Y5LV{X2lud4@XxwNW1#1_JtAUm+U z^u#KVI+E>5;2L%0_G0k{c?R*x!vsi5wnbd~dhAA%SgFfE(G<&QLb8(RqAXMK3D)^< zqA!bME?P5xpqc*s@iuvlKo0F(8gapoVsu1FlVx!o2BIKbfgl8MG1$k#{)w&mOM3kig%0&taMxADW@3sfN^}vSvQh2oA zyHbJj(9S;e`x;p^wdq%od60l%FRUN`!A7`UmdQu`61-?nli6Ag_70Lq^!3KzwE}cj7Y8?iv~#7HxN_O1agoQ7oC1&n7uZUFQMQ_WZn5C8$YQDL$&MZfcT| z2Gy!^|FuHGNGpW=6-rXiaBSJb-Zaoq9X8VuCctd&QujGkOT7;iZ?r7V-fTIedzx9d zSX#Z0Q{9f?A$&H)fD6Nio0VgdIgT??jON7O|#T|G+(26=;BQx~ag zl&Hg{74x5RCSYf%uFzrcEYQTKAP|!Z^mPzkQ@LyepL;bjn_l`n=B*YI@o}qaqN%{3 z;6iwEH+3={amK>i=bQOa5+WIfRuRYncJ^>WYvyjI zT!adgX7;S9uz2mUEyaW35<{rXt3{NR*X)JltWwiVn&5C(gezN}- zII^y=f7H~WZ>sKM2Fp?h(9mZC-5%D#o=~HbGI}o=bZe+YIYufKMQ_f#Z+2g4$Cx8C z`$D)`&QOwQt+LnmiTD1(=Jn((Q%ZwA%+#iuI+#c0Y_FS!2rbf3urAl|o#i5)e>53@SRF z$VJ*{LIGzw0?f?ioP6anlnqnj-E_#@fdCdep=|^If}sKuZ_1V6)w|fC9NN7!oj1qQ zzH&E#sMI}qoP{?6lDE@#*?Zjdc!q@r+MyW*!pL_$Ow}>;_#aAr%O6Hh6fp4L>bKbI z@xK6&EL1r}*%cZW@xg2KCpUDa(SLsM&1;~2e4fe{BC6f?oRE#19SMF1jD2-XHe`@+>0R`3+O7zHzys@}kkqbKfilJMKlid2sjo_B_W5i8^%3`;+00F23CV|H z^A~qZ6IHY%)-%m)w6)bFJmr@3@5*?4F4B&^!GLwh&uIYcH5tG5Oji#=z3OJrdih?@ zX?BaBDON7U-59Fdqpn)Wd{zD4KKCo*j8W6uC-Sw{vk+~}?{I$K-6h~p>fZrgb-(Rf zV&~;q%|4Duueme-hzOCC1c6Bz_+v?R?yiQ>l1b?@LtNj#n;m4cTOebm6SRtpYa1h| z8HFYII~LbI^z#R`yqc*{rEuDZZtcRE2}Vf78|4uCWrSEC8#9($6L!<|x)%$p_@LjV zpKI_Oov5V($+6{4oWoQK7k$Eq6-1`1yHAysUw`2C>{dbD*L=o*)G()5HU zA$-F+|Ij{ko!Y@UUxOcVVq|2mI%C$CZiWq5bah^N7sK$hTR`)ep*7Y*H0MQui)HDt zyv@}pR4aX4`19%BHfXUyIk8KH`pS9Nszhy7UFV{*l&Z|Cdn)feJFbk@4=P1S?8Sc2 zZC#3G)aWs3D3-a|=(bNBFBtaeT_k#_?x08AO1D%UoP&PWHzQph#)gUM7`lC4IH<}b zH^eNvc+e|b4WhUAQg5fz%b2+Gl@$ndik8OL)>Yl3uFUOY@qvA#J^7ECvyq3_!}ocO z^9OzFAX54XPO!3oYucxK8KROLYmdA)5@MP67-}>k1uz{8`qVRIY zrzTsG$Ergwt`#v-HyB~i?AeyO)%6dH4WclgFu&S^~x-OlclivW^|DGAW ze9k4}TD^tvqw(GkyCUh+g^k*Z(-Foy1ALbT1k3Z8FEW=-Ye}Q0I+kR%a+dp4)>M!= z^nTxSlYhoat}m!rv5qP+C(;lmme;rv4`K#1_tg3OjrMoGMu*s7OKp7;sLg?a`tEg^ z=VC@x=E%pEqMMqlY-PZURaesN(QLGmg(?Ok$1v!LPc@#HWy-RDx_SnEK2J6fMq z5d*V9yLS92XvX4jWj%weV*2RzfV=9jpyf8Mz};9KkGoq|V7*ogPfGEtrvWVrrviByoYeTM0PyumJrh_QBNJB8 z?RbA>`Jf~IEj@Ez?Q9N;a^F|{YgbUN>}3%)8MgVkC>8Cy>}=cNL(?lFYikSv!^aEE;VLoSrL6nd3`KSq zd@?X1K!AmeOo-9;YS*&(ch_)wrEucl(Ds5v|6%SG87@&L3O|$cTBZ9>E`VochMTXM zDT+#Ai%a>7k(D7<{n$10@#Ddom&9C$N+tulX-!xQ3Nk6x%Yh^dv5<(GFNar8z?#wm`K5B z60`B|^(rJgyMB2dTT^^l6mgW~3|L+G#uj0f>mWZze3pWogz1vS3m64PARbPvG zSiY8OQ5Z%Ksi&z@K^krcB6{$kS3&BkR)tzuJ2HcrC6C0=4 zm>*SDCxdVN{JAAQSw$u2DwXGXt=z!!a)Ud5=UQjv}(au8>u~S!Got?KIue?djIz& zOK!U?RwQTGQvpD!xo3Nkvj?gqW^2& zP0m*Aus7T0pLL~qy0IE^N*=4o{LbpE>(z`j<&Y4c1yp$O+$+(yl#7H7kj~K+cI7K0 zc-#V~(`m;^kleC0$m3!ekNIuSJ7CKTy6iY%LO?$UQ~j zdYd()IGQF_JMl8(ij+v&PA4mvH%XcMfWAFPU8v9BF7z}Laha7eup~NOeqhnaDsn@n zJa^QxP2o;x+sR&G{noNw4vV~vCVI!m-gu6<*0La!>%JF^NAZe9bR?@xjj41peq%J7 zyGuIRz&n(X5*QA7ciz-%6}M|bA6qvvy_8(&aq+LTeQgoC-6tllOar62v7SCe^pn-k zagi~AClsXdJ{q@5kkEsUO%f>S~} z9=!r@)YM~xypB`bSsHcTshi@+qM>eG_R9NU(s7yE-urXm_IxK&6ANm8RI)UDXv`{n$2#d&$QhB zz<7b$TSuY7nEBRUxA7!P->kXMqGM@c;}-EIwMx9rJle9aDKgST7xwMWEfvfD^)d1R zT}xewv1^cg6J01Y@bfFt?H4ae>r%?|x3BKO-#9-&*?Thzk@@0_X5%Gll}!=qHD*gb zzN*|P1-@rZ!3L~BN>G}3iQE@dS91^ZuE}0SWHTMu)~G^faukQ9{N$ma>_kfWD#E10 z3;;+kYzUg8`t<@<>KjSJBr%dSkCbb46IGG7XJn?ugqGidT)c#+3avSX>m#8f|02+{ zUSXd-MIL&!sAks*#PAeXtP3O0I?5RU5g3TkLA`&~NQvDE z{|lf00QotYbd_i$0R&33d{+{Za@naHUU6;Y-sC@|=6{h&{l&2*l|MR`jX)&#z?p7& z#3S5-IbcwyARoD)5?ofQYInvr!~^*K?)@!VPp1czjTT~kl|)|z!T|pDTq;pps!(sKL<5#$mG6rD$Ae{Rc*QX26oZU6a)PCpKd}=JdUf&XISbE+Y3k6Uib{EA zHBKoS-ne2X%2=+OQ)s2-OZss&>WR<2s|hq2+cfIxux}$4wr+Ao-9|r4xnw@paZx_x z8_qN3=e$Q@hU$qQm!oQ?+ai**+(NV4FRP&gIi2Lbbf+itcIC|pFI&}S<+Eyn3)ZK? z1v|GcXPwLkzbe4fRIJ)3TiO93S|o{=p^cp51%Y1Bu)( zq(-?9@^F3I8MfS8nc3Hc&em1)t`u9hTke$ZFX97Ji#hQ>#3c98)YR@XP2U0=Zr_lt z4o=8C7mhCoNsPA5$=vpt54)%z6FZioT+sc_Fhf@BeUxO@!F=-;WPM58JoF^+AL)oV zgLp)>m#ambbe)y-hOw=NtGUgg@|v)EgpPlx@NA~N*c*0nrM;0+tQA?H2X!h`*+$hS ze+yk+I@a6?xQo)XTltJd{wuiS>wHdoZVaNqe5*kTgFg4Mc=-pBz<`7NUD ziZOCN#tkkN4>ye7c_#W@yk65mO--Gak=njpc;%%LTw8W0uAxY#1}*`F9LG}|38o}; zV})Hc-iC)>F~V_+Cp;xP#XH0xmGPgKH~4q&^rp{il85$lXcU{bUFBJ<5q${z_Ld1= zRsa41u_xc{HMb-F*wXgYC0O{Gn^1-JV)D+D3F#M|>xoV}CW-KUD`xa za6zGnl#e&(Kgv8~C|J(3>^(6Dwo4}+JpNxI-={9zd?k9COgyKgq<7K|Dk~cx`nI^) zHK}Ng8(W^w6WVr;8rvZqk9k-~=iD7m3>u~^8^~zHh4CX$uGJsUU&yq6tXr+45_LYW zUsOFqmSd&2aX-x7E0)PYRnhO=8^aX`e&jL8yF`;QjU~!j5s}vF=_JvQ+N2lgx6RDE zIMiT-!)>bZxWS=TF+8rFE;{7L1ax|(p&1o>_D9#V>HMN;SgOI5$mhfMU8A_o+9H`_ z-^jdxk3>1ec%E!H2V4aOoj-m6L+N+qc2h+WuA43%qn0|z*%W@IR($qF)m(?mr%uF{ zW`A(=cz9FQ!pUy77cuG6(;~SdsrB$&p3AR0jMka0PnMM*?4AdPWXEeKkDdKu3+bow zXu4E`=(|lWpMT_7ug;~%P-@~!LOiIiytB&9-ViKf;HFgK)et0Kif7fsUMRmA%HZ_% z|L}I!VNrG6-^bg&RdAG01k_RK5D*X$P)P%jmKKptNy)*2p+yjoRGOhXhEk-vySuv? zc-JW2-oM}Tyw~;q^`4819Oj&}&)#dTz1H5JZ$qTz{&;>J`BcmEGjvV!bI0E*_fmF5 z8d!X?kqD;B7YM&B5Yp9UaHP>C6?cMQ?)~(5_t}V7_+YB(ebi&m>E&%zno+Iph1?$> zEI;;&cBI!a)#Pjubz~+R`nFt*i{l+Z?^N@_$G7J*y^4q~$L{Ch?V=r(b8>IxfEgF}{S2v>|(7S8-EQxDcIBF7|W4Uy9>||L+IVv}%j?|i_o(~Iw|Ic2fj>P9I zSN-l5KC-mw?nn@S@nD_wkKd8{k~cOQoViwS`tGgo&!WUO!%y!bspqJn1?w*^{>u%_ zH&|l@1}NugO?mL|a2FN*)AauuGY6kbx&;4AO&h-sf`DM?#f6X;g)i*!6A)LVyM-6f z_Ne|`OTgIoH0Q5XLte-`$>Ee7+8?-kr9u?d6|ZsPmtWfw;jin$B%^u7JKlGMdDOj! z-u?7rKToWIf9ne&#PhV@HAk=1RNnJ96@97gpyk;q1}XGOzW>(=M8xQL!wKutlhb?`~Njm4;BqL$qfA& z3p(s)2)ak$X9>4p1A^Nm`1K8iPx@45?i2r7xBKTLjmQ6Qh5tXVf$vUVf!O0L8=N1u zJ@fPPKaIppgDe1^B3x)$gTc^;H^4@lZu4)z^8hA_zrTO?-1=8Rm+X^YvVwwwA|fK- zSD)E+)>ZDD_>|+tOk{n6z*asQ%M;N=K9%*)(uae)FwlR=iB)|dy5anLO@x#0@Hjbv z#CvNo0$mnat_~cHgB8%z^YZe5USgIS6A|$oenh4N0w*^&_Zbd2?&Oy)DI}|O3HlI| zk?l__6a%VJpKq8cr#GTwMwzPzw&KIi-0p?j<=XYt|qK) zafP$}Y`noJ`cjxTttvqh=B{0@%HfiDSRX<_(>518v**jttMUCEY~0W5GtTHrY@6-Q zVt94-v4oMnKAtzlYi&&QerRP24+#QlRF6rq13=O(EiLir>5K;c;GiJeF0Fv0!G-JH z26&Vzx)*kc13a3Wo15=H9+y(fdtRpx7rBf8 z5ZEsOwg_xm2E0cp(9X)eD0g}W_T+`aY{ox73kHxEfbcc9wzh!avknB5n_AU=ot>Rj z|FN1RX)S!U3vo>ZEvBKNk*W>*Q(^0OWnGFn&)6gzp#*;TzCUy?Sy_zS@`?;>7IF$xr&FwA29&z0Wfp`bD^ZJ z&Vqb>chHq4>@2VkY_4-~WL1pg2lKxV8jTay1Xq27gBg&(b@&X-d~|DD8|*Qz11?-z z+LHSIwt}FbATY(+?+TfjnPp}G0nKi}VFO;GR|eRX3MxB80M{d%8yj!JI{;}XKQ0cx z+Z34&yaHgo_5$Auw)5|cUUtZ|U0X`*Ro@*Z>rgTGv{F%2q(kD_Qhrc6D{tv3CCV^}_d!qRl$tcJthbt^SQ7>mW{0FR;Qo076Dd>XC09 zV9pE;iG6@zB>9<5SO_pSoIr#DUIna{>=VFENqPeTKF`~5y&9MIRh3g;);rj-_k&g;B{|rFQ8%Ak;q}VGT3^*53nNq_?S}!wSZgX$`#t2 z!C4daZk51g5E=y-0SaJBc)S@r3t+AQ(giSl(Dwh^NK~x7UhSF)vzz~{6V~J<>3c7M zwi@==`?ILeTVVmD2Af_pMVo;=$7EmyL^9(XI{+!cX8kOyhZ-6~O_8Ex%qqFMO-RIN z9!q`r4e+}#7z}KXf5V4IHYAzrWF_|MgxQEZa^nM9FiY;>n_}4Y20*a1suw&XV-$B@ zJrTkWAON}nV_{-^+z~i1z!In`+~`lZF?f*Nx+@TaMj|nBp_r+u3vZm2Gl+7ldxzVC zR7p@nA-NX`-{MkS)OHIr%#Yjg)5*zV>;0d3a1ZkL&M_Q3xJITzJ_q@l(7L)h5(E>} zP%mb}H~~}#3mx4>PQy;;GnZh?woXQuA!iSato9Uj6hQlc^6gd$M8M2-7_u;^tSl@n zQeEO@Vu4{4ouAJW$j^(~+XYXIG%F)f(1ow(&Yx#@*j@#A0U0T&7X|mShk~rE>~%nc zTTV4c17HU!t)gPzCZ~V{>|H2_f&9#~gDd>9Z)5XgL?m=*JG59=stEdCX-pum!Lw&(3n$ z9-+6$!Hr|mnOzr+?+Clhi;H5h`K#Y#A%A5`V_&mP`7*8TbF2Xh&3rI3x~s4{z1sC` zeiry9Oh|#FV7tG(g2E1N4;Yt-r5)Tiev&ai@4rugrc5&+JG?2ln|)Thz@;E;CcnC5LL5^C zOgkN53w=m~>`uU&8&kl_aB*^iM+DN8sE7zBC#QcvfHi09?5tr@j@KkG3DeuV`sa>RzFl1t(DDcOGAkyG+y z^i+mK>@CCSV%64YN6QghfOMfd>F43pg!Ew&y<-u*u*MD#2KR z=ivuL6xgCR0T{EMo}TVxENZFGUa$Lg!O*va{|WV!w~I?6jcw9x9X9E^Ta(yU4tPRl zI=ZD-cc^{oIjE6M&*zCvKTi^TMYxqbc zZDw#q)L+SR+`yMvJm)$}K4ZSHTMiS~?b>(w0H}P|pzxLhXejWl)Uu02vIg=dG_mOx zs+sKcD%Jrg08h2V_JsupUm_uimWzpsx_S` z9M^{8G2QvO#(g}92M-^f^J#6BMbOdGUI6^by8e~3c+?%9EFtwMTs6;3>)yR@aFREC zhP}B`LTqeoJa+^SI&~vWl3uSPRxE;BjRBM=3<}r8i^d-tD1BtjZdUfydvM;T?M$%n zkBGo<%00Wn2nAv2h^O5z*lNfO72c3^tK3R!spDy}Nox^`stPBG&v!XaI(9pR;)8w=U38Du>w`eg`E%zaEZc!1mSx`#cKH|siSTEBhi$M+ zLN9=9_H7$b-wB7FXXH*12YSst@-9Q>;23`@6>4Y*-DJpqvfzPo3 zXe=);zbfdgOIIAAW*d**T?yYmy*{{R3*d^Y_*h}O_VmF_|L(wu6d!YOO(MG@@ru~Y zGhsiZ&Q$XFeZ=Rn!0Sj>5f}Bf02?p2$_PJ(H8;oNjLm6EfvvgMvA|z7)1B$Z@Np0k z^DRAFQqPJ|Ych&s$+hW7p<>v=(ob`<`hN6$sd~aJCt@DzW@PS4n%9;1u07yH=m)qLRS=qD0UWA9N}<(C9hZBhH)>ZivWXT6xRJ1%8sDMi=fQOY&=D`% zB(UA&Ei~%IJKtDeS6P`0jvjoiX^&A(A#l9F`|dA!qg2$?9O0c}V`E)xe`A`zbi|t~ zxK3n1JAy=ccV3i;ujj0;T1O_%v<89L(%wCK)zwCORS2OdmC+q(OPEWP zf>C5a$`sk&Hwq}U8Vxn7{~nqg5jn4quvx3*1=rgX^1dxhN!YJSD$>?i_g6=Jc7^>N zuhZF@+S*n^B>^qY&1W5eaRVmIvoY)9oj%}YGJ{$USW09%$RGhv2v0o)HZuSk@x z0$Wv*5M=DH#tP_-=vQTRs9LKM)xs?G=dsGe8@xf3tu-b(uB8N&q#EZWCQMJtF*e3{ z-`Q)zQI(cY4Du>$o&_2;vap+oXI-a{Xkp+%u+pn97j@59r8 zrKQiEI06z3FF1n*xqi1qV?vS%ZSF0N`ez2(5!^*x3x*I=NlzrXvzK{gcoX|1oxx`L z3sh~+zo>W`KFwf6m$KICshgeN~qb&Q)DNm_X0I_r$3En+Dy%;`wxVebL+R&(cY zwwamVLkK5_dKBK)tk$DI@77HHm9rAq%!b*_r!BV6M@B|UazvEt-Wnerwa`#Gcsi3U zR#-PVDc?F*kHzOZ@sV)nN8_*ehHT2r-Af1OA-Esk-t(;WqEOVA zrOmh6TpY7!5x(`K>)23zclmg9gd1k8-Yb*PtFb4u;zkHYIiWsC7v1D_yoXKrg|)ec z*kA2oTD)bfA`C&qEum9f<&$&x#gT+(e!cyw0&ek_+By8}q&EL0fw6t^tdHEt*VF_&&m!8uXa?s4lr%Z-b7Mh5*gb|xd zNAZ6@zgQr*G8O;)m0@LHlKaACWt$+UaEASbQo)6w;S1IQv~~Owev>Z|w!v~bdwt*b zXLj4HB`&xe`*lUz3}O5|+%V=GVbTLQN~-z0zy3h*g{Y@JQ(36sQ26=_MZ`<@oUR38 zpTaufZ`u$ywGJ9Q%!%@FB_CSNp|JSmUis9nKR$mgzR2lFxv=L$PMTcYydBZ*)GgK* z%Ir$}zM7O#Zr%C!W}a)*e48LDrh3vw=k+pUM7wjf9asPHje9Pn%B~PR5V|KjaCvtp zo@akrWlfevZ=W#kkX{?}6&ljEyJXG>muxEX(0I?owBm~0siG6sCrUm8)@({+Wd#

qSMdCSj(`7=jqjrQrYSUQ>PL?^^-F7k^BX+*%6n^CqD3Q-e|n)5)iG*P z^l~GmS^OkVnQ5acVmXyWN~FiR`~s!H%&FetGdBB$=DGx%`?G!4_QaNZn@bzDhLW5P zf+udgBmGN^T5s{ zcUl3mNODHv$q&Mh)_)i4D0${WFu2Pz%1=Ch+8d2NVRePT^xM$iCkZ~`KgCuH%pf4qXw4WVu({L z1vepK612ey+)beN@6~Y)0BGACJvuf9IUl)MPiR~qaBk7YSr#V-A$~r+Jnjl3uP0ch z>;F8YgQpjjncJ!=YnrQ;diapjDMI`+W1AKO1^ zb+AsRq-@9(x|-7|QaBO*(2pDUopI-I$WBb01-CB#-~s&S(}n&*O2P`r`b^XXZv+(Y zj?JVL8+WFDMJK8iZsVcYGs#$xh9sVSm7oQnaFDKgefJ6NQ`Fi~=12&r4ikKu*NdFo9}LUef}?f6rDYo4yfd z=O=?K#E&0BBP(dcHSnUV(+MeR5ad~5`8x4Sq5gfSAX}duH}!(d9z7DcjJNYr)r%dG z>2QL(R|Z$Ms;jF#EcG4u)vuh5iyOBBfyMqpnAXD#MzsUgl94uAh;07;+7wyv3hq}Uoq##!H0}-b_qSTU$PTn!6OdSgsTQPf+PI-0 z`0vB%RAX8Gwjmx}$_TB|oBI-@6m}%!!-qvcq^wNX%H2yo28qkHgB=CVP4 zrt0GjmZHt)8X8~IazN-4|2=4gqZ044RjaI?Oag(?V9 zX+fL-2(#nB3H1vJ*@ZzZ@tJK5`p!N~+vVg`0X7WiTX-gTXn4Sa%D8aI4TLAg)y*?*(cE_y_8?NG%bv30d~W1(%C43m0hA-iYx>+^#Os*U)z(sQ=(j?^egmnX^0Evp6r}w^e0(%U zJIL`|>T9oaK^F}&A~V36^;qqKgf(On`{5+zP?~W*Jq4b@=bk|zP|Klt*&FV#E=fV8FR%Q6zc$bL$DV*XSb^SL0^aciu867 zNAXHy2j9({C|aE+Y)-g2ky}JQxwV-HNA>N@yqgoxl0**X5`(>}iBlQKW~e;DB<2-J zOr$_X=i?0!(NS!Me0yw|ZME3aoh0z4czF$=6SH$t5{g3Y5AJbxV~0arInONGheKJa z7%U=WQbF#5#yTD3JMfTwAXF3aZ+?+lhC8_UAdSS*<6R$gOjmp+9|VhcK+$Hjo_qu^ zEr)_?kK_4%AXiUW@T@E^d(qu}EC^PEC#R=Zhh_lfhk6jl77M53x6k$A zHE6Zw2?b#p4o)DFgXy+gPt@gy!Rp(h^E0VI6I%+}+xPI|%;u2DzTm!)tdvC$NhT>j zP`3u5J6cd-BuCvi=$vsrah);QAL2eNZnwoL*HCD%%J1Y2oG=P(U8tT3Z+|>y9k{|q z$Y{S&@6;e;cjzFZy)8=iayt>N8^mcgS_y5H44w`IYWum;T(6 z?YezN#}7e4^hij$mz$;1EU?K};m-J00jkww8_5RN0-&@b8EyI?QL-}8Fo3l-v9ikk z>1dqLo4CD}^>{rc$oIc|nU-d=IC$?`Y<4<(o+Hi{0SP8s;RkqSMG9W4$3n)rL(bp@ zhCxeV6D&O-qJ!Yd0T7&@p8R_>XMdRN=aP_?t5f~(*|&(R$fg6``J;#nXFlbz^M>+Z z;_T#p$DtimqVd`~v6A61@i&Fl?EqP}0ShmoYz{+|h;Q|S&rj`wjIIL6RDJ(`U(*C+ zY$0$GkHOTt`arH1ggn5UCP;CEc(xEe*=3eTuMZU*Y@eEsOERgif&|4rGwwd(X_i>6?UghMWqjM$q|bOyJF~3FiV? z0ZM8(kifzBP0-JH45khmfQnt)qjV@v`}e&+iOs}LOJXpeh{tPfjH1xD^xKLXCdr@@QeDESOL~Mg`e*rhtoJ|IbMBFgwILtPrKQy2Lpf9g)*~e;sJNDOWc(U?%?TF71 zF%na8o4qJ=BPrJjsL2AS*^FHe2v#78Uk|Yw1USfZl$1=!Pf50y5D1OL{cfmW z>r?fCVJq5Qu7Nu^bKz!**t1Z0qENhofn@{vaq@&(N8m_9@+LvsrTps_ZqJRxTrvr~ z2DeFnLA4hjfNQ~IjzCn3XP1M?p8&NAdS+(4G9>2W@B40bt2w^kh8W+r#F1Te;q1Pz zs&E>02bmq;Aycp)a0BbOTN2k|(^yLF&bPA2Vw6(AVYILv9#GUYlTXGblS4KUoioe! zu;Z|DKyk3;i$}uYYe5!99TO5dFL1+eYSrZ({Iadb6yN)`u8qgc9q*w)Y$+b^sM~kp z6%ZFfY6iQBP0ueY_(w-4k|2a73Ir3K4;W(vaVI%C4~0a6`@-j}cQzIT@*0OsGmZO{ z@cuDgoQ<6wo;nuD_W14ug?N&1!g1aHsr-bBA|;lup?w$U^gHth>R4#9FG=A(Bp2*0 zs2&?2Q!N%2d-d2pgSkUuvlcL%vcKBRQ{4G+`wH+L?8B7v_ zby?^P1hC3KlWB3Ztfbg8L_{c%v=aUyb{*>Q!DxJrkc8v#>k?3Cg+WP;0NFPyo=}j* zi4$vMf;fm(Q>5+qZM@r3L<66EJt4;&Ls7^Sc`Beg^9%s*K1wp3^`D zlER)2x)PBXx=~GonkjUsqd}T&mDph}O1~=fKWUWw_a<(oKOg6=;H*v66|Wz|;4DCm zJQQA`tn_8(5Q4QQB=p$s10kzqT*3H1y_{Pn_3&*`&@R}41}X+kCJ@+q2{i!_%z@No z88{N57rUEFBB2mNyH`TZd3enhWP_oifCt@!3yuwU%!i-^uj7=d>m$X^#l^+UYz*(v z5-W)UAhbQdI>c>I*aSm7XTjr{**ySNnoba@f^yj%aB|=RfIV*r%U3|vT0aW9(qMEgJbQC5m&Rpi-1v3=HX`}=~Z5i5iju2aJtxW94d_qOW z_o26-^epUW{^oGweTC3Up*Wz3P((Kj5@Ge7zNV`Gz?Yx1oGvSxSNeg#Gy#YP; zSUVW)_lYzm47V>d8Y`S~%SfQn@FiNT@P~0w*Z78ijNu;arewrfFOx(jdEIP#ab{Cu z#Xn;}w{H$=KYT?j`$L440DzK39o#X54m%=WYsPFxeMCX^!Nd?0uaY3lf)Vi)oEqNuojbP&^>@CV*$ljB z1B78SW_eVgbe7R|2YNbQz@V%Qa=8$g4aP%w2jvM*s^O6lP#cJji=#zIK?+|s%f3)e zyag}d1#{d&SZbjTt4F|vT!e%*O!UFV$qI!?OhAYV;i+2WGOQygP+I_%NGOyrB0;49 z=0ZKNh$PRg|8#K|1($Izqyq-P*Mr2E2PH2jJNpVm^lbh$ln^=N#lT>d1WYr)V-0ep z$;CaT=Dc-l6#^(nnB3{qhv>cAfH_7n2-T5r00$GM8Tj{xRWAd#!v=1Mg6?eH`%-ic$9FK-|3WV^pBNi;b7zthk=9vCB241*btLu@Ni8WB z{jTCxyzcYSdiln~jw5eV5^tYfWWXSiJ})02BR(5(?XB@@ro`wW&ndE0U5L}Kk9!3o z=#m&+W+b&cZj>j0XwN`aWhMP%gjV>SXMU}wPYz5uM{K5H2YV;0Yz-u-B$d@13SCcOW)GNisBgot7RV~4e|0Ae zv`6AGL4kp<^aYCtP%y+NCv_9rlyLl#kdrVsGppys_>PT_+ZOi2;wVa9E00TD#xT%O z`#S@Y$)E)SQ6o77jh8R)2WP|E>*?uXFy&=s7wO6@>tGsa#gxnPPpZwbW{@|w;?|wC zu$gqp%aT1iXGkF@n()hCqd{@A+QLjURyHsElLY*C4PWxBrY>@s1~do*K@ zyOTC=EMCqnqSpSag3F=vJ?ptst&aW<%N-$^b2kbvvr%GN=SBO~iys)jIgDAq5;)%!KC+|23K+Pg}B}0c!1# zjj)5MHK-I8f_=pY*kIkx`GDa(_B&Hda66-w&j5@_PKD~TG#54g&6IqqMVm*$e$Grg z!cHoL3tL{Np&`0(^BF7)b0=kD!)FR1S9unk z#_T~IhTuznyZy?BBUf=qtJY|J7P~b7^1D9etmp4_`PHd1dZm^>D3q=9vwTsp$E2=M zTsmBEt!v6IawaEt?{&($iSjv;aN>YQwT{Cs)=VsPcdW|Cc{>UQk~fmPQUkxC>{&(@ zvyauKoUrDQFzhi|_tQ}5Uret8pK+gi)2dwy~=il1xG0p`XoBCV1I ziKv@5d)5UcLl1TP95T_H4q=}%Xyn*ElBXn9_lLv-JsfSoyAvA7(@n0>$Tgtzt^YOAt zN~Y?mkq14`^o_RF53ZScCiBZZ#o5%j&ikrs!D>FCMG05T2e@w$nR9r%nwVP*=u5s_ zU*P=x_{>GN8K0>BlcKi#0Kwu9zRTfC^GuegrRtJrLI=sgFERhBTR5FkC|O`)VB*i} zl?fs9Fg>`0*1!Ig8scPM_WLNNoc>d4(Dx;uJ$T_iUz`Sr{!11gf2Q2P-0$T_c>h;w z^2s&0RGruO%}M`0n63V672}`V`u*Dfnt4C*U#{Rk_CVgH{KFI*9L_duDWKV3kQB`z z?1=I$t&j=ZEfz@XT~?hdlZ#K>9oci*x4!T@@2&Swo`oPcm-4aXmQ`uuo8f{W3+X2z z8*|Qz$6Geq)l_0P9#zNOBgzWe-pZJ;yK->FpZ-(|oCYKy=VH9)mNOWFw%R>S*x?x z9&y!+bEnskC>EPc3{{<+qgOim*B(cY|9odL9e0}GMmFxT($*PMNLP0jn6Auc@fW8E z))*ho_eh*fs#HyVz7p}aVx7Yzkb-2KNAuZ5o2i*>rkURD>dacfVCxi5&%^<2%LjC0 z+#0X{e!7WFa7-vRd&;Q2mSn9ZT}HtK$3Iq?%+NzL+lz_qn?CEaJKReWOv-7}8)NB9 zh)ma?8D>j~l}VsVaNJveWo$c8UmC>qIC+l!U=h_I>d%VqT&G;VvK;Yg)Z*s+3DL)4 zyfbUu4ZWl5bR5jwb&8ablJ%?PA9Ro+cbgjLh2LABX5jQTvH9amv zXq>lG0LJ}tTrC{uQ8)aQheK;DvsCkVpSsexTE^T?_4aI!O(&5FMabHB9)yj_kO1xw z%2V?CLC1dhuTkUljt41posYe3WUZ4iThH5mYa5qw*!BwFPGvUEK$&B4t(>0RO9tCC|G5P@4iQMj`Mk=o- zp7;4YAb#HG()K`dypo`TZ-t3nbP%GSeR`^0oabG4sGw~xEfST#zBGQuSCCTUEf&^p z_sDR%p(nS4z3tgYrp6ForGhYP^EaOgHD`Fqw@45#__rNB-_1GKnyd#3Y+)Jt!#R;VUYGm!ljTp&30=R&B`%boRfvfg-}FK8 zi(0zcdyxsr4@g(S7!bDT;J1rWlRb=}W2`^#quCW7trcxEpGdv-F=X8|@>Ab*oc25F zIz#&tS-iqO67|cq(M*qp{(9{M86 z=EIgUVOW2VuW`yKo&LR{#_3dzvbQ5{H*}v&+kJZy8((_4aCRlSML6n*%ynckdRwPW zPEErL-RiFYMUPsXUHDzu>dfu*dxpWgto?7C}xgM+Cg zn$Kbv?P7x6E?T*D_E%FmnV0^BeeKC8I-N4L8FC;vPe zwT<0ejSbhz`nSscii_S#mk)+nN#4I|>PJUIP4>E0Sk7=RD8>TDL_2_J(#%NjvmL5C zWvR(i=*8B!&7RXhoXti}si;6pU&?XJBop1&Q1onPaaHH)@m=@!dU_@XIg+(y|H4N# zy%m0J5RJj|(DMrs^zHdH$ z-db(t$;$^ryIJ9AvB=IJPeeO6HM|Qm1}`^A7N?*Bd<1q2ZgLpx4>EHFahQ`7dd{+s z52*8PY+NLo$)_4V5u7vKLefKeXUxW^Rc^%i@?J?fD#>1Vn#r2bPIsV5l|A5bLRqA4 zQ(P9hV5@DQ@5XyxLgUYc!?19Lwm#usSNF%!lYJ-(b@sYwx=?$gh04J^5#7n?rLeG5 zJ9UE0hd0?KwIp<7YdL2;Wk*&h|N5H4q(7Ok5CUe2!rXb6g{6e9S6^INYN@##X4{jo zFrc~0snKfQv3`*O+1In`pxt|^0rhGC>?^#f%l_iBU&CdEHjG+#46VL!LwPoIWo^TY z;+{OJA5-i{V^H>)i=JRn(WXCkagr_j13jo-|MaiOZj2n|!SWI(a%rXPnTxX*Fj^b! zZQlqFS@yxq#V@Ju4d7RQ9#3e9i6RdD(rK`UCGB#CE=kLNeNAN<2>KOZIXCs?G4 z|Fa6^;6Dg9aeut%!S9@Sen5QTr(1mm?_2tTzM<6!{* z9CuY=mzpW5^h1cBo~XW-M&%H;ch9w4ie%i$in78pC;lAML;f4fBbXn)j#)p@P}W~o{*4nG1GKW@sHM2-$S_jH2vp&Tbpg=lZI!4#e<&P zhb>M7_L#V1G=3Er5ot=FTsMC7F1w(~*Ue|@e9kYp#s}ugxkQMS9*c1kAGiJW=5m+V z8o8Ik0%KhLgDjfU^DKw?SQR#i_!8use3W!`ZT@Pji@ZkNe_k^uyv*|oK}lJStS1kF z>6E)`d+m6T<6$q$&-d@QE}AB7i|1Qp6HQLkRxK5B@QFfVq^+FnCH^FFXQ;*R-u&W??v}fVe=b2wk7~)g&s+&qa(vqO`d)tSz9PO5* zIdAV%REN*Itsh^ll3^3R-mYbJn>a&}j;_hz+)0AL2)u1w#7_6uEtMQm$Vn#DWM}A^ zSep`Uarox-vo2^lZ(P8<`}$;ZXW+6rCN_GqOT>LNr^Ho?rL{OYid7{o(Ap%*W5XWcQxYMA({`Np6MC5sdxD{49?Z+$; zgS+EXMdH~nzgl%snG^dwhk(eA= zn@_r`O67EZ&IIG$o(m4dT388*1`|@y?&wAG(&X5ye!xuFx{B8ZSerdfnx4MwPDsU* zhBUmdX`Vv$6@gIF_=>9Q9`?i0#8?Jcr{_pyZ{e0})>|I8OK9gXyglwLEToGxvqb4^ zvo3r$Ii*a~^HE8O3pdLjVjFZ+Hvv>d@W){Dldfxah`fTxLL+56U38H8p!&LCQ&Z-7 z>@B$|YdNv|*9@kc(n%8B3j3FN@08btzbn@kmXr2=MdqM8?ueM1^b3I6SiH|JVcyXedD#@UHyvO4T)q3yd1=7nCAvy{A@I*$H! zx7(DvR2n)ndp-pld)(q5xWi*sO1<$>OBmVwHPoo`3lK~!Pnu&QQ@jIb0Zt&R_lONOrIph#NxwmCv3D;4DVPh*(NVA6U=ik3q zo9qdR=!%mo#}t0jx<1wa(r!y;ud|cAw;bV9)7`DQ4-8PbM+SNio~b`nG-F6;hz$?o zueoMNkURR*K9X2`B`oo*l?xE7EFT3(|(ovKxRGqmD-+8d5FpK;X`_T)#T(Dx5_qL z2r_-VVi>H0LL$$K9XHKBgZWJMIYd`_(rB~-n*{c6srSEIK2BM|1#N%x+6G+5`s^B7 zY{R%MXogZ^O~Q`BV8#C&z{))Po~FLIi;6`v?j{ey%N$iENPYGkLd(s4M&;#I0~)l) ziq;ou{cvHJ$m-hl!qtZ&YGLX^+K&FbHV@@kb`qZdO1L;Z64COW)zo}k^#J~hNPDzs zIi;48_mS|6#x@O_9Gte}x77{v>Vo_jI?C=Yj{#Yi_wi+AuNBnlLiQZJdXnHr${tvc zH2LACZw!U_O3||Zcp|rrW3ige6C`rMA#kN?A`kg7cmg%`$wS!vR6<2iWuAGTbn%m z+(UV@9X3Pd+EOF79QD4UXZuI0s+Zlwq|FVUpNH+So4D<2jITkgiSCGQ&Vc;>gt)fjdSfdnm#+loccm@1w9e}g~cJf-Vdi-?`~H?){lu6_H_H6 zY)}c8dlZcR(E%F_cHbZ64=-YgHk`AyW}-UkJ$q>IqxHCg=iw4zovHK)fqq2W^>51^ zE-LXZg{O)bXlYNN#l+B{qGI)__`LasTEl$deg@+cdO7JKZWokoxR=sS|4t0qqxmn0bX&rL8hUqV@JOe5I7glBAg`pRqN$E+wnx-uhI zyLO+#4gIPxVg?D*BB~Q;OvtN%@UEykH@|6RvYmUXnif-eLR^TDiXYl>fM#4N|Gj)! z?$nfAVfe3DWDqG0mUX1O=4+{z2uF%SIL{o4s*8;DlK_f^o)C@3JS?jVpYNjx<>e9j z@!}lfrE@$XHiC9|V}?6#Ea~Axn-?G_athqt;5k2^vEV&N38j1mnNQT>j2r!wngzox zo5rCI(%)w;zVEC_@y%D;^RBpNoTDa7k*XwzptI`e-dNE0vg3EgFTr&>HGttLGcT%Y zU|CeIeON)b?lSMq>QBJ!k{8yu-oqWxlZ5TdPo9x?N28a7V%z0d=eI8WT1tF{79O%m z@2jdBKKGv6BwQ%4#3r$?3Wu#`Y-}7|R5=`Z9rzb3sR*Y3f@y`RUUlWN!|iJw+A_yh z&DGk=+Y1dUx5In2sQN~qC}pP3Zns^7eVNS#pE+nc)XT^Py7qCyd*vv$!}~k8GWxVQ zM;y*iT-*C?8Vok58s_=Jz{uwtlxjYItguFiGn7};GB12%udam~9o-a)ilbdD4^}Mh zl6y*w_WV@t$#@CH9KW{tkAsso>lR%<`u=t2Q@xO9U769l<&UHvB@Co2Jfcg8U>F&1 zh$uUX-FPLztUv-;6mcn%__f&SJL$)r@}1gr4ynAMax7Wfyq4zUqofzlkJ5OZ?B16g zczBTp!sen z?~yW|#=LS-lU<&et-s#JcmzL)?oy7raq=y7bRJ(YzbMI4q%Kd3e7RPOG$S-GERENK zv@J$oe~U2z+}Ri<>ykF)R5qlcCbfmf`T)snf7}E=L&rK&FuqnGc`cBjHhs zaa)_M*P0dfFJH#k?Nbyw=sZy2GrzCto?J8U5EjmvYqtT*2c;?oE`AD_FlqiF9qNJ< zZ7GCjw_o&C{9Ls^bBaZ>LRs$n_XV53c4ey|p?en_7h-ib&c@0V&iPl)QmD_!b0WAc zaxu{$#DZNowLn4kskfT`;VZodU*tYQ;K@j;wXU#riR;s9?sr!EIQIBx4o(qwHy(4U z)}}~))qvESMw8DbQm}kGM6x_QKMh(1gv8AE%!{fU8a7N3jsz~lXU3?{eTMSs=Jw`f#Im>UET#74Vk6wq6b&Hb?r`tQm8u8r5OG;y&2H80 zjCE`w)1S*Vg%w5zVKVjgJl)tXrhI5iDDXbtIA*HPQKBnFfkqA}VtP;h z%aogx#ejI0)|8Zj@CG~?mWJ@4>6I->dx3V@8>~%+oD4360eO9L{QGxYyAp31J*Qha zMo6)`r>%e%Jb9+gOJ@^vS$Rd9^xCIv>XG5U0y=+~wRj@3dUU3J(l3UO|C*MTl-=Mg zZ~gLm>Z!8;9}cDcV1l#b)ET+p>rcHduV`uSTS-i~aoQBTe)DFUI$Iy}F~qR0dL7E1 z<5LQxJ->g5}i^;JEaHu?2%t{ z)__14>U>7hvDW3B{;K-AlzHD@nuP=*CJO7Pr@Ht?TbEm}^?wc@J9Z68Q1ZSKpYtC) zCo4-#%fwe_(%O5Yzza}E5T?qP8KbUz300auC2Obg^H)H%MBSKg5lCmD9|yxRcmyWc zU)CKsxbX9N{Gs2ayfFSxp(UiNyza&&t^oWth?4gbXMUrNf4&s|sbt|lWg%Oy#z0zH z#!J}?fjeXtGe4Q12R=UI9(RyZ&z;R-SxyS8O>tgBcre3yI)~&YTd8fSFY{M=6t2E85yh(io zUD!%XDG`XQ{S6G03i{yfz6701|10@ERK;>FZqG2pP-hjZh$YpWIpU?gXpPUDn;!fB zQYm^u=9y2w+Wu{bV>W6>YrV($9Zk71CN7dzn&n}_(A^XMsRCRlP`AE1hu(G;aS`x; zrgvWY(Y?^Hcpo`dx(Q2s>1gny1}j6(nUlpa3c6=@;_DL|_2rSJx%ecC(weJT{#i){g$tMi-}{?P*7>iZA!T{_5vNAUfw^#k~FV#y6b84?$9skR`;o|O8; zMgMymq#4U~)e)LqMdXC|##eF5r1d4o=)L54I3p!1CCI@VKHwLc>T-#WGt1OAXyMtE z4Qv2%>OMyq!)^Fb_mg#aJ?$+a%TB+^MT|dAzo)i4Ba|=nGRET;%!-`7!mjJzX(jL zZhssK=7%8%8TXc_?2v3vPAZ4RXSS8yo(;QZHpyLDjFHTvnj;CN7}9*)Om+Z; zfZo-oI%r#)k4*F`J8NHbDqg3%nPJ7}TAHXRV{8hx`&-kK?D$Bq3RcqBH){!ecyz^R}$bU|psl^iE!yiiFR!Jm)_0$I> z{~F$haJfmoIS}Wax6~E-RIzIP*_&=+E{eu z;)<7|799}F3d=H-xi__aegX3&H0Fa3s}kR%FSnj#XQf)KxlX;ydtcgq?eR<9BENIo zEht*(GciZ7JJ{K~UXntLNtFKU5)G0^`~A)#Sle#?^WJeq~^Kx4w6*BeD;$PYJT%L1|p3k;fh^=31 zy0@BizR~R=PGEnU+rHOs`K}HJ!)$Xm=}@a<=ji(OD!ummSSIVvwy?R`i+M#!cGa$UeCYAFA! zL2`e^Oh-BArg_hk!s~;3oHe<9s%rgwz0FPWw_EROD=u{?zTVxRZzir7q?LK2%Ch)r z_m#6fWm__LvgcB)-9yJp_J{XV%#Ye1+UPaZY+=3m(Dfz>%10xryxZ0#IGCqE^iG(Z zO;P4@O3V5opP(A)J2WkoGagxK%AWXIQiuKNS*S! zZMz8zYZxP^$RAO+HtNGg>92YX=XnOv(LlZDA`@4B6l$hRYj2;lkY| zl}}wf+GB_~xh2T=sOQ77twv1a(vt9H%9VL0fzKM}FVO`dk?TXO{>d6gO(+lurt;a4 z#wSh0O~USH?Nl;JO*o8e6?u6nJMulUl?v88UhADP8gIXvZJ1*B(hJ?0;w7!gh-n!W z=cY$yQw_miidv-+FTLE^Iau~t8oRIcc3&%arjHS*G0=a;nnNgc?Q#F%=<)Wk;}K_gP&-DB;P2!LhO-I+mGhhNI@YB)%z4NSNR^-aFHx<&+_FcYM7l*=))xTCBJg8 z)ynu@U0o{J8nZMWXqtp{$hHy8Xh;cXm4IG87$SV9Q_PhX<%n7HPG5JeL~W(cpL~ja z$1U;$jh;jf4d+sexU*u({B=@-&drhJq?GzzA}b1*=Y{OcRTMfalWD&_$vOBtm7(mS zzW$W~8n<@>onmLWT+yvLw|yxfBe!^@Hc9aD1Hx+#?X8Xa^OM9r%_LhQ;|W;`?_M1) zLuNC4pDT^B5hDY>KSTS0F6pdb6zIj7Prl^k7+KfC0MBq*5|?# z%jbd{H53^7a$FpztN#WW@g^=G-OF~9vGpst<#$2GGb%Bk2S_gDrE;w@G`vR>rnEC9 zc8VqG6?9EKv8JF+K-T+;dgv$ppVGcMDypdM_rM^CbP7s|(u{m3u+wSQWyZ765G=bRz;Ojbrg z5q!@CXVy@a{&;Zc-5`6`o)ni%8|&CmsbJ*{3xgZH-s=`SJROm)rVJ7MoVO%f`IF}- zHhX>EU{0zQX#1c^L4}%e6&ppN@ufd|e6}8ReoVbbtc%KdqS$3G$e;BXMCkoc%j8u< zlyZ7;xF!3uve+^L=!IySG){UY~? zl@&K12PtVVr(gFg*%s?YfGD0lc(o(ze*B+ z;|@4Gr#8=hyZru6xzBiD>dV&G(r1p-oB?wrP59McoiL}nCwV^SS4jL_pAIhOo|0)U zJLkQNB+J~ERQL>AZCFl<1sC9jS`ds^_7ZLX`gI+>qkp&`LpBROCGju0S*#C-r_CBY zT|}S3<@Lzn2ph~;^55Cn{G+tIJ^TH!7nU}*HRDE{fa$2-c{}FS5qlMZ4Mv(a-7n2E z=pah6PfC*$q^xgRMy|2L*+NxF*|g(F$Y=XOxlEEB9MosEJpe!K2G;u=4(;$YcN*Xc zI@6Xo8oGC04&m5M-#1!4lJfrSw`jbjc6L5i_)KY|r3B!_X$tUZ=)RlO`?Yv6@@=(9 z55iZK{w!!bD@0%_IHYH@z`?-y>ieSZWSk}>U-n-5@hRU)4GUoHg|BR{^cz^&z(O1j zm+xiBFzDiGX=qIB53WISp*FyFxAl^8Mo&XJ98pox0XR$$3$P)&5XU_o&e8HIWuu~O z2#Hb&V6sZQkl*a>;Ck$Ad!hGGC>1NV(V$-pF$d`{Fsi!sZ*lIHGtpojaoK z2Qg>Xl_>XDSxUbSsC$!X33-0qKZzP6(HsRN%4&efmDAWr9La7QBq1EyVyXY7!<092 z2CFd)0Q3u&xB#zC2Un)Tp#;@;7yw&Kisn)U2EZp!lN$&5v)NZaF0R3V7CK@nASR$> zDKq`qZFkqHNeVjzJZOchE5J8q+aOC29-G(oDLPG3D=wr-M5{z|$;H$Tl#H9#tA?hS zf#E26;G%IMtty8H8y(%{#L*Pt$8(Hf+wL9fjja?sp5sT+VY?F3S*|?CnC=SI0g)P7 zKDI=`oIVB~+z_NL@+TfLxFI*kCUiSSHHEmjPMVUSF5p?v@J)iuJK_eE?<%PRx}w`v zqbW~XqV}_$%Yg^B%R-upAid`R{9pIfm<|Nc9roX&^gj?F2bA%~9SI3m+g)?vo2r{7 zo#E+18KVfnguo;!1c5~%&^MNAS-9=(Q9Tn&3-kIqPck4%kDwVT)6>Vjg;30=t|RVT zs@yx5JHVH#HA)H`*o(Jc>Sb1_?&FiMx+vi*>t#rqrSF`Mh5`jE_$0yaw76r2to3I9^m8 z3lX*pCv>1n0^sRvMW+tMcr(Ta2V;T&&60yK7y_^UN3}-%8u820JaJUS%5uPJHtp6c zve^&v6GyMP*;#PzK=G;=P%UW*xNSjC(6aY7wqOzsDzgXk&v%Jo^B$q9Mfv~rlFwgfy7%>|g4N+`@d(@zhn zoa12nyhFJL(&>?(9wCA(uigS?1H7!fNG+2IInv*&-R{3rI}a`6R4}- zP120<`MbYJMZ8dhG#n2IEzqV@RSW|B7jZ5lnxa>L922Sw9;T)`q?JBuLSJ}`za;qM zKxBy{kYIi<7HL_5;wl;f3b%hDfaauggNf51_&!?V1?l_1d6pYIml7qUY_gm|e;;>L zdo{tm-fFP2q=10M0Gzpi!>(~jh^yz%vk4xj7l$yxj09TX^Z_yL3!WnJoZy^lD&I_0 zDndVIJXPCpaLsj%VOVDRL7O@rieG$6dk`bA!s@B3r zW6k>O;Z`H9-dPpOVj!NIwXn~V6WBOY0!)!Z^)d+GSuXn%8-Us!A{6G4^i?IHj zRbDRMO&AnbW;6mqSc(Aog#dK~ZzA7&e|C-wT}G_i5<^8dm*z;^bx~8VxWneAFek(o zOxNRJTTGy?syranVs61jAxlXpj*87k0w%k~cddTBCK@gD%_EJpx$lafx*!J&OwMki+6@_Lb<=|vPRfnh6}e*?L>&dHiKGPDsSqoBgHHJ71Ft7<*&H% z(W!y&QTNRzSxBC<^Aq;z8SXrFUz2Rl7HXnn@Uuw1VSOpqtels)*tKXR^g1@6G)$wY zC)!V#A+bJq!gDdaYBaM~)avO8T~kRoQ~(uE3!3?8X023+RM=(IDt&yu4jD>gmI|QN zyUlK)a#726>Gr$E$m`UzXL!iyR)K8(yG0#oiC($JDVOfn{Iwamga(_wonSV|1D*N8HxD(K z&8246$40tjj^gbNU)o{#D$TwziX?=y0-tsrs~)EQ_J?gOV&pURt@a^3_{8_FdOEIp zn#;0+lGW2!O?yJ}m?hh1)$)7|d;F{&$2RovG{(D-Peyb=yu?f0fH5z!<-_<2E&_#} zsCbUKXeq0PFZPjy4wOm2CmC&^@!V=IB@ifDv{g5Z&*>nSR zI2DeYr1#ie%39zw%gR%LiO;@Qpb3yBwp4|-&moMFb!$#7jpA&}LrqmEU+W)}>GXvTCq+Ocg%;DbRF0GD zc?4HAvN7eJ!+Id*|DXELmBP;{gD8L>ts`V7MLvN^#SQd~OicCEgwCa0`FhP9T!(BW zTXKdb73vpDw*8@e62+oX`_h6o9Gr%Yb$G%4tfMfhT;G$HRYl_pljBKrWxOzj5~B53 zgsQnR*D%4`C=Nd@^A-0^ZakBZZ|<1(!7nE{*@1ZYR=qeUQSXPhKvrU(7>)dCP3{|a zNWt}$v!hn_dRUW36Cana@4(%72v8&=4m{)_sFo}MI1|m*%*Le((sO^6SmCQoa}3A| zeRnb*bKJ@b?|bisy7pkRv}7~eHKoE2BZx2B{4oB}pAHa~Q2YT|6TN--MfYl*y)MA( zhkt2 z_Z@r}grZ`#ehyui!ciW3qtF!&y{o43aDu+!Y~W{-H`r6_v)l;bKuk35Q6YQJ>U6^sC-FU!iOcg&;#>Kdz=%C!>8=L`Qa*=0dIwg-Icj)r$7#NVrN zqtTQw(>}aDZ${O}x~AA={(ICYWRp%P53YtOtNU~(8jVpxS`ut4vjFWH729B7pU zOAXOa5N%Dh%C~blxs0C@{@E;i9j@#8hr%e?GS%HCT-L=8AL!mjkjm=zl%fvsl%Zgh zpr(OZ8T%w!^b1UgGSIQ31svy#2ghuGeM{fiBt6t`%pcqw&%Wl5;T6f<4~U z7-zj%3|AL%MVt-YzDxX(*?{tp@J ziIZ#RF<;%73Sj!H2Z-t+8k)5J85QBh{Hfgi9Iq_$W?ympsR^QHT#?fJfnLpYC|a7P z7GBfxEP0==7EL&tR`r2V4le+;v(GvGJ(^m|co{$Yvtpq{0q+YUPH`iuNYqgS68P)uV$fveoUXgvT^6EMmCzT|7 zyZK=)tSYK7^rlx2$aQKdmjO!lDYh9zcM{>?vtgIj`rOFZKET5kDDe-8*DX@+ti`V) z0t&-F6m>f|jed!9E2ReB+Sz~8r*zS95gL34Fc|gT(M;Cf@oEuW=W`qykJXmZSjFzr zy#%O#4%GIxkp50c1T$m>(#L57U$3u8+~M-?y2rfoTr2;_WTL2DI4Qx&QfgBf0G6A9 z58D~ElrSR1I)4>M8yL{Ej37C-jf%o(s~Ta5%Jo*T?Rozh|m4>ad--H1>npv z9YdX|QljC86nnkQJ>Y~iNQN6AhfyYD)7iEL=8vITuJ1C=8wsC_t>fh3v9#;DH{`J5 zO0%uZD~n^}I<<{1G=*UuYdr1@9ecjNX^n?vGK$Mm6*B6R zzflKPc=wHJ)W2t(ZqIQ1-_q4_rdWz30@(kcwQ!jIEC^kSugJOQz|LY=y0GN)^Ii6#>Q#~Cz{PK9k z=I5&9CrVM%I&tHJ1mm}N)UhN6lBbXTT5FJ9KPbRIjoac^Y@`V2zhrYm^XUP9uVPLt zfqL6$GJL?;4$OY0DgTo>#w0C!WyimWT~{ON3ILVPXV7X;EPl=r*HM-ToNjzjoyjcX z{j>Rd(DC)`yO7S#vI2F@_@M~r=dCi1!=y)^yFH4rzwVWm6yI4Y|&w4iwH_ioPYX+92l-*ap@TY_h3Z^2YQ?mk^Icm0Ie zNK%b*ry7RWDy(9Xg@8&;(H=peZC zz!}3?eYZu%Tyksr0YiIF#q-t!;tghr;*v0KHeEo~Lt5&Ej`DH0zKD)n6Mliwx(*H{N@9N7SP*YPgZ8kx6?^!i&!j#G;XDj zni&7Rt*of1rW5$XUz%RxYQPfIuoaFEnZ(kn7V;bFbOlAiHi;4p+L%SaAtdL{d$$YF z3x(~aX^lQmto~i@ooJEE`1W{17T-6u^SVbzUBh?tiR(qhnzv##H8f^-%p?BPC9mEr z{4B3frlaNC$zH@#)g@CDoey_TwCg3>g3J9b;=I4K7V2KwS@Kh7-i$P5tD^MN?r?!< ztX0g}!0K`6A1A>KF3zmBb<~GpTS;}1=W=4pXR};U&9pR7t;xv1&{Ew3epdHzg8ns^ z!DLGn(rakZA@zpLZVcUJt66O&HrSu=^|yb$NB@~~GR3NzG=kG4(STG{QPHKwM&?sx z`4{8Ylo6p8t7gZ7G$tkHzPiacY6}`(U%$6#1rR10Y-WA_;^Gc$=GMRZt%;6^S?FZ@ zo&M`YFy^PNZ1aJKIr=LHXI7BXXB_Y^fBl~y)AE-C`ESbNQcM9aVT_ZJFxX=y>{p6K zT7kF!vpL+Kwewg}agQsm^RXZ{eETlgUFqptVWFk)61c#Sm5z=l8gh%%(BW-}cZz|) z(3AcY9B6uO6U0ST378YkmR@lhExASoS{jp8L)TIRc{MRIBm3Eft5A;N1_Sv9`hTtG zU}iLPid~Z3Z*gW}cPLd9U2<$t$=ESDxkgFw$|(1rwX8x1+Q0j=T&o>eyklRiB!}j+ z{p*HH$^DTo@$iuUInv8>3!3L3FaH;U-ZDFb$}GzSGav6P%%p*QE&$O}ZEbDSaJ>IZ zLw^b_^!Da_ywWjlvD8o(5VN!-&18@D3XKW!Yapawku&*8uRDjLA{iwvMp5x#``-YZ zZXxf|e^S5ASxFH&BOC<92!eos1_9{`(wil80qGr7dhZ}5prR7Gp-NXegwTuBph$00 zLkmTE0-=Tuc{fLY_sSjb9pf9{AMYMGl5B9VHrJf#N_KFBsWPEArG})Ca++;+_i?r zS~si?GP^S2f1Ud)_$|Yu_!s}uA!sY6tR|^05%Y=HYB`mI#?Yg<@>yzWWc6FLtsIHJSn$mVb z4IwQPLU|>O=YT%_=d%7A+v51255HcW3vX(Wp1CC&p+JA;#JW6< zVs#WXHRJp@De01X&1+qgO0&hnqc@K%%+WnerhcA@2FSunZDO^zuZF9Vg7c@yz7f7` zBr0tx&Ow6VDE3^kh}?v83XI2O$LxRho#A9!xx(Ohs+}rVD}OUTpN!dOW%{#wmgDBu z;Sxr~pdze}umI}RE$r#;^pbJibDsR@vRf;&gJmXSL)CB|2Ua;!KC8usF)y^V5I9m3 zKCt)xh1ssUL_p3!#)ygMP%p29&5?=LMq+n)_*7qR_26C2` zo`q=SRfLeBpfcoO8*8$Ady4!)vrWiMBbdefdMA&A4|mX3#qCq3 zkhmfvO&lg+t!KH9m#^Dwx6~jj%jyo@ty^EgSE;Fye#wW+bQKXY?sue0wUqu_$2R_OY@&;h=eH25gCIKKFIcMjpXGodg+cW<4}p+-@86!TLHVF zR3vpju2|yUbFZI@boB?D$}ypHF85?B)5~C)dDX$ zO6DeQxZeGUxop?IN(U`RlLl8Qx1AtCtAJc1HJ|Cp zr1zILeYgP6)b-^Mw;H3y-&(oOS3=YYjX(m6cRe2$QWQ z-U9EUxD1K7_Z(@^YRl9o5PRVy%*++|<&sr~;?BXr$KZ60y!3Q~A@OJJ=&fO%!-f4E z7CKN)XC93QX+8RGg#j=4a&6f;-}8UY)~hYKRWeh)6Mxxm$%OQpv7MDLbTn0Vs?D@h5n*R&|GSZJ1$uAGL| z|Gs({E95(RzUOSAXyKn48|?~f>OL=&kJ>W|JA&Y;Qa=qVt??dS*bFLeI>wWcJKzf! z-YK2fw@uqg1rb&Ulo5Tfwv+Rmmxh42{Ku87#vrY+a~z`Gtb>3_?oS1^M$pfD+dV)@ z>zcZ3EUO%jC6%S9<|t(y_~C0YH7(V+T@18!K7N}HJ{J=>S{)Sov*DNwGKcNtUmlV& z$$R@Znbh%jb5LGE!9=A)D}5cyAB~ZaBv+Lazh@4l`4YNyQ~Pk+yJ58s6C-_DSy`8Z z2Fy1_mI&aVuGKn4hXk!fo)Uv8?|`06GWMnNno*|)JriwheUq8RMMRfP+4#xfAO+EG zx*1K^otB}4nDQa)AF;|-=I1YTv&QgtSk}w5;@yR}8P6o_!0{v<$c_Pw%ZXQ`oGmtNPMeDR_tr+(7!+`ROMCD_7ew zm|+W@eh7w@;^M`%nzfm$2<&)?I#f&JhB-skxcS)F!{{jGf?;Z@+#y?E!qRY+ zTCWZ-tFm(BDkg1vfS)mhr@12Hh99PJ?@Q0~=eL6T1x1`jOJ;giJj!Xtcy1SJ$|qqe zXL`Fi7zXkl+e7s994BWo=d`l=OJ%okMfLd(Nq3^#V-gZliX~jS4@@OlTu>SF){5Wo z-^CeVk+>UIX>Nm2$JOnQRlbwhA*qv}p$E;7Dqc>p8mEAcbT(xK!olOog)?($a%v_F z7BiRZ;#1t8z_VVCbbI}M;cuQpPyg;(&mglQY}%yHp;Ya_vi{l?->*Q1xb!OH&qIj|7AGDC#5?Hw|! zDykqVRz0{;LQ4c(0A?L=8J+Vrw(BY! zA#ax5YdpAV^M-6e^wK^?GjMR0qrg7uq+c*aqI~y;k%=~81;|1t>MDlwNRr$BH=F&` zce#N_Qo}v7?dM4bO(!B>G8Ax-k}RAh2C+WNC1!l_fa{{HKEdv%)vbvw)F6uX_Anw) zKR+WrT~AfKQ-I@Sl6#|NRo)P`wk5)QdU68~_o>I%>lb}IIyUA*@4S`j)J3XD#PEUd z^*Ky=<79iq!hMjiTy0y>{&RwO&HyKvHE8N7&bPVMLm2Qso848)vgVtQca6PAWw~PW z>*`)?Ri?-s@5md-WyAg^BdeF#x)xw+TGXn|HPpsvm!=PSf2mrvJ=)tlU;&|&%EK*T z)f|ZvcEHTIi$gEHk55garJ>PVB%a-R>oVb<>+|bV3U3PHJ8(>J^_z&&OiPIAmrf&v)4i^d3(m zVe7@}O+o&<#~n@k`{2#COtNk9dYGi}YIXOx=N(oS|Mf|b2xW*3_0Zb4$qM|MMb-F# zpg_iW$Udo;@of#q{ePp8g;Y7N#SxCun9t1J6zal zQ*GGpVI#x6~W1aePwUoY@E}s8g&4>;+8-{==%L^Mr|(b3gj4KQlPXNmYHqk zGZY*05N*xacV=R(T^U3*hBM@k);|X_c=m|TEi8n>zCL>@eQ>xR+}0n9uw|-quzKXJ z9gFeztGb(_!vU_Yt~P1niOT@L5Js1(eN4~KS5nL|c2zemGOU_#&wy!Z>#iv%(7@{t zqhk&NiPhtFb3~{Bv4+vH`lRKa3BkB_C#0>Bn-5uyC!8<30%s!kt-}I9%cQ|}wKdHT zuZ{ZTUkpbE!-sg~2x;AC%uCA;Pp7xh_wRK>$qlqPcg zczC5OOK3vd*pB)AoQAp2G2aua9q=xwt6RWZ0DzhY-2LcHQc#fZaML$n8Q;O$zR4=r zG4CIBzALI1lHgo+>MM!kO(eTO8q(H_Ix%VeA9j#K8XCjtEs*h$A50~4cZMgkXO{wqvy^6$xd)*ZwASM%3*qAC4meqj9Pp7n zW0XX>OVpfhi|xWk)fOkZp(mfU^qrlb>gibnKk*Svjx18HGe=^l*_CC(;cFNKa${;3 z@nz4tsnAdp$P#AiBtg~K4r0nvv7arsz%Rb+Oq4sKIgES7yyWs-Th zSn_FaL7!0PTElNW__&L{3Re+vua0lA19 zU#fTtE!ekdN~wyl9j}SiK~#*df?fC7&O>gytoawgflSwHa_ec%+8Cy0k%zPhS^er84F#eJbL%-*F7ww^NgRkE7>N(%0q*7E_WYoBAQS9r9DJEF+;jW zBI~Z8egEcsCi37*2mOC`oTSIiszzBKNa|r-DZ9;e{3x6C`i^8(s|sYzn2xcDc5J+9 zu}5p+INpCcx|i;b*T3+nZ2@$t;ssN~v13B%#)!xHA{r2oZs)(VEqS5Z9Gh zKv#=#^cs=P)NCCwAgdEMJ4m#c;yHK`Vb#~$+iUkx4si_x?CJh%0s90mP;huSEgayK zfDqB1ry5BHL<{Mi;UwS%Pk83C(YI(-#|O5p5a-H|tHX#)`LDTFhsl}s;*EVhB1IU& z`QmZoS)Dj*IiwB{Qh7}Szv7G4UQWY9LOP;%sLsr_Nx9XYAAo%TVk2z{KhfJD*|e|L z4(jZK{jhvYYI@i=(Q7y9faWuR2$Z_UMqV2l0x$#wl8Y{0x9o;ww!IKrcAere@#$1) zDrZ=|7H>o|RFidT7aeYm#f+rJ6Gf2~@e!t-*d(cKItrKXWFl-p$vGMfyCweA>gd~jZt3$rgG9g z@_^^HwYN|PnKbc8BVI7THqJV-AhiZ<30=CWx`-k0=Fqn_Y}wBuloJhn`qKlDqKd?v zhtWb&fIST;R*A5&B4Pj_=T2n*Wuiubdh&bF~61@X{XzBJO z=?CFD)h}4FYg_qJRxObr#*(3~JN)>s6L#gGa7M;P2)26Bb@%^d9cTe8BV$1|H zTkF0YlN4R+hE8qTuSnzP1mF7l5bEO@X;1Gx{a~{)rIw3}tDe+vZ^tf0!b7iOVKLq$ zO@C`+{fFhlhv(*FwG$FXIqc<-4L5h*sFtigdi+>GuxyuEq0sE^tCosq_;IknLZMH(|h)a3^EvhmfVO2?^kk3&UyTKy1Rp^vqva6SWqE^NR z0K)*E=2~Tl{xaS8s>aG$lnHgx&{yV)>lUG>Jg7g65_TGQ=_f60_}QNpaM_hK(F+1*-hNCu4IEBaCxfTn3_UGGS_ZN)JGS-ksLS#VqP@uA5` zj{LYVlHy{BCmY_l#5TOvt*re93)-bAEveVs|I9=!Qfu2Ey=uYF&`|zgLNH${lyV-K&iN2x@rJGzKNR9HIbuIU*rwDm_#+ zfb&@E_GzA510rW7$BKS~t;5)7cTjM!<6(~>e`-rSE~UyuYtuY$nI3@MnAGCddDn!a z=X~xI8M)ZFhBd>gqKQRmgJxN zgDX`g%e3g_^_jYRvv`@lF1j7;LPLYaAPha3XpAc!b*g8WUNco!ci9bK=SIlv4$k_d zOwPMvL^evGitQc6LHJtVLD*-|=kCJVgAi@IfP5l_w7~K?i-3xMJy*ci4A~l!OVg*1 z!s^F-EA=j0!w0khKSFm`6Pu|9#M`LkL_l;1Zj)H+&C4&Oh1-#txG0qwO8{wZUZpnbc(N7$Gj3Hv z^Ky1b+q(clrCo0WJ%29^T=~HCCE;3pYX8-<-IfWb{Xq@^g z#*WkfV0G8ncPF%MuOvgwhWtz)7t>mTV{ANX=m*lS3;DGbpsfZuv}pld$*AjiJJf*e zhIovLiMGYv>uMF$^MjqH{>M;L4Oa_`ELd2>K6}wUfa1QecNp46#sbCi0gB!=nLqH^ zTL`GVyR%3wziZwV2FqnEHQddMr=a6?p`)hW4Q`V?_VJAvwDADDX0~bUzU`GchJ>y= zgS@e$@oMf57R5R>{InoPeqm>`2$EMV?mJ)$R>$jeu^gH&vpbk;*6;>t-JRStAXBr) zq^WXKei<+tLnM$%;eS>s>%L5Z+6X6TbH&1FXlRD-jemqU@uMoi=;GF0~3Eo8{{8jwt9X(Luut<0*<8n&Y61O!T_pM6quHlrYoB<#ARBqnb}en%^HcEn_{etTrslQFO??L$;a(b3g5mU8^|U9jCuTTAzl zff59oM7$NTV+YI9c5-&>Xlc=a#&`-6yeJDja5^0ZCWIfBfJ|$=CN=qIw@8xX2g5NB z>PPG=FELiZ6ndFAcn$CG)<_Wgf?`821=trXzODClr-pl=>_=k1pmDvvECE3Djw$1h zEOM+0R!~rI?)H}UI~=-MHoELCz1LR$Gm1JjOrK>qgALqYq{iQy9vN9;Vg43oZZV(S zZ=aA*4Zb>>UJZ=8Y?AoR1YaS^>02BdAD{O1tjcY@nJAa}1Y+XrRynY&m8X@~3k9+_ zQLp2M+O;o7Q71q!alI?C>kQ4t_Ll@kMn=MTjJBcOwW>K=4>}%#s&h@9vNUwmIl+An zhJXZd1P?A*<+}pcok~uLi7{#o0fI(vscf1I;d_jD3y)#70pwfLqgITxL0;|x5c4zOV*N3!+jAp}|? zMHNKv()0|pZ#_u|ezCdhgejFp&JJ-*J2RQ0Pa|!V8HbpTYZue*s)X~?Q-=WLHA<-~ zhu09A0_Zx45L+EYNw>Jo6KQ4UhIrdN}l>x zzN!tQo>UOUg|k&{kLXW-kLOIRJjfX-Zm{L&Ba|c z0%Xs;EAnu;>L852l%~kwQ1fxOj=Fj^ApZe)!!P|@(|aQpYK#s<`o!EH%6=xGI5XQR zk(mmoJJIrEYJl)uw>>dfm?DwdU)G&?A+(LyeYolDf0E$UUtDJF=|XI$QoRBPGK=x` z#caW}D}OqO2hqc~sT%BQ3@Q%Bi=;1+Uj|aTK)DGYr5uM!`#wov13vcXe_GaTaz$`A zFaaFpdFMESoSHFRXxgvwYy2Ninw|i($YEG_U9Mc8JsHt zroBSMHC548Z5^G3#YM8e z{|2hSggyBi6>muSlxJq%mXw^ZzW?tHb8~ZmVp(5b2bvTmpEu-|A^)V<*4tZDP!JXp zV!I;uDtbDe=l-pusCVpdO-p}htt-j zYUYuZ0TdysKH4AdQ2WdJYDjXNl2ulX1{%nCIKJ}hTR+-kqfq$=W&SzubQt7El&{#| zrk(qtHMk-8fZHro&K5i2+92B@ARv&RpFdV^19&AVoP5NM5-d-v`c*SM2YP)H;} z|01fg%=1{U(7vK}6=l*EuKij#L}SNr+WGty;h-j)Vml$*lbWAKx+dCLc{+?8?NZZ5 zyC^5*(d}1mshj8_k=7soAI8yG4>wbYO$maRmYB81#m3eXPY$66+ZzSFsZzf49kJle zR+fP4l*GfFC-I$GlAG-uY3{3IgaX3;@$OXb$9qj>SDBelHaq!!3B|s<_&m)Vz-D}D z(|oXb&zs(-AXQ%n!$%n9C$_<6F+FNZ`N zCvuHerchI5scC7ZIy&s&l9CcDWV;7$nA5v7l&{7bSQO3dC&3XBJ-xjwGX8#!cw*03 zXLmOf@ZR6w-`#yXhEp?JCCRf-mM*w0j9zdUY3MavbB6g@UvmK#ImWnQXLWp0U(#={ zxcP;8A@$X(l1@yLFP}oh-It!-Pt7vRK_C!TJ;{JnEf@At%xz(Prl}QpqSIvEWhQZ0 zTM&!?5_6tWl^xn~vd-JAC5&F$Z?7vsP(JMXEwG?fPeEa!w8W(|Y1`he=oHu@MchN9 zCv=;Kot-^I!QH(awlYyYu%+YbS_*T=4X+k>e1yY4-!S*?mD&ZU->fAZ&KGzKYon4m znr7a|6`?}f-jY-ML7}lBN?D0_pT8-sHt_K|(8ryrylO7bGmED_vKGgoK1Vdi02gHq;fIqw93*)-82)b$WXG(a}+>tzy+PE`vzA zXr0N!%4!Cf>+Vzl%9n*-O*xK-oQi5~3Kh_8Xj!JOXvD_Bv5UtTaynvav>|+T>)*b7 z`SShy_qn;b2qp=xs+jTwi5@m^TENM^#9G~~nAOXddj}WcW%Anl(vmXx^*Ht(S_D@mI<42ET zC_kyFBv`tlpM_728ts{UxRQ#{4- z>zfM>W951KFRL6efUCKDL-Z4*qd>>bz`y`-z(BUSHd#+iP3^TdiG*~7gog6ch61|* zp@WREA|7q20qE6n;tN3i86tODN~n(_@z z^J_gCIk9%w8;j^k`B~pjk`5jlY2lb(1r?3@u6%MGVREy9Ko3^RdihN1^dOR&7D}lA zro*^)@>VRgLvHi&IWG5Sbp*$OORajDK|swzQIs$>9PpzUzpZYt>-zM=tfdGyq<)zk zgb#JJNyhy0LEhV-ph!j$M_>-Ty}j`%Ei@vL41|#YqG8$VduA<95(H$cnZ;b^_BLkS zTwUXM^j%w-yzWoX{bnRYt0*b!ic_4%SMr;bL;XGI4wHV?CW(t16-H78%*Mqb^zv=`VGwnbisU0rKD-J#~8-9k&NEniE+M>Dc3Qpk8wvS*~tN73cr* zCAB_(pG)OKj}cKpUj#LB?wl0%`Uk^UROZs6`e$AFL&&lDzFTJrG>7opn29lN_cYkS zXk)Jw*vqVr3GY*ftePyk0qNzE_`r7i_8k^rkulFp`HX8|ZE?JY$*MU9NL58eb+9$w zO2PwcE}T@!q>B+Se>^{6_rS2#^I znt)4(H|3G=Qp6iF_l>pd7WR3NjxL5N#ACk8KN2IwIx3kl2_^}~#EP1|D$QYW_!}&% za)^I}z|zkL#;=s#{x0GdMJH##OB3@-M@B|;O}ytbfszjOh|gL*-IV|QGg9f01Oq~i zHUw}awVtaQ8XElr0~r|^RDk_|Z=2|iU=UKOb{0Lw=mYNe-7 zKbRC1RoAqe_gmjjjRgXM0t9d+L$(CT!{2n7XxfeZj%p()+=TAMQyYYlXV%DYPKCO| z6}a@d4ZMQB6>&$dV&-E;*Nkgw%H;iYov|MXJ0e}*cQeKAOQ7J^sF=uAX!>8}mY{v? zdX$s}S!YTt!PJge638Cje!LMfeGVL@ks7D%t(QkL8Uis_8+{$&Pj{PH!r9X4S8eT& z^&X6Zez#TvEcw0~JJ;RfDBVdQN7JGGt@7i6#U)T5dm;SF6Evk&g)ANfUw6PHl4%8hQklPTuoxba_ zwJFChFsvC+gW$HA{JIvq;qhKJLrw0|iw?~ltqPC7z(kFEE)BocluC4Sk~L8E_a6qijKS}iHVI(*}47Ow}=(C>1fOs2;^ScIAMxa(5E6^=2DN9;PMlahQD9H z!@IX{flE;g&1Jr7k@8QRq?|X6%k~LzFjHQw%N3gQ|`j}eyrspmU<~ZP|N9T*jb_l zk^eX_ru`wI!Ggac@2L$WVBz@g5lL-v2XEE209Brf`1#Qjov{0FNpWu?Mmq~4TtUk# z1KC+6vP#}Q)n%nsVtf#~4;rb}NO!f^&PfB7;if^cr&c|Al65CbIyqk^qH|?EES{|i zx5W8W+(OkjKW^VSLTUUyp^&p`6d6;kX+cKC?6>qe2)nNk&u0uVC!QSJRu}}stDTyg zn}ge0!WlI+HDxmad%_z6K^8c0P-fYUokF2D3t;w>b^A}i7OFX#-bU?(oC7+#y1th* zFORxYITzubMJLb~=k<$bg?=H%S1t22x(^9(ot8GesY<1+q`Xkr3coVmF+|Xo5>zuS znbe#+yl>+tYi`%>Xfg5_Ue$6Snzkhcc>`03V67RFpV9*!>6k zQ4wcq$v>qNmq}B5P(>y3mejh7AvQZwtf5Nmq-o-vaQy83L=-sJNyN?QGtzqT^W4k~ zT4dltC`?I9A9l8&%{1&&=Qpb*I4mg?S+I+AeU#|YkE!9H{Nq>AMB()pew zj)D{N1Rj5CS;Nz2j2%?&=pN>2W0mABsN5!Wy*sHpsZ#BKt;DOE`>z4^Ra+;ZLq85U z0Ht)mB`rM(<-ULaJ`kcgIIMuQ2fv;G=Hyfy?O#a`<=ad~mX>ZlN$l;EXm}m)^5qxU z+jsAR6&xK4(Zz4x{N)vdC(^*beS2h4cE1WphIZ|Ln|(EF^~zIjGxN7Idnls{f(v|RAI79zXA-O!ryt7z;zMJga<|`9A((Svs1eGI-NxWcpoN>n ziZ~fJ2U+mdLMowL!mXN{wViO|Si@;`${&eLe63$Y9-lLCvokX<=G#-Km`8kLQOnr6 z0#(SBh9EUU735^&6lblT0<~MUo5kH0nxWVpS`p*v@!^GXC$7;#kO!*K_Br}zBfrxmwID$E_u}XmXWUnZ90~c8#_D8fkm3H4HTNAod3;SD+>aU5htwpF zg(?Jh8?2UxDtKvf1R4Yk>TKHdq_e%VTUhYp1+yaqtWei`6Sz~f%LEqv61TeLaC}~U z3kCihc6hoV6cLqNq`}>P$@X~u>v{KlK@289{cApwjU4>bT@2M>T6s61+2Wj(ncrA0 zgrL52)OT`L-N5SBJaotIrb$z(%8LY0nis=rAugw=Xs? zXM)Uv{q}9Yol&cQAq7_~01CK+oEr8QlIG{gtZdMJBJndr~CiWo1(ibbkZ1J+f7#f?KdMQp0Q8X1tF>DWgATv{PxtLPB{L zrVyjX4XSR;LBcY1Ex98LL^nHe%Pe*A1broK&vyu}EP;U;SD1pbnOkoHbwRLUWz8!86b@*6Vd-%aXfT{I%m`lazR?GJ8Ii0?)BQE{|l+%5#Rl7)RnGmD%ae zfh%-$Kui_ywNF^&-F;_L>bTT%HORK|7U=>U9%PyjsxT;? zoeZIP?wvc1Jc`$;T{4uk?nxsxPdX4nG+9;w;ZuXTzOf+-MfE?%uCr=#yD`{)=CbmL z?JP9X73Y@C{H71MKq>AHUCyx)0|;Qq~01(Q}Zl$rW^ z^)0)Gz#_*+raL)}gt$~AaX`uq%Z1F3M)k)oghihGB1MU=iw>CuK$m2W)&pVp?S=~$ zFz+bn#_WXy0s@MB0H3$em0(t)ck$vygDS^W91d|KB`^75-qLv0TBXA{R`1{dU=~ge za1<02p`oGUl@5r3W!s45!MuSeHs#IDO@ILas`lL7-GMB&`+0YwP?D5yW%CO~Wn~1w z4n{;oz+le;exF%MF58J^^V+tfd;SkP&mCL}RA1*UYB4!3GD7v81HzRw=L`7CnXZ9S zl-%^6$;cp0ZDkLB93LgAi8izH)6q$b(Xcrkof(#4?jXASHT6s<)~?CJ=1F+K78mQ$#~<2kzFG$ z=I#5_Yxv#+q)wG}U%G54Ek{cJ!-L;kde0kPwp#hOmp0LRK%qaE72&xdTw}_2<$=W1 z()Bc9u!gn7PnvjG`C!n#%`wEB05`2X@u10eUAKT-D;m2Y*{LU}0ouRC{Lx%yA4X!P{6tCQr<*l%+%w%_F^pNn6e``{9#NXj9` zfBV*z7mehP*S|G@;%g43f6;+M7Hj!V_gFoiyhTJKHai@xw7|aT*Wi)tZ!tNhNV!syC z_KB2U8U??2I+zqvB*;YyQjdibg;tCXSIF6eOfVi5vq~Uxn;F*6J?f}Dhg4FjeeM;v zZH2;HhxkqYlUBB>YYCfj_0?q?gg6#WzaoD_JavjrEk-- ziNn_{Kk_u6B>Pi+2Wql@eJ=n_|9cj2B5Wx-x|uo{^(A%aYRSGgVas(9LVYjPHo~?f z-iKJp|3g5W^WPPC6?%=$BE$0JXrI4fFTJmg^;ev3lwBQdK<^wT$<8_F@(P{2MAy7< zd(HzQGzYd%<5E)!6-)lrGJcgS0nyS3ZuTusp;BUkFq~p<@DO0y#>-oZ*Ahvk4J?!EZc|zCalp>$8Z$% zs1;3{c+s}BlQ+ZzMJ$$K9$~a6ulMR$KGMalX|eS&U>RrwilO0(LOb$OYf{OV#9DZy z=rI{7YU7kTrbowOL(I!BJ1mPSn~x59GiC~9YE4EY_dMTKl#k4b-jKvb_!u-9AD5`- zyZ+*C>9{RyF|juGC(mZ%?pIzW;qV2>Gss*{ocut_uhPSWj!(Ik&9}i#!=LVRl4iVk z0`qjG*NL+&$!AyVDr#?!_DK0QD;Qv--rY-R*yKQ$2?q|{%okx64c+RAf2d9o<{6j^ z1&Yk?`O%h?Kb0+EEOQ7MV%m++oBG8p^sF>som9(IRJ~}ITVo7<1sW1LP07y}sAr;Q z`yi5uu(ZvXPqfTB0iHV=Ryzm9&C~`{zFQ&Tpv2pf_fYaj%c{S^L>P;04e)Ide_mnN zy8iR63^$wXck7bu%2_U!XGwtm{i9bdu&Cc&orK=fxVO`CXpoWdi5K|SBWAC#S^R#O zz!Uj$ehlZYPhp1tOFKEmuHSM7QvPx__^L3%xNfMjAZe&t<91H&6RNO!w34Tc4O30R z>U^nYzZ}f^gZcsK#vb&u_4kA|1ut;pKYQiiz`JdaNS}Xr_vE1(3H@tn5`ll9B<3N4 z+vhq8s0U&OM;e4v*4EZQ3dP;OiqN>=mm!-WWd+0T zi+_cE!k`^1^g<~@fTg_q*YSW7@XXI5!;DIAJ19E-y+^9wc9XL7!^7+5w={v47wut#eCjtH0E}w{!2;Xehtkq)2%4kXGO?U^G__NZM$&&&g$G zOZ@hOzo)oS@ADI6tw4A;s&pgxUt7fCaMbYh^mKQQzrt?974C#Y;dx$x*37|99UUFp>?9;jOKcW?DR<-EJ-QX& z8Jm_Cu=#DWQBP1*bZO%yKuQjE_^n*Wjg5^J=m04Ju=0)nY^!lgoiqPCzVq6}ch^mx zub6#!cRB$2x4-}{1oD{w4GiGs_PGz7OXtI2@A%ISg5&9b%uR>KKc4eJlG2c(~=KJ*oKB~{%CFsUhP)Vnu%67AN>TeWw50&w2XX1=}4 z+n7395b2d6o*tm98G2E2do~#INM;@`&zt~gNf$cZn%-a9#dw+8eqVMkG^oJN-FWo* z#}x5o_Y$LU-SN&i7_fWi=XFAZF~At$-yx5dd(Md}m$-9~0F^?(mhZU=tk{Ror=QuR zk^gPez!v*$(>DcGItmrME?$4!VyHM;=e;HDH0ifK{q;GU1&|F6i0rIQ$p9I(`NX{I z{;@wOLEHm-LD_X6Tg69e^k;KuJdgg>i|>T(M}GA8y8_aJ3_wc*N6ZhNudc2FO4Hte znhf;o)iZLKGhA~ZrSZ;y)&5wOlc{Yn0@1nZi5m|O56`PjOiGe9KiFOd4##{mza9+1 zqlo?>{RY2NnKGd1lE?(0AVf+ZENw3j$j()vjvY4JR>!Mw1am+hJpHqGif4|FG&bAU z#m00ei)lhU*CtaK>}_qY!V?o2KsMIa3K);2AGbh?DOD>?!^6W1r>6uXB=ApZWT`Bz ztfpLAnZTRhDEMzjcJ=g_0FO=-vO76ilnK}=ORIlIf2Mroi*fYyjS7ljBQW!n3@p3J`RFEy<=F zm)Y~2mXg%azyQ$Wl$Dje0P+bCO#lhaqBD+d768`(Oavq=zIeM+Jv2vED5<-BAd*Ff z4{Tv!AuB7J7mjUq2KP@*nd$-G_5%PuK!sio+AsgTn|{4avxN0Uo$erjZ#m&M-B)sD zF46+zA}{e2T_cy~_FABEZf@@Ju`xjI1Ypl!QhNZf3T?B9H`XyU93LEX2PhPJLCdAu zZQC+)v3MAuuPvZaCmvX#6xBuS&?z2GC)!(Ef7$Ah0IVb3W(-g{Aejvf{-dQ9K1Ydm zg@8)#4bYxZIq|WvQhU>Zr}ol&asn9F-RIwmy9fdTAeUN9J5a>%$Mxky%L_&QJ2VGa zT~vPWeNR9er9^N6XZSJs+K4CY1AxYd$j?`mh63o%!NCEbRtJ?k>P(z~i3E`F^XD@F zos^iE2$%sfKIRf7uXjJ^{H5vz9|4If79}f`OGAA36 z`3EQBpOZwKhC7=Js}~X#quJSJi|Tg(SXr;5ufJas34Nj-W0eN2%*dFn9Ctj#nGyjj zPLk}hzHu_A$^T0^BQcuG=gkV;TAZg}N&T{W5AC+B`l}*)8Z1Do4CsJ3U=9f^$4l~R zGRNhxOr6+ef_X6Hr(4?L1ShT?l37{l34qUokl(7Bez`T=#DN_QpxwIV>?BBLb+tYf zJMbMN-t*T$)dyGvYNw)T88|pw?KFyE%X^`&j?yfVI z;PbuDTjx91b=Du7b*+ty`=0ZjbB-~7F)$>u<(RHbyw9Z}3lHbg6jMINy*=l-Ww7LU zoXbm`+L~fzuR1tbaXf4G@@)Pm`(-JiR0WBuGSXG*3&~=GVexbj0KA)KvkMHLnV6Xu zc*4WN{A)FGwcD`#U{&daC8jKlSG#Rv#xs!7uFtz`b*hcSAD1oAm^0MDDUG;mB~H#>Odn=$v=Pj4Ktr$ zLJ@`EE`uX$s{RZ5{=AibY#Swqxg-wHnUO2eC zZ6ncJaTYIL;1g6R+SKb7S0m4P`qLPZcuhZtJUxXMi8a+K5ER&wBb=Xtr7}b;qJvGz z72d8T{>g&rFwk~)|4xwHlcjCKF2|HuNo^_jdX5pL)6PxtXJ}FOFRvZ*-?(Yr5G}(- zjqNw{z1>H&kZ9h!gjb{Tbj@sVs=VUmtfj*FjCb}?k#Q8K2qF2>NHNjl+GSN#Mt}V> zhOMWhq{P+rs#~>~v~=RCQkH5J+&Ju%QW6qznI54LZbHb3^}Qx2&h)AK&SQpndurf@ zHciL=2_%SJBVXDY2Ha1HebbaOK&XS|_d?w;s48VOGLk{qCaSkEHMfa3o?_{>e#Pn7 zc2%M8Ri9nHYg^I%S5u$U0bpBYIq*Y?cuw5*5W* z*t|DCKOd@hV1D*^_hbUQ(hAD0$QrF&CG&2ngV(`8@G~tPp0z{sjAbRb4KucT*tUNj zk$=bC;=5KYUlSX)xXb_yDYYCQ!V2hz_+N9FBKd_(_z8@sUd0mJ%%mBC4meJQbw)@T)7GKSQkD4kVl z5$DBnj6BKbR0Lf=sz1D>Q_;@jVWO{6oTHs7R)0yfo;0HM46UWGjhuFx@Oob5O&*6vv%t^OI>%C?-WFCo^KF196wdG($8w5y21axo<4)Iju$7eFQuR*C%uu=a#a*IXtx_GtCLNItP|{_Y`sRuE=^9Ob zn?uL8oy1suF;|8Nwu+2uv!%_bee<@#26O2nF05`v9TQo}%`!!ma>t*gFKvVGJEr}T zRTUThi%K?_VY1RUQ$Tr2dNyfciDA>bSgRYO_w|C!nH5iA<~;F0G{dD=UWf z)$?YauLJ}<+|4X*M@v>FC6R`X?=zkqicVgg=|E z*N981k&)BGJBA5%q9qK!>sm>?Otaw|c^`WSIK%YwWMrE>ghY`g%HJ!B;+)vin(@N|u%+s5g+}&L! zr^Q^U*Ms~y)eW;7Votk^a=%B=Al_D~^CteFZ~FK`qUzpYyqd&63Ujg^$$h%i79lq> zwpIofhDi`Ivn{11V>iOvEgtvTF#ckGq%AKjt(j3(VgD_eW`V6R?QApo>AuePFsO6h zjFkZN(?*Vm;VoJa^}HS~H6DL~E>~)i?%o{1#(;artr{@_`}rLr0s>&Vt@X1EkBo3L zG6v#V&nE;U;}aA4R9rUw1*y@2)Zg9F;cwFiNzwKd7E@DGo;|{|aD7azm$vIMg6(QJ zFCE^!?tdDw{VeSd6`Yq5;zrRVhYC^}` z`smJTA;BD^+l`*FN71>%l>);k-QX^uzPdyZB^{!HwDl`1ErYSn3p-j1rt0s zjOzw?%(;)huQSQN$V~~tFl}G;;n~rdEjhKc^A7qY4gQI)a}eoi3C z^R~|J$ zqRVuj!z_@Kk_grO;9k~sy6aq3kP%W!W*pDa5kNU`_WacAgl4ZzvwcfXABJ>A@2%WI z#*}Nn!JD7hl(JtlJ&M(8Jr52J04zVgaoeqW*&@13a(W)p;B5%4(=#)B8Q;Rec+eQ$ zt8ANIAtTjYx>BMK+`z#?qe@_%U{lH8yWvt`&_A(zNO-$0y>Nv6^`EU88(W2Q&TD^( z&$is&HTEgOgmMUW5Ib9G&%Rr1Y_4hIO_$O8mpQ&hci=8(X1$S(s9c-$A>pfi>N$mP z!dnx6R2_ZmrHAn=uiN^qgoLjQE2njjjy{S|#CB;&5jbQjzs(bBO!=k#`cQ3NxZ$f9 z#0ALBk75Q-1OIlz%*Gs)e}()Y|1|2C_#>Or`!gOXv7Yr131Ux2?rp|#Qz=H#8QqgO z+d*G>y;BemQf$x3$S~~7*w`>Be{Im00c2v8;{_R>{RR=d->)sx%A;wT!h!-+0y>(j z?d#B&#n$E;3NU#K%V?R+w%+j##72XN!^Vj1DihwQV;2MxfCUxrj@WXrvjcJP+cgx& zanFqq3U?(7M4awoWB*+x5%q|u8b73a>lzvJ~vr3=&jTJ9~mO%tstyk$g= z>0FIUi5_JFi|}5s#|(y)*tNJQCoJ(c-RZ`R5CuV2CZ8c{8N@Ok4bwSaW)tw4VSo2jzZ_!V?Aa&mU-)c_H0 zTyH)A;9P5AQzNpOdssL)+}SB1Cgv9znNAaldSa&NfNnGZR8HJ5Wyo$j5(I5+ZLeLu znv$9-{G_v?4z@qj~X zVLNBxXzaI7XYiM8F;qAoAg&7xIWFC)wwU)4Q?`og`&^M-v0K}t%D<2ytF~umcsAC! z+CxEsxHnz>{rx+afk)0-U3}YZviXX#b{w!+i;DJcUi+2;Y*2XFMBO^iz)B*chMNR~ zW^fEVV?XAb^f!W8*!$X4MECFEf5jIw`UHVv2qvd$xuBh?jQ4vo7_ajcUpgO+dpquyb=7^StX+A2*9D%))d_B znO0X8qBUeWf}SU{H8x^TrnDD}vLtCj5Ptjthc3kF47L6#A1WH4z#)UI8c6_$?A-$1 zm3KAA*c!QXyPfux~s>G|LN1y!jrjY6?92IN#fj&ETie*X`c{!6-)Lj;uBAM{y0VOk-lH?TE`H(*yDS5R@WFQ zn^W89`rZ(%k6?t7E_!SU=L>gWgh3z`?#vOjcsMwJ8`99&SAP9sI@GVc`SZTHKSm1P zdO|8Acyi>mC?~82_zrb_IKFTVWeQu6d*#sS!O55y?){`9;kSDaV?Eq=lnos>m+Bd2 z6jjVxhKA(Et&q|4Z+%z4KiTbRKL@uS*gqx1X99wv<9aQ&Hy_b?p}%9MxT5#YWqIV3 zC?G)<1-9zh{ph`F6?=!9Hp|xu_Lb@|T@pWd1o^j$?urhL;^X$Q7ct1GZX_~apLhmQa4HKm;~JG4@P8|?<~1LM%d7P zb6GXAeT7sMDJ5*8rfqA%{2%}s%3?|)w{A9*+` z`tEslck>*&x{|H7Telh&HRXQAra5MNM_Iztq`;$X(ex{u%g$^AdzPdF@HP|E+e7Y@bK|6~Xn@%8^NvV#dJ*%vX}S^YpxW9*eh=TZYg|OWo%4C}>GQ2uaGonDgm>$ivq%0% zoW!K~Ma>DMXyot~$9Hh}m;o<`Q(Gn;M+`IK68Fj^#lpyNM+GNNsZsxl+3#Po>WRA;=gT8j|$S<_bG2*>(KqR>ctm#%K!^T?W>x`wc7|p;eX{V- z{Gn+boXCOp>HcJT{8EsD;rik69nQ7ykcvo5LXpxa%JIY+y}!gf6&N@_F+DyyG(Pg> zgeqVJTF|WPEqqkTy34(DUNK*H13}pus^;8QYhzjdIe=r#;~(y>Msrvi(yy2lA30RI z3q`VTe;_H=?R)}-nDvQTl-S4R1wONZj;?MoU|Z4FGoc+F;>TH&b-tN!y+|lwav%1$ z%}h_j*MYyk3aU~M5N|MzibEB~Rq7W}At)sJkw*f@Ut>$_B?#xW5`@HB#=rgi^Vy@d zSk&hiLygK>%o5r*pdeJyr79x>*H%mrt;u{tLG^N_CQgSVi~cDb5E1c zC&!w^|2M~9)j;n;{~W;jJnMnmC*dqCEQ}DBka&c+b?a7gDIhS95y_iD_I>b31GeNo z1gMl}>|E8-(h|o?(Mu#Mc;>?JKyHP0pB#^#a^(-b1(4<7CumNnWn2F&oZmmGF(%)R z3)*LZ24s5O3YcA6V;M@N$SD+LQFX^53&eI-oQAX!vPG6a5Jb>J`q1TNW?`|?mmODL zKCG87EiFwgm-9IP8wJkIg1kI%x(o^mS^?MvV8Z+N?>7Ze3I1$qYARP=g*yQ^*gW0N z#5a+8%;5n^AlXw?Y0CNb?Pe8hg7?H%pL?HCFgwVYlO^^vEX(~o)06`64 zluU`tc{6sIOy$V5X}$l1L~Z4h1yd;`)@1^g11~oso2S+&5C9}ivhwmNr%S7=*nXg| zx!F{fl_ffATsrLH;sTl^6fUs2u_5$CV&R{IJUs-x`9hbNjn==7vGETarM;Z=^yT;0 zh+aL$NE#}2I6hz*wL{?)sj@&f6g{2awgG?-p5z;7XxoU~shp5Ymo5Q*t(6N#0gHa{ z%MzKS3IN(%f`#81eUg`x1CFeqBC`)mOen~)-khOQC;&wTN^KQ=vXvCD4(ai7j2lr( zTG~;e<6$9OXiN%-E?Ui*^U#3z+c$;ss*V+LlQD`A24rWU$fjz&Yn zIfd@eux~PAHTeU|EK_+PWxw8NyHdQ$4_@sDQUm9I>#l9Oea8 z4#w{|?$uz5L}&}hJ-Ay0oY5`%d|$h%N3*8qq4OsO;uZy3w$#{19h+r`-WtY!EmjEj_a0 z(*G5v_9l!`0J$a53Y&z*3RF%^XPVLdhK7bfVoShnR|E$Ns(wfSABQ%W%B0E080myD zv$Ep&*73%}X_=l*b@%W10&(yyVh$HE95yV(zv>y@)7tdQ`-k&i6}^J zWp3_Ua7l1=CH$%+RZLs ze3n|M+5x+*nO>GEr9T#@+9QN0B2ejzW|n9Dk0G`O{~Y0p(p zAts=Fk`Gczj#NDO#7gn?JCte$$$V^{9gN{DcgAeh6KQ&KZJWJ(kZ4I>Aaxo1m7 z@70bQ1Fl{<>S!p_2m3$Nn{3hL^7W+5&9~XK&q`iS>w2@lMih7@c_V|Q651|H>oj)* zu&E*4%4lu0T)(gcN~@ce2~4AoYr81;&3BhiN(9LRKQ+0kCrjp(|dB(i6_LXkCO^x&^^*xlFWC8cS74(; zNJ(we=S*b`49J7Y)7wGNDxi*S?qdDX3VR^70F5EDXyV4PBFbawa6d3mH!Lwaqc>wd z*c*wczTJu4FGTd&Z0!4c)b5B}sW@Dr$S9hK=seq%zO(jV68^1^?xmNJE7B|2I5Q(8 zfxKSh-0zwuNSV;^C^8X&&Nk-g_#p$pIINl0Z2Tp3Dm%`hd|hL=C)e$J^U^$TG<}>= zwe=%&G;IBp@oLp{g$1i$f7vySu%i4;oVs-DJuf|O3?19S!qbR=HioHtp&Jv|m#yM0 zEiAUD(i$;6G-dj4_P{Skg{a)0j^3D&()bd3UqU3_F}$$-J|XMA_;aeE7#FuI<#a0i zTg1|Bl+5>8g{+U;tNvO~K*4cJdePIgGx?VXN4J;JNY#Qvyk=*b>sh8PTuW4T%L44^ zOKb+aa0?oWT_fIyB()cws;>31@tlFjPTA^d!GR>HjX;B_T=zDLm?k@r4+46R&9t$=ylFYWys^0~G=$%b>Zum<@2W-1QLHhF%!%fq{YUue4{^ z9RKR?>`a|AEE;zWH!PXA>-aO`i2|{t9WsVgkwHpv{0qyu$SJ2v>kRv# zZq;{THx%3d%GGP8*YK!+vub#prRgH0BL4#4^H6J6NX6dHTbUz2Z!JZj zzuO+Ya5zN5fPCq1ORiZ5b?XLmIoG9rwobujUYEp+QexOMt&y{t_C<&DItYi`4Jyk; zX(P&m`16CqHha-%Jn~O7tE@j8OtST5JrHs}{ZKUh_q_(C1-sUhZNs(sh zdO^hCgsq9fo9P!`tTubJOP3ZwFMzh8@eR zzG>a+CiEqx*|;)Gy12xE_Gw?Mo506y(Nx>WqH_Dk_a0Nh>t~C5sou#yM?Z9)xv7f6aqEP;52(ue=QqD2~og8_cQ$~FQ>Jw3fGDOQwp9{pXhB2-qHn3x)> z3yelHXo^C^!&CVYe)SNeEL)V$wuD;(Va@UDd7A3v)#7)C`Vk(V^Bprd){@#uYzVqE!G_e0TmHE5aVy!WrhjkE<*Dup;Ei_tOK>{_Y4MB7}@o)J1M(>pqq zmyRCHCC_Kw{8UVm_~JOpA^f()d-uF6gI+>DP!aa} z%IQ`Tk4OF6!r?@Iuxi+FUM5xcWv$PQG8d<6aMe)Hi+%j~d65Sx06jCZHA^ z1AxF1Dl~iW_-x@Ai&<#{*ffP>fK{s_Dp~$4rQPx?la(TJVJR1>@urm#9Wc(9@{&j` zO%VW4^FHp>wI=m!mmSsloL)yke6pXDbwm#1cf}uIaep7c6+wDXyylLwvLzEhTE>@uif8A6ZTw((DkTxd&qS}rq>Kk zmw?=q#kUe|xy6zxPkk+prq?sfD;KR^2x_nKYD5)uo%c#TNf*nFdN87OZQ|tY^O#a zO_4&n3$KG>m6MHNpfQ1mM)WF(ws;KsBDAZ8z_kQ8ddlTC`3_g0>jf!~^=kPpp7$!x z0%UF1 zGhObK5<|v6g1&MbZaYKNEdAanXwU;I2~e;*b4g~%!$wLHVC>N}R5qvK@dhe7D$p;`d|tOq_ebT^SfNI0o;LrS0O8-rpaR>vtH!TyILP=e`giy6 zZ;<@3;R=W}j-9&lhT~auNIw8W z77pAsG@7yp2w-=TuG)ibPh7jsn>VdlR*-?cyu&qRQEF>r1FJtLH+MUM0b+4KPA$+l z0V|O6W#<_KRM*I8_51f{5Q35QrU5!ObFjTwT3QNOD43533wu|OzvB7|C5z_tvmv@r z4QNp!Zjw+#M!WtGi$k!abu^>Jxub!bLz6Gu#xmmxtS-z`s~pI@OfgRqx+Wy=cVli} z&d(Z;kj_~0+0nl)EY9QTJ~zc&H%*b9`AL1U`9m~jV|&qvX$BZUC4cl?=S${#9S!6+(~tb>l*AK{@TBS)C%mJ1Ma^_n3;w zr?2!Jj7&sb_e&p9YA6z*i>EjrEwTR~6^NUZE#e4sPxef4pYYuJi5W&Tx}tjzVo)r< zN9RCYfX&2cWNTZfvKlo8x)knyo2#jbHmmRz8^UgWLK*j0vS@`LIruo&oo(Pu2>OMb zVz2YpFrKbyJL2DOl#`Q+8yV_ftGyHzA_F|YZn;2Ak=4ol;9;ecTHjffprF9B>EJ{E zmyT-^$%sz;$;PwA1!3}QG~!dCYbhwr1n4~LfpQeE-fAo(CCw`=J$*`5DU7c9p;i~H zDPm7oL%)Q0RoFCC-sVFQl_G)tV{~(At?s2Zt7d`y&qCqTAUuzn>QoLQ9@Cdf?I*(( z$EM4VI&HHYN>ffa8R z@9bprCja!hLw*)nT@oQ7#+QHS%Igo zlg-|N)Fp4+&HV{tz_$4KtucQ_`-#VoRGOAZOi#CM+_9a5+efS4!{Y+ZAJGHRQ${^- zUV)IiB=>B7K*d!H*xfb2?o{Vw^^O`cv|l&-$Rbv}ulZLgom4x{$tviW2<6=~cU5}% zu99NHp>y+1PQ(IDhW7GRg>&dG>~v{RR*jmGVrFuDWk*z7J8u=ByUHE@TS@Xl5&<9G zIR{I17spybJ?Q+L>S$5@j+Z;dOwS`=rI0Go7!cYO4zb7lnqtK;z<0Ot`8Z%3=dA63 zm+H%v=9zcd_m0mi`(N|Wyi|)9&7GC*%pP%3T3m5t6dmPgf=l7xl{-1T})cTB+d@KzJg(zeMeHQY~eC9yWdb1|VUxTsd3b;Z57!s|pgm znV+6azyfw7?8I0xxMewg#ekV9m4+FoBUal`tZ+Miac<ep*8Ij^G=ou$Ih691)d!Qq*>-a>pQlEC-LM7z~dVA4m-Y~yBA!MkO4K%npqxHp!H0^3UtQacz8CKiWibX-V&X?KmV`sJGCJC^Y_HxIimPV8%78? zgI17dpRL9G>u-L)wq~~4gw?&|S@K8Hc7glezce^|7kqeaIwH?5^#V8Yx86xHk_A;_ zPrC2U_oDq45vAMa%APS(PxBT9gbRYM9)CGEXjxlU>n~OMy|Gj+6>65fp z=;xJtw5}+eFlkM9cKb}bMccJ%F}TcTU0jSpl-}1Z65?&#$G#`m{E3LJhj4Wcpe9}y zAC&az*_on!ay_Q+$22g42-s_guKk%GX1H;U8_-~ci2ZJ~g z1tt#fTIpOS6p-Vd@-S7;yyM#UkR(U>`RflJ;e5jI{NIF6PVZg{Tb_^C`3yK^{|l8p z|K$7i#;y{I&b}iG{ODKeix-S@pzeSoz`RCbLl_<(W1r5qUc~u}hWx@O_T$dRYgnM% zik8;6tSor}#~ctebSqwvceP)`@uMqmffy>0<*y4kd^BW}iyyNo@7*X?%sBe}LIJ0$ z#>T%U*I6Uqq;`j)0LZ-je7SV#q|VU)fK5>cQ{1~jSN5Y1B@{zRx&8(z(+vP1{{|`7 z^SMh8Bz|4PK}S>h_Z$1K(yWVr&cZ*pj`^YO*st^nhE!OfHV>}AwRg;IkG#WFj z|GolB3P^iz?J)L)Q$*XM@At);@!CqZjBS}EqF~p9q%ZJDeIRAO=YCa*$g%|UFTU^p z{q<_((J$5j-O^0^eYP5?f1q6HQazrMLY~_nfw#K22o&lk0m&sL2Z2dK5v=Wog(~%u z5JW*`$aJb69qMzs9q}b4?BUDzl%z6y4_i4_2{pGl=TH3UxEVgm2|jR$9Pm8$s`dXw zigEooke|p}qcwsJo3ah`U_$^BB%>4`fH?ib25u1?vug1`qZ0^=V{N zzSqv{6jA&J$~5z6R1#$%;DoM*fbMZ{j7FcV!#^H9=qVgB+=cI|I2>MeI2;@W0^HKv+}=TNWt6_AbD9A28?oUbVmDEBjks*w6<<0{c71Qsx9PNPB(jcQ zYV?r!w3$0p;7Nl<=;ymHm@#^sC>x)Uzq+52Cl!O8TBM(^2c21@iwX<3rPC#g^|q7l zA$ijWd56T%HSW_=URQYa;Bf%0+>_A9450#A`^WF*Sg(F`_j|Z!b7DSYS}|+OZQA^( zb!vU)@IGjr8qZ?>KXacOu1YQKO|8AD;8lpr~6dm2=h)G~H_ND$UVJt=?CFYA4-#hWe ziM=5^R0XXGtBF6rXecv`l$1stMlDUvNGOW>>4)SO7RpFU%2EbxZEeAH8TDnwfF)tD zqKQCJ6!eSf=*Us8gUXbcs3@)K_Y399#1&;FCD6i&o*s3pMkyc1k@4BtbYx@=tGjJs zY`LAQt5CF>%7R)g)LTmdDJ9*0@NeTbW@Ym{ehIW6Ws4c5cHAmY?}28oXB}mt0+dR? z-0Tl`yotEv$;g}RBkgzk0X!TACLvUKuihs379rlDHO?9geW6OsphD(nI$%%$?l+t~ zJSx_5ii&n%I2ZwUR;6kTx;zBET8+CPbultoB#M)blytUuHVRE8%Fl1z*YW5P7Z)_N z>;syduCA^UyU6R;;e<|vESmWSx5o`V-twx&AWH1Eu;1_ta0~~m2SkL*NH8kv z9VR`eA#J@l-qxMO;H%M{P<-Dbv&ijk!EttPJcuOoc*+RUaG+a24hxr38bFQ@wyPV{ zO89GQ3Y4D(b50EX^Y~<5%>i6TYt{1U(&^GP z@_BO~lj$|Dpy(51Gyuuq-VV5hOOFg?!?_6!Gjek1Nj_jzSp53>9LsAOw5m(=;qF~s zT}Zx0K${{fkC)pm+v@94LdJ{iTAxO6hu>n`aO#GkOOSqPnZ7oF4QdOnyj)=G;oC9ycylTwh;@=K!IIPxG;V}Y7)Gs)4djr6sAiq|nnpIrPcK2`hSFkfer;)lzk1|0u zwOHj$o`V(0f$+SavYC(^HAX&k+_$&~=JCyncjM@6E-pvLd?(YV zk;YGw(52Gbq-zVZCO^Ps*eUjnZ|RC{WYF}gFX)LhV&yitRbY9_cgS_EaIU~j<&6-h z0xAsirOG48VM_T~!U{3`Inc+b(6-f&Xhj zWTgGpw0ycnjXOFV4&V|<0o5RC>H_(6ltpcSdmKNAY~PMnaFLSo_hT}Y-44xp+{X{J zl?g8zam_sNZ{@+$XOL_Gq%X|Yp~-9&3OIzqUv+~a=9 zB^v1ci~YnIQ@HeewRiHrO|sX>Os6Y#mytYuI4v2Itx{5D@m(P@-kTUpL3U8?^u6hD zNZ<2pyiJ+6rQ&8D&Mnq9`5Ce0dV+&d0fr;7bhWx__5-ai3C*3=O*Vu=UvrT@>HOMZ zWWH*$KddJv{D~(uI@?z{-DplnaWtpAB6oi-dBC}1!TnWeZ0+bkLddsVS;lOw{A7Uv z#AqVv93eW#uN1FLolI)O!DY042D66h$s3dTSfy?2E{mya0b(~-3gqmc^yT3;9AcdYm=*BgU5~&k&d;Y!NJA?jIY}rynp!@drpTJCmcLC zXV8(F-`JrUu0+z96)jKdBaEZlnFQWUj`gEh#&?*oDf24RG?J)av$OO@D2=x3wEs+; z^Vhd?<1x-SbXiu*{mJ;);#t^3CV}v8(#cBigK`J1(pZ!Y)Y4c~ zXwk2sOiSG5EoRe#$AAiGwKs18rSabKDv6!t#r2b$8P41+Jp1sE#_&>H$p+~1{U`V z_Y|aJ@$R$8BCq1)g1;T@J<44vk5@U=4)PmQer0`Gp<%yPN6k$SLRO09jsOQ~ z+@IXK2a@63Y#VQHdk{*hw(8$1aCuP$f8k>dU2X?;i&ulizQPNP1@%|>;k0@NL!#6E zrc{rpv)8EBIs798u-`EKo;zQ4S>%ouL4tw{}2$Mb_hXmeY0lRN^;M94dJfTaFb-?z+u&1D;Tyi+4XR%J&cgrIX@4{rA zU#E&C2W)vDHlzlloJYvfafx6eTx+qOicuM5Ng%J*PQ}j!e25wIALt$ZcB&)LvRI+& z#+QKmmdrCkE0aF$UnndM+i@3c?RNNP>CP4h~IDB46<| zPz4{%C-F{sQe=pdH&#Sjl-fhcKQ0+sH;Hnl2wbN5@}=+WA};VICviT0`1{<`ccPPs z>&Uw@f|+P9=ZIzLW;W*C_&t_y_=RUh z9oBT+xZ&@2G+`Hyw%ZWUD4^Vs+$Vuj@7r2mpEbA)KokM&>yX3HJO3tl^BYO+ zzXeVp3-BP(AMsswu7V_wg2U)dB)b$F*iu5yY73Ytpse#S%B-^yre(^jz!n8M7%;9G zw0-C)+1U0`z6Ox_!!ut1aaV3Ddund(CS-7uKlD%-pt_ox&BS@zo3Wei7;Aa^!|w+A zVvF9Bi+Q*YD^MU*3W5|LqRSQoEhMjlTDyb8Z-M3_2^LuU{VEc=_I#g74=$&8z79^CFHQ-`>^DSK(QRdFckuE?bZUL-W+8 zU>X@uAV@m`6?$@F;s=Uq7WDb^T5#o2EznAr`A;x{!P^PyAZ`-p?$gg?y;qIjGaEUP z@i%;2HW++_@Q%#babn?9=WODnRn|uLYoX#l>G<+?ZnG&*kpAYp-j{*y{)y|>Oyn&M zL%pLoH?MS)o9A%HB#j2h=C31SF}lAQ8P?(6ypm6@eZ!h+? zY0#4?0J)uwfpZ2TF<>?~y))#3CJMkT?gnNa6O*R4_M%zc9yqW)cSd=bKtc#L2WVy! z5kZjJZL^gHJuAu~{QZMD>)ZUl}oA;CFYSD0s5u> zE;)yKn?n7w3Uzvtcq&i^R7B2coS1Pc{wr6dK9<+v7 zBTdLhn3%Q)TOslxnw5tuZg%Rm1bra9Od|e*=Gh=q;5xqS7}Kf42oAdN8qGbAkR#SB z4Eon|ySyeDv&p|SVlRDLNRwD79Xi|)>fcW!K-Zl!xas0US;X|%;;Z}8c3`H3eS)Ne zdDnvC=+CbmZod2>XkWUEZ|Cio#Y+d=gnC9z?;Aato0& zYER3lOrnS0p5k-lXB=PrKq{INq1WwkYQWk>U!N#=58eBII`p^ve;pV%859N`jM0*IVmC87+O13iwH4w0`Z5Vda;#%8Qn+0uNi4J^aoHaM5jv8)nqNh@^;mDh?dL6ey_?~RwH?aVKcsCHMPCpzr-glC ze$=8`8n-ItP1v$tQbZsj9mbuXoV(pQ(~?wi4(1Xe|9F(Jno1VPkom#MP~hTK@wqa) zGc-};mVi!DDW#(0^qOZl{YiC7VH75hxSW=Oe$$R$5F((*OanQmuv?Y3%XHg^>TlN% zjs9Br;e{2Qd;NvCZo9mQ^_P5U;`xv)X9dlIQ!bbC@n~Nq{oRFw4W_!%_oBuhHS5`~ z#~nSd6jNVwyBN)sG1x0%ougjXlA)kSE`0TLcrd92pDL^<`f)T1yTeD8F@|s3h@8~7 zqYH`NT%G4&%G-L&!eVn1@%;7ex6@AwstIr5E7l_T5z?6)BVOG{vATfM+>#zN^I5=ckcvcqWnZ zKOG@M4wgkZX*;JvcoJTfQr7)bMhfGK`IX0CTRum`h;R(`i%yz$?c?oOm5!VdYQOf1 z-Fk{gTsycrT)^yUdCe*AEvie+iJmORGT-l2!-|Xdg`E$YQ&8KA;z`O!%47e1Wze~T zIv7WCfHl2;+Gz3Sl@e30;Hz9(qKl*Hr+ZiGJY|JH@g-RBA~PLW^~6Z@CY$xt%SX+( zO5KBGy&Fy)&U=D-XRG9b zw$HPteI!DOYc59>PNyNkdoLVRk`|{7Et_qRpMAj4+bU#y-9%QZZkcz5^KG4c?XBOJ z{&koAr{3o8w}1ZQL`*9EM|!qKXxE*?^#;GNvYdRt`|LNyuKlY8*muh{>ESxaZZkS@ z&&T4zn)8D{lHM^6*1mVF182pK3uD(%E5YaC;Af&`ET_qLmNvXKX$*X&=wmfQk;)Fz z?G$UN_sjq0$(WLyrjAdmfOda!5g+cj%G=Zy;)C1)gNJS@E(|YFv(tC7R zE=D+Qq<#^(ICmc}qSdF>$4o&!4;qa*(a6p5?)keYe}p$ws5kH1F2E-@f2C^T>@nP~ z;6q*iKVR8+;Fd-k7q{6NXZAw&E}Oo{v;FJM*W%p%gm?X)nB)8$KY#3u6o_M4@MdN> zy^5B^F)c$8efbv!>j1&XSC*!T`^$Vzm}7!izwZxw9NwV8jmnHMT(CjszH~wRrgJ38A0v-$BXv-fB1-=J(P}RN7GM72ru~b{QC6((a_M)Zg`pXz39_;@#NUyE5;X( zzFvIhE8&p8frfR#hM-a&hSB-P1uh@B#ru2cSDc9Iw4E?J_;z%;kKfz= zb>TXHBb0$-hyPwnFH_D0F46CoJM@{m&d9=2iXxqAK~t$=G~pzmN^D`J7pS`Oyh&e|LHcr7LGA|zn;B*woU%I-LbzBnt$H^{^|7Szufu1p0D#? zWPO3ZeAVVXkiLEEpd6}}EgL%wOE*bgpnpN9&`CM1rE|uju)FXv1P~IAyASsR-W&^` zubK-NoJM?J^?MuxdZbWpvkgHjb@^6@&hEtv-;)dMN40$tFSV~+ESGw4_wUQ>OnA@* zmTlXT^v@FxR}gU-b;V;=L(9T$_&fZiXPd_FOZ!e1>tk^)fH=Qar_L*X3Lxj7d`Cx~ zTaulBaDiV3qV`?!-vX(i;NYar`NpOu!Pl=RCMTgj1F(!u^FL7KRIyvXZ~@v@tOcA!LCLcsa`O0;CD7 zB%xjh=4uaLx<+*UCyAH(+^`=?(*&&?xS{HehE1I}xIBIUW*#-Und&qESmIVnF94M?`DCs!x>_9ERvjN(vBEdmSO$|L_-HypIz^)L$ zyXMAL;D}4Ua&h+1Ny0TUNF8+4G=+`nzgu&|=fX$m>g63IhN?F~K}j9j5N*J{49Kme z@Zlh{QFH}ZW?(6>eYDjs_~Yxj`tnhr83yxVdd+~F7(4~f zC#&)hN*1{o`Az_zZv=oU`zg8n8@(?^;0zQqm#*6^CmkL-P)qH$u(tbvGZZbfRgFQb z8%Dh(SW%sXh{*nW_n%m9vYPvDsgJB~_i=*$(^)W&x&DdmWLQZ26Ci_f1-llh)0d$H zk8#z`6_}F&HEX5ISq;yd3WK4eZ;~&}lNjfMZf!Qyh{_TZXFnBT($ygH+cjWc1@iW& z-fS0Y1k^RHW{$u~Z!(HrNCyfn@Y|>-<-)p!dF=wckx;Av`8jXT;ngwdB6-q+d9YW~ zsTr!Z#GbUQ)IIy}nPPe7vi>Whm-7svVZ;^=jZ~fu)$Z6-%digzFa_rksaiY;9rVB@ zhjvLJkb2zxR5T2=FHcOtQEdsIM*La;v@5@P82jUhNd3k0(C1Cj5#xX>1$+8X3JT<= z^cn?6K+Y8P{ip84&(H0s4g0b&+iAjaS{5SHZ-LY?sJGTnAQ-%g!JDmYRQ>7Gj#)Wj zVc_clO8)UK+VPHv1X4*8^PjV%v z#90{bSmMieXBvA}Vcnmo!S!rMI$|8phJojyx;rE*HobM>*8}knd2Z5ktUJ_y<~-3^ zVc>T5zN?omX}ou6N=%tV@<(aoSus7Hw(eB;wc3w4E+L*tx&>6d*K2)5- zSw52c<#2RxjG20d93!#y=;xK3_ctQc^*3#cv1i2xR1c5)M%SZ1>Yw$WU>?-xA|6&M z-6oEJ#l^*Xl5S3_M@TqbN_)^Qu0le!MQT|J>3O-iboBJl6bR+InN&EK@t7l8p-wm1 z>|b?|yX|I1efH=$HXxODzn9}u4s7R;NHe{$Nr(!Ll1Rhr^!)AAtZHLjAVJm=aa$Mr$5}B z-!&SLd}3uJxMawtZ2Yh;Kq1f?$QRUe%YQk=%&6wkGSkLpZ(50r9pHPMAo_GNtyVHq zhd)F~2X))Oh~zdjpY`C7nof0y`)@K~IZ#3%x9|s`!ro4c8!v#k!iYG3{(OCj(ftRn zJj(&Vh>D5|*a4krRqp*ltM;NjaY#ADUxvlRWVcY_T%dXv91?`A z6vvPbzIjK!_iUfXICYREwd}$vW*3$ToxyKX*|D&GME&(AcI520UtEBw6Xk;r?By7x z%EJw#dI9$_Hp}mXMg=r{_Y-9@_njtrV@{k~6|?KYGqgR_Q?x zuxWw_*0|6N^k%T%^sO(q0DK9N3x_tSRdcDT!YO2Yav18V>WbbG)di>tyfP;(_zq?Q zW1os3|9tfoWM|`XE`!4z_>m&PDQ0QMK`5c30xfv5C0K@R0{Bg@KzY>;#oEq?zg7r8 z$O7H6?JhCJOl8~LtHAcKDVU!K@lW5qYwbd^cTz zZQ`ihg(_GpvXKt4k<)k^rTKft9^Hw&cpCEp0ZAXEM%423d_4zD3L$FsCt-!vmMO5S zWdZ>~HS{jIz*L1I61-?@0rW(}^S&1*4&n+%RW1tF^KBp}egbVE`IuWK<+xLe(L`>d z&-RIQg~;b|yd^|!cA)ZGh59^5*QEZJ;0zLP)Ou@{_PEr>I4C;bQBbVd4c20=y;DnF1_cyW&LrL zsgtUgr;Uw+Gr0zK{nZ{5S=^6P-Tr<-hU`3eIlQ}%Z$vS?o1T`DV7`IJOfrr9-Ya9@ zWvVjG8suJ{(Yd)k5uPGf{q)}Eqn{ii%+n%U755BM%IpIQRyo*(0@nrHgTeZ(*{#-_ z$aS~?I%0qvAPud_AkN_AUi`ni3B$WD{73e-TA9erjvgnqQZDc<8vc&M zkBDKWuz)uhXzd8TzbsED?|^w7+E+k@7D|Qiav8>|8qbE3Ih{P*jLgLg?&*rzHv%9{ zV{`{k=F68Zm3Es#HBwo5$bJ>nt84<*iSPy5)_Ikc`qB$4S3bF4<)ud=(>wJOa5L;y zWV{|(waem?2{TC3AS(TKiM356$0sMfloJb;yO0T~vtK;(OG!WCvt(?HH$s0>j{|Trg!`dR{nXn_1(3+fA#_e|^=#)&wcfI({ws}w_D|c2n`y@94egciyD~!r5SxQ@ zp?)tL^&>)WY*WVfljujUyE;^L%mZq^$lA! z(>5pRyNjVaL`tKptpUp+zbD;&B4m@x9avmY(Z+C$p!Km7Xp@-Q-VR3hMqd!6lwSr| z#dunyylltGisrq~3dTY@!(&qYPCH&!bhly|-5<4oD2$Vm87ANZKe!rRG6N4%wL~6Tqrm!O?p(YqkEA0`51M3~3-m1H zDsqcTHOl^8f*mS5;ED{asC*SJ=J-x7@DzM?4jZslEUXR&)3rKNM?sF8cW)NdjgAq# zOv|w9W@h3Zc{*n>LTN1O#r(a%A#zn+ac1z`Zg#29`nIE3vta+S*4??F`(N@BGX<_~ z3*!<8S!apLsn}PMYWwWva}v!;$)~d;WYmfjNv{2UXixquzQ2}uThHE6_8u`;X?h4p z{cZ(i1kR1T#--g-zHuj}U!w#&fp`TkCb|D{ZsXX|H^&TGYctn86zm`EF~Od2HQV(@ zV@MG1Y~-?BT+~W$|j&PN$Od19Mq_=;mEc}-61}CQ7`m>&};Xyb*cE4~D^Y70Syl{%#`R$u*$wc*a z%%r_HBR#%954lILtt0-FNz$h`njP!AY`l@Pu}|-m`u<*x`^={prp2519KIn2r`FN2 z88dS)JEGt3?PAl@`BPu?*EcOnH8Eu{3qB&BFjj{KY78tHJlx&4w(-6G_JYkJ2DAM2 z`mBOMP2Hv&mpNI67(Npjo0-{_r~+nU!1lQ8e?DgPTV3d{zd8N-{{Xk10?n%w;Kex0 z1V}rU;3a73MSO%(EC$B@`|U%zV7ehStlM9&(qV0L0n}Zk`N!w_4Z!mM6imVDm_7y4 zDW$O%ES)ZJ{c8#t5k?8upa^S1_`&<24PNYIR`7AM;JSYL0_fJ_&(_vf#!#r}lK9ra zT}9pp%4erR{-6m4+y3Jhm`;Y9-QV0#A9?%>vd3WeW3NwvZ5%iCH-OzU?BiHGG;#W! zyJGH<>-A5A^Itdq{{?gIcwF|1^#KGO!otD|3Ly%^TkUzE9aB)y3#k+h&F=mTY$q1% zaxO8=6=Qb(&8dML!Q;Kw5`PGKhRo^t*ROWvqkNb<*UoeO;onTLO`>+w98+?*f6&`> z`;YBQ6pHcv%eL&uva1uw|L5D85#A*EM}V6iN)Lj?fH{E*jhDAK1Uz7A0?017yI^0r zX+uZ4KTUZ8f>)&=J^)4Xo&A}10n`dGu_1zNOV275O$9Uv+J)X?8?fDsXa$UCc)I8U z*3mW9E5RxPPE)9sYk%9=+1M=XIQTy9w|0i>hi&VhYbjIpgcT1Q%T(e!7k29Yw}cWz z*&j}9{r^%walwiCg9STUECEOGLWO|LE;v&PLf!?T{3)ypEFY-&@5tF*$GITY>Wd+& z0{z{NRynf|uGa_xUWba~p8P+sxVS4h^i6B)4D2xrDH758Yk_>mPSnBmwY&T^wWEwm zm(uA+b0L(IyN4CH@qS7orkjEEsMNBG+zp z=zmu~{Q2jirP}nMmwXhn_mjJAx!Ely?&A%DEz>{O2FmQs?+{p(jUsZLQFah!b{c?Y zS&cgs7(~UyjE#*m6^zx@#|Cn={l~bvrZ+%0OB?vOP;fw__5_zR@?f_4jgK9pQO*J} zjXF(jZ3=YIJW=mQ4!lt_1A~UvR=|}F^z`&X)Bk*RA0Q#1jKawTy%+;U=HqLEAHZq{ z)Wq`2%lUEvs_oBIQIwYM1fj0zXmvTc9^I(IGfxxvd6XoGHCBQ$1N`IJvrmD6*4qrA zSpruIq@yp-YXEi*;^%~j7yu{4JK&Ch*q<|8EI>`?F5qy%#eU=cKnK4Iy$}FMfifM0 zJiHFu3-Ofa|93#$WV4Bco4UHYgWaLoinF{raSE1;3N&1ZCR%Zkfg+)8@i7dhr{|yv zMy>#IBAb>NL06D159q%^NQhoPInHj>xugh%8bJMxi-xJ*jas+!1N?tfIHr)Pyn`A) zQAc8=uc>AFpmhT(=~*g8j8+2}Ik>Hu6TYKdz^Vbr@l#b@1z5k=fNj+2pBvsn$eRXs+F zFyJKde;*GPQi#b9BXJEQmkIX|=_laAFcSf3X`(Gq&b;d#{NaNISfBxv4VFm?80I$A zmF~G$f5ri*dtM<+)sQ2?yb8Sw(*~G4e47na&^iRpyz}eD)400g&dzA?n?mY9RS8l# z5pnU6V$10WV5-AIhns+s!O}t<#^Y(G(FHo@Kw?rJ39TQiMdMJ5GRhl-&L*h*LCRYW z{c#w00*Hj*2rSK_0wt$qsDA-P98zUCVHC2}E5N90P4F6ZHWb&0vM}v+Kq`lt6X&Q+ zV6E05d@__}KunNBujLac1|rpgK6|Wyq>+k_ghn7337)7nNLHmHxh$c|*5n3_d~`2! zK`wG06fA+w1b_)V@#U_Xvr-9sE z6mSo77(j_`^d$jA=S^M4k34Xs>EM`6OiV;n2K67yMEFxG|{h_dz?9PS5owzlYR4;wPkC7Bnz2s(P+mK}&pTf27S>Wz4RKY%L({ zY-2BdDqQwzotm0@7HXFjQ(p8zS11ABa#go#q=xwd)dQ}{f>-w^;=CEdp(+HGwQoYy zh&c-`9v&WKLSCLVOf7^JT)Ic-ef$8B4|$5TTN&uomGrN-hCzTyiAlB1j`y=&(PGHs zX~AQv@`HF!A5u<9B&2)LScO1T&s>{J_K(mjUWi6q0y&u2c32zrbO<}ZC%l0^R!i?>- z%;_Jxt!LP0%Ou?jS%WhjWKhb%=?r|_!iQH+yF`H^iP@%#Z_egL-y$db+H-9|b?p<`7JOz&))eudoqTYyN-^J8g_?(^31$k^kfPi!m}@ zKzNm$*|+FQudE!?$^ktmMPwWN1*0>a*h}(NOoe+nvj{;C!3Q`Sy6!VvfR z;Z+3+Fq|$Zunl8?I{Ic}uMr`hE;lnCH&o=gM2J$w{?5QONFgm-lr8QAum}Negt=1mPQNQavQStKf(_$~<^wIQDcFUaBGZ*s zR5UOL%DBt$qy?F^TdsI$tDbkuyZ4?7>Eyr{4z0$}cxy=RX6`OyL(4~OPj=}N9d%kS zMr8>gE0!b_y{!c3OifJCD+P(4fLcQulAYN6}T z&oy||OwsNjlE>CaXLNq+y$CcFdqSfrXt^#xkc3&N;1sD@pv(je4cKn=_4QqoanBZ) z8|C}g%%O|EIjfmhA_QNIynJEqGnmHusVi4FaXu=u^X8u}8u_GW3A27U>n)k2d&D?^ zhhmUb*wHJY&W=W}y9!DsWE2$az^Y>f;TC? zK!oyx$J8C-O>V0Mn0jexKzQ&ked&D<9DAb1g6eIzhkF)na&mGeo0BD@!ee5pzzzeJ zC0A4%JU1xW&Hx8^7WAolsY}mXm9T&mo7wRZ4(P+Ht4fJ-CSf&!hq&zFccVG z>je_C6pUFf)4YBCpgwjgqgqyL8nrT1ES}kypm*ct(9OsK$L>RaFR8}09|e7D;{$VP z@5{Z%3>xo#A_)D2hjT$WsyDby@X7@q$*jp~aottfJBS^5YWCS>q3_FPeEE(N%DZ<8 z9w)PY_n^|fVXUCswlMHQ?nJlu!bqC9Y4A8taE7I4|0h)(UAMg!8TCT!u zlm`edHaz|%_da>v)u=3wi~D9&Ab?qeBtyf)R8VOw0FA*>lSQZ<2tY3wrq|i^Xqpm? zGFw|(e8dwNG%ES1sHQ&=fkxWwXE33RxL7CxePL*F`;iS0Y=(9}K_(`q5g$xZ{kZA` zRaIjR0=FBt5V~jOThT8oT+o{=qyxN+C)2MonHGy{ligT=Hyh^WyT=l z^M{bEGZ`9)g59SRsv#}8i>Ny6%|_rGVoH?SPzjH}D>%2~K(!1R5@3klI5F=VfExl` zGFwpdwt>J6u?fbl5~;>)H>sZ_nZ&;opJWD&hS3+xYQE&Wg&?|&yuqNHAO~R^9LYj6 zptSATl45HU*;4gUTuSQdy6Ves&))$jr40mUyE{5Qo4h5#$G5INPs9h&2a%dh00x=e zyyftdP=z3oc&{%}F&e_(Wsoapg1V1n!4{|ELv8l7#6}k;Jt`C2vp+N$qEul<_i&;q}tj(R}g={x&M$G`<4WS zq1B$a)O=Rj-lMDW(iHDWT15^|7@hCPkM8i<%S*}IOO0;vA>7^FU<6zXh&=IRhVelZX0^jZ{kVI?`Ry)Y9MDIs3IFQ9xH!|*u?BkDSW8I z!^a;XLpsi34inb{AU&Cn&@0Yq%{r!8QFk_ZlgzTG5kvf^RpISBHT{)xJq*-Uitk@{ z+T)vt8kT3nTx$fF9Qhf=rl!(X{ivZl{sDzVbRB0Uv%@lbzbX|1YDD@Uei9|`J(Twq zGe1Uc8ivm1re0xkIeFSWj9{@JtXHf2CenW_LGSX#cqf~){0Bd4=tY8F@v7~%PSoQv z+zixK5ogwZ^;4UDDSATr;i`?2D@;!2HQrZki`K6N89V7UeT`fxU4*T6W-@uY} zo`P?5yc6ibM2wUN4Z@q$x>WeGb<-Q>&`B6?E$|+QX|wdDph&JU5*g)YqkQppC6n5WOf>#hsf1ip$RrHfImcwa0udhhcpO z6u)&>bieC+5zvu6&iAMAUxobQS^z|IkqQcDTkThA#*($I0_$u=0^6--L96ULb;KxN)mv8x*D%@v3sIqnV74J+#taQR}{ z8e3P;I;?RWZ`2|!jS*3>m8B@pK(fsLWtFKj`dq-8v*Xxtg(p|VbJvtj7}?AuBYh3m z&bX_0yh_b$^2kx+P1zY?d8)`#6^^2ct;72o*5;@}cN9cw4{BHT?08NClCS4QV^;oZ zQ{^H%#;CVvJFgmJjEKS37)&437!_|_hdfhTU;><)n?!yb9c*EQKeq%=tRWVMYDLV+ zK@4QZerHuuZOlpuWrDsL$mU2$dNcC!SOtBj;rK+u!>t750JAy@`qh`W9Qw@%K2Tn~ zt6^%VsK> zzCez60?+_-et~W0OE7DPDmREY!pNq|HO*T2ZpeYZ9+T7L<#EtIg~D|XP(Lw*3{ZLG zGgHtwnj+!^Mk=J}AO>O%MqsDt7C@~>E0GARZV65_$fK^Vub+Iy7O>dqRSXRc%?1Mo z${(>;oeg{r&L#_1y^Q>7oRGay`1843>_P!^C2qWnH~P}jXb|#wA%uz=k^VL>A0no3 z34F&YvaQ#puM8EkvuF+8??0YYXB*A6s-6P1L9#XKBDDXJL#xFM(Zj^<_zGVZJ#BM3 zRjbZ&V}Vs0Y)4+ZRb|whFtMwYe%betN0Vvv>-!otrXfA^&5LsK?J+-1FLQX`34_lu z2F}1RD+Qb#_$V}hYYhExk=eI_MzS33a=3WZ^66#FVNWtpO2O$Cz@yNZBTINLSNXS} zR(h;Ju&I%!2RXLt7T`rk06F3Z1T4}3X8^O}AS;CBiLw6X1oLzN#7s?B!9oOMf8sCQ z8!x2|^RO$l@)iPEng%X68Zi2tRQOsTz%dPy!FvuBd*k-wupKN|;GB3YK%{zzDF|lw zmqEDd;1!tr0g0UH z#gkuDgBk)M+hQj~T#Aw9Z26Jsc-c8;d2uK@Zb%%sN=h{pW^`8d^H6uUR1SO7AS8=4 zkYtkUA`$F)p~E4%Dk0soyqWp+13@*jRoh7?6pT?JK6~aai|?ok%n)%#9-e=B89i3M zb0@+~C1jw`tXPU3io@`7vvv&1EvL^ew}Ocad})aCeICa{$`26qq{V*9m!+j8{pgo^ z0G1A#*umgZp)ipytc{#!89=etc(4<-?lhm`3~n1R8eCjl#n6k4L`ciZ$|@tuH*jXJAX$yHi_o;2JWtDj0XZN-_P;!kEDZu7w%*aIKkmUgywmiJwCY*kmjRZ%2*sn%>9*Yq z_l8*r<(G{{9W)Bl9>d$5ZdM)OeY$WfSxmkc#xx-5Y~OL)Ys+n6pm$YQWk%%=SE)W@x9^+XTUXvnH)#K1fF9g_ zI=x%v^H+M>bEQGnHG1ZORi>0Kn|eoc>$?)I%9{a)2b}a?rKY`9t!3)mX>Uzv7A+14 zgvjg9oh7bBxz>(TVv*z_y3Vj_F!U)Nm?F7zlkX_UfUfVU3UuqhsDI9`<)1>+9m=8J zd~G2Ol%$ZfKEuHWLDgVMoLkdqgbeGvvb^;S)$-hkQ*D$+c7*j7$zE@%W3VSNuGCBP z4~?9gk|(o8CxW4)o4TRN7ETdqpX_E=68Dxo(hX5w&ZIFpEiKXv+lA`{0Z*T}bKX1B zFPi+gs5o-4^HS0EFrQ}c?1CG3DIe3;Ac5zJg58f6fl1t7pT7*Ij$rNGi&0XFfuqEq z7sIWmy4s=ZkF#$-&7CN%QbV^?t*N25sX3hlo~c%C?`Ge+P!?5bVYyTO z)Y6#-&v>k0s>J~sudT+wdKu+?L-pjqtB+c1IQpsKPKIuYxY6Fw?fRlr?`*Ws?WYLD z@$06yoanty)NaxGP}#;Fufj=M#`Bz-#flZ@6TXga9WK6CQ$x9jLIGK5-f<_n?*e(+ zKcD_x(IyKcU`@Z6dVd=OuW3)$`KrK)jtQ30{gt}~S6zO%2S$;f-3y~lp{eh&xni@& zoGBB#>SpB}rPnwS%(Gvng-inC!iHyb*Taaot;c#$#lwB3WK zvS^~Tolw5*(v5zv0gsx%C$&Y{4s5}m(F&W%CTU-^6a;d=Xh}Y=jqF76s8&{xx{h}j z>TZb%hF2YR%d=+FrB)!juSnSA2wh!$wv%EiD|JSumn5auB@QHE)dOv2!6%6_fOCSS zChSfS_JRCC?$*x|a)S2wl1_$w0&YuTW|TR)H2r0sCp`C$8)It5(F?rA^+a@63Hot% zD9Kig{e^}chKt^d@UR!p`Jd1r8e_-s+3z57M1%!eS`xR6mG?4Q^gbRxefc^60jFxk zDlco}K=nCetHC#(E5I8uEmNG$ZOoZ*-Aw8eK&>O{W2#svmK!P8k6YVmwD_a%KS{|S z6Hf+cLxMVv9Hq#LJew?B=bC)ZD@xemYUkQFFEvX3@%8+37!U-|j^Y|6|C-NSn`@<~ zKuziBp9W+^h-BV{-5W{Q1LP1a>ZWPolN*HTrP_D8U@1=5-Fi2+SqBrT*V^JMaDXaDU!Y&Ht%yi4*^;gZJle-1GVGfuDWn z>pkb(pI9t3Ka{hq?_MS1Xc^fX(!e#-u-|peE}dyIY&KgPEVAyuT*Zs3IK?p9&iyCK zighn~&UGQMvWu=}Ix6U6E_czmqT{Tu?iN$>*=JkmecHeN^W)t%dWu}l>AL#+x7!l7 z$HwzH2^W`nzZlCnoCchc*u;sy%*fm}@;6c`&o`x&?wLdcjqPRZKhkJEe&rYPWbt&H zvY_}+HWte*KDb>?j>Y&DJ+b8Oo{fphcz-c(m3pu5A$?Ync;3Nkf786GGM}>~WAM^Q zMM<{ddPAk8zeca^=&XP)Q{}g=u9D06=!yz#J4J?agq#xRz+hjFt|=YI#P&NnB2|iv z6p6=&H^&`l@%??raz2|qZjDPa%?D^!Y^!{vN4`zb~(&R9inTzt)v69S(f*-6&MO zYA;a`c(ZwksqI4|=Na$9XR}*Gp>5^g6_xe$@~*ra+RHrJiYBaTUw3Il$83cJu~+Q& zUy(Xqm?4*EEZ>}R6pZqF?BT$BnN+Xzq0pj%D>|6d%xSaUTkx>fqLs})PCipr|BA|5Sm9CO>(a&2W^LHij-$N}LMt|E%!fi45Ug2Md2?RI?|6oe4_leM z8p`)v6JE5vHs}`UoxCL9%Y80Z+~l)SXRx``NkrK9j_)z2YAIh!-u<3yxann69^UJR z31y$b%`NM=)jYUO=~>B-OH$HLP)VHl0qU(51Z*t?0ry{;as+#_tzYhMO)1|N)i8Xi z(LEOJ%!^DgI8>D?z9X~MMvK}El_ThD%@>#bPQf+sQLOIrO+61ICad$`2rp+78<@5v zU5`yW)YMhW*qRuwI?D2V@l7anx6rJ>AeD;edRY?GbzTYxy0t#Jb7xX#4GBo6}|4#=oTV(@zAJ zJ%QCC+I5-n&YJm1+3x)K;8BlM|Iy62b(l@D!%*bjs36bXAqUN!y{dw>&h?4|_me1z z@9x1nm)Fq|tjp59`J=8bMyvjms#}4sQroua;R*#W26y#cO6K%W)>UALkO_*?K_13t zjo6{B6>QIx^uy+cVq6@%0|nK$^04z36ZE)}7|M2TsV#U+ug19QueMj^9ljeEz4Dk) zJ_?4FwHcVUz`q*P4*Ip-VtFh|Dt7w3MC4AlOXZuph#?WHzBzs&7HFR#`mEs}->*JK z6rm?O>w1|yMM4{nuDax!jKwWAXL09AGO!lSE(&Xnt|C~p>peVXi^57)3iT;`_w2nJ zGDcOFv&yT<@3vqI58c!_S?f0)9;#Z(_Gl(wUzgmqB_%AVZYmA5!CU7pUVb~S7kCjb z%n|i2{f62s10!#q$?2v)&;NU8BfgQ5Ag4hUQrxZ zwcMC3t;xLO)lGRFuXU2|3DM%SyO>{fk1|EXk=*t^-HlWmelJqZXJ5)Ws-9XOunbeE zr{3{0Rf{>DYgr#D<@ zE>q_*!Da%48_w>H2dsfgn~ScO`rAJ(xSRzDxc~FG`o_zdsy_G$*;nnO0(R1#KO0r!ywz1Gs4>kR zOQ1%Cm$T}mLAb@V7pmy(8MkrS_4K{>HGfaQG2KDl81&AR)Uk81%OShQVck>n(A&g} z5oLE&1*sy4dA=+$QN7wx9K9Qtpe=3wJnXjWox+NpqN_1)|H4)#z04G@#%4XX$>jV< zy7%d-&AWakMCq{ojkA+P_3z(fH}PeZ>8OpIG13`@W;H9N$M+|p$c{^i8tz+nU1f@| z(NZJt&OW)%7|};A$r?I*yv-ERSMF5mJQ!lTeRJy*dk8IJbM1#dsmo4Bug0OqSlnjV zI@C@Z69Y5Ns}%BPltTr%$Sidt3dQ%o8e(qKZ42kO@-*3v_4TpCvUiTVt2jXl7t)t5 zl<{2P4Ub=CcxPuuMqaqHsF4y5bG#X+cFI{TlXR?sX6it~*k$wcq|FsF6MkKd)@60E z$BvoOLYltjS)R6qN=e+gdMC@SzJ(m&B_lr=mo;QvQVvAv&gll9!yfK(?-X&=-cn)m zC-7T+vG%@iCtF{p@yNA^ZwDXteG+d!wqomVl)U0*dH?sNXa_&3?dL!I;F+M+DSIJeP+Vg=I`88d z6?NL<_Ze5$#F5#iMtYm~tdE}N+Sw#PBA2p?#}+sIn7ML!*{c^L&A;WT1xCFt$9Hzj zJza8`DZ{Ebk;Z5x@8g&U58l!opQ~T`w2SyYQfH{3NWMjbuWq8^#CD0+<(~3rrhtCu zE!PtxV-A!|W?;mx#TuETB-GDvx70E>4~y~dvr)Zb;ko7yQ+jA zAqc6~=;NZHXgL04Y}!)|5VutKv$2sCY3wzNe=5)4^h~tLCOd85qr@|P zYUnq_!a^WI3C7u|Y+kyiS>L7BHRDh5`pEX~E^nxYwbCN$QQP&c5Xm z!H+p5d}wBz?Jf7--YOHXO-Wfi>QAQ+U zKYV*#<7p53q7}onNVkj~WK2DlRwN1~EW$Cr=*B91I@{c66HvVg@%Md;<|Ii{QNa^V z?>z5R;iGy8p$vkBRaakltmr7aieB`~`K$dBgFIm-RiR@-f=O#^@*}@4zwcugCC?6* zxt|xLuO~o$kB~`F{lszoUyKde}Q-(|QAr%)7 zZ%1xi?$2)WF&?TavhewQmfZDjG7euMdhWcdSM0XuDLT zEp(p?4bvhgR@GPylU+$|R|<1~&(Ggjqy3RvF4>bfZEHp?$Wi|^m9YocJsh0B&b`cD-N!CMS9TpZb{@GFJliX#@cdUdf z9Gk{NNUw{n7lDrG3SgkclU6fYyh1U2LnCT3wbS>{-nPM40`*@*Y{rx1T=qU1k3JUh zHWRbY6pu^@%_Ze`cuGE|TT)-RDpwS?)EGQA(KnJ-B(Qw9Gg2r#PV`HAXAwsxD)=ke zyZEQF1dZD&ou-Rp1&v2*Yu77gc>I%Px}DNYy-D(%cMhuwFS0(DI2y!EWSZ)aAIp;7 zGG5_s!f6fT2<(37ZNFVvNWWR?$;G}NJiDi?C6Cfm+}HcdH!d=YRD@}8u*yOBB>iA? zo#GQKS?48!@E84QSMIyK8XeqU%1ixl1RNjLv}v`3dFCQmL}+HMnUJodrtP47aV z4Wr?E#B7lRI`2wzzIIujSK7Wlsh^fU<8w-7JVE+-OiIX=o`lU~7b6qOalYpn0RUwV3SA8Ze>+;>dxx5@@GS%J#YPga{x7$-{ zGv3JSHGiQ(?{1ZeQb_`d_q-NDX3p@(&Db%yVavTu5+TLlah17Le|;*6a}nsr(!s6@ z14=_Af)<(<=)WRQh0(uSrUXV|rkdN9@Sv^PXm%2A*t_GK>&js(PK?iIP4rV0UX&Qy zJS(kVdBHJOg6P~Z84>l4N7AdW31s?}s=wdMBB2<_Se`}Bdm%aOUOZ)P&gYFL=h=Q( z?x-t@jMOcS6RpqNXsfDoR2O^iNtNM1XNz~S#pY0tBu)gVMOlg;)c{De*CCLgEgwl- zYT}9_gsaiKq$Fck)c^P1EftOB7fz=0tn|Z_?3pG3WkwG=TjOgEWwABP9CdoD$IC}{ zMf$`Ei|Il_#$GCun=Uz9!gmJ>sS&6P92q?O5lI7cvTL?Vji_e69TIx9)iBfe(h+(Y zJH|7CXSDQWvuiKrIOi-2`gz3Fy@T0dq=r1~OS-V9FT-E55N%V~eQvZN9tYNv;foyTuZBZSxin@RUu_g`yny%cv3;%(a_L-Xq=H6( zoTIRlqDx^@rH%^xgN4(gsY1^q`&z(m?5Ja0x}$97>PD1=y9zSl8!vkM>glm>;c*6O zeBHOWvFJsIn2MRSP{eiZ?(?=JjEnPalzdv; zdlBa!A7!K-xg@IWd3%ma-{Z-pFLBi=v9vHwG*6Pyn4?gWYxx+-&*pg+kJjT<>+_*y*Qk;I+mp zyU!Kt>keDljK?qav1KC=;bIq=?EJ=rkZhb7v<+**Np|8+E0ZcgvCWKi4nk`<#9>NF z0=IT=j%dW;UdUot2M)UsQ2@eT{#Ck`jRQg}r7i!-Su7QK;_3QLKbGx7)qaLsNP^}mY1dGZg|6XXNA zV3z5Fw^#UGSH%kb$80`M)*n57?((lO&26p+Ii_OLJ$|wFg?pKNMuCPKc&AuV+pi7I zU#I^Ev35!EXxZ_Pcp4#mR6}PyLH*<(UyoRE@NyU$6a7;j^77_NH%q6ahFfjEfsb_C zdSsu*UiF~?*KAdVIrn&go7%RQJy(^o{FcT_uuBobbV?^9L3V4%%XxDf^PI8pm)@K^ zHIM#k*hrnBI&veRgD*Mm3X$*IIqvlJ&3AL$!one{i!!!CNSm0Y%zmL;-BE$vhR|4jT z&7Z7fy8R9#iASB7P%bH7&g`E6Qi1@ISn5(1zUSzSj^LucqavkPdSPKk#!3lAa>Nw# zH~GpzL|3R$krYZ6_2c^%1r=M}g8Tb(n2mCiROo=cc5O(-_;~nZS6Q(3T>fbUEq;yS z1?KjMZx?P?1k+VCuHKbRU)$EjzB`)yVi`xVL@#l5!EakcQ?oK&*{P)1^o`tTwim-_ zx=ix03Jnk6$7cr&(c&?^1uSMsQm6w39mk}7<)SskX7k^h zZ3PGS`oB~B9%inFeeyJ5Ox%9#PnA+CmLD=+heUP;bw%*`{n;v?$|)%DqLdqRRrLtR$NnkD{X)2L9_jFONzexn}IN+8T|KrVDn5u5`%j$Cjo2`{grlXL;80 zRs9c`&%>I)`MI1%wll?+ToNa&6kp_Yiw>sA;>ll^C|TIuUif12{St%Kkdg{B!EnCg zxX9HNGTf(is)Qk8${R_MMe0{uYzMf@%$>W${VhvHRmRB2{^2#Ca1xMqh#yu<=WX6T z^3+ujbx@#b{r&?PnPnwGU>R&c|*kq0@Bcbb+G%SZ{f?cBz^Tsc5eFgblm2)%=P1D-gD;joKbzxbjwhR zOm{028`Q{c2$EaBZ-_Se8DvsdqJ7<~l$3X?VItAGsw(2B+2SNQj;$SbllTUlmSni_ zf39cF!B4$tWtQIOH~;Nb%FsQ`6pw*u=?c5P@VlMz02#JnrIZL-US5MMD^L7!?%6Fc zS=N3K&t?x<Ojh0S=GXW!2eE8>w5)cRunAP>gf)tgL3u+X!9d zZI;zTM6dGRD`wWF*^J=WD<5F{LapsQsW&by(>PlMEp>Z3r= z)qx)L)T@Q0fIgGqp&nHKv%+z6f*U32Rtlx`ok0z_<*Z-q%G=M9?KeL|^Jo_7JbP!E ztrS1EaPnNMNAQ@DK8dC~K6HX472rI-nSIEKWit!~&V z>zop+;rz!C&uvo?3c^9hmGB+mj2u&vkw7^0@SYMrda@T*`0~Y%b1>)nl5G6LeZ#0W zD$J8@vct*wRFnFO7*kRCoyM3MLu>p`cR_*qTa~MaJ{EkO9~vI}`EnDFcgHOKxH*{- zEE%!U84<3n@;D@$u53KE#D7@q)@VuM-dlsij*#ux!hJ@ul~>yj@E|<3vUw(J(k;!a z*YdsXI2gBg)t+E>Wks{pd!rFP2y5c*8Qk@CpA}z$#j-1)&ZMzU^2TTEp4aq|w?~FU zj=nrEG`!7gefttA?{I#1z%;4$W}kHdtV+i4{bdIgrr5yl|!-R(Jj^2Tww<8eXli{-&IR-{Fr(ce8NwDzBJQ*Vrt2mhmVdMC7(wEiQ;A{-rpB)r#S29a6Dk}ADvC4{DfcW za5tc_diN1;O!U4FmUz%+$zJg}d2O6EbUOe8xE6l2hxK9FJ}voIuLz7XS1)QF`J$=) zvKOCJjtav6i5K5%cU1i}&sot#UYe}uQ{9=YV!hgUvF-2amWqV0jzSy6?gxKbF5^9^6U)%GE4S$&F%4xg~W6923eEp|?5tB!c&s8IWL_aAM{wnExzUvC7oR9Cy3 zn@3F8+&iv2zJ3+AaJOod=YKC#ibIZNPAYj+g7+@?rxY29YIYhnrefS~V)}QukN)emBWiNK$tHNU=Ug zJjwK{Q~j_3*-`?2m@LLYDA%>})mJ*`DJP3v|3c++Ig@{n3yHtT`)1JT&M%*BDNbMZ z4Q*OCjD!;u$Vat5Y3iQ}00p*jX=sYy9PvMv#=39`$rgy5(+~4s!E45x`NSG2)f>-Yn{U zp@lz6>d@8Xa5SHir|6$N$|%J-tF$k0@kZUGiwu=rsZWb9{eDn$jgPM4aKt~JF+0om z!t(KdWl#BU+#ZS5crEF%_$D2lK;F}+OHx!=?kH5G4J6HG^(kPg{LtS(ENW2(MBB>9 zU9V?o*RM?=u($gR(ce^dhMIz7)0Os>+5_}=`OC_~l+CWGxXSZT#Ox4qi1Tk811B}+ zU&I*Q`Vvu1M$pOpGs>J;u&jjjIK>^~iTa4S(w8X6e^u9IGqyJOUodDOiZ11>4A;hl z1Fw&ZCI-@ymlJpT@p9&J;+jK$x^tzx>u^WgmU_!S<9EDaqeEL{o=kTA+kZ5q zo)NJB+ewiq8XcuB@0R!19IZf9^+w@3bFmJ2X(L>QE_uq!u1cjN)}UZ;AuD!owEL0n z5vS*a6P{0FebntD~=jZU-1n7F>9lQB(^ANtM$SHPu+}CxM~d;plB6ndX)UUFFxe zkN3&+aZ7v*ON$RxHw2Y01ZXwgiq@T=KQkNV&MQu6P>~T=pqQ@gH>b?Q}6k_YhDCexo1H?Jh^40 zarH2^`~XkzCg8~o1>fh6lEU8Y=Lzz}LxGK3ppCw|gDhjm`A@SAYqNPC9L!~G7wr{(g)QzAWxg&HP=M6RrXejo)(4tU$Iv%#1gCyyBlpC}?MSL`xmOd#LU^zEmX4A&S<%_<&%i4z1t^PAgi zHFO#$KJ^5!rK8H3Btx!Pm_N$s)zQp~^<_=A+_Vkj1Ex7E^mqF?X-ac<8+i-y`Ku25 zHm064I$T>M85U;u6Vh#up}%>m0&-b@;Oo204^+Ph=iJoO&g6~V*!mvzrCg8zbuu!2 ziC~yqCd{^zYv4HbNX40E92aa&l->sSbf4kt$$(ty>6x3rJLdt6Heoc#*vW5B@eqt1 zQ|goj;Yi_k6tl*$jj;Lz#$NI?5g@VNqKfq15 zMgo*Xop4_WY1jT#Kkiyb1<;xFoB#=y&-zR|XB5ihv#}J2u{sv|=z*~9F1NO8<&Mda z+uvSKFE_3k1gJDQnIvzSz`J%X+up*ShBPbR4Lt-o;GOCznNNJhEr{M%pqmnR^iI0R zZ|Pl4Ugx>1-pz0D^cg91)$va1toyEOXKHLWnKc_<#@p?z^Z{$K4JO6v5>&4Fa`sxs zd$xQP0MK=M;Sm*OA!EGc?Xz3xDzGQin2wB|1GvyhABh1QYvE}0j3#*+&Di$Yk6$P#HlT;SmlBVO(6dc z1ae?NzWp#34?;ks;z8FMl*Yqyay1pR&^aV_d?#{tAG2!zB%L& zwFaY5V8svf2>KTnn~hf*Oy&jFiKOiyi+TR3l9k!^{+K`v@AMQ})$UM<9Ox7qof=9T z?)1Fe(D{(z8_}UdYN73If2UG#P;Ngb?|-4;*?v*Hip`iHkGPZFswsaef-(b>$02M_ z-wkfOI`lGkusC8~t1%zVUODD&Vgdk4{aZ_9#jOjQtSXrGX|e@@P1!cp;3c~+i&X}b z*&Clc*Z1FUO~)$uLVlSB#YrD{oFZM}T|e>&G?>2r?x#lPcF=(xH@@EKX#A^ztP~p{ zKFNqLXpllA(?9fbh#$Q3-O{K@15=>0rp&yXOXQt<4Hr)&Uexcq0Wzoi*{z?1#Y%_j z-c3Y(9_564zBOH5S4p2w)@y?qdW2;};&b}-PKva9XA6bjzuhj~l70)BUO0^rqms7MIMYO7-~~F^Uo;4APs# z-w$zy()lz=7SgHgEhpY;zwYpj&XX9m!o{|T|ZQ}yjGG~iC13c64!3g zxG;B}s=DPEYW%)L!4<9==YL${2*IBkNUh6GY_Sk<`6@#zu9C{(Md#5shPozE#MqXq zOF^cUF$X|KfajVAXEVA^;p=1%yWJj@Dq@8GCfh7}nsGbJ_(`vDzKcRp;>z^e0tKrp zOY^1H)!L68yZJHCKu~A(N8E;ju~5^&rO=$A9DLq`uS4E0qrb<=sD|0at~4xZhC27X zAHX%~FjjGu;5Obtw&98Z(%@pvpj4~R_9wGqL4!&Fv?Ub_)u!nE+dFdGzMeZf1_xI9 z-ns(iyO9+MY8CQs4;0zcM}~NxIU%bwKboZU5W!N z3*@d$Ry%dPM@2|+0Z1`F>>_Ro=D!gMr0{Guy%nP~0efZG913w_35 zMdY)OpZ~`SeP*x_JJ{UpJY`dSL?~FpoIv_oot2@qS$i_|PX8tV%Iih5^y9C2Zt+;S zq8tCzE_3oK+9aXE(^5j)mlO-YDc<16g;iAq?prExkXs-FOY1aQX+JI1OVs+W7{m)G zm_r6GrB;%;7xDo5lo)U~(f*vp*4Sh8Eb{(+1O8I#>B&Nbsm98jare1HC6UyeVsCbv zHEEYI33=VWxi}&dthoWDZH@UM`=F62dS-+ggsadzw9AT9=`CKpCe$oWt?N`F0A%PgvCLa+^mJ^W07Ag~r_<|R9EFdPiMYr+Qf z9T3AWawRVv;b-&;gDOYLxP9bOh2N>z0f6JNU&lk}T_u~z>y3JBJv5(tVBggDNVO-e zFimR>3)ZH+dE!E2V{dGkMoD0^>q06B7OXp@-5`@@>dD^1lAXX8HtW;UZ9`Z3NH>4> zR_eu-#_psoN3y(4F#n#Ahu^f$rh*Mg@7teS&5qZEgXJ}|aoohPHKn&18-`|o3icHc z(HBpp(hiszle&nLh8I)K!`-;zm0d#>(stoZb{|~zv<=!?YuoR@A%-9uHfpm-bsqL@{?7-2#-!=zB=p)}u4kkAlngxbE0J!gTCT z!UYF#0c9guJYP2xF?t(IN5YdDryBIv9rW?1WYUVZeiHEvx2^|C54yIXG*lME)s)I&=Mcy7gZvHCawQn? zk1C+l^X`9QM-=4U2x5m|YLjU#cMz-j`;yB;aIM{KEfl+0cY`@y2s{6S_VJ^J&a}uv z>#l|k2{|C@(@mE+=5&Uw+v6JFkwKh|jtvQgP|Zc(6lU8mM+-F$^!@lS1=CC;mAib1 zi}j7mKvPE0%+iklgDxPTDA6}!l^&^D%ySElRUQ__yNK|Nvn}E=B>46(i>Ve z4UiQB039kjC-~cMtv23y3FvRNCPK-@d2AA~I*uLN&oq=;FlFbZo=kkt# zE2JxW=WR-M1bJEz^lJr!GAtQyP@LR0TzHB5L@ESo-lcFQ*cMkerjr9)(^vku5S~xS z0RytvWbgA52`lEs@t3vAZ6M!|Qx*lOGp>8{fqcqAmy+-u2~b;kQ&`HhrIlCs4#DSa^>8 zBm^Jepo@)I)^u!Iyg4pZt7$!h@w@||yp7p%H1nD5>C8vDkkBFND zg-tE94(QSM@<>1`e18b%p4D4Dp68GZM)&Lm_2=MRDd}905;hlSv_$xrxtUq+WMbI- zd~-v?3^b&nxnZWIVJ4_~HYrXbJo`hCSid5*xVxkb_2>3?&z`?4Z#iax#Km1MQfTei za?Ytf*PC4QKc+iD1u_bpYuuc*B&Z#09LMR)=4TLoAl33QO~_*uH-ZW-u-IHwN5?$& z=S{ewj)M^bR5ZfW0py3N&GXnUFvmZ4Dc6xwbt$(JzPHUUlk6@5Q6t&cuBava;R3k@ zZI1L8hu;wz9@1`|w#>JcLu)BT2%ML;$P8ppO|S@Cr2Che;*N5H!bD66+v%sV6LH<) z0F6a*Tmf%TtheR+iwT3hl!ww4d-P*pHy3f3LC3ya(k|U!a1}nFVSb8_;roY)m8!tv zUjUtH+mqJXF^N*jiCPu?EPT&Og~mq;%IDA2b;TYc|0ikenE@V3id%@~Pd2yfiE#;! zZ_vB*%b;sQ!x{;eGL6z6~ug;CitYRRxs|A<6ZHH zV`RNIR)O-z47Y$uDDRcdq<`P3 zd)q1l@;DRyW#U%oEuwQ)e$Fj+d14)NZQ>FSo|Rd8mh@!T!z#CcnWVA)hIpkDqeY*OPMWo7EbZVh8#`k4wS91 zn#G0Zafpze?|rZ0BV5=anDcf_sBos0)OF{T5s6FF!gWN0wTH~V1jYI~F!}N#Yq~+$ z_?59eI1h(3&zfHUr!#iZ zC_K>~BWBKfIt4n>Uk#uw(P`oYq}2T6{=9SIeylWy{~owrCU~~p6 z=T?4lwy2afz9fhn3}Ff3Ycp3!sE8grtnQY2)eh+zkgdR zEl|-gpHR3Smk|=t3+m)mkm|iH3gs)r_!0!I4lF65N}>og_T^etTU`;e85_!U)aiT_ zk@uW)#^9!iNwuhs2I?0mvXQhJu6@s44QWWnWYtpuN6$^(myAn5C@p1JCRsUebt_J%Ij4^XA)k7?&0Hq+$lcU6YRZ5Wm{ly(+eGb8KLfAR1d?b)xj z$iKrNkeK{0_80Q=$Pk-yn;uCsZbwpNncQ46Qw%3P)(rSS4m-P_;kfIv8ol}eL!Ji> z4enu+cSja0Mzm<>X_W*hwGjTQBQb@Haof}(6qD>JJol$;(q~Aa-#Q9V=JLT|>)X0n z#&9TUwajkba7v0G$zhYwg1Z4?V*E7hm^M4esm*fvuJ8cd`)Dd1a{9wKrAxWVclLeO zsx%2Fa1B@ZG*5n0UgpLO(Qo<$UqKzQgZ( z5`v^!VWY{q2H*OLSy=F;AMF{OA=Ookg$J1D4B!Mw9(Q?a;Dd4Z;H;P8o7MvQo5erB zy$3|RRD=ZR92mllD0FvYFRF9&UKSX5J^`9J`C<3D0#-3qR6EQ&L_Q7P zu8ci^s09uI+Wtm32_+S9+Wzo+bPd(}+Ph*Qz0-u4KRwOYg}67m^-o6)vk>^B<-(xA z(i;gpL%tpo%I0!n8CB^#^EaY$XX@TIk5@V{<}q3kfJ5+av;}xone4=A|FGKiVD4r5 zE4s5Rkz>`XfHXF$t9dl5T+Zr9XMY(dY%By7y@gric@V5Up{I<5jKQ=+LryS~(+eiX}Pg$`$Ez6^ba8R*X2m)*m=+dD0e zMMjSVTAqyQ6b511{cqyg2P#UUVF<^XeOF)Sm#5MA^Sy}z?vDjQ@@|ZOrZk5avx$K~ zA@``nX6wk2lZH$XC>3%dG{RmzO2WP^HrDRjL1@BshUQhe$2NqBdHMbE$8kFCOg#7-EP{nd5-N)XQVVrPge;5M9nDqVC2PMfWJ}FZp@f@;(G-_*+s~; zZo?qo<5<%YHG9^VH`N6AnzEJcK`iUB;uM<|z$r^Rmzb{DY{xipxO1m#&SYVI_#cAr14mmdkJ-9=2GoFH2QtJK3nB ze~RBcmFBql_jUeC`JcLK>H$WcB3P!#OZgm@0AJfM^rb0$4s_HjnH#S?emFNcXXzQc z!}1+uJ1;kKCBH6Z9TVaq3@3o{Hd!*u&NDoPYn^qaDZ4Oc8B0X4Z@qSWU%5*D`7G%M zjYA}zoPVt9T?nzZTa^#vR0lJWycBB3T)4(72bVf_^K3gUpB}vj(z$z>u z)a}f$K7WxIul}!2a(SuVoTps`mpB*(IZ>`9eb8>L$he7bJBd-+KeO8rg4+i{lQ!wm zd3{_R{r#!X;1ca!j$tZ{tM?uwnK*8Ay2IHeD1EVXs+4|j~5*(2MUhfoUsL2 z$3JzvC>xRcX;{6MTb5)13+jPsRUF2;aDs9db`eT-@?zuZ2}GG> zVd9c6aU=3R`&n4|_k{-7^x-kN-0dHYt)j_X!j8N2?@IeD=@t#vU1!Lio;}gmp{i?L z*zw-yLhcgPG=WV#yGpA?dW}Z__A>+V_y6Qb8kp?to;k)hR@R%w*HN!QaiD`fwY$vz zd>s_?AdoA)#^G+ZklQ;YpRz6O>yJrafNLmPLC!?GXEPFakHa#SA9y@~ z0!sRr98)fLQpMAEzFy~dFXeU0+7Z2d$E*(e&9Q9fF*3huFY*I}J_vcmLTOnat4y{7 zy`spwb%F83Uu&b0_ZFW@{z$pw3lkr9{%VJ#InGUHr`AmD7eIdyuu8I+XRN? z)dIqaO7s(_f~i)3rcTZhQPKcXcS_!ZEB_jSliqI{(*n;HNiU-l;Y*@wf4+vc zodsFhRt>5PH>Eg$3ETfJtk<9tin(p<*5 zDF9R%yR_-mcNmu&_v<)O4n?G2R*tUWZ%^Y6pGSWMR^5hxa$6ZCKvTQAw{%qc9+ubR ziyYI1&No7qZk;j&1ClFSx?-MlQcvsmqnWeMZbLZEON#EO^3D^CW{42^`UaN=W$kFG z8pi5zSKY~mn6VgoBJt^3rGX74-9H-`Z5141GN22VPV&6Lf^(lq6Yd5lFKBq7C-uYXVWTl%hiu;Y|ongRD1~|P#CQgeL z@S56^B6{=ct9Q30<8ZV#?b&ZB>Kv8&eCE{~NNROJ`EncL1jWh%~CU zZL}XV2ffv76jf*=Xabf8LoXmTw^+Q0LjFosjI1ubQkDo9e6jEu_6d}lS##*#;C_eN zk%pQe5;23SO4f;in*6F{n|t;R%J`l_#_YY{bCiG|d?C5_>SChcmTTn08{5Ls-% zHcstJfL+K60XAeQ_PQiY)Yvr$UpVX97Txpdt?3!AAcMxa6a5I|8b*(_8zbf6MU#={q?@IiS?lDq!&$h<$F9MkTUru zjlaif^=+*qHYD8u?%C%YW@G_WehY#Yo}~D*=EqdR`ub`F)?nJ?O<3N!gCwU6d2httQzrK6o z=zlp)&OE?C^7+(mE2wP64J_=}k>f?=Ck{Y!M*PY+NY*osY;WfW`|Exm_`D#vF;8); zC~g+@(xuHFOHl)d0=#mm5K^AkRAq2y*@9|+s!giblzl1tKcGLc;-+5#qno|g{yN>zjk;n0Y9yRU8P7lA>70+-0ZWJ)W z<)`qBjFlrF>wX$3K$_%!Z8I(?Htex;*)GTj=4AVOVCq!n6}dSka&aD1>DbtB(HocL z979dF_Uvip?TX)+R4QzQsa;bP`Ll?7;mwn|)^XJq``HRli}z!Pzo#$=;6%Ng)#b!+5EksLPXr=AYlpBS>yE@{??bb0<;~VvwjQ zH`335duKfTqerWFwN?)ZEE;Q4m8Pm&BH1JPWtV0(D-o)3-797j6*%A0>IdZw!tgER zx|It27&F_gwBuJ6)x+O;lB~~Iot}0kQ~S5|a{Ib*$SRPG zW^;#~a!eWSR`z2jq?3Dzx&J$ShfT;m3r5jXyND>nUtF095NSIq(PSHXeB|+Am}ugu++75PHkK1x5^O$mI-KYQ0J&W+ z3)1dHiX=QdZ118~`lTLeA-OT@*T%;K1L+>cSM~DqM|9!UgI-^5UjP8r=Qw8BdV$b} zh3h1AKiw7iO$jJ;JFi^laf}Uvimp}T(BwN8&DCAsiV>NYT(D1hkDJp^LYx$}Wg#@v zhNNvtGYxG7#waV|6$?J0QL7jmu<9=OHfV3kzEVO~S9J}A=t51G zt^~+6gXQowROjWOZ<~Lk+!7vQYb{ywU{?WP9{hMWG*<64h%FnB9ew zbQ5KL((|-v>Nd85yF}uD2K*8*AcIwao)tKk8z|qQ8bgE>%{xnS=Y*ye)`h_>-j+@vr{FlIh<11N@52>MDA{Bn%a5<+qZglP+L?KfXxoT zOD9QdE_l*&K&G}0-gEj2vEKzx>1f?-{c^Qa!5QWg*Mau7y3sX!%cPJUhT}<+`jhX+ z@?0maN~OF^{aqOyu6Jq$Hr{g`fb@MP~4?r5HI#t)!;GLlpugEaz~0LrCKI&uao6 zEk6`}KrXp_BaP)1zVa~EFtrSedQS;|LQLgG3ndMy1yW$y_G4pL@#Mrkss})Wd&ZE) zVNlKt#P(N!v+#Qc&khW7Nq3FX-Q)TOEm^q_qPjqdS+l@_f<*2v%;cH;?k7h-u0k_k zVc}yt{|3?82uV(|L@uPeA!rA+Ui5rM`GZsLp>gp=) zFZVcjFQA&2+3GZ&jKY;EH?|B}hc@a_+wOW?Z>&7?m(n`nI%gE#;Q|NMP^!4J=2BE0 zN^aWca6nM1kfNM-%P=$ZjoO`BDO&YO%;@hxoz2D14l&&hokcGO1ZI!1Jk~Rl3Q&!F z+u|aibADx^Z1J^P{cG9hzBzWOy@mBXNL@9GDx9%}Kg~T~@N(Uh`TVdHCXl;;ry|dD zRqaslOU_p1vWtxkct>u2Dm2oLj&+Pf2;Q2XZD$Lr=gt09L#5`= zq>KnOG_5CoCs7gB0rR&XE!MKw)V@8nsGjmm8~-qxOjWX-{$ZG7!&NN;YXx5SxZx^Dj@48zMc>OJtG}wZ|IxBj>;)Rji(h~AM7Du?wEL1_`!V~>W2{G>R z)f`|fc&T59nqYXw&5*&yvh4~^)#a*s6=4CIF~@`1`qnoE&<#Z48cRHH@cCh@X%&%E z(xT40?IHPgLcr`x@yZ@LTMYZqsNUg{tXn&}+PC$=LO#AB{mzm|sFq?Wj**>Nj~*&@ zFm>GkaU`AqXQgM}+pT9(^e1(=%2%_6 z-XWrM{p|2@H3A#7B(uOg@{E}>3*N%T+W4({N{bBsTMbMWe$gW;YEEbzh(U3e#E zHBUs17FBu-_mDybjqRz^3pzXfkeKC6LvAmmD&AKm&SdlU*;hY*<(Gmwbaqo1VCiRV@v62Vw@cHxTQ&-t(ram7gw7F>)9D*NoY=RZ zf|1p6ZME?216g;Los2X@;g5j&6H^9GlQXp0rzmc9`)DvG6PZ3u*o=o=pWYNFK*XqP4))XMBGTuW2&2ux!c9sMSO#CSQ-%vj>n$_b?% zqzwKc7Cd{!%8++(0yImEBle107pEBkq9zY6B$^msz(0AMTqLpPAaQ@x2=5+>`GXMW z5SfwQums1BqCe9?jUZ#ZdF<$v!x%IZA}qkkMFoMOvM7~WQT(cA z3%?{}wlehv^~=p?DX{2V13O(ao(!B#4>=;-8LU$K(9!%88r`Fe(@ARSfAeR>y-+~M8g=aP>QDrYgNHS0WbW;#9L>|#afiZ^7V5sa=-jB2zM2wbI z184Qy?Wze4&zlBdE&XM~Zhq*&v8jf>tZ#LMhQExd)xeYK4=3HqaXMFX(%}R8aqhd@WoY|NfLUdTqOOjwCGN-^35rf zOHtL;AF8zLe6)As$Qr&UInL2jnyDOW{!(1PWOjWMBih>sp7b%b#s)KvQTu3n-FC_K zTwzz&MOJGP)k$lsr{l7Qk`}j7(#g4JkC~UA!_*07A@z{Uu=@sf zwpTMz&fl~yr{k|tusEP%j*HwuPZNxGr{`Z95$d;lCkv%IbR$S2w_^JCk&5H3?Pw9fa9x&~8b(fk01{|YJ_liDWl!Z=|qdUwp^Il5z9WcUm;qY^#p&I(y z;?!Q992!p$!M~(RXse+=W(Ak7%!6zN>~l;LSH-&utO~ElR!0EqB9!Zsn0L-aSzdLc zJ58%ImAYW`S_bjT*YcBsvzRs?_LRHJ_w>4tZ#TOo+Z20+*?uESS4T)|&#Pyulxu{v zcopbYL_}haAoY*APMdkRZOb0TD7WyosyWJMzcE*Pqo_dtqGq4OdS#%Lr)=TnB}j@T z35=QyM=>BmlJ3X!AuC}N?zNYsJTE1Ccz7r-uB)jjb$xFo>UuCSoACfHA?v#B=2e)& z&c^0+G%Dx1qcOc(0=tl3hwf|I3EwCoF|ZIFk7T~BAE;8AH|Va zgH$UtuK-^F%K?UMzDkpkS777){WzS4nWb2%2cGOtZ&+R!pyMXgG~-^K(^Hkx6Yu_{ z5%|kGQUg0S^T4hWQ6zIVzhWsse$ARSC4_gh;N~mT>bOd+Z>hu^#1k0o6i6xX=+IUp zWB_5p@;?1gfhMY0LpRu!-+FS0aYvAG$MCSO^}yt9ERz+{VEz+{NYgiGuM!|PL3 z!J926E@FmP;tH#WMJPv_3R`jHgMG78nig3idK=#W(!q@CbEzrF+-Mj>&BJsHQZ2o4 zb-ni$SVvhnZ%3+v8|mc9kJ=f8@ltE}p`h2nFj`HVtsgiVALd__Jsk^&kDd^OYvv3D zGJvRyq$@e%g&Kc9clRgAu1e!fYTOYN6a@jZODf#`dDC%Hj47#S>BrN#x5Kf}eX?$R z7qb?KT>=nMZN4RQlOe^ms|jvu4F&E%finsP-E|9I1?%XpU&diOP6=&6zWD;zjkZVz zC)B4-wA4h!gNf_}y*IK^5pGPbe?6FYlNe)*X3yl>3z`)$d-f+L#PsyR2xHP}Y7ugR zsqG;Ei&f(OsIeS2iK)m|=fPyFN}10fLs3^p?v(1{s@zs~vT-w?v$xiI7KWRK2T}43(Yd;IE!SGq2 zy`iYS;aF#0Sx&H)lAV9Bd)%!m(sjh7-AO7vSj_bWF8tD-nych8m+gvQiaA)wMBPVn za0CmK{N$7eF2H#@FIpS5y2t$m#I)Oa$?RBf_a{Mj;KEj8tdxV$CZk^ry1t7zykt79ndLxEmiPO68B%Rg) zFJ1)QFWd!Azf#yt+o%;RM>3^)Znk8hyXhwlC%Wv%Y83u$8gxkZ@*^fV7%aw4#TiQp zmI(0!bQ8X5{l_gq*psN1Sc4uTd$_@G>SGu*jf5a{T)n?mYN+^Cc+i``+2Hrr%E$`T z3w&ecrg#NqYX{U4RjhTL?m_G2lzG-S?<+P4`4)TCJn=OL5;v0vDq=r9e9!AY`j-t(cDEj^oWZ{D=l-XJ&5*Wrg;R3Jm7Ygg=L z4$sOJ*FsWgpH|yp2hRnncOoNVstHKb&|w09?p{3!p((7uRc2+*`r@I#NDEe*2Eu%G zo!FBC8{xz}_=H^>A1WoI{hMc+O?i;u|*0%iA^V9HQ|yKUA>P?=Uq#8^?B8z_~Xk zVXoqhS}gGFSSnvlei?%Q+-Y`!^he$gE!b&IH>S2=-?yYAn99eH!RD5vj&ii_zqd{3*7A&o-D&cVS{rfVrQdftxi%q8kLDb#d8G@uOpkzq_U-IE+vw6m z;}WK|E|SUsi$c59R+rUtmL;g(L*ws00gy%whwW(A{s2vJNs~)I~Z_THeDDSz7l5i(ad@3Ana&2u2}|Pg>|91IGef#g7qpt9bRRf~PMBUK;Z2jfjHA z2?5W$`V}T^dT8i*@1V`AymcalReZ0$V&N~Z_cE0{856$)ZNGLWm8L*=>4J7c6sFPl zU*ArVNaobdF#hcTVUzls-v{ZyFT@c7|J98w5cEI10E-zUA|@#V_l0vYqr&^hwF}QtMKQ<&-a;zID(!WiJtPps?#c~&drtWxu zY;~!g!oJjdLOI%`s-)w4Pd{pA!2|C!m6xKM#l?x33p!jMkzQn2G4av? zhLM-~>gjkRWB;GzAFUVN-xX##j;W}heG2=}dSog%;uJ7ZZS&)KTkXbHdC=?JWtOjL z^XCKoJo+BIbbYBzG#WlCc5A!hrMLuZ*a5-=6=qRfY8M-s&VxV&y6zRJ3@kn$^Xrrr?fYPHGidFH6dc`E3#UMy~XwnZ*pI zFR#=w-@K*9tSAncCT1MfJT&@&=l}of=z%gVguil5=~*e|R-~=I)tc%$eHuMD+z{rZ zCoS65r6H=DK)A?zW4yFL*oDAE_}Ix!YxQJ3B11jM$)$+NiI)NbfvpaiD(O~Re^0_m z$-n7{|2L-+`Kc5y@gMhB`3gp@WYdx6;Rk~+@ivG|hmE$~7w52WlNPas$yc(oqD9W}l6| zY$GJhXkMg7oKG58r4lc;6j3q?{p!a!vi8!*V-{2E7CZ3j689>sEov7Elt}{Q&{kVb z6hfy(Q>wGpgu1Y5&&tb+mx?9CLv6cM+yMh*+RVWXuET^S`khu;aS9D;ZzON+LeORM zI1Edr0FVc2MR!0wi^|J~4Bdatg}*L$#xUxO=f1zr><(4~a;}Seu`#^s z3h>!rE}TNjh0fNQj>pHgOOqS6jiTpFjf5m(JJOwY-j-(Jyg|o;$f1X+m5N zF)zH#M@?0xFN_(=UcJat{8ik6{{k^~$VLCU(Igu`uxxZu;LPyS5aQbPRwoYgs_tjA zy3g+e!vf^MkLM_kU;o80J3ff{WQn-H@UQ!ISGTAa(oycuy&WHXwXyEFsi&jXGWp=9 z;W4vX#L=dpt{pNhTJLxomo;8Ib-W{{ z&f>2l#V{!9EvgS&WU;-Opp}}{?%U7Lp2dCohdCKe0JYOPd}r{bE|SH8h=vpk#6V{9 zTr;_55)#spJU;fDQ|G3`RW=S)P?f3d>obtC^;+I1c^2)>)CN|a82TCFc(QY5hizV( z?Q%i&lCfvtM9np(AL$F2niSDK*TC@u&o!N$)8O9QZi<&)^(w=%2TC?y ziiygXYlbBwT&A~FeaMRU%8Gfdp(JAr%xkqHWyaF_i_ePNkpPLl#qT%*2jG#aekcC$ z+f7F1HDWb*e;*2)Y)R^Miiz5Sks99V??#!D6YPj6(YZDD?!RY_7ci1H_`#69a4ZNZ z!Em+Vq*CQxlDd<0DAjQ>+Uku^D@19WiDMXXA%lk;EhgdDsp84%sWFNl&z>01;Adf&^d|pz%fG%U2K)m;{`a#FV*gC>f49a`{v!kayS4v6fnERIYO-eeH+{+f zv{}88L=_>4Q+@2L-ZdQ4I{Z6d|6h+X%m4e=8YewJg(Cc41%WSKee zfvnC*@J+Df-O4p@)Ez0+HY5^NNrHFOEiMn81RGf#1SM;{8n4^$1Dap9LVxSVXE@$I zD*lqO%@mijR+?z|1j@p4=tz5P&lPeXI(f`J&%{5BG8Tcz@#7T8jw5T6mUlADlcCuJ z{d^f6%RR`^H<73NMI39MCGhgqoMAM?4gOvEOw-3$97_j-Szt2%Fc^E|qK3n@D)zu? zrb=H+pmu9eTji4?d)=dg&wX^1?jaPv+hM|6D#W?(NBbt-VHQtt*T%(=?Q5%k>=rc6 zzi=HoFluFOkZQdsX9$x^A0fYMiwNO2uGfpFT$-%I+T=Bp&werd*TP;*;f#d4a<@C%}6X77N$wn*F51w z<&y;y8EbF^6^iV-U8oK3JMUK5?s0pYF}3lVGKy@z!ofcxG-Q-$Gc4N%OL7aIlSx z`0-)q+J*d5@z;L(b+`gV@WY1yD{xX{+3a^Sy6}oWp}i^OVb05|7;eGY-6`%U^c>%73`u{#0fySh^8wa;sfW249}@_%h7-e7fL4gY^PnY&wQQ)z_IoAZ)hrBdIq{;4qBglAHZnz0la~CyoBqTxN$nOAJA`9)(LAa9 zT9sfo=7L+2NVG@Y>|2`b&W|5JQnIi6Lbhh4b1VwR+?$7&nZik@FW&t18{ft1*S_E2 zU3RqZFHZ0Bk9E>kNcVAg-($?pXsn+0;ZFIltICt|K)>wvFj7V$t@?BI)BTt|^7 z@N^0fnpL}F3tFk?hc$`q2nU-YQq z{CUmQi~569wT5P7qY>D6=~{a)tfFA(sp|hQ_7z}JZe82hA|j4TmkyE=BHb#DgtXMq zHFP(aFtiBLA)UhzQbS93cS?76{u@2|p7;B||GNHtUFQhH%=7G6d+oLE`<`c2`HuGQ z-vF_h9rd))w9927)*ZC_PE&6ewOr5q?SMq7b*$Uvvbh0_jCdX1*|k~noQMKul&r@quUyALk zTDUT}t?(>zRD7|e#~XWW=JE9Va6I4<%!bIno=_yuuuyn;u)41-b`3Dh*=B4W>x?M# ze1Bpjt35CHLUxw%6m2=9mYyh+65#$sio}j=t6)VznTPw$CEPu#9QSwx@%Wkm7RT!A zm3uCyals80!Oq*n=_?4N=V25-9&EK!tL__5#r5S`m6i{ML>Nktzpy+w!`9#y_?(Pp z1WdPY$9BxL8(1^#?A@0Yk{xE%bkiwjL4TS*@Jp^_bvsfwULtkdN^S}kdJyVrW}O3e zRD0Cn35|mq$;MDBsr3I~1*;8vej^{#ZC&(xGXSeR^XB&q}f(v(Gm%+q#qrzZ}GoEJ9 zVN2z{q9K35xle7XSddeLqt@k;YJQ4ci4=z)t&pzjjrzWTwOb6OifnXG!vgrflJO`? z>_r;FTDkG2g|oN&x!LAk%l}=Hxnp34_Ib8v1BZGjE?0GOcR}ieP`+lk*YzjY$E*ta zyYfSCn6Y7GLBIcf!uo4vz83-zao%$a`%sT-jN7vleCD^W9dGZx`IK6NzpHbgK0+F% zOCuxQKkd6uyvSYqny1VE9%S>G!*STuIG!wHD!h(^cJbZT7vE$*6TB8*oRk}qnT?PA zO~#Pmt6EJ2{v-@`qY3Qmp*W|C^ej4oah((HjMLnv)njV*#|T6TtE;QMe@jIe7V%ow z)gY`w?}&I4O^CJBVkhgnX@rq7eLv);BQcx!s4Hl7)l;jZM$EzrwwXH(iL!EIx zddsq!{^s{z;_o!IKukwcMVy=SCz@WwB2LiRH$2QbHYN^vDkCKEcH->k&U^ z(KQ1ZlBqxnQV`) zGU~!IhS-8R-yyjB982#-xx9hOImRz!eGSR8bm8*Pnm6~)-TKDXVAYzuTZuNsE})(D zzK5gECF9vb;Nj-b?z3Hl82`3Ano|EnM_ffM{zpNH=lIord=SLk7YaZmL?v0u9A4@$ zx!78q`G}A_k+*Bbs>=)i+delPpiC$;b~``$v;9Ex-WXq;w~ly{DwC60*=A0HmakB%jN#@pO`5BiVN*BpLY4{%`N57G zddtGLrm?>eJAWuG1Tz2uGob6o89*f+h3_5_xBo?#A};I2H=gW%Zs-23Z^-L!75^1z zr)4I|!Z%HL8~7$TIP+zv8Z)M=l$x74o;-@3hDG!9Sn0pl%kOO$S~P?Gpmg=6o?6#gfsmC>`qRkegXg*bu${Tqlz>c-;}vvEAB z&u0m%dmV*iM0kBiz=>Qz;UdKPi}&34Eqt>!i#(X!M=HG!U*=ETfMdqn$sD$sdpvB! zB7y9l#lPt*dEeGOF7t8U{aJU#fZ@z&Jcn<@r6`0W&^RC(?GT8^9*1RuqCVqdZS8@f z%F#gz6~}kTA&cS-YS{82B{Af;R-I|+(`PHEgv}dnfXcN*p#@ub0f}bhH`kmvvpfeD zIHHYfS|g|+q;M0wL+icq7^IZuX!&+~f66kKHHv?c!IcodCEbxYWTDcma=Rnae99P` zp{vYgO=;gDhr2rw(m!oZR=Sn&CXt0{_g!{-$M`ZlT~hZ`$gyBj>PuY5_9dx&_Zl zcc+mfFFpS{@N2Km9G&dSi~?QfC+hb3@izPSw>i9B$bC&K4DH$YJk*O9yBhoV+s!vW z9Ym4cu{eHi$-t*yl_i1J4L2d>GjXCjvAlut8tX5lg@LiOG@F|nXgYpZHblm3q|Ac5 zj-bwMsg!3}B4FFMIzB`ta98Clyke(1#Z9|u|IG-4%N7Nr?8{rDkTH1i9xA2XAWgxG z06sspev8A6cf}X8=tS6drZ|uQ>pLQ9ZJUSEey)yQzeS)T?8*s)2Sf^%qh9a4){Ov> zvSO*Dg#DhcHkV0=h(>> ztnha@MDjkFjd>mQVCkOs!c4x(dia>vcel7(HYD@mYdiui5ZQ)H6BRO_oV*c@$9C>5 zjJtC(B3H>VZ|(Q*&of=h$1AQqP4L+0LPoi~kFOjq;Zu3`GOvsWGIsaYkl(45=~iUm z$yBWx&xu*&LxFL3qgpZA=)pzbRjl3H3-SUHyafpsQ%6s^kJZi3r++ROENW;5;8U4V zLJQaBr_}TZ5-J4n5DSy@iBcmo0cySY=hJEX00SU59#2lX;YxA-LD~tU%N60J zDr44hanp@7oE3=TL;lKpGopdGsZ8NCQ@iFv<3hf-ZtxR#-FulP=sHWDQubo)N*orD zcZmg_9G|;%w>;*WnQ1oHHI0^Vd% zKhW`G74J1Wt{XqI)!1r682xA#U^zMxRmCygF0#wS6ltYyMC+cpZVC0P>ruV%Y<7hlw(nUVL_7W4S4&A~VV)}>{hTd}MH8v(r4u%qgRzO7(DDg? zTM`8a{(H6(l6N+;YZy;5p4#Q|zV(N~20&uX833$zil$<5ONf(%vXGWU#wNAic=*xg z<*>nkChPuaU?|$aQ2IAyPX~)wTm&HFCuD7nnF*5bf8bNC$P2VS6A^ka^w`+yg)f$? z6CR6?`RDH+zfoR&K`SGHxAf&i3JMV?niEv@rCSWpj=jjx(>H2JF`hI01(ks}S{}c>dZqN zx@D01=CaLVEz9E0H{7)^m!%Sxb56oz$XH!GG}{T6Jdg)V9@EomdYW$t_?4gun?1=K zl`&+l`483}6oiM97Ut>NpJ!e>-9j{W22E{JUkgZ$s7=hhf0Rqz{k~Q6L}Ps9nO(f? zj+${-!@lOYA$BAUVTXE|Ep3>IT%S@UdwLj^Pu)B9vq9S>=R3q#WM+lPgdc* zwky@$w$1&T;;0hloRbinlhX7>EQ4<S9LO+Si_>;b?y1Sjpe__*!J&Mo2`|(|p0?oy{K$ za4+_WM$EBOJYrWVY`dFh#jI{!o$~FCR%201*M&zhLf51uRQADUyNqeIAXajmW4rH< zqWWOgy%A*1=0IwHP?=g2Gd1XK7neKsIh1kAb1}%`O{;#Y`)68dt~07f-3%k1QDsKB z!F$O9Gvy)&z4=3}Z#5)`552bp4z!hQ53Lc* zseq_FR)T-@5=QUke`TEGB=sOu`|OlGl-kTI@S0}WyqpEoR0pN!Oc`Ji^!^$3?MUOy zcLt5xjJp+bcawC!sXTbh$fBnalc^oG7J;AlrTK|=#Z%)Bm;F@tkj=B6LebSSDYX@$ z)VP)z>Xn`j!xQEg{ywb+&dvSvC+g$fV_-{OVMv&o&rd8YvYL`K-irAm@F+t&{;Bq& z@>v-_i(ZT)dmH4d_3Ua1^KmT;;fsU_K40S)cTyCGB+2hEOwnATmj=oAEJ<29?jGe?_zw zpr84(qfFKm4Z9_3?dMbiNA~HQ1u0n0b=Yb{KNovsW&M56LNuZaGu46$V?kEzA4h9! z)rEgAFELc5EP*=_F;JT9a=*aS<&dp~2rAOqa8- z&9gOl!6S|p{lzK$m2 zmzoMJjBFbj3qLIXWM5ycc`pz%x=!&U4el&)I5jm&{KbDO-H^17SrXjdtZau+qK$G zD}9XsFBb-mu&=J*2z#UQA7C8gt+4(LWKbd&&I6Wu+!SW`pDK@?M*m78Ep-X`vRR^} zqBD&k^&rq)jQ`|;-9exY{Zo_XZT*M;Ao?Kw=3$!k9PjjP^2w{-$;vduSIxWo`^dz= zAbRkhU4UUuAEEt(rEH^C_s^d({?oU&xczejkv3dVIZESeI3r8@$`v+Y_h8lXDB9{f zwC`uBG9gQEnvN~A%_~uOUW)xAW5kPDeYh|?HoB>0<8QVtd*XJdI+gE6qS-6^HR$@L z8kaM1GGuYt8GI2lxNi5-lxBrT=LYBmOf&DT<=q;9D3>FhpOC6wA+=Hr3<8wT@^IP; zARm|4R5G`e(26iYVN=!qTkIJHnc<2myHFSwd?!n>)CUuT_ZoYEsaUa{I(!E?nQr(!>s__?D$aGZMW)US@jadyz_Zq^ACCsC7H~dWU(Tr0 zHnH8eAD{EJkb_&zG=JhY*%IEbm=lso>o!PDp1Rb9N9o#zt?*(X|4^Nm=xg1PKZM;f zYI{68;L4L>8K=E#rRYA#Pdunk=ChreibQ6L_0;v0;sL^m1)Pqb=$NJus=_i>7N*A) zT2e$ysyFR?z*IFK-ji~=RPjd(;OZZs?gVHmUAh zwL|j+)$x3#7fdrY_N+vz=~;QhfDS`e8zL2D!L}h%l*A~G|JPBdG#D!O*oYumEiX=ZfAmDuo%dey$Z0G(a74eS!id+BxB6kpDt7yC3>&q?Il)_1C& zFl4A!y5I<_BmGsk>s3jL4AAQtO}Ck|;^YRVnf(3Y z$@$8dlcU1y(81<$>z4?KT>{dgx-=+d|p?enky{aS*Yr=i|j5cI6BK% z?`u&-`l!9yEuic68CF6<>3nG#p_e@=O))VsAGPR3WCgeKR|?z-(7kz;+2GmFkU&J3tsvli6*%FC#OQZt~+wK-d`z=tn7y2b6~~aR0!IhGWP0oLp5+ZDe?u9?!SS==R0$85tR%j(MF9rXBP-c`r&S z_Uc`{ETznvpOd?{K9Do1+@t{a@Z0$E$+G$t-jB;27&j?+uRK)WxZO$3Uh&5-{GLRT z-&LCaGtHBn|2)$5?@JB;^C{4r0pJ0`U8AGwFxYrk7i38YytNYtJsSHu72`Iedn!vW z?Vk_7{fxQ(B`-8IwELpTdDU3?nBVadT7STV=$fyHYWAFnPs8!z?ZT;={jMJY<|w|I&Kp`59r+Ma>UW0|%#O(-{4>n0V={ z1oS5`l+E#-0Yf#xn7>@;P5(!7a)l1-sbj%Q&UrVS{i3=L1E$Ec835e zJDZ)zUqwX)ALbYj?BY8+8=DHU9~6ft^sn#WKQ`@%CP^V)H*bI+~j^fDrQ7 zXC=_UbIQtGCo0^&e*J1}YOe~u~=s`D(unSnMlJNx2NtMet{npBC%#M-1f#1* z%c%1{lshg2@(8M-p;0ty?c1cM`rW82?)=2Cpk}pbLi*y_Yq=FYizu4a{SyFViJD&!CM6QnHZX>wdN!mkIH$ zyQtaicD&a&Fp!#@e9ggqCP>tvEwXR9mn?98++`h^k z>#|nT6il7Tx4pf+Ubfvn0CQPQi8A-`_Qv%492a+9RC&taSR)0d^|s&Svphj5Qgf6H ziu9I;?WJhu`3CYE0epRI?om?AQG$_>KSo4gBsha_0PM5q4ZTvP-THX>J1uc>?>=~2 zb2Aa2^Nx3_A1RML@CTp_{7s>HKf*+^-v0jU`(-P6Ez(M$_GuvB-r(Nq2m@YAuna#T zE^hL5Wuz)4bUxO_{2M5!N+lhg#`vlbEbSzO%z5F_(nVp1>#;q{JcFZz2S)pj0pABf zckb%AS5D;WwZgDs@N`Ju&k46T-`eBdyelZ)G~&S?-_Dr_Q_jd-{Is} z#Yda|nY8Cdw}bqdg!THfRL=K3&iBVN6!+<2Sv`xqWF0Q+<`_vDcUB71xt({?5cDc}dP8~o)m8}%0sd`!?SbDm-0St&S9-G)nE;qq*RtI=Ko83MXJk#_>3qd#GQ0YxzheszzgAaRIV( zbV}J@xa~##iH7!Jk~-sj>=)P%Vgdi8|CH&>_nu=t`slBcZ{XbsR>LrVs8@rD3B8a2 zb0On(d@a-twkRQs!PRw@Ma!L&f+^b8Xs+$O5sQwdCY2nmcjy%rQBkkqkf5MPP&PKU ziKB~RiyJY6x@Nmar_}WvN^G9c3x6={y;hA5-|RPkyN9sUWT*A<-N{D=#!7 z<8c&>XboqsF=gK$?pd_;5p!{IvCJtbD9FlM*-0g4*7=C;nFWs1oIT2nedOZIVV|9y zO|LWKk7GJ`5Z0oorsg7o4}?f>&%UQ;lnflidfq{y!sJ-A8g0x5J55G4+-uKg$vZ#~ zf3>&y)t~8T??dboo9veRuLZb#jg4i59`CJYrlq~6V2@e_zC|&Rhli)KD~T$25qK0x z(150ghQ_(++GvSEeo0EoT*-Vamcmy%2Zsrp4nAc~x7`MEj}r#Q3y(1oIxx>hC3vni5or@YN7wv7l1pd<1qak2(R~lFhS~*pUXysFq zQd0+ddt0M8P-$s2%*?XLm-6_;JRd&QGM297d7K|llsc@RO_QG=Y_%{E`_}_+ef;hQ znd4$|7&r-@SxwcZq@*C#iiTJBc6N5Grt95&Ux8Swsz|YC&B&%irXO9N;^gRBUuAIh zqpC?JP}1u8x9^fKw)0qM-T%`wj7Rs!GeeA7i6C8z6?A^S{{As+CZTBQH^Al|X7|BA z$oTpBxgIU1`ug|?B^4Jtj-l;7efsn&jPZ0c(1RXIsd5e2@8;?J zicg+CE^)v=<{J+#;N~Bpa%cbeC;yx@jsFt?$7@7Y_-dQ6s|gC$Z2YYGu-tJJG_hW! zip0gm?ds|RH|Sd4ov5S7IO=MT0p(1)lf*&1#7XN_qg|jnJvB8oJuMi6)#=sW-(OhB z?pvqBN+cN5GC+T3DU*(z)~eg5#tITFJxK z=MJq!WqMLFj+O`~?QZ5*>PsK5W%eyuassdRGUjII4Ore?k?X6k@GKMZPbjgpKZenN zycT=uu88CH=f#tR-((Dr*PB`m@`p^%*6NiXP=Zax#s`z z!@XQq1gN{5Y6s=8`8_7-v4+#3e#El-)7KXP;ji}g^{sV(Ch*%^i;S$2n$_JHbjMsq z$~%c0H~G$U;V9&6uN!*6gZK5~KOew&d-rz@0mCpU#Wp_Zd@?oBcTJb43g7W<7$kR@m3junSN>bOF z(aKqdQdZ-=^P2aq>NY!Uhl^o;V*PZ@VGL(XOGkq@?h-)JEA8ERRn#>HxhW&&1JdDI z#Ye*n#M$zQOmivcEG4GL%h595sj>H}q(?ZWl2&XEmm9m@NGEVu^ zZ>_hlFP?n&Hec-9YkH4w?w%^Aj%n#&DL&CiX-#rgsJlpV z@Uk5Be?r@&(rTJt&IIAa$L4{o^$u7B zOj;}GmI|(mrVe78gVk{fscQ3nF(qmz-$4B1&su(o{+m6>A`v$W@ce2J3HY%7vB#>n z3BL{=L*I^!$Y3b;K%lsQ0E>9)Vsfk`8x`%G=(UJLZ|1mxyasYl&VfC%$nfw5@^Pc$ z6>D!lZWWHU#+c8V<9O-M5~62l+$K7ze8jqPuBJAq_537=NhAM!_1&fU{bd^)iVBbN$B(0>9aca7 zSWvUrRa7+Y?yfwZ`C}E&8*_e1Qrq$H=aYm)qt_-BjM$oc{c0TrYaOB=+?DXm%RgrvaN%MPCmeh#*K`O6naccOWW&> z;kI#n>XY!18|xd+$u|HO$XBW!zQ6X?y+QfOkIMpS7@g~mO)uX}iax+NBJIF%>;7ba zBv{J+iTx7B!V3o4mJ0v@S|!d!awoHoi2PSuQh&?;y|A#bpg@HlFR_c@;&%l_MM8fO z-@1EC0H{dQMDk)jq<@1{@~Wv`#=S0Z;VHpg2!2vx{e1sF-K422=AYU9^)CPc{)&Qr zy-;QKKg&l2{8Jz0@8>L|8GA7?smv~C2m(r(srg4II(W+^77&oz;nnPv%n8y zmoxbVb3~pZMJi>jOhCHadE>Wh|775iS+kC?jr!}_|CbIQu*>a$W9HI9^Xn(HnHd?@ z+%Xg)eh38OW&j5p+mI=HIV>(FW^9?Z#uMAu+ncaaftZL$0)#~^E&d|9=H}+w+VM@j z4VAgh<>Omh)*q@tI?ndm__ATTM?dQJMaQ}LVi{9enk}B%-YGZ_)LqEsiD3++&KT0_ z+?@8Q&aTrA+*pt_7fic<>B#cr*O_!>w}CA__RF)|_WwXLPk#Vt<{H3iumCjUw{hFD zUXJ36hgQk6UKI>YLil;hMYK0&t_ecm-gNm7(;%f^$xzG`(%YVEyI!NL33IR1D@m^%N0^ZPv3Y={QrndjQcC`Gp)iAn*^ zH_FNv0rE76y{~bQ^1Hqf+S=asufH^GKRj#AcQ_LawLaY)5CA_6XgG`kSSXA^Ei^8U z2@0Yeb(Ggpt;3aYAf z(EHTVBssT%=>Sv>0n&E?*$2YdBErI_;73Rh7*FLu{ZH4qRL?w-^9mm*o+oiReu1P>e;-&`9G<;#NJ%(J zG)Lx&N=sb;I@8zJ2Se=xY~!Fs^DGY_L6gOZ??pux!`~O%I40<6X{}1st5eB&>?MV| zdwPPue0e_&XO2d&a&YVyr+Io_x6{r)+@9ZWkoxP_F1q=4`~q+^K0Mj(c=F`Qd|T9j zVF!TXva+%OrU(!^5C*5cuyaF!;JWDBH!YBGCnqO^C{8x27oPu|D-1!eS#c6!7>h~P z^Bb=Ln9foIh~v(WQ0M54A6Rbu6Dd+rLP=?93m^@an+{XA7MGT8$6_kG%gM*bcM8%y z6IIp7BUvQDD=DefASWH21ZS=Mq00>!qGV3B=0yOua@kuMdZLtdK;`hOg5lrkt$6p8 zhV~eAeZ6F!EVXBG|Jb5zv!4H0CYm!hBV(yEUJx#>0&*B(VPR)yXC0j>kml^Ijg=a9 zkQF3lB`4PzNn0cv z)VfO;|MX;DK$!y1SHCMk_#&Y~p>VQTQ_xo9Wt!)69Mn`#0ccw1i=WxRF=|%zKuY$1b%$ZSR*bKbYHv=Eyy&QPm6D(Wepgvscbit1b%e{R7;y8Y4;L_slN5cf^6zS=YsiA=U$yR%WqTtMvlGoCGR=kMS2E& z2F(3%wkMrVUisljfl1hr)ddKoleanDkXdu&%I_~CucVX``t4df0D??SO;=Wi@?HF> zpdgE_rqsvtoecw++tkY9;$r2=EYlV8`}gmA9QDNUA54BgRU9o5hJ-MvpjiQ{(z zN#?|k`tOzc=QeZKH$b>OK*lZ5@?4yF*Rq(^N^nIVEqK zW~0LK-OXys@AvH%L_dJ4HpQ_kz{vU*ZHfHn0P?VspCLKs+5O9r)Ya9k4(199QAvhB zgy!bvf_(=NRRG}v#|6BwTubBS`SvId8NQ7}S0oY%Ln#Y^1PGWk8jXH$C@v-jjz3PT zz3^}Xq!KYn1hcLudi)L8D;wisZo|b^)4)&CK=;QSCKls=Q?=f}{S6kF41nWZRZ$VZ zB(a}A-!*bO*=#u9MR|B5n)D*MY*GL^qh?XApFh-oVEsUtIjW64s?EG8?o0*53Rz$u zN3{VL_Jn~!@xzTfBzJ@I@}99hb3L%!z5Ltf-0ja80~?p;Wg?RgVPpW-+c=xq?c?Ubgni`h zW{*z&S6f9}XYdGjNuPEZmZSw-jr>xx!1C`yJ?>MlSI4vMEtCVb@-vfnmWj=;chEq+ z(Sv1;t?V*}$@%WC3eplnI*FNf9u7{aMLep+T!F?h!SnP9eWk7w)eUbz@t(;mrEs!`42(|A>WTwff&z=fMAT8c|n#X!T=w`esw zd+DpxGD;k+=cpi3#+cnQ!q*?N?_86iw7w~DT-NPc&6RsC9YPzVstcZuMq=1j-@h6) z;t;1I@q1DPo3!X@?+0Pt2mbEqPd(%X7Mm`kyq`A;&H&XC#W!yVj3lI`rDbGL>FFw^ z_GaG>B<(gPU)YJ~Ux259V7JYvrJ+#DIDj=pBqb#w!~ljIXpW7I{qDA~9tbepD=S6g z09=>xDh3>)uY`1TbSk8De&3ydw|P}D9TIxG=IO33Ora2AYB3R`E$Q1=C+SCaD!aT# zFBmD(b#PMgVN4NvFO@|j%38%!`%P(TVypY(vConBBfe^jJ;_&g?^l6Bxl$lpK`JWs zxBU~0ZP^{1SD;_;znXLk-S*8fdeFvfX!gCImC@N&NlzrIR99-bkfo8gB|PGO1a7J? zE~IE;dDo@4cwl||s2o-qD4Gs{kKfM)n3vdYZ3$)!=v-)7nT>%R4=$YSeC!qT-|Sk# z?^uCd1G@!JP(oTSv}1i+^f*d`9d%$$A*x60)v%C2-GIkl9ChOEk;2SS!#ApkGKe%FCJ2HHJ-RzSJ0rltmRRB4YT&;kMiCUsCa z0w=G9tgWqkrRZ3mJrk;h7bqzx%*@POMT3nbHD<}e&aNbMcF-UY%Xu(f?t;&t8Upf| zMvA~__kV1J;KYY#Uq&XJDATv=iE7}%y6z3Nt^hl+7uqX{6ICe zT{u-xC}nRkP`j4l6q@O)gC!z*N1)ejXG9WTGSP^=$n?qDC&^i!Tzsgn$>!cJrk(V1WmaZZd&X{FNcA7n?{v~w+0gDmli+MA5`9Hk^&)_Rj5+SU*P)R(fXY`lmm`$+lzcqW;H zy?tzZ8Q>MqPb<&y2We~uvNcd;TRsSS_2Sq5x5@d{$6U9)sbd$GmgGKw3OjdsRhq8~;IOSeH$7C(|lP^X|0sN5=V#hF`v z>8}6Ld6~AGp9#1J%rFe43n&+Z*0$?2OT(KA8JdAc@<|E2W2{*Q@&m0 zS;%Dnef?gLf}jMNe^Cuq2QeHQ&z|;(E17V2Og-(qsjC{z-_LYbM69%)7Z%DDP(D!B zwYqPFyP6bkHDQg`;;&qv@XrLXRc?h@!_RCA=+v&ugt>ZtZn>k|NI>%%93qR)&O(Eh zmF0qNS+&fyt(M~OnDiH|8rgFAHue@L^n~B~HBv1Eab#wQx7dpp)o8or_a-gxqb}co zKp^i9Kw8nu#|)AuWd((h@bI(%tzQI*x0+&XD@gL@TO;IwDQ_Rzua6_w%D4>MNeUDb zyXawMjvzV*>{97ZPhyZH0UyO~Hmp}P(q^y?Sg&Cjj$il1!1$Ro)HA+n2aL~GCx9ta zb2@9=8+afdf!$-Pqd_tMM>XMe&|@B)bW#b@K1d`YA)FaY4#iFfzHs(|>ZQQE7uMTl;X8wVs-)FG{>Z5~8oC*DtFqcn%w&kuP+%jCWf4q0CcrU#N29#uR%ygj4l)>ehgT4z>a5fT2=ogR;a8gdZ3f!1fUUOXKDv@Qati!zjy z4IEetxpOKwK!C74-{G;8p4tl$nrVK5Dg(rdEgEHR2WEb2wq+dZ6>J{ml@E~avUlW;;a5&#-?}}AwC+2uLm_m$#a7l$zC6}=kDy6Vho8Y7 zrqTFlo?Q3Ook*2GOrGD1;jiSdveG_g27>Rx<7#* zl#~PuUG|vv=NL72u(h=fJU~DwV2(epMYZ*Jc0P4I*Z>rh-Eu47h_4+c(n^U-2?@Cd_; z_IATAEz?y`!QN9lTve_RX;y?dsI{b^@JsB+^Nj3M2SbcgOAspd+x@NN%=FXRmwQ1(jEqA&eew^ z?KCyNw$0jjJ^p9mU+9BD8y5_N1U0yT>TxvunoItWEB7^6YH0LDbM?7HX}9_zy1c8s=obrx%S&ws=iqcTJwMj4|nr%0dM9}DLFt4shdw)<^Pr@`F)|F9h%zUj#r9yxq<>BLhy!UN7 zLHn=&V5kqvI9=D2WLy&XUkEj?;J$Ml3QIPC=9TJtJT$jsGa&Hve`+ueFKI^64Ar+9 zA;pFKB!uE@usIAy^-6~4jJ-RBW}>!)W_lv&F>mb~er z3-mQ`6BDcrYo#`QQ7y}e9^npz0EemxT|Rlb<=#2k$ZgB}Jh&gJ{J4$CytSBecS*@X z&99%8-Psld5eYPLb}OU@A;0JEI_@8h*xTxaVcEpc5HlP)Q|p6Kc86EiLSHmbTvqho z9W;Bg2}qQ)!^j<+x+#5V7%N>W%Wboa6+-%Z=Y^1K{e+nEWBh@z#QQE)!j^k*=wLA;ip^WfwTRcsomr8Jok;$kXNEnTZ}SgwTm_)%aZ(+7B0ux zbPa_4!7-2m=Ktu`+B=;`LY0#xB#~eV;ThOqnZya$#!6r=ohywaUG?)v8;A z>3nP>R+^Q%zI^I_kUzsnODSEct2Eo-tixIU87ahgk4(lc_4M<=60tt_=L8v_o`yb7 zlaW!@aCPlMIwiK9r^p%a($JHC_Jq&^d`|5^tq?S^lCAEbTEke-X1Y zkea;L)W6=rNxX<_7Sc}x<=gwl3g5Xd^>JW*dB=XAvo5WR8a_=_j%19Rw@o={7z8?a z#6HnLnT|%pH{X6*w1|UAB}{0YH6dx?{*-MOTXy|@6jb4!U2YTlz`)|TYV@}Qa8}eK zvh&zudKpuvk_2!XB$9!Vk&&ODzqjPJ9DVw4_yOZKo_S@FD04*fPa7WBHJFUm!6;ir z3QJUWU5}Fib?hbZ?C7yV=OVKz176a$*gXRomp@uR0(EB<6xnY}{pbId_Bwxd+9bVo z6^sbUrpUAWSLt`1_OCtXH+m2$iNgK9W~2X^4~(}jeEv`!lk|D0(r16N1eV;TJZtQKA99v=3wYh1VaU~iFk73f5TY{G_lo7dXf4|))~n&d0YT5$NtWd}YdxJHJ>I9+I~4 zxdkJ6S6A0Gn@dNXfOwxXlaUnjyCBm6t)kmbl-Ly>5<>7V)`yjijpu==h_&+tM^Oc+ zh@1b&M%96AG^dIM3xKG(;X> z_^Z?AHMyS-+QEnYi6}ZS2u}Xdqw*4k2`0{G;{HGNs2GBtRDccouXY_vU*YG@UFpiX z7-d0D9*pR;^mK8dFJHd=qI)VQP&5)?m5-`$=rDG803*co2Iy%4UqQd|;X+S)!CAQr z0nzQ|?atJ@7Oq_Lkx|n5^A{q8X?2X}mnrpuUh-f1gK5xIyh%T3n3$NTsYQW|hDaq4 z{0bnKE~D>n=)qeTJ{Q(GcS!R_vvt z+WN5n8qnM6U$pwaUi<$M+yadLIYC-x=JL`KNOA-G{ju#B|8)r<_#sV5%CZ3JsRK`J zVlv0ccbTXP?h6&^)6T+GjiFH`;wAdK*(3qxsre(ek3~{@a86ci3`dzaf`2=Ki;MmI zMc`S|mEs`zTunv6ytZHcgZzR&R}psp-^P7m4*zcdi&*%AUnr7ORO}lX;@2@8mMqBtyIaCpYBs=R!@#Y6=lAOIP0a&iLleggmc zo`}<`cUOQaH7h$CWcuW6@LqsavPQxrcHkx^uOuW$?lgSFybmBibZKcI3D-SR(wMJb zDWPP6so%dV$c6x!<<&U-0@bLI!(H3jqD&oaQ1x{Q3MOyL1h`S~%itt5)W?`GxXTZ{yBlkG#D6sr$JJ zsvPfT2IvTx%StYYd*oY3#>P%TvYMq@K*z~xGuL{9-1R9tJNsRrS$!2s?*v$@D?#9m zfEeNjx9Uj+Vywkf;hBZ5L?C5+8X-x{ZhDh!6`;XmV8}oR>j+e**F;>-4(BypH|r+N zfiCJ4<-VHZ| zaoNm)+X0E5@C@KQU})QPL~EQtWll~G5@`n1fG+!MF&q{^J(x+4*V1I54n_^~Nhzux zz}|o(O;74z5?_B$f3hMfV4IQn1yv&oAQilzH!oNmU63Gvb%22hAwQYCAuB6Oy$pmrz0P?Ttlx!LZ^P$F_8ZXe zLbH*xJ)l7CI6sj=^_UIi0oAQqp{a_!eIZdaAl}|z!vcAU5(?&aI2-0GCNF<+4;%<) z7OQ;=L($3vz>4BWjNzuH3d>D@gs#=VPJvc$yI=rT3d9%yKWf){19l(=RX`-0`OGKm zHpPDv;shtG@zUSOM}g+=-AiUj$!0w<1D-D+c#F(6EQ^46mby3!ZKm(;?tURmy?y-y zdTIc~{bS?fy=ki2P>s|Z>z}EmrG=Mnm_9=g8=xAxdT2ZHkf|)+8V{*5;K9Ck* zzg9ATeW^)E58%j9ouUzk)nU8qx|cEDzS9(_n17l(2H4QCLnqi-% z^}!9>A}OJOR;bQZ%qq^VUzrP&Z^cZSqXIC!maIm2bDFIxcORZsR_#-IT zpgYH>?=}rPV#gAr(DQQ~iW(XWP+-$Icb06q?YccZJy{5T$vghqsWZVXV$Ik_o!kpt zZk6_kgwI()Lt~{Duyd^+Zs7Fy_WF_Wp+$pRX7hjm92@}MiL`ftS{*=6z-{mh9T+}u zvv2xB34m0rfDu!6i0nV%%|CZ|Yhi2@8X_zz>MI7mJ3)JthJ<&J0h&@MDL;R^ma24& zj8`}Z~Rg>LrY-h$bQ>BfqDL{j2_Llqr<6i?*Jg zp&3g@3@x;rac){V1{DEjK0P!046WEB#k4NW_eR zbiXl_Zh4PegA{(=+G{22S?sJPtS0KF^?4`cM7t>wN}+>Bt7RgOiB$u3ma+fJy)m9% z#cuocZ3`IwJuG(9!5kl-TLp$V{{H?@9rh?0Q@PBxURr0<;-TT;46-z8Dr#zgU__&{ zR(DkBLjlEdAqEBs=xW#{Fah3-5$K2A1TF^)pbBVfZ!HN@$-)H5rBL9uy*A*+k3vY2 zI!GxE7hhLX%%WxHqeNw=CnTK&-*lVV!Q}8GVrF~9*am7J42|`7+Ui*xdQN>Hk2Otz zEDFolx3-Cq-Y5P#x$88-P2Ajbi(}o1N!R64Joo2Ls{Rgj6`TAV9hNYe1?HyhveB#KbATrFO~lDn4=#inv` z5OtgW8FPbw)I)Fj?dYq&hKKPJ4h$oGL?{Ko{0TB5&S(VqUdKMdg*O&VKyXZUm(g$D z3=W5bAKpg+z*4AI?^#`K?Nw=k-MJ_zNHRC345{1LDFlVUo217;4(>|KMJNn$PORfq z&+HoDP!@sPC5lLbQJOgnR&Z{-L->gibGQx``-OB7`PcWuCT15DmRlMB4{Pro6;&2> z52Ci(t$?B^NNh?`i3*aVN)aUIoFwNAMMg!X5CIVo$s(B|NX|jYS#r*T${uR}L2Q&6=_HsT$ho%?7@NR5SthNaR-%w=Pi zYg6B&+=dz^f+?i;f+9v-qiJsG>!~;0j;EQ)*a$?U5l!j#){$(R@|9wiLTZWz!+Fmz z9-dj!6sMIU6z#l%fBGLUeHK7oZaw&yG2@lan8p$##Yue z-R`YRlP{Hema=)(&n~|EH$eE47f1BCATA-{E$9Rjf31xd!Rop`s}RL*wr)b?0@cEi z*3R8{(U$p+cQVZlDlBltb>w6T7J+W|Z@pdb_#!-P$ z^O(RLwZtwF_jpW^LDRGC{_ZnF9qbx9s-;WQJsR3SoJX#`>k!@v^H|{hmsp~EoBSND zA&!XqEm)>oouP2(z(;WZ0gjw;r)d0{LX_ixVOV&pQDp1pdY8=E`8*u|#eMY=QSm}% z>1z@FtXrc{j7Q`+Ee*nk^v5#*gqdDKl>eKfa%wP3k-D|&Z&E*k6H@#}$-U%~f zQ;%)mTM1AhRJIz`D%hN&%qkilN%s1SCqsZYUKUxf_M}Vi}ZRp$sZBGc(G{%2tDUkC2%~ldpU)-%1<# z7$1)L1!aI6n2t7YIXyN!&m&O};!LKGHuN2da!A1DRFu5x(~qSgPE5ANVN z0RZ~fJM)oEI2P<^HFbMLfvj8ze)IdJ-YGZn!BTa~Ht7v#s;k!K+Nm$G{{S06wuFL&jy_&Zz zk>2VlSBFpwfKDpU?Fp!?S&ft^5ge|DRhB_`^sK%TA|8K{DG1nL)r152Ey6tW8srE- zQH#4R{$YrzaJaL(dmTFXm6)wHv48e9 zvpJ@gd5NE^BtL^U&igoDY5PDxxZ_66X2o@?UZRLQqrjbXb?UVfRP z{%7}if^C!Cj{Z1c;R$fHCbJ3<(jn=FDpnxjoyz^y#x^EKM$K|pM-ZR^K|@nh)5m}eFz;S!b-9>7F-c{Q^7V*8h%BZ_zDF%VJDtnIMgyT?w}3Iwd= z-tyDCFtGNPOR*%C+rl6+wX@3^uuvx9gABwru;ZpoOquJ38Ky^FPL7C@lHaK1z4k4L zmk($^-F-haHJk4}G(}A0FM_%K>Zo~=NkQ?;O7Ex>W_{;?Oqk_bd?iN^vpjNG8_-^N zZK4j4&_oq9ySa~2(qcIasm@#3xw=*UBMk>HThQ~dDj|_bZSB!(Q7N~@)o}+_A~5Qi zGN$$C!LN2Uus*(Pj##*{g|9*$-Tr;((_?S6fuZ5XECnd%7D-q;;WV73ng=OX?qVl#kdTrUB1IF&lk{ci5b{c$>i!?_>yy1&Pa)n$yr=A! zyysQ3GjC_I$54Q7h-C>^~zA4_`{iI5g+85zwF@d%hY3EDE!Z{LK+kQ*_L|rGnLYU-b;nb+jqZQlC?))h8uLGK`Y~$`{LGclJkv z!h_dKvN)O`O^6Lk-yQw7Ge;Lc6SdXl1Og6HsDU|(X_ys(@&gW%Yw zHKqK(kFpYir5TlI|5xy5zokfgKqa9ItiMb`LNbqZN~paS_a(>+9w>JvJ@`wBo7Ty8jA<`Nt-UEpo&WF#2y|El&)>3VBC(TVh7%;xH0yE8@1rl+6qwKI9n>BwQ8GVKYo^q#ohE-VL|FL* zjlt75r{BIY1PTA6o4JRF+np?D{ycF~Xnezcy`#QmCkD$MyR+(16W+-CeV73;vPqaI z^mstTLsu3lR@QhCKC;gW^so`@agMFi;y%3rB?LN|Dds$SgaPL+_!c^Gs1D7WMKr5^ z#Q&$y?;pK%N?>GUY#F1CiqX!ZhK&dw&vT#!szK5>y#*ywZliavRXdtR=XZ7PeH|5w zjam12W2uxR6dq_PnCA7KA^uMlbbs_m6z;z-B4ow z)I)k10~%#FYlsJWO&i6FAI5jQf>=GZ^YgURGT)Clfx8U}oEp6I!&FW{mZPPt z4Wy3=1569q3%%Cz|7k2<)gyu>F__X-E_!=2@&?J%8?NeuX7F>y^Hq6rSK*V6v(3$GpGUg}ju>H)@jzJp`?+5)APnJOBqaVf zt0*pAzDjYKB;0Dr)y>lnZi2~0o0V_G0^8uCDDmWZcX%CivrH8mYn}iXx)oMIF%9C^4bW)#)+4MN#~B< zO*jl}QSBF%xuOi5?NV5prX33$*%W&zf4K4=GW~nMs-~a(>;T(P$DKpw3G-fQzx=LA zEd9UNqUqFq_&;G%JezD?9Pa-$c0J**k2k#oa=+PeM@uJv)?G)@{QJf;uW+9{+rBN= zmp~h^hPpzUW)TfoVsubYka|Yj*njQ==|$Yde-M29nFW3hGfxVI_Ly0A6+(JAyW#59 z_iwnWd%wT_m`L#Zi)^CIoR`_W`tJT6o9Wp+MvGU(2q?U-n1Mk;<@poR)e9SGiu7J8 zC!)@9(&ORtOp7uI>Ye~aLHre=lRwW07(5mmzk2$~gL=0oU1Y4}I*`EcmHvSjcA8ec zuAM@&WznhhCo2a)Flcl#NbH6N`?oI%2^ktDLJh+AF-AvoF@!p!>~hr`-E7rGCtm6- zo0Tb{t1mCYXa6|?FV6k^5`Mgw7e`O&ri}AfnP#uh{=R~rzsHa7|G(#+Orig=T{}J- za2=<~`p>~7y?{N*d>R-fgg)1~xz#X~`wIhKze1&gh>D8p>Q#7u z$XHrjTzC9x!p6eF^5DS(Zvl3^arfE0-X7|Mc})?U;wVeYUP1-g;L*Fzjwikg8ymaH z$jG4L)i6CfrB}z(-G{7^_kO{bA>)&nelliUGUPJkrWfYqwz;IRH8y=H1f^Sv*Q2nrqn z0{#;d6NMvALm!3xPM$o;EO`wiZSX*}JN;$*6lZOVqtyvw$yxp=8kQ!SaH@S%we%|P zncEd1v|VYhU8&w}Tx-sa{-ygT{HN^a@>euZ$*%}C>{MuV)qDDgU9)+Nza%zL30y~9 zj)VX{JFmyCLwtv-oZQOc9u@@7#b7D{HX1*sF#(=RML;Jd!ho3@*AMRZHBRo`-AaIck`|rg1q4)dKePlqYa&->1H`5)la%Wk^#{HZhr?> zOd1>LUUGq7)14v*#T!p1m#>WKd9$Er+29}f^tf6-kiP(cYDk$iQ6r9?O#kLJFpR2X z6E{&OjNA3HRT;1T#GPHUnt@GQ;k=?pG*D2W!we)A!NJ#Y?9L4o(|-mn9^T6| znlf=U6_v2n`^gKX=drLs2!8{at5snsB%cc#g8)1>wzs2-HOc&qYS2U#ZJ@;k^#Q=7 za67MEy9W6azwJ~*W@hH8vxHFl00BpGLc(imK>-??xT96ToI$7n$v@OUIwx102`RX2 zAhiJV;^%9Yjta=5Xvv53e^>12jGckVDV_s%(h2ue&sg4g!d;NakwFAqTg3>m|> zZw7|#tsg#o0PfxB2pKs!_%}dl<(aoPICThk?mLeWCY&oBsQ^0*XzudzGBn^(@XcuF zUp4bz9j56l#3g@~nmUVbdwVL-gItpy93%2Q_MAe+4YO6PDi54=_Hx?=7!WsuO+mqg zm;^dB6a+GLDl4{%;kZ-w<5`EvdOlo|;Qo(l$Pmg)`FCgI&yz6oDOr+pTJ2+x4y}n? zAe|2t1ChrbWTSq-6x>+ethp>R>^$XPkqibKI2#UZZ4X=*c%qzfe(><&LnQLAb~>mR zQM-H#f-=zPh*kMUFBAy+x7HhP#ezS|E>6?kDRF0Gx!8GSat~WMw$(kg!nZYSkJDdd z(E@h`MEUNkjkl7eA~(SjrLhkHfLK9jUAZyf2_7@y#()EaB*V=IR-+Bj)Ug5=1h#dt zUjMU?JXxM!H+=f^Dfn&-m$|G#iq{9SXBdm++F2Cxn4YE_RioUUdcni&z6F43RD=!Y zyHi1Y^A|q8Vo5o8XzQ-@)e zHN8?#eTRt&zNQo^PeKnLIsuyrWJ`r(o^J37;1ObowLe-DJW@c;w1fx0dGiC_PzaAV zRL$1gU7y2gjClBt_SeK?76E!`ZfNKSn+mAg#6ZUjuv?5(Xf+16g82`039JSHG6HWb zP+)f!kQ`*qv>lyG&f?-C00zvrZ_mPdt6J%K0Q-ZW)1vB_!~dAn`+9$hHX*-d*HhuH zr&do5)ca7!$pz07a_ce{X6ClO&n9flYB@+G4&4Cj4xE8Yd7%d+_W9^xucHH;w+vb` zimSt_P!L+Y(3`O)+=TP-FpSOsua!r>>8&yRL(_H{+U2{jsB*E$1L^C>SD_2O4Itfe z2a-R%KUAOS=;_a0xBy^PUubQKQF{y}1;ze6CW~34K+-qF62OzjV2~+E`o0)(?9;+1 zXY4Lnd-1-J==l2e17kUd>$*u|B8WRuQpAwB%wbwdCWfy#FV6w?OPDT;m+9_>OERcr zDsE=&H_CXf##0}dz-}7ru^Ik$1(xw$FbBAaz=>kADi2+WilMRzG`+oahM5&-98-;G zrz6nQ0T9F6{wi9IJs_(&f9@P?vlvj`K#Ib!5PYzVs`Lv8xEcg$SV2k2F6gg-tM>P= zt>5Mv>x&u@+5p>ihu(MATPUXb7s7_mMG^;i82x*84l@C#fMfij!d~eU`%71@Y``32 zRQ+^+;D;?>Zn*L$IJT(C$=MFtk63N7gTcx|Mj>EueSIm&8X$F7Ph2EQ?%DwnlV1b> z*6@_ZNhxF|6d{SYZN~4%vXqvTfR-x`ELTjsk~T*7SAhR|IS30j4i&A`!r!1ScrmEl zZL6TL(60d}tNiv(!)mU)r5`GRNz&1-ox5Auu}Uygg7PM$AfmW^%8uwub+w1v6HtYM z>|z7zfOF9w6+}QKIDc;t>bDBCGD+zxNo7^LyRM1T(0ZaEs8Yg7QixhkUEurb_U6l% zFW`a&!Z2^@{i;4v3X1faz5=U}G)Pc-=g}K-87Eow%sR@ zt02LJ3nWB_+$==gU_o<7tWP8MVt=-NJ<{OI^S?ymv)iW!OB~I$w4%{q*1_N}!?#i+ zbb+7EsD(6No!=OG&VD9p-S9XFAAd_?5OCgK+_p6OkolYW%xvD*6k^LKo_HQIcg#D# zba&Vb>38Pyl=m&zIu%Chcls5U_jC+y-4Q%aSD4QT3CJe_x?YZjr4N2URI>QCC|3`*;!|(gdBnW1}JZXon+t5f#=a&^o7er8Opps(iDh38avwas=$Az zytnZJNF?HWAizsP=I@$VSV)SC%kSd5CuYKiP9K*C8a85LVsryo-^`!4po0O@dS$8C^M=U11N{bbPi z_Q#JOIQBLys35lFk3V;OX-{ZqZn|{7zlQoF;){jlyQYur&suFe=%!ffb+Y7Bojlj` zDly05IG1l+`HCSm>{&QFCzvNdTfT|FQF?dhhC4cJ;le9dn;5S8(ZW8_@#O_j8*tVA z-E}#6*;C;oVYe;^!GhPdRfLL?*ph-!9FFq6KTe+}_)*&yeh+6R1@qEPC0#WrGLiuS zs;AkJ2?(%orf?9?gRh8_(KlwcyRY!RyxCQ^E9V&#!zyQ%NgDSJ zL@M`MH(B(**NY4IAa(+*6+#C9cCk+J9osiS&~aO0 zANv!kF1QkXR-m1gz5N=Vge%Yp78Vw~HUM7SY4?f&nv1-=JQg_!@4I~IlCl@t|kf6nnl ziio^w0Ijpvil5H<*ge2f!(bW5xg`M$02g-vsUIR2X|Rlm75jx!EamVR1wKoDeUn1SgV=_4;v}(hs+SRNo2HIuis_l>oJF2!W zhxG_Oa}0)llr)q}Jf7gsKk;)=`Gl@4$9Lqr3GBkK^zA}f<}vsY;T+Fmg@C!%q(j&od#QWOKpoR%cY3wYwAmW69|WhJ`~LuH{+frKNiiwfS_-EmA&p1t+pwGBJxQ z&Nx{ogkrclHcsV&-~#_7lxPXhzj6CUT1K(BZ;0Wy(KlR^OfT(=8KvWj1KvQHhq^ zBNso`S;C=7Rlz2qZ=oQcQ=}%4XMMVq&}~nRc(`5G@roO}`bfvA*xg;QfgoBbinc1W z|LS)<>$y-u&Qs0Z-HJqXjEw6`g|iB(swsT31)++YHH*QdU1s?>t&aa4Ly*{qp2Z1B z0c@1K8H{#^Xt%u7a=2&(%$I%^YdHzYIV6$YW*ac*=n;^8tVOd!JVFGfR#LG%k&%%e zwV=q1>aU-$AtoV--IPyJOaq3CD6%&T$71i;la2Mu^*TE23UK-eTO`pAB=mDD&aQ}e_k~i>CJwu^2`n9r z`ngHN7Wn?l%VtV0w4{^m!O;o93VTcN)pya+9w0M#)o`z!udAnr76I{gY0k;;oxRNp9hcSk4lc}oaIKPX+? zOU*JN7lsJ?)|ol!1<16az6?TYAOb0Vak47hiv?SW(lPuVSPWaExD{L`nocX0?XNZi zaSdz?NRT(cPpWtka^MO=FxI;gBo{9dEvIxm?TczsdImeXtRz|w-VIIP9$uXZ`r_r= zFgsTm*)KBIhe8-j^yJBNLi+h@Eph2D;a9PfnNJlKv4uz8U%n{xV;vKS)`_dhu<#9*C`2cY$bP0 z1Xw9D3~`M*bV-3_=>tQm zR21KWB1}@KxV~f7Dgj~vl&E=OyKo!Z8&v?qJ4k;m2J=iI`w|VL;)i4?r#z`vZ!R~~ z_y`~=s5(}RL5vSM#QtJ#JMb95D8sEY4OchESu{XgHu2>JlsOSO_TQ!bYCkTv(IIe- zY&uZ$g}57R#tuQv^V-jED96yRa{}tRRJ#1XhQSSj*u1~Zo0#ZBcgF)z5}WCli2QH< zebXfh9pbSeXSm`@WSViOG3K!P4f-)1x4ndu*8xeO56VL}XUk(cTiKXvbbI=2;@+VQ zQk`jfPCx4kk-fvq{A@m^EwX$ObF3NDU`6M2o*&v(RM%Z3hUEs&ZMMwvM6+@B*bZ6S zR0R>4Ql-RR3a0%HFnii{z`KM*{5}#w9|A*?2Jp2TE~J@(Oe5cI%PP1r z&t&2Yuq$u={r7$D1-RF$IR?Is2K1orUbV!zTvDQjCx(*N6uan-)9b$xIP?Myb`Jy> z_}Aa$3=RMT6nt)g{`sRJBkSo*3VW2k|fwhn)%hfjgCHmz{3#7GQ31awhJ=oKvI^6NZ8a% z%zEg-%Z)W0=V}TyyO^C{xY`qm=jw@##j8GKmLKYu=toyL1Bd7k4<7$z=C6!=hzPbf z9>wg-#1mX3QV^5Gb}a-fbK6Qv^mtA`aFYn_TObO$um2wXyw(5~W%|8WL!m5FU27pR zGqbPn(YuWcy?B0s{GNTb!0j)CHVHrc`a;+H$FCj-vUhL4{6TLBSzi(2&YdjqPbeq= z^_YmLC}5jDFJ5fSy$cChhbac7RG#Rz^>v4_VxWBl!EBhF1umj5Ic1)UL8jyPd6=(&L007^po}FYgL4jGObXk1i|*a1b>UwVfi+GZb4sedviH2N zm}wf_#h~#YQ(isc0-HmV8;7NsriBd?Hske?q`rlX6h*UX2G(*GMdZ538Z6@*^Akiw z-l

zXS4d*%Ca3w)I-M>bi9Y6h(ht%b==f()HsdK=Upxzq*kPyxHA9BBdeiNwBEx z|2|w)#AUe9c=+P^bAU+?hv~2Qiv&kR47%;)T~6(p&&$iRk7^K2?n+EZh}z>~4wJz2 zK-FYy(1I#Z^ieYHty{7EW8i0LU)*ixC@=rk`u6d+-A9t8U3iYf?;>yIzPCKQzpPIc z*tO(v+lPL&&nd5aeRgCoW7B7BYXJ<2!@TMrRJpAjtBG!Z|D>D$gG?#H!wW-zYlI zk%5+^i00i}-d?9mIZ6~P4s>Ol;J3e4|I&E-jV$2(%J>~7f562ijL_|DpdQ6ujMt9Da~41YCli-A`tYB zGuu1m?fD!OseDDA|83;mzpH&NcmePV9@7w=k&%&sJtfYT_vFcaEt2Db^pY>j*-bEm zFOMbBLVNOusp(WEuX>OE?kAQD5|zXhGK-$SaILf=ap1+%KB@GSc))#;;rI(@rsoH zqlGqSz>ON>NYMa_IqbvZ7TL_1l$f=4{^kOt52JLaSOfdw`8HfNUucY)q(YXfomH5d zQJ}lVYA*DGj;o3IYF|%>iFt-~en@YK!=_cEiW4$}s-Ik5h$mX5(J97I zjqX)ok{NCw=x!tA;t4*f)4P@oF{}zYtvdh@?HK`6K5alu+tco`BiXYYS+dcSH39D; zp8m45v3Z*M^1JjJiU1;SFHEHbo{d_6DGF%A@9%yddoljD$HrkSnE|_n!g&kh-CZYr zhLDQ09ib#BaB9}|PwhDKN95%m9uEzn@S!Cx)7bRCIL#ei22;(O6;d#13;u{}D<{Rh zu_($v+@<|}xZi=+Y39?)7@VApjjh{)W3Tj`A2H^4Ws>w{32!9)oIyzzHb1sZo2_i0 zy@oP5_4*dSL4L?C!tm6!`Qcv!wg~T^f@T#RdmmNj>`58dlvcZM&%%zUoz3XaGB)n( zpD6Hrns)RTL3IjwSD2CEPT^J6{w!7hd>th(p_L66|MMvN@c&k!X1*QF&3w~vviSA5 z9rp-xvIwTjOr~PxoXv`?Rp0&cj#pUM=EVsn zWCu7i^5et%8-*9@%(%PkUg_^s$G@v;pW#09Kga*43Ua{z?L#sufkAeBjI&-5a8gtp ztqFM51UmY?(&Og_{HG56#4CX@irQHISr?MU&%;2gY9f z{LCNh$~daR&!b8(prUk8S!>8?_7jxqap%U(aUFhqki7U|JHQj1OY8cmDyDfIF8bwGGS!&}6Z$#N6oR9pwR%5#YdQLx zMYlyoUyF)ZTrMairdUi%#q!JDH?dg{4L;Qxr=7K)m#5Bwb#O$*7L?`2FyH28(c3rC zl6^BYv6z{WF;=2dMm7DWeyKo7skCIVNx{QFbr$`a%3MQK&2alY(CGh2&cGPz$f1O$ z#2*DbyIgeP}%oy+-|G9Gisi-%SAbXhD_mgPn=doI7;cDdK8 zZ|6YsAcjg5rB&|IjTZ89k$4nJtHpcEtnP|dXpdW%023nUa(&Wh!KdWZTBih7i$xP5 z_b3debm{)%0Szoe$1z)l8Sgh7gZISmA#xPy|A0E=-ABT&j(y$amnJ1s8+T_Lo_v+n zGEkM#@@Ap?{GPQXKMJ@em; zYzu-{&jwN~=N05%p=PDMa4)d+FRGvcb%(Dxj)Ca+_Bqvc{wy_hw-x?Kr1kM9pRz4h z%uU*0<8b329dF4SpO~YMw!%(T3LG6Bk6JW3l766Y@@vD!gNl*v;}_LNJSKziGALlrRDw3W(tqZ7-oxj-t zqFAxm`Zg>qtRm-2*VNb9jYoZ6rLikL>KTSz1)VC&Ar&QMQ#Qru8cqHMip=A9%=qQM zcBv#m@5;}LAMO^Q*|7X~Ejpy=^GkM`%M+W&h=Sgo5`J{vN-I42evX zKWj2F@@p9)IT`fJN&v6w;^Y$c{m@8Z?S?<;JN4w|-4)7gdt)=9hnfY(<~6Qf#(*Vwe`qWwN5=}asyGrMg#ieh>!3zP30K4mfNcZv8} z&wBFU!9PCuHpODvcXbk#d4~oGD~ei1EJ=ujG#7UwnJ0G5kqPBR3EMBGCpXYyt6#|5 ze-&`?&s}Xg4!!a4G@0{0KR`^{mIQon{CFL)cJmclJduaGEA`YV;hKToxNnVL>~iq` z%KmPBycV1}{ZAx-$4t_@|7AQ~GMd}r+kqcHp$`j-0Gn|LI;GnzRAj-i^bD$GSwBB1 zFXZbYk(ywIR-*ii>DS%it%#2-q-RiwpPh1s=-=)ZhJ0uC|!UvaO^(Q6;M^?DuNzI)Sb z+iQ4QDYZ;Gc?09VwykMt@}pAcLuvc@<>cc@&W!)x^=v1HvxmoWZS+8=aNjXeJbp){ zS!Wb(TkGCAWKGXw1~`$kjX-YtI1;Lzz{eb?5TSsZ*{L-eu2@gP`P(WHG$0W9;~1jfWYB%qOZ0 z9|L-ZV14}M4=<9FQEwr*0D*N}+<&=hgOCjViY@l-yT3YLz&&qCGq0N`nnj?lV|3ZX z5<3#Jd4M&;zjd_VsWoO}(fZw(GrCvzZjGp8kLvJdOfETYcP@H9B3srvMmwGpzfb6= zQ@rW@?{X7H64*7M=CW9`+2V%nj)3?tamr=@k_gR_5z1*LW7IK)kc8@=8;plU4qgaN zk;!hdGd8U~edjn+Ir?4bh8Pen)`+oAaV!a@%F&2mcIN(hzqZf*_Y>4NNpMs>+_5jS zg#8_N{5$~x*ZF^L?eUo>Q^4DUm!x_khu<gb}WbUv^E*&n|W^obng(t~5U(9go2bSz-So zJlA>p$A|#46BCm%ctKve^bN;nll9D7mBu6C2G-x76a08KjpKs2%c!~->@qfTYH>hK zt)#GCLIGNmmg4Ht@4#D#9ueQ;|_KC5i1#K=}i<(m1+Ny z6F5zik(qG>KdGs0JGQa>&6A80`Nc`}UBUhBZ9#RHTTlEbuqW<5|3QUUWbo<#57y!T z5;Gd#sr)iYWbGuc_(XaNtN!ag;Dw$174_*h%+5NMRc07eAI4X_OCB5BSa&Op;4Rr# z|CI>hjd$X#9^lw#L|S*JBxEGlPuW7`IX($Yt|jI|k&JWQR60-7^9nZlu*>-8&qZ<2 zcU=5%B}m3e1uf=Qc&}}&=l3vRCOBTkw+Hy|&6sx2Z`AH>DJgB^K1^P7B+qPUBJcfI z-kI0R>Um~L$P}7x((w(H<*!CRk+^a^(mV_9S|=&Tr{)*H4z#C}r4K|7C+bCN+QwOGFcNj;Hub7CU zg{xwedLqxD33tD2{JIK;J@{p!TUiXh(saCj@JnG4$E|hZK9#=0!a25xAQs<7A~~nz z+)o8*ndSLgNE%mt?fD;f;~TGcnl1Dzd!`J{%~lq)G!&u7b@SDk5Nfu*E~(XlY-)^@ zmd5J2Ux-P$Y*j@Z#8QT)EKRIO9P_&81S*1!#E}jb=p9uFg~G0I%2;BvyUO`&SNc?V zN0#4@f0fJj{d`8>BhnG26yDc7C=*7c5ty5ldcFT5={2(y3lnRTg;(cTZV11abtLoY z((38AGF@OLKjm}l`P^YS)(sASOtWrL8tB&^>FyX9xoq*Lq(7m5;GgrA9AQHWa@I^~W|xA}>H@0SnNFxniA$!S zZ#+M{GtW(zrP8#Zv7m10r8FB-t${aHh3)-@H%^= z|3Z&)d9lY-6p_B2LTB)GeOHQrj}h_I?!_6OD%8UbofU@(rq+@Y0`EpXLzXc`Q^sEo zRPR{mu^#)okQaQVI_zjy4nO!{dddN(uYa)Gg)W>Kd!3u1dALuC0CCkN1Dq z%WN+@9x`DoQNdn^n=(?T>w6dJDp`56lZtcml4{NDZy?Wjpu|NsKQ#y*nR^)uX+*XX zT{bFn(%Y3%xaV;>V{g9UXo((UW@3%Cdro?)Q%}h2E-$aOv!-CaqXTb@i&x*_b$^6B z_zt)1l_DPES8DJzcbd7)9B7G@FV|AH^3VD5GjH9}d6zYyA=foLFl8egE;ZHl5nJNw zpqMhpr}^;L?5m&fPJ@+E(f z%O85udeU-J)Dq>HVmvocb<@S=G5bw+z>(qL=4TM-Sa!%9&Xz;g(Ha}vN;IgW{lm<$ zUOnRB`peYtc(`YTy~a$;q0TJ-rk5^r$D_p`tMVMy~}c znZsWb52ymS$}@Zlh?S`=KONeu2&+q`R8JpPMkhv|!-(8`(M}nOK%5Ts$67McrC8() zXjNbS5;0gRC@8BYKf-6HTJv&0P3}>U)Y1s~`LxwQwg8#HN4|t|&&(@5t-q^GzZ=kh z=`VKf4p(Ph@u!ls?1PM0&s$>DHfe)=+SWJ(mr}uDvKtU;xoegSLESC6uWwA3rJKZs-G1ns60)|_a5ProqY{9_Qoh6YO1hLL?tLwTFA zI!hOWq8_JyK2!awb>W?*`re8Puka0a9u9_^aa(i*mKh8jSsE3lmU0_%L_59jeWwW= z88|ZF%+luwdU$-7UY>a3B_eNTUFl6OqMh_vOr@1(sYTb|;Jd4#DHXgEuMRP;)_o++ zoEaO=<%pcz#qm`s^UB*yJbn?kZYFB+I=uodmD^rjA4jIg(wf&&CKttr9;OH$iR2wu zulwH2<0H|L(5zDTHJ-peaw_7eFqX-sjMzIvc}}%6j#OEF#3`|yMqgZt_KQ+_G%Wk1 zBsEIzTF-qcvo&SeP#*H5#%eel_D-EoX4K*`v{28>&o$l%Gc1=ewOA-QD#?~0nc=Un zc{M=1ZD@Khm=$oi8`6mF>8MBUl;f9pL5J%6rRsS&5kXoywzr`&resB^7^Hcw0Ms z+3JX|meY3~0e;!O&at;oe@!gB?c-;7*r+vR(d$$=dZQOHX+@@;Z295`JwZ8rgJ}w=-_@ndL<%1Scv;kvJIEYL zn0ZE$j*pxNC=d>aUWv!79Ypi;C_<2ir~3wiQ!A)A>&vCfMy@}0Z6`jT zAu4`sS9kd8<@Jp*w3k%V>vH4K_ju7F$ioPtH8szk?(}Wm?*``VW{q!DdS@6G-q1xW z=s%<`HC%0|L6=R~@-w%yU&e>kw5}ktsV+dS8X=<;DvfEA4xn&j6i2>|_s$KL8OSY6 ztJ}0ZYu=JnbyV_L@1embG~@es8yFO-o6V6}nZ52s?$&Q`aD4q33JrT<*SvB%c%C=i zDo`{-E{v1HI@R17hA&xY6Y0hNy^yRl$b)P1m72Ww=cjldQl;}I@qIOS)FK2u<+cqr zKusQf1r^k!Lf1X?8bFB>6Zk{^a@t1KJZq=j{oTYLnoodY(a6= z00AKb`cFdBxkXv{hu%_7xYPsJjAc!;I$7dtX1G6id)54EP|@Asx3?5dK1Os-*=z{V z@-s^|2K$OBKioFnQ8LYrQ}-W4IoLmT+^UobXXO(R%Np|NUmd#Ow81rk_$cp-?q3Xe z;Kh{as7yozPhgJk?n{yLMXx{l`M^%xNVK6p-l3ZVx1OKD{v}mjR30P1#9Uaqr4TJ* zzx5iEMaIriV=})yH}?s3?c|2H^3uvB8IxJ1+yUZ*{iC;PW1H;sd!{dA^%C*UmT zZ5d9oCdhSamus*kKN3lkK}(xe4t$DF*qoNB7_|M=d;Kj1<)xG!2PfXM$-|e~R6#tt zwCFMsoaOGX<&-=?>@0A};@1>;B5~9h;_%IjZo`H#Zj5K=plFPv?%W z*ZOpF*eVe1&|j*uG(7zMYuundNjm;Gk}7K}T&Yw*P}ZgA{WtSlNr|;D&&+lQG7c%) zrBAy{CU>O>*N~>TjhC!=5K6VY=BNJU&X0%3%zqrLMqwgv>7*)ShDt2%cpN?|?4m(* z#Lb*nyRYEGPQ&+QVWT3?);~)Qqg<-$jDOfW+OxLvRGi2KM!J1{(ry2zgKHa)`Dj>G zkx1G?E%qi0o(k0uDxj)JlTm&0TgLGun5Xsd=VO9w)IUEJG30ESa>@BH{wTFt80mPE z=-ox{iS+Z^Ib}%XGlG{8#ESc|{H?o^S6n#V%DUih%x@vMD(~ZkS`6ubBdgP#@a;8u zy|nKIg01}5GLq${n%b}d?cBx_deQ8`*GGCqPu@v$M5jZxr~4+tO2&)4jh@XcADxHV z!nY|&-?qu6`_#NV?D?s;uW04|bdJnlnz4=C z^N6XQ7RL{4kQOH*%2#K9llY~CGCx-V=HYy@htXKW!oPQ# zDm$-ot2VAFX*2cue5k>CDpx6(o(FzH_;&P8Zc{{g}rc;6S3Ky)g<+i=f=yGU1h`21>3=tg&1aG{q+lW%pJQia0F z{D@e0<9$DBuZ_)_Nq)l(Cd!dH+s7&a;=}6h>17ogc~^{RpUa2^4ro~IY`F2htbftk zb}eas>#7n)X-kK+>HB)+${7PXYNMsX@$w{!)K`KCMYqT!1+T!f#H6Q9%vF**I|uCDr^;s9 zoDvh8r!NMb$`5Vk$Y_<83?L*}SX?&CcJ!2!S5#KHj))0PD3(tDvBqO=uB`d7Zz01X zFG{!-vD%y28qL2M((wNI4>VrJvY3#EhRXLnoh6e(ly5`*p3_nW*Kqg14DUgFIj_?h zpWZ?Z52m9+V>bg?F)8yk)ALqyIR_p?b?EUo@PyouchWuw8J4l4J`r1r6w2zFRPAF0 zo>$4QYRf9w4xX5}lgct9Y8}z(wv_pL`xCsnM1JVFOV2>F-w~N~L-;I%jB$W;XG*ur zbv+MCqrhuMe}0uBzk)R{cC`>kKM>^Sm2+`jXbF1Ys`XeYxV-+#WuvH&Ysp<{oRngu z?2GG>AM4)xLwI7zm||tIKo)e78Nu>63HLQA3VXYJu3f}LzJK;I$m9lJk5zpi7ERHi zyJUuSn(>zu3nVCadT91~r_AEy@2XMJddmM&pu|fCXJ)3W5+)`o0~VVorYcj_40uLi z0X#vYTpNp3XI9vwmk49@ns_Fa++{K;^kr~KD`P~*n$5JAyfT<~PBEF70@ISkO^=`t zrRU{nT-czpN_Q+7G1l_d#z+wc5_ZoMvUhW7X6BL`{LQ|QUot!XCQW-M7cYZ@xnFZ& z@SAyuDcfH&==fFw#A;tD6=_(Mi_^w>wOU!5T$$5XpV}SW1=~TlsK=jQj!MM$u%VXe zFoHdF_IAPE>pVo+1JPH_RSeJqHKbfsb9~xa`CX&E(`bTw*G_KR=RD9YnZAI-i$E6# z<|y3UxT06`XCFp<8ux}T7vqAD;GKnX&ENoouRf={aa0AIBqHffUT5jj6eMJxm_gO-Q`sRPVh^(N)luw)>VEkW=!rG7N92 zy0h(ygCnb44CYMSmev%OTT}-XQIDIiz4`f$#RT-Nt25E_CaM=B)sDYn-)nQ3mnUtn z?dWK4bZ(w3lxY6+(;F9S|E{`U^}VPud#D(`tIME7QeC{Jub9_LVR_G|^!KpEYy0rB^mhVt z+~z8ie}Q!>6S0}!T!7>MwIMsMb56axUZ?$*on=_-H31|E25b@j`ioZ|`Jy>*OcU(n zJt*YtpQBtg6|Q+t#@P_G!Q7WxFf%cGg2Y^TIpIOTqM7Ba8815p|BaaziU8B!4fxJh zcc_hnv2=gPPhDQFa1t&mB_)^U2EdqIfWzti463HzM_*m26G@BkY)B^PE-vRaf9bAH z#BxRY+P8*-QUP@pp#%%X26b6h8Y@|9-6LQAP*D_t=!MJ5p2l2Bzi)ih<)4v`E?;Em zz5eMJnKnIXwPyJ|6ahAzcD5m6r8IV=U(5sw*Eb0d@(?e-?LP>xU{hSQI!`Eip-9#L zjN9(`VRfojbVl^@#O&iAu@uK$gI6`JP(AYN^{-#LdH*82$EUdb|EYi+pBQKRKjGT0 ze;5P5K1KT8@0(PY5z^nE8m>l_g>n<#c_SmxG0`CH#Fk%(TZ*mWYqhJnYFz;6N4s^tsib%WpR??&xe+VjlPLwrD393`6_XA3$bHrdRd|CqeVg7 z@-bcyk*@GB*>)A^Fc-qY@n05hEe|MoXl50aYej2YqpW3$G=4Ytwej!f`X8FEIEr{r zC}J3h>Grzn3U4MRSqeDZ9nSbYM(XYzch{2#rgIk78%4n%r{FnzQq8GOS+|%t=-34w zDRmQCtN3H$;)jBXFUJPIPrBXIke85V_DEV=_WW{Y8ojtoySVtOvm0Sw56{TbP4BqB z*_jwv8W+|=d!hj0>$>v@{kkdkZd~hBs+CnlM9>`C&)UpdP*tUCZbL+a@x|P=fPOBN zN?%)=zE*j;M=V2i%v~l1-X^o9jXiL1UQS8>%C>aLwM|-5+}H2?AKtte`gTsEkJUk> zl<{e4Xh~^_SWcG8w2pyUr&915jZTASr!Y^rRNwB5G4*qF)?`1SW34%Fpn!axoY>|u zFLrmfP~=vv^m~IG5mae(YmZx~gsa5w{t)sz9#6Wm9AQ|`wJM+UYsS5WoS~eGg^s9% zT5hWdTd9uc!-7_xD@X}YI&=|5f3h1ihSRMLdBm34qq%y9cQ3kcRsCO_y>(nvU-a*d zg#m&i(keO#5|Yws3?UNI(nxoQG>Xz8QqoE{(p@4T-5`yCbVxTmYgGK+dwgW z&q&7)Ot)A+GVf$;C*)1|>L_mi9)Vn%XzbH~a9gIEnicr8U7R7$8s z9P6)HLX(vweR6a#T>6xeZsPltfpI9VY>uez>?XeJD%k9r1BlaY^cy+w*^~Dmu3>0;3x2HCIzxawq4mgUfZY;;dWW~gkUrgOi)$e2*s^UL<<2jX>pd1pmf_s*w zGh%22Q#L##4BQj;+cbOioa<4UnSyTSB62n_bVAePCr!zmZz0VQ2sO3JT*mw9Ba@tp zYC+*2c{)5j{EOzVAhFs)0Eu{_VEeT&=!+uMHZa{5GEoJma4c;*QAQ(FHB(5*6gj!#Ow zyK9HFAw1CI+w!F)L!EJ_Xq$a0D=lhPc(=%XuV*`}$OZ>A;Z_wm%zPRg{PZ&ROQT+{oUu#JW{;#{uv638 z{yoX9sVD0+MH$Ana)WVrKA~fVxhO4BC6RwA=P!TLLF33^Er|-6Xb)O3!ZxT$=6l($#VFI$pRJIq@p>d?;RE9iP8Fl+QCN0~s%$SIzvFQu?r6oY z@!Ua@)Xa;RZfP}U<6l>8as54lWhHc*e{3I3t;}5uJa~Jo=auU)BIYr+FWoUlj@GH+Y+BOso*=m=k;GN0 zKNiClFw>GC z9ss*p70w^0!?_aryXlc>cArg8Hp{F34g6t@UrjYQc1hOaNoW3Y`qhIf@8b*+{u=eq z$9HP((}!OnblF}bofe~EiS#v*kPpaaXJej^D`ieGcx2`4sTUC%_^v~4$XYTifIp(| zn3uooB-wU>@NbT9{8458Y($-Ol77Ti54S1*YMVQM)nOmCd~=<+#pgX#S%W4sNmZ%U z+qt04e-L}3w|}R4v)k#@%5+E%L-C_MKWmA~!C;i|~xHKDtifl(6>XZ~%AI+}eg?Vsb^zuKSSIwA9|>N980iw(wwEw})=qVSP-1 zZ7>-@V0*h72i^UJFa0K*G6dL|$vQK`wzX+?L0B7v!(kTFkJB5$`c9lJb24#$rOK4% zKCl}kHL)ZwLDNH)b^z zFQy!^C1c~5w%njzirt|7yX9?(W8J%MulS7})8i^tQ7!JLOd4aIkE7iS`x^D^fwl-;}J@278{lTTu{vG*o@ee;CwsNNH zyvn&3@;E8zk7S@=VFn9d)3~r*m37}S3{2Cgn&q0K?(vH@p7chmHfg`5rO2kO1=?Et zWL>=on<<%!Y>C+=M^}KbMgk!sAZUlvFEnlOp5uPX-nE~J{&8i zVQ^t2kCa#(Sew~D&$%S9JM3ag>aWi<93`IIztGx5qW+NY)?#tv7osm>1oc-;av9}j z)iMvRR#P>YV)))Zz^r~=W9oV;=Qn8_1LMrX{8FZY)fO#f!`IQOwWM{GT$|?egr318 zWfS$nK0l+9{rmdo1DVCT=^1MWN}Vbz!4R+R9Xf$&oxnPa8GHBS!BKv7Vq&};osR?e zBG#ibZ3~pg2Q@EShpuNx$@Q8UP?es1dR*$SXEc75HNS{|tIv*oHltXbc`<-=l!dh2 zNq*VLl+=6qWNa5hHU96{#VFOP?EJOwi7I79dtaXwot&@w%CrBuB>F9*1-eWw2Ei1= zbixLX2e*XwBDvg4!hU&O&+_s^bk;H%qZZ%%_;9i!VSNAhCfYutBIzu^L3&(b6%g$yffI2f3_{uxQbgDDl*IaB+>!zEO_tS@7K7)x-T1 zCb@!dB%?Q{rz8l%dA`*|GxTp28@QLm|6){P?PbF577sU|iQu2=7OYPeY!ipI?8!OE zMKAw(pif!>!?p8&R3-mkEb9J;uGRm`%JBc7^&&mb4wc#S@ugbt3$2Aj3N!-f{Ovkm4 zL&&bD@N|BMDQUg?{fdh7=2?eDPuBl=MQ5yS{56=u^YH34f(47-U)609@fG%!(=Zl! z6x`R9D(cCqGuVsmu+nhXEoSs-+1ITAOR;`2gxW2$mQ=wHA<7<{{h#a{xMPPnHL~;j z$zD#ojUX*%6YhUJr%7(a-9)Q8W}Q+yl&3!?898uWJ11C9=BOh7-k`=2IIAXKNlu;- zUtEKe9{NUwnYeu{0{YWaLsrjFJ*jh74Huj~_3RG|9ogb|c}7HQwLC4&e|jylGe6}4 z>R^bb++}3&aM30+N97t*b`&C%F!BoC>IlQV-D5ch6;@r5?DpeG>qyHWyrW61r$>{8 ziR^9kR|BicmD2PTpWFYt1eHJ81`Hzp)8`ksV@K(c75fI0p@cqZXK{OM4valKiEw)A z=pt#nHJYBxkEq#R45RYy+8=B{PB`s6%@+2N=6(?xSt(cW%BT$zh_F? zUeUIV*SDg|&r@l`R*25mgZIf6`37d!&cCLNzkP4%n( zO}Fx3)ZwlDAaDLF|3zMkss5U&$sP)$-Ip~_uf>z>)|J%P1%>EEPA`SD-K6;9NPKvi z*ewnF`N<`zzxz9K8L9ft*(HT^3&=VyY@Jo&RrKg*A@BaIE>*mL;svTx{oTW%{_I|Q z)>YX~u8QaYQyiOp)Eyf7|7?SLl;nvu)O5S5i9pP4F4!?Ye)L9PG0Bt4ZAhytI7nSd zRV(7dwKSsplg7{d!#FlU|s$mu>AS%@=9!L-cm%HCkr~<8K#lCved! zU1QSux-Xuu{F0`7SlKI|B+mGgx4C%MS0R!~^7fQF4MbPtrw(XiqbsN)Ec*0*20d)L z75408*p1_z^vkQ2f`Gb|tqlvv|)|Y1zkjlt9 zZe>_8qY~F*#Ns@eREN|k_^{#KtHGthM6w~0-y0|Ii&?8a`gb!i2F9a*2Ft&Je^G5$ z^VoYor!~y($#$7g1(_TXm@gmM9ZoQ>10-Us+xnfxBH8tEpJBGX&rQi@eLJS)o0qy{(kR}=A^x`|E6Ct z=wDA`?MW2Is1!NdeY*OBZ-$~KjdCSa{6GLNQP@5>d&Io_4dbT zqt3QG5rtn`>z@D3z!b>8(Y!0a4vZr;*rL{DKBj7Z9Xc83<5y3{?uCueUXQYq6Jxl$ zk}RJX&A%5Jf9}6B)ZcGa>l-=#PTLU;wFu}NC~HB$mcSq|2^A}Oap%lv%+2^t(1xV9 zQ*`6Cz7TQJ@4nK1`&~0qe?UgM*nxA{98a< z{XI#v`1NV2Q?Y07ouXh4eSQ*~SO0LbEInyX3`Rw&pd^M1wzzJ&s|C7WMozoZKC-mT z^Snbz>0X9V$B0Ox_P)AI?|SOj!*i~;lU)kd&S7CZ_Rp5{z3BMtl=?UGptnq~0;n*s zIi1g7>$Mm;hKGdUFB4!WBQ}JI$2f^v%T=?x=U_*=kB?>TQFt4Zt&m(BbQvlRK{^XIWWok_74`2R5s`|ErV9yu=7taXUN)8IaEpK(gzJw) zJFbI*Cv)KpX{85yASX9+PM5GadJo4h*%rQ?iy#+)GafuALw&9BZ%@?`{v@m_wY z;M}UW@W=9&abYZ5bO-#o2^mt{z}z|?!m5p~r!Y2Q8#jpFIZ1#@R7P4_lSxr~B>T#u z@Ibc&{~<`&B$?LzFhlEleR|TDrLqqg3ScTl11u;B6xjp6efwr9*aX6Fl^cl+c04gk zCMDXSyMw%njY$U3Dzu1=lg9|`T|eH-J~m*3$V7VjLbKYDg_hP72!^5A^>b&IOBd*W z&-&!`3%D}HTmoN7FUS(`|H(d0jWuqwIIp2Kqi2)BHSy!yC0q z0RRvS!U|}G614=;u}zEKG}(xx%`yW}kg<<30$nl+3JMSy!FAbr4Dv3Z?6m>fjvKAk zyt`9DNGWAy9@9}CUi!X(p{1i~_9{?}A)r^@$uxC(%5ZeFU4=%-qNYH!&>kR4Fw`bt zLpI1d)-^PYfojbsL*AS1D{vdo5TIoV`^Sl}@A<+QWwc0Bj){5jjDJ*hZhKi_AlLy~_m8Wd zvY8B1A<=4Cz=6${2SOP%9?ngk?;RcQU>m#bFZ2dz5Eq<^l33c_2L&-xd3kxlYu63{ zO%)Jx(a^x4Yh4SiYY2jSn}O%fodcPG6;O)YmN?#!0C}WO#IV~NEy21uVunU4k!vbh zKb?IkRIxzcc+rr%`g2G5NKI&cH}2~i^;Yx8&XOM-PCOor`0pG8n&u#&2o`V7>L!QT zEtz|?w3b}aAgl^3Ouqg1pq+#^3LqGFLG}_B)~4Iu-d-g+$W)kDLdfo5QI~(H#4@?| z$4+myGd=?GiLgIgjR6TFlb3%2Zwhz}d%0Sm#=68^3!R6G0rh6}DSc#OvhpX7&FIlS zdbQ`=r`-iUFTWwcesgorf3afR?4)_k5=rj{GRw;ThFcAM*JWjP($1Almy2G1-&^V5FUs3L6 z!a_vU)GzC^*k|Zwxy|e80{VJyFVTBRV5>=-w2-0LCC+Qh`f*}1Y|lA4YV9~4cc z0ve>1m7_vKZQ54RPcaC~tN_*tRKl{9OV$k4_6Pebw^%{!bw?K)dI@wtcfFFF_O`!m zZP^@FTTyz{{xIVJ0UPb|XYfndewnPThcxFS1`kiknSLq0YGm5W#l44r7$Wqg=U*`k z(V6;o8tuKCXyGvic9UTam!kteQcxiRxyzSrfgU(F?9oa1(%L}1vU#X&*y!hYn3nT)LL;V%z<5K8&^tQ@_S4Vy{r*7Lp|e%Nn6#&~I=c4Iu&1PnT;1l&sL>yOR5FI)S;Y6umv-CYRmH6L&y4X7ygeSs+gFU~YWgrm>#b!Tt*>V>_ zmJd+``|!T^^~p&}GVLTcx3$S#6F{M@XxYfJMF~(rm=0y2GjEz;VYYwgNG&QVs;oSE zuZm7jyLb6Vgj)6$^&+TiIVKh^a~QVK8}g}re&vF3;*?>u87m7kp-^Lb7Trl=pC&DV z)r1!a{4_aP**W~M=KG^GA}6%~o1h-5#J{|_Se|m%&)3%x)P+!5OM?Y_d%B?TlRvZs zf=ys|^r&R}09{yp26q_5iqPtO+T|#8Z@gM6nDO_QEwZ)V+7l|jq5k}Aly$_)_4HLQ zVbfhsIfj4NbOiy?3nC34LI#{049dm79BVo=T!=7BKF>}4Jh1GXtcOEZ9LHIlC%)G&!21uzB0L`7w&F(;U! z-Ez!@XVM=&W#HFQYz~%$?F#J zGfrfldYoJ$xo+dK_Y5Jubw_&Ds7`B{o93A_BX-4!O>*u06CBLCHFEU8?y2tVi>NwB zP3NYrgJ1^nC(ZtT7kVpRRjUvb&^p`pBPNKB^=H@S* z|89q<_jhBAU$PkfId}2w<|I|nsyk1E_Y-OaWo^0d8-{C%i2_59$cvzSyPjDsHsix> zijiEi)Z27Vz5H3*%lcUK)8b+(@ixlZlv`kG{&W1PI?AFN;=AZQ;OkDo*~m^H>d8zh zfRYSu*P?K$yhTUxB@Xwy%INqYnM=x zP=#CbU}@EC>(a^D1QfN$NcVD@g_B8_SdjFRT%Ko- zt^H7`L+v1gnb+%ryZP!5P}mn6`*!#}6hniDM)KE49?#ZArZ2p-$}+x?Kz6e*pa z>lxflIWdv062BLI)gcq{gMH!D=0*6rTJ-|T>d*7*tfXt+EaV2!C~}>7Xnd9w2FCQ9 z3`1XCJ{Cvk^Z*`J%!?)cRaVB*=fNCP8vEsnrZlV(KPu(JW2ee!c=0srNA>d z{(0g2@*nz^uvi@IZ*+zt^sf++ymBX5Zz)M4UjM2`V)prVcE`l_xio(Z-G5f{(M~P8 zd9$!13Ng9zZBA(YL@^_CWP!BM-}S{OovEgC+hvPFlCrf46F(OJ{hlw=-)O2}_Va~mrSqrb zo9m_;t-PE3ZlB!V&NUa#lh=R#j!F`ZsJu1Ii``V;A9K+b3nS|SFsP94p+0k2(DjsF zL=u%qaTx|?^?7^@_o$=ZCc=9V7pPQZqWm2^Xn>T?S-v<7#;KjljJGsj#vaX%()}<7 z$ZRJ$niBoP10s73#Kejeak^JW}hMlLX7 z9Hxh;ohHKjojm%vbt#^VU4iKFN%rUxm`2K9q!^kPs!lP={UTL=@HYta@238rip>5; zDxsBV8PokaOyxkc0J!ad(Xx`MiHS{#GG!3;-MceUVW7Tqn~Tf-j*_vF5jg)qLe5vr z1Q9=&Mlv8G_4rs6wrHa3z?KIi0gg8P}SpSUeu5u@@WaqNy?@6V1 za(T;cDMrU|<>mf>*J|=;zur~N_hq-y(?Um;;x2_N)2GITwcUf(fVtXUdf1(h4thj|wF6oM1m($&*j~~0kC7bRTqbbEnp@iz%#Tdu9d^Pv1FyY> zQ9gE8G+!|;tZj>Aqc(n-*Xh~r{q8m_aTjIMtl4b2KQV0xw)oU5B^dQ(C;()O5426I zga-hAyy*rCQSBd?$YpG~qCs&CMJX1+q9KN+2et>%DHodoyFxX4yH{Q2gb^wA>>2D<6-ufd9UTRl znLH9LK<7BaJthghCf{~R`v|NbK{3h>)C$qk&GY0+ii)T^fJ6itL89={e?t)cLitCJ z89U=wAWopPj*kwVe9q%szXBl>WdwLb5@XT;&_;H#Yq1T3wkVKvJcpgON9+Mh3|g@u zl5K*p12NJ^EfrNqQJ#^hX@H+!uMaS`@dDAY8PFFz zf=)m4$`#gl5;yo9cNU%QfP^w=55k%l(#LXgeZWT9%X$IsDFE8>@$f(vt`wl*OG8Dv zIyzYZ+Lmem_RXwpWPf`FdIuJ?F)aG;2aIr3RF#xykpK_&^73+umBFo@dy)$-D56td z?*Djx)W7=r2`lL9&p}8B?W>M|`|#ldM0vn<1z1n>Boykwl-CbQrumDO(N}9)QFWUNWle{}fYtLL-I>__R-MeaC4*y$Y zVh|c*qng0qCjOw3KOZg>m|Pc^{cj0Y`#0DIdfv@F>Gp%7$RL81^e*#R!14qa-W zkWo<3`_n=|P*6upYqsBHs8FmuC{rOX3795JHpf43T~e~LoM%6I0GsBbEWqbF0OGx1 zo&$oOe1H;>aRGnMqQ2%?IHe^p^!HXa7YN$eF~|! zD#}znTp8Qz8@N2|yuSq6*dLu80V4q1bhXU%bcdfI@|9~L*~&eJU?d3w)sFZ2KcTrA zdz%Ymz@$O1_npodz{Ncv`l%?@O-D`8+tV@R_%WS%_@A$iF`nx?FRVZ%0U(bF2ra@2 zMuRbx+38#^;05M!aj$b+xfCQ8%_SSlMueS~2B^kQLLHC#3GnfO8Ff#kG#?fdQfg{r z>{8Qa9F?oBrZI<$FjmKQ>e(ORf>t=dykpRG^LC!~mO8;E0JLnn0i&cUA>j@U4LWj$ zUV?+->uMQ@welZL^3xlz!4>;nl@}r$l36e}U0RKi{xq?e@|x(el~6U{ub-uS&=-%< zYhJ^2D70QA3mumiWnHpsp9@_x#jvkB~Huo%kB1@XX}=6N%Z~$KMoXr6oJ-4_^<`7@QcL#IaK8GCXmdKmv`pA~3hoE*H(t5R8|U z?L~(s&LlN$7*j!oh=YTJk!O_tmlnC1+{47zD8)g7d02FWUJsKFJ)}Iwy;*@_uQ0#sXkmo;Le~r?+ixO z^vYbGy+Wbgqa@bHsKUafgOtNbf-!;o!Q~aB4c3K*PWP0H1RkC;RFr4s z*WyiMpKcwgD#1+QCA|Bw=>7Fm*({YIhRr)EPO(Rf{h_y!5qb{W*B zb@L4ZfxEF1l@C#>7L~;BTvd5`hlUpoj~&*B5NX=-v2{Z*y8axD8G@RyrM2J)!lw=p z7IXgPMz+x!GzS99d;*e?I%r385>lm;y}yC9Gl<1%YJHN1U?(ycdb+7e1QE_@sWaOH zTJ;+|12C)bxEwhiDk27ei}4fH^t$Gwb4iR!H=2I5#4b9fb)RMYH0ZE+0zyjpCECUc zCXXz;(c*B6)Q>$o#fr*49&lmR}r_M<(Hk)5t^t!_+xqAsQJw`%Q%Vk-VF{? z$1tc?sOula+pv4hi)3z$VB%ccC@sMg^}^9Ss9DBCYDm+aE6jMQDf%E^E_PqtR{H@~ z3~z7aXnjzH-=2B?Z4v(OW3oTGPdj<^+Aqk~eSGqHrI}|s(~lLUqbP;#{V6H=Q(@5) zR9EYhxlk>E9%){&yY{4e;&I+mdd(t?17c~PM#dCIRhDp~a?eD;zeb`ZngeA1A|;^(RC%3iHIX9Z}n91=3y1xNA+D_jT~l#I5^YbOc}}@2A+aV;20C&}1`2ku8fo1k%?5C>`ClvvUSOPVpdJBm3Ec@tS#|>q$F? z6Vu`W^r9>SXB0oz_2+h1m`V@fNQ_@~kr{hmE+yNOEK#A{D@|q}?!eU~MMrIi2u~vt zx9Tq++mlSFVdK_OETnrZ+Y_-oXuMR$KhJTaE-y9IP*bfk04l_nZ|MNS5d2F zWgy~-u&`N)a+%E%_-~*Xec^I?cJ>(LFVH|@l><^R2d=-?F`Z0YgvzN+3}kACj1w^h z#25zzw4A!RyW@Q{{Fo#$m2^Lrail#!4IlY0s|R>q8_YG4*j4srDlI9h8yNcV;z=wa zJ8g5Qb4q$H-8dqUxKo-g%}z=yhDIbkdn6@=t(TbszHNbImX%a;%xs)aFjCSIJz20%+-ccVk zo3{*KZD!GhZQ0-p5qf*pMtu}&+6W7OAy$yqC+8A5axM0Sm>A2Drmy5hN{6zsXuaHI zCCUT^>1{jzIn$AXoV=@Rv zvP$P^4%Gpvb6EI(v25VyS1-KB4KNEKEr3v$mPS)CVNp^E>8phXC<1V^Qw?ZXk3HUv z{p80CGV%L8(h^?Zz)eh*2b}e0D?6H1Vr3mYYMcy~803->fgo%yn+?o47zFUu?${%i ze9;IK*K@hXl+mHGMlrPOU7oMyi@su8J>sAen0>s606)V z7M5=tOe&EojxJo@qZFLV5LUL@&7KIU_Y>l2K%wxKdh3f+ZA_^0wN=7?5z9)B*&4fa z?oK~#W2fPf)#=?B_)0-pXSUl;q#-3rsx+$BpP!Szq=}&RZ3$qlH`22^UhyLp8R91h z6g5RWWoE05*gWPTH@X+|0bC*oPU}2+^qS9Y#01ELx~8&re86i^0gM?C2i$7L zE1mW?qS?dfl#*g)Wzk?bBA0R|5B9 z+`SN-mKsS;J?6K1qwXQKV;7$(3`tj|{;X|#FcX=TNzm&Umwe>@Wz&nIM5kX$8Ox8X^q_et{H&l^rfby*~7AWL&Mqm2>pAN+&VC`sJ9G^n9n}N`x+M! zBx)`lN9ptJp~PFy>l?umzTX-HT3D&w^)&N67w0$<)Yg&=7h1kONhN4#_3HmFW2l!# z^%yl4PFv=`V3%8IQ_}pX;XJp$mST6DZy;sec=zeY`=z6q|J+&0O z+|&L7h``paY7GlE!8DxI@EII_z-a+7VBfJaB_+jq6WD!d-~l9e+1P+(E{HUN&;vm$YyYIa#9A(EiwUZ6J`vN$B!R_{qHq^5MaU3k}hr} zAkO~q4j}nTjE#9$JnBPKlp~WVItD8kzc)G^qn<5rh-O#Mrp)gjPp&B3Zxp3$)Czxh zCHM4;%aOMPADq4S`jJXbF=6J-48ll@-LH(@k1=eE;`RuO4RC&A{$T7KA2K)o zwh+pOZ87@Bevn=?1w}%DV4}I*khh}5z4#F0wg`l!zM@{PCkIdK%~+vKt{`ngTo$qt z<)jppb{~-;$Q&Lj>xdP5ekhnOS(V>E?=*pj@L`M7&HkA@ReO5(^dsw>QayyE7tLpi z5RspFv67^e|I7tj3`vpZe;QkadJ8cS=}q9 zk#9BGue1*sS+Z2qjZTz~-dZo0OPmwzC0_HEM@7Z-iCL=^amy?&F7vJN4!>_tFW-sN zu`GM6owcTJ^>1lpmc?yKt_@dEU9$LjQcpwJQRlh4hzzT7!Od8C6c3+iVj1N_(~4Wj zr>IMjKEWD3eb=)C@B;ldhsIJ5M2XAUmPI737ZkY)Mvu9gS)MC2q_>THpCn5g9k~)8 zu(dOIq_`U6@+@A;;V5i@Yrjb)q-o|*zI8CIaL+P<_N2XBEeFCEL(qd}++lDe_%2c~ zQxbLQy3paIN8}}euNjxO`#n9M+(;{@vXt}mRFu&abF((c^h`-ukZq_J^6f}vL2C59 zxy%f}qJo95d%C=aRg<4R-{||zxudkde(Y0u5zfdvj|`&BZ(!(e4RxH5_A-$As&_n| zh|tW+7ZvTP12N$LF6R0l71M|BwyRQtJQt0IXUN6|icRCR%X2@=Wler&em<9FxL#sS zvqT|DNW%I8v*2hVC8w#R-UL;Uq(I5VOD7f+N>o9?c1IA!)?>VI5NwVhm%jYu6?w0* z_0Z|P#g3%X8`5-mUsyi%*Ncyhd|1GX6(D5gHYsiHE-jjS7L=UWEaxy>6RJtk*5&M; z@FdhHf=4{8s5j+e8MBSzn?B`kIki4Hr!A#_kpPYc1oZ+F;5~OYx+Mw}P0+i%O669} zS;`{)6d|9`YJF{|VVjj~Yp>iOq_l1_!+a>@p+`eWg>(Pgz<3W2Tw?W>3z7><%k_JO zd%`MO%EJ9mjSK9fEhk*&HwxWTHu+*#BI_#*4cKj+n>J<}D_7_>LsvMu9wsc7(MI2} zlyy{H?}gUz^CAk=pI=$u{Zf!5m%R{KKTK(Q1uZf(y`GV#%uQ>en0tFJiBG=hrDDsV zS?mrD$0ZtN5Ggo0xE>s(Hq37BK~hy6UE1cVs`>bk&4;b&r(vE$#GxBBY^$}}LhoK_ z;a>FwP;=PrBkm1eaHGs`la(KkTiUXl+%(5o@v>y^G3iXP+HDoMt2UcnYOLeODMI*T zt#EL}FSf@dW1EA^sPTvKZd}r2zCvc1DvGVoB(bT<` zANq5oU3V495~cNN7iDcsyBm7dFrDjzWQ)^+7h=k9Br}^>_D0QXI?LuMWOgC81>c_3 z0^GzQ+Re4eQ)Fp~N^BbP+Zr*6wS7ZY{O_j=uLlp3madK<4aYLk0UroqCaL769BMTL};d1n`v?RbA-;?Z(c+@;p<{ zt>y~)-`DW*!Qcb`Y*Emao=Po~=i>PFhWlX0Vor4^{le~`__O=(b?p!HnCm2Xh_Yk_ zo*S`JlZCMglr6NCQu4jTQ7?6(MhhKv8x_iwZC^T$YRGHnq%(i_`_~>(3q3V|wp9Lm z%Pi~99+~ljskMTT!cZUK^lZBC761PKWB;Gx_Ndqh9=hYK15p$Sk66v}-#9;ftOE0Q z!6bqSAg=1p5q!Sw#w*CSXR=YF=cDwCma zH;c7YamPR)yE+KNK2T8J$Me^nzvtNCg5uMUXZq{JaM5RA<&3_(g&EXI^n&GF8R18) z%l=7}9^`CXirGJM@McPLMU+5eRDWlRT}&=~;ZE0#=hEZl*T)oxkUdx_M-b zi^0Ws&(Z9&5{7tVs^wsWF~Z&R+pllx+o!EQjY^Mo463s}zl{0#v%{%U^(sq|D`21V87tU*D|Nqml|LqF+^UE9$S#QyDdxjl5_9{Ir z=fa4}C@ob)fp+lgkqaH?V%w9YQeTgp>o|Ax?9|5wYPC$qn$r&gCe!gNPP?RFbWH99 zzEyq2FNra0&U41<3F&aN8ry!oiQLV>yxsHAA79C+pNfm)*pOFK{~2m9&i=25+DG}% z14iP>L9y~2+pFYQ6&Lt(JA8h@i*xXT_7}FR>Kk2*4FdRKuiv!{f148h*WJI{xchQ>kMwMd%eIr2t<3o{ ztPMGnqV&v6*<`FvgH7Gz`J#5`y+yTae=Z3AoLxc}CAhZr^n_ifA0hr<1}#_#FRt&u zbJG7ljnRhxe-=UC-B|Wx@bUJBDrRu4t++;@V9UbNk{pS2$GJg5NZ9`(e8gldUtb#1 zGJy6Z1H|}sI$GLhINICUBU{76!yzXJ5*H&SdytWVQp@MS{rT{LL%rJd_G9;y-W#Y5 ze=`1<$VkZL0@57nafZ9PEH-8(gBpXQV7UUL8b~w%SR!Ge`gJd|RZc}E`rW%LP;ATL zJoPQGyu2K84Uj?6($T@W!L{MLe4-ul=dP?i+6pxh&?%x>>@dSIDC8z&WVG<}K}8%I zl(*Cbz^%^^qtw0w{3SlV(V2C0<@M1+0Sk*YG`|}vkQ=QbAqOnLR48;jk0TzzG6$JV zxV0lhc7GXU6(HNv!2cc!{DzE+=pc*4&AtDYQ1wb6q!Xcvw@#16j?s$)z4M;0l$r;J?N8%+R4LG?11)ocgcZvkXj zR7z#idq*LY_)1&E$_OaH(6qjOe)xf~PBDg(MqLIyjH_en= zNX;r%uXzm>e?)%u=ny>AYXXR4L$(gMM`u{@pr+7NEU3SRu1|;DFDzwI0s-6_`9Uv%qCV*|%j}G^T&W)790!^eb zr~*j1XpA}L{rmUP%wQ;P+^;7!#m2?ucY~dX)e2W`gll;T>ltLxW;fxYpURV?B7LQe#=WbD=NfD4=p^w&QPl0*x0}4tsj4c#G<$?O z0E0K=#@$2MOzIU7P&JnniJAfi;iaf5m*b0BqwF?1gxC{GDnKiy% zK$vq}hlMMr8z)@b08`v?eltULdZW5`w1eb85mZ6S->Mk5D1IlLei?82_JzeNXGHS0v<14zk4u* z+QB3y3-YLl#pcbqO#)Cx)~+rNmfz-1o$givPKz)J@xLhs8t9|pOl^JG-8^(w z;9zG}+yEjYA-#3te2~$!f?ER)8$9-f+*1}MMr@IuU+d}$xTnZc)6?^w--Fa=a>Ow< z)d;<6`3@Kl-0RtynfQqRJnG*auoKhPj;>F?8VG|~+#on!M4KDj3uG4r8`o}o)FNJb zJ$UirMUzi`y`WEBSL1XhAm33E)3S6UPV>% zjc6Lw6nvaO-6GZhd~`O2AnJ=Rnq%AXD-hWRC{_w5x z_TJ~|zNEB>cx%5KD+~XNg?FCOO%NTRWWIg-Q&2>umlqPC$6Y4FOE!cA1Tphg1&Yy8 zQObZNmXiyOin{(XrMz4jI622Y+u&)mmFUW)SoozH`s_ZRZ~^Y zk>J$+l0u3UGw0k@TmF6=$jr=P-PLh~7nI@l1^b6z>t=)h%7~@jQqH79TT{h$jhumd0-%8ja5!S0FJnWv!!-`#$Unu0XC z3Hwna8vp9JWA@V+^7$dV6{@D{3aTb1<3o@Q{rT1<@IFG{*qDxn=F&ZBo&9!~qiz$n zu;yNE1L!KN?v9O%XQb*?*^ai%A7^(geEIT7l>C0O3g_zRsm95W*RQB2F%x{F8Rtt8 zXOh!o3_?WckhRQ|?J*naFNpC39A&L3e!P&J%yZ@<2aesWFHW|`1tu|K+x^5zlR3Od zRl{sOb2gnjex(`{yaNjThuJ9B7hf-jX^(aC+xdDj@NroW?!*R=+Wd44&Y^|^6c)@x zZCqhuQo>E+YA0dEEqFWCOQC6~(%el)Jc?iE*H3$6VObUPm(QmJ6yy)J4x-)P#SDa? zled44>fL!q&tN*0(hZm+q(BtL*!Va5`3@74IZ)%Hqgb8xtbFREl$7?rU@`1iDncOD zWw%)krFC?$i|O{M77tPhatH+oGR?%gER(^Oo%OQggJjQ^NOlS0a3)RnPQCmABLGyt zZ=bpf|vu@Ni=Bkbq<(C-^N z@MCw9Ci|E1-s^FMF^&Fs;bu0eO~uGW{@wb^A6TaAac0`*`Hf#5KjDkDU5LT$H`*T> zTykk|RLkz1rSz}-Bq8dejoW%LK1)8mhjVwKYB*;iJ(?Exnn|zzWAO!P(VGs2MlwaX%yu!(c@CagL}Vmn<L42Z?tesEkv7!t9$uFMz^~!G0*y> zU%MLP)FnT`g-ttoQ|FRTIge0lKF8d*w2iiqF|U*khSx3|{Jz-~HPZayi90M`KR;z8xt;W^N4_sLPQ zpGyQ!=eX6>S^63J_=w#GnRDomi&$8Owxx3WOBITr)gMAd7SLx)jG<;^sF2!id(7q7 zkKI|9={>kYpydPm%)&y@hu1fVX5VLKj&WG#g!myt+3s{CB;R2R&Kbl2sD1PrL_-Hb zI$ADEnfN7|#toA$6u-e_>~lf^qNk_fEh^6Hsp-PwD2NKlj z!w{vCp^>#hNqL;>4D^B!m5H%I)++8ap;Ko0XF-bC>(t!_n@A}u9(J?Og`d6NOp)5jDG^9qY@JQvag1g0FevgrPOm@`5!q{jKW3-PhkrSi!zO%Mi#U2 z>SZfZSdEdj?;&A*d6`7aQ^Ju(a#wt?O6rH2z9YR`~BsmQn+uKU}!ywHj| zc?!g^`K~0l7G2Qu8D08(qg$^^4 z2`Sa@2A3rG>U#S~;&5AM*$qZq!f46iz9@6EJ$u+?y@WL_of0_txYoGsY!ff@Q_K<& zknAo!t1cWoTDZwky16jya!!1LSAY%YzZL`$6=y0aUPPm}R)*|}IMLu;$`LJJ6N?^C;VQn`%eMs)oXX|Ws7p_`!CQ7G zZn>Cq#$W!YkFdy>3(@o}W>*7}#WF8Hd3wEajx-$P8iv{1tzsaE3QBHDq>+#YrIC{EP`Z(B zL=jMGBPAhha?u^qCEYFE-QBG3Sm56G{rvBE{_&3QTf?CXS#e!w9>+Z9oM(L{@nq3I z1a7!wX)LOfI4F^4f_x0fXCs9sO2E&Nk_up|w4ADq1@U3cIy6PJtnoy~4Kg>ykyly0 zA8$($GLU|xm|)>dVS5GQyA=-3JQ{HSqIqDJ*tW2<){Zn@?TTo=A<&3XG2E?=i~XhzkI@m9 zZpn#qk^h2Ho>My6C$1hX#M+mQL7ZP>o~U!>+HK_`(}>FM$F^IKWvX=Y>E zX+Z|UOh<Fy%cUPrts@Bv@%(w)u*SO#N5b-74srG>TI-&SE6jLl zaTFkb{6>Xq$a?(DKD}{eAf?LHY9!`lV+r?+6!!oA+3PO6vFYtgre-Xjo+r1StsaXc zQWd#xe;o=RHL%EH{U~hM$Enty+tYH!_e1i(_K@+#v7kICvS}B0v3{swH0=8O#fNKG zQTpmGx{rXoa~Mm;1>yM@)yU){o01dsCZ2d(VOr2;bB=Z?KT?jhxA~w{{-KvXfU(}w z%M(|Dz|b^py!4>XtKnwSq;~h9orq+9Z&vxQYqu zc_v5Daa`MDAw|R%mC@VKE|=Qpygm5rS-)!AjTvt(3R1K6f!dz2HGNC=6|Sn60y;uv zKK}`N(QGa(N~%lh@8bpCk2SiP4?Ci|hP!$;vvPtQx8}k6hZSm*z-wh(|7oQKjtg=- zWr)JOH(sqt)cf*eC0p0#WDTT*)I+=n-IPq+11ozWKkqyUUvMp4X`lSp5eTYZjTbth z>Kj_C%6DGufhLC_?aG2yq|AskuWbNn#!A9n9M;wzV5 zJWD)+w~AA#Ag3iOr-l9BP~qv~1^D^x23!4C$=sUt-z&a0t_>JE&i^N0a>6)w#Dx4s z|HnHcR~_aLV^^Kh6_tGJrWm8h3h*=lVgC0{yXa6L<8y;+N&dKxh*#;m)cXHNjo{xG zqS^TV-yoDJ*L3x+mwB5X0p1K4YIPt#NacKf{P^()3Ll5dVT46WjQyFM^IPs?jd0*O@?g;1cj zCp@9MK(_&7Cgr~Jck$$-zs3I#p^3-;ubc57|A-!bHO#wv|54pt{+CqOwK9*L7h+355LG(4>C&)PqI^PtG{KwYaE zGd?#J;A)u#6GqR6sEWnMhR7MIZ&C|p=yqh!aeoy9So#Mci5gf|<6^kQs={BV)1iox z2%;VrAQT+OZMQM5p{hFE*B1i3{FoSCFki?Ysz!-M0JQ`zy1Grb5D)JWz@CbksLYX{}=>~z01xiRFmF4 z^LB+L>T{lRsNg@uz)9C#EUqJY~ti5B= zA&@p7fZ+fKK|pnlSny*Wuwr=o_N_twz3f3}JNjKg^vVfm?X;0G8uIXStHsJM7*cY;W3ON*w% z)W=M@t1_F*E*ux zA%8X{wE{N&wXD|Am>A$)7W+0qLib^KQiauaxerjqPfZ|u`7rjc^Ev7R-i@c>X3GwY zj%Jj1E!nKy0oYztaC>J5+#Mh#fjzQr3Qli|CMX`s>^21Q*;zfDJP!gVsG8vTQjn49 z0EkFwK`;c6i~Yv9fcKrDwS9mRWfFub9zV7l9|sC9B(?w<12H!g1;E3r0Q;6#QbHgQ znE20Vsi;!NREPo*IXRXP?!gI`AFg36j}%)RSTrSy1w%*(+|i)w7w}H7nW(L;1#S-z zDxlNsRtlh~0PY2PAiJK zlYmNGRP>c56orsm7M7{sXTOLSnUAH(rkXw*p(`GpCU%G36#Q{fMtC8pPCp1BQ>ES7 zaJ6-QUY@4QR((zmBNBD!JXA7XDJtS9?{^7cYJm9xV%hBMc(3ZH-9m5{Y-i+2^*{oR zqB8-K|9dNl+jG;iN?VWv=EW%c1`pqK$!K<&gl?N1udnU;cV%j=Kf@iPegz1-haS+* zG;*{(kZU0G0kusGIE6qNfFoxc zbg1DwhpGYZZkK?EQd>(46jKmybaXnt#J_!Q?l09$#jtuv5zB#_@5(y?>x!T*3?HaBVcxhQ(}I#1 zK&=Fp6`1*K6ckB0Ihzp%F~AQ#0uu_m3D*NER$$))oPRe6reRwQ4;M%_AZs9Y*{osi z5KgL1xQ3kFWp)^Y8OVaKy zCE)M@yNoYGD3}B$pug}8WTn8NWi$Pm1Ni>>1QtbhSWwUuF-q_SwlqL&b_4q5*YT^v^;bI2z}16-R4!RUOhV!mU_&{8%pfn8{L1+H zaC7Qq%@+#s5N;I75WoSu_0bX--(Z?kxbhkxmzOc7K!St*v}Y#(aHgrb*%1^(V1vSX z_ifo2{I=;Y5-9bsL`VP)TNSW2_%LC=hf`IkYFdZV-EOM5skd~J3K;z&AI4g5bZ>pI z@!D({CiNb}S=GE8%JBF5Y+R8RK@`-);jsC4yS^Hjv4HLfJ~$eNTAFk$@A-Qm{Bx~; z7Z4nG*^%_F+X4$4Cd-jVa1!~{1U5X@pcVSu^!2GuwA2Dz0tozO0z=Mbh>@=bhMbP( z!b6*eXPbk^%n#BO(#IWU{DaB#01^SLq{>I#x}CLn%r1)QJG63nIk#{9_VG$Kjjyd_ z3^xZix65U@O^H+wB!i0Ex&X;2y?Kq2y`b!8pg32`RUDs}C#&LQw%8TdA}S^(rl}d@ z!@c$6UJHVV%g*%0iyA;?2C}s#=wF)Mh>8E<%h>=3~N@) zu`aOhz@s(Vrw!y{iHWEp4WL37ajNCGI59_SBvX8v!$e0X5_?rs>S2hhj)Qz{5Krvv zh70nfGg;Q;BvBp9I4^}&HM?=fG>eAsSSuTmNy+&yn}Ho4>iCn+pTKbit^=;0Q?IV^ z!kI+VaSkRp7#Ear56)~!#PYhrNRmTH?O{^ieHzt1Z&f0gn3-t?L_bn1cwbp0hw536U|^x++`kXSPS}FG`5qDx zW&Id@W-*>IYNrRaj}sGfu-T%LOa5aVHP)lEGtRt(D0Lf>{2=<7$V>L~k!tc2dw6<= z=rI&H4*be$HGFzYpO=0A^MpF%PK`|Odh78PO?Lpc=+ zxQ8Gof#WcwdBH15npXu@Qa<&55`W=*?V8``3R~y$M%JJXXUIZLCWVm89#m3Pq_$xA zKKZ`dy3y~^4j7~yvvtWkRE$rY?-#>*OuCIcpFjdp@LZr}5~|G)l{ciFU@?nmpNDPe zQxnj^NA^-_=WO$!YPkujl-mT&EiJ0wE#r`6yDNj8$ls!(q6qJJTa^hWrlzKjnF62R z!-5WUDNkBIeQE;j%IS;J=WG1hP@^Mda@t*lX=_((0vbP*!t@I`^55p&|EYrlDL1F; z`i=KOpamr&dZrptDS=ba?2GV#l2avP|K!xH@(w$4CVE!VXL@bY>(xj2Zm6AaA=`3mJdKw6g&xg1&E`;qa@(-=TfNba&dsc#@3kq#2Ju!_40=}XMcL%)8JOMJ{L+n9DCE50sW8L2-)x z5(TgV30rVB&YNiadeMC-IAm9yAD7y$gm9}Ir1np!R7y2}245Ja2 z_VM9iYSP9)V%)p=L#NP*^q?n(y1M=7TkrZe;OsE~D=*ZgrPBqV1ttktPQN%B`xjMi zFV-L{wcK=_V?%BUOI9fnzaik<`d7yc@I^7ZG^yTaz$*c%Rkv!UK)Fl74h^d6DjRGn z6I0V@Hj9`%A^kURrt61>ru41V)J9+=-&IaX{JYfJ)*T?)^Q}O4R5`?v@ZmzwQf2;x zf4{Mrk`7Zu>CtNh2%2>oZr?qsXK;H8%?p7F%Kj=95wnduyRdLFXoSzC+XRBtJttQ& zv1YKtDs{Y31*O*Y-Q>}%V7_E21apMMc`rlRY97?nVS}26nVFTD=?Zmt;1ulPn}473 zxK3nItK(^5s;;ep}4SJo1Og7RU!A8KoLSnx9QOX51hHmqVCI{NV2^u_tY z^K;w#O_!KoW%h*3S}+Rl-e48xT74EcViX|7>M2xE^wXpT(}5?eI^7db2ee82{0XD$ zVe?*}OIzv;MmJNFd+=)|Wp0ql=9u(PS)5aldRz0hJY;2lZ1k0B-Jpm(DjwU`DIvYE zThq1ebvFBq{=d`HNu41e1Y$kdIUuNfOi6h?z#BkJh;*QOkqYeP%*-{**Nk^n3XK(B zzmBVa;|O}UGiFGC{t#{nu?E3&=}=Ekk~oO?P?F%&3FO%X43V&qko#Xn)!nRsrVG%O&h-U>aB>0#w$08Ghz*&2Lh2hDaG)GrlP3kK zs}Bm~cL!59Tg`~b$m2t3a)_N^Ih4Jq#?60KG*}FGG}``4y#47}LFr$~*G9J9uZbLF8l^GZ3;J1EcRipYxg>YE-*khE zQ7vfz`al%l&C#?cXSs(vw$nhee}n;frPgdVldpMxS#^ zqXSA?$b?E(K4ZF3OmmHSaYekueUtZKA;qOK?|u8Zlf7EHPXR~Om=kXFeaWe&U^xyzD1lSUMqx<)l5u z$U{%Rw-LCSd@em^Q{AotwecBL2@R*S<3Id-L-Kd8?c|8N!Rv*sZTOtW-rf}>;&z0^ z)b}&7WhsU|Z)mSf)q$^GS(W;Ly!&)2FD*BHOC&w%0}*>FG6>a#rBC#ZeIhtr4Ry?n zQmBk4znch4R0`nsjLm)icNumwtYDW7i;Scug_!*O4Cf43UZJ?>-zWt5`G+YQz=p!- z=h}YtJ!>ne42;R0BSCUbICLB;tdIpfQZ53&J*ZG2q0Mg-tPWG0_B3)t8C$VcsJz`ZaTPtx5jP+7en{hMwQuG?w$NHdIJT3TBL67C zAcAh<-2G4L!sNUs3tfX78xwgwd?~k{618jCYqySgSGrZT_s7ecFg0pq6TE!HUAXbi zG;qXrv5I)s!g`Q$N{j(k*l ziDl=fg;SDy=1&`UO^TG}%`{3*TC0eDp2rNX6eR~p;W5ThnxoWJ7mw$di}gWyfC+!{ z0b@VV0bzZooB^mH7S#WDAMrR9xO-$ZHHqE@P1a`k9xyYSwSms{gbYmHWToVj6}DeP zuo>=C4x{#*wRCuR7_b)#QdD-ElM7jox^2G#={D%WSlQTon;_qeH+3n2I;Oh{G$Z_I z)y_Hyl-*BFqb1g%q7m*oFCl0LBHi`<(N6$#Jo%J`3L9TGNi=(aM?phX`aJ9)ir`3S z^;ecS7=tn5`GXV(h)A-vc3>rv5XkFvP!?ATYbii5RK)5f9-pBHHQ^dz-;ehJ>!~Nj zJeE8~5CMbB049-(ss;?vzYk>7S+6jB&q6lQ46D;!FRD`dZc*r?jD0PF^yJW`GMYh) zjx%RKu|Ve_9J#Pp%Brd{@$N@x=fc?8oYsN4OQCirHS+KD5@5i%b9~>asKD1MRKocz z`kmqhM~LGT6zdGQqIde4)Vi~3E7Ds7*Rg_r{5FIb**B65jdonB-WmkGRqa^p@aHNx zCsa6`eUkJp)##4VdB*{bul<6VMz?7;c95S;e#h>`X=@ifx65djgU@aakrgWHtNnNj z+!&GIqk#e-fY^|gWj>g7&$)$ryNT?48sHSbi+GWzJCI7U;`xJiDchhuat{`G*rrf9 z9+3R#!%bidI9F(g(rAbo0m^R$WtxzZvS)Hgtf&Yu67Vl-v=`?ypG8RQEOff>4;eu2v;|%ii3r5TaAA5~7Rp578@P6i6s!7;F ztivC~?c;Jy8FBg8+sG2lY(!e(*>qVzgK_xI46Y2l@!5_3#fEY?~Dg3&Vcxhp^|s7x99hV$vu zEg;9VH#Io`qy!5yASO5zGpS5;C@0Cuca~!yjMElYM}q+sY9!1>FCzhfR?$M+Xy%vB4aL@cEk?S zNty#@QqnigUnszVkQkGQdl4xO3HV0YMkP={?M|wIvxDZu_Mv&ZDqJ(ON8N*d>6=Viv>xY@h8$pHc4=MzkI)OrZ#P+{CBdaaR>64D@@+mSP{2a*Vh;q> zMO=f1EWCStasbgmga`z3fbPUWU<4W7dsMtW+mKZOT-of30i&|A8r*yz7Z>3NEN4DUT7Y64y^Hx1x)0^lTONRm>paBm zaTysBNwVPRBwZ5yC3dC!2=-4%_w&alVHJSYM?qC}dCw{M%a=amJ(zigo!yaqAo&ldhB)lEd>jQUSM)^0=enkNkVZ=C7Bt*5w~l@wutO`7e<&85Bz{2T3aYRndY1e@ zGsau>sSHSsi40AZOdn_99C)zgEZiize1R{3w}4WUG%@;N6H6Zt5C-^Jq{v1p^(lixkcH8(!Mioci; zI1E-mO<=nL03s@r;QEOHjr@1GQQpWh^_{hm&B(HLBSs<2j!Ky%StXo}ts|FsSqjXU z`RrBj-9K&`O7`(&t_-UHJSP75b{BqIi%+*{`$JoWiSv7c@*3YyeB}Q8*QKJpi?QX+ko6)S zh8!E#7*bNO+-$!b6;%v3dCbitsdCLRCsLuUMn62a1c{X*xT1qyp zUtaP7F>`Bw&!oBaGDh<2U~^22z@gLg_++mc1G~`ogvh)$Cw3NEYU|ZK3pE1 z`OcQ4Ka8O~Af4`Q*yCOsDSqu3+t(d8^;9K%TxCyKYJ8r}Tir!2g!Sxw1ZJsP)NNs3 z>BLe83ps%jpToothvCKPp$#_=T=gZ({-m<5-4rWk<1#KNedX%YJ*2MnR@d8xW$KZ% z-sfeLfF(-Ocv+0H*?qeO^OJ*MxQ~T6*&?Ah^=?|sgKENWiz{n(tS;6wx!eW z{o5hLcY>H~THpP41z-j6CJiLwQ3YHB5C`x+9bEY^hB}wfDI!lGYOP;{aPQfinU2Pi zhZ)k(M;zy_zoI*>g@5NrSQj$>m0STYOL-po;GE%;(E}B;qZfBm)1R>C859n#Ea~#! zO{H(`U4QDJPCu79I8tQNuB94fEYSWmZeBTdq`&V~x4cqJ*^#Z@r@Z)z6`Qz*wCOC| z8!IlaXdyg~d~-(mL$Em*gg=Iao|pN>?^Y{)=EP zilHQLHS2L|ch-gopgMQ@WOcYuTu?ZWh;UWday*T@=w#=?!5`p@&yGlyH@3EU`IRH9 zXjX6?{Cr48`;sZFJPnU;}jf z+J}B)jDHKLmH{)5&??Y7(H7~T%#3|06lE_ybQ$+ru7y7_---&8Y7|c;DX+R-mm-LE zok@{)jemkxB=j7V1;Eq`L5fBtAF~V8v?sIh$2GUMOlw0`(Qm@O$uYR?}OGb<)DjKP44m#eWYE2&_di3V)_f>zh zy_)qoWWGJ3+cxx|;xLw3lrvI9El*eNXT!aQpLf1?dtjlH+-C|2BWG#;7!{WhwV#!! zk=|Zma@OlymKQmm%(l~an?h{lHe1{ctUo14o(4Z6@H5~h4g8|l9BIoEZGFxVkbthE zmY3^Fd$i*iToHS=YJlwMP%&8ZuXH22xW*=)3q55+5~qw30fVonDEjM{ug}`PIGSB) zuMp5E6!O2?{XVR6?4x7795zOfXaQYG72Oa4e+qp{I(>$(r{T$SYZGK;cLK#5A?BC! zI1QhfNp_Q*>})NFUSbBQ2#{lclW}q5Ox)jHO`(gFQDM|$3^iSw7~$!=>nmt-(?A|a zLM+KU^2E*-lMruV&Wl+fo4|=G^ijt;PQAD>_TE|QxIpujLLugso%CxxALuG1dFs9) zcI_@`kku*|=a#;ETT;H?U_{{`B+BBMld{~?@*||5uF7#~(jq(Ryiv$e5C0A?TcP2+ zcfH;)>Umm=&>A_qd%I0tmS)5IjoQjXK=<+5no({GdX zvi_`0-fPCa!#&P7YJG}_-Cc{XU&;4Ptioop7+0Pln#qbFSfgx2z}2PXVk9xr_kXlixGI-Fab}Ge<`_nOS{j zyT>HwtKUZZE!OA3c-amY;)n=vY2X{J3t4NRp4_XGJ6T8WtvfT3+~fZ_R5CSKwxKdn zX*_GFN$ak#Lsm7m5Vn7K$|Ji$31wOKCyn}SfuB)6m2BrPvTi_c%D>faf`&1 zr(nKYgYT$42Dyr780gv5kFi1I{ z8JFi~GS#3b7ROR+J?U2RI-T^|$y~7WhQ*4-9TM8?igo`nR6>h}_Vr;KDnzROzPdP^ zLuUHjY+Os2G&hrQMbq?bZ!w{)wD|l@V=+H^6qlyDy0s9wpyt}J%S6Y97nbIdvi0ls zxgac~Co#Q!XD?~7wR%gLJU3_0wVoZUeHM9f6Z;cPR}xR%v<(lX zD!%m9+9jcUgHZa3erW7{w%gCr1~oVZFWhCTQ{VvAw3M zy{nk!k+zei<#z^#F!m+!B&Ir%g4+gOQfqTwvS(Gq)X!KH4KA;L@$&i;<`rq3TXzox z$-2YUVwdq&7YUa+vt7>LxA;;v&Rx7zJN_Pu)Hvx^nmnb^Ey}_%itIJhl1lzU;)zdM z99!u(_^1+9QA9(NumQR&{97&f_ zj;=*dwO-swwV**5>X?P!$iBc0U2*}Ws$V%V;`nw_YkkkyXm>ugE7;%Q>At&vupX^A zsvNCU>9Rc1vouBZac?1<;SPx)X^<#+WOJf4oC>dBCOOTr)zf6DL>Qci z#`b4c`HQ?cakf8?o2#~fE;>I?#;P6Y2zSq4;4=tkpl)U$7Nt1xDpe{z-_ILf;Ya^; z_;AJOPI7P+*QOvkNyKt{5xKW>9^Hpj-J&Fe6y_185F&!kx#4S(?)Hz&YGYpxoF9n-`ZnY% zvoH|%4=Rq2_w75*Enk#xrXW^V9BoY&S`qt~y;gnq@|t~ZFp3(z0(iw$4zy1mFW-(ic`rQ!fn}bIf&j_U!JLL5JX34Vx@7hb)iM~#+UrZyWm~^Kj#pm zTE|=E5E}#`&jq?O&g0hQcul44)uazR-(8CR5SOtlfXDz?Z_TWqyI0yi(O6@m^Fyh;z(kGXayL^$o7MI;IzNX{9WK6qftT zK6W4`eY#O= zVy$r*V^;t17`0&g0UsXse#q<`g_W|vix*;UFSeS@<8yWi-M9J7-8K-2PSr582-b}r zuZCOrc(@B$!WWnH?nY-V1l83)@$kYWVr{!8pHD0!+`k0>e$y9L>-}|!$2+GMNS<#* zcSvUEaFb>0>gqmxs_lsmx#Zw7QeDh!9kuVPUg8s{mAdZIS)_WQ3d~t$bN5I*`=yap zyOWmaab@{OC6woMGRmc9br#mL7Y&bAok)(lasp}EFt$Q-O(6wWSc-k6O_W)p@l4Ie zBcITwXG8``?AzF`XGd_Br4*D_Ardle7rL~yG~HzOLKvd2vKM5a%?9c<&k^~T!XzNb_* zlT}l59*dl`!0S+RYF2O`?jP#Ovu|C!sx%;BWNI;zzZ0r}KgexWlULjs>1;F?5>cUX zoh|!pZy0wSpJ1^$V8C)b@2YS1w~EcQ`04GbIFF*RZ!3y?1ylBzAOGT&!ceMMBV(28 zY&EYs@(N6rtE)O5GT`Fc*KnQD=r&y4I;nKntll4t4GjO58G8&*+~Y;LEi8NbtY zo{6t$emExZ&3pe_jM~L;MPqz)LEmi*j@6Rq(9`FRi5<&~X0mq3;qzoPYq)BStMTn6 z#z%?8jTJvZotxDWv{I#F z-GF^CI|-8XzRC+gEjM)7pfc=(-|^OWvE$0v|bggw(#+|VMy~H zf7@f={W_+(_v0YeWwi_%k_Z{{UT{fmd}{*xBZaeFmg{!61@hPWjemM-DSD@u3jGF7=v(v%fdVeuP|=f7^= zj+)szJ@U=g*o)ez$}Rh~&KOFZECn&Xtp*witIuV}~bEl($F-49CYE&M9(# zP8|e%%jO>IudK2=et#>hf=osk`;#D3gbg}L&=ZcikL}#ZFA&THQ)NdMeP90Ooi2$; z$bH4-T0?cg<%1O_kpwl_qw!L0mVLq`j2(8(--@+|TAXawQVh;5P+t4u3QNn!k4cEs zuW8FmB%X#W#~lY}8+--Aq!ALyDc`4AG%m~b+SnJsK*wK6;mEBob8MKH7+KC+aWPo~ zH6E1|;T7reY3DZTG&?6qjk4n!vO9!7i|1_AFHcR>w)*wn*70o*EPj)nq7%16K0s{I zr%rDvMkkA^oOmb+phrj??zw6$u5lcm%o~zpH|qvOOI`{F-`%|oswB{sY^Xi4{2Rub z{=9k>VoQPR2f91D3akszHy^Q>f|}TGp_7m|6f4eW5i;m8MTnB;Sq@>4`NKrZcQ|U- zdw`($0qp?~MvkwjmHwfcy@3fjNsLmC<GZ>SZrD5Vl4 z&7ap)Rpq7d{3~3EalWKt%-aA@&1%S5H42VO+U?tXg?xLwVW!$UoTpB2+B?3! zKQ-m?{{3|E=Cm~J0fo|huBFK<9h|PQ{+@whd7dON zD?e$fem-H3+&JOHBz#2G>dDt;7ludVvcFZ9{e8V7)z;eFabkkvq)fcUm>=CaYQUb@ zm?PEs;(E|5r7n7?p~&K37`6|a<$@jYR)mc$0Xj*+>O|gp{dJLngxe%68$;#rTrRaH zD6FDvAlrGXCN56b&c6TNL?^?OtVpK(oRsqH;)G&^>tIp;4!ujO!_6jLrcUnj3XYxZ zl#Ga{>|Mc92m$mJO&3PbuHLuoq(w_}BKJP^@MKp;k7*oCpQZL~jtUKigeX$nY~9}& z|2a5vw7qk%K2otCGLC=Qcw&IkLeGU$5*^AKIga&t9htR{2%xJC^5xo|-Amvybvl}o zWZ>g2)teE=G;>*s3mTI*Ui@^uQvTH|SrwjyHSJ#5%=R;|_+h(-or7?t=lX#KpM&;E zYvqwnb=&&z;a;8|GLv zCR0i;S}oa>7V4RpvjCD*K|FhBHn*xHF>UN7rXtMvKT)sI`%Sv zLqoH=*rTm1YJ>*8EZ3-r?dm)PoJ+MJ7L+nq(m)QVLxh8R{3U)lc`tve9)g1%FP}mb z_`io~{~f>4#;m}y`OaE~7F$JL4u}7zWe!bf^K9v5#)Wo=8MW&Fbwz0FCt>USc&@3r zQKkFh6m_|i z1h$QNJ38H>#d&GqaeYExXQ$%HMJUQMcPb~T8b3d)R#w1TXH6ebzQ%z_*f?VM7^r)7CIQVL55BIi`T(3>RX2Ip;jk$hyK9ji;r z;y)5@8MiD9h9G((uG_D2#eK;#iaz$B$cLM1QHtWa?j6vdnf&CbtRoAy(GxlXsn|9@ z!^@^45N8=g#Ty30XKMSS`EgXiH%L7FGYJn&4OI&ds++psmmg-^ua?>GWZ@nV-XXNo;s`w zHb$vwgpECE;hM6(%pBRlsvq7IY}4JhJwngTfcGh5cx*U`U{UlAT?&;JXXdi!wIxD= zb%{Iu;Sbz4<9FB!ZHZsri79freE-9d8yO;i1~CrD(xTqIi-F<(y1ya2{QKO*z{cp$ z;kWT0xQF^&IvbSamfG4rPe;x7#m2n+^=1ZCTd*S@RC>%krnz;CAN@?Aj;AgaRwK9Z z_f&Hvo}#)0l*n;cwSDe9Yvl12Z+`R*O?GQdz2>ioL?(cg#!YaCqLozgTvQc=UW|*j zMt`23>rJD+cSJ^)-{MREO;JHk+4A?j%dd;pxm=l8|9bfmi~g%@?i;7}-7i1H{JQ$O z3HhO&wAk|Z>6nv8UEJ!v45wl7lVgAK=TZ#20Wl`B(oSS!(o#nAbiY3T)*GF2#O`g3 zk66j~A7p&hw0iNZ`BkzFYU#6sOGrX0y<=(RJp*W_luDNu~y||&oPzAMhM3^df?jVL{BeqpfKOt zxvTDPhpui^84d&pj6utw7uM^v{ejb5{@A)GH`7E<4v?HLUXiUU>4EO+8~ZRpdKEri zC9blm$-IRxfSz=;HHX2vr}xn!!WQ_BK8^HIPxVxdx>*xQ^=R~^GHK{a_YYhz$EZVv z!^JjfHG{y9a=PZdBJRtN&T~b8#M8i@*heIRi>p-Rn7c5Y0>}TE?Vg|cCt=dEGG_tw z^eD_Tm)`KzfveM5AJ+Y3xdWOP~%1C1| zH%;`{PAYHFGaES$YSvDS#cv@(ad77A)trM#JUeE)^2;zoW*5%P-JOf1H@Y_S%T(dB zS0cIc=<1es<;qm4Z{32azLP!U$k(>a$|6wN-Id(A-_zA+^W2ooTjq+Gz01 z&zw7ipsh$%h{@)0Yd&o=_KANV>yktH31!! zbv(Tho!P_cT9H4;yQ9p+yf=>&q71?wn4h@*3P7?cQSr+{)^5rlKgIAR5XHoFiF*sU zoeV8p$M}sM0%Inz%GMdbA#J)HcYvAamMQPlKaR*B8M(9P$J@sijn^Ekw}+4wx&Q2| zRBQ0%2q%a|#VeMoHW;_}FcASxDc0JH?3DYMNyEq21Dp*14CNK-zkH>o4Ql(XHrm=S zDaQTrMuexJiqbLqo_l3#YVwn}eF7B%b%+$s%i$0pbpb%oX`1rdSW#A5r7VOF(i@}B zqJ>X*bnkx$q1sTu4x^F6{rl>&N*o;>A5J)Xc+LAp{&Ssw2ZD( z4E?dt=fAby`%PjVF!9s!w*i-5z^H4HkH@S?Jw0^YUmXb-ogI|#_p_F`xYBpc7p1Gm`Ci-cZ}(Tl z?f_zAKc3y2m=VwZtn&RxOZ<{Y%-cI$6&Z$YCyb*(;iC8Z0#58Mn@d9R|7+a#{~wTh z7;8xWrRkPquKM8R7gihH??i!n)jzOeIViHbZ#xGE&ZKk>x+Bs{aHkUPAseu*1GuU{m>ek7`@$__I< zKaYC((Mk9;;c4wk_Mu%ij8AW~72e4l z#*R|FOGVdWvqxyIp3khR7RLK?$J?A?&F`6f^2iDQW9A;pXZeqN`*FRztmU~bFJtB+ z{qa-UaAo*G{dopML4b~MC@dzH_g;23!MM9&qI73xtnZ*w^1i(vFU|NX*3q5_-!GSg zi~Z@v(cs)D&p~1pV10nB4 z?Go*xo{tW@orX}kc{xdDeoepcE$d%^el_|1|9$PDWwK~iD9L)QPE>Nh2PFVYDRzfV zZS=V0_VLV1i0IFj9fbp7^{5v`Nliua)IFlhey)B-TzBesf9GlW5$( z3+NPhzGdC=c&h8J%5#H}1i0Vi!D87;(P!nlPdyx(-~7JyrCxA~!uLiEGuk~L1GL9C zJJEFVHYqN90RHy=CtHH&mS6wT-CzFVn;-b_fpXT~-hQDJ>z}v0lQViF0%jMS1)zcj zQ!?`4<{2CD$H7D9AnrlTmhv5U2acbNAGmEePjFEr{D$_I8f?_+;g)ZDmL#&oFF!`R zC-P@d9_pAmlvgN{=7ddlRjiWBptylCn!sZ$vFv+I*$upIlr*HJDr$6yK%bQ$?`X1q@#; zkeM)R!m@BrVB(droZ7)a(YiV=lcu!AhLdrlXX@q9e0Rg(i(m^sI(Vc0yis@s?K%2u z6Safq)O)L0Meu6+2nh#sC5zrT-Og&y2*qEZ*(VC z=rhPJz!kRqd=Z#?D1A2wEhw=+!D~_6;UG~z0v;GTU{nBJq1c~(fC6_QOT!uT0kmpG zt0{UcY;47ag%DAKYEhK@i`giA3TU=bD(vIq*jhex36s!sRfZ1d4?Xwhde=9~}t9n2@<0`-}?C-~~W6>^@)os?c9i^xeWad*@a z%lb8gw6P;^2j!V$mCS<~nvt>bFp!}^ryxQ+3=%&O4S_D7zUp*)P^1NQsFw?fb&1AM zMuRAVtnOc^j-7xpu>O4}0Hi7@a1f&>NRC_p5C@+FzD*Cz-Ti$~Z)SqS+f9Wv9rp{UixK!P*a@vXZE z`qcA@=gvrNNtqW}GeE-}X*wShR2hjVZb&TeDy1ViKL zR)#8Wd|WPmC8a@1-mOpLRW#=RKm7r5ixR z2L^B1da23mD7+;bj$(w{0(U?GP_ohDaP=2oM|-DZyn)Zz#rD^jo00p%$h<_mt&om# zNg++yR#`=UjknM0$hMAxxyjOEzNf0~Vf5k-3=Du_ zWCon)llT@0+TiQp@9&QiHJnj_rD6>%5$AE3u7-(vKn)%8cRktnUq_v|N<<@#!+$us zzr zcxJoI(E#87k%DuSX&>RH3WS9J}mkHaQ zc*~V>9tSuJ2V@Q(Mcy6?9w`gB;F!WFm+B!`Ib`y*1p2!BirU2S+!aih?d)@EqCiYc zOe9ENA^%u2v*0pD0(&2H%g7PvXOM$SuYU%uu{#`0Y3Hn4*bbI%b|0~LLA>t{3poB* z%-RH)hFdB9fi`&`%724=(RgFqtLsG^{S+nXhPS2H7KtiOT0G7@;%r`wlOHLOUM}#_ z6puS6V)bG0Lxl>4I~a`edTzqGHwXLs;7R;FE^aF23}y)O2oAv7h0E~u^|dHnN)B|H zB_lp+5`p!l!M#3+TLpazc^ZohZq^mH11@$xzGJomxRj`=MS<)1HC~7$5e>`Vp1O5F+fS@POwf0sr^u&)_Jb{M4zvMVHC3 z?_9(dwgeX#K{2qQdGqeZtri%Jjd_QYM)JA>%)K#QVR6X(8)Rp@w{Nd4t-EVj>z{da zME5Z#BZ_S+Dz{3K#|IbME<0$_|EqjZ9;iI%wJNziH>Hug(!i*qQ6w-SX@y+_i(hHb z~T75zH4M#2_w;7OMcvzr+O~2E9lM3Y3v$RCXwkh|f zpcz63if!}*!Wx{62W!+rn7wQ`^PQwORz?o&FzogbRx-ar^WcHR&j4oS-I7F!-pNKp z`n$TY_`$-z4mO!Q5Us4apA^FZZ|h4-9C5xlxVY}y!Lh)4Ic((X=y10KUsTk^jFlZk zfs{uyEu*RmiHaV$U?qa{Xs0n;1Yu@)=m<3qu-GT}7>xbUx?bvyxrb3z$@Ap~!v;sL zfzoAruvY>i!Cns=US!Z3svXitKm!2O6t6jh7y!Fo;KQ#Rn{WOxBriP&8jF)z>`wrxx+9+4N1PERMgh z0Q%M;@%y)+@d1Unt6fYEoMjZWuT1;w_y;)dFt{#>Lr`{$q#TTlP|*t*eN=+n%^fDy z4U!D4s`C^BLX!d%?sU|L;0^tDv^wQY+<~(R4yyTSX~v+}I0c2T3Y zM^TUvBt33LrIii=r6ok^Mx|rZohnL)w6t`0Hwc2D+@+v34uxV)dG#Vd*a zDI*03Z5qyK2ov`fM;O^1ehl^Xj&h~v?#`Fh;lF@}6qG+G;j~!Zdf+-QnSSr|-S;6O zA@B-)eV^v|US3hl`}DYVAe7dIn?X^W0KU5R=GT%3;M z!YF?bHi8Rd4fw&cE=>8LuepW2eOt}*w7-|8@q40e2bC}S20~>yW1#;ZpY57d#$2MR za{!)p^9Bu=rfjj3U%<>{E^zDQ#65iEI4{ePA3f0Fn52Ia;l5l&K&G*wyuwWT2T>Hm zcwPd&(`%ZTS|;VLhgyt3OOwXWkMaGLU3Aa>qLfbBpH26qCe}*V#$2i0{axWso(%|S z@3YW+$sV!3AH^}E6sRMvVbj0J>Y>5Hc0?YPv1C2tj8O~lUj zC`~FIFY7PlpkVEKbf%P@aGf4cP+WaM_g&qSZh zG?Om~{$9fYn7gn4g5reWO%+D-7iinRUg}M(@$qq=Lp--?w-v_UFuW?E)r-)@<$IjJ}B={ zkXzRwYEj_KeEs{mqvEf&kRKJg`igh(F5NQJT_D*$@YX8LL&Ao~)bF2vBr4T@wu2_> zC78JB^u|x%i;owzQE5(E+_Fj9{MV%DNWwpW!5suCsVRN^ckqbGXPWoV)dOsnK42JP zUN{Q4DPI2ypJ(meY=Ex&I!^MY#pXmxXPw-_{H0sJ8*j@qhg2;Egy((fn~u-b{T)Rl zTj}^pP!C1OXJcmHn=5>B6W**G-%*v^@_gPfQ3*Qg63bV0-BS`A`Z1;5i8mEG*^~3w z-+xO}UC^m$=(Vxav3^8SswvcTkv>A{V2_1f_`9F5k)w3?##*!IYJ3FtIu7=f&Honq zX#`$xz23HC?P zXzL84uRFPe^r^OM^La2$Q9u;(x+7|jGq#hJt^VVBp?_@4acqE zCfXGGO06f3YKAIpHS(<)PE+r8dq&=3)CW2ij3k}LI;qSJuo^L?Wp@sBLM5IiG9_34 zVe5IHmy(*=gluHuckw+0%v8^N*AiVH3;KRou=^I^g#7}>TibDG zBF&Ip%t7$ujX5BnM2p-3Ilz3)rU1}L8ed-m83V|~0R%z$f3qK1i)K=wiu5|n zpNNH<)o8B!Ir;BE^{yCR->Q#dj{GiXU0lvzT=QCWP~{jTCz-xvT z&FZXRIwz&DQ}DiD{NLEk8}r7C$-_z{Th8m6torA#1B@l24^FJ z`?uaBl3SlYzf9nUuvu~kNv5HPcG|mp#|tdrk&Jda1a2x4wFqa%78Y(eA7meGKwGhv zHAvAg`ZKZ*M^8Hd7*(U-3X|A7CYuIdL?;XDjp-_;`P{GMV2>}7H2CiDU3WN1FY{)K z|5cSD{&)20*3-SRsw})tPFd+=%*Y6PF;v-ro)w67HAB9w6XVEQo(qyz5ES_qfV>ER z0It^%0hb0(2k1=Bv(1CRO9fDkCe2l%>g}%(0v^u6n**a4bw|g8=~hI*2#Ks9iGWa0 z9TfVoi0DK#K{encQ1%Q7*Snq2z1kI7h=}k*GOZ;su=(@+c1=d7= z`-?<|UA}nnV)gO|hWFPHlUPZ|S!v)cyk7<=Vo&4dULU+UMCI@B0EB2nra(&V;H&z+ z;9oyIX|lG@EPEq8{|<{Ds?3l`S&Rd0!uJ>SVq+hNWbrs) zO+@+r0jH)_8OR}mn#B8SJb;&JBdj!7g-Aku>;1)>4kigYoKc9I)lELcM74T7k)i2N z1qNOxd9eOd;`|QeRG^+?T)t5B&>s9#|1OWz($UcYDVCxYPg*1aAqSY}@yE8J> z*S)^xxq%Tl0Ur{BPCtP@@GZ#4AZg%&9vfCx1Wk&cdN2o*I#)v&RNG<10qNL>7HViM zUkgz-BG!v}s38QX2l$JEZ`0(7sMP-czCUJbz8{}YS67!Eg$lV^_u(2$FHmda{|m7v zu$u3C0oV^DQG)S-bo(beGkl834|g9Qt5F|;(<{e=2Opa>5!{fX+8kaS492CUF|f2k z7}TTD4ge7F{S~Te)eka=y`c32&?D&VYvX8Pv~QbcUHgK7L=50<5#I=D_5pa2vqM@k zzbFRY6{AOc&eglK=@=;HNpfH31G(#dP~QkZSA)7IB>lZ)P8?2)NHwL2D`OV&{&BAM6)DI`r39ny|nr}IYm*8WtpVqNYI0hLI7=noJ z29s=&#Y`t`F7j%&l%!O1v$F#*py1PLX}JL@1@iVmsLe28zCD&7k^%sMK+NM?^xhAq zHelBwI2Zg@d0AQ2ijOhuz>bqlF{L2fBP@Bn2xhdQ%6!>eRs;zDiR^)N5ojCuvBAxQ zu62Kr-=I=D3bLK#)o54hOFYhWz+G(- zTu>WC4-w%d8!0UiP^bb8jC>j;*J~i+2U8XlQS$ZjuYG~_|B3t+EV%v#fA-$nevk0@ zN}$yq$uW+O=7()cc^Qur|SY9i~N0!FQkc4@m|=E-M#U4Q<5_RkNKAmcH?`) z_8v`fY(-8sv$(`O@{`I{pao%}ww4w;28OL{=A=&`kLzr={S%+`6Qfo$$SDMmM;MjO z0Kz2!;w<-K92n691)J4dfWS|T!PnrWo@r)Ss061&2)!yNhrP|-O4mbeJuHb)%ZoL5&sxwrtk5TA z_pAS!#B;mnAcWRhG(~egUWtv`u>Ja^Se%8nUmMrnGB34i;Xn112TlOY3WKWNtzXBZfg;GO}ypZ|5;z&*x1MmXUJhFlpPRzFwp8< zLxc}Li-2m4hglBHTu_%82Iv-ajfH*5Y!7#)5ZiMQw*b|A#yfWs4#2?)xyGjMcK4y} zhJF*Of4_Y(W{qJwSZ{V`u5V6IGu85IuO(MdoKRvkg|qp`=8}^6_%zo_UJ3hNi&E5Z z*V;?WPa;1$&K1L>u>$)ti9-8uZ3>67{y`MVMci@GJtj^r)j-t}TivU73O(f$$AVk! zQ@ufOR)wec;Jh>Uitq6oIK?YDrI#>Y|0v8I&Ngn8*?>$~-qplz{mk@qEF2uz<=osg zH8tHu6ciMoar^S+%f4hB0RaJKV>PEKg(j;v0U46Ru23^Zqca=Fas5F%47O-k*EQt= zLPE0}yZB2KBqV7ax@TxW^M?Qz3@t#NqGW3#bw{469CQ}}GkG?_&dq&*c$R@|Aewa$ zq|K5gV@6v0Pu}4&lgvWaQXHF?h&8|A+wo6Q73_we;Xe*Jyi(d7;rLffQBn)L9wGPf z{p&w()7#}M4~$fFW?*f32**oZOus5xV<7cIp-^;q-ZUfYa8kjmuwl5pl~hkg-)(8n zT_JJ=21>c2RaCxE3Sn0&?Xq#jEi=ihiMIo{=VQ(4(oy)?nPszL5#R%8b$`U^PWX+LBS zUhP~y`aWt^&LW_ZW+F#d=#(V$h(|KQ`XrCow&lEe8!^U%Q&PRmVg{fGSSX$V=?N!F z!IoPXc2D~L(>y>WZh(OsF^@yZQ5Vw2Ppg;*G6YbSDN9MY&%?vR$k^$;2`Z2PAg8*2 z5>EnoC>kzRdgUTeyog#Ee32>xGCc6l!se#>b$tAIvdHaoL2PO=G8PSVe>R7y8k#B+ z9kh3>n6Y%4jHmkznE?^U_mU1MyeurNj&m$_;erX0CTZ?nWwRB(q_bGkEInLHUKTzP z2ERu{8&%n+5a562eLtNX%yHzIjtVZ)+Z@VuV4<)SO2 zeBn{r&9tv``$^gnc$=3y&Iuaa_9lsUXSNRyxT=bd-sLP-6`(Eqri)H^l4)+E+^kx} zig`&Nqa#*h6kWYiCV2bmqAT4|`B$AWazO#QGy(pMwkfV(w0;=Du)?2oJjvZ!>08V& zzr1=3dp?v^MoY=dR!%u@vbnc^dTsgQuEa;oR8i@ysEVx4=E;pX#o4@bgGzccRi-G- z5Cf$Y$Jy_oUUClw`egFy-J;$oG`Ea+As_Wl0o!&^)Et(KjtUC zS_(s>!af%dT{E}z>fP|u@M3yruR_8hlicYk!XWL}6tdEPwKUzJXcOldIR)!y9TELc zcMBR+BQQ2E#n79zApuPGa&N5({QoZeIr(XYV#Pr^e97QckL{d(^}_Wa z=KBT4W+riUbA7`fBV_kfXPmd%#XB3ZvD$uIS-{Wq%CnEaVY`mb7>-7_@#cq|EIH14 zTM{ZRv1YXp&>xlxPzLMJL+Or=j+K=akmv-D#}&6W0tg(N1GHf)yAruTTc-?aq|I-? zzFuMAg*0eDUqMjtLv>w^Y%>TlD5FF53Y7%!P*P49mYwYOHEp%g4UR;@SFkazBr*ep zF?(^#;x?jP%`H3W`tFenlL?gqR#pbKMY5iF3&#(``qKHbSj-qbs**^4f}S=KPqFAPo6~T@=PiwhHHa#&8X3M3 z?-zX|DoAsFlW!$L=f-QSl6d!X2HyN%F1m|kk1>M=`#MKCY5B(0nd=D;eIKe*qv&I3 zo{;u3v9s`z9y)8?q?#)pzQME@=U$6#EEb&3sh~x6-PY7t+E2%qmVGvZ-Q}M#e_s|) zy8S|dmJsCL3zZ=3(_svQdx~X?MS98-^Yhjsi^um zcsq4!|B+*_P0c+XA5l!pH5$}R5(O%+Ejnh<=COcXxKxl2iI2XEq(i@tEh21WowbXoOtm*L5qrK;BG2sornKQ*tqr!r*B=``VLojoqU)w))rmMf z!Dl`DwKH+*RgUl0x$!(Yg565lZ8WGvIxU*Co~&X=d;T9KMy^>Ki>@r)q@y%PZM)%g zYB6lL3zj~A9MQ2Ym#g4eN*X{Lq_N4drweGlUhno%+jAYrG(eNmD5nW(TKuML=FL|l z7&wA)rijszZrR*wmUEChu#9h{$n*>E)A~`0KA~@o30uArY8KE0cbDc`s2l7nD^Qq~jUIRVO-*Zd(ge3}ZXi^%CA04ZUy?V9${KFD?rmzpQL4?!%+yPtZp z(#EC|2UMq)3*5hbY_CG$=}c3j2In>%l&j?hew2w+Dn(F4Qlse9U?0WDby12L{8XHN zYpQZ}LUE%%rdRVi^L0=6yh}+jk`7|chigjx!r#-XHi^KNe5Y|_C&z5#UiaJxjWRvz z773>PW0lShy)P#XPq9|1gAs0j@AKtSy1ATuL*4Dixcy4y<1Wn%MQn@eGF26=sthQ3 zVIQ9FgN3WIElIu2JiDJ)y|U4Yq7z=tmRR3ID%DrM!tkTwk@;3D(Q$N&X2Bi4vyiHOMdjain)Dure%L~ZM5sz;(AYahrt~A)ZeUxL%O)Q zm!wD|)r!sLNaQoCDYdX#|Mq!7N1eUhiq<%pKfD20tWuMs9pK-P-<=9-sg#sGv(QP? zj-=wn)dzG?1m)sVF5N_w#m#|gY~2^tf?Z7RO0UaB51rLI-= zTYpXKXW4_$w&7psj`a29@~5Pac#PxAA4(D7Ci*?U&f<(UKeJaO9&8DlPai zPaN+K7r6~i>B*Yzs=Q^egej=`K4}YUfxIsuZ!-c9aYKcR{YW*(KM#!S5KViLPSK;aUKV#C_`vJiePjc~KGEA(?RnS?eWFci$XySYfKeNf)SEd(RxI^NZkx;ow*`NfyN)}* z=)ZS=r{R-97q#Ux4YjEBsvhYWysk>W)G6Yjf$=r!+%)rDSAnKVs*A8Cvz08oO%CU7 zncBDF6)8+lO4(6JkJ}ZtIZ6Ib9nX}c4qZ@0erCnWxGLjJtj8KLc*W-6Z#dcEO1|+~ zE`9vhuYr8$M7i_@kezFb)>&<{Iv&bX+hK!)4&Ub1^ku>+YEw~p>WoUlg3}<&&6lUrFYS(O<@TW4OC*|hLLe+Mfo!p&hCu`JxRRU`@usb7-Q8s z+Vw&>)oj@ZZ9bM2PF`M}fSF@V*$|EpJ9JtZV;ltv$FEvrBIKH55}~Y|{5;B&h)Q3sP&ohlgu&6_NQ+g=dPCSJHFJcBxFcXY)}b*5WT(rI04unA6iZQ(yb z9})RTc5=Opo=M#-)wZN<3X^ULM|QSh#4NW28hdlcLGP$ zH_9PwN3E*z`&MtG$E&u}ED!kYE*iK*smTX6R`nF^%%*gp@Di`mSoe#PUoR&ygURH< zmbTGFZp2BF_K=w6LV0;=yQ#&~uW5mn!MzqKJ;UvjR55Mx6KTWUA4wI426H6Sd-jNA zXbh=TWp+F2r*m%&ybNrR`F-EWVsL8!CnE>VzJd9MegvdNUryh`wBj-iuin%O!%Z*= zjn7t0W#7MU9Vi7m7wXrlts2K>YgeDT(DWDn!dozyoYwrdVqZynjw0Eig zok@4>ii5YRisUxdr`S7Mqqqcr6W7K0o*0n7F;7kJp+(6?bB4B!GtA0_(}$bj=9U>K zX+>Y5MgvsfQy$x_+8ezS^VIP^%l(;@_}{TEZSM^v2s3v~#Pr50ogS@}(R(zMbcT4| zZe5KTvs%DSz50~N%GyTREaLK|6C|1bk+X!;=TNH5_;Pl0;FVfhS_1hCd;!UazuS6J z7M?7MH(2UngI2FQCe{iRiTnv;%E#VYRt%)JqO;8h+vW$Dp8ratHa)P2P%RSnDL1tr zuku~|ga2Gfo*>X$JwOq&?Nx@iKFx2R=tF#gpGxc3dkT&2fBN}6#4JT#MN4wYd!RsZ*?(0tk zp5pw_Kw6wwt{AQSxeBKxn=Bk}{2$Un0z^!w;SP5Idxgu)<@ z08uOX@Lwh*$L4L87DP3A%PyBsO+~`%tLH$sftwWN`D5Q9jRe_A~xZ7+fr#vyp_P zfA|Fdd|cmuNEQElgZ)H{<)1GBdt^uI_W!xwj2Dm(KfVVhPCZ>+1VE)GBt`A9pQd~qAPIP?0?mJkR$F~ zyoHtLNkxKxdb9MCM$spyq6CtAkJhcL9j3FLcOy;~0`e?T z+1+o3-=*sRbk7sxkd%zd!y3_)nuASJ~TqDH)kA;IV{+Bw=QN zmK_urh)+Q=lC541GSRGVU;t=5RI&n5HruF!IXLl}6(uCxVPj+C<>jTN%>a)t@JtYj z&*iorg)$k?Z9w?e!NCDOB*;XA+A(eiusd6tnoQZr;1@kSJU~F#4d*8IrAy$IM1Wg7 z7!?zfEtR^vK7(%)4A^Q~nkP{CK)f;ICa1-HHnuJB&w{6f&~&jkdz6@)6ZJ-F|%go)4fX5&OIz{?Lq`UpTF$N?)I83GM&0F=egx(%DuUB%&E|EQL2yv6i? z+@v>wvqS!jU!ksnLTNO3&wvWKkW9S3RQws~}G4DZo#N()@$bYVXVhCnlL3(Bfbxm^GKWRHc*uA#J`BR$eYv8D-8@vJh9G%EH0v90I#;v|qDLwen*!;S3hr}o zC{p=?2g__%D(GIP%V#4x%C}I!+ykj=U-xZ!5~yJg0_N9CJ>ByRMA=27=ey zrRIU#2dLoz=xrla2XLL!c|dpS)b=IK%^MhAw1wV8?Zliq>@xq(t1;#-^5GheWuSR( zySzp*+OBGPeSKX#2sIH>_IXGcu8<`eD)*OL8zT-cpeCC4s~1W>QlNb~xxfQv-&&1N zk)ME#n)AV&x)(BLr0mX6BiQGX_Y|?S+l306^@XJuVGAHmy*dT{-ZgJdqu>>+7X~aZ zn4dZICPql<%je%lVUE(m78bYveQPFmM0^Rfa4jba#lwhryOV+6cFqV?#A`yE`)RUD z;<^pNWr;0J93=&t)#L(DKJKHy%t|kP7gmzs9R`NBskNP{YNU8`ReE!1KNi-Tpk}>0 zAVD2q+g7S4DEJlZk^g;7c$BH~7pZ#mx7WGuB6@l=R1IX@b~1{p^Yzl8I1R8LV1NMO z>svO@Sim#j!XnqRY*q*JBCw@`*~8Vpuk&~Y@jXN6`Wb_Uthg_f{@vT)`QPMQ00slz z!2$_cS&IZ3LHW(!+R2yv3L54VBVr$vqnfhN=#@8QZ5)U z$uBCB<9JT*?dwZVOZ&(&BlCckQ0&#KpWsr(0k)OqjJF`Ag-X1%crC8|4Js>dy<(yY z291UVAhG}5RRB$f~X_NQv2^OSg-@Bc$FR}Kr$S@c8Y5W)t{I#2W_llkKMj*=u+5xY9<{kY_ z#*Mhc-m>A=*ECPIGShbWTO zh<}ye~6eOY;>i&K%+ku^)voN=*r(K2=KkQ>) zp|OK_wsrNV1v8eTg0_F>)MK4Eh+tdaLd~rvR6;@mO!SES)hg|xex*fM+If6~t#;O| z^voPM;Q5%D_TZu(G9lCXt%XC}t zARA6-n9Gj2GleY5qwFeqV(A@zH%&u1Xw~|=%!b^{yYyRr2B^+iS5;fKt^Qg)uv2Zv zu_KbdQpruv**zkYZ`+$EORB!ZsuSEh_f(R+BYv91@7?hIDt24a*6&>iRp1kt`f0jiIg)nT>-iQ4sIzh6numsuRLtxF=7wND_4n}=huxXv<_Fc_0~*+L2L(h2D|Rwe z=g`nT@bwdA6AZywPdmR9rY9U%8`gvH12B<#7RhcpTxO}WyA0Z>28qEJNB3TDCJoqfayqW*Y(PcV@Lk^f}KE>88n!D6g&-dfeR?gA4;^4VAn+y=dS>BSXA^K(RxD0)4lqr8LPeCbPkC z*ggBt7>g5qYY0UwA;4-GtTDi}*^o6tszaAKxupK(15mbzN4)Q#wXGTgEC77L!x@-X z0}Gq5memhtsU`Vxs|Z z99VGd!w@DWB_SpWEm!H0gSQ%}B8WLhh0O{S1@rUsCCR|;>KgOVgNo ziC>_8YHfjg-N)H=h(AKA>}tE-)uhiNM{s+;;uj&!eKsng_QLN@8e7(`-eqv`oD(mCg)vD43H7v3i<_j zB{Z9KjE#-Sg&Gu&@}`F|>D;BH^xaMx;f?ny=)GLuZTMkz}n z7YYT|GooT*_JF0IWos*g5DP|%=WtsCx2d(YNP!mdtpOArDp7bg7yw2a>+64nv(PXg zmn0RCAfnJEZ=DD{p3? zqZ%JqOs3hxpE0WEK1x`0^Bf`CGI8Z0cO74x;GG`?l&J1PEuEnvN@_ul607y4v(40C zhPM5dRQKOlV!N~yLQn1d9k4YrU%(?EaL&9?0iJ{R%2}D1%r|Dc;j{n~rpq`u zk`fZKMX^`bfyfqP^AlY2c=b9H5HcWNZ-1tl3pU$fIAl z3zJRO%=qvjFa&ZNDUhHTyfxzDL`{PwXsU0yk}2RbVNVCHe$EwBY+ zZj-E(ZjzX`bSXXNv?6(4*;E$PuZ=PR8B}90t9iL{(Amu01cRh96}qjsT`EL2X{CuP-pS|xx9XP7O>PM^6WCQ&LCmgW5-XaBBvM{` zwGof?n8Uk{ufVG&<--u7B}xGTeg8T*Eju09;X)#2$ifcP8%Q!1X26;pu1z;9NKXfz zX6_uGmFlJE?$-Blm+2zYGdBrdNO8Zhek(R9<$`^%)O-Y|636}GLet>y!*RXDm@=Hh z$@1c0Pk8l*)TMt>ty}T_C+y@EMlb#4kw3Lzchhb<7O6Z`RgI`azej9d-K0D2ZO*bM;H}5%#ql&7 zOD;goW!D=Y*RNMq%Jg$ft)LbM>}moRHc66YnDx=s1usyYzYMX5u41y&iUb6H(Q&(6 zeI-iS5Q(RG1>Zlevl1nxQl&1y_wIJf&y~Z|%~oMf6n@GAc;maCDd47-S1uR76;s$9 zTrxz53h=io-K|zEnEsZq9`vHa|LN2;$G_{eRU6(d}*#H{6d2pT_? zt4BV}EubDS5jp6#v9hAejBOY&5JE| zJs083v_hQSHn&m|MKAnve*Syrk?&vm^@3X^3j@jeq|^|pEBRDq#eV7O`J(r~6?JN? zxAe9aP(T9)7qD?~mZpUUdY}jboKW zoqJj;Ur(c-G~~S0P*9wBVsEbM)KMH^?YtfMNbui}9wiIFilC4X0FlW$%p_pz1;y2v zV|=+l?5W@(k+K>KK$`_Z1{<+8(lKa<)D|&_C;A*;O7u=g`@hQ7lD(L20M^)92Uf4r zS0q|kZe|=h->Q9aDa67Jqb8L*I-`$Gc8vnB(x`Fpszi5(zemq3pP0;Q&4|2*$cKd- z#+nH$TGhSck7FtB&&6zI-wuv*=BCHxaT`1BCZ{YL<*Ig48~9R-eg1dt)udW3A_@dG ze)V$loxU$d*LELKUf&a=-Yh-}Xh@npaAS*BU?aZWpu)_bC}rHX(={Sx>>*0bvMi7K zwOMjm%&IUvHaB&*%Q`iSYf_!zq%r?^OrCiiKk*JF<0E=*t&)yw15BnSw%xB>coHHK zlY8f=?hmLb*?x{M{}%-@d3)jttpYdxBq0P&2j8+VKUS^N{@|wnOC7AUSm+#bU@zFrU@@N{6;VD-MI3> z1=?cr5$T+sxeW_LjeqZ*6ZWx;Em?8-O9|oPWp0CwTX2{D!PSx9t7-EEa>>7cfjo+2 zP(UO$BRx4Pl;4$?#+N_Vnb@#XF$#T_E}hPDj8^{ngEv;k>MYN+>T5_j%g$T8mawu% zJ4SIZcIsZ4{V$MGfOUdes+K3fJesc8FDz6_&9}$}!)M>ce^0K*G?DRYW_Do8zZ|J`{|)mG1(CfI(0 zJYrCUN4UtNwo}R7daPCA;^OnuU`IoH-=EQGYRqBMggryfK;ZZi-Z~R01Y#Jfs!?CR zG=xr{1D_agq4=+_k6-P%Q{n2{|^R@7G^WKb)TNziZcoz6{Uqz=V-9Ty?XVkc0LCq z<7%WN?5_X0-qODpQ3NdjE2~Yg`6D4Eg>V)+7cm%&bMpnH7hq~~vNeKjDL{;#V=52{ zjv+^MMxwv+1=d6PR;itxo!|@pJtiAmr4lhSWNbz%&!0aR zwstgHUl?Sgr++389ua|alSf5LYIXHH12DO5A(=d$i~p{Dc;R*`*o;7HEnw;&F2J!( z8+7d3F4tZ>;79{fklyO=CIB>%2pvqigMI-SL#KEUZkRp9sIhc0OyEaI!Ru6UGyxqG z$Zvs|O14s=KA@%K10?^s^{3vniVk;NjpRu0)(2vp=0Jho+c`7HhQNPSRUj`vpR1rY zrt8x~joqy!1{CD1|6R`TJa1al@Ux|$w2?ptpqgAVpz*g+AfN%ZppZF%`?*oE*Efi0 zhK-7a;%A1-i=9oLJ$q(kWCUha;N~MHM#|%`?LjKc%EAIX0pxF>4+2!=T`pYp^z-Aj zw7-Z$VW_8f2^$;!OpaPD8&S0YPDyrJny7@t0#ptF`%6!ML?H<`0s^IX@7}>%z(A+N zChtc;Hwpf)KGn{}05|fSv&~1ngqDxdK=eAd`@kmIEDm!4 zzp!38C?sU^_iwO31l~t|y?|C*3&i3GX7maGH;7nxOvWBV=Lx)_i&sg0kG_@pX~NEG z*!LW&B11z%@D?CHjo!Zh-_bs`zs_lzKns(BV;kCA_DNN!>A{1$C4wzohC02Q({fhE z`2h4f@Nom1U{2ZLnsP>&PyRq+!LdT%P9na4fPx>g**f zEPPzbN6Gi;=p3hG-5>!x<0cWts0AgE!%)P}v8BZkT#}*h<|extjBoYc@X|g*rP>k* zvw&N?qch;R6JRMB1L^<}U_)oR7NpVH)lS<#-Oet33JcqT783zsVHyT=Gc)bhNU(l7 z(>UaObnx~_^50V!25S0R6g$wiZh+%gB`m2@^QoThZiyJ4VO?hG1dw2hFEYwgSDu86fM5K140i6dQGkcPsHfOSy}iypd6MnP(#nlTC5PGISP zIkWL7bQ3_B0vbIV(9PnsSULrNg1B0~)<6ittD4941YdF*6i*B!X^Y=mS_+|;BT2%` z*B3fk)WFaaQCHdC-Ugdjs<2G^rAjt{$Tb-u=c58^OQ@NUvKs0hoqe&Y4=C7SGccm$ zo5U>iuDXg@#hI2ic2*?0W1l1oQUov8wg`0=V7ja8L1Oj$o-o1IIh!&(5z?X!D z`u1&p;s?efk0>QsUwwHGa!a7Pne*07v8p+>2n=>ja<41*<*iRf2%}0JJu$ zA$djpvMLXL1a>%%_ryP4!^6un9dCd*8_ea6p?YvQ_~#Dr2IQpN@e$I=X0YZY;VcUs zM`?5l$ZQm${094xiY0PcOp1iLBPa}D_j2i$fEj??&>?K4*--My0!8qP=Pq6xh3+rZ zYAe?gcuS8S9d5lHJ%6#*x(K1ak0P~c5}&8-eGu7@ml@3xk4^+6GsIPUMP1%Dd= zaBR`49JFu_!6rBA232}WiVYk!0>44{0_jK#{sjU=;A7pQy%s0v74)?XQ3GmkZ-?L- z4ev)ZftH3w45~jWCC@!R96%{g5e|NEwjGQX22?VZwh|hJcoU_-hCCvI1|`Y85wx1+ zc9j%vO%4#jVFeX>27x3xsHxEG$sZnWRaZkp1#E=S$wfms4wa(~_+;+i{#5I5OaIi5 z^i>~b_wKbhA4&^AcUKl-eNCli`t>@(8-&!R8v3Zf?-|))6l0JEIzJB&Tz`N*fO9)F zF#*ka|Fi6HG7JkH#wEZjb^h-5b!UzX?!*Q4gEvZq>|X+sy3eP| zP-izz=oUl;>K)(2nl?9x&e&-Hw?OU-&>0g9x3?QejaBRy3Sk{1)qsUU3O^lyDiH-3 zC<9fe2-HLAA2UOM8~VLavk4Ek6#7yR#6Ej`la*JcM1g*k9*YhH17~p_z0B_y^1gvv zvNI0yKKD@|VdUk!CX8G9J3vf86S6qiSx|WvGg7#vP^>!}N?bmZ6yb_&V3RM2Lfl!?!ARvIo2xa5SZuKiyugd;1fb*$YLmsgHL0)axrT|RU z8ms#=y`Ao?a5H#3CA)EbG&!DhN6m0`>{I>QB=?g1nC2T~VO?8kxV_w_J9KSorgK>F z_rQ7llYLydR*~@h9t>`zM1g)$rZ!RucK}|Vpn&2taRI>8e&KYU*S~%SDz=~8uTVTgfVj+7RIvaXeN``u;li++>}AKzcVe8GBCm7}8{XcqK$`@4&^9GfMOT z#Ezunlrm-H#0l$A^1E77nXs`WPr}U(JA4V5Bga>i+)22L#&8*KQd&X=8lF8gUQcbl7WQqTTGe zbutKV>2a9D5i#_Xp`Tqu??mQlq^o=X7 zu*}u1`n3>Q6zJ`S7DVSrlvEEmDQYxa(enBMM~MGzq4%yIN>#ETf}}C=4tS^EHhf^q zh-(W9bvi#j#gbiYD7D~J9is#$OZO)TL#(R$;r^IG6##`CCsqyK#PGhIkTAe@lKAM& z{=UNzCmO`6$RF9{sAV(2a8%k^avA@Whc4Rak)TG)UxQsU7PoHQI{XgkY9}4HX9dUw42Q4CWC)dG!^5MrClcZAlZ!-{) zNM_N(;Hu5ocY`dI&6OiuzhL+C*aNS@4h3b|X$CuPBWVTi)|X}+5hR9R)v~#kg=%j- zd~)lBepKy(zEM`T33GR(Y1_RBYF51zp>(1c<3YenqY@*Vjy6t{f21zD|b#wdg4_M_xc7{k zk%Pk)I+>Bm1x!=G!N<{qMU_W~d+-OamF&e<{@J5+ZYS;4P`KRlOu;H*$S2exeTOPb zQVEJQl~o88h-s9TxBNY3P1Ec*zg(g$a3^<3_TexSwU;?R>h-w2>T8w(h%a2rfsa1Ku$9m%Tdb(wA^0>H4 zJ`Co$Uy!xxJ~;7wRTKP*=@sA<3cQj2b_!!P4tF*cv^8&c-Mk>brM^ zIM@YA41uzPTAfor;(Xvg3B|%ji-iGMMW-J>P0sa1L_~xF(0_vm-yil!XhE4r9Nr-x zajyV(ZBf>Dese9FXKKXH(8+UJ>a0RT48#S0B@%k&0$1+7pWWQ)$*e|txP!!jJDuH`pgBGF!L#NqIJ8D0fJ$opzcx+m3$Tyh#z9f#ZA4KW-?y%@)Ls;2PuW z_PDLp)YOCu2n-w_HSS=zcTbpOd5KoSKYSBnm%i-t%TO~PaK>$+dYyOhF2VHX;7kbQ(_9CxS zQBk1+qM}r#q&oFo>)iZ~e8xfWvB8IQ3yPAI%kxk=p+>6hI_-ZDA`@-CFXOJ>_J7}D z%&N6A+?zVmO`*HwGaq8RbEQy8TCT*gxA&-6XP9(DR#2(ZbS&69y=(MHI!!LQm~esm zhnixJ`G@KhoXdq2JBu}&M+6Is(x7vf9fVkw|_|H1G7`s_2Tw>Tdo!@4;yZ?ntb^rb{4#evb0k!%@Msm~_vB{p( z%OuIJzJio0)cr>c?`XA9==TYwDu+_{D{ykPH0*9J*aK(E;~P`kBTQG<9TZT$WLc|* zpaAEr0JY(+RB1#g5s+Y9$h181sMd_d^@;}pUg)ohmV5;zWzo1fC^bIk(-y@xgtY!p zl?BQymC#Iq?_!b6;qeHj%y|r-d@|M%l#kf{N`<_Pg zO+y>Uxip>GbYD*q1sQ9gQxrZu$inlF=TKkjUta3ukFd_LZF*JgArj+t@91;1bGr<6 zA#3VU_oYf9+o9G2<&ZKZk~Z|8WpO5H0k{b&nwq~clDyoI-a&xw2!|qc_%M_BfRY85 z^`guLBqKSj1@;o`^$w7zLQ6$YVKP*SiBn1P)uRsTQO`chDdyIav~le$eQGNF+dp0!IP-G9fPHrEga<0>Pc$ z8ssq`Zi5)l3qqqi_wF6d%2(ULK_Zu>veUvg>P^Bj3ydpBu}-=qkLjv6d=^xjdH$(e zgU|2iifNbiVIyz8qdZOKu4tYoUFoaR!E?ngy(Jp&o@@OnB4YCQV_yQ1TR+VjhnMI> zKVB<2XGUfcd$-^@-h=KN;-(BG@=bY^idvM4ouW$Un4N=%e19nxX)1T>C-$!4P~;<$ znmpwrjg0Y*OI!A%j%|2jfSRIY}UIu z1b+iR%?^}R08E5#nKGD$(A5niqOfEFN`mXTCw{@F#%L3#B%ZZUFrc)! zqVQIW6DA1e==7Nrj1x*4cLrxS7K$drTNj?8TUg$1-gVrwX{im$$i3^uFF@AH9N&aO zjf`LK*O;6pKeIr$S?^sR6p2aODR6D3{=_o(Ysn)?Hq>aKc!fdd!AZUMM(!_msk_C%902fu3M|Op4)#|;+=H~^#at7!28mCn)uww8O9ML+colK(1lO-~L$nGiFh>)x z+;*pLJMMf%WKIwPRN!6$jgZZJQXDt8-w+uq=ff%IEVasAxT9UdH23dsZ#f@rRa87W zh^=jR%qs?7D58FWl)_&49CmW^%}}2oZ{>oh#;J?J zX$2yCZ%|;Vf|-|)P`_XWQCnPM;sosXJ)tv?`e!TvRp`OjV1VNlrhX-)3LrTL8Vfi8 zWFW*!35f$Y=WRFmpa-BP19&*FYM~m%)33WbQx?4rLbBm4+V+rh!>PJ%NtazVn+CA( z2*5R-aTZoIu*85j;S6e(^COkYG2bO(iWb~>_hS$3z#be@qGdP?CLm?uJZi$j!$U-R zp|w4s?P&FQ$30b}*s-RG!qYWgG_+vH2-N>HQ{VO(bsUk!2;6zdQNH9N%zqm91%Zou z38UwFLx1Su=X@8Gr*3%1Yu?Q8Jsue|p^iu`Ubb*r9~L*(YTHp_?~~^h%G>r0^7sbC z3_z`<)z#HOMrCLy5%|W4NiR

PTT}1%AXQr=- zF)IkD0eAt7G=X4KTU)PW^3>Eb)qMa6s$r)y@{w1y?1t?FsE|o%`f6&?wJ&Xc`wI&s znCr8tK6};*CssdnAhSeR z;p9ziyq>=Kus?xyN6-g0x!{a7)*wsQnmGTD$qx3#hgCg`n1Bucj=mL_&W3RQC)LGs zSth{uyVUyfemRrIQpy5NBP+SC?X*sC@CV{Z7d)SN-rf#FCieEisI7J!V=QBOGrzD#DbJ`gGfqB zH{7#}qW6E_nfuP2xz8}e?z2yv=RBvr=X~OTrcZrIpma@?h}H4StJp7sw4F{CRR|5@ zC&Zz7>sIHI1yzt3iDCWEA|Js{2RQ%^DcTyfAX&Kf==9sAP!CGtHcH=t| z%BxJe0%kRRs=T%@X3E}*9Sl{I94`#`KrMF#MoOQ0`2(`U{{( zy)-<)2|W})S`+u2V$t)pw?9s8^^HqspNAQ$dR;uO1};1lfKYvn@H_-UGFA7$#bUu66C=*>dAKQB;&Rm{dNWc8$=zz0S5Vx& zS6-M7{q^xJ1^*h0G3X{*)G%j{wvkpCvAu-LyT!-uqa4~q%#RWP^zN`jiiY6KDIhBI z5shwcY67FZPa5}f|JbPsh{n?;P{>CXH3TtZWQvYNO@_7e>*GUt)gXKoD@Y|*xOtls zx;B+OfEB}Xx)ON(QR8n1`9FsBRGMZ~r2f3z6%p4zE2^stE#kdKIC}*tdY-?^`6zs? z$|#pf%R+i|mk!jjQp2^Et?Y9*&6|t-9yg8+!&lO&QK`8yG;;w`#~F2%$Y{S{CpSA$ zk#I?a64zdV;ire6FZ4sYp0iSAY(xaM``g*;?-KNd#oo2M&LE|D)m$e18Qb(FO(NFsOG=B>B)OFRSgl z6yP!uKiFbQe}Ts2;jP)N!otCu7YSV-9`&5XM>tfqk1fecOLLPOxH0jq7lxll2Km5r zPL{0c6byox?PFPcRq-C-*0|3>$Gc_QrhnbP^?k)j8<*2ZB;Z2R)5bY{^A2T}d`opM z%ESba`O~@Cf8c6fL`+T3?#9jq*ppO$GeBU?RQ}n#5?}per-1$8!|DHn(R2U&KO7+J z5nd;hPV6`Ac2h2w?f4V;_2txBanbp~g`zX*q93ao){A)Ge0d3Z+iYyV-uE7N6#iM9 z9~tg5`Lvj5!QoFw>R}gsOYxu*04cs$-_u{)>#a6aBa9g`QL(vxDsOexdXa< zW#@YR>?ffoVO#`0?3wi))TZza-n~bQr1`8(Aqy|)?I~}gBnJnfZZG-JO);V2RvhGI zOB*{-=wL;SQnhg*6H`imW)M3SWne*(>yWotS$)X+?{lsL@C1CceLfMLhqGFmL-a zySXc;Yc}sUS{E~r6EvtL6pQv-759gQb>d%wAJ-EF%99s;pX$P!)5uNt0-NQ%LJYb^ z(#}Ehu$)=eH~v&cNrr zUy*aWwvJj$DV*GHi>&UM!(AF1ZD~=qoojZTmy2k$JTsRbE4b7Itl0TsT1K6ny5ei0 z!96#n9gX6l<)#gvOBi98u|?o_lRZ>{0!79S3X?Gy6iiNMk9CL^W>p_6^Jw3jDrC@$ z_h@=u4*27HjlFRU$h7Nml=wSaQOlX4y3x_0-04mo1-Uc!{vUDWQA1)GmprliS6)z>S_*XSyW@9FHMtPOGB*0_2UrCQ>==Ma)D88f(8MpZ*$xuRqBkoF zrdNJ1wGA37aaNkzv_7T%WJvMcxC;aSc&ogDYy_O;Ixs6*#s&wLb9Z z-QJNhCu9%^eK_uwo;w$~(SEe7sN3#HMmQ;sYf<&d4Zgy0RT|*QshoDR#?C)Fl;F-O zH>rQ*IQIN#GIxfrm~VfvHXV3WSb0os*tZ*}PLFOzcpes0((LvciOb1;b#q)}L2h|= z3w}9F7H*hQs9a=m^xX5KF`_S3j>%(C9R7h_W&oO84xu5-X%pKMZ8#GjycKoVcO6s}zCJ~dN(uZ)G8 zA^nTpw%*ZlZ*DT+e{I*!K{$G)mXIilzsP*_h_?K*7JR2_)oJ;7wcwp`jHH7Kqe{lS2-8oKE`zD zN;YwbG|V`ljlw#=KL_fDCRNQMMItg{c7L2dVxrX(JVHjb0WB%U>vCvolxD?7AzjaT zBQ?yek{htYjRtXgrZnHj&X!I2Xlk~s!5o^yFMGkhjo{w4<1I$CT#P1P_7bx94Kczr z9=nnq8Fod~8qje;i(273R-02@2Uo1??S4&}>=HOBVC?zN`{f9}ve`N{0M z-@12`Z)OA(-H=dLS}8c8=s}yeQO|jy_?)!4d0V#l3OUu}ouwqEF}G@(b=BS|d!4WO z$Xw@oNB0A37i^4-uxWJub^`tw-!?*Xl(}igt`n=^0+UGCkMWjJ7 z0)@YZ&x)655J^Dx|I>4p_eFsAqj>FI>g7sKzG49uOBIvm{f}fn!3xKX1bzufc^0OO zdg&9x(!G%k$LE4iAd{xxr7sBYeniOKeV+|@ZuE2vu3!7v-SZ|EE5+X|bi0WvGSW_5 zPV;;Gvg~ntm~Z3f{U+KUR|cCnc~UV)}`a-;e_9fSNGDI{lLccQ`Y8ouBYpM!lgWnj_s2A3@JM`lE2AH zKErx**1cg`VQE3{ovO+^)%AML@HW=`rlfX|pt!eVdr=8@670>^oGp!!#aYr2kw8@d z5$2KI_3ij}5WL2ioH}0{v->ql+{1h$vWmB;yX@K%{Btm!wU3eIA|b67WFZ+8i^0bZ zlytZe09+&t{mgqJU!-X6lHtGSGV-^1mK#b%~_5t zXfY*`G#Fd1kFla9%qWkWmEW%A<8M};#o;Ephoq)rShUdRCQM1vuW?_n{wTFgX zv>~7gA`KZa8d|sK*l-bwYF=;@^F&wLw9T5)Ci$w3sZ5xOFw(}P5q^`#<;N`K#K)MU z`$k~eR?+^!6mM`<#Y!ru#LcKH3&N!UDb=+`;P*f|JcmvDHZ6^XtNU1oVl!N}hz;0@pUWVUp5d`rt5MQ`nyxKv`Mip`|!4--d!EE40`iu9#$OGjMVj5p&#uQXcMpYd$*c2r5R+H_J21D!4qbf8=cWXiBt5xdY4_zH7dH@K)_`i!HX>+ZtE z%d7bf4n-;Mx$(%T!)}ATg?l4Z%C*Fo(sVc`F8Md**)HaTk&WI4W`=Zcu;NBSTty!f zZV-BZjCmrxHQ1iO-S#PQz7K5cmz$>#|- zmMMC`89RVWs!K$;C52-&cdx3x($BGHD-rX|U}EC9u;JYjyOWesGHk}Sb;EHsqB+{d z;7L|klc{Zb;9PJ2wDrznzGK|NJ8K;e%y5%`dnS3hXp@3s|*WXqT{2iuON6WtWyi~GM6Smmee2-$-7I%&S^l!i?U}qd=WKgsIJ3hQ$ z#aUPSJBh#lRV&qCog48$ESh)6ec{Fh)M@QzgmM0AEhWA^kP$%CRJ)Qqy1A@cM>Lv^ zphKZnbtSaFCzIH`E@t_0st*rAX0w}UUa+~byjI=#0M|SqME$$oiM4fN+K|OXSDp*f z4Ms5z831d(yCIM3X0E&sWPIzrb=Rjkeyn8sM@2ZM@f@~xV#&2-{2MZ zQe{%@g-J_GPI7$Ti*48=ULo(*oz4F5s??8lit0%70Q-y8?;aiBucANVK7}nF8 z&5s>7oHs$qiIuUl%Xd1>7R4eM4ellsoX&9dLcrbUn9{2s$X?#RzwxWXZ~KGZMZ?cJ ze3T{U@$Exag^;b&4r?8Li9vtv_4yp055bd`|67**58>fSMF0Chp}#|d<2^5MK=gku zGpp}+9I5(HNrbGc2*nf*E^10l2r{C7@{-z7Po&zfzsP-D{yb?OL( zpbN#VZ;n?m=Xj>Y?kgxNf`Y~RY&fje=kdb(_<017~1Q0`mJ_@ zV#@N-+1$BzqoE;IPp-2v43N)vp4`w!@Hn>ciDe)Y>gtr6A{n32wdjWOS*)KuwV6kk zs1bESe&J|29%=Y+m$kkNH|3!qQSI4(&^jOZVg$}}S6W)qzh3|9#kYlw*7OuTrNR3A z0#uUT4+e)+6W9kWt!+b_jGxwhfILJM^HLCcjz34%zXQ`L(bqkdEaCGW7Qy#O8EUL} zq6ugX3hy$rJ>|$ODn?_FwDbJiuf{!~Z-`IQEP7m_@S(sG_u_;FaRZe28h+vjt$AO! z(k?yK9?p)aR^vO3`EqPU>8UsH{W=!{m->H-YY|;KBNl9#nvU%3lG|}M_6N?jM1%QE z-@cLeI^GDBQHXRNRB??2JbhW%(-&iv`aeUw#>Q-%o`l+HhZH^!V7OhQp-v{~bmv=C z%$gl*SX*vcC_Zgy^t0X7Qzkm@P7}ss{SV>Fi9ju_u{6-~^sgs>y0Sl082V0J6cPSx za`&iU7vH3`k`)dSg!x;-{4*6ryZR>s5slqVn-Z7f{AgKOmBXZO=;(vLo$s|ptx)$k z?|ZtFFN6qcKJ+ijC=@X_Tkg)Y$fNGMY|f)n*~)axELG`B>)3ZDXLOw7zU|1p*z zSMt#y!{)gD{S-#zC5#4SQc3!|H?JkMXeU{@)%dA#Yk@za*+|RD<5kTqRn>9D-eT*4 zxvg%A>k@5x@ASEt+Zjv z4Rs{nmI&FgyQ%P28FEPqnl4{x788E+qUG@9kms$dzW^FGUy6aDCJ^bidn)FoSCLJa zbywu(OSia_0`cqj^bNk^$B&Zku;t48^|?w-mOQu@k8dA&fXvm~0D9i0U zM++5aQ2bkYGLGAv+NU~?=&xV%mtZTovocXAmcyU7dqA3?*K=Ao*|vrD8os#;(9l4k zYlR(l10{Vk5)_$=Jf2D=H9w{4Ln4FK8Z5w)p4e*UCGdVm=E{QYwTS-^2*miB=qEI0IdN><3?ZVi<=L+=x^72$O&Ku zQY7OwPNzd`rcg* zxhnw;L1%eb?v}l65U~$wTjZvuIk{9W&V3z^JU&(-dBn=p6j8ozZl@4wWSeXG^*XXO zP21_g58`&Jt)1=^$Rn4S-2AGVTgiE>F|->SAx{bK#dRUgV#-&yruCH;E!3}f@DPQV zs&Ge#gSl^u9XGd=+u*yqp|QO*zG4{DKz8}PpZn~>fs(1EdQHC0Gk++t%Zr%4%WNlB!UY*?foGByED6rxjhI)>5L*QY7W zLQrse-)%l3s$^H!jf0*x^u1wA8((_b=2aVzTW>DMHnTZ`wxJIiAUV~%7TcWtSjgm| z)K__B!^7Cusp|Cd%tc!V-6tuhqrzEJo()O(EcI|sp|q?QU5IK!nr&$T+1ld&tbG;$0r z$_QK|;4>#62ktV?u$%6_8iXL!NhqA=dqAlmXlh^Zm z!Wkdxx>-acFzaaYUYjj!ET^|mYH%<)&NW$iaZ#034Sm}>Toba={T2!p-rPNy(~GxW z;2X5_w4CeBb2r&*8(Cb!m$l8y{h`q~!aIQ+={(678>XS1DAUg zJ&k|uVPZ*-heGF&P4tP++neUflnU$j9vsGS{eEqUr?e6sQX@H?z$ z9M;36({b(OH%`O<%`(Sv-LYY$s~yI~Cw`iStDjHuxgKS6%O^1L@tu3;yqkE2<8hrb z_QKUakF%pLIyH)EBiF><-Ebm)^S_VJl=a{UGG5*uKC~!<{by13KdJEQk~kiHb1$`2 z>NZFb_$3T-e z4gP<7`A=tt-Tyx<2t7SsM7s*R=-}OxL{8$H@#Bj5%MTRuzkS8PSy#D_J0~L65u;cQ z{U#ZH%@&{gN?93!;c91*Q0a)1(OQ7g@sT0T25O`U$af1oXI^P>BrymR+#t?M<_fUR z+>%Jv%EWaX*!1At$0s+y(eJ$8--t6(aG>CkcZfR6zpzS^eaZi^S`0I*ukcAR_9F7r z#Hu-Zg}`4#YC=?0Swu;#EQ6`h!SqZN*F(;fDaRz`X8+|lavtXgd5JkCU9w$^WHOhh zpZnIh{;ibC{0_1p(-?~RKjM(!A?bc69j*r&kmk@()5d15DU;}OF?&qneI2`kj^^vw z;PH5)3%`B5#Ds``3~_}&nh$#PGR7yPFI{Ny>z5lr_4!nU!j6~z9F0eJT@NT);Nj7! z8@%e6;t5xIGp5m3YExX5#Fylo!g#}93JTNt%g;0L#=SQV+c?J{%Z`@U(q~EK!;%Zf zxA^u{R3Y>uMgjJJN#_WHKd+Em;caM$g`5Kgih{Ny!Tu!o?;N>~E93f*p-<@}=E__;jA+%teJyW(3SZM)w!c}yc^t@+|9o{jD^HAS z{_aU8anS!R3nzQ<^fUe+!Q}sY8TbCZ`bU0;G>pD~=f+=&2!-SIQ2l-xMZUIp(|;Cx z@eoYM88dES@RFgI;Uo#S;y+#ZTcWVL|A#vCXMJ9+V$&j17LS9Lhld9NWZPUaFaPI~ zXcm)n2dv@Fp0!rNn|0ikVxjqj_icIG6wECIUcx)T1>mDeVAtn!vk@Te>#7ePJiuLM zHiS#ejwAf+ssDm8Ks*|3RP$P1y1<>?wL-Ue$TEMxK|}=a>2<%qlZMAneDi

Nqe4(TIC&{~8C5(_p&dB(Ja^4wx(HjDVUHP)iYdF{^azs(?iC{aYU#(K$KGQEOsH z`)k0jye}TW9zAYg>GHWGL8x4|XCGcChk%KUKbw~T@mNS68PLDn1%w$u@Kev$cJ{Nv z62?Xsfml9O-5fadDh*-xm7xI>S#vZ`hr!>)cDx_bZc<-`fJ6<{q#WDo(+!Vps-J)R z6Jo?}wfC9#d+^_te(>UJ>7Nt8CPx*m07qC!SGN)vY*AX5uhdllwBhv-z*Ei3%Ol=WQ&$HpxstrRPUC!% z>(^sF4|n*GS$5Ws9(iLXuv$6;7L@!}!&oJ_H}# zNuaF(h*dyJGa5j{09j9X4<6$HyoMx{85Lrk!21KhOaXHIQQ#+juJ%{c4oE+~`+Re- ze`5oxfrK05=0LG0_=~G8`y_IT&cbN1PhM^bgR_z@^T>a4(Me% zju8NL7!d)e#BI^sGj)##048k>XjC~^+LVo8xs}4_oq@R;wgg}w9N4UHI*HJRzI*o# zmXZk|Ym{T#KJ7DQ(=l)WY?FYf_F&DgS1>fogNuMaumjLp7P%Ubl!}T2fT4T*__0Q= zK|Pi@R#6LRZ@4*o{col*7MbSPL*|O^!#QFNxiO_lhwFzgZ_;*B{Zi-WpH4NdE2TPl zA9ZiX-|2E8Td|LUC~>T>xI~Hk_l5hgNx1#4|y`O|W!ZfWRWXe%(nVUhi-< z3Wai5-06O*=qwTsblTl!gvUB23`PJ&6KjJfmS7E7-dL3Q=r$JmRL2v2lS1weEPL^@gv~(EXv85YxmlhnCTD@k>^UkGlKM= z!vw{&!whNEsq{Epw&kV@bBbtp2iO@##S7+l*GTJ2Y5#o*9J0{}xpnS4 zN4}+je9XiOVEf6v=NPYh#3�_!!uh8$b45S)p%TMwL@o68mjdMlk_Q(#gQuyu!k zK-JLVAuvd{ep{Q!Ux}r`+5&?+`)Qz=38;a)1%`kKH!Ji7sLjZ(8iM zo=gLX(-NvZ4Gj&uySu<0<<+ZGL}u33%Y$8u4hak8SV$a>&yVF`OE%2dX#qlwUnAhY zmO1c6%i&xq6BH#?0tfzpg_*{(^Rg}N9ruo%v{IPvWJ_1a|BpZ~f zU1Eae#cR)0RJVFm>I(kJ%fM%OzfutZeLD%1rd%gl`Iq-Qg5?;W(qD51=w=8Oqdk}F z$wPq^6~(d;qr$^u2X25S7PwP$W=YJLBTE5kk;oZ9fw9#zd(p)+7-MH_0yuv6#Zgu( zd$k?etjN2RYhgEY(HUbrJ6dopXlaq#Cu>aRer_%rw`PEQwKlN#*rEy8Q4&|zsILNFaD&2 zkqgwCL~^+YHlL?<((gg30_SlkwnAt-Peq7s4g+F1aIv7L#kNp0t;jw>43@K3damBZ z9(b|q03l{TDH{WjT)_8s1XXH&5hzP$W-f>yPKaZf$cIkYcxKiU&Qf;1F^H#dk50L) zi9FeFw{`El!wdLIfeYo3RwuCzC+gqa*=~7~mgCM$3mQU67%aJRZVE65To#-`BTW=eOdx~DFF7P=|Ud7z3?;`RmVMEeCtzPsq66WTX7Ll!CJBdl)^O1!B4yE4> zFKeFjegNbfcfM@B0m91w{OT?~0{9ao5=2;68fM@j3KUiWaI(h{2=FOfDgKBW+<2#S zM@2`+T@t4+a$x4^j_&Px_-?*=!eg;HnIL}i`F7pr7NJ$J#6R_kyY?<98M_a|fI9Ee z@KQ&Hax%ma^jLuU>LLq+O#ri4ZuhdRti=L_7y!)%OeMd@%k1SJ+C3eBm~whrni>EL z15R$-N3N_-VB3qe+xzwnQ2Wm%wDX&HX8@O&d?!C}GXW?n z&GMt!l8rj-){9}Xw&z;sDl!fe%?4vabwrm|rM?*1^AhK+ovZk}7*ZFu9iQ5w&w)TR ztvhqDF*BPrg);1lJR^LYoDZ}S97Mla6(5aX@dU_fAA|*9`UAcYU>3Rp{RL2kjER5+ zhhnf)oCOa+)=$ldAC-u$4FKsmG#OT^I3+EOTHyJ2tk+zHOJVyA^pk$g5%8TgGB%!j zyFQNv^?FsB|9TViDlu$vL~BIg^VE8Jk=K#wgmmB1{x79=;`iW?-&ob{!o9vhPrxGLc$oYXKyf%{F;*pc(DrUYq551K;!qc4JqD%DL}D;8Q|VbFouj( z0rS;PBG?ryLBZZayb4n*E8rjr>bY@#eQ6}S)4cv-?GRyt#d5b?UrwYP^)RSe7zWp+ z-Jvd~Az?<|gA(NOlh>mSO?hk-)o$P#X;?Z|XW2G6EuV&%-Kmp~)r&*EChZP&1~ zu?VpAb=fNKD%t0A_o>Dy@K#>Iy>tnR->@sH1Nx`A`CN5i&K(&w?ABWyhaKmp&xd{8 zeYzm}?Hff9${41tFuPZCxR~A#9eq*G`dzVrpfQK}PmX|T;e17JR0`q#%0DZwP7bd6 z)UcvFyLtHe_J+!e$Q7K-mluy-Joi5teBY17Ni2DV7(J4J^}Ri(d`X*do-C?lSfeMy z%v*Or6YA)sR~-4Jd;N@mD!h5xuU45?==1Tq9QLla;&JPE)fbY?@>cp++Q1c7bvBVE zLGjTM2Apmlmf9nvrUX7Pm5FHAqw2QoD|R{`$Djo-I{FQe1H%Ke1lwq2@YV|+?6q*) zt>fp&D>zr%Q(0{i9Bz?rSs&2xG(WUT6t*5N7xhsVWl@|is3s8~tuIi}XP3G5%pqyM zY(|eS{a5U=?(?br4&r-~PiL|VQ%9KokwF*2Eu~do$V=W_`=Ak7n`a)C2;FNI%B1a*-LF5mdL|8GL{$l zqHp|3G2S;4TnHO*D?#POnflJf-W)~N_hC`JCFolNHfx!jZ<;QiQ-&C(V<>m-C6bW! zQuU`=rwCuGv>$yFr4T<-!onoDxJSl@jEpehA%_@x>R`tZI;khuC-1~XSNbUcAtolP z>DOwjW3!a>CBW)tL#m*446hJ-6Jyf$myo&0YH3LdTNQbzt5XglrT=DMtI!ded&Y5| zgOP6jynjy2`w%0g@C?)vnk`uQpX<;4zJ7-J#cjF{di2iR#et}x_PN~lNx441Xqas* zACb{oZo`9|6@1$n(#-=0&t;|PX)GQe*dp_{K0AG31yFHEPq@~0w z{awBc)^aC&N!LguH1z~cx|vH`+oVfuv^n z=rsa~?i$omrkk7TR_Mj9O1r~U(y*Ow(RNDi@v_S~pQ2k?|2xuGJsAyqIT;+UHdpH_ z@v>sf_Atyl#;^nrwcO7fwKPH91!pgX-aM^0qS?XIYQZy7^ouLKXEf@NOz!>Mm9s1q zB2QNv@ZEkh6hJBw+&Fl%7t!rZa~9(u8;t;W@kzgn1Lsnsc3OO1Wu%SfPO>x+nv7j3 zvG`izyO~HGagt^C<&HrTl4sBzx6+B2h>4UTtJ>}{TP#I#=U_4qfBdB86&excm^#Ug zTa}zAu!`fI|Nl-ld<;Gi8(AI4Q$L0G=%G}~Jrfl$VWDuB3g0D|Um!H}+O;t3{K92D zxPKMUvW&DU(t=^<8+_*O=J(9W(5VF_D*SF>UCVD%;oG{tNel*UDiLT?r9L-3EdS1U z+2$o2uZ{>{8Q8(>2FdX}?aA}v3Ecr?w?KRBM8Mz8>-GaPi%r6HK_2<8qoF{Q`UkVq zI(U*4^2+PX+u_8B3-jMknmz1UOfZUB;uo!9AlyNj&lFX56jDBq6q+#bOUm#w;2>Y8 zkFL?@AjfwpOEXLP`u8??c!!>@jsZir@XzX#|3^w-)D-G_dh#&6>AyUa$EjzYAqfla zpgOHmRM*_Ozfg-Dx?EY(5@jx1Y%*w*Xk$(d?+lG%9TWM^xbIplJrElL1`MjIY<{J8 zB2f#)IVW>Yp53L=16)ZKOHFCv1_8Ak_#GiBI(_ya5+Zyy>Md9$DSgt|4XbAFZW*mi@yIi$I*Uce u2S}L@6)o?LyGS#)D%v + Conceptual IKET timeline + A conceptual timeline showing producer and consumer warp lanes with IKET ranges, markers, timestamps, and payloads. + + + + + + + What IKET can visualize inside a kernel + Conceptual illustration, not a run-iket output format. Actual trace viewing uses Perfetto; script analysis uses JSON. + + + + time + + Producer warp + Consumer warp 0 + Consumer warp 1 + + + + + + + + + + + + tma_issue + + tma_issue + + tma_issue + + + wait + + mma_tile + + epilogue + + mma_tile + + epilogue + + + wait + + mma_tile + + epilogue + + mma_tile + + epilogue + + + mark("next_tile") + + + + mark("wait_done") + + + Range start/end events record timestamps at IKET call sites. + Payload example: k_tile = 3 + + + + Range: duration reconstructed from paired start/end events + + Marker: one timestamped point event + Payload: optional runtime value attached to an event + + The actual Perfetto trace may contain many kernels, CTAs, warps, and tracks; this diagram only shows the mental model for user-defined IKET instrumentation. + diff --git a/operators/README.md b/operators/README.md new file mode 100644 index 000000000..8bccf848c --- /dev/null +++ b/operators/README.md @@ -0,0 +1,149 @@ +CUTLASS Operator API +==================== + +> [!NOTE] +> CUTLASS Operator API is currently in beta. All interfaces, names, +> and paths are subject to change. + +CUTLASS Operator API is a Python interface for integrating kernels written in CUTLASS +Python DSLs (like [CuTe DSL](https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/overview.html)) into your code. + +While DSLs focus on kernel authoring, CUTLASS Operator API focuses on ease of +managing and integrating those kernels into libraries that use CUTLASS. + +It views kernels as end-to-end "Operators" that execute an operation (like a GEMM), and provides two things: + +1. a **kernel-agnostic interface** for finding operators that support an + operation/operands, inspecting their properties, and executing them — the same + way regardless of which kernel each operator wraps. +2. a **registry** of officially maintained CUTLASS kernels exposed through that interface. + + +Example +------- + +```python +import cutlass.operators as ops +import torch + +A, B, out = (torch.randn(128, 128, device="cuda", dtype=torch.float16) for _ in range(3)) + +# Arguments express an operation, and the operands to it +args = ops.GemmArguments(A, B, out, accumulator_type=torch.float32) + +# Find operators that support our provided arguments, and can run on SM100. +# Returns a list of ``Operator``s that wrap ready-to-compile CuTe DSL kernels. +operators = ops.get_operators(args, target_sm="100") + +# JIT compile and execute one of the returned operators using our arguments +operators[0].run(args) +``` + +Why use it? +=========== + +Any software that uses kernels requires finding kernels that do what you want, wiring up glue code to call them, +and maintaining the integration as kernels evolve. +Without an integration layer, these tasks are manual, error-prone, and repeated for every kernel. +CUTLASS Operator API eases each of these tasks — as kernels rapidly evolve, you don't have +to choose between integration churn and adoption inertia. + +**"Which kernel does what I want?"** + +Without Operator API, finding the right kernel means reading kernel source code and manually +deducing support — which dtypes, layouts, tile sizes, arch features, etc. each kernel supports. Operator API +provides a simple ``get_operators(args)``: express the operation you want to run, and get all the +operators that support it. Each operator also exposes uniform metadata describing +its constraints and design features for more advanced inspection, instead of requiring you to deduce +from source code. + +**"How do I get newer, faster, or fixed kernels?"** + +Adopting kernels directly carries a maintenance burden to mirror bug fixes/optimizations into local copies and +monitor release notes for new relevant kernels. With Operator API, you **integrate once, +not perpetually**. New kernels, fixes, and optimizations land in the registry on each release, and you get +them automatically without changing your integration code — just upgrade ``nvidia-cutlass-operators``. + +**"How do I call this kernel with my torch tensors?"** + +Different kernels have different usage conventions — direct usage requires kernel-specific +glue code to convert your framework tensors to the kernel's expected inputs, set performance +options, and run preparation steps. Operator API wraps every kernel with a consistent interface: pass PyTorch (or any +DLPack-compatible) tensors directly into ``GemmArguments`` and call ``operator.run(args)``. That lets you swap operators without touching your call site. + +It also supports: + +- **Custom epilogue fusions with ease** — pass a plain Python function; Operator API lowers it onto + CuTe DSL's Epilogue Fusion Configuration (EFC) framework and fuses it into supported kernels. +- **Bring-your-own-kernel** — register your own CuTe DSL kernels so you can call them through the same + interfaces as pre-bundled ones from CUTLASS, with no separate integration path for in-house kernels. +- **Negligible runtime overhead** on top of invoking the underlying kernel directly. + + +How to use it? +============== + +Installation +------------ + +To use with PyTorch, install the ``nvidia-cutlass-operators[torch]`` package: + +```bash +pip install nvidia-cutlass-operators[torch] +``` + +Alternatively, choose which dependencies to install: + +```bash +# Install only nvidia-cutlass-operators core +pip install nvidia-cutlass-operators + +# Install all dependencies to develop, run tests, etc. +pip install nvidia-cutlass-operators[dev] +``` + +Next steps +---------- + +- [Basic GEMM tutorial](examples/000_gemm.ipynb) — a guided tour of the core interface concepts to run a minimal GEMM. +- More guided notebook tutorials in [examples/](examples/). +- [Full API reference](https://docs.nvidia.com/cutlass/latest/media/docs/operators/tutorials/index.html) for ``cutlass.operators``. + + +Supported and Upcoming Features +=============================== + +CUTLASS Operator API will support a wide range of functionality, configurations, and optimizations, all robustly tested. + +**Current support:** + +- Kernel coverage + + - Dense GEMMs (F32, F16, BF16, INT8) for Blackwell, Hopper, Ampere + - Preferred and fallback cluster shapes + - Static and dynamic scheduling + - Block-scaled GEMMs (NVFP4, MXFP4, MXFP8, mixed input precision) for Blackwell + - Grouped GEMM (Contiguous offset) for Blackwell + - Low-latency TGV GEMM for Blackwell + +- Custom epilogue fusions (e.g. activations, elementwise ops, aux load/store) +- CUDA Graph support +- Native support for PyTorch and other DLPack tensors +- Bring-your-own-kernel + + +**Upcoming support:** + +- Additional GEMM kernel coverage: Sparsity, performance optimizations, grouped GEMM variants, and more +- Ahead-of-time compilation +- JAX Graph support +- nvMatmulHeuristics support + +Community & Feedback +==================== + +We welcome contributions and feedback from the developer community. You can: + +- File bug reports, feature requests, documentation gaps, and upvote roadmap items on the [CUTLASS GitHub Issues page](https://github.com/NVIDIA/cutlass/issues). +- Contribute examples, tutorials, or improvements via pull request to the [CUTLASS repository](https://github.com/NVIDIA/cutlass). +- Ask questions and share ideas in the [#cutlass channel](https://discord.com/channels/1019361803752456192/1150868614921064590) on [NVIDIA Developer Discord](https://discord.gg/nvidiadeveloper). diff --git a/operators/cutlass/operators/__init__.py b/operators/cutlass/operators/__init__.py new file mode 100644 index 000000000..2d4200397 --- /dev/null +++ b/operators/cutlass/operators/__init__.py @@ -0,0 +1,136 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from cutlass.operators import mma, workspace +from cutlass.operators.arch import ArchPortability, TargetSm +from cutlass.operators.arguments import ( + DenseTensor, + EpilogueArguments, + GemmArguments, + GroupedGemmArguments, + Operand, + PerformanceControls, + RuntimeArguments, + ScaledOperand, + ScaleMode, + ScaleSwizzleMode, +) +from cutlass.operators.artifact import CompiledArtifact +from cutlass.operators.base import Operator +from cutlass.operators.config import GlobalOptions +from cutlass.operators.manifest import Manifest +from cutlass.operators.metadata import OperatorMetadata +from cutlass.operators.providers import ( + CuTeDSLProvider, + Provider, + available_providers, + register_provider, +) +from cutlass.operators.status import Status +from cutlass.operators.typing import NumericLike, TensorLike + +if TYPE_CHECKING: + from collections.abc import Callable + + +__version__ = "0.1.0" + + +def get_operators( + args: RuntimeArguments | None = None, + metadata_filter: Callable[[OperatorMetadata], bool] | None = None, + target_sm: TargetSm | str | None = None, + providers: list[Provider] | None = None, +) -> list[Operator]: + """Return Operators that match the given arguments, metadata filter, and target. + + Args: + args (RuntimeArguments | None): Runtime arguments describing the + operator invocation (e.g. :class:`GemmArguments`). When ``None``, + no argument-based filtering is applied. + metadata_filter (Callable[[OperatorMetadata], bool] | None): An + optional Callable that takes OperatorMetadata as input and returns + a boolean indicating if it should be considered for inclusion in + results. The result is an intersection of operators filtered by the + callable and by other parameters passed to this method. + target_sm (TargetSm | str | None): Compute capability to target (e.g. + ``"100a"`` or a :class:`TargetSm` instance). Filters Operators that + cannot run on this target. + providers (list[Provider] | None): Optional list of Providers to + restrict discovery to (e.g. ``[ops.CuTeDSLProvider]``). + + Returns: + list[Operator]: Operators matching all filters. + """ + return Manifest.get_operators(args, metadata_filter, target_sm, providers) + + +__all__ = [ + # Central class exposing the Operator interface + "Operator", + # Runtime arguments + "RuntimeArguments", + "PerformanceControls", + "EpilogueArguments", + "GemmArguments", + "GroupedGemmArguments", + # Operands + "Operand", + "DenseTensor", + "ScaledOperand", + "ScaleMode", + "ScaleSwizzleMode", + # Operator Discovery + "get_operators", + "Manifest", + "OperatorMetadata", + # Misc. core types + "Status", + "CompiledArtifact", + # Arch + "TargetSm", + "ArchPortability", + # Typing markers + "TensorLike", + "NumericLike", + # Configuration + "GlobalOptions", + # Provider management + "Provider", # base class for all Providers + "CuTeDSLProvider", # CuTeDSLProvider holding registry of kernels written in CuTe DSL + "available_providers", # global list of available, registered Providers + "register_provider", # decorator to register a new Provider + # Submodules to access less commonly used public surfaces + "mma", + "workspace", +] diff --git a/operators/cutlass/operators/arch.py b/operators/cutlass/operators/arch.py new file mode 100644 index 000000000..51eb8a287 --- /dev/null +++ b/operators/cutlass/operators/arch.py @@ -0,0 +1,219 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +import enum +from dataclasses import dataclass +from enum import auto as enum_auto +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Iterable + + from cutlass.operators.metadata import DesignMetadata, OperandsMetadata + + +@dataclass(init=False) +class TargetSm: + """Target compute capability & portability to compile for.""" + + cc: int + portability: ArchPortability + + def __init__( + self, + sm_str: str | None = None, + *, + cc: int | None = None, + portability: ArchPortability | None = None, + ): + """Create a ``TargetSm``. + + Args: + sm_str (str | None): String like "100", "100a", or "100f". + cc (int | None): Compute capability as integer (e.g., 100). + portability (ArchPortability | None): Architecture portability. If not + provided, defaults to ArchConditional if cc >= 90, else Portable. + + Note: + Must use either sm_str or cc/portability, not both. + + Example: + >>> TargetSm("100a") + >>> TargetSm(cc=100) + >>> TargetSm(cc=100, portability=ArchPortability.ArchConditional) + + Raises: + ValueError: If both sm_str and cc/portability are provided, or neither. + ValueError: If portability is not valid for the given cc. + """ + if sm_str is not None: + if cc is not None or portability is not None: + raise ValueError("Cannot specify sm_str with cc/portability") + self.cc, self.portability = self._parse_str(sm_str) + elif cc is not None: + self.cc = cc + self.portability = portability or ( + ArchPortability.ArchConditional + if cc >= 90 + else ArchPortability.Portable + ) + else: + raise ValueError("Must specify either sm_str or cc/portability") + + if self.portability == ArchPortability.ArchConditional and self.cc < 90: + raise ValueError( + f"TargetSm {self.__str__()} is invalid. ArchConditional targets must be Hopper (sm90) or newer" + ) + if self.portability == ArchPortability.FamilyPortable and self.cc < 100: + raise ValueError( + f"TargetSm {self.__str__()} is invalid. FamilyPortable targets must be Blackwell (sm100) or newer" + ) + + @classmethod + def ensure(cls, value: TargetSm | str | None) -> TargetSm | None: + """Coerce a string to ``TargetSm``, or return as-is if already one.""" + if value is None or isinstance(value, cls): + return value + if isinstance(value, str): + return cls(value) + raise TypeError(f"Expected TargetSm or str. Got {value} of type {type(value)}") + + @staticmethod + def get_supported_targets( + design: DesignMetadata, operands: OperandsMetadata + ) -> list[TargetSm]: + """Get the supported targets for a given design and operands. + + Raises: + ValueError: If ``design`` does not have a ``mma_instruction_type`` field. + """ + # For now, just return the supported instruction. + # It can be extended to any other rules that are introduced. + metadata_to_check = [] + + if not hasattr(design, "mma_instruction_type"): + raise ValueError( + "Cannot determine supported targets with missing mma_instruction_type field in " + f"design metadata: {design}" + ) + mma_supported_targets = design.mma_instruction_type.supported_targets( + design, operands + ) + + metadata_to_check.append(mma_supported_targets) + + # check against tile scheduler supported targets + if hasattr(design, "tile_scheduler") and design.tile_scheduler: + tile_scheduler_supported_targets = design.tile_scheduler.supported_targets( + design, operands + ) + metadata_to_check.append(tile_scheduler_supported_targets) + return list(set.intersection(*map(set, metadata_to_check))) + + def is_portable_to(self, other: TargetSm | str) -> bool: + """Check if this target can compile/run on architecture described by `other` TargetSm.""" + other = TargetSm.ensure(other) + match self.portability: + case ArchPortability.Portable: + return self.cc <= other.cc + case ArchPortability.FamilyPortable: + return self.major == other.major and self.cc <= other.cc + case ArchPortability.ArchConditional: + return self == other + case _: + raise NotImplementedError + + def supports_operators_from( + self, operator_targets: Iterable[TargetSm | str] + ) -> bool: + """Check if this target can compile/run operators designed for any of the given targets.""" + return any( + TargetSm.ensure(target).is_portable_to(self) for target in operator_targets + ) + + @property + def major(self) -> int: + """Get the major version of the target SM.""" + return self.cc // 10 + + @property + def minor(self) -> int: + """Get the minor version of the target SM.""" + return self.cc % 10 + + @staticmethod + def _parse_str(s: str) -> tuple[int, ArchPortability]: + """Parse a compile target string (e.g., "100", "100a", "100f").""" + portability_map = { + "a": ArchPortability.ArchConditional, + "f": ArchPortability.FamilyPortable, + } + + suffix = s[-1] + try: + if suffix.isdigit(): + return (int(s), ArchPortability.Portable) + return (int(s[:-1]), portability_map[suffix]) + except (ValueError, KeyError): + raise ValueError( + f"Invalid TargetSm: '{s}'. Expected form: '100', '100a', '100f'" + ) from None + + def __str__(self) -> str: + suffix = ( + self.portability.name[0].lower() + if self.portability != ArchPortability.Portable + else "" + ) + return f"{self.cc}{suffix}" + + def __hash__(self) -> int: + return hash((self.cc, self.portability)) + + +class ArchPortability(enum.Enum): + """Portability of a compiled Operator to other architectures/compute capabilities. + + See: + https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/compute-capabilities.html#feature-set-compiler-targets + """ + + Portable = enum_auto() + """Portable to future architectures (e.g. sm_100 without "a" or "f").""" + + FamilyPortable = enum_auto() + """Portable to future architectures within the same family (e.g. sm_100f). + Only applicable to Blackwell and newer architectures (cc >= 100). + """ + + ArchConditional = enum_auto() + """Not portable to any other architecture (e.g. sm_100a). + Only applicable to Hopper and newer architectures (cc >= 90). + """ diff --git a/operators/cutlass/operators/arguments/__init__.py b/operators/cutlass/operators/arguments/__init__.py new file mode 100644 index 000000000..5cd310284 --- /dev/null +++ b/operators/cutlass/operators/arguments/__init__.py @@ -0,0 +1,60 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +from .base import ( + Operand, + PerformanceControls, + RuntimeArguments, +) +from .epilogue import EpilogueArguments +from .gemm import GemmArguments, GemmProblemSize +from .grouped_gemm import GroupedGemmArguments +from .operand import ( + DenseTensor, + ScaledOperand, + ScaleMode, + ScaleSwizzleMode, +) + +__all__ = [ + # Top-level arguments & constituents + "RuntimeArguments", + "PerformanceControls", + "EpilogueArguments", + # Various RuntimeArguments subtypes + "GemmArguments", + "GemmProblemSize", + "GroupedGemmArguments", + # Operands & related classes + "Operand", + "DenseTensor", + "ScaledOperand", + "ScaleMode", + "ScaleSwizzleMode", +] diff --git a/operators/cutlass/operators/arguments/base.py b/operators/cutlass/operators/arguments/base.py new file mode 100644 index 000000000..60a396c2d --- /dev/null +++ b/operators/cutlass/operators/arguments/base.py @@ -0,0 +1,134 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +import copy +from abc import ABC, abstractmethod +from dataclasses import dataclass, field, fields +from typing import Any, final, get_type_hints + +from cutlass.operators.typing import NumericLike, TensorLike +from cutlass.operators.utils.dtype import to_cutlass_type +from cutlass.operators.utils.tensor import TensorWrapper + + +@dataclass +class RuntimeArguments: + """Describes the operands and all other arguments passed to an Operator at runtime. + + It contains runtime operands (usually tensors) passed to the operation, as well as + any custom epilogue fusions, and runtime performance controls. + + This is an abstract base class, whose subclass describes the operation type itself + (e.g. GemmArguments, GroupedGemmArguments, etc.). All operators that implement the same + operation type accept the same RuntimeArguments subclass. + """ + + performance: PerformanceControls | None = field(default=None, kw_only=True) + """Optional runtime performance controls passed to the Operator""" + + def _validate(self): + """Checks that the arguments are valid. + + This is run before all fields have been converted to TensorWrapper and cutlass.Numeric. + """ + + def __post_init__(self): + _convert_to_internal_types(self) + + +@dataclass +class PerformanceControls: + """Optional runtime performance controls passed to the Operator. + + Some operators may support performance options that can be controlled at runtime. + This class is the general container for all such controls. + """ + + +class Operand(ABC): + """Base class for all operands to Operators, which encapsulates one or more TensorLike objects. + + In the most basic case, an Operand enacapsulates a single tensor. + + In more complex cases, an Operand may encapsulate multiple tensors that encapsulate a single logical operand. + For instance, a :class:`~cutlass.operators.ScaledOperand` encapsulates a quantized and scale tensor, that together + reconstruct the logical value of the operand. + """ + + @final + def copy(self) -> Operand: + """Returns a copy of the operand. Does not copy the underlying tensor.""" + return self.__copy__() + + @abstractmethod + def __copy__(self) -> Operand: + """Returns a copy of the operand. Does not copy the underlying tensor.""" + raise NotImplementedError + + def _convert_to_internal_types(self, metadata: dict[str, Any] = None): + _convert_to_internal_types(self, metadata=metadata) + + +def _convert_to_internal_types(caller, metadata: dict[str, Any] = None): + """Converts fields of the caller to internal types. + + Current fields that are converted: + * ``TensorLike`` -> ``TensorWrapper`` + * ``NumericLike`` -> ``cutlass.Numeric`` + * Classes that implement ``_convert_to_internal_types`` -> their internal types + + Args: + caller (Any): The caller object to convert the fields of + metadata (dict[str, Any] | None): Additional metadata to be used for conversion + """ + type_hints = get_type_hints(type(caller)) + for f in fields(caller): + hint = type_hints.get(f.name) + value = getattr(caller, f.name) + + global_metadata = {} if metadata is None else copy.deepcopy(metadata) + global_metadata.update(f.metadata) + + if isinstance(value, TensorWrapper): + # No conversion needed + setattr(caller, f.name, value) + elif hint is TensorLike: + # Find all fields that are annotated as TensorLike, + # and wrap them in TensorWrapper + setattr(caller, f.name, TensorWrapper(value, **global_metadata)) + elif hint is NumericLike: + # Find all fields that are annotated as NumericLike, + # and convert them to cutlass.Numeric + setattr(caller, f.name, to_cutlass_type(value)) + elif hasattr(value, "_convert_to_internal_types"): + # If the field is an instance of a class that implements + # _convert_to_internal_types, convert it to internal types + value._convert_to_internal_types(metadata=global_metadata) + setattr(caller, f.name, value) diff --git a/operators/cutlass/operators/arguments/epilogue.py b/operators/cutlass/operators/arguments/epilogue.py new file mode 100644 index 000000000..29a748fcd --- /dev/null +++ b/operators/cutlass/operators/arguments/epilogue.py @@ -0,0 +1,196 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from collections import OrderedDict +from typing import TYPE_CHECKING + +from cutlass.operators.fusion import EmptyTensor, trace, trace_in_out +from cutlass.operators.utils.layout import LayoutType +from cutlass.operators.utils.tensor import TensorWrapper, is_torch_tensor + +if TYPE_CHECKING: + from collections.abc import Callable + + import cutlass.cute as cute + from cutlass.cute.typing import Numeric + + +class EpilogueArguments: + """Describes a user-defined epilogue that is fused on top of the operation described by the primary RuntimeArguments. + + An epilogue fusion is a custom function that performs tensor-level transformations on the + result of a matrix multiplication. This transformation is fused into the kernel's epilogue, + which stores the final output. + + ``EpilogueArguments`` encapsulates the epilogue function describing the transformation + along with its arguments. + + To support flexible definition of epilogues, ``EpilogueArguments`` is + defined generically as taking in an ``epilogue_fn`` and additional ``kwargs``. + + Under the hood, the AST for ``epilogue_fn`` is parsed to determine the + operands and outputs of the epilogue. ``kwargs`` must contain Tensors or scalars + for all operands and outputs in the provided epilogue. + + **Structure of ``epilogue_fn``** + + The ``epilogue_fn`` is a function describing the custom transformation + on the accumulator tensor (intermediate result before the epilogue). + + The general structure of these functions is: + + .. code-block:: python + + def custom_epi_name(accum, *args) -> TensorType | tuple[TensorType, ...]: + '''Compute the epilogue. + + # Args: + accum (TensorType): Result of the primary operation (e.g. + ``A @ B`` for a GEMM) before the epilogue. + *args: Additional tensors or scalars used in the epilogue + (e.g. aux tensors). + + # Returns: + At least one tensor resulting from the epilogue computation. + ''' + # Do some compute + return D # and potentially other values + + ``epilogue_fn`` must be a Python callable (or its string representation) + that **must** satisfy the following constraints: + + * Takes a first positional argument named ``accum`` -- the result of the + operation just before the epilogue. For a GEMM, ``accum = A @ B``. + * Returns at least one tensor resulting from the epilogue. Currently the + return list must contain at least one output named ``D``. + * Each argument following ``accum`` is a tensor or scalar to be loaded. + * Each variable in the return statement is a tensor or scalar to be stored. + * Operations are represented in static single assignment (SSA) form. This means + that each variable can be assigned exactly once. + + The underlying implementation of the epilogue in + the kernel will determine how operands are loaded and stored. + + **Structure of ``kwargs``** + + ``kwargs`` must contain sample Tensors or scalars for all operands and outputs + in the provided epilogue. + For example, with an epilogue of: + + .. code-block:: python + + def my_epi(accum, alpha, C, beta): + F = (accum * alpha) + (C * beta) + D = relu(F) + return D, F + + A user would need to construct epilogue arguments as follows: + + .. code-block:: python + + epi_args = EpilogueArguments( + my_epi, + alpha=..., C=..., beta=..., D=..., F=... + ) + """ + + def __init__( + self, + epilogue_fn: Callable | str, + **kwargs, + ): + epilogue_inputs: list[str] = [] + epilogue_outputs: list[str] = [] + if epilogue_fn is not None: + # Parse the epilogue_fn AST to get the required input and output arguments + epilogue_inputs, epilogue_outputs = trace_in_out(epilogue_fn) + + # Get required input and output arguments from kwargs. + # Deduplicate because a name can appear as both an input parameter + # and an output (e.g. ``def epi(accum, D): ... return D``). + self.tensors = OrderedDict() + for kw in dict.fromkeys(epilogue_inputs + epilogue_outputs): + if kw not in kwargs: + raise ValueError( + f"Argument {kw} is not provided in the kwargs of the EpilogueArguments constructor" + ) + self.tensors[kw] = kwargs[kw] + del kwargs[kw] + + if len(kwargs) > 0: + raise ValueError( + f"Unexpected keyword arguments for epilogue: {kwargs.keys()}" + ) + + self.epilogue_fn = epilogue_fn + + @property + def parameters(self) -> list[cute.Tensor | Numeric]: + """Returns the list of input and output parameters of the epilogue.""" + return list(self.tensors.values()) + + @property + def parameter_names(self) -> list[str]: + """Returns the list of names of the input and output parameters of the epilogue.""" + return list(self.tensors.keys()) + + def to_tensor_wrappers(self, permute: list[int] | None = None): + """Converts the input and output parameters of the epilogue to TensorWrappers.""" + for k, v in self.tensors.items(): + if is_torch_tensor(v): + if permute is not None: + v = v.permute(permute) + + self.tensors[k] = TensorWrapper(v) + + def trace(self, accumulator_shape: tuple[int, ...], accumulator_type: Numeric): + """Traces the epilogue function and generates an internal representation of the epilogue. + + Args: + accumulator_shape (tuple[int, ...]): The shape of the accumulator tensor. For example, for a GEMM, this would be the shape of the output tensor. + accumulator_type (Numeric): The datatype of the accumulator tensor. + """ + accumulator = EmptyTensor( + element=accumulator_type, + shape=accumulator_shape, + layout_tag=LayoutType.RowMajor, + ) + tensors_for_tracing = {**self.tensors, "accum": accumulator} + + # Parse the AST of the epilogue_fn again, this time with the set of required + # tensors. This pass converts the epilogue into an internal representation and + # performs a limited set of correctness checks (e.g., shape matches) + # + # Since all current providers are not based on C++ EVT, we do not need to convert + # the DAG to a tree. If a provider that tightly matches the C++ EVT template structure, + # this will need to be revisited. + self.traced_epilogue = trace( + self.epilogue_fn, tensors_for_tracing, requires_conversion_to_tree=False + ) diff --git a/operators/cutlass/operators/arguments/gemm.py b/operators/cutlass/operators/arguments/gemm.py new file mode 100644 index 000000000..73881066a --- /dev/null +++ b/operators/cutlass/operators/arguments/gemm.py @@ -0,0 +1,212 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from dataclasses import dataclass +from typing import NamedTuple + +from cutlass.operators.arguments.base import Operand, RuntimeArguments +from cutlass.operators.arguments.epilogue import EpilogueArguments +from cutlass.operators.arguments.operand import _operand_or_dense +from cutlass.operators.typing import NumericLike, TensorLike + + +def _logical_contraction_extent(operand: Operand, k_axis: int) -> int: + """Logical K extent of ``operand`` along ``k_axis``. + + ``_validate`` runs before ``TensorWrapper`` conversion, so PyTorch-backed + operands report physical shapes; sub-byte dtypes (e.g. ``float4_e2m1fn_x2``) + pack their K dimension, so scale by the storage packing factor. Falls back to + the physical extent for non-PyTorch operands (cute tensors are already + logical) and unknown dtypes. + """ + extent = operand.shape[k_axis] + try: + import torch + except ImportError: + return extent + dtype = getattr(operand, "dtype", None) + if not isinstance(dtype, torch.dtype): + return extent + from cutlass.operators.utils.dtype import ( + cutlass_type_from_torch_type, + torch_storage_packing_factor, + ) + + try: + return extent * torch_storage_packing_factor( + cutlass_type_from_torch_type(dtype) + ) + except KeyError: + return extent + + +class GemmProblemSize(NamedTuple): + """Problem size for a GEMM operation. + + A GEMM with problem size (M, N, K, L) has operands with the following shapes (batch, rows, columns): + + * ``A``: (L, M, K) + * ``B``: (L, K, N) + * ``out``: (L, M, N) + """ + + M: int + """Number of rows in A and out""" + N: int + """Number of columns in B and out""" + K: int + """Number of columns in A and rows in B""" + L: int + """Number of batches of matrix multiplications""" + + +@dataclass +class GemmArguments(RuntimeArguments): + """Arguments for a Generalized Matrix Multiplication (GEMM) operation: ``out = A @ B``. + + The tensors must be all rank-3 or all rank-2. + * ``L``: Number of batches + * ``M``: Number of rows in A and out + * ``K``: Number of columns in A and rows in B + * ``N``: Number of columns in B and out + + For convenience, construction for ``A``, ``B``, and ``out`` that are operands + to a dense GEMM can be passed in without wrapping them in ``DenseTensor``. + + .. code-block:: python + + GemmArguments(A, B, out, accumulator_type) + # is equivalent to: + GemmArguments(DenseTensor(A), DenseTensor(B), DenseTensor(out), accumulator_type) + + Other operand types must explicitly wrap tensors in a ``Operand`` subclass. + For example, a scaled GEMM can be constructed as: + + .. code-block:: python + + GemmArguments( + ScaledOperand(A, ScaleATensor, scale_mode, scale_swizzle), + ScaledOperand(B, ScaleBTensor, scale_mode, scale_swizzle), + out, # No need to wrap in a `DenseTensor` + accumulator_type, + ) + """ + + A: Operand + """Input tensor A of shape (L, M, K) or (M, K)""" + + B: Operand + """Input tensor B of shape (L, K, N) or (K, N)""" + + out: Operand + """Output tensor C of shape (L, M, N) or (M, N)""" + + accumulator_type: NumericLike + """Data type of the accumulator""" + + epilogue: EpilogueArguments | None + """Optional custom epilogue fusion to be performed after the GEMM""" + + def __init__( + self, + A: TensorLike | Operand, + B: TensorLike | Operand, + out: TensorLike | Operand, + accumulator_type: NumericLike, + epilogue: EpilogueArguments | None = None, + ): + self.A = _operand_or_dense(A).copy() + self.B = _operand_or_dense(B).copy() + self.out = _operand_or_dense(out).copy() + + self.accumulator_type = accumulator_type + self.epilogue = epilogue + super().__init__() + + @property + def problem_size(self) -> GemmProblemSize: + """Problem size for a GEMM operation.""" + M = self.A.shape[-2] + N = self.B.shape[-1] + K = self.A.shape[-1] + L = self.A.shape[0] if len(self.A.shape) == 3 else 1 + + return GemmProblemSize(M=M, N=N, K=K, L=L) + + def _validate(self): + """Checks that the arguments are valid.""" + if len(self.A.shape) < 2 or len(self.A.shape) > 3: + raise ValueError( + f"A must be a tensor of rank 2 or 3 (L=1, M, K), got {self.A.shape}" + ) + if len(self.B.shape) < 2 or len(self.B.shape) > 3: + raise ValueError( + f"B must be a tensor of rank 2 or 3 (L=1, K, N), got {self.B.shape}" + ) + if len(self.out.shape) < 2 or len(self.out.shape) > 3: + raise ValueError( + f"out must be a tensor of rank 2 or 3 (L=1, M, N), got {self.out.shape}" + ) + a_k = _logical_contraction_extent(self.A, -1) + b_k = _logical_contraction_extent(self.B, -2) + if a_k != b_k: + raise ValueError( + f"A's K dimension ({a_k}) must be equal to B's K dimension ({b_k}). A shape (L, M, K): {self.A.shape}, B shape (L, K, N): {self.B.shape}" + ) + if self.out.shape[-2] != self.A.shape[-2]: + raise ValueError( + f"out's M dimension ({self.out.shape[-2]}) must be equal to A's M dimension ({self.A.shape[-2]}). A shape (L, M, K): {self.A.shape}, out shape (L, M, N): {self.out.shape}" + ) + if self.out.shape[-1] != self.B.shape[-1]: + raise ValueError( + f"out's N dimension ({self.out.shape[-1]}) must be equal to B's N dimension ({self.B.shape[-1]}). B shape (L, K, N): {self.B.shape}, out shape (L, M, N): {self.out.shape}" + ) + if self.A.shape[:-2] != self.B.shape[:-2]: + raise ValueError( + f"A & B must have the same rank and batch dimension (if any). A shape (L, M, K): {self.A.shape}, B shape (L, K, N): {self.B.shape}" + ) + if self.out.shape[:-2] != self.A.shape[:-2]: + raise ValueError( + f"out & A must have the same rank and batch dimension (if any). out shape (L, M, N): {self.out.shape}, A shape (L, M, K): {self.A.shape}" + ) + + def _convert_epilogue(self): + """Converts the epilogue to an internal representation using internal types.""" + if self.epilogue is not None: + L = self.A.shape[0] if len(self.A.shape) == 3 else 1 + M, N = self.A.shape[-2], self.B.shape[-1] + accum_shape = (L, M, N) + self.epilogue.trace(accum_shape, self.accumulator_type) + self.epilogue.to_tensor_wrappers() + + def __post_init__(self): + self._validate() + self._convert_epilogue() + super().__post_init__() diff --git a/operators/cutlass/operators/arguments/grouped_gemm.py b/operators/cutlass/operators/arguments/grouped_gemm.py new file mode 100644 index 000000000..7b9df1ce8 --- /dev/null +++ b/operators/cutlass/operators/arguments/grouped_gemm.py @@ -0,0 +1,107 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from dataclasses import dataclass, field + +from cutlass.operators.arguments.base import Operand, RuntimeArguments +from cutlass.operators.arguments.epilogue import EpilogueArguments +from cutlass.operators.arguments.operand import _operand_or_dense +from cutlass.operators.typing import NumericLike, TensorLike + + +@dataclass +class GroupedGemmArguments(RuntimeArguments): + """Arguments for a grouped GEMM operation. + + A grouped GEMM performs a series of independent GEMM operations, where each GEMM + can have different matrix dimensions. The most basic formulation of a grouped GEMM is one that + takes lists of tensors for each of ``A``, ``B``, and ``out``, and computes the following: + + .. code-block:: python + + for i in range(problems_in_group): + out[i] = A[i] @ B[i] + + While the abstract formulation uses lists of tensors, in practice, the tensors are often packed + together contiguously in memory in a single tensor for each of ``A``, ``B``, and ``out``. + In this case, an ``offsets`` tensor delineates the ending positions of each problem in the group + within those contiguous tensors. + + Currently-supported variants of a grouped GEMM are: + + Contiguous offset 2D-3D grouped GEMM: + * ``A`` is a tensor of shape (TotalM, K) or (1, TotalM, K) + * ``B`` is a tensor of shape (problems_in_group, K, N) + * ``out`` is a tensor of shape (TotalM, N) or (1, TotalM, N) + * ``offsets`` is a tensor delineating the ending positions of each problem in the group. + + .. code-block:: python + + start = 0 + for i in range(problems_in_group): + end = offsets[i] + out[start:end, :] = A[start:end, :] @ B[i, :, :] + start = end + """ + + A: Operand + """Input tensor A""" + + B: Operand + """Input tensor B""" + + out: Operand + """Output tensor""" + + accumulator_type: NumericLike + """Data type of the accumulator""" + + offsets: Operand = field(metadata={"alignment_bytes": 4}) + """Offsets tensor delineating the ending positions of each problem in the group""" + + epilogue: EpilogueArguments | None + """Optional custom epilogue fusion to be performed after the GEMM""" + + def __init__( + self, + A: TensorLike | Operand, + B: TensorLike | Operand, + out: TensorLike | Operand, + accumulator_type: NumericLike, + offsets: TensorLike | Operand, + epilogue: EpilogueArguments | None = None, + ): + self.A = _operand_or_dense(A).copy() + self.B = _operand_or_dense(B).copy() + self.out = _operand_or_dense(out).copy() + self.accumulator_type = accumulator_type + self.offsets = _operand_or_dense(offsets).copy() + self.epilogue = epilogue + super().__init__() diff --git a/operators/cutlass/operators/arguments/operand.py b/operators/cutlass/operators/arguments/operand.py new file mode 100644 index 000000000..a605dc052 --- /dev/null +++ b/operators/cutlass/operators/arguments/operand.py @@ -0,0 +1,396 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +import enum +from dataclasses import dataclass +from enum import auto as enum_auto +from typing import Any + +from cutlass.operators.arguments.base import Operand +from cutlass.operators.typing import TensorLike +from cutlass.operators.utils.common import ceil_div, round_up + + +@dataclass +class DenseTensor(Operand): + """An operand encapsulating a simple, single dense tensor.""" + + tensor: TensorLike + + def __copy__(self) -> DenseTensor: + return DenseTensor(self.tensor) + + def __getattr__(self, attr: str) -> Any: + if hasattr(self.tensor, attr): + return getattr(self.tensor, attr) + else: + raise AttributeError( + f"'{self.__class__.__name__}' object has no attribute '{attr}'" + ) + + +@dataclass +class ScaledOperand(Operand): + """An operand whose logical value is a quantized tensor multiplied by a tensor of scale factors. + + A ``ScaledOperand`` encapsulates physical :attr:`quantized` and :attr:`scale` + tensors that together represent the logical value ``scale * quantized``: + each entry of :attr:`scale` broadcasts and multiplies a contiguous block + of :attr:`quantized`, where the block shape is given by :attr:`mode`. + + This Operand is typically used to express narrow-precision formats (such as + OCP MXFP8 / MXFP4 and NVIDIA NVFP4), which store the data as a quantized + narrow-precision tensor multiplied by a separate tensor of scale factors that recover + dynamic range. + + The scale :attr:`mode` and :attr:`swizzle` describe the physical + layout of the scale tensor. + + The :attr:`mode` describes how scale factors are broadcast over the quantized tensor, + and therefore determines its size. It's usually dictated by the data format you + seek to represent -- MXFP8 / MXFP4 / NVFP4 require specific scale modes. + See :class:`ScaleMode` for more details. + + The :attr:`swizzle` describes the in-memory layout of the scale tensor. + It's usually dictated by the hardware architecture -- for example, Blackwell block-scaled MMAs + require a specific swizzle layout, and most kernels using them require the input to be + pre-arranged in this layout. + See :class:`ScaleSwizzleMode` for more details. + + Example: + + .. code-block:: python + + quantized_A = torch.randn(M, K, dtype=torch.float8_e4m3fn, device="cuda") + scale_A = torch.randn(M, K // 32, dtype=torch.float8_e8m0fnu, device="cuda") + A = ScaledOperand( + quantized_A, scale_A, ScaleMode.Blockwise1x32, ScaleSwizzleMode.Swizzle32x4x4, + ) + + See Also: + - :class:`ScaleMode`: granularity of the scaling, and the + format-to-mode mapping. + - :class:`ScaleSwizzleMode`: in-memory layout of the scale tensor. + + References: + - `NVIDIA cuBLAS documentation for narrow precision data types `_ + """ + + quantized: DenseTensor + """Narrow-precision tensor holding the quantized values. + + The dequantized logical value is ``scale * quantized``: each scale factor broadcasts over a block of + this tensor (block shape given by :attr:`mode`) and multiplies its values. + """ + + scale: DenseTensor + """Tensor of scale factors. + + Each entry of this tensor broadcasts over a block of :attr:`quantized` and scales it. + Its shape and layout are determined together by :attr:`mode` and :attr:`swizzle`. + + This tensor must be contiguous and have exactly + :py:meth:`ScaledOperand.numel_scale` ``(quantized_shape, mode, swizzle)`` + elements. It must already be arranged in the layout named by :attr:`swizzle`. + Its shape is otherwise unconstrained and not validated. + """ + + mode: ScaleMode | tuple[int, ...] + """Block shape over which each scale factor is broadcast. + + Typically, a block of quantized-tensor elements share the same scale factor value. + The logical value of the scaled operand is the result of broadcasting + the scale factor along the block shape described by ``mode``, + and then multiplying it element-wise with the quantized tensor. + + This may be a :class:`ScaleMode` enum for commonly used block shapes, or + a bare ``(L, M, K)`` tuple for custom block shapes. + For instance, ``(1, 1, 32)`` means that each scale factor covers 32 elements + along the K axis, and 1 element along the L and M axes. + """ + + swizzle: ScaleSwizzleMode + """In-memory layout of the scale factor tensor. + + The scale factor tensor is sometimes required to be stored in a specific swizzled layout, + often dictated by the hardware MMA's contract. + """ + + def __init__( + self, + quantized: TensorLike | DenseTensor, + scale: TensorLike | DenseTensor, + mode: ScaleMode | tuple[int, ...], + swizzle: ScaleSwizzleMode, + ): + self.quantized = _operand_or_dense(quantized).copy() + self.scale = _operand_or_dense(scale).copy() + self.mode = mode + self.swizzle = swizzle + + def __copy__(self) -> ScaledOperand: + # quantized and scale must be recursively copied because they are Operands + return ScaledOperand( + self.quantized.copy(), self.scale.copy(), self.mode, self.swizzle + ) + + def __getattr__(self, attr: str) -> Any: + if hasattr(self.quantized, attr): + return getattr(self.quantized, attr) + else: + raise AttributeError( + f"'{self.__class__.__name__}' object has no attribute '{attr}'" + ) + + @staticmethod + def numel_scale( + quantized_shape: tuple[int, ...], + mode: ScaleMode | tuple[int, ...], + swizzle: ScaleSwizzleMode, + ) -> int: + """Return the number of elements expected in :attr:`scale` for an operand of ``quantized_shape``. + + ``quantized_shape`` is ``(L, outer, K)`` for a 3D operand, where ``outer`` + is ``M`` for an A-side scale and ``N`` for a B-side scale. One scale factor + covers ``V`` elements along ``K``, so the scale tensor has logical shape + ``(L, outer, ceil_div(K, V))``. + + `scale.numel()` must equal the value returned by this method. + The shape of the scale tensor is otherwise unconstrained. + + Returned value depends on ``swizzle``: + + - :attr:`ScaleSwizzleMode.SwizzleNone`: + ``L * outer * ceil_div(K, V)`` + - :attr:`ScaleSwizzleMode.Swizzle32x4x4`: + ``L * round_up(outer, 128) * round_up(ceil_div(K, V), 4)`` + + where ``V = ScaleMode.numel(mode)`` is the number of quantized-tensor + elements covered by one scale factor. + + Args: + quantized_shape (tuple[int, ...]): Logical shape of the operand the + scale will be paired with, as ``(L, outer, K)`` or + ``(outer, K)``. + mode (ScaleMode | tuple[int, ...]): Block shape of the scaling. + swizzle (ScaleSwizzleMode): In-memory layout of the scale tensor. + + Returns: + int: Required number of scale-factor elements. + + Raises: + ValueError: If ``quantized_shape`` is not rank 2 or 3, or ``swizzle`` + is not a recognised :class:`ScaleSwizzleMode`. + """ + if len(quantized_shape) == 3: + L, outer, K = quantized_shape + elif len(quantized_shape) == 2: + L = 1 + outer, K = quantized_shape + else: + raise ValueError( + f"quantized_shape must be rank 2 or 3, got rank {len(quantized_shape)}: {quantized_shape}" + ) + + V = ScaleMode.numel(mode) + K_blocks = ceil_div(K, V) + + match swizzle: + case ScaleSwizzleMode.SwizzleNone: + return L * outer * K_blocks + case ScaleSwizzleMode.Swizzle32x4x4: + return L * round_up(outer, 128) * round_up(K_blocks, 4) + case _: + raise ValueError(f"Unrecognised ScaleSwizzleMode: {swizzle}") + + +class ScaleMode(enum.Enum): + """An enum over commonly used block scaling modes for a :class:`ScaledOperand`. + + Each member's value is a ``(batch, row, col)`` tuple giving the + shape of the block of quantized-tensor elements that share a scale factor. + Each scale factor is broadcast & multiplied over this block of quantized-tensor + elements. + + For block shapes that are not enumerated here, bare ``(batch, row, col)`` + tuples are also accepted as :attr:`ScaledOperand.mode` + + See Also: + - :class:`ScaledOperand`: the operand that consumes this enum. + - :class:`ScaleSwizzleMode`: in-memory layout of the scale tensor. + """ + + Blockwise1x16 = (1, 1, 16) + """Scale over ``(1, 1, 16)`` block. Used by NVIDIA NVFP4 (FP8 E4M3 scale dtype).""" + + Blockwise1x32 = (1, 1, 32) + """Scale over ``(1, 1, 32)`` block. Used by OCP MXFP8 / MXFP4 (E8M0 scale dtype).""" + + @staticmethod + def compare( + mode1: ScaleMode | tuple[int, ...], mode2: ScaleMode | tuple[int, ...] + ) -> bool: + """Return whether two scale modes describe the same scaling. + + Allows mixed comparisons between :class:`ScaleMode` enums and bare + tuples, and tolerates differing tuple lengths as long as the longer + tuple has leading ``1`` s for the extra positions (so that, e.g., + ``(1, 1, 16) == (1, 16)``). + + Args: + mode1 (ScaleMode | tuple[int, ...]): First mode. + mode2 (ScaleMode | tuple[int, ...]): Second mode. + + Returns: + bool: ``True`` if ``mode1`` and ``mode2`` describe the same scaling. + """ + if isinstance(mode1, ScaleMode) and isinstance(mode2, ScaleMode): + return mode1.value == mode2.value + + # One of the two modes is a tuple. Use tuple comparison and allow + # for different lengths as long as long as the longer tuple contains + # only 1s for the extra leading positions (i.e., allow (1,1,16) == (1, 16)) + tuple1 = mode1 if isinstance(mode1, tuple) else mode1.value + tuple2 = mode2 if isinstance(mode2, tuple) else mode2.value + + if len(tuple1) == len(tuple2): + return tuple1 == tuple2 + elif len(tuple1) < len(tuple2): + padding = (1,) * (len(tuple2) - len(tuple1)) + return padding + tuple1 == tuple2 + else: + padding = (1,) * (len(tuple1) - len(tuple2)) + return padding + tuple2 == tuple1 + + def __eq__(self, other: ScaleMode | tuple[int, ...]) -> bool: + return ScaleMode.compare(self, other) + + def __ne__(self, other: ScaleMode | tuple[int, ...]) -> bool: + return not ScaleMode.compare(self, other) + + def __getitem__(self, index: int) -> int: + return self.value[index] + + def __len__(self) -> int: + return len(self.value) + + @staticmethod + def numel(scale: ScaleMode | tuple[int, ...]) -> int: + """Return the number of quantized-tensor elements scaled by one scale factor. + + For a scale mode of block shape ``(L, M, K)``, this is the block volume + ``L * M * K`` -- the number of contiguous elements in a quantized tensor + that share (and are multiplied by) the same scale factor. + + Args: + scale (ScaleMode | tuple[int, ...]): Scale mode whose tuple entries + give the block shape. + + Returns: + int: Product of the entries of ``scale`` (e.g. + ``numel(Blockwise1x32) == 32``). + """ + res = 1 + for val in scale: + res *= val + return res + + +class ScaleSwizzleMode(enum.Enum): + """In-memory layout of the scale tensor in a :class:`ScaledOperand`. + + Hardware MMAs that consume block-scaled inputs typically require their + scale tensors to be stored in a specific non-trivial layout. + + ``ScaleSwizzleMode`` indicates that scale factors inside a :class:`ScaledOperand` + are already stored in one of those specific layouts. + + Operator API validates that the scale tensor is contiguous and has the + element count required by ``(mode, swizzle)``, but it does not inspect + or rearrange the values. The producer of the scale tensor (typically a + quantizer) is responsible for writing values in the named layout. + + See Also: + - :class:`ScaledOperand`: the operand that consumes this enum. + - :class:`ScaleMode`: granularity of the scaling. + - PyTorch :class:`torch.nn.functional.SwizzleType`, which this enum closely resembles. + """ + + SwizzleNone = enum_auto() + """Scale tensor stored in the natural ordering implied by :class:`ScaleMode`. + + For an ``(L, M, K)`` operand at mode ``(1, 1, V)``, this is a tensor of + ``L * M * (K // V)`` scale factors with one value per ``mode`` block. + """ + + Swizzle32x4x4 = enum_auto() + """1D block-scaling layout dictated by Blackwell ``tcgen05.mma`` block-scaled MMAs. + + This is used for MXFP8 / MXFP4 / NVFP4 GEMMs on Blackwell. + Under this layout, each tile contains 128x4 scale factors, and the 128 rows are + interleaved in groups of 32, matching the warp-group structure of the Blackwell tensor core. + + For a quantized tensor of shape ``(L, M, K)`` and a scale vector size of V, + the scale tensor is required to have ``L * round_up(M, 128) * round_up(ceil_div(K, V), 4)`` + elements. + + References: + - `NVIDIA cuDNN: The 128x4 Tiled Layout for Block Scaling Factors `_ + - `NVIDIA cuBLAS: 1D Block Scaling Factors Layout `_ + """ + + +def _operand_or_dense( + operand: TensorLike | Operand, +) -> Operand: + """If operand is already a ``Operand``, return it. Otherwise, wrap it in a ``DenseTensor``. + + This is used for convenience interfaces to avoid needing to wrap dense tensors in a ``DenseTensor`` + class. + + Args: + operand (TensorLike | Operand): The operand to convert + + Returns: + Operand: The operand wrapped in a ``DenseTensor`` if it is a ``TensorLike``, otherwise the operand itself + + Raises: + ValueError: If the operand is not a ``Operand`` or ``TensorLike`` + """ + if isinstance(operand, Operand): + return operand + elif isinstance(operand, TensorLike): + return DenseTensor(operand) + else: + raise ValueError( + f"Expected operand to be a Operand or TensorLike, got {type(operand)}" + ) + + diff --git a/operators/cutlass/operators/artifact.py b/operators/cutlass/operators/artifact.py new file mode 100644 index 000000000..da7ac09b5 --- /dev/null +++ b/operators/cutlass/operators/artifact.py @@ -0,0 +1,55 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from cutlass.operators.arch import TargetSm + from cutlass.operators.base import Operator + + +@dataclass +class CompiledArtifact: + """Lightweight wrapper over the result of compiling a DSL kernel. + + For instance, when using CuTe DSL, it's a wrapper over the result of + ``cute.compile()``. + In addition, it stores which Operator and TargetSm an artifact was compiled for. + """ + + compiled_obj: Any + """The result of compiling the Operator""" + + operator_obj: Operator + """The Operator object on which `.compile()` was called to generate this artifact""" + + compiled_for: TargetSm + """The target for which the Operator was compiled""" diff --git a/operators/cutlass/operators/base.py b/operators/cutlass/operators/base.py new file mode 100644 index 000000000..c23c0254b --- /dev/null +++ b/operators/cutlass/operators/base.py @@ -0,0 +1,318 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +from __future__ import annotations + +import inspect +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, ClassVar, final + +import cuda.bindings.driver as cuda + +from cutlass.operators.arch import TargetSm +from cutlass.operators.status import Status +from cutlass.operators.workspace import AllocationRequirement + +if TYPE_CHECKING: + from collections.abc import Callable + + from cutlass.operators.arguments import EpilogueArguments, RuntimeArguments + from cutlass.operators.artifact import CompiledArtifact + from cutlass.operators.metadata import OperatorMetadata + from cutlass.operators.workspace import Workspace + + +class Operator(ABC): + """Central class that provides a uniform interface over DSL kernels. + + An ``Operator`` is Operator API's central interface of instantiating, + finding, and executing DSL kernels. + A **kernel** is the raw DSL code that does the computation; an + **Operator** wraps a kernel and exposes a uniform interface to it. + Each Operator subclass fills in the kernel-agnostic interface with + kernel-specific implementations for: + + * correctly extracting arguments and compile/run the kernel + * generating instances of the Operator + * exposing the kernel's properties via ``OperatorMetadata`` + + Typically, DSL kernels are written as classes that are parameterized (e.g. by + operand data types, tensor layouts, tile shapes, etc.) into a family + of related kernel configurations, much like a "template". Instantiating & + compiling a kernel class with concrete parameters produces a single runnable kernel. + + Likewise, an ``Operator`` subclass represents a family of parametrized kernels, + wrapped in the uniform interface. Each *instance* of that subclass represents + one fully-specified kernel configuration. + The subclass is responsible for generating valid configurations of this family from the + design space. All features that distinguish one Operator instance from another are + exposed via a read-only :attr:`metadata`. + """ + + supported_args_type: ClassVar[type[RuntimeArguments]] + """The subtype of ``RuntimeArguments`` that this Operator supports. + + This is used by providers to filter operator classes before calling ``generate_operators()``. + Subclasses must override this, failing which an error is raised when the subclass is defined. + """ + + designed_for_min_cc: ClassVar[int] + """Minimum compute capability that this operator class is designed for. + + It is assumed that an Operator cannot run on an architecture with compute capability + older than this. + It indicates a broad class-level minimum, meant for rough filtering. + However, it does not imply that the Operator can run on every newer architecture -- + ``operator.metadata.supported_targets`` provides finegrained ground-truth information about which + architectures each Operator instance can run on. + """ + + def __init_subclass__(cls, **kwargs): + """Validate that the concrete subclass defines required class-level attributes. + + Raises: + TypeError: If the subclass does not define ``supported_args_type`` or + ``designed_for_min_cc``. + """ + super().__init_subclass__(**kwargs) + if inspect.isabstract(cls): + return + + required_attrs = ["supported_args_type", "designed_for_min_cc"] + for attr in required_attrs: + if attr not in cls.__dict__: + raise TypeError( + f"Operator subclass {cls.__name__} must define '{attr}'." + ) + + def __init__(self, metadata: OperatorMetadata): + self._metadata = metadata + + @property + def metadata(self) -> OperatorMetadata: + """The read-only metadata for the operator. + + All features that distinguish one Operator from another (supported operand attributes, design + choices, supported targets, optional epilogue, ...) are exposed via metadata, and + can be used for inspection and filtering. + """ + return self._metadata + + @final + def supports( + self, args: RuntimeArguments, target_sm: TargetSm | str | None = None + ) -> Status: + """Return whether the operator can be compiled or run with the provided arguments. + + Args: + args (RuntimeArguments): Arguments with which the operator is to be compiled or run. + target_sm (TargetSm | str | None): An optional target compute capability to check support for. + + Returns: + Status: Status indicating whether the operator can be compiled or run with the provided arguments. + """ + target_sm = TargetSm.ensure(target_sm) + # Metadata should capture most common checks + if not (status := self.metadata.supports(args, target_sm=target_sm)): + return status + # Additional checks can be implemented in the subclass + return self._supports(args, target_sm=target_sm) + + @final + def compile( + self, args: RuntimeArguments, target_sm: TargetSm | str | None = None + ) -> CompiledArtifact: + """Compile the operator for the given arguments and target architecture. + + This is implicitly called by :meth:`run()`, but can also be called explicitly to obtain a ``CompiledArtifact`` + that and passed to :meth:`run()` later. + + Args: + args (RuntimeArguments): Arguments to compile the operator with. The operator will be compiled to support the types and layouts of tensors in these args, according to any static information present in it (layouts, dtypes, etc.). + target_sm (TargetSm | str | None): Target compute capability to compile the operator for. + + Returns: + CompiledArtifact: The compiled artifact. + """ + target_sm = TargetSm.ensure(target_sm) + return self._compile(args, target_sm) + + @final + def run( + self, + args: RuntimeArguments, + compiled_artifact: CompiledArtifact | None = None, + stream=None, + workspace: Workspace | None = None, + assume_supported_args: bool = False, + ) -> None: + """Execute the operator end-to-end with provided arguments. + + Check that args are supported, compile if needed, create a default stream if + needed, and launch the underlying kernel. + + Args: + args (RuntimeArguments): Arguments with which to run the operator. + compiled_artifact (CompiledArtifact | None): Compiled artifact returned from the :meth:`compile()`. If ``None``, the operator will first be compiled. + stream (cuda.CUstream, torch.cuda.Stream, or other stream-like object, or None): Stream to execute the operator on. If ``None``, the default/null stream ``cuda.CUstream(0)`` is used. + workspace (Workspace | None): Allocation of workspace at least as large as the workspace size returned from the :meth:`get_workspace_size()`. If the operator does not require workspace, this can be ``None``. If a workspace of inappropriate size is provided, the behavior is undefined and the operator may crash. + assume_supported_args (bool): By default, ``operator.supports(args)`` is called to check if the arguments are supported. If ``True``, this check is skipped. + """ + if not assume_supported_args and not (supports := self.supports(args)): + raise ValueError( + f"Operator does not support the provided arguments: {supports.error}" + ) + + compiled_artifact = compiled_artifact or self.compile(args) + if not stream: + stream = cuda.CUstream(0) + return self._run(args, compiled_artifact, stream, workspace) + + @final + def get_workspace_size(self, args: RuntimeArguments) -> AllocationRequirement: + """Return the workspace requirement for the operator for the given arguments. + + Args: + args (RuntimeArguments): Arguments to the operator to get the workspace requirement for. + + Returns: + AllocationRequirement: Requirement describing the size & alignment of the workspace required by the operator. + """ + return self._get_workspace_size(args) + + @final + def initialize_workspace( + self, args: RuntimeArguments, workspace: Workspace + ) -> None: + """Initialize the workspace for the operator. + + Args: + args (RuntimeArguments): Arguments to initialize the workspace for. + workspace (Workspace): Workspace to initialize. + """ + return self._initialize_workspace(args, workspace) + + @classmethod + @final + def generate_operators( + cls, + metadata_filter: Callable[[OperatorMetadata], bool], + epilogue_args: EpilogueArguments | None = None, + target_sm: TargetSm | str | None = None, + args: RuntimeArguments | None = None, + ) -> list[Operator]: + """Return a list of all supported operator configurations, optionally constrained to given parameters. + + Args: + metadata_filter (Callable[[OperatorMetadata], bool]): A Callable + that takes OperatorMetadata as input and returns a boolean + indicating if it should be considered for inclusion in results. + The result is an intersection of operators filtered by the + callable and by other parameters passed to this method. + epilogue_args (EpilogueArguments | None): Optional arguments to + pass to the operator's epilogue. + target_sm (TargetSm | str | None): Optionally restrict generation + to operators compatible with a given TargetSm. + args (RuntimeArguments | None): Optional arguments to filter + operators by. If None, all operators are considered. + + Returns: + list[Operator]: List of all supported operator configurations. + """ + target_sm = TargetSm.ensure(target_sm) + return cls._generate_operators(metadata_filter, epilogue_args, target_sm, args) + + ### + ### Methods below are to be implemented by concrete classes + ### + + def _supports( + self, args: RuntimeArguments, target_sm: TargetSm | None = None + ) -> Status: + """Perform any additional support checks beyond what metadata captures. + + Subclasses may override this to add checks that cannot be expressed in metadata. + Ideally, this should be empty and all checks should be captured in the metadata. + By default, no such checks are performed and the method trivially returns + ``Status.success()``. + + Args: + args (RuntimeArguments): Arguments to check support for. + target_sm (TargetSm | None): An optional target compute capability to check support for. + + Returns: + Status: Status indicating success, or reason why the operator does not support the provided arguments. + """ + return Status.success() + + @abstractmethod + def _compile( + self, args: RuntimeArguments, target_sm: TargetSm | None = None + ) -> CompiledArtifact: + """Compile the operator.""" + raise NotImplementedError + + @abstractmethod + def _run( + self, + args: RuntimeArguments, + compiled_artifact: CompiledArtifact, + stream, + workspace: Workspace | None = None, + ) -> None: + """Launch the underlying kernel for an already-compiled, validated invocation. + + A minimal version of ``run()`` that assumes ``args`` are supported and that a + valid ``compiled_artifact`` and ``stream`` are provided. Intended to be overridden + by subclasses as a minimal wrapper around the underlying kernel launch. + + Args: + args (RuntimeArguments): Arguments to run the operator with. It is assumed that ``operator.supports(args)`` is True. + compiled_artifact (CompiledArtifact): Compiled artifact returned from the ``compile()`` method. + stream (cuda.CUstream, torch.cuda.Stream, or other stream-like object): Stream to execute the operator on. + workspace (Workspace | None): Optional allocation of workspace at least as large as the workspace size returned from the ``get_workspace_size()`` method. If the operator does not require workspace, this can be ``None``. If a workspace of inappropriate size is provided, the behavior is undefined and the operator may crash. + """ + raise NotImplementedError + + @classmethod + @abstractmethod + def _generate_operators( + cls, + metadata_filter: Callable[[OperatorMetadata], bool], + epilogue_args: EpilogueArguments | None = None, + target_sm: TargetSm | None = None, + args: RuntimeArguments | None = None, + ) -> list[Operator]: + raise NotImplementedError + + def _get_workspace_size(self, args: RuntimeArguments) -> AllocationRequirement: + return AllocationRequirement.empty() + + def _initialize_workspace( + self, args: RuntimeArguments, workspace: Workspace + ) -> None: + raise NotImplementedError diff --git a/operators/cutlass/operators/config.py b/operators/cutlass/operators/config.py new file mode 100644 index 000000000..aee5804c2 --- /dev/null +++ b/operators/cutlass/operators/config.py @@ -0,0 +1,210 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import os +import warnings +from importlib.util import find_spec + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +class GlobalOptions: + """Singleton class that configures global options for CUTLASS Operator API. + + This can be used to enable or disable certain features of the API. For example, + the user can enable TVM-FFI support to allow for the use of framework-level tensors + at run time. + + Each option may also be initialized from an environment variable read the + first time ``GlobalOptions()`` is instantiated. Programmatic changes via the + property setters always override the environment value afterwards. + + Supported environment variables: + + - ``CUTLASS_OPERATORS_USE_TVM_FFI`` --> ``GlobalOptions().use_tvm_ffi`` + + Example: + + .. code-block:: python + + GlobalOptions().use_tvm_ffi = True + """ + + _instance = None + + def __new__(cls): + """Create a new singleton instance of the GlobalOptions class only once.""" + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance._options = { + "use_tvm_ffi": cls._init_use_tvm_ffi(), + } + return cls._instance + + @property + def use_tvm_ffi(self) -> bool: + """Check if TVM FFI support is enabled. + + Default: ``True`` if ``tvm_ffi`` is installed. + + Can be overridden at process start by setting the + ``CUTLASS_OPERATORS_USE_TVM_FFI`` environment variable to a boolean value + (``1``/``0``, ``true``/``false``, ``yes``/``no``, ``on``/``off``). + + When enabled, conversions from DLPack compatible tensors to + cute.Tensor is via TVM FFI. Additionally, invoking the compiled Operator + happens via TVM FFI. Both can offer significant (3x-10x) speedups. + + Dependencies: + + * Required: ``tvm_ffi`` (pip install apache-tvm-ffi). + * Optional: ``torch_c_dlpack_ext`` (pip install torch-c-dlpack-ext). + """ + return self._options["use_tvm_ffi"] + + @use_tvm_ffi.setter + def use_tvm_ffi(self, value: bool) -> None: + if value and not find_spec("tvm_ffi"): + raise ImportError( + "Attempted to enable TVM FFI, but it is not installed. " + "Please install it via `pip install apache-tvm-ffi`." + ) + self._options["use_tvm_ffi"] = value + + def save(self) -> dict: + """Save the current options to a dictionary.""" + return self._options.copy() + + def restore(self, inp: dict) -> None: + """Restore previously saved options from a dictionary.""" + self._options = inp.copy() + + @staticmethod + def _init_use_tvm_ffi() -> bool: + """Resolve the initial value of ``use_tvm_ffi``. + + Resolution order: + 1. Environment variable ``CUTLASS_OPERATORS_USE_TVM_FFI`` (if set). + 2. Whether the ``tvm_ffi`` package is importable. + + If the environment variable explicitly enables TVM FFI but the package + is not installed, raise ``ImportError`` so the user notices the + misconfiguration rather than silently losing the requested speedup. + + Cross-check against ``CUTE_DSL_ENABLE_TVM_FFI`` variable: if + both env vars are set and resolve to opposite booleans, raise + ``ValueError`` to avoid confusing behavior. + """ + api_env_value = _parse_bool_env(_ENV_USE_TVM_FFI) + dsl_env_value = _parse_bool_env(_ENV_DSL_ENABLE_TVM_FFI) + + if ( + dsl_env_value != api_env_value + and dsl_env_value is not None + and api_env_value is not None + ): + raise ValueError( + f"Conflicting environment variables: {_ENV_USE_TVM_FFI}={api_env_value!r} " + f"({'enables' if api_env_value else 'disables'} TVM FFI) with " + f"{_ENV_DSL_ENABLE_TVM_FFI}={dsl_env_value!r} " + f"({'enables' if dsl_env_value else 'disables'} TVM FFI). " + f"Set both to consistent values, or unset one to avoid confusing behavior." + ) + + has_tvm_ffi = find_spec("tvm_ffi") is not None + + match api_env_value: + case True: + if not has_tvm_ffi: + raise ImportError( + f"{_ENV_USE_TVM_FFI} is set to enable TVM FFI, but it " + "is not installed. Please install it via " + "`pip install apache-tvm-ffi`." + ) + return True + case False: + return False + case None: + if not has_tvm_ffi: + warnings.warn( + "TVM FFI is not installed, and will be disabled. The " + "host overhead for running kernels may be 3-10x higher " + "without it. It is strongly recommended to install it " + "via `pip install apache-tvm-ffi`. To silence this " + f"warning, set {_ENV_USE_TVM_FFI}=0.", + stacklevel=3, # warn -> _init_use_tvm_ffi -> __new__ -> caller + ) + return has_tvm_ffi + + +# --------------------------------------------------------------------------- +# Implementation details +# --------------------------------------------------------------------------- + + +# Environment variable that overrides the default value of the +# ``use_tvm_ffi`` option at the time ``GlobalOptions`` is first instantiated. +_ENV_USE_TVM_FFI = "CUTLASS_OPERATORS_USE_TVM_FFI" + +# CuTe DSL-side environment variable that controls TVM FFI in the runtime +# (e.g. ``cute.runtime.from_dlpack``). We don't own this variable but we read +# it for cross-validation against ``_ENV_USE_TVM_FFI``. +_ENV_DSL_ENABLE_TVM_FFI = "CUTE_DSL_ENABLE_TVM_FFI" + +_TRUE_VALUES = frozenset({"1", "true", "yes", "on"}) +_FALSE_VALUES = frozenset({"0", "false", "no", "off"}) + + +def _parse_bool_env(name: str) -> bool | None: + """Parse a boolean-valued environment variable. + + Returns: + ``True`` if set to ``1``/``true``/``yes``/``on`` (case-insensitive). + ``False`` if set to ``0``/``false``/``no``/``off`` (case-insensitive). + ``None`` if the variable is not set. + + Raises: + ValueError: If the variable is set but the value is unrecognized. + We fail-fast rather than silently picking a default so that typos + (e.g. ``USE_TVM_FFI=ture``) cannot quietly degrade performance. + """ + raw = os.environ.get(name) + if raw is None: + return None + normalized = raw.strip().lower() + if normalized in _TRUE_VALUES: + return True + if normalized in _FALSE_VALUES: + return False + raise ValueError( + f"Invalid value for environment variable {name}={raw!r}. " + f"Expected a boolean: one of " + f"{sorted(_TRUE_VALUES) + sorted(_FALSE_VALUES)} (case-insensitive)." + ) diff --git a/operators/cutlass/operators/fusion/__init__.py b/operators/cutlass/operators/fusion/__init__.py new file mode 100644 index 000000000..2bc8e597f --- /dev/null +++ b/operators/cutlass/operators/fusion/__init__.py @@ -0,0 +1,126 @@ +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +import ast +import textwrap +from typing import Callable, Union + +from cutlass.operators.fusion.ir.tensor import Tensor as EmptyTensor +from cutlass.operators.fusion.frontend import PythonASTFrontend, PythonASTInOutProcessor + + +def trace_in_out(fn: Union[Callable, str]) -> tuple[list[str], list[str]]: + """ + Trace the input function and returns the names of inputs and outputs + + :param fn: The function to trace. + :type fn: Union[Callable, str] + :return: A tuple of lists of input and output names. + :rtype: tuple[list[str], list[str]] + """ + processor = PythonASTInOutProcessor() + return processor.trace(fn) + + +def trace(fn, example_tensors, **kwargs): + """ + Trace `fn(**example_tensors)` and generates epilogue visitor + + :param fn or str: Python callable or string of the epilogue function + :param example_tensors: example inputs for fn + :type example_tensors: dict + + .. hightlight:: python + .. code-block:: python + + # Define epilogue function as Python callable + def example_fn(accum, C, alpha, beta, gamma): + D = ((accum + C) * alpha - gamma) / beta + return D + + # Define the example tensors + example_inputs = { + "accum": torch.empty(size=(6, 512, 512), dtype=torch.float16, device="cuda"), + "C": torch.empty(size=(6, 512, 512), dtype=torch.float16, device="cuda"), + "alpha": 1.5, + "beta": 0.5, + "gamma": 2.5, + "D": torch.empty(size=(6, 512, 512), dtype=torch.float16, device="cuda") + } + + # Generate the epilogue functor + epilogue_visitor = ops.fusion.trace(example_fn, example_inputs) + """ + if callable(fn): + + class EpilogueFunctor(PythonASTFrontend): + def __init__(self, cc=None, **kwargs): + # Since we are currently only using the trace() method for generating an + # intermiediate representation (which is not CC specific), we can hardcode the cc to 100 + if not cc: + cc = 100 + super().__init__(cc, **kwargs) + + pass + + setattr(EpilogueFunctor, "__call__", staticmethod(fn)) + + epilogue_functor = EpilogueFunctor(**kwargs) + epilogue_functor.trace(example_tensors) + return epilogue_functor + elif isinstance(fn, str): + + class EpilogueFunctor(PythonASTFrontend): + def __init__(self, cc=None, **kwargs): + self.source = textwrap.dedent(fn) + # Since we are currently only using the trace() method for generating an + # intermiediate representation (which is not CC specific), we can hardcode the cc to 100 + if not cc: + cc = 100 + super().__init__(cc, **kwargs) + + def parse(self, example_inputs) -> None: + self.example_inputs = example_inputs + self.ast = ast.parse(self.source) + self.visit(self.ast) + + epilogue_functor = EpilogueFunctor(**kwargs) + epilogue_functor.trace(example_tensors) + return epilogue_functor + else: + raise NotImplementedError("Expect a callable Python function") + + +__all__ = [ + "EmptyTensor", + "trace_in_out", + "trace", +] diff --git a/operators/cutlass/operators/fusion/activation.py b/operators/cutlass/operators/fusion/activation.py new file mode 100644 index 000000000..64985aaf9 --- /dev/null +++ b/operators/cutlass/operators/fusion/activation.py @@ -0,0 +1,233 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Helpers for common activation functions +""" + +import ctypes + +from cutlass.operators.fusion.ir.c_types import dtype2ctype, to_ctype_value +from cutlass.operators.fusion.library import ActivationOp +from cutlass.operators.utils.common import is_torch_available, is_numpy_available +from cutlass.operators.utils.tensor import is_numpy_tensor, is_torch_tensor + + +if is_torch_available(): + import torch + import torch.nn.functional as F + + +if is_numpy_available(): + import numpy as np + + +class ActivationFunctor: + """ + Base class for frequently used activation functions + """ + + @staticmethod + def numpy(x): + raise NotImplementedError() + + @staticmethod + def epilogue_output_op(element_epilogue): + c_element_epilogue = dtype2ctype[element_epilogue] + + class _EpilogueOutputOpParams(ctypes.Structure): + _fields_ = [ + ("alpha", c_element_epilogue), + ("beta", c_element_epilogue), + ("alpha_ptr", ctypes.c_void_p), + ("beta_ptr", ctypes.c_void_p), + ] + + def __init__(self, alpha, beta, *args) -> None: + self.alpha = to_ctype_value(alpha, element_epilogue) + self.beta = to_ctype_value(beta, element_epilogue) + + return _EpilogueOutputOpParams + + +class ActivationMeta(type): + @classmethod + def __call__(cls, x, *args): + if is_numpy_tensor(x): + return cls.numpy(x, *args) + elif is_torch_tensor(x): + return cls.torch(x, *args) + else: + raise NotImplementedError("Unsupported tensor type") + + @classmethod + def numpy(cls, *args): + raise NotImplementedError( + f"Numpy reference for {cls.__name__[:-4]} is not implemented." + ) + + @classmethod + def torch(cls, *args): + raise NotImplementedError( + f"PyTorch reference for {cls.__name__[:-4]} is not implemented." + ) + + +############################################################################## +# identity operator +class identityMeta(ActivationMeta): + @classmethod + def numpy(cls, x): + return x + + @classmethod + def torch(cls, x): + return x + + +class identity(ActivationFunctor, metaclass=identityMeta): + binding_type = ActivationOp.Identity + + +############################################################################## +# ReLu operator +class reluMeta(ActivationMeta): + @classmethod + def numpy(cls, x): + return np.where(x > 0, x, 0) + + @classmethod + def torch(cls, x): + return F.relu(x) + + +class relu(ActivationFunctor, metaclass=reluMeta): + binding_type = ActivationOp.ReLU + + +############################################################################## +# Leaky ReLu operator +class leakyReLUMeta(ActivationMeta): + @classmethod + def numpy(cls, x, leaky_alpha): + return np.maximum(x, 0) + np.minimum(x, 0) * leaky_alpha + + @classmethod + def torch(cls, x, leaky_alpha): + return F.leaky_relu(x, leaky_alpha) + + +class leaky_relu(ActivationFunctor, metaclass=leakyReLUMeta): + binding_type = ActivationOp.LeakyReLU + + +############################################################################## +# Tanh operator +class tanhMeta(ActivationMeta): + @classmethod + def numpy(cls, x): + return np.tanh(x) + + @classmethod + def torch(cls, x): + return torch.tanh(x) + + +class tanh(ActivationFunctor, metaclass=tanhMeta): + binding_type = ActivationOp.Tanh + + +############################################################################## +# Sigmoid operator +class sigmoidMeta(ActivationMeta): + @classmethod + def numpy(cls, x): + return 1.0 / (1.0 + np.exp(-x)) + + @classmethod + def torch(cls, x): + return F.sigmoid(x) + + +class sigmoid(ActivationFunctor, metaclass=sigmoidMeta): + binding_type = ActivationOp.Sigmoid + + +############################################################################## +# SiLu operator +class siluMeta(ActivationMeta): + @classmethod + def numpy(cls, x): + return x * sigmoidMeta.numpy(x) + + @classmethod + def torch(cls, x): + return F.silu(x) + + +class silu(ActivationFunctor, metaclass=siluMeta): + binding_type = ActivationOp.SiLU + + +############################################################################## +# Hardswish operator +class hardswishMeta(ActivationMeta): + @classmethod + def numpy(cls, x): + relu6 = np.minimum(np.maximum(x + 3.0, 0), 6.0) + return x * relu6 / 6.0 + + @classmethod + def torch(cls, x): + return F.hardswish(x) + + +class hardswish(ActivationFunctor, metaclass=hardswishMeta): + binding_type = ActivationOp.HardSwish + + +############################################################################## +# GELU operator +class geluMeta(ActivationMeta): + @classmethod + def numpy(cls, x): + from scipy.special import erf + + return 0.5 * x * (1 + erf(x / np.sqrt(2.0))) + + @classmethod + def torch(cls, x): + return F.gelu(x) + + +class gelu(ActivationFunctor, metaclass=geluMeta): + binding_type = ActivationOp.Gelu diff --git a/operators/cutlass/operators/fusion/backend/__init__.py b/operators/cutlass/operators/fusion/backend/__init__.py new file mode 100644 index 000000000..f2f84b88f --- /dev/null +++ b/operators/cutlass/operators/fusion/backend/__init__.py @@ -0,0 +1,45 @@ +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + +from cutlass.operators.fusion.backend.sm80_emitter import Sm80Emitter +import cutlass.operators.fusion.backend.sm80_nodes as sm80_nodes +from cutlass.operators.fusion.backend.sm90_emitter import Sm90Emitter +import cutlass.operators.fusion.backend.sm90_nodes as sm90_nodes +from cutlass.operators.fusion.backend.sm100_emitter import Sm100Emitter +import cutlass.operators.fusion.backend.sm100_nodes as sm100_nodes + +__all__ = [ + "Sm80Emitter", + "sm80_nodes", + "Sm90Emitter", + "sm90_nodes", + "Sm100Emitter", + "sm100_nodes", +] diff --git a/operators/cutlass/operators/fusion/backend/emitter_base.py b/operators/cutlass/operators/fusion/backend/emitter_base.py new file mode 100644 index 000000000..96c4a4bee --- /dev/null +++ b/operators/cutlass/operators/fusion/backend/emitter_base.py @@ -0,0 +1,165 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Base class for Epilogue Visitor Emitter +""" + +from cutlass.operators.fusion.library import DataTypeTag +from cutlass.operators.fusion.ir import TopoVisitorNode, DAGIR + + +class FusionCallbacks: + def __init__(self, dag_ir: DAGIR, cc: int, emit_CD=True) -> None: + """ + Emit the EVT fusion callbacks + :param dag_ir: the DAG IR holding the epilogue visitor + :param cc: compute capability + :param emit_CD: whether to emit nodes C & D as a part of the fusion callbacks + For Sm90, set emit_CD=False, as Tensor C & D are hardcoded in the collective API + so that their shared memory can be explicitly reused + For Sm89, set emit_CD=True as they are treated as normal AuxLoad & AuxStore nodes. + """ + self.dag_ir = dag_ir + self.emit_CD = emit_CD + self.cc = cc + self.evt_cc = 90 if cc >= 90 else cc + if self.cc < 90: + self.namespace = "threadblock" + else: + self.namespace = "fusion" + + # + # Helper functions + # + + def get_visitor_name(self, node: str): + """ + Get the visitor name + """ + meta = self.dag_ir.get_node_meta(node) + if not isinstance(meta, TopoVisitorNode) and self.dag_ir.in_degree(node) > 0: + return f"EVT{meta.name_camel}" + else: + return meta.name_camel + + def emit(self): + node_metas = self.dag_ir.node_metas_topological_order() + epilogue_str = "" + # Step 1: emit individual node type decl + # emit the EVT & DAG connector + for meta in node_metas: + if not meta.disabled: + epilogue_str += self.emit_node(meta) + if not self.emit_CD and meta.name == "D": + continue + if isinstance(meta, TopoVisitorNode): + epilogue_str += self.emit_dag(meta) + else: + epilogue_str += self.emit_evt(meta) + + # Step 2: post-processing & get callback name + if not self.emit_CD: + if not self.dag_ir.has_node("C"): + epilogue_str += "using ElementC = void;\nusing StrideC = StrideD;\n" + output_node = self.dag_ir.get_all_inputs("D")[0] + # The callback is the src of node D + callback_name = self.get_visitor_name(output_node) + else: + # The callback is the last node in the topological order + callback_name = self.get_visitor_name(node_metas[-1].name) + return epilogue_str, callback_name + + def emit_evt(self, node): + if self.dag_ir.in_degree(node.name) == 0: + return "" + + evt_tmp = f""" +using EVT{node.name_camel} = cutlass::epilogue::{self.namespace}::Sm{self.evt_cc}EVT< + {node.name_camel}, +""" + sorted_children = self.dag_ir.get_all_inputs(node.name) + evt_node_strs = [ + f" {self.get_visitor_name(child_name)}" for child_name in sorted_children + ] + evt_tmp += ",\n".join(evt_node_strs) + ">;\n" + + return evt_tmp + + def emit_dag(self, node): + subgraph = node.subgraph + subgraph_nodes = subgraph.nodes_topological_order() + # Emit the Edge Tuple + edge_tuples = "cute::tuple<\n" + for n in subgraph_nodes[:-1]: + in_edges = subgraph.in_edges(n) + edge_weights = [ + subgraph.get_edge_weight(edge[0], edge[1]) for edge in in_edges + ] + sorted_children = [ + edge[0] for _, edge in sorted(zip(edge_weights, in_edges)) + ] + edge_tuple = " cute::seq<" + edge_str = [str(subgraph_nodes.index(child)) for child in sorted_children] + edge_tuple += ", ".join(edge_str) + ">,\n" + + edge_tuples += edge_tuple + edge_tuples += " >" + + # Emit the node list + dag_nodes = "" + dag_node_strs = [] + for n in subgraph_nodes[:-1]: + n_meta = subgraph.get_node_meta(n) + if n_meta.disabled: + dag_node_strs.append(f" {self.get_visitor_name(n)}") + else: + dag_node_strs.append(f" {n_meta.name_camel}") + dag_nodes = ",\n".join(dag_node_strs) + + return f""" +using {node.name_camel} = cutlass::epilogue::{self.namespace}::Sm{self.evt_cc}TopologicalVisitor< + {DataTypeTag[node.subgraph.element_compute]}, + {edge_tuples}, +{dag_nodes} +>; +""" + + def emit_node(self, node): + if isinstance(node, TopoVisitorNode): + emission = "" + for meta in node.subgraph.node_metas_topological_order(): + if not meta.disabled: + emission += self.emit_node(meta) + return emission + else: + return node.underlying_impl.type_decl diff --git a/operators/cutlass/operators/fusion/backend/sm100_emitter.py b/operators/cutlass/operators/fusion/backend/sm100_emitter.py new file mode 100644 index 000000000..879640eb3 --- /dev/null +++ b/operators/cutlass/operators/fusion/backend/sm100_emitter.py @@ -0,0 +1,151 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Emitter for Sm100 Epilogue Visitor +""" + +from cutlass.operators.fusion.library import ( + DataType, + DataTypeTag, + EpilogueScheduleTag, + KernelScheduleSuffixes, + OpcodeClassTag, +) +from cutlass.operators.fusion.backend.emitter_base import FusionCallbacks + + +def to_blackwell_threadblock_shape(tile_description, cluster_shape, kernel_schedule): + blackwell_threadblock_shape = tile_description.threadblock_shape + is_2sm = ( + False + if kernel_schedule is None + else ("2sm" in KernelScheduleSuffixes[kernel_schedule]) + ) + if cluster_shape[0] > 0: + blackwell_threadblock_shape = [ + tile_description.threadblock_shape[0] // cluster_shape[0], + tile_description.threadblock_shape[1] // cluster_shape[1], + tile_description.threadblock_shape[2] // cluster_shape[2], + ] + if is_2sm: + blackwell_threadblock_shape[0] *= 2 + else: + blackwell_threadblock_shape = ( + tile_description.math_instruction.instruction_shape + ) + return blackwell_threadblock_shape, is_2sm + + +class Sm100CollectiveEpilogue: + def __init__( + self, + tile_description, + kernel_schedule, + epilogue_schedule, + element_accumulator, + element_d, + fusion_callbacks, + ) -> None: + self.cta_tile_mnk, _ = to_blackwell_threadblock_shape( + tile_description, tile_description.cluster_shape, kernel_schedule + ) + self.element_accumulator = element_accumulator + if fusion_callbacks.dag_ir.has_node("C"): + self.element_c = fusion_callbacks.dag_ir.get_node_meta("C").element + else: + self.element_c = DataType.void + self.element_d = element_d + self.schedule = epilogue_schedule + self.fusion_callbacks = fusion_callbacks + self.opclass = tile_description.math_instruction.opcode_class + + @property + def CtaTileMNK(self) -> str: + """ + The threadblock shape + """ + return f"cute::Shape<_{self.cta_tile_mnk[0]}, _{self.cta_tile_mnk[1]}, _{self.cta_tile_mnk[2]}>" + + @property + def EpilogueTileType(self) -> str: + """ + The epilogue tile type + """ + return "cutlass::epilogue::collective::EpilogueTileAuto" + + @property + def Schedule(self) -> str: + return EpilogueScheduleTag[self.schedule] + + def emit(self): + stride_D_str = self.fusion_callbacks.dag_ir.get_node_meta( + "D" + ).underlying_impl.stride_mnl + stride_C_str = stride_D_str + if self.fusion_callbacks.dag_ir.has_node("C"): + stride_C_str = self.fusion_callbacks.dag_ir.get_node_meta( + "C" + ).underlying_impl.stride_mnl + + callback_decl, callback_name = self.fusion_callbacks.emit() + return ( + callback_name, + f""" +using EpilogueDescriptor = cutlass::epilogue::collective::detail::Sm100EpilogueDescriptor< + {OpcodeClassTag[self.opclass]}, + {self.CtaTileMNK}, {self.EpilogueTileType}, + {DataTypeTag[self.element_accumulator]}, {DataTypeTag[self.element_c]}, {DataTypeTag[self.element_d]}, + {self.Schedule}, {stride_C_str}, {stride_D_str}, + false /* IsPerColScaleSupported */, + false /* IsBlockScaleSupported */ +>; +{callback_decl} +""", + ) + + +class Sm100Emitter: + def __init__(self, operation, graph) -> None: + fusion_callbacks = FusionCallbacks(graph, cc=100, emit_CD=False) + + self.collective_epilogue = Sm100CollectiveEpilogue( + tile_description=operation.tile_description, + kernel_schedule=operation.tile_description.kernel_schedule, + epilogue_schedule=operation.tile_description.epilogue_schedule, + element_accumulator=operation.tile_description.math_instruction.element_accumulator, + element_d=fusion_callbacks.dag_ir.get_node_meta("D").element, + fusion_callbacks=fusion_callbacks, + ) + + def emit(self): + return self.collective_epilogue.emit() diff --git a/operators/cutlass/operators/fusion/backend/sm100_nodes.py b/operators/cutlass/operators/fusion/backend/sm100_nodes.py new file mode 100644 index 000000000..98c737216 --- /dev/null +++ b/operators/cutlass/operators/fusion/backend/sm100_nodes.py @@ -0,0 +1,140 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + + +from cutlass.operators.fusion.ir import AuxLoadImpl, AuxStoreImpl +from cutlass.operators.fusion.library import DataTypeSize, DataTypeTag, FloatRoundStyleTag +import cutlass.operators.fusion.backend.sm90_nodes as sm90_nodes +from cutlass.operators.fusion.pycute import product + + +Sm100AccumulatorImpl = sm90_nodes.Sm90AccumulatorImpl +Sm100LoadSrcImpl = sm90_nodes.Sm90LoadSrcImpl +Sm100ScalarBroadcastImpl = sm90_nodes.Sm90ScalarBroadcastImpl +Sm100RowBroadcastImpl = sm90_nodes.Sm90RowBroadcastImpl +Sm100ColumnBroadcastImpl = sm90_nodes.Sm90ColumnBroadcastImpl +Sm100ComputeImpl = sm90_nodes.Sm90ComputeImpl +Sm100StoreDImpl = sm90_nodes.Sm90StoreDImpl +Sm100ColumnReductionImpl = sm90_nodes.Sm90ColumnReductionImpl +Sm100RowReductionImpl = sm90_nodes.Sm90RowReductionImpl +Sm100ScalarReductionImpl = sm90_nodes.Sm90ScalarReductionImpl + + +class Sm100AuxLoadImpl(AuxLoadImpl): + @property + def descriptor(self) -> str: + """ + Descriptor for Aux Load + """ + return f"{self.name_camel}Descriptor" + + def decl_descriptor(self) -> str: + """ + Declare the descriptor type + """ + return f"\nusing {self.descriptor} = cutlass::epilogue::collective::detail::Sm100AuxLoadDescriptor;\n" + + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = self.decl_descriptor() + self._type_decl += f""" +using {self.name_camel} = cutlass::epilogue::fusion::Sm90AuxLoad< + {self.descriptor}::Stages, typename {self.descriptor}::EpilogueTile, {DataTypeTag[self.element]}, + {self.stride_mnl}, typename {self.descriptor}::SmemLayoutAtom, typename {self.descriptor}::CopyOpS2R +>; +""" + return self._type_decl + + def get_smem_size( + self, cta_tile_mnk, epilogue_tile_mn, stages_c, stages_d, epi_tiles + ): + """ + Get the shared memory size based on epilogue_tile_mn, stages_c, and stages_d + """ + return ( + DataTypeSize[self.element] * stages_c * product(epilogue_tile_mn) // 8, + 128, + ) + + +class Sm100AuxStoreImpl(AuxStoreImpl): + @property + def descriptor(self) -> str: + """ + Descriptor for Aux Load + """ + return f"{self.name_camel}Descriptor" + + def decl_descriptor(self) -> str: + """ + Declare the descriptor type + """ + return f""" +using {self.descriptor} = cutlass::epilogue::collective::detail::Sm100AuxStoreDescriptor< + EpilogueDescriptor, {self.stride_mnl}, {DataTypeTag[self.element]} +>; +""" + + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = self.decl_descriptor() + self._type_decl += f""" +using {self.name_camel} = cutlass::epilogue::fusion::Sm90AuxStore< + {self.descriptor}::Stages, typename {self.descriptor}::EpilogueTile, {DataTypeTag[self.element]}, + {FloatRoundStyleTag[self.round_style]}, {self.stride_mnl}, typename {self.descriptor}::SmemLayoutAtom, + typename {self.descriptor}::CopyOpR2S +>; +""" + return self._type_decl + + def get_smem_size( + self, cta_tile_mnk, epilogue_tile_mn, stages_c, stages_d, epi_tiles + ): + """ + Get the shared memory size based on epilogue_tile_mn, stages_c, and stages_d + """ + return ( + DataTypeSize[self.element] * stages_d * product(epilogue_tile_mn) // 8, + 128, + ) diff --git a/operators/cutlass/operators/fusion/backend/sm80_emitter.py b/operators/cutlass/operators/fusion/backend/sm80_emitter.py new file mode 100644 index 000000000..8fabca66d --- /dev/null +++ b/operators/cutlass/operators/fusion/backend/sm80_emitter.py @@ -0,0 +1,46 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Emitter for Sm80 Epilogue Visitor +""" + +from cutlass.operators.fusion.backend.emitter_base import FusionCallbacks + + +class Sm80Emitter: + def __init__(self, operation, graph) -> None: + self.fusion_callbacks = FusionCallbacks(graph, cc=80) + + def emit(self): + callback_decl, callback_name = self.fusion_callbacks.emit() + return callback_name, callback_decl diff --git a/operators/cutlass/operators/fusion/backend/sm80_nodes.py b/operators/cutlass/operators/fusion/backend/sm80_nodes.py new file mode 100644 index 000000000..6b08d1f9d --- /dev/null +++ b/operators/cutlass/operators/fusion/backend/sm80_nodes.py @@ -0,0 +1,247 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +from cutlass.operators.fusion.library import ( + DataTypeTag, + FloatRoundStyleTag, + FunctionalOp, + op_tag, +) + +from cutlass.operators.fusion.ir import ( + # Load Node + AccumulatorImpl, + AuxLoadImpl, + ColumnBroadcastImpl, + LoadNode, + RowBroadcastImpl, + ScalarBroadcastImpl, + # Compute Node + ComputeImpl, + # Store Node + AuxStoreImpl, + ColumnReductionImpl, + RowReductionImpl, + ScalarReductionImpl, +) + + +class Sm80AccumulatorImpl(AccumulatorImpl): + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f"""\nusing {self.name_camel} = cutlass::epilogue::threadblock::VisitorAccFetch;\n""" + return self._type_decl + + +class Sm80AuxLoadImpl(AuxLoadImpl): + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f""" +using {self.name_camel} = cutlass::epilogue::threadblock::VisitorAuxLoad< + OutputTileThreadMap, {DataTypeTag[self.element]}, {self.stride_mnl} +>; +""" + return self._type_decl + + +class Sm80LoadSrcImpl(Sm80AuxLoadImpl): + pass + + +class Sm80ScalarBroadcastImpl(ScalarBroadcastImpl): + def __init__(self, node: LoadNode) -> None: + super().__init__(node) + self.broadcast_count = 1 + self.reduction_fn = FunctionalOp.Multiplies + + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f""" +using {self.name_camel} = cutlass::epilogue::threadblock::VisitorScalarBroadcast< + {DataTypeTag[self.element]}, {self.stride_mnl}, {self.broadcast_count}, {op_tag(self.reduction_fn)} +>; +""" + return self._type_decl + + +class Sm80RowBroadcastImpl(RowBroadcastImpl): + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f""" +using {self.name_camel} = cutlass::epilogue::threadblock::VisitorRowBroadcast< + OutputTileThreadMap, {DataTypeTag[self.element]}, + {self.stride_mnl} +>; +""" + return self._type_decl + + +class Sm80ColumnBroadcastImpl(ColumnBroadcastImpl): + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f""" +using {self.name_camel} = cutlass::epilogue::threadblock::VisitorColBroadcast< + OutputTileThreadMap, {DataTypeTag[self.element]}, + {self.stride_mnl} +>; +""" + return self._type_decl + + +class Sm80ComputeImpl(ComputeImpl): + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f""" +using {self.name_camel} = cutlass::epilogue::threadblock::VisitorCompute< + {op_tag(self.fn)}, {DataTypeTag[self.element_output]}, {DataTypeTag[self.element_compute]}, + {FloatRoundStyleTag[self.round_style]} +>; +""" + return self._type_decl + + +class Sm80AuxStoreImpl(AuxStoreImpl): + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f""" +using {self.name_camel} = cutlass::epilogue::threadblock::VisitorAuxStore< + OutputTileThreadMap, {DataTypeTag[self.element]}, {FloatRoundStyleTag[self.round_style]}, + {self.stride_mnl} +>; +""" + return self._type_decl + + +class Sm80StoreDImpl(Sm80AuxStoreImpl): + pass + + +class Sm80ColumnReductionImpl(ColumnReductionImpl): + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f""" +using {self.name_camel} = cutlass::epilogue::threadblock::VisitorColReduction< + {op_tag(self.reg_reduce_fn)}, {op_tag(self.gmem_reduce_fn)}, + OutputTileThreadMap, {DataTypeTag[self.element]}, + {DataTypeTag[self.element_compute]}, {FloatRoundStyleTag[self.round_style]}, + {self.stride_mnl} +>; +""" + return self._type_decl + + +class Sm80RowReductionImpl(RowReductionImpl): + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f""" +using {self.name_camel} = cutlass::epilogue::threadblock::VisitorRowReduction< + {op_tag(self.reg_reduce_fn)}, {op_tag(self.gmem_reduce_fn)}, + OutputTileThreadMap, {DataTypeTag[self.element]}, + {DataTypeTag[self.element_compute]}, {FloatRoundStyleTag[self.round_style]}, + {self.stride_mnl} +>; +""" + return self._type_decl + + +class Sm80ScalarReductionImpl(ScalarReductionImpl): + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f""" +using {self.name_camel} = cutlass::epilogue::threadblock::VisitorScalarReduction< + {op_tag(self.reg_reduce_fn)}, {op_tag(self.gmem_reduce_fn)}, + OutputTileThreadMap, {DataTypeTag[self.element]}, + {DataTypeTag[self.element_compute]}, {FloatRoundStyleTag[self.round_style]}, + {self.stride_mnl} +>; +""" + return self._type_decl diff --git a/operators/cutlass/operators/fusion/backend/sm90_emitter.py b/operators/cutlass/operators/fusion/backend/sm90_emitter.py new file mode 100644 index 000000000..04ce10740 --- /dev/null +++ b/operators/cutlass/operators/fusion/backend/sm90_emitter.py @@ -0,0 +1,97 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Emitter for Sm90 Epilogue Visitor +""" + +from cutlass.operators.fusion.library import DataTypeTag, EpilogueScheduleTag +from cutlass.operators.fusion.backend.emitter_base import FusionCallbacks + + +class CollectiveEpilogue: + def __init__( + self, tile_description, schedule, element_c, element_d, fusion_callbacks + ) -> None: + self.cta_tile_mnk = tile_description.threadblock_shape + self.element_c = element_c + self.element_d = element_d + self.schedule = schedule + self.fusion_callbacks = fusion_callbacks + + @property + def CtaTileMNK(self) -> str: + """ + The threadblock shape + """ + return f"cute::Shape<_{self.cta_tile_mnk[0]}, _{self.cta_tile_mnk[1]}, _{self.cta_tile_mnk[2]}>" + + @property + def EpilogueTileType(self) -> str: + """ + The epilogue tile type + """ + return "cutlass::epilogue::collective::EpilogueTileAuto" + + @property + def Schedule(self) -> str: + return EpilogueScheduleTag[self.schedule] + + def emit(self): + callback_decl, callback_name = self.fusion_callbacks.emit() + return ( + callback_name, + f""" +using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor< + {self.CtaTileMNK}, {self.EpilogueTileType}, + {DataTypeTag[self.element_c]}, {DataTypeTag[self.element_d]}, + {self.Schedule} +>; +{callback_decl} +""", + ) + + +class Sm90Emitter: + def __init__(self, operation, graph) -> None: + fusion_callbacks = FusionCallbacks(graph, cc=90, emit_CD=False) + + self.collective_epilogue = CollectiveEpilogue( + tile_description=operation.tile_description, + schedule=operation.tile_description.epilogue_schedule, + element_c=operation.C.element, + element_d=operation.D.element, + fusion_callbacks=fusion_callbacks, + ) + + def emit(self): + return self.collective_epilogue.emit() diff --git a/operators/cutlass/operators/fusion/backend/sm90_nodes.py b/operators/cutlass/operators/fusion/backend/sm90_nodes.py new file mode 100644 index 000000000..d6f66aa84 --- /dev/null +++ b/operators/cutlass/operators/fusion/backend/sm90_nodes.py @@ -0,0 +1,327 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + + +from cutlass.operators.fusion.library import ( + DataTypeSize, + DataTypeTag, + FloatRoundStyleTag, + FunctionalOp, + op_tag, +) +from cutlass.operators.fusion.ir import ( + # Load Node + AccumulatorImpl, + AuxLoadImpl, + ColumnBroadcastImpl, + LoadNode, + LoadSrcImpl, + RowBroadcastImpl, + ScalarBroadcastImpl, + # Compute Node + ComputeImpl, + # Store Node + AuxStoreImpl, + ColumnReductionImpl, + RowReductionImpl, + ScalarReductionImpl, + StoreDImpl, +) +from cutlass.operators.fusion.pycute import product + + +class Sm90AccumulatorImpl(AccumulatorImpl): + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f"""\nusing {self.name_camel} = cutlass::epilogue::fusion::Sm90AccFetch;\n""" + return self._type_decl + + +class Sm90LoadSrcImpl(LoadSrcImpl): + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f""" +using ElementC = {DataTypeTag[self.element]}; +using StrideC = {self.stride_mnl}; +using {self.name_camel} = cutlass::epilogue::fusion::Sm90SrcFetch<{DataTypeTag[self.element]}>; +""" + return self._type_decl + + +class Sm90AuxLoadImpl(AuxLoadImpl): + @property + def descriptor(self) -> str: + """ + Descriptor for Aux Load + """ + return f"{self.name_camel}Descriptor" + + def decl_descriptor(self) -> str: + """ + Declare the descriptor type + """ + return f"\nusing {self.descriptor} = cutlass::epilogue::collective::detail::AuxLoadDescriptor;\n" + + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = self.decl_descriptor() + self._type_decl += f""" +using {self.name_camel} = cutlass::epilogue::fusion::Sm90AuxLoad< + {self.descriptor}::Stages, typename {self.descriptor}::EpilogueTile, {DataTypeTag[self.element]}, + {self.stride_mnl}, typename {self.descriptor}::SmemLayoutAtom, typename {self.descriptor}::CopyOpS2R +>; +""" + return self._type_decl + + def get_smem_size( + self, cta_tile_mnk, epilogue_tile_mn, stages_c, stages_d, epi_tiles + ): + """ + Get the shared memory size based on epilogue_tile_mn, stages_c, and stages_d + """ + return ( + DataTypeSize[self.element] * stages_c * product(epilogue_tile_mn) // 8, + 128, + ) + + +class Sm90ScalarBroadcastImpl(ScalarBroadcastImpl): + def __init__(self, node: LoadNode) -> None: + super().__init__(node) + self.broadcast_count = 1 + self.reduction_fn = FunctionalOp.Multiplies + + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f""" +using {self.name_camel} = cutlass::epilogue::fusion::Sm90ScalarBroadcast< + {DataTypeTag[self.element]}, {self.stride_mnl}, {self.broadcast_count}, {op_tag(self.reduction_fn)} +>; +""" + return self._type_decl + + +class Sm90RowBroadcastImpl(RowBroadcastImpl): + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f""" +using {self.name_camel} = cutlass::epilogue::fusion::Sm90RowBroadcast< + 0 /*Stages*/, typename EpilogueDescriptor::TileShape, {DataTypeTag[self.element]}, {DataTypeTag[self.element_output]}, + {self.stride_mnl} +>; +""" + return self._type_decl + + +class Sm90ColumnBroadcastImpl(ColumnBroadcastImpl): + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f""" +using {self.name_camel} = cutlass::epilogue::fusion::Sm90ColBroadcast< + 0 /*Stages*/, typename EpilogueDescriptor::TileShape, {DataTypeTag[self.element]}, {DataTypeTag[self.element_output]}, + {self.stride_mnl} +>; +""" + return self._type_decl + + +class Sm90ComputeImpl(ComputeImpl): + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f""" +using {self.name_camel} = cutlass::epilogue::fusion::Sm90Compute< + {op_tag(self.fn)}, {DataTypeTag[self.element_output]}, {DataTypeTag[self.element_compute]}, + {FloatRoundStyleTag[self.round_style]} +>; +""" + return self._type_decl + + +class Sm90AuxStoreImpl(AuxStoreImpl): + @property + def descriptor(self) -> str: + """ + Descriptor for Aux Load + """ + return f"{self.name_camel}Descriptor" + + def decl_descriptor(self) -> str: + """ + Declare the descriptor type + """ + return f""" +using {self.descriptor} = cutlass::epilogue::collective::detail::AuxStoreDescriptor< + EpilogueDescriptor, {self.stride_mnl}, {DataTypeTag[self.element]} +>; +""" + + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = self.decl_descriptor() + self._type_decl += f""" +using {self.name_camel} = cutlass::epilogue::fusion::Sm90AuxStore< + {self.descriptor}::Stages, typename {self.descriptor}::EpilogueTile, {DataTypeTag[self.element]}, + {FloatRoundStyleTag[self.round_style]}, {self.stride_mnl}, typename {self.descriptor}::SmemLayoutAtom, + typename {self.descriptor}::CopyOpR2S +>; +""" + return self._type_decl + + def get_smem_size( + self, cta_tile_mnk, epilogue_tile_mn, stages_c, stages_d, epi_tiles + ): + """ + Get the shared memory size based on epilogue_tile_mn, stages_c, and stages_d + """ + return ( + DataTypeSize[self.element] * stages_d * product(epilogue_tile_mn) // 8, + 128, + ) + + +class Sm90StoreDImpl(StoreDImpl): + @property + def type_decl(self): + """ + Return the string defining the type + """ + return f""" +using ElementD = {DataTypeTag[self.element]}; +using StrideD = {self.stride_mnl}; +""" + + +class Sm90ColumnReductionImpl(ColumnReductionImpl): + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f""" +using {self.name_camel} = cutlass::epilogue::fusion::Sm90ColReduction< + {op_tag(self.reg_reduce_fn)}, {op_tag(self.reg_reduce_fn)}, {op_tag(self.gmem_reduce_fn)}, 0, + typename EpilogueDescriptor::TileShape, {DataTypeTag[self.element]}, + {DataTypeTag[self.element_compute]}, {FloatRoundStyleTag[self.round_style]}, + {self.stride_mnl} +>; +""" + return self._type_decl + + +class Sm90RowReductionImpl(RowReductionImpl): + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f""" +using {self.name_camel} = cutlass::epilogue::fusion::Sm90RowReduction< + {op_tag(self.reg_reduce_fn)}, {op_tag(self.reg_reduce_fn)}, {op_tag(self.gmem_reduce_fn)}, 0 /* Stages */, + typename EpilogueDescriptor::TileShape, {DataTypeTag[self.element]}, + {DataTypeTag[self.element_compute]}, {FloatRoundStyleTag[self.round_style]}, + {self.stride_mnl} +>; +""" + return self._type_decl + + +class Sm90ScalarReductionImpl(ScalarReductionImpl): + @property + def type_decl(self): + """ + Return the string defining the type + """ + if self._type_decl is not None: + return self._type_decl + + self._type_decl = f""" +using {self.name_camel} = cutlass::epilogue::fusion::Sm90ScalarReduction< + {op_tag(self.reg_reduce_fn)}, {op_tag(self.gmem_reduce_fn)}, + {DataTypeTag[self.element]}, {DataTypeTag[self.element_compute]}, + {FloatRoundStyleTag[self.round_style]}, {self.stride_mnl} +>; +""" + return self._type_decl diff --git a/operators/cutlass/operators/fusion/epilogue.py b/operators/cutlass/operators/fusion/epilogue.py new file mode 100644 index 000000000..de2f06c5a --- /dev/null +++ b/operators/cutlass/operators/fusion/epilogue.py @@ -0,0 +1,88 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Epilogue Visitor interface for compiling, and running visitor-based epilogue. +""" + +import cutlass.operators.fusion.backend +from cutlass.operators.fusion.passes.util import cc_map + + +class EpilogueFunctorVisitor: + """ + Apply an epilogue functor described by the epilogue EVT + + :param cc: compute capability + :param visitor_frontend: user-provide visitor frontend + + """ + + def __init__(self, cc: int, visitor, element_compute) -> None: + # Type of Emitter based on CC + self.emit_cls = getattr( + cutlass.operators.fusion.backend, f"Sm{cc_map[cc]}Emitter" + ) + + # Visitor Types + self.visitor = visitor + self.graph = visitor.dag_ir + + # Data types + self.element_epilogue = element_compute # element compute + self.element_output = self.graph.get_node_meta("D").underlying_impl.element + + # Epilogue Thread Type + if cc_map[cc] in [90, 100]: + self.arg_c_type = self.visitor.arg_c_type + self.arg_d_type = self.visitor.arg_d_type + + # Epilogue stages specialized for sm80 kernel + if cc == 80: + if hasattr(self.visitor, "epilogue_stages"): + self.epilogue_stages = self.visitor.epilogue_stages + assert self.epilogue_stages <= 2, ( + "Only supports Stages <=2 in SM80 Epilogue" + ) + + def emit(self, operation): + """ + Emit the C++ code + """ + emitter = self.emit_cls(operation, self.graph) + return emitter.emit() + + def get_smem_size(self, tile_description): + """ + Get the shared memory size in bytes + """ + return self.visitor.get_smem_size(tile_description) diff --git a/operators/cutlass/operators/fusion/evt_ops.py b/operators/cutlass/operators/fusion/evt_ops.py new file mode 100644 index 000000000..fed5bc74a --- /dev/null +++ b/operators/cutlass/operators/fusion/evt_ops.py @@ -0,0 +1,116 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Collection of builtin functions used for host reference in EVT +""" + +from cutlass.operators.utils.common import is_torch_available, is_numpy_available +from cutlass.operators.utils.tensor import is_numpy_tensor, is_torch_tensor + +if is_numpy_available(): + import numpy as np + +if is_torch_available(): + import torch + + +def multiply_add(x, y, z): + return x * y + z + + +def sum(x, dim): + if is_numpy_tensor(x): + return x.sum(axis=tuple(dim)) + elif is_torch_tensor(x): + return torch.sum(x, dim) + else: + raise TypeError(f"Unsupported tensor type: {type(x)}") + + +def max(x, dim): + if is_numpy_tensor(x): + return x.max(axis=tuple(dim)) + elif is_torch_tensor(x): + return torch.amax(x, dim) + else: + raise TypeError(f"Unsupported tensor type: {type(x)}") + + +def maximum(x, y): + if is_numpy_tensor(x): + return np.maximum(x, y) + elif is_torch_tensor(x): + return torch.maximum(x, torch.tensor(y)) + else: + raise TypeError(f"Unsupported tensor type: {type(x)}") + + +def minimum(x, y): + if is_numpy_tensor(x): + return np.minimum(x, y) + elif is_torch_tensor(x): + return torch.minimum(x, torch.tensor(y)) + else: + raise TypeError(f"Unsupported tensor type: {type(x)}") + + +def exp(x): + if is_numpy_tensor(x): + return np.exp(x) + elif is_torch_tensor(x): + return torch.exp(x) + else: + raise TypeError(f"Unsupported tensor type: {type(x)}") + + +############################################################################## +# Layout manipulation nodes +############################################################################## + + +def permute(x, indices: tuple): + if is_numpy_tensor(x): + return np.transpose(x, axes=indices) + elif is_torch_tensor(x): + return x.permute(*indices) + else: + raise TypeError(f"Unsupported tensor type: {type(x)}") + + +def reshape(x, new_shape: tuple): + if is_numpy_tensor(x): + return np.reshape(x, newshape=new_shape) + elif is_torch_tensor(x): + return x.view(new_shape) + else: + raise TypeError(f"Unsupported tensor type: {type(x)}") diff --git a/operators/cutlass/operators/fusion/frontend/__init__.py b/operators/cutlass/operators/fusion/frontend/__init__.py new file mode 100644 index 000000000..3ba0598da --- /dev/null +++ b/operators/cutlass/operators/fusion/frontend/__init__.py @@ -0,0 +1,41 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +from cutlass.operators.fusion.frontend.python_ast import ( + PythonASTFrontend, + PythonASTInOutProcessor, +) + +__all__ = [ + "PythonASTFrontend", + "PythonASTInOutProcessor", +] diff --git a/operators/cutlass/operators/fusion/frontend/frontend_base.py b/operators/cutlass/operators/fusion/frontend/frontend_base.py new file mode 100644 index 000000000..4638c3c4c --- /dev/null +++ b/operators/cutlass/operators/fusion/frontend/frontend_base.py @@ -0,0 +1,303 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Base class for Python EVT Frontend +""" + +from typing import Union + +from cutlass.operators.fusion.library import DataType +from cutlass.operators.fusion.ir import ( + ComputeNode, + DAGIR, + LayoutNode, + LoadNode, + StoreNode, +) +from cutlass.operators.fusion.passes import ( + EVTGraphDrawer, + EVTPassManager, + GetSmemSize, + PassDAG2Tree, + PassGetArgumentType, + PassGetImpl, + PassFixElementD, + PassLayoutManipulateElimination, + PassPreprocessRed, + PassShapeTypePropagation, +) +from cutlass.operators.fusion.passes.util import cc_map +from cutlass.operators.fusion.evt_ops import permute, reshape + + +class EVTFrontendBase: + layout_fns = {"permute": permute, "reshape": reshape} + + def __init__( + self, + cc, + element_compute=DataType.f32, + additional_passes=None, + requires_conversion_to_tree=True, + **kwargs, + ) -> None: + self.cc = cc + self.element_compute = element_compute + self.dag_ir = DAGIR(self.cc, self.element_compute) + self.compute_cnt = 0 + self.layout_cnt = 0 + self.imm_cnt = 0 + + if additional_passes is None: + additional_passes = [] + + passes = [ + PassPreprocessRed, + PassGetArgumentType, + PassShapeTypePropagation, + PassLayoutManipulateElimination, + PassGetImpl, + PassDAG2Tree, + PassFixElementD, + ] + additional_passes + + soft_dependencies = [] + if not requires_conversion_to_tree: + passes.remove(PassDAG2Tree) + soft_dependencies.append(PassDAG2Tree) + + self.pass_manager = EVTPassManager( + self.dag_ir, passes, soft_dependencies=soft_dependencies + ) + + if self.cc == 80: + self._epilogue_stages = 1 + else: + self._epilogue_stages = None + + @property + def epilogue_stages(self): + return self._epilogue_stages + + @epilogue_stages.setter + def epilogue_stages(self, stages): + self._epilogue_stages = stages + + def parse(self, *args, **kwargs): + raise NotImplementedError( + "The 'parse' function must be overloaded in frontend class" + ) + + def trace(self, *args, **kwargs): + # Parse the input + self.parse(*args, **kwargs) + + if not self.dag_ir.has_node("D"): + raise RuntimeError("Output node D is not found in the epilogue function.") + + # Verify the DAG IR to ensure that "D" is the output node with out_degree = 0 + if self.cc >= 90: + if self.dag_ir.out_degree("D") != 0: + raise RuntimeError( + f"On SM90 or higher, D is expected to be a output node with 0 users to " + f"enable smem reuse between C and D, but got {self.dag_ir.out_degree('D')}" + ) + + # Run the passes + self.pass_manager() + # Set the epilogue type + self.epilogue_thread_type = self.dag_ir.epilogue_thread_type + if cc_map[self.cc] in [90, 100]: + self.arg_c_type = self.dag_ir.arg_c_type + self.arg_d_type = self.dag_ir.arg_d_type + self.reduction_names = self.dag_ir.reduction_names + + # + # Helper functions for DAG IR manipulation + # + + def add_node(self, node): + self.dag_ir.add_node(node) + + def add_edge(self, src, tgt, weight=0): + self.dag_ir.add_edge(src, tgt, weight=weight) + + def set_tensor(self, node_name, example): + """ + Add an example tensor to node {node_name} in the DAG IR + """ + meta = self.dag_ir.get_node_meta(node_name) + meta.tensor = {"tensor": example} + + def set_store_tensor(self, node_name, example): + """ + Add an example tensor to node {node_name} in the DAG IR + """ + meta = self.dag_ir.get_node_meta(node_name) + meta.store_tensor = {"tensor": example} + + def mark_output(self, node_name): + """ + Mark a store node as output + """ + meta = self.dag_ir.get_node_meta(node_name) + # Allow accum to also be used as an output so that + # a user can include it in the return line + if not isinstance(meta, StoreNode) and not meta.name == "accum": + raise ValueError( + f"Only StoreNodes can be marked as output. " + f"Got {type(meta).__name__}: {node_name}" + ) + meta.is_output = True + + # Add node with specific type + + def add_load_node(self, name, example): + """ + Add a Load node to DAG IR + :param name: name of the loaded variable + :type name: str + :param example: example input + :type example: np.ndarray|torch.Tensor|cupy.ndarray|float + """ + if name is None: + raise ValueError("Name is not provided.") + if example is None: + raise ValueError(f"Example input for {name} is not provided.") + load_node = LoadNode(name) + load_node.tensor = {"tensor": example} + # Special logics for accumulator + if name == "accum": + if load_node.tensor.rank == 2: + new_shape = tuple( + [ + 1, + ] + + list(load_node.tensor.shape) + ) + load_node.tensor.broadcast(new_shape) + elif load_node.tensor.rank < 2 or load_node.tensor.rank > 3: + raise ValueError( + f"Expect example inputs for 'accum' be a rank-2 or rank-3 tensor. Got {load_node.tensor.shape}." + ) + self.add_node(load_node) + + def add_imm(self, value: Union[float, int]): + """ + Add an immediate scalar value to DAG IR + :param value: the value of the immediate scalar + :type value: float + """ + try: + value = float(value) + except Exception as e: + raise ValueError( + f"{type(value).__name__} cannot be converted to float." + ) from e + + name = f"imm_{value}_k{self.imm_cnt}".replace(".", "_") + self.imm_cnt += 1 + load_node = LoadNode(name) + load_node.tensor = {"tensor": value, "is_constant": True} + self.add_node(load_node) + return name + + def add_compute_node(self, op, name=None): + """ + Add a compute node. + :param op: the computation op + :param name: the node name (optional) + :type name: str + :return: the name of the compute node + """ + if name is None: + name = f"compute_{self.compute_cnt}" + self.compute_cnt += 1 + compute_node = ComputeNode( + name=name, + fn=op, + element_output=self.element_compute, + element_compute=self.element_compute, + ) + self.add_node(compute_node) + return compute_node.name + + def add_layout_node(self, op, kwargs, name=None): + """ + Add a layout node. + :param op: the layout op + :type op: evt_ops + :param name: the node name (optional) + :type name: str + :return: the name of the layout node + """ + if name is None: + name = f"layout_{self.layout_cnt}" + self.layout_cnt += 1 + layout_node = LayoutNode(name=name, fn=op, kwargs=kwargs) + self.add_node(layout_node) + return layout_node.name + + def add_store_node(self, name): + store_node = StoreNode(name) + self.add_node(store_node) + + # + # Visualization The DAG IR + # + + def visualize(self, name="dag_ir"): + """ + Visualize the dag ir with svg file + :param name: the name of the graph + """ + drawer = EVTGraphDrawer(self.dag_ir, name) + try: + for name, graph in drawer.get_dot_graph(): + graph.write_svg(f"./{name}.svg") + except: + raise RuntimeError( + "'dot' is not found in path. GraphDrawer is disabled. " + "Please install it with 'sudo apt-get install graphviz'." + ) + + # + # Get shared memory size + # + + def get_smem_size(self, tile_description): + """ + Get the shared memory size of the epilogue + """ + smem_size = GetSmemSize(self.dag_ir)(tile_description) + return smem_size diff --git a/operators/cutlass/operators/fusion/frontend/python_ast.py b/operators/cutlass/operators/fusion/frontend/python_ast.py new file mode 100644 index 000000000..fac2a625a --- /dev/null +++ b/operators/cutlass/operators/fusion/frontend/python_ast.py @@ -0,0 +1,261 @@ +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Python AST frontend that parses input into DAG IR +""" + +import ast +import inspect +import textwrap +from typing import Callable, Union + + +from cutlass.operators.fusion.activation import ( + identity, + relu, + tanh, + sigmoid, + silu, + hardswish, + gelu, +) +from cutlass.operators.fusion.frontend.frontend_base import EVTFrontendBase +from cutlass.operators.fusion.library import DataType, FunctionalOp + + +class PythonASTFrontend(EVTFrontendBase, ast.NodeVisitor): + def __init__(self, cc, element_compute=DataType.f32, **kwargs): + super().__init__(cc, element_compute, **kwargs) + # Flags + # If this state is True, visit_Constant returns values without creating imm node + self.no_imm = False + self.visiting_return = False + + def parse(self, example_inputs): + self.example_inputs = example_inputs + self.source = textwrap.dedent(inspect.getsource(self.__call__)) + self.ast = ast.parse(self.source) + self.visit(self.ast) + + # + # Helper functions + # + @staticmethod + def ast_op_to_bindings(op): + mapping = { + ast.Add: FunctionalOp.Plus, + ast.Sub: FunctionalOp.Minus, + ast.Mult: FunctionalOp.Multiplies, + ast.Div: FunctionalOp.Divides, + "maximum": FunctionalOp.Maximum, + "minimum": FunctionalOp.Minimum, + "identity": identity.binding_type, + "relu": relu.binding_type, + "tanh": tanh.binding_type, + "sigmoid": sigmoid.binding_type, + "silu": silu.binding_type, + "hardswish": hardswish.binding_type, + "gelu": gelu.binding_type, + "multiply_add": FunctionalOp.MultiplyAdd, + "sum": (FunctionalOp.Plus, FunctionalOp.AtomicAdd), + "max": (FunctionalOp.Maximum, FunctionalOp.AtomicMaximum), + "exp": FunctionalOp.Exp, + } + return mapping[op] + + # + # Visiting different node types + # + + def visit_FunctionDef(self, node: ast.FunctionDef): + # Visit args and register load nodes + for arg in node.args.args: + self.visit(arg) + for expr in node.body: + self.visit(expr) + + def visit_arg(self, node: ast.arg): + # Name of the argument + name = node.arg + try: + example_tensor = self.example_inputs[name] + except Exception as e: + raise RuntimeError(f"Example input for {name} is not provided.") from e + + self.add_load_node(name, example_tensor) + + def visit_Name(self, node: ast.Name): + return node.id + + def visit_Constant(self, node: ast.Constant): + if self.no_imm: + return node.value + else: + name = self.add_imm(node.value) + return name + + def visit_Tuple(self, node: ast.Tuple): + results = [] + for elt in node.elts: + results.append(self.visit(elt)) + return tuple(results) + + def visit_keyword(self, node: ast.keyword): + return {node.arg: self.visit(node.value)} + + def visit_BinOp(self, node: ast.BinOp): + if self.visiting_return: + raise SyntaxError("Return value cannot be an expression") + lhs = self.visit(node.left) + rhs = self.visit(node.right) + op = self.ast_op_to_bindings(type(node.op)) + name = self.add_compute_node(op) + + # Add edges + # The edge weights are used to sort the input args + self.add_edge(lhs, name, weight=0) + self.add_edge(rhs, name, weight=1) + return name + + def visit_Assign(self, node: ast.Assign): + target = self.visit(node.targets[0]) + value = self.visit(node.value) + # Create the assign node + self.add_store_node(target) + + # Add edges + self.add_edge(value, target) + return target + + def visit_Call(self, node: ast.Call): + if self.visiting_return: + raise SyntaxError("Return value cannot be an expression") + func = self.visit(node.func) + args = [self.visit(arg) for arg in node.args] + + if func in self.layout_fns.keys(): + # Parse kwargs + # By default, visiting imm automatically creates a load node + # However, in function call, keyword args are used to set + # specific function attributes such as indices for permute + # So no_imm is set to True temporarily + self.no_imm = True + kwargs = {} + for kw in node.keywords: + kwargs.update(self.visit(kw)) + self.no_imm = False + op = self.layout_fns[func] + name = self.add_layout_node(op, kwargs) + else: + op = self.ast_op_to_bindings(func) + name = self.add_compute_node(op) + + # Add edges + for idx, arg in enumerate(args): + self.add_edge(arg, name, weight=idx) + return name + + def visit_Return(self, node: ast.Return): + self.visiting_return = True + results = self.visit(node.value) + self.visiting_return = False + self.return_names = results + if not isinstance(results, tuple): + results = (results,) + for rst in results: + try: + example_tensor = self.example_inputs[rst] + except Exception as e: + raise RuntimeError(f"Example input for {rst} is not provided.") from e + self.set_store_tensor(rst, example_tensor) + self.mark_output(rst) + + +class PythonASTInOutProcessor(ast.NodeVisitor): + """ + Simple processor that traces the input AST and returns the names of inputs + and outputs + + :param source: The source code of the function to trace. + :type source: Union[Callable, str] + :return: A tuple of lists of input and output names. + :rtype: tuple[list[str], list[str]] + """ + + def trace(self, source: Union[Callable, str]) -> tuple[list[str], list[str]]: + if callable(source): + self.source = textwrap.dedent(inspect.getsource(source)) + if isinstance(source, str): + self.source = source + self.ast = ast.parse(self.source) + + self.inputs = [] + self.outputs = [] + self.visiting_return = False + + self.visit(self.ast) + if "accum" not in self.inputs: + raise ValueError("accum must be an input to the epilogue function") + + # Users do not need to specify accum. Remove it. + self.inputs.remove("accum") + return self.inputs, self.outputs + + def visit_FunctionDef(self, node: ast.FunctionDef): + for arg in node.args.args: + self.visit(arg) + for expr in node.body: + self.visit(expr) + + def visit_arg(self, node: ast.arg): + self.inputs.append(node.arg) + + def visit_Constant(self, node: ast.Constant): + if self.visiting_return: + raise SyntaxError("Return value cannot be a constant") + + def visit_Name(self, node: ast.Name): + return node.id + + def visit_Tuple(self, node: ast.Tuple): + results = [] + for elt in node.elts: + results.append(self.visit(elt)) + return tuple(results) + + def visit_Return(self, node: ast.Return): + self.visiting_return = True + results = self.visit(node.value) + self.visiting_return = False + if not isinstance(results, tuple): + results = (results,) + self.outputs.extend(results) diff --git a/operators/cutlass/operators/fusion/ir/__init__.py b/operators/cutlass/operators/fusion/ir/__init__.py new file mode 100644 index 000000000..a393076d8 --- /dev/null +++ b/operators/cutlass/operators/fusion/ir/__init__.py @@ -0,0 +1,75 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +from cutlass.operators.fusion.ir.compute_nodes import ComputeNode, ComputeImpl +from cutlass.operators.fusion.ir.dag_ir import DAGIR +from cutlass.operators.fusion.ir.layout_nodes import LayoutNode +from cutlass.operators.fusion.ir.load_nodes import ( + LoadNode, + AccumulatorImpl, + LoadSrcImpl, + AuxLoadImpl, + RowBroadcastImpl, + ColumnBroadcastImpl, + ScalarBroadcastImpl, +) +from cutlass.operators.fusion.ir.node import TopoVisitorNode, NoOpImpl +from cutlass.operators.fusion.ir.store_nodes import ( + StoreNode, + StoreDImpl, + AuxStoreImpl, + ColumnReductionImpl, + RowReductionImpl, + ScalarReductionImpl, +) + +__all__ = [ + "ComputeNode", + "ComputeImpl", + "DAGIR", + "LayoutNode", + "LoadNode", + "AccumulatorImpl", + "LoadSrcImpl", + "AuxLoadImpl", + "RowBroadcastImpl", + "ColumnBroadcastImpl", + "ScalarBroadcastImpl", + "TopoVisitorNode", + "NoOpImpl", + "StoreNode", + "StoreDImpl", + "AuxStoreImpl", + "ColumnReductionImpl", + "RowReductionImpl", + "ScalarReductionImpl", +] diff --git a/operators/cutlass/operators/fusion/ir/c_types.py b/operators/cutlass/operators/fusion/ir/c_types.py new file mode 100644 index 000000000..7232097e5 --- /dev/null +++ b/operators/cutlass/operators/fusion/ir/c_types.py @@ -0,0 +1,247 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +import ctypes + +from cutlass.operators.fusion.library import DataType +from cutlass.operators.utils.common import is_numpy_available +from cutlass.operators.utils.tensor import is_numpy_tensor, is_torch_tensor + + +dtype2ctype = { + DataType.f16: ctypes.c_uint16, + DataType.bf16: ctypes.c_uint16, + DataType.f32: ctypes.c_float, + DataType.f64: ctypes.c_double, + DataType.s8: ctypes.c_int8, + DataType.s32: ctypes.c_int32, +} + + +def get_scalar(value): + """ + Returns a scalar value from a container (e.g., np.ndarray) + """ + if is_numpy_tensor(value) or is_torch_tensor(value): + if value.size != 1: + raise Exception("Scalars used in epilogue must be of size 1") + return value.reshape(-1)[0] + else: + return value + + +def to_ctype_value(value, dtype): + """ + Converts ``value`` to the corresponding storage needed for the ctype that + will store ``value``. + """ + scalar = get_scalar(value) + if dtype == DataType.f16: + if is_numpy_available(): + import numpy as np + + # Convert f16 value into an integer + return int.from_bytes(np.float16(scalar).tobytes(), "little") + else: + raise NotImplementedError("Numpy is not available") + else: + return scalar + + + +class Empty(ctypes.Structure): + _fields_ = [] + + def __init__(self, *arg) -> None: + pass + + +class EmptyByte(ctypes.Structure): + _fields_ = [("byte", ctypes.c_byte)] + + def __init__(self, *arg) -> None: + pass + + +class EBO: + def __init__(self, index: int, type) -> None: + self.index = index + self.type = type + + def __eq__(self, other) -> bool: + if isinstance(other, EBO): + return self.index == other.index and self.type == other.type + return False + + def __hash__(self) -> int: + return hash((self.index, self.type)) + + def __ne__(self, other): + return not self.__eq__(other) + + def __str__(self) -> str: + return f"<{self.index}, {self.type}>" + + +def tuple_factory_(input_tuple, dtype, constants: list[int] = None): + """ + The factory function generating cute::Tuple with input tuple + :param input_tuple: the input tuple + :type input_tuple: tuple + :param dtype: the data type for non-constant values + :type dtype: str, "int32_t", "int", "int64_t" + :param constant: the values that will be treated as constants + :type constant: list[int] + + :return: ctype structure representing the cute::Tuple + :return: the empty base classes of the tuple + """ + + if constants is None: + constants = [0, 1] + + # The empty base classes of the current tuple + empty_bases = [] + # The first non empty base class + first_non_empty_base = None + # The ctype fields of the current tuple + ctype_fields = [] + + for idx, entry in enumerate(input_tuple): + # For nested tuples + if isinstance(entry, tuple): + sub_tuple_ctype, sub_empty_bases = tuple_factory_(entry, dtype, constants) + if ctypes.sizeof(sub_tuple_ctype) == 0: + # The empty tuple base class is also an empty EBO + empty_bases.append(EBO(idx, entry)) + else: + if first_non_empty_base is None: + first_non_empty_base = sub_empty_bases + ctype_fields.append((f"entry_{idx}", sub_tuple_ctype)) + else: + if entry in constants: + empty_bases.append(EBO(idx, entry)) + ctype_fields.append((f"entry_{idx}", Empty)) + else: + ctype_fields.append((f"entry_{idx}", dtype)) + if first_non_empty_base is None: + first_non_empty_base = [] + + # Create the ctype tuple + class TupleType(ctypes.Structure): + _fields_ = ctype_fields + + def __init__(self, args) -> None: + fields = self._fields_ + + assert len(fields) == len(args) + for field, arg in zip(fields, args): + name = field[0] + field_type = field[1] + setattr(self, name, field_type(arg)) + + return TupleType, empty_bases + + +def tuple_factory(input_tuple, dtype: str, constants=[0, 1]): + """ + The factory function generating cute::Tuple with input tuple + :param input_tuple: the input tuple + :type input_tuple: tuple + :param dtype: the data type for non-constant values + :type dtype: str, "int32_t", "int", "int64_t" + :param constant: the values that will be treated as constants + :type constant: list[int] + + :return: ctype structure representing the cute::Tuple + :return: the empty base classes of the tuple + """ + # Step 1: convert the dtype + if dtype == "int64_t": + dtype = ctypes.c_longlong + elif dtype in ["int", "int32_t"]: + dtype = ctypes.c_int32 + else: + raise NotImplementedError(f"Type {dtype} is not supported") + + tuple_type, _ = tuple_factory_(input_tuple, dtype, constants) + + if ctypes.sizeof(tuple_type) == 0: + return EmptyByte + return tuple_type + + +def visitor_factory(node_types, node_names): + """ + Creates the argument type of epilogue visitor type + + :param node_types: list of argument types under ctypes + :param node_names: list of argument names under str + + :return: tuple type in ctypes.Structure + """ + ctypes_field = [] + # Struct is used when number of nodes < 4 + # Because the Sm90VisitorImplBase has specification up to 4 nodes + # in `include/cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp` + if len(node_types) <= 4: + for idx, node_type in enumerate(node_types): + if ctypes.sizeof(node_type) == 0: + # Special case for empty struct + # 1 byte placeholder is used for correct alignment + ctypes_field.append((node_names[idx], ctypes.c_byte)) + else: + ctypes_field.append((node_names[idx], node_type)) + + class VisitorType(ctypes.Structure): + _fields_ = ctypes_field + + def __init__(self, kwargs) -> None: + for field in self._fields_: + fname, ftype = field + if ftype != ctypes.c_byte: + setattr(self, fname, ftype(kwargs)) + + # For cases with more than 4 nodes, tuple is used + else: + for idx, node_type in enumerate(node_types): + ctypes_field.append((node_names[idx], node_type)) + + class VisitorType(ctypes.Structure): + _fields_ = ctypes_field + + def __init__(self, kwargs) -> None: + for field in self._fields_: + fname, ftype = field + setattr(self, fname, ftype(kwargs)) + + return VisitorType diff --git a/operators/cutlass/operators/fusion/ir/compute_nodes.py b/operators/cutlass/operators/fusion/ir/compute_nodes.py new file mode 100644 index 000000000..914981919 --- /dev/null +++ b/operators/cutlass/operators/fusion/ir/compute_nodes.py @@ -0,0 +1,99 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Python registration for compute nodes in EVT +""" + +from cutlass.operators.fusion.ir.node import NodeBase, ImplBase +from cutlass.operators.fusion.library import FloatRoundStyle + + +class ComputeImplBase(ImplBase): + """ + Base class for compute implementation + """ + + def __init__(self, node) -> None: + super().__init__(node) + + +class ComputeImpl(ComputeImplBase): + """ + Implementation for Compute Node + """ + + def __init__(self, node) -> None: + super().__init__(node) + + self.fn = node.fn + self.element_output = node.element_output + self.element_compute = node.element_compute + self.round_style = node.round_style + + @staticmethod + def match(node, problem_size: tuple): + return True + + +class ComputeNode(NodeBase): + """ + Compute Node in DAG IR + """ + + possible_impls = [ + ComputeImpl, + ] + + def __init__( + self, + name: str, + fn, + element_output, + element_compute, + round_style=FloatRoundStyle.ToNearest, + ) -> None: + super().__init__(name) + self.op = "compute" + self.fn = fn + self.element_compute = element_compute + self.round_style = round_style + + def type_propagation(self, *args, **kwargs): + """ + Load node loads tensor under type `tensor.element` and returns an array of type `tensor.element`. + """ + self.element = self.element_compute + # In general, the compute nodes have element_output = element_compute + # In certain cases like producer of D it is overwritten by other passes + if not hasattr(self, "element_output"): + self.element_output = self.element diff --git a/operators/cutlass/operators/fusion/ir/dag_ir.py b/operators/cutlass/operators/fusion/ir/dag_ir.py new file mode 100644 index 000000000..0ac05baeb --- /dev/null +++ b/operators/cutlass/operators/fusion/ir/dag_ir.py @@ -0,0 +1,258 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +DAG IR used by Python EVT +""" + +import networkx as nx + + +from cutlass.operators.fusion.ir.compute_nodes import ComputeNode +from cutlass.operators.fusion.library import DataType +from cutlass.operators.fusion.ir.node import NodeBase +from cutlass.operators.fusion.library import ActivationOp + + +class DAGIR: + """ + ``DAGIR`` is the main data structure used in the EVT Intermediate Representation. + It consists of a series of ``Node`` s, each representing epilogue visitor nodes. + + In the DAGIR, ``node`` is an string of its name. ``node_meta`` is the underlying class of the node + """ + + def __init__(self, cc, element_compute=DataType.f32) -> None: + # The EVT DAGIR is managed through the nextworkX Digraph class + self._graph = nx.DiGraph() + + self.element_compute = element_compute + + self.reduction_names = [] + + self.cc = cc + + self.identity_counter = 0 + + # + # IR manipulator + # + + def add_node(self, meta: NodeBase): + """ + Add a node to dag ir + """ + if self.has_node(meta.name): + raise SyntaxError(f"Variable '{meta.name}' cannot be defined twice.") + self._graph.add_node(meta.name, meta=meta) + + def add_edge(self, src: str, dst: str, weight: int = 0): + """ + Add an edge src -> dst to dag ir with weight + """ + if not self.has_node(src): + raise SyntaxError(f"Variable '{src}' is undefined.") + if not self.has_node(dst): + raise SyntaxError(f"Variable '{dst}' is undefined.") + + if self._graph.has_edge(src, dst): + # The DiGraph doesn't support multiple edges between two nodes + # We insert an identity node in such case as a workaround + identity_name = f"autogen_identity_{self.identity_counter}" + self.identity_counter += 1 + compute_node = ComputeNode( + name=identity_name, + fn=ActivationOp.Identity, + element_output=self.element_compute, + element_compute=self.element_compute, + ) + self.add_node(compute_node) + self.add_edge(src, identity_name, 0) + self.add_edge(identity_name, dst, weight) + else: + self._graph.add_edge(src, dst, weight=weight) + + def remove_node(self, node: str): + """ + Remove node from dag ir + """ + self._graph.remove_node(node) + + def remove_edge(self, src: str, dst: str): + """ + Remove edge src -> dst + """ + self._graph.remove_edge(src, dst) + + # + # Helper functions for getting attrs + # + + def has_node(self, node: str) -> bool: + """ + Check if the node is in the graph + """ + return self._graph.has_node(node) + + def in_degree(self, node: str): + """ + Get the input degree of node + """ + return self._graph.in_degree(node) + + def in_edges(self, node: str): + """ + Get the input edges of node + """ + return [edge for edge in self._graph.in_edges(node)] + + def out_degree(self, node: str): + """ + Get the output degree of node + """ + return self._graph.out_degree(node) + + def out_edges(self, node: str): + """ + Get the output edges of node + """ + return [edge for edge in self._graph.out_edges(node)] + + def get_node_meta(self, node: str): + """ + Get the meta data of the node + """ + return self._graph.nodes[node]["meta"] + + def get_edge_weight(self, src, dst): + """ + Get the edge weight of edge src->dst + """ + return self._graph.get_edge_data(src, dst)["weight"] + + # + # High-level helper functions + # + + def all_reachable_nodes(self, node: str): + """ + Get all the nodes reachable from the current node (exclude) + """ + return list(nx.dfs_preorder_nodes(self._graph, source=node)) + + def get_users(self, node: str): + """ + Get all users of the current node + """ + return [edge[1] for edge in self.out_edges(node)] + + def get_all_inputs(self, node: str): + """ + Get all the input nodes sorted by edge weight + """ + in_edges = self.in_edges(node) + edge_weights = [self.get_edge_weight(*edge) for edge in in_edges] + return [edge[0] for _, edge in sorted(zip(edge_weights, in_edges))] + + def get_all_inputs_meta(self, node: str): + """ + Get all the input node metas sorted by edge weight + """ + return [ + self.get_node_meta(input_node) for input_node in self.get_all_inputs(node) + ] + + def replace_all_uses_with(self, node1, node2): + """ + Replace all uses of node1 with node2 + """ + for edge in self.out_edges(node1): + weight = self.get_edge_weight(*edge) + user = edge[1] + self.add_edge(node2, user, weight) + self.remove_edge(node1, user) + self.remove_node(node1) + + # + # Node accessor + # + def nodes_topological_order(self): + """ + Get the nodes in the unique lexicographical topological order + It generates a unique ordering of nodes by first sorting topologically + and then additionally by sorting lexicographically. + + Although topological_sort alone also works, this generates a unique key + for each epilogue visitor pattern and ensures the compilation cache can be reused. + :return: list[str] + """ + return list(nx.lexicographical_topological_sort(self._graph)) + + def node_metas_topological_order(self): + """ + Get the node metas in topological order + :return: list[NodeBase] + """ + return [self.get_node_meta(node) for node in self.nodes_topological_order()] + + @property + def nodes(self): + """ + Get all nodes + :return: list[str] + """ + return list(self._graph.nodes) + + @property + def nodes_meta(self): + """ + Get all node metas + :return: list[NodeBase] + """ + return [data[1]["meta"] for data in self._graph.nodes.data()] + + @property + def edges(self): + """ + Get all edges + :return: list[(str, str)] + """ + return list(self._graph.edges) + + # + # Path + # + def has_path(self, src: str, target: str) -> bool: + """ + Return True is a path exists from src to target + """ + return nx.has_path(self._graph, src, target) diff --git a/operators/cutlass/operators/fusion/ir/layout_algorithm.py b/operators/cutlass/operators/fusion/ir/layout_algorithm.py new file mode 100644 index 000000000..abb81707a --- /dev/null +++ b/operators/cutlass/operators/fusion/ir/layout_algorithm.py @@ -0,0 +1,362 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Layout algebras +""" + +from cutlass.operators.fusion.pycute import ( + Layout, + composition, + make_layout, + flatten, + product, +) + + +def _infer_split(old_shape, new_shape): + old_shape = _tuple_to_list(old_shape) + new_shape = _tuple_to_list(new_shape) + if len(old_shape) == 0 and len(new_shape) == 0: + return [] + if len(old_shape) == 0: + if product(tuple(new_shape)) != 1: + raise ValueError("Invalid reshape size") + else: + return new_shape + if len(new_shape) == 0: + if product(tuple(old_shape)) != 1: + raise ValueError("Invalid reshape size") + else: + return old_shape + # This is done recursively by only process the last dimension at each time + old_dim = old_shape[-1] + new_dim = new_shape[-1] + # Exact match + if old_dim == new_dim: + return _infer_split(old_shape[:-1], new_shape[:-1]) + [ + new_dim, + ] + # Needs split + if old_dim > new_dim and old_dim % new_dim == 0: + residual = old_dim // new_dim + return _infer_split( + old_shape[:-1] + [residual], + new_shape[:-1], + ) + [new_dim] + # Needs merge + if old_dim < new_dim and new_dim % old_dim == 0: + residual = new_dim // old_dim + return _infer_split( + old_shape[:-1], + new_shape[:-1] + [residual], + ) + [old_dim] + + raise NotImplementedError(f"Unsupported split: {old_shape} -> {new_shape}") + + +def _infer_merge(flatten_shape, shape): + flatten_shape = _tuple_to_list(flatten_shape) + shape = _tuple_to_list(shape) + idx_flat = 0 + merged_shape = [] + for dim in shape: + # Exact match + if dim == flatten_shape[idx_flat]: + merged_shape.append(dim) + idx_flat += 1 + # Need group + elif dim > flatten_shape[idx_flat] and dim % flatten_shape[idx_flat] == 0: + residual = dim + group = [] + while residual > 1: + group.append(flatten_shape[idx_flat]) + residual = residual // flatten_shape[idx_flat] + idx_flat += 1 + merged_shape.append(group) + else: + raise NotImplementedError(f"Unsupported merge: {flatten_shape} -> {shape}") + + return merged_shape + + +def _list_to_tuple(nested_list): + if isinstance(nested_list, list) or isinstance(nested_list, tuple): + return tuple(_list_to_tuple(item) for item in nested_list) + return nested_list + + +def _tuple_to_list(nested_tuple): + if isinstance(nested_tuple, list) or isinstance(nested_tuple, tuple): + return list(_tuple_to_list(item) for item in nested_tuple) + return nested_tuple + + +def _reverse_tuple(nested_tuple: tuple): + if isinstance(nested_tuple, tuple): + return tuple([_reverse_tuple(item) for item in nested_tuple][::-1]) + return nested_tuple + + +def _get_first_lhs_nonzero_stride(stride_list, idx): + for i in reversed(range(idx)): + if stride_list[i] != 0: + return i + else: + return None + + +def _get_first_rhs_nonzero_stride(stride_list, idx): + for i in range(idx + 1, len(stride_list)): + if stride_list[i] != 0: + return i + else: + return None + + +def reshape(layout, new_shape): + """ + General reshape of input layout. + It takes two steps: + 1. split the dimensions of the old layout + 2. merge the splitted dimensions according to the new shape + """ + # + # Step 1: Split the dimensions of the old layout + # + # 1.1 Flat old and new shape + old_flatten_shape = list(flatten(layout.shape)) + new_flatten_shape = list(flatten(new_shape)) + + # 1.2 Infer the flatten splitted shape + splitted_flatten_shape = _infer_split(old_flatten_shape, new_flatten_shape) + + # 1.3 Unflat the splitted shape based on the old shape + splited_shape = _infer_merge(splitted_flatten_shape, old_flatten_shape) + + # 1.4 Infer the type of each split + # If the split type is in row-major (R), the dimension list is reversed because + # the cute::composition only support column-major split + split_type = [] # the type of each split (ColumnMajor or RowMajor) + permuted_splitted_shape = [] + old_flatten_stride = list(flatten(layout.stride)) + for idx, dim in enumerate(splited_shape): + if not isinstance(dim, list): + permuted_splitted_shape.append(dim) + split_type.append("C") + else: + lhs_stride = _get_first_lhs_nonzero_stride(old_flatten_stride, idx) + rhs_stride = _get_first_rhs_nonzero_stride(old_flatten_stride, idx) + # Special case for single tuple + # Use column-major by default + if lhs_stride is None and rhs_stride is None: + permuted_splitted_shape.append(dim) + split_type.append("C") + else: + if lhs_stride is not None and rhs_stride is not None: + # We consider shape[idx]:stride[idx] + # Case 1: stride[idx - 1] <= stride[idx] <= stride[idx + 1]: column major + if ( + lhs_stride <= old_flatten_stride[idx] + and old_flatten_stride[idx] <= rhs_stride + ): + permuted_splitted_shape.append(dim) + split_type.append("C") + # Case 2: stride[idx - 1] > stride[idx] > stride[idx + 1]: row major + elif ( + lhs_stride > old_flatten_stride[idx] + and old_flatten_stride[idx] > rhs_stride + ): + permuted_splitted_shape.append([d for d in reversed(dim)]) + split_type.append("R") + # Case 3: stride[idx - 1] <= stride[idx] > stride[idx + 1]: concave + elif ( + lhs_stride <= old_flatten_stride[idx] + and old_flatten_stride[idx] > rhs_stride + ): + if lhs_stride >= rhs_stride: + permuted_splitted_shape.append(dim) + split_type.append("C") + else: + permuted_splitted_shape.append([d for d in reversed(dim)]) + split_type.append("R") + # Case 4: stride[idx - 1] > stride[idx] <= stride[idx + 1]: concave + elif ( + lhs_stride > old_flatten_stride[idx] + and old_flatten_stride[idx] <= rhs_stride + ): + if lhs_stride >= rhs_stride: + permuted_splitted_shape.append(dim) + split_type.append("C") + else: + permuted_splitted_shape.append([d for d in reversed(dim)]) + split_type.append("R") + else: + raise NotImplementedError() + elif lhs_stride is None: + # Case 1: dim's stride < dim+1's stride, expand in column major + if old_flatten_stride[idx] > rhs_stride: + permuted_splitted_shape.append([d for d in reversed(dim)]) + split_type.append("R") + else: + permuted_splitted_shape.append(dim) + split_type.append("C") + else: + # Case 1: dim's stride > dim-1's stride + if old_flatten_stride[idx] < lhs_stride: + permuted_splitted_shape.append([d for d in reversed(dim)]) + split_type.append("R") + else: + permuted_splitted_shape.append(dim) + split_type.append("C") + + # 1.4 Generate the splitted layout + permuted_splitted_layout = composition( + layout, Layout(_list_to_tuple(permuted_splitted_shape)) + ) + + # 1.5 Reverse the permutation in 1.4 before merge + splitted_shape = [] + splitted_stride = [] + for shape_dim, stride_dim, type in zip( + permuted_splitted_layout.shape, permuted_splitted_layout.stride, split_type + ): + if type == "C": + splitted_shape.append(shape_dim) + splitted_stride.append(stride_dim) + else: + splitted_shape.append(tuple([d for d in reversed(shape_dim)])) + splitted_stride.append(tuple([d for d in reversed(stride_dim)])) + splitted_layout = Layout(tuple(splitted_shape), tuple(splitted_stride)) + + # + # Step 2: Merge the splitted dimensions according to the new shape + # + # 2.1 Merge layout + merged_layout = composition(splitted_layout, Layout(new_shape)) + + # 2.2 Cleaning up + output_layout = composition(merged_layout, Layout(new_shape)) + return output_layout + + +def permutation(layout, permutation): + """ + Permute the layout + """ + new_shape = tuple([layout.shape[idx] for idx in permutation]) + new_stride = tuple([layout.stride[idx] for idx in permutation]) + return Layout(new_shape, new_stride) + + +def _broadcast(layout, new_shape): + if len(layout) == 1 and isinstance(new_shape, int): + old_dim = layout.shape + old_stride = layout.stride + new_dim = new_shape + if old_dim == new_dim: + return Layout(old_dim, old_stride) + elif old_dim == 1: + return Layout(new_dim, 0) + else: + raise NotImplementedError(f"Invalid Broadcast: {old_dim} -> {new_dim}") + + # Align the dimensions + old_shape = layout.shape + if isinstance(old_shape, int): + old_shape = (old_shape,) + sub_layouts = [ + layout, + ] + else: + sub_layouts = [sub_layout for sub_layout in layout] + rhs_broadcast_layouts = [Layout(1, 0)] * (len(new_shape) - len(old_shape)) + # Get the broadcasted layout + broadcast_layouts = [] + try: + layout = make_layout(*sub_layouts, *rhs_broadcast_layouts) + broadcast_layouts = [] + for idx, sub_layout in enumerate(layout): + broadcast_layouts.append(_broadcast(sub_layout, new_shape[idx])) + except NotImplementedError: + layout = make_layout(*rhs_broadcast_layouts, *sub_layouts) + for idx, sub_layout in enumerate(layout): + broadcast_layouts.append(_broadcast(sub_layout, new_shape[idx])) + return make_layout(*broadcast_layouts) + + +def broadcast(layout, new_shape): + """ + Broadcast the new layout based on the input shape + The broadcasted shape equals to the new shape + The stride of broadcasted dimensions are 0 + """ + return _broadcast(layout, new_shape) + + +def debroadcast(layout, dims): + """ + Squeeze the 0-stride + """ + for dim in dims: + if layout.stride[dim] != 0: + raise ValueError( + f"Dim{dim} cannot be debroadcasted as it has stride {layout.stride[dim]}" + ) + new_shape = tuple([s for idx, s in enumerate(layout.shape) if idx not in dims]) + new_stride = tuple([s for idx, s in enumerate(layout.stride) if idx not in dims]) + return Layout(new_shape, new_stride) + + +def canonicalization_(shapes, strides): + if isinstance(shapes, tuple): + c_shapes = [] + c_strides = [] + for shape, stride in zip(shapes, strides): + c_shape, c_stride = canonicalization_(shape, stride) + c_shapes.append(c_shape) + c_strides.append(c_stride) + return tuple(c_shapes), tuple(c_strides) + else: + if shapes == 1: + return 1, 0 + else: + return shapes, strides + + +def canonicalization(layout): + """ + Canonicalize the input layout + 1. set the stride of shape "1" to 0 + """ + new_shape, new_stride = canonicalization_(layout.shape, layout.stride) + return Layout(new_shape, new_stride) diff --git a/operators/cutlass/operators/fusion/ir/layout_nodes.py b/operators/cutlass/operators/fusion/ir/layout_nodes.py new file mode 100644 index 000000000..2d9d16eef --- /dev/null +++ b/operators/cutlass/operators/fusion/ir/layout_nodes.py @@ -0,0 +1,351 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Layout manipulation nodes and implementations + +The layout Nodes change the layout of intermediate nodes in epilogue visitor graph +""" + +from copy import deepcopy + +from cutlass.operators.utils.layout import LayoutType +from cutlass.operators.fusion.ir.layout_algorithm import _list_to_tuple, _tuple_to_list +from cutlass.operators.fusion.ir.node import NodeBase +from cutlass.operators.fusion.ir.tensor import Tensor +from cutlass.operators.fusion.pycute import product, flatten + + +class PermutationImpl: + """ + Detailed implementation and helper functions for permutation + """ + + def __init__(self, node) -> None: + assert "indices" in node.kwargs.keys() + self.indices = list(node.kwargs["indices"]) + self.inverse_indices = self.get_inverse_indices(self.indices) + + def get_inverse_impl(self): + inverse_impl = deepcopy(self) + inverse_impl.indices = self.inverse_indices + inverse_impl.inverse_indices = self.indices + return inverse_impl + + def update(self, shape): + num_dim = len(shape) + indices = self.indices + num_old_dim = len(indices) + # Add offset + for i, idx in enumerate(indices): + indices[i] = idx + num_dim - num_old_dim + # Add broadcast dims + for i in range(num_dim - num_old_dim): + indices = [i] + indices + + self.indices = indices + self.inverse_indices = self.get_inverse_indices(self.indices) + + def get_inverse_indices(self, indices): + """ + Get the indices for inverse permutation + """ + num_dim = len(indices) + inverse_indices = [0] * num_dim + for i in range(num_dim): + inverse_indices[indices[i]] = i + return inverse_indices + + def shape_propagation(self, input_node_meta): + input_shape = input_node_meta.tensor.shape + output_shape = tuple([input_shape[idx] for idx in self.indices]) + return output_shape + + def broadcast(self, shape, node_meta: NodeBase): + """ + Broadcast the inputs based on current shape + """ + self.update(shape) + inverse_shape = tuple([shape[idx] for idx in self.inverse_indices]) + node_meta.tensor.broadcast(inverse_shape) + + def apply_to_user(self, usr_meta: NodeBase): + """ + Propagate the permutation to the users of the current nodes + """ + usr_meta.tensor.permute(self.inverse_indices) + if hasattr(usr_meta, "store_tensor"): + if usr_meta.store_tensor is not None: + usr_meta.store_tensor.permute(self.inverse_indices) + + def apply_to_input(self, input_meta: NodeBase): + """ + Propagate the permutation to inputs of the current nodes + """ + input_meta.tensor.permute(self.indices) + if hasattr(input_meta, "store_tensor"): + if input_meta.store_tensor is not None: + input_meta.store_tensor.permute(self.indices) + + +class ReshapeImpl: + """ + Detailed implementation and helper functions for reshape + """ + + def __init__(self, node) -> None: + self.node = node + assert "new_shape" in node.kwargs.keys() + self.output_shape = _list_to_tuple(node.kwargs["new_shape"]) + + def get_inverse_impl(self): + inverse_impl = deepcopy(self) + inverse_impl.output_shape = self.input_shape + inverse_impl.input_shape = self.output_shape + return inverse_impl + + def shape_propagation(self, input_node_meta): + self.input_shape = input_node_meta.tensor.shape + return _list_to_tuple(self.output_shape) + + def broadcast(self, shape, node_meta: NodeBase): + """ + Broadcast the inputs based on current shape. + """ + # Step 1: infer split + flatten_split_shape = self.infer_split( + flatten(self.input_shape), flatten(self.output_shape) + ) + split_input_shape = self.infer_merge(flatten_split_shape, self.input_shape) + split_output_shape = self.infer_merge(flatten_split_shape, self.output_shape) + + # broadcast shape -> split_output_shape -> flatten_split_shape + if len(shape) - len(split_output_shape) > 0: + for _ in range(len(shape) - len(split_output_shape)): + split_output_shape = [1] + split_output_shape + flatten_split_shape = [1] + flatten_split_shape + split_input_shape = [1] + split_input_shape + broadcast_factor = [] + for dim, old_dim in zip(shape, split_output_shape): + if not isinstance(dim, list): + dim = [dim] + if not isinstance(old_dim, list): + old_dim = [old_dim] + if product(tuple(dim)) == product(tuple(old_dim)): + broadcast_factor += [1] * len(old_dim) + elif product(tuple(old_dim)) == 1: + assert len(dim) == 1 + broadcast_factor.append(dim[0]) + else: + raise NotImplementedError(f"Invalid Broadcast: {old_dim} -> {dim}") + + # flatten_split_shape -> split_input_shape + factor_idx = 0 + broadcast_split_input_shape = [] + for dim in split_input_shape: + if isinstance(dim, list): + new_dim = [] + for d in dim: + new_dim.append(d * broadcast_factor[factor_idx]) + factor_idx += 1 + broadcast_split_input_shape.append(new_dim) + else: + broadcast_split_input_shape.append(dim * broadcast_factor[factor_idx]) + factor_idx += 1 + broadcast_split_input_shape = _list_to_tuple(broadcast_split_input_shape) + node_meta.tensor.reshape(_list_to_tuple(split_input_shape)) + node_meta.tensor.broadcast(broadcast_split_input_shape) + # Last reshape op to clean up + broadcast_input_shape = tuple( + [product(dim) for dim in broadcast_split_input_shape] + ) + node_meta.tensor.reshape(broadcast_input_shape) + # Update the input shape and output shape + self.input_shape = _list_to_tuple(node_meta.tensor.shape) + self.output_shape = _list_to_tuple(shape) + + def apply_to_user(self, user_meta: NodeBase): + """ + Propagate the reshape to user nodes + """ + user_meta.tensor.reshape(tuple(self.input_shape)) + if hasattr(user_meta, "store_tensor"): + if user_meta.store_tensor is not None: + user_meta.store_tensor.reshape(tuple(self.input_shape)) + + def apply_to_input(self, input_meta: NodeBase): + """ + Propagate the reshape to input nodes + """ + input_meta.tensor.reshape(tuple(self.output_shape)) + if hasattr(input_meta, "store_tensor"): + if input_meta.store_tensor is not None: + input_meta.store_tensor.reshape(tuple(self.output_shape)) + + # + # Helper functions + # + + def infer_split(self, input_shape, output_shape): + """ + Infer the flatten splitted shape that can be merged to both input_shape and output_shape + """ + input_shape = _tuple_to_list(input_shape) + output_shape = _tuple_to_list(output_shape) + if len(input_shape) == 0 and len(output_shape) == 0: + return [] + if len(input_shape) == 0: + if product(tuple(output_shape)) != 1: + raise ValueError("Invalid reshape size") + else: + return output_shape + if len(output_shape) == 0: + if product(tuple(input_shape)) != 1: + raise ValueError("Invalid reshape size") + else: + return input_shape + # This is done recursively by only process the last dimension at each time + old_dim = input_shape[-1] + new_dim = output_shape[-1] + # Exact match + if old_dim == new_dim: + return self.infer_split(input_shape[:-1], output_shape[:-1]) + [new_dim] + # Needs split + if old_dim > new_dim and old_dim % new_dim == 0: + residual = old_dim // new_dim + return self.infer_split( + input_shape[:-1] + [residual], + output_shape[:-1], + ) + [new_dim] + # Needs merge + if old_dim < new_dim and new_dim % old_dim == 0: + residual = new_dim // old_dim + return self.infer_split( + input_shape[:-1], + output_shape[:-1] + [residual], + ) + [old_dim] + + raise NotImplementedError(f"Unsupported split: {input_shape} -> {output_shape}") + + def infer_merge(self, flatten_shape, shape): + flatten_shape = _tuple_to_list(flatten_shape) + shape = _tuple_to_list(shape) + idx_flat = len(flatten_shape) - 1 + merged_shape = [] + for dim in reversed(shape): + # Exact match + if dim == flatten_shape[idx_flat]: + merged_shape.append(dim) + idx_flat -= 1 + # need group + elif dim > flatten_shape[idx_flat] and dim % flatten_shape[idx_flat] == 0: + residual = dim + group = [] + while residual > 1: + group.append(flatten_shape[idx_flat]) + residual = residual // flatten_shape[idx_flat] + idx_flat -= 1 + merged_shape.append(group[::-1]) + else: + raise NotImplementedError( + f"Unsupported merge: {flatten_shape} -> {shape}" + ) + + return merged_shape[::-1] + + +class LayoutNode(NodeBase): + """ + Layout manipulation nodes + """ + + fn_to_impl = { + "permute": PermutationImpl, + "reshape": ReshapeImpl, + } + + def __init__(self, name: str, fn, kwargs: dict) -> None: + super().__init__(name) + self.op = "layout" + self.fn = fn + self.kwargs = kwargs + self.underlying_impl = self.fn_to_impl[self.fn.__name__](self) + + def get_inverse_node(self): + inverse_node = deepcopy(self) + inverse_node.underlying_impl = self.underlying_impl.get_inverse_impl() + return inverse_node + + def shape_propagation(self, input_node_metas): + if self._tensor is not None: + return + assert len(input_node_metas) == 1, "Layout node can only have one input node" + + output_shape = self.underlying_impl.shape_propagation(input_node_metas[0]) + + self._tensor = Tensor( + element=self.element_output, + shape=output_shape, + layout_tag=LayoutType.RowMajor, + ) + + return super().shape_propagation(input_node_metas) + + def type_propagation(self, input_node_metas: "list[NodeBase]"): + """ + The store nodes has element_output = element_input + """ + assert len(input_node_metas) == 1, "Layout node can only have one input node" + self.element_output = input_node_metas[0].element_output + + def broadcast_propagation(self, input_node_metas: "list[NodeBase]"): + """ + Propagate the broadcast in the reversed topological order + """ + if self.tensor is None: + raise RuntimeError(f"The tensor of node {self.name} is unknown.") + shape = self.tensor.shape + + for child in input_node_metas: + self.underlying_impl.broadcast(shape, child) + + def apply_to_user(self, usr_meta: NodeBase): + """ + Propagate the permutation to user nodes + """ + self.underlying_impl.apply_to_user(usr_meta) + + def apply_to_input(self, input_meta: NodeBase): + """ + Propagate the permutation to input nodes + """ + self.underlying_impl.apply_to_input(input_meta) diff --git a/operators/cutlass/operators/fusion/ir/load_nodes.py b/operators/cutlass/operators/fusion/ir/load_nodes.py new file mode 100644 index 000000000..0f638389c --- /dev/null +++ b/operators/cutlass/operators/fusion/ir/load_nodes.py @@ -0,0 +1,312 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Load nodes and implementations +""" + +import ctypes + +from cutlass.operators.fusion.ir.c_types import dtype2ctype, to_ctype_value, tuple_factory +from cutlass.operators.fusion.ir.node import NodeBase, ImplBase + + +class LoadImplBase(ImplBase): + """ + Base class for load node implementations + """ + + reserved_names = ["accum", "C"] + + def __init__(self, node) -> None: + super().__init__(node) + self.element = node.element + self.element_output = node.element_output + self.stride = node.tensor.stride + + +class AccumulatorImpl(LoadImplBase): + """ + Accumulator node implementation + """ + + @staticmethod + def match(node, problem_size: tuple): + return node.name == "accum" and node.tensor.shape == problem_size + + +class LoadSrcImpl(LoadImplBase): + """ + Load C implementation + """ + + @property + def name_camel(self) -> str: + return "TensorC" + + @property + def argument_type_c(self): + stride_mnl = self.get_stride_mnl() + tuple_type = tuple_factory(stride_mnl, self.stride_dtype) + + class _Argument(ctypes.Structure): + _fields_ = [("ptr_C", ctypes.c_void_p), ("stride_C", tuple_type)] + + def __init__(self, ptr) -> None: + self.ptr_C = ptr + self.stride_C = tuple_type(stride_mnl) + + return _Argument + + @staticmethod + def match(node, problem_size: tuple): + return node.name == "C" and node.tensor.shape == problem_size + + +class AuxLoadImpl(LoadImplBase): + """ + Load arbitrary tensor + """ + + @property + def argument_type(self): + stride_mnl = self.get_stride_mnl() + name = self.name + tuple_type = tuple_factory(stride_mnl, self.stride_dtype) + element_type = self.element + + class _Argument(ctypes.Structure): + _fields_ = [ + ("ptr_aux", ctypes.c_void_p), + ("null_default", dtype2ctype[element_type]), + ("dAux", tuple_type), + ] + + def __init__(self, kwargs) -> None: + ptr = kwargs[name] + self.ptr_aux = ptr + self.null_default = to_ctype_value(0, element_type) + self.dAux = tuple_type(stride_mnl) + + return _Argument + + @staticmethod + def match(node, problem_size: tuple): + if node.name in LoadImplBase.reserved_names: + return False + strideMN = node.tensor.stride[-2:] + if 1 not in strideMN or 0 in strideMN: + return False + return True + + +class RowBroadcastImpl(LoadImplBase): + """ + Broadcast a row vector + """ + + def __init__(self, node) -> None: + super().__init__(node) + self.stride_dtype = "int" + + @property + def argument_type(self): + stride_mnl = self.get_stride_mnl() + name = self.name + tuple_type = tuple_factory(stride_mnl, self.stride_dtype) + element_type = self.element + + class _Argument(ctypes.Structure): + _fields_ = [ + ("ptr_row", ctypes.c_void_p), + ("null_default", dtype2ctype[element_type]), + ("dRow", tuple_type), + ] + + def __init__(self, kwargs) -> None: + ptr = kwargs[name] + self.ptr_row = ptr + self.null_default = to_ctype_value(0, element_type) + self.dRow = tuple_type(stride_mnl) + + return _Argument + + @staticmethod + def match(node, problem_size: tuple): + if node.name in LoadImplBase.reserved_names: + return False + + strideMN = node.tensor.stride[-2:] + if strideMN == (0, 1): + return True + else: + return False + + +class ColumnBroadcastImpl(LoadImplBase): + """ + Broadcast a column vector + """ + + def __init__(self, node) -> None: + super().__init__(node) + self.stride_dtype = "int" + + @property + def argument_type(self): + stride_mnl = self.get_stride_mnl() + name = self.name + tuple_type = tuple_factory(stride_mnl, self.stride_dtype) + element_type = self.element + + class _Argument(ctypes.Structure): + _fields_ = [ + ("ptr_col", ctypes.c_void_p), + ("null_default", dtype2ctype[element_type]), + ("dCol", tuple_type), + ] + + def __init__(self, kwargs) -> None: + ptr = kwargs[name] + self.ptr_col = int(ptr) + self.null_default = to_ctype_value(0, element_type) + self.dCol = tuple_type(stride_mnl) + + return _Argument + + @staticmethod + def match(node, problem_size: tuple): + if node.name in LoadImplBase.reserved_names: + return False + + strideMN = node.tensor.stride[-2:] + if strideMN == (1, 0): + return True + else: + return False + + +class ScalarBroadcastImpl(LoadImplBase): + """ + Broadcast a scalar + """ + + def __init__(self, node) -> None: + super().__init__(node) + self.stride_dtype = "int" + + @property + def argument_type(self): + stride_mnl = self.get_stride_mnl() + name = self.name + tuple_type = tuple_factory(stride_mnl, self.stride_dtype) + element_type = self.element + + if self.tensor.is_constant: + value = self.tensor.value + + class _Argument(ctypes.Structure): + _fields_ = [ + ("scalars", dtype2ctype[element_type]), + ("scalar_ptrs", ctypes.c_void_p), + ("dScalar", tuple_type), + ] + + def __init__(self, kwargs) -> None: + self.scalars = to_ctype_value(value, element_type) + self.scalar_ptrs = 0 + self.dScalar = tuple_type(stride_mnl) + + else: + + class _Argument(ctypes.Structure): + _fields_ = [ + ("scalars", dtype2ctype[element_type]), + ("scalar_ptrs", ctypes.c_void_p), + ("dScalar", tuple_type), + ] + + def __init__(self, kwargs) -> None: + scalar_or_ptr = kwargs[name] + if isinstance(scalar_or_ptr, float): + self.scalars = to_ctype_value(scalar_or_ptr, element_type) + self.scalar_ptrs = 0 + else: + self.scalar_ptrs = int(scalar_or_ptr) + + self.dScalar = tuple_type(stride_mnl) + + return _Argument + + @staticmethod + def match(node, problem_size: tuple): + if node.name in LoadImplBase.reserved_names: + return False + + strideMN = node.tensor.stride[-2:] + if strideMN == (0, 0): + return True + else: + return False + + +class LoadNode(NodeBase): + """ + Load Node + """ + + cnt = 0 + possible_impls = [ + AccumulatorImpl, + LoadSrcImpl, + AuxLoadImpl, + RowBroadcastImpl, + ColumnBroadcastImpl, + ScalarBroadcastImpl, + ] + + def __init__(self, name: str) -> None: + if name is None: + name = f"load{LoadNode.cnt}" + LoadNode.cnt += 1 + super().__init__(name) + self.op = "load" + + def type_propagation(self, *args, **kwargs): + """ + Load node loads tensor under type `tensor.element` and returns an array of type `tensor.element`. + """ + if self.tensor is None: + raise RuntimeError(f"The tensor of node {self.name} is unknown.") + + self.element = self.tensor.element + self.element_output = self.tensor.element diff --git a/operators/cutlass/operators/fusion/ir/node.py b/operators/cutlass/operators/fusion/ir/node.py new file mode 100644 index 000000000..31af1004d --- /dev/null +++ b/operators/cutlass/operators/fusion/ir/node.py @@ -0,0 +1,330 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Base & visitor classes of DAGIR Nodes +""" + +import ctypes +from re import sub + +from cutlass.operators.utils.layout import LayoutType +from cutlass.operators.fusion.ir.layout_algorithm import _list_to_tuple, _reverse_tuple +from cutlass.operators.fusion.ir.tensor import Tensor + + +class TupleEmitter: + """ + Emit the cute tuple to C++ code + """ + + def __init__(self, stride_dtype): + self.stride_dtype = stride_dtype + + def emit(self, py_tuple): + if isinstance(py_tuple, int): + if py_tuple in [0, 1]: + return f"cute::Int<{py_tuple}>" + else: + return f"{self.stride_dtype}" + elif isinstance(py_tuple, tuple): + decl = "cute::Stride<" + for item in py_tuple: + decl += self.emit(item) + ", " + return decl[:-2] + ">" + else: + raise ValueError( + f"TupleEmitter.emit only accepts tuple or int, got {type(py_tuple).__name__}" + ) + + +class ImplBase: + """ + Base class for Node Implementation + """ + + def __init__(self, node) -> None: + self.node = node + self.name = node.name + self.tensor = node.tensor + self._type_decl = None + self.tuple_emitter = TupleEmitter("int64_t") + + @property + def stride_dtype(self): + return self.tuple_emitter.stride_dtype + + @stride_dtype.setter + def stride_dtype(self, stride_dtype): + self.tuple_emitter.stride_dtype = stride_dtype + + @staticmethod + def match(node, problem_size: tuple): + """ + Match function used in get_underlying_impl + """ + raise NotImplementedError("The `match` function is not defined.") + + @property + def argument_type(self): + """ + Default class for Argument Type + """ + + class _Argument(ctypes.Structure): + _fields_ = [] + + def __init__(self, *args, **kwargs) -> None: + pass + + return _Argument + + @property + def name_camel(self) -> str: + """ + Return the CamelCase name. + """ + return sub(r"(_|-)+", " ", self.name).title().replace(" ", "") + + @property + def stride_mnl(self): + """ + Typename StrideMNL + """ + stride = _list_to_tuple( + [self.stride[-2], self.stride[-1]] + + list(_reverse_tuple(tuple(self.stride[:-2]))) + ) + return self.tuple_emitter.emit(stride) + + def get_non_constant_stride(self, py_tuple): + if isinstance(py_tuple, int): + if py_tuple not in [0, 1]: + return py_tuple + else: + return None + non_constant_stride = [] + for item in py_tuple: + item_out = self.get_non_constant_stride(item) + if item_out: + non_constant_stride.append(item_out) + return tuple(non_constant_stride) + + def get_stride_mnl(self): + """ + Get the non-zero stride mnl. This is used in argument construction + """ + stride = _list_to_tuple( + [self.stride[-2], self.stride[-1]] + + list(_reverse_tuple(tuple(self.stride[:-2]))) + ) + return stride + + def get_smem_size(self, *args, **kwargs): + """ + Get the shared memory size and alignment of current node + """ + return (0, 1) + + +class NoOpImpl(ImplBase): + """ + The NoOpImpl does nothing but forward its input to users + """ + + def __init__(self, node) -> None: + super().__init__(node) + + @staticmethod + def match(node, problem_size: tuple): + if node.op == "store": + # Store that is not output is a No OP + return not node.is_output + + +class NodeBase: + """ + Base class of DAG Node + """ + + def __init__(self, name: str) -> None: + self.name = name + self.underlying_impl = None + + self._tensor = None + + # Whether the node is disabled for emit + self.disabled = False + + @property + def name_camel(self) -> str: + """ + Return the CamelCase name. + """ + return self.underlying_impl.name_camel + + @property + def tensor(self) -> Tensor: + """ + Return the output tensor (concept: cutlass.operators.fusion.ir.tensor) + """ + return self._tensor + + @tensor.setter + def tensor(self, kwargs): + """ + Setting the tensor + """ + self._tensor = Tensor(**kwargs) + + # + # Helper functions for type/shape propagation + # + + def shape_propagation(self, input_node_metas): + """ + Infer shape from input nodes + General Broadcasting Rules from NumPy + When operating on two arrays, we compare their shapes element-wise. + It starts with the trailing (i.e. rightmost) dimension and works its + way left. Two dimensions are compatible when + 1. they are equal + 2. one of them is 1 + """ + if self._tensor is not None: + return + + shape = None + for src in input_node_metas: + src_shape = src.tensor.shape + if shape is None: + shape = src_shape + else: + len_difference = len(shape) - len(src_shape) + if len_difference > 0: + for _ in range(len_difference): + src_shape = [1] + list(src_shape) + elif len_difference < 0: + for _ in range(-len_difference): + shape = [1] + list(shape) + broadcasted_shape = [] + # Infer broadcast shape + for shape_dim, src_dim in zip(reversed(shape), reversed(src_shape)): + if shape_dim == 1: + broadcasted_shape = [src_dim] + list(broadcasted_shape) + elif src_dim == 1: + broadcasted_shape = [shape_dim] + list(broadcasted_shape) + elif shape_dim == src_dim: + broadcasted_shape = [shape_dim] + list(broadcasted_shape) + else: + error_msg = "Dimension mismatch between " + for src_ in input_node_metas: + error_msg += f"{src_.name}{src_.tensor.shape}, " + error_msg = error_msg[:-2] + "." + raise RuntimeError(error_msg) + shape = tuple(broadcasted_shape) + + self._tensor = Tensor( + element=self.element_output, shape=shape, layout_tag=LayoutType.RowMajor + ) + + def type_propagation(self, *args, **kwargs): + """ + Each node is associated with two data types: `element` and `element_output`. + The `element_output` is the type of return array of the node. The `element` + has specific meaning for different node types. + * Load Node: data type of tensor in gmem + * Compute Node: element compute + * Store Node: data type of tensor in gmem + This function must be overloaded in the derived classes + """ + raise NotImplementedError( + f"Function `type_propagation` is not overloaded in {self.__class__.__name__}" + ) + + def broadcast_propagation(self, input_node_metas: "list[NodeBase]"): + """ + Propagate the broadcast in the reversed topological order. + For example: + C[L, M, N] = A[M, 1] + B[L, M, N] + After the broadcast propagation, it will be come + C[L, M, N] = A[L, M, N] + B[L, M, N] + and each tensor will have a proper stride accessing the underlying tensor + """ + if self.tensor is None: + raise RuntimeError(f"The tensor of node {self.name} is unknown.") + for child in input_node_metas: + child.tensor.broadcast(self.tensor.shape) + + def get_underlying_impl(self, problem_size: tuple): + """ + Get the underlying implementation of the current node. + """ + if self.tensor is None: + raise RuntimeError( + f"The Layout of node {self.name} is unknown. Please call PassShapeTypePropagation first." + ) + + for impl in self.possible_impls: + if impl.match(self, problem_size): + self.underlying_impl = impl(self) + break + + if self.underlying_impl is None: + raise NotImplementedError( + f"No matching op for node {self.name} with stride {self.tensor.stride}." + ) + + +# +# Visitor Nodes & Impls +# + + +class TopoVisitorImpl(ImplBase): + """ + Impl for topological visitor + """ + + def __init__(self, node) -> None: + super().__init__(node.output_node) + self.name = node.name + self.element_output = node.output_node.element_output + + +class TopoVisitorNode(NodeBase): + def __init__(self, name: str, subgraph, output_node) -> None: + super().__init__(name) + self.subgraph = subgraph + self.output_node = output_node + self.op = "dag" + self.underlying_impl = TopoVisitorImpl(self) diff --git a/operators/cutlass/operators/fusion/ir/store_nodes.py b/operators/cutlass/operators/fusion/ir/store_nodes.py new file mode 100644 index 000000000..a900e1477 --- /dev/null +++ b/operators/cutlass/operators/fusion/ir/store_nodes.py @@ -0,0 +1,276 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Store node and implementations +""" + +import ctypes + +from cutlass.operators.fusion.library import DataType, FloatRoundStyle, FunctionalOp +from cutlass.operators.fusion.ir.c_types import dtype2ctype, to_ctype_value, tuple_factory +from cutlass.operators.fusion.ir.node import NodeBase, ImplBase, NoOpImpl +from cutlass.operators.fusion.ir.tensor import Tensor + + +class StoreImplBase(ImplBase): + """ + Base class for store node implementation + """ + + reserved_names = ["D"] + + def __init__(self, node) -> None: + super().__init__(node) + self.element = node.element + self.element_output = node.element_output + self.stride = node.store_tensor.stride + + +class StoreDImpl(StoreImplBase): + """ + Store D implementation + """ + + @property + def argument_type_d(self): + stride_mnl = self.get_stride_mnl() + tuple_type = tuple_factory(stride_mnl, self.stride_dtype) + + class _Argument(ctypes.Structure): + _fields_ = [("ptr_D", ctypes.c_void_p), ("stride_D", tuple_type)] + + def __init__(self, ptr: int) -> None: + self.ptr_D = ptr + self.stride_D = tuple_type(stride_mnl) + + return _Argument + + @staticmethod + def match(node, problem_size: tuple): + if node.name == "D" and node.store_tensor.shape == problem_size: + return True + return False + + +class AuxStoreImpl(StoreImplBase): + def __init__(self, node) -> None: + super().__init__(node) + self.round_style = FloatRoundStyle.ToNearest + + @property + def argument_type(self): + stride_mnl = self.get_stride_mnl() + name = self.name + tuple_type = tuple_factory(stride_mnl, self.stride_dtype) + + class _Argument(ctypes.Structure): + _fields_ = [("ptr_aux", ctypes.c_void_p), ("dAux", tuple_type)] + + def __init__(self, kwargs) -> None: + ptr = kwargs[name] + self.ptr_aux = ptr + self.dAux = tuple_type(stride_mnl) + + return _Argument + + @staticmethod + def match(node, problem_size: tuple): + if not node.is_output: + return False + if node.name in StoreImplBase.reserved_names: + return False + + strideMN = node.store_tensor.stride[-2:] + if 1 not in strideMN or 0 in strideMN: + return False + return True + + +class ReductionImplBase(StoreImplBase): + def __init__(self, node) -> None: + super().__init__(node) + self.element = node.store_tensor.element + self.element_compute = node.element_compute + self.reg_reduce_fn = self.node.reg_reduce_fn + self.gmem_reduce_fn = self.node.gmem_reduce_fn + self.round_style = node.round_style + self.stride_dtype = "int" + + def get_reduce_identity(self): + """ + Return the reduction identity of the current reduce_fn + """ + maxes = { + DataType.f32: (2**31) - 1, + DataType.f16: (2**15), + DataType.s32: (2**31) - 1, + DataType.s8: (2**7) - 1, + } + mins = { + DataType.f32: -maxes[DataType.f32], + DataType.f16: -maxes[DataType.f16], + DataType.s32: -maxes[DataType.s32], + DataType.s8: -maxes[DataType.s8], + } + if self.reg_reduce_fn == FunctionalOp.Maximum: + if self.element_compute not in mins: + raise Exception(f"No min entry for data type {self.element_compute}") + return to_ctype_value(mins[self.element_compute], self.element_compute) + elif self.reg_reduce_fn == FunctionalOp.Multiplies: + return to_ctype_value(1.0, self.element_compute) + elif self.reg_reduce_fn == FunctionalOp.Minimum: + if self.element_compute not in maxes: + raise Exception(f"No max entry for data type {self.element_compute}") + return to_ctype_value(maxes[self.element_compute], self.element_compute) + else: + return to_ctype_value(0.0, self.element_compute) + + @property + def argument_type(self): + self.get_reduce_identity() + stride_mnl = self.get_stride_mnl() + name = self.name + tuple_type = tuple_factory(stride_mnl, self.stride_dtype) + element_compute = self.element_compute + reduce_identity = self.get_reduce_identity() + + class _Argument(ctypes.Structure): + _fields_ = [ + ("ptr", ctypes.c_void_p), + ("reduce_identity", dtype2ctype[element_compute]), + ("dMNL", tuple_type), + ] + + def __init__(self, kwargs) -> None: + ptr = kwargs[name] + self.ptr = ptr + self.reduce_identity = reduce_identity + self.dMNL = tuple_type(stride_mnl) + + return _Argument + + +class ColumnReductionImpl(ReductionImplBase): + @staticmethod + def match(node, problem_size: tuple): + if not node.is_output: + return False + if node.name in StoreImplBase.reserved_names: + return False + + strideMN = node.store_tensor.stride[-2:] + if strideMN == (1, 0): + return True + else: + return False + + +class RowReductionImpl(ReductionImplBase): + @staticmethod + def match(node, problem_size: tuple): + if not node.is_output: + return False + if node.name in StoreImplBase.reserved_names: + return False + + strideMN = node.store_tensor.stride[-2:] + if strideMN == (0, 1): + return True + else: + return False + + +class ScalarReductionImpl(ReductionImplBase): + @staticmethod + def match(node, problem_size: tuple): + if not node.is_output: + return False + if node.name in StoreImplBase.reserved_names: + return False + + strideMN = node.store_tensor.stride[-2:] + if strideMN == (0, 0): + return True + else: + return False + + +class StoreNode(NodeBase): + """ + Store node + """ + + possible_impls = [ + AuxStoreImpl, + RowReductionImpl, + ColumnReductionImpl, + ScalarReductionImpl, + NoOpImpl, + StoreDImpl, + ] + + def __init__(self, name: str) -> None: + super().__init__(name) + self.op = "store" + self.is_output = False + self._store_tensor = None + + @property + def store_tensor(self) -> Tensor: + """ + Return the output tensor (concept: cutlass.operators.fusion.ir.tensor) + """ + return self._store_tensor + + @store_tensor.setter + def store_tensor(self, kwargs): + """ + Setting the tensor + """ + self._store_tensor = Tensor(**kwargs) + + def type_propagation(self, input_node_metas: "list[NodeBase]"): + """ + The store nodes has element_output = element_input + """ + if self.is_output: + if self.store_tensor is None: + raise RuntimeError(f"The store tensor of node {self.name} is unknown.") + self.element = self.store_tensor.element + assert len(input_node_metas) == 1, "Store node can only have one input node" + self.element_output = input_node_metas[0].element_output + + def broadcast_propagation(self, input_node_metas: "list[NodeBase]"): + super().broadcast_propagation(input_node_metas) + if self.is_output: + self._store_tensor.broadcast(self.tensor.shape) diff --git a/operators/cutlass/operators/fusion/ir/tensor.py b/operators/cutlass/operators/fusion/ir/tensor.py new file mode 100644 index 000000000..a78afbe35 --- /dev/null +++ b/operators/cutlass/operators/fusion/ir/tensor.py @@ -0,0 +1,155 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +High-level class for tensor +""" + +from cutlass.operators.utils.layout import LayoutType +from cutlass.operators.fusion.library import ( + get_datatype_and_layout, + get_tensor_shape, + library_type, +) +from cutlass.operators.fusion.ir.layout_algorithm import ( + Layout, + broadcast, + canonicalization, + permutation, + reshape, + _reverse_tuple, +) + + +class Tensor: + """ + The tensor abstracts the data type + """ + + def __init__( + self, + tensor=None, + element=None, + shape=None, + stride=None, + layout_tag=None, + is_constant=False, + ) -> None: + if element is not None and tensor is not None: + raise Exception("Must not specify both element and tensor") + elif shape is not None and tensor is not None: + raise Exception("Must not specify both shape and tensor") + elif layout_tag is not None and tensor is not None: + raise Exception("Must not specify both layout_tag and tensor") + elif ( + element is None or (layout_tag is None and stride is None) or shape is None + ) and (tensor is None): + raise Exception( + "Must specify one of (element, shape, layout/stride) or (tensor)" + ) + elif stride is not None and tensor is not None: + raise Exception("Must not specify both stride and tensor") + elif stride is not None and layout_tag is not None: + raise Exception("Must not specify layout_tag when stride is provided") + + if isinstance(tensor, Tensor): + # Directly copy all the attributes + self.__dict__.update(vars(tensor)) + else: + if tensor is None: + self.element = library_type(element) + else: + self.element, layout_tag = get_datatype_and_layout(tensor) + shape = get_tensor_shape(tensor) + if stride is not None: + self.layout = Layout(shape[::-1], stride[::-1]) + else: + if layout_tag == LayoutType.RowMajor: + self.layout = Layout(shape[::-1]) + elif layout_tag == LayoutType.ColumnMajor: + self.layout = permutation( + Layout(shape), [idx for idx in reversed(range(len(shape)))] + ) + self.layout = canonicalization(self.layout) + + self.is_constant = is_constant + # Save the tensor value if it is constant + if is_constant and tensor is not None: + self.value = tensor + + @property + def shape(self): + """ + Returns the RowMajor layout shape + """ + return _reverse_tuple(self.layout.shape) + + @property + def stride(self): + """ + Returns the RowMajor layout stride + """ + return _reverse_tuple(self.layout.stride) + + @property + def rank(self): + """ + Returns the rank of the tensor + """ + return len(self.shape) + + # + # Layout Algorithms + # + + def broadcast(self, shape): + """ + Broadcast self.layout to shape + """ + assert isinstance(shape, tuple) + self.layout = broadcast(self.layout, _reverse_tuple(shape)) + + def reshape(self, shape): + """ + Reshape self.layout to shape + """ + assert isinstance(shape, tuple) + reverse_shape = _reverse_tuple(shape) + self.layout = reshape(self.layout, reverse_shape) + + def permute(self, indices): + """ + Permute self.layout according to indices + """ + length = len(indices) + indices = [length - idx - 1 for idx in indices] + self.layout = permutation(self.layout, indices[::-1]) diff --git a/operators/cutlass/operators/fusion/library.py b/operators/cutlass/operators/fusion/library.py new file mode 100644 index 000000000..87cd29121 --- /dev/null +++ b/operators/cutlass/operators/fusion/library.py @@ -0,0 +1,430 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Copies of enum types from cutlass_library that are used in the fusion frontend. +""" + +import enum +from enum import auto as enum_auto + +from cutlass.cute.typing import ( + Numeric, + Float32, + Float16, + BFloat16, + Int32, + Int8, + Uint8, + Uint16, + Uint32, + Uint64, + Int16, + Int64, + Float8E5M2, + Float8E4M3FN, +) + +from cutlass.operators.utils.common import is_torch_available, is_numpy_available +from cutlass.operators.utils.tensor import is_numpy_tensor, is_torch_tensor +from cutlass.operators.utils.layout import LayoutType + + +class DataType(enum.Enum): + void = enum_auto() # primarily used to disable C tensor for epilogues + b1 = enum_auto() + u2 = enum_auto() + u4 = enum_auto() + u8 = enum_auto() + u16 = enum_auto() + u32 = enum_auto() + u64 = enum_auto() + s2 = enum_auto() + s4 = enum_auto() + s8 = enum_auto() + s16 = enum_auto() + s32 = enum_auto() + s64 = enum_auto() + e4m3 = enum_auto() + e5m2 = enum_auto() + f8 = enum_auto() + f6 = enum_auto() + f4 = enum_auto() + e3m2 = enum_auto() + e2m3 = enum_auto() + e2m1 = enum_auto() + ue8m0 = enum_auto() + ue4m3 = enum_auto() + f16 = enum_auto() + bf16 = enum_auto() + f32 = enum_auto() + tf32 = enum_auto() + f64 = enum_auto() + cf16 = enum_auto() + cbf16 = enum_auto() + cf32 = enum_auto() + ctf32 = enum_auto() + cf64 = enum_auto() + cs2 = enum_auto() + cs4 = enum_auto() + cs8 = enum_auto() + cs16 = enum_auto() + cs32 = enum_auto() + cs64 = enum_auto() + cu2 = enum_auto() + cu4 = enum_auto() + cu8 = enum_auto() + cu16 = enum_auto() + cu32 = enum_auto() + cu64 = enum_auto() + invalid = enum_auto() + + +DataTypeSize = { + DataType.void: 0, + DataType.b1: 1, + DataType.u2: 2, + DataType.u4: 4, + DataType.u8: 8, + DataType.u16: 16, + DataType.u32: 32, + DataType.u64: 64, + DataType.s2: 2, + DataType.s4: 4, + DataType.s8: 8, + DataType.s16: 16, + DataType.s32: 32, + DataType.s64: 64, + DataType.e4m3: 8, + DataType.e5m2: 8, + DataType.f8: 8, + DataType.f6: 6, + DataType.f4: 4, + DataType.e2m3: 6, + DataType.e3m2: 6, + DataType.e2m1: 4, + DataType.ue8m0: 8, + DataType.ue4m3: 8, + DataType.f16: 16, + DataType.bf16: 16, + DataType.f32: 32, + DataType.tf32: 32, + DataType.f64: 64, + DataType.cf16: 32, + DataType.cbf16: 32, + DataType.cf32: 64, + DataType.ctf32: 32, + DataType.cf64: 128, + DataType.cu2: 4, + DataType.cu4: 8, + DataType.cu8: 16, + DataType.cu16: 32, + DataType.cu32: 64, + DataType.cu64: 128, + DataType.cs2: 4, + DataType.cs4: 8, + DataType.cs8: 16, + DataType.cs16: 32, + DataType.cs32: 64, + DataType.cs64: 128, +} + + +class EpilogueScheduleType(enum.Enum): + ScheduleAuto = enum_auto() + EpilogueTransposed = enum_auto() + NoSmemWarpSpecialized = enum_auto() + PtrArrayNoSmemWarpSpecialized = enum_auto() + NoSmemWarpSpecialized1Sm = enum_auto() + NoSmemWarpSpecialized2Sm = enum_auto() + FastF32NoSmemWarpSpecialized1Sm = enum_auto() + FastF32NoSmemWarpSpecialized2Sm = enum_auto() + BlockwiseNoSmemWarpSpecialized1Sm = enum_auto() + BlockwiseNoSmemWarpSpecialized2Sm = enum_auto() + PtrArrayNoSmemWarpSpecialized1Sm = enum_auto() + PtrArrayNoSmemWarpSpecialized2Sm = enum_auto() + PtrArrayFastF32NoSmemWarpSpecialized1Sm = enum_auto() + PtrArrayFastF32NoSmemWarpSpecialized2Sm = enum_auto() + PtrArrayBlockwiseNoSmemWarpSpecialized1Sm = enum_auto() + PtrArrayBlockwiseNoSmemWarpSpecialized2Sm = enum_auto() + TmaWarpSpecialized = enum_auto() + TmaWarpSpecializedCooperative = enum_auto() + TmaWarpSpecialized1Sm = enum_auto() + TmaWarpSpecialized2Sm = enum_auto() + PtrArrayTmaWarpSpecialized1Sm = enum_auto() + PtrArrayTmaWarpSpecialized2Sm = enum_auto() + PtrArrayTmaWarpSpecializedPingpong = enum_auto() + PtrArrayTmaWarpSpecializedCooperative = enum_auto() + TmaWarpSpecialized1SmNvf4 = enum_auto() + TmaWarpSpecialized2SmNvf4 = enum_auto() + TmaWarpSpecialized1SmMxf4 = enum_auto() + TmaWarpSpecialized2SmMxf4 = enum_auto() + TmaWarpSpecialized1SmMxf8f6f4 = enum_auto() + TmaWarpSpecialized2SmMxf8f6f4 = enum_auto() + SparseTmaWarpSpecializedCooperativeSm120 = enum_auto() + + +class ActivationOp(enum.Enum): + DGelu = enum_auto() + Gelu = enum_auto() + GeluTaylor = enum_auto() + HardSwish = enum_auto() + Identity = enum_auto() + LeakyReLU = enum_auto() + ReLU = enum_auto() + Sigmoid = enum_auto() + SiLU = enum_auto() + Tanh = enum_auto() + + +ActivationOpTag = { + ActivationOp.DGelu: "cutlass::epilogue::thread::dGELU", + ActivationOp.Gelu: "cutlass::epilogue::thread::GELU", + ActivationOp.GeluTaylor: "cutlass::epilogue::thread::GELU_taylor", + ActivationOp.HardSwish: "cutlass::epilogue::thread::HardSwish", + ActivationOp.Identity: "cutlass::epilogue::thread::Identity", + ActivationOp.LeakyReLU: "cutlass::epilogue::thread::LeakyReLU", + ActivationOp.ReLU: "cutlass::epilogue::thread::ReLu", + ActivationOp.Sigmoid: "cutlass::epilogue::thread::Sigmoid", + ActivationOp.SiLU: "cutlass::epilogue::thread::SiLu", + ActivationOp.Tanh: "cutlass::epilogue::thread::Tanh", +} + + +class FloatRoundStyle(enum.Enum): + ToNearest = enum_auto() + ToNearestSatfinite = enum_auto() + Indeterminate = enum_auto() + TowardZero = enum_auto() + TowardInfinity = enum_auto() + TowardNegInfinity = enum_auto() + HalfUlpTruncDntz = enum_auto() + HalfUlpTruncate = enum_auto() + + +class FunctionalOp(enum.Enum): + AtomicAdd = enum_auto() + AtomicMaximum = enum_auto() + Divides = enum_auto() + Maximum = enum_auto() + Minimum = enum_auto() + Minus = enum_auto() + Multiplies = enum_auto() + MultiplyAdd = enum_auto() + Plus = enum_auto() + Exp = enum_auto() + + +FunctionalOpTag = { + FunctionalOp.AtomicAdd: "cutlass::atomic_add", + FunctionalOp.AtomicMaximum: "cutlass::atomic_maximum", + FunctionalOp.Divides: "cutlass::divides", + FunctionalOp.Maximum: "cutlass::maximum", + FunctionalOp.Minimum: "cutlass::minimum", + FunctionalOp.Minus: "cutlass::minus", + FunctionalOp.Multiplies: "cutlass::multiplies", + FunctionalOp.MultiplyAdd: "cutlass::multiply_add", + FunctionalOp.Plus: "cutlass::plus", + FunctionalOp.Exp: "cutlass::fast_exp_op", +} + + +def op_tag(op) -> str: + """ + Dispatches `op` to the appropriate *Tag dictionary depending on whether + `op` is an ActivationOp or FunctionalOp. This is useful for cases in which + either type can be used. + + :param op: operation to emit a tag for + :type op: ActivationOp | FunctionalOp + + :return: tag corresponding to op + :rtype: str + """ + if isinstance(op, ActivationOp): + return ActivationOpTag[op] + elif isinstance(op, FunctionalOp): + return FunctionalOpTag[op] + else: + raise Exception( + f"Unexpected op type {op}. Must be one of ActivationOp or FunctionalOp." + ) + + +# The current EVT frontend also contains code needed for emitting C++ EVTs. +# C++ emission is not currently supported by cutlass.operators, but we still +# need many of the utilities surrounding it. These utilities make use of dictionaries +# from cutlass_library types to strings containing C++ code. We define placeholders +# for empty versions of these dictionaries that raise an error if they are used. +class _UnimplementedDict: + def __init__(self, name): + self.name = name + + def __getitem__(self, key): + raise NotImplementedError( + f"Dictionary {self.name} is not implemented. This code path should not have been reafched." + ) + + +DataTypeTag = _UnimplementedDict("DataTypeTag") +EpilogueScheduleTag = _UnimplementedDict("EpilogueScheduleTag") +FloatRoundStyleTag = _UnimplementedDict("FloatRoundStyleTag") +KernelScheduleSuffixes = _UnimplementedDict("KernelScheduleSuffixes") +OpcodeClassTag = _UnimplementedDict("OpcodeClassTag") + + +_torch_to_library_dict = None +if is_torch_available(): + import torch + + _torch_to_library_dict = { + torch.half: DataType.f16, + torch.float16: DataType.f16, + torch.bfloat16: DataType.bf16, + torch.float: DataType.f32, + torch.float32: DataType.f32, + } + + +def _tensor_from_numpy(np_tensor) -> tuple[DataType, LayoutType]: + dtype = library_type(np_tensor.dtype) + if np_tensor.flags.c_contiguous: + layout = LayoutType.RowMajor + elif np_tensor.flags.f_contiguous: + layout = LayoutType.ColumnMajor + return (dtype, layout) + + +def _tensor_from_torch(pt_tensor): + dtype = library_type(pt_tensor.dtype) + return (dtype, LayoutType.RowMajor) + + +def torch_library_type(inp) -> DataType: + if _torch_to_library_dict is None: + return None + return _torch_to_library_dict.get(inp, None) + + +def numpy_library_type(inp) -> DataType: + if is_numpy_available(): + import numpy as np + + if inp == np.float16: + return DataType.f16 + elif inp == np.float32: + return DataType.f32 + elif inp == np.float64: + return DataType.f64 + elif inp == np.int8: + return DataType.s8 + elif inp == np.int32: + return DataType.s32 + return None + + +def cutlass_library_type(inp: Numeric) -> DataType: + if inp == Float32: + return DataType.f32 + elif inp == Float16: + return DataType.f16 + elif inp == BFloat16: + return DataType.bf16 + elif inp == Int32: + return DataType.s32 + elif inp == Int8: + return DataType.s8 + elif inp == Uint8: + return DataType.u8 + elif inp == Uint16: + return DataType.u16 + elif inp == Uint32: + return DataType.u32 + elif inp == Uint64: + return DataType.u64 + elif inp == Int16: + return DataType.s16 + elif inp == Int64: + return DataType.s64 + elif inp == Float8E5M2: + return DataType.e5m2 + elif inp == Float8E4M3FN: + return DataType.e4m3 + else: + return None + + +def library_type(inp): + if inp in DataTypeSize: + return inp + + for cvt_fn in [ + numpy_library_type, + torch_library_type, + cutlass_library_type, + ]: + out = cvt_fn(inp) + if out is not None: + return out + + raise ValueError(f"No available conversion from type {inp} to a library type.") + + +def get_datatype_and_layout(tensor) -> tuple[DataType, LayoutType]: + if is_numpy_tensor(tensor): + return _tensor_from_numpy(tensor) + elif is_torch_tensor(tensor): + return _tensor_from_torch(tensor) + elif isinstance(tensor, float) or isinstance(tensor, int): + return (DataType.f32, LayoutType.RowMajor) + else: + raise Exception( + f"Unable to convert tensor of type {type(tensor)} to Python-bound CUTLASS datatype and layout." + ) + + +def get_tensor_shape(tensor, op="GEMM") -> tuple: + if is_numpy_tensor(tensor): + return tensor.shape + elif is_torch_tensor(tensor): + size = tensor.size() + if op == "CONV": + # PyTorch Tensors have shape NCHW + return (size[0], size[2], size[3], size[1]) + else: + return tuple(tensor.size()) + elif isinstance(tensor, float) or isinstance(tensor, int): + return (1,) + else: + raise Exception( + f"Unable to convert tensor of type {type(tensor)} to Python-bound CUTLASS datatype and layout." + ) diff --git a/operators/cutlass/operators/fusion/passes/__init__.py b/operators/cutlass/operators/fusion/passes/__init__.py new file mode 100644 index 000000000..7c21823d3 --- /dev/null +++ b/operators/cutlass/operators/fusion/passes/__init__.py @@ -0,0 +1,59 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +from cutlass.operators.fusion.passes.graph_drawer import EVTGraphDrawer +from cutlass.operators.fusion.passes.pass_argument_type import PassGetArgumentType +from cutlass.operators.fusion.passes.pass_dag_2_tree import PassDAG2Tree +from cutlass.operators.fusion.passes.pass_get_impl import PassGetImpl +from cutlass.operators.fusion.passes.pass_fix_element_d import PassFixElementD +from cutlass.operators.fusion.passes.pass_layout_elimination import ( + PassLayoutManipulateElimination, +) +from cutlass.operators.fusion.passes.pass_manager import EVTPassManager +from cutlass.operators.fusion.passes.pass_preprocess_red import PassPreprocessRed +from cutlass.operators.fusion.passes.pass_shape_type_propagation import ( + PassShapeTypePropagation, +) +from cutlass.operators.fusion.passes.smem_size_calculator import GetSmemSize + +__all__ = [ + "EVTGraphDrawer", + "PassGetArgumentType", + "PassDAG2Tree", + "PassGetImpl", + "PassFixElementD", + "PassLayoutManipulateElimination", + "EVTPassManager", + "PassPreprocessRed", + "PassShapeTypePropagation", + "GetSmemSize", +] diff --git a/operators/cutlass/operators/fusion/passes/graph_drawer.py b/operators/cutlass/operators/fusion/passes/graph_drawer.py new file mode 100644 index 000000000..33cdeeed8 --- /dev/null +++ b/operators/cutlass/operators/fusion/passes/graph_drawer.py @@ -0,0 +1,133 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + +from __future__ import annotations + + +from cutlass.operators.fusion.library import DataTypeTag +from cutlass.operators.fusion.ir.dag_ir import DAGIR + + +_COLOR_MAP = { + "load": '"AliceBlue"', + "compute": "LemonChiffon1", + "accumulator": "LightGrey", + "store": "PowderBlue", + "layout": "lightseagreen", + "dag": "darkorange", +} + + +class EVTGraphDrawer: + """ + Visualize a EVT DAGIR with graphviz + """ + + def __init__(self, graph: DAGIR, name: str): + self._name = name + self._dot_graphs = {} + + self._dot_graphs[name] = self._to_dot(graph, name) + + def _get_node_style(self, node): + template = { + "shape": "record", + "fillcolor": "#CAFFE3", + "style": '"filled,rounded"', + "fontcolor": "#000000", + } + if node.op in _COLOR_MAP: + template["fillcolor"] = _COLOR_MAP[node.op] + else: + raise NotImplementedError("unknown node op") + if node.disabled: + template["fontcolor"] = "grey" + template["fillcolor"] = "white" + return template + + def _get_node_label(self, node): + label = "{" + f"name={node.name}|op={node.op}" + if node.op == "layout": + label += f"|fn={node.fn.__name__}" + for key in node.kwargs: + label += f"|{key}={node.kwargs[key]}" + if node.underlying_impl is not None: + label += f"|impl={type(node.underlying_impl).__name__}" + if node.op == "load": + label += f"|element_output={DataTypeTag[node.underlying_impl.element]}" + elif node.op == "compute": + label += f"|element_compute={DataTypeTag[node.underlying_impl.element_compute]}|element_output={DataTypeTag[node.underlying_impl.element_output]}" + elif node.op == "store": + label += f"|element_store={DataTypeTag[node.underlying_impl.element]}|element_output={DataTypeTag[node.underlying_impl.element_output]}" + label += f"|element_output={DataTypeTag[node.underlying_impl.element_output]}" + if node.tensor is not None: + shape = node.tensor.shape + stride = node.tensor.stride + label += f"|shape={shape}|stride={stride}" + + if hasattr(node, "store_tensor") and node.store_tensor is not None: + store_shape = node.store_tensor.shape + store_stride = node.store_tensor.stride + label += f"|store_shape={store_shape}|store_stride={store_stride}" + + label += "}" + return label + + def _to_dot(self, graph: DAGIR, name: str): + import pydot + + dot_graph = pydot.Dot(name, rankdir="TB") + for node in graph.nodes_meta: + style = self._get_node_style(node) + label = self._get_node_label(node) + dot_node = pydot.Node(node.name, label=label, **style) + dot_graph.add_node(dot_node) + if node.op == "dag": + dot_subgraph = self._to_dot(node.subgraph, name=node.name) + self._dot_graphs[node.name] = dot_subgraph + + # Add edges + for src, dst in graph.edges: + weight = graph.get_edge_weight(src, dst) + dot_graph.add_edge(pydot.Edge(src, dst, label=weight)) + + return dot_graph + + def get_dot_graph(self) -> "pydot.Dot": + return [ + (key, self.get_dot_graph_by_name(key)) for key in self._dot_graphs.keys() + ] + + def get_dot_graph_by_name(self, name) -> "pydot.Dot": + return self._dot_graphs[name] + + def get_main_dot_graph(self) -> "pydot.Dot": + return self._dot_graphs[self._name] diff --git a/operators/cutlass/operators/fusion/passes/pass_argument_type.py b/operators/cutlass/operators/fusion/passes/pass_argument_type.py new file mode 100644 index 000000000..20b56f35d --- /dev/null +++ b/operators/cutlass/operators/fusion/passes/pass_argument_type.py @@ -0,0 +1,136 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Construct the epilogue visitor argument type +""" + +from cutlass.operators.fusion.ir import TopoVisitorNode +from cutlass.operators.fusion.ir.c_types import visitor_factory +from cutlass.operators.fusion.passes.pass_dag_2_tree import PassDAG2Tree +from cutlass.operators.fusion.passes.pass_get_impl import PassGetImpl +from cutlass.operators.fusion.passes.pass_manager import EVTPassBase +from cutlass.operators.fusion.passes.pass_shape_type_propagation import ( + PassShapeTypePropagation, +) +from cutlass.operators.fusion.passes.util import cc_map + + +class PassGetArgumentType(EVTPassBase): + """ + Construct the epilogue visitor argument type + """ + + dependencies = [ + PassShapeTypePropagation, # The Layout of all nodes must be set + PassDAG2Tree, # The type of each node must be set + PassGetImpl, # The DAG subgraphs must be set + ] + + def requires(self) -> None: + # Check "D" is in the node list + if cc_map[self.cc] in [90, 100] and (not self.dag_ir.has_node("D")): + raise SyntaxError( + "Sm90+ EVT requires the epilogue to have a returned tensor D, " + "but the variable 'D' is not found in the return values." + ) + + def call(self): + nodes = self.dag_ir.nodes_topological_order() + self.argument_types = {} + for node in nodes: + meta = self.dag_ir.get_node_meta(node) + if not meta.disabled: + self.argument_types[node] = meta.underlying_impl.argument_type + if node == "D" and cc_map[self.cc] in [90, 100]: + continue + if isinstance(meta, TopoVisitorNode): + self.get_dag_argument_type(node) + else: + self.get_evt_argument_type(node) + + self.cc_specific_method(self.set_argument_type)() + + def get_evt_argument_type(self, node): + # Sort the input nodes by edge weight + input_types = [ + self.argument_types[child] for child in self.dag_ir.get_all_inputs(node) + ] + if len(input_types) > 0: + self.argument_types[node] = visitor_factory( + input_types + [self.argument_types[node]], + self.dag_ir.get_all_inputs(node) + [node], + ) + + def get_dag_argument_type(self, node): + meta = self.dag_ir.get_node_meta(node) + subgraph = meta.subgraph + subgraph_nodes = subgraph.nodes_topological_order() + # Visit the unvisited nodes in subgraph + for n in subgraph_nodes: + M = subgraph.get_node_meta(n) + if M.disabled: + continue + else: + self.argument_types[n] = M.underlying_impl.argument_type + input_types = [self.argument_types[child] for child in subgraph_nodes[:-1]] + if len(input_types) > 0: + self.argument_types[node] = visitor_factory( + input_types, subgraph_nodes[:-1] + ) + + def set_argument_type(self): + pass + + def sm90_set_argument_type(self): + self.dag_ir.epilogue_thread_type = self.argument_types[ + self.dag_ir.get_all_inputs("D")[0] + ] + # Get the tensorD argument type + self.dag_ir.arg_d_type = self.dag_ir.get_node_meta( + "D" + ).underlying_impl.argument_type_d + + # Get the tensorC argument type + if self.dag_ir.has_node("C"): + self.dag_ir.arg_c_type = self.dag_ir.get_node_meta( + "C" + ).underlying_impl.argument_type_c + else: + self.dag_ir.arg_c_type = self.dag_ir.arg_d_type + + def sm100_set_argument_type(self): + self.sm90_set_argument_type() + + def sm80_set_argument_type(self): + nodes = self.dag_ir.nodes_topological_order() + self.dag_ir.epilogue_thread_type = self.argument_types[nodes[-1]] diff --git a/operators/cutlass/operators/fusion/passes/pass_dag_2_tree.py b/operators/cutlass/operators/fusion/passes/pass_dag_2_tree.py new file mode 100644 index 000000000..3ca79c474 --- /dev/null +++ b/operators/cutlass/operators/fusion/passes/pass_dag_2_tree.py @@ -0,0 +1,176 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Merge non-tree sub-graphs of the DAG IR into a single DAG. The fused DAG will be implemented +by the topological visitor, while the rest of the graph will be implemented with the tree visitor. +""" + +from copy import deepcopy + +from cutlass.operators.fusion.ir import DAGIR, TopoVisitorNode +from cutlass.operators.fusion.passes.pass_get_impl import PassGetImpl +from cutlass.operators.fusion.passes.pass_manager import EVTPassBase +from cutlass.operators.fusion.passes.pass_shape_type_propagation import ( + PassShapeTypePropagation, +) + + +class PassDAG2Tree(EVTPassBase): + """ + Convert the DAG IR to Tree by fusing subgraphs + """ + + dependencies = [PassShapeTypePropagation, PassGetImpl] + + def call(self): + # Step 1: find the nodes that have multiple parents + multi_parent_nodes = [] + + for node in self.dag_ir.nodes_topological_order(): + if self.dag_ir.out_degree(node) > 1: + multi_parent_nodes.append(node) + # Step 2: find the lowest common ancestor (LCA) of all its parents + for node in multi_parent_nodes: + # A multi-parent node could be already fused by the previous node + if not self.dag_ir.has_node(node): + continue + # A node uncovered by the previous fusions can have out degree change + # Case 1: it has <= 1 edges to the previously fused subgraph, no degree change + # Case 2: it has more than one edges to the previously fused subgraph, degree drops + if self.dag_ir.out_degree(node) <= 1: + continue + + # Otherwise, the node still + reachable_nodes = [] + # Complexity: O(Dout*N) + for parent in self.dag_ir.get_users(node): + reachable_nodes.append(set(self.dag_ir.all_reachable_nodes(parent))) + # get the common reachable objects + common_items = set.intersection(*reachable_nodes) + node_to_fuse = set.union(*reachable_nodes).difference(common_items) + + lca = None + # If common ancestor exists, find the lowest one + if len(common_items) > 0: + topo_order = self.dag_ir.nodes_topological_order() + topo_idx = -1 + for item in common_items: + if lca is None: + lca = item + topo_idx = topo_order.index(item) + else: + if topo_idx > topo_order.index(item): + lca = item + topo_idx = topo_order.index(item) + else: + # there is no common ancestor for all the parents, we pack all the reachable + # nodes into a single DAG node as a fallback. The lca should be the input node of + # one of the output nodes with out_degree = 0 + potential_output_nodes = [] + for node in node_to_fuse: + if self.dag_ir.out_degree(node) == 0: + potential_output_nodes.append(node) + if len(potential_output_nodes) == 0: + raise RuntimeError("No output node with out degree = 0 found.") + + output_node = None + if self.dag_ir.cc >= 90: + # For SM90+, the lca should be the input node of D + if not self.dag_ir.has_node("D"): + raise RuntimeError("D is not a node in the DAG IR.") + output_node = "D" + else: + output_node = potential_output_nodes[0] + + if output_node is None: + raise RuntimeError("No output node found.") + lca = self.dag_ir.get_all_inputs(output_node)[0] + node_to_fuse.remove(output_node) + + # The lca is the output node of the DAG node + # Get the nodes to be fused + node_to_fuse.add(lca) + # Get all the input nodes + all_input_nodes = [] + all_output_nodes = [] + for node in node_to_fuse: + all_input_nodes.append(set(self.dag_ir.get_all_inputs(node))) + all_output_nodes.append(set(self.dag_ir.get_users(node))) + all_input_nodes = set.union(*all_input_nodes) + all_output_nodes = set.union(*all_output_nodes) + + new_subgraph_nodes = set.union( + node_to_fuse, all_input_nodes, all_output_nodes + ) + + # Create the subgraph + subgraph_ = self.dag_ir._graph.subgraph(new_subgraph_nodes) + subgraph = DAGIR(self.dag_ir.cc) + for node in subgraph_.nodes: + meta = deepcopy(self.dag_ir.get_node_meta(node)) + if node not in node_to_fuse: + meta.disabled = True + subgraph.add_node(meta) + for edge in subgraph_.edges: + subgraph.add_edge( + edge[0], edge[1], self.dag_ir.get_edge_weight(edge[0], edge[1]) + ) + + # Create the fused node + dag_node = TopoVisitorNode( + name=f"dag_{lca}", + subgraph=subgraph, + output_node=self.dag_ir.get_node_meta(lca), + ) + self.dag_ir.add_node(dag_node) + + # Add input edges + for idx, node in enumerate(all_input_nodes): + self.dag_ir.add_edge(node, dag_node.name, weight=idx) + + # Replace all uses with DAG node (only 1 output node) + self.dag_ir.replace_all_uses_with(lca, dag_node.name) + + # Remove all fused nodes + node_to_fuse.remove(lca) + for node in node_to_fuse: + self.dag_ir.remove_node(node) + + def ensures(self) -> None: + # Ensure that after the pass, the resulting DAG becomes a tree + for node in self.dag_ir.nodes: + out_degree = self.dag_ir.out_degree(node) + if out_degree > 1: + raise RuntimeError( + f"PassDAG2Tree failed. Node {node} still have outdegree = {out_degree}" + ) diff --git a/operators/cutlass/operators/fusion/passes/pass_fix_element_d.py b/operators/cutlass/operators/fusion/passes/pass_fix_element_d.py new file mode 100644 index 000000000..ef232354f --- /dev/null +++ b/operators/cutlass/operators/fusion/passes/pass_fix_element_d.py @@ -0,0 +1,86 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Fix the element_output of producer of D. + +In Sm90 epilogue visitor, the node writing D to gmem does not have internal +element converter, so the compute node producing D must have element_output = type(D). +""" + +from cutlass.operators.fusion.passes.pass_layout_elimination import ( + PassLayoutManipulateElimination, +) +from cutlass.operators.fusion.passes.pass_manager import EVTPassBase + + +class PassFixElementD(EVTPassBase): + """ + In Sm90 epilogue visitor, the node writing D to gmem does not have internal + element converter, so the compute node producing D must have + element_output = type(D) + """ + + dependencies = [PassLayoutManipulateElimination] + + def get_producer(self, node: str, element_D, visited=None): + if visited is None: + visited = set() + + if node in visited: + raise RuntimeError( + f"Cycle detected while traversing to producer of D: {node}" + ) + visited.add(node) + + node_meta = self.dag_ir.get_node_meta(node) + if node_meta.op == "compute": + node_meta.element_output = element_D + elif node_meta.op == "store": + inputs = self.dag_ir.get_all_inputs(node) + if len(inputs) != 1: + raise RuntimeError( + f"Store node {node} has {len(inputs)} inputs, expected 1" + ) + self.get_producer(inputs[0], element_D, visited) + elif node_meta.op == "load": + node_meta.element_output = element_D + else: + raise NotImplementedError( + f"Unsupported node op: {node_meta.op} when getting producer of D" + ) + + def call(self): + if self.dag_ir.has_node("D"): + node_d_meta = self.dag_ir.get_node_meta("D") + element_D = node_d_meta.store_tensor.element + self.get_producer("D", element_D) diff --git a/operators/cutlass/operators/fusion/passes/pass_get_impl.py b/operators/cutlass/operators/fusion/passes/pass_get_impl.py new file mode 100644 index 000000000..bfd4ef8ef --- /dev/null +++ b/operators/cutlass/operators/fusion/passes/pass_get_impl.py @@ -0,0 +1,93 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Infer the underlying implement of each node. + +While the frontend only distinguish between Load/Store/Compute Node, +each of these nodes can have different underlying implementation based +on their layout. For instance, a LoadNode can be AuxLoad, Row/Col/Scalar broadcast, etc. +This pass infers the underlying impl of each node +""" + +import cutlass.operators.fusion.backend as evt_backend +from cutlass.operators.fusion.ir import DAGIR, LoadNode +from cutlass.operators.fusion.passes.pass_fix_element_d import PassFixElementD +from cutlass.operators.fusion.passes.pass_manager import EVTPassBase +from cutlass.operators.fusion.passes.pass_no_op_elimination import PassNoOpElimination +from cutlass.operators.fusion.passes.pass_shape_type_propagation import ( + PassShapeTypePropagation, +) +from cutlass.operators.fusion.passes.util import cc_map + + +class PassGetImpl(EVTPassBase): + """ + While the frontend only distinguish between Load/Store/Compute Node, + each of these nodes can have different underlying implementation based + on their layout. For instance, a LoadNode can be AuxLoad, Row/Col/Scalar broadcast, etc. + This pass infers the underlying impl of each node + """ + + dependencies = [ + PassShapeTypePropagation, # The shape and type info are required for inference + PassFixElementD, + ] + + def __init__(self, dag_ir: DAGIR) -> None: + super().__init__(dag_ir) + self.no_op_elimination = PassNoOpElimination(dag_ir) + + def requires(self) -> None: + # Verify "accum" is in the arg list + if not self.dag_ir.has_node("accum"): + raise SyntaxError("Cannot find 'accum' in the argument list.") + + def call(self): + # The loop structure of the epilogue is determined by the + # accumulator shape + accumulator: LoadNode = self.dag_ir.get_node_meta("accum") + problem_size = accumulator.tensor.shape + + for node_meta in self.dag_ir.node_metas_topological_order(): + node_meta.get_underlying_impl(problem_size) + + def ensures(self) -> None: + # Some nodes will be lowered to NoOp, eliminate them + self.no_op_elimination() + # Lower to cc-specific impl + for node_meta in self.dag_ir.nodes_meta: + node_impl_ccs = getattr(evt_backend, f"sm{cc_map[self.cc]}_nodes") + node_meta.underlying_impl = getattr( + node_impl_ccs, + f"Sm{cc_map[self.cc]}" + node_meta.underlying_impl.__class__.__name__, + )(node_meta) diff --git a/operators/cutlass/operators/fusion/passes/pass_layout_elimination.py b/operators/cutlass/operators/fusion/passes/pass_layout_elimination.py new file mode 100644 index 000000000..b80c4c067 --- /dev/null +++ b/operators/cutlass/operators/fusion/passes/pass_layout_elimination.py @@ -0,0 +1,230 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Eliminate layout manipulation nodes +""" + +from copy import deepcopy + +from cutlass.operators.fusion.ir import DAGIR, LayoutNode +from cutlass.operators.fusion.passes.pass_manager import EVTPassBase +from cutlass.operators.fusion.passes.pass_shape_type_propagation import ( + PassShapeTypePropagation, +) + + +class PassLayoutManipulateElimination(EVTPassBase): + """ + Eliminate layout manipulation nodes + """ + + dependencies = [PassShapeTypePropagation] + + def __init__(self, dag_ir: DAGIR) -> None: + super().__init__(dag_ir) + self.copy_cnt = 0 + + def call(self): + self.layout_nodes_worklist = self.get_all_layout_nodes() + # Run while loop utill all layout nodes are eliminated + while len(self.layout_nodes_worklist) > 0: + node = self.layout_nodes_worklist.pop(0) + # for node in layout_nodes: + # Step 1: get the propagation direction + direction = self.get_propagation_direction(node) + self.visited = [] + getattr(self, f"propagate_to_{direction}")( + self.dag_ir.get_node_meta(node), node + ) + # Eliminate the current node + input_node = self.dag_ir.get_all_inputs(node)[0] + self.dag_ir.replace_all_uses_with(node, input_node) + # layout_nodes = self.get_all_layout_nodes() + + def get_all_layout_nodes(self): + layout_nodes = [] + for node_meta in reversed(self.dag_ir.node_metas_topological_order()): + if isinstance(node_meta, LayoutNode): + layout_nodes.append(node_meta.name) + return layout_nodes + + def get_propagation_direction(self, node: str): + """ + The logic is propagating all layout nodes away from the accumulator node. + """ + self.visited = [] + self.get_influenced_users(node) + nodes_influenced_dir_users = self.visited + self.visited = [] + self.get_influenced_inputs(node) + nodes_influenced_dir_inputs = self.visited + + if ( + "accum" in nodes_influenced_dir_users + and "accum" not in nodes_influenced_dir_inputs + ): + return "inputs" + elif ( + "accum" not in nodes_influenced_dir_users + and "accum" in nodes_influenced_dir_inputs + ): + return "users" + else: + raise RuntimeError("Unsolved propagation direction") + + # Get all influenced nodes if we propagate along the user direction + def get_influenced_users(self, node: str): + if node in self.visited: + return + self.visited.append(node) + + users = self.dag_ir.get_users(node) + for user in users: + self.get_influenced_users(user) + user_inputs = [] + for user in users: + user_inputs.append(set(self.dag_ir.get_all_inputs(user))) + if len(user_inputs) > 0: + user_inputs = set.union(*user_inputs) + user_inputs.remove(node) + for input_node in user_inputs: + self.get_influenced_inputs(input_node) + + # Get all influenced nodes if we propagate along the input direction + def get_influenced_inputs(self, node: str): + if node in self.visited: + return + self.visited.append(node) + + inputs = self.dag_ir.get_all_inputs(node) + for input_node in inputs: + self.get_influenced_inputs(input_node) + input_users = [] + for input_node in inputs: + input_users.append(set(self.dag_ir.get_users(input_node))) + if len(input_users) > 0: + input_users = set.union(*input_users) + input_users.remove(node) + for user in input_users: + self.get_influenced_users(user) + + def add_copy_before(self, layout_node_meta: LayoutNode, target: str): + copied_node_meta = deepcopy(layout_node_meta) + copied_node = f"{copied_node_meta.name}_copy{self.copy_cnt}" + self.copy_cnt += 1 + copied_node_meta.name = copied_node + self.dag_ir.add_node(copied_node_meta) + # Add edges + target_inputs = self.dag_ir.get_all_inputs(target) + for src in target_inputs: + self.dag_ir.remove_edge(src, target) + self.dag_ir.add_edge(src, copied_node) + self.dag_ir.add_edge(copied_node, target) + self.layout_nodes_worklist.append(copied_node) + + def add_copy_after(self, layout_node_meta: LayoutNode, target: str): + copied_node_meta = deepcopy(layout_node_meta) + copied_node = f"{copied_node_meta.name}_copy{self.copy_cnt}" + self.copy_cnt += 1 + copied_node_meta.name = copied_node + self.dag_ir.add_node(copied_node_meta) + # Add edges + users = self.dag_ir.get_users(target) + for user in users: + self.dag_ir.remove_edge(target, user) + self.dag_ir.add_edge(copied_node, user) + self.dag_ir.add_edge(target, copied_node) + self.layout_nodes_worklist.append(copied_node) + + # Propagate the layout `node` along the user direction + def propagate_to_users(self, layout_node_meta: LayoutNode, node: str): + """ + Propagate layout node to users + """ + if node in self.visited: + # Avoid applying twice + return + self.visited.append(node) + + node_meta = self.dag_ir.get_node_meta(node) + if layout_node_meta.name != node: + if isinstance(node_meta, LayoutNode): + # Layout node is not transparent with layout node + self.add_copy_before(layout_node_meta, node) + return + else: + layout_node_meta.apply_to_user(node_meta) + + users = self.dag_ir.get_users(node) + user_inputs = [] + for user in users: + user_inputs.append(set(self.dag_ir.get_all_inputs(user))) + for user in users: + self.propagate_to_users(layout_node_meta, user) + if len(user_inputs) > 0: + user_inputs = set.union(*user_inputs) + user_inputs.remove(node) + for input_node in user_inputs: + self.propagate_to_inputs( + layout_node_meta.get_inverse_node(), input_node + ) + + # Propagate the layout `node` along the input direction + def propagate_to_inputs(self, layout_node_meta: LayoutNode, node: str): + """ + Propagate layout node to inputs + """ + if node in self.visited: + # Avoid applying twice + return + self.visited.append(node) + + node_meta = self.dag_ir.get_node_meta(node) + if layout_node_meta.name != node: + if isinstance(node_meta, LayoutNode): + # Layout node is not transparent with layout node + self.add_copy_after(layout_node_meta, node) + return + else: + layout_node_meta.apply_to_input(node_meta) + inputs = self.dag_ir.get_all_inputs(node) + input_users = [] + for input_node in inputs: + input_users.append(set(self.dag_ir.get_users(input_node))) + for input_node in inputs: + self.propagate_to_inputs(layout_node_meta, input_node) + if len(input_users) > 0: + input_users = set.union(*input_users) + input_users.remove(node) + for user in input_users: + self.propagate_to_users(layout_node_meta.get_inverse_node(), user) diff --git a/operators/cutlass/operators/fusion/passes/pass_manager.py b/operators/cutlass/operators/fusion/passes/pass_manager.py new file mode 100644 index 000000000..e75462f69 --- /dev/null +++ b/operators/cutlass/operators/fusion/passes/pass_manager.py @@ -0,0 +1,185 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Pass manager for DAG IR. +""" + +from typing import Any + +import networkx as nx + +from cutlass.operators.fusion.ir import DAGIR +from cutlass.operators.fusion.passes.util import cc_map + + +class EVTPassBase: + """ + Base class for EVT Passes + """ + + dependencies = [] + + def __init__(self, dag_ir: DAGIR) -> None: + self.dag_ir = dag_ir + self.cc = self.dag_ir.cc + + def requires(self) -> None: + """ + This function will be called before the pass is run. + """ + pass + + def call(self) -> None: + """ + The pass that is run through the self.dag_ir + """ + raise NotImplementedError( + f"__call__ is not overwritten in Pass {self.__class__.__name__}" + ) + + def ensures(self) -> None: + """ + This function will be called after the pass is run. + """ + pass + + def __call__(self) -> Any: + self.requires() + self.call() + self.ensures() + + def cc_specific_method(self, func): + """ + This enables defining function that behaves differently under different cc + The simplest example of using this function is the following + + .. highlight:: python + .. code-block:: python + + class ExamplePass(EVTPassBase): + + def call(sekf): + # This automatically select the smXX_func based on current cc + self.cc_specific_method(self.func)() + + # Interface func, can be empty + def func(self): + pass + + # Sm90 specific func + def sm90_func(self): + // sm90 specific method + return + + # Sm80 specific func + def sm80_func(self): + // sm80 specific method + return + """ + func_name = f"sm{cc_map[self.cc]}_{func.__name__}" + if hasattr(self, func_name): + return getattr(self, func_name) + else: + raise NotImplementedError( + f"func {func.__name__} is not overwritten for Sm{self.cc}" + ) + + +class EVTPassManager(nx.DiGraph): + """ + Topological-based Pass Manager. + Each registered pass has a list of dependencies. The pass manager organizes + the passes as a DAG and launch the compiler passes under topological order. + """ + + def __new__(cls, *args, **kwargs): + # NetworkX 3.0+ changed DiGraph.__new__ to accept fewer arguments. + # Override to accept and ignore extra arguments from __init__. + return super().__new__(cls) + + def __init__(self, dag_ir: DAGIR, pass_list, soft_dependencies=None): + super().__init__() + self.dag_ir = dag_ir + for pass_cls in pass_list: + self.add_pass(pass_cls) + + self.passes = pass_list + self.soft_dependencies = soft_dependencies + if self.soft_dependencies is None: + self.soft_dependencies = [] + + self.sorted_passes = self.schedule() + + def get_callable(self, pass_name): + """ + Return the callable of the pass + """ + return self.nodes[pass_name]["callable"] + + def add_pass(self, pass_cls): + """ + Add a pass to the pass manager + :param pass_cls: the class of pass + :type pass_cls: derived class of EVTPassBase + """ + name = pass_cls.__name__ + pass_callable = pass_cls(self.dag_ir) + self.add_node(name, callable=pass_callable) + + def schedule(self): + """ + Schedule the added passes under topological order + """ + # Add edges + for pass_name in self.nodes: + callable = self.get_callable(pass_name) + for dependency_cls in callable.dependencies: + if dependency_cls not in self.passes: + if dependency_cls not in self.soft_dependencies: + raise ValueError( + f"Pass {pass_name} depends on {dependency_cls.__name__}, which is not in the pass list" + ) + else: + continue + self.add_edge(dependency_cls.__name__, type(callable).__name__) + + # Topological sort + return list(nx.topological_sort(self)) + + def __call__(self) -> Any: + """ + Launch the registered passes + """ + for pass_name in self.sorted_passes: + callable = self.get_callable(pass_name) + callable() diff --git a/operators/cutlass/operators/fusion/passes/pass_no_op_elimination.py b/operators/cutlass/operators/fusion/passes/pass_no_op_elimination.py new file mode 100644 index 000000000..d242ee7ff --- /dev/null +++ b/operators/cutlass/operators/fusion/passes/pass_no_op_elimination.py @@ -0,0 +1,59 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +No op elimination node +""" + +from typing import Any + +from cutlass.operators.fusion.ir import NoOpImpl +from cutlass.operators.fusion.passes.pass_manager import EVTPassBase + + +class PassNoOpElimination(EVTPassBase): + """ + The dead node elimination pass removes nodes with NoOpImpl in DAG IR + """ + + dependencies = [] + + def call(self) -> Any: + for node in self.dag_ir.nodes_topological_order(): + node_meta = self.dag_ir.get_node_meta(node) + if isinstance(node_meta.underlying_impl, NoOpImpl): + inputs = self.dag_ir.get_all_inputs(node) + if len(inputs) != 1: + raise RuntimeError( + f"NoOp node {node} has {len(inputs)} inputs, expected 1" + ) + self.dag_ir.replace_all_uses_with(node, inputs[0]) diff --git a/operators/cutlass/operators/fusion/passes/pass_preprocess_red.py b/operators/cutlass/operators/fusion/passes/pass_preprocess_red.py new file mode 100644 index 000000000..24a5a6bf1 --- /dev/null +++ b/operators/cutlass/operators/fusion/passes/pass_preprocess_red.py @@ -0,0 +1,96 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Preprocess the reduction nodes. + +The parser treats reduction as Compute(op=(reg_reduce_fn, gmem_reduce_fn)) - Store() +This pass fuses these into a single store node, and then replaces all uses of the +current node with the new store node. +""" + +from cutlass.operators.fusion.ir import ComputeNode, StoreNode +from cutlass.operators.fusion.passes.pass_manager import EVTPassBase + + +class PassPreprocessRed(EVTPassBase): + """ + Preprocess red nodes + """ + + def call(self): + # Step 1: find the compute nodes with op=red + red_compute_nodes = [] + for node_meta in self.dag_ir.nodes_meta: + if isinstance(node_meta, ComputeNode) and type(node_meta.fn) == tuple: + # To keep the frontend simple, the reduction nodes + # are parsed into compute nodes by default + # The simple heuristic to distinguish between compute + # and reduction node is that compute node is a single function, + # while the reduction node is a tuple of functions for + # in-register reduction and atomic global memory reduction + red_compute_nodes.append(node_meta.name) + + # Step 2: for each compute, merge it with the succeeding store + for node in red_compute_nodes: + # Verify + users = self.dag_ir.get_users(node) + inputs = self.dag_ir.get_all_inputs(node) + # Has a single user + assert len(users) == 1 + assert len(inputs) == 1 + user = users[0] + input_node = inputs[0] + + user_meta = self.dag_ir.get_node_meta(user) + # Must be a store node + assert isinstance(user_meta, StoreNode) + # With output degree == 0 + assert self.dag_ir.out_degree(user) == 0 + # Register the reduce op + node_meta = self.dag_ir.get_node_meta(node) + user_meta.reg_reduce_fn, user_meta.gmem_reduce_fn = node_meta.fn + user_meta.element_compute = node_meta.element_compute + user_meta.round_style = node_meta.round_style + + # Replace all uses + self.dag_ir.remove_edge(input_node, node) + input_users = self.dag_ir.get_users(input_node) + for iu in input_users: + weight = self.dag_ir.get_edge_weight(input_node, iu) + self.dag_ir.add_edge(user, iu, weight) + self.dag_ir.remove_edge(input_node, iu) + self.dag_ir.add_edge(input_node, user) + self.dag_ir.remove_node(node) + + # Register the reduction name + self.dag_ir.reduction_names.append(user) diff --git a/operators/cutlass/operators/fusion/passes/pass_shape_type_propagation.py b/operators/cutlass/operators/fusion/passes/pass_shape_type_propagation.py new file mode 100644 index 000000000..d1f2b0357 --- /dev/null +++ b/operators/cutlass/operators/fusion/passes/pass_shape_type_propagation.py @@ -0,0 +1,60 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Shape and type propagation pass +""" + +from cutlass.operators.fusion.ir.node import NodeBase +from cutlass.operators.fusion.passes.pass_manager import EVTPassBase +from cutlass.operators.fusion.passes.pass_preprocess_red import PassPreprocessRed + + +class PassShapeTypePropagation(EVTPassBase): + """ + Propagate the shape and type of all nodes + """ + + dependencies = [PassPreprocessRed] + + def call(self): + # Propagate the node shape and type + for node in self.dag_ir.nodes_topological_order(): + node_meta: NodeBase = self.dag_ir.get_node_meta(node) + input_node_metas = self.dag_ir.get_all_inputs_meta(node) + node_meta.type_propagation(input_node_metas) + node_meta.shape_propagation(input_node_metas) + + for node in reversed(self.dag_ir.nodes_topological_order()): + node_meta: NodeBase = self.dag_ir.get_node_meta(node) + input_node_metas = self.dag_ir.get_all_inputs_meta(node) + node_meta.broadcast_propagation(input_node_metas) diff --git a/operators/cutlass/operators/fusion/passes/smem_size_calculator.py b/operators/cutlass/operators/fusion/passes/smem_size_calculator.py new file mode 100644 index 000000000..2540bf988 --- /dev/null +++ b/operators/cutlass/operators/fusion/passes/smem_size_calculator.py @@ -0,0 +1,363 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Compute the shared memory size in bytes +""" + +from math import gcd + +from cutlass.operators.fusion.library import DataType, DataTypeSize, EpilogueScheduleType +from cutlass.operators.fusion.ir import TopoVisitorNode, DAGIR +from cutlass.operators.fusion.pycute import flatten, shape_div, product + + +class GetSmemSize: + """ + Get the size in byte of shared memory used by the kernel + """ + + def __init__(self, dag_ir: DAGIR) -> None: + self.dag_ir = dag_ir + self.cc = self.dag_ir.cc + + # + # Sm90 epilogue specific + # + + def sm90_epilogue_tile(self, tile_description): + # Get the epilogue tile size + schedule = tile_description.epilogue_schedule + if schedule == EpilogueScheduleType.TmaWarpSpecialized: + element_d = self.dag_ir.get_node_meta("D").element + nperf = ( + 64 + if ( + DataTypeSize[element_d] == 8 + and tile_description.threadblock_shape[1] % 64 == 0 + ) + else 32 + ) + epi_tile_m = min(64, tile_description.threadblock_shape[0]) + epi_tile_n = gcd( + min(nperf, tile_description.threadblock_shape[1]), + tile_description.threadblock_shape[1], + ) + epilogue_tile_mn = (epi_tile_m, epi_tile_n) + elif schedule == EpilogueScheduleType.TmaWarpSpecializedCooperative: + epi_tile_m = min(128, tile_description.threadblock_shape[0]) + epi_tile_n = gcd( + min(32, tile_description.threadblock_shape[1]), + tile_description.threadblock_shape[1], + ) + epilogue_tile_mn = (epi_tile_m, epi_tile_n) + else: + raise NotImplementedError(f"Unsupported schedule: {schedule}") + + # Get the pipeline stages + stages_d = 2 + epi_tiles = product( + shape_div(tuple(tile_description.threadblock_shape)[:2], epilogue_tile_mn) + ) + if self.dag_ir.has_node("C"): + element_c = self.dag_ir.get_node_meta("C").element + else: + element_c = None + + element_d = self.dag_ir.get_node_meta("D").element + if element_c == element_d: + reuse_smem_c = True + else: + reuse_smem_c = False + stages_c = max(epi_tiles, stages_d + 1) if reuse_smem_c else epi_tiles + + # Record the epilogue tile + self.cta_tile_mnk = tuple(tile_description.threadblock_shape) + self.epilogue_tile_mn = epilogue_tile_mn + self.epi_tiles = epi_tiles + self.stages_c = stages_c + self.stages_d = stages_d + self.reuse_smem_c = reuse_smem_c + self.element_c = element_c + self.element_d = element_d + self.is_source_supported = element_c is not None + + def sm90_or_sm100_epilogue_smem_size(self, tile_description): + # Get the Fusion Storage + nodes = self.dag_ir.nodes_topological_order() + self.smem_types = {} + for node in nodes: + meta = self.dag_ir.get_node_meta(node) + if not meta.disabled: + self.smem_types[node] = meta.underlying_impl.get_smem_size( + self.cta_tile_mnk, + self.epilogue_tile_mn, + self.stages_c, + self.stages_d, + self.epi_tiles, + ) + if node == "D": + continue + if isinstance(meta, TopoVisitorNode): + self.get_dag_smem_type(node) + else: + self.get_evt_smem_type(node) + + thread_smem_size = self.smem_types[self.dag_ir.get_all_inputs("D")[0]][0] + # Get the Tensor Storage + tensors = [] + if self.is_source_supported: + smem_C = ( + DataTypeSize[self.element_c] + * product(self.epilogue_tile_mn) + * self.stages_c + // 8 + ) + tensors.append((smem_C, 128)) + else: + tensors.append((0, 1)) + if self.reuse_smem_c: + tensors.append((0, 128)) + else: + smem_D = ( + DataTypeSize[self.element_d] + * product(self.epilogue_tile_mn) + * self.stages_d + // 8 + ) + tensors.append((smem_D, 128)) + tensors.append((thread_smem_size, 128)) + + tensor_smem_size = self.get_struct_size(tensors) + # Get pipeline storage size + # sizeof(uint64_t * stages_c * 2), alignment of uint64_t + # 2 is for FullBarrier and EmptyBarrier + pipeline_smem_size = (8 * self.stages_c * 2, 8) + + # get SharedStorage size + smem_size = self.get_struct_size([tensor_smem_size, pipeline_smem_size]) + return smem_size[0] + + def sm90_epilogue_smem_size(self, tile_description): + """ + Compute the shared memory size of sm90 collective epilogue + """ + self.sm90_epilogue_tile(tile_description) + return self.sm90_or_sm100_epilogue_smem_size(tile_description) + + # + # Sm100 epilogue specific + # + + def sm100_epilogue_tile(self, tile_description): + cta_tile = ( + tile_description.blackwell_threadblock_shape[0], + tile_description.blackwell_threadblock_shape[1], + ) + mma_tile = cta_tile + + if tile_description.is_2sm: + cta_tile = (cta_tile[0] // 2, cta_tile[1]) + + if tile_description.is_2sm and mma_tile[0] == 128: + tmem_warps = (2, 2) + else: + tmem_warps = (4, 1) + + if self.dag_ir.has_node("C"): + element_c = self.dag_ir.get_node_meta("C").element + element_c_size = DataTypeSize[element_c] + else: + element_c = None + element_c_size = 0 + + element_d = self.dag_ir.get_node_meta("D").element + + DisableSource = ( + element_c is None + or not self.dag_ir.has_node("C") + or self.dag_ir.get_node_meta("C").element == DataType.void + ) + + CtaM = cta_tile[0] + CtaN = cta_tile[1] + WarpM = tmem_warps[0] + WarpN = tmem_warps[1] + MaxBits = max(element_c_size, DataTypeSize[element_d]) + DpFull = 32 + M = min(CtaM, DpFull * WarpM) + + if DisableSource: + # Epilogues w/o residual load are less sensitive to smem allocation + # Target a fixed amount of compute per epilogue iteration + if MaxBits == 4: + # Make epilogue tile larger to reduce the epilogue iterations. + # 64 is the experimental value. It will minimize epilogue iterations but keep the number of A/B buffers the same. + ComputeElts = 8192 + Nperf = ComputeElts // M + else: + ComputeElts = 4096 + Nperf = ComputeElts // M + else: + # Epilogues w/ residual load are more sensitive to smem allocation + # Target optimal smem distribution between epilogue+mainloop based on datatype+tilesize + if MaxBits == 32: + Nperf = 16 if CtaM > 64 and CtaN <= 128 else 32 + elif MaxBits == 16: + Nperf = 32 if CtaN <= 128 else 64 + else: + Nperf = 64 + + def is_m_major(layout): + return flatten(layout.stride[0]) == 1 + + if DisableSource or is_m_major(self.dag_ir.get_node_meta("C").tensor.layout): + N_min_C = 8 * WarpN + elif element_c_size == 6: + N_min_C = 128 * WarpN + else: + N_min_C = (128 // element_c_size) * WarpN + + if is_m_major(self.dag_ir.get_node_meta("D").tensor.layout): + N_min_D = 8 * WarpN + elif DataTypeSize[element_d] == 6: + N_min_D = 128 * WarpN + else: + N_min_D = (128 // DataTypeSize[element_d]) * WarpN + + N = min(CtaN, max(Nperf, N_min_C, N_min_D)) + + tile_m = M + tile_n_size = N // WarpN * WarpN + + epilogue_tile_mn = (tile_m, tile_n_size) + epi_tiles = product( + shape_div(tuple(tile_description.threadblock_shape)[:2], epilogue_tile_mn) + ) + + stages_d = min(epi_tiles, 2) + reuse_smem_c = element_c_size > 8 + + if reuse_smem_c: + stages_c = max(min(epi_tiles, 4), stages_d + 1) + else: + stages_c = min(epi_tiles, 4) + + # Record the epilogue tile + self.cta_tile_mnk = tuple(tile_description.threadblock_shape) + self.epilogue_tile_mn = epilogue_tile_mn + self.epi_tiles = epi_tiles + self.stages_c = stages_c + self.stages_d = stages_d + self.reuse_smem_c = reuse_smem_c + self.element_c = element_c + self.element_d = element_d + self.is_source_supported = not DisableSource + + def sm100_epilogue_smem_size(self, tile_description): + """ + Compute the shared memory size of sm100 collective epilogue + """ + self.sm100_epilogue_tile(tile_description) + return self.sm90_or_sm100_epilogue_smem_size(tile_description) + + def __call__(self, tile_description): + return getattr(self, f"sm{self.cc}_epilogue_smem_size")(tile_description) + + # + # Helper functions + # + + @staticmethod + def get_visitor_size(members: list, ebo: bool): + """ + Get the size of struct in bytes + """ + offset = 0 + max_alignment = 1 + if len(members) > 0: + # Get alignment + for _, alignment in members: + max_alignment = max(max_alignment, alignment) + + for type_size, _ in members: + if type_size != 0: + offset = ( + (offset + max_alignment - 1) // max_alignment + ) * max_alignment + if type_size == 0 and not ebo: + offset += 1 + else: + offset += type_size + offset = ((offset + max_alignment - 1) // max_alignment) * max_alignment + return (offset, max_alignment) + else: + # Struct size is at least 1 + return (1, 1) + + def get_struct_size(self, members: list): + """ + Get the size of struct in bytes + """ + return self.get_visitor_size(members, False) + + def get_evt_smem_type(self, node): + # Sort the input nodes by edge weight + input_types = [ + self.smem_types[child] for child in self.dag_ir.get_all_inputs(node) + ] + input_types.append(self.smem_types[node]) + if len(input_types) > 1: + ebo = len(input_types) > 4 + self.smem_types[node] = self.get_visitor_size(input_types, ebo) + + def get_dag_smem_type(self, node): + meta = self.dag_ir.get_node_meta(node) + subgraph = meta.subgraph + subgraph_nodes = subgraph.nodes_topological_order() + # Visit the unvisited nodes in subgraph + for n in subgraph_nodes: + M = subgraph.get_node_meta(n) + if M.disabled: + continue + else: + self.smem_types[n] = M.underlying_impl.get_smem_size( + self.cta_tile_mnk, + self.epilogue_tile_mn, + self.stages_c, + self.stages_d, + self.epi_tiles, + ) + input_types = [self.smem_types[child] for child in subgraph_nodes[:-1]] + if len(input_types) > 0: + ebo = len(input_types) > 4 + self.smem_types[node] = self.get_visitor_size(input_types, ebo) diff --git a/operators/cutlass/operators/fusion/passes/util.py b/operators/cutlass/operators/fusion/passes/util.py new file mode 100644 index 000000000..f8ca3d790 --- /dev/null +++ b/operators/cutlass/operators/fusion/passes/util.py @@ -0,0 +1,46 @@ + +# +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Utilities for passes +""" + +# Map from the CC of the kernel to the EVT implementation that the CC targets +cc_map = { + 80: 80, + 86: 80, + 89: 80, + 90: 90, + 100: 100, + 101: 100, + 103: 100, +} diff --git a/operators/cutlass/operators/fusion/pycute/__init__.py b/operators/cutlass/operators/fusion/pycute/__init__.py new file mode 100644 index 000000000..bdb630f03 --- /dev/null +++ b/operators/cutlass/operators/fusion/pycute/__init__.py @@ -0,0 +1,36 @@ + +# +# Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +from .int_tuple import * +from .layout import * +from .swizzle import * +from .typing import * diff --git a/operators/cutlass/operators/fusion/pycute/int_tuple.py b/operators/cutlass/operators/fusion/pycute/int_tuple.py new file mode 100644 index 000000000..a766f924f --- /dev/null +++ b/operators/cutlass/operators/fusion/pycute/int_tuple.py @@ -0,0 +1,229 @@ + +# +# Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Functions for manipulating IntTuples +""" + +from functools import reduce +from itertools import chain +from typing import Union +from .typing import Integer + + +def is_int(x): + return isinstance(x, Integer) + + +def is_tuple(x): + return isinstance(x, tuple) + + +def flatten(t): + if is_tuple(t): + if len(t) == 0: + return () + else: + return tuple(i for a in t for i in flatten(a)) + else: + return (t,) + + +def signum(a): + return bool(a > 0) - bool(a < 0) + + +def product(a): + if is_tuple(a): + return reduce(lambda val, elem: val * product(elem), a, 1) + else: + return a + + +def inner_product(a, b): + if is_tuple(a): # tuple tuple + assert len(a) == len(b) + return sum(inner_product(x, y) for x, y in zip(a, b)) + else: # "int" "int" + assert not is_tuple(b) + return a * b + + +def tuple_max(a): + if is_tuple(a): + return max(tuple_max(x) for x in a) + else: + return a + + +def elem_scale(a, b): + if is_tuple(a): + if is_tuple(b): # tuple tuple + assert len(a) == len(b) + return tuple(elem_scale(x, y) for x, y in zip(a, b)) + else: # tuple "int" + assert False # Error + else: + if is_tuple(b): # "int" tuple + return elem_scale(a, product(b)) + else: # "int" "int" + return a * b + + +# Inclusive prefix ceil div with output congruent to input a +def shape_div(a, b): + if is_tuple(a): + if is_tuple(b): # tuple tuple + assert len(a) == len(b) + return tuple(shape_div(x, y) for x, y in zip(a, b)) + else: # tuple "int" + # r = [shape_div(a[0],b)] + [shape_div(a[i],b := shape_div(b, product(a[i-1]))) for i in range(1,len(a))] + r = [] + for v in a: + r.append(shape_div(v, b)) + b = shape_div(b, product(v)) + return tuple(r) + else: + if is_tuple(b): # "int" tuple + return shape_div(a, product(b)) + else: # "int" "int" + assert a % b == 0 or b % a == 0 + return (a + b - 1) // b + + +# Exclusive prefix product with output congruent to input a +def prefix_product(a, init=1): + if is_tuple(a): + if is_tuple(init): # tuple tuple + assert len(a) == len(init) + return tuple(prefix_product(x, i) for x, i in zip(a, init)) + else: # tuple "int" + # r = [prefix_product(a[0],init)] + [prefix_product(a[i],init := init * product(a[i-1])) for i in range(1,len(a))] + r = [] + for v in a: + r.append(prefix_product(v, init)) + init = init * product(v) + return tuple(r) + else: + if is_tuple(init): # "int" tuple + assert False # Error + else: # "int" "int" + return init + + +def idx2crd(idx, shape, stride=None): + if stride is None: + stride = prefix_product(shape) + + if is_tuple(idx): + if is_tuple(shape): # tuple tuple tuple + assert len(idx) == len(shape) and len(idx) == len(stride) + return tuple(idx2crd(i, s, d) for i, s, d in zip(idx, shape, stride)) + else: # tuple "int" "int" + assert False # Error + else: + if is_tuple(shape): # "int" tuple tuple + assert len(shape) == len(stride) + return tuple(idx2crd(idx, s, d) for s, d in zip(shape, stride)) + else: # "int" "int" "int" + return (idx // stride) % shape + + +def crd2idx(crd, shape, stride=None): + if stride is None: + stride = prefix_product(shape) + + if is_tuple(crd): + if is_tuple(shape): # tuple tuple tuple + assert len(crd) == len(shape) and len(crd) == len(stride) + return sum(crd2idx(c, s, d) for c, s, d in zip(crd, shape, stride)) + else: # tuple "int" "int" + assert False, f"crd={crd}, shape={shape}" # Error + else: + if crd is None: + crd = 0 + + if is_tuple(shape): # "int" tuple tuple + assert len(shape) == len(stride) + result = 0 + for i in range(len(shape) - 1): + result += crd2idx(crd % product(shape[i]), shape[i], stride[i]) + crd = crd // product(shape[i]) + return result + crd2idx(crd, shape[-1], stride[-1]) + else: # "int" "int" "int" + return crd * stride + + +# Transform crd into the dst_shape's iteration space +def crd2crd(crd, dst_shape, src_shape=None): + if is_tuple(crd): + if is_tuple(dst_shape): # tuple tuple + assert len(crd) == len(dst_shape) + return tuple(crd2crd(x, y) for x, y in zip(crd, dst_shape)) + else: # tuple "int" + # Ambiguous unless we have src_shape + assert src_shape is not None + return crd2idx(crd, src_shape) + else: + if is_tuple(dst_shape): # "int" tuple + return idx2crd(crd, dst_shape) + else: # "int" "int" + assert crd < dst_shape + return crd + + +# Filter trg according to crd: keep only elements of trg that are paired with None +def slice_(crd: Union[None, tuple, int], trg: Union[tuple, int]): + if is_tuple(crd): + if is_tuple(trg): # tuple tuple + assert len(crd) == len(trg) + # match C++ behavior of `filter_tuple` using `tuple_cat(...)` + return tuple( + chain( + *filter(lambda x: x != (), [slice_(c, s) for c, s in zip(crd, trg)]) + ) + ) + else: + assert False # tuple "int" : Error + elif crd is None: + # match C++ behavior `return cute::tuple{b};` + return (trg,) + else: + return () + + +# Determine if None appears at any of an int_tuples' terminals +def has_none(a: Union[None, tuple, int]): + if is_tuple(a): + return any(has_none(v) for v in a) + else: + return a is None diff --git a/operators/cutlass/operators/fusion/pycute/layout.py b/operators/cutlass/operators/fusion/pycute/layout.py new file mode 100644 index 000000000..8dd61108f --- /dev/null +++ b/operators/cutlass/operators/fusion/pycute/layout.py @@ -0,0 +1,409 @@ +# +# Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Definition of CuTe Layouts and functions to manipulate them +""" + +from itertools import chain + +from .int_tuple import * + + +class LayoutBase: + pass + + +def is_layout(x): + return isinstance(x, LayoutBase) + + +class Layout(LayoutBase): + def __init__(self, _shape, _stride=None): + self.shape = _shape + if _stride is None: + self.stride = prefix_product(self.shape) + else: + self.stride = _stride + + # operator == + def __eq__(self, other): + return self.shape == other.shape and self.stride == other.stride + + # operator len(L) (len [rank] like tuples) + def __len__(self): + if is_tuple(self.shape): + return len(self.shape) + else: + return 1 + + # operator () (map coord to idx) + def __call__(self, *args): + """ + Map a logical coordinate to a linear index (Coord has no Underscore slice operators) + OR + Slice the layout and return the sublayout (Coord has an Underscore slice op) + + Follow the same behavior of `Layout::operator(Coord const&)` in cute C++ + """ + if has_none(args): + if len(args) == 1: + return Layout(slice_(args[0], self.shape), slice_(args[0], self.stride)) + else: + return Layout(slice_(args, self.shape), slice_(args, self.stride)) + else: + if len(args) == 1: + return crd2idx(args[0], self.shape, self.stride) + else: + return crd2idx(args, self.shape, self.stride) + + # operator [] (get-i like tuples) + def __getitem__(self, i): + if is_tuple(self.shape): + return Layout(self.shape[i], self.stride[i]) + else: + assert i == 0 + return Layout(self.shape, self.stride) + + # size(layout) Size of the domain + def size(self): + return product(self.shape) + + # cosize(layout) Size of the codomain + def cosize(self): + return self(self.size() - 1) + 1 + + # print and str + def __str__(self): + return f"{self.shape}:{self.stride}" + + # error msgs and representation + def __repr__(self): + return f"Layout({self.shape},{self.stride})" + + +# Make Layout from a list of layouts (each layout it's own mode in the result) +def make_layout(*layouts): + if len(layouts) == 1 and not is_layout(layouts[0]): + layouts = layouts[0] + + shape, stride = zip(*((a.shape, a.stride) for a in layouts)) + return Layout(shape, stride) + + +# Size of the domain +def size(layout): + if is_layout(layout): + return layout.size() + return product(layout) + + +# Size of the codomain +def cosize(layout): + return layout.cosize() + + +# Layout coalesce -- flatten and combine as many modes as possible while preserving the int-to-int function +def coalesce(layout, profile=None): + if is_tuple(profile): + assert len(layout) >= len(profile) + return make_layout( + chain( + (coalesce(layout[i], profile[i]) for i in range(0, len(profile))), + (layout[i] for i in range(len(profile), len(layout))), + ) + ) + + result_shape = [1] + result_stride = [0] + for shape, stride in zip(flatten(layout.shape), flatten(layout.stride)): + # skip their shape-1s + if shape == 1: + continue + # replace our shape-1 with anything + elif result_shape[-1] == 1: + result_shape[-1] = shape + result_stride[-1] = stride + # merge modes if the shape*stride match + elif result_shape[-1] * result_stride[-1] == stride: + result_shape[-1] = result_shape[-1] * shape + # append a new mode + else: + result_shape.append(shape) + result_stride.append(stride) + + if len(result_shape) == 1: + return Layout(result_shape[0], result_stride[0]) + else: + return Layout(tuple(result_shape), tuple(result_stride)) + + +# Layout filter -- replace all stride-0 modes with size-1 and then coalesce to remove them +def filter(layout, profile=None): + if is_tuple(profile): + assert len(layout) >= len(profile) + return make_layout( + chain( + (filter(layout[i], profile[i]) for i in range(0, len(profile))), + (layout[i] for i in range(len(profile), len(layout))), + ) + ) + + result_shape = [] + result_stride = [] + for shape, stride in zip(flatten(layout.shape), flatten(layout.stride)): + # skip their shape-1s and stride-0s + if not (shape == 1 or stride == 0): + result_shape.append(shape) + result_stride.append(stride) + + if len(result_shape) == 0: + return Layout(1, 0) + else: + return coalesce(Layout(tuple(result_shape), tuple(result_stride))) + + +# Layout composition +# Use tuples-of-layouts to perform this operation by-mode and None as no-op +def composition(layoutA, layoutB): + if layoutB is None: + return layoutA + elif is_int(layoutB): + return composition(layoutA, Layout(layoutB)) + elif is_tuple(layoutB): + assert len(layoutA) >= len(layoutB) + return make_layout( + chain( + (composition(layoutA[i], layoutB[i]) for i in range(0, len(layoutB))), + (layoutA[i] for i in range(len(layoutB), len(layoutA))), + ) + ) + elif is_tuple(layoutB.shape): + return make_layout(composition(layoutA, layoutB_i) for layoutB_i in layoutB) + + if layoutB.stride == 0: + return Layout(layoutB.shape, 0) + else: + result_shape = [] + result_stride = [] + rest_shape = layoutB.shape + rest_stride = layoutB.stride + flat_A = coalesce(layoutA) + for curr_shape, curr_stride in zip( + flatten(flat_A.shape)[:-1], flatten(flat_A.stride)[:-1] + ): + assert curr_shape % rest_stride == 0 or rest_stride % curr_shape == 0 + new_shape = min(max(1, curr_shape // rest_stride), rest_shape) + + if new_shape != 1: + result_shape.append(new_shape) + result_stride.append(rest_stride * curr_stride) + + rest_shape = rest_shape // new_shape + rest_stride = -( + -rest_stride // curr_shape + ) # Python exclusive impl: "//" is always floor div so == ceil_div(abs(rest_stride), curr_shape) * signum(rest_stride) + + if rest_shape != 1 or len(result_shape) == 0: + result_shape.append(rest_shape) + result_stride.append(rest_stride * flatten(flat_A.stride)[-1]) + + if len(result_shape) == 1: + return Layout(result_shape[0], result_stride[0]) + else: + return Layout(tuple(result_shape), tuple(result_stride)) + + +# Layout complement +def complement(layout, max_idx=1): + if is_int(layout): + return complement(Layout(layout)) + + result_shape = [] + result_stride = [] + current_idx = 1 + + sorted_DS = sorted(zip(flatten(layout.stride), flatten(layout.shape))) + for stride, shape in sorted_DS: + if stride == 0 or shape == 1: + continue + + in_bound = current_idx <= shape * stride + # To support symbolic value which can't be evaluated now + assert (type(in_bound) is not bool) or in_bound + + result_shape.append(stride // current_idx) + result_stride.append(current_idx) + current_idx = shape * stride + + result_shape.append((max_idx + current_idx - 1) // current_idx) # ceil_div + result_stride.append(current_idx) + + return coalesce(Layout(tuple(result_shape), tuple(result_stride))) + + +# Layout right inverse +def right_inverse(layout): + if layout is None: + return None + elif is_int(layout): + return Layout(layout) + + result_shape = [] + result_stride = [] + current_idx = 1 + + flat_shape = flatten(layout.shape) + flat_stride = flatten(layout.stride) + sorted_DSA = sorted(zip(flat_stride, flat_shape, prefix_product(flat_shape))) + for stride, shape, rstride in sorted_DSA: + if shape == 1: + continue + if current_idx != stride: + break + + result_shape.append(shape) + result_stride.append(rstride) + current_idx = shape * stride + + return coalesce(Layout(tuple(result_shape), tuple(result_stride))) + + +# Layout left inverse +def left_inverse(layout): + if layout is None: + return None + elif is_int(layout): + return Layout(layout) + return right_inverse(make_layout(layout, complement(layout))) + + +# Split a layout by the composition of B and the "rest" +# Use tuples-of-layouts to perform this operation by-mode and None as no-op +def logical_divide(layoutA, layoutB): + if layoutB is None: + return layoutA + elif is_int(layoutB): + return logical_divide(layoutA, Layout(layoutB)) + elif is_tuple(layoutB): + assert len(layoutA) >= len(layoutB) + return make_layout( + chain( + ( + logical_divide(layoutA[i], layoutB[i]) + for i in range(0, len(layoutB)) + ), + (layoutA[i] for i in range(len(layoutB), len(layoutA))), + ) + ) + + return composition( + layoutA, make_layout(layoutB, complement(layoutB, size(layoutA))) + ) + + +# Reproduce a layoutA over a layoutB +# Use tuples-of-layouts to perform this operation by-mode and None as no-op +def logical_product(layoutA, layoutB): + if layoutB is None: + return layoutA + elif is_int(layoutB): + return logical_divide(layoutA, Layout(layoutB)) + elif is_tuple(layoutB): + assert len(layoutA) >= len(layoutB) + return make_layout( + chain( + ( + logical_product(layoutA[i], layoutB[i]) + for i in range(0, len(layoutB)) + ), + (layoutA[i] for i in range(len(layoutB), len(layoutA))), + ) + ) + + return make_layout( + layoutA, + composition(complement(layoutA, size(layoutA) * cosize(layoutB)), layoutB), + ) + + +# Gather the modes from a hierarchical logical_divide or logical_product +def hier_unzip(splitter, layoutA, layoutB): + if layoutB is None: + return make_layout(Layout(1, 0), layoutA) + elif is_tuple(layoutB): + assert len(layoutA) >= len(layoutB) + # A layout with shape ((A,a),(B,b),(C,c)) + split = make_layout( + hier_unzip(splitter, layoutA[i], layoutB[i]) for i in range(0, len(layoutB)) + ) + # Gather to shape ((A,B,C,...),(a,b,c,...,y,z)) + return make_layout( + make_layout(split[i][0] for i in range(0, len(layoutB))), + make_layout( + chain( + (split[i][1] for i in range(0, len(layoutB))), + (layoutA[i] for i in range(len(layoutB), len(layoutA))), + ) + ), + ) + + # splitter must return a rank-2 layout + return splitter(layoutA, layoutB) + + +# Apply logical divide hierarchically and gather the split modes into two modes +def zipped_divide(layoutA, layoutB): + return hier_unzip(logical_divide, layoutA, layoutB) + + +# Perform logical divide hierarchically and gather tiles (B-layouts) into a new mode +def tiled_divide(layoutA, layoutB): + result = zipped_divide(layoutA, layoutB) + return make_layout([result[0]] + [result[1][i] for i in range(len(result[1]))]) + + +# Apply logical product hierarchically and gather the split modes into two modes +def zipped_product(layoutA, layoutB): + return hier_unzip(logical_product, layoutA, layoutB) + + +# Perform logical product hierarchically and gather tiles (B-layouts) into a new mode +def tiled_product(layoutA, layoutB): + result = zipped_product(layoutA, layoutB) + return make_layout([result[0]] + [result[1][i] for i in range(len(result[1]))]) + + +def slice_and_offset(crd: tuple, layout: Layout): + return ( + Layout(slice_(crd, layout.shape), slice_(crd, layout.stride)), + crd2idx(crd, layout.shape, layout.stride), + ) diff --git a/operators/cutlass/operators/fusion/pycute/swizzle.py b/operators/cutlass/operators/fusion/pycute/swizzle.py new file mode 100644 index 000000000..ca0745b92 --- /dev/null +++ b/operators/cutlass/operators/fusion/pycute/swizzle.py @@ -0,0 +1,133 @@ + +# +# Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +""" +Methods for layout swizzling +""" + +from .layout import * + + +def shiftr(a, s): + return a >> s if s > 0 else shiftl(a, -s) + + +def shiftl(a, s): + return a << s if s > 0 else shiftr(a, -s) + + +## A generic Swizzle functor +# 0bxxxxxxxxxxxxxxxYYYxxxxxxxZZZxxxx +# ^--^ Base is the number of least-sig bits to keep constant +# ^-^ ^-^ Bits is the number of bits in the mask +# ^---------^ Shift is the distance to shift the YYY mask +# (pos shifts YYY to the right, neg shifts YYY to the left) +# +# e.g. Given +# 0bxxxxxxxxxxxxxxxxYYxxxxxxxxxZZxxx +# the result is +# 0bxxxxxxxxxxxxxxxxYYxxxxxxxxxAAxxx where AA = ZZ xor YY +# +class Swizzle: + def __init__(self, bits, base, shift): + assert bits >= 0 + assert base >= 0 + assert abs(shift) >= bits + self.bits = bits + self.base = base + self.shift = shift + bit_msk = (1 << bits) - 1 + self.yyy_msk = bit_msk << (base + max(0, shift)) + self.zzz_msk = bit_msk << (base - min(0, shift)) + + # operator () (transform integer) + def __call__(self, offset): + return offset ^ shiftr(offset & self.yyy_msk, self.shift) + + # Size of the domain + def size(self): + return 1 << (self.bits + self.base + abs(self.shift)) + + # Size of the codomain + def cosize(self): + return self.size() + + # print and str + def __str__(self): + return f"SW_{self.bits}_{self.base}_{self.shift}" + + # error msgs and representation + def __repr__(self): + return f"Swizzle({self.bits},{self.base},{self.shift})" + + +class ComposedLayout(LayoutBase): + def __init__(self, layoutB, offset, layoutA): + self.layoutB = layoutB + self.offset = offset + self.layoutA = layoutA + + # operator == + def __eq__(self, other): + return ( + self.layoutB == other.layoutB + and self.offset == other.offset + and self.layoutA == other.layoutA + ) + + # operator len(L) (len [rank] like tuples) + def __len__(self): + return len(self.layoutA) + + # operator () (map coord to idx) + def __call__(self, *args): + return self.layoutB(self.offset + self.layoutA(*args)) + + # operator [] (get-i like tuples) + def __getitem__(self, i): + return ComposedLayout(self.layoutB, self.offset, self.layoutA[i]) + + # size(layout) Size of the domain + def size(self): + return size(self.layoutA) + + # cosize(layout) Size of the codomain + def cosize(self): + return cosize(self.layoutB) + + # print and str + def __str__(self): + return f"{self.layoutB} o {self.offset} o {self.layoutA}" + + # error msgs and representation + def __repr__(self): + return f"ComposedLayout({repr(self.layoutB)},{repr(self.offset)},{repr(self.layoutA)})" diff --git a/operators/cutlass/operators/fusion/pycute/typing.py b/operators/cutlass/operators/fusion/pycute/typing.py new file mode 100644 index 000000000..0bb7d95e9 --- /dev/null +++ b/operators/cutlass/operators/fusion/pycute/typing.py @@ -0,0 +1,42 @@ + +# +# Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +from abc import ABC + + +class Integer(ABC): + @classmethod + def __subclasshook__(cls, c): + if c in [bool, float]: + return False + + return issubclass(c, int) diff --git a/operators/cutlass/operators/manifest.py b/operators/cutlass/operators/manifest.py new file mode 100644 index 000000000..f97cbbc2a --- /dev/null +++ b/operators/cutlass/operators/manifest.py @@ -0,0 +1,153 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from cutlass.operators.arch import TargetSm +from cutlass.operators.providers import available_providers + +if TYPE_CHECKING: + from collections.abc import Callable + + from cutlass.operators.arguments import RuntimeArguments + from cutlass.operators.base import Operator + from cutlass.operators.metadata import OperatorMetadata + from cutlass.operators.providers import Provider + +_logger = logging.getLogger(__name__) + + +class Manifest: + """A Manifest holds a collection of generated Operators.""" + + def __init__(self) -> None: + self._candidate_operators = [] + + @staticmethod + def get_operators( + args: RuntimeArguments = None, + metadata_filter: Callable[[OperatorMetadata], bool] | None = None, + target_sm: TargetSm | str | None = None, + providers: list[Provider] = None, + ) -> list[Operator]: + """Get the operators that match the given arguments, metadata filter, compute capability and available providers. + + Args: + args (RuntimeArguments): The arguments of the operator + metadata_filter (Callable[[OperatorMetadata], bool]): A boolean function that takes in OperatorMetadata and returns whether + an Operator with this metadata should be included + target_sm (TargetSm | str | None): Optional target compute capability for which operators should be fetched + providers (list[Provider]): Optionally limit the search to the given providers. If None, all available providers are used. + + Returns: + list[Operator]: The operators that match the given arguments, metadata filter, and compute capability + """ + target_sm = TargetSm.ensure(target_sm) + providers = providers if providers is not None else available_providers.values() + filter_fn = metadata_filter or (lambda x: True) + + if args is None and target_sm is None: + full_filter_fn = filter_fn + else: + + def full_filter_fn(metadata: OperatorMetadata) -> bool: + if not filter_fn(metadata): + return False + if not (supports := metadata.supports(args, target_sm=target_sm)): + _logger.debug( + f"Rejecting operator {metadata.operator_name}. Reason: {supports.error}" + ) + return False + return True + + epilogue_args = getattr(args, "epilogue", None) + operators = [ + op + # Generate operators from all providers + for provider in providers + # Filter operators by metadata and args + for op in provider.generate_operators( + full_filter_fn, + target_sm=target_sm, + epilogue_args=epilogue_args, + args=args, + ) + # Do any additional checks on args that are not captured by metadata + if not args or op._supports(args, target_sm=target_sm) + ] + return operators + + def add_operators( + self, + operators: list[Operator], + ) -> None: + """Add the given operators to the Manifest's set of discovered operators. + + Args: + operators (list[Operator]): The operators to add to the Manifest + """ + self._candidate_operators.extend(operators) + + def filter_operators( + self, + args: RuntimeArguments = None, + metadata_filter: Callable[[OperatorMetadata], bool] | None = None, + target_sm: TargetSm | str | None = None, + providers: list[Provider] = None, + ) -> list[Operator]: + """Filter the operators that have been added to this Manifest. + + This is similar to ``get_operators``, but only considers operators previously added + to this Manifest using :meth:`add_operators`. + """ + target_sm = TargetSm.ensure(target_sm) + metadata_filter = metadata_filter or (lambda x: True) + + # Append provider filter to metadata_filter + if providers is not None: + + def full_metadata_filter(metadata: OperatorMetadata) -> bool: + return metadata.provider in providers and metadata_filter(metadata) + else: + full_metadata_filter = metadata_filter + + # Filter self._candidate_operators to those that support inputs & full_metadata_filter + return [ + op + for op in self._candidate_operators + if op.supports(args, target_sm=target_sm) + and full_metadata_filter(op.metadata) + ] + + @property + def operators(self) -> list[Operator]: + """Get the operators that have been added to this Manifest.""" + return self._candidate_operators diff --git a/operators/cutlass/operators/metadata/__init__.py b/operators/cutlass/operators/metadata/__init__.py new file mode 100644 index 000000000..c6281403f --- /dev/null +++ b/operators/cutlass/operators/metadata/__init__.py @@ -0,0 +1,73 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from .base import OperatorMetadata +from .design import ( + BLASDesignMetadata, + DesignMetadata, + Sm100DesignMetadata, +) +from .design.tile_scheduler import ( + CLCDynamicPersistentTileSchedulerMetadata, + StaticPersistentTileSchedulerMetadata, + TileSchedulerMetadata, +) +from .epilogue import EpilogueMetadata +from .operand_constraints import ( + DenseTensorConstraints, + OperandConstraints, + ScaledOperandConstraints, +) +from .operands import ( + GemmOperandsMetadata, + GroupedGemmOperandsMetadata, + OperandsMetadata, +) + +__all__ = [ + # Top-level metadata & immediate constituents + "OperatorMetadata", + "OperandsMetadata", + "DesignMetadata", + "EpilogueMetadata", + + # OperandsMetadata variants + "GemmOperandsMetadata", + "GroupedGemmOperandsMetadata", + # Operand constraints + "OperandConstraints", + "DenseTensorConstraints", + "ScaledOperandConstraints", + # Design metadata variants + "BLASDesignMetadata", + "Sm100DesignMetadata", + # Tile scheduler metadata + "TileSchedulerMetadata", + "CLCDynamicPersistentTileSchedulerMetadata", + "StaticPersistentTileSchedulerMetadata", +] diff --git a/operators/cutlass/operators/metadata/base.py b/operators/cutlass/operators/metadata/base.py new file mode 100644 index 000000000..dfc62a6ca --- /dev/null +++ b/operators/cutlass/operators/metadata/base.py @@ -0,0 +1,170 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from cutlass.operators.arch import TargetSm +from cutlass.operators.status import Status + +if TYPE_CHECKING: + from cutlass.operators.arguments import RuntimeArguments + from cutlass.operators.base import Operator + from cutlass.operators.metadata.design import DesignMetadata + from cutlass.operators.metadata.epilogue import EpilogueMetadata + from cutlass.operators.metadata.operands import OperandsMetadata + from cutlass.operators.providers import Provider + + +@dataclass +class OperatorMetadata: + """OperatorMetadata exposes a uniform schema to inspect the properties of an Operator. + + It allows users to programmatically and uniformly inspect properties of an + Operator without manually inspecting its source code, and without learning + a different schema of properties for each Operator. + + It describes fixed properties of the Operator in three categories: + + - :class:`~cutlass.operators.metadata.OperandsMetadata`: properties of the + operands passed to the Operator and the constraints it places on them. + It describes the conditions for an Operator to functionally support + (compute the correct result of) an operation. + - :class:`~cutlass.operators.metadata.DesignMetadata`: architectural, + performance, or design characteristics of the kernel. Useful for + inspecting and comparing otherwise functionally equivalent Operators. + - :class:`~cutlass.operators.metadata.EpilogueMetadata`: any custom epilogue + computation that was fused into the Operator. + + These properties are declared by the ``Operator`` itself when it is instantiated, + and are fixed after that. + + They may be determined by design, or the Operator's implementation may be parametrized + to generate instances supporting different values of a metadata field. + """ + + operator_name: str + """A unique human-readable Operator name. + + Important: + There are no backward-compatibility guarantees for these names, even for + patch releases. Users are strongly advised to not use these names for + filtering, and use metadata fields instead. + """ + + operator_class: type[Operator] + """The derived class of cutlass.operators.Operator that implements interfaces to expose a DSL kernel""" + + supported_targets: list[TargetSm] + """Supported targets that this Operator can be compiled for + + Only the earliest portable target in each architecture family is reported, + as support for future targets is implied. + + Examples: + * ``supported_targets=[TargetSm("100f"), TargetSm("120f")]`` implies support for ``["100a", "103a", ... , "120a", "121a", ...]`` for all members of ``100f`` & ``120f`` arch families + * ``supported_targets=[TargetSm("80")]`` implies support for ``["80","89","90","90a","100", ...]`` + + To check if this Operator is supported on a desired target_sm, use ``target_sm.supports_operators_from(supported_targets)``. + """ + + operands: OperandsMetadata + """Describes the per-operand requirements of the Operator: per-operand + constraints and any operand-independent choices (e.g. the accumulator + type for a GEMM).""" + + design: DesignMetadata | None = None + """Describes characteristics of the design & implementation of the Operator""" + + epilogue: EpilogueMetadata | None = None + """Describes the custom epilogue of the Operator.""" + + @property + def provider(self) -> Provider: + """The :class:`~cutlass.operators.Provider` that this Operator's class is registered with.""" + return self.operator_class._provider + + def supports( + self, args: RuntimeArguments | None, *, target_sm: TargetSm | str | None = None + ) -> Status: + """Check if the provided ``args``/``target_sm`` satisfy the properties described by this OperatorMetadata. + + Args: + args (RuntimeArguments | None): Check if Operator supports these RuntimeArguments. Can be explicitly + set to None, in order to check only target_sm. + target_sm (TargetSm | str | None): Check if Operator can be compiled/ran on given TargetSm architecture, if any. + + Returns: + Status: Status indicating whether the args/target_sm are supported, and reason for + failure if not. + + Raises: + ValueError: If both args and target_sm are None. + """ + if args is None and target_sm is None: + raise ValueError("Either args or target_sm must be provided") + + if target_sm is not None: + target_sm = TargetSm.ensure(target_sm) + if not target_sm.supports_operators_from(self.supported_targets): + return Status.fail( + f"Operator cannot be compiled for target {target_sm}. " + f"Supported targets: {self.supported_targets}" + ) + + if args is None: + return Status.success() + + # check args + + def supports_or_none(member, corresponding_arg, name: str) -> Status: + # If metadata is absent, accept only when the corresponding argument is also absent. + if member is None: + if corresponding_arg is None: + return Status.success() + return Status.fail( + f"{name} metadata is absent but argument is provided" + ) + return member.supports(args) + + if not (status := self.operands.supports(args)): + return status + + if not (status := supports_or_none(self.design, args.performance, "design")): + return status + + if not ( + status := supports_or_none( + self.epilogue, getattr(args, "epilogue", None), "epilogue" + ) + ): + return status + + return Status.success() diff --git a/operators/cutlass/operators/metadata/design/__init__.py b/operators/cutlass/operators/metadata/design/__init__.py new file mode 100644 index 000000000..c138693b3 --- /dev/null +++ b/operators/cutlass/operators/metadata/design/__init__.py @@ -0,0 +1,45 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from .base import BLASDesignMetadata, DesignMetadata +from .sm100 import Sm100DesignMetadata +from .tile_scheduler import ( + CLCDynamicPersistentTileSchedulerMetadata, + StaticPersistentTileSchedulerMetadata, + TileSchedulerMetadata, +) + +__all__ = [ + "DesignMetadata", + "BLASDesignMetadata", + "Sm100DesignMetadata", + # Tile scheduler metadata + "TileSchedulerMetadata", + "CLCDynamicPersistentTileSchedulerMetadata", + "StaticPersistentTileSchedulerMetadata", +] diff --git a/operators/cutlass/operators/metadata/design/base.py b/operators/cutlass/operators/metadata/design/base.py new file mode 100644 index 000000000..b3c45489c --- /dev/null +++ b/operators/cutlass/operators/metadata/design/base.py @@ -0,0 +1,109 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from cutlass.operators.status import Status + +if TYPE_CHECKING: + from cutlass.operators.arguments import RuntimeArguments + from cutlass.operators.mma import MmaInstruction + + +@dataclass +class DesignMetadata(ABC): + """Describes architectural, performance, or design characteristics of the kernel separate from its functional contribution. + + These properties do not directly alter the logical result of the Operator. + Instead, they distinguish the design characteristics of otherwise logically-equivalent Operators, + and can be used to guide Operator selection and performance characterization. + + This could include properties like: + + * inherent design choices like which MMA instruction it uses, or use of other architectural features + * its performance options, tile scheduling algorithm, etc. + * its tile shape, cluster shape, number of shared memory stages, etc. + """ + + @abstractmethod + def supports(self, args: RuntimeArguments) -> Status: + """Checks whether the provided ``args`` satisfy the properties described by this DesignMetadata. + + It validates that :class:`~cutlass.operators.PerformanceControls` within + ``args.performance`` are supported by the Operator with this DesignMetadata. + + ``DesignMetadata`` does not directly influence the logical result of the Operator, + and thus does not affect what logical operands can be supported in ``RuntimeArguments``. + + Args: + args (RuntimeArguments): The arguments to check support for. + + Returns: + Status: Status indicating whether the provided args satisfy the properties described by the design in this metadata + """ + + +@dataclass(kw_only=True) +class BLASDesignMetadata(DesignMetadata): + """Design metadata for a BLAS (basic-linear algebra subprogram) operation, like GEMM. + + It describes fields common to many kernels that perform BLAS-like operations, such as + tile/cluster shape, MMA instruction type, etc. + """ + + mma_instruction_type: type[MmaInstruction] + """The type of the MMA instruction used by the Operator, if any.""" + + tile_shape: tuple[int, ...] + """Shape of the threadblock-level tile""" + + cluster_shape: tuple[int, ...] + """Shape of the cluster used in terms of threadblocks""" + + num_smem_stages: int | None = None + """Number of shared memory buffers used to pipeline data loading and computation. + When None, the Operator determines this internally.""" + + def supports(self, args: RuntimeArguments) -> Status: + """Checks whether the provided args satisfy the properties described by the design in this metadata. + + Args: + args (RuntimeArguments): The arguments to check support for + + Returns: + Status: Status indicating whether the provided args satisfy the properties described by the design in this metadata + """ + if args.performance is not None: + return Status.fail( + "BLASDesignMetadata does not yet support performance controls" + ) + return Status.success() diff --git a/operators/cutlass/operators/metadata/design/sm100.py b/operators/cutlass/operators/metadata/design/sm100.py new file mode 100644 index 000000000..3a3d3dc57 --- /dev/null +++ b/operators/cutlass/operators/metadata/design/sm100.py @@ -0,0 +1,111 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from cutlass.operators.metadata.design.base import BLASDesignMetadata + +if TYPE_CHECKING: + from cutlass.operators.metadata.design.tile_scheduler import TileSchedulerMetadata + + +@dataclass +class Sm100DesignMetadata(BLASDesignMetadata): + """Design metadata for operators in the SM100 architecture family.""" + + use_2cta_mma: bool + """Whether to use a 2CTA MMA instruction""" + + use_tma_store: bool + """Whether to use TMA to store the results of the operation""" + + tile_scheduler: TileSchedulerMetadata | None = None + """Describes the tile scheduling strategy used by the Operator""" + + fallback_cluster_shape: tuple[int, ...] | None = None + """Fallback cluster shape for preferred cluster mode. + + When ``None`` (default), the Operator uses a single fixed cluster shape + (``BLASDesignMetadata.cluster_shape``). When set: + + * the Operator uses two cluster shapes: ``cluster_shape`` as the + preferred (larger) shape and ``fallback_cluster_shape`` as the smaller + fallback for tiles that cannot fill a full preferred cluster. + * Preferred cluster depth (Z dimension) must be the same as that of fallback cluster. + * Fallback cluster shape must evenly divide the preferred cluster shape. + """ + + def use_fallback_cluster(self) -> bool: + """Return True when this design uses a preferred (larger) and fallback (smaller) cluster shape.""" + return self.fallback_cluster_shape is not None + + def __post_init__(self): + # Enforce the preferred/fallback cluster size invariant. + if self.fallback_cluster_shape is None: + return + + # Both shapes must have the same rank. + if len(self.fallback_cluster_shape) != len(self.cluster_shape): + raise ValueError( + f"fallback_cluster_shape {self.fallback_cluster_shape} must have " + f"the same rank as cluster_shape {self.cluster_shape}" + ) + + # A fallback shape equal to preferred shape is not strictly incorrect, but + # defeats the purpose of specifying fallback cluster and so is likely unintended. + if self.fallback_cluster_shape == self.cluster_shape: + raise ValueError( + "fallback_cluster_shape should differ from cluster_shape " + "(use fallback_cluster_shape=None to express no preferred-cluster " + f"mode instead); got both = {self.cluster_shape}" + ) + + # Depth of both cluster shapes should match + if self.fallback_cluster_shape[-1] != self.cluster_shape[-1]: + raise ValueError( + "fallback_cluster_shape and cluster_shape must agree on the " + "trailing (batch) dim; got " + f"fallback_cluster_shape={self.fallback_cluster_shape}, " + f"cluster_shape={self.cluster_shape}" + ) + + # fallback cluster shape should evenly divide preferred cluster shape, and be smaller + for dim, (fallback, preferred) in enumerate( + zip(self.fallback_cluster_shape[:-1], self.cluster_shape[:-1]) + ): + if preferred % fallback != 0: + raise ValueError( + "fallback_cluster_shape must evenly divide cluster_shape " + "in every non-batch dim; got " + f"fallback_cluster_shape={self.fallback_cluster_shape}, " + f"cluster_shape={self.cluster_shape} " + f"(dim {dim}: {preferred} % {fallback} != 0)" + ) diff --git a/operators/cutlass/operators/metadata/design/tile_scheduler.py b/operators/cutlass/operators/metadata/design/tile_scheduler.py new file mode 100644 index 000000000..1568b9469 --- /dev/null +++ b/operators/cutlass/operators/metadata/design/tile_scheduler.py @@ -0,0 +1,91 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from cutlass.operators.arch import TargetSm + +if TYPE_CHECKING: + from cutlass.operators.metadata import DesignMetadata, OperandsMetadata + + +@dataclass(kw_only=True, frozen=True, slots=True) +class TileSchedulerMetadata(ABC): + """Models an abstract tile scheduler.""" + + @classmethod + @abstractmethod + def supported_targets( + cls, + design: DesignMetadata | None = None, + operands: OperandsMetadata | None = None, + ) -> list[TargetSm]: + """Returns all the supported targets for the tile scheduler. + + Considers optional design and operands metadata when determining + which targets are supported. + """ + raise NotImplementedError("Subclasses must implement this method") + + +@dataclass(kw_only=True, frozen=True, slots=True) +class CLCDynamicPersistentTileSchedulerMetadata(TileSchedulerMetadata): + """Metadata for the CLC dynamic persistent tile scheduler.""" + + @classmethod + def supported_targets( + cls, + design: DesignMetadata | None = None, + operands: OperandsMetadata | None = None, + ) -> list[TargetSm]: + """Returns supported targets for CLC dynamic persistent tile scheduler.""" + return [TargetSm("100f"), TargetSm("110f"), TargetSm("120")] + + +@dataclass(kw_only=True, frozen=True, slots=True) +class StaticPersistentTileSchedulerMetadata(TileSchedulerMetadata): + """Metadata for the static persistent tile scheduler.""" + + @classmethod + def supported_targets( + cls, + design: DesignMetadata | None = None, + operands: OperandsMetadata | None = None, + ) -> list[TargetSm]: + """Returns supported targets for static persistent tile scheduler.""" + return [ + TargetSm("80"), + TargetSm("90"), + TargetSm("100f"), + TargetSm("110f"), + TargetSm("120"), + ] diff --git a/operators/cutlass/operators/metadata/epilogue.py b/operators/cutlass/operators/metadata/epilogue.py new file mode 100644 index 000000000..c0031964c --- /dev/null +++ b/operators/cutlass/operators/metadata/epilogue.py @@ -0,0 +1,71 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from cutlass.operators.status import Status + +if TYPE_CHECKING: + import cutlass.cute as cute + from cutlass.cute.typing import Numeric + + from cutlass.operators.arguments import ( + EpilogueArguments, + RuntimeArguments, + ) + + +class EpilogueMetadata: + """Describes properties related to custom EpilogueArguments that an Operator may support.""" + + def __init__(self, epilogue_args: EpilogueArguments): + self.traced_epilogue = epilogue_args.traced_epilogue + self.tensors = epilogue_args.tensors + self.epilogue_fn = epilogue_args.epilogue_fn + + @classmethod + def from_args(cls, args: EpilogueArguments) -> EpilogueMetadata: + """Creates an EpilogueMetadata from an EpilogueArguments.""" + # For now, EpilogueArguments and EpilogueMetadata are the same + return cls(args) + + @property + def parameters(self) -> list[cute.Tensor | Numeric]: + """Returns the list of input and output parameters of the epilogue.""" + return list(self.tensors.values()) + + @property + def parameter_names(self) -> list[str]: + """Returns the list of names of the input and output parameters of the epilogue.""" + return list(self.tensors.keys()) + + def supports(self, args: RuntimeArguments) -> Status: + """Checks if the provided EpilogueArguments are supported by the Operator.""" + return Status.success() diff --git a/operators/cutlass/operators/metadata/operand_constraints.py b/operators/cutlass/operators/metadata/operand_constraints.py new file mode 100644 index 000000000..5ced00c09 --- /dev/null +++ b/operators/cutlass/operators/metadata/operand_constraints.py @@ -0,0 +1,255 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import cutlass.operators.utils as utils +from cutlass.operators.arguments import ( + DenseTensor, + Operand, + ScaledOperand, + ScaleMode, + ScaleSwizzleMode, +) +from cutlass.operators.status import Status + +if TYPE_CHECKING: + from cutlass.cute.typing import Numeric + + +class OperandConstraints(ABC): + """Abstract base class describing the constraints an Operator places on an Operand. + + An instance of this class describes the requirements an Operator + places on one :class:`~cutlass.operators.arguments.Operand`, such as + its data type, stride pattern, divisibility, and pointer alignment. + """ + + @abstractmethod + def supports(self, operand: Operand | OperandConstraints) -> Status: + """Check whether ``operand`` satisfies these constraints. + + Args: + operand (Operand | OperandConstraints): The runtime operand or + another constraints object to check against ``self``. + + Returns: + Status: ``Status.success()`` when ``operand`` is accepted, otherwise + ``Status.fail(...)`` with a human-readable reason. + """ + + +@dataclass +class DenseTensorConstraints(OperandConstraints): + """Constraints on a single :class:`~cutlass.operators.arguments.DenseTensor`. + + Describes the data type, stride pattern, divisibility, and pointer alignment + that an Operator requires a :class:`~cutlass.operators.arguments.DenseTensor` to satisfy. + """ + + dtype: Numeric + """The data type of the tensor""" + + stride: tuple[int, ...] | None + """The stride of the tensor. If none, any arbitrary stride is allowed.""" + + divisibility: int + """The divisibility requirement on a tensor's stride and shape elements""" + + ptr_alignment_bytes: int + """The alignment of the tensor's data pointer, in bytes + By default, it matches the number of bytes in stride/shape alignment. + """ + + def __init__( + self, + dtype: Numeric, + stride: tuple[int, ...], + divisibility: int, + ptr_alignment_bytes: int | None = None, + ): + self.dtype = dtype + self.stride = stride + self.divisibility = divisibility + self.ptr_alignment_bytes = ptr_alignment_bytes or ( + (divisibility * dtype.width) // 8 + ) + + def supports(self, operand: DenseTensor | DenseTensorConstraints) -> Status: + """Check whether ``operand`` satisfies the constraints described by this object. + + Args: + operand (Operand | OperandConstraints): Operand or peer + constraints object to validate. + + Returns: + Status: ``Status.success()`` if accepted, ``Status.fail(...)`` + otherwise. + """ + is_operand = isinstance(operand, DenseTensor) + is_constraint = isinstance(operand, DenseTensorConstraints) + if not (is_operand or is_constraint): + return Status.fail( + f"Expected DenseTensor or DenseTensorConstraints, got {type(operand)}" + ) + + if operand.dtype != self.dtype: + return Status.fail( + f"Expected element type {self.dtype}, got {operand.dtype}" + ) + + if self.stride is None: + major_mode_idx = ( + utils.tensor.major_mode(operand.shape, operand.stride) + if is_operand + else None + ) + else: + # If the operand is batched, the expected stride is the same as the self.stride + if len(self.stride) == len(operand.stride): + expected_stride = self.stride + # We can try if strides match without the batch mode + elif len(self.stride) - 1 == len(operand.stride): + expected_stride = self.stride[1:] + else: + return Status.fail( + f"Expected tensor with strides of rank {len(self.stride)} (batched) or " + f"{len(self.stride) - 1} (unbatched), got {len(operand.stride)} " + f"({operand.stride})" + ) + + # Strides are considered compatible if: + # 1. They have the same rank (checked above) + # 2. The major mode is the same + + # When setting stride from args, any modes of stride 1 and shape 1 + # are changed to have stride 0. Thus, there will only be one mode + # with stride 1. + major_mode_idx = expected_stride.index(1) + if operand.stride[major_mode_idx] != 1: + return Status.fail( + f"Expected stride[{major_mode_idx}] to be 1, got " + f"{operand.stride[major_mode_idx]} " + f"(strides: {operand.stride})" + ) + + # Check that divisibility constraints are met + if is_operand: + if not self._is_tuple_aligned( + operand.stride, self.divisibility, major_mode_idx + ): + return Status.fail( + f"Expected operand stride to be divisible by {self.divisibility} for " + f"all non-contiguous modes, got {operand.stride}" + ) + else: + # When comparing another DenseTensorConstraints, ensure its divisibility + # is a refinement (i.e. a multiple) of ours. + if operand.divisibility % self.divisibility != 0: + return Status.fail( + f"Expected operand divisibility {operand.divisibility} to be divisible " + f"by {self.divisibility}" + ) + + # Check data ptr alignment, if available + if is_operand and operand.data_ptr % self.ptr_alignment_bytes != 0: + return Status.fail( + f"Expected data pointer to be {self.ptr_alignment_bytes}B-aligned." + ) + + return Status.success() + + def _is_tuple_aligned( + self, tup: tuple[int, ...], divisibility: int, contiguous_dim: int + ) -> bool: + """Check that all elements of ``tup`` are divisible by ``divisibility``. + + The element at ``contiguous_dim`` is 1, and so is exempt from the check. + """ + return all( + t % divisibility == 0 for dim, t in enumerate(tup) if dim != contiguous_dim + ) + + +@dataclass +class ScaledOperandConstraints(OperandConstraints): + """Constraints on a scaled (quantized + scale-factor) operand. + + A scaled operand combines a quantized tensor and a scale tensor with a + chosen scale mode + scale swizzle. Constraints on each piece are carried + in :attr:`quantized` and :attr:`scale`; :meth:`supports` checks all four. + """ + + quantized: DenseTensorConstraints + """Constraints on the narrow-precision tensor holding the quantized values.""" + + scale: DenseTensorConstraints + """Constraints on the scale-factor tensor.""" + + mode: ScaleMode | tuple[int, ...] + """Descriptor of the elements over which each scale factor operates""" + + swizzle: ScaleSwizzleMode + """Swizzling pattern used within scale factor tensor""" + + def supports(self, operand: ScaledOperand | ScaledOperandConstraints) -> Status: + """Check whether ``operand`` satisfies these scaled-operand constraints.""" + if not isinstance(operand, (ScaledOperand, ScaledOperandConstraints)): + return Status.fail( + f"Expected ScaledOperand or ScaledOperandConstraints, got {type(operand)}" + ) + + if not (status := self.quantized.supports(operand.quantized)): + return status + + if not (status := self.scale.supports(operand.scale)): + return status + + if not ScaleMode.compare(self.mode, operand.mode): + return Status.fail(f"Expected scale mode {self.mode}, got {operand.mode}") + + if self.swizzle != operand.swizzle: + return Status.fail( + f"Expected scale swizzle mode {self.swizzle}, got {operand.swizzle}" + ) + + return Status.success() + + def __getattr__(self, attr: str) -> Any: + if hasattr(self.quantized, attr): + return getattr(self.quantized, attr) + else: + raise AttributeError( + f"'{self.__class__.__name__}' object has no attribute '{attr}'" + ) + + diff --git a/operators/cutlass/operators/metadata/operands/__init__.py b/operators/cutlass/operators/metadata/operands/__init__.py new file mode 100644 index 000000000..621cd5de5 --- /dev/null +++ b/operators/cutlass/operators/metadata/operands/__init__.py @@ -0,0 +1,37 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from .base import OperandsMetadata +from .gemm import GemmOperandsMetadata +from .grouped_gemm import GroupedGemmOperandsMetadata + +__all__ = [ + "OperandsMetadata", + "GemmOperandsMetadata", + "GroupedGemmOperandsMetadata", +] diff --git a/operators/cutlass/operators/metadata/operands/base.py b/operators/cutlass/operators/metadata/operands/base.py new file mode 100644 index 000000000..8c7ce2dd6 --- /dev/null +++ b/operators/cutlass/operators/metadata/operands/base.py @@ -0,0 +1,70 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from cutlass.operators.arguments import RuntimeArguments + from cutlass.operators.status import Status + + +@dataclass +class OperandsMetadata(ABC): + """Describes properties expected of the logical operands passed to an Operator. + + It describes the conditions under which an Operator can correctly execute + a given :class:`~cutlass.operators.RuntimeArguments`. + It mirrors the fields of a :class:`~cutlass.operators.RuntimeArguments` subclass implemented by + the Operator, and describes requirements placed on each field. For instance, + :class:`~cutlass.operators.metadata.GemmOperandsMetadata` mirrors + :class:`~cutlass.operators.GemmArguments`. + + | :class:`~cutlass.operators.RuntimeArguments` describes what operation the user wishes + to perform, independent of which Operator will perform it. + | ``OperandsMetadata`` describes conditions under which an :class:`~cutlass.operators.Operator` + can accept and correctly execute :class:`~cutlass.operators.RuntimeArguments` passed to it. + + If an Operator's ``OperandsMetadata`` supports the ``RuntimeArguments`` presented by a user, then + it produces functionally correct result for those RuntimeArguments. + """ + + @abstractmethod + def supports(self, args: RuntimeArguments) -> Status: + """Check whether ``args`` satisfy the operand requirements described by this metadata. + + Args: + args (RuntimeArguments): Runtime arguments to validate. + + Returns: + Status: Status indicating whether ``args`` satisfy the operand + requirements described by this metadata. + """ diff --git a/operators/cutlass/operators/metadata/operands/gemm.py b/operators/cutlass/operators/metadata/operands/gemm.py new file mode 100644 index 000000000..f25ca54dd --- /dev/null +++ b/operators/cutlass/operators/metadata/operands/gemm.py @@ -0,0 +1,90 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from cutlass.operators.arguments import GemmArguments, RuntimeArguments +from cutlass.operators.metadata.operands.base import OperandsMetadata +from cutlass.operators.status import Status + +if TYPE_CHECKING: + from cutlass.cute.typing import Numeric + + from cutlass.operators.metadata.operand_constraints import OperandConstraints + + +@dataclass +class GemmOperandsMetadata(OperandsMetadata): + """Describes conditions under which an Operation functionally supports :class:`~cutlass.operators.GemmArguments`.""" + + A: OperandConstraints + """Constraints on operand ``A`` of the GEMM""" + + B: OperandConstraints + """Constraints on operand ``B`` of the GEMM""" + + out: OperandConstraints + """Constraints on operand ``out`` of the GEMM""" + + accumulator_type: Numeric + """The required data type of the accumulator tensor""" + + def supports(self, other: GemmArguments | GemmOperandsMetadata) -> Status: + """Check whether ``other`` arguments or metadata are compatible with GEMM operands described by this metadata. + + Args: + other (GemmArguments | GemmOperandsMetadata): Either runtime + :class:`~cutlass.operators.GemmArguments` or a peer ``GemmOperandsMetadata`` to + validate. + + Returns: + Status: ``Status.success()`` if every operand and the accumulator + type match, otherwise ``Status.fail(...)`` with the first failing + operand's reason. + """ + if isinstance(other, RuntimeArguments) and not isinstance(other, GemmArguments): + return Status.fail(f"Expected GemmArguments, got {type(other)}") + + if not (status := self.A.supports(other.A)): + return Status.fail(f"Operand `A` is unsupported: {status.error}") + + if not (status := self.B.supports(other.B)): + return Status.fail(f"Operand `B` is unsupported: {status.error}") + + if not (status := self.out.supports(other.out)): + return Status.fail(f"Operand `out` is unsupported: {status.error}") + + if self.accumulator_type != other.accumulator_type: + return Status.fail( + f"Expected accumulator type {self.accumulator_type}, got {other.accumulator_type}" + ) + + return Status.success() diff --git a/operators/cutlass/operators/metadata/operands/grouped_gemm.py b/operators/cutlass/operators/metadata/operands/grouped_gemm.py new file mode 100644 index 000000000..e0768daba --- /dev/null +++ b/operators/cutlass/operators/metadata/operands/grouped_gemm.py @@ -0,0 +1,100 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from cutlass.operators.arguments import GroupedGemmArguments, RuntimeArguments +from cutlass.operators.metadata.operands.base import OperandsMetadata +from cutlass.operators.status import Status + +if TYPE_CHECKING: + from cutlass.cute.typing import Numeric + + from cutlass.operators.metadata.operand_constraints import OperandConstraints + + +@dataclass +class GroupedGemmOperandsMetadata(OperandsMetadata): + """Describes conditions under which an Operation functionally supports :class:`~cutlass.operators.GroupedGemmArguments`.""" + + A: OperandConstraints + """Constraints on operand ``A`` of the grouped GEMM""" + + B: OperandConstraints + """Constraints on operand ``B`` of the grouped GEMM""" + + out: OperandConstraints + """Constraints on operand ``out`` of the grouped GEMM""" + + offsets: OperandConstraints + """Constraints on operand ``offsets`` of the grouped GEMM.""" + + accumulator_type: Numeric + """The required data type of the accumulator tensor""" + + def supports( + self, other: GroupedGemmArguments | GroupedGemmOperandsMetadata + ) -> Status: + """Check whether ``other`` arguments or metadata are compatible with grouped-GEMM operands described by this metadata. + + Args: + other (GroupedGemmArguments | GroupedGemmOperandsMetadata): + Either runtime :class:`~cutlass.operators.GroupedGemmArguments` or a peer + ``GroupedGemmOperandsMetadata`` to validate. + + Returns: + Status: ``Status.success()`` if every operand, ``offsets``, and + the accumulator type match, otherwise ``Status.fail(...)`` with + the first failure. + """ + if isinstance(other, RuntimeArguments) and not isinstance( + other, GroupedGemmArguments + ): + return Status.fail(f"Expected GroupedGemmArguments, got {type(other)}") + + if not (status := self.A.supports(other.A)): + return Status.fail(f"Operand `A` is unsupported: {status.error}") + + if not (status := self.B.supports(other.B)): + return Status.fail(f"Operand `B` is unsupported: {status.error}") + + if not (status := self.out.supports(other.out)): + return Status.fail(f"Operand `out` is unsupported: {status.error}") + + if not (status := self.offsets.supports(other.offsets)): + return Status.fail(f"Operand `offsets` is unsupported: {status.error}") + + if self.accumulator_type != other.accumulator_type: + return Status.fail( + f"Expected accumulator type {self.accumulator_type}, got {other.accumulator_type}" + ) + + return Status.success() diff --git a/operators/cutlass/operators/mma.py b/operators/cutlass/operators/mma.py new file mode 100644 index 000000000..b8e8f522e --- /dev/null +++ b/operators/cutlass/operators/mma.py @@ -0,0 +1,166 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import cutlass + +from cutlass.operators.arch import TargetSm + +if TYPE_CHECKING: + from cutlass.operators.metadata import DesignMetadata, OperandsMetadata + + +class MmaInstruction: + """Models properties of an abstract MMA instruction.""" + + @staticmethod + def supported_targets( + design: DesignMetadata | None = None, + operands: OperandsMetadata | None = None, + ) -> list[TargetSm]: + """Return all the supported targets for the MMA instruction, given optional design and operands metadata.""" + raise NotImplementedError + + @staticmethod + def shape_k( + operands: OperandsMetadata, + ) -> int: + """Return the K dimension of the shape of matrix multiplication performed by this instruction.""" + raise NotImplementedError + + + +class BlackwellTcgen05Mma(MmaInstruction): + """Models the tcgen05.mma instruction for the Blackwell architecture.""" + + @staticmethod + def supported_targets( + design: DesignMetadata | None = None, + operands: OperandsMetadata | None = None, + ) -> list[TargetSm]: + """See :meth:`MmaInstruction.supported_targets`.""" + # See: + # https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#release-notes-a-spec-f-spec-ptx-feature-release-history + + # Most tcgen05.mma instructions are supported on the full family + default_targets = [TargetSm("100f"), TargetSm("110f")] + + if operands is None or design is None: + return default_targets + + A = getattr(operands, "A", None) + tile_shape = getattr(design, "tile_shape", None) + if A is None: + raise NotImplementedError( + f"Could not determine MMA operand A for {type(operands)}" + ) + if tile_shape is None: + raise NotImplementedError( + f"Could not determine MMA tile shape for {type(design)}" + ) + + if A.dtype in [cutlass.Int8, cutlass.Uint8]: + # Only 100a, 110a support int types with tcgen05.mma + return [TargetSm("100a"), TargetSm("110a")] + + # tile K == 96 is only supported on sm103a + if tile_shape[2] == 96: + return [TargetSm("103a")] + + return default_targets + + @staticmethod + def shape_k( + operands: OperandsMetadata, + ) -> int: + """See :meth:`MmaInstruction.shape_k`.""" + try: + A = operands.A + B = operands.B + except AttributeError as err: + raise AttributeError( + f"{type(operands)} does not have attributes for operands `A` and `B`" + ) from err + + if A.dtype in [cutlass.Float32, cutlass.TFloat32]: + return 8 + elif A.dtype in [cutlass.Float16, cutlass.BFloat16]: + return 16 + elif A.dtype in [cutlass.Int8, cutlass.Uint8]: + return 32 + elif A.dtype == cutlass.Float4E2M1FN and B.dtype == cutlass.Float4E2M1FN: + return 64 + elif A.dtype in [ + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Float6E3M2FN, + cutlass.Float6E2M3FN, + cutlass.Float4E2M1FN, + ]: + return 32 + else: + raise NotImplementedError(f"Unsupported operand dtype: {A.dtype}") + + +class BlackwellTcgen05MmaSparse(MmaInstruction): + """Models the tcgen05.mma.sp instruction for the Blackwell architecture.""" + + @staticmethod + def supported_targets( + design: DesignMetadata | None = None, + operands: OperandsMetadata | None = None, + ) -> list[TargetSm]: + """See :meth:`MmaInstruction.supported_targets`.""" + raise NotImplementedError + + +class HopperWgmma(MmaInstruction): + """Models the wgmma Tensor Core MMA instruction for the Hopper architecture.""" + + @staticmethod + def supported_targets( + design: DesignMetadata | None = None, + operands: OperandsMetadata | None = None, + ) -> list[TargetSm]: + """See :meth:`MmaInstruction.supported_targets`.""" + return [TargetSm("90a")] + + +class AmpereMma(MmaInstruction): + """Models the Tensor Core MMA instruction for the Ampere architecture.""" + + @staticmethod + def supported_targets( + design: DesignMetadata | None = None, + operands: OperandsMetadata | None = None, + ) -> list[TargetSm]: + """See :meth:`MmaInstruction.supported_targets`.""" + return [TargetSm(cc=80)] diff --git a/operators/cutlass/operators/providers/__init__.py b/operators/cutlass/operators/providers/__init__.py new file mode 100644 index 000000000..0e2592917 --- /dev/null +++ b/operators/cutlass/operators/providers/__init__.py @@ -0,0 +1,99 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import contextlib +import logging +from collections.abc import Callable + +from cutlass.operators.providers.provider import Provider + +_logger = logging.getLogger(__name__) + +available_providers: dict[str, type[Provider]] = {} +"""Dictionary of available providers by name.""" + + +def register_provider(name: str) -> Callable[[type[Provider]], type[Provider]]: + """ + Decorator used to register a provider class with the given name. + + Args: + name (str): The name of the provider + + Returns: + Callable[[Type[Provider]], Type[Provider]]: The wrapper function + """ + + def wrapper(provider_class: type[Provider]) -> type[Provider]: + """ + Wrapper function to register a provider class with the given name. + + Args: + provider_class (Type[Provider]): The provider class to register + + Returns: + Type[Provider]: The registered provider class + """ + global available_providers + available_providers[name] = provider_class + return provider_class + + return wrapper + + +# Import each provider submodule's ``*Provider`` class. The submodule +# registers its class as a side effect when its underlying DSL is importable +# and otherwise leaves the class undefined + +with contextlib.suppress(ImportError): + from cutlass.operators.providers.cutedsl import CuTeDSLProvider + + +def __getattr__(name: str) -> object: + """Give a clear ImportError for known but unavailable providers. + + For unknown attributes, raise the generic AttributeError. + """ + PROVIDER_INSTALL_HINTS = { + "CuTeDSLProvider": "install 'nvidia-cutlass-dsl' to enable it", + } + if name in PROVIDER_INSTALL_HINTS: + raise ImportError( + f"{name} is unavailable: " + f"{PROVIDER_INSTALL_HINTS[name]}. Check " + f"`available_providers` for runtime availability." + ) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +__all__ = [ + "available_providers", + "register_provider", + "Provider", + "CuTeDSLProvider", +] diff --git a/operators/cutlass/operators/providers/cutedsl/__init__.py b/operators/cutlass/operators/providers/cutedsl/__init__.py new file mode 100644 index 000000000..cc4a75f19 --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/__init__.py @@ -0,0 +1,66 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import logging + +from cutlass.operators.providers import register_provider +from cutlass.operators.providers.provider import Provider + + +def try_import() -> bool: + import importlib.util + + if importlib.util.find_spec("cutlass") is None: + logging.warning( + "'cutlass' could not be imported. " + "The cutedsl provider will not be available." + ) + return False + return True + + +available = try_import() + + +if available: + + @register_provider("cutedsl") + class CuTeDSLProvider(Provider): + """ + Provider for CuTe DSL (CUTLASS 4.x) based operators. + + CuTe DSL is a Python DSL fully isomorphic with CuTe C++ that compiles + to optimized CUDA kernels via MLIR and ptxas. + """ + + # Imports for side effects (Operator registration) + import cutlass.operators.providers.cutedsl.gemm # noqa: F401 + + __all__ = ["CuTeDSLProvider"] +else: + __all__ = [] diff --git a/operators/cutlass/operators/providers/cutedsl/evt/common_efc.py b/operators/cutlass/operators/providers/cutedsl/evt/common_efc.py new file mode 100644 index 000000000..2a146f9df --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/evt/common_efc.py @@ -0,0 +1,2189 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +# This is not to use module annotations from a future version but to change the type system to postpone the evaluation of annotations, +# about forward declaration and lazy type checking. +# See https://docs.python.org/3/library/__future__.html#future__.annotations and https://peps.python.org/pep-0563/. +from __future__ import annotations + +import dataclasses +import enum +import functools +import inspect +import logging +import os +import types +import typing + +import cutlass + +""" +CUTLASS EFC Framework +""" + +# List of available activation functions in the EFC Configuration class. This +# can be used for example in a CLI using this module. +ACTIVATION_FUNCTIONS = [ + "identity", + "relu", + "leaky_relu", + "tanh", + "sigmoid", + "silu", + "hardswish", + "gelu", +] + +# Some workaround for Python pre-3.12 f-string preventing the use of backslash. +NEWLINE_AND_TAB = "\n\t" +TAB = "\t" + +# To have some verbosity, set the CUTE_DSL_EFC_LOG_LEVEL environment variable to +# INFO or even DEBUG before launching this program. +if log_level := os.environ.get("CUTE_DSL_EFC_LOG_LEVEL", None): + logging.basicConfig(level=getattr(logging, log_level)) +logger = logging.getLogger(__name__) + + +def log(message: str): + """Helper function to log messages. Change logger.debug to another level + here if needed. + """ + logger.debug(message) + + +def if_debug(function): + """Execute a function if in debug mode.""" + if logger.isEnabledFor(logging.DEBUG): + function() + + +def if_info(function): + """Execute a function if in info mode.""" + if logger.isEnabledFor(logging.INFO): + function() + + +def mark_mlir(message: str): + """Put some message in MLIR output to make MLIR assembly clearer or trace execution.""" + if_debug(lambda: cutlass.cute.printf(f"mark_mlir: {message}")) + + +def trace_in_mlir(func): + """Decorator to trace function entry and exit in MLIR.""" + + @functools.wraps(func) + def wrapper(*args, **kwargs): + function_name = func.__name__ + mark_mlir(f"entering {function_name}") + result = func(*args, **kwargs) + mark_mlir(f"leaving {function_name}") + return result + + return wrapper + + +def create_named_epilogue(param_names, func): + """Create a wrapper function with specific parameter names that delegates to an implementation function. + + This function solves a common problem in the EFC (Epilogue Fusion Configuration) framework: + epilogue functions must have parameters with specific names (e.g., "alpha", "beta", "C", "D") + to match the EFC calling convention, but you may want to generate these functions + programmatically from a generic implementation. + + Instead of using string manipulation with exec() or eval() (which is insecure and breaks + tooling), this function uses Python's inspect module to create a proper function signature + that tools like debuggers, type checkers, and IDEs can understand. + + Args: + param_names: List of parameter names for the generated function + (e.g., ["alpha", "beta", "C", "x_factor"]) + func: Implementation function that accepts the same number of arguments as param_names. + The arguments will be passed in the order specified by param_names. + + Returns: + A new function with the specified parameter names that calls func with those + parameters in order. The wrapper preserves func's name and docstring, and + has a proper signature for introspection. + + Example: + # Generic implementation that doesn't care about parameter names + def compute(a, b, c): + return a + b * c + + # Create EFC-compliant function with required parameter names + epilogue = create_named_epilogue(["alpha", "X", "Y"], compute) + # Now epilogue(alpha=1, X=2, Y=3) calls compute(1, 2, 3) + # and inspect.signature(epilogue) shows the correct parameter names + + Use Case: + When programmatically generating epilogue functions with different tensor + configurations, you need each function to have the right parameter names + for the EFC framework to call them correctly with keyword arguments. + + """ + # Create Parameter objects for each parameter name, using standard Python argument binding. + parameters = [ + inspect.Parameter(name, inspect.Parameter.POSITIONAL_OR_KEYWORD) + for name in param_names + ] + + # Create a new signature with the custom parameter names + new_signature = inspect.Signature(parameters) + + # Create a wrapper function that accepts arguments according to the new signature + def wrapper(*args, **kwargs): + # Bind the provided arguments to our custom signature + bound = new_signature.bind(*args, **kwargs) + bound.apply_defaults() + + # Extract argument values in the order specified by param_names + ordered_args = [bound.arguments[name] for name in param_names] + + # Call the original function with the properly ordered arguments + return func(*ordered_args) + + # Assign the custom signature to the wrapper so introspection works correctly + wrapper.__signature__ = new_signature + wrapper.__name__ = getattr(func, "__name__", "generated_function") + wrapper.__doc__ = func.__doc__ + + return wrapper + + +def _validate_source_mode_map(source_mode_map, source_rank, target_rank): + """Validate the source_mode_map for remap_modes() and torch_tensor_remap_modes(). + + Checks that: + - Every element is either ``True`` (broadcast) or an ``int`` (source mode index). + - The map length matches the target rank. + - The number of non-broadcast entries does not exceed the source rank + (some source dims may be unused, e.g. size-1 dims that are implicitly + broadcast). + - No integer index exceeds the source rank. + - No integer index is negative. + - No integer index appears more than once. + """ + for i, v in enumerate(source_mode_map): + if v is not True and not isinstance(v, int): + raise TypeError( + f"source_mode_map[{i}] must be an int or True, got {type(v).__name__}: {v!r}" + ) + if target_rank != len(source_mode_map): + raise ValueError( + f"The target layout rank ({target_rank}) and the source_mode_map " + f"length ({len(source_mode_map)}) should be the same." + ) + # The number of non-True entries must not exceed source_rank — but can be + # less, allowing some source dims to be unused (e.g. size-1 broadcast dims + # in a tensor like (1, N) where dim 0 is ignored). + num_true = sum(v is True for v in source_mode_map) + num_used = len(source_mode_map) - num_true + if num_used > source_rank: + raise ValueError( + f"source_mode_map references {num_used} source dimensions but " + f"the source only has {source_rank}. " + f"The number of non-True entries must not exceed the source rank." + ) + int_indices = [v for v in source_mode_map if v is not True] + for v in int_indices: + if v < 0: + raise ValueError( + f"source_mode_map contains a negative index ({v}). " + f"All mode indices must be non-negative." + ) + if v >= source_rank: + raise ValueError( + f"source_mode_map contains index {v} which exceeds the source " + f"rank ({source_rank}). Valid indices are 0..{source_rank - 1}." + ) + if len(int_indices) != len(set(int_indices)): + duplicates = [v for v in int_indices if int_indices.count(v) > 1] + raise ValueError( + f"source_mode_map contains duplicate source mode indices: " + f"{sorted(set(duplicates))}. Each source mode should appear at most once." + ) + + +def remap_modes( + source: cutlass.cute.Tensor, + target_shape: tuple, + source_mode_map: tuple, + info_tensor_name="", + info_source_name="", +): + """Remap modes of a source tensor to match a target shape. + + The source_mode_map is a sequence with a length matching the rank of + the target shape. + + Each element is either an integer defining the mode of the source + tensor to map at this place or ``True`` to express that the source + tensor is broadcast for this mode. + """ + ss = source.stride + _validate_source_mode_map(source_mode_map, len(ss), len(target_shape)) + stride = tuple([0 if v is True else ss[v] for v in source_mode_map]) + log(f"remap_modes {stride = }") + if_info( + lambda: cutlass.cute.printf( + f"cute remap_modes stride {info_tensor_name} from {info_source_name} {{}} to {{}} (target_shape {target_shape})", + source.stride, + stride, + ) + ) + log( + f"remap_modes shape {info_tensor_name} from {info_source_name} {source.shape} to {target_shape = }" + ) + if_info( + lambda: cutlass.cute.printf( + f"cute remap_modes shape {info_tensor_name} from {info_source_name} {{}} to {target_shape = }", + source.shape, + ) + ) + broadcast_layout = cutlass.cute.make_layout(shape=target_shape, stride=stride) + return cutlass.cute.make_tensor(source.iterator, broadcast_layout) + + +def torch_tensor_remap_modes( + source: torch.Tensor, target_shape: tuple, source_mode_map: tuple +): + """Remap modes of a source tensor to match a target shape. + + This is the PyTorch equivalent of remap_modes() for cutlass.cute.Tensor. + + The source_mode_map is a sequence with a length matching the rank of + the target shape. + + Each element is either an integer defining the mode of the source + tensor to map at this place or ``True`` to express that the source + tensor is broadcast for this mode. + """ + import torch + + ss = source.stride() + _validate_source_mode_map(source_mode_map, len(ss), len(target_shape)) + stride = tuple([0 if v is True else ss[v] for v in source_mode_map]) + log(f"torch_tensor_remap_modes {stride = }") + log(f"torch_tensor_remap_modes {target_shape = }") + return torch.as_strided(source, size=target_shape, stride=stride) + + +class VariadicParameters: + """Minimal mixin wrapper for variadic parameters for @cute.jit/@cute.kernel + functions taking advantage that the DSL can ingest a recursive + combination of tuples and lists. + """ + + class _UnassignedArgument: + """Sentinel class to detect uninitialized arguments.""" + + def __repr__(self): + return "" + + def __init__(self, efc, parameter_names): + # Add local shortcuts to the efc and gemm objects + self.efc = efc + self.gemm = efc.gemm + self._create_parameter_class(parameter_names) + self._init_args_with_unassigned() + log(f"Initial {self.arg = }") + + def _create_parameter_class(self, parameter_names): + """Create a dataclass with the given parameter names. + + The dataclass uses __slots__ to avoid assigning wrong arguments by error + and has __getitem__/__setitem__ methods for dict-like access. + """ + fields = [(name, typing.Any) for name in parameter_names] + self._parameter_class = dataclasses.make_dataclass( + "Parameter", fields, slots=True + ) + self._add_item_accessors() + + def _add_item_accessors(self): + """Add __getitem__ and __setitem__ methods to the parameter class. + + This allows accessing attributes by name, e.g., arg["a"] or parameter["b"]. + """ + + def getitem(self, name): + """Access the dataclass attribute by name.""" + return getattr(self, name) + + def setitem(self, name, value): + """Set the dataclass attribute by name.""" + setattr(self, name, value) + + self._parameter_class.__getitem__ = getitem + self._parameter_class.__setitem__ = setitem + + def _init_args_with_unassigned(self): + """Initialize all arg attributes to _UnassignedArgument sentinel. + + This ensures any forgotten field will trigger a noticeable error. + """ + self.arg = self._parameter_class( + *( + [self._UnassignedArgument] + * len(dataclasses.fields(self._parameter_class)) + ) + ) + + def pack_arguments(self, *args, **kwargs): + """Pack the arguments as a tuple to pass them through a + @cute.jit/@cute.kernel call. + + If some arguments are provided, pack them and set self.arg, otherwise + just use the self.arg object by default. + + Return a tuple or None as an interface object since a + @cute.jit/@cute.kernel can ingest a recursive combination of tuples and + lists. + """ + if args or kwargs: + # Override the current argument object from the provided arguments, if any. + self.arg = self._parameter_class(*args, **kwargs) + # dataclasses.astuple(self.arg) breaks because it is recursive and + # applies a deepcopy incompatible with the DSL magic. Just generate 1 + # level of tuple of object references. + r = tuple(self.arg[f.name] for f in dataclasses.fields(self.arg)) + log( + f"pack_arguments {args = }{NEWLINE_AND_TAB}{kwargs = }{NEWLINE_AND_TAB}{self.arg = }{NEWLINE_AND_TAB}{r = }" + ) + # The DSL does not accept an empty tuple but can handle None. So + # remap to None in that case. + if not r: + return None + return r + + def unpack_parameters(self, p: tuple): + """Unpack the parameters inside a @cute.jit/@cute.kernel function. + + Assign all the self.parameter attributes. + """ + # Do the opposite mapping of None to an empty tuple to have the + # parameter constructor happy. + if p is None: + p = () + # Instantiate the dataclass holding the parameters from the + # individual parameter values. + self.parameter = self._parameter_class(*p) + log(f"unpack_parameters {p = }{NEWLINE_AND_TAB}{self.parameter = }") + + def instantiate_args(self): + """Create an arg attribute from the Parameter class to be used + as an alternative way to pass the arguments instead of using an + explicit pack_arguments(). + + All the arg attributes are initialized to a noticeable name so that + any forgotten field will trigger an error. + """ + self._init_args_with_unassigned() + + def extend_parameters(self, names): + """Extend the dataclass with new field names. + + This creates a new dataclass that inherits all existing fields and adds + the new ones. The existing arg/parameter values are preserved and new + fields are initialized to _UnassignedArgument. + + Args: + names: List of new parameter names to add to the dataclass. + """ + if not names: + return + + # Get existing field names and values from arg. + existing_fields = [f.name for f in dataclasses.fields(self._parameter_class)] + existing_arg_values = [self.arg[name] for name in existing_fields] + + # Also get existing parameter values if parameter exists. + has_parameter = hasattr(self, "parameter") + if has_parameter: + existing_param_values = [self.parameter[name] for name in existing_fields] + + # Create new dataclass with combined fields. + self._create_parameter_class(existing_fields + list(names)) + log(f"Extended parameters {dataclasses.fields(self._parameter_class) = }") + + # Reconstruct the arg object with existing values and new unassigned + # fields. + new_arg_values = existing_arg_values + [self._UnassignedArgument] * len(names) + log(f"Extended parameters {new_arg_values = }") + self.arg = self._parameter_class(*new_arg_values) + + # Also reconstruct parameter if it existed. + if has_parameter: + new_param_values = existing_param_values + [self._UnassignedArgument] * len( + names + ) + self.parameter = self._parameter_class(*new_param_values) + + log(f"Extended parameters with {names}, now {self.arg = }") + + +class EFC: + """Epilogue Fusion Configuration.""" + + # Helper functions for CuTe operations + @staticmethod + def maximum(x, y): + """Element-wise maximum of 2 CuTe tensors""" + x_type = x.element_type + y_type = y.element_type + if x_type is not y_type: + raise TypeError(f"Type mismatch: x is {x_type}, y is {y_type}") + return cutlass.cute.where(x > y, x, y) + + @staticmethod + def minimum(x, y): + """Element-wise minimum of 2 CuTe tensors""" + x_type = x.element_type + y_type = y.element_type + if x_type is not y_type: + raise TypeError(f"Type mismatch: x is {x_type}, y is {y_type}") + return cutlass.cute.where(x < y, x, y) + + class JIT(VariadicParameters): + """Handle Python/@cute.jit and its boundaries with Host.""" + + # All the following customization functions should go somewhere else in + # the long term, as part of a refactoring similar to CUTLASS + # collective/main loop/epilogue... + + # The following foreach_... functions are some refinement of the EFC + # ones to handle more precisely the broadcast tensors for example. + + def get_argument_or_remapped_tensor(self, name): + """A remapped tensor is not passed as a parameter but synthesized and stored in self.remapped_tensors.""" + if tensor := self.remapped_tensors.get(name): + return tensor + return self.parameter[name] + + def get_remapped_tensor_or_itself(self, name): + """For a remapped tensor name, return the remapped tensor view, otherwise return the tensor itself.""" + if self.efc.parameter_attributes[name].mapped_source: + return self.remapped_tensors[name] + return self.parameter[name] + + def foreach_tensor(self, function): + """Execute the given function for each supplemental tensor, + providing the tensor object from the JIT parameters. + """ + + def wrapper(tensor_name, attributes): + function( + self.get_argument_or_remapped_tensor(tensor_name), + tensor_name, + attributes, + ) + + self.efc.foreach_tensor(wrapper) + + def foreach_non_degenerate_tensor(self, function): + """Execute the given function for each supplemental tensor, + providing the tensor object from the JIT parameters. Skip the tensor + if it is a degenerate case as scalar in broadcast. + """ + + def wrapper(tensor_name, attributes): + if not attributes.degenerate_to_scalar: + function( + self.get_argument_or_remapped_tensor(tensor_name), + tensor_name, + attributes, + ) + + self.efc.foreach_tensor(wrapper) + + def foreach_remapped_tensor(self, function): + """Execute the given function for each remapped tensor, + providing the remapped tensor object if any. + """ + + def wrapper(tensor_name, attributes): + if not attributes.degenerate_to_scalar: + # Call the function only if it ends up as a real tensor. + function( + self.remapped_tensors[tensor_name], tensor_name, attributes + ) + + self.efc.foreach_remapped_tensor(wrapper) + + def foreach_remapped_tensor_as_scalar(self, function): + """Execute the given function for each remapped tensor, providing + the remapped tensor object only if it degenerates as a scalar + because all the modes are broadcast. + """ + + def wrapper(tensor_name, attributes): + if attributes.degenerate_to_scalar: + # Call the function only if it ends up as a scalar. + function( + self.remapped_tensors[tensor_name], tensor_name, attributes + ) + + self.efc.foreach_remapped_tensor(wrapper) + + def foreach_argument(self, function): + """Execute the given function for each supplemental argument, + providing the parameter value from the JIT parameters. + + Note that remapped tensors are not direct parameters but are + synthesized and accessed via self.remapped_tensors. + """ + + def wrapper(name, attributes): + log(f"JIT.foreach_argument {name = }") + function( + self.get_argument_or_remapped_tensor(name), + name, + attributes, + ) + + self.efc.foreach_argument(wrapper) + + @trace_in_mlir + def handle_remapping(self): + # Track the remapped tensors by their names. + self.efc.jit.remapped_tensors = {} + + def f(tensor_name, attributes): + log(f"handle_remapping of tensor {tensor_name}") + log(f"{TAB}mapped_source {attributes.mapped_source}") + log(f"{TAB}source_mode_map {attributes.source_mode_map}") + self.efc.jit.remapped_tensors[tensor_name] = remap_modes( + source=self.efc.jit.parameter[attributes.mapped_source], + target_shape=self.efc.output_shape, + source_mode_map=attributes.source_mode_map, + info_tensor_name=tensor_name, + info_source_name=attributes.mapped_source, + ) + log( + f"{TAB}self.efc.jit.remapped_tensors[{tensor_name}] = {self.efc.jit.remapped_tensors[tensor_name]}" + ) + + self.efc.foreach_remapped_tensor(f) + + @trace_in_mlir + def record_tensor_dtypes(self): + """It does not seem that the tma_tensor and tma_atom carry over the + element type, so, store it here for later use. + """ + self.tensor_dtype = {} + + def f(tensor, tensor_name, attributes): + self.tensor_dtype[tensor_name] = tensor.element_type + + self.foreach_tensor(f) + + @trace_in_mlir + def written_tensor_name_with_bigger_element_type(self): + """The type of the written tensor is used to compute a lot of + implementation details about tiling and so on in the kernel. + + The compilation phase has already checked that there is at least 1 + written tensor name. + + Return the name of the written tensor with the biggest + element_type. + + """ + return max( + (tensor_name for tensor_name in self.efc.written_tensor_names), + key=lambda tensor_name: self.tensor_dtype[tensor_name].width, + ) + + @trace_in_mlir + def read_tensor_name_with_bigger_element_type(self): + """The type of the read tensor is used to compute a lot of + implementation details about tiling and so on in the kernel. Return + the name of the read tensor with the biggest element_type, or None + if there is no read tensor. + """ + if self.efc.read_tensor_names: + return max( + (tensor_name for tensor_name in self.efc.read_tensor_names), + key=lambda tensor_name: self.tensor_dtype[tensor_name].width, + ) + return None + + @trace_in_mlir + def compute_stage(self): + """Get the contribution from the tensors used in the EFC to the + pipeline stage numbers. + """ + self.smem_size_of_read_tensors = 0 + self.smem_size_of_written_tensors = 0 + self.tensor_dtype = {} + + def f(tensor, tensor_name, attributes): + # It does not seem that the tma_tensor and tma_atom carry over + # the element type, so, store it here for later use. + self.tensor_dtype[tensor_name] = tensor.element_type + + smem_size_in_bytes_of_a_pipeline_stage = cutlass.cute.size_in_bytes( + tensor.element_type, + cutlass.utils.blackwell_helpers.make_smem_layout_epi( + tensor.element_type, self.gemm.d_layout, self.gemm.epi_tile, 1 + ), + ) + # Prepare the information to be asked soon, to memoize this + # loop. + if attributes.is_read: + self.smem_size_of_read_tensors += ( + smem_size_in_bytes_of_a_pipeline_stage + ) + log( + f"JIT.compute_stage smem storage to read tensor {tensor_name}: {smem_size_in_bytes_of_a_pipeline_stage} bytes" + ) + + if attributes.is_written: + self.smem_size_of_written_tensors += ( + smem_size_in_bytes_of_a_pipeline_stage + ) + log( + f"JIT.compute_stage smem storage to write tensor {tensor_name}: {smem_size_in_bytes_of_a_pipeline_stage} bytes" + ) + + self.foreach_non_degenerate_tensor(f) + + @trace_in_mlir + def smem_size_in_bytes_of_read_tensors(self): + """Get the contribution in a smem pipeline stage from the tensors + loaded in the EFC. + """ + log( + f"JIT.smem_size_in_bytes_of_read_tensors {self.smem_size_of_read_tensors = }" + ) + return self.smem_size_of_read_tensors + + @trace_in_mlir + def smem_size_in_bytes_of_written_tensors(self): + """Get the contribution in a smem pipeline stage from the tensors + stored in the EFC. + """ + log( + f"JIT.smem_size_in_bytes_of_written_tensors {self.smem_size_of_written_tensors = }" + ) + return self.smem_size_of_written_tensors + + @trace_in_mlir + def smem_layout(self): + """Get the smem_layout for the tensors used in the EFC.""" + self.smem_layout_staged_read = {} + self.smem_layout_staged_written = {} + + def f(tensor, tensor_name, attributes): + log(f"JIT.smem_layout {tensor_name} = {tensor!s}") + ref_tensor = self.get_argument_or_remapped_tensor(tensor_name) + tensor_layout = cutlass.utils.LayoutEnum.from_tensor(ref_tensor) + log(f"JIT.smem_layout tensor_layout[{tensor_name}] = {tensor_layout!s}") + + if attributes.is_read: + self.smem_layout_staged_read[tensor_name] = ( + cutlass.utils.blackwell_helpers.make_smem_layout_epi( + tensor.element_type, + tensor_layout, + self.gemm.epi_tile, + self.gemm.num_c_stage, + ) + ) + log(f"JIT.smem_layout read {self.gemm.num_c_stage = }") + log( + f"JIT.smem_layout read self.smem_layout_staged_read[{tensor_name}] = {self.smem_layout_staged_read[tensor_name]!s}" + ) + if attributes.is_written: + self.smem_layout_staged_written[tensor_name] = ( + cutlass.utils.blackwell_helpers.make_smem_layout_epi( + tensor.element_type, + tensor_layout, + self.gemm.epi_tile, + self.gemm.num_d_stage, + ) + ) + log(f"JIT.smem_layout written {self.gemm.num_d_stage = }") + log( + f"JIT.smem_layout written self.smem_layout_staged_written[{tensor_name}] = {self.smem_layout_staged_written[tensor_name]!s}" + ) + + self.foreach_non_degenerate_tensor(f) + + @trace_in_mlir + def create_tma_arguments(self): + """Set the TMA related arguments for the tensors used in the EFC.""" + # Make the difference for read/written to handle the case a tensor + # is both read and written. + self.total_tma_load_bytes = 0 # Used by the PipelineTmaAsync + self.tma_atom_read = {} + self.tma_tensor_read = {} + self.tma_atom_written = {} + self.tma_tensor_written = {} + + def f(tensor, tensor_name, attributes): + if attributes.is_read: + smem_layout = cutlass.cute.slice_( + self.smem_layout_staged_read[tensor_name], (None, None, 0) + ) + self.total_tma_load_bytes += cutlass.cute.size_in_bytes( + tensor.element_type, smem_layout + ) + ( + self.tma_atom_read[tensor_name], + self.tma_tensor_read[tensor_name], + ) = cutlass.cute.nvgpu.cpasync.make_tiled_tma_atom( + cutlass.cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(), + tensor, + smem_layout, + self.gemm.epi_tile, + ) + log( + f"JIT.create_tma_arguments tma_atom_read[{tensor_name}] = {self.tma_atom_read[tensor_name]!s}" + ) + log( + f"JIT.create_tma_arguments tma_tensor_read[{tensor_name}] = {self.tma_tensor_read[tensor_name]!s}" + ) + + if attributes.is_written: + smem_layout = cutlass.cute.slice_( + self.smem_layout_staged_written[tensor_name], (None, None, 0) + ) + ( + self.tma_atom_written[tensor_name], + self.tma_tensor_written[tensor_name], + ) = cutlass.cute.nvgpu.cpasync.make_tiled_tma_atom( + cutlass.cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp(), + tensor, + smem_layout, + self.gemm.epi_tile, + ) + log( + f"JIT.create_tma_arguments tma_atom_written[{tensor_name}] = {self.tma_atom_written[tensor_name]!s}" + ) + log( + f"JIT.create_tma_arguments tma_tensor_written[{tensor_name}] = {self.tma_tensor_written[tensor_name]!s}" + ) + + self.foreach_non_degenerate_tensor(f) + + @trace_in_mlir + def create_supplemental_arguments_for_kernel(self): + """Executed before launching the @cute.kernel function to set up the + supplemental arguments to pass to the @cute.kernel function. + + In the @cute.kernel example, the parameters like `X_tma_tensor_read` + or `Y_tma_tensor_written` correspond to `mX_mnl` and `mY_mnl`. + """ + argument_names = [] + + def compute_argument_names(value, name, attributes): + if not attributes.is_tensor or attributes.degenerate_to_scalar: + # Just propagate the dynamic scalar with the same name or a + # scalar parameter for the degenerate tensor broadcast. + argument_names.append(name) + else: + if attributes.is_read: + argument_names.append(f"{name}_tma_atom_read") + argument_names.append(f"{name}_tma_tensor_read") + argument_names.append(f"{name}_smem_layout_staged_read") + if attributes.is_written: + argument_names.append(f"{name}_tma_atom_written") + argument_names.append(f"{name}_tma_tensor_written") + argument_names.append(f"{name}_smem_layout_staged_written") + + self.foreach_argument(compute_argument_names) + # Create the @cute.kernel-side meta-programming infrastructure + # handling also the supplemental argument handling. + self.efc.kernel = EFC.Kernel(self.efc, argument_names) + + arg = self.efc.kernel.arg + + def populate_the_kernel_arguments(value, name, attributes): + if not attributes.is_tensor: + # Just propagate the dynamic scalar with the same name. + arg[name] = value + elif attributes.degenerate_to_scalar: + # The tensor broadcast is optimized out as a single scalar + # from the tensor source which is implicitly broadcast when + # used but since it might be in device memory, pass it + # unchanged so it can be dereferenced inside the kernel. + arg[name] = self.efc.jit.parameter[attributes.mapped_source] + else: + if attributes.is_read: + arg[f"{name}_tma_atom_read"] = self.tma_atom_read[name] + arg[f"{name}_tma_tensor_read"] = self.tma_tensor_read[name] + arg[f"{name}_smem_layout_staged_read"] = ( + self.smem_layout_staged_read[name] + ) + if attributes.is_written: + arg[f"{name}_tma_atom_written"] = self.tma_atom_written[name] + arg[f"{name}_tma_tensor_written"] = self.tma_tensor_written[ + name + ] + arg[f"{name}_smem_layout_staged_written"] = ( + self.smem_layout_staged_written[name] + ) + + self.foreach_argument(populate_the_kernel_arguments) + + class Kernel(VariadicParameters): + """Handle kernel part and @cute.jit/@cute.kernel boundaries.""" + + def foreach_tensor(self, function): + """Execute the given function for each supplemental tensor, skipping + the degenerate scalar broadcast cases. + """ + + def wrapper(tensor_name, attributes): + if not attributes.degenerate_to_scalar: + function( + tensor_name, + attributes, + ) + + self.efc.foreach_tensor(wrapper) + + def foreach_read_tensor(self, function): + """Execute the given function for each supplemental read tensor, + skipping the degenerate scalar broadcast cases. + """ + + def wrapper(tensor_name, attributes): + if not attributes.degenerate_to_scalar: + function( + tensor_name, + attributes, + ) + + self.efc.foreach_read_tensor(wrapper) + + def foreach_written_tensor(self, function): + """Execute the given function for each supplemental written tensor, + skipping the degenerate scalar broadcast cases. + """ + + def wrapper(tensor_name, attributes): + if not attributes.degenerate_to_scalar: + function( + tensor_name, + attributes, + ) + + self.efc.foreach_written_tensor(wrapper) + + @trace_in_mlir + def prefetch_tma_descriptors(self): + """Prefetch the TMA descriptors for the tensors used in the EFC.""" + + def f(tensor_name, attributes): + if attributes.is_read: + cutlass.cute.nvgpu.cpasync.prefetch_descriptor( + self.parameter[f"{tensor_name}_tma_atom_read"] + ) + + if attributes.is_written: + cutlass.cute.nvgpu.cpasync.prefetch_descriptor( + self.parameter[f"{tensor_name}_tma_atom_written"] + ) + + self.foreach_tensor(f) + + @trace_in_mlir + def allocate_smem(self): + """Allocate the shared memory for all the supplemental tensors.""" + self.smem_read = {} + self.smem_written = {} + + def f(tensor_name, attributes): + element_type = self.efc.jit.tensor_dtype[tensor_name] + if attributes.is_read: + smem_layout_staged = self.parameter[ + f"{tensor_name}_smem_layout_staged_read" + ] + self.smem_read[tensor_name] = self.gemm.smem.allocate_tensor( + element_type=element_type, + layout=smem_layout_staged.outer, + byte_alignment=self.gemm.buffer_align_bytes, + swizzle=smem_layout_staged.inner, + ) + if attributes.is_written: + smem_layout_staged = self.parameter[ + f"{tensor_name}_smem_layout_staged_written" + ] + self.smem_written[tensor_name] = self.gemm.smem.allocate_tensor( + element_type=element_type, + layout=smem_layout_staged.outer, + byte_alignment=self.gemm.buffer_align_bytes, + swizzle=smem_layout_staged.inner, + ) + + self.foreach_tensor(f) + + @trace_in_mlir + def partition_global_tensors_for_tiled_mma(self): + """Partition the global supplemental tensors for TiledMMA_C/D.""" + self.tCgC_read = {} + self.tCgD_written = {} + + def f(tensor_name, attributes): + if attributes.is_read: + # (bM, bN, loopM, loopN, loopL) + gC_mnl = cutlass.cute.local_tile( + self.parameter[f"{tensor_name}_tma_tensor_read"], + cutlass.cute.slice_(self.gemm.mma_tiler, (None, None, 0)), + (None, None, None), + ) + log( + f"Kernel.partition_global_tensors_for_tiled_mma: gC_mnl[{tensor_name}] = {gC_mnl!s}" + ) + # (MMA, MMA_M, MMA_N, loopM, loopN, loopL) + self.tCgC_read[tensor_name] = self.gemm.thr_mma.partition_C(gC_mnl) + log( + f"Kernel.partition_global_tensors_for_tiled_mma: self.tCgC_read[{tensor_name}] = {self.tCgC_read[tensor_name]!s}" + ) + + if attributes.is_written: + # (bM, bN, loopM, loopN, loopL) + gD_mnl = cutlass.cute.local_tile( + self.parameter[f"{tensor_name}_tma_tensor_written"], + cutlass.cute.slice_(self.gemm.mma_tiler, (None, None, 0)), + (None, None, None), + ) + log( + f"Kernel.partition_global_tensors_for_tiled_mma: gD_mnl[{tensor_name}] = {gD_mnl!s}" + ) + # (MMA, MMA_M, MMA_N, loopM, loopN, loopL) + self.tCgD_written[tensor_name] = self.gemm.thr_mma.partition_C( + gD_mnl + ) + log( + f"Kernel.partition_global_tensors_for_tiled_mma: self.tCgD_written[{tensor_name}] = {self.tCgD_written[tensor_name]!s}" + ) + + self.foreach_tensor(f) + + # The following functions are executed by the specialized warps for + # epilogue computation. + + @trace_in_mlir + def copy_and_partition_supplemental_rmem_tensors( + self, tiled_copy_t2r, tTR_rAcc, epi_tidx, epi_tile + ): + # Load tensor. + self.tiled_copy_s2r = {} + self.tSR_rC = {} + self.tSR_sC = {} + + # Store tensor. + self.tiled_copy_r2s = {} + self.tRS_rD = {} + self.tRS_sD = {} + self.bSG_sD = {} # ((ATOM_V, REST_V), EPI_M, EPI_N) + self.bSG_gD_partitioned = {} # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) + + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors: tiled_copy_t2r = {tiled_copy_t2r!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors: tTR_rAcc = {tTR_rAcc!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors: epi_tile = {epi_tile!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors: epi_tidx = {epi_tidx!s}" + ) + + def f(tensor_name, attributes): + element_type = self.efc.jit.tensor_dtype[tensor_name] + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors: element_type[{tensor_name}] = {element_type!s}" + ) + + if attributes.is_read: + tTR_rC = cutlass.cute.make_rmem_tensor(tTR_rAcc.shape, element_type) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors read: tTR_rC[{tensor_name}] = {tTR_rC!s}" + ) + + ( + self.tiled_copy_s2r[tensor_name], + self.tSR_rC[tensor_name], + self.tSR_sC[tensor_name], + ) = self.gemm.epilogue_smem_copy_and_partition_load( + tiled_copy_t2r, + tTR_rC, + epi_tidx, + self.smem_read[tensor_name], + ) + + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors read: self.tiled_copy_s2r[{tensor_name}] = {self.tiled_copy_s2r[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors read: self.tSR_rC[{tensor_name}] = {self.tSR_rC[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors read: self.tSR_sC[{tensor_name}] = {self.tSR_sC[tensor_name]!s}" + ) + + if attributes.is_written: + # (T2R, T2R_M, T2R_N) + tTR_rD = cutlass.cute.make_rmem_tensor(tTR_rAcc.shape, element_type) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: tTR_rD[{tensor_name}] = {tTR_rD!s}" + ) + + # utils.gemm.sm100.epilogue_smem_copy_and_partition uses + # explicitly "C" as the output matrix and introspects the + # gemm object while in this kernel "C" is used for read but + # "D" is for output according to the BLAS convention. + # So construct a minimal mock-up with the required + # information. + faux_gemm = types.SimpleNamespace() + faux_gemm.c_layout = self.gemm.d_layout + faux_gemm.c_dtype = self.gemm.d_dtype + faux_gemm.acc_dtype = self.gemm.acc_dtype + ( + self.tiled_copy_r2s[tensor_name], + self.tRS_rD[tensor_name], # (R2S, R2S_M, R2S_N) + self.tRS_sD[tensor_name], # (R2S, R2S_M, R2S_N) + ) = cutlass.utils.gemm.sm100.epilogue_smem_copy_and_partition( + faux_gemm, + tiled_copy_t2r, # (EPI_TILE_M, EPI_TILE_N) + tTR_rD, + epi_tidx, + self.smem_written[tensor_name], + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.smem_written[{tensor_name}] = {self.smem_written[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.tiled_copy_r2s[{tensor_name}] = {self.tiled_copy_r2s[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.tRS_rD[{tensor_name}] = {self.tRS_rD[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.tRS_sD[{tensor_name}] = {self.tRS_sD[tensor_name]!s}" + ) + ( + _, + self.bSG_sD[tensor_name], # ((ATOM_V, REST_V), EPI_M, EPI_N) + self.bSG_gD_partitioned[ + tensor_name + ], # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) + ) = self.gemm.epilogue_gmem_copy_and_partition( + epi_tidx, + self.parameter[f"{tensor_name}_tma_atom_written"], + self.tCgD_written[tensor_name], + epi_tile, + self.smem_written[tensor_name], + element_type, + ) + + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.parameter[{tensor_name}_tma_atom_written] = {self.parameter[f'{tensor_name}_tma_atom_written']!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.tCgD_written[{tensor_name}] = {self.tCgD_written[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.smem_written[{tensor_name}] = {self.smem_written[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: element_type = {element_type!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.bSG_sD[{tensor_name}] = {self.bSG_sD[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.bSG_gD_partitioned[{tensor_name}] = {self.bSG_gD_partitioned[tensor_name]!s}" + ) + + self.foreach_tensor(f) + + @trace_in_mlir + def slice_written_tensors_per_mma_tile_index(self, mma_tile_coord_mnl): + """Slice the supplemental written tensors per MMA tile index.""" + self.bSG_gD = {} # ((ATOM_V, REST_V), (EPI_M, EPI_N)) + + def f(tensor_name, attributes): + # ((ATOM_V, REST_V), EPI_M, EPI_N) + bSG_gD = self.bSG_gD_partitioned[tensor_name][ + ( + None, + None, + None, + *mma_tile_coord_mnl, + ) + ] + log( + f"Kernel.slice_written_tensors_per_mma_tile_index: bSG_gD[{tensor_name}] = {bSG_gD!s}" + ) + # Group the 2 last modes so the subtile_idx loop can iterate + # through it using 1-D indexing. + # ((ATOM_V, REST_V), (EPI_M, EPI_N)) + self.bSG_gD[tensor_name] = cutlass.cute.group_modes( + bSG_gD, 1, cutlass.cute.rank(bSG_gD) + ) + log( + f"Kernel.slice_written_tensors_per_mma_tile_index: self.bSG_gD[{tensor_name}] = {self.bSG_gD[tensor_name]!s}" + ) + + self.foreach_written_tensor(f) + + @trace_in_mlir + def load_tensors_from_smem_to_register(self, index): + """Load supplemental tensors from shared memory to register.""" + + def f(tensor_name, attributes): + cutlass.cute.copy( + self.tiled_copy_s2r[tensor_name], + self.tSR_sC[tensor_name][ + ( + None, + None, + None, + index, + ) + ], + self.tSR_rC[tensor_name], + ) + log( + f"Kernel.load_tensors_from_smem_to_register cutlass.cute.copy: self.tiled_copy_s2r[{tensor_name}] = {self.tiled_copy_s2r[tensor_name]!s}" + ) + log( + f"Kernel.load_tensors_from_smem_to_register cutlass.cute.copy: self.tSR_sC[{tensor_name}] = {self.tSR_sC[tensor_name]!s}" + ) + log( + f"Kernel.load_tensors_from_smem_to_register cutlass.cute.copy: self.tSR_rC[{tensor_name}] = {self.tSR_rC[tensor_name]!s}" + ) + + self.foreach_read_tensor(f) + + @trace_in_mlir + def epilogue_computation(self, epilogue_context): + """Execute the EFC epilogue.""" + epilogue_context.load = {} + epilogue_context.store = {} + + def broadcast_degenerate_tensor_load(tensor_name, attributes): + """Create the SSA Tensor for the degenerate tensor broadcast + read. + """ + if attributes.degenerate_to_scalar: + log( + f"Kernel.epilogue_computation broadcast_degenerate_tensor_load: broadcast scalar from {tensor_name} = {self.parameter[tensor_name]!s}" + ) + # If the broadcast tensor degenerates to a scalar, just + # dereference the scalar at index 0 and broadcast it to a + # TensorSSA similar to the accumulator. + epilogue_context.load[tensor_name] = cutlass.cute.full_like( + epilogue_context.acc_vec, + self.parameter[tensor_name][0], + self.efc.gemm.epi_dtype, + ) + + self.efc.foreach_remapped_tensor(broadcast_degenerate_tensor_load) + + def load_setup(tensor_name, attributes): + # Retile the read subtile to fit the accumulator subtile vector + # TV layout. + epilogue_context.load[tensor_name] = ( + self.tiled_copy_r2s[self.gemm.d_name_bigger] + .retile(self.tSR_rC[tensor_name]) + .load() + ) + log( + f"Kernel.epilogue_computation load_setup: {self.tiled_copy_r2s[self.gemm.d_name_bigger] = !s}" + ) + log( + f"Kernel.epilogue_computation load_setup: self.tSR_rC[{tensor_name}] = {self.tSR_rC[tensor_name]!s}" + ) + log( + f"Kernel.epilogue_computation load_setup: self.tiled_copy_r2s[self.gemm.d_name_bigger].retile(self.tSR_rC[{tensor_name}]) = {self.tiled_copy_r2s[self.gemm.d_name_bigger].retile(self.tSR_rC[tensor_name])!s}" + ) + log( + f"Kernel.epilogue_computation load_setup: epilogue_context.load[{tensor_name}] = {epilogue_context.load[tensor_name]!s}" + ) + + self.foreach_read_tensor(load_setup) + + def store_setup(tensor_name, attributes): + epilogue_context.store[tensor_name] = self.tRS_rD[tensor_name] + log( + f"Kernel.epilogue_computation store_setup: epilogue_context.store[{tensor_name}] = {epilogue_context.store[tensor_name]!s}" + ) + + self.foreach_written_tensor(store_setup) + + self.efc.specialized_epilogue(EFC.Phase.ThreadOperation, epilogue_context)() + + @trace_in_mlir + def store_written_tensors_to_smem(self, d_buffer): + """Store the EFC written tensors to shared memory.""" + + def f(tensor_name, attributes): + cutlass.cute.copy( + self.tiled_copy_r2s[tensor_name], + self.tRS_rD[tensor_name], + self.tRS_sD[tensor_name][(None, None, None, d_buffer)], + ) + log( + f"Kernel.store_written_tensors_to_smem cutlass.cute.copy: self.tiled_copy_r2s[{tensor_name}] = {self.tiled_copy_r2s[tensor_name]!s}" + ) + log( + f"Kernel.store_written_tensors_to_smem cutlass.cute.copy: self.tRS_rD[{tensor_name}] = {self.tRS_rD[tensor_name]!s}" + ) + log( + f"Kernel.store_written_tensors_to_smem cutlass.cute.copy: self.tRS_sD[{tensor_name}] = {self.tRS_sD[tensor_name]!s}" + ) + log( + f"Kernel.store_written_tensors_to_smem cutlass.cute.copy: self.tRS_sD[{tensor_name}][(None, None, None, d_buffer)] = {self.tRS_sD[tensor_name][(None, None, None, d_buffer)]!s}" + ) + + self.foreach_written_tensor(f) + + @trace_in_mlir + def tma_store_written_tensors_to_gmem(self, d_buffer, subtile_idx): + """Store with TMA the written EFC tensors to global memory.""" + + def f(tensor_name, attributes): + cutlass.cute.copy( + self.parameter[f"{tensor_name}_tma_atom_written"], + self.bSG_sD[tensor_name][(None, d_buffer)], + self.bSG_gD[tensor_name][(None, subtile_idx)], + ) + log( + f"Kernel.tma_store_written_tensors_to_gmem cutlass.cute.copy: self.parameter[{tensor_name}_tma_atom_written] = {self.parameter[f'{tensor_name}_tma_atom_written']!s}" + ) + log( + f"Kernel.tma_store_written_tensors_to_gmem cutlass.cute.copy: self.bSG_sD[{tensor_name}] = {self.bSG_sD[tensor_name]!s}" + ) + log( + f"Kernel.tma_store_written_tensors_to_gmem cutlass.cute.copy: self.bSG_sD[{tensor_name}][(None, d_buffer)] = {self.bSG_sD[tensor_name][(None, d_buffer)]!s}" + ) + log( + f"Kernel.tma_store_written_tensors_to_gmem cutlass.cute.copy: self.bSG_gD[{tensor_name}] = {self.bSG_gD[tensor_name]!s}" + ) + log( + f"Kernel.tma_store_written_tensors_to_gmem cutlass.cute.copy: self.bSG_gD[{tensor_name}][(None, subtile_idx)] = {self.bSG_gD[tensor_name][(None, subtile_idx)]!s}" + ) + + self.foreach_written_tensor(f) + + # The following functions are executed by the specialized warp for the + # epilogue load. + + @trace_in_mlir + def create_epilogue_subtile_tensors(self, tidx, epi_tile): + """Setup the pipelines reading the EFC supplemental tensors.""" + self.bGS_sC = {} + self.bGS_gC_partitioned = {} + + def f(tensor_name, attributes): + ( + _, + self.bGS_sC[tensor_name], + self.bGS_gC_partitioned[tensor_name], + ) = self.gemm.epilogue_gmem_copy_and_partition( + tidx, + self.parameter[f"{tensor_name}_tma_atom_read"], + self.tCgC_read[tensor_name], + epi_tile, + self.smem_read[tensor_name], + self.efc.jit.tensor_dtype[tensor_name], + ) + + self.foreach_read_tensor(f) + + @trace_in_mlir + def prepare_tensor_load_for_subtiles( + self, + mma_tile_coord_mnl, + ): + """Prepare the EFC tensors to be loaded by the subtiles and return the number of subtiles to compute.""" + self.bGS_gC = {} + # In the case there is no supplemental tensor to load in the + # epilogue: + self._subtile_cnt = 0 + + def f(tensor_name, attributes): + self.bGS_gC[tensor_name] = self.bGS_gC_partitioned[tensor_name][ + ( + None, + None, + None, + *mma_tile_coord_mnl, + ) + ] + self.bGS_gC[tensor_name] = cutlass.cute.group_modes( + self.bGS_gC[tensor_name], + 1, + cutlass.cute.rank(self.bGS_gC[tensor_name]), + ) + st_cnt = cutlass.cute.size(self.bGS_gC[tensor_name].shape, mode=[1]) + if self._subtile_cnt == 0: + # Keep the first loaded tensor as a reference. + self._subtile_cnt = st_cnt + if st_cnt != self._subtile_cnt: + raise NotImplementedError( + f"Subtile count mismatch: tensor '{self.efc.read_tensor_names[0]}' has {self._subtile_cnt} subtiles, " + f"but tensor '{tensor_name}' has {st_cnt} subtiles. All tensors must have the same subtile count." + ) + + self.foreach_read_tensor(f) + + return self._subtile_cnt + + @trace_in_mlir + def load_tensor_subtiles( + self, subtile_idx, c_pipeline, c_pipeline_producer_state + ): + """Load the subtiles of the EFC tensors.""" + + def f(tensor_name, attributes): + # Load supplemental tensor from global memory to shared memory. + cutlass.cute.copy( + self.parameter[f"{tensor_name}_tma_atom_read"], + self.bGS_gC[tensor_name][(None, subtile_idx)], + self.bGS_sC[tensor_name][(None, c_pipeline_producer_state.index)], + tma_bar_ptr=c_pipeline.producer_get_barrier( + c_pipeline_producer_state + ), + ) + + self.foreach_read_tensor(f) + + class Phase(enum.Enum): + ParameterAnalysis = enum.auto() + """Epilogue function during analysis of its parameters.""" + + ThreadOperation = enum.auto() + """Epilogue function used for computation.""" + + PyTorchEvaluation = enum.auto() + """Epilogue function used for verification on CPU with PyTorch.""" + + class _RemapModesAccessor: + """Helper to enable ``C.remap_modes[:,0,1]`` syntax via ``__getitem__``. + + Converts ``slice(None)`` (``:`` in subscript notation) to ``True`` + (broadcast) and passes integers through as mode indices. + """ + + def __init__(self, tensor): + self._tensor = tensor + + @staticmethod + def _convert_subscript_element(i, k): + """Convert a single subscript element to a source_mode_map entry. + + Only ``:`` (broadcast) and ``int`` (mode index) are valid. + Reject everything else early with a clear message referencing + the ``C.remap_modes[:, 0, 1]`` syntax. + """ + if isinstance(k, int): + return k + if isinstance(k, slice): + if k != slice(None): + raise TypeError( + f"remap_modes subscript element [{i}] is a " + f"non-trivial slice ({k.start}:{k.stop}:{k.step}). " + f"Only ':' (broadcast) is supported, " + f"e.g. C.remap_modes[:, 0, 1]." + ) + return True + raise TypeError( + f"remap_modes subscript element [{i}] has unsupported type " + f"{type(k).__name__}: {k!r}. Each element must be ':' " + f"(broadcast) or an int (mode index), " + f"e.g. C.remap_modes[:, 0, 1]." + ) + + def __getitem__(self, key): + if not isinstance(key, tuple): + key = (key,) + source_mode_map = tuple( + self._convert_subscript_element(i, k) for i, k in enumerate(key) + ) + return self._tensor.configuration.remap_modes( + source=self._tensor, + source_mode_map=source_mode_map, + ) + + class Tensor: + """A proxy object to be used as an EFC argument to introspect or execute + the epilogue configuration function in a given phase. + """ + + @dataclasses.dataclass + class ParameterAttributes: + """Store some characteristics of the epilogue parameters.""" + + is_tensor: bool # Tensor or scalar. + is_read: bool = False # This tensor is read. + is_written: bool = False # This tensor is written. + # The name of the tensor to be remapped. + mapped_source: str = None + # The remapping modes to define the new tensor in respect to the + # layout model. Each element is either an integer to express the + # source mode to map at its place or `True` to broadcast the source + # tensor for this mode. + source_mode_map: tuple = None + # True if the source_mode_map has only broadcast modes, so the + # tensor can degenerate to a single scalar as an optimization. + degenerate_to_scalar: bool = False + + def __init__( + self, + phase: typing.ForwardRef("EFC.Phase"), + name: str, + efc: EFC, + configuration, + ): + """Initialize a tensor proxy bound to an EFC phase and configuration. + + This proxy tracks how an epilogue parameter is used during different + phases of EFC execution and provides access to per-parameter metadata. + It also captures the current configuration context so later `load()` + and `store()` calls can route to the correct backend (analysis, + kernel-thread operation, or PyTorch evaluation). + + Args: + phase: Current EFC phase controlling the proxy behavior. + name: Parameter name as declared in the epilogue function. + efc: Owning EFC instance that holds shared state. + configuration: Phase-specific configuration with runtime args. + """ + self.phase = phase + self.name = name + self.efc = efc + self.configuration = configuration + self.attributes: EFC.Tensor.ParameterAttributes = efc.parameter_attributes[ + name + ] + log(f"Tensor {self.name = }") + + def load(self): + """Load a tensor value inside the EFC function. + + The behavior depends on the phase: + - ParameterAnalysis: mark the tensor as read and return a dummy value + - ThreadOperation: load from the epilogue context and cast to epi dtype + - PyTorchEvaluation: return the runtime tensor from evaluation inputs + """ + log(f"Tensor.load {self.name = } {self.phase = }") + match self.phase: + case EFC.Phase.ParameterAnalysis: + # Record that the tensor is read: + self.attributes.is_read = True + # Some value to have expression evaluation happy + return 1 + + case EFC.Phase.ThreadOperation: + return self.configuration.epilogue_context.load[self.name].to( + self.efc.gemm.epi_dtype + ) + + case EFC.Phase.PyTorchEvaluation: + return self.configuration.args[self.name] + + case _: + raise NotImplementedError( + f"load({self.name}) not implemented for phase {self.phase}" + ) + + def store(self, value): + """Store a tensor value inside the EFC function. + + The behavior depends on the phase: + - ParameterAnalysis: mark the tensor as written (no data movement) + - ThreadOperation: store into epilogue context with JIT dtype + - PyTorchEvaluation: copy into the runtime output tensor + """ + log(f"Tensor.store {self.name = } {self.phase = }") + match self.phase: + case EFC.Phase.ParameterAnalysis: + # Record that the tensor is written: + self.attributes.is_written = True + + case EFC.Phase.ThreadOperation: + tRS_rD = self.configuration.epilogue_context.store[self.name] + tRS_rD.store(value.to(self.efc.jit.tensor_dtype[self.name])) + + case EFC.Phase.PyTorchEvaluation: + # Assign the PyTorch tensor target with the given value. + self.configuration.args[self.name].copy_(value) + + case _: + raise NotImplementedError( + f"store({self.name}) not implemented for phase {self.phase}" + ) + + @property + def remap_modes(self): + """Remap modes of this tensor to match the output shape (m, n, l). + + Use subscript syntax with ``:`` for broadcast (``True``) and + integers for mode indices:: + + C.remap_modes[:,0,1].load() # broadcast along first dim + Y.remap_modes[1,0,2].store(v) # transpose + s.remap_modes[:,:,:].load() # scalar broadcast + + This is a ``@property`` so that ``C.remap_modes`` (without + parentheses) returns a helper object whose ``__getitem__`` + intercepts the ``[...]`` subscript. Without ``@property``, + ``C.remap_modes[:,0,1]`` would be a syntax error because + Python would see ``remap_modes`` as an unsubscripted method + object, not something that supports ``[]``. + """ + return EFC._RemapModesAccessor(self) + + class Configuration: + """Specialize the epilogue function provided by the user to be called in + the compute kernel customization point at a given phase. + """ + + def __init__(self, efc: EFC, phase: EFC.Phase, *args): + """According to the phase, args represent various things. For: + - EFC.Phase.ParameterAnalysis: + - args is empty and not used; + - EFC.Phase.ThreadOperation: + - args[0] is the epilogue_context; + - EFC.Phase.PyTorchEvaluation: + - args[0] is the matrix multiplication result as a torch.Tensor; + - args[1] is a dict mapping argument names to the arguments, like + the scalars or torch.Tensor. It is later extended when + remap_modes() is used. + """ + self.efc = efc + self.phase = phase + self._remap_cache = {} + match self.phase: + case EFC.Phase.ThreadOperation: + self.epilogue_context = args[0] + + case EFC.Phase.PyTorchEvaluation: + self.accum_value = args[0] + self.args = args[1] + + self.arguments = [ + # These are the original epilogue function parameters, it does + # not include synthetic arguments used for tensor broadcast for + # example. + self._argument(name) + for name in efc.epilogue_parameter_names + ] + + def _argument(self, name): + """Generate the argument used by the specialized epilogue with the + given name. + """ + if self.efc.parameter_attributes[name].is_tensor: + # Delegate the phase-related behavior to the proxy Tensor object + # itself. + return EFC.Tensor(self.phase, name, self.efc, self) + # Otherwise, we have a dynamic scalar parameter. + match self.phase: + case EFC.Phase.ParameterAnalysis: + # Use some dummy value during introspection phase. + return cutlass.Float32(42).to(self.efc.gemm.epi_dtype) + + case EFC.Phase.ThreadOperation: + # Return directly the real kernel parameter with the same name. + return cutlass.Float32(self.efc.kernel.parameter[name]).to( + self.efc.gemm.epi_dtype + ) + + case EFC.Phase.PyTorchEvaluation: + # Just use the provided argument. + return self.args[name] + + case _: + raise NotImplementedError( + f"argument({name}) not implemented for phase {self.phase}" + ) + + def __call__(self): + """Execute the epilogue provided by the end-user with some specific + arguments crafted for the current phase. + + Pass self as an argument, to be seen as `efc_config`, a way to + access the EFC instance and its properties. + """ + log( + f"Configuration.__call__ evaluate the user EFC function {self.phase = } {self.arguments = }" + ) + return self.efc.epilogue_function_configuration(self, *self.arguments) + + def accum(self): + """Provide the accumulator value to the user.""" + log(f"Configuration.accum {self.phase = }") + match self.phase: + case EFC.Phase.ParameterAnalysis: + log("Configuration.accum ParameterAnalysis 42!") + # The answer to anything. + return cutlass.Float32(42) + + case EFC.Phase.ThreadOperation: + log( + f"Configuration.accum ThreadOperation {self.epilogue_context.acc_vec = }" + ) + return self.epilogue_context.acc_vec + + case EFC.Phase.PyTorchEvaluation: + log(f"Configuration.accum PyTorchEvaluation {self.accum_value = }") + return self.accum_value + + case _: + raise NotImplementedError( + f"accum() not implemented for phase {self.phase}" + ) + + def remap_modes( + self, + source: EFC.Tensor, + source_mode_map: tuple, + ): + """Remap modes of a source tensor to match the output shape (m, n, l). + + The target shape is implicitly the output tensor shape derived + from the GEMM input matrices A(m, k, l) and B(n, k, l). + + The source_mode_map is a sequence with a length matching the + output rank. + + Each element is either an integer defining the mode of the source + tensor to map at this place or `True` to express that the source + tensor is broadcast for this mode. + """ + # Validate that source is an EFC.Tensor proxy, not a raw scalar + # or some other value passed by mistake. + if not isinstance(source, EFC.Tensor): + raise TypeError( + f"remap_modes() expects an EFC.Tensor proxy as source, " + f"got {type(source).__name__}: {source!r}. " + f"Use remap_modes on a tensor parameter like " + f"C.remap_modes[:, 0, 1], not on a scalar." + ) + # The source must be an actual tensor parameter, not a scalar + # like alpha or beta. + if not source.attributes.is_tensor: + raise ValueError( + f"remap_modes() source '{source.name}' is a scalar " + f"parameter, not a tensor. Only tensor parameters can " + f"be remapped." + ) + # Remapping an already-remapped tensor is likely a mistake: + # remap the original source instead. + if source.attributes.mapped_source is not None: + raise ValueError( + f"remap_modes() source '{source.name}' is already a " + f"remapped view of '{source.attributes.mapped_source}'. " + f"Remap the original tensor directly instead." + ) + # Validate source_mode_map element types, bounds and uniqueness + # early, before they propagate into harder-to-diagnose downstream + # errors. The implied source rank is the number of int entries + # (non-broadcast modes). + int_indices = [] + for i, v in enumerate(source_mode_map): + if v is True: + continue + if not isinstance(v, int): + raise TypeError( + f"source_mode_map[{i}] must be an int or True " + f"(broadcast), got {type(v).__name__}: {v!r}." + ) + int_indices.append(v) + # The effective source rank is the max index + 1 (not the count + # of indices), because some source dims may be unused (e.g. + # size-1 broadcast dims in a tensor like (1, N) where dim 0 is + # skipped via True). + source_rank = (max(int_indices) + 1) if int_indices else 0 + # Check bounds. Guard with ``v is not True`` because bool is a + # subclass of int in Python, so isinstance(True, int) is True. + for i, v in enumerate(source_mode_map): + if v is not True and isinstance(v, int) and not (0 <= v < source_rank): + raise ValueError( + f"source_mode_map[{i}] = {v} is out of range for a " + f"source tensor of rank {source_rank}. Valid indices " + f"are 0..{source_rank - 1}." + ) + if len(int_indices) != len(set(int_indices)): + duplicates = sorted( + set(v for v in int_indices if int_indices.count(v) > 1) + ) + raise ValueError( + f"source_mode_map contains duplicate source mode indices: " + f"{duplicates}. Each source mode must appear exactly once." + ) + log( + f"Configuration.remap_modes {source.name = } {source_mode_map = } {self.phase = }" + ) + # Construct a new tensor as a remapped view with a mangled name + # storing all the details. + mangled_name = f"{source.name}__remapped__{'__'.join([str(e) for e in source_mode_map])}" + # Return cached result if this exact remapping was already processed. + # This avoids overwriting the ParameterAttributes entry (which would + # lose is_read/is_written flags set by earlier .load()/.store() calls) + # and avoids redundant EFC.Tensor proxy creation and PyTorch remapping. + if mangled_name in self._remap_cache: + return self._remap_cache[mangled_name] + match self.phase: + case EFC.Phase.ParameterAnalysis: + log("EFC.remap_modes creates and registers a tensor") + # If the source_mode_map has only broadcast modes, it can be + # resolved to a scalar instead as an optimization. + broadcast_only = sum(v is True for v in source_mode_map) == len( + source_mode_map + ) + # Register the new tensor. + self.efc.parameter_attributes[mangled_name] = ( + EFC.Tensor.ParameterAttributes( + is_tensor=True, + mapped_source=source.name, + source_mode_map=source_mode_map, + degenerate_to_scalar=broadcast_only, + ) + ) + + case EFC.Phase.ThreadOperation: + log("EFC.remap_modes ThreadOperation returns Tensor proxy") + # Just use the proxy tensor as is, so .load() and .store() + # can work with it later. + + case EFC.Phase.PyTorchEvaluation: + # Do the real PyTorch tensor remapping on the source + # tensor using the accumulator shape as the output (m, n, l). + remapped = torch_tensor_remap_modes( + self.args[source.name], + tuple(self.accum_value.shape), + source_mode_map, + ) + log( + f"EFC.remap_modes PyTorchEvaluation self.args[{source.name}] = {self.args[source.name]}" + ) + log(f"EFC.remap_modes PyTorchEvaluation {remapped = }") + # Register the new tensor so it can be used later in an + # expression. + self.args[mangled_name] = remapped + + case _: + raise NotImplementedError( + f"remap_modes({mangled_name = }) not implemented for phase {self.phase}" + ) + + # Cache and return the proxy tensor with the mangled name so it can + # be used later in an expression. + tensor = EFC.Tensor(self.phase, mangled_name, self.efc, self) + self._remap_cache[mangled_name] = tensor + return tensor + + # Some helper functions for common operations. + + def maximum(self, x, y): + """Element-wise maximum of 2 tensors""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + return EFC.maximum(x, y) + case EFC.Phase.PyTorchEvaluation: + import torch + + return torch.maximum(x, y) + case _: + raise NotImplementedError( + f"maximum() not implemented for phase {self.phase}" + ) + + def minimum(self, x, y): + """Element-wise minimum of 2 tensors""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + return EFC.minimum(x, y) + case EFC.Phase.PyTorchEvaluation: + import torch + + return torch.minimum(x, y) + case _: + raise NotImplementedError( + f"minimum() not implemented for phase {self.phase}" + ) + + # Define some activation functions inspired by: + # - cutlass/python/cutlass_cppgen/epilogue/epilogue.py + # - cutlass/python/cutlass_cppgen/backend/epilogue.py + + def identity(self, x): + """Identity activation function: f(x) = x""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + return x + case EFC.Phase.PyTorchEvaluation: + return x + case _: + raise NotImplementedError( + f"identity() not implemented for phase {self.phase}" + ) + + def relu(self, x): + """ReLU activation function: f(x) = maximum(0, x)""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + return EFC.maximum(x, self.full_like(x, 0)) + case EFC.Phase.PyTorchEvaluation: + import torch + + return torch.nn.functional.relu(x) + case _: + raise NotImplementedError( + f"relu() not implemented for phase {self.phase}" + ) + + def leaky_relu(self, x, negative_slope=0.01): + """Leaky ReLU activation function: f(x) = maximum(0, x) + negative_slope * minimum(0, x)""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + # Use self.full_like to have all the computation done with + # the same type as x element type. + zero = self.full_like(x, 0) + return EFC.maximum(x, zero) + EFC.minimum(x, zero) * self.full_like( + x, negative_slope + ) + case EFC.Phase.PyTorchEvaluation: + import torch + + return torch.nn.functional.leaky_relu(x, negative_slope) + case _: + raise NotImplementedError( + f"leaky_relu() not implemented for phase {self.phase}" + ) + + def tanh(self, x): + """Hyperbolic tangent activation function""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + return cutlass.cute.tanh(x) + case EFC.Phase.PyTorchEvaluation: + import torch + + return torch.tanh(x) + case _: + raise NotImplementedError( + f"tanh() not implemented for phase {self.phase}" + ) + + def sigmoid(self, x): + """Sigmoid activation function: f(x) = 1 / (1 + exp(-x))""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + # Use self.full_like to have all the computation done with + # the same type as x element type. + # sigmoid(x) = 1 / (1 + exp(-x)) + return self.full_like(x, 1) / ( + self.full_like(x, 1) + cutlass.cute.exp(-x) + ) + case EFC.Phase.PyTorchEvaluation: + import torch + + return torch.sigmoid(x) + case _: + raise NotImplementedError( + f"sigmoid() not implemented for phase {self.phase}" + ) + + def silu(self, x): + """SiLU (Swish) activation function: f(x) = x * sigmoid(x)""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + # silu(x) = x * sigmoid(x) + return x * self.sigmoid(x) + case EFC.Phase.PyTorchEvaluation: + import torch + + return torch.nn.functional.silu(x) + case _: + raise NotImplementedError( + f"silu() not implemented for phase {self.phase}" + ) + + def hardswish(self, x): + """Hard Swish activation function: f(x) = x * relu6(x + 3) / 6""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + # Use self.full_like to have all the computation done with + # the same type as x element type. + # hardswish(x) = x * minimum(maximum(x + 3, 0), 6) / 6 + relu6 = EFC.minimum( + EFC.maximum(x + self.full_like(x, 3), self.full_like(x, 0)), + self.full_like(x, 6), + ) + return x * relu6 / self.full_like(x, 6) + case EFC.Phase.PyTorchEvaluation: + import torch + + return torch.nn.functional.hardswish(x) + case _: + raise NotImplementedError( + f"hardswish() not implemented for phase {self.phase}" + ) + + def gelu(self, x): + """GELU (Gaussian Error Linear Unit) activation function.""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + # Use self.full_like to have all the computation done with + # the same type as x element type. + # GELU approximation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) + # Using a simpler approximation for CuTe + sqrt_2_over_pi = self.full_like(x, 0.7978845608028654) + return ( + self.full_like(x, 0.5) + * x + * ( + self.full_like(x, 1) + + cutlass.cute.tanh( + sqrt_2_over_pi + * (x + self.full_like(x, 0.044715) * x * x * x) + ) + ) + ) + case EFC.Phase.PyTorchEvaluation: + import torch + + return torch.nn.functional.gelu(x) + case _: + raise NotImplementedError( + f"gelu() not implemented for phase {self.phase}" + ) + + def __getattr__(self, name): + """Called when an attribute or method is not found. + + Hijack this mechanism to dispatch/emulate functions like + cute.full_like() or torch.full_like() provided inside the epilogue + function as self.full_like(). + + This is required since the epilogue is used not only in a @cute.jit + or @cute.kernel but also executed in a normal context for analyzing + the epilogue content and even run in emulation with frameworks like + PyTorch. + """ + + def chameleon(self, *args, **kwargs): + """The great impostor method. + + """ + log(f"Configuration.chameleon {name = } {self.phase = }") + match self.phase: + case EFC.Phase.ParameterAnalysis: + # Just return a value to go on with the fake evaluation, in + # the case the function is expected to return a result. It + # will be ignored anyway in the opposite case. + return 1 + + case EFC.Phase.ThreadOperation: + # In the @cute.kernel context, just use the normal CuTe + # implementation. + return getattr(cutlass.cute, name)(*args, **kwargs) + + case EFC.Phase.PyTorchEvaluation: + # In the PyTorch context, call the equivalent function + # with the same name. + import torch + + return getattr(torch, name)(*args, **kwargs) + + case _: + raise NotImplementedError( + f"self.{name} not implemented for phase {self.phase}" + ) + + # Update the function name to match the requested attribute name. + chameleon.__name__ = name + # Return chameleon blessed as a bound method of self. + return types.MethodType(chameleon, self) + + def __init__( + self, + gemm, + epilogue_function_configuration, + ): + """Construct an EFC instance.""" + self.gemm = gemm + self.epilogue_function_configuration = epilogue_function_configuration + self.analyze_epilogue(epilogue_function_configuration) + + def analyze_epilogue(self, epilogue_function_configuration): + """Analyze the epilogue configuration function to extract its parameter + names. + """ + sig = inspect.signature(epilogue_function_configuration) + names = list(sig.parameters) + # Require the first parameter to be named "efc_config". This is very + # intrusive but at the same time some people got confused when they + # forgot this parameter. + if names[0] != "efc_config": + raise RuntimeError( + "The epilogue configuration function must take efc_config as an argument" + ) + + # Keep all the argument names but the first "efc_config" one. + self.epilogue_parameter_names = names[1:] + log(f"{self.epilogue_parameter_names = }") + + def compile(self, *supplemental_arguments): + """Compile with all the arguments to know the types during compilation + while hiding the epilogue details. + """ + if len(supplemental_arguments) != len(self.epilogue_parameter_names): + raise ValueError( + f"Expected {len(self.epilogue_parameter_names)} supplemental arguments " + f"({', '.join(self.epilogue_parameter_names)}), got {len(supplemental_arguments)}." + ) + # Update the active epilogue instance to use the new Parameter class + self.analyze_epilogue_with_arguments(supplemental_arguments) + # Create the metaprogramming objects for the @cute.jit and @cute.kernel + # parts. For now just forward all the parameters as is. + self.jit = EFC.JIT(self, self.epilogue_parameter_names) + if not self.written_tensor_names: + raise NotImplementedError( + "The epilogue requires at least one written tensor to do something useful." + ) + + def analyze_epilogue_with_arguments(self, supplemental_arguments): + # Keep track of the usage attributes of the tensors by their names. + self.parameter_attributes = {} + log(f"{self.analyze_epilogue_with_arguments}:") + for name, a in zip(self.epilogue_parameter_names, supplemental_arguments): + log(f"{name = } {a = }, {type(a) = }") + self.parameter_attributes[name] = EFC.Tensor.ParameterAttributes( + is_tensor=isinstance(a, cutlass.cute.Tensor) + ) + + # Evaluate the epilogue function for parameter analysis + self.specialized_epilogue(EFC.Phase.ParameterAnalysis)() + log( + f"tensor attributes = {NEWLINE_AND_TAB}{NEWLINE_AND_TAB.join([name + ':' + TAB + str(attributes) for name, attributes in self.parameter_attributes.items()])}" + ) + # Keep track of all the epilogue tensor use cases per name: + self.used_tensor_names = [] + self.read_tensor_names = [] + self.written_tensor_names = [] + self.remapped_tensor_names = [] + for name, attributes in self.parameter_attributes.items(): + attributes = self.parameter_attributes[name] + if not attributes.is_tensor: + continue + if attributes.is_read or attributes.is_written: + self.used_tensor_names.append(name) + if attributes.is_read: + self.read_tensor_names.append(name) + if attributes.is_written: + self.written_tensor_names.append(name) + if attributes.mapped_source: + self.remapped_tensor_names.append(name) + log(f"{TAB}{self.used_tensor_names = }") + log(f"{TAB}{self.read_tensor_names = }") + log(f"{TAB}{self.written_tensor_names = }") + log(f"{TAB}{self.remapped_tensor_names = }") + + def specialized_epilogue(self, phase: typing.ForwardRef("EFC.Phase"), *args): + """Construct a configuration of the epilogue specialized for a given + phase. The arguments are opaque and depend on the actual phase use. + """ + return EFC.Configuration(self, phase, *args) + + def foreach_argument(self, function): + """Execute the given function for each supplemental argument of the epilogue.""" + for name, attributes in self.parameter_attributes.items(): + function(name, attributes) + + def foreach_tensor(self, function): + """Execute the given function for each supplemental tensor.""" + for tensor_name in self.used_tensor_names: + attributes = self.parameter_attributes[tensor_name] + function(tensor_name, attributes) + + def foreach_read_tensor(self, function): + """Execute the given function for each supplemental read tensor.""" + for tensor_name in self.read_tensor_names: + attributes = self.parameter_attributes[tensor_name] + function(tensor_name, attributes) + + def foreach_written_tensor(self, function): + """Execute the given function for each supplemental written tensor.""" + for tensor_name in self.written_tensor_names: + attributes = self.parameter_attributes[tensor_name] + function(tensor_name, attributes) + + def foreach_remapped_tensor(self, function): + """Execute the given function for each supplemental remapped tensor.""" + for tensor_name in self.remapped_tensor_names: + attributes = self.parameter_attributes[tensor_name] + function(tensor_name, attributes) + + def evaluate_on_cpu(self, matrix_multiplication_ref, args: tuple): + """Evaluate the epilogue fusion configuration function on CPU for + validation using the precomputed matrix multiplication result. + + Use PyTorch for now but could be whatever. + """ + # A mapping from argument names to their values, scalars or torch.Tensor. + epilogue_args = dict(zip(self.epilogue_parameter_names, args)) + # Evaluate the epilogue with PyTorch. The tensor arguments which are + # stored are also evaluated and this is how some results are returned. + self.specialized_epilogue( + EFC.Phase.PyTorchEvaluation, + matrix_multiplication_ref, + epilogue_args, + )() diff --git a/operators/cutlass/operators/providers/cutedsl/evt/converter.py b/operators/cutlass/operators/providers/cutedsl/evt/converter.py new file mode 100644 index 000000000..c385103fb --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/evt/converter.py @@ -0,0 +1,470 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# + + +from collections.abc import Callable + +from cutlass.operators.fusion.ir import DAGIR, LoadNode, StoreNode, TopoVisitorNode +from cutlass.operators.fusion.ir.load_nodes import ( + ColumnBroadcastImpl, + RowBroadcastImpl, + ScalarBroadcastImpl, +) +from cutlass.operators.fusion.ir.store_nodes import ReductionImplBase +from cutlass.operators.fusion.library import ActivationOp, FunctionalOp +from cutlass.operators.providers.cutedsl.evt import common_efc +from cutlass.operators.status import Status + +OpToCuteImplStr = { + FunctionalOp.Exp: lambda x: f"exp({x})", + ActivationOp.ReLU: lambda x: f"relu({x})", + ActivationOp.Sigmoid: lambda x: f"sigmoid({x})", + ActivationOp.Tanh: lambda x: f"tanh({x})", + FunctionalOp.Divides: lambda x, y: f"({x} / {y})", + FunctionalOp.Multiplies: lambda x, y: f"({x} * {y})", + FunctionalOp.Plus: lambda x, y: f"({x} + {y})", + FunctionalOp.Minus: lambda x, y: f"({x} - {y})", +} + + +# Functions to eventually be run as part of the EFC function. +# Every function takes in `efc_config` as the first argument (even if it +# is not used). This is necssary for running analysis +# passes on the EFC function absent an MLIR context (which would be +# needed is we used `cute.exp` directly). + +OpToCuteImpl = { + FunctionalOp.Exp: lambda efc_config, x: efc_config.exp(x), + ActivationOp.ReLU: lambda efc_config, x: efc_config.relu(x), + ActivationOp.Sigmoid: lambda efc_config, x: efc_config.sigmoid(x), + ActivationOp.Tanh: lambda efc_config, x: efc_config.tanh(x), + FunctionalOp.Divides: lambda efc_config, x, y: x / y, + FunctionalOp.Multiplies: lambda efc_config, x, y: x * y, + FunctionalOp.Plus: lambda efc_config, x, y: x + y, + FunctionalOp.Minus: lambda efc_config, x, y: x - y, +} + + +def store(efc_config, x, y): + x.store(y) + + +def load(efc_config, x): + return x.load() + + +def remap_and_load(source_mode_map): + """Return a closure that applies remap_modes with the given source_mode_map, then loads. + + Calls ``Configuration.remap_modes()`` directly with a ``source_mode_map`` + tuple containing ``True`` for broadcast dimensions and integer indices for + real source dimensions. This bypasses the ``_RemapModesAccessor`` + subscript helper (which converts ``slice(None)`` → ``True``) since we + already have the canonical ``source_mode_map`` form. + """ + + def _op(efc_config, x): + return efc_config.remap_modes(source=x, source_mode_map=source_mode_map).load() + + return _op + + +def get_val(x): + return lambda efc_config: x + + +# The SM100 EFC kernel permutes supplemental epilogue tensors from the +# user-facing (L, M, N) order to the kernel-internal (M, N, L) order. +# This must match the ``cute.select(t.layout, ...)`` call in +# ``sm100_static_persistent_efc_impl.py``. +_EPILOGUE_PERMUTATION = (1, 2, 0) # (L, M, N) → (M, N, L) + + +def _build_source_mode_map(broadcast_stride, original_rank: int) -> tuple: + """Build an EFC ``source_mode_map`` for ``handle_remapping``. + + The ``source_mode_map`` is in kernel ``(M, N, L)`` order (matching + ``output_shape``) and each entry either: + + - ``True``: this target dimension is broadcast. + - An integer *i*: this target dimension maps to source dimension *i* + in the **original** (untouched) user tensor. + + Broadcast source tensors are passed to the kernel **as-is** — no + ``add_batch_mode``, no permute, no squeeze. ``remap_modes`` inside + the kernel does all the layout work. + + The DAGIR broadcast stride is in ``(L, M, N)`` order. The + ``_EPILOGUE_PERMUTATION`` maps DAGIR positions to kernel target + positions: target[t] corresponds to DAGIR dim + ``_EPILOGUE_PERMUTATION[t]``. + + Within the user tensor, NumPy broadcasting prepends dims from the left. + So the user's original dims occupy the **trailing** positions of the + DAGIR stride. The source dim index for a DAGIR position is its offset + from the first original dim. + + Args: + broadcast_stride: The tensor stride from the DAGIR, in user-facing + ``(L, M, N)`` order after broadcast propagation. May be an + ``int`` for rank-1 tensors. + original_rank: The rank of the user-provided tensor (before DAGIR + broadcast propagation expanded it). + + Examples (target is always ``(M, N, L)`` order):: + + # (N,) orig_rank=1 → (True, 0, True) + _build_source_mode_map((0, 0, 1), 1) + + # (1, N) orig_rank=2 → (True, 1, 0) + _build_source_mode_map((0, 0, 1), 2) + + # (M, 1) orig_rank=2 → (0, 1, True) + _build_source_mode_map((0, 1, 0), 2) + + # (L, M, 1) orig_rank=3 → (1, 2, 0) + _build_source_mode_map((5, 1, 0), 3) + """ + # Normalize: the DAGIR Layout stores stride as a plain int for rank 1. + if isinstance(broadcast_stride, int): + broadcast_stride = (broadcast_stride,) + + dagir_rank = len(broadcast_stride) + n_prepended = dagir_rank - original_rank + + # Build source_mode_map in DAGIR (L, M, N) order. + # - Prepended dims (don't exist in source) → True. + # - Original dims with stride 0 (size-1 broadcast dims) → True. + # These dims exist in the source tensor but carry no data; + # using True forces stride 0 in the target, avoiding the + # non-zero stride that PyTorch assigns to size-1 dims. + # - Original dims with non-zero stride → source dim index. + lmn_map = [] + source_dim = 0 + for i in range(dagir_rank): + if i < n_prepended: + # Dimension prepended by broadcast propagation. + lmn_map.append(True) + elif broadcast_stride[i] == 0: + # Original dim with stride 0 — size-1 broadcast dim. + # Use True to force stride 0, skipping this source dim. + lmn_map.append(True) + source_dim += 1 # Skip past this source dim. + else: + # Original dim with real data. + lmn_map.append(source_dim) + source_dim += 1 + + # Permute from DAGIR (L, M, N) to kernel (M, N, L) target order. + return tuple(lmn_map[i] for i in _EPILOGUE_PERMUTATION) + + +class EFCConverter: + """Helper class to translate from DAGIR to the CuTe DSL epilogue visitor tree (EVT) structure + + The CuTe DSL EVT structure is as follows for a (alpha * accum + beta * C) epilogue: + def epi(efc_config, C, alpha, beta, D): + C_val = C.load() + alpha_val = alpha.load() + beta_val = beta.load() + compute_0_val = alpha_val * efc_config.accum() + compute_1_val = beta_val * C_val + compute_2_val = compute_0_val + compute_1_val + D.store(compute_2_val) + """ + + @staticmethod + def convert( + dag_ir: DAGIR, + parameter_names: list[str], + parameter_tensors: dict | None = None, + ) -> Callable: + """Converts the DAGIR to a callable epilogue function supported by CuTe DSL EFC. + + :param parameter_tensors: Optional dict mapping parameter names to + their original user-provided tensor/scalar values. Used to + determine the original rank of broadcast tensors. + + The simplest way to do this would be to convert the DAGIR into the equivalent + string representation of the EFC function and call `exec` on it to get the callable. + However, this is generally not considered safe as it can potentially open avenues + for allowing arbitrary code execution. + + Instead, we define a generic configurable epilogue that we specialize based + on the DAGIR itself. The generic epilogue takes in a list of parameters + and executes a predefined sequence of operations. Each operation is executed + in order and places its result on a stack. + + Outside this generic epilogue, we define the sequence of operations that + are to be executed and determine where the sources for such operations live + (either their index in the parameter list or on the stack). + + Example: + ======= + Suppose we have the following epilogue: + ``` + def epi(accum, alpha): + D = accum * alpha + return D + ``` + A simple string representation of the epilogue in EFC format would be: + + ``` + def efc_epi(efc_config, alpha, D): + accum_val = efc_config.accum() + alpha_val = alpha.load() + temp = accum_val * alpha_val + D.store(temp) + ``` + + We would like to generalize this to a function with signature: + ``` + def efc_epi(efc_config, *parameters): + ``` + that can perform arbitrary operations. + + We would know from `parameter_names` that the list of parameters provided to the EFC epilogue + will be [alpha, D]. We then traverse the DAGIR and see the following operations in order: + + LoadNode(accum) + LoadNode(alpha) + ComputeNode(mul, alpha, accum) -> temp + StoreNode(temp, D) + + We can convert this into a series of corresponding operations that operate on variables + corresponding to indices. Indices represent either a position in the parameter list + or a position on the stack (index < len(parameters) implying a position in the parameter list). + In this example, [alpha, D] correspond to indices 0 and 1, respectively. + + Traversing each DAGIR node above, we get: + LoadNode(accum) -> efc_config.accum(), result is in stack[0] (index 2) + LoadNode(alpha) -> alpha.load() + -> parameter[0].load(), result is in stack[1] (index 3) + ComputeNode(mul, alpha, accum) -> stack[1] * stack[0] (index 3 * index 2), result is in stack[2] (index 4) + StoreNode(temp, D) -> D.store(stack[2]) + -> paramter[1].store(stack[2]) + + This is encoded as the following tuples: + ``` + ops = [ + # Load of accum omitted because it is performed automatically in the generic epilogue + (load, 0), + (mul, 3, 2), + (store, 1, 2), + ] + ``` + + The generic epilogue then simply performs: + ``` + def efc_epi(efc_config, *parameters): + stack = [efc_config.accum()] + for op in ops: + fn = op[0] + inputs = [get(idx) for idx in op[1:]] + stack.append(fn(efc_config, *inputs)) + ``` + Where the `get()` function is a helper that returns either a value from the parameter list or the stack, + depending on the index. + """ + # Provide a unique identifier for cases in which `accum` is also being written out + # We use an integer since parameter_names is a list of strings -- -1 is guaranteed + # not to be in the parameter_names list, but can still be used as a key in the dictionary. + accum_out_name = -1 + + # If 'accum' is in parameter_names, we know that it must be because + # the accumulator is also being written out + accum_out_loc = -1 + if "accum" in parameter_names: + # Rename the output version temporarily so as not to confuse with the + # input accumulator + accum_out_loc = parameter_names.index("accum") + parameter_names[accum_out_loc] = accum_out_name + + name_to_idx = {} + for idx, name in enumerate(parameter_names): + name_to_idx[name] = idx + + cur_idx = len(name_to_idx) + + def add_name(name: str): + nonlocal cur_idx + name_to_idx[name] = cur_idx + cur_idx += 1 + + def idx(name: str): + val = name_to_idx[name] + if isinstance(val, str): + return idx(val) + return val + + add_name("accum") + + # Each entry is a tuple containing a load/store/compute op and operands needed for it. + ops = [] + debug_string_ops = [] + # Track which parameters are broadcast sources. + broadcast_source_names = set() + + for meta in dag_ir.node_metas_topological_order(): + if isinstance(meta, LoadNode) and getattr(meta, "is_output", False): + assert meta.name == "accum" + # Handle the special case where the accumulator is also being written out + # This occurs in DAG IR as a load node with is_output = True + ops.append((store, idx(accum_out_name), idx("accum"))) + debug_string_ops.append("accum_out_name.store(accum)") + cur_idx += 1 + if isinstance(meta, LoadNode): + # Add new values to the stack for any operations that need to be .load'ed. + # This includes any non-scalar input parameters. + is_scalar = isinstance(meta.underlying_impl, ScalarBroadcastImpl) + is_param_scalar = is_scalar and meta.name in parameter_names + is_broadcast = isinstance( + meta.underlying_impl, (RowBroadcastImpl, ColumnBroadcastImpl) + ) + + if meta.name != "accum" and not is_param_scalar: + if is_scalar: + ops.append((get_val(meta.tensor.value),)) + add_name(meta.name) + elif is_broadcast: + # For broadcast tensors, use remap_modes to expand + # the tensor to the full output shape before loading. + # + # Broadcast parameters are passed to the kernel + # untouched — no add_batch_mode, no permute, no + # squeeze. remap_modes inside the kernel does all + # the layout work. + stride = meta.tensor.stride + if isinstance(stride, int): + stride = (stride,) + # Get original tensor rank to determine how many + # dims were prepended by broadcast propagation. + orig_rank = len(stride) # fallback + if parameter_tensors and meta.name in parameter_tensors: + t = parameter_tensors[meta.name] + if hasattr(t, "ndim"): + orig_rank = t.ndim + elif hasattr(t, "shape"): + orig_rank = len(t.shape) + source_mode_map = _build_source_mode_map(stride, orig_rank) + broadcast_source_names.add(meta.name) + ops.append((remap_and_load(source_mode_map), idx(meta.name))) + debug_string_ops.append( + f"{meta.name} = {meta.name}" + f".remap_modes[{source_mode_map}].load()" + ) + add_name(meta.name) + else: + ops.append((load, idx(meta.name))) + debug_string_ops.append(f"{meta.name} = {meta.name}.load()") + + # Update the entry in name_to_idx to the index of the loaded value + add_name(meta.name) + elif isinstance(meta, StoreNode): + children = dag_ir.get_all_inputs(meta.name) + if len(children) != 1: + raise ValueError( + f"Store node {meta.name} has {len(children)} children, but only one is supported" + ) + child = children[0] + ops.append((store, idx(meta.name), idx(child))) + debug_string_ops.append(f"{meta.name}.store({child})") + + # We want to map the following sequence: + # def epi(accum, x, y): + # out0 = accum * x + # out1 = out0 * y + # return out0, out1 + # To: + # x_val = x.load() + # y_val = y.load() + # c0 = efc_config.accum() * x_val + # out0.store(c0) + # c1 = c0 * y_val + # out1.store(c1) + # To support computation on top of out0, we need to remap it to c0 so that the next computation can use it. + name_to_idx[meta.name] = child + cur_idx += 1 + elif isinstance(meta, TopoVisitorNode): + raise ValueError(f"TopoVisitorNode {meta.name} is not supported") + else: + if dag_ir.in_degree(meta.name) == 0: + continue + + sorted_children = [ + idx(child_name) for child_name in dag_ir.get_all_inputs(meta.name) + ] + entry = (OpToCuteImpl[meta.fn], *sorted_children) + ops.append(entry) + debug_string_ops.append( + f"{meta.name} = {OpToCuteImplStr[meta.fn](*dag_ir.get_all_inputs(meta.name))}" + ) + add_name(meta.name) + + def epi(efc_config, *parameters): + stack = [efc_config.accum()] + + def get(idx: int): + if idx < len(parameters): + return parameters[idx] + else: + return stack[idx - len(parameters)] + + for i, op in enumerate(ops): + fn = op[0] + inputs = [get(idx) for idx in op[1:]] + stack.append(fn(efc_config, *inputs)) + + if accum_out_loc != -1: + # Restore the original name of the accumulator + parameter_names[accum_out_loc] = "accum" + + named_epi = common_efc.create_named_epilogue( + ["efc_config", *parameter_names], epi + ) + # Attach metadata about which parameters are broadcast sources so the + # kernel can squeeze them before entering the JIT context. + named_epi._broadcast_source_names = broadcast_source_names + return named_epi + + @staticmethod + def supports(dag_ir: DAGIR) -> Status: + """Checks if the DAGIR is supported by CuTe DSL EFC. + """ + # Currently do not support TopVisitorNode and reductions. + for meta in dag_ir.node_metas_topological_order(): + if isinstance(meta, TopoVisitorNode): + return Status.fail("TopoVisitorNode is not supported") + if isinstance(meta.underlying_impl, ReductionImplBase): + return Status.fail("ReductionImplBase is not supported") + return Status.success() + + @staticmethod + def identity_efc(efc_config, D): + D.store(efc_config.accum()) diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/__init__.py b/operators/cutlass/operators/providers/cutedsl/gemm/__init__.py new file mode 100644 index 000000000..065e2a7ca --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/__init__.py @@ -0,0 +1,43 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +# Importing is done for the side effect of registering the Operator classes with the provider. +# F401: Suppress linter warnings about unused import. +# I001: suppress linter warnings about import order, so we control Operator registration order +# E402: False warning about module-level imports not at top of file +# ruff: noqa: F401, I001, E402 + +import cutlass.operators.providers.cutedsl.gemm.sm100_persistent +import cutlass.operators.providers.cutedsl.gemm.sm100_static_persistent_efc +import cutlass.operators.providers.cutedsl.gemm.sm100_tgv_gemm +import cutlass.operators.providers.cutedsl.gemm.sm100_dense_blockscaled_static_persistent +import cutlass.operators.providers.cutedsl.gemm.sm100_contiguous_offset_2d3d_dense_gemm +import cutlass.operators.providers.cutedsl.gemm.sm100_mixed_input +import cutlass.operators.providers.cutedsl.gemm.sm100_persistent_preferred_cluster +import cutlass.operators.providers.cutedsl.gemm.sm90_static_persistent +import cutlass.operators.providers.cutedsl.gemm.sm80_tensorop_gemm diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/implementations/operator_helpers.py b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/operator_helpers.py new file mode 100644 index 000000000..171f88603 --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/operator_helpers.py @@ -0,0 +1,73 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""Shared helpers for CuTeDSL GEMM kernel implementations. + +This module hosts small, dependency-light utilities that are reused across +the persistent / preferred-cluster / EFC / etc. kernel implementations but +do not exist in the CuteDSL API, as well as associated mapper functions. +""" + +from __future__ import annotations + +import enum + +from cutlass.operators.metadata.design.tile_scheduler import ( + CLCDynamicPersistentTileSchedulerMetadata, + StaticPersistentTileSchedulerMetadata, + TileSchedulerMetadata, +) + + +class PersistentTileSchedulerType(enum.Enum): + """Identifies which persistent tile-scheduler dispatch the kernel uses. + + Members: + CLC: Cluster Launch Control dynamic persistent scheduler. + STATIC_PERSISTENT: Static persistent scheduler. Each SM in the + launch grid iterates over a fixed slice of tiles known at + launch time. + """ + + CLC = enum.auto() + STATIC_PERSISTENT = enum.auto() + + +def scheduler_metadata_to_enum( + metadata: TileSchedulerMetadata | None, +) -> PersistentTileSchedulerType: + """Map tile-scheduler metadata to its kernel-side dispatch tag.""" + if isinstance(metadata, CLCDynamicPersistentTileSchedulerMetadata): + return PersistentTileSchedulerType.CLC + if isinstance(metadata, StaticPersistentTileSchedulerMetadata): + return PersistentTileSchedulerType.STATIC_PERSISTENT + raise ValueError( + f"Unsupported tile-scheduler metadata: {metadata!r}. Expected an " + f"instance of {CLCDynamicPersistentTileSchedulerMetadata.__name__} or " + f"{StaticPersistentTileSchedulerMetadata.__name__}." + ) diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/implementations/scheduler.py b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/scheduler.py new file mode 100644 index 000000000..a5f7c3213 --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/scheduler.py @@ -0,0 +1,589 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass.cutlass_dsl import ( + Boolean, + Int32, + Integer, + dsl_user_op, + extract_mlir_values, + min, + new_from_mlir_values, +) + + +class WorkTileInfo: + """A class to represent information about a work tile. + + :ivar tile_idx: The index of the tile. + :type tile_idx: cute.Coord + :ivar is_valid_tile: Whether the tile is valid. + :type is_valid_tile: Boolean + """ + + def __init__(self, tile_idx: cute.Coord, is_valid_tile: Boolean): + self._tile_idx = tile_idx + self._is_valid_tile = Boolean(is_valid_tile) + + def __extract_mlir_values__(self) -> list[ir.Value]: + values = extract_mlir_values(self.tile_idx) + values.extend(extract_mlir_values(self.is_valid_tile)) + return values + + def __new_from_mlir_values__(self, values: list[ir.Value]) -> "WorkTileInfo": + assert len(values) == 4 + new_tile_idx = new_from_mlir_values(self._tile_idx, values[:-1]) + new_is_valid_tile = new_from_mlir_values(self._is_valid_tile, [values[-1]]) + return WorkTileInfo(new_tile_idx, new_is_valid_tile) + + @property + def is_valid_tile(self) -> Boolean: + """Check latest tile returned by the scheduler is valid or not. Any scheduling + requests after all tasks completed will return an invalid tile. + + :return: The validity of the tile. + :rtype: Boolean + """ + return self._is_valid_tile + + @property + def tile_idx(self) -> cute.Coord: + """Get the index of the tile. + + :return: The index of the tile. + :rtype: cute.Coord + """ + return self._tile_idx + + +class PersistentTileSchedulerParams: + """A class to represent parameters for a persistent tile scheduler. + + This class is designed to manage and compute the layout of clusters and tiles + in a batched gemm problem. + + :ivar cluster_shape_mn: Shape of the cluster in (m, n) dimensions (K dimension cta count must be 1). + :type cluster_shape_mn: tuple + :ivar problem_layout_ncluster_mnl: Layout of the problem in terms of + number of clusters in (m, n, l) dimensions. + :type problem_layout_ncluster_mnl: cute.Layout + """ + + @dsl_user_op + def __init__( + self, + problem_shape_ntile_mnl: cute.Shape, + cluster_shape_mnk: cute.Shape, + swizzle_size: int = 1, + raster_along_m: bool = True, + *, + loc=None, + ip=None, + ): + """Initializes the PersistentTileSchedulerParams with the given parameters. + + :param problem_shape_ntile_mnl: The shape of the problem in terms of + number of CTA (Cooperative Thread Array) in (m, n, l) dimensions. + :type problem_shape_ntile_mnl: cute.Shape + :param cluster_shape_mnk: The shape of the cluster in (m, n) dimensions. + :type cluster_shape_mnk: cute.Shape + :param swizzle_size: Swizzling size in the unit of cluster. 1 means no swizzle + :type swizzle_size: int + :param raster_along_m: Rasterization order of clusters. Only used when swizzle_size > 1. + True means along M, false means along N. + :type raster_along_m: bool + + :raises ValueError: If cluster_shape_k is not 1. + """ + if cluster_shape_mnk[2] != 1: + raise ValueError(f"unsupported cluster_shape_k {cluster_shape_mnk[2]}") + if swizzle_size < 1: + raise ValueError(f"expect swizzle_size >= 1, but get {swizzle_size}") + + self.problem_shape_ntile_mnl = problem_shape_ntile_mnl + # cluster_shape_mnk is kept for reconstruction + self._cluster_shape_mnk = cluster_shape_mnk + self.cluster_shape_mn = cluster_shape_mnk[:2] + self.swizzle_size = swizzle_size + self.raster_along_m = raster_along_m + self._loc = loc + + # By default, we follow m major (col-major) raster order, so make a col-major layout + self.problem_layout_ncluster_mnl = cute.make_layout( + cute.ceil_div( + self.problem_shape_ntile_mnl, cluster_shape_mnk[:2], loc=loc, ip=ip + ), + loc=loc, + ip=ip, + ) + + # Apply swizzle if swizzle_size > 1 + if swizzle_size > 1: + problem_shape_ncluster_mnl = cute.round_up( + self.problem_layout_ncluster_mnl.shape, + (1, swizzle_size, 1) if raster_along_m else (swizzle_size, 1, 1), + ) + + if raster_along_m: + self.problem_layout_ncluster_mnl = cute.make_layout( + ( + problem_shape_ncluster_mnl[0], + (swizzle_size, problem_shape_ncluster_mnl[1] // swizzle_size), + problem_shape_ncluster_mnl[2], + ), + stride=( + swizzle_size, + (1, swizzle_size * problem_shape_ncluster_mnl[0]), + problem_shape_ncluster_mnl[0] * problem_shape_ncluster_mnl[1], + ), + loc=loc, + ip=ip, + ) + else: + self.problem_layout_ncluster_mnl = cute.make_layout( + ( + (swizzle_size, problem_shape_ncluster_mnl[0] // swizzle_size), + problem_shape_ncluster_mnl[1], + problem_shape_ncluster_mnl[2], + ), + stride=( + (1, swizzle_size * problem_shape_ncluster_mnl[1]), + swizzle_size, + problem_shape_ncluster_mnl[0] * problem_shape_ncluster_mnl[1], + ), + loc=loc, + ip=ip, + ) + + # Create FastDivmod divisors (only when swizzle_size == 1 for correctness) + # FastDivmod assumes simple col-major layout, incompatible with swizzled layouts + if swizzle_size == 1: + problem_layout_size = cute.size( + self.problem_layout_ncluster_mnl, loc=loc, ip=ip + ) + cluster_count_m = self.problem_layout_ncluster_mnl.shape[0] + cluster_count_n = self.problem_layout_ncluster_mnl.shape[1] + + # batch_fdd: Used to map linear_idx to work_unit_id (handles persistent scheduling) + self.batch_fdd = cute.fast_divmod_create_divisor( + problem_layout_size, loc=loc, ip=ip + ) + + if raster_along_m: + cluster_count_major = cluster_count_m + cluster_count_minor = cluster_count_n + else: + cluster_count_major = cluster_count_n + cluster_count_minor = cluster_count_m + + # cluster_shape_major_fdd: Used to decode work_unit_id to cluster coordinates + self.cluster_shape_major_fdd = cute.fast_divmod_create_divisor( + cluster_count_major, loc=loc, ip=ip + ) + + # cluster_shape_minor_fdd: Used for the second level decomposition + self.cluster_shape_minor_fdd = cute.fast_divmod_create_divisor( + cluster_count_minor, loc=loc, ip=ip + ) + else: + # FastDivmod not applicable with swizzling, set to None + self.batch_fdd = None + self.cluster_shape_major_fdd = None + self.cluster_shape_minor_fdd = None + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + for obj in [ + self.problem_shape_ntile_mnl, + self._cluster_shape_mnk, + self.swizzle_size, + self.raster_along_m, + ]: + obj_values = extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + + # Add FastDivmod divisors to MLIR values for Host->Device transfer + # Only add non-None values to avoid MLIR type errors + fastdivmod_values = [] + fastdivmod_indices = [] # Track which FastDivmod objects are present + + for i, (fdd_name, fdd_obj) in enumerate( + [ + ("batch_fdd", self.batch_fdd), + ("cluster_shape_major_fdd", self.cluster_shape_major_fdd), + ("cluster_shape_minor_fdd", self.cluster_shape_minor_fdd), + ] + ): + if fdd_obj is not None: + # Extract MLIR values from FastDivmodDivisor objects + fdd_values = extract_mlir_values(fdd_obj) + fastdivmod_values.extend(fdd_values) + fastdivmod_indices.append(i) + + values += fastdivmod_values + self._values_pos.append( + len(fastdivmod_indices) + ) # Store count of FastDivmod objects, not values + self._fastdivmod_indices = fastdivmod_indices # Store for reconstruction + + return values + + def __new_from_mlir_values__(self, values): + obj_list = [] + values_copy = list(values) # Make a copy to avoid modifying original + + # Reconstruct original objects from MLIR values + for obj, n_items in zip( + [ + self.problem_shape_ntile_mnl, + self._cluster_shape_mnk, + self.swizzle_size, + self.raster_along_m, + ], + self._values_pos[:-1], # Exclude FastDivmod count + ): + obj_list.append(new_from_mlir_values(obj, values_copy[:n_items])) + values_copy = values_copy[n_items:] + + # Create new params object by calling __init__ with reconstructed values + # This properly recreates layouts and other derived attributes in the device context + new_params = PersistentTileSchedulerParams(*(tuple(obj_list)), loc=self._loc) + + # Restore FastDivmod divisors from remaining values + fdd_names = ["batch_fdd", "cluster_shape_major_fdd", "cluster_shape_minor_fdd"] + + if hasattr(self, "_fastdivmod_indices") and len(self._fastdivmod_indices) > 0: + # Override the FastDivmod divisors created by __init__ with reconstructed ones + for j, original_index in enumerate(self._fastdivmod_indices): + fdd_name = fdd_names[original_index] + # Get the original FastDivmodDivisor object + original_fdd = getattr(self, fdd_name) + if original_fdd is not None and j < len(values_copy): + # Each FastDivmodDivisor has 1 MLIR value + reconstructed_fdd = new_from_mlir_values( + original_fdd, [values_copy[j]] + ) + setattr(new_params, fdd_name, reconstructed_fdd) + + return new_params + + @dsl_user_op + def get_grid_shape( + self, max_active_clusters: Int32, *, loc=None, ip=None + ) -> tuple[Integer, Integer, Integer]: + """Computes the grid shape based on the maximum active clusters allowed. + + :param max_active_clusters: The maximum number of active clusters that + can run in one wave. + :type max_active_clusters: Int32 + + :return: A tuple containing the grid shape in (m, n, persistent_clusters). + - m: self.cluster_shape_m. + - n: self.cluster_shape_n. + - persistent_clusters: Number of persistent clusters that can run. + """ + # Total ctas in problem size + num_ctas_mnl = tuple( + cute.size(x) * y + for x, y in zip( + self.problem_layout_ncluster_mnl.shape, self.cluster_shape_mn + ) + ) + (self.problem_layout_ncluster_mnl.shape[2],) + + num_ctas_in_problem = cute.size(num_ctas_mnl, loc=loc, ip=ip) + + num_ctas_per_cluster = cute.size(self.cluster_shape_mn, loc=loc, ip=ip) + # Total ctas that can run in one wave + num_ctas_per_wave = max_active_clusters * num_ctas_per_cluster + + num_persistent_ctas = min(num_ctas_in_problem, num_ctas_per_wave) + num_persistent_clusters = num_persistent_ctas // num_ctas_per_cluster + + return (*self.cluster_shape_mn, num_persistent_clusters) + + +class StaticPersistentTileScheduler: + """A scheduler for static persistent tile execution in CUTLASS/CuTe kernels. + + :ivar params: Tile schedule related params, including cluster shape and problem_layout_ncluster_mnl + :type params: PersistentTileSchedulerParams + :ivar num_persistent_clusters: Number of persistent clusters that can be launched + :type num_persistent_clusters: Int32 + :ivar cta_id_in_cluster: ID of the CTA within its cluster + :type cta_id_in_cluster: cute.Coord + :ivar _num_tiles_executed: Counter for executed tiles + :type _num_tiles_executed: Int32 + :ivar _current_work_linear_idx: Current cluster index + :type _current_work_linear_idx: Int32 + """ + + def __init__( + self, + params: PersistentTileSchedulerParams, + num_persistent_clusters: Int32, + current_work_linear_idx: Int32, + cta_id_in_cluster: cute.Coord, + num_tiles_executed: Int32, + ): + """Initializes the StaticPersistentTileScheduler with the given parameters. + + :param params: Tile schedule related params, including cluster shape and problem_layout_ncluster_mnl. + :type params: PersistentTileSchedulerParams + :param num_persistent_clusters: Number of persistent clusters that can be launched. + :type num_persistent_clusters: Int32 + :param current_work_linear_idx: Current cluster index. + :type current_work_linear_idx: Int32 + :param cta_id_in_cluster: ID of the CTA within its cluster. + :type cta_id_in_cluster: cute.Coord + :param num_tiles_executed: Counter for executed tiles. + :type num_tiles_executed: Int32 + """ + self.params = params + self.num_persistent_clusters = num_persistent_clusters + self._current_work_linear_idx = current_work_linear_idx + self.cta_id_in_cluster = cta_id_in_cluster + self._num_tiles_executed = num_tiles_executed + + def __extract_mlir_values__(self) -> list[ir.Value]: + values = extract_mlir_values(self.num_persistent_clusters) + values.extend(extract_mlir_values(self._current_work_linear_idx)) + values.extend(extract_mlir_values(self.cta_id_in_cluster)) + values.extend(extract_mlir_values(self._num_tiles_executed)) + + # CRITICAL: Also extract FastDivmod divisors from params + values.extend(extract_mlir_values(self.params)) + + return values + + def __new_from_mlir_values__( + self, values: list[ir.Value] + ) -> "StaticPersistentTileScheduler": + assert len(values) >= 6 + new_num_persistent_clusters = new_from_mlir_values( + self.num_persistent_clusters, [values[0]] + ) + new_current_work_linear_idx = new_from_mlir_values( + self._current_work_linear_idx, [values[1]] + ) + new_cta_id_in_cluster = new_from_mlir_values( + self.cta_id_in_cluster, values[2:5] + ) + new_num_tiles_executed = new_from_mlir_values( + self._num_tiles_executed, [values[5]] + ) + + # Reconstruct params with FastDivmod divisors + params_values = values[6:] # Remaining values are from params + new_params = new_from_mlir_values(self.params, params_values) + + return StaticPersistentTileScheduler( + new_params, # Use reconstructed params with FastDivmod divisors + new_num_persistent_clusters, + new_current_work_linear_idx, + new_cta_id_in_cluster, + new_num_tiles_executed, + ) + + @staticmethod + @dsl_user_op + def create( + params: PersistentTileSchedulerParams, + block_idx: tuple[Integer, Integer, Integer], + grid_dim: tuple[Integer, Integer, Integer], + *, + loc=None, + ip=None, + ): + """Initialize the static persistent tile scheduler. + + :param params: Parameters for the persistent + tile scheduler. + :type params: PersistentTileSchedulerParams + :param block_idx: The 3d block index in the format (bidx, bidy, bidz). + :type block_idx: Tuple[Integer, Integer, Integer] + :param grid_dim: The 3d grid dimensions for kernel launch. + :type grid_dim: Tuple[Integer, Integer, Integer] + + :return: A StaticPersistentTileScheduler object. + :rtype: StaticPersistentTileScheduler + """ + # Calculate the number of persistent clusters by dividing the total grid size + # by the number of CTAs per cluster + num_persistent_clusters = cute.size(grid_dim, loc=loc, ip=ip) // cute.size( + params.cluster_shape_mn, loc=loc, ip=ip + ) + + bidx, bidy, bidz = block_idx + + # Initialize workload index equals to the cluster index in the grid + current_work_linear_idx = Int32(bidz) + + # CTA id in the cluster + cta_id_in_cluster = ( + Int32(bidx % params.cluster_shape_mn[0]), + Int32(bidy % params.cluster_shape_mn[1]), + Int32(0), + ) + # Initialize number of tiles executed to zero + num_tiles_executed = Int32(0) + return StaticPersistentTileScheduler( + params, + num_persistent_clusters, + current_work_linear_idx, + cta_id_in_cluster, + num_tiles_executed, + ) + + # called by host + @staticmethod + def get_grid_shape( + params: PersistentTileSchedulerParams, + max_active_clusters: Int32, + *, + loc=None, + ip=None, + ) -> tuple[Integer, Integer, Integer]: + """Calculates the grid shape to be launched on GPU using problem shape, + threadblock shape, and active cluster size. + + :param params: Parameters for grid shape calculation. + :type params: PersistentTileSchedulerParams + :param max_active_clusters: Maximum active clusters allowed. + :type max_active_clusters: Int32 + + :return: The calculated 3d grid shape. + :rtype: Tuple[Integer, Integer, Integer] + """ + return params.get_grid_shape(max_active_clusters, loc=loc, ip=ip) + + # private method + def _get_current_work_for_linear_idx( + self, current_work_linear_idx: Int32, *, loc=None, ip=None + ) -> WorkTileInfo: + """Compute current tile coord given current_work_linear_idx and cta_id_in_cluster. + + :param current_work_linear_idx: The linear index of the current work. + :type current_work_linear_idx: Int32 + + :return: An object containing information about the current tile coordinates + and validity status. + :rtype: WorkTileInfo + """ + is_valid = current_work_linear_idx < cute.size( + self.params.problem_layout_ncluster_mnl, loc=loc, ip=ip + ) + + # Choose coordinate calculation method based on swizzle configuration + if self.params.swizzle_size == 1: + # Use FastDivmod optimization for non-swizzled layouts + cur_cluster_coord = self._get_cluster_work_idx_with_fastdivmod( + current_work_linear_idx, loc=loc, ip=ip + ) + else: + # Use get_flat_coord for swizzled layouts (FastDivmod doesn't support them) + cur_cluster_coord = self.params.problem_layout_ncluster_mnl.get_flat_coord( + current_work_linear_idx, loc=loc, ip=ip + ) + + # cur_tile_coord is a tuple of i32 values + cur_tile_coord = tuple( + Int32(x) * Int32(z) + Int32(y) + for x, y, z in zip( + cur_cluster_coord, + self.cta_id_in_cluster, + (*self.params.cluster_shape_mn, Int32(1)), + ) + ) + + return WorkTileInfo(cur_tile_coord, is_valid) + + def _get_cluster_work_idx_with_fastdivmod( + self, current_work_linear_idx: Int32, *, loc=None, ip=None + ) -> tuple[Int32, Int32, Int32]: + """FastDivmod optimized CLUSTER coordinate calculation. + + CRITICAL: This should mimic problem_layout_ncluster_mnl.get_hier_coord() + which returns CLUSTER coordinates, not tile coordinates! + + :param current_work_linear_idx: Linear index in the work space + :type current_work_linear_idx: Int32 + :return: Cluster coordinates (m, n, l) or None if FastDivmod not available + :rtype: Tuple[Int32, Int32, Int32] or None + """ + # Step 1: Handle persistent scheduling - map linear_idx to work_unit_id + work_iteration, work_unit_id = divmod( + current_work_linear_idx, self.params.batch_fdd + ) + + # Step 2: Decode work_unit_id using FastDivmod objects + # The layout structure is: problem_layout_ncluster_mnl has shape (cluster_count_m, cluster_count_n, batch_count) + # work_unit_id needs to be decomposed into (batch_l, cluster_minor, cluster_major) in little-endian order + + # First, get cluster_major using cluster_shape_major_fdd + cluster_minor_batch, cluster_major = divmod( + work_unit_id, self.params.cluster_shape_major_fdd + ) + + # Then decode cluster_minor_batch to get cluster_minor and batch_l using FastDivmod + batch_l, cluster_minor = divmod( + cluster_minor_batch, self.params.cluster_shape_minor_fdd + ) + + if self.params.raster_along_m: + cluster_m = cluster_major + cluster_n = cluster_minor + else: + cluster_m = cluster_minor + cluster_n = cluster_major + + return (cluster_m, cluster_n, batch_l) + + @dsl_user_op + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + return self._get_current_work_for_linear_idx( + self._current_work_linear_idx, loc=loc, ip=ip + ) + + @dsl_user_op + def initial_work_tile_info(self, *, loc=None, ip=None) -> WorkTileInfo: + return self.get_current_work(loc=loc, ip=ip) + + @dsl_user_op + def advance_to_next_work(self, *, advance_count: int = 1, loc=None, ip=None): + self._current_work_linear_idx += Int32(advance_count) * Int32( + self.num_persistent_clusters + ) + self._num_tiles_executed += Int32(1) + + @property + def num_tiles_executed(self) -> Int32: + return self._num_tiles_executed diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_contiguous_offset_2d3d_dense_gemm_impl.py b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_contiguous_offset_2d3d_dense_gemm_impl.py new file mode 100644 index 000000000..762c45111 --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_contiguous_offset_2d3d_dense_gemm_impl.py @@ -0,0 +1,1976 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass._mlir import ir +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cutlass_dsl import ( + Int32, + Integer, + dsl_user_op, + new_from_mlir_values, +) +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +from . import scheduler as scheduler_utils + +""" +High-performance persistent contiguous grouped dense GEMM (C = A * B) example for the NVIDIA Blackwell architecture +using CUTE DSL. +- Matrix A is MxKx1, A can be row-major("K"), ValidM is composed of valid m in different groups +- Matrix B is NxKxL, B can be column-major("K"), L is grouped dimension +- Matrix C is MxNx1, C can be row-major("N"), ValidM is composed of valid m in different groups + +Matrix A/C Memory Layout Diagrams: + + ``` + Group 0 Group 1 Group 2 + -+---------+---------+---------+ + | | | | + K| ValidM0 | ValidM1 | ValidM2 | + | | | | + -+---------+---------+---------+ + |<- ValidM ->| + ``` + Note: the Group(L) dimension will be flatted into M dimension, and the rest Group(L) size is 1. +""" + + +class ContiguousOffsetGroupedGemmTileScheduler( + scheduler_utils.StaticPersistentTileScheduler +): + """Simplified scheduler for contiguous grouped GEMM problems. + + The traditional static persistent tile scheduler is designed for cases in which + each work tile starts at a multiple of the tile shape. This is not necessarily the case + for contiguous offset grouped GEMM. + + This scheduler assumes one of the dimensions in M or N is constant across problems in + the group. It then simply returns the cluster coordinate corresponding to the linearized + ID rasterized along this constant mode. + + It is the responsibility of the caller to determine the validity of the returned tile + and which starting position it maps to. + """ + + def __init__( + self, + params: scheduler_utils.PersistentTileSchedulerParams, + num_persistent_clusters: Int32, + current_work_linear_idx: Int32, + cta_id_in_cluster: cute.Coord, + num_tiles_executed: Int32, + ): + """Initializes the StaticPersistentTileScheduler with the given parameters. + + :param params: Tile schedule related params, including cluster shape and problem_layout_ncluster_mnl. + :type params: PersistentTileSchedulerParams + :param num_persistent_clusters: Number of persistent clusters that can be launched. + :type num_persistent_clusters: Int32 + :param current_work_linear_idx: Current cluster index. + :type current_work_linear_idx: Int32 + :param cta_id_in_cluster: ID of the CTA within its cluster. + :type cta_id_in_cluster: cute.Coord + :param num_tiles_executed: Counter for executed tiles. + :type num_tiles_executed: Int32 + """ + super().__init__( + params, + num_persistent_clusters, + current_work_linear_idx, + cta_id_in_cluster, + num_tiles_executed, + ) + + def __new_from_mlir_values__( + self, values: list[ir.Value] + ) -> "ContiguousOffsetGroupedGemmTileScheduler": + assert len(values) >= 6 + new_num_persistent_clusters = new_from_mlir_values( + self.num_persistent_clusters, [values[0]] + ) + new_current_work_linear_idx = new_from_mlir_values( + self._current_work_linear_idx, [values[1]] + ) + new_cta_id_in_cluster = new_from_mlir_values( + self.cta_id_in_cluster, values[2:5] + ) + new_num_tiles_executed = new_from_mlir_values( + self._num_tiles_executed, [values[5]] + ) + + # Reconstruct params with FastDivmod divisors (same as parent class) + params_values = values[6:] # Remaining values are from params + new_params = new_from_mlir_values(self.params, params_values) + + return ContiguousOffsetGroupedGemmTileScheduler( + new_params, # Use reconstructed params with FastDivmod divisors + new_num_persistent_clusters, + new_current_work_linear_idx, + new_cta_id_in_cluster, + new_num_tiles_executed, + ) + + @staticmethod + @dsl_user_op + def create( + params: scheduler_utils.PersistentTileSchedulerParams, + block_idx: tuple[Integer, Integer, Integer], + grid_dim: tuple[Integer, Integer, Integer], + *, + loc=None, + ip=None, + ): + """Initialize the static persistent tile scheduler. + + :param params: Parameters for the persistent + tile scheduler. + :type params: PersistentTileSchedulerParams + :param block_idx: The 3d block index in the format (bidx, bidy, bidz). + :type block_idx: Tuple[Integer, Integer, Integer] + :param grid_dim: The 3d grid dimensions for kernel launch. + :type grid_dim: Tuple[Integer, Integer, Integer] + + :return: A ContiguousOffsetGroupedGemmTileScheduler object. + :rtype: ContiguousOffsetGroupedGemmTileScheduler + """ + # Calculate the number of persistent clusters by dividing the total grid size + # by the number of CTAs per cluster + num_persistent_clusters = cute.size(grid_dim, loc=loc, ip=ip) // cute.size( + params.cluster_shape_mn, loc=loc, ip=ip + ) + + bidx, bidy, bidz = block_idx + + # Initialize workload index equals to the cluster index in the grid + current_work_linear_idx = Int32(bidz) + + # CTA id in the cluster + cta_id_in_cluster = ( + Int32(bidx % params.cluster_shape_mn[0]), + Int32(bidy % params.cluster_shape_mn[1]), + Int32(0), + ) + # Initialize number of tiles executed to zero + num_tiles_executed = Int32(0) + return ContiguousOffsetGroupedGemmTileScheduler( + params, + num_persistent_clusters, + current_work_linear_idx, + cta_id_in_cluster, + num_tiles_executed, + ) + + # private method + def _get_current_work_for_linear_idx( + self, current_work_linear_idx: Int32, *, loc=None, ip=None + ) -> utils.WorkTileInfo: + """Compute current tile coord given current_work_linear_idx. + + :param current_work_linear_idx: The linear index of the current work. + :type current_work_linear_idx: Int32 + + :return: An object containing information about the current tile coordinates + and validity status. + :rtype: WorkTileInfo + """ + cluster_coord_minor, cluster_coord_major = divmod( + current_work_linear_idx, self.params.cluster_shape_major_fdd + ) + if self.params.raster_along_m: + cluster_coord_m = cluster_coord_major + cluster_coord_n = cluster_coord_minor + else: + cluster_coord_m = cluster_coord_minor + cluster_coord_n = cluster_coord_major + + return utils.WorkTileInfo((cluster_coord_m, cluster_coord_n, Int32(0)), True) + + +class ContiguousOffset2D3DGemmDenseKernel: + """This class implements batched matrix multiplication (C = A * B) with support for fp8 (e4m3fn, e5m2) + and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization. + + :param acc_dtype: Data type for accumulation during computation + :type acc_dtype: type[cutlass.Numeric] + :param use_2cta_instrs: Whether to use CTA group 2 for advanced thread cooperation + :type use_2cta_instrs: bool + :param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N) + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing + :type cluster_shape_mn: Tuple[int, int] + :param offsets_are_end: Whether the offsets represent ending positions of each problem in the group. + :type offsets_are_end: bool + """ + + def __init__( + self, + acc_dtype: type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + offsets_are_end: bool = True, + ): + """Initializes the configuration for a Blackwell blockwise dense GEMM kernel. + + This configuration includes several key aspects: + + 1. MMA Instruction Settings (tcgen05): + - acc_dtype: Data types for MMA accumulator. + - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler. + - use_2cta_instrs: Boolean indicating if the tcgen05 MMA variant + with cta_group=2 should be used. + + 2. Cluster Shape: + - cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster. + + :param acc_dtype: Data type of the accumulator. + :type acc_dtype: type[cutlass.Numeric] + :param mma_tiler_mn: Tuple (M, N) shape of the MMA instruction. + :type mma_tiler_mn: Tuple[int, int] + :param use_2cta_instrs: Boolean, True to use cta_group=2 MMA variant. + :type use_2cta_instrs: bool + :param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster. + :type cluster_shape_mn: Tuple[int, int] + :param offsets_are_end: Whether the offsets represent ending positions of each problem in the group. + :type offsets_are_end: bool + """ + self.acc_dtype: type[cutlass.Numeric] = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + # K dimension is deferred in _setup_attributes + self.mma_tiler = (*mma_tiler_mn, 1) + self.offsets_are_end = offsets_are_end + self.cta_group = ( + tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + self.occupancy = 1 + # Set specialized warp ids + self.epilog_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.sched_warp_id = 6 + self.threads_per_warp = 32 + self.threads_per_cta = self.threads_per_warp * len( + ( + *self.epilog_warp_id, + self.mma_warp_id, + self.tma_warp_id, + self.sched_warp_id, + ) + ) + self.threads_wo_sched = self.threads_per_warp * len( + ( + *self.epilog_warp_id, + self.mma_warp_id, + self.tma_warp_id, + ) + ) + self.num_regs_uniform_warps = 32 + self.num_regs_sched_warps = 64 + self.num_regs_epilogue_warps = 208 + + # Set barrier for epilogue sync and tmem ptr sync + self.epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=32 * len(self.epilog_warp_id), + ) + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)), + ) + self.sched_sync_barrier = pipeline.NamedBarrier( + barrier_id=3, + num_threads=self.threads_per_warp, + ) + self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + # TMEM offset for final accumulator + self.tmem_final_offset = 384 + + def _setup_attributes(self): + """Set up configurations that are dependent on GEMM inputs + + This method configures various attributes based on the input tensor properties + (data types, leading dimensions) and kernel settings: + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B + - Computing epilogue subtile + - Setting up A/B/C stage counts in shared memory + - Computing A/B/C shared memory layout + - Computing tensor memory allocation columns + """ + # Configure tiled mma + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + # Compute mma/cluster/tile shapes + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.mma_tiler = ( + self.mma_tiler[0], + self.mma_tiler[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.mma_a_tiler = (1, self.mma_tiler[1], self.mma_tiler[2]) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cluster_tile_shape_mnk = ( + self.cluster_shape_mn[0] * self.cta_tile_shape_mnk[0], + self.cluster_shape_mn[1] * self.cta_tile_shape_mnk[1], + self.cta_tile_shape_mnk[2], + ) + + # Compute cluster layout + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + + # Compute number of multicast CTAs for A/B + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + self.cta_tile_shape_mnk_with_mcast = ( + self.cta_tile_shape_mnk[0] // self.num_mcast_ctas_a, + self.cta_tile_shape_mnk[1] // self.num_mcast_ctas_b, + self.cta_tile_shape_mnk[2], + ) + + # Compute epilogue subtile + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + + # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory + ( + self.num_acc_stage, + self.num_ab_stage, + self.num_c_stage, + self.num_tile_stage, + ) = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.epi_tile, + self.c_dtype, + self.c_layout, + self.num_smem_capacity, + self.occupancy, + ) + + # Compute A/B/C shared memory layout + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.b_dtype, + self.num_ab_stage, + ) + self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile, + self.num_c_stage, + ) + + # Compute the number of tensor memory allocation columns + self.num_tmem_alloc_cols = 512 + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + offset_mapping: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + """Execute the GEMM operation in steps: + - Setup static attributes before smem/grid/tma computation + - Setup TMA load/store atoms and tensors + - Compute grid size with regard to hardware constraints + - Define shared storage for kernel + - Launch the kernel synchronously + + :param a: Input tensor A + :type a: cute.Tensor + :param b: Input tensor B + :type b: cute.Tensor + :param c: Output tensor C + :type c: cute.Tensor + :param offset_mapping: Mapping from group index to m offset + :type gidx_mapping: cute.Tensor + :param max_active_clusters: Maximum number of active clusters + :type max_active_clusters: cutlass.Constexpr + :param stream: CUDA stream for asynchronous execution + :type stream: cuda.CUstream + :param epilogue_op: Optional elementwise lambda function to apply to the output tensor + :type epilogue_op: cutlass.Constexpr + :raises TypeError: If input data types are incompatible with the MMA instruction. + """ + + def add_batch_mode(tensor: cute.Tensor) -> cute.Tensor: + return cute.make_tensor( + tensor.iterator, + cute.prepend(tensor.layout, cute.make_layout(1), up_to_rank=3), + ) + + a = add_batch_mode(a) + b = add_batch_mode(b) + c = add_batch_mode(c) + + # Permute tensor modes from torch to cute convention + # A: (L, M, K) -> (M, K, L) + a = cute.make_tensor(a.iterator, cute.select(a.layout, [1, 2, 0])) + # B: (L, K, N) -> (N, K, L) + b = cute.make_tensor(b.iterator, cute.select(b.layout, [2, 1, 0])) + # C: (L, M, N) -> (M, N, L) + c = cute.make_tensor(c.iterator, cute.select(c.layout, [1, 2, 0])) + + # Setup static attributes before smem/grid/tma computation + self.a_dtype: type[cutlass.Numeric] = a.element_type + self.b_dtype: type[cutlass.Numeric] = b.element_type + self.c_dtype: type[cutlass.Numeric] = c.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + + # Setup group count + self.group_count = b.shape[2] + + # Check if input data types are compatible with MMA instruction + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = self._get_tma_atom_kind(atom_thr_size, self.is_a_mcast) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + + # Setup TMA load for B + b_op = self._get_tma_atom_kind(atom_thr_size, self.is_b_mcast) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size + + # Construct C tensor for contiguous layout + tensor_c = cute.make_tensor( + c.iterator, + cute.make_layout( + (c.shape[0], c.shape[1], c.shape[0]), + stride=(c.layout.stride[0], c.layout.stride[1], c.layout.stride[0]), + ), + ) + # Setup TMA store for C + tma_atom_c = None + tma_tensor_c = None + c_cta_v_layout = cute.composition( + cute.make_identity_layout(c.shape), self.epi_tile + ) + epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + tensor_c, + epi_smem_layout, + c_cta_v_layout, + ) + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + c, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters + ) + + self.buffer_align_bytes = 1024 + + c_smem_size = cute.cosize(self.c_smem_layout_staged.outer) + + # Define shared storage for kernel + @cute.struct + class SharedStorage: + # (start_m, tile_id_y, tile_id_z, margin_m) + sInfo: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int32, 4 * self.num_tile_stage], + # 1 byte alignment + 1, + ] + ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + sched_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_tile_stage * 2] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, + c_smem_size, + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + sA: cute.struct.Align[ + cute.struct.MemRange[ + self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer) + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + sB: cute.struct.Align[ + cute.struct.MemRange[ + self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer) + ], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + # Launch the kernel synchronously + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c, + tensor_c, + offset_mapping, + self.group_count, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + epilogue_op, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + min_blocks_per_mp=1, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + tensor_c: cute.Tensor, + offset_mapping: cute.Tensor, + group_count: cutlass.Int32, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: cute.Layout | cute.ComposedLayout | None, + epi_tile: cute.Tile, + tile_sched_params: scheduler_utils.PersistentTileSchedulerParams, + epilogue_op: cutlass.Constexpr, + ): + """GPU device kernel performing the Persistent batched GEMM computation. + """ + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + lane_idx = cute.arch.lane_idx() + + # + # Prefetch tma desc + # + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + # Coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + # + # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier + # + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # Initialize mainloop ab_pipeline (barrier) and states + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilog_warp_id) * ( + 2 if use_2cta_instrs else 1 + ) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Initialize tile info pipeline (barrier) and states + sched_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * 1, + ) + sched_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_wo_sched, + ) + sched_pipeline = pipeline.PipelineAsync.create( + barrier_storage=storage.sched_mbar_ptr.data_ptr(), + num_stages=self.num_tile_stage, + producer_group=sched_pipeline_producer_group, + consumer_group=sched_pipeline_consumer_group, + defer_sync=True, + ) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=self.tmem_alloc_barrier, + allocator_warp_id=self.epilog_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # + # Setup smem tensor A/B/C + # + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC = storage.sC.get_tensor( + c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner + ) + # (MMA, MMA_M, MMA_K, STAGE) + sA = storage.sA.get_tensor( + a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner + ) + # (MMA, MMA_N, MMA_K, STAGE) + sB = storage.sB.get_tensor( + b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner + ) + + # (start_m, tile_id_y, tile_id_z, margin_m) + info_layout = cute.make_layout((4, self.num_tile_stage), stride=(1, 4)) + sInfo = storage.sInfo.get_tensor(info_layout) + + # + # Compute multicast mask for A/B buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + # + # Local_tile partition global tensors + # + # (bM, bK, loopM, loopK, loopL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_a_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, loopN, loopK, loopL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bN, loopM, loopN, loopL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + gC_mnl_simt = cute.local_tile( + tensor_c, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # + # Partition global tensor for TiledMMA_A/B/C + # + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + # (MMA, MMA_M, MMA_K, loopM, loopK, loopL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_N, MMA_K, loopN, loopK, loopL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_N, loopM, loopN, loopL) + tCgC = thr_mma.partition_C(gC_mnl) + tCgC_simt = thr_mma.partition_C(gC_mnl_simt) + + # + # Partition global/shared tensor for TMA load A/B + # + # TMA load A partition_S/D + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA load B partition_S/D + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # + # Cluster wait before tensor memory alloc + # + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # + # Specialized Schedule warp + # + if warp_idx == self.sched_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_sched_warps) + # + # Persistent tile scheduling loop + # + tile_sched = ContiguousOffsetGroupedGemmTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + # First tile + work_tile = tile_sched.initial_work_tile_info() + + sched_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_tile_stage + ) + + last_search_m_block_end = cutlass.Int32(0) + next_group_m_block_start = cutlass.Int32(0) + next_group_m_block_end = cutlass.Int32(0) + next_group_idx = cutlass.Int32(0) + cur_m_offset = cutlass.Int32(0) + cur_m_start = cutlass.Int32(0) + + not_last_tile = cutlass.Boolean(1) + while not_last_tile: + sched_pipeline.producer_acquire(sched_producer_state) + + # initialize the tile info + cluster_tile_coord_mnl = work_tile.tile_idx + + # mn_tile offset + m_tile_offset = self.mma_tiler[0] * block_in_cluster_coord_vmnk[1] + n_tile_coord = ( + block_in_cluster_coord_vmnk[2] + + cluster_tile_coord_mnl[1] * self.cluster_shape_mn[1] + ) + + ( + last_search_m_block_end, + next_group_m_block_start, + next_group_m_block_end, + next_group_idx, + cur_m_offset, + cur_m_start, + ) = self.group_search( + group_count, + cluster_tile_coord_mnl[0], + last_search_m_block_end, + next_group_m_block_start, + next_group_m_block_end, + next_group_idx, + cur_m_offset, + offset_mapping, + ) + + sInfo_pipe = sInfo[(None, sched_producer_state.index)] + + not_last_tile = next_group_idx <= group_count + m_remaining_in_cur_group = ( + next_group_m_block_start - cur_m_start - m_tile_offset + ) + with cute.arch.elect_one(): + sInfo_pipe[0] = cur_m_start + m_tile_offset + sInfo_pipe[1] = n_tile_coord + sInfo_pipe[2] = next_group_idx - 1 + sInfo_pipe[3] = m_remaining_in_cur_group + + # fence view async shared + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.sched_sync_barrier.arrive_and_wait() + # commit tile info pipeline + sched_pipeline.producer_commit(sched_producer_state) + + # advance to next tile + sched_producer_state.advance() + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + sched_pipeline.producer_tail(sched_producer_state) + + # + # Specialized TMA load warp + # + if warp_idx == self.tma_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps) + + ab_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) + + sched_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + + # get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + sched_pipeline.consumer_wait(sched_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, sched_consumer_state.index)] + is_valid_tile = tile_info[2] < group_count + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + sched_pipeline.consumer_release(sched_consumer_state) + sched_consumer_state.advance() + + mcast_offset = ( + self.cta_tile_shape_mnk_with_mcast[0] * block_in_cluster_coord_vmnk[2] + + block_in_cluster_coord_vmnk[0] * self.cta_tile_shape_mnk[0] + ) + + while is_valid_tile: + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), loopK) + tAgA_slice = tAgA[(None, tile_info[0] + mcast_offset, None, 0)] + # ((atom_v, rest_v), loopK) + tBgB_slice = tBgB[(None, tile_info[1], None, tile_info[2])] + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + ab_producer_state.reset_count() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) + # + # Tma load loop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + tAgA_k = tAgA_slice[(None, ab_producer_state.count)] + tBgB_k = tBgB_slice[(None, ab_producer_state.count)] + tAsA_pipe = tAsA[(None, ab_producer_state.index)] + tBsB_pipe = tBsB[(None, ab_producer_state.index)] + + tma_bar = ab_pipeline.producer_get_barrier(ab_producer_state) + + # Conditionally wait for AB buffer empty + ab_pipeline.producer_acquire( + ab_producer_state, peek_ab_empty_status + ) + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 + ab_producer_state.advance() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) + + # TMA load A/B + cute.copy( + tma_atom_a, + tAgA_k, + tAsA_pipe, + tma_bar_ptr=tma_bar, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_k, + tBsB_pipe, + tma_bar_ptr=tma_bar, + mcast_mask=b_full_mcast_mask, + ) + + # + # Advance to next tile + # + sched_pipeline.consumer_wait(sched_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, sched_consumer_state.index)] + is_valid_tile = tile_info[2] < group_count + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + sched_pipeline.consumer_release(sched_consumer_state) + sched_consumer_state.advance() + + # + # Wait A/B buffer empty + # + ab_pipeline.producer_tail(ab_producer_state) + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps) + # + # Bar sync for retrieve tensor memory ptr from shared mem + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + ab_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_ab_stage + ) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + + sched_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + + # get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + sched_pipeline.consumer_wait(sched_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, sched_consumer_state.index)] + is_valid_tile = tile_info[2] < group_count + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + sched_pipeline.consumer_release(sched_consumer_state) + sched_consumer_state.advance() + + while is_valid_tile: + # Peek (try_wait) AB buffer full for k_tile = 0 + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait( + ab_consumer_state + ) + + # Wait for Acc buffer empty (once per tile) + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + + # Set tensor memory buffer for current tile + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + # + # Mma mainloop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + if is_leader_cta: + # Conditionally wait for AB buffer full + ab_pipeline.consumer_wait( + ab_consumer_state, peek_ab_full_status + ) + + num_kblocks = cute.size(tCrA, mode=[2]) + for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): + kblk_crd = (None, None, kblk_idx, ab_consumer_state.index) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblk_crd], + tCrB[kblk_crd], + tCtAcc, + ) + + # Enable accumulate after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + # Async arrive AB buffer empty + ab_pipeline.consumer_release(ab_consumer_state) + + # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 + ab_consumer_state.advance() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt: + if is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait( + ab_consumer_state + ) + + # + # Async arrive accumulator buffer full + # + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + sched_pipeline.consumer_wait(sched_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, sched_consumer_state.index)] + is_valid_tile = tile_info[2] < group_count + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + sched_pipeline.consumer_release(sched_consumer_state) + sched_consumer_state.advance() + + # + # Wait for accumulator buffer empty + # + acc_pipeline.producer_tail(acc_producer_state) + + # + # Specialized epilogue warps + # + if warp_idx <= self.epilog_warp_id[-1] and warp_idx >= self.epilog_warp_id[0]: + cute.arch.warpgroup_reg_alloc(self.num_regs_epilogue_warps) + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base_ = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Partition for epilogue + # + epi_tidx = tidx % 128 + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc, + ) = self.epilog_tmem_copy_and_partition( + epi_tidx, tCtAcc_base_, tCgC, epi_tile, use_2cta_instrs + ) + + tTR_rC = None + tiled_copy_r2s = None + simt_atom = None + tRS_rC = None + tRS_sC = None + bSG_sC = None + bSG_gC_partitioned = None + tTR_gC_partitioned = None + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition( + tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + ( + tma_atom_c, + bSG_sC, + bSG_gC_partitioned, + simt_atom, + tTR_gC_partitioned, + ) = self.epilog_gmem_copy_and_partition( + epi_tidx, tma_atom_c, tiled_copy_t2r, tCgC, tCgC_simt, epi_tile, sC + ) + + thr_mapping = cute.make_identity_tensor( + (self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1]) + ) + thr_mapping_mn = cute.flat_divide(thr_mapping, epi_tile) + + thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) + m_thr_offset = thr_copy_t2r.partition_D(thr_mapping_mn) + m_thr_offset = cute.group_modes(m_thr_offset, 3, cute.rank(m_thr_offset)) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + # Threads/warps participating in tma store pipeline + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilog_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, + producer_group=c_producer_group, + ) + + sched_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + sched_pipeline.consumer_wait(sched_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, sched_consumer_state.index)] + is_valid_tile = tile_info[2] < group_count + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + + num_prev_subtiles = cutlass.Int32(0) + + m_block_offset = block_in_cluster_coord_vmnk[0] * self.cta_tile_shape_mnk[0] + + while is_valid_tile: + n_block_offset = tile_info[1] * self.cta_tile_shape_mnk[1] + + # + # Slice to per mma tile index + # + bSG_gC = None + tTR_gC = None + # ((ATOM_V, REST_V), EPI_M, EPI_N) + bSG_gC = bSG_gC_partitioned[ + ( + None, + None, + None, + 0, + tile_info[1], + tile_info[0], + ) + ] + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) + tTR_gC = tTR_gC_partitioned[ + ( + None, + None, + None, + None, + None, + 0, + tile_info[1], + tile_info[0], + ) + ] + + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + + # + # Wait for accumulator buffer to be full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + tTR_gC = cute.group_modes(tTR_gC, 3, cute.rank(tTR_gC)) + + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + for subtile_idx in cutlass.range(subtile_cnt): + # + # Load accumulated result from tensor memory + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + if tile_info[3] >= m_block_offset + self.cta_tile_shape_mnk[0]: + # + # Convert to C type (TMA store path) + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) + tRS_rC.store(acc_vec) + + # + # Store C to shared memory + # + num_prev_subtiles = num_prev_subtiles + 1 + c_buffer = num_prev_subtiles % self.num_c_stage + cute.copy( + tiled_copy_r2s, + tRS_rC, + tRS_sC[(None, None, None, c_buffer)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.epilog_sync_barrier.arrive_and_wait() + + # + # TMA store C to global memory + # + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + else: + # + # Convert to C type (SIMT store path) + # + acc_vec = tTR_rAcc.load() + acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) + tTR_rC.store(acc_vec) + + tCpC = cute.make_rmem_tensor( + cute.make_layout(tTR_rC.shape), + cutlass.Boolean, + ) + m_thr_slice = m_thr_offset[(None, None, None, subtile_idx)] + + m_in_range = ( + m_thr_slice[(cute.size(m_thr_slice) - 1)][0] + + m_block_offset + < tile_info[3] + ) + n_in_range = ( + m_thr_slice[(cute.size(m_thr_slice) - 1)][1] + + n_block_offset + < tensor_c.shape[1] + ) + # fast-path + if ( + not ( + self.cta_tile_shape_mnk[0] == 64 + and not self.use_2cta_instrs + ) + and m_in_range + and n_in_range + ): + cute.copy( + simt_atom, + cute.flatten(tTR_rC), + cute.flatten(tTR_gC[(None, None, None, subtile_idx)]), + ) + else: + m_thr_slice = m_thr_offset[(None, None, None, subtile_idx)] + for i in cutlass.range(cute.size(tCpC), unroll_full=True): + tCpC[i] = ( + m_thr_slice[(i)][0] + m_block_offset < tile_info[3] + ) and ( + m_thr_slice[(i)][1] + n_block_offset + < tensor_c.shape[1] + ) + + # + # Store C to global memory + # + cute.copy( + simt_atom, + cute.flatten(tTR_rC), + cute.flatten(tTR_gC[(None, None, None, subtile_idx)]), + pred=cute.flatten(tCpC), + ) + + # + # Async arrive accumulator buffer empty + # + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Advance to next tile + # + sched_pipeline.consumer_release(sched_consumer_state) + sched_consumer_state.advance() + sched_pipeline.consumer_wait(sched_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, sched_consumer_state.index)] + is_valid_tile = tile_info[2] < group_count + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + + sched_pipeline.consumer_release(sched_consumer_state) + sched_consumer_state.advance() + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + # + # Wait for C store complete + # + c_pipeline.producer_tail() + + @cute.jit + def group_search( + self, + group_count: cutlass.Int32, + m_block: cutlass.Int32, + last_search_m_block_end: cutlass.Int32, + next_group_m_block_start: cutlass.Int32, + next_group_m_block_end: cutlass.Int32, + next_group_idx: cutlass.Int32, + cur_m_offset: cutlass.Int32, + offset_mapping: cute.Tensor, + ): + not_found = m_block >= next_group_m_block_end + new_next_m_boundary = cutlass.Int32(0) + if not_found: + next_group_idx = next_group_idx + 1 + while not_found and next_group_idx <= group_count: + idx = next_group_idx + if cutlass.const_expr(self.offsets_are_end): + idx = next_group_idx - 1 + new_next_m_boundary = offset_mapping[idx] + num_m_blocks = cute.ceil_div( + (new_next_m_boundary - next_group_m_block_start), + self.cluster_tile_shape_mnk[0], + ) + new_tile_count = num_m_blocks + next_group_m_block_end + not_found = m_block >= new_tile_count + + last_search_m_block_end = next_group_m_block_end + cur_m_offset = next_group_m_block_start + next_group_m_block_start = new_next_m_boundary + next_group_m_block_end = new_tile_count + if not_found: + next_group_idx = next_group_idx + 1 + + cur_m_start = cur_m_offset + self.cluster_tile_shape_mnk[0] * ( + m_block - last_search_m_block_end + ) + + return ( + last_search_m_block_end, + next_group_m_block_start, + next_group_m_block_end, + next_group_idx, + cur_m_offset, + cur_m_start, + ) + + def epilog_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tAcc: cute.Tensor, + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + use_2cta_instrs: cutlass.Boolean | bool, + ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param tAcc: The accumulator tensor to be copied and partitioned + :type tAcc: cute.Tensor + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param use_2cta_instrs: Whether use_2cta_instrs is enabled + :type use_2cta_instrs: bool + + :return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where: + - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + - tTR_tAcc: The partitioned accumulator tensor + - tTR_rAcc: The accumulated tensor in register used to hold t2r results + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + # Make tiledCopy for tensor memory load + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + epi_tile, + use_2cta_instrs, + ) + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) + tAcc_epi = cute.flat_divide( + tAcc[((None, None), 0, 0, None)], + epi_tile, + ) + # (EPI_TILE_M, EPI_TILE_N) + tiled_copy_t2r = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)] + ) + + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, loopM, loopN, loopL) + gC_mnl_epi = cute.flat_divide( + gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile + ) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, loopM, loopN, loopL) + tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) + + # (T2R, T2R_M, T2R_N) + tTR_rAcc = cute.make_rmem_tensor( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype + ) + + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc + + def epilog_smem_copy_and_partition( + self, + tiled_copy_t2r: cute.TiledCopy, + tTR_rC: cute.Tensor, + tidx: cutlass.Int32, + sC: cute.Tensor, + ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination). + + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + :type tiled_copy_t2r: cute.TiledCopy + :param tTR_rC: The partitioned accumulator tensor + :type tTR_rC: cute.Tensor + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + :type sepi: cute.Tensor + + :return: A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where: + - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) + - tRS_rC: The partitioned tensor C (register source) + - tRS_sC: The partitioned tensor C (smem destination) + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + copy_atom_r2s = sm100_utils.get_smem_store_op( + self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r + ) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + # (R2S, R2S_M, R2S_N, PIPE_D) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sC = thr_copy_r2s.partition_D(sC) + # (R2S, R2S_M, R2S_N) + tRS_rC = tiled_copy_r2s.retile(tTR_rC) + return tiled_copy_r2s, tRS_rC, tRS_sC + + def epilog_gmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tma_atom_c: cute.CopyAtom, + tiled_copy_t2r: cute.TiledCopy, + gC_mnl_tma: cute.Tensor, + gC_mnl_simt: cute.Tensor, + epi_tile: cute.Tile, + sC: cute.Tensor, + ) -> tuple[ + cute.CopyAtom, cute.Tensor, cute.Tensor, cute.CopyAtom, cute.Tensor, cute.Tensor + ]: + """Make tiledCopy for global memory store, then use it to: + - partition register array (source) and global memory (destination) for none TMA store version; + - partition shared memory (source) and global memory (destination) for TMA store version. + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param atom: The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version + :type atom: cute.CopyAtom or cute.TiledCopy + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + + :return: A tuple containing: + - For TMA store: (tma_atom_c, bSG_sC, bSG_gC) where: + - tma_atom_c: The TMA copy atom + - bSG_sC: The partitioned shared memory tensor C + - bSG_gC: The partitioned global tensor C + - For non-TMA store: (simt_atom, tTR_gC) where: + - simt_atom: The SIMT copy atom + - tTR_gC: The partitioned global tensor C + :rtype: Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor] + """ + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, loopM, loopN, loopL) + gC_epi_tma = cute.flat_divide( + gC_mnl_tma[((None, None), 0, 0, None, None, None)], epi_tile + ) + gC_epi_simt = cute.flat_divide( + gC_mnl_simt[((None, None), 0, 0, None, None, None)], epi_tile + ) + # TMA store + sC_for_tma_partition = cute.group_modes(sC, 0, 2) + gC_for_tma_partition = cute.group_modes(gC_epi_tma, 0, 2) + # ((ATOM_V, REST_V), EPI_M, EPI_N) + # ((ATOM_V, REST_V), EPI_M, EPI_N, loopM, loopN, loopL) + bSG_sC, bSG_gC = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + sC_for_tma_partition, + gC_for_tma_partition, + ) + # STG store + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, loopM, loopN, loopL) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_gC = thr_copy_t2r.partition_D(gC_epi_simt) + simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) + + return tma_atom_c, bSG_sC, bSG_gC, simt_atom, tTR_gC + + @staticmethod + def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: tuple[int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + epi_tile: cute.Tile, + c_dtype: type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + num_smem_capacity: int, + occupancy: int, + ) -> tuple[int, int, int, int]: + """Computes the number of stages for A/B/C operands based on heuristics. + + :param tiled_mma: The tiled MMA object defining the core computation. + :type tiled_mma: cute.TiledMma + :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler. + :type mma_tiler_mnk: tuple[int, int, int] + :param a_dtype: Data type of operand A. + :type a_dtype: type[cutlass.Numeric] + :param b_dtype: Data type of operand B. + :type b_dtype: type[cutlass.Numeric] + :param epi_tile: The epilogue tile shape. + :type epi_tile: cute.Tile + :param c_dtype: Data type of operand C (output). + :type c_dtype: type[cutlass.Numeric] + :param c_layout: Layout of operand C. + :type c_layout: utils.LayoutEnum + :param num_smem_capacity: Total available shared memory capacity in bytes. + :type num_smem_capacity: int + :param occupancy: Target number of CTAs per SM (occupancy). + :type occupancy: int + + :return: A tuple containing the computed number of stages for: + (ACC stages, A/B operand stages, C stages, tile info stages) + :rtype: tuple[int, int, int, int] + """ + num_acc_stage = 2 + + # Default C stages + num_c_stage = 2 + + # Default Tile info stages + num_tile_stage = 2 + + # Calculate smem layout and size for one stage of A, B, and C + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + tiled_mma, + mma_tiler_mnk, + a_dtype, + 1, # a tmp 1 stage is provided + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + tiled_mma, + mma_tiler_mnk, + b_dtype, + 1, # a tmp 1 stage is provided + ) + c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + 1, + ) + + ab_bytes_per_stage = cute.size_in_bytes( + a_dtype, a_smem_layout_stage_one + ) + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + # 1024B alignment + mbar_helpers_bytes = 1024 + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + c_bytes = c_bytes_per_stage * num_c_stage + + # Calculate A/B stages: + # Start with total smem per CTA (capacity / occupancy) + # Subtract reserved bytes and initial C stages bytes + # Divide remaining by bytes needed per A/B stage + num_ab_stage = ( + num_smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes) + ) // ab_bytes_per_stage + + # Refine epilogue stages: + # Calculate remaining smem after allocating for A/B stages and reserved bytes + # Add remaining unused smem to epilogue + num_c_stage += ( + num_smem_capacity + - occupancy * ab_bytes_per_stage * num_ab_stage + - occupancy * (mbar_helpers_bytes + c_bytes) + ) // (occupancy * c_bytes_per_stage) + return num_acc_stage, num_ab_stage, num_c_stage, num_tile_stage + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + max_active_clusters: cutlass.Constexpr, + ) -> tuple[scheduler_utils.PersistentTileSchedulerParams, tuple[int, int, int]]: + """Use persistent tile scheduler to compute the grid size for the output tensor C. + + :param c: The output tensor C + :type c: cute.Tensor + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + :param max_active_clusters: Maximum number of active clusters. + :type max_active_clusters: cutlass.Constexpr + + :return: A tuple containing: + - tile_sched_params: Parameters for the persistent tile scheduler. + - grid: Grid shape for kernel launch. + :rtype: Tuple[scheduler_utils.PersistentTileSchedulerParams, tuple[int, int, int]] + """ + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = scheduler_utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl, raster_along_m=False + ) + + grid = scheduler_utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + + return tile_sched_params, grid + + @staticmethod + def _get_tma_atom_kind( + atom_sm_cnt: cutlass.Int32, mcast: cutlass.Boolean + ) -> cpasync.CopyBulkTensorTileG2SMulticastOp | cpasync.CopyBulkTensorTileG2SOp: + """Select the appropriate TMA copy atom based on the number of SMs and the multicast flag. + + :param atom_sm_cnt: The number of SMs + :type atom_sm_cnt: cutlass.Int32 + :param mcast: The multicast flag + :type mcast: cutlass.Boolean + + :return: The appropriate TMA copy atom kind + :rtype: cpasync.CopyBulkTensorTileG2SMulticastOp or cpasync.CopyBulkTensorTileG2SOp + + :raise ValueError: If the atom_sm_cnt is invalid + """ + if atom_sm_cnt == 2 and mcast: + return cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.TWO) + elif atom_sm_cnt == 2 and not mcast: + return cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.TWO) + elif atom_sm_cnt == 1 and mcast: + return cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.ONE) + elif atom_sm_cnt == 1 and not mcast: + return cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE) + + raise ValueError(f"Invalid atom_sm_cnt: {atom_sm_cnt} and {mcast}") + + @staticmethod + def is_valid_dtypes( + ab_dtype: type[cutlass.Numeric], + acc_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + ) -> bool: + """Check if the dtypes are valid + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param acc_dtype: The data type of the accumulator + :type acc_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + + :return: True if the dtypes are valid, False otherwise + :rtype: bool + """ + is_valid = True + if ab_dtype not in { + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.BFloat16, + cutlass.Float16, + }: + is_valid = False + if acc_dtype not in {cutlass.Float32}: + is_valid = False + if c_dtype not in {cutlass.Float32, cutlass.Float16, cutlass.BFloat16}: + is_valid = False + return is_valid + + @staticmethod + def is_valid_mma_tiler_and_cluster_shape( + use_2cta_instrs: bool, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + ) -> bool: + """Check if the mma tiler and cluster shape are valid + + :param use_2cta_instrs: Whether to use 2 CTA groups + :type use_2cta_instrs: bool + :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster + :type cluster_shape_mn: Tuple[int, int] + + :return: True if the mma tiler and cluster shape are valid, False otherwise + :rtype: bool + """ + is_valid = True + # Skip invalid mma tile shape + if not ( + (not use_2cta_instrs and mma_tiler_mn[0] in [64, 128]) + or (use_2cta_instrs and mma_tiler_mn[0] in [128, 256]) + ): + is_valid = False + # Skip invalid mma tile n + if mma_tiler_mn[1] not in (128, 256): + is_valid = False + # Skip illegal cluster shape + if cluster_shape_mn[0] % (2 if use_2cta_instrs else 1) != 0: + is_valid = False + # Skip invalid cluster shape + is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if ( + cluster_shape_mn[0] * cluster_shape_mn[1] > 16 + or cluster_shape_mn[0] <= 0 + or cluster_shape_mn[1] <= 0 + or not is_power_of_2(cluster_shape_mn[0]) + or not is_power_of_2(cluster_shape_mn[1]) + ): + is_valid = False + return is_valid + + @staticmethod + def is_valid_tensor_alignment( + m: int, + n: int, + k: int, + l: int, + ab_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """Check if the tensor alignment is valid + + :param m: The number of rows in the A tensor + :type m: int + :param n: The number of columns in the B tensor + :type n: int + :param k: The number of columns in the A tensor + :type k: int + :param l: The number of columns in the C tensor + :type l: int + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param a_major: The major axis of the A tensor + :type a_major: str + :param b_major: The major axis of the B tensor + :type b_major: str + :param c_major: The major axis of the C tensor + :type c_major: str + + :return: True if the problem shape is valid, False otherwise + :rtype: bool + """ + is_valid = True + + def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // dtype.width + return num_major_elements % num_contiguous_elements == 0 + + if ( + not check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) + or not check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) + or not check_contigous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + ): + is_valid = False + return is_valid diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_dense_blockscaled_static_persistent_impl.py b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_dense_blockscaled_static_persistent_impl.py new file mode 100644 index 000000000..fafa4ed26 --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_dense_blockscaled_static_persistent_impl.py @@ -0,0 +1,2170 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from typing import Literal + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.nvgpu.common import OperandMajorMode +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +""" +This example provides an experimental implementation of the SM100 batched dense blockscaled GEMM kernel, please note that the APIs and implementation details related to this kernel may change in future releases. + +A high-performance persistent batched dense blockscaled GEMM example for the NVIDIA Blackwell SM100 architecture +using CUTE DSL. +- Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") for MXF8 input type and can only be row-major("K") for MXF4/NVF4 input type +- Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") for MXF8 input type and can only be row-major("K") for MXF4/NVF4 input type +- Matrix C is MxNxL, L is batch dimension, C can be row-major("N") or column-major("M") +- Matrix SFA layout is filled internally according to A shape and BlockScaledBasicChunk, which has M×ceil_div(K, sf_vec_size)×L elements respectively +- Matrix SFB layout is filled internally according to B shape and BlockScaledBasicChunk, which has N×ceil_div(K, sf_vec_size)×L elements respectively + +This GEMM kernel supports the following features: + - Utilizes Tensor Memory Access (TMA) for efficient memory operations + - Utilizes Blackwell's tcgen05.mma for matrix multiply-accumulate (MMA) operations (including 2cta mma instructions) + - Implements TMA multicast with cluster to reduce L2 memory traffic + - Support persistent tile scheduling to better overlap memory load/store with mma between tiles + - Support warp specialization to avoid explicit pipelining between mainloop load and mma + +This GEMM works as follows: +1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. +2. MMA warp: + - Load scale factor A/B from shared memory (SMEM) to tensor memory (TMEM) using tcgen05.cp instruction. + - Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction. +3. EPILOGUE warp: + - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. + - Type convert C matrix to output type. + - Optionally store C matrix from registers (RMEM) to shared memory (SMEM) to global memory (GMEM) with TMA operations, + or directly store C matrix from registers (RMEM) to global memory (GMEM) without TMA operations. + - Optionally accept an elementwise lambda function epilogue_op to apply to the output tensor: + e.g., relu can set epilogue_op = lambda x: cute.where(x > 0, x, cute.full_like(x, 0)) + +SM100 tcgen05.mma.kind.block_scale instructions operate as follows: +- Read matrix A from SMEM +- Read matrix B from SMEM +- Read scalefactor A from TMEM +- Read scalefactor B from TMEM +- Write accumulator to TMEM +The accumulator in TMEM must then be loaded to registers before writing back to GMEM. + +Input arguments to this example is shown below: + +.. code-block:: bash + + python examples/blackwell/dense_blockscaled_gemm_persistent.py \ + --ab_dtype Float4E2M1FN --sf_dtype Float8E8M0FNU --sf_vec_size 16 \ + --c_dtype Float16 \ + --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,1024,1 + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/dense_blockscaled_gemm_persistent.py \ + --ab_dtype Float4E2M1FN --sf_dtype Float8E8M0FNU --sf_vec_size 16 \ + --c_dtype Float16 \ + --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,1024,1 \ + --warmup_iterations 1 --iterations 10 --skip_ref_check + + +Constraints: +* Supported input data types: mxf8, mxf4, nvf4, and mixed MXFP4 x MXFP8 and MXFP8 x MXFP4 + see detailed valid dtype combinations in below PersistentDenseBlockScaledGemmKernel class documentation +* A and B may have the same data type, or be a mixed MXFP4 x MXFP8 pair + (e.g. Float4E2M1FN x Float8E4M3FN). Mixed inputs use MX scaling + (Float8E8M0FNU scale factor, sf_vec_size 32). +* Mma tiler M must be 128 or 256(use_2cta_instrs) +* Mma tiler N must be 64/128/192/256 +* Cluster shape M/N must be positive and power of 2, total cluster size <= 16 +* Cluster shape M must be multiple of 2 if Mma tiler M is 256(use_2cta_instrs) +* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, + i.e, number of elements is a multiple of 16 and 32 for Float8 and Float4, respectively. +""" + + +class PersistentDenseBlockScaledGemmKernel: + """This class implements batched matrix multiplication (C = A x SFA x B x SFB) with support for various data types + and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization. + + :param sf_vec_size: Scalefactor vector size. + :type sf_vec_size: int + :param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N) + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing + :type cluster_shape_mn: Tuple[int, int] + + :note: A and B may have the same element type, or form a mixed MXFP4 x MXFP8 + or MXFP8 x MXFP4 pair. + + :note: Supported combinations of A/B data types, SF data typs and SF vector size: + - MXF8: A/B: Float8E5M2/Float8E4M3FN + SF: Float8E8M0FNU + sf_vec_size: 32 + - MXF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU + sf_vec_size: 32 + - NVF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU/Float8E4M3FN + sf_vec_size: 16 + - Mixed MXFP4 x MXFP8: one of A/B is Float4E2M1FN and the other is + Float8E5M2/Float8E4M3FN + SF: Float8E8M0FNU + sf_vec_size: 32 + + :note: Supported accumulator data types: + - Float32 + + :note: Supported C data types: + - Float32 + - Float16/BFloat16 + - Float8E4M3FN/Float8E5M2 + :note: Constraints: + - MMA tiler M must be 128 or 256 (use_2cta_instrs) + - MMA tiler N must be 64/128/192/256 + - Cluster shape M must be multiple of 2 if Mma tiler M is 256 + - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 + - Also, Cluster shape M/N must be <= 4 for scale factor multicasts due to limited size of scale factors + + Example: + >>> gemm = PersistentDenseBlockScaledGemmKernel( + ... sf_vec_size=16, + ... mma_tiler_mn=(256, 128), + ... cluster_shape_mn=(2, 1) + ... ) + >>> gemm(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, max_active_clusters, stream) + """ + + def __init__( + self, + sf_vec_size: int, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + ): + """Initializes the configuration for a Blackwell dense GEMM kernel. + + This configuration includes several key aspects: + + 1. MMA Instruction Settings (tcgen05): + - acc_dtype: Data types for MMA accumulator, always set to Float32 + - sf_vec_size: Scalefactor A/B vector size. + - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler. + + 2. Cluster Shape: + - cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster. + + :param sf_vec_size: Scalefactor vector size. + :type sf_vec_size: int + :param mma_tiler_mn: Tuple (M, N) shape of the MMA instruction. + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster. + :type cluster_shape_mn: Tuple[int, int] + """ + self.acc_dtype = cutlass.Float32 + self.sf_vec_size = sf_vec_size + self.use_2cta_instrs = mma_tiler_mn[0] == 256 + self.cluster_shape_mn = cluster_shape_mn + # K dimension is deferred in _setup_attributes + self.mma_tiler = (*mma_tiler_mn, 1) + + self.cta_group = ( + tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + self.occupancy = 1 + # Set specialized warp ids + self.epilog_warp_id = ( + 0, + 1, + 2, + 3, + ) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.threads_per_warp = 32 + self.threads_per_cta = self.threads_per_warp * len( + (self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id) + ) + # Set barrier id for epilogue sync and tmem ptr sync + self.epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.threads_per_warp * len(self.epilog_warp_id), + ) + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=self.threads_per_warp + * len((self.mma_warp_id, *self.epilog_warp_id)), + ) + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") + + def _setup_attributes(self): + """Set up configurations that are dependent on GEMM inputs + + This method configures various attributes based on the input tensor properties + (data types, leading dimensions) and kernel settings: + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B/SFA/SFB + - Computing epilogue subtile + - Setting up A/B/SFA/SFB/C stage counts in shared memory + - Computing A/B/SFA/SFB/C shared memory layout + """ + # Compute mma instruction shapes + # (MMA_Tile_Shape_M, MMA_Tile_Shape_N, MMA_Inst_Shape_K) + self.mma_inst_shape_mn = ( + self.mma_tiler[0], + self.mma_tiler[1], + ) + # (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K) + self.mma_inst_shape_mn_sfb = ( + self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1), + cute.round_up(self.mma_inst_shape_mn[1], 128), + ) + + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape_mn, + ) + + tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mma_inst_shape_mn_sfb, + ) + + # Compute mma/cluster/tile shapes + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.mma_tiler = ( + self.mma_inst_shape_mn[0], + self.mma_inst_shape_mn[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.mma_tiler_sfb = ( + self.mma_inst_shape_mn_sfb[0], + self.mma_inst_shape_mn_sfb[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cta_tile_shape_mnk_sfb = ( + self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler_sfb[1], + self.mma_tiler_sfb[2], + ) + + # Compute cluster layout + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.cluster_layout_sfb_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma_sfb.thr_id.shape,), + ) + + # Compute number of multicast CTAs for A/B + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 + + # Compute epilogue subtile + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + self.epi_tile_n = cute.size(self.epi_tile[1]) + + # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.smem_alloc_a_dtype, + self.smem_alloc_b_dtype, + self.epi_tile, + self.c_dtype, + self.c_layout, + self.sf_dtype, + self.sf_vec_size, + self.smem_capacity, + self.occupancy, + ) + + # Compute A/B/SFA/SFB/C shared memory layout. A/B use the SMEM + # allocation dtype (Int8 for an unpacked sub-byte operand) so the + # layout cosize reflects the unpacked container size. + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.smem_alloc_a_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.smem_alloc_b_dtype, + self.num_ab_stage, + ) + self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile, + self.num_c_stage, + ) + + # Overlap and double buffer accumulator when num_acc_stage == 1 for cta_tile_n = 256 case + self.overlapping_accum = self.num_acc_stage == 1 + + # Compute number of TMEM columns for SFA/SFB/Accumulator + sf_atom_mn = 32 + self.num_sfa_tmem_cols = ( + self.cta_tile_shape_mnk[0] // sf_atom_mn + ) * mma_inst_tile_k + self.num_sfb_tmem_cols = ( + self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn + ) * mma_inst_tile_k + self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols + self.num_accumulator_tmem_cols = ( + self.cta_tile_shape_mnk[1] * self.num_acc_stage + if not self.overlapping_accum + else self.cta_tile_shape_mnk[1] * 2 - self.num_sf_tmem_cols + ) + + # Only when overlapping_accum is enabled, we need to release accumulator buffer early in epilogue + self.iter_acc_early_release_in_epilogue = ( + self.num_sf_tmem_cols // self.epi_tile_n + ) + + @cute.jit + def __call__( + self, + a_ptr: cute.Pointer, + b_ptr: cute.Pointer, + sfa_ptr: cute.Pointer, + sfb_ptr: cute.Pointer, + c_ptr: cute.Pointer, + layouts: cutlass.Constexpr[ + tuple[OperandMajorMode, OperandMajorMode, utils.LayoutEnum] + ], + problem_mnkl: tuple[int, int, int, int], + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + """Execute the GEMM operation in steps: + - Setup static attributes before smem/grid/tma computation + - Setup TMA load/store atoms and tensors + - Compute grid size with regard to hardware constraints + - Define shared storage for kernel + - Launch the kernel synchronously + + :param a_tensor: Input tensor A + :type a_tensor: cute.Tensor + :param b_tensor: Input tensor B + :type b_tensor: cute.Tensor + :param sfa_tensor: Scale factor tensor A + :type sfa_tensor: cute.Tensor + :param sfb_tensor: Scale factor tensor B + :type sfb_tensor: cute.Tensor + :param c_tensor: Output tensor C + :type c_tensor: cute.Tensor + :param max_active_clusters: Maximum number of active clusters + :type max_active_clusters: cutlass.Constexpr + :param stream: CUDA stream for asynchronous execution + :type stream: cuda.CUstream + :param epilogue_op: Optional elementwise lambda function to apply to the output tensor + :type epilogue_op: cutlass.Constexpr + :raises TypeError: If input data types are incompatible with the MMA instruction. + """ + # Setup static attributes before smem/grid/tma computation + self.a_dtype: type[cutlass.Numeric] = a_ptr.value_type + self.b_dtype: type[cutlass.Numeric] = b_ptr.value_type + self.sf_dtype: type[cutlass.Numeric] = sfa_ptr.value_type + self.c_dtype: type[cutlass.Numeric] = c_ptr.value_type + + # For mixed-precision (mxf8f6f4) inputs where A and B have different + # element widths (e.g. MXFP4 x MXFP8), TMA must unpack the narrower, + # sub-byte operand into 1B-per-element SMEM containers so that A and B + # share a uniform byte-per-element SMEM layout for the MMA. The MMA + # fragments are still typed by the real element dtype via ``tiled_mma``. + self.mxf8f6f4 = self.needs_unpack_tma(self.a_dtype, self.b_dtype) + self.smem_alloc_a_dtype = ( + cutlass.Int8 if (self.mxf8f6f4 and self.a_dtype.width < 8) else self.a_dtype + ) + self.smem_alloc_b_dtype = ( + cutlass.Int8 if (self.mxf8f6f4 and self.b_dtype.width < 8) else self.b_dtype + ) + + m, n, k, l = problem_mnkl + self.a_major_mode, self.b_major_mode, self.c_layout = layouts + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + a_layout = cute.make_ordered_layout((m, cute.assume(k, 32), l), order=(0, 1, 2)) + if cutlass.const_expr(self.a_major_mode == OperandMajorMode.K): + a_layout = cute.make_ordered_layout( + (cute.assume(m, 32), k, l), order=(1, 0, 2) + ) + b_layout = cute.make_ordered_layout((n, cute.assume(k, 32), l), order=(0, 1, 2)) + if cutlass.const_expr(self.b_major_mode == OperandMajorMode.K): + b_layout = cute.make_ordered_layout( + (cute.assume(n, 32), k, l), order=(1, 0, 2) + ) + c_layout = cute.make_ordered_layout((cute.assume(m, 32), n, l), order=(0, 1, 2)) + if cutlass.const_expr(self.c_layout == utils.LayoutEnum.ROW_MAJOR): + c_layout = cute.make_ordered_layout( + (m, cute.assume(n, 32), l), order=(1, 0, 2) + ) + a_tensor = cute.make_tensor(a_ptr, a_layout) + b_tensor = cute.make_tensor(b_ptr, b_layout) + c_tensor = cute.make_tensor(c_ptr, c_layout) + + # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout + # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL) + sfa_layout = blockscaled_utils.tile_atom_to_shape_SF( + a_tensor.shape, self.sf_vec_size + ) + sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout) + + # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL) + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF( + b_tensor.shape, self.sf_vec_size + ) + sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout) + + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape_mn, + ) + + tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mma_inst_shape_mn_sfb, + ) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a_tensor, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=self.smem_alloc_a_dtype + if (self.mxf8f6f4 and self.a_dtype.width < 8) + else None, + ) + + # Setup TMA load for B + b_op = sm100_utils.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, tiled_mma.thr_id + ) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b_tensor, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=self.smem_alloc_b_dtype + if (self.mxf8f6f4 and self.b_dtype.width < 8) + else None, + ) + + # Setup TMA load for SFA + sfa_op = sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + sfa_smem_layout = cute.slice_( + self.sfa_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + sfa_tensor, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # Setup TMA load for SFB + sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB( + self.cluster_shape_mn, tiled_mma.thr_id + ) + sfb_smem_layout = cute.slice_( + self.sfb_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + sfb_tensor, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Int16, + ) + + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): + x = tma_tensor_sfb.stride[0][1] + y = cute.ceil_div(tma_tensor_sfb.shape[0][1], 4) + + new_shape = ( + (tma_tensor_sfb.shape[0][0], ((2, 2), y)), + tma_tensor_sfb.shape[1], + tma_tensor_sfb.shape[2], + ) + # Use right multiplication for ScaledBasis (3 * x instead of x * 3) + x_times_3 = 3 * x + new_stride = ( + (tma_tensor_sfb.stride[0][0], ((x, x), x_times_3)), + tma_tensor_sfb.stride[1], + tma_tensor_sfb.stride[2], + ) + tma_tensor_sfb_new_layout = cute.make_layout(new_shape, stride=new_stride) + tma_tensor_sfb = cute.make_tensor( + tma_tensor_sfb.iterator, tma_tensor_sfb_new_layout + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) + sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + self.num_tma_load_bytes = ( + a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size + ) * atom_thr_size + + # Setup TMA store for C + epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + c_tensor, + epi_smem_layout, + self.epi_tile, + ) + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + c_tensor, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + max_active_clusters, + ) + + self.buffer_align_bytes = 1024 + + # Define shared storage for kernel + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] + ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] + acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, + cute.cosize(self.c_smem_layout_staged.outer), + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + sA: cute.struct.Align[ + cute.struct.MemRange[ + self.smem_alloc_a_dtype, + cute.cosize(self.a_smem_layout_staged.outer), + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + sB: cute.struct.Align[ + cute.struct.MemRange[ + self.smem_alloc_b_dtype, + cute.cosize(self.b_smem_layout_staged.outer), + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + sSFA: cute.struct.Align[ + cute.struct.MemRange[ + self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + sSFB: cute.struct.Align[ + cute.struct.MemRange[ + self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + # Launch the kernel synchronously + self.kernel( + tiled_mma, + tiled_mma_sfb, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_sfa, + tma_tensor_sfa, + tma_atom_sfb, + tma_tensor_sfb, + tma_atom_c, + tma_tensor_c, + self.cluster_layout_vmnk, + self.cluster_layout_sfb_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + epilogue_op, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + min_blocks_per_mp=1, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_sfb: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_sfa: cute.CopyAtom, + mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + cluster_layout_sfb_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + c_smem_layout_staged: cute.Layout | cute.ComposedLayout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + epilogue_op: cutlass.Constexpr, + ): + """GPU device kernel performing the Persistent batched GEMM computation.""" + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # + # Prefetch tma desc + # + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_sfa) + cpasync.prefetch_descriptor(tma_atom_sfb) + cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + # Coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + # + # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier + # + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # Initialize mainloop ab_pipeline (barrier) and states + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = ( + self.threads_per_warp + * len(self.epilog_warp_id) + * (2 if use_2cta_instrs else 1) + ) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=self.tmem_alloc_barrier, + allocator_warp_id=self.epilog_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # + # Setup smem tensor A/B/SFA/SFB/C + # + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC = storage.sC.get_tensor( + c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner + ) + # (MMA, MMA_M, MMA_K, STAGE) + sA = storage.sA.get_tensor( + a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner + ) + # (MMA, MMA_N, MMA_K, STAGE) + sB = storage.sB.get_tensor( + b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner + ) + # (MMA, MMA_M, MMA_K, STAGE) + sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged) + # (MMA, MMA_N, MMA_K, STAGE) + sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged) + + # + # Compute multicast mask for A/B/SFA/SFB buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + sfa_full_mcast_mask = None + sfb_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + sfa_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + sfb_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1 + ) + + # + # Local_tile partition global tensors + # + # (bM, bK, RestM, RestK, RestL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, RestN, RestK, RestL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bK, RestM, RestK, RestL) + gSFA_mkl = cute.local_tile( + mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, RestN, RestK, RestL) + gSFB_nkl = cute.local_tile( + mSFB_nkl, + cute.slice_(self.mma_tiler_sfb, (0, None, None)), + (None, None, None), + ) + # (bM, bN, RestM, RestN, RestL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # + # Partition global tensor for TiledMMA_A/B/C + # + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) + tCgC = thr_mma.partition_C(gC_mnl) + + # + # Partition global/shared tensor for TMA load A/B + # + # TMA load A partition_S/D + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA load B partition_S/D + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # TMA load SFA partition_S/D + sfa_cta_layout = a_cta_layout + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsSFA, tAgSFA = cute.nvgpu.cpasync.tma_partition( + tma_atom_sfa, + block_in_cluster_coord_vmnk[2], + sfa_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + # TMA load SFB partition_S/D + sfb_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + tBsSFB, tBgSFB = cute.nvgpu.cpasync.tma_partition( + tma_atom_sfb, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + if cutlass.const_expr(self.overlapping_accum): + num_acc_stage_overlapped = 2 + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, num_acc_stage_overlapped) + ) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = cute.make_tensor( + tCtAcc_fake.iterator, + cute.make_layout( + tCtAcc_fake.shape, + stride=( + tCtAcc_fake.stride[0], + tCtAcc_fake.stride[1], + tCtAcc_fake.stride[2], + (256 - self.num_sf_tmem_cols) * tCtAcc_fake.stride[0][1], + ), + ), + ) + else: + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # + # Cluster wait before tensor memory alloc + # + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # + # Specialized TMA load warp + # + if warp_idx == self.tma_warp_id: + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + ab_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), RestK) + tAgA_slice = tAgA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # ((atom_v, rest_v), RestK) + tBgB_slice = tBgB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] + + # ((atom_v, rest_v), RestK) + tAgSFA_slice = tAgSFA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + + slice_n = mma_tile_coord_mnl[1] + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): + slice_n = mma_tile_coord_mnl[1] // 2 + # ((atom_v, rest_v), RestK) + tBgSFB_slice = tBgSFB[(None, slice_n, None, mma_tile_coord_mnl[2])] + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + ab_producer_state.reset_count() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) + # + # Tma load loop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Conditionally wait for AB buffer empty + ab_pipeline.producer_acquire( + ab_producer_state, peek_ab_empty_status + ) + + # TMA load A/B/SFA/SFB + cute.copy( + tma_atom_a, + tAgA_slice[(None, ab_producer_state.count)], + tAsA[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, ab_producer_state.count)], + tBsB[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=b_full_mcast_mask, + ) + cute.copy( + tma_atom_sfa, + tAgSFA_slice[(None, ab_producer_state.count)], + tAsSFA[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=sfa_full_mcast_mask, + ) + cute.copy( + tma_atom_sfb, + tBgSFB_slice[(None, ab_producer_state.count)], + tBsSFB[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=sfb_full_mcast_mask, + ) + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 + ab_producer_state.advance() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait A/B buffer empty + # + ab_pipeline.producer_tail(ab_producer_state) + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + # + # Bar sync for retrieve tensor memory ptr from shared mem + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator/SFA/SFB tensor + # + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # Make accumulator tmem tensor + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + # Make SFA tmem tensor + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols, + dtype=self.sf_dtype, + ) + # (MMA, MMA_M, MMA_K) + tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) + + # Make SFB tmem tensor + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, + dtype=self.sf_dtype, + ) + # (MMA, MMA_N, MMA_K) + tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + # + # Partition for S2T copy of SFA/SFB + # + ( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t, + tCtSFA_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA) + ( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t, + tCtSFB_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + ab_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_ab_stage + ) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # Get accumulator stage index + if cutlass.const_expr(self.overlapping_accum): + acc_stage_index = acc_producer_state.phase ^ 1 + else: + acc_stage_index = acc_producer_state.index + + # Set tensor memory buffer for current tile + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)] + + # Peek (try_wait) AB buffer full for k_tile = 0 + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait( + ab_consumer_state + ) + + # + # Wait for accumulator buffer empty + # + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + + tCtSFB_mma = tCtSFB + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): + # If this is an ODD tile, shift the TMEM start address for cta_tile_shape_n=192 case by two words (ignores first 64 columns of SFB) + offset = ( + cutlass.Int32(2) + if mma_tile_coord_mnl[1] % 2 == 1 + else cutlass.Int32(0) + ) + shifted_ptr = cute.recast_ptr( + acc_tmem_ptr + + self.num_accumulator_tmem_cols + + self.num_sfa_tmem_cols + + offset, + dtype=self.sf_dtype, + ) + tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout) + elif cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): + # Move in increments of 64 columns of SFB + offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2) + shifted_ptr = cute.recast_ptr( + acc_tmem_ptr + + self.num_accumulator_tmem_cols + + self.num_sfa_tmem_cols + + offset, + dtype=self.sf_dtype, + ) + tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout) + + # + # Reset the ACCUMULATE field for each tile + # + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + # + # Mma mainloop + # + for k_tile in range(k_tile_cnt): + if is_leader_cta: + # Conditionally wait for AB buffer full + ab_pipeline.consumer_wait( + ab_consumer_state, peek_ab_full_status + ) + + # Copy SFA/SFB from smem to tmem + s2t_stage_coord = ( + None, + None, + None, + None, + ab_consumer_state.index, + ) + tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord] + tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord] + cute.copy( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t_staged, + tCtSFA_compact_s2t, + ) + cute.copy( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t_staged, + tCtSFB_compact_s2t, + ) + + # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB. + # Issue a single block-scaled MMA over the full K with the + # scale-factor tensors passed inline as part of the A/B + # operands. + tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + tile_crd = (None, None, None, ab_consumer_state.index) + cute.gemm( + tiled_mma, + tCtAcc, + [tCrA[tile_crd], tCtSFA], + [tCrB[tile_crd], tCtSFB_mma], + tCtAcc, + ) + + # Async arrive AB buffer empty + ab_pipeline.consumer_release(ab_consumer_state) + + # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 + ab_consumer_state.advance() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt: + if is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait( + ab_consumer_state + ) + + # + # Async arrive accumulator buffer full + # + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait for accumulator buffer empty + # + acc_pipeline.producer_tail(acc_producer_state) + # + # Specialized epilogue warps + # + if warp_idx < self.mma_warp_id: + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + # + # Partition for epilogue + # + epi_tidx = tidx + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc, + ) = self.epilog_tmem_copy_and_partition( + epi_tidx, tCtAcc_base, tCgC, epi_tile, use_2cta_instrs + ) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition( + tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + ( + tma_atom_c, + bSG_sC, + bSG_gC_partitioned, + ) = self.epilog_gmem_copy_and_partition( + epi_tidx, tma_atom_c, tCgC, epi_tile, sC + ) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + # Threads/warps participating in tma store pipeline + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.epilog_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, + producer_group=c_producer_group, + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # ((ATOM_V, REST_V), EPI_M, EPI_N) + bSG_gC = bSG_gC_partitioned[ + ( + None, + None, + None, + *mma_tile_coord_mnl, + ) + ] + + # Get accumulator stage index + if cutlass.const_expr(self.overlapping_accum): + acc_stage_index = acc_consumer_state.phase + reverse_subtile = True if acc_stage_index == 0 else False + else: + acc_stage_index = acc_consumer_state.index + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_stage_index) + ] + + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt + for subtile_idx in cutlass.range(subtile_cnt): + real_subtile_idx = subtile_idx + if cutlass.const_expr(self.overlapping_accum): + if reverse_subtile: + real_subtile_idx = ( + self.cta_tile_shape_mnk[1] // self.epi_tile_n + - 1 + - subtile_idx + ) + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # + # Async arrive accumulator buffer empty ealier when overlapping_accum is enabled + # + if cutlass.const_expr(self.overlapping_accum): + if subtile_idx == self.iter_acc_early_release_in_epilogue: + # Fence for TMEM load + cute.arch.fence_view_async_tmem_load() + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Convert to C type + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) + tRS_rC.store(acc_vec) + + # + # Store C to shared memory + # + c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage + cute.copy( + tiled_copy_r2s, + tRS_rC, + tRS_sC[(None, None, None, c_buffer)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.epilog_sync_barrier.arrive_and_wait() + + # + # TMA store C to global memory + # + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, real_subtile_idx)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + + # + # Async arrive accumulator buffer empty + # + if cutlass.const_expr(not self.overlapping_accum): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(acc_tmem_ptr) + # + # Wait for C store complete + # + c_pipeline.producer_tail() + + def mainloop_s2t_copy_and_partition( + self, + sSF: cute.Tensor, + tSF: cute.Tensor, + ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """Make tiledCopy for smem to tmem load for scale factor tensor, then use it to partition smem memory (source) and tensor memory (destination). + + :param sSF: The scale factor tensor in smem + :type sSF: cute.Tensor + :param tSF: The scale factor tensor in tmem + :type tSF: cute.Tensor + + :return: A tuple containing (tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) where: + - tiled_copy_s2t: The tiled copy operation for smem to tmem load for scale factor tensor(s2t) + - tCsSF_compact_s2t: The partitioned scale factor tensor in smem + - tSF_compact_s2t: The partitioned scale factor tensor in tmem + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + # (MMA, MMA_MN, MMA_K, STAGE) + tCsSF_compact = cute.filter_zeros(sSF) + # (MMA, MMA_MN, MMA_K) + tCtSF_compact = cute.filter_zeros(tSF) + + # Make S2T CopyAtom and tiledCopy + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), + self.sf_dtype, + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t, tCsSF_compact_s2t_ + ) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + + return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t + + def epilog_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tAcc: cute.Tensor, + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + use_2cta_instrs: cutlass.Boolean | bool, + ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param tAcc: The accumulator tensor to be copied and partitioned + :type tAcc: cute.Tensor + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param use_2cta_instrs: Whether use_2cta_instrs is enabled + :type use_2cta_instrs: bool + + :return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where: + - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + - tTR_tAcc: The partitioned accumulator tensor + - tTR_rAcc: The accumulated tensor in register used to hold t2r results + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + # Make tiledCopy for tensor memory load + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + epi_tile, + use_2cta_instrs, + ) + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) + tAcc_epi = cute.flat_divide( + tAcc[((None, None), 0, 0, None)], + epi_tile, + ) + # (EPI_TILE_M, EPI_TILE_N) + tiled_copy_t2r = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)] + ) + + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + gC_mnl_epi = cute.flat_divide( + gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile + ) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) + # (T2R, T2R_M, T2R_N) + tTR_rAcc = cute.make_rmem_tensor( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype + ) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc + + def epilog_smem_copy_and_partition( + self, + tiled_copy_t2r: cute.TiledCopy, + tTR_rC: cute.Tensor, + tidx: cutlass.Int32, + sC: cute.Tensor, + ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination). + + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + :type tiled_copy_t2r: cute.TiledCopy + :param tTR_rC: The partitioned accumulator tensor + :type tTR_rC: cute.Tensor + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + :type sepi: cute.Tensor + + :return: A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where: + - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) + - tRS_rC: The partitioned tensor C (register source) + - tRS_sC: The partitioned tensor C (smem destination) + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + copy_atom_r2s = sm100_utils.get_smem_store_op( + self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r + ) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + # (R2S, R2S_M, R2S_N, PIPE_D) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sC = thr_copy_r2s.partition_D(sC) + # (R2S, R2S_M, R2S_N) + tRS_rC = tiled_copy_r2s.retile(tTR_rC) + return tiled_copy_r2s, tRS_rC, tRS_sC + + def epilog_gmem_copy_and_partition( + self, + tidx: cutlass.Int32, + atom: cute.CopyAtom | cute.TiledCopy, + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + sC: cute.Tensor, + ) -> tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: + """Make tiledCopy for global memory store, then use it to: + partition shared memory (source) and global memory (destination) for TMA store version. + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param atom: The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version + :type atom: cute.CopyAtom or cute.TiledCopy + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + + :return: A tuple containing (tma_atom_c, bSG_sC, bSG_gC) where: + - tma_atom_c: The TMA copy atom + - bSG_sC: The partitioned shared memory tensor C + - bSG_gC: The partitioned global tensor C + :rtype: Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor] + """ + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + gC_epi = cute.flat_divide( + gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile + ) + + tma_atom_c = atom + sC_for_tma_partition = cute.group_modes(sC, 0, 2) + gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2) + # ((ATOM_V, REST_V), EPI_M, EPI_N) + # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) + bSG_sC, bSG_gC = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + sC_for_tma_partition, + gC_for_tma_partition, + ) + return tma_atom_c, bSG_sC, bSG_gC + + @staticmethod + def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: tuple[int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + epi_tile: cute.Tile, + c_dtype: type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + sf_dtype: type[cutlass.Numeric], + sf_vec_size: int, + smem_capacity: int, + occupancy: int, + ) -> tuple[int, int, int]: + """Computes the number of stages for A/B/C operands based on heuristics. + + :param tiled_mma: The tiled MMA object defining the core computation. + :type tiled_mma: cute.TiledMma + :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler. + :type mma_tiler_mnk: tuple[int, int, int] + :param a_dtype: Data type of operand A. + :type a_dtype: type[cutlass.Numeric] + :param b_dtype: Data type of operand B. + :type b_dtype: type[cutlass.Numeric] + :param epi_tile: The epilogue tile shape. + :type epi_tile: cute.Tile + :param c_dtype: Data type of operand C (output). + :type c_dtype: type[cutlass.Numeric] + :param c_layout: Layout enum of operand C. + :type c_layout: utils.LayoutEnum + :param sf_dtype: Data type of Scale factor. + :type sf_dtype: type[cutlass.Numeric] + :param sf_vec_size: Scale factor vector size. + :type sf_vec_size: int + :param smem_capacity: Total available shared memory capacity in bytes. + :type smem_capacity: int + :param occupancy: Target number of CTAs per SM (occupancy). + :type occupancy: int + + :return: A tuple containing the computed number of stages for: + (ACC stages, A/B operand stages, C stages) + :rtype: tuple[int, int, int] + """ + # ACC stages + num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2 + + # Default C stages + num_c_stage = 2 + + # Calculate smem layout and size for one stage of A, B, SFA, SFB and C + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + tiled_mma, + mma_tiler_mnk, + a_dtype, + 1, # a tmp 1 stage is provided + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + tiled_mma, + mma_tiler_mnk, + b_dtype, + 1, # a tmp 1 stage is provided + ) + sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + mma_tiler_mnk, + sf_vec_size, + 1, # a tmp 1 stage is provided + ) + sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + mma_tiler_mnk, + sf_vec_size, + 1, # a tmp 1 stage is provided + ) + + c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + 1, + ) + + ab_bytes_per_stage = ( + cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) + + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) + ) + mbar_helpers_bytes = 1024 + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + c_bytes = c_bytes_per_stage * num_c_stage + + # Calculate A/B/SFA/SFB stages: + # Start with total smem per CTA (capacity / occupancy) + # Subtract reserved bytes and initial C stages bytes + # Divide remaining by bytes needed per A/B/SFA/SFB stage + num_ab_stage = ( + smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes) + ) // ab_bytes_per_stage + + # Refine epilogue stages: + # Calculate remaining smem after allocating for A/B/SFA/SFB stages and reserved bytes + # Add remaining unused smem to epilogue + num_c_stage += ( + smem_capacity + - occupancy * ab_bytes_per_stage * num_ab_stage + - occupancy * (mbar_helpers_bytes + c_bytes) + ) // (occupancy * c_bytes_per_stage) + + return num_acc_stage, num_ab_stage, num_c_stage + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + max_active_clusters: cutlass.Constexpr, + ) -> tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]]: + """Use persistent tile scheduler to compute the grid size for the output tensor C. + + :param c: The output tensor C + :type c: cute.Tensor + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + :param max_active_clusters: Maximum number of active clusters. + :type max_active_clusters: cutlass.Constexpr + + :return: A tuple containing: + - tile_sched_params: Parameters for the persistent tile scheduler. + - grid: Grid shape for kernel launch. + :rtype: Tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]] + """ + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + + return tile_sched_params, grid + + @staticmethod + def needs_unpack_tma( + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + ) -> bool: + """Decide whether TMA must use the UNPACK_U8 variant for narrow operands. + + Unpack is required when the operand element widths differ (mxf8f6f4 + mixed precision, e.g. MXFP4 x MXFP8). A and B must share a uniform + byte-per-element SMEM layout for the MMA, so the narrower (sub-byte) + operand is unpacked into 1B-per-element containers in SMEM. When both + operands share the same width (e.g. f4 x f4 or f8 x f8) TMA can use the + natural packed format. + + :param a_dtype: Element data type of the A operand + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: Element data type of the B operand + :type b_dtype: Type[cutlass.Numeric] + + :return: True if the UNPACK_U8 TMA format must be used, False otherwise + :rtype: bool + """ + return a_dtype.width != b_dtype.width + + @staticmethod + def is_valid_dtypes_and_scale_factor_vec_size( + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + sf_dtype: type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: type[cutlass.Numeric], + ) -> bool: + """Check if the dtypes and sf_vec_size are valid combinations + + Supports same-dtype MXF8 / MXF4 / NVF4 as well as mixed-precision + MXFP4 x MXFP8 and MXFP8 x MXFP4, where A and B differ in element type. + Mixed-precision inputs use MX scaling (Float8E8M0FNU, sf_vec_size 32). + + :param a_dtype: The data type of the A operand + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: The data type of the B operand + :type b_dtype: Type[cutlass.Numeric] + :param sf_dtype: The data type of the scale factor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: The vector size of the scale factor + :type sf_vec_size: int + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + + :return: True if the dtypes and sf_vec_size are valid, False otherwise + :rtype: bool + """ + supported_ab_dtypes = { + cutlass.Float4E2M1FN, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + } + + # Check A/B element types + if a_dtype not in supported_ab_dtypes or b_dtype not in supported_ab_dtypes: + return False + + # Check valid sf_dtype + if sf_dtype not in {cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN}: + return False + + # sf_vec_size rules: + # * 16 (NVF4 / MXF4 fp4-pair) is only supported for Float4E2M1FN x + # Float4E2M1FN. + # * 32 (MX scaling) is required for every other A/B combination, + # including mixed MXFP4 x MXFP8, and requires sf_dtype + # Float8E8M0FNU. + both_fp4 = a_dtype is cutlass.Float4E2M1FN and b_dtype is cutlass.Float4E2M1FN + if sf_vec_size == 16: + if not both_fp4: + return False + elif sf_vec_size == 32: + if sf_dtype is not cutlass.Float8E8M0FNU: + return False + else: + return False + + # Check valid c_dtype + if c_dtype not in { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + }: + return False + + return True + + @staticmethod + def is_valid_layouts( + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + ) -> bool: + """Check if layouts and dtypes are valid combinations + + FP4 operands can only be K-major; the check is per operand so that a + mixed MXFP4 x MXFP8 GEMM only constrains the FP4 operand (the FP8 + operand may be M/N- or K-major). + + :param a_dtype: The data type of the A operand + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: The data type of the B operand + :type b_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param a_major: The major dimension of the A tensor + :type a_major: Literal["m", "k"] + :param b_major: The major dimension of the B tensor + :type b_major: Literal["n", "k"] + :param c_major: The major dimension of the C tensor + :type c_major: Literal["m", "n"] + + :return: True if the layouts are valid, False otherwise + :rtype: bool + """ + is_valid = True + + # FP4 operands can only be k-major (checked per operand) + if a_dtype is cutlass.Float4E2M1FN and a_major != "k": + is_valid = False + if b_dtype is cutlass.Float4E2M1FN and b_major != "k": + is_valid = False + return is_valid + + @staticmethod + def is_valid_mma_tiler_and_cluster_shape( + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + ) -> bool: + """Check if the mma tiler and cluster shape are valid + + :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster + :type cluster_shape_mn: Tuple[int, int] + + :return: True if the mma tiler and cluster shape are valid, False otherwise + :rtype: bool + """ + is_valid = True + # Skip invalid mma tile shape + if mma_tiler_mn[0] not in [128, 256]: + is_valid = False + if mma_tiler_mn[1] not in [64, 128, 192, 256]: + is_valid = False + # Skip illegal cluster shape + if cluster_shape_mn[0] % (2 if mma_tiler_mn[0] == 256 else 1) != 0: + is_valid = False + # Skip invalid cluster shape + is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if ( + cluster_shape_mn[0] * cluster_shape_mn[1] > 16 + or cluster_shape_mn[0] <= 0 + or cluster_shape_mn[1] <= 0 + # Special cluster shape check for scale factor multicasts. + # Due to limited size of scale factors, we can't multicast among more than 4 CTAs. + or cluster_shape_mn[0] > 4 + or cluster_shape_mn[1] > 4 + or not is_power_of_2(cluster_shape_mn[0]) + or not is_power_of_2(cluster_shape_mn[1]) + ): + is_valid = False + return is_valid + + @staticmethod + def is_valid_tensor_alignment( + m: int, + n: int, + k: int, + l: int, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + ) -> bool: + """Check if the tensor alignment is valid + + :param m: The number of rows in the A tensor + :type m: int + :param n: The number of columns in the B tensor + :type n: int + :param k: The number of columns in the A tensor + :type k: int + :param l: The number of columns in the C tensor + :type l: int + :param a_dtype: The data type of the A operand + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: The data type of the B operand + :type b_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param a_major: The major axis of the A tensor + :type a_major: Literal["m", "k"] + :param b_major: The major axis of the B tensor + :type b_major: Literal["n", "k"] + :param c_major: The major axis of the C tensor + :type c_major: Literal["m", "n"] + + :return: True if the problem shape is valid, False otherwise + :rtype: bool + """ + is_valid = True + + def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // dtype.width + return num_major_elements % num_contiguous_elements == 0 + + if ( + not check_contigous_16B_alignment(a_dtype, a_major == "m", (m, k, l)) + or not check_contigous_16B_alignment(b_dtype, b_major == "n", (n, k, l)) + or not check_contigous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + ): + is_valid = False + return is_valid + + @staticmethod + def can_implement( + mnkl: tuple[int, int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + sf_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + sf_vec_size: int, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + ) -> bool: + """Check if the gemm can be implemented + + :param mnkl: The problem size as a tuple (M, N, K, L). + :type mnkl: Tuple[int, int, int, int] + :param a_dtype: The data type of the A operand + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: The data type of the B operand + :type b_dtype: Type[cutlass.Numeric] + :param sf_dtype: The data type of the scale factor tensor + :type sf_dtype: Type[cutlass.Numeric] + :param a_major: The major axis of the A tensor + :type a_major: Literal["m", "k"] + :param b_major: The major axis of the B tensor + :type b_major: Literal["n", "k"] + :param c_major: The major axis of the C tensor + :type c_major: Literal["m", "n"] + :param sf_vec_size: The vector size + :type sf_vec_size: int + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster + :type cluster_shape_mn: Tuple[int, int] + :return: True if the gemm can be implemented, False otherwise + :rtype: bool + """ + # Unpack parameters + m, n, k, l = mnkl + kernel = PersistentDenseBlockScaledGemmKernel + can_implement = True + # Skip unsupported types + if not kernel.is_valid_dtypes_and_scale_factor_vec_size( + a_dtype, b_dtype, sf_dtype, sf_vec_size, c_dtype + ): + can_implement = False + # Skip unsupported layouts + if not kernel.is_valid_layouts( + a_dtype, b_dtype, c_dtype, a_major, b_major, c_major + ): + can_implement = False + # Skip invalid mma tile shape and cluster shape + if not kernel.is_valid_mma_tiler_and_cluster_shape( + mma_tiler_mn, cluster_shape_mn + ): + can_implement = False + # Skip illegal problem shape for load/store alignment + if not kernel.is_valid_tensor_alignment( + m, n, k, l, a_dtype, b_dtype, c_dtype, a_major, b_major, c_major + ): + can_implement = False + return can_implement diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_mixed_input_impl.py b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_mixed_input_impl.py new file mode 100644 index 000000000..796dd4db8 --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_mixed_input_impl.py @@ -0,0 +1,2058 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""Mixed-Input GEMM kernel implementation for NVIDIA Blackwell SM100 architecture. + +This module implements a warp-specialized persistent GEMM kernel that supports +mixed-precision inputs where tensor A uses narrow precision (Int4/Int8/Uint8) +and tensor B uses wide precision (BFloat16/Float16). + +Warp Specialization Architecture (12 warps, 384 threads per CTA): +- Warps 0-3 (Epilogue): Load accumulators from TMEM, convert, store to GMEM +- Warp 4 (MMA): Execute matrix multiply-accumulate operations +- Warp 5 (TMA Load A/B): Load A and B tensors from GMEM to SMEM via TMA +- Warp 6 (TMA Load Scale): Load scale tensor from GMEM to SMEM (convert-scale mode) +- Warp 7 (Idle): Reserved +- Warps 8-11 (Transform): Type conversion and optional scaling of A tensor + +Transform Modes: +- Convert-Only: C = type_convert(A) @ B (for Int8/Uint8) +- Convert-Scale: C = (type_convert(A) * scale) @ B (for Int4) + +Pipeline Architecture: +- a_load2trans_pipeline: TMA load A → Transform warps +- scale_load2trans_pipeline: TMA load scale → Transform warps (convert-scale mode) +- trans2mma_pipeline: Transform warps → MMA warp +- b_load2mma_pipeline: TMA load B → MMA warp +- acc_pipeline: MMA warp → Epilogue warps +""" + +from dataclasses import dataclass +from math import ceil, log2 + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.mixed_input_helpers as mixed_input_utils +from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05 +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +from cutlass.utils.mixed_input_helpers import TransformMode + +from cutlass.operators.utils.common import round_up + +# ============================================================================= +# Capacity computation for mixed-input GEMM kernels +# ============================================================================= + +# Constants for capacity calculations +SM100_SMEM_ALIGN = 1024 +_BYTES_PER_PIPELINE_STAGE = 16 + + +@dataclass +class _TmemConfig: + """TMEM allocation configuration.""" + + num_tmem_acc_col_per_stage: int + num_tmem_cols_a_per_stage: int + accumulator_stage_count: int + num_tmem_acc_cols: int + + +@dataclass +class _SmemConfig: + """SMEM allocation configuration: per-stage byte counts, stage counts, and totals.""" + + a_load_bytes_per_stage: int + b_load_bytes_per_stage: int + ab_load_bytes_per_stage: int + a_transform_bytes_per_stage: int + a_scale_bytes_per_stage: int + scale_load2trans_stage_count: int + c_bytes_per_stage: int + initial_c_stage_count: int + carveout_smem_bytes: int + + +@dataclass +class CapacityConfig: + """Result of capacity computation for mixed-input GEMM kernels. + + This dataclass contains all computed stage counts and TMEM allocations + for a given kernel configuration. It serves as the single source of truth + for capacity-related computations, used by both early validation checks + and compile-time kernel configuration. + """ + + # Stage counts + load2transform_stage_count: int + transform2mma_stage_count: int + accumulator_stage_count: int + scale_load2trans_stage_count: int + c_stage_count: int + + # TMEM allocations + num_tmem_acc_cols: int + num_tmem_a_cols: int + num_tmem_alloc_cols: int + num_tmem_acc_col_per_stage: int + num_tmem_cols_a_per_stage: int + + @property + def is_valid(self) -> bool: + """Check if the configuration has sufficient capacity.""" + return ( + self.load2transform_stage_count >= 2 + and self.transform2mma_stage_count >= 2 + and self.accumulator_stage_count >= 1 + ) + + +def _compute_tmem_config( + cta_tile_m: int, + cta_tile_n: int, + cta_tile_k: int, + use_2cta: bool, + mma_dtype_width: int, + transform_a_source_is_tmem: bool, + tmem_capacity: int, +) -> _TmemConfig: + """Compute TMEM allocation configuration. + + Determines accumulator stage count and TMEM column allocations based on + tile sizes and whether A uses TMEM. + """ + from cutlass.utils.blackwell_helpers import compute_acc_tmem_cols_per_stage + + # Compute accumulator TMEM columns per stage + num_tmem_acc_col_per_stage = compute_acc_tmem_cols_per_stage( + cta_tile_m, + cta_tile_n, + use_2cta, + mma_n=cta_tile_n, + transform_a_source_is_tmem=transform_a_source_is_tmem, + ) + + # Dynamic accumulator stage count heuristic + # When A uses TMEM, reduce accumulator stages to leave room for A + if transform_a_source_is_tmem: + if num_tmem_acc_col_per_stage < 128: + accumulator_stage_count = 3 + elif num_tmem_acc_col_per_stage < 256: + accumulator_stage_count = 2 + else: + accumulator_stage_count = 1 + else: + accumulator_stage_count = tmem_capacity // num_tmem_acc_col_per_stage + + num_elts_per_tmem_col = 32 // mma_dtype_width + num_tmem_cols_a_per_stage = ( + round_up(cta_tile_k // num_elts_per_tmem_col, 4) + if transform_a_source_is_tmem + else 0 + ) + + num_tmem_acc_cols = round_up( + accumulator_stage_count * num_tmem_acc_col_per_stage, 4 + ) + + return _TmemConfig( + num_tmem_acc_col_per_stage=num_tmem_acc_col_per_stage, + num_tmem_cols_a_per_stage=num_tmem_cols_a_per_stage, + accumulator_stage_count=accumulator_stage_count, + num_tmem_acc_cols=num_tmem_acc_cols, + ) + + +def _compute_smem_bytes_config( + cta_tile_m: int, + cta_tile_n: int, + cta_tile_k: int, + use_2cta: bool, + a_dtype_width: int, + b_dtype_width: int, + mma_dtype_width: int, + out_dtype_width: int, + transform_a_source_is_tmem: bool, + scale_mode_is_convert_only: bool, + scale_granularity_m: int, + scale_granularity_k: int, + use_tma_store: bool, + accumulator_stage_count: int, + smem_align: int, +) -> _SmemConfig: + """Compute SMEM bytes per stage for each buffer type. + + Calculates memory requirements for A/B loads, transform buffers, + scale tensors, and epilogue buffers. + """ + from cutlass.utils.blackwell_helpers import compute_epilogue_tile_size + + cta_divisor = 2 if use_2cta else 1 + + # A load bytes per stage + a_load_bytes_per_stage = round_up( + cta_tile_m * cta_tile_k * a_dtype_width // 8, smem_align + ) + + # B load bytes per stage (N divided by cta_divisor for 2CTA sharing) + b_n_for_smem = cta_tile_n // cta_divisor + b_load_bytes_per_stage = round_up( + b_n_for_smem * cta_tile_k * b_dtype_width // 8, smem_align + ) + + ab_load_bytes_per_stage = ( + a_load_bytes_per_stage + b_load_bytes_per_stage + 2 * _BYTES_PER_PIPELINE_STAGE + ) + + # Transformed A SMEM per stage (only when A source is SMEM) + a_transform_bytes_per_stage = ( + round_up(cta_tile_m * cta_tile_k * mma_dtype_width // 8, smem_align) + if not transform_a_source_is_tmem + else 0 + ) + _BYTES_PER_PIPELINE_STAGE + + # Scale bytes (for scaled mode) + if scale_mode_is_convert_only: + scale_load2trans_stage_count = 0 + a_scale_bytes_per_stage = 0 + else: + a_scale_k_mode = ( + max(cta_tile_k // scale_granularity_k, 1) if scale_granularity_k > 0 else 1 + ) + a_scale_m_mode = ( + max(cta_tile_m // scale_granularity_m, 1) if scale_granularity_m > 0 else 1 + ) + scale_load2trans_stage_count = 4 + a_scale_bytes_per_stage = round_up( + a_scale_m_mode * a_scale_k_mode * mma_dtype_width // 8, smem_align + ) + + a_scale_bytes = ( + a_scale_bytes_per_stage + _BYTES_PER_PIPELINE_STAGE + ) * scale_load2trans_stage_count + + # Epilogue bytes (c_bytes) + c_stage_count = 2 if use_tma_store else 0 + epi_tile_m, epi_tile_n = compute_epilogue_tile_size( + cta_tile_m, cta_tile_n, use_2cta, out_dtype_width + ) + c_bytes_per_stage = round_up( + epi_tile_m * epi_tile_n * out_dtype_width // 8, smem_align + ) + c_bytes = c_bytes_per_stage * c_stage_count + + # Carveout SMEM bytes (reserved for accumulator pipeline, scale, epilogue) + carveout_smem_bytes = ( + _BYTES_PER_PIPELINE_STAGE * accumulator_stage_count + a_scale_bytes + c_bytes + ) + + return _SmemConfig( + a_load_bytes_per_stage=a_load_bytes_per_stage, + b_load_bytes_per_stage=b_load_bytes_per_stage, + ab_load_bytes_per_stage=ab_load_bytes_per_stage, + a_transform_bytes_per_stage=a_transform_bytes_per_stage, + a_scale_bytes_per_stage=a_scale_bytes_per_stage, + scale_load2trans_stage_count=scale_load2trans_stage_count, + c_bytes_per_stage=c_bytes_per_stage, + initial_c_stage_count=c_stage_count, + carveout_smem_bytes=carveout_smem_bytes, + ) + + +def _compute_stage_counts( + tmem_config: _TmemConfig, + smem_config: _SmemConfig, + transform_a_source_is_tmem: bool, + use_tma_store: bool, + smem_capacity: int, + tmem_capacity: int, +) -> tuple[int, int, int]: + """Compute final stage counts based on TMEM and SMEM capacity limits. + + Returns (load2transform_stage_count, transform2mma_stage_count, c_stage_count). + """ + # TMEM-limited stage count for transform2mma + if transform_a_source_is_tmem and tmem_config.num_tmem_cols_a_per_stage > 0: + transform2mma_stage_count_tmem = ( + tmem_capacity - tmem_config.num_tmem_acc_cols + ) // tmem_config.num_tmem_cols_a_per_stage + else: + transform2mma_stage_count_tmem = float("inf") + + # SMEM-limited stage count for transform2mma + transform2mma_stage_count_smem = ( + smem_capacity - smem_config.carveout_smem_bytes + ) // (smem_config.ab_load_bytes_per_stage + smem_config.a_transform_bytes_per_stage) + + transform2mma_stage_count = int( + min(transform2mma_stage_count_tmem, transform2mma_stage_count_smem) + ) + + # Load2transform stage count + load2transform_stage_count = ( + smem_capacity + - smem_config.carveout_smem_bytes + - (transform2mma_stage_count * smem_config.a_transform_bytes_per_stage) + ) // smem_config.ab_load_bytes_per_stage + + # Increase c_stage_count with leftover SMEM if using TMA store + c_stage_count = smem_config.initial_c_stage_count + if use_tma_store and smem_config.c_bytes_per_stage > 0: + c_stage_count += ( + smem_capacity + - load2transform_stage_count * smem_config.ab_load_bytes_per_stage + - transform2mma_stage_count * smem_config.a_transform_bytes_per_stage + - smem_config.carveout_smem_bytes + ) // smem_config.c_bytes_per_stage + + return load2transform_stage_count, transform2mma_stage_count, c_stage_count + + +def compute_capacity_config( + cta_tile_m: int, + cta_tile_n: int, + cta_tile_k: int, + use_2cta: bool, + a_dtype_width: int, + b_dtype_width: int, + mma_dtype_width: int, + out_dtype_width: int, + transform_a_source_is_tmem: bool, + smem_align: int, + scale_mode_is_convert_only: bool = True, + scale_granularity_m: int = 0, + scale_granularity_k: int = 0, + use_tma_store: bool = True, + smem_capacity: int | None = None, + tmem_capacity: int | None = None, +) -> CapacityConfig: + """Compute capacity configuration for mixed-input GEMM kernels. + + This function computes all stage counts, TMEM allocations, and SMEM requirements + for a given kernel configuration. It is the single source of truth for capacity + computations, used by both early validation and compile-time configuration. + """ + from cutlass.cute.arch import get_max_tmem_alloc_cols + from cutlass.utils import get_smem_capacity_in_bytes + + # Get hardware capacities + if smem_capacity is None: + smem_capacity = get_smem_capacity_in_bytes("sm_100") + if tmem_capacity is None: + tmem_capacity = get_max_tmem_alloc_cols("sm_100") + + # Step 1: Compute TMEM allocations + tmem_config = _compute_tmem_config( + cta_tile_m, + cta_tile_n, + cta_tile_k, + use_2cta, + mma_dtype_width, + transform_a_source_is_tmem, + tmem_capacity, + ) + + # Step 2: Compute SMEM bytes per stage + smem_config = _compute_smem_bytes_config( + cta_tile_m, + cta_tile_n, + cta_tile_k, + use_2cta, + a_dtype_width, + b_dtype_width, + mma_dtype_width, + out_dtype_width, + transform_a_source_is_tmem, + scale_mode_is_convert_only, + scale_granularity_m, + scale_granularity_k, + use_tma_store, + tmem_config.accumulator_stage_count, + smem_align, + ) + + # Step 3: Compute final stage counts + ( + load2transform_stage_count, + transform2mma_stage_count, + c_stage_count, + ) = _compute_stage_counts( + tmem_config, + smem_config, + transform_a_source_is_tmem, + use_tma_store, + smem_capacity, + tmem_capacity, + ) + + # Compute TMEM A columns from stage count + num_tmem_a_cols = transform2mma_stage_count * tmem_config.num_tmem_cols_a_per_stage + + # Apply hardware allocation constraints: min 32 columns, power-of-2 total + from cutlass.cute.arch import get_min_tmem_alloc_cols + + num_tmem_alloc_cols = max( + get_min_tmem_alloc_cols("sm_100"), + 2 ** ceil(log2(tmem_config.num_tmem_acc_cols + num_tmem_a_cols)) + if (tmem_config.num_tmem_acc_cols + num_tmem_a_cols) > 0 + else 0, + ) + + return CapacityConfig( + load2transform_stage_count=load2transform_stage_count, + transform2mma_stage_count=transform2mma_stage_count, + accumulator_stage_count=tmem_config.accumulator_stage_count, + scale_load2trans_stage_count=smem_config.scale_load2trans_stage_count, + c_stage_count=c_stage_count, + num_tmem_acc_cols=tmem_config.num_tmem_acc_cols, + num_tmem_a_cols=num_tmem_a_cols, + num_tmem_alloc_cols=num_tmem_alloc_cols, + num_tmem_acc_col_per_stage=tmem_config.num_tmem_acc_col_per_stage, + num_tmem_cols_a_per_stage=tmem_config.num_tmem_cols_a_per_stage, + ) + + +class MixedInputGemmKernel: + """Mixed-Input GEMM kernel implementation for NVIDIA Blackwell SM100 architecture. + + This kernel supports GEMM operations where input tensors A and B have different + data types, with tensor A being transformed to the precision of tensor B before + matrix multiplication. + + :param scale_granularity_m: Number of elements sharing the same scale factor along M + :param scale_granularity_k: Number of elements sharing the same scale factor along K + :param acc_dtype: Data type for accumulation during computation + :param use_2cta_instrs: Whether to use CTA group 2 for advanced thread cooperation + :param mma_tiler_mn: Shape of the MMA tile (M, N) + :param cluster_shape_mn: Cluster dimensions (M, N) for parallel processing + :param use_tma_store: Whether to use TMA for storing results + :param shuffle_a: Whether to use shuffle intrinsic for int4-to-bf16 conversion + """ + + def __init__( + self, + scale_granularity_m: int, + scale_granularity_k: int, + acc_dtype: type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_tma_store: bool, + shuffle_a: bool, + ): + """Initializes the mixed-input GEMM kernel with specified configuration. + """ + self.scale_granularity_m = scale_granularity_m + self.scale_granularity_k = scale_granularity_k + + # Determine transform mode based on scale granularity + if scale_granularity_m == 0 and scale_granularity_k == 0: + self.scale_mode = TransformMode.ConvertOnly + else: + self.scale_mode = TransformMode.ConvertScale + + self.acc_dtype = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + self.mma_tiler = mma_tiler_mnk # Full 3-tuple (M, N, K) from user + self.use_tma_store = use_tma_store + self.shuffle_a = shuffle_a + self.arch = "sm_100" + + self.cta_group = ( + tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + # Set specialized warp ids + self.epilog_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.scale_tma_warp_id = 6 + self.idle_warp_id = 7 + self.transform_warp_id = (8, 9, 10, 11) + + # Register budgets for different warp types + self.num_regs_epilogue_warps = 192 + self.num_regs_mma_warp = 96 + self.num_regs_tma_warps = 96 + self.num_regs_transform_warps = 208 + self.num_regs_idle_warp = 24 + + self.threads_per_cta = 32 * ( + max( + ( + self.mma_warp_id, + self.tma_warp_id, + self.scale_tma_warp_id, + *self.epilog_warp_id, + *self.transform_warp_id, + ) + ) + + 1 + ) + + # Set barrier ids for synchronization + self.epilog_sync_barrier = pipeline.NamedBarrier( + 1, 32 * len(self.epilog_warp_id) + ) + self.tmem_ptr_sync_barrier = pipeline.NamedBarrier(2, self.threads_per_cta) + self.transform_sync_barrier = pipeline.NamedBarrier( + 3, 32 * len(self.transform_warp_id) + ) + + self.smem_buffer_align_bytes = SM100_SMEM_ALIGN + self.smem_capacity = utils.get_smem_capacity_in_bytes(self.arch) + + def _setup_attributes(self): + """Set up configurations that are dependent on GEMM inputs. + + This method configures various attributes based on the input tensor properties + (data types, leading dimensions) and kernel settings: + - Deduce where the transformed A tensor is stored (SMEM or TMEM) + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B + - Computing epilogue sub-tile + - Setting up stage counts for pipelines + - Computing memory layouts + - Computing tensor memory allocation columns + """ + # Deduce where the transformed A tensor is stored + self.transform_a_source = mixed_input_utils.get_transform_a_source( + self.a_major_mode + ) + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.mma_dtype, + self.mma_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + self.transform_a_source, + ) + + # Use user-provided mma_tiler directly (K dimension is already set) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + if cutlass.const_expr(self.use_tma_store): + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + else: + self.epi_tile = self.cta_tile_shape_mnk[:2] + + # Compute stages and TMEM columns + ( + self.num_load2trans_stage, + self.num_scale_load2trans_stage, + self.num_trans2mma_stage, + self.num_acc_stage, + self.num_c_stage, + self.num_acc_tmem_cols, + self.num_a_tmem_cols, + self.num_tmem_alloc_cols, + ) = self._compute_stages_and_tmem_cols( + tiled_mma, + self.cta_tile_shape_mnk, + self.a_dtype, + self.b_dtype, + self.c_dtype, + self.transform_a_source, + self.scale_granularity_m, + self.scale_granularity_k, + self.smem_buffer_align_bytes, + self.use_tma_store, + self.scale_mode, + ) + + # Get smem layout for C tensor when TMA store is enabled + self.c_smem_layout_staged = ( + sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile, + self.num_c_stage, + ) + if self.use_tma_store + else None + ) + + # Get smem layout for A, transformed A, and B + ( + self.smem_layout_a, + self.smem_layout_a_transform, + self.smem_layout_b, + ) = mixed_input_utils.compute_smem_layout( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.num_load2trans_stage, + self.num_trans2mma_stage, + ) + + # Get smem layout for scale tensor + self.smem_layout_scale_per_stage = None + self.smem_layout_scale = None + if self.scale_mode == TransformMode.ConvertScale: + ( + self.scale_tile_shape, + self.smem_layout_scale_per_stage, + self.smem_layout_scale, + ) = mixed_input_utils.get_smem_layout_scale( + self.mma_tiler, + self.use_2cta_instrs, + self.scale_granularity_m, + self.scale_granularity_k, + self.scale_major_mode, + self.a_scale_dtype, + self.num_scale_load2trans_stage, + ) + + @staticmethod + def _compute_stages_and_tmem_cols( + tiled_mma, + cta_tile_shape_mnk, + a_dtype, + b_dtype, + c_dtype, + transform_a_source, + scale_granularity_m, + scale_granularity_k, + smem_align, + use_tma_store, + scale_mode, + ): + """Compute pipeline stages and TMEM column allocation configurations. + + This method uses the shared compute_capacity_config function to compute + all stage counts and TMEM allocations, ensuring consistency with the + early validation check in _check_capacity. + """ + # Determine if A uses TMEM or SMEM + transform_a_source_is_tmem = transform_a_source == tcgen05.OperandSource.TMEM + + # Use the shared capacity computation function + config = compute_capacity_config( + cta_tile_m=cta_tile_shape_mnk[0], + cta_tile_n=cta_tile_shape_mnk[1], + cta_tile_k=cta_tile_shape_mnk[2], + use_2cta=(cute.size(tiled_mma.thr_id) == 2), + a_dtype_width=a_dtype.width, + b_dtype_width=b_dtype.width, + mma_dtype_width=tiled_mma.op.a_dtype.width, + out_dtype_width=c_dtype.width, + transform_a_source_is_tmem=transform_a_source_is_tmem, + scale_mode_is_convert_only=(scale_mode == TransformMode.ConvertOnly), + scale_granularity_m=scale_granularity_m, + scale_granularity_k=scale_granularity_k, + use_tma_store=use_tma_store, + smem_align=smem_align, + ) + + # Check capacity using the shared validation + if not config.is_valid: + raise ValueError("Not enough SMEM or TMEM capacity for selected tile size") + + return ( + config.load2transform_stage_count, + config.scale_load2trans_stage_count, + config.transform2mma_stage_count, + config.accumulator_stage_count, + config.c_stage_count, + config.num_tmem_acc_cols, + config.num_tmem_a_cols, + config.num_tmem_alloc_cols, + ) + + def _validate_inputs( + self, + a: cute.Tensor, + a_scale: cute.Tensor | None, + b: cute.Tensor, + c: cute.Tensor, + ) -> None: + """Validates input tensors and their properties.""" + # Validate scale tensor major mode for convert-scale mode + if self.scale_mode == TransformMode.ConvertScale: + if a_scale is None: + raise ValueError("Scale tensor required for convert-scale mode") + if ( + utils.LayoutEnum.from_tensor(a_scale).mma_major_mode() + != OperandMajorMode.MN + ): + raise ValueError("scale_major_mode must be M-major") + + @cute.jit + def __call__( + self, + a: cute.Tensor, + a_scale: cute.Tensor | None, + b: cute.Tensor, + c: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + """Executes the Mixed Input GEMM operation. + + This method sets up the kernel parameters, computes the grid size, + defines the shared storage, and launches the kernel. + """ + + # Add batch dimension if tensor is 2D (up_to_rank=3) + def add_batch_mode(tensor: cute.Tensor) -> cute.Tensor: + return cute.make_tensor( + tensor.iterator, + cute.prepend(tensor.layout, cute.make_layout(1), up_to_rank=3), + ) + + a = add_batch_mode(a) + b = add_batch_mode(b) + c = add_batch_mode(c) + + # Permute tensor modes from PyTorch to CuTe convention + # A: (L, M, K) -> (M, K, L) + a = cute.make_tensor(a.iterator, cute.select(a.layout, [1, 2, 0])) + # B: (L, K, N) -> (N, K, L) + b = cute.make_tensor(b.iterator, cute.select(b.layout, [2, 1, 0])) + # C: (L, M, N) -> (M, N, L) + c = cute.make_tensor(c.iterator, cute.select(c.layout, [1, 2, 0])) + + # Handle scale tensor for ConvertScale mode + # Use compile-time check since a_scale being None is determined at compile time + if cutlass.const_expr(self.scale_mode is TransformMode.ConvertScale): + # Scale: (L, M, K//scale_k) -> (M, K//scale_k, L) + a_scale = add_batch_mode(a_scale) + a_scale = cute.make_tensor( + a_scale.iterator, cute.select(a_scale.layout, [1, 2, 0]) + ) + + # Extract data types from tensors + self.a_dtype: type[cutlass.Numeric] = a.element_type + self.a_scale_dtype: type[cutlass.Numeric] = ( + a_scale.element_type + if self.scale_mode is TransformMode.ConvertScale + else None + ) + self.b_dtype: type[cutlass.Numeric] = b.element_type + self.c_dtype: type[cutlass.Numeric] = c.element_type + self.mma_dtype = self.b_dtype + + # Extract layout information + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.scale_major_mode = ( + utils.LayoutEnum.from_tensor(a_scale).mma_major_mode() + if self.scale_mode is TransformMode.ConvertScale + else None + ) + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + self.gmem_layout_scale = mixed_input_utils.get_gmem_layout_scale( + a.shape, + self.scale_granularity_m, + self.scale_granularity_k, + self.scale_major_mode, + ) + + # Validate inputs + self._validate_inputs(a, a_scale, b, c) + + # Setup attributes dependent on GEMM inputs + self._setup_attributes() + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.mma_dtype, + self.mma_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + self.transform_a_source, + ) + + # Set up TMA copy atoms + a_op = mixed_input_utils.get_tma_atom_kind( + self.is_a_mcast, self.use_2cta_instrs, False + ) + b_op = mixed_input_utils.get_tma_atom_kind( + self.is_b_mcast, self.use_2cta_instrs, True + ) + a_scale_op = a_op + + # Create TMA atoms for A, B, and scale + smem_layout_a_per_stage = cute.slice_(self.smem_layout_a, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + smem_layout_a_per_stage, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + + tma_atom_scale, tma_tensor_scale = None, None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + smem_layout_for_tma_atom = cute.get( + tiled_mma._thrfrg_A(self.smem_layout_scale_per_stage.outer), mode=[1] + ) + smem_layout_for_tma_atom = cute.dice( + smem_layout_for_tma_atom, + (1, (1,) * cute.rank(self.smem_layout_scale_per_stage.outer)), + ) + tma_atom_scale, tma_tensor_scale = cute.nvgpu.make_tiled_tma_atom_A( + a_scale_op, + cute.make_tensor(a_scale.iterator, self.gmem_layout_scale), + smem_layout_for_tma_atom, + (self.scale_tile_shape[0], 1, self.scale_tile_shape[1]), + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 + if a_scale.element_type is cutlass.Float32 + else None + ), + ) + + smem_layout_b_per_stage = cute.slice_(self.smem_layout_b, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + smem_layout_b_per_stage, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + + # Calculate copy sizes for TMA + a_copy_size = cute.size_in_bytes(self.a_dtype, smem_layout_a_per_stage) + b_copy_size = cute.size_in_bytes(self.b_dtype, smem_layout_b_per_stage) + a_scale_copy_size = ( + cute.size_in_bytes(self.a_scale_dtype, self.smem_layout_scale_per_stage) + if self.scale_mode is TransformMode.ConvertScale + else 0 + ) + + self.num_tma_load_bytes_a = a_copy_size + self.num_tma_load_bytes_b = b_copy_size * cute.size(tiled_mma.thr_id.shape) + self.num_tma_load_bytes_scale = a_scale_copy_size + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + c, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + max_active_clusters, + ) + + # Setup TMA atom for C (if using TMA store) + tma_atom_c = None + tma_tensor_c = None + if cutlass.const_expr(self.use_tma_store): + epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + c, + epi_smem_layout, + self.epi_tile, + ) + + # Define shared storage structure + @cute.struct + class SharedStorage: + a_load2trans_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + a_load2trans_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + a_scale_load2trans_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_scale_load2trans_stage + ] + a_scale_load2trans_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_scale_load2trans_stage + ] + a_trans2mma_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_trans2mma_stage + ] + a_trans2mma_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_trans2mma_stage + ] + b_load2mma_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + b_load2mma_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + smem_C: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, + cute.size(self.c_smem_layout_staged.outer) + if self.use_tma_store + else 0, + ], + self.smem_buffer_align_bytes, + ] + smem_A: cute.struct.Align[ + cute.struct.MemRange[self.a_dtype, cute.size(self.smem_layout_a.outer)], + self.smem_buffer_align_bytes, + ] + smem_B: cute.struct.Align[ + cute.struct.MemRange[self.b_dtype, cute.size(self.smem_layout_b.outer)], + self.smem_buffer_align_bytes, + ] + smem_A_transform: cute.struct.Align[ + cute.struct.MemRange[ + self.mma_dtype, + cute.size(self.smem_layout_a_transform.outer) + if self.transform_a_source == tcgen05.OperandSource.SMEM + else 0, + ], + self.smem_buffer_align_bytes, + ] + smem_A_scale: cute.struct.Align[ + cute.struct.MemRange[ + self.mma_dtype, + cute.size(self.smem_layout_scale.outer) + if self.scale_mode is TransformMode.ConvertScale + else 0, + ], + self.smem_buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + # Launch kernel + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_scale, + tma_tensor_scale, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c if self.use_tma_store else c, + self.cluster_layout_vmnk, + self.smem_layout_a, + self.smem_layout_scale, + self.smem_layout_a_transform, + self.smem_layout_b, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + min_blocks_per_mp=1, + stream=stream, + ) + return + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + max_active_clusters: cutlass.Constexpr, + ): + """Use persistent tile scheduler to compute the grid size for the output tensor C.""" + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + + return tile_sched_params, grid + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_s: cute.CopyAtom | None, + mS_mkl: cute.Tensor | None, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom | None, + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + a_smem_layout: cute.ComposedLayout, + scale_smem_layout: cute.ComposedLayout, + a_smem_layout_transform: cute.ComposedLayout, + b_smem_layout: cute.ComposedLayout, + c_smem_layout_staged: cute.ComposedLayout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + ): + """GPU device kernel performing the Persistent Mixed-Input GEMM computation. + """ + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, bidz = cute.arch.block_idx() + # Prefetch TMA descriptors + if warp_idx == self.epilog_warp_id[0]: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + cpasync.prefetch_descriptor(tma_atom_s) + if cutlass.const_expr(self.use_tma_store): + cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + bidx, bidy, bidz = cute.arch.block_idx() + # Compute how many k_tiles share the same scale + num_k_tiles_per_scale = self.scale_granularity_k // self.cta_tile_shape_mnk[2] + + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + tidx, _, _ = cute.arch.thread_idx() + + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # Initialize load2transform pipeline, which tracks the dependencies between TMA's loading + # of A and B, and the transformation of A and MMA's consumption + transform_thread_idx = ( + tidx - 32 * self.transform_warp_id[0] + if tidx >= 32 * self.transform_warp_id[0] + else tidx + ) + a_load2trans_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.a_load2trans_full_mbar_ptr.data_ptr(), + num_stages=self.num_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.num_mcast_ctas_a * len(self.transform_warp_id), + ), + tx_count=self.num_tma_load_bytes_a, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=transform_thread_idx, + mcast_mode_mn=(1, 0), # multicast for A will only happen on the M-mode + defer_sync=True, + ) + # Initialize scale_load2trans pipeline, which tracks the dependencies between TMA's loading + # of scale, and the transformation of A + scale_load2trans_pipeline = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + num_producers_a_scale = self.num_mcast_ctas_a + scale_load2trans_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.a_scale_load2trans_full_mbar_ptr.data_ptr(), + num_stages=self.num_scale_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + num_producers_a_scale + * len(self.transform_warp_id) + * num_k_tiles_per_scale, + ), + tx_count=self.num_tma_load_bytes_scale, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=transform_thread_idx, + mcast_mode_mn=( + 1, + 0, + ), # multicast for scale_a will only happen on the M-mode + defer_sync=True, + ) + # Initialize transform2mma pipeline, which tracks the dependencies between the transformation + # of A and MMA's consumption of transformed A + cta_v_size = cute.size(cluster_layout_vmnk, mode=[0]) + trans2mma_pipeline = pipeline.PipelineAsyncUmma.create( + barrier_storage=storage.a_trans2mma_full_mbar_ptr.data_ptr(), + num_stages=self.num_trans2mma_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.transform_warp_id) * cta_v_size, + ), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + # Initialize pipeline for tensor B load to MMA + # MMA warp informs TMA warp to proceed to load next tile of B tensor + b_load2mma_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.b_load2mma_full_mbar_ptr.data_ptr(), + num_stages=self.num_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_mcast_ctas_b + ), + tx_count=self.num_tma_load_bytes_b, + cta_layout_vmnk=cluster_layout_vmnk, + mcast_mode_mn=(0, 1), # multicast for B will only happen on the N-mode + defer_sync=True, + ) + # Initialize accumulator pipeline, which tracks the dependencies between + # MMA's computation of accumulators and epilogue warps' consumption of accumulators + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, cta_v_size * len(self.epilog_warp_id) + ), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=self.tmem_ptr_sync_barrier, + allocator_warp_id=self.epilog_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # Setup smem tensor A/scale/B/C + sC = ( + storage.smem_C.get_tensor( + c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner + ) + if self.use_tma_store + else None + ) + sA_input = storage.smem_A.get_tensor( + a_smem_layout.outer, swizzle=a_smem_layout.inner + ) + sS_input = ( + storage.smem_A_scale.get_tensor( + scale_smem_layout.outer, swizzle=scale_smem_layout.inner + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + sB_input = storage.smem_B.get_tensor( + b_smem_layout.outer, swizzle=b_smem_layout.inner + ) + sA_transform = None + # Get smem tensor for transformed A when transform_a_source is SMEM + if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.SMEM): + sA_transform = storage.smem_A_transform.get_tensor( + a_smem_layout_transform.outer, swizzle=a_smem_layout_transform.inner + ) + + # Compute multicast mask for A/B buffer full + a_full_mcast_mask = None + b_full_mcast_mask = None + s_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + # scale tensor share the same multicast mask with A tensor + s_full_mcast_mask = a_full_mcast_mask + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + # local_tile partition global tensors + # (bM, bK, loopM, loopK, loopL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bM, bK, loopM, loopK, loopL) + gS_mkl = ( + cute.local_tile( + mS_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + # (bN, bK, loopN, loopK, loopL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bN, loopM, loopN, loopL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # Partition global tensor for TiledMMA_A/B/C + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + # (MMA, MMA_M, MMA_K, loopM, loopK, loopL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_M, MMA_K, loopM, loopK, loopL) + tCgS = ( + thr_mma.partition_A(gS_mkl) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + # (MMA, MMA_N, MMA_K, loopN, loopK, loopL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_N, loopM, loopN, loopL) + tCgC = thr_mma.partition_C(gC_mnl) + + # Setup copy atom to load A from shared memory for further transformation + copy_atom_a_input = ( + cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self.a_dtype, num_bits_per_copy=32 + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + a_smem_shape = tiled_mma.partition_shape_A( + cute.dice(self.mma_tiler, (1, None, 1)) + ) + # Setup copy atom to store transformed A into tensor memory or shared memory + copy_atom_a_transform = mixed_input_utils.get_copy_atom_a_transform( + self.mma_dtype, + self.use_2cta_instrs, + self.transform_a_source, + a_smem_shape, + self.a_dtype, + ) + + # Partition global/shared tensor for TMA load A/B + # TMA load A partition_S/D + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA_input, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + + tCsS = None + tSsS = None + tSgS = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + thr_mma_leader_cta = tiled_mma.get_slice(0) + # (MMA, MMA_M, MMA_K, STAGE) + tCsS = thr_mma_leader_cta.partition_A(sS_input) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tSsS, tSgS = mixed_input_utils.scale_tma_partition( + tCsS, + tCgS, + tma_atom_s, + block_in_cluster_coord_vmnk, + a_cta_layout, + ) + + # TMA load B partition_S/D + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB_input, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB_input) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # Cluster wait before TMEM alloc and ensure pipelines are ready + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # TMEM allocation + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + # Get the pointer to the TMEM buffer + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + accumulators = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + tCrA = None + if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.TMEM): + # TMEM columns are 32-bit wide. For accumulators with element size < 32 bits + # (e.g., Float16), we need to scale the column offset to element offset. + # num_acc_tmem_cols is in column units, but pointer arithmetic uses element units. + column_to_element_scale = 32 // self.acc_dtype.width + tmem_ptr_transform = cute.recast_ptr( + accumulators.iterator + + self.num_acc_tmem_cols * column_to_element_scale, + dtype=self.mma_dtype, + ) + tCrA = cute.make_tensor( + tmem_ptr_transform, + tiled_mma.make_fragment_A(a_smem_layout_transform.outer).layout, + ) + else: + tCrA = tiled_mma.make_fragment_A(sA_transform) + + # Specialized TMA load warp for A/B tensor + if warp_idx == self.tma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_tma_warps) + # Persistent tile scheduling loop + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, (bidx, bidy, bidz), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + a_load2trans_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_load2trans_stage + ) + b_load2mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_load2trans_stage + ) + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + tAgA_slice = tAgA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + tBgB_slice = tBgB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] + + a_load2trans_producer_state.reset_count() + peek_load2trans_empty_status = cutlass.Boolean(1) + if a_load2trans_producer_state.count < k_tile_cnt: + peek_load2trans_empty_status = ( + a_load2trans_pipeline.producer_try_acquire( + a_load2trans_producer_state + ) + ) + b_load2mma_producer_state.reset_count() + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + a_load2trans_pipeline.producer_acquire( + a_load2trans_producer_state, peek_load2trans_empty_status + ) + b_load2mma_pipeline.producer_acquire(b_load2mma_producer_state) + # TMA load A/B + cute.copy( + tma_atom_a, + tAgA_slice[(None, a_load2trans_producer_state.count)], + tAsA[(None, a_load2trans_producer_state.index)], + tma_bar_ptr=a_load2trans_pipeline.producer_get_barrier( + a_load2trans_producer_state + ), + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, b_load2mma_producer_state.count)], + tBsB[(None, b_load2mma_producer_state.index)], + tma_bar_ptr=b_load2mma_pipeline.producer_get_barrier( + b_load2mma_producer_state + ), + mcast_mask=b_full_mcast_mask, + ) + a_load2trans_pipeline.producer_commit(a_load2trans_producer_state) + b_load2mma_pipeline.producer_commit(b_load2mma_producer_state) + a_load2trans_producer_state.advance() + b_load2mma_producer_state.advance() + if a_load2trans_producer_state.count < k_tile_cnt: + peek_load2trans_empty_status = ( + a_load2trans_pipeline.producer_try_acquire( + a_load2trans_producer_state + ) + ) + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # Wait A/B buffer empty + a_load2trans_pipeline.producer_tail(a_load2trans_producer_state) + b_load2mma_pipeline.producer_tail(b_load2mma_producer_state) + + # Specialized TMA load for scale tensor + if warp_idx == self.scale_tma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_tma_warps) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Persistent tile scheduling loop + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, (bidx, bidy, bidz), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + scale_load2trans_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_scale_load2trans_stage + ) + scale_k_tile_cnt = cute.size(mS_mkl.layout.shape[1][1]) + + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + # ((atom_v, rest_v), RestK) + tSgS_slice = tSgS[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # Filter zeros in rest mode + rest_filtered = cute.filter_zeros(tSgS_slice[(0, None)].layout) + tSgS_slice_filtered = cute.make_tensor( + tSgS_slice.iterator, + cute.make_layout( + (tSgS_slice.layout[0].shape, rest_filtered.shape), + stride=(tSgS_slice.layout[0].stride, rest_filtered.stride), + ), + ) + + scale_load2trans_producer_state.reset_count() + peek_scale_load2trans_empty_status = cutlass.Boolean(1) + if scale_load2trans_producer_state.count < scale_k_tile_cnt: + peek_scale_load2trans_empty_status = ( + scale_load2trans_pipeline.producer_try_acquire( + scale_load2trans_producer_state + ) + ) + for k_tile in cutlass.range(0, scale_k_tile_cnt, 1, unroll=1): + scale_load2trans_pipeline.producer_acquire( + scale_load2trans_producer_state, + peek_scale_load2trans_empty_status, + ) + # TMA load scale + cute.copy( + tma_atom_s, + tSgS_slice_filtered[ + (None, scale_load2trans_producer_state.count) + ], + tSsS[(None, scale_load2trans_producer_state.index)], + tma_bar_ptr=scale_load2trans_pipeline.producer_get_barrier( + scale_load2trans_producer_state + ), + mcast_mask=s_full_mcast_mask, + ) + + scale_load2trans_producer_state.advance() + peek_scale_load2trans_empty_status = cutlass.Boolean(1) + if scale_load2trans_producer_state.count < scale_k_tile_cnt: + peek_scale_load2trans_empty_status = ( + scale_load2trans_pipeline.producer_try_acquire( + scale_load2trans_producer_state + ) + ) + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # Wait scale buffer empty + scale_load2trans_pipeline.producer_tail(scale_load2trans_producer_state) + + # Specialized transform warps + if warp_idx >= self.transform_warp_id[0]: + cute.arch.setmaxregister_increase(self.num_regs_transform_warps) + transform_local_tidx = tidx - 32 * self.transform_warp_id[0] + # Partition tensors for transform input and output and set up the copy atom + # used for loading and storing transformed A tensor + ( + src_copy_a, + dst_copy_a, + tAsA_input, + tAsA_transform, + ) = mixed_input_utils.transform_partition( + self.transform_a_source, + self.scale_mode, + copy_atom_a_input, + copy_atom_a_transform, + sA_input, + ( + tCrA + if self.transform_a_source == tcgen05.OperandSource.TMEM + else sA_transform + ), + transform_local_tidx, + ) + # make rmem tensor for input A and transformed A + tArA = cute.make_rmem_tensor( + tAsA_input[(None, None, None, None, 0)].shape, tAsA_input.element_type + ) + tArA_transform = cute.make_rmem_tensor( + tAsA_input[(None, None, None, None, 0)].shape, self.mma_dtype + ) + # Partition scale tensor + smem_thr_copy_S = None + tSsS_trans = None + tSrS_copy = None + tSrS = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + smem_thr_copy_S, tSsS_trans, tSrS_copy, tSrS = ( + mixed_input_utils.scale_partition( + src_copy_a, tCsS, transform_local_tidx, self.mma_dtype + ) + ) + assert cute.size(tSrS, mode=[0]) == cute.size(tArA, mode=[0]), ( + "tSrS and tArA have different leading dimension" + ) + assert cute.size(tSrS) == cute.size(tArA), ( + "tSrS and tArA have different shape" + ) + # Deduce a sub-tile size and tile tensors + transform_tiler_size = min( + cute.size(cute.coalesce(tAsA_input.layout), mode=[0]), 64 + ) + transform_tiler = cute.make_layout(transform_tiler_size) + tArA_load = cute.flat_divide(tArA, transform_tiler) + tArA_load = cute.group_modes(tArA_load, 1, cute.rank(tArA_load)) + tSrS_load = ( + cute.flat_divide(tSrS, transform_tiler) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + tSrS_load = ( + cute.group_modes(tSrS_load, 1, cute.rank(tSrS_load)) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + tArA_transform_store = cute.flat_divide(tArA_transform, transform_tiler) + tArA_transform_store = cute.group_modes( + tArA_transform_store, 1, cute.rank(tArA_transform_store) + ) + + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, (bidx, bidy, bidz), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + a_load2trans_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + self.num_load2trans_stage, + ) + scale_load2trans_consumer_state = ( + pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + self.num_scale_load2trans_stage, + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + trans2mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, + self.num_trans2mma_stage, + ) + while work_tile.is_valid_tile: + a_load2trans_consumer_state.reset_count() + peek_load2trans_full_status = cutlass.Boolean(1) + if a_load2trans_consumer_state.count < k_tile_cnt: + peek_load2trans_full_status = ( + a_load2trans_pipeline.consumer_try_wait( + a_load2trans_consumer_state + ) + ) + peek_scale_load2trans_full_status = cutlass.Boolean(1) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + scale_load2trans_consumer_state.reset_count() + peek_scale_load2trans_full_status = ( + scale_load2trans_pipeline.consumer_try_wait( + scale_load2trans_consumer_state + ) + ) + trans2mma_producer_state.reset_count() + peek_trans2mma_empty_status = cutlass.Boolean(1) + if trans2mma_producer_state.count < k_tile_cnt: + peek_trans2mma_empty_status = ( + trans2mma_pipeline.producer_try_acquire( + trans2mma_producer_state + ) + ) + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + a_load2trans_pipeline.consumer_wait( + a_load2trans_consumer_state, peek_load2trans_full_status + ) + # Load A from shared memory + tAsA_input_slice = tAsA_input[ + (None, None, None, None, a_load2trans_consumer_state.index) + ] + tAsA_input_slice = cute.flat_divide( + tAsA_input_slice, transform_tiler + ) + tAsA_input_slice = cute.group_modes( + tAsA_input_slice, 1, cute.rank(tAsA_input_slice) + ) + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + scale_load2trans_pipeline.consumer_wait( + scale_load2trans_consumer_state, + peek_scale_load2trans_full_status, + ) + trans2mma_pipeline.producer_acquire( + trans2mma_producer_state, peek_trans2mma_empty_status + ) + # load scale tensor when needed + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + if k_tile % num_k_tiles_per_scale == 0: + tSsS_slice = tSsS_trans[ + ( + None, + None, + None, + None, + scale_load2trans_consumer_state.index, + ) + ] + tSsS_slice_filtered = cute.make_tensor( + tSsS_slice.iterator, + cute.filter_zeros(tSsS_slice.layout), + ) + cute.autovec_copy(tSsS_slice_filtered, tSrS_copy) + cur_scale_load2trans_consumer_state = ( + scale_load2trans_consumer_state.clone() + ) + if (k_tile + 1) % num_k_tiles_per_scale == 0: + scale_load2trans_consumer_state.advance() + + cur_a_load2trans_consumer_state = ( + a_load2trans_consumer_state.clone() + ) + for idx in cutlass.range_constexpr(cute.size(tArA_load, mode=[1])): + # Load A from shared memory + cute.autovec_copy( + tAsA_input_slice[(None, idx)], + tArA_load[(None, idx)], + ) + if cutlass.const_expr( + idx == cute.size(tArA_load, mode=[1]) - 1 + ): + a_load2trans_consumer_state.advance() + if a_load2trans_consumer_state.count < k_tile_cnt: + peek_load2trans_full_status = ( + a_load2trans_pipeline.consumer_try_wait( + a_load2trans_consumer_state + ) + ) + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + peek_scale_load2trans_full_status = ( + scale_load2trans_pipeline.consumer_try_wait( + scale_load2trans_consumer_state + ) + ) + # Convert it to mma dtype + tensor_transformed = mixed_input_utils.cvt_tensor_a( + tArA_load[(None, idx)], self.mma_dtype, self.shuffle_a + ) + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + scale = cute.TensorSSA( + tSrS_load[(None, idx)].load(), + tensor_transformed.shape, + self.mma_dtype, + ) + # Apply scale + tensor_transformed = tensor_transformed * scale + tArA_transform_store[(None, idx)].store(tensor_transformed) + # Store transformed A to tensor memory or shared memory + mixed_input_utils.store_transformed_a( + tArA_transform, + tAsA_transform[ + (None, None, None, None, trans2mma_producer_state.index) + ], + dst_copy_a, + ) + # Ensure all transform threads have finished the copy and reached the fence + self.transform_sync_barrier.arrive_and_wait() + if cutlass.const_expr( + self.transform_a_source == tcgen05.OperandSource.TMEM + ): + cute.arch.fence_view_async_tmem_store() + else: + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + # Signal the completion of transformation + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + scale_load2trans_pipeline.consumer_release( + cur_scale_load2trans_consumer_state + ) + a_load2trans_pipeline.consumer_release( + cur_a_load2trans_consumer_state + ) + # Signal the completion of transformation + trans2mma_pipeline.producer_commit(trans2mma_producer_state) + trans2mma_producer_state.advance() + if trans2mma_producer_state.count < k_tile_cnt: + peek_trans2mma_empty_status = ( + trans2mma_pipeline.producer_try_acquire( + trans2mma_producer_state + ) + ) + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # Wait a_transform buffer empty + trans2mma_pipeline.producer_tail(trans2mma_producer_state) + + # Specialized MMA warp + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_mma_warp) + tCtAcc_base = accumulators + # Persistent tile scheduling loop + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, (bidx, bidy, bidz), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + trans2mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_trans2mma_stage + ) + b_load2mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_load2trans_stage + ) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + b_load2mma_consumer_state.reset_count() + trans2mma_consumer_state.reset_count() + peek_trans2mma_full_status = cutlass.Boolean(1) + if is_leader_cta: + if trans2mma_consumer_state.count < k_tile_cnt: + peek_trans2mma_full_status = ( + trans2mma_pipeline.consumer_try_wait( + trans2mma_consumer_state + ) + ) + acc_pipeline.producer_acquire(acc_producer_state) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + # Mma mainloop + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + trans2mma_pipeline.consumer_wait( + trans2mma_consumer_state, peek_trans2mma_full_status + ) + b_load2mma_pipeline.consumer_wait(b_load2mma_consumer_state) + num_kblocks = cute.size(tCrA, mode=[2]) + for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): + kblock_coord_a = ( + None, + None, + kblock_idx, + trans2mma_consumer_state.index, + ) + kblock_coord_b = ( + None, + None, + kblock_idx, + b_load2mma_consumer_state.index, + ) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblock_coord_a], + tCrB[kblock_coord_b], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + trans2mma_pipeline.consumer_release(trans2mma_consumer_state) + b_load2mma_pipeline.consumer_release(b_load2mma_consumer_state) + trans2mma_consumer_state.advance() + b_load2mma_consumer_state.advance() + peek_trans2mma_full_status = cutlass.Boolean(1) + if trans2mma_consumer_state.count < k_tile_cnt: + peek_trans2mma_full_status = ( + trans2mma_pipeline.consumer_try_wait( + trans2mma_consumer_state + ) + ) + # Async arrive accumulator buffer full + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # Wait for accumulator buffer empty + acc_pipeline.producer_tail(acc_producer_state) + + # Specialized epilogue warps + if warp_idx < self.mma_warp_id: + cute.arch.setmaxregister_increase(self.num_regs_epilogue_warps) + epi_tidx = tidx + tCtAcc_base = accumulators + # Partition for epilogue + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc, + ) = mixed_input_utils.epilog_tmem_copy_and_partition( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + epi_tidx, + tCtAcc_base, + tCgC, + epi_tile, + self.use_2cta_instrs, + ) + + tTR_rC = None + tiled_copy_r2s = None + simt_atom = None + tRS_rC = None + tRS_sC = None + bSG_sC = None + bSG_gC_partitioned = None + tTR_gC_partitioned = None + if cutlass.const_expr(self.use_tma_store): + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = ( + mixed_input_utils.epilog_smem_copy_and_partition( + self.c_layout, + self.c_dtype, + self.acc_dtype, + tiled_copy_t2r, + tTR_rC, + epi_tidx, + sC, + ) + ) + ( + tma_atom_c, + bSG_sC, + bSG_gC_partitioned, + ) = self.epilog_gmem_copy_and_partition( + epi_tidx, tma_atom_c, tCgC, epi_tile, sC + ) + else: + ( + simt_atom, + tTR_rC, + tTR_gC_partitioned, + ) = self.epilog_gmem_copy_and_partition( + epi_tidx, tiled_copy_t2r, tCgC, epi_tile, sC + ) + # Persistent tile scheduling loop + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, (bidx, bidy, bidz), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + c_pipeline = None + if cutlass.const_expr(self.use_tma_store): + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilog_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, + producer_group=c_producer_group, + ) + + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + bSG_gC = None + tTR_gC = None + if cutlass.const_expr(self.use_tma_store): + bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)] + else: + tTR_gC = tTR_gC_partitioned[ + (None, None, None, None, None, *mma_tile_coord_mnl) + ] + + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + # Wait for accumulator buffer full + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + if cutlass.const_expr(self.use_tma_store): + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + else: + tTR_gC = cute.group_modes(tTR_gC, 3, cute.rank(tTR_gC)) + + # Store accumulator to global memory in subtiles + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt + for subtile_idx in cutlass.range(subtile_cnt): + # Load accumulator from tensor memory buffer to register + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + if cutlass.const_expr(self.use_tma_store): + # Convert to C type + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = acc_vec.to(self.c_dtype) + tRS_rC.store(acc_vec) + c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage + # Store C to shared memory + cute.copy( + tiled_copy_r2s, + tRS_rC, + tRS_sC[(None, None, None, c_buffer)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.epilog_sync_barrier.arrive_and_wait() + # TMA store C to global memory + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + ) + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + else: + # Convert to C type + acc_vec = tTR_rAcc.load() + acc_vec = acc_vec.to(self.c_dtype) + tTR_rC.store(acc_vec) + # Store C to global memory + cute.autovec_copy( + tTR_rC, tTR_gC[(None, None, None, subtile_idx)] + ) + # Async arrive accumulator buffer empty + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # Dealloc the tensor memory buffer + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + if cutlass.const_expr(self.use_tma_store): + c_pipeline.producer_tail() + + # Idle warp + if warp_idx == self.idle_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_idle_warp) + + def epilog_gmem_copy_and_partition( + self, + tidx: cutlass.Int32, + atom: cute.CopyAtom | cute.TiledCopy, + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + sC: cute.Tensor, + ) -> tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: + """Partitions source and destination tensors for a TMA store or SIMT store. + """ + if self.use_tma_store: + tma_atom_c, bSG_sC, bSG_gC, _, _ = ( + mixed_input_utils.epilog_gmem_copy_and_partition( + self.c_dtype, tidx, atom, None, gC_mnl, None, epi_tile, sC + ) + ) + return tma_atom_c, bSG_sC, bSG_gC + else: + _, _, _, simt_atom, tTR_gC = ( + mixed_input_utils.epilog_gmem_copy_and_partition( + self.c_dtype, tidx, None, atom, None, gC_mnl, epi_tile, sC + ) + ) + # (T2R, T2R_M, T2R_N) + tTR_rC = cute.make_rmem_tensor( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.c_dtype + ) + simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) + return simt_atom, tTR_rC, tTR_gC diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_persistent_impl.py b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_persistent_impl.py new file mode 100644 index 000000000..2605dcfe7 --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_persistent_impl.py @@ -0,0 +1,1277 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from collections.abc import Callable +from typing import Literal, Union + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cutlass_dsl import Int32 +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +from .operator_helpers import PersistentTileSchedulerType + +""" +A high-performance persistent batched dense GEMM example for the NVIDIA Blackwell SM100 architecture +using CUTE DSL. +- Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") +- Matrix B is KxNxL, L is batch dimension, B can be row-major("N") or column-major("K") + - Internally, this is treated as an NxKxL matrix +- Matrix C is MxNxL, L is batch dimension, C can be row-major("N") or column-major("M") + +This GEMM kernel supports the following features: + - Utilizes Tensor Memory Access (TMA) for efficient memory operations + - Utilizes Blackwell's tcgen05.mma for matrix multiply-accumulate (MMA) operations (including 2cta mma instructions) + - Implements TMA multicast with cluster to reduce L2 memory traffic + - Allows user to switch between + - CLC dynamic persistent tile scheduling to have near perfect load balancing + The CLC dynamic persistent scheduling technique performs dynamic loading balancing. + It has the ability to adapt available SMs rather than a statically selected number. To support this, + a new instruction is introduced to query for a new tile to compute. This new instruction is similar + to programmatic multicast in context of clusters in that the same starting tile ID for a given cluster + is broadcasted to all threadblocks in the cluster. + See `PTX documentation `. + - Static persistent tile scheduling + - Support warp specialization to avoid explicit pipelining between mainloop load and mma + +This GEMM works as follows: +1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. +2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction. +3. EPILOGUE warp: + - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. + - Type convert C matrix to output type. + - Optionally store C matrix from registers (RMEM) to shared memory (SMEM) to global memory (GMEM) with TMA operations, + or directly store C matrix from registers (RMEM) to global memory (GMEM) without TMA operations. + - Optionally accept an elementwise lambda function epilogue_op to apply to the output tensor: + e.g., relu can set epilogue_op = lambda x: cute.where(x > 0, x, cute.full_like(x, 0)) + +SM100 tcgen05.mma instructions operate as follows: +- Read matrix A from SMEM +- Read matrix B from SMEM +- Write accumulator to TMEM +The accumulator in TMEM must then be loaded to registers before writing back to GMEM. + +Constraints are: +* Supported input data types: fp16, bf16, tf32, int8, uint8, fp8 (e4m3fn, e5m2), + see detailed valid dtype combinations in below PersistentDenseGemmOperator class documentation +* A/B tensor must have the same data type +* Mma tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) +* Mma tiler N must be 32-256, step 32 +* Cluster shape M/N must be positive and power of 2, total cluster size <= 16 +* Cluster shape M must be multiple of 2 if use_2cta_instrs=True +* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, + i.e, number of elements is a multiple of 4, 8, and 16 for TFloat32, + Float16/BFloat16, and Int8/Uint8/Float8, respectively. +* OOB tiles are not allowed when TMA store is disabled +""" + + +def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: tuple[int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + smem_capacity: int, + occupancy: int, + use_tma_store: bool, + c_smem_layout: cute.Layout | None, +) -> tuple[int, int, int]: + """Computes the number of stages for A/B/C operands based on heuristics. + + :param tiled_mma: The tiled MMA object defining the core computation. + :type tiled_mma: cute.TiledMma + :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler. + :type mma_tiler_mnk: tuple[int, int, int] + :param a_dtype: Data type of operand A. + :type a_dtype: type[cutlass.Numeric] + :param b_dtype: Data type of operand B. + :type b_dtype: type[cutlass.Numeric] + :param c_dtype: Data type of operand C (output). + :type c_dtype: type[cutlass.Numeric] + :param smem_capacity: Total available shared memory capacity in bytes. + :type smem_capacity: int + :param occupancy: Target number of CTAs per SM (occupancy). + :type occupancy: int + :param use_tma_store: Whether TMA store is enabled. + :type use_tma_store: bool + :param c_smem_layout: Layout of C operand in shared memory, or None if not using TMA store. + :type c_smem_layout: Union[cute.Layout, None] + + :return: A tuple containing the computed number of stages for: + (ACC stages, A/B operand stages, C stages) + :rtype: tuple[int, int, int] + """ + # Default ACC stages + num_acc_stage = 2 + + # Default C stages + num_c_stage = 2 if use_tma_store else 0 + + # Calculate smem layout and size for one stage of A, B, and C with 1-stage + a_smem_layout_stage_one = utils.sm100.make_smem_layout_a( + tiled_mma, mma_tiler_mnk, a_dtype, 1 + ) + b_smem_layout_staged_one = utils.sm100.make_smem_layout_b( + tiled_mma, mma_tiler_mnk, b_dtype, 1 + ) + + ab_bytes_per_stage = cute.size_in_bytes( + a_dtype, a_smem_layout_stage_one + ) + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + mbar_helpers_bytes = 1024 + + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout) + c_bytes = c_bytes_per_stage * num_c_stage + + # Calculate A/B stages: + # Start with total smem per CTA (capacity / occupancy) + # Subtract reserved bytes and initial C stages bytes + # Divide remaining by bytes needed per A/B stage + num_ab_stage = ( + smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes) + ) // ab_bytes_per_stage + + # Refine epilogue stages: + # Calculate remaining smem after allocating for A/B stages and reserved bytes + # Add remaining unused smem to epilogue + if use_tma_store: + num_c_stage += ( + smem_capacity + - occupancy * ab_bytes_per_stage * num_ab_stage + - occupancy * (mbar_helpers_bytes + c_bytes) + ) // (occupancy * c_bytes_per_stage) + return num_acc_stage, num_ab_stage, num_c_stage + + +class PersistentDenseGemmKernel: + def __init__( + self, + acc_dtype: type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + use_tma_store: bool, + persistent_tile_scheduler_type: PersistentTileSchedulerType, + epilogue_op: Callable, + ): + """Initializes the configuration for a Blackwell dense GEMM kernel. + + This configuration includes several key aspects: + + 1. MMA Instruction Settings (tcgen05): + - acc_dtype: Data types for MMA accumulator. + - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler. + - use_2cta_instrs: Boolean indicating if the tcgen05 MMA variant + with cta_group=2 should be used. + + 2. Cluster Shape: + - cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster. + + 3. Output C tensor store mode: + - use_tma_store: Boolean indicating whether to use Tensor Memory Access (TMA) for storing results. + + :param acc_dtype: Data type of the accumulator. + :type acc_dtype: type[cutlass.Numeric] + :param mma_tiler_mn: tuple (M, N) shape of the MMA instruction. + :type mma_tiler_mn: tuple[int, int] + :param use_2cta_instrs: Boolean, True to use cta_group=2 MMA variant. + :type use_2cta_instrs: bool + :param cluster_shape_mn: tuple (ClusterM, ClusterN) shape of the cluster. + :type cluster_shape_mn: tuple[int, int] + :param use_tma_store: Use Tensor Memory Access (TMA) or normal store for output C tensor. + :type use_tma_store: bool + :param epilogue_op: Elementwise lambda function to apply to the output tensor. + :type epilogue_op: Callable + """ + self.acc_dtype: type[cutlass.Numeric] = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + # K dimension is deferred in _setup_attributes + self.mma_tiler_mn = mma_tiler_mn + self.mma_tiler = (*mma_tiler_mn, 1) + self.use_tma_store = use_tma_store + self.arch = "sm_100" + self.epilogue_op = epilogue_op + + self.cta_group = ( + tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + self.occupancy = 1 + # Set specialized warp ids + self.epilogue_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + + self.use_static_persistent_scheduler = False + self.use_clc_dynamic_scheduler = False + if persistent_tile_scheduler_type is PersistentTileSchedulerType.CLC: + self.use_clc_dynamic_scheduler = True + self.sched_warp_id = 6 + self.threads_per_cta = 32 * len( + ( + self.mma_warp_id, + self.tma_warp_id, + self.sched_warp_id, + *self.epilogue_warp_id, + ) + ) + # runtime values for swizzle size and raster along not yet supported + # and these parameters are set to use defaults + elif ( + persistent_tile_scheduler_type + is PersistentTileSchedulerType.STATIC_PERSISTENT + ): + self.use_static_persistent_scheduler = True + self.threads_per_cta = 32 * len( + (self.mma_warp_id, self.tma_warp_id, *self.epilogue_warp_id) + ) + else: + raise ValueError( + "Unsupported persistent tile scheduler type: " + f"{persistent_tile_scheduler_type!r}" + ) + + # Set barrier id for cta sync, epilogue sync and tmem ptr sync + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + self.tmem_dealloc_sync_bar_id = 3 + + self.smem_capacity = utils.get_smem_capacity_in_bytes(self.arch) + + def _setup_attributes(self): + """Set up configurations that are dependent on GEMM inputs + + This method configures various attributes based on the input tensor properties + (data types, leading dimensions) and kernel settings: + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B + - Computing epilogue subtile + - Setting up A/B/C stage counts in shared memory + - Computing A/B/C shared memory layout + - Computing tensor memory allocation columns + """ + # Configure tiled mma + tiled_mma = utils.sm100.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + # Compute mma/cluster/tile shapes + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 # kphase: number of MMA atoms along K before advancing to the next pipeline stage + self.mma_tiler = ( + self.mma_tiler[0], + self.mma_tiler[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + + # Compute cluster layout + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + + # Compute number of multicast CTAs for A/B + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + # Compute epilogue subtile + if cutlass.const_expr(self.use_tma_store): + self.epi_tile = utils.sm100.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + else: + self.epi_tile = self.cta_tile_shape_mnk[:2] + + c_smem_layout = None + if cutlass.const_expr(self.use_tma_store): + c_smem_layout = utils.sm100.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, 1 + ) + + # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = _compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.c_dtype, + self.smem_capacity, + self.occupancy, + self.use_tma_store, + c_smem_layout, + ) + + if self.use_clc_dynamic_scheduler: + # Setup clc stage by default, currently only single-stage CLC pipeline is supported + self.num_clc_stage = 1 + + # Compute A/B/C shared memory layout + self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( + tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage + ) + self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( + tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage + ) + + self.c_smem_layout_staged = None + if self.use_tma_store: + self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage + ) + + # Compute the number of tensor memory allocation columns + self.num_tmem_alloc_cols = self._compute_num_tmem_alloc_cols( + tiled_mma, self.mma_tiler, self.num_acc_stage + ) + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: cuda.CUstream, + max_active_clusters: cutlass.Constexpr = None, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + """Execute the GEMM operation in steps: + - Setup static attributes before smem/grid/tma computation + - Setup TMA load/store atoms and tensors + - Compute grid size with regard to hardware constraints + - Define shared storage for kernel + - Launch the kernel synchronously + + :param a: Input tensor A of shape (L, M, K) or (M, K) + :type a: cute.Tensor + :param b: Input tensor B of shape (L, K, N) or (K, N) + :type b: cute.Tensor + :param c: Output tensor C of shape (L, M, N) or (M, N) + :type c: cute.Tensor + :param max_active_clusters: Maximum number of active clusters + :type max_active_clusters: cutlass.Constexpr + :param stream: CUDA stream for asynchronous execution + :type stream: cuda.CUstream + :param epilogue_op: Optional elementwise lambda function to apply to the output tensor + :type epilogue_op: cutlass.Constexpr + :raises TypeError: If input data types are incompatible with the MMA instruction. + :raises AssertionError: If OOB (Out-Of-Bounds) tiles are present when TMA store is disabled. + """ + + def add_batch_mode(tensor: cute.Tensor) -> cute.Tensor: + return cute.make_tensor( + tensor.iterator, + cute.prepend(tensor.layout, cute.make_layout(1), up_to_rank=3), + ) + + # Normalise A/B/C to rank-3 by prepending a batch mode of size 1 when + # the input is 2D (non-batched). This allows the rest of the kernel to + # treat both batched (L, M, K) and non-batched (M, K) inputs uniformly + # as (1, M, K) without any special-casing downstream. + a = add_batch_mode(a) + b = add_batch_mode(b) + c = add_batch_mode(c) + + # Permute tensor modes from torch to cute convention + # A: (L, M, K) -> (M, K, L) + a = cute.make_tensor(a.iterator, cute.select(a.layout, [1, 2, 0])) + # B: (L, K, N) -> (N, K, L) + b = cute.make_tensor(b.iterator, cute.select(b.layout, [2, 1, 0])) + # C: (L, M, N) -> (M, N, L) + c = cute.make_tensor(c.iterator, cute.select(c.layout, [1, 2, 0])) + + # Setup static attributes before smem/grid/tma computation + self.a_dtype: type[cutlass.Numeric] = a.element_type + self.b_dtype: type[cutlass.Numeric] = b.element_type + self.c_dtype: type[cutlass.Numeric] = c.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + + # Check if input data types are compatible with MMA instruction + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + tiled_mma = utils.sm100.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = utils.sm100.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + + # Setup TMA load for B + b_op = utils.sm100.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, tiled_mma.thr_id + ) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size + + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + # Response size is 4B * 4 elements + self.num_clc_response_bytes = 16 + + # Setup TMA store for C + tma_atom_c = None + tma_tensor_c = None + if cutlass.const_expr(self.use_tma_store): + epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout, self.epi_tile + ) + + # Compute grid size + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + self.tile_sched_params, grid = ( + self._compute_grid_clc_dynamic_persistent_scheduler( + c, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + ) + ) + else: + self.tile_sched_params, grid = ( + self._compute_grid_static_persistent_scheduler( + c, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + max_active_clusters, + ) + ) + + # Launch the kernel synchronously + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c if self.use_tma_store else c, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + epilogue_op, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom | None, + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: cute.Layout | cute.ComposedLayout | None, + epi_tile: cute.Tile, + tile_sched_params: Union[ + utils.ClcDynamicPersistentTileSchedulerParams, + utils.PersistentTileSchedulerParams, + ], + epilogue_op: cutlass.Constexpr, + ): + """GPU device kernel performing the Persistent batched GEMM computation.""" + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # + # Prefetch tma desc + # + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + if cutlass.const_expr(self.use_tma_store): + cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + is_first_cta_in_cluster = cta_rank_in_cluster == 0 + + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + # Coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + # + # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier + # + # Define shared storage for kernel + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_ab_stage * 2 + ] + acc_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_acc_stage * 2 + ] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2] + clc_response_align_bytes = self.num_clc_response_bytes + clc_response: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int32, 4], + clc_response_align_bytes, + ] + else: + + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_ab_stage * 2 + ] + acc_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_acc_stage * 2 + ] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # Initialize mainloop ab_pipeline (barrier) and states + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilogue_warp_id) * ( + 2 if use_2cta_instrs else 1 + ) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + # Initialize clc_pipeline (barrier) and states + clc_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread + ) + cluster_size = cute.size(self.cluster_shape_mn) + num_clc_consumer_threads = 32 * ( + 1 + cluster_size * (1 + len(self.epilogue_warp_id) + 1) + ) + clc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_clc_consumer_threads + ) + clc_pipeline = pipeline.PipelineClcFetchAsync.create( + barrier_storage=storage.clc_mbar_ptr.data_ptr(), + num_stages=self.num_clc_stage, + producer_group=clc_pipeline_producer_group, + consumer_group=clc_pipeline_consumer_group, + tx_count=self.num_clc_response_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem_dealloc_barrier = None + if cutlass.const_expr(not self.use_tma_store): + tmem_dealloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_dealloc_sync_bar_id, + num_threads=32 * len(self.epilogue_warp_id), + ) + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) + + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + # Initial clc response pointer + clc_response_ptr = storage.clc_response.data_ptr() + clc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_clc_stage + ) + + # + # Setup smem tensor A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + sA = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + # (MMA, MMA_N, MMA_K, STAGE) + sB = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + + # + # Compute multicast mask for A/B buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + # + # Local_tile partition global tensors + # + # (bM, bK, RestM, RestK, RestL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, RestN, RestK, RestL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bN, RestM, RestN, RestL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # + # Partition global tensor for TiledMMA_A/B/C + # + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) + tCgC = thr_mma.partition_C(gC_mnl) + + # + # Partition global/shared tensor for TMA load A/B + # + # TMA load A partition_S/D + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA load B partition_S/D + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # + # Cluster wait before tensor memory alloc + # + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + + # + # Construct the scheduler + # + + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + tile_sched = utils.ClcDynamicPersistentTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + cute.arch.grid_dim(), + clc_response_ptr, + ) + else: + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + # + # Specialized TMA load warp + # + + if warp_idx == self.tma_warp_id: + # + # Persistent tile scheduling loop + # + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), RestK) + tAgA_slice = tAgA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # ((atom_v, rest_v), RestK) + tBgB_slice = tBgB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + # + # Tma load loop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Conditionally wait for AB buffer empty + handle = ab_producer.acquire_and_advance(peek_ab_empty_status) + + # TMA load A/B + cute.copy( + tma_atom_a, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=b_full_mcast_mask, + ) + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + + # + # Advance to next tile + # + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + else: + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait A/B buffer empty + # + ab_producer.tail() + + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + # + # Sched warp + # + if warp_idx == self.sched_warp_id and is_first_cta_in_cluster: + # + # Persistent tile scheduling loop + # + clc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.ProducerConsumer, self.num_clc_stage + ) + + while work_tile.is_valid_tile: + # + # Advance to next tile + # + clc_pipeline.producer_acquire(clc_producer_state) + mbarrier_addr = clc_pipeline.producer_get_barrier( + clc_producer_state + ) + tile_sched.advance_to_next_work(mbarrier_addr) + clc_producer_state.advance() + + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + clc_pipeline.producer_tail(clc_producer_state) + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Persistent tile scheduling loop + # + + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # Set tensor memory buffer for current tile + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + + # Peek (try_wait) AB buffer full for k_tile = 0 + ab_consumer.reset() + peek_ab_full_status = cutlass.Boolean(1) + if is_leader_cta: + peek_ab_full_status = ab_consumer.try_wait() + + # + # Wait for accumulator buffer empty + # + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + + # + # Reset the ACCUMULATE field for each tile + # + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + # + # Mma mainloop + # + for k_tile in range(k_tile_cnt): + if is_leader_cta: + # Conditionally wait for AB buffer full + handle = ab_consumer.wait_and_advance(peek_ab_full_status) + + # tCtAcc += tCrA * tCrB + num_kblocks = cute.size(tCrA, mode=[2]) + for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): + kblk_crd = (None, None, kblk_idx, handle.index) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblk_crd], + tCrB[kblk_crd], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + # Async arrive AB buffer empty + handle.release() + + # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 + peek_ab_full_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_full_status = ab_consumer.try_wait() + + # + # Async arrive accumulator buffer full + # + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + # + # Advance to next tile + # + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + else: + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait for accumulator buffer empty + # + acc_pipeline.producer_tail(acc_producer_state) + + sC = None + if cutlass.const_expr(self.use_tma_store): + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC = smem.allocate_tensor( + element_type=self.c_dtype, + layout=c_smem_layout_staged.outer, + byte_alignment=128, + swizzle=c_smem_layout_staged.inner, + ) + + # + # Specialized epilogue warps + # + if warp_idx < self.mma_warp_id: + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Persistent tile scheduling loop for epilogue + # + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + if cutlass.const_expr(self.use_tma_store): + assert tma_atom_c is not None and sC is not None + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilogue_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, producer_group=c_producer_group + ) + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + if cutlass.const_expr(not self.use_clc_dynamic_scheduler): + # + # Pre-advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + else: + # CLC epilogue warps don't call advance_to_next_work (only + # the sched warp does), so track the tile count locally for + # correct C-store pipeline buffer indexing. + tile_sched._num_tiles_executed += Int32(1) + + num_tiles_executed = tile_sched.num_tiles_executed + if cutlass.const_expr(self.use_tma_store): + acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( + self, + tidx, + warp_idx, + tma_atom_c, + tCtAcc_base, + sC, + tCgC, + epi_tile, + num_tiles_executed, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + c_pipeline, + ) + else: + acc_consumer_state = utils.gemm.sm100.epilogue( + self, + tidx, + tCtAcc_base, + tCgC, + epi_tile, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + ) + + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + # + # Advance to next tile + # + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + + if cutlass.const_expr(self.use_tma_store): + # Wait for C store complete + c_pipeline.producer_tail() + else: + # Synchronize before TMEM dealloc (done by the caller) + tmem_dealloc_barrier.arrive_and_wait() + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + @staticmethod + def _compute_grid_clc_dynamic_persistent_scheduler( + c: cute.Tensor, + cta_tile_shape_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + ) -> tuple[utils.ClcDynamicPersistentTileSchedulerParams, tuple[int, int, int]]: + """Use CLC dynamic persistent tile scheduler to compute the grid size for the output tensor C. + + :param c: The output tensor C + :type c: cute.Tensor + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + + :return: A tuple containing: + - tile_sched_params: Parameters for the CLC dynamic persistent tile scheduler. + - grid: Grid shape for kernel launch. + :rtype: tuple[utils.ClcDynamicPersistentTileSchedulerParams, tuple[int, int, int]] + """ + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + # swizzle_size and raster_along_m are set to use defaults + tile_sched_params = utils.ClcDynamicPersistentTileSchedulerParams( + num_ctas_mnl, + cluster_shape_mnl, + ) + grid = utils.ClcDynamicPersistentTileScheduler.get_grid_shape(tile_sched_params) + + return tile_sched_params, grid + + @staticmethod + def _compute_grid_static_persistent_scheduler( + c: cute.Tensor, + cta_tile_shape_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + max_active_clusters: cutlass.Constexpr, + ) -> tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]]: + """Use persistent tile scheduler to compute the grid size for the output tensor C. + + :param c: The output tensor C + :type c: cute.Tensor + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + :param max_active_clusters: Maximum number of active clusters. + :type max_active_clusters: cutlass.Constexpr + + :return: A tuple containing: + - tile_sched_params: Parameters for the persistent tile scheduler. + - grid: Grid shape for kernel launch. + :rtype: tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]] + """ + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + + return tile_sched_params, grid + + @staticmethod + def _compute_num_tmem_alloc_cols( + tiled_mma: cute.TiledMma, + mma_tiler: tuple[int, int, int], + num_acc_stage: int, + ) -> int: + """Compute the number of tensor memory allocation columns. + + :param tiled_mma: The tiled MMA object defining the core computation. + :type tiled_mma: cute.TiledMma + :param mma_tiler: The shape (M, N, K) of the MMA tile. + :type mma_tiler: tuple[int, int, int] + :param num_acc_stage: The stage of the accumulator tensor. + :type num_acc_stage: int + + :return: The number of tensor memory allocation columns. + :rtype: int + """ + acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, num_acc_stage)) + num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake) + + return num_tmem_alloc_cols diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_persistent_preferred_cluster_impl.py b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_persistent_preferred_cluster_impl.py new file mode 100644 index 000000000..2f84154cf --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_persistent_preferred_cluster_impl.py @@ -0,0 +1,1176 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from collections.abc import Callable +from typing import Union + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cutlass_dsl import Int32 +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +from .operator_helpers import PersistentTileSchedulerType +from .sm100_persistent_impl import PersistentDenseGemmKernel + +""" +A high-performance persistent batched dense GEMM with preferred cluster support for the NVIDIA +Blackwell SM100 architecture using CUTE DSL. + +Extends the base persistent dense GEMM (see ``sm100_persistent_impl.py``) with +support for two cluster shapes launched simultaneously: + +- **Preferred cluster** (larger): handles the majority of output tiles, benefiting from + wider TMA multicast and optional 2-CTA MMA instructions. +- **Fallback cluster** (smaller): handles the residual tiles that cannot fill a complete + preferred cluster, ensuring all SMs remain occupied. + +This GEMM works as follows: +1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using + TMA operations. The TMA atom is selected based on the cluster shape each CTA belongs to. +2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma. +3. EPILOGUE warp: + - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) via tcgen05.ld. + - Type convert C matrix to output type. + - Store C matrix from registers to global memory, optionally through SMEM via TMA. + - Optionally apply an elementwise epilogue function to the output. + +Constraints are: +* Same input data type support as PersistentDenseGemmKernel (fp16, bf16, tf32, int8, + uint8, fp8). +* A/B tensor must have the same data type. +* Mma tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True). +* Mma tiler N must be 32-256, step 32. +* Preferred cluster shape M/N must each be a power of 2 (≥ 1) and a multiple of the + corresponding fallback cluster shape M/N. Total cluster size ≤ 16. +* Cluster shape M must be a multiple of 2 when use_2cta_instrs=True, for both shapes. +* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, i.e. the + element count must be a multiple of 4 (TFloat32), 8 (Float16/BFloat16), or 16 + (Int8/Uint8/Float8) — matching the 16-byte TMA alignment requirement. +""" + + +class PersistentDenseGemmPreferredClusterKernel(PersistentDenseGemmKernel): + """Persistent dense GEMM with preferred cluster (preferred + fallback cluster shapes). + + Extends :class:`PersistentDenseGemmKernel` to support launching with two + cluster shapes simultaneously. A larger *preferred* cluster shape is used for + the bulk of the computation (better TMA multicast and optional 2-CTA MMA), while + a smaller *fallback* cluster shape handles leftover tiles. + """ + + def __init__( + self, + acc_dtype: type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: tuple[int, int], + preferred_cluster_shape_mn: tuple[int, int], + fallback_cluster_shape_mn: tuple[int, int], + use_tma_store: bool, + persistent_tile_scheduler_type: PersistentTileSchedulerType, + epilogue_op: Callable, + ): + """Initialise a preferred-cluster GEMM kernel with a preferred (larger) and a fallback + (smaller) cluster shape. The base class is constructed with the fallback shape + so that shared-memory layout and pipeline stage counts are sized conservatively; + preferred-cluster TMA atoms and multicast masks are derived separately in + :meth:`_setup_attributes`. + + :param acc_dtype: Accumulator data type. + :type acc_dtype: type[cutlass.Numeric] + :param use_2cta_instrs: Whether to use the tcgen05 cta_group=2 MMA variant. + :type use_2cta_instrs: bool + :param mma_tiler_mn: (M, N) MMA instruction tile shape. + :type mma_tiler_mn: tuple[int, int] + :param preferred_cluster_shape_mn: Preferred (larger) cluster shape (ClusterM, ClusterN). + :type preferred_cluster_shape_mn: tuple[int, int] + :param fallback_cluster_shape_mn: Fallback (smaller) cluster shape (ClusterM, ClusterN). + :type fallback_cluster_shape_mn: tuple[int, int] + :param use_tma_store: Whether to use TMA for storing the output C tensor. + :type use_tma_store: bool + :param persistent_tile_scheduler_type: Persistent tile scheduler variant. + :type persistent_tile_scheduler_type: PersistentTileSchedulerType + :param epilogue_op: Elementwise function applied to the output after GEMM. + :type epilogue_op: Callable + """ + # Initialise the base class with the fallback cluster shape so that all + # shared-memory and stage calculations are driven by the smaller shape. + # The preferred cluster uses the same smem layout; only the TMA atom and + # the cluster-pipeline synchronisation differ. + super().__init__( + acc_dtype=acc_dtype, + use_2cta_instrs=use_2cta_instrs, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=fallback_cluster_shape_mn, + use_tma_store=use_tma_store, + persistent_tile_scheduler_type=persistent_tile_scheduler_type, + epilogue_op=epilogue_op, + ) + self.preferred_cluster_shape_mn = preferred_cluster_shape_mn + self.fallback_cluster_shape_mn = fallback_cluster_shape_mn + # self.fallback_cluster_shape_mn == self.cluster_shape_mn (inherited), + # since the parent was initialised with fallback_cluster_shape_mn above. + assert self.fallback_cluster_shape_mn == self.cluster_shape_mn + + # ------------------------------------------------------------------ + # Setup helpers + # ------------------------------------------------------------------ + + def _setup_attributes(self): + """Extend base setup to also compute preferred-cluster layout and mcast info.""" + super()._setup_attributes() + + # Re-create tiled_mma (same call as in parent) to compute preferred layout. + # Attributes like self.a_dtype, self.a_major_mode, etc., are set by the + # parent's _setup_attributes caller (__call__) before _setup_attributes runs. + tiled_mma = utils.sm100.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + # Preferred cluster layout used for TMA multicast mask and pipeline sync. + self.preferred_cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.preferred_cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.num_preferred_mcast_ctas_a = cute.size( + self.preferred_cluster_layout_vmnk.shape[2] + ) + self.num_preferred_mcast_ctas_b = cute.size( + self.preferred_cluster_layout_vmnk.shape[1] + ) + self.is_preferred_a_mcast = self.num_preferred_mcast_ctas_a > 1 + self.is_preferred_b_mcast = self.num_preferred_mcast_ctas_b > 1 + + # Fallback attributes are simply those computed by the base class (which + # was initialised with fallback_cluster_shape_mn). + self.fallback_cluster_layout_vmnk = self.cluster_layout_vmnk + self.num_fallback_mcast_ctas_a = self.num_mcast_ctas_a + self.num_fallback_mcast_ctas_b = self.num_mcast_ctas_b + self.is_fallback_a_mcast = self.is_a_mcast + self.is_fallback_b_mcast = self.is_b_mcast + + # ------------------------------------------------------------------ + # Host-side JIT entry point + # ------------------------------------------------------------------ + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + max_active_preferred_clusters: cutlass.Constexpr = None, + max_active_fallback_clusters: cutlass.Constexpr = None, + stream: cuda.CUstream = None, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + """Execute preferred-cluster GEMM. + + Sets up TMA atoms for both preferred and fallback cluster shapes, computes + the launch grid, and dispatches the mega-kernel with both cluster shapes. + + :param a: Input tensor A (L, M, K) or (M, K), row- or column-major. + :type a: cute.Tensor + :param b: Input tensor B (L, K, N) or (K, N), row- or column-major. + :type b: cute.Tensor + :param c: Output tensor C (L, M, N) or (M, N), row- or column-major. + :type c: cute.Tensor + :param max_active_preferred_clusters: Maximum concurrently active preferred clusters. + :type max_active_preferred_clusters: cutlass.Constexpr + :param max_active_fallback_clusters: Maximum concurrently active fallback clusters. + :type max_active_fallback_clusters: cutlass.Constexpr + :param stream: CUDA stream. + :type stream: cuda.CUstream + :param epilogue_op: Optional elementwise epilogue applied to the accumulator. + :type epilogue_op: cutlass.Constexpr + """ + + def add_batch_mode(tensor: cute.Tensor) -> cute.Tensor: + return cute.make_tensor( + tensor.iterator, + cute.prepend(tensor.layout, cute.make_layout(1), up_to_rank=3), + ) + + # Normalise A/B/C to rank-3 by prepending a batch mode of size 1 when + # the input is 2D (non-batched). This allows the rest of the kernel to + # treat both batched (L, M, K) and non-batched (M, K) inputs uniformly + # as (1, M, K) without any special-casing downstream. + a = add_batch_mode(a) + b = add_batch_mode(b) + c = add_batch_mode(c) + + # Permute tensor modes from torch to cute convention + # A: (L, M, K) -> (M, K, L) + a = cute.make_tensor(a.iterator, cute.select(a.layout, [1, 2, 0])) + # B: (L, K, N) -> (N, K, L) + b = cute.make_tensor(b.iterator, cute.select(b.layout, [2, 1, 0])) + # C: (L, M, N) -> (M, N, L) + c = cute.make_tensor(c.iterator, cute.select(c.layout, [1, 2, 0])) + + # Capture dtype/layout info needed by _setup_attributes. + self.a_dtype: type[cutlass.Numeric] = a.element_type + self.b_dtype: type[cutlass.Numeric] = b.element_type + self.c_dtype: type[cutlass.Numeric] = c.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + + # Check if input data types are compatible with MMA instruction + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + tiled_mma = utils.sm100.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + + # ------------------------------------------------------------------ + # TMA atoms for the preferred cluster + # ------------------------------------------------------------------ + a_op_preferred = utils.sm100.cluster_shape_to_tma_atom_A( + self.preferred_cluster_shape_mn, tiled_mma.thr_id + ) + tma_atom_a_preferred, tma_tensor_a_preferred = cute.nvgpu.make_tiled_tma_atom_A( + a_op_preferred, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.preferred_cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + b_op_preferred = utils.sm100.cluster_shape_to_tma_atom_B( + self.preferred_cluster_shape_mn, tiled_mma.thr_id + ) + tma_atom_b_preferred, tma_tensor_b_preferred = cute.nvgpu.make_tiled_tma_atom_B( + b_op_preferred, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.preferred_cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + + # ------------------------------------------------------------------ + # TMA atoms for the fallback cluster + # ------------------------------------------------------------------ + a_op_fallback = utils.sm100.cluster_shape_to_tma_atom_A( + self.fallback_cluster_shape_mn, tiled_mma.thr_id + ) + tma_atom_a_fallback, tma_tensor_a_fallback = cute.nvgpu.make_tiled_tma_atom_A( + a_op_fallback, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.fallback_cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + b_op_fallback = utils.sm100.cluster_shape_to_tma_atom_B( + self.fallback_cluster_shape_mn, tiled_mma.thr_id + ) + tma_atom_b_fallback, tma_tensor_b_fallback = cute.nvgpu.make_tiled_tma_atom_B( + b_op_fallback, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.fallback_cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size + + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + # Response size is 4B * 4 elements + self.num_clc_response_bytes = 16 + + # ------------------------------------------------------------------ + # TMA atom for C output (cluster-independent) + # ------------------------------------------------------------------ + tma_atom_c = None + tma_tensor_c = None + if cutlass.const_expr(self.use_tma_store): + epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout, self.epi_tile + ) + + # ------------------------------------------------------------------ + # Grid computation for preferred and fallback clusters + # ------------------------------------------------------------------ + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + # CLC: each cluster shape gets its own scheduler params so that + # the dynamic load-balancer correctly scopes each cluster's tiles. + # The per-cluster CLC params/grid computation is identical to the + # base kernel's, so reuse the inherited helper rather than duplicating it. + self.fallback_tile_sched_params, _ = ( + self._compute_grid_clc_dynamic_persistent_scheduler( + c, + self.cta_tile_shape_mnk, + self.fallback_cluster_shape_mn, + ) + ) + self.preferred_tile_sched_params, preferred_grid = ( + self._compute_grid_clc_dynamic_persistent_scheduler( + c, + self.cta_tile_shape_mnk, + self.preferred_cluster_shape_mn, + ) + ) + else: + ( + self.preferred_tile_sched_params, + preferred_grid, + *_, + ) = self._compute_preferred_cluster_grid( + c, + self.cta_tile_shape_mnk, + self.preferred_cluster_shape_mn, + max_active_preferred_clusters, + self.fallback_cluster_shape_mn, + max_active_fallback_clusters, + ) + # Static scheduler derives tile coords from the shared (preferred-sized) + # grid, so fallback clusters MUST reuse the preferred params. + self.fallback_tile_sched_params = self.preferred_tile_sched_params + + # ------------------------------------------------------------------ + # Launch the mega-kernel with preferred cluster + fallback cluster + # ------------------------------------------------------------------ + self.kernel( + tiled_mma, + tma_atom_a_preferred, + tma_tensor_a_preferred, + tma_atom_a_fallback, + tma_tensor_a_fallback, + tma_atom_b_preferred, + tma_tensor_b_preferred, + tma_atom_b_fallback, + tma_tensor_b_fallback, + tma_atom_c, + tma_tensor_c if self.use_tma_store else c, + self.preferred_cluster_layout_vmnk, + self.fallback_cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.c_smem_layout_staged, + self.epi_tile, + self.preferred_tile_sched_params, + self.fallback_tile_sched_params, + epilogue_op, + ).launch( + grid=preferred_grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.preferred_cluster_shape_mn, 1), + fallback_cluster=(*self.fallback_cluster_shape_mn, 1), + stream=stream, + smem_merge_branch_allocs=True, + ) + + # ------------------------------------------------------------------ + # Cluster-specific device function JIT (called from kernel below) + # ------------------------------------------------------------------ + + @cute.jit + def cluster_specific_kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom | None, + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + tile_sched_params: Union[ + utils.ClcDynamicPersistentTileSchedulerParams, + utils.PersistentTileSchedulerParams, + ], + num_tma_producer: int, + is_a_mcast: bool, + is_b_mcast: bool, + epilogue_op: cutlass.Constexpr, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: cute.Layout | cute.ComposedLayout | None, + epi_tile: cute.Tile, + ): + """Cluster-specific inner kernel body. + + This function contains the full persistent GEMM loop parameterised by the + cluster layout and TMA atoms for a single cluster shape (either preferred or + fallback). It is called from :meth:`kernel` after branching on the actual + runtime cluster dimensions. + + :note: ``a_smem_layout_staged``, ``b_smem_layout_staged``, + ``c_smem_layout_staged``, and ``epi_tile`` must be passed explicitly + rather than captured via ``self.*`` because this function is called + from a ``@cute.kernel`` region which is isolated — MLIR SSA values + defined in the enclosing ``@cute.jit __call__`` scope cannot be + referenced across the isolation boundary. + + :param tiled_mma: Tiled MMA atom. + :type tiled_mma: cute.TiledMma + :param tma_atom_a: TMA copy atom for operand A (multicast or unicast). + :type tma_atom_a: cute.CopyAtom + :param mA_mkl: Global A tensor (M, K, L) exposed to TMA. + :type mA_mkl: cute.Tensor + :param tma_atom_b: TMA copy atom for operand B. + :type tma_atom_b: cute.CopyAtom + :param mB_nkl: Global B tensor (N, K, L) exposed to TMA. + :type mB_nkl: cute.Tensor + :param tma_atom_c: TMA copy atom for output C (S2G), or None. + :type tma_atom_c: cute.CopyAtom | None + :param mC_mnl: Global output C tensor (M, N, L). + :type mC_mnl: cute.Tensor + :param cluster_layout_vmnk: Cluster layout for this branch (preferred or fallback). + :type cluster_layout_vmnk: cute.Layout + :param tile_sched_params: Tile scheduler parameters (same for both branches). + :type tile_sched_params: utils.PersistentTileSchedulerParams + :param num_tma_producer: Number of TMA producer threads for AB pipeline. + :type num_tma_producer: int + :param is_a_mcast: True when A is multicasted across the cluster. + :type is_a_mcast: bool + :param is_b_mcast: True when B is multicasted across the cluster. + :type is_b_mcast: bool + :param epilogue_op: Elementwise epilogue function. + :type epilogue_op: cutlass.Constexpr + :param a_smem_layout_staged: Staged SMEM layout for A (passed from kernel param). + :type a_smem_layout_staged: cute.ComposedLayout + :param b_smem_layout_staged: Staged SMEM layout for B (passed from kernel param). + :type b_smem_layout_staged: cute.ComposedLayout + :param c_smem_layout_staged: Staged SMEM layout for C, or None. + :type c_smem_layout_staged: cute.Layout | cute.ComposedLayout | None + :param epi_tile: Epilogue tile shape (passed from kernel param). + :type epi_tile: cute.Tile + """ + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # Prefetch TMA descriptors + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + if cutlass.const_expr(self.use_tma_store): + cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + # CTA coordinates inside the cluster + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + is_first_cta_in_cluster = cta_rank_in_cluster == 0 + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + tidx, _, _ = cute.arch.thread_idx() + + # Shared memory barriers and helpers + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_ab_stage * 2 + ] + acc_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_acc_stage * 2 + ] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2] + clc_response_align_bytes = self.num_clc_response_bytes + clc_response: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int32, 4], + clc_response_align_bytes, + ] + + else: + + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_ab_stage * 2 + ] + acc_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_acc_stage * 2 + ] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # AB mainloop pipeline (TMA → SMEM → MMA) + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + + # Accumulator pipeline (MMA → TMEM → epilogue) + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilogue_warp_id) * ( + 2 if use_2cta_instrs else 1 + ) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # CLC fetch pipeline (CLC dynamic scheduler only) + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + clc_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread + ) + cluster_size = cute.size(cluster_layout_vmnk) + num_clc_consumer_threads = 32 * ( + 1 + cluster_size * (1 + len(self.epilogue_warp_id) + 1) + ) + clc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_clc_consumer_threads + ) + clc_pipeline = pipeline.PipelineClcFetchAsync.create( + barrier_storage=storage.clc_mbar_ptr.data_ptr(), + num_stages=self.num_clc_stage, + producer_group=clc_pipeline_producer_group, + consumer_group=clc_pipeline_consumer_group, + tx_count=self.num_clc_response_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem_dealloc_barrier = None + if cutlass.const_expr(not self.use_tma_store): + tmem_dealloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_dealloc_sync_bar_id, + num_threads=32 * len(self.epilogue_warp_id), + ) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, + ) + + pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) + + # CLC state (CLC dynamic scheduler only) + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + clc_response_ptr = storage.clc_response.data_ptr() + clc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_clc_stage + ) + + # SMEM tensors for A and B + # SMEM source address needs to be 128-byte aligned for TMA + sA = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sB = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + + # TMA multicast masks (needed for multicast TMA and 2-CTA MMA coordination) + a_full_mcast_mask = None + b_full_mcast_mask = None + if cutlass.const_expr(is_a_mcast or is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + # Global tensor tiles partitioned per CTA + # (bM, bK, RestM, RestK, RestL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, RestN, RestK, RestL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bN, RestM, RestN, RestL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) + tCgC = thr_mma.partition_C(gC_mnl) + + # TMA partition for A: map CTA coord → SMEM slot and global slice + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA partition for B + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # MMA register fragments + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + + # Construct the tile scheduler (shared across all warp roles) + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + tile_sched = utils.ClcDynamicPersistentTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + cute.arch.grid_dim(), + clc_response_ptr, + ) + else: + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + # ------------------------------------------------------------------ + # Warp-specialised TMA load loop + # ------------------------------------------------------------------ + if warp_idx == self.tma_warp_id: + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + tAgA_slice = tAgA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + tBgB_slice = tBgB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] + + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + handle = ab_producer.acquire_and_advance(peek_ab_empty_status) + + cute.copy( + tma_atom_a, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=b_full_mcast_mask, + ) + + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + else: + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + ab_producer.tail() + + # ------------------------------------------------------------------ + # Warp-specialised scheduler warp (CLC dynamic scheduler only) + # ------------------------------------------------------------------ + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + if warp_idx == self.sched_warp_id and is_first_cta_in_cluster: + clc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.ProducerConsumer, self.num_clc_stage + ) + + while work_tile.is_valid_tile: + clc_pipeline.producer_acquire(clc_producer_state) + mbarrier_addr = clc_pipeline.producer_get_barrier( + clc_producer_state + ) + tile_sched.advance_to_next_work(mbarrier_addr) + clc_producer_state.advance() + + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + + clc_pipeline.producer_tail(clc_producer_state) + + # ------------------------------------------------------------------ + # Warp-specialised MMA loop + # ------------------------------------------------------------------ + if warp_idx == self.mma_warp_id: + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + + ab_consumer.reset() + peek_ab_full_status = cutlass.Boolean(1) + if is_leader_cta: + peek_ab_full_status = ab_consumer.try_wait() + + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + for k_tile in range(k_tile_cnt): + if is_leader_cta: + handle = ab_consumer.wait_and_advance(peek_ab_full_status) + + num_kblocks = cute.size(tCrA, mode=[2]) + for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): + kblk_crd = (None, None, kblk_idx, handle.index) + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblk_crd], + tCrB[kblk_crd], + tCtAcc, + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + handle.release() + + peek_ab_full_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_full_status = ab_consumer.try_wait() + + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + else: + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + acc_pipeline.producer_tail(acc_producer_state) + + sC = None + if cutlass.const_expr(self.use_tma_store): + sC = smem.allocate_tensor( + element_type=self.c_dtype, + layout=c_smem_layout_staged.outer, + byte_alignment=128, + swizzle=c_smem_layout_staged.inner, + ) + + # ------------------------------------------------------------------ + # Warp-specialised epilogue loop + # ------------------------------------------------------------------ + if warp_idx < self.mma_warp_id: + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + if cutlass.const_expr(self.use_tma_store): + assert tma_atom_c is not None and sC is not None + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilogue_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, producer_group=c_producer_group + ) + + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + if cutlass.const_expr(not self.use_clc_dynamic_scheduler): + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + else: + # CLC epilogue warps don't call advance_to_next_work (only + # the sched warp does), so track the tile count locally for + # correct C-store pipeline buffer indexing. + tile_sched._num_tiles_executed += Int32(1) + + num_tiles_executed = tile_sched.num_tiles_executed + + if cutlass.const_expr(self.use_tma_store): + acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( + self, + tidx, + warp_idx, + tma_atom_c, + tCtAcc_base, + sC, + tCgC, + epi_tile, + num_tiles_executed, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + c_pipeline, + ) + else: + acc_consumer_state = utils.gemm.sm100.epilogue( + self, + tidx, + tCtAcc_base, + tCgC, + epi_tile, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + ) + + if cutlass.const_expr(self.use_clc_dynamic_scheduler): + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + + if cutlass.const_expr(self.use_tma_store): + c_pipeline.producer_tail() + else: + tmem_dealloc_barrier.arrive_and_wait() + + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # ------------------------------------------------------------------ + # GPU device kernel: routes to cluster_specific_kernel + # ------------------------------------------------------------------ + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a_preferred: cute.CopyAtom, + tma_tensor_a_preferred: cute.Tensor, + tma_atom_a_fallback: cute.CopyAtom, + tma_tensor_a_fallback: cute.Tensor, + tma_atom_b_preferred: cute.CopyAtom, + tma_tensor_b_preferred: cute.Tensor, + tma_atom_b_fallback: cute.CopyAtom, + tma_tensor_b_fallback: cute.Tensor, + tma_atom_c: cute.CopyAtom | None, + mC_mnl: cute.Tensor, + preferred_cluster_layout_vmnk: cute.Layout, + fallback_cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: cute.Layout | cute.ComposedLayout | None, + epi_tile: cute.Tile, + preferred_tile_sched_params: Union[ + utils.ClcDynamicPersistentTileSchedulerParams, + utils.PersistentTileSchedulerParams, + ], + fallback_tile_sched_params: Union[ + utils.ClcDynamicPersistentTileSchedulerParams, + utils.PersistentTileSchedulerParams, + ], + epilogue_op: cutlass.Constexpr, + ): + """Mega-kernel that dispatches to the appropriate cluster-specific body. + + At runtime, each CTA checks its actual cluster dimensions to determine + whether it belongs to a preferred or fallback cluster and calls + :meth:`cluster_specific_kernel` with the matching parameters. + + :note: ``smem_merge_branch_allocs=True`` must be set at launch so that + the shared-memory allocations from both branches are merged. + """ + cbdim_x, cbdim_y, cbdim_z = cute.arch.block_in_cluster_dim() + is_preferred_cluster = ( + cbdim_x == self.preferred_cluster_shape_mn[0] + and cbdim_y == self.preferred_cluster_shape_mn[1] + and cbdim_z == 1 + ) + + if is_preferred_cluster: + self.cluster_specific_kernel( + tiled_mma, + tma_atom_a_preferred, + tma_tensor_a_preferred, + tma_atom_b_preferred, + tma_tensor_b_preferred, + tma_atom_c, + mC_mnl, + preferred_cluster_layout_vmnk, + preferred_tile_sched_params, + self.num_preferred_mcast_ctas_a + self.num_preferred_mcast_ctas_b - 1, + self.is_preferred_a_mcast, + self.is_preferred_b_mcast, + epilogue_op, + a_smem_layout_staged, + b_smem_layout_staged, + c_smem_layout_staged, + epi_tile, + ) + else: + self.cluster_specific_kernel( + tiled_mma, + tma_atom_a_fallback, + tma_tensor_a_fallback, + tma_atom_b_fallback, + tma_tensor_b_fallback, + tma_atom_c, + mC_mnl, + fallback_cluster_layout_vmnk, + fallback_tile_sched_params, + self.num_fallback_mcast_ctas_a + self.num_fallback_mcast_ctas_b - 1, + self.is_fallback_a_mcast, + self.is_fallback_b_mcast, + epilogue_op, + a_smem_layout_staged, + b_smem_layout_staged, + c_smem_layout_staged, + epi_tile, + ) + + # ------------------------------------------------------------------ + # Grid computation for preferred cluster + # ------------------------------------------------------------------ + + @staticmethod + def _compute_preferred_cluster_grid( + c: cute.Tensor, + cta_tile_shape_mnk: tuple[int, int, int], + preferred_cluster_shape_mn: tuple[int, int], + max_active_preferred_clusters: cutlass.Constexpr, + fallback_cluster_shape_mn: tuple[int, int], + max_active_fallback_clusters: cutlass.Constexpr, + ) -> tuple[ + utils.PersistentTileSchedulerParams, + tuple[int, int, int], + utils.PersistentTileSchedulerParams, + tuple[int, int, int], + ]: + """Compute launch grids for the preferred and fallback cluster shapes. + + The preferred grid is sized so that its total CTA count does not exceed the + fallback grid's total CTA count, and is rounded up to a multiple of the + preferred cluster size. This ensures the hardware can fill the preferred + grid with a mix of preferred and fallback clusters. + + :param c: Output tensor C (M, N, L) used to determine tile count. + :type c: cute.Tensor + :param cta_tile_shape_mnk: CTA tile shape (M, N, K). + :type cta_tile_shape_mnk: tuple[int, int, int] + :param preferred_cluster_shape_mn: Preferred cluster shape (ClusterM, ClusterN). + :type preferred_cluster_shape_mn: tuple[int, int] + :param max_active_preferred_clusters: Maximum active preferred clusters. + :type max_active_preferred_clusters: cutlass.Constexpr + :param fallback_cluster_shape_mn: Fallback cluster shape (ClusterM, ClusterN). + :type fallback_cluster_shape_mn: tuple[int, int] + :param max_active_fallback_clusters: Maximum active fallback clusters. + :type max_active_fallback_clusters: cutlass.Constexpr + + :returns: (preferred_tile_sched_params, preferred_grid, + fallback_tile_sched_params, fallback_grid) + :rtype: tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int], + utils.PersistentTileSchedulerParams, tuple[int, int, int]] + """ + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + + # Compute scheduler params for preferred cluster + preferred_cluster_shape_mnl = (*preferred_cluster_shape_mn, 1) + preferred_tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, preferred_cluster_shape_mnl + ) + preferred_grid = utils.StaticPersistentTileScheduler.get_grid_shape( + preferred_tile_sched_params, max_active_preferred_clusters + ) + + # Compute scheduler params for fallback cluster + fallback_cluster_shape_mnl = (*fallback_cluster_shape_mn, 1) + fallback_tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, fallback_cluster_shape_mnl + ) + fallback_grid = utils.StaticPersistentTileScheduler.get_grid_shape( + fallback_tile_sched_params, max_active_fallback_clusters + ) + + # Round the preferred grid up so it is a multiple of the preferred cluster + preferred_grid = cute.round_up(preferred_grid, preferred_cluster_shape_mnl) + + # Cap the preferred launch grid so that its max CTA count ≤ the fallback + # grid's max CTA count (which approximates the GPU's total SM capacity). + # The preferred grid is the actual CUDA launch grid, so it cannot exceed + # the hardware's CTA capacity. At runtime, the vast majority of CTAs + # will run as preferred clusters; only the small boundary fraction that + # cannot fill a complete preferred cluster will run as fallback clusters. + # + # NOTE: when the problem is smaller than a single preferred cluster tile + # (M < preferred_cluster_m * cta_m OR N < preferred_cluster_n * cta_n), + # ``max_fallback_ctas < preferred_cluster_size`` and this floor + # division would yield 0, producing a zero-sized grid Z dim and + # triggering ``cudaErrorInvalidValue`` at launch. Such problem sizes + # must be rejected by the caller; see the host-side guard in + # ``PersistentDenseGemmPreferredClusterOperator._compile``. + preferred_cluster_size = ( + preferred_cluster_shape_mn[0] * preferred_cluster_shape_mn[1] + ) + max_fallback_ctas = fallback_grid[0] * fallback_grid[1] * fallback_grid[2] + max_preferred_clusters_count = max_fallback_ctas // preferred_cluster_size + preferred_grid = ( + preferred_grid[0], + preferred_grid[1], + max_preferred_clusters_count, + ) + + return ( + preferred_tile_sched_params, + preferred_grid, + fallback_tile_sched_params, + fallback_grid, + ) diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_static_persistent_efc_impl.py b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_static_persistent_efc_impl.py new file mode 100644 index 000000000..12e3b5a52 --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_static_persistent_efc_impl.py @@ -0,0 +1,1686 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import types +import typing + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.nvgpu import cpasync, tcgen05 + +import cutlass.operators.providers.cutedsl.evt.common_efc as common_efc +from cutlass.operators.providers.cutedsl.evt.common_efc import log + +""" +Common base infrastructure for high-performance persistent batched dense GEMM with custom epilogue fusion +for the NVIDIA Blackwell SM100 architecture using CUTE DSL and Epilogue Fusion Configuration (EFC). + +This module provides the PersistentDenseGemmEFCKernel base class that implements the core GEMM +functionality with support for custom epilogue operations. Subclasses define specific epilogue +configurations by providing an epilogue function that operates on the accumulator and supplemental +tensors. + +Key Features: + - Utilizes Tensor Memory Access (TMA) for efficient memory operations + - Utilizes Blackwell's tcgen05.mma for matrix multiply-accumulate (MMA) operations (including 2CTA MMA instructions) + - Implements TMA multicast with cluster to reduce L2 memory traffic + - Supports persistent tile scheduling to better overlap memory load/store with MMA across tiles + - Supports warp specialization to avoid explicit pipelining between mainloop load and MMA + - Uses Epilogue Fusion Configuration (EFC) to define custom epilogue operations + +GEMM Execution Flow: +1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. +2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction. +3. EPILOGUE warp (customizable via EFC): + - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. + - Load supplemental input tensors from GMEM to SMEM using TMA, then to RMEM. + - Execute custom epilogue function (defined by subclass) that: + * Accesses the accumulator via efc_config.accum() + * Reads from supplemental input tensors via tensor.load() + * Writes to supplemental output tensors via tensor.store() + * Can apply arbitrary element-wise operations, scaling, and fusion + - Store result tensors from RMEM to SMEM to GMEM with TMA operations. + +SM100 tcgen05.mma instructions operate as follows: +- Read matrix A from SMEM +- Read matrix B from SMEM +- Write accumulator to TMEM +The accumulator in TMEM must then be loaded to registers before writing back to GMEM. + +Base Tensor Dimensions: +- Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") +- Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") +- Supplemental tensors are MxNxL with layout matching the epilogue configuration + +Common Constraints: +* Supported input data types: fp16, bf16, tf32, int8, uint8, fp8 (e4m3fn, e5m2) +* A/B tensors must have the same data type +* MMA tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) +* MMA tiler N must be 32-256, step 32 +* Cluster shape M/N must be positive and power of 2, total cluster size <= 16 +* Cluster shape M must be multiple of 2 if use_2cta_instrs=True +* The contiguous dimension of all tensors must be at least 16 bytes aligned, + i.e., number of elements is a multiple of 4, 8, and 16 for TFloat32, + Float16/BFloat16, and Int8/Uint8/Float8, respectively. +* OOB tiles are not allowed when TMA store is disabled + +Subclass Examples: +- custom_epilogue_dense_gemm.py: Custom fused epilogue with multiple read/write tensors +- broadcast_custom_epilogue_dense_gemm.py: Broadcasting and mode remapping with transposed tensors +- activation_custom_epilogue_dense_gemm.py: Activation functions with Ada FP8 GEMM epilogue pattern +- synthetic_custom_epilogue_dense_gemm.py: Synthetic epilogue for testing with configurable tensor counts + +Writing an Epilogue Function (EFC User Guide) +============================================== + +The Epilogue Fusion Configuration (EFC) lets you define custom epilogue +operations as a plain Python function. The framework analyzes the +function signature, JIT-compiles it for the GPU kernel, and also +executes it on the CPU with PyTorch for automatic verification. + +Function Signature +------------------ + +The epilogue function must follow this contract: + + def epilogue(efc_config, , , ...): + ... + +- The first parameter **must** be named ``efc_config``. It is an + ``EFC.Configuration`` instance that provides access to the GEMM + accumulator, activation functions, and mode remapping. +- The remaining parameters are user-defined. Their names must match + the supplemental arguments passed to the GEMM ``compile()`` call. +- **Tensor** parameters correspond to ``cutlass.cute.Tensor`` values + and support ``.load()`` / ``.store()`` operations. +- **Scalar** parameters (``float``, ``int``, etc.) are used directly + in arithmetic expressions without ``.load()``. +- The function does **not** return a value. Results are written via + ``.store()`` on tensor parameters. At least one ``.store()`` call + is required for a useful computation. + +Accessing the Accumulator +~~~~~~~~~~~~~~~~~~~~~~~~~ + +``efc_config.accum()`` returns the GEMM accumulator (the result of +A * B). It can be called multiple times within the same epilogue:: + + def epilogue(efc_config, D): + D.store(efc_config.accum()) + +Loading and Storing Tensors +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Tensor parameters expose two operations: + +- ``tensor.load()`` — reads the tensor value (marks it as an input). +- ``tensor.store(value)`` — writes a computed value (marks it as an output). + +A tensor can be both loaded and stored (read-modify-write). When +using multiple tensors, all ``.load()`` calls should be issued before +any ``.store()`` calls:: + + def epilogue(efc_config, C, D, alpha, beta): + D.store(efc_config.accum() * alpha + C.load() * beta) + +Activation Functions +~~~~~~~~~~~~~~~~~~~~ + +The ``efc_config`` object provides built-in activation functions: + +- ``efc_config.identity(x)`` — f(x) = x +- ``efc_config.relu(x)`` — f(x) = max(0, x) +- ``efc_config.leaky_relu(x, negative_slope=0.01)`` +- ``efc_config.tanh(x)`` +- ``efc_config.sigmoid(x)`` — f(x) = 1 / (1 + exp(-x)) +- ``efc_config.silu(x)`` — f(x) = x * sigmoid(x) (Swish) +- ``efc_config.hardswish(x)`` — f(x) = x * relu6(x + 3) / 6 +- ``efc_config.gelu(x)`` + +Example:: + + def epilogue(efc_config, C, D, alpha, beta): + D.store(efc_config.relu(efc_config.accum() * alpha + C.load() * beta)) + +Element-wise Utilities +~~~~~~~~~~~~~~~~~~~~~~ + +- ``efc_config.maximum(x, y)`` — element-wise maximum +- ``efc_config.minimum(x, y)`` — element-wise minimum + +Additionally, any method called on ``efc_config`` that is not +explicitly defined is dispatched to ``cutlass.cute.()`` on GPU +and ``torch.()`` during PyTorch verification. This gives access +to functions like ``efc_config.abs(x)``, ``efc_config.exp(x)``, +``efc_config.sqrt(x)``, ``efc_config.clamp(x, min, max)``, +``efc_config.full_like(x, value)``, etc. + +Broadcasting and Mode Remapping +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +When supplemental tensors have fewer dimensions or a different mode +ordering than the output tensor, use ``tensor.remap_modes[...]`` to +broadcast or transpose them:: + + tensor.remap_modes[] + +The target shape is implicitly the output shape ``(m, n, l)`` derived +from the GEMM input matrices ``A(m, k, l)`` and ``B(n, k, l)``. + +The subscript uses Python's ``__getitem__`` syntax where each element +is either: + +- An **integer** — the index of the source mode to place at this + position. +- ``:`` — broadcast the source along this dimension (stride 0). + +The returned remapped tensor proxy supports ``.load()`` and +``.store()`` just like regular tensor parameters. + +Examples (output is (M, N, L)): + +- **Broadcast along M** — ``C`` is (N, L):: + + C.remap_modes[:, 0, 1].load() + +- **Broadcast along N** — ``X`` is (M, L):: + + X.remap_modes[0, :, 1].load() + +- **Scalar broadcast** — ``x_factor`` is a 0-d tensor:: + + x_factor.remap_modes[:, :, :].load() + +- **Transpose** — ``Y`` is (N, M, L):: + + Y.remap_modes[1, 0, 2].store(value) + +A remapped tensor can also be both loaded and stored for +read-modify-write on a transposed layout:: + + remapped_Y = Y.remap_modes[1, 0, 2] + remapped_Y.store(efc_config.accum() + remapped_Y.load()) + +Complete Examples +~~~~~~~~~~~~~~~~~ + +**Alpha-beta scaling** (simplest useful epilogue):: + + def epilogue(efc_config, C, D, alpha, beta): + D.store(efc_config.accum() * alpha + C.load() * beta) + +**Multi-tensor with activation**:: + + def epilogue(efc_config, C, D, alpha, beta, X, x_factor, Y): + Y.store(efc_config.accum()) + result = ( + efc_config.relu(efc_config.accum() * alpha + C.load() * beta) + + X.load() * x_factor + ) + D.store(result) + +**Ada FP8 GEMM pattern with dynamic activation**:: + + def epilogue(efc_config, C, Aux, alpha, beta, bias, + scale_a, scale_b, scale_c, D, leaky_relu_alpha): + aux_val = ( + (alpha * scale_a * scale_b) * efc_config.accum() + + (beta * scale_c) * C.load() + + bias + ) + Aux.store(aux_val) + D.store(efc_config.relu(aux_val)) + +**Broadcasting with mode remapping**:: + + def epilogue(efc_config, C, D, alpha, beta, X, x_factor, Y): + # Y is (N, M, L), output is (M, N, L): transpose via [1, 0, 2] + Y.remap_modes[1, 0, 2].store(efc_config.accum()) + # C is (N, L): broadcast along M via [:, 0, 1] + # X is (M, L): broadcast along N via [0, :, 1] + result = ( + efc_config.relu( + efc_config.accum() * alpha + + C.remap_modes[:, 0, 1].load() * beta + ) + + X.remap_modes[0, :, 1].load() + * x_factor.remap_modes[:, :, :].load() + ) + D.store(result) +""" + + +class PersistentDenseGemmEFCKernel: + """Base class for batched GEMM with custom epilogue fusion using EFC. + + This class provides the core infrastructure for persistent batched GEMM operations + with customizable epilogue fusion. Subclasses define specific epilogue behaviors + by providing an epilogue configuration function that describes operations on the + accumulator and supplemental tensors. + + The class handles: + - GEMM mainloop (A * B computation) + - TMA-based memory operations + - Warp specialization + - Persistent tile scheduling + - EFC (Epilogue Fusion Configuration) integration + - Tensor creation and validation + + :param acc_dtype: Data type for accumulation during computation + :type acc_dtype: type[cutlass.Numeric] + :param epi_dtype: Data type for epilogue operation + :type epi_dtype: type[cutlass.Numeric] + :param use_2cta_instrs: Whether to use CTA group 2 for advanced thread cooperation + :type use_2cta_instrs: bool + :param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N) + :type mma_tiler_mn: tuple[int, int] + :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing + :type cluster_shape_mn: tuple[int, int] + :param epilogue_function_configuration: Function defining the epilogue behavior via EFC + :type epilogue_function_configuration: Callable + + :note: Supported A/B data types: + - TFloat32 + - Float16/BFloat16 + - Int8/Uint8 + - Float8E4M3FN/Float8E5M2 + (A and B must have the same data type) + + :note: Supported accumulator data types: + - Float32 (for all floating point A/B data types) + - Float16 (only for fp16 and fp8 A/B data types) + - Int32 (only for uint8/int8 A/B data types) + + :note: Supported supplemental tensor data types (epilogue-dependent): + - Float32 (for float32 and int32 accumulator data types) + - Int32 (for float32 and int32 accumulator data types) + - Float16/BFloat16 (for fp16 and fp8 accumulator data types) + - Int8/Uint8 (for uint8/int8 accumulator data types) + - Float8E4M3FN/Float8E5M2 (for float32 accumulator data types) + + :note: Constraints: + - MMA tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) + - MMA tiler N must be 32-256, step 32 + - Cluster shape M must be multiple of 2 if use_2cta_instrs=True + - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 + """ + + def __init__( + self, + acc_dtype: type[cutlass.Numeric], + epi_dtype: type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + epilogue_function_configuration: typing.Callable, + ): + """Initializes the configuration for a Blackwell dense GEMM kernel with EFC. + + This configuration includes several key aspects: + + 1. MMA Instruction Settings (tcgen05): + - acc_dtype: Data type for MMA accumulator. + - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler. + - use_2cta_instrs: Boolean indicating if the tcgen05 MMA variant + with cta_group=2 should be used. + + 2. Cluster Shape: + - cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster. + + 3. Epilogue Configuration: + - epilogue_function_configuration: Defines custom epilogue behavior + that operates on accumulator and supplemental tensors. + + :param acc_dtype: Data type of the accumulator. + :type acc_dtype: type[cutlass.Numeric] + :param epi_dtype: Data type of the epilogue. + :type epi_dtype: type[cutlass.Numeric] + :param use_2cta_instrs: Boolean, True to use cta_group=2 MMA variant. + :type use_2cta_instrs: bool + :param mma_tiler_mn: tuple (M, N) shape of the MMA instruction. + :type mma_tiler_mn: tuple[int, int] + :param cluster_shape_mn: tuple (ClusterM, ClusterN) shape of the cluster. + :type cluster_shape_mn: tuple[int, int] + :param epilogue_function_configuration: Function defining epilogue behavior via EFC. + :type epilogue_function_configuration: Callable + """ + self.acc_dtype: type[cutlass.Numeric] = acc_dtype + self.epi_dtype: type[cutlass.Numeric] = epi_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + # K dimension is deferred in _setup_attributes + self.mma_tiler_mn = mma_tiler_mn + self.mma_tiler = (*mma_tiler_mn, 1) + self.arch = "sm_100" + + self.c_dtype = self.epi_dtype + + self.cta_group = ( + tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + self.occupancy = 1 + # Set specialized warp ids: + + # The warps responsible for computing the epilogue function and storing + # the results. + self.epilogue_warp_id = (0, 1, 2, 3) + # The warp responsible for computing the matrix multiplication. + self.mma_warp_id = 4 + # The warp responsible for loading the tensors A & B to feed the MMA. + self.tma_warp_id = 5 + # The warp responsible for loading the auxiliary tensors used in the epilogue. + self.epilogue_load_warp_id = 6 + self.threads_per_cta = 32 * len( + ( + self.mma_warp_id, + self.tma_warp_id, + *self.epilogue_warp_id, + self.epilogue_load_warp_id, + ) + ) + # Barrier ids for cta sync, epilogue sync and tmem ptr sync. + self.epilogue_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + self.tmem_dealloc_sync_bar_id = 3 + # Amount of available shared memory. + self.smem_capacity = utils.get_smem_capacity_in_bytes(self.arch) + + # Setup the EFC from the given function representing the epilogue + # configuration. + self.efc = common_efc.EFC(self, epilogue_function_configuration) + + def _create_tiled_mma(self): + """Make a tiled MMA atom with given data type, leading dimension, CTA + group and MMA tile shape. Use SMEM operand source for A. + """ + return utils.sm100.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + def _setup_attributes(self): + """Set up configurations that are dependent on GEMM inputs + + This method configures various attributes based on the input tensor properties + (data types, leading dimensions) and kernel settings: + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B + - Computing epilogue subtile + - Setting up A/B/C/D stage counts in shared memory + - Computing A/B/C/D shared memory layout + - Computing tensor memory allocation columns + """ + # Get the right tiled MMA. + self._tiled_mma = self._create_tiled_mma() + log(f"{self._tiled_mma = !s}") + + # Compute mma/cluster/tile shapes + mma_inst_shape_k = cute.size(self._tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + # Extend mma_tiler with k-dimension (MMA_M, MMA_N, MMA_K) + self.mma_tiler = ( + self.mma_tiler[0], + self.mma_tiler[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + log(f"{self.mma_tiler = !s}") + # CTA tiler with the 2CTA instruction correction. + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(self._tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + log(f"{self.cta_tile_shape_mnk = !s}") + # Compute cluster layout, V for the 2CTA instructions. + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (self._tiled_mma.thr_id.shape,), + ) + log(f"{cute.make_layout((*self.cluster_shape_mn, 1)) = !s}") + log(f"{self.cluster_layout_vmnk = !s}") + # Compute number of multicast CTAs for A/B + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + log(f"{self.num_mcast_ctas_a = }, {self.num_mcast_ctas_b = }") + log(f"{self.is_a_mcast = }, {self.is_b_mcast = }") + + # Compute epilogue (EPI_TILE_M, EPI_TILE_N) subtile of cta_tile_shape_mnk + # according to some heuristics. + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + layout_d=self.d_layout, + elem_ty_d=self.d_dtype, + layout_c=self.c_layout, + elem_ty_c=self.c_dtype, + ) + log(f"{self.epi_tile = !s}") + + # Setup A/B/C/D pipeline stage count in shared memory and ACC stage + # count in tensor memory. + self.compute_stages() + log(f"{self.num_acc_stage = }, {self.num_ab_stage = }, {self.num_c_stage = }") + # Compute A/B shared memory layout + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + self._tiled_mma, + self.mma_tiler, + self.a_dtype, + self.num_ab_stage, + ) + log(f"{self.a_smem_layout_staged = !s}") + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + self._tiled_mma, + self.mma_tiler, + self.b_dtype, + self.num_ab_stage, + ) + log(f"{self.b_smem_layout_staged = !s}") + # Get the smem_layout for the tensors used in the EFC. + self.efc.jit.smem_layout() + + # Compute the number of tensor memory allocation columns + self.compute_num_tmem_alloc_cols() + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + *supplemental_parameters, + ): + """Execute the GEMM operation in steps: + - Setup static attributes before smem/grid/tma computation + - Setup TMA load/store atoms and tensors + - Compute grid size with regard to hardware constraints + - Define shared storage for kernel + - Launch the kernel synchronously + + :param a: Input tensor A in (L, M, K) or (M, K) layout. + :type a: cute.Tensor + :param b: Input tensor B in (L, K, N) or (K, N) layout. + :type b: cute.Tensor + :param max_active_clusters: Maximum number of active clusters. + :type max_active_clusters: cutlass.Constexpr + :param stream: CUDA stream for asynchronous execution. + :type stream: cuda.CUstream + :param supplemental_parameters: Variadic epilogue parameters (tensors and scalars). + :raises TypeError: If input data types are incompatible with the MMA instruction. + :raises AssertionError: If OOB (Out-Of-Bounds) tiles are present when TMA store is disabled. + + """ + + def add_batch_mode(tensor: cute.Tensor) -> cute.Tensor: + return cute.make_tensor( + tensor.iterator, + cute.prepend(tensor.layout, cute.make_layout(1), up_to_rank=3), + ) + + a = add_batch_mode(a) + b = add_batch_mode(b) + + # Permute tensor modes from torch to cute convention + # A: (L, M, K) -> (M, K, L) + a = cute.make_tensor(a.iterator, cute.select(a.layout, [1, 2, 0])) + # B: (L, K, N) -> (N, K, L) + b = cute.make_tensor(b.iterator, cute.select(b.layout, [2, 1, 0])) + + # Identify parameters that are remap_modes sources (broadcast tensors). + # These skip add_batch_mode + permute entirely — they are passed as-is + # to the EFC and remap_modes inside the kernel handles all the layout + # work (broadcast expansion, dimension reordering). + remap_sources = { + attrs.mapped_source + for attrs in self.efc.parameter_attributes.values() + if attrs.mapped_source + } + + # Add batch mode and permute epilogue parameters, except for + # broadcast remap sources which are passed untouched. + param_names = self.efc.epilogue_parameter_names + supplemental_parameters = ( + # Broadcast source: pass through untouched. + t + if name in remap_sources + # Regular tensor: (L, M, N) -> rank 3 -> (M, N, L). + else cute.make_tensor( + add_batch_mode(t).iterator, + cute.select(add_batch_mode(t).layout, [1, 2, 0]), + ) + if isinstance(t, cute.Tensor) + # Scalar: pass through. + else t + for name, t in zip(param_names, supplemental_parameters) + ) + + # Process the variadic parameters. + self.efc.jit.unpack_parameters(supplemental_parameters) + + # - Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") + # - Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") + # - Supplemental tensors are MxNxL with layout matching the epilogue configuration + + # The output shape (m, n, l) is derived from A(m, k, l) and B(n, k, l). + self.efc.output_shape = (a.shape[0], b.shape[0], a.shape[2]) + self.efc.jit.handle_remapping() + + # Setup static attributes before smem/grid/tma computation + self.a_dtype: type[cutlass.Numeric] = a.element_type + self.b_dtype: type[cutlass.Numeric] = b.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + + # Gather all the auxiliary tensor element data types. + self.efc.jit.record_tensor_dtypes() + + # There is no D tensor to be used as a returned tensor. In the + # following, D is used more like a "store" concept. So use the + # written tensor with the biggest element_type to set up all the tiling + # heuristics and epilogue store pipeline. + self.d_name_bigger = self.efc.jit.written_tensor_name_with_bigger_element_type() + d = self.efc.jit.parameter[self.d_name_bigger] + self.d_dtype: type[cutlass.Numeric] = d.element_type + self.d_layout = utils.LayoutEnum.from_tensor(d) + log(f"d{self.d_name_bigger} = {d!s}") + + # C is the read tensor with the biggest element_type, if any, used by + # some heuristics for tiling. + self.c_dtype = None + self.c_layout = None + self.c_name_bigger = self.efc.jit.read_tensor_name_with_bigger_element_type() + if cutlass.const_expr(self.c_name_bigger): + # The tensor with the biggest data type might be a broadcast vector. + c = self.efc.jit.get_remapped_tensor_or_itself(self.c_name_bigger) + log(f"{self.c_name_bigger = } -> {c = !s}") + self.c_dtype = c.element_type + # If only broadcast tensors are read, there is no leading dimension + # defined, so use the same as the output. + try: + self.c_layout = utils.LayoutEnum.from_tensor(c) + except ValueError: + self.c_layout = self.d_layout + + # Check if input data types are compatible with MMA instruction + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Types must match: {self.a_dtype} != {self.b_dtype}") + + # Setup attributes that depend on gemm inputs + self._setup_attributes() + + atom_thr_size = cute.size(self._tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = utils.sm100.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, self._tiled_mma.thr_id + ) + log(f"{a_op = !s}") + # Get rid of the pipeline dimension. + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + log(f"{a_smem_layout = !s}") + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + self._tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + log(f"{tma_atom_a = !s}") + log(f"{tma_tensor_a = !s}") + # Setup TMA load for B + b_op = utils.sm100.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, self._tiled_mma.thr_id + ) + log(f"{b_op = !s}") + # Get rid of the pipeline dimension. + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + log(f"{b_smem_layout = !s}") + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + self._tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + log(f"{tma_atom_b = !s}") + log(f"{tma_tensor_b = !s}") + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size + log(f"{self.num_tma_load_bytes = }") + + # Set the TMA related arguments for the tensors used in the EFC. + self.efc.jit.create_tma_arguments() + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + d, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters + ) + + self.efc.jit.create_supplemental_arguments_for_kernel() + + # Launch the kernel synchronously + self.kernel( + self._tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + self.efc.kernel.pack_arguments(), + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + ) + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + supplemental_parameters: tuple, + ): + """GPU device kernel performing the Persistent batched GEMM computation. + """ + # Process the variadic parameters. + self.efc.kernel.unpack_parameters(supplemental_parameters) + + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # + # Prefetch TMA descriptors + # + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + # Prefetch the TMA descriptors for all the supplemental tensors. + self.efc.kernel.prefetch_tma_descriptors() + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + # Coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + # + # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier + # + self.buffer_align_bytes = 1024 + + # Define shared storage for kernel + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] + ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] + acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + # Barriers used by the supplemental load tensor pipeline. + c_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_c_stage] + c_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_c_stage] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + # (MMA, MMA_M, MMA_K, STAGE) + sA: cute.struct.Align[ + cute.struct.MemRange[ + self.a_dtype, cute.cosize(a_smem_layout_staged.outer) + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + sB: cute.struct.Align[ + cute.struct.MemRange[ + self.b_dtype, cute.cosize(b_smem_layout_staged.outer) + ], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + self.smem = utils.SmemAllocator() + storage = self.smem.allocate(self.shared_storage) + + # Allocate the shared memory for all the supplemental tensors. + self.efc.kernel.allocate_smem() + + # Initialize mainloop ab_pipeline (barrier) and states + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilogue_warp_id) * ( + 2 if self.use_2cta_instrs else 1 + ) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Load pipeline, used to load all the supplemental tensors of the + # epilogue. + c_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + c_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.epilogue_warp_id), + ) + c_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.c_full_mbar_ptr.data_ptr(), + num_stages=self.num_c_stage, + producer_group=c_producer_group, + consumer_group=c_consumer_group, + # Unlock the barrier when all the tensor bytes have been loaded. + tx_count=self.efc.jit.total_tma_load_bytes, + defer_sync=True, + ) + + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=self.use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, + ) + + # Cluster arrive after barrier init + pipeline.pipeline_init_arrive( + cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True + ) + + # + # Setup smem tensor A/B + # + # (MMA, MMA_M, MMA_K, STAGE) + sA = storage.sA.get_tensor( + a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner + ) + # (MMA, MMA_N, MMA_K, STAGE) + sB = storage.sB.get_tensor( + b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner + ) + + # + # Compute multicast mask for A/B buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + if cutlass.const_expr( + self.is_a_mcast or self.is_b_mcast or self.use_2cta_instrs + ): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + # + # Local_tile partition global tensors + # + # (bM, bK, loopM, loopK, loopL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, loopN, loopK, loopL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # + # Partition global tensor for TiledMMA_A/B/D + # + self.thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + # (MMA, MMA_M, MMA_K, loopM, loopK, loopL) + tCgA = self.thr_mma.partition_A(gA_mkl) + log(f"{tCgA = !s}") + # (MMA, MMA_N, MMA_K, loopN, loopK, loopL) + tCgB = self.thr_mma.partition_B(gB_nkl) + log(f"{tCgB = !s}") + # Create the local_tile gX_mnl for all the EFC supplemental tensors. + self.efc.kernel.partition_global_tensors_for_tiled_mma() + + # + # Partition global/shared tensor for TMA load A/B + # + # TMA load A partition_S/D + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), tiles_m, tiles_k, tiles_l) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA load B partition_S/D + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), tiles_n, tiles_k, tiles_l) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C/D + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + log(f"{tCrA = !s}") + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + log(f"{tCrB = !s}") + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + log(f"{acc_shape = !s}") + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + log(f"{tCtAcc_fake = !s}") + + # Named barriers + # + epilogue_sync_barrier = pipeline.NamedBarrier( + self.epilogue_sync_bar_id, 32 * len(self.epilogue_warp_id) + ) + + # + # Cluster wait before tensor memory alloc + # + pipeline.pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # + # Specialized TMA load warp + # + + if warp_idx == self.tma_warp_id: + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), loopK) + tAgA_slice = tAgA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # ((atom_v, rest_v), loopK) + tBgB_slice = tBgB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + # + # Tma load loop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Conditionally wait for AB buffer empty + handle = ab_producer.acquire_and_advance(peek_ab_empty_status) + + # TMA load A/B + cute.copy( + tma_atom_a, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=b_full_mcast_mask, + ) + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait A/B buffer empty + # + ab_producer.tail() + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + # + # Bar sync to retrieve tensor memory ptr from shared mem + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # Set tensor memory buffer for current tile + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + + # Peek (try_wait) AB buffer full for k_tile = 0 + ab_consumer.reset() + peek_ab_full_status = cutlass.Boolean(1) + if is_leader_cta: + peek_ab_full_status = ab_consumer.try_wait() + + # + # Wait for accumulator buffer empty + # + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + + # + # Reset the ACCUMULATE field for each tile + # + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + # + # Mma mainloop + # + for k_tile in range(k_tile_cnt): + if is_leader_cta: + # Conditionally wait for AB buffer full + handle = ab_consumer.wait_and_advance(peek_ab_full_status) + + # tCtAcc += tCrA * tCrB + tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + tile_crd = (None, None, None, handle.index) + cute.gemm( + tiled_mma, tCtAcc, tCrA[tile_crd], tCrB[tile_crd], tCtAcc + ) + + # Async arrive AB buffer empty + handle.release() + + # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 + peek_ab_full_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_full_status = ab_consumer.try_wait() + + # + # Async arrive accumulator buffer full + # + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait for accumulator buffer empty + # + acc_pipeline.producer_tail(acc_producer_state) + # + # Specialized epilogue warps + # + if warp_idx < self.mma_warp_id: + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Bar sync to retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + log(f"tmem_ptr = {tmem_ptr!s}") + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + log(f"tCtAcc_base = {tCtAcc_base!s}") + # + # Partition for epilogue + # + epi_tidx = tidx + tCgD = self.efc.kernel.tCgD_written[self.d_name_bigger] + log(f"tCgD (aka tCgD_written[{self.d_name_bigger}])= {tCgD!s}") + + ( + tiled_copy_t2r, # (EPI_TILE_M, EPI_TILE_N) + tTR_tAcc_base, # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) + tTR_rAcc, # (T2R, T2R_M, T2R_N) + ) = self.epilogue_tmem_copy_and_partition( + epi_tidx, tCtAcc_base, tCgD, epi_tile + ) + log(f"{tiled_copy_t2r = !s}") + log(f"{tTR_tAcc_base = !s}") + log(f"{tTR_rAcc = !s}") + # Copy and partition for the supplemental EFC tensors. + self.efc.kernel.copy_and_partition_supplemental_rmem_tensors( + tiled_copy_t2r, tTR_rAcc, epi_tidx, epi_tile + ) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + # Store D pipeline used for all the written tensors in the epilogue. + d_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilogue_warp_id), + ) + d_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_d_stage, + producer_group=d_producer_group, + ) + + c_pipeline_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_c_stage + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # Slice the supplemental written tensors per MMA tile index. + self.efc.kernel.slice_written_tensors_per_mma_tile_index( + mma_tile_coord_mnl + ) + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + log(f"tTR_tAcc = {tTR_tAcc!s}") + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + # Group together the EPI_M, EPI_M which are starting at group 3. + # (T2R, T2R_M, T2R_N, (EPI_M, EPI_M)) + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + log(f"group_modes tTR_tAcc = {tTR_tAcc!s}") + # + # Store accumulator to global memory in subtiles + # + # Use EPI_M*EPI_M to iterate using the 1-D coordinate. + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt + for subtile_idx in cutlass.range(subtile_cnt): + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + log(f"cute.copy tiled_copy_t2r = {tiled_copy_t2r!s}") + log(f"cute.copy tTR_tAcc_mn = {tTR_tAcc_mn!s}") + log(f"cute.copy tTR_rAcc = {tTR_rAcc!s}") + + # Wait for the EFC tensor loads to complete. The wait is + # blocking even if the tx_count is 0 when there is no tensor + # to load. So, only wait if there is something to read, + # i.e. smem_read is not empty. + if cutlass.const_expr(self.efc.kernel.smem_read): + c_pipeline.consumer_wait(c_pipeline_consumer_state) + + # Load supplemental tensors from shared memory to register. + self.efc.kernel.load_tensors_from_smem_to_register( + c_pipeline_consumer_state.index + ) + + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + c_pipeline.consumer_release(c_pipeline_consumer_state) + + # Advance pipeline states + c_pipeline_consumer_state.advance() + + # + # Perform epilogue op on accumulator. + # + tiled_copy_r2s = self.efc.kernel.tiled_copy_r2s[self.d_name_bigger] + log(f"tiled_copy_r2s = {tiled_copy_r2s!s}") + # Use a SimpleNamespace to pass easily some local content as + # an extensible class compatible with CuTe DSL + # implementation. + epilogue_context = types.SimpleNamespace() + # Load the accumulator cast to the epi_dtype used to do all + # the computations in the epilogue. + # Retile the accumulator subtile to fit the destination + # subtile vector TV layout. + epilogue_context.acc_vec = ( + tiled_copy_r2s.retile(tTR_rAcc).load().to(self.epi_dtype) + ) + log(f"before .retile tTR_rAcc = {tTR_rAcc!s}") + log( + f"tiled_copy_r2s.retile(tTR_rAcc) = {tiled_copy_r2s.retile(tTR_rAcc)!s}" + ) + log( + f"tiled_copy_r2s.retile(tTR_rAcc).load() = {tiled_copy_r2s.retile(tTR_rAcc).load()!s}" + ) + log(f"epilogue_context.acc_vec = {epilogue_context.acc_vec!s}") + + # Execute the EFC epilogue. + self.efc.kernel.epilogue_computation(epilogue_context) + d_buffer = (num_prev_subtiles + subtile_idx) % self.num_d_stage + + # Store the EFC written tensors to shared memory. + self.efc.kernel.store_written_tensors_to_smem(d_buffer) + + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + epilogue_sync_barrier.arrive_and_wait() + + # + # TMA store D to global memory + # + if warp_idx == self.epilogue_warp_id[0]: + # Store with TMA the written EFC tensors to global memory. + self.efc.kernel.tma_store_written_tensors_to_gmem( + d_buffer, subtile_idx + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + d_pipeline.producer_commit() + d_pipeline.producer_acquire() + + epilogue_sync_barrier.arrive_and_wait() + + # + # Async arrive accumulator buffer empty + # + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + epilogue_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + # + # Wait for D store complete + # + d_pipeline.producer_tail() + + # + # Specialized epilogue load warp + # + if warp_idx == self.epilogue_load_warp_id: + # Create the tiled tensors to be loaded in the epilogue. + self.efc.kernel.create_epilogue_subtile_tensors(tidx, epi_tile) + + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + c_pipeline_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_c_stage + ) + + # Setup the pipelines reading the EFC supplemental tensors. + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + # Prepare the EFC tensors to be loaded by the subtiles. + subtile_cnt = self.efc.kernel.prepare_tensor_load_for_subtiles( + mma_tile_coord_mnl, + ) + + # Assume the pipeline can work even in the case there is no + # tensor to load and so subtile_cnt is 0. + for subtile_idx in cutlass.range(subtile_cnt): + # Load C from global memory to shared memory. + c_pipeline.producer_acquire(c_pipeline_producer_state) + + # Load the subtiles of EFC tensors. + self.efc.kernel.load_tensor_subtiles( + subtile_idx, c_pipeline, c_pipeline_producer_state + ) + + c_pipeline_producer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait for the load buffer to be empty. + # + c_pipeline.producer_tail(c_pipeline_producer_state) + + def epilogue_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tAcc: cute.Tensor, + tCgC: cute.Tensor, + epi_tile: cute.Tile, + ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). + """ + # Make tiledCopy for tensor memory load + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.d_layout, # Take this as the reference layout for the epilogue tile. + self.epi_dtype, # But we get the accumulator as epi_dtype in the epilogue. + self.acc_dtype, + epi_tile, + self.use_2cta_instrs, + ) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) + tAcc_epi = cute.flat_divide( + tAcc[((None, None), 0, 0, None)], + epi_tile, + ) + # (EPI_TILE_M, EPI_TILE_N) + tiled_copy_t2r = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)] + ) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + tCgC_epi = cute.flat_divide( + tCgC[((None, None), 0, 0, None, None, None)], epi_tile + ) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + tTR_gC = thr_copy_t2r.partition_D(tCgC_epi) + # (T2R, T2R_M, T2R_N) + tTR_rAcc = cute.make_rmem_tensor( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype + ) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc + + def epilogue_smem_copy_and_partition_load( + self, + tiled_copy_t2r: cute.TiledCopy, + tTR_rC: cute.Tensor, + tidx: cutlass.Int32, + sC: cute.Tensor, + ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """Make tiledCopy for shared memory load, then use it to partition register array (destination) and shared memory (source). + + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + :type tiled_copy_t2r: cute.TiledCopy + :param tTR_rC: The partitioned accumulator tensor + :type tTR_rC: cute.Tensor + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + + :return: A tuple containing (tiled_copy_s2r, tSR_rC, tSR_sC) where: + - tiled_copy_s2r: The tiled copy operation for smem to register copy(s2r) + - tSR_rC: The partitioned tensor C (register destination) + - tSR_sC: The partitioned tensor C (smem source) + :rtype: tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + copy_atom_s2r = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) + tiled_copy_s2r = cute.make_tiled_copy_D(copy_atom_s2r, tiled_copy_t2r) + # (S2R, S2R_M, S2R_N, PIPE_C) + thr_copy_s2r = tiled_copy_s2r.get_slice(tidx) + tSR_sC = thr_copy_s2r.partition_D(sC) + # (S2R, S2R_M, S2R_N) + tSR_rC = tiled_copy_s2r.retile(tTR_rC) + return tiled_copy_s2r, tSR_rC, tSR_sC + + def epilogue_gmem_copy_and_partition( + self, + tidx: cutlass.Int32, + atom: cute.CopyAtom | cute.TiledCopy, + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + sC: cute.Tensor, + dtype: type[cutlass.Numeric], + ) -> tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: + """Make tiledCopy for global memory store, then use it to: + - partition register array (source) and global memory (destination) for non-TMA store version; + - partition shared memory (source) and global memory (destination) for TMA store version. + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param atom: The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for non-TMA store version + :type atom: cute.CopyAtom or cute.TiledCopy + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + + :return: A tuple containing either: + - For TMA store: (tma_atom_c, bSG_sC, bSG_gC) where: + - tma_atom_c: The TMA copy atom + - bSG_sC: The partitioned shared memory tensor C + - bSG_gC: The partitioned global tensor C + - For non-TMA store: (simt_atom, tTR_rC, tTR_gC) where: + - simt_atom: The SIMT copy atom + - tTR_rC: The register tensor C + - tTR_gC: The partitioned global tensor C + :rtype: tuple[cute.CopyAtom, cute.Tensor, cute.Tensor] + """ + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, tiles_m, tiles_n, tiles_l) + gC_epi = cute.flat_divide( + gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile + ) + + tma_atom_c = atom + sC_for_tma_partition = cute.group_modes(sC, 0, 2) + gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2) + # ((ATOM_V, REST_V), EPI_M, EPI_N) + # ((ATOM_V, REST_V), EPI_M, EPI_N, tiles_m, tiles_n, tiles_l) + bSG_sC, bSG_gC = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + sC_for_tma_partition, + gC_for_tma_partition, + ) + return tma_atom_c, bSG_sC, bSG_gC + + def compute_stages(self) -> None: + """Compute and set the number of stages for A/B/C/D operands. + + Uses instance attributes to compute and assign: + `self.num_acc_stage`, `self.num_ab_stage`, `self.num_c_stage`, + and `self.num_d_stage`. + """ + # Defaults + self.num_acc_stage = 2 + # To read the tensors needed for the epilogue: + self.num_c_stage = 2 + # To write the tensors produced by the epilogue: + self.num_d_stage = 2 + + # Calculate smem layout and size for one stage of A, B, C, and D + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + self._tiled_mma, self.mma_tiler, self.a_dtype, 1 + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + self._tiled_mma, self.mma_tiler, self.b_dtype, 1 + ) + + # Get the contribution from the tensors used in the EFC. + self.efc.jit.compute_stage() + + ab_bytes_per_stage = cute.size_in_bytes( + self.a_dtype, a_smem_layout_stage_one + ) + cute.size_in_bytes(self.b_dtype, b_smem_layout_staged_one) + mbar_helpers_bytes = 1024 + # Contribution from the tensors loaded in the EFC. + c_bytes_per_stage = self.efc.jit.smem_size_in_bytes_of_read_tensors() + c_bytes = c_bytes_per_stage * self.num_c_stage + # Contribution from the tensors stored in the EFC. There is at least 1 + # written tensor, so the following is strictly positive. + d_bytes_per_stage = self.efc.jit.smem_size_in_bytes_of_written_tensors() + d_bytes = d_bytes_per_stage * self.num_d_stage + + # Calculate A/B stages + self.num_ab_stage = ( + self.smem_capacity // self.occupancy + - (mbar_helpers_bytes + c_bytes + d_bytes) + ) // ab_bytes_per_stage + log(f"{common_efc.TAB}{self.num_ab_stage = }") + + if self.num_ab_stage <= 0: + raise MemoryError("Not enough smem capacity to allocate all the tensors.") + + # Refine epilogue stages: + # Calculate remaining smem after allocating for A/B stages and reserved bytes. + # Add remaining unused smem to epilogue. + self.num_d_stage += ( + self.smem_capacity + - self.occupancy * ab_bytes_per_stage * self.num_ab_stage + - self.occupancy * (mbar_helpers_bytes + c_bytes + d_bytes) + ) // (self.occupancy * d_bytes_per_stage) + log(f"{common_efc.TAB}new {self.num_d_stage = }") + + @staticmethod + def _compute_grid( + d: cute.Tensor, + cta_tile_shape_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + max_active_clusters: cutlass.Constexpr, + ) -> tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]]: + """Use persistent tile scheduler to compute the grid size for the output tensor D. + + :param d: The output tensor D + :type d: cute.Tensor + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + :param max_active_clusters: Maximum number of active clusters. + :type max_active_clusters: cutlass.Constexpr + + :return: A tuple containing: + - tile_sched_params: Parameters for the persistent tile scheduler. + - grid: Grid shape for kernel launch. + :rtype: tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]] + """ + log(f"compute_grid: {max_active_clusters = }") + d_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gd = cute.zipped_divide(d, tiler=d_shape) + num_ctas_mnl = gd[(0, (None, None, None))].shape + common_efc.if_debug( + lambda: cute.printf("compute_grid: num_ctas_mnl = {}", num_ctas_mnl) + ) + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + common_efc.if_debug(lambda: cute.printf("compute_grid: grid = {}", grid)) + return tile_sched_params, grid + + def compute_num_tmem_alloc_cols(self) -> None: + """Compute and set the number of tensor memory allocation columns. + + This method uses the instance attributes computed during setup to + determine the number of tensor memory allocation columns and stores + the result in `self.num_tmem_alloc_cols`. + """ + acc_shape = self._tiled_mma.partition_shape_C(self.mma_tiler[:2]) + log(f"compute_num_tmem_alloc_cols: {acc_shape = !s}") + tCtAcc_fake = self._tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + log(f"compute_num_tmem_alloc_cols: {tCtAcc_fake = !s}") + self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake) + log(f"compute_num_tmem_alloc_cols: {self.num_tmem_alloc_cols = }") diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_tgv_gemm_impl.py b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_tgv_gemm_impl.py new file mode 100644 index 000000000..f211a6c61 --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm100_tgv_gemm_impl.py @@ -0,0 +1,1085 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +"""Low-latency Blackwell GEMM kernel implementation (TGV). + +CUTE DSL translation of the TGV GEMM kernel (tgv_gemm.cuh). + +This is a low-latency Blackwell GEMM kernel: C = A * B +- A (M, K, L) with K contiguous (UmmaMajor::K) +- B (N, K, L) with K contiguous +- C (M, N, L) with M contiguous + +Features: +- TMA loads for A and B matrices (GMEM -> SMEM) +- tcgen05.mma for matrix multiply-accumulate in TMEM +- Multi-stage pipeline for overlapping TMA and MMA +- Non-TMA-store epilogue (TMEM -> RMEM -> GMEM with direct store) +- 1 SM mode (no 2CTA instructions) +- 1x1 cluster (no multicast) +- PDL (Programmatic Dependent Launch) support via griddepcontrol + +Default config: CTA_M=64, CTA_N=8, CTA_K=128, DMA_Stage=8 + TypeA=bf16, TypeB=bf16, TypeC=bf16, AccType=float + UmmaMajorA=Major::K, UmmaMajorB=Major::K + +* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, + i.e, number of elements is a multiple of 4, 8, and 16 for TFloat32, + Float16/BFloat16, and Int8/Uint8/Float8, respectively. +""" + +from typing import NamedTuple + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.nvgpu import cpasync, tcgen05 + + +class WorkTileInfo(NamedTuple): + """Which output tile this CTA processes. Matches C++ WorkTileInfo struct. + For non-persistent static scheduler, CTA id is the work tile info. + """ + + M_idx: cutlass.Int32 + N_idx: cutlass.Int32 + L_idx: cutlass.Int32 + K_idx_start: cutlass.Int32 # kblock range [K_idx_start, K_idx_end) + K_idx_end: cutlass.Int32 + + +class TgvGemmKernel: + """Low-latency Blackwell GEMM kernel, literal translation from tgv_gemm.cuh. + + Uses raw mbarriers for synchronization with 8-warp specialization: + warp 0 = DMA_A, warp 1 = DMA_B, warp 2 = MMA, warps 4-7 = EPILOG. + """ + + def __init__( + self, + acc_dtype: type[cutlass.Numeric] = cutlass.Float32, + cta_m: int = 64, + cta_n: int = 8, + cta_k: int = 128, + num_ab_stage: int = 8, + use_pdl: bool = False, + pdl_count: int = -1, + ): + self.acc_dtype = acc_dtype + self.cta_m = cta_m + self.cta_n = cta_n + self.cta_k = cta_k + self.num_ab_stage = num_ab_stage + self.use_pdl = use_pdl + self.pdl_count = pdl_count + + # Fixed configuration matching C++ kernel + self.threads_per_cta = 256 # 8 warps + self.use_2cta_instrs = False # 1 SM mode + self.cluster_shape_mn = (1, 1) # No multicast, 1x1 cluster + self.cta_group = tcgen05.CtaGroup.ONE + + def _setup_attributes(self): + """Set up derived config. Corresponds to C++ gemm_host() setup section.""" + mma_tiler_mn = (self.cta_m, self.cta_n) + # tiled_mma just contains a single mma atom/instruction, not stacking anything + # DSL equivalent of C++ sm100_make_1sm_trivial_tiled_mma: + # TiledMMA tiled_mma = sm100_make_1sm_trivial_tiled_mma, Shape<1,1,1>, + # UmmaMajorA, UmmaMajorB>(); + # For 1SM mode, tiled_mma.thr_id has shape (_1,), meaning 1 "thread" in the MMA partition + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + mma_tiler_mn, + ) + + # Derive mma_inst_tile_k = NumMma_K = CTA_K / Mma_K + # For bf16: Mma_K = 16, so CTA_K=128 -> mma_inst_tile_k = 8 + # This is the number of MMA instructions along the K dimension per CTA_K tile + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + self.mma_inst_tile_k = self.cta_k // mma_inst_shape_k + + # Maps a CTA's linear rank within its cluster to 4D (V, M, N, K) coordinates. + # Used by TMA partition to determine which CTAs share data for multicast: + # V = CTA index within one tiled_mma (2SM: V=2, 1SM: V=1) + # M, N = CTA position along cluster M/N dimensions + # K = always 1 (no cluster tiling along K) + # A is multicast along N (CTAs with same M share A tiles) → + # tma_partition uses the N-slice and block_in_cluster_coord_vmnk[2] + # B is multicast along M (CTAs with same N share B tiles) → + # tma_partition uses the M-slice and block_in_cluster_coord_vmnk[1] + # For 1SM + 1x1 cluster: shape is ((_1,), _1, _1, _1), all trivial. + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + + # Create SMEM layouts for A and B with swizzle + # C++ equivalent: UMMA::tile_to_mma_shape(SmemLayoutAtom, mma_shape) + # For bf16 K-major: swizzle atom is Sw<3,4,3> with shape (M8, K64) = 8x128B + # The atom is stacked/replicated to fill ((Mma_M, Mma_K), NumMma_M, NumMma_K, DMA_Stage) + # sA_layout: ((Mma_M, Mma_K), NumMma_M, NumMma_K, DMA_Stage) + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + (self.cta_m, self.cta_n, self.cta_k), + self.a_dtype, + self.num_ab_stage, + ) + # sB_layout: ((Mma_N, Mma_K), NumMma_N, NumMma_K, DMA_Stage) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + (self.cta_m, self.cta_n, self.cta_k), + self.b_dtype, + self.num_ab_stage, + ) + + # Accumulator shape for TMEM allocation size calculation + # acc_shape: ((Mma_M, Mma_N), NumMma_M, NumMma_N) + acc_shape = tiled_mma.partition_shape_C((self.cta_m, self.cta_n)) + tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape) + self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake) + + @cute.jit + def __call__( + self, + a: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K or M contiguous + b: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K or N contiguous + c: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M contiguous + stream: cuda.CUstream, + ): + """Host-side setup and kernel launch. Corresponds to C++ gemm_host().""" + self.a_dtype = a.element_type + self.b_dtype = b.element_type + self.c_dtype = c.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + + self._setup_attributes() + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + (self.cta_m, self.cta_n), + ) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) # 1 for 1SM + + # Setup TMA load for A + # C++: make_tma_atom(SM90_TMA_LOAD{}, mA, sA_layout(_,_,_,0), make_shape(CTA_M, CTA_K)) + # tma_tensor_a is the TMA coordinate tensor: (Gemm_M, Gemm_K, Gemm_L) + a_op = sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + # a_smem_layout: ((Mma_M, Mma_K), NumMma_M, NumMma_K) — 1 stage slice + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + (self.cta_m, self.cta_n, self.cta_k), + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # Setup TMA load for B + # C++: make_tma_atom(SM90_TMA_LOAD{}, mB, sB_layout(_,_,_,0), make_shape(CTA_N, CTA_K)) + # tma_tensor_b is the TMA coordinate tensor: (Gemm_N, Gemm_K, Gemm_L) + b_op = sm100_utils.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, tiled_mma.thr_id + ) + # b_smem_layout: ((Mma_N, Mma_K), NumMma_N, NumMma_K) — 1 stage slice + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + (self.cta_m, self.cta_n, self.cta_k), + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # Calculate total bytes that TMA will transfer per stage to track completion + # C++: int tma_transaction_bytes = sizeof(make_tensor_like(tAsA(_,0))) + # A and B have separate barriers, so we track bytes separately + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + self.tma_bytes_a = a_copy_size * atom_thr_size + self.tma_bytes_b = b_copy_size * atom_thr_size + + # Grid: dim3{ceil_div(M, CTA_M), ceil_div(N, CTA_N), L} + # Each CTA processes one (CTA_M, CTA_N) output tile + grid = ( + cute.ceil_div(c.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n), + c.layout.shape[2], + ) + + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, # tma_tensor_a is mA_mkl: (Gemm_M, Gemm_K, Gemm_L) + tma_atom_b, + tma_tensor_b, # tma_tensor_b is mB_nkl: (Gemm_N, Gemm_K, Gemm_L) + c, # mC_mnl: (Gemm_M, Gemm_N, Gemm_L) + self.cluster_layout_vmnk, + self.a_smem_layout_staged, # ((Mma_M, Mma_K), NumMma_M, NumMma_K, DMA_Stage) + self.b_smem_layout_staged, # ((Mma_N, Mma_K), NumMma_N, NumMma_K, DMA_Stage) + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], # 256 threads = 8 warps + cluster=(*self.cluster_shape_mn, 1), # 1x1 cluster, no multicast + stream=stream, + use_pdl=self.use_pdl, + ) + return + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L) — TMA coordinate tensor for A + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L) — TMA coordinate tensor for B + mC_mnl: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L) — output tensor in GMEM + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, # ((Mma_M, Mma_K), NumMma_M, NumMma_K, DMA_Stage) + b_smem_layout_staged: cute.ComposedLayout, # ((Mma_N, Mma_K), NumMma_N, NumMma_K, DMA_Stage) + ): + """GPU device kernel. 1:1 translation of C++ gemm_device(). + + Slim kernel: only SMEM alloc + barrier init + warp dispatch. + All tensor partitioning is done inside warp functions. + + 256 threads, 8 warps: + Warp 0: DMA_A - loads A tiles via TMA + Warp 1: DMA_B - loads B tiles via TMA + Warp 2: MMA - performs tcgen05.mma + Warp 3: unused + Warps 4-7: EPILOG - TMEM->RMEM->GMEM store + """ + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, bidz = cute.arch.block_idx() + + # Prefetch TMA descriptors (warp 0 only) + if warp_idx == 0: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + + # ============================================================ + # WorkTileInfo (C++: non-persistent static scheduler, CTA id = work tile) + # ============================================================ + # mma_tile_coord_v: which CTA within the tiled_mma (always 0 for 1SM mode) + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + work_tile_info = WorkTileInfo( + M_idx=bidx // cute.size(tiled_mma.thr_id.shape), + N_idx=bidy, + L_idx=bidz, + K_idx_start=cutlass.Int32(0), + K_idx_end=cute.ceil_div(cute.size(mA_mkl, mode=[1]), self.cta_k), + ) + k_tile_count = work_tile_info.K_idx_end - work_tile_info.K_idx_start + + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + # block_in_cluster_coord_vmnk: this CTA's (V, M, N, K) coordinate within its cluster. + # Used by tma_partition to set up TMA multicast masks: + # [2] (N) → A's tma_partition (A is multicast along N) + # [1] (M) → B's tma_partition (B is multicast along M) + # For 1x1 cluster: always (0, 0, 0, 0), trivially no multicast. + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + + # ============================================================ + # SMEM allocation (SharedStorage struct for barriers + TMEM ptr) + # ============================================================ + DMA_Stage = self.num_ab_stage + + # Shared storage struct (matches C++ SharedStorage) + # There are 2 kinds of barriers used here: + # 1. Transaction barriers (tma_mma_full): 64-bit in SMEM, support transaction byte + # counting for TMA completion tracking. set_barrier_transaction_bytes arrives and + # sets the expected TX count; TMA increments TX bytes as data arrives in SMEM. + # When both arrival count and TX count are met, the barrier phase flips. + # 2. Cluster barriers (all others): 64-bit in SMEM, simpler arrive/wait semantics. + # Used for thread-to-thread synchronization within/across warps. + @cute.struct + class SharedStorage: + # Barrier between TMA and MMA: TMA tells MMA the SMEM tile is ready/full + tma_mma_full_barrier: cute.struct.MemRange[cutlass.Int64, DMA_Stage] + # Barrier between MMA and TMA: MMA tells TMA the SMEM tile is consumed/empty + tma_mma_empty_barrier: cute.struct.MemRange[cutlass.Int64, DMA_Stage] + # Barrier between TMA_B and epilog: TMA_B tells epilog all activation loads are issued + tma_epilog_full_barrier: cutlass.Int64 + # Barrier between MMA and epilog: MMA tells epilog the accumulator is ready + mma_epilog_full_barrier: cutlass.Int64 + # Barrier between MMA and epilog: sync TMEM allocation/deallocation status + tmem_allocation_result_barrier: cutlass.Int64 + # Base pointer for TMEM allocation, MMA will write the allocated address here + tmem_base_ptr: cutlass.Int32 + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # Extract barrier pointers from struct + tma_mma_full_bar = storage.tma_mma_full_barrier.data_ptr() + tma_mma_empty_bar = storage.tma_mma_empty_barrier.data_ptr() + tma_epilog_full_bar = storage.tma_epilog_full_barrier.ptr + mma_epilog_full_bar = storage.mma_epilog_full_barrier.ptr + tmem_alloc_result_bar = storage.tmem_allocation_result_barrier.ptr + tmem_base_smem_ptr = storage.tmem_base_ptr.ptr + + # ============================================================ + # Barrier initialization — ALL threads reach here, elect_one inits + # C++: warp 0 initializes all barriers + # + # Phase initializes to 0, expected arrival count is set below. + # For transaction barriers (tma_mma_full): set_barrier_transaction_bytes will + # arrive and set TX count; TMA DRAM requests increment TX bytes. When both + # arrival count and TX count are met, the barrier phase flips. + # ============================================================ + if warp_idx == 0: + with cute.arch.elect_one(): + # Transaction barrier: 2 arrivals (DMA_A + DMA_B each set TX bytes) + for i in range(DMA_Stage): + cute.arch.mbarrier_init(tma_mma_full_bar + i, 2) + # Cluster barrier: 1 arrival (MMA warp signals SMEM slot consumed) + for i in range(DMA_Stage): + cute.arch.mbarrier_init(tma_mma_empty_bar + i, 1) + + # 32 threads/1 warp (DMA_B warp) arrive to signal epilog + cute.arch.mbarrier_init(tma_epilog_full_bar, 32) + # 1 thread (MMA warp via tcgen05.commit) signals epilog + cute.arch.mbarrier_init(mma_epilog_full_bar, 1) + # 32 (MMA warp) + 128 (EPILOG warps 4-7) = 160 arrivals for TMEM alloc sync + cute.arch.mbarrier_init(tmem_alloc_result_bar, 32 + 128) + + # Barrier initialization needs to be visible to all warps before proceeding + # C++: fence_barrier_init() + __syncthreads() + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + + # ============================================================ + # SMEM Tensor Allocation + # C++: SharedStorage contains alignas(128) ArrayEngine for A and B + # ============================================================ + # sA: ((Mma_M, Mma_K), NumMma_M, NumMma_K, DMA_Stage) with swizzle + sA = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + # sB: ((Mma_N, Mma_K), NumMma_N, NumMma_K, DMA_Stage) with swizzle + sB = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + + # ============================================================ + # Warp dispatch (C++: gemm_device warp_idx dispatch) + # Each warp specializes in a different role: + # Warp 0 (threads 0-31): DMA_A — loads A tiles via TMA + # Warp 1 (threads 32-63): DMA_B — loads B tiles via TMA + # Warp 2 (threads 64-95): MMA — performs tcgen05.mma + # Warp 3 (threads 96-127): unused + # Warps 4-7 (threads 128-255): EPILOG — TMEM->RMEM->GMEM store + # ============================================================ + if warp_idx == 0: + self.dma_a_warp( + tma_mma_full_bar, + tma_mma_empty_bar, + tma_atom_a, + mA_mkl, + sA, + tiled_mma, + cluster_layout_vmnk, + block_in_cluster_coord_vmnk, + mma_tile_coord_v, + work_tile_info, + k_tile_count, + ) + elif warp_idx == 1: + self.dma_b_warp( + tma_mma_full_bar, + tma_mma_empty_bar, + tma_epilog_full_bar, + tma_atom_b, + mB_nkl, + sB, + tiled_mma, + cluster_layout_vmnk, + block_in_cluster_coord_vmnk, + mma_tile_coord_v, + work_tile_info, + k_tile_count, + ) + elif warp_idx == 2: + self.mma_warp( + tma_mma_full_bar, + tma_mma_empty_bar, + mma_epilog_full_bar, + tmem_alloc_result_bar, + tiled_mma, + sA, + sB, + tmem_base_smem_ptr, + mma_tile_coord_v, + k_tile_count, + ) + elif warp_idx >= 4: + # Epilog tid is from 128 to 255, need to offset by -128 + # when getting the per-thread slice of the tiled_copy_t2r + epi_tid = tidx - 128 + self.epilog_warp( + tma_epilog_full_bar, + mma_epilog_full_bar, + tmem_alloc_result_bar, + tiled_mma, + mC_mnl, + tmem_base_smem_ptr, + epi_tid, + mma_tile_coord_v, + work_tile_info, + ) + + # Final sync — C++: __syncthreads() + cute.arch.barrier() + return + + @cute.jit + def dma_a_warp( + self, + tma_mma_full_bar, + tma_mma_empty_bar, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L) — TMA coordinate tensor + sA: cute.Tensor, # ((Mma_M, Mma_K), NumMma_M, NumMma_K, DMA_Stage) + tiled_mma: cute.TiledMma, + cluster_layout_vmnk: cute.Layout, + block_in_cluster_coord_vmnk: tuple, + mma_tile_coord_v: cutlass.Int32, + work_tile_info: tuple, # WorkTileInfo: (M_idx, N_idx, L_idx, K_idx_start, K_idx_end) + k_tile_count: cutlass.Int32, + ): + """DMA_A warp: loads A tiles via TMA. C++ DMA_A_warp() lines 204-318.""" + DMA_Stage = self.num_ab_stage + + # ---- Self-contained tensor partitioning (matches C++) ---- + # Partitioned tensors use tXgY naming convention: + # tX = partitioning pattern applied to tensor gY + # tC = tensor partitioned into the MMA shape, i.e. ((Mma_M, Mma_K), ...) + # tA = tensor partitioned into the TMA shape, i.e. (TMA, ...) + # g = gmem, s = smem, t = tmem, r = rmem + + # Tile mA_mkl (Gemm_M, Gemm_K, Gemm_L) into CTA-level tiles + # C++: local_tile(mA, make_shape(CTA_M, CTA_K), make_coord(M_idx, _, L_idx)) + # gA_mkl: (CTA_M, CTA_K, Tiles_M, Tiles_K, Gemm_L) — all tiles + gA_mkl = cute.local_tile(mA_mkl, (self.cta_m, self.cta_k), (None, None, None)) + # Partition gmem tensor into the MMA shape + # In 1SM mode, tiled_mma has 1 "thread", so get_slice(0) gives the only partition + # C++: ThrMMA cta_mma = tiled_mma.get_slice(0) + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + # tCgA: ((Mma_M, Mma_K), NumMma_M, NumMma_K, Tiles_M, Tiles_K, Gemm_L) + tCgA = thr_mma.partition_A(gA_mkl) + + # ---- TMA partition for cluster-level multicast ---- + # tma_partition sets up TMA descriptors that can multicast data to multiple CTAs + # in the same cluster that share the same tile. It takes 3 key arguments: + # 1. This CTA's coordinate within the multicast group + # 2. The multicast group layout (how many CTAs participate) + # 3. The SMEM/GMEM tensors to partition + # + # For A: all CTAs along the N dimension of the cluster share the same A tile + # (same M, same K, different N). So the multicast group is the N-slice of the + # cluster layout. For a 2x3 cluster, 3 CTAs along N would share each A tile. + # + # a_cta_layout: layout of CTAs that share A data = N-dimension of cluster. + # Extracted by slicing cluster_layout_vmnk at V=0, M=0, K=0, keeping N free. + # For 1x1 cluster: shape (_1,) — just this CTA, no multicast. + # block_in_cluster_coord_vmnk[2]: this CTA's position within that N multicast group. + # For 1x1 cluster: always 0. + # + # group_modes groups ((Mma_M, Mma_K), NumMma_M, NumMma_K) into a single mode + # so the TMA handles the entire SMEM tile in one copy call. + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], # this CTA's N coord within cluster + a_cta_layout, + cute.group_modes( + sA, 0, 3 + ), # (((Mma_M, Mma_K), NumMma_M, NumMma_K), DMA_Stage) + cute.group_modes( + tCgA, 0, 3 + ), # (((Mma_M, Mma_K), NumMma_M, NumMma_K), Tiles_M, Tiles_K, Gemm_L) + ) + # tAsA: ((TMA, NumTma_K), DMA_Stage) — SMEM destination for each pipeline stage + # tAgA: ((TMA, NumTma_K), Tiles_M, Tiles_K, Gemm_L) — GMEM source tiles + + # Slice to this CTA's M tile and batch index, keep K tiles and TMA modes free + # C++: implicit via work_tile_info.M_idx and work_tile_info.L_idx in local_tile + # tAgA after slice: ((TMA, NumTma_K), Tiles_K) + tAgA = tAgA[(None, work_tile_info.M_idx, None, work_tile_info.L_idx)] + + # ---- K-loop ---- + # Initial phase bit for tma_mma_empty_barrier is 0 (set by mbarrier_init), denoting SMEM slot is empty. + # We wait on the "old" phase bit of 1: when it flips to 0 (the initial value), the slot is empty. + # Example for DMA_Stage=2, barriers a and b: + # kblock 0: old phase=1, wait for a to flip to 0 (initial), slot a empty + # kblock 1: old phase=1, wait for b to flip to 0 (initial), slot b empty + # kblock 2: old phase=0, wait for a to flip to 1 (flipped once), slot a empty + # kblock 3: old phase=0, wait for b to flip to 1 (flipped once), slot b empty + # ... + tma_mma_empty_phase = cutlass.Int32(1) + pdl_count = self.pdl_count + + # Iterate over k-blocks + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % DMA_Stage + # Wait for MMA to signal that SMEM slot is empty, ready for next TMA load + # C++: wait_barrier(shared_storage.tma_mma_empty_barrier[k_tile % DMA_Stage], phase) + cute.arch.mbarrier_wait( + tma_mma_empty_bar + stage, + tma_mma_empty_phase, + ) + + # Set the barrier transaction bytes for tracking TMA completion + # C++: set_barrier_transaction_bytes(..., tma_transaction_bytes) + # Must be done by a single thread (elect_one) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_mma_full_bar + stage, + self.tma_bytes_a, + ) + # Issue TMA load: load A tile (CTA_M, CTA_K) from GMEM into SMEM + # C++: copy(tma_atom_A->with(barrier), tAgA(_,k_tile), tAsA(_,stage)) + # CRITICAL: cute.copy for TMA must be OUTSIDE elect_one()! + # TMA partition internally ensures only thread 0 issues cp.async.bulk. + # Putting cute.copy inside elect_one() causes GPU deadlock in DSL. + cute.copy( + tma_atom_a, + tAgA[(None, k_tile)], # GMEM source: ((TMA, NumTma_K),) for k_tile + tAsA[(None, stage)], # SMEM destination: ((TMA, NumTma_K),) for stage + tma_bar_ptr=tma_mma_full_bar + stage, + ) + + # For every DMA_Stage iterations, flip the phase bit to reuse barriers + # C++: tma_mma_empty_barrier_phase_bit ^= 1 + if (k_tile % DMA_Stage) == (DMA_Stage - 1): + tma_mma_empty_phase = tma_mma_empty_phase ^ 1 + + # PDL: do launch dependents at a specific k_tile count + # C++: if (k_tile == pdl_count) cutlass::arch::launch_dependent_grids() + if self.use_pdl: + if k_tile == pdl_count: + cute.arch.griddepcontrol_launch_dependents() + + # PDL: if pdl_count is -1, do launch dependents at the end + # C++: cutlass::arch::launch_dependent_grids() (unconditional at end) + if self.use_pdl: + cute.arch.griddepcontrol_launch_dependents() + + @cute.jit + def dma_b_warp( + self, + tma_mma_full_bar, + tma_mma_empty_bar, + tma_epilog_full_bar, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L) — TMA coordinate tensor + sB: cute.Tensor, # ((Mma_N, Mma_K), NumMma_N, NumMma_K, DMA_Stage) + tiled_mma: cute.TiledMma, + cluster_layout_vmnk: cute.Layout, + block_in_cluster_coord_vmnk: tuple, + mma_tile_coord_v: cutlass.Int32, + work_tile_info: tuple, # WorkTileInfo: (M_idx, N_idx, L_idx, K_idx_start, K_idx_end) + k_tile_count: cutlass.Int32, + ): + """DMA_B warp: loads B tiles via TMA. C++ DMA_B_warp() lines 327-390.""" + DMA_Stage = self.num_ab_stage + + # ---- Self-contained tensor partitioning (similar to DMA_A_warp) ---- + # Tile mB_nkl (Gemm_N, Gemm_K, Gemm_L) into CTA-level tiles + # C++: local_tile(mB, make_shape(CTA_N, CTA_K), make_coord(N_idx, _, L_idx)) + # gB_nkl: (CTA_N, CTA_K, Tiles_N, Tiles_K, Gemm_L) + gB_nkl = cute.local_tile(mB_nkl, (self.cta_n, self.cta_k), (None, None, None)) + # C++: ThrMMA cta_mma = tiled_mma.get_slice(0) + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + # tCgB: ((Mma_N, Mma_K), NumMma_N, NumMma_K, Tiles_N, Tiles_K, Gemm_L) + tCgB = thr_mma.partition_B(gB_nkl) + + # Same pattern as A, but for B the multicast group is along M. + # All CTAs along M share the same B tile (same N, same K, different M). + # b_cta_layout: layout of CTAs that share B data = M-dimension of cluster. + # Extracted by slicing cluster_layout_vmnk at V=0, N=0, K=0, keeping M free. + # For 1x1 cluster: shape (_1,) — just this CTA, no multicast. + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], # this CTA's M coord within cluster + b_cta_layout, + cute.group_modes( + sB, 0, 3 + ), # (((Mma_N, Mma_K), NumMma_N, NumMma_K), DMA_Stage) + cute.group_modes( + tCgB, 0, 3 + ), # (((Mma_N, Mma_K), NumMma_N, NumMma_K), Tiles_N, Tiles_K, Gemm_L) + ) + # tBsB: ((TMA, NumTma_K), DMA_Stage) + # tBgB: ((TMA, NumTma_K), Tiles_N, Tiles_K, Gemm_L) + + # Slice to this CTA's N tile and batch index + # tBgB after slice: ((TMA, NumTma_K), Tiles_K) + tBgB = tBgB[(None, work_tile_info.N_idx, None, work_tile_info.L_idx)] + + # PDL: only do GDC wait on B loading, B is activation + # C++: cutlass::arch::wait_on_dependent_grids() + if self.use_pdl: + cute.arch.griddepcontrol_wait() + + # ---- K-loop (same phase tracking as DMA_A_warp) ---- + tma_mma_empty_phase = cutlass.Int32(1) + + # Iterate over k-blocks + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % DMA_Stage + # Wait for MMA to signal SMEM slot is empty + cute.arch.mbarrier_wait( + tma_mma_empty_bar + stage, + tma_mma_empty_phase, + ) + + # Set transaction bytes and issue TMA load for B tile (CTA_N, CTA_K) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_mma_full_bar + stage, + self.tma_bytes_b, + ) + # CRITICAL: cute.copy for TMA must be OUTSIDE elect_one() + cute.copy( + tma_atom_b, + tBgB[(None, k_tile)], # GMEM source: ((TMA, NumTma_K),) for k_tile + tBsB[(None, stage)], # SMEM destination: ((TMA, NumTma_K),) for stage + tma_bar_ptr=tma_mma_full_bar + stage, + ) + + # Flip phase bit every DMA_Stage iterations + if (k_tile % DMA_Stage) == (DMA_Stage - 1): + tma_mma_empty_phase = tma_mma_empty_phase ^ 1 + + # All activation/B loads are issued, signal the epilog warp + # C++: arrive_barrier(shared_storage.tma_epilog_full_barrier) + cute.arch.mbarrier_arrive(tma_epilog_full_bar) + + @cute.jit + def mma_warp( + self, + tma_mma_full_bar, + tma_mma_empty_bar, + mma_epilog_full_bar, + tmem_alloc_result_bar, + tiled_mma: cute.TiledMma, + sA: cute.Tensor, # ((Mma_M, Mma_K), NumMma_M, NumMma_K, DMA_Stage) + sB: cute.Tensor, # ((Mma_N, Mma_K), NumMma_N, NumMma_K, DMA_Stage) + tmem_base_smem_ptr, + mma_tile_coord_v: cutlass.Int32, + k_tile_count: cutlass.Int32, + ): + """MMA warp: performs tcgen05.mma. C++ MMA_warp() lines 398-591.""" + DMA_Stage = self.num_ab_stage + + # ---- MMA Fragment Allocation ---- + # We allocate "fragments" which are SMEM descriptors that serve as inputs to cute.gemm. + # For tcgen05.mma operations: + # - Matrices A and B are sourced from SMEM + # - tCrA and tCrB provide descriptor views of sA and sB respectively + # - The first mode of each descriptor represents the SMEM for a single MMA instruction + # C++: tCrA = cta_mma.make_fragment_A(tCsA) + # tCrA: ((Mma_M, Mma_K), NumMma_M, NumMma_K, DMA_Stage) — SMEM descriptors for A + tCrA = tiled_mma.make_fragment_A(sA) + # tCrB: ((Mma_N, Mma_K), NumMma_N, NumMma_K, DMA_Stage) — SMEM descriptors for B + tCrB = tiled_mma.make_fragment_B(sB) + + # TMEM Accumulator shape + # On SM100, accumulators are stored exclusively in tensor memory (TMEM). + # make_fragment_C creates a TMEM tensor view with the appropriate layout. + # tCtAcc is a view of the accumulator tensor; its tmem base ptr is unset until alloc_tmem. + # acc_shape / tCtAcc: ((Mma_M, Mma_N), NumMma_M, NumMma_N) + # For bf16 M64 N8 K16: (((_16,_4),_8),_1,_1) — Mma_M=(16,4)=(Mma_M_per_subp, NumSubp), Mma_N=8 + acc_shape = tiled_mma.partition_shape_C((self.cta_m, self.cta_n)) + tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape) + + # ---- TMEM Allocation ---- + # Only use half of TMEM to allow overlapping with next CTA + # TMEM has 128 lanes, 512 columns, each word is 4B, 256KB total + # Our accumulator uses CTA_N * sizeof(AccType) columns + # C++: tmem_allocator.allocate(Sm100TmemCapacityColumns / 2, &shared_storage.tmem_base_ptr) + num_tmem_cols = 256 # SM100_TMEM_CAPACITY_COLUMNS / 2 + cute.arch.alloc_tmem(num_tmem_cols, tmem_base_smem_ptr) + + # Notify epilog warp that TMEM allocation is complete + # 32 threads (MMA warp) arrive; NO WAIT here — just arrives and continues + # C++: arrive_barrier(shared_storage.tmem_allocation_result_barrier) + cute.arch.mbarrier_arrive(tmem_alloc_result_bar) + + # Relinquish TMEM allocation lock early so that the next prefetch CTA can be launched + # C++: tmem_allocator.release_allocation_lock() + cute.arch.relinquish_tmem_alloc_permit() + + # Retrieve TMEM ptr from SMEM and create the accumulator tensor view + # C++: tCtAcc.data() = shared_storage.tmem_base_ptr + tmem_ptr = cute.arch.retrieve_tmem_ptr( + self.acc_dtype, + 16, + tmem_base_smem_ptr, + ) + # tCtAcc: ((Mma_M, Mma_N), NumMma_M, NumMma_N) — TMEM accumulator view + # For bf16 M64 N8: tmem layout is N-major; stride between dp lanes is 65536, + # stride between subpartitions is 2097152 + tCtAcc = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # ---- MMA mainloop with try_wait/wait overlap ---- + # This follows the CUTLASS pattern to overlap barrier waits with MMA instructions: + # for each kblock: + # if not waitComplete: blocking wait + # update stage_idx + # try_wait for next iteration (non-blocking, returns predicate) + # execute MMA + # try_wait provides best-effort wait: the mbarrier.try_wait returns true/false, + # and subsequent instructions (MMA) are NOT blocked, unlike the blocking wait_barrier + # which spin-loops until the barrier is ready + # penalty before MMA can issue) + # + # Initial phase bit for tma_mma_full_barrier is 0: + # kblock 0: old phase=0, wait for phase to flip to 1 (slot full) + # kblock 1: old phase=0, wait for phase to flip to 1 (slot full) + # ... after DMA_Stage iterations, phase flips + tma_mma_full_phase = cutlass.Int32(0) + # Explicitly manage stage_idx instead of k_tile % DMA_Stage to avoid + # expensive modulo when DMA_Stage is non-power-of-2 (tighter mainloop) + stage_idx = cutlass.Int32(0) + old_stage_idx = cutlass.Int32(0) + + # Peel off the first iteration: try_wait for the first stage (non-blocking) + # C++: bool waitComplete = try_wait_barrier(barrier[0], phase) + waitComplete = cute.arch.mbarrier_try_wait( + tma_mma_full_bar + stage_idx, tma_mma_full_phase + ) + + # Iterate over k-blocks + for k_tile in cutlass.range(k_tile_count): + # If try_wait failed, do a blocking wait + if ~waitComplete: + cute.arch.mbarrier_wait( + tma_mma_full_bar + stage_idx, tma_mma_full_phase + ) + + # Update stage index + old_stage_idx = stage_idx + stage_idx = stage_idx + 1 + # Flip phase bit every DMA_Stage iterations + if stage_idx == DMA_Stage: + tma_mma_full_phase = tma_mma_full_phase ^ 1 + stage_idx = cutlass.Int32(0) + + # Try wait for the next iteration (non-blocking, overlaps with MMA below) + if k_tile < (k_tile_count - 1): + waitComplete = cute.arch.mbarrier_try_wait( + tma_mma_full_bar + stage_idx, tma_mma_full_phase + ) + + # Manually unroll the NumMma_K loop so that each cute.gemm issues exactly + # one MMA instruction, and we can control the ACCUMULATE flag per instruction. + # First instruction of first k_tile: ACCUMULATE=False (C = A*B, clears TMEM) + # All subsequent instructions: ACCUMULATE=True (C += A*B) + # C++: for (k_block = 0; k_block < size<2>(tCrA); ++k_block) + # cute::gemm(tiled_mma, tCrA(_,_,k_block,stage), tCrB(_,_,k_block,stage), tCtAcc) + for k_block in range(self.mma_inst_tile_k): + if k_block == 0: + # First k_block of each k_tile: accumulate only if not the very first k_tile + # C++: tiled_mma.accumulate_ = (k_tile == 0) ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One + tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + else: + # Subsequent k_blocks always accumulate (C += A*B) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + # Execute a single Mma_M x Mma_N x Mma_K GEMM instruction + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[(None, None, k_block, old_stage_idx)], + tCrB[(None, None, k_block, old_stage_idx)], + tCtAcc, + ) + + # Notify DMA warps that all MMA issued prior to this tcgen05.commit are done, + # the SMEM slot is now consumed and can be reused for the next TMA load + # C++: cutlass::arch::umma_arrive(&shared_storage.tma_mma_empty_barrier[old_stage_idx]) + # CRITICAL: tcgen05.commit MUST be inside elect_one() — in C++, umma_arrive + # internally does elect_one_sync(), but DSL's tcgen05.commit has NO internal guard. + # Without elect_one(), all 32 threads commit (32x redundant tcgen05.commit instructions). + with cute.arch.elect_one(): + tcgen05.commit( + tma_mma_empty_bar + old_stage_idx, + None, + self.cta_group, + ) + + # Notify the epilog warp that MMA is done, the TMEM accumulator is ready to consume + # C++: cutlass::arch::umma_arrive(&shared_storage.mma_epilog_full_barrier) + with cute.arch.elect_one(): + tcgen05.commit( + mma_epilog_full_bar, + None, + self.cta_group, + ) + + # ---- TMEM dealloc protocol (matches C++ MMA_warp lines 581-591) ---- + # Wait for tmem deallocation signal from epilog warp + # C++: arrive_barrier(shared_storage.tmem_allocation_result_barrier) + cute.arch.mbarrier_arrive(tmem_alloc_result_bar) + # Initial phase bit = 1 since it's already flipped once for TMEM allocation + # It will flip to 0 when TMEM can be deallocated, so we wait for old phase of 1 + # C++: wait_barrier(shared_storage.tmem_allocation_result_barrier, 1) + cute.arch.mbarrier_wait(tmem_alloc_result_bar, 1) + # Deallocate TMEM + # C++: tmem_allocator.free(shared_storage.tmem_base_ptr, Sm100TmemCapacityColumns / 2) + cute.arch.dealloc_tmem(tmem_ptr, num_tmem_cols) + + @cute.jit + def epilog_warp( + self, + tma_epilog_full_bar, + mma_epilog_full_bar, + tmem_alloc_result_bar, + tiled_mma: cute.TiledMma, + mC_mnl: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L) — output tensor in GMEM + tmem_base_smem_ptr, + epi_tid: cutlass.Int32, # thread id within epilog warps (0-127) + mma_tile_coord_v: cutlass.Int32, + work_tile_info: tuple, # WorkTileInfo: (M_idx, N_idx, L_idx, K_idx_start, K_idx_end) + ): + """EPILOG warp: TMEM -> RMEM -> type convert -> GMEM. C++ EPILOG_warp() lines 600-838.""" + # ---- Self-contained tensor partitioning (matches C++) ---- + # Get the local tile of C for this CTA + # C++: local_tile(mC, make_shape(CTA_M, CTA_N), make_coord(M_idx, N_idx, L_idx)) + # gC_mnl: (CTA_M, CTA_N, Tiles_M, Tiles_N, Gemm_L) + gC_mnl = cute.local_tile(mC_mnl, (self.cta_m, self.cta_n), (None, None, None)) + # C++: ThrMMA cta_mma = tiled_mma.get_slice(0) + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + # tCgC: ((Mma_M, Mma_N), NumMma_M, NumMma_N, Tiles_M, Tiles_N, Gemm_L) + tCgC = thr_mma.partition_C(gC_mnl) + + # Since tCtAcc is a view of the accumulator tensor, it's safe to create a new view + # in the epilog warp too. The layout is the same as in the MMA warp. + # acc_shape / tCtAcc_fake: ((Mma_M, Mma_N), NumMma_M, NumMma_N) + acc_shape = tiled_mma.partition_shape_C((self.cta_m, self.cta_n)) + tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape) + + # ---- Arrive + wait tmem_allocation_result_barrier phase 0 ---- + # 128 threads from EPILOG warps (4-7) arrive + # C++: arrive_barrier(shared_storage.tmem_allocation_result_barrier) + cute.arch.mbarrier_arrive(tmem_alloc_result_bar) + # Wait for MMA warp to complete TMEM allocation + # Initial phase=0, it will flip to 1 when TMEM is allocated, so we wait for old phase 0 + # C++: wait_barrier(shared_storage.tmem_allocation_result_barrier, 0) + cute.arch.mbarrier_wait(tmem_alloc_result_bar, 0) + + # Update TMEM base ptr of the accumulator tensor view + # C++: tCtAcc.data() = shared_storage.tmem_base_ptr + tmem_ptr = cute.arch.retrieve_tmem_ptr( + self.acc_dtype, + 16, + tmem_base_smem_ptr, + ) + # tCtAcc: ((Mma_M, Mma_N), NumMma_M, NumMma_N) — TMEM accumulator view + tCtAcc = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # ---- Create tiled copy for TMEM -> RMEM ---- + # C++: TiledCopy tiled_t2r_copy = make_tmem_copy( + # TMEM::op_repeater(), tCtAcc) + # For M=64, M/N major output, we use SM100_TMEM_LOAD_16dp256b1x: + # - 16dp version because M=64, each sub-partition uses 16dp of TMEM in MMA + # - 256b = 32B per load, loads 16dp x 8col per tcgen05.ld instruction + # The op_repeater automatically figures out the tcgen05.ld repeat count based on CTA_N + copy_atom_t2r = sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), + self.c_layout, + self.c_dtype, + self.acc_dtype, + (self.cta_m, self.cta_n), + self.use_2cta_instrs, + ) + + # C++: TiledCopy tiled_t2r_copy = make_tmem_copy(op, tCtAcc) + tiled_copy_t2r = tcgen05.make_tmem_copy( + copy_atom_t2r, tCtAcc[((None, None), 0, 0)] + ) + # Epilog tid is 0-127 (threads 128-255 offset by -128) + # C++: ThrCopy thr_t2r_copy = tiled_t2r_copy.get_slice(tid) + thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tid) + + # tD means the partitioning pattern of tcgen05.ld (TMEM->RMEM copy) + # tDtAcc: per-subpartition view of the accumulator tensor (TMEM source) + # Shape: (CpyS, NumCpy_M, NumCpy_N) + # For bf16 M64 N8, tcgen05.ld atom SM100_TMEM_LOAD_16dp256b1x: + # CpyS = (CpyS_N, CpyS_M) = (8, 16) — 16dp x 8col per tcgen05.ld instruction + # NumCpy_M = 1, NumCpy_N = CTA_N / CpyS_N = 1 + # C++: tDtAcc = thr_t2r_copy.partition_S(tCtAcc) + tDtAcc = thr_copy_t2r.partition_S(tCtAcc[((None, None), 0, 0)]) + + # Partition tCgC for 2 reasons: + # 1. Get post-partition shape for allocating rmem space for the accumulator + # 2. Partition it for storing the result back to gmem (per-thread slice) + # tDgC represents which values in tCgC are stored in this thread's rmem + # after tcgen05.ld from TMEM to RMEM + # Shape: (CpyD, NumCpy_M, NumCpy_N, Tiles_M, Tiles_N, Gemm_L) + # For bf16 M64 N8: CpyD = (2,2) — 4 registers per thread per tcgen05.ld + # C++: tDgC = thr_t2r_copy.partition_D(tCgC) + # DSL tCgC has extra grid dims; partition with them, create rmem before slicing + tDgC = thr_copy_t2r.partition_D(tCgC[((None, None), 0, 0, None, None, None)]) + + # Allocate per-thread rmem space for the accumulator + # tDrAcc and tDrC have the same shape as the single-tile slice of tDgC + # Shape: (CpyD, NumCpy_M, NumCpy_N) + # C++: Tensor tDrAcc = make_tensor(shape(tDgC)) + # IMPORTANT: Create rmem tensors BEFORE slicing tDgC with work_tile_info, + # because slicing reduces rank and breaks shape extraction + epi_frag_shape = tDgC[(None, None, None, 0, 0, 0)].shape + # tDrAcc: (CpyD, NumCpy_M, NumCpy_N) — per-thread accumulator in rmem (AccType) + tDrAcc = cute.make_rmem_tensor(epi_frag_shape, self.acc_dtype) + # tDrC: (CpyD, NumCpy_M, NumCpy_N) — per-thread converted output in rmem (TypeC) + tDrC = cute.make_rmem_tensor(epi_frag_shape, self.c_dtype) + + # Slice gmem partition to this CTA's output tile + # tDgC after slice: (CpyD, NumCpy_M, NumCpy_N) — per-thread GMEM store destinations + tDgC = tDgC[ + ( + None, + None, + None, + work_tile_info.M_idx, + work_tile_info.N_idx, + work_tile_info.L_idx, + ) + ] + + # ---- Bounds-check predicate (C++: make_identity_tensor + elem_less) ---- + # An identity tensor maps (m, n, l) coordinates to the payload tuple (m, n, l). + # We use this to check if each thread's store coordinates are within bounds, + # because the problem shape may not be a multiple of the tile size. + # For TMA loads, OOB is handled automatically; for stg stores, we need explicit predicates. + # C++: Tensor coordC = make_identity_tensor(shape(mC)) + coordC = cute.make_identity_tensor(mC_mnl.shape) # (M, N, L) -> (m, n, l) + # Create the local tile of coordC, same tiling as gC + # C++: local_tile(coordC, make_shape(CTA_M, CTA_N), make_coord(M_idx, N_idx, L_idx)) + gCcoord = cute.local_tile(coordC, (self.cta_m, self.cta_n), (None, None, None)) + # tCcC: ((Mma_M, Mma_N), NumMma_M, NumMma_N, Tiles_M, Tiles_N, Gemm_L) + tCcC = thr_mma.partition_C(gCcoord) + # tDcC: (CpyD, NumCpy_M, NumCpy_N, Tiles_M, Tiles_N, Gemm_L) + # — per-thread coordinate tensor, same shape as tDgC but payload is (m,n,l) coord + tDcC = thr_copy_t2r.partition_D(tCcC[((None, None), 0, 0, None, None, None)]) + # Slice to this CTA's tile: (CpyD, NumCpy_M, NumCpy_N) + tDcC = tDcC[ + ( + None, + None, + None, + work_tile_info.M_idx, + work_tile_info.N_idx, + work_tile_info.L_idx, + ) + ] + # Construct predicate tensor: compare each coordinate with problem shape (M, N, L) + # tDpC(t) = true if coordinate is in-bounds, false if out-of-bounds + # C++: tDpC(t) = elem_less(tDcC(t), shape(mC)) + tDpC = cute.make_rmem_tensor( + tDcC.shape, cutlass.Boolean + ) # (CpyD, NumCpy_M, NumCpy_N) + for i in range(cute.size(tDpC)): + tDpC[i] = cute.elem_less(tDcC[i], mC_mnl.shape) + + # ---- Wait for TMA_B to finish ---- + # Initial phase=0, it flips to 1 when TMA_B warp is done + # C++: wait_barrier(shared_storage.tma_epilog_full_barrier, 0) + cute.arch.mbarrier_wait(tma_epilog_full_bar, 0) + + # ---- Wait for MMA to finish ---- + # Initial phase=0, it flips to 1 when MMA warp is done + # C++: wait_barrier(shared_storage.mma_epilog_full_barrier, 0) + cute.arch.mbarrier_wait(mma_epilog_full_bar, 0) + + # ---- Copy TMEM -> RMEM (single copy, no subtile loop) ---- + # The copy operation is: + # for each NumCpy_M: + # for each NumCpy_N: + # tcgen05.ld.16dp256bit.x1: CpyS -> CpyD + # Each tcgen05.ld instruction copies 16dp x 8col (CpyS) for each TMEM subpartition + # to 128 threads' rmem (4 registers (CpyD) per thread) + # C++: copy(tiled_t2r_copy, tDtAcc, tDrAcc) + cute.copy(tiled_copy_t2r, tDtAcc, tDrAcc) + + # ---- Signal MMA warp that TMEM read is done ---- + # Wait for tcgen05.ld to finish so we can safely deallocate TMEM + # C++: cutlass::arch::fence_view_async_tmem_load() + cute.arch.fence_view_async_tmem_load() + # C++: arrive_barrier(shared_storage.tmem_allocation_result_barrier) + cute.arch.mbarrier_arrive(tmem_alloc_result_bar) + + # ---- Type conversion (AccType -> TypeC) and predicated store to GMEM ---- + # C++: converter(tDrAcc) -> tDrC + # Convert accumulator data from AccType (e.g. fp32) to TypeC (e.g. bf16) + acc_vec = tDrAcc.load().to(self.c_dtype) + tDrC.store(acc_vec) + # Predicated store: rmem -> gmem + # The store is unfortunately uncoalesced because gmem addresses across threads + # are not contiguous, but this is acceptable for small batch sizes. + # C++: copy_if(tDpC, tDrC, tDgC) + # For each element: if tDpC[i] is true, store tDrC[i] to tDgC[i] via stg + cute.basic_copy_if(tDpC, tDrC, tDgC) diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm80_tensorop_gemm_impl.py b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm80_tensorop_gemm_impl.py new file mode 100644 index 000000000..3deafa7d5 --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm80_tensorop_gemm_impl.py @@ -0,0 +1,814 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import math + +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils + +""" +A high-performance batched dense GEMM example for the utilzing SM80 style TensorCore instructions +- Matrix A is LxMxK, L is batch dimension, A can be row-major("t") or column-major("n") +- Matrix B is LxKxN, L is batch dimension, B can be row-major("t") or column-major("n") +- Matrix C is LxMxN, L is batch dimension, C can be row-major("t") or column-major("n") + - Internally, this is treated as an MxKxL, NxKxL, MxNxL CuTe Tensors + +This GEMM kernel supports the following features: + - Utilizes SM80's tensor cores for matrix multiply-accumulate (MMA) operations + - Threadblock rasterization to improve data re-use + - Supports multi-stage pipeline to overlap computation and memory access + +Constraints: +* Supported input and output data types: fp16, bf16 +* Support accumulator data types: f32 +* Atom layout's MNK shape is set so that tile shape can be divided by MMA instruction shape +* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, +""" + + +class Sm80TensorOpGemmKernel: + def __init__( + self, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + acc_dtype: type[cutlass.Numeric], + ): + self.ab_dtype = a_dtype + self.c_dtype = c_dtype + self.acc_dtype = acc_dtype + self.cta_tiler = (128, 128, 32) + self.num_stages = 4 + self.atom_layout_mnk = (2, 2, 1) + atom_lay_M, atom_lay_N, atom_lay_K = self.atom_layout_mnk + self.num_threads = atom_lay_M * atom_lay_N * atom_lay_K * 32 + + self.bM, self.bN, self.bK = self.cta_tiler + self.mma_inst_shape = (16, 8, 16) + mmaM, mmaN, mmaK = self.mma_inst_shape + + assert a_dtype == b_dtype, "This example does not support different A, B dtypes" + assert self.bM % (atom_lay_M * mmaM) == 0, ( + "bM must be divisible by MMA instruction" + ) + assert self.bN % (atom_lay_N * mmaN) == 0, ( + "bN must be divisible by MMA instruction" + ) + assert atom_lay_K == 1, "this example does not support atom layout K > 1" + assert self.bK % mmaK == 0, "bK must be divisible by MMA instruction" + assert self.num_stages >= 3, "num_stages must be greater than or equal to 3" + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + def add_batch_mode(tensor: cute.Tensor) -> cute.Tensor: + return cute.make_tensor( + tensor.iterator, + cute.prepend(tensor.layout, cute.make_layout(1), up_to_rank=3), + ) + + a = add_batch_mode(a) + b = add_batch_mode(b) + c = add_batch_mode(c) + + # Permute tensor modes from torch to cute convention + # A: (L, M, K) -> (M, K, L) + a = cute.make_tensor(a.iterator, cute.select(a.layout, [1, 2, 0])) + # B: (L, K, N) -> (N, K, L) + b = cute.make_tensor(b.iterator, cute.select(b.layout, [2, 1, 0])) + # C: (L, M, N) -> (M, N, L) + c = cute.make_tensor(c.iterator, cute.select(c.layout, [1, 2, 0])) + + # The grid divides the problems's M, N, and L dimensions by the + # respective modes of the tile shape (bM, bN, 1). The K dimension is + # handled within a block via a multistage process. + + self.a_major_mode = utils.LayoutEnum.from_tensor(a) + self.b_major_mode = utils.LayoutEnum.from_tensor(b) + self.c_major_mode = utils.LayoutEnum.from_tensor(c) + + # /////////////////////////////////////////////////////////////////////////////// + # Shared memory layout: + # /////////////////////////////////////////////////////////////////////////////// + + # Creates a layout with the size required for the provided tile + # size and num stages (stages are used for K dimension) that is also + # sectioned into 64x8 or 8x32 layout atoms. The swizzle is set so that + # the atom for the shared memory -> register copy does not encounter + # bank conflicts + + # assume the input is 16B align + ab_copy_bits = 128 + sA_layout = self._make_smem_layout_AB( + a.element_type, + self.a_major_mode, + ab_copy_bits, + (self.cta_tiler[0], self.cta_tiler[2], self.num_stages), + ) + sB_layout = self._make_smem_layout_AB( + b.element_type, + self.b_major_mode, + ab_copy_bits, + (self.cta_tiler[1], self.cta_tiler[2], self.num_stages), + ) + + # Creates a similar layout but without num_stages or layout atoms + sC_layout = self._make_smem_layout_C( + c.element_type, + self.c_major_mode, + ab_copy_bits, + (self.cta_tiler[0], self.cta_tiler[1]), + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Tiled copy: + # The majorness of tA/tB/tC follows the majorness of gA/gB/gC, + # enabling merged accesses to global memory for faster data + # transfer between global and shared memory. + # /////////////////////////////////////////////////////////////////////////////// + + # Create a copy atom for a global to shared memory asynchronous copy + atom_async_copy = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(cache_mode=cute.nvgpu.LoadCacheMode.GLOBAL), + a.element_type, + num_bits_per_copy=ab_copy_bits, + ) + + # Create thread layouts for tiled copy from the copy atom where the + # thread layout simply follows the leading dimension of the tensor + tiled_copy_A = self._make_gmem_tiled_copy_AB( + atom_async_copy, a.element_type, self.a_major_mode, ab_copy_bits + ) + tiled_copy_B = self._make_gmem_tiled_copy_AB( + atom_async_copy, b.element_type, self.b_major_mode, ab_copy_bits + ) + + # Creates a synchronous copy atom and thread layouts for the epilogue + c_copy_bits = 128 + atom_sync_copy = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + c.element_type, + num_bits_per_copy=c_copy_bits, + ) + tiled_copy_C = self._make_gmem_tiled_copy_C( + atom_sync_copy, c.element_type, self.c_major_mode, c_copy_bits + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Tiled MMA + # /////////////////////////////////////////////////////////////////////////////// + + # Creates a mma atom with 16x8x16 shape for MNK + op = cute.nvgpu.warp.MmaF16BF16Op( + self.ab_dtype, self.acc_dtype, self.mma_inst_shape + ) + + permutation_mnk = ( + self.atom_layout_mnk[0] * self.mma_inst_shape[0], + # if atom layout's N-mode is 1, to leverage the largest coalesced + # shared memory -> register copy, set the tiled mma's N mode to 16 + self.atom_layout_mnk[1] * self.mma_inst_shape[1] * 2, + self.atom_layout_mnk[2] * self.mma_inst_shape[2], + ) + + # Created a tiled mma that tiles the atom according to specified layout. + # For a 2x2x1 atom layout, the mma atom is duplicated 4 times, twice + # across M and twice across N + tC = cute.make_layout(self.atom_layout_mnk) + tiled_mma = cute.make_tiled_mma( + op, + tC, + permutation_mnk=permutation_mnk, + ) + + # grid_dim: ((m + BLK_M - 1) // BLK_M, (n + BLK_N - 1) // BLK_N, l) + grid_dim = cute.ceil_div(c.shape, (self.bM, self.bN, 1)) + + # Add threadblock rasterization to improve re-use of data + raster_factor = 1 + grid_dim_n = cute.size(grid_dim[1]) + # Thresholds picked so that it doesn't cause too many no-op CTAs + if grid_dim_n > 5: + raster_factor = 8 + elif grid_dim_n > 2: + raster_factor = 4 + elif grid_dim_n > 1: + raster_factor = 2 + rasterization_remap_grid_dim = ( + cute.size(grid_dim[0]) * raster_factor, + (cute.size(grid_dim[1]) + raster_factor - 1) // raster_factor, + cute.size(grid_dim[2]), + ) + + self.kernel( + a, + b, + c, + sA_layout, + sB_layout, + sC_layout, + tiled_copy_A, + tiled_copy_B, + tiled_copy_C, + tiled_mma, + raster_factor, + epilogue_op, + ).launch( + grid=rasterization_remap_grid_dim, + block=[self.num_threads, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mA: cute.Tensor, + mB: cute.Tensor, + mC: cute.Tensor, + sA_layout: cute.ComposedLayout, + sB_layout: cute.ComposedLayout, + sC_layout: cute.ComposedLayout, + tiled_copy_A: cute.TiledCopy, + tiled_copy_B: cute.TiledCopy, + tiled_copy_C: cute.TiledCopy, + tiled_mma: cute.TiledMma, + rasterization_factor: cutlass.Int32, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + # Thread index, block index + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, bidz = cute.arch.block_idx() + grid_dim = cute.ceil_div(mC.shape, (self.bM, self.bN, 1)) + offset_tile_x, offset_tile_y = self.raster_tile( + bidx, bidy, rasterization_factor + ) + # Early exit if CTA is out of range + if grid_dim[0] <= offset_tile_x or grid_dim[1] <= offset_tile_y: + pass + else: + tiler_coord = (offset_tile_x, offset_tile_y, None) + + # /////////////////////////////////////////////////////////////////////////////// + # Get the appropriate tiles for this thread block. + # gA: (BLK_M, BLK_N, k), gB: (BLK_N, BLK_K, k), gC: (BLK_M, BLK_N) + # /////////////////////////////////////////////////////////////////////////////// + gA = cute.local_tile( + mA[None, None, bidz], + tiler=self.cta_tiler, + coord=tiler_coord, + proj=(1, None, 1), + ) + gB = cute.local_tile( + mB[None, None, bidz], + tiler=self.cta_tiler, + coord=tiler_coord, + proj=(None, 1, 1), + ) + gC = cute.local_tile( + mC[None, None, bidz], + tiler=self.cta_tiler, + coord=tiler_coord, + proj=(1, 1, None), + ) + + # By default, if the tensor k mode does not divide into the tile k + # size, then last tiles in the k dimension are irregular. + # Instead, make the first tiles irregular when k is irregular. + # This allows us to handle the irregular tile first to avoid + # checking for this condition within the mainloop. + + # residual_k is a negative number indicating the amount needed to + # shift the pointer by in dimension k + residual_k = cute.size(mA, mode=[1]) - cutlass.Int32(self.bK) * cute.size( + gA, mode=[2] + ) + + # move the pointer of gA/gB in the `-k` direction + gA = cute.domain_offset((0, residual_k, 0), gA) + gB = cute.domain_offset((0, residual_k, 0), gB) + # input is 16B aligned + gA = cute.make_tensor(gA.iterator.align(16), gA.layout) + gB = cute.make_tensor(gB.iterator.align(16), gB.layout) + + # Construct identity layout for sA and sB (mirrors global tensors, + # used for predication only) + mcA = cute.make_identity_tensor(mA.layout.shape) + mcB = cute.make_identity_tensor(mB.layout.shape) + cA = cute.local_tile( + mcA[None, None, bidz], + tiler=self.cta_tiler, + coord=tiler_coord, + proj=(1, None, 1), + ) + cB = cute.local_tile( + mcB[None, None, bidz], + tiler=self.cta_tiler, + coord=tiler_coord, + proj=(None, 1, 1), + ) + + cA = cute.domain_offset((0, residual_k, 0), cA) + cB = cute.domain_offset((0, residual_k, 0), cB) + + # /////////////////////////////////////////////////////////////////////////////// + # Create shared memory buffers and get the appropriate fragments for this thread. + # sA: (BLK_M, BLK_K, PIPE) , sB: (BLK_N, BLK_K, PIPE) + # tAgA: (CPY, CPY_M, CPY_K, k) , tBgB: (CPY, CPY_N, CPY_K, k) + # tAsA: (CPY, CPY_M, CPY_K, PIPE) , tBsB: (CPY, CPY_N, CPY_K, PIPE) + # /////////////////////////////////////////////////////////////////////////////// + @cute.struct + class SharedStorageAB: + a: cute.struct.Align[ + cute.struct.MemRange[mA.element_type, cute.cosize(sA_layout)], + 16, + ] + b: cute.struct.Align[ + cute.struct.MemRange[mB.element_type, cute.cosize(sB_layout)], + 16, + ] + + @cute.struct + class SharedStorageC: + c: cute.struct.Align[ + cute.struct.MemRange[mC.element_type, cute.cosize(sC_layout)], + 16, + ] + + # Shared memory buffer + smem = cutlass.utils.SmemAllocator() + # Shared memory allocated for operations with A, B will be + # overwritten for operations on C. This is to improve performance + # by reducing the size of shared memory requested by each block + storage = smem.allocate( + max(SharedStorageAB.size_in_bytes(), SharedStorageC.size_in_bytes()), + byte_alignment=16, + ) + sA = SharedStorageAB(storage).a.get_tensor(sA_layout) + sB = SharedStorageAB(storage).b.get_tensor(sB_layout) + sC = SharedStorageC(storage).c.get_tensor(sC_layout) + + thr_copy_A = tiled_copy_A.get_slice(tidx) + thr_copy_B = tiled_copy_B.get_slice(tidx) + thr_copy_C = tiled_copy_C.get_slice(tidx) + tAgA = thr_copy_A.partition_S(gA) + tAsA = thr_copy_A.partition_D(sA) + tBgB = thr_copy_B.partition_S(gB) + tBsB = thr_copy_B.partition_D(sB) + tCsC_epilogue = thr_copy_C.partition_S(sC) + tCgC_epilogue = thr_copy_C.partition_D(gC) + + # Repeat the partitioning with identity layouts + tAcA = thr_copy_A.partition_S(cA) + tBcB = thr_copy_B.partition_S(cB) + + # /////////////////////////////////////////////////////////////////////////////// + # Predicate: Mark indices that need to copy when problem_shape isn't a multiple + # of tile_shape + # /////////////////////////////////////////////////////////////////////////////// + + # For predication over the tensors A (M/K), B (N/K), and (in the + # epilogue) C (M/N), we will compute it in a fashion similar to an + # outer product. The predication along one of the dimensions is + # evaluated and stored in a predication tensor. Then, the + # predication for the remaining dimension is handled later via an + # if/else branch at the copy. + # For A and B, predication booleans along M/N are stored in a + # predication tensor and along K is handled via a if/else branch. + + # Allocate predicate tensors for M and N. Predication is checked + # at the granularity of a copy atom, so the predicate tensor does not + # need separate booleans for individual elements within a copy + # atom (for example, the elements of tAgA.shape[0][0].) + tApA = cute.make_rmem_tensor( + cute.make_layout( + ( + tAgA.shape[0][1], + cute.size(tAgA, mode=[1]), + cute.size(tAgA, mode=[2]), + ), + stride=(cute.size(tAgA, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + tBpB = cute.make_rmem_tensor( + cute.make_layout( + ( + tBsB.shape[0][1], + cute.size(tBsB, mode=[1]), + cute.size(tBsB, mode=[2]), + ), + stride=(cute.size(tBsB, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + # Set predicates for M/N bounds + for rest_v in range(tApA.shape[0]): + for m in range(tApA.shape[1]): + tApA[rest_v, m, 0] = cute.elem_less( + tAcA[(0, rest_v), m, 0, 0][0], mA.shape[0] + ) + for rest_v in range(tBpB.shape[0]): + for n in range(tBpB.shape[1]): + tBpB[rest_v, n, 0] = cute.elem_less( + tBcB[(0, rest_v), n, 0, 0][0], mB.shape[0] + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Prefetch Prologue + # /////////////////////////////////////////////////////////////////////////////// + # Clear the smem tiles to account for predicated off loads + tAsA.fill(0) + tBsB.fill(0) + cute.arch.sync_threads() + # Start async loads for the first k-tile. Here we take care of the k residue + # via if/else check along the k dimension. Because we shifted the identity tensor + # by the residue_k and because the identity tensor is a coord tensor, the + # values of any identity tensor element that is poison is less than -1 + num_smem_stages = cute.size(tAsA, mode=[3]) + k_tile_count = cute.size(tAgA, mode=[3]) + k_tile_index = cutlass.Int32(0) + + for k in range(tApA.shape[2]): + if cute.elem_less(cutlass.Int32(-1), tAcA[0, 0, k, 0][1]): + cute.copy( + tiled_copy_A, + tAgA[None, None, k, k_tile_index], + tAsA[None, None, k, 0], + pred=tApA[None, None, k], + ) + for k in range(tBpB.shape[2]): + if cute.elem_less(cutlass.Int32(-1), tBcB[0, 0, k, 0][1]): + cute.copy( + tiled_copy_B, + tBgB[None, None, k, k_tile_index], + tBsB[None, None, k, 0], + pred=tBpB[None, None, k], + ) + k_tile_index = k_tile_index + 1 + cute.arch.cp_async_commit_group() + + # Start async loads for rest of the k-tiles + for k_tile in range(1, num_smem_stages - 1): + if k_tile == k_tile_count: + tApA.fill(0) + tBpB.fill(0) + cute.copy( + tiled_copy_A, + tAgA[None, None, None, k_tile_index], + tAsA[None, None, None, k_tile], + pred=tApA, + ) + cute.copy( + tiled_copy_B, + tBgB[None, None, None, k_tile_index], + tBsB[None, None, None, k_tile], + pred=tBpB, + ) + k_tile_index = k_tile_index + 1 + cute.arch.cp_async_commit_group() + + # /////////////////////////////////////////////////////////////////////////////// + # Tile MMA compute thread partitions and allocate accumulators + # /////////////////////////////////////////////////////////////////////////////// + thr_mma = tiled_mma.get_slice(tidx) + tCsA = thr_mma.partition_A(sA) + tCsB = thr_mma.partition_B(sB) + tCsC = thr_mma.partition_C(sC) + tCgC = thr_mma.partition_C(gC) + tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0]) + tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0]) + tCrC = tiled_mma.make_fragment_C(tCgC) + # Clear the accumulator + tCrC.fill(0.0) + + # /////////////////////////////////////////////////////////////////////////////// + # Copy Atom A/B retiling + # /////////////////////////////////////////////////////////////////////////////// + + # Create the copy atoms for the copy from shared memory to register + atom_copy_s2r_A = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp( + self.a_major_mode != utils.LayoutEnum.ROW_MAJOR, 4 + ), + mA.element_type, + ) + atom_copy_s2r_B = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp( + self.b_major_mode != utils.LayoutEnum.ROW_MAJOR, 4 + ), + mB.element_type, + ) + + # Creates the tiled copy so that it matches the thread-value layout + # expected by the tiled mma + tiled_copy_s2r_A = cute.make_tiled_copy_A(atom_copy_s2r_A, tiled_mma) + tiled_copy_s2r_B = cute.make_tiled_copy_B(atom_copy_s2r_B, tiled_mma) + + thr_copy_ldmatrix_A = tiled_copy_s2r_A.get_slice(tidx) + thr_copy_ldmatrix_B = tiled_copy_s2r_B.get_slice(tidx) + tCsA_copy_view = thr_copy_ldmatrix_A.partition_S(sA) + tCrA_copy_view = thr_copy_ldmatrix_A.retile(tCrA) + tCsB_copy_view = thr_copy_ldmatrix_B.partition_S(sB) + tCrB_copy_view = thr_copy_ldmatrix_B.retile(tCrB) + + # Current pipe index in smem to read from / write to + smem_pipe_read = 0 + smem_pipe_write = num_smem_stages - 1 + + tCsA_p = tCsA_copy_view[None, None, None, smem_pipe_read] + tCsB_p = tCsB_copy_view[None, None, None, smem_pipe_read] + + # /////////////////////////////////////////////////////////////////////////////// + # PREFETCH register pipeline + # /////////////////////////////////////////////////////////////////////////////// + num_k_block = cute.size(tCrA, mode=[2]) + if num_k_block > 1: + # Wait until our first prefetched tile is loaded in + cute.arch.cp_async_wait_group(num_smem_stages - 2) + cute.arch.sync_threads() + # Prefetch the first k-block rmem from the first k-tile + cute.copy( + tiled_copy_s2r_A, + tCsA_p[None, None, 0], + tCrA_copy_view[None, None, 0], + ) + cute.copy( + tiled_copy_s2r_B, + tCsB_p[None, None, 0], + tCrB_copy_view[None, None, 0], + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Mainloop + # 1. Shared memory pipeline (gmem -> smem): + # The default smem pipeline depth is 3, meaning that for shared + # memory buffers, we allocate three times the size described by the + # CTA tiler. We prefetch 2 of these buffers before entering the main + # loop. Considering only the transfer from global memory to shared + # memory, the general structure of the mainloop is: + # (1) copy k-tile from gmem to smem; + # (2) perform gemm computation on k-tile; + # (3) wait for the next copy to finish. + # The `cute.arch.cp_async_wait_group(num_smem_stages - 2)` command + # waits for the number of unfinished 'copy' to be <= 1. The advantage + # of this approach is that it allows for simultaneous production + # (i.e., step (1)) and consumption (i.e., step (2)) of smem. + # A common misconception is to prefetch N buffers and rewrite + # the pipeline logic to wait on N-1 pending copies. The disadvantage + # of this approach is that it requires fully consuming a buffer in + # order to open an empty buffer for the next copy. + # 2. Register pipeline (smem -> register): + # Similarly, the register pipeline produces i+1, consumes i, and + # produces i+2... Notably, i and i+1 do not use the same register, + # eliminating dependencies on the same register for better parallelism. + # 3. Combining the smem and register pipelines results in the mainloop. + # /////////////////////////////////////////////////////////////////////////////// + for k_tile in range(k_tile_count): + for k_block in cutlass.range(num_k_block, unroll_full=True): + if k_block == num_k_block - 1: + tCsA_p = tCsA_copy_view[None, None, None, smem_pipe_read] + tCsB_p = tCsB_copy_view[None, None, None, smem_pipe_read] + cute.arch.cp_async_wait_group(num_smem_stages - 2) + cute.arch.sync_threads() + + # Load A, B from shared memory to registers for k_block + 1 + k_block_next = (k_block + 1) % num_k_block # static + cute.copy( + tiled_copy_s2r_A, + tCsA_p[None, None, k_block_next], + tCrA_copy_view[None, None, k_block_next], + ) + cute.copy( + tiled_copy_s2r_B, + tCsB_p[None, None, k_block_next], + tCrB_copy_view[None, None, k_block_next], + ) + + # Fetch next A and B and update smem pipeline read/write + if k_block == 0: + if k_tile + num_smem_stages - 1 < k_tile_count: + cute.copy( + tiled_copy_A, + tAgA[None, None, None, k_tile_index], + tAsA[None, None, None, smem_pipe_write], + pred=tApA, + ) + if k_tile + num_smem_stages - 1 < k_tile_count: + cute.copy( + tiled_copy_B, + tBgB[None, None, None, k_tile_index], + tBsB[None, None, None, smem_pipe_write], + pred=tBpB, + ) + k_tile_index = k_tile_index + 1 + cute.arch.cp_async_commit_group() + smem_pipe_write = smem_pipe_read + smem_pipe_read = smem_pipe_read + 1 + if smem_pipe_read == num_smem_stages: + smem_pipe_read = 0 + + # Thread-level register gemm for k_block + cute.gemm( + tiled_mma, + tCrC, + tCrA[None, None, k_block], + tCrB[None, None, k_block], + tCrC, + ) + + # Sync before epilogue + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + + # /////////////////////////////////////////////////////////////////////////////// + # Epilogue with fusion + # /////////////////////////////////////////////////////////////////////////////// + tCrD = cute.make_fragment_like(tCrC, self.c_dtype) + tCrD[None] = epilogue_op(tCrC.load()).to(self.c_dtype) + + # Copy results of D back to shared memory + cute.autovec_copy(tCrD, tCsC) + + # Create coord tensor for C + ceilM, ceilN, _ = cute.ceil_div(mC.shape, (self.bM, self.bN, 1)) + mcC = cute.make_identity_tensor( + ( + cute.size(ceilM) * self.cta_tiler[0], + cute.size(ceilN) * self.cta_tiler[1], + 1, + ) + ) + cC = cute.local_tile( + mcC[None, None, bidz], + tiler=self.cta_tiler, + coord=tiler_coord, + proj=(1, 1, None), + ) + tCcC = thr_copy_C.partition_S(cC) + + tCrC_epilogue = cute.make_fragment_like(tCsC_epilogue) + # Wait for all writes to shared memory to finish before starting copies + # using the new layouts + cute.arch.sync_threads() + cute.autovec_copy(tCsC_epilogue, tCrC_epilogue) + + # Create predication tensor for m + tCpC = cute.make_rmem_tensor( + cute.make_layout( + ( + tCgC_epilogue.shape[0][1], + cute.size(tCgC_epilogue, mode=[1]), + cute.size(tCgC_epilogue, mode=[2]), + ), + stride=(cute.size(tCgC_epilogue, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + for rest_v in range(tCpC.shape[0]): + for m in range(tCpC.shape[1]): + tCpC[rest_v, m, 0] = cute.elem_less( + tCcC[(0, rest_v), m, 0][0], mC.shape[0] + ) + + # Copy to global memory using better vectorization + for rest_v in range(tCpC.shape[0]): + for n in range(tCpC.shape[2]): + if cute.elem_less(tCcC[(0, rest_v), 0, n][1], mC.shape[1]): + cute.copy( + tiled_copy_C, + tCrC_epilogue[None, None, n], + tCgC_epilogue[None, None, n], + pred=tCpC[None, None, n], + ) + return + + def _make_smem_layout_AB(self, dtype, major_mode, copy_bits, smem_tiler): + major_mode_size = ( + smem_tiler[1] if major_mode == utils.LayoutEnum.ROW_MAJOR else smem_tiler[0] + ) + major_mode_size = min(64, major_mode_size) + + swizzle_bits = int(math.log2(major_mode_size * dtype.width // copy_bits)) + swizzle_bits = min(swizzle_bits, 3) + + layout_atom_outer = ( + cute.make_layout((8, major_mode_size), stride=(major_mode_size, 1)) + if major_mode == utils.LayoutEnum.ROW_MAJOR + else cute.make_layout((major_mode_size, 8), stride=(1, major_mode_size)) + ) + layout_atom = cute.make_composed_layout( + cute.make_swizzle(swizzle_bits, 3, 3), + 0, + layout_atom_outer, + ) + layout = cute.tile_to_shape(layout_atom, smem_tiler, (0, 1, 2)) + return layout + + def _make_smem_layout_C(self, dtype, major_mode, copy_bits, smem_tiler): + major_mode_size = ( + smem_tiler[1] if major_mode == utils.LayoutEnum.ROW_MAJOR else smem_tiler[0] + ) + + swizzle_bits = int(math.log2(major_mode_size * dtype.width // copy_bits)) + swizzle_bits = min(swizzle_bits, 3) + + layout_atom_outer = ( + cute.make_layout((8, major_mode_size), stride=(major_mode_size, 1)) + if major_mode == utils.LayoutEnum.ROW_MAJOR + else cute.make_layout((major_mode_size, 8), stride=(1, major_mode_size)) + ) + layout_atom = cute.make_composed_layout( + cute.make_swizzle(swizzle_bits, 3, 4), + 0, + layout_atom_outer, + ) + + # Due to the thread layout of the mma, remove swizzle in C to + # prevent shared memory fragments owned by an single thread from + # holding swizzles + if major_mode == utils.LayoutEnum.COL_MAJOR: + layout_atom = cute.make_composed_layout( + cute.make_swizzle(0, 3, 4), 0, layout_atom_outer + ) + layout = cute.tile_to_shape( + layout_atom, + smem_tiler, + (0, 1), + ) + return layout + + def _make_gmem_tiled_copy_AB(self, atom_copy, dtype, major_mode, copy_bits): + copy_elems = copy_bits // dtype.width + shape_dim_1 = cute.size(self.bK) // copy_elems + # thread layout for copy + thread_layout = cute.make_layout( + (self.num_threads // shape_dim_1, shape_dim_1), stride=(shape_dim_1, 1) + ) + if major_mode != utils.LayoutEnum.ROW_MAJOR: + shape_dim_0 = cute.size(self.bM) // copy_elems + thread_layout = cute.make_layout( + (shape_dim_0, self.num_threads // shape_dim_0), stride=(1, shape_dim_0) + ) + # Value layout for copy + value_layout = ( + cute.make_layout((1, copy_elems)) + if major_mode == utils.LayoutEnum.ROW_MAJOR + else cute.make_layout((copy_elems, 1)) + ) + return cute.make_tiled_copy_tv(atom_copy, thread_layout, value_layout) + + def _make_gmem_tiled_copy_C(self, atom_copy, dtype, major_mode, copy_bits): + copy_elems = copy_bits // dtype.width + shape_dim_1 = cute.size(self.bN) // copy_elems + # thread layout for copy + thread_layout = cute.make_layout( + (self.num_threads // shape_dim_1, shape_dim_1), stride=(shape_dim_1, 1) + ) + if major_mode != utils.LayoutEnum.ROW_MAJOR: + shape_dim_0 = cute.size(self.bM) // copy_elems + thread_layout = cute.make_layout( + (shape_dim_0, self.num_threads // shape_dim_0), stride=(1, shape_dim_0) + ) + value_layout = ( + cute.make_layout((1, copy_elems)) + if major_mode == utils.LayoutEnum.ROW_MAJOR + else cute.make_layout((copy_elems, 1)) + ) + return cute.make_tiled_copy_tv(atom_copy, thread_layout, value_layout) + + def raster_tile(self, i, j, f): + new_i = i // f + new_j = (i % f) + (j * f) + return (new_i, new_j) diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm90_static_persistent_impl.py b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm90_static_persistent_impl.py new file mode 100644 index 000000000..6d23ed8f0 --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/implementations/sm90_static_persistent_impl.py @@ -0,0 +1,1236 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""SM90 (Hopper) Persistent Dense GEMM Implementation. + +A high-performance batched dense GEMM (C = A * B) for NVIDIA Hopper architecture +using CuTe DSL. + +Features: +- TMA (Tensor Memory Access) for efficient memory operations +- WGMMA (Warp Group Matrix Multiply-Accumulate) instructions +- TMA multicast with cluster for reduced L2 traffic +- Persistent tile scheduling for better memory/compute overlap +- Warp specialization to avoid explicit pipelining + +Constraints: +- CTA tile M: 64 or 128 +- CTA tile N: 64, 128, or 256 +- CTA tile K: 64 (auto-computed from MMA instruction) +- 8-bit types require K-major layout +- Cluster size <= 4 +""" + +import math + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.hopper_helpers as sm90_utils +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + + +class Sm90StaticPersistentGemmKernel: + """SM90 Persistent GEMM implementation using WGMMA and TMA. + + This class implements batched matrix multiplication (C = A x B) with support + for various data types and architectural features specific to Hopper GPUs. + + :param acc_dtype: Data type for accumulation during computation + :param tile_shape_mn: Shape of the CTA tile (M, N) + :param cluster_shape_mn: Cluster dimensions (M, N) for parallel processing + """ + + def __init__( + self, + acc_dtype: type[cutlass.Numeric], + tile_shape_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + swizzle_size: int = 1, + raster_along_m: bool = True, + ): + self.acc_dtype = acc_dtype + + self.cluster_shape_mn = cluster_shape_mn + self.swizzle_size = swizzle_size + self.raster_along_m = raster_along_m + self.mma_inst_shape_mn = None + # K dimension is deferred in _setup_attributes + self.tile_shape_mnk = (*tile_shape_mn, 1) + # For large tile size, using two warp groups is preferred because using only one warp + # group may result in register spill + self.atom_layout_mnk = ( + (2, 1, 1) + if self.tile_shape_mnk[0] > 64 and self.tile_shape_mnk[1] > 128 + else (1, 1, 1) + ) + self.num_mcast_ctas_a = None + self.num_mcast_ctas_b = None + self.is_a_mcast = False + self.is_b_mcast = False + self.tiled_mma = None + + self.occupancy = 1 + self.num_dma_warp_groups = 1 + self.num_mma_warp_groups = math.prod(self.atom_layout_mnk) + self.num_warps_per_warp_group = 4 + self.num_threads_per_warp_group = self.num_warps_per_warp_group * 32 + self.threads_per_cta = ( + self.num_dma_warp_groups + self.num_mma_warp_groups + ) * self.num_threads_per_warp_group + self.load_warp_id = 0 + self.epi_store_warp_id = ( + self.num_dma_warp_groups * self.num_warps_per_warp_group + ) + self.load_register_requirement = 40 + self.mma_register_requirement = 232 + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_90") + + self.ab_stage = None + self.epi_stage = None + + self.a_smem_layout_staged = None + self.b_smem_layout_staged = None + self.epi_smem_layout_staged = None + self.epi_tile = None + + self.shared_storage = None + self.buffer_align_bytes = 1024 + + self.num_mma_threads = ( + self.num_mma_warp_groups * self.num_threads_per_warp_group + ) + self.epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, num_threads=self.num_mma_threads + ) + + def _setup_attributes(self): + """Set up configurations that are dependent on GEMM inputs + + This method configures various attributes based on the input tensor properties + (data types, leading dimensions) and kernel settings: + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B + - Computing epilogue subtile + - Setting up A/B/C stage counts in shared memory + - Computing A/B/C shared memory layout + """ + # check the cta tile shape + if self.tile_shape_mnk[0] not in [64, 128]: + raise ValueError("CTA tile shape M must be 64/128") + if self.tile_shape_mnk[1] not in [64, 128, 256]: + raise ValueError("CTA tile shape N must be 64/128/256") + + self.tiled_mma = sm90_utils.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_layout.sm90_mma_major_mode(), + self.b_layout.sm90_mma_major_mode(), + self.acc_dtype, + self.atom_layout_mnk, + tiler_mn=(64, self.tile_shape_mnk[1]), + ) + mma_inst_shape_k = cute.size(self.tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.tile_shape_mnk = ( + self.tile_shape_mnk[0], + self.tile_shape_mnk[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + + self.cta_layout_mnk = cute.make_layout((*self.cluster_shape_mn, 1)) + self.num_mcast_ctas_a = self.cluster_shape_mn[1] + self.num_mcast_ctas_b = self.cluster_shape_mn[0] + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + is_cooperative = self.atom_layout_mnk == (2, 1, 1) + self.epi_tile = self._sm90_compute_tile_shape_or_override( + self.tile_shape_mnk, self.c_dtype, is_cooperative=is_cooperative + ) + + # Compute stage before compute smem layout + self.ab_stage, self.epi_stage = self._compute_stages( + self.tile_shape_mnk, + self.a_dtype, + self.b_dtype, + self.epi_tile, + self.c_dtype, + self.smem_capacity, + self.occupancy, + ) + + ( + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.epi_smem_layout_staged, + ) = self._make_smem_layouts( + self.tile_shape_mnk, + self.epi_tile, + self.a_dtype, + self.a_layout, + self.b_dtype, + self.b_layout, + self.ab_stage, + self.c_dtype, + self.c_layout, + self.epi_stage, + ) + + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream, + ): + """Execute the GEMM operation in steps: + - Setup static attributes + - Setup TMA load/store atoms and tensors + - Compute grid size + - Define shared storage for kernel + - Launch the kernel synchronously + + :param a: Input tensor A + :type a: cute.Tensor + :param b: Input tensor B + :type b: cute.Tensor + :param c: Output tensor C + :type c: cute.Tensor + :param max_active_clusters: Maximum number of active clusters + :type max_active_clusters: cutlass.Constexpr + :param stream: CUDA stream for asynchronous execution + :type stream: cuda.CUstream + """ + + def add_batch_mode(tensor: cute.Tensor) -> cute.Tensor: + return cute.make_tensor( + tensor.iterator, + cute.prepend(tensor.layout, cute.make_layout(1), up_to_rank=3), + ) + + a = add_batch_mode(a) + b = add_batch_mode(b) + c = add_batch_mode(c) + + # Transform tensors to the convention expected by the Hopper kernel + # Operator API uses (L, M, K) / (L, K, N) / (L, M, N) format (batch first) + # The Hopper kernel expects (M, K, L) / (N, K, L) / (M, N, L) format (batch last) + # A: (L, M, K) -> (M, K, L) using select [1, 2, 0] + a = cute.make_tensor(a.iterator, cute.select(a.layout, [1, 2, 0])) + # B: (L, K, N) -> (N, K, L) using select [2, 1, 0] + b = cute.make_tensor(b.iterator, cute.select(b.layout, [2, 1, 0])) + # C: (L, M, N) -> (M, N, L) using select [1, 2, 0] + c = cute.make_tensor(c.iterator, cute.select(c.layout, [1, 2, 0])) + + # setup static attributes before smem/grid/tma computation + self.a_dtype = a.element_type + self.b_dtype = b.element_type + self.c_dtype = c.element_type + self.a_layout = utils.LayoutEnum.from_tensor(a) + self.b_layout = utils.LayoutEnum.from_tensor(b) + self.c_layout = utils.LayoutEnum.from_tensor(c) + + if cutlass.const_expr( + self.a_dtype.width == 16 and self.a_dtype != self.b_dtype + ): + raise TypeError(f"Type mismatch: {self.a_dtype} != {self.b_dtype}") + if cutlass.const_expr(self.a_dtype.width != self.b_dtype.width): + raise TypeError( + f"Type width mismatch: {self.a_dtype.width} != {self.b_dtype.width}" + ) + if cutlass.const_expr(self.a_dtype.width != 16 and self.a_dtype.width != 8): + raise TypeError("a_dtype should be float16, float8, or int8 ") + + self._setup_attributes() + + tma_atom_a, tma_tensor_a = self._make_tma_atoms_and_tensors( + a, + self.a_smem_layout_staged, + (self.tile_shape_mnk[0], self.tile_shape_mnk[2]), + self.cluster_shape_mn[1], + ) + + tma_atom_b, tma_tensor_b = self._make_tma_atoms_and_tensors( + b, + self.b_smem_layout_staged, + (self.tile_shape_mnk[1], self.tile_shape_mnk[2]), + self.cluster_shape_mn[0], + ) + + tma_atom_c, tma_tensor_c = self._make_tma_store_atoms_and_tensors( + c, + self.epi_smem_layout_staged, + self.epi_tile, + ) + + tile_sched_params, grid = self._compute_grid( + c, + self.tile_shape_mnk, + self.cluster_shape_mn, + self.swizzle_size, + self.raster_along_m, + max_active_clusters, + ) + + @cute.struct + class SharedStorage: + mainloop_pipeline_array_ptr: cute.struct.MemRange[ + cutlass.Int64, self.ab_stage * 2 + ] + sA: cute.struct.Align[ + cute.struct.MemRange[ + self.a_dtype, cute.cosize(self.a_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + sB: cute.struct.Align[ + cute.struct.MemRange[ + self.b_dtype, cute.cosize(self.b_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + sC: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, + cute.cosize(self.epi_smem_layout_staged), + ], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + # Launch the kernel synchronously + self.kernel( + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c, + self.tiled_mma, + self.cta_layout_mnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.epi_smem_layout_staged, + tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + min_blocks_per_mp=1, + stream=stream, + ) + return + + @cute.kernel + def kernel( + self, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + tiled_mma: cute.TiledMma, + cta_layout_mnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + epi_smem_layout_staged: cute.ComposedLayout, + tile_sched_params: utils.PersistentTileSchedulerParams, + ): + """GPU device kernel performing the batched GEMM computation. + + :param tma_atom_a: TMA copy atom for A tensor + :type tma_atom_a: cute.CopyAtom + :param mA_mkl: Input tensor A + :type mA_mkl: cute.Tensor + :param tma_atom_b: TMA copy atom for B tensor + :type tma_atom_b: cute.CopyAtom + :param mB_nkl: Input tensor B + :type mB_nkl: cute.Tensor + :param tma_atom_c: TMA copy atom for C tensor + :type tma_atom_c: cute.CopyAtom + :param mC_mnl: Output tensor C + :type mC_mnl: cute.Tensor + :param tiled_mma: Tiled MMA object + :type tiled_mma: cute.TiledMma + :param cta_layout_mnk: CTA layout + :type cta_layout_mnk: cute.Layout + :param a_smem_layout_staged: Shared memory layout for A + :type a_smem_layout_staged: cute.ComposedLayout + :param b_smem_layout_staged: Shared memory layout for B + :type b_smem_layout_staged: cute.ComposedLayout + :param epi_smem_layout_staged: Shared memory layout for epilogue + :type epi_smem_layout_staged: cute.ComposedLayout + :param tile_sched_params: Parameters for the persistent tile scheduler + :type tile_sched_params: utils.PersistentTileSchedulerParams + """ + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # Prefetch Tma desc + if warp_idx == 0: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_a) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_b) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_c) + + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + cluster_coord_mnk = cta_layout_mnk.get_flat_coord(cta_rank_in_cluster) + + a_mcast_mask = cute.make_layout_image_mask( + cta_layout_mnk, cluster_coord_mnk, mode=1 + ) + b_mcast_mask = cute.make_layout_image_mask( + cta_layout_mnk, cluster_coord_mnk, mode=0 + ) + + a_mcast_mask = a_mcast_mask if self.is_a_mcast else 0 + b_mcast_mask = b_mcast_mask if self.is_b_mcast else 0 + a_smem_layout = cute.slice_(a_smem_layout_staged, (None, None, 0)) + b_smem_layout = cute.slice_(b_smem_layout_staged, (None, None, 0)) + tma_copy_bytes = cute.size_in_bytes( + self.a_dtype, a_smem_layout + ) + cute.size_in_bytes(self.b_dtype, b_smem_layout) + + # Alloc and init AB full/empty + ACC full mbar (pipeline) + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # mbar arrays + mainloop_pipeline_array_ptr = storage.mainloop_pipeline_array_ptr.data_ptr() + + # Threads/warps participating in this pipeline + mainloop_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread + ) + # Each warp will constribute to the arrive count with the number of mcast size + mcast_size = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + consumer_arrive_cnt = ( + mcast_size * self.num_mma_warp_groups * self.num_warps_per_warp_group + ) + mainloop_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, consumer_arrive_cnt + ) + + mainloop_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=mainloop_pipeline_array_ptr, + num_stages=self.ab_stage, + producer_group=mainloop_pipeline_producer_group, + consumer_group=mainloop_pipeline_consumer_group, + tx_count=tma_copy_bytes, + cta_layout_vmnk=cute.make_layout((1, *cta_layout_mnk.shape)), + defer_sync=True, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # Generate smem tensor A/B + sA = storage.sA.get_tensor( + a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner + ) + sB = storage.sB.get_tensor( + b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner + ) + sC = storage.sC.get_tensor( + epi_smem_layout_staged.outer, swizzle=epi_smem_layout_staged.inner + ) + + # Local_tile partition global tensors + # (bM, bK, RestM, RestK, RestL) + gA_mkl = cute.local_tile( + mA_mkl, + cute.slice_(self.tile_shape_mnk, (None, 0, None)), + (None, None, None), + ) + # (bN, bK, RestN, RestK, RestL) + gB_nkl = cute.local_tile( + mB_nkl, + cute.slice_(self.tile_shape_mnk, (0, None, None)), + (None, None, None), + ) + # (bM, bN, RestM, RestN, RestL) + gC_mnl = cute.local_tile( + mC_mnl, + cute.slice_(self.tile_shape_mnk, (None, None, 0)), + (None, None, None), + ) + + # Partition shared tensor for TMA load A/B + # TMA load A partition_S/D + a_cta_layout = cute.make_layout(cute.slice_(cta_layout_mnk, (0, None, 0)).shape) + a_cta_crd = cluster_coord_mnk[1] + tAsA, tAgA = cute.nvgpu.cpasync.tma_partition( + tma_atom_a, + a_cta_crd, + a_cta_layout, + cute.group_modes(sA, 0, 2), + cute.group_modes(gA_mkl, 0, 2), + ) + + # TMA load B partition_S/D + b_cta_layout = cute.make_layout(cute.slice_(cta_layout_mnk, (None, 0, 0)).shape) + b_cta_crd = cluster_coord_mnk[0] + tBsB, tBgB = cute.nvgpu.cpasync.tma_partition( + tma_atom_b, + b_cta_crd, + b_cta_layout, + cute.group_modes(sB, 0, 2), + cute.group_modes(gB_nkl, 0, 2), + ) + + # Partition global tensor for TiledMMA_A/B/C + warp_group_idx = cute.arch.make_warp_uniform( + tidx // self.num_threads_per_warp_group + ) + mma_warp_group_thread_layout = cute.make_layout( + self.num_mma_warp_groups, stride=self.num_threads_per_warp_group + ) + thr_mma = tiled_mma.get_slice( + mma_warp_group_thread_layout(warp_group_idx - self.num_dma_warp_groups) + ) + + # Make fragments + tCsA = thr_mma.partition_A(sA) + tCsB = thr_mma.partition_B(sB) + tCrA = tiled_mma.make_fragment_A(tCsA) + tCrB = tiled_mma.make_fragment_B(tCsB) + + tCgC = thr_mma.partition_C(gC_mnl) + acc_shape = tCgC.shape[:3] + accumulators = cute.make_rmem_tensor(acc_shape, self.acc_dtype) + + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # Cluster wait for barrier init + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + is_dma_warp_group = warp_group_idx < self.num_dma_warp_groups + if is_dma_warp_group: + cute.arch.setmaxregister_decrease(self.load_register_requirement) + + if warp_idx == self.load_warp_id: + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + mainloop_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.ab_stage + ) + + while work_tile.is_valid_tile: + tile_coord_mnl = work_tile.tile_idx + tAgA_mkl = tAgA[(None, tile_coord_mnl[0], None, tile_coord_mnl[2])] + tBgB_nkl = tBgB[(None, tile_coord_mnl[1], None, tile_coord_mnl[2])] + + mainloop_producer_state.reset_count() + + for k_tile in range(k_tile_cnt): + # Conditionally wait for AB buffer empty + mainloop_pipeline.producer_acquire(mainloop_producer_state) + # Slice to global/shared memref to current k_tile + tAgA_k = tAgA_mkl[(None, mainloop_producer_state.count)] + tAsA_pipe = tAsA[(None, mainloop_producer_state.index)] + + tBgB_k = tBgB_nkl[(None, mainloop_producer_state.count)] + tBsB_pipe = tBsB[(None, mainloop_producer_state.index)] + + # TMA load A/B + cute.copy( + tma_atom_a, + tAgA_k, + tAsA_pipe, + tma_bar_ptr=mainloop_pipeline.producer_get_barrier( + mainloop_producer_state + ), + mcast_mask=a_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_k, + tBsB_pipe, + tma_bar_ptr=mainloop_pipeline.producer_get_barrier( + mainloop_producer_state + ), + mcast_mask=b_mcast_mask, + ) + + # Mainloop pipeline's producer commit is a NOP + mainloop_pipeline.producer_commit(mainloop_producer_state) + mainloop_producer_state.advance() + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + mainloop_pipeline.producer_tail(mainloop_producer_state) + + # MMA warp group + if not is_dma_warp_group: + cute.arch.setmaxregister_increase(self.mma_register_requirement) + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + mainloop_consumer_read_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.ab_stage + ) + mainloop_consumer_release_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.ab_stage + ) + + num_k_blocks = cute.size(tCrA, mode=[2]) + + # Partition for epilogue + copy_atom_r2s = sm90_utils.sm90_get_smem_store_op( + self.c_layout, + elem_ty_d=self.c_dtype, + elem_ty_acc=self.acc_dtype, + ) + + copy_atom_C = cute.make_copy_atom( + cute.nvgpu.warp.StMatrix8x8x16bOp( + self.c_layout.is_m_major_c(), + 4, + ), + self.c_dtype, + ) + + tiled_copy_C_Atom = cute.make_tiled_copy_C_atom(copy_atom_C, tiled_mma) + + tiled_copy_r2s = cute.make_tiled_copy_S( + copy_atom_r2s, + tiled_copy_C_Atom, + ) + + # (R2S, R2S_M, R2S_N, PIPE_D) + thr_copy_r2s = tiled_copy_r2s.get_slice( + tidx - self.num_dma_warp_groups * self.num_threads_per_warp_group + ) + # (t)hread-partition for (r)egister to (s)mem copy (tRS_) + tRS_sD = thr_copy_r2s.partition_D(sC) + # (R2S, R2S_M, R2S_N) + tRS_rAcc = tiled_copy_r2s.retile(accumulators) + + # Allocate D registers. + rD_shape = cute.shape(thr_copy_r2s.partition_S(sC)) + tRS_rD_layout = cute.make_layout(rD_shape[:3]) + tRS_rD = cute.make_rmem_tensor(tRS_rD_layout.shape, self.acc_dtype) + tRS_rD_out = cute.make_rmem_tensor(tRS_rD_layout.shape, self.c_dtype) + size_tRS_rD = cute.size(tRS_rD) + + k_pipe_mmas = 1 + prologue_mma_cnt = min(k_pipe_mmas, k_tile_cnt) + + # Initialize tma store pipeline + tma_store_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.num_mma_threads, + ) + tma_store_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.epi_stage, + producer_group=tma_store_producer_group, + ) + + while work_tile.is_valid_tile: + tile_coord_mnl = work_tile.tile_idx + gC_mnl_slice = gC_mnl[(None, None, *tile_coord_mnl)] + + # MAINLOOP + mainloop_consumer_read_state.reset_count() + mainloop_consumer_release_state.reset_count() + accumulators.fill(0.0) + tiled_mma.set(cute.nvgpu.warpgroup.Field.ACCUMULATE, True) + cute.nvgpu.warpgroup.fence() + + for k_tile in range(prologue_mma_cnt): + # Wait for TMA copies to complete + mainloop_pipeline.consumer_wait(mainloop_consumer_read_state) + # WGMMA + for k_block_idx in cutlass.range_constexpr(num_k_blocks): + k_block_coord = ( + None, + None, + k_block_idx, + mainloop_consumer_read_state.index, + ) + cute.gemm( + tiled_mma, + accumulators, + tCrA[k_block_coord], + tCrB[k_block_coord], + accumulators, + ) + + cute.nvgpu.warpgroup.commit_group() + mainloop_consumer_read_state.advance() + + for k_tile in range(prologue_mma_cnt, k_tile_cnt): + # Wait for TMA copies to complete + mainloop_pipeline.consumer_wait(mainloop_consumer_read_state) + # WGMMA + for k_block_idx in cutlass.range_constexpr(num_k_blocks): + k_block_coord = ( + None, + None, + k_block_idx, + mainloop_consumer_read_state.index, + ) + cute.gemm( + tiled_mma, + accumulators, + tCrA[k_block_coord], + tCrB[k_block_coord], + accumulators, + ) + + cute.nvgpu.warpgroup.commit_group() + # Wait on the wgmma barrier for WGMMA to complete + cute.nvgpu.warpgroup.wait_group(k_pipe_mmas) + + mainloop_pipeline.consumer_release(mainloop_consumer_release_state) + mainloop_consumer_release_state.advance() + mainloop_consumer_read_state.advance() + + cute.nvgpu.warpgroup.wait_group(0) + for k_tile in range(prologue_mma_cnt): + mainloop_pipeline.consumer_release(mainloop_consumer_release_state) + mainloop_consumer_release_state.advance() + + # Epilogue + tCgC_for_tma_partition = cute.zipped_divide(gC_mnl_slice, self.epi_tile) + + # thread(b)lock-partition for (s)mem to (g)mem copy (bSG_) + bSG_sD, bSG_gD = cute.nvgpu.cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + cute.group_modes(sC, 0, 2), + tCgC_for_tma_partition, + ) + + epi_tile_num = cute.size(tCgC_for_tma_partition, mode=[1]) + epi_tile_shape = tCgC_for_tma_partition.shape[1] + epi_tile_layout = cute.make_layout( + epi_tile_shape, stride=(epi_tile_shape[1], 1) + ) + + num_prev_epi_tiles = tile_sched.num_tiles_executed * epi_tile_num + for epi_idx in cutlass.range_constexpr(epi_tile_num): + # Copy from accumulators to D registers + for epi_v in cutlass.range_constexpr(size_tRS_rD): + tRS_rD[epi_v] = tRS_rAcc[epi_idx * size_tRS_rD + epi_v] + + # Type conversion + acc_vec = tRS_rD.load() + tRS_rD_out.store(acc_vec.to(self.c_dtype)) + + # Copy from D registers to shared memory + epi_buffer = (num_prev_epi_tiles + epi_idx) % cute.size( + tRS_sD, mode=[3] + ) + cute.copy( + tiled_copy_r2s, + tRS_rD_out, + tRS_sD[(None, None, None, epi_buffer)], + ) + + cute.arch.fence_proxy("async.shared", space="cta") + self.epilog_sync_barrier.arrive_and_wait() + + gmem_coord = epi_tile_layout.get_hier_coord(epi_idx) + # Copy from shared memory to global memory + if warp_idx == self.epi_store_warp_id: + cute.copy( + tma_atom_c, + bSG_sD[(None, epi_buffer)], + bSG_gD[(None, gmem_coord)], + ) + tma_store_pipeline.producer_commit() + tma_store_pipeline.producer_acquire() + + self.epilog_sync_barrier.arrive_and_wait() + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + tma_store_pipeline.producer_tail() + + # ========================================================================= + # Helper methods + # ========================================================================= + + @staticmethod + def _compute_stages( + tile_shape_mnk: tuple[int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + epi_tile: tuple[int, int], + c_dtype: type[cutlass.Numeric], + smem_capacity: int, + occupancy: int, + ) -> tuple[int, int]: + """Computes the number of stages for A/B/C operands based on heuristics. + + :param tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type tile_shape_mnk: tuple[int, int, int] + :param a_dtype: Data type of operand A. + :type a_dtype: type[cutlass.Numeric] + :param b_dtype: Data type of operand B. + :type b_dtype: type[cutlass.Numeric] + :param epi_tile: Epilogue tile shape + :type epi_tile: Tuple[int, int] + :param c_dtype: The data type of the output tensor + :type c_dtype: type[cutlass.Numeric] + :param smem_capacity: Total available shared memory capacity in bytes. + :type smem_capacity: int + :param occupancy: Target number of CTAs per SM (occupancy). + :type occupancy: int + + :return: A tuple containing the computed number of stages for: + (A/B operand stages, epilogue stages) + :rtype: tuple[int, int] + """ + a_shape = cute.slice_(tile_shape_mnk, (None, 0, None)) + b_shape = cute.slice_(tile_shape_mnk, (0, None, None)) + ab_bytes_per_stage = ( + cute.size(a_shape) * a_dtype.width // 8 + + cute.size(b_shape) * b_dtype.width // 8 + ) + c_bytes_per_stage = cute.size(epi_tile) * c_dtype.width // 8 + epi_stage = 4 + epi_bytes = c_bytes_per_stage * epi_stage + + mbar_helpers_bytes = 1024 + + ab_stage = ( + smem_capacity // occupancy - (mbar_helpers_bytes + epi_bytes) + ) // ab_bytes_per_stage + return ab_stage, epi_stage + + @staticmethod + def _sm90_compute_tile_shape_or_override( + tile_shape_mnk: tuple[int, int, int], + element_type: type[cutlass.Numeric], + is_cooperative: bool = False, + epi_tile_override: tuple[int, int] | None = None, + ) -> tuple[int, int]: + """Compute the epilogue tile shape or use override if provided. + + :param tile_shape_mnk: CTA tile shape (M,N,K) + :type tile_shape_mnk: Tuple[int, int, int] + :param element_type: Data type of elements + :type element_type: type[cutlass.Numeric] + :param is_cooperative: Whether to use cooperative approach + :type is_cooperative: bool + :param epi_tile_override: Optional override for epilogue tile shape + :type epi_tile_override: Tuple[int, int] or None + + :return: Computed epilogue tile shape + :rtype: Tuple[int, int] + """ + if epi_tile_override is not None: + return epi_tile_override + if is_cooperative: + tile_m = min(128, cute.size(tile_shape_mnk, mode=[0])) + tile_n = min(32, cute.size(tile_shape_mnk, mode=[1])) + return (tile_m, tile_n) + else: + n_perf = 64 if element_type.width == 8 else 32 + tile_m = min(64, cute.size(tile_shape_mnk, mode=[0])) + tile_n = min(n_perf, cute.size(tile_shape_mnk, mode=[1])) + return (tile_m, tile_n) + + @staticmethod + def _make_smem_layouts( + tile_shape_mnk: tuple[int, int, int], + epi_tile: tuple[int, int], + a_dtype: type[cutlass.Numeric], + a_layout: utils.LayoutEnum, + b_dtype: type[cutlass.Numeric], + b_layout: utils.LayoutEnum, + ab_stage: int, + c_dtype: type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + epi_stage: int, + ) -> tuple[cute.ComposedLayout, cute.ComposedLayout, cute.ComposedLayout]: + """Create shared memory layouts for A, B, and C tensors. + + :param tile_shape_mnk: CTA tile shape (M,N,K) + :type tile_shape_mnk: Tuple[int, int, int] + :param epi_tile: Epilogue tile shape + :type epi_tile: Tuple[int, int] + :param a_dtype: Data type for matrix A + :type a_dtype: type[cutlass.Numeric] + :param a_layout: Layout enum for matrix A + :type a_layout: utils.LayoutEnum + :param b_dtype: Data type for matrix B + :type b_dtype: type[cutlass.Numeric] + :param b_layout: Layout enum for matrix B + :type b_layout: utils.LayoutEnum + :param ab_stage: Number of stages for A/B tensors + :type ab_stage: int + :param c_dtype: Data type for output matrix C + :type c_dtype: type[cutlass.Numeric] + :param c_layout: Layout enum for the output matrix C + :type c_layout: utils.LayoutEnum + :param epi_stage: Number of epilogue stages + :type epi_stage: int + + :return: Tuple of shared memory layouts for A, B, and C + :rtype: Tuple[cute.ComposedLayout, cute.ComposedLayout, cute.ComposedLayout] + """ + a_smem_shape = cute.slice_(tile_shape_mnk, (None, 0, None)) + + a_is_k_major = ( + a_layout.sm90_mma_major_mode() == cute.nvgpu.common.OperandMajorMode.K + ) + b_is_k_major = ( + b_layout.sm90_mma_major_mode() == cute.nvgpu.common.OperandMajorMode.K + ) + a_major_mode_size = tile_shape_mnk[2 if a_is_k_major else 0] + a_smem_layout_atom = cute.nvgpu.warpgroup.make_smem_layout_atom( + sm90_utils.get_smem_layout_atom( + a_layout, + a_dtype, + a_major_mode_size, + ), + a_dtype, + ) + a_smem_layout_staged = cute.tile_to_shape( + a_smem_layout_atom, + cute.append(a_smem_shape, ab_stage), + order=(0, 1, 2) if a_is_k_major else (1, 0, 2), + ) + + b_smem_shape = cute.slice_(tile_shape_mnk, (0, None, None)) + + b_major_mode_size = tile_shape_mnk[2 if b_is_k_major else 1] + b_smem_layout_atom = cute.nvgpu.warpgroup.make_smem_layout_atom( + sm90_utils.get_smem_layout_atom( + b_layout, + b_dtype, + b_major_mode_size, + ), + b_dtype, + ) + b_smem_layout_staged = cute.tile_to_shape( + b_smem_layout_atom, + cute.append(b_smem_shape, ab_stage), + order=(0, 1, 2) if b_is_k_major else (1, 0, 2), + ) + + c_smem_shape = epi_tile + c_major_mode_size = epi_tile[1] if c_layout.is_n_major_c() else epi_tile[0] + c_smem_layout_atom = cute.nvgpu.warpgroup.make_smem_layout_atom( + sm90_utils.get_smem_layout_atom( + c_layout, + c_dtype, + c_major_mode_size, + ), + c_dtype, + ) + epi_smem_layout_staged = cute.tile_to_shape( + c_smem_layout_atom, + cute.append(c_smem_shape, epi_stage), + order=(1, 0, 2) if c_layout.is_m_major_c() else (0, 1, 2), + ) + + return a_smem_layout_staged, b_smem_layout_staged, epi_smem_layout_staged + + @staticmethod + def _compute_grid( + c: cute.Tensor, + tile_shape_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + swizzle_size: int, + raster_along_m: bool, + max_active_clusters: cutlass.Constexpr, + ) -> tuple[int, int, int]: + """Compute grid shape for the output tensor C. + + :param c: The output tensor C + :type c: cute.Tensor + :param tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type tile_shape_mnk: tuple[int, int, int] + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + :param max_active_clusters: Maximum number of active clusters. + :type max_active_clusters: cutlass.Constexpr + + :return: Grid shape for kernel launch. + :rtype: tuple[int, int, int] + """ + c_shape = cute.slice_(tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, + cluster_shape_mnl, + swizzle_size, + raster_along_m, + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + return tile_sched_params, grid + + @staticmethod + def _make_tma_store_atoms_and_tensors( + tensor_c: cute.Tensor, + epi_smem_layout_staged: cute.ComposedLayout, + epi_tile: tuple[int, int], + ) -> tuple[cute.CopyAtom, cute.Tensor]: + """Create TMA atoms and tensors for C tensor storage. + + :param tensor_c: Output tensor C + :type tensor_c: cute.Tensor + :param epi_smem_layout_staged: Shared memory layout for epilogue + :type epi_smem_layout_staged: cute.ComposedLayout + :param epi_tile: Epilogue tile shape + :type epi_tile: Tuple[int, int] + + :return: TMA atom and tensor for C + :rtype: Tuple[cute.CopyAtom, cute.Tensor] + """ + epi_smem_layout = cute.slice_(epi_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cute.nvgpu.cpasync.make_tiled_tma_atom( + cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp(), + tensor_c, + epi_smem_layout, + epi_tile, + ) + + return tma_atom_c, tma_tensor_c + + @staticmethod + def _make_tma_atoms_and_tensors( + tensor: cute.Tensor, + smem_layout_staged: cute.ComposedLayout, + smem_tile: tuple[int, int], + mcast_dim: int, + ) -> tuple[cute.CopyAtom, cute.Tensor]: + """Create TMA atoms and tensors for input tensors. + + :param tensor: Input tensor (A or B) + :type tensor: cute.Tensor + :param smem_layout_staged: Shared memory layout for the tensor + :type smem_layout_staged: cute.ComposedLayout + :param smem_tile: Shared memory tile shape + :type smem_tile: Tuple[int, int] + :param mcast_dim: Multicast dimension + :type mcast_dim: int + + :return: TMA atom and tensor + :rtype: Tuple[cute.CopyAtom, cute.Tensor] + """ + op = ( + cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() + if mcast_dim == 1 + else cute.nvgpu.cpasync.CopyBulkTensorTileG2SMulticastOp() + ) + + smem_layout = cute.slice_(smem_layout_staged, (None, None, 0)) + tma_atom, tma_tensor = cute.nvgpu.cpasync.make_tiled_tma_atom( + op, + tensor, + smem_layout, + smem_tile, + num_multicast=mcast_dim, + ) + return tma_atom, tma_tensor + + # ========================================================================= + # Validation utilities + # ========================================================================= + + @staticmethod + def is_valid_dtypes( + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + acc_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + a_major: str, + b_major: str, + ) -> bool: + """Check if the dtypes are valid + + :param a_dtype: The data type of tensor A + :type a_dtype: type[cutlass.Numeric] + :param b_dtype: The data type of tensor B + :type b_dtype: type[cutlass.Numeric] + :param acc_dtype: The data type of the accumulator + :type acc_dtype: type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: type[cutlass.Numeric] + :param a_major: major mode of tensor A + :type a_major: str + :param b_major: major mode of tensor B + :type b_major: str + + :return: True if the dtypes are valid, False otherwise + :rtype: bool + """ + is_valid = True + + valid_ab_dtypes = { + cutlass.Float16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Uint8, + cutlass.Int8, + } + if a_dtype not in valid_ab_dtypes: + is_valid = False + if b_dtype not in valid_ab_dtypes: + is_valid = False + + # make sure a_dtype == b_dtype for Float16 + if a_dtype.width == 16 and a_dtype != b_dtype: + is_valid = False + if a_dtype.width != b_dtype.width: + is_valid = False + if not a_dtype.is_same_kind(b_dtype): + is_valid = False + + # for 8-bit types, this implementation only supports k-major layout + if (a_dtype.width == 8 and a_major != "k") or ( + b_dtype.width == 8 and b_major != "k" + ): + is_valid = False + + # Define compatibility mapping between accumulator type and AB type + acc_ab_compatibility = { + cutlass.Float32: { + cutlass.Float16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }, + cutlass.Float16: { + cutlass.Float16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }, + cutlass.Int32: {cutlass.Uint8, cutlass.Int8}, + } + # Check compatibility between accumulator type and A type + if acc_dtype not in acc_ab_compatibility or a_dtype not in acc_ab_compatibility[acc_dtype]: + is_valid = False + + # Define compatibility mapping between accumulator type and C type + acc_c_compatibility = { + cutlass.Float32: { + cutlass.Float32, + cutlass.Float16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }, + cutlass.Float16: { + cutlass.Float32, + cutlass.Float16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }, + cutlass.Int32: { + cutlass.Float32, + cutlass.Float16, + cutlass.Int32, + cutlass.Int8, + cutlass.Uint8, + }, + } + # Check compatibility between accumulator type and C type + if c_dtype not in acc_c_compatibility.get(acc_dtype, set()): + is_valid = False + + return is_valid + + @staticmethod + def is_valid_tensor_alignment( + m: int, + n: int, + k: int, + l: int, + ab_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """Check if the tensor alignment is valid + + :param m: The number of rows in the A tensor + :type m: int + :param n: The number of columns in the B tensor + :type n: int + :param k: The number of columns in the A tensor + :type k: int + :param l: The number of columns in the C tensor + :type l: int + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: type[cutlass.Numeric] + :param a_major: The major axis of the A tensor + :type a_major: str + :param b_major: The major axis of the B tensor + :type b_major: str + :param c_major: The major axis of the C tensor + :type c_major: str + + :return: True if the problem shape is valid, False otherwise + :rtype: bool + """ + is_valid = True + + def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // dtype.width + return num_major_elements % num_contiguous_elements == 0 + + if ( + not check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) + or not check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) + or not check_contigous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + ): + is_valid = False + return is_valid diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/sm100_contiguous_offset_2d3d_dense_gemm.py b/operators/cutlass/operators/providers/cutedsl/gemm/sm100_contiguous_offset_2d3d_dense_gemm.py new file mode 100644 index 000000000..4a1022d2c --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/sm100_contiguous_offset_2d3d_dense_gemm.py @@ -0,0 +1,451 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +from __future__ import annotations + +import itertools +from typing import TYPE_CHECKING + +import cutlass + +from cutlass.operators.arch import TargetSm +from cutlass.operators.arguments import ( + DenseTensor, + EpilogueArguments, + GroupedGemmArguments, +) +from cutlass.operators.metadata import ( + DenseTensorConstraints, + DesignMetadata, + EpilogueMetadata, + GroupedGemmOperandsMetadata, + OperatorMetadata, + Sm100DesignMetadata, +) +from cutlass.operators.mma import BlackwellTcgen05Mma +from cutlass.operators.providers.cutedsl import CuTeDSLProvider +from cutlass.operators.providers.cutedsl.integration_utils.mma import ( + get_max_active_clusters, +) +from cutlass.operators.providers.cutedsl.operator import CuteDslOperator +from cutlass.operators.status import Status +from cutlass.operators.utils.common import tuple_to_string +from cutlass.operators.utils.device import to_cuda_stream +from cutlass.operators.utils.gemm import ( + tensor_output_shape_matches, + tensor_reduction_mode_matches, +) +from cutlass.operators.utils.tensor import ( + normalized_major_stride, + strides_to_layout_string, +) + +from .implementations.sm100_contiguous_offset_2d3d_dense_gemm_impl import ( + ContiguousOffset2D3DGemmDenseKernel, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Generator + + from cutlass.operators.artifact import CompiledArtifact + + +@CuTeDSLProvider.register +class ContiguousOffset2D3DGemmDenseOperator(CuteDslOperator): + supported_args_type = GroupedGemmArguments + designed_for_min_cc = 100 + + def __init__(self, metadata: OperatorMetadata): + super().__init__(metadata) + mma_tiler_mn = (metadata.design.tile_shape[0], metadata.design.tile_shape[1]) + cluster_shape_mn = ( + metadata.design.cluster_shape[0], + metadata.design.cluster_shape[1], + ) + self.impl = ContiguousOffset2D3DGemmDenseKernel( + metadata.operands.accumulator_type, + metadata.design.use_2cta_mma, + mma_tiler_mn, + cluster_shape_mn, + ) + + def _supports( + self, args: GroupedGemmArguments, target_sm: TargetSm | None = None + ) -> Status: + if not (status := tensor_reduction_mode_matches(args)): + return status + + if not (status := tensor_output_shape_matches(args)): + return status + + if len(args.A.shape) == 3 and args.A.shape[0] != 1: + return Status.fail("Operand A must have batch size 1.") + + if len(args.out.shape) == 3 and args.out.shape[0] != 1: + return Status.fail("out must have batch size 1.") + + ValidM, N = args.out.shape[-2:] + K = args.A.shape[-1] + group_count = args.B.shape[0] if len(args.B.shape) == 3 else 1 + + if args.A.shape[-2] != ValidM or args.A.shape[-1] != K: + return Status.fail(f"A must have shape ({ValidM}, {K}). Got {args.A.shape}") + if args.B.shape != (group_count, K, N): + return Status.fail( + f"B must have shape ({group_count}, {K}, {N}). Got {args.B.shape}" + ) + if args.out.shape[-2] != ValidM or args.out.shape[-1] != N: + return Status.fail( + f"out must have shape ({ValidM}, {N}). Got {args.out.shape}" + ) + + # + # Check offsets + # + if len(args.offsets.shape) != 1: + return Status.fail(f"offsets must be a 1D tensor. Got {args.offsets.shape}") + + if args.offsets.numel() != group_count: + return Status.fail( + f"offsets must have {group_count} elements when offset mode is End and group count is {group_count}." + ) + + return Status.success() + + def _compile( + self, args: GroupedGemmArguments, target_sm: TargetSm | None = None + ) -> CompiledArtifact: + stream = cutlass.cute.runtime.make_fake_stream() + max_active_clusters = get_max_active_clusters(self.impl.cluster_shape_mn) + + return self.cute_compile( + self.impl, + args.A.tensor, + args.B.tensor, + args.out.tensor, + args.offsets.tensor, + max_active_clusters, + stream, + target_sm=target_sm, + ) + + def _run( + self, + args: GroupedGemmArguments, + compiled_artifact: CompiledArtifact, + stream, + workspace=None, + ) -> None: + stream = to_cuda_stream(stream) + compiled_gemm = compiled_artifact.compiled_obj + self.cute_run( + compiled_gemm, + args.A.tensor, + args.B.tensor, + args.out.tensor, + args.offsets.tensor, + stream, + ) + + @staticmethod + def _valid_operands(operands: GroupedGemmOperandsMetadata) -> bool: + if operands.accumulator_type != cutlass.Float32: + return False + + if operands.A.stride[-2:].index(1) != 1: + # A must be k-major + return False + if operands.B.stride[-2:].index(1) != 0: + # B must be n-major + return False + if operands.out.stride[-2:].index(1) != 1: + # out must be n-major + return False + + return True + + @staticmethod + def _valid_design_metadata(design: DesignMetadata) -> bool: + use_2cta_mma = design.use_2cta_mma + mma_tiler_mn = design.tile_shape[:2] + cluster_shape_mn = design.cluster_shape[:2] + impl = ContiguousOffset2D3DGemmDenseKernel + return impl.is_valid_mma_tiler_and_cluster_shape( + use_2cta_mma, mma_tiler_mn, cluster_shape_mn + ) + + @staticmethod + def _valid_epilogue_metadata(epilogue: EpilogueMetadata | None) -> bool: + return epilogue is None + + @staticmethod + def _valid_metadata(metadata: OperatorMetadata) -> bool: + if not ContiguousOffset2D3DGemmDenseOperator._valid_operands( + metadata.operands + ): + return False + if not ContiguousOffset2D3DGemmDenseOperator._valid_design_metadata( + metadata.design + ): + return False + if not ContiguousOffset2D3DGemmDenseOperator._valid_epilogue_metadata( + metadata.epilogue + ): + return False + return True + + @staticmethod + def _construct_divisible_operands( + a_dtype: cutlass.Numeric, + b_dtype: cutlass.Numeric, + out_dtype: cutlass.Numeric, + acc_dtype: cutlass.Numeric, + stride_A: tuple[int, ...], + stride_B: tuple[int, ...], + stride_out: tuple[int, ...], + ) -> GroupedGemmOperandsMetadata: + """Constructs a GroupedGemmOperandsMetadata object based on the provided data types and strides and + with divisibility constraints set based on those required by the kernel. + """ + alignment_bytes = 16 + offsets_alignment_bytes = 4 + a_divisibility = alignment_bytes * 8 // a_dtype.width + b_divisibility = alignment_bytes * 8 // b_dtype.width + out_divisibility = alignment_bytes * 8 // out_dtype.width + offsets_divisibility = offsets_alignment_bytes * 8 // cutlass.Int32.width + + return GroupedGemmOperandsMetadata( + A=DenseTensorConstraints( + dtype=a_dtype, + stride=stride_A, + divisibility=a_divisibility, + ), + B=DenseTensorConstraints( + dtype=b_dtype, + stride=stride_B, + divisibility=b_divisibility, + ), + out=DenseTensorConstraints( + dtype=out_dtype, + stride=stride_out, + divisibility=out_divisibility, + ), + offsets=DenseTensorConstraints( + dtype=cutlass.Int32, + stride=(1,), + divisibility=offsets_divisibility, + ), + accumulator_type=acc_dtype, + ) + + @staticmethod + def _metadata_operands_from_args( + args: GroupedGemmArguments, + ) -> Generator[GroupedGemmOperandsMetadata, None, None]: + """Generator that yields all valid GroupedGemmOperandsMetadata that correspond to `args` and + that are supported by this Operator. + + Args: + args (GroupedGemmArguments): The arguments to generate operands from. + + Returns: + Generator[GroupedGemmOperandsMetadata, None, None]: A generator of valid GroupedGemmOperandsMetadata that correspond to `args` and + that are supported by this Operator. + """ + if not all( + isinstance(arg, DenseTensor) + for arg in [args.A, args.B, args.out, args.offsets] + ): + return + + if args.offsets.dtype != cutlass.Int32: + return + + exp_A_stride = (0, 0, 1) + exp_B_stride = (0, 1, 0) + exp_out_stride = (0, 0, 1) + + A_stride = normalized_major_stride( + args.A.shape, args.A.stride, prepend_zeros_to_rank=3 + ) + B_stride = normalized_major_stride( + args.B.shape, args.B.stride, prepend_zeros_to_rank=3 + ) + out_stride = normalized_major_stride( + args.out.shape, args.out.stride, prepend_zeros_to_rank=3 + ) + + if A_stride != exp_A_stride: + return + if B_stride != exp_B_stride: + return + if out_stride != exp_out_stride: + return + + operands = ContiguousOffset2D3DGemmDenseOperator._construct_divisible_operands( + args.A.dtype, + args.B.dtype, + args.out.dtype, + args.accumulator_type, + exp_A_stride, + exp_B_stride, + exp_out_stride, + ) + + if not ContiguousOffset2D3DGemmDenseOperator._valid_operands(operands): + return + + yield operands + + @staticmethod + def _metadata_operand_combinations() -> Generator[ + tuple[GroupedGemmOperandsMetadata, int], None, None + ]: + """Generator that yields all valid (GroupedGemmOperandsMetadata, sf_vec_size) combinations + based on the validation rules in _valid_operands. + """ + # Supported A/B data types (must be the same) + ab_dtypes = [ + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + cutlass.Float16, + cutlass.BFloat16, + ] + + out_dtypes = [ + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + ] + + row_major_stride = (0, 0, 1) + col_major_stride = (0, 1, 0) + + stride_A = row_major_stride + stride_B = col_major_stride + stride_out = row_major_stride + + acc_dtype = cutlass.Float32 + Impl = ContiguousOffset2D3DGemmDenseKernel + + for ab_dtype, out_dtype in itertools.product(ab_dtypes, out_dtypes): + if not Impl.is_valid_dtypes(ab_dtype, acc_dtype, out_dtype): + continue + + yield ContiguousOffset2D3DGemmDenseOperator._construct_divisible_operands( + ab_dtype, + ab_dtype, + out_dtype, + acc_dtype, + stride_A, + stride_B, + stride_out, + ) + + @classmethod + def _generate_operators( + cls, + metadata_filter: Callable[[OperatorMetadata], bool], + epilogue_args: EpilogueArguments = None, + target_sm: TargetSm | None = None, + args: GroupedGemmArguments = None, + ) -> list[ContiguousOffset2D3DGemmDenseOperator]: + """Returns a list of all possible configurations of ContiguousOffset2D3DGemmDenseOperator that + adhere to constraints passed in under kwargs. + """ + design_params = { + "mma_instruction_type": [BlackwellTcgen05Mma], + "use_2cta_mma": [True, False], + # Fill in tile shape K based on data type + "tile_shape": [ + (64, 128, None), + (128, 128, None), + (256, 128, None), + (256, 256, None), + ], + "cluster_shape": [(M, N, 1) for M in [1, 2, 4] for N in [1, 2, 4]], + "use_tma_store": [True], + } + + if epilogue_args is not None: + return [] + + from itertools import product + + # Get the list of tunable parameter names and their possible values + param_names = list(design_params.keys()) + + operator_list = [] + + if args is not None: + operands_generator = cls._metadata_operands_from_args(args) + else: + operands_generator = cls._metadata_operand_combinations() + + for operands in operands_generator: + if not cls._valid_operands(operands): + continue + + tile_k = BlackwellTcgen05Mma.shape_k(operands) * 4 + design_params["tile_shape"] = [ + (M, N, tile_k) for M, N, _ in design_params["tile_shape"] + ] + param_values = [design_params[name] for name in param_names] + for values in product(*param_values): + design = Sm100DesignMetadata(**dict(zip(param_names, values))) + operator_name = "cutedsl.ContiguousOffset2D3DGemmDenseOperator_sm100_{layout}_A{A}_B{B}_out{out}_acc{acc}_{num_cta}cta_cluster{cluster}_tile{tile}{_tma_store}".format( + layout=strides_to_layout_string( + operands.A.stride, operands.B.stride, operands.out.stride + ), + A=operands.A.dtype, + B=operands.B.dtype, + out=operands.out.dtype, + acc=operands.accumulator_type, + num_cta=("2" if design.use_2cta_mma else "1"), + cluster=tuple_to_string(design.cluster_shape), + tile=tuple_to_string(design.tile_shape), + _tma_store="_tma_store" if design.use_tma_store else "", + ) + + metadata = OperatorMetadata( + operands=operands, + design=design, + operator_name=operator_name, + operator_class=cls, + supported_targets=TargetSm.get_supported_targets(design, operands), + epilogue=None, + ) + + metadata_valid = ContiguousOffset2D3DGemmDenseOperator._valid_metadata( + metadata + ) + if metadata_valid and metadata_filter(metadata): + operator_list.append( + ContiguousOffset2D3DGemmDenseOperator(metadata) + ) + + return operator_list diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/sm100_dense_blockscaled_static_persistent.py b/operators/cutlass/operators/providers/cutedsl/gemm/sm100_dense_blockscaled_static_persistent.py new file mode 100644 index 000000000..74f8cd4c6 --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/sm100_dense_blockscaled_static_persistent.py @@ -0,0 +1,606 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +import itertools +from typing import TYPE_CHECKING + +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +from cutlass.cute.nvgpu.common import OperandMajorMode + +from cutlass.operators.arch import TargetSm +from cutlass.operators.arguments import ( + DenseTensor, + EpilogueArguments, + GemmArguments, + ScaledOperand, + ScaleMode, + ScaleSwizzleMode, +) +from cutlass.operators.metadata import ( + DenseTensorConstraints, + DesignMetadata, + EpilogueMetadata, + GemmOperandsMetadata, + OperatorMetadata, + ScaledOperandConstraints, + Sm100DesignMetadata, +) +from cutlass.operators.mma import BlackwellTcgen05Mma +from cutlass.operators.providers.cutedsl import CuTeDSLProvider +from cutlass.operators.providers.cutedsl import integration_utils as cutedsl_utils +from cutlass.operators.providers.cutedsl.operator import CuteDslOperator +from cutlass.operators.status import Status +from cutlass.operators.utils.common import tuple_to_string +from cutlass.operators.utils.device import to_cuda_stream +from cutlass.operators.utils.tensor import ( + normalized_major_stride, + strides_to_layout_string, +) + +from .implementations.sm100_dense_blockscaled_static_persistent_impl import ( + PersistentDenseBlockScaledGemmKernel, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Generator + + from cutlass.operators.artifact import CompiledArtifact + + +@CuTeDSLProvider.register +class PersistentDenseBlockScaledGemmOperator(CuteDslOperator): + supported_args_type = GemmArguments + designed_for_min_cc = 100 + + def __init__(self, metadata: OperatorMetadata): + super().__init__(metadata) + self.sf_vec_size = ScaleMode.numel(metadata.operands.A.mode) + mma_tiler_mn = (metadata.design.tile_shape[0], metadata.design.tile_shape[1]) + cluster_shape_mn = ( + metadata.design.cluster_shape[0], + metadata.design.cluster_shape[1], + ) + self.impl = PersistentDenseBlockScaledGemmKernel( + self.sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, + ) + + def _supports( + self, args: GemmArguments, target_sm: TargetSm | None = None + ) -> Status: + M, N = args.out.shape[-2:] + K = args.A.shape[-1] + L = args.A.shape[0] if len(args.A.shape) == 3 else 1 + + # Shapes of scale factor tensors are not enforced. It is expected that the + # data underlying the passed in scale factor tensor is in the layout named + # by the operand's swizzle. We check only that the number of elements is + # correct for that layout, given the operand's logical shape. + expected_sfa_elements = ScaledOperand.numel_scale( + (L, M, K), args.A.mode, args.A.swizzle + ) + expected_sfb_elements = ScaledOperand.numel_scale( + (L, N, K), args.B.mode, args.B.swizzle + ) + + if args.A.scale.numel() != expected_sfa_elements: + return Status.fail( + f"Scale factor A for tensor A of shape {args.A.shape} must have " + f"{expected_sfa_elements} elements. Scale factor A is of shape {args.A.scale.shape} " + f"and has {args.A.scale.numel()} elements." + ) + if args.B.scale.numel() != expected_sfb_elements: + return Status.fail( + f"Scale factor B for tensor B of shape {args.B.shape} must have " + f"{expected_sfb_elements} elements. Scale factor B is of shape {args.B.scale.shape} " + f"and has {args.B.scale.numel()} elements." + ) + + return Status.success() + + def _construct_pointers( + self, args: GemmArguments, nullptr: bool = False + ) -> tuple[cute.Pointer, cute.Pointer, cute.Pointer, cute.Pointer, cute.Pointer]: + gmem = cutlass.AddressSpace.gmem + a_ptr = cutedsl_utils.builders.make_ptr( + args.A.tensor, nullptr=nullptr, address_space=gmem, assumed_align=16 + ) + b_ptr = cutedsl_utils.builders.make_ptr( + args.B.tensor, nullptr=nullptr, address_space=gmem, assumed_align=16 + ) + out_ptr = cutedsl_utils.builders.make_ptr( + args.out.tensor, nullptr=nullptr, address_space=gmem, assumed_align=16 + ) + sfa_ptr = cutedsl_utils.builders.make_ptr( + args.A.scale, nullptr=nullptr, address_space=gmem, assumed_align=32 + ) + sfb_ptr = cutedsl_utils.builders.make_ptr( + args.B.scale, nullptr=nullptr, address_space=gmem, assumed_align=32 + ) + return a_ptr, b_ptr, out_ptr, sfa_ptr, sfb_ptr + + @staticmethod + def _major_modes( + args: GemmArguments | GemmOperandsMetadata, + ) -> tuple[ + tuple[OperandMajorMode, str], + tuple[OperandMajorMode, str], + tuple[utils.LayoutEnum, str], + ]: + # A, B, and out can be of rank 2 or 3. Extract the final two dimensions + # to determine the major mode + a_major_mode = ( + (OperandMajorMode.K, "k") + if args.A.stride[-2:].index(1) == 1 + else (OperandMajorMode.MN, "m") + ) + + b_major_mode = ( + (OperandMajorMode.K, "k") + if args.B.stride[-2:].index(1) == 0 + else (OperandMajorMode.MN, "n") + ) + + if args.out.stride[-2:].index(1) == 1: + out_layout = (utils.LayoutEnum.ROW_MAJOR, "n") + else: + out_layout = (utils.LayoutEnum.COL_MAJOR, "m") + + return a_major_mode, b_major_mode, out_layout + + def _compile( + self, args: GemmArguments, target_sm: TargetSm | None = None + ) -> CompiledArtifact: + stream = cutlass.cute.runtime.make_fake_stream() + max_active_clusters = cutedsl_utils.mma.get_max_active_clusters( + self.impl.cluster_shape_mn + ) + + a_ptr, b_ptr, out_ptr, sfa_ptr, sfb_ptr = self._construct_pointers( + args, nullptr=True + ) + fake_problem_shape = ( + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + ) + + def epilogue_op(x): + return x + + (a_major_mode, _), (b_major_mode, _), (out_layout, _) = ( + PersistentDenseBlockScaledGemmOperator._major_modes(args) + ) + + return self.cute_compile( + self.impl, + a_ptr, + b_ptr, + sfa_ptr, + sfb_ptr, + out_ptr, + (a_major_mode, b_major_mode, out_layout), + fake_problem_shape, + max_active_clusters, + stream, + epilogue_op, + target_sm=target_sm, + ) + + def _run( + self, + args: GemmArguments, + compiled_artifact: CompiledArtifact, + stream, + workspace=None, + ) -> None: + stream = to_cuda_stream(stream) + compiled_gemm = compiled_artifact.compiled_obj + + M, N = args.out.shape[-2:] + K = args.A.shape[-1] + L = args.A.shape[0] if len(args.A.shape) == 3 else 1 + + a_ptr, b_ptr, out_ptr, sfa_ptr, sfb_ptr = self._construct_pointers(args) + self.cute_run( + compiled_gemm, a_ptr, b_ptr, sfa_ptr, sfb_ptr, out_ptr, (M, N, K, L), stream + ) + + @staticmethod + def _valid_operands(operands: GemmOperandsMetadata) -> bool: + if not isinstance(operands, GemmOperandsMetadata): + return False + + # A and B may have the same element type or form a mixed MXFP4 x MXFP8 + # pair (handled below by is_valid_dtypes_and_scale_factor_vec_size). + # The scale factors A and B must share the same dtype, scale mode, and + # swizzle mode. + if operands.A.scale.dtype != operands.B.scale.dtype: + return False + + if operands.A.mode != operands.B.mode or operands.A.mode not in [ + ScaleMode.Blockwise1x16, + ScaleMode.Blockwise1x32, + ]: + return False + + if ( + operands.A.swizzle != operands.B.swizzle + or operands.A.swizzle != ScaleSwizzleMode.Swizzle32x4x4 + ): + return False + + if operands.accumulator_type != cutlass.Float32: + return False + + a_dtype = operands.A.dtype + b_dtype = operands.B.dtype + sf_dtype = operands.A.scale.dtype + out_dtype = operands.out.dtype + + impl = PersistentDenseBlockScaledGemmKernel + if not impl.is_valid_dtypes_and_scale_factor_vec_size( + a_dtype, b_dtype, sf_dtype, ScaleMode.numel(operands.A.mode), out_dtype + ): + return False + + (_, a_major), (_, b_major), (_, out_major) = ( + PersistentDenseBlockScaledGemmOperator._major_modes(operands) + ) + if not impl.is_valid_layouts( + a_dtype, b_dtype, out_dtype, a_major, b_major, out_major + ): + return False + + return True + + @staticmethod + def _valid_design_metadata(design: DesignMetadata) -> bool: + if not isinstance(design, Sm100DesignMetadata): + return False + + mma_tiler_mn = design.tile_shape[:2] + cluster_shape_mn = design.cluster_shape[:2] + impl = PersistentDenseBlockScaledGemmKernel + return impl.is_valid_mma_tiler_and_cluster_shape(mma_tiler_mn, cluster_shape_mn) + + @staticmethod + def _valid_epilogue_metadata(epilogue: EpilogueMetadata | None) -> bool: + return epilogue is None + + @staticmethod + def _valid_metadata(metadata: OperatorMetadata) -> bool: + if not PersistentDenseBlockScaledGemmOperator._valid_operands( + metadata.operands + ): + return False + if not PersistentDenseBlockScaledGemmOperator._valid_design_metadata( + metadata.design + ): + return False + if not PersistentDenseBlockScaledGemmOperator._valid_epilogue_metadata( + metadata.epilogue + ): + return False + return True + + @staticmethod + def _construct_divisible_operands( + a_dtype: cutlass.Numeric, + b_dtype: cutlass.Numeric, + sfa_dtype: cutlass.Numeric, + sfb_dtype: cutlass.Numeric, + out_dtype: cutlass.Numeric, + acc_dtype: cutlass.Numeric, + sfa_scale_mode: ScaleMode, + sfb_scale_mode: ScaleMode, + sfa_swizzle: ScaleSwizzleMode, + sfb_swizzle: ScaleSwizzleMode, + stride_A: tuple[int, ...], + stride_B: tuple[int, ...], + stride_out: tuple[int, ...], + ) -> GemmOperandsMetadata: + """Constructs a GemmOperandsMetadata object based on the provided data types and strides and + with divisibility constraints set based on those required by the kernel. + """ + alignment_bytes = 16 + a_divisibility = alignment_bytes * 8 // a_dtype.width + b_divisibility = alignment_bytes * 8 // b_dtype.width + out_divisibility = alignment_bytes * 8 // out_dtype.width + + # Mixed pairs (e.g. MXFP4 x MXFP8) load the sub-byte operand via the UNPACK + # TMA, whose contiguous dim must be a multiple of 128 elements. + # Encoding it as divisibility makes discovery reject unsupported shapes + unpack_contiguous_divisibility = 128 + if PersistentDenseBlockScaledGemmKernel.needs_unpack_tma(a_dtype, b_dtype): + if a_dtype.width < 8: + a_divisibility = max(a_divisibility, unpack_contiguous_divisibility) + if b_dtype.width < 8: + b_divisibility = max(b_divisibility, unpack_contiguous_divisibility) + + sf_alignment_bytes = 32 + sfa_divisibility = sf_alignment_bytes * 8 // sfa_dtype.width + sfb_divisibility = sf_alignment_bytes * 8 // sfb_dtype.width + + return GemmOperandsMetadata( + A=ScaledOperandConstraints( + quantized=DenseTensorConstraints( + dtype=a_dtype, + stride=stride_A, + divisibility=a_divisibility, + ), + scale=DenseTensorConstraints( + dtype=sfa_dtype, + stride=None, + divisibility=sfa_divisibility, + ), + mode=sfa_scale_mode, + swizzle=sfa_swizzle, + ), + B=ScaledOperandConstraints( + quantized=DenseTensorConstraints( + dtype=b_dtype, + stride=stride_B, + divisibility=b_divisibility, + ), + scale=DenseTensorConstraints( + dtype=sfb_dtype, + stride=None, + divisibility=sfb_divisibility, + ), + mode=sfb_scale_mode, + swizzle=sfb_swizzle, + ), + out=DenseTensorConstraints( + dtype=out_dtype, + stride=stride_out, + divisibility=out_divisibility, + ), + accumulator_type=acc_dtype, + ) + + @staticmethod + def _metadata_operands_from_args( + args: GemmArguments, + ) -> Generator[GemmOperandsMetadata, None, None]: + """Generator that yields all valid GemmOperandsMetadata that correspond to `args` and + that are supported by this Operator. + + :param args: The arguments to generate operands from. + :type args: GemmArguments + + :return: A generator of valid GemmOperandsMetadata that correspond to `args` and + that are supported by this Operator. + :rtype: Generator[GemmOperandsMetadata, None, None] + """ + if any(not isinstance(x, ScaledOperand) for x in [args.A, args.B]): + return + if not isinstance(args.out, DenseTensor): + return + if any(not isinstance(x.scale, DenseTensor) for x in [args.A, args.B]): + return + + valid_major_strides = {(0, 0, 1), (0, 1, 0)} + + A_stride = normalized_major_stride( + args.A.shape, args.A.stride, prepend_zeros_to_rank=3 + ) + B_stride = normalized_major_stride( + args.B.shape, args.B.stride, prepend_zeros_to_rank=3 + ) + out_stride = normalized_major_stride( + args.out.shape, args.out.stride, prepend_zeros_to_rank=3 + ) + + if any(x not in valid_major_strides for x in [A_stride, B_stride, out_stride]): + return + + yield PersistentDenseBlockScaledGemmOperator._construct_divisible_operands( + args.A.dtype, + args.B.dtype, + args.A.scale.dtype, + args.B.scale.dtype, + args.out.dtype, + args.accumulator_type, + args.A.mode, + args.B.mode, + args.A.swizzle, + args.B.swizzle, + A_stride, + B_stride, + out_stride, + ) + + @staticmethod + def _metadata_operand_combinations() -> Generator[ + tuple[GemmOperandsMetadata, int], None, None + ]: + """Generator that yields all valid (GemmOperandsMetadata, sf_vec_size) combinations + based on the validation rules in _valid_operands. + """ + valid_dtypes = [cutlass.Float8E5M2, cutlass.Float8E4M3FN, cutlass.Float4E2M1FN] + + out_dtypes = [ + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ] + + sf_dtypes = [ + cutlass.Float8E8M0FNU, + cutlass.Float8E4M3FN, + ] + + scale_modes = [ScaleMode.Blockwise1x16, ScaleMode.Blockwise1x32] + + row_major_stride = (0, 0, 1) + col_major_stride = (0, 1, 0) + + def major_str_a(major_tuple: tuple[int, int, int]) -> str: + return "k" if major_tuple == row_major_stride else "m" + + def major_str_b(major_tuple: tuple[int, int, int]) -> str: + return "n" if major_tuple == row_major_stride else "k" + + def major_str_out(major_tuple: tuple[int, int, int]) -> str: + return "n" if major_tuple == row_major_stride else "m" + + Impl = PersistentDenseBlockScaledGemmKernel + + for a_dtype, b_dtype, sf_dtype, scale_mode, out_dtype in itertools.product( + valid_dtypes, valid_dtypes, sf_dtypes, scale_modes, out_dtypes + ): + sf_vec_size = scale_mode[-1] + if not Impl.is_valid_dtypes_and_scale_factor_vec_size( + a_dtype, b_dtype, sf_dtype, sf_vec_size, out_dtype + ): + continue + + for stride_A, stride_B, stride_out in itertools.product( + [row_major_stride, col_major_stride], repeat=3 + ): + a_major = major_str_a(stride_A) + b_major = major_str_b(stride_B) + out_major = major_str_out(stride_out) + + if not Impl.is_valid_layouts( + a_dtype, b_dtype, out_dtype, a_major, b_major, out_major + ): + continue + + yield PersistentDenseBlockScaledGemmOperator._construct_divisible_operands( + a_dtype, + b_dtype, + sf_dtype, + sf_dtype, + out_dtype, + cutlass.Float32, + scale_mode, + scale_mode, + ScaleSwizzleMode.Swizzle32x4x4, + ScaleSwizzleMode.Swizzle32x4x4, + stride_A, + stride_B, + stride_out, + ) + + @classmethod + def _generate_operators( + cls, + metadata_filter: Callable[[OperatorMetadata], bool], + epilogue_args: EpilogueArguments = None, + target_sm: TargetSm | None = None, + args: GemmArguments = None, + ) -> list[PersistentDenseBlockScaledGemmOperator]: + """Returns a list of all possible configurations of PersistentDenseBlockScaledGemmOperator that + adhere to constraints passed in under kwargs. + """ + design_params = { + "mma_instruction_type": [BlackwellTcgen05Mma], + "use_2cta_mma": [True], + "tile_shape": [ + # Fill in tile shape K based on data type + (M, N, None) + for M in [128, 256] + for N in [64, 128, 192] + ], + "cluster_shape": [(M, N, 1) for M in [2, 4] for N in [1, 2, 4]], + "use_tma_store": [True], + } + + if epilogue_args is not None: + return [] + + from itertools import product + + # Get the list of tunable parameter names and their possible values + param_names = list(design_params.keys()) + + operator_list = [] + + if args is not None: + operands_generator = cls._metadata_operands_from_args(args) + else: + operands_generator = cls._metadata_operand_combinations() + + for operands in operands_generator: + if not cls._valid_operands(operands): + continue + + tile_k = BlackwellTcgen05Mma.shape_k(operands) * 4 + design_params["tile_shape"] = [ + (M, N, tile_k) for M, N, _ in design_params["tile_shape"] + ] + param_values = [design_params[name] for name in param_names] + for values in product(*param_values): + design = Sm100DesignMetadata(**dict(zip(param_names, values))) + operator_name = "cutedsl.PersistentDenseBlockScaledGemmOperator_sm100_{layout}_A{A}_B{B}_out{out}_SFA{SFA}_SFB{SFB}_acc{acc}_scale{scale_mode}_swizzle{scale_swizzle}_{num_cta}cta_cluster{cluster}_tile{tile}{_tma_store}".format( + layout=strides_to_layout_string( + operands.A.stride, operands.B.stride, operands.out.stride + ), + A=operands.A.dtype, + B=operands.B.dtype, + out=operands.out.dtype, + SFA=operands.A.scale.dtype, + SFB=operands.B.scale.dtype, + acc=operands.accumulator_type, + scale_mode=operands.A.mode, + scale_swizzle=operands.A.swizzle, + num_cta=("2" if design.use_2cta_mma else "1"), + cluster=tuple_to_string(design.cluster_shape), + tile=tuple_to_string(design.tile_shape), + _tma_store="_tma_store" if design.use_tma_store else "", + ) + + metadata = OperatorMetadata( + operands=operands, + design=design, + operator_name=operator_name, + operator_class=cls, + supported_targets=TargetSm.get_supported_targets(design, operands), + epilogue=None, + ) + + metadata_valid = PersistentDenseBlockScaledGemmOperator._valid_metadata( + metadata + ) + + if metadata_valid and metadata_filter(metadata): + operator_list.append( + PersistentDenseBlockScaledGemmOperator(metadata) + ) + + return operator_list diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/sm100_mixed_input.py b/operators/cutlass/operators/providers/cutedsl/gemm/sm100_mixed_input.py new file mode 100644 index 000000000..ac60ba7f3 --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/sm100_mixed_input.py @@ -0,0 +1,910 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""Mixed-Input GEMM Operator for the cutlass.operators. + +This module provides the API wrapper for the mixed-input GEMM Operator, which supports +GEMM operations where input tensors A and B have different data types: +- A: Narrow precision (Int4, Int8, Uint8) +- B: Wide precision (BFloat16, Float16) + +Transform modes: +- Convert-Only: C = type_convert(A) @ B (for Int8/Uint8) +- Convert-Scale: C = (type_convert(A) * scale) @ B (for Int4) + +Usage with GemmArguments: +- Convert-only mode (Int8/Uint8): + GemmArguments(A=A_int8, B=B_bf16, out=out, accumulator_type=torch.float32) + +- Convert-scale mode (Int4): + GemmArguments( + A=ScaledOperand(A_int4, scale, mode=(1, 128), swizzle=ScaleSwizzleMode.SwizzleNone), + B=B_bf16, + out=out, + accumulator_type=torch.float32 + ) +""" + +from __future__ import annotations + +import itertools +from typing import TYPE_CHECKING + +import cutlass + +from cutlass.operators.arch import TargetSm +from cutlass.operators.arguments import ( + DenseTensor, + EpilogueArguments, + GemmArguments, + ScaledOperand, + ScaleSwizzleMode, +) +from cutlass.operators.metadata import ( + DenseTensorConstraints, + DesignMetadata, + GemmOperandsMetadata, + OperatorMetadata, + ScaledOperandConstraints, + Sm100DesignMetadata, +) +from cutlass.operators.mma import BlackwellTcgen05Mma +from cutlass.operators.providers.cutedsl import CuTeDSLProvider +from cutlass.operators.providers.cutedsl.integration_utils.mma import ( + get_max_active_clusters, +) +from cutlass.operators.providers.cutedsl.operator import CuteDslOperator +from cutlass.operators.status import Status +from cutlass.operators.utils.common import tuple_to_string +from cutlass.operators.utils.device import to_cuda_stream +from cutlass.operators.utils.tensor import ( + normalized_major_stride, + strides_to_layout_string, +) + +from .implementations.sm100_mixed_input_impl import ( + SM100_SMEM_ALIGN, + MixedInputGemmKernel, + compute_capacity_config, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Generator + + from cutlass.operators.artifact import CompiledArtifact + + +def _get_scale_granularity(args: GemmArguments) -> tuple[int, int]: + """Extract scale granularity from GemmArguments. + + For ScaledOperand A, the mode tuple is (scale_granularity_m, scale_granularity_k). + For DenseTensor A, returns (0, 0) for convert-only mode. + + :param args: GemmArguments + :return: Tuple of (scale_granularity_m, scale_granularity_k) + """ + if isinstance(args.A, ScaledOperand): + # mode is (scale_granularity_m, scale_granularity_k) + return (args.A.mode[0], args.A.mode[1]) + return (0, 0) + + +def _get_a_tensor(args: GemmArguments): + """Get the quantized A tensor from GemmArguments.""" + if isinstance(args.A, ScaledOperand): + return args.A.quantized.tensor + return args.A.tensor + + +def _get_scale_tensor(args: GemmArguments): + """Get the scale tensor from GemmArguments, or None if not present.""" + if isinstance(args.A, ScaledOperand): + return args.A.scale.tensor + return None + + +@CuTeDSLProvider.register +class MixedInputGemmOperator(CuteDslOperator): + """Mixed-Input GEMM Operator for NVIDIA Blackwell SM100 architecture. + + This Operator supports GEMM operations where input tensors A and B have different + data types, with tensor A being transformed to the precision of tensor B before + matrix multiplication. + + Supported A data types (narrow precision): + - Int4 (requires scale tensor via ScaledOperand) + - Int8 + - Uint8 + + Supported B data types (wide precision): + - BFloat16 + - Float16 + + Supported accumulator data types: + - Float32 (required when B is BFloat16) + - Float16 (only when B is Float16) + + Supported output data types: + - BFloat16 + - Float16 + - Float32 + + Transform modes: + - Convert-Only: C = type_convert(A) @ B (for Int8/Uint8 with DenseTensor A) + - Convert-Scale: C = (type_convert(A) * scale) @ B (for Int4 with ScaledOperand A) + + Constraints: + - MMA tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) + - MMA tiler N must be 32-256, step 32 + - Cluster shape M must be multiple of 2 if use_2cta_instrs=True + - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 + - scale_granularity_k must be a multiple of mma_tiler_k + """ + + supported_args_type = GemmArguments + designed_for_min_cc = 100 + + def __init__(self, metadata: OperatorMetadata): + """Initialize the mixed-input GEMM Operator with specified metadata. + + :param metadata: Operator metadata containing operand and design information + """ + super().__init__(metadata) + + mma_tiler_mnk = metadata.design.tile_shape # Full 3-tuple (M, N, K) + cluster_shape_mn = ( + metadata.design.cluster_shape[0], + metadata.design.cluster_shape[1], + ) + + # Extract scale granularity from metadata + scale_granularity_m, scale_granularity_k = self._extract_scale_granularity( + metadata.operands + ) + + shuffle_a = self._compute_shuffle_a(metadata, scale_granularity_k) + + self.impl = MixedInputGemmKernel( + scale_granularity_m=scale_granularity_m, + scale_granularity_k=scale_granularity_k, + acc_dtype=metadata.operands.accumulator_type, + use_2cta_instrs=metadata.design.use_2cta_mma, + mma_tiler_mnk=mma_tiler_mnk, + cluster_shape_mn=cluster_shape_mn, + use_tma_store=metadata.design.use_tma_store, + shuffle_a=shuffle_a, + ) + + self._shuffle_a = shuffle_a + + @staticmethod + def _extract_scale_granularity( + operands: GemmOperandsMetadata, + ) -> tuple[int, int]: + """Extract scale granularity from operands metadata. + + For mixed-input GEMM, mode is a 2-tuple: (scale_granularity_m, scale_granularity_k) + For DenseTensor A (convert-only mode), returns (0, 0). + """ + if isinstance(operands.A, ScaledOperandConstraints): + mode = getattr(operands.A, "mode", None) + if isinstance(mode, (tuple, list)) and len(mode) >= 2: + return (mode[0], mode[1]) + return (0, 0) + + @staticmethod + def _compute_shuffle_a( + metadata: OperatorMetadata, scale_granularity_k: int + ) -> bool: + """Determine if shuffle should be used for A tensor transformation. + + Shuffle optimization enables a fast Int4→BFloat16 conversion path using the + hardware intrinsic `cvt_i4_bf16_intrinsic` (see: cutlass.cute.arch.numeric_conversion). + + This intrinsic processes 8 Int4 elements at a time. When shuffle is enabled, + the input elements are expected in a specific pattern: indices (0,1,2,3,4,5,6,7) + are reordered to (0,2,1,3,4,6,5,7). This pre-shuffled layout allows the intrinsic + to produce BFloat16 output in natural order without extra permutation (prmt) + instructions, improving performance. + + Shuffle optimization requires: + 1. A tensor is K-major (row-major for M×K) - data layout matches shuffle pattern + 2. A dtype is Int4 - intrinsic is specific to Int4 + 3. B dtype is BFloat16 - intrinsic only outputs BFloat16 (no Float16 equivalent) + 4. scale_granularity_k >= 8 - sufficient granularity for shuffle pattern + + The intrinsic operates on 8-element groups, but the TMA 16-byte alignment + for Int4 already guarantees K % 32 == 0, which subsumes K % 8. + + See: cutlass.utils.mixed_input_helpers.is_shuffle_a() + cutlass.cute.arch.numeric_conversion.cvt_i4_bf16_intrinsic() + + :param metadata: Operator metadata containing operand information + :param scale_granularity_k: Scale granularity along K dimension + :return: True if shuffle should be used, False otherwise + """ + # Get A dtype and stride from metadata + if isinstance(metadata.operands.A, ScaledOperandConstraints): + a_dtype = metadata.operands.A.quantized.dtype + a_stride = metadata.operands.A.quantized.stride + else: + a_dtype = metadata.operands.A.dtype + a_stride = metadata.operands.A.stride + + b_dtype = metadata.operands.B.dtype + + # Check conditions that can be verified at Operator generation time: + # 1. A dtype is Int4 + # 2. B dtype is BFloat16 (shuffle only supported for BFloat16, not Float16) + # 3. A is K-major + # 4. scale_granularity_k >= 8 + if a_dtype == cutlass.Int4 and b_dtype == cutlass.BFloat16: + # Check if A is K-major (last dimension contiguous) + # For 3D stride (batch, row, col): + # (0, 0, 1) means K dimension is contiguous = K-major + # (0, 1, 0) means M dimension is contiguous = M-major + is_k_major = a_stride == (0, 0, 1) + + if is_k_major and scale_granularity_k >= 8: + return True + return False + + def _supports( + self, args: GemmArguments, target_sm: TargetSm | None = None + ) -> Status: + """Check if the Operator supports the given arguments. + + Checks here are for problem-size-dependent constraints: K divisibility + by scale_granularity_k and OOB tile alignment. Operand compatibility + (dtype, stride, divisibility) and epilogue rejection are handled by the + base class via metadata.supports(); scale_granularity_k/tile_k + divisibility is handled by _valid_metadata at generation time. + + :param args: The GEMM arguments to check + :return: Status indicating support + """ + # Get tensors for dimension checks + a_tensor = _get_a_tensor(args) + b_tensor = args.B.tensor + + # Extract dimensions (tensors are in batch-first format: L, M, K or L, K, N) + k = a_tensor.shape[-1] # K dimension from A + + tile_m, tile_n, tile_k = self.metadata.design.tile_shape + use_2cta = self.metadata.design.use_2cta_mma + + # Validate scale granularity constraints that depend on problem size + _, scale_granularity_k = _get_scale_granularity(args) + if scale_granularity_k > 0: + # k must be divisible by scale_granularity_k + if k % scale_granularity_k != 0: + return Status.fail( + f"K dimension ({k}) must be divisible by " + f"scale_granularity_k ({scale_granularity_k})" + ) + # scale_granularity_k must be divisible by mma_tiler_k + if scale_granularity_k % tile_k != 0: + return Status.fail( + f"scale_granularity_k ({scale_granularity_k}) must be " + f"divisible by mma_tiler_k ({tile_k})" + ) + + # Validate OOB tile constraint when TMA store is disabled + # Without TMA store, out-of-bounds tiles are not supported + if not self.metadata.design.use_tma_store: + m = a_tensor.shape[-2] # M dimension from A + n = b_tensor.shape[-1] # N dimension from B + cta_tile_m = tile_m // (2 if use_2cta else 1) + cta_tile_n = tile_n + if m % cta_tile_m != 0: + return Status.fail( + f"When TMA store is disabled, M dimension ({m}) must be " + f"divisible by CTA tile M ({cta_tile_m})" + ) + if n % cta_tile_n != 0: + return Status.fail( + f"When TMA store is disabled, N dimension ({n}) must be " + f"divisible by CTA tile N ({cta_tile_n})" + ) + + return Status.success() + + def _compile( + self, args: GemmArguments, target_sm: TargetSm | None = None + ) -> CompiledArtifact: + """Compile the Operator for the given arguments. + + :param args: The GEMM arguments + :param target_sm: Target compute capability (optional) + :return: Compiled artifact + """ + stream = cutlass.cute.runtime.make_fake_stream() + max_active_clusters = get_max_active_clusters(self.impl.cluster_shape_mn) + + a_tensor = _get_a_tensor(args) + scale_tensor = _get_scale_tensor(args) + + return self.cute_compile( + self.impl, + a_tensor, + scale_tensor, + args.B.tensor, + args.out.tensor, + max_active_clusters, + stream, + target_sm=target_sm, + ) + + def _run( + self, + args: GemmArguments, + compiled_artifact: CompiledArtifact, + stream, + workspace=None, + ) -> None: + """Execute the compiled Operator. + + :param args: The GEMM arguments + :param compiled_artifact: The compiled artifact + :param stream: CUDA stream for execution + :param workspace: Optional workspace (unused) + """ + stream = to_cuda_stream(stream) + compiled_gemm = compiled_artifact.compiled_obj + + a_tensor = _get_a_tensor(args) + scale_tensor = _get_scale_tensor(args) + + self.cute_run( + compiled_gemm, + a_tensor, + scale_tensor, + args.B.tensor, + args.out.tensor, + stream, + ) + + @staticmethod + def _valid_operands(operands: GemmOperandsMetadata) -> bool: + """Check if the operand metadata represents a valid mixed-input configuration. + + :param operands: The operand metadata to validate + :return: True if valid, False otherwise + """ + # Get A dtype + if isinstance(operands.A, ScaledOperandConstraints): + a_dtype = operands.A.quantized.dtype + is_scaled = True + elif isinstance(operands.A, DenseTensorConstraints): + a_dtype = operands.A.dtype + is_scaled = False + else: + # Plain OperandConstraints - treat as dense + a_dtype = operands.A.dtype + is_scaled = False + + # A must be narrow precision + narrow_types = [cutlass.Int4, cutlass.Int8, cutlass.Uint8] + if a_dtype not in narrow_types: + return False + + # Get B dtype + if isinstance(operands.B, ScaledOperandConstraints): + # Mixed-input doesn't support scaled B + return False + elif isinstance(operands.B, DenseTensorConstraints): + b_dtype = operands.B.dtype + else: + b_dtype = operands.B.dtype + + # B must be wide precision + wide_types = [cutlass.BFloat16, cutlass.Float16] + if b_dtype not in wide_types: + return False + + # Accumulator type constraints: + # - BFloat16 B: accumulator must be Float32 + # - Float16 B: accumulator can be Float16 or Float32 + if b_dtype == cutlass.BFloat16 and operands.accumulator_type != cutlass.Float32: + return False + if b_dtype == cutlass.Float16 and operands.accumulator_type not in [ + cutlass.Float16, + cutlass.Float32, + ]: + return False + + # Get output dtype + if isinstance(operands.out, ScaledOperandConstraints): + return False # Output shouldn't be scaled + elif isinstance(operands.out, DenseTensorConstraints): + out_dtype = operands.out.dtype + else: + out_dtype = operands.out.dtype + + # Output must be compatible + valid_out_types = [cutlass.BFloat16, cutlass.Float16, cutlass.Float32] + if out_dtype not in valid_out_types: + return False + + # Accumulator/output type compatibility: + # - Float16 accumulator can output to 16-bit types (Float16, BFloat16) + # - Float16 accumulator cannot output to Float32 (no acc_bits=16, d_bits=32 case) + # - Float32 accumulator can output to any supported type + # See: cutlass/utils/blackwell_helpers.py acc_bits/d_bits combinations + if operands.accumulator_type == cutlass.Float16 and out_dtype not in [ + cutlass.Float16, + cutlass.BFloat16, + ]: + return False + + # Validate scale requirements based on A dtype + a_bit_width = a_dtype.width + + if a_bit_width == 8: # Int8 or Uint8 + # 8-bit types use convert-only mode (no scale) + if is_scaled: + return False + elif a_bit_width == 4: # Int4 + # Int4 requires scale + if not is_scaled: + return False + # Apply constraints from mixed_input_helpers.is_valid_scale_granularity + scale_m, scale_k = MixedInputGemmOperator._extract_scale_granularity( + operands + ) + # Constraint: scale_granularity_m must be 1 (per-row scaling) + if scale_m != 1: + return False + # Constraint: scale_granularity_k must be > 0 + if scale_k <= 0: + return False + + return True + + @staticmethod + def _valid_design(design: DesignMetadata) -> bool: + """Check if the design metadata represents a valid configuration. + + Validates hardware constraints of the SM100 tcgen05 MMA instruction: + - Cluster M must be divisible by 2 in 2CTA mode (CTAs work in pairs) + - Per-CTA tile M must be 64 or 128 (supported tcgen05 MMA tile sizes) + + :param design: The design metadata to validate + :return: True if valid, False otherwise + """ + tile_m = design.tile_shape[0] + cluster_m = design.cluster_shape[0] + + cta_divisor = 2 if design.use_2cta_mma else 1 + if cluster_m % cta_divisor != 0: + return False + + cta_tile_m = tile_m // cta_divisor + if cta_tile_m not in [64, 128]: + return False + + return True + + @staticmethod + def _check_capacity(metadata: OperatorMetadata) -> bool: + """Check if the Operator configuration fits within SMEM/TMEM capacity. + + This uses the shared compute_capacity_config function which is also + used by the compile-time code, ensuring the checks are identical. + + :param metadata: The operator metadata to validate + :return: True if capacity is sufficient, False otherwise + """ + design = metadata.design + operands = metadata.operands + + tile_m, tile_n, tile_k = design.tile_shape + use_2cta = design.use_2cta_mma + cta_divisor = 2 if use_2cta else 1 + + # Get A tensor attributes + if isinstance(operands.A, ScaledOperandConstraints): + a_dtype = operands.A.quantized.dtype + a_stride = operands.A.quantized.stride + scale_mode_is_convert_only = False + mode = getattr(operands.A, "mode", None) + if isinstance(mode, (tuple, list)) and len(mode) >= 2: + scale_granularity_m = mode[0] + scale_granularity_k = mode[1] + else: + scale_granularity_m = 0 + scale_granularity_k = 0 + else: + a_dtype = operands.A.dtype + a_stride = operands.A.stride + scale_mode_is_convert_only = True + scale_granularity_m = 0 + scale_granularity_k = 0 + + # Determine if A uses TMEM (K-major) or SMEM (M-major) + is_k_major = a_stride == (0, 0, 1) + + # Compute capacity configuration using the shared function + config = compute_capacity_config( + cta_tile_m=tile_m // cta_divisor, + cta_tile_n=tile_n, + cta_tile_k=tile_k, + use_2cta=use_2cta, + a_dtype_width=a_dtype.width, + b_dtype_width=operands.B.dtype.width, + mma_dtype_width=operands.B.dtype.width, # MMA operates in B's precision + out_dtype_width=operands.out.dtype.width, + transform_a_source_is_tmem=is_k_major, + scale_mode_is_convert_only=scale_mode_is_convert_only, + scale_granularity_m=scale_granularity_m, + scale_granularity_k=scale_granularity_k, + use_tma_store=design.use_tma_store, + smem_align=SM100_SMEM_ALIGN, + ) + + return config.is_valid + + @staticmethod + def _valid_metadata(metadata: OperatorMetadata) -> bool: + """Check if the operator metadata is valid. + + :param metadata: The operator metadata to validate + :return: True if valid, False otherwise + """ + if not MixedInputGemmOperator._valid_operands(metadata.operands): + return False + if not MixedInputGemmOperator._valid_design(metadata.design): + return False + + # Cross-validate scale granularity against tile shape + scale_m, scale_k = MixedInputGemmOperator._extract_scale_granularity( + metadata.operands + ) + tile_k = metadata.design.tile_shape[2] + if scale_k > 0 and scale_k % tile_k != 0: + return False + + if not MixedInputGemmOperator._check_capacity(metadata): + return False + return True + + @staticmethod + def _metadata_operands_from_args( + args: GemmArguments, + ) -> Generator[GemmOperandsMetadata, None, None]: + """Generator that yields all valid GemmOperandsMetadata that correspond to `args` and + that are supported by this Operator. + + :param args: The arguments to generate operands from. + :type args: GemmArguments + + :return: A generator of valid GemmOperandsMetadata that correspond to `args` and + that are supported by this Operator. + :rtype: Generator[GemmOperandsMetadata, None, None] + """ + alignment_bytes = 16 + col_major_stride = (0, 1, 0) + valid_major_strides = {(0, 0, 1), (0, 1, 0)} + + # --- A operand --- + if isinstance(args.A, ScaledOperand): + a_quantized = args.A.quantized + a_dtype = a_quantized.dtype + A_stride = normalized_major_stride( + a_quantized.shape, a_quantized.stride, prepend_zeros_to_rank=3 + ) + if A_stride is None or A_stride not in valid_major_strides: + return + mode = getattr(args.A, "mode", None) + if not isinstance(mode, (tuple, list)) or len(mode) < 2: + return + scale_m, scale_k = mode[0], mode[1] + a_attr = ScaledOperandConstraints( + quantized=DenseTensorConstraints( + dtype=a_dtype, + stride=A_stride, + divisibility=alignment_bytes * 8 // a_dtype.width, + ), + scale=DenseTensorConstraints( + dtype=args.A.scale.dtype, + stride=col_major_stride, + divisibility=alignment_bytes * 8 // args.A.scale.dtype.width, + ), + mode=(scale_m, scale_k), + swizzle=args.A.swizzle, + ) + elif isinstance(args.A, DenseTensor): + a_dtype = args.A.dtype + A_stride = normalized_major_stride( + args.A.shape, args.A.stride, prepend_zeros_to_rank=3 + ) + if A_stride is None or A_stride not in valid_major_strides: + return + a_attr = DenseTensorConstraints( + dtype=a_dtype, + stride=A_stride, + divisibility=alignment_bytes * 8 // a_dtype.width, + ) + else: + return + + # --- B operand --- + if not isinstance(args.B, DenseTensor): + return + B_stride = normalized_major_stride( + args.B.shape, args.B.stride, prepend_zeros_to_rank=3 + ) + if B_stride is None or B_stride not in valid_major_strides: + return + b_attr = DenseTensorConstraints( + dtype=args.B.dtype, + stride=B_stride, + divisibility=alignment_bytes * 8 // args.B.dtype.width, + ) + + # --- Out operand --- + out_stride = normalized_major_stride( + args.out.shape, args.out.stride, prepend_zeros_to_rank=3 + ) + if out_stride is None or out_stride not in valid_major_strides: + return + out_attr = DenseTensorConstraints( + dtype=args.out.dtype, + stride=out_stride, + divisibility=alignment_bytes * 8 // args.out.dtype.width, + ) + + operands = GemmOperandsMetadata( + A=a_attr, B=b_attr, out=out_attr, accumulator_type=args.accumulator_type + ) + + if not MixedInputGemmOperator._valid_operands(operands): + return + + yield operands + + @staticmethod + def _metadata_operand_combinations() -> Generator[GemmOperandsMetadata, None, None]: + """Generate all valid operand metadata combinations for mixed-input GEMM. + + :yield: Valid GemmOperandsMetadata instances + """ + narrow_dtypes = [cutlass.Int4, cutlass.Int8, cutlass.Uint8] + wide_dtypes = [cutlass.BFloat16, cutlass.Float16] + out_dtypes = [cutlass.BFloat16, cutlass.Float16, cutlass.Float32] + + # Scale granularities: (m, k) + # (0, 0) = convert-only mode (for Int8/Uint8) + # (1, k) = per-row scaling with k elements sharing a scale along K (for Int4) + # + # Note: Unlike the original example which accepts user-specified values, + # the API must enumerate operators at discovery time, so we pre-define + # common configurations. These values must satisfy: + # - scale_granularity_k % mma_tiler_k == 0 (validated at runtime) + # - k % scale_granularity_k == 0 (validated at runtime) + # + # Common values are multiples of typical mma_tiler_k values (64, 128, 256). + # Users needing other values can extend this list or create custom operators. + scale_configs = [ + (0, 0), # convert-only + (1, 64), # per-row, 64 elements per scale along K + (1, 128), # per-row, 128 elements per scale along K + (1, 256), # per-row, 256 elements per scale along K + (1, 512), # per-row, 512 elements per scale along K + ] + + # Layout strides (rank-3): (batch, row, col) + row_major_stride = (0, 0, 1) + col_major_stride = (0, 1, 0) + + layout_strides = [ + (row_major_stride, row_major_stride, row_major_stride), + (row_major_stride, col_major_stride, row_major_stride), + (col_major_stride, row_major_stride, row_major_stride), + (col_major_stride, col_major_stride, row_major_stride), + ] + + alignment_bytes = 16 + + for a_dtype in narrow_dtypes: + for b_dtype in wide_dtypes: + for out_dtype in out_dtypes: + for scale_m, scale_k in scale_configs: + for a_stride, b_stride, out_stride in layout_strides: + # Calculate divisibility based on alignment + a_divisibility = alignment_bytes * 8 // a_dtype.width + b_divisibility = alignment_bytes * 8 // b_dtype.width + out_divisibility = alignment_bytes * 8 // out_dtype.width + + # Int4 requires scale (convert-scale mode) + if a_dtype == cutlass.Int4: + if scale_m == 0 or scale_k == 0: + continue + + # A is a ScaledOperand + a_attr = ScaledOperandConstraints( + quantized=DenseTensorConstraints( + dtype=a_dtype, + stride=a_stride, + divisibility=a_divisibility, + ), + scale=DenseTensorConstraints( + dtype=b_dtype, # Scale has same type as B + stride=col_major_stride, # M-major + divisibility=b_divisibility, + ), + mode=(scale_m, scale_k), + swizzle=ScaleSwizzleMode.SwizzleNone, + ) + else: + # Int8/Uint8 use convert-only (no scale) + if scale_m != 0 or scale_k != 0: + continue + + # A is a DenseTensor + a_attr = DenseTensorConstraints( + dtype=a_dtype, + stride=a_stride, + divisibility=a_divisibility, + ) + + # B is always a DenseTensor for mixed-input + b_attr = DenseTensorConstraints( + dtype=b_dtype, + stride=b_stride, + divisibility=b_divisibility, + ) + + # Output is always a DenseTensor + out_attr = DenseTensorConstraints( + dtype=out_dtype, + stride=out_stride, + divisibility=out_divisibility, + ) + + # Determine valid accumulator types based on B dtype + # and acc/output compatibility + if b_dtype == cutlass.BFloat16: + valid_acc_types = [cutlass.Float32] + else: # Float16 B + # Float16 acc can output to 16-bit types (Float16, BFloat16) + # Float32 acc can output to any type + if out_dtype in [cutlass.Float16, cutlass.BFloat16]: + valid_acc_types = [cutlass.Float16, cutlass.Float32] + else: # Float32 output + valid_acc_types = [cutlass.Float32] + + for acc_type in valid_acc_types: + yield GemmOperandsMetadata( + A=a_attr, + B=b_attr, + out=out_attr, + accumulator_type=acc_type, + ) + + @classmethod + def _generate_operators( + cls, + metadata_filter: Callable[[OperatorMetadata], bool], + epilogue_args: EpilogueArguments = None, + target_sm: TargetSm | None = None, + args: GemmArguments = None, + ) -> list[MixedInputGemmOperator]: + """Generate all valid MixedInputGemmOperator configurations. + + :param metadata_filter: Filter function for metadata + :param epilogue_args: Epilogue arguments (unused, returns empty if provided) + :param target_sm: Target compute capability + :return: List of valid Operator instances + """ + # No epilogue support for mixed-input operators currently + if epilogue_args is not None: + return [] + + design_params = { + "mma_instruction_type": [BlackwellTcgen05Mma], + "use_2cta_mma": [True, False], + "tile_shape": [ + (M, N, K) + for M in [64, 128, 256] + for N in [32, 64, 128, 256] + for K in [64, 128, 256] + ], + "cluster_shape": [(M, N, 1) for M in [1, 2] for N in [1, 2]], + "use_tma_store": [ + True + ], # Non-TMA store path disabled due to OOB tile limitations + } + + param_names = list(design_params.keys()) + param_values = [design_params[name] for name in param_names] + + operator_list = [] + + if args is not None: + operands_generator = cls._metadata_operands_from_args(args) + else: + operands_generator = cls._metadata_operand_combinations() + + for operands in operands_generator: + for values in itertools.product(*param_values): + design = Sm100DesignMetadata(**dict(zip(param_names, values))) + + # Extract scale mode for naming + scale_m, scale_k = cls._extract_scale_granularity(operands) + scale_mode = f"scale{scale_m}x{scale_k}" if scale_m > 0 else "convert" + + # Get A dtype for naming + if isinstance(operands.A, ScaledOperandConstraints): + a_dtype = operands.A.quantized.dtype + a_stride = operands.A.quantized.stride + else: + a_dtype = operands.A.dtype + a_stride = operands.A.stride + + # Get B dtype + if isinstance(operands.B, DenseTensorConstraints): + b_dtype = operands.B.dtype + b_stride = operands.B.stride + else: + b_dtype = operands.B.dtype + b_stride = operands.B.stride + + # Get out dtype + if isinstance(operands.out, DenseTensorConstraints): + out_dtype = operands.out.dtype + out_stride = operands.out.stride + else: + out_dtype = operands.out.dtype + out_stride = operands.out.stride + + layout = strides_to_layout_string(a_stride, b_stride, out_stride) + + operator_name = ( + f"cutedsl.MixedInputGemmOperator_sm100_{scale_mode}_{layout}" + f"A{a_dtype}_B{b_dtype}_out{out_dtype}_" + f"acc{operands.accumulator_type}_" + f"{'2cta' if design.use_2cta_mma else '1cta'}_" + f"cluster{tuple_to_string(design.cluster_shape)}_" + f"tile{tuple_to_string(design.tile_shape)}" + f"{'_tma_store' if design.use_tma_store else ''}" + ) + + metadata = OperatorMetadata( + operands=operands, + design=design, + operator_name=operator_name, + operator_class=cls, + supported_targets=TargetSm.get_supported_targets(design, operands), + epilogue=None, + ) + + if cls._valid_metadata(metadata) and metadata_filter(metadata): + operator_list.append(cls(metadata)) + + return operator_list diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/sm100_persistent.py b/operators/cutlass/operators/providers/cutedsl/gemm/sm100_persistent.py new file mode 100644 index 000000000..ed44bdaed --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/sm100_persistent.py @@ -0,0 +1,494 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +import itertools +from typing import TYPE_CHECKING + +import cutlass + +import cutlass.operators.utils.generate as generate +from cutlass.operators.arch import TargetSm +from cutlass.operators.arguments import ( + DenseTensor, + EpilogueArguments, + GemmArguments, +) +from cutlass.operators.metadata import ( + GemmOperandsMetadata, + OperatorMetadata, + Sm100DesignMetadata, +) +from cutlass.operators.metadata.design.tile_scheduler import ( + CLCDynamicPersistentTileSchedulerMetadata, + StaticPersistentTileSchedulerMetadata, +) +from cutlass.operators.mma import BlackwellTcgen05Mma +from cutlass.operators.providers.cutedsl import CuTeDSLProvider +from cutlass.operators.providers.cutedsl.integration_utils.mma import ( + get_max_active_clusters, +) +from cutlass.operators.providers.cutedsl.operator import CuteDslOperator +from cutlass.operators.utils.common import tuple_to_string +from cutlass.operators.utils.device import to_cuda_stream +from cutlass.operators.utils.tensor import ( + normalized_major_stride, + strides_to_layout_string, +) + +from .implementations.operator_helpers import ( + scheduler_metadata_to_enum, +) +from .implementations.sm100_persistent_impl import ( + PersistentDenseGemmKernel, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Generator + + from cutlass.operators.artifact import CompiledArtifact + + +@CuTeDSLProvider.register +class PersistentDenseGemmOperator(CuteDslOperator): + """This class implements batched matrix multiplication (C = A @ B) with support for various data types + and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization. + + :note: In current version, A and B tensor must have the same data type + - i.e., Float8E4M3FN for A and Float8E5M2 for B is not supported + + :note: Supported A/B data types: + - Float16/BFloat16 + - Int8/Uint8 + - Float8E4M3FN/Float8E5M2 + + :note: Supported accumulator data types: + - Float32 (for all floating point A/B data types) + - Float16 (only for fp16 and fp8 A/B data types) + - Int32 (only for uint8/int8 A/B data types) + + :note: Supported C data types: + - Float32 (for float32 and int32 accumulator data types) + - Int32 (for float32 and int32 accumulator data types) + - Float16/BFloat16 (for fp16 and fp8 accumulator data types) + - Int8/Uint8 (for uint8/int8 accumulator data types) + - Float8E4M3FN/Float8E5M2 (for float32 accumulator data types) + + :note: Constraints: + - MMA tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) + - MMA tiler N must be 32-256, step 32 + - Cluster shape M must be multiple of 2 if use_2cta_instrs=True + - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 + """ + + supported_args_type = GemmArguments + designed_for_min_cc = 100 + + def __init__(self, metadata: OperatorMetadata): + super().__init__(metadata) + + def epilogue_op(x): + return x + + mma_tiler_mn = (metadata.design.tile_shape[0], metadata.design.tile_shape[1]) + cluster_shape_mn = ( + metadata.design.cluster_shape[0], + metadata.design.cluster_shape[1], + ) + + self.impl = PersistentDenseGemmKernel( + metadata.operands.accumulator_type, + metadata.design.use_2cta_mma, + mma_tiler_mn, + cluster_shape_mn, + metadata.design.use_tma_store, + scheduler_metadata_to_enum(metadata.design.tile_scheduler), + epilogue_op, + ) + + def _compile( + self, args: GemmArguments, target_sm: TargetSm | None = None + ) -> CompiledArtifact: + stream = cutlass.cute.runtime.make_fake_stream() + if self.impl.use_static_persistent_scheduler: + max_active_clusters = get_max_active_clusters(self.impl.cluster_shape_mn) + else: + max_active_clusters = None + + return self.cute_compile( + self.impl, + args.A.tensor, + args.B.tensor, + args.out.tensor, + stream, + max_active_clusters, + self.impl.epilogue_op, + target_sm=target_sm, + ) + + def _run( + self, + args: GemmArguments, + compiled_artifact: CompiledArtifact, + stream, + workspace=None, + ) -> None: + stream = to_cuda_stream(stream) + compiled_gemm = compiled_artifact.compiled_obj + self.cute_run( + compiled_gemm, args.A.tensor, args.B.tensor, args.out.tensor, stream + ) + + @staticmethod + def _valid_operands(operands: GemmOperandsMetadata) -> bool: + if not isinstance(operands, GemmOperandsMetadata): + return False + + # In current version, A and B tensor must have the same data type + # i.e., Float8E4M3FN for A and Float8E5M2 for B is not supported + if operands.A.dtype != operands.B.dtype: + return False + + abtype = operands.A.dtype + + # Supported A/B data types: + # - Float16/BFloat16 + # - Int8/Uint8 + # - Float8E4M3FN/Float8E5M2 + if abtype not in [ + cutlass.Float16, + cutlass.BFloat16, + cutlass.Int8, + cutlass.Uint8, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ]: + return False + + # Supported accumulator data types: + # - Float32 (for all floating point A/B data types) + # - Float16 (only for fp16 and fp8 A/B data types) + # - Int32 (only for uint8/int8 A/B data types) + if operands.accumulator_type == cutlass.Float32: + if not abtype.is_float: + return False + elif operands.accumulator_type == cutlass.Float16: + if abtype not in [ + cutlass.Float16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ]: + return False + elif operands.accumulator_type == cutlass.Int32: + if abtype not in [cutlass.Uint8, cutlass.Int8]: + return False + else: + return False + + # Supported out data types: + # - Float32 (for float32 and int32 accumulator data types) + # - Int32 (for float32 and int32 accumulator data types) + # - Float16/BFloat16 (for fp16 and fp8 accumulator data types) + # - Int8/Uint8 (for uint8/int8 accumulator data types) + # - Float8E4M3FN/Float8E5M2 (for float32 accumulator data types) + if operands.out.dtype == cutlass.Float32 or operands.out.dtype == cutlass.Int32: + if operands.accumulator_type not in [cutlass.Float32, cutlass.Int32]: + return False + elif ( + operands.out.dtype == cutlass.Float16 + or operands.out.dtype == cutlass.BFloat16 + ): + if operands.accumulator_type not in [ + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + ]: + return False + elif operands.out.dtype == cutlass.Int8 or operands.out.dtype == cutlass.Uint8: + if operands.accumulator_type not in [cutlass.Int32]: + return False + elif ( + operands.out.dtype == cutlass.Float8E4M3FN + or operands.out.dtype == cutlass.Float8E5M2 + ): + if operands.accumulator_type not in [cutlass.Float32]: + return False + else: + return False + + return True + + @staticmethod + def _metadata_operands_from_args( + args: GemmArguments, + ) -> Generator[GemmOperandsMetadata, None, None]: + """Generator that yields all valid GemmOperandsMetadata that correspond to `args` and + that are supported by this Operator. + + :param args: The arguments to generate operands from. + :type args: GemmArguments + + :return: A generator of valid GemmOperandsMetadata that correspond to `args` and + that are supported by this Operator. + :rtype: Generator[GemmOperandsMetadata, None, None] + """ + if any(not isinstance(arg, DenseTensor) for arg in [args.A, args.B, args.out]): + return + + if args.A.dtype != args.B.dtype: + return + + valid_major_strides = {(0, 0, 1), (0, 1, 0)} + + A_stride = normalized_major_stride( + args.A.shape, args.A.stride, prepend_zeros_to_rank=3 + ) + B_stride = normalized_major_stride( + args.B.shape, args.B.stride, prepend_zeros_to_rank=3 + ) + out_stride = normalized_major_stride( + args.out.shape, args.out.stride, prepend_zeros_to_rank=3 + ) + + if any(x not in valid_major_strides for x in [A_stride, B_stride, out_stride]): + return + + operands = generate.construct_dense_gemm_operands( + args.A.dtype, + args.B.dtype, + args.out.dtype, + args.accumulator_type, + A_stride, + B_stride, + out_stride, + alignment_bytes=16, + ) + + if not PersistentDenseGemmOperator._valid_operands(operands): + return + + yield operands + + @staticmethod + def _metadata_operand_combinations() -> Generator[GemmOperandsMetadata, None, None]: + """Generator that yields all valid GemmOperandsMetadata combinations + based on the validation rules in _valid_operands. + """ + # Supported A/B data types (must be the same) + ab_dtypes = [ + cutlass.Float16, + cutlass.BFloat16, + cutlass.Int8, + cutlass.Uint8, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ] + + row_major_stride = (0, 0, 1) + col_major_stride = (0, 1, 0) + + for ab_dtype in ab_dtypes: + # Determine valid accumulator types for this A/B dtype + valid_acc_dtypes = [] + + if ab_dtype.is_float: + valid_acc_dtypes.append(cutlass.Float32) + if ab_dtype in [ + cutlass.Float16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ]: + valid_acc_dtypes.append(cutlass.Float16) + else: # Int8, Uint8 + valid_acc_dtypes.append(cutlass.Int32) + + for acc_dtype in valid_acc_dtypes: + # Determine valid output types for this accumulator type + valid_out_dtypes = [] + + if acc_dtype == cutlass.Float32: + valid_out_dtypes.extend( + [ + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float32, + cutlass.Int32, + ] + ) + elif acc_dtype == cutlass.Int32: + valid_out_dtypes.extend( + [cutlass.Int8, cutlass.Uint8, cutlass.Float32, cutlass.Int32] + ) + elif acc_dtype == cutlass.Float16: + valid_out_dtypes.extend([cutlass.Float16]) + + for out_dtype in valid_out_dtypes: + for stride_A, stride_B, stride_out in itertools.product( + [row_major_stride, col_major_stride], repeat=3 + ): + operands = generate.construct_dense_gemm_operands( + ab_dtype, + ab_dtype, + out_dtype, + acc_dtype, + stride_A, + stride_B, + stride_out, + alignment_bytes=16, + ) + yield operands + + @staticmethod + def _valid_metadata(metadata: OperatorMetadata) -> bool: + if not PersistentDenseGemmOperator._valid_operands(metadata.operands): + return False + + design = metadata.design + if not isinstance(design, Sm100DesignMetadata): + return False + + cluster_size_m, cluster_size_n, _ = design.cluster_shape + + if cluster_size_m % 2 != 0 and cluster_size_m != 1: + return False + if cluster_size_n % 2 != 0 and cluster_size_n != 1: + return False + if cluster_size_m * cluster_size_n > 16: + return False + + tile = design.tile_shape + + # Constraints based on whether 2CTA instructions are used + if design.use_2cta_mma is not None: + if design.use_2cta_mma: + if cluster_size_m % 2 != 0: + return False + if tile is not None and tile[0] not in [ + 128, + 256, + ]: + return False + else: + if tile is not None and tile[0] not in [ + 64, + 128, + ]: + return False + + if tile is not None and tile[1] not in range(32, 257, 32): + return False + + if metadata.epilogue is not None: + return False + + return True + + @classmethod + def _generate_operators( + cls, + metadata_filter: Callable[[OperatorMetadata], bool], + epilogue_args: EpilogueArguments = None, + target_sm: TargetSm | None = None, + args: GemmArguments = None, + ) -> list[PersistentDenseGemmOperator]: + """Returns a list of all possible configurations of PersistentDenseGemmOperator that + adhere to constraints passed in under kwargs. + """ + design_params = { + "mma_instruction_type": [BlackwellTcgen05Mma], + "use_2cta_mma": [True, False], + # Fill in tile shape K based on data type + "tile_shape": [ + (M, N, None) for M in [64, 128, 256] for N in [32, 64, 128, 256] + ], + "cluster_shape": [ + (M, N, 1) for M in [1, 2, 4, 8, 16] for N in [1, 2, 4, 8, 16] + ], + "use_tma_store": [True], + "tile_scheduler": [ + CLCDynamicPersistentTileSchedulerMetadata(), + StaticPersistentTileSchedulerMetadata(), + ], + } + + if epilogue_args is not None: + return [] + + from itertools import product + + # Get the list of tunable parameter names and their possible values + param_names = list(design_params.keys()) + + operator_list = [] + + if args is not None: + operands_generator = cls._metadata_operands_from_args(args) + else: + operands_generator = cls._metadata_operand_combinations() + + for operands in operands_generator: + tile_k = BlackwellTcgen05Mma.shape_k(operands) * 4 # kphase=4 + design_params["tile_shape"] = [ + (M, N, tile_k) for M, N, _ in design_params["tile_shape"] + ] + param_values = [design_params[name] for name in param_names] + for values in product(*param_values): + design = Sm100DesignMetadata(**dict(zip(param_names, values))) + operator_name = "cutedsl.PersistentDenseGemmOperator_sm100_{layout}_A{A}_B{B}_out{out}_acc{acc}_{num_cta}cta_cluster{cluster}_tile{tile}_scheduler{scheduler}{_tma_store}".format( + layout=strides_to_layout_string( + operands.A.stride, operands.B.stride, operands.out.stride + ), + A=operands.A.dtype, + B=operands.B.dtype, + out=operands.out.dtype, + acc=operands.accumulator_type, + num_cta=("2" if design.use_2cta_mma else "1"), + cluster=tuple_to_string(design.cluster_shape), + tile=tuple_to_string(design.tile_shape), + scheduler=scheduler_metadata_to_enum(design.tile_scheduler).name, + _tma_store="_tma_store" if design.use_tma_store else "", + ) + + metadata = OperatorMetadata( + operands=operands, + design=design, + operator_name=operator_name, + operator_class=cls, + supported_targets=TargetSm.get_supported_targets(design, operands), + epilogue=None, + ) + + if PersistentDenseGemmOperator._valid_metadata( + metadata + ) and metadata_filter(metadata): + operator_list.append(PersistentDenseGemmOperator(metadata)) + + return operator_list diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/sm100_persistent_preferred_cluster.py b/operators/cutlass/operators/providers/cutedsl/gemm/sm100_persistent_preferred_cluster.py new file mode 100644 index 000000000..8403611ec --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/sm100_persistent_preferred_cluster.py @@ -0,0 +1,671 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""Operator API Operator class for preferred-cluster persistent dense GEMM on SM100. + +In preferred cluster mode, driver launches CTAs with two cluster shapes +simultaneously: + +- **Preferred cluster** (larger): main workload, benefits from TMA multicast + and optional 2-CTA MMA instructions. +- **Fallback cluster** (smaller): handles the residual tiles that cannot fill a + full preferred cluster, keeping all SMs busy. + +From the caller's perspective this Operator is identical to the non-preferred-cluster +:class:`PersistentDenseGemmOperator` — it accepts the same +:class:`~cutlass.operators.arguments.GemmArguments` and the same +:class:`~cutlass.operators.metadata.OperatorMetadata`. The distinction is expressed +solely through the :class:`~cutlass.operators.metadata.Sm100DesignMetadata` fields: + +- ``cluster_shape`` → preferred (larger) cluster shape. +- ``fallback_cluster_shape`` → fallback (smaller) cluster shape. + +When ``fallback_cluster_shape`` is ``None`` or equal to ``cluster_shape`` the +design metadata represents a regular single-cluster operator, handled by +:class:`PersistentDenseGemmOperator`. + +Usage example:: + + import torch + import cutlass.operators as ops + + M, N, K, L = 512, 512, 128, 1 + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=torch.float16) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=torch.float16) + D = torch.empty((L, M, N), device="cuda", dtype=torch.float16) + args = ops.GemmArguments(A, B, D, accumulator_type=torch.float32) + + operator = ops.get_operators( + args, + metadata_filter=lambda m: ( + isinstance(m.design, ops.metadata.design.sm100.Sm100DesignMetadata) + and m.design.use_fallback_cluster() + ), + )[0] + + operator.run(args) + torch.testing.assert_close(D, (A @ B).to(D.dtype)) + +""" + +from __future__ import annotations + +import itertools +from typing import TYPE_CHECKING + +import cutlass + +import cutlass.operators.utils.generate as generate +from cutlass.operators.arch import TargetSm +from cutlass.operators.arguments import ( + DenseTensor, + EpilogueArguments, + GemmArguments, +) +from cutlass.operators.metadata import ( + GemmOperandsMetadata, + OperatorMetadata, + Sm100DesignMetadata, +) +from cutlass.operators.metadata.design.tile_scheduler import ( + CLCDynamicPersistentTileSchedulerMetadata, + StaticPersistentTileSchedulerMetadata, +) +from cutlass.operators.mma import BlackwellTcgen05Mma +from cutlass.operators.providers.cutedsl import CuTeDSLProvider +from cutlass.operators.providers.cutedsl.integration_utils.mma import ( + get_max_active_clusters, +) +from cutlass.operators.providers.cutedsl.operator import CuteDslOperator +from cutlass.operators.utils.common import tuple_to_string +from cutlass.operators.utils.device import to_cuda_stream +from cutlass.operators.utils.tensor import ( + normalized_major_stride, + strides_to_layout_string, +) + +from .implementations.operator_helpers import scheduler_metadata_to_enum +from .implementations.sm100_persistent_preferred_cluster_impl import ( + PersistentDenseGemmPreferredClusterKernel, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Generator + + from cutlass.operators.artifact import CompiledArtifact + + +@CuTeDSLProvider.register +class PersistentDenseGemmPreferredClusterOperator(CuteDslOperator): + """Persistent batched dense GEMM with preferred cluster on SM100 Blackwell GPUs. + + Launches with two cluster shapes simultaneously: + + - The *preferred* (larger) cluster shape is used for the majority of the + output tiles. It enables larger TMA multicast regions and, optionally, + 2-CTA MMA instructions. + - The *fallback* (smaller) cluster shape handles the residual tiles whose + count is not a multiple of the preferred cluster size, avoiding idle SMs. + + :note: Preferred cluster shape M/N must each be a multiple of the fallback + cluster shape M/N. + :note: Both preferred and fallback cluster shapes must satisfy: + + - M and N are positive powers of 2. + - Total cluster size (M × N) ≤ 16. + - M is a multiple of 2 when use_2cta_mma=True. + + :note: Supported A/B data types: Float16, BFloat16, Int8, Uint8, + Float8E4M3FN, Float8E5M2. A and B must share the same type. + + :note: MMA tiler M constraints: + + - 64 or 128 when use_2cta_mma=False. + - 128 or 256 when use_2cta_mma=True. + + :note: MMA tiler N must be in the range [32, 256] with step 32. + """ + + supported_args_type = GemmArguments + designed_for_min_cc = 100 + + def __init__(self, metadata: OperatorMetadata): + super().__init__(metadata) + + def epilogue_op(x): + return x + + mma_tiler_mn = (metadata.design.tile_shape[0], metadata.design.tile_shape[1]) + preferred_cluster_shape_mn = ( + metadata.design.cluster_shape[0], + metadata.design.cluster_shape[1], + ) + fallback_cluster_shape_mn = ( + metadata.design.fallback_cluster_shape[0], + metadata.design.fallback_cluster_shape[1], + ) + + self.impl = PersistentDenseGemmPreferredClusterKernel( + acc_dtype=metadata.operands.accumulator_type, + use_2cta_instrs=metadata.design.use_2cta_mma, + mma_tiler_mn=mma_tiler_mn, + preferred_cluster_shape_mn=preferred_cluster_shape_mn, + fallback_cluster_shape_mn=fallback_cluster_shape_mn, + use_tma_store=metadata.design.use_tma_store, + epilogue_op=epilogue_op, + persistent_tile_scheduler_type=scheduler_metadata_to_enum( + metadata.design.tile_scheduler + ), + ) + + def _compile( + self, args: GemmArguments, target_sm: TargetSm | None = None + ) -> CompiledArtifact: + stream = cutlass.cute.runtime.make_fake_stream() + + if self.impl.use_static_persistent_scheduler: + max_active_preferred_clusters = get_max_active_clusters( + self.impl.preferred_cluster_shape_mn + ) + max_active_fallback_clusters = get_max_active_clusters( + self.impl.fallback_cluster_shape_mn + ) + else: + max_active_preferred_clusters = None + max_active_fallback_clusters = None + + return self.cute_compile( + self.impl, + args.A.tensor, + args.B.tensor, + args.out.tensor, + max_active_preferred_clusters, + max_active_fallback_clusters, + stream, + self.impl.epilogue_op, + target_sm=target_sm, + ) + + def _run( + self, + args: GemmArguments, + compiled_artifact: CompiledArtifact, + stream, + workspace=None, + ) -> None: + stream = to_cuda_stream(stream) + + compiled_gemm = compiled_artifact.compiled_obj + self.cute_run( + compiled_gemm, + args.A.tensor, + args.B.tensor, + args.out.tensor, + stream, + ) + + # ------------------------------------------------------------------ + # Operand validation + # ------------------------------------------------------------------ + + @staticmethod + def _valid_operands(operands: GemmOperandsMetadata) -> bool: + """Return True when *operands* are supported by this Operator.""" + if not isinstance(operands, GemmOperandsMetadata): + return False + + # In current version, A and B tensor must have the same data type + # i.e., Float8E4M3FN for A and Float8E5M2 for B is not supported + if operands.A.dtype != operands.B.dtype: + return False + + abtype = operands.A.dtype + + # Supported A/B data types: + # - Float16/BFloat16 + # - Int8/Uint8 + # - Float8E4M3FN/Float8E5M2 + if abtype not in [ + cutlass.Float16, + cutlass.BFloat16, + cutlass.Int8, + cutlass.Uint8, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ]: + return False + + # Supported accumulator data types: + # - Float32 (for all floating point A/B data types) + # - Float16 (only for fp16 and fp8 A/B data types) + # - Int32 (only for uint8/int8 A/B data types) + if operands.accumulator_type == cutlass.Float32: + if not abtype.is_float: + return False + elif operands.accumulator_type == cutlass.Float16: + if abtype not in [ + cutlass.Float16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ]: + return False + elif operands.accumulator_type == cutlass.Int32: + if abtype not in [cutlass.Uint8, cutlass.Int8]: + return False + else: + return False + + # Supported out data types: + # - Float32 (for float32 and int32 accumulator data types) + # - Int32 (for float32 and int32 accumulator data types) + # - Float16/BFloat16 (for fp16 and fp8 accumulator data types) + # - Int8/Uint8 (for uint8/int8 accumulator data types) + # - Float8E4M3FN/Float8E5M2 (for float32 accumulator data types) + if operands.out.dtype == cutlass.Float32 or operands.out.dtype == cutlass.Int32: + if operands.accumulator_type not in [cutlass.Float32, cutlass.Int32]: + return False + elif ( + operands.out.dtype == cutlass.Float16 + or operands.out.dtype == cutlass.BFloat16 + ): + if operands.accumulator_type not in [ + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + ]: + return False + elif operands.out.dtype == cutlass.Int8 or operands.out.dtype == cutlass.Uint8: + if operands.accumulator_type not in [cutlass.Int32]: + return False + elif ( + operands.out.dtype == cutlass.Float8E4M3FN + or operands.out.dtype == cutlass.Float8E5M2 + ): + if operands.accumulator_type not in [cutlass.Float32]: + return False + else: + return False + + return True + + @staticmethod + def _metadata_operands_from_args( + args: GemmArguments, + ) -> Generator[GemmOperandsMetadata, None, None]: + """Yield operand metadata derived from *args* that this Operator supports. + + Inspects the runtime tensor dtypes, strides, and accumulator type to + construct a :class:`GemmOperandsMetadata` with divisibility constraints + baked in. Yields nothing if *args* are incompatible (e.g. non-dense + tensors, mismatched A/B dtypes, or unsupported layout). + + :param args: Runtime GEMM arguments supplied by the caller. + :type args: GemmArguments + + :return: Generator of matching operand metadata (zero or one item). + :rtype: Generator[GemmOperandsMetadata, None, None] + """ + if any(not isinstance(arg, DenseTensor) for arg in [args.A, args.B, args.out]): + return + + if args.A.dtype != args.B.dtype: + return + + # (0, 0, 1) is row major + # (0, 1, 0) is col major + valid_major_strides = {(0, 0, 1), (0, 1, 0)} + + A_stride = normalized_major_stride( + args.A.shape, args.A.stride, prepend_zeros_to_rank=3 + ) + B_stride = normalized_major_stride( + args.B.shape, args.B.stride, prepend_zeros_to_rank=3 + ) + out_stride = normalized_major_stride( + args.out.shape, args.out.stride, prepend_zeros_to_rank=3 + ) + + if any(x not in valid_major_strides for x in [A_stride, B_stride, out_stride]): + return + + operands = generate.construct_dense_gemm_operands( + args.A.dtype, + args.B.dtype, + args.out.dtype, + args.accumulator_type, + A_stride, + B_stride, + out_stride, + alignment_bytes=16, + ) + + if not PersistentDenseGemmPreferredClusterOperator._valid_operands(operands): + return + + yield operands + + @staticmethod + def _metadata_operand_combinations() -> Generator[GemmOperandsMetadata, None, None]: + """Yield all valid operand metadata combinations supported by this Operator. + + Used when no runtime arguments are available (e.g. manifest enumeration). + Iterates over all supported A/B dtypes, accumulator types, output dtypes, + and row/col-major layout combinations, constructing a + :class:`GemmOperandsMetadata` with 16-byte alignment divisibility for each. + + :return: Generator of all valid operand metadata configurations. + :rtype: Generator[GemmOperandsMetadata, None, None] + """ + ab_dtypes = [ + cutlass.Float16, + cutlass.BFloat16, + cutlass.Int8, + cutlass.Uint8, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ] + + row_major_stride = (0, 0, 1) + col_major_stride = (0, 1, 0) + + for ab_dtype in ab_dtypes: + # Determine valid accumulator types for this A/B dtype + valid_acc_dtypes = [] + + if ab_dtype.is_float: + valid_acc_dtypes.append(cutlass.Float32) + if ab_dtype in [ + cutlass.Float16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ]: + valid_acc_dtypes.append(cutlass.Float16) + else: + valid_acc_dtypes.append(cutlass.Int32) + + for acc_dtype in valid_acc_dtypes: + # Determine valid output types for this accumulator type + valid_out_dtypes = [] + + if acc_dtype == cutlass.Float32: + valid_out_dtypes.extend( + [ + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float32, + cutlass.Int32, + ] + ) + elif acc_dtype == cutlass.Int32: + valid_out_dtypes.extend( + [cutlass.Int8, cutlass.Uint8, cutlass.Float32, cutlass.Int32] + ) + elif acc_dtype == cutlass.Float16: + valid_out_dtypes.extend([cutlass.Float16]) + + for out_dtype in valid_out_dtypes: + for stride_A, stride_B, stride_out in itertools.product( + [row_major_stride, col_major_stride], repeat=3 + ): + operands = generate.construct_dense_gemm_operands( + ab_dtype, + ab_dtype, + out_dtype, + acc_dtype, + stride_A, + stride_B, + stride_out, + alignment_bytes=16, + ) + yield operands + + # ------------------------------------------------------------------ + # Metadata validation + # ------------------------------------------------------------------ + + @staticmethod + def _valid_cluster_shape(cluster_size_m: int, cluster_size_n: int) -> bool: + """Return True when the cluster shape is structurally valid. + + A valid cluster shape must have M and N each be a power of 2 (≥ 1), + and the total cluster size M * N must not exceed 16. + + :param cluster_size_m: Cluster M dimension. + :type cluster_size_m: int + :param cluster_size_n: Cluster N dimension. + :type cluster_size_n: int + + :return: True if the cluster shape satisfies all hardware constraints. + :rtype: bool + """ + if cluster_size_m % 2 != 0 and cluster_size_m != 1: + return False + if cluster_size_n % 2 != 0 and cluster_size_n != 1: + return False + if cluster_size_m * cluster_size_n > 16: + return False + return True + + @staticmethod + def _valid_metadata(metadata: OperatorMetadata) -> bool: + if not PersistentDenseGemmPreferredClusterOperator._valid_operands( + metadata.operands + ): + return False + + design = metadata.design + if not isinstance(design, Sm100DesignMetadata): + return False + + if not design.use_fallback_cluster(): + return False + + if not isinstance( + design.tile_scheduler, + ( + CLCDynamicPersistentTileSchedulerMetadata, + StaticPersistentTileSchedulerMetadata, + ), + ): + return False + + preferred_m, preferred_n = design.cluster_shape[0], design.cluster_shape[1] + fallback_m, fallback_n = ( + design.fallback_cluster_shape[0], + design.fallback_cluster_shape[1], + ) + + if not PersistentDenseGemmPreferredClusterOperator._valid_cluster_shape( + preferred_m, preferred_n + ): + return False + if not PersistentDenseGemmPreferredClusterOperator._valid_cluster_shape( + fallback_m, fallback_n + ): + return False + + # Preferred cluster shape must be a multiple of fallback in each dimension. + if preferred_m % fallback_m != 0 or preferred_n % fallback_n != 0: + return False + + tile = design.tile_shape + + # 2-CTA MMA constraints propagate to both cluster shapes. + if design.use_2cta_mma is not None: + if design.use_2cta_mma: + if preferred_m % 2 != 0 or fallback_m % 2 != 0: + return False + if tile is not None and tile[0] not in [128, 256]: + return False + else: + if tile is not None and tile[0] not in [64, 128]: + return False + + if tile is not None and tile[1] not in range(32, 257, 32): + return False + + # Epilogue fusion is not yet supported. + if metadata.epilogue is not None: + return False + + return True + + # ------------------------------------------------------------------ + # Operator generation (called by the provider/Manifest) + # ------------------------------------------------------------------ + + @classmethod + def _generate_operators( + cls, + metadata_filter: Callable[[OperatorMetadata], bool], + epilogue_args: EpilogueArguments = None, + target_sm: TargetSm | None = None, + args: GemmArguments = None, + ) -> list[PersistentDenseGemmPreferredClusterOperator]: + """Return all valid preferred-cluster operator configurations matching *metadata_filter*. + + :param metadata_filter: Predicate applied to each candidate :class:`OperatorMetadata`. + :param epilogue_args: Not yet supported; passing a non-None value yields []. + :param target_sm: Optional target architecture filter. + :param args: Optional runtime arguments used to narrow the operand search space. + """ + # No epilogue fusion support for the operator class, so early return + if epilogue_args is not None: + return [] + + from itertools import product + + # Enumerate candidate cluster shape pairs. + # Preferred cluster must be a strict superset of fallback (preferred ≥ fallback + # per dimension) and must differ from fallback. + preferred_cluster_shapes = [ + (M, N, 1) for M in [1, 2, 4, 8, 16] for N in [1, 2, 4, 8, 16] + ] + fallback_cluster_shapes = [(2, 1, 1)] + cluster_shape_pairs = [ + (preferred, fallback) + for preferred in preferred_cluster_shapes + for fallback in fallback_cluster_shapes + if ( + preferred != fallback + and preferred[0] % fallback[0] == 0 + and preferred[1] % fallback[1] == 0 + ) + ] + + design_params = { + "mma_instruction_type": [BlackwellTcgen05Mma], + "use_2cta_mma": [True, False], + # Fill in tile shape K based on data type + "tile_shape": [ + (M, N, None) for M in [64, 128, 256] for N in [32, 64, 128, 256] + ], + # cluster_shape and fallback_cluster_shape are co-varied below. + "use_tma_store": [True], + "tile_scheduler": [ + CLCDynamicPersistentTileSchedulerMetadata(), + StaticPersistentTileSchedulerMetadata(), + ], + } + + param_names = list(design_params.keys()) + + operator_list = [] + + if args is not None: + operands_generator = cls._metadata_operands_from_args(args) + else: + operands_generator = cls._metadata_operand_combinations() + + for operands in operands_generator: + tile_k = BlackwellTcgen05Mma.shape_k(operands) * 4 # kphase=4 + design_params["tile_shape"] = [ + (M, N, tile_k) for M, N, _ in design_params["tile_shape"] + ] + param_values = [design_params[name] for name in param_names] + + for values in product(*param_values): + base_design = dict(zip(param_names, values)) + + for preferred_shape, fallback_shape in cluster_shape_pairs: + design = Sm100DesignMetadata( + **base_design, + cluster_shape=preferred_shape, + fallback_cluster_shape=fallback_shape, + ) + + operator_name = ( + "cutedsl.PersistentDenseGemmPreferredClusterOperator_sm100" + "_{layout}" + "_A{A}_B{B}_out{out}_acc{acc}" + "_{num_cta}cta" + "_preferred{preferred}_fallback{fallback}" + "_tile{tile}" + "_scheduler{scheduler}" + "{_tma_store}" + ).format( + layout=strides_to_layout_string( + operands.A.stride, + operands.B.stride, + operands.out.stride, + ), + A=operands.A.dtype, + B=operands.B.dtype, + out=operands.out.dtype, + acc=operands.accumulator_type, + num_cta=("2" if design.use_2cta_mma else "1"), + preferred=tuple_to_string(preferred_shape), + fallback=tuple_to_string(fallback_shape), + tile=tuple_to_string(design.tile_shape), + scheduler=scheduler_metadata_to_enum( + design.tile_scheduler + ).name, + _tma_store="_tma_store" if design.use_tma_store else "", + ) + + metadata = OperatorMetadata( + operands=operands, + design=design, + operator_name=operator_name, + operator_class=cls, + supported_targets=TargetSm.get_supported_targets( + design, operands + ), + epilogue=None, + ) + + if PersistentDenseGemmPreferredClusterOperator._valid_metadata( + metadata + ) and metadata_filter(metadata): + operator_list.append( + PersistentDenseGemmPreferredClusterOperator(metadata) + ) + + return operator_list diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/sm100_static_persistent_efc.py b/operators/cutlass/operators/providers/cutedsl/gemm/sm100_static_persistent_efc.py new file mode 100644 index 000000000..630075ae0 --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/sm100_static_persistent_efc.py @@ -0,0 +1,506 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +import itertools +from typing import TYPE_CHECKING + +import cutlass + +import cutlass.operators.utils.generate as generate +from cutlass.operators.arch import TargetSm +from cutlass.operators.arguments import ( + DenseTensor, + EpilogueArguments, + GemmArguments, + Operand, +) +from cutlass.operators.metadata import ( + EpilogueMetadata, + GemmOperandsMetadata, + OperatorMetadata, + Sm100DesignMetadata, +) +from cutlass.operators.mma import BlackwellTcgen05Mma +from cutlass.operators.providers.cutedsl import CuTeDSLProvider +from cutlass.operators.providers.cutedsl.evt.converter import EFCConverter +from cutlass.operators.providers.cutedsl.integration_utils.mma import ( + get_max_active_clusters, +) +from cutlass.operators.providers.cutedsl.operator import CuteDslOperator +from cutlass.operators.status import Status +from cutlass.operators.utils.common import tuple_to_string +from cutlass.operators.utils.device import to_cuda_stream +from cutlass.operators.utils.tensor import ( + TensorWrapper, + normalized_major_stride, + strides_to_layout_string, +) + +from .implementations.sm100_static_persistent_efc_impl import ( + PersistentDenseGemmEFCKernel, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Generator + + from cutlass.operators.artifact import CompiledArtifact + + +@CuTeDSLProvider.register +class PersistentDenseGemmEFCOperator(CuteDslOperator): + """Base class for batched GEMM with custom epilogue fusion using EFC. + + This class provides the core infrastructure for persistent batched GEMM operations + with customizable epilogue fusion. Subclasses define specific epilogue behaviors + by providing an epilogue configuration function that describes operations on the + accumulator and supplemental tensors. + + The class handles: + - GEMM mainloop (A * B computation) + - TMA-based memory operations + - Warp specialization + - Persistent tile scheduling + - EFC (Epilogue Fusion Configuration) integration + - Tensor creation and validation + + :param acc_dtype: Data type for accumulation during MMA computation + :type acc_dtype: type[cutlass.Numeric] + :param epi_dtype: Data type for epilogue operation + :type epi_dtype: type[cutlass.Numeric] + :param use_2cta_instrs: Whether to use CTA group 2 for 2CTA MMA instructions + :type use_2cta_instrs: bool + :param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N) + :type mma_tiler_mn: tuple[int, int] + :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing + :type cluster_shape_mn: tuple[int, int] + :param epilogue_configuration_function: Function defining the epilogue behavior via EFC + :type epilogue_configuration_function: Callable + + :note: Supported A/B data types: + - Float16/BFloat16 + - Int8/Uint8 + - Float8E4M3FN/Float8E5M2 + (A and B must have the same data type) + + :note: Supported accumulator data types: + - Float32 (for all floating point A/B data types) + - Float16 (only for fp16 and fp8 A/B data types) + - Int32 (only for uint8/int8 A/B data types) + + :note: Supported supplemental tensor data types (epilogue-dependent): + - Float32 (for float32 and int32 accumulator data types) + - Int32 (for float32 and int32 accumulator data types) + - Float16/BFloat16 (for fp16 and fp8 accumulator data types) + - Int8/Uint8 (for uint8/int8 accumulator data types) + - Float8E4M3FN/Float8E5M2 (for float32 accumulator data types) + + :note: Constraints: + - MMA tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) + - MMA tiler N must be 32-256, step 32 + - Cluster shape M must be multiple of 2 if use_2cta_instrs=True + - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 + """ + + supported_args_type = GemmArguments + designed_for_min_cc = 100 + + _valid_ab_acc_combinations = { + cutlass.Float16: [cutlass.Float16, cutlass.Float32], + cutlass.BFloat16: [cutlass.Float32], + cutlass.Uint8: [cutlass.Int32], + cutlass.Int8: [cutlass.Int32], + cutlass.Float8E4M3FN: [cutlass.Float16, cutlass.Float32], + cutlass.Float8E5M2: [cutlass.Float16, cutlass.Float32], + } + + def __init__(self, metadata: OperatorMetadata): + super().__init__(metadata) + + if metadata.epilogue is not None: + epilogue_op = EFCConverter.convert( + metadata.epilogue.traced_epilogue.dag_ir, + metadata.epilogue.parameter_names, + parameter_tensors=metadata.epilogue.tensors, + ) + else: + epilogue_op = EFCConverter.identity_efc + + mma_tiler_mn = (metadata.design.tile_shape[0], metadata.design.tile_shape[1]) + cluster_shape_mn = ( + metadata.design.cluster_shape[0], + metadata.design.cluster_shape[1], + ) + + self.impl = PersistentDenseGemmEFCKernel( + metadata.operands.accumulator_type, + metadata.operands.out.dtype, + metadata.design.use_2cta_mma, + mma_tiler_mn, + cluster_shape_mn, + epilogue_op, + ) + + @staticmethod + def _valid_fusion(fusion: EpilogueMetadata) -> Status: + if not isinstance(fusion, EpilogueMetadata): + return Status.fail("Unsupported epilogue argument type.") + if not (status := EFCConverter.supports(fusion.traced_epilogue.dag_ir)): + return status + return Status.success() + + @staticmethod + def _valid_operands(operands: GemmOperandsMetadata) -> Status: + if not isinstance(operands, GemmOperandsMetadata): + return False + + if operands.A.dtype != operands.B.dtype: + return False + + ab_dtype = operands.A.dtype + acc_dtype = operands.accumulator_type + + if ab_dtype not in PersistentDenseGemmEFCOperator._valid_ab_acc_combinations: + return False + + # Check compatibility between accumulator type and AB type + if ( + acc_dtype + not in PersistentDenseGemmEFCOperator._valid_ab_acc_combinations[ab_dtype] + ): + return False + + return True + + @staticmethod + def _valid_design(design: Sm100DesignMetadata) -> bool: + """Check if the design metadata is valid. + + :param design: The design metadata + :type design: Sm100DesignMetadata + + :return: True if the design is valid, False otherwise + :rtype: bool + """ + if not isinstance(design, Sm100DesignMetadata): + return False + + use_2cta_instrs = design.use_2cta_mma + mma_tiler_mn = design.tile_shape + cluster_shape_mn = design.cluster_shape + + # Check invalid mma tile shape M dimension + if not ( + (not use_2cta_instrs and mma_tiler_mn[0] in [64, 128]) + or (use_2cta_instrs and mma_tiler_mn[0] in [128, 256]) + ): + return False + + # Check invalid mma tile shape N dimension + if mma_tiler_mn[1] not in range(32, 257, 32): + return False + # Check illegal cluster shape M dimension + if cluster_shape_mn[0] % (2 if use_2cta_instrs else 1) != 0: + return False + + def is_power_of_2(x): + return x > 0 and (x & (x - 1)) == 0 + + # Check invalid cluster shape constraints + if cluster_shape_mn[0] * cluster_shape_mn[1] > 16: + return False + if cluster_shape_mn[0] <= 0 or cluster_shape_mn[1] <= 0: + return False + if not is_power_of_2(cluster_shape_mn[0]) or not is_power_of_2( + cluster_shape_mn[1] + ): + return False + + return True + + @staticmethod + def _valid_metadata(metadata: OperatorMetadata) -> bool: + if not PersistentDenseGemmEFCOperator._valid_operands(metadata.operands): + return False + + if not PersistentDenseGemmEFCOperator._valid_design(metadata.design): + return False + + return True + + @staticmethod + def _metadata_operands_from_args( + args: GemmArguments, + ) -> Generator[GemmOperandsMetadata, None, None]: + """Generator that yields all valid GemmOperandsMetadata that correspond to `args` and + that are supported by this Operator. + + :param args: The arguments to generate operands from. + :type args: GemmArguments + + :return: A generator of valid GemmOperandsMetadata that correspond to `args` and + that are supported by this Operator. + :rtype: Generator[GemmOperandsMetadata, None, None] + """ + if any(not isinstance(arg, DenseTensor) for arg in [args.A, args.B, args.out]): + return + + if args.A.dtype != args.B.dtype: + return + + valid_major_strides = {(0, 0, 1), (0, 1, 0)} + + A_stride = normalized_major_stride( + args.A.shape, args.A.stride, prepend_zeros_to_rank=3 + ) + B_stride = normalized_major_stride( + args.B.shape, args.B.stride, prepend_zeros_to_rank=3 + ) + out_stride = normalized_major_stride( + args.out.shape, args.out.stride, prepend_zeros_to_rank=3 + ) + + if any(x not in valid_major_strides for x in [A_stride, B_stride, out_stride]): + return + + operands = generate.construct_dense_gemm_operands( + args.A.dtype, + args.B.dtype, + args.out.dtype, + args.accumulator_type, + A_stride, + B_stride, + out_stride, + alignment_bytes=16, + ) + + if not PersistentDenseGemmEFCOperator._valid_operands(operands): + return + + yield operands + + @staticmethod + def _metadata_operand_combinations() -> Generator[GemmOperandsMetadata, None, None]: + """Generator that yields all valid GemmOperandsMetadata combinations + based on the validation rules in _valid_operands. + """ + row_major_stride = (0, 0, 1) + col_major_stride = (0, 1, 0) + + for ( + ab_dtype, + valid_acc_dtypes, + ) in PersistentDenseGemmEFCOperator._valid_ab_acc_combinations.items(): + for acc_dtype in valid_acc_dtypes: + # Determine valid output types for this accumulator type + valid_out_dtypes = [] + + if acc_dtype == cutlass.Float32: + valid_out_dtypes.extend( + [ + cutlass.Float32, + cutlass.Int32, + cutlass.Float16, + cutlass.BFloat16, + ] + ) + # Float8 output types only valid with Float32 accumulator + valid_out_dtypes.extend([cutlass.Float8E4M3FN, cutlass.Float8E5M2]) + elif acc_dtype == cutlass.Int32: + valid_out_dtypes.extend([cutlass.Float32, cutlass.Int32]) + # Integer output types only valid with Int32 accumulator + valid_out_dtypes.extend([cutlass.Int8, cutlass.Uint8]) + elif acc_dtype == cutlass.Float16: + valid_out_dtypes.extend([cutlass.Float16, cutlass.BFloat16]) + + for out_dtype in valid_out_dtypes: + for stride_A, stride_B, stride_out in itertools.product( + [row_major_stride, col_major_stride], repeat=3 + ): + yield generate.construct_dense_gemm_operands( + ab_dtype, + ab_dtype, + out_dtype, + acc_dtype, + stride_A, + stride_B, + stride_out, + alignment_bytes=16, + ) + + def _supports( + self, args: GemmArguments, target_sm: TargetSm | None = None + ) -> Status: + if args.epilogue is not None: + fusion_metadata = EpilogueMetadata.from_args(args.epilogue) + if not self._valid_fusion(fusion_metadata): + return Status.fail( + "Provided epilogue fusion is not supported by this Operator" + ) + + return Status.success() + + def _compile( + self, args: GemmArguments, target_sm: TargetSm | None = None + ) -> CompiledArtifact: + stream = cutlass.cute.runtime.make_fake_stream() + max_active_clusters = get_max_active_clusters(self.impl.cluster_shape_mn) + + if args.epilogue is not None: + epilogue_params = [ + e.compile_time_tensor if isinstance(e, TensorWrapper) else e + for e in args.epilogue.parameters + ] + else: + epilogue_params = [args.out.tensor.compile_time_tensor] + + # EFC needs special handling for supplemental arguments. + self.impl.efc.compile(*epilogue_params) + compiled_artifact = self.cute_compile( + self.impl, + args.A.tensor, + args.B.tensor, + max_active_clusters, + stream, + *epilogue_params, + target_sm=target_sm, + ) + + return compiled_artifact + + def _run( + self, + args: GemmArguments, + compiled_artifact: CompiledArtifact, + stream, + workspace=None, + ) -> None: + stream = to_cuda_stream(stream) + + epilogue_params = ( + args.epilogue.parameters if args.epilogue is not None else [args.out] + ) + + # Unwrap Operand wrappers (e.g. DenseTensor) to their + # underlying TensorWrapper so that cute_run can extract the + # runtime tensor for TVM FFI. + epilogue_params = [ + e.tensor if isinstance(e, Operand) else e for e in epilogue_params + ] + + compiled_gemm = compiled_artifact.compiled_obj + self.cute_run( + compiled_gemm, args.A.tensor, args.B.tensor, stream, *epilogue_params + ) + + @classmethod + def _generate_operators( + cls, + metadata_filter: Callable[[OperatorMetadata], bool], + epilogue_args: EpilogueArguments = None, + target_sm: TargetSm | None = None, + args: GemmArguments = None, + ) -> list[PersistentDenseGemmEFCOperator]: + """Returns a list of all possible configurations of PersistentDenseGemmEFCOperator that + adhere to constraints passed in under kwargs. + """ + design_params = { + "mma_instruction_type": [BlackwellTcgen05Mma], + "use_2cta_mma": [True, False], + # Fill in tile shape K based on data type + "tile_shape": [ + (M, N, None) for M in [64, 128, 256] for N in [32, 64, 128, 256] + ], + "cluster_shape": [ + (M, N, 1) for M in [1, 2, 4, 8, 16] for N in [1, 2, 4, 8, 16] + ], + "use_tma_store": [True], + } + + if epilogue_args is not None: + if not isinstance(epilogue_args, EpilogueArguments): + return [] + + epilogue_metadata = EpilogueMetadata.from_args(epilogue_args) + if not PersistentDenseGemmEFCOperator._valid_fusion(epilogue_metadata): + return [] + else: + epilogue_metadata = None + + from itertools import product + + param_names = list(design_params.keys()) + + operator_list = [] + + if args is not None: + operands_generator = cls._metadata_operands_from_args(args) + else: + operands_generator = cls._metadata_operand_combinations() + + for operands in operands_generator: + if not cls._valid_operands(operands): + continue + + tile_k = BlackwellTcgen05Mma.shape_k(operands) * 4 + design_params["tile_shape"] = [ + (M, N, tile_k) for M, N, _ in design_params["tile_shape"] + ] + param_values = [design_params[name] for name in param_names] + for values in product(*param_values): + design = Sm100DesignMetadata(**dict(zip(param_names, values))) + operator_name = "cutedsl.PersistentDenseGemmEFCOperator_sm100_{layout}_A{A}_B{B}_out{out}_acc{acc}_{num_cta}cta_cluster{cluster}_tile{tile}{_tma_store}".format( + layout=strides_to_layout_string( + operands.A.stride, operands.B.stride, operands.out.stride + ), + A=operands.A.dtype, + B=operands.B.dtype, + out=operands.out.dtype, + acc=operands.accumulator_type, + num_cta=("2" if design.use_2cta_mma else "1"), + cluster=tuple_to_string(design.cluster_shape), + tile=tuple_to_string(design.tile_shape), + _tma_store="_tma_store" if design.use_tma_store else "", + ) + + metadata = OperatorMetadata( + operands=operands, + design=design, + operator_name=operator_name, + operator_class=cls, + supported_targets=TargetSm.get_supported_targets(design, operands), + epilogue=epilogue_metadata, + ) + + if PersistentDenseGemmEFCOperator._valid_metadata( + metadata + ) and metadata_filter(metadata): + operator_list.append(PersistentDenseGemmEFCOperator(metadata)) + + return operator_list diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/sm100_tgv_gemm.py b/operators/cutlass/operators/providers/cutedsl/gemm/sm100_tgv_gemm.py new file mode 100644 index 000000000..d125372ef --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/sm100_tgv_gemm.py @@ -0,0 +1,420 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute + +import cutlass.operators.utils.generate as generate +from cutlass.operators.arch import TargetSm +from cutlass.operators.arguments import ( + DenseTensor, + EpilogueArguments, + GemmArguments, +) +from cutlass.operators.metadata import ( + GemmOperandsMetadata, + OperatorMetadata, + Sm100DesignMetadata, +) +from cutlass.operators.mma import BlackwellTcgen05Mma +from cutlass.operators.providers.cutedsl import CuTeDSLProvider +from cutlass.operators.providers.cutedsl.operator import CuteDslOperator +from cutlass.operators.utils.common import tuple_to_string +from cutlass.operators.utils.device import to_cuda_stream +from cutlass.operators.utils.tensor import ( + normalized_major_stride, + strides_to_layout_string, +) + +from .implementations.sm100_tgv_gemm_impl import TgvGemmKernel + +if TYPE_CHECKING: + from collections.abc import Callable, Generator + + from cutlass.operators.artifact import CompiledArtifact + + +@cute.jit +def _tgv_bmm( + gemm_op: cutlass.Constexpr, + a_api: cute.Tensor, + b_api: cute.Tensor, + c_api: cute.Tensor, + stream: cuda.CUstream, +): + """Swap API operands to TGV convention and reorder batch-first to batch-last. + + cutlass.operators convention: C = A @ B + A = activation (L, M, K) row-major (K-contiguous) + B = weight (L, K, N) col-major (K-contiguous) + C = output (L, M, N) row-major (N-contiguous) + + TGV convention: C = A * B (where A=weight, B=activation) + a = weight (M_tgv, K, L) K-contiguous [M_tgv = N_api] + b = activation (N_tgv, K, L) K-contiguous [N_tgv = M_api] + c = output (M_tgv, N_tgv, L) M_tgv-contiguous + + The mapping: + B_api (L,K,N) -> select([2,1,0]) -> (N,K,L) = TGV a + A_api (L,M,K) -> select([1,2,0]) -> (M,K,L) = TGV b + C_api (L,M,N) -> select([2,1,0]) -> (N,M,L) = TGV c + """ + + # Promote rank-2 tensors to rank-3 by prepending batch dim of 1. + # Same pattern as PersistentDenseGemmKernel.__call__. + def add_batch_mode(tensor: cute.Tensor) -> cute.Tensor: + return cute.make_tensor( + tensor.iterator, + cute.prepend(tensor.layout, cute.make_layout(1), up_to_rank=3), + ) + + a_api = add_batch_mode(a_api) # (M, K) -> (1, M, K) + b_api = add_batch_mode(b_api) # (K, N) -> (1, K, N) + c_api = add_batch_mode(c_api) # (M, N) -> (1, M, N) + + # Swap + mode reorder: API (L-first) -> TGV (L-last) with operand swap + a_tgv = cute.make_tensor( + b_api.iterator, cute.select(b_api.layout, [2, 1, 0]) + ) # (L,K,N) -> (N,K,L) = TGV a + b_tgv = cute.make_tensor( + a_api.iterator, cute.select(a_api.layout, [1, 2, 0]) + ) # (L,M,K) -> (M,K,L) = TGV b + c_tgv = cute.make_tensor( + c_api.iterator, cute.select(c_api.layout, [2, 1, 0]) + ) # (L,M,N) -> (N,M,L) = TGV c + + gemm_op(a_tgv, b_tgv, c_tgv, stream) + + +@CuTeDSLProvider.register +class TgvGemmOperator(CuteDslOperator): + """Low-latency Blackwell GEMM kernel (TGV) for narrow-batch inference. + + Computes C = A @ B where A=activation, B=weight through the standard + cutlass.operators GemmArguments interface. Internally swaps operands to match + TGV's native convention (weight @ activation). + + .. note:: Tile shape semantics with operand swap: + + TGV internally swaps A/B operands: TGV_A=weight (API's B), + TGV_B=activation (API's A). + + - ``tile_shape[0]`` = CTA_M tiles TGV's M dim = API's N (out_features) + - ``tile_shape[1]`` = CTA_N tiles TGV's N dim = API's M (batch/tokens) + - ``tile_shape[2]`` = CTA_K tiles the shared K (contraction) dim + + This means ``tile_shape[0]=64`` tiles the weight output feature dim, and + ``tile_shape[1]=8`` tiles the batch dim — which is why small CTA_N suits + small batch. + + .. note:: Layout requirements: + + - A (activation): row-major, K-contiguous, stride (0,0,1) + - B (weight): col-major, K-contiguous, stride (0,1,0) + - C (output): row-major, N-contiguous, stride (0,0,1) + + .. note:: Only emitted for small M (batch) <= 64. + + .. note:: Features: + + - 1 SM mode (no 2CTA instructions), 1x1 cluster + - TMA loads for A and B, non-TMA-store epilogue + - 8-warp specialization: DMA_A, DMA_B, MMA, EPILOG + - No epilogue fusion support + """ + + supported_args_type = GemmArguments + designed_for_min_cc = 100 + + def __init__(self, metadata: OperatorMetadata): + super().__init__(metadata) + + self.impl = TgvGemmKernel( + acc_dtype=metadata.operands.accumulator_type, + cta_m=metadata.design.tile_shape[0], + cta_n=metadata.design.tile_shape[1], + cta_k=metadata.design.tile_shape[2], + num_ab_stage=metadata.design.num_smem_stages, + ) + + def _compile( + self, args: GemmArguments, target_sm: TargetSm | None = None + ) -> CompiledArtifact: + stream = cutlass.cute.runtime.make_fake_stream() + + compiled_kernel = self.cute_compile( + _tgv_bmm, + self.impl, + args.A.tensor, + args.B.tensor, + args.out.tensor, + stream, + target_sm=target_sm, + ) + return compiled_kernel + + def _run( + self, + args: GemmArguments, + compiled_artifact: CompiledArtifact, + stream, + workspace=None, + ) -> None: + stream = to_cuda_stream(stream) + compiled_gemm = compiled_artifact.compiled_obj + self.cute_run( + compiled_gemm, + args.A.tensor, + args.B.tensor, + args.out.tensor, + stream, + ) + + @staticmethod + def _valid_operands(operands: GemmOperandsMetadata) -> bool: + if not isinstance(operands, GemmOperandsMetadata): + return False + + # A and B must have the same dtype + if operands.A.dtype != operands.B.dtype: + return False + + abtype = operands.A.dtype + + # Supported A/B dtypes: float16, bfloat16, fp8 + if abtype not in [ + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ]: + return False + + # Only Float32 accumulator for TGV + if operands.accumulator_type != cutlass.Float32: + return False + + if operands.out.dtype in [cutlass.Float8E4M3FN, cutlass.Float8E5M2]: + return False + if operands.out.dtype not in [abtype, cutlass.Float32]: + return False + + # Layout constraints: + # A (activation) must be row-major (K-contiguous): stride (0, 0, 1) + if operands.A.stride != (0, 0, 1): + return False + + # B (weight) must be col-major (K-contiguous): stride (0, 1, 0) + if operands.B.stride != (0, 1, 0): + return False + + # C (output) must be row-major (N-contiguous): stride (0, 0, 1) + if operands.out.stride != (0, 0, 1): + return False + + return True + + @staticmethod + def _metadata_operands_from_args( + args: GemmArguments, + ) -> Generator[GemmOperandsMetadata, None, None]: + if any(not isinstance(arg, DenseTensor) for arg in [args.A, args.B, args.out]): + return + + if args.A.dtype != args.B.dtype: + return + + A_stride = normalized_major_stride( + args.A.shape, args.A.stride, prepend_zeros_to_rank=3 + ) + B_stride = normalized_major_stride( + args.B.shape, args.B.stride, prepend_zeros_to_rank=3 + ) + out_stride = normalized_major_stride( + args.out.shape, args.out.stride, prepend_zeros_to_rank=3 + ) + + # TGV requires specific layouts + if A_stride != (0, 0, 1): + return + if B_stride != (0, 1, 0): + return + if out_stride != (0, 0, 1): + return + + operands = generate.construct_dense_gemm_operands( + args.A.dtype, + args.B.dtype, + args.out.dtype, + args.accumulator_type, + A_stride, + B_stride, + out_stride, + alignment_bytes=16, + ) + + yield operands + + @staticmethod + def _metadata_operand_combinations() -> Generator[GemmOperandsMetadata, None, None]: + ab_dtypes = [ + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ] + + # TGV requires fixed layouts + a_stride = (0, 0, 1) # row-major (K-contiguous) + b_stride = (0, 1, 0) # col-major (K-contiguous) + out_stride = (0, 0, 1) # row-major (N-contiguous) + + for ab_dtype in ab_dtypes: + # Only Float32 accumulator + acc_dtype = cutlass.Float32 + + # Output: same as input (if not FP8) or Float32. + valid_out_dtypes = [cutlass.Float32] + if ab_dtype not in [cutlass.Float8E4M3FN, cutlass.Float8E5M2]: + valid_out_dtypes.insert(0, ab_dtype) + for out_dtype in valid_out_dtypes: + operands = generate.construct_dense_gemm_operands( + ab_dtype, + ab_dtype, + out_dtype, + acc_dtype, + a_stride, + b_stride, + out_stride, + alignment_bytes=16, + ) + yield operands + + @staticmethod + def _valid_metadata(metadata: OperatorMetadata) -> bool: + if not TgvGemmOperator._valid_operands(metadata.operands): + return False + + design = metadata.design + if not isinstance(design, Sm100DesignMetadata): + return False + + # TGV is 1 SM mode, no 2CTA, no TMA store, 1x1 cluster + if design.use_2cta_mma: + return False + if design.use_tma_store: + return False + if design.cluster_shape != (1, 1, 1): + return False + + tile = design.tile_shape + if tile is not None: + # CTA_M (tiles API's N / out_features): 64 only for now + if tile[0] not in [64]: + return False + # CTA_N (tiles API's M / batch): small values + if tile[1] not in [8, 16, 32]: + return False + # CTA_K: 128 fixed + if tile[2] != 128: + return False + + if metadata.epilogue is not None: + return False + + return True + + @classmethod + def _generate_operators( + cls, + metadata_filter: Callable[[OperatorMetadata], bool], + epilogue_args: EpilogueArguments = None, + target_sm: TargetSm | None = None, + args: GemmArguments = None, + ) -> list[TgvGemmOperator]: + if epilogue_args is not None: + return [] + + design_params = { + "mma_instruction_type": [BlackwellTcgen05Mma], + "use_2cta_mma": [False], + "tile_shape": [(64, n, 128) for n in [8, 16, 32]], + "cluster_shape": [(1, 1, 1)], + "use_tma_store": [False], + "num_smem_stages": [8], + } + + from itertools import product + + param_names = list(design_params.keys()) + param_values = [design_params[name] for name in param_names] + + operator_list = [] + + if args is not None: + operands_generator = cls._metadata_operands_from_args(args) + else: + operands_generator = cls._metadata_operand_combinations() + + for operands in operands_generator: + for values in product(*param_values): + design = Sm100DesignMetadata(**dict(zip(param_names, values))) + operator_name = ( + "cutedsl.TgvGemmOperator_sm100_{layout}_A{A}_B{B}_out{out}_acc{acc}_tile{tile}" + ).format( + layout=strides_to_layout_string( + operands.A.stride, operands.B.stride, operands.out.stride + ), + A=operands.A.dtype, + B=operands.B.dtype, + out=operands.out.dtype, + acc=operands.accumulator_type, + tile=tuple_to_string(design.tile_shape), + ) + + metadata = OperatorMetadata( + operands=operands, + design=design, + operator_name=operator_name, + operator_class=cls, + supported_targets=TargetSm.get_supported_targets(design, operands), + epilogue=None, + ) + + if TgvGemmOperator._valid_metadata(metadata) and metadata_filter( + metadata + ): + operator_list.append(cls(metadata)) + + return operator_list diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/sm80_tensorop_gemm.py b/operators/cutlass/operators/providers/cutedsl/gemm/sm80_tensorop_gemm.py new file mode 100644 index 000000000..166b3e7ba --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/sm80_tensorop_gemm.py @@ -0,0 +1,316 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +import itertools +from typing import TYPE_CHECKING + +import cutlass + +import cutlass.operators.utils.generate as generate +from cutlass.operators.arch import TargetSm +from cutlass.operators.arguments import ( + DenseTensor, + EpilogueArguments, + GemmArguments, +) +from cutlass.operators.metadata import ( + BLASDesignMetadata, + GemmOperandsMetadata, + OperatorMetadata, +) +from cutlass.operators.mma import AmpereMma +from cutlass.operators.providers.cutedsl import CuTeDSLProvider +from cutlass.operators.providers.cutedsl.operator import CuteDslOperator +from cutlass.operators.utils.common import tuple_to_string +from cutlass.operators.utils.device import to_cuda_stream +from cutlass.operators.utils.tensor import ( + normalized_major_stride, + strides_to_layout_string, +) + +from .implementations.sm80_tensorop_gemm_impl import Sm80TensorOpGemmKernel + +if TYPE_CHECKING: + from collections.abc import Callable, Generator + + from cutlass.operators.artifact import CompiledArtifact + + +@CuTeDSLProvider.register +class Sm80TensorOpGemmOperator(CuteDslOperator): + """This class implements batched matrix multiplication (C = A @ B) + + :note: In current version, A and B tensor must have the same data type + + :note: Supported A/B data types: + - Float16/BFloat16 + + :note: Supported accumulator data types: + - Float32 (for all floating point A/B data types) + :note: Supported C data types: + - Float16/BFloat16 (same as A, B) + + :note: Constraints: + - MMA tiler M must be 64/128 + - MMA tiler N must be 64/128 + """ + + supported_args_type = GemmArguments + designed_for_min_cc = 80 + + def __init__(self, metadata: OperatorMetadata): + super().__init__(metadata) + self.impl = Sm80TensorOpGemmKernel( + metadata.operands.A.dtype, + metadata.operands.B.dtype, + metadata.operands.out.dtype, + metadata.operands.accumulator_type, + ) + assert self.impl.num_stages == metadata.design.num_smem_stages + + def _compile( + self, args: GemmArguments, target_sm: TargetSm | None = None + ) -> CompiledArtifact: + stream = cutlass.cute.runtime.make_fake_stream() + + return self.cute_compile( + self.impl, + args.A.tensor, + args.B.tensor, + args.out.tensor, + stream, + target_sm=target_sm, + ) + + def _run( + self, + args: GemmArguments, + compiled_artifact: CompiledArtifact, + stream, + workspace=None, + ) -> None: + stream = to_cuda_stream(stream) + compiled_gemm = compiled_artifact.compiled_obj + self.cute_run( + compiled_gemm, args.A.tensor, args.B.tensor, args.out.tensor, stream + ) + + @staticmethod + def _valid_operands(operands: GemmOperandsMetadata) -> bool: + if not isinstance(operands, GemmOperandsMetadata): + return False + + # In current version, A and B tensor must have the same data type + if operands.A.dtype != operands.B.dtype: + return False + + # Supported A/B data types: + if operands.A.dtype not in [cutlass.Float16, cutlass.BFloat16]: + return False + + # Supported accumulator data types: + if operands.accumulator_type not in [cutlass.Float32]: + return False + + # Supported out data types: + if operands.out.dtype not in [ + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + ]: + return False + + return True + + @staticmethod + def _metadata_operands_from_args( + args: GemmArguments, + ) -> Generator[GemmOperandsMetadata, None, None]: + """Generator that yields all valid GemmOperandsMetadata that correspond to `args` and + that are supported by this Operator. + + :param args: The arguments to generate operands from. + :type args: GemmArguments + + :return: A generator of valid GemmOperandsMetadata that correspond to `args` and + that are supported by this Operator. + :rtype: Generator[GemmOperandsMetadata, None, None] + """ + if any(not isinstance(arg, DenseTensor) for arg in [args.A, args.B, args.out]): + return + + if args.A.dtype != args.B.dtype: + return + + valid_major_strides = {(0, 0, 1), (0, 1, 0)} + + A_stride = normalized_major_stride( + args.A.shape, args.A.stride, prepend_zeros_to_rank=3 + ) + B_stride = normalized_major_stride( + args.B.shape, args.B.stride, prepend_zeros_to_rank=3 + ) + out_stride = normalized_major_stride( + args.out.shape, args.out.stride, prepend_zeros_to_rank=3 + ) + + if any(x not in valid_major_strides for x in [A_stride, B_stride, out_stride]): + return + + operands = generate.construct_dense_gemm_operands( + args.A.dtype, + args.B.dtype, + args.out.dtype, + args.accumulator_type, + A_stride, + B_stride, + out_stride, + alignment_bytes=16, + ) + + yield operands + + @staticmethod + def _metadata_operand_combinations() -> Generator[GemmOperandsMetadata, None, None]: + """Generator that yields all valid GemmOperandsMetadata combinations + based on the validation rules in _valid_operands. + """ + # Supported A/B data types (must be the same) + ab_dtypes = [cutlass.Float16, cutlass.BFloat16] + valid_acc_dtypes = [cutlass.Float32] + valid_out_dtypes = ab_dtypes + + # Torch conventions (L, M, K) and (L, K, N) + row_major_stride = (0, 0, 1) + col_major_stride = (0, 1, 0) + + for ab_dtype in ab_dtypes: + for out_dtype in valid_out_dtypes: + for acc_dtype in valid_acc_dtypes: + for stride_A, stride_B, stride_out in itertools.product( + [row_major_stride, col_major_stride], repeat=3 + ): + yield generate.construct_dense_gemm_operands( + ab_dtype, + ab_dtype, + out_dtype, + acc_dtype, + stride_A, + stride_B, + stride_out, + alignment_bytes=16, + ) + + @staticmethod + def _valid_metadata(metadata: OperatorMetadata) -> bool: + if not Sm80TensorOpGemmOperator._valid_operands(metadata.operands): + return False + + design = metadata.design + if not isinstance(design, BLASDesignMetadata): + return False + + if design.tile_shape is None: + return False + + # MMA tiler N must be 32/64/128/256 + if design.tile_shape[1] not in [32, 64, 128]: + return False + + # MMA tiler M must be 64/128/256 + if design.tile_shape[0] not in [64, 128]: + return False + + if metadata.epilogue is not None: + return False + + return True + + @classmethod + def _generate_operators( + cls, + metadata_filter: Callable[[OperatorMetadata], bool], + epilogue_args: EpilogueArguments = None, + target_sm: TargetSm | None = None, + args: GemmArguments = None, + ) -> list[Sm80TensorOpGemmOperator]: + """Returns a list of all possible configurations of Sm80TensorOpGemmOperator that + adhere to constraints passed in under kwargs. + """ + design_params = { + "mma_instruction_type": [AmpereMma], + "tile_shape": [(M, N, 32) for M in [64, 128] for N in [64, 128]], + "cluster_shape": [(1, 1, 1)], + "num_smem_stages": [4], + } + + if epilogue_args is not None: + return [] + + from itertools import product + + # Get the list of tunable parameter names and their possible values + param_names = list(design_params.keys()) + param_values = [design_params[name] for name in param_names] + + operator_list = [] + + if args is not None: + operands_generator = cls._metadata_operands_from_args(args) + else: + operands_generator = cls._metadata_operand_combinations() + + for operands in operands_generator: + for values in product(*param_values): + design = BLASDesignMetadata(**dict(zip(param_names, values))) + operator_name = "cutedsl.Sm80TensorOpGemmOperator_{layout}_A{A}_B{B}_out{out}_acc{acc}_tile{tile}".format( + layout=strides_to_layout_string( + operands.A.stride, operands.B.stride, operands.out.stride + ), + A=operands.A.dtype, + B=operands.B.dtype, + out=operands.out.dtype, + acc=operands.accumulator_type, + tile=tuple_to_string(design.tile_shape), + ) + + metadata = OperatorMetadata( + operands=operands, + design=design, + operator_name=operator_name, + operator_class=cls, + supported_targets=TargetSm.get_supported_targets(design, operands), + epilogue=None, + ) + + if cls._valid_metadata(metadata) and metadata_filter(metadata): + operator_list.append(cls(metadata)) + + return operator_list diff --git a/operators/cutlass/operators/providers/cutedsl/gemm/sm90_static_persistent.py b/operators/cutlass/operators/providers/cutedsl/gemm/sm90_static_persistent.py new file mode 100644 index 000000000..bfbbf36cc --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/gemm/sm90_static_persistent.py @@ -0,0 +1,441 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +import itertools +from typing import TYPE_CHECKING + +import cutlass + +import cutlass.operators.utils.generate as generate +from cutlass.operators.arch import TargetSm +from cutlass.operators.arguments import ( + DenseTensor, + EpilogueArguments, + GemmArguments, +) +from cutlass.operators.metadata import ( + BLASDesignMetadata, + GemmOperandsMetadata, + OperatorMetadata, +) +from cutlass.operators.mma import HopperWgmma +from cutlass.operators.providers.cutedsl import CuTeDSLProvider +from cutlass.operators.providers.cutedsl.integration_utils.mma import ( + get_max_active_clusters, +) +from cutlass.operators.providers.cutedsl.operator import CuteDslOperator +from cutlass.operators.utils.common import tuple_to_string +from cutlass.operators.utils.device import to_cuda_stream +from cutlass.operators.utils.tensor import ( + normalized_major_stride, + strides_to_layout_string, +) + +from .implementations.sm90_static_persistent_impl import Sm90StaticPersistentGemmKernel + +if TYPE_CHECKING: + from collections.abc import Callable, Generator + + from cutlass.operators.artifact import CompiledArtifact + + +@CuTeDSLProvider.register +class Sm90StaticPersistentGemmOperator(CuteDslOperator): + """SM90 (Hopper) Persistent GEMM Operator using WGMMA and TMA. + + This class implements batched matrix multiplication (C = A @ B) with support + for various data types and architectural features specific to Hopper GPUs + with persistent tile scheduling and warp specialization. + + :note: In current version, A and B tensor must have the same data type + + :note: Supported A/B data types: + - Float16/BFloat16 + - Int8/Uint8 + - Float8E4M3FN/Float8E5M2 + + :note: Supported accumulator data types: + - Float32 (for all floating point A/B data types) + - Float16 (only for fp16 and fp8 A/B data types) + - Int32 (only for uint8/int8 A/B data types) + + :note: Supported C data types: + - Float32 (for float32 and int32 accumulator data types) + - Int32 (for float32 and int32 accumulator data types) + - Float16/BFloat16 (for fp16 and fp8 accumulator data types) + - Int8/Uint8 (for uint8/int8 accumulator data types) + - Float8E4M3FN/Float8E5M2 (for float32 accumulator data types) + + :note: Constraints: + - CTA tile M must be 64 or 128 + - CTA tile N must be 64, 128, or 256 + - Cluster shape M/N must be positive and power of 2, total cluster size <= 4 + - 8-bit types (fp8, int8) require K-major layout + """ + + supported_args_type = GemmArguments + designed_for_min_cc = 90 + + def __init__(self, metadata: OperatorMetadata): + super().__init__(metadata) + + tile_shape_mn = (metadata.design.tile_shape[0], metadata.design.tile_shape[1]) + cluster_shape_mn = ( + metadata.design.cluster_shape[0], + metadata.design.cluster_shape[1], + ) + + self.impl = Sm90StaticPersistentGemmKernel( + metadata.operands.accumulator_type, + tile_shape_mn, + cluster_shape_mn, + ) + + def _compile( + self, args: GemmArguments, target_sm: TargetSm | None = None + ) -> CompiledArtifact: + stream = cutlass.cute.runtime.make_fake_stream() + max_active_clusters = get_max_active_clusters(self.impl.cluster_shape_mn) + + return self.cute_compile( + self.impl, + args.A.tensor, + args.B.tensor, + args.out.tensor, + max_active_clusters, + stream, + target_sm=target_sm, + ) + + def _run( + self, + args: GemmArguments, + compiled_artifact: CompiledArtifact, + stream, + workspace=None, + ) -> None: + stream = to_cuda_stream(stream) + compiled_gemm = compiled_artifact.compiled_obj + self.cute_run( + compiled_gemm, args.A.tensor, args.B.tensor, args.out.tensor, stream + ) + + @staticmethod + def _valid_operands(operands: GemmOperandsMetadata) -> bool: + if not isinstance(operands, GemmOperandsMetadata): + return False + + # A and B tensor must have the same data type + if operands.A.dtype != operands.B.dtype: + return False + + abtype = operands.A.dtype + + # Supported A/B data types + if abtype not in [ + cutlass.Float16, + cutlass.BFloat16, + cutlass.Int8, + cutlass.Uint8, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ]: + return False + + # Supported accumulator data types + if operands.accumulator_type == cutlass.Float32: + if not abtype.is_float: + return False + elif operands.accumulator_type == cutlass.Float16: + if abtype not in [ + cutlass.Float16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ]: + return False + elif operands.accumulator_type == cutlass.Int32: + if abtype not in [cutlass.Uint8, cutlass.Int8]: + return False + else: + return False + + # Supported out data types + if operands.out.dtype == cutlass.Float32 or operands.out.dtype == cutlass.Int32: + if operands.accumulator_type not in [cutlass.Float32, cutlass.Int32]: + return False + elif ( + operands.out.dtype == cutlass.Float16 + or operands.out.dtype == cutlass.BFloat16 + ): + if operands.accumulator_type not in [ + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + ]: + return False + elif operands.out.dtype == cutlass.Int8 or operands.out.dtype == cutlass.Uint8: + if operands.accumulator_type not in [cutlass.Int32]: + return False + elif ( + operands.out.dtype == cutlass.Float8E4M3FN + or operands.out.dtype == cutlass.Float8E5M2 + ): + if operands.accumulator_type not in [cutlass.Float32]: + return False + else: + return False + + return True + + @staticmethod + def _metadata_operands_from_args( + args: GemmArguments, + ) -> Generator[GemmOperandsMetadata, None, None]: + """Generator that yields all valid GemmOperandsMetadata that correspond to `args` and + that are supported by this Operator. + + :param args: The arguments to generate operands from. + :type args: GemmArguments + + :return: A generator of valid GemmOperandsMetadata that correspond to `args` and + that are supported by this Operator. + :rtype: Generator[GemmOperandsMetadata, None, None] + """ + if any(not isinstance(arg, DenseTensor) for arg in [args.A, args.B, args.out]): + return + + if args.A.dtype != args.B.dtype: + return + + valid_major_strides = {(0, 0, 1), (0, 1, 0)} + + A_stride = normalized_major_stride( + args.A.shape, args.A.stride, prepend_zeros_to_rank=3 + ) + B_stride = normalized_major_stride( + args.B.shape, args.B.stride, prepend_zeros_to_rank=3 + ) + out_stride = normalized_major_stride( + args.out.shape, args.out.stride, prepend_zeros_to_rank=3 + ) + + if any(x not in valid_major_strides for x in [A_stride, B_stride, out_stride]): + return + + yield generate.construct_dense_gemm_operands( + args.A.dtype, + args.B.dtype, + args.out.dtype, + args.accumulator_type, + A_stride, + B_stride, + out_stride, + alignment_bytes=16, + ) + + @staticmethod + def _metadata_operand_combinations() -> Generator[GemmOperandsMetadata, None, None]: + """Generator that yields all valid GemmOperandsMetadata combinations + based on the validation rules in _valid_operands. + """ + # Supported A/B data types (must be the same) + ab_dtypes = [ + cutlass.Float16, + cutlass.BFloat16, + cutlass.Int8, + cutlass.Uint8, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ] + + row_major_stride = (0, 0, 1) + col_major_stride = (0, 1, 0) + + for ab_dtype in ab_dtypes: + # Determine valid accumulator types for this A/B dtype + valid_acc_dtypes = [] + + if ab_dtype.is_float: + valid_acc_dtypes.append(cutlass.Float32) + if ab_dtype in [ + cutlass.Float16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ]: + valid_acc_dtypes.append(cutlass.Float16) + else: # Int8, Uint8 + valid_acc_dtypes.append(cutlass.Int32) + + for acc_dtype in valid_acc_dtypes: + # Determine valid output types for this accumulator type + valid_out_dtypes = [] + + if acc_dtype == cutlass.Float32: + valid_out_dtypes.extend( + [ + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float32, + cutlass.Int32, + ] + ) + elif acc_dtype == cutlass.Int32: + valid_out_dtypes.extend( + [cutlass.Int8, cutlass.Uint8, cutlass.Float32, cutlass.Int32] + ) + elif acc_dtype == cutlass.Float16: + valid_out_dtypes.extend([cutlass.Float16]) + + for out_dtype in valid_out_dtypes: + # For 8-bit types, only K-major layout is supported + # K-major for A means row-major (0, 0, 1) stride pattern + # K-major for B means col-major (0, 1, 0) stride pattern + if ab_dtype.width == 8: + # Only K-major layouts for 8-bit + stride_combinations = [ + (row_major_stride, col_major_stride, row_major_stride), + (row_major_stride, col_major_stride, col_major_stride), + ] + else: + # All layout combinations for 16-bit types + stride_combinations = list( + itertools.product( + [row_major_stride, col_major_stride], repeat=3 + ) + ) + + for stride_A, stride_B, stride_out in stride_combinations: + yield generate.construct_dense_gemm_operands( + ab_dtype, + ab_dtype, + out_dtype, + acc_dtype, + stride_A, + stride_B, + stride_out, + alignment_bytes=16, + ) + + @staticmethod + def _valid_metadata(metadata: OperatorMetadata) -> bool: + if not Sm90StaticPersistentGemmOperator._valid_operands(metadata.operands): + return False + + design = metadata.design + if not isinstance(design, BLASDesignMetadata): + return False + + cluster_size_m, cluster_size_n, _ = design.cluster_shape + + # Cluster constraints for SM90 + if cluster_size_m not in [1, 2]: + return False + if cluster_size_n not in [1, 2]: + return False + if cluster_size_m * cluster_size_n > 4: + return False + + tile = design.tile_shape + + # CTA tile M must be 64 or 128 + if tile is not None and tile[0] not in [64, 128]: + return False + + # CTA tile N must be 64, 128, or 256 + if tile is not None and tile[1] not in [64, 128, 256]: + return False + + if metadata.epilogue is not None: + return False + + return True + + @classmethod + def _generate_operators( + cls, + metadata_filter: Callable[[OperatorMetadata], bool], + epilogue_args: EpilogueArguments = None, + target_sm: TargetSm | None = None, + args: GemmArguments = None, + ) -> list[Sm90StaticPersistentGemmOperator]: + """Returns a list of all possible configurations of Sm90StaticPersistentGemmOperator + that adhere to constraints passed in under kwargs. + """ + design_params = { + "mma_instruction_type": [HopperWgmma], + "tile_shape": [(M, N, 64) for M in [64, 128] for N in [64, 128, 256]], + "cluster_shape": [(M, N, 1) for M in [1, 2] for N in [1, 2] if M * N <= 4], + } + + if epilogue_args is not None: + return [] + + from itertools import product + + param_names = list(design_params.keys()) + param_values = [design_params[name] for name in param_names] + + operator_list = [] + + if args is not None: + operands_generator = cls._metadata_operands_from_args(args) + else: + operands_generator = cls._metadata_operand_combinations() + + for operands in operands_generator: + for values in product(*param_values): + design = BLASDesignMetadata(**dict(zip(param_names, values))) + operator_name = "cutedsl.Sm90StaticPersistentGemmOperator_{layout}_A{A}_B{B}_out{out}_acc{acc}_cluster{cluster}_tile{tile}".format( + layout=strides_to_layout_string( + operands.A.stride, operands.B.stride, operands.out.stride + ), + A=operands.A.dtype, + B=operands.B.dtype, + out=operands.out.dtype, + acc=operands.accumulator_type, + cluster=tuple_to_string(design.cluster_shape), + tile=tuple_to_string(design.tile_shape), + ) + + metadata = OperatorMetadata( + operands=operands, + design=design, + operator_name=operator_name, + operator_class=cls, + supported_targets=TargetSm.get_supported_targets(design, operands), + epilogue=None, + ) + + if cls._valid_metadata(metadata) and metadata_filter(metadata): + operator_list.append(cls(metadata)) + + return operator_list diff --git a/operators/cutlass/operators/providers/cutedsl/integration_utils/__init__.py b/operators/cutlass/operators/providers/cutedsl/integration_utils/__init__.py new file mode 100644 index 000000000..946e9b107 --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/integration_utils/__init__.py @@ -0,0 +1,40 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +CuTe DSL-specific integration utilities for wiring operators into the cutlass.operators backend. + +Functions here should depend on the CuTe DSL framework (cutlass.cute / cutlass.utils). +""" + +from . import builders, mma + +__all__ = [ + "builders", + "mma", +] diff --git a/operators/cutlass/operators/providers/cutedsl/integration_utils/builders.py b/operators/cutlass/operators/providers/cutedsl/integration_utils/builders.py new file mode 100644 index 000000000..ffd3328fa --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/integration_utils/builders.py @@ -0,0 +1,97 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""CuTe DSL-specific integration utilities for wiring operators into the cutlass.operators backend. + +Functions here should depend on the CuTe DSL framework (cutlass.cute / cutlass.utils). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import cutlass +import cutlass.cute as cute +from cutlass.cute import nvgpu + +if TYPE_CHECKING: + from cutlass.operators.arguments.gemm import GemmArguments + from cutlass.operators.metadata.operands.gemm import GemmOperandsMetadata + from cutlass.operators.utils.tensor import TensorWrapper + + +def make_ptr( + tensor: TensorWrapper, + nullptr: bool = False, + address_space: cutlass.AddressSpace = cutlass.AddressSpace.gmem, + assumed_align: int | None = 16, +) -> cute.Pointer: + """Shim for cute.runtime.make_ptr that accepts a TensorWrapper. + + :param tensor: The tensor to create a pointer for + :param nullptr: Use nullptr if True, otherwise use data_ptr from TensorWrapper + :param address_space: The address space of the pointer + :param assumed_align: Assumed alignment in bytes + """ + + def ptr(x): + return 0 if nullptr else x.data_ptr + + return cute.runtime.make_ptr( + tensor.dtype, ptr(tensor), address_space, assumed_align + ) + + +@dataclass(frozen=True) +class GemmMajorModes: + A: nvgpu.OperandMajorMode + B: nvgpu.OperandMajorMode + out: nvgpu.OutputMajorMode + + +def gemm_major_modes(gemm: GemmArguments | GemmOperandsMetadata) -> GemmMajorModes: + a_major_mode = ( + nvgpu.OperandMajorMode.K + if gemm.A.stride[-2:].index(1) == 1 + else nvgpu.OperandMajorMode.MN + ) + + b_major_mode = ( + nvgpu.OperandMajorMode.K + if gemm.B.stride[-2:].index(1) == 0 + else nvgpu.OperandMajorMode.MN + ) + + out_major_mode = ( + nvgpu.OutputMajorMode.N + if gemm.out.stride[-2:].index(1) == 1 + else nvgpu.OutputMajorMode.M + ) + + return GemmMajorModes(A=a_major_mode, B=b_major_mode, out=out_major_mode) diff --git a/operators/cutlass/operators/providers/cutedsl/integration_utils/mma.py b/operators/cutlass/operators/providers/cutedsl/integration_utils/mma.py new file mode 100644 index 000000000..c5e79d77b --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/integration_utils/mma.py @@ -0,0 +1,49 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import cutlass.utils as utils + +if TYPE_CHECKING: + from cutlass.operators.utils.layout import TupleRank2Or3 + + +def get_max_active_clusters(cluster_shape: TupleRank2Or3) -> int: + """:param cluster_shape: The shape of the cluster. + :type cluster_shape: tuple[int, int] + + :returns: The maximum number clusters of the provided shape that can fit on the current GPU + :rtype: int + """ + return utils.HardwareInfo().get_max_active_clusters( + cluster_shape[0] * cluster_shape[1] + ) diff --git a/operators/cutlass/operators/providers/cutedsl/operator.py b/operators/cutlass/operators/providers/cutedsl/operator.py new file mode 100644 index 000000000..a4ab1338c --- /dev/null +++ b/operators/cutlass/operators/providers/cutedsl/operator.py @@ -0,0 +1,90 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +import cutlass.cute as cute +from cutlass.cutlass_dsl import CuTeDSL + +from cutlass.operators.arch import TargetSm +from cutlass.operators.artifact import CompiledArtifact +from cutlass.operators.base import Operator +from cutlass.operators.config import GlobalOptions +from cutlass.operators.utils.tensor import TensorWrapper + + +class CuteDslOperator(Operator): + """Base class for all CuTe DSL operators""" + + def cute_compile(self, entry_point_fn, *fn_args, target_sm: TargetSm | None = None): + """Compiles an operator using CuTe DSL compile. + + This method is intended to provide a provider-wide consistent implementation of compile() for + all CuTe DSL operators, respecting GlobalOptions. + + :param entry_point_fn: The operator's entry-point function to compile + :param fn_args: All arguments to pass to the operator's compilation + :param target_sm: Target compute capability to compile the operator for. + If not provided, the operator will be compiled for the compute capability of the current device. + :type target_sm: TargetSm | None + :return: The compiled artifact + """ + # If no target_sm specified, determine the same as what CuTe DSL does. + # Explicitly set it as a compilation option so that the compiled_for sm + # reported in CompiledArtifact is accurate. + target_sm = target_sm or TargetSm( + CuTeDSL._get_dsl().envar.arch.removeprefix("sm_") + ) + + compile_args = [ + x.compile_time_tensor if isinstance(x, TensorWrapper) else x + for x in fn_args + ] + + compiled_fn = cute.compile[ + cute.GPUArch(f"sm_{target_sm}"), + cute.EnableTVMFFI(GlobalOptions().use_tvm_ffi), + ](entry_point_fn, *compile_args) + + return CompiledArtifact( + compiled_obj=compiled_fn, + operator_obj=self, + compiled_for=target_sm, + ) + + def cute_run(self, entry_point_fn, *fn_args): + """Extracts runtime tensors from TensorWrappers and runs the operator. + + :param entry_point_fn: The operator's entry-point function to run + :param fn_args: All arguments to pass to the operator + :return: The result of running the operator + """ + run_args = [ + x.runtime_tensor if isinstance(x, TensorWrapper) else x for x in fn_args + ] + return entry_point_fn(*run_args) diff --git a/operators/cutlass/operators/providers/provider.py b/operators/cutlass/operators/providers/provider.py new file mode 100644 index 000000000..e452a69c3 --- /dev/null +++ b/operators/cutlass/operators/providers/provider.py @@ -0,0 +1,153 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from cutlass.operators.arch import TargetSm +from cutlass.operators.base import Operator + +if TYPE_CHECKING: + from collections.abc import Callable + + from cutlass.operators.arguments import EpilogueArguments, RuntimeArguments + from cutlass.operators.metadata import OperatorMetadata + + +class Provider: + """Base class for operator providers with common implementation for operator registration and generation. + + Concrete providers should inherit from this class to gain automatic operator registration and generation functionality. + """ + + # Keyed by supported argument type for efficient lookup. + # Each subclass gets its own dict via __init_subclass__. + _operator_classes_by_argtype: dict[type[RuntimeArguments], list[type[Operator]]] + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + # Each subclass gets its own dict to avoid sharing state + cls._operator_classes_by_argtype = {} + + @classmethod + def generate_operators( + cls, + metadata_filter: Callable[[OperatorMetadata], bool] | None = None, + epilogue_args: EpilogueArguments | None = None, + target_sm: TargetSm | str | None = None, + args: RuntimeArguments | None = None, + ) -> list[Operator]: + """Return a list of operators that match the given filters. + + Args: + metadata_filter (Callable[[OperatorMetadata], bool]): boolean filter function to apply to the metadata + epilogue_args (EpilogueArguments | None): the epilogue arguments + target_sm (TargetSm | str | None): compute capability of device for which operators should be compiled + args (RuntimeArguments | None): the arguments to filter operators by. If None, all operators are considered + + Returns: + list[Operator]: list of all supported operator configurations + """ + target_sm = TargetSm.ensure(target_sm) + + # Determine which Operator classes to iterate over + if args is not None: + # Only consider Operator classes that support the requested argument type + operator_classes_to_use = cls._operator_classes_by_argtype.get( + type(args), [] + ) + else: + # No filtering - use all Operator classes + operator_classes_to_use = [ + k + for operator_classes in cls._operator_classes_by_argtype.values() + for k in operator_classes + ] + if target_sm is not None: + operator_classes_to_use = [ + k + for k in operator_classes_to_use + if k.designed_for_min_cc <= target_sm.cc + ] + + import os + + args = args if os.getenv("DISABLE_ARGUMENT_FILTERING", "0") == "0" else None + + operators_for_provider = [] + for operator_cls in operator_classes_to_use: + operators_for_provider.extend( + operator_cls.generate_operators( + metadata_filter, + epilogue_args=epilogue_args, + target_sm=target_sm, + args=args, + ) + ) + + return operators_for_provider + + @classmethod + def register(cls, operator_class: type[Operator]) -> type[Operator]: + """Register an Operator class with this provider. + + Intended to be used as a class decorator: + + .. code-block:: python + + @MyProvider.register + class MyOperator(Operator): + supported_args_type = MyArguments + ... + + + Args: + operator_class (type[Operator]): the Operator class to register + + Returns: + type[Operator]: the registered Operator class (unchanged) + + Raises: + TypeError: if operator_class does not define supported_args_type + """ + if not issubclass(operator_class, Operator): + raise TypeError( + f"class {operator_class.__name__} must inherit from Operator to register with a Provider" + ) + + # Stash the provider on the class to later report it in + # OperatorMetadata.provider. + # Users should read `operator.metadata.provider`, not this attribute. + operator_class._provider = cls + + args_type = operator_class.supported_args_type + if args_type not in cls._operator_classes_by_argtype: + cls._operator_classes_by_argtype[args_type] = [] + cls._operator_classes_by_argtype[args_type].append(operator_class) + return operator_class diff --git a/operators/cutlass/operators/status.py b/operators/cutlass/operators/status.py new file mode 100644 index 000000000..e1b764256 --- /dev/null +++ b/operators/cutlass/operators/status.py @@ -0,0 +1,68 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class Status: + """A simple status class to wrap an optional exception.""" + + error: Exception | None = None + + def __bool__(self) -> bool: + """Return ``True`` if this represents success, ``False`` otherwise. + + Example: + >>> if not (status := operator.supports(args)): + ... raise status.error + """ + return self.error is None + + @classmethod + def success(cls) -> Status: + """Create a successful status.""" + return cls() + + @classmethod + def fail(cls, error: str | Exception) -> Status: + """Create a failed status with an error. + + Args: + error (str | Exception): Error string or ``Exception`` describing the failure + """ + if isinstance(error, str): + error = ValueError(error) + return cls(error=error) + + def raise_on_error(self) -> None: + """Raise the stored exception if this status represents a failure.""" + if self.error is not None: + raise self.error diff --git a/operators/cutlass/operators/typing.py b/operators/cutlass/operators/typing.py new file mode 100644 index 000000000..10a52e108 --- /dev/null +++ b/operators/cutlass/operators/typing.py @@ -0,0 +1,73 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""Type markers for CUTLASS Operator API field annotations.""" + +from __future__ import annotations + +from typing import Any, Protocol, TypeAlias, runtime_checkable + +import cutlass.cute as cute + +from cutlass.operators.utils.tensor import TensorWrapper + + +@runtime_checkable +class _SupportsDLPack(Protocol): + """Protocol for Tensor-like objects that support the DLPack protocol. + + CUTLASS Operator API supports tensor-like objects (e.g., torch.Tensor, cute.Tensor) that + implement ``__dlpack__`` and ``__dlpack_device__``. + """ + + def __dlpack__(self, *, stream: Any = None) -> Any: + """Return a DLPack capsule representing the tensor data.""" + ... + + def __dlpack_device__(self) -> tuple[int, int]: + """Return the device type and device id as a tuple.""" + ... + + +TensorLike: TypeAlias = _SupportsDLPack | cute.Tensor | TensorWrapper +"""Tensor-like inputs accepted by the Operator API. + +Includes: + * Any DLPack-compatible tensor (``torch.Tensor``, ``jax.Array``, ``numpy.ndarray``, ...). + * ``cutlass.cute.Tensor`` (CuTe DSL host tensor; does not implement DLPack but is handled natively by ``TensorWrapper``). + * ``cutlass.operators.utils.tensor.TensorWrapper`` (implements DLPack). +""" + + +class NumericLike(Protocol): + """Type marker for fields that accept numeric-like types. + + Fields annotated with NumericLike accept the following: + * ``cutlass.Numeric`` + * ``torch.dtype`` + """ diff --git a/operators/cutlass/operators/utils/__init__.py b/operators/cutlass/operators/utils/__init__.py new file mode 100644 index 000000000..a36670cbc --- /dev/null +++ b/operators/cutlass/operators/utils/__init__.py @@ -0,0 +1,39 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from . import common, device, dtype, gemm, generate, layout, tensor + +__all__ = [ + "common", + "device", + "dtype", + "generate", + "layout", + "tensor", + "gemm", +] diff --git a/operators/cutlass/operators/utils/common.py b/operators/cutlass/operators/utils/common.py new file mode 100644 index 000000000..a11759890 --- /dev/null +++ b/operators/cutlass/operators/utils/common.py @@ -0,0 +1,84 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""Utilities for CUTLASS Operator API.""" + +from __future__ import annotations + +import importlib +import logging + +_logger = logging.getLogger(__name__) + + +def ceil_div(a: int, b: int) -> int: + """Compute the ceiling division of a by b. + + :param a: The dividend. + :type a: int + :param b: The divisor. + :type b: int + + :return: The ceiling division of a by b. + :rtype: int + """ + return (a + b - 1) // b + + +def round_up(a: int, b: int) -> int: + """Round a up to the nearest multiple of b. + + :param a: The value to round up. + :type a: int + :param b: The multiple to round up to. + :type b: int + + :return: The value of a rounded up to the nearest multiple of b. + :rtype: int + """ + return ((a + b - 1) // b) * b + + +def is_numpy_available() -> bool: + """Check if numpy is available.""" + return importlib.util.find_spec("numpy") is not None + + +def is_torch_available() -> bool: + """Check if torch is available.""" + return importlib.util.find_spec("torch") is not None + + +def is_jax_available() -> bool: + """Check if jax is available.""" + return importlib.util.find_spec("jax") is not None + + +def tuple_to_string(tuple: tuple[int, ...]) -> str: + """Convert a tuple of integers to an 'x'-separated string (e.g., (2, 3, 4) -> '2x3x4').""" + return "x".join(str(x) for x in tuple) if len(tuple) > 1 else str(tuple[0]) diff --git a/operators/cutlass/operators/utils/device.py b/operators/cutlass/operators/utils/device.py new file mode 100644 index 000000000..1aa95366b --- /dev/null +++ b/operators/cutlass/operators/utils/device.py @@ -0,0 +1,247 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +import importlib +import logging +import os +from collections.abc import Iterable +from typing import TYPE_CHECKING, Any + +import cutlass.operators.utils.common as common +from cutlass.operators.arch import TargetSm +from cutlass.operators.config import GlobalOptions +from cutlass.operators.status import Status + +_logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + import cuda + import torch + + +def _lazy_import(mod_name: str) -> Any: + """Internal utility to lazily import a module only when needed.""" + + class Lazy: + def __getattr__(self, name: str) -> Any: + module = importlib.import_module(mod_name) + return getattr(module, name) + + return Lazy() + + +def check_cuda_errors(result: list) -> Any: + """Check whether ``result`` contains a CUDA error and raise it, otherwise return the rest. + + Args: + result (list): Result of a ``cuda`` method, of the form + ``(cudaError_t, result...)`` -- an error code followed by any method results. + + Returns: + Any: Non-error-code results from ``result``. Returns ``None`` when the only element + is the (success) error code, the single result when there is exactly one, or the + tuple of remaining results otherwise. + + Raises: + RuntimeError: If the leading error code in ``result`` is non-zero. + """ + # `result` is of the format : (cudaError_t, result...) + err = result[0] + if err.value: + _lazy_import("cuda.cuda") + cuda_bindings_runtime = _lazy_import("cuda.bindings.runtime") + raise RuntimeError( + f"CUDA error: {cuda_bindings_runtime.cudaGetErrorString(err)[1].decode('utf-8')}" + ) + + if len(result) == 1: + return None + elif len(result) == 2: + return result[1] + else: + return result[1:] + + +def device_cc(device: int = 0) -> int: + """Query the device with ID ``device`` and return its compute capability. + + Args: + device (int): ID of the device to query. + + Returns: + int: Compute capability of the queried device (e.g., ``80`` for SM80). + """ + _lazy_import("cuda.cuda") + cuda_bindings_runtime = _lazy_import("cuda.bindings.runtime") + deviceProp = check_cuda_errors( + cuda_bindings_runtime.cudaGetDeviceProperties(device) + ) + major = str(deviceProp.major) + minor = str(deviceProp.minor) + return int(major + minor) + + +def device_or_env_target_sm( + device: int = 0, env_var: str = "CUTE_DSL_ARCH" +) -> TargetSm: + """Return the target SM from environment variable ``env_var`` or device ``device``. + + The environment variable takes precedence over device detection. + + Args: + device (int): CUDA device ID to query when the environment variable is not set. + env_var (str): Name of the environment variable that overrides device detection. + Expected format is ``"sm_"`` (e.g. ``"sm_100a"``). + + Returns: + TargetSm: The target SM corresponding to the environment variable or the physical device. + + Raises: + ValueError: If the environment variable is set but cannot be parsed as a valid ``TargetSm``. + """ + cute_dsl_arch = os.environ.get(env_var, None) + if cute_dsl_arch: + try: + return TargetSm(cute_dsl_arch.removeprefix("sm_")) + except ValueError as e: + raise ValueError(f"Invalid CUTE_DSL_ARCH: {cute_dsl_arch}. {e}") from e + cc = device_cc(device) + return TargetSm(cc=cc) + + +def device_or_env_supports( + requested_targets: TargetSm | str | Iterable[TargetSm | str], +) -> Status: + """Check whether the current device (or ``CUTE_DSL_ARCH``) supports ``requested_targets``. + + Args: + requested_targets (TargetSm | str | Iterable[TargetSm | str]): + Operator target(s) to check compatibility for. + + Returns: + Status: ``Status.success()`` if the device or ``CUTE_DSL_ARCH`` supports + ``requested_targets``; otherwise ``Status.fail(...)`` with a descriptive message. + """ + try: + # Allow passing a single target directly. + if isinstance(requested_targets, str) or not isinstance( + requested_targets, Iterable + ): + requested_targets = (requested_targets,) + + device = device_or_env_target_sm() + if not device.supports_operators_from(requested_targets): + return Status.fail( + f"Device/CUTE_DSL_ARCH {device} does not support {requested_targets}." + ) + return Status.success() + except Exception as e: + return Status.fail(f"Failed to get device or CUTE_DSL_ARCH: {e}") + + + +def device_l2_cache_size(device: int | None = None) -> int: + """Return L2 cache size in bytes for the requested device. + + Args: + device (int | None): CUDA device ID to query. If ``None``, uses the current device. + + Returns: + int: L2 cache size in bytes. + """ + cuda_bindings_runtime = _lazy_import("cuda.bindings.runtime") + if device is None: + device = check_cuda_errors(cuda_bindings_runtime.cudaGetDevice()) + device_prop = check_cuda_errors( + cuda_bindings_runtime.cudaGetDeviceProperties(device) + ) + return device_prop.l2CacheSize + + +def to_cuda_stream( + stream: cuda.bindings.driver.CUstream | torch.cuda.Stream, + skip_if_ffi: bool = True, +) -> cuda.bindings.driver.CUstream: + """Convert the provided stream to a ``cuda.bindings.driver.CUstream``. + + Args: + stream (cuda.bindings.driver.CUstream | torch.cuda.Stream): + The stream to convert. + skip_if_ffi (bool): If ``True`` and TVM-FFI is enabled in + :class:`~cutlass.operators.config.GlobalOptions`, return ``stream`` unchanged + (TVM-FFI handles stream types directly, including raw int handles). + + Returns: + cuda.bindings.driver.CUstream: The converted stream. + + Raises: + ValueError: If ``stream`` is of an unsupported type. + """ + if skip_if_ffi and GlobalOptions().use_tvm_ffi: + # TVM-FFI can directly handle streams of various types, including raw int handles + return stream + + cuda = _lazy_import("cuda.bindings.driver") + if isinstance(stream, cuda.CUstream): + return stream + if common.is_torch_available(): + import torch + + if isinstance(stream, torch.cuda.Stream): + return cuda.CUstream(stream.cuda_stream) + raise ValueError(f"Unsupported stream type: {type(stream)}") + + +def get_cuda_version() -> tuple[int, int] | None: + """Get the CUDA runtime version as a ``(major, minor)`` tuple. + + Returns: + tuple[int, int] | None: CUDA version tuple ``(major, minor)``, or ``None`` + if the version cannot be determined. + + Example: + >>> get_cuda_version() + (12, 4) # CUDA 12.4 + """ + try: + cuda_bindings_runtime = _lazy_import("cuda.bindings.runtime") + err, version = cuda_bindings_runtime.cudaRuntimeGetVersion() + if err.value == 0: + # Version is encoded as major * 1000 + minor * 10 + major = version // 1000 + minor = (version % 1000) // 10 + return (major, minor) + else: + _logger.debug( + "cudaRuntimeGetVersion returned non-zero error code: %s", err.value + ) + except Exception as e: + _logger.debug("Failed to import cuda.bindings.runtime: %s", e) + return None diff --git a/operators/cutlass/operators/utils/dtype.py b/operators/cutlass/operators/utils/dtype.py new file mode 100644 index 000000000..34abfd275 --- /dev/null +++ b/operators/cutlass/operators/utils/dtype.py @@ -0,0 +1,213 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from cutlass.cute.typing import ( + BFloat16, + Boolean, + Float4E2M1FN, + Float8E4M3B11FNUZ, + Float8E4M3FN, + Float8E5M2, + Float8E8M0FNU, + Float16, + Float32, + Float64, + Int8, + Int16, + Int32, + Int64, + Numeric, + Uint8, + Uint16, + Uint64, + Float4E2M1FN, +) + +import cutlass.operators.utils.common as common + +if TYPE_CHECKING: + import jax.numpy as jnp + import torch + + +def get_torch_to_cutlass_dtype_map() -> dict[torch.dtype, type[Numeric]]: + """Return the mapping from torch dtypes to their CUTLASS ``Numeric`` types.""" + import torch + + torch_dtype_map = { + torch.float64: Float64, + torch.float32: Float32, + torch.float16: Float16, + torch.bfloat16: BFloat16, + torch.int64: Int64, + torch.int32: Int32, + torch.int16: Int16, + torch.int8: Int8, + torch.uint8: Uint8, + torch.uint16: Uint16, + torch.uint64: Uint64, + torch.bool: Boolean, + torch.float8_e5m2: Float8E5M2, + torch.float8_e4m3fn: Float8E4M3FN, + torch.float8_e4m3fnuz: Float8E4M3B11FNUZ, + torch.float8_e8m0fnu: Float8E8M0FNU, + # NB: float4_e2m1fn_x2 packs 2 FP4 values per byte, so torch shapes + # are physical (K/2). TensorWrapper converts to logical shapes/strides + # automatically; all operator-level code should go through TensorWrapper. + torch.float4_e2m1fn_x2: Float4E2M1FN, + } + return torch_dtype_map + + +def cutlass_type_from_torch_type(dtype) -> type[Numeric]: + """Convert a torch dtype to a cutlass dtype. + + Args: + dtype (torch.dtype): The torch dtype to convert. + + Returns: + The cutlass dtype. + + Raises: + KeyError: If the dtype is not supported. + """ + torch_dtype_map = get_torch_to_cutlass_dtype_map() + try: + return torch_dtype_map[dtype] + except KeyError as err: + raise KeyError(f"Unsupported dtype: {dtype}") from err + + +def torch_type_from_cutlass_type(dtype) -> torch.dtype: + """Convert a cutlass dtype to a torch dtype.""" + torch_dtype_map = { + value: key for key, value in get_torch_to_cutlass_dtype_map().items() + } + try: + return torch_dtype_map[dtype] + except KeyError as err: + raise KeyError(f"Unsupported dtype: {dtype}") from err + + +def get_jax_to_cutlass_dtype_map() -> dict[jnp.dtype, type[Numeric]]: + """Create a map between jax dtype and cutlass dtype.""" + import jax.numpy as jnp + + # jax tensor.dtype returns numpy.dtype instances, so use jnp.dtype() as keys + jax_dtype_map = { + jnp.dtype("float64"): Float64, + jnp.dtype("float32"): Float32, + jnp.dtype("float16"): Float16, + jnp.dtype("bfloat16"): BFloat16, + jnp.dtype("int64"): Int64, + jnp.dtype("int32"): Int32, + jnp.dtype("int16"): Int16, + jnp.dtype("int8"): Int8, + jnp.dtype("uint8"): Uint8, + jnp.dtype("uint16"): Uint16, + jnp.dtype("uint64"): Uint64, + jnp.dtype("bool"): Boolean, + jnp.float8_e5m2.dtype: Float8E5M2, + jnp.float8_e4m3fn.dtype: Float8E4M3FN, + jnp.float8_e8m0fnu.dtype: Float8E8M0FNU, + } + return jax_dtype_map + + +def cutlass_type_from_jax_type(dtype) -> type[Numeric]: + """Convert a jax dtype to a cutlass dtype.""" + jax_dtype_map = get_jax_to_cutlass_dtype_map() + try: + return jax_dtype_map[dtype] + except KeyError as err: + raise KeyError(f"Unsupported dtype: {dtype}") from err + + +def jax_type_from_cutlass_type(dtype) -> jnp.dtype: + """Convert a cutlass dtype to a jax dtype.""" + jax_dtype_map = { + value: key for key, value in get_jax_to_cutlass_dtype_map().items() + } + try: + return jax_dtype_map[dtype] + except KeyError as err: + raise KeyError(f"Unsupported dtype: {dtype}") from err + + +def torch_storage_packing_factor(dtype: type[Numeric]) -> int: + """Return the number of logical elements per PyTorch storage element. + + For most types this is 1. For subbyte types like Float4E2M1FN, PyTorch + uses a packed representation (``float4_e2m1fn_x2``) that stores 2 logical + elements per byte, so this returns 2. + + Args: + dtype: A CUTLASS numeric type. + + Returns: + int: The packing factor (logical elements per stored element). + """ + import torch + + torch_dtype = torch_type_from_cutlass_type(dtype) + torch_bits = torch_dtype.itemsize * 8 + return torch_bits // dtype.width + + +def to_cutlass_type(dtype) -> type[Numeric]: + """Convert a dtype to a cutlass dtype. + + Args: + dtype (torch.dtype | jnp.dtype | np.dtype): The dtype to convert + + Returns: + The cutlass dtype. + + Raises: + KeyError: If the dtype is not supported. + """ + if isinstance(dtype, type) and issubclass(dtype, Numeric): + return dtype + + converters = [] + if common.is_torch_available(): + converters.append(cutlass_type_from_torch_type) + if common.is_jax_available(): + converters.append(cutlass_type_from_jax_type) + + # Iterate through the available converters and return the first one that succeeds. + for converter in converters: + try: + return converter(dtype) + except KeyError: # noqa: PERF203 + continue + raise KeyError(f"Unsupported dtype: {dtype}") diff --git a/operators/cutlass/operators/utils/gemm.py b/operators/cutlass/operators/utils/gemm.py new file mode 100644 index 000000000..8bc26ab71 --- /dev/null +++ b/operators/cutlass/operators/utils/gemm.py @@ -0,0 +1,96 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""Framework-agnostic argument validation for GEMM kernel integration.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from cutlass.operators.status import Status + +if TYPE_CHECKING: + from cutlass.operators.arguments import RuntimeArguments + + +def strip_batch_mode(dims: tuple[int, ...]) -> tuple[int, int]: + """Strips the batch mode from the dimensions. + + :throws: ValueError if the dimensions are not [L, M, K] or [M, K]. + """ + if len(dims) == 3: + return (dims[1], dims[2]) + elif len(dims) == 2: + return (dims[0], dims[1]) + else: + raise ValueError(f"Invalid dimensions: {dims}") + + +def tensor_rank_2_or_3(args: RuntimeArguments) -> Status: + """Checks that the arguments are valid for a tensor of rank 2 or 3.""" + if len(args.A.shape) < 2 or len(args.A.shape) > 3: + return Status.fail( + f"A must be a tensor of rank 2 or 3 (L=1, M, K), got {args.A.shape}" + ) + if len(args.B.shape) < 2 or len(args.B.shape) > 3: + return Status.fail( + f"B must be a tensor of rank 2 or 3 (L=1, K, N), got {len(args.B.shape)}" + ) + if len(args.out.shape) < 2 or len(args.out.shape) > 3: + return Status.fail( + f"out must be a tensor of rank 2 or 3 (L=1, M, N), got {len(args.out.shape)}" + ) + return Status.success() + + +def tensor_reduction_mode_matches(args: RuntimeArguments) -> Status: + """Checks that A's K dimension matches B's K dimension.""" + if args.A.shape[-1] != args.B.shape[-2]: + return Status.fail( + f"A's K dimension ({args.A.shape[-1]}) must be equal to B's " + f"K dimension ({args.B.shape[-2]}). " + f"A shape (L, M, K): {args.A.shape}, B shape (L, K, N): {args.B.shape}" + ) + return Status.success() + + +def tensor_output_shape_matches(args: RuntimeArguments) -> Status: + """Checks that the output tensor's M and N dimensions match A and B respectively.""" + if args.out.shape[-2] != args.A.shape[-2]: + return Status.fail( + f"out's M dimension ({args.out.shape[-2]}) must be equal to A's " + f"M dimension ({args.A.shape[-2]}). " + f"A shape (L, M, K): {args.A.shape}, out shape (L, M, N): {args.out.shape}" + ) + if args.out.shape[-1] != args.B.shape[-1]: + return Status.fail( + f"out's N dimension ({args.out.shape[-1]}) must be equal to B's " + f"N dimension ({args.B.shape[-1]}). " + f"B shape (L, K, N): {args.B.shape}, out shape (L, M, N): {args.out.shape}" + ) + return Status.success() diff --git a/operators/cutlass/operators/utils/generate.py b/operators/cutlass/operators/utils/generate.py new file mode 100644 index 000000000..4d01cdcef --- /dev/null +++ b/operators/cutlass/operators/utils/generate.py @@ -0,0 +1,100 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""Utilities for assisting with operator metadata generation.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from cutlass.cute.typing import Numeric + + from cutlass.operators.metadata.operands.gemm import GemmOperandsMetadata + + +def construct_dense_gemm_operands( + a_dtype: Numeric, + b_dtype: Numeric, + out_dtype: Numeric, + acc_dtype: Numeric, + stride_A: tuple[int, ...], + stride_B: tuple[int, ...], + stride_out: tuple[int, ...], + alignment_bytes: int, +) -> GemmOperandsMetadata: + """Constructs a GemmOperandsMetadata object based on the provided data types and strides. + + Each operand is assumed to be a dense tensor. + + :param a_dtype: The data type of the A operand. + :type a_dtype: cutlass.Numeric + :param b_dtype: The data type of the B operand. + :type b_dtype: cutlass.Numeric + :param out_dtype: The data type of the out operand. + :type out_dtype: cutlass.Numeric + :param acc_dtype: The data type of the accumulator operand. + :type acc_dtype: cutlass.Numeric + :param stride_A: The stride of the A operand. + :type stride_A: tuple[int, ...] + :param stride_B: The stride of the B operand. + :type stride_B: tuple[int, ...] + :param stride_out: The stride of the out operand. + :type stride_out: tuple[int, ...] + :param alignment_bytes: The alignment of the operands, in bytes. + :type alignment_bytes: int + + :return: A GemmOperandsMetadata object. + :rtype: GemmOperandsMetadata + """ + import cutlass.operators.metadata as metadata + + a_divisibility = alignment_bytes * 8 // a_dtype.width + b_divisibility = alignment_bytes * 8 // b_dtype.width + out_divisibility = alignment_bytes * 8 // out_dtype.width + + return metadata.GemmOperandsMetadata( + A=metadata.DenseTensorConstraints( + dtype=a_dtype, + stride=stride_A, + divisibility=a_divisibility, + ), + B=metadata.DenseTensorConstraints( + dtype=b_dtype, + stride=stride_B, + divisibility=b_divisibility, + ), + out=metadata.DenseTensorConstraints( + dtype=out_dtype, + stride=stride_out, + divisibility=out_divisibility, + ), + accumulator_type=acc_dtype, + ) + + diff --git a/operators/cutlass/operators/utils/layout.py b/operators/cutlass/operators/utils/layout.py new file mode 100644 index 000000000..85c787493 --- /dev/null +++ b/operators/cutlass/operators/utils/layout.py @@ -0,0 +1,80 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +import enum +from enum import auto as enum_auto + +# All current Gemm shapes/strides in cutlass.operators are rank 2 or 3. +TupleRank2Or3 = tuple[int, int, int] | tuple[int, int] + + +def normalized_row_major() -> tuple[int, int, int]: + """Returns the normalized row major strides for a rank 3 tensor. + + Leading dimension has stride=1. All others are 0. + Example: shape (L=batch, M=row, K=column) => stride (0, 0, 1) + + Uses definition provided by normalized_major_stride. + """ + return (0, 0, 1) + + +def normalized_col_major() -> tuple[int, int, int]: + """Returns the normalized column major strides for a rank 3 tensor. + + Leading dimension has stride=1. All others are 0. + Example: shape (L=batch, M=row, K=column) => stride (0, 1, 0) + + Uses definition provided by normalized_major_stride. + """ + return (0, 1, 0) + + +class LayoutType(enum.Enum): + ColumnMajor = enum_auto() + RowMajor = enum_auto() + ColumnMajorInterleaved2 = enum_auto() + RowMajorInterleaved2 = enum_auto() + ColumnMajorInterleaved32 = enum_auto() + RowMajorInterleaved32 = enum_auto() + ColumnMajorInterleaved64 = enum_auto() + RowMajorInterleaved64 = enum_auto() + TensorNWC = enum_auto() + TensorNHWC = enum_auto() + TensorNDHWC = enum_auto() + TensorNCHW = enum_auto() + TensorNGHWC = enum_auto() + TensorNC32HW32 = enum_auto() + TensorNC64HW64 = enum_auto() + TensorC32RSK32 = enum_auto() + TensorC64RSK64 = enum_auto() + TensorKCS = enum_auto() + TensorKCSR = enum_auto() + TensorKCSRT = enum_auto() diff --git a/operators/cutlass/operators/utils/tensor.py b/operators/cutlass/operators/utils/tensor.py new file mode 100644 index 000000000..d8e60454b --- /dev/null +++ b/operators/cutlass/operators/utils/tensor.py @@ -0,0 +1,575 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from typing import Literal + + import cutlass.cute as cute + from cutlass.cute.typing import Numeric + + from cutlass.operators.typing import TensorLike + +import cuda.bindings.driver as cuda_driver + +import cutlass.cute as cute +from cutlass.cute.runtime import from_dlpack + +import cutlass.operators.utils.common as common +import cutlass.operators.utils.dtype as dtype_utils +from cutlass.operators.config import GlobalOptions + + +def major_mode(shape: tuple[int, ...], stride: tuple[int, ...]) -> int | None: + """Get the major mode of a tensor from its shape and stride. + + This is the first mode with stride[mode]==1 and shape[mode]>1. + If the tensor has a single element, the major mode is the final mode. + Modes with both stride of 1 and shape of 1 are treated + as though they have stride of 0. + + :param shape: The shape of the tensor. + :type shape: tuple[int, ...] + :param stride: The stride of the tensor. + :type stride: tuple[int, ...] + + :return: The major mode of the tensor, or None if the tensor has no major mode. + :rtype: int | None + """ + if all(s == 1 for s in shape): + return len(shape) - 1 + + if not any(s == 1 for s in stride): + return None + + updated_stride = [s if sz != 1 else 0 for s, sz in zip(stride, shape)] + if 1 not in updated_stride: + return None + return updated_stride.index(1) + + +def major_mode_from_tensor(tensor: TensorWrapper) -> int | None: + """Get the major mode of a tensor. + + See ``major_mode`` for more details. + + :param tensor: The tensor to get the major mode of. + :type tensor: TensorWrapper + + :return: The major mode of the tensor, or None if the tensor has no major mode. + :rtype: int | None + """ + return major_mode(tensor.shape, tensor.stride) + + +def normalized_major_stride( + shape: tuple[int, ...], stride: tuple[int, ...], prepend_zeros_to_rank: int = None +) -> tuple[int, ...] | None: + """Get the normalized major stride of a tensor. + + This is a tuple of rank equal to that of stride, + where the major mode has stride 1 and the other modes have stride 0. + + :param shape: The shape of the tensor. + :type shape: tuple[int, ...] + :param stride: The stride of the tensor. + :type stride: tuple[int, ...] + :param prepend_zeros_to_rank: If set, prepends zeros to the returned stride until the given rank + is reached. + :type prepend_zeros_to_rank: int | None + + :return: The normalized major stride of the tensor, or None if the tensor has no major mode. + :rtype: tuple[int, ...] | None + """ + major_mode_idx = major_mode(shape, stride) + if major_mode_idx is None: + return None + + extra_zeros = ( + tuple((0,) * max(prepend_zeros_to_rank - len(shape), 0)) + if prepend_zeros_to_rank is not None + else () + ) + + return extra_zeros + tuple( + 0 if i != major_mode_idx else 1 for i in range(len(stride)) + ) + + +def get_stride_order(stride: tuple[int, ...]) -> tuple[int, ...]: + """Return the order of the modes of a stride. + + The returned tuple contains indices of modes in `stride`, sorted from outermost to innermost. + + For example, for a stride (1024, 1, 512), the stride order is (0, 2, 1). + + :param stride: The stride of the tensor. + :type stride: tuple[int, ...] + + :return: The order of the modes of the stride. + :rtype: tuple[int, ...] + """ + # The code below performs an == reverse argsort on the stride: + # indices = range(len(stride)) + # Sort indices using comparison between stride[i] and stride[j] when + # sorting indices i and j. Sort in descending order. + # Example: For a stride (1024, 1, 512), the reverse argsort is (0, 2, 1). + return tuple(sorted(range(len(stride)), key=stride.__getitem__, reverse=True)) + + +def get_stride_rank(stride: tuple[int, ...]) -> tuple[int, ...]: + """Returns the rank of the each mode in the stride of a tensor. + + For a stride of rank N, the mode with the smallest stride will + have stride rank 0 and the mode with the largest stride will have stride rank N-1. + + For example, for a stride (1024, 1, 512), the stride rank is (2, 0, 1). + + :param stride: The stride of the tensor. + :type stride: tuple[int, ...] + + :return: The rank of the each mode in the stride of the tensor. + :rtype: tuple[int, ...] + """ + # The code below performs an argsort on the stride: + # indices = range(len(stride)) + # Sort indices using comparison between stride[i] and stride[j] when + # sorting indices i and j. + # Example: For a stride (1024, 1, 512), the argsorted is (1, 2, 0). + argsorted = tuple(sorted(range(len(stride)), key=stride.__getitem__)) + + # Set the stride order of each mode to the index of the mode in the + # argsorted tuple. + # Example: For a stride (1024, 1, 512), the argsorted is (1, 2, 0), + # so the stride rank is (2, 0, 1). + res = [-1] * len(stride) + for i, idx in enumerate(argsorted): + res[idx] = i + return tuple(res) + + +def is_torch_tensor(inp) -> bool: + """Check if the input is a torch tensor.""" + if common.is_torch_available(): + import torch + + return isinstance(inp, torch.Tensor) + return False + + +def is_torch_fake_tensor(inp) -> bool: + """Check if the input is a torch fake tensor.""" + if common.is_torch_available(): + import torch + + return isinstance(inp, torch._subclasses.fake_tensor.FakeTensor) + return False + + +def is_numpy_tensor(inp) -> bool: + """Check if the input is a numpy tensor.""" + if common.is_numpy_available(): + import numpy as np + + return isinstance(inp, np.ndarray) + return False + + +def is_jax_array(inp) -> bool: + """Check if the input is a JAX array.""" + if common.is_jax_available(): + import jax + + return isinstance(inp, jax.Array) + return False + + +def _compute_row_major_contiguous_stride(shape: tuple[int, ...]) -> tuple[int, ...]: + """Compute row-major (C-contiguous) strides from a shape tuple.""" + strides = [1] * len(shape) + for i in range(len(shape) - 2, -1, -1): + strides[i] = strides[i + 1] * shape[i + 1] + return tuple(strides) + + +def _validate_jax_array(tensor) -> None: + """Validate that a JAX array is suitable for use with CUTLASS operators. + + Checks that the array is on a single GPU device, is not sharded, and has + a row-major (C-contiguous) layout. + """ + import jax + + if tensor.is_deleted(): + raise ValueError("JAX array has been deleted and cannot be used.") + + if hasattr(tensor, "sharding") and not isinstance( + tensor.sharding, jax.sharding.SingleDeviceSharding + ): + raise ValueError( + f"Sharded JAX arrays are not supported. " + f"Got sharding type: {type(tensor.sharding).__name__}. " + f"Please use a single-device array." + ) + + device = tensor.devices().pop() + if device.platform != "gpu": + raise ValueError( + f"JAX array must reside on a GPU device, but got device: {device} " + f"(platform='{device.platform}'). Move the array to GPU first." + ) + + if not _is_jax_array_c_contiguous(tensor): + raise ValueError( + "Only row-major (C-contiguous) JAX arrays are currently supported." + ) + + +def _is_jax_array_c_contiguous(tensor) -> bool: + """Check whether a JAX array has C-contiguous (row-major) layout.""" + shape = tensor.shape + if any(s == 0 for s in shape): + return True + expected_size = 1 + for s in shape: + expected_size *= s + return tensor.size == expected_size + + +class TensorWrapper: + """Wrapper class for supporting compilation and execution both with and without TVM-FFI. + + When using TVM-FFI, one can pass a framework-level tensor (e.g., torch.Tensor) + to the JIT function at run time, but not at compile time. At compile time, one + must use a `_FakeTensor` to specify the tensor. + + When not using TVM-FFI, one passes in a cute.Tensor at both compile and run time. + + This class contains two key members: + - runtime_tensor: The tensor to use at run time. + - compile_time_tensor: The tensor to use at compile time. + + Users of this class should access each of these underlying members for execution + and runtime, respectively. Users of this class do not need to know whether TVM-FFI + is enabled. + """ + + _alignment_bytes: int + _runtime_tensor: TensorLike + compile_time_tensor: TensorLike + _shape: tuple[int, ...] + _stride: tuple[int, ...] + _data_ptr: int + + def __init__(self, tensor: Any, alignment_bytes: int = 16): + self._alignment_bytes = alignment_bytes + if isinstance(tensor, TensorWrapper): + # Instead of double-nested TensorWrapper, we grab runtime tensor and flatten + self._alignment_bytes = tensor._alignment_bytes + self._runtime_tensor = tensor._runtime_tensor + self._shape = tensor._shape + self._stride = tensor._stride + self._data_ptr = tensor._data_ptr + elif isinstance(tensor, cute.Tensor): + # Regardless of whether TVM-FFI is enabled, if the tensor passed in is a cute.Tensor, + # it can be used as the runtime tensor and compile time tensor. + self._runtime_tensor = tensor + self.compile_time_tensor = tensor + self._shape = tensor.shape + self._stride = tensor.stride + self._data_ptr = tensor.iterator._pointer + elif is_torch_fake_tensor(tensor) or ( + is_torch_tensor(tensor) and GlobalOptions().use_tvm_ffi + ): + if is_torch_fake_tensor(tensor): + self._data_ptr = 0 + self._runtime_tensor = None + else: + self._runtime_tensor = tensor + self._data_ptr = tensor.data_ptr() + + dtype = dtype_utils.cutlass_type_from_torch_type(tensor.dtype) + + # For packed sub-byte torch dtypes (e.g. float4_e2m1fn_x2 which + # stores 2 FP4 values per byte), convert from physical (packed) + # shape/stride to logical (unpacked) shape/stride so that + # TensorWrapper reports the same logical layout as from_dlpack. + packing = tensor.element_size() * 8 // dtype.width + if packing > 1: + phys_shape = tuple(tensor.shape) + phys_stride = tensor.stride() + contiguous_dim = phys_stride.index(1) + self._shape = tuple( + s * packing if i == contiguous_dim else s + for i, s in enumerate(phys_shape) + ) + self._stride = tuple( + s * packing if i != contiguous_dim else s + for i, s in enumerate(phys_stride) + ) + else: + self._shape = tuple(tensor.shape) + self._stride = tensor.stride() + stride_order = get_stride_rank(tensor.stride()) + major_mode_idx = stride_order.index(0) + shape = [cute.SymInt() for _ in range(tensor.dim())] + shape[major_mode_idx] = cute.SymInt( + divisibility=alignment_bytes * 8 // dtype.width + ) + + self.compile_time_tensor = cute.runtime.make_fake_compact_tensor( + dtype, + shape, + stride_order=stride_order, + assumed_align=alignment_bytes, + ) + + elif is_jax_array(tensor) and GlobalOptions().use_tvm_ffi: + _validate_jax_array(tensor) + self._shape = tuple(int(s) for s in tensor.shape) + self._stride = _compute_row_major_contiguous_stride(self._shape) + self._runtime_tensor = tensor + self._data_ptr = tensor.unsafe_buffer_pointer() + + dtype = dtype_utils.cutlass_type_from_jax_type(tensor.dtype) + stride_order = get_stride_rank(self._stride) + major_mode_idx = stride_order.index(0) + shape = [cute.SymInt() for _ in range(len(self._shape))] + shape[major_mode_idx] = cute.SymInt( + divisibility=alignment_bytes * 8 // dtype.width + ) + + self.compile_time_tensor = cute.runtime.make_fake_compact_tensor( + dtype, + shape, + stride_order=stride_order, + assumed_align=alignment_bytes, + ) + + elif GlobalOptions().use_tvm_ffi: + raise ValueError( + "TVM-FFI is currently only supported for torch tensors and JAX arrays." + ) + else: + # TVM-FFI is disabled and the tensor passed in is not a cute.Tensor or torch fake tensor, + # We must convert it to a cute.Tensor + if is_torch_tensor(tensor): + dtype = dtype_utils.to_cutlass_type(tensor.dtype) + shape = tuple(tensor.shape) + stride = tensor.stride() + elif is_jax_array(tensor): + _validate_jax_array(tensor) + dtype = dtype_utils.to_cutlass_type(tensor.dtype) + shape = tuple(int(s) for s in tensor.shape) + stride = _compute_row_major_contiguous_stride(shape) + else: + raise ValueError(f"Unsupported tensor type: {type(tensor)}") + + major_mode_idx = major_mode(shape, stride) + if major_mode_idx is None: + raise ValueError(f"Tensor has no major mode: {tensor}") + stride_order = get_stride_order(stride) + self._runtime_tensor = ( + from_dlpack( + tensor, + assumed_align=alignment_bytes, + ) + .mark_layout_dynamic(major_mode_idx) + .mark_compact_shape_dynamic( + mode=major_mode_idx, + divisibility=alignment_bytes * 8 // dtype.width, + stride_order=stride_order, + ) + ) + + self._shape = self._runtime_tensor.shape + self._stride = self._runtime_tensor.stride + self._data_ptr = self._runtime_tensor.iterator._pointer + + # Since the runtime tensor is now a cute.Tensor, we can use it at + # compile time as well + self.compile_time_tensor = self._runtime_tensor + + @property + def runtime_tensor(self): + if self._runtime_tensor is None: + raise ValueError( + "Attempting to access runtime tensor from argument constructed with a fake tensor." + ) + return self._runtime_tensor + + @property + def dtype(self) -> type[Numeric]: + """Return the Cute element type of the tensor.""" + return self.compile_time_tensor.element_type + + @property + def shape(self) -> tuple[int, ...]: + return self._shape + + @property + def stride(self) -> tuple[int, ...]: + return self._stride + + @property + def data_ptr(self) -> int: + return self._data_ptr + + def to( + self, + dtype: type[Numeric] | None = None, + device: Literal["cpu", "cuda"] | None = None, + ) -> TensorLike | None: + """Copy the tensor to the given dtype and/or device and return a new framework tensor. + + When ``_runtime_tensor`` is a framework tensor (torch / JAX) the copy + is delegated to the framework. When it is a ``cute.Tensor`` (non-TVM-FFI + path), data is copied from GPU memory via ``cuMemcpyDtoH`` / + ``cuMemcpyDtoD`` into a freshly allocated torch tensor. + + :param dtype: The dtype to convert the tensor to. + :param device: ``"cpu"`` or ``"cuda"``. + """ + if self._runtime_tensor is None: + return None + + if is_torch_tensor(self._runtime_tensor): + return self._runtime_tensor.to(dtype=dtype, device=device) + elif is_jax_array(self._runtime_tensor): + import jax + + new_tensor = self._runtime_tensor + if dtype is not None: + new_tensor = new_tensor.astype(dtype) + if device is not None: + device_str = "gpu" if device == "cuda" else "cpu" + jax_device = jax.devices(device_str)[0] + new_tensor = jax.device_put(new_tensor, jax_device) + + return new_tensor + elif isinstance(self._runtime_tensor, cute.Tensor): + device_str = "cuda" if device is None else device + new_tensor = self._copy_cute_tensor_to_torch(device=device_str) + if dtype is not None: + new_tensor = new_tensor.to(dtype=dtype) + return new_tensor + + return None + + def cuda(self) -> TensorLike | None: + return self.to(device="cuda") + + def cpu(self) -> TensorLike | None: + return self.to(device="cpu") + + def _copy_cute_tensor_to_torch( + self, device: Literal["cpu", "cuda"] + ) -> TensorLike | None: + """Copy ``_runtime_tensor`` (a cute.Tensor on GPU) into a new torch tensor. + + Uses ``cuMemcpyDtoH`` for device-to-host and ``cuMemcpyDtoD`` for + device-to-device copies so that the returned tensor owns its own + memory and is independent of the runtime tensor. + """ + import torch + + import cutlass.torch as cutlass_torch + + rt = self._runtime_tensor + num_bits = self.numel() * rt.element_type.width + assert num_bits % 8 == 0, ( + "Assumes that tensors are always byte-aligned even for subbyte types" + ) + num_bytes = num_bits // 8 + torch_dtype = cutlass_torch.dtype(rt.element_type) + + if device == "cpu": + dest = torch.empty_strided(rt.shape, rt.stride, dtype=torch_dtype) + (err,) = cuda_driver.cuMemcpyDtoH(dest.data_ptr(), rt.data_ptr, num_bytes) + elif device == "cuda": + dest = torch.empty_strided( + rt.shape, rt.stride, dtype=torch_dtype, device="cuda" + ) + (err,) = cuda_driver.cuMemcpyDtoD(dest.data_ptr(), rt.data_ptr, num_bytes) + + if err != cuda_driver.CUresult.CUDA_SUCCESS: + err_name = cuda_driver.cuGetErrorString(err)[1] + raise RuntimeError(f"CUDA memcpy to {device} failed: {err_name}") + return dest + + def numel(self) -> int: + num = self._shape[0] + for i in range(1, len(self._shape)): + num *= self._shape[i] + return num + + def __dlpack__(self, *, stream: Any = None) -> Any: + """Return a DLPack capsule representing the tensor data.""" + if self._runtime_tensor is None: + raise ValueError("Attempting to call dlpack with a fake tensor.") + if not hasattr(self._runtime_tensor, "__dlpack__"): + raise NotImplementedError( + f"{type(self._runtime_tensor).__name__} does not support DLPack export." + ) + return self._runtime_tensor.__dlpack__(stream=stream) + + def __dlpack_device__(self) -> tuple[int, int]: + """Return the device type and device id as a tuple.""" + if self._runtime_tensor is None: + raise ValueError("Attempting to call dlpack_device with a fake tensor.") + if not hasattr(self._runtime_tensor, "__dlpack_device__"): + raise NotImplementedError( + f"{type(self._runtime_tensor).__name__} does not support DLPack export." + ) + return self._runtime_tensor.__dlpack_device__() + + +def strides_to_layout_string(*strides: list[tuple[int, ...]]) -> str: + """Convert list[stride tuples] to layout string ('t' for row-major, 'n' for col-major). + + Example: + >>> strides_to_layout_string((0, 1, 0), (0, 0, 1), (0, 0, 1)) + 'ntt' + """ + row_major_stride_rank3 = (0, 0, 1) + col_major_stride_rank3 = (0, 1, 0) + row_major_stride_rank2 = (0, 1) + col_major_stride_rank2 = (1, 0) + + layout_string_map = { + row_major_stride_rank3: "t", + col_major_stride_rank3: "n", + row_major_stride_rank2: "t", + col_major_stride_rank2: "n", + } + + return "".join(layout_string_map[s] for s in strides) diff --git a/operators/cutlass/operators/workspace.py b/operators/cutlass/operators/workspace.py new file mode 100644 index 000000000..282309419 --- /dev/null +++ b/operators/cutlass/operators/workspace.py @@ -0,0 +1,67 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from cutlass.operators.typing import TensorLike + + +@dataclass +class Workspace: + """Container for a user-provided memory allocation used by an Operator. + + A workspace is a user-provided, contiguous block of memory some Operators may use to store intermediate results, + which may be used by the underlying kernel for any temporary storage, book keeping, intermediate results, semaphores etc. + The workspace is typically allocated by the user and passed to the Operator via the :meth:`Operator.run` method. + """ + + data: TensorLike | int + """Tensor holding workspace data, or an integer address of the start of the workspace.""" + + leading_zero_bytes: int = 0 + """Number of bytes from the beginning of the workspace have been set to zero by user.""" + + +@dataclass +class AllocationRequirement: + """Describes requirements for a user-provided memory allocation.""" + + size_bytes: int + """The minimal size (in bytes) that must be allocated assuming perfect alignment.""" + + ptr_alignment: int + """Alignment requirement for the pointer to the allocation in bytes.""" + + @classmethod + def empty(cls) -> AllocationRequirement: + """Returns a requirement for an empty allocation.""" + return cls(size_bytes=0, ptr_alignment=1) diff --git a/operators/examples/000_gemm.ipynb b/operators/examples/000_gemm.ipynb new file mode 100644 index 000000000..d5ae23323 --- /dev/null +++ b/operators/examples/000_gemm.ipynb @@ -0,0 +1,568 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "3dd45ef2", + "metadata": {}, + "source": [ + "# Basic GEMM using CUTLASS Operator API" + ] + }, + { + "cell_type": "markdown", + "id": "4709aa60", + "metadata": {}, + "source": [ + "**CUTLASS Operator API** serves to integrate kernels written using CUTLASS Python DSLs into your code. It provides a kernel-agnostic interface to discover, compile, and run CUTLASS Python kernels, and a registry of kernels exposed via that interface.\n", + "\n", + "This notebook walks through a minimal GEMM (Generalized Matrix-Matrix Multiplication) example, and introduces the core concepts of the API." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "f878d960-d175-4d84-b978-88afbd318850", + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "\n", + "import cutlass\n", + "\n", + "import cutlass.operators as ops\n", + "\n", + "if not (status := ops.utils.device.device_or_env_supports(\"80\")):\n", + " print(\n", + " f\"This notebook requires a GPU with compute capability >= 80.\\n{status.error}\"\n", + " )\n", + " import sys\n", + "\n", + " sys.exit(0)" + ] + }, + { + "cell_type": "markdown", + "id": "db91dab6", + "metadata": {}, + "source": [ + "## Running your first Operator" + ] + }, + { + "cell_type": "markdown", + "id": "90292548", + "metadata": {}, + "source": [ + "While DSLs focus on kernel authoring, CUTLASS Operator API focuses on ease of\n", + "managing and integrating those kernels into libraries that use CUTLASS.\n", + "\n", + "It views kernels as end-to-end **operators** that execute an **operation** (like GEMM, MoE GEMM).\n", + "We will specify a basic GEMM operation, find `Operator`s that support it, and comple and run one of them." + ] + }, + { + "cell_type": "markdown", + "id": "7b7b87b0", + "metadata": {}, + "source": [ + "### Setting up the operation arguments\n", + "\n", + "CUTLASS Operator API has first-class support for PyTorch and other DLPack-compatible tensors. We start by creating torch tensors that will be operands to a matrix multiplication." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "f550c4ea", + "metadata": {}, + "outputs": [], + "source": [ + "M, N, K, L = 128, 256, 64, 2\n", + "ab_type = torch.float16\n", + "out_type = torch.float16\n", + "acc_type = torch.float32\n", + "\n", + "A = torch.randint(-1, 2, (L, M, K), device=\"cuda\", dtype=ab_type)\n", + "B = torch.randint(-1, 2, (L, K, N), device=\"cuda\", dtype=ab_type)\n", + "out = torch.empty((L, M, N), device=\"cuda\", dtype=out_type)\n", + "\n", + "reference = (A @ B).to(out.dtype)" + ] + }, + { + "cell_type": "markdown", + "id": "6b6cb805", + "metadata": {}, + "source": [ + "We then create a `GemmArguments` object. This object specifies:\n", + "1. what logical **operation** do we want to perform (a GEMM)\n", + "2. on which **operands** we want to perform that operation (`A`, `B`, `out` as declared above)" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "b57690df", + "metadata": {}, + "outputs": [], + "source": [ + "args = ops.GemmArguments(A=A, B=B, out=out, accumulator_type=acc_type)" + ] + }, + { + "cell_type": "markdown", + "id": "67e5ddcf", + "metadata": {}, + "source": [ + "### Operator discovery\n", + "\n", + "We now need to find **operators** that can perform the operation we expressed in `args`.\n", + "\n", + "The simplest way to do so is to use `get_operators(args)`. The library has a pre-registered set of many operators. A subset of these support GEMM, and a further subset support the tensor dtypes, layouts, target arch, etc. we want. With `get_operators`, we query all `Operator`s that can functionally support (accept and give correct result for) our args.\n", + "\n", + "Any of the returned operators will be functionally equivalent, but they may have different design or performance characteristics. We arbitrarily pick the first of the returned operators to execute here" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "9872ad66", + "metadata": {}, + "outputs": [], + "source": [ + "operators = ops.get_operators(args, target_sm=\"80\")\n", + "assert operators, \"No operators found for the given arguments!\"\n", + "operator = operators[0]" + ] + }, + { + "cell_type": "markdown", + "id": "4c17693e", + "metadata": {}, + "source": [ + "#### Run the operator\n", + "\n", + "Running the operator is as simple as `operator.run(args)`.\n", + "\n", + "This implicitly JIT-compiles the operator, and launches it on the GPU device using our given arguments." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "baf4588a", + "metadata": {}, + "outputs": [], + "source": [ + "operator.run(args)\n", + "\n", + "torch.testing.assert_close(out, reference)" + ] + }, + { + "cell_type": "markdown", + "id": "630e9e4b", + "metadata": {}, + "source": [ + "## Understanding the core interfaces\n" + ] + }, + { + "cell_type": "markdown", + "id": "2d8b8e94", + "metadata": {}, + "source": [ + "### 1. `RuntimeArguments` / `GemmArguments`\n", + "\n", + "`RuntimeArguments` describes the operation that we want to perform, and all the runtime information we wish to pass to an operator. This includes primary runtime operands to the operation, but also any custom epilogue fusions and runtime performance options.\n", + "\n", + "We provide builtin subtypes of `RuntimeArguments` for common operations (e.g. GEMM, Grouped GEMM).\n", + "\n", + "For instance, `GemmArguments` is a type of `RuntimeArguments`:\n", + "\n", + "```python\n", + "@dataclass\n", + "class GemmArguments(RuntimeArguments):\n", + " A: Operand\n", + " B: Operand\n", + " out: Operand\n", + " accumulator_type: NumericLike\n", + "```\n", + "\n", + "`GemmArguments` conveys:\n", + "* We want to perform a GEMM operation (`out = A @ B`)\n", + "* We want to perform it for operands in `A, B, out`, with intermediate results stored as `accumulator_type`. In our case, the operands are individual tensors but later tutorials explore other operands.\n", + "* We can optionally set a custom epilogue that is fused on top of the base GEMM. Some operators also support some runtime performance controls which can be specified here. These will be discussed in detail in other tutorials.\n", + "\n", + "It is an operator-agnostic way to specify the desired functionality -- we are declaring that this is the operation we would like to perform, irrespective of which operator will be chosen to perform it.\n", + "\n", + "`RuntimeArguments` can be constructed from any tensor-like object. This includes `torch.Tensor`, `cute.Tensor`, or any other DLPack-compatible tensors." + ] + }, + { + "cell_type": "markdown", + "id": "e7eda0dd", + "metadata": {}, + "source": [ + "### 2. Operator Discovery\n", + "\n", + "There are several operators available in CUTLASS DSLs that are registered with, and discoverable via, the CUTLASS Operator API.\n", + "\n", + "This includes operators for various operations. For the same operation, there are different implementations that use different algorithms, optimizations, or architecture features. Even within the same implementation, there are several parametrized instances or configurations of it with different combinations of operand types, layouts, tile sizes, etc.\n", + "\n", + "In the previous step, we used `GemmArguments` to specify our desired GEMM in an operator-agnostic way. Now we find specific operators that can _fulfill_ that functionality." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "3b737131", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "A total of 184434 Operator instances are available.\n", + "Of these, 572 support the given arguments.\n", + "Of these, 548 support the given GPU architecture.\n", + "Picked operator with name: cutedsl.PersistentDenseGemmOperator_sm100_ttt_AFloat16_BFloat16_outFloat16_accFloat32_2cta_cluster2x1x1_tile128x32x64_tma_store\n" + ] + } + ], + "source": [ + "# get_operators() fetches all operators when called without args\n", + "all_operators = ops.get_operators()\n", + "print(f\"A total of {len(all_operators)} Operator instances are available.\")\n", + "\n", + "# we can limit the search to operators supporting given args\n", + "operators = ops.get_operators(args)\n", + "print(f\"Of these, {len(operators)} support the given arguments.\")\n", + "\n", + "# we can further limit the search to a specific GPU arch (here, we use current device)\n", + "target_sm = ops.utils.device.device_or_env_target_sm()\n", + "operators = ops.get_operators(args, target_sm=target_sm)\n", + "print(f\"Of these, {len(operators)} support the given GPU architecture.\")\n", + "\n", + "operator = operators[0]\n", + "print(f\"Picked operator with name: {operator.metadata.operator_name}\")" + ] + }, + { + "cell_type": "markdown", + "id": "252a4d38", + "metadata": {}, + "source": [ + "### 3. `Operator` execution" + ] + }, + { + "cell_type": "markdown", + "id": "574d004b", + "metadata": {}, + "source": [ + "Once we have selected an operator, we are now ready to execute it. We previously showed the simplest way to do this is `operator.run(args)`.\n", + "\n", + "This method does the following:\n", + "* verify that the Operator supports the given `args`\n", + "* JIT-compile the Operator\n", + "* launch the compiled Operator function\n", + "\n", + "Users can do these steps individually for more control:" + ] + }, + { + "cell_type": "markdown", + "id": "e8945aa6", + "metadata": {}, + "source": [ + "#### Verify that the Operator supports the given `args`\n", + "`operator.supports(args)` checks if the Operator supports the given `args`\n", + " * this is relevant if the operator was not picked just for these `args`" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "159fd610", + "metadata": {}, + "outputs": [], + "source": [ + "supported = operator.supports(args)\n", + "assert supported" + ] + }, + { + "cell_type": "markdown", + "id": "948689a8", + "metadata": {}, + "source": [ + "If the arguments are not supported by this Operator, `supports` returns a `Status` object explaining the error." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "2cfc9ea7", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Operand `A` is unsupported: Expected element type Float16, got BFloat16\n" + ] + } + ], + "source": [ + "unsupported_args = ops.GemmArguments(\n", + " A=A.to(torch.bfloat16), B=B, out=out, accumulator_type=acc_type\n", + ")\n", + "if not (status := operator.supports(unsupported_args)):\n", + " print(status.error)\n", + "\n", + "assert not status" + ] + }, + { + "cell_type": "markdown", + "id": "c2db8f20", + "metadata": {}, + "source": [ + "#### JIT compiling the operator\n", + "\n", + "`operator.compile(args)` compiles the operator, and returns a `CompiledArtifact`\n", + "\n", + "This compiled artifact is a lightweight wrapper over the result of compiling an Operator (e.g., via `cute.compile()`).\n", + "\n", + "For just-in-time compilation, we can use the compiled artifact straightaway.\n", + "In the future, we will support optionally serializing it for ahead-of-time compilation and deserialized in a different context.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "02e79eb8", + "metadata": {}, + "outputs": [], + "source": [ + "compiled_artifact = operator.compile(args)" + ] + }, + { + "cell_type": "markdown", + "id": "4dfb8d51", + "metadata": {}, + "source": [ + "#### Launching the compiled Operator function\n", + "`operator.run(args)` launches the compiled Operator function. The next example uses:\n", + " * the precompiled artifact\n", + " * a custom stream to launch to\n", + " * bypasses the supports check already performed above (`assume_supported_args=True`)." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "02398bf0", + "metadata": {}, + "outputs": [], + "source": [ + "# zero the output to avoid testing stale output\n", + "out.zero_()\n", + "\n", + "operator.run(\n", + " args,\n", + " compiled_artifact,\n", + " stream=torch.cuda.Stream(),\n", + " assume_supported_args=True,\n", + ")\n", + "torch.testing.assert_close(out, reference)" + ] + }, + { + "cell_type": "markdown", + "id": "7956813d", + "metadata": {}, + "source": [ + "Passing in a precompiled operator is critical to achieving good performance because it avoids\n", + "JIT compiling the operator on each invocation. JIT compilation always occurs when a precompiled\n", + "operator is not provided in the call to `operator.run()`." + ] + }, + { + "cell_type": "markdown", + "id": "c228495a", + "metadata": {}, + "source": [ + "#### Workspace Buffers" + ] + }, + { + "cell_type": "markdown", + "id": "f67eeb8f", + "metadata": {}, + "source": [ + "Some operators may also require a device \"workspace\". This is an additional buffer needed by some operators for book-keeping, temporary results, etc.\n", + "Its size can be queried using `operator.get_workspace_size(args)`. Most operators will have a workspace size of 0.\n", + "If an operator does have a non-zero workspace size, an additional buffer of at least that size must be provided. Without it, the Operator behavior is undefined." + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "5b2e2d56", + "metadata": {}, + "outputs": [], + "source": [ + "workspace_size = operator.get_workspace_size(args)\n", + "workspace = torch.empty(workspace_size.size_bytes, device=\"cuda\", dtype=torch.int8)\n", + "\n", + "out.zero_()\n", + "operator.run(args, compiled_artifact, stream=torch.cuda.Stream(), workspace=workspace)\n", + "torch.testing.assert_close(out, reference)" + ] + }, + { + "cell_type": "markdown", + "id": "baffaf12", + "metadata": {}, + "source": [ + "## Advanced: Filtering on Metadata" + ] + }, + { + "cell_type": "markdown", + "id": "86dd521a", + "metadata": {}, + "source": [ + "Using `RuntimeArguments` to search for supporting operators is a convenient way to discover operators: users directly specify their desired functionality, and `get_operators()` finds the supporting operators.\n", + "It covers all logical operands of any operation, as well as (in later examples) epilogue fusions, and performance controls.\n", + "\n", + "However, there may be cases where users want more advanced ways to query operators. These could be:\n", + "* when the desired properties may not be expressed in runtime controls\n", + " * the simplest scenario may be if you're searching searching for an Operator with a specific name, a specific class, etc.\n", + " * searching for operator's static properties such as tile size, cluster size, etc.\n", + "* when the `RuntimeArguments` are not available or you want to generate & pre-compile a broader set of operators\n", + "\n", + "For such cases, we provide a more advanced filtering based on `OperatorMetadata`" + ] + }, + { + "cell_type": "markdown", + "id": "ef54ae50", + "metadata": {}, + "source": [ + "`OperatorMetadata` captures a wide variety of properties of a `Operator`.\n", + "\n", + "These are properties of an Operator's functional support (like operand types, layouts, alignments), as well as architectural/design choices & performance characteristics (like tilze size, scheduling characteristics).\n", + "\n", + "Different operators may use different sub-classes of `metadata.operands`, `metadata.design`, `metadata.epilogue` for flexibility, which can also identify their characteristics.\n", + "\n", + "```python\n", + "@dataclass\n", + "class OperatorMetadata:\n", + " operator_name: str\n", + " operator_class: type[Operator]\n", + " supported_targets: list[TargetSm]\n", + " operands: OperandsMetadata\n", + " design: DesignMetadata | None = None\n", + " epilogue: EpilogueMetadata | None = None\n", + "```" + ] + }, + { + "cell_type": "markdown", + "id": "f9943888", + "metadata": {}, + "source": [ + "Every unique Operator instance can be distinguished by its metadata.\n", + "It can be used in filtering for operators in addition to the `RuntimeArguments`, by providing a custom `metadata_filter`.\n", + "\n", + "Here, we get all operators that support `args`, and have `metadata.design` of type `Sm100DesignMetadata`.\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "8717ac89", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Found 544 operators which support args & have Sm100DesignMetadata\n" + ] + } + ], + "source": [ + "operators = ops.get_operators(\n", + " args,\n", + " metadata_filter=lambda metadata: isinstance(\n", + " metadata.design, ops.metadata.Sm100DesignMetadata\n", + " ),\n", + ")\n", + "print(f\"Found {len(operators)} operators which support args & have Sm100DesignMetadata\")" + ] + }, + { + "cell_type": "markdown", + "id": "1d1f9124", + "metadata": {}, + "source": [ + "We can construct more advanced filters by leveraging duck-typing.\n", + "Additionally, we can get all the operators that match our filter, rather than supporting a fully-defined set of arguments.\n", + "This could be useful to pre-generate large set of operators not targeted to any one problem." + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "a76ec20f", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Found 16736 matching operators\n" + ] + } + ], + "source": [ + "def a_more_complex_filter(metadata: ops.OperatorMetadata) -> bool:\n", + " \"\"\"Find all GEMM operators that support Float16 A and 2-CTA MMA.\"\"\"\n", + " # Only look at GEMM operators\n", + " if not isinstance(metadata.operands, ops.metadata.GemmOperandsMetadata):\n", + " return False\n", + " # Only look at operators with A-type F16\n", + " if metadata.operands.A.dtype != cutlass.Float16:\n", + " return False\n", + " # Only look at operators with tile_shape[0] == 128\n", + " if getattr(metadata.design, \"tile_shape\", [None])[0] != 128:\n", + " return False\n", + " return True\n", + "\n", + "\n", + "# Look ma, no args! Fetch all operators that match the filter,\n", + "# instead of supporting a complete set of args\n", + "operators = ops.get_operators(\n", + " args=None,\n", + " metadata_filter=a_more_complex_filter,\n", + ")\n", + "print(f\"Found {len(operators)} matching operators\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/operators/examples/001_gemm_with_fused_epilogue.ipynb b/operators/examples/001_gemm_with_fused_epilogue.ipynb new file mode 100644 index 000000000..31c1d6aa6 --- /dev/null +++ b/operators/examples/001_gemm_with_fused_epilogue.ipynb @@ -0,0 +1,679 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "a31330d3", + "metadata": {}, + "source": [ + "# Custom epilogue fusions for GEMMs\n" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "bb450878", + "metadata": {}, + "outputs": [], + "source": [ + "import cutlass.operators as ops\n", + "\n", + "if not (status := ops.utils.device.device_or_env_supports(\"100f\")):\n", + " print(f\"This notebook requires a Blackwell GPU (sm_100f family).\\n{status.error}\")\n", + " import sys\n", + "\n", + " sys.exit(0)" + ] + }, + { + "cell_type": "markdown", + "id": "154e9d59", + "metadata": {}, + "source": [ + "CUTLASS Operator API provides flexible epilogue fusion support -- users can specify high-level tensor transformations that can be fused into an existing Operator. This is useful for fusing activations, storing auxiliary tensors for backprop, elementwise ops, broadcasts, etc.\n", + "\n", + "For those familiar with the legacy CUTLASS Python API's [epilogue visitor tree frontend](https://github.com/NVIDIA/cutlass/blob/a2439551c765c5393aebe557ee75d3a0412d2211/examples/python/deprecated/04_epilogue_visitor.ipynb), much of the interface is shared.\n", + "\n", + "CUTLASS Operator API enables one to express an epilogue using a function operating at the `torch.Tensor`-level, and has tooling to automatically add this to operators supporting the provided function. \n", + "\n", + "For example, in PyTorch one might write the following to compute a GEMM + epilogue:" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "e6d77d53", + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "\n", + "torch.manual_seed(2025)\n", + "\n", + "L, M, N, K = 1, 1024, 1024, 1024\n", + "A = torch.randint(-2, 3, (L, M, K), device=\"cuda\", dtype=torch.float16)\n", + "B = torch.randint(-2, 3, (L, K, N), device=\"cuda\", dtype=torch.float16)\n", + "C = torch.randint(-2, 3, (L, M, N), device=\"cuda\", dtype=torch.float16)\n", + "\n", + "\n", + "def my_epilogue(accum, C, alpha, beta, extra_scalar):\n", + " Aux = (alpha * accum) + (beta * C)\n", + " D = extra_scalar * Aux\n", + " return D, Aux\n", + "\n", + "\n", + "alpha, beta, extra_scalar = 1.0, 2.0, 0.5\n", + "D, Aux = my_epilogue(A @ B, C, alpha, beta, extra_scalar)" + ] + }, + { + "cell_type": "markdown", + "id": "66ee4dd1", + "metadata": {}, + "source": [ + "Now suppose you want to compose this `my_epilogue` computation with a CUTLASS Python kernel that performs a regular GEMM.\n", + "CUTLASS Operator API allows doing so, where supporting GEMM Operators can fuse the computation into the epilogue before storing the final output.\n", + "\n", + "To do so, one defines `EpilogueArguments` consisting of the epilogue function to compute (or a string representation of it) along with arguments corresponding to each input and output of the function (except for `accum`):" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "f079d9d6", + "metadata": {}, + "outputs": [], + "source": [ + "# Allocate buffers for D and Aux\n", + "D_, Aux_ = [\n", + " torch.empty((L, M, N), device=\"cuda\", dtype=torch.float16) for _ in range(2)\n", + "]\n", + "\n", + "epi_args = ops.EpilogueArguments(\n", + " my_epilogue,\n", + " C=C,\n", + " alpha=alpha,\n", + " beta=beta,\n", + " extra_scalar=extra_scalar,\n", + " D=D_,\n", + " Aux=Aux_,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "97ef8e8a", + "metadata": {}, + "source": [ + "These arguments can be added to `GemmArguments` and passed in to `get_operators()` for use when retrieving compatible operators:" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "60215c4e", + "metadata": {}, + "outputs": [], + "source": [ + "args = ops.GemmArguments(\n", + " A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args\n", + ")\n", + "target_sm = ops.utils.device.device_or_env_target_sm()\n", + "operators = ops.get_operators(args, target_sm=target_sm)\n", + "assert len(operators) > 0" + ] + }, + { + "cell_type": "markdown", + "id": "b0a7f9a2", + "metadata": {}, + "source": [ + "Each of the operators returned by `get_operators` can be compiled and executed just the same with these new arguments, as it was in examples without\n", + "epilogue fusion. For example, using the first operator:" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "150f3296", + "metadata": {}, + "outputs": [], + "source": [ + "operators[0].run(args)\n", + "\n", + "torch.testing.assert_close(D, D_)\n", + "torch.testing.assert_close(Aux, Aux_)" + ] + }, + { + "cell_type": "markdown", + "id": "f1a826e3", + "metadata": {}, + "source": [ + "## How the epilogue fusion API works\n", + "To support specifying an epilogue via a Python function, an Operator needs some mechanism to:\n", + "1. Detect the operations in the epilogue function\n", + "2. Determine if the Operator can support the operations\n", + "3. Emit code to perform these operations within the kernel\n", + "\n", + "Step 1 listed above does not depend on the Operator and its implementation (e.g., DSL), while steps 2 and 3 depend on the Operator and/or its implementation.\n", + "\n", + "CUTLASS Operator API separates these components so that step 1 takes place at the API level and steps 2 and 3 take place in the Operator. This process is visualized below. We will walk through each step in greater detail.\n", + "\n", + "```python\n", + " +------------------------------------+\n", + " | def epi(accum, alpha, beta, C): |\n", + " | D = (accum * alpha) + (beta * C) | 1. Define epilogue via a Python function\n", + " | return D |\n", + " +------------------------------------+\n", + " |\n", + " |\n", + " |\n", + " GemmArguments(..., 2. Pass epilogue function, operands, and outputs\n", + " epilogue=EpilogueArguments( to EpilogueArguments constructor,\n", + " epi, alpha=alpha, beta=beta, C=C)) and add this to the GemmArguments. Under the\n", + " | hood, this parses the Python AST of the\n", + " | epilogue function to produce a DAG of load,\n", + " | store, and compute nodes.\n", + " V\n", + " +-----------------------------------------+ \n", + " | Intermediate DAG representation |\n", + " | =============================== |\n", + " | |\n", + " | Store() |\n", + " | | |\n", + " | Add() |\n", + " | / \\ |\n", + " | / \\ |\n", + " | / \\ |\n", + " | Mul() Mul() |\n", + " | / \\ / \\ |\n", + " | AccFetch() | Load(C) \\ |\n", + " | | \\ |\n", + " | Load(alpha) Load(beta) |\n", + " | |\n", + " +-----------------------------------------+\n", + " / | \\\n", + " / | \\ 3. Individual Operator classes use the DAG representation\n", + " / | \\ to determine if the Operator class supports the DAG.\n", + " Operator 0 Operator 1 Operator 2 If so, the Operator class emits DSL-level operations\n", + " epilogue epilogue epilogue needed to compute the epilogue DAG alongside the\n", + " emitter emitter emitter basic operation of the kernel (e.g., GEMM).\n", + " | | |\n", + " | | |\n", + " V V V\n", + "```\n", + "\n", + "### Defining an epilogue via a Python function\n", + "Epilogue fusion patterns are defined by users in Python functions that perform Tensor-level transformations -- using a `torch.Tensor` (for example) -- on the resul of a matrix multiplication.\n", + "\n", + "The structure of these functions is as follows:\n", + "```python\n", + "def custom_epi_name(accum, *args) -> Union[TensorType, tuple[TensorType]]:\n", + " \"\"\"\n", + " :param accum: result of matrix multiplication, convolution, etc. before the epilogue\n", + " :type accum: TensorType\n", + " :param args: additional arguments to be used in the epilogue (e.g., aux tensors)\n", + " :type args: list[Union[TensorType, ScalarType]]\n", + "\n", + " :returns: at least one tensor resulting from the operation of the epilogue\n", + " :rtype: Union[TensorType, tuple[TensorType]]\n", + " \"\"\"\n", + " # Do some compute\n", + " return D # and potentially other values\n", + "```\n", + "\n", + "The user defines a custom epilogue via a Python function that **must** do at least the following:\n", + "1. Take in a first positional argument named `accum` that represents the result of operation just before the epilogue is to be performed. For example, in a GEMM, `accum = A @ B`.\n", + "2. Return at least one tensor that results from computing the epilogue. Currently, the return list must contain at least one output named `D`, though this constraint may be loosened in the future.\n", + "\n", + "Each additional argument following `accum` in the function definition is expected to be either a Tensor or scalar to be loaded. Each variable in the return statement represents a Tensor or scalar to be stored. The underlying implementation of the epilogue in the kernel will determine how operands are loaded and stored.\n", + "\n", + "Compute operations are represented in static single assignment (SSA) form.\n", + "This means that each variable can be assigned exactly once.\n", + "Operations currently supported are:\n", + "* Tensor-tensor elementwise addition, subtraction, multiplication, and division\n", + "* Scalar broadcasts via addition, subtraction, multiplication, and division\n", + "* Row broadcasts: load a bias of shape `(N,)` or `(1, N)` and broadcast across M\n", + "* Column broadcasts: load a bias of shape `(M,)` or `(M, 1)` and broadcast across N\n", + "* Combined row and column broadcasts in the same epilogue\n", + "* Binary elementwise minimum and maximum\n", + "* Predefined elementwise activation functions (e.g., ReLU, sigmoid, tanh)\n", + "\n", + "Operations that are not yet supported include:\n", + "* Reductions (row, column, scalar)\n", + "\n", + "If attempting to use unsupported operations, no operators will be found in the call to `get_operators`.\n", + "\n", + "Violations to SSA or use of unexpected operators will be flagged with an exception when parsing the AST of the custom epilogue.\n", + "\n", + "Examples of epilogues fitting these patterns are given below. We will show full, runnable examples at the end of this notebook.\n", + "```python\n", + "def relu_aux_store(accum, alpha, C):\n", + " # Note that the function definition itself does not indicate the types and\n", + " # ranks of alpha and C. Thus, one cannot tell whether the epilogue is performing\n", + " # broadcasts or elementwise operations until actual arguments or metadata are\n", + " # provided to the epilogue. See below for details.\n", + " F = (accum * alpha) + (C * 2.0) # Constant beta of 2.0\n", + " D = relu(F)\n", + " return D, F\n", + "\n", + "def aux_normalize(accum, aux):\n", + " D = accum / aux\n", + " return D\n", + "\n", + "def bias_add(accum, row_bias, col_bias):\n", + " # row_bias has shape (N,) and col_bias has shape (M, 1).\n", + " # The broadcast type is determined automatically from the tensor shapes\n", + " # provided to EpilogueArguments.\n", + " D = accum + row_bias + col_bias\n", + " return D\n", + "```\n", + "\n", + "Additional information about each operand and output must be provided by the user when constructing `EpilogueArguments`, as we will discuss below. This additional information is necessary for fully defining the operations being performed -- without knowledge of whether `alpha` is a scalar or a Tensor, we cannot determine whether multiplication by `alpha` is a broadcasted or elementwise operation.\n", + "\n", + "### Constructing epilogue arguments\n", + "`EpilogueArguments` encapsulate the arguments needed to determine the functional operation of a fused epilogue.\n", + "\n", + "A user must provide in the construction of `EpilogueArguments` tensors for all operands and outputs of the epilogue. However, unlike arguments for basic operations (e.g., GEMM), the full set of operands needed to be specified for an epilogue pattern depends upon the custom epilogue defined by the user.\n", + "\n", + "Therefore, `EpilogueArguments` is defined generically as taking in an `epilogue_fn` and additional `kwargs`. Under the hood, the AST for `epilogue_fn` is parsed to determine the operands and outputs of the epilogue. The user is required to provide in `kwargs` Tensors or scalars for all operands and outputs in the provided epilogue.\n", + "\n", + "For example, with an epilogue of:\n", + "```python\n", + "def my_epi(accum, alpha, C, beta):\n", + " F = (accum * alpha) + (C * beta)\n", + " D = relu(F)\n", + " return D, F\n", + "```\n", + "A user would need to construct epilogue arguments as follows:\n", + "```python\n", + "epi_args = EpilogueArguments(my_epi, alpha=..., C=..., beta=..., D=..., F=...)\n", + "```\n", + "\n", + "After verifying that all required operands and outputs are present, the constructor to `EpilogueArguments` will perform additional passes on the AST of `epilogue_fn` using the provided inputs to generate an internal DAG representing the epilogue. This DAG structure is attached to `EpilogueArguments` for use as they are passed through a call to `get_operators`.\n", + "\n", + "### Discovering operators that support the epilogue pattern\n", + "\n", + "The call to `get_operators(args)` will return any operators that support the provided `GemmArguments`.\n", + "Since the `GemmArguments` constructed above now include `EpilogueArguments`, returned operators must support the provided epilogue.\n", + "\n", + "Under the hood of `get_operators()`, each `Operator` class will determine in its `generate_operators()` method whether it supports the provided `EpilogueArguments`.\n", + "It can do so by traversing the DAG that resulted from the construction of `EpilogueArguments` to find the operations that compose the epilogue.\n", + "Assuming that the `Operator` can support the DAG, it must then add to the source for the kernel any operations needed to support the DAG.\n", + "An example of how this is done generically for an SM100 CuTe DSL GEMM is provided in [`cutlass/operators/providers/cutedsl/gemm/sm100_static_persistent_efc.py`](../cutlass/operators/providers/cutedsl/gemm/sm100_static_persistent_efc.py).\n", + "\n", + "## Example epilogues\n", + "We now provide various examples of adding custom epilogues to GEMM operators targeting SM100. A broader set of epilogue examples are available in [`test/integration/test_gemm_epilogue_fusion.py`](../test/integration/test_gemm_epilogue_fusion.py).\n", + "\n", + "### Auxiliary input and output tensors" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "171ac178", + "metadata": {}, + "outputs": [], + "source": [ + "from cutlass.operators.fusion.activation import relu\n", + "\n", + "\n", + "def relu_aux_store(accum, alpha, C):\n", + " F = (accum * alpha) + (C * 2.0) # Constant beta\n", + " D = relu(F)\n", + " return D, F\n", + "\n", + "\n", + "C = torch.randint(-2, 3, (L, M, N), device=\"cuda\", dtype=torch.float16)\n", + "alpha = 3.0\n", + "D = torch.empty((L, M, N), device=\"cuda\", dtype=torch.float16)\n", + "F = torch.empty((L, M, N), device=\"cuda\", dtype=torch.float16)\n", + "\n", + "epi_args = ops.EpilogueArguments(relu_aux_store, alpha=alpha, C=C, D=D, F=F)\n", + "args = ops.GemmArguments(\n", + " A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args\n", + ")\n", + "operators = ops.get_operators(args, target_sm=target_sm)\n", + "assert len(operators) > 0\n", + "operators[0].run(args)\n", + "\n", + "D_ref, F_ref = relu_aux_store(A @ B, alpha, C)\n", + "\n", + "torch.testing.assert_close(D, D_ref)\n", + "torch.testing.assert_close(F, F_ref)" + ] + }, + { + "cell_type": "markdown", + "id": "f947b403", + "metadata": {}, + "source": [ + "### Keyword functions and returning accumulator" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "62c2b49b", + "metadata": {}, + "outputs": [], + "source": [ + "def relu_scale_return_acc(accum, alpha, beta, C, scale):\n", + " F = relu((accum * alpha) + (C * beta))\n", + " D = F * scale\n", + " return D, F, accum\n", + "\n", + "\n", + "C = torch.randint(-2, 3, (L, M, N), device=\"cuda\", dtype=torch.float16)\n", + "alpha = 1.0\n", + "beta = 2.0\n", + "scale = 0.5\n", + "D = torch.empty((L, M, N), device=\"cuda\", dtype=torch.float16)\n", + "F = torch.empty((L, M, N), device=\"cuda\", dtype=torch.float16)\n", + "accum = torch.empty((L, M, N), device=\"cuda\", dtype=torch.float32)\n", + "\n", + "epi_args = ops.EpilogueArguments(\n", + " relu_scale_return_acc,\n", + " alpha=alpha,\n", + " beta=beta,\n", + " C=C,\n", + " scale=scale,\n", + " D=D,\n", + " F=F,\n", + " accum=accum,\n", + ")\n", + "args = ops.GemmArguments(\n", + " A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args\n", + ")\n", + "operators = ops.get_operators(args, target_sm=target_sm)\n", + "assert len(operators) > 0\n", + "operators[0].run(args)\n", + "\n", + "D_ref, F_ref, accum_ref = relu_scale_return_acc(A @ B, alpha, beta, C, scale)\n", + "\n", + "torch.testing.assert_close(D, D_ref)\n", + "torch.testing.assert_close(F, F_ref)\n", + "torch.testing.assert_close(accum, accum_ref.to(accum.dtype))" + ] + }, + { + "cell_type": "markdown", + "id": "c641911f", + "metadata": {}, + "source": [ + "### Passing a string representation of the function\n", + "`EpilogueArguments` can additionally be constructed using a string representation of the epilogue function:" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "5987bf44", + "metadata": {}, + "outputs": [], + "source": [ + "epi_str = \"def epi(accum, alpha, beta, C): F = (accum * alpha) + (C * beta); D = relu(F); return D, F\"\n", + "\n", + "C = torch.randint(-2, 3, (L, M, N), device=\"cuda\", dtype=torch.float16)\n", + "alpha = 1.0\n", + "beta = 0.5\n", + "D = torch.empty((L, M, N), device=\"cuda\", dtype=torch.float16)\n", + "F = torch.empty((L, M, N), device=\"cuda\", dtype=torch.float16)\n", + "\n", + "epi_args = ops.EpilogueArguments(epi_str, alpha=alpha, beta=beta, C=C, D=D, F=F)\n", + "args = ops.GemmArguments(\n", + " A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args\n", + ")\n", + "operators = ops.get_operators(args, target_sm=target_sm)\n", + "assert len(operators) > 0\n", + "operators[0].run(args)\n", + "\n", + "F_ref = (A @ B) * alpha + (C * beta)\n", + "D_ref = torch.relu(F_ref)\n", + "\n", + "torch.testing.assert_close(D, D_ref)\n", + "torch.testing.assert_close(F, F_ref)" + ] + }, + { + "cell_type": "markdown", + "id": "a7505a64", + "metadata": {}, + "source": [ + "### Row and column broadcasts\n", + "A common epilogue pattern is adding a per-row or per-column bias to the GEMM result. The broadcast type is determined automatically from the shape of the tensor provided to `EpilogueArguments`:\n", + "- **Row broadcast**: bias has shape `(N,)` or `(1, N)` — broadcast across all M rows.\n", + "- **Column broadcast**: bias has shape `(M, 1)` — broadcast across all N columns.\n", + "\n", + "Both types can be combined in the same epilogue:" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "320e3609", + "metadata": {}, + "outputs": [], + "source": [ + "def broadcast_epilogue(accum, row_bias, col_bias, alpha):\n", + " # row_bias shape (N,): same bias added to every row.\n", + " # col_bias shape (M, 1): same bias added to every column.\n", + " # The broadcast type is inferred from the shapes provided\n", + " # to EpilogueArguments.\n", + " D = (accum * alpha) + row_bias + col_bias\n", + " return D\n", + "\n", + "\n", + "# Row bias: one value per column (N).\n", + "row_bias = torch.randn(N, device=\"cuda\", dtype=torch.float16)\n", + "# Column bias: one value per row (M), with an extra trailing dimension.\n", + "col_bias = torch.randn(M, 1, device=\"cuda\", dtype=torch.float16)\n", + "alpha = 1.0\n", + "D = torch.empty((L, M, N), device=\"cuda\", dtype=torch.float16)\n", + "\n", + "epi_args = ops.EpilogueArguments(\n", + " broadcast_epilogue,\n", + " row_bias=row_bias,\n", + " col_bias=col_bias,\n", + " alpha=alpha,\n", + " D=D,\n", + ")\n", + "args = ops.GemmArguments(\n", + " A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args\n", + ")\n", + "operators = ops.get_operators(args, target_sm=target_sm)\n", + "assert len(operators) > 0\n", + "operators[0].run(args)\n", + "\n", + "D_ref = broadcast_epilogue(A @ B, row_bias, col_bias, alpha)\n", + "\n", + "torch.testing.assert_close(D, D_ref)" + ] + }, + { + "cell_type": "markdown", + "id": "e26a58a2", + "metadata": {}, + "source": [ + "### Failure examples\n", + "The following are examples of constructing `EpilogueArguments` that are expected to fail." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "1e3d0c89", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "accum must be an input to the epilogue function\n" + ] + } + ], + "source": [ + "####################################################\n", + "# Epilogues must take in an accumulator\n", + "####################################################\n", + "def fail_missing_accum(alpha, beta, C):\n", + " D = C * beta\n", + " return D\n", + "\n", + "\n", + "try:\n", + " epi_args = ops.EpilogueArguments(\n", + " fail_missing_accum,\n", + " alpha=alpha,\n", + " beta=beta,\n", + " C=C,\n", + " D=D,\n", + " )\n", + " args = ops.GemmArguments(\n", + " A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args\n", + " )\n", + "except Exception as e:\n", + " # \"accum must be an input to the epilogue function\"\n", + " print(e)" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "48a359f7", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Output node D is not found in the epilogue function.\n" + ] + } + ], + "source": [ + "####################################################\n", + "# Epilogues must return an output named D\n", + "####################################################\n", + "def fail_missing_D(accum, alpha, beta, C):\n", + " F = (accum * alpha) + (C * beta)\n", + " return F\n", + "\n", + "\n", + "try:\n", + " epi_args = ops.EpilogueArguments(fail_missing_D, alpha=alpha, beta=beta, C=C, F=F)\n", + " args = ops.GemmArguments(\n", + " A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args\n", + " )\n", + "except Exception as e:\n", + " # \"Output node D is not found in the epilogue function\n", + " print(e)" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "49d9ee94", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Variable 'tmp' cannot be defined twice.\n" + ] + } + ], + "source": [ + "####################################################\n", + "# Epilogues must use single-static assignment (SSA)\n", + "####################################################\n", + "def fail_ssa(accum):\n", + " tmp = accum * 2.0\n", + " # Redefine tmp, which violates SSA form.\n", + " tmp = tmp - 1.0\n", + " D = tmp / 4.0\n", + " return D, tmp\n", + "\n", + "\n", + "try:\n", + " epi_args = ops.EpilogueArguments(fail_ssa, D=D, tmp=F)\n", + " args = ops.GemmArguments(\n", + " A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args\n", + " )\n", + "except Exception as e:\n", + " # \"Variable 'tmp' cannot be defined twice.\"\n", + " print(e)" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "871bb727", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Argument D is not provided in the kwargs of the EpilogueArguments constructor\n", + "Argument alpha is not provided in the kwargs of the EpilogueArguments constructor\n" + ] + } + ], + "source": [ + "####################################################\n", + "# Must provide all operands and outputs to\n", + "# EpilogueArguments\n", + "####################################################\n", + "def my_epi(accum, alpha, beta, C):\n", + " F = (accum * alpha) + (C * beta)\n", + " D = relu(F)\n", + " return D\n", + "\n", + "\n", + "try:\n", + " # Missing D\n", + " epi_args = ops.EpilogueArguments(my_epi, alpha=alpha, beta=beta, C=C)\n", + " args = ops.GemmArguments(\n", + " A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args\n", + " )\n", + "except Exception as e:\n", + " # \"Argument D is not provided in the kwargs of the EpilogueArguments constructor\"\n", + " print(e)\n", + "\n", + "try:\n", + " # Missing alpha\n", + " epi_args = ops.EpilogueArguments(my_epi, beta=beta, C=C, D=D)\n", + " args = ops.GemmArguments(\n", + " A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args\n", + " )\n", + "except Exception as e:\n", + " # \"Argument alpha is not provided in the kwargs of the EpilogueArguments constructor\"\n", + " print(e)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.5" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/operators/examples/002_bring_your_own_kernel.ipynb b/operators/examples/002_bring_your_own_kernel.ipynb new file mode 100644 index 000000000..91feee8f6 --- /dev/null +++ b/operators/examples/002_bring_your_own_kernel.ipynb @@ -0,0 +1,663 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "578f2730", + "metadata": {}, + "source": [ + "# Adding new Operators to CUTLASS Operator API\n", + "CUTLASS Operator API is designed to make it easy for users to add their own kernels\n", + "so that it can be discovered and run as Operators under the uniform API. We welcome\n", + "contributions of kernels written in CUTLASS Python by \"bringing your own operator.\"\n", + "\n", + "This example shows how to add a CuTe DSL operator to the CUTLASS Operator API.\n", + "\n", + "## Bring your own implementation\n", + "Kernel authors wishing to add a CuTe DSL operator to Operator API likely already\n", + "have the kernel written in CuTe DSL, but have not yet implemented the API's needed\n", + "`Operator` interface. Within the API, we separate these components into the \"kernel\" or \"implementation\" itself,\n", + "which is written in CuTe DSL, and the `Operator` that implements API methods to expose the\n", + "underlying kernel via a uniform interface.\n", + "\n", + "For example, consider the following implementation of a simple FP64 GEMM kernel:" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "5a64b0be", + "metadata": {}, + "outputs": [], + "source": [ + "from collections.abc import Callable\n", + "\n", + "import cuda.bindings.driver as cuda\n", + "\n", + "import cutlass\n", + "import cutlass.cute as cute\n", + "\n", + "\n", + "class F64GemmKernelImplementation:\n", + " def __init__(self, cta_tile_shape_mn: tuple[int, int]):\n", + " self.cta_tile_shape_mn = cta_tile_shape_mn\n", + "\n", + " @cute.jit\n", + " def __call__(\n", + " self, a: cute.Tensor, b: cute.Tensor, out: cute.Tensor, stream: cuda.CUstream\n", + " ):\n", + " l, m, n = out.shape # noqa: E741 (L batch dim)\n", + " m_tiles = (m + self.cta_tile_shape_mn[0] - 1) // self.cta_tile_shape_mn[0]\n", + " n_tiles = (n + self.cta_tile_shape_mn[1] - 1) // self.cta_tile_shape_mn[1]\n", + "\n", + " grid = (m_tiles, n_tiles, l)\n", + " block = [self.cta_tile_shape_mn[0], self.cta_tile_shape_mn[1], 1]\n", + " self.kernel(a, b, out).launch(grid=grid, block=block, stream=stream)\n", + "\n", + " @cute.kernel\n", + " def kernel(self, a: cute.Tensor, b: cute.Tensor, out: cute.Tensor):\n", + " _L, M, N = out.shape\n", + " K = a.shape[-1]\n", + " m_tile, n_tile, l_idx = cute.arch.block_idx()\n", + " tidx, tidy, _ = cute.arch.thread_idx()\n", + "\n", + " m_idx = m_tile * self.cta_tile_shape_mn[0] + tidx\n", + " n_idx = n_tile * self.cta_tile_shape_mn[1] + tidy\n", + "\n", + " if m_idx < M and n_idx < N:\n", + " out[l_idx, m_idx, n_idx] = cutlass.Float64(0)\n", + " for k_idx in range(K):\n", + " out[l_idx, m_idx, n_idx] += (\n", + " a[l_idx, m_idx, k_idx] * b[l_idx, k_idx, n_idx]\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "36a08d4b", + "metadata": {}, + "source": [ + "The implementation is configurable via a `cta_tile_shape_mn` argument, which\n", + "controls the size of blocks and tiles in the M and N modes. A simple `cute.jit` function\n", + "computes the grid and block size for the input problem based on `cta_tile_shape_mn`,\n", + "and launches the kernel. The `cute.kernel` itself simply has each thread compute a single\n", + "output element of the matrix by taking a dot product.\n", + "\n", + "This implementation is not performant, but is kept simple for illustrative purposes." + ] + }, + { + "cell_type": "markdown", + "id": "a5d0e661", + "metadata": {}, + "source": [ + "## Defining interface methods\n", + "As it currently stands, this GEMM kernel implementation cannot be used via the\n", + "CUTLASS Operator API because it does not implement interface methods. Specifically, operators\n", + "within the CUTLASS Operator API must inherit from and implement the `cutlass.operators.Operator`\n", + "abstract class. This class has methods needed for many common operations\n", + "performed when compiling and executing DSL kernels.\n", + "\n", + "Certain providers (i.e., DSLs), such as CuTe DSL, provide an additional layer atop the\n", + "`Operator` class to add convenience utilities for Operators using that provider.\n", + "For example, the CuTe DSL provider in CUTLASS Operator API defines\n", + "`ops.providers.cutedsl.operator.CuteDslOperator`, which adds utilities surrounding\n", + "`cute.compile()` to add compile-time arguments needed for using TVM-FFI when\n", + "it is enabled.\n", + "\n", + "We will next walk through the steps in defining interface methods for this\n", + "implementation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1a2da869", + "metadata": {}, + "outputs": [], + "source": [ + "import itertools\n", + "\n", + "import cutlass.operators as ops" + ] + }, + { + "cell_type": "markdown", + "id": "86ae75cc", + "metadata": {}, + "source": [ + "We begin by defining a class to represent the Operator's interface.\n", + "As mentioned above, since this is a CuTe DSL Operator, our interface must\n", + "inherit from and implement `ops.providers.cutedsl.operator.CuteDslOperator`.\n", + "\n", + "The class must additionally be registered with the CuTe DSL provider\n", + "via the `@CuTeDSLProvider.register` decorator so that the class\n", + "can be considered when discovering Operators.\n", + "\n", + "Each `Operator` subclass must define two class-level attributes:\n", + "- `supported_args_type`: The `RuntimeArguments` subclass this Operator supports (e.g., `GemmArguments`). This allows providers to efficiently filter Operator classes before calling `generate_operators()`.\n", + "- `designed_for_min_cc`: Minimum compute capability required by this Operator class. This indicates a broad minimum that this Operator implementation as a whole is\n", + "designed for. Individual Operator instances present more finegrained information about which arch they support, under `OperatorMetadata.supported_targets`." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "3a86d138", + "metadata": {}, + "outputs": [], + "source": [ + "@ops.providers.cutedsl.CuTeDSLProvider.register\n", + "class F64GemmOperator(ops.providers.cutedsl.operator.CuteDslOperator):\n", + " # Declare which RuntimeArguments subclass this Operator supports.\n", + " supported_args_type = ops.GemmArguments\n", + "\n", + " # Minimum compute capability required by this Operator class.\n", + " designed_for_min_cc = 80\n", + "\n", + " # Empty versions of interface methods. These will be implemented later, interspersed\n", + " # with notebook markdown. Normally, one would define them inline with the class definition.\n", + " def __init__(self, metadata: ops.OperatorMetadata):\n", + " pass\n", + "\n", + " def _run(\n", + " self,\n", + " args: ops.GemmArguments,\n", + " artifact: ops.artifact.CompiledArtifact,\n", + " stream,\n", + " workspace=None,\n", + " ):\n", + " pass\n", + "\n", + " def _compile(\n", + " self, args: ops.GemmArguments, target_sm: ops.TargetSm | None = None\n", + " ) -> ops.artifact.CompiledArtifact:\n", + " pass\n", + "\n", + " @classmethod\n", + " def _generate_operators(\n", + " cls, metadata_filter, epilogue_args=None, target_sm=None, args=None\n", + " ) -> list[\"F64GemmOperator\"]:\n", + " pass\n", + "\n", + " def _supports(\n", + " self, args: ops.GemmArguments, target_sm: ops.TargetSm | None = None\n", + " ) -> ops.Status:\n", + " pass\n", + "\n", + " def _get_workspace_size(\n", + " self, args: ops.GemmArguments\n", + " ) -> ops.workspace.AllocationRequirement:\n", + " pass" + ] + }, + { + "cell_type": "markdown", + "id": "327e9e7c", + "metadata": {}, + "source": [ + "The `__init__` method of the class takes in a `OperatorMetadata` object\n", + "from which it extracts the `cta_tile_shape_mn`. This is used to construct\n", + "the Operator implementation object. We will discuss later how the `OperatorMetadata`\n", + "object passed in here is constructed:" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "785d1882", + "metadata": {}, + "outputs": [], + "source": [ + "def __init__(self, metadata: ops.OperatorMetadata):\n", + " # Using Python-2-style super() because we're defining this method outside of the class definition.\n", + " super(F64GemmOperator, self).__init__(metadata)\n", + " cta_tile_shape_mn = metadata.design.tile_shape[:2]\n", + " self.impl = F64GemmKernelImplementation(cta_tile_shape_mn)" + ] + }, + { + "cell_type": "markdown", + "id": "500a0030", + "metadata": {}, + "source": [ + "### Defining interfaces for compilation and execution\n", + "The interfaces needed for compilation and execution are simple.\n", + "\n", + "The `compile` method simply constructs a placeholder stream object\n", + "and passes that and relevant arguments to `self.cute_compile`. This\n", + "is a utility defined in the `CuteDSLOperator` abstract class that\n", + "passes in compilation flags needed for certain options to `cute.compile`\n", + "(e.g., TVM-FFI). The result is wrapped as a `CompiledArtifact`." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "63b4a129", + "metadata": {}, + "outputs": [], + "source": [ + "def _compile(\n", + " self, args: ops.GemmArguments, target_sm: ops.TargetSm | None = None\n", + ") -> ops.artifact.CompiledArtifact:\n", + " stream = cutlass.cute.runtime.make_fake_stream()\n", + " return self.cute_compile(\n", + " self.impl,\n", + " args.A.tensor,\n", + " args.B.tensor,\n", + " args.out.tensor,\n", + " stream,\n", + " target_sm=target_sm,\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "023127fd", + "metadata": {}, + "source": [ + "Users define the `_run` method rather than the top-level `run` method\n", + "(no leading underscore) that is used in interacting with operators. `_run` (1) extracts from `args`\n", + "the arguments needed to run the JIT function, and (2) calls the JIT function\n", + "passed in via `artifact` with these arguments." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "2ae7c009", + "metadata": {}, + "outputs": [], + "source": [ + "def _run(\n", + " self,\n", + " args: ops.GemmArguments,\n", + " artifact: ops.artifact.CompiledArtifact,\n", + " stream,\n", + " workspace=None,\n", + "):\n", + " stream = ops.utils.device.to_cuda_stream(stream)\n", + " compiled_gemm = artifact.compiled_obj\n", + " self.cute_run(compiled_gemm, args.A.tensor, args.B.tensor, args.out.tensor, stream)" + ] + }, + { + "cell_type": "markdown", + "id": "4052e5a0", + "metadata": {}, + "source": [ + "Finally, since this Operator does not require any device workspace,\n", + "we give it a simple `get_workspace_size` method that always returns an empty workspace. Note that this is the same as the base implementation of this method, and so it's not strictly required in this case. We show it here for demonstration of the required methods if workspace is used." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "968906ea", + "metadata": {}, + "outputs": [], + "source": [ + "def _get_workspace_size(\n", + " self, args: ops.GemmArguments\n", + ") -> ops.workspace.AllocationRequirement:\n", + " return ops.workspace.AllocationRequirement.empty()" + ] + }, + { + "cell_type": "markdown", + "id": "e245a319", + "metadata": {}, + "source": [ + "### Defining interfaces for Operator generation\n", + "We have implemented the interfaces needed for constructing the Operator\n", + "interface, compiling it, and running it. We now must implement methods for\n", + "generating the possible configurations of this Operator that the Operator\n", + "class itself supports. This will be used in Operator discovery (e.g., via\n", + "`ops.get_operators()`).\n", + "\n", + "To do so, we write the `generate_operators` method. This takes in a\n", + "binary function `metadata_filter`, epilogue arguments `epilogue_args`,\n", + "and a compute capability `cc`. It returns a list of all instances\n", + "of the Operator interface that support the `epilogue_args`, are compatible\n", + "with the given `cc`, and which pass the `metadata_filter`.\n", + "\n", + "The `Operator` class is responsible for defining what valid possible configurations (instances) of it can exist.\n", + "In this example, the valid configurations involve a cross-product of row/column-major strides and two preset tile shapes.\n", + "We create a nested loop over these options and create a `OperatorMetadata` corresponding to each unique configuration.\n", + "\n", + "The `generate_operators` method must additionally filter the generated operators by passing it through a `metadata_filter`.\n", + "This is a user-provided custom filter to filter generated metadata combinations. More information on `metadata_filter` is provided in other examples." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "47dc2f20", + "metadata": {}, + "outputs": [], + "source": [ + "@classmethod\n", + "def _generate_operators(\n", + " cls,\n", + " metadata_filter: Callable[[ops.OperatorMetadata], bool],\n", + " epilogue_args: ops.EpilogueArguments = None,\n", + " target_sm: ops.TargetSm | None = None,\n", + " args: ops.GemmArguments | None = None,\n", + ") -> list[\"F64GemmOperator\"]:\n", + " # The tile shapes this Operator supports/exposes\n", + " supported_tile_shapes = [(32, 32, 1), (16, 16, 1)]\n", + "\n", + " if epilogue_args is not None:\n", + " return []\n", + "\n", + " row_major_stride = (0, 0, 1)\n", + " col_major_stride = (0, 1, 0)\n", + " stride_combos = list(\n", + " itertools.product([row_major_stride, col_major_stride], repeat=3)\n", + " )\n", + "\n", + " # If args are provided, use the strides from the arguments\n", + " if args is not None:\n", + " if any(x.dtype != cutlass.Float64 for x in [args.A, args.B, args.out]):\n", + " return []\n", + " if any(not isinstance(x, ops.DenseTensor) for x in [args.A, args.B, args.out]):\n", + " return []\n", + "\n", + " A_stride = ops.utils.tensor.normalized_major_stride(\n", + " args.A.shape, args.A.stride, prepend_zeros_to_rank=3\n", + " )\n", + " B_stride = ops.utils.tensor.normalized_major_stride(\n", + " args.B.shape, args.B.stride, prepend_zeros_to_rank=3\n", + " )\n", + " out_stride = ops.utils.tensor.normalized_major_stride(\n", + " args.out.shape, args.out.stride, prepend_zeros_to_rank=3\n", + " )\n", + "\n", + " combo = (A_stride, B_stride, out_stride)\n", + " if combo not in stride_combos:\n", + " return []\n", + "\n", + " # Override stride_combos\n", + " stride_combos = [combo]\n", + "\n", + " divisibility = 1\n", + "\n", + " def stride_name(stride):\n", + " return \"T\" if stride == row_major_stride else \"N\"\n", + "\n", + " operators = []\n", + " for tile_shape in supported_tile_shapes:\n", + " for stride_A, stride_B, stride_out in stride_combos:\n", + " # Create OperandConstraints for A, B, and out tensors\n", + " a_attrs = ops.metadata.DenseTensorConstraints(\n", + " cutlass.Float64, stride_A, divisibility\n", + " )\n", + " b_attrs = ops.metadata.DenseTensorConstraints(\n", + " cutlass.Float64, stride_B, divisibility\n", + " )\n", + " out_attrs = ops.metadata.DenseTensorConstraints(\n", + " cutlass.Float64, stride_out, divisibility\n", + " )\n", + " layout_str = ops.utils.tensor.strides_to_layout_string(\n", + " stride_A, stride_B, stride_out\n", + " )\n", + "\n", + " design = ops.metadata.BLASDesignMetadata(\n", + " mma_instruction_type=ops.mma.AmpereMma,\n", + " tile_shape=tile_shape,\n", + " cluster_shape=(1, 1, 1),\n", + " )\n", + " operands = ops.metadata.GemmOperandsMetadata(\n", + " A=a_attrs, B=b_attrs, out=out_attrs, accumulator_type=cutlass.Float64\n", + " )\n", + " name = f\"F64GemmOperator_tile{tile_shape[0]}x{tile_shape[1]}_{layout_str}\"\n", + "\n", + " metadata = ops.OperatorMetadata(\n", + " operator_name=name,\n", + " operator_class=F64GemmOperator,\n", + " operands=operands,\n", + " design=design,\n", + " supported_targets=ops.TargetSm.get_supported_targets(design, operands),\n", + " )\n", + "\n", + " if metadata_filter(metadata):\n", + " operators.append(cls(metadata))\n", + "\n", + " return operators" + ] + }, + { + "cell_type": "markdown", + "id": "c7cdbc66", + "metadata": {}, + "source": [ + "We also add a method for indicating whether a Operator instance in question\n", + "supports a set of arguments. The top-level `Operator.supports` method will\n", + "already verify that the `args` and `target_sm` match the metadata with which\n", + "this `Operator` instance was constructed. Here, we can define additional\n", + "checks specific to this Operator. For instance, here we require that\n", + "all operands to be of rank 3:" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "54067d47", + "metadata": {}, + "outputs": [], + "source": [ + "def _supports(\n", + " self, args: ops.GemmArguments, target_sm: ops.TargetSm | None = None\n", + ") -> ops.Status:\n", + " if not (\n", + " len(args.A.shape) == 3 # A should be (L, M, K)\n", + " and len(args.B.shape) == 3 # B should be (L, K, N)\n", + " and len(args.out.shape) == 3 # out should be (L, M, N)\n", + " ):\n", + " return ops.Status.fail(\"All operands must be rank 3.\")\n", + " return ops.Status.success()" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "edaf2cba", + "metadata": {}, + "outputs": [], + "source": [ + "# Assign methods to the class because we interspersed notebook markdown\n", + "# with the class definition. This is not needed in a real implementation.\n", + "F64GemmOperator.__init__ = __init__\n", + "F64GemmOperator._compile = _compile\n", + "F64GemmOperator._run = _run\n", + "F64GemmOperator._supports = _supports\n", + "F64GemmOperator._generate_operators = _generate_operators\n", + "F64GemmOperator._get_workspace_size = _get_workspace_size" + ] + }, + { + "cell_type": "markdown", + "id": "c8fc84e9", + "metadata": {}, + "source": [ + "## Discovering instances of the Operator and using them\n", + "The CUTLASS Operator API is now prepared to discover instances of this\n", + "Operator interface just as was done in previous examples.\n", + "\n", + "We add a small modification of using a `metadata_filter`\n", + "to ensure that all returned operators are instances of the\n", + "`F64GemmOperator` class we just implemented. This is needed\n", + "only for example/testing purposes." + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "cec5431d", + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "\n", + "torch.manual_seed(2025)\n", + "\n", + "L, M, N, K = 1, 256, 1024, 128\n", + "A = torch.randint(-2, 3, (L, M, K), device=\"cuda\", dtype=torch.float64)\n", + "B = torch.randint(-2, 3, (L, K, N), device=\"cuda\", dtype=torch.float64)\n", + "out = torch.empty(L, M, N, device=\"cuda\", dtype=torch.float64)\n", + "\n", + "args = ops.GemmArguments(A, B, out, accumulator_type=torch.float64)\n", + "\n", + "\n", + "def is_f64gemm_operator(metadata):\n", + " return metadata.operator_class is F64GemmOperator\n", + "\n", + "\n", + "operators = ops.get_operators(args, metadata_filter=is_f64gemm_operator)" + ] + }, + { + "cell_type": "markdown", + "id": "50e81a7d", + "metadata": {}, + "source": [ + "We can print off the names of the first few operators to see that\n", + "they come from our recently-added Operator." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "cdb92b5e", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "F64GemmOperator_tile32x32_ttt\n", + "F64GemmOperator_tile16x16_ttt\n" + ] + } + ], + "source": [ + "print(operators[0].metadata.operator_name)\n", + "print(operators[1].metadata.operator_name)" + ] + }, + { + "cell_type": "markdown", + "id": "697ee3c3", + "metadata": {}, + "source": [ + "We can evaluate and test the correctness of an instance of our Operator:" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "f5486244", + "metadata": {}, + "outputs": [], + "source": [ + "operators[0].run(args)\n", + "torch.testing.assert_close(out, A @ B)" + ] + }, + { + "cell_type": "markdown", + "id": "8de96f7e", + "metadata": {}, + "source": [ + "We can also test the limits of our Operator's design space by providing a\n", + "metadata filter that expects a CTA tile size M of 256, which is not exposed\n", + "in the `generate_operators` method of our recently-added Operator. We expect\n", + "no operators of type `F64GemmOperator` to be returned." + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "917c74e3", + "metadata": {}, + "outputs": [], + "source": [ + "def my_filter(metadata):\n", + " return (\n", + " is_f64gemm_operator(metadata)\n", + " and isinstance(metadata.design, ops.metadata.BLASDesignMetadata)\n", + " and metadata.design.tile_shape[0] == 256\n", + " )\n", + "\n", + "\n", + "operators_ctam256 = ops.get_operators(args, metadata_filter=my_filter)\n", + "\n", + "# No operators should be found\n", + "assert len(operators_ctam256) == 0" + ] + }, + { + "cell_type": "markdown", + "id": "caa80a7d", + "metadata": {}, + "source": [ + "## A note on contributing operators to directory structure\n", + "This example showed how to define an Operator inline and add it to the\n", + "API for example purposes. This Operator doesn't necessarily need to live\n", + "within the API's source code.\n", + "\n", + "We welcome contributions of operators that do live within the CUTLASS\n", + "API's repository as well.\n", + "\n", + "Kernels in the repository are organized based on the \"provider\" in which they are\n", + "authored (i.e., the DSL). All operators corresponding to a given\n", + "provider live a directory corresponding to that provider under\n", + "`cutlass/operators/providers`. For example, CuTe DSL operators live\n", + "under `cutlass/operators/providers/cutedsl`.\n", + "\n", + "Each provider can organize operators differently. For CuTe DSL,\n", + "operators are further split based on their logical operation,\n", + "with GEMM operators under the `cutlass/operators/providers/cutedsl/gemm`\n", + "directory.\n", + "\n", + "We recommend separating the implementation of the Operator from\n", + "its interface not just by using separate classes, as done in\n", + "this example, but also by separating the implementation and\n", + "interface into separate files. This makes it easier to update\n", + "each without affecting the other.\n", + "\n", + "For example, CuTe DSL GEMM operators have the following organization:\n", + "```text\n", + "cutlass/operators/\n", + " providers/\n", + " cutedsl/\n", + " gemm/\n", + " sm100_persistent.py\n", + " implementations/\n", + " sm100_persistent_impl.py\n", + "```" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.5" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/operators/examples/003_host_latency_best_practices.ipynb b/operators/examples/003_host_latency_best_practices.ipynb new file mode 100644 index 000000000..79a6f2b28 --- /dev/null +++ b/operators/examples/003_host_latency_best_practices.ipynb @@ -0,0 +1,619 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "f97e61c9", + "metadata": {}, + "source": [ + "# Best practices for reducing host-side latency" + ] + }, + { + "cell_type": "markdown", + "id": "a7a9c63c", + "metadata": {}, + "source": [ + "Overall performance depends on both device performance (i.e., that of the operator) and host performance (i.e., that of the runtime).\n", + "This notebook focuses on the latter: techniques to minimize any overheads incurred from the CUTLASS API and underlying\n", + "DSL runtimes.\n", + "\n", + "This notebook does not discuss techniques for improving device-side performance. A future notebook may cover this topic.\n", + "\n", + "**Note**: Latency measurements can vary from system to system. You may see different results on your system than shown\n", + "in the pre-populated fields of this notebook." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "e3ca9e40", + "metadata": {}, + "outputs": [], + "source": [ + "import time\n", + "\n", + "import torch\n", + "\n", + "import cutlass.operators as ops" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "efaac09c", + "metadata": {}, + "outputs": [], + "source": [ + "if not (status := ops.utils.device.device_or_env_supports(\"80\")):\n", + " print(\n", + " f\"This notebook requires a GPU with compute capability >= 80.\\n{status.error}\"\n", + " )\n", + " import sys\n", + "\n", + " sys.exit(0)" + ] + }, + { + "cell_type": "markdown", + "id": "40de11ce", + "metadata": {}, + "source": [ + "We start with boilerplate initial setup to create tensors and pick an operator.\n", + "\n", + "For the purposes of this notebook, we use a very small GEMM size of M=N=K=128\n", + "and L=1. This small size is chosen to magnify the impact of host latency on\n", + "end-to-end performance so as to better illustrate the effect of the techniques\n", + "described below." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "b8c44947", + "metadata": {}, + "outputs": [], + "source": [ + "warmup_iterations = 10\n", + "profiling_iterations = 100\n", + "total_iterations = warmup_iterations + profiling_iterations\n", + "\n", + "# Use a small problem size to showcase host overheads\n", + "L, M, N, K = 1, 128, 128, 128\n", + "\n", + "# We use different operands in each iteration. Though not particularly relevant for\n", + "# host latency, this is a best practice when benchmarking GPU kernels to avoid\n", + "# unrealistic caching effects.\n", + "As = [\n", + " torch.randint(-1, 2, (M, K), device=\"cuda\", dtype=torch.float16)\n", + " for _ in range(total_iterations)\n", + "]\n", + "Bs = [\n", + " torch.randint(-1, 2, (K, N), device=\"cuda\", dtype=torch.float16)\n", + " for _ in range(total_iterations)\n", + "]\n", + "outs = [\n", + " torch.empty((M, N), device=\"cuda\", dtype=torch.float16)\n", + " for _ in range(total_iterations)\n", + "]\n", + "\n", + "# Construct arguments outside of the benchmarking loop. We will later also consider\n", + "# cases in which they are constructed inside the benchmarking loop.\n", + "args = [\n", + " ops.GemmArguments(A=As[i], B=Bs[i], out=outs[i], accumulator_type=torch.float32)\n", + " for i in range(total_iterations)\n", + "]\n", + "\n", + "references = [(As[i] @ Bs[i]).to(outs[i].dtype) for i in range(total_iterations)]\n", + "\n", + "target_sm = ops.utils.device.device_or_env_target_sm()\n", + "operators = ops.get_operators(args[0], target_sm=target_sm)\n", + "assert len(operators) > 0\n", + "\n", + "operator = operators[0]" + ] + }, + { + "cell_type": "markdown", + "id": "f2e7eece", + "metadata": {}, + "source": [ + "We next set up a basic benchmarking routine." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "2472eafa", + "metadata": {}, + "outputs": [], + "source": [ + "def benchmark(\n", + " label, code, warmup_it=warmup_iterations, profiling_it=profiling_iterations\n", + "):\n", + " total_it = warmup_it + profiling_it\n", + " assert total_it <= total_iterations, (\n", + " f\"Benchmark-local iteration count must be less than or equal to total iterations: {total_it} > {total_iterations}\"\n", + " )\n", + " # warmup\n", + " rets = [None] * total_it\n", + " for i in range(warmup_it):\n", + " rets[i] = code(i)\n", + " torch.cuda.synchronize()\n", + "\n", + " start = time.time()\n", + " for i in range(profiling_it):\n", + " idx = warmup_it + i\n", + " rets[idx] = code(idx)\n", + " torch.cuda.synchronize()\n", + " end = time.time()\n", + "\n", + " avg_time = (end - start) / profiling_it\n", + " print(f\"[{label:<30}] avg of {profiling_it} iterations: {avg_time:1.3e} seconds\")\n", + " return avg_time, rets" + ] + }, + { + "cell_type": "markdown", + "id": "4909a76b", + "metadata": {}, + "source": [ + "We now describe techniques for reducing host latency:\n", + "* Compile once, run many times\n", + "* Bypassing checks for argument-Operator compatibility\n", + "* Using [CUDA Graphs](https://developer.nvidia.com/blog/cuda-graphs/)\n", + "* Using [TVM FFI](https://tvm.apache.org/ffi/)\n", + "\n", + "These techniques are complementary and should be used together when applicable\n", + "for an application." + ] + }, + { + "cell_type": "markdown", + "id": "06495033", + "metadata": {}, + "source": [ + "### Compile once, run many times\n", + "The `operator.run` method takes in an optional `compiled_artifact` argument of type\n", + "`ops.artifact.CompiledArtifact`. When this argument is set, the operator\n", + "will directly use the precompiled function within `compiled_artifact`. When\n", + "it is not set, the call to `operator.run` will JIT compile the operator on each\n", + "invocation.\n", + "\n", + "Precompiling the operator is critical to achieving good performance." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "6de11f56", + "metadata": {}, + "outputs": [], + "source": [ + "stream = torch.cuda.current_stream()\n", + "\n", + "\n", + "def no_compiled_artifact(i: int):\n", + " return operator.run(args[i], stream=stream)\n", + "\n", + "\n", + "# Compile the Operator once, reuse for each iterations\n", + "compiled_artifact = operator.compile(args[0])\n", + "\n", + "\n", + "def with_compiled_artifact(i: int):\n", + " return operator.run(args[i], stream=stream, compiled_artifact=compiled_artifact)" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "350c9bd6", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[Without compiled artifact ] avg of 5 iterations: 1.474e+00 seconds\n", + "[With compiled artifact ] avg of 5 iterations: 1.988e-05 seconds\n" + ] + } + ], + "source": [ + "time_no_artifact, _ = benchmark(\n", + " \"Without compiled artifact\", no_compiled_artifact, warmup_it=2, profiling_it=5\n", + ")\n", + "time_w_artifact, _ = benchmark(\n", + " \"With compiled artifact\", with_compiled_artifact, warmup_it=2, profiling_it=5\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "5cfbc2d2", + "metadata": {}, + "source": [ + "### Bypassing checks for argument-Operator compatibility\n", + "By default, the call to `operator.run` will check if the Operator supports the provided arguments.\n", + "Under the hood, this invokes `operator.supports(args)`.\n", + "\n", + "While these checks are helpful for catching incompatible arguments, they are performed\n", + "in Python, and thus can add to host overhead.\n", + "\n", + "When confident that arguments will be compatible with an Operator, one should bypass\n", + "the `supports` check in `operator.run` by setting the optional `assume_supported_args`\n", + "argument to `True`." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "5b93dfae", + "metadata": {}, + "outputs": [], + "source": [ + "def with_supports_check(i: int):\n", + " return operator.run(\n", + " args[i],\n", + " compiled_artifact=compiled_artifact,\n", + " stream=stream,\n", + " assume_supported_args=False,\n", + " )\n", + "\n", + "\n", + "def without_supports_check(i: int):\n", + " return operator.run(\n", + " args[i],\n", + " compiled_artifact=compiled_artifact,\n", + " stream=stream,\n", + " assume_supported_args=True,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "b282f437", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[With supports check ] avg of 100 iterations: 2.418e-05 seconds\n", + "[Bypass supports check ] avg of 100 iterations: 6.518e-06 seconds\n", + "Speedup with skip supports: 3.71x\n" + ] + } + ], + "source": [ + "time_w_supports, _ = benchmark(\"With supports check\", with_supports_check)\n", + "time_wo_supports, _ = benchmark(\"Bypass supports check\", without_supports_check)\n", + "print(f\"Speedup with skip supports: {time_w_supports / time_wo_supports:.2f}x\")" + ] + }, + { + "cell_type": "markdown", + "id": "d74cb3e7", + "metadata": {}, + "source": [ + "### CUDA Graphs" + ] + }, + { + "cell_type": "markdown", + "id": "656d5e2c", + "metadata": {}, + "source": [ + "[CUDA Graphs](https://developer.nvidia.com/blog/cuda-graphs/) allow a sequence of GPU operations to be defined as a dependency graph and then launched as a single unit, significantly reducing CPU launch overhead and enabling whole-graph optimizations.\n", + "\n", + "CUTLASS API supports CUDA Graphs usage with PyTorch as usual.\n", + "\n", + "The operator compilation must happen outside the CUDA graph. Then, we create a graph using usual PyTorch idioms to launch a kernel several times on the graph's stream." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "e614509f", + "metadata": {}, + "outputs": [], + "source": [ + "num_launches = 20\n", + "\n", + "# Create a CUDA Graph to run our compiled Operator N times\n", + "g = torch.cuda.CUDAGraph()\n", + "with torch.cuda.graph(g):\n", + " ### NOTE! Operator compilation must happen outside the graph\n", + " ### operator.compile(args)\n", + "\n", + " # Run N iterations of our compiled Operator on the current stream\n", + " for i in range(num_launches):\n", + " operator.run(\n", + " args[i],\n", + " compiled_artifact=compiled_artifact,\n", + " stream=torch.cuda.current_stream(),\n", + " assume_supported_args=True,\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "8fc69c6e", + "metadata": {}, + "source": [ + "This records/captures all the kernel launches to the CUDA Stream associated with the graph `g`, without actually launching them.\n", + "Once captured, we can replay the graph.\n", + "\n", + "Note that graph replay will only replay the kernel launches placed on the graph's stream\n", + "* During graph capture, we must be careful to capture to the correct stream (`torch.cuda.current_stream()` under the graph context)\n", + "* Any other preparatory work on the host and arguments passed in from Python are cached during the capture. Changing them would require re-capturing the graph" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "d9c5d5c5", + "metadata": {}, + "outputs": [], + "source": [ + "# Replay captured graph and check first result\n", + "g.replay()\n", + "\n", + "torch.testing.assert_close(outs[0], references[0])" + ] + }, + { + "cell_type": "markdown", + "id": "388c8e02", + "metadata": {}, + "source": [ + "Let's compare the timing:" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "45d4e739", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[20 launches without CUDA Graph] avg of 1 iterations: 7.496e-04 seconds\n", + "[20 launches with CUDA Graph ] avg of 1 iterations: 1.972e-04 seconds\n", + "Speedup with CUDA Graph: 3.80x\n" + ] + } + ], + "source": [ + "def without_cuda_graph(x: int):\n", + " for i in range(num_launches):\n", + " operator.run(\n", + " args[i],\n", + " compiled_artifact=compiled_artifact,\n", + " stream=torch.cuda.current_stream(),\n", + " assume_supported_args=True,\n", + " )\n", + "\n", + "\n", + "def with_cuda_graph(x: int):\n", + " g.replay()\n", + "\n", + "\n", + "time_wo_cuda_graph, _ = benchmark(\n", + " f\"{num_launches} launches without CUDA Graph\",\n", + " without_cuda_graph,\n", + " warmup_it=0,\n", + " profiling_it=1,\n", + ")\n", + "time_w_cuda_graph, _ = benchmark(\n", + " f\"{num_launches} launches with CUDA Graph\",\n", + " with_cuda_graph,\n", + " warmup_it=0,\n", + " profiling_it=1,\n", + ")\n", + "\n", + "print(f\"Speedup with CUDA Graph: {time_wo_cuda_graph / time_w_cuda_graph:.2f}x\")" + ] + }, + { + "cell_type": "markdown", + "id": "fe5c3168", + "metadata": {}, + "source": [ + "### TVM FFI" + ] + }, + { + "cell_type": "markdown", + "id": "ee7f9fd2", + "metadata": {}, + "source": [ + "[Apache TVM FFI](https://tvm.apache.org/ffi/) is an open ABI and FFI for machine learning systems.\n", + "When available, CUTLASS API uses Apache TVM-FFI under the hood as its interface for invoking compiled DSL kernels from Python.\n", + "\n", + "TVM FFI is enabled by default in CUTLASS API, and is recommended for best performance." + ] + }, + { + "cell_type": "markdown", + "id": "1690bbed", + "metadata": {}, + "source": [ + "`ops.GlobalOptions().use_tvm_ffi` controls whether or not TVM-FFI will be used by CUTLASS API." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "993c60ae", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "True\n" + ] + } + ], + "source": [ + "print(ops.GlobalOptions().use_tvm_ffi)" + ] + }, + { + "cell_type": "markdown", + "id": "00ed9a40", + "metadata": {}, + "source": [ + "If for some reason you do not wish to use it, this section demonstrates how, you can set this to False. No other change is needed. The below code compares the performance with and without TVM-FFI." + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "e8f56be3", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[[TVM-FFI ON ] Create args ] avg of 100 iterations: 1.595e-04 seconds\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[[TVM-FFI ON ] Compile operator] avg of 5 iterations: 1.469e+00 seconds\n", + "[[TVM-FFI ON ] Launch host overhead ] avg of 100 iterations: 5.984e-06 seconds\n" + ] + } + ], + "source": [ + "original_use_tvm_ffi = ops.GlobalOptions().use_tvm_ffi\n", + "\n", + "ops.GlobalOptions().use_tvm_ffi = True\n", + "\n", + "\n", + "def run_iteration(i):\n", + " args = ops.GemmArguments(\n", + " A=As[i], B=Bs[i], out=outs[i], accumulator_type=torch.float16\n", + " )\n", + " return operator.run(\n", + " args,\n", + " compiled_artifact=compiled_artifact,\n", + " stream=torch.cuda.current_stream(),\n", + " assume_supported_args=True,\n", + " )\n", + "\n", + "\n", + "def create_arguments(i: int):\n", + " return ops.GemmArguments(\n", + " A=As[i], B=Bs[i], out=outs[i], accumulator_type=torch.float16\n", + " )\n", + "\n", + "\n", + "args_creation_on, args = benchmark(\"[TVM-FFI ON ] Create args\", create_arguments)\n", + "compilation_on, compiled = benchmark(\n", + " \"[TVM-FFI ON ] Compile operator\",\n", + " lambda i: operator.compile(args[i]),\n", + " warmup_it=2,\n", + " profiling_it=5,\n", + ")\n", + "compiled_artifact = compiled[0]\n", + "run_on, _ = benchmark(\n", + " \"[TVM-FFI ON ] Launch host overhead \",\n", + " lambda i: operator.run(\n", + " args[i],\n", + " compiled_artifact=compiled_artifact,\n", + " assume_supported_args=True,\n", + " stream=stream,\n", + " ),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "5a4c2db4", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[[TVM-FFI OFF ] Create args ] avg of 100 iterations: 2.151e-04 seconds\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[[TVM-FFI OFF ] Compile operator] avg of 5 iterations: 1.348e+00 seconds\n", + "[[TVM-FFI OFF ] Launch host overhead ] avg of 100 iterations: 5.973e-05 seconds\n" + ] + } + ], + "source": [ + "ops.GlobalOptions().use_tvm_ffi = False\n", + "args_creation_off, args = benchmark(\"[TVM-FFI OFF ] Create args\", create_arguments)\n", + "compilation_off, compiled = benchmark(\n", + " \"[TVM-FFI OFF ] Compile operator\",\n", + " lambda i: operator.compile(args[i]),\n", + " warmup_it=2,\n", + " profiling_it=5,\n", + ")\n", + "compiled_artifact = compiled[0]\n", + "run_off, _ = benchmark(\n", + " \"[TVM-FFI OFF ] Launch host overhead \",\n", + " lambda i: operator.run(\n", + " args[i],\n", + " compiled_artifact=compiled_artifact,\n", + " assume_supported_args=True,\n", + " stream=stream,\n", + " ),\n", + ")\n", + "\n", + "# Restore original setting\n", + "ops.GlobalOptions().use_tvm_ffi = original_use_tvm_ffi" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "17b43718", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Speedups with TVM-FFI: \n", + "Arg creation: 1.35x\n", + "Compilation: 0.92x\n", + "Launch host overhead: 9.98x\n" + ] + } + ], + "source": [ + "print(\"Speedups with TVM-FFI: \")\n", + "print(f\"Arg creation: {args_creation_off / args_creation_on:.2f}x\")\n", + "print(f\"Compilation: {compilation_off / compilation_on:.2f}x\")\n", + "print(f\"Launch host overhead: {run_off / run_on:.2f}x\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/operators/examples/004_fake_tensors.ipynb b/operators/examples/004_fake_tensors.ipynb new file mode 100644 index 000000000..102444062 --- /dev/null +++ b/operators/examples/004_fake_tensors.ipynb @@ -0,0 +1,167 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "4620d513", + "metadata": {}, + "source": [ + "# Using fake tensors with the CUTLASS Operator API\n", + "Fake tensors (e.g., [torch's FakeTensor](https://docs.pytorch.org/docs/2.8/torch.compiler_fake_tensor.html))\n", + "are useful for describing the properties of a tensor without actually allocating backing data.\n", + "\n", + "This example shows how fake tensors can be used within the CUTLASS Operator API\n", + "for discovering and compiling a GEMM Operator." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "d231b32e", + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "\n", + "import cutlass.operators as ops\n", + "\n", + "torch.manual_seed(2025)\n", + "\n", + "if not (status := ops.utils.device.device_or_env_supports(\"80\")):\n", + " print(\n", + " f\"This notebook requires a GPU with compute capability >= 80.\\n{status.error}\"\n", + " )\n", + " import sys\n", + "\n", + " sys.exit(0)" + ] + }, + { + "cell_type": "markdown", + "id": "f7af2d90", + "metadata": {}, + "source": [ + "We first set up operands `A`, `B`, and `out` in torch's `FakeTensorMode`.\n", + "These will have all the properties needed for CUTLASS Operator API to construct\n", + "the internal representations of tensors used for discovering and compiling\n", + "operators." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "9426b66f", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "FakeTensor(..., device='cuda:0', size=(128, 512), dtype=torch.float16)\n", + "FakeTensor(..., device='cuda:0', size=(512, 256), dtype=torch.float16)\n", + "FakeTensor(..., device='cuda:0', size=(128, 256), dtype=torch.float16)\n" + ] + } + ], + "source": [ + "M, N, K = 128, 256, 512\n", + "\n", + "with torch._subclasses.fake_tensor.FakeTensorMode():\n", + " A_fake = torch.randint(-2, 3, (M, K), device=\"cuda\", dtype=torch.float16)\n", + " B_fake = torch.randint(-2, 3, (K, N), device=\"cuda\", dtype=torch.float16)\n", + " out_fake = torch.empty(M, N, device=\"cuda\", dtype=torch.float16)\n", + "\n", + "print(A_fake)\n", + "print(B_fake)\n", + "print(out_fake)" + ] + }, + { + "cell_type": "markdown", + "id": "4f540c78", + "metadata": {}, + "source": [ + "We can now use these fake tensors to create `GemmArguments`, and use\n", + "these to discover and compile a compatible operator. Note that the same APIs are\n", + "used in creating `GemmArguments` as would be used if using\n", + "\"real\" tensors." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "e32b700d", + "metadata": {}, + "outputs": [], + "source": [ + "args_fake = ops.GemmArguments(A_fake, B_fake, out_fake, accumulator_type=torch.float32)\n", + "\n", + "target_sm = ops.utils.device.device_or_env_target_sm()\n", + "operators = ops.get_operators(args_fake, target_sm=target_sm)\n", + "assert len(operators) > 0\n", + "\n", + "operator = operators[0]\n", + "compiled_artifact = operator.compile(args_fake)" + ] + }, + { + "cell_type": "markdown", + "id": "07fff511", + "metadata": {}, + "source": [ + "The `operator` and `compiled_artifact` discovered using fake tensors\n", + "above can now used for running the operator using real tensors." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "b3034bf1", + "metadata": {}, + "outputs": [], + "source": [ + "# Create real tensors\n", + "A_real = torch.randint(-2, 3, (M, K), device=\"cuda\", dtype=torch.float16)\n", + "B_real = torch.randint(-2, 3, (K, N), device=\"cuda\", dtype=torch.float16)\n", + "out_real = torch.empty(M, N, device=\"cuda\", dtype=torch.float16)\n", + "\n", + "args_real = ops.GemmArguments(A_real, B_real, out_real, accumulator_type=torch.float32)\n", + "\n", + "# Run the operator using the compiled_artifact from resulting\n", + "# from compiling with fake tensors.\n", + "operator.run(args_real, compiled_artifact)" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "09871eca", + "metadata": {}, + "outputs": [], + "source": [ + "ref = A_real @ B_real\n", + "torch.testing.assert_close(out_real, ref)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.5" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/operators/examples/005_grouped_gemm_contiguous_offset.ipynb b/operators/examples/005_grouped_gemm_contiguous_offset.ipynb new file mode 100644 index 000000000..e5ae0b323 --- /dev/null +++ b/operators/examples/005_grouped_gemm_contiguous_offset.ipynb @@ -0,0 +1,246 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "91d43c2b", + "metadata": {}, + "source": [ + "# Grouped GEMM with contiguous tensors via the CUTLASS Operator API\n", + "\n", + "Note: this notebook requires a GPU with compute capability 100:" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "f671f602", + "metadata": {}, + "outputs": [], + "source": [ + "import cutlass.operators as ops\n", + "\n", + "if not (status := ops.utils.device.device_or_env_supports(\"100f\")):\n", + " print(f\"This notebook requires a Blackwell GPU (sm_100f family).\\n{status.error}\")\n", + " import sys\n", + "\n", + " sys.exit(0)" + ] + }, + { + "cell_type": "markdown", + "id": "bc4adf7d", + "metadata": {}, + "source": [ + "This notebook shows how to use the CUTLASS Operator API to discover, compile, and execute\n", + "operators supporting contiguous offset grouped GEMMs.\n", + "\n", + "In a \"contiguous offset\" grouped GEMM, `G` different problems are executed\n", + "in which problems differ only in the `M` mode. Their problem sizes are thus\n", + "represented as:\n", + "\n", + "```text\n", + "M0 x N x K\n", + "M1 x N x K\n", + "M2 x N x K\n", + "...\n", + "M(G-1) x N x K\n", + "```\n", + "\n", + "The grouped GEMM is referred to as \"contiguous\" because operands for different\n", + "problems in the group are contained within contiguous tensors.\n", + "\n", + "Rather than having `G` different tensors for each of operands `A` and `B`, tensors\n", + "for different problems in the group are packed together:\n", + "* `A` is of shape `(TotalM, K)`, where `TotalM` is the sum of all `M` mode sizes for problems in the group.\n", + "The `A` operands for each problem in the group are stacked along the `M` mode to form this input. More on this below.\n", + "* `B` is of shape `(G, K, N)`, where `B[i, :, :]` represents the GEMM `B` operand for the `i`th problem in the group.\n", + "\n", + "For example, with `G=3` (three problems in the group), with `M` mode sizes of M0, M1, and M2,\n", + "respectively, the tensor `A` would be laid out as follows:\n", + "\n", + "```text\n", + "\n", + " +----------------------------------+ ^ \n", + " | | | | \n", + " | A0 | M0 | \n", + " | | | | \n", + " |- - - - - - - - - - - -| | \n", + " | | | |\n", + " | | | TotalM \n", + " | A1 | M1 |\n", + " | | | |\n", + " | | | | \n", + " |- - - - - - - - - - - -| | \n", + " | A2 | M2 | \n", + " +----------------------------------+ v \n", + "```\n", + "\n", + "The extents of individual `A` operands packed within the overall contiguous offset `A` tensor\n", + "are provided by an auxiliary `offsets` vector of shape `(G,)`. `offsets[i]` indicates the ending\n", + "M coordinate (exclusive) for the `i`th `A` operand.\n", + "\n", + "Thus, for the example above, `offsets = [M0, M0 + M1, M0 + M1 + M2]`.\n", + "\n", + "The output of the operation is of shape `(TotalM, N)`. The `i`th output occupies `out[start:end, :]`,\n", + "where `start` and `end` are `offsets[i-1]` and `offsets[i]`, respectively (unless `i=0`, in which case\n", + "`start` is 0).\n", + "\n", + "The reference code below shows the computation of this Operator." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "6185f60a", + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "\n", + "\n", + "def reference_contiguous_offset_grouped_gemm(A, B, offsets, out_dtype):\n", + " G, K, N = B.shape\n", + " TotalM = A.shape[0]\n", + "\n", + " out = torch.empty((TotalM, N), dtype=out_dtype, device=A.device)\n", + "\n", + " start = 0\n", + " for i in range(G):\n", + " end = offsets[i]\n", + " out[start:end, :] = A[start:end, :] @ B[i, :, :]\n", + " start = end\n", + "\n", + " return out" + ] + }, + { + "cell_type": "markdown", + "id": "d0bf2f91", + "metadata": {}, + "source": [ + "## Contiguous offset grouped GEMM in PyTorch" + ] + }, + { + "cell_type": "markdown", + "id": "4308a6a2", + "metadata": {}, + "source": [ + "The same operation is performed by `torch`'s `torch._grouped_mm` (torch < 2.10)\n", + "and `torch.nn.functional.grouped_mm` (torch >= 2.10)." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "043906af", + "metadata": {}, + "outputs": [], + "source": [ + "TotalM = 8192\n", + "G = 12\n", + "K = 1024\n", + "N = 2048\n", + "\n", + "offsets = torch.arange(\n", + " TotalM // G, TotalM, TotalM // G, device=\"cuda\", dtype=torch.int32\n", + ")\n", + "offsets[-1] = TotalM\n", + "\n", + "A = torch.randint(-2, 3, (TotalM, K), device=\"cuda\", dtype=torch.bfloat16)\n", + "B = torch.randint(-2, 3, (G, N, K), device=\"cuda\", dtype=torch.bfloat16).permute(\n", + " 0, 2, 1\n", + ")\n", + "\n", + "out_torch = torch._grouped_mm(A, B, offsets, out_dtype=torch.bfloat16)\n", + "reference = reference_contiguous_offset_grouped_gemm(\n", + " A, B, offsets, out_dtype=torch.bfloat16\n", + ")\n", + "\n", + "torch.testing.assert_close(out_torch, reference)" + ] + }, + { + "cell_type": "markdown", + "id": "0d0e9479", + "metadata": {}, + "source": [ + "## Contiguous offset grouped GEMM in CUTLASS Operator API\n", + "\n", + "CUTLASS Operator API exposes this contiguous offset grouped GEMM via `GroupedGemmArguments`,\n", + "which are constructed similarly to `GemmArguments`, but take in an `offsets`\n", + "tensor as well:" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "ff8d3ef1", + "metadata": {}, + "outputs": [], + "source": [ + "out = torch.empty((TotalM, N), device=\"cuda\", dtype=torch.bfloat16)\n", + "\n", + "args = ops.GroupedGemmArguments(\n", + " A,\n", + " B,\n", + " out,\n", + " accumulator_type=torch.float32,\n", + " offsets=offsets,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "0dc6d1cb", + "metadata": {}, + "source": [ + "One can then use the same APIs for finding, compiling, and executing a\n", + "operator supporting this operation" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "80213e1e", + "metadata": {}, + "outputs": [], + "source": [ + "operators = ops.get_operators(args)\n", + "\n", + "assert operators, \"No operators found\"\n", + "\n", + "# Select the first operator found for simplicity\n", + "operator = operators[0]\n", + "\n", + "compiled_artifact = operator.compile(args)\n", + "\n", + "# Execute the operator\n", + "operator.run(args, compiled_artifact=compiled_artifact)\n", + "\n", + "torch.testing.assert_close(out, reference)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.5" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/operators/examples/006_block_scaled_gemm.ipynb b/operators/examples/006_block_scaled_gemm.ipynb new file mode 100644 index 000000000..b032b2dee --- /dev/null +++ b/operators/examples/006_block_scaled_gemm.ipynb @@ -0,0 +1,352 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "e6f682b3", + "metadata": {}, + "source": [ + "# Block-scaled GEMM using CUTLASS Operator API\n", + "\n", + "This notebook walks through using CUTLASS Operator API to run *block-scaled* GEMMs -- the operations behind matrix-multiplication for narrow-precision types (like MXFP8).\n", + "\n", + "It builds on [the basic GEMM tutorial](000_gemm.ipynb), which introduced the flow for a plain GEMM.\n", + "\n", + "**Key takeaway:** *the same flow works here*. A block-scaled GEMM is still logically `out = A @ B`, only the operands are stored differently. Most of this notebook explains how to construct an `Operand` describing block-scaling. After that, expressing the GEMM, finding a supported kernel, compiling, and running it is completely identical to the basic GEMM." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "403291e8", + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "\n", + "import cutlass.operators as ops\n", + "\n", + "if not (status := ops.utils.device.device_or_env_supports(\"100f\")):\n", + " print(f\"This notebook requires a Blackwell GPU (sm_100f family).\\n{status.error}\")\n", + " import sys\n", + "\n", + " sys.exit(0)" + ] + }, + { + "cell_type": "markdown", + "id": "64e4a31c", + "metadata": {}, + "source": [ + "## What is a block-scaled GEMM?\n", + "\n", + "In a plain GEMM, we compute `out = A @ B` and the operands carry their own values in a single tensor. For narrower formats -- such as 8-bit/4-bit floating point -- a single tensor can no longer represent both the small magnitudes and the large outliers of a typical activation or weight tensor. The block-scaled formulation **quantizes** the bulk of the data into a narrow-precision tensor, and recovers dynamic range by multiplying it with a smaller separate tensor of **scale factors**. Each scale factor scales a small, contiguous block of the quantized tensor: `out = (scale_A * quantized_A) @ (scale_B * quantized_B)`\n", + "\n", + "\n", + "A block of `V` elements in `quantized_A` share the same `scale_A` multiplier factor (likewise for `B`). So, the scale factor's size is `1/V` of the quantized tensor's.\n", + "\n", + "Blackwell Tensor Cores can apply the scale-and-matmul in a single MMA instruction in hardware, rather than as separate steps. The format we demonstrate here is **MXFP8**: FP8 E4M3FN quantized values (8-bit, 4 exponent / 3 mantissa) paired with E8M0 scales -- one scale per 32 K-elements -- following the OCP Microscaling Format spec." + ] + }, + { + "cell_type": "markdown", + "id": "8adef2a9", + "metadata": {}, + "source": [ + "## Operands: same `A @ B`, different storage\n", + "\n", + "CUTLASS Operator API draws a clean line between a *logical operand* of an operation (`A` in `out = A @ B`), and the *physical tensor(s)* that store its values (`A = scale_A * quantized_A`).\n", + "\n", + "The shared abstraction is `Operand`. It represents the logical operand, and encapsulates its constituent physical tensor(s). `DenseTensor` and `ScaledOperand` are 2 types of `Operand`s:\n", + "\n", + "* A plain GEMM computes `out = A @ B` and the `Operand`s in it are `DenseTensor`. The logical value of the operand is trivially stored in the tensor.\n", + "* A block-scaled GEMM still logically computes `out = A @ B`, and the `Operand`s are `ScaledOperand`s. The logical values of `A` & `B` need to be reconstructed by multiplying the scale factors over the quantized tensors, upon which we then perform matrix multiplication as usual.\n", + "\n", + "`GemmArguments` expect A & B to be `Operand`s, so both forms work with it -- either way, `GemmArguments` express a logical GEMM over A and B.\n", + "\n", + "**Once the operation is expressed, the discovery / compile / run flow is exactly what you saw in [the basic GEMM tutorial](000_gemm.ipynb).**\n", + "\n", + "From this point in the notebook, the bulk of the work is constructing the `ScaledOperand` correctly." + ] + }, + { + "cell_type": "markdown", + "id": "a3564625", + "metadata": {}, + "source": [ + "## Constructing `ScaledOperand`s\n", + "\n", + "A `ScaledOperand` has the following fields:\n", + "\n", + "* `quantized`: the narrow-precision tensor holding the quantized values.\n", + "* `scale`: a tensor of scale factors, one per block of `quantized`.\n", + "* `mode`: the *block shape* over which each scale factor is broadcast, given as an `(L, M, K)` tuple. `ops.ScaleMode.Blockwise1x32` is `(1, 1, 32)`, meaning \"per-element along L and M, every 32 elements along K\".\n", + "* `swizzle`: how the scale tensor is laid out in memory. `SwizzleNone` is the natural shape (`L, M, K // V`), one scale per `mode` block. `Swizzle32x4x4` is a special layout based on Blackwell MMA requirements, see `ScaleSwizzleMode` for details.\n", + "\n", + "The latter two enums map directly to `torch.nn.functional.scaled_mm`, for easy interop:\n", + "* `ScaleMode` is the analog of `torch.nn.functional.ScalingType` (with members like `BlockWise1x32`)\n", + "* `ScaleSwizzleMode` is the analog of `torch.nn.functional.SwizzleType`\n", + "\n", + "Like everything passed to `GemmArguments`, the fields above describe the operation and operands we have -- including their dtypes, block-scaling granularity, and swizzle layout -- independent of which operator will consume it. `get_operators(args)` then finds operators that support consuming operands described that way." + ] + }, + { + "cell_type": "markdown", + "id": "169d5949", + "metadata": {}, + "source": [ + "### Step 1: Choose a format and a problem size\n", + "\n", + "We pick MXFP8 format, which requires:\n", + "* quantized tensor of E4M3FN type\n", + "* scale factors of E8M0 type, shared by a block of 1x32 elements\n", + "\n", + "This fixes the dtypes of the quantized & scale tensors, and `ScaleMode`.\n", + "\n", + "\n", + "The `swizzle` tag is determined by Blackwell's `tcgen05.mma` block-scaled MMA instruction (see documentation from [cuDNN](https://nvidia.github.io/cudnn-frontend/mxfp8-scale-factor-128x4-layout/) and [cuBLAS](https://docs.nvidia.com/cuda/cublas/#d-block-scaling-factors-layout) for details). We tag our scale tensors as `Swizzle32x4x4` to declare that we are providing data arranged in this layout." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "6735270f", + "metadata": {}, + "outputs": [], + "source": [ + "M, N, K, L = 256, 512, 1024, 1\n", + "\n", + "ab_dtype = torch.float8_e4m3fn # narrow quantized type (MXFP8)\n", + "scale_dtype = torch.float8_e8m0fnu # E8M0 scale type per the OCP MX spec\n", + "out_dtype = torch.float32\n", + "\n", + "mode = ops.ScaleMode.Blockwise1x32 # one scale per 32 K-elements\n", + "swizzle = ops.ScaleSwizzleMode.Swizzle32x4x4 # tcgen05 block-scaled layout" + ] + }, + { + "cell_type": "markdown", + "id": "5dca9d12", + "metadata": {}, + "source": [ + "### Step 2: Build the quantized tensors\n", + "\n", + "The quantized tensors `quantized_A` and `quantized_B` are no different than a DenseTensor from a plain GEMM. For a `(L, M, K) @ (L, K, N) -> (L, M, N)` GEMM, the Blackwell tcgen05 MMA requires both `quantized_A` and `quantized_B` to be **K-major** -- that is, `quantized_A` (with K columns) must be row-major, and `quantized_B` (with K rows) must be column-major. We construct those tensors here." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "50ec843c", + "metadata": {}, + "outputs": [], + "source": [ + "quantized_A = torch.randint(-1, 2, (L, M, K), device=\"cuda\").to(ab_dtype)\n", + "quantized_B = torch.randint(-1, 2, (L, N, K), device=\"cuda\").to(ab_dtype).transpose(1, 2)\n", + "out = torch.empty((L, M, N), device=\"cuda\", dtype=out_dtype)" + ] + }, + { + "cell_type": "markdown", + "id": "d1f70d97", + "metadata": {}, + "source": [ + "### Step 3: Build the scale tensors\n", + "\n", + "Logically, each scale factor in `scale_A` covers a `(1, 1, 32)` block of `quantized_A`, so `scale_A` carries one value per row of `A` and one per group of 32 K-elements -- i.e. a logical shape of `(L, M, ceil(K / 32))`.\n", + "\n", + "`Swizzle32x4x4` packs those values into `128 × 4` tiles. The tile boundaries impose two alignment requirements on the logical shape:\n", + "\n", + "* the outer dimension (`M` for `scale_A`) is padded to a multiple of **128** -- the leading `32 × 4` in `Swizzle32x4x4`\n", + "* the inner dimension (K-blocks) is padded to a multiple of **4** -- the trailing `4` in the name\n", + "\n", + "Thus the required number of elements for `scale_A` is `L * round_up(M, 128) * round_up(ceil(K / 32), 4)`.\n", + "\n", + "`scale_B` is also handled similarly, with `M` replaced by `N`.\n", + "\n", + "The result of this calculation is also exposed as a helper, `ScaledOperand.numel_scale(quantized_shape, mode, swizzle)`, which we use below." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15fa8725", + "metadata": {}, + "outputs": [], + "source": [ + "scale_A_numel = ops.ScaledOperand.numel_scale((L, M, K), mode, swizzle)\n", + "scale_B_numel = ops.ScaledOperand.numel_scale((L, N, K), mode, swizzle)\n", + "\n", + "scale_A = torch.rand(scale_A_numel, device=\"cuda\").to(scale_dtype)\n", + "scale_B = torch.rand(scale_B_numel, device=\"cuda\").to(scale_dtype)" + ] + }, + { + "cell_type": "markdown", + "id": "0ee139c2", + "metadata": {}, + "source": [ + "**API contract.** Operator API only requires that the scale tensor is contiguous and has `numel_scale(...)` elements. The Python shape is otherwise free -- the 1D buffers above are valid input as-is. If a structured shape is more convenient, any reshape preserving the element count and contiguity also works:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ba7d2628", + "metadata": {}, + "outputs": [], + "source": [ + "from cutlass.operators.utils.common import round_up\n", + "\n", + "scale_A = scale_A.view(L, round_up(M, 128), -1)\n", + "print(f\"scale_A: shape={tuple(scale_A.shape)}, contiguous={scale_A.is_contiguous()}\")\n", + "print(f\"scale_B: shape={tuple(scale_B.shape)}, contiguous={scale_B.is_contiguous()}\")" + ] + }, + { + "cell_type": "markdown", + "id": "c431868b", + "metadata": {}, + "source": [ + "Most kernels for block-scaled GEMM expect that the scale storage is already laid out in the swizzled layout used by hardware. When we set `swizzle=Swizzle32x4x4`, we declare that we are providing input in this layout. In a real workflow, that layout is produced by a quantization step or an earlier kernel that writes the scale values directly into it.\n", + "\n", + "We don't explicitly perform that rearrangement in this tutorial. We'll use `torch._scaled_mm` for checking our output, which also expects data in the same swizzled layout. Since kernels interpret the storage the same way, we can cross-check our result without an explicit swizzle step." + ] + }, + { + "cell_type": "markdown", + "id": "c5fb63ec", + "metadata": {}, + "source": [ + "### Step 4: Pair the quantized tensor with its scale via `ScaledOperand`\n", + "\n", + "Now we package each quantized tensor with its scale tensor, `mode`, and `swizzle` into a `ScaledOperand`. Once paired, `GemmArguments` looks the same as in the plain GEMM tutorial -- the only difference is that `A` and `B` are `ScaledOperand`s rather than plain tensors." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "1aae7592", + "metadata": {}, + "outputs": [], + "source": [ + "args = ops.GemmArguments(\n", + " A=ops.ScaledOperand(quantized_A, scale_A, mode, swizzle),\n", + " B=ops.ScaledOperand(quantized_B, scale_B, mode, swizzle),\n", + " out=out, # same as DenseTensor(out)\n", + " accumulator_type=torch.float32,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "d7a2d2ae", + "metadata": {}, + "source": [ + "### Step 5: Discover, compile, and run\n", + "\n", + "**From this point on, everything is identical to any other operation!**\n", + "\n", + "Despite the different set of tensors passed to the underlying kernel, the interface with the Operator remains the same -- there is no separate `BlockScaledGemmArguments`, `block_scaled_get_operators`, or `run_block_scaled` API. We hand `args` to `get_operators(...)`, which walks the registered Operator catalog and finds the ones compatible with the operation we described -- same as any other operation.\n", + "\n", + "We describe **what** operation to perform (via `GemmArguments`); the API finds operators that know **how** to do it.\n", + "\n", + "We then pick the first one and run it." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "b0bee4ac", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Picked: cutedsl.PersistentDenseBlockScaledGemmOperator_sm100_tnt_AFloat8E4M3FN_BFloat8E4M3FN_outFloat32_SFAFloat8E8M0FNU_SFBFloat8E8M0FNU_accFloat32_scaleScaleMode.Blockwise1x32_swizzleScaleSwizzleMode.Swizzle32x4x4_2cta_cluster2x1x1_tile128x64x128_tma_store\n" + ] + } + ], + "source": [ + "target_sm = ops.utils.device.device_or_env_target_sm()\n", + "\n", + "operators = ops.get_operators(args, target_sm=target_sm)\n", + "assert operators, \"No block-scaled GEMM operator matched the given args\"\n", + "operator = operators[0]\n", + "print(f\"Picked: {operator.metadata.operator_name}\")\n", + "\n", + "operator.run(args)" + ] + }, + { + "cell_type": "markdown", + "id": "2388e225", + "metadata": {}, + "source": [ + "## Verifying against `torch._scaled_mm`\n", + "\n", + "To confirm the result, we cross-check `out` against PyTorch's reference implementation, `torch._scaled_mm`." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "c2e761c7", + "metadata": {}, + "outputs": [], + "source": [ + "def reference_scaled_mm(A, B, scale_A, scale_B, out_dtype):\n", + " L, M, N = A.shape[0], A.shape[1], B.shape[2]\n", + " scale_A = scale_A.reshape(L, -1)\n", + " scale_B = scale_B.reshape(L, -1)\n", + " out = torch.empty((L, M, N), device=A.device, dtype=out_dtype)\n", + " for L_idx in range(L):\n", + " out[L_idx] = torch._scaled_mm(\n", + " A[L_idx], B[L_idx], scale_a=scale_A[L_idx], scale_b=scale_B[L_idx], out_dtype=out_dtype\n", + " )\n", + " return out\n", + "\n", + "\n", + "reference = reference_scaled_mm(quantized_A, quantized_B, scale_A, scale_B, out_dtype)\n", + "torch.testing.assert_close(out, reference)" + ] + }, + { + "cell_type": "markdown", + "id": "a3a121a3", + "metadata": {}, + "source": [ + "## Recap\n", + "\n", + "You've run an MXFP8 block-scaled GEMM end-to-end through the same Operator API entry points used in the basic GEMM tutorial:\n", + "\n", + "1. **Describe the operation with `GemmArguments`.** Each slot is an `Operand` -- a logical role like \"the matrix `A`\" -- and accepts whatever storage form the operand needs: a `DenseTensor` (any DLPack-compatible tensor) or a `ScaledOperand` when the value is reconstructed from a quantized tensor + scale.\n", + "2. **Discover with `get_operators(args)`.** The Operator catalog filters down to the kernels that match the arguments you described.\n", + "3. **Compile and run with `operator.compile(args)` / `operator.run(args)`.** Identical to the basic GEMM flow.\n", + "\n", + "What changed from the basic to the block-scaled tutorial is *only how the operands are described*: `ScaledOperand` instead of plain tensors. We express the operation we want to perform, and Operator API maps it to a supported kernel. The same pattern -- describe `GemmArguments` -> `get_operators` -> `compile` / `run` -- works despite different kernel implementations." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.8" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/operators/pyproject.toml b/operators/pyproject.toml new file mode 100644 index 000000000..402ee9a5c --- /dev/null +++ b/operators/pyproject.toml @@ -0,0 +1,175 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +[build-system] +requires = ["setuptools>=77"] +build-backend = "setuptools.build_meta" + +[project] +name = "nvidia-cutlass-operators" +dynamic = ["version"] +description = "High-level discover/compile/execute API for CUTLASS Python kernels." +readme = "README.md" +license = "BSD-3-Clause" +authors = [{ name = "NVIDIA Corporation" }] +requires-python = ">=3.10" +keywords = ["cutlass", "cuda", "gpu", "gemm", "nvidia", "deep-learning"] +classifiers = [ + "Development Status :: 4 - Beta", + "Environment :: GPU :: NVIDIA CUDA :: 13", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "Operating System :: POSIX :: Linux", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Topic :: Scientific/Engineering", + "Topic :: Software Development :: Libraries", +] + +dependencies = [ + "nvidia-cutlass-dsl >= 4.6", + "apache-tvm-ffi >= 0.1.11", # Required for best performance. To opt-out: uninstall, set CUTLASS_OPERATORS_USE_TVM_FFI=0, or set cutlass.operators.config.GlobalOptions().use_tvm_ffi=False + "networkx >= 3.0.0", # Required by dag-ir.py +] + +[project.urls] +Homepage = "https://github.com/NVIDIA/cutlass" +Repository = "https://github.com/NVIDIA/cutlass" +Documentation = "https://docs.nvidia.com/cutlass/latest/" +Issues = "https://github.com/NVIDIA/cutlass/issues" + +[project.optional-dependencies] +torch = [ + "torch", + "torch-c-dlpack-ext", +] +jax = [ + # jax[cuda13] first appears in jax>=0.7.0, which dropped Python 3.10 + # (min is 3.11). So 3.10 gets the CUDA 12 wheels, 3.11+ gets CUDA 13. + "jax[cuda12]; python_version < '3.11'", + "jax[cuda13]; python_version >= '3.11'", +] +test = [ + "jupyter", + "nbconvert", + "nbformat", + "pytest", + "nvidia-cutlass-operators[torch]", + "nvidia-cutlass-operators[jax]", +] +dev = [ + "nvidia-cutlass-operators[test]", + "ruff", + "darker[ruff] == 3.0.0", +] + +# Discover the `cutlass` namespace package and all `cutlass.*` subpackages. +# The wrapper directory `cutlass/` has no `__init__.py` (implicit PEP 420 +# namespace package), so `namespaces = true` is required for setuptools to +# pick up `cutlass.operators`. +[tool.setuptools.packages.find] +where = ["."] +include = ["cutlass", "cutlass.*"] +namespaces = true + +[tool.setuptools.dynamic] +version = {attr = "cutlass.operators.__version__"} + +[tool.pytest.ini_options] +pythonpath = ["test", "."] + +[tool.ruff] +target-version = "py310" +src = ["cutlass/operators"] +line-length = 88 + +# Exclude fusion/ from all ruff checks +extend-exclude = [ + "**/cutlass/operators/fusion", + "**/cutlass/operators/providers/cutedsl/evt", + "**/cutlass/operators/providers/cutedsl/*/implementations", +] + +[tool.ruff.format] +line-ending = "lf" + +[tool.ruff.lint] +select = [ + "F", # pyflakes + "E", # pycodestyle errors + "W", # pycodestyle warnings + "I", # isort + "UP", # pyupgrade + "B", # flake8-bugbear + "SIM", # flake8-simplify + "C4", # flake8-comprehensions + "PTH", # flake8-use-pathlib + "FA", # flake8-future-annotations + "PERF", # perflint - performance anti-patterns + "FURB", # refurb - modern Python idioms + "PIE", # flake8-pie - misc improvements + "TCH", # flake8-type-checking + "PLE", # pylint errors + "RSE", # flake8-raise + "D", # pydocstyle +] +ignore = [ + "SIM103", # nedless-bool - can be noisy + "B905", # zip(strict=) - too strict + "E501", # line too long - delegated to `ruff format` + # --- Docstring rules intentionally relaxed (mirrors the PyTorch setup) --- + "D100", # Missing docstring in public module - too noisy at module level + "D104", # Missing docstring in public package - __init__.py is mostly re-exports + "D105", # Missing docstring in magic method - __repr__ etc. are usually self-evident + "D107", # Missing docstring in __init__ - we document at the class level + "D203", # 1 blank line required before class - conflicts with D211 (ruff "incompatible-default") + "D213", # Multi-line summary on second line - conflicts with D212 (ruff "incompatible-default") +] + +[tool.ruff.lint.flake8-type-checking] +# Subclasses of these base classes evaluate their field annotations at +# runtime via `typing.get_type_hints()` (see +# `cutlass.operators.arguments.base._convert_to_internal_types`). Imports +# that are syntactically only used in annotations on those subclasses are +# therefore *not* safe to move under `if TYPE_CHECKING:` — disable TC001/TC002 +# for them so we don't trip the lint while keeping the runtime contract. +runtime-evaluated-base-classes = [ + "cutlass.operators.arguments.base.RuntimeArguments", + "cutlass.operators.arguments.base.Operand", +] + +# Functions decorated with `@cute.jit` or `@cute.kernel` have their signature +# resolved at runtime +# Imports used only in annotations of such functions are therefore *not* safe to move under +# `if TYPE_CHECKING:`, even though ruff would otherwise suggest it. +runtime-evaluated-decorators = [ + "cutlass.cute.jit", + "cutlass.cute.kernel", +] + +[tool.ruff.lint.isort] +section-order = ["future", "standard-library", "third-party", "cutlass", "first-party", "local-folder"] +# `cutlass.operators` is our own code, but the top-level `cutlass` package is +# owned by `nvidia-cutlass-dsl`. We classify the operators API as first-party +# so that `import cutlass.operators as ops` and `from cutlass.operators...` +# imports group together and after the DSL imports. +known-first-party = ["cutlass.operators"] +known-local-folder = ["test_utils"] + +[tool.ruff.lint.isort.sections] +cutlass = ["cutlass"] + +[tool.ruff.lint.per-file-ignores] +"__init__.py" = ["D"] # D: package init docstrings not required +"**/test/**" = ["C408", "D"] # C408: dict(...) kwargs are allowed; D: tests don't need docstrings +"**/examples/**" = ["D"] # D: example notebooks/scripts are tutorial prose, not API surface +"**/cutlass/operators/providers/cutedsl/**" = ["D"] # D: discovered via `get_operators(...)`, not public API +[tool.ruff.lint.pydocstyle] +convention = "google" + +[tool.ruff.lint.pycodestyle] +max-line-length = 120 diff --git a/operators/test/__init__.py b/operators/test/__init__.py new file mode 100644 index 000000000..37df7d65a --- /dev/null +++ b/operators/test/__init__.py @@ -0,0 +1,27 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/operators/test/conftest.py b/operators/test/conftest.py new file mode 100644 index 000000000..97616bfce --- /dev/null +++ b/operators/test/conftest.py @@ -0,0 +1,104 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import os +import random +import sys +from pathlib import Path + +import pytest +import torch + +import cutlass.operators as ops + +from test_utils.common import _TEST_LEVEL_ENV_VAR + + + +# The ``--test-level`` CLI option is provided by the ``test_sharding`` +# plugin registered above. This hook bridges the CLI value to the +# ``CUTLASS_OPERATORS_TEST_LEVEL`` env var. +# +# Priority: +# 1. Explicit ``--test-level`` on the command line (wins if passed; note +# ``test_sharding`` sets its own default). +# 2. Pre-existing ``CUTLASS_OPERATORS_TEST_LEVEL`` env var (manual override). +# 3. Default ``"L0"``. +def pytest_configure(config): + cli_explicit = any( + arg == "--test-level" or arg.startswith("--test-level=") for arg in sys.argv + ) + test_level = ( + config.getoption("--test-level") + if cli_explicit + else os.environ.get(_TEST_LEVEL_ENV_VAR, "L0") + ) + os.environ[_TEST_LEVEL_ENV_VAR] = test_level + config.option.test_level = test_level + + +# +# Before each test, save the GlobalOptions dict +# After each test, restore the GlobalOptions dict +# +@pytest.fixture(autouse=True) +def save_restore_global_options(): + saved = ops.GlobalOptions().save() + yield + ops.GlobalOptions().restore(saved) + + +# +# Fixtures for toggling the TVM FFI global option and forcing its enablement or disablement +# +@pytest.fixture( + params=[True, False], ids=["use_tvm_ffi=True", "use_tvm_ffi=False"], autouse=False +) +def fixture_toggle_tvm_ffi(request): + ops.GlobalOptions().use_tvm_ffi = request.param + + +@pytest.fixture(autouse=False) +def fixture_enable_tvm_ffi(): + ops.GlobalOptions().use_tvm_ffi = True + + +@pytest.fixture(autouse=False) +def fixture_disable_tvm_ffi(): + ops.GlobalOptions().use_tvm_ffi = False + + +# +# Fixture for setting random seed +# +@pytest.fixture(autouse=True) +def fixture_random_seed(): + random.seed(2025) + torch.manual_seed(2025) + + diff --git a/operators/test/integration/test_blockscaled_gemm.py b/operators/test/integration/test_blockscaled_gemm.py new file mode 100644 index 000000000..75a23b3e9 --- /dev/null +++ b/operators/test/integration/test_blockscaled_gemm.py @@ -0,0 +1,1152 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import pytest +import torch + +import cutlass.operators as ops +from cutlass.operators.metadata import Sm100DesignMetadata +from cutlass.operators.providers.cutedsl.gemm.sm100_dense_blockscaled_static_persistent import ( + PersistentDenseBlockScaledGemmOperator, +) +from cutlass.operators.utils.device import device_or_env_target_sm + +import test_utils +from test_utils import assert_close_with_reference_conversion +from test_utils.gemm import ( + make_mxfp4_kmajor, + make_random_pow2_scale, + scaled_dense_gemm_reference, +) +from test_utils.reference_check import reference_scaled_mm + +torch.manual_seed(2025) + + +@pytest.mark.arch(sms=["100f"]) +def test_mxfp8_gemm_sm100_fake_tensor(fixture_enable_tvm_ffi): + import torch._functorch.config + + torch._functorch.config.fake_tensor_allow_unsafe_data_ptr_access = False + + M, N, K, L = 256, 512, 1024, 1 + ab_dtype = torch.float8_e4m3fn + c_dtype = torch.float32 + accumulator_type = torch.float32 + scale_dtype = torch.float8_e8m0fnu + scale = ops.ScaleMode.Blockwise1x32 + swizzle = ops.ScaleSwizzleMode.Swizzle32x4x4 + + with torch._subclasses.fake_tensor.FakeTensorMode(): + A = torch.randint(-1, 2, (L, M, K), device="cuda").to(ab_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=c_dtype) + B = torch.randint(-1, 2, (L, N, K), device="cuda").to(ab_dtype).transpose(1, 2) + SFA = torch.rand( + ops.ScaledOperand.numel_scale((L, M, K), scale, swizzle), + device="cuda", + ).to(scale_dtype) + SFB = torch.rand( + ops.ScaledOperand.numel_scale((L, N, K), scale, swizzle), + device="cuda", + ).to(scale_dtype) + + fake_args = ops.GemmArguments( + A=ops.ScaledOperand( + A, + SFA, + scale, + swizzle, + ), + B=ops.ScaledOperand( + B, + SFB, + scale, + swizzle, + ), + out=D, + accumulator_type=accumulator_type, + ) + operators = ops.get_operators(fake_args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operator = operators[0] + assert operator.supports(fake_args) + compiled_artifact = operator.compile(fake_args) + + A_real = torch.randint(-1, 2, (L, M, K), device="cuda").to(ab_dtype) + B_real = torch.randint(-1, 2, (L, N, K), device="cuda").to(ab_dtype).transpose(1, 2) + D_real = torch.empty((L, M, N), device="cuda", dtype=c_dtype) + SFA_real = torch.rand( + ops.ScaledOperand.numel_scale((L, M, K), scale, swizzle), + device="cuda", + ).to(scale_dtype) + SFB_real = torch.rand( + ops.ScaledOperand.numel_scale((L, N, K), scale, swizzle), + device="cuda", + ).to(scale_dtype) + args = ops.GemmArguments( + A=ops.ScaledOperand( + A_real, + SFA_real, + scale, + swizzle, + ), + B=ops.ScaledOperand( + B_real, + SFB_real, + scale, + swizzle, + ), + out=D_real, + accumulator_type=accumulator_type, + ) + operator.run(args, compiled_artifact=compiled_artifact, assume_supported_args=True) + + # torch._scaled_mm does not support f8e5m2 * f8e5m2 currently. + # Simply skip reference check in that case (but test that a CUTLASS Operator + # is found and runs) + if ab_dtype != torch.float8_e5m2: + reference = reference_scaled_mm(A_real, B_real, SFA_real, SFB_real, c_dtype) + assert_close_with_reference_conversion(D_real, reference, D_real.dtype) + + +@pytest.mark.parametrize( + "M, N, K", + [ + (256, 512, 1024), + (1024, 640, 512), + ], +) +@pytest.mark.parametrize( + "ab_dtype, c_dtype, accumulator_type, scale_dtype", + [ + (torch.float8_e4m3fn, torch.float32, torch.float32, torch.float8_e8m0fnu), + ], +) +@pytest.mark.arch(sms=["100f"]) +def test_mxfp8_gemm_sm100_2d( + M: int, + N: int, + K: int, + ab_dtype: torch.dtype, + c_dtype: torch.dtype, + accumulator_type: torch.dtype, + scale_dtype: torch.dtype, + fixture_enable_tvm_ffi, +): + """ + Tests valid MXFP8 GEMM cases in which A, B, and out are 2D tensors. + """ + # Create torch fp8 tensors for A and B + A = torch.randint(-1, 2, (M, K), device="cuda").to(ab_dtype) + D = torch.empty((M, N), device="cuda", dtype=c_dtype) + + # Transpose B because torch._scaled_mm expects B to be column-major + B = torch.randint(-1, 2, (N, K), device="cuda").to(ab_dtype).transpose(0, 1) + + scale = ops.ScaleMode.Blockwise1x32 + swizzle = ops.ScaleSwizzleMode.Swizzle32x4x4 + + SFA = torch.rand( + ops.ScaledOperand.numel_scale((M, K), scale, swizzle), + device="cuda", + ).to(scale_dtype) + SFB = torch.rand( + ops.ScaledOperand.numel_scale((N, K), scale, swizzle), + device="cuda", + ).to(scale_dtype) + + args = ops.GemmArguments( + A=ops.ScaledOperand( + A, + SFA, + scale, + swizzle, + ), + B=ops.ScaledOperand( + B, + SFB, + scale, + swizzle, + ), + out=D, + accumulator_type=accumulator_type, + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operator = operators[0] + assert operator.supports(args) + compiled_artifact = operator.compile(args) + operator.run(args, compiled_artifact=compiled_artifact, assume_supported_args=True) + + reference = reference_scaled_mm(A, B, SFA, SFB, c_dtype) + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.arch(sms=["100f"]) +@pytest.mark.parametrize( + "M, N, K", + [ + (256, 512, 256), + (1024, 640, 512), + ], +) +@pytest.mark.parametrize( + "a_dtype, b_dtype", + [ + (torch.float4_e2m1fn_x2, torch.float8_e4m3fn), + (torch.float8_e4m3fn, torch.float4_e2m1fn_x2), + ], + ids=["mxfp4xmxfp8", "mxfp8xmxfp4"], +) +def test_mxfp4_mxfp8_mixed_gemm_sm100_2d( + a_dtype: torch.dtype, + b_dtype: torch.dtype, + M: int, + N: int, + K: int, + fixture_enable_tvm_ffi, +): + """Tests valid mixed MXFP4 x MXFP8 GEMM cases with 2D A, B, and out. + + The mixed-precision analog of :func:`test_mxfp8_gemm_sm100_2d`. torch._scaled_mm + cannot reference FP4, so the reference uses + :func:`_emulated_scaled_dense_gemm_reference` (dequantize, unpack scales, + float32 matmul). Covers both directions: MXFP4 x MXFP8 (A is FP4) and + MXFP8 x MXFP4 (B is FP4). + """ + scale = ops.ScaleMode.Blockwise1x32 + swizzle = ops.ScaleSwizzleMode.Swizzle32x4x4 + + def make_kmajor(rows: int, dtype: torch.dtype) -> torch.Tensor: + """K-major operand: FP4 via make_mxfp4_kmajor, else small-integer FP8.""" + if dtype == torch.float4_e2m1fn_x2: + return make_mxfp4_kmajor(rows, K) + return torch.randint(-1, 2, (rows, K), device="cuda").to(dtype) + + A = make_kmajor(M, a_dtype) # (M, K), K-major + B = make_kmajor(N, b_dtype).transpose(0, 1) # (K, N), K-major + + D = torch.empty((M, N), device="cuda", dtype=torch.float32) + + SFA = make_random_pow2_scale(ops.ScaledOperand.numel_scale((M, K), scale, swizzle)) + SFB = make_random_pow2_scale(ops.ScaledOperand.numel_scale((N, K), scale, swizzle)) + + args = ops.GemmArguments( + A=ops.ScaledOperand(A, SFA, scale, swizzle), + B=ops.ScaledOperand(B, SFB, scale, swizzle), + out=D, + accumulator_type=torch.float32, + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operator = operators[0] + assert operator.supports(args) + compiled_artifact = operator.compile(args) + operator.run(args, compiled_artifact=compiled_artifact, assume_supported_args=True) + + # Integer operands and power-of-two scales make the accumulation exact. + reference = scaled_dense_gemm_reference(args, ref_as_acc_dtype=True) + if reference is not None: + test_utils.reference_check.assert_close_with_reference_conversion( + D, + reference, + D.dtype, + atols=0.0, + rtols=0.0, + ) + + +@pytest.mark.arch(sms=["100f"]) +def test_mxfp8_reshaped_scale_factors(fixture_enable_tvm_ffi): + """ + Tests for valid MXFP8 GEMM cases in which scale factors have correct numel but arbitrary shape + """ + M, N, K, L = 256, 512, 1024, 1 + scale = ops.ScaleMode.Blockwise1x32 + swizzle = ops.ScaleSwizzleMode.Swizzle32x4x4 + + # Create torch fp8 tensors for A and B + A = torch.randint(-1, 2, (L, M, K), device="cuda").to(torch.float8_e4m3fn) + B = ( + torch.randint(-1, 2, (L, N, K), device="cuda") + .to(torch.float8_e4m3fn) + .transpose(1, 2) + ) + D = torch.empty((L, M, N), device="cuda", dtype=torch.float32) + + SFA = torch.rand( + ops.ScaledOperand.numel_scale((L, M, K), scale, swizzle), + device="cuda", + ).to(torch.float8_e8m0fnu) + SFB = torch.rand( + ops.ScaledOperand.numel_scale((L, N, K), scale, swizzle), + device="cuda", + ).to(torch.float8_e8m0fnu) + + args = ops.GemmArguments( + # Pass in SFA and SFB viewed as (L, -1, 2, K) + A=ops.ScaledOperand( + A, + SFA.view(L, -1, 2, K), + scale, + swizzle, + ), + B=ops.ScaledOperand( + B, + SFB.view(L, -1, 2, K), + scale, + swizzle, + ), + out=D, + accumulator_type=torch.float32, + ) + + # Only test operators for which we know there will be an error if the scale factors are invalid + def metadata_filter(metadata: ops.OperatorMetadata): + return metadata.operator_class is PersistentDenseBlockScaledGemmOperator + + operators = ops.get_operators( + args, target_sm=device_or_env_target_sm(), metadata_filter=metadata_filter + ) + assert len(operators) > 0 + operator = operators[0] + assert operator.supports(args) + compiled_artifact = operator.compile(args) + operator.run(args, compiled_artifact=compiled_artifact, assume_supported_args=True) + + def transform_sf(x: torch.Tensor) -> torch.Tensor: + """Flattens scale factor tensors after indexing into the batch dimension""" + return x.view(-1) + + reference = reference_scaled_mm(A, B, SFA, SFB, torch.float32, transform_sf) + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.parametrize( + "discovery_scale_mode, runtime_scale_mode", + [ + ((1, 1, 32), (1, 1, 32)), + ((1, 32), (1, 1, 32)), + (ops.ScaleMode.Blockwise1x32, ops.ScaleMode.Blockwise1x32), + (ops.ScaleMode.Blockwise1x32, (1, 32)), + ], +) +@pytest.mark.arch(sms=["100f"]) +def test_mxfp8_tuple_scale_mode( + discovery_scale_mode: tuple[int, ...], + runtime_scale_mode: tuple[int, ...], + fixture_enable_tvm_ffi, +): + """ + Tests for valid MXFP8 GEMM cases in which A, B, and out are 3D tensors, and + the scale mode is specified as a tuple. + """ + M, N, K, L = 256, 512, 1024, 1 + + # Create torch fp8 tensors for A and B + A = torch.randint(-1, 2, (L, M, K), device="cuda").to(torch.float8_e4m3fn) + B = ( + torch.randint(-1, 2, (L, N, K), device="cuda") + .to(torch.float8_e4m3fn) + .transpose(1, 2) + ) + D = torch.empty((L, M, N), device="cuda", dtype=torch.float32) + + swizzle = ops.ScaleSwizzleMode.Swizzle32x4x4 + SFA = torch.rand( + ops.ScaledOperand.numel_scale((L, M, K), discovery_scale_mode, swizzle), + device="cuda", + ).to(torch.float8_e8m0fnu) + SFB = torch.rand( + ops.ScaledOperand.numel_scale((L, N, K), discovery_scale_mode, swizzle), + device="cuda", + ).to(torch.float8_e8m0fnu) + + discovery_args = ops.GemmArguments( + A=ops.ScaledOperand( + A, + SFA, + discovery_scale_mode, + swizzle, + ), + B=ops.ScaledOperand( + B, + SFB, + discovery_scale_mode, + swizzle, + ), + out=D, + accumulator_type=torch.float32, + ) + + # Only test operators for which we know there will be an error if the scale factors are invalid + def metadata_filter(metadata: ops.OperatorMetadata): + return metadata.operator_class is PersistentDenseBlockScaledGemmOperator + + operators = ops.get_operators( + discovery_args, + target_sm=device_or_env_target_sm(), + metadata_filter=metadata_filter, + ) + assert len(operators) > 0 + operator = operators[0] + compiled_artifact = operator.compile(discovery_args) + + runtime_args = ops.GemmArguments( + A=ops.ScaledOperand( + A, + SFA, + runtime_scale_mode, + swizzle, + ), + B=ops.ScaledOperand( + B, + SFB, + runtime_scale_mode, + swizzle, + ), + out=D, + accumulator_type=torch.float32, + ) + assert operator.supports(runtime_args) + operator.run( + runtime_args, compiled_artifact=compiled_artifact, assume_supported_args=True + ) + + reference = reference_scaled_mm(A, B, SFA, SFB, torch.float32) + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.arch(sms=["100f"]) +def test_mxfp8_incompatible_mode_or_swizzle(fixture_enable_tvm_ffi): + """ + Tests for correct flagging of incompatible scale modes or swizzles: + - Using no swizzle mode for MXFP8 + - Using an incompatible scale mode for MXFP8 + """ + M, N, K, L = 256, 512, 1024, 1 + scale = ops.ScaleMode.Blockwise1x32 + ab_dtype = torch.float8_e4m3fn + scale_dtype = torch.float8_e8m0fnu + + # Create torch fp8 tensors for A and B + A = torch.randint(-1, 2, (L, M, K), device="cuda").to(ab_dtype) + B = torch.randint(-1, 2, (L, N, K), device="cuda").to(ab_dtype).transpose(1, 2) + D = torch.empty((L, M, N), device="cuda", dtype=torch.float32) + + SFA = torch.rand( + ops.ScaledOperand.numel_scale( + (L, M, K), scale, ops.ScaleSwizzleMode.Swizzle32x4x4 + ), + device="cuda", + ).to(scale_dtype) + SFB = torch.rand( + ops.ScaledOperand.numel_scale( + (L, N, K), scale, ops.ScaleSwizzleMode.Swizzle32x4x4 + ), + device="cuda", + ).to(scale_dtype) + + args_bad = ops.GemmArguments( + # Use incompatible swizzle mode for SFA + A=ops.ScaledOperand( + A, + SFA, + scale, + ops.ScaleSwizzleMode.SwizzleNone, + ), + B=ops.ScaledOperand( + B, + SFB, + scale, + ops.ScaleSwizzleMode.Swizzle32x4x4, + ), + out=D, + accumulator_type=torch.float32, + ) + + # Only test operators for which the swizzle mode is incompatible + def metadata_filter(metadata: ops.OperatorMetadata): + return metadata.operator_class is PersistentDenseBlockScaledGemmOperator + + operators = ops.get_operators( + args_bad, target_sm=device_or_env_target_sm(), metadata_filter=metadata_filter + ) + assert len(operators) == 0 + + # Find operator using compatible swizzle mode, but pass in arguments at runtime with + # incompatible swizzle mode + args_good = ops.GemmArguments( + A=ops.ScaledOperand( + A, + SFA, + scale, + ops.ScaleSwizzleMode.Swizzle32x4x4, + ), + B=ops.ScaledOperand( + B, + SFB, + scale, + ops.ScaleSwizzleMode.Swizzle32x4x4, + ), + out=D, + accumulator_type=torch.float32, + ) + operators = ops.get_operators( + args_good, target_sm=device_or_env_target_sm(), metadata_filter=metadata_filter + ) + assert len(operators) > 0 + operator = operators[0] + + with pytest.raises(ValueError): + operator.run(args_bad) + + # Pass in arguments with an incompatible scale mode + args_bad_scale = ops.GemmArguments( + A=ops.ScaledOperand( + A, + SFA, + scale, + ops.ScaleSwizzleMode.Swizzle32x4x4, + ), + # Bad scale mode: SFB had scale mode Blockwise1x16 + B=ops.ScaledOperand( + B, + SFB, + ops.ScaleMode.Blockwise1x16, + ops.ScaleSwizzleMode.Swizzle32x4x4, + ), + out=D, + accumulator_type=torch.float32, + ) + with pytest.raises(ValueError): + operator.run(args_bad_scale) + + +@pytest.mark.arch(sms=["100f"]) +def test_mxfp8_missing_scale_factors(fixture_enable_tvm_ffi): + """ + Tests for correct flagging of a missing scale factor (both must be supplied) + """ + M, N, K, L = 256, 512, 1024, 1 + scale = ops.ScaleMode.Blockwise1x32 + swizzle = ops.ScaleSwizzleMode.Swizzle32x4x4 + + # Create torch fp8 tensors for A and B + A = torch.randint(-1, 2, (L, M, K), device="cuda").to(torch.float8_e4m3fn) + B = ( + torch.randint(-1, 2, (L, N, K), device="cuda") + .to(torch.float8_e4m3fn) + .transpose(1, 2) + ) + D = torch.empty((L, M, N), device="cuda", dtype=torch.float32) + SFA = torch.rand( + ops.ScaledOperand.numel_scale((L, M, K), scale, swizzle), + device="cuda", + ).to(torch.float8_e8m0fnu) + + # Construct arguments with only one scale factor + args = ops.GemmArguments( + A=ops.ScaledOperand( + A, + SFA, + scale, + swizzle, + ), + B=B, + out=D, + accumulator_type=torch.float32, + ) + + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + assert len(operators) == 0 + + +@pytest.mark.arch(sms=["100f"]) +def test_mxfp8_invalid_scale_factors(fixture_enable_tvm_ffi): + """ + Tests for correct flagging of invalid scale factors: + - Too large for the A tensor + - Too large for the B tensor + """ + M, N, K, L = 256, 512, 1024, 1 + scale = ops.ScaleMode.Blockwise1x32 + swizzle = ops.ScaleSwizzleMode.Swizzle32x4x4 + ab_dtype = torch.float8_e4m3fn + scale_dtype = torch.float8_e8m0fnu + + # Create torch fp8 tensors for A and B + A = torch.randint(-1, 2, (L, M, K), device="cuda").to(ab_dtype) + B = torch.randint(-1, 2, (L, N, K), device="cuda").to(ab_dtype).transpose(1, 2) + D = torch.empty((L, M, N), device="cuda", dtype=torch.float32) + + # Pad SFA past its expected element count. This makes it too large for the A tensor. + SFA = torch.rand( + ops.ScaledOperand.numel_scale((L, M, K), scale, swizzle) + 32, + device="cuda", + ).to(scale_dtype) + SFB = torch.rand( + ops.ScaledOperand.numel_scale((L, N, K), scale, swizzle), + device="cuda", + ).to(scale_dtype) + + args = ops.GemmArguments( + A=ops.ScaledOperand( + A, + SFA, + scale, + swizzle, + ), + B=ops.ScaledOperand( + B, + SFB, + scale, + swizzle, + ), + out=D, + accumulator_type=torch.float32, + ) + + # Only test operators for which we know there will be an error if the scale factors are invalid + def metadata_filter(metadata: ops.OperatorMetadata): + return metadata.operator_class is PersistentDenseBlockScaledGemmOperator + + operators = ops.get_operators( + args, target_sm=device_or_env_target_sm(), metadata_filter=metadata_filter + ) + assert len(operators) == 0 + + # Pad SFB past its expected element count. This makes it too large for the B tensor. + SFA = torch.rand( + ops.ScaledOperand.numel_scale((L, M, K), scale, swizzle), + device="cuda", + ).to(scale_dtype) + SFB = torch.rand( + ops.ScaledOperand.numel_scale((L, N, K), scale, swizzle) + 32, + device="cuda", + ).to(scale_dtype) + args = ops.GemmArguments( + A=ops.ScaledOperand( + A, + SFA, + scale, + swizzle, + ), + B=ops.ScaledOperand( + B, + SFB, + scale, + swizzle, + ), + out=D, + accumulator_type=torch.float32, + ) + operators = ops.get_operators( + args, target_sm=device_or_env_target_sm(), metadata_filter=metadata_filter + ) + assert len(operators) == 0 + + +@pytest.mark.arch(sms=["100f"]) +def test_mxfp8_gemm_sm100_invalid_epilogue(fixture_enable_tvm_ffi): + """ + Tests for correct flagging of invalid epilogues: currently, no MXFP8 operators support + epilogue fusions. + """ + M, N, K, L = 256, 512, 1024, 1 + + # Create torch fp8 tensors for A and B + A = torch.randint(-1, 2, (L, M, K), device="cuda").to(torch.float8_e4m3fn) + B = ( + torch.randint(-1, 2, (L, N, K), device="cuda") + .to(torch.float8_e4m3fn) + .transpose(1, 2) + ) + D = torch.empty((L, M, N), device="cuda", dtype=torch.float32) + + scale_dtype = torch.float8_e8m0fnu + scale = ops.ScaleMode.Blockwise1x32 + swizzle = ops.ScaleSwizzleMode.Swizzle32x4x4 + SFA = torch.rand( + ops.ScaledOperand.numel_scale((L, M, K), scale, swizzle), + device="cuda", + ).to(scale_dtype) + SFB = torch.rand( + ops.ScaledOperand.numel_scale((L, N, K), scale, swizzle), + device="cuda", + ).to(scale_dtype) + + C = torch.randint(-1, 2, (L, M, N), device="cuda").to(torch.float32) + + def epilogue(accum, alpha, beta, C): + D = alpha * accum + beta * C + return D + + epi_args = ops.EpilogueArguments(epilogue, alpha=1.0, beta=1.0, C=C, D=D) + + args = ops.GemmArguments( + A=ops.ScaledOperand( + A, + SFA, + scale, + swizzle, + ), + B=ops.ScaledOperand( + B, + SFB, + scale, + swizzle, + ), + out=D, + accumulator_type=torch.float32, + epilogue=epi_args, + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) == 0 + + +@pytest.mark.arch(sms=["100f"]) +def test_mxfp8_gemm_sm100_design_metadata(fixture_enable_tvm_ffi): + """ + Tests for correct filtering of operators based on design metadata specifications + (e.g., tile shape, cluster shape, etc.). + """ + M, N, K, L = 256, 512, 1024, 1 + + # Create torch fp8 tensors for A and B + A = torch.randint(-1, 2, (L, M, K), device="cuda").to(torch.float8_e4m3fn) + B = ( + torch.randint(-1, 2, (L, N, K), device="cuda") + .to(torch.float8_e4m3fn) + .transpose(1, 2) + ) + D = torch.empty((L, M, N), device="cuda", dtype=torch.float32) + + scale_dtype = torch.float8_e8m0fnu + scale = ops.ScaleMode.Blockwise1x32 + swizzle = ops.ScaleSwizzleMode.Swizzle32x4x4 + SFA = torch.rand( + ops.ScaledOperand.numel_scale((L, M, K), scale, swizzle), + device="cuda", + ).to(scale_dtype) + SFB = torch.rand( + ops.ScaledOperand.numel_scale((L, N, K), scale, swizzle), + device="cuda", + ).to(scale_dtype) + + args = ops.GemmArguments( + A=ops.ScaledOperand( + A, + SFA, + scale, + swizzle, + ), + B=ops.ScaledOperand( + B, + SFB, + scale, + swizzle, + ), + out=D, + accumulator_type=torch.float32, + ) + + def design_filter(metadata: ops.OperatorMetadata): + if not isinstance(metadata.design, Sm100DesignMetadata): + return False + return metadata.design.tile_shape[:2] == (256, 128) + + operators = ops.get_operators( + args, target_sm=device_or_env_target_sm(), metadata_filter=design_filter + ) + assert len(operators) > 0 + + for operator in operators: + assert design_filter(operator.metadata) + + operator = operators[0] + assert operator.supports(args) + compiled_artifact = operator.compile(args) + operator.run(args, compiled_artifact=compiled_artifact) + + reference = reference_scaled_mm(A, B, SFA, SFB, torch.float32) + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.parametrize( + "M, N, K", + [ + (256, 512, 512), + (512, 256, 512), + ], +) +@pytest.mark.arch(sms=["100f"]) +def test_nvfp4_gemm_sm100( + M: int, + N: int, + K: int, + fixture_enable_tvm_ffi, +): + """ + Tests NVFP4 GEMM using the PersistentDenseBlockScaledGemmOperator. + Float4E2M1FN requires K-major layout for both A and B. + PyTorch stores FP4 as float4_e2m1fn_x2 (2 values packed per byte). + """ + L = 1 + c_dtype = torch.float16 + accumulator_type = torch.float32 + scale_dtype = torch.float8_e4m3fn + scale_mode = ops.ScaleMode.Blockwise1x16 + + # FP4 tensors: K is the number of FP4 elements; PyTorch stores K//2 packed pairs + K_packed = K // 2 + A = torch.randint(0, 2, (L, M, K_packed), dtype=torch.uint8, device="cuda").view( + torch.float4_e2m1fn_x2 + ) + B = ( + torch.randint(0, 2, (L, N, K_packed), dtype=torch.uint8, device="cuda") + .view(torch.float4_e2m1fn_x2) + .transpose(1, 2) + ) + D = torch.empty((L, M, N), device="cuda", dtype=c_dtype) + + swizzle = ops.ScaleSwizzleMode.Swizzle32x4x4 + + SFA = torch.rand( + ops.ScaledOperand.numel_scale((L, M, K), scale_mode, swizzle), + device="cuda", + ).to(scale_dtype) + SFB = torch.rand( + ops.ScaledOperand.numel_scale((L, N, K), scale_mode, swizzle), + device="cuda", + ).to(scale_dtype) + + args = ops.GemmArguments( + A=ops.ScaledOperand(A, SFA, scale_mode, swizzle), + B=ops.ScaledOperand(B, SFB, scale_mode, swizzle), + out=D, + accumulator_type=accumulator_type, + ) + + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + assert len(operators) > 0, f"No NVFP4 kernels found for M={M}, N={N}, K={K}" + operator = operators[0] + assert operator.supports(args) + compiled_artifact = operator.compile(args) + operator.run(args, compiled_artifact=compiled_artifact, assume_supported_args=True) + + reference = scaled_dense_gemm_reference(args, ref_as_acc_dtype=True) + test_utils.reference_check.assert_close_with_reference_conversion( + D, reference, D.dtype, atols=1e-01, rtols=1e-02, + ) + + +@pytest.mark.arch(sms=["100f"]) +def test_nvfp4_gemm_sm100_batched(fixture_enable_tvm_ffi): + """ + Tests batched NVFP4 GEMM (L > 1) with Float4E2M1FN inputs and + Float8E4M3FN scale factors using Blockwise1x16 scale mode. + """ + M, N, K, L = 256, 512, 512, 2 + c_dtype = torch.float16 + accumulator_type = torch.float32 + scale_dtype = torch.float8_e4m3fn + scale_mode = ops.ScaleMode.Blockwise1x16 + + K_packed = K // 2 + A = torch.randint(0, 2, (L, M, K_packed), dtype=torch.uint8, device="cuda").view( + torch.float4_e2m1fn_x2 + ) + B = ( + torch.randint(0, 2, (L, N, K_packed), dtype=torch.uint8, device="cuda") + .view(torch.float4_e2m1fn_x2) + .transpose(1, 2) + ) + D = torch.empty((L, M, N), device="cuda", dtype=c_dtype) + + swizzle = ops.ScaleSwizzleMode.Swizzle32x4x4 + + SFA = torch.rand( + ops.ScaledOperand.numel_scale((L, M, K), scale_mode, swizzle), + device="cuda", + ).to(scale_dtype) + SFB = torch.rand( + ops.ScaledOperand.numel_scale((L, N, K), scale_mode, swizzle), + device="cuda", + ).to(scale_dtype) + + args = ops.GemmArguments( + A=ops.ScaledOperand(A, SFA, scale_mode, swizzle), + B=ops.ScaledOperand(B, SFB, scale_mode, swizzle), + out=D, + accumulator_type=accumulator_type, + ) + + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + assert len(operators) > 0 + operator = operators[0] + assert operator.supports(args) + compiled_artifact = operator.compile(args) + operator.run(args, compiled_artifact=compiled_artifact, assume_supported_args=True) + + reference = scaled_dense_gemm_reference(args, ref_as_acc_dtype=True) + test_utils.reference_check.assert_close_with_reference_conversion( + D, reference, D.dtype, atols=1e-01, rtols=1e-02, + ) + + +@pytest.mark.parametrize( + "M, N, K, scale_mode", + [ + (256, 512, 512, ops.ScaleMode.Blockwise1x16), + (512, 256, 512, ops.ScaleMode.Blockwise1x32), + ], +) +@pytest.mark.arch(sms=["100f"]) +def test_nvfp4_gemm_sm100_e8m0_scales( + M: int, + N: int, + K: int, + scale_mode: ops.ScaleMode, + fixture_enable_tvm_ffi, +): + """Tests NVFP4 GEMM with Float8E8M0FNU scale factors. + + torch._scaled_mm does not support FP4 + E8M0FNU, so this exercises the + emulated reference path (dequantize + unpack scales + float32 matmul). + """ + L = 1 + c_dtype = torch.float16 + accumulator_type = torch.float32 + scale_dtype = torch.float8_e8m0fnu + swizzle = ops.ScaleSwizzleMode.Swizzle32x4x4 + + K_packed = K // 2 + _FP4_BYTE_PATTERNS = torch.tensor( + [0x00, 0x02, 0x0A, 0x20, 0x22, 0x2A, 0xA0, 0xA2, 0xAA], + dtype=torch.uint8, + device="cuda", + ) + A = _FP4_BYTE_PATTERNS[ + torch.randint(0, len(_FP4_BYTE_PATTERNS), (L, M, K_packed), device="cuda") + ].view(torch.float4_e2m1fn_x2) + B = ( + _FP4_BYTE_PATTERNS[ + torch.randint(0, len(_FP4_BYTE_PATTERNS), (L, N, K_packed), device="cuda") + ] + .view(torch.float4_e2m1fn_x2) + .transpose(1, 2) + ) + D = torch.empty((L, M, N), device="cuda", dtype=c_dtype) + + SFA = torch.ones( + ops.ScaledOperand.numel_scale((L, M, K), scale_mode, swizzle), + device="cuda", + ).to(scale_dtype) + SFB = torch.ones( + ops.ScaledOperand.numel_scale((L, N, K), scale_mode, swizzle), + device="cuda", + ).to(scale_dtype) + + args = ops.GemmArguments( + A=ops.ScaledOperand(A, SFA, scale_mode, swizzle), + B=ops.ScaledOperand(B, SFB, scale_mode, swizzle), + out=D, + accumulator_type=accumulator_type, + ) + + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + assert len(operators) > 0, f"No NVFP4 kernels found for M={M}, N={N}, K={K}" + operator = operators[0] + assert operator.supports(args) + compiled_artifact = operator.compile(args) + operator.run(args, compiled_artifact=compiled_artifact, assume_supported_args=True) + + reference = scaled_dense_gemm_reference(args, ref_as_acc_dtype=True) + test_utils.reference_check.assert_close_with_reference_conversion( + D, reference, D.dtype, atols=1e-01, rtols=1e-02, + ) + + +def _make_nvfp4_tensors(M, N, K, L=1): + """Create standard NVFP4 tensors (K-major A/B, output D, scale factors).""" + K_packed = K // 2 + scale_mode = ops.ScaleMode.Blockwise1x16 + swizzle = ops.ScaleSwizzleMode.Swizzle32x4x4 + scale_dtype = torch.float8_e4m3fn + + A = torch.randint(0, 2, (L, M, K_packed), dtype=torch.uint8, device="cuda").view( + torch.float4_e2m1fn_x2 + ) + B = ( + torch.randint(0, 2, (L, N, K_packed), dtype=torch.uint8, device="cuda") + .view(torch.float4_e2m1fn_x2) + .transpose(1, 2) + ) + D = torch.empty((L, M, N), device="cuda", dtype=torch.float16) + + SFA = torch.rand( + ops.ScaledOperand.numel_scale((L, M, K), scale_mode, swizzle), + device="cuda", + ).to(scale_dtype) + SFB = torch.rand( + ops.ScaledOperand.numel_scale((L, N, K), scale_mode, swizzle), + device="cuda", + ).to(scale_dtype) + + return A, B, D, SFA, SFB, scale_mode + + +@pytest.mark.arch(sms=["100f"]) +def test_nvfp4_no_kernels_for_mn_major(fixture_enable_tvm_ffi): + """ + Tests that no kernels are found when FP4 operands use M-major layout, + since Float4E2M1FN only supports K-major layout for both A and B. + """ + M, N, K, L = 256, 512, 512, 1 + scale_mode = ops.ScaleMode.Blockwise1x16 + scale_dtype = torch.float8_e4m3fn + + K_packed = K // 2 + A = ( + torch.randint(0, 2, (L, K_packed, M), dtype=torch.uint8, device="cuda") + .view(torch.float4_e2m1fn_x2) + .transpose(1, 2) + ) + B = ( + torch.randint(0, 2, (L, N, K_packed), dtype=torch.uint8, device="cuda") + .view(torch.float4_e2m1fn_x2) + .transpose(1, 2) + ) + D = torch.empty((L, M, N), device="cuda", dtype=torch.float16) + + swizzle = ops.ScaleSwizzleMode.Swizzle32x4x4 + + SFA = torch.rand( + ops.ScaledOperand.numel_scale((L, M, K), scale_mode, swizzle), + device="cuda", + ).to(scale_dtype) + SFB = torch.rand( + ops.ScaledOperand.numel_scale((L, N, K), scale_mode, swizzle), + device="cuda", + ).to(scale_dtype) + + args = ops.GemmArguments( + A=ops.ScaledOperand(A, SFA, scale_mode, swizzle), + B=ops.ScaledOperand(B, SFB, scale_mode, swizzle), + out=D, + accumulator_type=torch.float32, + ) + + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + assert len(operators) == 0 + + +@pytest.mark.arch(sms=["100f"]) +def test_nvfp4_invalid_scale_factors(fixture_enable_tvm_ffi): + """ + Tests for correct flagging of invalid scale factors for NVFP4: + - SFA too large for the A tensor + - SFB too large for the B tensor + """ + M, N, K, L = 256, 512, 512, 1 + A, B, D, _, _, scale_mode = _make_nvfp4_tensors(M, N, K, L) + scale_dtype = torch.float8_e4m3fn + swizzle = ops.ScaleSwizzleMode.Swizzle32x4x4 + correct_sfa_numel = ops.ScaledOperand.numel_scale((L, M, K), scale_mode, swizzle) + correct_sfb_numel = ops.ScaledOperand.numel_scale((L, N, K), scale_mode, swizzle) + + # SFA with extra elements — too large for A + SFA_bad = torch.rand(correct_sfa_numel + 32, device="cuda").to(scale_dtype) + SFB_good = torch.rand(correct_sfb_numel, device="cuda").to(scale_dtype) + + args = ops.GemmArguments( + A=ops.ScaledOperand(A, SFA_bad, scale_mode, swizzle), + B=ops.ScaledOperand(B, SFB_good, scale_mode, swizzle), + out=D, + accumulator_type=torch.float32, + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + assert len(operators) == 0 + + # SFB with extra elements — too large for B + SFA_good = torch.rand(correct_sfa_numel, device="cuda").to(scale_dtype) + SFB_bad = torch.rand(correct_sfb_numel + 32, device="cuda").to(scale_dtype) + + args = ops.GemmArguments( + A=ops.ScaledOperand(A, SFA_good, scale_mode, swizzle), + B=ops.ScaledOperand(B, SFB_bad, scale_mode, swizzle), + out=D, + accumulator_type=torch.float32, + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + assert len(operators) == 0 + + +@pytest.mark.arch(sms=["100f"]) +def test_nvfp4_missing_scale_factors(fixture_enable_tvm_ffi): + """ + Tests that no blockscaled kernels are found when B is passed as a plain + DenseTensor (without scale factors) while A has scale factors. + """ + M, N, K, L = 256, 512, 512, 1 + A, B, D, SFA, _, scale_mode = _make_nvfp4_tensors(M, N, K, L) + + args = ops.GemmArguments( + A=ops.ScaledOperand(A, SFA, scale_mode, ops.ScaleSwizzleMode.Swizzle32x4x4), + B=B, + out=D, + accumulator_type=torch.float32, + ) + + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + assert len(operators) == 0 + + +@pytest.mark.arch(sms=["100f"]) +def test_nvfp4_incompatible_swizzle(fixture_enable_tvm_ffi): + """ + Tests that no kernels are found when NVFP4 operands use SwizzleNone + instead of the required Swizzle32x4x4. + """ + M, N, K, L = 256, 512, 512, 1 + A, B, D, SFA, SFB, scale_mode = _make_nvfp4_tensors(M, N, K, L) + + args = ops.GemmArguments( + A=ops.ScaledOperand(A, SFA, scale_mode, ops.ScaleSwizzleMode.SwizzleNone), + B=ops.ScaledOperand(B, SFB, scale_mode, ops.ScaleSwizzleMode.Swizzle32x4x4), + out=D, + accumulator_type=torch.float32, + ) + + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + assert len(operators) == 0 diff --git a/operators/test/integration/test_contiguous_offset_dense_gemm.py b/operators/test/integration/test_contiguous_offset_dense_gemm.py new file mode 100644 index 000000000..db1c44da2 --- /dev/null +++ b/operators/test/integration/test_contiguous_offset_dense_gemm.py @@ -0,0 +1,128 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import random + +import pytest +import torch + +import cutlass.operators as ops +from cutlass.operators.utils.device import device_or_env_supports + +from test_utils import assert_close_with_reference_conversion + +torch.manual_seed(2025) +random.seed(2025) + + +def test_incorrect_offset_length(): + """ + Offset tensors are required to have `problem_count` elements. + + Test that no operators are found when this is violated. + """ + + problem_count, m, n, k = 12, 8192, 128, 512 + + A = torch.empty((1, m, k), device="cuda", dtype=torch.float16) + B = torch.empty((problem_count, n, k), device="cuda", dtype=torch.float16).permute( + 0, 2, 1 + ) + out = torch.empty((1, m, n), device="cuda", dtype=torch.float32) + + # Incorrect: should have `problem_count` elements + offsets = torch.empty((problem_count + 1,), device="cuda", dtype=torch.int32) + + args = ops.GroupedGemmArguments( + A=A, + B=B, + out=out, + accumulator_type=torch.float32, + offsets=offsets, + ) + + operators = ops.get_operators(args, target_sm="100a") + assert len(operators) == 0 + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_contiguous_offset_dense_gemm_2d3d_fake_tensor(fixture_toggle_tvm_ffi): + import torch._functorch.config + + torch._functorch.config.fake_tensor_allow_unsafe_data_ptr_access = False + + M, N, K, L = 256, 512, 128, 2 + ab_dtype = torch.float16 + c_dtype = torch.float16 + accumulator_type = torch.float32 + + with torch._subclasses.fake_tensor.FakeTensorMode(): + A = torch.randint(-1, 2, (M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, N, K), device="cuda", dtype=ab_dtype).permute( + 0, 2, 1 + ) + out = torch.empty((M, N), device="cuda", dtype=c_dtype) + offsets = torch.empty((L,), device="cuda", dtype=torch.int32) + + fake_args = ops.GroupedGemmArguments( + A=A, B=B, out=out, accumulator_type=accumulator_type, offsets=offsets + ) + operators = ops.get_operators(fake_args, target_sm="100a") + + assert len(operators) > 0 + operator = operators[0] + compiled_artifact = operator.compile(fake_args) + + A_real = torch.randint(-1, 2, (M, K), device="cuda", dtype=ab_dtype) + B_real = torch.randint(-1, 2, (L, N, K), device="cuda", dtype=ab_dtype).permute( + 0, 2, 1 + ) + out_real = torch.empty((M, N), device="cuda", dtype=c_dtype) + offsets_real = torch.Tensor([128, 256]).to("cuda").to(torch.int32) + args = ops.GroupedGemmArguments( + A=A_real, + B=B_real, + out=out_real, + accumulator_type=accumulator_type, + offsets=offsets_real, + ) + operator.run(args, compiled_artifact=compiled_artifact) + + reference = torch._grouped_mm(A_real, B_real, offsets_real) + + # Move to CPU for comparison + out_real = out_real.cpu() + reference = reference.cpu() + assert_close_with_reference_conversion( + out_real, + reference, + out_real.dtype, + ) diff --git a/operators/test/integration/test_cuda_graph.py b/operators/test/integration/test_cuda_graph.py new file mode 100644 index 000000000..189e64cb0 --- /dev/null +++ b/operators/test/integration/test_cuda_graph.py @@ -0,0 +1,164 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +import pytest +import torch +from torch.cuda import current_stream + +import cutlass.operators as ops +from cutlass.operators.utils.device import device_or_env_supports + +from test_utils import assert_close_with_reference_conversion + + +@pytest.mark.parametrize( + "M, N, K", + [ + (256, 512, 1024), + ], +) +@pytest.mark.parametrize( + "ab_dtype, c_dtype, accumulator_type", + [ + (torch.float16, torch.float16, torch.float16), + ], +) +@pytest.mark.parametrize( + "n_iterations", + [ + 20, + ], +) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_sm100( + M: int, + N: int, + K: int, + ab_dtype: torch.dtype, + c_dtype: torch.dtype, + accumulator_type: torch.dtype, + n_iterations: int, + fixture_toggle_tvm_ffi, +): + A = torch.randint(-1, 2, (M, K), device="cuda").to(ab_dtype) + B = torch.randint(-1, 2, (K, N), device="cuda").to(ab_dtype) + D = torch.randint(-1, 2, (M, N), device="cuda").to(c_dtype) + + args = ops.GemmArguments(A=A, B=B, out=D, accumulator_type=accumulator_type) + operators = ops.get_operators(args, target_sm="100a") + + assert len(operators) > 0 + operator = operators[0] + + """ + Compile the Operator and capture CUDA Graph. + The Operator needs to be compiled outside the CUDA graph. + """ + + assert operator.supports(args) + compiled_artifact = operator.compile(args) + stream = torch.cuda.Stream() + + # Create a CUDA Graph to run our compiled Operator N times + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + # Run N iterations of our compiled Operator on the current stream + for _ in range(n_iterations): + operator.run( + args, + compiled_artifact=compiled_artifact, + stream=current_stream(), + assume_supported_args=True, + ) + + # Zero the output so we don't refcheck stale results + D.zero_() + + # Replay captured graph & check result + g.replay() + torch.cuda.synchronize() + + reference = A @ B + assert_close_with_reference_conversion(D, reference, D.dtype) + + """ + Run with & without graph capture to compare overhead + """ + + # Create CUDA events for measuring performance + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + + # Warmup the GPU + for _ in range(n_iterations): + operator.run( + args, + compiled_artifact=compiled_artifact, + stream=stream, + assume_supported_args=True, + ) + + # Run without CUDA graph and time it + start.record() + for _ in range(n_iterations): + operator.run( + args, + compiled_artifact=compiled_artifact, + stream=stream, + assume_supported_args=True, + ) + end.record() + torch.cuda.synchronize() + without_graph_time = start.elapsed_time(end) + + # Warmup again + for _ in range(n_iterations): + operator.run( + args, + compiled_artifact=compiled_artifact, + stream=stream, + assume_supported_args=True, + ) + + # Run with CUDA graph and time it + start.record() + g.replay() + end.record() + torch.cuda.synchronize() + with_graph_time = start.elapsed_time(end) + + percent_speedup = (without_graph_time - with_graph_time) / with_graph_time + print("-" * 80) + print(f"Number of launches : {n_iterations}") + print(f"Time without CUDA graph: {without_graph_time:.2f} ms") + print(f"Time with CUDA graph: {with_graph_time:.2f} ms") + print(f"Speedup : {percent_speedup * 100.0:.2f}%") diff --git a/operators/test/integration/test_cutlass_tensor.py b/operators/test/integration/test_cutlass_tensor.py new file mode 100644 index 000000000..731c223c9 --- /dev/null +++ b/operators/test/integration/test_cutlass_tensor.py @@ -0,0 +1,302 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import pytest +import torch + +import cutlass +from cutlass.cute.runtime import from_dlpack + +import cutlass.operators as ops +from cutlass.operators.utils.device import ( + device_or_env_supports, + device_or_env_target_sm, +) +from cutlass.operators.utils.dtype import to_cutlass_type +from cutlass.operators.utils.tensor import get_stride_order, major_mode + +from test_utils.reference_check import ( + assert_close_with_reference_conversion, + reference_scaled_mm, +) + +""" +Tests that cute.Tensor can be passed directly as RuntimeArguments to an Operator. +""" + + +@pytest.mark.arch(sms=["100f"]) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="sm80_tensorop_gemm_impl has an ICE because it attempts to compile using sm90+ cluster launch instructions", +) +def test_gemm_args_for_cute_tensor(fixture_toggle_tvm_ffi): + M, N, K, L = 256, 256, 256, 1 + + with torch.device("cuda"): + A_torch = torch.randint(-1, 2, (L, M, K), dtype=torch.float16) + B_torch = torch.randint(-1, 2, (L, K, N), dtype=torch.float16) + out_torch = torch.empty((L, M, N), dtype=torch.float16) + reference_output = A_torch @ B_torch + + A = ( + from_dlpack( + A_torch, assumed_align=16, enable_tvm_ffi=ops.GlobalOptions().use_tvm_ffi + ) + .mark_layout_dynamic(major_mode(A_torch.shape, A_torch.stride())) + .mark_compact_shape_dynamic( + mode=major_mode(A_torch.shape, A_torch.stride()), + divisibility=16 * 8 // to_cutlass_type(A_torch.dtype).width, + stride_order=get_stride_order(A_torch.stride()), + ) + ) + B = ( + from_dlpack( + B_torch, assumed_align=16, enable_tvm_ffi=ops.GlobalOptions().use_tvm_ffi + ) + .mark_layout_dynamic(major_mode(B_torch.shape, B_torch.stride())) + .mark_compact_shape_dynamic( + mode=major_mode(B_torch.shape, B_torch.stride()), + divisibility=16 * 8 // to_cutlass_type(B_torch.dtype).width, + stride_order=get_stride_order(B_torch.stride()), + ) + ) + out = ( + from_dlpack( + out_torch, assumed_align=16, enable_tvm_ffi=ops.GlobalOptions().use_tvm_ffi + ) + .mark_layout_dynamic(major_mode(out_torch.shape, out_torch.stride())) + .mark_compact_shape_dynamic( + mode=major_mode(out_torch.shape, out_torch.stride()), + divisibility=16 * 8 // to_cutlass_type(out_torch.dtype).width, + stride_order=get_stride_order(out_torch.stride()), + ) + ) + + accumulator_type = cutlass.Float32 + + args = ops.GemmArguments(A=A, B=B, out=out, accumulator_type=accumulator_type) + + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + assert len(operators) > 0 + operator = operators[0] + operator.run(args) + assert_close_with_reference_conversion( + args.out.tensor.cuda(), reference_output, args.out.dtype + ) + + +@pytest.mark.arch(sms=["100a"]) +@pytest.mark.skipif( + not device_or_env_supports("100a"), + reason="from_dlpack for scale factors on int: + return round_up(ceil_div(k, ops.ScaleMode.numel(scale)), 4) + + M, N, K, L = 1024, 1024, 1024, 1 + + scale_mode = ops.ScaleMode.Blockwise1x32 + torch_scale_dtype = torch.float8_e8m0fnu + + with torch.device("cuda"): + A_torch = torch.randint(-1, 2, (L, M, K)).to(torch.float8_e4m3fn) + scale_A_torch = torch.rand((L, round_up(M, 128), prep_k(K, scale_mode))).to( + torch_scale_dtype + ) + B_torch = ( + torch.randint(-1, 2, (L, K, N)).to(torch.float8_e4m3fn).transpose(1, 2) + ) + scale_B_torch = torch.rand((L, prep_k(K, scale_mode), round_up(N, 128))).to( + torch_scale_dtype + ) + out_torch = torch.empty((L, M, N), dtype=torch.float32) + + reference_output = reference_scaled_mm( + A_torch, B_torch, scale_A_torch, scale_B_torch, out_torch.dtype + ) + scale_dtype = to_cutlass_type(torch_scale_dtype) + + A = ( + from_dlpack( + A_torch, assumed_align=16, enable_tvm_ffi=ops.GlobalOptions().use_tvm_ffi + ) + .mark_layout_dynamic(major_mode(A_torch.shape, A_torch.stride())) + .mark_compact_shape_dynamic( + mode=major_mode(A_torch.shape, A_torch.stride()), + divisibility=16 * 8 // to_cutlass_type(A_torch.dtype).width, + stride_order=get_stride_order(A_torch.stride()), + ) + ) + + scale_A = ( + from_dlpack( + scale_A_torch, + assumed_align=32, + enable_tvm_ffi=ops.GlobalOptions().use_tvm_ffi, + ) + .mark_layout_dynamic(major_mode(scale_A_torch.shape, scale_A_torch.stride())) + .mark_compact_shape_dynamic( + mode=major_mode(scale_A_torch.shape, scale_A_torch.stride()), + divisibility=32 * 8 // scale_dtype.width, + stride_order=get_stride_order(scale_A_torch.stride()), + ) + ) + B = ( + from_dlpack( + B_torch, assumed_align=16, enable_tvm_ffi=ops.GlobalOptions().use_tvm_ffi + ) + .mark_layout_dynamic(major_mode(B_torch.shape, B_torch.stride())) + .mark_compact_shape_dynamic( + mode=major_mode(B_torch.shape, B_torch.stride()), + divisibility=16 * 8 // to_cutlass_type(B_torch.dtype).width, + stride_order=get_stride_order(B_torch.stride()), + ) + ) + scale_B = ( + from_dlpack( + scale_B_torch, + assumed_align=32, + enable_tvm_ffi=ops.GlobalOptions().use_tvm_ffi, + ) + .mark_layout_dynamic(major_mode(scale_B_torch.shape, scale_B_torch.stride())) + .mark_compact_shape_dynamic( + mode=major_mode(scale_B_torch.shape, scale_B_torch.stride()), + divisibility=32 * 8 // scale_dtype.width, + stride_order=get_stride_order(scale_B_torch.stride()), + ) + ) + out = ( + from_dlpack( + out_torch, assumed_align=16, enable_tvm_ffi=ops.GlobalOptions().use_tvm_ffi + ) + .mark_layout_dynamic(major_mode(out_torch.shape, out_torch.stride())) + .mark_compact_shape_dynamic( + mode=major_mode(out_torch.shape, out_torch.stride()), + divisibility=16 * 8 // to_cutlass_type(out_torch.dtype).width, + stride_order=get_stride_order(out_torch.stride()), + ) + ) + + accumulator_type = cutlass.Float32 + + args = ops.GemmArguments( + A=ops.ScaledOperand(A, scale_A, scale_mode, ops.ScaleSwizzleMode.Swizzle32x4x4), + B=ops.ScaledOperand(B, scale_B, scale_mode, ops.ScaleSwizzleMode.Swizzle32x4x4), + out=out, + accumulator_type=accumulator_type, + ) + + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + assert len(operators) > 0 + operator = operators[0] + operator.run(args) + + out_dtype = args.out.dtype + assert_close_with_reference_conversion( + args.out.tensor.cuda(), reference_output, out_dtype + ) + + +@pytest.mark.arch(sms=["100f"]) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Library only supports Grouped Gemm for >=sm100", +) +def test_grouped_gemm_args_for_cute_tensor(fixture_toggle_tvm_ffi): + groups = 2 + M = 256 + TotalM = groups * M + M_offsets = [M, TotalM] + N, K = 256, 256 + + with torch.device("cuda"): + A_torch = torch.randint(-1, 2, (TotalM, K), dtype=torch.float16) + B_torch = torch.randint(-1, 2, (groups, N, K), dtype=torch.float16).permute( + 0, 2, 1 + ) + out_torch = torch.empty((TotalM, N), dtype=torch.float16) + offsets_torch = torch.tensor(M_offsets, dtype=torch.int32) + + reference_output = torch._grouped_mm(A_torch, B_torch, offsets_torch) + + A = ( + from_dlpack( + A_torch, assumed_align=16, enable_tvm_ffi=ops.GlobalOptions().use_tvm_ffi + ) + .mark_layout_dynamic(major_mode(A_torch.shape, A_torch.stride())) + .mark_compact_shape_dynamic( + mode=major_mode(A_torch.shape, A_torch.stride()), + divisibility=16 * 8 // to_cutlass_type(A_torch.dtype).width, + stride_order=get_stride_order(A_torch.stride()), + ) + ) + B = ( + from_dlpack( + B_torch, assumed_align=16, enable_tvm_ffi=ops.GlobalOptions().use_tvm_ffi + ) + .mark_layout_dynamic(major_mode(B_torch.shape, B_torch.stride())) + .mark_compact_shape_dynamic( + mode=major_mode(B_torch.shape, B_torch.stride()), + divisibility=16 * 8 // to_cutlass_type(B_torch.dtype).width, + stride_order=get_stride_order(B_torch.stride()), + ) + ) + out = ( + from_dlpack( + out_torch, assumed_align=16, enable_tvm_ffi=ops.GlobalOptions().use_tvm_ffi + ) + .mark_layout_dynamic(major_mode(out_torch.shape, out_torch.stride())) + .mark_compact_shape_dynamic( + mode=major_mode(out_torch.shape, out_torch.stride()), + divisibility=16 * 8 // to_cutlass_type(out_torch.dtype).width, + stride_order=get_stride_order(out_torch.stride()), + ) + ) + offsets = from_dlpack( + offsets_torch, assumed_align=16, enable_tvm_ffi=ops.GlobalOptions().use_tvm_ffi + ).mark_layout_dynamic(major_mode(offsets_torch.shape, offsets_torch.stride())) + + accumulator_type = cutlass.Float32 + + args = ops.GroupedGemmArguments( + A=A, B=B, out=out, offsets=offsets, accumulator_type=accumulator_type + ) + + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + assert len(operators) > 0 + operator = operators[0] + operator.run(args) + assert_close_with_reference_conversion( + args.out.tensor.cuda(), reference_output, args.out.dtype + ) + + diff --git a/operators/test/integration/test_gemm.py b/operators/test/integration/test_gemm.py new file mode 100644 index 000000000..a33f76992 --- /dev/null +++ b/operators/test/integration/test_gemm.py @@ -0,0 +1,206 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +import logging +from pprint import pformat + +import pytest +import torch + +import cutlass + +import cutlass.operators as ops +from cutlass.operators.metadata import ( + Sm100DesignMetadata, +) +from cutlass.operators.utils.device import device_or_env_target_sm + +from test_utils import assert_close_with_reference_conversion + +torch.manual_seed(2025) +logger = logging.getLogger(__name__) + + +@pytest.mark.parametrize( + "M, N, K, L", + [ + (256, 512, 1024, 1), + (256, 512, 64, 1), + (256, 512, 64, 2), + ], +) +def test_gemm( + M: int, + N: int, + K: int, + L: int, + fixture_toggle_tvm_ffi, +): + device = "cuda" + ab_dtype = torch.float16 + c_dtype = torch.float16 + accumulator_type = torch.float32 + A = torch.randint(-1, 2, (L, M, K), device=device, dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device=device, dtype=ab_dtype) + D = torch.empty((L, M, N), device=device, dtype=c_dtype) + + args = ops.GemmArguments(A.cuda(), B.cuda(), D.cuda(), accumulator_type) + + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + assert len(operators) > 0 + operator = operators[0] + logger.debug(f"Picked operator: {operator.metadata.operator_name}") + logger.debug(f"Operator metadata:\n{pformat(operator.metadata)}") + operator.run(args) + + reference = A.float() @ B.float() + assert_close_with_reference_conversion( + D, + reference, + D.dtype, + ) + + +def test_gemm_2d(fixture_toggle_tvm_ffi): + device = "cuda" + ab_dtype = torch.float16 + c_dtype = torch.float16 + accumulator_type = torch.float32 + M = 256 + N = 512 + K = 128 + A = torch.randint(-1, 2, (M, K), device=device, dtype=ab_dtype) + B = torch.randint(-1, 2, (K, N), device=device, dtype=ab_dtype) + D = torch.empty((M, N), device=device, dtype=c_dtype) + + args = ops.GemmArguments(A.cuda(), B.cuda(), D.cuda(), accumulator_type) + + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operator = operators[0] + logger.debug(f"Picked operator: {operator.metadata.operator_name}") + logger.debug(f"Operator metadata:\n{pformat(operator.metadata)}") + operator.run(args) + + reference = A.float() @ B.float() + assert_close_with_reference_conversion( + D, + reference, + D.dtype, + ) + + +def test_no_gemms_available(): + M = N = K = 128 + L = 1 + A = torch.empty((L, M, K)).to(torch.float32) + B = torch.empty((L, K, N)).to(torch.float32) + D = torch.empty((L, M, N)).to(torch.float32) + + args = ops.GemmArguments(A, B, D, accumulator_type=torch.float32) + operators = ops.get_operators(args, target_sm="70") + + # There are currently no operators available for compute capability 70. + assert len(operators) == 0 + + +def test_metadata_filter(): + device = "cuda" + # Test supplying metadata filter only + def tile_size_m_filter(metadata: ops.OperatorMetadata) -> bool: + if not isinstance(metadata.design, Sm100DesignMetadata): + return False + return metadata.design.tile_shape[0] == 64 + + operators = ops.get_operators(metadata_filter=tile_size_m_filter) + for operator in operators: + assert operator.metadata.design.tile_shape[0] == 64, ( + f"Operator {operator.metadata.operator_name} has tile shape " + f"{operator.metadata.design.tile_shape}" + ) + + # Test supplying metadata filter and arguments + A = torch.randint(-1, 2, (1, 256, 256), device=device).to(torch.float16) + B = torch.randint(-1, 2, (1, 256, 256), device=device).to(torch.float16) + D = torch.empty((1, 256, 256), device=device).to(torch.float16) + + args = ops.GemmArguments( + A.cuda(), B.cuda(), D.cuda(), accumulator_type=torch.float16 + ) + operators = ops.get_operators(args=args, metadata_filter=tile_size_m_filter) + for operator in operators: + assert operator.metadata.design.tile_shape[0] == 64, ( + f"Operator {operator.metadata.operator_name} has tile shape " + f"{operator.metadata.design.tile_shape}" + ) + assert operator.metadata.operands.A.dtype == cutlass.Float16 + assert operator.metadata.operands.B.dtype == cutlass.Float16 + assert operator.metadata.operands.out.dtype == cutlass.Float16 + assert operator.metadata.operands.accumulator_type == cutlass.Float16 + + +def test_gemm_fake_tensor(fixture_toggle_tvm_ffi): + import torch._functorch.config + + torch._functorch.config.fake_tensor_allow_unsafe_data_ptr_access = False + + device = "cuda" + M, N, K, L = 256, 512, 128, 1 + ab_dtype = torch.float16 + c_dtype = torch.float16 + accumulator_type = torch.float32 + + with torch._subclasses.fake_tensor.FakeTensorMode(): + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=c_dtype) + + fake_args = ops.GemmArguments(A, B, D, accumulator_type) + + operators = ops.get_operators(fake_args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operator = operators[0] + + compiled_artifact = operator.compile(fake_args) + + A = torch.randint(-1, 2, (L, M, K), device=device, dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device=device, dtype=ab_dtype) + D = torch.empty((L, M, N), device=device, dtype=c_dtype) + args = ops.GemmArguments(A.cuda(), B.cuda(), D.cuda(), accumulator_type) + + operator.run(args, compiled_artifact=compiled_artifact) + + reference = A.float() @ B.float() + assert_close_with_reference_conversion( + D, + reference, + D.dtype, + ) diff --git a/operators/test/integration/test_gemm_dynamic.py b/operators/test/integration/test_gemm_dynamic.py new file mode 100644 index 000000000..6211c58af --- /dev/null +++ b/operators/test/integration/test_gemm_dynamic.py @@ -0,0 +1,128 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import logging +from pprint import pformat + +import pytest +import torch + +import cutlass.operators as ops +from cutlass.operators.metadata.design.tile_scheduler import ( + CLCDynamicPersistentTileSchedulerMetadata, + StaticPersistentTileSchedulerMetadata, +) +from cutlass.operators.utils.device import ( + device_or_env_supports, + device_or_env_target_sm, +) + +from test_utils import assert_close_with_reference_conversion + +torch.manual_seed(2025) +logger = logging.getLogger(__name__) + + +@pytest.mark.parametrize( + "M, N, K, L", + [ + (256, 512, 1024, 1), + (256, 512, 64, 1), + (256, 512, 64, 2), + ], +) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Currently requires compiling for an arch in sm_100f family", +) +def test_gemm_persistent_tile_schedulers( + M: int, + N: int, + K: int, + L: int, + fixture_toggle_tvm_ffi, +): + """Test static and CLC persistent tile schedulers for dense GEMM.""" + ab_dtype = torch.float16 + c_dtype = torch.float16 + accumulator_type = torch.float32 + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=c_dtype) + + reference = A @ B + + args = ops.GemmArguments(A, B, D, accumulator_type) + + def clc_dynamic_persistent_filter(metadata): + scheduler = getattr(metadata.design, "tile_scheduler", None) + if scheduler and isinstance( + scheduler, CLCDynamicPersistentTileSchedulerMetadata + ): + return True + return False + + def static_persistent_filter(metadata): + scheduler = getattr(metadata.design, "tile_scheduler", None) + if scheduler and isinstance(scheduler, StaticPersistentTileSchedulerMetadata): + return True + return False + + # Test one CLC-based tile scheduler kernel + operators = ops.get_operators( + args, + metadata_filter=clc_dynamic_persistent_filter, + target_sm=device_or_env_target_sm(), + ) + assert len(operators) > 0 + + operator = operators[0] + logger.debug(f"Picked operator: {operator.metadata.operator_name}") + logger.debug(f"Operator metadata:\n{pformat(operator.metadata)}") + + operator.run(args) + assert_close_with_reference_conversion(D, reference, D.dtype) + + # Reinitialize output tensor to ensure independent validation + D = torch.empty((L, M, N), device="cuda", dtype=c_dtype) + args = ops.GemmArguments(A, B, D, accumulator_type) + + operators = ops.get_operators( + args, + metadata_filter=static_persistent_filter, + target_sm=device_or_env_target_sm(), + ) + + assert len(operators) > 0 + operator = operators[0] + + logger.debug(f"Picked operator: {operator.metadata.operator_name}") + logger.debug(f"Operator metadata:\n{pformat(operator.metadata)}") + + operator.run(args) + assert_close_with_reference_conversion(D, reference, D.dtype) diff --git a/operators/test/integration/test_gemm_epilogue_fusion.py b/operators/test/integration/test_gemm_epilogue_fusion.py new file mode 100644 index 000000000..f3ce422e0 --- /dev/null +++ b/operators/test/integration/test_gemm_epilogue_fusion.py @@ -0,0 +1,1158 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +import pytest +import torch + +import cutlass.operators as ops +from cutlass.operators.providers.cutedsl.gemm.sm100_static_persistent_efc import ( + PersistentDenseGemmEFCOperator, +) +from cutlass.operators.utils.device import ( + device_or_env_supports, + device_or_env_target_sm, +) + +from test_utils import assert_close_with_reference_conversion +from test_utils.reference_check import get_torch_default_tolerances + +torch.manual_seed(2025) + + +def problem_sizes(): + """ + Problem sizes for tests + """ + return [ + (256, 512, 1024, 1), + (256, 512, 128, 1), + (256, 512, 128, 2), + ] + + +def base_data_types(): + """ + Data types for (ab, c, d, accumulator) + """ + return [ + (torch.float16, torch.float32, torch.float32, torch.float32), + (torch.float16, torch.float16, torch.float16, torch.float16), + (torch.bfloat16, torch.bfloat16, torch.bfloat16, torch.float32), + ] + + +# Unary operation strings and functions +identity = ("", lambda x: x) +relu = ("relu", torch.relu) +tanh = ("tanh", torch.tanh) +sigmoid = ("sigmoid", torch.sigmoid) +exp = ("exp", torch.exp) + +unary_ops = [identity, relu, tanh, sigmoid, exp] + + +# Binary operation strings and functions +add = (lambda a, b: f"{a} + {b}", lambda a, b: a + b) +sub = (lambda a, b: f"{a} - {b}", lambda a, b: a - b) +mul = (lambda a, b: f"{a} * {b}", lambda a, b: a * b) + +# Don't include divide in main binary ops due to issues with division by zero in refchecks +binary_ops = [add, sub, mul] + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_no_fusion(fixture_toggle_tvm_ffi): + """ + Tests EFC GEMM with no fusion provided. + + """ + M, N, K, L = 256, 512, 128, 2 + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=torch.float16) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=torch.float16) + D = torch.empty((L, M, N), device="cuda", dtype=torch.float16) + + def metadata_filter(metadata): + return metadata.operator_class is PersistentDenseGemmEFCOperator + + args = ops.GemmArguments(A, B, D, accumulator_type=torch.float16) + operators = ops.get_operators( + args, target_sm=device_or_env_target_sm(), metadata_filter=metadata_filter + ) + assert len(operators) > 0 + operators[0].run(args) + + reference = A @ B + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.parametrize("M, N, K, L", problem_sizes()) +# Restrict to D of float16 for now to avoid rounding error when converting torch f16 output to f32 +@pytest.mark.parametrize( + "ab_dtype, d_dtype, accumulator_type", + [(torch.float16, torch.float16, torch.float16)], +) +@pytest.mark.parametrize("unary_str, unary_op", unary_ops) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_unary( + M, N, K, L, ab_dtype, d_dtype, accumulator_type, unary_str, unary_op +): + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + def epi(accum): + D = unary_op(accum) + return D + + epi_str = f"def epi(accum): D = {unary_str}(accum); return D" + + epi_args = ops.EpilogueArguments(epi_str, D=D) + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference = epi(A @ B) + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.parametrize("M, N, K, L", [(256, 512, 128, 2)]) +# Restrict to D of float16 for now to avoid rounding error when converting torch f16 output to f32 +@pytest.mark.parametrize( + "ab_dtype, d_dtype, accumulator_type", + [(torch.float16, torch.float16, torch.float16)], +) +@pytest.mark.parametrize("unary_str, unary_op", [relu]) +@pytest.mark.parametrize("unary_str2, unary_op2", [sigmoid, tanh]) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_unary_composition( + M, + N, + K, + L, + ab_dtype, + d_dtype, + accumulator_type, + unary_str, + unary_op, + unary_str2, + unary_op2, +): + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + def epi(accum): + D = unary_op2(unary_op(accum)) + return D + + epi_str = f"def epi(accum): D = {unary_str2}({unary_str}(accum)); return D" + + epi_args = ops.EpilogueArguments(epi_str, D=D) + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference = epi(A @ B) + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.parametrize("M, N, K, L", problem_sizes()) +# Restrict to D of float16 for now to avoid rounding error when converting torch f16 output to f32 +@pytest.mark.parametrize( + "ab_dtype, d_dtype, accumulator_type", + [(torch.float16, torch.float16, torch.float16)], +) +# Restrict unary to identity and relu to avoid rounding errors +@pytest.mark.parametrize("unary_str, unary_op", [identity, relu]) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_unary_literal( + M, N, K, L, ab_dtype, d_dtype, accumulator_type, unary_str, unary_op +): + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + def epi(accum): + D = unary_op(accum) * 3.0 - 1.234 + return D + + epi_str = f"def epi(accum): D = {unary_str}(accum) * 3.0 - 1.234; return D" + + epi_args = ops.EpilogueArguments(epi_str, D=D) + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference = epi(A @ B) + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.parametrize("M, N, K, L", problem_sizes()) +@pytest.mark.parametrize( + "ab_dtype, c_dtype, d_dtype, accumulator_type", base_data_types() +) +@pytest.mark.parametrize("unary_str, unary_op", [identity, relu]) +@pytest.mark.parametrize("binary_str, binary_op", binary_ops) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_unary_binary_composition( + M, + N, + K, + L, + ab_dtype, + c_dtype, + d_dtype, + accumulator_type, + unary_str, + unary_op, + binary_str, + binary_op, +): + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + C = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + def epi(accum, C): + z = unary_op(accum) + D = binary_op(z, C) + return D + + epi_str = f"def epi(accum, C): z = {unary_str}(accum); D = {binary_str('z', 'C')}; return D" + + epi_args = ops.EpilogueArguments(epi_str, C=C, D=D) + + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference = epi(A @ B, C) + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.parametrize("M, N, K, L", problem_sizes()) +# Restrict to D of float16 for now to avoid rounding error when converting torch f16 output to f32 +@pytest.mark.parametrize( + "ab_dtype, d_dtype, accumulator_type", + [(torch.float16, torch.float16, torch.float16)], +) +@pytest.mark.parametrize("c0_dtype", [torch.float16, torch.float32]) +@pytest.mark.parametrize("c1_dtype", [torch.float16, torch.float32]) +@pytest.mark.parametrize("binary_str0, binary_op0", [add, sub]) +@pytest.mark.parametrize("binary_str1, binary_op1", [add, sub]) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_binary_binary_composition( + M, + N, + K, + L, + ab_dtype, + d_dtype, + accumulator_type, + c0_dtype, + c1_dtype, + binary_str0, + binary_op0, + binary_str1, + binary_op1, +): + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + C0 = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c0_dtype) + C1 = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c1_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + def epi(accum, C0, C1): + z = torch.relu(accum) + z1 = binary_op0(z, C0) + D = binary_op1(z1, C1) + return D + + epi_str = f"def epi(accum, C0, C1): z = relu(accum); z1 = {binary_str0('z', 'C0')}; D = {binary_str1('z1', 'C1')}; return D" + + epi_args = ops.EpilogueArguments(epi_str, C0=C0, C1=C1, D=D) + + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference = epi(A @ B, C0, C1) + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_division(): + M, N, K, L = 256, 512, 128, 2 + ab_dtype = torch.float16 + d_dtype = torch.float16 + accumulator_type = torch.float16 + + # Specifically initialize A and B with ones to avoid division by zero in refchecks + A = torch.ones((L, M, K), device="cuda", dtype=ab_dtype) + B = torch.ones((L, K, N), device="cuda", dtype=ab_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + ######################################################### + # Test division by a literal + ######################################################### + def epi(accum): + D = accum / 2.0 + return D + + epi_str = "def epi(accum): D = accum / 2.0; return D" + epi_args = ops.EpilogueArguments(epi_str, D=D) + + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference = epi(A @ B) + assert_close_with_reference_conversion(D, reference, D.dtype) + + ######################################################### + # Test division by an input + ######################################################### + def epi(accum, scalar): + D = accum / scalar + return D + + epi_str = "def epi(accum, scalar): D = accum / scalar; return D" + + scalar = 4.0 + epi_args = ops.EpilogueArguments(epi_str, scalar=scalar, D=D) + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference = epi(A @ B, scalar) + + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.parametrize("M, N, K, L", problem_sizes()) +# Restrict to D of float16 for now to avoid rounding error when converting torch f16 output to f32 +@pytest.mark.parametrize( + "ab_dtype, d_dtype, accumulator_type", + [(torch.float16, torch.float16, torch.float16)], +) +@pytest.mark.parametrize("unary_str, unary_op", [sigmoid, tanh]) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_unary_multi_output( + M, N, K, L, ab_dtype, d_dtype, accumulator_type, unary_str, unary_op +): + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + z = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + def epi(accum): + z0 = torch.relu(accum) + z = unary_op(z0) + D = z + z0 + return D, z + + epi_str = f"def epi(accum): z0 = relu(accum); z = {unary_str}(z0); D = z + z0; return D, z" + + epi_args = ops.EpilogueArguments(epi_str, z=z, D=D) + + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + acc = (A @ B).to(accumulator_type) + ref_D, ref_z = epi(acc.float()) + rtol, atol = get_torch_default_tolerances()[d_dtype] + assert_close_with_reference_conversion(D, ref_D, D.dtype, rtols=rtol, atols=atol) + assert_close_with_reference_conversion(z, ref_z, z.dtype, rtols=rtol, atols=atol) + + +@pytest.mark.parametrize("M, N, K, L", problem_sizes()) +# Restrict to D of float16 for now to avoid rounding error when converting torch f16 output to f32 +@pytest.mark.parametrize( + "ab_dtype, d_dtype, accumulator_type", + [(torch.float16, torch.float16, torch.float16)], +) +@pytest.mark.parametrize("c_dtype", [torch.float16, torch.float32]) +@pytest.mark.parametrize("binary_str, binary_op", binary_ops) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_binary_multi_output( + M, N, K, L, ab_dtype, d_dtype, accumulator_type, c_dtype, binary_str, binary_op +): + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + C = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + z = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + def epi(accum, C): + z0 = torch.relu(accum) + z = binary_op(z0, C) + D = z + z0 + return D, z + + epi_str = f"def epi(accum, C): z0 = relu(accum); z = {binary_str('z0', 'C')}; D = z + z0; return D, z" + + epi_args = ops.EpilogueArguments(epi_str, C=C, z=z, D=D) + + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + ref_D, ref_z = epi(A @ B, C) + assert_close_with_reference_conversion(D, ref_D, D.dtype) + assert_close_with_reference_conversion(z, ref_z, z.dtype) + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_return_acc(): + M, N, K, L = 256, 512, 128, 2 + ab_dtype = torch.float16 + c_dtype = torch.float32 + d_dtype = torch.float16 + accumulator_type = torch.float16 + + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + C = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + accum = torch.empty((L, M, N), device="cuda", dtype=accumulator_type) + + def epi(accum, C): + D = torch.relu(accum) + C + return D, accum + + epi_str = "def epi(accum, C): D = relu(accum) + C; return D, accum" + + epi_args = ops.EpilogueArguments(epi_str, C=C, D=D, accum=accum) + + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference, ref_accum = epi(A @ B, C) + assert_close_with_reference_conversion(D, reference, D.dtype) + assert_close_with_reference_conversion(accum, ref_accum, accum.dtype) + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_acc_as_multiple_input(): + M, N, K, L = 256, 512, 128, 2 + ab_dtype = torch.float16 + c_dtype = torch.float32 + d_dtype = torch.float16 + accumulator_type = torch.float16 + + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + C = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + ######################################################### + # Test binary op inside + ######################################################### + def epi(accum, C): + D = torch.relu(torch.relu(accum) * C) + accum + return D + + epi_str = "def epi(accum, C): D = relu(relu(accum) * C) + accum; return D" + + epi_args = ops.EpilogueArguments(epi_str, C=C, D=D) + + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference = epi(A @ B, C) + assert_close_with_reference_conversion(D, reference, D.dtype) + + ######################################################### + # Test unary op inside + ######################################################### + def epi(accum): + D = torch.relu(torch.sigmoid(torch.relu(accum))) + accum + return D + + epi_str = "def epi(accum): D = relu(sigmoid(relu(accum))) + accum; return D" + + epi_args = ops.EpilogueArguments(epi_str, D=D) + + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + acc = (A @ B).to(accumulator_type) + reference = epi(acc.to(c_dtype)) + rtol, atol = get_torch_default_tolerances()[d_dtype] + assert_close_with_reference_conversion( + D, reference, D.dtype, rtols=rtol, atols=atol + ) + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_matmul_input_as_aux(): + M, N, K, L = 1024, 1024, 1024, 2 + ab_dtype = torch.float16 + c_dtype = torch.float32 + d_dtype = torch.float16 + accumulator_type = torch.float16 + + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + C = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + ######################################################### + # Test binary op inside + ######################################################### + def epi(accum, C, A): + D = torch.sigmoid(torch.relu(accum) * C) + A + return D + + epi_str = "def epi(accum, C, A): D = sigmoid(relu(accum) * C) + A; return D" + + epi_args = ops.EpilogueArguments(epi_str, C=C, A=A, D=D) + + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference = epi(A @ B, C, A) + assert_close_with_reference_conversion(D, reference, D.dtype) + + ######################################################### + # Test unary op inside + ######################################################### + def epi(accum, A): + D = torch.tanh(torch.relu(accum)) + A + return D + + epi_str = "def epi(accum, A): D = tanh(relu(accum)) + A; return D" + + epi_args = ops.EpilogueArguments(epi_str, A=A, D=D) + + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference = epi(A @ B, A) + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.parametrize("M, N, K, L", problem_sizes()) +@pytest.mark.parametrize( + "ab_dtype, c_dtype, d_dtype, accumulator_type", base_data_types() +) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_alpha_beta( + M, N, K, L, ab_dtype, c_dtype, d_dtype, accumulator_type, fixture_toggle_tvm_ffi +): + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + C = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + def epi(accum, C, alpha, beta): + D = alpha * accum + beta * C + return D + + alpha = 0.5 + beta = 0.5 + epi_args = ops.EpilogueArguments(epi, C=C, alpha=alpha, beta=beta, D=D) + + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operator = operators[0] + + for a, b in [(0.5, 0.5), (1.0, 0.0), (0.0, 1.0)]: + epi_args = ops.EpilogueArguments(epi, C=C, alpha=a, beta=b, D=D) + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operator.run(args) + reference = epi(A @ B, C, a, b) + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_alpha_beta_fake_tensor(fixture_toggle_tvm_ffi): + import torch._functorch.config + + torch._functorch.config.fake_tensor_allow_unsafe_data_ptr_access = False + + M, N, K, L = 256, 512, 128, 2 + ab_dtype = torch.float16 + c_dtype = torch.float32 + d_dtype = torch.bfloat16 + accumulator_type = torch.float16 + + with torch._subclasses.fake_tensor.FakeTensorMode(): + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + C = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + def epi(accum, C, alpha, beta): + D = alpha * accum + beta * C + return D + + alpha = 0.5 + beta = 0.5 + epi_args = ops.EpilogueArguments(epi, C=C, alpha=alpha, beta=beta, D=D) + + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operator = operators[0] + + A_real = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B_real = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + C_real = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + D_real = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + for a, b in [(0.5, 0.5), (1.0, 0.0), (0.0, 1.0)]: + epi_args = ops.EpilogueArguments(epi, C=C_real, alpha=a, beta=b, D=D_real) + args = ops.GemmArguments( + A=A_real, + B=B_real, + out=D_real, + accumulator_type=accumulator_type, + epilogue=epi_args, + ) + operator.run(args) + reference = epi(A_real @ B_real, C_real, a, b) + assert_close_with_reference_conversion(D_real, reference, D_real.dtype) + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_big_epi(fixture_toggle_tvm_ffi): + M, N, K, L = 256, 512, 128, 2 + ab_dtype = torch.float16 + c_dtype = torch.float32 + d_dtype = torch.bfloat16 + accumulator_type = torch.float16 + + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + + In0 = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + In1 = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + In2 = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + In3 = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + In4 = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + In5 = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + In6 = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + In7 = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + + Out0 = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + Out1 = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + Out2 = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + Out3 = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + Out4 = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + Out5 = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + Out6 = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + sc0 = 1.0 + sc1 = 2.0 + sc2 = 3.0 + sc3 = 4.0 + sc4 = 5.0 + sc5 = 6.0 + sc6 = 7.0 + sc7 = 8.0 + + def epi( + accum, + In0, + In1, + In2, + In3, + In4, + In5, + In6, + In7, + sc0, + sc1, + sc2, + sc3, + sc4, + sc5, + sc6, + sc7, + ): + Out0 = accum * sc0 + In0 + Out1 = Out0 + In1 * sc1 + Out2 = Out1 - In2 * sc2 + Out3 = Out2 + In3 * sc3 + Out4 = Out3 - In4 * sc4 + Out5 = Out4 + In5 * sc5 + Out6 = Out5 - In6 * sc6 + D = Out6 + In7 * sc7 + return Out0, Out1, Out2, Out3, Out4, Out5, Out6, D + + epi_args = ops.EpilogueArguments( + epi, + In0=In0, + In1=In1, + In2=In2, + In3=In3, + In4=In4, + In5=In5, + In6=In6, + In7=In7, + Out0=Out0, + Out1=Out1, + Out2=Out2, + Out3=Out3, + Out4=Out4, + Out5=Out5, + Out6=Out6, + D=D, + sc0=sc0, + sc1=sc1, + sc2=sc2, + sc3=sc3, + sc4=sc4, + sc5=sc5, + sc6=sc6, + sc7=sc7, + ) + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference = epi( + A @ B, + In0, + In1, + In2, + In3, + In4, + In5, + In6, + In7, + sc0, + sc1, + sc2, + sc3, + sc4, + sc5, + sc6, + sc7, + ) + + for out, ref in zip([Out0, Out1, Out2, Out3, Out4, Out5, Out6, D], reference): + assert_close_with_reference_conversion(out, ref, out.dtype) + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_reduction_not_available(fixture_toggle_tvm_ffi): + """Reductions are not yet supported by the EFC converter.""" + M = 256 + N = 512 + K = 1024 + A = torch.randint(-1, 2, (M, K), device="cuda", dtype=torch.float16) + B = torch.randint(-1, 2, (K, N), device="cuda", dtype=torch.float16) + D = torch.empty((M, N), device="cuda", dtype=torch.float16) + row_max = torch.empty((N,), device="cuda", dtype=torch.float32) + + epi_str = ( + "def epi(accum): row_max = max(accum, dim=[0]); D = accum; return D, row_max" + ) + + epi_args = ops.EpilogueArguments(epi_str, row_max=row_max, D=D) + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=torch.float16, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) == 0 + + +@pytest.mark.parametrize("M, N, K, L", problem_sizes()) +@pytest.mark.parametrize( + "ab_dtype, d_dtype, accumulator_type", + [(torch.float16, torch.float16, torch.float16)], +) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_column_broadcast(M, N, K, L, ab_dtype, d_dtype, accumulator_type): + """Column broadcast: bias shape (L, M, 1) is broadcast along N.""" + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + bias = torch.randint(-1, 2, (L, M, 1), device="cuda", dtype=d_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + def epi(accum, bias): + D = accum + bias + return D + + epi_args = ops.EpilogueArguments(epi, bias=bias, D=D) + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference = epi(A @ B, bias) + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.parametrize("M, N, K, L", problem_sizes()) +@pytest.mark.parametrize( + "ab_dtype, d_dtype, accumulator_type", + [(torch.float16, torch.float16, torch.float16)], +) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_row_broadcast(M, N, K, L, ab_dtype, d_dtype, accumulator_type): + """Row broadcast: bias shape (N,) is broadcast along L and M.""" + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + bias = torch.randint(-1, 2, (N,), device="cuda", dtype=d_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + def epi(accum, bias): + D = accum + bias + return D + + epi_args = ops.EpilogueArguments(epi, bias=bias, D=D) + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference = epi(A @ B, bias) + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.parametrize("M, N, K, L", [(256, 512, 128, 2)]) +@pytest.mark.parametrize( + "ab_dtype, d_dtype, accumulator_type", + [(torch.float16, torch.float16, torch.float16)], +) +@pytest.mark.parametrize("unary_str, unary_op", [identity, relu]) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_broadcast_with_activation( + M, N, K, L, ab_dtype, d_dtype, accumulator_type, unary_str, unary_op +): + """Broadcast combined with activation and scalar: D = activation(alpha * accum + bias).""" + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + col_bias = torch.randint(-1, 2, (L, M, 1), device="cuda", dtype=d_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + def epi(accum, col_bias, alpha): + D = unary_op(alpha * accum + col_bias) + return D + + alpha = 0.5 + epi_str = f"def epi(accum, col_bias, alpha): D = {unary_str}(alpha * accum + col_bias); return D" + + epi_args = ops.EpilogueArguments(epi_str, col_bias=col_bias, alpha=alpha, D=D) + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference = epi(A @ B, col_bias, alpha) + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.parametrize("M, N, K, L", problem_sizes()) +@pytest.mark.parametrize( + "ab_dtype, d_dtype, accumulator_type", + [(torch.float16, torch.float16, torch.float16)], +) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_row_broadcast_2d(M, N, K, L, ab_dtype, d_dtype, accumulator_type): + """Row broadcast with rank-2 bias (1, N) — size-1 batch dim is broadcast.""" + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + bias = torch.randint(-1, 2, (1, N), device="cuda", dtype=d_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + def epi(accum, bias): + D = accum + bias + return D + + epi_args = ops.EpilogueArguments(epi, bias=bias, D=D) + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference = epi(A @ B, bias) + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.parametrize("M, N, K, L", problem_sizes()) +@pytest.mark.parametrize( + "ab_dtype, d_dtype, accumulator_type", + [(torch.float16, torch.float16, torch.float16)], +) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_column_broadcast_2d( + M, N, K, L, ab_dtype, d_dtype, accumulator_type +): + """Column broadcast with rank-2 bias (M, 1) — common bias-per-row pattern.""" + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + bias = torch.randint(-1, 2, (M, 1), device="cuda", dtype=d_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + def epi(accum, bias): + D = accum + bias + return D + + epi_args = ops.EpilogueArguments(epi, bias=bias, D=D) + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference = epi(A @ B, bias) + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.parametrize("M, N, K, L", [(256, 512, 128, 2)]) +@pytest.mark.parametrize( + "ab_dtype, d_dtype, accumulator_type", + [(torch.float16, torch.float16, torch.float16)], +) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_row_broadcast_with_scalar_multiply( + M, N, K, L, ab_dtype, d_dtype, accumulator_type +): + """Row broadcast combined with scalar multiply: D = accum * alpha + bias.""" + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + bias = torch.randint(-1, 2, (N,), device="cuda", dtype=d_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + def epi(accum, bias, alpha): + D = accum * alpha + bias + return D + + alpha = 0.5 + epi_args = ops.EpilogueArguments(epi, bias=bias, alpha=alpha, D=D) + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference = epi(A @ B, bias, alpha) + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.parametrize("M, N, K, L", [(256, 512, 128, 2)]) +@pytest.mark.parametrize( + "ab_dtype, d_dtype, accumulator_type", + [(torch.float16, torch.float16, torch.float16)], +) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_gemm_fusion_both_row_and_col_broadcast( + M, N, K, L, ab_dtype, d_dtype, accumulator_type +): + """Both row and column broadcast in a single epilogue: D = accum + row_bias + col_bias.""" + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + row_bias = torch.randint(-1, 2, (N,), device="cuda", dtype=d_dtype) + col_bias = torch.randint(-1, 2, (L, M, 1), device="cuda", dtype=d_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + epi_str = ( + "def epi(accum, row_bias, col_bias): D = accum + row_bias + col_bias; return D" + ) + + def epi_ref(accum, row_bias, col_bias): + return accum + row_bias + col_bias + + epi_args = ops.EpilogueArguments(epi_str, row_bias=row_bias, col_bias=col_bias, D=D) + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operators[0].run(args) + + reference = epi_ref(A @ B, row_bias, col_bias) + assert_close_with_reference_conversion(D, reference, D.dtype) diff --git a/operators/test/integration/test_gemm_jax.py b/operators/test/integration/test_gemm_jax.py new file mode 100644 index 000000000..85a18a599 --- /dev/null +++ b/operators/test/integration/test_gemm_jax.py @@ -0,0 +1,74 @@ +# Copyright (c) 2026 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +import logging +from pprint import pformat + +import jax +import jax.numpy as jnp +import pytest + +import cutlass + +import cutlass.operators as ops +from cutlass.operators.utils.device import device_or_env_target_sm + +from test_utils import assert_close_with_reference_conversion + +# The test exercises the JAX-to-GPU operator path, which requires a GPU backend. +pytestmark = pytest.mark.skipif( + not any(d.platform == "gpu" for d in jax.devices()), + reason="JAX has no GPU device available", +) + +logger = logging.getLogger(__name__) + + +@pytest.mark.parametrize("M, N, K", [(256, 512, 128)]) +def test_gemm_jax(M, N, K, fixture_toggle_tvm_ffi): + key = jax.random.PRNGKey(2025) + k1, k2 = jax.random.split(key) + + A = jax.random.normal(k1, (1, M, K), dtype=jnp.float16) + B = jax.random.normal(k2, (1, K, N), dtype=jnp.float16) + D = jnp.zeros((1, M, N), dtype=jnp.float16) + + args = ops.GemmArguments(A, B, D, accumulator_type=cutlass.Float32) + operators = ops.get_operators(args, target_sm=device_or_env_target_sm()) + + assert len(operators) > 0 + operator = operators[0] + logger.debug(f"Picked operator: {operator.metadata.operator_name}") + logger.debug(f"Operator metadata:\n{pformat(operator.metadata)}") + operator.run(args) + + reference = jnp.matmul(A, B) + assert_close_with_reference_conversion( + D, reference, cutlass.Float16, rtols=1e-2, atols=1e-2 + ) diff --git a/operators/test/integration/test_mixed_input_gemm.py b/operators/test/integration/test_mixed_input_gemm.py new file mode 100644 index 000000000..e54a5aaf4 --- /dev/null +++ b/operators/test/integration/test_mixed_input_gemm.py @@ -0,0 +1,274 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +Integration tests for Mixed-Input GEMM Operators. + +These tests verify the user-facing API for mixed-input GEMM operations where +tensor A uses narrow precision (Int8/Uint8) and tensor B uses wide precision +(BFloat16/Float16). + +Exhaustive correctness testing across all Operator variants, tile shapes, and +problem geometries is handled by the implementation tests in +test/implementation/test_gemm_operator.py, which includes MixedInputGemmOperator. + +These integration tests focus on: +- Basic end-to-end correctness with a representative configuration +- 2D tensor support +- API behaviors: Operator discovery, argument validation, metadata filtering +- Compile-and-reuse workflow +""" + +import logging + +import pytest +import torch + +import cutlass.operators as ops +from cutlass.operators.utils.device import device_or_env_supports + +from test_utils import assert_close_with_reference_conversion + +torch.manual_seed(2025) +logger = logging.getLogger(__name__) + + +# ============================================================================= +# Basic Correctness Tests (representative configurations only) +# ============================================================================= + + +@pytest.mark.parametrize( + "M, N, K, L", + [ + (256, 512, 256, 1), + (128, 128, 128, 2), + (256, 256, 512, 1), + ], +) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_mixed_input_gemm( + M: int, + N: int, + K: int, + L: int, + fixture_toggle_tvm_ffi, +): + """ + Test basic mixed-input GEMM with Int8 A and BFloat16 B. + """ + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=torch.int8) + B = torch.randint(-1, 2, (L, K, N), device="cuda").to(torch.bfloat16) + D = torch.empty((L, M, N), device="cuda", dtype=torch.bfloat16) + + args = ops.GemmArguments(A=A, B=B, out=D, accumulator_type=torch.float32) + operators = ops.get_operators(args) + + assert len(operators) > 0 + operator = operators[0] + logger.debug(f"Picked operator: {operator.metadata.operator_name}") + + assert operator.supports(args) + operator.run(args) + + reference = A.to(torch.float32) @ B.to(torch.float32) + assert_close_with_reference_conversion( + D, + reference, + D.dtype, + ) + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_mixed_input_gemm_2d(fixture_toggle_tvm_ffi): + """ + Test 2D tensors (non-batched case). + """ + M, N, K = 256, 512, 256 + A = torch.randint(-1, 2, (M, K), device="cuda", dtype=torch.int8) + B = torch.randint(-1, 2, (K, N), device="cuda").to(torch.bfloat16) + D = torch.empty((M, N), device="cuda", dtype=torch.bfloat16) + + args = ops.GemmArguments(A=A, B=B, out=D, accumulator_type=torch.float32) + operators = ops.get_operators(args) + + assert len(operators) > 0 + operator = operators[0] + logger.debug(f"Picked operator: {operator.metadata.operator_name}") + + assert operator.supports(args) + operator.run(args) + + reference = A.to(torch.float32) @ B.to(torch.float32) + assert_close_with_reference_conversion( + D, + reference, + D.dtype, + ) + + +# ============================================================================= +# Convert-Scale Mode Tests (Int4) +# ============================================================================= +# NOTE: These tests are skipped because PyTorch does not have native Int4 support. +# Testing Int4 requires packed Int4 tensors using CUTLASS utilities (like +# create_i4_tensor_and_scale from mixed_input_host_utils.py) which are not yet +# exposed through the cutlass.operators. To enable these tests: +# 1. Add Int4 tensor creation utilities to cutlass.operators +# 2. Or use cutlass.cute tensor creation with proper packing + + +@pytest.mark.skip( + reason="Int4 tensors require CUTLASS packed format not available through PyTorch" +) +def test_mixed_input_gemm_convert_scale_int4(): + """ + Test convert-scale mode with Int4 A tensors. + + Convert-scale mode: out = (type_convert(A) * scale) @ B + + This test is skipped because PyTorch does not support Int4 dtype. + Int4 tensors must be created using CUTLASS utilities with proper + 2-element packing into 8-bit storage. + """ + + +# ============================================================================= +# Edge Cases and Error Handling +# ============================================================================= + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_mixed_input_gemm_no_operators_for_unsupported_cc(): + """ + Test that no operators are returned for unsupported compute capabilities. + """ + M, N, K, L = 128, 128, 128, 1 + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=torch.int8) + B = torch.randint(-1, 2, (L, K, N), device="cuda").to(torch.bfloat16) + D = torch.empty((L, M, N), device="cuda", dtype=torch.bfloat16) + + args = ops.GemmArguments(A=A, B=B, out=D, accumulator_type=torch.float32) + + # Request operators for SM80 (not supported for mixed-input) + operators = ops.get_operators(args, target_sm="80") + + # Mixed-input GEMM is only supported on SM100+ + assert len(operators) == 0, "Expected no operators for SM80" + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_mixed_input_gemm_invalid_accumulator_for_bfloat16(): + """ + Test that operators correctly reject invalid accumulator type for BFloat16 B. + BFloat16 B requires Float32 accumulator. + """ + M, N, K, L = 128, 128, 128, 1 + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=torch.int8) + B = torch.randint(-1, 2, (L, K, N), device="cuda").to(torch.bfloat16) + D = torch.empty((L, M, N), device="cuda", dtype=torch.bfloat16) + + # Try Float16 accumulator with BFloat16 B (invalid) + args = ops.GemmArguments(A=A, B=B, out=D, accumulator_type=torch.float16) + + # Should return no operators (invalid accumulator for BFloat16) + operators = ops.get_operators(args) + assert len(operators) == 0, ( + "Expected no operators for Float16 accumulator with BFloat16 B" + ) + + +def test_mixed_input_gemm_argument_validation(): + """ + Test argument validation in GemmArguments for mixed-input. + """ + M, N, K, L = 128, 128, 128, 1 + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=torch.int8) + B = torch.randint(-1, 2, (L, K, N), device="cuda").to(torch.bfloat16) + D = torch.empty((L, M, N), device="cuda", dtype=torch.bfloat16) + + # Valid arguments should work + args = ops.GemmArguments(A=A, B=B, out=D, accumulator_type=torch.float32) + assert args is not None + + # Test shape mismatch - A and B must have matching K dimension + B_wrong_k = torch.randint(-1, 2, (L, K + 1, N), device="cuda").to(torch.bfloat16) + with pytest.raises(ValueError, match="K dimension"): + ops.GemmArguments(A=A, B=B_wrong_k, out=D, accumulator_type=torch.float32) + + +# ============================================================================= +# Operator Selection and Compilation Tests +# ============================================================================= + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_mixed_input_gemm_metadata_filter(fixture_toggle_tvm_ffi): + """ + Test that metadata filter correctly filters operators. + """ + M, N, K, L = 256, 256, 256, 1 + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=torch.int8) + B = torch.randint(-1, 2, (L, K, N), device="cuda").to(torch.bfloat16) + D = torch.empty((L, M, N), device="cuda", dtype=torch.bfloat16) + + args = ops.GemmArguments(A=A, B=B, out=D, accumulator_type=torch.float32) + + # Get all matching operators + all_operators = ops.get_operators(args) + + if len(all_operators) == 0: + pytest.skip("No operators available") + + # Filter by specific tile shape + def tile_filter(metadata): + return metadata.design is not None and metadata.design.tile_shape == ( + 128, + 128, + 64, + ) + + filtered_operators = ops.get_operators(args, metadata_filter=tile_filter) + + # Filtered list should be a subset + assert len(filtered_operators) <= len(all_operators) diff --git a/operators/test/integration/test_notebooks.py b/operators/test/integration/test_notebooks.py new file mode 100644 index 000000000..bc73377e0 --- /dev/null +++ b/operators/test/integration/test_notebooks.py @@ -0,0 +1,93 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import sys +from pathlib import Path + +import pytest + +import test_utils # isort: skip + + +@pytest.mark.parametrize( + "notebook_name, operator_targets", + [ + ("000_gemm.ipynb", ["80", "90a", "100f"]), + ("001_gemm_with_fused_epilogue.ipynb", ["100f"]), + ("002_bring_your_own_kernel.ipynb", ["80"]), + ("003_host_latency_best_practices.ipynb", ["80", "90a", "100f"]), + ("004_fake_tensors.ipynb", ["80"]), + ("005_grouped_gemm_contiguous_offset.ipynb", ["100a"]), + ("006_block_scaled_gemm.ipynb", ["100f"]), + ], +) +def test_notebooks(notebook_name, operator_targets, monkeypatch): + test_utils.skip_if_unsupported_target(operator_targets) + + notebook_dir = Path(__file__).parent / ".." / ".." / "examples" + full_notebook_path = notebook_dir / notebook_name + + import os + import subprocess + + import nbformat + from nbconvert.preprocessors import ExecutePreprocessor + + # Preload torch's libgomp and libc10 in the kernel subprocess so their + # static TLS is allocated before jupyter/jaxlib exhaust the budget. + # Without this the kernel's `import torch` fails with "cannot allocate + # memory in static TLS block". libgomp is the heaviest TLS user so it + # must come first. + from cutlass.operators.utils.common import is_torch_available + + if sys.platform.startswith("linux") and is_torch_available(): + import torch + + torch_lib = Path(torch.__file__).parent / "lib" + preload_libs = [] + for pattern in ("libgomp*.so*", "libc10.so"): + preload_libs.extend(sorted(str(p) for p in torch_lib.glob(pattern))) + if preload_libs: + existing = os.environ.get("LD_PRELOAD", "") + value = ":".join(preload_libs + ([existing] if existing else [])) + monkeypatch.setenv("LD_PRELOAD", value) + + # Register the current Python interpreter as the python3 kernel + subprocess.run( + [sys.executable, "-m", "ipykernel", "install", "--user", "--name", "python3"], + check=True, + capture_output=True, + ) + + with full_notebook_path.open() as file: + notebook = nbformat.read(file, as_version=4) + ep = ExecutePreprocessor(timeout=600, kernel_name="python3") + + # Execute the notebook. This call will error out on any assertions or errors in the + # notebook itself. Allow these to propagate up so the test will fail on notebook failure. + ep.preprocess(notebook, {"metadata": {"path": notebook_dir}}) diff --git a/operators/test/integration/test_overhead.py b/operators/test/integration/test_overhead.py new file mode 100644 index 000000000..d1e40bdb6 --- /dev/null +++ b/operators/test/integration/test_overhead.py @@ -0,0 +1,311 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +import time + +import pytest +import torch +from torch.cuda import current_stream + +import cutlass.operators as ops +from cutlass.operators.utils.device import ( + device_or_env_supports, + device_or_env_target_sm, + to_cuda_stream, +) +from cutlass.operators.utils.tensor import TensorWrapper + +from test_utils import assert_close_with_reference_conversion + + +def _benchmark_cpu(label, code, n_iterations): + """Measure CPU-side wall time per call for `code`""" + + # Warmup + for _ in range(n_iterations): + code() + # Sync before timed code + torch.cuda.synchronize() + + start = time.perf_counter() + for _ in range(n_iterations): + code() + # Skip synchronization here: we only measure host time + # independent of device code perf + end = time.perf_counter() + elapsed = end - start + + avg_us = elapsed / n_iterations * 1e6 + print(f"[{label:40}] avg of {n_iterations} iterations: {avg_us:8.1f} us") + return avg_us + + +def _extract_runtime_args(*fn_args): + """Extract runtime tensors from TensorWrappers, passing other args through.""" + return [x.runtime_tensor if isinstance(x, TensorWrapper) else x for x in fn_args] + + +def test_overhead(fixture_toggle_tvm_ffi): + """Compare CPU overhead of operator.run() vs direct compiled_obj() call.""" + print() + M, N, K = 256, 512, 1024 + A = torch.randint(-1, 2, (M, K), device="cuda", dtype=torch.float16) + B = torch.randint(-1, 2, (K, N), device="cuda", dtype=torch.float16) + D = torch.empty((M, N), device="cuda", dtype=torch.float16) + + args = ops.GemmArguments(A=A, B=B, out=D, accumulator_type=torch.float32) + operators = ops.get_operators(args, target_sm="80", providers=[ops.CuTeDSLProvider]) + assert len(operators) > 0 + operator = operators[0] + + assert operator.supports(args) + compiled_artifact = operator.compile(args) + stream = current_stream() + + n_iterations = 100 + + api_time = _benchmark_cpu( + "With CUTLASS Operator API (operator.run())", + lambda: operator.run( + args, + compiled_artifact=compiled_artifact, + stream=stream, + assume_supported_args=True, + ), + n_iterations, + ) + + compiled_obj = compiled_artifact.compiled_obj + direct_args = _extract_runtime_args( + args.A.tensor, + args.B.tensor, + args.out.tensor, + to_cuda_stream(stream), + ) + + _benchmark_cpu( + "Argument extraction", + lambda: _extract_runtime_args( + args.A.tensor, + args.B.tensor, + args.out.tensor, + to_cuda_stream(stream), + ), + n_iterations, + ) + direct_time = _benchmark_cpu( + "Direct CuTe DSL (compiled_obj())", + lambda: compiled_obj(*direct_args), + n_iterations, + ) + + overhead_us = api_time - direct_time + print(f" Overhead: {overhead_us:.1f} us over direct CuTe DSL") + + torch.cuda.synchronize() + reference = A @ B + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_overhead_big_epilogue(fixture_toggle_tvm_ffi): + """Compare CPU overhead with a large epilogue (8 aux inputs, 8 aux outputs).""" + print() + + M, N, K, L = 256, 512, 128, 2 + ab_dtype = torch.float16 + c_dtype = torch.float32 + d_dtype = torch.bfloat16 + accumulator_type = torch.float16 + + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=ab_dtype) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=ab_dtype) + + In0 = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + In1 = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + In2 = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + In3 = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + In4 = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + In5 = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + In6 = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + In7 = torch.randint(-1, 2, (L, M, N), device="cuda", dtype=c_dtype) + + Out0 = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + Out1 = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + Out2 = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + Out3 = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + Out4 = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + Out5 = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + Out6 = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + Out7 = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + D = torch.empty((L, M, N), device="cuda", dtype=d_dtype) + + sc0, sc1, sc2, sc3 = 1.0, 2.0, 3.0, 4.0 + sc4, sc5, sc6, sc7, sc8 = 5.0, 6.0, 7.0, 8.0, 9.0 + + def epi( + accum, + In0, + In1, + In2, + In3, + In4, + In5, + In6, + In7, + sc0, + sc1, + sc2, + sc3, + sc4, + sc5, + sc6, + sc7, + sc8, + ): + Out0 = accum * sc0 + In0 + Out1 = Out0 + In1 * sc1 + Out2 = Out1 - In2 * sc2 + Out3 = Out2 + In3 * sc3 + Out4 = Out3 - In4 * sc4 + Out5 = Out4 + In5 * sc5 + Out6 = Out5 - In6 * sc6 + Out7 = Out6 + In7 * sc7 + D = Out7 * sc8 + return Out0, Out1, Out2, Out3, Out4, Out5, Out6, Out7, D + + epi_args = ops.EpilogueArguments( + epi, + In0=In0, + In1=In1, + In2=In2, + In3=In3, + In4=In4, + In5=In5, + In6=In6, + In7=In7, + Out0=Out0, + Out1=Out1, + Out2=Out2, + Out3=Out3, + Out4=Out4, + Out5=Out5, + Out6=Out6, + Out7=Out7, + D=D, + sc0=sc0, + sc1=sc1, + sc2=sc2, + sc3=sc3, + sc4=sc4, + sc5=sc5, + sc6=sc6, + sc7=sc7, + sc8=sc8, + ) + args = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type, epilogue=epi_args + ) + operators = ops.get_operators( + args, target_sm=device_or_env_target_sm(), providers=[ops.CuTeDSLProvider] + ) + assert len(operators) > 0 + operator = operators[0] + + assert operator.supports(args) + compiled_artifact = operator.compile(args) + stream = current_stream() + + n_iterations = 100 + + api_time = _benchmark_cpu( + "With CUTLASS Operator API (operator.run())", + lambda: operator.run( + args, + compiled_artifact=compiled_artifact, + stream=stream, + assume_supported_args=True, + ), + n_iterations, + ) + + compiled_obj = compiled_artifact.compiled_obj + direct_args = _extract_runtime_args( + args.A.tensor, + args.B.tensor, + to_cuda_stream(stream), + *args.epilogue.parameters, + ) + + _benchmark_cpu( + "Argument extraction", + lambda: _extract_runtime_args( + args.A.tensor, + args.B.tensor, + to_cuda_stream(stream), + *args.epilogue.parameters, + ), + n_iterations, + ) + + direct_time = _benchmark_cpu( + "Direct CuTe DSL", + lambda: compiled_obj(*direct_args), + n_iterations, + ) + + overhead_us = api_time - direct_time + print(f" Overhead: {overhead_us:.1f} us over direct CuTe DSL") + + torch.cuda.synchronize() + reference = epi( + A @ B, + In0, + In1, + In2, + In3, + In4, + In5, + In6, + In7, + sc0, + sc1, + sc2, + sc3, + sc4, + sc5, + sc6, + sc7, + sc8, + ) + for out, ref in zip([Out0, Out1, Out2, Out3, Out4, Out5, Out6, Out7, D], reference): + assert_close_with_reference_conversion(out, ref, out.dtype) diff --git a/operators/test/integration/test_preferred_cluster_gemm.py b/operators/test/integration/test_preferred_cluster_gemm.py new file mode 100644 index 000000000..c6b2e4069 --- /dev/null +++ b/operators/test/integration/test_preferred_cluster_gemm.py @@ -0,0 +1,415 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +Integration tests for the preferred-cluster dense GEMM Operator. + +These tests verify the user-facing API for the Blackwell +``PersistentDenseGemmPreferredClusterOperator``, which uses a preferred (larger) +cluster shape with a fallback (smaller) cluster to let the driver fill the grid +with the preferred shape where SMs are available and fall back where they are +not. + +Exhaustive correctness testing across all Operator variants, tile shapes, and +cluster configurations is handled by the implementation tests in +``test/implementation/test_gemm_kernel.py``, which includes +``PersistentDenseGemmPreferredClusterOperator``. + +These integration tests focus on: +- Per-scheduler (CLC dynamic and static persistent) correctness +- 2D tensor support +- API behaviors: Operator discovery, metadata filtering, argument validation +- Compile-and-reuse workflow + +Problem sizes are chosen to be at least one preferred cluster tile +(``preferred_cluster_m × cta_m``, ``preferred_cluster_n × cta_n``) in each of M +and N. This is a hard requirement for the operator: smaller problems produce a +zero-sized launch grid because ``max_preferred_cluster_count`` floor-divides +the fallback CTA count by the preferred cluster size. +""" + +import logging +from pprint import pformat + +import pytest +import torch + +import cutlass.operators as ops +from cutlass.operators.metadata.design.sm100 import Sm100DesignMetadata +from cutlass.operators.metadata.design.tile_scheduler import ( + CLCDynamicPersistentTileSchedulerMetadata, + StaticPersistentTileSchedulerMetadata, + TileSchedulerMetadata, +) +from cutlass.operators.utils.device import ( + device_or_env_supports, + device_or_env_target_sm, +) + +from test_utils import assert_close_with_reference_conversion + +torch.manual_seed(2025) +logger = logging.getLogger(__name__) + + +# ============================================================================= +# Helpers +# ============================================================================= + + +def _preferred_cluster_filter( + metadata: ops.OperatorMetadata, + scheduler_type: type[TileSchedulerMetadata] | None = None, +) -> bool: + """Select a preferred-cluster Operator variant. + + Args: + metadata: Operator metadata to test. + scheduler_type: If given, only match this scheduler variant; otherwise + match any preferred-cluster variant (scheduler-agnostic). + + Returns: + bool: Whether the operator is a matching preferred-cluster variant. + """ + design = metadata.design + if not isinstance(design, Sm100DesignMetadata): + return False + if not design.use_fallback_cluster(): + return False + if scheduler_type is None: + return True + return isinstance(getattr(design, "tile_scheduler", None), scheduler_type) + + +# ============================================================================= +# 2D Tensor Support +# ============================================================================= + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_preferred_cluster_gemm_2d(fixture_toggle_tvm_ffi): + """2D (non-batched) tensor support. + + Exercises the ``add_batch_mode`` path in the Operator's ``__call__`` that + normalises rank-2 inputs to rank-3 before the main GEMM body. + """ + M, N, K = 512, 512, 128 + A = torch.randint(-1, 2, (M, K), device="cuda", dtype=torch.float16) + B = torch.randint(-1, 2, (K, N), device="cuda", dtype=torch.float16) + out = torch.empty((M, N), device="cuda", dtype=torch.float16) + + args = ops.GemmArguments(A=A, B=B, out=out, accumulator_type=torch.float32) + operators = ops.get_operators(args, metadata_filter=_preferred_cluster_filter) + + assert len(operators) > 0, "No preferred-cluster Operator matched the filter" + operator = operators[0] + logger.debug(f"Picked operator: {operator.metadata.operator_name}") + + assert operator.supports(args) + operator.run(args) + + reference = A.float() @ B.float() + assert_close_with_reference_conversion( + out, + reference, + out.dtype, + atols=0.0, + rtols=0.0, + msg=f"preferred-cluster 2D GEMM failed: M={M}, N={N}, K={K}", + ) + + +# ============================================================================= +# Per-Scheduler Correctness Tests +# ============================================================================= + + +@pytest.mark.parametrize( + "M, N, K, L", + [ + (256, 256, 256, 1), + (512, 512, 128, 1), + (512, 512, 128, 2), + ], +) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_preferred_cluster_gemm_clc_scheduler( + M: int, + N: int, + K: int, + L: int, + fixture_toggle_tvm_ffi, +): + """Preferred-cluster GEMM with CLC dynamic scheduler produces correct results.""" + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=torch.float16) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=torch.float16) + out = torch.empty((L, M, N), device="cuda", dtype=torch.float16) + + args = ops.GemmArguments(A=A, B=B, out=out, accumulator_type=torch.float32) + operators = ops.get_operators( + args, + metadata_filter=lambda m: _preferred_cluster_filter( + m, CLCDynamicPersistentTileSchedulerMetadata + ), + target_sm=device_or_env_target_sm(), + ) + + assert len(operators) > 0, "No preferred-cluster CLC operator matched the filter" + operator = operators[0] + logger.debug(f"Picked operator: {operator.metadata.operator_name}") + logger.debug(f"Operator metadata:\n{pformat(operator.metadata)}") + + assert operator.supports(args) + operator.run(args) + + reference = A.float() @ B.float() + assert_close_with_reference_conversion( + out, + reference, + out.dtype, + atols=0.0, + rtols=0.0, + msg=f"CLC preferred-cluster GEMM failed: M={M}, N={N}, K={K}, L={L}", + ) + + +@pytest.mark.parametrize( + "M, N, K, L", + [ + (256, 256, 256, 1), + (512, 512, 128, 1), + (512, 512, 128, 2), + ], +) +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_preferred_cluster_gemm_static_scheduler( + M: int, + N: int, + K: int, + L: int, + fixture_toggle_tvm_ffi, +): + """Preferred-cluster GEMM with static persistent scheduler produces correct results.""" + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=torch.float16) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=torch.float16) + out = torch.empty((L, M, N), device="cuda", dtype=torch.float16) + + args = ops.GemmArguments(A=A, B=B, out=out, accumulator_type=torch.float32) + operators = ops.get_operators( + args, + metadata_filter=lambda m: _preferred_cluster_filter( + m, StaticPersistentTileSchedulerMetadata + ), + target_sm=device_or_env_target_sm(), + ) + + assert len(operators) > 0, "No preferred-cluster static operator matched the filter" + operator = operators[0] + logger.debug(f"Picked operator: {operator.metadata.operator_name}") + logger.debug(f"Operator metadata:\n{pformat(operator.metadata)}") + + assert operator.supports(args) + operator.run(args) + + reference = A.float() @ B.float() + assert_close_with_reference_conversion( + out, + reference, + out.dtype, + atols=0.0, + rtols=0.0, + msg=f"Static preferred-cluster GEMM failed: M={M}, N={N}, K={K}, L={L}", + ) + + +# ============================================================================= +# API Behaviors: Operator Discovery and Metadata Filtering +# ============================================================================= + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_preferred_cluster_gemm_discovery(): + """Verify the Operator class is registered and discoverable via metadata filter.""" + M, N, K, L = 512, 512, 128, 1 + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=torch.float16) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=torch.float16) + out = torch.empty((L, M, N), device="cuda", dtype=torch.float16) + + args = ops.GemmArguments(A=A, B=B, out=out, accumulator_type=torch.float32) + + operators = ops.get_operators(args, metadata_filter=_preferred_cluster_filter) + assert len(operators) > 0, ( + "Expected at least one preferred-cluster Operator variant to be registered" + ) + + for operator in operators: + design = operator.metadata.design + assert isinstance(design, Sm100DesignMetadata) + assert design.use_fallback_cluster() + assert design.fallback_cluster_shape is not None + assert design.fallback_cluster_shape != design.cluster_shape + assert isinstance( + design.tile_scheduler, + ( + CLCDynamicPersistentTileSchedulerMetadata, + StaticPersistentTileSchedulerMetadata, + ), + ) + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_preferred_cluster_both_schedulers_discoverable(): + """Both CLC and static scheduler preferred-cluster variants are registered.""" + M, N, K, L = 512, 512, 128, 1 + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=torch.float16) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=torch.float16) + out = torch.empty((L, M, N), device="cuda", dtype=torch.float16) + + args = ops.GemmArguments(A=A, B=B, out=out, accumulator_type=torch.float32) + + clc_ops = ops.get_operators( + args, + metadata_filter=lambda m: _preferred_cluster_filter( + m, CLCDynamicPersistentTileSchedulerMetadata + ), + target_sm=device_or_env_target_sm(), + ) + static_ops = ops.get_operators( + args, + metadata_filter=lambda m: _preferred_cluster_filter( + m, StaticPersistentTileSchedulerMetadata + ), + target_sm=device_or_env_target_sm(), + ) + + assert len(clc_ops) > 0, ( + "Expected at least one preferred-cluster CLC variant to be registered" + ) + assert len(static_ops) > 0, ( + "Expected at least one preferred-cluster static variant to be registered" + ) + + for operator in clc_ops: + assert isinstance( + operator.metadata.design.tile_scheduler, + CLCDynamicPersistentTileSchedulerMetadata, + ) + + for operator in static_ops: + assert isinstance( + operator.metadata.design.tile_scheduler, + StaticPersistentTileSchedulerMetadata, + ) + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_preferred_cluster_gemm_not_available_on_unsupported_cc(): + """No preferred-cluster operators are offered for pre-Blackwell targets.""" + M, N, K, L = 256, 256, 128, 1 + A = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=torch.float16) + B = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=torch.float16) + out = torch.empty((L, M, N), device="cuda", dtype=torch.float16) + + args = ops.GemmArguments(A=A, B=B, out=out, accumulator_type=torch.float32) + + # SM90 (Hopper) does not have the preferred-cluster launch mechanism. + operators = ops.get_operators( + args, metadata_filter=_preferred_cluster_filter, target_sm="90" + ) + assert len(operators) == 0, "Expected no preferred-cluster operators for SM90" + + +# ============================================================================= +# Compile-and-Reuse Workflow +# ============================================================================= + + +@pytest.mark.skipif( + not device_or_env_supports("100f"), + reason="Requires compute capability 100 and to be compiled with sm_100a or sm_100f", +) +def test_preferred_cluster_gemm_compile_and_reuse(fixture_toggle_tvm_ffi): + """Compile once, run twice with different input data. + + Exercises the same compiled artifact across multiple ``run`` invocations, + which is the typical inference-serving pattern. + """ + M, N, K, L = 512, 512, 128, 1 + A1 = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=torch.float16) + B1 = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=torch.float16) + out1 = torch.empty((L, M, N), device="cuda", dtype=torch.float16) + + args1 = ops.GemmArguments(A=A1, B=B1, out=out1, accumulator_type=torch.float32) + operators = ops.get_operators(args1, metadata_filter=_preferred_cluster_filter) + assert len(operators) > 0 + operator = operators[0] + + compiled_artifact = operator.compile(args1) + operator.run(args1, compiled_artifact=compiled_artifact) + assert_close_with_reference_conversion( + out1, + A1.float() @ B1.float(), + out1.dtype, + atols=0.0, + rtols=0.0, + msg="first run after compile failed", + ) + + # Reuse the compiled artifact with different inputs of the same shape. + A2 = torch.randint(-1, 2, (L, M, K), device="cuda", dtype=torch.float16) + B2 = torch.randint(-1, 2, (L, K, N), device="cuda", dtype=torch.float16) + out1 = torch.empty((L, M, N), device="cuda", dtype=torch.float16) + args2 = ops.GemmArguments(A=A2, B=B2, out=out1, accumulator_type=torch.float32) + + operator.run(args2, compiled_artifact=compiled_artifact) + assert_close_with_reference_conversion( + out1, + A2.float() @ B2.float(), + out1.dtype, + atols=0.0, + rtols=0.0, + msg="reused compile artifact failed on second inputs", + ) diff --git a/operators/test/integration/test_tgv_gemm.py b/operators/test/integration/test_tgv_gemm.py new file mode 100644 index 000000000..5ed5ca2ec --- /dev/null +++ b/operators/test/integration/test_tgv_gemm.py @@ -0,0 +1,89 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""Integration tests for the TGV low-latency GEMM Operator. + +TGV is optimized for narrow-batch decode inference (small M). +It requires specific layouts: + A (activation): row-major (K-contiguous) + B (weight): col-major (K-contiguous) + C (output): row-major (N-contiguous) +""" + +import logging + +import torch + +import cutlass.operators as ops +from cutlass.operators.providers.cutedsl.gemm.sm100_tgv_gemm import TgvGemmOperator +from cutlass.operators.utils.device import device_or_env_target_sm + +torch.manual_seed(2026) +logger = logging.getLogger(__name__) + + +def _make_col_major_B(K, N, L=None, dtype=torch.float16): + """Create a col-major B tensor (K-contiguous) as expected by TGV. + + Allocates as (N, K) row-major then transposes, so K has stride 1. + """ + if L is not None: + B = torch.randint(-1, 2, (L, N, K), device="cuda", dtype=dtype) + return B.transpose(-1, -2) # (L, K, N) with K stride=1 + B = torch.randint(-1, 2, (N, K), device="cuda", dtype=dtype) + return B.t() # (K, N) with K stride=1 + + +def _get_tgv_operators(args): + """Get only TGV operators from the manifest.""" + + def metadata_filter(metadata): + return metadata.operator_class is TgvGemmOperator + + return ops.get_operators( + args, + target_sm=device_or_env_target_sm(), + metadata_filter=metadata_filter, + ) + + +def test_tgv_not_returned_for_row_major_B(): + """TGV Operator should NOT be returned when B is row-major (N-contiguous).""" + M, K, N = 16, 256, 128 + dtype = torch.float16 + A = torch.randint(-1, 2, (M, K), device="cuda", dtype=dtype) + # Standard row-major B — NOT K-contiguous + B = torch.randint(-1, 2, (K, N), device="cuda", dtype=dtype) + D = torch.empty((M, N), device="cuda", dtype=dtype) + + args = ops.GemmArguments(A, B, D, torch.float32) + tgv_operators = _get_tgv_operators(args) + + assert len(tgv_operators) == 0, ( + f"Expected no TGV operators for row-major B, got {len(tgv_operators)}" + ) diff --git a/operators/test/integration/test_tvm_ffi.py b/operators/test/integration/test_tvm_ffi.py new file mode 100644 index 000000000..53e940521 --- /dev/null +++ b/operators/test/integration/test_tvm_ffi.py @@ -0,0 +1,230 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +import cuda.bindings.driver as cuda +import pytest +import torch +from torch.cuda import current_stream + +import cutlass.operators as ops + +from test_utils import assert_close_with_reference_conversion + +# Skip this entire module if tvm_ffi is not available, rather than producing a +# collection-time ImportError. +tvm_ffi = pytest.importorskip("tvm_ffi") + + +def _benchmark(label, code, n_iterations): + """Helper function to benchmark code execution.""" + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + + # Warmup + for _ in range(n_iterations): + code() + + start.record() + for _ in range(n_iterations): + code() + end.record() + torch.cuda.synchronize() + + avg_time = start.elapsed_time(end) / n_iterations + print(f"[{label:20}] avg of {n_iterations} iterations: {avg_time:1.3e} ms") + return avg_time + + +@pytest.mark.parametrize("M, N, K", [(256, 512, 1024)]) +@pytest.mark.parametrize( + "ab_dtype, c_dtype, accumulator_type", + [(torch.float16, torch.float16, torch.float32)], +) +def test_option_takes_effect( + M: int, + N: int, + K: int, + ab_dtype: torch.dtype, + c_dtype: torch.dtype, + accumulator_type: torch.dtype, + fixture_toggle_tvm_ffi, +): + """Test refcheck passes and compiled object type matches TVM FFI setting.""" + A = torch.randint(-1, 2, (M, K), device="cuda").to(ab_dtype) + B = torch.randint(-1, 2, (K, N), device="cuda").to(ab_dtype) + D = torch.empty((M, N), device="cuda").to(c_dtype) + + args = ops.GemmArguments(A=A, B=B, out=D, accumulator_type=accumulator_type) + operators = ops.get_operators(args, target_sm="80") + + assert len(operators) > 0 + operator = operators[0] + + # Compile and verify the compiled object type + assert operator.supports(args) + compiled_artifact = operator.compile(args) + + use_tvm_ffi = ops.GlobalOptions().use_tvm_ffi + obj_to_check = compiled_artifact.compiled_obj + if use_tvm_ffi: + obj_to_check = tvm_ffi.convert(obj_to_check) + is_tvm_ffi_function = isinstance(obj_to_check, tvm_ffi.Function) + assert is_tvm_ffi_function == use_tvm_ffi, ( + "Compiled object should be tvm_ffi.Function if and only if TVM FFI is enabled." + ) + + # Run and verify results + operator.run( + args, + compiled_artifact=compiled_artifact, + stream=current_stream(), + assume_supported_args=True, + ) + + reference = A @ B + assert_close_with_reference_conversion(D, reference, D.dtype) + + +@pytest.mark.parametrize("M, N, K", [(256, 512, 1024)]) +@pytest.mark.parametrize( + "ab_dtype, c_dtype, accumulator_type", + [(torch.float16, torch.float16, torch.float32)], +) +def test_cuda_graph( + M: int, + N: int, + K: int, + ab_dtype: torch.dtype, + c_dtype: torch.dtype, + accumulator_type: torch.dtype, + fixture_enable_tvm_ffi, +): + """Test that CUDA graphs work correctly with TVM FFI enabled.""" + A = torch.randint(-1, 2, (M, K), device="cuda").to(ab_dtype) + B = torch.randint(-1, 2, (K, N), device="cuda").to(ab_dtype) + D = torch.empty((M, N), device="cuda").to(c_dtype) + + args = ops.GemmArguments(A=A, B=B, out=D, accumulator_type=accumulator_type) + operators = ops.get_operators(args, target_sm="80") + + assert len(operators) > 0 + operator = operators[0] + + # Compile the Operator (must be done outside CUDA graph capture) + assert operator.supports(args) + compiled_artifact = operator.compile(args) + + # Capture CUDA graph + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + operator.run( + args, + compiled_artifact=compiled_artifact, + stream=current_stream(), + assume_supported_args=True, + ) + + # Zero output and replay the graph + D.zero_() + g.replay() + torch.cuda.synchronize() + + # Verify results + reference = A @ B + assert_close_with_reference_conversion( + D, reference, D.dtype, msg="Refcheck failed with CUDA graph + TVM FFI!" + ) + + +@pytest.mark.parametrize("M, N, K", [(256, 512, 1024)]) +@pytest.mark.parametrize( + "ab_dtype, c_dtype, accumulator_type", + [(torch.float16, torch.float16, torch.float32)], +) +@pytest.mark.parametrize("n_iterations", [50]) +def test_speedup_comparison( + M: int, + N: int, + K: int, + ab_dtype: torch.dtype, + c_dtype: torch.dtype, + accumulator_type: torch.dtype, + n_iterations: int, +): + """Compare performance between TVM FFI enabled and disabled to measure speedup.""" + print() + + A = torch.randint(-1, 2, (M, K), device="cuda").to(ab_dtype) + B = torch.randint(-1, 2, (K, N), device="cuda").to(ab_dtype) + D = torch.empty((M, N), device="cuda").to(c_dtype) + + # Compile and run with TVM FFI enabled + ops.GlobalOptions().use_tvm_ffi = True + args_with_tvm_ffi = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type + ) + operators = ops.get_operators(args_with_tvm_ffi, target_sm="80") + + assert len(operators) > 0 + operator = operators[0] + + compiled_artifact_with_tvm_ffi = operator.compile(args_with_tvm_ffi) + + avg_time_with_tvm_ffi = _benchmark( + "Run with TVM-FFI", + lambda: operator.run( + args_with_tvm_ffi, + compiled_artifact=compiled_artifact_with_tvm_ffi, + stream=current_stream(), + assume_supported_args=True, + ), + n_iterations, + ) + + # Compile and run with TVM FFI disabled + ops.GlobalOptions().use_tvm_ffi = False + args_without_tvm_ffi = ops.GemmArguments( + A=A, B=B, out=D, accumulator_type=accumulator_type + ) + compiled_artifact_without_tvm_ffi = operator.compile(args_without_tvm_ffi) + + stream = cuda.CUstream(current_stream().cuda_stream) + avg_time_without_tvm_ffi = _benchmark( + "Run without TVM-FFI", + lambda: operator.run( + args_without_tvm_ffi, + compiled_artifact=compiled_artifact_without_tvm_ffi, + stream=stream, + assume_supported_args=True, + ), + n_iterations, + ) + + speedup = avg_time_without_tvm_ffi / avg_time_with_tvm_ffi + print(f"Speedup with TVM-FFI: {speedup:.3f}x") diff --git a/operators/test/test_utils/__init__.py b/operators/test/test_utils/__init__.py new file mode 100644 index 000000000..a2a287072 --- /dev/null +++ b/operators/test/test_utils/__init__.py @@ -0,0 +1,40 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from . import allocation, gemm # noqa: F401 +from .common import get_test_level, operators_for_classes, skip_if_unsupported_target +from .reference_check import assert_close_with_reference_conversion + +__all__ = [ + "allocation", + "gemm", + "get_test_level", + "operators_for_classes", + "skip_if_unsupported_target", + "assert_close_with_reference_conversion", +] diff --git a/operators/test/test_utils/allocation.py b/operators/test/test_utils/allocation.py new file mode 100644 index 000000000..3ef03b475 --- /dev/null +++ b/operators/test/test_utils/allocation.py @@ -0,0 +1,511 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import random +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Literal + +import torch + +import cutlass +import cutlass.torch + +import cutlass.operators as ops +from cutlass.operators.metadata import ( + DenseTensorConstraints, + OperandConstraints, + ScaledOperandConstraints, +) +from cutlass.operators.utils.common import ceil_div, round_up +from cutlass.operators.utils.dtype import torch_storage_packing_factor +from cutlass.utils.gemm.tensor_utils import create_scale_factor_tensor + + +class RNGStateManager: + """ + Helper for saving and restoring the RNG state. This is useful to make sure + that an RNG state is set to a specific value for a given block of code and + does not affect calling operations. + """ + + def __init__(self, seed: int = None): + self._seed = seed + self._rng_state = None + self._cuda_rng_state = None + + def __enter__(self): + if self._seed is not None: + self._rng_state = torch.get_rng_state() + self._cuda_rng_state = torch.cuda.get_rng_state() + torch.manual_seed(self._seed) + + def __exit__(self, exc_type, exc_value, traceback): + if self._seed is not None: + torch.set_rng_state(self._rng_state) + torch.cuda.set_rng_state(self._cuda_rng_state) + + +################################################################################ +# Initialization modes +################################################################################ + + +@dataclass +class InitializationMode(ABC): + """ + An initialization mode for a tensor. + """ + + def create( + self, dtype, shape: tuple[int, ...], seed: int = 2025, device="cuda" + ) -> torch.Tensor: + """ + Creates a tensor of the given datatype, shape, and device. Simple wrapper around the + child-defined _create method that ensures that the RNG state is set to the given seed. + + :param dtype: The datatype of the tensor + :type dtype: cutlass.Numeric + :param shape: The shape of the tensor + :type shape: tuple[int, ...] + :param seed: The seed to use for the RNG state + :type seed: int + :param device: The device to use for the tensor + :type device: str + + :return: A tensor of the given datatype, shape, and device + :rtype: torch.Tensor + """ + with RNGStateManager(seed): + return self._create(dtype, shape, device) + + @abstractmethod + def _create(self, dtype, shape: tuple[int, ...], device="cuda") -> torch.Tensor: + """ + Creates a tensor of the given datatype, shape, and device. + """ + + +@dataclass +class RandomIntInitializationMode(InitializationMode): + """ + An initialization mode for a tensor that is initialized with random integers + in range [min_val, max_val] and scaled by the global scale factor. + """ + + min_val: int + max_val: int + global_scale: float = 1 + + def _create(self, dtype, shape: tuple[int, ...], device="cuda") -> torch.Tensor: + """ + Creates a tensor of the given datatype, shape, and device. + Also scales the tensor by the global scale factor if provided. + """ + if dtype == torch.float4_e2m1fn_x2: + # Packed FP4: each byte holds 2 FP4 values (low + high nibble). + # Use the same valid byte patterns as CuTe DSL's + # create_gemm_tensor_torch — each byte encodes a pair from {0, 1, -1}. + _FP4_BYTE_PATTERNS = torch.tensor( + [0x00, 0x02, 0x0A, 0x20, 0x22, 0x2A, 0xA0, 0xA2, 0xAA], + dtype=torch.uint8, + device=device, + ) + return _FP4_BYTE_PATTERNS[ + torch.randint(0, len(_FP4_BYTE_PATTERNS), shape, device=device) + ].view(torch.float4_e2m1fn_x2) + scaled_tensor = self.global_scale * torch.randint( + self.min_val, self.max_val, shape, device=device + ) + return scaled_tensor.to(dtype) + + @classmethod + def from_dtypes( + cls, dtype: cutlass.Numeric, out_dtype: cutlass.Numeric + ) -> "RandomIntInitializationMode": + """ + Creates a RandomIntInitializationMode from the given datatypes. + + Range is determined based on the bitwidth of the `dtype` and `out_dtype` -- + the lower the bitwidth of these types, the lower the range of values to avoid overflow + and rounding error. + + Args: + dtype (cutlass.Numeric): The datatype of the tensor + out_dtype (cutlass.Numeric): The datatype of the output tensor + + Returns: + A RandomIntInitializationMode + + Raises: + ValueError: If the dtype is not supported + """ + if dtype.is_float: + # [-2, -1, 0, 1, 2] + min_val = -2 + max_val = 3 # exclusive + sf = 1 + + if dtype.width <= 8 or out_dtype.width <= 8: + # [-1, -0.5, 0, 0.5, 1] + sf = 0.5 + + return cls(min_val=min_val, max_val=max_val, global_scale=sf) + elif dtype.is_integer: + min_val = -2 + max_val = 3 + + if dtype.width <= 8 or out_dtype.width <= 8: + min_val = -1 + max_val = 2 # exclusive + + if dtype in {cutlass.Uint8, cutlass.Uint16, cutlass.Uint32, cutlass.Uint64}: + min_val = 0 + + return cls(min_val=min_val, max_val=max_val) + else: + raise ValueError(f"Unsupported dtype: {dtype}") + + +@dataclass +class EmptyInitializationMode(InitializationMode): + """ + An initialization mode for a tensor that is initialized with empty values. + """ + + def _create(self, dtype, shape: tuple[int, ...], device="cuda") -> torch.Tensor: + """ + Creates a tensor of the given datatype, shape, and device. + """ + return torch.empty(shape, dtype=dtype, device=device) + + +################################################################################ +# MemPool +################################################################################ + + +class MemPool: + """ + A pool of tensors used for test cases. The goals of the pool are the following: + - Reduce reallocations of tensors across test cases + - Reduce reinitialization of tensors across test cases + - Allow for returning tensors with arbitrary starting offsets (to stress pointer alignment) + - Allow for returning non-compact tensors + + Toward these goals, the pool maintains a set of tensors, with one tensor corresponding to a + combination of datatype and InitializationMode. + + Callers request a tensor of a given datatype, shape, and InitializationMode. + + If the pool contains a tensor that is at least as large as is needed by the caller and + with the same datatype and InitializationMode, that tensor is returned by reusing + the existing tensor. This is done only for InitializationModes that are position independent + -- if a caller requests a tensor to be initialized sequentially starting from X, returning + a subset of that tensor may not satisfy the caller's requirements. In these cases, a new + tensor may be reinitialized. + + Tensors are intended to be under the ownership of the MemPool. Explicitly deallocating + memory backing the pool is expected to lead to an error. + """ + + def __init__(self, device: Literal["cpu", "cuda"] = "cuda"): + """ + :param device: The torch device to allocate pool tensors on. Defaults to "cuda". + :type device: str + """ + self.device = device + self.tensors = {} + + def get( + self, + dtype: cutlass.Numeric, + shape: tuple[int, ...], + stride: tuple[int, ...], + initialization_mode: InitializationMode, + random_offset_divisibility: int = None, + ) -> torch.Tensor: + """ + Gets a tensor of the given datatype, shape, and stride, and initialized with the given initialization mode. + If `random_offset_divisibility` is provided, the tensor is offset by a random amount + (divisible by `random_offset_divisibility`) to stress dependence on higher pointer alignments + than stated as needed by the Operator. + + :param dtype: The datatype of the tensor + :type dtype: cutlass.Numeric + :param shape: The shape of the tensor + :type shape: tuple[int, ...] + :param stride: The stride of the tensor + :type stride: tuple[int, ...] + :param initialization_mode: The initialization mode of the tensor + :type initialization_mode: InitializationMode + :param random_offset_divisibility: The divisibility of the random offset + :type random_offset_divisibility: int + :return: A tensor of the given datatype, shape, and stride, and initialized with the given initialization mode + :rtype: torch.Tensor + """ + if random_offset_divisibility is not None and random_offset_divisibility <= 0: + raise ValueError( + f"Expected random_offset_divisibility to be greater than 0, got {random_offset_divisibility}" + ) + + if len(shape) != len(stride): + raise ValueError( + f"Expected shape and stride to have the same length, got {len(shape)} and {len(stride)}" + ) + + size_required = sum(shape[i] * stride[i] for i in range(len(shape))) + if size_required == 0: + raise ValueError( + f"Expected size_required to be greater than 0, got {size_required}" + ) + + key = str(dtype) + str(initialization_mode) + torch_dtype = cutlass.torch.dtype(dtype) + tensor = self.tensors.get(key, torch.tensor([], dtype=torch_dtype)) + + if tensor.numel() < 2 * size_required: + tensor = initialization_mode.create( + torch_dtype, (2 * size_required,), device=self.device + ) + self.tensors[key] = tensor + + if random_offset_divisibility is not None: + offset = random.randrange( + random_offset_divisibility, + size_required - 1, + random_offset_divisibility, + ) + else: + offset = 0 + return torch.as_strided(tensor[offset : offset + size_required], shape, stride) + + +################################################################################ +# Tensor creation from metadata +################################################################################ + + +def dense_tensor_from_metadata( + mempool: MemPool, + attr: DenseTensorConstraints, + shape: tuple[int, ...], + tensor_name: str, + initialization_mode: InitializationMode, + randomly_offset: bool = True, +) -> ops.DenseTensor: + """ + Creates a dense tensor from metadata. + + :param mempool: The memory pool to use + :type mempool: MemPool + :param attr: The constraints of the tensor + :type attr: cutlass.operators.metadata.DenseTensorConstraints + :param shape: The shape of the tensor + :type shape: tuple[int, ...] + :param tensor_name: The name of the tensor + :type tensor_name: str + :param initialization_mode: The initialization mode to use + :type initialization_mode: InitializationMode + :param randomly_offset: Whether to randomly offset the tensor by its divisibility factor + :type randomly_offset: bool + + :return: A dense tensor + :rtype: cutlass.operators.arguments.DenseTensor + """ + if not isinstance(attr, DenseTensorConstraints): + raise ValueError(f"Expected DenseTensorConstraints, got {type(attr)}") + + # Plan to allow for non-compact strides in the future. + # For now, assume that one of the final two modes has stride 1, + # that the other of the final two modes has the next lowest stride, + # and that remaining modes have strides that increase in right-to-left + # order (which matches the default in torch). + if attr.stride is None or len(attr.stride) == 1: + stride = [ + 1, + ] + next_stride = shape[-1] + else: + # Build up stride starting from the final modes + stride = [shape[-1], 1] if attr.stride[-2:][1] == 1 else [1, shape[-2]] + next_stride = shape[-1] * shape[-2] + + remaining_modes = len(shape) - len(stride) + for i in range(remaining_modes, 0, -1): + stride.insert(0, next_stride) + next_stride *= shape[i] + + stride = tuple(stride) + + tensor = mempool.get( + attr.dtype, + shape, + stride, + initialization_mode, + attr.divisibility if randomly_offset else None, + ) + return ops.DenseTensor(tensor) + + +def scaled_tensor_from_metadata( + mempool: MemPool, + attr: ScaledOperandConstraints, + shape: tuple[int, int, int], + tensor_name: str, + initialization_mode: InitializationMode, + randomly_offset: bool = True, +) -> ops.ScaledOperand: + """ + Creates a scaled tensor from metadata. + + :param mempool: The memory pool to use + :type mempool: MemPool + :param attr: The constraints of the tensor + :type attr: cutlass.operators.metadata.ScaledOperandConstraints + :param shape: The shape of the tensor + :type shape: tuple[int, int, int] + :param tensor_name: The name of the tensor + :type tensor_name: str + :param initialization_mode: The initialization mode to use + :type initialization_mode: InitializationMode + :param randomly_offset: Whether to randomly offset the tensor by its divisibility factor + :type randomly_offset: bool + + :return: A scaled tensor + :rtype: cutlass.operators.arguments.ScaledOperand + """ + if len(shape) != 3: + raise ValueError(f"Expected shape to be of length 3, got {len(shape)}") + + if not isinstance(attr, ScaledOperandConstraints): + raise ValueError(f"Expected ScaledOperandConstraints, got {type(attr)}") + + if tensor_name not in ["A", "B"]: + raise ValueError(f"Expected tensor_name to be 'A' or 'B', got {tensor_name}") + + # Adjust shape for packed subbyte types: float4_e2m1fn_x2 packs 2 FP4 values + # per byte, so the physical K dimension is K // 2. + quantized_shape = shape + if attr.quantized.dtype == cutlass.Float4E2M1FN: + if tensor_name == "A": + quantized_shape = (shape[0], shape[1], shape[2] // 2) + else: + quantized_shape = (shape[0], shape[1] // 2, shape[2]) + + quantized = create( + mempool, attr.quantized, quantized_shape, tensor_name, initialization_mode, randomly_offset + ) + + # ScaledOperand.numel_scale takes the quantized shape as (L, outer, K), where + # outer=M for an A-side scale and outer=N for a B-side scale. The B quantized + # tensor is stored as (L, K, N), so we permute the inner dims to match. + if tensor_name == "A": + (L, outer, K) = shape + else: + (L, K, outer) = shape + + sf_vec_size = ops.ScaleMode.numel(attr.mode) + + if attr.swizzle == ops.ScaleSwizzleMode.Swizzle32x4x4: + # Use CuTe DSL's create_scale_factor_tensor to produce the packed + # on-device layout (sf_gpu). The emulated reference in gemm.py + # calls unpack_scale_factors directly on the scale tensor from args, + # so no pre-unpacked data needs to be threaded through here. + _sf_ref, _sf_cute, sf_gpu = create_scale_factor_tensor( + outer, K, L, sf_vec_size, attr.scale.dtype + ) + scale = ops.DenseTensor(sf_gpu.flatten()) + return ops.ScaledOperand(quantized, scale, attr.mode, attr.swizzle) + + total_tensor_size = ops.ScaledOperand.numel_scale( + (L, outer, K), attr.mode, attr.swizzle + ) + + # Block-scaled GEMM hardware treats scale factors as non-negative + # magnitudes (consistent with CuTe DSL conventions). + if ( + isinstance(initialization_mode, RandomIntInitializationMode) + and initialization_mode.min_val < 0 + ): + raise ValueError( + f"Scale factor initialization requires non-negative min_val, " + f"got {initialization_mode.min_val}. Block-scaled GEMM hardware " + f"treats scale factors as non-negative magnitudes." + ) + + scale = create( + mempool, + attr.scale, + (total_tensor_size,), + tensor_name, + initialization_mode, + randomly_offset, + ) + return ops.ScaledOperand(quantized, scale, attr.mode, attr.swizzle) + + +def create( + mempool: MemPool, + attr: OperandConstraints, + shape: tuple[int, int, int], + tensor_name: str, + initialization_mode: InitializationMode, + randomly_offset: bool = True, +) -> ops.Operand: + """ + Creates a tensor from metadata. + + :param mempool: The memory pool to use + :type mempool: MemPool + :param attr: The constraints of the tensor + :type attr: cutlass.operators.metadata.OperandConstraints + :param shape: The shape of the tensor + :type shape: tuple[int, int, int] + :param tensor_name: The name of the tensor + :type tensor_name: str + :param initialization_mode: The initialization mode to use + :type initialization_mode: InitializationMode + :param randomly_offset: Whether to randomly offset the tensor by its divisibility factor + :type randomly_offset: bool + + :return: A tensor + :rtype: cutlass.operators.arguments.Operand + """ + # Perform exact type matches to avoid issues in case of inheritance + match attr: + case DenseTensorConstraints(): + return dense_tensor_from_metadata( + mempool, attr, shape, tensor_name, initialization_mode, randomly_offset + ) + case ScaledOperandConstraints(): + return scaled_tensor_from_metadata( + mempool, attr, shape, tensor_name, initialization_mode, randomly_offset + ) + case _: + raise ValueError(f"Unhandled operand constraints type: {type(attr)}") diff --git a/operators/test/test_utils/common.py b/operators/test/test_utils/common.py new file mode 100644 index 000000000..674ca0010 --- /dev/null +++ b/operators/test/test_utils/common.py @@ -0,0 +1,214 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +import os +import random +from collections import defaultdict +from typing import TYPE_CHECKING + +import pytest + +import cutlass.operators as ops +from cutlass.operators.utils.device import device_or_env_target_sm + +if TYPE_CHECKING: + from collections.abc import Callable, Iterable + +# Environment variable used to propagate the test level ("L0" / "L1" / "L2") +# to helpers invoked during test collection +_TEST_LEVEL_ENV_VAR = "CUTLASS_OPERATORS_TEST_LEVEL" + + +def get_test_level() -> str: + """Return the test level set via ``CUTLASS_OPERATORS_TEST_LEVEL`` (defaults to ``"L0"``).""" + return os.environ.get(_TEST_LEVEL_ENV_VAR, "L0") + + +def skip_if_unsupported_target(operator_targets: Iterable[ops.TargetSm | str]) -> None: + """ + Skips the test if the current device or CUTE_DSL_ARCH does not support the requested Operator targets. + + :param operator_targets: Set of Operator targets to check compatibility for + :type operator_targets: Iterable[TargetSm | str] + """ + operator_targets = {ops.TargetSm.ensure(sm) for sm in operator_targets} + + device_sm = device_or_env_target_sm() + if not device_sm.supports_operators_from(operator_targets): + pytest.skip( + f"Current device {device_sm} does not support requested Operator targets {operator_targets}" + ) + + + +def _operators_for_class( + operator_class: type[ops.Operator], + extra_filter: Callable | None = None, + seed: int | None = None, + test_level: str | None = None, + num_sample: int | None = None, +) -> list[ops.Operator]: + """ + Returns a list of operators to test of type `operator_class`. + + For L2 tests, the full set of possible operators are returned. + + For L0 or L1 tests, a random sample of operators are returned. The sample is chosen so that there + is at least one operator for every unique operands metadata and for every unique design metadata. + L0 and L1 use distinct default seeds so that their sampled operator sets differ, + maximising combined coverage when both levels are run. + + :param operator_class: The class of Operator to test + :type operator_class: type[cutlass.operators.Operator] + :param extra_filter: An extra filter function to apply to the metadata filter + :type extra_filter: Callable | None + :param seed: Override for the RNG seed. When ``None``, the level-specific + default (``L0=2025``, ``L1=2026``) is used. + :type seed: int | None + :param num_sample: Override for the per-group sample count at L0/L1. When + ``None``, the level-specific default (``L0=1``, ``L1=3``) is used. + :type num_sample: int | None + + :return: A list of operators to test + :rtype: list[cutlass.operators.Operator] + """ + if extra_filter is None: + + def extra_filter(x): + return True + + def metadata_filter(metadata): + return metadata.operator_class is operator_class and extra_filter(metadata) + + operators = ops.get_operators( + args=None, + metadata_filter=metadata_filter, + ) + + if len(operators) == 0: + raise ValueError(f"No operators found for Operator class {operator_class}") + + if test_level is None: + # Read the test level from the ``CUTLASS_OPERATORS_TEST_LEVEL`` env + # var (defaults to ``"L0"`` when unset). + test_level = os.environ.get("CUTLASS_OPERATORS_TEST_LEVEL", "L0") + + valid_test_levels = ["L0", "L1", "L2"] + if test_level not in valid_test_levels: + raise ValueError( + f"Invalid test level: {test_level}. Must be one of {valid_test_levels}." + ) + + if test_level == "L2": + return operators + + # Seed the RNG with a level-specific default when the caller didn't + # override, so L0 and L1 explore disjoint regions of the Operator space. + seed = {"L0": 2025, "L1": 2026}[test_level] if seed is None else seed + random.seed(seed) + + # L0 / L1 tests. Group operators by OperandsMetadata and DesignMetadata + # so that we can sample from distinct categories + operators_by_operands: defaultdict[str, list] = defaultdict(list) + operators_by_design: defaultdict[str, list] = defaultdict(list) + for operator in operators: + operators_by_operands[str(operator.metadata.operands)].append(operator) + operators_by_design[str(operator.metadata.design)].append(operator) + + # Sample operators from each unique operand metadata. Caller-supplied + # ``num_sample`` wins; otherwise fall back to the per-level default. + num_sample = {"L0": 1, "L1": 3}[test_level] if num_sample is None else num_sample + sampled_operators = [] + for operand_operators in operators_by_operands.values(): + sampled_operators.extend( + random.sample(operand_operators, min(num_sample, len(operand_operators))) + ) + + # Sample operators from each unique design metadata. First remove design metadata that + # has already been covered by operators previously sampled + designs_already_sampled = { + str(operator.metadata.design) for operator in sampled_operators + } + for design in designs_already_sampled: + assert design in operators_by_design + del operators_by_design[design] + + for design_operators in operators_by_design.values(): + sampled_operators.extend( + random.sample(design_operators, min(num_sample, len(design_operators))) + ) + + return sampled_operators + + +def operators_for_classes( + operator_classes: Iterable[type[ops.Operator]], + extra_filters: dict[type[ops.Operator], Callable] | None = None, + seed: int | None = None, + test_level: str | None = None, + num_sample: int | None = None, +) -> list[ops.Operator]: + """ + Takes a list of Operator classes and returns instances of those classes with sufficient + test coverage for the given test level. + + :param operator_classes: A list of Operator classes to test + :type operator_classes: Iterable[type[cutlass.operators.Operator]] + :param extra_filters: A dictionary of extra filter functions for each Operator class + :type extra_filters: dict[type[cutlass.operators.Operator], Callable] | None + :param seed: Override for the RNG seed. When ``None``, a level-specific default is used + (L0 and L1 use distinct seeds to cover disjoint Operator subsets). + :type seed: int | None + :param test_level: The test level ("L0", "L1", "L2") + :type test_level: str + :param num_sample: Override for the per-group sample count at L0/L1. When + ``None``, the level-specific default (``L0=1``, ``L1=3``) is used. + :type num_sample: int | None + + :return: A list of operators to test + :rtype: list[cutlass.operators.Operator] + """ + if extra_filters is None: + extra_filters = {} + + operators = [ + operator + for operator_class in operator_classes + for operator in _operators_for_class( + operator_class, + extra_filters.get(operator_class), + seed, + test_level, + num_sample, + ) + ] + + operators.sort(key=lambda x: x.metadata.operator_name) + return operators diff --git a/operators/test/test_utils/gemm.py b/operators/test/test_utils/gemm.py new file mode 100644 index 000000000..9a08cf91e --- /dev/null +++ b/operators/test/test_utils/gemm.py @@ -0,0 +1,403 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import torch + +import cutlass +import cutlass.torch +from cutlass.utils.gemm.tensor_utils import decode_float4e2m1fn, unpack_scale_factors + +import cutlass.operators as ops +from cutlass.operators.utils.common import ceil_div +from cutlass.operators.utils.device import device_or_env_target_sm +from cutlass.operators.utils.dtype import torch_storage_packing_factor + +from . import reference_check + +################################################################################ +# Reference computation +################################################################################ + + +def clamp_to_finite_range(reference: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """ + Clamps the reference tensor to the range of the given dtype if the + dtype is determined to be a narrow precision dtype. + """ + if dtype not in reference_check.get_clamp_torch_dtypes(): + return reference + if dtype.is_floating_point: + info = torch.finfo(dtype) + else: + raise ValueError(f"Unsupported dtype: {dtype}") + return torch.clamp(reference, max=info.max, min=info.min) + + +def dense_gemm_reference( + args: ops.GemmArguments, ref_as_acc_dtype: bool = False +) -> torch.Tensor: + """ + Computes the reference result for a dense GEMM operation. + + Args: + args (cutlass.operators.arguments.GemmArguments): The arguments for the GEMM operation + ref_as_acc_dtype (bool): Option to return reference tensor as accumulator dtype. The default + behavior is to return the reference tensor as the output dtype. + + Returns: + The reference result for the GEMM operation + """ + if not isinstance(args.A, ops.DenseTensor): + raise ValueError(f"Expected args.A to be a DenseTensor, got {type(args.A)}") + if not isinstance(args.B, ops.DenseTensor): + raise ValueError(f"Expected args.B to be a DenseTensor, got {type(args.B)}") + if not isinstance(args.out, ops.DenseTensor): + raise ValueError(f"Expected args.out to be a DenseTensor, got {type(args.out)}") + + A_tensor = args.A.tensor.runtime_tensor + match A_tensor: + case torch.Tensor(): + B_tensor = args.B.tensor.runtime_tensor + assert isinstance(B_tensor, torch.Tensor) + reference = A_tensor.to(torch.float32) @ B_tensor.to(torch.float32) + if not ref_as_acc_dtype: + # pytorch and numpy, following the IEEE standard, converts to Inf or NaN + # but our kernels use satfinite. To make them the same output we have to + # clamp on narrow precision output dtypes. + torch_dtype = cutlass.torch.dtype(args.out.dtype) + reference = clamp_to_finite_range(reference, torch_dtype).to( + torch_dtype + ) + return reference + case _: + raise NotImplementedError(f"Unsupported tensor type: {type(A_tensor)}") + + +def _decode_fp4_packed(tensor: torch.Tensor, pack_dim: int) -> torch.Tensor: + """Decode a ``float4_e2m1fn_x2`` tensor to float32 via :func:`decode_float4e2m1fn`. + + CuTe DSL's :func:`decode_float4e2m1fn` expects 2x-overallocated ``uint8`` + tensors (one byte per logical element) as produced by + ``create_gemm_tensor_torch``. Operator API tests use PyTorch's + tightly-packed ``float4_e2m1fn_x2`` (two values per byte, ``K_packed = + K / 2``). This adapter flattens, pads to the over-allocated size, decodes, + and reshapes back. + """ + u8 = tensor.view(torch.uint8) + ndim = u8.dim() + + if pack_dim != ndim - 1: + perm = [i for i in range(ndim) if i != pack_dim] + [pack_dim] + u8 = u8.permute(perm).contiguous() + inv_perm = [0] * ndim + for i, p in enumerate(perm): + inv_perm[p] = i + else: + inv_perm = list(range(ndim)) + + total = u8.numel() + padded = torch.zeros(1, 2 * total, 1, dtype=torch.uint8, device=tensor.device) + padded[0, :total, 0] = u8.flatten() + decoded = decode_float4e2m1fn(padded).flatten() + + out_shape = list(u8.shape) + out_shape[-1] *= 2 + return decoded.reshape(out_shape).permute(inv_perm) + + +def _emulated_scaled_dense_gemm_reference( + args: ops.GemmArguments, + ref_as_acc_dtype: bool = False, +) -> torch.Tensor: + """Emulated block-scaled GEMM reference: dequantize, apply scales, matmul. + + Works for any data type / scale factor combination without relying on + ``torch._scaled_mm``. All arithmetic is performed in float32. + + The computation is:: + + A_scaled = A_f32 * expand(SFA) # (L, M, K) + B_scaled = B_f32 * expand(SFB) # (L, K, N) + D = A_scaled @ B_scaled # (L, M, N) + + When block-scaled output is needed in the future, a ``scale_out`` parameter + can be added here without restructuring the computation. + """ + A = args.A.quantized.tensor.runtime_tensor + B = args.B.quantized.tensor.runtime_tensor + scale_A = args.A.scale.tensor.runtime_tensor + scale_B = args.B.scale.tensor.runtime_tensor + out_dtype = cutlass.torch.dtype(args.out.dtype) + + M, N = args.out.shape[-2:] + packing_A = torch_storage_packing_factor(args.A.quantized.dtype) + K = A.shape[-1] * packing_A + is_2d = A.dim() == 2 + L = 1 if is_2d else A.shape[0] + sf_vec_size = ops.ScaleMode.numel(args.A.mode) + + # --- dequantize A and B to float32 --- + # float4_e2m1fn_x2 is a sub-byte type: .to(float32) triggers a CUDA + # device-side assert in PyTorch's fetch_and_cast. Decode via CuTe DSL's + # decode_float4e2m1fn (adapted for PyTorch's tightly-packed format). + # Each operand is checked independently to support mixed FP4 x FP8. + A_f32 = ( + _decode_fp4_packed(A, pack_dim=A.dim() - 1) + if A.dtype == torch.float4_e2m1fn_x2 + else A.to(torch.float32) + ) + B_f32 = ( + _decode_fp4_packed(B, pack_dim=B.dim() - 2) + if B.dtype == torch.float4_e2m1fn_x2 + else B.to(torch.float32) + ) + + if is_2d: + A_f32 = A_f32.unsqueeze(0) # (1, M, K) + B_f32 = B_f32.unsqueeze(0) # (1, K, N) + + # --- obtain dense scale factors --- + # unpack_scale_factors inverts the Swizzle32x4x4 layout directly from + # the scale tensor in args, returning (MN, K, L) float32. + # Permute to (L, MN, K) to match our (L, M, K) / (L, K, N) convention. + device = A_f32.device + sfa_expanded = ( + unpack_scale_factors(scale_A.to(torch.float32), sf_vec_size, M, K, L) + .to(device) + .permute(2, 0, 1)[:, :M, :K] + ) # (L, M, K) + sfb_expanded = ( + unpack_scale_factors(scale_B.to(torch.float32), sf_vec_size, N, K, L) + .to(device) + .permute(2, 0, 1)[:, :N, :K] + ) # (L, N, K) + + # --- apply scale factors and matmul --- + A_scaled = A_f32 * sfa_expanded # (L, M, K) + B_scaled = B_f32 * sfb_expanded.transpose(1, 2) # (L, K, N) + reference = A_scaled @ B_scaled # (L, M, N) + + if is_2d: + reference = reference.squeeze(0) + + if not ref_as_acc_dtype: + reference = clamp_to_finite_range(reference, out_dtype).to(out_dtype) + + return reference + + +def scaled_dense_gemm_reference( + args: ops.GemmArguments, ref_as_acc_dtype: bool = False +) -> torch.Tensor: + """ + Computes the reference result for a scaled dense GEMM operation. + + Args: + args (cutlass.operators.arguments.GemmArguments): The arguments for the GEMM operation + ref_as_acc_dtype (bool): Option to return reference tensor as accumulator dtype. The default + behavior is to return the reference tensor as the output dtype. + + Returns: + The reference result for the scaled dense GEMM operation + """ + + if not isinstance(args.A, ops.ScaledOperand): + raise ValueError(f"Expected args.A to be a ScaledOperand, got {type(args.A)}") + if not isinstance(args.B, ops.ScaledOperand): + raise ValueError(f"Expected args.B to be a ScaledOperand, got {type(args.B)}") + if not isinstance(args.out, ops.DenseTensor): + raise ValueError(f"Expected args.out to be a DenseTensor, got {type(args.out)}") + if not isinstance(args.A.scale, ops.DenseTensor): + raise ValueError( + f"Expected args.A.scale to be a DenseTensor, got {type(args.A.scale)}" + ) + if not isinstance(args.B.scale, ops.DenseTensor): + raise ValueError( + f"Expected args.B.scale to be a DenseTensor, got {type(args.B.scale)}" + ) + + A = args.A.quantized.tensor.runtime_tensor + B = args.B.quantized.tensor.runtime_tensor + scale_A = args.A.scale.tensor.runtime_tensor + scale_B = args.B.scale.tensor.runtime_tensor + acc_dtype = cutlass.torch.dtype(args.accumulator_type) + out_dtype = cutlass.torch.dtype(args.out.dtype) + + if not isinstance(A, torch.Tensor): + raise NotImplementedError(f"Unsupported tensor type: {type(A)}") + if not isinstance(B, torch.Tensor): + raise NotImplementedError(f"Unsupported tensor type: {type(B)}") + if not isinstance(scale_A, torch.Tensor): + raise NotImplementedError(f"Unsupported tensor type: {type(scale_A)}") + if not isinstance(scale_B, torch.Tensor): + raise NotImplementedError(f"Unsupported tensor type: {type(scale_B)}") + if not isinstance(out_dtype, torch.dtype): + raise NotImplementedError(f"Unsupported dtype type: {type(out_dtype)}") + + # torch._scaled_mm supports FP4 only with Float8E4M3FN scale factors + # (matching CuTe DSL's is_emulated_dtype). Configs torch can't handle (e.g. + # FP4 + Float8E8M0FNU, mixed FP4 x FP8) use the emulated reference, which + # supports 2D/3D operands and arbitrary N -- returning here avoids the 3D-only + # N-padding below. Each operand is checked independently for mixed configs. + use_emulated = ( + A.dtype == torch.float4_e2m1fn_x2 and scale_A.dtype != torch.float8_e4m3fn + ) or (B.dtype == torch.float4_e2m1fn_x2 and scale_B.dtype != torch.float8_e4m3fn) + if use_emulated: + return _emulated_scaled_dense_gemm_reference(args, ref_as_acc_dtype) + + # torch._scaled_mm currently requires K and N to be divisible by 16. + # Given that torch also only supports TN layout for now, and data types are + # FP8 or smaller, K divisibility >= (alignment_bytes * 8 // dtype.width) = (16 * 8 // 8) = 16. + # Thus, we only need to pad N. (Operands here are 3D.) + N = B.shape[-1] + padded_N = ceil_div(N, 16) * 16 + if padded_N != N: + N_pad = padded_N - N + if B.dtype == torch.float4_e2m1fn_x2: + # Packed FP4: pad at byte level since float32 round-trip changes shape. + # float4_e2m1fn_x2 and int8 share 1-byte element size, so view is safe. + B_padded = torch.zeros( + (B.shape[0], B.shape[1], padded_N), + dtype=torch.int8, + device=B.device, + ) + B_padded[:, :, :N] = B.view(dtype=torch.int8) + B_padded = B_padded.view(dtype=torch.float4_e2m1fn_x2) + else: + B_padded = ( + torch.nn.functional.pad( + B.transpose(1, 2).to(torch.float32), + (0, 0, 0, N_pad), + mode="constant", + value=0, + ) + .to(B.dtype) + .transpose(1, 2) + ) + else: + B_padded = B + + # torch._scaled_mm does not support batch mode. Iterate through each + # problem in the batch. + L, M, N = A.shape[0], A.shape[1], B.shape[2] + scale_A = scale_A.view(L, -1).contiguous() + scale_B = scale_B.view(L, -1).contiguous() + reference = torch.empty((L, M, padded_N), device=A.device, dtype=acc_dtype) + + for l_idx in range(L): + # Use out type of F32 and then convert to out_dtype due to cuBLAS + # errors occasionally thrown with F8 types. + # See https://github.com/pytorch/pytorch/issues/160816 + try: + reference[l_idx, :, :] = torch._scaled_mm( + A[l_idx, :, :], + B_padded[l_idx, :, :], + scale_a=scale_A[l_idx, :], + scale_b=scale_B[l_idx, :], + out_dtype=acc_dtype, + ) + except RuntimeError: + # cuBLAS may not have an algorithm for this configuration; + # fall back to emulated reference. + return _emulated_scaled_dense_gemm_reference(args, ref_as_acc_dtype) + + if not ref_as_acc_dtype: + reference = clamp_to_finite_range(reference, out_dtype).to(out_dtype) + + # Remove padding from reference + reference = reference[:, :, :N] + + return reference + + +def reference(args: ops.GemmArguments, ref_as_acc_dtype: bool = False) -> torch.Tensor: + """ + Computes the reference result for a GEMM operation. + + Args: + args (cutlass.operators.arguments.GemmArguments): The arguments for the GEMM operation + ref_as_acc_dtype (bool): Option to return reference tensor as accumulator dtype. The default + behavior is to return the reference tensor as the output dtype. + + Returns: + The reference result for the GEMM operation + """ + if ( + isinstance(args.A, ops.DenseTensor) + and isinstance(args.B, ops.DenseTensor) + and isinstance(args.out, ops.DenseTensor) + ): + return dense_gemm_reference(args, ref_as_acc_dtype=ref_as_acc_dtype) + elif ( + isinstance(args.A, ops.ScaledOperand) + and isinstance(args.B, ops.ScaledOperand) + and isinstance(args.out, ops.DenseTensor) + ): + return scaled_dense_gemm_reference(args, ref_as_acc_dtype=ref_as_acc_dtype) + else: + raise ValueError(f"No reference implementation found for {args}") + + +# Byte patterns whose two FP4 (E2M1) nibbles each decode to a small integer in +# {0, +1, -1} (0x0 -> 0.0, 0x2 -> 1.0, 0xA -> -1.0). +_FP4_INT_BYTE_PATTERNS = [0x00, 0x02, 0x0A, 0x20, 0x22, 0x2A, 0xA0, 0xA2, 0xAA] + + +def make_mxfp4_kmajor(rows: int, k: int) -> torch.Tensor: + """Build a K-major, tightly-packed MXFP4 operand with integer values in {-1, 0, 1}. + + Returns a ``(rows, k // 2)`` ``torch.float4_e2m1fn_x2`` CUDA tensor (logical + ``(rows, k)``, K-major). With integer operands and power-of-two scales (see + make_random_pow2_scale every dequantized product and partial sum is + an exact integer in float32, so the kernel and scaled_dense_gemm_reference + accumulate bit-identically. + + Args: + rows (int): The non-K (M or N) dimension size. + k (int): The K dimension size. + + Raises: + ValueError: If ``k`` is odd. + """ + if k % 2 != 0: + raise ValueError(f"K must be even to tightly pack FP4, got {k}") + patterns = torch.tensor(_FP4_INT_BYTE_PATTERNS, dtype=torch.uint8, device="cuda") + idx = torch.randint(0, patterns.numel(), (rows, k // 2), device="cuda") + return patterns[idx].view(torch.float4_e2m1fn_x2) + + +def make_random_pow2_scale(numel: int) -> torch.Tensor: + """Random E8M0 scale factors drawn from the powers of two ``{1, 2, 4}``. + + Power-of-two scales have no fractional bits, so scaling integer operands keeps + every product and partial sum an exact integer in float32. Accumulation is + then order-independent. + """ + lut = torch.tensor([1.0, 2.0, 4.0], device="cuda") + idx = torch.randint(0, lut.numel(), (numel,), device="cuda") + return lut[idx].to(torch.float8_e8m0fnu) diff --git a/operators/test/test_utils/reference_check.py b/operators/test/test_utils/reference_check.py new file mode 100644 index 000000000..072039a85 --- /dev/null +++ b/operators/test/test_utils/reference_check.py @@ -0,0 +1,454 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from __future__ import annotations + +from enum import IntEnum +from functools import partial +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from collections.abc import Callable + + import jax.numpy as jnp + import numpy as np + + import cutlass + + ArrayLike = torch.Tensor | np.ndarray | jnp.ndarray + DtypeLike = cutlass.Numeric | torch.dtype | np.dtype | jnp.dtype +import cutlass +import cutlass.torch + +import cutlass.operators.utils.dtype as dtype_utils +from cutlass.operators.utils.tensor import ( + is_jax_array, + is_numpy_tensor, + is_torch_tensor, +) + + +class ReferenceCheckWarning(UserWarning): + """Warning related to reference check numerics.""" + + +class ClampMode(IntEnum): + NoClamp = 0 # bool(False) -> no clamp + Clamp = 1 # bool(True) -> clamp + ClampOnRestrictedDtypes = 2 # Clamp on narrow precision dtype + + +def get_clamp_torch_dtypes(): + """Get the dtypes that will be clamped under ClampOnRestrictedDtypes. + These dtypes are narrow precision float datatypes that when used + in kernels as output dtypes, typically use the saturation-to-finite + modifier in mma ops. Therefore, to maintain same behavior between + external numerical libraries and operators, we restrict the dtypes' behavior + to clamp the reference tensor output to its max/min finite range. + """ + import torch + + clamp_torch_dtypes = { + torch.float8_e4m3fn, + torch.float8_e5m2, + } + return clamp_torch_dtypes + + +def get_clamp_jax_dtypes(): + """Get the dtypes that will be clamped under ClampOnRestrictedDtypes.""" + import jax.numpy as jnp + + clamp_jax_dtypes = { + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + } + return clamp_jax_dtypes + + +def get_clamp_cutlass_dtypes(): + """Get the cutlass dtypes that will be clamped under ClampOnRestrictedDtypes.""" + clamp_cutlass_dtypes = { + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + } + return clamp_cutlass_dtypes + + +def get_torch_default_tolerances(): + """PyTorch per-dtype default ``(rtol, atol)`` tolerances. + + Snapshotted from + https://docs.pytorch.org/docs/stable/testing.html#torch.testing.assert_close + """ + import torch + + return { + torch.float16: (1e-3, 1e-5), + torch.bfloat16: (1.6e-2, 1e-5), + torch.float32: (1.3e-6, 1e-5), + torch.float64: (1e-7, 1e-7), + } + + +def check_as_numpy_or_jax( + out: np.ndarray | jnp.ndarray, + reference: np.ndarray | jnp.ndarray, + output_dtype: np.dtype | jnp.dtype | cutlass.Numeric, + rtol: float, + atol: float, + equal_nan: bool, + msg: str, + skip_check: bool, + clamp_mode: ClampMode, +) -> np.ndarray: + """Helper functions for assert_close_with_reference_conversion when the + reference uses numpy or jax arrays. Returns converted ref as numpy array. + """ + import jax.numpy as jnp + import numpy as np + + convert_fn = np.ndarray.astype + base = np.float32 + + safe_emulate_fn = partial(convert_fn, dtype=base) + + # convert from jax to numpy array + # (numpy and jax share dtypes if jax is imported with numpy) + reference = np.asarray(reference) + out = np.asarray(out) + + if isinstance(output_dtype, type) and issubclass(output_dtype, cutlass.Numeric): + embedded_dtype = dtype_utils.jax_type_from_cutlass_type(output_dtype) + else: + # raises a TypeError if cannot convert + embedded_dtype = jnp.dtype(output_dtype) + + if clamp_mode == ClampMode.ClampOnRestrictedDtypes: + clamp = embedded_dtype in get_clamp_jax_dtypes() + elif clamp_mode == ClampMode.Clamp: + clamp = True + elif clamp_mode == ClampMode.NoClamp: + clamp = False + else: + raise ValueError(f"Unsupported clamp mode: {clamp_mode}") + + # Check that out tensor dtype matches expected dtype + if out.dtype != embedded_dtype: + raise ValueError( + f"Result tensor dtype {out.dtype} does not match target dtype {output_dtype} ({embedded_dtype})" + ) + + if clamp: + if jnp.issubdtype(embedded_dtype, jnp.floating): + embedded_info = jnp.finfo(embedded_dtype) + elif jnp.issubdtype(embedded_dtype, jnp.integer): + embedded_info = jnp.iinfo(embedded_dtype) + else: + raise ValueError(f"Unsupported dtype {embedded_dtype}") + + # clamp the value of the reference tensor to the finite range of the output dtype + reference = np.clip(reference, min=embedded_info.min, max=embedded_info.max) + + # Convert reference tensor if it doesn't match expected dtype + if reference.dtype != embedded_dtype: + reference = convert_fn(reference, embedded_dtype) + + if not skip_check: + np.testing.assert_allclose( + safe_emulate_fn(out), + safe_emulate_fn(reference), + rtol=rtol, + atol=atol, + equal_nan=equal_nan, + err_msg=msg, + ) + return reference + + +def check_as_torch( + out: torch.Tensor, + reference: torch.Tensor, + output_dtype: torch.dtype, + rtol: float, + atol: float, + equal_nan: bool, + msg: str, + skip_check: bool, + clamp_mode: ClampMode, +) -> torch.Tensor: + """Helper functions for assert_close_with_reference_conversion when the + reference uses PyTorch tensors. Returns converted ref as torch tensor. + """ + import torch + + convert_fn = torch.Tensor.to + base = torch.float32 + safe_emulate_fn = partial(convert_fn, dtype=base) + + # Normalize output dtype to torch dtype + if isinstance(output_dtype, torch.dtype): + embedded_dtype = output_dtype + elif isinstance(output_dtype, type) and issubclass(output_dtype, cutlass.Numeric): + embedded_dtype = cutlass.torch.dtype(output_dtype) + else: + raise ValueError( + f"Output dtype requires a torch.dtype or cutlass.Numeric, got {output_dtype}" + ) + + if clamp_mode == ClampMode.ClampOnRestrictedDtypes: + clamp = embedded_dtype in get_clamp_torch_dtypes() + elif clamp_mode == ClampMode.Clamp: + clamp = True + elif clamp_mode == ClampMode.NoClamp: + clamp = False + else: + raise ValueError(f"Unsupported clamp mode: {clamp_mode}") + + # Check that out tensor dtype matches expected dtype + if out.dtype != embedded_dtype: + raise ValueError( + f"Result tensor dtype {out.dtype} does not match target dtype {output_dtype} ({embedded_dtype})" + ) + + if clamp: + if embedded_dtype.is_floating_point: + embedded_info = torch.finfo(embedded_dtype) + elif embedded_dtype.is_complex: + raise ValueError(f"Complex dtype {embedded_dtype} is not supported") + else: + embedded_info = torch.iinfo(embedded_dtype) + + # clamp the value of the reference tensor to the finite range of the output dtype + reference = torch.clamp(reference, min=embedded_info.min, max=embedded_info.max) + + # Convert reference tensor if it doesn't match expected dtype + if reference.dtype != embedded_dtype: + reference = convert_fn(reference, embedded_dtype) + + # Do comparisons in float32 because testing.assert_close and testing.assert_allclose + # is not safe, doesn't support atol and rtol on lower-bitwidth datatypes + if not skip_check: + torch.testing.assert_close( + safe_emulate_fn(out), + safe_emulate_fn(reference), + rtol=rtol, + atol=atol, + equal_nan=equal_nan, + msg=msg, + ) + return reference + + +def assert_close_with_reference_conversion( + result_tensors: ArrayLike | list[ArrayLike], + accumulator_reference_tensor: ArrayLike | list[ArrayLike], + output_dtypes: DtypeLike | list[DtypeLike], + rtols: float | list[float] = None, + atols: float | list[float] = None, + equal_nan: bool = False, + msg: str = None, + skip_check: bool = False, + clamp: ClampMode | int = ClampMode.ClampOnRestrictedDtypes, +) -> torch.Tensor | np.ndarray | list[torch.Tensor | np.ndarray]: + """ + Asserts that the result tensor is similar to the reference tensor. + This function automatically handles: + - Converting the reference tensor to the correct output dtype + - Handling overflow as saturation instead of NaN or Inf + When reference checking multiple tensors at a time, they must be passed in the same order. + + Args: + result_tensors (ArrayLike | list[ArrayLike]): The output tensors to compare against the reference. Can be a torch.Tensor, np.ndarray, or jnp.ndarray + or a list of these types. + accumulator_reference_tensor (ArrayLike | list[ArrayLike]): The reference tensor in accumulator dtype. Must be the same type as the result tensor. + output_dtypes (DtypeLike | list[DtypeLike]): The dtypes of the output, can be a cutlass.Numeric, torch.dtype, np.dtype, or jnp.dtype + or a list of these types. + rtols (float | list[float]): The relative tolerance to use. Optional, if no value is set then will default to 0.0. + atols (float | list[float]): The absolute tolerance to use. Optional, if no value is set then will default to 0.0. + equal_nan (bool): If True, will compare NaN values as equal + msg (str): The message to display if the assertion fails + skip_check (bool): If True, will skip the ref check and return the reference tensors directly + clamp (ClampMode | int): Configures whether to clamp the reference tensor to the range of the output dtype + + Returns: + The reference tensors as torch.Tensor or np.array, after being converted to output dtype + + Raises: + ValueError: If the number of result tensors, references, and output dtypes are not the same + ValueError: If the number of rtol or atol values is >1 but not the same as the number of result tensors + ValueError: If the reference tensor and result tensor types don't match + ValueError: If the reference tensor is not a torch tensor, numpy array, or jax array + ValueError: If the clamp mode is not supported + """ + + # Argument list preprocessing + accumulator_reference_tensor = ( + [accumulator_reference_tensor] + if not isinstance(accumulator_reference_tensor, list) + else accumulator_reference_tensor + ) + result_tensors = ( + [result_tensors] if not isinstance(result_tensors, list) else result_tensors + ) + output_dtypes = ( + [output_dtypes] if not isinstance(output_dtypes, list) else output_dtypes + ) + rtols = [rtols] if not isinstance(rtols, list) else rtols + atols = [atols] if not isinstance(atols, list) else atols + + num_to_check = len(result_tensors) + dtypes_match = num_to_check == len(output_dtypes) + rtols_match = num_to_check == len(rtols) + atols_match = num_to_check == len(atols) + references_match = num_to_check == len(accumulator_reference_tensor) + + if not all([dtypes_match, references_match]): + raise ValueError( + "The number of result tensors, references, and output dtypes must be the same" + ) + + if not rtols_match: + if len(rtols) == 1: + rtols = ( + rtols * num_to_check + ) # expands length of list to match number of tensors + else: + raise ValueError( + "The number of rtol values must be the same as the number of result tensors" + ) + if not atols_match: + if len(atols) == 1: + atols = ( + atols * num_to_check + ) # expands length of list to match number of tensors + else: + raise ValueError( + "The number of atol values must be the same as the number of result tensors" + ) + + clamp_mode = ClampMode(int(clamp)) + + # When rtol/atol are not explicitly passed, they default to 0.0 (exact + # match). Tests should compute the reference at the kernel's precision + # rather than relax tolerances; the exception is hardware fast-math + # transcendentals (tanh/exp/sigmoid), which need the dtype-default + # tolerance from get_torch_default_tolerances. + # + # Call reference check for each set of reference, result, output_dtype, and rtol/atol + for i, (reference, out, output_dtype, rtol, atol) in enumerate( + zip(accumulator_reference_tensor, result_tensors, output_dtypes, rtols, atols) + ): + rtol = rtol if rtol is not None else 0.0 + atol = atol if atol is not None else 0.0 + + if is_torch_tensor(reference): + if is_numpy_tensor(out) or is_jax_array(out): + raise ValueError( + "Reference tensor is a torch tensor and result tensor is not, this is not supported" + ) + accumulator_reference_tensor[i] = check_as_torch( + out, + reference, + output_dtype, + rtol, + atol, + equal_nan, + msg, + skip_check, + clamp_mode, + ) + elif is_numpy_tensor(reference) or is_jax_array(reference): + if is_torch_tensor(out): + raise ValueError( + "Reference tensor is a jax array and result tensor is a torch tensor, this is not supported" + ) + accumulator_reference_tensor[i] = check_as_numpy_or_jax( + out, + reference, + output_dtype, + rtol, + atol, + equal_nan, + msg, + skip_check, + clamp_mode, + ) + else: + raise ValueError(f"Unsupported tensor type: {type(reference)}") + + return accumulator_reference_tensor + + +def reference_scaled_mm( + A: torch.Tensor, + B: torch.Tensor, + scale_A: torch.Tensor, + scale_B: torch.Tensor, + out_dtype: torch.dtype, + transform_sf: Callable[[torch.Tensor], torch.Tensor] = lambda x: x, +): + """ + Computes a reference scaled mm operation. Currently, torch._scaled_mm does not support batch mode. + When a batch mode is present, this function iterates through each problem in the batch. + + :param A: The A tensor + :type A: torch.Tensor + :param B: The B tensor + :type B: torch.Tensor + :param scale_A: The scale factor tensor for operand A + :type scale_A: torch.Tensor + :param scale_B: The scale factor tensor for operand B + :type scale_B: torch.Tensor + :param out_dtype: The output dtype + :type out_dtype: torch.dtype + :param transform_sf: A function to transform the scale factor tensors to the correct shape for the scaled mm operation + :type transform_sf: Callable + :return: The reference scaled mm operation + """ + if len(A.shape) == 2: + return torch._scaled_mm( + A, B, scale_a=scale_A, scale_b=scale_B, out_dtype=out_dtype + ) + else: + # torch._scaled_mm does not support batch mode. Iterate through each problem in the batch + L, M, N = A.shape[0], A.shape[1], B.shape[2] + scale_A = scale_A.view(L, -1) + scale_B = scale_B.view(L, -1) + reference = torch.empty((L, M, N), device=A.device, dtype=out_dtype) + for l_idx in range(L): + reference[l_idx] = torch._scaled_mm( + A[l_idx], + B[l_idx], + scale_a=transform_sf(scale_A[l_idx]), + scale_b=transform_sf(scale_B[l_idx]), + out_dtype=out_dtype, + ) + return reference diff --git a/operators/test/unit/test_arguments.py b/operators/test/unit/test_arguments.py new file mode 100644 index 000000000..df1baa752 --- /dev/null +++ b/operators/test/unit/test_arguments.py @@ -0,0 +1,153 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""Unit tests for :mod:`cutlass.operators.arguments`.""" + +import pytest + +import cutlass.operators as ops + + +class TestNumelScale: + """Tests for :py:meth:`ops.ScaledOperand.numel_scale`.""" + + def test_swizzle_none_blockwise1x32(self): + # Natural ordering: L * outer * ceil_div(K, 32). No padding. + # 1 * 256 * (1024 / 32) = 256 * 32 = 8192. + assert ops.ScaledOperand.numel_scale( + (1, 256, 1024), + ops.ScaleMode.Blockwise1x32, + ops.ScaleSwizzleMode.SwizzleNone, + ) == (1 * 256 * (1024 // 32)) + + def test_swizzle_none_blockwise1x16(self): + # 2 * 64 * ceil_div(48, 16) = 2 * 64 * 3 = 384. + assert ops.ScaledOperand.numel_scale( + (2, 64, 48), + ops.ScaleMode.Blockwise1x16, + ops.ScaleSwizzleMode.SwizzleNone, + ) == (2 * 64 * (48 // 16)) + + def test_swizzle_none_partial_k_block(self): + # K=33 with V=32 -> ceil_div(33, 32) = 2 K-blocks. + assert ops.ScaledOperand.numel_scale( + (1, 16, 33), + ops.ScaleMode.Blockwise1x32, + ops.ScaleSwizzleMode.SwizzleNone, + ) == 1 * 16 * ((33 + 32 - 1) // 32) + + def test_swizzle32x4x4_aligned(self): + # Aligned shape: M is a multiple of 128, K-blocks is a multiple of 4. + # 1 * 256 * (1024 / 32) = 8192. round_up is a no-op here. + assert ops.ScaledOperand.numel_scale( + (1, 256, 1024), + ops.ScaleMode.Blockwise1x32, + ops.ScaleSwizzleMode.Swizzle32x4x4, + ) == 1 * 256 * (1024 // 32) + + def test_swizzle32x4x4_pads_outer(self): + # M=200 -> round_up(200, 128) = 256. K=1024 -> 32 K-blocks (already + # a multiple of 4). 1 * 256 * 32 = 8192. + assert ( + ops.ScaledOperand.numel_scale( + (1, 200, 1024), + ops.ScaleMode.Blockwise1x32, + ops.ScaleSwizzleMode.Swizzle32x4x4, + ) + == 8192 + ) + + def test_swizzle32x4x4_pads_inner(self): + # K=160 -> ceil_div(160, 32) = 5 K-blocks -> round_up(5, 4) = 8. + # 1 * 256 * 8 = 2048. + assert ( + ops.ScaledOperand.numel_scale( + (1, 256, 160), + ops.ScaleMode.Blockwise1x32, + ops.ScaleSwizzleMode.Swizzle32x4x4, + ) + == 2048 + ) + + def test_swizzle32x4x4_pads_both(self): + # M=200 -> 256, K-blocks = ceil_div(160, 32) = 5 -> 8. + # L=2 * 256 * 8 = 4096. + assert ( + ops.ScaledOperand.numel_scale( + (2, 200, 160), + ops.ScaleMode.Blockwise1x32, + ops.ScaleSwizzleMode.Swizzle32x4x4, + ) + == 4096 + ) + + def test_swizzle32x4x4_blockwise1x16(self): + # NVFP4-style: V=16. K=128 -> 8 K-blocks (a multiple of 4). + # 1 * 128 * 8 = 1024. + assert ( + ops.ScaledOperand.numel_scale( + (1, 128, 128), + ops.ScaleMode.Blockwise1x16, + ops.ScaleSwizzleMode.Swizzle32x4x4, + ) + == 1024 + ) + + def test_tuple_mode(self): + # Bare tuple mode is accepted (and equivalent to the named enum). + as_enum = ops.ScaledOperand.numel_scale( + (1, 200, 160), + ops.ScaleMode.Blockwise1x32, + ops.ScaleSwizzleMode.Swizzle32x4x4, + ) + as_tuple = ops.ScaledOperand.numel_scale( + (1, 200, 160), + (1, 1, 32), + ops.ScaleSwizzleMode.Swizzle32x4x4, + ) + assert as_enum == as_tuple + + def test_2d_quantized_shape(self): + # Rank-2 quantized shape is treated as L=1. + assert ops.ScaledOperand.numel_scale( + (256, 1024), + ops.ScaleMode.Blockwise1x32, + ops.ScaleSwizzleMode.Swizzle32x4x4, + ) == ops.ScaledOperand.numel_scale( + (1, 256, 1024), + ops.ScaleMode.Blockwise1x32, + ops.ScaleSwizzleMode.Swizzle32x4x4, + ) + + def test_invalid_rank_raises(self): + with pytest.raises(ValueError, match="rank 2 or 3"): + ops.ScaledOperand.numel_scale( + (1, 1, 256, 1024), + ops.ScaleMode.Blockwise1x32, + ops.ScaleSwizzleMode.Swizzle32x4x4, + ) diff --git a/operators/test/unit/test_config.py b/operators/test/unit/test_config.py new file mode 100644 index 000000000..ed878b554 --- /dev/null +++ b/operators/test/unit/test_config.py @@ -0,0 +1,183 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""Unit tests for ``cutlass.operators.config.GlobalOptions``. + +Covers environment-variable based initialization, programmatic overrides via +the property setters, and graceful behavior when the optional ``tvm_ffi`` +dependency is missing. +""" + +from importlib.util import find_spec + +import pytest + +import cutlass.operators as ops +from cutlass.operators import config +from cutlass.operators.config import _ENV_USE_TVM_FFI, _parse_bool_env + + +@pytest.fixture +def reset_singleton(monkeypatch): + """Reset the ``GlobalOptions`` singleton so ``__new__`` re-runs init logic. + + ``monkeypatch.setattr`` restores ``_instance`` after the test, so other + tests keep observing the original singleton (preserving the ``save`` / + ``restore`` autouse fixture in ``conftest.py``). + """ + monkeypatch.setattr(ops.GlobalOptions, "_instance", None) + monkeypatch.delenv(_ENV_USE_TVM_FFI, raising=False) + + +# --------------------------------------------------------------------------- +# _parse_bool_env +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "raw, expected", + [ + ("1", True), + ("true", True), + ("True", True), + ("TRUE", True), + ("yes", True), + ("YES", True), + ("on", True), + (" on ", True), + ("0", False), + ("false", False), + ("False", False), + ("FALSE", False), + ("no", False), + ("NO", False), + ("off", False), + ], +) +def test_parse_bool_env_recognized_values(monkeypatch, raw, expected): + monkeypatch.setenv("CUTLASS_OPERATORS_TEST_FLAG", raw) + assert _parse_bool_env("CUTLASS_OPERATORS_TEST_FLAG") is expected + + +def test_parse_bool_env_unset_returns_none(monkeypatch): + monkeypatch.delenv("CUTLASS_OPERATORS_TEST_FLAG", raising=False) + assert _parse_bool_env("CUTLASS_OPERATORS_TEST_FLAG") is None + + +@pytest.mark.parametrize("raw", ["maybe", "ture", "2", "", "yes please"]) +def test_parse_bool_env_rejects_unrecognized_values(monkeypatch, raw): + """Typos must fail loudly so they don't silently pick a default.""" + monkeypatch.setenv("CUTLASS_OPERATORS_TEST_FLAG", raw) + with pytest.raises(ValueError, match="Invalid value"): + _parse_bool_env("CUTLASS_OPERATORS_TEST_FLAG") + + +# --------------------------------------------------------------------------- +# GlobalOptions environment-variable initialization +# --------------------------------------------------------------------------- + + +def test_env_var_disable_is_respected(monkeypatch, reset_singleton): + """Setting the env var to a falsy value disables ``use_tvm_ffi`` even if + ``tvm_ffi`` is installed.""" + monkeypatch.setenv(_ENV_USE_TVM_FFI, "0") + assert ops.GlobalOptions().use_tvm_ffi is False + + +@pytest.mark.skipif( + find_spec("tvm_ffi") is None, reason="apache-tvm-ffi is not installed" +) +def test_env_var_enable_is_respected(monkeypatch, reset_singleton): + monkeypatch.setenv(_ENV_USE_TVM_FFI, "1") + assert ops.GlobalOptions().use_tvm_ffi is True + + +def test_env_var_invalid_value_raises(monkeypatch, reset_singleton): + monkeypatch.setenv(_ENV_USE_TVM_FFI, "maybe") + with pytest.raises(ValueError, match="Invalid value"): + ops.GlobalOptions() + + +def test_env_var_enable_without_tvm_ffi_raises(monkeypatch, reset_singleton): + """If the user explicitly opts in via env var but ``tvm_ffi`` is missing, + fail fast so the misconfiguration is visible.""" + monkeypatch.setattr(config, "find_spec", lambda name: None) + monkeypatch.setenv(_ENV_USE_TVM_FFI, "1") + with pytest.raises(ImportError, match=_ENV_USE_TVM_FFI): + ops.GlobalOptions() + + +def test_env_var_disable_without_tvm_ffi_does_not_warn( + monkeypatch, reset_singleton, recwarn +): + """If the user explicitly opts out via env var, do not warn about the + missing optional dependency: they made a deliberate choice.""" + monkeypatch.setattr(config, "find_spec", lambda name: None) + monkeypatch.setenv(_ENV_USE_TVM_FFI, "0") + assert ops.GlobalOptions().use_tvm_ffi is False + tvm_warnings = [w for w in recwarn.list if "TVM FFI" in str(w.message)] + assert tvm_warnings == [], ( + "Expected no TVM FFI warning when user explicitly disables via env var, " + f"got: {[str(w.message) for w in tvm_warnings]}" + ) + + +# --------------------------------------------------------------------------- +# GlobalOptions default behavior (env var unset) +# --------------------------------------------------------------------------- + + +def test_default_warns_when_tvm_ffi_missing(monkeypatch, reset_singleton): + """Without env var, a missing ``tvm_ffi`` is reported via UserWarning so + users notice the performance hit but the API still functions.""" + monkeypatch.setattr(config, "find_spec", lambda name: None) + with pytest.warns(UserWarning, match="TVM FFI is not installed"): + opts = ops.GlobalOptions() + assert opts.use_tvm_ffi is False + + +@pytest.mark.skipif( + find_spec("tvm_ffi") is None, reason="apache-tvm-ffi is not installed" +) +def test_default_enabled_when_tvm_ffi_present(reset_singleton): + assert ops.GlobalOptions().use_tvm_ffi is True + + +# --------------------------------------------------------------------------- +# Setter overrides env var +# --------------------------------------------------------------------------- + + +def test_setter_overrides_env_var(monkeypatch, reset_singleton): + """Programmatic assignment should win over the env var seed value.""" + monkeypatch.setenv(_ENV_USE_TVM_FFI, "0") + opts = ops.GlobalOptions() + assert opts.use_tvm_ffi is False + if find_spec("tvm_ffi") is not None: + opts.use_tvm_ffi = True + assert opts.use_tvm_ffi is True diff --git a/operators/test/unit/test_fusion_arguments_parsing.py b/operators/test/unit/test_fusion_arguments_parsing.py new file mode 100644 index 000000000..bb7c68d83 --- /dev/null +++ b/operators/test/unit/test_fusion_arguments_parsing.py @@ -0,0 +1,94 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +import pytest +import torch + +import cutlass.operators as ops + +pytestmark = pytest.mark.arch("80") + + +def test_basic_keywords(): + def epi(accum, C, alpha, beta): + D = (alpha * accum) + (beta * C) + return D + + ops.EpilogueArguments( + epilogue_fn=epi, + D=torch.randn(10, 10), + C=torch.randn(10, 10), + alpha=1.0, + beta=1.0, + ) + + +def test_missing_keywords(): + epifn = """ +def epi(accum, C, alpha, beta): + D = (alpha * accum) + (beta * C) + F = relu(D) + return D, F +""" + try: + # Missing F + ops.EpilogueArguments( + epilogue_fn=epifn, + D=torch.randn(10, 10), + C=torch.randn(10, 10), + alpha=1.0, + beta=1.0, + ) + except ValueError as e: + assert "F" in str(e) + else: + raise AssertionError("Failed to catch missing keyword") + + +def test_extra_keywords(): + epifn = """ +def epi(accum, C, alpha, beta): + D = (alpha * accum) + (beta * C) + F = relu(D) + return D, F +""" + try: + ops.EpilogueArguments( + epilogue_fn=epifn, + D=torch.randn(10, 10), + C=torch.randn(10, 10), + alpha=1.0, + beta=1.0, + F=torch.randn(10, 10), + gamma=3.0, + ) + except ValueError as e: + assert "gamma" in str(e) + else: + raise AssertionError("Failed to catch extra keyword") diff --git a/operators/test/unit/test_manifest.py b/operators/test/unit/test_manifest.py new file mode 100644 index 000000000..93be53e84 --- /dev/null +++ b/operators/test/unit/test_manifest.py @@ -0,0 +1,167 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +Unit tests for cutlass.operators.get_operators() and Manifest functionality. +""" + +from __future__ import annotations + +import pprint +from typing import TYPE_CHECKING + +import pytest +import torch + +import cutlass.operators as ops + +if TYPE_CHECKING: + from collections.abc import Callable + + +# ----------------------------------------------------------------------------- +# Test fixtures and helpers +# ----------------------------------------------------------------------------- + + +def make_simple_gemm_args() -> ops.GemmArguments: + """Create simple GemmArguments for testing.""" + A = torch.empty((2, 128, 64), dtype=torch.float16) + B = torch.empty((2, 64, 128), dtype=torch.float16) + out = torch.empty((2, 128, 128), dtype=torch.float16) + return ops.GemmArguments(A=A, B=B, out=out, accumulator_type=torch.float32) + + +def make_epilogue_gemm_args() -> ops.GemmArguments: + """Create GemmArguments with custom epilogue for testing.""" + + def epi(accum): + D = accum * 2.0 + return D + + A = torch.empty((2, 128, 64), dtype=torch.float16) + B = torch.empty((2, 64, 128), dtype=torch.float16) + out = torch.empty((2, 128, 128), dtype=torch.float16) + epi_args = ops.EpilogueArguments(epi, D=out) + return ops.GemmArguments( + A=A, B=B, out=out, accumulator_type=torch.float32, epilogue=epi_args + ) + + +def hash_based_metadata_filter(metadata: ops.OperatorMetadata) -> bool: + """Filter that selects ~50% of operators based on name hash.""" + return hash(metadata.operator_name) % 2 == 0 + + +# ----------------------------------------------------------------------------- +# Test parameters (use factories to avoid creating tensors at import time) +# ----------------------------------------------------------------------------- + +ARGS_FACTORIES = [ + pytest.param(lambda: None, id="args=None"), + pytest.param(make_simple_gemm_args, id="args=simple_gemm"), + pytest.param(make_epilogue_gemm_args, id="args=epilogue_gemm"), +] + +TARGET_SM_KINDS = [ + pytest.param(None, id="target_sm=None"), + pytest.param("90", id="target_sm=90_portable"), + pytest.param("103a", id="target_sm=103a_arch_conditional"), + pytest.param("120f", id="target_sm=120f_family_portable"), +] + +METADATA_FILTERS = [ + pytest.param(None, id="filter=None"), + pytest.param(hash_based_metadata_filter, id="filter=hash_based"), +] + + +# ----------------------------------------------------------------------------- +# Tests +# ----------------------------------------------------------------------------- + + +@pytest.mark.parametrize("args_factory", ARGS_FACTORIES) +@pytest.mark.parametrize("target_sm", TARGET_SM_KINDS) +@pytest.mark.parametrize("metadata_filter", METADATA_FILTERS) +def test_get_operators_supports_invariant( + args_factory: Callable[[], ops.GemmArguments | None], + target_sm: str | None, + metadata_filter: Callable[[ops.OperatorMetadata], bool] | None, +): + """All operators from get_operators() must support the provided args, target_sm, and metadata_filter.""" + args = args_factory() # Create fresh args for each test + operators = ops.get_operators( + args=args, target_sm=target_sm, metadata_filter=metadata_filter + ) + + for k in operators: + # Verify metadata_filter was respected + if metadata_filter is not None: + assert metadata_filter(k.metadata), ( + f"Operator {k.metadata.operator_name} does not satisfy the requested metadata_filter" + ) + + # Verify target_sm compatibility + if target_sm is not None: + assert ops.TargetSm(target_sm).supports_operators_from( + k.metadata.supported_targets + ), f"Operator {k.metadata.operator_name} is not supported on {target_sm}" + + # Verify supports() returns True for the query parameters + if args is not None: + status = k.supports(args, target_sm=target_sm) + assert status, ( + f"Operator {k.metadata.operator_name} returned by get_operators " + f"does not support args: {status.error}" + ) + + +def test_unsupported_target_sm_returns_no_operators(): + """No arch supports cc=0, expect no supported operators to be returned.""" + args = make_simple_gemm_args() + operators = ops.get_operators(args=args, target_sm="0") + assert operators == [], f"Expected no operators for cc=0, got {operators}" + + +def test_all_unique_operator_names(): + """All operator names should be unique. Torch Inductor is using a cache based on operator name.""" + operator_map: dict[str, ops.Operator] = {} + operators = ops.get_operators() + + for k in operators: + name = k.metadata.operator_name + if name in operator_map: + existing = pprint.pformat(operator_map[name].metadata) + duplicate = pprint.pformat(k.metadata) + pytest.fail( + f"Duplicate operator name: {name}\n\n" + f"--- First operator ---\n{existing}\n\n" + f"--- Duplicate operator ---\n{duplicate}" + ) + operator_map[k.metadata.operator_name] = k diff --git a/operators/test/unit/test_metadata.py b/operators/test/unit/test_metadata.py new file mode 100644 index 000000000..946d0748f --- /dev/null +++ b/operators/test/unit/test_metadata.py @@ -0,0 +1,204 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from dataclasses import dataclass + +import pytest +import torch + +import cutlass +from cutlass import cute + +import cutlass.operators as ops +from cutlass.operators.arguments.base import Operand +from cutlass.operators.arguments.operand import _operand_or_dense +from cutlass.operators.metadata import ( + DenseTensorConstraints, + OperandConstraints, + OperandsMetadata, +) +from cutlass.operators.providers.cutedsl.operator import CuteDslOperator +from cutlass.operators.status import Status +from cutlass.operators.typing import TensorLike + + +@dataclass +class DummyArguments(ops.RuntimeArguments): + A: Operand + B: Operand + out: Operand + + def __init__( + self, + A: TensorLike | Operand, + B: TensorLike | Operand, + out: TensorLike | Operand, + ): + self.A = _operand_or_dense(A).copy() + self.B = _operand_or_dense(B).copy() + self.out = _operand_or_dense(out).copy() + super().__init__() + + +@dataclass +class DummyOperandsMetadata(OperandsMetadata): + A: OperandConstraints + B: OperandConstraints + out: OperandConstraints + + def supports(self, other: DummyArguments) -> Status: + return all( + [ + self.A.supports(other.A), + self.B.supports(other.B), + self.out.supports(other.out), + ] + ) + + +class NoopOperatorForTesting(CuteDslOperator): + supported_args_type = DummyArguments + designed_for_min_cc = 80 + + @cute.jit + def impl(self, A, B, out, stream): + cute.printf("Called kernel from host successfully!") + return + + def _compile(self, args: DummyArguments, target_sm: ops.TargetSm | None = None): + stream = cute.runtime.make_fake_stream() + return self.cute_compile( + self.impl, + args.A.tensor, + args.B.tensor, + args.out.tensor, + stream, + target_sm=target_sm, + ) + + def _run( + self, + args: DummyArguments, + compiled_artifact, + stream, + workspace=None, + ): + self.cute_run( + compiled_artifact.compiled_obj, + args.A.tensor, + args.B.tensor, + args.out.tensor, + stream, + ) + + @classmethod + def _generate_operators(cls, _filter, _epilogue_args, _target_sm, _args): + attrs = DenseTensorConstraints( + stride=(0, 1), + dtype=cutlass.Float16, + divisibility=8, + ) + metadata = ops.OperatorMetadata( + operator_name="NoopOperatorForTesting", + operator_class=cls, + supported_targets=[ops.TargetSm("80")], + operands=DummyOperandsMetadata( + A=attrs, + B=attrs, + out=attrs, + ), + ) + return [NoopOperatorForTesting(metadata)] + + +operator = NoopOperatorForTesting.generate_operators(metadata_filter=None)[0] + + +def test_perfectly_aligned(): + divisibility = operator.metadata.operands.A.divisibility + A, B, out = [ + torch.randn(divisibility, divisibility * 2, dtype=torch.float16, device="cuda") + for _ in range(3) + ] + args = DummyArguments(A=A, B=B, out=out) + operator.run(args) + + +def test_overaligned(): + A, B, out = [ + torch.randn(1024, 1024, dtype=torch.float16, device="cuda") for _ in range(3) + ] + args = DummyArguments(A=A, B=B, out=out) + operator.run(args) + + +def _check_misaligned_args(error_match_string: str, **tensors): + """Helper to test various misalignment errors are properly caught. + + With TVM-FFI: + args creation may succeed, but operator.supports must fail. + TVM-FFI should still catch errors if user bypasses supports. + Without TVM-FFI: + error must be caught early during argument creation itself. + """ + if ops.GlobalOptions().use_tvm_ffi: + args = DummyArguments(**tensors) + assert not operator.supports(args), "Unsupported args should be rejected" + + with pytest.raises(Exception, match=error_match_string): + operator.run(args, assume_supported_args=True) + else: + with pytest.raises(Exception, match=error_match_string): + DummyArguments(**tensors) + + +def test_underaligned(fixture_toggle_tvm_ffi): + divisibility = operator.metadata.operands.A.divisibility + A, B, out = [ + torch.randn( + divisibility + divisibility // 2, + divisibility + divisibility // 4, + dtype=torch.float16, + device="cuda", + ) + for _ in range(3) + ] + _check_misaligned_args("divisible", A=A, B=B, out=out) + + +def test_ptr_misaligned(fixture_toggle_tvm_ffi): + rows = operator.metadata.operands.A.divisibility * 4 + cols = rows + offset = 117 + + A = torch.randn(rows * cols + offset, dtype=torch.float16, device="cuda") + B = torch.randn(rows, cols, dtype=torch.float16, device="cuda") + out = torch.randn(rows, cols, dtype=torch.float16, device="cuda") + A_offset = torch.as_strided(A[offset:], (rows, cols), (cols, 1)) + + _check_misaligned_args("align", A=A_offset, B=B, out=out) diff --git a/operators/test/unit/test_provider.py b/operators/test/unit/test_provider.py new file mode 100644 index 000000000..a883317eb --- /dev/null +++ b/operators/test/unit/test_provider.py @@ -0,0 +1,271 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +Unit tests for Provider class and argument type filtering. +""" + +import pytest +import torch + +import cutlass.operators as ops +from cutlass.operators.providers.provider import Provider + + +class TestProvider(Provider): + """Test provider for unit testing.""" + + +def args_for_type(args_type: type[ops.RuntimeArguments]) -> ops.RuntimeArguments: + """Construct arguments for a given argument type.""" + if args_type is ops.GemmArguments: + return ops.GemmArguments( + A=torch.empty((1, 1, 1), dtype=torch.float16), + B=torch.empty((1, 1, 1), dtype=torch.float16), + out=torch.empty((1, 1, 1), dtype=torch.float16), + accumulator_type=torch.float16, + ) + elif args_type is ops.GroupedGemmArguments: + return ops.GroupedGemmArguments( + A=torch.empty((1, 1, 1), dtype=torch.float16), + B=torch.empty((1, 1, 1), dtype=torch.float16), + out=torch.empty((1, 1, 1), dtype=torch.float16), + offsets=torch.empty((1,), dtype=torch.int32), + accumulator_type=torch.float16, + ) + else: + raise ValueError(f"Unsupported argument type: {args_type}") + + +class TestProviderArgumentTypeFiltering: + """Tests for filtering Operator classes by argument type.""" + + def setup_method(self): + """Reset the test provider's Operator classes before each test.""" + TestProvider._operator_classes_by_argtype = {} + + def test_register_operator_class_with_supported_args_type(self): + """Operator classes with supported_args_type should register successfully.""" + + @TestProvider.register + class MockGemmOperator(ops.Operator): + supported_args_type = ops.GemmArguments + + @classmethod + def generate_operators( + cls, metadata_filter, epilogue_args, target_sm, args + ): + return ["gemm_operator"] + + assert ops.GemmArguments in TestProvider._operator_classes_by_argtype + assert ( + MockGemmOperator + in TestProvider._operator_classes_by_argtype[ops.GemmArguments] + ) + + def test_generate_operators_filters_by_args_type(self): + """generate_operators should only call Operator classes matching args_type.""" + + @TestProvider.register + class MockGemmOperator(ops.Operator): + supported_args_type = ops.GemmArguments + + @classmethod + def generate_operators( + cls, metadata_filter, epilogue_args, target_sm, args + ): + return ["gemm_operator"] + + @TestProvider.register + class MockGroupedGemmOperator(ops.Operator): + supported_args_type = ops.GroupedGemmArguments + + @classmethod + def generate_operators( + cls, metadata_filter, epilogue_args, target_sm, args + ): + return ["grouped_gemm_operator"] + + # Filter by GemmArguments - should only get gemm operators + gemm_operators = TestProvider.generate_operators( + metadata_filter=None, args=args_for_type(ops.GemmArguments) + ) + assert gemm_operators == ["gemm_operator"] + + # Filter by GroupedGemmArguments - should only get grouped gemm operators + grouped_operators = TestProvider.generate_operators( + metadata_filter=None, args=args_for_type(ops.GroupedGemmArguments) + ) + assert grouped_operators == ["grouped_gemm_operator"] + + def test_generate_operators_returns_all_when_args_type_is_none(self): + """generate_operators should return all operators when args_type is None.""" + + @TestProvider.register + class MockGemmOperator(ops.Operator): + supported_args_type = ops.GemmArguments + + @classmethod + def generate_operators( + cls, metadata_filter, epilogue_args, target_sm, args + ): + return ["gemm_operator"] + + @TestProvider.register + class MockGroupedGemmOperator(ops.Operator): + supported_args_type = ops.GroupedGemmArguments + + @classmethod + def generate_operators( + cls, metadata_filter, epilogue_args, target_sm, args + ): + return ["grouped_gemm_operator"] + + # No args filter - should get all operators + all_operators = TestProvider.generate_operators() + assert set(all_operators) == {"gemm_operator", "grouped_gemm_operator"} + + def test_multiple_operators_same_args_type(self): + """Multiple Operator classes with the same args_type should all be returned.""" + + @TestProvider.register + class MockGemmOperator1(ops.Operator): + supported_args_type = ops.GemmArguments + + @classmethod + def generate_operators( + cls, metadata_filter, epilogue_args, target_sm, args + ): + return ["gemm_operator_1"] + + @TestProvider.register + class MockGemmOperator2(ops.Operator): + supported_args_type = ops.GemmArguments + + @classmethod + def generate_operators( + cls, metadata_filter, epilogue_args, target_sm, args + ): + return ["gemm_operator_2"] + + operators = TestProvider.generate_operators( + metadata_filter=None, args=args_for_type(ops.GemmArguments) + ) + assert set(operators) == {"gemm_operator_1", "gemm_operator_2"} + + +class TestProviderSubclassIsolation: + """Tests that different provider subclasses have isolated state.""" + + def test_subclasses_have_separate_operator_dicts(self): + """Each provider subclass should have its own _operator_classes_by_argtype.""" + + class ProviderA(Provider): + pass + + class ProviderB(Provider): + pass + + @ProviderA.register + class MockOperatorA(ops.Operator): + supported_args_type = ops.GemmArguments + + @classmethod + def generate_operators( + cls, metadata_filter, epilogue_args, target_sm, args + ): + return ["operator_a"] + + @ProviderB.register + class MockOperatorB(ops.Operator): + supported_args_type = ops.GemmArguments + + @classmethod + def generate_operators( + cls, metadata_filter, epilogue_args, target_sm, args + ): + return ["operator_b"] + + # Each provider should only have its own operator + operators_a = ProviderA.generate_operators( + metadata_filter=None, args=args_for_type(ops.GemmArguments) + ) + operators_b = ProviderB.generate_operators( + metadata_filter=None, args=args_for_type(ops.GemmArguments) + ) + + assert operators_a == ["operator_a"] + assert operators_b == ["operator_b"] + + +class TestOperatorSubclassRequirements: + """Tests for required class attributes on concrete Operator subclasses.""" + + def test_missing_supported_args_type_raises_at_class_definition(self): + with pytest.raises(TypeError, match="supported_args_type"): + + class MissingSupportedArgsTypeOperator(ops.Operator): + designed_for_min_cc = 80 + + def _compile(self, args, target_sm=None): + raise NotImplementedError + + def _run(self, args, compiled_artifact, stream, workspace=None): + raise NotImplementedError + + @classmethod + def _generate_operators( + cls, + metadata_filter, + epilogue_args=None, + target_sm=None, + args=None, + ): + return [] + + def test_missing_designed_for_min_cc_raises_at_class_definition(self): + with pytest.raises(TypeError, match="designed_for_min_cc"): + + class MissingMinCcOperator(ops.Operator): + supported_args_type = ops.GemmArguments + + def _compile(self, args, target_sm=None): + raise NotImplementedError + + def _run(self, args, compiled_artifact, stream, workspace=None): + raise NotImplementedError + + @classmethod + def _generate_operators( + cls, + metadata_filter, + epilogue_args=None, + target_sm=None, + args=None, + ): + return [] diff --git a/operators/test/unit/test_reference_check.py b/operators/test/unit/test_reference_check.py new file mode 100644 index 000000000..9e7586029 --- /dev/null +++ b/operators/test/unit/test_reference_check.py @@ -0,0 +1,477 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +Unit tests for assert_close_with_reference_conversion function. Does not require GPU. +""" + +import jax.numpy as jnp +import numpy as np +import pytest +import torch + +import cutlass + +from test_utils.reference_check import ClampMode, assert_close_with_reference_conversion + + +class TestAssertCloseWithReferenceConversionTorch: + def test_basic_float32_match(self): + """Test basic comparison with matching float32 tensors.""" + reference = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) + result = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) + assert_close_with_reference_conversion( + result, reference, output_dtypes=cutlass.Float32 + ) + + def test_float32_with_tolerance(self): + """Test comparison with rtols and atols parameters.""" + reference = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) + result = torch.tensor([1.001, 2.002, 3.003], dtype=torch.float32) + # Test that values within tolerance pass + assert_close_with_reference_conversion( + result, + reference, + output_dtypes=cutlass.Float32, + rtols=[1e-2], + atols=[1e-2], + ) + # Test that values outside tolerance raise AssertionError + # This test passes if AssertionError is raised (expected behavior) + result_outside_tolerance = torch.tensor([1.1, 2.1, 3.1], dtype=torch.float32) + with pytest.raises(AssertionError): + assert_close_with_reference_conversion( + result_outside_tolerance, + reference, + output_dtypes=cutlass.Float32, + rtols=[1e-2], + atols=[1e-2], + ) + + def test_cutlass_float16_dtype_conversion(self): + """Test dtype conversion from float32 reference to cutlass.Float16 result.""" + reference = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) + result = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float16) + assert_close_with_reference_conversion( + result, reference, output_dtypes=cutlass.Float16 + ) + + def test_torch_float16_dtype_conversion(self): + """Test dtype conversion from float32 reference to torch.float16 result.""" + reference = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) + result = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float16) + assert_close_with_reference_conversion( + result, reference, output_dtypes=torch.float16 + ) + + def test_int_dtype(self): + """Test comparison with integer dtype.""" + reference = torch.tensor([1, 2, 3], dtype=torch.int32) + result = torch.tensor([1, 2, 3], dtype=torch.int32) + assert_close_with_reference_conversion( + result, reference, output_dtypes=cutlass.Int32 + ) + + def test_multiple_tensors(self): + """Test comparison with multiple tensors.""" + ref1 = torch.tensor([1.0, 2.0], dtype=torch.float32) + ref2 = torch.tensor([3.0, 4.0], dtype=torch.float32) + res1 = torch.tensor([1.0, 2.0], dtype=torch.float32) + res2 = torch.tensor([3.0, 4.0], dtype=torch.float32) + assert_close_with_reference_conversion( + [res1, res2], + [ref1, ref2], + output_dtypes=[cutlass.Float32, cutlass.Float32], + ) + + def test_multiple_tensors_different_dtypes(self): + """Test multiple tensors with different dtypes.""" + ref1 = torch.tensor([1.0, 2.0], dtype=torch.float32) + ref2 = torch.tensor([3, 4], dtype=torch.int32) + res1 = torch.tensor([1.0, 2.0], dtype=torch.float16) + res2 = torch.tensor([3, 4], dtype=torch.int32) + assert_close_with_reference_conversion( + [res1, res2], + [ref1, ref2], + output_dtypes=[cutlass.Float16, cutlass.Int32], + ) + + def test_overflow_saturation_float16(self): + """Test that overflow values are clamped to dtype range.""" + # Create reference with values outside float16 range + reference = torch.tensor([1e10, -1e10, 1.0], dtype=torch.float32) + # Result should be saturated to float16 max/min + result = torch.tensor([65504.0, -65504.0, 1.0], dtype=torch.float16) + assert_close_with_reference_conversion( + result, reference, output_dtypes=cutlass.Float16, clamp=True + ) + + def test_custom_message(self): + """Test that custom error message is passed through.""" + reference = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) + result = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) + # Should not raise, but message should be set + assert_close_with_reference_conversion( + result, + reference, + output_dtypes=cutlass.Float32, + msg="Custom test message", + ) + + def test_rtols_atols_expansion(self): + """Test that single rtols/atols values are expanded to match tensor count.""" + ref1 = torch.tensor([1.0], dtype=torch.float32) + ref2 = torch.tensor([2.0], dtype=torch.float32) + res1 = torch.tensor([1.001], dtype=torch.float32) + res2 = torch.tensor([2.001], dtype=torch.float32) + # Single rtols/atols should be expanded to both tensors + assert_close_with_reference_conversion( + [res1, res2], + [ref1, ref2], + output_dtypes=[cutlass.Float32, cutlass.Float32], + rtols=[1e-2], # Single value + atols=[1e-2], # Single value + ) + + def test_mismatched_dtype_raises(self): + """Test that mismatched result dtype raises ValueError.""" + reference = torch.tensor([1.0, 2.0], dtype=torch.float32) + result = torch.tensor([1.0, 2.0], dtype=torch.float32) # Wrong dtype in result + with pytest.raises(ValueError, match="does not match target dtype"): + assert_close_with_reference_conversion( + result, + reference, + output_dtypes=cutlass.Float16, # not float32 + ) + + def test_wrong_number_of_tensors_raises(self): + """Test that mismatched number of tensors raises ValueError.""" + reference = torch.tensor([1.0], dtype=torch.float32) + result1 = torch.tensor([1.0], dtype=torch.float32) + result2 = torch.tensor([2.0], dtype=torch.float32) + with pytest.raises( + ValueError, + match="number of result tensors, references, and output dtypes must be the same", + ): + assert_close_with_reference_conversion( + [result1, result2], reference, output_dtypes=cutlass.Float32 + ) + + def test_wrong_number_of_rtols_raises(self): + """Test that wrong number of rtols values raises ValueError.""" + ref1 = torch.tensor([1.0], dtype=torch.float32) + ref2 = torch.tensor([2.0], dtype=torch.float32) + res1 = torch.tensor([1.0], dtype=torch.float32) + res2 = torch.tensor([2.0], dtype=torch.float32) + with pytest.raises(ValueError, match="number of rtol values must be the same"): + assert_close_with_reference_conversion( + [res1, res2], + [ref1, ref2], + output_dtypes=[cutlass.Float32, cutlass.Float32], + rtols=[1e-2, 1e-2, 1e-2], # Wrong count + ) + + def test_wrong_number_of_atols_raises(self): + """Test that wrong number of atols values raises ValueError.""" + ref1 = torch.tensor([1.0], dtype=torch.float32) + ref2 = torch.tensor([2.0], dtype=torch.float32) + res1 = torch.tensor([1.0], dtype=torch.float32) + res2 = torch.tensor([2.0], dtype=torch.float32) + with pytest.raises(ValueError, match="number of atol values must be the same"): + assert_close_with_reference_conversion( + [res1, res2], + [ref1, ref2], + output_dtypes=[cutlass.Float32, cutlass.Float32], + atols=[1e-2, 1e-2, 1e-2], # Wrong count + ) + + def test_invalid_output_dtype_raises(self): + """Test that invalid output dtype type raises ValueError.""" + reference = torch.tensor([1.0, 2.0], dtype=torch.float32) + result = torch.tensor([1.0, 2.0], dtype=torch.float32) + with pytest.raises( + ValueError, match="Output dtype requires a torch.dtype or cutlass.Numeric" + ): + assert_close_with_reference_conversion( + result, reference, output_dtypes=["invalid"] + ) + + +class TestAssertCloseWithReferenceConversionNumpy: + def test_np_array_basic(self): + """Test basic comparison with numpy arrays.""" + reference = np.array([1.0, 2.0, 3.0], dtype=np.float32) + result = np.array([1.0, 2.0, 3.0], dtype=np.float32) + assert_close_with_reference_conversion( + result, reference, output_dtypes=cutlass.Float32 + ) + + def test_cutlass_float16_dtype_conversion(self): + """Test dtype conversion from float32 reference to float16 result.""" + reference = np.array([1.0, 2.0, 3.0], dtype=np.float32) + result = np.array([1.0, 2.0, 3.0], dtype=np.float16) + assert_close_with_reference_conversion( + result, reference, output_dtypes=cutlass.Float16 + ) + + def test_overflow_float16(self): + """Test that overflow values are not clamped with non-narrow datatype.""" + # Create reference with values outside float16 range + reference = np.array([1e10, -1e10, 1.0], dtype=np.float32) + result = np.array([float("inf"), float("-inf"), 1.0], dtype=np.float16) + assert_close_with_reference_conversion( + result, reference, output_dtypes=cutlass.Float16 + ) + + def test_overflow_saturation_float8_e4m3fn(self): + """Test that overflow values are clamped to dtype range.""" + # Create reference with values outside float8_e4m3fn range + reference = np.array([1e10, -1e10, 1.0], dtype=np.float32) + # Result should be saturated to float8_e4m3fn max/min + result = np.array([448.0, -448.0, 1.0], dtype=jnp.float8_e4m3fn) + assert_close_with_reference_conversion( + result, reference, output_dtypes=jnp.float8_e4m3fn + ) + + def test_numpy_float16_dtype_conversion(self): + """Test using np.dtype directly instead of cutlass.Numeric.""" + reference = np.array([1.0, 2.0], dtype=np.float32) + result = np.array([1.0, 2.0], dtype=np.float16) + assert_close_with_reference_conversion( + result, reference, output_dtypes=np.float16 + ) + + def test_mismatched_tensor_types_raises(self): + """Test that mixing numpy and torch tensors raises ValueError.""" + reference = np.array([1.0, 2.0], dtype=np.float32) + result = torch.tensor([1.0, 2.0], dtype=torch.float32) + with pytest.raises( + ValueError, match="jax array and result tensor is a torch tensor" + ): + assert_close_with_reference_conversion( + result, reference, output_dtypes=cutlass.Float32 + ) + + +class TestAssertCloseWithReferenceConversionJax: + def test_jax_array_float8_e4m3fn_with_cutlass_dtype(self): + """Test jax float8_e4m3fn arrays with cutlass.Float8E4M3FN output dtype.""" + reference = jnp.array([1.0, -1.0, 0.5], dtype=jnp.float32) + result = jnp.array([1.0, -1.0, 0.5], dtype=jnp.float8_e4m3fn) + assert_close_with_reference_conversion( + result, reference, output_dtypes=cutlass.Float8E4M3FN + ) + + def test_jax_array_float8_e5m2_with_cutlass_dtype(self): + """Test jax float8_e5m2 arrays with cutlass.Float8E5M2 output dtype.""" + reference = jnp.array([1.0, -1.0, 0.5], dtype=jnp.float32) + result = jnp.array([1.0, -1.0, 0.5], dtype=jnp.float8_e5m2) + assert_close_with_reference_conversion( + result, reference, output_dtypes=cutlass.Float8E5M2 + ) + + def test_jax_array_overflow_clamp_float8_e4m3fn(self): + """Test overflow clamping for jax float8_e4m3fn arrays.""" + reference = jnp.array([500.0, -500.0, 1.0], dtype=jnp.float32) + result = jnp.array([448.0, -448.0, 1.0], dtype=jnp.float8_e4m3fn) + assert_close_with_reference_conversion( + result, reference, output_dtypes=cutlass.Float8E4M3FN + ) + + +class TestAssertCloseWithReferenceConversionEdgeCases: + def test_none_rtols_atols_defaults(self): + """Test that None rtols/atols default to 0.0.""" + reference = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) + result = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) + # Should work with exact match when rtols/atols are None + assert_close_with_reference_conversion( + result, + reference, + output_dtypes=cutlass.Float32, + rtols=None, + atols=None, + ) + result_outside_tolerance = torch.tensor([1.1, 2.1, 3.1], dtype=torch.float32) + with pytest.raises(AssertionError): + assert_close_with_reference_conversion( + result_outside_tolerance, + reference, + output_dtypes=cutlass.Float32, + rtols=None, + atols=None, + ) + + def test_empty_tensor(self): + """Test comparison with empty tensors.""" + reference = torch.tensor([], dtype=torch.float32) + result = torch.tensor([], dtype=torch.float32) + assert_close_with_reference_conversion( + result, reference, output_dtypes=cutlass.Float32 + ) + + def test_torch_equals_nan(self): + """Test NaN equality does not raise on assertion.""" + reference = torch.tensor([1.0, float("nan"), 3.0], dtype=torch.float16) + result = torch.tensor([1.0, float("nan"), 3.0], dtype=torch.float16) + assert_close_with_reference_conversion( + result, + reference, + output_dtypes=cutlass.Float16, + clamp=False, + equal_nan=True, + ) + + def test_numpy_equals_nan(self): + """Test NaN equality does not raise on assertion.""" + reference = np.array([1.0, float("nan"), 3.0], dtype=np.float32) + result = np.array([1.0, float("nan"), 3.0], dtype=np.float32) + assert_close_with_reference_conversion( + result, + reference, + output_dtypes=cutlass.Float32, + clamp=False, + equal_nan=True, + ) + + def test_clamp_mode_enum_clamp(self): + """Test that ClampMode.Clamp enum value works correctly.""" + # Create reference with values outside float16 range + reference = torch.tensor([1e10, -1e10, 1.0], dtype=torch.float32) + # Result should be saturated to float16 max/min + result = torch.tensor([65504.0, -65504.0, 1.0], dtype=torch.float16) + assert_close_with_reference_conversion( + result, reference, output_dtypes=cutlass.Float16, clamp=ClampMode.Clamp + ) + + def test_clamp_mode_enum_no_clamp(self): + """Test that ClampMode.NoClamp enum value works correctly.""" + reference = torch.tensor([1.0, float("inf"), 3.0], dtype=torch.float32) + result = torch.tensor([1.0, float("inf"), 3.0], dtype=torch.float32) + assert_close_with_reference_conversion( + result, reference, output_dtypes=cutlass.Float32, clamp=ClampMode.NoClamp + ) + + def test_restricted_clamp_mode_clamp_on_float16(self): + """Test that ClampMode.ClampOnRestrictedDtypes doesn't clamp on float16.""" + # float16 is a restricted dtype, so should not clamp + reference = torch.tensor([1.0, float("inf"), 3.0], dtype=torch.float32) + result = torch.tensor([1.0, float("inf"), 3.0], dtype=torch.float16) + # Should not raise ValueError when clamp=ClampMode.ClampOnRestrictedDtypes with float16 + assert_close_with_reference_conversion( + result, + reference, + output_dtypes=cutlass.Float16, + clamp=ClampMode.ClampOnRestrictedDtypes, + ) + + def test_restricted_clamp_mode_clamp_on_int8(self): + """Test that ClampMode.ClampOnRestrictedDtypes doesn't clamp on int8.""" + # int8 is a restricted dtype, so should not clamp + reference = torch.tensor([1000, -1000, 1], dtype=torch.int32) + # When not clamped, overflow values will wrap or become invalid + result = torch.tensor([-24, 24, 1], dtype=torch.int8) + # Should not clamp overflow values for int8 + assert_close_with_reference_conversion( + result, + reference, + output_dtypes=cutlass.Int8, + clamp=ClampMode.ClampOnRestrictedDtypes, + ) + + def test_restricted_clamp_mode_clamp_on_uint8(self): + """Test that ClampMode.ClampOnRestrictedDtypes doesn't clamp on uint8.""" + # uint8 is a restricted dtype, so should not clamp + reference = torch.tensor([1000, 500, 1], dtype=torch.int32) + # When not clamped, overflow values will wrap + result = torch.tensor([232, 244, 1], dtype=torch.uint8) + # Should not clamp overflow values for uint8 + assert_close_with_reference_conversion( + result, + reference, + output_dtypes=cutlass.Uint8, + clamp=ClampMode.ClampOnRestrictedDtypes, + ) + + def test_restricted_clamp_mode_clamp_on_float8_e4m3fn(self): + """Test that ClampMode.ClampOnRestrictedDtypes DOES clamp on float8_e4m3fn (not in NOCLAMP list).""" + # float8_e4m3fn is NOT in NOCLAMP list, so should clamp + reference = torch.tensor([500.0, -500.0, 1.0], dtype=torch.float32) + # Result should be clamped to float8_e4m3fn range + result = torch.tensor([448.0, -448.0, 1.0], dtype=torch.float8_e4m3fn) + # Should clamp overflow values for float8_e4m3fn + assert_close_with_reference_conversion( + result, + reference, + output_dtypes=cutlass.Float8E4M3FN, + clamp=ClampMode.ClampOnRestrictedDtypes, + ) + + def test_restricted_clamp_mode_clamp_on_float8_e5m2(self): + """Test that ClampMode.ClampOnRestrictedDtypes DOES clamp on float8_e5m2 (not in NOCLAMP list).""" + # float8_e5m2 is NOT in NOCLAMP list, so should clamp + reference = torch.tensor([100000.0, -100000.0, 1.0], dtype=torch.float32) + # Result should be clamped to float8_e5m2 range + result = torch.tensor([57344.0, -57344.0, 1.0], dtype=torch.float8_e5m2) + # Should clamp overflow values for float8_e5m2 + assert_close_with_reference_conversion( + result, + reference, + output_dtypes=cutlass.Float8E5M2, + clamp=ClampMode.ClampOnRestrictedDtypes, + ) + + def test_clamp_mode_bool_true(self): + """Test that ClampMode can be converted from boolean True.""" + reference = torch.tensor([100000.0, -100000.0, 1.0], dtype=torch.float32) + result = torch.tensor([57344.0, -57344.0, 1.0], dtype=torch.float8_e5m2) + assert_close_with_reference_conversion( + result, reference, output_dtypes=cutlass.Float8E5M2, clamp=True + ) + + def test_clamp_mode_bool_false(self): + """Test that ClampMode can be converted from boolean False.""" + reference = torch.tensor([1000.0, -1000.0, 1.0], dtype=torch.float32) + result = torch.tensor( + [float("nan"), -float("nan"), 1.0], dtype=torch.float8_e4m3fn + ) + assert_close_with_reference_conversion( + result, + reference, + output_dtypes=cutlass.Float8E4M3FN, + clamp=False, + equal_nan=True, + ) + + def test_numpy_reference_jax_result_float32(self): + """Test numpy reference array with jax result array.""" + reference = np.array([1.0, 2.0, 3.0], dtype=np.float32) + result = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float32) + assert_close_with_reference_conversion( + result, reference, output_dtypes=cutlass.Float32 + ) diff --git a/operators/test/unit/test_tensor_wrapper.py b/operators/test/unit/test_tensor_wrapper.py new file mode 100644 index 000000000..726c8ddae --- /dev/null +++ b/operators/test/unit/test_tensor_wrapper.py @@ -0,0 +1,145 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""Unit tests for TensorWrapper with packed sub-byte types (float4_e2m1fn_x2).""" + +import pytest +import torch + +import cutlass +import cutlass.operators as ops +from cutlass.operators.utils.tensor import TensorWrapper + + +@pytest.fixture(autouse=True) +def _enable_tvm_ffi(): + """TensorWrapper's FP4 packing logic lives in the TVM-FFI path.""" + prev = ops.GlobalOptions().use_tvm_ffi + ops.GlobalOptions().use_tvm_ffi = True + yield + ops.GlobalOptions().use_tvm_ffi = prev + + +class TestTensorWrapperFP4Packing: + """Tests that TensorWrapper correctly reports logical shape/stride for + float4_e2m1fn_x2 tensors, which pack 2 FP4 values per byte. + """ + + def _make_fp4_tensor(self, phys_shape: tuple[int, ...]) -> torch.Tensor: + """Create a float4_e2m1fn_x2 tensor with given physical shape.""" + return torch.randint( + 0, 256, phys_shape, dtype=torch.uint8, device="cpu" + ).view(torch.float4_e2m1fn_x2) + + def test_2d_row_major(self): + """2D row-major (M, K_phys): logical K should be 2x physical.""" + M, K_phys = 64, 128 + t = self._make_fp4_tensor((M, K_phys)) + assert t.stride() == (K_phys, 1) + + wrapper = TensorWrapper(t, alignment_bytes=16) + assert wrapper.shape == (M, K_phys * 2) + assert wrapper.stride == (K_phys * 2, 1) + + def test_2d_col_major(self): + """2D col-major (K_phys, N): logical K should be 2x physical.""" + K_phys, N = 128, 64 + # Create (N, K_phys) contiguous then transpose to get col-major: + # shape (K_phys, N) with stride (1, K_phys) + t = torch.randint( + 0, 256, (N, K_phys), dtype=torch.uint8, device="cpu" + ).view(torch.float4_e2m1fn_x2).t() + assert t.shape == (K_phys, N) + assert t.stride() == (1, K_phys) + + wrapper = TensorWrapper(t, alignment_bytes=16) + assert wrapper.shape == (K_phys * 2, N) + assert wrapper.stride == (1, K_phys * 2) + + def test_3d_row_major(self): + """3D row-major (L, M, K_phys): logical K should be 2x physical.""" + L, M, K_phys = 2, 64, 128 + t = self._make_fp4_tensor((L, M, K_phys)) + assert t.stride() == (M * K_phys, K_phys, 1) + + wrapper = TensorWrapper(t, alignment_bytes=16) + assert wrapper.shape == (L, M, K_phys * 2) + assert wrapper.stride == (M * K_phys * 2, K_phys * 2, 1) + + def test_3d_col_major_innermost(self): + """3D with stride-1 on dim 1 (e.g. B tensor stored as (L, K_phys, N) + but physically contiguous on K): logical K should be 2x physical.""" + L, K_phys, N = 2, 128, 64 + # Create (L, N, K_phys) contiguous then transpose dims 1,2 + base = torch.randint( + 0, 256, (L, N, K_phys), dtype=torch.uint8, device="cpu" + ).view(torch.float4_e2m1fn_x2) + t = base.transpose(1, 2) + # t.shape = (L, K_phys, N), t.stride = (N*K_phys, 1, K_phys) + assert t.shape == (L, K_phys, N) + assert t.stride() == (N * K_phys, 1, K_phys) + + wrapper = TensorWrapper(t, alignment_bytes=16) + assert wrapper.shape == (L, K_phys * 2, N) + assert wrapper.stride == (N * K_phys * 2, 1, K_phys * 2) + + def test_batch_1(self): + """3D with L=1 should still produce correct logical shape.""" + L, M, K_phys = 1, 32, 256 + t = self._make_fp4_tensor((L, M, K_phys)) + + wrapper = TensorWrapper(t, alignment_bytes=16) + assert wrapper.shape == (L, M, K_phys * 2) + assert wrapper.stride == (M * K_phys * 2, K_phys * 2, 1) + + def test_dtype_is_logical(self): + """TensorWrapper.dtype should be the logical unpacked type.""" + t = self._make_fp4_tensor((32, 64)) + wrapper = TensorWrapper(t, alignment_bytes=16) + assert wrapper.dtype == cutlass.Float4E2M1FN + + +class TestTensorWrapperNonPackedTypes: + """Sanity checks: non-packed types should pass through shape/stride unchanged.""" + + @pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.bfloat16]) + def test_standard_dtypes_2d(self, dtype): + M, K = 64, 128 + t = torch.randn(M, K, dtype=dtype) + wrapper = TensorWrapper(t, alignment_bytes=16) + assert wrapper.shape == (M, K) + assert wrapper.stride == (K, 1) + + def test_float8_e4m3fn_3d(self): + L, M, K = 2, 64, 128 + t = torch.randint(0, 127, (L, M, K), dtype=torch.uint8).view( + torch.float8_e4m3fn + ) + wrapper = TensorWrapper(t, alignment_bytes=16) + assert wrapper.shape == (L, M, K) + assert wrapper.stride == (M * K, K, 1) diff --git a/python/CuTeDSL/_mlir_helpers/__init__.py b/python/CuTeDSL/_mlir_helpers/__init__.py new file mode 100644 index 000000000..33fc30738 --- /dev/null +++ b/python/CuTeDSL/_mlir_helpers/__init__.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +""" +This module provides MLIR Dialect helper functions +""" + +from . import nanobind_compat # must be first: patches ir.Block, ir.InsertionPoint +from . import arith +from . import vector +from .lru_cache_ir import lru_cache_ir +from .op import dsl_user_op + +__all__ = ["arith", "vector", "lru_cache_ir", "dsl_user_op"] + +try: + from . import gpu + + __all__.extend(["gpu"]) +except ImportError: + pass diff --git a/python/CuTeDSL/_mlir_helpers/arith.py b/python/CuTeDSL/_mlir_helpers/arith.py new file mode 100644 index 000000000..5a781a3f1 --- /dev/null +++ b/python/CuTeDSL/_mlir_helpers/arith.py @@ -0,0 +1,1068 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +""" +This module provides MLIR Arith Dialect helper functions +""" + +import array +from typing import Any, Callable, Optional, Union + +import numpy as np + +from ..base_dsl.common import * +from ..base_dsl.common import DSLRuntimeError, DSLNotImplemented +from .._mlir import ir +from .._mlir.extras import types as T +from .._mlir.dialects import arith, math, builtin +from .op import dsl_user_op + +from .lru_cache_ir import lru_cache_ir + +# ============================================================================= +# Arith Dialect Helper functions +# ============================================================================= + + +def recast_type(src_type: ir.Type, res_elem_type: ir.Type) -> ir.Type: + if isinstance(src_type, T.VectorType): + if src_type.scalable: + res_type = T.vector( + *src_type.shape, + res_elem_type, + scalable=src_type.scalable, + scalable_dims=src_type.scalable_dims, + ) + else: + res_type = T.vector(*src_type.shape, res_elem_type) + elif isinstance(src_type, T.RankedTensorType): + res_type = T.RankedTensorType.get( + element_type=res_elem_type, shape=src_type.shape, strides=src_type.strides + ) + elif isinstance(src_type, T.UnrankedTensorType): + res_type = T.UnrankedTensorType.get(element_type=res_elem_type) + elif isinstance(src_type, T.MemRefType): + res_type = T.MemRefType.get( + element_type=res_elem_type, shape=src_type.shape, strides=src_type.strides + ) + else: + res_type = res_elem_type + return res_type + + +def is_scalar(ty: ir.Type) -> bool: + return not isinstance( + ty, (T.VectorType, T.RankedTensorType, T.UnrankedTensorType, T.MemRefType) + ) + + +def element_type(ty: ir.Type) -> ir.Type: + if not is_scalar(ty): + return ty.element_type + else: + return ty + + +def is_narrow_precision(ty: ir.Type) -> bool: + narrow_types = { + T.f8E8M0FNU(), + T.f8E4M3FN(), + T.f8E4M3(), + T.f8E5M2(), + T.f8E4M3B11FNUZ(), + T.f4E2M1FN(), + T.f6E3M2FN(), + T.f6E2M3FN(), + } + return ty in narrow_types + + +def is_float_type(ty: ir.Type) -> bool: + return ( + isinstance(ty, ir.FloatType) + # TODO-upstream: prediction is not correct. Patch here and fix in upstream later + or is_narrow_precision(ty) + or ty in (T.bf16(), T.tf32()) + ) + + +def is_integer_like_type(ty: ir.Type) -> bool: + return isinstance(ty, (ir.IntegerType, ir.IndexType)) + + +def _python_int_for_integer_attr(value: int, width: int, signed: bool) -> int: + """Coerce a Python int to a value MLIR's IntegerAttr.get accepts. + + MLIR's Python bindings only accept integers in the signed range of the + target width. Unsigned DSL constants may be supplied either as canonical + unsigned literals, e.g. 0xffffffffffffffff, or as the equivalent signed + storage representation, e.g. -1. + """ + if width <= 0: + raise ValueError(f"Invalid integer width: {width}") + + signed_lo = -(1 << (width - 1)) + signed_hi = (1 << (width - 1)) - 1 + + if signed: + if not (signed_lo <= value <= signed_hi): + raise ValueError(f"Signed integer literal {value} does not fit in i{width}") + return value + + max_unsigned = (1 << width) - 1 + + if signed_lo <= value < 0: + return value + + if not (0 <= value <= max_unsigned): + raise ValueError( + f"Unsigned integer literal {value} does not fit in unsigned i{width}" + ) + + if value > signed_hi: + return value - (1 << width) + + return value + + +def bitcast( + src: ir.Value, + res_elem_type: ir.Type, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> ir.Value: + res_type = recast_type(src.type, res_elem_type) + return arith.bitcast(res_type, src, loc=loc, ip=ip) + + +def cvtf( + src: ir.Value, + res_elem_type: ir.Type, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> ir.Value: + src_elem_type = element_type(src.type) + + if res_elem_type == src_elem_type: + return src + + res_type = recast_type(src.type, res_elem_type) + + # Treat TF32 as F32 and use i32 as intermediate data + # TODO-upstream: update arith to support tf32 <-> f32 conversion + if src_elem_type == T.tf32(): + # tf32 -> i32 + tmp_type = recast_type(src.type, T.i32()) + src = builtin.unrealized_conversion_cast([tmp_type], [src], loc=loc, ip=ip) + # i32 -> f32 + src = bitcast(src, T.f32(), loc=loc, ip=ip) + # f32 -> X with `cvtf` recursively + return cvtf(src, res_elem_type, loc=loc, ip=ip) + + if res_elem_type == T.tf32(): + # X -> f32 with `cvtf`` recursively + tmp = cvtf(src, T.f32(), loc=loc, ip=ip) + # f32 -> i32 + tmp = bitcast(tmp, T.i32(), loc=loc, ip=ip) + # i32 -> tf32 + return builtin.unrealized_conversion_cast([res_type], [tmp], loc=loc, ip=ip) + + if res_elem_type.width > src_elem_type.width: + return arith.extf(res_type, src, loc=loc, ip=ip) + else: + # bf16 <-> f16: both are 16-bit, arith.truncf requires strict narrowing. + # Route through f32 intermediate. + if (src_elem_type == T.f16() and res_elem_type == T.bf16()) or ( + src_elem_type == T.bf16() and res_elem_type == T.f16() + ): + tmp_type = recast_type(src.type, T.f32()) + tmp = arith.extf(tmp_type, src, loc=loc, ip=ip) + return arith.truncf(res_type, tmp, loc=loc, ip=ip) + + # E8M0 requires upward rounding; all others default to to_nearest_even + roundingmode = ( + arith.RoundingMode.upward if res_elem_type == T.f8E8M0FNU() else None + ) + return arith.truncf(res_type, src, roundingmode=roundingmode, loc=loc, ip=ip) + + +def fptoi( + src: ir.Value, + signed: Union[bool, None], + res_elem_type: ir.Type, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> ir.Value: + res_type = recast_type(src.type, res_elem_type) + # TODO-upstream: update arith to support this kind of conversion + if element_type(src.type) in (T.tf32(), T.bf16()): + src = cvtf(src, T.f32(), loc=loc, ip=ip) + + if signed != False: # noqa: E712 + return arith.fptosi(res_type, src, loc=loc, ip=ip) + else: + return arith.fptoui(res_type, src, loc=loc, ip=ip) + + +def itofp( + src: ir.Value, + signed: Union[bool, None], + res_elem_type: ir.Type, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> ir.Value: + res_type = recast_type(src.type, res_elem_type) + + orig_res_type = res_type + # TODO-upstream: update arith to support this kind of conversion + if res_elem_type in (T.tf32(), T.bf16()): + res_type = recast_type(src.type, T.f32()) + + if signed != False and element_type(src.type).width > 1: # noqa: E712 + res = arith.sitofp(res_type, src, loc=loc, ip=ip) + else: + res = arith.uitofp(res_type, src, loc=loc, ip=ip) + + if orig_res_type == res_type: + return res + + return cvtf(res, element_type(orig_res_type), loc=loc, ip=ip) + + +def int_to_int( + a: ir.Value, + dst_elem_type: Any, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> ir.Value: + src_signed = a.signed + dst_signed = dst_elem_type.signed + src_width = element_type(a.type).width + dst_width = dst_elem_type.width + + dst_mlir_type = recast_type(a.type, dst_elem_type.mlir_type) + + if dst_width == src_width: + return a + elif src_signed != False and not dst_signed: # noqa: E712 + # Signed -> Unsigned + if dst_width > src_width: + return arith.extui(dst_mlir_type, a, loc=loc, ip=ip) + else: + return arith.trunci(dst_mlir_type, a, loc=loc, ip=ip) + elif src_signed == dst_signed: + # Same signedness + if dst_width > src_width: + if src_signed != False and src_width > 1: # noqa: E712 + return arith.extsi(dst_mlir_type, a, loc=loc, ip=ip) + else: + return arith.extui(dst_mlir_type, a, loc=loc, ip=ip) + else: + return arith.trunci(dst_mlir_type, a, loc=loc, ip=ip) + else: + # Unsigned -> Signed + if dst_width > src_width: + return arith.extui(dst_mlir_type, a, loc=loc, ip=ip) + else: + # For truncation from unsigned to signed, we need to handle overflow + # First truncate to the target width + trunc = arith.trunci(dst_mlir_type, a, loc=loc, ip=ip) + # Then reinterpret as signed + if dst_signed: + return arith.bitcast(dst_mlir_type, trunc, loc=loc, ip=ip) + return trunc + + +# ============================================================================= +# Arith Ops Emitter Helpers +# - assuming type of lhs and rhs match each other +# - op name matches python module operator +# ============================================================================= + + +def _cast( + res_elem_ty: ir.Type, + src: ir.Value, + is_signed: Optional[bool] = None, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> ir.Value: + """ + This function provides simplified interface to upstream op builder + arith.truncf(T.vector(shape, new_type), src) + + is simplified as because it's element-wise op which can't change shape + arith.truncf(new_type, src) + """ + if isinstance(src, ir.Value): + src_ty = src.type + else: + src_ty = type(src).mlir_type + src = src.ir_value() + + src_elem_ty = element_type(src_ty) + + if src_elem_ty == res_elem_ty: + return src + elif is_float_type(src_elem_ty) and is_float_type(res_elem_ty): + # float-to-float + return cvtf(src, res_elem_ty, loc=loc, ip=ip) + elif is_integer_like_type(src_elem_ty) and is_integer_like_type(res_elem_ty): + if src_elem_ty.width >= res_elem_ty.width: + cast_op = arith.trunci + else: + if is_signed != False: # noqa: E712 + cast_op = arith.extsi + else: + cast_op = arith.extui + + res_ty = recast_type(src_ty, res_elem_ty) + return cast_op(res_ty, src, loc=loc, ip=ip) + elif is_float_type(src_elem_ty) and is_integer_like_type(res_elem_ty): + return fptoi(src, is_signed, res_elem_ty, loc=loc, ip=ip) + elif is_integer_like_type(src_elem_ty) and is_float_type(res_elem_ty): + return itofp(src, is_signed, res_elem_ty, loc=loc, ip=ip) + else: + raise DSLRuntimeError( + f"cast from {src_elem_ty} to {res_elem_ty} is not supported" + ) + + +@lru_cache_ir() +def const( + value: Union[int, float, bool, np.ndarray], + ty: Optional[Union[ir.Type, "NumericMeta"]] = None, # type: ignore[name-defined] + *, + signed: Union[bool, None] = None, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> ir.Value: + """ + Generates dynamic expression for constant values. + """ + from ..base_dsl.typing import Numeric, NumericMeta + from ..base_dsl.dsl import is_dynamic_expression + from ..base_dsl.utils.numpy import _numpy_type_to_mlir_type + + if isinstance(value, Numeric): + value = value.value + + # Early return + if is_dynamic_expression(value) and ( + isinstance(value.type, type(value.type)) # type: ignore[union-attr] + or isinstance(value.type, type(T.bool())) # type: ignore[union-attr] + ): + return value + + # Assume type + if ty is None: + if isinstance(value, float): + ty = T.f32() + elif isinstance(value, bool): + ty = T.bool() + elif isinstance(value, int): + ty = T.i32() + elif isinstance(value, np.ndarray): + ty = T.vector(*value.shape, _numpy_type_to_mlir_type(value.dtype)) + value = array.array(value.dtype.kind, value.flatten().tolist()) # type: ignore[assignment] + else: + raise DSLNotImplemented(f"{type(value)} is not supported") + elif isinstance(ty, NumericMeta): + if signed is None: + signed = getattr(ty, "signed", None) + if isinstance(value, int) and signed is False: + value = _python_int_for_integer_attr(value, ty.width, signed=False) + ty = ty.mlir_type + elif isinstance(ty, ir.Type): + if isinstance(ty, ir.RankedTensorType) or isinstance(ty, ir.VectorType): + elem_ty = ty.element_type + if isinstance(elem_ty, ir.IntegerType): + if isinstance(value, int) and signed is False: + value = _python_int_for_integer_attr( + value, elem_ty.width, signed=False + ) + attr = ir.IntegerAttr.get(elem_ty, value) + else: + attr = ir.FloatAttr.get(elem_ty, value) + value = ir.DenseElementsAttr.get_splat(ty, attr) + elif isinstance(ty, ir.FloatType) and isinstance(value, (bool, int)): + value = float(value) + elif is_integer_like_type(ty) and isinstance(value, float): + value = int(value) + else: + raise DSLNotImplemented(f"type {ty} is not supported") + + if isinstance(value, int) and isinstance(ty, ir.IntegerType) and signed is False: + value = _python_int_for_integer_attr(value, ty.width, signed=False) + + return arith.constant(ty, value, loc=loc, ip=ip) + + +def _dispatch_to_rhs_r_op(op: Callable[..., "ArithValue"]) -> Callable[..., Any]: + """Decorator that dispatches to the right-hand-side's reverse operation. + + If the other operand is not an ArithValue or is a subclass (more specific) + of ArithValue, this allows proper method resolution for binary operations. + """ + + def wrapper( + self: "ArithValue", other: Union[int, float, bool, "ArithValue"], **kwargs: Any + ) -> Any: + if not isinstance(other, ArithValue): + if not isinstance(other, (int, float, bool)): + return NotImplemented + + return op(self, other, **kwargs) + + return wrapper + + +def _binary_op(op: Callable[..., "ArithValue"]) -> Callable[..., "ArithValue"]: + """ + Decorator to check if the 'other' argument is an ArithValue. + If not, returns NotImplemented. + """ + + def wrapper( + self: "ArithValue", other: Union[int, float, bool, "ArithValue"], **kwargs: Any + ) -> "ArithValue": + if isinstance(other, (int, float, bool)): + other = const(other, self.type, signed=self.signed).with_signedness( + self.signed + ) + + return op(self, other, **kwargs) + + return wrapper + + +# Operator overloading +@ir.register_value_caster(ir.Float4E2M1FNType.static_typeid) +@ir.register_value_caster(ir.Float6E2M3FNType.static_typeid) +@ir.register_value_caster(ir.Float6E3M2FNType.static_typeid) +@ir.register_value_caster(ir.Float8E4M3FNType.static_typeid) +@ir.register_value_caster(ir.Float8E4M3B11FNUZType.static_typeid) +@ir.register_value_caster(ir.Float8E5M2Type.static_typeid) +@ir.register_value_caster(ir.Float8E4M3Type.static_typeid) +@ir.register_value_caster(ir.Float8E8M0FNUType.static_typeid) +@ir.register_value_caster(ir.BF16Type.static_typeid) +@ir.register_value_caster(ir.F16Type.static_typeid) +@ir.register_value_caster(ir.FloatTF32Type.static_typeid) +@ir.register_value_caster(ir.F32Type.static_typeid) +@ir.register_value_caster(ir.F64Type.static_typeid) +@ir.register_value_caster(ir.IntegerType.static_typeid) +@ir.register_value_caster(ir.RankedTensorType.static_typeid) +class ArithValue(ir.Value): + """Overloads operators for MLIR's Arith dialects binary operations.""" + + @dsl_user_op + def __init__( + self, + v: Union[int, ir.Value], + signed: Union[bool, None] = None, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> None: + if isinstance(v, int): + v = const(v, self.type, signed=signed, loc=loc, ip=ip) + super().__init__(v) + + elem_ty = element_type(self.type) + self.is_float = isinstance(elem_ty, ir.FloatType) + # arith dialect consider `1` in `i1` as `-1`, treat it as unsigned for DSL + self.signed = signed and elem_ty.width > 1 + + @dsl_user_op + def ir_value( + self, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> ir.Value: + return self + + def with_signedness(self, signed: Union[bool, None]) -> "ArithValue": + return type(self)(self, signed) + + @dsl_user_op + def __neg__( + self, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + if self.type == T.bool(): + raise TypeError( + "Negation, the operator `-` is not supported for boolean type" + ) + + if self.is_float: + return arith.negf(self, loc=loc, ip=ip) + else: + c0 = arith.constant(self.type, 0, loc=loc, ip=ip) + return arith.subi(c0, self, loc=loc, ip=ip) + + @dsl_user_op + @_binary_op + def __pow__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + if self.is_float and other.is_float: + return math.powf(self, other, loc=loc, ip=ip) + elif self.is_float and not other.is_float: + return math.fpowi(self, other, loc=loc, ip=ip) + elif not self.is_float and other.is_float: + lhs = itofp(self, self.signed, T.f32(), loc=loc, ip=ip) + rhs = cvtf(other, T.f32(), loc=loc, ip=ip) + return math.powf(lhs, rhs, loc=loc, ip=ip) + elif not self.is_float and not other.is_float: + return math.ipowi(self, other, loc=loc, ip=ip) + else: + raise DSLNotImplemented(f"Unsupported '{self} ** {other}'") + + @dsl_user_op + @_binary_op + def __rpow__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + return other.__pow__(self, loc=loc, ip=ip) + + # arith operators + @dsl_user_op + @_dispatch_to_rhs_r_op + @_binary_op + def __add__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + if self.is_float: + return arith.addf(self, other, loc=loc, ip=ip) + else: + return arith.addi(self, other, loc=loc, ip=ip) + + @dsl_user_op + @_dispatch_to_rhs_r_op + @_binary_op + def __sub__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + if self.is_float: + return arith.subf(self, other, loc=loc, ip=ip) + else: + return arith.subi(self, other, loc=loc, ip=ip) + + @dsl_user_op + @_dispatch_to_rhs_r_op + @_binary_op + def __mul__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + if self.is_float: + return arith.mulf(self, other, loc=loc, ip=ip) + else: + return arith.muli(self, other, loc=loc, ip=ip) + + @dsl_user_op + @_dispatch_to_rhs_r_op + @_binary_op + def __truediv__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + if self.is_float: + return arith.divf(self, other, loc=loc, ip=ip) + else: + lhs = itofp(self, self.signed, T.f32(), loc=loc, ip=ip) + rhs = itofp(other, other.signed, T.f32(), loc=loc, ip=ip) + return arith.divf(lhs, rhs, loc=loc, ip=ip) + + @dsl_user_op + @_dispatch_to_rhs_r_op + @_binary_op + def __floordiv__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + if self.is_float: + q = arith.divf(self, other, loc=loc, ip=ip) + return math.floor(q, loc=loc, ip=ip) + elif self.signed != False: # noqa: E712 + return arith.floordivsi(self, other, loc=loc, ip=ip) + else: + return arith.divui(self, other, loc=loc, ip=ip) + + @dsl_user_op + @_dispatch_to_rhs_r_op + @_binary_op + def __mod__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + if self.is_float: + return arith.remf(self, other, loc=loc, ip=ip) + elif self.signed != False: # noqa: E712 + return arith.remsi(self, other, loc=loc, ip=ip) + else: + return arith.remui(self, other, loc=loc, ip=ip) + + @dsl_user_op + @_binary_op + def __radd__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + return other.__add__(self, loc=loc, ip=ip) + + @dsl_user_op + @_binary_op + def __rsub__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + return other.__sub__(self, loc=loc, ip=ip) + + @dsl_user_op + @_binary_op + def __rmul__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + return other.__mul__(self, loc=loc, ip=ip) + + @dsl_user_op + @_binary_op + def __rtruediv__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + return other.__truediv__(self, loc=loc, ip=ip) + + @dsl_user_op + @_binary_op + def __rfloordiv__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + return other.__floordiv__(self, loc=loc, ip=ip) + + @dsl_user_op + @_binary_op + def __rmod__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + return other.__mod__(self, loc=loc, ip=ip) + + # Comparison operators (comparison doesn't have right-hand-side variants) + @dsl_user_op + @_dispatch_to_rhs_r_op + @_binary_op + def __lt__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + if self.is_float: + return arith.cmpf(arith.CmpFPredicate.OLT, self, other, loc=loc, ip=ip) + elif self.signed != False: # noqa: E712 + return arith.cmpi(arith.CmpIPredicate.slt, self, other, loc=loc, ip=ip) + else: + return arith.cmpi(arith.CmpIPredicate.ult, self, other, loc=loc, ip=ip) + + @dsl_user_op + @_dispatch_to_rhs_r_op + @_binary_op + def __le__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + if self.is_float: + return arith.cmpf(arith.CmpFPredicate.OLE, self, other, loc=loc, ip=ip) + elif self.signed != False: # noqa: E712 + return arith.cmpi(arith.CmpIPredicate.sle, self, other, loc=loc, ip=ip) + else: + return arith.cmpi(arith.CmpIPredicate.ule, self, other, loc=loc, ip=ip) + + @dsl_user_op + @_dispatch_to_rhs_r_op + @_binary_op + def __eq__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + if self.is_float: + return arith.cmpf(arith.CmpFPredicate.OEQ, self, other, loc=loc, ip=ip) + else: + return arith.cmpi(arith.CmpIPredicate.eq, self, other, loc=loc, ip=ip) + + @dsl_user_op + @_dispatch_to_rhs_r_op + @_binary_op + def __ne__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + if self.is_float: + # In Python, bool(float("nan")) is True, so use unordered comparison here + return arith.cmpf(arith.CmpFPredicate.UNE, self, other, loc=loc, ip=ip) + else: + return arith.cmpi(arith.CmpIPredicate.ne, self, other, loc=loc, ip=ip) + + @dsl_user_op + @_dispatch_to_rhs_r_op + @_binary_op + def __gt__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + if self.is_float: + return arith.cmpf(arith.CmpFPredicate.OGT, self, other, loc=loc, ip=ip) + elif self.signed != False: # noqa: E712 + return arith.cmpi(arith.CmpIPredicate.sgt, self, other, loc=loc, ip=ip) + else: + return arith.cmpi(arith.CmpIPredicate.ugt, self, other, loc=loc, ip=ip) + + @dsl_user_op + @_dispatch_to_rhs_r_op + @_binary_op + def __ge__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + if self.is_float: + return arith.cmpf(arith.CmpFPredicate.OGE, self, other, loc=loc, ip=ip) + elif self.signed != False: # noqa: E712 + return arith.cmpi(arith.CmpIPredicate.sge, self, other, loc=loc, ip=ip) + else: + return arith.cmpi(arith.CmpIPredicate.uge, self, other, loc=loc, ip=ip) + + # Unary operators + @dsl_user_op + def __abs__( + self, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + if self.is_float: + return math.absf(self, loc=loc, ip=ip) + else: + return math.absi(self, loc=loc, ip=ip) + + @dsl_user_op + def __invert__( + self, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + # For i1 (boolean) types, the all-ones value is 1, not -1. + # Using -1 with i1 vectors causes arith.constant to produce a + # type mismatch (e.g. vector<32xi0> instead of vector<32xi1>). + all_ones = 1 if element_type(self.type).width == 1 else -1 + return arith.xori(self, const(all_ones, self.type)) + + # Bitwise operations + @dsl_user_op + @_dispatch_to_rhs_r_op + @_binary_op + def __and__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + return arith.andi(self, other, loc=loc, ip=ip) + + @dsl_user_op + @_dispatch_to_rhs_r_op + @_binary_op + def __or__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + return arith.ori(self, other, loc=loc, ip=ip) + + @dsl_user_op + @_dispatch_to_rhs_r_op + @_binary_op + def __xor__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + return arith.xori(self, other, loc=loc, ip=ip) + + @dsl_user_op + @_dispatch_to_rhs_r_op + @_binary_op + def __rshift__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + if self.signed != False: # noqa: E712 + return arith.shrsi(self, other, loc=loc, ip=ip) + else: + return arith.shrui(self, other, loc=loc, ip=ip) + + @dsl_user_op + @_dispatch_to_rhs_r_op + @_binary_op + def __lshift__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + return arith.shli(self, other, loc=loc, ip=ip) + + @dsl_user_op + @_binary_op + def __rand__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + return arith.andi(other, self, loc=loc, ip=ip) + + @dsl_user_op + @_binary_op + def __ror__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + return arith.ori(other, self, loc=loc, ip=ip) + + @dsl_user_op + @_binary_op + def __rxor__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + return arith.xori(other, self, loc=loc, ip=ip) + + @dsl_user_op + @_binary_op + def __rrshift__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + return other.__rshift__(self, loc=loc, ip=ip) + + @dsl_user_op + @_binary_op + def __rlshift__( + self, + other: "ArithValue", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "ArithValue": + return other.__lshift__(self, loc=loc, ip=ip) + + def __hash__(self) -> int: + return super().__hash__() + + def __str__(self) -> str: + return "?" + + def __repr__(self) -> str: + return self.__str__() + + +def _min( + lhs: ir.Value, + rhs: ir.Value, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> ir.Value: + """ + This function provides a unified interface for building arith min + + Assuming the operands have the same type + """ + from ..base_dsl.dsl import is_dynamic_expression + + if not is_dynamic_expression(lhs): + if not is_dynamic_expression(rhs): + return min(lhs, rhs) + else: + lhs = arith.constant(rhs.type, lhs, loc=loc, ip=ip) + else: + if not is_dynamic_expression(rhs): + rhs = arith.constant(lhs.type, rhs, loc=loc, ip=ip) + + # Handle vector types + if isinstance(lhs.type, T.VectorType): + elem_type = lhs.type.element_type + if is_integer_like_type(elem_type): + assert hasattr(lhs, "signed"), ( + "Should have attribute `signed`, must be a bug" + ) + if lhs.signed != False: # noqa: E712 + return arith.minsi(lhs, rhs, loc=loc, ip=ip) + else: + return arith.minui(lhs, rhs, loc=loc, ip=ip) + else: + return arith.minimumf(lhs, rhs, loc=loc, ip=ip) + elif is_integer_like_type(lhs.type): + assert hasattr(lhs, "signed"), "Should have attribute `signed`, must be a bug" + if lhs.signed != False: # noqa: E712 + return arith.minsi(lhs, rhs, loc=loc, ip=ip) + else: + return arith.minui(lhs, rhs, loc=loc, ip=ip) + else: + return arith.minimumf(lhs, rhs, loc=loc, ip=ip) + + +def _max( + lhs: ir.Value, + rhs: ir.Value, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> ir.Value: + """ + This function provides a unified interface for building arith max + + Assuming the operands have the same type + """ + from ..base_dsl.dsl import is_dynamic_expression + + if not is_dynamic_expression(lhs): + if not is_dynamic_expression(rhs): + return max(lhs, rhs) + else: + lhs = arith.constant(rhs.type, lhs, loc=loc, ip=ip) + else: + if not is_dynamic_expression(rhs): + rhs = arith.constant(lhs.type, rhs, loc=loc, ip=ip) + # Handle vector types + if isinstance(lhs.type, T.VectorType): + elem_type = lhs.type.element_type + if isinstance(elem_type, ir.IntegerType): + assert hasattr(lhs, "signed"), ( + "Should have attribute `signed`, must be a bug" + ) + if lhs.signed != False: # noqa: E712 + return arith.maxsi(lhs, rhs, loc=loc, ip=ip) + else: + return arith.maxui(lhs, rhs, loc=loc, ip=ip) + else: + return arith.maximumf(lhs, rhs, loc=loc, ip=ip) + elif is_integer_like_type(lhs.type): + assert hasattr(lhs, "signed"), "Should have attribute `signed`, must be a bug" + if lhs.signed != False: # noqa: E712 + return arith.maxsi(lhs, rhs, loc=loc, ip=ip) + else: + return arith.maxui(lhs, rhs, loc=loc, ip=ip) + else: + return arith.maximumf(lhs, rhs, loc=loc, ip=ip) + + +def __getattr__(name: str) -> object: + if name == "Vector": + from .vector import Vector + + return Vector + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/CuTeDSL/_mlir_helpers/dialect_proxy.py b/python/CuTeDSL/_mlir_helpers/dialect_proxy.py new file mode 100644 index 000000000..b573a6cfd --- /dev/null +++ b/python/CuTeDSL/_mlir_helpers/dialect_proxy.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +import enum +import types +from collections.abc import Callable + + +class DialectAutoConvertProxy: + """ + Proxy that wraps a raw MLIR dialect module, auto-converting DSL types + (anything with an ``.ir_value()`` method) to ``ir.Value`` when calling + dialect operations. + + This enables users to write cleaner code without explicit + ``.ir_value()`` calls:: + + # Before (raw dialect module): + nvvm.shfl_sync(T.i32(), Int32(mask).ir_value(), ...) + + # After (proxied): + nvvm.shfl_sync(T.i32(), Int32(mask), ...) + + Non-callable attributes and enum classes are passed through unchanged + so that attribute access like ``nvvm.ShflKind.idx`` still works. + + Parameters + ---------- + dialect_module + The raw MLIR dialect module to wrap + (e.g. ``cutlass._mlir.dialects.nvvm``). + """ + + def __init__(self, dialect_module: types.ModuleType) -> None: + self._module = dialect_module + self._wrapped_cache: dict[str, Callable[..., object]] = {} + + @staticmethod + def _convert_arg( + arg: object, + loc: object | None, + ip: object | None, + ) -> object: + """Convert DSL value wrappers to ir.Value; pass primitives/containers through. + + Numeric primitives (int/float/bool) are immediates/attributes, not + operands, so they pass through unchanged. + DSL ``Numeric`` types are not int/float subclasses, so genuine operands + still convert below. + """ + if isinstance(arg, (bool, int, float)): + return arg + if hasattr(arg, "ir_value") and callable(arg.ir_value): + try: + return arg.ir_value(loc=loc, ip=ip) + except TypeError: + # Some ir_value() methods (e.g. Array) don't accept loc/ip. + return arg.ir_value() + if isinstance(arg, (list, tuple)): + converted = [ + DialectAutoConvertProxy._convert_arg(item, loc, ip) for item in arg + ] + return type(arg)(converted) + return arg + + def __getattr__(self, name: str) -> object: + attr = getattr(self._module, name) + + # Non-callable attributes and enum classes pass through + # unchanged. Enum classes need attribute access (e.g. + # ShflKind.idx), but MLIR operation classes should be + # wrapped for argument conversion. + if not callable(attr) or isinstance(attr, enum.EnumMeta): + return attr + + # Use cache for wrapped callables + if name not in self._wrapped_cache: + + def _make_wrapper( + func: Callable[..., object], + ) -> Callable[..., object]: + def wrapped( + *args: object, + loc: object | None = None, + ip: object | None = None, + **kwargs: object, + ) -> object: + converted_args = tuple( + DialectAutoConvertProxy._convert_arg(arg, loc, ip) + for arg in args + ) + converted_kwargs = { + k: DialectAutoConvertProxy._convert_arg(v, loc, ip) + for k, v in kwargs.items() + } + return func( + *converted_args, + loc=loc, + ip=ip, + **converted_kwargs, + ) + + return wrapped + + self._wrapped_cache[name] = _make_wrapper(attr) + + return self._wrapped_cache[name] + + def __dir__(self) -> list[str]: + return dir(self._module) diff --git a/python/CuTeDSL/_mlir_helpers/gpu.py b/python/CuTeDSL/_mlir_helpers/gpu.py new file mode 100644 index 000000000..566f175ea --- /dev/null +++ b/python/CuTeDSL/_mlir_helpers/gpu.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +""" +This module provides MLIR GPU Dialect helper functions +""" + +from .._mlir import ir +from .._mlir.dialects import gpu, arith, scf +from .._mlir.extras import types as _T + +from ..base_dsl.common import * + +# ============================================================================= +# GPU Dialect Helper functions +# ============================================================================= + + +def create_async_token() -> ir.Value: + token_ty = gpu.AsyncTokenType.get() + token = gpu.wait(token_ty, []) + return token + + +def printf(fmt: str, *args: ir.Value, threadNumber: int = -1) -> None: + """Generate gpu.printf OP predicated on threadNumber""" + type_formats = [] + for arg in args: + ty_format = None + if isinstance(arg.type, ir.IndexType): + ty_format = "%llu" + if isinstance(arg.type, ir.IntegerType): + width = ir.IntegerType(arg.type).width + if width == 64: + ty_format = "%llu" + elif width == 32: + ty_format = "%d" + elif width == 1: + ty_format = "%i" + if isinstance(arg.type, ir.F32Type): + ty_format = "%f" + if ty_format is None: + raise DSLNotImplemented(arg.type) + type_formats.append(ty_format) + if threadNumber == -1: + gpu.printf(fmt.format(*type_formats) + "\n", args) + if threadNumber != -1: + tidx = gpu.thread_id(gpu.Dimension.x) + predicate = arith.cmpi( + arith.CmpIPredicate.eq, tidx, arith.constant(_T.index(), threadNumber) + ) + if_op = scf.IfOp(predicate) + with ir.InsertionPoint(if_op.then_block): + gpu.printf(fmt.format(*type_formats) + "\n", args) + scf.yield_([]) diff --git a/python/CuTeDSL/_mlir_helpers/lru_cache_ir.py b/python/CuTeDSL/_mlir_helpers/lru_cache_ir.py new file mode 100644 index 000000000..2c34c1ca5 --- /dev/null +++ b/python/CuTeDSL/_mlir_helpers/lru_cache_ir.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +""" +This module provides @lru_cache_ir +It extends functools.lru_cache with IR Context awareness. + +Example usage: +from cutlass import ir +from lru_cache_ir import lru_cache_ir + +@lru_cache_ir(ir, maxsize=128, typed=False) +def make_layout(...): +... + +""" + +from functools import lru_cache, wraps +from typing import Any, Callable + +from .._mlir import ir +from .nanobind_compat import has_insertion_point + + +def get_ir_context(func: Any) -> Any: + """ + Return the context for given func called under ir. + Currently the context includes MLIRContext and InsertionPoint. + """ + if not ir: + return None + try: + ctx = ir.Context.current + except (ValueError, RuntimeError): + return None + if has_insertion_point(): + return (ctx, ir.InsertionPoint.current) + return (ctx,) + + +def lru_cache_ir(maxsize: int = 128, typed: bool = True) -> Callable[..., Any]: + """ + Applies an LRU cache to a given function, with awareness of IR context. + + Usage is similar to functools.lru_cache while taking `ir` as required argument. + + :param ir: The IR object from which to derive the context by `get_ir_context` + :param maxsize: Max cache size, same as functools.lru_cache + :param typed: Whether params are type-sensitive, default to True as IR is type-sensitive + """ + + def decorator(func: Callable[..., Any]) -> Callable[..., Any]: + # Use functools.lru_cache with a custom wrapper to control the key generation + @lru_cache(maxsize=maxsize, typed=typed) + def cached_func(context: Any, *args: Any, **kwargs: Any) -> Any: + return func(*args, **kwargs) + + @wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + try: + # Call the cached function with the context + return cached_func(get_ir_context(func), *args, **kwargs) + except (RuntimeError, TypeError): + return func(*args, **kwargs) + + # Expose cache-related methods for introspection + wrapper.cache_clear = cached_func.cache_clear # type: ignore[attr-defined] + wrapper.cache_info = cached_func.cache_info # type: ignore[attr-defined] + return wrapper + + return decorator diff --git a/python/CuTeDSL/_mlir_helpers/math.py b/python/CuTeDSL/_mlir_helpers/math.py new file mode 100644 index 000000000..75cf14e95 --- /dev/null +++ b/python/CuTeDSL/_mlir_helpers/math.py @@ -0,0 +1,2213 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary + +""" +Unified NVIDIA GPU math API. + +This module exposes the full set of NVIDIA GPU math capabilities as a +single, stable Python surface. Each call is dispatched to the cleanest +layer that currently provides the requested semantics: + + 1. MLIR ``math`` / ``arith`` dialect ops — target-agnostic, + vectorizable, pattern-match friendly. The default and + ``fastmath`` paths use this layer. + 2. ``nvvm.*`` dialect ops — first-class MLIR ops for a subset of + NVIDIA-specific capabilities (``nvvm.fabs``, ``nvvm.fmin`` / + ``fmax``, ``nvvm.rcp_approx_ftz_f``). Used where available. + 3. ``llvm.call_intrinsic "llvm.nvvm.*"`` — direct LLVM NVVM + intrinsics for PTX modifier combinations (approx, ftz, + specific rounding modes, ``div.full``) that don't yet have + first-class MLIR ops. + +Why this module exists +---------------------- +Upstream MLIR's ``math`` / ``arith`` dialects are chartered +target-agnostic and intentionally don't carry NVIDIA-specific attributes +(approx, ftz, rounding mode, ``div.full``). The NVVM dialect +promotes a selected subset to first-class ops but deliberately +doesn't enumerate every PTX combination. That leaves a layer — the +full NVIDIA GPU math capability surface, presented to end users — +that neither upstream owner wants to fill. This module fills it. + +As upstream ``math`` / ``arith`` gain rounding-mode attributes +(in progress) and the NVVM dialect promotes more PTX variants, +intrinsic-call paths here shrink. The user-facing API stays stable. + +Supported operand types +----------------------- + - Numeric scalars (Float32, Float64, Int32, etc.) + - ArithValue / Vector (MLIR value types; subclasses like + ``TensorSSA`` go through the Vector dispatch path automatically, + and :meth:`Vector._wrap_like` preserves subclass-specific metadata + on per-element results) + - Python ``float`` / ``int`` / ``bool`` literals, promoted to + ``Float32`` / ``Int32`` via :func:`_coerce_operand`. + +Consumers that need to layer their own framework-specific types (e.g., +memory-tracked array wrappers) can add a thin adapter that unwraps to a +supported type, delegates here, and rewraps the result. +""" + +from enum import Enum +from typing import Any, Optional, Union + +from .._mlir import ir +from .._mlir.extras import types as T +from .._mlir.dialects import arith, math as math_dialect, llvm, nvgpu, nvvm +from .arith import ArithValue, is_float_type, element_type +from .op import dsl_user_op +from .vector import Vector + +from ..base_dsl.typing import Numeric, Float32, Int32 + +# ============================================================================= +# Type alias +# ============================================================================= + +# Accepted operand types. ``Vector`` covers CuTeDSL's ``TensorSSA`` via +# inheritance (e.g., ``TensorSSA``). Additional framework-specific wrappers +# should be layered as thin adapters in their own package and delegate here. +# Python ``float``/``int``/``bool`` literals are promoted to ``Float32`` / +# ``Int32`` automatically so callers can write ``mul(x, 2.0)`` or +# ``add(x, 1)`` without manual wrapping. +MathOperand = Union[Numeric, Vector, float, int, bool] + + +# ============================================================================= +# Enums +# ============================================================================= + + +class RoundingMode(str, Enum): + """IEEE 754 rounding modes for floating-point operations.""" + + NEAREST_EVEN = "rn" + ZERO = "rz" + UP = "rp" + DOWN = "rm" + + +# ============================================================================= +# Internal helpers +# ============================================================================= + + +def _coerce_operand(x: MathOperand) -> MathOperand: + """Promote Python numeric literals to Numeric scalars. + + Python ``bool``/``int`` -> :class:`Int32`, ``float`` -> :class:`Float32`. + ``bool`` is checked before ``int`` because ``bool`` subclasses ``int``. + Non-literal operands (Numeric, ArithValue, Vector) pass through + unchanged. + """ + if isinstance(x, bool): + return Int32(int(x)) + if isinstance(x, float): + return Float32(x) + if isinstance(x, int): + return Int32(x) + return x + + +def _numeric_type_name(x: MathOperand) -> str: + """Return the scalar element type name of a coerced operand. + + Returns the class name of the scalar type (e.g. ``"Float32"``, + ``"Int32"``). For ``Vector`` the dtype's class name is returned + (``type(x.dtype).__name__``). + """ + if isinstance(x, Vector): + return type(x.dtype).__name__ + return type(x).__name__ + + +def _validate_same_numeric_type(op_name: str, *operands: MathOperand) -> None: + """Raise :class:`TypeError` if operands don't share one numeric type. + + Python literals are promoted via :func:`_coerce_operand` before the + check, so the comparison is purely on Numeric / Vector element types + after promotion. + + Mixing typed values (e.g. ``fma(Float64_val, Float32_val, Int32_val)``) + would otherwise produce cryptic MLIR type-mismatch errors deep in the + compiler. This helper surfaces the problem immediately. + + We promote Python literals (for ergonomics) but do *not* auto-promote + heterogeneous typed values — if a user mixes Float32 and Float16, they + need to convert explicitly. + """ + coerced = [_coerce_operand(op) for op in operands] + # Raw ArithValue operands (e.g. reduction results threaded through an + # threaded through a task loop) don't carry a Numeric class, so comparing + # ``type(op).__name__`` would produce false positives. Vector subclasses + # ArithValue here, so we must exclude it explicitly. + if any(isinstance(op, ArithValue) and not isinstance(op, Vector) for op in coerced): + return + type_names = [_numeric_type_name(op) for op in coerced] + if len(set(type_names)) > 1: + raise TypeError( + f"All operands to {op_name} must have the same numeric type, " + f"got {', '.join(type_names)}" + ) + + +def _check_vector_consistency(op_name: str, *operands: MathOperand) -> None: + """Validate that vector-like and scalar operands aren't mixed. + + MLIR math/arith ops require all operands to share a shape. We forbid + mixing a Vector with a Numeric scalar early so the user sees a clear + TypeError instead of a cryptic MLIR error later. + """ + coerced = [_coerce_operand(op) for op in operands] + is_vector_like = [isinstance(op, Vector) for op in coerced] + if any(is_vector_like) and not all(is_vector_like): + vector_type = next( + type(op).__name__ for op, is_v in zip(coerced, is_vector_like) if is_v + ) + scalar_type = next( + type(op).__name__ for op, is_v in zip(coerced, is_vector_like) if not is_v + ) + raise TypeError( + f"{op_name}: Expected all inputs to be of type {vector_type}, " + f"but got {scalar_type} for another operand" + ) + + +def _get_ir_value( + x: MathOperand, *, ip: Optional[ir.InsertionPoint] = None +) -> ir.Value: + """Extract MLIR IR value from any supported operand type.""" + x = _coerce_operand(x) + if isinstance(x, ir.Value): + # Vector (and its subclasses) are already ir.Values — use directly. + return x + if not isinstance(x, Numeric): + raise TypeError(f"Expected a Vector or Numeric, but got {type(x).__name__}") + # Numeric (Float32, etc.) — has ir_value() method + return x.ir_value(ip=ip) + + +def _get_element_type(x: MathOperand) -> ir.Type: + """Get the scalar element type of an operand.""" + v = _get_ir_value(x) + return element_type(v.type) + + +def _is_float(x: MathOperand) -> bool: + """Check if operand is a floating-point type.""" + return is_float_type(_get_element_type(x)) + + +def _is_f32(x: MathOperand) -> bool: + return _get_element_type(x) == T.f32() + + +def _is_f64(x: MathOperand) -> bool: + return _get_element_type(x) == T.f64() + + +def _is_unsigned_int(x: MathOperand) -> bool: + """Check if operand is an unsigned integer type.""" + return not _is_float(x) and getattr(x, "signed", True) is False + + +def _wrap_result(original: MathOperand, result_ir: ir.Value) -> MathOperand: + """Wrap an MLIR result back into the original operand's type.""" + if isinstance(original, Vector): + # Polymorphic wrap: Vector subclasses (e.g., cute.TensorSSA) use + # ``_wrap_like`` to preserve their own metadata (CuTe nested shape, + # layout) when a math op returns a new SSA value. Fall back to a + # plain Vector construction when the base-class hook isn't present + # (compat with older Vector that doesn't yet define ``_wrap_like``). + if hasattr(original, "_wrap_like"): + return original._wrap_like(result_ir) + return Vector(result_ir, dtype=original._dtype) + if isinstance(original, ArithValue): + return result_ir # ArithValue IS an ir.Value + # Coerce Python literals so we return a Numeric subclass instance instead + # of attempting to construct e.g. ``float(ir.Value)``. + original = _coerce_operand(original) + # Numeric — wrap in the same Numeric subclass + return type(original)(result_ir) + + +def _get_fastmath_flag(fastmath: bool) -> arith.FastMathFlags: + return arith.FastMathFlags.fast if fastmath else arith.FastMathFlags.none + + +def _get_type_suffix(x: MathOperand) -> str: + """Get NVVM type suffix: 'f' for f32, 'd' for f64.""" + if _is_f32(x): + return "f" + elif _is_f64(x): + return "d" + else: + raise TypeError("NVVM math intrinsics require f32 or f64 scalar operands") + + +def _get_llvm_type(x: MathOperand) -> ir.Type: + """Get LLVM result type for NVVM intrinsic.""" + if _is_f32(x): + return T.f32() + elif _is_f64(x): + return T.f64() + else: + raise TypeError("NVVM math intrinsics require f32 or f64 scalar operands") + + +def _const_one( + ty: ir.Type, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> ir.Value: + """Constant 1.0 of a scalar-float or vector-of-float type.""" + if isinstance(ty, ir.VectorType): + attr = ir.DenseElementsAttr.get_splat( + ty, ir.FloatAttr.get(ty.element_type, 1.0) + ) + return arith.constant(ty, attr, loc=loc, ip=ip) + return arith.constant(ty, 1.0, loc=loc, ip=ip) + + +# ============================================================================= +# Validation +# ============================================================================= + + +def _validate_fastmath_exclusive( + op_name: str, + fastmath: bool, + *, + approx: bool = False, + ftz: bool = False, + rounding: Optional[RoundingMode] = None, + full: bool = False, +) -> None: + """Raise if fastmath (math/arith dialect) is combined with NVVM attributes. + + Three valid modes: + - No flags: strict MLIR math/arith op + - fastmath=True: MLIR op with fastmath attribute + - approx/ftz/rounding: NVVM intrinsic + + fastmath and approx/ftz/rounding are mutually exclusive because + they target different MLIR dialects. + """ + nvvm_attrs = [] + if approx: + nvvm_attrs.append("approx") + if ftz: + nvvm_attrs.append("ftz") + if rounding is not None: + nvvm_attrs.append("rounding") + if full: + nvvm_attrs.append("full") + if fastmath and nvvm_attrs: + attr_str = "/".join(nvvm_attrs) + raise ValueError( + f"{op_name}: fastmath and {attr_str} are mutually exclusive. " + f"fastmath uses math/arith dialect (compiler decides lowering); " + f"{attr_str} emit NVVM intrinsics (exact instruction control)." + ) + + +def _validate_ftz_requires_approx(op_name: str, approx: bool, ftz: bool) -> None: + """Raise if ftz is set without approx on transcendentals.""" + if ftz and not approx: + raise ValueError( + f"{op_name}: ftz requires approx=True for transcendental " + f"functions. There is no non-approximate ftz variant." + ) + + +# ============================================================================= +# NVVM intrinsic helpers +# +# Callers land here only when neither arith/math nor the NVVM dialect has a +# matching op for the requested (op, attributes, type) combination. Dialect-op +# coverage audit at time of writing: +# - Scalar float add/sub/mul/fma with rounding+ftz: NVVM has only +# *Packed*F32x2 variants; no scalar ops (nvvm.MulOp is integer). +# - Scalar transcendental .approx (sin/cos/ex2/lg2/tanh/rsqrt/sqrt): no +# NVVM dialect ops. +# - Scalar div.approx / div.full / div.: no NVVM dialect ops. +# - Scalar sqrt.: no NVVM dialect ops. +# - Scalar rcp. / rcp.approx (no ftz): no NVVM dialect ops. +# Ops that *do* have dialect coverage are used directly: nvvm.fabs (abs.ftz), +# nvvm.fmin/fmax (min/max.ftz), nvvm.rcp_approx_ftz_f (rcp scalar f32 +# approx+ftz), nvgpu.rcp (vector rcp approx[.ftz]). +# ============================================================================= + +# NVVM intrinsics that use overloaded types (no .f/.d suffix). +_NVVM_OVERLOADED_OPS = frozenset({"ex2", "tanh"}) + + +def _call_nvvm_unary( + x: MathOperand, + op_name: str, + rounding: Optional[RoundingMode] = None, + ftz: bool = False, + approx: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Call a unary NVVM intrinsic.""" + if isinstance(x, Vector): + # NVVM intrinsics are scalar-only. For vectors (and Arrays that wrap + # Vectors), we can't emit them without per-element unrolling (which + # defeats vectorization). Direct users to the math/arith path (via + # fastmath=True) so the MLIR pipeline can vectorize properly. + attr = "approx" if approx else ("rounding" if rounding is not None else "ftz") + raise TypeError( + f"{op_name}: {attr} is scalar-only (NVVM intrinsics don't " + f"vectorize). For vector inputs, use fastmath=True — the " + f"math/arith dialect lowering will be vectorized by the compiler." + ) + + x_ir = _get_ir_value(x, ip=ip) + result_type = _get_llvm_type(x) + type_suffix = _get_type_suffix(x) + + intrinsic_name = f"llvm.nvvm.{op_name}" + if approx: + intrinsic_name += ".approx" + elif rounding is not None: + intrinsic_name += f".{rounding.value}" + elif ftz: + # ftz without explicit rounding: default to rn + intrinsic_name += ".rn" + if ftz: + intrinsic_name += ".ftz" + if op_name not in _NVVM_OVERLOADED_OPS: + intrinsic_name += f".{type_suffix}" + + result = llvm.call_intrinsic( + result_type, + intrinsic_name, + [x_ir], + [], + [], + loc=loc, + ip=ip, + ) + return _wrap_result(x, result) + + +def _call_nvvm_binary( + x: MathOperand, + y: MathOperand, + op_name: str, + rounding: Optional[RoundingMode] = None, + ftz: bool = False, + approx: bool = False, + full: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Call a binary NVVM intrinsic.""" + if isinstance(x, Vector) or isinstance(y, Vector): + # See _call_nvvm_unary: scalar-only; direct users to fastmath for + # vectorizable math/arith lowering. + attr = ( + "approx" + if approx + else ("full" if full else ("rounding" if rounding is not None else "ftz")) + ) + raise TypeError( + f"{op_name}: {attr} is scalar-only (NVVM intrinsics don't " + f"vectorize). For vector inputs, use fastmath=True — the " + f"math/arith dialect lowering will be vectorized by the compiler." + ) + x_ir = _get_ir_value(x, ip=ip) + y_ir = _get_ir_value(y, ip=ip) + result_type = _get_llvm_type(x) + type_suffix = _get_type_suffix(x) + + intrinsic_name = f"llvm.nvvm.{op_name}" + if approx: + intrinsic_name += ".approx" + elif full: + intrinsic_name += ".full" + elif rounding is not None: + intrinsic_name += f".{rounding.value}" + elif ftz: + # ftz without explicit rounding: default to rn + intrinsic_name += ".rn" + if ftz: + intrinsic_name += ".ftz" + if not full: + intrinsic_name += f".{type_suffix}" + + result = llvm.call_intrinsic( + result_type, + intrinsic_name, + [x_ir, y_ir], + [], + [], + loc=loc, + ip=ip, + ) + return _wrap_result(x, result) + + +def _call_nvvm_ternary( + a: MathOperand, + b: MathOperand, + c: MathOperand, + op_name: str, + rounding: Optional[RoundingMode] = None, + ftz: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Call a ternary NVVM intrinsic (e.g., fma).""" + if any(isinstance(v, Vector) for v in (a, b, c)): + # See _call_nvvm_unary: scalar-only; direct users to fastmath for + # vectorizable math/arith lowering. + attr = "rounding" if rounding is not None else "ftz" + raise TypeError( + f"{op_name}: {attr} is scalar-only (NVVM intrinsics don't " + f"vectorize). For vector inputs, use fastmath=True — the " + f"math/arith dialect lowering will be vectorized by the compiler." + ) + a_ir = _get_ir_value(a, ip=ip) + b_ir = _get_ir_value(b, ip=ip) + c_ir = _get_ir_value(c, ip=ip) + result_type = _get_llvm_type(a) + type_suffix = _get_type_suffix(a) + + intrinsic_name = f"llvm.nvvm.{op_name}" + if rounding is not None: + intrinsic_name += f".{rounding.value}" + else: + intrinsic_name += ".rn" + if ftz: + intrinsic_name += ".ftz" + intrinsic_name += f".{type_suffix}" + + result = llvm.call_intrinsic( + result_type, + intrinsic_name, + [a_ir, b_ir, c_ir], + [], + [], + loc=loc, + ip=ip, + ) + return _wrap_result(a, result) + + +def _needs_nvvm_intrinsic(rounding: Optional[RoundingMode], ftz: bool) -> bool: + """Check if we need NVVM intrinsics (any explicit rounding or ftz).""" + if ftz: + return True + if rounding is not None: + return True + return False + + +# ============================================================================= +# Unary math helper +# ============================================================================= + + +def _unary_math_op( + x: MathOperand, + float_op: Any, + int_op: Any = None, + fastmath: bool = False, + op_name: str = "math_op", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Apply a unary MLIR math operation.""" + x_ir = _get_ir_value(x, ip=ip) + fastmath_flag = _get_fastmath_flag(fastmath) + + if _is_float(x): + result = float_op(x_ir, fastmath=fastmath_flag, loc=loc, ip=ip) + else: + if int_op is None: + raise TypeError(f"{op_name} is only supported for floating-point types") + result = int_op(x_ir, loc=loc, ip=ip) + + return _wrap_result(x, result) + + +# ============================================================================= +# Transcendentals with NVVM approx intrinsics +# ============================================================================= + + +@dsl_user_op +def sin( + x: MathOperand, + fastmath: bool = False, + approx: bool = False, + ftz: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise sine of the input operand. + + :param a: Input operand (in radians) + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the sine of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = sin(a) # Compute sine + """ + _validate_fastmath_exclusive("sin", fastmath, approx=approx, ftz=ftz) + _validate_ftz_requires_approx("sin", approx, ftz) + if approx: + return _call_nvvm_unary(x, "sin", approx=True, ftz=ftz, loc=loc, ip=ip) + return _unary_math_op(x, math_dialect.sin, None, fastmath, "sin", loc=loc, ip=ip) + + +@dsl_user_op +def cos( + x: MathOperand, + fastmath: bool = False, + approx: bool = False, + ftz: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise cosine of the input operand. + + :param a: Input operand (in radians) + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the cosine of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = cos(a) # Compute cosine + """ + _validate_fastmath_exclusive("cos", fastmath, approx=approx, ftz=ftz) + _validate_ftz_requires_approx("cos", approx, ftz) + if approx: + return _call_nvvm_unary(x, "cos", approx=True, ftz=ftz, loc=loc, ip=ip) + return _unary_math_op(x, math_dialect.cos, None, fastmath, "cos", loc=loc, ip=ip) + + +@dsl_user_op +def exp( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise exponential of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the exponential of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = exp(a) # Compute exponential + """ + return _unary_math_op(x, math_dialect.exp, None, fastmath, "exp", loc=loc, ip=ip) + + +@dsl_user_op +def exp2( + x: MathOperand, + fastmath: bool = False, + approx: bool = False, + ftz: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise base-2 exponential of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing 2 raised to the power of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = exp2(a) # Compute 2^x + """ + _validate_fastmath_exclusive("exp2", fastmath, approx=approx, ftz=ftz) + _validate_ftz_requires_approx("exp2", approx, ftz) + if approx: + return _call_nvvm_unary(x, "ex2", approx=True, ftz=ftz, loc=loc, ip=ip) + return _unary_math_op(x, math_dialect.exp2, None, fastmath, "exp2", loc=loc, ip=ip) + + +@dsl_user_op +def log( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise natural logarithm of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the natural logarithm of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = log(a) # Compute natural logarithm + """ + return _unary_math_op(x, math_dialect.log, None, fastmath, "log", loc=loc, ip=ip) + + +@dsl_user_op +def log2( + x: MathOperand, + fastmath: bool = False, + approx: bool = False, + ftz: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise base-2 logarithm of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the base-2 logarithm of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = log2(a) # Compute log base 2 + """ + _validate_fastmath_exclusive("log2", fastmath, approx=approx, ftz=ftz) + _validate_ftz_requires_approx("log2", approx, ftz) + if approx: + return _call_nvvm_unary(x, "lg2", approx=True, ftz=ftz, loc=loc, ip=ip) + return _unary_math_op(x, math_dialect.log2, None, fastmath, "log2", loc=loc, ip=ip) + + +@dsl_user_op +def log10( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise base-10 logarithm of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the base-10 logarithm of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = log10(a) # Compute log base 10 + """ + return _unary_math_op( + x, math_dialect.log10, None, fastmath, "log10", loc=loc, ip=ip + ) + + +@dsl_user_op +def tanh( + x: MathOperand, + fastmath: bool = False, + approx: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise hyperbolic tangent of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the hyperbolic tangent of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = tanh(a) # Compute hyperbolic tangent + """ + _validate_fastmath_exclusive("tanh", fastmath, approx=approx) + if approx: + return _call_nvvm_unary(x, "tanh", approx=True, loc=loc, ip=ip) + return _unary_math_op(x, math_dialect.tanh, None, fastmath, "tanh", loc=loc, ip=ip) + + +@dsl_user_op +def rsqrt( + x: MathOperand, + fastmath: bool = False, + approx: bool = False, + ftz: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise reciprocal square root of the input operand. + + Computes 1/√x element-wise. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the reciprocal square root of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = rsqrt(a) # Compute 1/√x + """ + _validate_fastmath_exclusive("rsqrt", fastmath, approx=approx, ftz=ftz) + _validate_ftz_requires_approx("rsqrt", approx, ftz) + if approx: + return _call_nvvm_unary(x, "rsqrt", approx=True, ftz=ftz, loc=loc, ip=ip) + return _unary_math_op( + x, math_dialect.rsqrt, None, fastmath, "rsqrt", loc=loc, ip=ip + ) + + +@dsl_user_op +def sqrt( + x: MathOperand, + fastmath: bool = False, + approx: bool = False, + ftz: bool = False, + rounding: Optional[RoundingMode] = None, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise square root of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the square root of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = sqrt(a) # Compute square root + """ + _validate_fastmath_exclusive( + "sqrt", fastmath, approx=approx, ftz=ftz, rounding=rounding + ) + if approx: + return _call_nvvm_unary(x, "sqrt", approx=True, ftz=ftz, loc=loc, ip=ip) + if _needs_nvvm_intrinsic(rounding, ftz): + return _call_nvvm_unary(x, "sqrt", rounding=rounding, ftz=ftz, loc=loc, ip=ip) + return _unary_math_op(x, math_dialect.sqrt, None, fastmath, "sqrt", loc=loc, ip=ip) + + +# ============================================================================= +# MLIR-only transcendentals (no NVVM approx intrinsic) +# ============================================================================= + + +@dsl_user_op +def erf( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise error function of the input operand. + + The error function is defined as: + erf(x) = 2/√π ∫[0 to x] exp(-t²) dt + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the error function value for each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = erf(a) # Compute error function + """ + return _unary_math_op(x, math_dialect.erf, None, fastmath, "erf", loc=loc, ip=ip) + + +@dsl_user_op +def tan( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise tangent of the input operand. + + :param a: Input operand (in radians) + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the tangent of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = tan(a) # Compute tangent + """ + return _unary_math_op(x, math_dialect.tan, None, fastmath, "tan", loc=loc, ip=ip) + + +@dsl_user_op +def acos( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise arc cosine of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the arc cosine of each element in input operand + :rtype: MathOperand + + Example: + + .. code-block:: + + y = acos(a) # Compute arc cosine + """ + return _unary_math_op(x, math_dialect.acos, None, fastmath, "acos", loc=loc, ip=ip) + + +@dsl_user_op +def asin( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise arc sine of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the arc sine of each element in input operand + :rtype: MathOperand + + Example: + + .. code-block:: + + y = asin(a) # Compute arc sine + """ + return _unary_math_op(x, math_dialect.asin, None, fastmath, "asin", loc=loc, ip=ip) + + +@dsl_user_op +def atan( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise arc tangent of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the arc tangent of each element in input operand + :rtype: MathOperand + + Example: + + .. code-block:: + + y = atan(a) # Compute arc tangent + """ + return _unary_math_op(x, math_dialect.atan, None, fastmath, "atan", loc=loc, ip=ip) + + +@dsl_user_op +def atan2( + y: MathOperand, + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise arc tangent of two tensors. + + Computes atan2(a, b) element-wise. The function atan2(a, b) is the angle in radians + between the positive x-axis and the point given by the coordinates (b, a). + + :param a: First input operand (y-coordinates) + :type a: MathOperand + :param b: Second input operand (x-coordinates) + :type b: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the arc tangent of a/b element-wise + :rtype: MathOperand + + Example: + + .. code-block:: + + theta = atan2(y, x) # Compute angles + """ + _check_vector_consistency("atan2", y, x) + _validate_same_numeric_type("atan2", y, x) + y_ir = _get_ir_value(y, ip=ip) + x_ir = _get_ir_value(x, ip=ip) + fastmath_flag = _get_fastmath_flag(fastmath) + result = math_dialect.atan2(y_ir, x_ir, fastmath=fastmath_flag, loc=loc, ip=ip) + return _wrap_result(y, result) + + +@dsl_user_op +def sinh( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise hyperbolic sine of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the hyperbolic sine of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = sinh(a) # Compute hyperbolic sine + """ + return _unary_math_op(x, math_dialect.sinh, None, fastmath, "sinh", loc=loc, ip=ip) + + +@dsl_user_op +def cosh( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise hyperbolic cosine of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the hyperbolic cosine of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = cosh(a) # Compute hyperbolic cosine + """ + return _unary_math_op(x, math_dialect.cosh, None, fastmath, "cosh", loc=loc, ip=ip) + + +@dsl_user_op +def pow( + base: MathOperand, + exponent: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise power of the input tensors. + + :param a: Input operand (base) + :type a: MathOperand + :param b: Input operand (exponent) + :type b: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing a raised to the power b for each element + :rtype: MathOperand + + Example: + + .. code-block:: + + z = pow(a, b) # Compute a^b + """ + _check_vector_consistency("pow", base, exponent) + _validate_same_numeric_type("pow", base, exponent) + base_ir = _get_ir_value(base, ip=ip) + exp_ir = _get_ir_value(exponent, ip=ip) + fastmath_flag = _get_fastmath_flag(fastmath) + result = math_dialect.powf(base_ir, exp_ir, fastmath=fastmath_flag, loc=loc, ip=ip) + return _wrap_result(base, result) + + +@dsl_user_op +def cbrt( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise cube root of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the cube root of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = cbrt(a) # Compute cube root + """ + return _unary_math_op(x, math_dialect.cbrt, None, fastmath, "cbrt", loc=loc, ip=ip) + + +@dsl_user_op +def expm1( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise exp(x) - 1 of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing exp(x) - 1 of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = expm1(a) # Compute exp(x) - 1 + """ + return _unary_math_op( + x, math_dialect.expm1, None, fastmath, "expm1", loc=loc, ip=ip + ) + + +@dsl_user_op +def log1p( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise log(1 + x) of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing log(1 + x) of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = log1p(a) # Compute log(1 + x) + """ + return _unary_math_op( + x, math_dialect.log1p, None, fastmath, "log1p", loc=loc, ip=ip + ) + + +# ============================================================================= +# Absolute value, sign, rounding +# ============================================================================= + + +@dsl_user_op +def abs( + x: MathOperand, + fastmath: bool = False, + ftz: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Absolute value (float or integer). + + :param ftz: Flush denormals to zero (float only, uses nvvm.fabs) + """ + _validate_fastmath_exclusive("abs", fastmath, ftz=ftz) + if ftz and _is_float(x): + x_ir = _get_ir_value(x, ip=ip) + result = nvvm.fabs(x_ir, ftz=True, loc=loc, ip=ip) + return _wrap_result(x, result) + return _unary_math_op( + x, + math_dialect.absf, + math_dialect.absi, + fastmath, + "abs", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def copysign( + mag: MathOperand, + sign: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise copysign, combining magnitude of a with sign of b. + + :param a: Input operand (magnitude source) + :type a: MathOperand + :param b: Input operand (sign source) + :type b: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the copysign of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + z = copysign(a, b) # Copy sign of b to magnitude of a + """ + _check_vector_consistency("copysign", mag, sign) + _validate_same_numeric_type("copysign", mag, sign) + mag_ir = _get_ir_value(mag, ip=ip) + sign_ir = _get_ir_value(sign, ip=ip) + fastmath_flag = _get_fastmath_flag(fastmath) + result = math_dialect.copysign( + mag_ir, sign_ir, fastmath=fastmath_flag, loc=loc, ip=ip + ) + return _wrap_result(mag, result) + + +@dsl_user_op +def ceil( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise ceiling of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the ceiling of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = ceil(a) # Compute ceiling + """ + return _unary_math_op(x, math_dialect.ceil, None, fastmath, "ceil", loc=loc, ip=ip) + + +@dsl_user_op +def floor( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise floor of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the largest integer less than or equal to each element in input operand + :rtype: MathOperand + + Example: + + .. code-block:: + + y = floor(a) # Compute floor + """ + return _unary_math_op( + x, math_dialect.floor, None, fastmath, "floor", loc=loc, ip=ip + ) + + +@dsl_user_op +def round( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise round to nearest integer (ties away from zero) of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the rounded value of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = round(a) # Round to nearest integer + """ + return _unary_math_op( + x, math_dialect.round, None, fastmath, "round", loc=loc, ip=ip + ) + + +@dsl_user_op +def roundeven( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise round to nearest integer (ties to even) of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the rounded value of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = roundeven(a) # Round to nearest integer (ties to even) + """ + return _unary_math_op( + x, math_dialect.roundeven, None, fastmath, "roundeven", loc=loc, ip=ip + ) + + +@dsl_user_op +def trunc( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise truncation toward zero of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the truncated value of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = trunc(a) # Truncate toward zero + """ + return _unary_math_op( + x, math_dialect.trunc, None, fastmath, "trunc", loc=loc, ip=ip + ) + + +# ============================================================================= +# Min / Max +# ============================================================================= + + +@dsl_user_op +def min( + x: MathOperand, + y: MathOperand, + propagate_nan: bool = False, + ftz: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Element-wise minimum. + + :param propagate_nan: If True, NaN propagates (IEEE 754 minimum). + If False, NaN is ignored (IEEE 754 minimumNumber). + :param ftz: Flush denormals to zero (float only, uses nvvm.fmin). + """ + _check_vector_consistency("min", x, y) + _validate_same_numeric_type("min", x, y) + x_ir = _get_ir_value(x, ip=ip) + y_ir = _get_ir_value(y, ip=ip) + if _is_float(x): + if ftz: + result = nvvm.fmin(x_ir, y_ir, ftz=True, nan=propagate_nan, loc=loc, ip=ip) + elif propagate_nan: + result = arith.minimumf(x_ir, y_ir, loc=loc, ip=ip) + else: + result = arith.minnumf(x_ir, y_ir, loc=loc, ip=ip) + elif _is_unsigned_int(x): + result = arith.minui(x_ir, y_ir, loc=loc, ip=ip) + else: + result = arith.minsi(x_ir, y_ir, loc=loc, ip=ip) + return _wrap_result(x, result) + + +@dsl_user_op +def max( + x: MathOperand, + y: MathOperand, + propagate_nan: bool = False, + ftz: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Element-wise maximum. + + :param propagate_nan: If True, NaN propagates (IEEE 754 maximum). + If False, NaN is ignored (IEEE 754 maximumNumber). + :param ftz: Flush denormals to zero (float only, uses nvvm.fmax). + """ + _check_vector_consistency("max", x, y) + _validate_same_numeric_type("max", x, y) + x_ir = _get_ir_value(x, ip=ip) + y_ir = _get_ir_value(y, ip=ip) + if _is_float(x): + if ftz: + result = nvvm.fmax(x_ir, y_ir, ftz=True, nan=propagate_nan, loc=loc, ip=ip) + elif propagate_nan: + result = arith.maximumf(x_ir, y_ir, loc=loc, ip=ip) + else: + result = arith.maxnumf(x_ir, y_ir, loc=loc, ip=ip) + elif _is_unsigned_int(x): + result = arith.maxui(x_ir, y_ir, loc=loc, ip=ip) + else: + result = arith.maxsi(x_ir, y_ir, loc=loc, ip=ip) + return _wrap_result(x, result) + + +# ============================================================================= +# FMA +# ============================================================================= + + +@dsl_user_op +def fma( + a: MathOperand, + b: MathOperand, + c: MathOperand, + fastmath: bool = False, + rounding: Optional[RoundingMode] = None, + ftz: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Fused multiply-add: a * b + c.""" + _check_vector_consistency("fma", a, b, c) + _validate_same_numeric_type("fma", a, b, c) + _validate_fastmath_exclusive("fma", fastmath, ftz=ftz, rounding=rounding) + if _needs_nvvm_intrinsic(rounding, ftz): + return _call_nvvm_ternary(a, b, c, "fma", rounding, ftz, loc=loc, ip=ip) + a_ir = _get_ir_value(a, ip=ip) + b_ir = _get_ir_value(b, ip=ip) + c_ir = _get_ir_value(c, ip=ip) + fastmath_flag = _get_fastmath_flag(fastmath) + result = math_dialect.fma(a_ir, b_ir, c_ir, fastmath=fastmath_flag, loc=loc, ip=ip) + return _wrap_result(a, result) + + +# ============================================================================= +# Arithmetic with fastmath / rounding / ftz +# ============================================================================= + + +@dsl_user_op +def add( + x: MathOperand, + y: MathOperand, + fastmath: bool = False, + rounding: Optional[RoundingMode] = None, + ftz: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Floating-point addition.""" + _check_vector_consistency("add", x, y) + _validate_same_numeric_type("add", x, y) + _validate_fastmath_exclusive("add", fastmath, ftz=ftz, rounding=rounding) + if _needs_nvvm_intrinsic(rounding, ftz): + return _call_nvvm_binary(x, y, "add", rounding, ftz, loc=loc, ip=ip) + x_ir = _get_ir_value(x, ip=ip) + y_ir = _get_ir_value(y, ip=ip) + fastmath_flag = _get_fastmath_flag(fastmath) + result = arith.addf(x_ir, y_ir, fastmath=fastmath_flag, loc=loc, ip=ip) + return _wrap_result(x, result) + + +@dsl_user_op +def sub( + x: MathOperand, + y: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Floating-point subtraction. + + Note: NVVM sub uses a different intrinsic convention (rounding as an + integer arg), so rounding/ftz are not supported here. Callers that + need explicit rounding control on subtraction should emit the + PTX-level ``sub.[.ftz].f32`` intrinsic directly. + """ + _check_vector_consistency("sub", x, y) + _validate_same_numeric_type("sub", x, y) + x_ir = _get_ir_value(x, ip=ip) + y_ir = _get_ir_value(y, ip=ip) + fastmath_flag = _get_fastmath_flag(fastmath) + result = arith.subf(x_ir, y_ir, fastmath=fastmath_flag, loc=loc, ip=ip) + return _wrap_result(x, result) + + +@dsl_user_op +def mul( + x: MathOperand, + y: MathOperand, + fastmath: bool = False, + rounding: Optional[RoundingMode] = None, + ftz: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Floating-point multiplication.""" + _check_vector_consistency("mul", x, y) + _validate_same_numeric_type("mul", x, y) + _validate_fastmath_exclusive("mul", fastmath, ftz=ftz, rounding=rounding) + if _needs_nvvm_intrinsic(rounding, ftz): + return _call_nvvm_binary(x, y, "mul", rounding, ftz, loc=loc, ip=ip) + x_ir = _get_ir_value(x, ip=ip) + y_ir = _get_ir_value(y, ip=ip) + fastmath_flag = _get_fastmath_flag(fastmath) + result = arith.mulf(x_ir, y_ir, fastmath=fastmath_flag, loc=loc, ip=ip) + return _wrap_result(x, result) + + +@dsl_user_op +def div( + x: MathOperand, + y: MathOperand, + fastmath: bool = False, + approx: bool = False, + ftz: bool = False, + full: bool = False, + rounding: Optional[RoundingMode] = None, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Floating-point division.""" + _check_vector_consistency("div", x, y) + _validate_same_numeric_type("div", x, y) + _validate_fastmath_exclusive( + "div", fastmath, approx=approx, ftz=ftz, rounding=rounding, full=full + ) + if fastmath: + x_ir = _get_ir_value(x, ip=ip) + y_ir = _get_ir_value(y, ip=ip) + fastmath_flag = _get_fastmath_flag(True) + result = arith.divf(x_ir, y_ir, fastmath=fastmath_flag, loc=loc, ip=ip) + return _wrap_result(x, result) + if approx: + return _call_nvvm_binary(x, y, "div", approx=True, ftz=ftz, loc=loc, ip=ip) + if full: + return _call_nvvm_binary(x, y, "div", full=True, ftz=ftz, loc=loc, ip=ip) + if _needs_nvvm_intrinsic(rounding, ftz): + return _call_nvvm_binary(x, y, "div", rounding, ftz, loc=loc, ip=ip) + x_ir = _get_ir_value(x, ip=ip) + y_ir = _get_ir_value(y, ip=ip) + result = arith.divf(x_ir, y_ir, loc=loc, ip=ip) + return _wrap_result(x, result) + + +@dsl_user_op +def rcp( + x: MathOperand, + fastmath: bool = False, + approx: bool = False, + rounding: Optional[RoundingMode] = None, + ftz: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Reciprocal (1/x).""" + _validate_fastmath_exclusive( + "rcp", fastmath, approx=approx, ftz=ftz, rounding=rounding + ) + if fastmath: + x_ir = _get_ir_value(x, ip=ip) + one = _const_one(x_ir.type, loc=loc, ip=ip) + fastmath_flag = _get_fastmath_flag(True) + result = arith.divf(one, x_ir, fastmath=fastmath_flag, loc=loc, ip=ip) + return _wrap_result(x, result) + if approx: + if isinstance(x, Vector): + x_ir = _get_ir_value(x, ip=ip) + result = nvgpu.rcp( + x_ir, + rounding=nvgpu.RcpRoundingMode.APPROX, + ftz=ftz or None, + loc=loc, + ip=ip, + ) + return _wrap_result(x, result) + # Scalar approx+ftz on f32 has a first-class NVVM dialect op. + # Prefer it over llvm.call_intrinsic for better pattern-match/ + # canonicalization visibility. + if ftz and _is_f32(x): + x_ir = _get_ir_value(x, ip=ip) + result = nvvm.rcp_approx_ftz_f(x_ir, loc=loc, ip=ip) + return _wrap_result(x, result) + return _call_nvvm_unary(x, "rcp", approx=True, ftz=ftz, loc=loc, ip=ip) + if _needs_nvvm_intrinsic(rounding, ftz): + return _call_nvvm_unary(x, "rcp", rounding=rounding, ftz=ftz, loc=loc, ip=ip) + # Default semantics: IEEE round-to-nearest reciprocal. + # Scalar f32/f64: emit nvvm.rcp.rn.{f32,f64} (one PTX rcp instruction). + # Vector or scalar element type without an NVVM scalar-rcp intrinsic + # (e.g., f16): lower to arith.divf(1, x) — same IEEE semantics, + # vectorizable by the compiler. + if isinstance(x, Vector) or not (_is_f32(x) or _is_f64(x)): + x_ir = _get_ir_value(x, ip=ip) + one = _const_one(x_ir.type, loc=loc, ip=ip) + result = arith.divf(one, x_ir, loc=loc, ip=ip) + return _wrap_result(x, result) + return _call_nvvm_unary( + x, "rcp", rounding=RoundingMode.NEAREST_EVEN, loc=loc, ip=ip + ) + + +# ============================================================================= +# Inverse hyperbolic +# ============================================================================= + + +@dsl_user_op +def acosh( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise inverse hyperbolic cosine of the input operand. + + :param a: Input operand (must be >= 1.0) + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the inverse hyperbolic cosine of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = acosh(a) # Compute inverse hyperbolic cosine + """ + return _unary_math_op( + x, math_dialect.acosh, None, fastmath, "acosh", loc=loc, ip=ip + ) + + +@dsl_user_op +def asinh( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise inverse hyperbolic sine of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the inverse hyperbolic sine of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = asinh(a) # Compute inverse hyperbolic sine + """ + return _unary_math_op( + x, math_dialect.asinh, None, fastmath, "asinh", loc=loc, ip=ip + ) + + +@dsl_user_op +def atanh( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise inverse hyperbolic tangent of the input operand. + + :param a: Input operand (must be in (-1, 1)) + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the inverse hyperbolic tangent of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = atanh(a) # Compute inverse hyperbolic tangent + """ + return _unary_math_op( + x, math_dialect.atanh, None, fastmath, "atanh", loc=loc, ip=ip + ) + + +# ============================================================================= +# Complementary error function +# ============================================================================= + + +@dsl_user_op +def erfc( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Compute element-wise complementary error function of the input operand. + + :param a: Input operand + :type a: MathOperand + :param fastmath: Enable fast math optimizations, defaults to False + :type fastmath: bool, optional + :param loc: Source location information, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for IR generation, defaults to None + :type ip: Optional[InsertionPoint] + :return: Result containing the complementary error function of each element + :rtype: MathOperand + + Example: + + .. code-block:: + + y = erfc(a) # Compute complementary error function + """ + return _unary_math_op(x, math_dialect.erfc, None, fastmath, "erfc", loc=loc, ip=ip) + + +# ============================================================================= +# Combined sin+cos +# ============================================================================= + + +@dsl_user_op +def sincos( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> tuple: + """Combined sine and cosine: returns (sin(x), cos(x)). + + More efficient than separate sin() and cos() calls. + """ + x_ir = _get_ir_value(x, ip=ip) + fastmath_flag = _get_fastmath_flag(fastmath) + op = math_dialect.SincosOp(x_ir, fastmath=fastmath_flag, loc=loc, ip=ip) + return (_wrap_result(x, op.sin), _wrap_result(x, op.cos)) + + +# ============================================================================= +# Power variants +# ============================================================================= + + +@dsl_user_op +def fpowi( + base: MathOperand, + exponent: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Float raised to integer power: base^exponent (exponent is integer).""" + base_ir = _get_ir_value(base, ip=ip) + exp_ir = _get_ir_value(exponent, ip=ip) + fastmath_flag = _get_fastmath_flag(fastmath) + result = math_dialect.fpowi(base_ir, exp_ir, fastmath=fastmath_flag, loc=loc, ip=ip) + return _wrap_result(base, result) + + +@dsl_user_op +def ipowi( + base: MathOperand, + exponent: MathOperand, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Integer raised to integer power: base^exponent.""" + base_ir = _get_ir_value(base, ip=ip) + exp_ir = _get_ir_value(exponent, ip=ip) + result = math_dialect.ipowi(base_ir, exp_ir, loc=loc, ip=ip) + return _wrap_result(base, result) + + +# ============================================================================= +# Clamp (native math dialect op) +# ============================================================================= + + +@dsl_user_op +def clamp( + x: MathOperand, + lo: MathOperand, + hi: MathOperand, + fastmath: bool = False, + propagate_nan: bool = False, + ftz: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Clamp value to range [lo, hi]. + + Lowered unconditionally as ``max(lo, min(x, hi))``. The math dialect's + ``clampf`` op currently has no LLVM translation interface registered, so + a direct lowering fails at JIT time on scalar f16 inputs (and any other + type). Composition via min/max picks up the right per-type lowering + (``arith.minnumf`` / ``arith.maximumf`` / ``nvvm.fmin`` / ``nvvm.fmax``). + """ + _validate_fastmath_exclusive("clamp", fastmath, ftz=ftz) + inner = min(x, hi, propagate_nan=propagate_nan, ftz=ftz, loc=loc, ip=ip) + return max(lo, inner, propagate_nan=propagate_nan, ftz=ftz, loc=loc, ip=ip) + + +# ============================================================================= +# FP classification predicates +# ============================================================================= + + +@dsl_user_op +def isnan( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Test if value is NaN. Returns i1.""" + x_ir = _get_ir_value(x, ip=ip) + fastmath_flag = _get_fastmath_flag(fastmath) + return math_dialect.isnan(x_ir, fastmath=fastmath_flag, loc=loc, ip=ip) + + +@dsl_user_op +def isinf( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Test if value is infinite. Returns i1.""" + x_ir = _get_ir_value(x, ip=ip) + fastmath_flag = _get_fastmath_flag(fastmath) + return math_dialect.isinf(x_ir, fastmath=fastmath_flag, loc=loc, ip=ip) + + +@dsl_user_op +def isfinite( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Test if value is finite. Returns i1.""" + x_ir = _get_ir_value(x, ip=ip) + fastmath_flag = _get_fastmath_flag(fastmath) + return math_dialect.isfinite(x_ir, fastmath=fastmath_flag, loc=loc, ip=ip) + + +@dsl_user_op +def isnormal( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Test if value is normal (not zero, subnormal, inf, or NaN). Returns i1.""" + x_ir = _get_ir_value(x, ip=ip) + fastmath_flag = _get_fastmath_flag(fastmath) + return math_dialect.isnormal(x_ir, fastmath=fastmath_flag, loc=loc, ip=ip) + + +# ============================================================================= +# Integer math +# ============================================================================= + + +@dsl_user_op +def absi( + x: MathOperand, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Integer absolute value.""" + x_ir = _get_ir_value(x, ip=ip) + result = math_dialect.absi(x_ir, loc=loc, ip=ip) + return _wrap_result(x, result) + + +# ============================================================================= +# Negation and remainder (arith dialect) +# ============================================================================= + + +@dsl_user_op +def neg( + x: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Floating-point negation.""" + x_ir = _get_ir_value(x, ip=ip) + fastmath_flag = _get_fastmath_flag(fastmath) + result = arith.negf(x_ir, fastmath=fastmath_flag, loc=loc, ip=ip) + return _wrap_result(x, result) + + +@dsl_user_op +def rem( + x: MathOperand, + y: MathOperand, + fastmath: bool = False, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> MathOperand: + """Floating-point remainder.""" + x_ir = _get_ir_value(x, ip=ip) + y_ir = _get_ir_value(y, ip=ip) + fastmath_flag = _get_fastmath_flag(fastmath) + result = arith.remf(x_ir, y_ir, fastmath=fastmath_flag, loc=loc, ip=ip) + return _wrap_result(x, result) + + +# ============================================================================= +# __all__ +# ============================================================================= + +__all__ = [ + "RoundingMode", + "MathOperand", + # Transcendentals with NVVM approx (hardware instructions) + "sin", + "cos", + "exp2", + "log2", + "tanh", + "rsqrt", + "sqrt", + # Transcendentals without hardware approx + "exp", + "log", + "log10", + # Trigonometric + "tan", + "acos", + "asin", + "atan", + "atan2", + # Hyperbolic + "sinh", + "cosh", + "acosh", + "asinh", + "atanh", + # Error functions + "erf", + "erfc", + # Power / root / exponential + "pow", + "fpowi", + "ipowi", + "cbrt", + "expm1", + "log1p", + # Combined + "sincos", + # Absolute value, sign, rounding + "abs", + "absi", + "copysign", + "neg", + "ceil", + "floor", + "round", + "roundeven", + "trunc", + # Clamp + "clamp", + # Min / Max + "min", + "max", + # FP predicates + "isnan", + "isinf", + "isfinite", + "isnormal", + # FMA + "fma", + # Arithmetic + "add", + "sub", + "mul", + "div", + "rem", + "rcp", +] diff --git a/python/CuTeDSL/_mlir_helpers/nanobind_compat.py b/python/CuTeDSL/_mlir_helpers/nanobind_compat.py new file mode 100644 index 000000000..860c9ec4c --- /dev/null +++ b/python/CuTeDSL/_mlir_helpers/nanobind_compat.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +""" +Compatibility patches for nanobind-based MLIR Python bindings. + +The LLVM prebuilt is compiled with -fno-exceptions. Several nanobind +binding methods signal errors or end-of-iteration by throwing C++ +exceptions (nb::stop_iteration, nb::value_error, ...). With +-fno-exceptions the throw cannot unwind and calls std::terminate/abort. +The patches below replace those code paths with alternatives +that never reach the throwing C++ code. +""" + +from threading import local +from typing import Any, Iterator + +from .._mlir import ir + +_thread_local = local() + + +# -- InsertionPoint.current -------------------------------------------------- +# ir.InsertionPoint.current throws nb::value_error when no insertion +# point is active. Track enter/exit at the Python level so callers can +# check before touching the C++ property. + + +def has_insertion_point() -> bool: + return getattr(_thread_local, "ip_depth", 0) > 0 + + +_orig_ip_enter = ir.InsertionPoint.__enter__ +_orig_ip_exit = ir.InsertionPoint.__exit__ + + +def _tracked_ip_enter(self: ir.InsertionPoint) -> ir.InsertionPoint: + _thread_local.ip_depth = getattr(_thread_local, "ip_depth", 0) + 1 + return _orig_ip_enter(self) + + +def _tracked_ip_exit(self: ir.InsertionPoint, *args: Any) -> None: + _thread_local.ip_depth = getattr(_thread_local, "ip_depth", 1) - 1 + _orig_ip_exit(self, *args) + + +ir.InsertionPoint.__enter__ = _tracked_ip_enter +ir.InsertionPoint.__exit__ = _tracked_ip_exit + + +# -- Block.__iter__ ----------------------------------------------------------- +# PyOperationIterator.__next__ throws nb::stop_iteration() at +# end-of-block. Replace with indexed access via Block.operations +# whose __len__ and __getitem__ use the C API without throwing. + + +def _safe_block_iter(self: ir.Block) -> Iterator[Any]: + ops = self.operations + for i in range(len(ops)): + yield ops[i] + + +ir.Block.__iter__ = _safe_block_iter diff --git a/python/CuTeDSL/_mlir_helpers/op.py b/python/CuTeDSL/_mlir_helpers/op.py new file mode 100644 index 000000000..8ed06b069 --- /dev/null +++ b/python/CuTeDSL/_mlir_helpers/op.py @@ -0,0 +1,389 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +""" +This module provides MLIR's OP helper functions +""" + +import inspect +import importlib.util +import os +import types +import warnings +from functools import wraps, lru_cache +from typing import Any, Callable, TYPE_CHECKING + +from .._mlir import ir +from ..base_dsl.common import DSLRuntimeError, DSLOperationBuildError +from ..base_dsl.utils.logger import log +from ..base_dsl.utils.stacktrace import walk_to_top_module + +# The DSL package root is empty by default. +_DSL_PACKAGE_ROOT: str | None = "" +_DSL_PACKAGE_ROOTS: tuple[str, ...] | None = None + +# Whether location tracking is enabled. +_ENABLE_FRAME_FILTERING: bool = False + +# When True, dsl_user_op attributes ops to the closest frame (including DSL +# library code) instead of skipping up to the user's call site. Enabled when +# debugging mode is ON so library developers can see where inside the DSL an op +# was built. +_INCLUDE_LIB_FRAME: bool = False + +_SOURCE_FRAMEWORK_DIRS = ( + "_mlir_helpers", + "base_dsl", + "cutlass_dsl", +) +_FRAMEWORK_MODULE_PREFIXES = ( + "cutlass", +) + +# Trace-time MLIR op verification — see `dsl_user_op` docstring. +# Default resolved lazily from CUTE_DSL_VERIFY_TRACE on first read to +# avoid a circular import with base_dsl.env_manager (which transitively +# imports base_dsl.runtime, _mlir_helpers, and back into this module). +_VERIFY_TRACE_ENABLED: bool | None = None + + +def _verify_trace_enabled() -> bool: + global _VERIFY_TRACE_ENABLED + if _VERIFY_TRACE_ENABLED is None: + from ..base_dsl.env_manager import get_bool_env_var + + _VERIFY_TRACE_ENABLED = get_bool_env_var("CUTE_DSL_VERIFY_TRACE", False) + return _VERIFY_TRACE_ENABLED + + +def set_verify_trace(enabled: bool) -> None: + """Toggle trace-time MLIR op verification at runtime. + + Useful for hot benchmarking loops where the per-op `verify()` + cost matters. + """ + global _VERIFY_TRACE_ENABLED + _VERIFY_TRACE_ENABLED = bool(enabled) + + +def get_verify_trace() -> bool: + """Return whether trace-time verification is currently enabled.""" + return _verify_trace_enabled() + + +def _verify_new_block_ops(snap_block: Any, snap_n_ops: int) -> None: + """Verify ops appended to ``snap_block`` after ``snap_n_ops``.""" + ops = snap_block.operations + new_count = len(ops) - snap_n_ops + if new_count <= 0: + return + + # OperationList materialization is O(block size). Most DSL wrappers append + # one op, so walk the new tail by negative index instead of list-slicing the + # whole block on every wrapper call. + for offset in range(new_count, 0, -1): + ops[-offset].verify() + + +def _set_enable_frame_filtering(enable: bool) -> None: + """Set whether location tracking is enabled.""" + global _ENABLE_FRAME_FILTERING + _ENABLE_FRAME_FILTERING = enable + + +def _set_include_lib_frame(enable: bool) -> None: + """Set whether ops are attributed to the closest (library) frame. + + When enabled, `_find_user_frame` returns the immediate caller of the op + wrapper instead of walking up past DSL library frames. Used when debugging + mode is ON for internal DSL development. + """ + global _INCLUDE_LIB_FRAME + _INCLUDE_LIB_FRAME = enable + + +@lru_cache(maxsize=256) +def _is_framework_frame(filename: str) -> bool: + """Check if a frame's filename belongs to DSL library code.""" + global _DSL_PACKAGE_ROOT, _DSL_PACKAGE_ROOTS + if _DSL_PACKAGE_ROOTS is None: + roots: list[str] = [] + + module_root = walk_to_top_module(os.path.dirname(os.path.abspath(__file__))) + + _DSL_PACKAGE_ROOT = module_root + if _DSL_PACKAGE_ROOT is not None: + roots.append(os.path.abspath(_DSL_PACKAGE_ROOT)) + + for package_name in _FRAMEWORK_MODULE_PREFIXES: + spec = importlib.util.find_spec(package_name) + if spec is None: + continue + if spec.submodule_search_locations: + roots.extend( + os.path.abspath(path) for path in spec.submodule_search_locations + ) + elif spec.origin: + roots.append(os.path.dirname(os.path.abspath(spec.origin))) + + _DSL_PACKAGE_ROOTS = tuple(dict.fromkeys(roots)) + + if not _DSL_PACKAGE_ROOTS: + return False + + filename = os.path.abspath(filename) + return any( + filename == root or filename.startswith(root + os.sep) + for root in _DSL_PACKAGE_ROOTS + ) + + +def _is_framework_module(module_name: str) -> bool: + if not module_name: + return False + framework_prefixes = _FRAMEWORK_MODULE_PREFIXES + + return any( + module_name == prefix or module_name.startswith(prefix + ".") + for prefix in framework_prefixes + ) + + +def _is_framework_stack_frame(frame: types.FrameType) -> bool: + module_name = frame.f_globals.get("__name__", "") + return _is_framework_module(module_name) or _is_framework_frame( + frame.f_code.co_filename + ) + + +def _find_user_frame(start_frame: types.FrameType | None) -> types.FrameType | None: + """Walk up the call stack from start_frame to find the first user (non-library) frame. + + Returns the first frame whose file is not under the DSL package root. + Falls back to start_frame if no user frame is found. + + Uses linear search with early termination - returns immediately upon finding the + first user frame. This is O(m) where m is the position of the first user frame + (typically 3-10). This is optimal - no algorithm can do better for finding the + first occurrence in an unsorted sequence. + + The main optimization is @lru_cache on _is_framework_frame(), which caches filename + checks across all DSL operations. + """ + if _INCLUDE_LIB_FRAME: + return start_frame + if not _ENABLE_FRAME_FILTERING: + return start_frame + + frame = start_frame + while frame is not None: + if not _is_framework_stack_frame(frame): + return frame + frame = frame.f_back + # Fallback: if everything is framework code, use the original caller + return start_frame + + +def dsl_user_op(opFunc: Callable[..., Any]) -> Callable[..., Any]: + """Decorator for user-facing DSL op wrappers. + + Responsibilities: + + 1. Attach a source-location to every MLIR op built by ``opFunc`` so that + diagnostics and IR dumps point back at the user's Python call site. + 2. Run trace-time MLIR verification on each newly-built op so verifier + errors surface at the call site rather than at module-verify time. + + Verification snapshots ``InsertionPoint.current.block`` and its op + count before invoking ``opFunc``, then calls ``OpView.verify()`` on + every op appended during the call. ``verify()`` recurses through + regions, so a wrapper that opens an ``scf.if`` covers the body via + the outer op. + + Disabled by default; enable with the ``CUTE_DSL_VERIFY_TRACE`` env + var (``1`` / ``true`` / ``on`` / ``yes``) or ``set_verify_trace(True)``. + + :param opFunc: The user-facing API function. + :type opFunc: Callable + :return: The wrapped user-facing API function. + :rtype: Callable + """ + + @wraps(opFunc) + def wrapper(*args: Any, **kwargs: Any) -> Any: + # Pop loc= from kwargs so callers that still pass it don't break. + # We no longer forward it — LOC_TRACEBACKS captures full stacks automatically. + loc: Any = kwargs.pop("loc", None) + frameInfo = None + verifier_error = False + + if loc is None and ir.Context.current is not None: + cur_frame = inspect.currentframe() + assert cur_frame is not None + frame = _find_user_frame(cur_frame.f_back) + del cur_frame # break self-ref + assert frame is not None + frameInfo = inspect.getframeinfo(frame) + try: + # In Python < 3.11, getframeinfo returns a NamedTuple without positions + if not hasattr(frameInfo, "positions"): + file_loc = ir.Location.file( + frameInfo.filename, + frameInfo.lineno, + 0, + ) + else: + file_loc = ir.Location.file( + frameInfo.filename, + frameInfo.positions.lineno, # type: ignore[attr-defined] + frameInfo.positions.col_offset or 0, # type: ignore[attr-defined] + ) + loc = ir.Location.name( + ( + "".join([c.strip() for c in frameInfo.code_context]) + if frameInfo.code_context + else frameInfo.function + ), + childLoc=file_loc, + ) + except RuntimeError: + # No MLIR context available (e.g. validation-only call + # outside a kernel). Proceed with loc=None so that the + # wrapped function's own validation can still fire. + pass + + is_init = False + # Snapshot the current insertion block so we can verify newly-built + # ops after opFunc returns. The dialect Python bindings strip + # `.result` for value-producing ops before our wrapper sees the + # return value, so we cannot rely on the return value alone — + # block-diff sees every op regardless. + snap_block: Any = None + snap_n_ops: int = 0 + if _verify_trace_enabled() and not is_init: + try: + snap_block = ir.InsertionPoint.current.block + snap_n_ops = len(snap_block.operations) + except ValueError: + # No active InsertionPoint (e.g. validation-only call + # outside a kernel) — skip trace-time verification. + snap_block = None + + try: + res_or_list = opFunc(*args, **kwargs, loc=loc) + verifier_error = True + # Trace-time verification: walk every op appended to the + # snapshotted block during opFunc. OpView.verify() is + # recursive — for ops with regions (scf.if, scf.for, ...) it + # verifies the body too, so we don't need to descend. + # + # Skip when opFunc returns a context manager (e.g. elect_one, + # if_generate): the op's body is filled by the surrounding + # `with` block after opFunc returns, so verifying here would + # see an empty region and fail. The body still gets covered + # at module-verify time. + is_cm = hasattr(res_or_list, "__enter__") and hasattr( + res_or_list, "__exit__" + ) + # Fast path: if opFunc didn't append anything to the block, + # skip materializing the operation list — it's O(N) in the + # block size and adds up across thousands of wrapper calls. + if ( + snap_block is not None + and not is_cm + and len(snap_block.operations) > snap_n_ops + ): + _verify_new_block_ops(snap_block, snap_n_ops) + # Defensive fallback: cover the rare wrapper that returns an + # OpView directly. Block-diff already verified it if it was + # appended to snap_block, so this is a no-op in the common + # case; the branch only matters for cross-block builders or + # calls with no active insertion point. + if hasattr(res_or_list, "verify"): + res_or_list.verify() + + except DSLOperationBuildError as e: + # Nested DSLOperationError + raise DSLOperationBuildError( + message=e.message, cause=e, frameInfo=frameInfo + ) + except Exception as e: + # Check if it's a decorator config error first + func_name = getattr(opFunc, "__name__", str(opFunc)) + if "unexpected keyword argument 'loc'" in str(e): + raise DSLRuntimeError( + f"Function '{func_name}' decorated with @dsl_user_op does not accept the required 'loc' parameter.", + suggestion=[ + f"1. Add 'loc=None' as a keyword-only parameter to {func_name}:", + f" def {func_name}(..., *, loc=None):", + "", + "2. Remove the @dsl_user_op decorator if location tracking is not needed", + ], + cause=e, + ) from e + if verifier_error: + raise DSLOperationBuildError( + message="Operation verification failed", + cause=e, + frameInfo=frameInfo, + auto_translate=False, + ) + + raise e + + return res_or_list + + return wrapper + + +def _get_call_arg_names(frameInfo: Any) -> dict[str | int, str]: + """Extract argument names from the source line (best-effort).""" + if frameInfo is None or not frameInfo.code_context: + return {} + import ast as _ast + + line = "".join(frameInfo.code_context).strip() + call_node: _ast.expr | None = None + try: + expr_tree = _ast.parse(line, mode="eval") + call_node = expr_tree.body + except SyntaxError: + try: + mod_tree = _ast.parse(line, mode="exec") + if not mod_tree.body: + return {} + stmt = mod_tree.body[0] + if isinstance(stmt, _ast.Expr): + call_node = stmt.value + elif isinstance(stmt, _ast.Assign): + call_node = stmt.value + elif isinstance(stmt, _ast.Return): + call_node = stmt.value + elif isinstance(stmt, _ast.AugAssign): + call_node = stmt.value + else: + return {} + except SyntaxError: + return {} + + # Unwrap chained calls: ``type(self)(self, signed)`` parses as a Call + # whose func is another Call. We want the outermost (rightmost) call. + if not isinstance(call_node, _ast.Call): + return {} + names: dict[str | int, str] = {} + for i, arg in enumerate(call_node.args): + names[i] = _ast.unparse(arg) + for kw in call_node.keywords: + if kw.arg: + names[kw.arg] = _ast.unparse(kw.value) + return names + + diff --git a/python/CuTeDSL/_mlir_helpers/vector.py b/python/CuTeDSL/_mlir_helpers/vector.py new file mode 100644 index 000000000..5a79c109a --- /dev/null +++ b/python/CuTeDSL/_mlir_helpers/vector.py @@ -0,0 +1,1358 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +""" +This module provides MLIR Vector Dialect helper classes. +""" + +import array +import builtins +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any, Literal, Optional, Type, Union, cast + +from .._mlir import ir +from .._mlir.extras import types as T +from .._mlir.dialects import arith, vector, llvm +from .op import dsl_user_op +from .arith import ( + ArithValue, + _cast, + const, + cvtf, + element_type, + fptoi, + int_to_int, + itofp, +) + +if TYPE_CHECKING: + from ..base_dsl.typing import Int32, Numeric + +# ============================================================================= +# Vector Type - Immutable on registers +# ============================================================================= + + +@ir.register_value_caster(ir.VectorType.static_typeid) +class Vector(ArithValue): + """Wrap an MLIR ``vector`` register value with DSL type information. + + Provides element extraction (``vec[i]`` / ``vec[a:b]``), element-wise + arithmetic (``+``, ``-``, ``*``, ``/``), type conversion (:meth:`to`), + and bit-reinterpretation (:meth:`bitcast`) on top of a raw MLIR vector. + + Vectors live entirely in registers — they carry no memory address and do + not support in-place element assignment. + + Registered as the MLIR value caster for :class:`ir.VectorType`, so any + op that returns a vector automatically produces a ``Vector`` instance. + + :param v: Underlying MLIR vector value. + :type v: ir.Value + :param dtype: DSL element type (e.g. ``Float32``, ``Int32``). + Inferred from the MLIR element type when omitted. + :type dtype: type, optional + """ + + _dtype: Type["Numeric"] + _mlir_type: ir.Type + _shape: "tuple[int, ...]" + + # Cache parameterized subclasses so ``Vector[T, N] is Vector[T, N]``. + _parameterized_cache: "dict[tuple, type]" = {} + + @classmethod + def __class_getitem__(cls, params: "tuple[type, int | tuple[int, ...]]") -> type: + """Return a parameterized Vector subclass with a ``mlir_type`` descriptor. + + ``Vector[Int32, 4].mlir_type`` returns ``vector<4xi32>`` and + ``Vector[Float32, (4, 8)].mlir_type`` returns ``vector<4x8xf32>``, + matching the dual type-descriptor / value-constructor pattern of + scalar types like ``Int32``. Follows the same approach as + ``Pointer.__class_getitem__``. + """ + element_type, shape = params + + from ..base_dsl.typing import Numeric + + if not (isinstance(element_type, type) and issubclass(element_type, Numeric)): + raise TypeError( + f"Vector element type must be a Numeric type (Integer or Float), " + f"got {element_type!r}" + ) + if isinstance(shape, int): + shape = (shape,) + shape = tuple(shape) + if not shape or any(d <= 0 for d in shape): + raise ValueError( + f"Vector shape dimensions must be positive integers, got {shape}" + ) + key = (cls, element_type, shape) + if key not in cls._parameterized_cache: + shape_str = "x".join(str(d) for d in shape) + + class _Parameterized(cls): # type: ignore[valid-type, misc] + """Vector subclass with an ``mlir_type`` type descriptor.""" + + class mlir_type: # noqa: N801 + def __get__( + self, obj: object, objtype: Optional[type] = None + ) -> ir.Type: + return ir.VectorType.get(list(shape), element_type.mlir_type) # type: ignore[arg-type, attr-defined] + + mlir_type = mlir_type() # type: ignore[misc, assignment] + + @staticmethod + def __get_mlir_types__() -> list: + """Return MLIR types list — compatible with FFI ``_to_mlir_types``.""" + return [ir.VectorType.get(list(shape), element_type.mlir_type)] # type: ignore[arg-type, attr-defined] + + @staticmethod + def isinstance(value: object) -> bool: + """Check if an ``ir.Value`` matches this parameterized vector type.""" + if not builtins.isinstance(value, ir.Value): + return False + ty = value.type + if not builtins.isinstance(ty, ir.VectorType): + return False + return ( + list(ty.shape) == list(shape) # type: ignore[arg-type] + and ty.element_type == element_type.mlir_type # type: ignore[attr-defined] + ) + + _Parameterized.__name__ = f"Vector[{element_type.__name__}, {shape_str}]" + _Parameterized.__qualname__ = _Parameterized.__name__ + cls._parameterized_cache[key] = _Parameterized + return cls._parameterized_cache[key] + + def __init__( + self, + v: "ir.Value", + *, + dtype: "Type[Numeric] | None" = None, + loc: "ir.Location | None" = None, + ip: "ir.InsertionPoint | None" = None, + ) -> None: + # Derive signedness from dtype for ArithValue base + signed = getattr(dtype, "signed", None) + super().__init__(v, signed, loc=loc, ip=ip) + + # Infer dtype from MLIR element type if not provided + if dtype is None: + from ..base_dsl.common import DSLRuntimeError + from ..base_dsl.typing import Numeric + + try: + dtype = Numeric.from_mlir_type(self.type.element_type) + except DSLRuntimeError as exc: + try: + llvm.PointerType(self.type.element_type) + except Exception: + raise exc from None + dtype = None + # Pointer element vectors are transient lowering values; numeric vector + # APIs still require a Numeric dtype and will fail if used on them. + self._dtype = cast("Type[Numeric]", dtype) + self._mlir_type = ( + dtype.mlir_type if dtype is not None else self.type.element_type + ) + + # Shape is always derived from MLIR vector type + self._shape = tuple(self.type.shape) + + @staticmethod + @dsl_user_op + def from_elements( + scalars: tuple, + dtype: Type["Numeric"], + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> "Vector": + """Build a 1-D ``Vector`` from a tuple of scalar values.""" + if len(scalars) == 0: + raise ValueError("Vector.from_elements requires at least one element") + elem_ty = dtype.mlir_type + elements = [_value_to_ir_value(s, elem_ty, loc=loc, ip=ip) for s in scalars] + vec_type = ir.VectorType.get([len(scalars)], elem_ty) + ir_val = vector.from_elements(vec_type, elements, loc=loc, ip=ip) + return Vector(ir_val, dtype=dtype, loc=loc, ip=ip) + + # ========================================================================= + # DSL Infrastructure + # ========================================================================= + + def __extract_mlir_values__(self) -> list: + return [self] + + def __extract_mlir_attributes__(self) -> list: + return [ir.DictAttr.get({})] + + def __new_from_mlir_values__(self, values: list) -> "Vector": + return self._wrap_like(values[0]) + + def _wrap_like(self, result_ir: "ir.Value") -> "Vector": + """Construct a new instance of the same type from a result IR value. + + Subclasses (e.g., :class:`TensorSSA`) override to preserve their own + metadata (CuTe nested shape, layout). The math foundation uses this + for polymorphic result wrapping so a math op on a TensorSSA returns + a TensorSSA, not a Vector. + """ + return type(self)(result_ir, dtype=self._dtype) + + def _wrap_result( + self, + result_ir: "ir.Value", + *, + dtype: "type | None" = None, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "Vector": + """Construct a result wrapper, optionally with a different dtype.""" + return type(self)( + result_ir, + dtype=self._dtype if dtype is None else dtype, + loc=loc, + ip=ip, + ) + + def with_signedness(self, signed: Union[bool, None]) -> "Vector": + """Override ArithValue.with_signedness for keyword-only __init__.""" + new_vec = self._wrap_like(self) + elem_ty = self.type.element_type + new_vec.signed = ( + signed + and ir.IntegerType.isinstance(elem_ty) + and ir.IntegerType(elem_ty).width > 1 + ) + return new_vec + + # ========================================================================= + # Properties + # ========================================================================= + + @property + def dtype(self) -> Type["Numeric"]: + """The DSL element type (e.g., Float32, Int32).""" + return self._dtype + + @property + def shape(self) -> "tuple[int, ...]": + """The logical shape of the vector array (1D, 2D, or 3D).""" + return self._shape + + @property + def _count(self) -> int: + """Total number of elements (product of shape dimensions).""" + result = 1 + for dim in self._shape: + result *= dim + return result + + def numel(self) -> int: + """Total number of elements (product of all shape dimensions).""" + return self._count + + @staticmethod + def _normalize_static_slice( + idx: slice, + extent: int, + *, + error_context: str, + ) -> "tuple[int, int, int, list[int]]": + """Normalize a Python slice and reject unsupported empty/reverse forms.""" + if not all( + isinstance(v, int) or v is None for v in (idx.start, idx.stop, idx.step) + ): + context = f" {error_context}" if error_context else "" + raise TypeError( + f"Vector slice indices must be static ints{context}; " + f"got start={idx.start}, stop={idx.stop}, step={idx.step}" + ) + try: + start, stop, step = idx.indices(extent) + except ValueError as exc: + raise ValueError("Vector slices require a non-zero step") from exc + if step <= 0: + raise ValueError("Vector slices require a positive step") + + positions = list(range(start, stop, step)) + if not positions: + raise IndexError("Empty vector slices are not supported") + if any(pos < 0 or pos >= extent for pos in positions): + raise IndexError( + f"Vector slice produced out-of-bounds indices {positions} " + f"for extent {extent}" + ) + return start, stop, step, positions + + # Vector has no memory space - it's always in registers + # The space property is intentionally not exposed on Vector + + def ir_value( + self, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> ir.Value: + """Return the underlying MLIR vector value.""" + return self + + # ========================================================================= + # Indexing Operations + # ========================================================================= + + def _compute_linear_index( + self, + indices: "tuple[Union[int, Int32], ...]", # type: ignore[name-defined] + ) -> "Union[int, Int32]": # type: ignore[name-defined] + """Compute linear index from multi-dimensional indices (row-major order).""" + if len(indices) != len(self._shape): + raise IndexError( + f"Expected {len(self._shape)} indices for shape {self._shape}, " + f"got {len(indices)}" + ) + + # Check if all indices are static (Python ints) + all_static = all(isinstance(i, int) for i in indices) + + if all_static: + # Static computation + linear = 0 + stride = 1 + for i in range(len(self._shape) - 1, -1, -1): + linear += indices[i] * stride + stride *= self._shape[i] + return linear + else: + from ..base_dsl.typing import Int32 + + # Dynamic computation using Int32 arithmetic + linear = Int32(0) # type: ignore[assignment] + stride = 1 + for i in range(len(self._shape) - 1, -1, -1): + idx = indices[i] if isinstance(indices[i], Int32) else Int32(indices[i]) + linear = linear + idx * Int32(stride) + stride *= self._shape[i] + return linear + + def __getitem__( + self, + idx: "Union[int, Int32, tuple, slice]", # type: ignore[name-defined] + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> object: + """Extract an element or a contiguous sub-vector. + + Supports three indexing modes: + + * **Scalar index** — returns a single DSL scalar value:: + + elem = vec[i] # static int or Int32 + + * **1-D slice** — all bounds must be static Python ``int``s:: + + sub = vec[start:stop] # stride defaults to 1 + sub = vec[start:stop:stride] # explicit stride + + * **Multi-dimensional slice** — one entry per dimension, all bounds + must be static ``int``s. An integer in a multi-dim slice is treated + as a size-1 slice (the dimension is kept):: + + sub = mat[r0:r1, c0:c1] # 2-D: rows r0:r1, cols c0:c1 + sub = mat[:, c0:c1] # 2-D: all rows, cols c0:c1 + sub = mat[0, c0:c1] # 2-D: row 0 (size 1), cols c0:c1 + + Slices use ``vector.extract_strided_slice`` internally; dynamic + (MLIR-value) slice bounds are **not** supported. + + :param idx: Element index (int or Int32), a slice, or a tuple of + ints/slices for multi-dimensional access. + :type idx: int or Int32 or tuple or slice + :return: A scalar DSL value (for element indexing) or a new + :class:`Vector` (for slice indexing). + :rtype: Numeric or Vector + :raises TypeError: If slice bounds are not static Python ints. + :raises IndexError: If the number of dimensions in a multi-dim index + does not match :attr:`shape`. + """ + from ..base_dsl.utils.logger import log + + # Slice → extract_strided_slice (step==1) or vector.shuffle (step>1) + if isinstance(idx, slice): + if len(self._shape) != 1: + raise TypeError( + "Use per-dimension indexing for multi-dimensional vectors, " + "e.g. vec[a:b, :]" + ) + start, _stop, step, positions = self._normalize_static_slice( + idx, self._count, error_context="" + ) + size = len(positions) + result_ty = ir.VectorType.get([size], self._mlir_type) + if step == 1: + result = vector.extract_strided_slice( + result_ty, self, [start], [size], [step], loc=loc, ip=ip + ) + else: + # vector.extract_strided_slice requires stride==1; use shuffle instead + result = vector.shuffle(self, self, positions, loc=loc, ip=ip) + return self._wrap_result(result, loc=loc, ip=ip) + + # Multi-dimensional slice: tuple containing at least one slice object + if isinstance(idx, tuple) and any(isinstance(i, slice) for i in idx): + if len(idx) != len(self._shape): + raise IndexError( + f"Expected {len(self._shape)} indices for shape {self._shape}, " + f"got {len(idx)}" + ) + offsets: "list[int]" = [] + sizes: "list[int]" = [] + strides: "list[int]" = [] + for dim, (i, dim_size) in enumerate(zip(idx, self._shape)): + if isinstance(i, slice): + start, _stop, step, positions = self._normalize_static_slice( + i, dim_size, error_context=f"in dimension {dim}" + ) + if step != 1: + raise NotImplementedError( + f"Multi-dimensional strided slice (step={step}) is not supported; " + "use step=1 for multi-dimensional slices" + ) + offsets.append(start) + sizes.append(len(positions)) + strides.append(1) + elif isinstance(i, int): + # Integer index: treated as a size-1 slice (rank is preserved) + if i < 0: + i += dim_size + if i < 0 or i >= dim_size: + raise IndexError( + f"Vector index {i} out of bounds for dimension {dim} " + f"with size {dim_size}" + ) + offsets.append(i) + sizes.append(1) + strides.append(1) + else: + raise TypeError( + f"Vector multi-dimensional slice: dimension {dim} index must be " + f"a static int or slice, got {type(i).__name__}" + ) + result_ty = ir.VectorType.get(sizes, self._mlir_type) + result = vector.extract_strided_slice( + result_ty, self, offsets, sizes, strides, loc=loc, ip=ip + ) + return self._wrap_result(result, loc=loc, ip=ip) + + # Normalize to tuple + if not isinstance(idx, tuple): + indices = (idx,) + else: + indices = idx + + if len(indices) != len(self._shape): + raise IndexError( + f"Expected {len(self._shape)} indices for shape {self._shape}, " + f"got {len(indices)}" + ) + + if all(isinstance(i, int) for i in indices): + static_indices: "list[int]" = [] + for dim, (i, dim_size) in enumerate(zip(indices, self._shape)): + if i < 0: + i += dim_size + if i < 0 or i >= dim_size: + raise IndexError( + f"Vector index {i} out of bounds for dimension {dim} " + f"with size {dim_size}" + ) + static_indices.append(i) + log().info( + f"Vector.__getitem__: idx={idx}, static={static_indices}, " + f"dtype={self._dtype}, shape={self._shape}" + ) + elem = vector.extract(self, [], static_indices, loc=loc, ip=ip) + else: + if len(self._shape) != 1: + raise NotImplementedError( + "Dynamic Vector indexing is currently only supported for 1-D vectors" + ) + linear_idx = self._compute_linear_index(indices) + log().info( + f"Vector.__getitem__: idx={idx}, linear={linear_idx}, " + f"dtype={self._dtype}, shape={self._shape}" + ) + if isinstance(linear_idx, int): + from ..base_dsl.typing import Int32 + + linear_idx = Int32(linear_idx) + elem = llvm.extractelement(self, linear_idx.ir_value(), loc=loc, ip=ip) + + return self._dtype(elem) + + @dsl_user_op + def to_elements( + self, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "tuple[Numeric, ...]": + """Extract every vector lane as scalar DSL values. + + This is useful when a vectorized operation should define many scalar + SSA values that are then consumed independently. + """ + raw_elements = vector.to_elements(self.ir_value(loc=loc, ip=ip)) + if self._count == 1 and not isinstance(raw_elements, Sequence): + raw_elements = (raw_elements,) + return tuple(self._dtype(elem, loc=loc, ip=ip) for elem in raw_elements) + + def __setitem__( + self, + idx: "Union[int, Int32, tuple]", # type: ignore[name-defined] + value: object, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> None: + """ + Vector element assignment is not supported. + + Vectors are immutable register values. Use one of these alternatives: + + 1. Use a frontend memory-backed tensor or array abstraction for + mutable storage. + + 2. Use full() to create vectors with initial values: + vec = cute.full((4,), 1.0, dtype=cutlass.Float32).to_vector() + """ + raise TypeError( + "Vector is immutable. Element assignment (vec[i] = value) is not supported. " + ) + + # ========================================================================= + # Arithmetic Operations + # ========================================================================= + + def _is_float_type(self) -> bool: + """Check if this vector contains floating-point elements.""" + return self._dtype.is_float + + # Arithmetic operators delegate to ArithValue and wrap results with the + # same concrete Vector subclass. + + def to( + self, + dtype: "type", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "Vector": + """Convert the vector elements to a different numeric type. + + :param dtype: Target DSL element type (e.g. ``Float16``, ``Int32``). + :type dtype: Type[Numeric] + :return: A new :class:`Vector` with the same shape and elements cast + to ``dtype``. + :rtype: Vector + :raises TypeError: If ``dtype`` is not a subclass of ``Numeric``. + + Example:: + + vec_f32 = cute.full((4,), 1.5, dtype=cutlass.Float32).to_vector() + vec_i32 = vec_f32.to(cutlass.Int32) # fp → int truncation + vec_f16 = vec_f32.to(cutlass.Float16) # fp32 → fp16 narrowing + """ + from inspect import isclass + from ..base_dsl.typing import Numeric, Integer + + if dtype is ir.Value: + return self + + if not isclass(dtype) or not issubclass(dtype, Numeric): + raise TypeError(f"dtype must be a type of Numeric, but got {type(dtype)}") + + src_dtype = self._dtype + if src_dtype == dtype: + return self + + # maybe_downcast handles narrow precision types, with_signedness sets signedness + src = self.maybe_downcast().with_signedness(self.signed) + + if src_dtype.is_float and dtype.is_float: + res_vect = cvtf(src, dtype.mlir_type, loc=loc, ip=ip) + elif src_dtype.is_float and issubclass(dtype, Integer): + res_vect = fptoi(src, dtype.signed, dtype.mlir_type, loc=loc, ip=ip) + elif issubclass(src_dtype, Integer) and dtype.is_float: + res_vect = itofp(src, src_dtype.signed, dtype.mlir_type, loc=loc, ip=ip) + else: + res_vect = int_to_int(src, dtype, loc=loc, ip=ip) + + return self._wrap_result(res_vect, dtype=dtype, loc=loc, ip=ip) + + @dsl_user_op + def bitcast( + self, + dtype: "type", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "Vector": + """Reinterpret the vector bits as a different element type. + + The total bit width is preserved; the element count adjusts + proportionally. For example, ``vector<4xi32>`` bitcast to + ``Float16`` yields ``vector<8xf16>`` (4 × 32 = 8 × 16 bits). + + :param dtype: Target DSL element type (e.g. ``Float32``, ``Float16``). + :type dtype: Type[Numeric] + :return: A new :class:`Vector` with bits reinterpreted as ``dtype``. + :rtype: Vector + :raises TypeError: If ``dtype`` is not a subclass of ``Numeric``. + """ + from inspect import isclass + from ..base_dsl.typing import Numeric + + if not isclass(dtype) or not issubclass(dtype, Numeric): + raise TypeError(f"dtype must be a Numeric type, but got {dtype}") + if dtype is self._dtype: + return self + total_bits = self._count * self._dtype.width + if total_bits % dtype.width != 0: + raise ValueError( + f"Cannot bitcast {self.type} to {dtype}: " + "total bit width must match exactly" + ) + new_count = total_bits // dtype.width + target_vec_ty = T.vector(new_count, dtype.mlir_type) + res_vec = vector.bitcast(target_vec_ty, self, loc=loc, ip=ip) + return self._wrap_result(res_vec, dtype=dtype, loc=loc, ip=ip) + + def _broadcast_scalar_operand( + self, + other: Any, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> Any: + """If *other* is a runtime scalar, splat it to this Vector's shape. + + Python ``int`` / ``float`` / ``bool`` literals pass through + (``_binary_op`` will coerce via ``const`` with self.type, which + already produces a vector-shaped constant). Runtime scalar + ``ArithValue`` / ``Numeric`` values don't: they retain their 0-dim + MLIR type and the arith dialect rejects `arith.mulf(vec, scalar)`. + Broadcast them via ``vector.broadcast`` so ``vec +/-/*/÷ scalar`` + behaves like a full-vector splat of that scalar. + """ + if isinstance(other, (int, float, bool)): + return other + # Normalise to ir.Value — Numeric wrappers expose .ir_value(); bare + # ir.Value / ArithValue already is one. + if hasattr(other, "ir_value"): + scalar_ir = other.ir_value(loc=loc, ip=ip) + else: + scalar_ir = other + # Runtime Vector operand — already shape-matched; leave alone. + if isinstance(scalar_ir.type, ir.VectorType): + return other + # Runtime scalar — element type must match the Vector's element type + # before `vector.broadcast`, otherwise the op verifier rejects the + # splat (``result.elementType != operand.type``). Cast first when the + # runtime scalar was produced at a different precision (e.g. a + # Float32 constant into an fp16 Vector, or an Int32 index into an + # Int64 Vector). `_cast` is a no-op when types already match. + vec_elem = self.type.element_type + if scalar_ir.type != vec_elem: + scalar_ir = _cast(vec_elem, scalar_ir, loc=loc, ip=ip) + splat = vector.broadcast(self.type, scalar_ir, loc=loc, ip=ip) + return self._wrap_like(splat) + + @dsl_user_op + def __add__( + self, + other: "Vector", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "Vector": + other = self._broadcast_scalar_operand(other, loc=loc, ip=ip) + result = super().__add__(other, loc=loc, ip=ip) + return self._wrap_like(result) + + @dsl_user_op + def __radd__( + self, + other: "Vector", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "Vector": + other = self._broadcast_scalar_operand(other, loc=loc, ip=ip) + result = super().__radd__(other, loc=loc, ip=ip) + return self._wrap_like(result) + + @dsl_user_op + def __sub__( + self, + other: "Vector", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "Vector": + other = self._broadcast_scalar_operand(other, loc=loc, ip=ip) + result = super().__sub__(other, loc=loc, ip=ip) + return self._wrap_like(result) + + @dsl_user_op + def __rsub__( + self, + other: "Vector", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "Vector": + other = self._broadcast_scalar_operand(other, loc=loc, ip=ip) + result = super().__rsub__(other, loc=loc, ip=ip) + return self._wrap_like(result) + + @dsl_user_op + def __mul__( + self, + other: "Vector", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "Vector": + other = self._broadcast_scalar_operand(other, loc=loc, ip=ip) + result = super().__mul__(other, loc=loc, ip=ip) + return self._wrap_like(result) + + @dsl_user_op + def __rmul__( + self, + other: "Vector", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "Vector": + other = self._broadcast_scalar_operand(other, loc=loc, ip=ip) + result = super().__rmul__(other, loc=loc, ip=ip) + return self._wrap_like(result) + + @dsl_user_op + def __truediv__( + self, + other: "Vector", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "Vector": + other = self._broadcast_scalar_operand(other, loc=loc, ip=ip) + result = super().__truediv__(other, loc=loc, ip=ip) + return self._wrap_like(result) + + @dsl_user_op + def __rtruediv__( + self, + other: "Vector", + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> "Vector": + other = self._broadcast_scalar_operand(other, loc=loc, ip=ip) + result = super().__rtruediv__(other, loc=loc, ip=ip) + return self._wrap_like(result) + + # ========================================================================= + # Reduction — vector.reduction wrapper + # ========================================================================= + + _REDUCE_KINDS = { + "add": lambda self: vector.CombiningKind.ADD, + "mul": lambda self: vector.CombiningKind.MUL, + "min": lambda self: ( + vector.CombiningKind.MINNUMF + if self._is_float_type() + else ( + vector.CombiningKind.MINSI + if getattr(self._dtype, "signed", True) + else vector.CombiningKind.MINUI + ) + ), + "max": lambda self: ( + vector.CombiningKind.MAXNUMF + if self._is_float_type() + else ( + vector.CombiningKind.MAXSI + if getattr(self._dtype, "signed", True) + else vector.CombiningKind.MAXUI + ) + ), + } + + @dsl_user_op + def reduce( + self, + op: Literal["add", "mul", "min", "max"] = "add", + *, + dim: Optional[Union[int, list[int]]] = None, + acc: Any = None, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> Any: + """Reduce the vector using the specified combining operation. + + When ``dim`` is ``None`` (default), reduces **all** dimensions to a + scalar via ``vector.reduction``. When ``dim`` is an int or list of + ints, reduces only those dimensions via ``vector.multi_reduction``, + returning a lower-rank :class:`Vector`. + + :param op: Reduction operation — one of ``"add"``, ``"mul"``, + ``"min"``, ``"max"``. For ``"min"``/``"max"`` the combining + kind adapts automatically to the element type (float vs signed + vs unsigned integer). + :param dim: Dimension(s) to reduce. ``None`` reduces all dims to a + scalar. An int or list of ints reduces only those dims. + :param acc: Optional accumulator. For scalar reduction a scalar value; + for multi-dim reduction a vector matching the result shape. + :return: Scalar (when ``dim is None``) or :class:`Vector` (when + ``dim`` is specified). + + Examples: + + .. code-block:: python + + v = cute.full((4,), 3.0, dtype=cutlass.Float32).to_vector() + v.reduce("add") # 12.0 (scalar) + + m = cute.full((4, 8), 1.0, dtype=cutlass.Float32).to_vector() + m.reduce("add", dim=1) # vector<4xf32>, each element = 8.0 + m.reduce("add", dim=0) # vector<8xf32>, each element = 4.0 + + .. note:: + + This method operates on a ``Vector`` value. If a higher-level API + in a downstream library returns a different SSA wrapper with its + own ``reduce(...)`` method and a different signature, call that + library's ``.to_vector()`` (or equivalent) to get a plain + ``Vector`` first so this 1-arg form applies. + + .. note:: + + ``Vector.reduce`` builds an MLIR ``vector.reduction`` over the + elements of one register vector. It is not the warp-collective + ``nvvm.redux_sync`` API. If a backend or target-specific + lowering maps a reduction to PTX ``redux.sync``, PTX legality still + applies: integer/bitwise ``redux.sync`` forms require ``sm_80`` or + higher, while ``redux.sync`` ``.f32`` min/max support was added in + PTX ISA 8.6 and is limited to ``sm_100a`` plus the ``sm_100f`` + family support added in PTX ISA 8.8. For examples that must remain + portable to generic ``sm_120`` targets, prefer an explicit scalar + fold or shuffle tree for ``Float32`` min/max instead of relying on + a lowering that may choose ``redux.sync.f32``. + """ + kind_fn = self._REDUCE_KINDS.get(op) + if kind_fn is None: + raise ValueError( + f"Unknown reduction op {op!r}. " + f"Expected one of: {', '.join(sorted(self._REDUCE_KINDS))}" + ) + kind = kind_fn(self) + vec_ty = ir.VectorType(self.type) + elem_ty = vec_ty.element_type + + ndim = len(vec_ty.shape) + + if dim is None and ndim == 1: + # 1-D full reduction to scalar — wrap in _dtype so type info is preserved + raw = vector.reduction(elem_ty, kind, self, acc=acc, loc=loc, ip=ip) + return self._dtype(raw) + + # Multi-dimension reduction + if dim is None: + # Reduce all dims for N-D vector + reduction_dims = list(range(ndim)) + else: + reduction_dims = [dim] if isinstance(dim, int) else list(dim) + if not reduction_dims: + raise ValueError("Reduction dims must be non-empty") + if len(set(reduction_dims)) != len(reduction_dims): + raise ValueError(f"Reduction dims must be unique, got {reduction_dims}") + for d in reduction_dims: + if d < 0 or d >= ndim: + raise ValueError(f"Reduction dim {d} out of range for {ndim}-D vector") + + # Compute result shape and build accumulator + result_shape = [ + s for i, s in enumerate(vec_ty.shape) if i not in reduction_dims + ] + if not result_shape: + # All dims reduced on N-D vector — reduce to 1-D first, then scalar + # Keep dim 0 to get a 1-D vector via multi_reduction + partial_dims = list(range(1, ndim)) + partial_shape = [vec_ty.shape[0]] + partial_acc_ty = ir.VectorType.get(partial_shape, elem_ty) + init_map = { + "add": 0.0 if self._is_float_type() else 0, + "mul": 1.0 if self._is_float_type() else 1, + } + if op in init_map: + p_init = init_map[op] + else: + import math + + if self._is_float_type(): + p_init = math.inf if op == "min" else -math.inf + else: + width = self._dtype.width + signed = getattr(self._dtype, "signed", True) + if op == "min": + p_init = (1 << (width - 1)) - 1 if signed else (1 << width) - 1 + else: + p_init = -(1 << (width - 1)) if signed else 0 + p_scalar = arith.constant(elem_ty, p_init, loc=loc, ip=ip) + p_acc = vector.broadcast(partial_acc_ty, p_scalar, loc=loc, ip=ip) + vec_1d = vector.multi_reduction( + kind, self, acc=p_acc, reduction_dims=partial_dims, loc=loc, ip=ip + ) + raw = vector.reduction(elem_ty, kind, vec_1d, acc=acc, loc=loc, ip=ip) + return self._dtype(raw) + + if acc is None: + # Build identity accumulator for the result shape + init_map = { + "add": 0.0 if self._is_float_type() else 0, + "mul": 1.0 if self._is_float_type() else 1, + } + if op in init_map: + init_val = init_map[op] + else: + # For min/max, use the first slice as initial value — + # MLIR requires an acc, so broadcast a neutral constant. + # Use extreme values: for float min → +inf, max → -inf + import math + + if self._is_float_type(): + init_val = math.inf if op == "min" else -math.inf + else: + # Integer: use max/min representable value + width = self._dtype.width + signed = getattr(self._dtype, "signed", True) + if op == "min": + init_val = ( + (1 << (width - 1)) - 1 if signed else (1 << width) - 1 + ) + else: # max + init_val = -(1 << (width - 1)) if signed else 0 + + acc_ty = ir.VectorType.get(result_shape, elem_ty) + init_scalar = arith.constant(elem_ty, init_val, loc=loc, ip=ip) + acc = vector.broadcast(acc_ty, init_scalar, loc=loc, ip=ip) + + result = vector.multi_reduction( + kind, self, acc=acc, reduction_dims=reduction_dims, loc=loc, ip=ip + ) + return self._wrap_result(result, loc=loc, ip=ip) + + +def _shape_tuple(shape: Union[int, Sequence[int]]) -> tuple[int, ...]: + if isinstance(shape, int): + shape = (shape,) + return tuple(shape) + + +def _scalar_to_ir_value( + scalar: object, + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, +) -> ir.Value: + from ..base_dsl.typing import Numeric + + if isinstance(scalar, Numeric): + return scalar.ir_value(loc=loc, ip=ip) + if isinstance(scalar, ArithValue): + return scalar.ir_value(loc=loc, ip=ip) + if isinstance(scalar, ir.Value): + return scalar + if isinstance(scalar, (int, float, bool)): + return const(scalar, loc=loc, ip=ip) + raise ValueError(f"Expected scalar value, got {scalar}") + + +def _value_to_ir_value( + value: object, + elem_type: ir.Type, + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, +) -> ir.Value: + from ..base_dsl.typing import Numeric + + if isinstance(value, (int, float, bool)): + return const(value, elem_type, loc=loc, ip=ip) + + if isinstance(value, Numeric): + value = value.ir_value(loc=loc, ip=ip) + elif isinstance(value, ArithValue): + value = value.ir_value(loc=loc, ip=ip) + + if isinstance(value, ir.Value): + if value.type != elem_type: + value = _cast(elem_type, value, loc=loc, ip=ip) + return value + + raise NotImplementedError( + f"Expected a Python scalar, Numeric, ArithValue, or ir.Value, " + f"got {type(value).__name__!r}" + ) + + +def _infer_element_type(value: object) -> ir.Type: + from ..base_dsl.typing import Numeric + + if isinstance(value, float): + return T.f32() + if isinstance(value, bool): + return T.bool() + if isinstance(value, int): + return T.i32() + if isinstance(value, Numeric): + return value.dtype.mlir_type + if isinstance(value, (ArithValue, ir.Value)): + return element_type(value.type) + raise NotImplementedError( + f"Cannot infer element type from value of type {type(value).__name__!r}" + ) + + +@dsl_user_op +def create_mask( + shape: Sequence[int], + dim_sizes: Sequence[Union[int, "Numeric"]], + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, +) -> Vector: + """Create a boolean ``Vector`` mask for ``shape`` and dynamic bounds. + + ``mask[i]`` is true when each index is less than the corresponding + exclusive bound in ``dim_sizes``. For a keep-through-last-valid pattern, + pass ``last_valid + 1``. + """ + if not isinstance(dim_sizes, Sequence): + raise TypeError("dim_sizes must be a sequence of integers") + + shape = _shape_tuple(shape) + dim_size_values = [_scalar_to_ir_value(s, loc=loc, ip=ip) for s in dim_sizes] + bool_ty = T.bool() + res_ty = ir.VectorType.get(list(shape), element_type=bool_ty) + i64_ty = ir.IntegerType.get_signless(64) + + def _cast_bound_to_i64(s: ir.Value) -> ir.Value: + if ir.IntegerType.isinstance(s.type): + width = ir.IntegerType(s.type).width + if width == 64: + return s + if width > 64: + return arith.trunci(i64_ty, s, loc=loc, ip=ip) + if getattr(s, "signed", True) is False: + return arith.extui(i64_ty, s, loc=loc, ip=ip) + return arith.extsi(i64_ty, s, loc=loc, ip=ip) + if ir.IndexType.isinstance(s.type): + return arith.index_cast(i64_ty, s, loc=loc, ip=ip) + raise TypeError( + f"Mask dimension sizes must be integer or index values, got {s.type}" + ) + + def _build_1d_mask(n: int, s: ir.Value) -> ir.Value: + s = _cast_bound_to_i64(s) + iota_ty = ir.VectorType.get([n], element_type=i64_ty) + iota_attr = ir.DenseElementsAttr.get( + array.array("q", list(range(n))), type=iota_ty + ) + iota = arith.constant(iota_ty, iota_attr, loc=loc, ip=ip) + splat_ty = ir.VectorType.get([n], element_type=i64_ty) + dim_splat = vector.broadcast(splat_ty, s, loc=loc, ip=ip) + return arith.cmpi(arith.CmpIPredicate.slt, iota, dim_splat, loc=loc, ip=ip) + + if len(shape) == 1: + mask = _build_1d_mask(shape[0], dim_size_values[0]) + elif len(shape) == 2: + m, n = shape + col_mask = _build_1d_mask(n, dim_size_values[1]) + col_bc_ty = ir.VectorType.get([m, n], element_type=bool_ty) + col_bc = vector.broadcast(col_bc_ty, col_mask, loc=loc, ip=ip) + + d0 = _cast_bound_to_i64(dim_size_values[0]) + + d0_bc_ty = ir.VectorType.get([m, n], element_type=i64_ty) + d0_bc = vector.broadcast(d0_bc_ty, d0, loc=loc, ip=ip) + row_indices = [] + for row in range(m): + row_indices.extend([row] * n) + row_iota_2d_ty = ir.VectorType.get([m, n], element_type=i64_ty) + row_iota_2d_attr = ir.DenseElementsAttr.get( + array.array("q", row_indices), type=row_iota_2d_ty + ) + row_iota_2d = arith.constant(row_iota_2d_ty, row_iota_2d_attr, loc=loc, ip=ip) + row_bc = arith.cmpi(arith.CmpIPredicate.slt, row_iota_2d, d0_bc, loc=loc, ip=ip) + mask = arith.andi(row_bc, col_bc, loc=loc, ip=ip) + else: + index_ty = T.index() + dim_size_indices = [ + arith.index_cast(index_ty, s, loc=loc, ip=ip) + if not ir.IndexType.isinstance(s.type) + else s + for s in dim_size_values + ] + mask = vector.create_mask(res_ty, dim_size_indices, loc=loc, ip=ip) + return mask + + +@dsl_user_op +def broadcast_to( + src: Union["Numeric", Vector, int, float, bool, ir.Value], + target_shape: Sequence[int], + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, +) -> Vector: + """Broadcast a scalar or narrower ``Vector`` to ``target_shape``.""" + from ..base_dsl.typing import Numeric + + target_shape = _shape_tuple(target_shape) + if isinstance(src, (int, float, bool)): + elem_type = _infer_element_type(src) + src = const(src, elem_type, loc=loc, ip=ip) + elif isinstance(src, Numeric): + src = src.ir_value(loc=loc, ip=ip) + elem_type = element_type(src.type) + elif isinstance(src, Vector): + elem_type = src.dtype.mlir_type + elif isinstance(src, ArithValue): + elem_type = element_type(src.type) + src = src.ir_value(loc=loc, ip=ip) + elif isinstance(src, ir.Value): + elem_type = element_type(src.type) + else: + raise NotImplementedError( + f"broadcast_to does not support src of type {type(src).__name__!r}" + ) + + target_type = T.vector(*target_shape, elem_type) + return vector.broadcast(target_type, src, loc=loc, ip=ip) + + +@dsl_user_op +def full( + shape: Sequence[int], + fill_value: Union["Numeric", int, float, bool, ir.Value], + dtype: Optional[Type["Numeric"]] = None, + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, +) -> Vector: + """Return a new ``Vector`` of ``shape`` filled with ``fill_value``.""" + shape = _shape_tuple(shape) + elem_type = _infer_element_type(fill_value) if dtype is None else dtype.mlir_type + fill_val = _value_to_ir_value(fill_value, elem_type, loc=loc, ip=ip) + res_type = T.vector(*shape, elem_type) + return vector.broadcast(res_type, fill_val, loc=loc, ip=ip) + + +@dsl_user_op +def full_like( + a: Vector, + fill_value: Union["Numeric", int, float, bool, ir.Value], + dtype: Optional[Type["Numeric"]] = None, + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, +) -> Vector: + """Return a ``Vector`` filled with ``fill_value`` and shaped like ``a``.""" + return full(a.shape, fill_value, dtype or a.dtype, loc=loc, ip=ip) + + +@dsl_user_op +def zeros_like( + a: Vector, + dtype: Optional[Type["Numeric"]] = None, + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, +) -> Vector: + """Return a zero-filled ``Vector`` with the same shape as ``a``.""" + return full_like(a, 0, dtype or a.dtype, loc=loc, ip=ip) + + +@dsl_user_op +def ones_like( + a: Vector, + dtype: Optional[Type["Numeric"]] = None, + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, +) -> Vector: + """Return a one-filled ``Vector`` with the same shape as ``a``.""" + return full_like(a, 1, dtype or a.dtype, loc=loc, ip=ip) + + +@dsl_user_op +def empty_like( + a: Vector, + dtype: Optional[Type["Numeric"]] = None, + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, +) -> Vector: + """Return an uninitialized ``Vector`` shaped like ``a``. + + This currently returns zeros until the DSL has a distinct undef/poison + helper for register vectors. + """ + return full_like(a, 0, dtype or a.dtype, loc=loc, ip=ip) + + +@dsl_user_op +def where( + cond: Vector, + x: Union[Vector, "Numeric", int, float, bool, ir.Value], + y: Union[Vector, "Numeric", int, float, bool, ir.Value], + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, +) -> Vector: + """Return elements chosen from ``x`` or ``y`` depending on ``cond``.""" + result_dtype = None + if isinstance(x, Vector): + result_dtype = x.dtype + elif isinstance(y, Vector): + result_dtype = y.dtype + + def _promote(v: object) -> object: + if isinstance(v, Vector): + return v + if isinstance(v, (int, float, bool)): + return full(cond.shape, v, result_dtype, loc=loc, ip=ip) + return broadcast_to(v, cond.shape, loc=loc, ip=ip) + + return arith.select(cond, _promote(x), _promote(y), loc=loc, ip=ip) + + +@dsl_user_op +def any_( + x: Vector, + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, +) -> "Numeric": + """Return true if any vector element is non-zero.""" + from ..base_dsl.typing import Boolean + + zeros = zeros_like(x, loc=loc, ip=ip) + is_true = x != zeros + return Boolean( + vector.reduction(T.bool(), vector.CombiningKind.OR, is_true, loc=loc, ip=ip) + ) + + +@dsl_user_op +def all_( + x: Vector, + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, +) -> "Numeric": + """Return true if every vector element is non-zero.""" + from ..base_dsl.typing import Boolean + + zeros = zeros_like(x, loc=loc, ip=ip) + is_true = x != zeros + return Boolean( + vector.reduction(T.bool(), vector.CombiningKind.AND, is_true, loc=loc, ip=ip) + ) + + +@dsl_user_op +def outerproduct( + a: Vector, + b: Vector, + acc: Optional[Vector] = None, + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, +) -> Vector: + """Rank-1 outer product of two 1-D vectors with optional accumulation.""" + if len(a.shape) != 1: + raise ValueError(f"'a' must be 1-D, got shape {a.shape}") + if len(b.shape) != 1: + raise ValueError(f"'b' must be 1-D, got shape {b.shape}") + + result_type = ir.VectorType.get([a.shape[0], b.shape[0]], element_type(a.type)) + return vector.outerproduct( + result_type, + a, + b, + acc=acc, + kind=vector.CombiningKind.ADD if acc is not None else None, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def print_nd_vector( + src: Vector, + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, +) -> None: + """Print the contents of an N-D ``Vector`` at runtime. + + Device-side debug printf: each thread that reaches this call emits + its own line. Intended for correctness investigation, **not** for + production kernels: device ``printf`` serialises across the warp + and adds a large fixed overhead per call. + + :param src: Vector to print. N-D shapes are reshaped to 1-D + column-major internally, then printed via ``cute.print_tensor``. + :type src: Vector + + :sync-class: Per-thread debug printf. Every thread that reaches + this call emits one line. Output from different threads + interleaves non-deterministically. + :elect-safe: **Yes**. Wrap in ``if nvvm.elect_sync():`` to get one + line per warp instead of 32; wrap in ``if tid == 0:`` for one + line per CTA. + :device: All CUDA architectures. + :side-effects: writes to stdout via device ``printf``. Per-call + latency is large and serialises the warp; remove before + benchmarking or shipping. + """ + from .. import cute + from ..cute import tensor + + size = 1 + for dim in src.shape: + size *= dim + + tmp = vector.shape_cast( + ir.VectorType.get([size], src.dtype.mlir_type), src, loc=loc, ip=ip + ) + tmp = tensor._row2col(tmp, shape=tuple(src.shape), loc=loc, ip=ip) + tmp = tensor.TensorSSA(tmp, tuple(src.shape), src.dtype, loc=loc, ip=ip) + + cute.print_tensor(tmp, loc=loc, ip=ip) diff --git a/python/CuTeDSL/cutlass/__init__.py b/python/CuTeDSL/cutlass/__init__.py index 19bd2dd78..d89ad48ee 100644 --- a/python/CuTeDSL/cutlass/__init__.py +++ b/python/CuTeDSL/cutlass/__init__.py @@ -46,6 +46,7 @@ _common._get_cuda_version = _get_cuda_version # Import CUDA version from base_dsl from .base_dsl.version_info import CUDA_VERSION +from .address_space import AddressSpace from .cutlass_dsl import ( Constexpr, @@ -74,7 +75,7 @@ from .cutlass_dsl import ( # Data types dtype, # Provides conversions to types inheriting from NumericType DSLRuntimeError, - DSLAstPreprocessorError, + DSLUserCodeError, JitArgAdapterRegistry, # Construction utilities for user-defined classes extract_mlir_values, @@ -84,13 +85,24 @@ from .cutlass_dsl import ( ) from .cute.typing import * +from .base_dsl.typing import Pointer as Pointer # type: ignore[assignment] +from .base_dsl.typing import ( + TypedPointer, + Uint128, + align, +) +from .base_dsl.pointer import inttoptr +from .base_dsl.swizzle import Swizzle, apply_swizzle, load_swizzled, store_swizzled +from ._mlir_helpers.vector import Vector # Utilities not belonging to CuTe +from .base_dsl import vector as vector + +from . import runtime as runtime from . import utils as utils from . import pipeline as pipeline from . import testing as testing - -# Used as internal symbol +# Package-private symbol used by exported aliases below. from . import cutlass_dsl as _dsl # Aliases @@ -104,3 +116,40 @@ from . import jax as jax CACHE_FILE = "compiled_cache.db" +# --- cutlass namespace extension gate ---------------------------------------- +# `cutlass` is shipped as a regular package (this __init__.py) +# rather than a PEP 420 namespace, so this file fully controls +# what goes in the `cutlass.*` namespace. +# +# A small, explicitly-named set of NVIDIA-owned sibling distributions outside of +# nvidia-cutlass-dsl are permitted to ship subpackages under `cutlass.*`. +# They are listed below. For each one that is actually installed, we resolve where +# pip placed its `cutlass/` portion on disk and append that directory to +# cutlass.__path__, so that `import cutlass.` finds it via the standard +# import system. +# +# Anything NOT in this allowlist is silently ignored, even if it ships +# files at /cutlass//: the regular-package rule keeps it +# invisible. Adding a new approved sibling requires a release of +# nvidia-cutlass-dsl. +_APPROVED_CUTLASS_SIBLINGS = ("nvidia-cutlass-operators",) + + +def _extend_cutlass_path_for_approved_siblings() -> None: + import importlib.metadata as _md + from pathlib import Path + + for dist_name in _APPROVED_CUTLASS_SIBLINGS: + try: + dist = _md.distribution(dist_name) + except _md.PackageNotFoundError: + continue + + sibling_path = str(dist.locate_file("cutlass")) + if Path(sibling_path).is_dir() and sibling_path not in __path__: + __path__.append(sibling_path) + + +_extend_cutlass_path_for_approved_siblings() +del _extend_cutlass_path_for_approved_siblings +# --- end namespace extension gate -------------------------------------------- diff --git a/python/CuTeDSL/cutlass/_pth_hook.py b/python/CuTeDSL/cutlass/_pth_hook.py index 7f270d931..eba9dc4eb 100644 --- a/python/CuTeDSL/cutlass/_pth_hook.py +++ b/python/CuTeDSL/cutlass/_pth_hook.py @@ -13,7 +13,7 @@ This script sets up the editable install environment: 1. Sets CUTE_DSL_LIBS environment variable -2. Installs the custom editable finder for cutlass._mlir and DSL modules +2. Installs the custom editable finder for cutlass._mlir and legacy DSL modules The .pth file calls setup() with paths configured during installation. """ diff --git a/python/CuTeDSL/cutlass/address_space.py b/python/CuTeDSL/cutlass/address_space.py new file mode 100644 index 000000000..2639cefec --- /dev/null +++ b/python/CuTeDSL/cutlass/address_space.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +from cutlass.base_dsl.address_space import AddressSpace + +__all__ = ["AddressSpace"] diff --git a/python/CuTeDSL/cutlass/base_dsl/__init__.py b/python/CuTeDSL/cutlass/base_dsl/__init__.py index 15a02660c..234ee6de7 100644 --- a/python/CuTeDSL/cutlass/base_dsl/__init__.py +++ b/python/CuTeDSL/cutlass/base_dsl/__init__.py @@ -10,7 +10,9 @@ # is strictly prohibited. # Local module imports +from . import vector as vector from .dsl import * +from .pointer import * from .runtime import * from .._mlir_helpers import lru_cache_ir, dsl_user_op from .env_manager import get_str_env_var, detect_gpu_arch diff --git a/python/CuTeDSL/cutlass/base_dsl/address_space.py b/python/CuTeDSL/cutlass/base_dsl/address_space.py new file mode 100644 index 000000000..5738bd945 --- /dev/null +++ b/python/CuTeDSL/cutlass/base_dsl/address_space.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +from enum import IntEnum + + +class AddressSpace(IntEnum): + """Public CUTLASS address-space enum. + + Values match the CuTe MLIR dialect address-space encoding, but this enum is + intentionally defined in Python so public APIs do not depend on a dialect + binding object. + """ + + generic = 0 + gmem = 1 + smem = 3 + rmem = 5 + tmem = 6 + dsmem = 7 + cmem = 4 + def __str__(self) -> str: + return self.name diff --git a/python/CuTeDSL/cutlass/base_dsl/arch.py b/python/CuTeDSL/cutlass/base_dsl/arch.py index 174048634..bfdb67879 100644 --- a/python/CuTeDSL/cutlass/base_dsl/arch.py +++ b/python/CuTeDSL/cutlass/base_dsl/arch.py @@ -10,59 +10,12 @@ # is strictly prohibited. from collections.abc import Callable -from enum import Enum, EnumMeta -import re -from typing import Any +from enum import Enum + +from .version_info import CUDA_VERSION -class ArchMeta(EnumMeta): - """ - Custom metaclass for Arch enum that supports dynamic aliases based on CUDA version. - - - If cuda_version >= 13.0: sm_101/sm_101a/sm_101f are aliases of sm_110/sm_110a/sm_110f, use sm_110 as the canonical name - - Otherwise: sm_110/sm_110a/sm_110f are aliases of sm_101/sm_101a/sm_101f, use sm_101 as the canonical name - """ - - _arch_aliases: dict[str, str] = {} - - def __new__( - mcs, name: str, bases: tuple[type, ...], namespace: dict[str, Any] - ) -> "ArchMeta": - cls = super().__new__(mcs, name, bases, namespace) # type: ignore[arg-type] - from .version_info import CUDA_VERSION - - if CUDA_VERSION.major >= 13: - # sm_101 -> sm_110, use sm_110 as the canonical name - mcs._arch_aliases = { - "sm_101": "sm_110", - "sm_101a": "sm_110a", - "sm_101f": "sm_110f", - } - else: - # sm_110 -> sm_101, use sm_101 as the canonical name - mcs._arch_aliases = { - "sm_110": "sm_101", - "sm_110a": "sm_101a", - "sm_110f": "sm_101f", - } - return cls - - def __getattribute__(cls, name: str) -> Any: - # Use type.__getattribute__ to avoid recursion when accessing _arch_aliases - aliases = type.__getattribute__(cls, "_arch_aliases") - if name in aliases: - # Redirect to the target member - return type.__getattribute__(cls, aliases[name]) - return super().__getattribute__(name) - - def __getitem__(cls, name: str) -> "Arch": # type: ignore[override] - # Support Arch["sm_101"] style access - if name in cls._arch_aliases: - return super().__getitem__(cls._arch_aliases[name]) - return super().__getitem__(name) - - -class Arch(Enum, metaclass=ArchMeta): +class Arch(Enum): # sm_arch = (major, minor, suffix) # Ampere sm_80 = (8, 0, "") @@ -77,15 +30,23 @@ class Arch(Enum, metaclass=ArchMeta): sm_100 = (10, 0, "") sm_100a = (10, 0, "a") sm_100f = (10, 0, "f") - sm_101 = (10, 1, "") - sm_101a = (10, 1, "a") - sm_101f = (10, 1, "f") + if CUDA_VERSION.major >= 13: + sm_110 = (11, 0, "") + sm_110a = (11, 0, "a") + sm_110f = (11, 0, "f") + sm_101 = sm_110 + sm_101a = sm_110a + sm_101f = sm_110f + else: + sm_101 = (10, 1, "") # type: ignore[misc] + sm_101a = (10, 1, "a") # type: ignore[misc] + sm_101f = (10, 1, "f") # type: ignore[misc] + sm_110 = sm_101 # type: ignore[misc] + sm_110a = sm_101a # type: ignore[misc] + sm_110f = sm_101f # type: ignore[misc] sm_103 = (10, 3, "") sm_103a = (10, 3, "a") sm_103f = (10, 3, "f") - sm_110 = (11, 0, "") - sm_110a = (11, 0, "a") - sm_110f = (11, 0, "f") sm_120 = (12, 0, "") sm_120a = (12, 0, "a") sm_120f = (12, 0, "f") @@ -112,7 +73,7 @@ class Arch(Enum, metaclass=ArchMeta): @classmethod def BlackwellArchs(cls) -> tuple["Arch", ...]: - return ( + archs = ( Arch.sm_100, Arch.sm_100a, Arch.sm_100f, @@ -132,6 +93,7 @@ class Arch(Enum, metaclass=ArchMeta): Arch.sm_121a, Arch.sm_121f, ) + return tuple(dict.fromkeys(archs)) def __str__(self) -> str: return self.name @@ -168,7 +130,7 @@ class Arch(Enum, metaclass=ArchMeta): """ # sm_101 is renamed to sm_110, sm_101f is family of sm_110f, but is not family of sm_100f if self in [Arch.sm_101a, Arch.sm_101f]: - return arch.major == 11 and arch.minor >= 0 + return arch in [Arch.sm_101, Arch.sm_101a, Arch.sm_101f] return ( self.major == arch.major @@ -176,6 +138,28 @@ class Arch(Enum, metaclass=ArchMeta): and self.suffix in ["a", "f"] ) + def can_run_binary_built_for(self, target: "Arch") -> bool: + """Whether hardware of this (runtime) arch can execute a cubin built for ``target``. + + Arch-conditional targets (``a`` suffix, e.g. ``sm_100a``) run only on their exact + arch. Family-portable (``f``) and baseline (no suffix) targets run on any hardware + in the same major family whose minor is at least the target's minor (e.g. an + ``sm_100f`` binary runs on ``sm_100a`` and ``sm_103a``). + """ + if self == target: + return True + if self.major != target.major: + return False + # sm_101 is the renamed sm_110: a distinct family that only shares the (10,1) + # major.minor space with sm_10x when CUDA_VERSION.major < 13. Keep it from + # matching sm_100/sm_103 binaries in either direction (mirrors is_family_of). + sm_101_family = (Arch.sm_101, Arch.sm_101a, Arch.sm_101f) + if (self in sm_101_family) != (target in sm_101_family): + return False + if target.suffix == "a": + return self == target + return self.minor >= target.minor + def __lt__(self, other: object) -> bool: if not isinstance(other, Arch): return NotImplemented diff --git a/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py b/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py index 220278fa4..a12fd0153 100644 --- a/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py +++ b/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py @@ -28,6 +28,7 @@ import builtins from .utils.logger import log from .common import * +from .diagnostics import DiagId from .env_manager import get_str_env_var @@ -388,7 +389,7 @@ def range_dynamic(*args: Any, **kwargs: Any) -> None: raise DSLRuntimeError("range_dynamic should be always preprocessed to IR") -def range_constexpr(*args: Any) -> None: +def range_constexpr(*args: Any) -> range: raise DSLRuntimeError("range_constexpr should be preprocessed by preprocessor.") @@ -416,11 +417,9 @@ def const_expr(expression: Any) -> Any: failed = True if failed: - raise DSLRuntimeError( - f"The function `const_expr({expression})` received a dynamic expression (non compile-time constant).", - context={ - "If your expression depends on dynamic values": "Remove `const_expr()`", - }, + raise DSLUserCodeError( + DiagId.PHASE_REQUIRES_CONSTANT, + what="`const_expr()`", ) return expression @@ -455,17 +454,17 @@ def assert_executor(test: Any, msg: str | None = None) -> None: if not fail: assert test, msg else: - raise DSLRuntimeError( - "Only constexpr (Python Value) is allowed here, but got non-constexpr (IR Values) expression.", - suggestion="Please replace with runtime assert.", + raise DSLUserCodeError( + DiagId.PHASE_REQUIRES_CONSTANT, + what="`assert`", ) def bool_cast(value: Any) -> bool: if executor._is_dynamic_expression(value): # type: ignore[misc] - raise DSLRuntimeError( - "Only constexpr (Python Value) is allowed here, but got non-constexpr (IR Values) expression.", - suggestion="Please explicitly convert to boolean with expressions like comparison.", + raise DSLUserCodeError( + DiagId.PHASE_REQUIRES_CONSTANT, + what="Explicit boolean conversion", ) return bool(value) @@ -538,9 +537,8 @@ def range_value_check(*args: Any) -> tuple[int, int, int]: return (start, end, step) except: - raise DSLRuntimeError( - "`range_constexpr` requires constexpr (compile-time constant) for all arguments.", - suggestion="Use `range` instead of `range_constexpr`.", + raise DSLUserCodeError( + DiagId.ARG_NON_CONSTANT, ) @@ -572,9 +570,9 @@ def cf_symbol_check(symbol: Any) -> None: failed = True if failed: - raise DSLRuntimeError( - f"Incorrect `{name}` is used.", - suggestion=f"Please avoid overriding `{name}` from DSL package.", + raise DSLUserCodeError( + DiagId.CALL_WRONG_IMPORT, + name=name, ) @@ -589,8 +587,9 @@ def redirect_builtin_function(fcn: Any) -> Any: if isinstance(fcn, BuiltinFunctionType): if fcn in [builtins.exec, builtins.eval]: - raise DSLRuntimeError( - f"Built-in function `{fcn.__name__}` is not supported in DSL.", + raise DSLUserCodeError( + DiagId.UNSUP_BUILTIN, + name=fcn.__name__, ) if executor._builtin_redirector: return executor._builtin_redirector(fcn) @@ -660,9 +659,10 @@ def closure_check( if inspect.isfunction(value) or inspect.ismethod(value): closure_check([value], _visited) continue - raise DSLRuntimeError( - f"Function `{closure.__name__}` is a closure that captures variable `{name}` and is not supported in dynamic control flow", - suggestion="Please explicitly pass in captured variables as arguments", + raise DSLUserCodeError( + DiagId.SCOPE_CLOSURE_CAPTURE, + func_name=closure.__name__, + var_name=name, ) @@ -710,13 +710,13 @@ class FormattedValue: stacklevel=3, ) if self.format_spec is not None and len(self.format_spec) > 1: - raise DSLRuntimeError( - "At most one format specifier is supported for dynamic expressions", + raise DSLUserCodeError( + DiagId.UNSUP_FSTRING_FORMAT, ) if self.format_spec is not None: if self.format_spec[0].startswith(("<", ">", "^")): - raise DSLRuntimeError( - "Alignment specifier is not supported for dynamic expressions", + raise DSLUserCodeError( + DiagId.UNSUP_FSTRING_ALIGN, ) return f"%{self.format_spec[0]}", self.value return "{}", self.value @@ -754,7 +754,7 @@ def fstring_decompose( followed by all dynamic argument values. Raises: - DSLAstPreprocessorError: If an unsupported component type is encountered. + DSLUserCodeError: If an unsupported component type is encountered. """ format_string = "" dynamic_args = [] @@ -768,9 +768,7 @@ def fstring_decompose( elif isinstance(component, str): format_string += component else: - raise DSLAstPreprocessorError( - f"Unsupported component type in f-string: {type(component)}", - ) + raise DSLUserCodeError(DiagId.UNSUP_FSTRING) return (format_string, *dynamic_args) diff --git a/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py b/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py index 7ca91fda8..3c4e45856 100644 --- a/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py +++ b/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py @@ -49,6 +49,7 @@ from copy import deepcopy from itertools import chain from .common import * +from .diagnostics import DiagId from .utils.logger import log @@ -761,10 +762,29 @@ class DSLPreprocessor(ast.NodeTransformer): try: file_name = inspect.getsourcefile(function_pointer) or "" lines, start_line = inspect.getsourcelines(function_pointer) - dedented_source = textwrap.dedent("".join(lines)) + raw_source = "".join(lines) + dedented_source = textwrap.dedent(raw_source) tree = ast.parse(dedented_source, filename=file_name) # Bump the line numbers so they match the real source file ast.increment_lineno(tree, start_line - 1) + # ``textwrap.dedent`` stripped a constant leading-whitespace prefix, + # so ast column offsets are relative to the dedented source. Shift + # them back to the original file's columns so diagnostics underline + # the right place. The stripped width is identical on every non-blank + # line; derive it from the first line that still has content. + _col_shift = 0 + for _raw, _ded in zip(raw_source.split("\n"), dedented_source.split("\n")): + if _ded.strip(): + _col_shift = (len(_raw) - len(_raw.lstrip())) - ( + len(_ded) - len(_ded.lstrip()) + ) + break + if _col_shift: + for _node in ast.walk(tree): + if getattr(_node, "col_offset", None) is not None: + _node.col_offset += _col_shift # type: ignore[attr-defined] + if getattr(_node, "end_col_offset", None) is not None: + _node.end_col_offset += _col_shift # type: ignore[attr-defined] except Exception: # Under REPL mode, there is no way to get source of a function object, error out raise DSLRuntimeError( @@ -921,21 +941,20 @@ class DSLPreprocessor(ast.NodeTransformer): checker.generic_visit(tree) if not checker.has_early_exit: return - raise DSLAstPreprocessorError( - message=f"Early exit ({checker.early_exit_type}) is not allowed in `{self.session_data.function_name}`" - + ( - f" in `{self.session_data.class_name}`" - if self.session_data.class_name - else "" - ), + where = f"`{self.session_data.function_name}`" + ( + f" in `{self.session_data.class_name}`" + if self.session_data.class_name + else "" + ) + offender = checker.early_exit_node or tree + raise DSLUserCodeError( + DiagId.UNSUP_EARLY_EXIT, filename=self.session_data.file_name, - snippet=ast.unparse(tree), - suggestion=( - "If predicates are constant expression, write like " - "`if const_expr(...)` or `for ... in range_constexpr(...)`. " - "In that case, early exit will be executed by Python " - "interpreter, so it's supported." - ), + lineno=getattr(offender, "lineno", None), + col_offset=getattr(offender, "col_offset", None), + end_col_offset=getattr(offender, "end_col_offset", None), + kind=checker.early_exit_type, + where=where, ) def is_node_constexpr(self, node: ast.If | ast.While) -> bool: @@ -1132,9 +1151,12 @@ class DSLPreprocessor(ast.NodeTransformer): elif len(args) == 3: return self.visit(args[0]), self.visit(args[1]), self.visit(args[2]), True else: - raise DSLAstPreprocessorError( - "Unsupported number of arguments in range", + raise DSLUserCodeError( + DiagId.UNSUP_RANGE_ARGS, filename=self.session_data.file_name, + lineno=getattr(iter_node, "lineno", None), + col_offset=getattr(iter_node, "col_offset", None), + end_col_offset=getattr(iter_node, "end_col_offset", None), ) def extract_unroll_args(self, iter_node: ast.Call) -> tuple[ast.expr, ast.expr]: @@ -1198,6 +1220,7 @@ class DSLPreprocessor(ast.NodeTransformer): # Create the loop body transformed_body: list[ast.stmt] = [] + body_prep_stmts = self._prepare_loop_body_vars(node, write_args) with Region(self.session_data, new_value=transformed_body): for stmt in node.body: transformed_stmt = self.visit( @@ -1208,6 +1231,9 @@ class DSLPreprocessor(ast.NodeTransformer): else: transformed_body.append(transformed_stmt) + if body_prep_stmts: + transformed_body[:0] = body_prep_stmts + # Handle the return for a single iterated argument correctly if len(write_args) == 0: transformed_body.append(ast.Return()) @@ -1312,8 +1338,9 @@ class DSLPreprocessor(ast.NodeTransformer): ) self.session_data.import_top_module = True else: - # BoolOp should be either And or Or - raise DSLAstPreprocessorError( + # BoolOp should be either And or Or -- reaching here is a compiler + # bug (the AST grammar only produces And/Or), not an author mistake. + raise DSLRuntimeError( f"Unsupported boolean operation: {node.op}", filename=self.session_data.file_name, snippet=ast.unparse(node), @@ -1668,10 +1695,12 @@ class DSLPreprocessor(ast.NodeTransformer): # Check for early exit and raise exception self.check_early_exit(node, "for") if node.orelse: - raise DSLAstPreprocessorError( - "dynamic for loop with else is not supported", + raise DSLUserCodeError( + DiagId.UNSUP_LOOP_ELSE, filename=self.session_data.file_name, - snippet=ast.unparse(node), + lineno=node.lineno, + col_offset=node.col_offset, + end_col_offset=getattr(node.iter, "end_col_offset", None), ) # Get loop target variable name @@ -1860,10 +1889,10 @@ class DSLPreprocessor(ast.NodeTransformer): elif isinstance(component, ast.FormattedValue): elements.append(self.processFormattedValue(component)) else: - raise DSLAstPreprocessorError( - f"Unsupported component type in f-string: {type(component)}", + raise DSLUserCodeError( + DiagId.UNSUP_FSTRING, filename=self.session_data.file_name, - snippet=ast.unparse(component), + lineno=getattr(component, "lineno", None), ) call = ast.Call( func=_create_module_attribute( @@ -2045,7 +2074,13 @@ class DSLPreprocessor(ast.NodeTransformer): node, ) elif node.id == "_" and isLoad: - raise DSLAstPreprocessorError("Read '_' is not allowed") + raise DSLUserCodeError( + DiagId.UNSUP_READ_UNDERSCORE, + filename=self.session_data.file_name, + lineno=getattr(node, "lineno", None), + col_offset=getattr(node, "col_offset", None), + end_col_offset=getattr(node, "end_col_offset", None), + ) else: self.generic_visit(node) return node @@ -2127,9 +2162,10 @@ class DSLPreprocessor(ast.NodeTransformer): and dsl_decorator_index < len(decorator_list) - 1 ): decorator = ast.unparse(decorator_list[dsl_decorator_index]) - raise DSLAstPreprocessorError( - f"`{decorator}` decorator must be the inner most decorator", - suggestion=f"Please move the `{decorator}` decorator to the inner most position", + raise DSLUserCodeError( + DiagId.UNSUP_DECORATOR_ORDER, + filename=self.session_data.file_name, + decorator=decorator, ) return dsl_decorator_index is not None @@ -2158,9 +2194,12 @@ class DSLPreprocessor(ast.NodeTransformer): return new_decorator_list def visit_Global(self, node: ast.Global) -> None: - raise DSLAstPreprocessorError( - "`global` is not supported in DSL", - suggestion="Please explicitly pass in global variables as arguments", + raise DSLUserCodeError( + DiagId.UNSUP_GLOBAL, + filename=self.session_data.file_name, + lineno=getattr(node, "lineno", None), + col_offset=getattr(node, "col_offset", None), + end_col_offset=getattr(node, "end_col_offset", None), ) def visit_Nonlocal(self, node: ast.Nonlocal) -> ast.Nonlocal: @@ -2169,12 +2208,14 @@ class DSLPreprocessor(ast.NodeTransformer): intersect = nonlocal_names.intersections(active_symbols) for name in node.names: if name not in intersect: - raise DSLRuntimeError( - ( - f"`{ast.unparse(node)}` is referring to `{name}` which is not tracked by current JIT context, " - "this is not supported in DSL" - ), - suggestion="Please explicitly pass in nonlocal variables as arguments", + raise DSLUserCodeError( + DiagId.UNSUP_NONLOCAL, + filename=self.session_data.file_name, + lineno=getattr(node, "lineno", None), + col_offset=getattr(node, "col_offset", None), + end_col_offset=getattr(node, "end_col_offset", None), + stmt=ast.unparse(node), + name=name, ) self.generic_visit(node) return node @@ -2231,6 +2272,17 @@ class DSLPreprocessor(ast.NodeTransformer): if self.is_node_constexpr(node): return self._handle_constexpr_while(node) + # A staged ``while ... else:`` cannot capture its else-clause (parity with + # the staged ``for ... else`` rejection below). + if node.orelse: + raise DSLUserCodeError( + DiagId.UNSUP_LOOP_ELSE, + filename=self.session_data.file_name, + lineno=node.lineno, + col_offset=node.col_offset, + end_col_offset=getattr(node.test, "end_col_offset", None), + ) + active_symbols = self.session_data.scope_manager.get_active_symbols() active_callables = self.session_data.scope_manager.get_active_callables() @@ -2814,6 +2866,15 @@ class DSLPreprocessor(ast.NodeTransformer): """ return [] # No preparation needed in base class + def _prepare_loop_body_vars( + self, + node: "ast.For | ast.While", + write_args: list[str], + ) -> list[ast.stmt]: + """Prepend statements at the top of a for/while body. + """ + return [] + def create_while_function( self, func_name: str, @@ -2921,6 +2982,8 @@ class DSLPreprocessor(ast.NodeTransformer): test_expr, ) + body_prep_stmts = self._prepare_loop_body_vars(node, write_args) + # Section: while_after_block FunctionDef, which contains loop body while_after_stmts: list[ast.stmt] = [] with Region(self.session_data, new_value=while_after_stmts): @@ -2932,6 +2995,10 @@ class DSLPreprocessor(ast.NodeTransformer): while_after_stmts.extend(transformed_stmt) else: while_after_stmts.append(transformed_stmt) + + if body_prep_stmts: + while_after_stmts[:0] = body_prep_stmts + while_after_stmts.append(ast.Return(value=yield_args_ast_name_list)) while_after_block = ast.copy_location( diff --git a/python/CuTeDSL/cutlass/base_dsl/common.py b/python/CuTeDSL/cutlass/base_dsl/common.py index e3069e13b..774c19254 100644 --- a/python/CuTeDSL/cutlass/base_dsl/common.py +++ b/python/CuTeDSL/cutlass/base_dsl/common.py @@ -10,14 +10,22 @@ # is strictly prohibited. import inspect +import contextlib +import contextvars import os import subprocess import sys import types -from typing import Any, Dict, Optional, Union +from typing import Any, Dict, Generator, Optional, Union from functools import total_ordering from dataclasses import dataclass +from .diagnostics import ( + Colors, + find_user_source_location as _find_user_source_location, + render_user_diagnostic as _render_user_diagnostic, +) + """ This module provides a Exception classes DSL class for any Dialect. """ @@ -26,17 +34,36 @@ This module provides a Exception classes DSL class for any Dialect. # Store the original exception hook _original_excepthook = sys.excepthook -# Store registered environment manager (set by DSL singleton) -_registered_env_manager = None +_active_env_manager: contextvars.ContextVar[Any] = contextvars.ContextVar( + "active_env_manager", default=None +) -def register_env_manager(env_manager: Any) -> None: - """Register an EnvironmentVarManager instance for use by exception handling. +def get_current_env_manager() -> Any: + """Return the env manager for the active DSL context, if any. - Called by DSL singleton when it initializes. + Some shared runtime helpers do not have direct access to a ``BaseDSL`` + instance. Use the active context so those helpers observe the DSL that is + currently preprocessing/tracing. """ - global _registered_env_manager - _registered_env_manager = env_manager + return _active_env_manager.get() + + +@contextlib.contextmanager +def active_env_manager(env_manager: Any) -> Generator[None, None, None]: + """Temporarily make ``env_manager`` the current DSL env manager.""" + token = _active_env_manager.set(env_manager) + try: + yield + except Exception as e: + try: + if not hasattr(e, "_dsl_env_manager"): + setattr(e, "_dsl_env_manager", env_manager) + except (AttributeError, TypeError): + pass + raise + finally: + _active_env_manager.reset(token) @@ -52,14 +79,19 @@ def _dsl_excepthook( """ # Check if show_stacktrace is enabled via registered env manager show_stacktrace = False - if _registered_env_manager is not None: - show_stacktrace = getattr(_registered_env_manager, "show_stacktrace", False) + env_manager = getattr(exc_value, "_dsl_env_manager", None) + if env_manager is None: + env_manager = get_current_env_manager() + if env_manager is not None: + show_stacktrace = getattr(env_manager, "show_stacktrace", False) # Check if it's a DSL operation error (by name to avoid circular import issues) if exc_type.__name__ in ( "DSLOperationError", "DSLOperationBuildError", "DSLUserCodeError", + "CompilerDiagnosticError", + "DSLRuntimeError", ): if show_stacktrace: # Show full traceback in verbose mode @@ -77,18 +109,6 @@ def _dsl_excepthook( sys.excepthook = _dsl_excepthook -# Add color codes at the top of the file after imports -class Colors: - """ANSI color codes for error messages""" - - RED = "\033[91m" - YELLOW = "\033[93m" - BLUE = "\033[94m" - GREEN = "\033[92m" - BOLD = "\033[1m" - RESET = "\033[0m" - - # ============================================================================= # DSL Exceptions # ============================================================================= @@ -130,49 +150,22 @@ class DSLBaseError(Exception): return f"Caused exception: {self.cause}" return "" + # Subclasses may set this to True to render the "compiler bug, please + # report" envelope instead of the "here is your mistake" block. See + # ``_render_user_diagnostic``. + _is_internal: bool = False + def _format_message(self) -> str: + """Format via the single shared user-diagnostic renderer. + + Every DSL error renders through ``_render_user_diagnostic`` so the + output looks identical no matter which layer (AST pre-processing, + tracing, or runtime) raised it -- there is exactly one rendering + mechanism to change. Subclasses with genuinely different needs (e.g. + ``DSLOperationBuildError``, which walks an exception chain) may still + override this, but the default and recommended path is the shared one. """ - Formats the complete error message with available metadata. - Override this in subclasses if you want to change formatting logic. - """ - parts = [f"{self.__class__.__name__}: {self.message}"] - - if self.error_code is not None: - parts.append(f"{Colors.BOLD}Error Code:{Colors.RESET} {self.error_code}\n") - - if self.line is not None: - parts.append(f" Line: {self.line}") - - if self.filename is not None: - parts.append(f" File: {self.filename}") - - if self.snippet: - # Optionally truncate long snippets for readability - parts.append(f" Snippet: \n {self.snippet}") - - cause = self._generate_cause() - if cause: - parts.append(cause) - - if self.context: - if isinstance(self.context, dict): - parts.append(f"{Colors.BLUE}🔍 Additional Context:{Colors.RESET}\n") - for key, value in self.context.items(): - parts.append(f" {key}: {value}") - else: - parts.append( - f"{Colors.BLUE}🔍 Additional Context:{Colors.RESET} {self.context}" - ) - - if self.suggestion: - parts.append(f"{Colors.GREEN}💡 Suggestions:{Colors.RESET}") - if isinstance(self.suggestion, (list, tuple)): - for suggestion in self.suggestion: - parts.append(f" {Colors.GREEN}{suggestion}{Colors.RESET}") - else: - parts.append(f" {self.suggestion}") - - return "\n".join(parts) + return _render_user_diagnostic(self) class DSLSubprocessCallError(DSLBaseError): @@ -196,13 +189,29 @@ class DSLSubprocessCallError(DSLBaseError): class DSLRuntimeError(DSLBaseError): - """ - Raised when an error occurs during JIT-time code generation in the DSL. + """An internal / compiler error -- NOT the author's fault. + + Rendered as the "compiler bug, please report" envelope (see + ``_render_user_diagnostic``): a ``DSLRuntimeError`` means the DSL hit a + "should never happen" / failed-to-build-IR / wrapped-backend situation that + the kernel author cannot fix. For mistakes in the author's kernel raise + ``DSLUserCodeError`` with a ``DiagId`` instead, so the author gets a + "here is your mistake + how to fix it" message. """ - # Inherits all logic from DSLBaseError; override methods if you need - # specialized behavior or formatting for runtime errors. - pass + _is_internal = True + + +_ARCH_RELATED_CUDA_ERRORS = frozenset( + { + "CUDA_ERROR_INVALID_SOURCE", + "CUDA_ERROR_NO_BINARY_FOR_GPU", + "CUDA_ERROR_INVALID_PTX", + "CUDA_ERROR_UNSUPPORTED_PTX_VERSION", + "CUDA_ERROR_NO_DEVICE", + "CUDA_ERROR_INVALID_DEVICE", + } +) def _get_friendly_cuda_error_message( @@ -222,8 +231,9 @@ def _get_friendly_cuda_error_message( ): error_name = error_name[2:-1] - # Add target architecture info - target_arch = os.getenv("CUTE_DSL_ARCH", "unknown") + env_manager = get_current_env_manager() + target_arch = env_manager.arch if env_manager is not None else "unknown" + arch_is_relevant = error_name in _ARCH_RELATED_CUDA_ERRORS additional_info = { "CUDA_ERROR_INVALID_SOURCE": ( @@ -248,6 +258,10 @@ def _get_friendly_cuda_error_message( f"{Colors.RED}⚠️ Invalid parameter passed to CUDA operation.{Colors.RESET}\n\n" f"{Colors.YELLOW}This is likely a bug - please report it with:{Colors.RESET}" ), + "CUDA_ERROR_ILLEGAL_INSTRUCTION": ( + f"{Colors.RED}❌ A running kernel executed an illegal instruction.{Colors.RESET}\n\n" + f"{Colors.YELLOW}This is a fault inside the kernel, not an architecture mismatch.{Colors.RESET}\n\n" + ), "CUDA_ERROR_INVALID_CLUSTER_SIZE": ( f"{Colors.RED}❌ Invalid cluster size.{Colors.RESET}\n\n" ), @@ -266,6 +280,11 @@ def _get_friendly_cuda_error_message( "CUDA_ERROR_NO_BINARY_FOR_GPU": ( "Set env CUTE_DSL_ARCH to match your GPU architecture", ), + "CUDA_ERROR_ILLEGAL_INSTRUCTION": ( + "1. Inspect the kernel for an in-device fault (out-of-bounds access, " + "failed device-side assert, or an unsupported instruction at the launch config)", + "2. Re-run under compute-sanitizer to localize the faulting instruction", + ), "CUDA_ERROR_OUT_OF_MEMORY": ( "1. Reduce batch size", "2. Reduce model size", @@ -301,9 +320,8 @@ def _get_friendly_cuda_error_message( debug_info = f"\n- {Colors.BOLD}Error name: {error_name}\n" debug_info += f"- Error code: {error_code}\n" debug_info += f"- CUDA_TOOLKIT_PATH: {os.getenv('CUDA_TOOLKIT_PATH', 'not set')}\n" - debug_info += ( - f"- Target SM ARCH: {os.getenv('CUTE_DSL_ARCH', 'not set')}{Colors.RESET}\n" - ) + arch_source = "from env manager" if env_manager is not None else "unspecified" + debug_info += f"- Target SM ARCH: {target_arch} ({arch_source}){Colors.RESET}\n" try: # Get GPU information using CUDA Python API @@ -311,10 +329,21 @@ def _get_friendly_cuda_error_message( gpu_info = get_device_info() debug_info += gpu_info.pretty_str() - if target_arch and gpu_info.compatible_archs: + incompatible = bool( + target_arch + and gpu_info.compatible_archs + and target_arch not in gpu_info.compatible_archs + ) + + if ( + (arch_is_relevant or incompatible) + and target_arch + and target_arch != "unknown" + and gpu_info.compatible_archs + ): debug_info += f"\n{Colors.BOLD}Compatibility Check:{Colors.RESET}\n" - if target_arch not in gpu_info.compatible_archs: + if incompatible: debug_info += ( f"{Colors.RED}❌ Error: Target SM ARCH {target_arch} is not compatible\n" f"💡 Please use one of SM ARCHs: " @@ -337,17 +366,6 @@ def _get_friendly_cuda_error_message( return message, debug_info, error_suggestions.get(error_name, "") -def _get_cuda_error_name_from_code(error_code: int) -> Union[str, bytes]: - try: - # Avoid circular dependency. - from .runtime import cuda as cuda_helpers - - cu_result = cuda_helpers.cuda.CUresult(error_code) - return cuda_helpers._cudaGetErrorEnum(cu_result) - except (ValueError, AttributeError): - return f"" - - class DSLCudaRuntimeError(DSLBaseError): """ Raised when an error occurs during CUDA runtime code generation in the DSL. @@ -367,25 +385,70 @@ class DSLCudaRuntimeError(DSLBaseError): ) -def create_cuda_runtime_error( - error_code: int, cause: BaseException | None = None -) -> DSLCudaRuntimeError: - """Create a DSLCudaRuntimeError from a raw CUDA integer error code.""" - error = DSLCudaRuntimeError(error_code, _get_cuda_error_name_from_code(error_code)) - if cause is not None: - error.__cause__ = cause - error.__suppress_context__ = True - return error +class DSLWarning(UserWarning): + """A non-fatal author-facing warning, rendered like a ``DSLUserCodeError``. + The author's code is not wrong, only at risk (e.g. an implicit promotion + that can silently bite later), so this renders the shared diagnostic block + with a yellow ``[Warning]`` header instead of ``[Error]``. It subclasses + ``UserWarning`` so it flows through the standard ``warnings`` module + (filterable, deduplicated); prefer raising it via + ``diagnostics.report_warning(WarnId.X, ...)``. -class DSLAstPreprocessorError(DSLBaseError): - """ - Raised when an error occurs during AST preprocessing or visiting in the DSL. + The first positional argument is normally a :class:`WarnId` (the warnings + catalog -- separate from the :class:`DiagId` error catalog); a free-form + string is also accepted. """ - # Same approach: You could override _format_message if you want - # to emphasize AST node details or anything specific to preprocessing. - pass + _severity = "warning" + _is_internal = False + error_code = None + + def __init__( + self, + warn_or_message: Any, + filename: Optional[str] = None, + lineno: Optional[int] = None, + snippet: Optional[str] = None, + suggestion: Optional[Union[str, list]] = None, + context: Optional[Union[Dict[str, Any], str]] = None, + cause: Optional[BaseException] = None, + **fields: Any, + ) -> None: + from .diagnostics import WarnId # lazy: avoid an import cycle + + self.warn_id: Optional["WarnId"] = None + self.code: Optional[str] = None + if isinstance(warn_or_message, WarnId): + self.warn_id = warn_or_message + self.code = warn_or_message.code + message, catalog_fixes = warn_or_message.fill(**fields) + if suggestion is None: + suggestion = list(catalog_fixes) + else: + if fields: + raise TypeError( + "DSLWarning received template fields " + f"{sorted(fields)} but the first argument is a plain " + "string, not a WarnId." + ) + message = warn_or_message + + self.col: Optional[int] = None + self.end_col: Optional[int] = None + if filename is None and lineno is None: + filename, lineno, self.col, self.end_col = _find_user_source_location() + self.message = message + self.filename = filename + self.line = lineno + self.suggestion = suggestion + self.context = context + self.cause = cause + self.snippet = snippet + super().__init__(_render_user_diagnostic(self)) + + def _generate_cause(self) -> str: + return f"Caused exception: {self.cause}" if self.cause else "" class DSLNotImplemented(DSLBaseError): @@ -453,13 +516,27 @@ def translate_mlir_nanobind_error(exc: BaseException) -> str: class DSLUserCodeError(DSLBaseError): - """Raised when an error is detected in user DSL code. + """Raised when an error is detected in the author's kernel code. - Covers mutation violations, scope errors, type mismatches, and similar - user-facing diagnostics. Takes explicit ``filename`` and ``lineno`` -- - no ``inspect.stack()`` magic inside the class. + Covers mutation/phase violations, scope errors, type mismatches, + unsupported constructs, configuration mistakes, and similar author-facing + diagnostics. ``filename``/``lineno`` may be passed explicitly; when they + are not, the nearest author frame on the call stack is used so the error + still points at the user's code (see ``_find_user_source_location``). - Usage:: + The first positional argument is normally a :class:`DiagId` from the error + catalog; the catalog message and fixes are filled from the keyword + ``**fields`` and the stable code is appended automatically:: + + raise DSLUserCodeError( + DiagId.TYPE_UNSTABLE_JOIN, + filename="/path/to/user.py", + lineno=42, + var="accum", old_type="Int32", new_type="Float32", + ) + + A free-form string message is still accepted for one-off diagnostics that + do not yet have a catalog entry:: raise DSLUserCodeError( "Scope Error: variable `a` escapes its scope", @@ -471,17 +548,40 @@ class DSLUserCodeError(DSLBaseError): def __init__( self, - message: str, + diag_or_message: Any, filename: Optional[str] = None, lineno: Optional[int] = None, col_offset: Optional[int] = None, + end_col_offset: Optional[int] = None, cause: Optional[BaseException] = None, suggestion: Optional[Union[str, list]] = None, context: Optional[Union[Dict[str, Any], str]] = None, + snippet: Optional[str] = None, + **fields: Any, ) -> None: - snippet = None - if filename and lineno: - snippet = self._read_source_snippet(filename, lineno, col_offset) + from .diagnostics import DiagId # lazy: avoid an import cycle + + self.diag_id: Optional["DiagId"] = None + self.code: Optional[str] = None + if isinstance(diag_or_message, DiagId): + self.diag_id = diag_or_message + self.code = diag_or_message.code + message, catalog_fixes = diag_or_message.fill(**fields) + if suggestion is None: + suggestion = list(catalog_fixes) + else: + if fields: + raise TypeError( + "DSLUserCodeError received template fields " + f"{sorted(fields)} but the first argument is a plain " + "string, not a DiagId." + ) + message = diag_or_message + + self.col = col_offset + self.end_col = end_col_offset + if filename is None and lineno is None: + filename, lineno, self.col, self.end_col = _find_user_source_location() super().__init__( message, @@ -493,68 +593,6 @@ class DSLUserCodeError(DSLBaseError): context=context, ) - @staticmethod - def _read_source_snippet( - filename: str, - lineno: int, - col_offset: Optional[int] = None, - ) -> Optional[str]: - """Read a single source line and format it as a snippet.""" - try: - import linecache - - code_line = linecache.getline(filename, lineno).rstrip() - if not code_line: - return None - snippet = f" {lineno:4d} | {code_line}" - if col_offset is not None: - snippet += f"\n | {' ' * col_offset}^" - return snippet - except Exception: # noqa: BLE001 — best-effort snippet - return None - - def _format_message(self) -> str: - """Format a rich error message with code snippet and suggestions.""" - parts = [] - - parts.append( - f"\n{Colors.RED}{Colors.BOLD}[Error] {self.message}{Colors.RESET}\n" - ) - - if self.snippet and self.filename: - loc = f"{self.filename}:{self.line}" if self.line else self.filename - parts.append(f"{Colors.BLUE}Code:{Colors.RESET}") - parts.append(f"--> {Colors.BLUE}{loc}{Colors.RESET}") - parts.append(self.snippet) - parts.append("") - - if self.cause: - parts.append(f"{Colors.BLUE}Cause:{Colors.RESET} {self.cause}") - parts.append("") - - if self.context: - if isinstance(self.context, dict): - parts.append(f"{Colors.BLUE}Additional Context:{Colors.RESET}") - for key, value in self.context.items(): - parts.append(f" {key}: {value}") - else: - parts.append( - f"{Colors.BLUE}Additional Context:{Colors.RESET} {self.context}" - ) - parts.append("") - - if self.suggestion: - parts.append(f"{Colors.GREEN}Suggestion:{Colors.RESET}") - if isinstance(self.suggestion, (list, tuple)): - for s in self.suggestion: - parts.append(f" {Colors.GREEN}{s}{Colors.RESET}") - else: - parts.append(f" {self.suggestion}") - parts.append("") - - parts.append("=" * 100) - return "\n".join(parts) - class DSLOperationBuildError(DSLBaseError): """ @@ -745,8 +783,12 @@ def _coerce_to_cuda_version( return value if isinstance(value, str): return DSLCudaVersion(value) - raise DSLRuntimeError( - f"{param_name} must be a DSLCudaVersion or str, got {type(value).__name__}" + from .diagnostics import DiagId # lazy: avoid an import cycle + + raise DSLUserCodeError( + DiagId.CONFIG_VERSION_TYPE, + param_name=param_name, + got_type=type(value).__name__, ) @@ -797,6 +839,7 @@ def target_version( """ # Avoid circular dependency from .version_info import CUDA_VERSION + from .diagnostics import DiagId # Coerce all version parameters to DSLCudaVersion at the start exact_v = _coerce_to_cuda_version(exact_version, "exact_version") @@ -807,18 +850,14 @@ def target_version( is_range_check = min_v is not None or max_v is not None is_exact_version_check = exact_v is not None if is_range_check and is_exact_version_check: - raise DSLRuntimeError( - "Cannot use exact_version and [min_version, max_version] check at the same time" - ) + raise DSLUserCodeError(DiagId.CONFIG_VERSION_CONFLICT) if is_range_check: if min_v is None and max_v is None: - raise DSLRuntimeError( - "min_version and max_version cannot be None at the same time" - ) + raise DSLUserCodeError(DiagId.CONFIG_VERSION_MISSING) if min_v is not None and max_v is not None: if min_v > max_v: - raise DSLRuntimeError("min_version must be less than max_version") + raise DSLUserCodeError(DiagId.CONFIG_VERSION_RANGE_INVALID) result = True if min_v is not None: @@ -829,6 +868,4 @@ def target_version( elif is_exact_version_check: return CUDA_VERSION == exact_v else: - raise DSLRuntimeError( - "either exact_version, min_version, or max_version must be provided" - ) + raise DSLUserCodeError(DiagId.CONFIG_VERSION_REQUIRED) diff --git a/python/CuTeDSL/cutlass/base_dsl/compiler.py b/python/CuTeDSL/cutlass/base_dsl/compiler.py index 8f78e9df9..a0c4db3b9 100644 --- a/python/CuTeDSL/cutlass/base_dsl/compiler.py +++ b/python/CuTeDSL/cutlass/base_dsl/compiler.py @@ -17,12 +17,15 @@ and executes it using MLIR's ExecutionEngine. from typing import Any import collections.abc +import contextlib +import importlib.util import os import re import sys import inspect import types -from .common import DSLRuntimeError +from .common import DSLBaseError, DSLUserCodeError +from . import diagnostics as _diagnostics from .utils.logger import log from .env_manager import EnvironmentVarManager @@ -37,58 +40,32 @@ from .._mlir import ir # ============================================================================= -class CompilationError(RuntimeError): - """Custom error class for compilation failures""" - - # Add ANSI color codes - RED = "\033[91m" - YELLOW = "\033[93m" - BLUE = "\033[94m" - GREEN = "\033[92m" - BOLD = "\033[1m" - RESET = "\033[0m" +class CompilerDiagnosticError(DSLBaseError): + """Compilation failed with Python-rendered compiler diagnostics.""" def __init__( self, - message: str, - nvvm_error: str | None = None, - ir_context: str | None = None, - arch: str | None = None, + formatted: str, + raw_error: str = "", + *, + backend: str = "", + nvvm_error: str = "", + ir_context: str = "", + arch: str = "", ) -> None: + self.formatted = formatted + self.raw_error = raw_error + self.backend = backend self.nvvm_error = nvvm_error self.ir_context = ir_context self.arch = arch - # Call parent with formatted error to avoid showing class name - super().__init__("") # Empty string to avoid class name + super().__init__(formatted) - # Store formatted error for str() representation - self._formatted_error = self._format_error() + def _format_message(self) -> str: + return self.message def __str__(self) -> str: - """Override string representation to avoid showing class name""" - return self._formatted_error - - def __repr__(self) -> str: - """Override repr representation to avoid showing class name""" - return self._formatted_error - - def _format_error(self) -> str: - if not self.nvvm_error: - return str(self.args[0]) - - return f"""NVVM Compilation Error: ----------------------- - -{self.BLUE}⚙️ Current Settings:{self.RESET} -{self.BOLD}- Target Architecture: {self.arch}{self.RESET} - -IR Context (truncated): -{self.ir_context} - -{self.YELLOW}💡 Possible Solutions:{self.RESET} -{self.GREEN}1. Check if CUDA_TOOLKIT_PATH is set correctly -2. Verify target architecture ({self.arch}) is supported by your CUDA toolkit -3. Make sure CUDA toolkit version matches the target architecture requirements{self.RESET}""" + return self.formatted class Compiler: @@ -99,10 +76,11 @@ class Compiler: self.execution_engine = execution_engine self._post_compile_hook: collections.abc.Callable[[Any], None] | None = None - def _process_error(self, error_msg: str) -> tuple[str | None, str | None]: + def _process_error(self, error_msg: str) -> tuple[str | None, str | None, str]: """Process error message to extract NVVM error and IR context""" nvvm_error = None ir_msg = "" + location = _diagnostics.extract_compiler_location(error_msg) if "NVVM_ERROR" in error_msg: # Extract the specific NVVM error @@ -128,35 +106,69 @@ class Compiler: else: ir_msg = ir_section - return nvvm_error, ir_msg + return nvvm_error, ir_msg, location def compile( self, module: ir.Module, pipeline: str, arch: str = "", + remark_filter: str = "", + warnings_filter: str = "", + remark_output: str = "", + collect_compiler_diagnostics: bool = False, enable_debug_info: bool = False, enable_verifier: bool = False, ) -> None: """Compiles the module by invoking the pipeline.""" + diagnostic_session = _diagnostics.CompilerDiagnosticSession( + ir.Context.current, + remark_filter=remark_filter, + warnings_filter=warnings_filter, + remark_output=remark_output, + collect_diagnostics=collect_compiler_diagnostics, + ) try: pm = self.passmanager.PassManager.parse(pipeline) pm.enable_verifier(enable_verifier) - pm.run(module.operation) + # Enable remark streaming if configured + diagnostic_session.enable() + + with diagnostic_session.collecting(): + pm.run(module.operation) + formatted = diagnostic_session.format_success() + if formatted: + print(formatted, file=sys.stderr) except Exception as e: + if diagnostic_session.collect_diagnostics: + formatted = diagnostic_session.format_failure(str(e)) + if formatted: + raise CompilerDiagnosticError(formatted, raw_error=str(e)) from e error_msg = str(e) - nvvm_error, ir_msg = self._process_error(error_msg) + nvvm_error, ir_msg, location = self._process_error(error_msg) if nvvm_error: - raise CompilationError( - error_msg, + ir_context = ir_msg or "" + formatted = diagnostic_session.format_backend_failure( + raw_error=error_msg, nvvm_error=nvvm_error, - ir_context=ir_msg, + ir_context=ir_context, + arch=arch, + location=location, + ) + raise CompilerDiagnosticError( + formatted, + raw_error=error_msg, + backend="nvvm", + nvvm_error=nvvm_error, + ir_context=ir_context, arch=arch, ) from e raise e finally: + # Finalize remark output after passes complete + diagnostic_session.finalize() pass if self._post_compile_hook: @@ -181,6 +193,10 @@ class Compiler: opt_level: int = 2, arch: str = "", enable_debug_info: bool = False, + remark_filter: str = "", + warnings_filter: str = "", + remark_output: str = "", + collect_compiler_diagnostics: bool = False, enable_verifier: bool = False, ) -> Any: """Compiles and jits the module.""" @@ -188,6 +204,10 @@ class Compiler: module, pipeline, arch, + remark_filter=remark_filter, + warnings_filter=warnings_filter, + remark_output=remark_output, + collect_compiler_diagnostics=collect_compiler_diagnostics, enable_debug_info=enable_debug_info, enable_verifier=enable_verifier, ) @@ -253,7 +273,7 @@ class BooleanCompileOption(CompileOption): class StringCompileOption(CompileOption): def serialize(self) -> str: - if self._value: + if self._value and self.__class__.option_name: self._value = self._value.strip("'") return f"{self.__class__.option_name}='{self._value}'" return "" @@ -273,7 +293,7 @@ class BooleanBasedFileDumpOption(CompileOption): self._dump_path = path def serialize(self) -> str: - if self._value: + if self._value and self.__class__.option_name: assert self._dump_path, ( f"Dump path is not set for {self.__class__.__name__}" ) @@ -291,14 +311,13 @@ class OptLevel(CompileOption): def __init__(self, val: int) -> None: if val < 0 or val > 3: - raise DSLRuntimeError(f"Invalid OPT_LEVEL: {val}, valid range is [0, 3]") + raise DSLUserCodeError( + _diagnostics.DiagId.CONFIG_INVALID_OPT_LEVEL, val=val + ) super().__init__(val) -_NVVM_DIAG_SUB_ALIASES: dict[str, str] = { -} - _SR_DISABLE_PATTERNS: dict[str, str] = { } @@ -317,6 +336,35 @@ class ExtraCompilerOpts(CompileOption): return self._value +def _ensure_ptxas_verbose(options: str) -> str: + import shlex + + stripped = (options or "").strip() + if len(stripped) >= 2 and stripped[0] == stripped[-1] and stripped[0] in "'\"": + stripped = stripped[1:-1] + try: + tokens = shlex.split(stripped) if stripped else [] + except ValueError: + tokens = stripped.split() + if any(token in ("-v", "--verbose") for token in tokens): + return stripped + return f"{stripped} -v".strip() + + +def _ensure_compiler_diagnostic_selector(options: str, selector: str) -> str: + match = re.search(r"(^|\s)diagnostic=([^\s]+)", options) + if not match: + return f"{options} diagnostic={selector}".strip() + + selectors = {item for item in match.group(2).split(",") if item} + if selector in selectors: + return options + + selectors.add(selector) + replacement = f"{match.group(1)}diagnostic={','.join(sorted(selectors))}" + return options[: match.start()] + replacement + options[match.end() :] + + class PtxasOptions(StringCompileOption): option_name = "ptx-options" @@ -356,6 +404,16 @@ class KeepPTX(BooleanBasedFileDumpOption): +class KeepSASS(BooleanBasedFileDumpOption): + pass + + + +class NvdisasmOptions(StringCompileOption): + def __init__(self, val: str = "-g -c") -> None: + super().__init__(val) + + class FlattenLocsOutputJsonPath(StringCompileOption): """Path to the FlattenLocs JSON sidecar (vloc id -> original source chain). @@ -520,7 +578,9 @@ class HostTarget(EmptyCompileOption): try: self._triple, self._cpu, self._features = _parse_host_target(val) except ValueError as exc: - raise DSLRuntimeError(str(exc)) from exc + raise DSLUserCodeError( + _diagnostics.DiagId.CONFIG_INVALID_HOST_TARGET, error=str(exc) + ) from exc @property def value(self) -> str: @@ -547,6 +607,48 @@ class HostTarget(EmptyCompileOption): return self._features +class RemarkFilter(EmptyCompileOption): + """Regex filter for remark categories e.g. 'Memory|Algorithm' or '.*' for all. + This option is not serialized into the pipeline string; it configures + the MLIR context remark engine directly.""" + + option_name = "remark-filter" + + def __init__(self, val: str = "") -> None: + super().__init__(val) + + +class WarningsFilter(EmptyCompileOption): + """Checker domains whose WARNINGS the DSL displays, e.g. 'nvvm' or 'nvvm,ptx'. + + Set by the user-facing ``--warnings{}`` compile option. Errors are + always shown; warnings are shown only for domains listed here. Not + serialized into the pipeline string; consumed by the Python diagnostic + renderer to gate warning visibility.""" + + option_name = "warnings-filter" + + def __init__(self, val: str = "") -> None: + super().__init__(val) + + +class CollectCompilerDiagnostics(BooleanCompileOption): + option_name = "collect-compiler-diagnostics" + + def serialize(self) -> str: + return "" + + +class RemarkOutput(EmptyCompileOption): + """Output file path for YAML remark format. + This option is not serialized into the pipeline string; it configures + the MLIR context remark engine directly.""" + + option_name = "remark-output" + + def __init__(self, val: str = "") -> None: + super().__init__(val) + class CompileOptions: """ @@ -570,6 +672,8 @@ class CompileOptions: GenerateLineInfo: GenerateLineInfo(False), KeepCUBIN: KeepCUBIN(False), KeepPTX: KeepPTX(False), + KeepSASS: KeepSASS(False), + NvdisasmOptions: NvdisasmOptions(), GPUArch: GPUArch(""), FlattenLocsOutputJsonPath: FlattenLocsOutputJsonPath(""), LinkLibraries: LinkLibraries(""), @@ -577,7 +681,13 @@ class CompileOptions: DeviceTarget: DeviceTarget(False), DumpDir: DumpDir(""), HostTarget: HostTarget(""), + # Remark options + RemarkFilter: RemarkFilter(""), + WarningsFilter: WarningsFilter(""), + RemarkOutput: RemarkOutput(""), + CollectCompilerDiagnostics: CollectCompilerDiagnostics(False), } + self._ptxas_diagnostics_enabled = False if options is not None: self._update(options) @@ -585,7 +695,9 @@ class CompileOptions: def _update(self, options: "CompileOption | tuple[CompileOption, ...]") -> None: def _validate_and_update_option(option: CompileOption) -> None: if type(option) not in self.options: - raise DSLRuntimeError(f"Invalid compile option: {option}") + raise DSLUserCodeError( + _diagnostics.DiagId.CONFIG_UNKNOWN_COMPILE_OPTION, option=option + ) self.options[type(option)] = option if isinstance(options, tuple): @@ -601,6 +713,13 @@ class CompileOptions: object's options accordingly. Valid forms (comma- or space-separated, ``--`` prefix optional):: + # Errors are always shown and fail compilation -- no flag needed. + # --warnings / --remarks are opt-in and non-fatal; a {} + # selector shows only that category, bare shows all categories. + --warnings # show all warnings + --warnings{nvvm} # show only nvvm-category warnings + --remarks # show all remarks + --remarks{nvvm} # show only nvvm (sync) remarks --iket # enable IKET (In-Kernel Event Tracing) instrumentation :param opt_str: Compact option string to parse. @@ -664,8 +783,32 @@ class CompileOptions: f"or use the bare name to enable a boolean option" ) - if name == "nvvm-diag": - pass + if name in ("warnings", "remarks"): + if val_str is not None: + raise ValueError( + f"{name} expects selector braces, e.g. {name}{{nvvm}}" + ) + valid_selectors = { + "nvvm", + } + selectors = ( + set(valid_selectors) + if sub_str is None + else {item.strip() for item in sub_str.split(",") if item.strip()} + ) + unknown = selectors - valid_selectors + if unknown: + valid_list = ", ".join(sorted(valid_selectors)) + unknown_list = ", ".join(sorted(unknown)) + raise ValueError( + f"{name} supports selectors {{{valid_list}}}; " + f"unknown selector(s): {unknown_list}" + ) + # --warnings{} / --remarks{} enable the C++ collection passes via + # the internal `diagnostic=` selector. Errors are always-on + # regardless; these flags only control which severity the DSL + # displays (warnings vs remarks). + raw_opts.append(f"diagnostic={','.join(sorted(selectors))}") elif name in _PERF_ENABLE_OPTIONS and sub_str is None: # Emit the MLIR flag for this pass enable token directly. raw_opts.append(f"{_PERF_ENABLE_OPTIONS[name]}=true") @@ -688,8 +831,9 @@ class CompileOptions: ) else: if not val: - raise DSLRuntimeError( - f"Option '{key}' requires a value (e.g. {key}=)" + raise DSLUserCodeError( + _diagnostics.DiagId.CONFIG_OPTION_REQUIRES_VALUE, + key=key, ) opt.value = val else: @@ -708,8 +852,12 @@ class CompileOptions: self.options[KeepPTX].value = True if envar.keep_cubin: self.options[KeepCUBIN].value = True + if envar.keep_sass: + self.options[KeepSASS].value = True if envar.compiler_opt: self._apply_opt_string(envar.compiler_opt) + if envar.debug and "diagnostic=" not in self.options[ExtraCompilerOpts].value: + self._apply_opt_string("warnings{nvvm}") if envar.enable_assertions: self.options[EnableAssertions].value = True if envar.lineinfo: @@ -749,6 +897,31 @@ class CompileOptions: dump_dir, f"{function_name}.{arch}.cubin", ) + keep_sass = self.options[KeepSASS] + assert isinstance(keep_sass, KeepSASS) + if keep_sass.value: + assert dump_dir is not None + keep_sass.dump_path = os.path.join( + dump_dir, + f"{function_name}.sass", + ) + _need_cubin_on_disk = True + if _need_cubin_on_disk and not keep_cubin.value: + keep_cubin.value = True + keep_cubin.dump_path = os.path.join(dump_dir, f"{function_name}") + keep_cubin.full_cubin_path = os.path.join( + dump_dir, + f"{function_name}.{arch}.cubin", + ) + if envar.remarks and not self.options[RemarkFilter].value: + self.options[RemarkFilter].value = envar.remarks + if envar.remarks and not self.options[RemarkOutput].value: + assert dump_dir is not None + self.options[RemarkOutput].value = os.path.join( + dump_dir, + f"{function_name}_remarks.yaml", + ) + @property def generate_line_info(self) -> bool: return self.options[GenerateLineInfo].value @@ -791,16 +964,43 @@ class CompileOptions: assert isinstance(keep_cubin, KeepCUBIN) return keep_cubin.full_cubin_path if keep_cubin.value else None + @property + def dump_sass_path(self) -> str | None: + keep_sass = self.options[KeepSASS] + assert isinstance(keep_sass, KeepSASS) + return keep_sass.dump_path if keep_sass.value else None + + @property + def remark_filter(self) -> str: + print_remark_all = ".*" + return self.options[RemarkFilter].value or print_remark_all + + @property + def warnings_filter(self) -> str: + """Comma-separated checker domains whose warnings the DSL shows + (set by --warnings{}). Empty means no warnings are displayed.""" + return self.options[WarningsFilter].value or "" + + @property + def remark_output(self) -> str: + return self.options[RemarkOutput].value or "" + + @property + def remarks_enabled(self) -> bool: + return self.remark_output != "" + + @property + def collect_compiler_diagnostics(self) -> bool: + # Collection (and the nvvm checker) is enabled by --warnings{} / + # --remarks{}. When set, the checker runs and renders its diagnostics; + # errors among them are always fatal + shown (no separate error flag). + return bool(self.options[CollectCompilerDiagnostics].value) + @property def enable_tvm_ffi(self) -> bool: ret = self.options[EnableTVMFFI].value - if ret: - try: - import tvm_ffi - except ModuleNotFoundError: - raise DSLRuntimeError( - "TVM FFI is not installed, please install it via `pip install apache-tvm-ffi`" - ) + if ret and importlib.util.find_spec("tvm_ffi") is None: + raise DSLUserCodeError(_diagnostics.DiagId.CONFIG_MISSING_TVM_FFI) return ret def to_str(self) -> str: @@ -808,6 +1008,7 @@ class CompileOptions: Generate a string representation of all compilation options which will be used in pipeline options. """ + self._finalize_derived_options() flattend_options = "" for option in self.options.values(): flattend_options += option.serialize() + " " @@ -815,12 +1016,28 @@ class CompileOptions: log().info("`cute.compile` CompileOptions: options=" + flattend_options) return flattend_options + def _finalize_derived_options(self) -> None: + if not self._ptxas_diagnostics_enabled: + return + self.options[CollectCompilerDiagnostics].value = True + extra_options = self.options[ExtraCompilerOpts] + assert isinstance(extra_options, ExtraCompilerOpts) + extra_options.value = _ensure_compiler_diagnostic_selector( + extra_options.value, "ptx" + ) + ptxas_options = self.options[PtxasOptions] + assert isinstance(ptxas_options, PtxasOptions) + ptxas_options.value = _ensure_ptxas_verbose(ptxas_options.value) + def _extract_compact_options( options: str, ) -> "tuple[CompileOptions | None, str]": """Split *options* into compact tokens and legacy tokens. + Compact tokens such as ``diagnostic`` and ``iket`` are applied immediately + via ``CompileOptions._apply_opt_string``; the remaining legacy tokens are + returned as a plain string for argparse. Returns: (base_compile_options or None, remaining legacy options string). When the input is *pure* compact, the fully-configured CompileOptions @@ -830,6 +1047,8 @@ def _extract_compact_options( _COMPACT_NAMES: frozenset[str] = frozenset( { + "warnings", + "remarks", "iket", } | set(_PERF_ENABLE_OPTIONS) @@ -839,6 +1058,7 @@ def _extract_compact_options( bare = ( (token[2:] if token.startswith("--") else token).split(",")[0].split("{")[0] ) + bare = bare.split("=")[0] return bare in _COMPACT_NAMES stripped = options.strip() if options else "" @@ -856,6 +1076,7 @@ def _extract_compact_options( _LEGACY_VALUE_KEYS: frozenset[str] = frozenset( { + "--nvdisasm-options", "--ptxas-options", "--link-libraries", "--gpu-arch", @@ -888,7 +1109,7 @@ def _extract_compact_options( base = CompileOptions() base._apply_opt_string(compact_str) - return base, " ".join(legacy_tokens) + return base, " ".join(shlex.quote(token) for token in legacy_tokens) # This is a temp function to preserve backward compatibility. @@ -911,6 +1132,8 @@ def _parse_compile_options_from_str(options: str) -> CompileOptions: "generate_line_info": GenerateLineInfo, "keep_cubin": KeepCUBIN, "keep_ptx": KeepPTX, + "keep_sass": KeepSASS, + "nvdisasm_options": NvdisasmOptions, "gpu_arch": GPUArch, "enable_tvm_ffi": EnableTVMFFI, "dump_dir": DumpDir, @@ -927,6 +1150,8 @@ def _parse_compile_options_from_str(options: str) -> CompileOptions: parser.add_argument("--generate-line-info", action="store_true", default=False) parser.add_argument("--keep-cubin", action="store_true", default=False) parser.add_argument("--keep-ptx", action="store_true", default=False) + parser.add_argument("--keep-sass", action="store_true", default=False) + parser.add_argument("--nvdisasm-options", type=str, default="-g -c") parser.add_argument("--ptxas-options", type=str, default="") parser.add_argument("--gpu-arch", type=str, default="") parser.add_argument("--enable-tvm-ffi", action="store_true", default=False) @@ -948,8 +1173,8 @@ def _parse_compile_options_from_str(options: str) -> CompileOptions: compile_options.options[option_cls].value = value except SystemExit as e: # catch argparse error and raise as DSLRuntimeError - raise DSLRuntimeError( - f"Invalid compile options: '{options}'. Please check the option values and format." + raise DSLUserCodeError( + _diagnostics.DiagId.CONFIG_MALFORMED_COMPILE_OPTIONS, options=options ) from e return compile_options @@ -984,6 +1209,10 @@ class CompileCallable: ``CUTE_DSL_COMPILER_OPT``. Keep this docstring focused on the compile contract; use ``write-kernel/references/compiler-options.md`` as the authoritative catalog of option tokens and examples. + + ``cute.compile(..., trace_finalize_hooks=hook_or_hooks)`` temporarily + registers callbacks for that compile only. Hooks run after tracing and + before module hashing, and are removed even if compilation fails. """ def __init__(self, options: Any = None) -> None: @@ -1076,14 +1305,15 @@ class CompileCallable: decorated with ``@cute.jit``. """ if func is None: - raise DSLRuntimeError("Function is not set or invalid.") + raise DSLUserCodeError(_diagnostics.DiagId.CALL_FUNCTION_NOT_PROVIDED) if not callable(func): - raise DSLRuntimeError("Object is not callable.") + raise DSLUserCodeError(_diagnostics.DiagId.CALL_NOT_CALLABLE) # Pop the migration-aid kwarg before it leaks into the executor # call: the rest of the pipeline does not know about it. is_experimental_requested = kwargs.pop("is_experimental", False) + finalize_hook = kwargs.pop("trace_finalize_hooks", None) kwargs["compile_only"] = True kwargs["no_cache"] = True @@ -1105,10 +1335,7 @@ class CompileCallable: # Get the actual function from the class definition func = func.__call__.__func__ else: - raise DSLRuntimeError( - "Invalid function type, only function, method and module are supported, but got", - func, - ) + raise DSLUserCodeError(_diagnostics.DiagId.CALL_UNSUPPORTED_CALLABLE_TYPE) func_name_prefix = getattr(func, "_name_prefix", None) if func_name_prefix: @@ -1123,9 +1350,7 @@ class CompileCallable: BaseDSL._lazy_initialize_dsl(func) if not hasattr(func, "_dsl_object"): - raise DSLRuntimeError( - f"Function {func} is not decorated with jit decorator." - ) + raise DSLUserCodeError(_diagnostics.DiagId.CALL_MISSING_JIT_DECORATOR) # Validate the migration-aid ``is_experimental`` kwarg against # the routing already baked into the function by its jit/kernel @@ -1140,14 +1365,7 @@ class CompileCallable: if is_experimental_requested and not getattr( func._dsl_object, "_is_experimental_dsl", False ): - raise DSLRuntimeError( - "cute.compile(is_experimental=True) was called on a function " - f"routed through {type(func._dsl_object).__name__}, which is " - "not an experimental DSL. Decorate the function with " - "@cute.jit(is_experimental=True) (and any nested @cute.kernel " - "with is_experimental=True) before compiling, or drop the " - "is_experimental flag from the cute.compile call." - ) + raise DSLUserCodeError(_diagnostics.DiagId.CALL_EXPERIMENTAL_MISMATCH) # process compile options, extract the options and remove them from the kwargs options = kwargs.pop("options", None) @@ -1160,17 +1378,23 @@ class CompileCallable: compile_options = self._compile_options func._dsl_object.compile_options = compile_options - # Preprocess the function if not already preprocessed - func._dsl_object._preprocess_and_replace_code(func) + if finalize_hook is None: + hook_context = contextlib.nullcontext() + else: + hook_context = func._dsl_object.trace_finalize_hooks(finalize_hook) - # Route based on DeviceTarget option: compiles as __device__ function. - if compile_options.options[DeviceTarget].value: - # Device functions are always relocatable objects. - compile_options.options[RDC].value = True - # Force artifact dumping so .o and .ptx are available after compilation. - compile_options.options[KeepPTX].value = True - compile_options.options[KeepCUBIN].value = True - return func._dsl_object._device_func(func, *args, **kwargs) + with hook_context: + # Preprocess the function if not already preprocessed + func._dsl_object._preprocess_and_replace_code(func) - # Default: host wrapper + kernel - return func._dsl_object._func(func, *args, **kwargs) + # Route based on DeviceTarget option: compiles as __device__ function. + if compile_options.options[DeviceTarget].value: + # Device functions are always relocatable objects. + compile_options.options[RDC].value = True + # Force artifact dumping so .o and .ptx are available after compilation. + compile_options.options[KeepPTX].value = True + compile_options.options[KeepCUBIN].value = True + return func._dsl_object._device_func(func, *args, **kwargs) + + # Default: host wrapper + kernel + return func._dsl_object._func(func, *args, **kwargs) diff --git a/python/CuTeDSL/cutlass/base_dsl/diagnostics.py b/python/CuTeDSL/cutlass/base_dsl/diagnostics.py new file mode 100644 index 000000000..3051a1e4e --- /dev/null +++ b/python/CuTeDSL/cutlass/base_dsl/diagnostics.py @@ -0,0 +1,1857 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +"""User-facing DSL error catalog. + +One entry per author-facing mistake. The enum **member name** is the stable +error code (e.g. ``TYPE_UNSTABLE_JOIN``) -- there are no hand-written numbers to +keep in sync. The member **value** is ``(message_template, (fix_line, ...))``, +written in the author's words, never in compiler/IR terms:: + + raise DSLUserCodeError(DiagId.TYPE_UNSTABLE_JOIN, filename=fn, lineno=ln, + var="count", old_type="Int32", new_type="Float32") + +The exception consumes the ``DiagId`` directly (see ``DiagId.format``) and +renders through ``DSLUserCodeError`` (snippet + suggestion box), so an error +raised while tracing looks the same as one raised while reading code. + +Auto-doc readiness (the generator is intentionally *not* built yet): every entry +exposes ``.code`` (the name), ``.category`` (the name's prefix), ``.message`` and +``.fix``. A future doc generator just does:: + + for d in DiagId: # group by d.category, emit d.code/ + ... # d.message/d.fix + +so no separate metadata table is needed. +""" + +import enum +import linecache +import re +import sys +import textwrap +import types +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, Sequence + +__all__ = [ + "Colors", + "CompilerDiagnostic", + "CompilerDiagnosticSession", + "DiagId", + "WarnId", + "report_warning", + "META_VALUE", + "STAGED_VALUE", +] + +# Canonical author-facing names for the two value phases. Defined once here and +# auto-injected into every message as {meta}/{staged} (see DiagId.format), so the +# wording is consistent and can be changed in a single place. +META_VALUE = "Python value (Meta value)" +STAGED_VALUE = "Runtime value (Staged value)" + + +class Colors: + """Shared ANSI color codes for DSL and compiler diagnostics.""" + + RED = "\033[91m" + YELLOW = "\033[93m" + BLUE = "\033[94m" + GREEN = "\033[92m" + CYAN = "\033[96m" + BOLD = "\033[1m" + RESET = "\033[0m" + + +_Colors = Colors + + +@dataclass(frozen=True) +class CompilerDiagnostic: + severity: str + message: str + code: str = "" + namespace: str = "" + name: str = "" + location: str = "" + reason: str = "" + suggestion: str = "" + ptx_ref: str = "" + ptx_url: str = "" + notes: tuple[str, ...] = () + + +@dataclass(frozen=True) +class _CompilerDiagnosticInfo: + code: str + namespace: str = "" + ptx_ref: str = "" + ptx_url: str = "" + + +_PY_LOC_RE = re.compile(r'"(?P[^"]+?\.py)":(?P\d+):(?P\d+)') + +_PTX_DOC_BASE = "https://docs.nvidia.com/cuda/parallel-thread-execution/index.html" +_COMPILER_DIAG_TEXT_WIDTH = 100 +_COMPILER_CONTEXT_LINES = 2 +_DEVICE_BINARY_SERIALIZATION_DIAGNOSTIC = "DeviceBinarySerializationFailed" + + +def _nvvm_info(code: str, ptx_ref: str, ptx_anchor: str) -> _CompilerDiagnosticInfo: + return _CompilerDiagnosticInfo( + code=code, + namespace="nvvm-diag", + ptx_ref=ptx_ref, + ptx_url=_PTX_DOC_BASE + ptx_anchor, + ) + + +_COMPILER_DIAGNOSTICS: dict[str, _CompilerDiagnosticInfo] = { + "NvvmDiagMbarrierArriveNeedsElect": _nvvm_info( + "C3/C4", + "mbarrier arrive and elect.sync", + "#parallel-synchronization-and-communication-instructions-mbarrier-arrive", + ), + "NvvmDiagMbarrierCountMismatch": _nvvm_info( + "C3", + "mbarrier.init", + "#parallel-synchronization-and-communication-instructions-mbarrier-init", + ), + "NvvmDiagMbarrierCountMismatchWarning": _nvvm_info( + "C3", + "mbarrier.init", + "#parallel-synchronization-and-communication-instructions-mbarrier-init", + ), + "NvvmDiagExpectTxNeedsCompletion": _nvvm_info( + "C5", + "mbarrier.arrive.expect_tx", + "#parallel-synchronization-and-communication-instructions-" + "mbarrier-expect-tx-operation", + ), + "NvvmDiagCta2TmaMulticastMayHang": _nvvm_info( + "C5", + "cp.async.bulk.tensor cta_group::2 completion", + "#data-movement-and-conversion-instructions-cp-async-bulk-tensor", + ), + "NvvmDiagTcgen05LdNeedsCommit": _nvvm_info( + "C7", + "tcgen05 memory consistency", + "#tcgen05-memory-consistency-model-canonical-sync-patterns", + ), + "NvvmDiagTcgen05LdMissingAfterFence": _nvvm_info( + "C7", + "tcgen05 memory consistency", + "#tcgen05-memory-consistency-model-canonical-sync-patterns", + ), + "NvvmDiagTcgen05CollectiveNeedsFullWarp": _nvvm_info( + "C10", + "tcgen05 relinquish_alloc_permit", + "#tcgen05-instructions-tcgen05-alloc-dealloc-relinquish-alloc-permit", + ), + "NvvmDiagTcgen05CommitNeedsElect": _nvvm_info( + "C11", + "tcgen05 commit", + "#tcgen-async-sync-operations-commit", + ), + "NvvmDiagPartialWarpElectSync": _nvvm_info( + "C13", + "elect.sync", + "#parallel-synchronization-and-communication-instructions-elect-sync", + ), + _DEVICE_BINARY_SERIALIZATION_DIAGNOSTIC: _CompilerDiagnosticInfo(""), +} + +_GENERIC_MLIR_DIAGNOSTIC = "MLIRDiagnostic" + + +class CompilerDiagnosticCollector: + """Scoped collector for regular MLIR diagnostics emitted during pass runs.""" + + def __init__(self, context: Any, enabled: bool) -> None: + self.context = context + self.enabled = enabled + self.diagnostics: list[CompilerDiagnostic] = [] + self._handler: Any | None = None + self._previous_emit_error_diagnostics: bool | None = None + + def __enter__(self) -> "CompilerDiagnosticCollector": + if not self.enabled or not hasattr(self.context, "attach_diagnostic_handler"): + return self + if hasattr(self.context, "emit_error_diagnostics"): + self._previous_emit_error_diagnostics = bool( + self.context.emit_error_diagnostics + ) + self.context.emit_error_diagnostics = True + self._handler = self.context.attach_diagnostic_handler(self._handle) + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: types.TracebackType | None, + ) -> None: + if self._handler is not None: + self._handler.detach() + self._handler = None + if self._previous_emit_error_diagnostics is not None: + self.context.emit_error_diagnostics = self._previous_emit_error_diagnostics + self._previous_emit_error_diagnostics = None + + def _handle(self, diagnostic: Any) -> bool: + try: + self.diagnostics.append(_mlir_diagnostic_from_callback(diagnostic)) + except Exception: # noqa: BLE001 -- diagnostics must not break compile + return False + return True + + +class CompilerDiagnosticSession: + """Own compiler diagnostic collection and rendering for one pass-manager run.""" + + def __init__( + self, + context: Any, + *, + remark_filter: str = "", + warnings_filter: str = "", + remark_output: str = "", + collect_diagnostics: bool = False, + ) -> None: + self.context = context + self.remark_filter = remark_filter + # Checker domains whose warnings the DSL shows (from --warnings{}). + # Errors always show; warnings are gated to these domains; remarks are + # gated at collection by the remark filter. + self.warnings_filter = warnings_filter + self.remark_output = remark_output + self.collect_diagnostics = bool(collect_diagnostics and not remark_output) + self.enabled = bool(remark_output) or self.collect_diagnostics + self._collector = CompilerDiagnosticCollector( + self.context, self.collect_diagnostics + ) + + def enable(self) -> None: + if not self.enabled: + return + enable_compiler_diagnostics( + self.context, + remark_filter=self.remark_filter, + remark_output=self.remark_output, + collect_diagnostics=self.collect_diagnostics, + ) + + def collecting(self) -> CompilerDiagnosticCollector: + return self._collector + + def format_success(self) -> str: + if not self.collect_diagnostics: + return "" + return format_compiler_diagnostics(self._collected_diagnostics()) + + def format_failure(self, raw_error: str) -> str: + if not self.collect_diagnostics: + return "" + return format_compiler_failure_diagnostics( + self._collected_diagnostics(), raw_error + ) + + def format_backend_failure( + self, + *, + raw_error: str, + nvvm_error: str = "", + ir_context: str = "", + arch: str = "", + location: str = "", + ) -> str: + return format_compiler_backend_failure( + raw_error=raw_error, + nvvm_error=nvvm_error, + ir_context=ir_context, + arch=arch, + location=location, + ) + + def finalize(self) -> None: + if self.enabled: + finalize_compiler_diagnostics(self.context) + + def _collected_diagnostics(self) -> list[CompilerDiagnostic]: + diagnostics = collect_compiler_diagnostics( + self.context, extra_diagnostics=self._collector.diagnostics + ) + diagnostics = _filter_warning_visibility(diagnostics, self.warnings_filter) + return _filter_remark_visibility(diagnostics, self.remark_filter) + + +def _warning_domain(name: str) -> str: + """Map a warning's name to its checker domain, for --warnings{} gating.""" + if name.startswith("NvvmDiag"): + return "nvvm" + return "" + + +def _filter_warning_visibility( + diagnostics: Sequence[CompilerDiagnostic], warnings_filter: str +) -> list[CompilerDiagnostic]: + """Drop warnings whose checker domain is not enabled via --warnings{}. + + Errors and remarks always pass: errors are unconditional, and remarks are + already gated at collection by the remark filter. A warning is shown only + when its domain is listed in ``warnings_filter`` (or it is '.*' / 'all').""" + domains = {d.strip() for d in warnings_filter.split(",") if d.strip()} + show_all = bool(domains & {".*", "all"}) + out: list[CompilerDiagnostic] = [] + for diag in diagnostics: + if diag.severity != "warning" or show_all: + out.append(diag) + continue + domain = _warning_domain(diag.name) + # Domain-tagged warnings show when their domain is requested; untagged + # (generic) warnings show whenever any --warnings{} was passed. + if domain in domains or (not domain and domains): + out.append(diag) + return out + + +def _filter_remark_visibility( + diagnostics: Sequence[CompilerDiagnostic], remark_filter: str +) -> list[CompilerDiagnostic]: + """Drop opt-in remarks unless --remarks{} requested them. + + Errors and warnings always pass (warnings are already gated by + _filter_warning_visibility). Opt-in remarks are shown only when --remarks{} + set a concrete category: the engine then narrowed them to those categories + at collection. When no --remarks{} was passed the filter defaults to the + match-all sentinel (``.*``) — which the engine cannot use to suppress at + collection time — so we treat that default as "not requested" and drop the + remarks here. Errors/warnings (always-on) are unaffected.""" + requested = remark_filter.strip() + if requested and requested not in (".*", "all"): + return list(diagnostics) + return [diag for diag in diagnostics if diag.severity != "remark"] + + +def _load_diagnostic_bindings() -> Any | None: + try: + from .._mlir._mlir_libs._cutlass_ir import _diagnostic + except ImportError: + return None + return _diagnostic + + +def _mlir_diagnostic_from_callback(diagnostic: Any) -> CompilerDiagnostic: + notes = tuple( + str(getattr(note, "message", "")) for note in getattr(diagnostic, "notes", ()) + ) + severity = str(getattr(diagnostic, "severity", "")) + return CompilerDiagnostic( + severity="warning" if severity.endswith(".WARNING") else "error", + message=str(diagnostic.message), + name=_GENERIC_MLIR_DIAGNOSTIC, + location=str(diagnostic.location), + notes=tuple(note for note in notes if note), + ) + + +def enable_compiler_diagnostics( + context: Any, + *, + remark_filter: str = "", + remark_output: str = "", + collect_diagnostics: bool = False, +) -> None: + """Enable YAML remark output or structured diagnostic collection.""" + + bindings = _load_diagnostic_bindings() + if bindings is None: + from .utils.logger import log + + log().warning( + "CutlassIR diagnostic bindings not found; compiler diagnostics " + "and remarks are unavailable" + ) + return + + filter_str = remark_filter or ".*" + if collect_diagnostics: + # Collecting path (a checker is enabled via warnings{}/remarks{}): + # errors/warnings are always-on and collected regardless of any filter + # (warnings are gated later at display by --warnings{}). A non-empty + # remark filter narrows opt-in remarks to the requested categories, but + # an empty filter is read as match-all by the engine. So opt-in remarks + # are gated at display instead (see _filter_remark_visibility): they are + # shown only when --remarks{} was passed. + opt_filter = remark_filter or "" + if opt_filter == "ptxas": + opt_filter = "ptxas.*" + bindings.enable_remark_collection( + context, + policy="all", + passed_filter=opt_filter, + missed_filter=opt_filter, + failed_filter=opt_filter, + analysis_filter=opt_filter, + ) + elif remark_output: + bindings.enable_remarks( + context, + policy="final", + output_file=remark_output, + passed_filter=filter_str, + missed_filter=filter_str, + failed_filter=filter_str, + analysis_filter=filter_str, + ) + + +def _dedupe_compiler_diagnostics( + diagnostics: Sequence[CompilerDiagnostic], +) -> list[CompilerDiagnostic]: + richer_keys = { + (diag.severity, diag.location) + for diag in diagnostics + if diag.name != _GENERIC_MLIR_DIAGNOSTIC + } + seen: set[tuple[str, str, str, str]] = set() + deduped: list[CompilerDiagnostic] = [] + for diag in diagnostics: + if ( + diag.name == _GENERIC_MLIR_DIAGNOSTIC + and (diag.severity, diag.location) in richer_keys + ): + continue + key = (diag.severity, diag.name, diag.message, diag.location) + if key in seen: + continue + seen.add(key) + deduped.append(diag) + return deduped + + +def collect_compiler_diagnostics( + context: Any, extra_diagnostics: Sequence[CompilerDiagnostic] = () +) -> list[CompilerDiagnostic]: + """Collect structured compiler diagnostics from the current MLIR context.""" + + bindings = _load_diagnostic_bindings() + diagnostics: list[CompilerDiagnostic] = [] + if bindings is not None: + for item in bindings.get_collected_remarks(context): + diagnostic = compiler_diagnostic_from_remark(item) + if diagnostic is not None: + diagnostics.append(diagnostic) + diagnostics.extend(extra_diagnostics) + return _dedupe_compiler_diagnostics(diagnostics) + + +def finalize_compiler_diagnostics(context: Any) -> None: + """Finalize the MLIR remark engine, if diagnostic bindings are available.""" + + bindings = _load_diagnostic_bindings() + if bindings is None: + return + bindings.finalize_remarks(context) + + +def find_user_source_location() -> tuple[ + str | None, int | None, int | None, int | None +]: + """Best-effort author location: ``(filename, line, col, end_col)``. + + Walks out to the nearest call-stack frame *outside* the DSL package (frames + classified by module name -- robust across source / build-tree / installed + layouts), so a diagnostic raised deep inside the DSL still points at the + author's own line. On Python 3.11+ the column span of the instruction the + author frame is suspended at is recovered via ``co_positions`` so the caret + can underline it. Returns ``(None, None, None, None)`` when no author frame + is on the stack. + """ + try: + frame: types.FrameType | None = sys._getframe(1) + except Exception: # noqa: BLE001 + return None, None, None, None + stdlib: frozenset[str] = getattr(sys, "stdlib_module_names", frozenset()) + try: + while frame is not None: + mod = frame.f_globals.get("__name__", "") or "" + top = mod.split(".", 1)[0] + fn = frame.f_code.co_filename + is_internal = ( + mod == "cutlass" + or mod.startswith("cutlass.") + or top in stdlib + or fn.startswith("<") + ) + if not is_internal: + line, col, end_col = frame.f_lineno, None, None + try: # exact column span of the suspended instruction (3.11+) + positions = list(frame.f_code.co_positions()) # type: ignore[attr-defined] + idx = frame.f_lasti // 2 + if 0 <= idx < len(positions): + sl, el, sc, ec = positions[idx] + if sl is not None: + line = sl + if sc is not None and el == sl: + col, end_col = sc, ec + except Exception: # noqa: BLE001 -- best-effort column + pass + return fn, line, col, end_col + frame = frame.f_back + finally: + del frame + return None, None, None, None + + +def render_code_frame( + filename: str | None, + line: int | None, + col: int | None = None, + end_col: int | None = None, +) -> str | None: + """Code frame: ``--> file:line:col`` + gutter + source lines + caret. + + Shows the error line preceded by up to two lines of context. The ``^`` + underlines ``[col, end_col)`` on the error line when a column span is known, + otherwise it points at the first non-blank character of that line. + Best-effort -- returns just the ``-->`` location line if the source cannot + be read. + """ + if not filename or not line: + return None + frame = _format_user_source_frame(filename, line, col, end_col) + return "\n".join(frame) if frame else None + + +def render_user_diagnostic(err: Any) -> str: + """Render a DSL user diagnostic or warning. + + Used by ``DSLBaseError`` and ``DSLWarning`` so AST pre-processing, tracing, + runtime errors, and warnings share one source-frame and suggestion format. + """ + parts = [] + cause_text = err._generate_cause() + code = getattr(err, "code", None) + frame = render_code_frame( + err.filename, err.line, getattr(err, "col", None), getattr(err, "end_col", None) + ) + + if getattr(err, "_is_internal", False): + return _format_internal_error_diagnostic(err, frame, cause_text) + + severity = "warning" if getattr(err, "_severity", "error") == "warning" else "error" + parts.append( + _format_diagnostic_headline( + severity, + err.message, + code=code or "", + bold=True, + leading_newline=True, + ) + ) + + if frame: + parts.append(frame) + + if getattr(err, "error_code", None) is not None: + parts.extend(_format_user_labeled_text("note", f"error code {err.error_code}")) + if cause_text: + parts.extend(_format_user_labeled_text("note", cause_text)) + if err.context: + if isinstance(err.context, dict): + for key, value in err.context.items(): + parts.extend(_format_user_labeled_text("note", f"{key}: {value}")) + else: + parts.extend(_format_user_labeled_text("note", str(err.context))) + if err.suggestion: + fixes = ( + err.suggestion + if isinstance(err.suggestion, (list, tuple)) + else [err.suggestion] + ) + for s in fixes: + parts.extend(_format_labeled_text("suggestion", str(s))) + + parts.append("") + return "\n".join(parts) + + +def _format_internal_error_diagnostic( + err: Any, frame: str | None, cause_text: str +) -> str: + """Format internal compiler errors with a user-facing bug-report envelope.""" + # Internal errors (compiler bugs) get a "please report" envelope instead of + # a "here's your mistake + fix" block -- they are not the author's fault. + parts = [ + f"\n{_Colors.RED}{_Colors.BOLD}[Internal Error] The compiler hit a " + f"problem it could not trace back to your code.{_Colors.RESET}", + "This is a bug in the DSL, not a mistake in your kernel.\n", + ] + if frame: + parts.append(f"{_Colors.BLUE}Closest point in your code:{_Colors.RESET}") + parts.append(frame) + parts.append("") + if err.message: + parts.append(f"{_Colors.BLUE}Detail:{_Colors.RESET} {err.message}") + parts.append("") + if cause_text: + parts.append(f"{_Colors.BLUE}Cause:{_Colors.RESET} {cause_text}") + parts.append("") + parts.append(f"{_Colors.GREEN}What to do:{_Colors.RESET}") + parts.append( + f" {_Colors.GREEN}Please report this with the snippet above and your " + f"kernel.{_Colors.RESET}" + ) + parts.append( + f" {_Colors.GREEN}Re-run with CUTE_DSL_SHOW_STACKTRACE=1 to include the " + f"full technical detail.{_Colors.RESET}" + ) + parts.append("") + parts.append("=" * 100) + return "\n".join(parts) + + +def _nearest_function_name(lines: list[str], line_no: int) -> str | None: + for idx in range(min(line_no - 1, len(lines) - 1), -1, -1): + match = re.match( + r"\s*(?:async\s+)?def\s+([A-Za-z_][A-Za-z0-9_]*)\(", + lines[idx], + ) + if match: + return match.group(1) + return None + + +def _read_source_lines(file_path: str | Path) -> list[str]: + filename = str(file_path) + try: + lines = linecache.getlines(filename) + except Exception: # noqa: BLE001 + lines = [] + if lines: + return [line.rstrip("\n") for line in lines] + try: + return Path(filename).read_text(encoding="utf-8").splitlines() + except (OSError, UnicodeDecodeError): + return [] + + +def _display_column(col: int | None, *, col_zero_based: bool) -> int | None: + if col is None: + return None + return col + 1 if col_zero_based else col + + +def _caret_column(source_line: str, col: int | None, *, col_zero_based: bool) -> int: + if col is None: + return len(source_line) - len(source_line.lstrip()) + return max(col if col_zero_based else col - 1, 0) + + +def _caret_span( + source_line: str, + col: int | None, + end_col: int | None, + *, + col_zero_based: bool, +) -> int: + if col is None or end_col is None: + return 1 + start = _caret_column(source_line, col, col_zero_based=col_zero_based) + end = end_col if col_zero_based else end_col - 1 + end = min(end, len(source_line)) + return max(1, end - start) if end > start else 1 + + +def _format_source_location( + filename: str, + line: int, + col: int | None, + *, + absolute_path: bool, + col_zero_based: bool, +) -> str: + display_path = str(Path(filename).resolve()) if absolute_path else filename + display_col = _display_column(col, col_zero_based=col_zero_based) + loc = f"{display_path}:{line}" + if display_col is not None: + loc += f":{display_col}" + return loc + + +def _format_user_source_frame( + filename: str, line: int, col: int | None, end_col: int | None +) -> list[str]: + source_lines = _read_source_lines(filename) + loc = _format_source_location( + filename, line, col, absolute_path=False, col_zero_based=True + ) + width = len(str(line)) + pad = " " * width + lines = [f"{pad}{_Colors.BLUE}-->{_Colors.RESET} {loc}"] + if not (1 <= line <= len(source_lines)): + return lines + source_line = source_lines[line - 1] + if not source_line.strip(): + return lines + caret_col = _caret_column(source_line, col, col_zero_based=True) + span = _caret_span(source_line, col, end_col, col_zero_based=True) + caret = f"{_Colors.RED}{'^' * span}{_Colors.RESET}" + lines.append(f"{pad} |") + for current in range(max(1, line - 2), line): + context_line = source_lines[current - 1] + lines.append(f"{current:>{width}} | {context_line}") + lines.append(f"{line} | {source_line}") + lines.append(f"{pad} | {' ' * caret_col}{caret}") + return lines + + +def _format_compiler_source_frame(filename: str, line: int, col: int) -> list[str]: + source_lines = _read_source_lines(filename) + loc = _format_source_location( + filename, line, col, absolute_path=True, col_zero_based=False + ) + lines = ["", f" --> {loc}"] + fn_name = _nearest_function_name(source_lines, line) + if fn_name: + lines.append( + f" in function `{_Colors.GREEN}{fn_name}{_Colors.RESET}(...)`:" + ) + if not source_lines: + return lines + + start = max(1, line - _COMPILER_CONTEXT_LINES) + end = min(len(source_lines), line + _COMPILER_CONTEXT_LINES) + width = len(str(end)) + lines.append(" |") + for current in range(start, end + 1): + source_line = source_lines[current - 1] + prefix = ">" if current == line else " " + lines.append(f"{prefix} {current:{width}d} | {source_line}") + if current == line: + caret_col = _caret_column(source_line, col, col_zero_based=False) + lines.append(f" {' ' * width} | {' ' * caret_col}^") + return lines + + +_DIAGNOSTIC_LABEL_COLORS = { + "error": _Colors.RED, + "warning": _Colors.YELLOW, + "remark": _Colors.CYAN, + "suggestion": _Colors.GREEN, + "note": _Colors.BLUE, +} + + +def _diagnostic_label_prefixes( + label: str, *, marker: str = "", prefix: str = " " +) -> tuple[str, str]: + plain_prefix = f"{prefix}{marker}{label}: " + color = _DIAGNOSTIC_LABEL_COLORS.get(label, "") + colored_label = f"{marker}{label}:" + colored_prefix = ( + f"{prefix}{color}{colored_label}{_Colors.RESET} " if color else plain_prefix + ) + return plain_prefix, colored_prefix + + +def _diagnostic_label_indent(label: str, *, marker: str = "") -> str: + plain_prefix, _ = _diagnostic_label_prefixes(label, marker=marker) + return " " * len(plain_prefix) + + +def _format_labeled_text( + label: str, text: str, *, marker: str = "", wrap: bool = True +) -> list[str]: + plain_prefix, colored_prefix = _diagnostic_label_prefixes(label, marker=marker) + if not wrap: + if not text: + return [colored_prefix.rstrip()] + return [colored_prefix + text] + + wrapped = textwrap.wrap( + text, + width=max(20, _COMPILER_DIAG_TEXT_WIDTH - len(plain_prefix)), + break_long_words=False, + break_on_hyphens=False, + ) + if not wrapped: + return [colored_prefix.rstrip()] + + lines = [colored_prefix + wrapped[0]] + lines.extend( + _diagnostic_label_indent(label, marker=marker) + line for line in wrapped[1:] + ) + return lines + + +def _format_user_labeled_text(label: str, text: str) -> list[str]: + return _format_labeled_text(label, text, marker="= ", wrap=False) + + +def _format_labeled_block(label: str, heading: str, body: str) -> list[str]: + lines = _format_labeled_text(label, heading) + indent = _diagnostic_label_indent(label) + width = max(20, _COMPILER_DIAG_TEXT_WIDTH - len(indent)) + for raw_line in body.splitlines(): + if not raw_line: + lines.append(indent.rstrip()) + continue + wrapped = textwrap.wrap( + raw_line, + width=width, + break_long_words=False, + break_on_hyphens=False, + ) + lines.extend(indent + line for line in wrapped) + return lines + + +def _diagnostic_marker(severity: str, code: str = "", namespace: str = "") -> str: + marker = severity + if code: + marker = ( + f"{severity}[{namespace}:{code}]" if namespace else f"{severity}[{code}]" + ) + return marker + + +def _format_diagnostic_headline( + severity: str, + message: str, + *, + code: str = "", + namespace: str = "", + bold: bool = False, + leading_newline: bool = False, +) -> str: + marker = _diagnostic_marker(severity, code, namespace) + color = _DIAGNOSTIC_LABEL_COLORS.get(severity, "") + style = f"{color}{_Colors.BOLD if bold else ''}" + prefix = "\n" if leading_newline else "" + if style: + return f"{prefix}{style}{marker}:{_Colors.RESET} {message}" + return f"{prefix}{marker}: {message}" + + +def _format_compiler_diagnostic_headline( + severity: str, message: str, code: str = "", namespace: str = "" +) -> str: + namespace = namespace or ("nvvm-diag" if code.startswith("C") else "") + return _format_diagnostic_headline( + severity, message, code=code, namespace=namespace + ) + + +def _format_compiler_location(location: str) -> list[str]: + if not location: + return [] + + loc_match = _PY_LOC_RE.search(location) + if not loc_match: + return ["", f" --> {location}"] + + return _format_compiler_source_frame( + loc_match.group("file"), + int(loc_match.group("line")), + int(loc_match.group("col")), + ) + + +def _format_compiler_ptx_reference(ptx_ref: str, ptx_url: str) -> list[str]: + if not ptx_ref and not ptx_url: + return [] + + if ptx_ref and ptx_url: + return [ + *_format_labeled_text("note", f"PTX ISA {ptx_ref}"), + _diagnostic_label_indent("note") + f"docs: {ptx_url}", + ] + if ptx_ref: + return _format_labeled_text("note", f"PTX ISA {ptx_ref}") + return _format_labeled_text("note", f"docs: {ptx_url}") + + +def _format_compiler_diagnostic_body(diag: CompilerDiagnostic) -> list[str]: + lines: list[str] = [] + if diag.reason: + reason_label = ( + diag.severity + if diag.severity in ("error", "warning", "remark") + else "error" + ) + lines.extend(_format_labeled_text(reason_label, diag.reason)) + if diag.suggestion: + lines.extend(_format_labeled_text("suggestion", diag.suggestion)) + for note in diag.notes: + lines.extend(_format_labeled_text("note", note)) + lines.extend(_format_compiler_ptx_reference(diag.ptx_ref, diag.ptx_url)) + return lines + + +def _is_concise_nvvm_backend_error(nvvm_error: str) -> bool: + stripped = nvvm_error.strip() + return "\n" not in stripped and stripped.startswith("Cannot select: intrinsic ") + + +def extract_compiler_location(text: str) -> str: + """Return the first Python source location embedded in compiler text.""" + + loc_match = _PY_LOC_RE.search(text) + return loc_match.group(0) if loc_match else "" + + +def format_compiler_backend_failure( + *, + raw_error: str, + nvvm_error: str = "", + ir_context: str = "", + arch: str = "", + location: str = "", +) -> str: + """Render an unstructured backend failure with compiler diagnostic styling.""" + + is_nvvm_failure = bool(nvvm_error) + concise_nvvm_error = is_nvvm_failure and _is_concise_nvvm_backend_error(nvvm_error) + headline = ( + f"NVVM backend compilation failed: {nvvm_error.strip()}" + if concise_nvvm_error + else "NVVM backend compilation failed" + if is_nvvm_failure + else "compiler backend failed" + ) + lines = [_format_compiler_diagnostic_headline("error", headline)] + lines.extend(_format_compiler_location(location)) + reason = ( + "" + if concise_nvvm_error + else "libNVVM failed while compiling generated device IR." + if is_nvvm_failure + else "the compiler backend failed while compiling generated device code." + ) + if reason: + lines.extend(_format_labeled_text("error", reason)) + + labeled_blocks: list[tuple[str, str]] = [] + if nvvm_error and not concise_nvvm_error: + labeled_blocks.append(("backend log:", nvvm_error.strip())) + elif raw_error and not concise_nvvm_error: + labeled_blocks.append(("raw compiler error:", raw_error.strip())) + if ir_context and not concise_nvvm_error: + labeled_blocks.append(("IR context (truncated):", ir_context.strip())) + + if arch: + lines.extend(_format_labeled_text("note", f"target architecture: {arch}")) + for heading, body in labeled_blocks: + if body: + lines.extend(_format_labeled_block("note", heading, body)) + + suggestion = "check that CUDA_TOOLKIT_PATH is set correctly" + if arch: + suggestion += f" and that the CUDA toolkit supports target architecture {arch}" + suggestion += "; rerun with IR dump flags if backend context is needed" + lines.extend(_format_labeled_text("suggestion", suggestion)) + return "\n".join(lines) + + +def compiler_diagnostic_from_remark( + item: Mapping[str, Any], +) -> CompilerDiagnostic | None: + diag_class = str(item.get("diag_class", "")) + name = str(item.get("name", "")) + if diag_class == "Error": + severity = "error" + elif diag_class == "Warning": + severity = "warning" + elif diag_class == "OptRemark": + # Opt-in remarks (perf / optimization / analysis). Rendered through the + # same path as errors/warnings, labeled remark[]. Which + # remarks reach here is gated by the remark filter at collection time + # (see enable_compiler_diagnostics), so this does not flood by default. + severity = "remark" + else: + return None + + if severity == "remark": + # Remarks are keyed by their category (e.g. "ptxas"), not the + # error/warning code table. + code = str(item.get("category", "")) or name + namespace = "" + ptx_ref = "" + ptx_url = "" + else: + info = _COMPILER_DIAGNOSTICS.get(name) + code = info.code if info else str(item.get("remark_id", "")) + namespace = info.namespace if info else "" + ptx_ref = info.ptx_ref if info else "" + ptx_url = info.ptx_url if info else "" + return CompilerDiagnostic( + severity=severity, + message=str(item.get("message", "")), + code=code, + namespace=namespace, + name=name, + location=str(item.get("location", "")), + reason=str(item.get("reason", "")), + suggestion=str(item.get("suggestion", "")), + ptx_ref=ptx_ref, + ptx_url=ptx_url, + ) + + +def format_compiler_diagnostics(diagnostics: Sequence[CompilerDiagnostic]) -> str: + blocks: list[str] = [] + for diag in diagnostics: + lines = [ + _format_compiler_diagnostic_headline( + diag.severity, diag.message, diag.code, diag.namespace + ) + ] + lines.extend(_format_compiler_location(diag.location)) + lines.extend(_format_compiler_diagnostic_body(diag)) + blocks.append("\n".join(lines)) + return "\n\n".join(blocks) + + +def format_compiler_failure_diagnostics( + diagnostics: Sequence[CompilerDiagnostic], raw_error: str +) -> str: + diagnostics = _dedupe_compiler_diagnostics(diagnostics) + formatted = format_compiler_diagnostics(diagnostics) + if not formatted: + return "" + if raw_error and any( + diag.name == _DEVICE_BINARY_SERIALIZATION_DIAGNOSTIC for diag in diagnostics + ): + raw_block = "\n".join( + _format_labeled_block("note", "raw compiler error:", raw_error.strip()) + ) + return f"{formatted}\n{raw_block}" + if any(diag.severity == "error" for diag in diagnostics) or not raw_error: + return formatted + failure_headline = _format_compiler_diagnostic_headline( + "error", "compiler pass failed before a structured error was collected" + ) + return f"{formatted}\n\n{failure_headline}\n{raw_error}" + + +class _MissingField: + """Renders a not-supplied ``{field}`` back as the literal placeholder. + + A diagnostic must never crash the compile it is trying to explain, so a + template referencing a field the call site forgot degrades to showing the + raw ``{field}`` (caught by tests) instead of raising ``KeyError``. + """ + + __slots__ = ("key",) + + def __init__(self, key: str) -> None: + self.key = key + + def __format__(self, spec: str) -> str: + return "{" + self.key + (":" + spec if spec else "") + "}" + + def __str__(self) -> str: + return "{" + self.key + "}" + + +class _SafeFields(dict): + def __missing__(self, key: str) -> "_MissingField": + return _MissingField(key) + + +class _DiagMixin: + """Shared behaviour for the diagnostic catalogs (:class:`DiagId` errors and + :class:`WarnId` warnings). + + Both are ``enum.Enum``s whose member **name** is the stable code and whose + value is ``(message_template, (fix_line, ...))``. This mix-in supplies the + code/category/message/fix accessors and ``_fill`` (template substitution); + each enum adds its own ``format``/``format_warning`` for the code label. + """ + + # ``name`` and ``value`` are supplied by ``enum.Enum`` at runtime (this + # mix-in is combined with ``enum.Enum`` in DiagId/WarnId); declared here so + # the accessor bodies below type-check. + name: str + value: Any # the (message_template, (fix_line, ...)) tuple, per enum member + + @property + def code(self) -> str: + """The stable, human-readable code shown to the user (the name).""" + return self.name + + @property + def category(self) -> str: + """High-level group, derived from the name prefix (PHASE / TYPE / ...).""" + return self.name.split("_", 1)[0] + + @property + def message(self) -> str: + return self.value[0] + + @property + def fix(self) -> tuple[str, ...]: + return self.value[1] + + def fill(self, **fields: Any) -> tuple[str, tuple[str, ...]]: + """Fill the templates and return ``(message, fixes)``. + + No code suffix -- the stable code (``.code``) is shown separately in the + ``error[CODE]:`` / ``warning[CODE]:`` header. The two + value-phase names :data:`META_VALUE` / :data:`STAGED_VALUE` are always + available as ``{meta}`` / ``{staged}``; a field a call site forgot + degrades to the literal ``{field}`` instead of crashing (see + :class:`_SafeFields`).""" + fields.setdefault("meta", META_VALUE) + fields.setdefault("staged", STAGED_VALUE) + safe = _SafeFields(fields) + message = self.message.format_map(safe) + fixes = tuple(f.format_map(safe) for f in self.fix) + return message, fixes + + +class DiagId(_DiagMixin, enum.Enum): + """User-facing diagnostics: member name == stable code, value == (message, fix). + + Naming convention ``_`` -- CATEGORY is one of + PHASE / TYPE / SCOPE / CONTAINER / UNSUP / ARG / CALL / CONFIG / LAUNCH / + TENSOR / AOT and is what ``.category`` returns. Keep the prefix when adding + entries. Internal/compiler errors do not live here -- raise + ``DSLRuntimeError`` for those. + """ + + PHASE_ASSIGN_PYTHON_TO_TRACKED = ( + "`{var}` is a {staged}, but a {meta} is being assigned to it inside a " + "for/while/if. Only a {staged} can be used there.", + ( + "Wrap the value in a runtime type, e.g. Int32(...), Float32(...), or " + "Boolean(...).", + ), + ) + PHASE_MUTATE_PYTHON = ( + "`{var}` is a {meta}, but it is being changed inside a for/while/if whose " + "path is decided at run time. A {meta} is fixed once your code is read and " + "cannot change there.", + ( + "If `{var}` must change there, create it as a {staged} first, e.g. " + "`{var} = Int32(0)`, then update it.", + "If it is a fixed setting, set it once before the for/while/if.", + ), + ) + PHASE_PYTHON_THEN_TRACKED = ( + "`{var}` was defined as a {meta} (type {old_type}, value {old_value}) but " + "is being assigned a {staged} inside a for/while/if. Give it a runtime type " + "from the start so it has one consistent type.", + ( + "Initialize `{var}` as a {staged} before the for/while/if, e.g. " + "`{var} = Int32(0)`.", + ), + ) + TYPE_UNSTABLE_JOIN = ( + "`{var}` has type `{old_type}` on one path and `{new_type}` on another. " + "Where the paths come back together, both versions must be the same type.", + ( + "Make every assignment to `{var}` produce the same type.", + "If you need a conversion, convert it explicitly inside the branch so the " + "result is consistently one type.", + ), + ) + SCOPE_READ_NEVER_SET = ( + "A variable is read on a path where it was never given a value. It must be " + "set before it is read, on every path that can reach this point.", + ( + "Set the variable before the for/while/if, or set it on every branch " + "that can reach this read.", + ), + ) + CONTAINER_TUPLE_LENGTH_CHANGED = ( + "`{var}` is a tuple with {old} item(s) on one pass and {new} on another. A " + "tuple must keep the same number of items every pass so its parts can be " + "followed.", + ("Keep `{var}` the same length on every branch and every loop pass.",), + ) + CONTAINER_STRUCTURE_CHANGED = ( + "`{var}` has a different structure at the end of this `{op_type}` than at " + "the start. A value carried through a `{op_type}` must keep the same shape " + "(same fields and the same number of parts) on every pass.", + ( + "Give `{var}` the same structure on every branch and every pass of the " + "`{op_type}`.", + ), + ) + CONTAINER_OBJECT_FIELD_MISMATCH = ( + "`{var}` is replaced inside a for/while/if with an object that has a " + "different set of fields ({detail}). Replacement only works when the new " + "object has exactly the same fields.", + ( + "Update the individual fields of `{var}` instead of replacing the whole " + "object, or make the new object carry exactly the same fields.", + ), + ) + CONTAINER_OBJECT_REPLACED = ( + "`{var}` is replaced as a whole object inside a for/while/if, but it holds " + "runtime values that cannot be followed through a whole-object replacement.", + ( + "Update the individual fields of `{var}` instead, e.g. " + "`{var}.field = new_value`.", + ), + ) + PHASE_CONVERSION_FAILED = ( + "`{var}` could not be turned into a runtime {new_type} value (its current " + "value {old_value} is not compatible). Give it a compatible runtime type " + "before the loop.", + ( + "Create `{var}` as a {staged} of a compatible type before the " + "for/while/if, e.g. `{var} = {new_type}(...)`.", + ), + ) + UNSUP_LOOP_ELSE = ( + "A `for`/`while` loop with an `else:` clause is not supported in compiled " + "code.", + ("Remove the `else:` and put its code after the loop.",), + ) + # --- AST-preprocessing (compile-time) user errors --- + UNSUP_EARLY_EXIT = ( + "Early exit ({kind}) is not allowed in {where}. A staged for/while/if runs " + "as a whole, so it cannot return/break/continue out early.", + ( + "If the condition is decided at compile time, write `if const_expr(...)` " + "or `for ... in range_constexpr(...)` -- then the early exit runs in " + "Python and is allowed.", + ), + ) + UNSUP_RANGE_ARGS = ( + "A staged `range(...)` takes 1 to 3 arguments.", + ( + "Call it as `range(stop)`, `range(start, stop)`, or " + "`range(start, stop, step)`.", + ), + ) + UNSUP_FSTRING = ( + "This kind of expression can't be used inside an f-string in compiled code.", + ( + "Compute the value into a variable first, then reference that variable in " + "the f-string.", + ), + ) + UNSUP_READ_UNDERSCORE = ( + "`_` is a throwaway name and cannot be read.", + ("Give the value a real name if you need to read it.",), + ) + UNSUP_DECORATOR_ORDER = ( + "The `{decorator}` decorator must be the innermost decorator (closest to " + "`def`).", + ("Move `{decorator}` directly above the function definition.",), + ) + UNSUP_GLOBAL = ( + "`global` is not supported in a compiled function.", + ("Pass the value in as a function argument instead.",), + ) + UNSUP_NONLOCAL = ( + "`{stmt}` refers to `{name}`, which is not available inside the compiled " + "function.", + ("Pass `{name}` in as a function argument instead.",), + ) + UNSUP_MIXED_ASSIGN_TARGETS = ( + "An assignment that mixes plain names, subscripts, and tuple targets on a " + "single line is not supported here.", + ("Split it into separate assignment statements, one kind of target per line.",), + ) + + # ===================================================================== + # Migrated from raw DSLRuntimeError raises across the DSL (author + # mistakes: config, types, arguments, calls, launch, tensors, ...). + # ===================================================================== + # --- AOT --- + AOT_MISSING_RETURN_TYPE = ( + "The exported function must declare its return type.", + ( + "Add a return type annotation to the function definition, e.g. `-> Int32:` or `-> Float32:`.", + ), + ) + AOT_UNSUPPORTED_RETURN_TYPE = ( + "The return type {return_type} is not supported for AOT export. Use Int32, Int64, Float32, or Float64.", + ( + "Change the function's return type annotation to a supported type: Int32, Int64, Float32, or Float64.", + ), + ) + # --- ARG --- + ARG_BIND_FAILED = ( + "The arguments do not match the function signature.", + ("Check that you are passing the correct number and types of arguments.",), + ) + ARG_COUNT_MISMATCH = ( + "Passed {got} argument(s) to FFI function, but it expects {expected}", + ("Check the function signature and provide the correct number of arguments",), + ) + ARG_FOR_LOOP_STEP_NOT_INT = ( + "Loop bounds and step must be integers, not {kind}", + ( + "Pass integer values for start, stop, and step to for_yield().", + "If you need a fractional step, multiply by an integer scale and adjust your loop logic.", + ), + ) + ARG_INVALID_ALIGNMENT = ( + "Alignment value must be a positive power of 2", + ("Pass a positive power of 2 (e.g., 1, 2, 4, 8, 16) as the alignment value",), + ) + ARG_NON_CONSTANT = ( + "`range_constexpr()` requires all arguments to be Python integers known at code-read time.", + ( + "Use `range(start, stop, step)` instead if your bounds are {staged}.", + "Ensure all arguments are Python int literals or const Python variables.", + ), + ) + ARG_NOT_A_TENSOR = ( + "Argument `{arg_name}` must be a tensor (array or object supporting __dlpack__), but you passed something that is not.", + ( + "Pass a tensor or array-like object (numpy array, GPU tensor, etc.) for `{arg_name}`.", + ), + ) + ARG_NOT_MARSHALABLE = ( + "Argument `{arg_name}` of type `{arg_type}` cannot be passed to the kernel: it exposes neither a C-pointer interface nor a registered JIT argument adapter, so it would reach the kernel as zero arguments.", + ( + "Pass a tensor, a numeric value, or a pointer the DSL recognizes for `{arg_name}`.", + "Or register a JIT argument adapter for type `{arg_type}`.", + ), + ) + ARG_NOT_NUMERIC = ( + "Argument `{arg_name}` expects a numeric value, but you passed a `{arg_type}`.", + ( + "Pass a numeric value (an int, a float, or a DSL numeric type) for `{arg_name}`.", + ), + ) + ARG_POINTER_NEGATIVE = ( + "Pointer address must be non-negative (got {address})", + ("Pass a non-negative address, or use nullptr() for null pointers",), + ) + ARG_SPLIT_MODE_INVALID = ( + "The split_mode must be one of the valid CudaDlcSplitMode enum values.", + ( + "Use cuda_helpers.CudaDlcSplitMode enum: NONE, ANY, SINGLE_DLCC, or DUAL_DLCC.", + ), + ) + ARG_TENSOR_NOT_ON_DEVICE = ( + "The tensor `{arg_name}` must be in GPU memory, but it is currently on the host. Move it to the GPU before passing it to the kernel.", + ( + "Transfer the tensor to GPU memory using your framework's GPU allocation (e.g., `.cuda()` for PyTorch, `cp.array()` for CuPy).", + ), + ) + ARG_TENSOR_REQUIRED = ( + "`{arg_name}` is declared as a `{type_kind}`, but you passed something that is not a tensor. Pass a tensor (array/DLPack-compatible object).", + ( + "Pass a tensor or array-like object (numpy array, GPU tensor, etc.) for `{arg_name}`.", + ), + ) + ARG_TYPE_MISMATCH = ( + "Pointer alignment must be created using align(...), not a plain integer", + ("Wrap the alignment value with align(), e.g. align(4) instead of just 4",), + ) + ARG_UNORDERED_CONTAINER = ( + "Arguments cannot be sets because sets are unordered and the compiler needs to preserve argument order.", + ( + "Convert the set to a list or tuple, e.g. `list(my_set)` or `tuple(my_set)`.", + ), + ) + ARG_UNSUPPORTED_C_EXPORT = ( + "Argument {arg_name} has type {arg_type}, which cannot be exported to C. Only numeric types (Int32, Float32, etc.) and CUstream are supported.", + ( + "Use a supported scalar type (Int8/16/32/64, Uint8/16/32/64, Float32/64, BFloat16, Float16, Boolean) or CUstream.", + "Check that your function arguments have correct type annotations.", + ), + ) + ARG_UNSUPPORTED_TYPE = ( + "Argument {arg_name} (position {num}) cannot be converted to a {phase_label} for function '{function_name}'. Its type `{arg_type}` is not supported.", + ( + "If the value is known at compile time, annotate it: `{arg_name}: Constexpr`.", + "Otherwise, implement the `JitArgument` or `DynamicExpression` protocol for type `{arg_type}`, or register a custom argument adapter.", + ), + ) + ARG_WORKSPACE_COUNT_MISMATCH = ( + "The kernel call has {got} extra workspace arguments, but {expected} are required.", + ("Provide the correct number of workspace pointers when calling the kernel.",), + ) + ARG_WRONG_TYPE = ( + "Loop `{name}` must be an integer type, not `{dtype}`. Loop bounds (start, stop, step) must all be integers.", + ( + "Use integer values for `range()` or loop bounds.", + "Convert the value to an integer type, e.g., `int({name})`.", + ), + ) + # --- CALL --- + CALL_BUILTIN_KWARGS_UNSUPPORTED = ( + "The built-in function '{fcn}' does not support keyword arguments.", + ("Remove keyword arguments or use positional arguments only.",), + ) + CALL_DUPLICATE_ARGUMENT = ( + "Got multiple values for argument {argument_name}", + ( + "Pass the value either as a positional argument or as a keyword argument, not both", + ), + ) + CALL_EXPERIMENTAL_MISMATCH = ( + "The is_experimental flag does not match the function's DSL routing. Either remove the flag or use @cute.jit(is_experimental=True).", + ( + "If using experimental DSL, decorate the function with @cute.jit(is_experimental=True).", + "Otherwise, remove the is_experimental=True flag from cute.compile().", + ), + ) + CALL_FFI_NO_KWARGS = ( + "FFI functions do not accept keyword arguments", + ("Use only positional arguments when calling FFI functions",), + ) + CALL_FFI_NO_MATCH = ( + "No compatible overload found for the arguments passed to this FFI function", + ( + "Check that your argument types match one of the function's overloads", + "Use the function's type annotations to see what types are accepted", + ), + ) + CALL_FUNCTION_NOT_PROVIDED = ( + "No function was provided to compile. Pass a callable decorated with @cute.jit.", + ("Ensure you pass a valid @cute.jit-decorated function to cute.compile().",), + ) + CALL_MISSING_ARG = ( + "Required argument `{name}` is missing in the call to `{function_name}`.", + ("Pass a value for `{name}`.",), + ) + CALL_MISSING_ARGS = ( + "Missing required arguments: {missing}", + ("Provide all required arguments without defaults",), + ) + CALL_MISSING_JIT_DECORATOR = ( + "Function must be decorated with @cute.jit before passing it to cute.compile().", + ("Add @cute.jit decorator to the function definition.",), + ) + CALL_NOT_CALLABLE = ( + "The provided object is not callable. Pass a function or method decorated with @cute.jit.", + ( + "Ensure the first argument to cute.compile() is a callable (function, method, or callable object).", + ), + ) + CALL_SIGNATURE_MISMATCH = ( + "Arguments do not match the function signature (got {provided} args, {provided_kw} kwargs)", + ("Verify that all arguments match the expected signature",), + ) + CALL_TOO_MANY_ARGS = ( + "Too many positional arguments (expected {expected}, got {provided})", + ("Remove extra positional arguments or pass them as keyword arguments",), + ) + CALL_UNEXPECTED_KWARG = ( + "Unexpected keyword argument: {argument_name}", + ( + "Remove the argument or check the function signature for the correct parameter name", + ), + ) + CALL_UNSUPPORTED_CALLABLE_TYPE = ( + "The provided object is of an unsupported callable type. Only functions, methods, and callable instances are supported.", + ( + "Pass a regular function, bound method, or callable instance (with __call__ method).", + ), + ) + CALL_WRONG_IMPORT = ( + "`{name}` was imported from a different module; you must use the one from the DSL package.", + ( + "Remove any local definition or import of `{name}`.", + "Use the version from the DSL (cutlass.range, cutlass.if_, etc.).", + ), + ) + # --- CONFIG --- + CONFIG_ATTRIBUTES_INVALID_TYPE = ( + "`attributes=` must be a dict, or a callable that returns a dict.", + ( + "Pass a dict of string keys to string (or attribute) values, or a callable returning one.", + ), + ) + CONFIG_ATTRIBUTES_UNSUPPORTED = ( + "The `@kernel` decorator with `attributes=` is not supported by this DSL. Only experimental CuTe supports this feature.", + ( + "Remove the `attributes=` parameter from the @kernel decorator.", + "Or use `@cute.experimental.kernel(attributes=...)` if you need kernel-level attributes.", + ), + ) + CONFIG_ATTR_KEY_UNSUPPORTED = ( + "The attribute key '{key}' is not supported by this DSL.", + ("Use one of the allowed attribute keys for your DSL.",), + ) + CONFIG_ATTR_VALUE_TYPE_INVALID = ( + "The value for attribute '{key}' must be a string (or attribute) value.", + ("Use a string value for attribute '{key}'.",), + ) + CONFIG_BINARY_LOAD_PREFIX_REQUIRED = ( + "Binary loading requires a function prefix, but none was provided.", + ( + "Provide a prefix when using load_from_binary=True.", + "If loading from a compiled binary, the prefix must match the function name used when the binary was created.", + ), + ) + CONFIG_CLUSTER_DYNAMIC = ( + "The cluster dimension must use static/Constexpr values, not runtime values.", + ( + "Use `Constexpr` annotation or a Python constant for cluster/fallback_cluster.", + ), + ) + CONFIG_EXPORT_NO_RUNTIME_ARGS = ( + "Cannot export a function with no runtime arguments. All of the function's parameters are compile-time constants, but a C interface needs at least one runtime parameter to call.", + ( + "Remove the constexpr annotation from at least one parameter so it becomes a runtime argument.", + "If all parameters must be compile-time constants, you cannot export this function as a standalone C library.", + ), + ) + CONFIG_FALLBACK_CLUSTER_REQUIRES_CLUSTER = ( + "You cannot set `fallback_cluster` without also setting `cluster`.", + ("If using fallback_cluster, also set cluster in the LaunchConfig.",), + ) + CONFIG_INCOMPATIBLE_FLAGS = ( + "TVM-FFI export (--enable-tvm-ffi) and host-target export (--host-target) cannot be used together.", + ( + "Choose one export mode: use --enable-tvm-ffi for GPU executables, or --host-target for CPU-only.", + "For CPU-only export, remove the --enable-tvm-ffi flag.", + ), + ) + CONFIG_INVALID_HOST_TARGET = ( + "Invalid --host-target specification: {error}", + ( + "Use a known preset (e.g. 'linux-aarch64') or the long form: 'llvm -mtriple= [-mcpu=] [-mattr=]'.", + ), + ) + CONFIG_INVALID_OPT_LEVEL = ( + "Optimization level must be between 0 and 3, but got {val}.", + ( + "Use a valid optimization level: 0 (no optimization), 1, 2, or 3 (maximum optimization).", + ), + ) + CONFIG_INVALID_VALUE = ( + "Loop attribute `prefetch_stages` must be 0 or greater; got `{value}`.", + ("Use a non-negative integer for `prefetch_stages`.",), + ) + CONFIG_MALFORMED_COMPILE_OPTIONS = ( + "Compiler options string is malformed: '{options}'. Check the format and values.", + ( + "Verify the compile options string syntax and ensure values are properly formatted.", + ), + ) + CONFIG_MAX_THREADS_DYNAMIC = ( + "The `max_number_threads` launch config parameter must use static/Constexpr values, not runtime values.", + ("Use `Constexpr` annotation or a Python constant for max_number_threads.",), + ) + CONFIG_MISSING_NVDISASM = ( + "{vars} requires the 'nvdisasm' tool to write SASS output, but it is not found in PATH.", + ( + "Install the CUDA Toolkit from https://developer.nvidia.com/cuda-downloads; if CUDA is installed, add its bin directory to PATH: export PATH=/usr/local/cuda/bin:$PATH", + ), + ) + CONFIG_MISSING_TVM_FFI = ( + "TVM FFI is not installed, but TVM FFI support was requested.", + ("Install TVM FFI with `pip install apache-tvm-ffi`, or disable TVM FFI.",), + ) + CONFIG_OPTION_REQUIRES_VALUE = ( + "Option '{key}' requires a value. Provide it in the form {key}=.", + ("Specify a value for the option, e.g. {key}=.",), + ) + CONFIG_RESOLVER_INVALID_RETURN = ( + "The resolver function must return either a dict or None.", + ("Return a dict of string keys to string (or attribute) values, or None.",), + ) + CONFIG_UNKNOWN_COMPILE_OPTION = ( + "Unrecognized compile option: {option}. Please check that the option is spelled correctly.", + ("Verify the compile option name and use only supported options.",), + ) + CONFIG_UNSUPPORTED_ARCH = ( + "Unsupported GPU architecture '{arch}'. This kernel is only compatible with: {expected_archs}.", + ( + "Set `CUTE_DSL_ARCH` environment variable or pass a compatible `gpu_arch` to compile options.", + ), + ) + CONFIG_VERSION_CONFLICT = ( + "Cannot use both exact_version and min/max version range at the same time. Choose one approach.", + ( + 'Either use exact_version="X.Y" for a single version,', + "or use min_version and/or max_version for a range, but not both.", + ), + ) + CONFIG_VERSION_MISSING = ( + "When checking a version range, at least min_version or max_version must be provided.", + ( + 'Set min_version="X.Y" to check minimum,', + 'set max_version="X.Y" to check maximum,', + "or set both for a range between them.", + ), + ) + CONFIG_VERSION_RANGE_INVALID = ( + "The minimum version cannot be greater than or equal to the maximum version. Swap them or fix the range.", + ( + 'Ensure min_version < max_version, e.g., min_version="12.0", max_version="13.0"', + ), + ) + CONFIG_VERSION_REQUIRED = ( + "At least one version parameter is required.", + ( + 'Provide exact_version="X.Y" for a specific version,', + 'or min_version="X.Y" and/or max_version="X.Y" for a version range.', + ), + ) + CONFIG_VERSION_TYPE = ( + '{param_name} must be a version string (e.g., "12.3") or DSLCudaVersion, not {got_type}', + ( + 'Pass a version string like "12.3" or use DSLCudaVersion("{version_string}")', + ), + ) + # --- LAUNCH --- + LAUNCH_INVALID_CLUSTER = ( + "Launch cluster must have exactly 3 dimensions.", + ("Provide cluster as a 3D vector, e.g. `cluster=[1, 1, 1]`.",), + ) + LAUNCH_INVALID_DIMENSION = ( + "Launch {name} dimension at position {idx} must be an integer, but got {arg_type}.", + ("Ensure all elements of {name} are integers, e.g. `{name}=[1, 256, 1]`.",), + ) + LAUNCH_INVALID_FALLBACK = ( + "Launch fallback_cluster must have exactly 3 dimensions.", + ( + "Provide fallback_cluster as a 3D vector, e.g. `fallback_cluster=[1, 1, 1]`.", + ), + ) + LAUNCH_INVALID_GRID = ( + "Launch {name} must have at most 3 dimensions, but got {count}.", + ("Reduce the {name} to at most 3 elements, e.g. `{name}=[x, y, z]`.",), + ) + LAUNCH_MISSING_ARG = ( + "Required argument `{name}` is missing when launching kernel `{kernel_name}`.", + ("Pass a value for `{name}` to the kernel launch.",), + ) + # --- PHASE --- + PHASE_CONDITIONAL_NOT_DYNAMIC = ( + "The condition must be a {staged} value, not a {meta}.", + ("Ensure the condition is a runtime value (Boolean, Int32, etc.).",), + ) + PHASE_DYNAMIC_INDEX = ( + "Cannot use a {staged} value as a list index or for loop range in plain Python code", + ( + "Mark the loop as dynamic with @dynamic_expr or @range_dynamic", + "Decorate the parent function with @jit so the indexing happens in staged code where {staged} values are allowed", + ), + ) + PHASE_DYNAMIC_TO_STATIC_BOOL = ( + "Cannot use a {staged} value where a {meta} boolean is required (e.g. in plain Python control flow)", + ( + "Decorate the function with @jit and ensure preprocess is enabled so control flow is staged", + "Or move the boolean check inside a staged for/while/if block where {staged} values are allowed", + ), + ) + PHASE_REQUIRES_CONSTANT = ( + "{what} requires a value known at code-read time, but you gave it a {staged}.", + ( + "Make the value a Python constant (a literal or a value computed only from Python values).", + "If the value is only known at run time, use the runtime form instead (for example a runtime `if`/loop or a runtime assert).", + ), + ) + # --- SCOPE --- + SCOPE_CLOSURE_CAPTURE = ( + "Function `{func_name}` captures variable `{var_name}`, which is not supported in staged for/while/if.", + ( + "Pass `{var_name}` as an explicit function argument instead of relying on closure capture.", + "Define the function inside the loop/if, or refactor to avoid the closure.", + ), + ) + # --- TENSOR --- + TENSOR_ALREADY_ALLOCATED = ( + "This tensor is already allocated. Call deallocate() before allocating again.", + ( + "Call deallocate() before allocating the same tensor again.", + "Or create a new tensor instead of reusing the same one.", + ), + ) + TENSOR_FRAMEWORK_MANAGED = ( + "This tensor is managed by the framework and cannot be manually {action}.", + ("Let the framework manage this tensor; remove the manual {action} call.",), + ) + TENSOR_NOT_ALLOCATED = ( + "This tensor has no device memory yet -- it must be allocated before this operation.", + ( + "Call `allocate()` (and `copy_to_gpu()` if you need its contents) before using it here.", + ), + ) + # --- TYPE --- + TYPE_CLUSTER_NOT_INT = ( + "The {config_name} dimensions must be integers.", + ("Provide integer values for {config_name}.",), + ) + TYPE_CONDITIONAL_BRANCH_MISMATCH = ( + "The two branches of this conditional expression produce different types. Both branches must produce the same type.", + ("Make both branches produce the same type, converting explicitly if needed.",), + ) + TYPE_CONDITIONAL_VALUE_COUNT = ( + "The conditional expression branches must all return exactly one value.", + ("Ensure all branches (if/else) return exactly one value.",), + ) + TYPE_DEVICE_FUNC_RETURN_COUNT = ( + "The device function returned {count} value(s), but its return type expects exactly one.", + ("Return exactly one value matching the declared return type.",), + ) + TYPE_DEVICE_FUNC_RETURN_INVALID = ( + "The return type annotation must be a DSL type (like Float32, Int32) or a @native_struct.", + ( + "Use a DSL type (Float32, Int32, etc.) or @native_struct for the return annotation.", + ), + ) + TYPE_DEVICE_FUNC_RETURN_NONE = ( + "The device function has a return type annotation but returned None.", + ( + "Either add a return statement with a value, or remove the return type annotation.", + ), + ) + TYPE_DYNAMIC_EXPR_UNSUPPORTED = ( + "A value carried through this `{body_name}` is a plain {meta} (a `{py_type}`) " + "that cannot be turned into a {staged}, so the `{body_name}` cannot run at " + "run time.", + ( + "Make it a {staged} before the `{body_name}`, e.g. Int32(...), " + "Float32(...), or Boolean(...).", + "Or keep the `{body_name}` fully compile-time with `range_constexpr(...)` " + "/ `const_expr(...)` so the value can stay a {meta}.", + ), + ) + TYPE_ITER_ARGS_INVALID = ( + "The iter_args elements must be convertible to runtime values.", + ( + "Ensure all iter_args elements are DSL types or implement DynamicExpression.", + ), + ) + TYPE_LOOP_PARAM_NOT_INT = ( + "The loop parameter must be an integer, not a float.", + ("Use integer values for start, stop, and step in for_generate().",), + ) + TYPE_MISSING_ANNOTATION = ( + "`{arg_name}` must have a type annotation. Annotate it with `IRConst`, `IRValue`, `Pointer`, `TensorView`, or `Constexpr`.", + ( + "Add a type annotation to the parameter, e.g. `def foo({arg_name}: IRValue[Float32]):`.", + ), + ) + TYPE_PARAM_NOT_INT = ( + "The parameter must be an integer.", + ("Use an integer value for this parameter.",), + ) + TYPE_RETURN_TYPE_NOT_DSL = ( + "The return_types must all be DSL types (Int32, Float32, Boolean, etc.).", + ("Use DSL types for return_types.",), + ) + TYPE_UNSUPPORTED_C_ARG = ( + "Argument type {arg_type} is not supported for C header export. Only numeric types (Int32, Float32, etc.) are supported.", + ( + "Use a supported numeric type (Int8/16/32/64, Uint8/16/32/64, Float32/64, BFloat16, Float16, Boolean).", + "Check that your function signature has correct type annotations.", + ), + ) + TYPE_UNSUPPORTED_C_EXPORT = ( + "Type {type_name} cannot be exported to a C header. Only numeric types (Int32, Float32, etc.) and @native_struct classes are supported.", + ( + "Use a supported scalar type (Int8/16/32/64, Uint8/16/32/64, Float32/64, BFloat16, Float16, Boolean) or a @native_struct.", + "If you need to pass a custom type, define it as a @native_struct with supported fields.", + ), + ) + TYPE_UNSUPPORTED_TENSOR = ( + "`{func}` does not support this tensor type.", + ("Pass a tensor of a supported type, or convert it first.",), + ) + TYPE_WHILE_INPUTS_INVALID = ( + "The while loop inputs must be convertible to runtime values.", + ("Ensure all inputs are DSL types or implement DynamicExpression.",), + ) + # --- UNSUP --- + UNSUP_ARCH = ( + "The `vectorize` attribute requires compute capability 10.0 or higher; your target is `{arch}`.", + ("Remove the `vectorize` attribute, or build for sm_100 or higher.",), + ) + UNSUP_BUILTIN = ( + "The built-in function `{name}` is not allowed in compiled code.", + ( + "Avoid using `{name}` in your kernel code.", + "If you need dynamic code, refactor to avoid it or pre-compute at setup time.", + ), + ) + UNSUP_BUILTIN_FUNCTION = ( + "The built-in function '{fcn}' is not supported by the DSL.", + ("Use a supported built-in or equivalent DSL function.",), + ) + UNSUP_COMPARISON_OPERATOR = ( + "The comparison operator '{op}' is not supported.", + ("Use a supported comparison operator (==, !=, <, >, <=, >=).",), + ) + UNSUP_FSTRING_ALIGN = ( + "F-string alignment specifiers (<, >, ^) are not supported for {staged} values.", + ( + "Remove the alignment specifier or format the value before the f-string.", + 'Use a simple format like f"{value:d}" without alignment.', + ), + ) + UNSUP_FSTRING_FORMAT = ( + "F-string format specifiers for {staged} values can include at most one specifier.", + ( + 'Use a single format specifier, e.g., f"{value:.2f}".', + "Compute the formatted value into a variable first if you need multiple transformations.", + ), + ) + UNSUP_IN_OPERATOR = ( + "The `in` operator is not supported between these values.", + ("Use a supported comparison, or restructure the check to avoid `in`.",), + ) + + +class WarnId(_DiagMixin, enum.Enum): + """Author-facing **warnings** (non-fatal) -- a separate catalog from the + :class:`DiagId` errors so the two namespaces never collide. + + Same shape as ``DiagId`` (member name == stable code, value == + ``(message, fix)``) and rendered through the same block, but with a yellow + ``[Warning]`` header and a ``(warning: CODE)`` label. Raised via + :func:`report_warning` / ``DSLWarning(WarnId.X, ...)``. + """ + + PHASE_AUTO_PROMOTED_TO_STAGED = ( + "`{var}` (a {meta}, value {value}) was automatically turned into a {staged} " + "so it can be updated inside this for/while/if. If you read `{var}` AFTER " + "the for/while/if, you may still see its original fixed value, not the " + "updated one.", + ( + "To be safe, make `{var}` a {staged} before the for/while/if, e.g. " + "`{var} = {type}({var})`.", + ), + ) + + +def report_warning( + warn_id: "WarnId", + *, + filename: str | None = None, + lineno: int | None = None, + stacklevel: int = 2, + **fields: Any, +) -> None: + """Emit *warn_id* as a non-fatal :class:`DSLWarning` (does not raise). + + The warning analogue of ``raise DSLUserCodeError(DiagId.X, ...)``: it + renders through the one shared diagnostic block (yellow ``[Warning]`` + header) and flows through the standard ``warnings`` module so it is + filterable and deduplicated. + """ + import warnings + + from .common import DSLWarning # lazy: avoid an import cycle + + warnings.warn( + DSLWarning(warn_id, filename=filename, lineno=lineno, **fields), + stacklevel=stacklevel, + ) diff --git a/python/CuTeDSL/cutlass/base_dsl/dsl.py b/python/CuTeDSL/cutlass/base_dsl/dsl.py index 657f8f695..ed3551b55 100644 --- a/python/CuTeDSL/cutlass/base_dsl/dsl.py +++ b/python/CuTeDSL/cutlass/base_dsl/dsl.py @@ -24,10 +24,12 @@ import os import io import sys import errno +import tempfile import re import inspect import argparse import hashlib +from contextvars import ContextVar from contextlib import contextmanager from functools import lru_cache, wraps from collections import namedtuple, OrderedDict @@ -42,7 +44,7 @@ from typing import ( get_origin, get_args, ) -from collections.abc import Callable +from collections.abc import Callable, Iterable from types import SimpleNamespace, UnionType if TYPE_CHECKING: @@ -52,10 +54,10 @@ import warnings import threading from . import typing as t -from .env_manager import EnvironmentVarManager, is_cutlass_family_dsl_prefix -from .compiler import CompileOptions, LinkLibraries +from .env_manager import EnvironmentVarManager, dump_sass, is_cutlass_family_dsl_prefix +from .compiler import CompileOptions, CompilerDiagnosticError, LinkLibraries from .ast_helpers import DSLOptimizationWarning -from .common import register_env_manager +from .common import DSLRuntimeError, active_env_manager # ============================================================================= # Local module imports @@ -75,6 +77,7 @@ from .runtime.jit_arg_adapters import ( from .ast_preprocessor import DSLPreprocessor from .common import * +from .diagnostics import DiagId from .typing import ( Constexpr, get_c_pointers, @@ -83,7 +86,12 @@ from .typing import ( implements_dynamic_expression, implements_jit_argument, ) -from .._mlir_helpers.op import _set_enable_frame_filtering +from .._mlir_helpers.op import ( + _set_enable_frame_filtering, + _set_include_lib_frame, + get_verify_trace, + set_verify_trace, +) # ============================================================================= # MLIR modules @@ -171,10 +179,8 @@ def extract_mlir_values(obj: object, *, structured: bool = False) -> Any: for k, v in obj.__dict__.items(): res.extend(extract_mlir_values(v)) elif isinstance(obj, set): - raise DSLRuntimeError( - "Sets are not supported in extract_mlir_values to ensure order preservation", - context="The DSL attempted to generate JIT function argument(s) for an argument of type set but failed.", - suggestion="Consider using a list or tuple instead", + raise DSLUserCodeError( + DiagId.ARG_UNORDERED_CONTAINER, ) elif isinstance(obj, ir.Value): res = [obj] @@ -453,11 +459,7 @@ def extract_mlir_attributes(obj: object) -> list[Any]: # - exact-type primitives (int/float/bool/str); use `type(v) in (...)` # so that subclass instances carrying their own DSL hooks (e.g. # `numpy.float64`) still get recursed into - if ( - v is None - or isinstance(v, type) - or type(v) in (int, float, bool, str) - ): + if v is None or isinstance(v, type) or type(v) in (int, float, bool, str): continue ftype = f.type origin = get_origin(ftype) if not isinstance(ftype, str) else None @@ -466,10 +468,8 @@ def extract_mlir_attributes(obj: object) -> list[Any]: res.extend(extract_mlir_attributes(v)) # Can't call is_dynamic_expression as _is_dynamic_expression depends on extract_mlir_values elif isinstance(obj, set): - raise DSLRuntimeError( - "Sets are not supported in extract_mlir_values to ensure order preservation", - context="The DSL attempted to generate JIT function argument(s) for an argument of type set but failed.", - suggestion="Consider using a list or tuple instead", + raise DSLUserCodeError( + DiagId.ARG_UNORDERED_CONTAINER, ) elif isinstance(obj, ir.Value): res = [ir.DictAttr.get({})] @@ -550,10 +550,8 @@ def new_from_mlir_values(obj: Any, values: Any, *, structured: bool = False) -> values = values[n_items:] return ns elif isinstance(obj, set): - raise DSLRuntimeError( - "Sets are not supported in new_from_mlir_values to ensure order preservation", - context="The DSL attempted to generate JIT function argument(s) for an argument of type set but failed.", - suggestion="Consider using a list or tuple instead", + raise DSLUserCodeError( + DiagId.ARG_UNORDERED_CONTAINER, ) elif is_dynamic_expression(obj): if len(values) == 0: @@ -592,6 +590,8 @@ class DSLSingletonMeta(type): _instances: ClassVar[dict] = {} _lock: ClassVar[threading.Lock] = threading.Lock() + _optimization_warnings_enabled: ClassVar[bool] = False + _stacktrace_filter_disabled: ClassVar[bool] = False def __call__(cls, *args: Any, **kwargs: Any) -> Any: with cls._lock: @@ -688,7 +688,6 @@ class BaseDSL(metaclass=DSLSingletonMeta): self.num_kernels: int = 0 # Read environment variables self.envar: EnvironmentVarManager = self._env_class(self.name) - register_env_manager(self.envar) self.enable_preprocessor: bool = preprocess # This cache uses hash of original ir and env as key, allows dump/load to/from file. Enabled by default self.jit_cache: JitCacheDict = JitCacheDict( @@ -699,8 +698,21 @@ class BaseDSL(metaclass=DSLSingletonMeta): self.device_jit_decorator_name: str = f"@{BaseDSL.kernel.__name__}" # set warning - if not self.envar.enable_optimization_warnings: - # By default, optimization warnings are disabled + # + # Multiple DSL singletons may be initialized with different env var + # prefixes. The optimization-warning filter is global, so a later + # DSL instance must not blindly add an "ignore" filter when an + # earlier instance already opted *in* via its own env var (e.g. + # CUTE_DSL_ENABLE_OPTIMIZATION_WARNINGS=1). + if self.envar.enable_optimization_warnings: + if not DSLSingletonMeta._optimization_warnings_enabled: + DSLSingletonMeta._optimization_warnings_enabled = True + warnings.filters = [ + f + for f in warnings.filters + if not (f[0] == "ignore" and f[2] is DSLOptimizationWarning) + ] + elif not DSLSingletonMeta._optimization_warnings_enabled: warnings.filterwarnings("ignore", category=DSLOptimizationWarning) if self.envar.warnings_as_errors: warnings.filterwarnings("error") @@ -715,6 +727,13 @@ class BaseDSL(metaclass=DSLSingletonMeta): self.launch_inner_count: int = 0 # initialize default compile options self.compile_options: CompileOptions = CompileOptions() + # Synchronous callbacks run after tracing and before module hashing. + # Signature: hook(owner, module, function_name). Hooks may mutate the + # finalized ir.Module; hook exceptions are wrapped by the caller. + self._trace_finalize_hooks: list[Callable[[Any, ir.Module, str], None]] = [] + self._scoped_trace_finalize_hooks: ContextVar[ + tuple[Callable[[Any, ir.Module, str], None], ...] + ] = ContextVar(f"{self.name}_trace_finalize_hooks", default=()) if preprocess: preprocessor: DSLPreprocessor = DSLPreprocessor(dsl_package_name) @@ -730,7 +749,16 @@ class BaseDSL(metaclass=DSLSingletonMeta): self.cache_misses: int = 0 # Hook excepthook - if self.envar.filter_stacktrace: + # + # Guard: if a prior DSL instance explicitly disabled filtering + # (e.g. CUTE_DSL_FILTER_STACKTRACE=0), do not let a later + # instance with a different prefix re-enable it. + if not self.envar.filter_stacktrace: + DSLSingletonMeta._stacktrace_filter_disabled = True + if ( + self.envar.filter_stacktrace + and not DSLSingletonMeta._stacktrace_filter_disabled + ): origin_excepthook = sys.excepthook module_dir = walk_to_top_module(os.path.dirname(os.path.abspath(__file__))) @@ -838,14 +866,15 @@ class BaseDSL(metaclass=DSLSingletonMeta): BaseDSL._preprocess_and_replace_code(func) custom_name = getattr(jit_wrapper, "_name_prefix", None) - if custom_name: - return getattr(func._dsl_object, executor_name)( - func, *args, **kwargs, _name_prefix=custom_name - ) - else: - return getattr(func._dsl_object, executor_name)( - func, *args, **kwargs - ) + with active_env_manager(func._dsl_object.envar): + if custom_name: + return getattr(func._dsl_object, executor_name)( + func, *args, **kwargs, _name_prefix=custom_name + ) + else: + return getattr(func._dsl_object, executor_name)( + func, *args, **kwargs + ) def set_name_prefix(name: str) -> None: jit_wrapper._name_prefix = name # type: ignore[attr-defined] @@ -910,23 +939,6 @@ class BaseDSL(metaclass=DSLSingletonMeta): return pipeline return None - @staticmethod - def log_additions( - func_type: Any, operands: Any = None, types: Any = None, arg_attrs: Any = None - ) -> None: - if operands is not None and operands != []: - log().debug( - f"Added {func_type} operands: [%s]", ", ".join(map(str, operands)) - ) - if types is not None: - log().debug( - f"Added {func_type} arg_types: [%s]", ", ".join(map(str, types)) - ) - if arg_attrs is not None: - log().debug( - f"Added {func_type} arg_attrs: [%s]", ", ".join(map(str, arg_attrs)) - ) - def mangle_name( self, function_name: str, args: tuple[Any, ...], sig: inspect.Signature ) -> str: @@ -1041,7 +1053,6 @@ class BaseDSL(metaclass=DSLSingletonMeta): else: ir_arg = ir_arg[0] - self.log_additions(ir_arg) return ir_arg, iv_block_args fop_args = list(fop.regions[0].blocks[0].arguments) @@ -1196,6 +1207,7 @@ class BaseDSL(metaclass=DSLSingletonMeta): ) annotation_markers = [] for sub_ty in candidate_sub_types: + # Annotated[T, marker] at the top — bare annotated argument. ty, *markers = ( get_args(sub_ty) if get_origin(sub_ty) is Annotated else (sub_ty,) ) @@ -1203,6 +1215,26 @@ class BaseDSL(metaclass=DSLSingletonMeta): annotation_markers = markers break + # List[Annotated[T, marker]] / Tuple[Annotated[T, marker], ...]: + # peel one container layer so per-element markers (e.g. + # cuda.grid_constant) survive list/tuple args. + container_origin = get_origin(sub_ty) + if ( + container_origin in (list, tuple) + and isinstance(arg, (list, tuple)) + and arg + ): + container_args = get_args(sub_ty) + if container_args and get_origin(container_args[0]) is Annotated: + inner_ty, *inner_markers = get_args(container_args[0]) + if ( + inner_markers + and isinstance(inner_ty, type) + and all(isinstance(e, inner_ty) for e in arg) + ): + annotation_markers = inner_markers + break + log().debug("Processing [%d] Argument [%s : %s]", i, arg_name, spec_ty) # Implicitly convert into Numeric type if possible @@ -1249,25 +1281,48 @@ class BaseDSL(metaclass=DSLSingletonMeta): jit_arg_attr.extend(extract_mlir_attributes(arg)) if not jit_arg_type or not jit_exec_arg: - # when it is compile only, we don't have to prepare the executable arguments. - if (is_host and (compile_only or implements_jit_argument(arg))) or ( - not is_host and implements_dynamic_expression(arg) + # Compile-only placeholders can provide MLIR signature types + # without runtime execution arguments, e.g. FakeTensor. + is_compile_only_type_placeholder = ( + is_host and compile_only and jit_arg_type and not jit_exec_arg + ) + if ( + is_compile_only_type_placeholder + or (is_host and implements_jit_argument(arg)) + or (not is_host and implements_dynamic_expression(arg)) ): pass else: - raise DSLRuntimeError( - f"failed to generate argument #{i + 1} ({arg_name}) for JIT function '{function_name}'.", - context={ - f"Argument {arg_name}": "The DSL attempted to convert it into Dynamic Expression (aka MLIR values) but failed.", - "Call-site argument value": arg, - "Call-site argument type": type(arg), - }, - suggestion=f"Consider annotating the argument with `{arg_name} : Constexpr` " - "if it's a value known at compile-time. " - f"Otherwise, implement the {'`JitArgument`' if is_host else '`DynamicExpression`'} " - f"protocol or register a custom JIT argument adapter for type `{type(arg)}` to " - "enable dynamic value conversion at runtime.", - ) + if ( + is_host + and compile_only + and not implements_jit_argument(arg) + and not is_argument_constexpr( + arg, spec_ty, arg_name, i, func + ) + ): + warning_msg, warning_suggestions = ( + DiagId.ARG_UNSUPPORTED_TYPE.fill( + num=i + 1, + arg_name=arg_name, + arg_type=type(arg), + phase_label="JitArgument", + function_name=function_name, + ) + ) + warning_msg = " ".join((warning_msg, *warning_suggestions)) + warnings.warn(warning_msg, UserWarning, stacklevel=3) + else: + raise DSLUserCodeError( + DiagId.ARG_UNSUPPORTED_TYPE, + num=i + 1, + arg_name=arg_name, + arg_type=type(arg), + phase_label=( + "JitArgument" if is_host else "DynamicExpression" + ), + function_name=function_name, + ) if jit_arg_type is not None: assert jit_exec_arg is not None and jit_arg_attr is not None @@ -1313,9 +1368,6 @@ class BaseDSL(metaclass=DSLSingletonMeta): compile_only=compile_only, ) - log().debug("Execution Arguments: %s", ", ".join(map(str, exe_args))) - log().debug("Types: %s", ", ".join(map(str, types))) - assert ( compile_only or self.envar.enable_tvm_ffi or len(exe_args) == len(types) ), "expects the same number of arguments and function parameters" @@ -1351,13 +1403,18 @@ class BaseDSL(metaclass=DSLSingletonMeta): dim = [dim] if len(dim) > 3: - raise DSLRuntimeError( - f"Expected {name} dimension to be less than or equal to 3, but got {len(dim)}" + raise DSLUserCodeError( + DiagId.LAUNCH_INVALID_GRID, + name=name, + count=len(dim), ) for idx, e in enumerate(dim): if not isinstance(e, (Integer, int)): - raise DSLRuntimeError( - f"Expected integer for {name} dimension at index {idx}, but got {type(e)}" + raise DSLUserCodeError( + DiagId.LAUNCH_INVALID_DIMENSION, + name=name, + idx=idx, + arg_type=type(e), ) # Pad with 1s to 3-dim vector for grid or block dimensions @@ -1371,13 +1428,13 @@ class BaseDSL(metaclass=DSLSingletonMeta): if self.cluster is None: self.cluster = [None, None, None] elif len(self.cluster) != 3: - raise DSLRuntimeError("Expect 3d cluster!") + raise DSLUserCodeError(DiagId.LAUNCH_INVALID_CLUSTER) self.has_fallback_cluster = self.fallback_cluster is not None if self.fallback_cluster is None: self.fallback_cluster = [None, None, None] elif len(self.fallback_cluster) != 3: - raise DSLRuntimeError("Expect 3d fallback_cluster!") + raise DSLUserCodeError(DiagId.LAUNCH_INVALID_FALLBACK) def has_max_number_threads(self) -> bool: """Check if max_number_threads is given by user""" @@ -1400,9 +1457,24 @@ class BaseDSL(metaclass=DSLSingletonMeta): args, _ = parser.parse_known_args() ctx = ir.Context.current + compiler_opt = os.environ.get("CUTE_DSL_COMPILER_OPT", "") + if args.diagnostic is None and ( + self.compile_options.collect_compiler_diagnostics + or "diagnostic" in compiler_opt + ): + return def callback(d: Any) -> None: - print(f" [{self.name} Diagnostic] : {d.message}") + if self.compile_options.collect_compiler_diagnostics: + return + try: + message = d.message + except BaseException: + try: + message = str(d) + except BaseException: + return + print(f" [{self.name} Diagnostic] : {message}") ctx.attach_diagnostic_handler(callback) @@ -1483,6 +1555,12 @@ class BaseDSL(metaclass=DSLSingletonMeta): pipeline, shared_libs=shared_libs, arch=self.envar.arch, + remark_filter=self.compile_options.remark_filter, + warnings_filter=self.compile_options.warnings_filter, + remark_output=self.compile_options.remark_output, + collect_compiler_diagnostics=( + self.compile_options.collect_compiler_diagnostics + ), enable_debug_info=enable_debug_info, ) @@ -1495,13 +1573,27 @@ class BaseDSL(metaclass=DSLSingletonMeta): print(redirect_stdout.getvalue(), file=sys.stdout, end="") print(redirect_stderr.getvalue(), file=sys.stderr, end="") + sass_path = self.compile_options.dump_sass_path + if sass_path: + self._dump_sass_artifact(sass_path) + return kernel + except CompilerDiagnosticError: + raise except Exception as e: raise DSLRuntimeError("🧊🧊🧊 ICE 🧊🧊🧊", cause=e) finally: pass + def _dump_sass_artifact(self, sass_path: str) -> None: + cubin_path = self.compile_options.full_cubin_path + assert cubin_path is not None + from .compiler import NvdisasmOptions + + flags = self.compile_options.options[NvdisasmOptions].value + dump_sass(cubin_path, sass_path, flags) + def preprocess_pipeline(self, pipeline: str, arch: str) -> str: options = { self.pass_sm_arch_name: arch, @@ -1579,6 +1671,94 @@ class BaseDSL(metaclass=DSLSingletonMeta): ) return module_hash + def register_trace_finalize_hook( + self, hook: Callable[[Any, ir.Module, str], None] + ) -> None: + """Register a callback run after tracing and before hashing the module. + + The hook is called synchronously as ``hook(owner, module, + function_name)``, where ``owner`` is this DSL instance, ``module`` is + the finalized ``ir.Module``, and ``function_name`` identifies the trace. + Hooks may inspect or annotate the module before its cache key is + computed. Passing ``None`` raises ``DSLRuntimeError``; registering the + same hook object more than once is ignored. + """ + if hook is None: + raise DSLRuntimeError("Trace finalize hook must not be None.") + if not callable(hook): + raise DSLRuntimeError("Trace finalize hook must be callable.") + if hook not in self._trace_finalize_hooks: + self._trace_finalize_hooks.append(hook) + + @contextmanager + def trace_finalize_hooks( + self, + hooks: Callable[[Any, ir.Module, str], None] + | Iterable[Callable[[Any, ir.Module, str], None]], + ) -> Generator[None, Any, None]: + """Temporarily register trace-finalize hooks in the current context. + + Args: + hooks: A single hook or iterable of hooks. Each hook is called as + ``hook(owner, module, function_name)``, where ``owner`` is this + DSL instance, ``module`` is the finalized ``ir.Module``, and + ``function_name`` identifies the trace. + + Scoped hooks are stored in ``_scoped_trace_finalize_hooks`` for the + duration of the context, preserving order and ignoring duplicates. The + context manager restores the previous hook state when the ``with`` block + exits. + + Raises: + DSLRuntimeError: If ``hooks`` is neither callable nor iterable, or + if any hook entry is ``None`` or not callable. + """ + scoped_hooks: tuple[Callable[[Any, ir.Module, str], None], ...] + if callable(hooks): + scoped_hooks = (hooks,) + else: + try: + scoped_hooks = tuple(hooks) + except TypeError as e: + raise DSLRuntimeError( + "Trace finalize hooks must be callable or iterable." + ) from e + + for hook in scoped_hooks: + if hook is None: + raise DSLRuntimeError("Trace finalize hook must not be None.") + if not callable(hook): + raise DSLRuntimeError("Trace finalize hook must be callable.") + + current_hooks = self._scoped_trace_finalize_hooks.get() + combined_hooks = list(current_hooks) + for hook in scoped_hooks: + if hook not in combined_hooks: + combined_hooks.append(hook) + token = self._scoped_trace_finalize_hooks.set(tuple(combined_hooks)) + try: + yield + finally: + self._scoped_trace_finalize_hooks.reset(token) + + def _run_trace_finalize_hooks(self, module: ir.Module, function_name: str) -> None: + hooks = list(self._trace_finalize_hooks) + for hook in self._scoped_trace_finalize_hooks.get(): + if hook not in hooks: + hooks.append(hook) + + for hook in hooks: + try: + hook(self, module, function_name) + except Exception as e: + hook_name = getattr( + hook, "__qualname__", getattr(hook, "__name__", repr(hook)) + ) + # DSLRuntimeError inherits DSLBaseError, which formats ``cause``. + raise DSLRuntimeError( + f"Trace finalize hook failed: {hook_name}", cause=e + ) from e + def build_module(self, module: ir.Module, function_name: str) -> ir.Module: """ Build the MLIR module, verify and return the module @@ -1712,12 +1892,75 @@ class BaseDSL(metaclass=DSLSingletonMeta): module, result = self.profiler(build_ir_module)() else: module, result = build_ir_module() + self._run_trace_finalize_hooks(module, function_name) module_hash = self.get_module_hash(module, function_name) module = self.build_module(module, function_name) return module, module_hash, result + def _runtime_can_run_compiled( + self, runtime_arch: str | None, compiled_arch: str | None + ) -> bool: + """Whether the runtime arch can execute a binary built for ``compiled_arch``. + + Base behavior requires an exact arch match. DSLs with a richer arch model + (e.g. CuTe DSL) override this to allow family-portable targets. + """ + return runtime_arch == compiled_arch + + def _lookup_jit_entry(self, engine: Any, function_name: str) -> Any: + """Resolve the JIT'd entry symbol, turning a failed lookup into a clear error. + + On the MLIR ExecutionEngine path, looking up the entry symbol lazily + materializes it. If the runtime library is resolvable as a file but does + not export the symbols the module needs (e.g. an incompatible + ``{name}_LIBS`` ``.so``), ORC reports ``Symbols not found: [...]`` on the + process stderr and the lookup returns null, which the upstream wrapper + surfaces as an opaque ``RuntimeError: Unknown function ``. The + real diagnostic is captured from the duplicated stderr fd here and + re-raised with a hint pointing at the runtime-library environment + variable. + """ + lookup = ( + self.profiler(engine.lookup) + if self.envar.jit_time_profiling + else engine.lookup + ) + + sys.stderr.flush() + saved_stderr_fd = os.dup(2) + with tempfile.TemporaryFile() as capture: + os.dup2(capture.fileno(), 2) + try: + try: + capi_func = lookup(function_name) + finally: + sys.stderr.flush() + os.dup2(saved_stderr_fd, 2) + os.close(saved_stderr_fd) + capture.seek(0) + captured = capture.read().decode("utf-8", "replace") + except RuntimeError as e: + raise self._jit_lookup_error(captured) from e + if captured: + sys.stderr.write(captured) + sys.stderr.flush() + return capi_func + + def _jit_lookup_error(self, captured_stderr: str) -> DSLUserCodeError: + libs_env = f"{self.name}_LIBS" + diagnostic = captured_stderr.strip() + return DSLUserCodeError( + f"The {self.name} runtime library could not be resolved or is " + "incompatible.", + suggestion=( + f"Check that {libs_env} points to a runtime library that " + "exports the symbols this module needs." + ), + context=diagnostic or None, + ) + def compile_and_cache( self, module: ir.Module, @@ -1742,8 +1985,12 @@ class BaseDSL(metaclass=DSLSingletonMeta): if not self.compile_options.gpu_arch else self.compile_options.gpu_arch ) - # If no gpu kernels or compile_gpu_arch is same as the arch from the environment variable, generate a JIT engine. Otherwise, only do the compilation. - gen_jit_engine = self.num_kernels == 0 or compile_gpu_arch == self.envar.arch + # Build a JIT engine when the runtime arch can execute a binary built for compile_gpu_arch; + # a genuine cross-compile gets no engine and must be exported. _runtime_can_run_compiled is + # overridden per DSL: the base requires an exact arch match, CuTe DSL adds family-portability. + gen_jit_engine = self.num_kernels == 0 or self._runtime_can_run_compiled( + self.envar.arch, compile_gpu_arch + ) if no_jit_engine: gen_jit_engine = False # Preprocess the pipeline. @@ -1791,11 +2038,23 @@ class BaseDSL(metaclass=DSLSingletonMeta): self.profiler(self.compiler_provider.compile)( module, pipeline, + remark_filter=self.compile_options.remark_filter, + warnings_filter=self.compile_options.warnings_filter, + remark_output=self.compile_options.remark_output, + collect_compiler_diagnostics=( + self.compile_options.collect_compiler_diagnostics + ), ) else: self.compiler_provider.compile( module, pipeline, + remark_filter=self.compile_options.remark_filter, + warnings_filter=self.compile_options.warnings_filter, + remark_output=self.compile_options.remark_output, + collect_compiler_diagnostics=( + self.compile_options.collect_compiler_diagnostics + ), ) engine = None else: @@ -1816,10 +2075,7 @@ class BaseDSL(metaclass=DSLSingletonMeta): if gen_jit_engine else None ) - if self.envar.jit_time_profiling: - capi_func = self.profiler(engine.lookup)(function_name) if engine else None - else: - capi_func = engine.lookup(function_name) if engine else None + capi_func = self._lookup_jit_entry(engine, function_name) if engine else None fn = func_type( module, @@ -1833,6 +2089,7 @@ class BaseDSL(metaclass=DSLSingletonMeta): jit_function_artifacts=JitFunctionArtifacts( PTX=self.compile_options.full_ptx_path, CUBIN=self.compile_options.full_cubin_path, + SASS=self.compile_options.dump_sass_path, MLIR=( str(self.dump_mlir_path) if (self.envar.keep_ir or self.envar.keep_ir_clean) @@ -1935,10 +2192,22 @@ class BaseDSL(metaclass=DSLSingletonMeta): _loc_tb_ctx = None if _loc_tb_depth > 0: try: - _loc_tb_ctx = ir.loc_tracebacks(max_depth=_loc_tb_depth) + # New LLVM (>= upstream PR #192310 cherry-pick): NAMELOC_WRAP + # preserves Scope/WarpScope NameLocs around the generated + # traceback chain so profiling annotations survive to PTX/SASS. + _loc_tb_ctx = ir.loc_tracebacks( + max_depth=_loc_tb_depth, + current_loc_actn=ir.CurrentLocAction.NAMELOC_WRAP, + ) + except (TypeError, AttributeError): + # Older LLVM without the composition kwargs / enums: + # still give the user basic tracebacks. + try: + _loc_tb_ctx = ir.loc_tracebacks(max_depth=_loc_tb_depth) + except (ValueError, TypeError, AttributeError): + _loc_tb_ctx = None + if _loc_tb_ctx is not None: _loc_tb_ctx.__enter__() - except (ValueError, TypeError, AttributeError): - pass try: # Convert input arguments to MLIR arguments @@ -1992,10 +2261,30 @@ class BaseDSL(metaclass=DSLSingletonMeta): buf = io.BytesIO() module.operation.write_bytecode(buf) - return _cutlass_ir.PreCompiledMlirArtifact.from_bitcode( + artifact = _cutlass_ir.PreCompiledMlirArtifact.from_bitcode( buf.getvalue() ) + # Metadata is the whole point of the precompiled-MLIR + # artifact (it is the calling-convention contract the + # cutlass_compiler ABI wrappers are generated from). Let any + # failure from build_function_metadata propagate with its + # real cause and traceback rather than wrapping it -- the + # raised error already identifies the offending parameter, + # and swallowing/re-wrapping only obscures it. + from cutlass.cute.metadata import build_function_metadata + + artifact.metadata.append( + build_function_metadata( + function_name=function_name, + signature=sig, + args=args, + kwonlyargs=kwonlyargs, + display_name=funcBody.__name__, + ) + ) + return artifact + # dryrun is used to only generate IR if self.envar.dryrun: return result @@ -2077,6 +2366,16 @@ class BaseDSL(metaclass=DSLSingletonMeta): def run_preprocessor( self, original_function: Any, + ) -> Any: + # Preprocessing runs before jit_wrapper enters its call-time context. + with active_env_manager(self.envar): + return self._run_preprocessor_impl( + original_function, + ) + + def _run_preprocessor_impl( + self, + original_function: Any, ) -> Any: function_name = original_function.__name__ self.funcBody = original_function @@ -2103,7 +2402,6 @@ class BaseDSL(metaclass=DSLSingletonMeta): ) original_function._preprocessed = True - return preprocessor_session.exec( original_function.__name__, original_function, @@ -2163,8 +2461,10 @@ class BaseDSL(metaclass=DSLSingletonMeta): param.default is inspect.Parameter.empty and param.name not in bound_args.arguments ): - raise DSLRuntimeError( - f"Missing required argument in `{function_name}`: '{param.name}'" + raise DSLUserCodeError( + DiagId.CALL_MISSING_ARG, + name=param.name, + function_name=function_name, ) return has_varargs @@ -2303,6 +2603,7 @@ class BaseDSL(metaclass=DSLSingletonMeta): if not no_cache and ( self.envar.keep_ptx or self.envar.keep_cubin + or self.envar.keep_sass ): no_cache = True self.print_warning("Cache is disabled as user wants to generate PTX/ASM.") @@ -2336,10 +2637,18 @@ class BaseDSL(metaclass=DSLSingletonMeta): function_name = f"{func_name_prefix}_{function_name}" self.compile_options.apply_envar_settings(self.envar, function_name) - if not self.compile_options.generate_line_info: + track_source_locations = ( + self.compile_options.generate_line_info + ) + if not track_source_locations: self.decorator_location = None - # Enable frame filtering if line info is enabled - _set_enable_frame_filtering(self.compile_options.generate_line_info) + # Enable frame filtering when diagnostics or line info need user frames. + _set_enable_frame_filtering(track_source_locations) + # Debug mode: attribute ops to the closest (library) frame and turn on + # trace-time MLIR op verification so malformed ops fail at the call + # site. An explicit CUTE_DSL_VERIFY_TRACE=1 is preserved. + _set_include_lib_frame(self.envar.debug) + set_verify_trace(get_verify_trace() or self.envar.debug) return self._CompilationSetup( function_name=function_name, @@ -2365,6 +2674,14 @@ class BaseDSL(metaclass=DSLSingletonMeta): 5. Operator overloading (a + b --> arith.addi a, b) 6. Generates GPU kernel function with GPU module and kernel attributes baked """ + # Keep this guard even though jit_wrapper also enters the env context: + # compile/device paths may call _func directly. + with active_env_manager(self.envar): + return self._func_impl(funcBody, *args, **kwargs) + + def _func_impl( + self, funcBody: Callable[..., Any], *args: Any, **kwargs: Any + ) -> Any: if ir.Context.current is None: pass elif ir.InsertionPoint.current is not None: @@ -2506,6 +2823,13 @@ class BaseDSL(metaclass=DSLSingletonMeta): :return: A named tuple containing the launch function and function return, the kernel name and the MLIR module. """ + # Kernel generation may be reached from driver/device-only paths. + with active_env_manager(self.envar): + return self._generate_kernel_module_impl(kernel_generator) + + def _generate_kernel_module_impl( + self, kernel_generator: Callable[..., Any] + ) -> tuple[Any, str, ir.Module]: ret = None with ir.Context(), self.get_ir_location() as loc: @@ -2555,10 +2879,6 @@ class BaseDSL(metaclass=DSLSingletonMeta): ) ) - log().debug("Final kernel_operands: %s", ", ".join(map(str, kernel_operands))) - log().debug("Final kernel_arg_types: %s", ", ".join(map(str, kernel_arg_types))) - log().debug("Final kernel_arg_attrs: %s", ", ".join(map(str, kernel_arg_attrs))) - assert len(kernel_operands) == len(kernel_arg_types) == len(kernel_arg_attrs), ( "Size of kernel_operands, kernel_arg_types and kernel_arg_attrs must be equal" ) @@ -2617,8 +2937,10 @@ class BaseDSL(metaclass=DSLSingletonMeta): for name in argNames: value = kwargs.pop(name, None) if assertIfNone and value is None: - raise DSLRuntimeError( - f"{name} is required for {kernel_name}" + raise DSLUserCodeError( + DiagId.LAUNCH_MISSING_ARG, + name=name, + kernel_name=kernel_name, ) extracted.append(value) @@ -2749,7 +3071,8 @@ class BaseDSL(metaclass=DSLSingletonMeta): arch = self.get_arch_enum() if not criterion(arch): - raise DSLRuntimeError( - f"invalid arch, expected one of {Arch.filter(criterion)}, but got {arch}.", - suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture", + raise DSLUserCodeError( + DiagId.CONFIG_UNSUPPORTED_ARCH, + arch=arch, + expected_archs=Arch.filter(criterion), ) diff --git a/python/CuTeDSL/cutlass/base_dsl/env_manager.py b/python/CuTeDSL/cutlass/base_dsl/env_manager.py index 0107ab9d1..e54a4a686 100644 --- a/python/CuTeDSL/cutlass/base_dsl/env_manager.py +++ b/python/CuTeDSL/cutlass/base_dsl/env_manager.py @@ -28,7 +28,8 @@ from pathlib import Path from functools import lru_cache from ..base_dsl.runtime.cuda import get_compute_capability_major_minor -from .common import DSLRuntimeError +from .common import DSLUserCodeError +from .diagnostics import DiagId from .utils.logger import log from .cache_helpers import get_default_file_dump_root @@ -339,6 +340,79 @@ def _find_cuda_home() -> str | None: return cuda_home + +@lru_cache(maxsize=1) +def _find_nvdisasm_binary() -> str: + """Return absolute path to the bundled nvdisasm binary.""" + from importlib import metadata + + try: + dist = metadata.distribution("nvidia-cuda-nvdisasm") + except metadata.PackageNotFoundError: + dist = None + if dist is not None and dist.files is not None: + for entry in dist.files: + if entry.name == "nvdisasm": + binpath = Path(str(dist.locate_file(entry))) + if binpath.is_file(): + return str(binpath) + raise DSLUserCodeError( + "nvdisasm binary not found inside the nvidia-cuda-nvdisasm wheel.", + suggestion="\n".join( + [ + "nvidia-cuda-nvdisasm is a runtime dependency of nvidia-cutlass-dsl", + "and should have been installed automatically. Try:", + " • pip install --force-reinstall nvidia-cuda-nvdisasm", + ] + ), + ) + + +def _check_nvdisasm_available( + prefix: str, + keep_sass: bool, +) -> None: + """Check that nvdisasm is available when SASS output is requested.""" + try: + _find_nvdisasm_binary() + return + except DSLUserCodeError: + pass + enabled_vars = [] + if keep_sass: + enabled_vars.append(f"{prefix}_KEEP_SASS") + vars_str = " and ".join(enabled_vars) + raise DSLUserCodeError(DiagId.CONFIG_MISSING_NVDISASM, vars=vars_str) + + +def dump_sass( + cubin_path: str, + sass_path: str | None, + flags: str, +) -> None: + """Disassemble a CUBIN file into SASS. + + If sass_path is None, the SASS is written to stderr. + Otherwise it is written to the given file path. + """ + import shlex + import subprocess + + if not cubin_path or not os.path.exists(cubin_path): + raise DSLUserCodeError( + f"Cannot dump SASS: CUBIN file does not exist at {cubin_path!r}." + ) + + nvdisasm = _find_nvdisasm_binary() + tokens = [nvdisasm, *shlex.split(flags), cubin_path] + + if sass_path: + with open(sass_path, "w") as sass_file: + subprocess.run(tokens, stdout=sass_file, check=True) + else: + subprocess.run(tokens, stdout=sys.stderr, check=True) + + def get_prefix_dsl_libs(prefix: str) -> str | None: """ Returns get_str_env_var('{prefix}_LIBS') if set. @@ -442,13 +516,22 @@ class EnvironmentVarManager(LogEnvironmentManager): ir-debug — Raw IR before any passes ptx — PTX assembly cubin — CUBIN binary + sass — SASS disassembly all — all of the above Example: CUTE_DSL_KEEP=ir,ptx # Deprecated — use [DSL_NAME]_KEEP instead: - [DSL_NAME]_KEEP_IR: (deprecated) use KEEP=ir-debug - [DSL_NAME]_KEEP_PTX: (deprecated) use KEEP=ptx - [DSL_NAME]_KEEP_CUBIN: (deprecated) use KEEP=cubin + - [DSL_NAME]_KEEP_SASS: (deprecated) use KEEP=sass Other options: + - [DSL_NAME]_DEBUG: Master debug switch for DSL developers (default: False). + When True, raises the default of LINEINFO, SHOW_STACKTRACE (and disables + FILTER_STACKTRACE) and ENABLE_OPTIMIZATION_WARNINGS; + attributes MLIR ops to the closest (library) frame; and runs full + per-launch argument validation. These defaults remain independently + overridable by their own env vars. Trace-time MLIR op verification + (CUTE_DSL_VERIFY_TRACE) is always enabled while debug mode is on. - [DSL_NAME]_SHOW_STACKTRACE: Show full stack traces on failure (default: False) - [DSL_NAME]_LINEINFO: Compile with `--lineinfo` enabling developer tools such as the profiler and debugger (default: False) - [DSL_NAME]_LOG_LEVEL: Logging level to set, for LOG_TO_CONSOLE or LOG_TO_FILE (default: 1). @@ -465,9 +548,16 @@ class EnvironmentVarManager(LogEnvironmentManager): - [DSL_NAME]_ENABLE_TVM_FFI: Enable TVM-FFI or not (default: False) - [DSL_NAME]_LOC_TRACEBACKS: Maximum depth of location tracebacks (default: 0) - [DSL_NAME]_COMPILER_OPT: Compact compiler option string (default: ""). - Controls compiler passes and diagnostic checks. Forms accepted: + Errors always show and fail compilation (no flag needed); warnings and + remarks are opt-in and non-fatal. A {} selector shows only that + category; bare shows all categories. Forms accepted: + warnings — show all warnings + warnings{nvvm} — show only nvvm-category warnings + remarks — show all remarks + remarks{nvvm} — show only nvvm (sync) remarks iket — enable IKET (In-Kernel Event Tracing) instrumentation Examples: + CUTE_DSL_COMPILER_OPT="warnings{nvvm}" CUTE_DSL_COMPILER_OPT="iket" The same option strings are accepted by cute.compile(..., options=...). @@ -476,14 +566,25 @@ class EnvironmentVarManager(LogEnvironmentManager): def __init__(self, prefix: str = "DSL") -> None: super().__init__(prefix) + # Master debug switch for DSL developers. When True, it raises the + # default of a curated set of diagnostic/correctness settings below + # (lineinfo, stacktrace, optimization warnings, IR verification). + # Each of those settings remains independently overridable by its own + # env var, so debugging mode only changes their defaults. + self.debug = get_bool_env_var(f"{prefix}_DEBUG", False) + # Printing options self.print_after_preprocessor = get_bool_env_var( f"{prefix}_PRINT_AFTER_PREPROCESSOR", False ) self.print_ir = get_bool_env_var(f"{prefix}_PRINT_IR", False) - self.filter_stacktrace = get_bool_env_var(f"{prefix}_FILTER_STACKTRACE", True) - self.show_stacktrace = get_bool_env_var(f"{prefix}_SHOW_STACKTRACE", False) - self.lineinfo = get_bool_env_var(f"{prefix}_LINEINFO", False) + # SHOW_STACKTRACE (and DEBUG) show the full, unfiltered traceback, so + # internal-frame filtering is disabled by default in either mode. + self.show_stacktrace = get_bool_env_var(f"{prefix}_SHOW_STACKTRACE", self.debug) + self.filter_stacktrace = get_bool_env_var( + f"{prefix}_FILTER_STACKTRACE", not (self.debug or self.show_stacktrace) + ) + self.lineinfo = get_bool_env_var(f"{prefix}_LINEINFO", self.debug) self.no_cache = get_bool_env_var(f"{prefix}_NO_CACHE", False) self.jit_cache_max_elems = get_int_or_none_env_var( f"{prefix}_JIT_CACHE_MAX_ELEMS", None @@ -529,6 +630,13 @@ class EnvironmentVarManager(LogEnvironmentManager): ) _keep_tokens.add("cubin") + if get_bool_env_var(f"{prefix}_KEEP_SASS", False): + warnings.warn( + f"{prefix}_KEEP_SASS is deprecated; use {prefix}_KEEP=sass instead.", + DeprecationWarning, + stacklevel=2, + ) + _keep_tokens.add("sass") self.keep_tokens: frozenset[str] = frozenset(_keep_tokens) # Derived boolean attributes — used by compiler.py and dsl.py. @@ -538,6 +646,15 @@ class EnvironmentVarManager(LogEnvironmentManager): self.keep_ir: bool = "ir-debug" in self.keep_tokens self.keep_ptx: bool = "ptx" in self.keep_tokens self.keep_cubin: bool = "cubin" in self.keep_tokens + self.keep_sass: bool = "sass" in self.keep_tokens + check_sass = self.keep_sass + if check_sass: + _check_nvdisasm_available( + prefix, + self.keep_sass, + ) + self.remarks = get_str_env_var(f"{prefix}_REMARKS", "") + # Other options self.dryrun = get_bool_env_var(f"{prefix}_DRYRUN", False) self.arch = get_str_env_var(f"{prefix}_ARCH", detect_gpu_arch(prefix)) @@ -546,7 +663,7 @@ class EnvironmentVarManager(LogEnvironmentManager): ) self.warnings_ignore = get_bool_env_var(f"{prefix}_WARNINGS_IGNORE", False) self.enable_optimization_warnings = get_bool_env_var( - f"{prefix}_ENABLE_OPTIMIZATION_WARNINGS", False + f"{prefix}_ENABLE_OPTIMIZATION_WARNINGS", self.debug ) self.disable_file_caching = get_bool_env_var( f"{prefix}_DISABLE_FILE_CACHING", False diff --git a/python/CuTeDSL/cutlass/base_dsl/export/c_header_generator.py b/python/CuTeDSL/cutlass/base_dsl/export/c_header_generator.py index 8d54d25f1..afdd1029d 100644 --- a/python/CuTeDSL/cutlass/base_dsl/export/c_header_generator.py +++ b/python/CuTeDSL/cutlass/base_dsl/export/c_header_generator.py @@ -28,7 +28,8 @@ from ..typing import ( TFloat32, ) from ..dsl import is_dynamic_expression -from ..common import DSLRuntimeError +from ..common import DSLRuntimeError, DSLUserCodeError +from ..diagnostics import DiagId from ..jit_executor import ExecutionArgs from ..._mlir import ir @@ -138,9 +139,7 @@ class CHeaderGenerator: return c.strip() if hasattr(ann, "_field_names"): return ann.__name__ - raise DSLRuntimeError( - f"Unsupported type annotation for C header generation: {ann}" - ) + raise DSLUserCodeError(DiagId.TYPE_UNSUPPORTED_C_EXPORT, type_name=ann) @classmethod def generate_struct_typedef(cls, struct_cls: Any, device: bool = False) -> str: @@ -273,9 +272,42 @@ class CHeaderGenerator: """ if arg_type in self.numeric_to_c_type: return self.numeric_to_c_type[arg_type] + arg_name - raise DSLRuntimeError( - f"Unsupported argument type for c function argument generation: {arg_type}" + raise DSLUserCodeError(DiagId.TYPE_UNSUPPORTED_C_ARG, arg_type=arg_type) + + def _generate_tensor_element_typedef( + self, + symbol_prefix: str, + arg_name: str, + n_dyn_shape: int, + n_dyn_stride: int, + stride_type: str, + shape_type: str | None = None, + ) -> str: + """ + Format a per-element tensor descriptor typedef (`{prefix}_Tensor_{name}_t`). + Reused by both scalar-Tensor and list/tuple-of-Tensor dispatch branches in + subclasses so the on-wire descriptor layout stays identical. When + ``shape_type`` is ``None`` it defaults to ``stride_type`` so callers that + only pass ``stride_type`` get a struct whose shape and stride fields share + the same dynamic-int width (matching the runtime memref descriptor built + by ``cute::abi::MemRef::build``). + """ + shape_type = shape_type if shape_type is not None else stride_type + dyn_shape_decl = ( + f"\n {shape_type} dynamic_shapes[{n_dyn_shape}];" + if n_dyn_shape > 0 + else "" ) + dyn_stride_decl = ( + f"\n {stride_type} dynamic_strides[{n_dyn_stride}];" + if n_dyn_stride > 0 + else "" + ) + return f""" +typedef struct {{ + void *data;{dyn_shape_decl}{dyn_stride_decl} +}} {symbol_prefix}_Tensor_{arg_name}_t; +""" def _generate_check_cuda(self, dsl_name: str) -> str: check_cuda = ( @@ -385,8 +417,10 @@ static inline void {symbol_prefix}_Kernel_Module_Unload({symbol_prefix}_Kernel_M arguments.append("CUstream " + arg_name) packed_args.append("&" + arg_name) else: - raise DSLRuntimeError( - f"Unsupported argument for c function argument generation: {arg} with type {arg_type}" + raise DSLUserCodeError( + DiagId.ARG_UNSUPPORTED_C_EXPORT, + arg_name=arg_name, + arg_type=arg_type, ) return arguments, packed_args, declarations diff --git a/python/CuTeDSL/cutlass/base_dsl/export/export.py b/python/CuTeDSL/cutlass/base_dsl/export/export.py index c99dc9ef7..23af0f0b8 100644 --- a/python/CuTeDSL/cutlass/base_dsl/export/export.py +++ b/python/CuTeDSL/cutlass/base_dsl/export/export.py @@ -13,7 +13,8 @@ import inspect import io import os -from ..common import DSLRuntimeError +from ..common import DSLRuntimeError, DSLUserCodeError +from ..diagnostics import DiagId from ..._mlir import ir from ..._mlir.dialects import llvm @@ -151,9 +152,7 @@ def encode_metadata_into_ir_module( @param object_file_version: The version of the object file. """ if not signature: - raise DSLRuntimeError( - "signature is empty, please set the signature for the python jit function." - ) + raise DSLUserCodeError(DiagId.CONFIG_EXPORT_NO_RUNTIME_ARGS) version = object_file_version + c_string_suffix signature_bytes = signature_processor.dumps(signature) diff --git a/python/CuTeDSL/cutlass/base_dsl/export/external_binary_module.py b/python/CuTeDSL/cutlass/base_dsl/export/external_binary_module.py index b6f9c9987..36a1b7e29 100644 --- a/python/CuTeDSL/cutlass/base_dsl/export/external_binary_module.py +++ b/python/CuTeDSL/cutlass/base_dsl/export/external_binary_module.py @@ -17,7 +17,8 @@ from collections.abc import Callable from inspect import Signature, Parameter from typing import Any, cast -from ..common import DSLRuntimeError +from ..common import DSLRuntimeError, DSLUserCodeError +from ..diagnostics import DiagId from ...base_dsl.dsl import BaseDSL from ...base_dsl.typing import Int32, Int64, Float32, Float64 from .export import SignatureProcessor, decode_metadata_from_execution_engine @@ -27,7 +28,7 @@ def _get_ctypes_return_type(signature: Signature) -> Any: """Get the ctypes return type from the signature.""" return_type = signature.return_annotation if return_type is Parameter.empty: - raise DSLRuntimeError("Return type is not specified for AOT compiled function.") + raise DSLUserCodeError(DiagId.AOT_MISSING_RETURN_TYPE) type_to_ctype = { Int32: ctypes.c_int32, Int64: ctypes.c_int64, @@ -38,7 +39,9 @@ def _get_ctypes_return_type(signature: Signature) -> Any: ctype = type_to_ctype.get(return_type) if ctype is None: - raise DSLRuntimeError(f"Unsupported return type for AOT loading: {return_type}") + raise DSLUserCodeError( + DiagId.AOT_UNSUPPORTED_RETURN_TYPE, return_type=return_type + ) return ctype diff --git a/python/CuTeDSL/cutlass/base_dsl/ffi.py b/python/CuTeDSL/cutlass/base_dsl/ffi.py index ba0c6fbe6..d8ebe91fb 100644 --- a/python/CuTeDSL/cutlass/base_dsl/ffi.py +++ b/python/CuTeDSL/cutlass/base_dsl/ffi.py @@ -42,7 +42,6 @@ import inspect from dataclasses import dataclass import string -# Typing.get_overloads requires Python 3.11+; fall back to typing_extensions on 3.10. if sys.version_info >= (3, 11): from typing import get_overloads else: @@ -54,7 +53,8 @@ from .._mlir.dialects import func, gpu, llvm from . import typing as t from .typing import get_mlir_types, NumericMeta, as_numeric from .dsl import extract_mlir_values -from .common import DSLRuntimeError +from .common import DSLRuntimeError, DSLUserCodeError +from .diagnostics import DiagId @dataclass(frozen=True) @@ -360,7 +360,7 @@ class ExternCallHandler: break if not matched: - raise DSLRuntimeError("failed to find matching overload for call to ffi") + raise DSLUserCodeError(DiagId.CALL_FFI_NO_MATCH) assert params_types is not None assert return_types is not None @@ -654,10 +654,7 @@ class FFI: """ if kwargs: - raise DSLRuntimeError( - "Keyword arguments are not supported for FFI calls", - suggestion="Use positional arguments only", - ) + raise DSLUserCodeError(DiagId.CALL_FFI_NO_KWARGS) # Get the current insertion point and operation try: @@ -671,9 +668,10 @@ class FFI: module_op, insertion_region = self._get_prototype_region(current_op) if len(args) != len(self.params_types): - raise DSLRuntimeError( - f"Number of arguments mismatch, expected {len(self.params_types)}, got {len(args)}", - suggestion="Make sure the number of arguments matches the number of parameters", + raise DSLUserCodeError( + DiagId.ARG_COUNT_MISMATCH, + got=len(args), + expected=len(self.params_types), ) # Extract the arguments to MLIR values diff --git a/python/CuTeDSL/cutlass/base_dsl/jit_executor.py b/python/CuTeDSL/cutlass/base_dsl/jit_executor.py index 1d32fc810..d36a4853e 100644 --- a/python/CuTeDSL/cutlass/base_dsl/jit_executor.py +++ b/python/CuTeDSL/cutlass/base_dsl/jit_executor.py @@ -10,7 +10,27 @@ # is strictly prohibited. """ -This module provides jit executor related classes +This module provides jit executor related classes. + +Pointer-address runtime arguments are opt-in. A compile-time example such as +``nullptr(dtype, space)`` marks the corresponding runtime argument as a +pointer slot, and +``ExecutionArgs.record_pointer_arg_specs_from_compile_args`` records that +metadata in ``_pointer_address_arg_specs``. Runtime calls may then pass an +integer address, ``ctypes.c_void_p``, a ctypes pointer object, or +``nullptr(...)`` for a null address. Native JIT execution packs raw +addresses into stable ``ctypes.c_void_p`` storage via ``_RuntimePointerArg``; +``nullptr(...)`` already provides that storage itself and is retained as a +keepalive when used as a temporary runtime argument. TVM FFI receives raw +integer addresses instead. + +Plain integers are treated as pointer addresses only when the function +annotation is pointer-shaped. This avoids silently repacking ordinary scalar +integer parameters. Nested list/tuple compile-time arguments preserve the +sequence shape; mismatched runtime sequence lengths are left unchanged. +Arguments beyond the filtered runtime signature do not have annotation +metadata, so plain integer tail arguments stay scalar. Explicit ctypes pointer +tail arguments are still packed as pointer slots. """ import abc @@ -18,7 +38,7 @@ import array import ctypes import inspect import io -from typing import Any, NamedTuple, TYPE_CHECKING, ClassVar +from typing import Any, NamedTuple, TYPE_CHECKING, ClassVar, cast, get_args, get_origin from collections.abc import Callable, Sequence import weakref import threading @@ -35,9 +55,10 @@ from .._mlir.dialects import llvm from . import typing as t from .common import ( DSLRuntimeError, + DSLUserCodeError, DSLCudaRuntimeError, - create_cuda_runtime_error, ) +from .diagnostics import DiagId from .runtime import cuda as cuda_helpers from .runtime.jit_arg_adapters import JitArgAdapterRegistry, is_arg_annotation_constexpr from .typing import get_c_pointers @@ -176,17 +197,296 @@ class KwargsWrapperSpec(NamedTuple): kwonly_defaults: dict[str, Any] -class ExecutionArgs: - """Helper that wraps the function signature spec to filter execution and compile time arguments.""" +def _validate_pointer_address(address: int) -> None: + if address < 0: + raise DSLUserCodeError( + DiagId.ARG_POINTER_NEGATIVE, + address=address, + ) - def __init__(self, signature: inspect.Signature, function_name: str) -> None: + +class _RuntimePointerArg: + """JIT argument wrapper for a raw pointer address supplied from Python.""" + + def __init__(self, address: int) -> None: + _validate_pointer_address(address) + self._desc = ctypes.c_void_p(address) + self._c_pointers_cache = [ctypes.addressof(self._desc)] + + def __c_pointers__(self) -> list[int]: + return self._c_pointers_cache + + +PointerAddressArgSpec = bool | list[object] | tuple[object, ...] | None + + +def _is_ctypes_pointer(arg: object) -> bool: + return isinstance(arg, (ctypes.c_void_p, ctypes._Pointer)) + + +def _is_pointer_annotation(annotation: object) -> bool: + if annotation is None: + return False + + if type(annotation).__name__ == "TypedPointer": + return True + + if any(_is_pointer_annotation(arg) for arg in get_args(annotation)): + return True + + return getattr(annotation, "__name__", None) == "Pointer" and getattr( + annotation, "__module__", "" + ).endswith( + ( + "cutlass.base_dsl.typing", + "cutlass.cute.typing", + "cutlass.cute.core", + ) + ) + + +def _sequence_element_annotations( + annotation: object, length: int, *, is_tuple: bool +) -> list[object]: + args = get_args(annotation) + if not args: + return [None] * length + + origin = get_origin(annotation) + if origin is tuple: + if not is_tuple: + return [None] * length + if len(args) == 2 and args[1] is Ellipsis: + return [args[0]] * length + if len(args) == length: + return list(args) + return [None] * length + + if origin in (list, Sequence): + return [args[0]] * length + + return [None] * length + + +def _python_pointer_address(arg: object) -> int | None: + if isinstance(arg, bool): + return None + if getattr(arg, "_python_pointer_address_arg", False) is True: + return 0 + if isinstance(arg, int): + return arg + if isinstance(arg, ctypes.c_void_p): + return 0 if arg.value is None else int(arg.value) + if isinstance(arg, ctypes._Pointer): + value = ctypes.cast(arg, ctypes.c_void_p).value + return 0 if value is None else int(value) + return None + + +def _pointer_address_spec_from_compile_arg( + arg: object, annotation: object = None +) -> PointerAddressArgSpec: + if getattr(arg, "_python_pointer_address_arg", False) is True: + return True + + if isinstance(arg, (list, tuple)): + child_annotations = _sequence_element_annotations( + annotation, len(arg), is_tuple=isinstance(arg, tuple) + ) + child_specs = [ + _pointer_address_spec_from_compile_arg(x, child_annotations[index]) + for index, x in enumerate(arg) + ] + if any(spec is not None for spec in child_specs): + return cast( + PointerAddressArgSpec, + tuple(child_specs) if isinstance(arg, tuple) else child_specs, + ) + + address = _python_pointer_address(arg) + if address is not None: + if _is_ctypes_pointer(arg) or _is_pointer_annotation(annotation): + return True + return None + + return None + + +def _convert_python_pointer_arg( + arg: object, spec: PointerAddressArgSpec, *, for_tvm_ffi: bool = False +) -> tuple[object, bool]: + if spec is None: + return arg, False + + if spec is True: + if ( + not for_tvm_ffi + and getattr(arg, "_python_pointer_address_arg", False) is True + ): + return arg, True + address = _python_pointer_address(arg) + if address is None: + return arg, False + _validate_pointer_address(address) + if for_tvm_ffi: + return address, True + return _RuntimePointerArg(address), True + + if isinstance(spec, (list, tuple)) and isinstance(arg, (list, tuple)): + if len(arg) != len(spec): + return arg, False + changed = False + converted = [] + for item, item_spec in zip(arg, spec): + converted_item, item_changed = _convert_python_pointer_arg( + item, cast(PointerAddressArgSpec, item_spec), for_tvm_ffi=for_tvm_ffi + ) + converted.append(converted_item) + changed = changed or item_changed + if not changed: + return arg, False + return tuple(converted) if isinstance(arg, tuple) else converted, True + + return arg, False + + +class ExecutionArgs: + """Runtime argument binder for compiled JIT functions. + + Besides normal signature rectification and scalar casting, this class owns + the pointer-address conversion contract. + ``record_pointer_arg_specs_from_compile_args`` derives a lightweight spec + tree from compile-time arguments, and + ``generate_execution_args`` / ``convert_python_pointer_args_for_tvm_ffi`` use + that spec to decide where Python pointer-like values may replace runtime + pointer objects. Extra tail arguments are outside the signature and have no + ``name_to_index`` entry or compile-example spec; only explicit ctypes + pointer objects are repacked there. + """ + + def __init__( + self, + signature: inspect.Signature, + function_name: str, + full_arg_check: bool = False, + ) -> None: self.function_name = function_name self.signature = self.filter_runtime_signature(signature) self.original_signature = signature self._missing = object() self._meta = self._build_meta() + self._pointer_address_arg_specs: list[PointerAddressArgSpec] = [ + None + ] * self._meta.arg_count + self._has_pointer_address_arg_specs = False + # When True (set when debugging mode is ON), generate_execution_args + # runs thorough per-argument validation that is otherwise skipped to + # keep the launch path fast. + self._full_arg_check = full_arg_check self._tls = threading.local() + @property + def has_pointer_address_arg_specs(self) -> bool: + """Whether this compiled signature has pointer-address conversion slots.""" + return self._has_pointer_address_arg_specs + + def record_pointer_arg_specs_from_compile_args( + self, args: tuple[Any, ...], kwargs: dict[str, Any] + ) -> None: + """Record pointer conversion metadata from compile-time arguments. + + The resulting ``_pointer_address_arg_specs`` list is indexed like the + filtered runtime signature. Each leaf is ``True`` for an argument slot + that may accept a raw pointer address, ``None`` for ordinary arguments, + or a nested list/tuple spec for sequence arguments. + + ``_pointer_address_spec_from_compile_arg`` marks explicit pointer + arguments such as ``nullptr(...)`` and ctypes pointer objects. Plain + ``int`` compile-time arguments are marked only when the corresponding + annotation is pointer-shaped; this keeps scalar integer parameters from + being packed as pointer slots. For mixed tuple annotations, element + annotations are matched positionally. For generic sequence annotations + such as + ``list[cutlass.Pointer]`` or ``Sequence[cutlass.Pointer]``, the inner annotation + is applied to each element. + + ``args``/``kwargs`` are first rectified through the function signature + unless they already exactly match the positional runtime shape. That + path may raise the usual argument-binding ``DSLRuntimeError`` cases: + missing, unexpected, or duplicate arguments. + """ + if not args and not kwargs: + return + + if not kwargs and len(args) == self._meta.arg_count: + input_args: Sequence[Any] = args + else: + input_args = self.get_rectified_args(args, kwargs) + + self._pointer_address_arg_specs = [ + _pointer_address_spec_from_compile_arg( + arg, self._meta.annotated_types[index] + ) + for index, arg in enumerate(input_args) + ] + self._has_pointer_address_arg_specs = any( + spec is not None for spec in self._pointer_address_arg_specs + ) + + def convert_python_pointer_args_for_tvm_ffi( + self, args: tuple[Any, ...], kwargs: dict[str, Any] + ) -> tuple[tuple[Any, ...], dict[str, Any]]: + """Convert pointer-like Python arguments for TVM FFI calls. + + TVM FFI expects pointer values as integer addresses rather than native + JIT pointer-slot storage. This method walks positional arguments and + keyword arguments using ``_pointer_address_arg_specs`` and + ``_convert_python_pointer_arg``. ``name_to_index`` maps keyword names + back to the signature slot whose compile-time spec controls conversion. + + Returns the original ``args`` and ``kwargs`` objects when nothing + changed. If any conversion happens, returns a new positional tuple and, + only when needed, a copied kwargs dict. Sequence shapes are preserved + only when runtime list/tuple lengths match the compile-time arguments; + a mismatched or unsupported value is left unchanged. Negative pointer + addresses raise ``DSLRuntimeError``. + """ + if not self._has_pointer_address_arg_specs: + return args, kwargs + + converted_args = list(args) + changed = False + positional_count = min(len(converted_args), self._meta.arg_count) + for index in range(positional_count): + converted, did_convert = _convert_python_pointer_arg( + converted_args[index], + self._pointer_address_arg_specs[index], + for_tvm_ffi=True, + ) + if did_convert: + converted_args[index] = converted + changed = True + + converted_kwargs = kwargs + for name, value in kwargs.items(): + kw_index = self._meta.name_to_index.get(name) + if kw_index is None or kw_index >= len(self._pointer_address_arg_specs): + continue + converted, did_convert = _convert_python_pointer_arg( + value, + self._pointer_address_arg_specs[kw_index], + for_tvm_ffi=True, + ) + if did_convert: + if converted_kwargs is kwargs: + converted_kwargs = dict(kwargs) + converted_kwargs[name] = converted + changed = True + + if not changed: + return args, kwargs + return tuple(converted_args), converted_kwargs + def _build_meta(self) -> ArgMeta: """ Precompute metadata for the fast-path execution. @@ -243,16 +543,10 @@ class ExecutionArgs: """ pos_count = len(self._meta.pos_names) if len(args) > pos_count: - raise DSLRuntimeError( - "input args/kwargs length does not match runtime function signature!", - context={ - "input args length": len(args), - "input kwargs length": len(kwargs), - "function signature args length": len(self._meta.pos_names), - "function signature kwonlyargs length": len( - self._meta.kwonly_names - ), - }, + raise DSLUserCodeError( + DiagId.CALL_TOO_MANY_ARGS, + expected=len(self._meta.pos_names), + provided=len(args), ) # Start with every slot marked missing, we overwrite as values/defaults bind @@ -278,14 +572,14 @@ class ExecutionArgs: for name, value in kwargs.items(): idx = self._meta.name_to_index.get(name) # type: ignore[assignment] if idx is None: - raise DSLRuntimeError( - "unexpected keyword argument", - context={"argument_name": name}, + raise DSLUserCodeError( + DiagId.CALL_UNEXPECTED_KWARG, + argument_name=name, ) if idx < pos_len: - raise DSLRuntimeError( - "multiple values for argument", - context={"argument_name": name}, + raise DSLUserCodeError( + DiagId.CALL_DUPLICATE_ARGUMENT, + argument_name=name, ) rectified[idx] = value @@ -295,15 +589,9 @@ class ExecutionArgs: for i, name in enumerate(self._meta.all_names) if rectified[i] is self._missing ] - raise DSLRuntimeError( - "input args/kwargs length does not match runtime function signature!", - context={ - "missing": missing_args, - "function signature args length": len(self._meta.pos_names), - "function signature kwonlyargs length": len( - self._meta.kwonly_names - ), - }, + raise DSLUserCodeError( + DiagId.CALL_MISSING_ARGS, + missing=missing_args, ) return rectified @@ -334,7 +622,18 @@ class ExecutionArgs: else: input_args = self.get_rectified_args(args, kwargs) + if self._full_arg_check: + self._validate_args_full(input_args) + for index, arg in enumerate(input_args): + pointer_arg, is_pointer_address = _convert_python_pointer_arg( + arg, self._pointer_address_arg_specs[index] + ) + if is_pointer_address: + adapted_args.append(pointer_arg) + exe_arg_chunks[index] = get_c_pointers(pointer_arg) + continue + cptr_method = getattr(arg, "__c_pointers__", None) if cptr_method is not None: exe_arg_chunks[index] = cptr_method() @@ -366,10 +665,92 @@ class ExecutionArgs: if cptr_method is not None: exe_args.extend(cptr_method()) else: - exe_args.append(arg) + address = ( + _python_pointer_address(arg) if _is_ctypes_pointer(arg) else None + ) + if address is None: + exe_args.append(arg) + else: + pointer_arg = _RuntimePointerArg(address) + adapted_args.append(pointer_arg) + exe_args.extend(pointer_arg.__c_pointers__()) return exe_args, adapted_args + def _arg_marshals_to_pointers(self, arg: Any, spec: PointerAddressArgSpec) -> bool: + """Whether ``arg`` will contribute at least one C pointer at launch. + + Mirrors the marshalling decisions in ``generate_execution_args``: + pointer-address slots, ``__c_pointers__`` providers, registered JIT + argument adapters, and (recursively) sequences. An argument that + matches none of these is silently dropped by ``get_c_pointers`` (it + returns ``[]``), so the kernel would receive the wrong arguments. + """ + # Pointer-address slot: converted to a _RuntimePointerArg when the value + # carries a pointer address (raw int / ctypes pointer / nullptr marker, + # which _python_pointer_address already reports). When it does not, + # generate_execution_args falls through to the branches below, so we do + # too rather than rejecting a value that still marshals. + if spec is True and _python_pointer_address(arg) is not None: + return True + if hasattr(arg, "__c_pointers__"): + return True + try: + if JitArgAdapterRegistry.get_registered_adapter(arg) is not None: + return True + except Exception: + # Classification is best-effort; if adapter lookup itself trips, + # defer to the normal launch path rather than raising here. + return True + if isinstance(arg, (list, tuple)): + if isinstance(spec, (list, tuple)) and len(spec) == len(arg): + return all( + self._arg_marshals_to_pointers( + item, cast(PointerAddressArgSpec, item_spec) + ) + for item, item_spec in zip(arg, spec) + ) + return all(self._arg_marshals_to_pointers(item, None) for item in arg) + return False + + def _validate_args_full(self, input_args: Sequence[Any]) -> None: + """Thorough per-argument validation, run when debugging mode is ON. + + The normal launch path skips these checks to minimize overhead and + trusts the caller to pass well-formed arguments. When debugging mode is + ON we instead validate each argument against the recorded specs and + raise a precise error rather than silently mis-marshalling (an argument + that yields zero C pointers) or crashing inside the C call. + + Note: mismatched runtime sequence lengths are intentionally tolerated + (see the module docstring) and are not flagged here. + """ + meta = self._meta + for index, arg in enumerate(input_args): + name = meta.all_names[index] if index < len(meta.all_names) else f"#{index}" + spec = self._pointer_address_arg_specs[index] + + # Numeric scalars must be number-like so the later t.cast succeeds; + # only reject clearly non-numeric values to avoid false positives on + # the DSL's own numeric wrapper types. + if meta.numeric_flags[index]: + if arg is None or isinstance(arg, (str, bytes, list, tuple, dict, set)): + raise DSLUserCodeError( + DiagId.ARG_NOT_NUMERIC, + arg_name=name, + arg_type=type(arg).__name__, + context={"position": index}, + ) + continue + + if not self._arg_marshals_to_pointers(arg, spec): + raise DSLUserCodeError( + DiagId.ARG_NOT_MARSHALABLE, + arg_name=name, + arg_type=type(arg).__name__, + context={"position": index}, + ) + def get_kwargs_wrapper_spec( self, exclude_arg_names: Sequence[str] = () ) -> KwargsWrapperSpec: @@ -462,14 +843,10 @@ class ExecutionArgs: ) if len(runtime_args) != pos_count or len(runtime_kwargs) != kw_count: - raise DSLRuntimeError( - "input args/kwargs length does not match runtime function signature!", - context={ - "input args length": len(runtime_args), - "input kwargs length": len(runtime_kwargs), - "function signature args length": pos_count, - "function signature kwonlyargs length": kw_count, - }, + raise DSLUserCodeError( + DiagId.CALL_SIGNATURE_MISMATCH, + provided=len(runtime_args), + provided_kw=len(runtime_kwargs), ) return runtime_args + tuple(runtime_kwargs.values()) @@ -697,7 +1074,7 @@ class JitExecutor: error_code = self.cuda_result.value # type: ignore[union-attr] if error_code == 0: return error_code - raise create_cuda_runtime_error(error_code) + raise cuda_helpers.create_cuda_runtime_error(error_code) except DSLCudaRuntimeError as e: raise e except Exception as e: @@ -716,6 +1093,7 @@ class JitFunctionArtifacts: PTX: str | None CUBIN: str | bytes | None + SASS: str | None MLIR: str | None # Device compilation artifacts (set when DeviceTarget is enabled) device_header: str | None = None @@ -733,7 +1111,13 @@ class JitFunctionArtifacts: with open(self.CUBIN, "rb") as f: self.CUBIN = f.read() except (IOError, OSError) as e: - raise DSLRuntimeError(f"Failed to read CUBIN file '{self.CUBIN}': {e}") # type: ignore[str-bytes-safe] + raise DSLRuntimeError(f"Failed to read CUBIN file {self.CUBIN!r}: {e}") + if self.SASS is not None and os.path.exists(self.SASS): + try: + with open(self.SASS, "r") as f: + self.SASS = f.read() + except (IOError, OSError) as e: + raise DSLRuntimeError(f"Failed to read SASS file '{self.SASS}': {e}") if self.MLIR is not None and os.path.exists(self.MLIR): try: with open(self.MLIR, "r") as f: @@ -801,7 +1185,19 @@ class JitCompiledFunction: self.function_name = function_name self.kernel_info = kernel_info if kernel_info is not None else dict[str, Any]() if signature is not None: - self.execution_args = ExecutionArgs(signature, self.function_name) + # Enable thorough per-launch argument validation when debugging mode + # is ON. Read it from the DSL's env manager so the launch path reads + # a single stored boolean instead of the env on every call; when + # export_provider is unset, validation stays off. + full_arg_check = False + if self.export_provider is not None: + full_arg_check = self.export_provider.dsl._get_dsl().envar.debug + self.execution_args = ExecutionArgs( + signature, self.function_name, full_arg_check=full_arg_check + ) + self.execution_args.record_pointer_arg_specs_from_compile_args( + tuple(dynamic_args or ()), dynamic_kwargs or {} + ) self.jit_time_profiling = jit_time_profiling assert ( @@ -844,6 +1240,11 @@ class JitCompiledFunction: """Returns the CUBIN data of the JIT-compiled function.""" return self.artifacts.CUBIN if self.artifacts is not None else None + @property + def __sass__(self) -> str | None: + """Returns the SASS code of the JIT-compiled function.""" + return self.artifacts.SASS if self.artifacts is not None else None + @property def __mlir__(self) -> str | None: """Returns the MLIR code of the JIT-compiled function.""" diff --git a/python/CuTeDSL/cutlass/base_dsl/leaf_utils.py b/python/CuTeDSL/cutlass/base_dsl/leaf_utils.py index ed630e877..e0e65c23e 100644 --- a/python/CuTeDSL/cutlass/base_dsl/leaf_utils.py +++ b/python/CuTeDSL/cutlass/base_dsl/leaf_utils.py @@ -239,7 +239,7 @@ class LeafInfo: self.obj.value = new_vals[0] return - # Case 3: raw ir.Value (including subclasses like ArithValue, Pointer) + # Case 3: raw ir.Value, including subclasses like ArithValue. if isinstance(self.obj, ir.Value): if len(new_vals) >= 1: self._replace_in_parent(new_vals[0]) diff --git a/python/CuTeDSL/cutlass/base_dsl/pointer.py b/python/CuTeDSL/cutlass/base_dsl/pointer.py new file mode 100644 index 000000000..7a9dbb88e --- /dev/null +++ b/python/CuTeDSL/cutlass/base_dsl/pointer.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +from typing import Type + +from .._mlir import ir +from .._mlir.dialects import llvm +from .._mlir_helpers.arith import ArithValue +from .._mlir_helpers.op import dsl_user_op +from .address_space import AddressSpace +from .typing import Int32, Int64, Integer, Numeric, Pointer, _normalize_address_space + + +@dsl_user_op +def inttoptr( + value: int | Integer | ArithValue | ir.Value, + mem_space: AddressSpace | int, + dtype: Type[Numeric], + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, +) -> Pointer: + """Convert an integer address to a typed ``Pointer``. + + :param value: Integer address value. + :param mem_space: Target pointer address space. + :param dtype: Pointer element type. + :return: Typed pointer in ``mem_space``. + """ + space = _normalize_address_space(mem_space) + if isinstance(value, int): + address_dtype = ( + Int64 if space in (AddressSpace.generic, AddressSpace.gmem) else Int32 + ) + value = address_dtype(value).ir_value(loc=loc, ip=ip) + elif isinstance(value, Integer): + value = value.ir_value(loc=loc, ip=ip) + elif isinstance(value, ArithValue): + value = value.ir_value(loc=loc, ip=ip) + elif not isinstance(value, ir.Value): + raise TypeError(f"Unsupported value type for inttoptr: {type(value)}") + + res_val = llvm.inttoptr(llvm.PointerType.get(space.value), value, loc=loc, ip=ip) + return Pointer._from_raw_ptr(res_val, dtype) + + +__all__ = ["inttoptr"] diff --git a/python/CuTeDSL/cutlass/base_dsl/runtime/cuda.py b/python/CuTeDSL/cutlass/base_dsl/runtime/cuda.py index cd5e3ee7e..91f72381b 100644 --- a/python/CuTeDSL/cutlass/base_dsl/runtime/cuda.py +++ b/python/CuTeDSL/cutlass/base_dsl/runtime/cuda.py @@ -60,6 +60,56 @@ def _cudaGetErrorEnum(error: Any) -> Any: raise DSLRuntimeError("Unknown error type: {}".format(error)) +def get_cuda_error_name_from_code( + error_code: int, + cuda_error_type: type[ + cuda.CUresult | cudart.cudaError_t | nvrtc.nvrtcResult + ] = cuda.CUresult, +) -> Union[str, bytes]: + """ + Return the CUDA error name for a raw integer error code. + :param error_code: The raw integer CUDA error code. + :type error_code: int + :param cuda_error_type: The CUDA enum class used to decode the error code. + :type cuda_error_type: type(cuda.CUresult or cudart.cudaError_t or nvrtc.nvrtcResult) + :return: The decoded CUDA error name, or a string describing an unknown code. + :rtype: str or bytes + """ + try: + cu_err = cuda_error_type(error_code) + return _cudaGetErrorEnum(cu_err) + except (ValueError, AttributeError): + return f"" + + +def create_cuda_runtime_error( + error_code: int, + cuda_error_type: type[ + cuda.CUresult | cudart.cudaError_t | nvrtc.nvrtcResult + ] = cuda.CUresult, + cause: BaseException | None = None, +) -> DSLCudaRuntimeError: + """ + Create a DSLCudaRuntimeError from a raw CUDA integer error code. + :param error_code: The raw integer CUDA error code. + :type error_code: int + :param cuda_error_type: The CUDA enum class used to decode the error code. + :type cuda_error_type: type(cuda.CUresult or cudart.cudaError_t or nvrtc.nvrtcResult) + :param cause: The original exception that caused this runtime error. + :type cause: BaseException or None + :return: The DSL CUDA runtime error. + :rtype: DSLCudaRuntimeError + """ + error = DSLCudaRuntimeError( + error_code, + get_cuda_error_name_from_code(error_code, cuda_error_type), + ) + if cause is not None: + error.__cause__ = cause + error.__suppress_context__ = True + return error + + def _get_gpu_arch_info(major: int, minor: int) -> tuple[str, str, list[str]]: """ Get GPU architecture information and compatibility details. @@ -82,17 +132,15 @@ def _get_gpu_arch_info(major: int, minor: int) -> tuple[str, str, list[str]]: (8, 7): ("Ampere", "sm_87", ["sm_87", "sm_86", "sm_80"]), # A10, A40 (9, 0): ("Hopper", "sm_90a", ["sm_90a"]), # H100 (10, 0): ("Blackwell", "sm_100a", ["sm_100a"]), # B200 + (10, 1): ("Thor", "sm_101a", ["sm_101a"]), # Thor (pre 13.0) (10, 3): ("Blackwell", "sm_103a", ["sm_103a"]), + (11, 0): ("Thor", "sm_110a", ["sm_110a"]), # Thor (post 13.0) (12, 0): ( "Blackwell", "sm_120a", ["sm_120a"], ), # RTX PRO 6000 / RTX 50 Series - (12, 1): ( - "Blackwell", - "sm_121a", - ["sm_121a"], - ), # DGX Spark + (12, 1): ("Blackwell", "sm_121a", ["sm_121a"]), # DGX Spark } return gpu_arch_map.get( (major, minor), ("Unknown", f"sm_{major}{minor}", [f"sm_{major}{minor}"]) @@ -566,7 +614,7 @@ def load_library_data(cubin_data: bytes | int) -> Any: # Load module data if isinstance(cubin_data, bytes): cubin_data = np.char.array(cubin_data).ctypes.data - _log().info(f"cuLibraryLoadData {cubin_data}") + _log().info(f"cuLibraryLoadData {cubin_data!r}") library = checkCudaErrors( cuda.cuLibraryLoadData(cubin_data, None, None, 0, None, None, 0) @@ -845,6 +893,7 @@ def default_stream() -> Any: return cuda.CUstream(0) + @lru_cache(maxsize=1) def get_driver_version() -> Any: """ diff --git a/python/CuTeDSL/cutlass/base_dsl/runtime/device_tensor.py b/python/CuTeDSL/cutlass/base_dsl/runtime/device_tensor.py index 184ec7e0f..b1169bafc 100644 --- a/python/CuTeDSL/cutlass/base_dsl/runtime/device_tensor.py +++ b/python/CuTeDSL/cutlass/base_dsl/runtime/device_tensor.py @@ -15,6 +15,7 @@ from typing import Any from . import cuda as cuda_helpers from .tensor_descriptor import * from ..common import * +from ..diagnostics import DiagId def allocate(tensor: TensorDescriptor, stream: Any = None) -> None: @@ -22,11 +23,9 @@ def allocate(tensor: TensorDescriptor, stream: Any = None) -> None: Allocates GPU memory """ if tensor._check_is_managed_by_framework(): - raise DSLRuntimeError( - "GPU tensors are managed by the framework and cannot be modified." - ) + raise DSLUserCodeError(DiagId.TENSOR_FRAMEWORK_MANAGED, action="allocated") if not tensor.device_pointer is None: - raise DSLRuntimeError("Tensor is already allocated on the device.") + raise DSLUserCodeError(DiagId.TENSOR_ALREADY_ALLOCATED) tensor.device_pointer = cuda_helpers.allocate(tensor.size_in_bytes, stream) @@ -38,11 +37,9 @@ def deallocate(tensor: TensorDescriptor, stream: Any = None) -> None: Deallocates GPU memory """ if tensor._check_is_managed_by_framework(): - raise DSLRuntimeError( - "GPU tensors are managed by the framework and cannot be modified." - ) + raise DSLUserCodeError(DiagId.TENSOR_FRAMEWORK_MANAGED, action="deallocated") if tensor.device_pointer is None: - raise DSLRuntimeError("Tensor is not allocated on the device.") + raise DSLUserCodeError(DiagId.TENSOR_NOT_ALLOCATED) log().info( "Deallocating done tensor=[%s] dev_ptr=[%s]", tensor, tensor.device_pointer @@ -78,11 +75,9 @@ def copy_from_gpu( """ log().info("copyout tensor=[%s] dev_ptr=[%s]", tensor, tensor.device_pointer) if tensor._check_is_managed_by_framework(): - raise DSLRuntimeError( - "GPU tensors are managed by the framework and cannot be modified." - ) + raise DSLUserCodeError(DiagId.TENSOR_FRAMEWORK_MANAGED, action="copied") if tensor.device_pointer is None: - raise DSLRuntimeError("Tensor is not allocated on the device.") + raise DSLUserCodeError(DiagId.TENSOR_NOT_ALLOCATED) cuda_helpers.memcpy_d2h( tensor.data_ptr, tensor.device_pointer, tensor.size_in_bytes, stream @@ -106,7 +101,7 @@ def to_gpu(tensor: Any, stream: Any = None) -> TensorDescriptor: copy_to_gpu(new_tensor, stream=stream) return new_tensor - raise DSLRuntimeError("Unsupported type") + raise DSLUserCodeError(DiagId.TYPE_UNSUPPORTED_TENSOR, func="to_gpu()") def from_gpu(tensor: Any, stream: Any = None) -> TensorDescriptor: @@ -123,4 +118,4 @@ def from_gpu(tensor: Any, stream: Any = None) -> TensorDescriptor: copy_from_gpu(new_tensor, stream=stream) return new_tensor - raise DSLRuntimeError("Unsupported type") + raise DSLUserCodeError(DiagId.TYPE_UNSUPPORTED_TENSOR, func="from_gpu()") diff --git a/python/CuTeDSL/cutlass/base_dsl/runtime/jit_arg_adapters.py b/python/CuTeDSL/cutlass/base_dsl/runtime/jit_arg_adapters.py index 58f2745a8..be3143cb1 100644 --- a/python/CuTeDSL/cutlass/base_dsl/runtime/jit_arg_adapters.py +++ b/python/CuTeDSL/cutlass/base_dsl/runtime/jit_arg_adapters.py @@ -207,7 +207,7 @@ class DefaultDataclassAdapter: self._arg = arg for f in fields(arg): # type: ignore[arg-type] arg_field = getattr(arg, f.name) - if not is_constexpr_field(f): + if not is_constexpr_field(f, arg): if isinstance(f.type, NumericMeta) and not isinstance( arg_field, f.type ): @@ -241,7 +241,7 @@ class DefaultDataclassAdapter: kwargs = {} idx = 0 for f in fields(self._arg): # type: ignore[arg-type] - if is_constexpr_field(f): + if is_constexpr_field(f, self._arg): kwargs[f.name] = getattr(self._arg, f.name) else: kwargs[f.name] = new_from_mlir_values( diff --git a/python/CuTeDSL/cutlass/base_dsl/swizzle.py b/python/CuTeDSL/cutlass/base_dsl/swizzle.py new file mode 100644 index 000000000..de02726eb --- /dev/null +++ b/python/CuTeDSL/cutlass/base_dsl/swizzle.py @@ -0,0 +1,374 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +""" +``cutlass.base_dsl.swizzle`` — SMEM XOR swizzle library. + +**Recommended usage** — import ``cutlass`` and use the ``cutlass.`` namespace:: + + import cutlass + sw = cutlass.Swizzle(3, 4, 3) + sw = cutlass.Swizzle.from_name("s128b") + vec = (smem_ptr + offset).load_swizzled(sw, count=64) + +This file is the base-DSL implementation home. + +See ``write-dsl-kernel/swizzle.md`` for full documentation and worked examples. +""" + +from dataclasses import dataclass +from typing import cast + +from . import dsl_user_op +from .._mlir import ir as _ir + + +# ============================================================================= +# Swizzle descriptor (pure Python — no MLIR dependencies) +# ============================================================================= + + +@dataclass(frozen=True) +class Swizzle: + """CuTe-style XOR swizzle descriptor for SMEM address permutation. + + Encodes a 3-parameter XOR swizzle applied to SMEM byte addresses:: + + mask = (1 << bbits) - 1 + yyy = (addr >> (mbase + sshift)) & mask # upper bitfield + swizzled_addr = addr ^ (yyy << mbase) # XOR lower with upper + + Common presets (for use with ``tcgen05.mma kind::f16``):: + + s128b = Swizzle(3, 4, 3) # 128B XOR — required for SM100 tcgen05.mma + s64b = Swizzle(2, 4, 3) + s32b = Swizzle(1, 4, 3) + none = Swizzle(0, 0, 0) # identity — no-op + + :param bbits: Number of bits in the XOR mask (width of swizzle + region). ``3`` → 8-element swizzle period. + :param mbase: Least-significant bits kept constant (not affected by + XOR). ``4`` → bottom 16B granule is unchanged. + :param sshift: Shift of the upper mask from ``mbase`` (positive = + shift right). ``3`` → upper bits start at + ``mbase + sshift = 7``. + """ + + bbits: int + mbase: int + sshift: int + + @classmethod + def from_name(cls, name: str) -> "Swizzle": + """Construct a ``Swizzle`` from its canonical name string. + + Name → Swizzle(bbits, mbase, sshift): + + ========= ======================= + Name Swizzle + ========= ======================= + ``none`` ``Swizzle(0, 0, 0)`` + ``s32b`` ``Swizzle(1, 4, 3)`` + ``s64b`` ``Swizzle(2, 4, 3)`` + ``s128b`` ``Swizzle(3, 4, 3)`` + ========= ======================= + + :param name: Swizzle name string — one of ``"none"``, + ``"s32b"``, ``"s64b"``, ``"s128b"``. + :return: Corresponding ``Swizzle`` instance. + :raises ValueError: if ``name`` is not one of the supported + preset names. + + .. code-block:: python + + sw = Swizzle.from_name("s128b") + # sw == Swizzle(3, 4, 3) + """ + mapping = { + "none": cls(0, 0, 0), + "s32b": cls(1, 4, 3), + "s64b": cls(2, 4, 3), + "s128b": cls(3, 4, 3), + } + try: + return mapping[name] + except KeyError as exc: + choices = list(mapping) + raise ValueError( + f"Unsupported swizzle name {name!r}. Supported: {choices}." + ) from exc + + def to(self, target_type: type[object]) -> object: + """Convert this swizzle descriptor to another swizzle representation. + + ``cutlass.Swizzle`` is the semantic SMEM layout descriptor used by + author code. TMA tensor maps and tcgen05 descriptors use separate + low-level enum encodings, so convert at the API boundary instead of + carrying multiple encodings through kernel configuration. + """ + if target_type is Swizzle: + return self + + from cutlass.experimental.cuda import TensorMapSwizzle # noqa: PLC0415 + + mapping: dict[Swizzle, object] + if target_type is TensorMapSwizzle: + mapping = { + Swizzle(0, 0, 0): TensorMapSwizzle.none, + Swizzle(1, 4, 3): TensorMapSwizzle.s32b, + Swizzle(2, 4, 3): TensorMapSwizzle.s64b, + Swizzle(3, 4, 3): TensorMapSwizzle.s128b, + } + else: + raise TypeError(f"Unsupported swizzle conversion target: {target_type!r}") + try: + return mapping[self] + except KeyError as exc: + raise ValueError(f"No {target_type.__name__} preset for {self!r}.") from exc + + +# ============================================================================= +# Standalone swizzle library functions +# ============================================================================= + + +@dsl_user_op +def apply_swizzle( + ptr: object, + swizzle: Swizzle, + *, + loc: _ir.Location | None = None, + ip: _ir.InsertionPoint | None = None, +) -> object: + """Apply a CuTe-style XOR swizzle to a base-DSL Pointer's SMEM address. + + This low-level helper returns the physically swizzled pointer for op + families that need the address itself, such as per-thread + ``cp_async_shared_global`` destinations. For ordinary SMEM reads and writes, + prefer :func:`load_swizzled` and :func:`store_swizzled` so vector width is + preserved. + """ + if swizzle.bbits == 0: + return ptr + + from .typing import Pointer # noqa: PLC0415 + from .._mlir.dialects import cute as _cute_ir # noqa: PLC0415 + + base_ptr = cast(Pointer, ptr) + sw_ty = _ir.Type.parse( + f'!cute.swizzle<"S<{swizzle.bbits},{swizzle.mbase},{swizzle.sshift}>">' + ) + llvm_ptr = base_ptr.llvm_ptr + op = _cute_ir.ApplySwizzleOp( + ptr=llvm_ptr, + swizzle=sw_ty.attribute, + results=[llvm_ptr.type], + loc=loc, + ip=ip, + ) + return Pointer._from_raw_ptr(op.result, base_ptr.dtype) + + +@dsl_user_op +def load_swizzled( + ptr: object, + swizzle: Swizzle, + alignment: int | None = None, + *, + count: int | None = None, + loc: _ir.Location | None = None, + ip: _ir.InsertionPoint | None = None, +) -> object: + """Load scalar or vector from a swizzled SMEM base-DSL Pointer. + + Builds a CuTe-typed pointer with the swizzle encoded in its type, then + issues a CuTe load. Equivalent to ``ptr.load_swizzled(swizzle, ...)``. + + **Prefer wide vectors for efficiency.** The CuTe tensor that + implements the load remaps addresses over whole 128-bit vectors + (16 B, 8 fp16), not individual elements. Any ``count`` value is + correct — CuTe handles the XOR for scalar reads (``count=None``), + ``count=8`` atoms, single-period loads, and counts that span + multiple swizzle periods. The preference is about cost: scalar + and ``count=8`` reads emit N× the memory ops of a single wide + load, block ILP, and produce more instructions for the same + bytes. Use ``count = `` (16 / 32 / 64 fp16 for + s32b / s64b / s128b) or the full per-thread row span unless you + have a concrete measured reason to fragment. + + :param ptr: SMEM base-DSL ``Pointer`` to load from. + :param swizzle: ``Swizzle`` descriptor matching the SMEM layout. + :param alignment: Optional byte alignment hint for the load. + :param count: Element count. ``None`` → scalar load. Any integer + ``N ≥ 1`` → a ``Vector`` of length ``N`` (even ``count=1`` + returns a 1-element Vector, not a scalar). Very large counts + (≥ 256) can stress the MLIR lowering; if you see compile-time + regressions, drop back to one period per load. + :return: Scalar ``Numeric`` when ``count is None``; otherwise a + base-DSL ``Vector[ptr.dtype, count]`` so callers can use + ``vec.reduce(op)``, ``vec.to(dtype)``, element-wise arithmetic, + and pair with :func:`store_swizzled` without any intermediate + conversions. + + .. code-block:: python + + sw = cutlass.Swizzle.from_name("s128b") + vec = (smem_ptr + offset).load_swizzled(sw, count=64) # one period, fp16 + + :sync-class: Per-thread — single-thread SMEM load; not a collective. + :elect-safe: N/A — no divergent hardware side effect. + :device: SM80+ (swizzle kinds depend on what downstream op consumes + them; for ``tcgen05.mma kind::f16`` on SM100 the swizzle must be + ``SWIZZLE_128B``, i.e. ``Swizzle.from_name("s128b")``). + :side-effects: issues a swizzle-aware CuTe load; writes into the + register file (as scalar or 1-D Vector). + :peer: Pairs with :func:`store_swizzled` for an SMEM round-trip; the + vector returned here can be fed back verbatim. Swizzle must + match the writer's (TMA ``TensorMapSwizzle.sXb`` or a prior + :func:`store_swizzled` with the same ``Swizzle``). Keep the same + logical swizzle name across the TMA descriptor, + ``load_swizzled`` / ``store_swizzled``, and tcgen05 descriptor; + note that those APIs use different integer encodings for the + same logical swizzle. + """ + from .typing import Pointer # noqa: PLC0415 + import cutlass.cute as cute # noqa: PLC0415 + + base_ptr = cast(Pointer, ptr) + swizzle_obj = cute.make_swizzle( + swizzle.bbits, swizzle.mbase, swizzle.sshift, loc=loc, ip=ip + ) + cute_ptr = cute.make_ptr( + base_ptr.dtype, + base_ptr, + assumed_align=alignment, + swizzle_=swizzle_obj, + loc=loc, + ip=ip, + ) + if count is None: + return cute_ptr.load(loc=loc, ip=ip) + # Shape must be a tuple — passing a bare int here causes downstream + # make_tensor / layout machinery to mis-flatten for count > 8 and + # produces "`load_swizzled` count must be 8" tutorial folklore. + cute_tensor = cute.make_tensor(cute_ptr, (count,), loc=loc, ip=ip) + # TensorSSA.to_vector() returns a base-DSL Vector so the result works + # natively with Vector.reduce(op), Vector.to(dtype), element-wise ops, + # and round-trips through store_swizzled without manual conversion. + return cute_tensor.load(loc=loc, ip=ip).to_vector(loc=loc, ip=ip) + + +@dsl_user_op +def store_swizzled( + ptr: object, + value: object, + swizzle: Swizzle, + alignment: int | None = None, + *, + loc: _ir.Location | None = None, + ip: _ir.InsertionPoint | None = None, +) -> None: + """Store a scalar or 1-D vector to a swizzled SMEM base-DSL Pointer. + + Builds a CuTe-typed pointer with the swizzle encoded in its type, then + issues a CuTe store. Equivalent to ``ptr.store_swizzled(value, swizzle, ...)``. + + **Prefer wide vectors.** The CuTe tensor that implements the store + remaps addresses over whole 128-bit vectors (16 B, 8 fp16), not + individual elements. Pair a ``load_swizzled(sw, count=N)`` with a + ``store_swizzled(vec, sw)`` of the same width — one SMEM load, + one SMEM store, compute in registers between them. Storing + scalars one at a time on a swizzled pointer is legal (a dedicated + scalar path handles the XOR), but inflates SMEM traffic by N×; + there is no reason to fragment a vector return from + :func:`load_swizzled` into per-element stores. + + :param ptr: SMEM base-DSL ``Pointer`` to store to. + :param value: Either a scalar ``Numeric`` / ``ArithValue`` / Python + ``int`` / ``float`` (one element), or a 1-D ``Vector`` of + values. The element type is cast to ``ptr.dtype`` + automatically. Mirrors :func:`load_swizzled`: a scalar return + there pairs with a scalar ``value`` here; a ``Vector[T, N]`` + (any ``N ≥ 1``) return pairs with a matching ``Vector[T, N]`` + here. + :param swizzle: ``Swizzle`` descriptor matching the SMEM layout. + :param alignment: Optional byte alignment hint for the store. + :raises ValueError: if ``value`` is a multi-dimensional Vector + (only scalar or 1-D is accepted). + + .. code-block:: python + + # Canonical round-trip (preferred) + sw = cutlass.Swizzle.from_name("s128b") + vec = (smem_ptr + offset).load_swizzled(sw, count=64) + vec = (vec.to(cutlass.Float32) * 2.0).to(cutlass.Float16) + (smem_ptr + offset).store_swizzled(vec, sw) + + :sync-class: Per-thread — single-thread SMEM store; not a collective. + :elect-safe: N/A — no divergent hardware side effect. + :device: SM80+ (swizzle kinds depend on what downstream op consumes + them; for ``tcgen05.mma kind::f16`` on SM100 the swizzle must be + ``SWIZZLE_128B``, i.e. ``Swizzle.from_name("s128b")``). + :side-effects: issues a swizzle-aware CuTe store; writes into SMEM + at the XOR'd physical address. The ``Swizzle`` passed here + defines the physical layout any subsequent reader must assume. + :peer: Pairs with :func:`load_swizzled` — use the same ``Swizzle`` + on both ends of a register-residency round-trip. A TMA reader + consuming the tile must be built with a matching + ``TensorMapSwizzle``. + """ + from .typing import Pointer # noqa: PLC0415 + import cutlass.cute as cute # noqa: PLC0415 + + base_ptr = cast(Pointer, ptr) + # Detect scalar vs 1-D Vector. Scalars (Python int/float, Numeric, bare + # ArithValue) have no `.shape`; Vector exposes `.shape = (N,)` for N ≥ 1. + value_shape = getattr(value, "shape", None) + if value_shape is None or not isinstance(value_shape, tuple): + # Scalar path — route through the cute swizzle op + scalar store. + # The cute.make_ptr / cute_ptr.store path used below for vectors + # lowers a scalar store through `unrealized_conversion_cast` back + # to !llvm.ptr<3> before emitting llvm.store, which drops the + # swizzle-type annotation and stores at the unswizzled address. + swizzled_ptr = cast(Pointer, apply_swizzle(base_ptr, swizzle, loc=loc, ip=ip)) + swizzled_ptr.store(value, loc=loc, ip=ip) + return + swizzle_obj = cute.make_swizzle( + swizzle.bbits, swizzle.mbase, swizzle.sshift, loc=loc, ip=ip + ) + cute_ptr = cute.make_ptr( + base_ptr.dtype, + base_ptr, + assumed_align=alignment, + swizzle_=swizzle_obj, + loc=loc, + ip=ip, + ) + if len(value_shape) != 1: + raise ValueError( + "store_swizzled only supports scalar or 1-D Vector, " + f"got shape {value_shape}" + ) + count = value_shape[0] + cute_tensor = cute.make_tensor(cute_ptr, (count,), loc=loc, ip=ip) + # Internal: convert Vector -> TensorSSA and cast dtype before issuing the + # cute-side store. This is an implementation detail; callers pass a Vector. + cute_tensor.store( + cute.TensorSSA.from_vector(value).to(base_ptr.dtype), loc=loc, ip=ip + ) + + +__all__ = [ + "Swizzle", + "apply_swizzle", + "load_swizzled", + "store_swizzled", +] diff --git a/python/CuTeDSL/cutlass/base_dsl/torch.py b/python/CuTeDSL/cutlass/base_dsl/torch.py index 3bd7dbeeb..465611bc4 100644 --- a/python/CuTeDSL/cutlass/base_dsl/torch.py +++ b/python/CuTeDSL/cutlass/base_dsl/torch.py @@ -24,6 +24,7 @@ from .typing import ( Float8E5M2, Float8E8M0FNU, Float4E2M1FN, + Float4E2M1FNx2, ) @@ -41,6 +42,7 @@ def dtype(ty: Type[Numeric]) -> "torch.dtype": Float8E4M3FN: torch.float8_e4m3fn, Float8E4M3B11FNUZ: torch.float8_e4m3fnuz, Float4E2M1FN: torch.float4_e2m1fn_x2, + Float4E2M1FNx2: torch.float4_e2m1fn_x2, } # float8_e8m0fnu is introduced in latest version of torch diff --git a/python/CuTeDSL/cutlass/base_dsl/typing.py b/python/CuTeDSL/cutlass/base_dsl/typing.py index db38b9b4b..50e38d1aa 100644 --- a/python/CuTeDSL/cutlass/base_dsl/typing.py +++ b/python/CuTeDSL/cutlass/base_dsl/typing.py @@ -10,6 +10,7 @@ # is strictly prohibited. import ctypes +from abc import abstractmethod from itertools import chain import numpy as np import operator @@ -21,6 +22,8 @@ from typing import ( Generic, Optional, Protocol, + Literal, + TypeAlias, Union, Any, Type, @@ -31,13 +34,18 @@ from typing import ( from .common import * from .common import DSLRuntimeError as DSLRuntimeError +from .diagnostics import DiagId from .._mlir_helpers import arith as arith_helper from .._mlir_helpers.arith import ArithValue +from .._mlir_helpers.vector import Vector from .._mlir_helpers.op import dsl_user_op from .._mlir import ir from .._mlir.extras import types as T -from .._mlir.dialects import arith +from .._mlir.dialects import arith, llvm, nvvm, vector + +from .address_space import AddressSpace + # ============================================================================= # Dynamic Expression Protocol # ============================================================================= @@ -230,10 +238,8 @@ def get_c_pointers(obj: Any) -> list[ctypes.c_void_p]: elif isinstance(obj, (tuple, list)): return list(chain.from_iterable(get_c_pointers(x) for x in obj)) elif isinstance(obj, set): - raise DSLRuntimeError( - "Sets are not supported in get_c_pointers to ensure order preservation", - context="The DSL attempted to generate JIT function argument(s) for an argument of type set but failed.", - suggestion="Consider using a list or tuple instead", + raise DSLUserCodeError( + DiagId.ARG_UNORDERED_CONTAINER, ) return [] @@ -251,10 +257,8 @@ def get_mlir_types(obj: Any) -> list[ir.Type]: elif isinstance(obj, (tuple, list)): return sum((get_mlir_types(x) for x in obj), []) elif isinstance(obj, set): - raise DSLRuntimeError( - "Sets are not supported in get_mlir_types to ensure order preservation", - context="The DSL attempted to generate JIT function argument(s) for an argument of type set but failed.", - suggestion="Consider using a list or tuple instead", + raise DSLUserCodeError( + DiagId.ARG_UNORDERED_CONTAINER, ) return [] @@ -418,10 +422,16 @@ class NumericMeta(DslType): return cls._np_dtype @property - def is_integer(cls) -> bool: ... # type: ignore[empty-body] + @abstractmethod + def is_integer(cls) -> bool: ... @property - def is_float(cls) -> bool: ... # type: ignore[empty-body] + @abstractmethod + def is_float(cls) -> bool: ... + + @property + @abstractmethod + def zero(cls) -> Union[int, float]: ... def is_same_kind(cls, other: Type) -> bool: return cls.is_integer == other.is_integer or cls.is_float == other.is_float @@ -1003,6 +1013,15 @@ class Numeric(metaclass=NumericMeta, is_abstract=True): ip: Optional[ir.InsertionPoint] = None, ) -> "Numeric": ... + @overload + def to( + self, + dtype: Type[bool], + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> bool: ... + @overload def to( self, @@ -1021,15 +1040,6 @@ class Numeric(metaclass=NumericMeta, is_abstract=True): ip: Optional[ir.InsertionPoint] = None, ) -> float: ... - @overload - def to( # type: ignore[overload-cannot-match] - self, - dtype: Type[bool], - *, - loc: Optional[ir.Location] = None, - ip: Optional[ir.InsertionPoint] = None, - ) -> bool: ... - @overload def to( self, @@ -1095,9 +1105,7 @@ class Numeric(metaclass=NumericMeta, is_abstract=True): return dtype(self) elif dtype is ir.Value: if isinstance(self.value, (int, float, bool)): - res = arith_helper.const( - self.value, self.dtype.mlir_type, loc=loc, ip=ip - ) + res = arith_helper.const(self.value, type(self), loc=loc, ip=ip) elif isinstance(self.value, ir.Value): res = self.value else: @@ -1149,9 +1157,6 @@ class Numeric(metaclass=NumericMeta, is_abstract=True): result = arith.bitcast(dtype.mlir_type, ir_val, loc=loc, ip=ip) return dtype(result) - @property - def zero(self) -> "Numeric": ... # type: ignore[empty-body] - def __dsl_not__( self, *, @@ -1306,23 +1311,13 @@ class Numeric(metaclass=NumericMeta, is_abstract=True): if isinstance(self.value, (int, float, bool)): return bool(self.value) else: - raise DSLRuntimeError( - f"Unable to convert dynamic `{type(self).__name__}` value to bool at compile time.", - suggestion=[ - "Decorate the parent function with `jit` decorator and with `preprocess` enabled.", - "Ensure not using patterns that DSL does not support.", - "Otherwise, please file a bug report.", - ], - ) + raise DSLUserCodeError(DiagId.PHASE_DYNAMIC_TO_STATIC_BOOL) def __index__(self) -> int: if isinstance(self.value, (int, float, bool)): return self.value # type: ignore[return-value] else: - raise DSLRuntimeError( - f"'{type(self.value)}' object cannot be interpreted as an integer", - suggestion="Mark the loop as dynamic with `dynamic_expr` or `range_dynamic` and decorate the parent function with `jit` decorator", - ) + raise DSLUserCodeError(DiagId.PHASE_DYNAMIC_INDEX) def __neg__( self, @@ -1910,7 +1905,6 @@ class Float(Numeric, metaclass=FloatMeta, mlir_type=T.f32, is_abstract=True): ip: Optional[ir.InsertionPoint] = None, ) -> None: ty = type(self) - if isinstance(x, (bool, int, float)): # Why we need to convert x to with numpy? # np_dtype = ty.numpy_dtype @@ -2177,10 +2171,42 @@ class Float6E3M2FN(Float, metaclass=FloatMeta, width=6, mlir_type=T.f6E3M2FN): . class Float6E2M3FN(Float, metaclass=FloatMeta, width=6, mlir_type=T.f6E2M3FN): ... +# Packed narrow-float views. For FP4x2, the MLIR element type is the packed +# storage container (`i8`) so generic pointer vector loads/stores move packed +# bytes correctly. The class identity carries the floating-point packing +# semantics (`x2` in the name). `width` is the width of one packed tensor +# element: 8 bits for FP4x2. Naming follows torch: `float4_e2m1fn_x2`. +class Float4E2M1FNx2( + Float, + metaclass=FloatMeta, + width=8, + mlir_type=T.i8, +): + """Packed FP4 E2M1 — 2 elements per byte (matches ``torch.float4_e2m1fn_x2``). + + Shape and strides on any layout carrying this dtype are interpreted + in **fp4x2 tensor-element units**. One tensor element is already one + packed storage unit, so ``create_tensor_map_tiled_from_tensor`` uses + ``width == 8`` directly when converting stride units for TMA. + + ``width`` is the packed 8-bit tensor-element width and ``mlir_type`` is + the packed storage type ``i8``. Internal helpers that still need scalar + FP4 lane precision treat this packed dtype specially by class identity. + + Use this dtype when the input is already organized in packed fp4x2 + storage units (for example a ``torch.uint8`` buffer viewed as + ``torch.float4_e2m1fn_x2``) or when the kernel allocates layouts + directly from packed extents. + + """ + + def _element_precision_width(dtype: Type["Numeric"]) -> int: """Return scalar lane precision for packed narrow-float view dtypes.""" + if dtype is Float4E2M1FNx2: + return 4 return dtype.width @@ -2190,6 +2216,7 @@ _unsupported_dst_float_types = [ Float4E2M1FN, Float6E3M2FN, Float6E2M3FN, + Float4E2M1FNx2, ] @@ -2218,6 +2245,7 @@ ALL_DTYPES = { Float4E2M1FN, Float6E2M3FN, Float6E3M2FN, + Float4E2M1FNx2, } __STR_TO_DTYPE__ = {dt.__name__: dt for dt in ALL_DTYPES} @@ -2278,88 +2306,824 @@ else: pass +class _GridConstantMarker: + """Type marker that indicates a kernel argument as CUDA grid constant.""" + + def __extract_mlir_attributes__(self) -> list[ir.DictAttr]: + return [ir.DictAttr.get({"cuda.grid_constant": ir.UnitAttr.get()})] + + +# Singleton instance of the grid constant marker +grid_constant: _GridConstantMarker = _GridConstantMarker() + +# Type alias for ``Annotated[TY, grid_constant]`` +GridConstant: TypeAlias = Annotated[TY, grid_constant] + + class align(int): def __new__(cls, value: int) -> "align": if value <= 0 or (value & (value - 1)) != 0: - raise DSLRuntimeError("expects align be power of 2 as positive value") + raise DSLUserCodeError(DiagId.ARG_INVALID_ALIGNMENT) return super().__new__(cls, value) def __str__(self) -> str: return f"align({super().__str__()})" -class PointerMeta(DslType): - _value_type: Any - _align: Any - def __new__( - cls, - name: str, - bases: tuple, - attrs: dict, - value_type: Any = Int32, - align_: Any = align(1), - ) -> Any: - new_cls = super().__new__( - cls, - name, - bases, - attrs, - mlir_type=lambda: getattr(ir, "UnrankedMemRefType").get( - value_type.mlir_type, getattr(ir, "Attribute").parse("0") - ), - ) - new_cls._value_type = value_type - new_cls._align = align_ - return new_cls +class TypedPointer: + """Type annotation object for a ``Pointer`` element type and memory space. - def __eq__(cls, other: Any) -> bool: - if not isinstance(other, PointerMeta): - return False - return ( - cls._value_type == other._value_type - and cls._align._value == other._align._value - ) # Compare alignment values + ``Pointer[dtype, space]`` returns a ``TypedPointer`` for use in kernel/JIT + signatures. It is not an LLVM pointer value itself; it carries the metadata + needed to derive the corresponding MLIR argument type. - def __hash__(cls) -> int: - return hash((cls._value_type, cls._align._value)) # Hash alignment value + :param dtype: Element type such as ``Float32`` or ``Int8``. + :param space: Memory space such as ``cutlass.AddressSpace.gmem`` or its integer + address-space value. - def __getitem__(cls, params: Any) -> Type["Pointer"]: - value_type, align_ = params + Example:: - if not isinstance(align_, align): - raise DSLRuntimeError(f"expects align but got {align_}") - - # Create new class with proper name and parameters - new_cls = type( - f"Pointer[{value_type.__name__}, {align_}]", - (Pointer,), - {}, - value_type=value_type, - align_=align_, # Pass alignment to __new__ - ) - return new_cls - - def __str__(cls) -> str: - return f"ptr<{cls._value_type}, {cls._align}>" - - -class Pointer(metaclass=PointerMeta): - """ - A pointer to a memory location. - - Examples: - - def foo(a : Pointer[Int32, align=8]): + def kernel(data: cutlass.Pointer[Float32, cutlass.AddressSpace.gmem]): ... + ``__get_mlir_types__`` reads ``space.value`` when present, otherwise treats + ``space`` as the raw LLVM address-space integer. """ - def __init__(self, value: Any) -> None: - self.value = value + def __init__(self, dtype: Any, space: Any): + self.dtype = dtype + self.space = space + + def __repr__(self) -> str: + return f"TypedPointer[{self.dtype}, {self.space}]" + + def __get_mlir_types__(self) -> list[ir.Type]: + addrspace = getattr(self.space, "value", self.space) + return [llvm.PointerType.get(int(addrspace))] + + +def _normalize_address_space(space: Any) -> AddressSpace: + if isinstance(space, AddressSpace): + return space + if isinstance(space, bool): + raise TypeError( + f"space=bool is not allowed; pass a cutlass.AddressSpace member " + f"or its int value. Got {space!r}." + ) + if isinstance(space, int): + try: + return AddressSpace(space) + except ValueError: + valid = sorted(int(m) for m in AddressSpace) + raise ValueError( + f"space={space!r} is not a valid cutlass.AddressSpace int. " + f"Valid ints: {valid}." + ) from None + raise TypeError( + f"space must be a cutlass.AddressSpace member or int; got " + f"{type(space).__name__} ({space!r})." + ) + + +MemOrdering = Literal[ + "not_atomic", + "monotonic", + "acquire", + "release", + "acq_rel", + "seq_cst", +] +MemScope = Literal["cta", "cluster", "gpu", "sys"] +SharedSpace = Literal["cta", "cluster"] +LoadCacheModifier = Literal["ca", "cg", "cs", "lu", "cv"] +StoreCacheModifier = Literal["wb", "cg", "cs", "wt"] +EvictPriority = Literal["first", "last", "normal", "unchanged", "noallocate"] +L2PrefetchSize = Literal["size_64b", "size_128b", "size_256b"] +L1EvictKind = Literal[ + "evict_normal", + "evict_first", + "evict_last", + "evict_no_allocate", + "evict_unchanged", +] + +MLIR_DYNAMIC_INDEX = -(2**31) + + +def _to_llvm_compatible_type(elem_type: ir.Type) -> ir.Type: + """Convert non-LLVM-compatible element types to same-width integer types.""" + compatible_float_widths = {16, 32, 64, 80, 128} + if isinstance(elem_type, ir.FloatType): + width = elem_type.width + if width not in compatible_float_widths: + return ir.IntegerType.get_signless(width) + return elem_type + + +def _gep( + base: ir.Value, + elem_type: ir.Type, + *, + static_indices: list[int] | None = None, + dynamic_indices: list[ir.Value] | None = None, + loc: object = None, + ip: object = None, +) -> ir.Value: + """Helper for LLVM getelementptr operations.""" + if static_indices is None: + static_indices = [] + if dynamic_indices is None: + dynamic_indices = [] + + return llvm.getelementptr( + base.type, + base, + dynamic_indices, + static_indices, + _to_llvm_compatible_type(elem_type), + no_wrap_flags="None", + loc=loc, + ip=ip, + ) + + +class Pointer(ir.Value): + """An ``llvm.ptr`` value with element dtype metadata. + + ``Pointer`` is the canonical low-level DSL pointer type in the ``cutlass`` + namespace. It subclasses ``ir.Value`` so it can be passed directly to MLIR + ops that require a pointer operand. + """ + + _dtype: Type[Numeric] + _mlir_type: ir.Type + _addrspace: int + _base: ir.Value + + def __init__( + self, + base: ir.Value, + *, + dtype: Type[Numeric] | None = None, + space: AddressSpace | int | None = None, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> None: + if hasattr(base, "to_llvm_ptr"): + base = base.to_llvm_ptr(loc=loc, ip=ip) + elif hasattr(base, "ir_value") and not isinstance(base, ir.Value): + base = base.ir_value(loc=loc, ip=ip) + if not isinstance(base, ir.Value): + raise TypeError(f"Pointer expects an MLIR value, got {type(base).__name__}") + + super().__init__(base) + self._base = base + self._dtype = dtype or Int8 + self._mlir_type = self._dtype.mlir_type + + if space is not None: + self._addrspace = _normalize_address_space(space).value + else: + try: + self._addrspace = llvm.PointerType(base.type).address_space + except Exception: + self._addrspace = AddressSpace.generic.value + + @classmethod + def _from_raw_ptr( + cls, + value: ir.Value, + dtype: Type[Numeric] | None = None, + ) -> "Pointer": + return cls(value, dtype=dtype or Int8) + + def __class_getitem__(cls, args: Any) -> TypedPointer: + if not isinstance(args, tuple) or len(args) != 2: + raise DSLRuntimeError("Pointer[...] expects (dtype, memory_space)") + return TypedPointer(*args) + + def ir_value( + self, + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> ir.Value: + return self + + @property + def llvm_ptr(self) -> ir.Value: + return self + + def to_llvm_ptr( + self, + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> ir.Value: + return self + + def __extract_mlir_values__(self) -> list[ir.Value]: + return [ir.Value(self)] + + def __new_from_mlir_values__(self, values: list[ir.Value]) -> "Pointer": + return Pointer._from_raw_ptr(values[0], self._dtype) def __str__(self) -> str: - return f"{self.value} : {type(self)}" + return f"ptr" + + def __repr__(self) -> str: + return self.__str__() + + @property + def dtype(self) -> Type[Numeric]: + return self._dtype + + @property + def natural_alignment(self) -> int: + return max(1, self._dtype.width // 8) + + @property + def alignment(self) -> int: + return self.natural_alignment + + @property + def max_alignment(self) -> int: + return self.natural_alignment + + @property + def value_type(self) -> Type[Numeric]: + return self._dtype + + @property + def space(self) -> AddressSpace: + try: + return AddressSpace(self._addrspace) + except ValueError: + return AddressSpace.generic + + @property + def memspace(self) -> AddressSpace: + return self.space + + @property + def mlir_type(self) -> ir.Type: + return self.type + + def _effective_alignment(self, alignment: int | None) -> int: + return alignment if alignment is not None else self.natural_alignment + + def _prepare_store_value( + self, + value: Any, + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> Any: + if isinstance(value, Vector): + if value._dtype is not self._dtype: + raise TypeError(f"Expected dtype {self._dtype}, got {value._dtype}") + return value + if self._dtype.isinstance(value): + if isinstance(value, (int, float, bool)): + return arith_helper.const(value, loc=loc, ip=ip) + if isinstance(value, Numeric): + return value.ir_value(loc=loc, ip=ip) + return value + if hasattr(value, "_base") and hasattr(value, "_shape"): + raise TypeError( + "Cannot store Array directly. Use ptr.store(array.load()) instead." + ) + raise TypeError(f"Expected value to be Scalar or Vector, got {type(value)}") + + @dsl_user_op + def tospace( + self, + space: AddressSpace | int, + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> "Pointer": + target = _normalize_address_space(space) + if target.value == self._addrspace: + return self + if self._addrspace != 0 and target.value != 0: + raise ValueError( + "cannot cast from non-generic memory space to another non-generic memory space" + ) + res_ptr = llvm.addrspacecast( + llvm.PointerType.get(target.value), + self._base, + loc=loc, + ip=ip, + ) + return Pointer(res_ptr, dtype=self._dtype, space=target) + + def _as_tmem_offset( + self, + offset: Union[int, Integer, ir.Value], + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> "Int32": + if isinstance(offset, int): + return Int32(offset) + if isinstance(offset, Integer): + return Int32(offset, loc=loc, ip=ip) + if isinstance(offset, ir.Value) and ir.IntegerType.isinstance(offset.type): + return Int32(offset, loc=loc, ip=ip) + raise ValueError( + f"Expects static or dynamic integer value, but got {type(offset)}" + ) + + def _gep_tmem( + self, + offset: Union[int, Integer, ir.Value], + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> "Pointer": + # TMEM uses a packed i32 address token, so offsets are raw address + # increments rather than LLVM GEP element indices. + base_addr = self.toint(Int32, loc=loc, ip=ip) + offset_addr = self._as_tmem_offset(offset, loc=loc, ip=ip) + return self._inttoptr( + base_addr + offset_addr, + self._addrspace, + self._dtype, + loc=loc, + ip=ip, + ) + + def _gep( + self, + offset: Union[int, Integer, ir.Value], + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> "Pointer": + if self._addrspace == AddressSpace.tmem.value: + return self._gep_tmem(offset, loc=loc, ip=ip) + + if isinstance(offset, int): + static_indices = [offset] + dyn_indices = [] + elif isinstance(offset, Integer): + static_indices = [MLIR_DYNAMIC_INDEX] + dyn_indices = [offset.ir_value(loc=loc, ip=ip)] + elif isinstance(offset, ir.Value) and ir.IntegerType.isinstance(offset.type): + static_indices = [MLIR_DYNAMIC_INDEX] + dyn_indices = [offset] + else: + raise ValueError( + f"Expects static or dynamic integer value, but got {type(offset)}" + ) + + res_ptr = _gep( + self._base, + self._mlir_type, + static_indices=static_indices, + dynamic_indices=dyn_indices, + loc=loc, + ip=ip, + ) + return Pointer(res_ptr, dtype=self._dtype, space=self.space) + + @dsl_user_op + def toint( + self, + dtype: Optional[Type[Integer]] = None, + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> Integer: + addrspace = self._addrspace + if dtype is None: + dtype = Int64 + if addrspace == 6: + return dtype(llvm.ptrtoint(Int32.mlir_type, self._base, loc=loc, ip=ip)) + return dtype(llvm.ptrtoint(dtype.mlir_type, self._base, loc=loc, ip=ip)) + + @dsl_user_op + def load( + self, + alignment: Optional[int] = None, + *, + count: Optional[int] = None, + is_volatile: bool = False, + is_invariant: bool = False, + is_invariant_group: bool = False, + ordering: Any = "not_atomic", + syncscope: Any | None = None, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> Any: + effective_alignment = self._effective_alignment(alignment) + if count is None: + res_ty = self._mlir_type + else: + res_ty = ir.VectorType.get([count], self._mlir_type) + + res = llvm.load( + res_ty, + self._base, + alignment=effective_alignment, + volatile_=is_volatile, + nontemporal=False, + invariant=is_invariant, + invariant_group=is_invariant_group, + ordering=( + llvm.AtomicOrdering(ordering) if ordering != "not_atomic" else None + ), + syncscope=syncscope, + loc=loc, + ip=ip, + ) + if count is None: + return self._dtype(res, loc=loc, ip=ip) + if hasattr(res, "ir_value"): + res = res.ir_value() + return Vector(res, dtype=self._dtype) + + def nvvm_load_ext( + self, + *, + count: Optional[int] = None, + cache_modifier: Any | None = None, + evict: Any | None = None, + scope: Any | None = None, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> Any: + if count is None: + res_ty = self._mlir_type + else: + res_ty = ir.VectorType.get([count], self._mlir_type) + + kwargs: dict = {} + if cache_modifier is not None: + name = ( + cache_modifier + if isinstance(cache_modifier, str) + else cache_modifier.name + ) + kwargs["cache_modifier"] = getattr( + nvvm.LoadCacheModifierExtKind, name.upper() + ) + if evict is not None: + name = evict if isinstance(evict, str) else evict.name + kwargs["evict"] = getattr(nvvm.EvictKind, name.upper()) + if scope is not None: + name = scope if isinstance(scope, str) else scope.name + kwargs["scope"] = getattr(nvvm.MemScopeKind, name.upper()) + + res = nvvm.load_ext(res_ty, self._base, **kwargs, loc=loc, ip=ip) + if count is None: + return self._dtype(res, loc=loc, ip=ip) + if hasattr(res, "ir_value"): + res = res.ir_value() + return Vector(res, dtype=self._dtype) + + def store( + self, + value: Any, + *, + alignment: Optional[int] = None, + is_volatile: bool = False, + is_invariant_group: bool = False, + ordering: Any = "not_atomic", + syncscope: Any | None = None, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> None: + effective_alignment = self._effective_alignment(alignment) + ir_value = self._prepare_store_value(value, loc=loc, ip=ip) + return llvm.store( + ir_value, + self._base, + alignment=effective_alignment, + volatile_=is_volatile, + nontemporal=False, + invariant_group=is_invariant_group, + ordering=( + llvm.AtomicOrdering(ordering) if ordering != "not_atomic" else None + ), + syncscope=syncscope if syncscope else None, + loc=loc, + ip=ip, + ) + + def nvvm_store_ext( + self, + value: Any, + *, + cache_modifier: Any | None = None, + evict: Any | None = None, + scope: Any | None = None, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> None: + ir_value = self._prepare_store_value(value, loc=loc, ip=ip) + kwargs: dict = {} + if cache_modifier is not None: + name = ( + cache_modifier + if isinstance(cache_modifier, str) + else cache_modifier.name + ) + kwargs["cache_modifier"] = getattr( + nvvm.StoreCacheModifierKind, name.upper() + ) + if evict is not None: + name = evict if isinstance(evict, str) else evict.name + kwargs["evict"] = getattr(nvvm.EvictKind, name.upper()) + if scope is not None: + name = scope if isinstance(scope, str) else scope.name + kwargs["scope"] = getattr(nvvm.MemScopeKind, name.upper()) + + return nvvm.store_ext(ir_value, self._base, **kwargs, loc=loc, ip=ip) + + @dsl_user_op + def masked_load( + self, + mask: Any, + pass_thru: Any | None = None, + *, + alignment: Optional[int] = None, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> Any: + from cutlass._mlir_helpers.vector import Vector # noqa: PLC0415 + + if not isinstance(mask, Vector): + raise TypeError(f"mask must be a Vector, got {type(mask)}") + mask_elem_type = mask._mlir_type + if not ( + ir.IntegerType.isinstance(mask_elem_type) + and ir.IntegerType(mask_elem_type).width == 1 + ): + raise TypeError( + f"mask must be a vector of i1, got vector of {mask_elem_type}" + ) + + vec_len = mask.shape[0] + if pass_thru is not None: + if not isinstance(pass_thru, Vector): + raise TypeError(f"pass_thru must be a Vector, got {type(pass_thru)}") + if pass_thru.shape[0] != vec_len: + raise ValueError( + f"pass_thru length ({pass_thru.shape[0]}) must match mask length ({vec_len})" + ) + if pass_thru.dtype is not self.dtype: + raise TypeError( + f"pass_thru dtype ({pass_thru.dtype}) must match pointer dtype ({self.dtype})" + ) + + res_ty = ir.VectorType.get([vec_len], self.dtype.mlir_type) + res = llvm.intr_masked_load( + res_ty, + self._base, + mask, + pass_thru=pass_thru, + alignment=alignment or self.natural_alignment, + loc=loc, + ip=ip, + ) + if isinstance(res, Vector): + return res + if hasattr(res, "ir_value"): + res = res.ir_value() + return Vector(res, dtype=self.dtype) + + @dsl_user_op + def masked_store( + self, + value: Any, + mask: Any, + *, + alignment: Optional[int] = None, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> None: + from cutlass._mlir_helpers.vector import Vector # noqa: PLC0415 + + if not isinstance(value, Vector): + raise TypeError(f"value must be a Vector, got {type(value)}") + if not isinstance(mask, Vector): + raise TypeError(f"mask must be a Vector, got {type(mask)}") + + mask_elem_type = mask._mlir_type + if not ( + ir.IntegerType.isinstance(mask_elem_type) + and ir.IntegerType(mask_elem_type).width == 1 + ): + raise TypeError( + f"mask must be a vector of i1, got vector of {mask_elem_type}" + ) + if value.shape[0] != mask.shape[0]: + raise ValueError( + f"value and mask must have the same length, got {value.shape[0]} and {mask.shape[0]}" + ) + if value.dtype is not self.dtype: + raise TypeError( + f"value dtype ({value.dtype}) must match pointer dtype ({self.dtype})" + ) + + llvm.intr_masked_store( + value, + self._base, + mask, + alignment=alignment or self.natural_alignment, + loc=loc, + ip=ip, + ) + + @dsl_user_op + def masked_gather( + self, + offsets: Any, + mask: Any, + pass_thru: Any | None = None, + *, + alignment: Optional[int] = None, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> Any: + from cutlass._mlir_helpers.vector import Vector # noqa: PLC0415 + + if not isinstance(offsets, Vector): + raise TypeError(f"offsets must be a Vector, got {type(offsets)}") + if not isinstance(mask, Vector): + raise TypeError(f"mask must be a Vector, got {type(mask)}") + if offsets.shape[0] != mask.shape[0]: + raise ValueError( + f"offsets and mask must have the same length, got {offsets.shape[0]} and {mask.shape[0]}" + ) + mask_elem_type = mask._mlir_type + if not ( + ir.IntegerType.isinstance(mask_elem_type) + and ir.IntegerType(mask_elem_type).width == 1 + ): + raise TypeError( + f"mask must be a vector of i1, got vector of {mask_elem_type}" + ) + + vec_len = offsets.shape[0] + vec_ptr_type = ir.VectorType.get([vec_len], self.mlir_type) + base_ptr_vec = vector.broadcast(vec_ptr_type, self._base, loc=loc, ip=ip) + ptrs = llvm.getelementptr( + vec_ptr_type, + base_ptr_vec, + [offsets], + [MLIR_DYNAMIC_INDEX], + self.dtype.mlir_type, + no_wrap_flags="None", + loc=loc, + ip=ip, + ) + result_type = ir.VectorType.get([vec_len], self.dtype.mlir_type) + result = llvm.intr_masked_gather( + result_type, + ptrs, + mask, + [pass_thru] if pass_thru is not None else [], + alignment=alignment or self.natural_alignment, + loc=loc, + ip=ip, + ) + if isinstance(result, Vector): + return result + if hasattr(result, "ir_value"): + result = result.ir_value() + return Vector(result, dtype=self.dtype) + + @dsl_user_op + def __add__( + self, + offset: Union[int, Integer, ir.Value], + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> "Pointer": + return self._gep(offset, loc=loc, ip=ip) + + @dsl_user_op + def __sub__( + self, + offset: Union[int, Integer, ir.Value], + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> "Pointer": + if isinstance(offset, int): + neg_offset = -offset + else: + neg_offset = Int32(0) - offset + return self._gep(neg_offset, loc=loc, ip=ip) + + @dsl_user_op + def __radd__( + self, + offset: Union[int, Integer, ir.Value], + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> "Pointer": + return self._gep(offset, loc=loc, ip=ip) + + @dsl_user_op + def __iadd__( + self, + offset: Union[int, Integer, ir.Value], + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> "Pointer": + return self._gep(offset, loc=loc, ip=ip) + + @dsl_user_op + def __isub__( + self, + offset: Union[int, Integer, ir.Value], + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> "Pointer": + return self.__sub__(offset, loc=loc, ip=ip) + + def _inttoptr( + self, + value: Any, + addrspace: int, + dtype: Type[Numeric], + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> "Pointer": + if isinstance(value, int): + value = (Int64 if addrspace in (0, 1) else Int32)(value) + if hasattr(value, "ir_value"): + value = value.ir_value(loc=loc, ip=ip) + res_val = llvm.inttoptr(llvm.PointerType.get(addrspace), value, loc=loc, ip=ip) + return Pointer._from_raw_ptr(res_val, dtype) + + @dsl_user_op + def __and__( + self, + mask: Union[int, Integer, ir.Value], + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> "Pointer": + masked_int = self.toint(loc=loc, ip=ip) & mask + return self._inttoptr(masked_int, self._addrspace, self.dtype, loc=loc, ip=ip) + + @dsl_user_op + def __rand__( + self, + mask: Union[int, Integer, ir.Value], + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> "Pointer": + return self.__and__(mask, loc=loc, ip=ip) + + def _validate_scalar_index(self, idx: Any) -> Any: + if isinstance(idx, tuple): + raise TypeError( + "Pointer tuple indexing is not supported; use explicit " + "load(..., alignment=...) or store(..., alignment=...) calls" + ) + if isinstance(idx, slice): + raise TypeError( + "Pointer slices are not supported; use explicit " + "load(count=...) or store(...) calls" + ) + return idx + + def __getitem__( + self, + idx: Any, + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> Any: + offset = self._validate_scalar_index(idx) + ptr = self._gep(offset, loc=loc, ip=ip) + return ptr.load(loc=loc, ip=ip) + + @dsl_user_op + def __setitem__( + self, + idx: Any, + value: Any, + *, + loc: ir.Location | None = None, + ip: ir.InsertionPoint | None = None, + ) -> None: + offset = self._validate_scalar_index(idx) + ptr = self._gep(offset, loc=loc, ip=ip) + return ptr.store(value, loc=loc, ip=ip) class IRConst(Generic[TY]): @@ -2468,9 +3232,20 @@ __all__ = [ "Float4E2M1FN", "Float6E2M3FN", "Float6E3M2FN", + "Float4E2M1FNx2", "as_numeric", "align", + "AddressSpace", "Pointer", + "TypedPointer", + "MemOrdering", + "MemScope", + "SharedSpace", + "LoadCacheModifier", + "StoreCacheModifier", + "EvictPriority", + "L2PrefetchSize", + "L1EvictKind", "dtype", "Constexpr", "IRConst", @@ -2478,3 +3253,5 @@ __all__ = [ "IRVariadic", "implicitDowncastNumericType", ] + + diff --git a/python/CuTeDSL/cutlass/base_dsl/utils/tree_utils.py b/python/CuTeDSL/cutlass/base_dsl/utils/tree_utils.py index f4a973b00..91aeff0bc 100644 --- a/python/CuTeDSL/cutlass/base_dsl/utils/tree_utils.py +++ b/python/CuTeDSL/cutlass/base_dsl/utils/tree_utils.py @@ -10,13 +10,20 @@ # is strictly prohibited. from collections.abc import Callable, Iterable, Iterator -from typing import Any, NamedTuple, get_origin +from typing import Any, NamedTuple, get_origin, get_type_hints import dataclasses +import functools import itertools as it from types import SimpleNamespace -from ..typing import as_numeric, Numeric, Constexpr, implements_dynamic_expression +from ..typing import ( + as_numeric, + Numeric, + Constexpr, + implements_dynamic_expression, + implements_jit_argument, +) from ..._mlir_helpers.arith import ArithValue from ..common import DSLBaseError from ..._mlir import ir @@ -165,15 +172,36 @@ def is_namedtuple_instance(x: Any) -> bool: ) -def is_constexpr_field(field: dataclasses.Field) -> bool: +@functools.lru_cache(maxsize=None) +def _resolved_type_hints(cls: type) -> dict: + """Resolve a dataclass's annotations, including PEP 563 string forms. + + Cached per class -- ``get_type_hints`` walks the MRO and the defining + module's globals, too costly to repeat on every flatten / adapt call. + Returns ``{}`` if resolution fails (e.g. an unresolvable forward + reference), so callers fall back to the raw annotation. """ - Check if a field is a constexpr field. + try: + return get_type_hints(cls) + except Exception: + return {} + + +def is_constexpr_field(field: dataclasses.Field, owner: Any = None) -> bool: + """Check if a dataclass field is annotated ``Constexpr`` / ``Constexpr[T]``. + + Under ``from __future__ import annotations`` (PEP 563) ``field.type`` is a + string (e.g. ``"Constexpr[int]"``); pass ``owner`` -- the dataclass type or + an instance of it -- so the annotation is resolved via ``get_type_hints`` + before the check. Without ``owner`` only already-resolved annotations are + recognized (a stringized one is otherwise misread as non-constexpr, which + silently keeps a constexpr field as a runtime argument). """ - if field.type is Constexpr: - return True - elif get_origin(field.type) is Constexpr: - return True - return False + annotation = field.type + if isinstance(annotation, str) and owner is not None: + cls = owner if isinstance(owner, type) else type(owner) + annotation = _resolved_type_hints(cls).get(field.name, annotation) + return annotation is Constexpr or get_origin(annotation) is Constexpr # ============================================================================= @@ -261,7 +289,7 @@ def extract_dataclass_members(x: Any) -> tuple[list[str], list[Any], list[str]]: members = [] constexpr_fields = [] for field in dataclasses.fields(x): - if is_constexpr_field(field): + if is_constexpr_field(field, x): constexpr_fields.append(field.name) fields.remove(field.name) v = getattr(x, field.name) @@ -759,11 +787,49 @@ def _tree_flatten( elif isinstance(x, ArithValue): return [x], [ir.DictAttr.get({})], create_leaf_for_value(x, is_numeric=True) + elif ( + implements_dynamic_expression(x) + and implements_jit_argument(x) + and x.__extract_mlir_values__() is None + ): + if return_ir_values: + raise DSLTreeFlattenError( + "Flatten Error: children is None", + get_fully_qualified_class_name(x), + ) + return ( + [x], + [ir.DictAttr.get({})], + create_leaf_for_value( + x, + node_metadata=SimpleNamespace(is_dynamic_expression=1, original_obj=x), + ir_type_str="unknown", + ), + ) + elif implements_dynamic_expression(x) and isinstance(x, ir.Value): - # Only for ir.Value subclasses (e.g. Pointer). Check before plain ir.Value + # Only for ir.Value subclasses. Check before plain ir.Value # so they are unflattened via __new_from_mlir_values__. Other dynamic # expressions (e.g. TmemAllocator with 2 values) use the registered/node path. - v = _flatten_mlir_values(x.__extract_mlir_values__()) + extracted = x.__extract_mlir_values__() + if extracted is None: + if return_ir_values: + raise DSLTreeFlattenError( + "Flatten Error: children is None", + get_fully_qualified_class_name(x), + ) + return ( + [x], + [ir.DictAttr.get({})], + create_leaf_for_value( + x, + node_metadata=SimpleNamespace( + is_dynamic_expression=1, original_obj=x + ), + ir_type_str="unknown", + ), + ) + v = _flatten_mlir_values(extracted) a = ( [ir.DictAttr.get({})] if not hasattr(x, "__extract_mlir_attributes__") diff --git a/python/CuTeDSL/cutlass/base_dsl/vector.py b/python/CuTeDSL/cutlass/base_dsl/vector.py new file mode 100644 index 000000000..1702dd472 --- /dev/null +++ b/python/CuTeDSL/cutlass/base_dsl/vector.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +"""Vector helper APIs for base DSL register vectors.""" + +from .._mlir_helpers.vector import ( + all_, + any_, + broadcast_to, + create_mask, + empty_like, + full, + full_like, + ones_like, + outerproduct, + print_nd_vector, + where, + zeros_like, +) + +__all__ = [ + "all_", + "any_", + "broadcast_to", + "create_mask", + "empty_like", + "full", + "full_like", + "ones_like", + "outerproduct", + "print_nd_vector", + "where", + "zeros_like", +] diff --git a/python/CuTeDSL/cutlass/compiler/__init__.py b/python/CuTeDSL/cutlass/compiler/__init__.py index 4990b2f8d..8c334c7f3 100644 --- a/python/CuTeDSL/cutlass/compiler/__init__.py +++ b/python/CuTeDSL/cutlass/compiler/__init__.py @@ -10,19 +10,44 @@ # is strictly prohibited. from .._mlir._mlir_libs._cutlass_ir import ( + # Artifact and compiler surfaces. ArtifactType, CompilationArtifact, - ParamKind, - ParamInfo, - FunctionInfo, PreCompiledMlirArtifact, CompiledMlirArtifact, LlvmIrArtifact, ObjectArtifact, - ExecutableFunction, - Executor, Compiler, CuteCompiler, - serialize_compilation_artifact, + Executor, + ExecutableFunction, + # Metadata enums. + Abi, + DTypeCode, + BindingTypeId, + AddressSpace, + # Metadata value types. + FunctionMetadata, + Binding, + DType, + DimSymbol, + SymbolId, + Dim, + Var, + Const, + Shape, + Tensor, + Stream, + EnvStream, + Pointer, + Unit, + Tuple, + # Bound list classes returned by .shape / .strides / .params / .functions. + DimList, + DimSymbolList, + BindingList, + FunctionMetadataList, + # Serialization helpers. deserialize_compilation_artifact, + serialize_compilation_artifact, ) diff --git a/python/CuTeDSL/cutlass/cute/__init__.py b/python/CuTeDSL/cutlass/cute/__init__.py index 91e215cd2..89777742c 100644 --- a/python/CuTeDSL/cutlass/cute/__init__.py +++ b/python/CuTeDSL/cutlass/cute/__init__.py @@ -12,8 +12,14 @@ from collections.abc import Callable from typing import Any -# Use the auto-generated enum AddressSpace -from cutlass._mlir.dialects.cute import AddressSpace, CacheEvictionPriority +# Keep AddressSpace aligned with the public cutlass namespace. +from cutlass.address_space import AddressSpace as _AddressSpace +from cutlass._mlir.dialects.cute_nvgpu import CacheEvictionPriority + +# Keep ``cute.AddressSpace`` as a quiet compatibility alias. Documentation marks +# this spelling as scheduled for deprecation; runtime warnings will be enabled in +# a later release. +AddressSpace = _AddressSpace # Explicitly import types that might be directly used by other modules. # This is a fix for using Sphinx to generate documentation @@ -212,7 +218,7 @@ from . import experimental # Export all math ops without "math." from .math import * -# Used as internal symbol +# Package-private symbol used by exported aliases below. from .. import cutlass_dsl as _dsl from .ffi import ffi, extern, BitCode, ConstValue, mangle @@ -230,6 +236,8 @@ EnableAssertions = _dsl.EnableAssertions GenerateLineInfo = _dsl.GenerateLineInfo KeepCUBIN = _dsl.KeepCUBIN KeepPTX = _dsl.KeepPTX +KeepSASS = _dsl.KeepSASS +NvdisasmOptions = _dsl.NvdisasmOptions GPUArch = _dsl.GPUArch LinkLibraries = _dsl.LinkLibraries EnableTVMFFI = _dsl.EnableTVMFFI @@ -242,6 +250,9 @@ make_native_struct = _dsl.make_native_struct # factory for dynamic struct types from . import _tvm_ffi_args_spec_converter _tvm_ffi_args_spec_converter.attach_args_spec_converter(_dsl.CuTeDSL._get_dsl()) +_tvm_ffi_args_spec_converter.attach_args_spec_converter( + _dsl.CuteExperimentalDSL._get_dsl() +) # Explicitly export all symbols for documentation generation __all__ = [ @@ -251,6 +262,7 @@ __all__ = [ "CacheEvictionPriority", "Tensor", "Layout", + "Numeric", "ComposedLayout", "Swizzle", "E", diff --git a/python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py b/python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py index 90a07d982..849a30a16 100644 --- a/python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py +++ b/python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py @@ -10,13 +10,13 @@ # is strictly prohibited. +from cutlass.address_space import AddressSpace from dataclasses import is_dataclass, fields as dataclass_fields from cutlass.base_dsl.tvm_ffi_builder import spec from cutlass.base_dsl.jit_executor import ExecutionArgs from cutlass.base_dsl.common import DSLRuntimeError from cutlass.base_dsl.utils.tree_utils import is_constexpr_field from cutlass.cutlass_dsl import is_cute_algebra_type -from cutlass._mlir.dialects import cute as _cute_ir from .runtime import _FakeStream from .typing import Tensor, Pointer, SymInt from .typing import ( @@ -42,6 +42,7 @@ from .typing import ( Float8E4M3FN, Float8E8M0FNU, Float4E2M1FN, + Float4E2M1FNx2, ) import cuda.bindings.driver as cuda @@ -80,6 +81,9 @@ NumericToTVMFFIDtype = { Float4E2M1FN: "float4_e2m1fn", Float6E2M3FN: "float6_e2m3fn", Float6E3M2FN: "float6_e3m2fn", + # Packed FP4 maps to the tvm-ffi packed dtype (matches torch's + # ``float4_e2m1fn_x2``; tvm-ffi spells it without the underscore). + Float4E2M1FNx2: "float4_e2m1fnx2", } _UNSUPPORTED_TVM_FFI_NUMERIC_TYPES = set[Any]() @@ -124,17 +128,17 @@ AcceptableNumericTypesForScalar = [ def _get_llvm_address_space_from_memspace( - memspace: _cute_ir.AddressSpace, + memspace: AddressSpace, ) -> Optional[int]: - if memspace == _cute_ir.AddressSpace.gmem: + if memspace == AddressSpace.gmem: return 1 return None def _is_gpu_memspace( - memspace: _cute_ir.AddressSpace, + memspace: AddressSpace, ) -> bool: - return memspace != _cute_ir.AddressSpace.generic + return memspace != AddressSpace.generic def _native_struct_type(value: Any) -> Optional[type]: @@ -337,7 +341,7 @@ def _convert_single_arg( assert isinstance(arg_shape, tuple) assert isinstance(arg_stride, tuple) shapes: List[Any] = [] - for i, dyn_mask in enumerate(arg.dynamic_shapes_mask): # type: ignore[attr-defined] + for i, dyn_mask in enumerate(arg.dynamic_shapes_mask): if not dyn_mask: shapes.append(arg_shape[i]) elif isinstance(arg_shape[i], SymInt): @@ -350,7 +354,7 @@ def _convert_single_arg( ) strides: List[Any] = [] - for i, dyn_mask in enumerate(arg.dynamic_strides_mask): # type: ignore[attr-defined] + for i, dyn_mask in enumerate(arg.dynamic_strides_mask): if not dyn_mask: strides.append(arg_stride[i]) elif isinstance(arg_stride[i], SymInt): @@ -391,7 +395,7 @@ def _convert_single_arg( tvm_ffi_cute_tensor = spec.Tensor( arg_name, shapes, - _numeric_to_tvm_ffi_dtype(arg.element_type), # type: ignore[arg-type] + _numeric_to_tvm_ffi_dtype(arg.element_type), strides=strides, data_alignment=arg._assumed_align, # type: ignore[attr-defined] device_type=device_type, @@ -563,9 +567,9 @@ def _tvm_ffi_args_spec_converter( wrapper_extra_exclude_arg_names.append(arg_name) # Covers both plain ``op: AddOp`` and ``op: AddOp | MulOp`` since we # check the runtime type when the annotation isn't itself a dataclass. - if ( - isinstance(arg_type, type) and is_dataclass(arg_type) - ) or is_dataclass(type(arg)): + if (isinstance(arg_type, type) and is_dataclass(arg_type)) or is_dataclass( + type(arg) + ): map_dataclass_to_tuple.append(arg_name) kwargs_wrapper_spec = exec_args.get_kwargs_wrapper_spec( wrapper_extra_exclude_arg_names diff --git a/python/CuTeDSL/cutlass/cute/arch/__init__.py b/python/CuTeDSL/cutlass/cute/arch/__init__.py index 545b5b965..be27bf860 100644 --- a/python/CuTeDSL/cutlass/cute/arch/__init__.py +++ b/python/CuTeDSL/cutlass/cute/arch/__init__.py @@ -81,6 +81,8 @@ __all__ = [ "vote_uni_sync", "warp_redux_sync", "atomic_max_float32", # Deprecated: use atomic_fmax + "atomic_fmax", + "atomic_fmin", "atomic_add", "atomic_and", "atomic_or", @@ -101,6 +103,12 @@ __all__ = [ "setmaxregister_increase", "setmaxregister_decrease", "fma_packed_f32x2", + "fma_f16", + "fma_bf16", + "fma_packed_f16x2", + "fma_packed_bf16x2", + "fma_packed_f32x2_f16x2_f32x2_f32x2", + "fma_packed_f32x2_bf16x2_f32x2_f32x2", "mul_packed_f32x2", "add_packed_f32x2", "sub_packed_f32x2", @@ -113,10 +121,18 @@ __all__ = [ "cvt_i8_bf16", "cvt_i8x2_to_bf16x2", "cvt_i8x4_to_bf16x4", + "cvt_f32_tf32", "cvt_f32x2_bf16x2", "warp_redux_sync", "smid", "nsmid", + "total_smem_size", + "aggr_smem_size", + "gridid", + "nwarpid", + "warpsize", + "globaltimer", + "globaltimer_lo", "clock", "clock64", "match_sync", diff --git a/python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py b/python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py index e66013863..c5f182a81 100644 --- a/python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py +++ b/python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py @@ -9,7 +9,7 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. -from typing import Optional +from typing import Optional, Type from cutlass.base_dsl.arch import Arch from cutlass.base_dsl.common import DSLRuntimeError @@ -30,9 +30,10 @@ from .nvvm_wrappers import ( cvt_i4x2_to_bf16x2, cvt_i4_bf16, cvt_f32_bf16, + cvt_i4x8_to_mxf8x8, sext_unpacked_i4x4_to_i8x4, ) -from ..typing import Int4, Int8, Float32, BFloat16, Int32 +from ..typing import Int4, Int8, Float32, BFloat16, Int32, Numeric @dsl_user_op @@ -257,6 +258,69 @@ def cvt_i4_bf16_intrinsic( return vec_dst +@dsl_user_op +def cvt_i4_mxf8_intrinsic( + vec_i4: ir.Value, + length: int, + dst_type: Type[Numeric], + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> ir.Value: + """Convert an int4 vector to FP8 operand bytes for MXF8 MMA. + + This helper is used by mixed-input kernels that stage int4 weights as FP8 + values before issuing ``tcgen05.mma``. It processes the input in 8-element + chunks via :func:`cvt_i4x8_to_mxf8x8` and returns a vector with the same + logical length as ``vec_i4``. + + :param vec_i4: Input 1D vector of int4 values. + :type vec_i4: vector of int4 + :param length: Number of int4 elements in ``vec_i4``. Must be a multiple of 8. + :type length: int + :param dst_type: FP8 destination numeric type, typically ``cutlass.Float8E4M3FN``. + :type dst_type: type + :return: Output 1D vector with ``length`` FP8 elements. + :rtype: vector + """ + # n.b. Build the intermediate vector as int8 bytes because LLVM conversion + # can fail when materializing MXFP8 vector constants directly. + assert length % 8 == 0, f"Invalid {length=}" + dst_mlir_type = dst_type.mlir_type + src_pos = 0 + vec_i4x8_type = ir.VectorType.get([8], Int4.mlir_type, loc=loc) + vec_dst_type = ir.VectorType.get([length], dst_mlir_type, loc=loc) + vec_dst_type_i8 = ir.VectorType.get([length], Int8.mlir_type, loc=loc) + zero_attr = ir.IntegerAttr.get(Int8.mlir_type, 0) + + vec_dst = arith.ConstantOp( + vec_dst_type_i8, + ir.DenseElementsAttr.get_splat(vec_dst_type_i8, zero_attr), + loc=loc, + ip=ip, + ).result + + # try to use vectorized version + if length >= 8: + num_vec8 = length // 8 + for _ in range(num_vec8): + vec_i4x8 = vector.extract_strided_slice( + vec_i4x8_type, vec_i4, [src_pos], [8], [1], loc=loc, ip=ip + ) + vec_mxf8x8 = cvt_i4x8_to_mxf8x8( + vec_i4x8, dst_type=Int8.mlir_type, loc=loc, ip=ip + ) + vec_dst = vector.insert_strided_slice( + vec_mxf8x8, vec_dst, [src_pos], [1], loc=loc, ip=ip + ) + src_pos += 8 + length -= 8 + + if dst_mlir_type != Int8.mlir_type: + vec_dst = vector.bitcast(vec_dst_type, vec_dst, loc=loc, ip=ip) + return vec_dst + + @dsl_user_op def sext_unpacked_i4_i8_intrinsic( vec_unpacked_i4: ir.Value, diff --git a/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py b/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py index 46b49d8fe..a9f0573a2 100644 --- a/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py +++ b/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py @@ -23,6 +23,7 @@ from cutlass._mlir import ir from cutlass._mlir.dialects import arith, builtin, llvm, math, nvvm as _nvvm_raw, vector from .constants import WARP_SIZE +from .elect import make_warp_uniform from ..core import size @@ -171,7 +172,6 @@ def _enhance_enum_with_str_mapping(enum_class: Any) -> Any: str_to_enum_map[str_repr] = member # Add from_str class method - @classmethod # type: ignore[misc] def from_str(cls: Any, s: Any) -> Any: """ Convert a string literal to the corresponding enum member. @@ -204,7 +204,7 @@ def _enhance_enum_with_str_mapping(enum_class: Any) -> Any: ) return str_to_enum_map[s] - enum_class.from_str = from_str + enum_class.from_str = classmethod(from_str) return enum_class @@ -331,7 +331,7 @@ def warp_idx( ntid_x = Int32(nvvm.read_ptx_sreg_ntid_x(T.i32(), loc=loc, ip=ip)) ntid_y = Int32(nvvm.read_ptx_sreg_ntid_y(T.i32(), loc=loc, ip=ip)) tid = tid_x + tid_y * ntid_x + tid_z * ntid_x * ntid_y - return tid // WARP_SIZE + return make_warp_uniform(tid // WARP_SIZE, loc=loc, ip=ip) @dsl_user_op @@ -488,6 +488,119 @@ def dynamic_smem_size( return Int32(nvvm.read_ptx_sreg_dynamic_smem_size(T.i32(), loc=loc, ip=ip)) +@dsl_user_op +def total_smem_size( + *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> Int32: + """ + Returns the total shared memory reserved for the CTA, in bytes. + + This is the static plus dynamic shared-memory size for the CTA. + + See https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-total-smem-size + + :return: Total CTA shared-memory size in bytes + :rtype: Int32 + """ + return Int32(nvvm.read_ptx_sreg_total_smem_size(T.i32(), loc=loc, ip=ip)) + + +@dsl_user_op +def aggr_smem_size( + *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> Int32: + """ + Returns the aggregate shared memory across the CTA's cluster, in bytes (sm_90+). + + See https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-aggr-smem-size + + :return: Aggregate cluster shared-memory size in bytes + :rtype: Int32 + """ + return Int32(nvvm.read_ptx_sreg_aggr_smem_size(T.i32(), loc=loc, ip=ip)) + + +@dsl_user_op +def gridid( + *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> Int32: + """ + Returns a unique identifier for the current grid launch. + + See https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-gridid + + :return: Grid launch identifier + :rtype: Int32 + """ + return Int32(nvvm.read_ptx_sreg_gridid(T.i32(), loc=loc, ip=ip)) + + +@dsl_user_op +def nwarpid( + *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> Int32: + """ + Returns the maximum number of warp slots (warp IDs) per SM. + + Defines the valid range for physical_warp_id() as [0, nwarpid() - 1]. + + See https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-nwarpid + + :return: Number of warp slots per SM + :rtype: Int32 + """ + return Int32(nvvm.read_ptx_sreg_nwarpid(T.i32(), loc=loc, ip=ip)) + + +@dsl_user_op +def warpsize( + *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> Int32: + """ + Returns the warp size in threads as a runtime register read (``%WARP_SZ``). + + Prefer the compile-time ``WARP_SIZE`` constant unless a runtime read is + specifically required; both are 32 on current hardware. + + See https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-warpid-nwarpid-warpsize + + :return: Warp size in threads + :rtype: Int32 + """ + return Int32(nvvm.read_ptx_sreg_warpsize(T.i32(), loc=loc, ip=ip)) + + +@dsl_user_op +def globaltimer( + *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> Int64: + """ + Returns the 64-bit global nanosecond timer (``%globaltimer``). + + Intended for use by NVIDIA tools: the behavior is target-specific and may + change or be removed on future architectures. + + See https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-globaltimer-globaltimer-lo-globaltimer-hi + + :return: Global timer value in nanoseconds + :rtype: Int64 + """ + return Int64(nvvm.read_ptx_sreg_globaltimer(T.i64(), loc=loc, ip=ip)) + + +@dsl_user_op +def globaltimer_lo( + *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> Int32: + """ + Returns the low 32 bits of the global nanosecond timer (``%globaltimer_lo``). + + :return: Low 32 bits of the global timer + :rtype: Int32 + """ + return Int32(nvvm.read_ptx_sreg_globaltimer_lo(T.i32(), loc=loc, ip=ip)) + + @dsl_user_op def shuffle_sync_op( value: Union[Numeric, "TensorSSA"], # type: ignore[name-defined] # noqa: F821 @@ -1308,13 +1421,21 @@ def fmax( a: Union[float, Float32], b: Union[float, Float32], *, + abs: bool = False, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None, ) -> Float32: + """Floating-point max via ``nvvm.fmax`` (FMNMX, not ``arith.max``/SEL). + + :param abs: When True, lower to the xorsign-abs form + ``sign(a ^ b) * max(|a|, |b|)`` (FMNMX.XORSIGN.ABS); default False is a + plain max. NaN / signed-zero behavior follows the underlying NVVM op. + """ return Float32( nvvm.fmax( Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip), + abs=abs, loc=loc, ip=ip, ) @@ -1326,13 +1447,21 @@ def fmin( a: Union[float, Float32], b: Union[float, Float32], *, + abs: bool = False, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None, ) -> Float32: + """Floating-point min via ``nvvm.fmin`` (FMNMX, not ``arith.min``/SEL). + + :param abs: When True, lower to the xorsign-abs form + ``sign(a ^ b) * min(|a|, |b|)`` (FMNMX.XORSIGN.ABS); default False is a + plain min. NaN / signed-zero behavior follows the underlying NVVM op. + """ return Float32( nvvm.fmin( Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip), + abs=abs, loc=loc, ip=ip, ) @@ -1527,6 +1656,30 @@ def cvt_f32_bf16( return bf16_val +@dsl_user_op +def cvt_f32_tf32( + src_f32: Union[float, Float32, ir.Value], + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> ir.Value: + """Convert one f32 value to a TF32 MMA operand register.""" + src_val = ( + src_f32 + if isinstance(src_f32, ir.Value) + else Float32(src_f32).ir_value(loc=loc, ip=ip) + ) + return llvm.inline_asm( + Int32.mlir_type, + [src_val], + "cvt.rna.tf32.f32 $0, $1;", + "=r,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + # Convert vector of 4 int8 values to vector of 4 float32 values @dsl_user_op def cvt_i8x4_to_f32x4( @@ -2011,6 +2164,101 @@ def cvt_i4x8_to_bf16x8( return vec_bf16x8 +@dsl_user_op +def cvt_i4x8_to_mxf8x8_impl( + src_i32: ir.Value, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> ir.Value: + """Convert eight packed int4 values to two int32 words of FP8 bytes. + + ``src_i32`` holds eight signed int4 values, one nibble per element. The + inline PTX uses ``prmt.b32`` with positive and negative E4M3 lookup tables + to produce the raw FP8 byte encodings in two 32-bit words. + + :param src_i32: Packed int4 source value containing eight 4-bit elements. + :type src_i32: i32 + :return: Vector of two i32 values, each containing four packed FP8 bytes. + :rtype: vector<2xi32> + """ + sign_mask = arith.constant(Int32.mlir_type, 0x88888888, loc=loc, ip=ip) + lut_idx_mask = arith.constant(Int32.mlir_type, 0x77777777, loc=loc, ip=ip) + final_prmt_base = arith.constant(Int32.mlir_type, 0x32103210, loc=loc, ip=ip) + one = arith.constant(Int32.mlir_type, 1, loc=loc, ip=ip) + sixteen = arith.constant(Int32.mlir_type, 16, loc=loc, ip=ip) + + POS_E4M3s_REG1 = arith.constant(Int32.mlir_type, 0x44403800, loc=loc, ip=ip) + POS_E4M3s_REG2 = arith.constant(Int32.mlir_type, 0x4E4C4A48, loc=loc, ip=ip) + NEG_E4M3s_REG1 = arith.constant(Int32.mlir_type, 0xCACCCED0, loc=loc, ip=ip) + NEG_E4M3s_REG2 = arith.constant(Int32.mlir_type, 0xB8C0C4C8, loc=loc, ip=ip) + + sign = arith.andi(src_i32, sign_mask, loc=loc, ip=ip) + sign = arith.shrui(sign, one, loc=loc, ip=ip) + + lut_idx = arith.andi(src_i32, lut_idx_mask, loc=loc, ip=ip) + + rsts = [] + for i in range(2): + final_prmt_idx = arith.ori(sign, final_prmt_base, loc=loc, ip=ip) + rst = llvm.inline_asm( + Int32.mlir_type, + [ + POS_E4M3s_REG1, + POS_E4M3s_REG2, + NEG_E4M3s_REG1, + NEG_E4M3s_REG2, + lut_idx, + final_prmt_idx, + ], + """{\n\t + .reg .b32 pos_f8s, neg_f8s;\n\t + prmt.b32 pos_f8s, $1, $2, $5;\n\t + prmt.b32 neg_f8s, $3, $4, $5;\n\t + prmt.b32 $0, pos_f8s, neg_f8s, $6;\n\t + }""", + "=r,n,n,n,n,r,r", + ) + rsts.append(rst) + + sign = arith.shrui(sign, sixteen, loc=loc, ip=ip) + lut_idx = arith.shrui(lut_idx, sixteen, loc=loc, ip=ip) + + vec_type = ir.VectorType.get([2], Int32.mlir_type, loc=loc) + return vector.from_elements(vec_type, rsts, loc=loc, ip=ip) + + +@dsl_user_op +def cvt_i4x8_to_mxf8x8( + src_vec8: ir.Value, + *, + dst_type: ir.Type, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> ir.Value: + """Convert eight int4 values to FP8 operand bytes for MXF8 MMA. + + The output values are generated as raw bytes by + :func:`cvt_i4x8_to_mxf8x8_impl`, then bitcast to ``vector<8 x dst_type>``. + Callers normally pass ``Int8.mlir_type`` and treat the result as raw bytes; + higher-level helpers may bitcast those bytes to an FP8 vector afterward. + This keeps the conversion instruction sequence small and avoids building + intermediate FP8 constants directly. + + :param src_vec8: Source vector with eight int4 elements. + :type src_vec8: vector<8xi4> + :param dst_type: MLIR element type for the output bytes, usually ``Int8.mlir_type``. + :type dst_type: ir.Type + :return: Vector with eight FP8 elements. + :rtype: vector<8 x dst_type> + """ + # Pack 8 int4 values into 1 int32 value and fill upper bits with 0. + src_i32 = llvm.bitcast(Int32.mlir_type, src_vec8, loc=loc, ip=ip) + rst_i32x2 = cvt_i4x8_to_mxf8x8_impl(src_i32, loc=loc, ip=ip) + vec_mxf8x8_type = ir.VectorType.get([8], dst_type, loc=loc) + return llvm.bitcast(vec_mxf8x8_type, rst_i32x2, loc=loc, ip=ip) + + # Sign extend 4 int4 unpacked in 8b containers @dsl_user_op def sext_unpacked_i4x4_to_i8x4( @@ -2627,24 +2875,22 @@ def atomic_fmax( loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None, ) -> Float32: - """ - Implementation of atomic fmax using integer bitcast. + """Atomically apply fmax through integer-bitcast atomics. - Works for +inf, -inf, and signbit-0 nans including canonical nan. - Atomically maxes `val` to the value at memory location `ptr` and returns the old value. + This wrapper handles ``+inf``, ``-inf``, and sign-bit-zero NaNs, + including canonical NaN. - :param ptr: Pointer to memory location. Supports: - - ir.Value (LLVM pointer) - - cute.ptr (_Pointer instance) - :param val: value to max + :param ptr: Pointer to the memory location. + :type ptr: Union[ir.Value, Pointer] + :param val: Value to combine with memory. :type val: Float32 - :param sign_bit: Indicates the sign bit of `val` if known beforehand, e.g. abs vals - :type sign_bit: Optional[bool] - :param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel") - :type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] - :param scope: Memory scope ("gpu", "cta", "cluster", "sys") - :type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] - :return: Old value at memory location + :param sign_bit: Known sign bit of ``val``, defaults to None. + :type sign_bit: Optional[bool], optional + :param sem: Memory semantic, defaults to None. + :type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]], optional + :param scope: Memory scope, defaults to None. + :type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]], optional + :return: Old value at ``ptr``. :rtype: Float32 """ intval = llvm.bitcast(T.i32(), val.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) @@ -2676,6 +2922,64 @@ def atomic_fmax( ) +@dsl_user_op +def atomic_fmin( + ptr: Union[ir.Value, Pointer], + val: Float32, + *, + sign_bit: Optional[bool] = None, + sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None, + scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Float32: + """Atomically apply fmin through integer-bitcast atomics. + + This wrapper handles ``+inf``, ``-inf``, and sign-bit-zero NaNs, + including canonical NaN. + + :param ptr: Pointer to the memory location. + :type ptr: Union[ir.Value, Pointer] + :param val: Value to combine with memory. + :type val: Float32 + :param sign_bit: Known sign bit of ``val``, defaults to None. + :type sign_bit: Optional[bool], optional + :param sem: Memory semantic, defaults to None. + :type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]], optional + :param scope: Memory scope, defaults to None. + :type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]], optional + :return: Old value at ``ptr``. + :rtype: Float32 + """ + intval = llvm.bitcast(T.i32(), val.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + then_body = lambda: atomic_max( + ptr, Uint32(intval), sem=sem, scope=scope, loc=loc, ip=ip + ) + else_body = lambda: atomic_min( + ptr, Int32(intval), sem=sem, scope=scope, loc=loc, ip=ip + ) + + if sign_bit is None: + old_intval = cutlass_dsl.if_generate( + Int32(intval) < 0, + then_body, + else_body, + [], + [Int32], + loc=loc, + ip=ip, + ) + elif sign_bit: + old_intval = then_body() + else: + old_intval = else_body() + + assert not isinstance(old_intval, list) + return Float32( + llvm.bitcast(T.f32(), old_intval.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + ) + + @deprecated("atomic_max_float32 is deprecated, use atomic_fmax instead") def atomic_max_float32( ptr: Union[ir.Value, Pointer], diff --git a/python/CuTeDSL/cutlass/cute/arch/smem.py b/python/CuTeDSL/cutlass/cute/arch/smem.py index c4a24af3d..95957e56e 100644 --- a/python/CuTeDSL/cutlass/cute/arch/smem.py +++ b/python/CuTeDSL/cutlass/cute/arch/smem.py @@ -9,6 +9,7 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. +from cutlass.address_space import AddressSpace from typing import Optional, Type from cutlass.cutlass_dsl import T, dsl_user_op @@ -20,8 +21,6 @@ from cutlass._mlir.dialects import nvvm, llvm from ..typing import Int, Int32, Pointer, Numeric, NumericMeta -_AddressSpace = _cute_ir.AddressSpace - @dsl_user_op def alloc_smem( @@ -53,9 +52,7 @@ def alloc_smem( if alignment is None: # Default alignment based on the element type's width alignment = element_type.width // 8 - ptr_ty = _cute_ir.PtrType.get( - element_type.mlir_type, _cute_ir.AddressSpace.smem, alignment - ) + ptr_ty = _cute_ir.PtrType.get(element_type.mlir_type, AddressSpace.smem, alignment) return _cute_nvgpu_ir.arch_alloc_smem( ptr=ptr_ty, input=ir.IntegerAttr.get(T.i32(), size_in_elems), @@ -93,7 +90,7 @@ def get_dyn_smem( alignment = element_type.width // 8 ptr_ty = _cute_ir.PtrType.get( element_type.mlir_type, - _cute_ir.AddressSpace.smem, + AddressSpace.smem, alignment, ) return _cute_nvgpu_ir.arch_get_dyn_smem(ptr=ptr_ty, loc=loc, ip=ip) @@ -124,17 +121,36 @@ def map_dsmem_ptr( """ Maps a shared memory pointer to a remote CTA's distributed shared memory. - :param smem_ptr: A pointer in SMEM + :param smem_ptr: A pointer in SMEM. :type smem_ptr: Pointer :param cta_rank_in_cluster: The CTA in cluster to map to :type cta_rank_in_cluster: Int - :return: The remote shared memory CuTe pointer + :return: A DSMEM CuTe pointer to the remote CTA's shared memory. :rtype: Pointer + :raises TypeError: If smem_ptr is not a CuTe pointer. + :raises ValueError: If smem_ptr is not in the SMEM address space. + """ + if not isinstance(smem_ptr, Pointer): + raise TypeError(f"smem_ptr must be a Pointer, but got {type(smem_ptr)}") + # Public API validation: map_dsmem_ptr accepts only local SMEM pointers. + if smem_ptr.memspace != AddressSpace.smem: + raise ValueError( + "map_dsmem_ptr expects smem_ptr to be an SMEM pointer, " + f"but got {smem_ptr.memspace}" + ) + + mlir_type = smem_ptr.dtype.mlir_type + memspace = AddressSpace.dsmem + alignment = smem_ptr.alignment + is_swizzled = smem_ptr.type.is_swizzled + swizzle = smem_ptr.type.swizzle_type.attribute if is_swizzled else None + ptr_type = _cute_ir.PtrType.get(mlir_type, memspace, alignment, swizzle) + dsmem_llvm_ptr = nvvm.mapa( - llvm.PointerType.get(_cute_ir.AddressSpace.dsmem), + llvm.PointerType.get(memspace), smem_ptr.to_llvm_ptr(loc=loc, ip=ip), Int32(cta_rank_in_cluster).ir_value(loc=loc, ip=ip), loc=loc, @@ -142,10 +158,10 @@ def map_dsmem_ptr( ) intptr = llvm.ptrtoint(T.i32(), dsmem_llvm_ptr, loc=loc, ip=ip) - aligned_ty = _cute_ir.ConstrainedIntType.get(smem_ptr.alignment, 32) + aligned_ty = _cute_ir.ConstrainedIntType.get(alignment, 32) aligned_intptr = _cute_ir.assume(aligned_ty, intptr, loc=loc, ip=ip) - return _cute_ir.inttoptr(smem_ptr.type, aligned_intptr, loc=loc, ip=ip) + return _cute_ir.inttoptr(ptr_type, aligned_intptr, loc=loc, ip=ip) @dsl_user_op @@ -170,8 +186,41 @@ def store_async_dsmem( :param mbar_ptr: Mbarrier pointer in this CTA's shared memory. Mapped to the peer CTA via ``nvvm.mapa``. :param peer_cta_rank: Target CTA rank in the cluster. + :raises TypeError: If `smem_ptr` or `mbar_ptr` is not a CuTe pointer. + :raises ValueError: If `value` is not a scalar or 2/4-tuple, or if + `smem_ptr`/`mbar_ptr` is not in the SMEM address + space, or if `smem_ptr` is not aligned to + `4 * len(value)` bytes. """ - dsmem_ptr_ty = llvm.PointerType.get(_AddressSpace.dsmem) + if not isinstance(smem_ptr, Pointer): + raise TypeError(f"smem_ptr must be a Pointer, but got {type(smem_ptr)}") + if not isinstance(mbar_ptr, Pointer): + raise TypeError(f"mbar_ptr must be a Pointer, but got {type(mbar_ptr)}") + if smem_ptr.memspace != AddressSpace.smem: + raise ValueError( + "store_async_dsmem expects smem_ptr to be an SMEM pointer, " + f"but got {smem_ptr.memspace}" + ) + if mbar_ptr.memspace != AddressSpace.smem: + raise ValueError( + "store_async_dsmem expects mbar_ptr to be an SMEM pointer, " + f"but got {mbar_ptr.memspace}" + ) + + values = value if isinstance(value, tuple) else (value,) + n = len(values) + if n not in [1, 2, 4]: + raise ValueError( + f"store_async_dsmem: value must be a scalar or a tuple of 2 or 4 i32s, got {n}" + ) + required_align = n * 4 + if smem_ptr.alignment % required_align != 0: + raise ValueError( + f"store_async_dsmem: v{n}.b32 requires {required_align}-byte aligned pointer, " + f"but smem_ptr has alignment {smem_ptr.alignment}" + ) + + dsmem_ptr_ty = llvm.PointerType.get(AddressSpace.dsmem) cta_rank_ir = Int32(peer_cta_rank).ir_value(loc=loc, ip=ip) dsmem_addr = nvvm.mapa( diff --git a/python/CuTeDSL/cutlass/cute/arch/tmem.py b/python/CuTeDSL/cutlass/cute/arch/tmem.py index 28e37e6fc..27cf0099c 100644 --- a/python/CuTeDSL/cutlass/cute/arch/tmem.py +++ b/python/CuTeDSL/cutlass/cute/arch/tmem.py @@ -9,6 +9,7 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. +from cutlass.address_space import AddressSpace from typing import Optional, Type from cutlass.cutlass_dsl import dsl_user_op @@ -102,9 +103,7 @@ def retrieve_tmem_ptr( f"element_type must be a type of Numeric, but got {element_type}" ) - res_ty = _cute_ir.PtrType.get( - element_type.mlir_type, _cute_ir.AddressSpace.tmem, alignment - ) + res_ty = _cute_ir.PtrType.get(element_type.mlir_type, AddressSpace.tmem, alignment) return _cute_nvgpu_ir.arch_sm100_retrieve_tmem_ptr( res_ty, ptr_to_buffer_holding_addr.value, diff --git a/python/CuTeDSL/cutlass/cute/atom.py b/python/CuTeDSL/cutlass/cute/atom.py index c5f3642b2..01676ef26 100644 --- a/python/CuTeDSL/cutlass/cute/atom.py +++ b/python/CuTeDSL/cutlass/cute/atom.py @@ -10,7 +10,7 @@ # is strictly prohibited. from abc import ABC, ABCMeta, abstractmethod -from typing import Type, Union, Optional, Any, overload, List, Tuple +from typing import ClassVar, Type, Union, Optional, Any, overload, List, Tuple from .typing import ( Shape, @@ -67,6 +67,9 @@ class MmaOp(Op, metaclass=ABCMeta): MMA Operation abstract base class. """ + # Concrete MmaOp subclasses set this to a human-readable label used by __str__. + descriptive_name: ClassVar[str] + @abstractmethod def _make_trait( self, @@ -996,7 +999,7 @@ def make_copy_atom( :param op: The Copy Operation to construct an Atom for :type op: CopyOp - :param copy_internal_type: An internal data type used to construct the source/destination layouts in unit of tensor elements + :param copy_internal_type: Element type used to construct the source/destination layouts in units of tensor elements :type copy_internal_type: Type[Numeric] :return: The Copy Atom :rtype: CopyAtom diff --git a/python/CuTeDSL/cutlass/cute/core.py b/python/CuTeDSL/cutlass/cute/core.py index c2648551c..a3a765148 100644 --- a/python/CuTeDSL/cutlass/cute/core.py +++ b/python/CuTeDSL/cutlass/cute/core.py @@ -10,10 +10,22 @@ # is strictly prohibited. +from collections.abc import Iterable from functools import partial, reduce import inspect from inspect import isclass -from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union, overload +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Tuple, + Type, + Union, + cast, + overload, +) from typing_extensions import deprecated @@ -41,10 +53,11 @@ from cutlass.cutlass_dsl import ( lru_cache_ir, not_, ) +from cutlass.base_dsl.typing import Int8 as _BaseInt8 +from cutlass.base_dsl.typing import Pointer as _BasePointer from .tuple import find_if, flatten_to_tuple, product_each, transform_leaf, unwrap, wrap from .typing import ( - AddressSpace, Boolean, ComposedLayout, Coord, @@ -66,10 +79,14 @@ from .typing import ( Tiler, XTuple, _element_precision_width, + _as_address_space, + _to_mlir_address_space, is_int_tuple, is_integer, + AddressSpace, ) + __all__ = [ # Classes "IntValue", @@ -2061,7 +2078,7 @@ def printf( This function provides printf-style formatted printing capabilities. It can print values directly or format them using C-style format strings. The function supports printing various types including - layouts, numeric values, tensors, and other CuTe objects. + layouts, numeric values, tensors, pointers, and other CuTe objects. The function accepts either: 1. A list of values to print directly @@ -2112,38 +2129,47 @@ def printf( fmt = "{}" + ", {}" * (len(args) - 1) + end def process_arg(arg: Any) -> Any: - arg0 = arg.value if isinstance(arg, Numeric) else arg + if isinstance(arg, Numeric): + return arg.ir_value(loc=loc, ip=ip) - if isinstance(arg0, ir.Value): - return arg.ir_value(loc=loc, ip=ip) if isinstance(arg, Numeric) else arg0 - elif isinstance(arg0, bool): - return const(arg0, Boolean) - elif isinstance(arg0, int): - return const(arg0, Int32) - elif isinstance(arg0, float): - return const(arg0, Float32) - elif has_underscore(arg0): + if isinstance(arg, _BasePointer): + ptr = make_ptr( + arg.dtype, + arg.to_llvm_ptr(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return ptr.value + elif isinstance(arg, ir.Value): + return arg + elif isinstance(arg, bool): + return const(arg, Boolean) + elif isinstance(arg, int): + return const(arg, Int32) + elif isinstance(arg, float): + return const(arg, Float32) + elif has_underscore(arg): # Assume it's a coordinate - return _pack_coord(arg0) # type: ignore[arg-type] - elif has_scaled_basis(arg0): + return _pack_coord(arg) # type: ignore[arg-type] + elif has_scaled_basis(arg): # Assume it's a stride - return _pack_stride(arg0) # type: ignore[arg-type] - elif is_int_tuple(arg0): - return _pack_int_tuple(arg0) # type: ignore[arg-type] - elif isinstance(arg0, tuple): + return _pack_stride(arg) # type: ignore[arg-type] + elif is_int_tuple(arg): + return _pack_int_tuple(arg) # type: ignore[arg-type] + elif isinstance(arg, tuple): # Assume it's a tile - return _pack_tile(arg0) - elif isinstance(arg0, _Tensor): - arg0._check_can_load_store() - if isinstance(arg0.layout, ComposedLayout) and isinstance( - arg0.layout.inner, Swizzle + return _pack_tile(arg) + elif isinstance(arg, _Tensor): + arg._check_can_load_store() + if isinstance(arg.layout, ComposedLayout) and isinstance( + arg.layout.inner, Swizzle ): raise NotImplementedError( "tensor with swizzled layout (PISL) is not supported in printf, please use swizzled pointer (PDSL) instead" ) - return arg0.value - elif isinstance(arg0, (_Pointer, _ComposedLayout)): - return arg0.value + return arg.value + elif isinstance(arg, (_Pointer, _ComposedLayout)): + return arg.value else: raise TypeError(f"unsupported argument type in printf, got {type(arg)}") @@ -3567,7 +3593,7 @@ def round_up(a: IntTuple, b: IntTuple) -> IntTuple: @dsl_user_op def make_layout( - shape: Shape, + shape: Union[Shape, Iterable[Layout]], *, stride: Union[Stride, None] = None, loc: Optional[ir.Location] = None, @@ -3578,9 +3604,14 @@ def make_layout( A Layout in CuTe represents the mapping between logical and physical coordinates of a tensor. This function creates a Layout object that defines how tensor elements are arranged in memory. - :param shape: Shape of the layout defining the size of each mode - :type shape: Shape - :param stride: Optional stride values for each mode, defaults to None + As an alternative to a shape, an iterable of :class:`Layout` objects may be + passed, in which case each layout becomes a separate mode of the result (the + ``stride`` argument is ignored). This mirrors CuTe's variadic + ``make_layout(layoutA, layoutB, ...)``. + + :param shape: Shape of the layout defining the size of each mode, or an iterable of Layout objects to concatenate (each becomes a mode) + :type shape: Union[Shape, Iterable[Layout]] + :param stride: Optional stride values for each mode, defaults to None (ignored when shape is an iterable of layouts) :type stride: Union[Stride, None] :param loc: Source location information, defaults to None :type loc: Optional[Location] @@ -3605,6 +3636,11 @@ def make_layout( # Create a layout with custom strides layout = make_layout((2,2,2), stride=(4,1,2)) # layout with strides (4,1,2) + # Concatenate layouts: each becomes a mode of the result + mode0 = make_layout(64, stride=1) + mode1 = make_layout(128, stride=64) + combined = make_layout([mode0, mode1]) # (64,128):(1,64) + Note: - If stride is not provided, a default compact left-most stride is computed based on the shape - The resulting layout maps logical coordinates to physical memory locations @@ -3617,11 +3653,31 @@ def make_layout( - Stride is keyword only argument to improve readability, e.g. * make_layout((3,4), (1,4)) can be confusing with make_layout(((3,4), (1,4))) * make_layout((3,4), stride=(1,4)) is more readable + - When passing an iterable of layouts, each layout becomes a separate mode """ + # Concatenation form: an iterable of Layouts, each becoming a mode. Strings + # and bytes are iterable too, so exclude them; ints are not iterable and so + # fall through to the normal shape path unchanged. + if isinstance(shape, Iterable) and not isinstance(shape, (str, bytes)): + # tuples/lists are reusable; materialize only one-shot iterables (generators). + seq = shape if isinstance(shape, (list, tuple)) else list(shape) + if seq and all(isinstance(elem, Layout) for elem in seq): + # all() above guarantees every element is a Layout; mypy can't infer + # that from a runtime predicate, so narrow explicitly. + layouts = cast("List[Layout]", seq) + return make_layout( + tuple(layout.shape for layout in layouts), + stride=tuple(layout.stride for layout in layouts), + loc=loc, + ip=ip, + ) + # Not layouts: keep the materialized form so generators still work below. + shape = cast("Shape", seq) + if stride is not None and not is_congruent(shape, stride): raise ValueError("shape and stride must be congruent") - shape_val = _pack_shape(shape, loc=loc, ip=ip) + shape_val = _pack_shape(cast("Shape", shape), loc=loc, ip=ip) if stride is not None: stride_val = _pack_stride(stride, loc=loc, ip=ip) layout_ty = _cute_ir.LayoutType.get(shape_val.type, stride_val.type) @@ -4316,7 +4372,9 @@ def recast_ptr( value_ty = cvt_ty or ptr.type.value_type swizzle_attr = swizzle_.type.attribute if swizzle_ is not None else None - res_ty = _cute_ir.PtrType.get(value_ty, ptr.memspace, ptr.alignment, swizzle_attr) + res_ty = _cute_ir.PtrType.get( + value_ty, _to_mlir_address_space(ptr.memspace), ptr.alignment, swizzle_attr + ) return _cute_ir.recast_iter(res_ty, ptr.value, loc=loc, ip=ip) @@ -4346,8 +4404,7 @@ def make_ptr( if mem_space is None: mem_space = AddressSpace.generic - if not isinstance(mem_space, AddressSpace): - raise TypeError(f"expects mem_space to be an AddressSpace, but got {mem_space}") + mem_space = _as_address_space(mem_space) # TMEM addresses are 32b wide is_tmem = mem_space == AddressSpace.tmem @@ -4372,7 +4429,9 @@ def make_ptr( data_ty = cvt_type swizzle = swizzle_.type.attribute if swizzle_ is not None else None - ptr_ty = _cute_ir.PtrType.get(data_ty, mem_space, assumed_align, swizzle) + ptr_ty = _cute_ir.PtrType.get( + data_ty, _to_mlir_address_space(mem_space), assumed_align, swizzle + ) ptr = _cute_ir.inttoptr(ptr_ty, aligned_intptr, loc=loc, ip=ip) ptr._dtype = dtype return ptr @@ -5804,14 +5863,6 @@ class struct: :return: The MLIR value of the underlying pointer. :rtype: ir.Value """ - import warnings - - with warnings.catch_warnings(): - warnings.simplefilter("always") - warnings.warn( - "Use explicit `struct.scalar.ptr` for pointer instead.", - DeprecationWarning, - ) return self._ptr.value # util func for base dsl scalar types diff --git a/python/CuTeDSL/cutlass/cute/experimental/algorithm.py b/python/CuTeDSL/cutlass/cute/experimental/algorithm.py index 019d20b24..a3db330f1 100644 --- a/python/CuTeDSL/cutlass/cute/experimental/algorithm.py +++ b/python/CuTeDSL/cutlass/cute/experimental/algorithm.py @@ -9,6 +9,7 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. +from cutlass.address_space import AddressSpace from typing import Optional from cutlass import cute @@ -80,14 +81,14 @@ def partition_and_copy( src_partitioned = src dst_partitioned = dst tid_x = tiled_copy.thr_idx - if src.memspace != cute.AddressSpace.rmem: + if src.memspace != AddressSpace.rmem: src_partitioned = partition( src, tid_x, layout_tv=tiled_copy.layout_src_tv_tiled, tiler=cute.core._pack_tile(tiled_copy.tiler_mn), ) - if dst.memspace != cute.AddressSpace.rmem: + if dst.memspace != AddressSpace.rmem: dst_partitioned = partition( dst, tid_x, @@ -115,16 +116,13 @@ def partition_and_copy( src.memspace, dst.memspace, ) in [ - (cute.AddressSpace.rmem, cute.AddressSpace.smem), - (cute.AddressSpace.smem, cute.AddressSpace.rmem), - (cute.AddressSpace.rmem, cute.AddressSpace.gmem), - (cute.AddressSpace.gmem, cute.AddressSpace.rmem), + (AddressSpace.rmem, AddressSpace.smem), + (AddressSpace.smem, AddressSpace.rmem), + (AddressSpace.rmem, AddressSpace.gmem), + (AddressSpace.gmem, AddressSpace.rmem), ]: simt_auto_vec_copy(src_partitioned, dst_partitioned, loc=loc, ip=ip) - elif ( - src.memspace == cute.AddressSpace.gmem - and dst.memspace == cute.AddressSpace.smem - ): + elif src.memspace == AddressSpace.gmem and dst.memspace == AddressSpace.smem: simt_auto_vec_copy( src_partitioned, dst_partitioned, async_op=True, loc=loc, ip=ip ) diff --git a/python/CuTeDSL/cutlass/cute/experimental/core.py b/python/CuTeDSL/cutlass/cute/experimental/core.py index 43b0a41b8..e564837bc 100644 --- a/python/CuTeDSL/cutlass/cute/experimental/core.py +++ b/python/CuTeDSL/cutlass/cute/experimental/core.py @@ -17,8 +17,7 @@ from cutlass._mlir import ir from cutlass.cutlass_dsl import lru_cache_ir from cutlass._mlir.dialects.core import OperationTypeEnum from cutlass import cute -from cutlass.cute.core import _pack_coord -from cutlass.cute.typing import Boolean, Coord, Tensor +from cutlass.cute.typing import Boolean class _SupportsIrValue(Protocol): @@ -30,7 +29,7 @@ class _SupportsIrValue(Protocol): ) -> ir.Value: ... -SkipWaitToken: TypeAlias = bool | ir.Value | _SupportsIrValue +SkipWaitToken: TypeAlias = Union[bool, ir.Value, _SupportsIrValue] # A producer/consumer role may be a single operation type or a non-empty # sequence of them, for pipelines where a single stage is written (or read) @@ -71,30 +70,6 @@ def get_mbarrier( return cutlass_lir_ir.GetMbarrierOp(stage_token, loc=loc, ip=ip) -@dsl_user_op -def domain_offset( - coord: Coord, - tensor: Tensor, - loc: Optional[ir.Location] = None, - ip: Optional[ir.InsertionPoint] = None, -) -> Tensor: - """Shift a tensor's logical domain while preserving its descriptor fields. - - This is the IR-preserving counterpart to ``cute.domain_offset``. It keeps - the source-level domain-offset semantics explicit until the IR partitioning - passes can apply the offset to generated coordinate tensors. - """ - coord_value = _pack_coord(coord, loc=loc, ip=ip) - op = cutlass_lir_ir.DomainOffsetOp( - tensor.value, - coord_value, - results=[tensor.value.type], - loc=loc, - ip=ip, - ) - return op.result - - @ir.register_value_caster(cutlass_lir_ir.PipelineStateType.get_static_typeid()) class PipelineState(ir.Value): def __init__(self, value: ir.Value) -> None: diff --git a/python/CuTeDSL/cutlass/cute/experimental/math.py b/python/CuTeDSL/cutlass/cute/experimental/math.py index 1198e1f31..f8ec9708b 100644 --- a/python/CuTeDSL/cutlass/cute/experimental/math.py +++ b/python/CuTeDSL/cutlass/cute/experimental/math.py @@ -82,3 +82,38 @@ def dot( ) +@dsl_user_op +def dot_sparse( + mma_atom: cute.MmaAtom, + a: cute.Tensor, + e: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> None: + """Emit a sparse MMA operation. + + :param mma_atom: Regular sparse MMA atom. The op derives the metadata + operand expected by the atom from ``e`` during lowering. + :param a: Sparse A operand partitioned as rank-3 ``(V, Rest, K)``. + :param e: Sparse metadata operand partitioned as rank-3 ``(V, Rest, K)``. + :param b: Dense B operand partitioned as rank-3 ``(V, Rest, K)``. + :param c: Accumulator/result operand. + :param loc: Optional MLIR location. + :param ip: Optional MLIR insertion point. + """ + a = _ensure_rank3(a, loc, ip) + e = _ensure_rank3(e, loc, ip) + b = _ensure_rank3(b, loc, ip) + cutlass_lir.DotSparseOp( + mma_atom._unpack(), + a.value, + e.value, + b.value, + c.value, + loc=loc, + ip=ip, + ) + + diff --git a/python/CuTeDSL/cutlass/cute/experimental/memory.py b/python/CuTeDSL/cutlass/cute/experimental/memory.py index 9723bf42e..a2bd36100 100644 --- a/python/CuTeDSL/cutlass/cute/experimental/memory.py +++ b/python/CuTeDSL/cutlass/cute/experimental/memory.py @@ -9,7 +9,8 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. -from typing import Type, Optional +from cutlass.address_space import AddressSpace +from typing import Type, Optional, Union from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir from cutlass._mlir import ir from cutlass.cutlass_dsl import dsl_user_op @@ -111,8 +112,8 @@ def _emit_tma_load( @dsl_user_op def allocate( type: Type[cute.Numeric], - address_space: cute.AddressSpace, - layout: cute.Layout | cute.ComposedLayout, + address_space: AddressSpace, + layout: Union[cute.Layout, cute.ComposedLayout], alignment: cute.Int32, is2cta: bool = False, loc: Optional[ir.Location] = None, @@ -188,7 +189,7 @@ def tma_load( Required for multicast loads and for selecting SM100_TMA_LOAD_2SM. vmnk_layout: layout describing the cluster configuration. Required when tma_operation_type is multicast. multicast_mode (optional): multicast projection mode. Required when tma_operation_type is multicast. - internal_type (optional): selects the TMA transfer's internal element encoding used by hardware. + internal_type (optional): selects the TMA transfer element-format override. Does not change src/dst memref types. For structured sparsity, use base storage types: Float16 for 2:4 FP16 sparse element type, Uint8 for 8:1 uint8 sparse element type. """ @@ -274,7 +275,7 @@ def tma_store( """ Copies a tensor from a Buffer to a tensor pointed to by a !cute.memref. - internal_type (optional): selects the TMA transfer's internal element encoding used by hardware. + internal_type (optional): selects the TMA transfer element-format override. Does not change src/dst memref types. For structured sparsity, use base storage types: Float16 for 2:4 FP16 sparse element type, Uint8 for 8:1 uint8 sparse element type. """ @@ -320,13 +321,12 @@ def tma_reduce_store( kinds: ADD / MIN / MAX / INC / DEC / AND / OR / XOR. :param cta_v_map: optional CTA-v map layout. If omitted, the layout is inferred from the destination GMEM memref. - :param internal_type: optional override for the TMA transfer's internal - element encoding; see :func:`tma_store` for semantics. + :param internal_type: optional override for the TMA transfer element + format; see :func:`tma_store` for semantics. """ if not isinstance(kind, cute.ReductionKind): raise TypeError( - f"tma_reduce_store: unsupported kind {kind!r}; " - f"expected cute.ReductionKind" + f"tma_reduce_store: unsupported kind {kind!r}; expected cute.ReductionKind" ) kwargs = { "kind": ir.Attribute.parse(f"#cute_nvgpu.tma_reduce_kind<{kind.name}>"), diff --git a/python/CuTeDSL/cutlass/cute/experimental/pipeline.py b/python/CuTeDSL/cutlass/cute/experimental/pipeline.py index 93280ea23..3f32452d1 100644 --- a/python/CuTeDSL/cutlass/cute/experimental/pipeline.py +++ b/python/CuTeDSL/cutlass/cute/experimental/pipeline.py @@ -132,7 +132,7 @@ class GenericPipelineBase: """Acquire a producer stage and return its stage token/index. When `token` is provided, reuse the preceding `producer_try_acquire()` - result and keep the internal state at the acquired stage so a following + result and keep the stored state at the acquired stage so a following `producer_commit_and_advance()` retires the same stage. """ if token is None: @@ -186,9 +186,9 @@ class GenericPipelineBase: return self def producer_commit_and_advance(self) -> "GenericPipelineBase": - """Combined producer commit + advance with automatic elect_one using internal state.""" + """Combined producer commit + advance with automatic elect_one using stored state.""" self.producer_commit() - # Update internal state in-place for better performance + # Update state in-place for better performance self.producer_state = pipeline_advance_iterator( self.raw_pipeline, self.producer_state ) @@ -200,7 +200,7 @@ class GenericPipelineBase: """Wait for a consumer stage and return its stage token/index. When `token` is provided, reuse the preceding `consumer_try_wait()` - result and keep the internal state at the consumed stage so a following + result and keep the stored state at the consumed stage so a following `consumer_release_and_advance()` retires the same stage. """ if token is None: @@ -230,9 +230,9 @@ class GenericPipelineBase: return self def consumer_release_and_advance(self) -> "GenericPipelineBase": - """Combined consumer release + advance with automatic elect_one using internal state.""" + """Combined consumer release + advance with automatic elect_one using stored state.""" self.consumer_release() - # Update internal state in-place for better performance + # Update state in-place for better performance self.consumer_state = pipeline_advance_iterator( self.raw_pipeline, self.consumer_state ) @@ -466,7 +466,7 @@ class TMAToUMMAPipeline(GenericPipelineBase): mma_operation_type: OperationTypeEnum, cluster_layout_vmnk: cute.Layout, ) -> "TMAToUMMAPipeline": - """Internal: compute TMA multicast masks from cluster layout.""" + """Compute TMA multicast masks from cluster layout.""" tma_mcast_proj_A = 2 # multicast across CTAs in same row tma_mcast_proj_B = 1 # multicast across CTAs in same column @@ -1051,7 +1051,7 @@ class TMAStorePipeline: return tail_impl() def _barrier(self) -> None: - """Internal barrier synchronization.""" + """Synchronize the pipeline barrier.""" cute.arch.barrier( barrier_id=self.barrier_id, number_of_threads=self.arv_count, diff --git a/python/CuTeDSL/cutlass/cute/experimental/utils.py b/python/CuTeDSL/cutlass/cute/experimental/utils.py index 3cada8b98..11780c2d1 100644 --- a/python/CuTeDSL/cutlass/cute/experimental/utils.py +++ b/python/CuTeDSL/cutlass/cute/experimental/utils.py @@ -9,7 +9,8 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. -from typing import Callable, Optional, Tuple, Union +from cutlass.address_space import AddressSpace +from typing import TYPE_CHECKING, Callable, Optional, Tuple, Union import cutlass from cutlass import cute @@ -22,6 +23,9 @@ from ... import cutlass_dsl as _dsl from ..arch.constants import THREADS_PER_WARPGROUP from .pipeline import TMAStorePipeline, TMAToUMMAPipeline +if TYPE_CHECKING: + from cutlass.utils.layout import LayoutEnum + def get_cta_v_map_ab( gmem_tensor: cute.Tensor, @@ -270,13 +274,13 @@ def epilogue_tma_store( ) rmem_acc_buffer = allocate( acc_dtype, - cute.AddressSpace.rmem, + AddressSpace.rmem, acc_d_rmem_layout, alignment=32, ) rmem_d_buffer = allocate( d_dtype, - cute.AddressSpace.rmem, + AddressSpace.rmem, acc_d_rmem_layout, alignment=32, ) @@ -289,7 +293,7 @@ def epilogue_tma_store( ) smem_d_buffer = allocate( d_dtype, - cute.AddressSpace.smem, + AddressSpace.smem, d_smem_layout_staged, alignment=1024, ) diff --git a/python/CuTeDSL/cutlass/cute/export/aot_config.py b/python/CuTeDSL/cutlass/cute/export/aot_config.py index bd91bee6e..59dde4310 100644 --- a/python/CuTeDSL/cutlass/cute/export/aot_config.py +++ b/python/CuTeDSL/cutlass/cute/export/aot_config.py @@ -96,7 +96,7 @@ def get_libdir(target: str = "") -> str: library subtree exists, so callers (and the CLI) surface a non-zero exit instead of silently returning ``""``. """ - from ..runtime import find_runtime_libraries + from cutlass.runtime import find_runtime_libraries libs = find_runtime_libraries(enable_tvm_ffi=False) default_dir = str(Path(libs[0]).parent) if libs else "" @@ -121,7 +121,7 @@ def get_libs(enable_tvm_ffi: bool = False, target: str = "") -> str: :return: Space-separated -l flags (e.g., "-lcuda_dialect_runtime -ltvm_ffi") :rtype: str """ - from ..runtime import find_runtime_libraries + from cutlass.runtime import find_runtime_libraries libs = find_runtime_libraries(enable_tvm_ffi=enable_tvm_ffi) # Convert full paths to -l flags @@ -144,7 +144,7 @@ def get_lib_paths(enable_tvm_ffi: bool = False) -> list[str]: :return: List of full library paths :rtype: list[str] """ - from ..runtime import find_runtime_libraries + from cutlass.runtime import find_runtime_libraries return find_runtime_libraries(enable_tvm_ffi=enable_tvm_ffi) diff --git a/python/CuTeDSL/cutlass/cute/export/c_header_generator.py b/python/CuTeDSL/cutlass/cute/export/c_header_generator.py index 5b5de7ced..d004e43bf 100644 --- a/python/CuTeDSL/cutlass/cute/export/c_header_generator.py +++ b/python/CuTeDSL/cutlass/cute/export/c_header_generator.py @@ -169,26 +169,59 @@ static inline void {symbol_prefix}_Kernel_Module_Unload({symbol_prefix}_Kernel_M arguments.append(f"void *{arg_name}") packed_args.append("&" + arg_name) elif isinstance(arg, Tensor): - dynamic_shapes = ( - f"\n int32_t dynamic_shapes[{sum(arg.dynamic_shapes_mask)}];" # type: ignore[attr-defined] - if sum(arg.dynamic_shapes_mask) > 0 # type: ignore[attr-defined] - else "" - ) stride_type = "int32_t" if arg._use_32bit_stride else "int64_t" # type: ignore[attr-defined] - dynamic_strides = ( - f"\n {stride_type} dynamic_strides[{sum(arg.dynamic_strides_mask)}];" # type: ignore[attr-defined] - if sum(arg.dynamic_strides_mask) > 0 # type: ignore[attr-defined] - else "" - ) declarations.append( - f""" -typedef struct {{ - void *data;{dynamic_shapes}{dynamic_strides} -}} {symbol_prefix}_Tensor_{arg_name}_t; -""" + self._generate_tensor_element_typedef( + symbol_prefix, + arg_name, + sum(arg.dynamic_shapes_mask), + sum(arg.dynamic_strides_mask), + stride_type, + shape_type="int32_t", + ) ) arguments.append(f"{symbol_prefix}_Tensor_{arg_name}_t *{arg_name}") packed_args.append(arg_name) + elif ( + isinstance(arg, (list, tuple)) + and len(arg) > 0 + and all(isinstance(e, Tensor) for e in arg) + ): + first = arg[0] + for e in arg[1:]: + if ( + e.dynamic_shapes_mask != first.dynamic_shapes_mask + or e.dynamic_strides_mask != first.dynamic_strides_mask + or e._use_32bit_stride != first._use_32bit_stride # type: ignore[attr-defined] + ): + raise DSLRuntimeError( + f"Inconsistent tensor descriptors in list/tuple argument {arg_name}: " + f"every element must share the same dynamic_shapes_mask, " + f"dynamic_strides_mask, and _use_32bit_stride" + ) + stride_type = "int32_t" if first._use_32bit_stride else "int64_t" # type: ignore[attr-defined] + declarations.append( + self._generate_tensor_element_typedef( + symbol_prefix, + arg_name, + sum(first.dynamic_shapes_mask), + sum(first.dynamic_strides_mask), + stride_type, + shape_type="int32_t", + ) + ) + n = len(arg) + declarations.append( + f""" +typedef struct {{ + {symbol_prefix}_Tensor_{arg_name}_t elements[{n}]; + int32_t count; +}} {symbol_prefix}_TensorList_{arg_name}_t; +""" + ) + arguments.append(f"{symbol_prefix}_TensorList_{arg_name}_t *{arg_name}") + for i in range(n): + packed_args.append(f"&{arg_name}->elements[{i}]") # Generate basic numeric types elif isinstance(arg_type, NumericMeta): arguments.append(self._generate_numeric_argument(arg_name, arg_type)) # type: ignore[arg-type] diff --git a/python/CuTeDSL/cutlass/cute/math.py b/python/CuTeDSL/cutlass/cute/math.py index 608ad0c40..f65aa86ab 100644 --- a/python/CuTeDSL/cutlass/cute/math.py +++ b/python/CuTeDSL/cutlass/cute/math.py @@ -9,102 +9,13 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. -"""CuTeDSL math API — thin re-export from :mod:`cutlass._mlir_helpers.math`. -All math implementations live in the foundation module. This file -curates the CuTeDSL-side public surface (historical transcendentals, -rounding, abs / copysign) and aliases ``abs`` to ``absf`` for backwards -compatibility. +# Re-export non-callable public symbols straight through. +from cutlass._mlir_helpers.math import * # noqa: F401 -``TensorSSA`` inherits from ``Vector``, so the foundation's Vector -dispatch path handles CuTeDSL tensor values directly — no -CuTeDSL-specific wrapper is needed here. Per-element math ops preserve -a ``TensorSSA``'s CuTe nested shape via the ``_wrap_like`` polymorphic -hook on the Vector base class. -""" +from cutlass._mlir_helpers import math as _base_math -from cutlass._mlir_helpers.math import ( - # Trigonometric - sin, - cos, - tan, - acos, - asin, - atan, - atan2, - # Hyperbolic - sinh, - cosh, - tanh, - acosh, - asinh, - atanh, - # Exponential / logarithmic - exp, - exp2, - expm1, - log, - log2, - log10, - log1p, - # Error functions - erf, - erfc, - # Power / root - pow, - cbrt, - sqrt, - rsqrt, - # Rounding - ceil, - floor, - round, - roundeven, - trunc, - # Sign - copysign, -) -from cutlass._mlir_helpers.math import abs as absf # historical CuTeDSL name +absf = _base_math.abs -__all__ = [ - # Trigonometric - "sin", - "cos", - "tan", - "acos", - "asin", - "atan", - "atan2", - # Hyperbolic - "sinh", - "cosh", - "tanh", - "acosh", - "asinh", - "atanh", - # Exponential / logarithmic - "exp", - "exp2", - "expm1", - "log", - "log2", - "log10", - "log1p", - # Error functions - "erf", - "erfc", - # Power / root - "pow", - "cbrt", - "sqrt", - "rsqrt", - # Rounding - "ceil", - "floor", - "round", - "roundeven", - "trunc", - # Abs / sign - "absf", - "copysign", -] +# Same public surface as the foundation, plus the ``absf`` alias. +__all__ = [*_base_math.__all__, "absf"] diff --git a/python/CuTeDSL/cutlass/cute/metadata.py b/python/CuTeDSL/cutlass/cute/metadata.py new file mode 100644 index 000000000..f163b1789 --- /dev/null +++ b/python/CuTeDSL/cutlass/cute/metadata.py @@ -0,0 +1,814 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +"""Build typed ``cutlass.compiler.metadata`` values directly from a CuTe DSL +kernel signature. + +This walks the live ``(signature, args, kwonlyargs)`` of a ``@cute.jit`` kernel +straight into the typed C++ metadata hierarchy carried by +``CompilationArtifact``, reading the DSL-native argument types directly, so +downstream consumers (codegen, AOT wrapper generation, diagnostics) can +introspect the calling convention. + +Its only MLIR dependency is the ``cute`` dialect's ``AddressSpace`` enum, used to +classify a tensor/pointer's memory space; it builds no IR and needs no active +MLIR context. +""" + +import inspect +import struct +from dataclasses import fields as dataclass_fields, is_dataclass +from types import UnionType +from typing import ( + Any, + Dict, + Iterator, + List, + Optional, + SupportsInt, + Tuple, + Type, + Union, + cast, + get_args, + get_origin, + get_type_hints, +) + +import cuda.bindings.driver as cuda + +import cutlass.compiler as md +from cutlass._mlir.dialects import cute as _cute_ir +from cutlass.base_dsl.common import DSLRuntimeError +from cutlass.base_dsl.runtime.jit_arg_adapters import is_arg_annotation_constexpr +from cutlass.base_dsl.typing import Constexpr +from cutlass.compiler import DType, DTypeCode +from cutlass.cutlass_dsl import is_cute_algebra_type +from .runtime import _FakeStream +from .typing import ( + Numeric, + Pointer, + SymInt, + Tensor, + Boolean, + Int4, + Int8, + Uint8, + Int16, + Uint16, + Int32, + Uint32, + Int64, + Uint64, + Float16, + BFloat16, + Float32, + TFloat32, + Float64, + Float6E2M3FN, + Float6E3M2FN, + Float8E5M2, + Float8E4M3FN, + Float8E8M0FNU, + Float4E2M1FN, +) + + +# --------------------------------------------------------------------------- +# DType mapping +# --------------------------------------------------------------------------- + +# Maps each supported CuTe ``Numeric`` subclass to a factory producing its +# metadata ``DType``. Mapping the subclass directly is exact: ``DType.from_string`` +# only accepts canonical lowercase names like ``"int32"``, which ``str(Int32)`` +# (``"Int32"``) is not. +_NUMERIC_TO_DTYPE_FACTORY = { + Boolean: DType.bool_, + # No named factory for 4-bit ints; construct from (code, bits) directly. + Int4: lambda: DType(DTypeCode.Int, 4), + Int8: DType.int8, + Uint8: DType.uint8, + Int16: DType.int16, + Uint16: DType.uint16, + Int32: DType.int32, + Uint32: DType.uint32, + Int64: DType.int64, + Uint64: DType.uint64, + Float16: DType.float16, + BFloat16: DType.bfloat16, + Float32: DType.float32, + # TF32 has no distinct metadata DTypeCode; represent it as float32. This + # drops the "this is TF32" distinction, which metadata cannot express. + TFloat32: DType.float32, + Float64: DType.float64, + Float8E5M2: DType.float8_e5m2, + Float8E4M3FN: DType.float8_e4m3fn, + Float8E8M0FNU: DType.float8_e8m0fnu, + Float6E2M3FN: DType.float6_e2m3fn, + Float6E3M2FN: DType.float6_e3m2fn, + Float4E2M1FN: DType.float4_e2m1fn, +} + +# Numeric subclasses accepted as a plain scalar parameter: standard integer and +# float widths only -- sub-byte and narrow-float types appear only as tensor +# element types, never as scalar args. +_ACCEPTABLE_SCALAR_NUMERICS = frozenset( + { + Boolean, + Int8, + Uint8, + Int16, + Uint16, + Int32, + Uint32, + Int64, + Uint64, + Float16, + BFloat16, + Float32, + TFloat32, + Float64, + } +) + + +def _numeric_to_dtype(dtype: Type[Numeric]) -> DType: + """Map a CuTe ``Numeric`` subclass to its metadata ``DType``. + + Raises ``DSLRuntimeError`` for unsupported dtypes (packed FP6x4) and for any + subclass missing from the factory table. A drift guard asserts the mapped + ``DType.bits`` equals the subclass ``width``, catching silent divergence if + either side's dtype set changes (skipped for packed dtypes, whose storage + width differs from the unpacked element bits by design). + """ + try: + factory = _NUMERIC_TO_DTYPE_FACTORY[dtype] + except KeyError as exc: + raise DSLRuntimeError( + f"Metadata does not support dtype {dtype.__name__}." + ) from exc + + result = factory() + return result + + +# --------------------------------------------------------------------------- +# Memory space mapping (the lone MLIR enum dependency) +# --------------------------------------------------------------------------- + +# DLPack ``DLDeviceType`` values (dlpack.h). A tensor's ``device_kind`` is a +# compile-time type-level property: GPU memory => CUDA, otherwise host => CPU. +_KDL_CPU = 1 +_KDL_CUDA = 2 + + +def _memspace_to_device_kind(memspace: "_cute_ir.AddressSpace") -> int: + """Classify a tensor's device kind from its memory space: anything other than + the generic space is GPU memory.""" + return _KDL_CPU if memspace == _cute_ir.AddressSpace.generic else _KDL_CUDA + + +def _memspace_to_address_space(memspace: "_cute_ir.AddressSpace") -> "md.AddressSpace": + """Map a CuTe memory space to a metadata pointer address space: gmem -> global, + everything else -> generic.""" + if memspace == _cute_ir.AddressSpace.gmem: + return md.AddressSpace.Global + return md.AddressSpace.Generic + + +# DLPack dtype names that DType.from_string cannot parse directly: a packed FP4 +# name denotes the same unpacked element encoding. +_DLPACK_DTYPE_ALIASES: dict[str, str] = { +} + + +def _dlpack_dtype_to_dtype(dlpack_dtype: Any) -> DType: + """Map a real tensor's backing DLPack dtype to its metadata ``DType``. The + DLPack name is canonical lowercase, so ``DType.from_string`` parses it + directly once a packed-FP4 name is resolved to its unpacked element.""" + name = str(dlpack_dtype) + return DType.from_string(_DLPACK_DTYPE_ALIASES.get(name, name)) + + +def _same_dtype(a: DType, b: DType) -> bool: + return a.code == b.code and a.bits == b.bits and a.lanes == b.lanes + + +# DLPack DLDeviceType for the runtime device a real tensor lives on. +_DEVICE_TYPE_TO_KIND = {"cpu": _KDL_CPU, "cuda": _KDL_CUDA} + + +def _device_type_to_device_kind(device_type: Any) -> int: + return _DEVICE_TYPE_TO_KIND.get(str(device_type), _KDL_CUDA) + + +def _normalize_divisibility(divisibility: Optional[int]) -> Optional[int]: + """Collapse a trivial divisibility to ``None``. A ``SymInt`` defaults its + divisibility to 1 ("divisible by 1" is no constraint); metadata expresses + "no constraint" as ``None``, so emitting 1 would assert a noise ``% 1 == 0``.""" + return None if divisibility in (None, 1) else divisibility + + +# --------------------------------------------------------------------------- +# Dim symbol table +# --------------------------------------------------------------------------- + + +class _SymIntId: + """Identity key for a ``SymInt``: the same instance shared across two tensor + dimensions means those extents are the same runtime value. ``SymInt.__eq__`` + honours ``(width, divisibility, symbol)``, so value-equality would wrongly + merge distinct extents -- dedup must key on object identity.""" + + __slots__ = ("sym_int",) + + def __init__(self, sym_int: SymInt) -> None: + self.sym_int = sym_int + + def __hash__(self) -> int: + return id(self.sym_int) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, _SymIntId): + return NotImplemented + return self.sym_int is other.sym_int + + +class _DimSymbolTable: + """Accumulates the per-function ``DimSymbol`` table while the walk builds dims. + + A shared ``SymInt`` instance resolves to one ``DimSymbol`` (so two ``Dim``s + referencing the same ``SymbolId`` recover the "same extent" relationship); + an anonymous dynamic dim gets a fresh symbol. ``SymbolId`` equals the + symbol's index in ``symbols`` (first-encounter order). + """ + + def __init__(self) -> None: + self.symbols: List["md.DimSymbol"] = [] + self._symint_to_id: Dict[_SymIntId, "md.SymbolId"] = {} + self._num_shape = 0 + self._num_stride = 0 + + def _next_name(self, is_stride: bool) -> str: + if is_stride: + name = f"s{self._num_stride}" + self._num_stride += 1 + else: + name = f"n{self._num_shape}" + self._num_shape += 1 + return name + + def _append( + self, name: str, bits: int, divisibility: Optional[int] + ) -> "md.SymbolId": + sid = md.SymbolId(len(self.symbols)) + self.symbols.append( + md.DimSymbol(name=name, bits=bits, divisibility=divisibility) + ) + return sid + + def dim_for_symint(self, sym: SymInt, *, is_stride: bool) -> "md.Dim": + key = _SymIntId(sym) + existing = self._symint_to_id.get(key) + if existing is not None: + return md.Dim(existing) + name = sym.symbol or self._next_name(is_stride) + sid = self._append(name, sym.width, _normalize_divisibility(sym.divisibility)) + self._symint_to_id[key] = sid + return md.Dim(sid) + + def dim_for_anonymous(self, *, bits: int, is_stride: bool) -> "md.Dim": + return md.Dim(self._append(self._next_name(is_stride), bits, None)) + + +# --------------------------------------------------------------------------- +# Per-argument conversion +# --------------------------------------------------------------------------- + + +def _const_binding(name: str, arg: Any) -> "md.Binding": + """Build a ``Const``/``Unit`` for a compile-time constexpr value, pinning its + exact bit pattern for the wrapper to assert against.""" + if arg is None: + return md.Unit(name=name) + # bool before int (bool is a subclass of int). + if isinstance(arg, bool): + return md.Const(name=name, dtype=DType.bool_(), value=1 if arg else 0) + if isinstance(arg, int): + if not -(2**63) <= arg < 2**64: + raise DSLRuntimeError( + f"Constexpr int {arg} for {name!r} does not fit in 64 bits." + ) + # Const carries a raw u64 bit pattern that a boundary equality check + # reads back; wrap a negative value to its two's-complement u64. A + # constexpr int is always int64 -- its value, not a narrow width, is + # what gets asserted. + return md.Const( + name=name, dtype=DType.int64(), value=int(arg) & 0xFFFFFFFFFFFFFFFF + ) + if isinstance(arg, float): + bits = struct.unpack(" bool: + """Whether ``_const_binding`` can pin this constexpr value as a ``Const`` / + ``Unit`` -- only a scalar ``None`` / ``bool`` / ``int`` / ``float`` (``bool`` + is an ``int`` subclass). A non-scalar compile-time constexpr (a tuple, a + type, a CuTe layout, a Numeric, ...) has no metadata representation, so it is + pure compile-time specialization and is dropped from the calling convention + (as it always was) rather than forced into an unrepresentable ``Const``.""" + return value is None or isinstance(value, (int, float)) + + +def _tensor_binding(arg: "Tensor", name: str, table: _DimSymbolTable) -> "md.Tensor": + """Convert a live CuTe tensor value to a metadata ``Tensor``. + + The ABI ``dtype`` is the type each cell is passed as at the boundary. For a + real tensor that is the backing DLPack storage dtype; CuTe may set + ``element_type`` to a different logical type (e.g. fp4 cells reinterpreted + from ``uint8`` storage), recorded as ``element_dtype`` when it differs. A + fake (descriptor-only) tensor has no separate storage, so its declared + ``element_type`` is both the storage and logical type. + """ + shape = arg.shape + stride = arg.stride + assert isinstance(shape, tuple) and isinstance(stride, tuple) + + shape_dims: List["md.Dim"] = [] + for i, dynamic in enumerate(arg.dynamic_shapes_mask): + dim = shape[i] + if not dynamic: + shape_dims.append(md.Dim(int(cast(SupportsInt, dim)))) + elif isinstance(dim, SymInt): + shape_dims.append(table.dim_for_symint(dim, is_stride=False)) + else: + shape_dims.append(table.dim_for_anonymous(bits=32, is_stride=False)) + + use_32bit_stride = getattr(arg, "_use_32bit_stride", False) + stride_bits = 32 if use_32bit_stride else 64 + stride_dims: List["md.Dim"] = [] + for i, dynamic in enumerate(arg.dynamic_strides_mask): + s = stride[i] + if not dynamic: + stride_dims.append(md.Dim(int(cast(SupportsInt, s)))) + elif isinstance(s, SymInt): + stride_dims.append(table.dim_for_symint(s, is_stride=True)) + else: + stride_dims.append( + table.dim_for_anonymous(bits=stride_bits, is_stride=True) + ) + + if hasattr(arg, "_tvm_ffi_tensor"): + storage_dtype = _dlpack_dtype_to_dtype(arg._tvm_ffi_tensor.dtype) + logical_dtype = _numeric_to_dtype(arg.element_type) + element_dtype = ( + None if _same_dtype(logical_dtype, storage_dtype) else logical_dtype + ) + device_kind = _device_type_to_device_kind(arg._tvm_ffi_tensor.device.type) + else: + storage_dtype = _numeric_to_dtype(arg.element_type) + element_dtype = None + device_kind = _memspace_to_device_kind(arg.memspace) + + return md.Tensor( + name=name, + shape=shape_dims, + dtype=storage_dtype, + element_dtype=element_dtype, + strides=stride_dims, + data_address_space=_memspace_to_address_space(arg.memspace), + device_kind=device_kind, + data_alignment=getattr(arg, "_assumed_align", None), + ) + + +def _cute_shape_binding(arg: Any, name: str, table: _DimSymbolTable) -> "md.Shape": + """Convert a flat CuTe shape/stride/coord argument to a metadata ``Shape``. + + A metadata ``Shape`` is a flat list of ``Dim``s, so a hierarchical + (nested-tuple) shape has no representation here and is rejected. + """ + dims: List["md.Dim"] = [] + for elem in arg: + if isinstance(elem, bool): + raise DSLRuntimeError(f"Unexpected bool in cute shape arg {name!r}.") + if isinstance(elem, int): + dims.append(md.Dim(int(elem))) + elif isinstance(elem, SymInt): + dims.append(table.dim_for_symint(elem, is_stride=False)) + elif isinstance(elem, Numeric): + dims.append(table.dim_for_anonymous(bits=elem.dtype.width, is_stride=False)) + elif isinstance(elem, tuple): + raise DSLRuntimeError( + f"Nested cute shape arguments are not supported in metadata " + f"for {name!r}." + ) + else: + raise DSLRuntimeError( + f"Unexpected cute shape element {type(elem).__name__} in {name!r}." + ) + return md.Shape(name=name, dims=dims) + + +def _to_binding( + arg: Any, + name: str, + arg_type: Any, + table: _DimSymbolTable, + *, + is_constexpr: bool = False, +) -> "md.Binding": + """Convert one live argument to a metadata ``Binding``. + + Branch order is load-bearing: bool before int (bool is an int subclass), + value-typed scalars before annotation-typed, NamedTuple before generic tuple. + """ + if is_constexpr: + return _const_binding(name, arg) + + if arg is None: + return md.Unit(name=name) + + # Scalar Numeric, classified by runtime value then by annotation. + if isinstance(arg, Numeric) and arg.dtype in _ACCEPTABLE_SCALAR_NUMERICS: + return md.Var(name=name, dtype=_numeric_to_dtype(arg.dtype)) + if arg_type in _ACCEPTABLE_SCALAR_NUMERICS: + return md.Var(name=name, dtype=_numeric_to_dtype(arg_type)) + + # native_struct: an opaque handle passed by pointer. + if ( + _native_struct_type(arg_type) is not None + or _native_struct_type(arg) is not None + ): + return md.Var(name=name, dtype=DType.handle()) + + if is_cute_algebra_type(arg_type): + return _cute_shape_binding(arg, name, table) + + # Standalone SymInt scalar: a runtime integer, never a shared dim symbol. + if isinstance(arg, SymInt): + dtype = DType.int32() if arg.width == 32 else DType.int64() + return md.Var( + name=name, + dtype=dtype, + divisibility=_normalize_divisibility(arg.divisibility), + ) + + if isinstance(arg, Tensor): + return _tensor_binding(arg, name, table) + + if isinstance(arg, Pointer) or arg_type == Pointer: + space = md.AddressSpace.Generic + if hasattr(arg, "memspace"): + space = _memspace_to_address_space(arg.memspace) + return md.Pointer(name=name, address_space=space) + + if isinstance(arg, _FakeStream): + if arg.use_tvm_ffi_env_stream: + return md.EnvStream(name=name) + return md.Stream(name=name) + + if isinstance(arg, cuda.CUstream): + return md.Stream(name=name) + + # NamedTuple: normalize to a Tuple by field order. + if arg_type is not inspect.Parameter.empty and hasattr(arg_type, "_fields"): + type_hints = get_type_hints(arg_type) + element_types = [type_hints[field] for field in arg_type._fields] + return _tuple_binding(arg, name, element_types, table) + + # Tuple[X, Y, ...] annotation. + if arg_type is not inspect.Parameter.empty and get_origin(arg_type) is tuple: + return _tuple_binding(arg, name, list(get_args(arg_type)), table) + + # Plain tuple/list without element annotations. + if isinstance(arg, (tuple, list)): + return _tuple_binding(arg, name, [None] * len(arg), table) + + # Unannotated runtime scalars (bool before int). + if isinstance(arg, bool): + return md.Var(name=name, dtype=DType.bool_()) + if isinstance(arg, int): + return md.Var(name=name, dtype=DType.int32()) + if isinstance(arg, float): + return md.Var(name=name, dtype=DType.float32()) + + # Dataclass: each field becomes a tuple element; constexpr fields are pinned. + if ( + is_dataclass(arg_type) + if (arg_type is not None and arg_type is not inspect.Parameter.empty) + else is_dataclass(type(arg)) + ): + dc_type = ( + arg_type + if ( + arg_type is not None + and arg_type is not inspect.Parameter.empty + and is_dataclass(arg_type) + ) + else type(arg) + ) + if not isinstance(arg, dc_type): + raise DSLRuntimeError( + f"Expected {dc_type.__name__} for argument {name}, got {type(arg)}" + ) + # Field annotations are strings under PEP 563 (`from __future__ import + # annotations`); resolve them the same way the NamedTuple branch above + # does, so a nested dataclass or a Constexpr field is recognized rather + # than mistaken for an unsupported scalar. Fall back to the raw + # `f.type` if resolution fails (e.g. a locally-scoped class). + try: + field_types = get_type_hints(dc_type) + except Exception: + field_types = {} + values = [] + for f in dataclass_fields(dc_type): + field_type = field_types.get(f.name, f.type) + # Mirror is_constexpr_field's check, but on the resolved type (its + # raw `f.type` is a string under PEP 563). + field_is_constexpr = ( + field_type is Constexpr or get_origin(field_type) is Constexpr + ) + values.append( + _to_binding( + getattr(arg, f.name), + f"{name}.{f.name}", + field_type, + table, + is_constexpr=field_is_constexpr, + ) + ) + return md.Tuple(name=name, values=values) + + if arg_type is not None and ( + get_origin(arg_type) is UnionType or get_origin(arg_type) is Union + ): + for member_type in get_args(arg_type): + try: + return _to_binding(arg, name, member_type, table) + except DSLRuntimeError: + continue + raise DSLRuntimeError( + f"Unsupported argument {name!r} of type {type(arg)} for union " + f"{arg_type}: no member matched." + ) + + raise DSLRuntimeError( + f"Unsupported argument {name!r} of type {type(arg)} (annotation {arg_type})." + ) + + +def _tuple_binding( + arg: Any, + name: str, + element_types: List[Any], + table: _DimSymbolTable, +) -> "md.Tuple": + if not isinstance(arg, (tuple, list)): + raise DSLRuntimeError( + f"Expected tuple/list for argument {name}, got {type(arg)}" + ) + if len(arg) != len(element_types): + raise DSLRuntimeError( + f"Tuple length mismatch for argument {name}: expected " + f"{len(element_types)}, got {len(arg)}" + ) + values = [ + _to_binding(elem, f"{name}[{i}]", element_types[i], table) + for i, elem in enumerate(arg) + ] + return md.Tuple(name=name, values=values) + + +def _native_struct_type(value: Any) -> Optional[type]: + """Return the native_struct class a value or annotation represents, else None.""" + if value is None or value is inspect.Parameter.empty: + return None + if isinstance(value, type): + return value if hasattr(value, "_struct_type") else None + value_type = type(value) + return value_type if hasattr(value_type, "_struct_type") else None + + +# --------------------------------------------------------------------------- +# Signature -> metadata: parameter / return / arg-index helpers +# --------------------------------------------------------------------------- + + +def _rectify_runtime_params( + signature: inspect.Signature, args: Any, kwonlyargs: Any +) -> Iterator[Tuple[inspect.Parameter, Any, bool]]: + """Yield ``(parameter, value, is_constexpr)`` for every top-level signature + param, in order. + + Binds the call args to the signature and applies defaults. A compile-time + ``Constexpr`` param whose value is a representable scalar (``None`` / ``int`` + / ``bool`` / ``float``) is NOT dropped -- it is yielded with + ``is_constexpr=True`` so ``_to_binding`` pins it as a top-level ``Const`` / + ``Unit`` (the baked literal). Such a Const then gets ``abi_arg_index = None`` + and ``ir_arg_index = None`` from ``_assign_arg_indices``: it is recorded in + the metadata calling convention but forwarded to neither the ABI wrapper nor + the kernel (it is baked in). A non-scalar constexpr (a tuple, a type, a CuTe + layout, ...) has no metadata ``Const`` representation, so it is dropped (pure + compile-time specialization, as it always was). The yielded parameters and + values stay aligned. + """ + bound = signature.bind_partial(*args, **kwonlyargs) + bound.apply_defaults() + for index, (name, parameter) in enumerate(signature.parameters.items()): + is_constexpr = is_arg_annotation_constexpr( + parameter.annotation, name, index, None + ) + value = bound.arguments[name] + if is_constexpr and not _is_const_representable(value): + continue + yield parameter, value, is_constexpr + + +def _return_binding(return_annotation: Any, table: _DimSymbolTable) -> "md.Binding": + """Map a function's return annotation to its logical metadata return type. + + A missing / ``None`` / ``NoneType`` annotation is void (``Unit``). Otherwise + the return has no runtime value, so it is dispatched on the annotation alone + through ``_to_binding``: the empty sentinel skips its value-driven branches so + the annotation path runs (e.g. a scalar Numeric -> ``Var``), reusing the same + type-to-binding mapping as parameters. + """ + if ( + return_annotation is inspect.Signature.empty + or return_annotation is None + or return_annotation is type(None) + ): + return md.Unit() + # The return is dispatched on its annotation alone (an empty-sentinel value). + # `_to_binding`'s composite / value-dependent branches (tuple, NamedTuple, + # dataclass, cute algebra) would index that sentinel and crash. A composite + # return is not a representable metadata return anyway -- the ABIs carry only + # void or a scalar -- so reject it up front with a clear message. + if ( + is_cute_algebra_type(return_annotation) + or get_origin(return_annotation) is tuple + or hasattr(return_annotation, "_fields") # NamedTuple + or (isinstance(return_annotation, type) and is_dataclass(return_annotation)) + ): + raise DSLRuntimeError( + f"Unsupported return annotation {return_annotation!r}: a kernel " + f"metadata return must be void or a scalar; composite returns " + f"(tuple / dataclass / cute algebra) are not representable." + ) + return _to_binding(inspect.Parameter.empty, "", return_annotation, table) + + +def _assign_arg_indices(params: list) -> None: + """Assign each binding's two argument indices from the user's signature. + Both are the *generator's* to produce; the compiler only consumes them and + must never re-derive or hardcode them. + + ``ir_arg_index`` -- the IR-function (kernel C-interface) argument index. + Walks in declaration order (recursing into tuples) handing a running index + to every forwarded leaf (Var / Tensor / Stream / Pointer / EnvStream). + Stays ``None`` for what feeds no kernel arg: a baked ``Constexpr`` (-> Const), + a void Unit, a Shape (binds symbols), and the structural tuple itself. + + ``abi_arg_index`` -- the slot index in the ABI call surface. TOP-LEVEL + params only: each occupies one slot (a Tuple is a single slot; its elements + stay ``None`` and are located by tree order within that slot's sub-array) + EXCEPT ``EnvStream`` (stream fetched out-of-band), a top-level ``Const`` (a + baked value), and a top-level ``Unit`` (a baked void / absent optional) -- + none occupy an ABI slot. + """ + ir = 0 + abi = 0 + + def visit_ir(binding: Any) -> None: + nonlocal ir + if isinstance(binding, md.Tuple): + for elem in binding.values: + visit_ir(elem) + elif isinstance( + binding, (md.Var, md.Tensor, md.Stream, md.Pointer, md.EnvStream) + ): + binding.ir_arg_index = ir + ir += 1 + elif isinstance(binding, (md.Const, md.Shape, md.Unit)): + # Const / Shape / Unit forward no kernel argument -- a baked + # constant, a symbol-binding shape, or void -- so they get no + # ir_arg_index and do not advance the counter. + pass + else: + # Fail loudly on a binding type added later but not classified here, + # rather than silently dropping its ir_arg_index. + raise ValueError( + f"_assign_arg_indices: unexpected binding type {type(binding).__name__}" + ) + + for param in params: + visit_ir(param) + # These three top-level bindings occupy NO ABI slot -- they stay + # abi_arg_index=None, and the data-driven wrapper/Executor skip any param + # lacking one, so the ABI surface is exactly the abi-bearing params: + # * EnvStream -- its stream is fetched out-of-band. + # * Const -- a baked compile-time value (recorded, not forwarded). + # * Unit -- a baked void / absent optional (Unit is just a Const + # whose value is "nothing"), so it is treated the same. + if not isinstance(param, (md.EnvStream, md.Const, md.Unit)): + param.abi_arg_index = abi + abi += 1 + + +# --------------------------------------------------------------------------- +# Public entry points +# --------------------------------------------------------------------------- + + +def populate_function_metadata( + fn: "md.FunctionMetadata", + *, + function_name: str, + signature: inspect.Signature, + args: Any, + kwonlyargs: Any, + display_name: Optional[str] = None, +) -> None: + """Populate ``fn`` in place from a live CuTe DSL kernel signature. + + Use with ``FunctionMetadataList.emplace()`` to avoid the deep clone that + ``list.append(FunctionMetadata(...))`` would incur:: + + fn = artifact.metadata.emplace() + populate_function_metadata(fn, function_name=..., signature=..., ...) + """ + # Build all new field values locally before touching ``fn``. If any + # conversion raises, ``fn`` stays in its prior state rather than being left + # half-converted (e.g. a populated ``dim_symbol_table`` with stale ``params``). + table = _DimSymbolTable() + new_params = [ + _to_binding( + value, + parameter.name, + parameter.annotation, + table, + is_constexpr=is_constexpr, + ) + for parameter, value, is_constexpr in _rectify_runtime_params( + signature, args, kwonlyargs + ) + ] + # The generator produces both argument-index maps (metadata -> IR and + # metadata -> ABI slot) from the user's signature. The compiler only + # consumes `ir_arg_index` / `abi_arg_index`; it never re-derives them. + _assign_arg_indices(new_params) + new_symbols = table.symbols + # The function's logical return type, taken from the signature: void (no + # annotation / -> None) is Unit, a scalar return is a Var. This metadata + # records the logical calling convention; the machine-ABI return (e.g. a + # CUDA status code) is the wrapper's separate realization, not modeled here. + new_ret = _return_binding(signature.return_annotation, table) + + fn.symbol_name = function_name + fn.display_name = display_name or function_name + fn.dim_symbol_table.clear() + for symbol in new_symbols: + fn.dim_symbol_table.append(symbol) + fn.params.clear() + for param in new_params: + fn.params.append(param) + fn.ret = new_ret + + +def build_function_metadata( + *, + function_name: str, + signature: inspect.Signature, + args: Any, + kwonlyargs: Any, + display_name: Optional[str] = None, +) -> "md.FunctionMetadata": + """Build a fresh ``FunctionMetadata`` from a live CuTe DSL kernel signature. + + Convenience for one-off use; for population destined for a list slot prefer + ``populate_function_metadata`` + ``FunctionMetadataList.emplace()``. + """ + fn = md.FunctionMetadata() + populate_function_metadata( + fn, + function_name=function_name, + signature=signature, + args=args, + kwonlyargs=kwonlyargs, + display_name=display_name, + ) + return fn diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/helpers.py b/python/CuTeDSL/cutlass/cute/nvgpu/helpers.py index 3efe52d9b..121fbb38d 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/helpers.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/helpers.py @@ -106,8 +106,8 @@ def make_tiled_tma_atom_A( :type tiled_mma: atom.TiledMma :param cluster_shape_vmnk: The Cluster-level shape in VMNK dimensions :type cluster_shape_vmnk: Shape - :param internal_type: An optional parameter for the internal data type to when element - type does not match the copy type + :param internal_type: Optional element-format override used when the + tensor element type does not match the copy type :type internal_type: Type[Numeric] :return: A TmaInfo containing the Copy Atom, TMA tensor, and SMEM layout :rtype: TmaInfo @@ -254,8 +254,8 @@ def make_tiled_tma_atom_B( :type tiled_mma: core.TiledMma :param cluster_shape_vmnk: The Cluster-level shape in VMNK dimensions :type cluster_shape_vmnk: Shape - :param internal_type: An optional parameter for the internal data type to when element - type does not match the copy type + :param internal_type: Optional element-format override used when the + tensor element type does not match the copy type :type internal_type: Type[Numeric] :return: A TmaInfo containing the Copy Atom, TMA tensor, and SMEM layout :rtype: TmaInfo @@ -416,8 +416,8 @@ def make_im2col_tma_atom_A( :type dilation_dhw: Tuple[int, int, int] :param cluster_shape_vmnk: The Cluster-level shape in VMNK dimensions :type cluster_shape_vmnk: Shape - :param internal_type: An optional parameter for the internal data type to when element - type does not match the copy type + :param internal_type: Optional element-format override used when the + tensor element type does not match the copy type :type internal_type: Type[Numeric] :return: A TmaInfo containing the Copy Atom, TMA tensor, and SMEM layout :rtype: TmaInfo diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py index a4cab9074..39318a3f9 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py @@ -67,6 +67,19 @@ _F8F6F4_TYPES = [ #################################################################################################### +def _tcgen05_arch_suggestion(arch: Arch, warp_op_hint: str) -> str: + """tcgen05 (TMEM / UMMA / tcgen05.alloc) is datacenter-Blackwell-only hardware; + consumer Blackwell (sm_120 / sm_121) lacks it, so redirect those users to the + warp-level mma.sync equivalents instead of the generic CUTE_DSL_ARCH hint.""" + if arch.is_family_of(Arch.sm_120f) or arch.is_family_of(Arch.sm_121f): + return ( + f"tcgen05 MMA requires datacenter Blackwell (sm_100 / sm_103). On consumer " + f"Blackwell (sm_120 / sm_121) use the warp-level {warp_op_hint}; see the " + f"blackwell_geforce dense_gemm / blockscaled_gemm examples." + ) + return "Ensure env CUTE_DSL_ARCH matches your GPU architecture" + + class Tcgen05MmaOp(atom.MmaOp): """ Base class for all tcgen05 MMA operations. @@ -295,7 +308,7 @@ class MmaOp(Tcgen05MmaOp): raise OpError( self, f"expects arch to be one of {self.admissible_archs}, but got {arch}", - suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture", + suggestion=_tcgen05_arch_suggestion(arch, "cute.nvgpu.warp.MmaF16BF16Op"), ) # Verify that the user provided enum values if not isinstance(self.cta_group, CtaGroup): @@ -380,7 +393,7 @@ class MmaOp(Tcgen05MmaOp): def __str__(self) -> str: return ( - self.__class__.descriptive_name # type: ignore + self.__class__.descriptive_name + f"\n A data type = {self.a_dtype}" + f"\n B data type = {self.b_dtype}" + f"\n Accumulator data type = {self.acc_dtype}" @@ -514,6 +527,7 @@ class BlockScaledMmaOp(Tcgen05MmaOp): admissible_archs = [ Arch.sm_100a, + Arch.sm_100f, Arch.sm_103a, Arch.sm_110a, ] @@ -525,7 +539,9 @@ class BlockScaledMmaOp(Tcgen05MmaOp): raise OpError( self, f"expects arch to be one of {self.admissible_archs}, but got {arch}", - suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture", + suggestion=_tcgen05_arch_suggestion( + arch, "cute.nvgpu.warp.MmaMXF4Op / MmaMXF4NVF4Op / MmaMXF8Op" + ), ) # Verify that the user provided enum values if not isinstance(self.cta_group, CtaGroup): @@ -594,7 +610,7 @@ class BlockScaledMmaOp(Tcgen05MmaOp): def __str__(self) -> str: return ( - self.__class__.descriptive_name # type: ignore + self.__class__.descriptive_name + f"\n A data type = {self.a_dtype}" + f"\n B data type = {self.b_dtype}" + f"\n Accumulator data type = {self.acc_dtype}" @@ -1575,10 +1591,10 @@ class MmaMXF8Op(BlockScaledMmaOp): Boolean(False).ir_value(loc=loc, ip=ip), Boolean(False).ir_value(loc=loc, ip=ip), core.make_ptr( - self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip + self.sf_dtype, 0, AddressSpace.tmem, loc=loc, ip=ip ).value, core.make_ptr( - self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip + self.sf_dtype, 0, AddressSpace.tmem, loc=loc, ip=ip ).value, ], loc=loc, @@ -1725,10 +1741,10 @@ class MmaMXF8F6F4Op(BlockScaledMmaOp): Boolean(False).ir_value(loc=loc, ip=ip), Boolean(False).ir_value(loc=loc, ip=ip), core.make_ptr( - self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip + self.sf_dtype, 0, AddressSpace.tmem, loc=loc, ip=ip ).value, core.make_ptr( - self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip + self.sf_dtype, 0, AddressSpace.tmem, loc=loc, ip=ip ).value, ], loc=loc, @@ -1858,10 +1874,10 @@ class MmaMXF4Op(BlockScaledMmaOp): Boolean(False).ir_value(loc=loc, ip=ip), Boolean(False).ir_value(loc=loc, ip=ip), core.make_ptr( - self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip + self.sf_dtype, 0, AddressSpace.tmem, loc=loc, ip=ip ).value, core.make_ptr( - self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip + self.sf_dtype, 0, AddressSpace.tmem, loc=loc, ip=ip ).value, ], loc=loc, @@ -1998,10 +2014,10 @@ class MmaMXF4NVF4Op(BlockScaledMmaOp): Boolean(False).ir_value(loc=loc, ip=ip), Boolean(False).ir_value(loc=loc, ip=ip), core.make_ptr( - self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip + self.sf_dtype, 0, AddressSpace.tmem, loc=loc, ip=ip ).value, core.make_ptr( - self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip + self.sf_dtype, 0, AddressSpace.tmem, loc=loc, ip=ip ).value, ], loc=loc, @@ -2133,10 +2149,10 @@ class SM103MmaMXF4Op(BlockScaledMmaOp): Boolean(False).ir_value(loc=loc, ip=ip), Boolean(False).ir_value(loc=loc, ip=ip), core.make_ptr( - self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip + self.sf_dtype, 0, AddressSpace.tmem, loc=loc, ip=ip ).value, core.make_ptr( - self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip + self.sf_dtype, 0, AddressSpace.tmem, loc=loc, ip=ip ).value, ], loc=loc, @@ -2271,10 +2287,10 @@ class SM103MmaMXF4NVF4Op(BlockScaledMmaOp): Boolean(False).ir_value(loc=loc, ip=ip), Boolean(False).ir_value(loc=loc, ip=ip), core.make_ptr( - self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip + self.sf_dtype, 0, AddressSpace.tmem, loc=loc, ip=ip ).value, core.make_ptr( - self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip + self.sf_dtype, 0, AddressSpace.tmem, loc=loc, ip=ip ).value, ], loc=loc, diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/warp/__init__.py b/python/CuTeDSL/cutlass/cute/nvgpu/warp/__init__.py index 2c0c32be6..2ce59d6d5 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/warp/__init__.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/warp/__init__.py @@ -18,6 +18,7 @@ __all__ = [ # mma.py "Field", "MmaF16BF16Op", + "MmaTF32Op", "MmaFP8Op", "MmaMXF4Op", "MmaMXF4NVF4Op", diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py b/python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py index 9df4d9d03..f58673bc4 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py @@ -27,6 +27,7 @@ from ...typing import ( Float16, BFloat16, Float32, + TFloat32, Numeric, Pointer, ) @@ -162,6 +163,70 @@ class MmaF16BF16Trait(Trait): pass +@dataclass(frozen=True) +class MmaTF32Op(WarpMmaOp): + """TF32 warp-level MMA operation. + + This wraps ``mma.sync.aligned.m16n8k{K}.row.col.f32.tf32.tf32.f32``. + Operands are TF32 and accumulation is F32. + """ + + shape_mnk: Shape + + def __post_init__(self) -> None: + if self.shape_mnk not in [(16, 8, 4), (16, 8, 8)]: + raise OpError( + self, + "expects the 'shape_mnk' Op parameter to be one of (16,8,4) or (16,8,8)", + ) + + def _make_trait( + self, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + **kwargs: Any, + ) -> "MmaTF32Trait": + shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip) + ty = _cute_nvgpu_ir.MmaAtomSM80Type.get( + shape_mnk.type.attribute, + TFloat32.mlir_type, + TFloat32.mlir_type, + Float32.mlir_type, + ) + return MmaTF32Trait(make_atom(ty, loc=loc, ip=ip)) + + def __str__(self) -> str: + return ( + "warp-level TF32 MMA Operation" + + "\n A/B data type = TFloat32" + + "\n Accumulator data type = Float32" + + f"\n Instruction shape MNK = {self.shape_mnk}" + ) + + def _verify_fragment_A( + self, + input: Tensor, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> bool: + return True + + def _verify_fragment_B( + self, + input: Tensor, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> bool: + return True + + +class MmaTF32Trait(Trait): + pass + + @dataclass(frozen=True) class MmaFP8Op(WarpMmaOp): """ diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/warpgroup/mma.py b/python/CuTeDSL/cutlass/cute/nvgpu/warpgroup/mma.py index db8b7f9af..472f6c01e 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/warpgroup/mma.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/warpgroup/mma.py @@ -214,7 +214,7 @@ class MmaOp(WarpGroupMmaOp): def __str__(self) -> str: return ( - self.__class__.descriptive_name # type: ignore + self.__class__.descriptive_name + f"\n A data type = {self.a_dtype}" + f"\n B data type = {self.b_dtype}" + f"\n Accumulator data type = {self.acc_dtype}" diff --git a/python/CuTeDSL/cutlass/cute/runtime.py b/python/CuTeDSL/cutlass/cute/runtime.py index 56454582e..3ae652082 100644 --- a/python/CuTeDSL/cutlass/cute/runtime.py +++ b/python/CuTeDSL/cutlass/cute/runtime.py @@ -11,9 +11,7 @@ import ctypes -import sys import math -from pathlib import Path from functools import lru_cache import itertools import operator @@ -21,13 +19,18 @@ from typing import Any, Union, Optional, Type, List, NoReturn # MLIR modules imports from cutlass._mlir import ir -from cutlass.base_dsl.env_manager import get_prefix_dsl_libs import cutlass._mlir.dialects.cute as _cute_ir -import cutlass._mlir.dialects.cuda as _cuda_dialect from cutlass.cutlass_dsl import JitArgAdapterRegistry, CuTeDSL as _CuTeDSL from cutlass.base_dsl.common import DSLRuntimeError -from cutlass.base_dsl.export import ExternalBinaryModule + +# Compatibility re-exports for existing `cutlass.cute.runtime.*` users. +from cutlass.runtime import ( + _FakeStream as _FakeStream, + find_runtime_libraries as find_runtime_libraries, + load_module as load_module, + make_fake_stream as make_fake_stream, +) # Local modules imports from .typing import ( @@ -41,6 +44,8 @@ from .typing import ( TFloat32, Shape, Stride, + _as_address_space, + _to_mlir_address_space, ) from . import core from .tensor import _Tensor as CoreTensor @@ -55,7 +60,7 @@ class _Pointer(Pointer): :param dtype: Data type of the elements pointed to :type dtype: Type :param mem_space: Memory space where the pointer resides, defaults to generic - :type mem_space: _cute_ir.AddressSpace, optional + :type mem_space: AddressSpace, optional :param assumed_align: Assumed alignment of input pointer in bytes, defaults to None :type assumed_align: int, optional @@ -71,12 +76,12 @@ class _Pointer(Pointer): self, pointer: int, dtype: Type[Numeric], - mem_space: _cute_ir.AddressSpace = _cute_ir.AddressSpace.generic, + mem_space: AddressSpace = AddressSpace.generic, assumed_align: Optional[int] = None, ) -> None: self._pointer = pointer self._dtype = dtype - self._addr_space = mem_space + self._addr_space = _as_address_space(mem_space) if assumed_align is None: self._assumed_align = dtype.width // 8 @@ -112,7 +117,9 @@ class _Pointer(Pointer): @property def mlir_type(self) -> ir.Type: return _cute_ir.PtrType.get( - self._dtype.mlir_type, self._addr_space, self._assumed_align + self._dtype.mlir_type, + _to_mlir_address_space(self._addr_space), + self._assumed_align, ) @property @@ -309,10 +316,14 @@ class _Tensor(Tensor): """ self._dtype = new_type + @property + def dtype(self) -> Type[Numeric]: + return self.element_type + @property def memspace(self) -> AddressSpace: self.load_dltensor() - return self._dltensor_wrapper.address_space + return _as_address_space(self._dltensor_wrapper.address_space) @property def size_in_bytes(self) -> int: @@ -542,10 +553,18 @@ class _FakeTensor(Tensor): def __getitem__(self, crd: object) -> NoReturn: raise DSLRuntimeError("runtime._FakeTensor is not indexable") - @property # type: ignore[misc] + @property def element_type(self) -> Type[Numeric]: return self._typed_tensor.element_type + @element_type.setter + def element_type(self, new_type: Type[Numeric]) -> None: + self._typed_tensor.element_type = new_type + + @property + def dtype(self) -> Type[Numeric]: + return self.element_type + @property def memspace(self) -> AddressSpace: return self._typed_tensor.memspace @@ -755,52 +774,6 @@ def make_fake_tensor( ) -class _FakeStream: - """A fake stream that can be used as a placeholder for a stream in compilation. - - When use_tvm_ffi_env_stream is True and the function is compiled with TVM-FFI, - the argument will be skipped from the function signature and we pass in - this value through the environment stream obtained from caller context - (e.g. torch.cuda.current_stream()). - """ - - use_tvm_ffi_env_stream: bool - - def __init__(self, *, use_tvm_ffi_env_stream: bool = False) -> None: - self.use_tvm_ffi_env_stream = use_tvm_ffi_env_stream - - def __str__(self) -> str: - return "FakeStream" - - def __repr__(self) -> str: - return self.__str__() - - def __new_from_mlir_values__(self, values: List[object]) -> object: - assert len(values) == 1 - return values[0] - - def __c_pointers__(self) -> List[int]: - return [0] - - def __get_mlir_types__(self) -> List[ir.Type]: - return [_cuda_dialect.StreamType.get()] - - -def make_fake_stream(*, use_tvm_ffi_env_stream: bool = False) -> _FakeStream: - """Create a fake stream that can be used as a placeholder for a stream in compilation. - - When use_tvm_ffi_env_stream is True and the function is compiled with TVM-FFI, - the argument will be skipped from the function signature and we pass in - this value through the environment stream obtained from caller context - (e.g. torch.cuda.current_stream()). This can speedup the calling process - since we no longer need to do stream query in python. - - :param use_tvm_ffi_env_stream: Whether to skip this parameter use environment stream instead. - :type use_tvm_ffi_env_stream: bool - """ - return _FakeStream(use_tvm_ffi_env_stream=use_tvm_ffi_env_stream) - - def from_dlpack( tensor_dlpack: object, assumed_align: Optional[int] = None, @@ -957,70 +930,6 @@ class TensorAdapter: return self._mlir_types_cache -def find_runtime_libraries(*, enable_tvm_ffi: bool = True) -> List[str]: - """ - Find the runtime libraries that needs to be available for loading modules. - - :param enable_tvm_ffi: Whether to enable TVM-FFI. - :type enable_tvm_ffi: bool, optional - :return: A list of runtime libraries that needs to be available for loading modules. - :rtype: list - """ - - def _get_cute_dsl_runtime_path() -> Optional[str]: - libs = get_prefix_dsl_libs("CUTE_DSL") - if libs is None: - return None - - # check if the separator is ; for windows - if sys.platform.startswith("win32") and ";" in libs: - libs = libs.split(";") # type: ignore[assignment] - else: - libs = libs.split(":") # type: ignore[assignment] - - for path in libs: - if path.endswith("libcute_dsl_runtime.so"): - return path - - return None - - libs = [] - cute_dsl_runtime_path = _get_cute_dsl_runtime_path() - if cute_dsl_runtime_path: - libs.append(cute_dsl_runtime_path) - - if enable_tvm_ffi: - import tvm_ffi - - libs.append(tvm_ffi.libinfo.find_libtvm_ffi()) - - return libs - -# cache to load runtime libraries so they can be found by the DSO loader -_LOAD_MODULE_LIBS_CACHE: list[Any] = [] - - -def load_module( - file_path: str, *, enable_tvm_ffi: bool = False -) -> ExternalBinaryModule: - """Load a module from a file path. - - :param file_path: The path to the module file - :type file_path: str - :param enable_tvm_ffi: Whether to enable TVM-FFI, defaults to True. When True, the module will be loaded as a TVM-FFI module. - :type enable_tvm_ffi: bool, optional - :return: A module object - :rtype: module - """ - if len(_LOAD_MODULE_LIBS_CACHE) == 0: - # ensure the runtime libraries are loaded so they can be found by the DSO loader - # no need to load tvm_ffi library here since it will be loaded by tvm_ffi package. - for path in find_runtime_libraries(enable_tvm_ffi=False): - if Path(path).exists(): - _LOAD_MODULE_LIBS_CACHE.append(ctypes.CDLL(path)) - - return ExternalBinaryModule(file_path, enable_tvm_ffi=enable_tvm_ffi) - # ------------------------------------------------------------------------- # Try to register_jit_arg_adapter for TensorAdapter # ------------------------------------------------------------------------- diff --git a/python/CuTeDSL/cutlass/cute/tensor.py b/python/CuTeDSL/cutlass/cute/tensor.py index d43dc1d4f..8b7581041 100644 --- a/python/CuTeDSL/cutlass/cute/tensor.py +++ b/python/CuTeDSL/cutlass/cute/tensor.py @@ -31,7 +31,8 @@ from cutlass._mlir.dialects import vector, arith, llvm from cutlass._mlir.dialects.cute import ( SparseElemType as _SparseElemType, ) -from cutlass._mlir_helpers.arith import Vector +from cutlass._mlir_helpers.vector import Vector +from cutlass.base_dsl.typing import Pointer as _BasePointer from .typing import ( Numeric, @@ -76,6 +77,7 @@ from .core import ( flatten, has_underscore, make_layout, + make_ptr, select, slice_, crd2idx, @@ -146,6 +148,8 @@ class _Tensor(Tensor): subtensor = tensor[None, 0] # or subtensor = tensor[(None, 0)] """ + _dtype: Union[Type[Numeric], Type[IntTuple], None] + @dsl_user_op def __init__( self, @@ -465,28 +469,39 @@ class _Tensor(Tensor): return self.layout.stride @property - def leading_dim(self) -> Union[int, Tuple[int], None]: + def leading_dim(self) -> Union[int, Tuple[int, ...], None]: """Get the leading dimension of this Tensor. :return: The index or indices of the first mode (from left to right) with stride 1 - :rtype: Union[int, Tuple[int], None] + :rtype: Union[int, Tuple[int, ...], None] :returns: - int: Single leading dimension index if found - - Tuple[int]: Tuple of indices for nested leading dimensions + - Tuple[int, ...]: Tuple of indices for nested leading dimensions - None: If no leading dimension is found :postcondition: ``get(self.stride(), mode=self.leading_dim()) == 1 if self.leading_dim() != None else True`` """ - return leading_dim(self.shape, self.stride) # type: ignore[return-value] + return leading_dim(self.shape, self.stride) @property - def dtype(self) -> Type[Numeric]: - return self._dtype # type: ignore[return-value] + def dtype(self) -> Union[Type[Numeric], Type[IntTuple], None]: # type: ignore[override] + return self._dtype - @property # type: ignore[misc] + @property @lru_cache_ir() - def element_type(self) -> Union[Type[Numeric], Type[IntTuple]]: - return self._dtype # type: ignore[return-value] + def element_type(self) -> Union[Type[Numeric], Type[IntTuple], None]: + return self._dtype + + @element_type.setter + def element_type( + self, new_type: Union[Type[Numeric], Type[IntTuple], None] + ) -> None: + # Compile-time _Tensor's dtype is bound to the underlying MLIR ir.Value; + # mutating _dtype here would desync from self.value.type, and the + # @lru_cache_ir() on the getter would serve stale values. + raise NotImplementedError( + "Compile-time _Tensor does not support setting element_type" + ) @property @lru_cache_ir() @@ -797,6 +812,15 @@ def make_tensor( elif isinstance(iterator, Pointer): iterator = iterator.value res_ty = _cute_ir.MemRefType.get(iterator.type, layout.type) # type: ignore[union-attr] + elif isinstance(iterator, _BasePointer): + iterator = make_ptr( + iterator.dtype, + iterator.to_llvm_ptr(loc=loc, ip=ip), + assumed_align=iterator.max_alignment, + loc=loc, + ip=ip, + ).value + res_ty = _cute_ir.MemRefType.get(iterator.type, layout.type) # type: ignore[union-attr] elif isinstance(iterator, ir.Value) and isinstance( iterator.type, _cute_nvgpu_ir.SmemDescType, @@ -994,7 +1018,7 @@ def make_rmem_tensor_like( compact_layout = make_layout(src.shape, loc=loc, ip=ip) src_layout = _cute_ir.make_layout_like(compact_layout, loc=loc, ip=ip) else: - res_dtype = dtype or src.element_type # type: ignore[assignment] + res_dtype = dtype or src.element_type src_layout = src.layout elif isinstance(src, TensorSSA): res_dtype = dtype or src.element_type @@ -1109,7 +1133,7 @@ def recast_tensor( if src.element_type is Boolean: src_width = 8 else: - src_width = src.element_type.width # type: ignore[union-attr] + src_width = src.element_type.width src_iter = recast_ptr(src.iterator, dtype=dtype, loc=loc, ip=ip) src_layout = recast_layout(dst_width, src_width, src.layout, loc=loc, ip=ip) diff --git a/python/CuTeDSL/cutlass/cute/typing.py b/python/CuTeDSL/cutlass/cute/typing.py index e406abbcb..6da5c4484 100644 --- a/python/CuTeDSL/cutlass/cute/typing.py +++ b/python/CuTeDSL/cutlass/cute/typing.py @@ -55,6 +55,7 @@ from cutlass.base_dsl.typing import ( Float4E2M1FN, Float6E2M3FN, Float6E3M2FN, + Float4E2M1FNx2, as_numeric, ) @@ -62,7 +63,8 @@ _element_precision_width = _base_typing._element_precision_width from cutlass._mlir import ir import cutlass._mlir.dialects.cute as _cute_ir -from cutlass._mlir.dialects.cute import AddressSpace, ConstrainedIntType +from cutlass._mlir.dialects.cute import ConstrainedIntType +from cutlass.address_space import AddressSpace if TYPE_CHECKING: from cutlass.cute.core import ScaledBasis, Swizzle @@ -74,6 +76,24 @@ else: Int = Union[int, Integer] +def _as_address_space(memspace: AddressSpace | int) -> AddressSpace: + if isinstance(memspace, AddressSpace): + return memspace + if isinstance(memspace, bool): + raise TypeError("expects mem_space to be an AddressSpace, not bool") + value = getattr(memspace, "value", memspace) + try: + return AddressSpace(int(value)) + except (TypeError, ValueError): + raise TypeError( + f"expects mem_space to be an AddressSpace, but got {memspace}" + ) from None + + +def _to_mlir_address_space(memspace: AddressSpace | int) -> _cute_ir.AddressSpace: + return _cute_ir.AddressSpace(int(_as_address_space(memspace))) + + class SymInt: r"""A symbolic integer for runtime-bound dimensions of an AOT-compiled ``@cute.jit`` function. @@ -241,6 +261,33 @@ class SymInt: def __rmul__(self, other: "int | SymInt") -> "SymInt": return self.__mul__(other) + def __floordiv__(self, other: "int | SymInt") -> "SymInt": + """SymInt // int | SymInt: floor-divide a symbolic integer. + + The result is always a fresh ``SymInt`` (the quotient is symbolic + since ``self`` is). Its divisibility is only as strong as can be + *proven* from the divisibility contract: + + * ``self // k`` (``k`` an ``int``): if ``self``'s divisibility is a + multiple of ``k`` the division is exact on the contract, so the + quotient is a multiple of ``divisibility // k``. Otherwise nothing + is provable and the result advertises ``divisibility=1``. + * ``self // other`` (``other`` a ``SymInt``): ``other``'s runtime + value is only known to be *some* multiple of its divisibility, so + the quotient's alignment is not provable — ``divisibility=1``. + """ + if isinstance(other, int): + if other == 0: + raise ZeroDivisionError("integer division or modulo by zero") + result_width = self._width + if self._divisibility % other == 0: + return SymInt(result_width, divisibility=self._divisibility // other) + return SymInt(result_width, divisibility=1) + elif isinstance(other, SymInt): + return SymInt(width=max(self._width, other._width), divisibility=1) + else: + return NotImplemented + def __c_pointers__(self) -> List[int | None]: return [ctypes.c_void_p(0).value] @@ -489,7 +536,7 @@ XTuple = Union[Any, Tuple["XTuple", ...]] class Pointer(ABC): """ - Abstract base class for CuTe jit function and runtime _Pointer + Abstract base class for CuTe jit function and runtime _Pointer. """ value: ir.Value @@ -509,7 +556,7 @@ class Pointer(ABC): ... @property - def memspace(self) -> AddressSpace: + def memspace(self) -> AddressSpace: # type: ignore[empty-body] """The memory address space of this pointer. Implemented by subclasses.""" ... @@ -542,6 +589,14 @@ class Pointer(ABC): """Get the LLVM pointer representation (loc/ip propagated). Implemented by subclasses.""" raise NotImplementedError + def ir_value( + self, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> ir.Value: + return self.to_llvm_ptr(loc=loc, ip=ip) + def toint( self, *, @@ -629,7 +684,7 @@ class TypedTensor: self._dtype = dtype self._shape = shape self._stride = stride - self._memspace = memspace + self._memspace = _as_address_space(memspace) self._assumed_align = assumed_align if assumed_align is None: # use the bytes width of the element dtype. The alignment is at least one byte align. @@ -639,6 +694,12 @@ class TypedTensor: def element_type(self) -> Type[Numeric]: return self._dtype + @element_type.setter + def element_type(self, new_type: Type[Numeric]) -> None: + # TypedTensor is a compile-time type descriptor that's intentionally + # read-only. + raise NotImplementedError + @property def shape(self) -> Shape: return self._shape @@ -671,7 +732,9 @@ class TypedTensor: # Boolean types are stored as i8 in memory elem_type = T.i8() if self._dtype.width == 1 else self._dtype.mlir_type - ptr_ty = _cute_ir.PtrType.get(elem_type, self._memspace, self._assumed_align) + ptr_ty = _cute_ir.PtrType.get( + elem_type, _to_mlir_address_space(self._memspace), self._assumed_align + ) return _cute_ir.MemRefType.get(ptr_ty, layout_ty) @@ -764,14 +827,15 @@ class Tensor(ABC): @property @abstractmethod - def element_type(self) -> Union[Type[Numeric], Type[IntTuple]]: ... + def element_type(self) -> Type[Numeric]: ... @element_type.setter - def element_type(self, new_type: Union[Type[Numeric], Type[IntTuple]]) -> None: + def element_type(self, new_type: Type[Numeric]) -> None: """Implemented by subclasses.""" raise NotImplementedError @property + @abstractmethod def dtype(self) -> Type[Numeric]: """The element data type. Implemented by subclasses.""" raise NotImplementedError @@ -813,6 +877,16 @@ class Tensor(ABC): @abstractmethod def stride(self) -> Stride: ... + @property + def dynamic_shapes_mask(self) -> Tuple[int, ...]: + """Implemented by subclasses.""" + raise NotImplementedError + + @property + def dynamic_strides_mask(self) -> Tuple[int, ...]: + """Implemented by subclasses.""" + raise NotImplementedError + def load( self, *, @@ -922,6 +996,7 @@ __all__ = [ "Float4E2M1FN", "Float6E2M3FN", "Float6E3M2FN", + "Float4E2M1FNx2", "IntTuple", "ScaledBasis", "Coord", diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/__init__.py b/python/CuTeDSL/cutlass/cutlass_dsl/__init__.py index ff7c89523..43c071aaf 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/__init__.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/__init__.py @@ -51,6 +51,8 @@ from ..base_dsl.compiler import ( GenerateLineInfo, KeepCUBIN, KeepPTX, + KeepSASS, + NvdisasmOptions, GPUArch, LinkLibraries, EnableTVMFFI, diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py b/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py index a8c2dd774..842c9cdb8 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py @@ -27,12 +27,14 @@ import cuda.bindings.driver as cuda_driver from ..base_dsl.jit_executor import ( JitExecutor, JitCompiledFunction, + JitModule, ExecutionArgs, JitFunctionArtifacts, ) from ..base_dsl.compiler import HostTarget from ..base_dsl.utils.logger import log -from ..base_dsl.common import DSLRuntimeError +from ..base_dsl.common import DSLRuntimeError, DSLUserCodeError +from ..base_dsl.diagnostics import DiagId from ..base_dsl.typing import Int32 from ..base_dsl.runtime.cuda import checkCudaErrors @@ -75,6 +77,7 @@ class CudaDialectJitModule: class CudaDialectJitCompiledFunction(JitCompiledFunction): """Holds a compiled function and its module.""" + jit_module: "JitModule | CudaDialectJitModule | None" # type: ignore[assignment] device_header: Optional[str] device_object_path: Optional[str] device_ptx_path: Optional[str] @@ -156,18 +159,18 @@ class CudaDialectJitCompiledFunction(JitCompiledFunction): raise DSLRuntimeError("cuda_load not found") cuda_load = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(cuda_load) - cuda_init_args = [pointer_to_pointer_to_library, pointer_to_err] + cuda_init_args: List[Any] = [pointer_to_pointer_to_library, pointer_to_err] packed_args = (ctypes.c_void_p * len(cuda_init_args))() for i in range(len(cuda_init_args)): - packed_args[i] = ctypes.cast(cuda_init_args[i], ctypes.c_void_p) # type: ignore[arg-type] + packed_args[i] = ctypes.cast(cuda_init_args[i], ctypes.c_void_p) cuda_init(packed_args) checkCudaErrors((cuda_runtime.cudaError_t(err.value),)) - cuda_load_args = [pointer_to_library, pointer_to_err] + cuda_load_args: List[Any] = [pointer_to_library, pointer_to_err] packed_args = (ctypes.c_void_p * len(cuda_load_args))() for i in range(len(cuda_load_args)): - packed_args[i] = ctypes.cast(cuda_load_args[i], ctypes.c_void_p) # type: ignore[arg-type] + packed_args[i] = ctypes.cast(cuda_load_args[i], ctypes.c_void_p) cuda_load(packed_args) checkCudaErrors((cuda_runtime.cudaError_t(err.value),)) @@ -191,7 +194,7 @@ class CudaDialectJitCompiledFunction(JitCompiledFunction): # and cuda_load_to_device from the engine which are defined in CudaToLLVM. if self.load_from_binary: if self.prefix is None: - raise DSLRuntimeError("prefix is required to be set for binary loading") + raise DSLUserCodeError(DiagId.CONFIG_BINARY_LOAD_PREFIX_REQUIRED) cuda_init = self.engine.lookup(f"_mlir_{self.prefix}_cuda_init") if cuda_init is None: raise DSLRuntimeError(f"cuda_init not found for prefix {self.prefix}") @@ -227,10 +230,10 @@ class CudaDialectJitCompiledFunction(JitCompiledFunction): err = ctypes.c_int32(0) pointer_to_err = ctypes.pointer(err) - cuda_init_args = [pointer_to_pointer_to_library, pointer_to_err] + cuda_init_args: List[Any] = [pointer_to_pointer_to_library, pointer_to_err] packed_args = (ctypes.c_void_p * len(cuda_init_args))() for i in range(len(cuda_init_args)): - packed_args[i] = ctypes.cast(cuda_init_args[i], ctypes.c_void_p) # type: ignore[arg-type] + packed_args[i] = ctypes.cast(cuda_init_args[i], ctypes.c_void_p) cuda_init(packed_args) checkCudaErrors((cuda_runtime.cudaError_t(err.value),)) @@ -238,14 +241,14 @@ class CudaDialectJitCompiledFunction(JitCompiledFunction): device_id = ctypes.c_int32(0) pointer_to_device_id = ctypes.pointer(device_id) - cuda_load_args = [ + cuda_load_args: List[Any] = [ pointer_to_pointer_to_library, pointer_to_device_id, pointer_to_err, ] packed_args = (ctypes.c_void_p * len(cuda_load_args))() for i, arg in enumerate(cuda_load_args): - packed_args[i] = ctypes.cast(arg, ctypes.c_void_p) # type: ignore[arg-type] + packed_args[i] = ctypes.cast(arg, ctypes.c_void_p) for dev in range(self.num_devices): device_id.value = dev @@ -279,7 +282,7 @@ class CudaDialectJitCompiledFunction(JitCompiledFunction): and self.jit_module.is_unloaded() ): cuda_library = self._load_cuda_library() if self.has_gpu_module else [] - self.jit_module = CudaDialectJitModule( # type: ignore[assignment] + self.jit_module = CudaDialectJitModule( self.engine, self.capi_func, self.execution_args, diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py b/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py index 875f57d7a..7be0662b9 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py @@ -34,7 +34,6 @@ from typing import ( ) import functools import inspect -import pkgutil from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import fields from math import ceil @@ -59,8 +58,8 @@ from ..base_dsl.dsl import ( ) from ..base_dsl.typing import * from ..base_dsl.typing import ( - DynamicExpression, get_mlir_types, + Pointer, Int32, Int64, Int8, @@ -73,7 +72,13 @@ from ..base_dsl.typing import ( get_c_pointers, cast, ) -from ..base_dsl.common import DSLRuntimeError, DSLNotImplemented +from ..base_dsl.common import ( + DSLRuntimeError, + DSLUserCodeError, + DSLNotImplemented, + active_env_manager, +) +from ..base_dsl.diagnostics import DiagId from ..base_dsl.utils.logger import log from ..base_dsl.utils.tree_utils import ( Leaf, @@ -85,7 +90,7 @@ from ..base_dsl.utils.tree_utils import ( ) from ..base_dsl.leaf_utils import is_frozen_dataclass from ..base_dsl.runtime.jit_arg_adapters import is_arg_annotation_constexpr -from ..base_dsl.jit_executor import ExecutionArgs # noqa: F401 +from ..base_dsl.jit_executor import ExecutionArgs, _is_pointer_annotation # noqa: F401 from ..base_dsl.runtime import cuda as cuda_helpers from .cuda_stream_adapter import CudaDriverStreamAdapter, CudaRuntimeStreamAdapter # noqa: F401 from .cuda_library_adapter import CudaDialectLibraryAdapter # noqa: F401 @@ -210,6 +215,42 @@ def is_cute_algebra_type(arg_spec: object) -> bool: return False +def _is_cutlass_pointer_annotation(annotation: object) -> bool: + if get_origin(annotation) is Annotated: + annotation = get_args(annotation)[0] + return type(annotation).__name__ == "TypedPointer" or ( + getattr(annotation, "__name__", None) == "Pointer" + and getattr(annotation, "__module__", "").endswith("cutlass.base_dsl.typing") + ) + + +def _is_cute_pointer_like(arg: object) -> bool: + module = type(arg).__module__ + return ( + module.startswith("cutlass.cute") + and hasattr(arg, "dtype") + and hasattr(arg, "memspace") + and ( + hasattr(arg, "__get_mlir_types__") + or hasattr(arg, "__extract_mlir_values__") + ) + ) + + +def _cutlass_pointer_dtype_addrspace( + annotation: object, + arg: object, +) -> tuple[type[Numeric], int]: + if type(annotation).__name__ == "TypedPointer": + dtype = annotation.dtype # type: ignore[attr-defined] + space = annotation.space # type: ignore[attr-defined] + else: + dtype = getattr(arg, "dtype", Int8) + space = getattr(arg, "memspace", 0) + addrspace = getattr(space, "value", space) + return dtype, int(addrspace) + + def _build_kernel_attrs(config: BaseDSL.LaunchConfig) -> dict: """Build launch-time CUDA function attrs that are known before compilation.""" ATTR_SMEM_CARVEOUT = cuda_helpers.cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT @@ -254,13 +295,8 @@ class CutlassBaseDSL(BaseDSL): def attach_and_decorate(func: Callable[..., None]) -> Callable[..., None]: if target_cls._KERNEL_ATTR_SPEC_FIELD is None: - raise DSLRuntimeError( - f"{target_cls.__name__} does not support kernel-level 'attributes='.", - suggestion=( - "Only DSLs that set _KERNEL_ATTR_SPEC_FIELD support the " - "'attributes=' kwarg on @kernel; remove the argument or use " - "a DSL that supports it." - ), + raise DSLUserCodeError( + DiagId.CONFIG_ATTRIBUTES_UNSUPPORTED, ) setattr(func, target_cls._KERNEL_ATTR_SPEC_FIELD, attr_spec) return kernel_decorator(func) @@ -283,6 +319,7 @@ class CutlassBaseDSL(BaseDSL): device_compilation_only=device_compilation_only, preprocess=preprocess, ) + self._smem_partition_num: Optional[int] = 2 self._smem_usage_tracker: Optional[tuple] = None # extra function to convert cute arguments to tvm ffi spec params # this needs to be reverse registered because the arg convention @@ -334,15 +371,13 @@ class CutlassBaseDSL(BaseDSL): if resolved is None: return {} if not isinstance(resolved, dict): - raise DSLRuntimeError( - "Kernel attribute resolver must return a dict or None.", - suggestion="Return a dict[str, str | ir.Attribute], or None for no attributes.", + raise DSLUserCodeError( + DiagId.CONFIG_RESOLVER_INVALID_RETURN, ) return resolved - raise DSLRuntimeError( - f"Unsupported kernel decorator attributes spec type: {type(attr_spec)}", - suggestion="Use a dict or a callable returning dict[str, str | ir.Attribute].", + raise DSLUserCodeError( + DiagId.CONFIG_ATTRIBUTES_INVALID_TYPE, ) def _collect_extra_kernel_value_attrs( @@ -361,8 +396,9 @@ class CutlassBaseDSL(BaseDSL): suggestion = f"Use one of the allowed keys: {allowed_keys}." else: suggestion = f"No extra kernel function attributes are supported for '{self.name}'." - raise DSLRuntimeError( - f"Unsupported kernel function attribute key '{key}'.", + raise DSLUserCodeError( + DiagId.CONFIG_ATTR_KEY_UNSUPPORTED, + key=key, suggestion=suggestion, ) if isinstance(value, ir.Attribute): @@ -370,9 +406,9 @@ class CutlassBaseDSL(BaseDSL): elif isinstance(value, str): converted_attrs[key] = ir.StringAttr.get(value) else: - raise DSLRuntimeError( - f"Unsupported kernel function attribute value type for '{key}': {type(value)}", - suggestion="Use str or ir.Attribute as the attribute value.", + raise DSLUserCodeError( + DiagId.CONFIG_ATTR_VALUE_TYPE_INVALID, + key=key, ) return converted_attrs @@ -456,9 +492,8 @@ class CutlassBaseDSL(BaseDSL): max_ntid_str = ", ".join(map(str, config.max_number_threads)) for value in config.max_number_threads: if is_dynamic_expression(value): - raise DSLRuntimeError( - f"User defined max number threads `{max_ntid_str}` contains dynamic expression, cannot generate `nvvm.maxntid`", - suggestion="Consider using `Constexpr` annotation or Python constant", + raise DSLUserCodeError( + DiagId.CONFIG_MAX_THREADS_DYNAMIC, ) # Use user-defined maxntid ret["nvvm.maxntid"] = ir.Attribute.parse(f"array") @@ -484,20 +519,27 @@ class CutlassBaseDSL(BaseDSL): if config.smem_merge_branch_allocs: ret["smem.merge_branch_allocs"] = ir.UnitAttr.get() + # Set smem partition num + if self._smem_partition_num is not None: + ret["smem.partition_num"] = ir.IntegerAttr.get( + ir.IntegerType.get_signless(32), self._smem_partition_num + ) + return ret @staticmethod def _normalize_static_cluster_dim(dim: Any, config_name: str) -> int: if is_dynamic_expression(dim): - raise DSLRuntimeError( - f"`LaunchConfig.{config_name}` contains dynamic expression, cannot materialize compiler-visible cluster metadata", - suggestion="Consider using `Constexpr` annotation or Python constant", + raise DSLUserCodeError( + DiagId.CONFIG_CLUSTER_DYNAMIC, + config_name=config_name, ) value = dim.value if isinstance(dim, Integer) else dim if isinstance(value, bool) or not isinstance(value, Integral): - raise DSLRuntimeError( - f"`LaunchConfig.{config_name}` must contain integer dimensions, but got {type(value)}" + raise DSLUserCodeError( + DiagId.TYPE_CLUSTER_NOT_INT, + config_name=config_name, ) return int(value) @@ -517,8 +559,8 @@ class CutlassBaseDSL(BaseDSL): cls, config: BaseDSL.LaunchConfig ) -> dict[str, ir.Attribute]: if config.has_fallback_cluster and not config.has_cluster: - raise DSLRuntimeError( - "`LaunchConfig.fallback_cluster` requires `LaunchConfig.cluster` to also be set" + raise DSLUserCodeError( + DiagId.CONFIG_FALLBACK_CLUSTER_REQUIRES_CLUSTER, ) if config.has_fallback_cluster: @@ -647,19 +689,51 @@ class CutlassBaseDSL(BaseDSL): ) from e files.append((giant_dso_name, so_path, so_size)) - for lib in pkgutil.walk_packages([dsl_path], prefix="cutlass."): - spec = lib.module_finder.find_spec(lib.name) # type: ignore[call-arg] - if not spec or not spec.origin: + # Walk the filesystem to collect .py files for hashing. We deliberately + # avoid pkgutil.walk_packages here because it imports every module it + # discovers — and modules with top-level decorator side effects (e.g. + # _mlir_helpers/arith.py's @ir.register_value_caster) blow up if they + # end up reachable via two import paths in the same process. The hash + # only needs file contents, not importability. + for entry in sorted(os.listdir(dsl_path)): + sub = os.path.join(dsl_path, entry) + if not os.path.isdir(sub): continue - path = spec.origin - try: - size = os.path.getsize(path) - except Exception as e: - raise DSLRuntimeError( - f"Failed to read module file {lib.name}. The file may not exist or may not be readable." - "Please re-install the package." - ) from e - files.append((lib.name, path, size)) + if not os.path.isfile(os.path.join(sub, "__init__.py")): + # Skip non-package directories (build/, examples/, doc/, ...). + # This matches pkgutil.walk_packages, which only descends into + # regular packages found on the given path list. + continue + for root, dirs, fnames in os.walk(sub): + # Only descend into subdirectories that are themselves + # regular packages — matches pkgutil.walk_packages, which + # treats dirs without __init__.py as outside the import + # namespace and skips them. + dirs[:] = [ + d + for d in dirs + if d != "__pycache__" + and os.path.isfile(os.path.join(root, d, "__init__.py")) + ] + for fname in fnames: + if not fname.endswith(".py"): + continue + path = os.path.join(root, fname) + rel = os.path.relpath(path, dsl_path) + no_ext = rel[: -len(".py")] + init_suffix = os.sep + "__init__" + if no_ext.endswith(init_suffix): + no_ext = no_ext[: -len(init_suffix)] + key = "cutlass." + no_ext.replace(os.sep, ".") + try: + size = os.path.getsize(path) + except OSError as e: + raise DSLRuntimeError( + f"Failed to read module file {key}. " + "The file may not exist or may not be readable. " + "Please re-install the package." + ) from e + files.append((key, path, size)) # Submit chunks to a job queue chunk_size = 1 << 24 # 16 MB (tuned) @@ -717,6 +791,24 @@ class CutlassBaseDSL(BaseDSL): zero = arith.ConstantOp(return_type, 0).result return [zero] + def _runtime_can_run_compiled( + self, runtime_arch: Optional[str], compiled_arch: Optional[str] + ) -> bool: + """Allow a family-portable target (e.g. sm_100f) to run on its family + hardware (e.g. sm_100a), in addition to an exact arch match.""" + if runtime_arch == compiled_arch: + return True + if runtime_arch is None or compiled_arch is None: + return False + from ..base_dsl.arch import Arch + + try: + return Arch.from_string(runtime_arch).can_run_binary_built_for( + Arch.from_string(compiled_arch) + ) + except KeyError: + return False + def compile_and_cache( self, module: ir.Module, @@ -1251,6 +1343,10 @@ class CutlassBaseDSL(BaseDSL): if not any( (ty is Any) or (isinstance(ty, type) and isinstance(arg, ty)) + or ( + _is_cutlass_pointer_annotation(ty) + and _is_cute_pointer_like(arg) + ) or (get_origin(ty) is tuple and isinstance(arg, tuple)) or ( get_origin(ty) is Annotated and isinstance(arg, get_args(ty)[0]) @@ -1271,6 +1367,16 @@ class CutlassBaseDSL(BaseDSL): # allow FakeStream to be passed as CUDA stream pass + elif getattr( + arg, "_python_pointer_address_arg", False + ) is True and _is_pointer_annotation(arg_annotation): + pass + + elif _is_cutlass_pointer_annotation( + arg_annotation + ) and _is_cute_pointer_like(arg): + pass + elif isinstance(arg_annotation, type): # Handle simple type annotations if not isinstance(arg, arg_annotation) and arg is not None: @@ -1311,7 +1417,17 @@ class CutlassBaseDSL(BaseDSL): assert jit_arg_attr is not None assert jit_exec_arg is not None # Handle DSL specific types - if is_cute_algebra_type(arg_spec): + if _is_cutlass_pointer_annotation(arg_spec) and _is_cute_pointer_like(arg): + _dtype, addrspace = _cutlass_pointer_dtype_addrspace(arg_spec, arg) + from cutlass._mlir.dialects import llvm as llvm_dialect + + jit_arg_type.append(llvm_dialect.PointerType.get(addrspace)) + jit_arg_attr.append(default_attr) + if is_host: + jit_exec_arg.extend(get_c_pointers(arg)) + else: + jit_exec_arg.append(arg.to_llvm_ptr()) + elif is_cute_algebra_type(arg_spec): dyn_vals = extract_mlir_values(arg) if dyn_vals: # Handle dynamic types @@ -1366,7 +1482,12 @@ class CutlassBaseDSL(BaseDSL): ) if not ir_arg: # Handling DSL specific types - if is_cute_algebra_type(arg_spec): + if _is_cutlass_pointer_annotation(arg_spec) and _is_cute_pointer_like(arg): + dtype, addrspace = _cutlass_pointer_dtype_addrspace(arg_spec, arg) + blk_args = fop_args[iv_block_args : iv_block_args + 1] + ir_arg.append(Pointer(blk_args[0], dtype=dtype, space=addrspace)) + iv_block_args += 1 + elif is_cute_algebra_type(arg_spec): n_args = len(get_mlir_types(arg)) blk_args = fop_args[iv_block_args : iv_block_args + n_args] ir_arg.append(new_from_mlir_values(arg, blk_args)) @@ -1535,6 +1656,13 @@ class CuTeDSL(CutlassBaseDSL): this generates only a gpu.module containing a cuda.func with no host-side wrapper and no launch op. Supports return values. """ + # Device compilation can be invoked directly by cute.compile. + with active_env_manager(self.envar): + return self._device_func_impl(funcBody, *args, **kwargs) + + def _device_func_impl( + self, funcBody: Callable[..., Any], *args: Any, **kwargs: Any + ) -> Any: if ir.Context.current is not None and ir.InsertionPoint.current is not None: return funcBody(*args, **kwargs) @@ -1640,9 +1768,8 @@ class CuTeDSL(CutlassBaseDSL): elif hasattr(ret_annotation, "_struct_type"): ret_types = [ret_annotation._struct_type] else: - raise DSLRuntimeError( - f"Device function return type annotation must be a DSL type " - f"(e.g. Float32) or @native_struct, got {ret_annotation}" + raise DSLUserCodeError( + DiagId.TYPE_DEVICE_FUNC_RETURN_INVALID, ) loc = self.get_ir_location(setup.location) @@ -1677,18 +1804,17 @@ class CuTeDSL(CutlassBaseDSL): # Generate return op if ret_types: if result is None: - raise DSLRuntimeError( - f"Device function '{funcBody.__name__}' has return " - f"type annotation but returned None" + raise DSLUserCodeError( + DiagId.TYPE_DEVICE_FUNC_RETURN_NONE, ) if hasattr(result, "ir_value"): ret_val = result.ir_value() elif hasattr(result, "__extract_mlir_values__"): extracted_vals = result.__extract_mlir_values__() if len(extracted_vals) != 1: - raise DSLRuntimeError( - f"Device function '{funcBody.__name__}' returned " - f"{len(extracted_vals)} MLIR values; expected 1" + raise DSLUserCodeError( + DiagId.TYPE_DEVICE_FUNC_RETURN_COUNT, + count=len(extracted_vals), ) ret_val = extracted_vals[0] else: @@ -1710,6 +1836,7 @@ class CuTeDSL(CutlassBaseDSL): pointer_types=ptr_types, ) + self._run_trace_finalize_hooks(module, setup.function_name) module = self.build_module(module, setup.function_name) # dryrun: generate IR and header, skip compilation @@ -1775,12 +1902,10 @@ class _CuteExperimentalJitCompiledFunction(CudaDialectJitCompiledFunction): n = self.execution_args._meta.arg_count n_extra = builtins.max(0, len(args) - n) if n_extra != self.total_added_arguments: - raise DSLRuntimeError( - "Wrong number of extra workspace arguments", - context={ - "expected": self.total_added_arguments, - "got": n_extra, - }, + raise DSLUserCodeError( + DiagId.ARG_WORKSPACE_COUNT_MISMATCH, + expected=self.total_added_arguments, + got=n_extra, ) return super().__call__(*args, **kwargs) @@ -1791,8 +1916,9 @@ class _CuteExperimentalJitCompiledFunction(CudaDialectJitCompiledFunction): class CuteExperimentalDSL(CutlassBaseDSL): + _KERNEL_EXTRA_ARGS_ATTR = "lir.kernel_extra_args" _ALLOWED_EXTRA_KERNEL_VALUE_ATTRS: frozenset[str] = frozenset( - {"lir.tma_update_mode"} + {"lir.tma_update_mode", "lir.tma_override_mode"} ) _KERNEL_ATTR_SPEC_FIELD: Optional[str] = "_cute_experimental_kernel_attributes" # Marks this class as the deprecated "experimental" CuTe DSL so @@ -1931,7 +2057,7 @@ class CuteExperimentalDSL(CutlassBaseDSL): original_function_name=original_function_name, funcBody=funcBody, ) - # Extract the kernel_extra_args attribute written by FuncFinalizationPass + # Extract the lir.kernel_extra_args attribute written by FuncFinalizationPass # and store it on the compiled function for later retrieval. # Maps kernel name → number of extra workspace pointer args added to the # host entry point signature. @@ -1939,13 +2065,15 @@ class CuteExperimentalDSL(CutlassBaseDSL): fn.total_added_arguments = 0 if fn.ir_module is not None: attrs = fn.ir_module.operation.attributes - if "kernel_extra_args" in attrs: - for named in ir.DictAttr(attrs["kernel_extra_args"]): + attr_name = self._KERNEL_EXTRA_ARGS_ATTR + if attr_name in attrs: + for named in ir.DictAttr(attrs[attr_name]): fn.kernel_extra_args[named.name] = ir.IntegerAttr(named.attr).value fn.total_added_arguments = fn.kernel_extra_args.pop( "total_added_arguments", 0 ) - fn.__class__ = CuteExperimentalDSL.JitCompiledFunction + if type(fn) is CudaDialectJitCompiledFunction: + fn.__class__ = CuteExperimentalDSL.JitCompiledFunction return fn @@ -2001,8 +2129,8 @@ class KernelLauncher: try: sig.bind(*func_args, **func_kwargs) except TypeError as e: - raise DSLRuntimeError( - f"Failed to bind arguments to function `{funcBody.__name__}` with signature `{sig}`", + raise DSLUserCodeError( + DiagId.ARG_BIND_FAILED, cause=e, ) @@ -2194,19 +2322,16 @@ def unpack_to_irvalue( remove_read_only_frozen_dataclass(mixed_values, full_write_args_count) ) except DSLTreeFlattenError as e: - raise DSLRuntimeError( - f"The '{body_name}' statement encountered a user-defined Python object, which cannot be automatically converted into an dynamic expression.", - context={ - e.message: ( - f"All expressions within '{body_name}' must be dynamic expressions, " - "mixing Python objects and dynamic expressions is not supported. " - "The DSL failed to convert the Python object into dynamic expressions." - ) - }, - suggestion=( - f"Please ensure '{e.type_str}' implements the '{DynamicExpression.__name__}' or mark with `dataclass`, " - f"so it can be treated as a valid dynamic expression or mark '{body_name}' as a constant expression if conditions are Python objects." - ), + # Strip the "builtins." prefix so the author sees a plain type name. + py_type = str(e.type_str).rsplit(".", 1)[-1] + # Keep a specific reason (e.g. "`MyClass(...)` has extra field `a`") as a + # note, but drop the generic internal "Flatten Error" marker. + detail = str(e.message) if e.message else "" + raise DSLUserCodeError( + DiagId.TYPE_DYNAMIC_EXPR_UNSUPPORTED, + body_name=body_name, + py_type=py_type, + context=detail if detail and "Flatten Error" not in detail else None, ) # DEBUG: Print unpacked values after tree_flatten @@ -2634,6 +2759,11 @@ def all_(iterable: Iterable[Union[Numeric, ir.Value, bool]]) -> Boolean: conditions = [a > 0, b < 10, c != 0] result = all_(conditions) # Returns True if all conditions are met """ + from .._mlir_helpers.vector import Vector, all_ as vector_all_ + + if isinstance(iterable, Vector): + return vector_all_(iterable) + bool_iterable = [Boolean(i) for i in iterable] reducer = lambda lhs, rhs: ( lhs.__dsl_and__(rhs) if hasattr(lhs, "__dsl_and__") else lhs @@ -2664,6 +2794,11 @@ def any_(iterable: Iterable[Union[Numeric, ir.Value, bool]]) -> Boolean: conditions = [a > 10, b < 0, c != 0] result = any_(conditions) # Returns True if any condition is met """ + from .._mlir_helpers.vector import Vector, any_ as vector_any_ + + if isinstance(iterable, Vector): + return vector_any_(iterable) + bool_iterable = [Boolean(i) for i in iterable] reducer = lambda lhs, rhs: ( lhs.__dsl_or__(rhs) if hasattr(lhs, "__dsl_or__") else lhs @@ -2688,13 +2823,13 @@ def select_( if len(value) == 1: return value[0] else: - raise DSLRuntimeError( - "Conditional expression must have exactly one value in all expressions" + raise DSLUserCodeError( + DiagId.TYPE_CONDITIONAL_VALUE_COUNT, ) return value if not is_dynamic_expression(cond): - raise DSLRuntimeError("Conditional expression must be dynamic") + raise DSLUserCodeError(DiagId.PHASE_CONDITIONAL_NOT_DYNAMIC) # Extract MLIR values cond = extract_mlir_values(cond) @@ -2759,7 +2894,7 @@ def for_generate( if isinstance(p, int): p = const(p) elif isinstance(p, float): - raise DSLRuntimeError(f"{p=} must be int.") + raise DSLUserCodeError(DiagId.TYPE_LOOP_PARAM_NOT_INT) elif isinstance(p, Integer): p = p.ir_value() params[i] = p @@ -2768,12 +2903,12 @@ def for_generate( def _createI32Attr(value: Union[Int32, int]) -> ir.IntegerAttr: if not isinstance(value, int): - raise DSLRuntimeError("value must be int.") + raise DSLUserCodeError(DiagId.TYPE_PARAM_NOT_INT) return ir.IntegerAttr.get(ir.IntegerType.get_signless(32), value) ir_iter_args = extract_mlir_values(iter_args) if iter_args is not None else None if not _validate_iter_args_structure(iter_args, ir_iter_args): - raise DSLRuntimeError("iter_args: Elements should be extractable as ir.Value.") + raise DSLUserCodeError(DiagId.TYPE_ITER_ARGS_INVALID) for_op = scf.ForOp(start, stop, step, ir_iter_args, loc=loc, ip=ip) if unroll is not None: for_op.attributes["loop_annotation"] = unroll @@ -2862,7 +2997,7 @@ def if_generate( if return_types is not None: for t in return_types: if not isinstance(t, DslType): - raise DSLRuntimeError(f"{t=} must be a DslType.") + raise DSLUserCodeError(DiagId.TYPE_RETURN_TYPE_NOT_DSL) mlir_return_types.append(t.mlir_type) # type: ignore[attr-defined] # Determine whether there's an else branch. @@ -2942,7 +3077,7 @@ class WhileLoopContext: self.input_ir_values = extract_mlir_values(inputs) if not _validate_iter_args_structure(inputs, self.input_ir_values): - raise DSLRuntimeError("inputs: Elements should be extractable as ir.Value.") + raise DSLUserCodeError(DiagId.TYPE_WHILE_INPUTS_INVALID) self.condition = condition self.input_ir_types = [i.type for i in self.input_ir_values] @@ -3038,9 +3173,7 @@ def in_(lhs: object, rhs: Any) -> Union[bool, Boolean]: return lhs in rhs if not isinstance(rhs, Sequence): - raise DSLRuntimeError( - f"'in' not supported between instances of {type(lhs)} and {type(rhs)}" - ) + raise DSLUserCodeError(DiagId.UNSUP_IN_OPERATOR) return any_(equal(lhs, r) for r in rhs) @@ -3070,7 +3203,7 @@ def _lte_gte( else: return not_(lhs < rhs) else: - raise DSLRuntimeError(f"Unsupported comparison operator: {op}") + raise DSLUserCodeError(DiagId.UNSUP_COMPARISON_OPERATOR, op=op) if not is_dynamic_expression(lhs) and not is_dynamic_expression(rhs): return native_lte_gte(lhs, rhs, op) @@ -3186,7 +3319,7 @@ def _compare_dispatch( elif op == "<=": return less_equal(lhs, rhs) else: - raise DSLRuntimeError(f"Unsupported comparison operator: {op}") + raise DSLUserCodeError(DiagId.UNSUP_COMPARISON_OPERATOR, op=op) def _compare_executor( @@ -3214,8 +3347,9 @@ def _builtin_redirector(fcn: Callable[..., object]) -> Callable[..., object]: if is_dynamic_expression(args): if kwargs: # Redirected built-ins do not support keyword arguments - raise DSLRuntimeError( - f"Unsupported keyword arguments for built-in function: {fcn}" + raise DSLUserCodeError( + DiagId.CALL_BUILTIN_KWARGS_UNSUPPORTED, + fcn=fcn, ) if fcn is builtins.max: return max(*args) @@ -3226,7 +3360,7 @@ def _builtin_redirector(fcn: Callable[..., object]) -> Callable[..., object]: elif fcn is builtins.all: return all_(*args) else: - raise DSLRuntimeError(f"Unsupported built-in function: {fcn}") + raise DSLUserCodeError(DiagId.UNSUP_BUILTIN_FUNCTION, fcn=fcn) else: return fcn(*args, **kwargs) diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/cutlass_ast_decorators.py b/python/CuTeDSL/cutlass/cutlass_dsl/cutlass_ast_decorators.py index 6b8087d34..350c7b132 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/cutlass_ast_decorators.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/cutlass_ast_decorators.py @@ -16,7 +16,12 @@ from cutlass._mlir import ir from cutlass._mlir.dialects import scf from collections.abc import Sequence -from ..base_dsl.common import DSLRuntimeError, DSLNotImplemented +from ..base_dsl.common import ( + DSLRuntimeError, + DSLUserCodeError, + DSLNotImplemented, +) +from ..base_dsl.diagnostics import DiagId from ..base_dsl.dsl import is_dynamic_expression from .._mlir_helpers.arith import ArithValue from ..base_dsl.ast_helpers import * # noqa: F401,F403 @@ -166,14 +171,12 @@ class ScfGenerator: new_type_name = new_type.__name__ if type_mismatch: - raise DSLRuntimeError( - f"`{arg_name}` is {old_type_name} prior to this `{op_type_name}`, " - f"and update to {new_type_name} inside of this `{op_type_name}` is not supported.", - suggestion=( - f"Please avoid changing type inside a dynamic `{op_type_name}`, " - f"or change to compile-time control flow by marking this `{op_type_name}` with " - f"`{'range_constexpr' if op_type_name == 'for' else 'const_expr'}`." - ), + raise DSLUserCodeError( + DiagId.TYPE_UNSTABLE_JOIN, + var=arg_name, + old_type=old_type_name, + new_type=new_type_name, + op_type=op_type_name, ) def scf_execute_dynamic( @@ -245,13 +248,10 @@ class ScfGenerator: ) ) - raise DSLRuntimeError( - f"`{filterd_arg_names[mismatch]}` is structured different after this `{op_type_name}`.", - suggestion=( - f"Please avoid changing type structure inside a dynamic `{op_type_name}`, " - f"or change to compile-time control flow by marking this `{op_type_name}` with " - f"`{'range_constexpr' if op_type_name == 'for' else 'const_expr'}`." - ), + raise DSLUserCodeError( + DiagId.CONTAINER_STRUCTURE_CHANGED, + var=filterd_arg_names[mismatch], + op_type=op_type_name, ) scf.YieldOp(region_values) @@ -273,10 +273,11 @@ class ScfGenerator: def _attr_const_check(attr: object, expected_type: type, attr_name: str) -> None: - # Use strict type equality to prevent `bool` being accepted where `int` is required. - if is_dynamic_expression(attr) or type(attr) is not expected_type: - raise DSLRuntimeError( - f"loop attribute `{attr_name}` must be a Python value of type `{expected_type.__name__}`, got `{type(attr).__name__}`." + raw = attr + if is_dynamic_expression(attr) or type(raw) is not expected_type: + raise DSLUserCodeError( + DiagId.PHASE_ASSIGN_PYTHON_TO_TRACKED, + attr_name=attr_name, ) @@ -314,8 +315,10 @@ def _loop_execute_range_dynamic( step_n = as_numeric(step) for name, n in (("start", start_n), ("stop", stop_n), ("step", step_n)): if not n.dtype.is_integer: - raise DSLRuntimeError( - f"dynamic loop `{name}` must be an integer type, got {n.dtype}" + raise DSLUserCodeError( + DiagId.ARG_WRONG_TYPE, + name=name, + dtype=n.dtype, ) # Promote to a common integer type using pairwise type promotion _, _, tmp_dtype = _binary_op_type_promote(start_n, stop_n) @@ -344,8 +347,9 @@ def _loop_execute_range_dynamic( ir.IntegerType.get_signless(32), prefetch_stages ) else: - raise DSLRuntimeError( - f"loop attribute `prefetch_stages` must be non-negative, got `{prefetch_stages}`." + raise DSLUserCodeError( + DiagId.CONFIG_INVALID_VALUE, + value=prefetch_stages, ) log().debug("prefetch_stages attribute: %s", prefetch_stages_attr) @@ -355,8 +359,9 @@ def _loop_execute_range_dynamic( arch = cutlass_dsl.CuTeDSL._get_dsl().get_arch_enum() if arch < Arch.sm_100: - raise DSLRuntimeError( - f"vectorize is supported for sm_100 and above, got {arch}." + raise DSLUserCodeError( + DiagId.UNSUP_ARCH, + arch=arch, ) _attr_const_check(vectorize, bool, "vectorize") vectorize_attr = ir.BoolAttr.get(True) @@ -743,8 +748,8 @@ def _ifexp_execute_dynamic( assert isinstance(then_tree, PyTreeDef) assert isinstance(else_tree, PyTreeDef) if check_tree_equal(then_tree, else_tree) != -1: - raise DSLRuntimeError( - "Then and else blocks of ifexp return different types" + raise DSLUserCodeError( + DiagId.TYPE_CONDITIONAL_BRANCH_MISMATCH, ) # Collect result types for the SCF IfOp diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py b/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py index eaa6cbea8..fb2f445bf 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py @@ -27,8 +27,10 @@ from cutlass.base_dsl.jit_executor import JitExecutor from cutlass.base_dsl.common import ( DSLRuntimeError, DSLCudaRuntimeError, - _get_cuda_error_name_from_code, + DSLUserCodeError, ) +from cutlass.base_dsl.diagnostics import DiagId +from cutlass.base_dsl.runtime import cuda as cuda_helpers import tvm_ffi @@ -41,7 +43,12 @@ class CUDADialectError(DSLCudaRuntimeError): def __init__(self, message: str) -> None: self.raw_tvm_ffi_message = message error_code = CUDADialectError._parse_cuda_dialect_error_code(message) - super().__init__(error_code, _get_cuda_error_name_from_code(error_code)) + super().__init__( + error_code, + cuda_helpers.get_cuda_error_name_from_code( + error_code, cuda_helpers.cudart.cudaError_t + ), + ) def _format_message(self) -> str: message = super()._format_message() @@ -508,6 +515,8 @@ def _get_format_from_object_file_path(object_file_path: str) -> str: class TVMFFIJitCompiledFunctionBase(CudaDialectJitCompiledFunction): """Base class for TVM FFI compiled function.""" + engine: Any + def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) @@ -539,11 +548,7 @@ class TVMFFIJitCompiledFunctionBase(CudaDialectJitCompiledFunction): :param export_only_tvm_ffi_symbols: Only export TVM FFI symbols (hide all others). """ if self.host_target.value: - raise DSLRuntimeError( - "Host cross-compile via TVM-FFI is not supported. " - "Drop --enable-tvm-ffi to use the plain AOT export path " - "with --host-target." - ) + raise DSLUserCodeError(DiagId.CONFIG_INCOMPATIBLE_FLAGS) assert function_name is not None internal_symbol_prefix = "__cute_internal_" + function_name mod = self.ir_module @@ -611,6 +616,18 @@ class TVMFFIJitCompiledFunction(tvm_ffi.Function, TVMFFIJitCompiledFunctionBase) # move the handle from the tvm_ffi.Function to the current instance self.__move_handle_from__(tvm_ffi_function) + def __call__(self, *args: Any) -> Any: + if self.__chandle__() == 0: + raise DSLRuntimeError( + "TVM FFI function is not initialized." + " Was this function compiled for a different architecture?" + ) + if self.execution_args.has_pointer_address_arg_specs: + args, _ = self.execution_args.convert_python_pointer_args_for_tvm_ffi( + args, {} + ) + return tvm_ffi.Function.__call__(self, *args) + class TVMFFIJitCompiledFunctionWithKwargs(TVMFFIJitCompiledFunctionBase): """TVM FFI Function with kwargs wrapper support""" @@ -660,6 +677,11 @@ class TVMFFIJitCompiledFunctionWithKwargs(TVMFFIJitCompiledFunctionBase): "TVM FFI function is not initialized." " Was this function compiled for a different architecture?" ) + + if self.execution_args.has_pointer_address_arg_specs: + args, kwargs = self.execution_args.convert_python_pointer_args_for_tvm_ffi( + args, kwargs + ) return self._kwargs_wrapper(*args, **kwargs) def __tvm_ffi_object__(self) -> Optional["tvm_ffi.Function"]: diff --git a/python/CuTeDSL/cutlass/jax/types.py b/python/CuTeDSL/cutlass/jax/types.py index 953534ecc..693691e5c 100644 --- a/python/CuTeDSL/cutlass/jax/types.py +++ b/python/CuTeDSL/cutlass/jax/types.py @@ -20,8 +20,8 @@ import cutlass import cutlass.cute as cute from cutlass.cute.core import IntValue from cutlass.cute.runtime import from_dlpack as _from_dlpack -from cutlass.cute import AddressSpace from cutlass._mlir import ir +from cutlass.address_space import AddressSpace from cutlass._mlir.dialects import llvm, arith JAX_DTYPE_TO_CUTLASS_DTYPE = { @@ -54,6 +54,10 @@ DEFAULT_CUTLASS_DEVICE_MEMSPACE = AddressSpace.gmem DEFAULT_CUTLASS_DEVICE_BUFFER_ALIGNMENT = 256 +def _llvm_pointer_type(address_space: AddressSpace = AddressSpace.generic) -> ir.Type: + return llvm.PointerType.get(int(address_space)) + + @jax.tree_util.register_dataclass @dataclass(frozen=True) class TensorSpec: @@ -512,7 +516,7 @@ class JaxArrayValue(JaxArray): ) -> tuple[ir.Value, ...]: i64 = ir.IntegerType.get_signless(64) shape_array = llvm.extractvalue( - llvm.PointerType.get(), + _llvm_pointer_type(), ffi_buffer, [1], loc=loc, @@ -522,7 +526,7 @@ class JaxArrayValue(JaxArray): shape_i64 = [] for i in range(len(self.shape)): r = llvm.getelementptr( - llvm.PointerType.get(), + _llvm_pointer_type(), shape_array, [], no_wrap_flags=0, @@ -542,16 +546,20 @@ class JaxArrayValue(JaxArray): loc: ir.Location | None = None, ip: ir.InsertionPoint | None = None, ) -> ir.Value: - raw_ptr = llvm.extractvalue( - llvm.PointerType.get(), + data_ptr = llvm.extractvalue( + _llvm_pointer_type(), ffi_buffer, [0], loc=loc, ip=ip, ) + if self.mem_space != AddressSpace.generic: + data_ptr = llvm.addrspacecast( + _llvm_pointer_type(self.mem_space), data_ptr, loc=loc, ip=ip + ) return cute.make_ptr( self.dtype, - raw_ptr, + data_ptr, self.mem_space, assumed_align=self.assumed_align, loc=loc, @@ -562,7 +570,7 @@ class JaxArrayValue(JaxArray): self, *, loc: ir.Location | None = None, ip: ir.InsertionPoint | None = None ) -> ir.Value: ffi_buffer_type = llvm.StructType.get_literal( - [llvm.PointerType.get(), llvm.PointerType.get()] + [_llvm_pointer_type(), _llvm_pointer_type()] ) ffi_buffer = llvm.load(ffi_buffer_type, self.value, loc=loc, ip=ip) @@ -614,7 +622,7 @@ class JaxTracedArray(JaxArray): def __get_mlir_types__(self) -> list[ir.Type]: # Struct passed as opaque object. - return [llvm.PointerType.get()] + return [_llvm_pointer_type()] def __new_from_mlir_values__(self, values: ir.Value) -> JaxArrayValue: return JaxArrayValue( diff --git a/python/CuTeDSL/cutlass/pipeline/__init__.py b/python/CuTeDSL/cutlass/pipeline/__init__.py index b6b56bd9a..316697118 100644 --- a/python/CuTeDSL/cutlass/pipeline/__init__.py +++ b/python/CuTeDSL/cutlass/pipeline/__init__.py @@ -20,6 +20,7 @@ from .helpers import ( TmaStoreFence, PipelineUserType, PipelineState, + alloc_reserved_mbarrier, make_pipeline_state, pipeline_init_arrive, pipeline_init_wait, diff --git a/python/CuTeDSL/cutlass/pipeline/helpers.py b/python/CuTeDSL/cutlass/pipeline/helpers.py index c174823c8..97ead4a25 100644 --- a/python/CuTeDSL/cutlass/pipeline/helpers.py +++ b/python/CuTeDSL/cutlass/pipeline/helpers.py @@ -20,7 +20,53 @@ import cutlass.cute as cute from cutlass._mlir import ir from cutlass.base_dsl.arch import Arch from cutlass.cute.arch.constants import WARP_SIZE -from cutlass.cutlass_dsl import CuTeDSL, Boolean, Int32, if_generate, dsl_user_op +from cutlass.cutlass_dsl import CuTeDSL, Boolean, Int32, Int64, if_generate, dsl_user_op +from cutlass.utils import SmemAllocator, SmemPartition + + +############################################################################## +# Pipeline util +############################################################################## + + +@dsl_user_op +def alloc_reserved_mbarrier( + num_stages: int, + *, + multiplier: int = 2, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> cute.Pointer: + """Allocate shared memory for pipeline mbarriers at reserved low address in smem. + + Allocates `num_stages * multiplier` elements of `Int64` (8-byte aligned). + Reserved placement keeps this storage at the low address of the smem region + so it remains valid if shared memory is later resized , which can release + memory outside the reserved block. Pipeline classes (e.g. `PipelineAsync`, + `PipelineOrder`) use this when `barrier_storage=None` to create their + mbarrier storage. + + :param num_stages: Number of pipeline stages (barrier slots needed per "side"); + total elements allocated is `num_stages * multiplier`. + :type num_stages: int + :param multiplier: Number of barrier arrays to reserve space for; default 2 + for pipelines that use separate full and empty barriers. + :type multiplier: int + :param loc: Optional source location for emitted IR. + :param ip: Optional insertion point for emitted IR. + :return: Pointer to the allocated shared memory, suitable for use as + `barrier_storage` in pipeline `create()` methods. + :rtype: cute.Pointer + """ + barrier_storage = SmemAllocator().allocate_array( + Int64, + num_stages * multiplier, + byte_alignment=8, + partition=SmemPartition.RESERVED, + loc=loc, + ip=ip, + ) + return barrier_storage ############################################################################## @@ -54,8 +100,20 @@ class CooperativeGroup: """ def __init__( - self, agent: Agent, size: Union[int, Int32] = 1, alignment: Optional[int] = None + self, + agent: Agent, + size: Union[int, Int32] = 1, + alignment: Optional[int] = None, ): + if alignment is not None: + warnings.warn( + "The 'alignment' parameter of CooperativeGroup's constructor is " + "deprecated and will be removed in a subsequent release, please " + "remove it from your code.", + DeprecationWarning, + stacklevel=2, + ) + if agent in [ Agent.Thread, Agent.Warp, @@ -458,10 +516,9 @@ class MbarrierArray(SyncObject): *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None, - ) -> Optional[tuple]: + ) -> None: """ Wait on mbarrier. - uses mbarrier_wait and returns None. :param index: Index of the mbarrier in the array :param phase: Phase/parity to wait for (0 or 1) @@ -470,7 +527,6 @@ class MbarrierArray(SyncObject): cute.arch.mbarrier_wait( self.get_barrier(index, loc=loc, ip=ip), phase, loc=loc, ip=ip ) - return None @dsl_user_op def arrive_and_wait( diff --git a/python/CuTeDSL/cutlass/pipeline/profiling.py b/python/CuTeDSL/cutlass/pipeline/profiling.py index 4aca1fce7..2f8191bc7 100644 --- a/python/CuTeDSL/cutlass/pipeline/profiling.py +++ b/python/CuTeDSL/cutlass/pipeline/profiling.py @@ -55,7 +55,7 @@ def dump_profiling_metadata(dump_dir: str, extra: Optional[dict] = None) -> None :param dump_dir: Output directory. :type dump_dir: str - :param extra: Additional fields to merge. + :param extra: Additional fields to merge (e.g., TS schedule). :type extra: dict, optional :raises RuntimeError: If the output file cannot be written. """ diff --git a/python/CuTeDSL/cutlass/pipeline/sm100.py b/python/CuTeDSL/cutlass/pipeline/sm100.py index 3d1d132f1..d5df4073a 100644 --- a/python/CuTeDSL/cutlass/pipeline/sm100.py +++ b/python/CuTeDSL/cutlass/pipeline/sm100.py @@ -9,6 +9,7 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. +import warnings from dataclasses import dataclass from typing import TYPE_CHECKING, Literal, Optional, cast @@ -19,6 +20,7 @@ import cutlass.cute as cute from cutlass.cute.arch.constants import WARP_SIZE from cutlass.cute.core import is_static from cutlass.cutlass_dsl import BaseDSL, Boolean, Int32, if_generate, dsl_user_op + from cutlass.pipeline import ( Agent, CooperativeGroup, @@ -29,6 +31,7 @@ from cutlass.pipeline import ( TmaStoreFence, PipelineState, PipelineAsync, + alloc_reserved_mbarrier, agent_sync, ) from cutlass.pipeline.helpers import _get_thread_arrive_count @@ -213,6 +216,9 @@ class PipelineTmaUmma(PipelineAsync): :return: A new PipelineTmaUmma instance configured with the provided parameters :rtype: PipelineTmaUmma """ + # Create barrier storage at reserved low address of smem + if barrier_storage is None: + barrier_storage = alloc_reserved_mbarrier(num_stages) if not isinstance(barrier_storage, cute.Pointer): raise TypeError( f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" @@ -489,6 +495,9 @@ class PipelineAsyncUmma(PipelineAsync): :return: A new PipelineAsyncUmma instance configured with the provided parameters :rtype: PipelineAsyncUmma """ + # Create barrier storage at reserved low address of smem + if barrier_storage is None: + barrier_storage = alloc_reserved_mbarrier(num_stages) if not isinstance(barrier_storage, cute.Pointer): raise TypeError( f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" @@ -659,6 +668,9 @@ class PipelineUmmaAsync(PipelineAsync): :return: New instance of ``PipelineUmmaAsync`` :rtype: PipelineUmmaAsync """ + # Create barrier storage at reserved low address of smem + if barrier_storage is None: + barrier_storage = alloc_reserved_mbarrier(num_stages) if not isinstance(barrier_storage, cute.Pointer): raise TypeError( f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" @@ -775,7 +787,7 @@ class PipelineUmmaAsync(PipelineAsync): if is_leader_cta: # Assume state contains that next useful buffer # So we only need to advance to num_stages - 1 times to last used buffer - for i in cutlass.range_constexpr(self.num_stages - 1): # type: ignore[func-returns-value] + for i in cutlass.range_constexpr(self.num_stages - 1): state.advance(loc=loc, ip=ip) self.producer_acquire(state, loc=loc, ip=ip) @@ -849,6 +861,9 @@ class PipelineClcFetchAsync(PipelineAsync): :param consumer_mask: Mask for signaling arrives for the consumer agent, defaults to ``None`` :type consumer_mask: Int32, optional """ + # Create barrier storage at reserved low address of smem + if barrier_storage is None: + barrier_storage = alloc_reserved_mbarrier(num_stages) if not isinstance(barrier_storage, cute.Pointer): raise TypeError( f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" @@ -1027,6 +1042,7 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync): tidx: Optional[Int32] = None, enable_multicast_signaling: bool = False, defer_sync: bool = False, + force_deprecated_per_lane_signaling: Optional[bool] = None, name: str = "", ) -> "PipelineTmaMultiConsumersAsync": """ @@ -1051,7 +1067,12 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync): :type tidx: Int32 | None :param enable_multicast_signaling: See docstring in PipelineTmaAsync.create() for details :type enable_multicast_signaling: bool, optional + :param force_deprecated_per_lane_signaling: **Deprecated.** Set ``False`` if your arrive count is a multiple of ``WARP_SIZE`` and you do not want the legacy fallback. Leave unset otherwise. + :type force_deprecated_per_lane_signaling: bool | None """ + # Create barrier storage at reserved low address of smem + if barrier_storage is None: + barrier_storage = alloc_reserved_mbarrier(num_stages) if not isinstance(barrier_storage, cute.Pointer): raise TypeError( f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" @@ -1073,6 +1094,11 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync): "UMMA and AsyncThread consumer groups must be the same agent" ) + if enable_multicast_signaling and force_deprecated_per_lane_signaling: + raise ValueError( + "enable_multicast_signaling unsupported with force_deprecated_per_lane_signaling" + ) + if enable_multicast_signaling: consumer_thread_arrive_cnt_umma = _get_thread_arrive_count( consumer_group_umma @@ -1119,6 +1145,47 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync): thread_consumer_group_umma = consumer_group_umma thread_consumer_group_async = consumer_group_async + def _looks_legacy(cg: CooperativeGroup) -> bool: + return ( + cg.agent is Agent.Thread + and cg.size >= WARP_SIZE + and cg.size % WARP_SIZE == 0 + ) + + if force_deprecated_per_lane_signaling is None: + force_deprecated_per_lane_signaling = not enable_multicast_signaling and ( + _looks_legacy(thread_consumer_group_umma) + or _looks_legacy(thread_consumer_group_async) + ) + if force_deprecated_per_lane_signaling: + warnings.warn( + "PipelineTmaMultiConsumersAsync now expects only one lane " + "per warp to arrive on the empty barrier, but consumer_group " + "arrive count is a multiple of WARP_SIZE - this suggests you are expecting the" + "legacy behavior where every lane arrives. Falling back to " + "the legacy path. Divide the arrive count by WARP_SIZE to " + "migrate, or pass force_deprecated_per_lane_signaling=False to opt out of " + "this auto-detection.", + DeprecationWarning, + stacklevel=2, + ) + elif force_deprecated_per_lane_signaling: + warnings.warn( + "force_deprecated_per_lane_signaling=True is deprecated.", + DeprecationWarning, + stacklevel=2, + ) + + if force_deprecated_per_lane_signaling: + cluster_size = ( + cute.size(cta_layout_vmnk) if cta_layout_vmnk is not None else 1 + ) + if cluster_size > 1: + raise ValueError( + "force_deprecated_per_lane_signaling does not support multicast " + "(cluster_size > 1)." + ) + consumer_group = CooperativeGroup( thread_consumer_group_umma.agent, thread_consumer_group_umma.size + thread_consumer_group_async.size, @@ -1153,14 +1220,18 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync): ) # Compute AsyncThread consumer signaling (dst_rank + is_signaling_thread) - if tidx is None: - tidx, _, _ = cute.arch.thread_idx() - ( - dst_rank_async, - is_signaling_thread, - ) = PipelineTmaMultiConsumersAsync._init_empty_barrier_arrive_signal_2sm( - cta_layout_vmnk, tidx, mcast_mode_mn - ) + if force_deprecated_per_lane_signaling: + dst_rank_async = None + is_signaling_thread = Boolean(True) + else: + if tidx is None: + tidx, _, _ = cute.arch.thread_idx() + ( + dst_rank_async, + is_signaling_thread, + ) = PipelineTmaMultiConsumersAsync._init_empty_barrier_arrive_signal_2sm( + cta_layout_vmnk, tidx, mcast_mode_mn + ) if cute.size(cta_layout_vmnk) == 1: # No mcast mask if not using clusters @@ -1300,10 +1371,11 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync): ip: Optional[ir.InsertionPoint] = None, ) -> None: """Consumer waits for full barrier to be signaled.""" - _wait_fn = self.sync_object_full.wait if_generate( try_wait_token is None or try_wait_token == 0, - lambda: _wait_fn(state.index, state.phase, loc=loc, ip=ip), + lambda: self.sync_object_full.wait( + state.index, state.phase, loc=loc, ip=ip + ), loc=loc, ip=ip, ) @@ -1317,8 +1389,7 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync): ip: Optional[ir.InsertionPoint] = None, ) -> Boolean: """Non-blocking check if data is ready.""" - _try_wait_fn = self.sync_object_full.try_wait # type: ignore[attr-defined] - return _try_wait_fn(state.index, state.phase, loc=loc, ip=ip) + return self.sync_object_full.try_wait(state.index, state.phase, loc=loc, ip=ip) # type: ignore[attr-defined] @dsl_user_op def consumer_release( @@ -1334,8 +1405,13 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync): state.index, self.consumer_mask, self.cta_group, loc=loc, ip=ip ) elif op_type == PipelineOp.AsyncThread: - self.sync_object_empty_async.arrive( - state.index, self.consumer_mask, loc=loc, ip=ip + if_generate( + self.is_signaling_thread, + lambda: self.sync_object_empty_async.arrive( + state.index, self.consumer_dst_rank_async, loc=loc, ip=ip + ), + loc=loc, + ip=ip, ) else: raise ValueError(f"Invalid PipelineOp specified. op_type:{op_type}") diff --git a/python/CuTeDSL/cutlass/pipeline/sm90.py b/python/CuTeDSL/cutlass/pipeline/sm90.py index 774af8e28..548b22c47 100644 --- a/python/CuTeDSL/cutlass/pipeline/sm90.py +++ b/python/CuTeDSL/cutlass/pipeline/sm90.py @@ -29,6 +29,7 @@ from cutlass.pipeline import ( PipelineUserType, SyncObject, TmaStoreFence, + alloc_reserved_mbarrier, make_pipeline_state, agent_sync, ) @@ -193,6 +194,9 @@ class PipelineAsync: :return: A new ``PipelineAsync`` instance :rtype: PipelineAsync """ + # Create barrier storage at reserved low address of smem + if barrier_storage is None: + barrier_storage = alloc_reserved_mbarrier(num_stages) if not isinstance(barrier_storage, cute.Pointer): raise TypeError( f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" @@ -394,6 +398,7 @@ class PipelineCpAsync(PipelineAsync): @staticmethod def create( # type: ignore[override] + *, barrier_storage: cute.Pointer, num_stages: Int32, producer_group: CooperativeGroup, @@ -420,6 +425,9 @@ class PipelineCpAsync(PipelineAsync): :return: A new ``PipelineCpAsync`` instance configured with the provided parameters :rtype: PipelineCpAsync """ + # Create barrier storage at reserved low address of smem + if barrier_storage is None: + barrier_storage = alloc_reserved_mbarrier(num_stages) producer_type = PipelineOp.AsyncLoad consumer_type = PipelineOp.AsyncThread @@ -571,6 +579,9 @@ class PipelineTmaAsync(PipelineAsync): :return: New ``PipelineTmaAsync`` instance :rtype: PipelineTmaAsync """ + # Create barrier storage at reserved low address of smem + if barrier_storage is None: + barrier_storage = alloc_reserved_mbarrier(num_stages) if not isinstance(barrier_storage, cute.Pointer): raise TypeError( f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" @@ -852,14 +863,18 @@ class PipelineOrder: @staticmethod def create( - barrier_storage: cute.Pointer, + *, depth: int, length: int, group_id: int, producer_group: CooperativeGroup, + barrier_storage: Optional[cute.Pointer] = None, defer_sync: bool = False, name: str = "", ) -> "PipelineOrder": + # Create barrier storage at reserved low address of smem + if barrier_storage is None: + barrier_storage = alloc_reserved_mbarrier(depth * length) if not isinstance(barrier_storage, cute.Pointer): raise TypeError( f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" diff --git a/python/CuTeDSL/cutlass/runtime.py b/python/CuTeDSL/cutlass/runtime.py new file mode 100644 index 000000000..e6b927aea --- /dev/null +++ b/python/CuTeDSL/cutlass/runtime.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +import ctypes +import sys +from pathlib import Path +from typing import Any + +from cutlass._mlir import ir +import cutlass._mlir.dialects.cuda as _cuda_dialect +from cutlass.base_dsl.env_manager import get_prefix_dsl_libs +from cutlass.base_dsl.export import ExternalBinaryModule + + +class _FakeStream: + """A fake stream that can be used as a placeholder for a stream in compilation. + + When use_tvm_ffi_env_stream is True and the function is compiled with TVM-FFI, + the argument will be skipped from the function signature and we pass in + this value through the environment stream obtained from caller context + (e.g. torch.cuda.current_stream()). + """ + + use_tvm_ffi_env_stream: bool + + def __init__(self, *, use_tvm_ffi_env_stream: bool = False) -> None: + self.use_tvm_ffi_env_stream = use_tvm_ffi_env_stream + + def __str__(self) -> str: + return "FakeStream" + + def __repr__(self) -> str: + return self.__str__() + + def __new_from_mlir_values__(self, values: list[object]) -> object: + assert len(values) == 1 + return values[0] + + def __c_pointers__(self) -> list[int]: + return [0] + + def __get_mlir_types__(self) -> list[ir.Type]: + return [_cuda_dialect.StreamType.get()] + + +def make_fake_stream(*, use_tvm_ffi_env_stream: bool = False) -> _FakeStream: + """Create a fake stream that can be used as a placeholder for a stream in compilation. + + When use_tvm_ffi_env_stream is True and the function is compiled with TVM-FFI, + the argument will be skipped from the function signature and we pass in + this value through the environment stream obtained from caller context + (e.g. torch.cuda.current_stream()). This can speedup the calling process + since we no longer need to do stream query in python. + + :param use_tvm_ffi_env_stream: Whether to skip this parameter use environment stream instead. + :type use_tvm_ffi_env_stream: bool + """ + return _FakeStream(use_tvm_ffi_env_stream=use_tvm_ffi_env_stream) + + +def find_runtime_libraries(*, enable_tvm_ffi: bool = True) -> list[str]: + """ + Find the runtime libraries that needs to be available for loading modules. + + :param enable_tvm_ffi: Whether to enable TVM-FFI. + :type enable_tvm_ffi: bool, optional + :return: A list of runtime libraries that needs to be available for loading modules. + :rtype: list + """ + + def _get_cute_dsl_runtime_path() -> str | None: + libs = get_prefix_dsl_libs("CUTE_DSL") + if libs is None: + return None + + # check if the separator is ; for windows + if sys.platform.startswith("win32") and ";" in libs: + libs = libs.split(";") # type: ignore[assignment] + else: + libs = libs.split(":") # type: ignore[assignment] + + for path in libs: + if path.endswith("libcute_dsl_runtime.so"): + return path + + return None + + libs = [] + cute_dsl_runtime_path = _get_cute_dsl_runtime_path() + if cute_dsl_runtime_path: + libs.append(cute_dsl_runtime_path) + + if enable_tvm_ffi: + import tvm_ffi + + libs.append(tvm_ffi.libinfo.find_libtvm_ffi()) + + return libs + +# cache to load runtime libraries so they can be found by the DSO loader +_LOAD_MODULE_LIBS_CACHE: list[Any] = [] + + +def load_module( + file_path: str, *, enable_tvm_ffi: bool = False +) -> ExternalBinaryModule: + """Load a module from a file path. + + :param file_path: The path to the module file + :type file_path: str + :param enable_tvm_ffi: Whether to enable TVM-FFI, defaults to False. When True, the module will be loaded as a TVM-FFI module. + :type enable_tvm_ffi: bool, optional + :return: A module object + :rtype: module + """ + if len(_LOAD_MODULE_LIBS_CACHE) == 0: + # ensure the runtime libraries are loaded so they can be found by the DSO loader + # no need to load tvm_ffi library here since it will be loaded by tvm_ffi package. + for path in find_runtime_libraries(enable_tvm_ffi=False): + if Path(path).exists(): + _LOAD_MODULE_LIBS_CACHE.append(ctypes.CDLL(path)) + + return ExternalBinaryModule(file_path, enable_tvm_ffi=enable_tvm_ffi) diff --git a/python/CuTeDSL/cutlass/testing.py b/python/CuTeDSL/cutlass/testing.py index 21e27e741..0a2bd7f75 100644 --- a/python/CuTeDSL/cutlass/testing.py +++ b/python/CuTeDSL/cutlass/testing.py @@ -237,7 +237,11 @@ def _does_kernel_use_stream( ) -> bool: """ This function checks if the kernel uses the provided non-default stream. - It does this by capturing the stream and then checking if any kernels were launched. + It does this by capturing the stream and then checking whether any work was + recorded on it. A kernel that launches on a different stream records no nodes; + on more recent drivers such a foreign-stream launch additionally fails + fast during capture. Both outcomes are reported as "does not use the stream" + (return False) rather than propagated, so callers can raise a clear error. Note: the function accepts positional/keyword arguments for the kernel in non-unpacked form (as tuple/dict, respectively) to avoid name clashes with function's own arguments (e.g. stream). @@ -263,24 +267,27 @@ def _does_kernel_use_stream( ) _cuda_success(err, "Error on stream capture") + launch_error = None try: kernel(*args, **kwargs) - except Exception: - # Always end the capture even on failure to avoid zombie capture state - # that would poison all subsequent graph capture operations in the process. - try: - cuda_runtime.cudaStreamEndCapture(stream) - except Exception: - pass - raise + except Exception as exc: + launch_error = exc err, graph = cuda_runtime.cudaStreamEndCapture(stream) - _cuda_success(err, "Error on stream capture") + if err != cuda_runtime.cudaError_t.cudaSuccess: + if launch_error is not None: + raise launch_error + _cuda_success(err, "Error on stream capture") - # Get number of nodes in warmup graph to check it matches what is expected err, _, num_nodes = cuda_runtime.cudaGraphGetNodes(graph) _cuda_success(err, "Error on querying graph") - return num_nodes > 0 + + if num_nodes > 0: + if launch_error is not None: + raise launch_error + return True + + return False def benchmark( @@ -391,6 +398,17 @@ def benchmark( "workspace_generator and/or kernel_arguments should use JitArguments type" ) + if ( + not use_cuda_graphs + and int(stream) != int(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT) + and not _does_kernel_use_stream( + callable, stream, workspaces[0].args, workspaces[0].kwargs + ) + ): + raise ValueError( + "CUDA stream passed to benchmark does not match the stream the kernel was launched in" + ) + # use memset to flush L2 cache after workspace h2d copies if workspace_count > 1: from cutlass.utils import HardwareInfo @@ -469,15 +487,6 @@ def benchmark( f"kernel_launcher must be callable, got {type(kernel_launcher).__name__}" ) - if int(stream) != int( - cuda_driver.CUstream_flags.CU_STREAM_DEFAULT - ) and not _does_kernel_use_stream( - callable, stream, workspaces[0].args, workspaces[0].kwargs - ): - raise ValueError( - "CUDA stream passed to benchmark does not match the stream the kernel was launched in" - ) - err = cuda_driver.cuEventRecord(start_event, stream) _cuda_success(err, "Error on recording start event") diff --git a/python/CuTeDSL/cutlass/utils/README.md b/python/CuTeDSL/cutlass/utils/README.md new file mode 100644 index 000000000..2773729c2 --- /dev/null +++ b/python/CuTeDSL/cutlass/utils/README.md @@ -0,0 +1,637 @@ +# Utilities for Kernel Authoring + +> **Experimental Notice**: The utilities in this folder are experimental and subject to breaking changes. APIs may evolve as we refine the design based on user feedback and new hardware features. + +## Table of Contents + +- [Overview](#overview) +- [Prerequisites](#prerequisites) +- [Quick Start](#quick-start) +- [Block-Level Copy Utilities](#block-level-copy-utilities-blockpy) + - [Block Copy Overview](#block-copy-overview) + - [`block_copy`](#block_copy) + - [TMA Copy](#tma-copy) + - [Compiler-Driven TMA Multicast IR rewriting](#compiler-driven-tma-multicast-ir-rewriting) + - [S2T Copy (SMEM to TMEM)](#s2t-copy-smem-to-tmem) +- [TmaInfo Container](#tmainfo-container) +- [Tensor Memory Allocator](#tensor-memory-allocator-tmem_allocatorpy) + - [TMEM Allocator Overview](#tmem-allocator-overview) + - [`TmemAllocator`](#tmemallocator) + - [`TmemBufferPool`](#tmembufferpool) + - [Utility Functions](#utility-functions) +- [Other Utilities](#other-utilities) +- [When to Use What](#when-to-use-what) +- [Examples](#examples) +- [Troubleshooting](#troubleshooting) +- [Limitations](#limitations) + +--- + +## Overview + +This folder provides high-level building blocks that simplify common patterns in GPU kernel development with CuTe DSL. These utilities handle boilerplate code for you, letting you focus on your kernel's core logic. + +**Key utilities include:** + +| Utility | File | Purpose | +|---------|------|---------| +| Block-level copy | [`block.py`](./block.py) | Simplified tile copies (TMA, S2T) without manual partitioning | +| TMEM allocator | [`tmem_allocator.py`](./tmem_allocator.py) | Tensor memory management for Blackwell GPUs | +| Static tile scheduler | [`static_persistent_tile_scheduler.py`](./static_persistent_tile_scheduler.py) | Persistent kernel scheduling with static work distribution | +| Dynamic tile scheduler | [`dynamic_persistent_tile_scheduler.py`](./dynamic_persistent_tile_scheduler.py) | Persistent kernel scheduling with dynamic work stealing | +| Grouped GEMM helpers | [`grouped_gemm_tile_scheduler_helper.py`](./grouped_gemm_tile_scheduler_helper.py), [`tensormap_manager.py`](./tensormap_manager.py) | Utilities for batched/grouped GEMM kernels | + +--- + +## Prerequisites + +### Hardware Requirements + +- **Hopper GPUs (SM90)**: TMA-based `block_copy` operations +- **Blackwell GPUs (SM100+)**: Full support including TMEM allocation and S2T copies + +### Software Requirements + +```python +import cutlass +import cutlass.cute as cute +from cutlass import utils, pipeline +from cutlass.cute.nvgpu import cpasync, tcgen05 +``` + +--- + +## Quick Start + +Here's a minimal example showing how to use `block_copy` for TMA loads in a GEMM mainloop: + +```python +import cutlass +import cutlass.cute as cute +from cutlass import utils + +# Inside your kernel function... + +# 1. Create TMA atoms for A and B operands (configured during kernel setup) +# tma_atom_a, tma_atom_b = ... # Your TMA copy atoms + +# 2. Get your tiles from global and shared memory +# gA, sA = ... # Global and shared memory tensors for A +# gB, sB = ... # Global and shared memory tensors for B + +# 3. Load tiles using block_copy (much simpler than manual partitioning!) +for k_tile_idx in range(num_k_tiles): + # Acquire pipeline stage + pipeline_state = mainloop_pipeline.producer_acquire(producer_state) + + # TMA load A tile: group_modes makes the first mode represent the TMATile + utils.block_copy( + tma_atom_a, + cute.group_modes(gA[(None, None, k_tile_idx)], 0, 2), + cute.group_modes(sA[(None, None, producer_state.index)], 0, 2), + tma_bar_ptr=pipeline_state.barrier, + ) + + # TMA load B tile + utils.block_copy( + tma_atom_b, + cute.group_modes(gB[(None, None, k_tile_idx)], 0, 2), + cute.group_modes(sB[(None, None, producer_state.index)], 0, 2), + tma_bar_ptr=pipeline_state.barrier, + ) + + mainloop_pipeline.producer_commit(producer_state) + producer_state.advance() +``` + +For complete working examples, see the [Examples](#examples) section. + +--- + +## Block-Level Copy Utilities (`block.py`) + +### Block Copy Overview + +The block API provides a simplified interface for tile-level copy operations. Instead of manually: +1. Computing CTA coordinates +2. Calling partitioning functions +3. Handling copy masks +4. Invoking `cute.copy` + +You can use `block_copy` to perform all of this in a single call. + +**Benefits:** +- Works directly on tile-shaped layouts (e.g., `(TILE_M, TILE_K)` for A, `(TILE_N, TILE_K)` for B) +- Integrates with TMA pipelines and cluster multicast on Hopper/Blackwell +- Handles compiler-driven IR rewriting for TMA multicast automatically + +### `block_copy` + +```python +def block_copy( + tiled_copy: TiledCopy, + src: Tensor, + dst: Tensor, + *, + tma_bar_ptr=None, # Required for TMA loads (G2S) + tma_multicast=None, # Optional: enables compiler-driven multicast + **kwargs, +) -> None +``` + +#### Parameters + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `tiled_copy` | `TiledCopy` | Yes | The copy atom configured for your use case | +| `src` | `Tensor` | Yes | Source tensor (must have TMATile as first mode for TMA) | +| `dst` | `Tensor` | Yes | Destination tensor (must have TMATile as first mode for TMA) | +| `tma_bar_ptr` | `Pointer` | For G2S | TMA barrier pointer (required for global-to-shared loads) | +| `tma_multicast` | `dict` | No | Configuration for compiler-driven TMA multicast | + +#### What `block_copy` Handles For You + +**For TMA copies:** +- Calls `tma_partition()` with default CTA coordinates +- Applies `filter_zeros()` to remove layout padding +- Determines copy direction from the copy op type +- Passes through `tma_multicast` and `tma_bar_ptr` arguments + +**For S2T copies:** +- Filters zeros from source and destination tensors +- Gets thread copy slice (using thread 0) +- Partitions source and destination +- Creates SMEM descriptor tensor automatically + +#### Supported Copy Kinds + +| Copy Kind | Copy Op | Description | +|-----------|---------|-------------| +| TMA G2S | `CopyBulkTensorTileG2SOp` | Global to shared memory via TMA | +| TMA G2S Multicast | `CopyBulkTensorTileG2SMulticastOp` | G2S with cluster multicast | +| TMA S2G | TMA store ops | Shared to global memory via TMA | +| S2T | `tcgen05.Cp*Op` | Shared memory to tensor memory (Blackwell) | + +### TMA Copy + +For TMA-based copies, tensors must have their **first mode representing the TMATile**. Use `group_modes` to prepare your tensors: + +```python +# Given a rank-2 tile tensor with layout (TILE_M, TILE_N) +src_grouped = cute.group_modes(src_tile, 0, 2) # Now has layout (TMATile,) +dst_grouped = cute.group_modes(dst_tile, 0, 2) # Now has layout (TMATile,) + +# TMA load (Global → Shared) - requires barrier pointer +utils.block_copy( + tma_atom, + src_grouped, + dst_grouped, + tma_bar_ptr=barrier_ptr, # Required! +) + +# TMA store (Shared → Global) - no barrier needed +utils.block_copy( + tma_atom_store, + smem_grouped, + gmem_grouped, +) +``` + +### Compiler-Driven TMA Multicast IR rewriting + +For kernels using 2D clusters (multiple CTAs cooperating), you can enable **compiler-driven TMA multicast** to broadcast data efficiently across the cluster. The compiler automatically: +- Updates the CopyOp to its multicast version +- Calculates the multicast mask +- Partitions tensors considering multicast effects + +#### `tma_multicast` Configuration + +| Key | Type | Required | Description | +|-----|------|----------|-------------| +| `cluster_shape` | `tuple[int, int]` | Yes | 2D cluster shape `(cluster_m, cluster_n)` | +| `multicast_dim` | `"M"` or `"N"` | Yes | Which dimension to multicast along | +| `use_2cta_mma_inst` | `bool` | No | Use 2CTA MMA instructions (default: `False`) | + +**Multicast dimension explained:** +- `"M"`: CTAs along the M dimension share this data (use for **A operand** in GEMM, which has shape `(M, K)`) +- `"N"`: CTAs along the N dimension share this data (use for **B operand** in GEMM, which has shape `(N, K)`) + +#### Example: TMA Multicast with 2x1 Cluster + +```python +cluster_shape_m, cluster_shape_n = 2, 1 + +# tma_atom_a is obtained from TMA descriptor setup (e.g., tma_info.atom) +# group_modes(tensor, 0, 2) collapses modes [0,1] into a single TMATile mode +# This transforms shape (TILE_M, TILE_K) into (TMATile,) + +# Load A with multicast along M dimension +# A has shape (M, K), so CTAs along M share the same A tile +utils.block_copy( + tma_atom_a, + cute.group_modes(gA[(None, None, k_tile_idx)], 0, 2), # Slice then group + cute.group_modes(sA[(None, None, stage_idx)], 0, 2), + tma_multicast={ + "cluster_shape": (cluster_shape_m, cluster_shape_n), + "multicast_dim": "M", # A operand uses M dimension + "use_2cta_mma_inst": True, + }, + tma_bar_ptr=barrier_ptr, +) + +# Load B with multicast along N dimension +# B has shape (N, K), so CTAs along N share the same B tile +utils.block_copy( + tma_atom_b, + cute.group_modes(gB[(None, None, k_tile_idx)], 0, 2), + cute.group_modes(sB[(None, None, stage_idx)], 0, 2), + tma_multicast={ + "cluster_shape": (cluster_shape_m, cluster_shape_n), + "multicast_dim": "N", # B operand uses N dimension + "use_2cta_mma_inst": True, + }, + tma_bar_ptr=barrier_ptr, +) +``` + +> **Important**: When using `tma_multicast`, you still need to configure the pipeline's arrive count correctly to account for the reduced number of TMA producers. + +### S2T Copy (SMEM to TMEM) + +On Blackwell GPUs, you can copy data from shared memory to tensor memory (TMEM) using S2T copy atoms: + +```python +# S2T copy (SMEM to TMEM) - Blackwell only +# Note: Unlike TMA copies, S2T does NOT require group_modes + +# Create S2T copy atom with appropriate operation and CTA group +copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE), # CtaGroup.TWO for 2-CTA MMA + sf_dtype, +) + +# Copy directly without group_modes +utils.block_copy(copy_atom_s2t, smem_tensor, tmem_tensor) +``` + +The function automatically handles filtering, partitioning, and SMEM descriptor creation. + +--- + +## TmaInfo Container + +The `TmaInfo` class is a container that bundles a TMA Copy Atom with its associated SMEM layout and TMA tensor. It is returned by the TMA atom creation functions: + +- `make_tiled_tma_atom()` - General-purpose TMA atom creation +- `make_tiled_tma_atom_A()` - TMA atom for A operand (MK projections) in tcgen05 MMAs +- `make_tiled_tma_atom_B()` - TMA atom for B operand (NK projections) in tcgen05 MMAs + +### Properties + +| Property | Type | Description | +|----------|------|-------------| +| `atom` | `CopyAtom` | The TMA Copy Atom for use with copy operations | +| `tma_tensor` | `Tensor` | TMA tensor that maps logical GMEM coordinates to TMA-consumable coordinates | +| `smem_layout` | `Layout` | The SMEM layout used to construct the TMA descriptor can be either staged or non-staged | + +### Backward Compatibility + +`TmaInfo` supports tuple unpacking for backward compatibility with code that expects a `(atom, tma_tensor)` tuple: + +```python +# Legacy style: tuple unpacking +atom, tma_tensor = make_tiled_tma_atom(op, gmem_tensor, smem_layout, cta_tiler) + +# New style: access via container (recommended) +tma_info = make_tiled_tma_atom(op, gmem_tensor, smem_layout_staged, cta_tiler) +atom = tma_info.atom +tma_tensor = tma_info.tma_tensor +smem_layout_staged = tma_info.smem_layout # Only available via container +``` + +### Usage Example + +```python +from cutlass.cute.nvgpu import cpasync + +# Create TMA atom using make_tiled_tma_atom +tma_info = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + gmem_tensor, + smem_layout_staged, + cta_tiler, +) + +# Access the Copy Atom for use with block_copy +utils.block_copy( + tma_info.atom, + cute.group_modes(gA[(None, None, k_tile_idx)], 0, 2), + cute.group_modes(sA[(None, None, stage_idx)], 0, 2), + tma_bar_ptr=barrier_ptr, +) + +# Access the SMEM layout if needed (e.g., for verification or debugging) +print(f"SMEM layout: {tma_info.smem_layout}") +``` + +--- + +## Tensor Memory Allocator (`tmem_allocator.py`) + +### TMEM Allocator Overview + +Blackwell GPUs (SM100+) introduce **Tensor Memory (TMEM)**, a specialized memory for accumulator storage in matrix operations. The TMEM allocator utilities help you: + +1. **Allocate** TMEM columns from the allocator warp +2. **Synchronize** across warps before using the allocated memory +3. **Sub-allocate** regions for different tensors without manual offset calculations +4. **Deallocate** TMEM when done (with 2-CTA synchronization support) + +### `TmemAllocator` + +The main class for managing TMEM allocation and deallocation. + +#### Constructor + +```python +TmemAllocator( + alloc_result_dst_smem_ptr: cute.Pointer, # SMEM location to store allocated address + barrier_for_retrieve: pipeline.NamedBarrier, # Barrier for synchronization + allocator_warp_id: int = 0, # Which warp performs allocation + is_two_cta: bool = False, # Enable 2-CTA coordination + num_allocated_columns: int = 0, # Initial allocation count + two_cta_tmem_dealloc_mbar_ptr: cute.Pointer = None, # Required if is_two_cta=True + arch: str = "sm_100", # Target architecture +) +``` + +#### Basic Usage Pattern + +```python +from cutlass import utils, pipeline + +# 1. Setup: allocate SMEM for holding the TMEM address +smem = cutlass.utils.SmemAllocator() +storage = smem.allocate(SharedStorage) # SharedStorage has tmem_holding_buf field + +# 2. Create barrier for synchronization +tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=threads_per_cta, +) + +# 3. Create allocator +tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, +) + +# 4. Allocate TMEM columns (must be power of 2, multiple of 32) +num_tmem_cols = 512 +tmem.allocate(num_tmem_cols) + +# 5. Wait for allocation to complete +tmem.wait_for_alloc() + +# 6. Retrieve pointer (all warps can call this) +tmem_ptr = tmem.retrieve_ptr(dtype=cutlass.Float32) + +# 7. Create tensor with the pointer +tCtAcc = cute.make_tensor(tmem_ptr, acc_layout) + +# ... use tCtAcc in your kernel ... + +# 8. Deallocate when done +tmem.free(tmem_ptr) +``` + +#### Key Methods + +| Method | Description | +|--------|-------------| +| `allocate(num_columns)` | Allocate TMEM columns (called by allocator warp) | +| `wait_for_alloc()` | Synchronize all warps after allocation | +| `retrieve_ptr(dtype)` | Get pointer to allocated TMEM | +| `reserve(num_columns)` | Allocate and return a `TmemBufferPool` for sub-allocation | +| `relinquish_alloc_permit()` | Signal that no more allocations will be made | +| `free(tmem_ptr, num_columns)` | Deallocate TMEM | + +#### Allocation Constraints + +- Number of columns must be a **power of 2** (32, 64, 128, 256, 512, ...) +- Number of columns must be a **multiple of 32** +- Maximum depends on architecture (typically 512 for SM100) + +### `TmemBufferPool` + +For kernels that need multiple TMEM tensors (e.g., accumulator + scale factors), `TmemBufferPool` simplifies sub-allocation from a reserved TMEM region. + +#### Usage + +```python +# Reserve a large TMEM region and get a pool for sub-allocation +tmem_total_size = 512 # Total columns needed +tmem_pool = tmem_allocator.reserve(tmem_total_size) + +# Sub-allocate tensors - no manual offset calculations needed! +tCtAcc = tmem_pool.allocate_tensor(acc_layout, cutlass.Float32) +tCtSFA = tmem_pool.allocate_tensor(scale_a_layout, sf_dtype) +tCtSFB = tmem_pool.allocate_tensor(scale_b_layout, sf_dtype) + +# Or get just a pointer if you need custom tensor creation +sfb_ptr = tmem_pool.allocate(scale_b_layout, sf_dtype) +tCtSFB = cute.make_tensor(sfb_ptr, scale_b_layout) +``` + +#### Properties + +| Property | Description | +|----------|-------------| +| `base_ptr` | Base pointer of the pool | +| `total_cols` | Total columns in the pool | +| `current_offset` | Current allocation offset | +| `remaining_cols` | Columns still available | + +### Utility Functions + +#### `get_num_tmem_alloc_cols` + +Calculate the total TMEM columns needed for given tensors: + +```python +from cutlass.utils.tmem_allocator import get_num_tmem_alloc_cols + +# Calculate required columns (rounded to power of 2) +num_cols = get_num_tmem_alloc_cols( + [tCtAcc, tCtSFA, tCtSFB], + rounding=True, + arch="sm_100", +) +``` + +#### `compute_tmem_cols_from_layout` + +Compute TMEM columns required for a layout (without creating a tensor): + +```python +from cutlass.utils.tmem_allocator import compute_tmem_cols_from_layout + +num_cols = compute_tmem_cols_from_layout(acc_layout, cutlass.Float32) +``` + +--- + +## Other Utilities + +| File | Description | +|------|-------------| +| [`static_persistent_tile_scheduler.py`](./static_persistent_tile_scheduler.py) | Static work distribution for persistent kernels | +| [`dynamic_persistent_tile_scheduler.py`](./dynamic_persistent_tile_scheduler.py) | Dynamic work stealing for load-balanced scheduling | +| [`grouped_gemm_tile_scheduler_helper.py`](./grouped_gemm_tile_scheduler_helper.py) | Helpers for grouped/batched GEMM | +| [`tensormap_manager.py`](./tensormap_manager.py) | TMA descriptor management for grouped GEMM | +| [`smem_allocator.py`](./smem_allocator.py) | Shared memory allocation utilities | +| [`blackwell_helpers.py`](./blackwell_helpers.py) | Blackwell-specific helper functions | +| [`hopper_helpers.py`](./hopper_helpers.py) | Hopper-specific helper functions | + +--- + +## When to Use What + +### Block Copy Decision Matrix + +| Your Situation | Recommendation | +|----------------|----------------| +| Standard tile-level TMA load/store | Use `block_copy` - handles partitioning automatically | +| TMA with cluster multicast | Use `block_copy` + `tma_multicast` - compiler handles mask calculation | +| SMEM to TMEM copy on Blackwell | Use `block_copy` with S2T atom - handles partitioning and SMEM desc creation automatically | +| Custom CTA layouts (non-default coordinates) | Use `cute.copy` directly | +| Non-trivial CTA coordinates needed | Use `cute.copy` directly | +| Non-TMA, non-S2T copy operations | Use `cute.copy` directly | + +### TMEM Allocation Decision Matrix + +| Your Situation | Recommendation | +|----------------|----------------| +| Single TMEM tensor needed | Use `TmemAllocator` directly | +| Multiple TMEM tensors from same region | Use `TmemAllocator.reserve()` → `TmemBufferPool` | +| 2-CTA MMA with shared TMEM | Use `TmemAllocator` with `is_two_cta=True` | +| Need to calculate TMEM size upfront | Use `get_num_tmem_alloc_cols()` | + +--- + +## Examples + +The `examples/block_api/` directory contains complete working examples: + +### Basic Examples + +| Example | Description | Key Concepts | +|---------|-------------|--------------| +| [`blackwell/tutorial_gemm/fp16_gemm_0.py`](../examples/block_api/blackwell/tutorial_gemm/fp16_gemm_0.py) | Simple GEMM mainloop | `block_copy`, `TmemAllocator` | +| [`hopper/elementwise_apply_tma.py`](../examples/hopper/elementwise_apply_tma.py) | TMA for elementwise ops | Basic `block_copy` usage | + +### Intermediate Examples + +| Example | Description | Key Concepts | +|---------|-------------|--------------| +| [`blackwell/tutorial_gemm/fp16_gemm_1.py`](../examples/block_api/blackwell/tutorial_gemm/fp16_gemm_1.py) | GEMM with TMA multicast | `tma_multicast`, 2-CTA MMA | +| [`blackwell/tutorial_gemm/fp16_gemm_2.py`](../examples/block_api/blackwell/tutorial_gemm/fp16_gemm_2.py) | Full pipeline + epilogue | Advanced pipeline integration | + +### Advanced Examples + +| Example | Description | Key Concepts | +|---------|-------------|--------------| +| [`blackwell/tutorial_gemm/nvfp4_gemm_0.py`](../examples/blackwell/tutorial_gemm/nvfp4_gemm_0.py) | NVFP4 block-scaled GEMM | S2T copy, scale factors | +| [`blackwell/dense_gemm_persistent.py`](../examples/blackwell/dense_gemm_persistent.py) | Persistent GEMM | Full production pattern | +| [`hopper/dense_gemm_persistent.py`](../examples/hopper/dense_gemm_persistent.py) | Hopper persistent GEMM | WGMMA with pipelines | + +### Comparing Block API vs. Manual Implementation + +Each tutorial GEMM example has a non-block-API counterpart with the same filename under `examples/blackwell/tutorial_gemm/`: + +```text +examples/block_api/blackwell/tutorial_gemm/fp16_gemm_0.py ←→ examples/blackwell/tutorial_gemm/fp16_gemm_0.py +examples/block_api/blackwell/tutorial_gemm/fp16_gemm_1.py ←→ examples/blackwell/tutorial_gemm/fp16_gemm_1.py +examples/block_api/blackwell/tutorial_gemm/fp16_gemm_2.py ←→ examples/blackwell/tutorial_gemm/fp16_gemm_2.py +``` + +Comparing these pairs shows how the block API simplifies implementation. + +--- + +## Troubleshooting + +### `block_copy`-Specific Errors + +#### `ValueError: "Argument tma_bar_ptr is required."` + +- **When**: TMA load (G2S) called without barrier pointer +- **Fix**: Pass `tma_bar_ptr=barrier_ptr` for all TMA load operations + +```python +utils.block_copy(tma_atom, src, dst, tma_bar_ptr=barrier_ptr) +``` + +#### `ValueError: "block_copy with tma_multicast expects CopyBulkTensorTileG2SOp for compiler-driven multicast"` + +- **When**: Using `tma_multicast` with a non-G2S copy operation (e.g., TMA store) +- **Fix**: Only use `tma_multicast` with TMA load operations; the compiler will upgrade to multicast automatically + +#### `NotImplementedError: "Copy op {name} is not supported yet."` + +- **When**: Using an unsupported copy operation type +- **Fix**: Currently only `TmaCopyOp` and `_S2TCopyBase` are supported; use `cute.copy` directly for other types + +### Other Common Errors + +#### `AssertionError: "num_columns must be multiple of 32 and power of two..."` + +**Cause**: Invalid TMEM allocation size. + +**Solution**: TMEM columns must be: +- Greater than 0 +- A multiple of 32 +- A power of 2 (32, 64, 128, 256, 512, ...) +- Not exceeding maximum capacity (typically 512) + +#### `AssertionError: "Cannot allocate N columns, only M remaining"` + +**Cause**: `TmemBufferPool` doesn't have enough space. + +**Solution**: Either: +1. Reserve a larger TMEM region initially +2. Use `get_num_tmem_alloc_cols()` to calculate the required size upfront + +#### Pipeline arrive count mismatch with TMA multicast + +**Cause**: When using `tma_multicast`, the number of TMA producers changes. + +**Solution**: Manually adjust your pipeline's arrive count. The compiler handles the copy operation but not the pipeline configuration: +```python +num_mcast_ctas_a = cluster_shape_m # Multicast along M +num_mcast_ctas_b = cluster_shape_n # Multicast along N +num_tma_producer = num_mcast_ctas_a + num_mcast_ctas_b - 1 +``` + +--- + +## Limitations + +### `block_copy` Limitations + +- **Supported copy kinds**: Only TMA-based copies and S2T copies are currently supported +- **Tensor layout requirement**: For TMA copies, `src` and `dst` must have TMATile as the first mode; call `group_modes(tensor, 0, 2)` on rank-2 tiles +- **Block-level pattern**: Assumes standard block-level usage; custom CTA layouts may require direct `cute.copy` calls + +### `TmemAllocator` Limitations + +- **Blackwell only**: TMEM is only available on Blackwell (SM100+) GPUs +- **Allocation constraints**: Column count must be power of 2 and multiple of 32 +- **Single allocator warp**: Only one warp (default: warp 0) performs allocations +- **Manual pipeline coordination**: For 2-CTA mode, you must provide the deallocation mbarrier + +### General Limitations + +- All utilities in this folder are **experimental** and may have breaking changes +- Some features are architecture-specific (check hardware requirements) +- These utilities are designed for common patterns; edge cases may require lower-level APIs diff --git a/python/CuTeDSL/cutlass/utils/__init__.py b/python/CuTeDSL/cutlass/utils/__init__.py index bf454cd44..8ee18c6eb 100644 --- a/python/CuTeDSL/cutlass/utils/__init__.py +++ b/python/CuTeDSL/cutlass/utils/__init__.py @@ -76,6 +76,7 @@ from .tensormap_manager import ( from .smem_allocator import ( SmemAllocator, + SmemPartition, get_smem_capacity_in_bytes, get_kernel_smem_size, ) @@ -112,6 +113,7 @@ from .mixed_input_helpers import ( contiguous_group_search, make_contiguous_group_work_tile_info, cvt_tensor_a, + cvt_tensor_a_mxf8, store_transformed_a, ) @@ -120,7 +122,12 @@ from . import distributed from . import hopper_helpers as sm90 from . import blackwell_helpers as sm100 -from .print_latex import print_latex, print_latex_tv +from .print_latex import ( + print_latex, + print_latex_tv, + PALETTES, + Band, +) from .tensor_helpers import ( is_fp8_dtype, @@ -187,6 +194,8 @@ __all__ = [ "ClcDynamicPersistentTileScheduler", "print_latex", "print_latex_tv", + "PALETTES", + "Band", "is_fp8_dtype", "create_cute_tensor_for_fp8", "distributed", diff --git a/python/CuTeDSL/cutlass/utils/blackwell_helpers.py b/python/CuTeDSL/cutlass/utils/blackwell_helpers.py index 1b2ccb72f..2f9afb8c4 100644 --- a/python/CuTeDSL/cutlass/utils/blackwell_helpers.py +++ b/python/CuTeDSL/cutlass/utils/blackwell_helpers.py @@ -20,6 +20,7 @@ from cutlass.cutlass_dsl import ( Uint8, Int8, Int32, + Boolean, Float8E4M3FN, Float8E5M2, Float6E3M2FN, @@ -940,6 +941,9 @@ def make_smem_layout_epi( :rtype: Union[cute.Layout, cute.ComposedLayout] """ + if epi_dtype is Boolean: + epi_dtype = Int8 + epilog_shape = cute.product_each( cute.shape(epi_tile, loc=loc, ip=ip), loc=loc, ip=ip ) @@ -1835,7 +1839,11 @@ def compute_epilogue_tile_size( else: (warp_m, warp_n) = (2, 2) if (cta_tile_m == 64 and use_2cta) else (4, 1) disable_source = elem_width_c is None - max_bits = elem_width_d if disable_source else max(elem_width_c, elem_width_d) # type: ignore[type-var] + if disable_source: + max_bits = elem_width_d + else: + assert elem_width_c is not None + max_bits = max(elem_width_c, elem_width_d) # -- Step 2: tile_m --------------------------------------------------- # 32 datapaths per subpartition (hardware constant); cap so each warp diff --git a/python/CuTeDSL/cutlass/utils/gemm/tensor_utils.py b/python/CuTeDSL/cutlass/utils/gemm/tensor_utils.py index 6c60c8ce1..9529784f2 100644 --- a/python/CuTeDSL/cutlass/utils/gemm/tensor_utils.py +++ b/python/CuTeDSL/cutlass/utils/gemm/tensor_utils.py @@ -213,6 +213,68 @@ def get_gemm_tensors( return A, B, D, A_cute, B_cute, D_cute +def unpack_scale_factors( + sf: torch.Tensor, sf_vec_size: int, MN: int, K: int, L: int +) -> torch.Tensor: + """Unpack a scale-factor tensor from BlockScaledBasicChunk layout to a + dense ``(MN, K, L)`` tensor. + + The on-device SF layout packs scale factors into 512-byte atoms with + a specific index mapping: 128-row MN tiles, 4-element K groups, and + an interleaved ``(mn0, mn1, k1)`` addressing within each atom. This + function inverts that mapping so that the output element at position + ``(m, k, l)`` holds the scale factor that applies to element + ``(m, k, l)`` in A/B, ready for direct element-wise multiplication. + + :param sf: Scale-factor tensor in packed layout, shape + ``(L, m_padded, k_padded)``. + :type sf: torch.Tensor + :param sf_vec_size: Number of K-elements per scale-factor block. + :type sf_vec_size: int + :param MN: Logical (unpadded) MN dimension. + :type MN: int + :param K: Logical (unpadded) K dimension. + :type K: int + :param L: Batch dimension. + :type L: int + :return: Dense FP32 tensor of shape ``(MN, K, L)`` with unpacked + scale factors. + :rtype: torch.Tensor + """ # noqa: D205 + + def index_map() -> torch.Tensor: + ATOM = 512 + ATOM_MN = 128 + ATOM_K = 4 + # Number of interleaved MN lanes within each 512-byte atom. + INTERLEAVE_FACTOR = 32 + INTERLEAVE_STRIDE = ATOM // INTERLEAVE_FACTOR # 16 + K_TILE = ATOM_K * sf_vec_size + + k_tiles = (K + K_TILE - 1) // K_TILE + mn_tiles = (MN + ATOM_MN - 1) // ATOM_MN + sf_per_l = ATOM * mn_tiles * k_tiles + + m, k = torch.meshgrid( + torch.arange(MN, device=sf.device), + torch.arange(K, device=sf.device), + indexing="ij", + ) + + base = ( + (m // ATOM_MN) * (ATOM * k_tiles) + + (k // K_TILE) * ATOM + + INTERLEAVE_STRIDE * (m % INTERLEAVE_FACTOR) + + ATOM_K * ((m % ATOM_MN) // INTERLEAVE_FACTOR) + + ((k // sf_vec_size) % ATOM_K) + ) + + l_offsets = torch.arange(L, device=sf.device)[:, None, None] * sf_per_l + return base.unsqueeze(0) + l_offsets + + return sf.flatten()[index_map()].permute(1, 2, 0) + + def create_scale_factor_tensor( MN: int, K: int, L: int, sf_vec_size: int, sf_dtype: Type[Numeric] ) -> Tuple[torch.Tensor, cute.Tensor, torch.Tensor]: @@ -249,68 +311,6 @@ def create_scale_factor_tensor( backing PyTorch CUDA tensor for pointer-array construction. :rtype: tuple[torch.Tensor, cutlass.cute.Tensor, torch.Tensor] """ - - def unpack_scale_factors( - sf: torch.Tensor, sf_vec_size: int, MN: int, K: int, L: int - ) -> torch.Tensor: - """ - Unpack a scale-factor tensor from BlockScaledBasicChunk layout to a - dense ``(MN, K, L)`` tensor. - - The on-device SF layout packs scale factors into 512-byte atoms with - a specific index mapping: 128-row MN tiles, 4-element K groups, and - an interleaved ``(mn0, mn1, k1)`` addressing within each atom. This - function inverts that mapping so that the output element at position - ``(m, k, l)`` holds the scale factor that applies to element - ``(m, k, l)`` in A/B, ready for direct element-wise multiplication. - - :param sf: Scale-factor tensor in packed layout, shape - ``(L, m_padded, k_padded)``. - :type sf: torch.Tensor - :param sf_vec_size: Number of K-elements per scale-factor block. - :type sf_vec_size: int - :param MN: Logical (unpadded) MN dimension. - :type MN: int - :param K: Logical (unpadded) K dimension. - :type K: int - :param L: Batch dimension. - :type L: int - :return: Dense FP32 tensor of shape ``(MN, K, L)`` with unpacked - scale factors. - :rtype: torch.Tensor - """ - - def index_map() -> torch.Tensor: - ATOM = 512 - ATOM_MN = 128 - ATOM_K = 4 - DATA_PATHS = 32 - DATA_PATH_STRIDE = ATOM // DATA_PATHS # 16 - K_TILE = ATOM_K * sf_vec_size - - k_tiles = (K + K_TILE - 1) // K_TILE - mn_tiles = (MN + ATOM_MN - 1) // ATOM_MN - sf_per_l = ATOM * mn_tiles * k_tiles - - m, k = torch.meshgrid( - torch.arange(MN, device=sf.device), - torch.arange(K, device=sf.device), - indexing="ij", - ) - - base = ( - (m // ATOM_MN) * (ATOM * k_tiles) - + (k // K_TILE) * ATOM - + DATA_PATH_STRIDE * (m % DATA_PATHS) - + ATOM_K * ((m % ATOM_MN) // DATA_PATHS) - + ((k // sf_vec_size) % ATOM_K) - ) - - l_offsets = torch.arange(L, device=sf.device)[:, None, None] * sf_per_l - return base.unsqueeze(0) + l_offsets - - return sf.flatten()[index_map()].permute(1, 2, 0) - ATOM_MN = 128 ATOM_K = 4 @@ -369,11 +369,12 @@ def decode_float4e2m1fn(u8: torch.Tensor) -> torch.Tensor: -6.0, ], dtype=torch.float32, + device=u8.device, ) MN, K, L = u8.shape flat = u8.permute(2, 0, 1).flatten() - idx = torch.arange(u8.numel()) + idx = torch.arange(u8.numel(), device=u8.device) byte_idx = idx // 2 shift = (idx % 2) * 4 return lut[(flat[byte_idx] >> shift) & 0xF].view(L, MN, K).permute(1, 2, 0) diff --git a/python/CuTeDSL/cutlass/utils/grouped_gemm_tile_scheduler_helper.py b/python/CuTeDSL/cutlass/utils/grouped_gemm_tile_scheduler_helper.py deleted file mode 100644 index 6956985c9..000000000 --- a/python/CuTeDSL/cutlass/utils/grouped_gemm_tile_scheduler_helper.py +++ /dev/null @@ -1,481 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: LicenseRef-NvidiaProprietary -# -# Use of this software is governed by the terms and conditions of the -# NVIDIA End User License Agreement (EULA), available at: -# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html -# -# Any use, reproduction, disclosure, or distribution of this software -# and related documentation outside the scope permitted by the EULA -# is strictly prohibited. - -from typing import List, Tuple - -import cutlass.cute as cute -from cutlass.cutlass_dsl import ( - Int32, - extract_mlir_values, - new_from_mlir_values, - const_expr, -) -from cutlass._mlir import ir - -from cutlass.utils.static_persistent_tile_scheduler import PersistentTileSchedulerParams - - -class GroupSearchResult: - """ - The result of the group search for grouped gemm. - - :param group_idx: The result group index - :type group_idx: Int32 - :param cta_tile_idx_m: CTA tile index along M dimension after rasterization - :type cta_tile_idx_m: Int32 - :param cta_tile_idx_n: CTA tile index along N dimension after rasterization - :type cta_tile_idx_n: Int32 - :param problem_shape_m: The M dimension of the gemm problem - :type problem_shape_m: Int32 - :param problem_shape_n: The N dimension of the gemm problem - :type problem_shape_n: Int32 - :param problem_shape_k: The K dimension of the gemm problem - :type problem_shape_k: Int32 - :param cta_tile_count_k: Number of tiles along K dimension - :type cta_tile_count_k: Int32 - """ - - def __init__( - self, - group_idx: Int32, - cta_tile_idx_m: Int32, - cta_tile_idx_n: Int32, - problem_shape_m: Int32, - problem_shape_n: Int32, - problem_shape_k: Int32, - cta_tile_count_k: Int32, - ) -> None: - self.group_idx = group_idx - self.cta_tile_idx_m = cta_tile_idx_m - self.cta_tile_idx_n = cta_tile_idx_n - self.problem_shape_m = problem_shape_m - self.problem_shape_n = problem_shape_n - self.problem_shape_k = problem_shape_k - self.cta_tile_count_k = cta_tile_count_k - - def __extract_mlir_values__(self) -> List[ir.Value]: - values = extract_mlir_values(self.group_idx) - values.extend(extract_mlir_values(self.cta_tile_idx_m)) - values.extend(extract_mlir_values(self.cta_tile_idx_n)) - values.extend(extract_mlir_values(self.problem_shape_m)) - values.extend(extract_mlir_values(self.problem_shape_n)) - values.extend(extract_mlir_values(self.problem_shape_k)) - values.extend(extract_mlir_values(self.cta_tile_count_k)) - return values - - def __new_from_mlir_values__(self, values: List[ir.Value]) -> "GroupSearchResult": - assert len(values) == 7 - return GroupSearchResult(*tuple(values)) - - -class GroupedGemmGroupSearchState: - """ - The state of group index search for grouped gemm. - - The state will be initialized once and updated in every round of group index search. - - :param start_group_idx: The group idx to start the search with - :type start_group_idx: Int32 - :param tile_count_prev_group: Number of tiles before the matched group - :type tile_count_prev_group: Int32 - :param tile_count_searched: Number of tiles we have searched. When the matched group is found, - it records the number of tiles including the matched group - :type tile_count_searched: Int32 - """ - - def __init__( - self, - start_group_idx: Int32, - tile_count_prev_group: Int32, - tile_count_searched: Int32, - ) -> None: - self.start_group_idx = start_group_idx - self.tile_count_prev_group = tile_count_prev_group - self.tile_count_searched = tile_count_searched - - def __extract_mlir_values__(self) -> List[ir.Value]: - values = extract_mlir_values(self.start_group_idx) - values.extend(extract_mlir_values(self.tile_count_prev_group)) - values.extend(extract_mlir_values(self.tile_count_searched)) - return values - - def __new_from_mlir_values__( - self, values: List[ir.Value] - ) -> "GroupedGemmGroupSearchState": - start_group_idx = new_from_mlir_values(self.start_group_idx, [values[0]]) - tile_count_prev_group = new_from_mlir_values( - self.tile_count_prev_group, [values[1]] - ) - tile_count_searched = new_from_mlir_values( - self.tile_count_searched, [values[2]] - ) - return GroupedGemmGroupSearchState( - start_group_idx, tile_count_prev_group, tile_count_searched - ) - - -def create_initial_search_state() -> GroupedGemmGroupSearchState: - """ - Create an initial search state for grouped gemm. - - :return: A new search state with initial values - :rtype: GroupedGemmGroupSearchState - """ - return GroupedGemmGroupSearchState( - start_group_idx=Int32(0), - tile_count_prev_group=Int32(0), - tile_count_searched=Int32(0), - ) - - -class GroupedGemmTileSchedulerHelper: - """ - A helper to translate the raw block index (x, y, z) from tile scheduler to real CTA tile index for grouped gemm. - - :param group_count: Number of groups in current grouped gemm problem - :type group_count: int - :param tile_sched_params: Parameter used to create the tile scheduler this helper works with - :type tile_sched_params: PersistentTileSchedulerParams - :param cluster_tile_shape_mnk: The shape of cluster tile as (m, n, k) - :type cluster_tile_shape_mnk: tuple[int, int, int] - :param search_state: The initial search state - :type search_state: GroupedGemmGroupSearchState - """ - - def __init__( - self, - group_count: int, - tile_sched_params: PersistentTileSchedulerParams, - cluster_tile_shape_mnk: tuple[int, int, int], - search_state: GroupedGemmGroupSearchState, - ) -> None: - self.tile_sched_params = tile_sched_params - self.group_count = group_count - self.lane_idx = cute.arch.lane_idx() - self.cluster_tile_shape_mnk = cluster_tile_shape_mnk - self.search_state = search_state - - def __extract_mlir_values__(self) -> List[ir.Value]: - values = extract_mlir_values(self.tile_sched_params) - values.extend(extract_mlir_values(self.search_state)) - return values - - def __new_from_mlir_values__( - self, values: List[ir.Value] - ) -> "GroupedGemmTileSchedulerHelper": - # Reconstruct tile_sched_params and determine how many values it consumed. - # NOTE: tile_sched_params may contain FastDivmod divisors (when swizzle_size == 1), - # which adds extra MLIR values. - params_values = extract_mlir_values(self.tile_sched_params) - n_params_values = len(params_values) - tile_sched_params = new_from_mlir_values( - self.tile_sched_params, values[:n_params_values] - ) - - # Reconstruct search_state from remaining values - search_state = new_from_mlir_values(self.search_state, values[n_params_values:]) - - return GroupedGemmTileSchedulerHelper( - self.group_count, - tile_sched_params, - self.cluster_tile_shape_mnk, - search_state, - ) - - def delinearize_z( - self, - cta_tile_coord: tuple, - problem_shape_mnkl: cute.Tensor, - ) -> GroupSearchResult: - """ - Delinearize the linear z index and return GroupSearchResult. - - This function should be used by warps that need to know the CTA tile index on M and N dimensions. - - :param cta_tile_coord: The raw CTA coordinate from tile scheduler - :type cta_tile_coord: tuple of Int32 - :param problem_shape_mnkl: Tensor containing gemm problem size (M, N, K, L) for each group - :type problem_shape_mnkl: cute.Tensor - :return: The search result containing group index and tile coordinates - :rtype: GroupSearchResult - """ - # delinear the z coord - linear_idx = cta_tile_coord[2] - group_idx, problem_mnkl = self._group_search_and_load_problem_shape( - linear_idx, - problem_shape_mnkl, - self.search_state.start_group_idx, - self.search_state.tile_count_prev_group, - ) - # linear index local to current group - cluster_tile_idx_in_current_group = ( - linear_idx - self.search_state.tile_count_prev_group - ) - cluster_count_m, cluster_count_n, cluster_count_k = cute.ceil_div( - (problem_mnkl[0], problem_mnkl[1], problem_mnkl[2]), - ( - self.cluster_tile_shape_mnk[0], - self.cluster_tile_shape_mnk[1], - self.cluster_tile_shape_mnk[2], - ), - ) - # decompose to get indices on M and N - cta_tile_idx_m, cta_tile_idx_n = self._compute_cta_tile_coord( - cluster_tile_idx_in_current_group, - cta_tile_coord, - cluster_count_m, - cluster_count_n, - ) - return GroupSearchResult( - group_idx, - cta_tile_idx_m, - cta_tile_idx_n, - problem_mnkl[0], - problem_mnkl[1], - problem_mnkl[2], - cluster_count_k, - ) - - def search_cluster_tile_count_k( - self, - cta_tile_coord: tuple, - problem_shape_mnkl: cute.Tensor, - ) -> Tuple[Int32, Int32]: - """ - Search the matched group for given linear index and compute the number of tiles along K dimension for the matched group. - - This function should be used by warps that are only interested in the number of tiles along K dimension. - - :param cta_tile_coord: The raw CTA coordinate from tile scheduler - :type cta_tile_coord: tuple of Int32 - :param problem_shape_mnkl: Tensor containing gemm problem size (M, N, K, L) for all groups - :type problem_shape_mnkl: cute.Tensor - :return: A tuple containing cluster count along K dimension and the group index - :rtype: Tuple[Int32, Int32] - """ - group_idx, problem_mnk = self._group_search_and_load_problem_shape( - cta_tile_coord[2], - problem_shape_mnkl, - self.search_state.start_group_idx, - self.search_state.tile_count_prev_group, - ) - cluster_count_k = ( - problem_mnk[2] + self.cluster_tile_shape_mnk[2] - 1 - ) // self.cluster_tile_shape_mnk[2] - return cluster_count_k, group_idx - - @cute.jit - def _prefix_sum(self, value_per_thread: Int32) -> Int32: - """ - Perform prefix sum within a full warp. - - :param value_per_thread: The value for this thread to contribute to the prefix sum - :type value_per_thread: Int32 - :return: The prefix sum result for this thread - :rtype: Int32 - """ - clamp_value = 0 - idx = 1 - sum_per_thread = value_per_thread - while const_expr(idx < cute.arch.WARP_SIZE): - value = cute.arch.shuffle_sync_up( - sum_per_thread, idx, mask_and_clamp=clamp_value - ) - if self.lane_idx >= idx: - sum_per_thread += value - idx = idx << 1 - return sum_per_thread - - def _get_problem_for_group( - self, problem_shape_mnkl: cute.Tensor, group_idx: Int32 - ) -> cute.Tensor: - """ - Load gemm problem (m,n,k,l) for the specified group from global memory to register. - - :param problem_shape_mnkl: Tensor in global memory with layout (group_count, 4):(4, 1) - :type problem_shape_mnkl: cute.Tensor - :param group_idx: The index of the group to load - :type group_idx: Int32 - :return: The problem shape tensor for the specified group - :rtype: cute.Tensor - """ - cur_problem_mnkl = cute.make_rmem_tensor( - cute.make_layout(4), problem_shape_mnkl.element_type - ) - cute.autovec_copy(problem_shape_mnkl[(group_idx, None)], cur_problem_mnkl) - return cur_problem_mnkl - - def _get_cluster_tile_count_mn(self, problem_shape: cute.Tensor) -> Int32: - """ - Compute total cluster count. - - :param problem_shape: Tensor containing problem shape (m, n, k, l) - :type problem_shape: cute.Tensor - :return: The total cluster tile count for M and N dimensions - :rtype: Int32 - """ - cur_ntile_m = ( - problem_shape[0] + self.cluster_tile_shape_mnk[0] - 1 - ) // self.cluster_tile_shape_mnk[0] - cur_ntile_n = ( - problem_shape[1] + self.cluster_tile_shape_mnk[1] - 1 - ) // self.cluster_tile_shape_mnk[1] - cur_ntile_mn = cur_ntile_m * cur_ntile_n - return cur_ntile_mn - - def _compute_cta_tile_coord( - self, - cluster_tile_idx: Int32, - cta_tile_coord_in_cluster: tuple, - cluster_tile_count_m: Int32, - cluster_tile_count_n: Int32, - ) -> tuple: - """ - Compute CTA tile indices along M and N dimensions based on the linear index within a group. - - It uses the AlongM mode to decompose the linear index onto M and N dimensions. - - :param cluster_tile_idx: The linear index within a group - :type cluster_tile_idx: Int32 - :param cta_tile_coord_in_cluster: CTA indices along M and N dimensions within a cluster - :type cta_tile_coord_in_cluster: tuple of Int32 - :param cluster_tile_count_m: The number of clusters along M dimension of the matched group - :type cluster_tile_count_m: Int32 - :param cluster_tile_count_n: The number of clusters along N dimension of the matched group - :type cluster_tile_count_n: Int32 - :return: A tuple containing CTA tile indices along M and N dimensions - :rtype: tuple of (Int32, Int32) - """ - cluster_layout_mn = cute.make_layout( - (cluster_tile_count_m, cluster_tile_count_n) - ) - (mi, ni) = cluster_layout_mn.get_hier_coord(cluster_tile_idx) - cta_tile_idx_m = ( - mi * self.tile_sched_params.cluster_shape_mn[0] - + cta_tile_coord_in_cluster[0] - ) - cta_tile_idx_n = ( - ni * self.tile_sched_params.cluster_shape_mn[1] - + cta_tile_coord_in_cluster[1] - ) - return (cta_tile_idx_m, cta_tile_idx_n) - - @cute.jit - def _group_search( - self, - linear_idx: Int32, - problem_shape_mnkl: cute.Tensor, - init_group_idx: Int32, - init_tile_count_searched: Int32, - ) -> GroupedGemmGroupSearchState: - """ - Search which group the linear index belongs to. - - :param linear_idx: The linear index to be decomposed - :type linear_idx: Int32 - :param problem_shape_mnkl: Tensor containing gemm problem size (M, N, K, L) for all groups - :type problem_shape_mnkl: cute.Tensor - :param init_group_idx: The group idx to start the search with - :type init_group_idx: Int32 - :param init_tile_count_searched: The number of tiles we have searched - :type init_tile_count_searched: Int32 - :return: The updated search state - :rtype: GroupedGemmGroupSearchState - """ - c_0 = Int32(0).ir_value() - last_lane_idx = cute.arch.WARP_SIZE - 1 - - tile_count_searched = init_tile_count_searched - start_group_idx = init_group_idx - not_found = linear_idx >= tile_count_searched - tile_count_prev_group = self.search_state.tile_count_prev_group - while not_found: - # get group to search for current lane - cur_group_idx = start_group_idx + self.lane_idx - # check if the group to be checked is out of range - inside_group_bound = cur_group_idx < self.group_count - cur_ntile_mn = c_0 - if inside_group_bound: - # get problem size of current group - cur_problem_mnkl = self._get_problem_for_group( - problem_shape_mnkl, cur_group_idx - ) - cur_ntile_mn = self._get_cluster_tile_count_mn(cur_problem_mnkl) - # compute tile count from beginning to current group(included) - total_cluster_tile_count_ps_per_thread = self._prefix_sum(cur_ntile_mn) - cluster_tile_count_end_per_thread = ( - total_cluster_tile_count_ps_per_thread + tile_count_searched - ) - - group_not_in_window = linear_idx >= cluster_tile_count_end_per_thread - hitted_group_idx_in_search_window = cute.arch.popc( - cute.arch.vote_ballot_sync(group_not_in_window) - ) - not_found = hitted_group_idx_in_search_window == cute.arch.WARP_SIZE - start_group_idx = hitted_group_idx_in_search_window + start_group_idx - hit_the_1st_problem_in_search_window = ( - hitted_group_idx_in_search_window == c_0 - ) - tile_count_prev_group = tile_count_searched - if hit_the_1st_problem_in_search_window == False: - tile_count_prev_group = cute.arch.shuffle_sync( - cluster_tile_count_end_per_thread, - hitted_group_idx_in_search_window - 1, - ) - - # If no matched group, then get new_cluster_tile_count_end from last lane - # Otherwise, get new_cluster_tile_count_end from the hitted group - lane_idx_for_cluster_tile_count_end = hitted_group_idx_in_search_window - if not_found: - lane_idx_for_cluster_tile_count_end = last_lane_idx - tile_count_searched = cute.arch.shuffle_sync( - cluster_tile_count_end_per_thread, - lane_idx_for_cluster_tile_count_end, - ) - - return GroupedGemmGroupSearchState( - start_group_idx, - tile_count_prev_group, - tile_count_searched, - ) - - def _group_search_and_load_problem_shape( - self, - linear_idx: Int32, - problem_shape_mnkl: cute.Tensor, - start_group_idx: Int32, - tile_count_searched: Int32, - ) -> Tuple[Int32, cute.Tensor]: - """ - Perform group search and load problem shape for the matched group. - - :param linear_idx: The linear index to be decomposed - :type linear_idx: Int32 - :param problem_shape_mnkl: Tensor containing gemm problem size (M, N, K, L) for all groups - :type problem_shape_mnkl: cute.Tensor - :param start_group_idx: The group idx to start the search with - :type start_group_idx: Int32 - :param tile_count_searched: The number of tiles we have searched - :type tile_count_searched: Int32 - :return: A tuple containing the final group index and the problem shape tensor - :rtype: Tuple[Int32, cute.Tensor] - """ - self.search_state = self._group_search( - linear_idx, - problem_shape_mnkl, - start_group_idx, - tile_count_searched, - ) - # get final group search state - final_group_idx = self.search_state.start_group_idx - # let's revisit if it's better to broadcast problem_shape_mnk in group_search - problem_mnkl = self._get_problem_for_group(problem_shape_mnkl, final_group_idx) - return final_group_idx, problem_mnkl diff --git a/python/CuTeDSL/cutlass/utils/mixed_input_helpers.py b/python/CuTeDSL/cutlass/utils/mixed_input_helpers.py index 2ff0dbc3a..5065ffe2e 100644 --- a/python/CuTeDSL/cutlass/utils/mixed_input_helpers.py +++ b/python/CuTeDSL/cutlass/utils/mixed_input_helpers.py @@ -18,12 +18,19 @@ from typing import Any, Optional, Type, Union import cutlass import cutlass.cute as cute from cutlass._mlir import ir +from cutlass._mlir.dialects import ( + arith as mlir_arith, + llvm as mlir_llvm, + vector as mlir_vector, +) from cutlass.cutlass_dsl import ( Boolean, Numeric, + dsl_user_op, extract_mlir_values, new_from_mlir_values, ) +from cutlass.cute.typing import Int32 as _Int32, Int4 as _Int4, Int8 as _Int8 from cutlass.utils.layout import LayoutEnum import cutlass.utils.blackwell_helpers as sm100_utils from cutlass.cute.nvgpu import cpasync, tcgen05 @@ -309,6 +316,10 @@ def epilog_gmem_copy_and_partition( # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) tTR_gC = thr_copy_t2r.partition_D(gC_epi_simt) + # Grouped kernels rebase this tensor by runtime group offsets after + # partitioning. Let the copy verifier/lowering choose the safe vector + # width at the final copy site, where the actual pointer alignment is + # known. simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), c_dtype) return tma_atom_c, bSG_sC, bSG_gC, simt_atom, tTR_gC @@ -1013,6 +1024,16 @@ def cvt_tensor_a( return rst +def cvt_tensor_a_mxf8(src: cute.Tensor, dtype: type[cutlass.Numeric]) -> cute.TensorSSA: + """Convert an int4 A-tile slice to mxfp8 via :func:`cute.arch.cvt_i4_mxf8_intrinsic`.""" + rst = src.load() + return cute.TensorSSA.from_vector( + cute.arch.cvt_i4_mxf8_intrinsic(rst, cute.size(rst.shape), dtype), + shape=rst.shape, + dtype=dtype, + ) + + def store_transformed_a( src_a: cute.Tensor, dst_a: cute.Tensor, copy_atom_a: Optional[cute.CopyAtom] ) -> None: @@ -1023,3 +1044,141 @@ def store_transformed_a( cute.copy(copy_atom_a, src_a, dst_a) else: cute.autovec_copy(src_a, dst_a) + + +# --------------------------------------------------------------------------- +# Biased Int4 -> FP8 conversion (5 instructions per int32 vs ~13 for LUT) +# +# Encoding: biased_nibble = int4_nibble XOR 0b1000 = int4 + 8 (unsigned [0,15]). +# FP8 E4M3FN bytes 0x00-0x0F are perfectly linear: FP8(v) = v * 2^-9. +# So FP8_biased(int4) = (int4 + 8) * 2^-9 for all int4 in [-8, 7]. +# +# The result is biased by 8/512 per element, requiring an epilog correction: +# result_true[m, n] = 512 * acc_biased[m, n] +# - 8 * sum_blk scale_A[m, blk] * rowsum_B[blk, n] +# --------------------------------------------------------------------------- + + +@dsl_user_op +def cvt_i4x8_to_mxf8x8_biased( + src_vec8: ir.Value, + *, + dst_type: ir.Type, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> ir.Value: + """Convert 8 packed int4 values to 8 FP8 bytes via biased bitcast. + + 5 instructions per int32 -- uses explicit ``lop3.b32`` to fuse the + ``(src ^ mask_xor) & mask_and`` pattern into a single 3-input LOP3 + (LUT 0x28). ptxas does not perform this fusion automatically. + + 1. lo = lop3(src, 0x88888888, 0x0F0F0F0F, 0x28) -- (src^0x88)&0x0F + 2. hi_raw = src >> 4 + 3. hi = lop3(hi_raw, 0x08080808, 0x0F0F0F0F, 0x28) -- (hi_raw^0x08)&0x0F + 4. out0 = prmt(lo, hi, 0x5140) -- FP8 bytes [n3, n2, n1, n0] + 5. out1 = prmt(lo, hi, 0x7362) -- FP8 bytes [n7, n6, n5, n4] + + PTX LOP3 LUT bit position is ``(a<<2) | (b<<1) | c`` (a is MSB). + Truth table for ``F(a,b,c) = (a^b)&c``: bits 3 and 5 set -> LUT = 0x28. + """ + src_i32 = mlir_llvm.bitcast(_Int32.mlir_type, src_vec8, loc=loc, ip=ip) + four = mlir_arith.constant(_Int32.mlir_type, 4, loc=loc, ip=ip) + + # lo = (src ^ 0x88) & 0x0F via fused LOP3.LUT. + lo = mlir_llvm.inline_asm( + _Int32.mlir_type, + [src_i32], + "lop3.b32 $0, $1, 0x88888888, 0x0F0F0F0F, 0x28;", + "=r,r", + ) + + # hi_raw = src >> 4 (unsigned). + hi_raw = mlir_arith.shrui(src_i32, four, loc=loc, ip=ip) + + # hi = (hi_raw ^ 0x08) & 0x0F via the same 0x28 truth table. + # XOR mask is 0x08 (not 0x88) because after the right-shift only bit 3 + # of each byte needs flipping -- the high nibbles of the original src + # land there post-shift. + hi = mlir_llvm.inline_asm( + _Int32.mlir_type, + [hi_raw], + "lop3.b32 $0, $1, 0x08080808, 0x0F0F0F0F, 0x28;", + "=r,r", + ) + + out0 = mlir_llvm.inline_asm( + _Int32.mlir_type, + [lo, hi], + "prmt.b32 $0, $1, $2, 0x5140;", + "=r,r,r", + ) + out1 = mlir_llvm.inline_asm( + _Int32.mlir_type, + [lo, hi], + "prmt.b32 $0, $1, $2, 0x7362;", + "=r,r,r", + ) + + vec_i32x2_type = ir.VectorType.get([2], _Int32.mlir_type, loc=loc) + rst_i32x2 = mlir_vector.from_elements(vec_i32x2_type, [out0, out1], loc=loc, ip=ip) + vec_mxf8x8_type = ir.VectorType.get([8], dst_type, loc=loc) + return mlir_llvm.bitcast(vec_mxf8x8_type, rst_i32x2, loc=loc, ip=ip) + + +@dsl_user_op +def cvt_i4_mxf8_biased_intrinsic( + vec_i4: ir.Value, + length: int, + dst_type: Type[Numeric], + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> ir.Value: + """Loop over 8-element int4 chunks and apply biased FP8 conversion to each.""" + assert length % 8 == 0, f"Invalid {length=}" + src_pos = 0 + orig_length = length + vec_i4x8_type = ir.VectorType.get([8], _Int4.mlir_type, loc=loc) + vec_dst_type_i8 = ir.VectorType.get([orig_length], _Int8.mlir_type, loc=loc) + zero_attr = ir.IntegerAttr.get(_Int8.mlir_type, 0) + vec_dst = mlir_arith.ConstantOp( + vec_dst_type_i8, + ir.DenseElementsAttr.get_splat(vec_dst_type_i8, zero_attr), + loc=loc, + ip=ip, + ).result + + num_vec8 = orig_length // 8 + for _ in range(num_vec8): + vec_i4x8 = mlir_vector.extract_strided_slice( + vec_i4x8_type, vec_i4, [src_pos], [8], [1], loc=loc, ip=ip + ) + vec_mxf8x8 = cvt_i4x8_to_mxf8x8_biased( + vec_i4x8, dst_type=_Int8.mlir_type, loc=loc, ip=ip + ) + vec_dst = mlir_vector.insert_strided_slice( + vec_mxf8x8, vec_dst, [src_pos], [1], loc=loc, ip=ip + ) + src_pos += 8 + length -= 8 + + return vec_dst + + +def cvt_tensor_a_biased( + src: cute.Tensor, dtype: type[cutlass.Numeric] +) -> cute.TensorSSA: + """Drop-in replacement for :func:`cvt_tensor_a` using the biased bitcast above. + + Converts an int4 A-tile slice to FP8. The accumulator the kernel collects + from ``cute.gemm(...)`` is biased; the epilog applies the correction + result_true = 512 * acc_biased - 8 * sum_blk scale_A[m, blk] * rowsum_B[blk, n] + """ + rst = src.load() + rst = cute.TensorSSA.from_vector( + cvt_i4_mxf8_biased_intrinsic(rst, cute.size(rst.shape), dtype), + shape=rst.shape, + dtype=dtype, + ) + return rst diff --git a/python/CuTeDSL/cutlass/utils/print_latex.py b/python/CuTeDSL/cutlass/utils/print_latex.py index ac87201ed..110d356ea 100644 --- a/python/CuTeDSL/cutlass/utils/print_latex.py +++ b/python/CuTeDSL/cutlass/utils/print_latex.py @@ -9,7 +9,9 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. -from typing import Callable, Optional, Union +import colorsys +from types import MappingProxyType +from typing import Callable, List, NamedTuple, Optional, Union from ..cute import ( Layout, @@ -23,21 +25,26 @@ from ..cute import ( ) from ..cute.typing import IntTuple -__all__ = ["print_latex", "print_latex_tv"] +__all__ = [ + "print_latex", + "print_latex_tv", + "PALETTES", + "Band", + "tikz_color_rainbow", + "tikz_color_tv", + "tikz_escape", +] + + +# Eight grays as percent-black, ordered so adjacent indices contrast +# strongly (light/dark interleaved) rather than forming a smooth ramp. +# Shared by the plain ``tikz_color_bwx8`` and the ``"bw"`` / ``"bw_dual"`` +# TV palettes. +_BWX8_PERCENT = (0, 40, 20, 60, 10, 50, 30, 70) def tikz_color_bwx8(idx: int) -> str: - color_map = [ - "black!00", - "black!40", - "black!20", - "black!60", - "black!10", - "black!50", - "black!30", - "black!70", - ] - return color_map[idx % 8] + return "black!%02d" % _BWX8_PERCENT[idx % 8] def tikz_color_white(idx: int) -> str: @@ -45,6 +52,8 @@ def tikz_color_white(idx: int) -> str: def tikz_color_tv(tid: int, vid: int) -> str: + """Eight-pastel TV palette, cycled by ``tid % 8`` (``vid`` unused). + The default ``color`` for ``print_latex_tv``.""" color_map = [ "{rgb,255:red,175;green,175;blue,255}", "{rgb,255:red,175;green,255;blue,175}", @@ -58,20 +67,196 @@ def tikz_color_tv(tid: int, vid: int) -> str: return color_map[tid % 8] +def tikz_color_rainbow(index: int, count: int, saturation: float = 1.0) -> str: + """Hue-wheel color for ``index`` in ``[0, count)``. + + The reusable "rainbow" primitive: hue = ``index / count`` with value + pinned at 1, so every step is a distinct color. (A physical rainbow + has no purple->red wrap; this is the full HSV wheel, so the last + indices loop back toward red.) ``count`` clamps to >= 1 so a lone + index stays well-defined. ``saturation`` < 1 washes the color toward + white -- the ``"rainbow_dual"`` palette uses that to render the value + band paler than the thread band so the two are easy to tell apart. + + The palette factories below build on this primitive, spreading hues + over the wheel by thread (and, for ``"rainbow_dual"``, by value). + """ + count = max(count, 1) + hue = (index % count) / count + r, g, b = colorsys.hsv_to_rgb(hue, saturation, 1.0) + return "{rgb,255:red,%d;green,%d;blue,%d}" % ( + round(r * 255), + round(g * 255), + round(b * 255), + ) + + +class Band(NamedTuple): + """One horizontal fill band of a TV cell. ``lo`` and ``hi`` are + fractions in ``[0, 1]`` along the screen-vertical M axis (0 = the + cell's top edge, 1 = its bottom edge); ``color`` is a TikZ fill spec. + A single ``Band(0.0, 1.0, ...)`` is an ordinary one-fill cell; the + ``"rainbow_dual"`` palette stacks two half-height bands.""" + + lo: float + hi: float + color: str + + +# A *palette* is a factory ``palette(num_tid, num_vid) -> cell`` whose +# returned ``cell(tid, vid)`` honors the original ``color(tid, vid)`` +# contract and returns either a single TikZ fill string (one fill spanning +# the cell) or a list of ``Band``s (stacked horizontal fills). Building +# the palette as a closure lets it capture the thread / value counts (and +# any other parameters) without widening the per-cell signature -- the +# counts are needed to spread hues evenly over the wheel. + + +def _palette_pastel(num_tid: int, num_vid: int) -> Callable[[int, int], str]: + """Eight-pastel palette (the default), cycled by thread.""" + return lambda tid, vid: tikz_color_tv(tid, vid) + + +def _palette_rainbow(num_tid: int, num_vid: int) -> Callable[[int, int], str]: + """Single full-cell fill, hue spread over the wheel by thread.""" + return lambda tid, vid: tikz_color_rainbow(tid, num_tid) + + +def _palette_rainbow_dual( + num_tid: int, num_vid: int +) -> Callable[[int, int], List[Band]]: + """Two stacked bands: the thread hue (full saturation) over the value + hue (reduced saturation, paler), so ``tid`` and ``vid`` read as two + distinct hue channels rather than the weak saturation fade of a + single hue, at the cost of two fills per cell.""" + + def cell(tid: int, vid: int) -> List[Band]: + return [ + Band(0.0, 0.5, tikz_color_rainbow(tid, num_tid)), + Band(0.5, 1.0, tikz_color_rainbow(vid, num_vid, saturation=0.45)), + ] + + return cell + + +# Monochrome counterparts of the color palettes above -- for B&W printing, +# papers, or a colorblind-safe fallback. + + +def _palette_white(num_tid: int, num_vid: int) -> Callable[[int, int], str]: + """Blank cells: only the T#/V# labels and the grid show.""" + return lambda tid, vid: "white" + + +def _palette_bw(num_tid: int, num_vid: int) -> Callable[[int, int], str]: + """Monochrome analogue of ``"pastel"``: the eight-gray bwx8 ramp, + cycled by thread (``vid`` unused).""" + return lambda tid, vid: tikz_color_bwx8(tid) + + +def _palette_bw_dual(num_tid: int, num_vid: int) -> Callable[[int, int], List[Band]]: + """Monochrome analogue of ``"rainbow_dual"``: the thread gray over a + lightened value gray, so ``tid`` and ``vid`` read as two bands. The + value band is scaled toward white -- the role the desaturation plays + in ``"rainbow_dual"`` -- so the bottom band stays consistently paler + than the top.""" + + def cell(tid: int, vid: int) -> List[Band]: + return [ + Band(0.0, 0.5, tikz_color_bwx8(tid)), + Band(0.5, 1.0, "black!%d" % round(0.45 * _BWX8_PERCENT[vid % 8])), + ] + + return cell + + +# Named palettes selectable by string -- the "map of palettes". The +# ``palette`` argument of ``print_latex_tv`` also accepts a factory +# directly for one-off palettes. Exposed as a read-only ``MappingProxyType`` +# so a caller cannot mutate the shared registry and alter rendering +# process-wide; register a one-off palette by passing the factory directly. +PALETTES = MappingProxyType( + { + "pastel": _palette_pastel, + "rainbow": _palette_rainbow, + "rainbow_dual": _palette_rainbow_dual, + "white": _palette_white, + "bw": _palette_bw, + "bw_dual": _palette_bw_dual, + } +) + + +def tikz_escape(s: str) -> str: + """Escape the LaTeX specials that show up in free-form title text + (class paths, layout strings): ``\\``, ``_``, ``#``, ``&``, ``%``, + ``$``, braces, ``^``, ``~``.""" + repl = { + "\\": r"\textbackslash{}", + "_": r"\_", + "#": r"\#", + "&": r"\&", + "%": r"\%", + "$": r"\$", + "{": r"\{", + "}": r"\}", + "^": r"\textasciicircum{}", + "~": r"\textasciitilde{}", + } + return "".join(repl.get(c, c) for c in s) + + +def _latex_emitter( + render_func: Callable[[str], None] | None, +) -> "tuple[Callable[..., None], Callable[[], None]]": + """Build the output sink shared by ``print_latex`` and + ``print_latex_tv``. + + Returns a pair ``(emit, finish)``: + + - ``emit(*args)`` writes one chunk of LaTeX. With no ``render_func`` + it prints to stdout (the standalone-document path); otherwise it + buffers the space-joined chunk for later delivery. + - ``finish()`` hands the whole buffered ``{tikzpicture}`` to + ``render_func`` in a single call (the embed-in-a-notebook path); + it is a no-op when printing. + + Factoring this out keeps the two public functions from each + re-declaring the identical buffer / dispatch logic (the form this + started in before the helper was extracted). + """ + buffer: List[str] = [] + + def emit(*args: object) -> None: + if render_func is None: + print(*args) + else: + buffer.append(" ".join(str(arg) for arg in args)) + + def finish() -> None: + if render_func is not None: + render_func(" ".join(buffer)) + + return emit, finish + + def print_latex( x: Union[Layout, ComposedLayout], *, color: Callable = tikz_color_bwx8, - render_func: Optional[Callable[[str], None]] = None, + render_func: Callable[[str], None] | None = None, ) -> None: """ Prints a layout. + :param x: A layout :type x: Union[Layout, ComposedLayout] :param color: A function that returns TiKZ colors :type color: Callable - :param render_func: An user provided function to render the latex output, which only includes tikz picture section. If None, it will print to stdout. - :type render_func: Optional[Callable] + :param render_func: Optional callback fed the ``{tikzpicture}`` body + (without the standalone-document wrapper) in a single call, + instead of printing a full LaTeX document to stdout. + :type render_func: Callable[[str], None] | None """ if not is_static(x): @@ -84,20 +269,16 @@ def print_latex( else: layout = x - latex_output = [] - - def print_or_append(*args): - if render_func is not None: - latex_output.append(" ".join(str(arg) for arg in args)) - else: - print(*args) + emit, finish = _latex_emitter(render_func) + # The standalone-document wrapper is only useful for the print path; + # the render hook receives just the {tikzpicture} body. if render_func is None: - print_or_append("%% Layout: {}", layout) - print_or_append("\\documentclass[convert]{standalone}") - print_or_append("\\usepackage{tikz}") - print_or_append("\\begin{document}") - print_or_append( + emit("%% Layout: {}", layout) + emit("\\documentclass[convert]{standalone}") + emit("\\usepackage{tikz}") + emit("\\begin{document}") + emit( "\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},every node/.style={minimum size=1cm, outer sep=0pt}]" ) @@ -106,24 +287,22 @@ def print_latex( for m in range(M): for n in range(N): idx = layout((m, n)) - print_or_append("\\node[fill=") - print_or_append(color(idx)) - print_or_append("] at (%d,%d) {%d};\n" % (m, n, idx)) - print_or_append( + emit("\\node[fill=") + emit(color(idx)) + emit("] at (%d,%d) {%d};\n" % (m, n, idx)) + emit( "\\draw[color=black,thick,shift={(-0.5,-0.5)}] (0,0) grid (%d,%d);\n\n" % (M, N) ) for m in range(M): - print_or_append("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n" % (m, -1, m)) + emit("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n" % (m, -1, m)) for n in range(N): - print_or_append("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n" % (-1, n, n)) + emit("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n" % (-1, n, n)) ## Footer - print_or_append("\\end{tikzpicture}") + emit("\\end{tikzpicture}") if render_func is None: - print_or_append("\\end{document}") - - if render_func is not None: - render_func(" ".join(latex_output)) + emit("\\end{document}") + finish() def print_latex_tv( @@ -131,38 +310,64 @@ def print_latex_tv( tile_mn: Union[IntTuple, Layout], *, color: Callable = tikz_color_tv, - render_func: Optional[Callable[[str], None]] = None, + palette: Optional[Union[str, Callable]] = None, + title: Optional[str] = None, + axis_labels: bool = False, + render_func: Callable[[str], None] | None = None, ) -> None: """ Prints a tv layout for a tile M N. Everything must be static. + :param layout_tv: A static thread value layout :type layout_tv: Union[Layout, ComposedLayout] :param tile_mn: A static M N tile :type tile_mn: Union[IntTuple, Layout] - :param color: A function that returns TiKZ colors + :param color: A function ``color(tid, vid) -> str`` returning a TikZ + fill color for the cell owned by thread ``tid`` value ``vid``. + Used when ``palette`` is ``None``; ignored otherwise. :type color: Callable - :param render_func: An user provided function to render the latex output, which only includes tikz picture section. If None, it will print to stdout. - :type render_func: Optional[Callable] + :param palette: Optional richer cell coloring that supersedes + ``color``. Either the name of a built-in palette (a key of + ``PALETTES``: the color ``"pastel"``, ``"rainbow"``, + ``"rainbow_dual"`` or the monochrome ``"white"``, ``"bw"``, + ``"bw_dual"``) or a factory ``palette(num_tid, num_vid) -> cell`` + where + ``cell(tid, vid)`` returns either a TikZ fill string (one fill + spanning the cell) or a list of ``Band``s (stacked horizontal + fills). The factory form lets a palette capture the thread / + value counts -- needed to spread hues evenly over the wheel -- + without widening the per-cell ``(tid, vid)`` contract. + :type palette: Optional[Union[str, Callable]] + :param title: Optional title drawn above the figure, one line per + ``\\n``-separated segment (e.g. the operator / function / + tensor that produced this layout). LaTeX specials are escaped. + :type title: Optional[str] + :param axis_labels: When ``True``, annotate the M (row, downward) + and N (column, rightward) axis directions. The picture's + coordinate basis runs M down and N right; these labels make + that orientation explicit so a reader can tell whether a + thread's values are contiguous along the memory-major axis. + :type axis_labels: bool + :param render_func: Optional callback fed the ``{tikzpicture}`` body + (without the standalone-document wrapper) in a single call, + instead of printing a full LaTeX document to stdout. + :type render_func: Callable[[str], None] | None """ if not is_static(layout_tv) or not is_static(tile_mn): raise ValueError("Layout tv and tile_mn must be static") if rank(layout_tv) != 2: raise ValueError("Require layout_tv to be rank 2") - latex_output = [] - - def print_or_append(*args): - if render_func is not None: - latex_output.append(" ".join(str(arg) for arg in args)) - else: - print(*args) + emit, finish = _latex_emitter(render_func) + # The standalone-document wrapper is only useful for the print path; + # the render hook receives just the {tikzpicture} body. if render_func is None: - print_or_append("%% Layout TV: {}", layout_tv) - print_or_append("\\documentclass[convert]{standalone}") - print_or_append("\\usepackage{tikz}") - print_or_append("\\begin{document}") - print_or_append( + emit("%% Layout TV: {}", layout_tv) + emit("\\documentclass[convert]{standalone}") + emit("\\usepackage{tikz}") + emit("\\begin{document}") + emit( "\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},every node/.style={minimum size=1cm, outer sep=0pt}]\n" ) @@ -172,30 +377,94 @@ def print_latex_tv( M, N = product_each(tile_mn.shape) # type: ignore[union-attr] filled = [[False for n in range(N)] for m in range(M)] - for tid in range(size(layout_tv, mode=[0])): - for vid in range(size(layout_tv, mode=[1])): + num_tid = size(layout_tv, mode=[0]) + num_vid = size(layout_tv, mode=[1]) + + # Resolve a per-cell ``cell(tid, vid)`` function. Without a + # ``palette`` we use ``color`` directly (it already has the + # ``(tid, vid)`` contract). A ``palette`` is a factory called once + # with the thread / value counts, so it captures them in a closure + # instead of taking them per cell. + if palette is None: + cell = color + else: + if isinstance(palette, str): + if palette not in PALETTES: + raise ValueError( + "Unknown palette %r; choose one of %s or pass a factory" + % (palette, sorted(PALETTES)) + ) + palette = PALETTES[palette] + cell = palette(num_tid, num_vid) + + for tid in range(num_tid): + for vid in range(num_vid): idx = layout_tv((tid, vid)) m = (idx // tile_mn.stride[0]) % tile_mn.shape[0] # type: ignore[operator, union-attr, index] n = (idx // tile_mn.stride[1]) % tile_mn.shape[1] # type: ignore[operator, union-attr, index] - if not filled[m][n]: - filled[m][n] = True - print_or_append( - "\\node[fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n" - % (color(tid, vid), m, n, tid, vid) + if filled[m][n]: + continue + filled[m][n] = True + # A cell is one or more horizontal bands stacked along the + # screen-vertical M axis. Normalize a bare fill string to a + # single full-height band so both cases share one code path. + bands = cell(tid, vid) + if isinstance(bands, str): + bands = [Band(0.0, 1.0, bands)] + for band in bands: + # Cell (m, n) spans [m-0.5, m+0.5] x [n-0.5, n+0.5]; the + # band covers the M-fraction [lo, hi] (0 = top edge at + # m-0.5). Drawn as a filled rectangle. + emit( + "\\path[fill=%s] (%g,%g) rectangle (%g,%g);\n" + % ( + band.color, + m - 0.5 + band.lo, + n - 0.5, + m - 0.5 + band.hi, + n + 0.5, + ) ) + # Fill-less T#/V# label on top of the band(s). + emit("\\node at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n" % (m, n, tid, vid)) - print_or_append( + emit( "\\draw[color=black,thick,shift={(-0.5,-0.5)}] (0,0) grid (%d,%d);\n\n" % (M, N) ) for m in range(M): - print_or_append("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n" % (m, -1, m)) + emit("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n" % (m, -1, m)) for n in range(N): - print_or_append("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n" % (-1, n, n)) + emit("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n" % (-1, n, n)) + + # Optional axis-direction labels and title. The basis maps a data + # coord (a, b) to screen (b, -a), so the index labels sit just + # outside the grid (column indices at data x=-1, row indices at + # data y=-1); place these annotations just past those index labels + # so they sit close to the numbering without overlapping it. + if axis_labels: + emit("\\node at (%g,%g) {\\Large N (col) $\\rightarrow$};\n" % (-1.7, N / 2.0)) + # Rotate the M label by pi/2 so it reads alongside the (vertical, + # downward-running) M axis instead of across it. + emit( + "\\node[rotate=90] at (%g,%g) {\\Large $\\leftarrow$ M (row)};\n" + % (M / 2.0, -2.0) + ) + if title: + body = " \\\\ ".join( + "\\texttt{" + tikz_escape(line) + "}" for line in title.split("\n") + ) + # Anchor the title by its south (bottom) edge so extra title + # lines grow it upward without changing its gap from the N-axis + # label below. ``anchor`` is a canvas-space offset, unaffected + # by the picture's M-down/N-right coordinate basis, so south is + # always the visual bottom. + emit( + "\\node[anchor=south] at (%g,%g) {\\shortstack{%s}};\n" + % (-2.2, N / 2.0, body) + ) ## Footer - print_or_append("\\end{tikzpicture}") + emit("\\end{tikzpicture}") if render_func is None: - print_or_append("\\end{document}") - - if render_func is not None: - render_func(" ".join(latex_output)) + emit("\\end{document}") + finish() diff --git a/python/CuTeDSL/cutlass/utils/smem_allocator.py b/python/CuTeDSL/cutlass/utils/smem_allocator.py index 67889fa2b..57fc83ba1 100644 --- a/python/CuTeDSL/cutlass/utils/smem_allocator.py +++ b/python/CuTeDSL/cutlass/utils/smem_allocator.py @@ -9,12 +9,14 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. +from enum import IntEnum from typing import Any, Callable, Optional, Type, Union, overload from typing_extensions import deprecated import inspect import cutlass.cute as cute from cutlass.cute.tensor import _Tensor +from cutlass.address_space import AddressSpace from cutlass.cutlass_dsl import ( SMEM_CAPACITY_MAP, CuTeDSL, @@ -52,6 +54,20 @@ def _extract_struct_fields(struct_type: cute.struct) -> list[tuple[str, int, int return fields +class SmemPartition(IntEnum): + """Shared memory partition types. + + Defines the partition ownership of a shared memory region. + + Attributes: + RESERVED: Partition reserved for system use (value: 0). + USER: Partition available for user allocation (value: 1). + """ + + RESERVED = 0 + USER = 1 + + class SmemAllocator: """A helper class for managing shared memory allocation on GPU. @@ -134,9 +150,61 @@ class SmemAllocator: """ pass + @dsl_user_op + def calculate_partition_size( + self, + partition: SmemPartition, + *, + cumulative: bool = False, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> Int32: + """Get the size of shared memory allocation at given smem partition. + + :param partition: The smem partition to query + :type partition: SmemPartition + :param cumulative: Whether to return the cumulative size of all partitions up to and including the given partition + :type cumulative: bool, optional + :param loc: Source location information for debugging, defaults to None + :type loc: Optional[ir.Location] + :param ip: Insertion point for MLIR operations, defaults to None + :type ip: Optional[ir.InsertionPoint] + """ + return Int32( + _cute_ir.smem_partition_size( + partition_id=partition.value, + cumulative=cumulative, + loc=loc, + ip=ip, + ) + ) + + @dsl_user_op + def calculate_total_usage( + self, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, + ) -> Int32: + """Get total kernel smem usage calculated by allocator. + + :param loc: Source location information for debugging, defaults to None + :type loc: Optional[ir.Location] + :param ip: Insertion point for MLIR operations, defaults to None + :type ip: Optional[ir.InsertionPoint] + """ + return Int32( + _cute_ir.smem_partition_size( + partition_id=SmemPartition.USER.value, + cumulative=True, + loc=loc, + ip=ip, + ) + ) + @property @deprecated( - "Internal var `_allocated_bytes` is deprecated, use public API `arch.dynamic_smem_size()` instead." + "Private attribute `_allocated_bytes` is deprecated, use public API `arch.dynamic_smem_size()` instead." ) def _allocated_bytes(self) -> Int32: return cute.arch.dynamic_smem_size() @@ -150,6 +218,7 @@ class SmemAllocator: swizzle: Optional[cute.Swizzle] = None, struct_fields: Optional[list[tuple[str, int, int]]] = None, *, + partition: SmemPartition = SmemPartition.USER, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None, ) -> cute.Pointer: @@ -159,24 +228,29 @@ class SmemAllocator: Returns: cute.Pointer: An iterator (pointer) to the allocated shared memory. """ + # use init_byte_alignment for first allocation + byte_alignment = max(byte_alignment, dtype.width // 8) assert byte_alignment <= 1024, "max shared memory alignment limit to 1024 bytes" assert cute.is_static(layout), "shared memory allocation must be static layout" # allocate using cute.memref.alloca swizzle = swizzle.type.attribute if swizzle is not None else None mlir_type = Int8.mlir_type if dtype is Boolean else dtype.mlir_type ptr_ty = _cute_ir.PtrType.get( - mlir_type, cute.AddressSpace.smem, byte_alignment, swizzle + mlir_type, AddressSpace.smem, byte_alignment, swizzle ) res_ty = _cute_ir.MemRefType.get(ptr_ty, layout.type) memref = _cute_ir.memref_alloca(res_ty, layout=None, loc=loc, ip=ip) + # Attach attributes to allocate to smem partition + if partition is not None: + memref.value.owner.attributes["smem.partition_id"] = ir.IntegerAttr.get( + ir.IntegerType.get_signless(32), partition.value + ) # Attach struct field metadata as MLIR attributes if struct_fields: - from cutlass._mlir import ir as _ir - field_attrs = [] for name, size, offset in struct_fields: - field_attrs.append(_ir.StringAttr.get(f"{name}:{size}:{offset}")) - memref.value.owner.attributes["smem.struct_fields"] = _ir.ArrayAttr.get( + field_attrs.append(ir.StringAttr.get(f"{name}:{size}:{offset}")) + memref.value.owner.attributes["smem.struct_fields"] = ir.ArrayAttr.get( field_attrs ) return _cute_ir.get_iter(memref.value, loc=loc, ip=ip) @@ -185,8 +259,9 @@ class SmemAllocator: def allocate( self, size_or_type: int, - byte_alignment: int, + byte_alignment: int = 1, *, + partition: SmemPartition = SmemPartition.USER, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None, ) -> cute.Pointer: ... @@ -195,8 +270,9 @@ class SmemAllocator: def allocate( self, size_or_type: Type[Numeric], - byte_alignment: int, + byte_alignment: int = 1, *, + partition: SmemPartition = SmemPartition.USER, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None, ) -> cute.Pointer: ... @@ -205,8 +281,9 @@ class SmemAllocator: def allocate( # type: ignore[overload-cannot-match] self, size_or_type: cute.struct, - byte_alignment: int, + byte_alignment: int = 1, *, + partition: SmemPartition = SmemPartition.USER, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None, ) -> cute.Pointer: ... @@ -217,6 +294,7 @@ class SmemAllocator: size_or_type: Any, byte_alignment: int = 1, *, + partition: SmemPartition = SmemPartition.USER, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None, ) -> cute.Pointer: @@ -256,6 +334,7 @@ class SmemAllocator: Int8, alignment, struct_fields=struct_fields, + partition=partition, loc=loc, ip=ip, ) @@ -268,7 +347,9 @@ class SmemAllocator: ): element_width = size_or_type.width if size_or_type is not Boolean else 8 size_in_bytes = cute.ceil_div(element_width, 8) - base_ptr = self.allocate(size_in_bytes, byte_alignment, loc=loc, ip=ip) + base_ptr = self.allocate( + size_in_bytes, byte_alignment, partition=partition, loc=loc, ip=ip + ) return cute.recast_ptr(base_ptr, dtype=size_or_type, loc=loc, ip=ip) else: raise TypeError( @@ -282,7 +363,9 @@ class SmemAllocator: raise ValueError("`byte_alignment` must be at least 1") layout = cute.make_layout(size_in_bytes) - return self._smem_alloca(layout, Int8, byte_alignment, loc=loc, ip=ip) + return self._smem_alloca( + layout, Int8, byte_alignment, partition=partition, loc=loc, ip=ip + ) @dsl_user_op def allocate_array( @@ -293,6 +376,7 @@ class SmemAllocator: num_elems: int = 1, *, byte_alignment: int = 1, + partition: SmemPartition = SmemPartition.USER, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None, ) -> cute.Pointer: @@ -322,7 +406,11 @@ class SmemAllocator: element_width = element_type.width if element_type is not Boolean else 8 byte_alignment = max(byte_alignment, element_width // 8) ptr = self.allocate( - element_width * num_elems // 8, byte_alignment, loc=loc, ip=ip + element_width * num_elems // 8, + byte_alignment, + partition=partition, + loc=loc, + ip=ip, ) return cute.recast_ptr(ptr, dtype=element_type, loc=loc, ip=ip) @@ -335,6 +423,7 @@ class SmemAllocator: byte_alignment: int = 1, swizzle: Optional[cute.Swizzle] = None, *, + partition: SmemPartition = SmemPartition.USER, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None, ) -> cute.Tensor: @@ -399,7 +488,9 @@ class SmemAllocator: raise ValueError("invalid allocation that is not byte-aligned") num_bytes = cosize_in_bits // 8 - ptr = self.allocate(num_bytes, byte_alignment, loc=loc, ip=ip) + ptr = self.allocate( + num_bytes, byte_alignment, partition=partition, loc=loc, ip=ip + ) ptr = cute.recast_ptr(ptr, swizzle, dtype=element_type, loc=loc, ip=ip) tensor = cute.make_tensor(ptr, layout, loc=loc, ip=ip) return _Tensor(tensor, dtype=element_type, loc=loc, ip=ip) @@ -453,4 +544,4 @@ def get_kernel_smem_size( f"get_kernel_smem_size: kernel '{kernel.__name__}' has not been " "traced yet — call .launch() before get_kernel_smem_size()." ) - return _cute_ir.kernel_smem_size(sym, loc=loc, ip=ip) + return _cute_ir.kernel_smem_size(kernel_name=sym, loc=loc, ip=ip) diff --git a/python/CuTeDSL/cutlass/utils/tensormap_manager.py b/python/CuTeDSL/cutlass/utils/tensormap_manager.py index fc52cf9da..cf67d9cb3 100644 --- a/python/CuTeDSL/cutlass/utils/tensormap_manager.py +++ b/python/CuTeDSL/cutlass/utils/tensormap_manager.py @@ -16,11 +16,11 @@ from typing import Optional, Tuple import cutlass._mlir.dialects.cute as _cute_ir import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir from cutlass._mlir import ir +from cutlass.address_space import AddressSpace from cutlass.cutlass_dsl import dsl_user_op import cutlass.cute as cute from cutlass import const_expr -from cutlass.cute.core import AddressSpace as _CuteAddressSpace from cutlass.cute.core import make_ptr as _cute_make_ptr @@ -54,14 +54,14 @@ class TensorMapManager: def get_tensormap_ptr( self, ptr: cute.Pointer, - address_space: _cute_ir.AddressSpace = _cute_ir.AddressSpace.gmem, + address_space: AddressSpace = AddressSpace.gmem, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None, ) -> cute.Pointer: if address_space not in [ - _cute_ir.AddressSpace.gmem, - _cute_ir.AddressSpace.generic, + AddressSpace.gmem, + AddressSpace.generic, ]: raise ValueError(f"Invalid address space: {address_space} for tensormap") @@ -148,7 +148,7 @@ class TensorMapManager: _cute_make_ptr( p.dtype, cute.arch.make_warp_uniform(p.toint(), loc=loc, ip=ip), - mem_space=_CuteAddressSpace.smem, + mem_space=AddressSpace.smem, assumed_align=p.alignment, ) for p in tensormap_smem_ptr diff --git a/python/CuTeDSL/cutlass/utils/tmem_allocator.py b/python/CuTeDSL/cutlass/utils/tmem_allocator.py index 8b840a52e..7f538047a 100644 --- a/python/CuTeDSL/cutlass/utils/tmem_allocator.py +++ b/python/CuTeDSL/cutlass/utils/tmem_allocator.py @@ -18,6 +18,8 @@ from cutlass.base_dsl.arch import Arch from cutlass.cutlass_dsl import ( Numeric, Float32, + Int32, + Int64, Boolean, extract_mlir_values, new_from_mlir_values, @@ -29,6 +31,7 @@ from cutlass._mlir import ir from cutlass.cute.nvgpu.tcgen05 import find_tmem_tensor_col_offset from cutlass.cute.arch import get_max_tmem_alloc_cols, get_min_tmem_alloc_cols from cutlass.cute.arch.constants import WARP_SIZE +from cutlass.utils import SmemAllocator, SmemPartition _TMEM_COL_MASK = 0x0000FFFF @@ -300,13 +303,13 @@ class TmemAllocator: @dsl_user_op def __init__( self, - alloc_result_dst_smem_ptr: cute.Pointer, + alloc_result_dst_smem_ptr: Optional[cute.Pointer] = None, + *, barrier_for_retrieve: pipeline.NamedBarrier, allocator_warp_id: int = 0, is_two_cta: bool = False, num_allocated_columns: int = 0, two_cta_tmem_dealloc_mbar_ptr: Optional[cute.Pointer] = None, - *, arch: str = "sm_100", initialize_mbarrier: bool = True, loc: Optional[ir.Location] = None, @@ -324,8 +327,12 @@ class TmemAllocator: If `is_two_cta` is set to True, this will initialize the mbarrier pointer required for tensor memory deallocation across two CTAs. - :param alloc_result_dst_smem_ptr: The shared memory pointer that holds the base address of allocated tensor memory. - :type alloc_result_dst_smem_ptr: cute.Pointer + Auto-allocated smem pointers: If `alloc_result_dst` and `two_cta_tmem_dealloc` ptrs are omitted, + the allocator creates them to low address of the smem region, + so they can be treated as reserved partition allocations and survive kernel smem resize if needed. + + :param alloc_result_dst_smem_ptr: Shared memory pointer holding the base address of allocated tensor memory. If None, the allocator auto-allocates it in the reserved address. + :type alloc_result_dst_smem_ptr: Optional[cute.Pointer] :param barrier_for_retrieve: The named barrier for retrieving the tensor memory pointer. :type barrier_for_retrieve: pipeline.NamedBarrier :param allocator_warp_id: The warp ID of the allocator warp, defaults to 0. @@ -334,8 +341,8 @@ class TmemAllocator: :type is_two_cta: bool, optional :param num_allocated_columns: The number of columns allocated in tensor memory, defaults to 0. :type num_allocated_columns: int, optional - :param two_cta_tmem_dealloc_mbar_ptr: The mbarrier pointer required for two-CTA tensor memory deallocation, optional. - :type two_cta_tmem_dealloc_mbar_ptr: cute.Pointer, optional + :param two_cta_tmem_dealloc_mbar_ptr: Mbarrier pointer for two-CTA tensor memory deallocation. If None and is_two_cta, the allocator auto-allocates it in the reserved address. + :type two_cta_tmem_dealloc_mbar_ptr: Optional[cute.Pointer] :param initialize_mbarrier: Whether to initialize the mbarrier for two cta, defaults to True. :type initialize_mbarrier: bool, optional :param loc: Optional codegen location for debugging and error reporting. @@ -343,9 +350,35 @@ class TmemAllocator: :param ip: Optional insertion point for codegen. :type ip: Any, optional - :raises AssertionError: If two_cta_tmem_dealloc_mbar_ptr is None while is_two_cta is True. + :raises ValueError: If only provided one of required smem ptr in two cta. """ - # TODO: automatically maintain a smem address + # automatically maintain smem ptr at reserved low address + smem = SmemAllocator() + if is_two_cta and ( + (two_cta_tmem_dealloc_mbar_ptr is None) + != (alloc_result_dst_smem_ptr is None) + ): + # Partial control (one user, one auto) would leave lifecycle ambiguous + raise ValueError( + "Mixed smem ptr ownership is not allowed for two-CTA: " + "omit both (auto via TmemAllocator) or provide both (manual)." + ) + if alloc_result_dst_smem_ptr is None: + alloc_result_dst_smem_ptr = smem.allocate( + Int32, + byte_alignment=8, + partition=SmemPartition.RESERVED, + loc=loc, + ip=ip, + ) + if (two_cta_tmem_dealloc_mbar_ptr is None) and is_two_cta: + two_cta_tmem_dealloc_mbar_ptr = smem.allocate( + Int64, + byte_alignment=8, + partition=SmemPartition.RESERVED, + loc=loc, + ip=ip, + ) self._alloc_result_dst_smem_ptr = alloc_result_dst_smem_ptr self._allocator_warp_id = allocator_warp_id self._is_two_cta = is_two_cta @@ -379,15 +412,15 @@ class TmemAllocator: else None ) return TmemAllocator( - new_alloc_result_dst_smem_ptr, - pipeline.NamedBarrier( + alloc_result_dst_smem_ptr=new_alloc_result_dst_smem_ptr, + barrier_for_retrieve=pipeline.NamedBarrier( barrier_id=self._barrier_for_retrieve.barrier_id, num_threads=self._barrier_for_retrieve.num_threads, ), - self._allocator_warp_id, - self._is_two_cta, - self._num_allocated_columns, - new_two_cta_tmem_dealloc_mbar_ptr, + allocator_warp_id=self._allocator_warp_id, + is_two_cta=self._is_two_cta, + num_allocated_columns=self._num_allocated_columns, + two_cta_tmem_dealloc_mbar_ptr=new_two_cta_tmem_dealloc_mbar_ptr, arch=self._arch, # Preserve the architecture parameter initialize_mbarrier=False, ) diff --git a/python/CuTeDSL/requirements-cu13.txt b/python/CuTeDSL/requirements-cu13.txt index 961756b78..561815af1 100644 --- a/python/CuTeDSL/requirements-cu13.txt +++ b/python/CuTeDSL/requirements-cu13.txt @@ -1,3 +1,3 @@ # Use `pip install -r requirements-cu13.txt` with the present file to install a # wheel consistent with the present state of the github repository -nvidia-cutlass-dsl[cu13]==4.6.0.dev0 +nvidia-cutlass-dsl[cu13]==4.6.0 diff --git a/python/CuTeDSL/requirements.txt b/python/CuTeDSL/requirements.txt index 985c46ef0..2f5b51e0f 100644 --- a/python/CuTeDSL/requirements.txt +++ b/python/CuTeDSL/requirements.txt @@ -1,3 +1,3 @@ # Use `pip install -r requirements.txt` with the present file to install a # wheel consistent with the present state of the github repository -nvidia-cutlass-dsl==4.6.0.dev0 +nvidia-cutlass-dsl==4.6.0 diff --git a/python/cutlass_library/generator.py b/python/cutlass_library/generator.py index 37665e906..fa68104a4 100644 --- a/python/cutlass_library/generator.py +++ b/python/cutlass_library/generator.py @@ -11214,9 +11214,21 @@ def GenerateSM120_TensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cud [16, 8, 32] ] - tile_sizes = [ - [128, 32, 128], - [128, 64, 128], + # Cooperative supports TileN >= 8 (see sm120_blockscaled_mma_builder.inl AtomLayoutMNK). + # Pingpong uses AtomLayout Shape<_2,_2,_1>, giving a natural TiledMma N of 16, + # so pingpong tiles start at N = 16. + tile_sizes_cooperative = [ + [128, 8, 128], + [128, 16, 128], + [128, 32, 128], + [128, 64, 128], + [128, 128, 128] + ] + + tile_sizes_pingpong = [ + [128, 16, 128], + [128, 32, 128], + [128, 64, 128], [128, 128, 128] ] @@ -11273,11 +11285,6 @@ def GenerateSM120_TensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cud ) for math_inst in math_instructions: - tile_descriptions = [] - for tile_size in tile_sizes: - tile_descriptions.append( - TileDescription(tile_size, 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) - data_types = [ { "a_type" : math_inst.element_a, @@ -11325,12 +11332,19 @@ def GenerateSM120_TensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cud for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] - for data_type, kernel_schedule in product(data_types, kernel_schedules): - CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type, - [[kernel_schedule, EpilogueScheduleType.ScheduleAuto]], - tile_schedulers = tile_schedulers(data_type["sfd_type"], kernel_schedule), - gemm_kind = gemm_kind - ) + for kernel_schedule in kernel_schedules: + tile_sizes = tile_sizes_pingpong if is_pingpong(kernel_schedule) else tile_sizes_cooperative + tile_descriptions = [] + for tile_size in tile_sizes: + tile_descriptions.append( + TileDescription(tile_size, 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) + + for data_type in data_types: + CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type, + [[kernel_schedule, EpilogueScheduleType.ScheduleAuto]], + tile_schedulers = tile_schedulers(data_type["sfd_type"], kernel_schedule), + gemm_kind = gemm_kind + ) def GenerateSM120_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.BlockScaledUniversal3x): # SM120 MMA with with F4 + block scale @@ -11348,21 +11362,30 @@ def GenerateSM120_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_versio [16, 8, 64] ] + # Cooperative supports TileN >= 8 (see sm120_blockscaled_mma_builder.inl AtomLayoutMNK). + # Pingpong uses AtomLayout Shape<_2,_2,_1>, giving a natural TiledMma N of 16, + # so pingpong tiles start at N = 16. tile_sizes_cooperative = [ - [128, 32, 128], - [128, 32, 256], - [128, 64, 128], - [128, 64, 256], + [128, 8, 128], + [128, 8, 256], + [128, 16, 128], + [128, 16, 256], + [128, 32, 128], + [128, 32, 256], + [128, 64, 128], + [128, 64, 256], [128, 128, 128], [128, 128, 256], [256, 128, 128] ] tile_sizes_pingpong = [ - [128, 32, 128], - [128, 32, 256], - [128, 64, 128], - [128, 64, 256], + [128, 16, 128], + [128, 16, 256], + [128, 32, 128], + [128, 32, 256], + [128, 64, 128], + [128, 64, 256], [128, 128, 128], [128, 128, 256] ] diff --git a/test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/CMakeLists.txt b/test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/CMakeLists.txt index 20657c53e..0a83176c0 100644 --- a/test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/CMakeLists.txt +++ b/test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/CMakeLists.txt @@ -71,6 +71,7 @@ cutlass_test_unit_gemm_device_add_executable( cutlass_test_unit_gemm_device_add_executable( cutlass_test_unit_bs_grouped_gemm_device_tensorop_sm120 sm120_bs_gemm_nvf4_nvf4_f32_nvf4_group_gemm_fusion.cu + sm120_bs_gemm_nvf4_nvf4_f32_bf16_group_gemm.cu sm120_bs_gemm_mxf8_mxf4_f32_group_gemm_fusion.cu ) diff --git a/test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/sm120_bs_gemm_mxf4_mxf4_f32_f32.cu b/test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/sm120_bs_gemm_mxf4_mxf4_f32_f32.cu index 9c88a368d..3ab98f095 100644 --- a/test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/sm120_bs_gemm_mxf4_mxf4_f32_f32.cu +++ b/test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/sm120_bs_gemm_mxf4_mxf4_f32_f32.cu @@ -248,4 +248,132 @@ TEST(SM120_Device_Blockscaled_Gemm_mxf4t_mxf4n_f32n_tensor_op_f32, 128x32x256) { EXPECT_TRUE(result); } +namespace kernel_4 { + using ElementA = cutlass::float_e2m1_t; + using ElementB = cutlass::float_e2m1_t; + using ElementC = float; + using ElementD = float; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementSF = cutlass::float_ue8m0_t; + + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using LayoutD = cutlass::layout::ColumnMajor; + + using ElementPairA = cutlass::mx_float4_t; + using ElementPairB = cutlass::mx_float4_t; + + static constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits::value; // Align to 16 bytes. + static constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits::value; // Align to 16 bytes. + static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; + static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; + + using TileShape = Shape<_128,_16,_256>; + using ClusterShape = Shape<_1,_1,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp, + TileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, AlignmentC, + ElementD, LayoutD, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp, + ElementPairA, LayoutA, AlignmentA, + ElementPairB, LayoutB, AlignmentB, + ElementAccumulator, + TileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecializedCooperative + >::CollectiveOp; + + template + struct dummy { + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + }; + using GemmKernel = typename dummy::GemmKernel; + using Gemm = typename dummy::Gemm; + +} // kernel_4 + +TEST(SM120_Device_Blockscaled_Gemm_mxf4t_mxf4n_f32n_tensor_op_f32, 128x16x256) { + bool result = test::gemm::device::TestSmall(1.0, 0.5); + EXPECT_TRUE(result); +} + +namespace kernel_5 { + using ElementA = cutlass::float_e2m1_t; + using ElementB = cutlass::float_e2m1_t; + using ElementC = float; + using ElementD = float; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementSF = cutlass::float_ue8m0_t; + + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using LayoutD = cutlass::layout::ColumnMajor; + + using ElementPairA = cutlass::mx_float4_t; + using ElementPairB = cutlass::mx_float4_t; + + static constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits::value; // Align to 16 bytes. + static constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits::value; // Align to 16 bytes. + static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; + static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; + + using TileShape = Shape<_128,_8,_256>; + using ClusterShape = Shape<_1,_1,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp, + TileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, AlignmentC, + ElementD, LayoutD, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp, + ElementPairA, LayoutA, AlignmentA, + ElementPairB, LayoutB, AlignmentB, + ElementAccumulator, + TileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecializedCooperative + >::CollectiveOp; + + template + struct dummy { + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + }; + using GemmKernel = typename dummy::GemmKernel; + using Gemm = typename dummy::Gemm; + +} // kernel_5 + +TEST(SM120_Device_Blockscaled_Gemm_mxf4t_mxf4n_f32n_tensor_op_f32, 128x8x256) { + bool result = test::gemm::device::TestSmall(1.0, 0.5); + EXPECT_TRUE(result); +} + #endif // (defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED)) diff --git a/test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/sm120_bs_gemm_nvf4_nvf4_f32_bf16_group_gemm.cu b/test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/sm120_bs_gemm_nvf4_nvf4_f32_bf16_group_gemm.cu new file mode 100644 index 000000000..013a647b1 --- /dev/null +++ b/test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/sm120_bs_gemm_nvf4_nvf4_f32_bf16_group_gemm.cu @@ -0,0 +1,351 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + + +/*! \file + \brief Tests for device-wide grouped GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/default_epilogue.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/epilogue/thread/activation.h" + +#include "../../../common/cutlass_unit_test.h" +#include "../gemm_testbed_3x_ptr_array.hpp" + + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) + +TEST(SM120_Device_Gemm_e2m1t_e2m1n_bf16n_tensorop_group_pingpong, tile_128x32x128) { + using ElementInput = float_e2m1_t; + using ElementA = cutlass::nv_float4_t; + using ElementB = cutlass::nv_float4_t; + using ElementC = cutlass::bfloat16_t; + using ElementD = cutlass::bfloat16_t; + using ElementCompute = float; + using ElementAccumulator = float; + using ElementSF = cutlass::float_ue4m3_t; + using GmemLayoutA = cutlass::layout::RowMajor; + using GmemLayoutB = cutlass::layout::ColumnMajor; + using GmemLayoutC = cutlass::layout::ColumnMajor; + using TileShape_MNK = Shape<_128,_32,_128>; + using ClusterShape_MNK = Shape<_1,_1,_1>; + + constexpr int AlignmentA = 128 / cutlass::sizeof_bits::value; + constexpr int AlignmentB = 128 / cutlass::sizeof_bits::value; + constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; + constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; + + // + // Construct CollectiveEpilogue + // + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp, + TileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, GmemLayoutC *, AlignmentC, + ElementD, GmemLayoutC *, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto + >::CollectiveOp; + + // + // Construct CollectiveMainloop + // + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp, + ElementA, GmemLayoutA *, AlignmentA, + ElementB, GmemLayoutB *, AlignmentB, + ElementAccumulator, + TileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpong + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cutlass::gemm::GroupProblemShape>, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmall(1.0, 0.5); + EXPECT_TRUE(pass); +} + +TEST(SM120_Device_Gemm_e2m1t_e2m1n_bf16n_tensorop_group_cooperative, tile_128x32x128) { + using ElementInput = float_e2m1_t; + using ElementA = cutlass::nv_float4_t; + using ElementB = cutlass::nv_float4_t; + using ElementC = cutlass::bfloat16_t; + using ElementD = cutlass::bfloat16_t; + using ElementCompute = float; + using ElementAccumulator = float; + using ElementSF = cutlass::float_ue4m3_t; + using GmemLayoutA = cutlass::layout::RowMajor; + using GmemLayoutB = cutlass::layout::ColumnMajor; + using GmemLayoutC = cutlass::layout::ColumnMajor; + using TileShape_MNK = Shape<_128,_32,_128>; + using ClusterShape_MNK = Shape<_1,_1,_1>; + + constexpr int AlignmentA = 128 / cutlass::sizeof_bits::value; + constexpr int AlignmentB = 128 / cutlass::sizeof_bits::value; + constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; + constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; + + // + // Construct CollectiveEpilogue + // + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp, + TileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, GmemLayoutC *, AlignmentC, + ElementD, GmemLayoutC *, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto + >::CollectiveOp; + + // + // Construct CollectiveMainloop + // + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp, + ElementA, GmemLayoutA *, AlignmentA, + ElementB, GmemLayoutB *, AlignmentB, + ElementAccumulator, + TileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperative + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cutlass::gemm::GroupProblemShape>, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmall(1.0, 0.5); + EXPECT_TRUE(pass); +} + +TEST(SM120_Device_Gemm_e2m1t_e2m1n_bf16n_tensorop_group_pingpong, tile_128x16x128) { + using ElementInput = float_e2m1_t; + using ElementA = cutlass::nv_float4_t; + using ElementB = cutlass::nv_float4_t; + using ElementC = cutlass::bfloat16_t; + using ElementD = cutlass::bfloat16_t; + using ElementCompute = float; + using ElementAccumulator = float; + using ElementSF = cutlass::float_ue4m3_t; + using GmemLayoutA = cutlass::layout::RowMajor; + using GmemLayoutB = cutlass::layout::ColumnMajor; + using GmemLayoutC = cutlass::layout::ColumnMajor; + using TileShape_MNK = Shape<_128,_16,_128>; + using ClusterShape_MNK = Shape<_1,_1,_1>; + + constexpr int AlignmentA = 128 / cutlass::sizeof_bits::value; + constexpr int AlignmentB = 128 / cutlass::sizeof_bits::value; + constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; + constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; + + // + // Construct CollectiveEpilogue + // + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp, + TileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, GmemLayoutC *, AlignmentC, + ElementD, GmemLayoutC *, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto + >::CollectiveOp; + + // + // Construct CollectiveMainloop + // + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp, + ElementA, GmemLayoutA *, AlignmentA, + ElementB, GmemLayoutB *, AlignmentB, + ElementAccumulator, + TileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpong + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cutlass::gemm::GroupProblemShape>, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmall(1.0, 0.5); + EXPECT_TRUE(pass); +} + +TEST(SM120_Device_Gemm_e2m1t_e2m1n_bf16n_tensorop_group_cooperative, tile_128x16x128) { + using ElementInput = float_e2m1_t; + using ElementA = cutlass::nv_float4_t; + using ElementB = cutlass::nv_float4_t; + using ElementC = cutlass::bfloat16_t; + using ElementD = cutlass::bfloat16_t; + using ElementCompute = float; + using ElementAccumulator = float; + using ElementSF = cutlass::float_ue4m3_t; + using GmemLayoutA = cutlass::layout::RowMajor; + using GmemLayoutB = cutlass::layout::ColumnMajor; + using GmemLayoutC = cutlass::layout::ColumnMajor; + using TileShape_MNK = Shape<_128,_16,_128>; + using ClusterShape_MNK = Shape<_1,_1,_1>; + + constexpr int AlignmentA = 128 / cutlass::sizeof_bits::value; + constexpr int AlignmentB = 128 / cutlass::sizeof_bits::value; + constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; + constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; + + // + // Construct CollectiveEpilogue + // + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp, + TileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, GmemLayoutC *, AlignmentC, + ElementD, GmemLayoutC *, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto + >::CollectiveOp; + + // + // Construct CollectiveMainloop + // + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp, + ElementA, GmemLayoutA *, AlignmentA, + ElementB, GmemLayoutB *, AlignmentB, + ElementAccumulator, + TileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperative + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cutlass::gemm::GroupProblemShape>, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmall(1.0, 0.5); + EXPECT_TRUE(pass); +} + +TEST(SM120_Device_Gemm_e2m1t_e2m1n_bf16n_tensorop_group_cooperative, tile_128x8x128) { + using ElementInput = float_e2m1_t; + using ElementA = cutlass::nv_float4_t; + using ElementB = cutlass::nv_float4_t; + using ElementC = cutlass::bfloat16_t; + using ElementD = cutlass::bfloat16_t; + using ElementCompute = float; + using ElementAccumulator = float; + using ElementSF = cutlass::float_ue4m3_t; + using GmemLayoutA = cutlass::layout::RowMajor; + using GmemLayoutB = cutlass::layout::ColumnMajor; + using GmemLayoutC = cutlass::layout::ColumnMajor; + using TileShape_MNK = Shape<_128,_8,_128>; + using ClusterShape_MNK = Shape<_1,_1,_1>; + + constexpr int AlignmentA = 128 / cutlass::sizeof_bits::value; + constexpr int AlignmentB = 128 / cutlass::sizeof_bits::value; + constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; + constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; + + // + // Construct CollectiveEpilogue + // + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp, + TileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, GmemLayoutC *, AlignmentC, + ElementD, GmemLayoutC *, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto + >::CollectiveOp; + + // + // Construct CollectiveMainloop + // + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp, + ElementA, GmemLayoutA *, AlignmentA, + ElementB, GmemLayoutB *, AlignmentB, + ElementAccumulator, + TileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperative + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cutlass::gemm::GroupProblemShape>, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmall(1.0, 0.5); + EXPECT_TRUE(pass); +} +#endif // #if defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) diff --git a/test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/sm120_bs_gemm_nvf4_nvf4_f32_nvf4_group_gemm_fusion.cu b/test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/sm120_bs_gemm_nvf4_nvf4_f32_nvf4_group_gemm_fusion.cu index 8c4285126..ffe548c00 100644 --- a/test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/sm120_bs_gemm_nvf4_nvf4_f32_nvf4_group_gemm_fusion.cu +++ b/test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/sm120_bs_gemm_nvf4_nvf4_f32_nvf4_group_gemm_fusion.cu @@ -204,7 +204,6 @@ TEST(SM120_Device_Gemm_e2m1t_e2m1n_e2m1t_tensorop_f32_epilogue_VS16_group_pingpo EXPECT_TRUE(pass); } - TEST(SM120_Device_Gemm_e2m1t_e2m1n_e2m1t_tensorop_f32_epilogue_VS16_group_pingpong, silu_row_sf) { using ElementInput = float_e2m1_t; using ElementA = cutlass::nv_float4_t; @@ -352,8 +351,6 @@ TEST(SM120_Device_Gemm_e2m1t_e2m1n_e2m1t_tensorop_f32_epilogue_VS16_group_cooper EXPECT_TRUE(pass); } - - TEST(SM120_Device_Gemm_e2m1t_e2m1n_e2m1t_tensorop_f32_epilogue_VS16_group_cooperative, silu_row_sf) { using ElementInput = float_e2m1_t; using ElementA = cutlass::nv_float4_t; diff --git a/test/unit/gemm/device/sm120_tensorop_gemm/CMakeLists.txt b/test/unit/gemm/device/sm120_tensorop_gemm/CMakeLists.txt index b990b5e52..b82d7383c 100644 --- a/test/unit/gemm/device/sm120_tensorop_gemm/CMakeLists.txt +++ b/test/unit/gemm/device/sm120_tensorop_gemm/CMakeLists.txt @@ -26,7 +26,7 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -if (CUTLASS_NVCC_ARCHS MATCHES 120a) +if (CUTLASS_NVCC_ARCHS MATCHES 120a|120f|121a) add_custom_target( cutlass_test_unit_gemm_device_sm120_tensorop diff --git a/test/unit/gemm/device/sm120_tensorop_gemm/sm120_gemm_f8_f8_f32_tensor_op_group_gemm.cu b/test/unit/gemm/device/sm120_tensorop_gemm/sm120_gemm_f8_f8_f32_tensor_op_group_gemm.cu index b9b91312f..db4b82209 100644 --- a/test/unit/gemm/device/sm120_tensorop_gemm/sm120_gemm_f8_f8_f32_tensor_op_group_gemm.cu +++ b/test/unit/gemm/device/sm120_tensorop_gemm/sm120_gemm_f8_f8_f32_tensor_op_group_gemm.cu @@ -30,12 +30,7 @@ **************************************************************************************************/ /*! \file - \brief Tests for SM120 device-wide ptr-array grouped GEMM (non-blockscaled FP8). - - Exercises CollectiveMma and the - CollectiveBuilder specializations for KernelPtrArrayTmaWarpSpecializedPingpongSm120 - and KernelPtrArrayTmaWarpSpecializedCooperativeSm120 added for SM120 consumer - Blackwell (GB10 / DGX Spark, RTX 5090-class) grouped GEMM (MoE expert dispatch). + \brief Tests for SM120 device-wide ptr-array grouped GEMM. Layout constraint: only TN (RowMajor A, ColumnMajor B) is supported by the SM120 array TMA collective builder. diff --git a/test/unit/nvrtc/thread/nvrtc_config.in b/test/unit/nvrtc/thread/nvrtc_config.in index 6291b93bd..15306fdf7 100644 --- a/test/unit/nvrtc/thread/nvrtc_config.in +++ b/test/unit/nvrtc/thread/nvrtc_config.in @@ -1,3 +1,4 @@ #pragma once #define CUDA_INCLUDE_DIR "@CUDA_TOOLKIT_ROOT_DIR@/include" +#define CUDA_CCCL_INCLUDE_DIR "@CUTLASS_CUDA_CCCL_INCLUDE_DIR@" diff --git a/test/unit/nvrtc/thread/nvrtc_contraction.cu b/test/unit/nvrtc/thread/nvrtc_contraction.cu index 1c48d06f1..089ecdb6e 100644 --- a/test/unit/nvrtc/thread/nvrtc_contraction.cu +++ b/test/unit/nvrtc/thread/nvrtc_contraction.cu @@ -52,7 +52,7 @@ TEST(SM90_nvrtc_kernel, Contraction) { "-arch=sm_90", "-I" CUDA_INCLUDE_DIR, #if (__CUDACC_VER_MAJOR__ >= 13) - "-I" CUDA_INCLUDE_DIR "/cccl", + "-I" CUDA_CCCL_INCLUDE_DIR, #endif // __CUDACC_VER_MAJOR__ >= 13 };