From 85f5dd5d3dd52875e77fe9446c15cea11f9932c0 Mon Sep 17 00:00:00 2001 From: sepcnt <30561671+sepcnt@users.noreply.github.com> Date: Sun, 24 May 2026 16:15:27 +0800 Subject: [PATCH 1/6] Add block-scale MMA builtin Expose a TileLang builtin for Blackwell fp4/fp6 block-scale MMA, add the CUDA codegen lowering, and provide small helper builtins used by warp-local bf16 stores. --- src/cuda/codegen/codegen_cuda.cc | 101 +++++++ src/cuda/codegen/codegen_cuda.h | 2 + src/cuda/codegen/ptx.cc | 220 ++++++++++++++ src/cuda/codegen/ptx.h | 29 ++ src/op/builtin.cc | 5 + src/op/builtin.h | 20 ++ .../cuda/instruction/mma_blockscale.h | 275 ++++++++++++++++++ tilelang/language/__init__.py | 1 + tilelang/language/ast/ir.py | 2 + tilelang/language/builtin.py | 5 + tilelang/language/tir/ir.py | 1 + tilelang/language/tir/op.py | 54 ++++ 12 files changed, 715 insertions(+) create mode 100644 src/tl_templates/cuda/instruction/mma_blockscale.h diff --git a/src/cuda/codegen/codegen_cuda.cc b/src/cuda/codegen/codegen_cuda.cc index ff3f3d70c0..91048c81a4 100644 --- a/src/cuda/codegen/codegen_cuda.cc +++ b/src/cuda/codegen/codegen_cuda.cc @@ -601,6 +601,10 @@ std::string CodeGenTileLangCUDA::Finish() { if (need_mma_instruction_h_) { decl_stream << "#include \n"; } + if (need_mma_blockscale_instruction_h_) { + decl_stream + << "#include \n"; + } if (need_wgmma_instruction_h_) { decl_stream << "#include \n"; } @@ -2759,6 +2763,103 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) { replacer.register_rule("(C_ptr)", c_ref); replacer.register_rule("(C_offset)", c_bias); this->stream << replacer.rewrite(mma_call); + } else if (op->op.same_as(tl::ptx_mma_blockscale())) { + // arg 0: shape: mXnXkX + // arg 1: A layout: row/col + // arg 2: B layout: row/col + // arg 3: A precision: e2m1, ... + // arg 4: B precision: e2m1, ... + // arg 5: C precision: fp32, ... + // arg 6: scale vector width + // arg 7: A multiplicand + // arg 8: A multiplicand index + // arg 9: B multiplicand + // arg 10: B multiplicand index + // arg 11: C accumulator + // arg 12: C accumulator index + // arg 13: packed A scale vector + // arg 14: packed B scale vector + // arg 15: A scale selector + // arg 16: B scale selector + ICHECK_EQ(op->args.size(), 17U) + << "ptx_mma_blockscale expects 17 arguments"; + std::string shape = Downcast(op->args[0])->value; + std::string A_layout = Downcast(op->args[1])->value; + std::string B_layout = Downcast(op->args[2])->value; + std::string A_dtype = Downcast(op->args[3])->value; + std::string B_dtype = Downcast(op->args[4])->value; + std::string C_dtype = Downcast(op->args[5])->value; + int scale_vec = Downcast(op->args[6])->value; + constexpr int arg_base = 7; + std::string a_ref = this->PrintExpr(op->args[arg_base + 0]); + PrimExpr a_offset_expr = op->args[arg_base + 1]; + std::string b_ref = this->PrintExpr(op->args[arg_base + 2]); + PrimExpr b_offset_expr = op->args[arg_base + 3]; + std::string c_ref = this->PrintExpr(op->args[arg_base + 4]); + std::string c_offset = this->PrintExpr(op->args[arg_base + 5]); + std::string sfa = this->PrintExpr(op->args[arg_base + 6]); + std::string sfb = this->PrintExpr(op->args[arg_base + 7]); + std::string id_a = this->PrintExpr(op->args[arg_base + 8]); + std::string id_b = this->PrintExpr(op->args[arg_base + 9]); + auto dtype_a_enum = tl::codegen::ptx::DTypeFromString(A_dtype); + auto dtype_b_enum = tl::codegen::ptx::DTypeFromString(B_dtype); + auto dtype_c_enum = tl::codegen::ptx::DTypeFromString(C_dtype); + // The Python intrinsic accepts logical fp4 element offsets. Pointer + // arithmetic below is byte-addressed for packed e2m1 operands. + if (dtype_a_enum == tl::codegen::ptx::DataType::kFloat4_e2m1fn) { + a_offset_expr = arith::Analyzer().Simplify(truncdiv(a_offset_expr, 2)); + } + if (dtype_b_enum == tl::codegen::ptx::DataType::kFloat4_e2m1fn) { + b_offset_expr = arith::Analyzer().Simplify(truncdiv(b_offset_expr, 2)); + } + std::string a_offset = this->PrintExpr(a_offset_expr); + std::string b_offset = this->PrintExpr(b_offset_expr); + auto [m, n, k] = tl::codegen::ptx::ParseMMAShape(shape); + tl::codegen::ptx::GetMMABlockScaleInfo( + dtype_a_enum, dtype_b_enum, dtype_c_enum, m, n, k, + A_layout == "row" ? false : true, B_layout == "row" ? false : true, + scale_vec); + + need_mma_blockscale_instruction_h_ = true; + this->PrintIndent(); + std::string mma_call = + "tl::mma_sync_blockscale<(AType), (BType), (CType), (M), (N), (K), " + "(TransA), (TransB), (ScaleVec)>(" + "reinterpret_cast<(CRegType)*>((C_ptr) + (C_offset)), " + "reinterpret_cast((A_ptr) + (A_offset)), " + "reinterpret_cast((B_ptr) + (B_offset)), " + "(SFA), (SFB), static_cast((IdA)), " + "static_cast((IdB)));\n"; + tl::codegen::Replacer replacer; + replacer.register_rule("(AType)", + tl::codegen::ptx::DTypeEnumToString(dtype_a_enum)); + replacer.register_rule("(BType)", + tl::codegen::ptx::DTypeEnumToString(dtype_b_enum)); + replacer.register_rule("(CType)", + tl::codegen::ptx::DTypeEnumToString(dtype_c_enum)); + replacer.register_rule("(M)", std::to_string(m)); + replacer.register_rule("(N)", std::to_string(n)); + replacer.register_rule("(K)", std::to_string(k)); + replacer.register_rule("(TransA)", A_layout == "row" ? "false" : "true"); + replacer.register_rule("(TransB)", B_layout == "row" ? "false" : "true"); + replacer.register_rule("(ScaleVec)", std::to_string(scale_vec)); + replacer.register_rule("(ARegType)", + tl::codegen::GetMMARegisterType(dtype_a_enum)); + replacer.register_rule("(BRegType)", + tl::codegen::GetMMARegisterType(dtype_b_enum)); + replacer.register_rule("(CRegType)", + tl::codegen::GetMMARegisterType(dtype_c_enum)); + replacer.register_rule("(A_ptr)", a_ref); + replacer.register_rule("(A_offset)", a_offset); + replacer.register_rule("(B_ptr)", b_ref); + replacer.register_rule("(B_offset)", b_offset); + replacer.register_rule("(C_ptr)", c_ref); + replacer.register_rule("(C_offset)", c_offset); + replacer.register_rule("(SFA)", sfa); + replacer.register_rule("(SFB)", sfb); + replacer.register_rule("(IdA)", id_a); + replacer.register_rule("(IdB)", id_b); + this->stream << replacer.rewrite(mma_call); } else if (op->op.same_as(tl::tma_store_cluster())) { ICHECK_EQ(op->args.size(), 5U) << "tma_store_cluster requires 5 args"; this->PrintIndent(); diff --git a/src/cuda/codegen/codegen_cuda.h b/src/cuda/codegen/codegen_cuda.h index 54367cef6c..f6ace98564 100644 --- a/src/cuda/codegen/codegen_cuda.h +++ b/src/cuda/codegen/codegen_cuda.h @@ -114,6 +114,8 @@ class CodeGenTileLangCUDA final : public CodeGenC { bool need_mma_h_{false}; // whether need tl mma instruction header bool need_mma_instruction_h_{false}; + // whether need tl block-scaled mma instruction header + bool need_mma_blockscale_instruction_h_{false}; // whether need tl wgmma instruction header bool need_wgmma_instruction_h_{false}; // whether need tl tcgen05mma instruction header diff --git a/src/cuda/codegen/ptx.cc b/src/cuda/codegen/ptx.cc index 5ee92917bc..d8893a5b00 100644 --- a/src/cuda/codegen/ptx.cc +++ b/src/cuda/codegen/ptx.cc @@ -139,6 +139,122 @@ inline std::string DTypeToString(DataType dtype) { return dtype_str[static_cast(dtype)]; } +static const DataType valid_mxf8f6f4_blockscale_dtypes[] = { + DataType::kFloat4_e2m1fn, DataType::kFloat6_e2m3fn, + DataType::kFloat6_e3m2fn, DataType::kFloat8_e4m3, DataType::kFloat8_e5m2}; + +static bool IsValidMXF8F6F4BlockScaleDType(DataType dtype) { + for (DataType valid_dtype : valid_mxf8f6f4_blockscale_dtypes) { + if (dtype == valid_dtype) { + return true; + } + } + return false; +} + +static std::string DTypeToBlockScaleMMAString(DataType dtype) { + ICHECK(IsValidMXF8F6F4BlockScaleDType(dtype)) + << DTypeToString(dtype) << " is not a block-scaled MMA operand type."; + std::string enum_str = enum_to_str[static_cast(dtype)]; + size_t suffix_pos = enum_str.find('_'); + ICHECK(suffix_pos != std::string::npos) + << enum_str << " is not a block-scaled MMA operand data type."; + + std::string operand = enum_str.substr(suffix_pos + 1); + if (operand.size() >= 2 && + operand.compare(operand.size() - 2, 2, "fn") == 0) { + operand.resize(operand.size() - 2); + } + return operand; +} + +static std::string ScaleVecToString(int scale_vec) { + return std::to_string(scale_vec) + "X"; +} + +struct MMABlockScaleConfig { + explicit MMABlockScaleConfig(int m, int n, int k, DataType dtype_a, + DataType dtype_b, DataType dtype_c, bool trans_a, + bool trans_b, int scale_vec, const char *kind, + const char *scale_dtype) + : m(m), n(n), k(k), dtype_a(dtype_a), dtype_b(dtype_b), dtype_c(dtype_c), + trans_a(trans_a), trans_b(trans_b), scale_vec(scale_vec), kind(kind), + scale_dtype(scale_dtype) {} + + bool Matches(DataType other_dtype_a, DataType other_dtype_b, + DataType other_dtype_c, int other_m, int other_n, int other_k, + bool other_trans_a, bool other_trans_b, + int other_scale_vec) const { + return m == other_m && n == other_n && k == other_k && + dtype_a == other_dtype_a && dtype_b == other_dtype_b && + dtype_c == other_dtype_c && trans_a == other_trans_a && + trans_b == other_trans_b && scale_vec == other_scale_vec; + } + + MMABlockScaleInfo Info() const { + return {kind, ScaleVecToString(scale_vec), scale_dtype}; + } + + std::string ToString() const { + return std::string("kind::") + kind + + ".block_scale.scale_vec::" + ScaleVecToString(scale_vec) + ".m" + + std::to_string(m) + "n" + std::to_string(n) + "k" + + std::to_string(k) + "." + (trans_a ? "col" : "row") + "." + + (trans_b ? "col" : "row") + ".f32." + + DTypeToBlockScaleMMAString(dtype_a) + "." + + DTypeToBlockScaleMMAString(dtype_b) + ".f32." + scale_dtype; + } + + int m, n, k; + DataType dtype_a, dtype_b, dtype_c; + bool trans_a, trans_b; + int scale_vec; + const char *kind; + const char *scale_dtype; +}; + +static std::vector MakeValidMMABlockScaleConfigs() { + std::vector configs; + configs.reserve(27); + for (DataType dtype_a : valid_mxf8f6f4_blockscale_dtypes) { + for (DataType dtype_b : valid_mxf8f6f4_blockscale_dtypes) { + configs.emplace_back(16, 8, 32, dtype_a, dtype_b, DataType::kFloat32, + false, true, 1, "mxf8f6f4", "ue8m0"); + } + } + + configs.emplace_back(16, 8, 64, DataType::kFloat4_e2m1fn, + DataType::kFloat4_e2m1fn, DataType::kFloat32, false, + true, 2, "mxf4nvf4", "ue8m0"); + configs.emplace_back(16, 8, 64, DataType::kFloat4_e2m1fn, + DataType::kFloat4_e2m1fn, DataType::kFloat32, false, + true, 4, "mxf4nvf4", "ue4m3"); + return configs; +} + +static const std::vector valid_mma_blockscale_configs = + MakeValidMMABlockScaleConfigs(); + +MMABlockScaleInfo GetMMABlockScaleInfo(DataType dtype_a, DataType dtype_b, + DataType dtype_c, int m, int n, int k, + bool trans_a, bool trans_b, + int scale_vec) { + for (const MMABlockScaleConfig &valid_config : valid_mma_blockscale_configs) { + if (valid_config.Matches(dtype_a, dtype_b, dtype_c, m, n, k, trans_a, + trans_b, scale_vec)) { + return valid_config.Info(); + } + } + + LOG(FATAL) << "Unsupported SM120 block-scaled MMA configuration: m" << m + << "n" << n << "k" << k << ", dtype_a=" << DTypeToString(dtype_a) + << ", dtype_b=" << DTypeToString(dtype_b) + << ", dtype_c=" << DTypeToString(dtype_c) + << ", layout=" << (trans_a ? "col" : "row") << "." + << (trans_b ? "col" : "row") << ", scale_vec=" << scale_vec; + return {"", "", ""}; +} + /*! * \brief Get the number of bits of given PTX data type. */ @@ -1021,6 +1137,52 @@ GetMMAOperands(int m, int n, int k, ptx::DataType dtype_a, return std::make_tuple(templates.str(), inputs.str(), outputs.str()); } +inline std::tuple +GetMMABlockScaleOperands() { + std::stringstream templates, inputs, outputs; + + int arg_counter = 0; + templates << "{%" << arg_counter++ << ", %" << arg_counter++ << ", %" + << arg_counter++ << ", %" << arg_counter++ << "}, "; + templates << "{%" << arg_counter++ << ", %" << arg_counter++ << ", %" + << arg_counter++ << ", %" << arg_counter++ << "}, "; + templates << "{%" << arg_counter++ << ", %" << arg_counter++ << "}, "; + templates << "{%" << arg_counter++ << ", %" << arg_counter++ << ", %" + << arg_counter++ << ", %" << arg_counter++ << "}, "; + templates << "{%" << arg_counter++ << "}, {%" << arg_counter++ << ", %" + << arg_counter++ << "}, "; + templates << "{%" << arg_counter++ << "}, {%" << arg_counter++ << ", %" + << arg_counter++ << "}"; + + for (int i = 0; i < 4; ++i) { + if (i != 0) { + outputs << ", "; + } + outputs << "\"=f\"(((float *)({D}))[" << i << "])"; + } + + for (int i = 0; i < 4; ++i) { + if (i != 0) { + inputs << ", "; + } + inputs << "\"r\"(((unsigned *)({A}))[" << i << "])"; + } + for (int i = 0; i < 2; ++i) { + inputs << ", \"r\"(((unsigned *)({B}))[" << i << "])"; + } + for (int i = 0; i < 4; ++i) { + inputs << ", \"f\"(((float *)({C}))[" << i << "])"; + } + inputs << ", \"r\"(uint32_t(({SFA})))"; + inputs << ", \"h\"((short)0)"; + inputs << ", \"h\"((short)({IdA}))"; + inputs << ", \"r\"(uint32_t(({SFB})))"; + inputs << ", \"h\"((short)0)"; + inputs << ", \"h\"((short)({IdB}))"; + + return std::make_tuple(templates.str(), inputs.str(), outputs.str()); +} + inline std::tuple GetWGMMAOperands(int m, int n, int k, ptx::DataType dtype_a, ptx::DataType dtype_b, ptx::DataType dtype_c, bool sparse, @@ -1203,6 +1365,64 @@ PrintMMAAssembly(const std::string &shape, const std::string &A_layout, return asm_code; } +std::string PrintMMABlockScaleAssembly( + const std::string &shape, const std::string &A_layout, + const std::string &B_layout, const std::string &A_dtype, + const std::string &B_dtype, const std::string &C_dtype, int scale_vec, + const std::string &a_ptr, const std::string &a_elem_offset, + const std::string &b_ptr, const std::string &b_elem_offset, + const std::string &c_ptr, const std::string &c_elem_offset, + const std::string &sfa, const std::string &sfb, const std::string &id_a, + const std::string &id_b) { + ptx::DataType dtype_a = ptx::DTypeFromString(A_dtype), + dtype_b = ptx::DTypeFromString(B_dtype), + dtype_c = ptx::DTypeFromString(C_dtype); + ptx::LayoutType layout_a = ptx::LayoutTypeFromString(A_layout), + layout_b = ptx::LayoutTypeFromString(B_layout); + auto [m, n, k] = ptx::ParseMMAShape(shape); + ptx::MMABlockScaleInfo info = ptx::GetMMABlockScaleInfo( + dtype_a, dtype_b, dtype_c, m, n, k, + layout_a == ptx::LayoutType::kColumnMajor, + layout_b == ptx::LayoutType::kColumnMajor, scale_vec); + + std::string asm_code = R"( + { + __asm__ __volatile__( + "mma.sync.aligned.kind::{kind}.block_scale.scale_vec::{scale_vec}.{shape}.{alayout}.{blayout}{.dtype}{.atype}{.btype}{.ctype}.{scale_dtype} " + "{templates};\n" + : {outputs} + : {inputs}); + } +)"; + auto [templates_str, inputs_str, outputs_str] = GetMMABlockScaleOperands(); + + Replacer replacer; + replacer.register_rule("{kind}", info.kind); + replacer.register_rule("{scale_vec}", info.scale_vec); + replacer.register_rule("{shape}", shape); + replacer.register_rule("{alayout}", A_layout); + replacer.register_rule("{blayout}", B_layout); + replacer.register_rule("{.atype}", ptx::DTypeToString(dtype_a)); + replacer.register_rule("{.btype}", ptx::DTypeToString(dtype_b)); + replacer.register_rule("{.ctype}", ptx::DTypeToString(dtype_c)); + replacer.register_rule("{.dtype}", ptx::DTypeToString(dtype_c)); + replacer.register_rule("{scale_dtype}", info.scale_dtype); + replacer.register_rule("{templates}", templates_str); + replacer.register_rule("{outputs}", outputs_str); + replacer.register_rule("{inputs}", inputs_str); + asm_code = replacer.rewrite(asm_code); + replacer.empty_rules(); + replacer.register_rule("{A}", a_ptr + " + " + a_elem_offset); + replacer.register_rule("{B}", b_ptr + " + " + b_elem_offset); + replacer.register_rule("{C}", c_ptr + " + " + c_elem_offset); + replacer.register_rule("{D}", c_ptr + " + " + c_elem_offset); + replacer.register_rule("{SFA}", sfa); + replacer.register_rule("{SFB}", sfb); + replacer.register_rule("{IdA}", id_a); + replacer.register_rule("{IdB}", id_b); + return replacer.rewrite(asm_code); +} + std::string PrintWGMMAAssembly(const std::string &shape, const bool &a_is_k_major, const bool &b_is_k_major, const std::string &A_dtype, diff --git a/src/cuda/codegen/ptx.h b/src/cuda/codegen/ptx.h index 6085734baf..bb2cbe27b9 100644 --- a/src/cuda/codegen/ptx.h +++ b/src/cuda/codegen/ptx.h @@ -86,6 +86,25 @@ std::string DTypeEnumToString(const DataType &dtype); */ std::string DTypeEnumToString(const std::string &dtype); +/*! + * \brief Get the PTX suffix string for a data type, e.g. ".e2m1". + */ +std::string DTypeToString(DataType dtype); + +struct MMABlockScaleInfo { + std::string kind; + std::string scale_vec; + std::string scale_dtype; +}; + +/*! + * \brief Validate and describe an SM120 dense block-scaled MMA configuration. + */ +MMABlockScaleInfo GetMMABlockScaleInfo(DataType dtype_a, DataType dtype_b, + DataType dtype_c, int m, int n, int k, + bool trans_a, bool trans_b, + int scale_vec); + /*! * \brief Parse MMA shape from string. */ @@ -154,6 +173,16 @@ PrintMMAAssembly(const std::string &shape, const std::string &A_layout, const std::string &sparsity_selector, const std::string &bit_op, bool sparse, bool saturate); +std::string PrintMMABlockScaleAssembly( + const std::string &shape, const std::string &A_layout, + const std::string &B_layout, const std::string &A_dtype, + const std::string &B_dtype, const std::string &C_dtype, int scale_vec, + const std::string &a_ptr, const std::string &a_offset, + const std::string &b_ptr, const std::string &b_offset, + const std::string &c_ptr, const std::string &c_offset, + const std::string &sfa, const std::string &sfb, const std::string &id_a, + const std::string &id_b); + /*! * \brief Print WGMMA assembly string given parameters. * \param shape The shape string mMnNkK diff --git a/src/op/builtin.cc b/src/op/builtin.cc index d9924577fd..cb9c68b834 100644 --- a/src/op/builtin.cc +++ b/src/op/builtin.cc @@ -286,6 +286,11 @@ TIR_DEFINE_TL_BUILTIN(ptx_mma_sm70) .set_attr("TCallEffectKind", Integer(CallEffectKind::kOpaque)); +TIR_DEFINE_TL_BUILTIN(ptx_mma_blockscale) + .set_num_inputs(17) + .set_attr("TCallEffectKind", + Integer(CallEffectKind::kOpaque)); + TIR_DEFINE_TL_BUILTIN(ptx_ldmatrix) .set_num_inputs(4) .set_attr("TCallEffectKind", diff --git a/src/op/builtin.h b/src/op/builtin.h index 5cdf7f37d2..38d8fda461 100644 --- a/src/op/builtin.h +++ b/src/op/builtin.h @@ -491,6 +491,26 @@ TVM_DLL const Op &ptx_deallocate_tensor_memory(); */ TVM_DLL const Op &ptx_mma_sm70(); +/*! + * \brief TileLang intrinsic for PTX block-scaled MMA instructions. + * + * SM120 currently supports mxf8f6f4 scale_vec=1, m16n8k32 for all pairs of + * e2m1/e3m2/e2m3/e4m3/e5m2 inputs with ue8m0 scale factors, plus mxf4nvf4 + * scale_vec=2/4, m16n8k64 for e2m1 x e2m1 inputs with ue8m0/ue4m3 scale + * factors. + * + * void ptx_mma_blockscale(StringImm shape, StringImm A_layout, + * StringImm B_layout, StringImm A_dtype, + * StringImm B_dtype, StringImm C_dtype, + * Integer scale_vec, + * Var multiplicand_a, Expr a_index, + * Var multiplicand_b, Expr b_index, + * Var accumulator, Expr c_index, + * Expr scale_a, Expr scale_b, + * Expr scale_id_a, Expr scale_id_b); + */ +TVM_DLL const Op &ptx_mma_blockscale(); + /*! * \brief tvm intrinsics for ldmatrix * diff --git a/src/tl_templates/cuda/instruction/mma_blockscale.h b/src/tl_templates/cuda/instruction/mma_blockscale.h new file mode 100644 index 0000000000..4c5232ef7d --- /dev/null +++ b/src/tl_templates/cuda/instruction/mma_blockscale.h @@ -0,0 +1,275 @@ +#pragma once + +#include "../common.h" +#include +#ifndef __CUDACC_RTC__ +#include +#include +#endif + +namespace tl { + +#ifndef TL_ALWAYS_FALSE_V_DEFINED +#define TL_ALWAYS_FALSE_V_DEFINED +template inline constexpr bool always_false_v = false; +#endif + +namespace detail { + +template struct MmaBlockScaleImplTraits { + using DReg = std::remove_extent_t; + using AReg = std::remove_extent_t; + using BReg = std::remove_extent_t; + using CReg = std::remove_extent_t; + using SFAReg = std::remove_extent_t; + using SFBReg = std::remove_extent_t; + + static constexpr int kDRegs = std::extent_v; + static constexpr int kARegs = std::extent_v; + static constexpr int kBRegs = std::extent_v; + static constexpr int kCRegs = std::extent_v; + static constexpr int kSFARegs = std::extent_v; + static constexpr int kSFBRegs = std::extent_v; +}; + +template +TL_DEVICE void call_fma_blockscale_impl( + typename MmaBlockScaleImplTraits::DReg *d, + const typename MmaBlockScaleImplTraits::AReg *a, + const typename MmaBlockScaleImplTraits::BReg *b, + const typename MmaBlockScaleImplTraits::CReg *c, + const typename MmaBlockScaleImplTraits::SFAReg *sfa, + const typename MmaBlockScaleImplTraits::SFBReg *sfb, uint16_t id_a, + uint16_t id_b, std::index_sequence, std::index_sequence, + std::index_sequence, std::index_sequence, + std::index_sequence, std::index_sequence) { + Impl::fma(d[DIdx]..., a[AIdx]..., b[BIdx]..., c[CIdx]..., sfa[SFAIdx]..., + sfb[SFBIdx]..., id_a, id_b); +} + +template +TL_DEVICE void +call_fma_blockscale(typename MmaBlockScaleImplTraits::DReg *d, + const typename MmaBlockScaleImplTraits::AReg *a, + const typename MmaBlockScaleImplTraits::BReg *b, + const typename MmaBlockScaleImplTraits::CReg *c, + const typename MmaBlockScaleImplTraits::SFAReg *sfa, + const typename MmaBlockScaleImplTraits::SFBReg *sfb, + uint16_t id_a, uint16_t id_b) { + call_fma_blockscale_impl( + d, a, b, c, sfa, sfb, id_a, id_b, + std::make_index_sequence::kDRegs>{}, + std::make_index_sequence::kARegs>{}, + std::make_index_sequence::kBRegs>{}, + std::make_index_sequence::kCRegs>{}, + std::make_index_sequence::kSFARegs>{}, + std::make_index_sequence::kSFBRegs>{}); +} + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1200 +#define TL_SM120_BLOCKSCALE_ASM(KindAsm, ScaleVecAsm, ShapeAsm, AAsm, BAsm, \ + SFAsm, d0, d1, d2, d3, a0, a1, a2, a3, b0, b1, \ + c0, c1, c2, c3, sfa, sfb, id_a, id_b) \ + asm volatile("mma.sync.aligned.kind::" KindAsm \ + ".block_scale.scale_vec::" ScaleVecAsm "." ShapeAsm \ + ".row.col.f32." AAsm "." BAsm ".f32." SFAsm " " \ + "{%0, %1, %2, %3}, " \ + "{%4, %5, %6, %7}, " \ + "{%8, %9}, " \ + "{%10, %11, %12, %13}, " \ + "{%14}, {%15, %16}, " \ + "{%17}, {%18, %19};" \ + : "=f"(d0), "=f"(d1), "=f"(d2), "=f"(d3) \ + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), \ + "f"(c0), "f"(c1), "f"(c2), "f"(c3), "r"(uint32_t(sfa)), \ + "h"((short)0), "h"((short)id_a), "r"(uint32_t(sfb)), \ + "h"((short)0), "h"((short)id_b)); +#else +#define TL_SM120_BLOCKSCALE_ASM(KindAsm, ScaleVecAsm, ShapeAsm, AAsm, BAsm, \ + SFAsm, d0, d1, d2, d3, a0, a1, a2, a3, b0, b1, \ + c0, c1, c2, c3, sfa, sfb, id_a, id_b) \ + CUTE_INVALID_CONTROL_PATH( \ + "Attempting to use SM120 block-scaled MMA without SM120+"); +#endif + +#define TL_DEFINE_SM120_BLOCKSCALE_IMPL( \ + ImplName, KindAsm, ScaleVecAsm, ShapeAsm, AAsm, BAsm, SFAsm, ScaleRegType) \ + struct ImplName { \ + using DRegisters = float[4]; \ + using ARegisters = uint32_t[4]; \ + using BRegisters = uint32_t[2]; \ + using CRegisters = float[4]; \ + using SFARegisters = ScaleRegType[1]; \ + using SFBRegisters = ScaleRegType[1]; \ + \ + static TL_DEVICE void \ + fma(float &d0, float &d1, float &d2, float &d3, uint32_t const &a0, \ + uint32_t const &a1, uint32_t const &a2, uint32_t const &a3, \ + uint32_t const &b0, uint32_t const &b1, float const &c0, \ + float const &c1, float const &c2, float const &c3, \ + ScaleRegType const &sfa, ScaleRegType const &sfb, uint16_t id_a, \ + uint16_t id_b) { \ + TL_SM120_BLOCKSCALE_ASM(KindAsm, ScaleVecAsm, ShapeAsm, AAsm, BAsm, \ + SFAsm, d0, d1, d2, d3, a0, a1, a2, a3, b0, b1, \ + c0, c1, c2, c3, sfa, sfb, id_a, id_b) \ + } \ + }; + +#define TL_SM120_MXF8F6F4_IMPL_NAME_(AName, BName) \ + SM120_16x8x32_F32##AName##BName##F32_UE8M0_TN + +#define TL_SM120_MXF8F6F4_IMPL_NAME_EXPANDED(AName, BName) \ + TL_SM120_MXF8F6F4_IMPL_NAME_(AName, BName) + +#define TL_SM120_MXF8F6F4_IMPL_NAME(ATypeEnum, BTypeEnum) \ + TL_SM120_MXF8F6F4_IMPL_NAME_EXPANDED(TL_BLOCKSCALE_DTYPE_ASM(ATypeEnum), \ + TL_BLOCKSCALE_DTYPE_ASM(BTypeEnum)) + +#define TL_SM120_MXF4NVF4_IMPL_NAME_(ScaleName) \ + SM120_16x8x64_F32_e2m1_e2m1_F32_##ScaleName##_TN + +#define TL_BLOCKSCALE_DTYPE_ASM_kFloat4_e2m1fn e2m1 +#define TL_BLOCKSCALE_DTYPE_ASM_kFloat6_e2m3fn e2m3 +#define TL_BLOCKSCALE_DTYPE_ASM_kFloat6_e3m2fn e3m2 +#define TL_BLOCKSCALE_DTYPE_ASM_kFloat8_e4m3 e4m3 +#define TL_BLOCKSCALE_DTYPE_ASM_kFloat8_e5m2 e5m2 + +#define TL_BLOCKSCALE_DTYPE_ASM_(TypeEnum) TL_BLOCKSCALE_DTYPE_ASM_##TypeEnum +#define TL_BLOCKSCALE_DTYPE_ASM(TypeEnum) TL_BLOCKSCALE_DTYPE_ASM_(TypeEnum) + +#define TL_SM120_BLOCKSCALE_DTYPE_ROWS(M) \ + M(kFloat4_e2m1fn) \ + M(kFloat6_e2m3fn) \ + M(kFloat6_e3m2fn) \ + M(kFloat8_e4m3) \ + M(kFloat8_e5m2) + +#define TL_SM120_BLOCKSCALE_DTYPE_COLS(M, ATypeEnum) \ + M(ATypeEnum, kFloat4_e2m1fn) \ + M(ATypeEnum, kFloat6_e2m3fn) \ + M(ATypeEnum, kFloat6_e3m2fn) \ + M(ATypeEnum, kFloat8_e4m3) \ + M(ATypeEnum, kFloat8_e5m2) + +#define TL_BLOCKSCALE_STRINGIZE_(X) #X +#define TL_BLOCKSCALE_STRINGIZE(X) TL_BLOCKSCALE_STRINGIZE_(X) + +#define TL_DEFINE_SM120_MXF8F6F4_IMPL(ATypeEnum, BTypeEnum) \ + TL_DEFINE_SM120_BLOCKSCALE_IMPL( \ + TL_SM120_MXF8F6F4_IMPL_NAME(ATypeEnum, BTypeEnum), "mxf8f6f4", "1X", \ + "m16n8k32", TL_BLOCKSCALE_STRINGIZE(TL_BLOCKSCALE_DTYPE_ASM(ATypeEnum)), \ + TL_BLOCKSCALE_STRINGIZE(TL_BLOCKSCALE_DTYPE_ASM(BTypeEnum)), "ue8m0", \ + uint8_t) \ + TL_DEFINE_MMA_BLOCKSCALE_DISPATCHER( \ + ATypeEnum, BTypeEnum, kFloat32, 16, 8, 32, false, true, 1, \ + tl::detail::TL_SM120_MXF8F6F4_IMPL_NAME(ATypeEnum, BTypeEnum)) + +#define TL_DEFINE_SM120_MXF4NVF4_IMPL(ScaleVecValue, ScaleVecAsm, ScaleName) \ + TL_DEFINE_SM120_BLOCKSCALE_IMPL(TL_SM120_MXF4NVF4_IMPL_NAME_(ScaleName), \ + "mxf4nvf4", ScaleVecAsm, "m16n8k64", "e2m1", \ + "e2m1", TL_BLOCKSCALE_STRINGIZE(ScaleName), \ + uint32_t) \ + TL_DEFINE_MMA_BLOCKSCALE_DISPATCHER( \ + kFloat4_e2m1fn, kFloat4_e2m1fn, kFloat32, 16, 8, 64, false, true, \ + ScaleVecValue, tl::detail::TL_SM120_MXF4NVF4_IMPL_NAME_(ScaleName)) + +template +struct MmaBlockScaleDispatcher { + using CRegType = void; + using ARegType = void; + using BRegType = void; + using SFARegType = void; + using SFBRegType = void; + + static TL_DEVICE void exec(CRegType *, const ARegType *, const BRegType *, + const CRegType *, const SFARegType *, + const SFBRegType *, uint16_t, uint16_t) { + static_assert(always_false_v>, + "tl::mma_sync_blockscale: unsupported configuration"); + } +}; + +#define TL_DEFINE_MMA_BLOCKSCALE_DISPATCHER( \ + ATypeEnum, BTypeEnum, CTypeEnum, MValue, NValue, KValue, TransAValue, \ + TransBValue, ScaleVecValue, ImplType) \ + template <> \ + struct MmaBlockScaleDispatcher { \ + using Impl = ImplType; \ + using Traits = MmaBlockScaleImplTraits; \ + using CRegType = typename Traits::DReg; \ + using ARegType = typename Traits::AReg; \ + using BRegType = typename Traits::BReg; \ + using SFARegType = typename Traits::SFAReg; \ + using SFBRegType = typename Traits::SFBReg; \ + static_assert( \ + std::is_same_v, \ + "tl::mma_sync_blockscale requires matching accumulator/output regs"); \ + static TL_DEVICE void exec(CRegType *d, const ARegType *a, \ + const BRegType *b, const CRegType *c, \ + const SFARegType *sfa, const SFBRegType *sfb, \ + uint16_t id_a, uint16_t id_b) { \ + call_fma_blockscale(d, a, b, c, sfa, sfb, id_a, id_b); \ + } \ + }; + +#define TL_DEFINE_SM120_MXF8F6F4_ROW(ATypeEnum) \ + TL_SM120_BLOCKSCALE_DTYPE_COLS(TL_DEFINE_SM120_MXF8F6F4_IMPL, ATypeEnum) + +// SM120 dense block-scaled MMA, mxf8f6f4 scale_vec::1X, ue8m0 scale factors. +TL_SM120_BLOCKSCALE_DTYPE_ROWS(TL_DEFINE_SM120_MXF8F6F4_ROW) + +// SM120 dense block-scaled MMA, mxf4nvf4 e2m1 x e2m1. +TL_DEFINE_SM120_MXF4NVF4_IMPL(2, "2X", ue8m0) +TL_DEFINE_SM120_MXF4NVF4_IMPL(4, "4X", ue4m3) + +#undef TL_DEFINE_MMA_BLOCKSCALE_DISPATCHER +#undef TL_DEFINE_SM120_MXF8F6F4_ROW +#undef TL_DEFINE_SM120_MXF8F6F4_IMPL +#undef TL_DEFINE_SM120_MXF4NVF4_IMPL +#undef TL_SM120_MXF4NVF4_IMPL_NAME_ +#undef TL_SM120_MXF8F6F4_IMPL_NAME +#undef TL_SM120_MXF8F6F4_IMPL_NAME_EXPANDED +#undef TL_SM120_MXF8F6F4_IMPL_NAME_ +#undef TL_BLOCKSCALE_STRINGIZE +#undef TL_BLOCKSCALE_STRINGIZE_ +#undef TL_SM120_BLOCKSCALE_DTYPE_COLS +#undef TL_SM120_BLOCKSCALE_DTYPE_ROWS +#undef TL_BLOCKSCALE_DTYPE_ASM +#undef TL_BLOCKSCALE_DTYPE_ASM_ +#undef TL_BLOCKSCALE_DTYPE_ASM_kFloat8_e5m2 +#undef TL_BLOCKSCALE_DTYPE_ASM_kFloat8_e4m3 +#undef TL_BLOCKSCALE_DTYPE_ASM_kFloat6_e3m2fn +#undef TL_BLOCKSCALE_DTYPE_ASM_kFloat6_e2m3fn +#undef TL_BLOCKSCALE_DTYPE_ASM_kFloat4_e2m1fn +#undef TL_DEFINE_SM120_BLOCKSCALE_IMPL +#undef TL_SM120_BLOCKSCALE_ASM + +} // namespace detail + +template +TL_DEVICE void mma_sync_blockscale( + typename detail::MmaBlockScaleDispatcher< + AType, BType, CType, M, N, K, TransA, TransB, ScaleVec>::CRegType *c, + const typename detail::MmaBlockScaleDispatcher< + AType, BType, CType, M, N, K, TransA, TransB, ScaleVec>::ARegType *a, + const typename detail::MmaBlockScaleDispatcher< + AType, BType, CType, M, N, K, TransA, TransB, ScaleVec>::BRegType *b, + typename detail::MmaBlockScaleDispatcher< + AType, BType, CType, M, N, K, TransA, TransB, ScaleVec>::SFARegType sfa, + typename detail::MmaBlockScaleDispatcher< + AType, BType, CType, M, N, K, TransA, TransB, ScaleVec>::SFBRegType sfb, + uint16_t id_a, uint16_t id_b) { + using Dispatcher = + detail::MmaBlockScaleDispatcher; + static_assert(!std::is_void_v, + "tl::mma_sync_blockscale: unsupported configuration"); + Dispatcher::exec(c, a, b, c, &sfa, &sfb, id_a, id_b); +} + +} // namespace tl diff --git a/tilelang/language/__init__.py b/tilelang/language/__init__.py index 33f050240c..6e523c8cd5 100644 --- a/tilelang/language/__init__.py +++ b/tilelang/language/__init__.py @@ -121,6 +121,7 @@ from .builtin import __ffs as __ffs # noqa: F401 from .builtin import ds_read_tr16_b64 as ds_read_tr16_b64 # noqa: F401 from .builtin import ds_read_tr8_b64 as ds_read_tr8_b64 # noqa: F401 +from .builtin import pack_b16 as pack_b16 # noqa: F401 from .builtin import ldg32 as ldg32 # noqa: F401 from .builtin import ldg64 as ldg64 # noqa: F401 from .builtin import ldg128 as ldg128 # noqa: F401 diff --git a/tilelang/language/ast/ir.py b/tilelang/language/ast/ir.py index 04f953c865..0eb16ad81a 100644 --- a/tilelang/language/ast/ir.py +++ b/tilelang/language/ast/ir.py @@ -1882,6 +1882,7 @@ def wrapped(*args, **kwargs): call_pure_extern = _dtype_forward(_tir_op.call_pure_extern) ptx_mma = _dtype_forward(_tir_op.ptx_mma) ptx_mma_sp = _dtype_forward(_tir_op.ptx_mma_sp) +ptx_mma_blockscale = _dtype_forward(_tir_op.ptx_mma_blockscale) ptx_wgmma_ss = _dtype_forward(_tir_op.ptx_wgmma_ss) ptx_wgmma_rs = _dtype_forward(_tir_op.ptx_wgmma_rs) ptx_wgmma_sp_ss = _dtype_forward(_tir_op.ptx_wgmma_sp_ss) @@ -2137,6 +2138,7 @@ def wrapped(*args, **kwargs): "tvm_warp_activemask", "ptx_mma", "ptx_mma_sp", + "ptx_mma_blockscale", "ptx_wgmma_ss", "ptx_wgmma_rs", "ptx_wgmma_sp_ss", diff --git a/tilelang/language/builtin.py b/tilelang/language/builtin.py index cc8fa8c344..fee0ecb273 100644 --- a/tilelang/language/builtin.py +++ b/tilelang/language/builtin.py @@ -1499,6 +1499,11 @@ def ds_read_tr8_b64(src: BufferLikeType) -> PrimExpr: return tirx.call_intrin("uint32x2", tirx.op.Op.get("tl.ds_read_tr8_b64"), ptr) +def pack_b16(v0: PrimExpr, v1: PrimExpr) -> PrimExpr: + """Pack two b16 values into one uint32 lane.""" + return tirx.call_intrin("uint32", tirx.op.Op.get("tl.pack_b16"), v0, v1) + + def ldg32(src: BufferLikeType, pred: PrimExpr = None) -> PrimExpr: """Load 32 bits (4 bytes) from global memory using explicit PTX instructions. diff --git a/tilelang/language/tir/ir.py b/tilelang/language/tir/ir.py index 74d2af4a9a..cf3caabf1a 100644 --- a/tilelang/language/tir/ir.py +++ b/tilelang/language/tir/ir.py @@ -289,6 +289,7 @@ def wrapped(*args, **kwargs): call_pure_extern = _dtype_forward(_tir_op.call_pure_extern) ptx_mma = _dtype_forward(_tir_op.ptx_mma) ptx_mma_sp = _dtype_forward(_tir_op.ptx_mma_sp) +ptx_mma_blockscale = _dtype_forward(_tir_op.ptx_mma_blockscale) ptx_wgmma_ss = _dtype_forward(_tir_op.ptx_wgmma_ss) ptx_wgmma_rs = _dtype_forward(_tir_op.ptx_wgmma_rs) ptx_wgmma_sp_ss = _dtype_forward(_tir_op.ptx_wgmma_sp_ss) diff --git a/tilelang/language/tir/op.py b/tilelang/language/tir/op.py index bd3053dd26..1c7684e2cb 100644 --- a/tilelang/language/tir/op.py +++ b/tilelang/language/tir/op.py @@ -1076,6 +1076,60 @@ def ptx_mma_sp( ) +def ptx_mma_blockscale( + C_dtype, + shape, + A_layout, + B_layout, + A_dtype, + B_dtype, + scale_vec, + multiplicand_a, + a_index, + multiplicand_b, + b_index, + accumulator, + c_index, + scale_a, + scale_b, + scale_id_a, + scale_id_b, +): + """TileLang intrinsic for PTX block-scaled MMA instructions. + + ``C_dtype`` is the accumulator/result precision, mirroring the + ``accumulator`` operand. The intrinsic itself is side-effect only and emits + a ``handle``-typed TIR call, like TCGEN05 MMA intrinsics. + The CUDA backend maps the shape/layout/dtype tuple to the matching + ``tl::mma_sync_blockscale`` dispatcher specialization. On SM120, + ``scale_vec=1`` selects ``mxf8f6f4``/``ue8m0`` ``m16n8k32`` variants for + e2m1/e3m2/e2m3/e4m3/e5m2 inputs, while ``scale_vec=2`` and ``scale_vec=4`` + select ``mxf4nvf4`` ``m16n8k64`` e2m1 inputs with ``ue8m0`` and ``ue4m3`` + scale factors respectively. + """ + return call_intrin( + "handle", + _tvm_op.Op.get("tl.ptx_mma_blockscale"), + shape, + A_layout, + B_layout, + A_dtype, + B_dtype, + C_dtype, + scale_vec, + multiplicand_a, + a_index, + multiplicand_b, + b_index, + accumulator, + c_index, + scale_a, + scale_b, + scale_id_a, + scale_id_b, + ) + + def ptx_wgmma_ss( dtype, wgmma_prefix, From 3b4c1c949ca5caf9f107b111a64ff62d3cad0e63 Mon Sep 17 00:00:00 2001 From: sepcnt <30561671+sepcnt@users.noreply.github.com> Date: Thu, 4 Jun 2026 22:29:15 +0800 Subject: [PATCH 2/6] Add Sage Attention 3 and relevant support --- examples/sage_attention_sm120/__init__.py | 5 + .../sage_attention_sm120/sageattn3_fp4.py | 700 ++++++++++++++++++ examples/sage_attention_sm120/tir_helpers.py | 111 +++ src/cuda/codegen/codegen_cuda.cc | 41 +- src/cuda/op/copy.cc | 415 ++++++++--- src/cuda/op/copy_analysis.cc | 51 +- .../annotate_warp_group_reg_alloc.cc | 31 + src/cuda/transform/lower_hopper_intrin.cc | 21 +- src/layout/gemm_layouts.cc | 38 + src/layout/layout.cc | 27 +- src/layout/layout.h | 1 + src/op/builtin.cc | 5 + src/op/builtin.h | 18 + src/op/utils.h | 11 + src/tl_templates/cuda/copy.h | 41 + src/transform/layout_inference.cc | 25 + src/transform/lower_tile_op.cc | 33 + .../merge_shared_memory_allocations.cc | 114 +-- src/transform/storage_rewrite.cc | 23 +- src/transform/thread_storage_sync.cc | 367 +++++++-- .../language/test_tilelang_language_ldg.py | 27 + .../test_tilelang_language_tma_copy.py | 244 +++++- ...st_tilelang_language_tma_gather_scatter.py | 35 + .../language/test_tilelang_language_view.py | 177 +++++ ...tilelang_transform_disable_memory_reuse.py | 59 +- tilelang/language/__init__.py | 2 + tilelang/language/allocate.py | 44 +- tilelang/language/annotations.py | 17 + tilelang/language/builtin.py | 15 + tilelang/language/copy_op.py | 46 +- tilelang/language/customize.py | 18 +- tilelang/language/proxy.py | 23 +- tilelang/layout/__init__.py | 1 + tilelang/layout/swizzle.py | 11 + 34 files changed, 2522 insertions(+), 275 deletions(-) create mode 100644 examples/sage_attention_sm120/__init__.py create mode 100644 examples/sage_attention_sm120/sageattn3_fp4.py create mode 100644 examples/sage_attention_sm120/tir_helpers.py diff --git a/examples/sage_attention_sm120/__init__.py b/examples/sage_attention_sm120/__init__.py new file mode 100644 index 0000000000..fdd81814ff --- /dev/null +++ b/examples/sage_attention_sm120/__init__.py @@ -0,0 +1,5 @@ +"""SM120 SageAttention3 TileLang examples.""" + +from examples.sage_attention_sm120.sageattn3_fp4 import sage3_packed_fp4_attention_raw_kernel + +__all__ = ["sage3_packed_fp4_attention_raw_kernel"] diff --git a/examples/sage_attention_sm120/sageattn3_fp4.py b/examples/sage_attention_sm120/sageattn3_fp4.py new file mode 100644 index 0000000000..b9dc6c3d24 --- /dev/null +++ b/examples/sage_attention_sm120/sageattn3_fp4.py @@ -0,0 +1,700 @@ +"""SageAttention3 FP4 raw-core TileLang kernel for NVIDIA SM120.""" + +from __future__ import annotations + +import tilelang +import tilelang.language as T +from tilelang.carver.arch import driver + +from examples.sage_attention_sm120.tir_helpers import ( + mma_m16n32k64_blockscale_f32, + pack_cast2_u32, +) + + +_LOG2_E = 1.4426950408889634 +_F32_NEG_INF = -3.4028234663852886e38 +_SAGE3_PV_INV_SCALE_LOG2 = -11.392317422778762 # log2(1 / (448 * 6)) +_SAGE3_FP4_INV_SCALE_LOG2 = -2.584962500721156 # log2(1 / 6) + + +@T.macro +def _sage3_load_q_scale_word_shared(SFQ_words, row, ko): + row_local = row % 64 + return T.lds32(SFQ_words[ko, row // 64, row_local % 16, row_local // 16]) + + +@T.macro +def _sage3_load_k_scale_word_shared_stage(SFK_words, stage, physical_col, ko): + physical_local = physical_col % 64 + return T.lds32(SFK_words[stage, ko, physical_col // 64, physical_local % 16, physical_local // 16]) + + +@T.macro +def _sage3_pack_fp4_pair_byte(v0, v1): + return pack_cast2_u32(T.float4_e2m1fn, T.uint8, v0, v1) & T.Cast("uint32", 0xFF) + + +@T.macro +def _sage3_load_v_scale_word_stage(SFV_words, stage, dim, ko): + dim_local = dim % 64 + return T.lds32(SFV_words[stage, ko, dim // 64, dim_local % 16, dim_local // 16]) + + +@T.macro +def _sage3_ldmatrix_x4_u8_q_rowmajor_qregs(tile, row_tile, ko, q_regs, q_regs_ko, q_regs_mi, lane): + matrix_group = lane >> 3 + row_in_matrix = lane & 0x7 + half = matrix_group >> 1 + row_pair = matrix_group & 0x1 + row = row_tile * 16 + row_pair * 8 + row_in_matrix + col = ko * 64 + half * 32 + T.ptx_ldmatrix( + T.bool(False), + 4, + T.access_ptr(tile[row, col], "r", extent=32), + T.access_ptr(q_regs[q_regs_ko, q_regs_mi, 0], "w", extent=4), + ) + + +@T.macro +def _sage3_ldmatrix_x4_u8_b_rowmajor_stage_into(tile, stage, n32, ko, tile_half, regs, reg_offset, lane): + matrix_group = lane >> 3 + row_in_matrix = lane & 0x7 + reg_word = tile_half * 4 + matrix_group + e = reg_word >> 1 + half = reg_word & 0x1 + row = n32 * 32 + e * 8 + row_in_matrix + col = ko * 64 + half * 32 + T.ptx_ldmatrix( + T.bool(False), + 4, + T.access_ptr(tile[stage, row, col], "r", extent=32), + T.access_ptr(regs[reg_offset], "w", extent=4), + ) + + +@tilelang.jit +def sage3_packed_fp4_attention_raw_kernel( + query_tokens: int, + kv_tokens: int, + valid_k_tokens: int, + q_heads: int, + kv_heads: int, + head_dim: int, + *, + dtype: str = "bfloat16", + block_n: int = 128, +): + """SageAttention3 raw-core ABI implemented in pure TileLang. + + External inputs intentionally mirror ``fp4attn_cuda.fwd``: + + - ``Q``/``K``: uint8 packed FP4, shape ``[1,H,N,D/2]``. + - ``Vt``: uint8 packed transposed FP4, shape ``[1,H,D,N/2]``. + - ``SFQ``/``SFK``/``SFVt``: e4m3 scale tensors with Sage3's physical + swizzled layout but the same public 4D shapes. + - ``DeltaS``: float32 block-mean correction, shape + ``[1,H,ceil(query_tokens/128),kv_tokens]``. + - ``Out``/``LSE``: the raw-core padded output tensors returned by + ``fp4attn_cuda.fwd``. + + This version keeps the raw input/output contract and Sage3 scale/K layouts, + uses SM120 register FP4 blockscale MMA for both QK and PV, and has one + load path: packed FP4 global views are TMA-loaded into CUTLASS-compatible + packed FP4 shared memory and consumed with LDSM. + + Like the official kernel, the QK accumulator stays in Sage3's permuted + ("physical") K-token order end to end: K rows are LDSM-loaded row-major + (conflict-free), DeltaS seeds the accumulator using the physical->logical + column mapping, and each thread then owns 8 logically-consecutive P + columns so FP4 quantization of P is thread-local with no cross-lane + shuffles. + """ + if query_tokens <= 0 or kv_tokens <= 0 or valid_k_tokens <= 0: + raise ValueError(f"query_tokens, kv_tokens, and valid_k_tokens must be positive; got {query_tokens}, {kv_tokens}, {valid_k_tokens}") + if valid_k_tokens > kv_tokens: + raise ValueError(f"valid_k_tokens must be <= kv_tokens, got {valid_k_tokens} > {kv_tokens}") + if q_heads <= 0 or kv_heads <= 0 or head_dim <= 0: + raise ValueError(f"q_heads, kv_heads, and head_dim must be positive; got {q_heads}, {kv_heads}, {head_dim}") + if q_heads != kv_heads: + raise ValueError(f"SageAttn3 raw core currently requires q_heads == kv_heads, got {q_heads}, {kv_heads}") + if head_dim != 128: + raise ValueError(f"sage3_packed_fp4_attention_raw fixes head_dim=128, got {head_dim}") + if dtype != "bfloat16": + raise ValueError(f"sage3_packed_fp4_attention_raw currently supports bfloat16 output, got {dtype}") + if block_n != 128: + raise ValueError(f"sage3_packed_fp4_attention_raw fixes block_n=128, got {block_n}") + if query_tokens % 128 != 0 or kv_tokens % 128 != 0: + raise ValueError( + f"SageAttn3 raw inputs must be preprocessed/padded to 128-token blocks; got query_tokens={query_tokens}, kv_tokens={kv_tokens}" + ) + if kv_tokens % 16 != 0: + raise ValueError(f"kv_tokens must be a multiple of 16 for V scale groups, got {kv_tokens}") + + num_qh = q_heads + num_kvh = kv_heads + q_tokens = query_tokens + k_tokens = kv_tokens + k_valid = valid_k_tokens + block_M = 128 + block_N = block_n + q_blocks = T.ceildiv(q_tokens, block_M) + kv_blocks = T.ceildiv(k_valid, block_N) + kv_stage0_blocks = (kv_blocks + 1) // 2 + kv_stage1_blocks = kv_blocks // 2 + qk_scale_cols = head_dim // 16 + v_scale_cols = k_tokens // 16 + scale = (1.0 / head_dim) ** 0.5 * _LOG2_E + accum_dtype = "float32" + has_k_tail = (k_valid % block_N) != 0 + full_kv_blocks = k_valid // block_N + compute_threads = 256 + warp_count = compute_threads // 32 + warp_M_tiles = 8 // warp_count + warp_M = warp_M_tiles * 16 + warp_N16_tiles = 8 + warp_N32_tiles = 4 + qk_acc_count = warp_M_tiles * warp_N32_tiles * 16 + launch_sms = driver.get_num_sms() + if launch_sms <= 0: + raise ValueError(f"driver.get_num_sms() must be positive, got {launch_sms}") + + @T.prim_func + def main( + Q: T.Tensor((1, num_qh, q_tokens, 64), "uint8"), + K: T.Tensor((1, num_kvh, k_tokens, 64), "uint8"), + Vt: T.Tensor((1, num_kvh, 128, k_tokens // 2), "uint8"), + SFQ: T.Tensor((1, num_qh, q_tokens, qk_scale_cols), "float8_e4m3fn"), + SFK: T.Tensor((1, num_kvh, k_tokens, qk_scale_cols), "float8_e4m3fn"), + SFVt: T.Tensor((1, num_kvh, 128, v_scale_cols), "float8_e4m3fn"), + DeltaS: T.Tensor((1, num_qh, q_blocks, k_tokens), "float32"), + Out: T.Tensor((1, num_qh, q_tokens, 128), "bfloat16"), + LSE: T.Tensor((1, num_qh, q_tokens), "float32"), + ): + T.annotate_compile_flags(["--use_fast_math"]) + Q_fp4 = T.view(Q, (1, num_qh, q_tokens, 128), dtype=T.float4_e2m1fn) + K_fp4 = T.view(K, (1, num_kvh, k_tokens, 128), dtype=T.float4_e2m1fn) + Vt_fp4 = T.view(Vt, (1, num_kvh, 128, k_tokens), dtype=T.float4_e2m1fn) + SFQ_tma = T.view( + SFQ, + (num_qh, qk_scale_cols // 4, q_tokens // 64, 16, 8), + dtype=T.uint16, + strides=( + (q_tokens // 64) * ((qk_scale_cols // 4) * 16 * 8), + 16 * 8, + (qk_scale_cols // 4) * 16 * 8, + 8, + 1, + ), + ) + SFK_tma = T.view( + SFK, + (num_kvh, qk_scale_cols // 4, k_tokens // 64, 16, 8), + dtype=T.uint16, + strides=( + (k_tokens // 64) * ((qk_scale_cols // 4) * 16 * 8), + 16 * 8, + (qk_scale_cols // 4) * 16 * 8, + 8, + 1, + ), + ) + SFV_tma = T.view( + SFVt, + (num_kvh, v_scale_cols // 4, 128 // 64, 16, 8), + dtype=T.uint16, + strides=( + (128 // 64) * ((v_scale_cols // 4) * 16 * 8), + 16 * 8, + (v_scale_cols // 4) * 16 * 8, + 8, + 1, + ), + ) + + with T.Kernel(launch_sms, threads=384) as block_id: + Q_shared = T.alloc_shared((block_M, 128), T.float4_e2m1fn) + T.annotate_layout({Q_shared: tilelang.layout.make_sm120_fp4_smem_layout(Q_shared)}) + SFQ_shared_storage = T.alloc_shared((qk_scale_cols // 4, block_M // 64, 16, 4, 4), "uint8") + SFQ_shared_tma = T.view(SFQ_shared_storage, (qk_scale_cols // 4, block_M // 64, 16, 8), dtype=T.uint16) + SFQ_shared_words = T.view(SFQ_shared_storage, (qk_scale_cols // 4, block_M // 64, 16, 4), dtype=T.uint32) + K_shared = T.alloc_shared((2, block_N, 128), T.float4_e2m1fn) + T.annotate_layout({K_shared: tilelang.layout.make_sm120_fp4_smem_layout(K_shared)}) + SFK_shared_storage = T.alloc_shared((2, qk_scale_cols // 4, block_N // 64, 16, 4, 4), "uint8") + SFK_shared_tma = T.view( + SFK_shared_storage, + (2, qk_scale_cols // 4, block_N // 64, 16, 8), + dtype=T.uint16, + ) + SFK_shared_words = T.view( + SFK_shared_storage, + (2, qk_scale_cols // 4, block_N // 64, 16, 4), + dtype=T.uint32, + ) + DS_shared = T.alloc_shared((2, block_N), "float32") + Vt_shared = T.alloc_shared((2, 128, block_N), T.float4_e2m1fn) + T.annotate_layout({Vt_shared: tilelang.layout.make_sm120_fp4_smem_layout(Vt_shared)}) + SFV_shared_storage = T.alloc_shared((2, (block_N // 16) // 4, 128 // 64, 16, 4, 4), "uint8") + SFV_shared_tma = T.view( + SFV_shared_storage, + (2, (block_N // 16) // 4, 128 // 64, 16, 8), + dtype=T.uint16, + ) + SFV_shared_words = T.view( + SFV_shared_storage, + (2, (block_N // 16) // 4, 128 // 64, 16, 4), + dtype=T.uint32, + ) + O_shared = T.alloc_shared((block_M, 128), "bfloat16") + q_ready = T.alloc_barrier(32) + q_empty = T.alloc_barrier(256) + kv_full = T.alloc_barrier([32, 32]) + kv_empty = T.alloc_barrier([256, 256]) + o_ready = T.alloc_barrier(256) + o_empty = T.alloc_barrier(32) + tx = T.get_thread_binding() + producer_tx = tx - 256 + producer_warp_role = producer_tx >> 5 + consumer_tx = tx + lane = consumer_tx & 31 + warp = consumer_tx >> 5 + g = lane >> 2 + sublane = lane & 0x3 + scale_lane = consumer_tx & 31 + scale_warp = consumer_tx >> 5 + scale_g = scale_lane >> 2 + scale_sublane = scale_lane & 0x3 + with T.ws(2): + T.annotate_producer_reg_dealloc(24) + if producer_warp_role == 0: + for qb, qh in T.Persistent( + [q_blocks, num_qh], + launch_sms, + block_id, + group_size=1, + ): + tile_linear = qh * q_blocks + qb + tile_iter = (tile_linear - block_id) // launch_sms + tile_phase = tile_iter & 1 + + T.mbarrier_wait_parity(q_empty, tile_phase ^ 1) + T.tma_copy( + Q_fp4[0, qh, qb * block_M : (qb + 1) * block_M, 0:128], + Q_shared, + barrier=q_ready, + leader_scope_threads=32, + ) + T.tma_copy( + SFQ_tma[ + qh, + 0 : qk_scale_cols // 4, + qb * (block_M // 64) : (qb + 1) * (block_M // 64), + 0:16, + 0:8, + ], + SFQ_shared_tma, + barrier=q_ready, + leader_scope_threads=32, + ) + T.mbarrier_arrive(q_ready) + + for kb in T.unroll(kv_blocks, unroll_factor=1): + stage = kb & 1 + stage_phase0 = (tile_iter * kv_stage0_blocks + (kb >> 1)) & 1 + stage_phase1 = (tile_iter * kv_stage1_blocks + (kb >> 1)) & 1 + phase = T.if_then_else(stage == 0, stage_phase0, stage_phase1) + T.mbarrier_wait_parity(kv_empty[stage], phase ^ 1) + T.tma_copy( + DeltaS[0, qh, qb, kb * block_N : (kb + 1) * block_N], + DS_shared[stage, 0:block_N], + barrier=kv_full[stage], + leader_scope_threads=32, + ) + T.tma_copy( + SFK_tma[ + qh, + 0 : qk_scale_cols // 4, + kb * (block_N // 64) : (kb + 1) * (block_N // 64), + 0:16, + 0:8, + ], + SFK_shared_tma[stage, 0 : qk_scale_cols // 4, 0 : block_N // 64, 0:16, 0:8], + barrier=kv_full[stage], + leader_scope_threads=32, + ) + T.tma_copy( + K_fp4[0, qh, kb * block_N : (kb + 1) * block_N, 0:128], + K_shared[stage, 0:block_N, 0:128], + barrier=kv_full[stage], + leader_scope_threads=32, + ) + T.tma_copy( + SFV_tma[ + qh, + kb * ((block_N // 16) // 4) : (kb + 1) * ((block_N // 16) // 4), + 0 : 128 // 64, + 0:16, + 0:8, + ], + SFV_shared_tma[stage, 0 : (block_N // 16) // 4, 0 : 128 // 64, 0:16, 0:8], + barrier=kv_full[stage], + leader_scope_threads=32, + ) + T.tma_copy( + Vt_fp4[0, qh, 0:128, kb * block_N : (kb + 1) * block_N], + Vt_shared[stage, 0:128, 0:block_N], + barrier=kv_full[stage], + leader_scope_threads=32, + ) + T.mbarrier_arrive(kv_full[stage]) + elif producer_warp_role == 1: + for qb, qh in T.Persistent( + [q_blocks, num_qh], + launch_sms, + block_id, + group_size=1, + ): + tile_linear = qh * q_blocks + qb + tile_iter = (tile_linear - block_id) // launch_sms + tile_phase = tile_iter & 1 + + T.mbarrier_wait_parity(o_ready, tile_phase) + T.tma_copy( + O_shared, + Out[0, qh, qb * block_M : (qb + 1) * block_M, 0:128], + leader_scope_threads=32, + ) + T.tma_store_wait() + T.mbarrier_arrive(o_empty) + + with T.ws(0, 1): + T.annotate_consumer_reg_alloc(240) + q_regs = T.alloc_local((2, warp_M_tiles, 4), "uint32", role_scoped=True) + a_regs = T.alloc_local((warp_M_tiles, 4), "uint32", role_scoped=True) + b_regs = T.alloc_local((8,), "uint32", role_scoped=True) + q_scale_frag = T.alloc_fragment((2 * 4,), T.float8_e4m3fn, role_scoped=True) + b_scale_frag = T.alloc_fragment((4 * 4,), T.float8_e4m3fn, role_scoped=True) + q_scale_regs = T.view(q_scale_frag, (2,), dtype=T.uint32) + b_scale_regs = T.view(b_scale_frag, (4,), dtype=T.uint32) + qk_acc = T.alloc_local((qk_acc_count,), accum_dtype, role_scoped=True) + pv_acc = T.alloc_local((qk_acc_count,), accum_dtype, role_scoped=True) + o_acc = T.alloc_local((qk_acc_count,), accum_dtype, role_scoped=True) + row_scale_acc = T.alloc_local((warp_M_tiles, 2), accum_dtype, role_scoped=True) + row_max_acc = T.alloc_local((warp_M_tiles, 2), accum_dtype, role_scoped=True) + row_sum_acc = T.alloc_local((warp_M_tiles, 2), accum_dtype, role_scoped=True) + scale_f = T.Cast(accum_dtype, scale) + for qb, qh in T.Persistent( + [q_blocks, num_qh], + launch_sms, + block_id, + group_size=1, + ): + tile_linear = qh * q_blocks + qb + tile_iter = (tile_linear - block_id) // launch_sms + tile_phase = tile_iter & 1 + + T.mbarrier_wait_parity(q_ready, tile_phase) + T.clear(o_acc) + for mi in T.unroll(warp_M_tiles): + row_max_acc[mi, 0] = T.Cast(accum_dtype, _F32_NEG_INF) + row_max_acc[mi, 1] = T.Cast(accum_dtype, _F32_NEG_INF) + row_sum_acc[mi, 0] = T.Cast(accum_dtype, 0.0) + row_sum_acc[mi, 1] = T.Cast(accum_dtype, 0.0) + + for ko in T.unroll(2): + scale_slot = scale_lane & 0x3 + scale_row_pair = scale_slot & 0x1 + scale_row = scale_warp * warp_M + scale_g + scale_row_pair * 8 + q_scale_regs[ko] = _sage3_load_q_scale_word_shared(SFQ_shared_words, scale_row, ko) + + for mi in T.unroll(warp_M_tiles): + _sage3_ldmatrix_x4_u8_q_rowmajor_qregs( + Q_shared, + warp * warp_M_tiles + mi, + ko, + q_regs, + ko, + mi, + lane, + ) + + T.mbarrier_arrive(q_empty) + + for kb in T.serial(kv_blocks): + stage = kb & 1 + stage_phase0 = (tile_iter * kv_stage0_blocks + (kb >> 1)) & 1 + stage_phase1 = (tile_iter * kv_stage1_blocks + (kb >> 1)) & 1 + phase = T.if_then_else(stage == 0, stage_phase0, stage_phase1) + T.mbarrier_wait_parity(kv_full[stage], phase) + + # The QK accumulator stays in Sage3's physical (permuted) + # column order: thread `sublane` holds logical tokens + # nj*32 + sublane*8 .. +7 contiguously. DeltaS seeds the + # accumulator before the MMA, exactly like the official + # mainloop. + ds_vals = T.alloc_local((8,), accum_dtype, role_scoped=True) + for nj in T.unroll(warp_N32_tiles): + if has_k_tail: + if kb < full_kv_blocks: + for jj in T.unroll(8): + ds_vals[jj] = DS_shared[stage, nj * 32 + sublane * 8 + jj] + else: + for jj in T.unroll(8): + ds_vals[jj] = T.if_then_else( + kb * block_N + nj * 32 + sublane * 8 + jj < k_valid, + DS_shared[stage, nj * 32 + sublane * 8 + jj], + T.Cast(accum_dtype, _F32_NEG_INF), + ) + else: + for jj in T.unroll(8): + ds_vals[jj] = DS_shared[stage, nj * 32 + sublane * 8 + jj] + for mi in T.unroll(warp_M_tiles): + off = (mi * warp_N32_tiles + nj) * 16 + for jj in T.unroll(8): + qk_acc[off + (jj // 2) * 4 + (jj % 2)] = ds_vals[jj] + qk_acc[off + (jj // 2) * 4 + (jj % 2) + 2] = ds_vals[jj] + + for ko in T.unroll(2): + for p in T.unroll(4): + physical_col = p * 32 + scale_sublane * 8 + scale_g + b_scale_regs[p] = _sage3_load_k_scale_word_shared_stage(SFK_shared_words, stage, physical_col, ko) + + for nj in T.unroll(warp_N32_tiles): + _sage3_ldmatrix_x4_u8_b_rowmajor_stage_into(K_shared, stage, nj, ko, 0, b_regs, 0, lane) + _sage3_ldmatrix_x4_u8_b_rowmajor_stage_into(K_shared, stage, nj, ko, 1, b_regs, 4, lane) + for mi in T.unroll(warp_M_tiles): + mma_m16n32k64_blockscale_f32( + q_regs.data, + (ko * warp_M_tiles + mi) * 8, + b_regs.data, + 0, + qk_acc, + (mi * warp_N32_tiles + nj) * 16, + q_scale_regs[ko], + b_scale_regs[nj], + mi, + 0, + ) + + # Per-16-token scale-group maxes are taken on the raw + # accumulator (thread's 8 values plus the sublane^1 + # partner) and transformed analytically after exp2, so + # softmax needs a single pass like the official kernel. + absmax0 = T.alloc_local((warp_M_tiles, warp_N32_tiles), accum_dtype, role_scoped=True) + absmax1 = T.alloc_local((warp_M_tiles, warp_N32_tiles), accum_dtype, role_scoped=True) + for mi in T.unroll(warp_M_tiles): + for nj in T.unroll(warp_N32_tiles): + off = (mi * warp_N32_tiles + nj) * 16 + m0 = T.max( + T.max(T.max(qk_acc[off + 0], qk_acc[off + 1]), T.max(qk_acc[off + 4], qk_acc[off + 5])), + T.max(T.max(qk_acc[off + 8], qk_acc[off + 9]), T.max(qk_acc[off + 12], qk_acc[off + 13])), + ) + m1 = T.max( + T.max(T.max(qk_acc[off + 2], qk_acc[off + 3]), T.max(qk_acc[off + 6], qk_acc[off + 7])), + T.max(T.max(qk_acc[off + 10], qk_acc[off + 11]), T.max(qk_acc[off + 14], qk_acc[off + 15])), + ) + absmax0[mi, nj] = T.max(m0, T.shfl_xor(m0, 1, width=4)) + absmax1[mi, nj] = T.max(m1, T.shfl_xor(m1, 1, width=4)) + max0_local = T.max(T.max(absmax0[mi, 0], absmax0[mi, 1]), T.max(absmax0[mi, 2], absmax0[mi, 3])) + max1_local = T.max(T.max(absmax1[mi, 0], absmax1[mi, 1]), T.max(absmax1[mi, 2], absmax1[mi, 3])) + max0_group = T.max(max0_local, T.shfl_xor(max0_local, 2, width=4)) + max1_group = T.max(max1_local, T.shfl_xor(max1_local, 2, width=4)) + prev0 = row_max_acc[mi, 0] + prev1 = row_max_acc[mi, 1] + new0 = T.max(prev0, max0_group) + new1 = T.max(prev1, max1_group) + row_scale_acc[mi, 0] = T.exp2((prev0 - new0) * scale_f) + row_scale_acc[mi, 1] = T.exp2((prev1 - new1) * scale_f) + row_max_acc[mi, 0] = new0 + row_max_acc[mi, 1] = new1 + + for mi in T.unroll(warp_M_tiles): + max_scaled0 = T.alloc_var( + accum_dtype, + init=row_max_acc[mi, 0] * scale_f + T.Cast(accum_dtype, _SAGE3_PV_INV_SCALE_LOG2), + role_scoped=True, + ) + max_scaled1 = T.alloc_var( + accum_dtype, + init=row_max_acc[mi, 1] * scale_f + T.Cast(accum_dtype, _SAGE3_PV_INV_SCALE_LOG2), + role_scoped=True, + ) + sum0_local = T.alloc_var(accum_dtype, init=T.Cast(accum_dtype, 0.0), role_scoped=True) + sum1_local = T.alloc_var(accum_dtype, init=T.Cast(accum_dtype, 0.0), role_scoped=True) + for nj in T.unroll(warp_N32_tiles): + off = (mi * warp_N32_tiles + nj) * 16 + for aa in T.unroll(4): + atom = off + aa * 4 + p00 = T.exp2(qk_acc[atom + 0] * scale_f - max_scaled0) + p01 = T.exp2(qk_acc[atom + 1] * scale_f - max_scaled0) + p10 = T.exp2(qk_acc[atom + 2] * scale_f - max_scaled1) + p11 = T.exp2(qk_acc[atom + 3] * scale_f - max_scaled1) + qk_acc[atom + 0] = p00 + qk_acc[atom + 1] = p01 + qk_acc[atom + 2] = p10 + qk_acc[atom + 3] = p11 + sum0_local += p00 + p01 + sum1_local += p10 + p11 + absmax0[mi, nj] = T.exp2( + absmax0[mi, nj] * scale_f - max_scaled0 + T.Cast(accum_dtype, _SAGE3_FP4_INV_SCALE_LOG2) + ) + absmax1[mi, nj] = T.exp2( + absmax1[mi, nj] * scale_f - max_scaled1 + T.Cast(accum_dtype, _SAGE3_FP4_INV_SCALE_LOG2) + ) + row_sum_acc[mi, 0] = row_sum_acc[mi, 0] * row_scale_acc[mi, 0] + sum0_local + row_sum_acc[mi, 1] = row_sum_acc[mi, 1] * row_scale_acc[mi, 1] + sum1_local + + T.clear(pv_acc) + + for ko in T.unroll(2): + # SFP word assembly: each thread owns the e4m3 scales + # for k-groups ko*4 + (sublane>>1) and +2; the missing + # byte lanes come from the sublane^2 partner. + sel_lo = T.Select((sublane & 1) != 0, absmax1[0, ko * 2], absmax0[0, ko * 2]) + sel_hi = T.Select((sublane & 1) != 0, absmax1[0, ko * 2 + 1], absmax0[0, ko * 2 + 1]) + sfp_pair = pack_cast2_u32(T.float8_e4m3fn, T.uint16, sel_lo, sel_hi) + sfp_own = T.alloc_var( + "uint32", + init=((sfp_pair & T.Cast("uint32", 0xFF)) | ((sfp_pair & T.Cast("uint32", 0xFF00)) << T.Cast("uint32", 8))) + << T.Cast("uint32", (sublane >> 1) * 8), + role_scoped=True, + ) + scale_a = T.alloc_var( + "uint32", + init=sfp_own | T.shfl_xor(sfp_own, 2, width=4), + role_scoped=True, + ) + + for mi in T.unroll(warp_M_tiles): + for half in T.unroll(2): + nj_pack = ko * 2 + half + off = (mi * warp_N32_tiles + nj_pack) * 16 + inv_pmax0 = T.alloc_var( + "float32", + init=T.Cast("float32", 1.0) / absmax0[mi, nj_pack], + role_scoped=True, + ) + inv_pmax1 = T.alloc_var( + "float32", + init=T.Cast("float32", 1.0) / absmax1[mi, nj_pack], + role_scoped=True, + ) + a_regs[mi, half * 2 + 0] = ( + _sage3_pack_fp4_pair_byte(qk_acc[off + 0] * inv_pmax0, qk_acc[off + 1] * inv_pmax0) + | ( + _sage3_pack_fp4_pair_byte(qk_acc[off + 4] * inv_pmax0, qk_acc[off + 5] * inv_pmax0) + << T.Cast("uint32", 8) + ) + | ( + _sage3_pack_fp4_pair_byte(qk_acc[off + 8] * inv_pmax0, qk_acc[off + 9] * inv_pmax0) + << T.Cast("uint32", 16) + ) + | ( + _sage3_pack_fp4_pair_byte(qk_acc[off + 12] * inv_pmax0, qk_acc[off + 13] * inv_pmax0) + << T.Cast("uint32", 24) + ) + ) + a_regs[mi, half * 2 + 1] = ( + _sage3_pack_fp4_pair_byte(qk_acc[off + 2] * inv_pmax1, qk_acc[off + 3] * inv_pmax1) + | ( + _sage3_pack_fp4_pair_byte(qk_acc[off + 6] * inv_pmax1, qk_acc[off + 7] * inv_pmax1) + << T.Cast("uint32", 8) + ) + | ( + _sage3_pack_fp4_pair_byte(qk_acc[off + 10] * inv_pmax1, qk_acc[off + 11] * inv_pmax1) + << T.Cast("uint32", 16) + ) + | ( + _sage3_pack_fp4_pair_byte(qk_acc[off + 14] * inv_pmax1, qk_acc[off + 15] * inv_pmax1) + << T.Cast("uint32", 24) + ) + ) + + for p in T.unroll(4): + logical_dim = p * 32 + scale_sublane * 8 + scale_g + b_scale_regs[p] = _sage3_load_v_scale_word_stage(SFV_shared_words, stage, logical_dim, ko) + + for nj in T.unroll(warp_N32_tiles): + _sage3_ldmatrix_x4_u8_b_rowmajor_stage_into(Vt_shared, stage, nj, ko, 0, b_regs, 0, lane) + _sage3_ldmatrix_x4_u8_b_rowmajor_stage_into(Vt_shared, stage, nj, ko, 1, b_regs, 4, lane) + for mi in T.unroll(warp_M_tiles): + mma_m16n32k64_blockscale_f32( + a_regs.data, + mi * 8, + b_regs.data, + 0, + pv_acc, + (mi * warp_N32_tiles + nj) * 16, + scale_a, + b_scale_regs[nj], + mi, + 0, + ) + + for mi in T.unroll(warp_M_tiles): + scale0 = row_scale_acc[mi, 0] + scale1 = row_scale_acc[mi, 1] + for nj in T.unroll(warp_N32_tiles): + off = (mi * warp_N32_tiles + nj) * 16 + for p in T.unroll(4): + atom = off + p * 4 + o_acc[atom + 0] = o_acc[atom + 0] * scale0 + pv_acc[atom + 0] + o_acc[atom + 1] = o_acc[atom + 1] * scale0 + pv_acc[atom + 1] + o_acc[atom + 2] = o_acc[atom + 2] * scale1 + pv_acc[atom + 2] + o_acc[atom + 3] = o_acc[atom + 3] * scale1 + pv_acc[atom + 3] + + T.mbarrier_arrive(kv_empty[stage]) + T.mbarrier_wait_parity(o_empty, tile_phase ^ 1) + for mi in T.unroll(warp_M_tiles): + row0 = warp * warp_M + mi * 16 + g + row1 = row0 + 8 + sum0_pair = T.alloc_var( + accum_dtype, + init=row_sum_acc[mi, 0] + T.shfl_xor(row_sum_acc[mi, 0], 1, width=4), + role_scoped=True, + ) + sum1_pair = T.alloc_var( + accum_dtype, + init=row_sum_acc[mi, 1] + T.shfl_xor(row_sum_acc[mi, 1], 1, width=4), + role_scoped=True, + ) + denom0 = T.alloc_var( + accum_dtype, + init=sum0_pair + T.shfl_xor(sum0_pair, 2, width=4), + role_scoped=True, + ) + denom1 = T.alloc_var( + accum_dtype, + init=sum1_pair + T.shfl_xor(sum1_pair, 2, width=4), + role_scoped=True, + ) + inv_sum0 = T.Select( + (denom0 != T.Cast(accum_dtype, 0.0)) and (denom0 == denom0), + T.Cast(accum_dtype, 1.0) / denom0, + T.Cast(accum_dtype, 0.0), + ) + inv_sum1 = T.Select( + (denom1 != T.Cast(accum_dtype, 0.0)) and (denom1 == denom1), + T.Cast(accum_dtype, 1.0) / denom1, + T.Cast(accum_dtype, 0.0), + ) + for nj in T.unroll(warp_N16_tiles): + dim0 = nj * 16 + sublane * 2 + dim1 = dim0 + 1 + dim2 = dim0 + 8 + dim3 = dim2 + 1 + off = (mi * warp_N32_tiles + (nj >> 1)) * 16 + atom0 = off + (nj & 1) * 8 + atom1 = atom0 + 4 + O_shared[row0, dim0] = T.Cast("bfloat16", o_acc[atom0 + 0] * inv_sum0) + O_shared[row0, dim1] = T.Cast("bfloat16", o_acc[atom0 + 1] * inv_sum0) + O_shared[row1, dim0] = T.Cast("bfloat16", o_acc[atom0 + 2] * inv_sum1) + O_shared[row1, dim1] = T.Cast("bfloat16", o_acc[atom0 + 3] * inv_sum1) + O_shared[row0, dim2] = T.Cast("bfloat16", o_acc[atom1 + 0] * inv_sum0) + O_shared[row0, dim3] = T.Cast("bfloat16", o_acc[atom1 + 1] * inv_sum0) + O_shared[row1, dim2] = T.Cast("bfloat16", o_acc[atom1 + 2] * inv_sum1) + O_shared[row1, dim3] = T.Cast("bfloat16", o_acc[atom1 + 3] * inv_sum1) + T.mbarrier_arrive(o_ready) + + return main diff --git a/examples/sage_attention_sm120/tir_helpers.py b/examples/sage_attention_sm120/tir_helpers.py new file mode 100644 index 0000000000..94b6a96a62 --- /dev/null +++ b/examples/sage_attention_sm120/tir_helpers.py @@ -0,0 +1,111 @@ +"""Reusable TileLang TIR helpers for the SM120 SageAttention3 example.""" + +from __future__ import annotations + +import tilelang.language as T + + +@T.macro +def pack_cast2_u32(target_dtype, storage_dtype, v0, v1): + """Pack two scalar values by vector-casting them to a packed dtype.""" + return T.Cast( + T.uint32, + T.reinterpret( + T.Cast(T.dtype(target_dtype).with_lanes(2), T.Shuffle([v0, v1], [0, 1])), + storage_dtype, + ), + ) + + +@T.macro +def mma_m16n32k64_blockscale_f32( + a_regs, + a_offset, + b_regs, + b_offset, + acc, + c_offset, + scale_a, + scale_b, + scale_id_a, + scale_id_b_base, +) -> None: + """Emit four m16n8k64 FP4 MMAs in contiguous n8-atom register order.""" + scale_a_reg = T.alloc_var("uint32", init=scale_a, role_scoped=True) + scale_b_reg = T.alloc_var("uint32", init=scale_b, role_scoped=True) + T.ptx_mma_blockscale( + "float32", + "m16n8k64", + "row", + "col", + "e2m1", + "e2m1", + 4, + a_regs, + a_offset, + b_regs, + b_offset, + acc.data, + c_offset, + scale_a_reg, + scale_b_reg, + scale_id_a, + scale_id_b_base, + ) + T.ptx_mma_blockscale( + "float32", + "m16n8k64", + "row", + "col", + "e2m1", + "e2m1", + 4, + a_regs, + a_offset, + b_regs, + b_offset + 4, + acc.data, + c_offset + 4, + scale_a_reg, + scale_b_reg, + scale_id_a, + scale_id_b_base + 1, + ) + T.ptx_mma_blockscale( + "float32", + "m16n8k64", + "row", + "col", + "e2m1", + "e2m1", + 4, + a_regs, + a_offset, + b_regs, + b_offset + 8, + acc.data, + c_offset + 8, + scale_a_reg, + scale_b_reg, + scale_id_a, + scale_id_b_base + 2, + ) + T.ptx_mma_blockscale( + "float32", + "m16n8k64", + "row", + "col", + "e2m1", + "e2m1", + 4, + a_regs, + a_offset, + b_regs, + b_offset + 12, + acc.data, + c_offset + 12, + scale_a_reg, + scale_b_reg, + scale_id_a, + scale_id_b_base + 3, + ) diff --git a/src/cuda/codegen/codegen_cuda.cc b/src/cuda/codegen/codegen_cuda.cc index 91048c81a4..9c4e1d998a 100644 --- a/src/cuda/codegen/codegen_cuda.cc +++ b/src/cuda/codegen/codegen_cuda.cc @@ -2220,6 +2220,26 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) { this->stream << ss.str(); this->stream << ");\n"; }; + if (op->op.same_as(builtin::address_of())) { + ICHECK_EQ(op->args.size(), 1U); + const BufferLoadNode *load = op->args[0].as(); + if (load && load->dtype.is_float4() && load->dtype.is_scalar()) { + ICHECK_EQ(load->indices.size(), 1U) + << "CodeGenTileLangCUDA only supports flat fp4 address_of."; + enable_fp4_ = true; + const VarNode *buffer_var = load->buffer->data.get(); + std::string vid = GetVarID(buffer_var); + auto packed_it = fp4_packed_buffers_.find(buffer_var); + if (packed_it != fp4_packed_buffers_.end()) { + vid = packed_it->second; + } + PrimExpr packed_offset = + arith::Analyzer().Simplify(truncdiv(load->indices[0], 2)); + os << "(&(((fp4_e2_2_t*)" << vid << ")[" << PrintExpr(packed_offset) + << "]))"; + return; + } + } if (op->op.same_as(tl::max_nan()) || op->op.same_as(tl::min_nan())) { ICHECK_EQ(op->args.size(), 2); const bool is_max = op->op.same_as(tl::max_nan()); @@ -3811,6 +3831,21 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) { this->PrintExpr(op->args[0], os); } os << ")"; + } else if (op->op.same_as(tl::lds32())) { + // Explicit 32-bit shared memory load: load_shared_32(ptr) or + // load_shared_32_conditional(ptr, pred) + ICHECK(!op->args.empty()) + << "T.lds32 expects a shared-memory pointer argument."; + if (op->args.size() > 1) { + os << "tl::load_shared_32_conditional("; + this->PrintExpr(op->args[0], os); + os << ", "; + this->PrintExpr(op->args[1], os); + } else { + os << "tl::load_shared_32("; + this->PrintExpr(op->args[0], os); + } + os << ")"; } else if (op->op.same_as(tl::ldg64())) { // Explicit 64-bit global memory load: load_global_64(ptr) or // load_global_64_conditional(ptr, pred) @@ -4646,6 +4681,10 @@ void CodeGenTileLangCUDA::VisitStmt_(const AllocBufferNode *op) { (alloc_dtype == DataType::Int(4) || alloc_dtype == DataType::UInt(4)) && alloc_dtype.is_scalar() && scope == "local"; if (!is_fp4_scalar_local && !is_int4_scalar_local) { + int bits_per_elem = alloc_dtype.bits() * alloc_dtype.lanes(); + if (scope == "local" && bits_per_elem > 0 && bits_per_elem < 32) { + stream << "alignas(4) "; + } PrintStorageScope(scope, stream); if (is_float4_unpacked_shared) { stream << "uint8_t"; @@ -4689,7 +4728,7 @@ void CodeGenTileLangCUDA::VisitStmt_(const AllocBufferNode *op) { } else { if (alloc_dtype.is_float4_e2m1fn() && alloc_dtype.is_scalar()) { auto vid_packed = vid + "_packed"; - stream << "fp4_e2_2_t " << vid_packed << '[' + stream << "alignas(4) fp4_e2_2_t " << vid_packed << '[' << (constant_size + 1) / 2 << "];\n"; fp4_packed_buffers_[op->buffer->data.get()] = vid_packed; } else { diff --git a/src/cuda/op/copy.cc b/src/cuda/op/copy.cc index 7dbe3cdd67..b73650c59a 100644 --- a/src/cuda/op/copy.cc +++ b/src/cuda/op/copy.cc @@ -5,6 +5,7 @@ #include "op/copy.h" #include "support/check.h" +#include #include #include #include @@ -53,6 +54,148 @@ int TMAPayloadElementBits(DataType dtype) { return dtype.bits(); } +int TMATransactionElementBits(DataType dtype) { + return TMAPayloadElementBits(dtype); +} + +PrimExpr BufferBaseAddressWithElemOffset(const Buffer &buffer) { + PrimExpr base = buffer->data; + if (buffer->elem_offset.defined() && !is_zero(buffer->elem_offset)) { + PrimExpr byte_offset = + buffer->elem_offset * + IntImm(buffer->elem_offset.dtype(), buffer->dtype.bytes()); + base = Call(DataType::Handle(), builtin::handle_add_byte_offset(), + {base, byte_offset}); + } + return base; +} + +PrimExpr RowMajorOffset(const Array &indices, + const Array &shape) { + ICHECK_EQ(indices.size(), shape.size()) + << "indices/shape rank mismatch: " << indices.size() << " vs. " + << shape.size(); + PrimExpr offset = Integer(0); + PrimExpr stride = Integer(1); + for (int i = static_cast(shape.size()) - 1; i >= 0; --i) { + offset += indices[i] * stride; + stride *= shape[i]; + } + return offset; +} + +PrimExpr AddByteOffset(PrimExpr base, PrimExpr byte_offset) { + if (is_zero(byte_offset)) { + return base; + } + return Call(DataType::Handle(), builtin::handle_add_byte_offset(), + {std::move(base), std::move(byte_offset)}); +} + +PrimExpr SharedTmaAccessPtr(const Buffer &buffer, int access_mask, + PrimExpr elem_offset, PrimExpr extent) { + if (buffer->dtype.is_float4_e2m1_unpacked()) { + // float4_e2m1_unpacked is backed by one byte per FP4 element in SMEM. + // Buffer::access_ptr would use the logical FP4 pointer type and halve the + // element offset, which points later TMA stages at the wrong byte address. + PrimExpr byte_offset = cast(DataType::Int(64), elem_offset); + return AddByteOffset(BufferBaseAddressWithElemOffset(buffer), byte_offset); + } + return buffer.access_ptr(access_mask, DataType::Handle(), 1, elem_offset, + extent); +} + +std::vector BufferStridesInOriginalOrder(const Buffer &buffer) { + std::vector strides; + strides.reserve(buffer->shape.size()); + if (!buffer->strides.empty()) { + for (const PrimExpr &stride : buffer->strides) { + strides.push_back(stride); + } + return strides; + } + + PrimExpr stride = Integer(1); + for (int i = static_cast(buffer->shape.size()) - 1; i >= 0; --i) { + strides.insert(strides.begin(), stride); + stride *= buffer->shape[i]; + } + return strides; +} + +struct ProjectedTmaGlobalView { + Array shape; + Array stride; + Array coords; + Array box; + PrimExpr base_elem_offset; + std::vector original_axes; +}; + +ProjectedTmaGlobalView +ProjectGlobalTmaView(const Buffer &buffer, const Array &ranges, + const std::vector &strides, + arith::Analyzer *analyzer) { + ICHECK_EQ(buffer->shape.size(), ranges.size()) + << "ProjectGlobalTmaView: buffer/range rank mismatch for " + << buffer->name; + ICHECK_EQ(buffer->shape.size(), strides.size()) + << "ProjectGlobalTmaView: buffer/stride rank mismatch for " + << buffer->name; + + const int ndim = static_cast(ranges.size()); + ICHECK_GT(ndim, 0); + std::vector keep(ndim, false); + bool has_non_unit_extent = false; + for (int i = 0; i < ndim; ++i) { + if (!analyzer->CanProveEqual(ranges[i]->extent, 1)) { + keep[i] = true; + has_non_unit_extent = true; + } + } + if (!has_non_unit_extent) { + keep[ndim - 1] = true; + } + + // TMA's innermost descriptor dimension has implicit unit stride. If all + // non-unit axes are outside a fixed innermost axis, keep the fixed inner axes + // in the descriptor instead of folding them away. + int fastest_kept = -1; + for (int i = ndim - 1; i >= 0; --i) { + if (keep[i]) { + fastest_kept = i; + break; + } + } + ICHECK_GE(fastest_kept, 0); + if (!analyzer->CanProveEqual(strides[fastest_kept], 1)) { + for (int i = fastest_kept + 1; i < ndim; ++i) { + keep[i] = true; + } + } + + ProjectedTmaGlobalView view; + view.base_elem_offset = Integer(0); + for (int i = 0; i < ndim; ++i) { + // The tensor-map descriptor is created outside the device loop nest. Only + // compile-time fixed coordinates can be folded into its base address; + // dynamic fixed coordinates must stay as descriptor coordinates. + bool can_fold_fixed_axis = ranges[i]->min.as() != nullptr; + if (keep[i] || !can_fold_fixed_axis) { + view.shape.push_back(buffer->shape[i]); + view.stride.push_back(strides[i]); + view.coords.push_back(ranges[i]->min); + view.box.push_back(ranges[i]->extent); + view.original_axes.push_back(i); + } else { + view.base_elem_offset += ranges[i]->min * strides[i]; + } + } + ICHECK_GE(view.shape.size(), 1); + ICHECK_LE(view.shape.size(), 5); + return view; +} + PrimExpr TMABytesFromElements(PrimExpr elements, int bits) { PrimExpr elements_i64 = cast(DataType::Int(64), elements); if (bits % 8 == 0) { @@ -207,6 +350,42 @@ bool GetNoImplicitAsyncCommitWait(const CopyNode &op) { return GetBoolAnnotation(op, attr::kAsyncCopyNoImplicitCommitWait); } +bool HasPreferInstruction(const CopyNode &op, const char *instruction) { + if (auto val = op.annotations.Get("prefer_instruction")) { + auto str = val->as(); + ICHECK(str) + << "T.copy prefer_instruction annotation must be a string, but got " + << val.value().GetTypeKey(); + return str->value == instruction; + } + return false; +} + +DataType ParseTmaElementDType(std::string dtype) { + // Reuse TVM's dtype parser so TMA annotations accept the same spelling as + // Buffer dtypes. Keep the previous non-standard alias for compatibility. + if (dtype == "float8_e5m2fn") { + dtype = "float8_e5m2"; + } + DataType parsed(ffi::StringToDLDataType(dtype)); + ICHECK(parsed.is_scalar()) + << "T.tma_copy tma_element_dtype must be scalar, but got " << parsed; + return parsed; +} + +DataType GetTmaElementDType(const CopyNode &op, DataType default_dtype) { + if (auto val = op.annotations.Get("tma_element_dtype")) { + auto str = val->as(); + ICHECK(str) << "T.tma_copy tma_element_dtype annotation must be a string, " + "but got " + << val.value().GetTypeKey(); + DataType dtype = ParseTmaElementDType(str->value); + to_CUtensorMapDataType(dtype); + return dtype; + } + return default_dtype; +} + PrimExpr GetLeaderScopeThreads(const CopyNode &op, const LowerArgs &T) { if (auto val = op.annotations.Get("leader_scope_threads")) { auto int_val = val->as(); @@ -526,7 +705,9 @@ LayoutMap Copy::InferLayout(const CopyNode &op, const LayoutInferArgs &T, } if (copy_inst == CopyInst::kBulkLoad || copy_inst == CopyInst::kBulkStore || copy_inst == CopyInst::kBulkLoad1D || - copy_inst == CopyInst::kBulkStore1D) { + copy_inst == CopyInst::kBulkStore1D || + copy_inst == CopyInst::kBulkLoadGather4 || + copy_inst == CopyInst::kBulkStoreScatter4) { return InferBulkLayout(op, T, level, copy_inst); } @@ -621,10 +802,14 @@ LayoutMap Copy::InferBulkLayout(const CopyNode &op, const LayoutInferArgs &T, bool is_tma_1d = copy_inst == CopyInst::kBulkLoad1D || copy_inst == CopyInst::kBulkStore1D; - bool is_load = - copy_inst == CopyInst::kBulkLoad || copy_inst == CopyInst::kBulkLoad1D; - bool is_store = - copy_inst == CopyInst::kBulkStore || copy_inst == CopyInst::kBulkStore1D; + bool is_gather_scatter = copy_inst == CopyInst::kBulkLoadGather4 || + copy_inst == CopyInst::kBulkStoreScatter4; + bool is_load = copy_inst == CopyInst::kBulkLoad || + copy_inst == CopyInst::kBulkLoad1D || + copy_inst == CopyInst::kBulkLoadGather4; + bool is_store = copy_inst == CopyInst::kBulkStore || + copy_inst == CopyInst::kBulkStore1D || + copy_inst == CopyInst::kBulkStoreScatter4; Buffer shared_tensor = is_load ? op.dst : op.src; Array shared_range = is_load ? op.dst_range : op.src_range; @@ -649,6 +834,11 @@ LayoutMap Copy::InferBulkLayout(const CopyNode &op, const LayoutInferArgs &T, thread_extent, T.thread_bounds, result_map); } + if (is_gather_scatter && + IsFP4PackedToUnpackedStorageCopy(op.src->dtype, op.dst->dtype)) { + return result_map; + } + if (is_tma_1d) { // 1D TMA requires contiguous shared memory. Do not infer a swizzled shared // layout here, otherwise final instruction selection may fall back to @@ -1030,17 +1220,24 @@ Stmt Copy::LowerLDSM(const CopyNode &op, const LowerArgs &T, ICHECK(copy_inst == CopyInst::kLDSM || copy_inst == CopyInst::kSTSM) << "Invalid copy inst " << static_cast(copy_inst); bool is_ldmatrix = copy_inst == CopyInst::kLDSM; + const char *instruction = is_ldmatrix ? "ldsm" : "stsm"; + bool strict_matrix_copy = HasPreferInstruction(op, instruction); + auto fallback = [&](const std::string &reason) -> Stmt { + ICHECK(!strict_matrix_copy) << "T.copy prefer_instruction=\"" << instruction + << "\" could not be honored: " << reason; + return LowerNormal(op, T, analyzer); + }; Array loop_vars = op.MakeIterVars(); if (loop_vars.size() < 2) { - return LowerNormal(op, T, analyzer); + return fallback("matrix copy lowering requires at least two copy axes."); } for (const auto &iv : loop_vars) analyzer->Bind(iv->var, iv->dom); PrimExpr src_predicate = op.MakePredicate(analyzer, loop_vars, src->shape, 0); PrimExpr dst_predicate = op.MakePredicate(analyzer, loop_vars, dst->shape, 1); if (src_predicate.defined() || dst_predicate.defined()) { - return LowerNormal(op, T, analyzer); + return fallback("matrix copy lowering does not support predicates."); } Buffer shared_tensor = is_ldmatrix ? src : dst; @@ -1055,7 +1252,7 @@ Stmt Copy::LowerLDSM(const CopyNode &op, const LowerArgs &T, } } if (!is_full_range) { - return LowerNormal(op, T, analyzer); + return fallback("local fragment copy region is not a full-range copy."); } Array local_indices = @@ -1065,7 +1262,7 @@ Stmt Copy::LowerLDSM(const CopyNode &op, const LowerArgs &T, local_layout->Forward(local_indices); local_tensor = T.buffer_remap[local_tensor]; if (local_layout->OutputDim() != 1) { - return LowerNormal(op, T, analyzer); + return fallback("local fragment layout is not one-dimensional."); } Array shared_indices = @@ -1092,21 +1289,23 @@ Stmt Copy::LowerLDSM(const CopyNode &op, const LowerArgs &T, row_var->dom->extent, 2, analyzer)) { is_transposed = true; } else { - return LowerNormal(op, T, analyzer); + return fallback( + "local fragment layout is incompatible with 8x8 matrix copy."); } if (shared_tensor->dtype.bytes() != 2) { - return LowerNormal(op, T, analyzer); + return fallback("shared tensor dtype is not 16-bit."); } PrimExpr flattened_indice = shared_tensor.OffsetOf(shared_indices).back(); if (!IndicesCanVectorize(flattened_indice, loop_vars.back()->var, loop_vars.back()->dom->extent, 8, analyzer)) { - return LowerNormal(op, T, analyzer); + return fallback( + "shared tensor address is not vectorizable for matrix copy."); } for (size_t i = 0; i < dst_range.size(); i++) { if (!is_zero(dst_range[i]->min) || !analyzer->CanProveEqual(dst_range[i]->extent, dst->shape[i])) - return LowerNormal(op, T, analyzer); + return fallback("destination range is not a full tensor range."); } PrimExpr extent = local_tensor->shape[0]; @@ -1148,7 +1347,7 @@ Stmt Copy::LowerLDSM(const CopyNode &op, const LowerArgs &T, if (is_ldmatrix) { if (local_tensor->dtype != shared_tensor->dtype) { - return LowerNormal(op, T, analyzer); + return fallback("ldmatrix requires matching shared and fragment dtypes."); } PrimExpr local_addr = Call(DataType::Handle(), tl::access_ptr(), @@ -1412,41 +1611,7 @@ Stmt Copy::LowerBulk(const CopyNode &op, const LowerArgs &T, Array shared_indices; for (auto r : shared_range) shared_indices.push_back(r->min); - std::vector shared_strides; - PrimExpr shared_stride = 1; - for (size_t i = 0; i < shared_tensor->shape.size(); i++) { - auto s = shared_tensor->shape[shared_tensor->shape.size() - i - 1]; - shared_strides.insert(shared_strides.begin(), shared_stride); - shared_stride *= s; - } - - Array global_indices; - for (auto r : global_range) { - global_indices.push_back(r->min); - } - std::vector global_strides; - PrimExpr global_stride = 1; - for (size_t i = 0; i < global_tensor->shape.size(); i++) { - auto s = global_tensor->shape[global_tensor->shape.size() - i - 1]; - global_strides.insert(global_strides.begin(), global_stride); - global_stride *= s; - } - - ICHECK(shared_strides.size() == shared_indices.size()) - << "shared_strides.size() != shared_indices.size()" - << shared_strides.size() << " " << shared_indices.size(); - PrimExpr shared_offset = 0; - for (size_t i = 0; i < shared_indices.size(); i++) { - shared_offset += shared_indices[i] * shared_strides[i]; - } - PrimExpr global_offset = 0; - for (size_t i = 0; i < global_indices.size(); i++) { - global_offset += global_indices[i] * global_strides[i]; - } - - TMADesc desc; - desc.rank = global_tensor->shape.size(); - ICHECK(desc.rank >= 1 && desc.rank <= 5) << desc.rank; + PrimExpr shared_offset = RowMajorOffset(shared_indices, shared_tensor->shape); ICHECK( IsValidTMADtypePair(is_load, global_tensor->dtype, shared_tensor->dtype)) @@ -1454,26 +1619,29 @@ Stmt Copy::LowerBulk(const CopyNode &op, const LowerArgs &T, << shared_tensor->name << " with incompatible data type " << global_tensor->dtype << " and " << shared_tensor->dtype; + DataType tma_element_dtype = GetTmaElementDType(op, global_tensor->dtype); + const bool fp4_packed_to_byte_smem = + is_load && + IsFP4PackedToByteStorageCopy(tma_element_dtype, shared_tensor->dtype); + std::vector global_strides = + BufferStridesInOriginalOrder(global_tensor); + ProjectedTmaGlobalView global_view = ProjectGlobalTmaView( + global_tensor, global_range, global_strides, analyzer); + + TMADesc desc; + desc.rank = global_view.shape.size(); + ICHECK(desc.rank >= 1 && desc.rank <= 5) << desc.rank; desc.data_type = - TensorMapDataTypeForTMA(global_tensor->dtype, shared_tensor->dtype); - desc.global_addr = global_tensor->data; - desc.global_shape = ReverseArray(global_tensor->shape); - Array global_coords = - ReverseArray(global_range.Map([](Range r) { return r->min; })); - if (!global_tensor->strides.empty()) { - desc.global_stride = ReverseArray(global_tensor->strides); - } else { - PrimExpr stride = 1; - desc.global_stride.reserve(desc.rank); - for (size_t i = 0; i < desc.rank; i++) { - desc.global_stride.push_back(stride); - stride *= desc.global_shape[i]; - } - } + TensorMapDataTypeForTMA(tma_element_dtype, shared_tensor->dtype); + desc.global_addr = AddByteOffset( + BufferBaseAddressWithElemOffset(global_tensor), + TMABytesFromElements(global_view.base_elem_offset, tma_element_dtype)); + desc.global_shape = ReverseArray(global_view.shape); + Array global_coords = ReverseArray(global_view.coords); + desc.global_stride = ReverseArray(global_view.stride); ICHECK(is_one(desc.global_stride[0])) << desc.global_stride; - desc.global_stride = desc.global_stride.Map([&](PrimExpr e) { - return TMAGlobalBytesFromElements(e, global_tensor->dtype); - }); + desc.global_stride = desc.global_stride.Map( + [&](PrimExpr e) { return TMABytesFromElements(e, tma_element_dtype); }); for (size_t i{1}; i < desc.global_stride.size(); i++) { auto stride = desc.global_stride[i].as(); if (stride != nullptr) { @@ -1485,31 +1653,41 @@ Stmt Copy::LowerBulk(const CopyNode &op, const LowerArgs &T, } } - auto s_range_idx = 0; - for (size_t i = 0; i < global_range.size(); i++) { - auto g_range = global_range[i]; - if (is_one(g_range->extent)) { + size_t s_range_idx = 0; + for (size_t i = 0; i < global_view.box.size(); ++i) { + PrimExpr g_extent = global_view.box[i]; + if (analyzer->CanProveEqual(g_extent, 1)) { continue; } - while (is_one(shared_range[s_range_idx]->extent) && - s_range_idx < shared_range.size()) { + while (s_range_idx < shared_range.size() && + analyzer->CanProveEqual(shared_range[s_range_idx]->extent, 1)) { s_range_idx++; } if (s_range_idx >= shared_range.size()) { - LOG(FATAL) << "TMA bulk copy cannot support a global range of " - << global_range << ", shared_range " << shared_range; + LOG(FATAL) << "TMA bulk copy cannot support projected global box " + << global_view.box << ", shared_range " << shared_range; } auto s_range = shared_range[s_range_idx]; s_range_idx++; - ICHECK(StructuralEqual()(g_range->extent, s_range->extent)) - << global_tensor->name << "[" << i << "] is illegal, " - << global_tensor->name << "[" << i << "] = " << g_range->extent << ", " - << shared_tensor->name << "[" << s_range_idx - << "] = " << s_range->extent; + bool extent_match = StructuralEqual()(g_extent, s_range->extent); + if (!extent_match && fp4_packed_to_byte_smem && + global_view.original_axes[i] + 1 == + static_cast(global_tensor->shape.size())) { + PrimExpr global_bytes = TMABytesFromElements(g_extent, tma_element_dtype); + PrimExpr shared_bytes = + TMABytesFromElements(s_range->extent, shared_tensor->dtype); + extent_match = analyzer->CanProveEqual(global_bytes, shared_bytes); + } + + ICHECK(extent_match) << "Projected TMA extent from " << global_tensor->name + << "[" << global_view.original_axes[i] + << "] is illegal, " << global_tensor->name << "[" + << global_view.original_axes[i] << "] = " << g_extent + << ", " << shared_tensor->name << "[" + << (s_range_idx - 1) << "] = " << s_range->extent; } - desc.smem_box = - ReverseArray(global_range.Map([](Range r) { return r->extent; })); + desc.smem_box = ReverseArray(global_view.box); desc.smem_stride = Array(desc.rank, PrimExpr(1)); desc.l2_promotion = static_cast(CU_TENSOR_MAP_L2_PROMOTION_L2_128B); @@ -1523,6 +1701,8 @@ Stmt Copy::LowerBulk(const CopyNode &op, const LowerArgs &T, << "shared_tensor: " << shared_tensor->name << " not found in buffer_remap"; shared_tensor = T.buffer_remap.at(shared_tensor); + shared_offset = RowMajorOffset(shared_layout->Forward(shared_indices), + shared_tensor->shape); } if (!shared_layout.defined()) { desc.swizzle = static_cast(CU_TENSOR_MAP_SWIZZLE_NONE); @@ -1578,11 +1758,18 @@ Stmt Copy::LowerBulk(const CopyNode &op, const LowerArgs &T, return fallback_to_normal("non-constant inner box dimension"); } int instruction_dim = *inner_box_dim; - DataType smem_elem_dtype = shared_tensor->dtype; - if (desc.swizzle == static_cast(CU_TENSOR_MAP_SWIZZLE_64B)) { - instruction_dim = TMAElementsForBytes(64, smem_elem_dtype); + if (shared_tensor->dtype.is_float4_e2m1_unpacked()) { + // CUDA encodes unpacked FP4 SMEM TMA as 16U4_ALIGN16B. For this + // descriptor format the inner box dimension is specified in logical FP4 + // elements and must remain 128 even when the shared-memory swizzle span is + // 128B. Deriving the dimension from the swizzle byte span would produce + // 256 elements for CU_TENSOR_MAP_SWIZZLE_128B and incorrectly reject the + // same descriptor shape used by CuTe/CUTLASS. + instruction_dim = 128; + } else if (desc.swizzle == static_cast(CU_TENSOR_MAP_SWIZZLE_64B)) { + instruction_dim = TMAElementsForBytes(64, tma_element_dtype); } else if (desc.swizzle == static_cast(CU_TENSOR_MAP_SWIZZLE_128B)) { - instruction_dim = TMAElementsForBytes(128, smem_elem_dtype); + instruction_dim = TMAElementsForBytes(128, tma_element_dtype); } if (instruction_dim > 256) { ICHECK((*inner_box_dim) % 256 == 0) @@ -1594,8 +1781,7 @@ Stmt Copy::LowerBulk(const CopyNode &op, const LowerArgs &T, << " is not divisible by instruction_dim: " << instruction_dim; desc.smem_box.Set(0, PrimExpr(instruction_dim)); - int inner_box_dim_ = - TMABytesFromElements(instruction_dim, shared_tensor->dtype); + int inner_box_dim_ = TMABytesFromElements(instruction_dim, tma_element_dtype); struct SwizzleCheck { int swizzle; @@ -1654,6 +1840,10 @@ Stmt Copy::LowerBulk(const CopyNode &op, const LowerArgs &T, PrimExpr total_elements = 1; for (auto e : desc.smem_box) total_elements *= e; + PrimExpr shared_access_elements = + fp4_packed_to_byte_smem + ? TMABytesFromElements(total_elements, tma_element_dtype) + : total_elements; auto build_multicast_args = [&](const Array ®ular_args) { Array mc_args; @@ -1672,9 +1862,10 @@ Stmt Copy::LowerBulk(const CopyNode &op, const LowerArgs &T, Var loop_var("i"); int loop_extent = (*inner_box_dim) / instruction_dim; - PrimExpr shared_addr = shared_tensor.access_ptr( - is_load ? 2 : 1, DataType::Handle(), 1, - shared_offset + total_elements * loop_var, total_elements); + PrimExpr shared_addr = + SharedTmaAccessPtr(shared_tensor, is_load ? 2 : 1, + shared_offset + shared_access_elements * loop_var, + shared_access_elements); args.push_back(shared_addr); global_coords.Set(0, global_coords[0] + instruction_dim * loop_var); for (auto coord : global_coords) @@ -1709,8 +1900,8 @@ Stmt Copy::LowerBulk(const CopyNode &op, const LowerArgs &T, multicast_copy, regular_or_noop); } } else { - PrimExpr shared_addr = shared_tensor.access_ptr( - is_load ? 2 : 1, DataType::Handle(), 1, shared_offset, total_elements); + PrimExpr shared_addr = SharedTmaAccessPtr( + shared_tensor, is_load ? 2 : 1, shared_offset, shared_access_elements); args.push_back(shared_addr); for (auto coord : global_coords) args.push_back(coord); @@ -1762,10 +1953,10 @@ Stmt Copy::LowerBulk(const CopyNode &op, const LowerArgs &T, if ((*inner_box_dim) != instruction_dim) { int loop_extent = (*inner_box_dim) / instruction_dim; total_bytes = TMATransactionBytesFromElements( - total_elements * loop_extent, shared_tensor->dtype); + total_elements * loop_extent, tma_element_dtype); } else { total_bytes = - TMATransactionBytesFromElements(total_elements, shared_tensor->dtype); + TMATransactionBytesFromElements(total_elements, tma_element_dtype); } Stmt barrier_before_tma_stmt; @@ -1860,7 +2051,11 @@ Stmt Copy::LowerBulkGather4(const CopyNode &op, const LowerArgs &T, ICHECK(shared_lead != nullptr && *shared_lead == 4) << "tma_gather4/scatter4 shared tile leading dim must be 4, got " << shared_tensor->shape[0]; - ICHECK_EQ(global_tensor->dtype, shared_tensor->dtype); + ICHECK( + IsValidTMADtypePair(is_load, global_tensor->dtype, shared_tensor->dtype)) + << "tma_gather4/scatter4 between buffer " << global_tensor->name + << " and " << shared_tensor->name << " has incompatible data type " + << global_tensor->dtype << " and " << shared_tensor->dtype; Array rows = GetGather4Rows(op); PrimExpr col = GetGather4Col(op); @@ -1869,8 +2064,10 @@ Stmt Copy::LowerBulkGather4(const CopyNode &op, const LowerArgs &T, TMADesc desc; desc.rank = 2; - desc.data_type = to_CUtensorMapDataType(global_tensor->dtype); - desc.global_addr = global_tensor->data; + DataType tma_element_dtype = GetTmaElementDType(op, global_tensor->dtype); + desc.data_type = + TensorMapDataTypeForTMA(tma_element_dtype, shared_tensor->dtype); + desc.global_addr = BufferBaseAddressWithElemOffset(global_tensor); desc.global_shape = ReverseArray(global_tensor->shape); if (!global_tensor->strides.empty()) { @@ -1887,7 +2084,7 @@ Stmt Copy::LowerBulkGather4(const CopyNode &op, const LowerArgs &T, << "tma_gather4/scatter4 requires unit innermost global stride, got " << desc.global_stride; desc.global_stride = desc.global_stride.Map([&](PrimExpr e) { - return TMAGlobalBytesFromElements(e, global_tensor->dtype); + return TMAGlobalBytesFromElements(e, tma_element_dtype); }); for (size_t i = 1; i < desc.global_stride.size(); ++i) { if (auto stride = desc.global_stride[i].as()) { @@ -1956,8 +2153,8 @@ Stmt Copy::LowerBulkGather4(const CopyNode &op, const LowerArgs &T, PrimExpr total_elements = 4 * K_box; PrimExpr smem_addr = - shared_tensor.access_ptr(is_load ? 2 : 1, DataType::Handle(), 1, - IntImm(DataType::Int(32), 0), total_elements); + SharedTmaAccessPtr(shared_tensor, is_load ? 2 : 1, + IntImm(DataType::Int(32), 0), total_elements); Array args; args.push_back(create_descriptor); @@ -2044,10 +2241,18 @@ Stmt Copy::LowerBulk1D(const CopyNode &op, const LowerArgs &T, } PrimExpr elements = analyzer->Simplify(shared_elements); - PrimExpr shared_addr = shared_tensor.access_ptr( - is_load ? 2 : 1, DataType::Handle(), 1, shared_offset, elements); + PrimExpr shared_addr = SharedTmaAccessPtr(shared_tensor, is_load ? 2 : 1, + shared_offset, elements); PrimExpr global_addr = global_tensor.access_ptr( is_load ? 1 : 2, DataType::Handle(), 1, global_offset, elements); + DataType tma_element_dtype = GetTmaElementDType(op, global_tensor->dtype); + if (op.annotations.Get("tma_element_dtype")) { + ICHECK_EQ(TMATransactionElementBits(tma_element_dtype), + TMATransactionElementBits(global_tensor->dtype)) + << "T.tma_copy tma_element_dtype changes the transaction width in " + "the 1D TMA path, which has no tensor-map descriptor. Use a buffer " + "with matching logical dtype or a multi-dimensional TMA copy."; + } int barrier_base_id = -1; PrimExpr mbar_handle; @@ -2068,7 +2273,7 @@ Stmt Copy::LowerBulk1D(const CopyNode &op, const LowerArgs &T, Stmt tma_copy; PrimExpr total_bytes = - TMATransactionBytesFromElements(elements, shared_tensor->dtype); + TMATransactionBytesFromElements(elements, tma_element_dtype); if (is_load) { PrimExpr mbar_arg = barrier_base_id >= 0 ? mbar_handle : PrimExpr(0); tma_copy = Evaluate(Call(DataType::Handle(), tma_load(), @@ -2154,7 +2359,7 @@ Stmt Im2Col::Lower(const Im2ColOpNode &op, const LowerArgs &T, TMAIm2ColDesc desc; desc.rank = src->shape.size(); desc.data_type = to_CUtensorMapDataType(src->dtype); - desc.global_addr = src->data; + desc.global_addr = BufferBaseAddressWithElemOffset(src); desc.global_shape = ReverseArray(src->shape); if (!src->strides.empty()) { diff --git a/src/cuda/op/copy_analysis.cc b/src/cuda/op/copy_analysis.cc index 7a6c292263..d3ee4812c1 100644 --- a/src/cuda/op/copy_analysis.cc +++ b/src/cuda/op/copy_analysis.cc @@ -94,6 +94,8 @@ enum class PreferredCopyInstruction { kAuto, kTMA, kCPAsync, + kLDSM, + kSTSM, kSync, }; @@ -103,6 +105,8 @@ constexpr std::pair kPreferredCopyInstructions[] = { {"tma", PreferredCopyInstruction::kTMA}, {"cp_async", PreferredCopyInstruction::kCPAsync}, + {"ldsm", PreferredCopyInstruction::kLDSM}, + {"stsm", PreferredCopyInstruction::kSTSM}, {"sync", PreferredCopyInstruction::kSync}, }; @@ -126,7 +130,8 @@ ParsePreferredCopyInstruction(const std::string &prefer) { } } LOG(FATAL) << "Unsupported T.copy prefer_instruction=\"" << prefer - << "\". Expected one of: \"tma\", \"cp_async\", \"sync\"."; + << "\". Expected one of: \"tma\", \"cp_async\", \"ldsm\", " + "\"stsm\", \"sync\"."; return PreferredCopyInstruction::kAuto; } @@ -462,6 +467,15 @@ std::string MakeTmaUnavailableReason(const CopyNode &op) { << " (scope=" << op.dst.scope() << ", dtype=" << op.dst->dtype << ")."; return oss.str(); } + if (op.src->dtype.is_float4_e2m1fn() && op.dst->dtype == DataType::UInt(8)) { + oss << "T.tma_copy() supports packed FP4 global -> uint8 shared only as a " + "descriptor-based TMA load where the shared uint8 buffer is " + "physical byte storage. Got src=" + << op.src->name << " (scope=" << op.src.scope() + << ", dtype=" << op.src->dtype << "), dst=" << op.dst->name + << " (scope=" << op.dst.scope() << ", dtype=" << op.dst->dtype << ")."; + return oss.str(); + } oss << "T.tma_copy() requires TMA-capable target and global<->shared copy " "pattern, but TMA is not available for src=" << op.src->name << ", dst=" << op.dst->name; @@ -687,6 +701,28 @@ CopyInstSelection SelectCopyInstForLowering(const CopyNode &op, facts.async_unavailable_reason); } + if (facts.prefer_instruction == PreferredCopyInstruction::kLDSM) { + if (!facts.can_ldsm) { + return Unsupported( + "T.copy prefer_instruction=\"ldsm\" requires a shared->fragment copy " + "on a target with ldmatrix support. Got src=" + + std::string(op.src->name) + " (scope=" + op.src.scope() + "), dst=" + + std::string(op.dst->name) + " (scope=" + op.dst.scope() + ")."); + } + return Supported(CopyInst::kLDSM); + } + + if (facts.prefer_instruction == PreferredCopyInstruction::kSTSM) { + if (!facts.can_stsm) { + return Unsupported( + "T.copy prefer_instruction=\"stsm\" requires a fragment->shared copy " + "on a target with stmatrix support. Got src=" + + std::string(op.src->name) + " (scope=" + op.src.scope() + "), dst=" + + std::string(op.dst->name) + " (scope=" + op.dst.scope() + ")."); + } + return Supported(CopyInst::kSTSM); + } + if (facts.prefer_instruction == PreferredCopyInstruction::kSync) { return Supported(SelectSyncLikeInst(facts)); } @@ -731,6 +767,19 @@ CopyInstSelection ClassifyWarpSpecializedProducerCopy(const CopyNode &op, : Unsupported(facts.async_unavailable_reason); } + if (facts.prefer_instruction == PreferredCopyInstruction::kLDSM) { + return facts.can_ldsm + ? Supported(CopyInst::kLDSM) + : Unsupported("T.copy prefer_instruction=\"ldsm\" could not be " + "honored for the producer copy."); + } + if (facts.prefer_instruction == PreferredCopyInstruction::kSTSM) { + return facts.can_stsm + ? Supported(CopyInst::kSTSM) + : Unsupported("T.copy prefer_instruction=\"stsm\" could not be " + "honored for the producer copy."); + } + if (!facts.disable_tma) { CopyInst inst = SelectTmaInst(facts, /*allow_load=*/true, /*allow_store=*/false, diff --git a/src/cuda/transform/annotate_warp_group_reg_alloc.cc b/src/cuda/transform/annotate_warp_group_reg_alloc.cc index 1bd73e6f57..7369ea65b6 100644 --- a/src/cuda/transform/annotate_warp_group_reg_alloc.cc +++ b/src/cuda/transform/annotate_warp_group_reg_alloc.cc @@ -26,6 +26,16 @@ using namespace ffi; namespace { +bool IsManualWarpSpecializationScope(const AttrStmtNode *op) { + return op->attr_key == "warp_specialize"; +} + +Evaluate MakeSetMaxNRegCall(int reg_count, int is_inc) { + return Evaluate(Call(DataType::Handle(), set_max_nreg(), + {IntImm(DataType::Int(32), reg_count), + IntImm(DataType::Int(32), is_inc)})); +} + template Stmt RewriteWarpSpecializationBody(const Stmt &stmt, F &&rewrite_if, bool *rewrote) { @@ -203,6 +213,20 @@ class SetMaxNRegInjector : public StmtExprMutator { call->op.same_as(set_max_nreg())) { return StmtExprMutator::VisitStmt_(op); } + if (in_manual_warp_specialization_ && + call->op.same_as(annotate_producer_reg_dealloc())) { + auto reg_hint = call->args[0].as()->value; + ICHECK(reg_hint <= 240 && reg_hint >= 24) + << "Invalid reg hint: " << reg_hint; + return MakeSetMaxNRegCall(reg_hint, 0); + } + if (in_manual_warp_specialization_ && + call->op.same_as(annotate_consumer_reg_alloc())) { + auto reg_hint = call->args[0].as()->value; + ICHECK(reg_hint <= 240 && reg_hint >= 24) + << "Invalid reg hint: " << reg_hint; + return MakeSetMaxNRegCall(reg_hint, 1); + } if (call->op.same_as(annotate_producer_reg_dealloc()) || call->op.same_as(annotate_consumer_reg_alloc()) || call->op.same_as(no_set_max_nreg())) { @@ -226,6 +250,12 @@ class SetMaxNRegInjector : public StmtExprMutator { } thread_iv_ = {}; return attr_stmt; + } else if (IsManualWarpSpecializationScope(op)) { + bool old_in_manual_ws = in_manual_warp_specialization_; + in_manual_warp_specialization_ = true; + AttrStmt attr_stmt = Downcast(StmtExprMutator::VisitStmt_(op)); + in_manual_warp_specialization_ = old_in_manual_ws; + return attr_stmt; } else if (op->attr_key == attr::kWarpSpecializationScope) { bool rewrote_ws_body = false; auto rewrite_if = [&](const IfThenElse &if_then_else) -> Stmt { @@ -362,6 +392,7 @@ class SetMaxNRegInjector : public StmtExprMutator { Array nreg_; bool preserve_explicit_set_max_nreg_{false}; + bool in_manual_warp_specialization_{false}; IterVar thread_iv_; Optional updated_thread_extent_; bool need_update_thread_extent_ = false; diff --git a/src/cuda/transform/lower_hopper_intrin.cc b/src/cuda/transform/lower_hopper_intrin.cc index 7ff134aca7..1c17ab812e 100644 --- a/src/cuda/transform/lower_hopper_intrin.cc +++ b/src/cuda/transform/lower_hopper_intrin.cc @@ -13,6 +13,7 @@ #include #include +#include #include #include @@ -141,8 +142,8 @@ class LowerHopperIntrin : public StmtExprMutator { return; } for (auto &desc_init : desc_inits_) { - if (!desc_init.emitted && - desc_init.base_var == alloc->buffer->data.get()) { + if (!desc_init.emitted && desc_init.base_var.defined() && + desc_init.base_var.value().same_as(alloc->buffer->data)) { result->push_back(desc_init.stmt); desc_init.emitted = true; } @@ -186,15 +187,22 @@ class LowerHopperIntrin : public StmtExprMutator { if (iter != desc_map_.end()) { var = iter->second; } else { - String name = call->args[2].as().value()->name_hint; + Optional base_var; + std::string name; + if (const auto *base_var_node = call->args[2].as()) { + base_var = GetRef(base_var_node); + name = base_var.value()->name_hint; + } else { + name = "tma_offset_" + std::to_string(desc_counter_++); + } var = Var(name + "_desc", PointerType(PrimType(cuTensorMapType()), "grid_constant")); Call call_ref = GetRef(call); desc_map_[call_ref] = var; Array init_desc_args = MakeInitDescArgs(call_ref, var); init_desc_arg_map_.Set(var, init_desc_args); - desc_inits_.push_back({call->args[2].as().value().get(), - MakeInitDescStmt(var, init_desc_args), false}); + desc_inits_.push_back( + {base_var, MakeInitDescStmt(var, init_desc_args), false}); prefetch_calls_.push_back( Evaluate(Call(DataType::Handle(), builtin::call_extern(), {StringImm("tl::prefetch_tma_descriptor"), var}))); @@ -231,7 +239,7 @@ class LowerHopperIntrin : public StmtExprMutator { private: struct DescInit { - const VarNode *base_var; + Optional base_var; Stmt stmt; bool emitted; }; @@ -265,6 +273,7 @@ class LowerHopperIntrin : public StmtExprMutator { std::unordered_map desc_map_; std::vector desc_inits_; Map> init_desc_arg_map_; + int desc_counter_{0}; LowerHopperIntrin(bool disable_shuffle_elect) : disable_shuffle_elect_(disable_shuffle_elect) {} bool disable_shuffle_elect_; diff --git a/src/layout/gemm_layouts.cc b/src/layout/gemm_layouts.cc index 5ff74cf749..ba8d2ba379 100644 --- a/src/layout/gemm_layouts.cc +++ b/src/layout/gemm_layouts.cc @@ -912,6 +912,36 @@ Layout makeGemmABLayoutSm100(int mat_stride, int mat_continuous, int continuity, TILELANG_COMPILER_UNREACHABLE(); // to prevent compiler warning } +static Layout makeSm120RrSmemSelectorLayout(int mat_stride, int major_size, + int element_size) { + ICHECK(element_size <= 8) + << "SM120 RR shared-memory selector supports <=8-bit elements, got " + << element_size; + int vector_size = 128 / element_size; + + // Mirror CUTLASS detail::sm120_rr_smem_selector: + // Layout_K_SW128_Atom -> Layout_K_SW64_Atom -> + // Layout_K_SW32_Atom -> Layout_K_INTER_Atom. + if (mat_stride % 8 == 0) { + if (major_size % (vector_size * 8) == 0) + return MakeFullBankSwizzleLayout2D(mat_stride, major_size, element_size); + if (major_size % (vector_size * 4) == 0) + return MakeHalfBankSwizzleLayout2D(mat_stride, major_size, element_size); + if (major_size % (vector_size * 2) == 0) + return MakeQuarterBankSwizzleLayout2D(mat_stride, major_size, + element_size); + } + + if (major_size % vector_size == 0) + return makeLinearLayout( + Array{Integer(mat_stride), Integer(major_size)}); + + ICHECK(0) << "Unsupported SM120 RR shared-memory layout with stride=" + << mat_stride << ", major_size=" << major_size + << ", element_size=" << element_size; + TILELANG_COMPILER_UNREACHABLE(); +} + Layout makeGemmABLayoutCDNA(int stride, int continuous, int element_size, int kPack) { return makeMatrixCoreSwizzleLayout(stride, continuous, element_size, kPack); @@ -962,6 +992,14 @@ Layout makeTcgen05mmaSwizzledLayout(const Buffer &buffer, int continuity, return ExpandLayout2D(base, buffer); } +Layout makeSm120Fp4SmemLayout(const Buffer &buffer) { + auto info = GetSwizzleShapeInfoChecked(buffer); + auto base = makeSm120RrSmemSelectorLayout(static_cast(info.stride), + static_cast(info.continuous), + info.element_size); + return ExpandLayout2D(base, buffer); +} + SwizzleMode DetectSwizzleMode(const Layout &layout, const Buffer &buffer) { SwizzleShapeInfo info; if (!TryGetSwizzleShapeInfo(buffer, &info)) { diff --git a/src/layout/layout.cc b/src/layout/layout.cc index c59c3abae4..266d23edfb 100644 --- a/src/layout/layout.cc +++ b/src/layout/layout.cc @@ -109,6 +109,13 @@ PrimExpr RestoreInputPlaceholders(const PrimExpr &expr, return restored; } +bool NeedsExplicitPackLane(int64_t elem_bits) { + // Byte and sub-byte typed aliases need an explicit lane when they are packed + // inside a wider storage element. Without it, e.g. uint16 -> uint8 maps both + // byte lanes back to the same physical layout coordinate. + return elem_bits <= 8; +} + Layout TryPackedSubtypeReshape(const LayoutNode *layout_node, const Array &shape, arith::Analyzer *analyzer, @@ -125,12 +132,12 @@ Layout TryPackedSubtypeReshape(const LayoutNode *layout_node, const Array &input_shape = layout_node->InputShape(); - // Narrower target element, e.g. uint8 -> fp4. + // Narrower target element, e.g. uint16 -> uint8 or uint8 -> fp4. // One old logical element now contains `pack_factor` new logical elements. // The generic flat-index reshape would lose this packed-storage structure, so // we materialize it as an extra trailing output dimension ("pack lane"). if (*old_elem_bits > *new_elem_bits && *old_elem_bits % *new_elem_bits == 0 && - *new_elem_bits < 8) { + NeedsExplicitPackLane(*new_elem_bits)) { int64_t pack_factor = *old_elem_bits / *new_elem_bits; Array new_vars = CreateReshapeVars(shape, analyzer); PrimExpr flat_index = ComputeFlatIndex(shape, new_vars); @@ -147,12 +154,12 @@ Layout TryPackedSubtypeReshape(const LayoutNode *layout_node, return Layout(shape, new_forward_index); } - // Wider target element, e.g. fp4 -> uint8. + // Wider target element, e.g. uint8 -> uint16 or fp4 -> uint8. // This is only valid if the current layout already exposes the packed // sub-elements as its last output dimension. We collapse that trailing pack // lane back into the logical element index of the wider dtype. if (*old_elem_bits < *new_elem_bits && *new_elem_bits % *old_elem_bits == 0 && - *old_elem_bits < 8) { + NeedsExplicitPackLane(*old_elem_bits)) { int64_t pack_factor = *new_elem_bits / *old_elem_bits; Array output_shape = layout_node->OutputShape(); if (output_shape.empty() || @@ -198,11 +205,11 @@ Fragment TryPackedSubtypeReshape(const FragmentNode *fragment_node, const Array &input_shape = fragment_node->InputShape(); - // Same idea as Layout::Reshape above: preserve packed sub-byte storage by - // making the pack lane explicit in the fragment mapping instead of silently - // flattening it away. + // Same idea as Layout::Reshape above: preserve packed byte/sub-byte storage + // by making the pack lane explicit in the fragment mapping instead of + // silently flattening it away. if (*old_elem_bits > *new_elem_bits && *old_elem_bits % *new_elem_bits == 0 && - *new_elem_bits < 8) { + NeedsExplicitPackLane(*new_elem_bits)) { int64_t pack_factor = *old_elem_bits / *new_elem_bits; Array new_vars = CreateReshapeVars(shape, analyzer); PrimExpr flat_index = ComputeFlatIndex(shape, new_vars); @@ -231,7 +238,7 @@ Fragment TryPackedSubtypeReshape(const FragmentNode *fragment_node, } if (*old_elem_bits < *new_elem_bits && *new_elem_bits % *old_elem_bits == 0 && - *old_elem_bits < 8) { + NeedsExplicitPackLane(*old_elem_bits)) { int64_t pack_factor = *new_elem_bits / *old_elem_bits; Array output_shape = fragment_node->OutputShape(); if (output_shape.empty() || @@ -1147,6 +1154,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { [](const Buffer &buffer, int continuity, bool k_inner) { return makeTcgen05mmaSwizzledLayout(buffer, continuity, k_inner); }) + .def("tl.make_sm120_fp4_smem_layout", + [](const Buffer &buffer) { return makeSm120Fp4SmemLayout(buffer); }) .def("tl.make_full_bank_swizzled_layout", [](const Buffer &buffer) { return makeFullBankSwizzleLayout(buffer); diff --git a/src/layout/layout.h b/src/layout/layout.h index 6147afbb82..a3d65e233a 100644 --- a/src/layout/layout.h +++ b/src/layout/layout.h @@ -283,6 +283,7 @@ Layout makeWgmmaSwizzledLayout(const Buffer &buffer, int continuity = -1, bool k_inner = true); Layout makeTcgen05mmaSwizzledLayout(const Buffer &buffer, int continuity = -1, bool k_inner = true); +Layout makeSm120Fp4SmemLayout(const Buffer &buffer); Layout makeFullBankSwizzleLayout(const Buffer &buffer); Layout makeHalfBankSwizzleLayout(const Buffer &buffer); Layout makeQuarterBankSwizzleLayout(const Buffer &buffer); diff --git a/src/op/builtin.cc b/src/op/builtin.cc index cb9c68b834..63c2923beb 100644 --- a/src/op/builtin.cc +++ b/src/op/builtin.cc @@ -729,6 +729,11 @@ TIR_DEFINE_TL_BUILTIN(__ffs).set_num_inputs(1).set_attr( TIR_DEFINE_TL_BUILTIN(ldg32).set_num_inputs(-1).set_attr( "TCallEffectKind", Integer(CallEffectKind::kPure)); +// lds32(address, predicate(optional)) -> 32-bit value +// Shared memory load with 32-bit vector width. +TIR_DEFINE_TL_BUILTIN(lds32).set_num_inputs(-1).set_attr( + "TCallEffectKind", Integer(CallEffectKind::kPure)); + // ldg64(address, predicate(optional)) -> 64-bit value // Global memory load with 64-bit vector width TIR_DEFINE_TL_BUILTIN(ldg64).set_num_inputs(-1).set_attr( diff --git a/src/op/builtin.h b/src/op/builtin.h index 38d8fda461..4cdba72e25 100644 --- a/src/op/builtin.h +++ b/src/op/builtin.h @@ -48,6 +48,11 @@ static constexpr const char *kAsyncCopyNoImplicitCommitWait = static constexpr const char *kPipelineMbarPhaseExpr = "tl.pipeline_mbar_phase_expr"; static constexpr const char *kLocalVarInit = "tl.local_var_init"; +// Allocation annotation requesting StorageRewrite to keep the allocation at +// the lexical statement where it was emitted. This is intended for manual +// warp-specialized code where branch-local register arrays must not be hoisted +// ahead of setmaxnreg or out of a producer/consumer role branch. +static constexpr const char *kRoleScopedAlloc = "tl.role_scoped_alloc"; // A PrimFunc-level attribute carrying a list of handle Vars // that must NOT be marked with the restrict qualifier in codegen. // Type: ffi::Array @@ -1275,6 +1280,19 @@ TVM_DLL const Op &__ffs(); */ TVM_DLL const Op &ldg32(); +/*! + * \brief tilelang intrinsic for shared memory load with 32-bit vector width. + * + * This op loads 32 bits (4 bytes) from shared memory using an explicit + * PTX ld.shared.u32 instruction. The source pointer must be 4-byte aligned. + * An optional predicate skips the load and returns zero when false. + * + * Usage from TVMScript: + * y = T.lds32(smem_u8[i]) + * y = T.lds32(smem_u8[i], pred=i < N) + */ +TVM_DLL const Op &lds32(); + /*! * \brief tilelang intrinsic for global memory load with 64-bit vector width. * diff --git a/src/op/utils.h b/src/op/utils.h index e80a8fc2c6..a14713d0ca 100644 --- a/src/op/utils.h +++ b/src/op/utils.h @@ -121,8 +121,19 @@ inline bool IsFP4PackedToUnpackedStorageCopy(DataType global_dtype, shared_dtype.is_float4_e2m1_unpacked(); } +// True when global packed FP4 is copied into byte-addressable packed FP4 SMEM. +// The TensorMap descriptor still uses the packed FP4 element type; the uint8 +// shared buffer is only the physical byte storage viewed later by LDSM helpers. +inline bool IsFP4PackedToByteStorageCopy(DataType global_dtype, + DataType shared_dtype) { + return global_dtype.is_float4_e2m1fn() && shared_dtype == DataType::UInt(8); +} + inline bool IsValidTMALoadDtypePair(DataType global_dtype, DataType shared_dtype) { + if (IsFP4PackedToByteStorageCopy(global_dtype, shared_dtype)) { + return true; + } if (global_dtype.is_float4_e2m1_unpacked() || shared_dtype.is_float4_e2m1_unpacked()) { return IsFP4PackedToUnpackedStorageCopy(global_dtype, shared_dtype); diff --git a/src/tl_templates/cuda/copy.h b/src/tl_templates/cuda/copy.h index 3a98286520..c592c6d0a5 100644 --- a/src/tl_templates/cuda/copy.h +++ b/src/tl_templates/cuda/copy.h @@ -146,6 +146,34 @@ template struct global_load { } }; +// Shared memory load intrinsics with explicit vector widths +// Following the same wrapper/specialization shape as global_load. + +// Primary template declaration +template struct shared_load; + +// lds32: Load 32 bits (4 bytes) from shared memory +template struct shared_load { + TL_DEVICE shared_load(AccessType &D, void const *ptr) { + unsigned &data = reinterpret_cast(D); + uint32_t smem_int_ptr = smem_ptr_to_uint(ptr); + asm volatile("ld.shared.u32 %0, [%1];\n" : "=r"(data) : "r"(smem_int_ptr)); + } + + TL_DEVICE shared_load(AccessType &D, void const *ptr, bool pred_guard) { + unsigned &data = reinterpret_cast(D); + uint32_t smem_int_ptr = smem_ptr_to_uint(ptr); + asm volatile("{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %2, 0;\n" + " mov.b32 %0, %3;\n" + " @p ld.shared.u32 %0, [%1];\n" + "}\n" + : "=r"(data) + : "r"(smem_int_ptr), "r"((int)pred_guard), "r"(data)); + } +}; + // Convenience wrapper functions for direct use // load_global_32: Load 32 bits, return uint32_t TL_DEVICE uint32_t load_global_32(const void *ptr) { @@ -154,6 +182,13 @@ TL_DEVICE uint32_t load_global_32(const void *ptr) { return ret; } +// load_shared_32: Load one aligned 32-bit word from shared memory. +TL_DEVICE uint32_t load_shared_32(const void *ptr) { + uint32_t ret{}; + shared_load(ret, ptr); + return ret; +} + // load_global_64: Load 64 bits, return uint64_t TL_DEVICE uint2 load_global_64(const void *ptr) { uint2 ret{}; @@ -175,6 +210,12 @@ TL_DEVICE uint32_t load_global_32_conditional(const void *ptr, bool pred) { return ret; } +TL_DEVICE uint32_t load_shared_32_conditional(const void *ptr, bool pred) { + uint32_t ret{}; + shared_load(ret, ptr, pred); + return ret; +} + TL_DEVICE uint2 load_global_64_conditional(const void *ptr, bool pred) { uint2 ret{}; global_load(ret, ptr, pred); diff --git a/src/transform/layout_inference.cc b/src/transform/layout_inference.cc index 36a8a6c836..0334429fe7 100644 --- a/src/transform/layout_inference.cc +++ b/src/transform/layout_inference.cc @@ -797,6 +797,31 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { } } } + if (op->annotations.count("layout_view_map")) { + auto maybe_map = op->annotations.Get("layout_view_map") + ->as>(); + ICHECK(maybe_map.has_value()) + << "layout_view_map must be a Map"; + for (const auto &[key, val] : maybe_map.value()) { + auto maybe_buffer = key.as(); + ICHECK(maybe_buffer.has_value()) + << "layout_view_map key must be a Buffer"; + auto maybe_layout = val.as(); + ICHECK(maybe_layout.has_value()) + << "layout_view_map value must be a Layout"; + Buffer buffer = maybe_buffer.value(); + Layout layout = maybe_layout.value(); + if (ShapesEqual(layout->InputShape(), buffer->shape, &analyzer_)) { + annotated_layout_map_.Set(buffer, layout); + } else { + auto reshaped_layout = + layout->Reshape(buffer->shape, &analyzer_, + Integer(GetElementStorageBits(buffer->dtype)), + Integer(GetElementStorageBits(buffer->dtype))); + annotated_layout_map_.Set(buffer, reshaped_layout); + } + } + } } void VisitStmt_(const AttrStmtNode *op) final { diff --git a/src/transform/lower_tile_op.cc b/src/transform/lower_tile_op.cc index edbfadab62..d861e5aa2a 100644 --- a/src/transform/lower_tile_op.cc +++ b/src/transform/lower_tile_op.cc @@ -746,6 +746,39 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { return result; } + Call RewriteAccessPtrArg(Call call, int arg_index) { + PrimExpr access_ptr = call->args[arg_index]; + Call access_ptr_call = Downcast(access_ptr); + + if (access_ptr_call->op.same_as(builtin::address_of())) { + Optional resolved = ResolveBufferLoad(access_ptr_call->args[0]); + ICHECK(resolved.defined()) + << "Invalid address_of argument for permuted layout: " + << access_ptr_call->args[0]; + PrimExpr load_expr = resolved.value(); + if (!load_expr.same_as(access_ptr_call->args[0])) { + auto call_node = call.CopyOnWrite(); + call_node->args.Set(arg_index, + Call(access_ptr_call->dtype, access_ptr_call->op, + {load_expr}, access_ptr_call->annotations, + access_ptr_call->span)); + access_ptr_call = Downcast(call->args[arg_index]); + access_ptr = call->args[arg_index]; + } + } else if (!access_ptr_call->op.same_as(builtin::tvm_access_ptr()) && + !access_ptr_call->op.same_as(tl::access_ptr())) { + LOG(FATAL) << "Invalid access ptr for permuted layout: " << access_ptr; + } + + auto new_access_ptr = + HandleAccessPtrAndOffset(access_ptr, std::nullopt, call->dtype); + if (new_access_ptr.rewritten) { + auto new_call = call.CopyOnWrite(); + new_call->args.Set(arg_index, new_access_ptr.expr); + } + return call; + } + Optional ResolveBufferLoad(const PrimExpr &expr) const { if (expr->IsInstance()) { return expr; diff --git a/src/transform/merge_shared_memory_allocations.cc b/src/transform/merge_shared_memory_allocations.cc index d2cae3b49e..58fe690182 100644 --- a/src/transform/merge_shared_memory_allocations.cc +++ b/src/transform/merge_shared_memory_allocations.cc @@ -72,6 +72,68 @@ static bool IsStaticSharedMemory(Var buffer_var) { storage_scope.tag.empty(); } +static int64_t StorageBitsPerElement(DataType dtype) { + int64_t bits = static_cast(dtype.bits()) * dtype.lanes(); + ICHECK_GT(bits, 0) << "Invalid shared-memory dtype: " << dtype; + return bits; +} + +static DataType StorageIndexDType(DataType dtype) { + return (dtype.is_int() || dtype.is_uint()) ? dtype : DataType::Int(32); +} + +static PrimExpr StorageElementsForShape(const Array &shape) { + DataType size_dtype = + shape.empty() ? DataType::Int(32) : StorageIndexDType(shape[0].dtype()); + PrimExpr elements = make_const(size_dtype, 1); + for (const PrimExpr &extent : shape) { + PrimExpr e = extent; + if (e.dtype() != size_dtype) { + e = cast(size_dtype, e); + } + elements = elements * e; + } + return elements; +} + +static PrimExpr StorageBytesForElements(PrimExpr elements, DataType dtype) { + int64_t bits_per_element = StorageBitsPerElement(dtype); + DataType size_dtype = StorageIndexDType(elements.dtype()); + if (elements.dtype() != size_dtype) { + elements = cast(size_dtype, elements); + } + PrimExpr bits = elements * make_const(size_dtype, bits_per_element); + if (bits_per_element % 8 == 0) { + return indexdiv(bits, make_const(size_dtype, 8)); + } + return indexdiv(bits + make_const(size_dtype, 7), make_const(size_dtype, 8)); +} + +static PrimExpr StorageBytesForBuffer(const Buffer &buffer) { + return StorageBytesForElements(StorageElementsForShape(buffer->shape), + buffer->dtype); +} + +static int64_t StorageBytesForConstantElements(int64_t elements, + DataType dtype) { + ICHECK_GE(elements, 0); + int64_t bits = elements * StorageBitsPerElement(dtype); + return (bits + 7) / 8; +} + +static PrimExpr ByteOffsetToElementOffset(PrimExpr byte_offset, + DataType dtype) { + int64_t bits_per_element = StorageBitsPerElement(dtype); + if (bits_per_element % 8 == 0) { + return indexdiv(byte_offset, bits_per_element / 8); + } + ICHECK_EQ(8 % bits_per_element, 0) + << "Cannot represent byte offset as an element offset for dtype " + << dtype; + DataType offset_dtype = byte_offset.dtype(); + return byte_offset * make_const(offset_dtype, 8 / bits_per_element); +} + /*! * \brief collect the mapping from the buffer var to its allocate */ @@ -571,25 +633,7 @@ class SharedMemoryRewriter : public StmtExprMutator { for (const VarNode *var : sorted_vars) { const AllocBufferNode *alloc = shmem_allocs_.at(var); - int64_t bytes_per_elem = static_cast( - alloc->buffer->dtype.bytes() * alloc->buffer->dtype.lanes()); - - DataType size_dtype = DataType::Int(32); - if (!alloc->buffer->shape.empty()) { - size_dtype = alloc->buffer->shape[0].dtype(); - } - if (!size_dtype.is_int() && !size_dtype.is_uint()) { - size_dtype = DataType::Int(32); - } - - PrimExpr size_expr = make_const(size_dtype, bytes_per_elem); - for (const PrimExpr &extent : alloc->buffer->shape) { - PrimExpr e = extent; - if (e.dtype() != size_dtype) { - e = cast(size_dtype, e); - } - size_expr = size_expr * e; - } + PrimExpr size_expr = StorageBytesForBuffer(alloc->buffer); int alignment = align_bytes_; auto align_it = shmem_alignment_map_.find(var); @@ -628,9 +672,7 @@ class SharedMemoryRewriter : public StmtExprMutator { auto alloc_it = shmem_allocs_.find(buffer_var_node); if (alloc_it != shmem_allocs_.end()) { const AllocBufferNode *alloc = alloc_it->second; - PrimExpr buffer_size_bytes = alloc->buffer->shape[0] * - alloc->buffer->dtype.bytes() * - alloc->buffer->dtype.lanes(); + PrimExpr buffer_size_bytes = StorageBytesForBuffer(alloc->buffer); LOG(DEBUG) << " Buffer: " << buffer_var_node->name_hint << " (Type: " << alloc->buffer->dtype << ")" << ", Start Offset: " << byte_offset @@ -813,7 +855,7 @@ class SharedMemoryRewriter : public StmtExprMutator { auto it = buffer_byte_offsets_.find(buffer_var.get()); ICHECK(it != buffer_byte_offsets_.end()) << "buffer_var = " << buffer_var->name_hint << ", dtype = " << dtype; - return indexdiv(it->second, dtype.bytes() * dtype.lanes()); + return ByteOffsetToElementOffset(it->second, dtype); } bool HasBufferOffset(const Var &buffer_var) { @@ -1419,30 +1461,14 @@ class SharedMemoryRewriter : public StmtExprMutator { } const AllocBufferNode *alloc = shmem_allocs_.at(var); - int64_t bytes_per_elem = static_cast( - alloc->buffer->dtype.bytes() * alloc->buffer->dtype.lanes()); - DataType size_dtype = DataType::Int(32); - if (!alloc->buffer->shape.empty()) { - size_dtype = alloc->buffer->shape[0].dtype(); - } - if (!size_dtype.is_int() && !size_dtype.is_uint()) { - size_dtype = DataType::Int(32); - } - - PrimExpr size_expr = make_const(size_dtype, bytes_per_elem); - for (const PrimExpr &extent : alloc->buffer->shape) { - PrimExpr e = extent; - if (e.dtype() != size_dtype) { - e = cast(size_dtype, e); - } - size_expr = size_expr * e; - } - info.size_dtype = size_dtype; - info.size_expr = size_expr; + PrimExpr elements = StorageElementsForShape(alloc->buffer->shape); + info.size_dtype = elements.dtype(); + info.size_expr = StorageBytesForElements(elements, alloc->buffer->dtype); auto const_size = GetRef(alloc).ConstantAllocationSize(); if (const_size.has_value()) { - info.const_size_bytes = const_size.value() * bytes_per_elem; + info.const_size_bytes = StorageBytesForConstantElements( + const_size.value(), alloc->buffer->dtype); } buf_infos.push_back(std::move(info)); diff --git a/src/transform/storage_rewrite.cc b/src/transform/storage_rewrite.cc index 3ed7081a2e..acb9983ab3 100644 --- a/src/transform/storage_rewrite.cc +++ b/src/transform/storage_rewrite.cc @@ -127,6 +127,8 @@ class LinearAccessPatternFinder final : public StmtExprVisitor { size_t level{0}; // allocation stmt const AllocBufferNode *alloc{nullptr}; + // Optional exact statement where the allocation should be re-emitted. + const Object *attach_scope_override{nullptr}; }; void VisitStmt_(const AllocBufferNode *op) final { @@ -136,6 +138,10 @@ class LinearAccessPatternFinder final : public StmtExprVisitor { AllocEntry entry; entry.alloc = op; entry.level = level; + if (op->annotations.find(tl::attr::kRoleScopedAlloc) != + op->annotations.end()) { + entry.attach_scope_override = op; + } // Since StorageRewrite occurs after StorageFlatten/FlattenBuffer, // all allocations specify the extent of physical dimensions, and // is 1 for flat memory spaces. @@ -583,6 +589,10 @@ class StoragePlanRewriter : public StmtExprMutator { } Stmt VisitStmt_(const AllocBufferNode *op) final { + if (attach_map_.count(op)) { + auto &svec = attach_map_[op]; + return MakeAttach(svec, Evaluate(0)); + } // AllocBuffer combines allocation and buffer declaration. // Storage rewrite may merge this allocation with others. if (auto it = alloc_map_.find(op->buffer->data.get()); @@ -939,6 +949,9 @@ class StoragePlanRewriter : public StmtExprMutator { const AllocBufferNode *alloc = entry.alloc; auto storage_scope = StorageScope::Create(GetPtrStorageScope(GetRef(var))); + const Object *attach_scope = entry.attach_scope_override != nullptr + ? entry.attach_scope_override + : effective_scope(storage_scope); StorageEntry *dst_entry = nullptr; // inplace detection if (detect_inplace) { @@ -949,8 +962,7 @@ class StoragePlanRewriter : public StmtExprMutator { InplaceOpVerifier visitor; StorageEntry *src_entry = alloc_map_.at(src); if (src_entry->scope == storage_scope && - src_entry->attach_scope_ == - effective_scope(storage_scope) && + src_entry->attach_scope_ == attach_scope && src_entry->elem_type == alloc->buffer->dtype.element_of() && visitor.Check(s.stmt, var, src)) { int64_t const_size = GetRef(alloc) @@ -970,10 +982,9 @@ class StoragePlanRewriter : public StmtExprMutator { } } if (dst_entry == nullptr) { - dst_entry = - FindAlloc(alloc, effective_scope(storage_scope), storage_scope, - entry.num_physical_dimensions, enable_reuse, - reuse_require_exact_matched_dtype); + dst_entry = FindAlloc(alloc, attach_scope, storage_scope, + entry.num_physical_dimensions, enable_reuse, + reuse_require_exact_matched_dtype); } dst_entry->allocs.emplace_back(alloc); alloc_map_[var] = dst_entry; diff --git a/src/transform/thread_storage_sync.cc b/src/transform/thread_storage_sync.cc index c605fce110..372e7affe9 100644 --- a/src/transform/thread_storage_sync.cc +++ b/src/transform/thread_storage_sync.cc @@ -536,6 +536,7 @@ class ConditionThreadPropertyChecker : public IRMutatorWithAnalyzer { PrimExpr VisitExpr_(const CallNode *op) final { if (op->op.same_as(builtin::tvm_access_ptr()) || + op->op.same_as(tl::access_ptr()) || op->op.same_as(builtin::address_of())) { current_.depends_on_runtime = true; // Do not mark local-scope tvm_access_ptr loads as non-block-uniform @@ -592,6 +593,9 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { Buffer buffer_name; /*! \brief The access data type */ DataType dtype; + /*! \brief Lower bound of the allocation alias in bytes, if known. */ + bool has_alias_byte_lower_bound = false; + PrimExpr alias_byte_lower_bound; /*! \brief The touched access range * * Has one IntSet for each index in the buffer being accessed. @@ -607,6 +611,14 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { * (e.g., inside a TMA load) and therefore multiple writes * among themselves should not force barriers between them. */ bool is_async_copy = false; + /*! \brief Whether this access is part of a TMA load/store operation. + * + * TMA load completion is governed by mbarrier waits, while TMA store + * completion is governed by tma_store_wait. A TMA store still needs prior + * normal shared-memory writes from other threads to be CTA-synchronized + * before the leader issues the async read. + */ + bool is_tma_access = false; /*! \brief Whether this access is part of an atomic RMW (e.g., atomic_add). * * If both sides of a dependency are atomic, we should not insert a thread @@ -634,18 +646,22 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { return buffer_var->type_annotation.as() && GetPtrStorageScope(buffer_var) == "shared.dyn"; } - PrimExpr AliasElemOffset(Var buffer_var, DataType dtype, - DataType index_dtype) const { + PrimExpr AliasByteOffset(Var buffer_var, DataType index_dtype) const { auto it = shared_memory_alias_byte_offsets_.find(buffer_var.get()); if (it == shared_memory_alias_byte_offsets_.end()) { return make_const(index_dtype, 0); } - int elem_bytes = dtype.bytes() * dtype.lanes(); - ICHECK_GT(elem_bytes, 0); PrimExpr byte_offset = it->second; if (byte_offset.dtype() != index_dtype) { byte_offset = Cast(index_dtype, byte_offset); } + return byte_offset; + } + PrimExpr AliasElemOffset(Var buffer_var, DataType dtype, + DataType index_dtype) const { + int elem_bytes = dtype.bytes() * dtype.lanes(); + ICHECK_GT(elem_bytes, 0); + PrimExpr byte_offset = AliasByteOffset(std::move(buffer_var), index_dtype); return indexdiv(byte_offset, make_const(index_dtype, elem_bytes)); } PrimExpr AddAliasElemOffset(Var buffer_var, DataType dtype, @@ -664,6 +680,72 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { } return index + elem_offset; } + DataType AccessPtrElementDType(const CallNode *op) const { + ICHECK(!op->args.empty()); + const auto *type_annotation = op->args[0].as(); + ICHECK(type_annotation != nullptr) + << "Expected tvm_access_ptr dtype argument to be a type annotation, " + << "but got " << op->args[0]; + static const Op &type_annotation_op = Op::Get("tirx.type_annotation"); + ICHECK(type_annotation->op.same_as(type_annotation_op)) + << "Expected tvm_access_ptr dtype argument to be tirx.type_annotation, " + << "but got " << type_annotation->op; + DataType dtype = type_annotation->dtype; + ICHECK(!dtype.is_handle()) + << "Expected tvm_access_ptr element dtype, but got handle"; + return dtype; + } + void AppendAccessPtrBufferLoad(const BufferLoadNode *load, + PrimExpr elem_extent, int64_t flag) { + Buffer buffer = load->buffer; + Var buf = buffer->data; + buffer_data_to_buffer_.Set(GetRef(buf.get()), buffer); + StorageScope scope = GetScope(buf); + if (!Enabled(buf.get(), scope)) { + return; + } + ICHECK(allow_append_); + + AccessEntry e{.cset = {constr_stack_}}; + e.threads = env_threads(); + e.dtype = load->dtype.element_of(); + e.buffer = buf; + e.buffer_name = buffer; + e.has_alias_byte_lower_bound = true; + e.alias_byte_lower_bound = AliasByteOffset(buf, DataType::Int(64)); + e.is_pointer_access = true; + e.is_async_copy = (tma_depth_ > 0 || cp_async_depth_ > 0); + e.is_tma_access = (tma_depth_ > 0); + e.is_atomic = (atomic_dst_ptr_depth_ > 0); + e.scope = scope; + + ICHECK_EQ(buffer->shape.size(), load->indices.size()); + for (size_t i = 0; i < buffer->shape.size(); ++i) { + PrimExpr min = AddAliasElemOffset(buf, buffer->dtype, load->indices[i]); + e.buffer_ranges.push_back( + Range::FromMinExtent(min, make_const(buffer->shape[i].dtype(), 1))); + } + + PrimExpr linear_offset = make_const(DataType::Int(64), 0); + if (!load->indices.empty()) { + linear_offset = buffer.OffsetOf(load->indices).back(); + linear_offset = AddAliasElemOffset(buf, buffer->dtype, linear_offset); + } + if (elem_extent.dtype() != linear_offset.dtype()) { + elem_extent = Cast(linear_offset.dtype(), elem_extent); + } + e.touched = {arith::IntSet::FromRange( + Range::FromMinExtent(linear_offset, elem_extent))}; + + if (flag & 1) { + e.type = kRead; + curr_stmt_.access.emplace_back(e); + } + if (flag & 2) { + e.type = kWrite; + curr_stmt_.access.emplace_back(e); + } + } void RecordSharedMemoryAlias(const Var &alias_var, const PrimExpr &value) { const auto *call = value.as(); if (call == nullptr || @@ -707,6 +789,8 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { e.buffer = buf; e.buffer_name = op->buffer; e.dtype = op->dtype.element_of(); + e.has_alias_byte_lower_bound = true; + e.alias_byte_lower_bound = AliasByteOffset(buf, DataType::Int(64)); for (const auto &index : op->indices) { PrimExpr physical_index = AddAliasElemOffset(buf, op->buffer->dtype, index); @@ -734,6 +818,8 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { e.buffer = buf; e.buffer_name = op->buffer; e.dtype = op->value.dtype().element_of(); + e.has_alias_byte_lower_bound = true; + e.alias_byte_lower_bound = AliasByteOffset(buf, DataType::Int(64)); for (const auto &index : op->indices) { PrimExpr physical_index = AddAliasElemOffset(buf, op->buffer->dtype, index); @@ -836,6 +922,12 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { new_touched.push_back(arith::EvalSet(touched, relax_map)); } e.touched = std::move(new_touched); + if (e.is_pointer_access) { + // buffer_ranges describe the single access_ptr statement. After + // loop relaxation, touched is the loop-wide byte range; keeping the + // per-iteration range can make disjointness proofs unsound. + e.buffer_ranges.clear(); + } } } } @@ -988,18 +1080,22 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { } void VisitExpr_(const CallNode *op) final { - // Mark async TMA load context so that tvm_access_ptr within the call - // can be tagged accordingly. - auto is_tma_load = [&]() { + // Mark async TMA context so that tvm_access_ptr within the call can be + // tagged accordingly. TMA load/store ordering is represented by mbarrier + // waits and tma_store_wait, not by storage_sync. + auto is_tma_access = [&]() { if (auto opt = op->op.as()) { const Op &call_op = opt.value(); return call_op.same_as(tl::tma_load()) || call_op.same_as(tl::tma_load_im2col()) || - call_op.same_as(tl::tma_load_multicast()); + call_op.same_as(tl::tma_load_multicast()) || + call_op.same_as(tl::tma_store()) || + call_op.same_as(tl::tma_load_gather4()) || + call_op.same_as(tl::tma_store_scatter4()); } return false; }(); - if (is_tma_load) { + if (is_tma_access) { tma_depth_++; for (const auto &a : op->args) { this->VisitExpr(a); @@ -1084,12 +1180,17 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { e.buffer = Downcast(buffer->data); e.buffer_name = buffer; e.buffer_ranges = buffer_ranges; + e.has_alias_byte_lower_bound = true; + e.alias_byte_lower_bound = + AliasByteOffset(GetRef(buffer_var), DataType::Int(64)); for (const auto &index : load->indices) { PrimExpr physical_index = AddAliasElemOffset( GetRef(buffer_var), buffer->dtype, index); e.touched.push_back(arith::IntSet::Vector(physical_index)); } e.is_pointer_access = true; + e.is_async_copy = (tma_depth_ > 0 || cp_async_depth_ > 0); + e.is_tma_access = (tma_depth_ > 0); e.is_atomic = (atomic_dst_ptr_depth_ > 0); e.type = kRead; e.scope = scope; @@ -1099,9 +1200,20 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { } else { ConstrVisitor::VisitExpr_(op); } + } else if (op->op.same_as(tl::access_ptr())) { + ICHECK_EQ(op->args.size(), 3U); + const auto *load = op->args[0].as(); + const auto *flag = op->args[2].as(); + ICHECK(load) << "tl.access_ptr base must be a BufferLoad, but got " + << op->args[0]; + ICHECK(flag) << "tl.access_ptr rw_mask must be an IntImm, but got " + << op->args[2]; + AppendAccessPtrBufferLoad(load, op->args[1], flag->value); + ConstrVisitor::VisitExpr_(load); + this->VisitExpr(op->args[1]); } else if (op->op.same_as(builtin::tvm_access_ptr())) { ICHECK_EQ(op->args.size(), 5U); - DataType dtype = op->args[0].dtype(); + DataType dtype = AccessPtrElementDType(op); const VarNode *buffer_var = op->args[1].as(); PrimExpr offset = AddAliasElemOffset(GetRef(buffer_var), dtype, op->args[2]); @@ -1157,6 +1269,9 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { e.dtype = dtype; e.buffer = GetRef(buffer_var); e.buffer_ranges = buffer_ranges; + e.has_alias_byte_lower_bound = true; + e.alias_byte_lower_bound = + AliasByteOffset(GetRef(buffer_var), DataType::Int(64)); e.is_pointer_access = true; e.is_atomic = (atomic_dst_ptr_depth_ > 0); e.touched = { @@ -1165,11 +1280,13 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { if (flag->value & 1) { e.type = kRead; e.is_async_copy = (tma_depth_ > 0 || cp_async_depth_ > 0); + e.is_tma_access = (tma_depth_ > 0); curr_stmt_.access.emplace_back(e); } if (flag->value & 2) { e.type = kWrite; e.is_async_copy = (tma_depth_ > 0 || cp_async_depth_ > 0); + e.is_tma_access = (tma_depth_ > 0); curr_stmt_.access.emplace_back(e); } } @@ -1432,49 +1549,125 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { return; syncs_inserted_.insert(obj); } - bool PointerAccessIsDisjoint(const AccessEntry &lhs, const AccessEntry &rhs) { - if (lhs.touched.size() != 1 || rhs.touched.size() != 1) { + bool AccessByteRangeIsDisjoint(const AccessEntry &lhs, const AccessEntry &rhs, + const ForNode *loop) { + auto is_integer_index = [](DataType dtype) { + return dtype.is_int() || dtype.is_uint(); + }; + auto element_bytes = [](DataType dtype) { + int bytes = dtype.bytes() * dtype.lanes(); + ICHECK_GT(bytes, 0) << "Expected concrete element dtype, but got " + << dtype; + return bytes; + }; + auto to_i64 = [](PrimExpr expr) { + DataType dtype = expr.dtype(); + if (dtype != DataType::Int(64)) { + expr = Cast(DataType::Int(64), expr); + } + return expr; + }; + auto byte_range = [&](PrimExpr elem_min, PrimExpr elem_extent, + DataType dtype) { + PrimExpr bytes = make_const(DataType::Int(64), element_bytes(dtype)); + PrimExpr min = to_i64(elem_min) * bytes; + PrimExpr max = (to_i64(elem_min + elem_extent) * bytes) - + make_const(DataType::Int(64), 1); + return std::pair{min, max}; + }; + auto collect_ranges = [&](const AccessEntry &access) { + std::vector> ranges; + if (access.buffer_ranges.size() == 1) { + const Range &range = access.buffer_ranges[0]; + if (is_integer_index(range->min.dtype()) && + is_integer_index(range->extent.dtype())) { + ranges.push_back(byte_range(range->min, range->extent, access.dtype)); + } + } + if (access.touched.size() == 1) { + PrimExpr min = access.touched[0].min(); + PrimExpr max = access.touched[0].max(); + if (is_integer_index(min.dtype()) && is_integer_index(max.dtype())) { + PrimExpr extent = max - min + make_const(min.dtype(), 1); + ranges.push_back(byte_range(min, extent, access.dtype)); + } + } + if (access.buffer_name.defined() && access.buffer_name->data.defined()) { + PrimExpr elems = make_const(DataType::Int(64), 1); + for (const PrimExpr &dim : access.buffer_name->shape) { + elems = elems * to_i64(dim); + } + Var data = Downcast(access.buffer_name->data); + PrimExpr min = AliasByteOffset(data, DataType::Int(64)); + PrimExpr max = min + + elems * make_const(DataType::Int(64), + access.buffer_name->dtype.bytes()) - + make_const(DataType::Int(64), 1); + ranges.push_back({min, max}); + } + return ranges; + }; + auto lhs_ranges = collect_ranges(lhs); + auto rhs_ranges = collect_ranges(rhs); + if (lhs_ranges.empty() || rhs_ranges.empty()) { return false; } - ConstrSet prev_cset{lhs.cset}; - ConstrSet curr_cset{rhs.cset}; - arith::Analyzer analyzer; - struct ThreadVarInfo { - const char *name_prev; - const char *name_curr; - } thread_vars[] = { - {"tx1", "tx2"}, - {"ty1", "ty2"}, - {"tz1", "tz2"}, + auto prove_disjoint = [&](PrimExpr lhs_min, PrimExpr lhs_max, + PrimExpr rhs_min, PrimExpr rhs_max) { + ConstrSet prev_cset{lhs.cset}; + ConstrSet curr_cset{rhs.cset}; + arith::Analyzer analyzer; + + if (loop != nullptr) { + PrimExpr step = make_const(loop->loop_var.dtype(), 1); + Map loop_shift_sub = { + {loop->loop_var, loop->loop_var + step}}; + rhs_min = Substitute(rhs_min, loop_shift_sub); + rhs_max = Substitute(rhs_max, loop_shift_sub); + curr_cset = curr_cset.Substitute(loop_shift_sub); + } + + struct ThreadVarInfo { + const char *name_prev; + const char *name_curr; + } thread_vars[] = { + {"tx1", "tx2"}, + {"ty1", "ty2"}, + {"tz1", "tz2"}, + }; + for (unsigned idx = 0; idx != 3; ++idx) { + auto &info = thread_vars[idx]; + Var old_prev_var = lhs.threads[lhs.threads.size() + idx - 3]->var; + Var old_curr_var = rhs.threads[rhs.threads.size() + idx - 3]->var; + Var prev_var(info.name_prev, old_prev_var.dtype()); + Var curr_var(info.name_curr, old_curr_var.dtype()); + lhs_min = Substitute(lhs_min, {{old_prev_var, prev_var}}); + lhs_max = Substitute(lhs_max, {{old_prev_var, prev_var}}); + prev_cset = prev_cset.Substitute({{old_prev_var, prev_var}}); + rhs_min = Substitute(rhs_min, {{old_curr_var, curr_var}}); + rhs_max = Substitute(rhs_max, {{old_curr_var, curr_var}}); + curr_cset = curr_cset.Substitute({{old_curr_var, curr_var}}); + } + prev_cset.Populate(analyzer); + curr_cset.Populate(analyzer); + lhs_min = analyzer.Simplify(lhs_min); + lhs_max = analyzer.Simplify(lhs_max); + rhs_min = analyzer.Simplify(rhs_min); + rhs_max = analyzer.Simplify(rhs_max); + + return analyzer.CanProve(lhs_max < rhs_min, + arith::ProofStrength::kSymbolicBound) || + analyzer.CanProve(rhs_max < lhs_min, + arith::ProofStrength::kSymbolicBound); }; - PrimExpr lhs_min = analyzer.Simplify(lhs.touched[0].min()); - PrimExpr lhs_max = analyzer.Simplify(lhs.touched[0].max()); - PrimExpr rhs_min = analyzer.Simplify(rhs.touched[0].min()); - PrimExpr rhs_max = analyzer.Simplify(rhs.touched[0].max()); - for (unsigned idx = 0; idx != 3; ++idx) { - auto &info = thread_vars[idx]; - Var old_prev_var = lhs.threads[lhs.threads.size() + idx - 3]->var; - Var old_curr_var = rhs.threads[rhs.threads.size() + idx - 3]->var; - Var prev_var(info.name_prev, old_prev_var.dtype()); - Var curr_var(info.name_curr, old_curr_var.dtype()); - lhs_min = Substitute(lhs_min, {{old_prev_var, prev_var}}); - lhs_max = Substitute(lhs_max, {{old_prev_var, prev_var}}); - prev_cset = prev_cset.Substitute({{old_prev_var, prev_var}}); - rhs_min = Substitute(rhs_min, {{old_curr_var, curr_var}}); - rhs_max = Substitute(rhs_max, {{old_curr_var, curr_var}}); - curr_cset = curr_cset.Substitute({{old_curr_var, curr_var}}); - } - prev_cset.Populate(analyzer); - curr_cset.Populate(analyzer); - - if (analyzer.CanProve(lhs_max < rhs_min, - arith::ProofStrength::kSymbolicBound)) { - return true; - } - if (analyzer.CanProve(rhs_max < lhs_min, - arith::ProofStrength::kSymbolicBound)) { - return true; + for (const auto &lhs_range : lhs_ranges) { + for (const auto &rhs_range : rhs_ranges) { + if (prove_disjoint(lhs_range.first, lhs_range.second, rhs_range.first, + rhs_range.second)) { + return true; + } + } } return false; } @@ -1486,6 +1679,10 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { output << " Buffer: " << access.buffer << "\n"; output << " Buffer Name: " << access.buffer_name << "\n"; output << " Data Type: " << access.dtype << "\n"; + if (access.has_alias_byte_lower_bound) { + output << " Alias Byte Lower Bound: " << access.alias_byte_lower_bound + << "\n"; + } std::string type_str; switch (access.type) { @@ -1561,6 +1758,7 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { output << "is_pointer_access=" << (access.is_pointer_access ? "true" : "false"); output << ", is_async_copy=" << (access.is_async_copy ? "true" : "false"); + output << ", is_tma_access=" << (access.is_tma_access ? "true" : "false"); LOG(WARNING) << output.str(); } @@ -1585,10 +1783,29 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { */ bool FindConflict(const AccessEntry &prev, const AccessEntry &curr, const ForNode *loop) { - // Special case: ignore conflicts between async-copy writes (e.g., TMA - // loads into shared memory). Multiple async writes do not require - // interspersed barriers among themselves. We still respect conflicts with - // reads to ensure visibility before consumption. + // TMA load visibility is governed by the associated mbarrier wait, and TMA + // store completion is governed by tma_store_wait. The important exceptions + // are normal shared-memory writes before a TMA store reads from shared + // memory, and normal writes after that TMA read. The first barrier makes + // all producer writes visible before the leader issues the async store; the + // second makes all threads wait until the leader-side tma_store_wait has + // completed before any thread reuses the source region. + if (prev.is_tma_access || curr.is_tma_access) { + bool normal_write_before_tma_read = + curr.is_tma_access && curr.type == kRead && !prev.is_tma_access && + prev.type == kWrite; + bool tma_read_before_normal_write = + prev.is_tma_access && prev.type == kRead && !curr.is_tma_access && + curr.type == kWrite; + if (normal_write_before_tma_read || tma_read_before_normal_write) { + // Fall through to regular conflict checks below. + } else { + return false; + } + } + + // Special case: ignore conflicts between async-copy writes. Multiple async + // writes do not require interspersed barriers among themselves. if (prev.type == kWrite && curr.type == kWrite && prev.is_async_copy && curr.is_async_copy) { return false; @@ -1605,6 +1822,41 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { return false; } + auto allocation_before_alias = [&](const AccessEntry &allocation, + const AccessEntry &alias) { + if (!allocation.buffer_name.defined() || + !allocation.buffer_name->data.defined() || + !alias.has_alias_byte_lower_bound) { + return false; + } + auto to_i64 = [](PrimExpr expr) { + if (expr.dtype() != DataType::Int(64)) { + expr = Cast(DataType::Int(64), expr); + } + return expr; + }; + PrimExpr elems = make_const(DataType::Int(64), 1); + for (const PrimExpr &dim : allocation.buffer_name->shape) { + elems = elems * to_i64(dim); + } + Var data = Downcast(allocation.buffer_name->data); + PrimExpr byte_min = AliasByteOffset(data, DataType::Int(64)); + PrimExpr byte_max = + byte_min + + elems * make_const(DataType::Int(64), + allocation.buffer_name->dtype.bytes()) - + make_const(DataType::Int(64), 1); + arith::Analyzer analyzer; + ConstrSet cset{allocation.cset}; + cset.Populate(analyzer); + return analyzer.CanProve(byte_max < alias.alias_byte_lower_bound, + arith::ProofStrength::kSymbolicBound); + }; + if (allocation_before_alias(prev, curr) || + allocation_before_alias(curr, prev)) { + return false; + } + if (prev.buffer_indices.size() != curr.buffer_indices.size()) { // They are not the same indices, should be conflict. return true; @@ -1612,10 +1864,11 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { if (prev.is_pointer_access || curr.is_pointer_access) { // For accesses created via tvm_access_ptr we may still be able to prove - // disjointness using their byte ranges. If both sides expose a touched - // interval and we can show they don't overlap, skip the conflict. - if (prev.is_pointer_access && curr.is_pointer_access && - PointerAccessIsDisjoint(prev, curr)) { + // disjointness using their byte ranges. This also matters when the + // other side is a normal BufferLoad/BufferStore: shared.dyn allocation + // merging keeps byte offsets in touched ranges, so pointer-vs-buffer + // comparisons should not be treated as unconditional aliases. + if (AccessByteRangeIsDisjoint(prev, curr, loop)) { return false; } // Otherwise fall back to the conservative answer: treat them as diff --git a/testing/python/language/test_tilelang_language_ldg.py b/testing/python/language/test_tilelang_language_ldg.py index 47d82c52d2..dce88e7913 100644 --- a/testing/python/language/test_tilelang_language_ldg.py +++ b/testing/python/language/test_tilelang_language_ldg.py @@ -157,6 +157,33 @@ def ldg32_pred_kernel(X, Y): torch.testing.assert_close(Y, Y_ref, atol=1e-5, rtol=1e-5) +@tilelang.testing.requires_cuda +def test_lds32_predicated_codegen(): + """Test that lds32 with predicate generates tl::load_shared_32_conditional(ptr, pred).""" + + @tilelang.jit + def lds32_pred_kernel(Y): + Y: T.Tensor[[32], T.uint32] + + with T.Kernel(1, threads=32): + tx = T.get_thread_binding() + scratch = T.alloc_shared((32,), T.uint32) + scratch[tx] = T.Cast(T.uint32, tx + 1) + Y[tx] = T.lds32(scratch[tx], pred=tx < 16) + + Y = torch.empty(32, dtype=torch.uint32, device="cuda") + + lds32_pred_kernel(Y) + src = lds32_pred_kernel.get_kernel_source() + print("=== lds32 predicated codegen ===") + print(src) + assert "load_shared_32_conditional" in src, "Expected load_shared_32_conditional call in generated CUDA source" + + expected = torch.arange(1, 33, dtype=torch.int64) + expected[16:] = 0 + torch.testing.assert_close(Y.cpu().to(torch.int64), expected) + + @tilelang.testing.requires_cuda def test_ldg64_predicated_codegen(): """Test that ldg64 with predicate generates tl::load_global_64_conditional(ptr, pred) in CUDA source.""" diff --git a/testing/python/language/test_tilelang_language_tma_copy.py b/testing/python/language/test_tilelang_language_tma_copy.py index a87e61992d..5743cec26c 100644 --- a/testing/python/language/test_tilelang_language_tma_copy.py +++ b/testing/python/language/test_tilelang_language_tma_copy.py @@ -304,6 +304,94 @@ def main( return main +def fp4_tma_copy_unpacked_smem_full_swizzle_load(M=128, N=256, block_M=64, block_N=128): + from tilelang.layout import make_full_bank_swizzled_layout + + @T.prim_func + def main( + A: T.Tensor((M, N), T.float4_e2m1fn), + ): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): + A_shared = T.alloc_shared((block_M, block_N), T.float4_e2m1_unpacked) + T.annotate_layout({A_shared: make_full_bank_swizzled_layout(A_shared)}) + mbar = T.alloc_barrier(128) + T.tma_copy( + A[by * block_M, bx * block_N], + A_shared, + barrier=mbar, + ) + T.barrier_arrive(mbar) + T.mbarrier_wait_parity(mbar, 0) + + return main + + +def fp4_tma_copy_packed_smem_sm120_fp4_layout(M=128, N=256, block_M=64, block_N=128): + from tilelang.layout import make_sm120_fp4_smem_layout + + @T.prim_func + def main( + A: T.Tensor((M, N), T.float4_e2m1fn), + ): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): + A_shared = T.alloc_shared((block_M, block_N), T.float4_e2m1fn) + T.annotate_layout({A_shared: make_sm120_fp4_smem_layout(A_shared)}) + mbar = T.alloc_barrier(128) + T.tma_copy( + A[by * block_M, bx * block_N], + A_shared, + barrier=mbar, + ) + T.barrier_arrive(mbar) + T.mbarrier_wait_parity(mbar, 0) + + return main + + +def fp4_tma_copy_packed_global_to_uint8_smem_sm120_fp4_layout(M=128, N=256, block_M=64, block_N=128): + from tilelang.layout import make_sm120_fp4_smem_layout + + @T.prim_func + def main( + A: T.Tensor((M, N), T.float4_e2m1fn), + ): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): + A_shared = T.alloc_shared((block_M, block_N // 2), "uint8") + T.annotate_layout({A_shared: make_sm120_fp4_smem_layout(A_shared)}) + mbar = T.alloc_barrier(128) + T.tma_copy( + A[by * block_M : (by + 1) * block_M, bx * block_N : (bx + 1) * block_N], + A_shared, + barrier=mbar, + ) + T.barrier_arrive(mbar) + T.mbarrier_wait_parity(mbar, 0) + + return main + + +def fp4_tma_copy_unpacked_smem_sm120_fp4_layout(M=128, N=256, block_M=64, block_N=128): + from tilelang.layout import make_sm120_fp4_smem_layout + + @T.prim_func + def main( + A: T.Tensor((M, N), T.float4_e2m1fn), + ): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): + A_shared = T.alloc_shared((block_M, block_N), T.float4_e2m1_unpacked) + T.annotate_layout({A_shared: make_sm120_fp4_smem_layout(A_shared)}) + mbar = T.alloc_barrier(128) + T.tma_copy( + A[by * block_M, bx * block_N], + A_shared, + barrier=mbar, + ) + T.barrier_arrive(mbar) + T.mbarrier_wait_parity(mbar, 0) + + return main + + def fp4_tma_copy_unpacked_smem_store(M=128, N=256, block_M=64, block_N=128): @T.prim_func def main( @@ -326,15 +414,25 @@ def main( def _fp4_tma_descriptor_init_block(host_source, desc_name): - marker = f"[0].v_ptr) = {desc_name};" start = host_source.find(marker) - assert start >= 0, f"Missing {desc_name} TensorMap initialization" + assert start >= 0, f"Missing {desc_name} TensorMap initialization; available descriptors: {_fp4_tma_descriptor_names(host_source)}" end = host_source.find("TVMFFIFunctionCall(__tvm_tensormap_create_tiled_packed", start) assert end >= 0, f"Missing {desc_name} TensorMap creation call" return host_source[start:end] +def _fp4_tma_descriptor_names(host_source): + import re + + names = [] + for match in re.finditer(r"\[0\]\.v_ptr\)\s*=\s*(\w+_desc);", host_source): + name = match.group(1) + if name not in names: + names.append(name) + return names + + def _fp4_tma_stack_int(block, index): import re @@ -343,7 +441,7 @@ def _fp4_tma_stack_int(block, index): return int(match.group(1)) -def _assert_fp4_packed_tma_descriptor(host_source, desc_name): +def _assert_fp4_packed_tma_descriptor(host_source, desc_name, *, expect_swizzle=0): expected_tma_args = { 1: 13, # CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B 2: 2, @@ -356,7 +454,7 @@ def _assert_fp4_packed_tma_descriptor(host_source, desc_name): 10: 1, 11: 1, 12: 0, - 13: 2, # CU_TENSOR_MAP_SWIZZLE_64B + 13: expect_swizzle, 14: 2, 15: 0, } @@ -389,7 +487,9 @@ def run_fp4_tma_copy_roundtrip(): assert "tl::tma_load" in device_source assert "tl::tma_store" in device_source assert host_source.count("__tvm_tensormap_create_tiled") >= 2 - for desc_name in ("A_desc", "B_desc"): + desc_names = _fp4_tma_descriptor_names(host_source) + assert len(desc_names) == 2, f"Expected two TMA descriptors, got {desc_names}" + for desc_name in desc_names: _assert_fp4_packed_tma_descriptor(host_source, desc_name) a = torch.randint(-128, 128, (M, N // 2), device="cuda", dtype=torch.int8) @@ -398,6 +498,8 @@ def run_fp4_tma_copy_roundtrip(): def test_fp4_unpacksmem_tma_descriptor_uses_align16b(): + import re + program = fp4_tma_copy_unpacked_smem_load() artifact = tilelang.lower( program, @@ -408,8 +510,70 @@ def test_fp4_unpacksmem_tma_descriptor_uses_align16b(): device_ir = str(artifact.device_mod) assert 'T.handle("float4_e2m1fn", "global")' in host_ir assert 'A_shared = T.alloc_buffer((8192,), "custom[float4_e2m1_unpacked]"' in device_ir - assert '["__tvm_tensormap_create_tiled", A_desc, 14,' in host_ir - assert 'T.call_packed("__tvm_tensormap_create_tiled", A_desc, 14,' in host_ir + assert re.search(r'\["__tvm_tensormap_create_tiled", \w+, 14,', host_ir) + assert re.search(r'T\.call_packed\("__tvm_tensormap_create_tiled", \w+, 14,', host_ir) + + +def test_fp4_unpacksmem_full_swizzle_tma_descriptor_uses_align16b_128b_swizzle(): + program = fp4_tma_copy_unpacked_smem_full_swizzle_load() + artifact = tilelang.lower( + program, + target={"kind": "cuda", "arch": "sm_100"}, + enable_device_compile=False, + ) + host_ir = str(artifact.host_mod) + device_ir = str(artifact.device_mod) + assert 'A_shared = T.alloc_buffer((8192,), "custom[float4_e2m1_unpacked]"' in device_ir + assert "T.tma_load(tma_offset_0_desc" in device_ir + assert ( + '"__tvm_tensormap_create_tiled", tma_offset_0_desc, 14, 2, ' + "T.handle_add_byte_offset(A, 0), 256, 128, 1, 128, 128, 64, 1, 1, 0, 3, 2, 0" + ) in host_ir + + +def test_sm120_fp4_smem_layout_selects_packed_sw64_for_k128(): + program = fp4_tma_copy_packed_smem_sm120_fp4_layout() + artifact = tilelang.lower( + program, + target={"kind": "cuda", "arch": "sm_100"}, + enable_device_compile=False, + ) + host_ir = str(artifact.host_mod) + assert ( + '"__tvm_tensormap_create_tiled", tma_offset_0_desc, 13, 2, ' + "T.handle_add_byte_offset(A, 0), 256, 128, 1, 128, 128, 64, 1, 1, 0, 2, 2, 0" + ) in host_ir + + +def test_fp4_packed_global_to_uint8_smem_tma_uses_packed_descriptor_sw64_for_k128(): + program = fp4_tma_copy_packed_global_to_uint8_smem_sm120_fp4_layout() + artifact = tilelang.lower( + program, + target={"kind": "cuda", "arch": "sm_100"}, + enable_device_compile=False, + ) + host_ir = str(artifact.host_mod) + device_ir = str(artifact.device_mod) + assert 'A_shared = T.alloc_buffer((4096,), "uint8"' in device_ir + assert "T.tma_load(tma_offset_0_desc" in device_ir + assert ( + '"__tvm_tensormap_create_tiled", tma_offset_0_desc, 13, 2, ' + "T.handle_add_byte_offset(A, 0), 256, 128, 1, 128, 128, 64, 1, 1, 0, 2, 2, 0" + ) in host_ir + + +def test_sm120_fp4_smem_layout_selects_unpacksmem_sw128_for_k128(): + program = fp4_tma_copy_unpacked_smem_sm120_fp4_layout() + artifact = tilelang.lower( + program, + target={"kind": "cuda", "arch": "sm_100"}, + enable_device_compile=False, + ) + host_ir = str(artifact.host_mod) + assert ( + '"__tvm_tensormap_create_tiled", tma_offset_0_desc, 14, 2, ' + "T.handle_add_byte_offset(A, 0), 256, 128, 1, 128, 128, 64, 1, 1, 0, 3, 2, 0" + ) in host_ir def test_fp4_unpacksmem_tma_store_is_rejected(): @@ -422,6 +586,68 @@ def test_fp4_unpacksmem_tma_store_is_rejected(): ) +def tma_copy_projected_global_view_kernel(): + @T.prim_func + def main( + A: T.Tensor((1, 4, 128, 128), T.float16), + B: T.Tensor((4,), T.float16), + ): + with T.Kernel(4, threads=128) as head: + A_shared = T.alloc_shared((128, 64), T.float16) + mbar = T.alloc_barrier(128) + T.tma_copy(A[0, head, 0:128, 0:64], A_shared, barrier=mbar) + T.barrier_arrive(mbar) + T.mbarrier_wait_parity(mbar, 0) + if T.get_thread_binding() == 0: + B[head] = A_shared[0, 0] + + return main + + +@tilelang.testing.requires_cuda +@tilelang.testing.requires_cuda_compute_version_ge(9, 0) +def test_tma_copy_projects_static_axes_but_keeps_dynamic_unit_axes(): + import re + + kernel = tilelang.compile( + tma_copy_projected_global_view_kernel(), + out_idx=[1], + pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True}, + ) + device_source = kernel.get_kernel_source() + host_source = kernel.get_host_source() + + marker = "[0].v_ptr) = A_desc;" + start = host_source.find(marker) + assert start >= 0, "Missing A_desc TensorMap initialization" + end = host_source.find("TVMFFIFunctionCall(__tvm_tensormap_create_tiled_packed", start) + assert end >= 0, "Missing A_desc TensorMap creation call" + block = host_source[start:end] + + def stack_int(index): + match = re.search(rf"\[{index}\]\.v_int64\)\s*=\s*\(\(int64_t\)(-?\d+)\);", block) + assert match, f"Missing stack[{index}] integer assignment in:\n{block}" + return int(match.group(1)) + + expected_tma_args = { + 1: 6, # CU_TENSOR_MAP_DATA_TYPE_FLOAT16 + 2: 3, # dynamic head plus the two non-unit tile axes + 4: 128, # global_shape[0], reversed innermost dimension + 5: 128, # global_shape[1] + 6: 4, # global_shape[2], dynamic head dimension kept in descriptor + 7: 2, # raw innermost stride, ignored by CUDA encode + 8: 256, # row stride in bytes: 128 fp16 elements + 9: 32768, # head stride in bytes: 128 * 128 fp16 elements + 10: 64, # smem_box[0] + 11: 128, # smem_box[1] + 12: 1, # smem_box[2], dynamic unit axis + } + for index, expected in expected_tma_args.items(): + assert stack_int(index) == expected + + assert "tl::tma_load" in device_source + + def test_copy_prefer_tma_lowers_as_synchronous_tma_load(): @T.prim_func def main(x: T.Tensor((128, 32), T.float32)): @@ -454,7 +680,9 @@ def run_fp4_tma_copy_unpacked_smem_load(): device_source = kernel.get_kernel_source() host_source = kernel.get_host_source() assert "CUtensorMap" in device_source - _assert_fp4_unpacked_tma_descriptor(host_source, "A_desc") + desc_names = _fp4_tma_descriptor_names(host_source) + assert len(desc_names) == 1, f"Expected one TMA descriptor, got {desc_names}" + _assert_fp4_unpacked_tma_descriptor(host_source, desc_names[0]) # 64 x 128 logical FP4 elems -> 4096 transaction bytes (4b/elem), not 8192. assert "expect_transaction(4096)" in device_source diff --git a/testing/python/language/test_tilelang_language_tma_gather_scatter.py b/testing/python/language/test_tilelang_language_tma_gather_scatter.py index 50634c672d..58ffd0ae90 100644 --- a/testing/python/language/test_tilelang_language_tma_gather_scatter.py +++ b/testing/python/language/test_tilelang_language_tma_gather_scatter.py @@ -87,6 +87,41 @@ def run_gather_scatter(N=64, K=64, K_box=64): torch.testing.assert_close(Dst, expected) +def gather4_fp4_unpack_program(N: int = 64, K_box: int = 128): + @T.prim_func + def main( + Src: T.Tensor((N, K_box), T.float4_e2m1fn), + ): + with T.Kernel(1, 1, threads=128): + smem = T.alloc_shared((4, K_box), T.float4_e2m1_unpacked) + mbar = T.alloc_barrier(1) + + if T.shuffle_elect(128): + T.mbarrier_expect_tx(mbar, T.tma_gather4_bytes(K_box, "float4_e2m1_unpacked")) + T.tma_gather4(Src, smem, 0, [0, 7, 23, 42], barrier=mbar) + T.barrier_arrive(mbar) + T.mbarrier_wait_parity(mbar, 0) + + return main + + +def test_gather4_fp4_unpack_descriptor_codegen(): + import re + + artifact = tilelang.lower( + gather4_fp4_unpack_program(), + target={"kind": "cuda", "arch": "sm_100"}, + enable_device_compile=False, + ) + host_ir = str(artifact.host_mod) + device_ir = str(artifact.device_mod) + assert 'T.handle("float4_e2m1fn", "global")' in host_ir + assert 'smem = T.alloc_buffer((512,), "custom[float4_e2m1_unpacked]"' in device_ir + assert 'T.call_extern("handle", "tl.tma_load_gather4"' in device_ir + assert re.search(r'\["__tvm_tensormap_create_tiled", \w+, 14,', host_ir) + assert re.search(r"T\.mbarrier_expect_tx\([^,]+, 256\)", device_ir) + + @requires_sm100 def test_gather_scatter_basic(): run_gather_scatter(N=64, K=64, K_box=64) diff --git a/testing/python/language/test_tilelang_language_view.py b/testing/python/language/test_tilelang_language_view.py index c3f0395cc1..6acd951082 100644 --- a/testing/python/language/test_tilelang_language_view.py +++ b/testing/python/language/test_tilelang_language_view.py @@ -1,6 +1,7 @@ import tilelang.language as T from tilelang import tvm as tvm import tilelang.testing +import tilelang.layout import tilelang as tl import pytest import torch @@ -91,6 +92,182 @@ def test_view_subbyte_dtype_change(): assert A_viewed.data.same_as(A.data) +def test_view_accepts_explicit_strides(): + A = tvm.tirx.decl_buffer((4, 8), "float8_e4m3fn", name="A", strides=[16, 1]) + A_viewed = T.view(A, (4, 2), dtype=T.uint32, strides=(8, 1), elem_offset=4) + assert str(A_viewed.dtype) == "uint32" + assert tuple(int(dim) for dim in A_viewed.shape) == (4, 2) + assert tuple(int(stride) for stride in A_viewed.strides) == (8, 1) + assert int(A_viewed.elem_offset) == 4 + assert A_viewed.data.same_as(A.data) + + +def _int_tuple(values): + return tuple(int(value) for value in values) + + +def test_layout_reshape_preserves_packed_subtype_lane(): + layout = T.Layout((2, 2, 8), lambda a, b, c: [a, b, c]) + u8_view = layout.reshape((2, 2, 16), 16, 8) + assert _int_tuple(u8_view.get_output_shape()) == (2, 2, 8, 2) + assert _int_tuple(u8_view.map_forward_index([0, 0, 0])) == (0, 0, 0, 0) + assert _int_tuple(u8_view.map_forward_index([0, 0, 1])) == (0, 0, 0, 1) + assert _int_tuple(u8_view.map_forward_index([0, 0, 2])) == (0, 0, 1, 0) + + u16_view = u8_view.reshape((2, 2, 8), 8, 16) + assert _int_tuple(u16_view.get_output_shape()) == (2, 2, 8) + assert _int_tuple(u16_view.map_forward_index([0, 0, 1])) == (0, 0, 1) + + +def test_fragment_reshape_preserves_packed_subtype_lane(): + fragment = T.Fragment( + (2, 2, 8), + forward_thread_fn=lambda a, b, c: a * 2 + b, + forward_index_fn=lambda a, b, c: c, + ) + u8_view = fragment.reshape((2, 2, 16), 16, 8) + assert _int_tuple(u8_view.get_output_shape()) == (8, 2) + assert _int_tuple(u8_view.map_forward_index([0, 0, 0])) == (0, 0) + assert _int_tuple(u8_view.map_forward_index([0, 0, 1])) == (0, 1) + assert _int_tuple(u8_view.map_forward_index([0, 0, 2])) == (1, 0) + assert _int_tuple(u8_view.map_forward_thread([1, 1, 15])) == (3,) + + u16_view = u8_view.reshape((2, 2, 8), 8, 16) + assert _int_tuple(u16_view.get_output_shape()) == (8,) + assert _int_tuple(u16_view.map_forward_index([0, 0, 1])) == (1,) + assert _int_tuple(u16_view.map_forward_thread([1, 1, 7])) == (3,) + + +def annotated_fragment_layout_on_dtype_changing_view_test(): + @T.prim_func + def main(B: T.Tensor((4,), T.uint32)): + with T.Kernel(1, threads=32) as _: + sf = T.alloc_fragment((16,), T.float8_e4m3fn) + sf_words = T.view(sf, (4,), dtype=T.uint32) + T.annotate_layout( + { + sf_words: T.Fragment( + (4,), + forward_thread_fn=lambda i: i, + forward_index_fn=lambda i: 0, + ) + } + ) + + for i in T.Parallel(4): + sf_words[i] = T.Cast(T.uint32, i + 1) + B[i] = sf_words[i] + + return main + + +@tilelang.testing.requires_cuda +def test_annotated_fragment_layout_on_dtype_changing_view_compile(): + program = annotated_fragment_layout_on_dtype_changing_view_test() + kernel = tl.compile(program, out_idx=-1) + src = kernel.get_kernel_source() + assert "uint sf_words[1]" in src + out = kernel() + torch.testing.assert_close(out.cpu(), torch.tensor([1, 2, 3, 4], dtype=torch.uint32)) + + +def layout_view_ldmatrix_pointer_test(): + @T.prim_func + def main( + A: T.Tensor((16, 16), T.float16), + B: T.Tensor((4,), T.uint32), + ): + with T.Kernel(1, threads=32) as _: + S = T.alloc_shared((16, 16), T.float16) + V = T.view(S, (16, 16), dtype=T.float16) + R = T.alloc_local((4,), T.uint32) + T.annotate_layout_view({V: T.Layout((16, 16), lambda i, j: [j, i])}) + + tx = T.get_thread_binding() + S[tx // 16, tx % 16] = A[tx // 16, tx % 16] + T.ptx_ldmatrix( + T.bool(False), + 4, + T.access_ptr(V[0, 1], "r", extent=16), + T.access_ptr(R[0], "w", extent=4), + ) + if tx < 4: + B[tx] = R[tx] + + return main + + +def test_layout_view_applies_to_ldmatrix_access_ptr(): + artifact = tl.lower( + layout_view_ldmatrix_pointer_test(), + target="cuda", + enable_device_compile=False, + ) + source = str(artifact.kernel_source) + assert "tl::ptx_ldmatrix_x4((&(V[16])), (&(R[0])))" in source + + +def replicated_fragment_word_view_test(): + @T.prim_func + def main(B: T.Tensor((32, 4), T.uint32)): + with T.Kernel(1, threads=32) as _: + p = T.alloc_fragment((32,), T.float4_e2m1fn) + p_words = T.view(p, (4,), dtype=T.uint32) + T.annotate_layout({p_words: tilelang.layout.make_fully_replicated_layout_fragment(p_words, 32)}) + tx = T.get_thread_binding() + p_words[0] = T.Cast(T.uint32, tx + 1) + p_words[1] = T.Cast(T.uint32, tx + 2) + p_words[2] = T.Cast(T.uint32, tx + 3) + p_words[3] = T.Cast(T.uint32, tx + 4) + B[tx, 0] = p_words[0] + B[tx, 1] = p_words[1] + B[tx, 2] = p_words[2] + B[tx, 3] = p_words[3] + + return main + + +def ws_consumer_fragment_word_view_thread_range_test(): + @T.prim_func + def main(B: T.Tensor((384,), T.uint32)): + with T.Kernel(1, threads=384) as _: + tx = T.get_thread_binding() + role_tx = tx + scratch = T.alloc_shared((32,), T.uint32) + with T.ws(0): + if tx < 32: + scratch[tx] = T.Cast(T.uint32, tx) + with T.ws(1, 2): + sf = T.alloc_fragment((8,), T.float8_e4m3fn, role_scoped=True) + sf_words = T.view(sf, (2,), dtype=T.uint32) + sf_words[0] = T.lds32(scratch[role_tx & 31]) + B[tx] = sf_words[0] + + return main + + +@tilelang.testing.requires_cuda +@tilelang.testing.requires_cuda_compute_version_ge(10, 0) +def test_replicated_fragment_fp4_word_view_runs(): + program = replicated_fragment_word_view_test() + kernel = tl.compile(program, out_idx=-1) + src = kernel.get_kernel_source() + assert "uint p_words[4]" in src + out = kernel().cpu().to(torch.int64) + expected = torch.arange(32, dtype=torch.int64).reshape(32, 1) + torch.tensor([1, 2, 3, 4], dtype=torch.int64) + torch.testing.assert_close(out, expected) + + +@tilelang.testing.requires_cuda +@tilelang.testing.requires_cuda_compute_version_ge(10, 0) +def test_ws_consumer_fragment_word_view_uses_consumer_thread_range(): + program = ws_consumer_fragment_word_view_thread_range_test() + kernel = tl.compile(program, out_idx=-1) + src = kernel.get_kernel_source() + assert "load_shared_32" in src + assert "threadIdx.x) - 256" not in src + + def fp4_to_uint8_view_test(rows_per_cta=16, mask_k=256): @T.prim_func def main( diff --git a/testing/python/transform/test_tilelang_transform_disable_memory_reuse.py b/testing/python/transform/test_tilelang_transform_disable_memory_reuse.py index 25a5fba466..156fe44fa0 100644 --- a/testing/python/transform/test_tilelang_transform_disable_memory_reuse.py +++ b/testing/python/transform/test_tilelang_transform_disable_memory_reuse.py @@ -17,6 +17,19 @@ N = 1024 +def _extract_merged_smem_offsets(src: str) -> list[int]: + """Extract merged shared-memory byte offsets from generated CUDA code.""" + # Alias-preserving lowering emits: + # void* a_shared = ((void*)((char*)buf_dyn_shmem + 0)); + alias_pattern = r"void\*\s+\w+\s*=\s*\(\(void\*\)\(\(char\*\)buf_dyn_shmem\s*\+\s*(\d+)\)\);" + # Direct merged-buffer lowering emits access patterns such as: + # buf_dyn_shmem)[1024]) + direct_pattern = r"buf_dyn_shmem\)\[(\d+)\]" + offsets = {int(m) for m in re.findall(alias_pattern, src)} + offsets.update(int(m) for m in re.findall(direct_pattern, src)) + return sorted(offsets) + + def _make_data_integrity_kernel(): @tilelang.jit( @@ -124,20 +137,8 @@ def test_disable_reuse_no_overlap(): src_no_reuse = kernel_no_reuse.get_kernel_source() src_reuse = kernel_reuse.get_kernel_source() - def extract_smem_offsets(src: str) -> list[int]: - """Extract merged shared-memory offsets from generated code.""" - # Alias-preserving lowering emits: - # void* a_shared = ((void*)((char*)buf_dyn_shmem + 0)); - alias_pattern = r"void\*\s+\w+\s*=\s*\(\(void\*\)\(\(char\*\)buf_dyn_shmem\s*\+\s*(\d+)\)\);" - # Direct merged-buffer lowering emits access patterns such as: - # buf_dyn_shmem)[1024]) - direct_pattern = r"buf_dyn_shmem\)\[(\d+)\]" - offsets = {int(m) for m in re.findall(alias_pattern, src)} - offsets.update(int(m) for m in re.findall(direct_pattern, src)) - return sorted(offsets) - - offsets_no_reuse = extract_smem_offsets(src_no_reuse) - offsets_reuse = extract_smem_offsets(src_reuse) + offsets_no_reuse = _extract_merged_smem_offsets(src_no_reuse) + offsets_reuse = _extract_merged_smem_offsets(src_reuse) # With reuse disabled: must have at least 2 distinct offsets (buffers not merged) assert len(offsets_no_reuse) >= 2, f"Expected >=2 distinct smem offsets with reuse disabled, got {offsets_no_reuse}" @@ -146,7 +147,37 @@ def extract_smem_offsets(src: str) -> list[int]: assert len(offsets_reuse) == 1, f"Expected 1 smem offset with reuse enabled, got {offsets_reuse}" +@tilelang.testing.requires_cuda +def test_disable_reuse_fp4_uses_packed_storage_size(): + """FP4 shared-memory merge planning should count two logical elements per byte.""" + + @T.prim_func + def main(C: T.Tensor((1,), T.uint8)): + with T.Kernel(1, threads=128): + a_shared = T.alloc_shared((128,), T.float4_e2m1fn) + b_shared = T.alloc_shared((128,), T.float4_e2m1fn) + a_bytes = T.view(a_shared, (64,), dtype=T.uint8) + b_bytes = T.view(b_shared, (64,), dtype=T.uint8) + + for i in T.Parallel(64): + a_bytes[i] = T.uint8(1) + for i in T.Parallel(64): + b_bytes[i] = a_bytes[i] + C[0] = b_bytes[0] + + kernel = tilelang.compile( + main, + out_idx=[0], + target="cuda", + pass_configs={PassConfigKey.TL_DISABLE_SHARED_MEMORY_REUSE: True}, + ) + offsets = _extract_merged_smem_offsets(kernel.get_kernel_source()) + assert 0 in offsets, f"Expected first FP4 shared buffer at byte offset 0, got {offsets}" + assert 64 in offsets, f"Expected second FP4 shared buffer at byte offset 64, got {offsets}" + + if __name__ == "__main__": test_disable_reuse_data_integrity() test_disable_reuse_no_overlap() + test_disable_reuse_fp4_uses_packed_storage_size() print("All tests passed!") diff --git a/tilelang/language/__init__.py b/tilelang/language/__init__.py index 6e523c8cd5..0863bd70cd 100644 --- a/tilelang/language/__init__.py +++ b/tilelang/language/__init__.py @@ -123,6 +123,7 @@ from .builtin import ds_read_tr8_b64 as ds_read_tr8_b64 # noqa: F401 from .builtin import pack_b16 as pack_b16 # noqa: F401 from .builtin import ldg32 as ldg32 # noqa: F401 +from .builtin import lds32 as lds32 # noqa: F401 from .builtin import ldg64 as ldg64 # noqa: F401 from .builtin import ldg128 as ldg128 # noqa: F401 from .builtin import ldg256 as ldg256 # noqa: F401 @@ -145,6 +146,7 @@ from .annotations import ( # noqa: F401 use_swizzle, annotate_layout, + annotate_layout_view, annotate_safe_value, annotate_l2_hit_ratio, annotate_restrict_buffers, diff --git a/tilelang/language/allocate.py b/tilelang/language/allocate.py index 7a9c264dcc..e280c68919 100644 --- a/tilelang/language/allocate.py +++ b/tilelang/language/allocate.py @@ -31,6 +31,9 @@ from .proxy import Tensor, ptr as _ptr_sentinel +_ROLE_SCOPED_ALLOC_ATTR = "tl.role_scoped_alloc" + + def alloc_shared(shape: ShapeType, dtype: DType, scope="shared.dyn") -> Buffer: """Allocate a shared memory buffer for inter-thread communication. @@ -49,43 +52,63 @@ def alloc_shared(shape: ShapeType, dtype: DType, scope="shared.dyn") -> Buffer: return T.sblock_alloc_buffer(shape, dtype, scope=scope) -def alloc_local(shape: ShapeType, dtype: DType, scope="local") -> Buffer: +def alloc_local(shape: ShapeType, dtype: DType, scope="local", role_scoped: bool = False) -> Buffer: """Allocate a local memory buffer for thread-private storage. Args: shape (tuple): The shape of the buffer to allocate dtype (str): The data type of the buffer (e.g., 'float32', 'int32') scope (str, optional): The memory scope. Defaults to "local" + role_scoped (bool, optional): Emit the allocation at the current + statement position instead of as an SBlock allocation. Returns: T.Buffer: A TVM buffer object allocated in local memory """ + if role_scoped: + return T.alloc_buffer(shape, dtype, scope=scope, annotations={_ROLE_SCOPED_ALLOC_ATTR: True}) return T.sblock_alloc_buffer(shape, dtype, scope=scope) -def alloc_fragment(shape: ShapeType, dtype: DType, scope="local.fragment") -> Buffer: +def alloc_fragment(shape: ShapeType, dtype: DType, scope="local.fragment", role_scoped: bool = False) -> Buffer: """Allocate a fragment memory buffer for specialized operations. Args: shape (tuple): The shape of the buffer to allocate dtype (str): The data type of the buffer (e.g., 'float32', 'int32') scope (str, optional): The memory scope. Defaults to "local.fragment" + role_scoped (bool, optional): Emit the allocation at the current + statement position instead of as an SBlock allocation. Returns: T.Buffer: A TVM buffer object allocated in fragment memory """ + if role_scoped: + return T.alloc_buffer(shape, dtype, scope=scope, annotations={_ROLE_SCOPED_ALLOC_ATTR: True}) return T.sblock_alloc_buffer(shape, dtype, scope=scope) @overload -def alloc_var(dtype: DType, init: PrimExpr | int | float, scope: str = "local.var") -> Buffer: ... +def alloc_var(dtype: DType, init: PrimExpr | int | float, scope: str = "local.var", *, role_scoped: bool = False) -> Buffer: ... @overload -def alloc_var(dtype: DType, scope: str = "local.var", *, init: PrimExpr | int | float | None = None) -> Buffer: ... - - -def alloc_var(dtype: DType, *args, scope: str = "local.var", init: PrimExpr | int | float | None = None) -> Buffer: +def alloc_var( + dtype: DType, + scope: str = "local.var", + *, + init: PrimExpr | int | float | None = None, + role_scoped: bool = False, +) -> Buffer: ... + + +def alloc_var( + dtype: DType, + *args, + scope: str = "local.var", + init: PrimExpr | int | float | None = None, + role_scoped: bool = False, +) -> Buffer: """Allocate a single-element variable buffer. Args: @@ -100,6 +123,8 @@ def alloc_var(dtype: DType, *args, scope: str = "local.var", init: PrimExpr | in init (PrimExpr, optional): The optional initializer value. When provided, the generated code will initialize the variable with this value instead of defaulting to zero. + role_scoped (bool, optional): Emit the allocation at the current + statement position instead of as an SBlock allocation. Examples: a = T.alloc_var('int32', 1) # var with init 1 a = T.alloc_var('int32', 'local.var') # var with local.var scope @@ -136,7 +161,10 @@ def alloc_var(dtype: DType, *args, scope: str = "local.var", init: PrimExpr | in if dtype is _ptr_sentinel: dtype = _dtypes.int64 - buffer = T.sblock_alloc_buffer([1], dtype, scope=parsed_scope) + if role_scoped: + buffer = T.alloc_buffer([1], dtype, scope=parsed_scope, annotations={_ROLE_SCOPED_ALLOC_ATTR: True}) + else: + buffer = T.sblock_alloc_buffer([1], dtype, scope=parsed_scope) if parsed_init is not None: # Always use T.buffer_store for reliable initialisation across all # backends. The sblock_attr("tl.local_var_init") path feeds into the diff --git a/tilelang/language/annotations.py b/tilelang/language/annotations.py index 0a9ba83011..7aa67624bc 100644 --- a/tilelang/language/annotations.py +++ b/tilelang/language/annotations.py @@ -11,6 +11,7 @@ __all__ = [ "use_swizzle", "annotate_layout", + "annotate_layout_view", "annotate_safe_value", "annotate_l2_hit_ratio", "annotate_restrict_buffers", @@ -42,6 +43,22 @@ def annotate_layout(layout_map: dict): return sblock_attr({"layout_map": _layout_map}) +def annotate_layout_view(layout_map: dict): + """Annotate layouts for exact Buffer views without data-var broadcasting.""" + _layout_map = {} + for buffer, layout in layout_map.items(): + if is_fragment(buffer): + assert isinstance(layout, Fragment), f"for Fragment {buffer}, layout must be a Fragment, but got {type(layout)}" + if isinstance(layout, Layout): + _layout_map[buffer] = layout + elif isinstance(layout, Callable): + _layout_map[buffer] = Layout(buffer.shape, layout) + else: + raise ValueError(f"Invalid layout: {layout}") + + return sblock_attr({"layout_view_map": _layout_map}) + + def annotate_safe_value(safe_value_map: dict): """Annotate the safe value of the buffer.""" _safe_value_map = {} diff --git a/tilelang/language/builtin.py b/tilelang/language/builtin.py index fee0ecb273..25ccdd9988 100644 --- a/tilelang/language/builtin.py +++ b/tilelang/language/builtin.py @@ -1530,6 +1530,21 @@ def ldg32(src: BufferLikeType, pred: PrimExpr = None) -> PrimExpr: return tirx.call_intrin("uint32", tirx.op.Op.get("tl.ldg32"), ptr, pred) +def lds32(src: BufferLikeType, pred: PrimExpr = None) -> PrimExpr: + """Load one aligned 32-bit word from shared memory. + + Usage: `T.lds32(smem_u8[i])` emits `tl::load_shared_32(ptr)`. + `T.lds32(smem_u8[i], pred=i < N)` emits a predicated shared load. + """ + if not isinstance(src, BufferLikeTypeTuple): + raise TypeError(f"T.lds32 expects Buffer, BufferRegion, or BufferLoad. Got {type(src)}: {src}") + ptr = retrieve_ptr(src, access_type="r") + if pred is None: + return tirx.call_intrin("uint32", tirx.op.Op.get("tl.lds32"), ptr) + else: + return tirx.call_intrin("uint32", tirx.op.Op.get("tl.lds32"), ptr, pred) + + def ldg64(src: BufferLikeType, pred: PrimExpr = None) -> PrimExpr: """Load 64 bits (8 bytes) from global memory using explicit PTX instructions. diff --git a/tilelang/language/copy_op.py b/tilelang/language/copy_op.py index 6258a13a25..37528e63ee 100644 --- a/tilelang/language/copy_op.py +++ b/tilelang/language/copy_op.py @@ -3,14 +3,14 @@ from __future__ import annotations from typing import Literal, Any -from tilelang._typing import BufferLikeType +from tilelang._typing import BufferLikeType, DType from tilelang.utils.language import ( to_buffer_region, legalize_pairwise_extents, ) from tilelang.utils.deprecated import deprecated +from tilelang.language.dtypes import get_tvm_dtype from tilelang.language.utils import get_extent, buffer_region_to_tile_region -import tvm from tvm import ir, tirx @@ -71,10 +71,12 @@ def copy( disable_tma (bool, keyword-only): Whether to disable TMA acceleration. Defaults to False. eviction_policy (Optional[str], keyword-only): Cache eviction policy. Defaults to None. prefer_instruction (Optional[str], keyword-only): Backend-specific preferred lowering - instruction category. For CUDA, recognized values include "tma", "cp_async", and - "sync". For "tma", T.copy keeps synchronous copy semantics; global -> shared copies - lower through TMA with an automatically allocated barrier and wait when constraints - are satisfied. + instruction category. For CUDA, recognized values include "tma", "cp_async", + "ldsm", "stsm", and "sync". For "tma", T.copy keeps synchronous copy semantics; + global -> shared copies lower through TMA with an automatically allocated barrier + and wait when constraints are satisfied. "ldsm" and "stsm" are strict requests: + lowering fails instead of silently using normal SIMT copy when ldmatrix/stmatrix + constraints are not met. annotations (Optional[dict], keyword-only): Additional annotations dict. If provided, coalesced_width, disable_tma, eviction_policy, and prefer_instruction can also be specified here. @@ -237,6 +239,7 @@ def tma_copy( *, barrier=None, leader_scope_threads: int | None = None, + tma_element_dtype: DType | None = None, eviction_policy: Literal["evict_normal", "evict_first", "evict_last"] | None = None, annotations: dict | None = None, ) -> tirx.PrimExpr | tirx.Stmt: @@ -255,6 +258,9 @@ def tma_copy( FP4 unpacked shared-memory storage is load-only for TMA: packed global ``float4_e2m1fn`` may be loaded into unpacked shared ``float4_e2m1_unpacked``, but the reverse TMA store is not supported. + Packed global ``float4_e2m1fn`` may also be loaded into ``uint8`` shared + byte storage; the TensorMap descriptor remains packed FP4 and the uint8 + shared buffer is only the physical byte-addressable storage. Args: src: Source memory region (global or shared) @@ -264,7 +270,11 @@ def tma_copy( The TMA load will arrive at this barrier with expected byte count. The user must wait on the same barrier via T.mbarrier_wait_parity(). leader_scope_threads: Number of threads in each TMA leader-election scope - (e.g., 32 for per-warp). Defaults to the thread extend in the current context if not specified. + (e.g., 32 for per-warp). Defaults to the full block size if not specified. + tma_element_dtype: Optional descriptor element type for TMA. When set, + descriptor data type, descriptor strides, and transaction byte + counts use this dtype while the source/destination buffers keep + their normal TileLang dtypes. eviction_policy: Cache eviction policy. Defaults to None. annotations: Additional annotations dict. Values in annotations take precedence over individual arguments. @@ -303,6 +313,11 @@ def tma_copy( if "leader_scope_threads" not in ann: ann["leader_scope_threads"] = leader_scope_threads + if "tma_element_dtype" not in ann and tma_element_dtype is not None: + ann["tma_element_dtype"] = tirx.StringImm(str(get_tvm_dtype(tma_element_dtype))) + if isinstance(ann.get("tma_element_dtype"), str): + ann["tma_element_dtype"] = tirx.StringImm(ann["tma_element_dtype"]) + if "eviction_policy" not in ann and eviction_policy is not None: eviction_policy_map = {"evict_normal": 0, "evict_first": 1, "evict_last": 2} ann["eviction_policy"] = eviction_policy_map[eviction_policy] @@ -322,10 +337,17 @@ def tma_copy( "float32", "float64", "bfloat16", + "float4_e2m1fn", + "float4_e2m1_unpacked", + "custom[float4_e2m1_unpacked]8", } ) +def _is_fp4_unpack_tma_load(src_dtype: DType, dst_dtype: DType) -> bool: + return get_tvm_dtype(src_dtype).is_float4_e2m1fn() and get_tvm_dtype(dst_dtype).is_float4_e2m1_unpacked() + + def tma_gather4( src: tirx.Buffer, dst: tirx.Buffer, @@ -361,7 +383,7 @@ def tma_gather4( raise ValueError(f"tma_gather4 expects rank-2 global buffer, got {len(src.shape)}") if len(dst.shape) != 2: raise ValueError(f"tma_gather4 expects rank-2 shared buffer (4 x K_box), got {len(dst.shape)}") - if src.dtype != dst.dtype: + if src.dtype != dst.dtype and not _is_fp4_unpack_tma_load(src.dtype, dst.dtype): raise ValueError(f"tma_gather4 dtype mismatch: src={src.dtype}, dst={dst.dtype}") if not (isinstance(dst.shape[0], int) and dst.shape[0] == 4) and not (hasattr(dst.shape[0], "value") and int(dst.shape[0].value) == 4): raise ValueError(f"tma_gather4 shared tile leading dim must be 4, got {dst.shape[0]}") @@ -414,9 +436,13 @@ def tma_gather4_bytes(K_box, dtype: str) -> int: """Transaction byte count for a 4-row gather4 of width ``K_box``. Pass to ``T.mbarrier_expect_tx`` immediately before ``T.tma_gather4``. """ - if dtype not in _TMA_SUPPORTED_DTYPES: + dtype_str = str(dtype) + if dtype_str not in _TMA_SUPPORTED_DTYPES: raise ValueError(f"Unsupported dtype: {dtype}") - dt = tvm.DataType(dtype) + if dtype_str == "float4_e2m1_unpacked": + dt = get_tvm_dtype("custom[float4_e2m1_unpacked]8") + else: + dt = get_tvm_dtype(dtype) if dt.is_float4_e2m1_unpacked(): elem_bits = 4 else: diff --git a/tilelang/language/customize.py b/tilelang/language/customize.py index cdb82027b6..a72e0b2fca 100644 --- a/tilelang/language/customize.py +++ b/tilelang/language/customize.py @@ -57,20 +57,30 @@ def reshape(src: Buffer, shape: ShapeType) -> Buffer: assert prim_expr_equal(bits_product(shape, src.dtype), bits_product(src.shape, src.dtype)), ( f"T.reshape/view shape check failed. src {src} src.shape: {src.shape}, src.dtype: {src.dtype}, target shape: {shape}, target dtype: {src.dtype}" ) - return T.Tensor(shape, src.dtype, src.data) + return T.Tensor(shape, src.dtype, data=src.data, scope=src.scope()) -def view(src: Buffer, shape: ShapeType | None = None, dtype: DType | None = None) -> Buffer: +def view( + src: Buffer, + shape: ShapeType | None = None, + dtype: DType | None = None, + strides: tuple[PrimExpr, ...] | None = None, + elem_offset: PrimExpr | None = None, +) -> Buffer: """Return a Tensor view of the input buffer with an optional new shape and dtype. - If `shape` is None the source buffer's shape is used; if `dtype` is None the source buffer's dtype is used. The returned buffer shares the same underlying data as `src` (no copy). + If `shape` is None the source buffer's shape is used; if `dtype` is None the source buffer's dtype is used. The returned buffer shares the same underlying data as `src` (no copy). Explicit `strides` and `elem_offset` may be supplied to describe a non-row-major logical view. """ if shape is None: shape = src.shape if dtype is None: dtype = src.dtype + if strides is not None and len(shape) != len(strides): + raise ValueError("Invalid shape/strides' dimensions") assert prim_expr_equal(bits_product(shape, dtype), bits_product(src.shape, src.dtype)), "T.reshape/view shape check failed." - return T.Tensor(shape, dtype, src.data) + if elem_offset is None: + elem_offset = src.elem_offset + return T.Tensor(shape, dtype, data=src.data, scope=src.scope(), strides=strides, elem_offset=elem_offset) def loop_break() -> PrimExpr: diff --git a/tilelang/language/proxy.py b/tilelang/language/proxy.py index b91adbd357..52718b88e3 100644 --- a/tilelang/language/proxy.py +++ b/tilelang/language/proxy.py @@ -162,10 +162,29 @@ def _construct_strides(shape: tuple[Any]): strides.append(s) return tuple(reversed(strides)) - def __call__(self, shape: ShapeType | PrimExpr | int, dtype: DType = "float32", data=None, scope=None) -> tirx.Buffer: + def __call__( + self, + shape: ShapeType | PrimExpr | int, + dtype: DType = "float32", + data=None, + scope=None, + strides: tuple[Any, ...] | None = None, + elem_offset=None, + ) -> tirx.Buffer: if isinstance(shape, (int, PrimExpr)): shape = (shape,) - return super().__call__(shape, dtype=dtype, strides=TensorProxy._construct_strides(shape), data=data, scope=scope) + if strides is None: + strides = TensorProxy._construct_strides(shape) + elif len(shape) != len(strides): + raise ValueError("Invalid shape/strides' dimensions") + return super().__call__( + shape, + dtype=dtype, + strides=strides, + data=data, + elem_offset=elem_offset, + scope=scope, + ) class StridedTensorProxy(BaseTensorProxy): diff --git a/tilelang/layout/__init__.py b/tilelang/layout/__init__.py index ae50e86cb4..8af17b3713 100644 --- a/tilelang/layout/__init__.py +++ b/tilelang/layout/__init__.py @@ -8,6 +8,7 @@ make_volta_swizzled_layout, # noqa: F401 make_wgmma_swizzled_layout, # noqa: F401 make_tcgen05mma_swizzled_layout, # noqa: F401 + make_sm120_fp4_smem_layout, # noqa: F401 make_full_bank_swizzled_layout, # noqa: F401 make_half_bank_swizzled_layout, # noqa: F401 make_quarter_bank_swizzled_layout, # noqa: F401 diff --git a/tilelang/layout/swizzle.py b/tilelang/layout/swizzle.py index 07ce091258..8ab8d140c9 100644 --- a/tilelang/layout/swizzle.py +++ b/tilelang/layout/swizzle.py @@ -90,6 +90,17 @@ def make_tcgen05mma_swizzled_layout(buffer: BufferLikeType, continuity: int = No return _ffi_api.make_tcgen05mma_swizzled_layout(buf, continuity, k_major) +def make_sm120_fp4_smem_layout(buffer: BufferLikeType): + """CUTLASS-compatible SM120 FP4 K-major shared-memory layout. + + Mirrors CUTLASS ``sm120_rr_smem_selector`` for FP4 operands. With a + 128-element K dimension, packed FP4 selects 64B swizzle while byte-slot + unpacked FP4 selects 128B swizzle. + """ + buf, _, _ = _get_buffer_info(buffer) + return _ffi_api.make_sm120_fp4_smem_layout(buf) + + # swizzle 128B def make_full_bank_swizzled_layout(buffer: BufferLikeType): """ From 52b6aaca685a588ae17a9d6e971a7a68400815d5 Mon Sep 17 00:00:00 2001 From: sepcnt <30561671+sepcnt@users.noreply.github.com> Date: Fri, 5 Jun 2026 14:19:30 +0800 Subject: [PATCH 3/6] Update ReduceLowerer to improve batch handling and remove redundant checks --- src/backend/common/op/reduce.h | 133 ++++++++++++++++++--------- src/cuda/codegen/codegen_cuda.cc | 9 ++ src/cuda/codegen/codegen_cuda.h | 1 + src/transform/thread_storage_sync.cc | 13 --- 4 files changed, 99 insertions(+), 57 deletions(-) diff --git a/src/backend/common/op/reduce.h b/src/backend/common/op/reduce.h index 3389693511..ec7994ee67 100644 --- a/src/backend/common/op/reduce.h +++ b/src/backend/common/op/reduce.h @@ -24,6 +24,7 @@ #include #include +#include #include #include #include @@ -100,43 +101,44 @@ inline int GetPreferedVectorizedSize(DataType dt) { return 1; } -inline PrimExpr MakeInitValue(const ReduceOpNode &op, int vsize = 1) { - auto dst_dtype = op.dst->dtype; +inline PrimExpr MakeInitValue(const ReduceOpNode &op, DataType dtype, + int vsize = 1) { + auto dst_dtype = dtype; auto is_int = dst_dtype.is_int(); bool is_uint = dst_dtype.is_uint(); auto bits = dst_dtype.bits(); PrimExpr scalar; if (op.type->isSum() || op.type->isAbsSum()) { - scalar = make_zero(op.dst->dtype); + scalar = make_zero(dst_dtype); } else if (op.type->isMax()) { if (is_int) { - scalar = make_const(op.dst->dtype, SignedMin(bits)); + scalar = make_const(dst_dtype, SignedMin(bits)); } else if (is_uint) { - scalar = make_const(op.dst->dtype, 0); + scalar = make_const(dst_dtype, 0); } else { - scalar = make_const(op.dst->dtype, -INFINITY); + scalar = make_const(dst_dtype, -INFINITY); } } else if (op.type->isMin()) { if (is_int) { - scalar = make_const(op.dst->dtype, SignedMax(bits)); + scalar = make_const(dst_dtype, SignedMax(bits)); } else if (is_uint) { - scalar = make_const(op.dst->dtype, UnsignedMax(bits)); + scalar = make_const(dst_dtype, UnsignedMax(bits)); } else { - scalar = make_const(op.dst->dtype, INFINITY); + scalar = make_const(dst_dtype, INFINITY); } } else if (op.type->isAbsMax()) { - scalar = make_const(op.dst->dtype, 0); + scalar = make_const(dst_dtype, 0); } else if (op.type->isBitAnd()) { if (is_int) { - scalar = make_const(op.dst->dtype, -1); + scalar = make_const(dst_dtype, -1); } else if (is_uint) { - scalar = make_const(op.dst->dtype, UnsignedMax(bits)); + scalar = make_const(dst_dtype, UnsignedMax(bits)); } else { - scalar = make_const(op.dst->dtype, -INFINITY); + scalar = make_const(dst_dtype, -INFINITY); } } else if (op.type->isBitOr() || op.type->isBitXor()) { - scalar = make_zero(op.dst->dtype); + scalar = make_zero(dst_dtype); } else { LOG(FATAL) << "Unsupported reduce type: " << op.type->type; scalar = PrimExpr(); @@ -147,6 +149,10 @@ inline PrimExpr MakeInitValue(const ReduceOpNode &op, int vsize = 1) { return Broadcast(scalar, vsize); } +inline PrimExpr MakeInitValue(const ReduceOpNode &op, int vsize = 1) { + return MakeInitValue(op, op.dst->dtype, vsize); +} + inline PrimExpr MakeReduce(const ReduceOpNode &op, int vsize, const PrimExpr &acc, const PrimExpr &b) { if (vsize != 1 && vsize != 2) { @@ -247,6 +253,26 @@ inline std::optional MakeCodegenReducer(const ReduceOpNode &op, return std::nullopt; } +inline bool CanUsePackedReducer(const ReduceOpNode &op, DataType dtype, + int vsize) { + if (vsize <= 1) { + return false; + } + if (!MakeCodegenReducer(op, vsize).has_value()) { + return false; + } + if ((dtype.is_bfloat16() || dtype.is_float16()) && + (op.type->isSum() || op.type->isAbsSum())) { + return false; + } + return true; +} + +inline bool UseFloatAccumulator(const ReduceOpNode &op, DataType dtype) { + return (dtype.is_bfloat16() || dtype.is_float16()) && + (op.type->isSum() || op.type->isAbsSum()); +} + inline bool CanUsePackedRamp(const PrimExpr &index, const Var &var, int vsize, arith::Analyzer *analyzer) { ICHECK_GT(vsize, 1); @@ -368,6 +394,8 @@ template struct ReduceLowerer { auto clear_buffer = dst_buffer; auto need_duplicate = false; auto need_update = false; + bool use_float_accumulator = + reduce::UseFloatAccumulator(op, dst_buffer->dtype); if ((op.type->isSum() || op.type->isAbsSum()) && !op.clear) { need_duplicate = true; need_update = true; @@ -388,17 +416,38 @@ template struct ReduceLowerer { red_layout->ReplicateExtent())) { need_duplicate = true; } + if (use_float_accumulator) { + need_duplicate = true; + } ICHECK(!analyzer->CanProve(dst_layout->ReplicateExtent() > red_layout->ReplicateExtent())) << "Inconsistent layouts between src and dst in ReduceOp: " << "dst_layout=" << dst_layout << "red_layout=" << red_layout; if (need_duplicate) { - clear_buffer = decl_buffer(red_layout->OutputShape(), dst_buffer->dtype, + DataType clear_dtype = + use_float_accumulator ? DataType::Float(32) : dst_buffer->dtype; + clear_buffer = decl_buffer(red_layout->OutputShape(), clear_dtype, dst_buffer->name + "_clear", GetPtrStorageScope(dst_buffer->data)); } + auto make_dst_update = [&](const Array &dst_idx, + const Array &red_idx) -> PrimExpr { + PrimExpr value = BufferLoad(clear_buffer, red_idx); + if (need_update) { + PrimExpr dst_value = BufferLoad(dst_buffer, dst_idx); + if (dst_value->dtype != value->dtype) { + dst_value = Cast(value->dtype, dst_value); + } + value = reduce::MakeUpdate(op, dst_value, value); + } + if (value->dtype != dst_buffer->dtype) { + value = Cast(dst_buffer->dtype, value); + } + return value; + }; + Array src_indice_compressed; Array src_var_compressed; for (size_t i = 0; i < src_layout->OutputDim(); ++i) { @@ -419,7 +468,7 @@ template struct ReduceLowerer { if (vsize > 1 && !src_var_compressed.empty()) { auto *ext = src_var_compressed.back()->dom->extent.as(); if (ext && ext->value >= vsize && ext->value % vsize == 0 && - reduce::MakeCodegenReducer(op, vsize).has_value() && + reduce::CanUsePackedReducer(op, clear_buffer->dtype, vsize) && reduce::CanUsePackedRamp(src_indice_compressed.back(), src_var_compressed.back()->var, vsize, analyzer)) { @@ -436,9 +485,10 @@ template struct ReduceLowerer { if (require_init || (need_duplicate && (op.type->isMax() || op.type->isMin() || op.type->isAbsMax()))) { - local_body.push_back(BufferStore(clear_buffer_packed, - reduce::MakeInitValue(op, vsize), - red_indices)); + local_body.push_back(BufferStore( + clear_buffer_packed, + reduce::MakeInitValue(op, clear_buffer->dtype, vsize), + red_indices)); } const auto *ext_int = @@ -497,8 +547,9 @@ template struct ReduceLowerer { if (require_init || (need_duplicate && (op.type->isMax() || op.type->isMin() || op.type->isAbsMax()))) { - stmts.push_back(BufferStore(clear_buffer, reduce::MakeInitValue(op), - red_indices)); + stmts.push_back(BufferStore( + clear_buffer, reduce::MakeInitValue(op, clear_buffer->dtype), + red_indices)); } Stmt reduce_local = BufferStore( @@ -522,20 +573,22 @@ template struct ReduceLowerer { auto iter_sum = arith::NormalizeToIterSum(src_thread, ToVMap(src_vars), analyzer); - const int batch = op.batch; + int batch = op.batch; + int64_t physical_output_elems = 1; + for (const auto &s : clear_buffer->shape) { + const int64_t *p = as_const_int(s); + ICHECK(p != nullptr) << "ReduceOp: batch > 1 requires compile-time " + "constant output shape"; + physical_output_elems *= *p; + } if (batch > 1) { - int64_t N_total = 1; - for (const auto &s : clear_buffer->shape) { - const int64_t *p = as_const_int(s); - ICHECK(p != nullptr) << "ReduceOp: batch > 1 requires compile-time " - "constant output shape"; - N_total *= *p; - } - ICHECK_LE(batch, N_total) + ICHECK_GT(physical_output_elems, 0); + batch = + static_cast(std::min(batch, physical_output_elems)); + ICHECK_EQ(physical_output_elems % batch, 0) << "ReduceOp: batch=" << batch - << " exceeds per-thread output element count N=" << N_total; - ICHECK_EQ(N_total % batch, 0) << "ReduceOp: batch=" << batch - << " must evenly divide N=" << N_total; + << " must evenly divide per-thread output element count N=" + << physical_output_elems; } bool use_batch = batch > 1; @@ -595,7 +648,7 @@ template struct ReduceLowerer { Impl::GetPreferedVectorizedSize(clear_buffer->dtype, T.target); bool can_batch_pack = vsize > 1 && batch >= vsize && batch % vsize == 0 && - reduce::MakeCodegenReducer(op, vsize).has_value(); + reduce::CanUsePackedReducer(op, clear_buffer->dtype, vsize); int eff_batch = can_batch_pack ? (batch / vsize) : batch; std::string reducer = reduce::MakeCodegenReducer(op, can_batch_pack ? vsize : 1) @@ -740,11 +793,7 @@ template struct ReduceLowerer { predicate = analyzer->Simplify(predicate); } - PrimExpr update = - need_update - ? reduce::MakeUpdate(op, BufferLoad(dst_buffer, post_dst_idx), - BufferLoad(clear_buffer, post_red_idx)) - : BufferLoad(clear_buffer, post_red_idx); + PrimExpr update = make_dst_update(post_dst_idx, post_red_idx); auto store = BufferStore(dst_buffer, update, post_dst_idx); Stmt post_body; if (analyzer->CanProve(predicate)) { @@ -813,11 +862,7 @@ template struct ReduceLowerer { predicate = analyzer->Simplify(predicate); } if (need_duplicate) { - PrimExpr update = - need_update - ? reduce::MakeUpdate(op, BufferLoad(dst_buffer, dst_indices), - BufferLoad(clear_buffer, red_indices)) - : BufferLoad(clear_buffer, red_indices); + PrimExpr update = make_dst_update(dst_indices, red_indices); auto store = BufferStore(dst_buffer, update, dst_indices); if (analyzer->CanProve(predicate)) { stmts.push_back(store); diff --git a/src/cuda/codegen/codegen_cuda.cc b/src/cuda/codegen/codegen_cuda.cc index 9c4e1d998a..1ad3f3866a 100644 --- a/src/cuda/codegen/codegen_cuda.cc +++ b/src/cuda/codegen/codegen_cuda.cc @@ -1420,6 +1420,15 @@ void CodeGenTileLangCUDA::PrintVecElemStore(const std::string &vec, DataType t, } } +void CodeGenTileLangCUDA::PrintVecConstructor(DataType t, std::ostream &os) { + int lanes = t.lanes(); + if (t.is_float() && t.bits() == 32 && lanes >= 2 && lanes <= 4) { + os << "make_float" << lanes; + return; + } + CodeGenC::PrintVecConstructor(t, os); +} + void CodeGenTileLangCUDA::PrintStorageSync(const CallNode *op) { auto args = op->args; const std::string &sync = args[0].as()->value; diff --git a/src/cuda/codegen/codegen_cuda.h b/src/cuda/codegen/codegen_cuda.h index f6ace98564..f14a2bc5a3 100644 --- a/src/cuda/codegen/codegen_cuda.h +++ b/src/cuda/codegen/codegen_cuda.h @@ -39,6 +39,7 @@ class CodeGenTileLangCUDA final : public CodeGenC { std::ostream &os) final; // NOLINT(*) void PrintVecElemStore(const std::string &vec, DataType t, int i, const std::string &value) final; + void PrintVecConstructor(DataType t, std::ostream &os) final; std::string GetVecLoad(DataType t, const BufferNode *buffer, PrimExpr base) final; void PrintVecStore(const BufferNode *buffer, DataType t, PrimExpr base, diff --git a/src/transform/thread_storage_sync.cc b/src/transform/thread_storage_sync.cc index 372e7affe9..8c751e1266 100644 --- a/src/transform/thread_storage_sync.cc +++ b/src/transform/thread_storage_sync.cc @@ -1592,19 +1592,6 @@ struct TileLangThreadSyncPlanner : public ConstrVisitor { ranges.push_back(byte_range(min, extent, access.dtype)); } } - if (access.buffer_name.defined() && access.buffer_name->data.defined()) { - PrimExpr elems = make_const(DataType::Int(64), 1); - for (const PrimExpr &dim : access.buffer_name->shape) { - elems = elems * to_i64(dim); - } - Var data = Downcast(access.buffer_name->data); - PrimExpr min = AliasByteOffset(data, DataType::Int(64)); - PrimExpr max = min + - elems * make_const(DataType::Int(64), - access.buffer_name->dtype.bytes()) - - make_const(DataType::Int(64), 1); - ranges.push_back({min, max}); - } return ranges; }; auto lhs_ranges = collect_ranges(lhs); From 59dc6cbcb332b9f1b23c284722f5d029b0fe8a62 Mon Sep 17 00:00:00 2001 From: sepcnt <30561671+sepcnt@users.noreply.github.com> Date: Fri, 5 Jun 2026 14:20:41 +0800 Subject: [PATCH 4/6] Add handle byte offset support in TMA descriptor parsing for nvrtc --- tilelang/jit/adapter/nvrtc/wrapper.py | 38 ++++++- tilelang/jit/adapter/utils.py | 140 ++++++++++++++++++-------- tilelang/jit/adapter/wrapper.py | 21 +++- 3 files changed, 152 insertions(+), 47 deletions(-) diff --git a/tilelang/jit/adapter/nvrtc/wrapper.py b/tilelang/jit/adapter/nvrtc/wrapper.py index 2863fc2bd2..3dd6de7bbc 100644 --- a/tilelang/jit/adapter/nvrtc/wrapper.py +++ b/tilelang/jit/adapter/nvrtc/wrapper.py @@ -22,7 +22,13 @@ from tilelang import tvm as tvm from tilelang.jit.adapter.wrapper import TLCUDASourceWrapper -from tilelang.jit.adapter.utils import match_declare_kernel, pythonic_expr, parse_function_call_args, parse_tma_descriptor_args +from tilelang.jit.adapter.utils import ( + is_handle_add_byte_offset_call, + match_declare_kernel, + pythonic_expr, + parse_function_call_args, + parse_tma_descriptor_args, +) PREDEF_HOST_FUNC_PY = """ from cuda.bindings.driver import ( @@ -49,6 +55,9 @@ _function_names = {} +def _tl_device_ptr(tensor, byte_offset=0): + return tensor.data_ptr() + byte_offset + def call({}): {} """ @@ -56,7 +65,7 @@ def call({}): TMA_DESC_INIT_FUNC_PY = """ {0}_type = CUtensorMapDataType({1}) {0}_tensorRank = {2} - {0}_globalAddress = {3}.data_ptr() + {0}_globalAddress = {3} {0}_globalDim = [{4}] {0}_globalStride = [{5}][1:] {0}_boxDim = [{6}] @@ -87,7 +96,7 @@ def call({}): TMA_IM2COL_DESC_INIT_FUNC_PY = """ {0}_type = CUtensorMapDataType({1}) {0}_tensorRank = {2} - {0}_globalAddress = {3}.data_ptr() + {0}_globalAddress = {3} {0}_globalDim = [{4}] {0}_globalStride = [{5}][1:] {0}_elementStrides = [{6}] @@ -288,6 +297,21 @@ def _pythonic_expr(self, expr: tvm.tirx.PrimExpr) -> str: """ return pythonic_expr(expr, self._TYPE_MAP, ignore_cast=True, floor_div_op="//") + def _pythonic_handle_expr(self, expr: tvm.tirx.PrimExpr) -> str: + """Convert a TIR handle expression to a Python device pointer expression.""" + if isinstance(expr, tvm.tirx.Var): + return f"_tl_device_ptr({expr.name})" + if isinstance(expr, tvm.tirx.Cast): + return self._pythonic_handle_expr(expr.value) + if is_handle_add_byte_offset_call(expr): + if len(expr.args) != 2: + raise ValueError(f"handle_add_byte_offset expects 2 arguments, got {len(expr.args)}") + base, byte_offset = expr.args + if isinstance(base, tvm.tirx.Var): + return f"_tl_device_ptr({base.name}, {self._pythonic_expr(byte_offset)})" + return f"{self._pythonic_handle_expr(base)} + {self._pythonic_expr(byte_offset)}" + raise ValueError(f"Unsupported TMA global address expression: {expr}") + def create_dispatch_func(self, code, function_informations): """Generate Python dispatch function that launches multiple CUDA kernels. @@ -501,7 +525,13 @@ def generate_tma_descriptor_args(self, desc_name_map: dict[str, str], desc_name_ return tma_descriptor_init # Parse TMA descriptor arguments using the common utility - parsed_params = parse_tma_descriptor_args(self.tma_descriptor_args, desc_name_map, desc_name_var_map, self._pythonic_expr) + parsed_params = parse_tma_descriptor_args( + self.tma_descriptor_args, + desc_name_map, + desc_name_var_map, + self._pythonic_expr, + self._pythonic_handle_expr, + ) # Generate Python code from parsed parameters for params in parsed_params: diff --git a/tilelang/jit/adapter/utils.py b/tilelang/jit/adapter/utils.py index df8dff461b..1a323c97be 100644 --- a/tilelang/jit/adapter/utils.py +++ b/tilelang/jit/adapter/utils.py @@ -293,36 +293,80 @@ def _visitor(node): return next(iter(node_to_result_map[expr]), "") -def maybe_desc_name(name: str, matches: list[str], i: int, desc_name_map: dict[str, str] | None = None) -> bool: - """ - Check if a parameter name corresponds to a TMA descriptor. - - Args: - name: The parameter name to check. - matches: List of all matched parameter names. - i: Index of the current match. - desc_name_map: Optional mapping to store descriptor name relationships. +HANDLE_ADD_BYTE_OFFSET_OP = tvm.ir.Op.get("tirx.handle_add_byte_offset") + + +def split_function_parameters(declaration: str) -> list[str]: + """Return top-level C/CUDA parameter declarations from the function parameter list.""" + end = declaration.rfind(")") + start = -1 + depth = 0 + for index in range(end, -1, -1): + char = declaration[index] + if char == ")": + depth += 1 + elif char == "(": + depth -= 1 + if depth == 0: + start = index + break + if start == -1 or end <= start: + return [] - Returns: - True if the parameter is a TMA descriptor. - """ - match = matches[i] - if not (match == name + "_desc" or match.startswith(name + "_desc_")): - return False - desc_decls = [] - if desc_name_map is not None: - desc_name_map[match] = name - if i > 0: - desc_decls.append(matches[i - 1]) - if i < len(matches) - 1: - desc_decls.append(matches[i + 1]) - return any([decl == "CUtensorMap" for decl in desc_decls]) + params = [] + depth = 0 + current = [] + for char in declaration[start + 1 : end]: + if char == "," and depth == 0: + param = "".join(current).strip() + if param: + params.append(param) + current = [] + continue + current.append(char) + if char in "([{<": + depth += 1 + elif char in ")]}>": + depth = max(depth - 1, 0) + param = "".join(current).strip() + if param: + params.append(param) + return params + + +def parse_c_parameter(param: str) -> tuple[str, str] | None: + """Parse one C/CUDA parameter declaration into (name, type).""" + param = param.strip() + if not param or param == "void": + return None + match = re.search(r"([A-Za-z_]\w*)\s*(?:\[[^\]]*\])?\s*$", param) + if match is None: + return None + name = match.group(1) + param_type = param[: match.start(1)].strip() + return name, param_type + + +def is_tma_descriptor_type(param_type: str) -> bool: + return re.search(r"\bCUtensorMap\b", param_type) is not None + + +def resolve_descriptor_source_name(param_name: str, function_args: list[dict[str, str]]) -> str: + for arg in function_args: + name = arg["name"] + if param_name == name + "_desc" or param_name.startswith(name + "_desc_"): + return name + return param_name + + +def is_handle_add_byte_offset_call(expr: Any) -> bool: + return isinstance(expr, tvm.tirx.Call) and expr.op.same_as(HANDLE_ADD_BYTE_OFFSET_OP) def parse_function_call_args( declaration: str, function_args: list[dict[str, str]], - function_params: list[Any], + function_params: list[Any] | None, desc_name_map: dict[str, str] | None = None, desc_name_var_map: dict[str, tvm.tirx.Var] | None = None, transform_arg: Callable[[str, str], Any] | None = None, @@ -341,25 +385,32 @@ def parse_function_call_args( Returns: List of parsed call arguments. """ - pattern = r"[,\s]*(?:\w+\s*\*+\s*__restrict__\s+)?(\w+)" - matches = re.findall(pattern, declaration) + params = [] + for param in split_function_parameters(declaration): + parsed = parse_c_parameter(param) + if parsed is not None: + params.append(parsed) + + function_arg_map = {arg["name"]: arg for arg in function_args} call_args = [] - for i, match in enumerate(matches): - for arg in function_args: - if arg["name"] == match: - if transform_arg is not None: - call_args.append(transform_arg(match, arg["type"])) - else: - call_args.append(match) - elif maybe_desc_name(arg["name"], matches, i, desc_name_map): - if transform_arg is not None: - call_args.append(transform_arg(match, "None")) - else: - call_args.append(match) - if desc_name_var_map is not None and function_params is not None: - assert len(call_args) <= len(function_params), f"Too many arguments: {len(call_args)} > {len(function_params)}" - desc_name_var_map[match] = function_params[len(call_args) - 1] + for param_name, param_type in params: + if param_name in function_arg_map: + arg = function_arg_map[param_name] + if transform_arg is not None: + call_args.append(transform_arg(param_name, arg["type"])) + else: + call_args.append(param_name) + elif is_tma_descriptor_type(param_type): + if transform_arg is not None: + call_args.append(transform_arg(param_name, "None")) + else: + call_args.append(param_name) + if desc_name_map is not None: + desc_name_map[param_name] = resolve_descriptor_source_name(param_name, function_args) + if desc_name_var_map is not None and function_params is not None: + assert len(call_args) <= len(function_params), f"Too many arguments: {len(call_args)} > {len(function_params)}" + desc_name_var_map[param_name] = function_params[len(call_args) - 1] return call_args @@ -398,6 +449,7 @@ def parse_tma_descriptor_args( desc_name_map: dict[str, str], desc_name_var_map: dict[str, tvm.tirx.Var], pythonic_expr_func: Callable[[Any], str], + global_address_expr_func: Callable[[Any], str] | None = None, ) -> list[TMADescriptorParams]: """ Parse TMA descriptor arguments into structured parameters. @@ -407,6 +459,7 @@ def parse_tma_descriptor_args( desc_name_map: Mapping from descriptor handles to parameter names. desc_name_var_map: Mapping from descriptor handles to TVM variables. pythonic_expr_func: Function to convert TVM expressions to strings. + global_address_expr_func: Optional function for handle expressions. Returns: List of parsed TMA descriptor parameters. @@ -439,7 +492,10 @@ def parse_tma_descriptor_args( if not isinstance(tensor_rank, int) or tensor_rank <= 0: raise ValueError(f"Invalid tensor_rank: {tensor_rank}. Must be a positive integer") - global_address = pythonic_expr_func(global_address) + if global_address_expr_func is None: + global_address = pythonic_expr_func(global_address) + else: + global_address = global_address_expr_func(global_address) params = TMADescriptorParams(handle_name, dtype, tensor_rank, global_address, is_img2col) if not is_img2col: diff --git a/tilelang/jit/adapter/wrapper.py b/tilelang/jit/adapter/wrapper.py index 29dd61ddd9..982a58f2f3 100644 --- a/tilelang/jit/adapter/wrapper.py +++ b/tilelang/jit/adapter/wrapper.py @@ -14,6 +14,7 @@ is_hip_target, is_cpu_target, get_annotated_mod, + is_handle_add_byte_offset_call, pythonic_expr, parse_function_call_args, parse_tma_descriptor_args, @@ -265,6 +266,18 @@ def _pythonic_expr(self, expr: tvm.tirx.PrimExpr) -> str: # and '//' is not a valid operator in C/C++. return pythonic_expr(expr, self._TYPE_MAP, floor_div_op="/") + def _cxx_handle_expr(self, expr: tvm.tirx.PrimExpr) -> str: + if isinstance(expr, tvm.tirx.Var): + return expr.name + if isinstance(expr, tvm.tirx.Cast): + return self._cxx_handle_expr(expr.value) + if is_handle_add_byte_offset_call(expr): + if len(expr.args) != 2: + raise ValueError(f"handle_add_byte_offset expects 2 arguments, got {len(expr.args)}") + base, byte_offset = expr.args + return f"(reinterpret_cast({self._cxx_handle_expr(base)}) + {self._pythonic_expr(byte_offset)})" + raise ValueError(f"Unsupported TMA global address expression: {expr}") + def _lookup_type(self, dtype: str | Any) -> str: key = dtype if isinstance(dtype, str) else str(dtype) result = self._TYPE_MAP.get(key) @@ -406,7 +419,13 @@ def generate_tma_descriptor_args(self, desc_name_map: dict[str, str], desc_name_ return tma_descriptor_init # Parse TMA descriptor arguments using the common utility - parsed_params = parse_tma_descriptor_args(self.tma_descriptor_args, desc_name_map, desc_name_var_map, self._pythonic_expr) + parsed_params = parse_tma_descriptor_args( + self.tma_descriptor_args, + desc_name_map, + desc_name_var_map, + self._pythonic_expr, + self._cxx_handle_expr, + ) # Generate C++ code from parsed parameters for params in parsed_params: From 516c6640f635dacf1782a6c3d7b5f16595e9b453 Mon Sep 17 00:00:00 2001 From: sepcnt <30561671+sepcnt@users.noreply.github.com> Date: Thu, 11 Jun 2026 21:58:27 +0800 Subject: [PATCH 5/6] Refactor ptx_mma_blockscaled for compatibility with #2324 --- examples/sage_attention_sm120/tir_helpers.py | 16 +- src/cuda/codegen/codegen_cuda.cc | 98 ++++++---- src/cuda/codegen/codegen_cuda.h | 2 +- src/op/builtin.cc | 4 +- src/op/builtin.h | 17 +- .../{mma_blockscale.h => mma_blockscaled.h} | 177 +++++++++--------- tilelang/language/ast/ir.py | 4 +- tilelang/language/tir/ir.py | 2 +- tilelang/language/tir/op.py | 41 ++-- 9 files changed, 205 insertions(+), 156 deletions(-) rename src/tl_templates/cuda/instruction/{mma_blockscale.h => mma_blockscaled.h} (62%) diff --git a/examples/sage_attention_sm120/tir_helpers.py b/examples/sage_attention_sm120/tir_helpers.py index 94b6a96a62..dc28582589 100644 --- a/examples/sage_attention_sm120/tir_helpers.py +++ b/examples/sage_attention_sm120/tir_helpers.py @@ -33,13 +33,15 @@ def mma_m16n32k64_blockscale_f32( """Emit four m16n8k64 FP4 MMAs in contiguous n8-atom register order.""" scale_a_reg = T.alloc_var("uint32", init=scale_a, role_scoped=True) scale_b_reg = T.alloc_var("uint32", init=scale_b, role_scoped=True) - T.ptx_mma_blockscale( + T.ptx_mma_blockscaled( "float32", "m16n8k64", "row", "col", "e2m1", "e2m1", + "float32", + "float8_e4m3", 4, a_regs, a_offset, @@ -52,13 +54,15 @@ def mma_m16n32k64_blockscale_f32( scale_id_a, scale_id_b_base, ) - T.ptx_mma_blockscale( + T.ptx_mma_blockscaled( "float32", "m16n8k64", "row", "col", "e2m1", "e2m1", + "float32", + "float8_e4m3", 4, a_regs, a_offset, @@ -71,13 +75,15 @@ def mma_m16n32k64_blockscale_f32( scale_id_a, scale_id_b_base + 1, ) - T.ptx_mma_blockscale( + T.ptx_mma_blockscaled( "float32", "m16n8k64", "row", "col", "e2m1", "e2m1", + "float32", + "float8_e4m3", 4, a_regs, a_offset, @@ -90,13 +96,15 @@ def mma_m16n32k64_blockscale_f32( scale_id_a, scale_id_b_base + 2, ) - T.ptx_mma_blockscale( + T.ptx_mma_blockscaled( "float32", "m16n8k64", "row", "col", "e2m1", "e2m1", + "float32", + "float8_e4m3", 4, a_regs, a_offset, diff --git a/src/cuda/codegen/codegen_cuda.cc b/src/cuda/codegen/codegen_cuda.cc index 1ad3f3866a..002ec9380f 100644 --- a/src/cuda/codegen/codegen_cuda.cc +++ b/src/cuda/codegen/codegen_cuda.cc @@ -12,6 +12,8 @@ #include #include +#include +#include #include #include #include @@ -601,9 +603,9 @@ std::string CodeGenTileLangCUDA::Finish() { if (need_mma_instruction_h_) { decl_stream << "#include \n"; } - if (need_mma_blockscale_instruction_h_) { + if (need_mma_blockscaled_instruction_h_) { decl_stream - << "#include \n"; + << "#include \n"; } if (need_wgmma_instruction_h_) { decl_stream << "#include \n"; @@ -2792,34 +2794,38 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) { replacer.register_rule("(C_ptr)", c_ref); replacer.register_rule("(C_offset)", c_bias); this->stream << replacer.rewrite(mma_call); - } else if (op->op.same_as(tl::ptx_mma_blockscale())) { - // arg 0: shape: mXnXkX - // arg 1: A layout: row/col - // arg 2: B layout: row/col - // arg 3: A precision: e2m1, ... - // arg 4: B precision: e2m1, ... - // arg 5: C precision: fp32, ... - // arg 6: scale vector width - // arg 7: A multiplicand - // arg 8: A multiplicand index - // arg 9: B multiplicand - // arg 10: B multiplicand index - // arg 11: C accumulator - // arg 12: C accumulator index - // arg 13: packed A scale vector - // arg 14: packed B scale vector - // arg 15: A scale selector - // arg 16: B scale selector - ICHECK_EQ(op->args.size(), 17U) - << "ptx_mma_blockscale expects 17 arguments"; - std::string shape = Downcast(op->args[0])->value; - std::string A_layout = Downcast(op->args[1])->value; - std::string B_layout = Downcast(op->args[2])->value; - std::string A_dtype = Downcast(op->args[3])->value; - std::string B_dtype = Downcast(op->args[4])->value; - std::string C_dtype = Downcast(op->args[5])->value; - int scale_vec = Downcast(op->args[6])->value; - constexpr int arg_base = 7; + } else if (op->op.same_as(tl::ptx_mma_blockscaled())) { + // arg 0: dtype/result precision + // arg 1: shape: mXnXkX + // arg 2: A layout: row/col + // arg 3: B layout: row/col + // arg 4: A precision: e2m1, ... + // arg 5: B precision: e2m1, ... + // arg 6: C precision: fp32, ... + // arg 7: scale factor precision: ue8m0/ue4m3 or aliases + // arg 8: scale vector width + // arg 9: A multiplicand + // arg 10: A multiplicand index + // arg 11: B multiplicand + // arg 12: B multiplicand index + // arg 13: C accumulator + // arg 14: C accumulator index + // arg 15: packed A scale vector + // arg 16: packed B scale vector + // arg 17: optional A scale selector + // arg 18: optional B scale selector + ICHECK(op->args.size() == 17U || op->args.size() == 19U) + << "ptx_mma_blockscaled expects 17 or 19 arguments"; + std::string dtype = Downcast(op->args[0])->value; + std::string shape = Downcast(op->args[1])->value; + std::string A_layout = Downcast(op->args[2])->value; + std::string B_layout = Downcast(op->args[3])->value; + std::string A_dtype = Downcast(op->args[4])->value; + std::string B_dtype = Downcast(op->args[5])->value; + std::string C_dtype = Downcast(op->args[6])->value; + std::string SF_dtype = Downcast(op->args[7])->value; + int scale_vec = Downcast(op->args[8])->value; + constexpr int arg_base = 9; std::string a_ref = this->PrintExpr(op->args[arg_base + 0]); PrimExpr a_offset_expr = op->args[arg_base + 1]; std::string b_ref = this->PrintExpr(op->args[arg_base + 2]); @@ -2828,11 +2834,16 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) { std::string c_offset = this->PrintExpr(op->args[arg_base + 5]); std::string sfa = this->PrintExpr(op->args[arg_base + 6]); std::string sfb = this->PrintExpr(op->args[arg_base + 7]); - std::string id_a = this->PrintExpr(op->args[arg_base + 8]); - std::string id_b = this->PrintExpr(op->args[arg_base + 9]); + std::string id_a = + op->args.size() == 19U ? this->PrintExpr(op->args[arg_base + 8]) : "0"; + std::string id_b = + op->args.size() == 19U ? this->PrintExpr(op->args[arg_base + 9]) : "0"; auto dtype_a_enum = tl::codegen::ptx::DTypeFromString(A_dtype); auto dtype_b_enum = tl::codegen::ptx::DTypeFromString(B_dtype); auto dtype_c_enum = tl::codegen::ptx::DTypeFromString(C_dtype); + auto dtype_enum = tl::codegen::ptx::DTypeFromString(dtype); + ICHECK(dtype_enum == dtype_c_enum) + << "ptx_mma_blockscaled dtype must match C_dtype"; // The Python intrinsic accepts logical fp4 element offsets. Pointer // arithmetic below is byte-addressed for packed e2m1 operands. if (dtype_a_enum == tl::codegen::ptx::DataType::kFloat4_e2m1fn) { @@ -2844,15 +2855,32 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) { std::string a_offset = this->PrintExpr(a_offset_expr); std::string b_offset = this->PrintExpr(b_offset_expr); auto [m, n, k] = tl::codegen::ptx::ParseMMAShape(shape); - tl::codegen::ptx::GetMMABlockScaleInfo( + auto blockscale_info = tl::codegen::ptx::GetMMABlockScaleInfo( dtype_a_enum, dtype_b_enum, dtype_c_enum, m, n, k, A_layout == "row" ? false : true, B_layout == "row" ? false : true, scale_vec); + auto normalize_sf_dtype = [](std::string value) { + std::transform( + value.begin(), value.end(), value.begin(), + [](unsigned char ch) { return static_cast(std::tolower(ch)); }); + if (value == "ue8m0" || value == "uint8" || value == "uint8_t" || + value == "u8") { + return std::string("ue8m0"); + } + if (value == "ue4m3" || value == "e4m3" || value == "float8_e4m3" || + value == "float8_e4m3fn") { + return std::string("ue4m3"); + } + return value; + }; + ICHECK_EQ(normalize_sf_dtype(SF_dtype), blockscale_info.scale_dtype) + << "ptx_mma_blockscaled SF_dtype does not match the selected " + "block-scaled MMA configuration"; - need_mma_blockscale_instruction_h_ = true; + need_mma_blockscaled_instruction_h_ = true; this->PrintIndent(); std::string mma_call = - "tl::mma_sync_blockscale<(AType), (BType), (CType), (M), (N), (K), " + "tl::mma_sync_blockscaled<(AType), (BType), (CType), (M), (N), (K), " "(TransA), (TransB), (ScaleVec)>(" "reinterpret_cast<(CRegType)*>((C_ptr) + (C_offset)), " "reinterpret_cast((A_ptr) + (A_offset)), " diff --git a/src/cuda/codegen/codegen_cuda.h b/src/cuda/codegen/codegen_cuda.h index f14a2bc5a3..05f3dc3745 100644 --- a/src/cuda/codegen/codegen_cuda.h +++ b/src/cuda/codegen/codegen_cuda.h @@ -116,7 +116,7 @@ class CodeGenTileLangCUDA final : public CodeGenC { // whether need tl mma instruction header bool need_mma_instruction_h_{false}; // whether need tl block-scaled mma instruction header - bool need_mma_blockscale_instruction_h_{false}; + bool need_mma_blockscaled_instruction_h_{false}; // whether need tl wgmma instruction header bool need_wgmma_instruction_h_{false}; // whether need tl tcgen05mma instruction header diff --git a/src/op/builtin.cc b/src/op/builtin.cc index 63c2923beb..7f3f39a95e 100644 --- a/src/op/builtin.cc +++ b/src/op/builtin.cc @@ -286,8 +286,8 @@ TIR_DEFINE_TL_BUILTIN(ptx_mma_sm70) .set_attr("TCallEffectKind", Integer(CallEffectKind::kOpaque)); -TIR_DEFINE_TL_BUILTIN(ptx_mma_blockscale) - .set_num_inputs(17) +TIR_DEFINE_TL_BUILTIN(ptx_mma_blockscaled) + .set_num_inputs(-1) .set_attr("TCallEffectKind", Integer(CallEffectKind::kOpaque)); diff --git a/src/op/builtin.h b/src/op/builtin.h index 4cdba72e25..5caf0e4acf 100644 --- a/src/op/builtin.h +++ b/src/op/builtin.h @@ -504,17 +504,20 @@ TVM_DLL const Op &ptx_mma_sm70(); * scale_vec=2/4, m16n8k64 for e2m1 x e2m1 inputs with ue8m0/ue4m3 scale * factors. * - * void ptx_mma_blockscale(StringImm shape, StringImm A_layout, - * StringImm B_layout, StringImm A_dtype, - * StringImm B_dtype, StringImm C_dtype, - * Integer scale_vec, + * void ptx_mma_blockscaled(StringImm dtype, StringImm shape, + * StringImm A_layout, StringImm B_layout, + * StringImm A_dtype, StringImm B_dtype, + * StringImm C_dtype, StringImm SF_dtype, + * Integer scale_vec_size, * Var multiplicand_a, Expr a_index, * Var multiplicand_b, Expr b_index, * Var accumulator, Expr c_index, - * Expr scale_a, Expr scale_b, - * Expr scale_id_a, Expr scale_id_b); + * Expr scale_a, Expr scale_b); + * + * TileLang also accepts optional Expr scale_id_a, Expr scale_id_b arguments + * after scale_b for selecting entries inside packed scale registers. */ -TVM_DLL const Op &ptx_mma_blockscale(); +TVM_DLL const Op &ptx_mma_blockscaled(); /*! * \brief tvm intrinsics for ldmatrix diff --git a/src/tl_templates/cuda/instruction/mma_blockscale.h b/src/tl_templates/cuda/instruction/mma_blockscaled.h similarity index 62% rename from src/tl_templates/cuda/instruction/mma_blockscale.h rename to src/tl_templates/cuda/instruction/mma_blockscaled.h index 4c5232ef7d..ecbb4dd6ea 100644 --- a/src/tl_templates/cuda/instruction/mma_blockscale.h +++ b/src/tl_templates/cuda/instruction/mma_blockscaled.h @@ -16,7 +16,7 @@ template inline constexpr bool always_false_v = false; namespace detail { -template struct MmaBlockScaleImplTraits { +template struct MmaBlockScaledImplTraits { using DReg = std::remove_extent_t; using AReg = std::remove_extent_t; using BReg = std::remove_extent_t; @@ -34,13 +34,13 @@ template struct MmaBlockScaleImplTraits { template -TL_DEVICE void call_fma_blockscale_impl( - typename MmaBlockScaleImplTraits::DReg *d, - const typename MmaBlockScaleImplTraits::AReg *a, - const typename MmaBlockScaleImplTraits::BReg *b, - const typename MmaBlockScaleImplTraits::CReg *c, - const typename MmaBlockScaleImplTraits::SFAReg *sfa, - const typename MmaBlockScaleImplTraits::SFBReg *sfb, uint16_t id_a, +TL_DEVICE void call_fma_blockscaled_impl( + typename MmaBlockScaledImplTraits::DReg *d, + const typename MmaBlockScaledImplTraits::AReg *a, + const typename MmaBlockScaledImplTraits::BReg *b, + const typename MmaBlockScaledImplTraits::CReg *c, + const typename MmaBlockScaledImplTraits::SFAReg *sfa, + const typename MmaBlockScaledImplTraits::SFBReg *sfb, uint16_t id_a, uint16_t id_b, std::index_sequence, std::index_sequence, std::index_sequence, std::index_sequence, std::index_sequence, std::index_sequence) { @@ -50,27 +50,27 @@ TL_DEVICE void call_fma_blockscale_impl( template TL_DEVICE void -call_fma_blockscale(typename MmaBlockScaleImplTraits::DReg *d, - const typename MmaBlockScaleImplTraits::AReg *a, - const typename MmaBlockScaleImplTraits::BReg *b, - const typename MmaBlockScaleImplTraits::CReg *c, - const typename MmaBlockScaleImplTraits::SFAReg *sfa, - const typename MmaBlockScaleImplTraits::SFBReg *sfb, - uint16_t id_a, uint16_t id_b) { - call_fma_blockscale_impl( +call_fma_blockscaled(typename MmaBlockScaledImplTraits::DReg *d, + const typename MmaBlockScaledImplTraits::AReg *a, + const typename MmaBlockScaledImplTraits::BReg *b, + const typename MmaBlockScaledImplTraits::CReg *c, + const typename MmaBlockScaledImplTraits::SFAReg *sfa, + const typename MmaBlockScaledImplTraits::SFBReg *sfb, + uint16_t id_a, uint16_t id_b) { + call_fma_blockscaled_impl( d, a, b, c, sfa, sfb, id_a, id_b, - std::make_index_sequence::kDRegs>{}, - std::make_index_sequence::kARegs>{}, - std::make_index_sequence::kBRegs>{}, - std::make_index_sequence::kCRegs>{}, - std::make_index_sequence::kSFARegs>{}, - std::make_index_sequence::kSFBRegs>{}); + std::make_index_sequence::kDRegs>{}, + std::make_index_sequence::kARegs>{}, + std::make_index_sequence::kBRegs>{}, + std::make_index_sequence::kCRegs>{}, + std::make_index_sequence::kSFARegs>{}, + std::make_index_sequence::kSFBRegs>{}); } #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1200 -#define TL_SM120_BLOCKSCALE_ASM(KindAsm, ScaleVecAsm, ShapeAsm, AAsm, BAsm, \ - SFAsm, d0, d1, d2, d3, a0, a1, a2, a3, b0, b1, \ - c0, c1, c2, c3, sfa, sfb, id_a, id_b) \ +#define TL_SM120_BLOCKSCALED_ASM(KindAsm, ScaleVecAsm, ShapeAsm, AAsm, BAsm, \ + SFAsm, d0, d1, d2, d3, a0, a1, a2, a3, b0, \ + b1, c0, c1, c2, c3, sfa, sfb, id_a, id_b) \ asm volatile("mma.sync.aligned.kind::" KindAsm \ ".block_scale.scale_vec::" ScaleVecAsm "." ShapeAsm \ ".row.col.f32." AAsm "." BAsm ".f32." SFAsm " " \ @@ -86,14 +86,14 @@ call_fma_blockscale(typename MmaBlockScaleImplTraits::DReg *d, "h"((short)0), "h"((short)id_a), "r"(uint32_t(sfb)), \ "h"((short)0), "h"((short)id_b)); #else -#define TL_SM120_BLOCKSCALE_ASM(KindAsm, ScaleVecAsm, ShapeAsm, AAsm, BAsm, \ - SFAsm, d0, d1, d2, d3, a0, a1, a2, a3, b0, b1, \ - c0, c1, c2, c3, sfa, sfb, id_a, id_b) \ +#define TL_SM120_BLOCKSCALED_ASM(KindAsm, ScaleVecAsm, ShapeAsm, AAsm, BAsm, \ + SFAsm, d0, d1, d2, d3, a0, a1, a2, a3, b0, \ + b1, c0, c1, c2, c3, sfa, sfb, id_a, id_b) \ CUTE_INVALID_CONTROL_PATH( \ "Attempting to use SM120 block-scaled MMA without SM120+"); #endif -#define TL_DEFINE_SM120_BLOCKSCALE_IMPL( \ +#define TL_DEFINE_SM120_BLOCKSCALED_IMPL( \ ImplName, KindAsm, ScaleVecAsm, ShapeAsm, AAsm, BAsm, SFAsm, ScaleRegType) \ struct ImplName { \ using DRegisters = float[4]; \ @@ -110,9 +110,9 @@ call_fma_blockscale(typename MmaBlockScaleImplTraits::DReg *d, float const &c1, float const &c2, float const &c3, \ ScaleRegType const &sfa, ScaleRegType const &sfb, uint16_t id_a, \ uint16_t id_b) { \ - TL_SM120_BLOCKSCALE_ASM(KindAsm, ScaleVecAsm, ShapeAsm, AAsm, BAsm, \ - SFAsm, d0, d1, d2, d3, a0, a1, a2, a3, b0, b1, \ - c0, c1, c2, c3, sfa, sfb, id_a, id_b) \ + TL_SM120_BLOCKSCALED_ASM(KindAsm, ScaleVecAsm, ShapeAsm, AAsm, BAsm, \ + SFAsm, d0, d1, d2, d3, a0, a1, a2, a3, b0, b1, \ + c0, c1, c2, c3, sfa, sfb, id_a, id_b) \ } \ }; @@ -123,60 +123,61 @@ call_fma_blockscale(typename MmaBlockScaleImplTraits::DReg *d, TL_SM120_MXF8F6F4_IMPL_NAME_(AName, BName) #define TL_SM120_MXF8F6F4_IMPL_NAME(ATypeEnum, BTypeEnum) \ - TL_SM120_MXF8F6F4_IMPL_NAME_EXPANDED(TL_BLOCKSCALE_DTYPE_ASM(ATypeEnum), \ - TL_BLOCKSCALE_DTYPE_ASM(BTypeEnum)) + TL_SM120_MXF8F6F4_IMPL_NAME_EXPANDED(TL_BLOCKSCALED_DTYPE_ASM(ATypeEnum), \ + TL_BLOCKSCALED_DTYPE_ASM(BTypeEnum)) #define TL_SM120_MXF4NVF4_IMPL_NAME_(ScaleName) \ SM120_16x8x64_F32_e2m1_e2m1_F32_##ScaleName##_TN -#define TL_BLOCKSCALE_DTYPE_ASM_kFloat4_e2m1fn e2m1 -#define TL_BLOCKSCALE_DTYPE_ASM_kFloat6_e2m3fn e2m3 -#define TL_BLOCKSCALE_DTYPE_ASM_kFloat6_e3m2fn e3m2 -#define TL_BLOCKSCALE_DTYPE_ASM_kFloat8_e4m3 e4m3 -#define TL_BLOCKSCALE_DTYPE_ASM_kFloat8_e5m2 e5m2 +#define TL_BLOCKSCALED_DTYPE_ASM_kFloat4_e2m1fn e2m1 +#define TL_BLOCKSCALED_DTYPE_ASM_kFloat6_e2m3fn e2m3 +#define TL_BLOCKSCALED_DTYPE_ASM_kFloat6_e3m2fn e3m2 +#define TL_BLOCKSCALED_DTYPE_ASM_kFloat8_e4m3 e4m3 +#define TL_BLOCKSCALED_DTYPE_ASM_kFloat8_e5m2 e5m2 -#define TL_BLOCKSCALE_DTYPE_ASM_(TypeEnum) TL_BLOCKSCALE_DTYPE_ASM_##TypeEnum -#define TL_BLOCKSCALE_DTYPE_ASM(TypeEnum) TL_BLOCKSCALE_DTYPE_ASM_(TypeEnum) +#define TL_BLOCKSCALED_DTYPE_ASM_(TypeEnum) TL_BLOCKSCALED_DTYPE_ASM_##TypeEnum +#define TL_BLOCKSCALED_DTYPE_ASM(TypeEnum) TL_BLOCKSCALED_DTYPE_ASM_(TypeEnum) -#define TL_SM120_BLOCKSCALE_DTYPE_ROWS(M) \ +#define TL_SM120_BLOCKSCALED_DTYPE_ROWS(M) \ M(kFloat4_e2m1fn) \ M(kFloat6_e2m3fn) \ M(kFloat6_e3m2fn) \ M(kFloat8_e4m3) \ M(kFloat8_e5m2) -#define TL_SM120_BLOCKSCALE_DTYPE_COLS(M, ATypeEnum) \ +#define TL_SM120_BLOCKSCALED_DTYPE_COLS(M, ATypeEnum) \ M(ATypeEnum, kFloat4_e2m1fn) \ M(ATypeEnum, kFloat6_e2m3fn) \ M(ATypeEnum, kFloat6_e3m2fn) \ M(ATypeEnum, kFloat8_e4m3) \ M(ATypeEnum, kFloat8_e5m2) -#define TL_BLOCKSCALE_STRINGIZE_(X) #X -#define TL_BLOCKSCALE_STRINGIZE(X) TL_BLOCKSCALE_STRINGIZE_(X) +#define TL_BLOCKSCALED_STRINGIZE_(X) #X +#define TL_BLOCKSCALED_STRINGIZE(X) TL_BLOCKSCALED_STRINGIZE_(X) #define TL_DEFINE_SM120_MXF8F6F4_IMPL(ATypeEnum, BTypeEnum) \ - TL_DEFINE_SM120_BLOCKSCALE_IMPL( \ + TL_DEFINE_SM120_BLOCKSCALED_IMPL( \ TL_SM120_MXF8F6F4_IMPL_NAME(ATypeEnum, BTypeEnum), "mxf8f6f4", "1X", \ - "m16n8k32", TL_BLOCKSCALE_STRINGIZE(TL_BLOCKSCALE_DTYPE_ASM(ATypeEnum)), \ - TL_BLOCKSCALE_STRINGIZE(TL_BLOCKSCALE_DTYPE_ASM(BTypeEnum)), "ue8m0", \ + "m16n8k32", \ + TL_BLOCKSCALED_STRINGIZE(TL_BLOCKSCALED_DTYPE_ASM(ATypeEnum)), \ + TL_BLOCKSCALED_STRINGIZE(TL_BLOCKSCALED_DTYPE_ASM(BTypeEnum)), "ue8m0", \ uint8_t) \ - TL_DEFINE_MMA_BLOCKSCALE_DISPATCHER( \ + TL_DEFINE_MMA_BLOCKSCALED_DISPATCHER( \ ATypeEnum, BTypeEnum, kFloat32, 16, 8, 32, false, true, 1, \ tl::detail::TL_SM120_MXF8F6F4_IMPL_NAME(ATypeEnum, BTypeEnum)) #define TL_DEFINE_SM120_MXF4NVF4_IMPL(ScaleVecValue, ScaleVecAsm, ScaleName) \ - TL_DEFINE_SM120_BLOCKSCALE_IMPL(TL_SM120_MXF4NVF4_IMPL_NAME_(ScaleName), \ - "mxf4nvf4", ScaleVecAsm, "m16n8k64", "e2m1", \ - "e2m1", TL_BLOCKSCALE_STRINGIZE(ScaleName), \ - uint32_t) \ - TL_DEFINE_MMA_BLOCKSCALE_DISPATCHER( \ + TL_DEFINE_SM120_BLOCKSCALED_IMPL( \ + TL_SM120_MXF4NVF4_IMPL_NAME_(ScaleName), "mxf4nvf4", ScaleVecAsm, \ + "m16n8k64", "e2m1", "e2m1", TL_BLOCKSCALED_STRINGIZE(ScaleName), \ + uint32_t) \ + TL_DEFINE_MMA_BLOCKSCALED_DISPATCHER( \ kFloat4_e2m1fn, kFloat4_e2m1fn, kFloat32, 16, 8, 64, false, true, \ ScaleVecValue, tl::detail::TL_SM120_MXF4NVF4_IMPL_NAME_(ScaleName)) template -struct MmaBlockScaleDispatcher { +struct MmaBlockScaledDispatcher { using CRegType = void; using ARegType = void; using BRegType = void; @@ -187,19 +188,19 @@ struct MmaBlockScaleDispatcher { const CRegType *, const SFARegType *, const SFBRegType *, uint16_t, uint16_t) { static_assert(always_false_v>, - "tl::mma_sync_blockscale: unsupported configuration"); + "tl::mma_sync_blockscaled: unsupported configuration"); } }; -#define TL_DEFINE_MMA_BLOCKSCALE_DISPATCHER( \ +#define TL_DEFINE_MMA_BLOCKSCALED_DISPATCHER( \ ATypeEnum, BTypeEnum, CTypeEnum, MValue, NValue, KValue, TransAValue, \ TransBValue, ScaleVecValue, ImplType) \ template <> \ - struct MmaBlockScaleDispatcher { \ + struct MmaBlockScaledDispatcher { \ using Impl = ImplType; \ - using Traits = MmaBlockScaleImplTraits; \ + using Traits = MmaBlockScaledImplTraits; \ using CRegType = typename Traits::DReg; \ using ARegType = typename Traits::AReg; \ using BRegType = typename Traits::BReg; \ @@ -207,26 +208,26 @@ struct MmaBlockScaleDispatcher { using SFBRegType = typename Traits::SFBReg; \ static_assert( \ std::is_same_v, \ - "tl::mma_sync_blockscale requires matching accumulator/output regs"); \ + "tl::mma_sync_blockscaled requires matching accumulator/output regs"); \ static TL_DEVICE void exec(CRegType *d, const ARegType *a, \ const BRegType *b, const CRegType *c, \ const SFARegType *sfa, const SFBRegType *sfb, \ uint16_t id_a, uint16_t id_b) { \ - call_fma_blockscale(d, a, b, c, sfa, sfb, id_a, id_b); \ + call_fma_blockscaled(d, a, b, c, sfa, sfb, id_a, id_b); \ } \ }; #define TL_DEFINE_SM120_MXF8F6F4_ROW(ATypeEnum) \ - TL_SM120_BLOCKSCALE_DTYPE_COLS(TL_DEFINE_SM120_MXF8F6F4_IMPL, ATypeEnum) + TL_SM120_BLOCKSCALED_DTYPE_COLS(TL_DEFINE_SM120_MXF8F6F4_IMPL, ATypeEnum) // SM120 dense block-scaled MMA, mxf8f6f4 scale_vec::1X, ue8m0 scale factors. -TL_SM120_BLOCKSCALE_DTYPE_ROWS(TL_DEFINE_SM120_MXF8F6F4_ROW) +TL_SM120_BLOCKSCALED_DTYPE_ROWS(TL_DEFINE_SM120_MXF8F6F4_ROW) // SM120 dense block-scaled MMA, mxf4nvf4 e2m1 x e2m1. TL_DEFINE_SM120_MXF4NVF4_IMPL(2, "2X", ue8m0) TL_DEFINE_SM120_MXF4NVF4_IMPL(4, "4X", ue4m3) -#undef TL_DEFINE_MMA_BLOCKSCALE_DISPATCHER +#undef TL_DEFINE_MMA_BLOCKSCALED_DISPATCHER #undef TL_DEFINE_SM120_MXF8F6F4_ROW #undef TL_DEFINE_SM120_MXF8F6F4_IMPL #undef TL_DEFINE_SM120_MXF4NVF4_IMPL @@ -234,41 +235,41 @@ TL_DEFINE_SM120_MXF4NVF4_IMPL(4, "4X", ue4m3) #undef TL_SM120_MXF8F6F4_IMPL_NAME #undef TL_SM120_MXF8F6F4_IMPL_NAME_EXPANDED #undef TL_SM120_MXF8F6F4_IMPL_NAME_ -#undef TL_BLOCKSCALE_STRINGIZE -#undef TL_BLOCKSCALE_STRINGIZE_ -#undef TL_SM120_BLOCKSCALE_DTYPE_COLS -#undef TL_SM120_BLOCKSCALE_DTYPE_ROWS -#undef TL_BLOCKSCALE_DTYPE_ASM -#undef TL_BLOCKSCALE_DTYPE_ASM_ -#undef TL_BLOCKSCALE_DTYPE_ASM_kFloat8_e5m2 -#undef TL_BLOCKSCALE_DTYPE_ASM_kFloat8_e4m3 -#undef TL_BLOCKSCALE_DTYPE_ASM_kFloat6_e3m2fn -#undef TL_BLOCKSCALE_DTYPE_ASM_kFloat6_e2m3fn -#undef TL_BLOCKSCALE_DTYPE_ASM_kFloat4_e2m1fn -#undef TL_DEFINE_SM120_BLOCKSCALE_IMPL -#undef TL_SM120_BLOCKSCALE_ASM +#undef TL_BLOCKSCALED_STRINGIZE +#undef TL_BLOCKSCALED_STRINGIZE_ +#undef TL_SM120_BLOCKSCALED_DTYPE_COLS +#undef TL_SM120_BLOCKSCALED_DTYPE_ROWS +#undef TL_BLOCKSCALED_DTYPE_ASM +#undef TL_BLOCKSCALED_DTYPE_ASM_ +#undef TL_BLOCKSCALED_DTYPE_ASM_kFloat8_e5m2 +#undef TL_BLOCKSCALED_DTYPE_ASM_kFloat8_e4m3 +#undef TL_BLOCKSCALED_DTYPE_ASM_kFloat6_e3m2fn +#undef TL_BLOCKSCALED_DTYPE_ASM_kFloat6_e2m3fn +#undef TL_BLOCKSCALED_DTYPE_ASM_kFloat4_e2m1fn +#undef TL_DEFINE_SM120_BLOCKSCALED_IMPL +#undef TL_SM120_BLOCKSCALED_ASM } // namespace detail template -TL_DEVICE void mma_sync_blockscale( - typename detail::MmaBlockScaleDispatcher< +TL_DEVICE void mma_sync_blockscaled( + typename detail::MmaBlockScaledDispatcher< AType, BType, CType, M, N, K, TransA, TransB, ScaleVec>::CRegType *c, - const typename detail::MmaBlockScaleDispatcher< + const typename detail::MmaBlockScaledDispatcher< AType, BType, CType, M, N, K, TransA, TransB, ScaleVec>::ARegType *a, - const typename detail::MmaBlockScaleDispatcher< + const typename detail::MmaBlockScaledDispatcher< AType, BType, CType, M, N, K, TransA, TransB, ScaleVec>::BRegType *b, - typename detail::MmaBlockScaleDispatcher< + typename detail::MmaBlockScaledDispatcher< AType, BType, CType, M, N, K, TransA, TransB, ScaleVec>::SFARegType sfa, - typename detail::MmaBlockScaleDispatcher< + typename detail::MmaBlockScaledDispatcher< AType, BType, CType, M, N, K, TransA, TransB, ScaleVec>::SFBRegType sfb, uint16_t id_a, uint16_t id_b) { using Dispatcher = - detail::MmaBlockScaleDispatcher; + detail::MmaBlockScaledDispatcher; static_assert(!std::is_void_v, - "tl::mma_sync_blockscale: unsupported configuration"); + "tl::mma_sync_blockscaled: unsupported configuration"); Dispatcher::exec(c, a, b, c, &sfa, &sfb, id_a, id_b); } diff --git a/tilelang/language/ast/ir.py b/tilelang/language/ast/ir.py index 0eb16ad81a..7117f3c07a 100644 --- a/tilelang/language/ast/ir.py +++ b/tilelang/language/ast/ir.py @@ -1882,7 +1882,7 @@ def wrapped(*args, **kwargs): call_pure_extern = _dtype_forward(_tir_op.call_pure_extern) ptx_mma = _dtype_forward(_tir_op.ptx_mma) ptx_mma_sp = _dtype_forward(_tir_op.ptx_mma_sp) -ptx_mma_blockscale = _dtype_forward(_tir_op.ptx_mma_blockscale) +ptx_mma_blockscaled = _dtype_forward(_tir_op.ptx_mma_blockscaled) ptx_wgmma_ss = _dtype_forward(_tir_op.ptx_wgmma_ss) ptx_wgmma_rs = _dtype_forward(_tir_op.ptx_wgmma_rs) ptx_wgmma_sp_ss = _dtype_forward(_tir_op.ptx_wgmma_sp_ss) @@ -2138,7 +2138,7 @@ def wrapped(*args, **kwargs): "tvm_warp_activemask", "ptx_mma", "ptx_mma_sp", - "ptx_mma_blockscale", + "ptx_mma_blockscaled", "ptx_wgmma_ss", "ptx_wgmma_rs", "ptx_wgmma_sp_ss", diff --git a/tilelang/language/tir/ir.py b/tilelang/language/tir/ir.py index cf3caabf1a..36cb01fc44 100644 --- a/tilelang/language/tir/ir.py +++ b/tilelang/language/tir/ir.py @@ -289,7 +289,7 @@ def wrapped(*args, **kwargs): call_pure_extern = _dtype_forward(_tir_op.call_pure_extern) ptx_mma = _dtype_forward(_tir_op.ptx_mma) ptx_mma_sp = _dtype_forward(_tir_op.ptx_mma_sp) -ptx_mma_blockscale = _dtype_forward(_tir_op.ptx_mma_blockscale) +ptx_mma_blockscaled = _dtype_forward(_tir_op.ptx_mma_blockscaled) ptx_wgmma_ss = _dtype_forward(_tir_op.ptx_wgmma_ss) ptx_wgmma_rs = _dtype_forward(_tir_op.ptx_wgmma_rs) ptx_wgmma_sp_ss = _dtype_forward(_tir_op.ptx_wgmma_sp_ss) diff --git a/tilelang/language/tir/op.py b/tilelang/language/tir/op.py index 1c7684e2cb..147ab0c3d2 100644 --- a/tilelang/language/tir/op.py +++ b/tilelang/language/tir/op.py @@ -1076,24 +1076,26 @@ def ptx_mma_sp( ) -def ptx_mma_blockscale( - C_dtype, +def ptx_mma_blockscaled( + dtype, shape, A_layout, B_layout, A_dtype, B_dtype, - scale_vec, + C_dtype, + SF_dtype, + scale_vec_size, multiplicand_a, a_index, multiplicand_b, b_index, accumulator, c_index, - scale_a, - scale_b, - scale_id_a, - scale_id_b, + scale_factor_a, + scale_factor_b, + scale_id_a=None, + scale_id_b=None, ): """TileLang intrinsic for PTX block-scaled MMA instructions. @@ -1101,32 +1103,39 @@ def ptx_mma_blockscale( ``accumulator`` operand. The intrinsic itself is side-effect only and emits a ``handle``-typed TIR call, like TCGEN05 MMA intrinsics. The CUDA backend maps the shape/layout/dtype tuple to the matching - ``tl::mma_sync_blockscale`` dispatcher specialization. On SM120, + ``tl::mma_sync_blockscaled`` dispatcher specialization. On SM120, ``scale_vec=1`` selects ``mxf8f6f4``/``ue8m0`` ``m16n8k32`` variants for e2m1/e3m2/e2m3/e4m3/e5m2 inputs, while ``scale_vec=2`` and ``scale_vec=4`` select ``mxf4nvf4`` ``m16n8k64`` e2m1 inputs with ``ue8m0`` and ``ue4m3`` scale factors respectively. """ - return call_intrin( - "handle", - _tvm_op.Op.get("tl.ptx_mma_blockscale"), + args = [ + dtype, shape, A_layout, B_layout, A_dtype, B_dtype, C_dtype, - scale_vec, + SF_dtype, + IntImm("int32", scale_vec_size), multiplicand_a, a_index, multiplicand_b, b_index, accumulator, c_index, - scale_a, - scale_b, - scale_id_a, - scale_id_b, + scale_factor_a, + scale_factor_b, + ] + if (scale_id_a is None) != (scale_id_b is None): + raise ValueError("scale_id_a and scale_id_b must be provided together") + if scale_id_a is not None: + args.extend([scale_id_a, scale_id_b]) + return call_intrin( + "handle", + _tvm_op.Op.get("tl.ptx_mma_blockscaled"), + *args, ) From 4551f64b07d912e03f3ebdfd56a7ce5f558601e0 Mon Sep 17 00:00:00 2001 From: sepcnt <30561671+sepcnt@users.noreply.github.com> Date: Fri, 12 Jun 2026 15:30:37 +0800 Subject: [PATCH 6/6] Fix tma load_gather4 and scatter4 --- src/transform/lower_tile_op.cc | 4 +++- .../language/test_tilelang_language_tma_gather_scatter.py | 2 +- tilelang/language/copy_op.py | 4 ++-- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/transform/lower_tile_op.cc b/src/transform/lower_tile_op.cc index d861e5aa2a..13102f2a97 100644 --- a/src/transform/lower_tile_op.cc +++ b/src/transform/lower_tile_op.cc @@ -835,7 +835,9 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer { if (op->op.same_as(tl::tma_load()) || op->op.same_as(tl::tma_load_im2col()) || op->op.same_as(tl::tma_load_multicast()) || - op->op.same_as(tl::tma_store())) { + op->op.same_as(tl::tma_store()) || + op->op.same_as(tl::tma_load_gather4()) || + op->op.same_as(tl::tma_store_scatter4())) { // skip tma related calls, as they were transformed implicitly. has_tma_ = true; in_tma_context_ = true; diff --git a/testing/python/language/test_tilelang_language_tma_gather_scatter.py b/testing/python/language/test_tilelang_language_tma_gather_scatter.py index 58ffd0ae90..a677b0bd53 100644 --- a/testing/python/language/test_tilelang_language_tma_gather_scatter.py +++ b/testing/python/language/test_tilelang_language_tma_gather_scatter.py @@ -117,7 +117,7 @@ def test_gather4_fp4_unpack_descriptor_codegen(): device_ir = str(artifact.device_mod) assert 'T.handle("float4_e2m1fn", "global")' in host_ir assert 'smem = T.alloc_buffer((512,), "custom[float4_e2m1_unpacked]"' in device_ir - assert 'T.call_extern("handle", "tl.tma_load_gather4"' in device_ir + assert "T.tma_load_gather4(" in device_ir assert re.search(r'\["__tvm_tensormap_create_tiled", \w+, 14,', host_ir) assert re.search(r"T\.mbarrier_expect_tx\([^,]+, 256\)", device_ir) diff --git a/tilelang/language/copy_op.py b/tilelang/language/copy_op.py index 37528e63ee..6f6b196d62 100644 --- a/tilelang/language/copy_op.py +++ b/tilelang/language/copy_op.py @@ -391,7 +391,7 @@ def tma_gather4( inner = src.strides[1] if not ((isinstance(inner, int) and inner == 1) or (hasattr(inner, "value") and int(inner.value) == 1)): raise ValueError(f"tma_gather4 requires unit innermost global stride, got {inner}") - rows = list(rows) + rows = [tirx.IntImm("int32", row) if isinstance(row, int) else row for row in rows] if len(rows) != 4: raise ValueError(f"tma_gather4 expects exactly 4 row indices, got {len(rows)}") if swizzle not in (None, "none", 0): @@ -486,7 +486,7 @@ def tma_scatter4( inner = dst.strides[1] if not ((isinstance(inner, int) and inner == 1) or (hasattr(inner, "value") and int(inner.value) == 1)): raise ValueError(f"tma_scatter4 requires unit innermost global stride, got {inner}") - rows = list(rows) + rows = [tirx.IntImm("int32", row) if isinstance(row, int) else row for row in rows] if len(rows) != 4: raise ValueError(f"tma_scatter4 expects exactly 4 row indices, got {len(rows)}") if swizzle not in (None, "none", 0):