From 17c0288d4fe68810d812537d181cce98fa0a8ba1 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 28 Jul 2026 23:47:09 -0700 Subject: [PATCH 01/59] Add CUDA compute capabilities 8.9, 9.0, 10.0 and 12.0 The highest we could target was 8.6, so on anything newer we emitted sm_86 PTX and left it to the driver to JIT. Now Ada, Hopper and both Blackwells can be named, and host target detection picks them up. It makes no measurable difference to the tensor core matmul, which is what prompted this: the kernel only uses instructions that have been available since 7.0, so there's nothing for a newer target to do. cuBLAS on the same device runs an sm_80 kernel for the same reason. Co-Authored-By: Claude Opus 5 --- src/CodeGen_PTX_Dev.cpp | 20 ++++++++++++++-- src/Target.cpp | 48 +++++++++++++++++++++++++++++++++++-- src/Target.h | 4 ++++ src/runtime/HalideRuntime.h | 6 ++++- 4 files changed, 73 insertions(+), 5 deletions(-) diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 0414b33fd243..1a04fbcfa770 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -551,7 +551,15 @@ void CodeGen_PTX_Dev::codegen_vector_reduce(const VectorReduce *op, const Expr & } string CodeGen_PTX_Dev::mcpu_target() const { - if (target.has_feature(Target::CUDACapability86)) { + if (target.has_feature(Target::CUDACapability120)) { + return "sm_120"; + } else if (target.has_feature(Target::CUDACapability100)) { + return "sm_100"; + } else if (target.has_feature(Target::CUDACapability90)) { + return "sm_90"; + } else if (target.has_feature(Target::CUDACapability89)) { + return "sm_89"; + } else if (target.has_feature(Target::CUDACapability86)) { return "sm_86"; } else if (target.has_feature(Target::CUDACapability80)) { return "sm_80"; @@ -579,7 +587,15 @@ string CodeGen_PTX_Dev::mcpu_tune() const { } string CodeGen_PTX_Dev::mattrs() const { - if (target.has_feature(Target::CUDACapability86)) { + if (target.has_feature(Target::CUDACapability120)) { + return "+ptx87"; + } else if (target.has_feature(Target::CUDACapability100)) { + return "+ptx86"; + } else if (target.has_feature(Target::CUDACapability90)) { + return "+ptx78"; + } else if (target.has_feature(Target::CUDACapability89)) { + return "+ptx78"; + } else if (target.has_feature(Target::CUDACapability86)) { return "+ptx71"; } else if (target.has_feature(Target::CUDACapability80)) { return "+ptx70"; diff --git a/src/Target.cpp b/src/Target.cpp index 048a4383acd3..3a5856c681fe 100644 --- a/src/Target.cpp +++ b/src/Target.cpp @@ -642,8 +642,16 @@ Target::Feature calculate_host_cuda_capability(Target t) { return Target::CUDACapability75; } else if (ver < 86) { return Target::CUDACapability80; - } else { + } else if (ver < 89) { return Target::CUDACapability86; + } else if (ver < 90) { + return Target::CUDACapability89; + } else if (ver < 100) { + return Target::CUDACapability90; + } else if (ver < 120) { + return Target::CUDACapability100; + } else { + return Target::CUDACapability120; } } @@ -778,6 +786,10 @@ const std::map feature_name_map = { {"cuda_capability_75", Target::CUDACapability75}, {"cuda_capability_80", Target::CUDACapability80}, {"cuda_capability_86", Target::CUDACapability86}, + {"cuda_capability_89", Target::CUDACapability89}, + {"cuda_capability_90", Target::CUDACapability90}, + {"cuda_capability_100", Target::CUDACapability100}, + {"cuda_capability_120", Target::CUDACapability120}, {"opencl", Target::OpenCL}, {"cl_doubles", Target::CLDoubles}, {"cl_half", Target::CLHalf}, @@ -1008,7 +1020,11 @@ bool merge_string(Target &t, const std::string &target) { !t.has_feature(Target::CUDACapability70) && !t.has_feature(Target::CUDACapability75) && !t.has_feature(Target::CUDACapability80) && - !t.has_feature(Target::CUDACapability86)) { + !t.has_feature(Target::CUDACapability86) && + !t.has_feature(Target::CUDACapability89) && + !t.has_feature(Target::CUDACapability90) && + !t.has_feature(Target::CUDACapability100) && + !t.has_feature(Target::CUDACapability120)) { // Detect host cuda capability t.set_feature(get_host_cuda_capability(t)); } @@ -1545,6 +1561,18 @@ int Target::get_cuda_capability_lower_bound() const { if (has_feature(Target::CUDACapability86)) { return 86; } + if (has_feature(Target::CUDACapability89)) { + return 89; + } + if (has_feature(Target::CUDACapability90)) { + return 90; + } + if (has_feature(Target::CUDACapability100)) { + return 100; + } + if (has_feature(Target::CUDACapability120)) { + return 120; + } return 20; } @@ -1875,6 +1903,10 @@ bool Target::get_runtime_compatible_target(const Target &other, Target &result) CUDACapability75, CUDACapability80, CUDACapability86, + CUDACapability89, + CUDACapability90, + CUDACapability100, + CUDACapability120, HVX_v62, HVX_v65, @@ -2013,6 +2045,18 @@ bool Target::get_runtime_compatible_target(const Target &other, Target &result) if (cuda_capability < 86) { output.features.reset(CUDACapability86); } + if (cuda_capability < 89) { + output.features.reset(CUDACapability89); + } + if (cuda_capability < 90) { + output.features.reset(CUDACapability90); + } + if (cuda_capability < 100) { + output.features.reset(CUDACapability100); + } + if (cuda_capability < 120) { + output.features.reset(CUDACapability120); + } // Pick tight lower bound for Vulkan capability. Use fall-through to clear redundant features int vulkan_a = get_vulkan_capability_lower_bound(); diff --git a/src/Target.h b/src/Target.h index aab31554cb37..e0c11c57af8b 100644 --- a/src/Target.h +++ b/src/Target.h @@ -108,6 +108,10 @@ struct Target { CUDACapability75 = halide_target_feature_cuda_capability75, CUDACapability80 = halide_target_feature_cuda_capability80, CUDACapability86 = halide_target_feature_cuda_capability86, + CUDACapability89 = halide_target_feature_cuda_capability89, + CUDACapability90 = halide_target_feature_cuda_capability90, + CUDACapability100 = halide_target_feature_cuda_capability100, + CUDACapability120 = halide_target_feature_cuda_capability120, OpenCL = halide_target_feature_opencl, CLDoubles = halide_target_feature_cl_doubles, CLHalf = halide_target_feature_cl_half, diff --git a/src/runtime/HalideRuntime.h b/src/runtime/HalideRuntime.h index d35d52978d41..2f18755085f2 100644 --- a/src/runtime/HalideRuntime.h +++ b/src/runtime/HalideRuntime.h @@ -1466,7 +1466,11 @@ typedef enum halide_target_feature_t { halide_target_feature_cuda_capability70, ///< Enable CUDA compute capability 7.0 (Volta) halide_target_feature_cuda_capability75, ///< Enable CUDA compute capability 7.5 (Turing) halide_target_feature_cuda_capability80, ///< Enable CUDA compute capability 8.0 (Ampere) - halide_target_feature_cuda_capability86, ///< Enable CUDA compute capability 8.6 (Ampere) + halide_target_feature_cuda_capability86, ///< Enable CUDA compute capability 8.6 (Ampere) + halide_target_feature_cuda_capability89, ///< Enable CUDA compute capability 8.9 (Ada) + halide_target_feature_cuda_capability90, ///< Enable CUDA compute capability 9.0 (Hopper) + halide_target_feature_cuda_capability100, ///< Enable CUDA compute capability 10.0 (Blackwell) + halide_target_feature_cuda_capability120, ///< Enable CUDA compute capability 12.0 (Blackwell) halide_target_feature_opencl, ///< Enable the OpenCL runtime. halide_target_feature_cl_doubles, ///< Enable double support on OpenCL targets From 398506f1e10b1a4e9becb9b8f06bdc20c9959404 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Thu, 30 Jul 2026 11:31:29 -0700 Subject: [PATCH 02/59] Expose the new CUDA capabilities to the Python bindings The comment above halide_target_feature_t lists three places to keep in sync when adding a feature, and PyEnums.cpp was the one I missed. Co-Authored-By: Claude Opus 5 --- python_bindings/src/halide/halide_/PyEnums.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python_bindings/src/halide/halide_/PyEnums.cpp b/python_bindings/src/halide/halide_/PyEnums.cpp index c0f68be20fb3..b18fe519db4f 100644 --- a/python_bindings/src/halide/halide_/PyEnums.cpp +++ b/python_bindings/src/halide/halide_/PyEnums.cpp @@ -140,6 +140,10 @@ void define_enums(py::module &m) { .value("CUDACapability75", Target::Feature::CUDACapability75) .value("CUDACapability80", Target::Feature::CUDACapability80) .value("CUDACapability86", Target::Feature::CUDACapability86) + .value("CUDACapability89", Target::Feature::CUDACapability89) + .value("CUDACapability90", Target::Feature::CUDACapability90) + .value("CUDACapability100", Target::Feature::CUDACapability100) + .value("CUDACapability120", Target::Feature::CUDACapability120) .value("OpenCL", Target::Feature::OpenCL) .value("CLDoubles", Target::Feature::CLDoubles) .value("CLHalf", Target::Feature::CLHalf) From fac42093f9c01497bf7264aa8cbda371e30bb39f Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Thu, 30 Jul 2026 11:51:42 -0700 Subject: [PATCH 03/59] Let the driver choose the CUDA register count The cuda runtime capped kernels at 64 registers per thread when loading a module, with an HL_CUDA_JIT_MAX_REGISTERS environment variable as an escape hatch. Capping registers trades spilling against occupancy, and ptxas has more information about the kernel and the device than we do. Two apps set the escape hatch to 256 to compensate. Both are removed here; on an sm_86 device the 256 cap is now slightly slower than letting the driver decide. Co-Authored-By: Claude Opus 5 --- apps/conv_layer/process.cpp | 9 --------- apps/cuda_mat_mul/Makefile | 2 +- src/runtime/cuda.cpp | 16 ++++------------ 3 files changed, 5 insertions(+), 22 deletions(-) diff --git a/apps/conv_layer/process.cpp b/apps/conv_layer/process.cpp index 1a0eecc4d38a..89d25423a310 100644 --- a/apps/conv_layer/process.cpp +++ b/apps/conv_layer/process.cpp @@ -43,15 +43,6 @@ int main(int argc, char **argv) { Buffer output(CO, W, H, N); -// This is necessary to get the PTX compiler to do a good -// job. TODO: This should be a scheduling directive or a runtime -// function. -#ifdef _WIN32 - _putenv_s("HL_CUDA_JIT_MAX_REGISTERS", "256"); -#else - setenv("HL_CUDA_JIT_MAX_REGISTERS", "256", 1); -#endif - conv_layer(input, filter, bias, output); // Timing code diff --git a/apps/cuda_mat_mul/Makefile b/apps/cuda_mat_mul/Makefile index 2467a7ca9cc9..e0dfb78900fe 100644 --- a/apps/cuda_mat_mul/Makefile +++ b/apps/cuda_mat_mul/Makefile @@ -22,7 +22,7 @@ $(BIN)/%/runner: runner.cpp $(BIN)/%/mat_mul.a $(CXX) $(CXXFLAGS) -I$(BIN)/$* -Wall $^ -o $@ $(LDFLAGS) -lcudart -lcublas test: $(BIN)/$(HL_TARGET)/runner - HL_CUDA_JIT_MAX_REGISTERS=256 $^ $(MATRIX_SIZE) + $^ $(MATRIX_SIZE) clean: rm -rf $(BIN) diff --git a/src/runtime/cuda.cpp b/src/runtime/cuda.cpp index 90a4b9d47cfd..c4349aac63fd 100644 --- a/src/runtime/cuda.cpp +++ b/src/runtime/cuda.cpp @@ -526,19 +526,11 @@ WEAK int validate_device_pointer(void *user_context, halide_buffer_t *buf, size_ WEAK CUmodule compile_kernel(void *user_context, const char *ptx_src, int size) { debug(user_context) << "CUDA: compile_kernel cuModuleLoadData " << (void *)ptx_src << ", " << size << " -> "; - CUjit_option options[] = {CU_JIT_MAX_REGISTERS}; - unsigned int max_regs_per_thread = 64; - - // A hack to enable control over max register count for - // testing. This should be surfaced in the schedule somehow - // instead. - char *regs = getenv("HL_CUDA_JIT_MAX_REGISTERS"); - if (regs) { - max_regs_per_thread = atoi(regs); - } - void *optionValues[] = {(void *)(uintptr_t)max_regs_per_thread}; + // Let the driver pick the register count. Capping it trades spilling + // against occupancy, and the driver has more information about the kernel + // and the device than we do. CUmodule loaded_module; - CUresult err = cuModuleLoadDataEx(&loaded_module, ptx_src, 1, options, optionValues); + CUresult err = cuModuleLoadData(&loaded_module, ptx_src); if (err != CUDA_SUCCESS) { error(user_context) << "CUDA: cuModuleLoadData failed: " From 7aa69baca2974a7b36ff2bd6b497d285bcf413ac Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Thu, 30 Jul 2026 12:38:40 -0700 Subject: [PATCH 04/59] Simplify a comment in the cuda runtime Co-Authored-By: Claude Opus 5 --- src/runtime/cuda.cpp | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/runtime/cuda.cpp b/src/runtime/cuda.cpp index c4349aac63fd..2d0e83f9a477 100644 --- a/src/runtime/cuda.cpp +++ b/src/runtime/cuda.cpp @@ -526,9 +526,7 @@ WEAK int validate_device_pointer(void *user_context, halide_buffer_t *buf, size_ WEAK CUmodule compile_kernel(void *user_context, const char *ptx_src, int size) { debug(user_context) << "CUDA: compile_kernel cuModuleLoadData " << (void *)ptx_src << ", " << size << " -> "; - // Let the driver pick the register count. Capping it trades spilling - // against occupancy, and the driver has more information about the kernel - // and the device than we do. + // Use driver defaults for all JIT options. CUmodule loaded_module; CUresult err = cuModuleLoadData(&loaded_module, ptx_src); From 2c7aa914d634a73e3f80676092eb5aed0c625cb5 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Thu, 30 Jul 2026 13:07:19 -0700 Subject: [PATCH 05/59] Move subtile analysis out of the AMX pass and into MultiRamp get_subtile partitions accesses to a tile-memory allocation between the distinct sub-tiles it holds, and is_load_of_multiramp digs a load out from under the casts, broadcasts and lane permutations that can wrap it. Neither is specific to AMX, so move them next to the MultiRamp machinery they are built on. This drops 126 lines from ExtractTileOperations. Supporting this, is_multiramp learns to see through a shuffle of a single vector when the shuffle is a reshaping rather than a gather: either a transpose, or any permutation of a one-dimensional multiramp whose lane indices are themselves a multiramp. MultiRamp::transpose applies the former to a multiramp directly, splitting a dim in two where the transpose falls inside one. Routing the AMX operands through is_load_of_multiramp lets that pass match loads it previously missed, and lets it check the operand types against the type actually loaded from memory rather than the type of the expression wrapping it. Co-Authored-By: Claude Opus 5 --- src/ExtractTileOperations.cpp | 145 ++------------- src/MultiRamp.cpp | 218 +++++++++++++++++++++++ src/MultiRamp.h | 30 ++++ test/correctness/multiramp.cpp | 132 ++++++++++++++ test/correctness/tiled_matmul_errors.cpp | 2 +- 5 files changed, 400 insertions(+), 127 deletions(-) diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp index 2a8ebd842fbf..d835fde34568 100644 --- a/src/ExtractTileOperations.cpp +++ b/src/ExtractTileOperations.cpp @@ -121,11 +121,6 @@ Matmul convert_to_matmul(const Store *op, const string &new_name) { return fail("the vector reduction is not of a widening multiply"); } - if (lhs.type().bits() != 8 || - rhs.type().bits() != 8) { - return fail("the vector reduction operand or result types are not supported"); - } - } else { // Lower a widening_mul intrinsic, as they can be used but aren't lifted to for bf16. Expr reduce_value = simplify(lower_intrinsics(reduce->value)); @@ -137,65 +132,31 @@ Matmul convert_to_matmul(const Store *op, const string &new_name) { rhs = mul->b; } - // There may be a broadcast next (it can get hoisted outside of other ops) - auto debroadcast = [](Expr &e) -> int { - if (const Broadcast *b = e.as()) { - int lanes = b->lanes; - e = b->value; - return lanes; - } else { - return 1; - } - }; - int lhs_broadcast = debroadcast(lhs); - int rhs_broadcast = debroadcast(rhs); - - // Unpack the casts, if it was a direct multiply. This should only happen - // for floats (the integer branch above already extracted the cast inputs - // from the widening_mul intrinsic). - if (reduce->type.is_float()) { - const auto *lhs_cast = lhs.as(); - const auto *rhs_cast = rhs.as(); - if (!lhs_cast || !rhs_cast) { - return fail("the vector reduction is not of a widening multiply"); - } - lhs = lhs_cast->value; - rhs = rhs_cast->value; - if (!lhs.type().is_bfloat() || - rhs.type().element_of() != lhs.type().element_of()) { - return fail("the vector reduction operand or result types are not supported"); - } - } - - // Underneath all of this must be a load + // Underneath all of this must be a load, though it may be wrapped in a + // broadcast over the dimension it doesn't depend on, in the widening cast + // (for floats - the integer branch above already extracted the cast inputs + // from the widening_mul intrinsic), and in a lane permutation. // TODO: What if we want to multiply by the same matrix multiple times? It might be a let binding. - const auto *lhs_load = lhs.as(); - const auto *rhs_load = rhs.as(); + MultiRamp lhs_mr, rhs_mr; + Scope empty_scope; + const auto *lhs_load = is_load_of_multiramp(lhs, empty_scope, &lhs_mr); + const auto *rhs_load = is_load_of_multiramp(rhs, empty_scope, &rhs_mr); if (!lhs_load || !rhs_load) { - return fail("the matrix multiply operands are not loads"); + return fail("the matrix multiply operands are not loads with affine indices"); } // The loads must be unpredicated if (!is_const_one(lhs_load->predicate) || !is_const_one(rhs_load->predicate)) { return fail("the matrix multiply operands are predicated loads"); } - // Now we analyze the load indices as multiramps - MultiRamp lhs_mr, rhs_mr; - Scope empty_scope; - if (!is_multiramp(lhs_load->index, empty_scope, &lhs_mr) || - !is_multiramp(rhs_load->index, empty_scope, &rhs_mr)) { - return fail("the matrix multiply loads indices are not affine"); - } - - // Add back on any broadcasts as a stride-0 outer dim. - auto add_broadcast = [](MultiRamp &mr, int extent) { - if (extent > 1) { - mr.strides.push_back(make_zero(mr.base.type())); - mr.lanes.push_back(extent); + if (reduce->type.is_int_or_uint()) { + if (lhs_load->type.bits() != 8 || rhs_load->type.bits() != 8) { + return fail("the vector reduction operand or result types are not supported"); } - }; - add_broadcast(lhs_mr, lhs_broadcast); - add_broadcast(rhs_mr, rhs_broadcast); + } else if (!lhs_load->type.is_bfloat() || + rhs_load->type.element_of() != lhs_load->type.element_of()) { + return fail("the vector reduction operand or result types are not supported"); + } // In a matrix multiply with row-major inputs and outputs, the algorithm // looks like: @@ -234,7 +195,6 @@ Matmul convert_to_matmul(const Store *op, const string &new_name) { // Now deduce LHS and RHS. First some helpers. auto swap_sides = [&]() { - std::swap(lhs, rhs); std::swap(lhs_mr, rhs_mr); std::swap(lhs_load, rhs_load); }; @@ -309,8 +269,8 @@ Matmul convert_to_matmul(const Store *op, const string &new_name) { { Type t = op->value.type(); bool result_ok = t.bytes() * I * J <= 1024; - bool lhs_ok = lhs.type().bytes() * I * K <= 1024; - bool rhs_ok = rhs.type().bytes() * K * J <= 1024; + bool lhs_ok = lhs_load->type.bytes() * I * K <= 1024; + bool rhs_ok = rhs_load->type.bytes() * K * J <= 1024; if (!result_ok || !lhs_ok || !rhs_ok) { return fail("one more more matrices are too large to fit in AMX registers (more than 1024 bytes)"); } @@ -428,76 +388,9 @@ class ExtractTileOperations : public IRMutator { // registers as 2D sub-tiles. This map tracks those. std::vector amx_subtiles; - // Returns a unique subtile index for a load or store index, or -1 if it - // overlaps with an existing subtile, or is otherwise poorly behaved. - int get_subtile(const Expr &index) { - MultiRamp mr; - if (!is_multiramp(index, Scope::empty_scope(), &mr)) { - user_error << "Access to AMX tile not affine: " << index << "\n"; - } - if (!can_prove(mr.alias_free())) { - // What are you doing? - user_error << "Access to AMX tile may have duplicated lanes: " << index << "\n"; - } - if (amx_subtiles.empty()) { - amx_subtiles.push_back(std::move(mr)); - return 0; - } - - // All strides and lanes must match across all subtiles, or we give up. - const MultiRamp &first = amx_subtiles[0]; - if (mr.dimensions() != first.dimensions()) { - user_error - << "Access to AMX tile does not have the same shape as other accesses to the same memory."; - return -1; - } - for (int i = 0; i < first.dimensions(); i++) { - if (!can_prove(mr.strides[i] == first.strides[i]) || - mr.lanes[i] != first.lanes[i]) { - user_error - << "Access to AMX tile has different size and strides to other " - << "accesses to the same memory. All accesses must have the same " - << "subtile size and strides: " << index; - } - } - - // Now check for disjointedness - // Add a synthetic dimension, the purpose of which will become clear. - mr.strides.emplace_back(); - mr.lanes.push_back(2); - for (int i = 0; i < (int)amx_subtiles.size(); i++) { - auto &other = amx_subtiles[i]; - // One of two things must be true: - // 1) All of the lanes of mr equal the corresponding lane of - // other. We've already checked the strides and lanes, so it's just - // a matter of checking the base. - if (can_prove(mr.base == other.base)) { - return i; - } - - // 2) None of the lanes or mr equal any of the lanes of other. To do - // this we'll construct a combined mr that can be either 'mr' or - // 'other', and ask if it's alias-free. This is what the synthetic - // dimension was for. - mr.strides.back() = mr.base - other.base; - if (!can_prove(mr.alias_free())) { - user_error - << "Failed to prove access to AMX does not partially overlap " - << "another distinct access: " << index; - return -1; - } - } - - // Didn't already exist and didn't alias with anything. - mr.strides.pop_back(); - mr.lanes.pop_back(); - amx_subtiles.push_back(std::move(mr)); - return (int)amx_subtiles.size() - 1; - } - // Returns an index expression for a given load or store index. user_asserts if impossible std::string get_subtile_name(const Expr &index) { - int idx = get_subtile(index); + int idx = get_subtile(index, "AMX tile", &amx_subtiles); internal_assert(idx >= 0); // errors handled already return amx_name + std::to_string(idx); } diff --git a/src/MultiRamp.cpp b/src/MultiRamp.cpp index 3bd52e14d99e..695032f7af87 100644 --- a/src/MultiRamp.cpp +++ b/src/MultiRamp.cpp @@ -410,6 +410,52 @@ std::optional unbroadcast(const Expr &e) { } } +// Recognize a list of constant lane indices as a MultiRamp of constants, which +// is what it means for a shuffle to be a reshaping of its input rather than an +// arbitrary gather. +bool multiramp_of_constants(const std::vector &idx, Type t, MultiRamp *result) { + const int n = (int)idx.size(); + internal_assert(n > 0); + if (n == 1) { + *result = MultiRamp(make_const(t, idx[0]), {}, {}); + return true; + } + + // The innermost dim is the longest prefix that's an arithmetic progression. + const int stride = idx[1] - idx[0]; + int extent = 1; + while (extent < n && idx[extent] == idx[0] + extent * stride) { + extent++; + } + if (n % extent) { + return false; + } + + // Every block of that length has to be the same progression. + std::vector starts; + starts.reserve(n / extent); + for (int b = 0; b < n; b += extent) { + for (int j = 0; j < extent; j++) { + if (idx[b + j] != idx[b] + j * stride) { + return false; + } + } + starts.push_back(idx[b]); + } + + MultiRamp outer; + if (!multiramp_of_constants(starts, t, &outer)) { + return false; + } + + std::vector strides{make_const(t, stride)}; + strides.insert(strides.end(), outer.strides.begin(), outer.strides.end()); + std::vector lanes{extent}; + lanes.insert(lanes.end(), outer.lanes.begin(), outer.lanes.end()); + *result = MultiRamp(outer.base, strides, lanes); + return true; +} + // Internal is_multiramp. May leave *result in a partial state on failure; // the public is_multiramp below protects callers by only committing on // success. Recursive calls go through the public wrapper, so each branch @@ -429,6 +475,25 @@ bool is_multiramp_impl(const Expr &e, const Scope &scope, MultiRamp *resul result->strides.push_back(make_zero(elem_t)); result->lanes.push_back(b->lanes); return true; + } else if (const Shuffle *s = e.as(); s && s->vectors.size() == 1) { + if (s->is_transpose() && is_multiramp(s->vectors[0], scope, result)) { + return result->transpose(s->transpose_factor()); + } + // Any other shuffle of a single vector is a reshaping of it if the lane + // indices are themselves a multiramp, but we can only say what the + // result is if the values being shuffled are an affine function of the + // lane index, i.e. the input is one-dimensional. This is the shape that + // flatten_nested_ramps leaves a strided load in. + MultiRamp inner, perm; + if (is_multiramp(s->vectors[0], scope, &inner) && + inner.dimensions() == 1 && + multiramp_of_constants(s->indices, inner.base.type(), &perm)) { + perm.mul(inner.strides[0]); + perm.base = simplify(perm.base + inner.base); + *result = perm; + return true; + } + return false; } else if (const Ramp *r = e.as()) { if (auto stride = unbroadcast(r->stride)) { if (is_multiramp(r->base, scope, result)) { @@ -478,6 +543,7 @@ bool is_multiramp_impl(const Expr &e, const Scope &scope, MultiRamp *resul } } } + return false; } } // namespace @@ -493,6 +559,114 @@ bool is_multiramp(const Expr &e, const Scope &scope, MultiRamp *result) { return false; } +namespace { + +// Strip the casts, broadcasts and lane permutations off a load, moving the +// ones that rearrange lanes onto a copy of the load's index, where +// is_multiramp can make sense of them. A broadcast of a load is a load of a +// broadcast of the index, and likewise for a lane permutation. +const Load *peel_load(const Expr &e, Expr *index) { + if (const Cast *cast = e.as()) { + return peel_load(cast->value, index); + } else if (const Broadcast *broadcast = e.as()) { + const Load *load = peel_load(broadcast->value, index); + if (load) { + *index = Broadcast::make(*index, broadcast->lanes); + } + return load; + } else if (const Shuffle *shuffle = e.as()) { + if (shuffle->vectors.size() != 1) { + return nullptr; + } + const Load *load = peel_load(shuffle->vectors[0], index); + if (load) { + // Shuffling the values loaded is the same as shuffling the + // addresses loaded from. + *index = Shuffle::make({*index}, shuffle->indices); + } + return load; + } else if (const Load *load = e.as()) { + *index = load->index; + return load; + } + return nullptr; +} + +} // namespace + +int get_subtile(const Expr &index, const std::string &description, + std::vector *subtiles) { + MultiRamp mr; + if (!is_multiramp(index, Scope::empty_scope(), &mr)) { + user_error << "Access to " << description << " not affine: " << index << "\n"; + } + if (!can_prove(mr.alias_free())) { + user_error << "Access to " << description << " may have duplicated lanes: " + << index << "\n"; + } + if (subtiles->empty()) { + subtiles->push_back(std::move(mr)); + return 0; + } + + // All strides and lanes must match across all subtiles, or we give up. + const MultiRamp &first = (*subtiles)[0]; + if (mr.dimensions() != first.dimensions()) { + user_error << "Access to " << description << " does not have the same shape as " + << "other accesses to the same memory."; + return -1; + } + for (int i = 0; i < first.dimensions(); i++) { + if (!can_prove(mr.strides[i] == first.strides[i]) || + mr.lanes[i] != first.lanes[i]) { + user_error << "Access to " << description << " has different size and strides " + << "to other accesses to the same memory. All accesses must have " + << "the same subtile size and strides: " << index; + } + } + + // Now check for disjointedness. Add a synthetic dimension, the purpose of + // which will become clear. + mr.strides.emplace_back(); + mr.lanes.push_back(2); + for (int i = 0; i < (int)subtiles->size(); i++) { + const MultiRamp &other = (*subtiles)[i]; + // One of two things must be true: + // 1) All of the lanes of mr equal the corresponding lane of other. + // We've already checked the strides and lanes, so it's just a matter of + // checking the base. + if (can_prove(mr.base == other.base)) { + return i; + } + + // 2) None of the lanes of mr equal any of the lanes of other. To do + // this we construct a combined mr that can be either 'mr' or 'other', + // and ask if it's alias-free. This is what the synthetic dimension was + // for. + mr.strides.back() = mr.base - other.base; + if (!can_prove(mr.alias_free())) { + user_error << "Failed to prove access to " << description << " does not " + << "partially overlap another distinct access: " << index; + return -1; + } + } + + // Didn't already exist and didn't alias with anything. + mr.strides.pop_back(); + mr.lanes.pop_back(); + subtiles->push_back(std::move(mr)); + return (int)subtiles->size() - 1; +} + +const Load *is_load_of_multiramp(const Expr &e, const Scope &scope, MultiRamp *result) { + Expr index; + const Load *load = peel_load(e, &index); + if (load && is_multiramp(index, scope, result)) { + return load; + } + return nullptr; +} + Expr MultiRamp::operator==(const MultiRamp &other) const { // Construct the difference, and check if all strides are zero. MultiRamp diff = other; @@ -577,6 +751,50 @@ std::vector MultiRamp::alias_free_slice() { return peeled; } +bool MultiRamp::transpose(int cols) { + // Refine the dims so that some prefix of them accounts for exactly `cols` + // lanes, then move that prefix outwards. + std::vector shape; + size_t prefix = 0; + int prefix_lanes = 1; + for (int l : lanes) { + if (prefix_lanes == cols) { + shape.push_back(l); + continue; + } + if (cols % prefix_lanes) { + return false; + } + int remaining = cols / prefix_lanes; + if (l <= remaining) { + if (remaining % l) { + return false; + } + shape.push_back(l); + prefix_lanes *= l; + } else { + if (l % remaining) { + return false; + } + // This dim spans the split, so break it in two. + shape.push_back(remaining); + shape.push_back(l / remaining); + prefix_lanes = cols; + } + prefix++; + } + + std::vector new_strides; + if (prefix_lanes != cols || !strides_for_shape(shape, &new_strides)) { + return false; + } + + std::rotate(shape.begin(), shape.begin() + prefix, shape.end()); + std::rotate(new_strides.begin(), new_strides.begin() + prefix, new_strides.end()); + *this = MultiRamp(base, new_strides, shape); + return true; +} + int MultiRamp::rotate_stride_one_innermost() { int k = -1; for (int i = 0; i < dimensions(); i++) { diff --git a/src/MultiRamp.h b/src/MultiRamp.h index 583d257d93f4..74949d947008 100644 --- a/src/MultiRamp.h +++ b/src/MultiRamp.h @@ -14,6 +14,7 @@ namespace Internal { class IRMutator; class IRVisitor; +struct Load; /** A multi-dimensional ramp. I.e. a ramp of ramps of ramps of ramps... * @@ -151,6 +152,13 @@ struct MultiRamp { * a vector in the old lane order from one in the new order. */ int rotate_stride_one_innermost(); + /** Rearrange the lanes the way Shuffle::make_transpose(e, cols) does: + * view them as a row-major matrix with `cols` columns and transpose it, + * which moves the innermost `cols` lanes to the outside. Returns false, + * leaving *this undefined, if the dims can't be refactored to split off a + * prefix of exactly `cols` lanes. */ + bool transpose(int cols); + /** The dimensionality. May be lower than you expected, because this * gets flattened when possible by the operations above. */ int dimensions() const; @@ -210,10 +218,32 @@ struct MultiRamp { const std::vector &pos) const; }; +/** Locate the subtile of a tile-memory allocation that an access refers to. + * Such an allocation may hold several tile registers as disjoint sub-tiles, and + * each one becomes its own allocation, so the accesses have to be partitioned + * between them. + * + * `subtiles` accumulates the distinct subtiles found so far. Returns the index + * within it of the one `index` refers to, appending a new one if this is the + * first access to it. Every subtile must have the same shape, and any two must + * be either identical or disjoint - a partial overlap can't be expressed as + * separate tile registers. Anything else is a user error, reported in terms of + * `description`, which should name the kind of memory (e.g. "AMX tile"). */ +int get_subtile(const Expr &index, const std::string &description, + std::vector *subtiles); + /** Check if a vector Expr is a multiramp, and assign to result if so. * Returns false and leaves *result untouched if not. */ bool is_multiramp(const Expr &e, const Scope &scope, MultiRamp *result); +/** Check if an Expr is a load of a multiramp, possibly wrapped in casts, in + * broadcasts over dimensions the load doesn't depend on, and in the lane + * permutations the simplifier introduces when it rewrites a strided load as a + * dense load followed by a transpose. Returns the Load node underneath, with + * *result describing the addresses it loads in the lane order of `e`, or null + * if `e` isn't of that form (in which case *result is untouched). */ +const Load *is_load_of_multiramp(const Expr &e, const Scope &scope, MultiRamp *result); + } // namespace Internal } // namespace Halide diff --git a/test/correctness/multiramp.cpp b/test/correctness/multiramp.cpp index c7b8aae195db..af94c2b413c5 100644 --- a/test/correctness/multiramp.cpp +++ b/test/correctness/multiramp.cpp @@ -1,6 +1,7 @@ #include "Halide.h" #include +#include #include #include @@ -583,6 +584,124 @@ void check_reject_non_multiramp_sum() { CHECK(!is_multiramp(sum, scope, &m), "reject coprime-shape add"); } +// ---- MultiRamp::transpose ------------------------------------------------ + +// Mirror Shuffle::make_transpose: view v as a row-major matrix with `cols` +// columns and transpose it. +std::vector transpose_vec(const std::vector &v, int cols) { + int rows = (int)v.size() / cols; + std::vector result(v.size()); + for (int j = 0; j < cols; j++) { + for (int i = 0; i < rows; i++) { + result[j * rows + i] = v[i * cols + j]; + } + } + return result; +} + +void check_transpose_on_dim_boundary() { + // cols matches the innermost dim exactly, so no dim needs splitting. + MultiRamp A{0, {1, 100}, {4, 3}}; + auto want = transpose_vec(expand(A), 4); + CHECK(A.transpose(4), "transpose on a dim boundary"); + CHECK_SEQ(expand(A), want, "transpose on a dim boundary values"); +} + +void check_transpose_splits_a_dim() { + // cols lands in the middle of the only dim, which must be broken in two. + MultiRamp A{0, {1}, {8}}; + auto want = transpose_vec(expand(A), 4); + CHECK(A.transpose(4), "transpose splitting a dim"); + CHECK_SEQ(expand(A), want, "transpose splitting a dim values"); +} + +void check_transpose_spans_dims() { + // cols covers the two innermost dims together. + MultiRamp A{0, {1, 10, 100}, {2, 2, 3}}; + auto want = transpose_vec(expand(A), 4); + CHECK(A.transpose(4), "transpose spanning dims"); + CHECK_SEQ(expand(A), want, "transpose spanning dims values"); +} + +void check_transpose_rejects_indivisible() { + // The innermost dim of 3 can't be cut to leave a prefix of exactly 2 lanes. + MultiRamp A{0, {1, 100}, {3, 4}}; + CHECK(!A.transpose(2), "transpose with an indivisible split rejected"); +} + +// ---- Shuffles ------------------------------------------------------------ + +void check_recognize_transpose_shuffle() { + Expr e = Shuffle::make_transpose(Ramp::make(Expr(0), Expr(1), 12), 4); + Scope scope; + MultiRamp m; + CHECK(is_multiramp(e, scope, &m), "recognize a transpose shuffle"); + std::vector in(12); + std::iota(in.begin(), in.end(), 0); + CHECK_SEQ(expand(m), transpose_vec(in, 4), "transpose shuffle values"); +} + +void check_recognize_reshaping_shuffle() { + // A shuffle of a 1D ramp whose lane indices are themselves a multiramp is + // a reshaping of the ramp rather than an arbitrary gather. Indices + // [0,2,4,6,1,3,5,7] have shape (4,2) and strides (2,1). + Expr r = Ramp::make(Expr(0), Expr(3), 8); + Expr e = Shuffle::make({r}, {0, 2, 4, 6, 1, 3, 5, 7}); + Scope scope; + MultiRamp m; + CHECK(is_multiramp(e, scope, &m), "recognize a reshaping shuffle"); + CHECK_SEQ_LIT(expand(m), "reshaping shuffle values", + 0, 6, 12, 18, 3, 9, 15, 21); +} + +void check_reject_gather_shuffle() { + // Lane indices that aren't a multiramp really are a gather. + Expr r = Ramp::make(Expr(0), Expr(3), 8); + Expr e = Shuffle::make({r}, {0, 3, 1, 7, 2, 5, 4, 6}); + Scope scope; + MultiRamp m; + CHECK(!is_multiramp(e, scope, &m), "reject a gather shuffle"); +} + +// ---- is_load_of_multiramp ------------------------------------------------ + +Expr make_test_load(int lanes) { + return Load::make(Int(16, lanes), "buf", Ramp::make(Expr(0), Expr(1), lanes), + Buffer<>(), Parameter(), const_true(lanes), ModulusRemainder()); +} + +void check_load_under_cast_and_broadcast() { + // A broadcast of a load is a load at a stride-zero outer dim. + Expr e = Broadcast::make(Cast::make(Int(32, 4), make_test_load(4)), 3); + Scope scope; + MultiRamp m; + const Load *load = is_load_of_multiramp(e, scope, &m); + CHECK(load && load->name == "buf", "see through cast and broadcast"); + if (load) { + CHECK_SEQ_LIT(expand(m), "broadcast load addresses", + 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3); + } +} + +void check_load_under_shuffle() { + // Permuting the loaded values permutes the addresses loaded from. + Expr e = Shuffle::make_transpose(make_test_load(4), 2); + Scope scope; + MultiRamp m; + const Load *load = is_load_of_multiramp(e, scope, &m); + CHECK(load && load->name == "buf", "see through a lane permutation"); + if (load) { + CHECK_SEQ_LIT(expand(m), "permuted load addresses", 0, 2, 1, 3); + } +} + +void check_reject_non_load() { + Scope scope; + MultiRamp m; + CHECK(!is_load_of_multiramp(Ramp::make(Expr(0), Expr(1), 4), scope, &m), + "reject an Expr with no load underneath"); +} + } // namespace int main(int argc, char **argv) { @@ -631,6 +750,19 @@ int main(int argc, char **argv) { check_roundtrips(); check_reject_non_multiramp_sum(); + check_transpose_on_dim_boundary(); + check_transpose_splits_a_dim(); + check_transpose_spans_dims(); + check_transpose_rejects_indivisible(); + + check_recognize_transpose_shuffle(); + check_recognize_reshaping_shuffle(); + check_reject_gather_shuffle(); + + check_load_under_cast_and_broadcast(); + check_load_under_shuffle(); + check_reject_non_load(); + if (failures) { printf("%d failures\n", failures); return 1; diff --git a/test/correctness/tiled_matmul_errors.cpp b/test/correctness/tiled_matmul_errors.cpp index fd2e85a3f041..d60292eaec77 100644 --- a/test/correctness/tiled_matmul_errors.cpp +++ b/test/correctness/tiled_matmul_errors.cpp @@ -113,7 +113,7 @@ void scenario_naive_rhs() { // A gather-style matmul with an indirect row index — natural for sparse / // pruned matmul, indirect attention, embedding lookups. The LHS load index // goes through a table lookup, so the multiramp lift fails. Triggers the -// "loads indices are not affine" path in convert_to_matmul. +// "not loads with affine indices" path in convert_to_matmul. void scenario_indirect() { Buffer A(64, 64); Buffer B(4, 64, 16); From 435bf78621ad51f8c90da4dfddd1aa81c0aa12f1 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Thu, 30 Jul 2026 13:30:27 -0700 Subject: [PATCH 06/59] Don't let a cast hide the type AMX multiplies at is_load_of_multiramp stripped any number of casts off a load and returned only the Load underneath, so a caller had no way to see that the values had been cast. ExtractTileOperations then typed its tile registers from the load, and the signedness of those types picks which of the four integer tdpb instructions runs. A pipeline multiplying uint8 buffers reinterpreted as int8 compiled to tdpbuud - an unsigned multiply for a signed algorithm. Peel at most one cast, so the element type of the original Expr and the type of the returned Load together say whether the values were cast and to what, and have the AMX integer path reject a mismatch. The float path still reads the load's type, which is what it wants: the bf16 to f32 widening cast is expected there, and tile_load takes the bf16. Co-Authored-By: Claude Opus 5 --- src/ExtractTileOperations.cpp | 11 ++++++++++- src/MultiRamp.cpp | 19 +++++++++++++------ src/MultiRamp.h | 9 +++++++-- test/correctness/tiled_matmul_errors.cpp | 21 +++++++++++++++++++++ 4 files changed, 51 insertions(+), 9 deletions(-) diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp index d835fde34568..47d19f13de07 100644 --- a/src/ExtractTileOperations.cpp +++ b/src/ExtractTileOperations.cpp @@ -135,7 +135,8 @@ Matmul convert_to_matmul(const Store *op, const string &new_name) { // Underneath all of this must be a load, though it may be wrapped in a // broadcast over the dimension it doesn't depend on, in the widening cast // (for floats - the integer branch above already extracted the cast inputs - // from the widening_mul intrinsic), and in a lane permutation. + // from the widening_mul intrinsic), and in a lane permutation. The cast is + // checked against the load's type below. // TODO: What if we want to multiply by the same matrix multiple times? It might be a let binding. MultiRamp lhs_mr, rhs_mr; Scope empty_scope; @@ -150,6 +151,14 @@ Matmul convert_to_matmul(const Store *op, const string &new_name) { } if (reduce->type.is_int_or_uint()) { + // The tile registers are typed by what was loaded, and the signedness + // of those types picks which of the four integer tdpb instructions + // gets used, so a cast between the load and the multiply would change + // the meaning of the multiply without changing the instruction. + if (lhs.type().element_of() != lhs_load->type.element_of() || + rhs.type().element_of() != rhs_load->type.element_of()) { + return fail("the matrix multiply operands are cast after being loaded"); + } if (lhs_load->type.bits() != 8 || rhs_load->type.bits() != 8) { return fail("the vector reduction operand or result types are not supported"); } diff --git a/src/MultiRamp.cpp b/src/MultiRamp.cpp index 695032f7af87..715b995cdf3c 100644 --- a/src/MultiRamp.cpp +++ b/src/MultiRamp.cpp @@ -561,15 +561,22 @@ bool is_multiramp(const Expr &e, const Scope &scope, MultiRamp *result) { namespace { -// Strip the casts, broadcasts and lane permutations off a load, moving the +// Strip the cast, broadcasts and lane permutations off a load, moving the // ones that rearrange lanes onto a copy of the load's index, where // is_multiramp can make sense of them. A broadcast of a load is a load of a // broadcast of the index, and likewise for a lane permutation. -const Load *peel_load(const Expr &e, Expr *index) { +// +// At most one cast is peeled, so that comparing the element type of the +// original Expr against the type of the Load tells the caller whether the +// values were cast, and to what. +const Load *peel_load(const Expr &e, Expr *index, bool cast_allowed) { if (const Cast *cast = e.as()) { - return peel_load(cast->value, index); + if (!cast_allowed) { + return nullptr; + } + return peel_load(cast->value, index, false); } else if (const Broadcast *broadcast = e.as()) { - const Load *load = peel_load(broadcast->value, index); + const Load *load = peel_load(broadcast->value, index, cast_allowed); if (load) { *index = Broadcast::make(*index, broadcast->lanes); } @@ -578,7 +585,7 @@ const Load *peel_load(const Expr &e, Expr *index) { if (shuffle->vectors.size() != 1) { return nullptr; } - const Load *load = peel_load(shuffle->vectors[0], index); + const Load *load = peel_load(shuffle->vectors[0], index, cast_allowed); if (load) { // Shuffling the values loaded is the same as shuffling the // addresses loaded from. @@ -660,7 +667,7 @@ int get_subtile(const Expr &index, const std::string &description, const Load *is_load_of_multiramp(const Expr &e, const Scope &scope, MultiRamp *result) { Expr index; - const Load *load = peel_load(e, &index); + const Load *load = peel_load(e, &index, true); if (load && is_multiramp(index, scope, result)) { return load; } diff --git a/src/MultiRamp.h b/src/MultiRamp.h index 74949d947008..befd2c130c3f 100644 --- a/src/MultiRamp.h +++ b/src/MultiRamp.h @@ -236,12 +236,17 @@ int get_subtile(const Expr &index, const std::string &description, * Returns false and leaves *result untouched if not. */ bool is_multiramp(const Expr &e, const Scope &scope, MultiRamp *result); -/** Check if an Expr is a load of a multiramp, possibly wrapped in casts, in +/** Check if an Expr is a load of a multiramp, possibly wrapped in a cast, in * broadcasts over dimensions the load doesn't depend on, and in the lane * permutations the simplifier introduces when it rewrites a strided load as a * dense load followed by a transpose. Returns the Load node underneath, with * *result describing the addresses it loads in the lane order of `e`, or null - * if `e` isn't of that form (in which case *result is untouched). */ + * if `e` isn't of that form (in which case *result is untouched). + * + * At most one cast is seen through, so `e.type().element_of()` and the + * returned Load's type together tell you whether the loaded values were cast, + * and to what. Callers that care about the type the values actually have - + * including its signedness - must use the former, not the Load's type. */ const Load *is_load_of_multiramp(const Expr &e, const Scope &scope, MultiRamp *result); } // namespace Internal diff --git a/test/correctness/tiled_matmul_errors.cpp b/test/correctness/tiled_matmul_errors.cpp index d60292eaec77..4cf486df291f 100644 --- a/test/correctness/tiled_matmul_errors.cpp +++ b/test/correctness/tiled_matmul_errors.cpp @@ -130,6 +130,26 @@ void scenario_indirect() { mm.in().compile_jit(amx_target); } +// A matmul over uint8 buffers whose values are reinterpreted as int8 before +// multiplying. AMX types its tile registers by what was loaded, and their +// signedness picks which of the four integer tdpb instructions runs, so +// accepting this would run an unsigned multiply for a signed algorithm. +// Triggers the "cast after being loaded" path in convert_to_matmul. +void scenario_sign_changing_cast() { + Buffer A(64, 64); + Buffer B(4, 64, 16); + Var x("x"), y("y"); + RDom r(0, 64, "r"); + + Func mm("matmul_sign_cast"); + mm(x, y) = cast(0); + mm(x, y) += + cast(cast(A(r, y))) * + cast(cast(B(r % 4, x, r / 4))); + schedule_matmul(mm, r.x, 8, 8, 8); + mm.in().compile_jit(amx_target); +} + // A 1D convolution of a 2D signal with per-row kernels, aggressively // vectorized. Structurally a sum-of-widening-multiplies with a contiguous // inner K, but the LHS depends on x, k, and y simultaneously (no broadcast @@ -307,6 +327,7 @@ int main(int argc, char **argv) { failures += !expect_user_error("bad_result_type", scenario_bad_result_type); failures += !expect_user_error("naive_rhs", scenario_naive_rhs); failures += !expect_user_error("indirect", scenario_indirect); + failures += !expect_user_error("sign_changing_cast", scenario_sign_changing_cast); failures += !expect_user_error("conv1d", scenario_conv1d); failures += !expect_user_error("no_matmul", scenario_no_matmul); failures += !expect_user_error("widening_16bit", scenario_widening_16bit); From 7d3307bea4bc6e0449453accd174b7f329bbb1e3 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Thu, 30 Jul 2026 13:33:32 -0700 Subject: [PATCH 07/59] Recognize a transposing shuffle as just another reshaping shuffle A transpose's shuffle mask is itself a multiramp of constants - transposing n lanes with c columns gives a mask of shape (n/c, c) and strides (c, 1) - so the general reshaping-shuffle case in is_multiramp already covers it. Drop the separate is_transpose case and MultiRamp::transpose, which existed only to serve it and had no other caller. That leaves one rule for shuffles of a single vector: it is a reshaping rather than a gather if the mask is a multiramp of constants, and we can say what it reshapes to when the shuffled vector is one-dimensional. Co-Authored-By: Claude Opus 5 --- src/MultiRamp.cpp | 58 ++++------------------------------ src/MultiRamp.h | 7 ---- test/correctness/multiramp.cpp | 39 +---------------------- 3 files changed, 7 insertions(+), 97 deletions(-) diff --git a/src/MultiRamp.cpp b/src/MultiRamp.cpp index 715b995cdf3c..d9281d87ee40 100644 --- a/src/MultiRamp.cpp +++ b/src/MultiRamp.cpp @@ -476,14 +476,12 @@ bool is_multiramp_impl(const Expr &e, const Scope &scope, MultiRamp *resul result->lanes.push_back(b->lanes); return true; } else if (const Shuffle *s = e.as(); s && s->vectors.size() == 1) { - if (s->is_transpose() && is_multiramp(s->vectors[0], scope, result)) { - return result->transpose(s->transpose_factor()); - } - // Any other shuffle of a single vector is a reshaping of it if the lane - // indices are themselves a multiramp, but we can only say what the - // result is if the values being shuffled are an affine function of the - // lane index, i.e. the input is one-dimensional. This is the shape that - // flatten_nested_ramps leaves a strided load in. + // A shuffle of a single vector is a reshaping of it, rather than a + // gather, if the lane indices are themselves a multiramp. That covers + // transposes, whose masks are multiramps of constants. But we can only + // say what the result is if the values being shuffled are an affine + // function of the lane index, i.e. the input is one-dimensional. This + // is the shape that flatten_nested_ramps leaves a strided load in. MultiRamp inner, perm; if (is_multiramp(s->vectors[0], scope, &inner) && inner.dimensions() == 1 && @@ -758,50 +756,6 @@ std::vector MultiRamp::alias_free_slice() { return peeled; } -bool MultiRamp::transpose(int cols) { - // Refine the dims so that some prefix of them accounts for exactly `cols` - // lanes, then move that prefix outwards. - std::vector shape; - size_t prefix = 0; - int prefix_lanes = 1; - for (int l : lanes) { - if (prefix_lanes == cols) { - shape.push_back(l); - continue; - } - if (cols % prefix_lanes) { - return false; - } - int remaining = cols / prefix_lanes; - if (l <= remaining) { - if (remaining % l) { - return false; - } - shape.push_back(l); - prefix_lanes *= l; - } else { - if (l % remaining) { - return false; - } - // This dim spans the split, so break it in two. - shape.push_back(remaining); - shape.push_back(l / remaining); - prefix_lanes = cols; - } - prefix++; - } - - std::vector new_strides; - if (prefix_lanes != cols || !strides_for_shape(shape, &new_strides)) { - return false; - } - - std::rotate(shape.begin(), shape.begin() + prefix, shape.end()); - std::rotate(new_strides.begin(), new_strides.begin() + prefix, new_strides.end()); - *this = MultiRamp(base, new_strides, shape); - return true; -} - int MultiRamp::rotate_stride_one_innermost() { int k = -1; for (int i = 0; i < dimensions(); i++) { diff --git a/src/MultiRamp.h b/src/MultiRamp.h index befd2c130c3f..94653491255a 100644 --- a/src/MultiRamp.h +++ b/src/MultiRamp.h @@ -152,13 +152,6 @@ struct MultiRamp { * a vector in the old lane order from one in the new order. */ int rotate_stride_one_innermost(); - /** Rearrange the lanes the way Shuffle::make_transpose(e, cols) does: - * view them as a row-major matrix with `cols` columns and transpose it, - * which moves the innermost `cols` lanes to the outside. Returns false, - * leaving *this undefined, if the dims can't be refactored to split off a - * prefix of exactly `cols` lanes. */ - bool transpose(int cols); - /** The dimensionality. May be lower than you expected, because this * gets flattened when possible by the operations above. */ int dimensions() const; diff --git a/test/correctness/multiramp.cpp b/test/correctness/multiramp.cpp index af94c2b413c5..099ed1eaf56d 100644 --- a/test/correctness/multiramp.cpp +++ b/test/correctness/multiramp.cpp @@ -584,7 +584,7 @@ void check_reject_non_multiramp_sum() { CHECK(!is_multiramp(sum, scope, &m), "reject coprime-shape add"); } -// ---- MultiRamp::transpose ------------------------------------------------ +// ---- Shuffles ------------------------------------------------------------ // Mirror Shuffle::make_transpose: view v as a row-major matrix with `cols` // columns and transpose it. @@ -599,38 +599,6 @@ std::vector transpose_vec(const std::vector &v, int cols) { return result; } -void check_transpose_on_dim_boundary() { - // cols matches the innermost dim exactly, so no dim needs splitting. - MultiRamp A{0, {1, 100}, {4, 3}}; - auto want = transpose_vec(expand(A), 4); - CHECK(A.transpose(4), "transpose on a dim boundary"); - CHECK_SEQ(expand(A), want, "transpose on a dim boundary values"); -} - -void check_transpose_splits_a_dim() { - // cols lands in the middle of the only dim, which must be broken in two. - MultiRamp A{0, {1}, {8}}; - auto want = transpose_vec(expand(A), 4); - CHECK(A.transpose(4), "transpose splitting a dim"); - CHECK_SEQ(expand(A), want, "transpose splitting a dim values"); -} - -void check_transpose_spans_dims() { - // cols covers the two innermost dims together. - MultiRamp A{0, {1, 10, 100}, {2, 2, 3}}; - auto want = transpose_vec(expand(A), 4); - CHECK(A.transpose(4), "transpose spanning dims"); - CHECK_SEQ(expand(A), want, "transpose spanning dims values"); -} - -void check_transpose_rejects_indivisible() { - // The innermost dim of 3 can't be cut to leave a prefix of exactly 2 lanes. - MultiRamp A{0, {1, 100}, {3, 4}}; - CHECK(!A.transpose(2), "transpose with an indivisible split rejected"); -} - -// ---- Shuffles ------------------------------------------------------------ - void check_recognize_transpose_shuffle() { Expr e = Shuffle::make_transpose(Ramp::make(Expr(0), Expr(1), 12), 4); Scope scope; @@ -750,11 +718,6 @@ int main(int argc, char **argv) { check_roundtrips(); check_reject_non_multiramp_sum(); - check_transpose_on_dim_boundary(); - check_transpose_splits_a_dim(); - check_transpose_spans_dims(); - check_transpose_rejects_indivisible(); - check_recognize_transpose_shuffle(); check_recognize_reshaping_shuffle(); check_reject_gather_shuffle(); From 8e26fb5ccfeef004b9f7130b9f98a6294dba0c49 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Thu, 30 Jul 2026 14:42:05 -0700 Subject: [PATCH 08/59] Simplify division of a nested vector by a broadcast Dividing a vector by a broadcast only folds when the two are vectorized the same way, which they aren't when the numerator is a nested vector. Add rules for the two shapes that come up: a broadcast numerator, where the division can be pushed inwards until the two line up, and a ramp whose lanes are each repeated by an inner broadcast, where repeating a lane doesn't change the set of values so the existing first-and-last-lane argument still applies. While here, replace the can_prove predicate on the rule being generalized with a structural one. Matching the base as a multiple of the denominator is enough to know the quotient is uniform, given the ramp doesn't span far enough to reach the next multiple. The structural form also generalizes it: the stride may be any non-negative constant, the base's multiplier need only be a multiple of the denominator rather than equal to it, and the base may be affine rather than linear. Co-Authored-By: Claude Opus 5 --- src/Simplify_Div.cpp | 31 ++++++++++++++++++++++++++++--- test/correctness/simplify.cpp | 20 ++++++++++++++++++++ 2 files changed, 48 insertions(+), 3 deletions(-) diff --git a/src/Simplify_Div.cpp b/src/Simplify_Div.cpp index 159e7bab802f..4098f6f027e7 100644 --- a/src/Simplify_Div.cpp +++ b/src/Simplify_Div.cpp @@ -204,9 +204,34 @@ Expr Simplify::visit(const Div *op, ExprInfo *info) { (op->type.is_float() && rewrite(x / c0, x * fold(1 / c0))))) || (no_overflow_int(op->type) && (rewrite(ramp(x, c0, lanes) / broadcast(c1, lanes), ramp(x / c1, fold(c0 / c1), lanes), (c0 % c1 == 0)) || - rewrite(ramp(x, c0, lanes) / broadcast(c1, lanes), broadcast(x / c1, lanes), - // First and last lanes are the same when... - can_prove((x % c1 + c0 * (lanes - 1)) / c1 == 0, this)))) || + // Every lane gives the same quotient when the base is a multiple of + // the denominator and the ramp doesn't span far enough to reach the + // next one. In the affine case the offset just has to leave room + // within its own multiple for the rest of the ramp. + rewrite(ramp(x * c1, c0, lanes) / broadcast(c2, lanes), + broadcast(x * fold(c1 / c2), lanes), + c2 > 0 && c1 % c2 == 0 && c0 >= 0 && c0 * (lanes - 1) < c2) || + rewrite(ramp(x * c1 + c5, c0, lanes) / broadcast(c2, lanes), + broadcast(x * fold(c1 / c2) + fold(c5 / c2), lanes), + c2 > 0 && c1 % c2 == 0 && c0 >= 0 && + c5 % c2 + c0 * (lanes - 1) < c2) || + // The rules above require the numerator and the denominator to be + // vectorized the same way, which they aren't when the numerator is + // a nested vector. Push the division inwards to line them up. + rewrite(broadcast(x, c0) / broadcast(y, c1), + broadcast(x / broadcast(y, fold(c1 / c0)), c0), + c0 < c1 && c1 % c0 == 0) || + // The same argument, for a ramp whose lanes are each repeated by an + // inner broadcast. Repeating a lane doesn't change the set of + // values, so the span condition is unchanged. + rewrite(ramp(broadcast(x * c1, c3), broadcast(c0, c3), c4) / broadcast(c2, lanes), + broadcast(x * fold(c1 / c2), lanes), + c2 > 0 && c1 % c2 == 0 && c0 >= 0 && c0 * (c4 - 1) < c2 && + c3 * c4 == lanes) || + rewrite(ramp(broadcast(x * c1 + c5, c3), broadcast(c0, c3), c4) / broadcast(c2, lanes), + broadcast(x * fold(c1 / c2) + fold(c5 / c2), lanes), + c2 > 0 && c1 % c2 == 0 && c0 >= 0 && + c5 % c2 + c0 * (c4 - 1) < c2 && c3 * c4 == lanes))) || (no_overflow_scalar_int(op->type) && (rewrite(x / -1, -x) || (denominator_non_zero && rewrite(c0 / y, select(y < 0, fold(-c0), c0), c0 == -1)) || diff --git a/test/correctness/simplify.cpp b/test/correctness/simplify.cpp index 42dbdb480961..5b9b36766bf2 100644 --- a/test/correctness/simplify.cpp +++ b/test/correctness/simplify.cpp @@ -606,6 +606,26 @@ void check_vectors() { check(ramp(ramp(cast(x), cast(-1), 4), cast(UInt(8, 4), -4), 3), ramp(cast(x), cast(-1), 12)); + // Dividing a nested vector by a broadcast should give the same answer as + // dividing the equivalent flat one, even though the lanes are laid out + // differently. + check(ramp(broadcast(x * 16, 16), broadcast(1, 16), 16) / broadcast(16, 256), + broadcast(x, 256)); + check(broadcast(ramp(broadcast(x * 16, 16), broadcast(1, 16), 16), 16) / broadcast(16, 4096), + broadcast(x, 4096)); + check(broadcast(ramp(x, 1, 4), 2) / broadcast(y, 8), + broadcast(ramp(x, 1, 4) / broadcast(y, 4), 2)); + + // An affine base works too, as long as the offset leaves room within its + // own multiple of the denominator for the rest of the ramp. + check(ramp(x * 16 + 4, 1, 4) / broadcast(16, 4), broadcast(x, 4)); + check(ramp(x * 16 - 4, 1, 4) / broadcast(16, 4), broadcast(x + (-1), 4)); + check(ramp(broadcast(x * 16 + 4, 16), broadcast(1, 16), 4) / broadcast(16, 64), + broadcast(x, 64)); + // ... but not when it spills over into the next one. + check(ramp(x * 16 + 14, 1, 4) / broadcast(16, 4), + ramp(x * 16 + 14, 1, 4) / broadcast(16, 4)); + // Any linear combination of simple ramps and broadcasts should // reduce to a single ramp or broadcast. std::mt19937 rng(0); From 5d72da5f2a4aef1a989e4fb9683c1528f43d249b Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Wed, 29 Jul 2026 00:21:13 -0700 Subject: [PATCH 09/59] Use cp.async for global to shared copies in the PTX backend A store into shared memory whose value is a plain load from global memory is emitted as an asynchronous copy on sm_80 and later, which moves the data without routing it through registers. The copies issued in a producer are waited for at the end of it. This takes the shared-memory tensor core matmul from 31.4 to 39.5 TFlop/s at 2048^3 on an RTX 5060 Ti. Co-Authored-By: Claude Opus 5 --- src/CodeGen_PTX_Dev.cpp | 114 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 114 insertions(+) diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 3f026feedc80..9c84c89be2f6 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -80,6 +80,7 @@ class CodeGen_PTX_Dev : public CodeGen_LLVM, public CodeGen_GPU_Dev { void visit(const Load *) override; void visit(const Store *) override; void visit(const Atomic *) override; + void visit(const ProducerConsumer *) override; void codegen_vector_reduce(const VectorReduce *op, const Expr &init) override; // @} @@ -101,6 +102,21 @@ class CodeGen_PTX_Dev : public CodeGen_LLVM, public CodeGen_GPU_Dev { * intrinsic functions to call to get them. */ std::string simt_intrinsic(const std::string &name); + /** The memory type of each allocation made inside the kernel, so that + * copies into shared memory can be recognized. */ + Scope alloc_memory_type; + + /** Whether we're inside a producer node, which is where the wait for any + * asynchronous copies gets emitted, and whether any have been issued in + * it. */ + bool in_producer = false; + bool issued_async_copy = false; + + /** Try to emit a store into shared memory as an asynchronous copy, which + * moves the data straight from global memory without routing it through + * registers. Returns false if this store isn't one we can do that for. */ + bool codegen_async_copy(const Store *op); + bool supports_atomic_add(const Type &t) const override; }; @@ -328,6 +344,7 @@ void CodeGen_PTX_Dev::visit(const For *loop) { void CodeGen_PTX_Dev::visit(const Allocate *alloc) { user_assert(!alloc->new_expr.defined()) << "Allocate node inside PTX kernel has custom new expression.\n" << "(Memoization is not supported inside GPU kernels at present.)\n"; + ScopedBinding bind(alloc_memory_type, alloc->name, alloc->memory_type); if (alloc->memory_type == MemoryType::GPUShared) { // PTX uses zero in address space 3 as the base address for shared memory Value *shared_base = Constant::getNullValue(PointerType::get(*context, 3)); @@ -388,6 +405,99 @@ void CodeGen_PTX_Dev::visit(const Load *op) { CodeGen_LLVM::visit(op); } +// A copy from global memory into shared memory can be done by the hardware +// without going through registers, which saves the load, the store, and the +// registers in between. The copy is asynchronous, so it has to be waited for +// before the data is used; that happens at the end of the producer. +bool CodeGen_PTX_Dev::codegen_async_copy(const Store *op) { + // Asynchronous copies need something to wait for them, which only happens + // at the end of a producer. + if (!in_producer || emit_atomic_stores || + target.get_cuda_capability_lower_bound() < 80) { + return false; + } + + // The destination must be shared memory, and the source must be a plain + // load from something we didn't allocate in here, which is to say global + // memory. + const MemoryType *dst_memory_type = alloc_memory_type.find(op->name); + if (!dst_memory_type || *dst_memory_type != MemoryType::GPUShared) { + return false; + } + const Load *src = op->value.as(); + if (!src || alloc_memory_type.contains(src->name)) { + return false; + } + if (!is_const_one(op->predicate) || !is_const_one(src->predicate)) { + return false; + } + + // The hardware copies 4, 8 or 16 bytes at a time, from and to consecutive + // addresses. + const Type t = op->value.type(); + const int bytes = t.bytes() * t.lanes(); + if (!(bytes == 4 || bytes == 8 || bytes == 16)) { + return false; + } + Expr dst_base = op->index, src_base = src->index; + if (t.lanes() > 1) { + // Shared allocations are given an offset into one big block after the + // last simplification pass, so the indices need simplifying here. + dst_base = strided_ramp_base(simplify(op->index)); + src_base = strided_ramp_base(simplify(src->index)); + if (!dst_base.defined() || !src_base.defined()) { + return false; + } + } + Value *dst = codegen_buffer_pointer(op->name, t.element_of(), dst_base); + Value *src_ptr = codegen_buffer_pointer(src->name, t.element_of(), src_base); + + // Shared allocations are already in the shared address space. The source is + // a kernel argument, so it's global, but it comes in as a generic pointer. + llvm::Type *shared_ptr_t = PointerType::get(*context, 3); + llvm::Type *global_ptr_t = PointerType::get(*context, 1); + if (dst->getType() != shared_ptr_t) { + return false; + } + src_ptr = builder->CreateAddrSpaceCast(src_ptr, global_ptr_t); + + std::ostringstream name; + name << "llvm.nvvm.cp.async.ca.shared.global." << bytes; + llvm::Intrinsic::ID id = llvm::Intrinsic::lookupIntrinsicID(name.str()); + internal_assert(id != llvm::Intrinsic::not_intrinsic) + << "Could not find the nvvm intrinsic " << name.str() << "\n"; + llvm::Function *fn = llvm::Intrinsic::getOrInsertDeclaration(module.get(), id); + builder->CreateCall(fn, {dst, src_ptr}); + issued_async_copy = true; + return true; +} + +void CodeGen_PTX_Dev::visit(const ProducerConsumer *op) { + if (!op->is_producer) { + CodeGen_LLVM::visit(op); + return; + } + + ScopedValue old_issued(issued_async_copy, false); + ScopedValue old_in(in_producer, true); + codegen(op->body); + if (issued_async_copy) { + // Everything issued in here has to have landed before the values are + // used, which is after this producer. + for (const char *intrin : {"llvm.nvvm.cp.async.commit.group", + "llvm.nvvm.cp.async.wait.group"}) { + llvm::Intrinsic::ID id = llvm::Intrinsic::lookupIntrinsicID(intrin); + internal_assert(id != llvm::Intrinsic::not_intrinsic); + llvm::Function *fn = llvm::Intrinsic::getOrInsertDeclaration(module.get(), id); + vector args; + if (fn->getFunctionType()->getNumParams() == 1) { + args.push_back(ConstantInt::get(i32_t, 0)); + } + builder->CreateCall(fn, args); + } + } +} + void CodeGen_PTX_Dev::visit(const Store *op) { // Issue atomic store if we are inside an Atomic node. if (emit_atomic_stores) { @@ -395,6 +505,10 @@ void CodeGen_PTX_Dev::visit(const Store *op) { user_assert(op->value.type().bits() >= 32) << "CUDA: 8-bit or 16-bit atomics are not supported.\n"; } + if (codegen_async_copy(op)) { + return; + } + // Do aligned 4-wide 32-bit stores as a single i128 store. const Ramp *r = op->index.as(); // TODO: lanes >= 4, not lanes == 4 From 474331b05a3d7dad6b1fc7f6b997c4a50a5f9e15 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Wed, 29 Jul 2026 00:27:58 -0700 Subject: [PATCH 10/59] Sink redundant GPU thread barriers If nothing between a barrier and the next one at the same level loads what was stored before it, the later barrier can do the job of both. The fence types of the elided barrier are added to the one that remains. Co-Authored-By: Claude Opus 5 --- src/CodeGen_PTX_Dev.cpp | 39 ++++++++----- src/FuseGPUThreadLoops.cpp | 110 +++++++++++++++++++++++++++++++++++++ 2 files changed, 136 insertions(+), 13 deletions(-) diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 9c84c89be2f6..0c706db90813 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -117,6 +117,9 @@ class CodeGen_PTX_Dev : public CodeGen_LLVM, public CodeGen_GPU_Dev { * registers. Returns false if this store isn't one we can do that for. */ bool codegen_async_copy(const Store *op); + /** Wait for any asynchronous copies issued so far to have landed. */ + void wait_for_async_copies(); + bool supports_atomic_add(const Type &t) const override; }; @@ -287,6 +290,10 @@ void CodeGen_PTX_Dev::visit(const Call *op) { // arguments internal_assert(op->args.size() == 1) << "gpu_thread_barrier() intrinsic must specify memory fence type.\n"; + // A barrier tells other threads the shared memory this thread wrote is + // ready, so any asynchronous copies must have landed by now. + wait_for_async_copies(); + auto fence_type_ptr = as_const_int(op->args[0]); internal_assert(fence_type_ptr) << "gpu_thread_barrier() parameter is not a constant integer.\n"; @@ -481,21 +488,27 @@ void CodeGen_PTX_Dev::visit(const ProducerConsumer *op) { ScopedValue old_issued(issued_async_copy, false); ScopedValue old_in(in_producer, true); codegen(op->body); - if (issued_async_copy) { - // Everything issued in here has to have landed before the values are - // used, which is after this producer. - for (const char *intrin : {"llvm.nvvm.cp.async.commit.group", - "llvm.nvvm.cp.async.wait.group"}) { - llvm::Intrinsic::ID id = llvm::Intrinsic::lookupIntrinsicID(intrin); - internal_assert(id != llvm::Intrinsic::not_intrinsic); - llvm::Function *fn = llvm::Intrinsic::getOrInsertDeclaration(module.get(), id); - vector args; - if (fn->getFunctionType()->getNumParams() == 1) { - args.push_back(ConstantInt::get(i32_t, 0)); - } - builder->CreateCall(fn, args); + // Everything issued in here has to have landed before the values are used, + // which is after this producer. + wait_for_async_copies(); +} + +void CodeGen_PTX_Dev::wait_for_async_copies() { + if (!issued_async_copy) { + return; + } + for (const char *intrin : {"llvm.nvvm.cp.async.commit.group", + "llvm.nvvm.cp.async.wait.group"}) { + llvm::Intrinsic::ID id = llvm::Intrinsic::lookupIntrinsicID(intrin); + internal_assert(id != llvm::Intrinsic::not_intrinsic); + llvm::Function *fn = llvm::Intrinsic::getOrInsertDeclaration(module.get(), id); + vector args; + if (fn->getFunctionType()->getNumParams() == 1) { + args.push_back(ConstantInt::get(i32_t, 0)); } + builder->CreateCall(fn, args); } + issued_async_copy = false; } void CodeGen_PTX_Dev::visit(const Store *op) { diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index 4cd99492de5f..8aad78c18db4 100644 --- a/src/FuseGPUThreadLoops.cpp +++ b/src/FuseGPUThreadLoops.cpp @@ -1240,6 +1240,98 @@ class ExtractRegisterAllocations : public IRMutator { bool has_thread_loop = false; }; +// Gather the names loaded from before the first barrier at the top level of a +// statement. Barriers under a loop or an if don't count, because they don't +// necessarily run before the loads that follow. +class LoadsBeforeBarrier : public IRVisitor { + using IRVisitor::visit; + + bool at_top_level = true; + + void visit(const Block *op) override { + op->first.accept(this); + if (!found_barrier && op->rest.defined()) { + op->rest.accept(this); + } + } + + void visit(const For *op) override { + ScopedValue s(at_top_level, false); + IRVisitor::visit(op); + } + + void visit(const IfThenElse *op) override { + ScopedValue s(at_top_level, false); + IRVisitor::visit(op); + } + + void visit(const Evaluate *op) override { + const Call *c = op->value.as(); + if (at_top_level && c && c->is_intrinsic(Call::gpu_thread_barrier)) { + found_barrier = true; + } else { + IRVisitor::visit(op); + } + } + + void visit(const Load *op) override { + loads.insert(op->name); + IRVisitor::visit(op); + } + +public: + bool found_barrier = false; + std::set loads; +}; + +// Add fence types to the first barrier at the top level of a statement, so +// that it can stand in for one that would otherwise have preceded it. +class WidenFirstBarrier : public IRMutator { + using IRMutator::visit; + + bool at_top_level = true; + int mask; + + Stmt visit(const Block *op) override { + Stmt first = mutate(op->first); + if (done || !op->rest.defined()) { + return Block::make(first, op->rest); + } + return Block::make(first, mutate(op->rest)); + } + + Stmt visit(const For *op) override { + ScopedValue s(at_top_level, false); + return op; + } + + Stmt visit(const IfThenElse *op) override { + ScopedValue s(at_top_level, false); + return op; + } + + Stmt visit(const Evaluate *op) override { + const Call *c = op->value.as(); + if (at_top_level && !done && c && c->is_intrinsic(Call::gpu_thread_barrier)) { + done = true; + auto old_mask = as_const_int(c->args[0]); + internal_assert(old_mask); + return Evaluate::make(Call::make(Int(32), Call::gpu_thread_barrier, + {IntImm::make(Int(32), *old_mask | mask)}, + Call::Intrinsic)); + } + return op; + } + +public: + using IRMutator::mutate; + + bool done = false; + WidenFirstBarrier(int mask) + : mask(mask) { + } +}; + class InjectThreadBarriers : public IRMutator { protected: bool in_threads = false, injected_barrier; @@ -1380,6 +1472,24 @@ class InjectThreadBarriers : public IRMutator { break; } } + // If nothing in rest reads what first wrote until after a barrier + // of its own, that barrier can stand in for this one. + LoadsBeforeBarrier lbb; + rest.accept(&lbb); + bool needed = false; + for (const auto &st : shared_stores) { + needed |= lbb.loads.count(st) > 0; + } + for (const auto &st : device_stores) { + needed |= lbb.loads.count(st) > 0; + } + if (!needed && lbb.found_barrier) { + WidenFirstBarrier widen(mask); + rest = widen.mutate(rest); + internal_assert(widen.done); + injected_barrier = true; + return Block::make(first, rest); + } injected_barrier = true; return Block::make({first, make_barrier(mask), rest}); } else { From 2efd51198518c9086fcecceae6b8346e0f66c695 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Thu, 30 Jul 2026 14:52:21 -0700 Subject: [PATCH 11/59] Add MemoryType::GPUSharedAsync A Func stored in GPUSharedAsync goes in GPU shared memory, but is written by an asynchronous copy instruction that moves the data straight from global memory without routing it through registers. Asking for the memory type is a promise that every store to it is a copy the hardware can make that way, so a store that isn't reports what the copy engine requires and how a schedule usually satisfies it, rather than quietly falling back to a load and a store. Co-Authored-By: Claude Opus 5 --- src/CodeGen_D3D12Compute_Dev.cpp | 2 +- src/CodeGen_Metal_Dev.cpp | 4 +- src/CodeGen_OpenCL_Dev.cpp | 4 +- src/CodeGen_PTX_Dev.cpp | 103 +++++++++++++++++++++++++--- src/CodeGen_Vulkan_Dev.cpp | 2 +- src/CodeGen_WebGPU_Dev.cpp | 2 +- src/Deserialization.cpp | 2 + src/Expr.h | 13 ++++ src/FuseGPUThreadLoops.cpp | 6 +- src/IRPrinter.cpp | 3 + src/LowerWarpShuffles.cpp | 2 +- src/OffloadGPULoops.cpp | 2 +- src/Serialization.cpp | 2 + src/halide_ir.fbs | 1 + test/correctness/gpu_async_copy.cpp | 55 +++++++++++++++ 15 files changed, 183 insertions(+), 20 deletions(-) create mode 100644 test/correctness/gpu_async_copy.cpp diff --git a/src/CodeGen_D3D12Compute_Dev.cpp b/src/CodeGen_D3D12Compute_Dev.cpp index d916329c89c9..62bea34a6963 100644 --- a/src/CodeGen_D3D12Compute_Dev.cpp +++ b/src/CodeGen_D3D12Compute_Dev.cpp @@ -1406,7 +1406,7 @@ void CodeGen_D3D12Compute_Dev::CodeGen_D3D12Compute_C::visit(const Select *op) { } bool is_shared_allocation(const Allocate *op) { - return op->memory_type == MemoryType::GPUShared; + return is_gpu_shared(op->memory_type); } void CodeGen_D3D12Compute_Dev::CodeGen_D3D12Compute_C::visit(const Allocate *op) { diff --git a/src/CodeGen_Metal_Dev.cpp b/src/CodeGen_Metal_Dev.cpp index 61fbb8602c10..08145925aa10 100644 --- a/src/CodeGen_Metal_Dev.cpp +++ b/src/CodeGen_Metal_Dev.cpp @@ -539,7 +539,7 @@ void CodeGen_Metal_Dev::CodeGen_Metal_C::visit(const Select *op) { void CodeGen_Metal_Dev::CodeGen_Metal_C::visit(const Allocate *op) { - if (op->memory_type == MemoryType::GPUShared) { + if (is_gpu_shared(op->memory_type)) { // Already handled op->body.accept(this); } else { @@ -776,7 +776,7 @@ void CodeGen_Metal_Dev::CodeGen_Metal_C::add_kernel(const Stmt &s, const Allocate *shared_alloc = nullptr; shared_name = "__shared"; visit_with(s, [&](auto *self, const Allocate *op) { - if (op->memory_type == MemoryType::GPUShared) { + if (is_gpu_shared(op->memory_type)) { internal_assert(shared_alloc == nullptr) << "Found multiple shared allocations in metal kernel\n"; shared_alloc = op; diff --git a/src/CodeGen_OpenCL_Dev.cpp b/src/CodeGen_OpenCL_Dev.cpp index 807d75444ed4..229b79a6e184 100644 --- a/src/CodeGen_OpenCL_Dev.cpp +++ b/src/CodeGen_OpenCL_Dev.cpp @@ -780,7 +780,7 @@ void CodeGen_OpenCL_Dev::CodeGen_OpenCL_C::visit(const Allocate *op) { user_assert(!op->new_expr.defined()) << "Allocate node inside OpenCL kernel has custom new expression.\n" << "(Memoization is not supported inside GPU kernels at present.)\n"; - if (op->memory_type == MemoryType::GPUShared) { + if (is_gpu_shared(op->memory_type)) { // Already handled op->body.accept(this); } else { @@ -1063,7 +1063,7 @@ void CodeGen_OpenCL_Dev::CodeGen_OpenCL_C::add_kernel(Stmt s, const Allocate *shared_alloc = nullptr; shared_name = "__shared"; visit_with(s, [&](auto *self, const Allocate *op) { - if (op->memory_type == MemoryType::GPUShared) { + if (is_gpu_shared(op->memory_type)) { internal_assert(shared_alloc == nullptr) << "Found multiple shared allocations in metal kernel\n"; shared_alloc = op; diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 0c706db90813..1433fef168cc 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -14,6 +14,7 @@ #include "IRPrinter.h" #include "LLVM_Headers.h" #include "LLVM_Runtime_Linker.h" +#include "ModulusRemainder.h" #include "Simplify.h" #include "Solve.h" #include "Target.h" @@ -115,7 +116,7 @@ class CodeGen_PTX_Dev : public CodeGen_LLVM, public CodeGen_GPU_Dev { /** Try to emit a store into shared memory as an asynchronous copy, which * moves the data straight from global memory without routing it through * registers. Returns false if this store isn't one we can do that for. */ - bool codegen_async_copy(const Store *op); + bool codegen_async_copy(const Store *op, const char **reason); /** Wait for any asynchronous copies issued so far to have landed. */ void wait_for_async_copies(); @@ -352,7 +353,7 @@ void CodeGen_PTX_Dev::visit(const Allocate *alloc) { user_assert(!alloc->new_expr.defined()) << "Allocate node inside PTX kernel has custom new expression.\n" << "(Memoization is not supported inside GPU kernels at present.)\n"; ScopedBinding bind(alloc_memory_type, alloc->name, alloc->memory_type); - if (alloc->memory_type == MemoryType::GPUShared) { + if (is_gpu_shared(alloc->memory_type)) { // PTX uses zero in address space 3 as the base address for shared memory Value *shared_base = Constant::getNullValue(PointerType::get(*context, 3)); sym_push(alloc->name, shared_base); @@ -416,11 +417,19 @@ void CodeGen_PTX_Dev::visit(const Load *op) { // without going through registers, which saves the load, the store, and the // registers in between. The copy is asynchronous, so it has to be waited for // before the data is used; that happens at the end of the producer. -bool CodeGen_PTX_Dev::codegen_async_copy(const Store *op) { +bool CodeGen_PTX_Dev::codegen_async_copy(const Store *op, const char **reason) { + if (target.get_cuda_capability_lower_bound() < 80) { + *reason = "asynchronous copies require CUDA compute capability 8.0 or above"; + return false; + } + if (emit_atomic_stores) { + *reason = "the store is atomic"; + return false; + } // Asynchronous copies need something to wait for them, which only happens - // at the end of a producer. - if (!in_producer || emit_atomic_stores || - target.get_cuda_capability_lower_bound() < 80) { + // inside a producer. + if (!in_producer) { + *reason = "the store is not inside a produce node"; return false; } @@ -428,14 +437,19 @@ bool CodeGen_PTX_Dev::codegen_async_copy(const Store *op) { // load from something we didn't allocate in here, which is to say global // memory. const MemoryType *dst_memory_type = alloc_memory_type.find(op->name); - if (!dst_memory_type || *dst_memory_type != MemoryType::GPUShared) { + if (!dst_memory_type || !is_gpu_shared(*dst_memory_type)) { + *reason = "the destination is not in shared memory"; return false; } const Load *src = op->value.as(); if (!src || alloc_memory_type.contains(src->name)) { + *reason = "the value stored is not a load from a buffer outside the kernel. " + "An asynchronous copy moves bytes untouched, so the Func must be a " + "plain copy - no cast, no arithmetic, and no boundary condition"; return false; } if (!is_const_one(op->predicate) || !is_const_one(src->predicate)) { + *reason = "the load or the store is predicated"; return false; } @@ -444,6 +458,8 @@ bool CodeGen_PTX_Dev::codegen_async_copy(const Store *op) { const Type t = op->value.type(); const int bytes = t.bytes() * t.lanes(); if (!(bytes == 4 || bytes == 8 || bytes == 16)) { + *reason = "each thread must copy 4, 8 or 16 bytes at a time. Vectorize the " + "copy along its dense dimension by that many bytes' worth"; return false; } Expr dst_base = op->index, src_base = src->index; @@ -453,9 +469,35 @@ bool CodeGen_PTX_Dev::codegen_async_copy(const Store *op) { dst_base = strided_ramp_base(simplify(op->index)); src_base = strided_ramp_base(simplify(src->index)); if (!dst_base.defined() || !src_base.defined()) { + *reason = "the source and the destination are not both indexed densely"; + return false; + } + } + // The hardware needs both addresses aligned to the width of the copy. A + // stride that isn't a multiple of it - which is what an odd align_storage + // produces - shows up here as an unprovable alignment. + if (t.lanes() > 1) { + // Use the alignment lowering worked out, which knows what the loop + // variables in the index are multiples of. + // Only the destination is checked. Its alignment is what align_storage + // controls, so it is the one a schedule can get wrong. A source in a + // buffer whose strides are not known until runtime has no provable + // alignment either way, and rejecting those would fail schedules that + // are fine. + auto aligned = [&](const ModulusRemainder &a, const Expr &base) { + auto ok = [&](const ModulusRemainder &m) { + return m.modulus % t.lanes() == 0 && m.remainder % t.lanes() == 0; + }; + return ok(a) || ok(modulus_remainder(base)); + }; + if (!aligned(op->alignment, dst_base)) { + *reason = "the destination is not known to be aligned to the width of " + "the copy. Any align_storage on this Func has to be a multiple " + "of the number of elements each thread copies"; return false; } } + Value *dst = codegen_buffer_pointer(op->name, t.element_of(), dst_base); Value *src_ptr = codegen_buffer_pointer(src->name, t.element_of(), src_base); @@ -464,6 +506,7 @@ bool CodeGen_PTX_Dev::codegen_async_copy(const Store *op) { llvm::Type *shared_ptr_t = PointerType::get(*context, 3); llvm::Type *global_ptr_t = PointerType::get(*context, 1); if (dst->getType() != shared_ptr_t) { + *reason = "the destination did not end up in the shared address space"; return false; } src_ptr = builder->CreateAddrSpaceCast(src_ptr, global_ptr_t); @@ -518,8 +561,50 @@ void CodeGen_PTX_Dev::visit(const Store *op) { user_assert(op->value.type().bits() >= 32) << "CUDA: 8-bit or 16-bit atomics are not supported.\n"; } - if (codegen_async_copy(op)) { - return; + { + const char *reason = ""; + if (codegen_async_copy(op, &reason)) { + return; + } + // Asking for this memory type is a promise that the stores to it are + // copies the hardware can make asynchronously. If one isn't, say so + // rather than quietly emitting a load and a store instead. + const MemoryType *t = alloc_memory_type.find(op->name); + if (t && *t == MemoryType::GPUSharedAsync) { + user_error + << op->name << " is scheduled in GPUSharedAsync memory, but this " + << "store to it cannot be done with an asynchronous copy, because " + << reason << ".\n\n" + << "An asynchronous copy moves bytes from global memory into shared " + << "memory without routing them through registers. It requires that " + << "the Func is a plain copy of a buffer or another Func - no cast, " + << "arithmetic, or boundary condition, because the bytes move " + << "untouched - and that each thread stores a dense vector of 4, 8 or " + << "16 bytes, aligned to its own size at both ends. It needs CUDA " + << "compute capability 8.0 or above.\n\n" + << "The alignment of the destination is set by align_storage, which " + << "fixes the stride of the staged Func. Padding it to avoid bank " + << "conflicts is usually a good idea, but the padded stride has to " + << "stay a multiple of the vector width or the rows stop being " + << "aligned enough to copy into.\n\n" + << "The usual way to get a Func that is a plain copy is Func::in, " + << "which makes a wrapper that does nothing but hold a staged copy of " + << "something. Vectorizing its dense dimension by a whole number of " + << "bytes gives each thread one copy to issue, and its other " + << "dimensions are spread over the threads of the block as usual:\n\n" + << " A.in()\n" + << " .compute_at(consumer, r)\n" + << " .store_in(MemoryType::GPUSharedAsync)\n" + << " .tile(x, y, xi, yi, 256, 8) // 256 = 8 elements x 32 threads\n" + << " .vectorize(xi, 8) // 8 halves is 16 bytes\n" + << " .gpu_threads(xi, yi);\n\n" + << "The tile has to divide into the thread counts the block already " + << "has: its width over the vector width is the number of threads in " + << "x, and its height the number in y. Anything left over is covered " + << "by the serial loops the tile leaves outside.\n\n" + << "The store that could not be made asynchronous was:\n" + << Stmt(op); + } } // Do aligned 4-wide 32-bit stores as a single i128 store. diff --git a/src/CodeGen_Vulkan_Dev.cpp b/src/CodeGen_Vulkan_Dev.cpp index 2918eba7b85b..b7d8dd844667 100644 --- a/src/CodeGen_Vulkan_Dev.cpp +++ b/src/CodeGen_Vulkan_Dev.cpp @@ -1877,7 +1877,7 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::visit(const Allocate *op) { uint32_t array_size = 0; SpvStorageClass storage_class = SpvStorageClassGeneric; - if (op->memory_type == MemoryType::GPUShared) { + if (is_gpu_shared(op->memory_type)) { // Allocation of shared memory must be declared at global scope storage_class = SpvStorageClassWorkgroup; // shared across workgroup diff --git a/src/CodeGen_WebGPU_Dev.cpp b/src/CodeGen_WebGPU_Dev.cpp index 39a88337af1d..23db34fa02ad 100644 --- a/src/CodeGen_WebGPU_Dev.cpp +++ b/src/CodeGen_WebGPU_Dev.cpp @@ -410,7 +410,7 @@ void CodeGen_WebGPU_Dev::CodeGen_WGSL::add_kernel( } void CodeGen_WebGPU_Dev::CodeGen_WGSL::visit(const Allocate *op) { - if (op->memory_type == MemoryType::GPUShared) { + if (is_gpu_shared(op->memory_type)) { internal_assert(!workgroup_allocations.count(op->name)); workgroup_allocations.insert({op->name, op}); op->body.accept(this); diff --git a/src/Deserialization.cpp b/src/Deserialization.cpp index 20e665067763..72e975a12a41 100644 --- a/src/Deserialization.cpp +++ b/src/Deserialization.cpp @@ -180,6 +180,8 @@ MemoryType Deserializer::deserialize_memory_type(Serialize::MemoryType memory_ty return MemoryType::Register; case Serialize::MemoryType::GPUShared: return MemoryType::GPUShared; + case Serialize::MemoryType::GPUSharedAsync: + return MemoryType::GPUSharedAsync; case Serialize::MemoryType::GPUTexture: return MemoryType::GPUTexture; case Serialize::MemoryType::LockedCache: diff --git a/src/Expr.h b/src/Expr.h index 0f71b860f4b9..5a800e7bd625 100644 --- a/src/Expr.h +++ b/src/Expr.h @@ -405,8 +405,21 @@ enum class MemoryType { /** AMX Tile register for X86. Any data that would be used in an AMX matrix * multiplication must first be loaded into an AMX tile register. */ AMXTile, + + /** GPU shared memory, written by an asynchronous copy instruction that + * moves data straight from global memory without routing it through + * registers. The only stores allowed to such an allocation are unpredicated + * copies of a whole number of 4, 8 or 16 byte chunks from a densely indexed + * global buffer, because that is all the hardware can do. On GPU APIs with + * no such instruction this is ordinary shared memory. */ + GPUSharedAsync, }; +/** Whether a MemoryType places an allocation in GPU shared memory. */ +inline bool is_gpu_shared(MemoryType t) { + return t == MemoryType::GPUShared || t == MemoryType::GPUSharedAsync; +} + /** Whether a MemoryType is backed by tile-shaped storage with native 2D * loads and stores. Generic optimizations that rewrite the structure of * loads or stores (deinterleaving, strided-load staging, etc.) should diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index 8aad78c18db4..5d3d0d9481ac 100644 --- a/src/FuseGPUThreadLoops.cpp +++ b/src/FuseGPUThreadLoops.cpp @@ -461,7 +461,7 @@ class ExtractSharedAndHeapAllocations : public IRMutator { if ((fixed_size_thread_allocation && op->memory_type != MemoryType::Heap && - op->memory_type != MemoryType::GPUShared && + !is_gpu_shared(op->memory_type) && op->memory_type != MemoryType::GPUTexture) || op->memory_type == MemoryType::Register || op->memory_type == MemoryType::Stack) { @@ -470,7 +470,7 @@ class ExtractSharedAndHeapAllocations : public IRMutator { } user_assert(op->memory_type == MemoryType::Auto || - op->memory_type == MemoryType::GPUShared || + is_gpu_shared(op->memory_type) || op->memory_type == MemoryType::GPUTexture || op->memory_type == MemoryType::Heap) << "Allocation " << op->name << " must live in shared or heap memory, " @@ -1396,6 +1396,7 @@ class InjectThreadBarriers : public IRMutator { debug(4) << "Encountered store to " << op->name << "\n"; auto mem_type = memory_type_for_name(op->name); switch (mem_type) { + case MemoryType::GPUSharedAsync: case MemoryType::GPUShared: debug(4) << " memory type is shared\n"; shared_stores.insert(op->name); @@ -1421,6 +1422,7 @@ class InjectThreadBarriers : public IRMutator { debug(4) << "Encountered load from " << op->name << "\n"; auto mem_type = memory_type_for_name(op->name); switch (mem_type) { + case MemoryType::GPUSharedAsync: case MemoryType::GPUShared: debug(4) << " memory type is shared\n"; shared_loads.insert(op->name); diff --git a/src/IRPrinter.cpp b/src/IRPrinter.cpp index 8dd035e12865..33451ded3b72 100644 --- a/src/IRPrinter.cpp +++ b/src/IRPrinter.cpp @@ -154,6 +154,9 @@ std::ostream &operator<<(std::ostream &out, const MemoryType &t) { case MemoryType::Register: out << "Register"; break; + case MemoryType::GPUSharedAsync: + out << "GPUSharedAsync"; + break; case MemoryType::GPUShared: out << "GPUShared"; break; diff --git a/src/LowerWarpShuffles.cpp b/src/LowerWarpShuffles.cpp index fedde5ba2166..75f52e1d65cc 100644 --- a/src/LowerWarpShuffles.cpp +++ b/src/LowerWarpShuffles.cpp @@ -654,7 +654,7 @@ class LowerWarpShuffles : public IRMutator { Stmt visit(const Allocate *op) override { if (this_lane.defined() || - op->memory_type == MemoryType::GPUShared || + is_gpu_shared(op->memory_type) || op->memory_type == MemoryType::Heap) { // Not a warp-level allocation. Warp-level storage is per-lane // register storage; shared and heap (global) memory are never diff --git a/src/OffloadGPULoops.cpp b/src/OffloadGPULoops.cpp index d06345601183..ff33e4a3be12 100644 --- a/src/OffloadGPULoops.cpp +++ b/src/OffloadGPULoops.cpp @@ -79,7 +79,7 @@ class ExtractBounds : public IRVisitor { user_assert(!allocate->new_expr.defined()) << "Allocate node inside GPU kernel has custom new expression.\n" << "(Memoization is not supported inside GPU kernels at present.)\n"; - if (allocate->memory_type == MemoryType::GPUShared) { + if (is_gpu_shared(allocate->memory_type)) { internal_assert(allocate->extents.size() == 1); shared_mem_size += allocate->extents[0] * allocate->type.bytes(); found_shared = true; diff --git a/src/Serialization.cpp b/src/Serialization.cpp index dc201a8c818f..c460d35e8eb4 100644 --- a/src/Serialization.cpp +++ b/src/Serialization.cpp @@ -150,6 +150,8 @@ Serialize::MemoryType Serializer::serialize_memory_type(const MemoryType &memory return Serialize::MemoryType::Register; case MemoryType::GPUShared: return Serialize::MemoryType::GPUShared; + case MemoryType::GPUSharedAsync: + return Serialize::MemoryType::GPUSharedAsync; case MemoryType::GPUTexture: return Serialize::MemoryType::GPUTexture; case MemoryType::LockedCache: diff --git a/src/halide_ir.fbs b/src/halide_ir.fbs index 30d56f5af73e..9967296986f3 100644 --- a/src/halide_ir.fbs +++ b/src/halide_ir.fbs @@ -116,6 +116,7 @@ enum MemoryType: byte { LockedCache, VTCM, AMXTile, + GPUSharedAsync, } table Range { diff --git a/test/correctness/gpu_async_copy.cpp b/test/correctness/gpu_async_copy.cpp new file mode 100644 index 000000000000..8b39f21e5cad --- /dev/null +++ b/test/correctness/gpu_async_copy.cpp @@ -0,0 +1,55 @@ +#include "Halide.h" +#include + +using namespace Halide; + +// A Func stored in GPUSharedAsync is staged into shared memory by the copy +// engine rather than by loading it into registers and storing it back out. +int main(int argc, char **argv) { + Target target = get_jit_target_from_environment(); + if (!target.has_feature(Target::CUDA)) { + printf("[SKIP] No CUDA target enabled.\n"); + return 0; + } + if (target.get_cuda_capability_lower_bound() < 80) { + printf("[SKIP] Asynchronous copies need compute capability 8.0 or above.\n"); + return 0; + } + + const int W = 256, H = 64; + + Buffer input(W, H); + input.fill([](int x, int y) { return (float)(x + y * 3); }); + + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + + // The staged Func has to be a plain copy - the hardware moves the bytes + // untouched - so any arithmetic goes in the consumer. + stage(x, y) = input(x, y); + out(x, y) = stage(x, y) * 2.f + 1.f; + + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, 4); + + Buffer result(W, H); + out.realize(result); + result.copy_to_host(); + + for (int j = 0; j < H; j++) { + for (int i = 0; i < W; i++) { + const float correct = (float)(i + j * 3) * 2.f + 1.f; + if (result(i, j) != correct) { + printf("result(%d, %d) = %f instead of %f\n", + i, j, result(i, j), correct); + return 1; + } + } + } + + printf("Success!\n"); + return 0; +} From 565a68e851ec607d2c63bc87ece806e3ef17cc26 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Thu, 30 Jul 2026 14:52:21 -0700 Subject: [PATCH 12/59] Expose the remaining memory types to the Python bindings AMXTile was never added, and GPUSharedAsync is new. List them in the same order as the enum in Expr.h so that a missing one is easier to spot. Co-Authored-By: Claude Opus 5 --- python_bindings/src/halide/halide_/PyEnums.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python_bindings/src/halide/halide_/PyEnums.cpp b/python_bindings/src/halide/halide_/PyEnums.cpp index c0f68be20fb3..e9abff7d4883 100644 --- a/python_bindings/src/halide/halide_/PyEnums.cpp +++ b/python_bindings/src/halide/halide_/PyEnums.cpp @@ -48,7 +48,9 @@ void define_enums(py::module &m) { .value("GPUShared", MemoryType::GPUShared) .value("GPUTexture", MemoryType::GPUTexture) .value("LockedCache", MemoryType::LockedCache) - .value("VTCM", MemoryType::VTCM); + .value("VTCM", MemoryType::VTCM) + .value("AMXTile", MemoryType::AMXTile) + .value("GPUSharedAsync", MemoryType::GPUSharedAsync); py::enum_(m, "NameMangling") .value("Default", NameMangling::Default) From cd865e9cc807a57186ed14311f3c1c9746cc51d5 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Thu, 30 Jul 2026 15:08:51 -0700 Subject: [PATCH 13/59] Test the async copy constraints, including the ways to break them Cover each of the three copy widths the hardware supports at several element sizes, plus the shapes a staged input takes: a two-dimensional tile, a padded stride, two inputs staged into one kernel, and a Func::in wrapper. The error test checks that breaking each constraint produces a user error saying which one, rather than crashing or quietly falling back to a load and a store. Co-Authored-By: Claude Opus 5 --- test/correctness/gpu_async_copy.cpp | 194 ++++++++++++++++++--- test/correctness/gpu_async_copy_errors.cpp | 187 ++++++++++++++++++++ 2 files changed, 357 insertions(+), 24 deletions(-) create mode 100644 test/correctness/gpu_async_copy_errors.cpp diff --git a/test/correctness/gpu_async_copy.cpp b/test/correctness/gpu_async_copy.cpp index 8b39f21e5cad..0de8fe76c296 100644 --- a/test/correctness/gpu_async_copy.cpp +++ b/test/correctness/gpu_async_copy.cpp @@ -5,51 +5,197 @@ using namespace Halide; // A Func stored in GPUSharedAsync is staged into shared memory by the copy // engine rather than by loading it into registers and storing it back out. -int main(int argc, char **argv) { - Target target = get_jit_target_from_environment(); - if (!target.has_feature(Target::CUDA)) { - printf("[SKIP] No CUDA target enabled.\n"); - return 0; - } - if (target.get_cuda_capability_lower_bound() < 80) { - printf("[SKIP] Asynchronous copies need compute capability 8.0 or above.\n"); - return 0; +// The copy engine moves 4, 8 or 16 bytes per thread, so the cases below cover +// each of those widths at several element sizes, as well as the shapes a +// staged input tends to take: a two-dimensional tile, more than one input +// staged into the same kernel, and a wrapper made with Func::in. + +namespace { + +int failures = 0; + +template +void check_result(const char *name, Buffer &result, Buffer &input) { + for (int y = 0; y < result.height(); y++) { + for (int x = 0; x < result.width(); x++) { + T correct = (T)(input(x, y) * 2); + if (result(x, y) != correct) { + printf("[%s] result(%d, %d) = %f instead of %f\n", + name, x, y, (double)result(x, y), (double)correct); + failures++; + return; + } + } } + printf("[%s] OK\n", name); +} - const int W = 256, H = 64; +template +Buffer make_input(int w, int h) { + Buffer in(w, h); + in.fill([](int x, int y) { return (T)((x + y * 3) % 32); }); + return in; +} - Buffer input(W, H); - input.fill([](int x, int y) { return (float)(x + y * 3); }); +// Stage an input through shared memory with the copy engine, vectorized by +// `vec` elements, and check the consumer sees the right values. +template +void test_width(const char *name, int vec) { + const int W = 256, H = 32; + Buffer input = make_input(W, H); Var x("x"), y("y"), xi("xi"), yi("yi"); Func stage("stage"), out("out"); + stage(x, y) = input(x, y); + out(x, y) = cast(stage(x, y) * 2); + + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, vec); - // The staged Func has to be a plain copy - the hardware moves the bytes - // untouched - so any arithmetic goes in the consumer. + Buffer result(W, H); + out.realize(result); + result.copy_to_host(); + check_result(name, result, input); +} + +// The staged tile spread over threads in both dimensions. +void test_2d_tile() { + const int W = 256, H = 64; + Buffer input = make_input(W, H); + + Var x("x"), y("y"), xi("xi"), yi("yi"), xii("xii"); + Func stage("stage"), out("out"); + stage(x, y) = input(x, y); + out(x, y) = stage(x, y) * 2; + + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .split(x, x, xii, 4) + .vectorize(xii) + .gpu_threads(x, y); + + Buffer result(W, H); + out.realize(result); + result.copy_to_host(); + check_result("2d_tile", result, input); +} + +// Padding the rows to dodge bank conflicts, with a stride that stays a +// multiple of the vector width so the rows are still aligned. +void test_padded_storage() { + const int W = 256, H = 32; + Buffer input = make_input(W, H); + + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); stage(x, y) = input(x, y); - out(x, y) = stage(x, y) * 2.f + 1.f; + out(x, y) = stage(x, y) * 2; out.gpu_tile(x, y, xi, yi, 64, 8); stage.compute_at(out, x) .store_in(MemoryType::GPUSharedAsync) + .align_storage(x, 8) .gpu_threads(y) .vectorize(x, 4); Buffer result(W, H); out.realize(result); result.copy_to_host(); + check_result("padded_storage", result, input); +} - for (int j = 0; j < H; j++) { - for (int i = 0; i < W; i++) { - const float correct = (float)(i + j * 3) * 2.f + 1.f; - if (result(i, j) != correct) { - printf("result(%d, %d) = %f instead of %f\n", - i, j, result(i, j), correct); - return 1; - } - } +// Two inputs staged into the same kernel, so more than one copy is in flight +// before the barrier that waits for them. +void test_two_inputs() { + const int W = 256, H = 32; + Buffer a = make_input(W, H); + Buffer b = make_input(W, H); + + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func sa("sa"), sb("sb"), out("out"); + sa(x, y) = a(x, y); + sb(x, y) = b(x, y); + out(x, y) = sa(x, y) + sb(x, y); + + out.gpu_tile(x, y, xi, yi, 64, 8); + sa.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, 4); + sb.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, 4); + + Buffer result(W, H); + out.realize(result); + result.copy_to_host(); + check_result("two_inputs", result, a); +} + +// Func::in is the idiomatic way to get a Func that is a plain copy, and is +// what the error message points users at. +void test_wrapper() { + const int W = 256, H = 32; + Buffer input = make_input(W, H); + + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func in_f("in_f"), out("out"); + in_f(x, y) = input(x, y); + out(x, y) = in_f(x, y) * 2; + + out.gpu_tile(x, y, xi, yi, 64, 8); + in_f.in() + .compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, 4); + in_f.compute_root(); + + Buffer result(W, H); + out.realize(result); + result.copy_to_host(); + check_result("wrapper", result, input); +} + +} // namespace + +int main(int argc, char **argv) { + Target target = get_jit_target_from_environment(); + if (!target.has_feature(Target::CUDA)) { + printf("[SKIP] No CUDA target enabled.\n"); + return 0; } + if (target.get_cuda_capability_lower_bound() < 80) { + printf("[SKIP] Asynchronous copies need compute capability 8.0 or above.\n"); + return 0; + } + + // Each of the three copy widths the hardware supports, at several element + // sizes. + test_width("f32_x1_4b", 1); + test_width("f32_x2_8b", 2); + test_width("f32_x4_16b", 4); + test_width("u8_x4_4b", 4); + test_width("u8_x8_8b", 8); + test_width("u8_x16_16b", 16); + test_width("u16_x2_4b", 2); + test_width("u16_x8_16b", 8); + test_width("i32_x4_16b", 4); + test_2d_tile(); + test_padded_storage(); + test_two_inputs(); + test_wrapper(); + + if (failures) { + printf("%d case(s) failed\n", failures); + return 1; + } printf("Success!\n"); return 0; } diff --git a/test/correctness/gpu_async_copy_errors.cpp b/test/correctness/gpu_async_copy_errors.cpp new file mode 100644 index 000000000000..0cb8276534fa --- /dev/null +++ b/test/correctness/gpu_async_copy_errors.cpp @@ -0,0 +1,187 @@ +// Exercises the user-facing error paths for MemoryType::GPUSharedAsync. Asking +// for that memory type is a promise that every store to the allocation is a +// copy the copy engine can make, so each scenario below breaks one of those +// requirements and should be told exactly which one. +// +// The test verifies that each scenario produces a Halide::CompileError (a user +// error) rather than crashing, hitting an internal assert, or quietly falling +// back to a load and a store. + +#include "Halide.h" +#include + +#if HALIDE_WITH_EXCEPTIONS + +using namespace Halide; + +namespace { + +// Compiling is enough to reach the error - no device is needed. +Target async_target() { + return get_host_target() + .with_feature(Target::CUDA) + .with_feature(Target::CUDACapability80); +} + +// Run `body` and assert it produces a Halide user error mentioning `substring`. +template +bool expect_user_error(const char *name, const char *substring, F body) { + try { + body(); + } catch (const CompileError &e) { + std::string msg = e.what(); + if (msg.find(substring) == std::string::npos) { + printf("[%s] FAIL: error did not mention \"%s\":\n%s\n", + name, substring, msg.c_str()); + return false; + } + printf("[%s] OK\n", name); + return true; + } catch (...) { + printf("[%s] FAIL: expected a CompileError but got a different exception\n", name); + return false; + } + printf("[%s] FAIL: expected a user error but none was raised\n", name); + return false; +} + +Buffer input_f32() { + Buffer in(256, 64); + in.fill(0.f); + return in; +} + +// The copy engine moves bytes untouched, so the staged Func has to be a plain +// copy. Doing arithmetic in it means the value stored isn't a load at all. +void scenario_not_a_copy() { + Buffer in = input_f32(); + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = in(x, y) * 2.f; + out(x, y) = stage(x, y); + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, 4); + out.compile_jit(async_target()); +} + +// The source has to live outside the kernel. Staging something already +// computed into registers or shared memory isn't a global-to-shared copy. +void scenario_source_inside_kernel() { + Buffer in = input_f32(); + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func producer("producer"), stage("stage"), out("out"); + producer(x, y) = in(x, y) + 1.f; + stage(x, y) = producer(x, y); + out(x, y) = stage(x, y); + out.gpu_tile(x, y, xi, yi, 64, 8); + producer.compute_at(out, x).gpu_threads(y); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, 4); + out.compile_jit(async_target()); +} + +// The hardware copies 4, 8 or 16 bytes per thread. Three floats is 12. +void scenario_bad_vector_width() { + Buffer in = input_f32(); + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = in(x, y); + out(x, y) = stage(x, y); + out.gpu_tile(x, y, xi, yi, 48, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, 3); + out.compile_jit(async_target()); +} + +// A single byte per thread is below the minimum too. +void scenario_too_narrow() { + Buffer in(256, 64); + in.fill(0); + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = in(x, y); + out(x, y) = stage(x, y); + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(x, y); + out.compile_jit(async_target()); +} + +// Each copy has to be contiguous in both the source and the destination, so a +// strided read of the input can't be done this way. A strided load is staged +// as a shuffle of dense loads, so what the peephole sees is a value that isn't +// a plain load. +void scenario_strided_source() { + Buffer in(512, 64); + in.fill(0.f); + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = in(2 * x, y); + out(x, y) = stage(x, y); + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, 4); + out.compile_jit(async_target()); +} + +// The destination address of each copy has to be aligned to its width. Padding +// the rows to a stride that isn't a multiple of the vector width breaks that +// for every row after the first. +void scenario_misaligned_destination() { + Buffer in = input_f32(); + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = in(x, y); + out(x, y) = stage(x, y); + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .align_storage(x, 6) + .gpu_threads(y) + .vectorize(x, 4); + out.compile_jit(async_target()); +} + +} // namespace + +int main(int argc, char **argv) { + if (!Halide::exceptions_enabled()) { + printf("[SKIP] Halide was compiled without exceptions.\n"); + return 0; + } + + int failures = 0; + + failures += !expect_user_error("not_a_copy", "not a load", scenario_not_a_copy); + failures += !expect_user_error("source_inside_kernel", "not a load", scenario_source_inside_kernel); + failures += !expect_user_error("bad_vector_width", "4, 8 or 16 bytes", scenario_bad_vector_width); + failures += !expect_user_error("too_narrow", "4, 8 or 16 bytes", scenario_too_narrow); + failures += !expect_user_error("strided_source", "not a load", scenario_strided_source); + failures += !expect_user_error("misaligned_destination", "aligned", scenario_misaligned_destination); + + if (failures != 0) { + printf("%d scenario(s) did not produce the expected user error\n", failures); + return 1; + } + printf("Success!\n"); + return 0; +} + +#else // HALIDE_WITH_EXCEPTIONS + +int main(int argc, char **argv) { + printf("[SKIP] Halide was compiled without exceptions.\n"); + return 0; +} + +#endif // HALIDE_WITH_EXCEPTIONS From c7eb195c731a6eab0a82e0de0c3157e5a60d1c72 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Thu, 30 Jul 2026 17:11:00 -0700 Subject: [PATCH 14/59] Only copy asynchronously when the schedule asks for it The peephole fired on any store into shared memory that matched, so a Func in plain GPUShared got an asynchronous copy too, and there was no way to ask for the synchronous version. That made the memory type only control whether a store that didn't match was an error, rather than whether the copy was asynchronous at all. Require the destination to be GPUSharedAsync, so that GPUShared and GPUSharedAsync are the two ways to ask for the two lowerings, and a schedule can compare them. Co-Authored-By: Claude Opus 5 --- src/CodeGen_PTX_Dev.cpp | 12 +++++++----- test/correctness/gpu_async_copy.cpp | 29 +++++++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 5 deletions(-) diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 1433fef168cc..70802d92779e 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -433,12 +433,14 @@ bool CodeGen_PTX_Dev::codegen_async_copy(const Store *op, const char **reason) { return false; } - // The destination must be shared memory, and the source must be a plain - // load from something we didn't allocate in here, which is to say global - // memory. + // The destination must be shared memory that was asked for asynchronously, + // and the source must be a plain load from something we didn't allocate in + // here, which is to say global memory. Stores to plain GPUShared that + // happen to match the pattern are left synchronous, so that a schedule can + // ask for either one. const MemoryType *dst_memory_type = alloc_memory_type.find(op->name); - if (!dst_memory_type || !is_gpu_shared(*dst_memory_type)) { - *reason = "the destination is not in shared memory"; + if (!dst_memory_type || *dst_memory_type != MemoryType::GPUSharedAsync) { + *reason = "the destination is not stored in MemoryType::GPUSharedAsync"; return false; } const Load *src = op->value.as(); diff --git a/test/correctness/gpu_async_copy.cpp b/test/correctness/gpu_async_copy.cpp index 0de8fe76c296..403068cbbdad 100644 --- a/test/correctness/gpu_async_copy.cpp +++ b/test/correctness/gpu_async_copy.cpp @@ -162,6 +162,34 @@ void test_wrapper() { check_result("wrapper", result, input); } +// The same staging in plain GPUShared must stay synchronous, so that a +// schedule can ask for either one and compare them. +void test_opt_out() { + const int W = 256, H = 32; + Buffer input = make_input(W, H); + + auto build = [&](MemoryType mt) { + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = input(x, y); + out(x, y) = stage(x, y) * 2; + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(mt) + .gpu_threads(y) + .vectorize(x, 4); + Buffer result(W, H); + out.realize(result); + result.copy_to_host(); + return result; + }; + + Buffer shared = build(MemoryType::GPUShared); + check_result("opt_out_shared", shared, input); + Buffer async = build(MemoryType::GPUSharedAsync); + check_result("opt_out_async", async, input); +} + } // namespace int main(int argc, char **argv) { @@ -191,6 +219,7 @@ int main(int argc, char **argv) { test_padded_storage(); test_two_inputs(); test_wrapper(); + test_opt_out(); if (failures) { printf("%d case(s) failed\n", failures); From 97cd537dc961bd92c18f9f367dda9639ce51775b Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Thu, 30 Jul 2026 17:36:29 -0700 Subject: [PATCH 15/59] Say which requirement an async copy failed to meet A store whose value wasn't a plain load was always reported as the Func not being a copy, but that also caught two cases where it is a copy and something else is wrong. A source read with a stride is broken into a shuffle of dense loads before it reaches here, and a source computed elsewhere in the kernel is a load from the wrong place. Report those as themselves. Also test the predicated case, which needs an align_storage to keep the destination aligned, or the alignment requirement is what fails first. Co-Authored-By: Claude Opus 5 --- src/CodeGen_PTX_Dev.cpp | 20 +++++++++++++++- test/correctness/gpu_async_copy_errors.cpp | 27 ++++++++++++++++++---- 2 files changed, 41 insertions(+), 6 deletions(-) diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 70802d92779e..889c6529e9f6 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -444,12 +444,30 @@ bool CodeGen_PTX_Dev::codegen_async_copy(const Store *op, const char **reason) { return false; } const Load *src = op->value.as(); - if (!src || alloc_memory_type.contains(src->name)) { + if (!src) { + // A load that isn't dense is broken up into a shuffle of dense loads + // well before we get here, so say what that means for the copy rather + // than describing it as not being a load. + Expr value = op->value; + while (const Let *let = value.as()) { + value = let->body; + } + if (const Shuffle *s = value.as(); + s && !s->vectors.empty()) { + *reason = "the source is not read densely. Each copy moves one run of " + "bytes, so the Func must read its source with a stride of one"; + return false; + } *reason = "the value stored is not a load from a buffer outside the kernel. " "An asynchronous copy moves bytes untouched, so the Func must be a " "plain copy - no cast, no arithmetic, and no boundary condition"; return false; } + if (alloc_memory_type.contains(src->name)) { + *reason = "the value stored is loaded from another allocation inside the " + "kernel. An asynchronous copy reads from global memory"; + return false; + } if (!is_const_one(op->predicate) || !is_const_one(src->predicate)) { *reason = "the load or the store is predicated"; return false; diff --git a/test/correctness/gpu_async_copy_errors.cpp b/test/correctness/gpu_async_copy_errors.cpp index 0cb8276534fa..b2e15d8c761f 100644 --- a/test/correctness/gpu_async_copy_errors.cpp +++ b/test/correctness/gpu_async_copy_errors.cpp @@ -116,9 +116,7 @@ void scenario_too_narrow() { } // Each copy has to be contiguous in both the source and the destination, so a -// strided read of the input can't be done this way. A strided load is staged -// as a shuffle of dense loads, so what the peephole sees is a value that isn't -// a plain load. +// strided read of the input can't be done this way. void scenario_strided_source() { Buffer in(512, 64); in.fill(0.f); @@ -134,6 +132,24 @@ void scenario_strided_source() { out.compile_jit(async_target()); } +// A predicated tail means some lanes are masked off, which the copy engine +// can't express. +void scenario_predicated() { + Buffer in(256, 64); + in.fill(0.f); + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = in(x, y); + out(x, y) = stage(x, y); + out.gpu_tile(x, y, xi, yi, 66, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .align_storage(x, 4) + .gpu_threads(y) + .vectorize(x, 4, TailStrategy::Predicate); + out.compile_jit(async_target()); +} + // The destination address of each copy has to be aligned to its width. Padding // the rows to a stride that isn't a multiple of the vector width breaks that // for every row after the first. @@ -163,10 +179,11 @@ int main(int argc, char **argv) { int failures = 0; failures += !expect_user_error("not_a_copy", "not a load", scenario_not_a_copy); - failures += !expect_user_error("source_inside_kernel", "not a load", scenario_source_inside_kernel); + failures += !expect_user_error("source_inside_kernel", "another allocation inside the", scenario_source_inside_kernel); failures += !expect_user_error("bad_vector_width", "4, 8 or 16 bytes", scenario_bad_vector_width); failures += !expect_user_error("too_narrow", "4, 8 or 16 bytes", scenario_too_narrow); - failures += !expect_user_error("strided_source", "not a load", scenario_strided_source); + failures += !expect_user_error("strided_source", "not read densely", scenario_strided_source); + failures += !expect_user_error("predicated", "predicated", scenario_predicated); failures += !expect_user_error("misaligned_destination", "aligned", scenario_misaligned_destination); if (failures != 0) { From baa9e0402c255f2dae11d2c266f982e75e4bf152 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Thu, 30 Jul 2026 17:50:20 -0700 Subject: [PATCH 16/59] Test that the destination of an async copy must be dense too The density of the two ends of the copy is checked by strided_ramp_base, whose default stride of one is doing the work. That is easy to misread as extracting an address, so say so where it is called. A source read with a stride never reaches that check, because it is broken into a shuffle of dense loads first and fails the earlier test for being a plain load. Storing the staged Func in the opposite order to the one it is read in is what reaches it, so test that. Co-Authored-By: Claude Opus 5 --- src/CodeGen_PTX_Dev.cpp | 2 ++ test/correctness/gpu_async_copy_errors.cpp | 19 +++++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 889c6529e9f6..78c0b91d3579 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -486,6 +486,8 @@ bool CodeGen_PTX_Dev::codegen_async_copy(const Store *op, const char **reason) { if (t.lanes() > 1) { // Shared allocations are given an offset into one big block after the // last simplification pass, so the indices need simplifying here. + // strided_ramp_base returns undefined unless the stride is exactly + // one, so this checks the density as well as finding the address. dst_base = strided_ramp_base(simplify(op->index)); src_base = strided_ramp_base(simplify(src->index)); if (!dst_base.defined() || !src_base.defined()) { diff --git a/test/correctness/gpu_async_copy_errors.cpp b/test/correctness/gpu_async_copy_errors.cpp index b2e15d8c761f..cf533c16ec09 100644 --- a/test/correctness/gpu_async_copy_errors.cpp +++ b/test/correctness/gpu_async_copy_errors.cpp @@ -150,6 +150,24 @@ void scenario_predicated() { out.compile_jit(async_target()); } +// Each copy has to be dense at both ends. Storing the staged Func in the +// opposite order to the one it is read in leaves the source dense but the +// destination strided. +void scenario_strided_destination() { + Buffer in = input_f32(); + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = in(x, y); + out(x, y) = stage(x, y); + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .reorder_storage(y, x) + .gpu_threads(y) + .vectorize(x, 4); + out.compile_jit(async_target()); +} + // The destination address of each copy has to be aligned to its width. Padding // the rows to a stride that isn't a multiple of the vector width breaks that // for every row after the first. @@ -184,6 +202,7 @@ int main(int argc, char **argv) { failures += !expect_user_error("too_narrow", "4, 8 or 16 bytes", scenario_too_narrow); failures += !expect_user_error("strided_source", "not read densely", scenario_strided_source); failures += !expect_user_error("predicated", "predicated", scenario_predicated); + failures += !expect_user_error("strided_destination", "densely", scenario_strided_destination); failures += !expect_user_error("misaligned_destination", "aligned", scenario_misaligned_destination); if (failures != 0) { From 999bd5fbf63dc1159ca0654ab149fff26dd8acba Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 28 Jul 2026 15:43:47 -0700 Subject: [PATCH 17/59] Add NVIDIA tensor core support via a WMMAAccumulator memory type This is the GPU counterpart to the existing AMXTile support. Scheduling a matmul accumulator with .store_in(MemoryType::WMMAAccumulator) makes a vanilla Halide matrix multiply compile to wmma instructions. The new pass, extract_wmma_operations, recognizes the three operations that a tensor core accumulator supports - initialization from zero or from a matrix in memory, accumulation of a matrix multiply, and copying the result out - and rewrites them as intrinsics. Anything else is an error, because the layout of a fragment across the registers of a warp isn't architecturally specified, so those instructions are the only way in or out. Nothing in the schedule says the accumulator is spread over a warp, so the pass also introduces the loop over the 32 lanes. The fragment intrinsics are pure functions of a matrix value and the lane. The matrix is named by a Load of the whole of it, with the lanes in row-major order regardless of how it's laid out in memory, so that the passes that track uses of an allocation see the access and its true footprint. Its layout in memory and the distance between its rows or columns are recovered from the strides of that Load's index. Copying an accumulator out is a predicated store of the whole matrix, where the predicate says which entries this lane holds. Supported: float16 operands, float32 or float16 accumulators, all three tile shapes, row- and column-major operands and results, several accumulator fragments per warp, several warps per block, and operands staged through shared memory. Along the way: - is_multiramp learns to see through the lane permutations the simplifier and flatten_nested_ramps introduce, so those passes are free to rewrite the accesses and the backend puts them back. - The simplifier lifts broadcasts out of pure elementwise calls, so a lane-uniform value stays recognizable as one. - The subtile partitioning that AMX was doing moves to MultiRamp, shared between the two passes. - The CUDA runtime no longer caps registers per thread at 64. That was costing the tensor core matmul about 1.7x, and cost the two apps that set HL_CUDA_JIT_MAX_REGISTERS to work around it. apps/tensorcore_matmul reaches 40 TFlop/s on an RTX 5060 Ti, against 8 for the best non-tensor-core schedule and 50 for cuBLAS. apps/tensorcore_resize ports the block-sparse resampling algorithm, in which a resize becomes a dense matrix multiply. Its cudaonly schedule works; the tensorcore one is blocked on a simplifier gap documented in its README. Co-Authored-By: Claude Opus 5 --- Makefile | 2 + apps/tensorcore_matmul/Makefile | 38 ++ apps/tensorcore_matmul/matmul_generator.cpp | 133 +++++ apps/tensorcore_matmul/runner.cpp | 84 +++ apps/tensorcore_resize/Makefile | 31 + apps/tensorcore_resize/README.md | 34 ++ apps/tensorcore_resize/resize_generator.cpp | 334 +++++++++++ apps/tensorcore_resize/runner.cpp | 72 +++ src/CMakeLists.txt | 2 + src/CanonicalizeGPUVars.cpp | 26 +- src/CodeGen_PTX_Dev.cpp | 188 ++++++ src/Deserialization.cpp | 2 + src/Expr.h | 9 +- src/ExtractWMMAOperations.cpp | 618 ++++++++++++++++++++ src/ExtractWMMAOperations.h | 94 +++ src/FlattenNestedRamps.cpp | 23 + src/FuseGPUThreadLoops.cpp | 8 +- src/IR.cpp | 6 + src/IR.h | 31 + src/IRPrinter.cpp | 3 + src/Lower.cpp | 7 + src/LowerWarpShuffles.cpp | 11 +- src/Serialization.cpp | 2 + src/Simplify_Call.cpp | 41 ++ src/Simplify_Stmts.cpp | 1 + src/halide_ir.fbs | 1 + test/correctness/wmma_matmul.cpp | 262 +++++++++ 27 files changed, 2055 insertions(+), 8 deletions(-) create mode 100644 apps/tensorcore_matmul/Makefile create mode 100644 apps/tensorcore_matmul/matmul_generator.cpp create mode 100644 apps/tensorcore_matmul/runner.cpp create mode 100644 apps/tensorcore_resize/Makefile create mode 100644 apps/tensorcore_resize/README.md create mode 100644 apps/tensorcore_resize/resize_generator.cpp create mode 100644 apps/tensorcore_resize/runner.cpp create mode 100644 src/ExtractWMMAOperations.cpp create mode 100644 src/ExtractWMMAOperations.h create mode 100644 test/correctness/wmma_matmul.cpp diff --git a/Makefile b/Makefile index e15b978d3eb4..3fb6529b8fe9 100644 --- a/Makefile +++ b/Makefile @@ -508,6 +508,7 @@ SOURCE_FILES = \ Error.cpp \ Expr.cpp \ ExtractTileOperations.cpp \ + ExtractWMMAOperations.cpp \ FastIntegerDivide.cpp \ FindCalls.cpp \ FindIntrinsics.cpp \ @@ -710,6 +711,7 @@ HEADER_FILES = \ Extern.h \ ExternFuncArgument.h \ ExtractTileOperations.h \ + ExtractWMMAOperations.h \ FastIntegerDivide.h \ FindCalls.h \ FindIntrinsics.h \ diff --git a/apps/tensorcore_matmul/Makefile b/apps/tensorcore_matmul/Makefile new file mode 100644 index 000000000000..0ce43fc73816 --- /dev/null +++ b/apps/tensorcore_matmul/Makefile @@ -0,0 +1,38 @@ +include ../support/Makefile.inc + +MATMUL_M ?= 1024 +MATMUL_N ?= 1024 +MATMUL_K ?= 1024 + +# The wmma instructions require compute capability 7.0 or above. +TENSORCORE_TARGET ?= host-cuda-cuda_capability_80 + +all: $(BIN)/$(HL_TARGET)/runner + +$(GENERATOR_BIN)/matmul.generator: matmul_generator.cpp $(GENERATOR_DEPS) + @mkdir -p $(@D) + $(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LIBHALIDE_LDFLAGS) + +$(BIN)/%/matmul_cudaonly.a: $(GENERATOR_BIN)/matmul.generator + @mkdir -p $(@D) + $^ -g matmul -f matmul_cudaonly -e $(GENERATOR_OUTPUTS) -o $(@D) \ + target=$(TENSORCORE_TARGET) gpu_schedule=cudaonly \ + M=$(MATMUL_M) N=$(MATMUL_N) K=$(MATMUL_K) + +$(BIN)/%/matmul_tensorcore.a: $(GENERATOR_BIN)/matmul.generator + @mkdir -p $(@D) + $^ -g matmul -f matmul_tensorcore -e $(GENERATOR_OUTPUTS) -o $(@D) \ + target=$(TENSORCORE_TARGET) gpu_schedule=tensorcore \ + M=$(MATMUL_M) N=$(MATMUL_N) K=$(MATMUL_K) + +$(BIN)/%/runner: runner.cpp $(BIN)/%/matmul_cudaonly.a $(BIN)/%/matmul_tensorcore.a + @mkdir -p $(@D) + $(CXX) $(CXXFLAGS) -I$(BIN)/$* -Wall \ + -DMATMUL_M=$(MATMUL_M) -DMATMUL_N=$(MATMUL_N) -DMATMUL_K=$(MATMUL_K) \ + $^ -o $@ $(LDFLAGS) $(LIBHALIDE_LDFLAGS) + +test: $(BIN)/$(HL_TARGET)/runner + $^ + +clean: + rm -rf $(BIN) diff --git a/apps/tensorcore_matmul/matmul_generator.cpp b/apps/tensorcore_matmul/matmul_generator.cpp new file mode 100644 index 000000000000..c879551dd343 --- /dev/null +++ b/apps/tensorcore_matmul/matmul_generator.cpp @@ -0,0 +1,133 @@ +#include "Halide.h" + +namespace { + +using namespace Halide; + +enum class Schedule { + CUDA, + TensorCore, +}; + +class MatMul : public Halide::Generator { +public: + GeneratorParam gpu_schedule{ + "gpu_schedule", Schedule::TensorCore, + {{"cudaonly", Schedule::CUDA}, {"tensorcore", Schedule::TensorCore}}}; + + GeneratorParam M{"M", 1024}; + GeneratorParam N{"N", 1024}; + GeneratorParam K{"K", 1024}; + + // How many tensor core tiles of accumulator each warp holds, and how many + // warps there are per block. + GeneratorParam tiles_x{"tiles_x", 5}; + GeneratorParam tiles_y{"tiles_y", 4}; + GeneratorParam warps{"warps", 4}; + + Input> matA{"matA"}; // K x M + Input> matB{"matB"}; // N x K + + Output> output{"output"}; + + void generate() { + k = RDom(0, K, "k"); + + prod(x, y) = 0.f; + prod(x, y) += cast(matA(k, y)) * cast(matB(x, k)); + + output(x, y) = prod(x, y); + } + + void schedule() { + matA.dim(0).set_bounds(0, K).set_stride(1); + matA.dim(1).set_bounds(0, M).set_stride(K); + matB.dim(0).set_bounds(0, N).set_stride(1); + matB.dim(1).set_bounds(0, K).set_stride(N); + output.dim(0).set_bounds(0, N).set_stride(1); + output.dim(1).set_bounds(0, M).set_stride(N); + + if (gpu_schedule == Schedule::CUDA) { + // Schedule taken from the cuda_mat_mul app. + Var xi, yi, xii, yii; + + output.bound(x, 0, N) + .bound(y, 0, M) + .tile(x, y, xi, yi, 64, 16) + .tile(xi, yi, xii, yii, 4, 8) + .gpu_blocks(x, y) + .gpu_threads(xi, yi) + .unroll(xii) + .unroll(yii); + + prod.compute_at(output, xi) + .vectorize(x) + .unroll(y) + .update() + .reorder(x, y, k) + .vectorize(x) + .unroll(y) + .unroll(k, 8); + + matA.in().compute_at(prod, k).vectorize(_0).unroll(_1); + matB.in().compute_at(prod, k).vectorize(_0).unroll(_1); + } else { + // The tensor core tile shape, and how many of them each warp + // accumulates at once. Each operand tile loaded feeds tiles_x (or + // tiles_y) multiplies, so this is what gets us reuse out of the + // loads. + const int tile_x = 16, tile_y = 16, tile_k = 16; + + Var xi("xi"), yi("yi"), xt("xt"), mmxi("mmxi"), mmyi("mmyi"); + Var rxi("rxi"), ryi("ryi"); + RVar rro("rro"), rri("rri"); + + output.bound(x, 0, N) + .bound(y, 0, M) + .split(x, x, xi, tile_x * tiles_x * warps) + .split(xi, xt, xi, tile_x * tiles_x) + .split(xi, xi, mmxi, tile_x) + .split(y, y, yi, tile_y * tiles_y) + .split(yi, yi, mmyi, tile_y) + .gpu_blocks(x, y) + .gpu_threads(xt) + .reorder(mmxi, mmyi, xi, yi, xt, x, y) + .unroll(xi) + .unroll(yi) + .vectorize(mmxi) + .vectorize(mmyi); + + // The accumulators live in tensor core registers for the whole + // reduction, and are written out to memory once at the end. + prod.compute_at(output, xt) + .store_in(MemoryType::WMMAAccumulator) + .split(x, x, rxi, tile_x) + .split(y, y, ryi, tile_y) + .vectorize(rxi) + .vectorize(ryi) + .unroll(x) + .unroll(y); + + prod.update() + .split(x, x, rxi, tile_x) + .split(y, y, ryi, tile_y) + .split(k, rro, rri, tile_k) + .reorder(rri, rxi, ryi, x, y, rro) + .unroll(x) + .unroll(y) + .atomic() + .vectorize(rri) + .vectorize(rxi) + .vectorize(ryi); + } + } + +private: + Var x{"x"}, y{"y"}; + RDom k; + Func prod{"prod"}; +}; + +} // namespace + +HALIDE_REGISTER_GENERATOR(MatMul, matmul) diff --git a/apps/tensorcore_matmul/runner.cpp b/apps/tensorcore_matmul/runner.cpp new file mode 100644 index 000000000000..9a98333b42f9 --- /dev/null +++ b/apps/tensorcore_matmul/runner.cpp @@ -0,0 +1,84 @@ +#include "Halide.h" +#include "HalideBuffer.h" +#include "HalideRuntimeCuda.h" +#include "halide_benchmark.h" +#include + +#include "matmul_cudaonly.h" +#include "matmul_tensorcore.h" + +using Halide::float16_t; +using Halide::Runtime::Buffer; +using Halide::Tools::benchmark; + +namespace { + +constexpr int M = MATMUL_M, N = MATMUL_N, K = MATMUL_K; + +bool check(const Buffer &A, + const Buffer &B, + const Buffer &C, + const char *name) { + for (int y = 0; y < M; y += 97) { + for (int x = 0; x < N; x += 89) { + float ref = 0.f; + for (int k = 0; k < K; k++) { + ref += (float)A(k, y) * (float)B(x, k); + } + if (std::abs(C(x, y) - ref) > 1e-2f * std::max(1.f, std::abs(ref))) { + printf("%s: bad result at %d %d: %f != %f\n", name, x, y, C(x, y), ref); + return false; + } + } + } + return true; +} + +} // namespace + +int main(int argc, char **argv) { + const auto *interface = halide_cuda_device_interface(); + int major, minor; + if (interface->compute_capability(nullptr, &major, &minor) != 0 || + major * 10 + minor < 70) { + printf("[SKIP] Tensor cores require CUDA compute capability 7.0 or above.\n"); + return 0; + } + + Buffer A(K, M), B(N, K); + A.fill([]() { return float16_t(((float)rand() / RAND_MAX) - 0.5f); }); + B.fill([]() { return float16_t(((float)rand() / RAND_MAX) - 0.5f); }); + + Buffer C_cuda(N, M), C_tensorcore(N, M); + + matmul_cudaonly(A, B, C_cuda); + C_cuda.copy_to_host(); + if (!check(A, B, C_cuda, "cudaonly")) { + return 1; + } + + matmul_tensorcore(A, B, C_tensorcore); + C_tensorcore.copy_to_host(); + if (!check(A, B, C_tensorcore, "tensorcore")) { + return 1; + } + + // Two flops (a multiply and an add) per element of the reduction. + const double flops = 2.0 * M * N * K; + + double t_cuda = benchmark([&]() { + matmul_cudaonly(A, B, C_cuda); + C_cuda.device_sync(); + }); + double t_tensorcore = benchmark([&]() { + matmul_tensorcore(A, B, C_tensorcore); + C_tensorcore.device_sync(); + }); + + printf("cuda only: %8.3f ms %8.1f GFlop/s\n", t_cuda * 1e3, flops / t_cuda * 1e-9); + printf("tensor core: %8.3f ms %8.1f GFlop/s\n", t_tensorcore * 1e3, flops / t_tensorcore * 1e-9); + printf("speed-up: %8.2fx\n", t_cuda / t_tensorcore); + + printf("Success!\n"); + return 0; +} diff --git a/apps/tensorcore_resize/Makefile b/apps/tensorcore_resize/Makefile new file mode 100644 index 000000000000..307138395969 --- /dev/null +++ b/apps/tensorcore_resize/Makefile @@ -0,0 +1,31 @@ +include ../support/Makefile.inc + + +# The wmma instructions require compute capability 7.0 or above. +TENSORCORE_TARGET ?= host-cuda-cuda_capability_80 + +all: $(BIN)/$(HL_TARGET)/runner + +$(GENERATOR_BIN)/resize.generator: resize_generator.cpp $(GENERATOR_DEPS) + @mkdir -p $(@D) + $(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LIBHALIDE_LDFLAGS) + +$(BIN)/%/resize_cudaonly.a: $(GENERATOR_BIN)/resize.generator + @mkdir -p $(@D) + $^ -g resize -f resize_cudaonly -e $(GENERATOR_OUTPUTS) -o $(@D) \ + target=$(TENSORCORE_TARGET) gpu_schedule=cudaonly + +$(BIN)/%/resize_tensorcore.a: $(GENERATOR_BIN)/resize.generator + @mkdir -p $(@D) + $^ -g resize -f resize_tensorcore -e $(GENERATOR_OUTPUTS) -o $(@D) \ + target=$(TENSORCORE_TARGET) gpu_schedule=tensorcore + +$(BIN)/%/runner: runner.cpp $(BIN)/%/resize_cudaonly.a $(BIN)/%/resize_tensorcore.a + @mkdir -p $(@D) + $(CXX) $(CXXFLAGS) -I$(BIN)/$* -Wall $^ -o $@ $(LDFLAGS) $(LIBHALIDE_LDFLAGS) + +test: $(BIN)/$(HL_TARGET)/runner + $^ + +clean: + rm -rf $(BIN) diff --git a/apps/tensorcore_resize/README.md b/apps/tensorcore_resize/README.md new file mode 100644 index 000000000000..74388196a60c --- /dev/null +++ b/apps/tensorcore_resize/README.md @@ -0,0 +1,34 @@ +# Tensor core resize + +Resampling an image is a linear operator, so it can be written as a matrix +multiply. The matrix is enormous and almost entirely zero, so you never want to +materialize it, but each of its rows has a small number of contiguous +non-zeros. If neighbouring rows are made to share a starting column (which just +means storing a few more zeros), then a block of 16 rows of it is a small dense +matrix, and the inner loop becomes a matrix multiply. That's a large speed-up +even without tensor cores, and it lets us use tensor cores when we have them. + +This is the algorithmic difference between this generator and the one in +`apps/resize`, which starts each row of the matrix at its own column. + +## Status + +The `cudaonly` schedule works. The `tensorcore` schedule currently only works +for the first of the two stages (the resample in y). The resample in x fails +with: + +``` +Matrix multiply not recognized. [...] the matrix multiply operands are not +loads with affine indices. +``` + +The load index for that stage contains `begin_of((x / 16) * 16)`, i.e. the +starting column of this block of 16 rows of the matrix. That subexpression is +uniform across the 16 lanes of a tile, but the simplifier leaves it as +`ceil_f32` applied to `(ramp(block * 16, 1, 16) / 16) * 16` rather than folding +the divide and multiply away into a broadcast, so `is_multiramp` can't see that +it is uniform and the index isn't recognized as affine. + +Fixing this needs the simplifier to fold `ramp(a * k, 1, k) / k` down to +`broadcast(a, k)` when the ramp is nested inside another vector, after which the +lane-uniform recognition in `is_multiramp` handles the rest. diff --git a/apps/tensorcore_resize/resize_generator.cpp b/apps/tensorcore_resize/resize_generator.cpp new file mode 100644 index 000000000000..7794f5f35e26 --- /dev/null +++ b/apps/tensorcore_resize/resize_generator.cpp @@ -0,0 +1,334 @@ +#include "Halide.h" + +namespace { + +using namespace Halide; + +enum class InterpolationType { + Box, + Linear, + Cubic, + Lanczos, +}; + +enum class Schedule { + CUDA, + TensorCore, +}; + +Expr kernel_box(Expr x) { + Expr xx = abs(x); + return select(xx <= 0.5f, 1.0f, 0.0f); +} + +Expr kernel_linear(Expr x) { + Expr xx = abs(x); + return select(xx < 1.0f, 1.0f - xx, 0.0f); +} + +Expr kernel_cubic(Expr x) { + Expr xx = abs(x); + Expr xx2 = xx * xx; + Expr xx3 = xx2 * xx; + float a = -0.5f; + + return select(xx < 1.0f, (a + 2.0f) * xx3 - (a + 3.0f) * xx2 + 1, + select(xx < 2.0f, a * xx3 - 5 * a * xx2 + 8 * a * xx - 4.0f * a, + 0.0f)); +} + +Expr sinc(Expr x) { + x *= 3.14159265359f; + return sin(x) / x; +} + +constexpr int lanczos_lobes = 3; + +Expr kernel_lanczos(Expr x) { + Expr value = sinc(x) * sinc(x / lanczos_lobes); + // Take care of the singularity at zero + value = select(x == 0.0f, 1.0f, value); + // Clamp to zero out of bounds + value = select(x > lanczos_lobes || x < -lanczos_lobes, 0.0f, value); + return value; +} + +struct KernelInfo { + const char *name; + int taps; + Expr (*kernel)(Expr); +}; + +const KernelInfo kernel_info[] = { + {"box", 1, kernel_box}, + {"linear", 2, kernel_linear}, + {"cubic", 4, kernel_cubic}, + {"lanczos", 2 * lanczos_lobes, kernel_lanczos}}; + +// Resampling an image is a linear operator, so it can be written as a matrix +// multiply. The matrix is enormous and almost entirely zero, so you never want +// to materialize it, but each row of it has a small number of contiguous +// non-zeros, and if we let neighbouring rows share a starting column then a +// block of 16 rows of it is a small dense matrix. That makes the inner loop a +// matrix multiply, which is a large speed-up even without tensor cores, and +// lets us use tensor cores when we have them. +class Resize : public Halide::Generator { +public: + GeneratorParam interpolation_type{ + "interpolation_type", InterpolationType::Lanczos, + {{"box", InterpolationType::Box}, + {"linear", InterpolationType::Linear}, + {"cubic", InterpolationType::Cubic}, + {"lanczos", InterpolationType::Lanczos}}}; + + // If we statically know whether we're upsampling or downsampling, we can + // generate different pipelines (we want to reorder the resample in x and + // in y). + GeneratorParam upsample{"upsample", false}; + + GeneratorParam gpu_schedule{ + "gpu_schedule", Schedule::TensorCore, + {{"cudaonly", Schedule::CUDA}, {"tensorcore", Schedule::TensorCore}}}; + + Input> input{"input"}; + Input scale_factor{"scale_factor"}; + Output> output{"output"}; + + // The size of the blocks of the resampling matrix that we treat as dense. + static constexpr int tile = 16; + + void generate() { + // Invert the scale factor in a single place, to avoid getting slightly + // different ratios showing up in different places. + Expr inverse_scale_factor = 1.0f / scale_factor; + + // For downscaling, widen the interpolation kernel to perform lowpass + // filtering. + Expr kernel_scaling = upsample ? Expr(1.0f) : scale_factor; + Expr inverse_kernel_scaling = upsample ? Expr(1.0f) : inverse_scale_factor; + + const KernelInfo &info = kernel_info[(int)(InterpolationType)interpolation_type]; + + Expr kernel_radius = 0.5f * info.taps * inverse_kernel_scaling; + Expr kernel_taps = cast(ceil(info.taps * inverse_kernel_scaling)); + + // The (non-integer) coordinates in the source image. + Expr sourcex = (x + 0.5f) * inverse_scale_factor - 0.5f; + Expr sourcey = (y + 0.5f) * inverse_scale_factor - 0.5f; + + // For a given output coordinate, the first input coordinate it depends + // on. We can start a row of the matrix at any column we like as long + // as we store enough columns, so we use the same starting column for + // each group of `tile` rows. + auto begin_of = [&](Expr coord) { + return cast(ceil((coord + 0.5f) * inverse_scale_factor - 0.5f - kernel_radius)); + }; + + Expr beginx = begin_of((x / tile) * tile); + Expr beginy = begin_of((y / tile) * tile); + + // Moving the start of each row back like that means each row has to + // cover a longer contiguous region. + Expr extra_zeros = begin_of(tile) - begin_of(0); + + // Round the number of columns up to the next multiple of the tile size + // too, so that the reduction is a whole number of tiles. + Expr span = ((kernel_taps + extra_zeros + tile - 1) / tile) * tile; + + // Don't go off the end of the image. Those columns would be zero + // anyway. + beginx = clamp(beginx, 0, input.width() - span); + beginy = clamp(beginy, 0, input.height() - span); + + r = RDom(0, span, "r"); + + as_float(x, y, c) = cast(input(x, y, c)); + + unnormalized_kernel_x(x, k) = info.kernel((k + beginx - sourcex) * kernel_scaling); + unnormalized_kernel_y(y, k) = info.kernel((k + beginy - sourcey) * kernel_scaling); + + kernel_sum_x(x) += unnormalized_kernel_x(x, r); + kernel_sum_y(y) += unnormalized_kernel_y(y, r); + + kernel_x(x, k) = cast(unnormalized_kernel_x(x, k) / kernel_sum_x(x)); + kernel_y(y, k) = cast(unnormalized_kernel_y(y, k) / kernel_sum_y(y)); + + resized_y(x, y, c) += kernel_y(y, r) * as_float(x, r + beginy, c); + resized_x(x, y, c) += kernel_x(x, r) * resized_y(r + beginx, y, c); + + output(x, y, c) = clamp(resized_x(x, y, c), cast(0.f), cast(1.f)); + } + + void schedule() { + Var xi("xi"), yi("yi"), ki("ki"), xii("xii"), yii("yii"), xo("xo"), z("z"); + + // Precompute the sparse matrices. These are tiny compared to the + // image, so the schedule barely matters. + kernel_x.compute_root().gpu_tile(x, k, xi, ki, 32, 8); + unnormalized_kernel_x.compute_root().gpu_tile(x, k, xi, ki, 32, 8); + kernel_sum_x.in().compute_root().gpu_tile(x, xi, 32); + + kernel_y.compute_root().gpu_tile(y, k, yi, ki, 32, 8); + unnormalized_kernel_y.compute_root().gpu_tile(y, k, yi, ki, 32, 8); + kernel_sum_y.in().compute_root().gpu_tile(y, yi, 32); + + output.compute_root() + .align_bounds(x, tile) + .align_bounds(y, tile); + + if (gpu_schedule == Schedule::CUDA) { + // Resampling in y is the expensive stage for large downsamples. + // The load from the kernel doesn't depend on x or c, and the load + // from the image doesn't depend on y % tile, so we schedule it + // like a matrix multiply. + resized_y.in() + .compute_root() + .align_bounds(x, tile) + .align_bounds(y, tile) + .reorder(c, x, y) + .unroll(c) + .gpu_tile(x, y, xi, yi, 32, 16, TailStrategy::RoundUp) + .tile(xi, yi, xii, yii, 2, 4) + .unroll(xii) + .unroll(yii); + resized_y + .compute_at(resized_y.in(), xi) + .unroll(c) + .unroll(x) + .unroll(y) + .update() + .reorder(x, y, c, r) + .unroll(c) + .unroll(x) + .unroll(y); + as_float.compute_at(resized_y, c).vectorize(x).vectorize(y); + kernel_y.in().compute_at(resized_y, r).vectorize(y).vectorize(k); + + // After downsampling in y it's hard to fill the machine, so use + // smaller tiles and map color channels to gpu threads. + output + .gpu_threads(c) + .gpu_tile(x, y, xi, yi, 32, 4, TailStrategy::RoundUp) + .reorder(xi, yi, c, x, y) + .tile(xi, yi, xii, yii, 2, 2) + .vectorize(xii) + .unroll(yii); + + resized_x + .compute_at(output, xi) + .unroll(c) + .unroll(x) + .unroll(y) + .update() + .reorder(x, y, c, r) + .unroll(c) + .unroll(x) + .unroll(y); + resized_y.in().in().compute_at(resized_x, c).vectorize(y); + kernel_x.in().compute_at(resized_x, r).vectorize(x).vectorize(k); + } else { + // The tensor core instructions want the reduction dimension of + // each operand dense in memory. + kernel_x.reorder_storage(k, x); + kernel_y.reorder_storage(k, y); + + Var xio("xio"); + resized_y.in() + .compute_root() + .align_bounds(x, tile) + .align_bounds(y, tile) + .tile(x, y, xi, yi, 32, 16, TailStrategy::RoundUp) + .unroll(c) + .split(xi, xi, xii, 32) + .split(xi, xio, xi, 1) + .gpu_threads(xio) + .split(yi, yi, yii, 8) + .reorder(xii, yii, c, yi, xi, xio, x, y) + .vectorize(xii) + .vectorize(yii) + .unroll(yi) + .unroll(xi) + .gpu_blocks(x, y); + + // An 8x32 tile of accumulator, reducing 16 taps at a time. + resized_y.compute_at(resized_y.in(), xio) + .store_in(MemoryType::WMMAAccumulator) + .unroll(c) + .vectorize(x, 32) + .unroll(x) + .vectorize(y, 8) + .unroll(y) + .update() + .atomic() + .unroll(c) + .vectorize(x, 32) + .unroll(x) + .vectorize(y, 8) + .unroll(y) + .vectorize(r, tile) + .reorder(y, c, x, r); + + output + .tile(x, y, xi, yi, tile, tile, TailStrategy::RoundUp) + .reorder(yi, xi, x, y, c) + .gpu_blocks(x, y, c) + .split(yi, yi, yii, 2) + .fuse(xi, yii, z) + .gpu_lanes(z) + .unroll(yi); + + resized_x.in() + .compute_at(output, x) + .vectorize(x) + .vectorize(y); + + RVar ri("ri"), ro("ro"); + resized_x + .store_in(MemoryType::WMMAAccumulator) + .compute_at(resized_x.in(), c) + .vectorize(x) + .vectorize(y) + .update() + .atomic() + .split(r, ro, ri, tile) + .reorder(ri, x, y, ro) + .vectorize(x) + .vectorize(y) + .vectorize(ri); + + // An extra layer of staging, because we're not necessarily aligned + // in x. + resized_y.in() + .in() + .compute_at(output, x) + .store_in(MemoryType::GPUShared) + .split(x, xo, xi, 32, TailStrategy::RoundUp) + .gpu_lanes(xi); + } + + output.dim(0).set_min(0); + output.dim(1).set_min(0); + output.dim(2).set_bounds(0, 3); + input.dim(0).set_min(0); + input.dim(1).set_min(0); + input.dim(2).set_bounds(0, 3); + } + +private: + Var x{"x"}, y{"y"}, c{"c"}, k{"k"}; + RDom r; + + Func as_float{"as_float"}, + resized_x{"resized_x"}, + resized_y{"resized_y"}, + unnormalized_kernel_x{"unnormalized_kernel_x"}, + unnormalized_kernel_y{"unnormalized_kernel_y"}, + kernel_x{"kernel_x"}, + kernel_y{"kernel_y"}, + kernel_sum_x{"kernel_sum_x"}, + kernel_sum_y{"kernel_sum_y"}; +}; + +} // namespace + +HALIDE_REGISTER_GENERATOR(Resize, resize) diff --git a/apps/tensorcore_resize/runner.cpp b/apps/tensorcore_resize/runner.cpp new file mode 100644 index 000000000000..21429b538fb1 --- /dev/null +++ b/apps/tensorcore_resize/runner.cpp @@ -0,0 +1,72 @@ +#include "Halide.h" +#include "HalideBuffer.h" +#include "HalideRuntimeCuda.h" +#include "halide_benchmark.h" +#include + +#include "resize_cudaonly.h" +#include "resize_tensorcore.h" + +using Halide::float16_t; +using Halide::Runtime::Buffer; +using Halide::Tools::benchmark; + +int main(int argc, char **argv) { + const auto *interface = halide_cuda_device_interface(); + int major, minor; + if (interface->compute_capability(nullptr, &major, &minor) != 0 || + major * 10 + minor < 70) { + printf("[SKIP] Tensor cores require CUDA compute capability 7.0 or above.\n"); + return 0; + } + + const int in_w = 3840, in_h = 2160; + const float scale_factor = 0.25f; + const int out_w = (int)(in_w * scale_factor), out_h = (int)(in_h * scale_factor); + + Buffer input(in_w, in_h, 3); + input.fill([]() { return float16_t((float)rand() / RAND_MAX); }); + + Buffer out_cuda(out_w, out_h, 3), out_tensorcore(out_w, out_h, 3); + + resize_cudaonly(input, scale_factor, out_cuda); + resize_tensorcore(input, scale_factor, out_tensorcore); + out_cuda.copy_to_host(); + out_tensorcore.copy_to_host(); + + // The two schedules compute the same thing, but accumulate in a different + // order in half precision, so only compare them approximately. + int bad = 0; + for (int c = 0; c < 3; c++) { + for (int y = 0; y < out_h; y++) { + for (int x = 0; x < out_w; x++) { + float a = (float)out_cuda(x, y, c), b = (float)out_tensorcore(x, y, c); + if (std::abs(a - b) > 5e-3f) { + if (bad++ < 10) { + printf("Mismatch at %d %d %d: %f != %f\n", x, y, c, a, b); + } + } + } + } + } + if (bad) { + printf("Failed with %d mismatches\n", bad); + return 1; + } + + double t_cuda = benchmark([&]() { + resize_cudaonly(input, scale_factor, out_cuda); + out_cuda.device_sync(); + }); + double t_tensorcore = benchmark([&]() { + resize_tensorcore(input, scale_factor, out_tensorcore); + out_tensorcore.device_sync(); + }); + + printf("cuda only: %8.3f ms\n", t_cuda * 1e3); + printf("tensor core: %8.3f ms\n", t_tensorcore * 1e3); + printf("speed-up: %8.2fx\n", t_cuda / t_tensorcore); + + printf("Success!\n"); + return 0; +} diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 7c2fc39512a4..a2d211cbb235 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -118,6 +118,7 @@ target_sources( Extern.h ExternFuncArgument.h ExtractTileOperations.h + ExtractWMMAOperations.h FastIntegerDivide.h FindCalls.h FindIntrinsics.h @@ -300,6 +301,7 @@ target_sources( Error.cpp Expr.cpp ExtractTileOperations.cpp + ExtractWMMAOperations.cpp FastIntegerDivide.cpp FindCalls.cpp FindIntrinsics.cpp diff --git a/src/CanonicalizeGPUVars.cpp b/src/CanonicalizeGPUVars.cpp index 609323f8e5dd..25e2ed6781bf 100644 --- a/src/CanonicalizeGPUVars.cpp +++ b/src/CanonicalizeGPUVars.cpp @@ -6,6 +6,7 @@ #include "IR.h" #include "IRMutator.h" #include "Substitute.h" +#include "Util.h" namespace Halide { namespace Internal { @@ -38,12 +39,19 @@ class CountGPUBlocksThreads : public IRVisitor { // we're inside of, respectively. Lanes loops also count as threads loops. int nb = 0, nt = 0, nl = 0; + // Whether we're already inside a lane dimension. Every lane dimension maps + // to the innermost thread dimension, so one nested inside another is the + // same dimension, not a new one. + bool in_lanes = false; + void visit(const For *op) override { // Figure out how much to increment each counter by based on the loop // type. int db = op->for_type == ForType::GPUBlock; - int dl = op->for_type == ForType::GPULane; + int dl = (op->for_type == ForType::GPULane) && !in_lanes; int dt = op->for_type == ForType::GPUThread; + ScopedValue old_in_lanes(in_lanes, + in_lanes || op->for_type == ForType::GPULane); // The threads counter includes lanes loops dt += dl; @@ -67,6 +75,22 @@ class CountGPUBlocksThreads : public IRVisitor { nt -= dt; } + void visit(const Realize *op) override { + // extract_wmma_operations will wrap the statements that touch this + // allocation in loops over the lanes of a warp, so count it as a lane + // dimension. + const bool wmma = op->memory_type == MemoryType::WMMAAccumulator; + int dl = wmma && !in_lanes; + ScopedValue old_in_lanes(in_lanes, in_lanes || wmma); + nl += dl; + nt += dl; + nlanes = std::max(nl, nlanes); + nthreads = std::max(nt, nthreads); + IRVisitor::visit(op); + nl -= dl; + nt -= dl; + } + public: // The maximum values hit by the counters above, which tells us the nesting // depth of each type of loop within a Stmt. diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 962c650a77d5..78f51138ee26 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -7,6 +7,7 @@ #include "ConciseCasts.h" #include "Debug.h" #include "ExprUsesVar.h" +#include "ExtractWMMAOperations.h" #include "IREquality.h" #include "IRMatch.h" #include "IRMutator.h" @@ -15,6 +16,7 @@ #include "LLVM_Headers.h" #include "LLVM_Runtime_Linker.h" #include "ModulusRemainder.h" +#include "MultiRamp.h" #include "Simplify.h" #include "Solve.h" #include "Target.h" @@ -121,6 +123,17 @@ class CodeGen_PTX_Dev : public CodeGen_LLVM, public CodeGen_GPU_Dev { /** Wait for any asynchronous copies issued so far to have landed. */ void wait_for_async_copies(); + /** Emit calls to the nvvm warp-level matrix multiply-accumulate + * intrinsics that drive the tensor cores. */ + // @{ + void codegen_wmma(const Call *op); + void codegen_wmma_store(const Store *op); + void split_fragment(const Expr &e, std::vector &args); + llvm::Value *call_wmma_intrinsic(const std::string &name, + const std::vector &args, + const std::vector &overloads); + // @} + bool supports_atomic_add(const Type &t) const override; }; @@ -312,6 +325,16 @@ void CodeGen_PTX_Dev::visit(const Call *op) { return; } + if (is_wmma_intrinsic(op)) { + codegen_wmma(op); + return; + } + + internal_assert(!op->is_intrinsic(Call::wmma_fragment_to_matrix_d) && + !op->is_intrinsic(Call::wmma_lane_owns)) + << "A tensor core accumulator store was broken apart during lowering. " + << op->name << " only has meaning as part of one.\n"; + // TODO: It would be better if CodeGen_LLVM could handle overloaded intrin calls by default. value = call_overloaded_intrin(op->type, op->name, op->args); if (!value) { @@ -319,6 +342,166 @@ void CodeGen_PTX_Dev::visit(const Call *op) { } } +namespace { + +WMMAMatrixLayout matrix_in_memory(const string &name, const MultiRamp &mr, int rows, int cols) { + WMMAMatrixLayout result; + user_assert(wmma_matrix_layout(mr, rows, cols, &result)) + << "The memory a tensor core instruction moves a matrix of " << name + << " to or from is not a dense tile by the time it reaches the backend. " + << "This happens when the allocation is striped across threads, which " + << "occurs for a shared memory allocation made inside the loop over GPU " + << "threads. Compute it at a loop outside the threads instead.\n"; + return result; +} + +} // namespace + +void CodeGen_PTX_Dev::split_fragment(const Expr &e, vector &args) { + // One llvm value per 32-bit register. + Value *v = codegen(e); + const int lanes_per_reg = 32 / e.type().bits(); + for (int i = 0; i < e.type().lanes() / lanes_per_reg; i++) { + args.push_back(lanes_per_reg == 1 ? + builder->CreateExtractElement(v, i) : + slice_vector(v, i * lanes_per_reg, lanes_per_reg)); + } +} + +Value *CodeGen_PTX_Dev::call_wmma_intrinsic(const std::string &name, + const vector &args, + const vector &overloads) { + llvm::Intrinsic::ID id = llvm::Intrinsic::lookupIntrinsicID(name); + internal_assert(id != llvm::Intrinsic::not_intrinsic) + << "Could not find the nvvm intrinsic " << name << "\n"; + llvm::Function *fn = llvm::Intrinsic::getOrInsertDeclaration(module.get(), id, overloads); + return builder->CreateCall(fn, args); +} + +void CodeGen_PTX_Dev::codegen_wmma(const Call *op) { + // The nvvm wmma intrinsics take and return fragments as a flat list of + // 32-bit registers, packaged up as a literal struct. We represent them in + // Halide IR as vectors, so most of the work here is repacking. + auto get_int_arg = [&](int i) { + auto v = as_const_int(op->args[i]); + internal_assert(v) << "Expected a constant integer argument to " << op->name << "\n"; + return (int)*v; + }; + + const int M = get_int_arg(0), N = get_int_arg(1), K = get_int_arg(2); + const char *layouts[] = {"row", "col"}; + + std::ostringstream name; + name << "llvm.nvvm.wmma.m" << M << "n" << N << "k" << K << "."; + + vector args; + vector overloads; + + if (op->is_intrinsic(Call::wmma_mma)) { + // The two type suffixes are the types of the d and c operands, which + // for us are always the same. + const char *suffix = op->type.bits() == 32 ? "f32" : "f16"; + name << "mma." << layouts[get_int_arg(3)] << "." << layouts[get_int_arg(4)] + << "." << suffix << "." << suffix; + split_fragment(op->args[5], args); + split_fragment(op->args[6], args); + split_fragment(op->args[7], args); + } else { + // The a operand is M x K, the b operand is K x N, and the accumulator + // is M x N. + const bool is_a = op->is_intrinsic(Call::wmma_matrix_to_fragment_a); + const bool is_b = op->is_intrinsic(Call::wmma_matrix_to_fragment_b); + // The simplifier is free to have rewritten the load of the matrix into + // a dense load followed by a transpose, which is what happens to a + // column-major matrix. is_load_of_multiramp undoes that. + const Expr &arg = op->args[wmma_matrix_arg(op)]; + MultiRamp mr; + const Load *matrix = is_load_of_multiramp(arg, Scope::empty_scope(), &mr); + user_assert(matrix && matrix->type.element_of() == arg.type().element_of()) + << "The matrix a tensor core instruction takes a fragment out of is not a " + << "load with an affine index by the time it reaches the backend.\n"; + WMMAMatrixLayout mem = matrix_in_memory(matrix->name, mr, + is_b ? K : M, is_a ? K : N); + // The a and b operands are always 16-bit; an accumulator may be either. + const char *type_suffix = + is_a || is_b ? "f16" : (op->type.bits() == 32 ? "f32" : "f16"); + name << "load." << (is_a ? "a" : is_b ? "b" : + "c") + << "." << (mem.row_major ? "row" : "col") << ".stride." << type_suffix; + + Value *ptr = codegen_buffer_pointer(matrix->name, matrix->type.element_of(), mem.base); + overloads.push_back(ptr->getType()); + args.push_back(ptr); + args.push_back(codegen(cast(Int(32), mem.stride))); + } + + Value *result = call_wmma_intrinsic(name.str(), args, overloads); + + // Reassemble the returned struct into a Halide vector. + llvm::Type *result_type = llvm_type_of(op->type); + const int num_regs = op->type.bits() * op->type.lanes() / 32; + if (op->type.bits() == 32) { + value = UndefValue::get(result_type); + for (int i = 0; i < num_regs; i++) { + value = builder->CreateInsertElement(value, builder->CreateExtractValue(result, i), i); + } + } else { + vector regs; + regs.reserve(num_regs); + for (int i = 0; i < num_regs; i++) { + regs.push_back(builder->CreateExtractValue(result, i)); + } + value = concat_vectors(regs); + } + internal_assert(value->getType() == result_type) + << "Unexpected result type from " << name.str() << "\n"; +} + +void CodeGen_PTX_Dev::codegen_wmma_store(const Store *op) { + // Each lane writes the entries of the matrix that it holds, which the + // predicate describes and which one wmma store instruction does for the + // whole warp. + // A store to a column-major matrix arrives as a dense store of the + // transpose of it, because the simplifier rewrites stores to make them + // dense. Undoing that gives back the store as the extraction pass wrote it, + // and the column-major layout falls out of the index as usual. + Expr index; + const Call *inflate = peel_store_permutations(op, &index).as(); + internal_assert(inflate && inflate->args.size() == 4); + Expr predicate = op->predicate; + while (const Shuffle *shuffle = predicate.as()) { + predicate = shuffle->vectors[0]; + } + internal_assert(predicate.as() && + predicate.as()->is_intrinsic(Call::wmma_lane_owns)) + << "A store of a tensor core accumulator lost its predicate\n"; + + auto get_int_arg = [&](int i) { + auto v = as_const_int(inflate->args[i]); + internal_assert(v); + return (int)*v; + }; + const int M = get_int_arg(0), N = get_int_arg(1), K = get_int_arg(2); + const Expr &fragment = inflate->args[3]; + + MultiRamp mr; + internal_assert(is_multiramp(index, Scope::empty_scope(), &mr)); + WMMAMatrixLayout mem = matrix_in_memory(op->name, mr, M, N); + + std::ostringstream name; + name << "llvm.nvvm.wmma.m" << M << "n" << N << "k" << K << ".store.d." + << (mem.row_major ? "row" : "col") << ".stride." + << (fragment.type().bits() == 32 ? "f32" : "f16"); + + Value *ptr = codegen_buffer_pointer(op->name, op->value.type().element_of(), mem.base); + vector args{ptr}; + vector overloads{ptr->getType()}; + split_fragment(fragment, args); + args.push_back(codegen(cast(Int(32), mem.stride))); + + call_wmma_intrinsic(name.str(), args, overloads); +} + string CodeGen_PTX_Dev::simt_intrinsic(const string &name) { if (ends_with(name, gpu_thread_name(0))) { return "llvm.nvvm.read.ptx.sreg.tid.x"; @@ -577,6 +760,11 @@ void CodeGen_PTX_Dev::wait_for_async_copies() { } void CodeGen_PTX_Dev::visit(const Store *op) { + if (is_wmma_matrix_store(op)) { + codegen_wmma_store(op); + return; + } + // Issue atomic store if we are inside an Atomic node. if (emit_atomic_stores) { user_assert(is_const_one(op->predicate)) << "Atomic update does not support predicated store.\n"; diff --git a/src/Deserialization.cpp b/src/Deserialization.cpp index 72e975a12a41..016df76cf42a 100644 --- a/src/Deserialization.cpp +++ b/src/Deserialization.cpp @@ -190,6 +190,8 @@ MemoryType Deserializer::deserialize_memory_type(Serialize::MemoryType memory_ty return MemoryType::VTCM; case Serialize::MemoryType::AMXTile: return MemoryType::AMXTile; + case Serialize::MemoryType::WMMAAccumulator: + return MemoryType::WMMAAccumulator; default: user_error << "unknown memory type " << (int)memory_type << "\n"; return MemoryType::Auto; diff --git a/src/Expr.h b/src/Expr.h index 5a800e7bd625..5c22e27d911b 100644 --- a/src/Expr.h +++ b/src/Expr.h @@ -413,6 +413,13 @@ enum class MemoryType { * global buffer, because that is all the hardware can do. On GPU APIs with * no such instruction this is ordinary shared memory. */ GPUSharedAsync, + + /** An NVIDIA tensor core accumulator fragment. The storage is striped + * across the registers of the 32 lanes of a warp in a layout that is not + * architecturally specified, so the only legal accesses are the ones + * recognized by the WMMA lowering pass: zero-initialization, accumulation + * of a matrix multiply, and copying the tile out to memory. */ + WMMAAccumulator, }; /** Whether a MemoryType places an allocation in GPU shared memory. */ @@ -427,7 +434,7 @@ inline bool is_gpu_shared(MemoryType t) { * dedicated lowering pass that requires the original 2D-shaped loads * and stores to remain intact. */ inline bool is_tile_memory_type(MemoryType t) { - return t == MemoryType::AMXTile; + return t == MemoryType::AMXTile || t == MemoryType::WMMAAccumulator; } namespace Internal { diff --git a/src/ExtractWMMAOperations.cpp b/src/ExtractWMMAOperations.cpp new file mode 100644 index 000000000000..fa34b7cb7beb --- /dev/null +++ b/src/ExtractWMMAOperations.cpp @@ -0,0 +1,618 @@ +#include "ExtractWMMAOperations.h" + +#include "CanonicalizeGPUVars.h" +#include "FindIntrinsics.h" +#include "IREquality.h" +#include "IRMutator.h" +#include "IROperator.h" +#include "MultiRamp.h" +#include "Simplify.h" +#include "Util.h" + +/** \file Support extraction of NVIDIA tensor core (wmma) instructions. + * + * The wmma instructions are warp-level: the 32 lanes of a warp cooperate to + * hold a tile of a matrix, and each instruction is executed by the warp as a + * whole. The layout of the tile across the lanes' registers is not specified + * by the architecture, so the only way to get data into or out of a tile is + * with the wmma load and store instructions, which move a tile between the + * registers of a warp and a 2D array in memory. + * + * Accordingly this pass recognizes exactly three operations on a + * WMMAAccumulator allocation: + * + * 1) Zero-initialization. This one is layout-independent, so it stays a plain + * store of zero to the (shrunken) allocation. + * 2) Accumulation of a matrix multiply, which becomes a pair of wmma loads + * feeding a wmma mma. + * 3) Copying a tile out to memory, which becomes a wmma store. + * + * Anything else is an error. Each of those operations is wrapped in a loop over + * the 32 lanes of a warp, because nothing in the schedule says the tile is + * spread over a warp - that's a consequence of asking for tensor core storage. + */ + +namespace Halide { +namespace Internal { + +using std::string; +using std::vector; + +namespace { + +// The tile shapes the hardware supports, in the order we try them. +struct Shape { + int M, N, K; +}; +const Shape supported_shapes[] = {{16, 16, 16}, {32, 8, 16}, {8, 32, 16}}; + +// A matrix operand is stored either with its rows contiguous or its columns +// contiguous. +enum class Layout { + Row, + Col, +}; + +// Every warp is 32 lanes, and every tile shape we support holds 256 elements, +// so each lane holds 8 of them. +constexpr int warp_lanes = 32; +constexpr int fragment_elements = 8; + +// The Halide type we use to represent an a or b fragment. Every operand +// fragment is 8 32-bit registers per lane, which is more matrix elements than +// there are for some shapes, because the hardware replicates elements across +// lanes for those. +Type fragment_type(Type element_type) { + return element_type.with_lanes(16); +} + +// The Halide type we use to represent an accumulator fragment. +Type accumulator_fragment_type(Type element_type) { + return element_type.with_lanes(fragment_elements); +} + +// One operand of a matrix multiply, described in the canonical [K, N, M] +// (innermost first) coordinate system. +struct Operand { + const Load *load = nullptr; + MultiRamp mr; + vector strides; +}; + +// Try to interpret an operand as the left-hand side of the multiply: an M x K +// matrix that doesn't depend on the N coordinate. +bool is_lhs(const Operand &op, Layout *layout, Expr *stride) { + if (!is_const_zero(op.strides[1])) { + return false; + } + if (is_const_one(op.strides[0])) { + *layout = Layout::Row; + *stride = op.strides[2]; + } else if (is_const_one(op.strides[2])) { + *layout = Layout::Col; + *stride = op.strides[0]; + } else { + return false; + } + return true; +} + +// Try to interpret an operand as the right-hand side of the multiply: a K x N +// matrix that doesn't depend on the M coordinate. +bool is_rhs(const Operand &op, Layout *layout, Expr *stride) { + if (!is_const_zero(op.strides[2])) { + return false; + } + if (is_const_one(op.strides[1])) { + *layout = Layout::Row; + *stride = op.strides[0]; + } else if (is_const_one(op.strides[0])) { + *layout = Layout::Col; + *stride = op.strides[1]; + } else { + return false; + } + return true; +} + +// The wmma instructions are warp-level, so every statement that touches an +// accumulator has to be executed by all 32 lanes of a warp. Nothing in the +// schedule says so - it's a consequence of asking for tensor core storage - so +// the loop over lanes is introduced here. +Stmt in_lane_loop(Stmt s) { + return For::make(unique_name("wmma_lane") + gpu_thread_name(0), + 0, warp_lanes - 1, ForType::GPULane, Partition::Never, + DeviceAPI::CUDA, std::move(s)); +} + +// The matrix a wmma load or store moves, described as a Load of every element +// of it, with the lanes in row-major order regardless of how it is laid out in +// memory. A row-major matrix has consecutive columns one element apart and its +// rows stride apart; a column-major one is the other way around, which changes +// the strides of the index but not the order of the lanes. So the argument is +// the matrix value, and the memory layout is recoverable from the index. +Expr make_matrix_index(const Expr &base, int rows, int cols, Layout layout, + const Expr &stride) { + Expr one = make_one(base.type()); + Expr col_stride = layout == Layout::Row ? one : stride; + Expr row_stride = layout == Layout::Row ? stride : one; + return Ramp::make(Ramp::make(base, col_stride, cols), + Broadcast::make(row_stride, cols), rows); +} + +Expr make_matrix_address(const string &name, Type element_type, const Expr &base, + int rows, int cols, Layout layout, const Expr &stride, + const Buffer<> &image, const Parameter ¶m) { + Expr index = make_matrix_index(base, rows, cols, layout, stride); + const int lanes = rows * cols; + return Load::make(element_type.with_lanes(lanes), name, index, image, param, + const_true(lanes), ModulusRemainder()); +} + +// The shape of the matrix each fragment is taken out of. +void fragment_matrix_shape(Call::IntrinsicOp intrin, const Shape &shape, + int *rows, int *cols) { + *rows = intrin == Call::wmma_matrix_to_fragment_b ? shape.K : shape.M; + *cols = intrin == Call::wmma_matrix_to_fragment_a ? shape.K : shape.N; +} + +Expr make_matrix_to_fragment(Call::IntrinsicOp intrin, const Shape &shape, Layout layout, + const Load *load, const Expr &base, const Expr &stride) { + int rows, cols; + fragment_matrix_shape(intrin, shape, &rows, &cols); + Expr address = make_matrix_address(load->name, load->type.element_of(), base, + rows, cols, layout, stride, load->image, load->param); + Type type = intrin == Call::wmma_matrix_to_fragment_c ? + accumulator_fragment_type(load->type.element_of()) : + fragment_type(load->type.element_of()); + return Call::make(type, intrin, {shape.M, shape.N, shape.K, std::move(address)}, + Call::Intrinsic); +} + +struct Matmul { + Stmt stmt; + Shape shape; +}; + +Matmul convert_to_matmul(const Store *op, const string &new_name) { + // We expect the pattern: + // + // out[idx] = reduce_add(widen(lhs[multiramp]) * widen(rhs[multiramp])) + out[idx] + // + // Though either operand may have been hoisted out to a broadcast or had a + // lane permutation left on it by vectorization. + + auto fail = [&](const char *reason) -> Matmul { + user_error << "Matrix multiply not recognized. Store to a WMMAAccumulator " + << "allocation must be a zero-initialization or a sum of a vector " + << "reduce op and a load from the same allocation. In the following " + << "store, " << reason << ".\n" + << Stmt(op); + return Matmul{}; + }; + + // Peel lets + vector> peeled_lets; + Expr value = op->value; + while (const Let *let = value.as()) { + peeled_lets.emplace_back(let->name, let->value); + value = let->body; + } + + // The RHS must be an add + const auto *add = value.as(); + if (!add) { + return fail("the right-hand-side is not an add"); + } + + // The add must be between a vector reduce and a load. The simplifier will + // have placed the vector reduce to the left, due to canonicalization of + // commutative ops. + const auto *reduce = add->a.as(); + if (!reduce || reduce->op != VectorReduce::Add) { + return fail("the right-hand-side is not a vector reduction plus a load"); + } + + // The load must be to the same addresses as the store (i.e. this is a +=) + const auto *load = add->b.as(); + if (!load || load->name != op->name || !equal(load->index, op->index)) { + return fail("the right-hand-side load is not from the same address as the store"); + } + + // There must be no predicate on the load or store + if (!is_const_one(load->predicate) || !is_const_one(op->predicate)) { + return fail("the load or store is predicated"); + } + + if (!reduce->type.is_float() || + !(reduce->type.bits() == 32 || reduce->type.bits() == 16)) { + return fail("the accumulator type is not 32-bit or 16-bit float"); + } + + // The vector reduce must be of a widening multiply. FindIntrinsics does + // not lift float widening muls, so we just expect a multiply of two casts. + Expr reduce_value = simplify(lower_intrinsics(reduce->value)); + const auto *mul = reduce_value.as(); + if (!mul) { + return fail("the vector reduction is not of a multiply"); + } + // Under the casts, broadcasts and lane permutations that vectorization may + // have left on each operand there must be a load. + Scope empty_scope; + Operand lhs_op, rhs_op; + lhs_op.load = is_load_of_multiramp(mul->a, empty_scope, &lhs_op.mr); + rhs_op.load = is_load_of_multiramp(mul->b, empty_scope, &rhs_op.mr); + if (!lhs_op.load || !rhs_op.load) { + return fail("the matrix multiply operands are not loads with affine indices"); + } + if (!is_const_one(lhs_op.load->predicate) || !is_const_one(rhs_op.load->predicate)) { + return fail("the matrix multiply operands are predicated loads"); + } + if (lhs_op.load->type.element_of() != Float(16) || + rhs_op.load->type.element_of() != Float(16)) { + return fail("the matrix multiply operands are not both float16"); + } + + // In a matrix multiply with row-major inputs and outputs, the algorithm + // looks like: + // + // C(j, i) += A(k, i) * B(j, k) + // + // (Recall that for matrices where the rows are stored densely in memory, + // Halide is indexed col-major.) The canonical loop nest order, from + // innermost out, is k, j, i, which is the coordinate system we compare the + // operands' access patterns in. In that coordinate system, i indexes the M + // rows of the output, j indexes its N columns, and k indexes the reduction. + int MN = reduce->type.lanes(); + int K = reduce->value.type().lanes() / MN; + + // Deduce which operand is which and what tile shape this is by trying each + // supported shape and seeing which one the access patterns fit. + const Shape *shape = nullptr; + Layout lhs_layout = Layout::Row, rhs_layout = Layout::Row; + Expr lda, ldb; + for (const Shape &candidate : supported_shapes) { + if (candidate.M * candidate.N != MN || candidate.K != K) { + continue; + } + vector canonical_shape{candidate.K, candidate.N, candidate.M}; + if (!lhs_op.mr.strides_for_shape(canonical_shape, &lhs_op.strides) || + !rhs_op.mr.strides_for_shape(canonical_shape, &rhs_op.strides)) { + continue; + } + if (is_lhs(lhs_op, &lhs_layout, &lda) && is_rhs(rhs_op, &rhs_layout, &ldb)) { + shape = &candidate; + break; + } + if (is_lhs(rhs_op, &lhs_layout, &lda) && is_rhs(lhs_op, &rhs_layout, &ldb)) { + std::swap(lhs_op, rhs_op); + shape = &candidate; + break; + } + } + + if (!shape) { + return fail("the operands' access patterns do not describe a matrix " + "multiply of a tile shape the tensor cores support (16x16x16, " + "32x8x16, or 8x32x16)"); + } + + // Build the wmma intrinsics. + Expr a = make_matrix_to_fragment(Call::wmma_matrix_to_fragment_a, *shape, lhs_layout, + lhs_op.load, lhs_op.mr.base, lda); + Expr b = make_matrix_to_fragment(Call::wmma_matrix_to_fragment_b, *shape, rhs_layout, + rhs_op.load, rhs_op.mr.base, ldb); + + Type acc_type = accumulator_fragment_type(reduce->type.element_of()); + Expr frag_idx = Ramp::make(0, 1, fragment_elements); + Expr c = Load::make(acc_type, new_name, frag_idx, {}, {}, + const_true(fragment_elements), {}); + + Expr mma = Call::make(acc_type, Call::wmma_mma, + {shape->M, shape->N, shape->K, + (int)lhs_layout, (int)rhs_layout, + std::move(a), std::move(b), std::move(c)}, + Call::Intrinsic); + + Stmt store = in_lane_loop( + Store::make(new_name, std::move(mma), frag_idx, Parameter(), + const_true(fragment_elements), ModulusRemainder())); + for (auto &[name, v] : reverse_view(peeled_lets)) { + store = LetStmt::make(name, std::move(v), store); + } + return {std::move(store), *shape}; +} + +// Whether a store to an accumulator is its initialization, as opposed to a +// matrix multiply accumulating into it. +bool is_initialization(const Store *op, const string &tile_name) { + if (is_const_zero(op->value)) { + return true; + } + MultiRamp mr; + const Load *load = is_load_of_multiramp(op->value, Scope::empty_scope(), &mr); + return load && load->name != tile_name; +} + +Stmt convert_to_init(const Store *op, const string &new_name, const Shape &shape) { + Type element_type = op->value.type().element_of(); + Expr value; + if (is_const_zero(op->value)) { + // Zeroing an accumulator is layout-independent, so it doesn't need an + // instruction - the registers just get set to zero. + value = make_zero(accumulator_fragment_type(element_type)); + } else { + auto fail = [&](const char *reason) { + user_error << "Initialization of a tensor core accumulator not supported. " + << reason << ".\n" + << Stmt(op); + return Expr{}; + }; + + MultiRamp mr; + const Load *matrix = is_load_of_multiramp(op->value, Scope::empty_scope(), &mr); + internal_assert(matrix); // is_initialization checked this + if (matrix->type.element_of() != element_type) { + value = fail("An accumulator can only be initialized from a matrix of the " + "same type, because the hardware does not convert on the way in"); + } else if (!is_const_one(matrix->predicate)) { + value = fail("The load is predicated"); + } else { + WMMAMatrixLayout mem; + if (!wmma_matrix_layout(mr, shape.M, shape.N, &mem)) { + value = fail("The matrix loaded from is not a dense tile of the right shape"); + } else { + value = make_matrix_to_fragment( + Call::wmma_matrix_to_fragment_c, shape, + mem.row_major ? Layout::Row : Layout::Col, matrix, mem.base, mem.stride); + } + } + } + Expr frag_idx = Ramp::make(0, 1, fragment_elements); + return in_lane_loop( + Store::make(new_name, std::move(value), frag_idx, Parameter(), + const_true(fragment_elements), ModulusRemainder())); +} + +Stmt convert_to_tile_store(const Store *op, const Expr &store_index, + const string &new_name, const Shape &shape) { + auto fail = [&](const char *reason) { + user_error << "Store of a tensor core accumulator to memory not supported. " + << reason << ".\n" + << Stmt(op); + return Stmt{}; + }; + + if (!is_const_one(op->predicate)) { + return fail("The store has a predicate"); + } + MultiRamp mr; + if (!is_multiramp(store_index, Scope::empty_scope(), &mr)) { + return fail("The store index is not affine"); + } + WMMAMatrixLayout mem; + if (!wmma_matrix_layout(mr, shape.M, shape.N, &mem)) { + return fail("The store is not to a dense tile of the deduced matrix shape"); + } + Layout layout = mem.row_major ? Layout::Row : Layout::Col; + + // Each lane writes the entries of the matrix that it holds. The index + // enumerates the matrix in row-major order, which is the lane order of the + // accumulator. + Expr index = make_matrix_index(mem.base, shape.M, shape.N, layout, mem.stride); + Type element_type = op->value.type().element_of(); + Expr frag = Load::make(accumulator_fragment_type(element_type), new_name, + Ramp::make(0, 1, fragment_elements), {}, {}, + const_true(fragment_elements), {}); + const int lanes = shape.M * shape.N; + Expr matrix = Call::make(element_type.with_lanes(lanes), Call::wmma_fragment_to_matrix_d, + {shape.M, shape.N, shape.K, std::move(frag)}, + Call::Intrinsic); + Expr owned = Call::make(UInt(1, lanes), Call::wmma_lane_owns, + {shape.M, shape.N, shape.K}, Call::Intrinsic); + return in_lane_loop(Store::make(op->name, std::move(matrix), std::move(index), + op->param, std::move(owned), ModulusRemainder(), + op->is_streaming)); +} + +class ExtractWMMAOperations : public IRMutator { + using IRMutator::visit; + + string tile_name; + string wmma_name; + int pass = 0; + bool in_allocate = false; + bool found_shape = false; + Shape shape{}; + + // A WMMAAccumulator allocation may hold several accumulator fragments as + // 2D sub-tiles. This tracks them. + vector subtiles; + + string get_subtile_name(const Expr &index) { + int idx = Halide::Internal::get_subtile(index, "tensor core accumulator", &subtiles); + internal_assert(idx >= 0); // errors handled already + return wmma_name + std::to_string(idx); + } + + Stmt visit(const Allocate *op) override { + if (op->memory_type != MemoryType::WMMAAccumulator) { + return IRMutator::visit(op); + } + + user_assert(op->type == Float(32) || op->type == Float(16)) + << "Tensor core accumulators must hold 32-bit or 16-bit floats, but " + << op->name << " holds " << op->type << ".\n"; + + user_assert(!in_allocate) + << "Already in a tensor core accumulator allocation at the allocation for " + << op->name << ". We do not currently support multiple nested tensor core " + << "matrix multiplies."; + + ScopedValue old_wmma_name(wmma_name, op->name + ".wmma."); + ScopedValue old_tile_name(tile_name, op->name); + ScopedValue old_in_alloc(in_allocate, true); + + // In the first pass we recognize the matrix multiplies, which is what + // tells us the tile shape. In the second we recognize the + // zero-initializations and the stores out to memory, both of which + // need to know the shape. + pass = 0; + Stmt body = mutate(op->body); + user_assert(found_shape) + << op->name << " is stored in WMMAAccumulator memory, but no matrix " + << "multiply operation was found that stores to it, so the shape of the " + << "tile was unable to be determined.\n"; + pass = 1; + body = mutate(body); + + // Each fragment is one accumulator's worth of storage per lane. The + // allocations stay outside the loops over lanes, because a loop over + // lanes is a loop over threads, and register allocations outside a + // thread loop already get replicated per thread. + for (int i = 0; i < (int)subtiles.size(); i++) { + body = Allocate::make(wmma_name + std::to_string(i), op->type, + MemoryType::WMMAAccumulator, {fragment_elements}, + const_true(), body); + } + return body; + } + + Stmt visit(const Free *op) override { + if (op->name != tile_name) { + return op; + } + Stmt s; + for (int i = 0; i < (int)subtiles.size(); i++) { + Stmt f = Free::make(wmma_name + std::to_string(i)); + s = s.defined() ? Block::make(std::move(s), std::move(f)) : std::move(f); + } + return s; + } + + Stmt visit(const ProducerConsumer *op) override { + if (op->name != tile_name) { + return IRMutator::visit(op); + } + return ProducerConsumer::make(wmma_name, op->is_producer, mutate(op->body)); + } + + Expr visit(const Load *op) override { + user_assert(op->name != tile_name) + << "Tensor core accumulator " << tile_name + << " used outside a tensor core instruction"; + return IRMutator::visit(op); + } + + Stmt visit(const Store *op) override { + // There are three operations on an accumulator: + // 1) Zero-initialization + // 2) Matrix multiply + // 3) Stores to memory + // + // The matrix multiply is what tells us the tile shape, so we recognize + // those in the first pass and the other two in the second. + + if (op->name != tile_name) { + Expr store_index; + const Load *load = peel_store_permutations(op, &store_index).as(); + if (load && load->name == tile_name) { + return pass == 1 ? + convert_to_tile_store(op, store_index, + get_subtile_name(load->index), shape) : + Stmt(op); + } + // Not a copy of a tile out to memory. Recurse, so that any use of + // the accumulator buried in here gets reported as an error. + return IRMutator::visit(op); + } + + string subtile_name = get_subtile_name(op->index); + + if (is_initialization(op, tile_name)) { + return pass == 1 ? convert_to_init(op, subtile_name, shape) : Stmt(op); + } + + if (pass == 1) { + return op; + } + + Matmul matmul = convert_to_matmul(op, subtile_name); + user_assert(!found_shape || + (matmul.shape.M == shape.M && + matmul.shape.N == shape.N && + matmul.shape.K == shape.K)) + << "Found inconsistent tile shapes for a WMMAAccumulator allocation across " + << "multiple matrix multiplies that store to it."; + shape = matmul.shape; + found_shape = true; + return matmul.stmt; + } +}; + +} // namespace + +Stmt extract_wmma_operations(const Stmt &s) { + return ExtractWMMAOperations()(s); +} + +bool is_wmma_intrinsic(const Call *op) { + return (op->is_intrinsic(Call::wmma_matrix_to_fragment_a) || + op->is_intrinsic(Call::wmma_matrix_to_fragment_b) || + op->is_intrinsic(Call::wmma_matrix_to_fragment_c) || + op->is_intrinsic(Call::wmma_mma)); +} + +Expr peel_store_permutations(const Store *op, Expr *index) { + Expr value = op->value; + *index = op->index; + while (const Shuffle *shuffle = value.as()) { + if (!shuffle->is_transpose()) { + break; + } + // Transposing both sides of the store cancels out. + const int rows = value.type().lanes() / shuffle->transpose_factor(); + *index = Shuffle::make_transpose(*index, rows); + value = shuffle->vectors[0]; + } + return value; +} + +bool wmma_matrix_layout(const MultiRamp &mr, int rows, int cols, + WMMAMatrixLayout *result) { + std::vector strides; + if (mr.total_lanes() != rows * cols || + !mr.strides_for_shape({cols, rows}, &strides)) { + return false; + } + // A row-major matrix has its columns one element apart. + if (is_const_one(strides[0])) { + result->row_major = true; + result->stride = strides[1]; + } else if (is_const_one(strides[1])) { + result->row_major = false; + result->stride = strides[0]; + } else { + return false; + } + result->base = mr.base; + return true; +} + +bool is_wmma_matrix_store(const Store *op) { + Expr index; + const Call *value = peel_store_permutations(op, &index).as(); + return value && value->is_intrinsic(Call::wmma_fragment_to_matrix_d); +} + +int wmma_matrix_arg(const Call *op) { + if (op->is_intrinsic(Call::wmma_matrix_to_fragment_a) || + op->is_intrinsic(Call::wmma_matrix_to_fragment_b) || + op->is_intrinsic(Call::wmma_matrix_to_fragment_c)) { + return 3; + } + return -1; +} + +} // namespace Internal +} // namespace Halide diff --git a/src/ExtractWMMAOperations.h b/src/ExtractWMMAOperations.h new file mode 100644 index 000000000000..d492021cbac5 --- /dev/null +++ b/src/ExtractWMMAOperations.h @@ -0,0 +1,94 @@ +#ifndef HALIDE_EXTRACT_WMMA_OPERATIONS_H +#define HALIDE_EXTRACT_WMMA_OPERATIONS_H + +/** \file + * Defines the lowering pass that injects calls to the warp-level matrix + * multiply-accumulate intrinsics that drive NVIDIA tensor cores. + */ + +#include "Expr.h" +#include "MultiRamp.h" + +namespace Halide { +namespace Internal { + +struct Call; +struct Store; + +/** Rewrite matrix multiplies that accumulate into WMMAAccumulator memory as + * calls to the wmma intrinsics understood by the PTX backend, and wrap them in + * a loop over the 32 lanes of a warp. */ +Stmt extract_wmma_operations(const Stmt &s); + +/** Whether a Call is one of the wmma intrinsics produced by the pass above. */ +bool is_wmma_intrinsic(const Call *op); + +/** The index of the argument of a wmma intrinsic that gives the matrix it + * moves to or from memory, or -1 if it doesn't have one (which is the case for + * wmma_mma, whose operands are fragments rather than matrices). + * + * That argument is a Load of the whole matrix, with its lanes in row-major + * order regardless of how the matrix is laid out in memory. The layout in + * memory, and the distance between its rows or columns, are recoverable from + * the index of that Load: whichever dimension has stride one is the dense one. + * + * Being a Load rather than a bare handle to the allocation means that the + * passes which track uses of an allocation (dead allocation removal, shared + * memory liveness and packing, closure construction) see the access, and see + * its true footprint. It is never actually loaded - codegen takes its address + * - so passes that rewrite the *structure* of a load must leave it alone, or + * the tile shape will no longer be recoverable. */ +int wmma_matrix_arg(const Call *op); + +/** Whether a Store is the copy of a tensor core accumulator out to memory + * produced by the pass above. Such a store writes the whole matrix, with each + * lane of the warp writing the entries it holds: + * + * out[matrix] = wmma_fragment_to_matrix(M, N, K, fragment) + * with predicate wmma_lane_owns(M, N, K) + * + * wmma_fragment_to_matrix_d permutes this lane's fragment up into a whole + * matrix, leaving the entries the lane doesn't hold undefined, and + * wmma_lane_owns says whether this lane holds each entry of the matrix in its + * fragment. Both are opaque, because the mapping from matrix entry to lane + * isn't architecturally specified. The layout of the matrix in memory is + * recoverable from the index, as it is for a load. + * + * Passes that rewrite the structure of a store must leave these alone. */ +bool is_wmma_matrix_store(const Store *op); + +/** Peel the lane permutations the simplifier may have applied to the value of a + * store, moving the inverse of each onto the index, and return what's left of + * the value. Transposing both sides of a store cancels out, so this recovers + * the store as it was written. + * + * The simplifier rewrites a store whose index isn't dense into a dense store of + * a transposed value, so a wmma store to a column-major matrix arrives at the + * backend as a dense store of the transpose of that matrix. That's the same + * thing as a column-major store, which is one of the instructions, so the + * backend just undoes the rewrite rather than preventing it. */ +Expr peel_store_permutations(const Store *op, Expr *index); + +/** How a matrix a wmma instruction moves is laid out in memory. */ +struct WMMAMatrixLayout { + /** The address of its first entry. */ + Expr base; + /** The distance between its rows if it's row-major, or between its + * columns if it's column-major, in elements. */ + Expr stride; + bool row_major; +}; + +/** Read how a matrix is laid out in memory off the access pattern of the load + * or store of it: whichever dimension has stride one is the dense one, and the + * other stride is the distance between its rows or columns. The extraction pass + * uses this to build the access pattern, and the backend to read it back, so + * they can't disagree about the convention. Returns false if the access isn't a + * dense tile of the given shape. */ +bool wmma_matrix_layout(const MultiRamp &mr, int rows, int cols, + WMMAMatrixLayout *result); + +} // namespace Internal +} // namespace Halide + +#endif diff --git a/src/FlattenNestedRamps.cpp b/src/FlattenNestedRamps.cpp index 2cb552ec58d7..2e9e28a8eac8 100644 --- a/src/FlattenNestedRamps.cpp +++ b/src/FlattenNestedRamps.cpp @@ -2,6 +2,7 @@ #include "Bounds.h" #include "CSE.h" #include "Deinterleave.h" +#include "ExtractWMMAOperations.h" #include "IRMutator.h" #include "IROperator.h" #include "MultiRamp.h" @@ -63,6 +64,20 @@ class FlattenRamps : public IRMutator { return Shuffle::make_slice(v, n * inner_lanes, 1, inner_lanes); } + Expr visit(const Call *op) override { + if (wmma_matrix_arg(op) >= 0) { + // This pass turns a load of a nested ramp into a concatenation of + // dense loads, which is several load nodes rather than one, so + // unlike a reshaping shuffle it can't be undone by the backend. The + // matrix a wmma instruction takes a fragment out of is never + // actually loaded - only its address is used - so there's nothing + // to gain by rewriting it anyway. The other arguments are scalar + // constants, so there's nothing to do to them either. + return op; + } + return IRMutator::visit(op); + } + Expr visit(const Load *op) override { // Convert a load of a bounded span of indices into a shuffle // of a dense or strided load if possible. @@ -178,6 +193,14 @@ class FlattenRamps : public IRMutator { } Stmt visit(const Store *op) override { + if (is_wmma_matrix_store(op)) { + // One wmma store instruction writes the whole matrix for the whole + // warp, so it must not be broken into a store per row. Unlike the + // load side, this can't be undone afterwards, because it would have + // become several statements. + return op; + } + // If the index is a multiramp, unroll into a sequence of per-inner-ramp // stores, for the same reason as the Load visitor above. if (op->index.type().is_vector()) { diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index 5d3d0d9481ac..89097e579201 100644 --- a/src/FuseGPUThreadLoops.cpp +++ b/src/FuseGPUThreadLoops.cpp @@ -464,7 +464,8 @@ class ExtractSharedAndHeapAllocations : public IRMutator { !is_gpu_shared(op->memory_type) && op->memory_type != MemoryType::GPUTexture) || op->memory_type == MemoryType::Register || - op->memory_type == MemoryType::Stack) { + op->memory_type == MemoryType::Stack || + op->memory_type == MemoryType::WMMAAccumulator) { // These allocations go in register or local memory return IRMutator::visit(op); } @@ -1141,7 +1142,8 @@ class ExtractRegisterAllocations : public IRMutator { user_assert(op->memory_type == MemoryType::Stack || op->memory_type == MemoryType::Register || op->memory_type == MemoryType::Heap || - op->memory_type == MemoryType::Auto) + op->memory_type == MemoryType::Auto || + op->memory_type == MemoryType::WMMAAccumulator) << "Allocation " << op->name << " is scheduled inside a loop over GPU threads, so " << "it must live in stack memory, heap memory, or registers. " << "Shared allocations at this loop level are not yet supported.\n"; @@ -1412,6 +1414,7 @@ class InjectThreadBarriers : public IRMutator { case MemoryType::LockedCache: case MemoryType::VTCM: case MemoryType::AMXTile: + case MemoryType::WMMAAccumulator: break; } @@ -1438,6 +1441,7 @@ class InjectThreadBarriers : public IRMutator { case MemoryType::LockedCache: case MemoryType::VTCM: case MemoryType::AMXTile: + case MemoryType::WMMAAccumulator: break; } diff --git a/src/IR.cpp b/src/IR.cpp index 586292179f21..7969764ae69c 100644 --- a/src/IR.cpp +++ b/src/IR.cpp @@ -887,6 +887,12 @@ constexpr const char *intrinsic_op_names[] = { "widening_shift_left", "widening_shift_right", "widening_sub", + "wmma_fragment_to_matrix_d", + "wmma_lane_owns", + "wmma_matrix_to_fragment_a", + "wmma_matrix_to_fragment_b", + "wmma_matrix_to_fragment_c", + "wmma_mma", // keep-sorted end }; diff --git a/src/IR.h b/src/IR.h index c1a7d4430cf0..0dd3031f8500 100644 --- a/src/IR.h +++ b/src/IR.h @@ -868,6 +868,37 @@ struct Call : public ExprNode { widening_shift_left, widening_shift_right, widening_sub, + // Permute this lane's tensor core accumulator (d) fragment up into a + // whole matrix, leaving the entries this lane doesn't hold undefined. + // wmma_fragment_to_matrix_d(M, N, K, fragment) + wmma_fragment_to_matrix_d, + // Whether this lane holds each entry of an M x N tensor core + // accumulator in its fragment. Used as the predicate of the store that + // copies an accumulator out to memory. + // wmma_lane_owns(M, N, K) + wmma_lane_owns, + // Take this lane's share of a tensor core fragment out of a matrix. + // The a operand is M x K, the b operand is K x N, and the accumulator + // (c) is M x N. The mapping from matrix entry to lane is opaque, + // because it isn't architecturally specified. + // + // The matrix argument is a Load of the whole of it, with its lanes in + // row-major order; the layout in memory and the distance between its + // rows or columns are recoverable from the index. The hardware can only + // do this as part of a memory read, which is why the argument has to be + // a Load rather than an arbitrary matrix value. + // wmma_matrix_to_fragment_a(M, N, K, matrix) + // @{ + wmma_matrix_to_fragment_a, + wmma_matrix_to_fragment_b, + wmma_matrix_to_fragment_c, + // @} + // A tensor core matrix multiply-accumulate. The layouts say how the a + // and b fragments were taken out of their matrices, which changes how + // the hardware arranges them in registers. The accumulator's fragment + // layout doesn't depend on how it was loaded, so it isn't a parameter. + // wmma_mma(M, N, K, a_layout, b_layout, a, b, c) + wmma_mma, // keep-sorted end IntrinsicOpCount // Sentinel: keep last. }; diff --git a/src/IRPrinter.cpp b/src/IRPrinter.cpp index 33451ded3b72..e7d9c4c54b32 100644 --- a/src/IRPrinter.cpp +++ b/src/IRPrinter.cpp @@ -172,6 +172,9 @@ std::ostream &operator<<(std::ostream &out, const MemoryType &t) { case MemoryType::AMXTile: out << "AMXTile"; break; + case MemoryType::WMMAAccumulator: + out << "WMMAAccumulator"; + break; } return out; } diff --git a/src/Lower.cpp b/src/Lower.cpp index 7a5c46caa41b..b486d21b00b3 100644 --- a/src/Lower.cpp +++ b/src/Lower.cpp @@ -27,6 +27,7 @@ #include "Deinterleave.h" #include "EarlyFree.h" #include "ExtractTileOperations.h" +#include "ExtractWMMAOperations.h" #include "FindCalls.h" #include "FindIntrinsics.h" #include "FlattenNestedRamps.h" @@ -354,6 +355,12 @@ void lower_impl(const vector &output_funcs, s = simplify(s); log("Lowering after vectorizing:", s); + if (t.has_feature(Target::CUDA)) { + debug(1) << "Extracting tensor core operations...\n"; + s = extract_wmma_operations(s); + log("Lowering after extracting tensor core operations:", s); + } + if (t.has_gpu_feature() || t.has_feature(Target::Vulkan)) { debug(1) << "Injecting per-block gpu synchronization...\n"; diff --git a/src/LowerWarpShuffles.cpp b/src/LowerWarpShuffles.cpp index 75f52e1d65cc..8ce449c577f5 100644 --- a/src/LowerWarpShuffles.cpp +++ b/src/LowerWarpShuffles.cpp @@ -655,10 +655,13 @@ class LowerWarpShuffles : public IRMutator { Stmt visit(const Allocate *op) override { if (this_lane.defined() || is_gpu_shared(op->memory_type) || - op->memory_type == MemoryType::Heap) { - // Not a warp-level allocation. Warp-level storage is per-lane - // register storage; shared and heap (global) memory are never - // striped across lanes. + op->memory_type == MemoryType::Heap || + op->memory_type == MemoryType::WMMAAccumulator) { + // Not an allocation for us to stripe. Warp-level storage is + // per-lane register storage; shared and heap (global) memory are + // never striped across lanes, and tensor core accumulators are + // already striped across lanes in a layout that only the wmma + // instructions understand. return IRMutator::visit(op); } else { // Pick up this allocation and deposit it inside the loop over lanes at reduced size. diff --git a/src/Serialization.cpp b/src/Serialization.cpp index c460d35e8eb4..acc91504f022 100644 --- a/src/Serialization.cpp +++ b/src/Serialization.cpp @@ -160,6 +160,8 @@ Serialize::MemoryType Serializer::serialize_memory_type(const MemoryType &memory return Serialize::MemoryType::VTCM; case MemoryType::AMXTile: return Serialize::MemoryType::AMXTile; + case MemoryType::WMMAAccumulator: + return Serialize::MemoryType::WMMAAccumulator; default: user_error << "Unsupported memory type\n"; return Serialize::MemoryType::Auto; diff --git a/src/Simplify_Call.cpp b/src/Simplify_Call.cpp index d63caa2e6c39..97951f1d8148 100644 --- a/src/Simplify_Call.cpp +++ b/src/Simplify_Call.cpp @@ -59,6 +59,38 @@ Expr Simplify::visit(const Call *op, ExprInfo *info) { info->cast_to(op->type); } + // A pure elementwise call of broadcasts is a broadcast of the call. Lifting + // the broadcast out lets the surrounding expression see that the value is + // the same in every lane. The clauses below that special-case particular + // functions have already mutated their arguments, so they can use this + // directly; anything that falls through to the general case at the bottom + // gets the same treatment there. + auto lift_broadcast_out = [&](const std::vector &args) -> Expr { + if (!op->type.is_vector() || !op->is_pure()) { + return Expr(); + } + const int lanes = op->type.lanes(); + std::vector scalar_args; + scalar_args.reserve(args.size()); + for (const Expr &arg : args) { + if (arg.type().is_scalar()) { + scalar_args.push_back(arg); + } else if (const Broadcast *b = arg.as(); + b && b->lanes == lanes && b->value.type().is_scalar()) { + scalar_args.push_back(b->value); + } else { + return Expr(); + } + } + if (scalar_args.empty()) { + return Expr(); + } + Expr scalar = Call::make(op->type.element_of(), op->name, scalar_args, + op->call_type, op->func, op->value_index, + op->image, op->param); + return mutate(Broadcast::make(std::move(scalar), lanes), info); + }; + if (op->is_intrinsic(Call::unreachable)) { in_unreachable = true; return op; @@ -778,6 +810,8 @@ Expr Simplify::visit(const Call *op, ExprInfo *info) { if (auto f = as_const_float(arg)) { auto fn = it->second; return make_const(arg.type(), fn(*f), info); + } else if (Expr e = lift_broadcast_out({arg}); e.defined()) { + return e; } else if (arg.same_as(op->args[0])) { return op; } else { @@ -810,6 +844,8 @@ Expr Simplify::visit(const Call *op, ExprInfo *info) { if (auto f = as_const_float(arg)) { auto fn = it->second; return make_const(arg.type(), fn(*f), info); + } else if (Expr e = lift_broadcast_out({arg}); e.defined()) { + return e; } else if (call && (call->call_type == Call::PureExtern || call->call_type == Call::PureIntrinsic) && (it = pure_externs_truncation.find(call->name)) != pure_externs_truncation.end()) { // For any combination of these integer-valued functions, we can @@ -863,6 +899,11 @@ Expr Simplify::visit(const Call *op, ExprInfo *info) { // No else: we want to fall thru from the PureExtern clause. { auto [new_args, changed] = mutate_with_changes(op->args); + + if (Expr e = lift_broadcast_out(new_args); e.defined()) { + return e; + } + if (!changed) { return op; } else { diff --git a/src/Simplify_Stmts.cpp b/src/Simplify_Stmts.cpp index 381ebb5a6177..c656f8e714ba 100644 --- a/src/Simplify_Stmts.cpp +++ b/src/Simplify_Stmts.cpp @@ -320,6 +320,7 @@ Stmt Simplify::visit(const Provide *op) { } Stmt Simplify::visit(const Store *op) { + found_buffer_reference(op->name); Expr predicate = mutate(op->predicate, nullptr); diff --git a/src/halide_ir.fbs b/src/halide_ir.fbs index 9967296986f3..b5786596512b 100644 --- a/src/halide_ir.fbs +++ b/src/halide_ir.fbs @@ -117,6 +117,7 @@ enum MemoryType: byte { VTCM, AMXTile, GPUSharedAsync, + WMMAAccumulator, } table Range { diff --git a/test/correctness/wmma_matmul.cpp b/test/correctness/wmma_matmul.cpp new file mode 100644 index 000000000000..868db7f6a245 --- /dev/null +++ b/test/correctness/wmma_matmul.cpp @@ -0,0 +1,262 @@ +#include "Halide.h" +#include + +using namespace Halide; + +namespace { + +struct Params { + // The size of the matrices. + int M = 64, N = 64, K = 64; + // The tensor core tile shape to use. + int tile_m = 16, tile_n = 16, tile_k = 16; + // How many tiles of accumulator each warp holds, in each dimension. + int tiles_m = 1, tiles_n = 1; + // How many warps per block. + int warps = 1; + // Whether each input matrix is stored with its rows or its columns dense. + bool a_transposed = false, b_transposed = false; + // Whether to accumulate in half precision instead of single precision. + bool half_accumulator = false; + // Whether to stage the operand tiles through shared memory inside the + // reduction loop. + bool stage_in_shared = false; + // Whether to start the accumulator from a matrix in memory rather than + // from zero. + bool init_from_memory = false; + // Whether the output is stored with its columns dense rather than its rows. + bool out_transposed = false; +}; + +std::ostream &operator<<(std::ostream &s, const Params &p) { + return s << p.M << "x" << p.N << "x" << p.K + << " in " << p.tile_m << "x" << p.tile_n << "x" << p.tile_k + << " tiles, " << p.tiles_m << "x" << p.tiles_n << " tiles per warp, " + << p.warps << " warps, a" << (p.a_transposed ? "T" : "") + << " b" << (p.b_transposed ? "T" : "") + << (p.half_accumulator ? ", f16 accumulator" : "") + << (p.stage_in_shared ? ", staged through shared" : "") + << (p.init_from_memory ? ", accumulator initialized from memory" : "") + << (p.out_transposed ? ", transposed output" : ""); +} + +void fill(Buffer &buf) { + buf.fill([]() { + return float16_t(((float)rand() / RAND_MAX) - 0.5f); + }); +} + +bool test(const Params &p) { + // Halide indexes matrices with the dense dimension first, so A(k, y) is a + // row-major M x K matrix, and A(y, k) is a column-major one. + Buffer A(p.a_transposed ? p.M : p.K, p.a_transposed ? p.K : p.M); + Buffer B(p.b_transposed ? p.K : p.N, p.b_transposed ? p.N : p.K); + fill(A); + fill(B); + + Var x("x"), y("y"), kk("kk"), xx("xx"), yy("yy"); + RDom k(0, p.K, "k"); + Func prod("prod"), out("out"), Af("Af"), Bf("Bf"), init("init"); + + // These are inlined unless we're staging through shared memory. + Af(kk, yy) = p.a_transposed ? A(yy, kk) : A(kk, yy); + Bf(xx, kk) = p.b_transposed ? B(kk, xx) : B(xx, kk); + Expr a = Af(k, y); + Expr b = Bf(x, k); + // With a half-precision accumulator the tile is copied out to memory as + // float16, so the output has to be float16 too. + // The accumulator either starts at zero, or at a matrix that already exists + // in memory, which the hardware can load straight into the fragments. + Type acc_type = p.half_accumulator ? Float(16) : Float(32); + init(x, y) = cast(acc_type, (x * 3 + y) % 7) * cast(acc_type, 0.25f); + if (p.half_accumulator) { + prod(x, y) = p.init_from_memory ? init(x, y) : cast(0.f); + prod(x, y) += a * b; + } else { + prod(x, y) = p.init_from_memory ? init(x, y) : Expr(0.f); + prod(x, y) += cast(a) * cast(b); + } + out(x, y) = prod(x, y); + + Var xi("xi"), yi("yi"), xt("xt"), mmxi("mmxi"), mmyi("mmyi"); + Var rxi("rxi"), ryi("ryi"); + RVar rro("rro"), rri("rri"); + + // Each block computes a tile of the output. Within a block, each warp + // computes a strip of that tile, and each warp's strip is broken into + // tensor core tiles, which are the innermost two (vectorized) dimensions. + out.bound(x, 0, p.N) + .bound(y, 0, p.M) + .split(x, x, xi, p.tile_n * p.tiles_n * p.warps) + .split(xi, xt, xi, p.tile_n * p.tiles_n) + .split(xi, xi, mmxi, p.tile_n) + .split(y, y, yi, p.tile_m * p.tiles_m) + .split(yi, yi, mmyi, p.tile_m) + .gpu_blocks(x, y) + .reorder(mmxi, mmyi, xi, yi, xt, x, y) + .unroll(xi) + .unroll(yi) + .vectorize(mmxi) + .vectorize(mmyi); + if (p.warps > 1) { + // With one warp there's no need for a loop over warps, and leaving it + // serial keeps anything computed inside it out of the thread loops. + out.gpu_threads(xt); + } + + prod.compute_at(out, xt) + .store_in(MemoryType::WMMAAccumulator) + .split(x, x, rxi, p.tile_n) + .split(y, y, ryi, p.tile_m) + .vectorize(rxi) + .vectorize(ryi) + .unroll(x) + .unroll(y); + + prod.update() + .split(x, x, rxi, p.tile_n) + .split(y, y, ryi, p.tile_m) + .split(k, rro, rri, p.tile_k) + .reorder(rri, rxi, ryi, x, y, rro) + .unroll(x) + .unroll(y) + .atomic() + .vectorize(rri) + .vectorize(rxi) + .vectorize(ryi); + + if (p.init_from_memory) { + Var ixi("ixi"), iyi("iyi"); + init.compute_root().gpu_tile(x, y, ixi, iyi, 16, 16); + } + + if (p.stage_in_shared) { + Var t("t"), ti("ti"), to("to"); + Af.compute_at(prod, rro) + .store_in(MemoryType::GPUShared) + .fuse(kk, yy, t) + .split(t, to, ti, 32) + .gpu_lanes(ti); + Bf.compute_at(prod, rro) + .store_in(MemoryType::GPUShared) + .fuse(xx, kk, t) + .split(t, to, ti, 32) + .gpu_lanes(ti); + } + + if (p.out_transposed) { + out.output_buffer().dim(0).set_stride(p.M).dim(1).set_stride(1); + } + + // A transposed output is a view of a buffer with the dimensions swapped, so + // its columns are dense in memory instead of its rows. + Buffer result_storage(p.out_transposed ? p.M : p.N, + p.out_transposed ? p.N : p.M); + Buffer result = + p.out_transposed ? result_storage.transposed(0, 1) : result_storage; + Buffer result_half(p.N, p.M); + auto get = [&](int i, int j) { + return p.half_accumulator ? (float)result_half(i, j) : result(i, j); + }; + if (p.half_accumulator) { + out.realize(result_half); + result_half.copy_to_host(); + } else { + out.realize(result); + result.copy_to_host(); + } + + for (int j = 0; j < p.M; j++) { + for (int i = 0; i < p.N; i++) { + float ref = p.init_from_memory ? (float)(((i * 3 + j) % 7) * 0.25f) : 0.f; + for (int l = 0; l < p.K; l++) { + ref += (float)(p.a_transposed ? A(j, l) : A(l, j)) * + (float)(p.b_transposed ? B(l, i) : B(i, l)); + } + // The accumulation happens in a different order on the GPU, and + // the inputs are half-precision, so allow some slack. + float tolerance = p.half_accumulator ? 5e-2f : 1e-2f; + if (std::abs(get(i, j) - ref) > tolerance * std::max(1.f, std::abs(ref))) { + std::cerr << "Mismatch at " << i << ", " << j << ": " + << get(i, j) << " != " << ref << "\n" + << "For matmul of " << p << "\n"; + return false; + } + } + } + return true; +} + +} // namespace + +int main(int argc, char **argv) { + Target target = get_jit_target_from_environment(); + if (!target.has_feature(Target::CUDA)) { + printf("[SKIP] WMMA matrix multiplies require CUDA.\n"); + return 0; + } + if (target.get_cuda_capability_lower_bound() < 70) { + printf("[SKIP] WMMA matrix multiplies require CUDA compute capability 7.0 or above.\n"); + return 0; + } + + std::vector params; + + // The simplest possible case. + params.push_back({}); + + // Each of the other supported tile shapes. + params.push_back({.tile_m = 32, .tile_n = 8}); + params.push_back({.tile_m = 8, .tile_n = 32}); + + // Each combination of input layouts. + params.push_back({.a_transposed = true}); + params.push_back({.b_transposed = true}); + params.push_back({.a_transposed = true, .b_transposed = true}); + + // Several accumulator fragments live at once, which is what you want in + // practice so that each operand load feeds more than one multiply. + params.push_back({.tiles_m = 2, .tiles_n = 2}); + + // More than one warp per block. + params.push_back({.warps = 2}); + params.push_back({.tiles_m = 2, .tiles_n = 2, .warps = 2}); + + // Accumulating in half precision, which uses half as many registers per + // accumulator fragment. + params.push_back({.half_accumulator = true}); + params.push_back({.tiles_m = 2, .tiles_n = 2, .half_accumulator = true}); + + // Operand tiles staged through shared memory inside the reduction loop. + // This only works if the loop over lanes wraps the individual wmma + // statements rather than the whole accumulator allocation. + params.push_back({.stage_in_shared = true}); + params.push_back({.tiles_m = 2, .tiles_n = 2, .stage_in_shared = true}); + params.push_back({.half_accumulator = true, .stage_in_shared = true}); + + // Accumulators that start from a matrix already in memory rather than from + // zero, which uses the load-into-the-accumulator-fragment instruction. + params.push_back({.init_from_memory = true}); + params.push_back({.tiles_m = 2, .tiles_n = 2, .init_from_memory = true}); + params.push_back({.half_accumulator = true, .init_from_memory = true}); + params.push_back({.stage_in_shared = true, .init_from_memory = true}); + + // A column-major output, which the accumulator is stored to with the other + // layout of the store instruction. + params.push_back({.out_transposed = true}); + params.push_back({.tiles_m = 2, .tiles_n = 2, .out_transposed = true}); + + // A reduction that isn't a whole number of tiles per unrolled step, and + // matrices that aren't square. + params.push_back({.M = 32, .N = 128, .K = 256}); + + for (const Params &p : params) { + if (!test(p)) { + printf("Failed!\n"); + return 1; + } + } + + printf("Success!\n"); + return 0; +} From 0be212c56ee792a91b5ce8330196b831cc08a0ca Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 28 Jul 2026 16:11:42 -0700 Subject: [PATCH 18/59] Let tensor core accumulators live outside the loop over warps An accumulator scheduled at block level, with the reduction loop above the loop over warps, is the shape a matmul needs in order to stage its operand panels into shared memory once per block and have every warp reuse them. Three things were in the way: - The index of an access to the accumulator then depends on which warp is doing it. That dependence selects between the per-thread copies of the allocation rather than between subtiles within one, so it's substituted away before working out which subtile an access refers to. - Halide sees an accumulator outside the thread loops as shared between threads, so it keeps the atomic node around the update. Codegen then scalarizes the store, one lane at a time, and the wmma intrinsics fall apart. An accumulator is per-thread register storage, so there's nothing to race with, and the pass drops the atomic. - The loop over warps was being given the innermost thread dimension, which the loop over lanes needs. Inside an accumulator allocation, a thread loop is one dimension further out than it looks. Also tune apps/tensorcore_matmul, which now unrolls a couple of reduction steps to get more operand loads in flight. Together with the register cap removal this takes it from 20.7 to 26.1 TFlop/s at 1024, 32.6 to 38.8 at 2048, and 40.2 to 39.6 at 4096, against cuBLAS at 43.3, 49.6 and 50.0. Staging the operands through shared memory is now expressible, and there's a test for it, but it isn't yet a win: it peaks at 37.4 TFlop/s against 41.2 for loading the operands straight from global memory and letting L2 do the reuse. The barrier between staging and computing serializes the two, and recovering that needs double buffering, which Halide expresses with async and ring_buffer. Those lower to Fork nodes and semaphores, which have no meaning in device code, so that's a separate piece of work. Co-Authored-By: Claude Opus 5 --- apps/tensorcore_matmul/matmul_generator.cpp | 11 +- src/CanonicalizeGPUVars.cpp | 26 ++++- src/CodeGen_PTX_Dev.cpp | 16 +-- src/ExtractWMMAOperations.cpp | 43 +++++++- test/correctness/wmma_matmul.cpp | 114 ++++++++++++++++++++ 5 files changed, 197 insertions(+), 13 deletions(-) diff --git a/apps/tensorcore_matmul/matmul_generator.cpp b/apps/tensorcore_matmul/matmul_generator.cpp index c879551dd343..e5bfa3844a05 100644 --- a/apps/tensorcore_matmul/matmul_generator.cpp +++ b/apps/tensorcore_matmul/matmul_generator.cpp @@ -23,7 +23,10 @@ class MatMul : public Halide::Generator { // warps there are per block. GeneratorParam tiles_x{"tiles_x", 5}; GeneratorParam tiles_y{"tiles_y", 4}; - GeneratorParam warps{"warps", 4}; + GeneratorParam warps{"warps", 2}; + // How many reduction steps to unroll, which puts more operand loads in + // flight at once. + GeneratorParam k_unroll{"k_unroll", 2}; Input> matA{"matA"}; // K x M Input> matB{"matB"}; // N x K @@ -80,7 +83,7 @@ class MatMul : public Halide::Generator { Var xi("xi"), yi("yi"), xt("xt"), mmxi("mmxi"), mmyi("mmyi"); Var rxi("rxi"), ryi("ryi"); - RVar rro("rro"), rri("rri"); + RVar rro("rro"), rri("rri"), rru("rru"); output.bound(x, 0, N) .bound(y, 0, M) @@ -112,9 +115,11 @@ class MatMul : public Halide::Generator { .split(x, x, rxi, tile_x) .split(y, y, ryi, tile_y) .split(k, rro, rri, tile_k) - .reorder(rri, rxi, ryi, x, y, rro) + .split(rro, rro, rru, k_unroll) + .reorder(rri, rxi, ryi, x, y, rru, rro) .unroll(x) .unroll(y) + .unroll(rru) .atomic() .vectorize(rri) .vectorize(rxi) diff --git a/src/CanonicalizeGPUVars.cpp b/src/CanonicalizeGPUVars.cpp index 25e2ed6781bf..f844b3ff50ec 100644 --- a/src/CanonicalizeGPUVars.cpp +++ b/src/CanonicalizeGPUVars.cpp @@ -37,7 +37,7 @@ class CountGPUBlocksThreads : public IRVisitor { // Counters that track the number of blocks, threads, and lanes loops that // we're inside of, respectively. Lanes loops also count as threads loops. - int nb = 0, nt = 0, nl = 0; + int nb = 0, nt = 0, nl = 0, nto = 0; // Whether we're already inside a lane dimension. Every lane dimension maps // to the innermost thread dimension, so one nested inside another is the @@ -60,11 +60,13 @@ class CountGPUBlocksThreads : public IRVisitor { nb += db; nl += dl; nt += dt; + nto += op->for_type == ForType::GPUThread; // Update the maximum counter values seen. nblocks = std::max(nb, nblocks); nthreads = std::max(nt, nthreads); nlanes = std::max(nl, nlanes); + nthreads_excluding_lanes = std::max(nto, nthreads_excluding_lanes); // Visit the body IRVisitor::visit(op); @@ -73,6 +75,7 @@ class CountGPUBlocksThreads : public IRVisitor { nb -= db; nl -= dl; nt -= dt; + nto -= op->for_type == ForType::GPUThread; } void visit(const Realize *op) override { @@ -97,13 +100,30 @@ class CountGPUBlocksThreads : public IRVisitor { int nblocks = 0; int nthreads = 0; int nlanes = 0; + // Threads not counting lanes, which is what a thread loop's depth is when + // the lane dimension comes from an enclosing tensor core allocation rather + // than from anything in this loop's body. + int nthreads_excluding_lanes = 0; }; class CanonicalizeGPUVars : public IRMutator { map gpu_vars; + // Whether we're inside a tensor core accumulator allocation. + // extract_wmma_operations will introduce a loop over the lanes of a warp + // somewhere inside it, which takes the innermost thread dimension, so the + // thread loops in here are one dimension further out than they look. + bool in_wmma_alloc = false; + using IRMutator::visit; + Stmt visit(const Realize *op) override { + ScopedValue old(in_wmma_alloc, + in_wmma_alloc || + op->memory_type == MemoryType::WMMAAccumulator); + return IRMutator::visit(op); + } + std::string find_replacement(const std::string &suffix, const std::string &name) { vector v = split_string(name, suffix); internal_assert(v.size() == 2); @@ -131,7 +151,9 @@ class CanonicalizeGPUVars : public IRMutator { name += gpu_block_name(counter.nblocks); debug(5) << "Replacing " << op->name << " with GPU block name " << name << "\n"; } else if (op->for_type == ForType::GPUThread) { - name += gpu_thread_name(counter.nthreads); + name += gpu_thread_name(in_wmma_alloc ? + counter.nthreads_excluding_lanes + 1 : + counter.nthreads); debug(5) << "Replacing " << op->name << " with GPU thread name " << name << "\n"; } else if (op->for_type == ForType::GPULane) { name += gpu_thread_name(0); diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 78f51138ee26..82c67ed4d53d 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -344,14 +344,16 @@ void CodeGen_PTX_Dev::visit(const Call *op) { namespace { -WMMAMatrixLayout matrix_in_memory(const string &name, const MultiRamp &mr, int rows, int cols) { +WMMAMatrixLayout matrix_in_memory(const string &name, const MultiRamp &mr, int rows, int cols, const Expr &access = Expr()) { WMMAMatrixLayout result; user_assert(wmma_matrix_layout(mr, rows, cols, &result)) - << "The memory a tensor core instruction moves a matrix of " << name - << " to or from is not a dense tile by the time it reaches the backend. " - << "This happens when the allocation is striped across threads, which " - << "occurs for a shared memory allocation made inside the loop over GPU " - << "threads. Compute it at a loop outside the threads instead.\n"; + << "The memory a tensor core instruction moves a " << rows << "x" << cols + << " matrix of " << name << " to or from is not a dense tile by the time it " + << "reaches the backend. One cause is a shared memory allocation made inside " + << "the loop over GPU threads, which gets striped across them; compute it at " + << "a loop outside the threads instead. The addresses accessed are:\n" + << mr.to_expr() << "\nThe access is:\n" + << access << "\n"; return result; } @@ -421,7 +423,7 @@ void CodeGen_PTX_Dev::codegen_wmma(const Call *op) { << "The matrix a tensor core instruction takes a fragment out of is not a " << "load with an affine index by the time it reaches the backend.\n"; WMMAMatrixLayout mem = matrix_in_memory(matrix->name, mr, - is_b ? K : M, is_a ? K : N); + is_b ? K : M, is_a ? K : N, arg); // The a and b operands are always 16-bit; an accumulator may be either. const char *type_suffix = is_a || is_b ? "f16" : (op->type.bits() == 32 ? "f32" : "f16"); diff --git a/src/ExtractWMMAOperations.cpp b/src/ExtractWMMAOperations.cpp index fa34b7cb7beb..385b9d791b1c 100644 --- a/src/ExtractWMMAOperations.cpp +++ b/src/ExtractWMMAOperations.cpp @@ -7,6 +7,7 @@ #include "IROperator.h" #include "MultiRamp.h" #include "Simplify.h" +#include "Substitute.h" #include "Util.h" /** \file Support extraction of NVIDIA tensor core (wmma) instructions. @@ -429,12 +430,38 @@ class ExtractWMMAOperations : public IRMutator { // 2D sub-tiles. This tracks them. vector subtiles; + // The loops over GPU blocks, threads, and lanes we're inside of. + vector gpu_loop_vars; + + // An accumulator allocation may sit outside the loops over GPU threads, in + // which case each thread gets its own copy of it and only ever touches its + // own slice. Any dependence of the index on the thread is selecting between + // those copies, not between subtiles within one, so drop it. + Expr index_within_thread(const Expr &index) { + Expr idx = index; + for (const string &v : gpu_loop_vars) { + idx = substitute(v, 0, idx); + } + return simplify(idx); + } + string get_subtile_name(const Expr &index) { - int idx = Halide::Internal::get_subtile(index, "tensor core accumulator", &subtiles); + int idx = Halide::Internal::get_subtile(index_within_thread(index), + "tensor core accumulator", &subtiles); internal_assert(idx >= 0); // errors handled already return wmma_name + std::to_string(idx); } + Stmt visit(const For *op) override { + if (!is_gpu(op->for_type)) { + return IRMutator::visit(op); + } + gpu_loop_vars.push_back(op->name); + Stmt s = IRMutator::visit(op); + gpu_loop_vars.pop_back(); + return s; + } + Stmt visit(const Allocate *op) override { if (op->memory_type != MemoryType::WMMAAccumulator) { return IRMutator::visit(op); @@ -478,6 +505,20 @@ class ExtractWMMAOperations : public IRMutator { return body; } + Stmt visit(const Atomic *op) override { + if (op->producer_name == tile_name) { + // A tensor core accumulator is per-thread register storage, so + // there's nothing for another thread to race with. The atomic is + // there because the accumulator is scheduled outside the loops over + // threads, which makes it look shared. + user_assert(op->mutex_name.empty()) + << "Accumulating into a tensor core accumulator should not need a " + << "mutex.\n"; + return mutate(op->body); + } + return IRMutator::visit(op); + } + Stmt visit(const Free *op) override { if (op->name != tile_name) { return op; diff --git a/test/correctness/wmma_matmul.cpp b/test/correctness/wmma_matmul.cpp index 868db7f6a245..d2e9ba9a4538 100644 --- a/test/correctness/wmma_matmul.cpp +++ b/test/correctness/wmma_matmul.cpp @@ -187,6 +187,115 @@ bool test(const Params &p) { return true; } +// The accumulators can also live outside the loop over warps, which is what +// lets the reduction loop sit above it, so that the operand panels can be +// staged into shared memory once per block and reused by every warp. +bool test_block_level_accumulator() { + const int M = 128, N = 128, K = 128; + const int tile = 16, tiles_x = 2, tiles_y = 2, warps = 2, bk = 32; + const int block_x = tile * tiles_x * warps, block_y = tile * tiles_y; + + Buffer A(K, M), B(N, K); + fill(A); + fill(B); + + Var x("x"), y("y"), kk("kk"), yy("yy"), xx("xx"); + RDom k(0, K, "k"); + Func prod("prod"), out("out"), As("As"), Bs("Bs"); + + As(kk, yy) = A(kk, yy); + Bs(xx, kk) = B(xx, kk); + prod(x, y) = 0.f; + prod(x, y) += cast(As(k, y)) * cast(Bs(x, k)); + out(x, y) = prod(x, y); + + Var xi("xi"), xt("xt"), yi("yi"), mmxi("mmxi"), mmyi("mmyi"); + Var rxi("rxi"), ryi("ryi"), xw("xw"), t("t"), ti("ti"), tw("tw"), to("to"); + Var kko("kko"), kki("kki"), xxo("xxo"), xxi("xxi"); + RVar ko("ko"), ki("ki"), rri("rri"); + + out.bound(x, 0, N).bound(y, 0, M) + .split(x, x, xi, block_x) + .split(xi, xt, xi, tile * tiles_x) + .split(xi, xi, mmxi, tile) + .split(y, y, yi, block_y) + .split(yi, yi, mmyi, tile) + .gpu_blocks(x, y) + .gpu_threads(xt) + .reorder(mmxi, mmyi, xi, yi, xt, x, y) + .unroll(xi) + .unroll(yi) + .vectorize(mmxi) + .vectorize(mmyi); + + prod.compute_at(out, x) + .store_in(MemoryType::WMMAAccumulator) + .split(x, xw, xi, tile * tiles_x) + .split(xi, xi, rxi, tile) + .split(y, y, ryi, tile) + .reorder(rxi, ryi, xi, y, xw) + .gpu_threads(xw) + .vectorize(rxi) + .vectorize(ryi) + .unroll(xi) + .unroll(y); + + prod.update() + .split(k, ko, ki, bk) + .split(x, xw, xi, tile * tiles_x) + .split(xi, xi, rxi, tile) + .split(y, y, ryi, tile) + .split(ki, ki, rri, tile) + .reorder(rri, rxi, ryi, xi, y, ki, xw, ko) + .gpu_threads(xw) + .unroll(xi) + .unroll(y) + .unroll(ki) + .atomic() + .vectorize(rri) + .vectorize(rxi) + .vectorize(ryi); + + As.compute_at(prod, ko) + .store_in(MemoryType::GPUShared) + .split(kk, kko, kki, 8) + .fuse(kko, yy, t) + .split(t, t, ti, 32) + .split(t, to, tw, warps) + .gpu_lanes(ti) + .gpu_threads(tw) + .vectorize(kki); + Bs.compute_at(prod, ko) + .store_in(MemoryType::GPUShared) + .split(xx, xxo, xxi, 8) + .fuse(xxo, kk, t) + .split(t, t, ti, 32) + .split(t, to, tw, warps) + .gpu_lanes(ti) + .gpu_threads(tw) + .vectorize(xxi); + + Buffer result(N, M); + out.realize(result); + result.copy_to_host(); + + for (int j = 0; j < M; j++) { + for (int i = 0; i < N; i++) { + float ref = 0.f; + for (int l = 0; l < K; l++) { + ref += (float)A(l, j) * (float)B(i, l); + } + if (std::abs(result(i, j) - ref) > 1e-2f * std::max(1.f, std::abs(ref))) { + std::cerr << "Mismatch at " << i << ", " << j << ": " + << result(i, j) << " != " << ref << "\n" + << "For a block-level accumulator staged through shared memory\n"; + return false; + } + } + } + return true; +} + } // namespace int main(int argc, char **argv) { @@ -257,6 +366,11 @@ int main(int argc, char **argv) { } } + if (!test_block_level_accumulator()) { + printf("Failed!\n"); + return 1; + } + printf("Success!\n"); return 0; } From 0468a0951f86e4c5e9e8c2e0c6255d95b4f4389a Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Wed, 29 Jul 2026 01:01:07 -0700 Subject: [PATCH 19/59] Give each tensor core accumulator its own shape, and enable the resize app The tile shape and the set of sub-tiles were tracked across the whole pass rather than per allocation, so a pipeline with two accumulators of different shapes was rejected. apps/tensorcore_resize has one of each, and now works: on an RTX 5060 Ti it downsamples a 3840x2160 image by 4x in 0.246 ms against 0.359 ms for the cuda-only schedule. Both of that app's schedules round the output up to a multiple of the tile size, so the runner gives them an output buffer of that size. Co-Authored-By: Claude Opus 5 --- apps/tensorcore_resize/README.md | 23 +++++------------------ apps/tensorcore_resize/runner.cpp | 8 +++++++- src/ExtractWMMAOperations.cpp | 11 ++++++++++- 3 files changed, 22 insertions(+), 20 deletions(-) diff --git a/apps/tensorcore_resize/README.md b/apps/tensorcore_resize/README.md index 74388196a60c..fc0dd391a5ce 100644 --- a/apps/tensorcore_resize/README.md +++ b/apps/tensorcore_resize/README.md @@ -13,22 +13,9 @@ This is the algorithmic difference between this generator and the one in ## Status -The `cudaonly` schedule works. The `tensorcore` schedule currently only works -for the first of the two stages (the resample in y). The resample in x fails -with: +Both schedules work. On an RTX 5060 Ti, downsampling a 3840x2160 image by 4x +with a Lanczos kernel takes 0.359 ms with the `cudaonly` schedule and 0.246 ms +with the `tensorcore` one, a 1.46x speed-up. -``` -Matrix multiply not recognized. [...] the matrix multiply operands are not -loads with affine indices. -``` - -The load index for that stage contains `begin_of((x / 16) * 16)`, i.e. the -starting column of this block of 16 rows of the matrix. That subexpression is -uniform across the 16 lanes of a tile, but the simplifier leaves it as -`ceil_f32` applied to `(ramp(block * 16, 1, 16) / 16) * 16` rather than folding -the divide and multiply away into a broadcast, so `is_multiramp` can't see that -it is uniform and the index isn't recognized as affine. - -Fixing this needs the simplifier to fold `ramp(a * k, 1, k) / k` down to -`broadcast(a, k)` when the ramp is nested inside another vector, after which the -lane-uniform recognition in `is_multiramp` handles the rest. +Both schedules round the output size up to a multiple of 16, so the runner +gives them an output buffer of that size. diff --git a/apps/tensorcore_resize/runner.cpp b/apps/tensorcore_resize/runner.cpp index 21429b538fb1..b96cb7e710de 100644 --- a/apps/tensorcore_resize/runner.cpp +++ b/apps/tensorcore_resize/runner.cpp @@ -27,7 +27,13 @@ int main(int argc, char **argv) { Buffer input(in_w, in_h, 3); input.fill([]() { return float16_t((float)rand() / RAND_MAX); }); - Buffer out_cuda(out_w, out_h, 3), out_tensorcore(out_w, out_h, 3); + // Both schedules work on whole 16x16 tiles of the output, so round the + // output size up to a multiple of that. + const int tile = 16; + const int buf_w = ((out_w + tile - 1) / tile) * tile; + const int buf_h = ((out_h + tile - 1) / tile) * tile; + + Buffer out_cuda(buf_w, buf_h, 3), out_tensorcore(buf_w, buf_h, 3); resize_cudaonly(input, scale_factor, out_cuda); resize_tensorcore(input, scale_factor, out_tensorcore); diff --git a/src/ExtractWMMAOperations.cpp b/src/ExtractWMMAOperations.cpp index 385b9d791b1c..76a4a9624d17 100644 --- a/src/ExtractWMMAOperations.cpp +++ b/src/ExtractWMMAOperations.cpp @@ -480,6 +480,12 @@ class ExtractWMMAOperations : public IRMutator { ScopedValue old_tile_name(tile_name, op->name); ScopedValue old_in_alloc(in_allocate, true); + // Each accumulator allocation has its own shape and its own set of + // sub-tiles. + ScopedValue old_found_shape(found_shape, false); + ScopedValue old_shape(shape, Shape{}); + ScopedValue> old_subtiles(subtiles, {}); + // In the first pass we recognize the matrix multiplies, which is what // tells us the tile shape. In the second we recognize the // zero-initializations and the stores out to memory, both of which @@ -584,7 +590,10 @@ class ExtractWMMAOperations : public IRMutator { matmul.shape.N == shape.N && matmul.shape.K == shape.K)) << "Found inconsistent tile shapes for a WMMAAccumulator allocation across " - << "multiple matrix multiplies that store to it."; + << "multiple matrix multiplies that store to it: " + << shape.M << "x" << shape.N << "x" << shape.K << " vs " + << matmul.shape.M << "x" << matmul.shape.N << "x" << matmul.shape.K + << "."; shape = matmul.shape; found_shape = true; return matmul.stmt; From dd2d9b2537245b5a4ed6c167689d4d0b25df0185 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Wed, 29 Jul 2026 09:21:56 -0700 Subject: [PATCH 20/59] Rename WMMAAccumulator to WMMAFragment and let operands be staged in it The memory type now covers all three matrices of a multiply, with the role inferred from use: an allocation accumulated into by a matrix multiply is the accumulator, and one read as an operand is that operand. An operand staged this way is loaded into fragment registers where the schedule says to compute it, and reused by every multiply that reads it, which is a hoist that nothing downstream can do when the loop isn't unrolled. The pass now tracks a scope of fragments across two passes over the whole statement rather than one allocation at a time, so fragments can nest and several multiplies can be in flight at once. Co-Authored-By: Claude Opus 5 --- apps/tensorcore_matmul/matmul_generator.cpp | 128 ++-- apps/tensorcore_resize/resize_generator.cpp | 4 +- src/CanonicalizeGPUVars.cpp | 4 +- src/Deserialization.cpp | 4 +- src/Expr.h | 17 +- src/ExtractWMMAOperations.cpp | 609 ++++++++++++-------- src/ExtractWMMAOperations.h | 2 +- src/FuseGPUThreadLoops.cpp | 8 +- src/IRPrinter.cpp | 4 +- src/LowerWarpShuffles.cpp | 2 +- src/Serialization.cpp | 4 +- src/halide_ir.fbs | 2 +- test/correctness/wmma_matmul.cpp | 206 ++++++- 13 files changed, 702 insertions(+), 292 deletions(-) diff --git a/apps/tensorcore_matmul/matmul_generator.cpp b/apps/tensorcore_matmul/matmul_generator.cpp index e5bfa3844a05..54fe01ffc605 100644 --- a/apps/tensorcore_matmul/matmul_generator.cpp +++ b/apps/tensorcore_matmul/matmul_generator.cpp @@ -20,13 +20,18 @@ class MatMul : public Halide::Generator { GeneratorParam K{"K", 1024}; // How many tensor core tiles of accumulator each warp holds, and how many - // warps there are per block. - GeneratorParam tiles_x{"tiles_x", 5}; - GeneratorParam tiles_y{"tiles_y", 4}; - GeneratorParam warps{"warps", 2}; - // How many reduction steps to unroll, which puts more operand loads in - // flight at once. - GeneratorParam k_unroll{"k_unroll", 2}; + // warps there are per block in each dimension. + GeneratorParam tiles_x{"tiles_x", 4}; + GeneratorParam tiles_y{"tiles_y", 5}; + GeneratorParam warps_x{"warps_x", 2}; + GeneratorParam warps_y{"warps_y", 1}; + // How much of the reduction is staged in shared memory at a time. + GeneratorParam block_k{"block_k", 32}; + // Extra elements per row of the shared panels, which spreads consecutive + // rows across different banks. A multiple of eight keeps the rows aligned + // enough for the widest asynchronous copy. + GeneratorParam pad_a{"pad_a", 8}; + GeneratorParam pad_b{"pad_b", 8}; Input> matA{"matA"}; // K x M Input> matB{"matB"}; // N x K @@ -36,8 +41,14 @@ class MatMul : public Halide::Generator { void generate() { k = RDom(0, K, "k"); + // Wrappers for the operands, so that the tensor core schedule can + // stage them through shared memory. Left inline, they are just matA + // and matB. + As(kk, y) = matA(kk, y); + Bs(x, kk) = matB(x, kk); + prod(x, y) = 0.f; - prod(x, y) += cast(matA(k, y)) * cast(matB(x, k)); + prod(x, y) += cast(As(k, y)) * cast(Bs(x, k)); output(x, y) = prod(x, y); } @@ -79,58 +90,103 @@ class MatMul : public Halide::Generator { // accumulates at once. Each operand tile loaded feeds tiles_x (or // tiles_y) multiplies, so this is what gets us reuse out of the // loads. - const int tile_x = 16, tile_y = 16, tile_k = 16; + const int tile = 16; + const int block_x = tile * tiles_x * warps_x; + const int block_y = tile * tiles_y * warps_y; - Var xi("xi"), yi("yi"), xt("xt"), mmxi("mmxi"), mmyi("mmyi"); - Var rxi("rxi"), ryi("ryi"); - RVar rro("rro"), rri("rri"), rru("rru"); + Var xi("xi"), yi("yi"), xt("xt"), yt("yt"), mmxi("mmxi"), mmyi("mmyi"); + Var xw("xw"), yw("yw"), rxi("rxi"), ryi("ryi"); + RVar ko("ko"), ki("ki"), rri("rri"); output.bound(x, 0, N) .bound(y, 0, M) - .split(x, x, xi, tile_x * tiles_x * warps) - .split(xi, xt, xi, tile_x * tiles_x) - .split(xi, xi, mmxi, tile_x) - .split(y, y, yi, tile_y * tiles_y) - .split(yi, yi, mmyi, tile_y) + .split(x, x, xi, block_x) + .split(xi, xt, xi, tile * tiles_x) + .split(xi, xi, mmxi, tile) + .split(y, y, yi, block_y) + .split(yi, yt, yi, tile * tiles_y) + .split(yi, yi, mmyi, tile) .gpu_blocks(x, y) - .gpu_threads(xt) - .reorder(mmxi, mmyi, xi, yi, xt, x, y) + .gpu_threads(xt, yt) + .reorder(mmxi, mmyi, xi, yi, xt, yt, x, y) .unroll(xi) .unroll(yi) .vectorize(mmxi) .vectorize(mmyi); // The accumulators live in tensor core registers for the whole - // reduction, and are written out to memory once at the end. - prod.compute_at(output, xt) - .store_in(MemoryType::WMMAAccumulator) - .split(x, x, rxi, tile_x) - .split(y, y, ryi, tile_y) + // reduction, and are written out to memory once at the end. They + // sit at block level so that the reduction loop can be above the + // loop over warps, which lets every warp share one staged panel. + prod.compute_at(output, x) + .store_in(MemoryType::WMMAFragment) + .split(x, xw, xi, tile * tiles_x) + .split(xi, xi, rxi, tile) + .split(y, yw, yi, tile * tiles_y) + .split(yi, yi, ryi, tile) + .reorder(rxi, ryi, xi, yi, xw, yw) + .gpu_threads(xw, yw) .vectorize(rxi) .vectorize(ryi) - .unroll(x) - .unroll(y); + .unroll(xi) + .unroll(yi); prod.update() - .split(x, x, rxi, tile_x) - .split(y, y, ryi, tile_y) - .split(k, rro, rri, tile_k) - .split(rro, rro, rru, k_unroll) - .reorder(rri, rxi, ryi, x, y, rru, rro) - .unroll(x) - .unroll(y) - .unroll(rru) + .split(k, ko, ki, block_k) + .split(x, xw, xi, tile * tiles_x) + .split(xi, xi, rxi, tile) + .split(y, yw, yi, tile * tiles_y) + .split(yi, yi, ryi, tile) + .split(ki, ki, rri, tile) + .reorder(rri, rxi, ryi, xi, yi, ki, xw, yw, ko) + .gpu_threads(xw, yw) + .unroll(xi) + .unroll(yi) + .unroll(ki) .atomic() .vectorize(rri) .vectorize(rxi) .vectorize(ryi); + + // Stage the operand panels into shared memory once per reduction + // step, to be shared by every warp in the block. Each thread moves + // sixteen bytes at a time along the dense dimension, so that the + // reads from global memory coalesce and the writes to shared + // memory can be done as asynchronous copies. + const int vec = 8; + Var kko("kko"), kki("kki"), xxo("xxo"), xxi("xxi"); + Var t("t"), ti("ti"), tw("tw"), tw2("tw2"), to("to"); + + As.compute_at(prod, ko) + .store_in(MemoryType::GPUShared) + .align_storage(kk, (int)block_k + (int)pad_a) + .split(kk, kko, kki, vec) + .fuse(kko, y, t) + .split(t, t, ti, 32) + .split(t, t, tw, warps_x) + .split(t, to, tw2, warps_y) + .gpu_lanes(ti) + .gpu_threads(tw, tw2) + .vectorize(kki); + + Bs.compute_at(prod, ko) + .store_in(MemoryType::GPUShared) + .align_storage(x, block_x + (int)pad_b) + .split(x, xxo, xxi, vec) + .fuse(xxo, kk, t) + .split(t, t, ti, 32) + .split(t, t, tw, warps_x) + .split(t, to, tw2, warps_y) + .gpu_lanes(ti) + .gpu_threads(tw, tw2) + .vectorize(xxi); } } private: - Var x{"x"}, y{"y"}; + Var x{"x"}, y{"y"}, kk{"kk"}; RDom k; - Func prod{"prod"}; + Func prod{"prod"}, As{"As"}, Bs{"Bs"}; }; } // namespace diff --git a/apps/tensorcore_resize/resize_generator.cpp b/apps/tensorcore_resize/resize_generator.cpp index 7794f5f35e26..d434cfcffcf1 100644 --- a/apps/tensorcore_resize/resize_generator.cpp +++ b/apps/tensorcore_resize/resize_generator.cpp @@ -252,7 +252,7 @@ class Resize : public Halide::Generator { // An 8x32 tile of accumulator, reducing 16 taps at a time. resized_y.compute_at(resized_y.in(), xio) - .store_in(MemoryType::WMMAAccumulator) + .store_in(MemoryType::WMMAFragment) .unroll(c) .vectorize(x, 32) .unroll(x) @@ -284,7 +284,7 @@ class Resize : public Halide::Generator { RVar ri("ri"), ro("ro"); resized_x - .store_in(MemoryType::WMMAAccumulator) + .store_in(MemoryType::WMMAFragment) .compute_at(resized_x.in(), c) .vectorize(x) .vectorize(y) diff --git a/src/CanonicalizeGPUVars.cpp b/src/CanonicalizeGPUVars.cpp index f844b3ff50ec..41dee1c8111a 100644 --- a/src/CanonicalizeGPUVars.cpp +++ b/src/CanonicalizeGPUVars.cpp @@ -82,7 +82,7 @@ class CountGPUBlocksThreads : public IRVisitor { // extract_wmma_operations will wrap the statements that touch this // allocation in loops over the lanes of a warp, so count it as a lane // dimension. - const bool wmma = op->memory_type == MemoryType::WMMAAccumulator; + const bool wmma = op->memory_type == MemoryType::WMMAFragment; int dl = wmma && !in_lanes; ScopedValue old_in_lanes(in_lanes, in_lanes || wmma); nl += dl; @@ -120,7 +120,7 @@ class CanonicalizeGPUVars : public IRMutator { Stmt visit(const Realize *op) override { ScopedValue old(in_wmma_alloc, in_wmma_alloc || - op->memory_type == MemoryType::WMMAAccumulator); + op->memory_type == MemoryType::WMMAFragment); return IRMutator::visit(op); } diff --git a/src/Deserialization.cpp b/src/Deserialization.cpp index 016df76cf42a..341194622b33 100644 --- a/src/Deserialization.cpp +++ b/src/Deserialization.cpp @@ -190,8 +190,8 @@ MemoryType Deserializer::deserialize_memory_type(Serialize::MemoryType memory_ty return MemoryType::VTCM; case Serialize::MemoryType::AMXTile: return MemoryType::AMXTile; - case Serialize::MemoryType::WMMAAccumulator: - return MemoryType::WMMAAccumulator; + case Serialize::MemoryType::WMMAFragment: + return MemoryType::WMMAFragment; default: user_error << "unknown memory type " << (int)memory_type << "\n"; return MemoryType::Auto; diff --git a/src/Expr.h b/src/Expr.h index 5c22e27d911b..679c76e1c019 100644 --- a/src/Expr.h +++ b/src/Expr.h @@ -414,12 +414,17 @@ enum class MemoryType { * no such instruction this is ordinary shared memory. */ GPUSharedAsync, - /** An NVIDIA tensor core accumulator fragment. The storage is striped - * across the registers of the 32 lanes of a warp in a layout that is not + /** An NVIDIA tensor core matrix fragment. The storage is striped across + * the registers of the 32 lanes of a warp in a layout that is not * architecturally specified, so the only legal accesses are the ones - * recognized by the WMMA lowering pass: zero-initialization, accumulation - * of a matrix multiply, and copying the tile out to memory. */ - WMMAAccumulator, + * recognized by the WMMA lowering pass. Which of the three roles a + * fragment plays - the accumulator, or either operand of the multiply - + * follows from how it is used, and determines what those accesses are. + * An accumulator can be zero-initialized, initialized from a matrix in + * memory, accumulated into by a matrix multiply, and copied back out to + * memory. An operand can be filled from a matrix in memory and read by a + * matrix multiply. */ + WMMAFragment, }; /** Whether a MemoryType places an allocation in GPU shared memory. */ @@ -434,7 +439,7 @@ inline bool is_gpu_shared(MemoryType t) { * dedicated lowering pass that requires the original 2D-shaped loads * and stores to remain intact. */ inline bool is_tile_memory_type(MemoryType t) { - return t == MemoryType::AMXTile || t == MemoryType::WMMAAccumulator; + return t == MemoryType::AMXTile || t == MemoryType::WMMAFragment; } namespace Internal { diff --git a/src/ExtractWMMAOperations.cpp b/src/ExtractWMMAOperations.cpp index 76a4a9624d17..67bfcd35f532 100644 --- a/src/ExtractWMMAOperations.cpp +++ b/src/ExtractWMMAOperations.cpp @@ -1,5 +1,7 @@ #include "ExtractWMMAOperations.h" +#include + #include "CanonicalizeGPUVars.h" #include "FindIntrinsics.h" #include "IREquality.h" @@ -19,18 +21,24 @@ * with the wmma load and store instructions, which move a tile between the * registers of a warp and a 2D array in memory. * - * Accordingly this pass recognizes exactly three operations on a - * WMMAAccumulator allocation: + * A WMMAFragment allocation holds one of the three matrices of a multiply, and + * which one follows from how it is used. An allocation accumulated into by a + * matrix multiply is the accumulator; one read as an operand of a multiply is + * that operand. The role determines which accesses are legal: * - * 1) Zero-initialization. This one is layout-independent, so it stays a plain - * store of zero to the (shrunken) allocation. - * 2) Accumulation of a matrix multiply, which becomes a pair of wmma loads - * feeding a wmma mma. - * 3) Copying a tile out to memory, which becomes a wmma store. + * - An accumulator may be zero-initialized, filled from a matrix in memory, + * accumulated into by a matrix multiply, and copied back out to memory. + * - An operand may be filled from a matrix in memory and read by a matrix + * multiply. * - * Anything else is an error. Each of those operations is wrapped in a loop over - * the 32 lanes of a warp, because nothing in the schedule says the tile is + * Anything else is an error. Every one of those operations is wrapped in a loop + * over the 32 lanes of a warp, because nothing in the schedule says the tile is * spread over a warp - that's a consequence of asking for tensor core storage. + * + * Operands don't have to be staged in fragments. If a multiply reads its + * operand straight out of shared or global memory, the load into registers is + * synthesized at the multiply instead. Staging one explicitly is how a schedule + * says to load it once and reuse it across several multiplies. */ namespace Halide { @@ -54,22 +62,48 @@ enum class Layout { Col, }; -// Every warp is 32 lanes, and every tile shape we support holds 256 elements, -// so each lane holds 8 of them. -constexpr int warp_lanes = 32; -constexpr int fragment_elements = 8; - -// The Halide type we use to represent an a or b fragment. Every operand -// fragment is 8 32-bit registers per lane, which is more matrix elements than -// there are for some shapes, because the hardware replicates elements across -// lanes for those. -Type fragment_type(Type element_type) { - return element_type.with_lanes(16); +// Which of the three matrices of the multiply a fragment holds. +enum class Role { + Unknown, + A, + B, + Accumulator, +}; + +Call::IntrinsicOp intrinsic_for_role(Role role) { + switch (role) { + case Role::A: + return Call::wmma_matrix_to_fragment_a; + case Role::B: + return Call::wmma_matrix_to_fragment_b; + default: + return Call::wmma_matrix_to_fragment_c; + } } -// The Halide type we use to represent an accumulator fragment. -Type accumulator_fragment_type(Type element_type) { - return element_type.with_lanes(fragment_elements); +const char *role_name(Role role) { + switch (role) { + case Role::A: + return "the first operand"; + case Role::B: + return "the second operand"; + case Role::Accumulator: + return "the accumulator"; + default: + return "no part"; + } +} + +// Every warp is 32 lanes. An accumulator tile holds 256 elements, so each lane +// holds 8 of them. The hardware hands back operand fragments as 8 32-bit +// registers per lane whatever the shape, replicating elements across lanes for +// the shapes that hold fewer, so those are 16 16-bit elements per lane. +constexpr int warp_lanes = 32; +constexpr int accumulator_elements = 8; +constexpr int operand_elements = 16; + +int elements_per_lane(Role role) { + return role == Role::Accumulator ? accumulator_elements : operand_elements; } // One operand of a matrix multiply, described in the canonical [K, N, M] @@ -151,31 +185,35 @@ Expr make_matrix_address(const string &name, Type element_type, const Expr &base } // The shape of the matrix each fragment is taken out of. -void fragment_matrix_shape(Call::IntrinsicOp intrin, const Shape &shape, - int *rows, int *cols) { - *rows = intrin == Call::wmma_matrix_to_fragment_b ? shape.K : shape.M; - *cols = intrin == Call::wmma_matrix_to_fragment_a ? shape.K : shape.N; +void fragment_matrix_shape(Role role, const Shape &shape, int *rows, int *cols) { + *rows = role == Role::B ? shape.K : shape.M; + *cols = role == Role::A ? shape.K : shape.N; } -Expr make_matrix_to_fragment(Call::IntrinsicOp intrin, const Shape &shape, Layout layout, +Expr make_matrix_to_fragment(Role role, const Shape &shape, Layout layout, const Load *load, const Expr &base, const Expr &stride) { int rows, cols; - fragment_matrix_shape(intrin, shape, &rows, &cols); + fragment_matrix_shape(role, shape, &rows, &cols); Expr address = make_matrix_address(load->name, load->type.element_of(), base, rows, cols, layout, stride, load->image, load->param); - Type type = intrin == Call::wmma_matrix_to_fragment_c ? - accumulator_fragment_type(load->type.element_of()) : - fragment_type(load->type.element_of()); - return Call::make(type, intrin, {shape.M, shape.N, shape.K, std::move(address)}, + Type type = load->type.element_of().with_lanes(elements_per_lane(role)); + return Call::make(type, intrinsic_for_role(role), + {shape.M, shape.N, shape.K, std::move(address)}, Call::Intrinsic); } -struct Matmul { - Stmt stmt; +// A store to a fragment recognized as the accumulation of a matrix multiply, +// broken down into the pieces the multiply is built from. +struct MatmulInfo { Shape shape; + Operand lhs, rhs; + Layout lhs_layout = Layout::Row, rhs_layout = Layout::Row; + Expr lda, ldb; + Type accumulator_type; + vector> peeled_lets; }; -Matmul convert_to_matmul(const Store *op, const string &new_name) { +MatmulInfo analyze_matmul(const Store *op) { // We expect the pattern: // // out[idx] = reduce_add(widen(lhs[multiramp]) * widen(rhs[multiramp])) + out[idx] @@ -183,20 +221,21 @@ Matmul convert_to_matmul(const Store *op, const string &new_name) { // Though either operand may have been hoisted out to a broadcast or had a // lane permutation left on it by vectorization. - auto fail = [&](const char *reason) -> Matmul { - user_error << "Matrix multiply not recognized. Store to a WMMAAccumulator " - << "allocation must be a zero-initialization or a sum of a vector " - << "reduce op and a load from the same allocation. In the following " - << "store, " << reason << ".\n" + auto fail = [&](const char *reason) -> MatmulInfo { + user_error << "Matrix multiply not recognized. Store to a WMMAFragment " + << "allocation must be a zero-initialization, a fill from a matrix " + << "in memory, or a sum of a vector reduce op and a load from the " + << "same allocation. In the following store, " << reason << ".\n" << Stmt(op); - return Matmul{}; + return MatmulInfo{}; }; + MatmulInfo info; + // Peel lets - vector> peeled_lets; Expr value = op->value; while (const Let *let = value.as()) { - peeled_lets.emplace_back(let->name, let->value); + info.peeled_lets.emplace_back(let->name, let->value); value = let->body; } @@ -229,6 +268,7 @@ Matmul convert_to_matmul(const Store *op, const string &new_name) { !(reduce->type.bits() == 32 || reduce->type.bits() == 16)) { return fail("the accumulator type is not 32-bit or 16-bit float"); } + info.accumulator_type = reduce->type.element_of(); // The vector reduce must be of a widening multiply. FindIntrinsics does // not lift float widening muls, so we just expect a multiply of two casts. @@ -240,17 +280,16 @@ Matmul convert_to_matmul(const Store *op, const string &new_name) { // Under the casts, broadcasts and lane permutations that vectorization may // have left on each operand there must be a load. Scope empty_scope; - Operand lhs_op, rhs_op; - lhs_op.load = is_load_of_multiramp(mul->a, empty_scope, &lhs_op.mr); - rhs_op.load = is_load_of_multiramp(mul->b, empty_scope, &rhs_op.mr); - if (!lhs_op.load || !rhs_op.load) { + info.lhs.load = is_load_of_multiramp(mul->a, empty_scope, &info.lhs.mr); + info.rhs.load = is_load_of_multiramp(mul->b, empty_scope, &info.rhs.mr); + if (!info.lhs.load || !info.rhs.load) { return fail("the matrix multiply operands are not loads with affine indices"); } - if (!is_const_one(lhs_op.load->predicate) || !is_const_one(rhs_op.load->predicate)) { + if (!is_const_one(info.lhs.load->predicate) || !is_const_one(info.rhs.load->predicate)) { return fail("the matrix multiply operands are predicated loads"); } - if (lhs_op.load->type.element_of() != Float(16) || - rhs_op.load->type.element_of() != Float(16)) { + if (info.lhs.load->type.element_of() != Float(16) || + info.rhs.load->type.element_of() != Float(16)) { return fail("the matrix multiply operands are not both float16"); } @@ -270,23 +309,23 @@ Matmul convert_to_matmul(const Store *op, const string &new_name) { // Deduce which operand is which and what tile shape this is by trying each // supported shape and seeing which one the access patterns fit. const Shape *shape = nullptr; - Layout lhs_layout = Layout::Row, rhs_layout = Layout::Row; - Expr lda, ldb; for (const Shape &candidate : supported_shapes) { if (candidate.M * candidate.N != MN || candidate.K != K) { continue; } vector canonical_shape{candidate.K, candidate.N, candidate.M}; - if (!lhs_op.mr.strides_for_shape(canonical_shape, &lhs_op.strides) || - !rhs_op.mr.strides_for_shape(canonical_shape, &rhs_op.strides)) { + if (!info.lhs.mr.strides_for_shape(canonical_shape, &info.lhs.strides) || + !info.rhs.mr.strides_for_shape(canonical_shape, &info.rhs.strides)) { continue; } - if (is_lhs(lhs_op, &lhs_layout, &lda) && is_rhs(rhs_op, &rhs_layout, &ldb)) { + if (is_lhs(info.lhs, &info.lhs_layout, &info.lda) && + is_rhs(info.rhs, &info.rhs_layout, &info.ldb)) { shape = &candidate; break; } - if (is_lhs(rhs_op, &lhs_layout, &lda) && is_rhs(lhs_op, &rhs_layout, &ldb)) { - std::swap(lhs_op, rhs_op); + if (is_lhs(info.rhs, &info.lhs_layout, &info.lda) && + is_rhs(info.lhs, &info.rhs_layout, &info.ldb)) { + std::swap(info.lhs, info.rhs); shape = &candidate; break; } @@ -297,82 +336,19 @@ Matmul convert_to_matmul(const Store *op, const string &new_name) { "multiply of a tile shape the tensor cores support (16x16x16, " "32x8x16, or 8x32x16)"); } - - // Build the wmma intrinsics. - Expr a = make_matrix_to_fragment(Call::wmma_matrix_to_fragment_a, *shape, lhs_layout, - lhs_op.load, lhs_op.mr.base, lda); - Expr b = make_matrix_to_fragment(Call::wmma_matrix_to_fragment_b, *shape, rhs_layout, - rhs_op.load, rhs_op.mr.base, ldb); - - Type acc_type = accumulator_fragment_type(reduce->type.element_of()); - Expr frag_idx = Ramp::make(0, 1, fragment_elements); - Expr c = Load::make(acc_type, new_name, frag_idx, {}, {}, - const_true(fragment_elements), {}); - - Expr mma = Call::make(acc_type, Call::wmma_mma, - {shape->M, shape->N, shape->K, - (int)lhs_layout, (int)rhs_layout, - std::move(a), std::move(b), std::move(c)}, - Call::Intrinsic); - - Stmt store = in_lane_loop( - Store::make(new_name, std::move(mma), frag_idx, Parameter(), - const_true(fragment_elements), ModulusRemainder())); - for (auto &[name, v] : reverse_view(peeled_lets)) { - store = LetStmt::make(name, std::move(v), store); - } - return {std::move(store), *shape}; + info.shape = *shape; + return info; } -// Whether a store to an accumulator is its initialization, as opposed to a -// matrix multiply accumulating into it. -bool is_initialization(const Store *op, const string &tile_name) { +// Whether a store to a fragment fills it, as opposed to a matrix multiply +// accumulating into it. +bool is_fill(const Store *op) { if (is_const_zero(op->value)) { return true; } MultiRamp mr; const Load *load = is_load_of_multiramp(op->value, Scope::empty_scope(), &mr); - return load && load->name != tile_name; -} - -Stmt convert_to_init(const Store *op, const string &new_name, const Shape &shape) { - Type element_type = op->value.type().element_of(); - Expr value; - if (is_const_zero(op->value)) { - // Zeroing an accumulator is layout-independent, so it doesn't need an - // instruction - the registers just get set to zero. - value = make_zero(accumulator_fragment_type(element_type)); - } else { - auto fail = [&](const char *reason) { - user_error << "Initialization of a tensor core accumulator not supported. " - << reason << ".\n" - << Stmt(op); - return Expr{}; - }; - - MultiRamp mr; - const Load *matrix = is_load_of_multiramp(op->value, Scope::empty_scope(), &mr); - internal_assert(matrix); // is_initialization checked this - if (matrix->type.element_of() != element_type) { - value = fail("An accumulator can only be initialized from a matrix of the " - "same type, because the hardware does not convert on the way in"); - } else if (!is_const_one(matrix->predicate)) { - value = fail("The load is predicated"); - } else { - WMMAMatrixLayout mem; - if (!wmma_matrix_layout(mr, shape.M, shape.N, &mem)) { - value = fail("The matrix loaded from is not a dense tile of the right shape"); - } else { - value = make_matrix_to_fragment( - Call::wmma_matrix_to_fragment_c, shape, - mem.row_major ? Layout::Row : Layout::Col, matrix, mem.base, mem.stride); - } - } - } - Expr frag_idx = Ramp::make(0, 1, fragment_elements); - return in_lane_loop( - Store::make(new_name, std::move(value), frag_idx, Parameter(), - const_true(fragment_elements), ModulusRemainder())); + return load && load->name != op->name; } Stmt convert_to_tile_store(const Store *op, const Expr &store_index, @@ -402,9 +378,9 @@ Stmt convert_to_tile_store(const Store *op, const Expr &store_index, // accumulator. Expr index = make_matrix_index(mem.base, shape.M, shape.N, layout, mem.stride); Type element_type = op->value.type().element_of(); - Expr frag = Load::make(accumulator_fragment_type(element_type), new_name, - Ramp::make(0, 1, fragment_elements), {}, {}, - const_true(fragment_elements), {}); + Expr frag = Load::make(element_type.with_lanes(accumulator_elements), new_name, + Ramp::make(0, 1, accumulator_elements), {}, {}, + const_true(accumulator_elements), {}); const int lanes = shape.M * shape.N; Expr matrix = Call::make(element_type.with_lanes(lanes), Call::wmma_fragment_to_matrix_d, {shape.M, shape.N, shape.K, std::move(frag)}, @@ -416,24 +392,53 @@ Stmt convert_to_tile_store(const Store *op, const Expr &store_index, op->is_streaming)); } +// Everything we learn about one WMMAFragment allocation from the way it is +// used. The role and the shape come from the matrix multiplies it takes part +// in, so neither is known until those have been found. +struct Fragment { + string name; + // The prefix of the names of the per-fragment allocations this one becomes. + string fragment_name; + Type element_type; + Role role = Role::Unknown; + Shape shape{}; + bool found_shape = false; + // An allocation may hold several fragments as disjoint sub-tiles, each of + // which becomes its own allocation. + vector subtiles; + + Type value_type() const { + return element_type.with_lanes(elements_per_lane(role)); + } +}; + class ExtractWMMAOperations : public IRMutator { using IRMutator::visit; - string tile_name; - string wmma_name; - int pass = 0; - bool in_allocate = false; - bool found_shape = false; - Shape shape{}; + // Everything we've learned about each fragment allocation, and the names of + // the ones we're currently inside of. The records outlive the first pass so + // that the second one can use them. + std::map fragments; + vector in_scope; - // A WMMAAccumulator allocation may hold several accumulator fragments as - // 2D sub-tiles. This tracks them. - vector subtiles; + // In the first pass we recognize the matrix multiplies, which is what tells + // us what role each fragment plays and what shape it is. In the second we + // rewrite everything, which needs to know both. + int pass = 0; // The loops over GPU blocks, threads, and lanes we're inside of. vector gpu_loop_vars; - // An accumulator allocation may sit outside the loops over GPU threads, in + Fragment *find_fragment(const string &name) { + for (const string &n : in_scope) { + if (n == name) { + return &fragments[name]; + } + } + return nullptr; + } + + // A fragment allocation may sit outside the loops over GPU threads, in // which case each thread gets its own copy of it and only ever touches its // own slice. Any dependence of the index on the thread is selecting between // those copies, not between subtiles within one, so drop it. @@ -445,11 +450,155 @@ class ExtractWMMAOperations : public IRMutator { return simplify(idx); } - string get_subtile_name(const Expr &index) { - int idx = Halide::Internal::get_subtile(index_within_thread(index), - "tensor core accumulator", &subtiles); + string subtile_name(Fragment *f, const Expr &index) { + int idx = get_subtile(index_within_thread(index), + "tensor core fragment", &f->subtiles); internal_assert(idx >= 0); // errors handled already - return wmma_name + std::to_string(idx); + return f->fragment_name + std::to_string(idx); + } + + void set_role(Fragment *f, Role role) { + user_assert(f->role == Role::Unknown || f->role == role) + << "The tensor core fragment " << f->name << " is used as both " + << role_name(f->role) << " and " << role_name(role) << " of a matrix " + << "multiply. Those are held in registers in different layouts, so a " + << "fragment can only play one of those roles. Stage it through memory " + << "in between.\n"; + f->role = role; + } + + void set_shape(Fragment *f, const Shape &shape) { + user_assert(!f->found_shape || + (shape.M == f->shape.M && shape.N == f->shape.N && + shape.K == f->shape.K)) + << "Found inconsistent tile shapes for the tensor core fragment " + << f->name << " across the matrix multiplies that use it: " + << f->shape.M << "x" << f->shape.N << "x" << f->shape.K << " vs " + << shape.M << "x" << shape.N << "x" << shape.K << "."; + f->shape = shape; + f->found_shape = true; + } + + // Which subtile of a fragment an access refers to, as an index over just + // the tile. An operand is read by the multiply at the shape of the whole + // reduction, with one axis broadcast, so its index has to be projected back + // down to the tile before it can be compared against the fill that wrote + // it. + string operand_subtile_name(Fragment *f, const Expr &base, Role role, + const Shape &shape, Layout layout, const Expr &stride) { + int rows, cols; + fragment_matrix_shape(role, shape, &rows, &cols); + return subtile_name(f, make_matrix_index(base, rows, cols, layout, stride)); + } + + // Note what a matrix multiply tells us about an operand staged in a + // fragment. An operand read straight out of memory tells us nothing, + // because its load gets synthesized at the multiply. + void record_operand(const Operand &operand, Role role, const Shape &shape, + Layout layout, const Expr &stride) { + if (Fragment *f = find_fragment(operand.load->name)) { + set_role(f, role); + set_shape(f, shape); + operand_subtile_name(f, operand.mr.base, role, shape, layout, stride); + } + } + + // The value a matrix multiply uses for one of its operands: the fragment it + // was staged in, or a load synthesized here if it wasn't staged. + Expr operand_value(const Operand &operand, Role role, const Shape &shape, + Layout layout, const Expr &stride) { + if (Fragment *f = find_fragment(operand.load->name)) { + const int lanes = elements_per_lane(role); + const string name = + operand_subtile_name(f, operand.mr.base, role, shape, layout, stride); + return Load::make(f->value_type(), name, Ramp::make(0, 1, lanes), {}, {}, + const_true(lanes), {}); + } + return make_matrix_to_fragment(role, shape, layout, operand.load, + operand.mr.base, stride); + } + + Stmt convert_to_fill(const Store *op, Fragment *f) { + int rows, cols; + fragment_matrix_shape(f->role, f->shape, &rows, &cols); + MultiRamp dest_mr; + WMMAMatrixLayout dest; + user_assert(is_multiramp(op->index, Scope::empty_scope(), &dest_mr) && + wmma_matrix_layout(dest_mr, rows, cols, &dest)) + << "A tensor core fragment must be filled a whole tile at a time, but " + << "this fill is not to a dense " << rows << "x" << cols << " tile of " + << f->name << ".\n" + << Stmt(op); + const string name = subtile_name( + f, make_matrix_index(dest.base, rows, cols, + dest.row_major ? Layout::Row : Layout::Col, dest.stride)); + const int lanes = elements_per_lane(f->role); + Expr value; + if (is_const_zero(op->value)) { + // Zeroing a fragment is layout-independent, so it doesn't need an + // instruction - the registers just get set to zero. + value = make_zero(f->value_type()); + } else { + auto fail = [&](const char *reason) { + user_error << "Fill of a tensor core fragment not supported. " + << reason << ".\n" + << Stmt(op); + return Expr{}; + }; + + MultiRamp mr; + const Load *matrix = + is_load_of_multiramp(op->value, Scope::empty_scope(), &mr); + internal_assert(matrix); // is_fill checked this + int rows, cols; + fragment_matrix_shape(f->role, f->shape, &rows, &cols); + WMMAMatrixLayout mem; + if (find_fragment(matrix->name)) { + value = fail("A fragment can only be filled from a matrix in memory, " + "not from another fragment, because the layout in " + "registers is not known"); + } else if (matrix->type.element_of() != f->element_type) { + value = fail("A fragment can only be filled from a matrix of the same " + "type, because the hardware does not convert on the way in"); + } else if (!is_const_one(matrix->predicate)) { + value = fail("The load is predicated"); + } else if (!wmma_matrix_layout(mr, rows, cols, &mem)) { + value = fail("The matrix loaded from is not a dense tile of the right " + "shape"); + } else { + value = make_matrix_to_fragment( + f->role, f->shape, mem.row_major ? Layout::Row : Layout::Col, + matrix, mem.base, mem.stride); + } + } + return in_lane_loop( + Store::make(name, std::move(value), Ramp::make(0, 1, lanes), Parameter(), + const_true(lanes), ModulusRemainder())); + } + + Stmt convert_to_matmul(const Store *op, Fragment *f, const MatmulInfo &info) { + Expr a = operand_value(info.lhs, Role::A, info.shape, info.lhs_layout, info.lda); + Expr b = operand_value(info.rhs, Role::B, info.shape, info.rhs_layout, info.ldb); + + Type acc_type = info.accumulator_type.with_lanes(accumulator_elements); + Expr frag_idx = Ramp::make(0, 1, accumulator_elements); + const string name = subtile_name(f, op->index); + Expr c = Load::make(acc_type, name, frag_idx, {}, {}, + const_true(accumulator_elements), {}); + + Expr mma = Call::make(acc_type, Call::wmma_mma, + {info.shape.M, info.shape.N, info.shape.K, + (int)info.lhs_layout, (int)info.rhs_layout, + std::move(a), std::move(b), std::move(c)}, + Call::Intrinsic); + + Stmt store = in_lane_loop( + Store::make(name, std::move(mma), frag_idx, Parameter(), + const_true(accumulator_elements), ModulusRemainder())); + for (const auto &[let_name, v] : reverse_view(info.peeled_lets)) { + store = LetStmt::make(let_name, v, store); + } + return store; } Stmt visit(const For *op) override { @@ -463,62 +612,53 @@ class ExtractWMMAOperations : public IRMutator { } Stmt visit(const Allocate *op) override { - if (op->memory_type != MemoryType::WMMAAccumulator) { + if (op->memory_type != MemoryType::WMMAFragment) { return IRMutator::visit(op); } user_assert(op->type == Float(32) || op->type == Float(16)) - << "Tensor core accumulators must hold 32-bit or 16-bit floats, but " + << "Tensor core fragments must hold 32-bit or 16-bit floats, but " << op->name << " holds " << op->type << ".\n"; - user_assert(!in_allocate) - << "Already in a tensor core accumulator allocation at the allocation for " - << op->name << ". We do not currently support multiple nested tensor core " - << "matrix multiplies."; - - ScopedValue old_wmma_name(wmma_name, op->name + ".wmma."); - ScopedValue old_tile_name(tile_name, op->name); - ScopedValue old_in_alloc(in_allocate, true); - - // Each accumulator allocation has its own shape and its own set of - // sub-tiles. - ScopedValue old_found_shape(found_shape, false); - ScopedValue old_shape(shape, Shape{}); - ScopedValue> old_subtiles(subtiles, {}); - - // In the first pass we recognize the matrix multiplies, which is what - // tells us the tile shape. In the second we recognize the - // zero-initializations and the stores out to memory, both of which - // need to know the shape. - pass = 0; + Fragment &f = fragments[op->name]; + if (pass == 0) { + f.name = op->name; + f.fragment_name = op->name + ".wmma."; + f.element_type = op->type; + } + + in_scope.push_back(op->name); Stmt body = mutate(op->body); - user_assert(found_shape) - << op->name << " is stored in WMMAAccumulator memory, but no matrix " - << "multiply operation was found that stores to it, so the shape of the " - << "tile was unable to be determined.\n"; - pass = 1; - body = mutate(body); - - // Each fragment is one accumulator's worth of storage per lane. The - // allocations stay outside the loops over lanes, because a loop over - // lanes is a loop over threads, and register allocations outside a - // thread loop already get replicated per thread. - for (int i = 0; i < (int)subtiles.size(); i++) { - body = Allocate::make(wmma_name + std::to_string(i), op->type, - MemoryType::WMMAAccumulator, {fragment_elements}, + in_scope.pop_back(); + + if (pass == 0) { + user_assert(f.role != Role::Unknown) + << op->name << " is stored in WMMAFragment memory, but no matrix " + << "multiply was found that accumulates into it or reads it as an " + << "operand, so we can't tell what layout it should have.\n"; + return op; + } + + // Each fragment is one tile's worth of storage per lane. The allocations + // stay outside the loops over lanes, because a loop over lanes is a loop + // over threads, and register allocations outside a thread loop already + // get replicated per thread. + for (int i = 0; i < (int)f.subtiles.size(); i++) { + body = Allocate::make(f.fragment_name + std::to_string(i), f.element_type, + MemoryType::WMMAFragment, {elements_per_lane(f.role)}, const_true(), body); } return body; } Stmt visit(const Atomic *op) override { - if (op->producer_name == tile_name) { - // A tensor core accumulator is per-thread register storage, so - // there's nothing for another thread to race with. The atomic is - // there because the accumulator is scheduled outside the loops over - // threads, which makes it look shared. + if (find_fragment(op->producer_name)) { + // A tensor core fragment is per-thread register storage, so there's + // nothing for another thread to race with. The atomic is there + // because the fragment is scheduled outside the loops over threads, + // which makes it look shared. user_assert(op->mutex_name.empty()) - << "Accumulating into a tensor core accumulator should not need a " + << "Accumulating into a tensor core fragment should not need a " << "mutex.\n"; return mutate(op->body); } @@ -526,84 +666,93 @@ class ExtractWMMAOperations : public IRMutator { } Stmt visit(const Free *op) override { - if (op->name != tile_name) { + Fragment *f = find_fragment(op->name); + if (!f || pass == 0) { return op; } Stmt s; - for (int i = 0; i < (int)subtiles.size(); i++) { - Stmt f = Free::make(wmma_name + std::to_string(i)); - s = s.defined() ? Block::make(std::move(s), std::move(f)) : std::move(f); + for (int i = 0; i < (int)f->subtiles.size(); i++) { + Stmt free = Free::make(f->fragment_name + std::to_string(i)); + s = s.defined() ? Block::make(std::move(s), std::move(free)) : std::move(free); } return s; } Stmt visit(const ProducerConsumer *op) override { - if (op->name != tile_name) { + Fragment *f = find_fragment(op->name); + if (!f) { return IRMutator::visit(op); } - return ProducerConsumer::make(wmma_name, op->is_producer, mutate(op->body)); + return ProducerConsumer::make(f->fragment_name, op->is_producer, mutate(op->body)); } Expr visit(const Load *op) override { - user_assert(op->name != tile_name) - << "Tensor core accumulator " << tile_name - << " used outside a tensor core instruction"; + user_assert(!find_fragment(op->name)) + << "The tensor core fragment " << op->name + << " is used outside a tensor core instruction.\n"; return IRMutator::visit(op); } Stmt visit(const Store *op) override { - // There are three operations on an accumulator: - // 1) Zero-initialization - // 2) Matrix multiply - // 3) Stores to memory - // - // The matrix multiply is what tells us the tile shape, so we recognize - // those in the first pass and the other two in the second. - - if (op->name != tile_name) { + Fragment *f = find_fragment(op->name); + + if (!f) { + // A store to memory of a load from a fragment copies a tile out. Expr store_index; const Load *load = peel_store_permutations(op, &store_index).as(); - if (load && load->name == tile_name) { - return pass == 1 ? - convert_to_tile_store(op, store_index, - get_subtile_name(load->index), shape) : - Stmt(op); + Fragment *src = load ? find_fragment(load->name) : nullptr; + if (src) { + if (pass == 0) { + subtile_name(src, load->index); + return op; + } + user_assert(src->role == Role::Accumulator) + << "The tensor core fragment " << src->name << " is copied out to " + << "memory, but it holds " << role_name(src->role) << " of a matrix " + << "multiply. Only an accumulator can be copied out.\n"; + return convert_to_tile_store(op, store_index, + subtile_name(src, load->index), src->shape); } - // Not a copy of a tile out to memory. Recurse, so that any use of - // the accumulator buried in here gets reported as an error. + // Not a copy of a tile out to memory. Recurse, so that any use of a + // fragment buried in here gets reported as an error. return IRMutator::visit(op); } - string subtile_name = get_subtile_name(op->index); - - if (is_initialization(op, tile_name)) { - return pass == 1 ? convert_to_init(op, subtile_name, shape) : Stmt(op); + if (is_fill(op)) { + // A fill needs to know the role and the shape, so it waits for the + // second pass. It tells us neither, so there's nothing to record in + // the first. + return pass == 0 ? Stmt(op) : convert_to_fill(op, f); } - if (pass == 1) { + MatmulInfo info = analyze_matmul(op); + if (pass == 0) { + set_role(f, Role::Accumulator); + set_shape(f, info.shape); + subtile_name(f, op->index); + record_operand(info.lhs, Role::A, info.shape, info.lhs_layout, info.lda); + record_operand(info.rhs, Role::B, info.shape, info.rhs_layout, info.ldb); return op; } + return convert_to_matmul(op, f, info); + } - Matmul matmul = convert_to_matmul(op, subtile_name); - user_assert(!found_shape || - (matmul.shape.M == shape.M && - matmul.shape.N == shape.N && - matmul.shape.K == shape.K)) - << "Found inconsistent tile shapes for a WMMAAccumulator allocation across " - << "multiple matrix multiplies that store to it: " - << shape.M << "x" << shape.N << "x" << shape.K << " vs " - << matmul.shape.M << "x" << matmul.shape.N << "x" << matmul.shape.K - << "."; - shape = matmul.shape; - found_shape = true; - return matmul.stmt; +public: + void next_pass() { + pass = 1; } }; } // namespace Stmt extract_wmma_operations(const Stmt &s) { - return ExtractWMMAOperations()(s); + ExtractWMMAOperations mutator; + // The first pass only looks. What it learns about a fragment from one use + // of it is needed at all the others, including the ones it has already + // walked past. + mutator(s); + mutator.next_pass(); + return mutator(s); } bool is_wmma_intrinsic(const Call *op) { diff --git a/src/ExtractWMMAOperations.h b/src/ExtractWMMAOperations.h index d492021cbac5..79023f6193c8 100644 --- a/src/ExtractWMMAOperations.h +++ b/src/ExtractWMMAOperations.h @@ -15,7 +15,7 @@ namespace Internal { struct Call; struct Store; -/** Rewrite matrix multiplies that accumulate into WMMAAccumulator memory as +/** Rewrite matrix multiplies that accumulate into WMMAFragment memory as * calls to the wmma intrinsics understood by the PTX backend, and wrap them in * a loop over the 32 lanes of a warp. */ Stmt extract_wmma_operations(const Stmt &s); diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index 89097e579201..809bf784bd78 100644 --- a/src/FuseGPUThreadLoops.cpp +++ b/src/FuseGPUThreadLoops.cpp @@ -465,7 +465,7 @@ class ExtractSharedAndHeapAllocations : public IRMutator { op->memory_type != MemoryType::GPUTexture) || op->memory_type == MemoryType::Register || op->memory_type == MemoryType::Stack || - op->memory_type == MemoryType::WMMAAccumulator) { + op->memory_type == MemoryType::WMMAFragment) { // These allocations go in register or local memory return IRMutator::visit(op); } @@ -1143,7 +1143,7 @@ class ExtractRegisterAllocations : public IRMutator { op->memory_type == MemoryType::Register || op->memory_type == MemoryType::Heap || op->memory_type == MemoryType::Auto || - op->memory_type == MemoryType::WMMAAccumulator) + op->memory_type == MemoryType::WMMAFragment) << "Allocation " << op->name << " is scheduled inside a loop over GPU threads, so " << "it must live in stack memory, heap memory, or registers. " << "Shared allocations at this loop level are not yet supported.\n"; @@ -1414,7 +1414,7 @@ class InjectThreadBarriers : public IRMutator { case MemoryType::LockedCache: case MemoryType::VTCM: case MemoryType::AMXTile: - case MemoryType::WMMAAccumulator: + case MemoryType::WMMAFragment: break; } @@ -1441,7 +1441,7 @@ class InjectThreadBarriers : public IRMutator { case MemoryType::LockedCache: case MemoryType::VTCM: case MemoryType::AMXTile: - case MemoryType::WMMAAccumulator: + case MemoryType::WMMAFragment: break; } diff --git a/src/IRPrinter.cpp b/src/IRPrinter.cpp index e7d9c4c54b32..97bbbe8b82fa 100644 --- a/src/IRPrinter.cpp +++ b/src/IRPrinter.cpp @@ -172,8 +172,8 @@ std::ostream &operator<<(std::ostream &out, const MemoryType &t) { case MemoryType::AMXTile: out << "AMXTile"; break; - case MemoryType::WMMAAccumulator: - out << "WMMAAccumulator"; + case MemoryType::WMMAFragment: + out << "WMMAFragment"; break; } return out; diff --git a/src/LowerWarpShuffles.cpp b/src/LowerWarpShuffles.cpp index 8ce449c577f5..ce13d807b521 100644 --- a/src/LowerWarpShuffles.cpp +++ b/src/LowerWarpShuffles.cpp @@ -656,7 +656,7 @@ class LowerWarpShuffles : public IRMutator { if (this_lane.defined() || is_gpu_shared(op->memory_type) || op->memory_type == MemoryType::Heap || - op->memory_type == MemoryType::WMMAAccumulator) { + op->memory_type == MemoryType::WMMAFragment) { // Not an allocation for us to stripe. Warp-level storage is // per-lane register storage; shared and heap (global) memory are // never striped across lanes, and tensor core accumulators are diff --git a/src/Serialization.cpp b/src/Serialization.cpp index acc91504f022..bc85d7ce1683 100644 --- a/src/Serialization.cpp +++ b/src/Serialization.cpp @@ -160,8 +160,8 @@ Serialize::MemoryType Serializer::serialize_memory_type(const MemoryType &memory return Serialize::MemoryType::VTCM; case MemoryType::AMXTile: return Serialize::MemoryType::AMXTile; - case MemoryType::WMMAAccumulator: - return Serialize::MemoryType::WMMAAccumulator; + case MemoryType::WMMAFragment: + return Serialize::MemoryType::WMMAFragment; default: user_error << "Unsupported memory type\n"; return Serialize::MemoryType::Auto; diff --git a/src/halide_ir.fbs b/src/halide_ir.fbs index b5786596512b..59895465134a 100644 --- a/src/halide_ir.fbs +++ b/src/halide_ir.fbs @@ -117,7 +117,7 @@ enum MemoryType: byte { VTCM, AMXTile, GPUSharedAsync, - WMMAAccumulator, + WMMAFragment, } table Range { diff --git a/test/correctness/wmma_matmul.cpp b/test/correctness/wmma_matmul.cpp index d2e9ba9a4538..437bca94f434 100644 --- a/test/correctness/wmma_matmul.cpp +++ b/test/correctness/wmma_matmul.cpp @@ -105,7 +105,7 @@ bool test(const Params &p) { } prod.compute_at(out, xt) - .store_in(MemoryType::WMMAAccumulator) + .store_in(MemoryType::WMMAFragment) .split(x, x, rxi, p.tile_n) .split(y, y, ryi, p.tile_m) .vectorize(rxi) @@ -229,7 +229,7 @@ bool test_block_level_accumulator() { .vectorize(mmyi); prod.compute_at(out, x) - .store_in(MemoryType::WMMAAccumulator) + .store_in(MemoryType::WMMAFragment) .split(x, xw, xi, tile * tiles_x) .split(xi, xi, rxi, tile) .split(y, y, ryi, tile) @@ -296,6 +296,204 @@ bool test_block_level_accumulator() { return true; } +// Stage the operand tiles into fragment registers, so that each fragment +// loaded feeds several multiplies. Where each staging happens says how much +// reuse we get out of it. +bool test_staged_operands() { + const int M = 128, N = 128, K = 64; + const int tile = 16, tiles_x = 2, tiles_y = 2, bk = 32; + + Buffer A(K, M), B(N, K); + fill(A); + fill(B); + + Var x("x"), y("y"), kk("kk"), yy("yy"), xx("xx"); + RDom k(0, K, "k"); + Func prod("prod"), out("out"), Am("Am"), Bm("Bm"); + + Am(kk, yy) = A(kk, yy); + Bm(xx, kk) = B(xx, kk); + prod(x, y) = 0.f; + prod(x, y) += cast(Am(k, y)) * cast(Bm(x, k)); + out(x, y) = prod(x, y); + + Var xi("xi"), yi("yi"), mmxi("mmxi"), mmyi("mmyi"), rxi("rxi"), ryi("ryi"); + Var kko("kko"), kki("kki"), yyo("yyo"), yyi("yyi"), xxo("xxo"), xxi("xxi"); + RVar ko("ko"), ki("ki"), rri("rri"); + + out.bound(x, 0, N) + .bound(y, 0, M) + .split(x, x, xi, tile * tiles_x) + .split(xi, xi, mmxi, tile) + .split(y, y, yi, tile * tiles_y) + .split(yi, yi, mmyi, tile) + .gpu_blocks(x, y) + .reorder(mmxi, mmyi, xi, yi, x, y) + .unroll(xi) + .unroll(yi) + .vectorize(mmxi) + .vectorize(mmyi); + + prod.compute_at(out, x) + .store_in(MemoryType::WMMAFragment) + .split(x, x, rxi, tile) + .split(y, y, ryi, tile) + .vectorize(rxi) + .vectorize(ryi) + .unroll(x) + .unroll(y); + + // Loop nest of the update, outermost first: ko, ki, y, x. + prod.update() + .split(k, ko, ki, bk) + .split(x, x, rxi, tile) + .split(y, y, ryi, tile) + .split(ki, ki, rri, tile) + .reorder(rri, rxi, ryi, x, y, ki, ko) + .unroll(x) + .unroll(y) + .unroll(ki) + .atomic() + .vectorize(rri) + .vectorize(rxi) + .vectorize(ryi); + + // One a fragment per row of tiles, live across the loop over columns. + Am.compute_at(prod, y) + .store_in(MemoryType::WMMAFragment) + .split(kk, kko, kki, tile) + .split(yy, yyo, yyi, tile) + .reorder(kki, yyi, kko, yyo) + .unroll(kko) + .unroll(yyo) + .vectorize(kki) + .vectorize(yyi); + + // All the b fragments at once, live across the loops over both. + Bm.compute_at(prod, ki) + .store_in(MemoryType::WMMAFragment) + .split(xx, xxo, xxi, tile) + .split(kk, kko, kki, tile) + .reorder(xxi, kki, xxo, kko) + .unroll(xxo) + .unroll(kko) + .vectorize(xxi) + .vectorize(kki); + + Buffer result(N, M); + out.realize(result); + result.copy_to_host(); + + for (int j = 0; j < M; j++) { + for (int i = 0; i < N; i++) { + float ref = 0.f; + for (int l = 0; l < K; l++) { + ref += (float)A(l, j) * (float)B(i, l); + } + if (std::abs(result(i, j) - ref) > 1e-2f * std::max(1.f, std::abs(ref))) { + std::cerr << "Mismatch at " << i << ", " << j << ": " + << result(i, j) << " != " << ref << "\n" + << "For operands staged in fragment registers\n"; + return false; + } + } + } + return true; +} + +// A batch of matrix multiplies that all share the same left-hand side. Its +// fragments only need loading once for the whole batch, which is what staging +// them outside the loop over the batch does. The loop isn't unrolled, so +// nothing downstream of here could hoist them out of it. +bool test_operand_hoisted_out_of_loop() { + const int M = 32, N = 32, K = 16, batch = 4; + const int tile = 16; + + Buffer A(K, M), B(N, K, batch); + fill(A); + fill(B); + + Var x("x"), y("y"), n("n"), kk("kk"), yy("yy"); + RDom k(0, K, "k"); + Func prod("prod"), out("out"), Am("Am"); + + Am(kk, yy) = A(kk, yy); + prod(x, y, n) = 0.f; + prod(x, y, n) += cast(Am(k, y)) * cast(B(x, k, n)); + out(x, y, n) = prod(x, y, n); + + Var xi("xi"), yi("yi"), mmxi("mmxi"), mmyi("mmyi"), rxi("rxi"), ryi("ryi"); + Var kko("kko"), kki("kki"), yyo("yyo"), yyi("yyi"); + RVar rro("rro"), rri("rri"); + + out.bound(x, 0, N) + .bound(y, 0, M) + .bound(n, 0, batch) + .split(x, x, xi, N) + .split(xi, xi, mmxi, tile) + .split(y, y, yi, M) + .split(yi, yi, mmyi, tile) + .gpu_blocks(x, y) + .reorder(mmxi, mmyi, xi, yi, n, x, y) + .unroll(xi) + .unroll(yi) + .vectorize(mmxi) + .vectorize(mmyi); + + prod.compute_at(out, n) + .store_in(MemoryType::WMMAFragment) + .split(x, x, rxi, tile) + .split(y, y, ryi, tile) + .vectorize(rxi) + .vectorize(ryi) + .unroll(x) + .unroll(y); + + prod.update() + .split(x, x, rxi, tile) + .split(y, y, ryi, tile) + .split(k, rro, rri, tile) + .reorder(rri, rxi, ryi, x, y, rro) + .unroll(x) + .unroll(y) + .atomic() + .vectorize(rri) + .vectorize(rxi) + .vectorize(ryi); + + Am.compute_at(out, x) + .store_in(MemoryType::WMMAFragment) + .split(kk, kko, kki, tile) + .split(yy, yyo, yyi, tile) + .reorder(kki, yyi, kko, yyo) + .unroll(kko) + .unroll(yyo) + .vectorize(kki) + .vectorize(yyi); + + Buffer result(N, M, batch); + out.realize(result); + result.copy_to_host(); + + for (int b = 0; b < batch; b++) { + for (int j = 0; j < M; j++) { + for (int i = 0; i < N; i++) { + float ref = 0.f; + for (int l = 0; l < K; l++) { + ref += (float)A(l, j) * (float)B(i, l, b); + } + if (std::abs(result(i, j, b) - ref) > 1e-2f * std::max(1.f, std::abs(ref))) { + std::cerr << "Mismatch at " << i << ", " << j << ", " << b << ": " + << result(i, j, b) << " != " << ref << "\n" + << "For an operand staged outside a loop over a batch\n"; + return false; + } + } + } + } + return true; +} + } // namespace int main(int argc, char **argv) { @@ -366,7 +564,9 @@ int main(int argc, char **argv) { } } - if (!test_block_level_accumulator()) { + if (!test_block_level_accumulator() || + !test_staged_operands() || + !test_operand_hoisted_out_of_loop()) { printf("Failed!\n"); return 1; } From f72d389daf65cbc0fac5551d2fff1cfe63406ef1 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Wed, 29 Jul 2026 09:30:31 -0700 Subject: [PATCH 21/59] Stop tensor core fragments becoming vectors in the PTX backend Two things were costing the tensor core matmul instructions it didn't need to spend. Each producer waited for its own asynchronous copies, so a block staging two operand panels paid the latency of the first before it had even issued the second. Wait instead at the barrier or the load that needs the data, so every copy a block issues can be in flight at once. Fragments also round-tripped through their allocations as Halide vectors. NVPTX holds a wide float vector in pairs of registers, so every tensor core instruction unpacked its accumulator into eight registers and packed the result back up again - about 117 instructions per reduction step per warp, which was most of the instruction stream. Read and write fragments a register at a time instead, so they never become vectors. At 2048^3 on an RTX 5060 Ti this takes the best tensor core matmul from 40.0 to 46.0 TFlop/s. cuBLAS is 49.6. Co-Authored-By: Claude Opus 5 --- src/CodeGen_PTX_Dev.cpp | 144 +++++++++++++++++++++++++++++++--------- 1 file changed, 113 insertions(+), 31 deletions(-) diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 82c67ed4d53d..5a824bee753b 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -1,3 +1,5 @@ +#include + #include "CodeGen_PTX_Dev.h" #include "CSE.h" #include "CanonicalizeGPUVars.h" @@ -109,11 +111,12 @@ class CodeGen_PTX_Dev : public CodeGen_LLVM, public CodeGen_GPU_Dev { * copies into shared memory can be recognized. */ Scope alloc_memory_type; - /** Whether we're inside a producer node, which is where the wait for any - * asynchronous copies gets emitted, and whether any have been issued in - * it. */ + /** Whether we're inside a producer node, and the destinations of any + * asynchronous copies issued but not yet waited for. Waiting is deferred to + * the point where the data is actually needed, so that copies to several + * destinations can all be in flight at once. */ bool in_producer = false; - bool issued_async_copy = false; + std::set pending_async_copies; /** Try to emit a store into shared memory as an asynchronous copy, which * moves the data straight from global memory without routing it through @@ -127,8 +130,22 @@ class CodeGen_PTX_Dev : public CodeGen_LLVM, public CodeGen_GPU_Dev { * intrinsics that drive the tensor cores. */ // @{ void codegen_wmma(const Call *op); + llvm::Value *codegen_wmma_raw(const Call *op); void codegen_wmma_store(const Store *op); void split_fragment(const Expr &e, std::vector &args); + // @} + + /** A fragment is a fixed number of 32-bit registers per lane. Keeping it in + * that form all the way to and from its allocation matters, because NVPTX + * holds a wide float vector in pairs of registers, and packing and + * unpacking one around every tensor core instruction costs more + * instructions than the instructions themselves. + */ + // @{ + bool is_fragment_alloc(const std::string &name); + llvm::Type *fragment_reg_type(Type t); + llvm::Value *fragment_reg_ptr(const std::string &name, Type t, int i); + void codegen_fragment_store(const Store *op); llvm::Value *call_wmma_intrinsic(const std::string &name, const std::vector &args, const std::vector &overloads); @@ -361,6 +378,17 @@ WMMAMatrixLayout matrix_in_memory(const string &name, const MultiRamp &mr, int r void CodeGen_PTX_Dev::split_fragment(const Expr &e, vector &args) { // One llvm value per 32-bit register. + const int num_regs = e.type().bits() * e.type().lanes() / 32; + if (const Load *load = e.as()) { + if (is_fragment_alloc(load->name)) { + llvm::Type *reg_type = fragment_reg_type(e.type()); + for (int i = 0; i < num_regs; i++) { + args.push_back(builder->CreateAlignedLoad( + reg_type, fragment_reg_ptr(load->name, e.type(), i), llvm::Align(4))); + } + return; + } + } Value *v = codegen(e); const int lanes_per_reg = 32 / e.type().bits(); for (int i = 0; i < e.type().lanes() / lanes_per_reg; i++) { @@ -381,6 +409,74 @@ Value *CodeGen_PTX_Dev::call_wmma_intrinsic(const std::string &name, } void CodeGen_PTX_Dev::codegen_wmma(const Call *op) { + Value *result = codegen_wmma_raw(op); + + // Reassemble the returned struct into a Halide vector. + llvm::Type *result_type = llvm_type_of(op->type); + const int num_regs = op->type.bits() * op->type.lanes() / 32; + if (op->type.bits() == 32) { + value = UndefValue::get(result_type); + for (int i = 0; i < num_regs; i++) { + value = builder->CreateInsertElement(value, builder->CreateExtractValue(result, i), i); + } + } else { + vector regs; + regs.reserve(num_regs); + for (int i = 0; i < num_regs; i++) { + regs.push_back(builder->CreateExtractValue(result, i)); + } + value = concat_vectors(regs); + } + internal_assert(value->getType() == result_type) + << "Unexpected result type from a tensor core instruction\n"; +} + +bool CodeGen_PTX_Dev::is_fragment_alloc(const std::string &name) { + const MemoryType *t = alloc_memory_type.find(name); + return t && *t == MemoryType::WMMAFragment; +} + +llvm::Type *CodeGen_PTX_Dev::fragment_reg_type(Type t) { + return t.bits() == 32 ? llvm_type_of(t.element_of()) : + get_vector_type(llvm_type_of(t.element_of()), 32 / t.bits()); +} + +llvm::Value *CodeGen_PTX_Dev::fragment_reg_ptr(const std::string &name, Type t, int i) { + return codegen_buffer_pointer(name, t.element_of(), Expr(i * (32 / t.bits()))); +} + +void CodeGen_PTX_Dev::codegen_fragment_store(const Store *op) { + const Type t = op->value.type(); + const int num_regs = t.bits() * t.lanes() / 32; + llvm::Type *reg_type = fragment_reg_type(t); + + const Call *call = op->value.as(); + if (call && is_wmma_intrinsic(call)) { + // Take the registers straight out of the struct the instruction + // returns, without ever making a vector of them. + Value *result = codegen_wmma_raw(call); + for (int i = 0; i < num_regs; i++) { + builder->CreateAlignedStore(builder->CreateExtractValue(result, i), + fragment_reg_ptr(op->name, t, i), llvm::Align(4)); + } + return; + } + + // Anything else (a zero-initialization, say) does become a vector, but it + // is still written a register at a time so that the allocation only ever + // sees register-sized accesses. + Value *v = codegen(op->value); + const int lanes_per_reg = 32 / t.bits(); + for (int i = 0; i < num_regs; i++) { + Value *reg = lanes_per_reg == 1 ? + builder->CreateExtractElement(v, i) : + slice_vector(v, i * lanes_per_reg, lanes_per_reg); + internal_assert(reg->getType() == reg_type); + builder->CreateAlignedStore(reg, fragment_reg_ptr(op->name, t, i), llvm::Align(4)); + } +} + +llvm::Value *CodeGen_PTX_Dev::codegen_wmma_raw(const Call *op) { // The nvvm wmma intrinsics take and return fragments as a flat list of // 32-bit registers, packaged up as a literal struct. We represent them in // Halide IR as vectors, so most of the work here is repacking. @@ -437,26 +533,7 @@ void CodeGen_PTX_Dev::codegen_wmma(const Call *op) { args.push_back(codegen(cast(Int(32), mem.stride))); } - Value *result = call_wmma_intrinsic(name.str(), args, overloads); - - // Reassemble the returned struct into a Halide vector. - llvm::Type *result_type = llvm_type_of(op->type); - const int num_regs = op->type.bits() * op->type.lanes() / 32; - if (op->type.bits() == 32) { - value = UndefValue::get(result_type); - for (int i = 0; i < num_regs; i++) { - value = builder->CreateInsertElement(value, builder->CreateExtractValue(result, i), i); - } - } else { - vector regs; - regs.reserve(num_regs); - for (int i = 0; i < num_regs; i++) { - regs.push_back(builder->CreateExtractValue(result, i)); - } - value = concat_vectors(regs); - } - internal_assert(value->getType() == result_type) - << "Unexpected result type from " << name.str() << "\n"; + return call_wmma_intrinsic(name.str(), args, overloads); } void CodeGen_PTX_Dev::codegen_wmma_store(const Store *op) { @@ -579,6 +656,10 @@ void CodeGen_PTX_Dev::visit(const AssertStmt *op) { } void CodeGen_PTX_Dev::visit(const Load *op) { + if (pending_async_copies.count(op->name)) { + // This thread is about to read something it copied asynchronously. + wait_for_async_copies(); + } // Do aligned 4-wide 32-bit loads as a single i128 load. const Ramp *r = op->index.as(); @@ -725,7 +806,7 @@ bool CodeGen_PTX_Dev::codegen_async_copy(const Store *op, const char **reason) { << "Could not find the nvvm intrinsic " << name.str() << "\n"; llvm::Function *fn = llvm::Intrinsic::getOrInsertDeclaration(module.get(), id); builder->CreateCall(fn, {dst, src_ptr}); - issued_async_copy = true; + pending_async_copies.insert(op->name); return true; } @@ -735,16 +816,12 @@ void CodeGen_PTX_Dev::visit(const ProducerConsumer *op) { return; } - ScopedValue old_issued(issued_async_copy, false); ScopedValue old_in(in_producer, true); codegen(op->body); - // Everything issued in here has to have landed before the values are used, - // which is after this producer. - wait_for_async_copies(); } void CodeGen_PTX_Dev::wait_for_async_copies() { - if (!issued_async_copy) { + if (pending_async_copies.empty()) { return; } for (const char *intrin : {"llvm.nvvm.cp.async.commit.group", @@ -758,7 +835,7 @@ void CodeGen_PTX_Dev::wait_for_async_copies() { } builder->CreateCall(fn, args); } - issued_async_copy = false; + pending_async_copies.clear(); } void CodeGen_PTX_Dev::visit(const Store *op) { @@ -773,6 +850,11 @@ void CodeGen_PTX_Dev::visit(const Store *op) { user_assert(op->value.type().bits() >= 32) << "CUDA: 8-bit or 16-bit atomics are not supported.\n"; } + if (is_fragment_alloc(op->name)) { + codegen_fragment_store(op); + return; + } + { const char *reason = ""; if (codegen_async_copy(op, &reason)) { From bb7456bd1ac70441eb17ac93310956a5b45ecee0 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Wed, 29 Jul 2026 11:30:56 -0700 Subject: [PATCH 22/59] Tell ptxas how large a CUDA thread block is Without the annotation it has to assume the largest block the hardware supports and allocate registers for that. Measured no difference on the tensor core matmul, whose register count is pinned by its live accumulators, but it is information the backend should be passing on. Co-Authored-By: Claude Opus 5 --- src/CodeGen_PTX_Dev.cpp | 49 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 5a824bee753b..001c59d258fd 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -176,6 +176,35 @@ Type CodeGen_PTX_Dev::upgrade_type_for_storage(const Type &t) const { return CodeGen_LLVM::upgrade_type_for_storage(t); } + +namespace { + +// The size of the thread block a kernel will be launched with. ptxas allocates +// registers on the assumption that a block may be as large as the hardware +// allows unless we tell it otherwise, which costs occupancy. +class BlockSize : public IRVisitor { + using IRVisitor::visit; + + void visit(const For *op) override { + for (int i = 0; i < 3; i++) { + if (ends_with(op->name, gpu_thread_name(i))) { + if (auto e = as_const_int(simplify(op->extent()))) { + extent[i] = std::max(extent[i], (int)*e); + } else { + known = false; + } + } + } + IRVisitor::visit(op); + } + +public: + int extent[3] = {1, 1, 1}; + bool known = true; +}; + +} // namespace + void CodeGen_PTX_Dev::add_kernel(Stmt stmt, const std::string &name, const std::vector &args) { @@ -254,6 +283,26 @@ void CodeGen_PTX_Dev::add_kernel(Stmt stmt, module->getOrInsertNamedMetadata("nvvm.annotations")->addOperand(md_node); + // Tell ptxas how large the thread block is. Without this it has to assume + // the largest block the hardware supports, and allocates registers for that + // rather than for the block we actually launch. + BlockSize block_size; + stmt.accept(&block_size); + if (block_size.known) { + const char *annotation[] = {"maxntidx", "maxntidy", "maxntidz"}; + for (int i = 0; i < 3; i++) { + llvm::Metadata *args[] = { + llvm::ValueAsMetadata::get(function), + MDString::get(*context, annotation[i]), + llvm::ValueAsMetadata::get(ConstantInt::get(i32_t, block_size.extent[i]))}; + module->getOrInsertNamedMetadata("nvvm.annotations") + ->addOperand(MDNode::get(*context, args)); + } + debug(2) << "Kernel " << name << " has block size " + << block_size.extent[0] << "x" << block_size.extent[1] + << "x" << block_size.extent[2] << "\n"; + } + // Now verify the function is ok verifyFunction(*function); From adae5de8cb6ca3fd0891ff9d62f995434e68c1fd Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Wed, 29 Jul 2026 09:32:17 -0700 Subject: [PATCH 23/59] Retune apps/tensorcore_matmul for the cheaper fragment handling With fragments no longer packed into and out of vectors, a wider tile per warp pays off. At 1024^3 / 2048^3 / 4096^3 on an RTX 5060 Ti the tensor core schedule goes from 27.1 / 39.3 / 40.3 to 31.3 / 45.2 / 46.2 TFlop/s, against cuBLAS at 43.3 / 49.6 / 50.0. Co-Authored-By: Claude Opus 5 --- apps/tensorcore_matmul/matmul_generator.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/apps/tensorcore_matmul/matmul_generator.cpp b/apps/tensorcore_matmul/matmul_generator.cpp index 54fe01ffc605..89752fb773c4 100644 --- a/apps/tensorcore_matmul/matmul_generator.cpp +++ b/apps/tensorcore_matmul/matmul_generator.cpp @@ -21,8 +21,8 @@ class MatMul : public Halide::Generator { // How many tensor core tiles of accumulator each warp holds, and how many // warps there are per block in each dimension. - GeneratorParam tiles_x{"tiles_x", 4}; - GeneratorParam tiles_y{"tiles_y", 5}; + GeneratorParam tiles_x{"tiles_x", 5}; + GeneratorParam tiles_y{"tiles_y", 4}; GeneratorParam warps_x{"warps_x", 2}; GeneratorParam warps_y{"warps_y", 1}; // How much of the reduction is staged in shared memory at a time. @@ -31,7 +31,7 @@ class MatMul : public Halide::Generator { // rows across different banks. A multiple of eight keeps the rows aligned // enough for the widest asynchronous copy. GeneratorParam pad_a{"pad_a", 8}; - GeneratorParam pad_b{"pad_b", 8}; + GeneratorParam pad_b{"pad_b", 16}; Input> matA{"matA"}; // K x M Input> matB{"matB"}; // N x K From 8f6e08dd24b12a2cd0d5e7e6e8381e4e5c71769e Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Wed, 29 Jul 2026 09:38:07 -0700 Subject: [PATCH 24/59] Pad the shared B panel to a stride that actually avoids bank conflicts Of the strides that keep the rows aligned enough for an asynchronous copy, the one in use was among the worse ones: it left 2.2M shared load bank conflicts against 23K for the next multiple of eight up. Co-Authored-By: Claude Opus 5 --- apps/tensorcore_matmul/matmul_generator.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/tensorcore_matmul/matmul_generator.cpp b/apps/tensorcore_matmul/matmul_generator.cpp index 89752fb773c4..b0a0bef948ca 100644 --- a/apps/tensorcore_matmul/matmul_generator.cpp +++ b/apps/tensorcore_matmul/matmul_generator.cpp @@ -31,7 +31,7 @@ class MatMul : public Halide::Generator { // rows across different banks. A multiple of eight keeps the rows aligned // enough for the widest asynchronous copy. GeneratorParam pad_a{"pad_a", 8}; - GeneratorParam pad_b{"pad_b", 16}; + GeneratorParam pad_b{"pad_b", 24}; Input> matA{"matA"}; // K x M Input> matB{"matB"}; // N x K From 6d6c505b718f464cee478b02b6cc054a52185b9d Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Thu, 30 Jul 2026 10:58:07 -0700 Subject: [PATCH 25/59] Pick the block shape and its padding by problem size A 160x64 block leaves only a few dozen blocks to cover 36 SMs at 1024^3, and 160 does not divide 1024, so the last block in each row is ragged. A 128x32 block fixes both, and takes that size from 31.1 to 37.7 TFlop/s - 72% of cuBLAS to 88%. The padding has to be chosen with the shape rather than fixed, because it is what keeps consecutive rows of the staged panel in different banks. Carrying the 160-wide block's padding over to a 128-wide one costs 30%: 26.1 against 37.5 TFlop/s. The explicit tile parameters still override, and a search over shapes found nothing that beat the existing choice at 2048 or 4096. Co-Authored-By: Claude Opus 5 --- apps/tensorcore_matmul/matmul_generator.cpp | 57 ++++++++++++++------- 1 file changed, 39 insertions(+), 18 deletions(-) diff --git a/apps/tensorcore_matmul/matmul_generator.cpp b/apps/tensorcore_matmul/matmul_generator.cpp index b0a0bef948ca..1ac0efa0397b 100644 --- a/apps/tensorcore_matmul/matmul_generator.cpp +++ b/apps/tensorcore_matmul/matmul_generator.cpp @@ -1,5 +1,7 @@ #include "Halide.h" +#include + namespace { using namespace Halide; @@ -21,10 +23,13 @@ class MatMul : public Halide::Generator { // How many tensor core tiles of accumulator each warp holds, and how many // warps there are per block in each dimension. - GeneratorParam tiles_x{"tiles_x", 5}; - GeneratorParam tiles_y{"tiles_y", 4}; - GeneratorParam warps_x{"warps_x", 2}; - GeneratorParam warps_y{"warps_y", 1}; + // Zero means pick a shape based on the problem size. The best block gets + // smaller as the matrices do, because a large one leaves too few blocks to + // fill the machine. + GeneratorParam tiles_x{"tiles_x", 0}; + GeneratorParam tiles_y{"tiles_y", 0}; + GeneratorParam warps_x{"warps_x", 0}; + GeneratorParam warps_y{"warps_y", 0}; // How much of the reduction is staged in shared memory at a time. GeneratorParam block_k{"block_k", 32}; // Extra elements per row of the shared panels, which spreads consecutive @@ -91,8 +96,24 @@ class MatMul : public Halide::Generator { // tiles_y) multiplies, so this is what gets us reuse out of the // loads. const int tile = 16; - const int block_x = tile * tiles_x * warps_x; - const int block_y = tile * tiles_y * warps_y; + int tx = tiles_x, ty = tiles_y, wx = warps_x, wy = warps_y; + int pb = pad_b; + if (tx == 0 || ty == 0 || wx == 0 || wy == 0) { + // The padding goes with the shape: it is what keeps consecutive + // rows of the staged panel in different banks, so the right + // amount depends on how wide the panel is. + const int n = std::min({(int)M, (int)N, (int)K}); + if (n <= 1024) { + // Small problems need small blocks: a 160x64 block leaves + // only a few dozen of them to cover 36 SMs, and 160 does + // not divide 1024 so the last one in each row is ragged. + tx = 8, ty = 2, wx = 1, wy = 1, pb = 8; + } else { + tx = 5, ty = 4, wx = 2, wy = 1, pb = 24; + } + } + const int block_x = tile * tx * wx; + const int block_y = tile * ty * wy; Var xi("xi"), yi("yi"), xt("xt"), yt("yt"), mmxi("mmxi"), mmyi("mmyi"); Var xw("xw"), yw("yw"), rxi("rxi"), ryi("ryi"); @@ -101,10 +122,10 @@ class MatMul : public Halide::Generator { output.bound(x, 0, N) .bound(y, 0, M) .split(x, x, xi, block_x) - .split(xi, xt, xi, tile * tiles_x) + .split(xi, xt, xi, tile * tx) .split(xi, xi, mmxi, tile) .split(y, y, yi, block_y) - .split(yi, yt, yi, tile * tiles_y) + .split(yi, yt, yi, tile * ty) .split(yi, yi, mmyi, tile) .gpu_blocks(x, y) .gpu_threads(xt, yt) @@ -120,9 +141,9 @@ class MatMul : public Halide::Generator { // loop over warps, which lets every warp share one staged panel. prod.compute_at(output, x) .store_in(MemoryType::WMMAFragment) - .split(x, xw, xi, tile * tiles_x) + .split(x, xw, xi, tile * tx) .split(xi, xi, rxi, tile) - .split(y, yw, yi, tile * tiles_y) + .split(y, yw, yi, tile * ty) .split(yi, yi, ryi, tile) .reorder(rxi, ryi, xi, yi, xw, yw) .gpu_threads(xw, yw) @@ -133,9 +154,9 @@ class MatMul : public Halide::Generator { prod.update() .split(k, ko, ki, block_k) - .split(x, xw, xi, tile * tiles_x) + .split(x, xw, xi, tile * tx) .split(xi, xi, rxi, tile) - .split(y, yw, yi, tile * tiles_y) + .split(y, yw, yi, tile * ty) .split(yi, yi, ryi, tile) .split(ki, ki, rri, tile) .reorder(rri, rxi, ryi, xi, yi, ki, xw, yw, ko) @@ -163,20 +184,20 @@ class MatMul : public Halide::Generator { .split(kk, kko, kki, vec) .fuse(kko, y, t) .split(t, t, ti, 32) - .split(t, t, tw, warps_x) - .split(t, to, tw2, warps_y) + .split(t, t, tw, wx) + .split(t, to, tw2, wy) .gpu_lanes(ti) .gpu_threads(tw, tw2) .vectorize(kki); Bs.compute_at(prod, ko) - .store_in(MemoryType::GPUShared) - .align_storage(x, block_x + (int)pad_b) + .store_in(MemoryType::GPUSharedAsync) + .align_storage(x, block_x + pb) .split(x, xxo, xxi, vec) .fuse(xxo, kk, t) .split(t, t, ti, 32) - .split(t, t, tw, warps_x) - .split(t, to, tw2, warps_y) + .split(t, t, tw, wx) + .split(t, to, tw2, wy) .gpu_lanes(ti) .gpu_threads(tw, tw2) .vectorize(xxi); From dd504b2725839b4f2afa71d4f28ff55a2e1934bd Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Fri, 31 Jul 2026 09:28:12 -0700 Subject: [PATCH 26/59] Check the AMX error tests reach the error they are named for expect_user_error accepted any CompileError, so a scenario could pass on an unrelated one - a schedule Halide rejects before lowering looks the same as the lowering error the scenario is for. Give it the message to look for, and name the path each scenario exercises. That leaves one of the checks in get_subtile untested, so add a case for it: two matmuls into the same allocation whose tiles have the same rank but different extents. Also drop the returns after the user_errors there, which throw. Co-Authored-By: Claude Opus 5 --- src/ExtractTileOperations.cpp | 1 - test/correctness/tiled_matmul_errors.cpp | 70 ++++++++++++++++++++---- 2 files changed, 58 insertions(+), 13 deletions(-) diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp index 47d19f13de07..26ab57d5621a 100644 --- a/src/ExtractTileOperations.cpp +++ b/src/ExtractTileOperations.cpp @@ -400,7 +400,6 @@ class ExtractTileOperations : public IRMutator { // Returns an index expression for a given load or store index. user_asserts if impossible std::string get_subtile_name(const Expr &index) { int idx = get_subtile(index, "AMX tile", &amx_subtiles); - internal_assert(idx >= 0); // errors handled already return amx_name + std::to_string(idx); } diff --git a/test/correctness/tiled_matmul_errors.cpp b/test/correctness/tiled_matmul_errors.cpp index 4cf486df291f..751ca4f607e7 100644 --- a/test/correctness/tiled_matmul_errors.cpp +++ b/test/correctness/tiled_matmul_errors.cpp @@ -19,10 +19,16 @@ const Target amx_target("x86-64-linux-avx512_sapphirerapids"); // Run `body` and assert it produces a Halide user error. template -bool expect_user_error(const char *name, F body) { +bool expect_user_error(const char *name, const char *substring, F body) { try { body(); } catch (const CompileError &e) { + std::string msg = e.what(); + if (msg.find(substring) == std::string::npos) { + printf("[%s] FAIL: error did not mention \"%s\":\n%s\n", + name, substring, msg.c_str()); + return false; + } printf("[%s] OK: %s\n", name, e.what()); return true; } catch (...) { @@ -205,6 +211,45 @@ void scenario_widening_16bit() { mm.in().compile_jit(amx_target); } +// Two matmuls into the same allocation with tiles of the same rank but +// different extents, so they disagree about the shape of a tile register. +void scenario_mismatched_strides() { + Buffer A(64, 64), C(64, 64); + Buffer B(4, 64, 16), D(4, 64, 16); + Var x("x"), y("y"); + RDom r1(0, 64, "r1"), r2(0, 64, "r2"); + + Func mm("matmul_mismatched"); + mm(x, y) = cast(0); + mm(x, y) += cast(A(r1, y)) * cast(B(r1 % 4, x, r1 / 4)); + mm(x, y) += cast(C(r2, y)) * cast(D(r2 % 4, x, r2 / 4)); + + Var rxi("rxi"), ryi("ryi"); + RVar rri("rri"), rro("rro"); + mm.compute_at(mm.in(), x).store_in(MemoryType::AMXTile); + mm.update(0) + .tile(x, y, rxi, ryi, 8, 4, TailStrategy::GuardWithIf) + .split(r1.x, rro, rri, 8) + .reorder(rri, rxi, ryi, rro, x, y) + .atomic() + .vectorize(rri) + .vectorize(rxi) + .vectorize(ryi); + mm.update(1) + .tile(x, y, rxi, ryi, 4, 8, TailStrategy::GuardWithIf) + .split(r2.x, rro, rri, 8) + .reorder(rri, rxi, ryi, rro, x, y) + .atomic() + .vectorize(rri) + .vectorize(rxi) + .vectorize(ryi); + Var ixi("ixi"), iyi("iyi"); + mm.compute_at(mm.in(), x).tile(x, y, ixi, iyi, 8, 8).vectorize(ixi).vectorize(iyi); + Var mmxi("mmxi"), mmyi("mmyi"); + mm.in().tile(x, y, mmxi, mmyi, 8, 8).vectorize(mmxi).vectorize(mmyi); + mm.in().compile_jit(amx_target); +} + // A user gives the same Func two update definitions that each store into // the same AMXTile allocation but with different tile sizes (e.g. a fast // path for the bulk of K and a smaller fallback). The matcher requires @@ -323,17 +368,18 @@ int main(int argc, char **argv) { int failures = 0; - failures += !expect_user_error("too_large", scenario_too_large); - failures += !expect_user_error("bad_result_type", scenario_bad_result_type); - failures += !expect_user_error("naive_rhs", scenario_naive_rhs); - failures += !expect_user_error("indirect", scenario_indirect); - failures += !expect_user_error("sign_changing_cast", scenario_sign_changing_cast); - failures += !expect_user_error("conv1d", scenario_conv1d); - failures += !expect_user_error("no_matmul", scenario_no_matmul); - failures += !expect_user_error("widening_16bit", scenario_widening_16bit); - failures += !expect_user_error("inconsistent_tiles", scenario_inconsistent_tiles); - failures += !expect_user_error("not_a_matmul_pattern", scenario_not_a_matmul_pattern); - failures += !expect_user_error("matmul_by_constant", scenario_matmul_by_constant); + failures += !expect_user_error("too_large", "too large to fit in", scenario_too_large); + failures += !expect_user_error("bad_result_type", "must yield 32-bit integers", scenario_bad_result_type); + failures += !expect_user_error("naive_rhs", "storage layout for a matrix multiply operand is unsupported", scenario_naive_rhs); + failures += !expect_user_error("indirect", "not loads with affine indices", scenario_indirect); + failures += !expect_user_error("sign_changing_cast", "cast after being loaded", scenario_sign_changing_cast); + failures += !expect_user_error("conv1d", "storage layout for a matrix multiply operand is unsupported", scenario_conv1d); + failures += !expect_user_error("no_matmul", "no matrix multiply operation was found", scenario_no_matmul); + failures += !expect_user_error("widening_16bit", "operand or result types are not supported", scenario_widening_16bit); + failures += !expect_user_error("mismatched_strides", "has different size and strides", scenario_mismatched_strides); + failures += !expect_user_error("inconsistent_tiles", "does not have the same shape", scenario_inconsistent_tiles); + failures += !expect_user_error("not_a_matmul_pattern", "operand or result types are not supported", scenario_not_a_matmul_pattern); + failures += !expect_user_error("matmul_by_constant", "not loads with affine indices", scenario_matmul_by_constant); if (failures != 0) { printf("%d scenario(s) failed to produce a user-facing CompileError\n", failures); From 29a287e01bae1ef393692cf0c109e4e97e620fc7 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Fri, 31 Jul 2026 09:28:31 -0700 Subject: [PATCH 27/59] Recognize any permutation of the lanes as a reshaping shuffle is_multiramp treated a shuffle of a single vector as a reshaping only when the vector was one-dimensional, which is the shape flatten_nested_ramps leaves a strided load in. A tile load that gets transposed is not that shape, so add MultiRamp::shuffle, which handles any mask whose lanes are a permutation of [0, total_lanes()) - that is, whose strides are the prefix products of its lane counts in some order. It works by refining the mask's dims and the multiramp's against each other until both are groupings of a common shape, then reordering that shape the way the mask asks for. MultiRamp::add was already doing the same walk to line two multiramps up, and was also duplicating strides_for_shape to rescale the strides as it went, so share the walk and let strides_for_shape do the rest. That makes add a good deal shorter, and makes the failure paths that a common refinement rules out into internal errors rather than silent rejections. Co-Authored-By: Claude Opus 5 --- src/MultiRamp.cpp | 189 +++++++++++++++++++++------------ src/MultiRamp.h | 13 +++ test/correctness/multiramp.cpp | 83 +++++++++++++++ 3 files changed, 219 insertions(+), 66 deletions(-) diff --git a/src/MultiRamp.cpp b/src/MultiRamp.cpp index d9281d87ee40..d7277511e81a 100644 --- a/src/MultiRamp.cpp +++ b/src/MultiRamp.cpp @@ -34,6 +34,40 @@ void collapse_adjacent_dims(MultiRamp *m) { } } +// Walk two shapes innermost-out, emitting the coarsest shape that groups into +// both of them. Returns false if there is no such shape, which happens when a +// pair of lane counts is coprime. If from_a is non-null it receives, for each +// refined dim, the index of the dim of `a` it came from. +bool common_refinement(const std::vector &a, const std::vector &b, + std::vector *refined, std::vector *from_a) { + refined->clear(); + if (from_a) { + from_a->clear(); + } + size_t i = 0, j = 0; + int x = a.empty() ? 1 : a[0]; + int y = b.empty() ? 1 : b[0]; + while (i < a.size() && j < b.size()) { + int g = gcd(x, y); + if (g == 1) { + return false; + } + refined->push_back(g); + if (from_a) { + from_a->push_back((int)i); + } + x /= g; + y /= g; + if (x == 1 && ++i < a.size()) { + x = a[i]; + } + if (y == 1 && ++j < b.size()) { + y = b[j]; + } + } + return true; +} + } // namespace MultiRamp::MultiRamp(Expr base, std::vector strides, std::vector lanes) @@ -62,68 +96,32 @@ void MultiRamp::mul(const Expr &e) { // common refinement (the sum is not a multiramp). Adding multiramps with // different total lane counts is a caller error and triggers an assertion. bool MultiRamp::add(const MultiRamp &other) { - // We walk through both ramps' dimensions innermost-to-outermost, consuming - // gcd(a_lanes, b_lanes) of lanes at a time. When a dimension is only - // partially consumed, the remaining part of that dimension corresponds to - // an "outer" sub-dim in the refined shape and its stride must be scaled - // by the factor just consumed. + // Refine the two shapes until they agree, express both in that shape, and + // add the strides elementwise. internal_assert(total_lanes() == other.total_lanes()) << "MultiRamp::add: total lane counts must match (" << total_lanes() << " vs " << other.total_lanes() << ")"; + Expr new_base = simplify(base + other.base); if (lanes.empty()) { // Both are 0-dim scalars. - base = simplify(base + other.base); + base = new_base; return true; } - MultiRamp result; - result.base = simplify(base + other.base); - size_t ai = 0, bi = 0; - int a_lanes = lanes[0], b_lanes = other.lanes[0]; - Expr a_stride = strides[0], b_stride = other.strides[0]; - while (true) { - int next_lanes = gcd(a_lanes, b_lanes); - if (next_lanes == 1) { - // The two next lanes are coprime, e.g: - // [0, 1, 2, 100, 101, 102] + [0, 1, 100, 101, 200, 201] - // which has no common refinement. - return false; - } - result.strides.emplace_back(simplify(a_stride + b_stride)); - result.lanes.push_back(next_lanes); - a_lanes /= next_lanes; - b_lanes /= next_lanes; - bool a_done = false, b_done = false; - if (a_lanes == 1) { - ai++; - if (ai >= lanes.size()) { - a_done = true; - } else { - a_lanes = lanes[ai]; - a_stride = strides[ai]; - } - } else { - // Remaining portion of current A-dim has a scaled stride. - a_stride = simplify(a_stride * next_lanes); - } - if (b_lanes == 1) { - bi++; - if (bi >= other.lanes.size()) { - b_done = true; - } else { - b_lanes = other.lanes[bi]; - b_stride = other.strides[bi]; - } - } else { - b_stride = simplify(b_stride * next_lanes); - } - if (a_done && b_done) { - collapse_adjacent_dims(&result); - *this = std::move(result); - return true; - } - // The up-front lane-count check ensures both sides always exhaust - // together, so neither side should be done here. + std::vector refined; + if (!common_refinement(lanes, other.lanes, &refined, nullptr)) { + // e.g. [0, 1, 2, 100, 101, 102] + [0, 1, 100, 101, 200, 201] + return false; + } + std::vector a_strides, b_strides; + bool ok = (strides_for_shape(refined, &a_strides) && + other.strides_for_shape(refined, &b_strides)); + internal_assert(ok) << "a common refinement should group into both shapes\n"; + std::vector new_strides(refined.size()); + for (size_t i = 0; i < refined.size(); i++) { + new_strides[i] = simplify(a_strides[i] + b_strides[i]); } + *this = MultiRamp(new_base, new_strides, refined); + return true; } bool MultiRamp::strides_for_shape(const std::vector &target_lanes, @@ -477,18 +475,14 @@ bool is_multiramp_impl(const Expr &e, const Scope &scope, MultiRamp *resul return true; } else if (const Shuffle *s = e.as(); s && s->vectors.size() == 1) { // A shuffle of a single vector is a reshaping of it, rather than a - // gather, if the lane indices are themselves a multiramp. That covers - // transposes, whose masks are multiramps of constants. But we can only - // say what the result is if the values being shuffled are an affine - // function of the lane index, i.e. the input is one-dimensional. This - // is the shape that flatten_nested_ramps leaves a strided load in. + // gather, when its lane indices are themselves a multiramp that + // permutes them. This is the shape flatten_nested_ramps leaves a + // strided load in, and what a transpose of a tile looks like. MultiRamp inner, perm; if (is_multiramp(s->vectors[0], scope, &inner) && - inner.dimensions() == 1 && - multiramp_of_constants(s->indices, inner.base.type(), &perm)) { - perm.mul(inner.strides[0]); - perm.base = simplify(perm.base + inner.base); - *result = perm; + multiramp_of_constants(s->indices, inner.base.type(), &perm) && + inner.shuffle(perm)) { + *result = inner; return true; } return false; @@ -619,7 +613,6 @@ int get_subtile(const Expr &index, const std::string &description, if (mr.dimensions() != first.dimensions()) { user_error << "Access to " << description << " does not have the same shape as " << "other accesses to the same memory."; - return -1; } for (int i = 0; i < first.dimensions(); i++) { if (!can_prove(mr.strides[i] == first.strides[i]) || @@ -652,7 +645,6 @@ int get_subtile(const Expr &index, const std::string &description, if (!can_prove(mr.alias_free())) { user_error << "Failed to prove access to " << description << " does not " << "partially overlap another distinct access: " << index; - return -1; } } @@ -756,6 +748,71 @@ std::vector MultiRamp::alias_free_slice() { return peeled; } +bool MultiRamp::shuffle(const MultiRamp &mask) { + if (mask.total_lanes() != total_lanes()) { + return false; + } + // A permutation of [0, n) starts at zero and steps by constants. + auto base = as_const_int(simplify(mask.base)); + if (!base || *base != 0) { + return false; + } + std::vector mask_strides; + for (const Expr &e : mask.strides) { + auto s = as_const_int(simplify(e)); + if (!s) { + return false; + } + mask_strides.push_back(*s); + } + + // Sorting the mask's dims by stride has to give the prefix products of + // their lane counts. Anything else walks the lane index in a way that + // isn't a reshaping of it. + std::vector order(mask_strides.size()); + std::iota(order.begin(), order.end(), 0); + std::sort(order.begin(), order.end(), + [&](int a, int b) { return mask_strides[a] < mask_strides[b]; }); + std::vector mask_shape; + int64_t expected = 1; + for (int d : order) { + if (mask_strides[d] != expected) { + return false; + } + mask_shape.push_back(mask.lanes[d]); + expected *= mask.lanes[d]; + } + + // Refine the mask's shape and ours against each other, tracking which of + // the mask's dims each refined dim came from. + std::vector refined, from_mask_dim; + if (!common_refinement(mask_shape, lanes, &refined, &from_mask_dim)) { + return false; + } + // A common refinement groups into our dims, so this cannot fail. + std::vector refined_strides; + bool ok = strides_for_shape(refined, &refined_strides); + internal_assert(ok) << "refined shape does not group into the multiramp's dims\n"; + + // Emit the refined dims grouped and ordered the way the mask reads them. + std::vector sorted_pos(order.size()); + for (int k = 0; k < (int)order.size(); k++) { + sorted_pos[order[k]] = k; + } + std::vector new_strides; + std::vector new_lanes; + for (int d = 0; d < mask.dimensions(); d++) { + for (size_t k = 0; k < refined.size(); k++) { + if (from_mask_dim[k] == sorted_pos[d]) { + new_strides.push_back(refined_strides[k]); + new_lanes.push_back(refined[k]); + } + } + } + *this = MultiRamp(this->base, new_strides, new_lanes); + return true; +} + int MultiRamp::rotate_stride_one_innermost() { int k = -1; for (int i = 0; i < dimensions(); i++) { diff --git a/src/MultiRamp.h b/src/MultiRamp.h index 94653491255a..a720b04fcdd8 100644 --- a/src/MultiRamp.h +++ b/src/MultiRamp.h @@ -152,6 +152,19 @@ struct MultiRamp { * a vector in the old lane order from one in the new order. */ int rotate_stride_one_innermost(); + /** Permute the lanes by a mask. `mask` gives, for each lane of the + * result, which lane of *this to take, so its lanes must be a permutation + * of [0, total_lanes()). That is the case exactly when its base is zero + * and its strides are the prefix products of its lane counts in some + * order, which makes it a reshaping of the lane index rather than a + * gather - a transpose being the two-dimensional example. + * + * The mask's dims and ours are refined against each other until both are + * groupings of a common shape, and the result is that shape reordered the + * way the mask asks for. Returns false, leaving *this unchanged, if the + * mask isn't a permutation or the two shapes have no common refinement. */ + bool shuffle(const MultiRamp &mask); + /** The dimensionality. May be lower than you expected, because this * gets flattened when possible by the operations above. */ int dimensions() const; diff --git a/test/correctness/multiramp.cpp b/test/correctness/multiramp.cpp index 099ed1eaf56d..8f325f293c10 100644 --- a/test/correctness/multiramp.cpp +++ b/test/correctness/multiramp.cpp @@ -599,6 +599,84 @@ std::vector transpose_vec(const std::vector &v, int cols) { return result; } +// A mask whose strides are the prefix products of its lane counts, in some +// order, is a permutation of the lanes. Build one from a dim order. +MultiRamp permutation_mask(const std::vector &shape, const std::vector &order, + Type t) { + std::vector flat(shape.size()); + int64_t s = 1; + for (size_t i = 0; i < shape.size(); i++) { + flat[i] = s; + s *= shape[i]; + } + std::vector strides; + std::vector lanes; + for (int d : order) { + strides.push_back(make_const(t, flat[d])); + lanes.push_back(shape[d]); + } + return MultiRamp(make_zero(t), strides, lanes); +} + +// Apply a mask the slow way, for comparison. +std::vector apply_mask(const std::vector &v, const MultiRamp &mask) { + std::vector result; + for (int i : expand(mask)) { + result.push_back(v[i]); + } + return result; +} + +void check_shuffle_case(const MultiRamp &m, const std::vector &shape, + const std::vector &order, const char *msg) { + MultiRamp a = m; + MultiRamp mask = permutation_mask(shape, order, Int(32)); + auto want = apply_mask(expand(a), mask); + if (!a.shuffle(mask)) { + printf("FAIL: %s: shuffle returned false\n", msg); + failures++; + return; + } + check_seq(expand(a), want, msg, __LINE__); +} + +void check_shuffles() { + // The two-dimensional transpose the wmma pass does to an operand tile. + check_shuffle_case(MultiRamp{0, {1, 64}, {16, 16}}, {16, 16}, {1, 0}, + "transpose of a 2D tile"); + // A transpose that falls inside a dim, so it has to be split. + check_shuffle_case(MultiRamp{0, {1}, {8}}, {4, 2}, {1, 0}, + "transpose splitting a dim"); + // A permutation whose dims collapse to something coarser than the + // multiramp's, so both shapes need refining before they line up. + check_shuffle_case(MultiRamp{0, {1, 10, 100}, {2, 2, 3}}, {2, 2, 3}, {2, 0, 1}, + "permutation needing a common refinement"); + // A broadcast dim has stride zero, which never merges with its neighbour. + check_shuffle_case(MultiRamp{0, {1, 64, 0}, {16, 16, 16}}, {16, 256}, {1, 0}, + "transpose of a broadcast tile"); +} + +void check_shuffle_rejects_gather() { + // Strides that aren't the prefix products of the lane counts don't + // describe a permutation. + MultiRamp A{0, {1, 100}, {4, 4}}; + MultiRamp mask{0, {1, 2}, {4, 4}}; + CHECK(!A.shuffle(mask), "shuffle rejects a non-permutation mask"); +} + +void check_shuffle_rejects_coprime_shapes() { + // Lane counts of 3 and 2 innermost have no common refinement. + MultiRamp A{0, {1, 100}, {2, 3}}; + MultiRamp mask = permutation_mask({3, 2}, {1, 0}, Int(32)); + CHECK(!A.shuffle(mask), "shuffle rejects coprime shapes"); +} + +void check_shuffle_rejects_wrong_size() { + MultiRamp A{0, {1, 100}, {4, 4}}; + MultiRamp mask = permutation_mask({2, 2}, {1, 0}, Int(32)); + CHECK(!A.shuffle(mask), "shuffle rejects a mask of the wrong size"); +} + void check_recognize_transpose_shuffle() { Expr e = Shuffle::make_transpose(Ramp::make(Expr(0), Expr(1), 12), 4); Scope scope; @@ -718,6 +796,11 @@ int main(int argc, char **argv) { check_roundtrips(); check_reject_non_multiramp_sum(); + check_shuffles(); + check_shuffle_rejects_gather(); + check_shuffle_rejects_coprime_shapes(); + check_shuffle_rejects_wrong_size(); + check_recognize_transpose_shuffle(); check_recognize_reshaping_shuffle(); check_reject_gather_shuffle(); From 582b342f4b937ea65e29ba924f2d8e35bc59f26f Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Fri, 31 Jul 2026 09:43:47 -0700 Subject: [PATCH 28/59] Expose WMMAFragment to the Python bindings Co-Authored-By: Claude Opus 5 --- apps/tensorcore_matmul/matmul_generator.cpp | 2 +- python_bindings/src/halide/halide_/PyEnums.cpp | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/apps/tensorcore_matmul/matmul_generator.cpp b/apps/tensorcore_matmul/matmul_generator.cpp index 1ac0efa0397b..d906558ca7ad 100644 --- a/apps/tensorcore_matmul/matmul_generator.cpp +++ b/apps/tensorcore_matmul/matmul_generator.cpp @@ -179,7 +179,7 @@ class MatMul : public Halide::Generator { Var t("t"), ti("ti"), tw("tw"), tw2("tw2"), to("to"); As.compute_at(prod, ko) - .store_in(MemoryType::GPUShared) + .store_in(MemoryType::GPUSharedAsync) .align_storage(kk, (int)block_k + (int)pad_a) .split(kk, kko, kki, vec) .fuse(kko, y, t) diff --git a/python_bindings/src/halide/halide_/PyEnums.cpp b/python_bindings/src/halide/halide_/PyEnums.cpp index 8547ca6fef67..fa40d9d2f7a3 100644 --- a/python_bindings/src/halide/halide_/PyEnums.cpp +++ b/python_bindings/src/halide/halide_/PyEnums.cpp @@ -50,7 +50,8 @@ void define_enums(py::module &m) { .value("LockedCache", MemoryType::LockedCache) .value("VTCM", MemoryType::VTCM) .value("AMXTile", MemoryType::AMXTile) - .value("GPUSharedAsync", MemoryType::GPUSharedAsync); + .value("GPUSharedAsync", MemoryType::GPUSharedAsync) + .value("WMMAFragment", MemoryType::WMMAFragment); py::enum_(m, "NameMangling") .value("Default", NameMangling::Default) From a083b48bc72d1f77734b970bed5aa34eab421803 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Fri, 31 Jul 2026 11:07:15 -0700 Subject: [PATCH 29/59] Make asynchronous copies a property of the store, not the memory type A kernel could only ever have one shared memory allocation, because the PTX backend gives them all base address zero in address space 3. Adding GPUSharedAsync broke that: shared allocations are clustered by memory type, so a kernel that staged one Func synchronously and another asynchronously got two clusters, and they overlapped. The consumer read whatever the other copy had written. Being filled by the copy engine is a property of the stores to an allocation, not of the memory it lives in, so lower it as one. A pass inside fuse_gpu_thread_loops rewrites GPUSharedAsync allocations to ordinary shared memory, wraps the value of each store to them in a cuda_bypass_registers intrinsic, and awaits the copies at the end of the producer. The allocations then all cluster together again, which also lets them share space by lifetime the way they always could within a memory type. The intrinsic carries the group its copies belong to, so that consuming one Func doesn't wait for the copies into another. Waits are FIFO - cp.async can only wait for all but the newest N groups - so the backend tracks the committed groups and lowers a group to its position in that list. A barrier still waits for everything, because it publishes shared memory to the rest of the block. Since allocations are now packed together whatever their memory type, a group has to end on a 16 byte boundary for whatever follows it to be alignable enough to copy into, and the intrinsic carries the Func name so an error can name the Func rather than the packed allocation. Co-Authored-By: Claude Opus 5 --- src/CodeGen_PTX_Dev.cpp | 157 ++++++++++++++++++++-------- src/FuseGPUThreadLoops.cpp | 99 +++++++++++++++++- src/IR.cpp | 2 + src/IR.h | 14 +++ test/correctness/gpu_async_copy.cpp | 59 ++++++++++- 5 files changed, 286 insertions(+), 45 deletions(-) diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 78c0b91d3579..ade91487ba46 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -17,6 +17,7 @@ #include "ModulusRemainder.h" #include "Simplify.h" #include "Solve.h" +#include "Substitute.h" #include "Target.h" #include @@ -107,19 +108,32 @@ class CodeGen_PTX_Dev : public CodeGen_LLVM, public CodeGen_GPU_Dev { * copies into shared memory can be recognized. */ Scope alloc_memory_type; - /** Whether we're inside a producer node, which is where the wait for any - * asynchronous copies gets emitted, and whether any have been issued in - * it. */ + /** Whether we're inside a producer node. */ bool in_producer = false; - bool issued_async_copy = false; + + /** The groups of asynchronous copies committed so far, oldest first, and + * the group of any copies issued since the last commit. Waits are FIFO - + * the hardware can only wait for all but the newest N groups - so a + * group's position here is what determines the N we emit for it. */ + std::vector committed_groups; + int uncommitted_group = -1; /** Try to emit a store into shared memory as an asynchronous copy, which * moves the data straight from global memory without routing it through * registers. Returns false if this store isn't one we can do that for. */ bool codegen_async_copy(const Store *op, const char **reason); - /** Wait for any asynchronous copies issued so far to have landed. */ - void wait_for_async_copies(); + /** Close the current group of asynchronous copies, if there is one. */ + void commit_copies(); + + /** Emit a wait that leaves at most n groups of copies outstanding. */ + void emit_copy_wait(int n); + + /** Wait for the asynchronous copies in the given group to have landed. */ + void await_copies(int group); + + /** Wait for every asynchronous copy issued so far to have landed. */ + void await_all_copies(); bool supports_atomic_add(const Type &t) const override; }; @@ -284,6 +298,15 @@ void CodeGen_PTX_Dev::init_module() { } void CodeGen_PTX_Dev::visit(const Call *op) { + if (op->is_intrinsic(Call::cuda_await_copies)) { + internal_assert(op->args.size() == 1); + auto group = as_const_int(op->args[0]); + internal_assert(group) << "cuda_await_copies group is not a constant integer\n"; + await_copies((int)*group); + value = ConstantInt::get(i32_t, 0); + return; + } + if (op->is_intrinsic(Call::gpu_thread_barrier)) { // Even though we always insert a __syncthreads equivalent // (which has both a device and shared memory fence) @@ -293,7 +316,7 @@ void CodeGen_PTX_Dev::visit(const Call *op) { // A barrier tells other threads the shared memory this thread wrote is // ready, so any asynchronous copies must have landed by now. - wait_for_async_copies(); + await_all_copies(); auto fence_type_ptr = as_const_int(op->args[0]); internal_assert(fence_type_ptr) << "gpu_thread_barrier() parameter is not a constant integer.\n"; @@ -413,6 +436,15 @@ void CodeGen_PTX_Dev::visit(const Load *op) { CodeGen_LLVM::visit(op); } +// The name of the Func a marked store belonged to, for error messages. The +// store itself is named after the packed allocation it ended up in. +std::string async_copy_func_name(const Call *marker) { + internal_assert(marker->args.size() == 3); + const StringImm *name = marker->args[2].as(); + internal_assert(name) << "cuda_bypass_registers name is not a string\n"; + return name->value; +} + // A copy from global memory into shared memory can be done by the hardware // without going through registers, which saves the load, the store, and the // registers in between. The copy is asynchronous, so it has to be waited for @@ -433,25 +465,32 @@ bool CodeGen_PTX_Dev::codegen_async_copy(const Store *op, const char **reason) { return false; } - // The destination must be shared memory that was asked for asynchronously, - // and the source must be a plain load from something we didn't allocate in - // here, which is to say global memory. Stores to plain GPUShared that - // happen to match the pattern are left synchronous, so that a schedule can - // ask for either one. - const MemoryType *dst_memory_type = alloc_memory_type.find(op->name); - if (!dst_memory_type || *dst_memory_type != MemoryType::GPUSharedAsync) { - *reason = "the destination is not stored in MemoryType::GPUSharedAsync"; + // The value must be one the schedule asked to have moved without passing + // through registers, and it must be a plain load from something we didn't + // allocate in here, which is to say global memory. An ordinary store that + // happens to match the pattern is left synchronous. + // CSE and LICM lift common subexpressions of a stored value into Lets + // around it, which would hide the marker. Substituting them back in leaves + // an expression that stands alone, which is what the analysis below and + // codegen_buffer_pointer both need. It has to be held in a local, because + // everything below points into it. + const Expr stored = substitute_in_all_lets(op->value); + const Call *marker = stored.as(); + if (!(marker && marker->is_intrinsic(Call::cuda_bypass_registers))) { + *reason = "the store was not marked as an asynchronous copy"; return false; } - const Load *src = op->value.as(); + internal_assert(marker->args.size() == 3); + const Expr &copied = marker->args[0]; + auto group = as_const_int(marker->args[1]); + internal_assert(group) << "cuda_bypass_registers group is not a constant integer\n"; + + const Load *src = copied.as(); if (!src) { // A load that isn't dense is broken up into a shuffle of dense loads // well before we get here, so say what that means for the copy rather // than describing it as not being a load. - Expr value = op->value; - while (const Let *let = value.as()) { - value = let->body; - } + Expr value = copied; if (const Shuffle *s = value.as(); s && !s->vectors.empty()) { *reason = "the source is not read densely. Each copy moves one run of " @@ -475,7 +514,7 @@ bool CodeGen_PTX_Dev::codegen_async_copy(const Store *op, const char **reason) { // The hardware copies 4, 8 or 16 bytes at a time, from and to consecutive // addresses. - const Type t = op->value.type(); + const Type t = copied.type(); const int bytes = t.bytes() * t.lanes(); if (!(bytes == 4 || bytes == 8 || bytes == 16)) { *reason = "each thread must copy 4, 8 or 16 bytes at a time. Vectorize the " @@ -539,8 +578,13 @@ bool CodeGen_PTX_Dev::codegen_async_copy(const Store *op, const char **reason) { internal_assert(id != llvm::Intrinsic::not_intrinsic) << "Could not find the nvvm intrinsic " << name.str() << "\n"; llvm::Function *fn = llvm::Intrinsic::getOrInsertDeclaration(module.get(), id); + // Copies are committed in groups, so close the previous group before + // starting one for a different batch. + if (uncommitted_group != -1 && uncommitted_group != (int)*group) { + commit_copies(); + } builder->CreateCall(fn, {dst, src_ptr}); - issued_async_copy = true; + uncommitted_group = (int)*group; return true; } @@ -550,30 +594,58 @@ void CodeGen_PTX_Dev::visit(const ProducerConsumer *op) { return; } - ScopedValue old_issued(issued_async_copy, false); ScopedValue old_in(in_producer, true); codegen(op->body); - // Everything issued in here has to have landed before the values are used, - // which is after this producer. - wait_for_async_copies(); } -void CodeGen_PTX_Dev::wait_for_async_copies() { - if (!issued_async_copy) { +void CodeGen_PTX_Dev::commit_copies() { + if (uncommitted_group == -1) { return; } - for (const char *intrin : {"llvm.nvvm.cp.async.commit.group", - "llvm.nvvm.cp.async.wait.group"}) { - llvm::Intrinsic::ID id = llvm::Intrinsic::lookupIntrinsicID(intrin); - internal_assert(id != llvm::Intrinsic::not_intrinsic); - llvm::Function *fn = llvm::Intrinsic::getOrInsertDeclaration(module.get(), id); - vector args; - if (fn->getFunctionType()->getNumParams() == 1) { - args.push_back(ConstantInt::get(i32_t, 0)); + llvm::Intrinsic::ID id = + llvm::Intrinsic::lookupIntrinsicID("llvm.nvvm.cp.async.commit.group"); + internal_assert(id != llvm::Intrinsic::not_intrinsic); + builder->CreateCall(llvm::Intrinsic::getOrInsertDeclaration(module.get(), id), {}); + committed_groups.push_back(uncommitted_group); + uncommitted_group = -1; +} + +void CodeGen_PTX_Dev::emit_copy_wait(int n) { + llvm::Intrinsic::ID id = + llvm::Intrinsic::lookupIntrinsicID("llvm.nvvm.cp.async.wait.group"); + internal_assert(id != llvm::Intrinsic::not_intrinsic); + llvm::Function *fn = llvm::Intrinsic::getOrInsertDeclaration(module.get(), id); + vector args; + if (fn->getFunctionType()->getNumParams() == 1) { + args.push_back(ConstantInt::get(i32_t, n)); + } + builder->CreateCall(fn, args); +} + +void CodeGen_PTX_Dev::await_copies(int group) { + commit_copies(); + // The wait is FIFO, so waiting for this group means letting everything + // committed after it stay outstanding. Searching from the newest end finds + // the most recent batch with this group, which is the one just issued. + for (size_t i = committed_groups.size(); i > 0; i--) { + if (committed_groups[i - 1] != group) { + continue; } - builder->CreateCall(fn, args); + emit_copy_wait((int)(committed_groups.size() - i)); + committed_groups.erase(committed_groups.begin(), + committed_groups.begin() + i); + return; + } + // Nothing from that group is outstanding, so there is nothing to wait for. +} + +void CodeGen_PTX_Dev::await_all_copies() { + commit_copies(); + if (committed_groups.empty()) { + return; } - issued_async_copy = false; + emit_copy_wait(0); + committed_groups.clear(); } void CodeGen_PTX_Dev::visit(const Store *op) { @@ -588,13 +660,14 @@ void CodeGen_PTX_Dev::visit(const Store *op) { if (codegen_async_copy(op, &reason)) { return; } - // Asking for this memory type is a promise that the stores to it are + // Asking for that memory type is a promise that the stores to it are // copies the hardware can make asynchronously. If one isn't, say so // rather than quietly emitting a load and a store instead. - const MemoryType *t = alloc_memory_type.find(op->name); - if (t && *t == MemoryType::GPUSharedAsync) { + const Expr stored = substitute_in_all_lets(op->value); + const Call *marker = stored.as(); + if (marker && marker->is_intrinsic(Call::cuda_bypass_registers)) { user_error - << op->name << " is scheduled in GPUSharedAsync memory, but this " + << async_copy_func_name(marker) << " is scheduled in GPUSharedAsync memory, but this " << "store to it cannot be done with an asynchronous copy, because " << reason << ".\n\n" << "An asynchronous copy moves bytes from global memory into shared " diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index 5d3d0d9481ac..59b0c64be848 100644 --- a/src/FuseGPUThreadLoops.cpp +++ b/src/FuseGPUThreadLoops.cpp @@ -30,6 +30,97 @@ using std::vector; namespace { +// Being filled by the copy engine is a property of the stores to an +// allocation, not of the memory it lives in, so MemoryType::GPUSharedAsync is +// only the schedule's way of saying it. Rewrite such allocations to ordinary +// shared memory, wrapping the value of each store to them in a +// cuda_bypass_registers intrinsic and awaiting the copies where the data is +// consumed. Everything below then sees one kind of shared memory, which is +// what lets all of a kernel's shared allocations be packed together. +class MarkAsyncCopies : public IRMutator { + using IRMutator::visit; + + // The allocations being filled by the copy engine, and the group each + // one's copies belong to. Copies in a group are awaited together. + map groups; + int next_group = 0; + + // Only CUDA has a copy engine to drive. Elsewhere the allocation still + // becomes ordinary shared memory, but the stores to it stay ordinary too. + DeviceAPI device_api = DeviceAPI::None; + + Stmt visit(const For *op) override { + ScopedValue d(device_api, op->device_api == DeviceAPI::None ? + device_api : + op->device_api); + return IRMutator::visit(op); + } + + Stmt visit(const Allocate *op) override { + if (op->memory_type != MemoryType::GPUSharedAsync) { + return IRMutator::visit(op); + } + + // One group per allocation, so that consuming one Func doesn't wait + // for the copies into another. + auto [it, inserted] = groups.emplace(op->name, next_group++); + internal_assert(inserted) + << "Two asynchronously copied allocations are both named " << op->name << "\n"; + + Stmt body = mutate(op->body); + groups.erase(it); + + return Allocate::make(op->name, op->type, MemoryType::GPUShared, + op->extents, mutate(op->condition), body, + op->new_expr, op->free_function, op->padding); + } + + Stmt visit(const Store *op) override { + auto it = groups.find(op->name); + if (it == groups.end() || device_api != DeviceAPI::CUDA) { + return IRMutator::visit(op); + } + // The Func name is carried along because the allocations get packed + // together below, after which the store no longer knows which Func it + // belongs to, and that is the name an error has to name. + Expr value = Call::make(op->value.type(), Call::cuda_bypass_registers, + {mutate(op->value), it->second, StringImm::make(op->name)}, + Call::Intrinsic); + return Store::make(op->name, value, mutate(op->index), op->param, + mutate(op->predicate), op->alignment, op->is_streaming); + } + + Stmt visit(const ProducerConsumer *op) override { + Stmt body = mutate(op->body); + auto it = groups.find(op->name); + if (op->is_producer && it != groups.end() && device_api == DeviceAPI::CUDA) { + // At the end of the producer, which is before the barrier that + // publishes the data to the rest of the block, and before any of + // it is read. + Expr wait = Call::make(Int(32), Call::cuda_await_copies, + {it->second}, Call::Intrinsic); + body = Block::make(body, Evaluate::make(wait)); + } + return op->with(body); + } + +public: + using IRMutator::mutate; +}; + +// The copy engine moves up to 16 bytes at a time, and needs its destination +// aligned to the width of the copy. Allocations are packed one after another, +// so a group has to end on a 16-byte boundary for whatever follows it to be +// copyable into. Round a group's size up accordingly, given the size in bytes +// of the units it is measured in. +Expr round_up_group_size(const Expr &size, int unit_bytes) { + const int alignment = 16; + if (unit_bytes >= alignment) { + return size; + } + return align_up(size, alignment / unit_bytes); +} + class ExtractBlockSize : public IRVisitor { protected: Expr block_extent[3], block_count[3]; @@ -872,7 +963,7 @@ class ExtractSharedAndHeapAllocations : public IRMutator { int ratio = alloc.widest_type.bytes() / alloc_type.bytes(); internal_assert(ratio != 0) << "alloc_type should have been at most as wide as the widest type in group\n"; - total_size += alloc.max_size * ratio; + total_size += round_up_group_size(alloc.max_size * ratio, alloc_type.bytes()); } // Upgrade the alloc type to the widest type found, and @@ -947,7 +1038,8 @@ class ExtractSharedAndHeapAllocations : public IRMutator { offset = Variable::make(Int(32), name + "." + std::to_string(i - 1) + ".offset"); int ratio = (widest_type.bytes() / cluster[i - 1].widest_type.bytes()); internal_assert(ratio != 0); - offset += simplify((cluster[i - 1].max_size + ratio - 1) / ratio); + offset += simplify(round_up_group_size( + (cluster[i - 1].max_size + ratio - 1) / ratio, widest_type.bytes())); } else { if (memory_type == MemoryType::Heap) { // One slice of a larger global allocation @@ -1743,6 +1835,9 @@ Stmt fuse_gpu_thread_loops(Stmt s) { // into the innermost GPU block. FuseGPUThreadLoops would then // merge the predicate into the merged GPU thread. s = NormalizeIfStatements()(s); + // Must run before the allocations are packed together, because packing + // them relies on there being only one kind of shared memory. + s = MarkAsyncCopies().mutate(s); s = FuseGPUThreadLoops()(s); s = ZeroGPULoopMins()(s); return s; diff --git a/src/IR.cpp b/src/IR.cpp index 586292179f21..29deac32161d 100644 --- a/src/IR.cpp +++ b/src/IR.cpp @@ -803,6 +803,8 @@ constexpr const char *intrinsic_op_names[] = { "concat_bits", "count_leading_zeros", "count_trailing_zeros", + "cuda_await_copies", + "cuda_bypass_registers", "debug_to_file", "declare_box_touched", "div_round_to_zero", diff --git a/src/IR.h b/src/IR.h index c1a7d4430cf0..111e158faa6f 100644 --- a/src/IR.h +++ b/src/IR.h @@ -716,6 +716,20 @@ struct Call : public ExprNode { concat_bits, count_leading_zeros, count_trailing_zeros, + // cuda_await_copies(group) waits for every asynchronous copy in the + // given group to have landed in shared memory. Appears in an Evaluate + // node at the point where the copied data is first read. + cuda_await_copies, + // cuda_bypass_registers(value, group, func_name) marks a value that must be moved + // to its destination without being materialized in registers, which the + // CUDA backend does with the copy engine. It is only valid as the whole + // value of a Store, and promises that the store is a copy the copy + // engine can make: a dense, aligned, unpredicated run of 4, 8 or 16 + // bytes read from outside the kernel. `group` names a batch of such + // copies that are waited for together by cuda_await_copies, and + // `func_name` is the Func the store belonged to before the shared + // allocations were packed together, for error messages. + cuda_bypass_registers, debug_to_file, // Declares that a box region of an allocation has been touched (used by bounds inference) declare_box_touched, diff --git a/test/correctness/gpu_async_copy.cpp b/test/correctness/gpu_async_copy.cpp index 403068cbbdad..80bd32a2f20e 100644 --- a/test/correctness/gpu_async_copy.cpp +++ b/test/correctness/gpu_async_copy.cpp @@ -8,7 +8,8 @@ using namespace Halide; // The copy engine moves 4, 8 or 16 bytes per thread, so the cases below cover // each of those widths at several element sizes, as well as the shapes a // staged input tends to take: a two-dimensional tile, more than one input -// staged into the same kernel, and a wrapper made with Func::in. +// staged into the same kernel, a kernel that mixes the two staging modes, and +// a wrapper made with Func::in. namespace { @@ -137,6 +138,36 @@ void test_two_inputs() { check_result("two_inputs", result, a); } +// One input staged by the copy engine and one staged the ordinary way, in +// either order. The two get separate shared allocations, which have to be +// placed at different addresses within the block's shared memory. +void test_mixed(const char *name, MemoryType first, MemoryType second) { + const int W = 256, H = 32; + Buffer a = make_input(W, H); + Buffer b = make_input(W, H); + + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func sa("sa"), sb("sb"), out("out"); + sa(x, y) = a(x, y); + sb(x, y) = b(x, y); + out(x, y) = sa(x, y) + sb(x, y); + + out.gpu_tile(x, y, xi, yi, 64, 8); + sa.compute_at(out, x) + .store_in(first) + .gpu_threads(y) + .vectorize(x, 4); + sb.compute_at(out, x) + .store_in(second) + .gpu_threads(y) + .vectorize(x, 4); + + Buffer result(W, H); + out.realize(result); + result.copy_to_host(); + check_result(name, result, a); +} + // Func::in is the idiomatic way to get a Func that is a plain copy, and is // what the error message points users at. void test_wrapper() { @@ -190,9 +221,33 @@ void test_opt_out() { check_result("opt_out_async", async, input); } +// GPU APIs with no copy engine have to treat this memory type as ordinary +// shared memory rather than failing to compile. Only compiling is needed, so +// this runs whether or not such a device is present. +void test_other_gpu_apis() { + const int W = 256, H = 32; + Buffer input = make_input(W, H); + + for (Target::Feature api : {Target::OpenCL, Target::Metal}) { + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = input(x, y); + out(x, y) = stage(x, y) * 2; + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, 4); + out.compile_to_module({}, "f", get_host_target().with_feature(api)); + } + printf("[other_gpu_apis] OK\n"); +} + } // namespace int main(int argc, char **argv) { + test_other_gpu_apis(); + Target target = get_jit_target_from_environment(); if (!target.has_feature(Target::CUDA)) { printf("[SKIP] No CUDA target enabled.\n"); @@ -218,6 +273,8 @@ int main(int argc, char **argv) { test_2d_tile(); test_padded_storage(); test_two_inputs(); + test_mixed("mixed_sync_first", MemoryType::GPUShared, MemoryType::GPUSharedAsync); + test_mixed("mixed_async_first", MemoryType::GPUSharedAsync, MemoryType::GPUShared); test_wrapper(); test_opt_out(); From b01053172ee4008bbea73c04d327a61f46de6117 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Fri, 31 Jul 2026 11:13:02 -0700 Subject: [PATCH 30/59] Add the asynchronous copy tests to the CMake build The Makefile globs the test directory, so these have been running there since they were written, but CMake lists each test explicitly and they were never added. That is what CI builds from. Co-Authored-By: Claude Opus 5 --- test/correctness/CMakeLists.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index 141c14ff1788..f58a90655695 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -135,6 +135,8 @@ tests( gpu_allocation_cache.cpp gpu_arg_types.cpp gpu_assertion_in_kernel.cpp + gpu_async_copy.cpp + gpu_async_copy_errors.cpp gpu_bounds_inference_failure.cpp gpu_condition_lifting.cpp gpu_cpu_simultaneous_read.cpp From a1914d6f81119215eaea1375c3540cef3330b27f Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Fri, 31 Jul 2026 11:29:19 -0700 Subject: [PATCH 31/59] Move GPU barrier sinking out to its own branch Eliding a barrier by widening a later one is an optimization for kernels with more than one shared allocation, and has nothing to do with asynchronous copies - it fires just as often on kernels that have none. It also arrived with no tests, and a barrier wrongly elided is a race rather than a wrong answer, so it needs its own. The wait for outstanding copies at a barrier stays here, because a barrier publishing shared memory to the rest of the block does require that any copies into it have landed. Co-Authored-By: Claude Opus 5 --- src/FuseGPUThreadLoops.cpp | 110 ------------------------------------- 1 file changed, 110 deletions(-) diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index 59b0c64be848..9b2fd91b2f79 100644 --- a/src/FuseGPUThreadLoops.cpp +++ b/src/FuseGPUThreadLoops.cpp @@ -1332,98 +1332,6 @@ class ExtractRegisterAllocations : public IRMutator { bool has_thread_loop = false; }; -// Gather the names loaded from before the first barrier at the top level of a -// statement. Barriers under a loop or an if don't count, because they don't -// necessarily run before the loads that follow. -class LoadsBeforeBarrier : public IRVisitor { - using IRVisitor::visit; - - bool at_top_level = true; - - void visit(const Block *op) override { - op->first.accept(this); - if (!found_barrier && op->rest.defined()) { - op->rest.accept(this); - } - } - - void visit(const For *op) override { - ScopedValue s(at_top_level, false); - IRVisitor::visit(op); - } - - void visit(const IfThenElse *op) override { - ScopedValue s(at_top_level, false); - IRVisitor::visit(op); - } - - void visit(const Evaluate *op) override { - const Call *c = op->value.as(); - if (at_top_level && c && c->is_intrinsic(Call::gpu_thread_barrier)) { - found_barrier = true; - } else { - IRVisitor::visit(op); - } - } - - void visit(const Load *op) override { - loads.insert(op->name); - IRVisitor::visit(op); - } - -public: - bool found_barrier = false; - std::set loads; -}; - -// Add fence types to the first barrier at the top level of a statement, so -// that it can stand in for one that would otherwise have preceded it. -class WidenFirstBarrier : public IRMutator { - using IRMutator::visit; - - bool at_top_level = true; - int mask; - - Stmt visit(const Block *op) override { - Stmt first = mutate(op->first); - if (done || !op->rest.defined()) { - return Block::make(first, op->rest); - } - return Block::make(first, mutate(op->rest)); - } - - Stmt visit(const For *op) override { - ScopedValue s(at_top_level, false); - return op; - } - - Stmt visit(const IfThenElse *op) override { - ScopedValue s(at_top_level, false); - return op; - } - - Stmt visit(const Evaluate *op) override { - const Call *c = op->value.as(); - if (at_top_level && !done && c && c->is_intrinsic(Call::gpu_thread_barrier)) { - done = true; - auto old_mask = as_const_int(c->args[0]); - internal_assert(old_mask); - return Evaluate::make(Call::make(Int(32), Call::gpu_thread_barrier, - {IntImm::make(Int(32), *old_mask | mask)}, - Call::Intrinsic)); - } - return op; - } - -public: - using IRMutator::mutate; - - bool done = false; - WidenFirstBarrier(int mask) - : mask(mask) { - } -}; - class InjectThreadBarriers : public IRMutator { protected: bool in_threads = false, injected_barrier; @@ -1566,24 +1474,6 @@ class InjectThreadBarriers : public IRMutator { break; } } - // If nothing in rest reads what first wrote until after a barrier - // of its own, that barrier can stand in for this one. - LoadsBeforeBarrier lbb; - rest.accept(&lbb); - bool needed = false; - for (const auto &st : shared_stores) { - needed |= lbb.loads.count(st) > 0; - } - for (const auto &st : device_stores) { - needed |= lbb.loads.count(st) > 0; - } - if (!needed && lbb.found_barrier) { - WidenFirstBarrier widen(mask); - rest = widen.mutate(rest); - internal_assert(widen.done); - injected_barrier = true; - return Block::make(first, rest); - } injected_barrier = true; return Block::make({first, make_barrier(mask), rest}); } else { From 218c95ab964097e7dabcc5ee7a27920c5f3d70ca Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Fri, 31 Jul 2026 13:07:42 -0700 Subject: [PATCH 32/59] Merge the tensorcore matmul app into cuda_mat_mul The two apps computed the same product - cuda_mat_mul's A(x,r)*B(r,y) is the tensorcore app's Bs(x,k)*As(k,y) with the operands swapped - and the tensorcore app's second schedule was a copy of cuda_mat_mul's. Keep one app with one generator. The operands are now untyped, so their type is a generator param, and it is what picks the schedule: half precision gets the tensor cores, and anything else gets the schedule that accumulates in ordinary registers. Both accumulate in single precision, so the two are directly comparable. The wrappers around the operands keep the operand type rather than the accumulator type, so that half precision operands are staged through shared memory as half precision and reach the tensor cores as such. For 1024x1024 on this machine: 6789 GFlop/s for float, 37849 GFlop/s for half, against 14300 GFlop/s for cublas at single precision. The tensorcore app had no CMakeLists at all, so it was only ever built by its Makefile. Folding it in gets it into the CMake build. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/CMakeLists.txt | 14 +- apps/cuda_mat_mul/Makefile | 17 +- apps/cuda_mat_mul/mat_mul_generator.cpp | 238 ++++++++++++++++---- apps/cuda_mat_mul/runner.cpp | 107 ++++++--- apps/tensorcore_matmul/Makefile | 38 ---- apps/tensorcore_matmul/matmul_generator.cpp | 215 ------------------ apps/tensorcore_matmul/runner.cpp | 84 ------- 7 files changed, 303 insertions(+), 410 deletions(-) delete mode 100644 apps/tensorcore_matmul/Makefile delete mode 100644 apps/tensorcore_matmul/matmul_generator.cpp delete mode 100644 apps/tensorcore_matmul/runner.cpp diff --git a/apps/cuda_mat_mul/CMakeLists.txt b/apps/cuda_mat_mul/CMakeLists.txt index 803f28c5ecdb..3de2599826be 100644 --- a/apps/cuda_mat_mul/CMakeLists.txt +++ b/apps/cuda_mat_mul/CMakeLists.txt @@ -27,12 +27,20 @@ find_package(Halide REQUIRED) # Generator add_halide_generator(mat_mul.generator SOURCES mat_mul_generator.cpp) -# Filters -add_halide_library(mat_mul FROM mat_mul.generator FEATURES cuda cuda_capability_50 PARAMS size=1024) +# Filters. The operand type picks the schedule: half precision gets the tensor +# cores, which need compute capability 7.0 or above, and float gets a schedule +# that accumulates in ordinary registers. +add_halide_library(mat_mul FROM mat_mul.generator + FEATURES cuda cuda_capability_50 + PARAMS size=1024 A.type=float32 B.type=float32) +add_halide_library(mat_mul_f16 FROM mat_mul.generator + GENERATOR mat_mul + FEATURES cuda cuda_capability_80 + PARAMS size=1024 A.type=float16 B.type=float16) # Main executable add_executable(runner runner.cpp) -target_link_libraries(runner PRIVATE mat_mul Halide::Tools CUDA::cudart CUDA::cublas) +target_link_libraries(runner PRIVATE mat_mul mat_mul_f16 Halide::Tools CUDA::cudart CUDA::cublas) # Test that the app actually works! add_test(NAME mat_mul COMMAND runner) diff --git a/apps/cuda_mat_mul/Makefile b/apps/cuda_mat_mul/Makefile index e0dfb78900fe..7f898815d4fe 100644 --- a/apps/cuda_mat_mul/Makefile +++ b/apps/cuda_mat_mul/Makefile @@ -2,11 +2,16 @@ include ../support/Makefile.inc MATRIX_SIZE ?= 1024 -CUDA_SDK ?= /usr/local/cuda-10.0 +CUDA_SDK ?= /usr/local/cuda CXXFLAGS += -I $(CUDA_SDK)/include LDFLAGS += -L $(CUDA_SDK)/lib64 -Wl,-rpath,$(CUDA_SDK)/lib64 +# The float variant runs anywhere, but the half variant is scheduled onto the +# tensor cores, which need compute capability 7.0 or above. +FLOAT_TARGET ?= host-cuda-cuda_capability_50 +HALF_TARGET ?= host-cuda-cuda_capability_80 + all: $(BIN)/$(HL_TARGET)/runner $(GENERATOR_BIN)/mat_mul.generator: mat_mul_generator.cpp $(GENERATOR_DEPS) @@ -15,9 +20,15 @@ $(GENERATOR_BIN)/mat_mul.generator: mat_mul_generator.cpp $(GENERATOR_DEPS) $(BIN)/%/mat_mul.a: $(GENERATOR_BIN)/mat_mul.generator @mkdir -p $(@D) - $^ -g mat_mul -e $(GENERATOR_OUTPUTS) -o $(@D) target=host-cuda-cuda_capability_50 size=$(MATRIX_SIZE) + $^ -g mat_mul -f mat_mul -e $(GENERATOR_OUTPUTS) -o $(@D) \ + target=$(FLOAT_TARGET) size=$(MATRIX_SIZE) A.type=float32 B.type=float32 + +$(BIN)/%/mat_mul_f16.a: $(GENERATOR_BIN)/mat_mul.generator + @mkdir -p $(@D) + $^ -g mat_mul -f mat_mul_f16 -e $(GENERATOR_OUTPUTS) -o $(@D) \ + target=$(HALF_TARGET) size=$(MATRIX_SIZE) A.type=float16 B.type=float16 -$(BIN)/%/runner: runner.cpp $(BIN)/%/mat_mul.a +$(BIN)/%/runner: runner.cpp $(BIN)/%/mat_mul.a $(BIN)/%/mat_mul_f16.a @mkdir -p $(@D) $(CXX) $(CXXFLAGS) -I$(BIN)/$* -Wall $^ -o $@ $(LDFLAGS) -lcudart -lcublas diff --git a/apps/cuda_mat_mul/mat_mul_generator.cpp b/apps/cuda_mat_mul/mat_mul_generator.cpp index 6f2cb17c8cd6..d5ce125c3632 100644 --- a/apps/cuda_mat_mul/mat_mul_generator.cpp +++ b/apps/cuda_mat_mul/mat_mul_generator.cpp @@ -1,5 +1,7 @@ #include "Halide.h" +#include + using namespace Halide; namespace { @@ -12,61 +14,219 @@ void set_alignment_and_bounds(OutputImageParam p, int size) { .set_stride(size); } +// A square matrix multiply, scheduled two ways. The operands are untyped, so +// their type is a generator param, and it is what picks the schedule: half +// precision operands get the tensor cores, and anything else gets a schedule +// that keeps the accumulator in ordinary registers. The product is always +// accumulated and returned in single precision. class MatMul : public Halide::Generator { public: GeneratorParam size{"size", 1024}; - Input> A{"A"}; - Input> B{"B"}; + + // How many tensor core tiles of accumulator each warp holds, and how many + // warps there are per block in each dimension. Zero means pick a shape + // based on the problem size. The best block gets smaller as the matrices + // do, because a large one leaves too few blocks to fill the machine. + GeneratorParam tiles_x{"tiles_x", 0}; + GeneratorParam tiles_y{"tiles_y", 0}; + GeneratorParam warps_x{"warps_x", 0}; + GeneratorParam warps_y{"warps_y", 0}; + // How much of the reduction is staged in shared memory at a time. + GeneratorParam block_r{"block_r", 32}; + // Extra elements per row of the shared panels, which spreads consecutive + // rows across different banks. A multiple of eight keeps the rows aligned + // enough for the widest asynchronous copy. + GeneratorParam pad_a{"pad_a", 8}; + GeneratorParam pad_b{"pad_b", 24}; + + Input> A{"A"}; + Input> B{"B"}; Output> out{"out"}; + // Tensor cores multiply half precision operands into a single precision + // accumulator, so asking for half precision inputs is what asks for them. + bool use_tensor_cores() const { + return A.type() == Float(16); + } + void generate() { - // 688 us on an RTX 2060 - // cublas is 512 us on the same card + r = RDom(0, size, "r"); + + // Wrappers for the operands, so that the tensor core schedule can + // stage them through shared memory. Left inline, they are just A + // and B. They keep the operand type, so that half precision operands + // are staged as half precision and reach the tensor cores as such - + // the widening to the accumulator type happens at the multiply. + As(x, rr) = A(x, rr); + Bs(rr, y) = B(rr, y); - Var x("x"), y("y"), p("p"); + prod(x, y) = 0.f; + prod(x, y) += cast(As(x, r)) * cast(Bs(r, y)); - Func prod("prod"); - RDom r(0, size); - prod(x, y) += A(x, r) * B(r, y); out(x, y) = prod(x, y); + } - Var xi, yi, xio, xii, yii, xo, yo, x_pair, xiio, ty; - RVar rxo, rxi; - - if (!using_autoscheduler()) { - out.bound(x, 0, size) - .bound(y, 0, size) - .tile(x, y, xi, yi, 64, 16) - .tile(xi, yi, xii, yii, 4, 8) - .gpu_blocks(x, y) - .gpu_threads(xi, yi) - .unroll(xii) - .unroll(yii); - prod.compute_at(out, xi) - .vectorize(x) - .unroll(y) - .update() - .reorder(x, y, r) - .vectorize(x) - .unroll(y) - .unroll(r, 8); - A.in().compute_at(prod, r).vectorize(_0).unroll(_1); - B.in().compute_at(prod, r).vectorize(_0).unroll(_1); - - set_alignment_and_bounds(A, size); - set_alignment_and_bounds(B, size); - set_alignment_and_bounds(out, size); - } else { + void schedule() { + if (using_autoscheduler()) { A.dim(0).set_estimate(0, size).dim(1).set_estimate(0, size); B.dim(0).set_estimate(0, size).dim(1).set_estimate(0, size); + out.bound(x, 0, size).bound(y, 0, size); + return; + } + + set_alignment_and_bounds(A, size); + set_alignment_and_bounds(B, size); + set_alignment_and_bounds(out, size); + + out.bound(x, 0, size).bound(y, 0, size); + + if (use_tensor_cores()) { + schedule_tensor_cores(); + } else { + schedule_cuda(); + } + } + +private: + // 688 us for 1024x1024 floats on an RTX 2060, where cublas is 512 us. + void schedule_cuda() { + Var xi, yi, xii, yii; + + out.tile(x, y, xi, yi, 64, 16) + .tile(xi, yi, xii, yii, 4, 8) + .gpu_blocks(x, y) + .gpu_threads(xi, yi) + .unroll(xii) + .unroll(yii); + + prod.compute_at(out, xi) + .vectorize(x) + .unroll(y) + .update() + .reorder(x, y, r) + .vectorize(x) + .unroll(y) + .unroll(r, 8); + + As.compute_at(prod, r).vectorize(x).unroll(rr); + Bs.compute_at(prod, r).vectorize(rr).unroll(y); + } + + void schedule_tensor_cores() { + // The tensor core tile shape, and how many of them each warp + // accumulates at once. Each operand tile loaded feeds tiles_x (or + // tiles_y) multiplies, so this is what gets us reuse out of the loads. + const int tile = 16; + int tx = tiles_x, ty = tiles_y, wx = warps_x, wy = warps_y; + int pb = pad_b; + if (tx == 0 || ty == 0 || wx == 0 || wy == 0) { + // The padding goes with the shape: it is what keeps consecutive + // rows of the staged panel in different banks, so the right amount + // depends on how wide the panel is. + if ((int)size <= 1024) { + // Small problems need small blocks: a 160x64 block leaves only + // a few dozen of them to cover 36 SMs, and 160 does not divide + // 1024 so the last one in each row is ragged. + tx = 8, ty = 2, wx = 1, wy = 1, pb = 8; + } else { + tx = 5, ty = 4, wx = 2, wy = 1, pb = 24; + } } + const int block_x = tile * tx * wx; + const int block_y = tile * ty * wy; - // Always specify bounds for outputs, whether autoscheduled or not - out - .bound(x, 0, size) - .bound(y, 0, size); + Var xi("xi"), yi("yi"), xt("xt"), yt("yt"), mmxi("mmxi"), mmyi("mmyi"); + Var xw("xw"), yw("yw"), rxi("rxi"), ryi("ryi"); + RVar ro("ro"), ri("ri"), rri("rri"); + + out.split(x, x, xi, block_x) + .split(xi, xt, xi, tile * tx) + .split(xi, xi, mmxi, tile) + .split(y, y, yi, block_y) + .split(yi, yt, yi, tile * ty) + .split(yi, yi, mmyi, tile) + .gpu_blocks(x, y) + .gpu_threads(xt, yt) + .reorder(mmxi, mmyi, xi, yi, xt, yt, x, y) + .unroll(xi) + .unroll(yi) + .vectorize(mmxi) + .vectorize(mmyi); + + // The accumulators live in tensor core registers for the whole + // reduction, and are written out to memory once at the end. They sit + // at block level so that the reduction loop can be above the loop over + // warps, which lets every warp share one staged panel. + prod.compute_at(out, x) + .store_in(MemoryType::WMMAFragment) + .split(x, xw, xi, tile * tx) + .split(xi, xi, rxi, tile) + .split(y, yw, yi, tile * ty) + .split(yi, yi, ryi, tile) + .reorder(rxi, ryi, xi, yi, xw, yw) + .gpu_threads(xw, yw) + .vectorize(rxi) + .vectorize(ryi) + .unroll(xi) + .unroll(yi); + + prod.update() + .split(r, ro, ri, block_r) + .split(x, xw, xi, tile * tx) + .split(xi, xi, rxi, tile) + .split(y, yw, yi, tile * ty) + .split(yi, yi, ryi, tile) + .split(ri, ri, rri, tile) + .reorder(rri, rxi, ryi, xi, yi, ri, xw, yw, ro) + .gpu_threads(xw, yw) + .unroll(xi) + .unroll(yi) + .unroll(ri) + .atomic() + .vectorize(rri) + .vectorize(rxi) + .vectorize(ryi); + + // Stage the operand panels into shared memory once per reduction step, + // to be shared by every warp in the block. Each thread moves sixteen + // bytes at a time along the dense dimension, so that the reads from + // global memory coalesce and the writes to shared memory can be done + // as asynchronous copies. + const int vec = 8; + Var rro("rro"), rrv("rrv"), xxo("xxo"), xxi("xxi"); + Var t("t"), ti("ti"), tw("tw"), tw2("tw2"), to("to"); + + // Bs is dense in the reduction dimension. + Bs.compute_at(prod, ro) + .store_in(MemoryType::GPUSharedAsync) + .align_storage(rr, (int)block_r + (int)pad_a) + .split(rr, rro, rrv, vec) + .fuse(rro, y, t) + .split(t, t, ti, 32) + .split(t, t, tw, wx) + .split(t, to, tw2, wy) + .gpu_lanes(ti) + .gpu_threads(tw, tw2) + .vectorize(rrv); + + // As is dense in x. + As.compute_at(prod, ro) + .store_in(MemoryType::GPUSharedAsync) + .align_storage(x, block_x + pb) + .split(x, xxo, xxi, vec) + .fuse(xxo, rr, t) + .split(t, t, ti, 32) + .split(t, t, tw, wx) + .split(t, to, tw2, wy) + .gpu_lanes(ti) + .gpu_threads(tw, tw2) + .vectorize(xxi); } + + Var x{"x"}, y{"y"}, rr{"rr"}; + RDom r; + Func prod{"prod"}, As{"As"}, Bs{"Bs"}; }; } // namespace diff --git a/apps/cuda_mat_mul/runner.cpp b/apps/cuda_mat_mul/runner.cpp index 898496632802..257b76da35e1 100644 --- a/apps/cuda_mat_mul/runner.cpp +++ b/apps/cuda_mat_mul/runner.cpp @@ -1,17 +1,55 @@ #include "HalideBuffer.h" #include "HalideRuntimeCuda.h" #include "halide_benchmark.h" -#include "mat_mul.h" +#include #include #include #include +#include "mat_mul.h" +#include "mat_mul_f16.h" + using Halide::Runtime::Buffer; using Halide::Tools::benchmark; +namespace { + +// The same matrix multiply is compiled twice from one generator: once with +// float operands, which get a schedule that accumulates in ordinary registers, +// and once with half operands, which get the tensor cores. Both accumulate in +// and return single precision, so the two are directly comparable. + +template +bool check(const Buffer &A, const Buffer &B, + const Buffer &C, int size, const char *name) { + // Spot check on strides that are coprime with the tile sizes, so the + // samples land at varying offsets within a tile. + for (int y = 0; y < size; y += 97) { + for (int x = 0; x < size; x += 89) { + float correct = 0.f; + for (int k = 0; k < size; k++) { + correct += (float)A(x, k) * (float)B(k, y); + } + // The operands are small integers, which are exact in both float + // and half, and the accumulator is single precision either way, so + // the answer should be exact. + if (C(x, y) != correct) { + printf("%s: bad result at %d %d: %f != %f\n", + name, x, y, C(x, y), correct); + return false; + } + } + } + return true; +} + +double gflops(int size, double seconds) { + return 2.0 * size * size * size / seconds * 1e-9; +} + +} // namespace + int main(int argc, char **argv) { - // Our Generator is compiled using cuda_capability_50; if the system running this - // test doesn't have at least that, quietly skip the test. const auto *interface = halide_cuda_device_interface(); assert(interface->compute_capability != nullptr); int major, minor; @@ -28,41 +66,54 @@ int main(int argc, char **argv) { size = atoi(argv[1]); } - // Check correctness using small-integer matrices - if (1) { + { Buffer A(size, size), B(size, size), C(size, size); - A.for_each_value([](float &v) { v = (rand() & 3) - 1; }); - B.for_each_value([](float &v) { v = (rand() & 3) - 1; }); + A.for_each_value([](float &v) { v = (float)((rand() & 3) - 1); }); + B.for_each_value([](float &v) { v = (float)((rand() & 3) - 1); }); A.set_host_dirty(); B.set_host_dirty(); mat_mul(A, B, C); C.copy_to_host(); - for (int y = 0; y < size; y++) { - for (int x = 0; x < size; x++) { - float correct = 0.f; - for (int k = 0; k < size; k++) { - correct += A(x, k) * B(k, y); - } - float actual = C(x, y); - if (correct != actual) { - printf("%d %d: %f vs %f\n", x, y, correct, actual); - return -1; - } - } + if (!check(A, B, C, size, "float")) { + return 1; } - } - // Benchmark it - { - Buffer A(size, size), B(size, size), C(size, size); - double t = Halide::Tools::benchmark(5, 5, [&]() { + double t = benchmark(5, 5, [&]() { mat_mul(A, B, C); C.device_sync(); }); - printf("Halide time: %f\n", t); + printf("Halide float: %f s (%.1f GFlop/s)\n", t, gflops(size, t)); + } + + // The half variant is scheduled onto the tensor cores. + if (ver < 70) { + printf("[SKIP] Tensor cores require compute capability 7.0 or above; " + "this system has %d.%d.\n", + major, minor); + } else { + // _Float16 rather than Halide::float16_t, so that this stays a + // runtime-only program that doesn't link the compiler. + Buffer<_Float16, 2> A(size, size), B(size, size); + Buffer C(size, size); + A.for_each_value([](_Float16 &v) { v = (_Float16)((rand() & 3) - 1); }); + B.for_each_value([](_Float16 &v) { v = (_Float16)((rand() & 3) - 1); }); + A.set_host_dirty(); + B.set_host_dirty(); + mat_mul_f16(A, B, C); + C.copy_to_host(); + if (!check(A, B, C, size, "half")) { + return 1; + } + + double t = benchmark(5, 5, [&]() { + mat_mul_f16(A, B, C); + C.device_sync(); + }); + printf("Halide half (tensor cores): %f s (%.1f GFlop/s)\n", + t, gflops(size, t)); } - // Benchmark cublas + // Benchmark cublas at single precision, for reference. #ifdef _MSC_VER // https://github.com/halide/Halide/issues/5053 printf("Skipping cublas on Windows; see https://github.com/halide/Halide/issues/5053\n"); @@ -75,7 +126,7 @@ int main(int argc, char **argv) { cublasHandle_t handle; cublasCreate(&handle); float alpha = 1.0f, beta = 1.0f; - double t = Halide::Tools::benchmark(5, 5, [&]() { + double t = benchmark(5, 5, [&]() { cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, size, size, size, &alpha, A, size, B, size, &beta, C, size); cudaDeviceSynchronize(); @@ -84,7 +135,7 @@ int main(int argc, char **argv) { cudaFree(B); cudaFree(C); cublasDestroy(handle); - printf("cublas time: %f\n", t); + printf("cublas float: %f s (%.1f GFlop/s)\n", t, gflops(size, t)); } #endif diff --git a/apps/tensorcore_matmul/Makefile b/apps/tensorcore_matmul/Makefile deleted file mode 100644 index 0ce43fc73816..000000000000 --- a/apps/tensorcore_matmul/Makefile +++ /dev/null @@ -1,38 +0,0 @@ -include ../support/Makefile.inc - -MATMUL_M ?= 1024 -MATMUL_N ?= 1024 -MATMUL_K ?= 1024 - -# The wmma instructions require compute capability 7.0 or above. -TENSORCORE_TARGET ?= host-cuda-cuda_capability_80 - -all: $(BIN)/$(HL_TARGET)/runner - -$(GENERATOR_BIN)/matmul.generator: matmul_generator.cpp $(GENERATOR_DEPS) - @mkdir -p $(@D) - $(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LIBHALIDE_LDFLAGS) - -$(BIN)/%/matmul_cudaonly.a: $(GENERATOR_BIN)/matmul.generator - @mkdir -p $(@D) - $^ -g matmul -f matmul_cudaonly -e $(GENERATOR_OUTPUTS) -o $(@D) \ - target=$(TENSORCORE_TARGET) gpu_schedule=cudaonly \ - M=$(MATMUL_M) N=$(MATMUL_N) K=$(MATMUL_K) - -$(BIN)/%/matmul_tensorcore.a: $(GENERATOR_BIN)/matmul.generator - @mkdir -p $(@D) - $^ -g matmul -f matmul_tensorcore -e $(GENERATOR_OUTPUTS) -o $(@D) \ - target=$(TENSORCORE_TARGET) gpu_schedule=tensorcore \ - M=$(MATMUL_M) N=$(MATMUL_N) K=$(MATMUL_K) - -$(BIN)/%/runner: runner.cpp $(BIN)/%/matmul_cudaonly.a $(BIN)/%/matmul_tensorcore.a - @mkdir -p $(@D) - $(CXX) $(CXXFLAGS) -I$(BIN)/$* -Wall \ - -DMATMUL_M=$(MATMUL_M) -DMATMUL_N=$(MATMUL_N) -DMATMUL_K=$(MATMUL_K) \ - $^ -o $@ $(LDFLAGS) $(LIBHALIDE_LDFLAGS) - -test: $(BIN)/$(HL_TARGET)/runner - $^ - -clean: - rm -rf $(BIN) diff --git a/apps/tensorcore_matmul/matmul_generator.cpp b/apps/tensorcore_matmul/matmul_generator.cpp deleted file mode 100644 index d906558ca7ad..000000000000 --- a/apps/tensorcore_matmul/matmul_generator.cpp +++ /dev/null @@ -1,215 +0,0 @@ -#include "Halide.h" - -#include - -namespace { - -using namespace Halide; - -enum class Schedule { - CUDA, - TensorCore, -}; - -class MatMul : public Halide::Generator { -public: - GeneratorParam gpu_schedule{ - "gpu_schedule", Schedule::TensorCore, - {{"cudaonly", Schedule::CUDA}, {"tensorcore", Schedule::TensorCore}}}; - - GeneratorParam M{"M", 1024}; - GeneratorParam N{"N", 1024}; - GeneratorParam K{"K", 1024}; - - // How many tensor core tiles of accumulator each warp holds, and how many - // warps there are per block in each dimension. - // Zero means pick a shape based on the problem size. The best block gets - // smaller as the matrices do, because a large one leaves too few blocks to - // fill the machine. - GeneratorParam tiles_x{"tiles_x", 0}; - GeneratorParam tiles_y{"tiles_y", 0}; - GeneratorParam warps_x{"warps_x", 0}; - GeneratorParam warps_y{"warps_y", 0}; - // How much of the reduction is staged in shared memory at a time. - GeneratorParam block_k{"block_k", 32}; - // Extra elements per row of the shared panels, which spreads consecutive - // rows across different banks. A multiple of eight keeps the rows aligned - // enough for the widest asynchronous copy. - GeneratorParam pad_a{"pad_a", 8}; - GeneratorParam pad_b{"pad_b", 24}; - - Input> matA{"matA"}; // K x M - Input> matB{"matB"}; // N x K - - Output> output{"output"}; - - void generate() { - k = RDom(0, K, "k"); - - // Wrappers for the operands, so that the tensor core schedule can - // stage them through shared memory. Left inline, they are just matA - // and matB. - As(kk, y) = matA(kk, y); - Bs(x, kk) = matB(x, kk); - - prod(x, y) = 0.f; - prod(x, y) += cast(As(k, y)) * cast(Bs(x, k)); - - output(x, y) = prod(x, y); - } - - void schedule() { - matA.dim(0).set_bounds(0, K).set_stride(1); - matA.dim(1).set_bounds(0, M).set_stride(K); - matB.dim(0).set_bounds(0, N).set_stride(1); - matB.dim(1).set_bounds(0, K).set_stride(N); - output.dim(0).set_bounds(0, N).set_stride(1); - output.dim(1).set_bounds(0, M).set_stride(N); - - if (gpu_schedule == Schedule::CUDA) { - // Schedule taken from the cuda_mat_mul app. - Var xi, yi, xii, yii; - - output.bound(x, 0, N) - .bound(y, 0, M) - .tile(x, y, xi, yi, 64, 16) - .tile(xi, yi, xii, yii, 4, 8) - .gpu_blocks(x, y) - .gpu_threads(xi, yi) - .unroll(xii) - .unroll(yii); - - prod.compute_at(output, xi) - .vectorize(x) - .unroll(y) - .update() - .reorder(x, y, k) - .vectorize(x) - .unroll(y) - .unroll(k, 8); - - matA.in().compute_at(prod, k).vectorize(_0).unroll(_1); - matB.in().compute_at(prod, k).vectorize(_0).unroll(_1); - } else { - // The tensor core tile shape, and how many of them each warp - // accumulates at once. Each operand tile loaded feeds tiles_x (or - // tiles_y) multiplies, so this is what gets us reuse out of the - // loads. - const int tile = 16; - int tx = tiles_x, ty = tiles_y, wx = warps_x, wy = warps_y; - int pb = pad_b; - if (tx == 0 || ty == 0 || wx == 0 || wy == 0) { - // The padding goes with the shape: it is what keeps consecutive - // rows of the staged panel in different banks, so the right - // amount depends on how wide the panel is. - const int n = std::min({(int)M, (int)N, (int)K}); - if (n <= 1024) { - // Small problems need small blocks: a 160x64 block leaves - // only a few dozen of them to cover 36 SMs, and 160 does - // not divide 1024 so the last one in each row is ragged. - tx = 8, ty = 2, wx = 1, wy = 1, pb = 8; - } else { - tx = 5, ty = 4, wx = 2, wy = 1, pb = 24; - } - } - const int block_x = tile * tx * wx; - const int block_y = tile * ty * wy; - - Var xi("xi"), yi("yi"), xt("xt"), yt("yt"), mmxi("mmxi"), mmyi("mmyi"); - Var xw("xw"), yw("yw"), rxi("rxi"), ryi("ryi"); - RVar ko("ko"), ki("ki"), rri("rri"); - - output.bound(x, 0, N) - .bound(y, 0, M) - .split(x, x, xi, block_x) - .split(xi, xt, xi, tile * tx) - .split(xi, xi, mmxi, tile) - .split(y, y, yi, block_y) - .split(yi, yt, yi, tile * ty) - .split(yi, yi, mmyi, tile) - .gpu_blocks(x, y) - .gpu_threads(xt, yt) - .reorder(mmxi, mmyi, xi, yi, xt, yt, x, y) - .unroll(xi) - .unroll(yi) - .vectorize(mmxi) - .vectorize(mmyi); - - // The accumulators live in tensor core registers for the whole - // reduction, and are written out to memory once at the end. They - // sit at block level so that the reduction loop can be above the - // loop over warps, which lets every warp share one staged panel. - prod.compute_at(output, x) - .store_in(MemoryType::WMMAFragment) - .split(x, xw, xi, tile * tx) - .split(xi, xi, rxi, tile) - .split(y, yw, yi, tile * ty) - .split(yi, yi, ryi, tile) - .reorder(rxi, ryi, xi, yi, xw, yw) - .gpu_threads(xw, yw) - .vectorize(rxi) - .vectorize(ryi) - .unroll(xi) - .unroll(yi); - - prod.update() - .split(k, ko, ki, block_k) - .split(x, xw, xi, tile * tx) - .split(xi, xi, rxi, tile) - .split(y, yw, yi, tile * ty) - .split(yi, yi, ryi, tile) - .split(ki, ki, rri, tile) - .reorder(rri, rxi, ryi, xi, yi, ki, xw, yw, ko) - .gpu_threads(xw, yw) - .unroll(xi) - .unroll(yi) - .unroll(ki) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi); - - // Stage the operand panels into shared memory once per reduction - // step, to be shared by every warp in the block. Each thread moves - // sixteen bytes at a time along the dense dimension, so that the - // reads from global memory coalesce and the writes to shared - // memory can be done as asynchronous copies. - const int vec = 8; - Var kko("kko"), kki("kki"), xxo("xxo"), xxi("xxi"); - Var t("t"), ti("ti"), tw("tw"), tw2("tw2"), to("to"); - - As.compute_at(prod, ko) - .store_in(MemoryType::GPUSharedAsync) - .align_storage(kk, (int)block_k + (int)pad_a) - .split(kk, kko, kki, vec) - .fuse(kko, y, t) - .split(t, t, ti, 32) - .split(t, t, tw, wx) - .split(t, to, tw2, wy) - .gpu_lanes(ti) - .gpu_threads(tw, tw2) - .vectorize(kki); - - Bs.compute_at(prod, ko) - .store_in(MemoryType::GPUSharedAsync) - .align_storage(x, block_x + pb) - .split(x, xxo, xxi, vec) - .fuse(xxo, kk, t) - .split(t, t, ti, 32) - .split(t, t, tw, wx) - .split(t, to, tw2, wy) - .gpu_lanes(ti) - .gpu_threads(tw, tw2) - .vectorize(xxi); - } - } - -private: - Var x{"x"}, y{"y"}, kk{"kk"}; - RDom k; - Func prod{"prod"}, As{"As"}, Bs{"Bs"}; -}; - -} // namespace - -HALIDE_REGISTER_GENERATOR(MatMul, matmul) diff --git a/apps/tensorcore_matmul/runner.cpp b/apps/tensorcore_matmul/runner.cpp deleted file mode 100644 index 9a98333b42f9..000000000000 --- a/apps/tensorcore_matmul/runner.cpp +++ /dev/null @@ -1,84 +0,0 @@ -#include "Halide.h" -#include "HalideBuffer.h" -#include "HalideRuntimeCuda.h" -#include "halide_benchmark.h" -#include - -#include "matmul_cudaonly.h" -#include "matmul_tensorcore.h" - -using Halide::float16_t; -using Halide::Runtime::Buffer; -using Halide::Tools::benchmark; - -namespace { - -constexpr int M = MATMUL_M, N = MATMUL_N, K = MATMUL_K; - -bool check(const Buffer &A, - const Buffer &B, - const Buffer &C, - const char *name) { - for (int y = 0; y < M; y += 97) { - for (int x = 0; x < N; x += 89) { - float ref = 0.f; - for (int k = 0; k < K; k++) { - ref += (float)A(k, y) * (float)B(x, k); - } - if (std::abs(C(x, y) - ref) > 1e-2f * std::max(1.f, std::abs(ref))) { - printf("%s: bad result at %d %d: %f != %f\n", name, x, y, C(x, y), ref); - return false; - } - } - } - return true; -} - -} // namespace - -int main(int argc, char **argv) { - const auto *interface = halide_cuda_device_interface(); - int major, minor; - if (interface->compute_capability(nullptr, &major, &minor) != 0 || - major * 10 + minor < 70) { - printf("[SKIP] Tensor cores require CUDA compute capability 7.0 or above.\n"); - return 0; - } - - Buffer A(K, M), B(N, K); - A.fill([]() { return float16_t(((float)rand() / RAND_MAX) - 0.5f); }); - B.fill([]() { return float16_t(((float)rand() / RAND_MAX) - 0.5f); }); - - Buffer C_cuda(N, M), C_tensorcore(N, M); - - matmul_cudaonly(A, B, C_cuda); - C_cuda.copy_to_host(); - if (!check(A, B, C_cuda, "cudaonly")) { - return 1; - } - - matmul_tensorcore(A, B, C_tensorcore); - C_tensorcore.copy_to_host(); - if (!check(A, B, C_tensorcore, "tensorcore")) { - return 1; - } - - // Two flops (a multiply and an add) per element of the reduction. - const double flops = 2.0 * M * N * K; - - double t_cuda = benchmark([&]() { - matmul_cudaonly(A, B, C_cuda); - C_cuda.device_sync(); - }); - double t_tensorcore = benchmark([&]() { - matmul_tensorcore(A, B, C_tensorcore); - C_tensorcore.device_sync(); - }); - - printf("cuda only: %8.3f ms %8.1f GFlop/s\n", t_cuda * 1e3, flops / t_cuda * 1e-9); - printf("tensor core: %8.3f ms %8.1f GFlop/s\n", t_tensorcore * 1e3, flops / t_tensorcore * 1e-9); - printf("speed-up: %8.2fx\n", t_cuda / t_tensorcore); - - printf("Success!\n"); - return 0; -} From bc1a36499cd80b70ca701528c00cb6fc4a3a7343 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Fri, 31 Jul 2026 14:38:08 -0700 Subject: [PATCH 33/59] Use Func::in for the matmul operand wrappers The staged copies of the operands are a scheduling concern, so make them with Func::in in the schedule rather than by writing wrapper Funcs into the algorithm. This is what the app did before the tensorcore schedule was folded in, and it leaves the algorithm as just the multiply. Also refresh the performance comments, which quoted an RTX 2060. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/mat_mul_generator.cpp | 57 ++++++++----------------- 1 file changed, 18 insertions(+), 39 deletions(-) diff --git a/apps/cuda_mat_mul/mat_mul_generator.cpp b/apps/cuda_mat_mul/mat_mul_generator.cpp index d5ce125c3632..4b7f13637c8b 100644 --- a/apps/cuda_mat_mul/mat_mul_generator.cpp +++ b/apps/cuda_mat_mul/mat_mul_generator.cpp @@ -53,16 +53,12 @@ class MatMul : public Halide::Generator { void generate() { r = RDom(0, size, "r"); - // Wrappers for the operands, so that the tensor core schedule can - // stage them through shared memory. Left inline, they are just A - // and B. They keep the operand type, so that half precision operands - // are staged as half precision and reach the tensor cores as such - - // the widening to the accumulator type happens at the multiply. - As(x, rr) = A(x, rr); - Bs(rr, y) = B(rr, y); - prod(x, y) = 0.f; - prod(x, y) += cast(As(x, r)) * cast(Bs(r, y)); + // The widening to the accumulator type happens here, at the multiply, + // rather than in the operand wrappers the schedule stages through + // shared memory. That way half precision operands are staged as half + // precision and reach the tensor cores as such. + prod(x, y) += cast(A(x, r)) * cast(B(r, y)); out(x, y) = prod(x, y); } @@ -89,7 +85,8 @@ class MatMul : public Halide::Generator { } private: - // 688 us for 1024x1024 floats on an RTX 2060, where cublas is 512 us. + // 315 us for 1024x1024 on an RTX 5060 Ti, where cublas is 150 us at the + // same precision. void schedule_cuda() { Var xi, yi, xii, yii; @@ -109,10 +106,12 @@ class MatMul : public Halide::Generator { .unroll(y) .unroll(r, 8); - As.compute_at(prod, r).vectorize(x).unroll(rr); - Bs.compute_at(prod, r).vectorize(rr).unroll(y); + A.in().compute_at(prod, r).vectorize(_0).unroll(_1); + B.in().compute_at(prod, r).vectorize(_0).unroll(_1); } + // 57 us for 1024x1024 on an RTX 5060 Ti, which is 5.5x the schedule above + // and 2.6x cublas at single precision. void schedule_tensor_cores() { // The tensor core tile shape, and how many of them each warp // accumulates at once. Each operand tile loaded feeds tiles_x (or @@ -197,36 +196,16 @@ class MatMul : public Halide::Generator { Var rro("rro"), rrv("rrv"), xxo("xxo"), xxi("xxi"); Var t("t"), ti("ti"), tw("tw"), tw2("tw2"), to("to"); - // Bs is dense in the reduction dimension. - Bs.compute_at(prod, ro) - .store_in(MemoryType::GPUSharedAsync) - .align_storage(rr, (int)block_r + (int)pad_a) - .split(rr, rro, rrv, vec) - .fuse(rro, y, t) - .split(t, t, ti, 32) - .split(t, t, tw, wx) - .split(t, to, tw2, wy) - .gpu_lanes(ti) - .gpu_threads(tw, tw2) - .vectorize(rrv); - - // As is dense in x. - As.compute_at(prod, ro) - .store_in(MemoryType::GPUSharedAsync) - .align_storage(x, block_x + pb) - .split(x, xxo, xxi, vec) - .fuse(xxo, rr, t) - .split(t, t, ti, 32) - .split(t, t, tw, wx) - .split(t, to, tw2, wy) - .gpu_lanes(ti) - .gpu_threads(tw, tw2) - .vectorize(xxi); + // B.in() is dense in the reduction dimension, which is its _0. + B.in().compute_at(prod, ro).store_in(MemoryType::GPUSharedAsync).align_storage(_0, (int)block_r + (int)pad_a).split(_0, rro, rrv, vec).fuse(rro, _1, t).split(t, t, ti, 32).split(t, t, tw, wx).split(t, to, tw2, wy).gpu_lanes(ti).gpu_threads(tw, tw2).vectorize(rrv); + + // A.in() is dense in x, which is its _0. + A.in().compute_at(prod, ro).store_in(MemoryType::GPUSharedAsync).align_storage(_0, block_x + pb).split(_0, xxo, xxi, vec).fuse(xxo, _1, t).split(t, t, ti, 32).split(t, t, tw, wx).split(t, to, tw2, wy).gpu_lanes(ti).gpu_threads(tw, tw2).vectorize(xxi); } - Var x{"x"}, y{"y"}, rr{"rr"}; + Var x{"x"}, y{"y"}; RDom r; - Func prod{"prod"}, As{"As"}, Bs{"Bs"}; + Func prod{"prod"}; }; } // namespace From 3b2c1c580a1ff5a29592a01a68137765818f2140 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Fri, 31 Jul 2026 14:51:30 -0700 Subject: [PATCH 34/59] clang-format the tensor core sources These predate this branch's cleanup and were never run through clang-format, which CI checks. The repo sets ColumnLimit to zero, so hand-wrapped initializer lists get collapsed onto one line. Co-Authored-By: Claude Opus 5 --- apps/tensorcore_resize/resize_generator.cpp | 9 ++------- src/CodeGen_PTX_Dev.cpp | 3 +-- test/correctness/wmma_matmul.cpp | 14 +------------- 3 files changed, 4 insertions(+), 22 deletions(-) diff --git a/apps/tensorcore_resize/resize_generator.cpp b/apps/tensorcore_resize/resize_generator.cpp index d434cfcffcf1..ebc9ae49dbc3 100644 --- a/apps/tensorcore_resize/resize_generator.cpp +++ b/apps/tensorcore_resize/resize_generator.cpp @@ -75,11 +75,7 @@ const KernelInfo kernel_info[] = { class Resize : public Halide::Generator { public: GeneratorParam interpolation_type{ - "interpolation_type", InterpolationType::Lanczos, - {{"box", InterpolationType::Box}, - {"linear", InterpolationType::Linear}, - {"cubic", InterpolationType::Cubic}, - {"lanczos", InterpolationType::Lanczos}}}; + "interpolation_type", InterpolationType::Lanczos, {{"box", InterpolationType::Box}, {"linear", InterpolationType::Linear}, {"cubic", InterpolationType::Cubic}, {"lanczos", InterpolationType::Lanczos}}}; // If we statically know whether we're upsampling or downsampling, we can // generate different pipelines (we want to reorder the resample in x and @@ -87,8 +83,7 @@ class Resize : public Halide::Generator { GeneratorParam upsample{"upsample", false}; GeneratorParam gpu_schedule{ - "gpu_schedule", Schedule::TensorCore, - {{"cudaonly", Schedule::CUDA}, {"tensorcore", Schedule::TensorCore}}}; + "gpu_schedule", Schedule::TensorCore, {{"cudaonly", Schedule::CUDA}, {"tensorcore", Schedule::TensorCore}}}; Input> input{"input"}; Input scale_factor{"scale_factor"}; diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 368beed16264..db2e1d3da91c 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -1,11 +1,11 @@ #include -#include "CodeGen_PTX_Dev.h" #include "CSE.h" #include "CanonicalizeGPUVars.h" #include "CodeGen_GPU_Dev.h" #include "CodeGen_Internal.h" #include "CodeGen_LLVM.h" +#include "CodeGen_PTX_Dev.h" #include "ConciseCasts.h" #include "Debug.h" #include "ExprUsesVar.h" @@ -189,7 +189,6 @@ Type CodeGen_PTX_Dev::upgrade_type_for_storage(const Type &t) const { return CodeGen_LLVM::upgrade_type_for_storage(t); } - namespace { // The size of the thread block a kernel will be launched with. ptxas allocates diff --git a/test/correctness/wmma_matmul.cpp b/test/correctness/wmma_matmul.cpp index 437bca94f434..821a4335f97f 100644 --- a/test/correctness/wmma_matmul.cpp +++ b/test/correctness/wmma_matmul.cpp @@ -214,19 +214,7 @@ bool test_block_level_accumulator() { Var kko("kko"), kki("kki"), xxo("xxo"), xxi("xxi"); RVar ko("ko"), ki("ki"), rri("rri"); - out.bound(x, 0, N).bound(y, 0, M) - .split(x, x, xi, block_x) - .split(xi, xt, xi, tile * tiles_x) - .split(xi, xi, mmxi, tile) - .split(y, y, yi, block_y) - .split(yi, yi, mmyi, tile) - .gpu_blocks(x, y) - .gpu_threads(xt) - .reorder(mmxi, mmyi, xi, yi, xt, x, y) - .unroll(xi) - .unroll(yi) - .vectorize(mmxi) - .vectorize(mmyi); + out.bound(x, 0, N).bound(y, 0, M).split(x, x, xi, block_x).split(xi, xt, xi, tile * tiles_x).split(xi, xi, mmxi, tile).split(y, y, yi, block_y).split(yi, yi, mmyi, tile).gpu_blocks(x, y).gpu_threads(xt).reorder(mmxi, mmyi, xi, yi, xt, x, y).unroll(xi).unroll(yi).vectorize(mmxi).vectorize(mmyi); prod.compute_at(out, x) .store_in(MemoryType::WMMAFragment) From adf78ff5902b35ca6425c252eb8ad9f98f08334d Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Fri, 31 Jul 2026 14:55:56 -0700 Subject: [PATCH 35/59] Benchmark cublas at half precision too The interesting comparison for the tensor core schedule is against cublas doing the same thing, which is cublasGemmEx with half precision operands and a single precision accumulator - cublasHgemm accumulates in half, so it would not be comparable. Time a batch of launches with one synchronization at the end rather than synchronizing after each. cublas does a heuristic lookup on the host for every call, and synchronizing per launch measures that instead of letting it overlap with the GPU, which understated cublas by about nine percent. Halide runs on its own CUDA context, so it needs its own sync rather than cudaDeviceSynchronize. Touch cublas's buffers before timing anything. They were benchmarked straight out of cudaMalloc, so the measurement could have included faulting them in, and the operands were whatever happened to be in memory, which risks denormals. The schedule is slower than cublas at every size, by more as the matrices grow: 96% of its throughput at 1024, 94% at 2048, 84% at 4096. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/mat_mul_generator.cpp | 12 +++- apps/cuda_mat_mul/runner.cpp | 89 +++++++++++++++++++------ 2 files changed, 78 insertions(+), 23 deletions(-) diff --git a/apps/cuda_mat_mul/mat_mul_generator.cpp b/apps/cuda_mat_mul/mat_mul_generator.cpp index 4b7f13637c8b..5c8d04bd3567 100644 --- a/apps/cuda_mat_mul/mat_mul_generator.cpp +++ b/apps/cuda_mat_mul/mat_mul_generator.cpp @@ -85,7 +85,7 @@ class MatMul : public Halide::Generator { } private: - // 315 us for 1024x1024 on an RTX 5060 Ti, where cublas is 150 us at the + // 311 us for 1024x1024 on an RTX 5060 Ti, where cublas is 146 us at the // same precision. void schedule_cuda() { Var xi, yi, xii, yii; @@ -110,8 +110,14 @@ class MatMul : public Halide::Generator { B.in().compute_at(prod, r).vectorize(_0).unroll(_1); } - // 57 us for 1024x1024 on an RTX 5060 Ti, which is 5.5x the schedule above - // and 2.6x cublas at single precision. + // 53 us for 1024x1024 on an RTX 5060 Ti, which is 5.9x the schedule above + // and 2.8x cublas at single precision. + // + // Against cublas doing the same thing - half precision operands into a + // single precision accumulator - this is slower at every size: 96% of its + // throughput at 1024, 94% at 2048, and 84% at 4096. Our throughput peaks + // at 2048 and falls off after it while cublas keeps climbing, so the block + // shapes chosen below are the thing to revisit for large matrices. void schedule_tensor_cores() { // The tensor core tile shape, and how many of them each warp // accumulates at once. Each operand tile loaded feeds tiles_x (or diff --git a/apps/cuda_mat_mul/runner.cpp b/apps/cuda_mat_mul/runner.cpp index 257b76da35e1..dcf9474efe1f 100644 --- a/apps/cuda_mat_mul/runner.cpp +++ b/apps/cuda_mat_mul/runner.cpp @@ -10,7 +10,6 @@ #include "mat_mul_f16.h" using Halide::Runtime::Buffer; -using Halide::Tools::benchmark; namespace { @@ -47,6 +46,36 @@ double gflops(int size, double seconds) { return 2.0 * size * size * size / seconds * 1e-9; } +// Time a batch of launches with a single synchronization at the end, rather +// than synchronizing after each one. Both implementations queue work +// asynchronously, and cublas in particular does a heuristic lookup on the host +// for every call, so synchronizing per launch measures that host work instead +// of letting it overlap with the GPU. +// `sync` has to match the launcher: Halide runs on its own CUDA context, so +// cudaDeviceSynchronize does not wait for it. +template +double bench_batched(F &&launch, S &&sync) { + const int samples = 5, iterations = 5; + for (int i = 0; i < iterations; i++) { + launch(); + } + sync(); + double best = 0; + for (int s = 0; s < samples; s++) { + auto t0 = Halide::Tools::benchmark_now(); + for (int i = 0; i < iterations; i++) { + launch(); + } + sync(); + auto t1 = Halide::Tools::benchmark_now(); + double t = Halide::Tools::benchmark_duration_seconds(t0, t1) / iterations; + if (s == 0 || t < best) { + best = t; + } + } + return best; +} + } // namespace int main(int argc, char **argv) { @@ -78,10 +107,8 @@ int main(int argc, char **argv) { return 1; } - double t = benchmark(5, 5, [&]() { - mat_mul(A, B, C); - C.device_sync(); - }); + double t = bench_batched([&]() { mat_mul(A, B, C); }, + [&]() { C.device_sync(); }); printf("Halide float: %f s (%.1f GFlop/s)\n", t, gflops(size, t)); } @@ -105,37 +132,59 @@ int main(int argc, char **argv) { return 1; } - double t = benchmark(5, 5, [&]() { - mat_mul_f16(A, B, C); - C.device_sync(); - }); + double t = bench_batched([&]() { mat_mul_f16(A, B, C); }, + [&]() { C.device_sync(); }); printf("Halide half (tensor cores): %f s (%.1f GFlop/s)\n", t, gflops(size, t)); } - // Benchmark cublas at single precision, for reference. + // Benchmark cublas for reference, at both precisions. The half precision + // one accumulates in single precision, matching what the Halide pipeline + // does, so the two are comparable. #ifdef _MSC_VER // https://github.com/halide/Halide/issues/5053 printf("Skipping cublas on Windows; see https://github.com/halide/Halide/issues/5053\n"); #else { - float *A, *B, *C; - cudaMalloc((void **)&A, size * size * 4); - cudaMalloc((void **)&B, size * size * 4); - cudaMalloc((void **)&C, size * size * 4); + void *A, *B, *C; + cudaMalloc(&A, (size_t)size * size * 4); + cudaMalloc(&B, (size_t)size * size * 4); + cudaMalloc(&C, (size_t)size * size * 4); + // Touch the memory before timing anything, so that no part of the + // benchmark pays for faulting it in, and so that the operands are + // definite values rather than whatever was there. This byte pattern is + // a normal number read either as float or as half, which matters + // because denormals can be slow. + cudaMemset(A, 0x3c, (size_t)size * size * 4); + cudaMemset(B, 0x3c, (size_t)size * size * 4); + cudaMemset(C, 0, (size_t)size * size * 4); cublasHandle_t handle; cublasCreate(&handle); float alpha = 1.0f, beta = 1.0f; - double t = benchmark(5, 5, [&]() { - cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, - size, size, size, &alpha, A, size, B, size, &beta, C, size); - cudaDeviceSynchronize(); - }); + + double t = bench_batched([&]() { cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, + size, size, size, &alpha, (const float *)A, size, + (const float *)B, size, &beta, (float *)C, size); }, + []() { cudaDeviceSynchronize(); }); + printf("cublas float: %f s (%.1f GFlop/s)\n", t, gflops(size, t)); + + if (ver >= 70) { + // Half precision operands into a single precision accumulator, + // which is what the tensor cores do natively. + t = bench_batched([&]() { cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, + size, size, size, &alpha, + A, CUDA_R_16F, size, + B, CUDA_R_16F, size, &beta, + C, CUDA_R_32F, size, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT); }, + []() { cudaDeviceSynchronize(); }); + printf("cublas half: %f s (%.1f GFlop/s)\n", t, gflops(size, t)); + } + cudaFree(A); cudaFree(B); cudaFree(C); cublasDestroy(handle); - printf("cublas float: %f s (%.1f GFlop/s)\n", t, gflops(size, t)); } #endif From 5143ece5f765457359af312faf7abb52dc3fa0ef Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Fri, 31 Jul 2026 15:24:07 -0700 Subject: [PATCH 36/59] Retune the tensor core block shapes The shapes were picked for 1024x1024 and one guess above it, and the guess was a bad one: at 4096 the schedule ran at 84% of cublas, and its own throughput fell off from its peak at 2048 rather than climbing. Sweeping 96 shapes at each of 1024, 2048 and 4096 gives a shape per size, each 4% to 11% better at its own size than either of the others. There is no trend worth extrapolating from, so these are three measured points. What changes is how much accumulator a warp holds: 4096 wants a small one spread over four warps, where occupancy hides the memory latency better than a large accumulator's reuse does. A staging depth of 32 and a pad of 8 win at every size, so pad_b's default drops from 24 to 8 and the per-size override goes away. That takes 4096 from 84% of cublas to 97%, and leaves the schedule at 95% to 97% of it across the three sizes. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/mat_mul_generator.cpp | 31 +++++++++++++------------ 1 file changed, 16 insertions(+), 15 deletions(-) diff --git a/apps/cuda_mat_mul/mat_mul_generator.cpp b/apps/cuda_mat_mul/mat_mul_generator.cpp index 5c8d04bd3567..1a6177438f80 100644 --- a/apps/cuda_mat_mul/mat_mul_generator.cpp +++ b/apps/cuda_mat_mul/mat_mul_generator.cpp @@ -37,7 +37,7 @@ class MatMul : public Halide::Generator { // rows across different banks. A multiple of eight keeps the rows aligned // enough for the widest asynchronous copy. GeneratorParam pad_a{"pad_a", 8}; - GeneratorParam pad_b{"pad_b", 24}; + GeneratorParam pad_b{"pad_b", 8}; Input> A{"A"}; Input> B{"B"}; @@ -114,28 +114,29 @@ class MatMul : public Halide::Generator { // and 2.8x cublas at single precision. // // Against cublas doing the same thing - half precision operands into a - // single precision accumulator - this is slower at every size: 96% of its - // throughput at 1024, 94% at 2048, and 84% at 4096. Our throughput peaks - // at 2048 and falls off after it while cublas keeps climbing, so the block - // shapes chosen below are the thing to revisit for large matrices. + // single precision accumulator - this is a little slower at every size: + // 96% of its throughput at 1024, 95% at 2048, and 97% at 4096, with the + // block shapes below picked per size by measurement. void schedule_tensor_cores() { // The tensor core tile shape, and how many of them each warp // accumulates at once. Each operand tile loaded feeds tiles_x (or // tiles_y) multiplies, so this is what gets us reuse out of the loads. const int tile = 16; int tx = tiles_x, ty = tiles_y, wx = warps_x, wy = warps_y; - int pb = pad_b; if (tx == 0 || ty == 0 || wx == 0 || wy == 0) { - // The padding goes with the shape: it is what keeps consecutive - // rows of the staged panel in different banks, so the right amount - // depends on how wide the panel is. + // Measured on an RTX 5060 Ti. These do not follow a trend worth + // extrapolating from, so they are three measured points rather + // than a formula, and each is 4% to 11% better at its own size + // than either of the others would be. What changes is how much + // accumulator a warp holds: at 4096 a small one spread over four + // warps beats a large one, because the occupancy hides the memory + // latency better than a large accumulator's reuse does. if ((int)size <= 1024) { - // Small problems need small blocks: a 160x64 block leaves only - // a few dozen of them to cover 36 SMs, and 160 does not divide - // 1024 so the last one in each row is ragged. - tx = 8, ty = 2, wx = 1, wy = 1, pb = 8; + tx = 8, ty = 2, wx = 1, wy = 1; + } else if ((int)size <= 2048) { + tx = 5, ty = 4, wx = 2, wy = 2; } else { - tx = 5, ty = 4, wx = 2, wy = 1, pb = 24; + tx = 4, ty = 2, wx = 2, wy = 2; } } const int block_x = tile * tx * wx; @@ -206,7 +207,7 @@ class MatMul : public Halide::Generator { B.in().compute_at(prod, ro).store_in(MemoryType::GPUSharedAsync).align_storage(_0, (int)block_r + (int)pad_a).split(_0, rro, rrv, vec).fuse(rro, _1, t).split(t, t, ti, 32).split(t, t, tw, wx).split(t, to, tw2, wy).gpu_lanes(ti).gpu_threads(tw, tw2).vectorize(rrv); // A.in() is dense in x, which is its _0. - A.in().compute_at(prod, ro).store_in(MemoryType::GPUSharedAsync).align_storage(_0, block_x + pb).split(_0, xxo, xxi, vec).fuse(xxo, _1, t).split(t, t, ti, 32).split(t, t, tw, wx).split(t, to, tw2, wy).gpu_lanes(ti).gpu_threads(tw, tw2).vectorize(xxi); + A.in().compute_at(prod, ro).store_in(MemoryType::GPUSharedAsync).align_storage(_0, block_x + pad_b).split(_0, xxo, xxi, vec).fuse(xxo, _1, t).split(t, t, ti, 32).split(t, t, tw, wx).split(t, to, tw2, wy).gpu_lanes(ti).gpu_threads(tw, tw2).vectorize(xxi); } Var x{"x"}, y{"y"}; From 473ad084fcdf13442badef4f8f58f455462bc115 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Mon, 3 Aug 2026 11:43:50 -0700 Subject: [PATCH 37/59] Record the measured times for all four configurations One table of twelve numbers - two implementations, two precisions, three sizes - rather than a couple of ratios scattered across the schedules. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/mat_mul_generator.cpp | 27 ++++++++++++++++--------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/apps/cuda_mat_mul/mat_mul_generator.cpp b/apps/cuda_mat_mul/mat_mul_generator.cpp index 1a6177438f80..1d02d7e230d1 100644 --- a/apps/cuda_mat_mul/mat_mul_generator.cpp +++ b/apps/cuda_mat_mul/mat_mul_generator.cpp @@ -19,6 +19,20 @@ void set_alignment_and_bounds(OutputImageParam p, int size) { // precision operands get the tensor cores, and anything else gets a schedule // that keeps the accumulator in ordinary registers. The product is always // accumulated and returned in single precision. +// +// Time for one multiply on an RTX 5060 Ti, against cublas doing the same +// thing at each precision (cublasSgemm, and cublasGemmEx with half precision +// operands and a single precision accumulator): +// +// 1024 2048 4096 +// Halide f32 311 us 1675 us 18727 us +// cublas f32 147 us 1029 us 7813 us +// Halide f16 53 us 365 us 2784 us +// cublas f16 51 us 347 us 2699 us +// +// So the tensor core schedule is within a few percent of cublas, and the +// float one is about half its speed - the tensor cores are what this app has +// been tuned for. class MatMul : public Halide::Generator { public: GeneratorParam size{"size", 1024}; @@ -85,8 +99,7 @@ class MatMul : public Halide::Generator { } private: - // 311 us for 1024x1024 on an RTX 5060 Ti, where cublas is 146 us at the - // same precision. + // The float schedule. See the table above for how it does. void schedule_cuda() { Var xi, yi, xii, yii; @@ -110,13 +123,9 @@ class MatMul : public Halide::Generator { B.in().compute_at(prod, r).vectorize(_0).unroll(_1); } - // 53 us for 1024x1024 on an RTX 5060 Ti, which is 5.9x the schedule above - // and 2.8x cublas at single precision. - // - // Against cublas doing the same thing - half precision operands into a - // single precision accumulator - this is a little slower at every size: - // 96% of its throughput at 1024, 95% at 2048, and 97% at 4096, with the - // block shapes below picked per size by measurement. + // The tensor core schedule, which reaches 96%, 95% and 97% of cublas at + // the three sizes in the table above. The block shapes below are picked + // per size by measurement. void schedule_tensor_cores() { // The tensor core tile shape, and how many of them each warp // accumulates at once. Each operand tile loaded feeds tiles_x (or From dd758619ac2e5a55c5c807d134716bfc4ac7e99d Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Mon, 3 Aug 2026 11:46:57 -0700 Subject: [PATCH 38/59] Take align_up's argument by const reference align_up is now instantiated with Expr as well as integral types, and clang-tidy's performance-unnecessary-value-param fires on the by-value parameter for types with non-trivial copies. Co-Authored-By: Claude Opus 5 --- src/Util.h | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/Util.h b/src/Util.h index 01d89eb3b6ff..6c04daed65f4 100644 --- a/src/Util.h +++ b/src/Util.h @@ -605,8 +605,10 @@ inline bool is_power_of_two(int64_t x) { return (x & (x - 1)) == 0; } +/** Round x up to the next multiple of n. Works for integral types and for + * Expr. */ template -inline T align_up(T x, int n) { +inline T align_up(const T &x, int n) { return (x + n - 1) / n * n; } From 10fd10d7ca145e7d79fa5051f384e230002744da Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Mon, 3 Aug 2026 13:04:19 -0700 Subject: [PATCH 39/59] Support bfloat and eight-bit integer tensor core multiplies The tensor cores multiply brain floats into single precision and eight-bit integers into 32-bit ones, as well as halves. The pieces that were specific to halves: Fragment sizes were two constants. Only the accumulator one is universal - an accumulator tile holds eight elements per lane whatever its type and shape. Operand fragments hold exactly their share of the matrix, except half precision ones, which are a fixed sixteen elements per lane whatever the shape, with the hardware replicating across lanes when the shape holds fewer. So the operand size is now derived from the shape and type rather than assumed. The element type in the intrinsic names was guessed from the bit width, which would have silently emitted a half precision instruction for brain floats. It is now an explicit table that errors on anything unsupported, and a second table gives the operand and accumulator combinations that have an instruction. The intrinsics take a fragment as 32-bit registers. Halide's half vectors match that signature directly, but brain floats and bytes have to be bitcast to i32. A fragment that is a single register comes back as that register rather than as a struct holding one of them. wmma_matmul covers the new types across every shape, layout, and staging option it already covered halves for. Its operands are now small integers, which every type here represents exactly and whose dot products stay exact in the accumulator, so it checks for equality rather than to a tolerance. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/CMakeLists.txt | 15 ++- apps/cuda_mat_mul/Makefile | 21 +++- apps/cuda_mat_mul/mat_mul_generator.cpp | 88 +++++++++++----- apps/cuda_mat_mul/runner.cpp | 89 +++++++++++++++-- src/CodeGen_PTX_Dev.cpp | 79 ++++++++++++--- src/ExtractWMMAOperations.cpp | 92 +++++++++++------ test/correctness/wmma_matmul.cpp | 127 +++++++++++++++++------- 7 files changed, 397 insertions(+), 114 deletions(-) diff --git a/apps/cuda_mat_mul/CMakeLists.txt b/apps/cuda_mat_mul/CMakeLists.txt index 3de2599826be..3ddbe6169036 100644 --- a/apps/cuda_mat_mul/CMakeLists.txt +++ b/apps/cuda_mat_mul/CMakeLists.txt @@ -32,15 +32,24 @@ add_halide_generator(mat_mul.generator SOURCES mat_mul_generator.cpp) # that accumulates in ordinary registers. add_halide_library(mat_mul FROM mat_mul.generator FEATURES cuda cuda_capability_50 - PARAMS size=1024 A.type=float32 B.type=float32) + PARAMS size=1024 A.type=float32 B.type=float32 out.type=float32) add_halide_library(mat_mul_f16 FROM mat_mul.generator GENERATOR mat_mul FEATURES cuda cuda_capability_80 - PARAMS size=1024 A.type=float16 B.type=float16) + PARAMS size=1024 A.type=float16 B.type=float16 out.type=float32) + +add_halide_library(mat_mul_bf16 FROM mat_mul.generator + GENERATOR mat_mul + FEATURES cuda cuda_capability_80 + PARAMS size=1024 A.type=bfloat16 B.type=bfloat16 out.type=float32) +add_halide_library(mat_mul_u8 FROM mat_mul.generator + GENERATOR mat_mul + FEATURES cuda cuda_capability_80 + PARAMS size=1024 A.type=uint8 B.type=uint8 out.type=int32) # Main executable add_executable(runner runner.cpp) -target_link_libraries(runner PRIVATE mat_mul mat_mul_f16 Halide::Tools CUDA::cudart CUDA::cublas) +target_link_libraries(runner PRIVATE mat_mul mat_mul_f16 mat_mul_bf16 mat_mul_u8 Halide::Tools CUDA::cudart CUDA::cublas) # Test that the app actually works! add_test(NAME mat_mul COMMAND runner) diff --git a/apps/cuda_mat_mul/Makefile b/apps/cuda_mat_mul/Makefile index 7f898815d4fe..9234b95a4f46 100644 --- a/apps/cuda_mat_mul/Makefile +++ b/apps/cuda_mat_mul/Makefile @@ -21,14 +21,29 @@ $(GENERATOR_BIN)/mat_mul.generator: mat_mul_generator.cpp $(GENERATOR_DEPS) $(BIN)/%/mat_mul.a: $(GENERATOR_BIN)/mat_mul.generator @mkdir -p $(@D) $^ -g mat_mul -f mat_mul -e $(GENERATOR_OUTPUTS) -o $(@D) \ - target=$(FLOAT_TARGET) size=$(MATRIX_SIZE) A.type=float32 B.type=float32 + target=$(FLOAT_TARGET) size=$(MATRIX_SIZE) \ + A.type=float32 B.type=float32 out.type=float32 $(BIN)/%/mat_mul_f16.a: $(GENERATOR_BIN)/mat_mul.generator @mkdir -p $(@D) $^ -g mat_mul -f mat_mul_f16 -e $(GENERATOR_OUTPUTS) -o $(@D) \ - target=$(HALF_TARGET) size=$(MATRIX_SIZE) A.type=float16 B.type=float16 + target=$(HALF_TARGET) size=$(MATRIX_SIZE) \ + A.type=float16 B.type=float16 out.type=float32 -$(BIN)/%/runner: runner.cpp $(BIN)/%/mat_mul.a $(BIN)/%/mat_mul_f16.a +$(BIN)/%/mat_mul_bf16.a: $(GENERATOR_BIN)/mat_mul.generator + @mkdir -p $(@D) + $^ -g mat_mul -f mat_mul_bf16 -e $(GENERATOR_OUTPUTS) -o $(@D) \ + target=$(HALF_TARGET) size=$(MATRIX_SIZE) \ + A.type=bfloat16 B.type=bfloat16 out.type=float32 + +$(BIN)/%/mat_mul_u8.a: $(GENERATOR_BIN)/mat_mul.generator + @mkdir -p $(@D) + $^ -g mat_mul -f mat_mul_u8 -e $(GENERATOR_OUTPUTS) -o $(@D) \ + target=$(HALF_TARGET) size=$(MATRIX_SIZE) \ + A.type=uint8 B.type=uint8 out.type=int32 + +$(BIN)/%/runner: runner.cpp $(BIN)/%/mat_mul.a $(BIN)/%/mat_mul_f16.a \ + $(BIN)/%/mat_mul_bf16.a $(BIN)/%/mat_mul_u8.a @mkdir -p $(@D) $(CXX) $(CXXFLAGS) -I$(BIN)/$* -Wall $^ -o $@ $(LDFLAGS) -lcudart -lcublas diff --git a/apps/cuda_mat_mul/mat_mul_generator.cpp b/apps/cuda_mat_mul/mat_mul_generator.cpp index 1d02d7e230d1..c18890a1db99 100644 --- a/apps/cuda_mat_mul/mat_mul_generator.cpp +++ b/apps/cuda_mat_mul/mat_mul_generator.cpp @@ -17,22 +17,28 @@ void set_alignment_and_bounds(OutputImageParam p, int size) { // A square matrix multiply, scheduled two ways. The operands are untyped, so // their type is a generator param, and it is what picks the schedule: half // precision operands get the tensor cores, and anything else gets a schedule -// that keeps the accumulator in ordinary registers. The product is always -// accumulated and returned in single precision. +// that keeps the accumulator in ordinary registers. The tensor cores multiply +// halves, brain floats, or eight-bit integers; the output type is also the +// accumulator type, so it picks between the accumulators an operand type can +// pair with. // // Time for one multiply on an RTX 5060 Ti, against cublas doing the same -// thing at each precision (cublasSgemm, and cublasGemmEx with half precision -// operands and a single precision accumulator): +// thing (cublasSgemm, and cublasGemmEx with half precision operands and a +// single precision accumulator): // -// 1024 2048 4096 -// Halide f32 311 us 1675 us 18727 us -// cublas f32 147 us 1029 us 7813 us -// Halide f16 53 us 365 us 2784 us -// cublas f16 51 us 347 us 2699 us +// 1024 2048 4096 +// Halide f32 310 us 1667 us 18842 us +// cublas f32 147 us 1030 us 7869 us +// Halide f16 53 us 371 us 2797 us +// Halide bf16 53 us 365 us 2796 us +// Halide u8 51 us 247 us 2424 us +// cublas f16 51 us 347 us 2699 us // -// So the tensor core schedule is within a few percent of cublas, and the -// float one is about half its speed - the tensor cores are what this app has -// been tuned for. +// The float schedule is about half of cublas, and the half precision one is +// within a few percent of it. Eight-bit operands are the fastest of the lot - +// they halve the traffic through shared memory, and beat cublas at half +// precision by 40% at 2048 - which makes them the interesting case for +// imaging, where the inputs are usually bytes to begin with. class MatMul : public Halide::Generator { public: GeneratorParam size{"size", 1024}; @@ -48,31 +54,62 @@ class MatMul : public Halide::Generator { // How much of the reduction is staged in shared memory at a time. GeneratorParam block_r{"block_r", 32}; // Extra elements per row of the shared panels, which spreads consecutive - // rows across different banks. A multiple of eight keeps the rows aligned - // enough for the widest asynchronous copy. - GeneratorParam pad_a{"pad_a", 8}; - GeneratorParam pad_b{"pad_b", 8}; + // rows across different banks. Zero means pad by sixteen bytes, which is + // the least that keeps each row aligned both for the widest asynchronous + // copy and for the tensor core loads, whose matrix addresses have to be + // sixteen byte aligned. How many elements that is depends on the operand + // type, which is why it is not just a number here. + GeneratorParam pad_a{"pad_a", 0}; + GeneratorParam pad_b{"pad_b", 0}; Input> A{"A"}; Input> B{"B"}; - Output> out{"out"}; + // The output type is also the accumulator type - there is no tensor core + // store from a half precision fragment into single precision memory - so + // asking for a half precision output is what asks to accumulate in half. + // That halves the registers the accumulator takes, but summing a long + // reduction in half precision loses accuracy badly, so it is only worth + // asking for when the numerics of the problem allow it. + Output> out{"out"}; - // Tensor cores multiply half precision operands into a single precision - // accumulator, so asking for half precision inputs is what asks for them. + // The tensor cores multiply 16-bit floats or 8-bit integers, so asking for + // one of those operand types is what asks for them. Anything else gets the + // schedule that accumulates in ordinary registers. bool use_tensor_cores() const { - return A.type() == Float(16); + Type t = A.type(); + return t == Float(16) || t == BFloat(16) || t == Int(8) || t == UInt(8); + } + + // The accumulator each operand type pairs with. Half precision can also + // accumulate into halves, which is what asking for a half output does. + Type natural_accumulator() const { + Type t = A.type(); + if (t == Int(8) || t == UInt(8)) { + return Int(32); + } + return Float(32); } void generate() { + _halide_user_assert(A.type() == B.type()) + << "The two operands must have the same type, but they are " + << A.type() << " and " << B.type() << ".\n"; + _halide_user_assert(out.type() == natural_accumulator() || + (out.type() == Float(16) && A.type() == Float(16))) + << "A " << A.type() << " matrix multiply accumulates into " + << natural_accumulator() + << (A.type() == Float(16) ? " or float16" : "") + << ", but a " << out.type() << " output was asked for.\n"; r = RDom(0, size, "r"); - prod(x, y) = 0.f; + Type acc = out.type(); + prod(x, y) = cast(acc, 0); // The widening to the accumulator type happens here, at the multiply, // rather than in the operand wrappers the schedule stages through // shared memory. That way half precision operands are staged as half // precision and reach the tensor cores as such. - prod(x, y) += cast(A(x, r)) * cast(B(r, y)); + prod(x, y) += cast(acc, A(x, r)) * cast(acc, B(r, y)); out(x, y) = prod(x, y); } @@ -151,6 +188,9 @@ class MatMul : public Halide::Generator { const int block_x = tile * tx * wx; const int block_y = tile * ty * wy; + const int pa = pad_a ? (int)pad_a : 16 / A.type().bytes(); + const int pb = pad_b ? (int)pad_b : 16 / A.type().bytes(); + Var xi("xi"), yi("yi"), xt("xt"), yt("yt"), mmxi("mmxi"), mmyi("mmyi"); Var xw("xw"), yw("yw"), rxi("rxi"), ryi("ryi"); RVar ro("ro"), ri("ri"), rri("rri"); @@ -213,10 +253,10 @@ class MatMul : public Halide::Generator { Var t("t"), ti("ti"), tw("tw"), tw2("tw2"), to("to"); // B.in() is dense in the reduction dimension, which is its _0. - B.in().compute_at(prod, ro).store_in(MemoryType::GPUSharedAsync).align_storage(_0, (int)block_r + (int)pad_a).split(_0, rro, rrv, vec).fuse(rro, _1, t).split(t, t, ti, 32).split(t, t, tw, wx).split(t, to, tw2, wy).gpu_lanes(ti).gpu_threads(tw, tw2).vectorize(rrv); + B.in().compute_at(prod, ro).store_in(MemoryType::GPUSharedAsync).align_storage(_0, (int)block_r + pa).split(_0, rro, rrv, vec).fuse(rro, _1, t).split(t, t, ti, 32).split(t, t, tw, wx).split(t, to, tw2, wy).gpu_lanes(ti).gpu_threads(tw, tw2).vectorize(rrv); // A.in() is dense in x, which is its _0. - A.in().compute_at(prod, ro).store_in(MemoryType::GPUSharedAsync).align_storage(_0, block_x + pad_b).split(_0, xxo, xxi, vec).fuse(xxo, _1, t).split(t, t, ti, 32).split(t, t, tw, wx).split(t, to, tw2, wy).gpu_lanes(ti).gpu_threads(tw, tw2).vectorize(xxi); + A.in().compute_at(prod, ro).store_in(MemoryType::GPUSharedAsync).align_storage(_0, block_x + pb).split(_0, xxo, xxi, vec).fuse(xxo, _1, t).split(t, t, ti, 32).split(t, t, tw, wx).split(t, to, tw2, wy).gpu_lanes(ti).gpu_threads(tw, tw2).vectorize(xxi); } Var x{"x"}, y{"y"}; diff --git a/apps/cuda_mat_mul/runner.cpp b/apps/cuda_mat_mul/runner.cpp index dcf9474efe1f..3c9d68cee278 100644 --- a/apps/cuda_mat_mul/runner.cpp +++ b/apps/cuda_mat_mul/runner.cpp @@ -3,11 +3,14 @@ #include "halide_benchmark.h" #include #include +#include #include #include #include "mat_mul.h" +#include "mat_mul_bf16.h" #include "mat_mul_f16.h" +#include "mat_mul_u8.h" using Halide::Runtime::Buffer; @@ -18,23 +21,23 @@ namespace { // and once with half operands, which get the tensor cores. Both accumulate in // and return single precision, so the two are directly comparable. -template +template bool check(const Buffer &A, const Buffer &B, - const Buffer &C, int size, const char *name) { + const Buffer &C, int size, const char *name) { // Spot check on strides that are coprime with the tile sizes, so the // samples land at varying offsets within a tile. for (int y = 0; y < size; y += 97) { for (int x = 0; x < size; x += 89) { - float correct = 0.f; + double correct = 0; for (int k = 0; k < size; k++) { - correct += (float)A(x, k) * (float)B(k, y); + correct += (double)A(x, k) * (double)B(k, y); } // The operands are small integers, which are exact in both float // and half, and the accumulator is single precision either way, so // the answer should be exact. - if (C(x, y) != correct) { + if ((double)C(x, y) != correct) { printf("%s: bad result at %d %d: %f != %f\n", - name, x, y, C(x, y), correct); + name, x, y, (double)C(x, y), correct); return false; } } @@ -42,6 +45,22 @@ bool check(const Buffer &A, const Buffer &B, return true; } +// There is no C++ type for bfloat16 here, so the buffer carries the type at +// runtime and these convert. A bfloat is the top half of a float, so for the +// small integers this uses the conversion is exact in both directions. +uint16_t to_bf16(float f) { + uint32_t bits; + memcpy(&bits, &f, 4); + return (uint16_t)(bits >> 16); +} + +float from_bf16(uint16_t h) { + uint32_t bits = (uint32_t)h << 16; + float f; + memcpy(&f, &bits, 4); + return f; +} + double gflops(int size, double seconds) { return 2.0 * size * size * size / seconds * 1e-9; } @@ -138,6 +157,64 @@ int main(int argc, char **argv) { t, gflops(size, t)); } + // The other operand types the tensor cores multiply. Brain floats + // accumulate into single precision like halves do, and eight-bit integers + // into 32-bit ones, which is the interesting case for imaging. + if (ver >= 80) { + { + const halide_type_t bf16(halide_type_bfloat, 16); + Buffer A(bf16, size, size), B(bf16, size, size); + Buffer C(size, size); + // The buffer carries its type at runtime, so index the raw + // storage rather than going through a typed view. + uint16_t *Ap = (uint16_t *)A.data(), *Bp = (uint16_t *)B.data(); + auto Af = [&](int i, int j) { return from_bf16(Ap[j * size + i]); }; + auto Bf = [&](int i, int j) { return from_bf16(Bp[j * size + i]); }; + for (int i = 0; i < size * size; i++) { + Ap[i] = to_bf16((float)((rand() & 3) - 1)); + Bp[i] = to_bf16((float)((rand() & 3) - 1)); + } + A.set_host_dirty(); + B.set_host_dirty(); + mat_mul_bf16(A, B, C); + C.copy_to_host(); + for (int y = 0; y < size; y += 97) { + for (int x = 0; x < size; x += 89) { + double correct = 0; + for (int k = 0; k < size; k++) { + correct += (double)Af(x, k) * (double)Bf(k, y); + } + if ((double)C(x, y) != correct) { + printf("bfloat: bad result at %d %d: %f != %f\n", + x, y, (double)C(x, y), correct); + return 1; + } + } + } + double t = bench_batched([&]() { mat_mul_bf16(A, B, C); }, + [&]() { C.device_sync(); }); + printf("Halide bfloat (tensor cores): %f s (%.1f GFlop/s)\n", + t, gflops(size, t)); + } + { + Buffer A(size, size), B(size, size); + Buffer C(size, size); + A.for_each_value([](uint8_t &v) { v = (uint8_t)(rand() & 3); }); + B.for_each_value([](uint8_t &v) { v = (uint8_t)(rand() & 3); }); + A.set_host_dirty(); + B.set_host_dirty(); + mat_mul_u8(A, B, C); + C.copy_to_host(); + if (!check(A, B, C, size, "uint8")) { + return 1; + } + double t = bench_batched([&]() { mat_mul_u8(A, B, C); }, + [&]() { C.device_sync(); }); + printf("Halide uint8 (tensor cores): %f s (%.1f GFlop/s)\n", + t, gflops(size, t)); + } + } + // Benchmark cublas for reference, at both precisions. The half precision // one accumulates in single precision, matching what the Halide pipeline // does, so the two are comparable. diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index db2e1d3da91c..82aa3c993324 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -446,6 +446,14 @@ WMMAMatrixLayout matrix_in_memory(const string &name, const MultiRamp &mr, int r } // namespace +// The nvvm intrinsics take and return a fragment as 32-bit registers. Halide +// represents half precision as llvm's half, which is what the intrinsics for +// it take directly, but bfloat and the eight-bit integers have no llvm vector +// type in the signature - those intrinsics take the lanes packed into an i32. +bool wmma_reg_is_packed_i32(Type t) { + return t.bits() < 32 && t.element_of() != Float(16); +} + void CodeGen_PTX_Dev::split_fragment(const Expr &e, vector &args) { // One llvm value per 32-bit register. const int num_regs = e.type().bits() * e.type().lanes() / 32; @@ -462,9 +470,13 @@ void CodeGen_PTX_Dev::split_fragment(const Expr &e, vector &args) { Value *v = codegen(e); const int lanes_per_reg = 32 / e.type().bits(); for (int i = 0; i < e.type().lanes() / lanes_per_reg; i++) { - args.push_back(lanes_per_reg == 1 ? - builder->CreateExtractElement(v, i) : - slice_vector(v, i * lanes_per_reg, lanes_per_reg)); + Value *reg = lanes_per_reg == 1 ? + builder->CreateExtractElement(v, i) : + slice_vector(v, i * lanes_per_reg, lanes_per_reg); + if (wmma_reg_is_packed_i32(e.type())) { + reg = builder->CreateBitCast(reg, i32_t); + } + args.push_back(reg); } } @@ -484,16 +496,29 @@ void CodeGen_PTX_Dev::codegen_wmma(const Call *op) { // Reassemble the returned struct into a Halide vector. llvm::Type *result_type = llvm_type_of(op->type); const int num_regs = op->type.bits() * op->type.lanes() / 32; + // A fragment that is a single register comes back as that register rather + // than as a struct holding one of them. + auto get_reg = [&](int i) { + return result->getType()->isStructTy() ? + builder->CreateExtractValue(result, i) : + result; + }; if (op->type.bits() == 32) { value = UndefValue::get(result_type); for (int i = 0; i < num_regs; i++) { - value = builder->CreateInsertElement(value, builder->CreateExtractValue(result, i), i); + value = builder->CreateInsertElement(value, get_reg(i), i); } } else { vector regs; regs.reserve(num_regs); for (int i = 0; i < num_regs; i++) { - regs.push_back(builder->CreateExtractValue(result, i)); + Value *reg = get_reg(i); + if (wmma_reg_is_packed_i32(op->type)) { + reg = builder->CreateBitCast( + reg, get_vector_type(llvm_type_of(op->type.element_of()), + 32 / op->type.bits())); + } + regs.push_back(reg); } value = concat_vectors(regs); } @@ -546,6 +571,28 @@ void CodeGen_PTX_Dev::codegen_fragment_store(const Store *op) { } } +// The element type the wmma intrinsic names use for a Halide type. The +// hardware takes 16-bit floats or 8-bit integers as multiplicands, and +// accumulates the first into 16 or 32-bit floats and the second into 32-bit +// integers. +std::string wmma_type_suffix(Type t) { + if (t == Float(16)) { + return "f16"; + } else if (t == BFloat(16)) { + return "bf16"; + } else if (t == Float(32)) { + return "f32"; + } else if (t == Int(8)) { + return "s8"; + } else if (t == UInt(8)) { + return "u8"; + } else if (t == Int(32)) { + return "s32"; + } + user_error << "There is no tensor core instruction for " << t << ".\n"; + return ""; +} + llvm::Value *CodeGen_PTX_Dev::codegen_wmma_raw(const Call *op) { // The nvvm wmma intrinsics take and return fragments as a flat list of // 32-bit registers, packaged up as a literal struct. We represent them in @@ -566,11 +613,19 @@ llvm::Value *CodeGen_PTX_Dev::codegen_wmma_raw(const Call *op) { vector overloads; if (op->is_intrinsic(Call::wmma_mma)) { - // The two type suffixes are the types of the d and c operands, which - // for us are always the same. - const char *suffix = op->type.bits() == 32 ? "f32" : "f16"; + // Half precision multiplicands name the instruction after the d and c + // operands, which for us are always the same type. Everything else + // names it after the multiplicands, which are in args 5 and 6. + const Type operand_type = op->args[5].type().element_of(); + std::string signature; + if (operand_type == Float(16)) { + const std::string acc = wmma_type_suffix(op->type.element_of()); + signature = acc + "." + acc; + } else { + signature = wmma_type_suffix(operand_type); + } name << "mma." << layouts[get_int_arg(3)] << "." << layouts[get_int_arg(4)] - << "." << suffix << "." << suffix; + << "." << signature; split_fragment(op->args[5], args); split_fragment(op->args[6], args); split_fragment(op->args[7], args); @@ -590,9 +645,7 @@ llvm::Value *CodeGen_PTX_Dev::codegen_wmma_raw(const Call *op) { << "load with an affine index by the time it reaches the backend.\n"; WMMAMatrixLayout mem = matrix_in_memory(matrix->name, mr, is_b ? K : M, is_a ? K : N, arg); - // The a and b operands are always 16-bit; an accumulator may be either. - const char *type_suffix = - is_a || is_b ? "f16" : (op->type.bits() == 32 ? "f32" : "f16"); + const std::string type_suffix = wmma_type_suffix(op->type.element_of()); name << "load." << (is_a ? "a" : is_b ? "b" : "c") << "." << (mem.row_major ? "row" : "col") << ".stride." << type_suffix; @@ -640,7 +693,7 @@ void CodeGen_PTX_Dev::codegen_wmma_store(const Store *op) { std::ostringstream name; name << "llvm.nvvm.wmma.m" << M << "n" << N << "k" << K << ".store.d." << (mem.row_major ? "row" : "col") << ".stride." - << (fragment.type().bits() == 32 ? "f32" : "f16"); + << wmma_type_suffix(fragment.type().element_of()); Value *ptr = codegen_buffer_pointer(op->name, op->value.type().element_of(), mem.base); vector args{ptr}; diff --git a/src/ExtractWMMAOperations.cpp b/src/ExtractWMMAOperations.cpp index 67bfcd35f532..80bd0d044595 100644 --- a/src/ExtractWMMAOperations.cpp +++ b/src/ExtractWMMAOperations.cpp @@ -94,16 +94,50 @@ const char *role_name(Role role) { } } -// Every warp is 32 lanes. An accumulator tile holds 256 elements, so each lane -// holds 8 of them. The hardware hands back operand fragments as 8 32-bit -// registers per lane whatever the shape, replicating elements across lanes for -// the shapes that hold fewer, so those are 16 16-bit elements per lane. +// Every warp is 32 lanes, and an accumulator tile holds 256 elements whatever +// its shape, so each lane holds 8 of them whatever their type. constexpr int warp_lanes = 32; constexpr int accumulator_elements = 8; -constexpr int operand_elements = 16; +// Half precision operand fragments are eight 32-bit registers per lane +// whatever the shape, which is sixteen elements. For the shapes that hold +// fewer than that the hardware replicates elements across lanes. +constexpr int half_operand_elements = 16; + +// The multiplicand and accumulator type combinations the tensor cores have an +// instruction for. Half precision multiplicands accumulate into halves or +// floats, brain floats accumulate into floats, and eight-bit integers +// accumulate into 32-bit ones. The multiplicands must match each other: at +// this shape there is no instruction for mixed signedness. +bool wmma_types_supported(Type operand, Type accumulator) { + if (operand == Float(16)) { + return accumulator == Float(16) || accumulator == Float(32); + } else if (operand == BFloat(16)) { + return accumulator == Float(32); + } else if (operand == Int(8) || operand == UInt(8)) { + return accumulator == Int(32); + } + return false; +} + +// The shape of the matrix each fragment is taken out of. +void fragment_matrix_shape(Role role, const Shape &shape, int *rows, int *cols) { + *rows = role == Role::B ? shape.K : shape.M; + *cols = role == Role::A ? shape.K : shape.N; +} -int elements_per_lane(Role role) { - return role == Role::Accumulator ? accumulator_elements : operand_elements; +// How many elements of a fragment each lane holds. Every type but half +// precision holds exactly its share of the matrix, so this follows from the +// shape; half precision is the fixed size above however small the shape is. +int elements_per_lane(Role role, const Shape &shape, Type t) { + if (role == Role::Accumulator) { + return accumulator_elements; + } + if (t == Float(16)) { + return half_operand_elements; + } + int rows, cols; + fragment_matrix_shape(role, shape, &rows, &cols); + return rows * cols / warp_lanes; } // One operand of a matrix multiply, described in the canonical [K, N, M] @@ -184,19 +218,14 @@ Expr make_matrix_address(const string &name, Type element_type, const Expr &base const_true(lanes), ModulusRemainder()); } -// The shape of the matrix each fragment is taken out of. -void fragment_matrix_shape(Role role, const Shape &shape, int *rows, int *cols) { - *rows = role == Role::B ? shape.K : shape.M; - *cols = role == Role::A ? shape.K : shape.N; -} - Expr make_matrix_to_fragment(Role role, const Shape &shape, Layout layout, const Load *load, const Expr &base, const Expr &stride) { int rows, cols; fragment_matrix_shape(role, shape, &rows, &cols); Expr address = make_matrix_address(load->name, load->type.element_of(), base, rows, cols, layout, stride, load->image, load->param); - Type type = load->type.element_of().with_lanes(elements_per_lane(role)); + Type type = load->type.element_of().with_lanes( + elements_per_lane(role, shape, load->type.element_of())); return Call::make(type, intrinsic_for_role(role), {shape.M, shape.N, shape.K, std::move(address)}, Call::Intrinsic); @@ -264,11 +293,12 @@ MatmulInfo analyze_matmul(const Store *op) { return fail("the load or store is predicated"); } - if (!reduce->type.is_float() || - !(reduce->type.bits() == 32 || reduce->type.bits() == 16)) { - return fail("the accumulator type is not 32-bit or 16-bit float"); - } info.accumulator_type = reduce->type.element_of(); + if (!(info.accumulator_type == Float(32) || + info.accumulator_type == Float(16) || + info.accumulator_type == Int(32))) { + return fail("the accumulator type is not float32, float16, or int32"); + } // The vector reduce must be of a widening multiply. FindIntrinsics does // not lift float widening muls, so we just expect a multiply of two casts. @@ -288,9 +318,13 @@ MatmulInfo analyze_matmul(const Store *op) { if (!is_const_one(info.lhs.load->predicate) || !is_const_one(info.rhs.load->predicate)) { return fail("the matrix multiply operands are predicated loads"); } - if (info.lhs.load->type.element_of() != Float(16) || - info.rhs.load->type.element_of() != Float(16)) { - return fail("the matrix multiply operands are not both float16"); + const Type operand_type = info.lhs.load->type.element_of(); + if (info.rhs.load->type.element_of() != operand_type) { + return fail("the matrix multiply operands do not have the same type"); + } + if (!wmma_types_supported(operand_type, info.accumulator_type)) { + return fail("there is no tensor core instruction that multiplies these " + "operands into an accumulator of this type"); } // In a matrix multiply with row-major inputs and outputs, the algorithm @@ -408,7 +442,8 @@ struct Fragment { vector subtiles; Type value_type() const { - return element_type.with_lanes(elements_per_lane(role)); + return element_type.with_lanes( + elements_per_lane(role, shape, element_type)); } }; @@ -508,7 +543,7 @@ class ExtractWMMAOperations : public IRMutator { Expr operand_value(const Operand &operand, Role role, const Shape &shape, Layout layout, const Expr &stride) { if (Fragment *f = find_fragment(operand.load->name)) { - const int lanes = elements_per_lane(role); + const int lanes = f->value_type().lanes(); const string name = operand_subtile_name(f, operand.mr.base, role, shape, layout, stride); return Load::make(f->value_type(), name, Ramp::make(0, 1, lanes), {}, {}, @@ -532,7 +567,7 @@ class ExtractWMMAOperations : public IRMutator { const string name = subtile_name( f, make_matrix_index(dest.base, rows, cols, dest.row_major ? Layout::Row : Layout::Col, dest.stride)); - const int lanes = elements_per_lane(f->role); + const int lanes = f->value_type().lanes(); Expr value; if (is_const_zero(op->value)) { // Zeroing a fragment is layout-independent, so it doesn't need an @@ -616,9 +651,10 @@ class ExtractWMMAOperations : public IRMutator { return IRMutator::visit(op); } - user_assert(op->type == Float(32) || op->type == Float(16)) - << "Tensor core fragments must hold 32-bit or 16-bit floats, but " - << op->name << " holds " << op->type << ".\n"; + user_assert(op->type == Float(32) || op->type == Float(16) || + op->type == Int(32)) + << "Tensor core fragments must hold 32-bit or 16-bit floats, or " + << "32-bit integers, but " << op->name << " holds " << op->type << ".\n"; Fragment &f = fragments[op->name]; if (pass == 0) { @@ -645,7 +681,7 @@ class ExtractWMMAOperations : public IRMutator { // get replicated per thread. for (int i = 0; i < (int)f.subtiles.size(); i++) { body = Allocate::make(f.fragment_name + std::to_string(i), f.element_type, - MemoryType::WMMAFragment, {elements_per_lane(f.role)}, + MemoryType::WMMAFragment, {f.value_type().lanes()}, const_true(), body); } return body; diff --git a/test/correctness/wmma_matmul.cpp b/test/correctness/wmma_matmul.cpp index 821a4335f97f..56cf83664c88 100644 --- a/test/correctness/wmma_matmul.cpp +++ b/test/correctness/wmma_matmul.cpp @@ -6,6 +6,10 @@ using namespace Halide; namespace { struct Params { + // The type of the two matrices being multiplied. The accumulator type + // follows from it: the integer instructions accumulate into int32, and the + // float ones into float32 unless half_accumulator asks otherwise. + Type operand = Float(16); // The size of the matrices. int M = 64, N = 64, K = 64; // The tensor core tile shape to use. @@ -17,6 +21,7 @@ struct Params { // Whether each input matrix is stored with its rows or its columns dense. bool a_transposed = false, b_transposed = false; // Whether to accumulate in half precision instead of single precision. + // Only available with half precision operands. bool half_accumulator = false; // Whether to stage the operand tiles through shared memory inside the // reduction loop. @@ -34,23 +39,69 @@ std::ostream &operator<<(std::ostream &s, const Params &p) { << " tiles, " << p.tiles_m << "x" << p.tiles_n << " tiles per warp, " << p.warps << " warps, a" << (p.a_transposed ? "T" : "") << " b" << (p.b_transposed ? "T" : "") + << " of " << p.operand << (p.half_accumulator ? ", f16 accumulator" : "") << (p.stage_in_shared ? ", staged through shared" : "") << (p.init_from_memory ? ", accumulator initialized from memory" : "") << (p.out_transposed ? ", transposed output" : ""); } -void fill(Buffer &buf) { - buf.fill([]() { - return float16_t(((float)rand() / RAND_MAX) - 0.5f); - }); +// Small non-negative integers, which every type here represents exactly, and +// which keep the dot products small enough to be exact in the accumulator too. +// That lets the results be checked for equality rather than to a tolerance. +void fill(Buffer<> buf) { + Type t = buf.type(); + for (int y = 0; y < buf.height(); y++) { + for (int x = 0; x < buf.width(); x++) { + int v = rand() & 3; + if (t == Float(16)) { + buf.as()(x, y) = float16_t((float)v); + } else if (t == BFloat(16)) { + buf.as()(x, y) = bfloat16_t((float)v); + } else if (t == Int(8)) { + buf.as()(x, y) = (int8_t)v; + } else if (t == UInt(8)) { + buf.as()(x, y) = (uint8_t)v; + } else { + assert(false && "unhandled operand type"); + } + } + } +} + +// Read an element of a buffer of any of the types this test uses. +double read(const Buffer<> &buf, int x, int y) { + Type t = buf.type(); + if (t == Float(16)) { + return (double)Buffer(buf)(x, y); + } else if (t == BFloat(16)) { + return (double)Buffer(buf)(x, y); + } else if (t == Float(32)) { + return Buffer(buf)(x, y); + } else if (t == Int(8)) { + return Buffer(buf)(x, y); + } else if (t == UInt(8)) { + return Buffer(buf)(x, y); + } else if (t == Int(32)) { + return Buffer(buf)(x, y); + } + assert(false && "unhandled buffer type"); + return 0; +} + +// The accumulator the hardware pairs with a given operand type. +Type accumulator_for(const Params &p) { + if (p.operand.is_int() || p.operand.is_uint()) { + return Int(32); + } + return p.half_accumulator ? Float(16) : Float(32); } bool test(const Params &p) { // Halide indexes matrices with the dense dimension first, so A(k, y) is a // row-major M x K matrix, and A(y, k) is a column-major one. - Buffer A(p.a_transposed ? p.M : p.K, p.a_transposed ? p.K : p.M); - Buffer B(p.b_transposed ? p.K : p.N, p.b_transposed ? p.N : p.K); + Buffer<> A(p.operand, p.a_transposed ? p.M : p.K, p.a_transposed ? p.K : p.M); + Buffer<> B(p.operand, p.b_transposed ? p.K : p.N, p.b_transposed ? p.N : p.K); fill(A); fill(B); @@ -67,15 +118,10 @@ bool test(const Params &p) { // float16, so the output has to be float16 too. // The accumulator either starts at zero, or at a matrix that already exists // in memory, which the hardware can load straight into the fragments. - Type acc_type = p.half_accumulator ? Float(16) : Float(32); - init(x, y) = cast(acc_type, (x * 3 + y) % 7) * cast(acc_type, 0.25f); - if (p.half_accumulator) { - prod(x, y) = p.init_from_memory ? init(x, y) : cast(0.f); - prod(x, y) += a * b; - } else { - prod(x, y) = p.init_from_memory ? init(x, y) : Expr(0.f); - prod(x, y) += cast(a) * cast(b); - } + Type acc_type = accumulator_for(p); + init(x, y) = cast(acc_type, (x * 3 + y) % 7); + prod(x, y) = p.init_from_memory ? init(x, y) : cast(acc_type, 0); + prod(x, y) += cast(acc_type, a) * cast(acc_type, b); out(x, y) = prod(x, y); Var xi("xi"), yi("yi"), xt("xt"), mmxi("mmxi"), mmyi("mmyi"); @@ -150,35 +196,26 @@ bool test(const Params &p) { // A transposed output is a view of a buffer with the dimensions swapped, so // its columns are dense in memory instead of its rows. - Buffer result_storage(p.out_transposed ? p.M : p.N, - p.out_transposed ? p.N : p.M); - Buffer result = + Buffer<> result_storage(acc_type, p.out_transposed ? p.M : p.N, + p.out_transposed ? p.N : p.M); + Buffer<> result = p.out_transposed ? result_storage.transposed(0, 1) : result_storage; - Buffer result_half(p.N, p.M); - auto get = [&](int i, int j) { - return p.half_accumulator ? (float)result_half(i, j) : result(i, j); - }; - if (p.half_accumulator) { - out.realize(result_half); - result_half.copy_to_host(); - } else { - out.realize(result); - result.copy_to_host(); - } + out.realize(result); + result.copy_to_host(); for (int j = 0; j < p.M; j++) { for (int i = 0; i < p.N; i++) { - float ref = p.init_from_memory ? (float)(((i * 3 + j) % 7) * 0.25f) : 0.f; + double ref = p.init_from_memory ? (double)((i * 3 + j) % 7) : 0.0; for (int l = 0; l < p.K; l++) { - ref += (float)(p.a_transposed ? A(j, l) : A(l, j)) * - (float)(p.b_transposed ? B(l, i) : B(i, l)); + ref += read(A, p.a_transposed ? j : l, p.a_transposed ? l : j) * + read(B, p.b_transposed ? l : i, p.b_transposed ? i : l); } - // The accumulation happens in a different order on the GPU, and - // the inputs are half-precision, so allow some slack. - float tolerance = p.half_accumulator ? 5e-2f : 1e-2f; - if (std::abs(get(i, j) - ref) > tolerance * std::max(1.f, std::abs(ref))) { + // The operands are small integers and the dot products stay small + // enough to be exact in every accumulator type here, so the answer + // should be exact however the GPU orders the accumulation. + if (read(result, i, j) != ref) { std::cerr << "Mismatch at " << i << ", " << j << ": " - << get(i, j) << " != " << ref << "\n" + << read(result, i, j) << " != " << ref << "\n" << "For matmul of " << p << "\n"; return false; } @@ -522,6 +559,22 @@ int main(int argc, char **argv) { params.push_back({.half_accumulator = true}); params.push_back({.tiles_m = 2, .tiles_n = 2, .half_accumulator = true}); + // The other operand types the hardware multiplies. Brain floats accumulate + // into single precision, and eight-bit integers into 32-bit ones, which is + // the interesting case for imaging. + for (Type t : {BFloat(16), Int(8), UInt(8)}) { + params.push_back({.operand = t}); + params.push_back({.operand = t, .tile_m = 32, .tile_n = 8}); + params.push_back({.operand = t, .tile_m = 8, .tile_n = 32}); + params.push_back({.operand = t, .tiles_m = 2, .tiles_n = 2}); + params.push_back({.operand = t, .warps = 2}); + params.push_back({.operand = t, .a_transposed = true}); + params.push_back({.operand = t, .b_transposed = true}); + params.push_back({.operand = t, .stage_in_shared = true}); + params.push_back({.operand = t, .init_from_memory = true}); + params.push_back({.operand = t, .out_transposed = true}); + } + // Operand tiles staged through shared memory inside the reduction loop. // This only works if the loop over lanes wraps the individual wmma // statements rather than the whole accumulator allocation. From 8573b89bdbf5cc6057adf3f7d7c45c76b4f02a72 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Mon, 3 Aug 2026 15:41:04 -0700 Subject: [PATCH 40/59] Retune the block shapes for every operand type The shapes were tuned for half precision operands accumulated into single precision, and the other types were borrowing them. Sweeping sixty shapes at each size and operand type gives a shape per pair, and the operand type turns out to matter as much as the size. Bytes want a tall block spread over four warps, where the 16-bit types want a wide one over fewer: byte operands make the loads cheap enough to pay for a much larger accumulator. That is worth 45% at 1024 and 43% at 4096 over the shapes they were borrowing. Brain floats pick out exactly the same shapes as halves at every size, so they share a row. Also builds the half-into-half variant, which the app could not previously ask for. Its operands are sparse zeros and ones so that the dot products stay under 2048, the largest integer half precision represents exactly, which keeps the check exact. The comment now records all seven measured configurations at all three sizes rather than a couple of ratios. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/CMakeLists.txt | 6 +- apps/cuda_mat_mul/Makefile | 9 ++- apps/cuda_mat_mul/mat_mul_generator.cpp | 80 +++++++++++++++++-------- apps/cuda_mat_mul/runner.cpp | 23 +++++++ 4 files changed, 91 insertions(+), 27 deletions(-) diff --git a/apps/cuda_mat_mul/CMakeLists.txt b/apps/cuda_mat_mul/CMakeLists.txt index 3ddbe6169036..23f2f56edc77 100644 --- a/apps/cuda_mat_mul/CMakeLists.txt +++ b/apps/cuda_mat_mul/CMakeLists.txt @@ -38,6 +38,10 @@ add_halide_library(mat_mul_f16 FROM mat_mul.generator FEATURES cuda cuda_capability_80 PARAMS size=1024 A.type=float16 B.type=float16 out.type=float32) +add_halide_library(mat_mul_f16_acc16 FROM mat_mul.generator + GENERATOR mat_mul + FEATURES cuda cuda_capability_80 + PARAMS size=1024 A.type=float16 B.type=float16 out.type=float16) add_halide_library(mat_mul_bf16 FROM mat_mul.generator GENERATOR mat_mul FEATURES cuda cuda_capability_80 @@ -49,7 +53,7 @@ add_halide_library(mat_mul_u8 FROM mat_mul.generator # Main executable add_executable(runner runner.cpp) -target_link_libraries(runner PRIVATE mat_mul mat_mul_f16 mat_mul_bf16 mat_mul_u8 Halide::Tools CUDA::cudart CUDA::cublas) +target_link_libraries(runner PRIVATE mat_mul mat_mul_f16 mat_mul_f16_acc16 mat_mul_bf16 mat_mul_u8 Halide::Tools CUDA::cudart CUDA::cublas) # Test that the app actually works! add_test(NAME mat_mul COMMAND runner) diff --git a/apps/cuda_mat_mul/Makefile b/apps/cuda_mat_mul/Makefile index 9234b95a4f46..d94b24d5133c 100644 --- a/apps/cuda_mat_mul/Makefile +++ b/apps/cuda_mat_mul/Makefile @@ -30,6 +30,12 @@ $(BIN)/%/mat_mul_f16.a: $(GENERATOR_BIN)/mat_mul.generator target=$(HALF_TARGET) size=$(MATRIX_SIZE) \ A.type=float16 B.type=float16 out.type=float32 +$(BIN)/%/mat_mul_f16_acc16.a: $(GENERATOR_BIN)/mat_mul.generator + @mkdir -p $(@D) + $^ -g mat_mul -f mat_mul_f16_acc16 -e $(GENERATOR_OUTPUTS) -o $(@D) \ + target=$(HALF_TARGET) size=$(MATRIX_SIZE) \ + A.type=float16 B.type=float16 out.type=float16 + $(BIN)/%/mat_mul_bf16.a: $(GENERATOR_BIN)/mat_mul.generator @mkdir -p $(@D) $^ -g mat_mul -f mat_mul_bf16 -e $(GENERATOR_OUTPUTS) -o $(@D) \ @@ -43,7 +49,8 @@ $(BIN)/%/mat_mul_u8.a: $(GENERATOR_BIN)/mat_mul.generator A.type=uint8 B.type=uint8 out.type=int32 $(BIN)/%/runner: runner.cpp $(BIN)/%/mat_mul.a $(BIN)/%/mat_mul_f16.a \ - $(BIN)/%/mat_mul_bf16.a $(BIN)/%/mat_mul_u8.a + $(BIN)/%/mat_mul_bf16.a $(BIN)/%/mat_mul_u8.a \ + $(BIN)/%/mat_mul_f16_acc16.a @mkdir -p $(@D) $(CXX) $(CXXFLAGS) -I$(BIN)/$* -Wall $^ -o $@ $(LDFLAGS) -lcudart -lcublas diff --git a/apps/cuda_mat_mul/mat_mul_generator.cpp b/apps/cuda_mat_mul/mat_mul_generator.cpp index c18890a1db99..49ffe7b7458e 100644 --- a/apps/cuda_mat_mul/mat_mul_generator.cpp +++ b/apps/cuda_mat_mul/mat_mul_generator.cpp @@ -24,21 +24,26 @@ void set_alignment_and_bounds(OutputImageParam p, int size) { // // Time for one multiply on an RTX 5060 Ti, against cublas doing the same // thing (cublasSgemm, and cublasGemmEx with half precision operands and a -// single precision accumulator): +// single precision accumulator). The block shapes below were picked by +// sweeping sixty of them at each size and operand type. // -// 1024 2048 4096 -// Halide f32 310 us 1667 us 18842 us -// cublas f32 147 us 1030 us 7869 us -// Halide f16 53 us 371 us 2797 us -// Halide bf16 53 us 365 us 2796 us -// Halide u8 51 us 247 us 2424 us -// cublas f16 51 us 347 us 2699 us +// 1024 2048 4096 +// Halide f32 313 us 1679 us 18796 us +// cublas f32 148 us 1034 us 7938 us +// Halide f16 -> f32 53 us 368 us 2811 us +// Halide bf16 -> f32 54 us 367 us 2810 us +// cublas f16 -> f32 52 us 350 us 2724 us +// Halide f16 -> f16 36 us 199 us 1592 us +// Halide u8 -> i32 35 us 220 us 1697 us +// +// The float schedule is about half of cublas. Against cublas doing the same +// thing, half and brain float land within a few percent of it. The two rows +// below that are doing less work per multiply, so they are not comparable to +// the ones above: a half accumulator halves the registers it takes, and +// eight-bit operands halve the traffic through shared memory. Both are around +// 1.7x the single precision accumulator here, which makes bytes the +// interesting case for imaging, where the inputs are usually bytes anyway. // -// The float schedule is about half of cublas, and the half precision one is -// within a few percent of it. Eight-bit operands are the fastest of the lot - -// they halve the traffic through shared memory, and beat cublas at half -// precision by 40% at 2048 - which makes them the interesting case for -// imaging, where the inputs are usually bytes to begin with. class MatMul : public Halide::Generator { public: GeneratorParam size{"size", 1024}; @@ -170,19 +175,44 @@ class MatMul : public Halide::Generator { const int tile = 16; int tx = tiles_x, ty = tiles_y, wx = warps_x, wy = warps_y; if (tx == 0 || ty == 0 || wx == 0 || wy == 0) { - // Measured on an RTX 5060 Ti. These do not follow a trend worth - // extrapolating from, so they are three measured points rather - // than a formula, and each is 4% to 11% better at its own size - // than either of the others would be. What changes is how much - // accumulator a warp holds: at 4096 a small one spread over four - // warps beats a large one, because the occupancy hides the memory - // latency better than a large accumulator's reuse does. - if ((int)size <= 1024) { - tx = 8, ty = 2, wx = 1, wy = 1; - } else if ((int)size <= 2048) { - tx = 5, ty = 4, wx = 2, wy = 2; + // Measured on an RTX 5060 Ti by sweeping sixty shapes at each size + // and operand type. These do not follow a trend worth + // extrapolating from, so they are measured points rather than a + // formula. What moves between them is how much accumulator a warp + // holds and how many warps share a staged panel, and the operand + // type matters as much as the size: bytes make the operand loads + // cheap enough to pay for a much larger accumulator, so they want + // a tall block spread over four warps, where the 16-bit types want + // a wide one over fewer. + // + // Brain floats pick out exactly the same shapes as halves at + // every size, so they share a row here. + const bool bytes = A.type().bits() == 8; + const bool half_accumulator = out.type() == Float(16); + if (bytes) { + if ((int)size <= 1024) { + tx = 2, ty = 8, wx = 4, wy = 1; + } else if ((int)size <= 2048) { + tx = 2, ty = 10, wx = 4, wy = 1; + } else { + tx = 2, ty = 8, wx = 4, wy = 1; + } + } else if (half_accumulator) { + if ((int)size <= 1024) { + tx = 4, ty = 4, wx = 2, wy = 1; + } else if ((int)size <= 2048) { + tx = 5, ty = 4, wx = 2, wy = 2; + } else { + tx = 4, ty = 4, wx = 2, wy = 2; + } } else { - tx = 4, ty = 2, wx = 2, wy = 2; + if ((int)size <= 1024) { + tx = 8, ty = 2, wx = 1, wy = 1; + } else if ((int)size <= 2048) { + tx = 5, ty = 4, wx = 1, wy = 2; + } else { + tx = 4, ty = 2, wx = 2, wy = 2; + } } } const int block_x = tile * tx * wx; diff --git a/apps/cuda_mat_mul/runner.cpp b/apps/cuda_mat_mul/runner.cpp index 3c9d68cee278..93e793232dd4 100644 --- a/apps/cuda_mat_mul/runner.cpp +++ b/apps/cuda_mat_mul/runner.cpp @@ -10,6 +10,7 @@ #include "mat_mul.h" #include "mat_mul_bf16.h" #include "mat_mul_f16.h" +#include "mat_mul_f16_acc16.h" #include "mat_mul_u8.h" using Halide::Runtime::Buffer; @@ -157,6 +158,28 @@ int main(int argc, char **argv) { t, gflops(size, t)); } + // Half precision operands accumulated into half precision, which halves + // the registers the accumulator needs. The operands here are sparse zeros + // and ones, so the dot products stay small enough to be exact even in a + // half accumulator, whose integers run out at 2048. + if (ver >= 70) { + Buffer<_Float16, 2> A(size, size), B(size, size); + Buffer<_Float16, 2> C(size, size); + A.for_each_value([](_Float16 &v) { v = (_Float16)((rand() & 3) == 0); }); + B.for_each_value([](_Float16 &v) { v = (_Float16)((rand() & 3) == 0); }); + A.set_host_dirty(); + B.set_host_dirty(); + mat_mul_f16_acc16(A, B, C); + C.copy_to_host(); + if (!check(A, B, C, size, "half into half")) { + return 1; + } + double t = bench_batched([&]() { mat_mul_f16_acc16(A, B, C); }, + [&]() { C.device_sync(); }); + printf("Halide half into half (tensor cores): %f s (%.1f GFlop/s)\n", + t, gflops(size, t)); + } + // The other operand types the tensor cores multiply. Brain floats // accumulate into single precision like halves do, and eight-bit integers // into 32-bit ones, which is the interesting case for imaging. From c2d8a11c185624a6357c43b2fc8c54887b917fa6 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Mon, 3 Aug 2026 15:44:09 -0700 Subject: [PATCH 41/59] Benchmark cublas at the other operand and accumulator types cublas does brain floats, eight-bit integers, and a half accumulator too, so the table can compare like with like at every pair of types rather than leaving the new ones without a reference. Brain floats land where halves do, within a few percent of cublas, and the half accumulator is ahead of it at 2048. Eight bit is the weak one: 55% of cublas at 1024 and 62% at 4096. cublas is not using the wmma instructions there - eight-bit operands have an mma shape with twice the reduction depth per instruction, which this schedule cannot reach. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/mat_mul_generator.cpp | 43 +++++++++++++++---------- apps/cuda_mat_mul/runner.cpp | 37 +++++++++++++++++++++ 2 files changed, 63 insertions(+), 17 deletions(-) diff --git a/apps/cuda_mat_mul/mat_mul_generator.cpp b/apps/cuda_mat_mul/mat_mul_generator.cpp index 49ffe7b7458e..d4622b392972 100644 --- a/apps/cuda_mat_mul/mat_mul_generator.cpp +++ b/apps/cuda_mat_mul/mat_mul_generator.cpp @@ -23,26 +23,35 @@ void set_alignment_and_bounds(OutputImageParam p, int size) { // pair with. // // Time for one multiply on an RTX 5060 Ti, against cublas doing the same -// thing (cublasSgemm, and cublasGemmEx with half precision operands and a -// single precision accumulator). The block shapes below were picked by -// sweeping sixty of them at each size and operand type. +// thing at each pair of types. The block shapes below were picked by sweeping +// sixty of them at each size and operand type. // // 1024 2048 4096 -// Halide f32 313 us 1679 us 18796 us -// cublas f32 148 us 1034 us 7938 us -// Halide f16 -> f32 53 us 368 us 2811 us -// Halide bf16 -> f32 54 us 367 us 2810 us -// cublas f16 -> f32 52 us 350 us 2724 us -// Halide f16 -> f16 36 us 199 us 1592 us -// Halide u8 -> i32 35 us 220 us 1697 us +// Halide f32 312 us 1670 us 18812 us +// cublas f32 148 us 1026 us 7818 us // -// The float schedule is about half of cublas. Against cublas doing the same -// thing, half and brain float land within a few percent of it. The two rows -// below that are doing less work per multiply, so they are not comparable to -// the ones above: a half accumulator halves the registers it takes, and -// eight-bit operands halve the traffic through shared memory. Both are around -// 1.7x the single precision accumulator here, which makes bytes the -// interesting case for imaging, where the inputs are usually bytes anyway. +// Halide f16 -> f32 53 us 367 us 2810 us +// cublas f16 -> f32 51 us 350 us 2714 us +// +// Halide bf16 -> f32 53 us 367 us 2801 us +// cublas bf16 -> f32 51 us 350 us 2714 us +// +// Halide f16 -> f16 36 us 198 us 1589 us +// cublas f16 -> f16 31 us 228 us 1566 us +// +// Halide u8 -> i32 35 us 219 us 1691 us +// cublas s8 -> i32 20 us 140 us 1055 us +// +// Rows within a pair of types are comparable to each other; rows in different +// pairs are not, because a narrower accumulator or narrower operands are less +// work. The 16-bit float schedules land within a few percent of cublas, and +// the half accumulator one is ahead of it at 2048. The float schedule is +// about half of cublas, and the eight-bit one is the weakest: 55% of it at +// 1024 and 62% at 4096. cublas is not using the wmma instructions there - +// eight-bit operands have an mma shape with twice the reduction depth per +// instruction, which this schedule has no way to reach. cublas takes signed +// operands where the variant here takes unsigned ones; the hardware runs both +// at the same rate. // class MatMul : public Halide::Generator { public: diff --git a/apps/cuda_mat_mul/runner.cpp b/apps/cuda_mat_mul/runner.cpp index 93e793232dd4..43eab225aaad 100644 --- a/apps/cuda_mat_mul/runner.cpp +++ b/apps/cuda_mat_mul/runner.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include "mat_mul.h" @@ -279,6 +280,42 @@ int main(int argc, char **argv) { CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT); }, []() { cudaDeviceSynchronize(); }); printf("cublas half: %f s (%.1f GFlop/s)\n", t, gflops(size, t)); + + // Brain floats, also into a single precision accumulator. + t = bench_batched([&]() { cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, + size, size, size, &alpha, + A, CUDA_R_16BF, size, + B, CUDA_R_16BF, size, &beta, + C, CUDA_R_32F, size, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT); }, + []() { cudaDeviceSynchronize(); }); + printf("cublas bfloat: %f s (%.1f GFlop/s)\n", t, gflops(size, t)); + + // Eight-bit integers into a 32-bit accumulator. cublas takes + // signed operands here where the Halide variant above takes + // unsigned ones; the hardware runs both at the same rate. + int32_t ialpha = 1, ibeta = 1; + t = bench_batched([&]() { cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, + size, size, size, &ialpha, + A, CUDA_R_8I, size, + B, CUDA_R_8I, size, &ibeta, + C, CUDA_R_32I, size, + CUBLAS_COMPUTE_32I, CUBLAS_GEMM_DEFAULT); }, + []() { cudaDeviceSynchronize(); }); + printf("cublas int8: %f s (%.1f GFlop/s)\n", t, gflops(size, t)); + + // Halves into a half accumulator, which is what the half output + // variant above does. + __half halpha = __float2half(1.f), hbeta = __float2half(1.f); + t = bench_batched([&]() { cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, + size, size, size, &halpha, + A, CUDA_R_16F, size, + B, CUDA_R_16F, size, &hbeta, + C, CUDA_R_16F, size, + CUBLAS_COMPUTE_16F, CUBLAS_GEMM_DEFAULT); }, + []() { cudaDeviceSynchronize(); }); + printf("cublas half into half: %f s (%.1f GFlop/s)\n", + t, gflops(size, t)); } cudaFree(A); From fc1d0f3d9fb255839c2c37ffa31d77e5f71dd65a Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Mon, 3 Aug 2026 16:03:30 -0700 Subject: [PATCH 42/59] Widen the staged copy for byte operands, and retune them The staging vector was eight elements, which is the widest asynchronous copy for a 16-bit operand but half of it for a byte. Making it sixteen bytes rather than eight elements is the same element-versus-byte confusion the padding had. With that fixed, and with the staging depth swept as well - bytes fit twice the reduction depth in the same shared memory, and mostly want to - the eight-bit variant goes from 61 to 63 TFlop/s at 1024 and from 81 to 90 at 4096. That leaves it at 58% to 69% of cublas rather than 55% to 62%. The gap is still real, but it was not the wmma instructions, as the comment claimed before this: it was a schedule tuned for 16-bit operands. The table is now in GFlop/s with a peak column, which shows how far each row is from what the hardware could do rather than only from cublas. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/mat_mul_generator.cpp | 76 +++++++++++++++---------- 1 file changed, 45 insertions(+), 31 deletions(-) diff --git a/apps/cuda_mat_mul/mat_mul_generator.cpp b/apps/cuda_mat_mul/mat_mul_generator.cpp index d4622b392972..e25ac9a0b599 100644 --- a/apps/cuda_mat_mul/mat_mul_generator.cpp +++ b/apps/cuda_mat_mul/mat_mul_generator.cpp @@ -22,36 +22,39 @@ void set_alignment_and_bounds(OutputImageParam p, int size) { // accumulator type, so it picks between the accumulators an operand type can // pair with. // -// Time for one multiply on an RTX 5060 Ti, against cublas doing the same -// thing at each pair of types. The block shapes below were picked by sweeping -// sixty of them at each size and operand type. +// GFlop/s on an RTX 5060 Ti, against cublas doing the same thing at each +// pair of types, and against the peak the hardware could reach. The block +// shapes below were picked by sweeping them at each size and operand type. // -// 1024 2048 4096 -// Halide f32 312 us 1670 us 18812 us -// cublas f32 148 us 1026 us 7818 us +// 1024 2048 4096 peak +// Halide f32 6878 10294 7298 28500 +// cublas f32 14503 16574 17449 // -// Halide f16 -> f32 53 us 367 us 2810 us -// cublas f16 -> f32 51 us 350 us 2714 us +// Halide f16 -> f32 40070 46671 48888 56980 +// cublas f16 -> f32 41658 49069 50440 // -// Halide bf16 -> f32 53 us 367 us 2801 us -// cublas bf16 -> f32 51 us 350 us 2714 us +// Halide bf16 -> f32 40025 46681 48692 56980 +// cublas bf16 -> f32 41664 49054 50437 // -// Halide f16 -> f16 36 us 198 us 1589 us -// cublas f16 -> f16 31 us 228 us 1566 us +// Halide f16 -> f16 60349 86502 86599 113960 +// cublas f16 -> f16 69221 75073 87073 // -// Halide u8 -> i32 35 us 219 us 1691 us -// cublas s8 -> i32 20 us 140 us 1055 us +// Halide u8 -> i32 62869 82391 89641 227920 +// cublas s8 -> i32 107868 122203 129970 +// +// The peak column is 36 SMs times the 3090 MHz maximum clock times the rate +// per SM per clock, which is 256 flops for the cuda cores, and 512, 1024 and +// 2048 for the tensor cores at each accumulator width. Those per-SM rates are +// the consumer part pattern rather than something measured here, so treat the +// column as a scale rather than a number. // // Rows within a pair of types are comparable to each other; rows in different // pairs are not, because a narrower accumulator or narrower operands are less -// work. The 16-bit float schedules land within a few percent of cublas, and -// the half accumulator one is ahead of it at 2048. The float schedule is -// about half of cublas, and the eight-bit one is the weakest: 55% of it at -// 1024 and 62% at 4096. cublas is not using the wmma instructions there - -// eight-bit operands have an mma shape with twice the reduction depth per -// instruction, which this schedule has no way to reach. cublas takes signed -// operands where the variant here takes unsigned ones; the hardware runs both -// at the same rate. +// work. The two 16-bit float schedules land within a few percent of cublas, +// and the half accumulator one is ahead of it at 2048. The eight-bit one is +// the weakest against cublas, at 58% to 69% of it, though it is the fastest +// thing here in absolute terms - which is what makes bytes interesting for +// imaging, where the inputs are usually bytes anyway. // class MatMul : public Halide::Generator { public: @@ -65,8 +68,9 @@ class MatMul : public Halide::Generator { GeneratorParam tiles_y{"tiles_y", 0}; GeneratorParam warps_x{"warps_x", 0}; GeneratorParam warps_y{"warps_y", 0}; - // How much of the reduction is staged in shared memory at a time. - GeneratorParam block_r{"block_r", 32}; + // How much of the reduction is staged in shared memory at a time. Zero + // means pick it along with the block shape below. + GeneratorParam block_r{"block_r", 0}; // Extra elements per row of the shared panels, which spreads consecutive // rows across different banks. Zero means pad by sixteen bytes, which is // the least that keeps each row aligned both for the widest asynchronous @@ -183,6 +187,9 @@ class MatMul : public Halide::Generator { // tiles_y) multiplies, so this is what gets us reuse out of the loads. const int tile = 16; int tx = tiles_x, ty = tiles_y, wx = warps_x, wy = warps_y; + // The staging depth goes with the shape, so an explicit block_r only + // wins if it was asked for. + int br = 32; if (tx == 0 || ty == 0 || wx == 0 || wy == 0) { // Measured on an RTX 5060 Ti by sweeping sixty shapes at each size // and operand type. These do not follow a trend worth @@ -192,19 +199,21 @@ class MatMul : public Halide::Generator { // type matters as much as the size: bytes make the operand loads // cheap enough to pay for a much larger accumulator, so they want // a tall block spread over four warps, where the 16-bit types want - // a wide one over fewer. + // a wide one over fewer, with a deeper staged panel to match. // // Brain floats pick out exactly the same shapes as halves at // every size, so they share a row here. const bool bytes = A.type().bits() == 8; const bool half_accumulator = out.type() == Float(16); if (bytes) { + // Bytes stage twice the reduction depth in the same shared + // memory, and mostly want to. if ((int)size <= 1024) { - tx = 2, ty = 8, wx = 4, wy = 1; + tx = 2, ty = 8, wx = 2, wy = 1, br = 64; } else if ((int)size <= 2048) { - tx = 2, ty = 10, wx = 4, wy = 1; + tx = 2, ty = 10, wx = 2, wy = 1, br = 32; } else { - tx = 2, ty = 8, wx = 4, wy = 1; + tx = 2, ty = 8, wx = 2, wy = 1, br = 64; } } else if (half_accumulator) { if ((int)size <= 1024) { @@ -224,6 +233,9 @@ class MatMul : public Halide::Generator { } } } + if (block_r) { + br = block_r; + } const int block_x = tile * tx * wx; const int block_y = tile * ty * wy; @@ -266,7 +278,7 @@ class MatMul : public Halide::Generator { .unroll(yi); prod.update() - .split(r, ro, ri, block_r) + .split(r, ro, ri, br) .split(x, xw, xi, tile * tx) .split(xi, xi, rxi, tile) .split(y, yw, yi, tile * ty) @@ -287,12 +299,14 @@ class MatMul : public Halide::Generator { // bytes at a time along the dense dimension, so that the reads from // global memory coalesce and the writes to shared memory can be done // as asynchronous copies. - const int vec = 8; + // Each thread moves sixteen bytes, the widest asynchronous copy the + // hardware has. How many elements that is depends on the operand type. + const int vec = 16 / A.type().bytes(); Var rro("rro"), rrv("rrv"), xxo("xxo"), xxi("xxi"); Var t("t"), ti("ti"), tw("tw"), tw2("tw2"), to("to"); // B.in() is dense in the reduction dimension, which is its _0. - B.in().compute_at(prod, ro).store_in(MemoryType::GPUSharedAsync).align_storage(_0, (int)block_r + pa).split(_0, rro, rrv, vec).fuse(rro, _1, t).split(t, t, ti, 32).split(t, t, tw, wx).split(t, to, tw2, wy).gpu_lanes(ti).gpu_threads(tw, tw2).vectorize(rrv); + B.in().compute_at(prod, ro).store_in(MemoryType::GPUSharedAsync).align_storage(_0, br + pa).split(_0, rro, rrv, vec).fuse(rro, _1, t).split(t, t, ti, 32).split(t, t, tw, wx).split(t, to, tw2, wy).gpu_lanes(ti).gpu_threads(tw, tw2).vectorize(rrv); // A.in() is dense in x, which is its _0. A.in().compute_at(prod, ro).store_in(MemoryType::GPUSharedAsync).align_storage(_0, block_x + pb).split(_0, xxo, xxi, vec).fuse(xxo, _1, t).split(t, t, ti, 32).split(t, t, tw, wx).split(t, to, tw2, wy).gpu_lanes(ti).gpu_threads(tw, tw2).vectorize(xxi); From a966cc8df63c6f7b9ef42ddf531ab16e32367dd1 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Mon, 3 Aug 2026 16:09:16 -0700 Subject: [PATCH 43/59] Measure the instruction ceilings, and compare against those The previous comment measured the schedules against a peak computed from the maximum clock and an assumed rate per SM per clock. Both inputs were wrong: the part averages 2817 MHz while benchmarking rather than its 3090 maximum, and wmma multiplies bytes at the same rate it multiplies halves into halves, not twice it. Issuing back-to-back wmma instructions out of registers with no memory traffic gives the real ceilings: 51541, 99626 and 100650 GOP/s. Against those the schedules reach 95%, 87% and 89%, and the half-into-half one matches cublas exactly. That is a much better account of them than 39% of a peak they could never reach. It also settles the eight-bit gap, which an earlier commit message got wrong in both directions. cublas is 29% past the wmma ceiling, so it cannot be using wmma. The mma instructions reach 188355 GOP/s at the same shape, 1.87x. Both do the same 8192 ops per instruction - the earlier claim that the deeper reduction was more work per instruction was wrong - so what differs is purely how fast the two families issue. The gap is structural after all, but not for the reason first given. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/mat_mul_generator.cpp | 45 ++++++++++++++----------- 1 file changed, 25 insertions(+), 20 deletions(-) diff --git a/apps/cuda_mat_mul/mat_mul_generator.cpp b/apps/cuda_mat_mul/mat_mul_generator.cpp index e25ac9a0b599..e9a2237e4702 100644 --- a/apps/cuda_mat_mul/mat_mul_generator.cpp +++ b/apps/cuda_mat_mul/mat_mul_generator.cpp @@ -23,38 +23,43 @@ void set_alignment_and_bounds(OutputImageParam p, int size) { // pair with. // // GFlop/s on an RTX 5060 Ti, against cublas doing the same thing at each -// pair of types, and against the peak the hardware could reach. The block -// shapes below were picked by sweeping them at each size and operand type. +// pair of types, and against the ceiling of the instructions this schedule +// uses. The block shapes below were picked by sweeping them at each size and +// operand type. // -// 1024 2048 4096 peak -// Halide f32 6878 10294 7298 28500 +// 1024 2048 4096 ceiling +// Halide f32 6878 10294 7298 25960 // cublas f32 14503 16574 17449 // -// Halide f16 -> f32 40070 46671 48888 56980 +// Halide f16 -> f32 40070 46671 48888 51541 // cublas f16 -> f32 41658 49069 50440 // -// Halide bf16 -> f32 40025 46681 48692 56980 +// Halide bf16 -> f32 40025 46681 48692 51541 // cublas bf16 -> f32 41664 49054 50437 // -// Halide f16 -> f16 60349 86502 86599 113960 +// Halide f16 -> f16 60349 86502 86599 99626 // cublas f16 -> f16 69221 75073 87073 // -// Halide u8 -> i32 62869 82391 89641 227920 +// Halide u8 -> i32 62869 82391 89641 100650 // cublas s8 -> i32 107868 122203 129970 // -// The peak column is 36 SMs times the 3090 MHz maximum clock times the rate -// per SM per clock, which is 256 flops for the cuda cores, and 512, 1024 and -// 2048 for the tensor cores at each accumulator width. Those per-SM rates are -// the consumer part pattern rather than something measured here, so treat the -// column as a scale rather than a number. +// The ceiling for the tensor core rows is measured, by issuing back-to-back +// wmma instructions out of registers with no memory traffic at all. The one +// for the float row is 36 SMs times the 2817 MHz this part averages while +// benchmarking times 256 flops per SM per clock, which is what the cuda cores +// do. Rows within a pair of types are comparable to each other; rows in +// different pairs are not, because a narrower accumulator or narrower +// operands are less work. // -// Rows within a pair of types are comparable to each other; rows in different -// pairs are not, because a narrower accumulator or narrower operands are less -// work. The two 16-bit float schedules land within a few percent of cublas, -// and the half accumulator one is ahead of it at 2048. The eight-bit one is -// the weakest against cublas, at 58% to 69% of it, though it is the fastest -// thing here in absolute terms - which is what makes bytes interesting for -// imaging, where the inputs are usually bytes anyway. +// So the schedules reach 95%, 87% and 89% of what their instructions can do, +// and the last of those matches cublas at f16 -> f16 exactly. The one row +// that does not is eight-bit, where cublas is 29% past the ceiling of the +// instruction used here: wmma multiplies bytes at the same rate it multiplies +// halves into halves, whereas the mma instructions reach 188355 GOP/s at the +// same shape - 1.87x - so cublas must be using those. Both instructions do +// the same 8192 ops each; what differs is how fast they issue. Reaching that +// needs mma rather than wmma, which is a different fragment layout and not +// something this schedule can express. // class MatMul : public Halide::Generator { public: From 4207b5218ecde416a2e9cf8dfc2d1c04eed71c7b Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Mon, 3 Aug 2026 16:14:17 -0700 Subject: [PATCH 44/59] Say why the single precision row is the slow one It reaches 28% of what the cuda cores can do where the tensor core schedules reach 87% to 95% of their instructions, and the difference is one missing scheduling primitive rather than an old schedule. Sharing a staged panel across a block needs the reduction chunk loop above the thread loops and the accumulator below them, but an accumulator that outlives the chunk loop lands at block level, where a Register allocation is sized for the whole block tile and spills. The tensor core schedules only get around it because a WMMAFragment allocation at block level is already per-lane. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/mat_mul_generator.cpp | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/apps/cuda_mat_mul/mat_mul_generator.cpp b/apps/cuda_mat_mul/mat_mul_generator.cpp index e9a2237e4702..640a90f097fd 100644 --- a/apps/cuda_mat_mul/mat_mul_generator.cpp +++ b/apps/cuda_mat_mul/mat_mul_generator.cpp @@ -51,7 +51,20 @@ void set_alignment_and_bounds(OutputImageParam p, int size) { // different pairs are not, because a narrower accumulator or narrower // operands are less work. // -// So the schedules reach 95%, 87% and 89% of what their instructions can do, +// The float row is the one far from its ceiling, at 28%, and that is a +// missing scheduling primitive rather than a bad schedule. A fast single +// precision matmul wants the loop over chunks of the reduction above the +// thread loops, so that one staged panel serves the whole block, and the +// accumulator below them in registers, living across that loop. An +// accumulator that outlives the chunk loop has to be declared outside it, +// which puts it at block level, and a block level MemoryType::Register +// allocation is sized for the whole block tile and spills to local memory. +// So this schedule stages its operands per thread instead, and shares +// nothing across the block. The tensor core schedules get around it only +// because a WMMAFragment allocation at block level is already per-lane. +// +// So the tensor core schedules reach 95%, 87% and 89% of what their +// instructions can do, // and the last of those matches cublas at f16 -> f16 exactly. The one row // that does not is eight-bit, where cublas is 29% past the ceiling of the // instruction used here: wmma multiplies bytes at the same rate it multiplies From ad40d8561b4626994ef49da98704883aae56098e Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Mon, 3 Aug 2026 16:17:52 -0700 Subject: [PATCH 45/59] Tighten the matmul comments, and fix the stale ones Four had gone stale as the generator grew: the header still said only half precision reached the tensor cores, tiles_x said the best block gets smaller as the matrices grow, the block shape comment said bytes want four warps when they want two, and the tensor core schedule quoted ratios against cublas from before the retune. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/mat_mul_generator.cpp | 141 +++++++++--------------- 1 file changed, 55 insertions(+), 86 deletions(-) diff --git a/apps/cuda_mat_mul/mat_mul_generator.cpp b/apps/cuda_mat_mul/mat_mul_generator.cpp index 640a90f097fd..8256c78de06e 100644 --- a/apps/cuda_mat_mul/mat_mul_generator.cpp +++ b/apps/cuda_mat_mul/mat_mul_generator.cpp @@ -14,18 +14,16 @@ void set_alignment_and_bounds(OutputImageParam p, int size) { .set_stride(size); } -// A square matrix multiply, scheduled two ways. The operands are untyped, so -// their type is a generator param, and it is what picks the schedule: half -// precision operands get the tensor cores, and anything else gets a schedule -// that keeps the accumulator in ordinary registers. The tensor cores multiply -// halves, brain floats, or eight-bit integers; the output type is also the -// accumulator type, so it picks between the accumulators an operand type can -// pair with. +// A square matrix multiply. The operands are untyped, so their type is a +// generator param, and it picks the schedule: the tensor cores multiply +// halves, brain floats and bytes, and anything else accumulates in ordinary +// registers. The output type is the accumulator type, so it picks between the +// accumulators an operand type can pair with. // -// GFlop/s on an RTX 5060 Ti, against cublas doing the same thing at each -// pair of types, and against the ceiling of the instructions this schedule -// uses. The block shapes below were picked by sweeping them at each size and -// operand type. +// GFlop/s on an RTX 5060 Ti, against cublas and against the ceiling of the +// instructions used. Rows within a pair of types are comparable; rows in +// different pairs are not, because narrower operands or a narrower +// accumulator are less work. // // 1024 2048 4096 ceiling // Halide f32 6878 10294 7298 25960 @@ -43,81 +41,64 @@ void set_alignment_and_bounds(OutputImageParam p, int size) { // Halide u8 -> i32 62869 82391 89641 100650 // cublas s8 -> i32 107868 122203 129970 // -// The ceiling for the tensor core rows is measured, by issuing back-to-back -// wmma instructions out of registers with no memory traffic at all. The one -// for the float row is 36 SMs times the 2817 MHz this part averages while -// benchmarking times 256 flops per SM per clock, which is what the cuda cores -// do. Rows within a pair of types are comparable to each other; rows in -// different pairs are not, because a narrower accumulator or narrower -// operands are less work. +// The tensor core ceilings are measured, by issuing wmma instructions back to +// back out of registers. The float one is 36 SMs times the 2817 MHz this part +// averages while benchmarking times the 256 flops per SM per clock the cuda +// cores do. // -// The float row is the one far from its ceiling, at 28%, and that is a -// missing scheduling primitive rather than a bad schedule. A fast single -// precision matmul wants the loop over chunks of the reduction above the -// thread loops, so that one staged panel serves the whole block, and the -// accumulator below them in registers, living across that loop. An -// accumulator that outlives the chunk loop has to be declared outside it, -// which puts it at block level, and a block level MemoryType::Register -// allocation is sized for the whole block tile and spills to local memory. -// So this schedule stages its operands per thread instead, and shares -// nothing across the block. The tensor core schedules get around it only -// because a WMMAFragment allocation at block level is already per-lane. +// The tensor core schedules reach 87% to 95% of their ceilings, and match +// cublas at f16 -> f16. Two rows fall short for reasons outside the schedule. // -// So the tensor core schedules reach 95%, 87% and 89% of what their -// instructions can do, -// and the last of those matches cublas at f16 -> f16 exactly. The one row -// that does not is eight-bit, where cublas is 29% past the ceiling of the -// instruction used here: wmma multiplies bytes at the same rate it multiplies -// halves into halves, whereas the mma instructions reach 188355 GOP/s at the -// same shape - 1.87x - so cublas must be using those. Both instructions do -// the same 8192 ops each; what differs is how fast they issue. Reaching that -// needs mma rather than wmma, which is a different fragment layout and not -// something this schedule can express. +// At eight bits cublas is 29% past the ceiling, so it is not using wmma, which +// multiplies bytes no faster than it multiplies halves into halves. The mma +// instructions reach 188355 GOP/s at the same shape, 1.87x, for the same 8192 +// ops per instruction. Reaching that needs mma's fragment layout, which this +// schedule cannot express. +// +// The float row is at 28% because sharing a staged panel across a block needs +// the reduction chunk loop above the thread loops and the accumulator below +// them, living across it. Such an accumulator lands at block level, where a +// Register allocation is sized for the whole block tile and spills. So it +// stages per thread and shares nothing. The tensor core schedules only escape +// this because a WMMAFragment allocation at block level is already per-lane. // class MatMul : public Halide::Generator { public: GeneratorParam size{"size", 1024}; // How many tensor core tiles of accumulator each warp holds, and how many - // warps there are per block in each dimension. Zero means pick a shape - // based on the problem size. The best block gets smaller as the matrices - // do, because a large one leaves too few blocks to fill the machine. + // warps per block in each dimension. Zero means use the measured shapes + // below. GeneratorParam tiles_x{"tiles_x", 0}; GeneratorParam tiles_y{"tiles_y", 0}; GeneratorParam warps_x{"warps_x", 0}; GeneratorParam warps_y{"warps_y", 0}; // How much of the reduction is staged in shared memory at a time. Zero - // means pick it along with the block shape below. + // means use the depth that goes with the shape below. GeneratorParam block_r{"block_r", 0}; - // Extra elements per row of the shared panels, which spreads consecutive - // rows across different banks. Zero means pad by sixteen bytes, which is - // the least that keeps each row aligned both for the widest asynchronous - // copy and for the tensor core loads, whose matrix addresses have to be - // sixteen byte aligned. How many elements that is depends on the operand - // type, which is why it is not just a number here. + // Extra elements per row of the shared panels, to spread consecutive rows + // across banks. Zero means sixteen bytes, the least that keeps each row + // aligned for both the widest asynchronous copy and the tensor core loads. + // That is a different number of elements per operand type. GeneratorParam pad_a{"pad_a", 0}; GeneratorParam pad_b{"pad_b", 0}; Input> A{"A"}; Input> B{"B"}; - // The output type is also the accumulator type - there is no tensor core - // store from a half precision fragment into single precision memory - so - // asking for a half precision output is what asks to accumulate in half. - // That halves the registers the accumulator takes, but summing a long - // reduction in half precision loses accuracy badly, so it is only worth - // asking for when the numerics of the problem allow it. + // The output type is the accumulator type - there is no tensor core store + // from a half fragment into single precision memory. A half accumulator + // halves the registers it takes, but loses accuracy over a long reduction, + // so only ask for it when the numerics allow. Output> out{"out"}; - // The tensor cores multiply 16-bit floats or 8-bit integers, so asking for - // one of those operand types is what asks for them. Anything else gets the - // schedule that accumulates in ordinary registers. + // Asking for a type the tensor cores multiply is what asks for them. bool use_tensor_cores() const { Type t = A.type(); return t == Float(16) || t == BFloat(16) || t == Int(8) || t == UInt(8); } - // The accumulator each operand type pairs with. Half precision can also + // The accumulator each operand type pairs with. Halves can also // accumulate into halves, which is what asking for a half output does. Type natural_accumulator() const { Type t = A.type(); @@ -141,10 +122,9 @@ class MatMul : public Halide::Generator { Type acc = out.type(); prod(x, y) = cast(acc, 0); - // The widening to the accumulator type happens here, at the multiply, - // rather than in the operand wrappers the schedule stages through - // shared memory. That way half precision operands are staged as half - // precision and reach the tensor cores as such. + // The widening happens here, at the multiply, rather than in the + // operand wrappers the schedule stages through shared memory, so that + // narrow operands are staged narrow and reach the tensor cores so. prod(x, y) += cast(acc, A(x, r)) * cast(acc, B(r, y)); out(x, y) = prod(x, y); @@ -172,7 +152,7 @@ class MatMul : public Halide::Generator { } private: - // The float schedule. See the table above for how it does. + // See the table above. void schedule_cuda() { Var xi, yi, xii, yii; @@ -196,36 +176,25 @@ class MatMul : public Halide::Generator { B.in().compute_at(prod, r).vectorize(_0).unroll(_1); } - // The tensor core schedule, which reaches 96%, 95% and 97% of cublas at - // the three sizes in the table above. The block shapes below are picked - // per size by measurement. + // See the table above. void schedule_tensor_cores() { - // The tensor core tile shape, and how many of them each warp - // accumulates at once. Each operand tile loaded feeds tiles_x (or - // tiles_y) multiplies, so this is what gets us reuse out of the loads. + // Each operand tile loaded feeds tiles_x (or tiles_y) multiplies, so + // the tile counts are what get reuse out of the loads. const int tile = 16; int tx = tiles_x, ty = tiles_y, wx = warps_x, wy = warps_y; - // The staging depth goes with the shape, so an explicit block_r only - // wins if it was asked for. + // The staging depth goes with the shape below unless asked for. int br = 32; if (tx == 0 || ty == 0 || wx == 0 || wy == 0) { - // Measured on an RTX 5060 Ti by sweeping sixty shapes at each size - // and operand type. These do not follow a trend worth - // extrapolating from, so they are measured points rather than a - // formula. What moves between them is how much accumulator a warp - // holds and how many warps share a staged panel, and the operand - // type matters as much as the size: bytes make the operand loads - // cheap enough to pay for a much larger accumulator, so they want - // a tall block spread over four warps, where the 16-bit types want - // a wide one over fewer, with a deeper staged panel to match. - // - // Brain floats pick out exactly the same shapes as halves at - // every size, so they share a row here. + // Measured on an RTX 5060 Ti by sweeping shapes at each size and + // operand type. There is no trend worth extrapolating from, so + // these are measured points. The operand type matters as much as + // the size: bytes want a tall block and a deep staged panel, where + // the 16-bit types want a wide block and a shallow one. Brain + // floats pick out the same shapes as halves at every size, so they + // share a row. const bool bytes = A.type().bits() == 8; const bool half_accumulator = out.type() == Float(16); if (bytes) { - // Bytes stage twice the reduction depth in the same shared - // memory, and mostly want to. if ((int)size <= 1024) { tx = 2, ty = 8, wx = 2, wy = 1, br = 64; } else if ((int)size <= 2048) { From 15ba249814e9e2c6c36d10bb2f23d058ba7b150b Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Mon, 3 Aug 2026 16:35:03 -0700 Subject: [PATCH 46/59] Break the staging schedules back across lines Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/mat_mul_generator.cpp | 33 ++++++++++++++++++++----- 1 file changed, 27 insertions(+), 6 deletions(-) diff --git a/apps/cuda_mat_mul/mat_mul_generator.cpp b/apps/cuda_mat_mul/mat_mul_generator.cpp index 8256c78de06e..0435280ea43e 100644 --- a/apps/cuda_mat_mul/mat_mul_generator.cpp +++ b/apps/cuda_mat_mul/mat_mul_generator.cpp @@ -16,7 +16,7 @@ void set_alignment_and_bounds(OutputImageParam p, int size) { // A square matrix multiply. The operands are untyped, so their type is a // generator param, and it picks the schedule: the tensor cores multiply -// halves, brain floats and bytes, and anything else accumulates in ordinary +// halves, bfloats and bytes, and anything else accumulates in ordinary // registers. The output type is the accumulator type, so it picks between the // accumulators an operand type can pair with. // @@ -189,9 +189,8 @@ class MatMul : public Halide::Generator { // operand type. There is no trend worth extrapolating from, so // these are measured points. The operand type matters as much as // the size: bytes want a tall block and a deep staged panel, where - // the 16-bit types want a wide block and a shallow one. Brain - // floats pick out the same shapes as halves at every size, so they - // share a row. + // the 16-bit types want a wide block and a shallow one. Bfloats pick + // out the same shapes as halves at every size, so they share a row. const bool bytes = A.type().bits() == 8; const bool half_accumulator = out.type() == Float(16); if (bytes) { @@ -293,10 +292,32 @@ class MatMul : public Halide::Generator { Var t("t"), ti("ti"), tw("tw"), tw2("tw2"), to("to"); // B.in() is dense in the reduction dimension, which is its _0. - B.in().compute_at(prod, ro).store_in(MemoryType::GPUSharedAsync).align_storage(_0, br + pa).split(_0, rro, rrv, vec).fuse(rro, _1, t).split(t, t, ti, 32).split(t, t, tw, wx).split(t, to, tw2, wy).gpu_lanes(ti).gpu_threads(tw, tw2).vectorize(rrv); + B.in() + .compute_at(prod, ro) + .store_in(MemoryType::GPUSharedAsync) + .align_storage(_0, br + pa) + .split(_0, rro, rrv, vec) + .fuse(rro, _1, t) + .split(t, t, ti, 32) + .split(t, t, tw, wx) + .split(t, to, tw2, wy) + .gpu_lanes(ti) + .gpu_threads(tw, tw2) + .vectorize(rrv); // A.in() is dense in x, which is its _0. - A.in().compute_at(prod, ro).store_in(MemoryType::GPUSharedAsync).align_storage(_0, block_x + pb).split(_0, xxo, xxi, vec).fuse(xxo, _1, t).split(t, t, ti, 32).split(t, t, tw, wx).split(t, to, tw2, wy).gpu_lanes(ti).gpu_threads(tw, tw2).vectorize(xxi); + A.in() + .compute_at(prod, ro) + .store_in(MemoryType::GPUSharedAsync) + .align_storage(_0, block_x + pb) + .split(_0, xxo, xxi, vec) + .fuse(xxo, _1, t) + .split(t, t, ti, 32) + .split(t, t, tw, wx) + .split(t, to, tw2, wy) + .gpu_lanes(ti) + .gpu_threads(tw, tw2) + .vectorize(xxi); } Var x{"x"}, y{"y"}; From fa13cb26dd24cf3f67dfa38ab6b80bdfdd251fbe Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 4 Aug 2026 09:32:44 -0700 Subject: [PATCH 47/59] Use the benchmark helper for the batched timing Running the filter five times inside the lambda and synchronizing once at the end amortizes the per-launch synchronization just as well as the hand-rolled sampling loop did, and leaves the sampling to the helper. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/runner.cpp | 37 +++++++++++++----------------------- 1 file changed, 13 insertions(+), 24 deletions(-) diff --git a/apps/cuda_mat_mul/runner.cpp b/apps/cuda_mat_mul/runner.cpp index 43eab225aaad..6451c453f645 100644 --- a/apps/cuda_mat_mul/runner.cpp +++ b/apps/cuda_mat_mul/runner.cpp @@ -67,34 +67,23 @@ double gflops(int size, double seconds) { return 2.0 * size * size * size / seconds * 1e-9; } -// Time a batch of launches with a single synchronization at the end, rather -// than synchronizing after each one. Both implementations queue work -// asynchronously, and cublas in particular does a heuristic lookup on the host -// for every call, so synchronizing per launch measures that host work instead -// of letting it overlap with the GPU. +// Time one call of the filter. Both implementations queue work +// asynchronously, and cublas does a heuristic lookup on the host for every +// call, so synchronizing per launch would measure that host work rather than +// letting it overlap with the GPU. Batch the launches instead, and sync once. // `sync` has to match the launcher: Halide runs on its own CUDA context, so // cudaDeviceSynchronize does not wait for it. template double bench_batched(F &&launch, S &&sync) { - const int samples = 5, iterations = 5; - for (int i = 0; i < iterations; i++) { - launch(); - } - sync(); - double best = 0; - for (int s = 0; s < samples; s++) { - auto t0 = Halide::Tools::benchmark_now(); - for (int i = 0; i < iterations; i++) { - launch(); - } - sync(); - auto t1 = Halide::Tools::benchmark_now(); - double t = Halide::Tools::benchmark_duration_seconds(t0, t1) / iterations; - if (s == 0 || t < best) { - best = t; - } - } - return best; + const int batch = 5; + return Halide::Tools::benchmark(5, 1, + [&]() { + for (int i = 0; i < batch; i++) { + launch(); + } + sync(); + }) / + batch; } } // namespace From e39f114eccc605838561fafcc7adc0b57632d5ae Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 4 Aug 2026 09:38:35 -0700 Subject: [PATCH 48/59] Note that the benchmark data was checked for compressibility The operands are small integers so the results can be checked exactly, which also makes them very compressible - three quarters of them are zero. Measuring against dense random operands instead moves nothing by more than a couple of percent, on either side of the comparison. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/runner.cpp | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/apps/cuda_mat_mul/runner.cpp b/apps/cuda_mat_mul/runner.cpp index 6451c453f645..db170f77304f 100644 --- a/apps/cuda_mat_mul/runner.cpp +++ b/apps/cuda_mat_mul/runner.cpp @@ -18,6 +18,13 @@ using Halide::Runtime::Buffer; namespace { +// The operands are small integers so that every dot product is exact in every +// accumulator here, which lets the results be checked for equality rather than +// to a tolerance. That makes them unusually compressible, so the numbers were +// checked against dense random operands too: no configuration moved by more +// than a couple of percent, which is what you would expect of a multiply that +// is issue-bound rather than waiting on memory. +// // The same matrix multiply is compiled twice from one generator: once with // float operands, which get a schedule that accumulates in ordinary registers, // and once with half operands, which get the tensor cores. Both accumulate in From 8ba8604b55e1d11fb9884f0006098780b5e5743f Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 4 Aug 2026 09:43:14 -0700 Subject: [PATCH 49/59] Hand cublas the same device buffers as the Halide filters The two sides had their own allocations, so they were multiplying different data - Halide's small integers against cublas's memset pattern. Fill a Halide buffer, let running the filter copy it down, and pass halide_cuda_get_device_ptr of it to cublas, so both multiply the same numbers. That drops the separate cudaMalloc and memset entirely. Doing that made the two implementations pair up naturally, so each pair of types is now one function that runs the filter, checks it, times it, and times cublas beside it. The runner is 64 lines shorter. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/runner.cpp | 428 +++++++++++++++-------------------- 1 file changed, 182 insertions(+), 246 deletions(-) diff --git a/apps/cuda_mat_mul/runner.cpp b/apps/cuda_mat_mul/runner.cpp index db170f77304f..6ab9db5201f3 100644 --- a/apps/cuda_mat_mul/runner.cpp +++ b/apps/cuda_mat_mul/runner.cpp @@ -18,35 +18,107 @@ using Halide::Runtime::Buffer; namespace { +// The same matrix multiply is compiled from one generator at each pair of +// operand and accumulator types, and each is compared against cublas doing the +// same thing. Both are handed the same device buffers - the Halide side fills +// them and copies them down, and cublas is given the pointers out of them - so +// the two see identical data. +// // The operands are small integers so that every dot product is exact in every // accumulator here, which lets the results be checked for equality rather than // to a tolerance. That makes them unusually compressible, so the numbers were // checked against dense random operands too: no configuration moved by more // than a couple of percent, which is what you would expect of a multiply that // is issue-bound rather than waiting on memory. -// -// The same matrix multiply is compiled twice from one generator: once with -// float operands, which get a schedule that accumulates in ordinary registers, -// and once with half operands, which get the tensor cores. Both accumulate in -// and return single precision, so the two are directly comparable. -template -bool check(const Buffer &A, const Buffer &B, - const Buffer &C, int size, const char *name) { - // Spot check on strides that are coprime with the tile sizes, so the - // samples land at varying offsets within a tile. +// Time one call. Both implementations queue work asynchronously, and cublas +// does a heuristic lookup on the host for every call, so synchronizing per +// launch would measure that host work rather than letting it overlap with the +// GPU. Batch the launches instead, and sync once. `sync` has to match the +// launcher: Halide runs on its own CUDA context, so cudaDeviceSynchronize does +// not wait for it. +template +double bench(F &&launch, S &&sync) { + const int batch = 5; + return Halide::Tools::benchmark(5, 1, + [&]() { + for (int i = 0; i < batch; i++) { + launch(); + } + sync(); + }) / + batch; +} + +double gflops(int size, double seconds) { + return 2.0 * size * size * size / seconds * 1e-9; +} + +// Read or write an element of a buffer of any of the types used here, so that +// the ones with no C++ equivalent are handled too. +double element(const halide_buffer_t *b, size_t i) { + halide_type_t t = b->type; + if (t == halide_type_t(halide_type_float, 16)) { + return (double)((const _Float16 *)b->host)[i]; + } else if (t == halide_type_t(halide_type_bfloat, 16)) { + // A bfloat is the top half of a float. + uint32_t bits = (uint32_t)((const uint16_t *)b->host)[i] << 16; + float f; + memcpy(&f, &bits, 4); + return f; + } else if (t == halide_type_t(halide_type_float, 32)) { + return ((const float *)b->host)[i]; + } else if (t == halide_type_t(halide_type_uint, 8)) { + return ((const uint8_t *)b->host)[i]; + } else if (t == halide_type_t(halide_type_int, 32)) { + return ((const int32_t *)b->host)[i]; + } + fprintf(stderr, "unhandled buffer type\n"); + exit(1); +} + +void set_element(halide_buffer_t *b, size_t i, int v) { + halide_type_t t = b->type; + if (t == halide_type_t(halide_type_float, 16)) { + ((_Float16 *)b->host)[i] = (_Float16)v; + } else if (t == halide_type_t(halide_type_bfloat, 16)) { + float f = (float)v; + uint32_t bits; + memcpy(&bits, &f, 4); + ((uint16_t *)b->host)[i] = (uint16_t)(bits >> 16); + } else if (t == halide_type_t(halide_type_float, 32)) { + ((float *)b->host)[i] = (float)v; + } else if (t == halide_type_t(halide_type_uint, 8)) { + ((uint8_t *)b->host)[i] = (uint8_t)v; + } else { + fprintf(stderr, "unhandled operand type\n"); + exit(1); + } +} + +void fill_ints(Buffer &b, int modulus) { + size_t n = (size_t)b.width() * b.height(); + for (size_t i = 0; i < n; i++) { + set_element(b.raw_buffer(), i, rand() % modulus); + } + b.set_host_dirty(); +} + +// Spot check on strides coprime with the tile sizes, so the samples land at +// varying offsets within a tile. +bool check(Buffer &Ab, Buffer &Bb, Buffer &Cb, + int size, const char *name) { for (int y = 0; y < size; y += 97) { for (int x = 0; x < size; x += 89) { double correct = 0; for (int k = 0; k < size; k++) { - correct += (double)A(x, k) * (double)B(k, y); + correct += element(Ab.raw_buffer(), (size_t)k * size + x) * + element(Bb.raw_buffer(), (size_t)y * size + k); } - // The operands are small integers, which are exact in both float - // and half, and the accumulator is single precision either way, so - // the answer should be exact. - if ((double)C(x, y) != correct) { + double got = element(Cb.raw_buffer(), (size_t)y * size + x); + if (got != correct) { printf("%s: bad result at %d %d: %f != %f\n", - name, x, y, (double)C(x, y), correct); + name, x, y, got, correct); return false; } } @@ -54,44 +126,41 @@ bool check(const Buffer &A, const Buffer &B, return true; } -// There is no C++ type for bfloat16 here, so the buffer carries the type at -// runtime and these convert. A bfloat is the top half of a float, so for the -// small integers this uses the conversion is exact in both directions. -uint16_t to_bf16(float f) { - uint32_t bits; - memcpy(&bits, &f, 4); - return (uint16_t)(bits >> 16); -} +// One row of the table: run the Halide filter, check it, time it, then time +// cublas on the same device buffers. +template +bool row(const char *name, int size, halide_type_t a_type, halide_type_t c_type, + int modulus, Filter filter, Cublas cublas) { + Buffer Ab(a_type, size, size), Bb(a_type, size, size); + Buffer Cb(c_type, size, size), Cb_cublas(c_type, size, size); + fill_ints(Ab, modulus); + fill_ints(Bb, modulus); -float from_bf16(uint16_t h) { - uint32_t bits = (uint32_t)h << 16; - float f; - memcpy(&f, &bits, 4); - return f; -} + if (filter(Ab.raw_buffer(), Bb.raw_buffer(), Cb.raw_buffer()) != 0) { + printf("%s: filter returned an error\n", name); + return false; + } + Cb.copy_to_host(); + if (!check(Ab, Bb, Cb, size, name)) { + return false; + } + double t = bench([&]() { filter(Ab.raw_buffer(), Bb.raw_buffer(), Cb.raw_buffer()); }, + [&]() { Cb.device_sync(); }); + printf(" Halide %-12s %9.0f GFlop/s\n", name, gflops(size, t)); -double gflops(int size, double seconds) { - return 2.0 * size * size * size / seconds * 1e-9; + // Running the filter left the operands on the device. Hand cublas the + // same ones, and somewhere of its own to put an answer that is only timed. + Cb_cublas.device_malloc(halide_cuda_device_interface()); + void *Ad = (void *)halide_cuda_get_device_ptr(nullptr, Ab.raw_buffer()); + void *Bd = (void *)halide_cuda_get_device_ptr(nullptr, Bb.raw_buffer()); + void *Cd = (void *)halide_cuda_get_device_ptr(nullptr, Cb_cublas.raw_buffer()); + t = bench([&]() { cublas(Ad, Bd, Cd, size); }, + []() { cudaDeviceSynchronize(); }); + printf(" cublas %-12s %9.0f GFlop/s\n", name, gflops(size, t)); + return true; } -// Time one call of the filter. Both implementations queue work -// asynchronously, and cublas does a heuristic lookup on the host for every -// call, so synchronizing per launch would measure that host work rather than -// letting it overlap with the GPU. Batch the launches instead, and sync once. -// `sync` has to match the launcher: Halide runs on its own CUDA context, so -// cudaDeviceSynchronize does not wait for it. -template -double bench_batched(F &&launch, S &&sync) { - const int batch = 5; - return Halide::Tools::benchmark(5, 1, - [&]() { - for (int i = 0; i < batch; i++) { - launch(); - } - sync(); - }) / - batch; -} +cublasHandle_t handle; } // namespace @@ -112,215 +181,82 @@ int main(int argc, char **argv) { size = atoi(argv[1]); } - { - Buffer A(size, size), B(size, size), C(size, size); - A.for_each_value([](float &v) { v = (float)((rand() & 3) - 1); }); - B.for_each_value([](float &v) { v = (float)((rand() & 3) - 1); }); - A.set_host_dirty(); - B.set_host_dirty(); - mat_mul(A, B, C); - C.copy_to_host(); - if (!check(A, B, C, size, "float")) { - return 1; - } + cublasCreate(&handle); + static float alpha = 1.0f, beta = 1.0f; + int failures = 0; - double t = bench_batched([&]() { mat_mul(A, B, C); }, - [&]() { C.device_sync(); }); - printf("Halide float: %f s (%.1f GFlop/s)\n", t, gflops(size, t)); - } + // A gemm at one pair of types, given how cublas should read the buffers. + auto gemm_ex = [](cudaDataType at, cudaDataType ct, cublasComputeType_t comp, + const void *al, const void *be) { + return [=](void *A, void *B, void *C, int n) { + cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, n, n, n, al, + A, at, n, B, at, n, be, C, ct, n, comp, + CUBLAS_GEMM_DEFAULT); + }; + }; + + failures += !row( + "f32 -> f32", size, halide_type_t(halide_type_float, 32), + halide_type_t(halide_type_float, 32), 4, + [](halide_buffer_t *A, halide_buffer_t *B, halide_buffer_t *C) { + return mat_mul(A, B, C); + }, + [](void *A, void *B, void *C, int n) { + cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, n, n, n, &alpha, + (const float *)A, n, (const float *)B, n, &beta, + (float *)C, n); + }); - // The half variant is scheduled onto the tensor cores. if (ver < 70) { printf("[SKIP] Tensor cores require compute capability 7.0 or above; " "this system has %d.%d.\n", major, minor); } else { - // _Float16 rather than Halide::float16_t, so that this stays a - // runtime-only program that doesn't link the compiler. - Buffer<_Float16, 2> A(size, size), B(size, size); - Buffer C(size, size); - A.for_each_value([](_Float16 &v) { v = (_Float16)((rand() & 3) - 1); }); - B.for_each_value([](_Float16 &v) { v = (_Float16)((rand() & 3) - 1); }); - A.set_host_dirty(); - B.set_host_dirty(); - mat_mul_f16(A, B, C); - C.copy_to_host(); - if (!check(A, B, C, size, "half")) { - return 1; - } + failures += !row( + "f16 -> f32", size, halide_type_t(halide_type_float, 16), + halide_type_t(halide_type_float, 32), 4, + [](halide_buffer_t *A, halide_buffer_t *B, halide_buffer_t *C) { + return mat_mul_f16(A, B, C); + }, + gemm_ex(CUDA_R_16F, CUDA_R_32F, CUBLAS_COMPUTE_32F, &alpha, &beta)); - double t = bench_batched([&]() { mat_mul_f16(A, B, C); }, - [&]() { C.device_sync(); }); - printf("Halide half (tensor cores): %f s (%.1f GFlop/s)\n", - t, gflops(size, t)); - } + failures += !row( + "bf16 -> f32", size, halide_type_t(halide_type_bfloat, 16), + halide_type_t(halide_type_float, 32), 4, + [](halide_buffer_t *A, halide_buffer_t *B, halide_buffer_t *C) { + return mat_mul_bf16(A, B, C); + }, + gemm_ex(CUDA_R_16BF, CUDA_R_32F, CUBLAS_COMPUTE_32F, &alpha, &beta)); - // Half precision operands accumulated into half precision, which halves - // the registers the accumulator needs. The operands here are sparse zeros - // and ones, so the dot products stay small enough to be exact even in a - // half accumulator, whose integers run out at 2048. - if (ver >= 70) { - Buffer<_Float16, 2> A(size, size), B(size, size); - Buffer<_Float16, 2> C(size, size); - A.for_each_value([](_Float16 &v) { v = (_Float16)((rand() & 3) == 0); }); - B.for_each_value([](_Float16 &v) { v = (_Float16)((rand() & 3) == 0); }); - A.set_host_dirty(); - B.set_host_dirty(); - mat_mul_f16_acc16(A, B, C); - C.copy_to_host(); - if (!check(A, B, C, size, "half into half")) { - return 1; - } - double t = bench_batched([&]() { mat_mul_f16_acc16(A, B, C); }, - [&]() { C.device_sync(); }); - printf("Halide half into half (tensor cores): %f s (%.1f GFlop/s)\n", - t, gflops(size, t)); - } + // Zeros and ones, so that the dot products stay under 2048, the + // largest integer half precision represents exactly. + static __half halpha = __float2half(1.f), hbeta = __float2half(1.f); + failures += !row( + "f16 -> f16", size, halide_type_t(halide_type_float, 16), + halide_type_t(halide_type_float, 16), 2, + [](halide_buffer_t *A, halide_buffer_t *B, halide_buffer_t *C) { + return mat_mul_f16_acc16(A, B, C); + }, + gemm_ex(CUDA_R_16F, CUDA_R_16F, CUBLAS_COMPUTE_16F, &halpha, &hbeta)); - // The other operand types the tensor cores multiply. Brain floats - // accumulate into single precision like halves do, and eight-bit integers - // into 32-bit ones, which is the interesting case for imaging. - if (ver >= 80) { - { - const halide_type_t bf16(halide_type_bfloat, 16); - Buffer A(bf16, size, size), B(bf16, size, size); - Buffer C(size, size); - // The buffer carries its type at runtime, so index the raw - // storage rather than going through a typed view. - uint16_t *Ap = (uint16_t *)A.data(), *Bp = (uint16_t *)B.data(); - auto Af = [&](int i, int j) { return from_bf16(Ap[j * size + i]); }; - auto Bf = [&](int i, int j) { return from_bf16(Bp[j * size + i]); }; - for (int i = 0; i < size * size; i++) { - Ap[i] = to_bf16((float)((rand() & 3) - 1)); - Bp[i] = to_bf16((float)((rand() & 3) - 1)); - } - A.set_host_dirty(); - B.set_host_dirty(); - mat_mul_bf16(A, B, C); - C.copy_to_host(); - for (int y = 0; y < size; y += 97) { - for (int x = 0; x < size; x += 89) { - double correct = 0; - for (int k = 0; k < size; k++) { - correct += (double)Af(x, k) * (double)Bf(k, y); - } - if ((double)C(x, y) != correct) { - printf("bfloat: bad result at %d %d: %f != %f\n", - x, y, (double)C(x, y), correct); - return 1; - } - } - } - double t = bench_batched([&]() { mat_mul_bf16(A, B, C); }, - [&]() { C.device_sync(); }); - printf("Halide bfloat (tensor cores): %f s (%.1f GFlop/s)\n", - t, gflops(size, t)); - } - { - Buffer A(size, size), B(size, size); - Buffer C(size, size); - A.for_each_value([](uint8_t &v) { v = (uint8_t)(rand() & 3); }); - B.for_each_value([](uint8_t &v) { v = (uint8_t)(rand() & 3); }); - A.set_host_dirty(); - B.set_host_dirty(); - mat_mul_u8(A, B, C); - C.copy_to_host(); - if (!check(A, B, C, size, "uint8")) { - return 1; - } - double t = bench_batched([&]() { mat_mul_u8(A, B, C); }, - [&]() { C.device_sync(); }); - printf("Halide uint8 (tensor cores): %f s (%.1f GFlop/s)\n", - t, gflops(size, t)); - } + // cublas takes signed bytes where the Halide variant takes unsigned + // ones. The hardware runs both at the same rate, and these values are + // small enough to mean the same thing either way. + static int32_t ialpha = 1, ibeta = 1; + failures += !row( + "u8 -> i32", size, halide_type_t(halide_type_uint, 8), + halide_type_t(halide_type_int, 32), 4, + [](halide_buffer_t *A, halide_buffer_t *B, halide_buffer_t *C) { + return mat_mul_u8(A, B, C); + }, + gemm_ex(CUDA_R_8I, CUDA_R_32I, CUBLAS_COMPUTE_32I, &ialpha, &ibeta)); } - // Benchmark cublas for reference, at both precisions. The half precision - // one accumulates in single precision, matching what the Halide pipeline - // does, so the two are comparable. -#ifdef _MSC_VER - // https://github.com/halide/Halide/issues/5053 - printf("Skipping cublas on Windows; see https://github.com/halide/Halide/issues/5053\n"); -#else - { - void *A, *B, *C; - cudaMalloc(&A, (size_t)size * size * 4); - cudaMalloc(&B, (size_t)size * size * 4); - cudaMalloc(&C, (size_t)size * size * 4); - // Touch the memory before timing anything, so that no part of the - // benchmark pays for faulting it in, and so that the operands are - // definite values rather than whatever was there. This byte pattern is - // a normal number read either as float or as half, which matters - // because denormals can be slow. - cudaMemset(A, 0x3c, (size_t)size * size * 4); - cudaMemset(B, 0x3c, (size_t)size * size * 4); - cudaMemset(C, 0, (size_t)size * size * 4); - cublasHandle_t handle; - cublasCreate(&handle); - float alpha = 1.0f, beta = 1.0f; - - double t = bench_batched([&]() { cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, - size, size, size, &alpha, (const float *)A, size, - (const float *)B, size, &beta, (float *)C, size); }, - []() { cudaDeviceSynchronize(); }); - printf("cublas float: %f s (%.1f GFlop/s)\n", t, gflops(size, t)); - - if (ver >= 70) { - // Half precision operands into a single precision accumulator, - // which is what the tensor cores do natively. - t = bench_batched([&]() { cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, - size, size, size, &alpha, - A, CUDA_R_16F, size, - B, CUDA_R_16F, size, &beta, - C, CUDA_R_32F, size, - CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT); }, - []() { cudaDeviceSynchronize(); }); - printf("cublas half: %f s (%.1f GFlop/s)\n", t, gflops(size, t)); - - // Brain floats, also into a single precision accumulator. - t = bench_batched([&]() { cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, - size, size, size, &alpha, - A, CUDA_R_16BF, size, - B, CUDA_R_16BF, size, &beta, - C, CUDA_R_32F, size, - CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT); }, - []() { cudaDeviceSynchronize(); }); - printf("cublas bfloat: %f s (%.1f GFlop/s)\n", t, gflops(size, t)); - - // Eight-bit integers into a 32-bit accumulator. cublas takes - // signed operands here where the Halide variant above takes - // unsigned ones; the hardware runs both at the same rate. - int32_t ialpha = 1, ibeta = 1; - t = bench_batched([&]() { cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, - size, size, size, &ialpha, - A, CUDA_R_8I, size, - B, CUDA_R_8I, size, &ibeta, - C, CUDA_R_32I, size, - CUBLAS_COMPUTE_32I, CUBLAS_GEMM_DEFAULT); }, - []() { cudaDeviceSynchronize(); }); - printf("cublas int8: %f s (%.1f GFlop/s)\n", t, gflops(size, t)); - - // Halves into a half accumulator, which is what the half output - // variant above does. - __half halpha = __float2half(1.f), hbeta = __float2half(1.f); - t = bench_batched([&]() { cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, - size, size, size, &halpha, - A, CUDA_R_16F, size, - B, CUDA_R_16F, size, &hbeta, - C, CUDA_R_16F, size, - CUBLAS_COMPUTE_16F, CUBLAS_GEMM_DEFAULT); }, - []() { cudaDeviceSynchronize(); }); - printf("cublas half into half: %f s (%.1f GFlop/s)\n", - t, gflops(size, t)); - } - - cudaFree(A); - cudaFree(B); - cudaFree(C); - cublasDestroy(handle); + cublasDestroy(handle); + if (failures) { + printf("%d configuration(s) failed\n", failures); + return 1; } -#endif - printf("Success!\n"); return 0; } From 77022a56893bb9c73bf6f47fe7a5642d5342027e Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 4 Aug 2026 10:02:49 -0700 Subject: [PATCH 50/59] Use Halide's float16 types in the runner, and link libHalide for them The runner was carrying its own element get and set for buffers of a type with no C++ equivalent, which is really just bfloat16. Halide already has usable float16_t and bfloat16_t, and Float16.h already declares halide_type_of for both, so the buffers can be plain typed ones and the dispatch goes away. The cost is linking libHalide into a program that otherwise only needs the generated code, which a TODO records. Float16.h depends on nothing but HalideRuntime.h and its implementation touches nothing in the compiler, so making those definitions inline and shipping the header - or moving the types to the runtime - would remove the need. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/CMakeLists.txt | 2 +- apps/cuda_mat_mul/Makefile | 2 +- apps/cuda_mat_mul/mat_mul_generator.cpp | 27 ++--- apps/cuda_mat_mul/runner.cpp | 134 +++++++++--------------- 4 files changed, 67 insertions(+), 98 deletions(-) diff --git a/apps/cuda_mat_mul/CMakeLists.txt b/apps/cuda_mat_mul/CMakeLists.txt index 23f2f56edc77..5995bbfd6100 100644 --- a/apps/cuda_mat_mul/CMakeLists.txt +++ b/apps/cuda_mat_mul/CMakeLists.txt @@ -53,7 +53,7 @@ add_halide_library(mat_mul_u8 FROM mat_mul.generator # Main executable add_executable(runner runner.cpp) -target_link_libraries(runner PRIVATE mat_mul mat_mul_f16 mat_mul_f16_acc16 mat_mul_bf16 mat_mul_u8 Halide::Tools CUDA::cudart CUDA::cublas) +target_link_libraries(runner PRIVATE mat_mul mat_mul_f16 mat_mul_f16_acc16 mat_mul_bf16 mat_mul_u8 Halide::Halide Halide::Tools CUDA::cudart CUDA::cublas) # Test that the app actually works! add_test(NAME mat_mul COMMAND runner) diff --git a/apps/cuda_mat_mul/Makefile b/apps/cuda_mat_mul/Makefile index d94b24d5133c..5f25604d9a02 100644 --- a/apps/cuda_mat_mul/Makefile +++ b/apps/cuda_mat_mul/Makefile @@ -52,7 +52,7 @@ $(BIN)/%/runner: runner.cpp $(BIN)/%/mat_mul.a $(BIN)/%/mat_mul_f16.a \ $(BIN)/%/mat_mul_bf16.a $(BIN)/%/mat_mul_u8.a \ $(BIN)/%/mat_mul_f16_acc16.a @mkdir -p $(@D) - $(CXX) $(CXXFLAGS) -I$(BIN)/$* -Wall $^ -o $@ $(LDFLAGS) -lcudart -lcublas + $(CXX) $(CXXFLAGS) -I$(BIN)/$* -Wall $^ -o $@ $(LDFLAGS) $(LIBHALIDE_LDFLAGS) -lcudart -lcublas test: $(BIN)/$(HL_TARGET)/runner $^ $(MATRIX_SIZE) diff --git a/apps/cuda_mat_mul/mat_mul_generator.cpp b/apps/cuda_mat_mul/mat_mul_generator.cpp index 0435280ea43e..db86d17225a4 100644 --- a/apps/cuda_mat_mul/mat_mul_generator.cpp +++ b/apps/cuda_mat_mul/mat_mul_generator.cpp @@ -26,30 +26,31 @@ void set_alignment_and_bounds(OutputImageParam p, int size) { // accumulator are less work. // // 1024 2048 4096 ceiling -// Halide f32 6878 10294 7298 25960 -// cublas f32 14503 16574 17449 +// Halide f32 6937 10199 7306 25960 +// cublas f32 14688 16782 17545 // -// Halide f16 -> f32 40070 46671 48888 51541 -// cublas f16 -> f32 41658 49069 50440 +// Halide f16 -> f32 40396 47019 49314 51541 +// cublas f16 -> f32 42324 49596 51212 // -// Halide bf16 -> f32 40025 46681 48692 51541 -// cublas bf16 -> f32 41664 49054 50437 +// Halide bf16 -> f32 40347 47007 49129 51541 +// cublas bf16 -> f32 42315 49601 51179 // -// Halide f16 -> f16 60349 86502 86599 99626 -// cublas f16 -> f16 69221 75073 87073 +// Halide f16 -> f16 60745 87075 86973 99626 +// cublas f16 -> f16 73466 76443 90182 // -// Halide u8 -> i32 62869 82391 89641 100650 -// cublas s8 -> i32 107868 122203 129970 +// Halide u8 -> i32 63588 83223 90390 100650 +// cublas u8 -> i32 120109 128990 140195 // // The tensor core ceilings are measured, by issuing wmma instructions back to // back out of registers. The float one is 36 SMs times the 2817 MHz this part // averages while benchmarking times the 256 flops per SM per clock the cuda // cores do. // -// The tensor core schedules reach 87% to 95% of their ceilings, and match -// cublas at f16 -> f16. Two rows fall short for reasons outside the schedule. +// The tensor core schedules reach 87% to 95% of their ceilings, and beat +// cublas at f16 -> f16 at 2048. Two rows fall short for reasons outside the +// schedule. // -// At eight bits cublas is 29% past the ceiling, so it is not using wmma, which +// At eight bits cublas is 39% past the ceiling, so it is not using wmma, which // multiplies bytes no faster than it multiplies halves into halves. The mma // instructions reach 188355 GOP/s at the same shape, 1.87x, for the same 8192 // ops per instruction. Reaching that needs mma's fragment layout, which this diff --git a/apps/cuda_mat_mul/runner.cpp b/apps/cuda_mat_mul/runner.cpp index 6ab9db5201f3..c5df75a578b3 100644 --- a/apps/cuda_mat_mul/runner.cpp +++ b/apps/cuda_mat_mul/runner.cpp @@ -1,9 +1,16 @@ +// TODO: this links libHalide purely to get Halide::float16_t and +// Halide::bfloat16_t, which an AOT program should not have to do. Float16.h +// depends on nothing but HalideRuntime.h and already declares halide_type_of +// for both types, but its bodies live in Float16.cpp inside libHalide, and the +// header is not distributed. Making those definitions inline and shipping the +// header, or moving the types to the runtime, would let this go back to +// linking only the generated code. +#include "Halide.h" + #include "HalideBuffer.h" #include "HalideRuntimeCuda.h" #include "halide_benchmark.h" -#include #include -#include #include #include #include @@ -14,6 +21,8 @@ #include "mat_mul_f16_acc16.h" #include "mat_mul_u8.h" +using Halide::bfloat16_t; +using Halide::float16_t; using Halide::Runtime::Buffer; namespace { @@ -22,7 +31,7 @@ namespace { // operand and accumulator types, and each is compared against cublas doing the // same thing. Both are handed the same device buffers - the Halide side fills // them and copies them down, and cublas is given the pointers out of them - so -// the two see identical data. +// the two see identical data, and both answers are checked. // // The operands are small integers so that every dot product is exact in every // accumulator here, which lets the results be checked for equality rather than @@ -54,71 +63,26 @@ double gflops(int size, double seconds) { return 2.0 * size * size * size / seconds * 1e-9; } -// Read or write an element of a buffer of any of the types used here, so that -// the ones with no C++ equivalent are handled too. -double element(const halide_buffer_t *b, size_t i) { - halide_type_t t = b->type; - if (t == halide_type_t(halide_type_float, 16)) { - return (double)((const _Float16 *)b->host)[i]; - } else if (t == halide_type_t(halide_type_bfloat, 16)) { - // A bfloat is the top half of a float. - uint32_t bits = (uint32_t)((const uint16_t *)b->host)[i] << 16; - float f; - memcpy(&f, &bits, 4); - return f; - } else if (t == halide_type_t(halide_type_float, 32)) { - return ((const float *)b->host)[i]; - } else if (t == halide_type_t(halide_type_uint, 8)) { - return ((const uint8_t *)b->host)[i]; - } else if (t == halide_type_t(halide_type_int, 32)) { - return ((const int32_t *)b->host)[i]; - } - fprintf(stderr, "unhandled buffer type\n"); - exit(1); -} - -void set_element(halide_buffer_t *b, size_t i, int v) { - halide_type_t t = b->type; - if (t == halide_type_t(halide_type_float, 16)) { - ((_Float16 *)b->host)[i] = (_Float16)v; - } else if (t == halide_type_t(halide_type_bfloat, 16)) { - float f = (float)v; - uint32_t bits; - memcpy(&bits, &f, 4); - ((uint16_t *)b->host)[i] = (uint16_t)(bits >> 16); - } else if (t == halide_type_t(halide_type_float, 32)) { - ((float *)b->host)[i] = (float)v; - } else if (t == halide_type_t(halide_type_uint, 8)) { - ((uint8_t *)b->host)[i] = (uint8_t)v; - } else { - fprintf(stderr, "unhandled operand type\n"); - exit(1); - } -} - -void fill_ints(Buffer &b, int modulus) { - size_t n = (size_t)b.width() * b.height(); - for (size_t i = 0; i < n; i++) { - set_element(b.raw_buffer(), i, rand() % modulus); - } +template +void fill_ints(Buffer &b, int modulus) { + b.for_each_value([&](T &v) { v = (T)(rand() % modulus); }); b.set_host_dirty(); } // Spot check on strides coprime with the tile sizes, so the samples land at // varying offsets within a tile. -bool check(Buffer &Ab, Buffer &Bb, Buffer &Cb, +template +bool check(Buffer &Ab, Buffer &Bb, Buffer &Cb, int size, const char *name) { for (int y = 0; y < size; y += 97) { for (int x = 0; x < size; x += 89) { double correct = 0; for (int k = 0; k < size; k++) { - correct += element(Ab.raw_buffer(), (size_t)k * size + x) * - element(Bb.raw_buffer(), (size_t)y * size + k); + correct += (double)Ab(x, k) * (double)Bb(k, y); } - double got = element(Cb.raw_buffer(), (size_t)y * size + x); - if (got != correct) { + if ((double)Cb(x, y) != correct) { printf("%s: bad result at %d %d: %f != %f\n", - name, x, y, got, correct); + name, x, y, (double)Cb(x, y), correct); return false; } } @@ -128,11 +92,10 @@ bool check(Buffer &Ab, Buffer &Bb, Buffer &Cb, // One row of the table: run the Halide filter, check it, time it, then time // cublas on the same device buffers. -template -bool row(const char *name, int size, halide_type_t a_type, halide_type_t c_type, - int modulus, Filter filter, Cublas cublas) { - Buffer Ab(a_type, size, size), Bb(a_type, size, size); - Buffer Cb(c_type, size, size), Cb_cublas(c_type, size, size); +template +bool row(const char *name, int size, int modulus, Filter filter, Cublas cublas) { + Buffer Ab(size, size), Bb(size, size); + Buffer Cb(size, size); fill_ints(Ab, modulus); fill_ints(Bb, modulus); @@ -148,12 +111,22 @@ bool row(const char *name, int size, halide_type_t a_type, halide_type_t c_type, [&]() { Cb.device_sync(); }); printf(" Halide %-12s %9.0f GFlop/s\n", name, gflops(size, t)); - // Running the filter left the operands on the device. Hand cublas the - // same ones, and somewhere of its own to put an answer that is only timed. - Cb_cublas.device_malloc(halide_cuda_device_interface()); + // Running the filter left everything on the device, so cublas gets the + // same buffers, output included. Halide's dense-first-dimension layout is + // what cublas calls column major with a leading dimension of the size, so + // no transposing is needed, and the answer is checked the same way to + // confirm that. void *Ad = (void *)halide_cuda_get_device_ptr(nullptr, Ab.raw_buffer()); void *Bd = (void *)halide_cuda_get_device_ptr(nullptr, Bb.raw_buffer()); - void *Cd = (void *)halide_cuda_get_device_ptr(nullptr, Cb_cublas.raw_buffer()); + void *Cd = (void *)halide_cuda_get_device_ptr(nullptr, Cb.raw_buffer()); + cublas(Ad, Bd, Cd, size); + cudaDeviceSynchronize(); + Cb.set_device_dirty(); + Cb.copy_to_host(); + if (!check(Ab, Bb, Cb, size, name)) { + printf(" (that was cublas, not Halide)\n"); + return false; + } t = bench([&]() { cublas(Ad, Bd, Cd, size); }, []() { cudaDeviceSynchronize(); }); printf(" cublas %-12s %9.0f GFlop/s\n", name, gflops(size, t)); @@ -182,7 +155,7 @@ int main(int argc, char **argv) { } cublasCreate(&handle); - static float alpha = 1.0f, beta = 1.0f; + static float alpha = 1.0f, beta = 0.0f; int failures = 0; // A gemm at one pair of types, given how cublas should read the buffers. @@ -195,9 +168,8 @@ int main(int argc, char **argv) { }; }; - failures += !row( - "f32 -> f32", size, halide_type_t(halide_type_float, 32), - halide_type_t(halide_type_float, 32), 4, + failures += !row( + "f32 -> f32", size, 4, [](halide_buffer_t *A, halide_buffer_t *B, halide_buffer_t *C) { return mat_mul(A, B, C); }, @@ -212,17 +184,15 @@ int main(int argc, char **argv) { "this system has %d.%d.\n", major, minor); } else { - failures += !row( - "f16 -> f32", size, halide_type_t(halide_type_float, 16), - halide_type_t(halide_type_float, 32), 4, + failures += !row( + "f16 -> f32", size, 4, [](halide_buffer_t *A, halide_buffer_t *B, halide_buffer_t *C) { return mat_mul_f16(A, B, C); }, gemm_ex(CUDA_R_16F, CUDA_R_32F, CUBLAS_COMPUTE_32F, &alpha, &beta)); - failures += !row( - "bf16 -> f32", size, halide_type_t(halide_type_bfloat, 16), - halide_type_t(halide_type_float, 32), 4, + failures += !row( + "bf16 -> f32", size, 4, [](halide_buffer_t *A, halide_buffer_t *B, halide_buffer_t *C) { return mat_mul_bf16(A, B, C); }, @@ -230,10 +200,9 @@ int main(int argc, char **argv) { // Zeros and ones, so that the dot products stay under 2048, the // largest integer half precision represents exactly. - static __half halpha = __float2half(1.f), hbeta = __float2half(1.f); - failures += !row( - "f16 -> f16", size, halide_type_t(halide_type_float, 16), - halide_type_t(halide_type_float, 16), 2, + static __half halpha = __float2half(1.f), hbeta = __float2half(0.f); + failures += !row( + "f16 -> f16", size, 2, [](halide_buffer_t *A, halide_buffer_t *B, halide_buffer_t *C) { return mat_mul_f16_acc16(A, B, C); }, @@ -242,10 +211,9 @@ int main(int argc, char **argv) { // cublas takes signed bytes where the Halide variant takes unsigned // ones. The hardware runs both at the same rate, and these values are // small enough to mean the same thing either way. - static int32_t ialpha = 1, ibeta = 1; - failures += !row( - "u8 -> i32", size, halide_type_t(halide_type_uint, 8), - halide_type_t(halide_type_int, 32), 4, + static int32_t ialpha = 1, ibeta = 0; + failures += !row( + "u8 -> i32", size, 4, [](halide_buffer_t *A, halide_buffer_t *B, halide_buffer_t *C) { return mat_mul_u8(A, B, C); }, From 81be1dc780a68ecebf66ff7e4442103be8ac6fd7 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 4 Aug 2026 10:05:26 -0700 Subject: [PATCH 51/59] Record that cublas is asked to write rather than accumulate Its beta was one, so it read the output and added to it, which the filters here do not do. Zero makes the two do the same work, and is worth up to ten percent to cublas at the widths where the output is large relative to the operands. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/mat_mul_generator.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/apps/cuda_mat_mul/mat_mul_generator.cpp b/apps/cuda_mat_mul/mat_mul_generator.cpp index db86d17225a4..a84cb3afa682 100644 --- a/apps/cuda_mat_mul/mat_mul_generator.cpp +++ b/apps/cuda_mat_mul/mat_mul_generator.cpp @@ -44,9 +44,10 @@ void set_alignment_and_bounds(OutputImageParam p, int size) { // The tensor core ceilings are measured, by issuing wmma instructions back to // back out of registers. The float one is 36 SMs times the 2817 MHz this part // averages while benchmarking times the 256 flops per SM per clock the cuda -// cores do. +// cores do. cublas is asked for a beta of zero, so that it writes its output +// rather than accumulating onto it, which is what the filters here do. // -// The tensor core schedules reach 87% to 95% of their ceilings, and beat +// The tensor core schedules reach 87% to 96% of their ceilings, and beat // cublas at f16 -> f16 at 2048. Two rows fall short for reasons outside the // schedule. // From b437f75570bc26d8df02ec682dc0cb41c870037d Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 4 Aug 2026 16:40:18 -0700 Subject: [PATCH 52/59] Unify tile storage as MemoryType::Tile and add tile scheduling directives MemoryType::AMXTile and MemoryType::WMMAFragment both describe storage for a matrix tile in whatever form a target's matrix unit keeps one, so merge them into MemoryType::Tile. AMXTile remains as a deprecated alias. WMMAFragment was never released, so it is simply removed. Add tile_init, tile_load, tile_store and tile_matmul to Stage and Func. They are sugar over the existing directives: each reorders the dimensions that make up a tile to be innermost and vectorizes them, tile_matmul additionally marking the stage atomic so that a reduction dimension can be vectorized. Everything but tile_store also sets the memory type, since those are the ones that produce a tile. The order of a tile's dimensions is fixed by the instruction rather than free to schedule, so imposing it here lets schedules drop it: reorder only permutes the dimensions it names among the positions they already hold, so a file-local helper names every dimension to move the tile ones inwards. The same four directives now describe both targets' matrix units, lowering to tilezero/tileloadd/tdpbf16ps/tilestored on x86 and to the corresponding wmma operations on CUDA. Rename test/correctness/tiled_matmul.cpp to amx_matmul.cpp to say which matrix unit it covers, and register wmma_matmul.cpp in the correctness CMakeLists, where it was missing. In apps/cuda_mat_mul, express the paired splits as tile calls, drop the reorder arguments the directives now supply, and give the tensor core dimensions the names mmx, mmy and mmr so that Vars are no longer spelled like RVars. Co-Authored-By: Claude Opus 5 --- apps/cuda_mat_mul/mat_mul_generator.cpp | 80 +++++------- apps/tensorcore_resize/resize_generator.cpp | 4 +- .../src/halide/halide_/PyEnums.cpp | 5 +- src/CanonicalizeGPUVars.cpp | 4 +- src/CodeGen_PTX_Dev.cpp | 2 +- src/CodeGen_X86.cpp | 8 +- src/Deserialization.cpp | 6 +- src/Expr.h | 31 ++--- src/ExtractTileOperations.cpp | 8 +- src/ExtractTileOperations.h | 2 +- src/ExtractWMMAOperations.cpp | 12 +- src/ExtractWMMAOperations.h | 2 +- src/Func.cpp | 72 +++++++++++ src/Func.h | 52 ++++++++ src/FuseGPUThreadLoops.cpp | 10 +- src/IRPrinter.cpp | 7 +- src/LowerWarpShuffles.cpp | 2 +- src/Serialization.cpp | 6 +- src/halide_ir.fbs | 3 +- test/correctness/CMakeLists.txt | 3 +- .../{tiled_matmul.cpp => amx_matmul.cpp} | 51 ++++---- test/correctness/tiled_matmul_errors.cpp | 18 +-- test/correctness/wmma_matmul.cpp | 118 +++++++----------- test/performance/tiled_matmul.cpp | 4 +- 24 files changed, 292 insertions(+), 218 deletions(-) rename test/correctness/{tiled_matmul.cpp => amx_matmul.cpp} (91%) diff --git a/apps/cuda_mat_mul/mat_mul_generator.cpp b/apps/cuda_mat_mul/mat_mul_generator.cpp index a84cb3afa682..02b9ad95b45e 100644 --- a/apps/cuda_mat_mul/mat_mul_generator.cpp +++ b/apps/cuda_mat_mul/mat_mul_generator.cpp @@ -62,7 +62,7 @@ void set_alignment_and_bounds(OutputImageParam p, int size) { // them, living across it. Such an accumulator lands at block level, where a // Register allocation is sized for the whole block tile and spills. So it // stages per thread and shares nothing. The tensor core schedules only escape -// this because a WMMAFragment allocation at block level is already per-lane. +// this because a tile allocation at block level is already per-lane. // class MatMul : public Halide::Generator { public: @@ -230,57 +230,45 @@ class MatMul : public Halide::Generator { const int pa = pad_a ? (int)pad_a : 16 / A.type().bytes(); const int pb = pad_b ? (int)pad_b : 16 / A.type().bytes(); - Var xi("xi"), yi("yi"), xt("xt"), yt("yt"), mmxi("mmxi"), mmyi("mmyi"); - Var xw("xw"), yw("yw"), rxi("rxi"), ryi("ryi"); - RVar ro("ro"), ri("ri"), rri("rri"); + // out and prod are tiled the same way: blocks of warps, each warp + // holding several tensor core tiles. mmx, mmy and mmr are the + // dimensions of one tensor core operation. + Var xi("xi"), yi("yi"), xw("xw"), yw("yw"); + Var mmx("mmx"), mmy("mmy"); + RVar ro("ro"), ri("ri"), mmr("mmr"); - out.split(x, x, xi, block_x) - .split(xi, xt, xi, tile * tx) - .split(xi, xi, mmxi, tile) - .split(y, y, yi, block_y) - .split(yi, yt, yi, tile * ty) - .split(yi, yi, mmyi, tile) + out.tile(x, y, xi, yi, block_x, block_y) + .tile(xi, yi, xw, yw, xi, yi, tile * tx, tile * ty) + .tile(xi, yi, mmx, mmy, tile, tile) .gpu_blocks(x, y) - .gpu_threads(xt, yt) - .reorder(mmxi, mmyi, xi, yi, xt, yt, x, y) + .gpu_threads(xw, yw) .unroll(xi) .unroll(yi) - .vectorize(mmxi) - .vectorize(mmyi); + .tile_store(mmx, mmy); // The accumulators live in tensor core registers for the whole // reduction, and are written out to memory once at the end. They sit // at block level so that the reduction loop can be above the loop over // warps, which lets every warp share one staged panel. prod.compute_at(out, x) - .store_in(MemoryType::WMMAFragment) - .split(x, xw, xi, tile * tx) - .split(xi, xi, rxi, tile) - .split(y, yw, yi, tile * ty) - .split(yi, yi, ryi, tile) - .reorder(rxi, ryi, xi, yi, xw, yw) + .tile(x, y, xw, yw, xi, yi, tile * tx, tile * ty) + .tile(xi, yi, mmx, mmy, tile, tile) .gpu_threads(xw, yw) - .vectorize(rxi) - .vectorize(ryi) + .tile_init(mmx, mmy) .unroll(xi) .unroll(yi); prod.update() .split(r, ro, ri, br) - .split(x, xw, xi, tile * tx) - .split(xi, xi, rxi, tile) - .split(y, yw, yi, tile * ty) - .split(yi, yi, ryi, tile) - .split(ri, ri, rri, tile) - .reorder(rri, rxi, ryi, xi, yi, ri, xw, yw, ro) + .split(ri, ri, mmr, tile) + .tile(x, y, xw, yw, xi, yi, tile * tx, tile * ty) + .tile(xi, yi, mmx, mmy, tile, tile) + .reorder(xi, yi, ri, xw, yw, ro) .gpu_threads(xw, yw) .unroll(xi) .unroll(yi) .unroll(ri) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi); + .tile_matmul(mmr, mmx, mmy); // Stage the operand panels into shared memory once per reduction step, // to be shared by every warp in the block. Each thread moves sixteen @@ -290,36 +278,36 @@ class MatMul : public Halide::Generator { // Each thread moves sixteen bytes, the widest asynchronous copy the // hardware has. How many elements that is depends on the operand type. const int vec = 16 / A.type().bytes(); - Var rro("rro"), rrv("rrv"), xxo("xxo"), xxi("xxi"); - Var t("t"), ti("ti"), tw("tw"), tw2("tw2"), to("to"); + Var ko("ko"), kv("kv"), xo("xo"), xv("xv"); + Var t("t"), ti("ti"), to("to"); // B.in() is dense in the reduction dimension, which is its _0. B.in() .compute_at(prod, ro) .store_in(MemoryType::GPUSharedAsync) .align_storage(_0, br + pa) - .split(_0, rro, rrv, vec) - .fuse(rro, _1, t) + .split(_0, ko, kv, vec) + .fuse(ko, _1, t) .split(t, t, ti, 32) - .split(t, t, tw, wx) - .split(t, to, tw2, wy) + .split(t, t, xw, wx) + .split(t, to, yw, wy) .gpu_lanes(ti) - .gpu_threads(tw, tw2) - .vectorize(rrv); + .gpu_threads(xw, yw) + .vectorize(kv); // A.in() is dense in x, which is its _0. A.in() .compute_at(prod, ro) .store_in(MemoryType::GPUSharedAsync) .align_storage(_0, block_x + pb) - .split(_0, xxo, xxi, vec) - .fuse(xxo, _1, t) + .split(_0, xo, xv, vec) + .fuse(xo, _1, t) .split(t, t, ti, 32) - .split(t, t, tw, wx) - .split(t, to, tw2, wy) + .split(t, t, xw, wx) + .split(t, to, yw, wy) .gpu_lanes(ti) - .gpu_threads(tw, tw2) - .vectorize(xxi); + .gpu_threads(xw, yw) + .vectorize(xv); } Var x{"x"}, y{"y"}; diff --git a/apps/tensorcore_resize/resize_generator.cpp b/apps/tensorcore_resize/resize_generator.cpp index ebc9ae49dbc3..433ec74ec333 100644 --- a/apps/tensorcore_resize/resize_generator.cpp +++ b/apps/tensorcore_resize/resize_generator.cpp @@ -247,7 +247,7 @@ class Resize : public Halide::Generator { // An 8x32 tile of accumulator, reducing 16 taps at a time. resized_y.compute_at(resized_y.in(), xio) - .store_in(MemoryType::WMMAFragment) + .store_in(MemoryType::Tile) .unroll(c) .vectorize(x, 32) .unroll(x) @@ -279,7 +279,7 @@ class Resize : public Halide::Generator { RVar ri("ri"), ro("ro"); resized_x - .store_in(MemoryType::WMMAFragment) + .store_in(MemoryType::Tile) .compute_at(resized_x.in(), c) .vectorize(x) .vectorize(y) diff --git a/python_bindings/src/halide/halide_/PyEnums.cpp b/python_bindings/src/halide/halide_/PyEnums.cpp index e929c1bfe928..28c9b63c9057 100644 --- a/python_bindings/src/halide/halide_/PyEnums.cpp +++ b/python_bindings/src/halide/halide_/PyEnums.cpp @@ -50,9 +50,10 @@ void define_enums(py::module &m) { .value("GPUTexture", MemoryType::GPUTexture) .value("LockedCache", MemoryType::LockedCache) .value("VTCM", MemoryType::VTCM) - .value("AMXTile", MemoryType::AMXTile) + .value("Tile", MemoryType::Tile) .value("GPUSharedAsync", MemoryType::GPUSharedAsync) - .value("WMMAFragment", MemoryType::WMMAFragment); + // Deprecated alias for Tile. + .value("AMXTile", MemoryType::Tile); py::enum_(m, "NameMangling") .value("Default", NameMangling::Default) diff --git a/src/CanonicalizeGPUVars.cpp b/src/CanonicalizeGPUVars.cpp index 41dee1c8111a..4e02304a37ac 100644 --- a/src/CanonicalizeGPUVars.cpp +++ b/src/CanonicalizeGPUVars.cpp @@ -82,7 +82,7 @@ class CountGPUBlocksThreads : public IRVisitor { // extract_wmma_operations will wrap the statements that touch this // allocation in loops over the lanes of a warp, so count it as a lane // dimension. - const bool wmma = op->memory_type == MemoryType::WMMAFragment; + const bool wmma = op->memory_type == MemoryType::Tile; int dl = wmma && !in_lanes; ScopedValue old_in_lanes(in_lanes, in_lanes || wmma); nl += dl; @@ -120,7 +120,7 @@ class CanonicalizeGPUVars : public IRMutator { Stmt visit(const Realize *op) override { ScopedValue old(in_wmma_alloc, in_wmma_alloc || - op->memory_type == MemoryType::WMMAFragment); + op->memory_type == MemoryType::Tile); return IRMutator::visit(op); } diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 82aa3c993324..6d0bc63dcad8 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -528,7 +528,7 @@ void CodeGen_PTX_Dev::codegen_wmma(const Call *op) { bool CodeGen_PTX_Dev::is_fragment_alloc(const std::string &name) { const MemoryType *t = alloc_memory_type.find(name); - return t && *t == MemoryType::WMMAFragment; + return t && *t == MemoryType::Tile; } llvm::Type *CodeGen_PTX_Dev::fragment_reg_type(Type t) { diff --git a/src/CodeGen_X86.cpp b/src/CodeGen_X86.cpp index f88b67a68c91..d8c6351b89c7 100644 --- a/src/CodeGen_X86.cpp +++ b/src/CodeGen_X86.cpp @@ -1648,9 +1648,9 @@ void CodeGen_X86::visit(const Allocate *op) { void CodeGen_X86::visit(const Load *op) { if (const auto *mt = mem_type.find(op->name)) { - if (*mt == MemoryType::AMXTile) { + if (*mt == MemoryType::Tile) { const Ramp *ramp = op->index.as(); - internal_assert(ramp) << "Expected AMXTile to have index ramp\n"; + internal_assert(ramp) << "Expected a tile to have index ramp\n"; Value *ptr = codegen_buffer_pointer(op->name, op->type, ramp->base); LoadInst *load = builder->CreateAlignedLoad(llvm_type_of(upgrade_type_for_storage(op->type)), ptr, llvm::Align(op->type.bytes())); add_tbaa_metadata(load, op->name, op->index); @@ -1663,11 +1663,11 @@ void CodeGen_X86::visit(const Load *op) { void CodeGen_X86::visit(const Store *op) { if (const auto *mt = mem_type.find(op->name)) { - if (*mt == MemoryType::AMXTile) { + if (*mt == MemoryType::Tile) { Value *val = codegen(op->value); Halide::Type value_type = op->value.type(); const Ramp *ramp = op->index.as(); - internal_assert(ramp) << "Expected AMXTile to have index ramp\n"; + internal_assert(ramp) << "Expected a tile to have index ramp\n"; Value *ptr = codegen_buffer_pointer(op->name, value_type, ramp->base); StoreInst *store = builder->CreateAlignedStore(val, ptr, llvm::Align(value_type.bytes())); add_tbaa_metadata(store, op->name, op->index); diff --git a/src/Deserialization.cpp b/src/Deserialization.cpp index 140bcf1ecaae..54bdfa5f631e 100644 --- a/src/Deserialization.cpp +++ b/src/Deserialization.cpp @@ -188,10 +188,8 @@ MemoryType Deserializer::deserialize_memory_type(Serialize::MemoryType memory_ty return MemoryType::LockedCache; case Serialize::MemoryType::VTCM: return MemoryType::VTCM; - case Serialize::MemoryType::AMXTile: - return MemoryType::AMXTile; - case Serialize::MemoryType::WMMAFragment: - return MemoryType::WMMAFragment; + case Serialize::MemoryType::Tile: + return MemoryType::Tile; default: user_error << "unknown memory type " << (int)memory_type << "\n"; return MemoryType::Auto; diff --git a/src/Expr.h b/src/Expr.h index 679c76e1c019..665148c247aa 100644 --- a/src/Expr.h +++ b/src/Expr.h @@ -402,9 +402,18 @@ enum class MemoryType { * on Hexagon */ VTCM, - /** AMX Tile register for X86. Any data that would be used in an AMX matrix - * multiplication must first be loaded into an AMX tile register. */ - AMXTile, + /** Storage for a matrix tile, in whatever form the target's matrix unit + * keeps one: an AMX tile register on x86, or an NVIDIA tensor core + * fragment striped across the registers of the 32 lanes of a warp. Either + * way the layout is not architecturally specified, so the only legal + * accesses are the whole-tile ones a matrix unit provides - filling a + * tile, loading one from memory, storing one back, and multiplying a pair + * of them into a third. Which of those a given tile takes part in, and so + * which role it plays in the multiply, follows from how it is used. + * + * Schedule these with tile_init, tile_load, tile_store and tile_matmul, + * which set this memory type where it is needed. */ + Tile, /** GPU shared memory, written by an asynchronous copy instruction that * moves data straight from global memory without routing it through @@ -414,17 +423,9 @@ enum class MemoryType { * no such instruction this is ordinary shared memory. */ GPUSharedAsync, - /** An NVIDIA tensor core matrix fragment. The storage is striped across - * the registers of the 32 lanes of a warp in a layout that is not - * architecturally specified, so the only legal accesses are the ones - * recognized by the WMMA lowering pass. Which of the three roles a - * fragment plays - the accumulator, or either operand of the multiply - - * follows from how it is used, and determines what those accesses are. - * An accumulator can be zero-initialized, initialized from a matrix in - * memory, accumulated into by a matrix multiply, and copied back out to - * memory. An operand can be filled from a matrix in memory and read by a - * matrix multiply. */ - WMMAFragment, + /** Deprecated alias for Tile, which covers the tile storage of every + * target's matrix unit rather than just x86's. */ + AMXTile = Tile, }; /** Whether a MemoryType places an allocation in GPU shared memory. */ @@ -439,7 +440,7 @@ inline bool is_gpu_shared(MemoryType t) { * dedicated lowering pass that requires the original 2D-shaped loads * and stores to remain intact. */ inline bool is_tile_memory_type(MemoryType t) { - return t == MemoryType::AMXTile || t == MemoryType::WMMAFragment; + return t == MemoryType::Tile; } namespace Internal { diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp index 26ab57d5621a..af05db944aef 100644 --- a/src/ExtractTileOperations.cpp +++ b/src/ExtractTileOperations.cpp @@ -393,7 +393,7 @@ class ExtractTileOperations : public IRMutator { int found_J = -1; int found_K = -1; - // An AMXTile allocation may represent multiple AMX accumulator + // A tile allocation may represent multiple AMX accumulator // registers as 2D sub-tiles. This map tracks those. std::vector amx_subtiles; @@ -404,7 +404,7 @@ class ExtractTileOperations : public IRMutator { } Stmt visit(const Allocate *op) override { - if (op->memory_type == MemoryType::AMXTile) { + if (op->memory_type == MemoryType::Tile) { user_assert( (op->type.is_int() && op->type.bits() == 32) || (op->type.is_float() && op->type.bits() == 32)) @@ -422,7 +422,7 @@ class ExtractTileOperations : public IRMutator { pass = 0; body = mutate(body); user_assert(found_I >= 0 && found_J >= 0 && found_K >= 0) - << op->name << " is stored in AMXTile memory, but no matrix multiply " + << op->name << " is stored in Tile memory, but no matrix multiply " << "operation was found that stores to it, so the shape of the tile " << "was unable to be determined.\n"; pass = 1; @@ -430,7 +430,7 @@ class ExtractTileOperations : public IRMutator { for (int i = 0; i < (int)amx_subtiles.size(); i++) { body = Allocate::make(amx_name + std::to_string(i), op->type.element_of(), - MemoryType::AMXTile, {256}, const_true(), body); + MemoryType::Tile, {256}, const_true(), body); } return body; } diff --git a/src/ExtractTileOperations.h b/src/ExtractTileOperations.h index 918e3b1b9940..38da92329fca 100644 --- a/src/ExtractTileOperations.h +++ b/src/ExtractTileOperations.h @@ -11,7 +11,7 @@ namespace Halide { namespace Internal { -/** Rewrite any AMX tile operations that have been stored in the AMXTile memory +/** Rewrite any AMX tile operations that have been stored in the Tile memory * type as intrinsic calls, to be used in the X86 backend. */ Stmt extract_tile_operations(const Stmt &s); diff --git a/src/ExtractWMMAOperations.cpp b/src/ExtractWMMAOperations.cpp index 80bd0d044595..8212c922278d 100644 --- a/src/ExtractWMMAOperations.cpp +++ b/src/ExtractWMMAOperations.cpp @@ -21,7 +21,7 @@ * with the wmma load and store instructions, which move a tile between the * registers of a warp and a 2D array in memory. * - * A WMMAFragment allocation holds one of the three matrices of a multiply, and + * A tile allocation holds one of the three matrices of a multiply, and * which one follows from how it is used. An allocation accumulated into by a * matrix multiply is the accumulator; one read as an operand of a multiply is * that operand. The role determines which accesses are legal: @@ -251,7 +251,7 @@ MatmulInfo analyze_matmul(const Store *op) { // lane permutation left on it by vectorization. auto fail = [&](const char *reason) -> MatmulInfo { - user_error << "Matrix multiply not recognized. Store to a WMMAFragment " + user_error << "Matrix multiply not recognized. Store to a Tile " << "allocation must be a zero-initialization, a fill from a matrix " << "in memory, or a sum of a vector reduce op and a load from the " << "same allocation. In the following store, " << reason << ".\n" @@ -426,7 +426,7 @@ Stmt convert_to_tile_store(const Store *op, const Expr &store_index, op->is_streaming)); } -// Everything we learn about one WMMAFragment allocation from the way it is +// Everything we learn about one tile allocation from the way it is // used. The role and the shape come from the matrix multiplies it takes part // in, so neither is known until those have been found. struct Fragment { @@ -647,7 +647,7 @@ class ExtractWMMAOperations : public IRMutator { } Stmt visit(const Allocate *op) override { - if (op->memory_type != MemoryType::WMMAFragment) { + if (op->memory_type != MemoryType::Tile) { return IRMutator::visit(op); } @@ -669,7 +669,7 @@ class ExtractWMMAOperations : public IRMutator { if (pass == 0) { user_assert(f.role != Role::Unknown) - << op->name << " is stored in WMMAFragment memory, but no matrix " + << op->name << " is stored in Tile memory, but no matrix " << "multiply was found that accumulates into it or reads it as an " << "operand, so we can't tell what layout it should have.\n"; return op; @@ -681,7 +681,7 @@ class ExtractWMMAOperations : public IRMutator { // get replicated per thread. for (int i = 0; i < (int)f.subtiles.size(); i++) { body = Allocate::make(f.fragment_name + std::to_string(i), f.element_type, - MemoryType::WMMAFragment, {f.value_type().lanes()}, + MemoryType::Tile, {f.value_type().lanes()}, const_true(), body); } return body; diff --git a/src/ExtractWMMAOperations.h b/src/ExtractWMMAOperations.h index 79023f6193c8..283ea8deb533 100644 --- a/src/ExtractWMMAOperations.h +++ b/src/ExtractWMMAOperations.h @@ -15,7 +15,7 @@ namespace Internal { struct Call; struct Store; -/** Rewrite matrix multiplies that accumulate into WMMAFragment memory as +/** Rewrite matrix multiplies that accumulate into Tile memory as * calls to the wmma intrinsics understood by the PTX backend, and wrap them in * a loop over the 32 lanes of a warp. */ Stmt extract_wmma_operations(const Stmt &s); diff --git a/src/Func.cpp b/src/Func.cpp index 6382b0c8ec26..fabfa33c4904 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -1712,6 +1712,49 @@ Stage &Stage::vectorize(const VarOrRVar &var) { return *this; } +namespace { + +// Reorder the given dimensions to be the innermost loops of a stage, in the +// order given, leaving the relative order of the others unchanged. Naming +// every dimension makes reorder produce this exact total order, rather than +// just permuting the named ones among the positions they already hold. The +// outermost sentinel is left off so that it stays put. +Stage &reorder_innermost(Stage &stage, const std::vector &vars) { + std::vector order = vars; + for (const VarOrRVar &v : stage.split_vars()) { + if (v.name() == Var::outermost().name()) { + continue; + } + bool named = std::any_of(vars.begin(), vars.end(), + [&](const VarOrRVar &w) { return w.name() == v.name(); }); + if (!named) { + order.push_back(v); + } + } + return stage.reorder(order); +} + +} // namespace + +Stage &Stage::tile_init(const VarOrRVar &x, const VarOrRVar &y) { + return reorder_innermost(*this, {x, y}).vectorize(x).vectorize(y); +} + +Stage &Stage::tile_load(const VarOrRVar &x, const VarOrRVar &y) { + return reorder_innermost(*this, {x, y}).vectorize(x).vectorize(y); +} + +Stage &Stage::tile_store(const VarOrRVar &x, const VarOrRVar &y) { + return reorder_innermost(*this, {x, y}).vectorize(x).vectorize(y); +} + +Stage &Stage::tile_matmul(const VarOrRVar &r, const VarOrRVar &x, const VarOrRVar &y) { + // Vectorizing a reduction dimension needs the stage marked atomic, which + // is asking to reassociate the sum rather than asking for atomic memory + // operations. + return reorder_innermost(atomic(), {r, x, y}).vectorize(r).vectorize(x).vectorize(y); +} + Stage &Stage::unroll(const VarOrRVar &var) { set_dim_type(var, ForType::Unrolled); return *this; @@ -2574,6 +2617,35 @@ Func &Func::parallel(const VarOrRVar &var) { return *this; } +Func &Func::tile_init(const VarOrRVar &x, const VarOrRVar &y) { + invalidate_cache(); + store_in(MemoryType::Tile); + Stage(func, func.definition(), 0).tile_init(x, y); + return *this; +} + +Func &Func::tile_load(const VarOrRVar &x, const VarOrRVar &y) { + invalidate_cache(); + store_in(MemoryType::Tile); + Stage(func, func.definition(), 0).tile_load(x, y); + return *this; +} + +// A tile_store reads a tile and writes somewhere more general, so this Func is +// the destination rather than the tile. +Func &Func::tile_store(const VarOrRVar &x, const VarOrRVar &y) { + invalidate_cache(); + Stage(func, func.definition(), 0).tile_store(x, y); + return *this; +} + +Func &Func::tile_matmul(const VarOrRVar &r, const VarOrRVar &x, const VarOrRVar &y) { + invalidate_cache(); + store_in(MemoryType::Tile); + Stage(func, func.definition(), 0).tile_matmul(r, x, y); + return *this; +} + Func &Func::vectorize(const VarOrRVar &var) { invalidate_cache(); Stage(func, func.definition(), 0).vectorize(var); diff --git a/src/Func.h b/src/Func.h index 46516e3e572e..703e22df9247 100644 --- a/src/Func.h +++ b/src/Func.h @@ -499,6 +499,47 @@ class Stage { } // @} + /** Schedule this stage as one operation of the target's matrix unit, + * which computes on whole tiles held in MemoryType::Tile. There is one + * method per thing such a unit can do: + * + * - tile_init fills a tile with a value uniform across it, usually zero. + * - tile_load copies a tile in from more general memory. + * - tile_store copies a tile back out to more general memory. On targets + * whose tiles live in general purpose registers this may also fold in + * elementwise work whose other operands are uniform across the tile; on + * targets whose tiles are a separate register file, such as x86, it + * cannot, because no elementwise instruction can read a tile. + * - tile_matmul multiplies a pair of tiles into this one, taking the + * reduction dimension first. Its operands come either from memory, in + * which case they are staged into tiles implicitly, or from Funcs held + * in MemoryType::Tile. + * + * The dimensions passed are the ones that make up a tile, innermost + * first. They are reordered into that order and vectorized. That is all + * these do: they are sugar over reorder and vectorize, plus atomic for the + * multiply, which is what permits a reduction dimension to be vectorized. + * Asking for it is asking to reassociate the sum, which a matrix unit does + * anyway, rather than asking for atomic memory operations. + * + * Which dimensions make up a tile, and their order, is fixed by the + * instruction rather than free to schedule, so these move them to be the + * innermost loops in the order given. The surrounding loops keep their + * relative order, and can be reordered separately either side of these + * calls. There is no need to name the tile dimensions in such a reorder. + * + * On CUDA these are warp-level, with an inner loop over the 32 lanes of a + * warp that the schedule does not name. That loop becomes the innermost + * GPU thread dimension, so schedule these outside the innermost thread + * loops of any other Func computed alongside them, and give those Funcs a + * 32-wide gpu_lanes loop of their own to match. */ + // @{ + Stage &tile_init(const VarOrRVar &x, const VarOrRVar &y); + Stage &tile_load(const VarOrRVar &x, const VarOrRVar &y); + Stage &tile_store(const VarOrRVar &x, const VarOrRVar &y); + Stage &tile_matmul(const VarOrRVar &r, const VarOrRVar &x, const VarOrRVar &y); + // @} + /** Get the Vars and RVars of this definition, from innermost out, with * splits applied. This represents all the potentially-valid compute_at * sites for this Stage. The RVars returned will be symbolic and not tied to @@ -1521,6 +1562,17 @@ class Func { * innermost one. */ Func &vectorize(const VarOrRVar &var); + /** Schedule this Func as one operation of the target's matrix unit. See + * Stage::tile_init. Everything but tile_store also sets this Func's + * memory type to MemoryType::Tile, since those are the ones that produce + * a tile; tile_store reads one and writes somewhere more general. */ + // @{ + Func &tile_init(const VarOrRVar &x, const VarOrRVar &y); + Func &tile_load(const VarOrRVar &x, const VarOrRVar &y); + Func &tile_store(const VarOrRVar &x, const VarOrRVar &y); + Func &tile_matmul(const VarOrRVar &r, const VarOrRVar &x, const VarOrRVar &y); + // @} + /** Mark a dimension to be completely unrolled. The dimension * should have constant extent - e.g. because it is the inner * dimension following a split by a constant factor. For most uses diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index 489c59a1c245..1ba063969ae9 100644 --- a/src/FuseGPUThreadLoops.cpp +++ b/src/FuseGPUThreadLoops.cpp @@ -329,7 +329,7 @@ bool allocation_goes_to_registers(const Allocate *op, bool in_threads) { op->memory_type != MemoryType::GPUTexture) || op->memory_type == MemoryType::Register || op->memory_type == MemoryType::Stack || - op->memory_type == MemoryType::WMMAFragment; + op->memory_type == MemoryType::Tile; } // Rename an allocation and all of its loads, stores, and frees. Relies on the @@ -1259,7 +1259,7 @@ class ExtractRegisterAllocations : public IRMutator { op->memory_type == MemoryType::Register || op->memory_type == MemoryType::Heap || op->memory_type == MemoryType::Auto || - op->memory_type == MemoryType::WMMAFragment) + op->memory_type == MemoryType::Tile) << "Allocation " << op->name << " is scheduled inside a loop over GPU threads, so " << "it must live in stack memory, heap memory, or registers. " << "Shared allocations at this loop level are not yet supported.\n"; @@ -1476,8 +1476,7 @@ class InjectThreadBarriers : public IRMutator { case MemoryType::Register: case MemoryType::LockedCache: case MemoryType::VTCM: - case MemoryType::AMXTile: - case MemoryType::WMMAFragment: + case MemoryType::Tile: break; } @@ -1503,8 +1502,7 @@ class InjectThreadBarriers : public IRMutator { case MemoryType::Register: case MemoryType::LockedCache: case MemoryType::VTCM: - case MemoryType::AMXTile: - case MemoryType::WMMAFragment: + case MemoryType::Tile: break; } diff --git a/src/IRPrinter.cpp b/src/IRPrinter.cpp index ca747297a9ff..ea7b8d39e1e1 100644 --- a/src/IRPrinter.cpp +++ b/src/IRPrinter.cpp @@ -172,11 +172,8 @@ std::ostream &operator<<(std::ostream &out, const MemoryType &t) { case MemoryType::VTCM: out << "VTCM"; break; - case MemoryType::AMXTile: - out << "AMXTile"; - break; - case MemoryType::WMMAFragment: - out << "WMMAFragment"; + case MemoryType::Tile: + out << "Tile"; break; } return out; diff --git a/src/LowerWarpShuffles.cpp b/src/LowerWarpShuffles.cpp index ce13d807b521..7f796776c3ee 100644 --- a/src/LowerWarpShuffles.cpp +++ b/src/LowerWarpShuffles.cpp @@ -656,7 +656,7 @@ class LowerWarpShuffles : public IRMutator { if (this_lane.defined() || is_gpu_shared(op->memory_type) || op->memory_type == MemoryType::Heap || - op->memory_type == MemoryType::WMMAFragment) { + op->memory_type == MemoryType::Tile) { // Not an allocation for us to stripe. Warp-level storage is // per-lane register storage; shared and heap (global) memory are // never striped across lanes, and tensor core accumulators are diff --git a/src/Serialization.cpp b/src/Serialization.cpp index 9e78527614b0..9c050e4a95b3 100644 --- a/src/Serialization.cpp +++ b/src/Serialization.cpp @@ -158,10 +158,8 @@ Serialize::MemoryType Serializer::serialize_memory_type(const MemoryType &memory return Serialize::MemoryType::LockedCache; case MemoryType::VTCM: return Serialize::MemoryType::VTCM; - case MemoryType::AMXTile: - return Serialize::MemoryType::AMXTile; - case MemoryType::WMMAFragment: - return Serialize::MemoryType::WMMAFragment; + case MemoryType::Tile: + return Serialize::MemoryType::Tile; default: user_error << "Unsupported memory type\n"; return Serialize::MemoryType::Auto; diff --git a/src/halide_ir.fbs b/src/halide_ir.fbs index 1f2b10f627be..4715678f9c1f 100644 --- a/src/halide_ir.fbs +++ b/src/halide_ir.fbs @@ -116,9 +116,8 @@ enum MemoryType: byte { GPUTexture, LockedCache, VTCM, - AMXTile, + Tile, GPUSharedAsync, - WMMAFragment, } table Range { diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index 058e31116a76..7c774b2d2bf3 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -5,6 +5,7 @@ tests( SOURCES # keep-sorted start case=no align_bounds.cpp + amx_matmul.cpp argmax.cpp async_device_copy.cpp autodiff.cpp @@ -334,7 +335,6 @@ tests( sve_codegen_reinterpret.cpp target.cpp target_query.cpp - tiled_matmul.cpp tiled_matmul_errors.cpp tracing.cpp tracing_bounds.cpp @@ -378,6 +378,7 @@ tests( vectorized_load_from_vectorized_allocation.cpp vectorized_reduction_bug.cpp widening_reduction.cpp + wmma_matmul.cpp # keep-sorted end ) diff --git a/test/correctness/tiled_matmul.cpp b/test/correctness/amx_matmul.cpp similarity index 91% rename from test/correctness/tiled_matmul.cpp rename to test/correctness/amx_matmul.cpp index ea44064902db..0161968dda78 100644 --- a/test/correctness/tiled_matmul.cpp +++ b/test/correctness/amx_matmul.cpp @@ -113,40 +113,37 @@ bool matmul(int col, int row, int acc, int tile_x, int tile_y, int tile_r, bool // loads. But if you go too big you'll run out of tile registers and // compilation will fail (The LLVM AMX register allocator will spill, but it // seems to be fussy about it). Doing this also tests the case of more than - // one matrix multiply operation applied to a single AMXTile allocation. + // one matrix multiply operation applied to a single tile allocation. int outer_tile_x = col > tile_x ? 2 : 1, outer_tile_y = row > tile_y ? 2 : 1; mm.compute_at(mm.in(), x) - .store_in(MemoryType::AMXTile) .update() .tile(x, y, rxi, ryi, tile_x, tile_y, TailStrategy::GuardWithIf) .split(r, rro, rri, tile_r) - .reorder(rri, rxi, ryi, rro, x, y) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi) + .reorder(rro, x, y) + .tile_matmul(rri, rxi, ryi) .tile(x, y, xi, yi, outer_tile_x, outer_tile_y) - .reorder(rri, rxi, ryi, xi, yi, rro, x, y) + .reorder(xi, yi, rro, x, y) .unroll(xi) .unroll(yi); Var ixi("ixi"), iyi("iyi"); mm.compute_at(mm.in(), x) .tile(x, y, ixi, iyi, tile_x, tile_y) - .vectorize(ixi) - .vectorize(iyi) + .tile_init(ixi, iyi) .unroll(x) .unroll(y); - // schedule the consumer - Var mmxi("mmxi"), mmyi("mmyi"); + // Schedule the consumer, which is what copies each tile back out. The + // outer tiling matches the group of tiles the accumulator holds, and the + // inner one is a single tile. + Var mmxo("mmxo"), mmyo("mmyo"), mmxi("mmxi"), mmyi("mmyi"); mm.in() - .tile(x, y, mmxi, mmyi, tile_x * outer_tile_x, tile_y * outer_tile_y) - .vectorize(mmxi, tile_x) - .vectorize(mmyi, tile_y) - .unroll(mmxi) - .unroll(mmyi); + .tile(x, y, mmxo, mmyo, tile_x * outer_tile_x, tile_y * outer_tile_y) + .tile(mmxo, mmyo, mmxi, mmyi, tile_x, tile_y) + .tile_store(mmxi, mmyi) + .unroll(mmxo) + .unroll(mmyo); Func result = mm.in(); @@ -160,7 +157,7 @@ bool matmul(int col, int row, int acc, int tile_x, int tile_y, int tile_r, bool result.realize(out); } else { // Just compile it to see if anything crashes - result.compile_to_assembly(Internal::get_test_tmp_dir() + "tiled_matmul.s", + result.compile_to_assembly(Internal::get_test_tmp_dir() + "amx_matmul.s", {A_buf, B_buf}, Target{"x86-64-linux-avx512_sapphirerapids-no_asserts-no_runtime-no_bounds_query"}); return true; } @@ -211,28 +208,22 @@ bool matmul_bf16(int col, int row, int acc, int tile_x, int tile_y, int tile_r, RVar rri("rri"), rro("rro"); mm.compute_at(mm.in(), x) - .store_in(MemoryType::AMXTile) .update() .tile(x, y, rxi, ryi, tile_x, tile_y, TailStrategy::GuardWithIf) .split(r.x, rro, rri, tile_r) .reorder({rri, rxi, ryi, rro, x, y}) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi); + .tile_matmul(rri, rxi, ryi); Var ixi("ixi"), iyi("iyi"); mm.compute_at(mm.in(), x) .tile(x, y, ixi, iyi, tile_x, tile_y) - .vectorize(ixi) - .vectorize(iyi); + .tile_init(ixi, iyi); // schedule the consumer Var mmxi("mmxi"), mmyi("mmyi"); mm.in() .tile(x, y, mmxi, mmyi, tile_x, tile_y) - .vectorize(mmxi) - .vectorize(mmyi); + .tile_store(mmxi, mmyi); Func result = mm.in(); @@ -242,15 +233,15 @@ bool matmul_bf16(int col, int row, int acc, int tile_x, int tile_y, int tile_r, Buffer out(col, row); // Uncomment to check the asm - // result.compile_to_llvm_assembly(Internal::get_test_tmp_dir() + "tiled_matmul_bf16.ll", {A, B}, target); - // result.compile_to_assembly(Internal::get_test_tmp_dir() + "tiled_matmul.s", {A, B}, target); + // result.compile_to_llvm_assembly(Internal::get_test_tmp_dir() + "amx_matmul_bf16.ll", {A, B}, target); + // result.compile_to_assembly(Internal::get_test_tmp_dir() + "amx_matmul.s", {A, B}, target); Target target = get_jit_target_from_environment(); if (target.has_feature(Target::AVX512_SapphireRapids)) { result.realize(out); } else { // Just compile it to see if anything crashes - result.compile_to_assembly(Internal::get_test_tmp_dir() + "tiled_matmul.s", {A, B}, Target{"x86-64-linux-avx512_sapphirerapids"}); + result.compile_to_assembly(Internal::get_test_tmp_dir() + "amx_matmul.s", {A, B}, Target{"x86-64-linux-avx512_sapphirerapids"}); return true; } diff --git a/test/correctness/tiled_matmul_errors.cpp b/test/correctness/tiled_matmul_errors.cpp index 751ca4f607e7..ba11bc84da70 100644 --- a/test/correctness/tiled_matmul_errors.cpp +++ b/test/correctness/tiled_matmul_errors.cpp @@ -45,7 +45,7 @@ void schedule_matmul(Func mm, RVar r, int tile_x, int tile_y, int tile_r) { Var x("x"), y("y"), rxi("rxi"), ryi("ryi"); RVar rri("rri"), rro("rro"); mm.compute_at(mm.in(), x) - .store_in(MemoryType::AMXTile) + .store_in(MemoryType::Tile) .update() .tile(x, y, rxi, ryi, tile_x, tile_y, TailStrategy::GuardWithIf) .split(r, rro, rri, tile_r) @@ -83,7 +83,7 @@ void scenario_too_large() { mm.in().compile_jit(amx_target); } -// AMXTile allocated for a non-i32/f32 result. AMX always accumulates into +// A tile allocated for a non-i32/f32 result. AMX always accumulates into // 32-bit registers, so we reject. Triggers the user_assert in // visit(Allocate). void scenario_bad_result_type() { @@ -177,7 +177,7 @@ void scenario_conv1d() { conv.in().compile_jit(amx_target); } -// A non-matmul value scheduled into an AMXTile allocation by mistake. +// A non-matmul value scheduled into a tile allocation by mistake. // Triggers the "no matrix multiply was found" assertion. void scenario_no_matmul() { Var x("x"), y("y"), xo("xo"), yo("yo"), xi("xi"), yi("yi"); @@ -185,7 +185,7 @@ void scenario_no_matmul() { Func f("f"); f(x, y) = 0; f.compute_at(f.in(), xo) - .store_in(MemoryType::AMXTile) + .store_in(MemoryType::Tile) .vectorize(x, 8) .vectorize(y, 8); @@ -226,7 +226,7 @@ void scenario_mismatched_strides() { Var rxi("rxi"), ryi("ryi"); RVar rri("rri"), rro("rro"); - mm.compute_at(mm.in(), x).store_in(MemoryType::AMXTile); + mm.compute_at(mm.in(), x).store_in(MemoryType::Tile); mm.update(0) .tile(x, y, rxi, ryi, 8, 4, TailStrategy::GuardWithIf) .split(r1.x, rro, rri, 8) @@ -251,7 +251,7 @@ void scenario_mismatched_strides() { } // A user gives the same Func two update definitions that each store into -// the same AMXTile allocation but with different tile sizes (e.g. a fast +// the same tile allocation but with different tile sizes (e.g. a fast // path for the bulk of K and a smaller fallback). The matcher requires // every matmul touching a given allocation to agree on tile dimensions. void scenario_inconsistent_tiles() { @@ -268,7 +268,7 @@ void scenario_inconsistent_tiles() { Var rxi("rxi"), ryi("ryi"); RVar rri("rri"), rro("rro"); - mm.compute_at(mm.in(), x).store_in(MemoryType::AMXTile); + mm.compute_at(mm.in(), x).store_in(MemoryType::Tile); mm.update(0) .tile(x, y, rxi, ryi, 8, 8, TailStrategy::GuardWithIf) .split(r1.x, rro, rri, 8) @@ -301,7 +301,7 @@ void scenario_inconsistent_tiles() { mm.in().compile_jit(amx_target); } -// A reduction inside an AMXTile that's a sum-of-something-else, not a +// A reduction inside a tile that's a sum-of-something-else, not a // vector_reduce_add of a widening multiply. Here we accumulate a // non-multiplied value, which produces a Store whose RHS is not a // vector-reduce-of-multiply. @@ -318,7 +318,7 @@ void scenario_not_a_matmul_pattern() { Var rxi("rxi"), ryi("ryi"); RVar rri("rri"), rro("rro"); mm.compute_at(mm.in(), x) - .store_in(MemoryType::AMXTile) + .store_in(MemoryType::Tile) .update() .tile(x, y, rxi, ryi, 8, 8, TailStrategy::GuardWithIf) .split(r.x, rro, rri, 8) diff --git a/test/correctness/wmma_matmul.cpp b/test/correctness/wmma_matmul.cpp index 56cf83664c88..bdf9c594f8de 100644 --- a/test/correctness/wmma_matmul.cpp +++ b/test/correctness/wmma_matmul.cpp @@ -135,15 +135,13 @@ bool test(const Params &p) { .bound(y, 0, p.M) .split(x, x, xi, p.tile_n * p.tiles_n * p.warps) .split(xi, xt, xi, p.tile_n * p.tiles_n) - .split(xi, xi, mmxi, p.tile_n) .split(y, y, yi, p.tile_m * p.tiles_m) - .split(yi, yi, mmyi, p.tile_m) + .tile(xi, yi, mmxi, mmyi, p.tile_n, p.tile_m) .gpu_blocks(x, y) - .reorder(mmxi, mmyi, xi, yi, xt, x, y) + .reorder(xi, yi, xt, x, y) .unroll(xi) .unroll(yi) - .vectorize(mmxi) - .vectorize(mmyi); + .tile_store(mmxi, mmyi); if (p.warps > 1) { // With one warp there's no need for a loop over warps, and leaving it // serial keeps anything computed inside it out of the thread loops. @@ -151,25 +149,24 @@ bool test(const Params &p) { } prod.compute_at(out, xt) - .store_in(MemoryType::WMMAFragment) - .split(x, x, rxi, p.tile_n) - .split(y, y, ryi, p.tile_m) - .vectorize(rxi) - .vectorize(ryi) + .tile(x, y, rxi, ryi, p.tile_n, p.tile_m) .unroll(x) .unroll(y); + // The accumulator either starts at a constant or is loaded from memory, + // and those are different tile instructions. + if (p.init_from_memory) { + prod.tile_load(rxi, ryi); + } else { + prod.tile_init(rxi, ryi); + } prod.update() - .split(x, x, rxi, p.tile_n) - .split(y, y, ryi, p.tile_m) + .tile(x, y, rxi, ryi, p.tile_n, p.tile_m) .split(k, rro, rri, p.tile_k) - .reorder(rri, rxi, ryi, x, y, rro) + .reorder(x, y, rro) .unroll(x) .unroll(y) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi); + .tile_matmul(rri, rxi, ryi); if (p.init_from_memory) { Var ixi("ixi"), iyi("iyi"); @@ -251,17 +248,25 @@ bool test_block_level_accumulator() { Var kko("kko"), kki("kki"), xxo("xxo"), xxi("xxi"); RVar ko("ko"), ki("ki"), rri("rri"); - out.bound(x, 0, N).bound(y, 0, M).split(x, x, xi, block_x).split(xi, xt, xi, tile * tiles_x).split(xi, xi, mmxi, tile).split(y, y, yi, block_y).split(yi, yi, mmyi, tile).gpu_blocks(x, y).gpu_threads(xt).reorder(mmxi, mmyi, xi, yi, xt, x, y).unroll(xi).unroll(yi).vectorize(mmxi).vectorize(mmyi); + out.bound(x, 0, N) + .bound(y, 0, M) + .tile(x, y, xi, yi, block_x, block_y) + .split(xi, xt, xi, tile * tiles_x) + .tile(xi, yi, mmxi, mmyi, tile, tile) + .gpu_blocks(x, y) + .gpu_threads(xt) + .reorder(xi, yi, xt, x, y) + .unroll(xi) + .unroll(yi) + .tile_store(mmxi, mmyi); prod.compute_at(out, x) - .store_in(MemoryType::WMMAFragment) .split(x, xw, xi, tile * tiles_x) .split(xi, xi, rxi, tile) .split(y, y, ryi, tile) - .reorder(rxi, ryi, xi, y, xw) + .reorder(xi, y, xw) .gpu_threads(xw) - .vectorize(rxi) - .vectorize(ryi) + .tile_init(rxi, ryi) .unroll(xi) .unroll(y); @@ -271,15 +276,12 @@ bool test_block_level_accumulator() { .split(xi, xi, rxi, tile) .split(y, y, ryi, tile) .split(ki, ki, rri, tile) - .reorder(rri, rxi, ryi, xi, y, ki, xw, ko) + .reorder(xi, y, ki, xw, ko) .gpu_threads(xw) .unroll(xi) .unroll(y) .unroll(ki) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi); + .tile_matmul(rri, rxi, ryi); As.compute_at(prod, ko) .store_in(MemoryType::GPUShared) @@ -348,62 +350,49 @@ bool test_staged_operands() { out.bound(x, 0, N) .bound(y, 0, M) - .split(x, x, xi, tile * tiles_x) - .split(xi, xi, mmxi, tile) - .split(y, y, yi, tile * tiles_y) - .split(yi, yi, mmyi, tile) + .tile(x, y, xi, yi, tile * tiles_x, tile * tiles_y) + .tile(xi, yi, mmxi, mmyi, tile, tile) .gpu_blocks(x, y) - .reorder(mmxi, mmyi, xi, yi, x, y) + .reorder(xi, yi, x, y) .unroll(xi) .unroll(yi) .vectorize(mmxi) .vectorize(mmyi); prod.compute_at(out, x) - .store_in(MemoryType::WMMAFragment) - .split(x, x, rxi, tile) - .split(y, y, ryi, tile) - .vectorize(rxi) - .vectorize(ryi) + .tile(x, y, rxi, ryi, tile, tile) + .tile_init(rxi, ryi) .unroll(x) .unroll(y); // Loop nest of the update, outermost first: ko, ki, y, x. prod.update() .split(k, ko, ki, bk) - .split(x, x, rxi, tile) - .split(y, y, ryi, tile) + .tile(x, y, rxi, ryi, tile, tile) .split(ki, ki, rri, tile) - .reorder(rri, rxi, ryi, x, y, ki, ko) + .reorder(x, y, ki, ko) .unroll(x) .unroll(y) .unroll(ki) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi); + .tile_matmul(rri, rxi, ryi); // One a fragment per row of tiles, live across the loop over columns. Am.compute_at(prod, y) - .store_in(MemoryType::WMMAFragment) .split(kk, kko, kki, tile) .split(yy, yyo, yyi, tile) .reorder(kki, yyi, kko, yyo) .unroll(kko) .unroll(yyo) - .vectorize(kki) - .vectorize(yyi); + .tile_load(kki, yyi); // All the b fragments at once, live across the loops over both. Bm.compute_at(prod, ki) - .store_in(MemoryType::WMMAFragment) .split(xx, xxo, xxi, tile) .split(kk, kko, kki, tile) .reorder(xxi, kki, xxo, kko) .unroll(xxo) .unroll(kko) - .vectorize(xxi) - .vectorize(kki); + .tile_load(xxi, kki); Buffer result(N, M); out.realize(result); @@ -454,47 +443,36 @@ bool test_operand_hoisted_out_of_loop() { out.bound(x, 0, N) .bound(y, 0, M) .bound(n, 0, batch) - .split(x, x, xi, N) - .split(xi, xi, mmxi, tile) - .split(y, y, yi, M) - .split(yi, yi, mmyi, tile) + .tile(x, y, xi, yi, N, M) + .tile(xi, yi, mmxi, mmyi, tile, tile) .gpu_blocks(x, y) - .reorder(mmxi, mmyi, xi, yi, n, x, y) + .reorder(xi, yi, n, x, y) .unroll(xi) .unroll(yi) .vectorize(mmxi) .vectorize(mmyi); prod.compute_at(out, n) - .store_in(MemoryType::WMMAFragment) - .split(x, x, rxi, tile) - .split(y, y, ryi, tile) - .vectorize(rxi) - .vectorize(ryi) + .tile(x, y, rxi, ryi, tile, tile) + .tile_init(rxi, ryi) .unroll(x) .unroll(y); prod.update() - .split(x, x, rxi, tile) - .split(y, y, ryi, tile) + .tile(x, y, rxi, ryi, tile, tile) .split(k, rro, rri, tile) - .reorder(rri, rxi, ryi, x, y, rro) + .reorder(x, y, rro) .unroll(x) .unroll(y) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi); + .tile_matmul(rri, rxi, ryi); Am.compute_at(out, x) - .store_in(MemoryType::WMMAFragment) .split(kk, kko, kki, tile) .split(yy, yyo, yyi, tile) .reorder(kki, yyi, kko, yyo) .unroll(kko) .unroll(yyo) - .vectorize(kki) - .vectorize(yyi); + .tile_load(kki, yyi); Buffer result(N, M, batch); out.realize(result); diff --git a/test/performance/tiled_matmul.cpp b/test/performance/tiled_matmul.cpp index 916aa354eb57..c990ca052e32 100644 --- a/test/performance/tiled_matmul.cpp +++ b/test/performance/tiled_matmul.cpp @@ -107,7 +107,7 @@ bool matmul(Halide::Target target) { Var rxi("rxi"), ryi("ryi"); RVar rri("rri"), rro("rro"); mm.compute_at(mm.in(), y) - .store_in(MemoryType::AMXTile) + .store_in(MemoryType::Tile) .update() // Split into (x,y) tile .tile(y, x, ryi, rxi, tile_y, tile_x, TailStrategy::GuardWithIf) @@ -196,7 +196,7 @@ bool matmul_bf16(Halide::Target target) { RVar rri("rri"), rro("rro"); mm.compute_at(mm.in(), x) - .store_in(MemoryType::AMXTile) + .store_in(MemoryType::Tile) .update() .tile(x, y, rxi, ryi, tile_x, tile_y, TailStrategy::GuardWithIf) .split(r.x, rro, rri, tile_r) From 1f7875bd92ff2f6794e0c48b1527c60990232b57 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 4 Aug 2026 16:42:55 -0700 Subject: [PATCH 53/59] Remove the tensorcore_resize app Its schedule needs adjusting for the tile directives, so it will return in a later PR. Co-Authored-By: Claude Opus 5 --- apps/tensorcore_resize/Makefile | 31 -- apps/tensorcore_resize/README.md | 21 -- apps/tensorcore_resize/resize_generator.cpp | 329 -------------------- apps/tensorcore_resize/runner.cpp | 78 ----- 4 files changed, 459 deletions(-) delete mode 100644 apps/tensorcore_resize/Makefile delete mode 100644 apps/tensorcore_resize/README.md delete mode 100644 apps/tensorcore_resize/resize_generator.cpp delete mode 100644 apps/tensorcore_resize/runner.cpp diff --git a/apps/tensorcore_resize/Makefile b/apps/tensorcore_resize/Makefile deleted file mode 100644 index 307138395969..000000000000 --- a/apps/tensorcore_resize/Makefile +++ /dev/null @@ -1,31 +0,0 @@ -include ../support/Makefile.inc - - -# The wmma instructions require compute capability 7.0 or above. -TENSORCORE_TARGET ?= host-cuda-cuda_capability_80 - -all: $(BIN)/$(HL_TARGET)/runner - -$(GENERATOR_BIN)/resize.generator: resize_generator.cpp $(GENERATOR_DEPS) - @mkdir -p $(@D) - $(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LIBHALIDE_LDFLAGS) - -$(BIN)/%/resize_cudaonly.a: $(GENERATOR_BIN)/resize.generator - @mkdir -p $(@D) - $^ -g resize -f resize_cudaonly -e $(GENERATOR_OUTPUTS) -o $(@D) \ - target=$(TENSORCORE_TARGET) gpu_schedule=cudaonly - -$(BIN)/%/resize_tensorcore.a: $(GENERATOR_BIN)/resize.generator - @mkdir -p $(@D) - $^ -g resize -f resize_tensorcore -e $(GENERATOR_OUTPUTS) -o $(@D) \ - target=$(TENSORCORE_TARGET) gpu_schedule=tensorcore - -$(BIN)/%/runner: runner.cpp $(BIN)/%/resize_cudaonly.a $(BIN)/%/resize_tensorcore.a - @mkdir -p $(@D) - $(CXX) $(CXXFLAGS) -I$(BIN)/$* -Wall $^ -o $@ $(LDFLAGS) $(LIBHALIDE_LDFLAGS) - -test: $(BIN)/$(HL_TARGET)/runner - $^ - -clean: - rm -rf $(BIN) diff --git a/apps/tensorcore_resize/README.md b/apps/tensorcore_resize/README.md deleted file mode 100644 index fc0dd391a5ce..000000000000 --- a/apps/tensorcore_resize/README.md +++ /dev/null @@ -1,21 +0,0 @@ -# Tensor core resize - -Resampling an image is a linear operator, so it can be written as a matrix -multiply. The matrix is enormous and almost entirely zero, so you never want to -materialize it, but each of its rows has a small number of contiguous -non-zeros. If neighbouring rows are made to share a starting column (which just -means storing a few more zeros), then a block of 16 rows of it is a small dense -matrix, and the inner loop becomes a matrix multiply. That's a large speed-up -even without tensor cores, and it lets us use tensor cores when we have them. - -This is the algorithmic difference between this generator and the one in -`apps/resize`, which starts each row of the matrix at its own column. - -## Status - -Both schedules work. On an RTX 5060 Ti, downsampling a 3840x2160 image by 4x -with a Lanczos kernel takes 0.359 ms with the `cudaonly` schedule and 0.246 ms -with the `tensorcore` one, a 1.46x speed-up. - -Both schedules round the output size up to a multiple of 16, so the runner -gives them an output buffer of that size. diff --git a/apps/tensorcore_resize/resize_generator.cpp b/apps/tensorcore_resize/resize_generator.cpp deleted file mode 100644 index 433ec74ec333..000000000000 --- a/apps/tensorcore_resize/resize_generator.cpp +++ /dev/null @@ -1,329 +0,0 @@ -#include "Halide.h" - -namespace { - -using namespace Halide; - -enum class InterpolationType { - Box, - Linear, - Cubic, - Lanczos, -}; - -enum class Schedule { - CUDA, - TensorCore, -}; - -Expr kernel_box(Expr x) { - Expr xx = abs(x); - return select(xx <= 0.5f, 1.0f, 0.0f); -} - -Expr kernel_linear(Expr x) { - Expr xx = abs(x); - return select(xx < 1.0f, 1.0f - xx, 0.0f); -} - -Expr kernel_cubic(Expr x) { - Expr xx = abs(x); - Expr xx2 = xx * xx; - Expr xx3 = xx2 * xx; - float a = -0.5f; - - return select(xx < 1.0f, (a + 2.0f) * xx3 - (a + 3.0f) * xx2 + 1, - select(xx < 2.0f, a * xx3 - 5 * a * xx2 + 8 * a * xx - 4.0f * a, - 0.0f)); -} - -Expr sinc(Expr x) { - x *= 3.14159265359f; - return sin(x) / x; -} - -constexpr int lanczos_lobes = 3; - -Expr kernel_lanczos(Expr x) { - Expr value = sinc(x) * sinc(x / lanczos_lobes); - // Take care of the singularity at zero - value = select(x == 0.0f, 1.0f, value); - // Clamp to zero out of bounds - value = select(x > lanczos_lobes || x < -lanczos_lobes, 0.0f, value); - return value; -} - -struct KernelInfo { - const char *name; - int taps; - Expr (*kernel)(Expr); -}; - -const KernelInfo kernel_info[] = { - {"box", 1, kernel_box}, - {"linear", 2, kernel_linear}, - {"cubic", 4, kernel_cubic}, - {"lanczos", 2 * lanczos_lobes, kernel_lanczos}}; - -// Resampling an image is a linear operator, so it can be written as a matrix -// multiply. The matrix is enormous and almost entirely zero, so you never want -// to materialize it, but each row of it has a small number of contiguous -// non-zeros, and if we let neighbouring rows share a starting column then a -// block of 16 rows of it is a small dense matrix. That makes the inner loop a -// matrix multiply, which is a large speed-up even without tensor cores, and -// lets us use tensor cores when we have them. -class Resize : public Halide::Generator { -public: - GeneratorParam interpolation_type{ - "interpolation_type", InterpolationType::Lanczos, {{"box", InterpolationType::Box}, {"linear", InterpolationType::Linear}, {"cubic", InterpolationType::Cubic}, {"lanczos", InterpolationType::Lanczos}}}; - - // If we statically know whether we're upsampling or downsampling, we can - // generate different pipelines (we want to reorder the resample in x and - // in y). - GeneratorParam upsample{"upsample", false}; - - GeneratorParam gpu_schedule{ - "gpu_schedule", Schedule::TensorCore, {{"cudaonly", Schedule::CUDA}, {"tensorcore", Schedule::TensorCore}}}; - - Input> input{"input"}; - Input scale_factor{"scale_factor"}; - Output> output{"output"}; - - // The size of the blocks of the resampling matrix that we treat as dense. - static constexpr int tile = 16; - - void generate() { - // Invert the scale factor in a single place, to avoid getting slightly - // different ratios showing up in different places. - Expr inverse_scale_factor = 1.0f / scale_factor; - - // For downscaling, widen the interpolation kernel to perform lowpass - // filtering. - Expr kernel_scaling = upsample ? Expr(1.0f) : scale_factor; - Expr inverse_kernel_scaling = upsample ? Expr(1.0f) : inverse_scale_factor; - - const KernelInfo &info = kernel_info[(int)(InterpolationType)interpolation_type]; - - Expr kernel_radius = 0.5f * info.taps * inverse_kernel_scaling; - Expr kernel_taps = cast(ceil(info.taps * inverse_kernel_scaling)); - - // The (non-integer) coordinates in the source image. - Expr sourcex = (x + 0.5f) * inverse_scale_factor - 0.5f; - Expr sourcey = (y + 0.5f) * inverse_scale_factor - 0.5f; - - // For a given output coordinate, the first input coordinate it depends - // on. We can start a row of the matrix at any column we like as long - // as we store enough columns, so we use the same starting column for - // each group of `tile` rows. - auto begin_of = [&](Expr coord) { - return cast(ceil((coord + 0.5f) * inverse_scale_factor - 0.5f - kernel_radius)); - }; - - Expr beginx = begin_of((x / tile) * tile); - Expr beginy = begin_of((y / tile) * tile); - - // Moving the start of each row back like that means each row has to - // cover a longer contiguous region. - Expr extra_zeros = begin_of(tile) - begin_of(0); - - // Round the number of columns up to the next multiple of the tile size - // too, so that the reduction is a whole number of tiles. - Expr span = ((kernel_taps + extra_zeros + tile - 1) / tile) * tile; - - // Don't go off the end of the image. Those columns would be zero - // anyway. - beginx = clamp(beginx, 0, input.width() - span); - beginy = clamp(beginy, 0, input.height() - span); - - r = RDom(0, span, "r"); - - as_float(x, y, c) = cast(input(x, y, c)); - - unnormalized_kernel_x(x, k) = info.kernel((k + beginx - sourcex) * kernel_scaling); - unnormalized_kernel_y(y, k) = info.kernel((k + beginy - sourcey) * kernel_scaling); - - kernel_sum_x(x) += unnormalized_kernel_x(x, r); - kernel_sum_y(y) += unnormalized_kernel_y(y, r); - - kernel_x(x, k) = cast(unnormalized_kernel_x(x, k) / kernel_sum_x(x)); - kernel_y(y, k) = cast(unnormalized_kernel_y(y, k) / kernel_sum_y(y)); - - resized_y(x, y, c) += kernel_y(y, r) * as_float(x, r + beginy, c); - resized_x(x, y, c) += kernel_x(x, r) * resized_y(r + beginx, y, c); - - output(x, y, c) = clamp(resized_x(x, y, c), cast(0.f), cast(1.f)); - } - - void schedule() { - Var xi("xi"), yi("yi"), ki("ki"), xii("xii"), yii("yii"), xo("xo"), z("z"); - - // Precompute the sparse matrices. These are tiny compared to the - // image, so the schedule barely matters. - kernel_x.compute_root().gpu_tile(x, k, xi, ki, 32, 8); - unnormalized_kernel_x.compute_root().gpu_tile(x, k, xi, ki, 32, 8); - kernel_sum_x.in().compute_root().gpu_tile(x, xi, 32); - - kernel_y.compute_root().gpu_tile(y, k, yi, ki, 32, 8); - unnormalized_kernel_y.compute_root().gpu_tile(y, k, yi, ki, 32, 8); - kernel_sum_y.in().compute_root().gpu_tile(y, yi, 32); - - output.compute_root() - .align_bounds(x, tile) - .align_bounds(y, tile); - - if (gpu_schedule == Schedule::CUDA) { - // Resampling in y is the expensive stage for large downsamples. - // The load from the kernel doesn't depend on x or c, and the load - // from the image doesn't depend on y % tile, so we schedule it - // like a matrix multiply. - resized_y.in() - .compute_root() - .align_bounds(x, tile) - .align_bounds(y, tile) - .reorder(c, x, y) - .unroll(c) - .gpu_tile(x, y, xi, yi, 32, 16, TailStrategy::RoundUp) - .tile(xi, yi, xii, yii, 2, 4) - .unroll(xii) - .unroll(yii); - resized_y - .compute_at(resized_y.in(), xi) - .unroll(c) - .unroll(x) - .unroll(y) - .update() - .reorder(x, y, c, r) - .unroll(c) - .unroll(x) - .unroll(y); - as_float.compute_at(resized_y, c).vectorize(x).vectorize(y); - kernel_y.in().compute_at(resized_y, r).vectorize(y).vectorize(k); - - // After downsampling in y it's hard to fill the machine, so use - // smaller tiles and map color channels to gpu threads. - output - .gpu_threads(c) - .gpu_tile(x, y, xi, yi, 32, 4, TailStrategy::RoundUp) - .reorder(xi, yi, c, x, y) - .tile(xi, yi, xii, yii, 2, 2) - .vectorize(xii) - .unroll(yii); - - resized_x - .compute_at(output, xi) - .unroll(c) - .unroll(x) - .unroll(y) - .update() - .reorder(x, y, c, r) - .unroll(c) - .unroll(x) - .unroll(y); - resized_y.in().in().compute_at(resized_x, c).vectorize(y); - kernel_x.in().compute_at(resized_x, r).vectorize(x).vectorize(k); - } else { - // The tensor core instructions want the reduction dimension of - // each operand dense in memory. - kernel_x.reorder_storage(k, x); - kernel_y.reorder_storage(k, y); - - Var xio("xio"); - resized_y.in() - .compute_root() - .align_bounds(x, tile) - .align_bounds(y, tile) - .tile(x, y, xi, yi, 32, 16, TailStrategy::RoundUp) - .unroll(c) - .split(xi, xi, xii, 32) - .split(xi, xio, xi, 1) - .gpu_threads(xio) - .split(yi, yi, yii, 8) - .reorder(xii, yii, c, yi, xi, xio, x, y) - .vectorize(xii) - .vectorize(yii) - .unroll(yi) - .unroll(xi) - .gpu_blocks(x, y); - - // An 8x32 tile of accumulator, reducing 16 taps at a time. - resized_y.compute_at(resized_y.in(), xio) - .store_in(MemoryType::Tile) - .unroll(c) - .vectorize(x, 32) - .unroll(x) - .vectorize(y, 8) - .unroll(y) - .update() - .atomic() - .unroll(c) - .vectorize(x, 32) - .unroll(x) - .vectorize(y, 8) - .unroll(y) - .vectorize(r, tile) - .reorder(y, c, x, r); - - output - .tile(x, y, xi, yi, tile, tile, TailStrategy::RoundUp) - .reorder(yi, xi, x, y, c) - .gpu_blocks(x, y, c) - .split(yi, yi, yii, 2) - .fuse(xi, yii, z) - .gpu_lanes(z) - .unroll(yi); - - resized_x.in() - .compute_at(output, x) - .vectorize(x) - .vectorize(y); - - RVar ri("ri"), ro("ro"); - resized_x - .store_in(MemoryType::Tile) - .compute_at(resized_x.in(), c) - .vectorize(x) - .vectorize(y) - .update() - .atomic() - .split(r, ro, ri, tile) - .reorder(ri, x, y, ro) - .vectorize(x) - .vectorize(y) - .vectorize(ri); - - // An extra layer of staging, because we're not necessarily aligned - // in x. - resized_y.in() - .in() - .compute_at(output, x) - .store_in(MemoryType::GPUShared) - .split(x, xo, xi, 32, TailStrategy::RoundUp) - .gpu_lanes(xi); - } - - output.dim(0).set_min(0); - output.dim(1).set_min(0); - output.dim(2).set_bounds(0, 3); - input.dim(0).set_min(0); - input.dim(1).set_min(0); - input.dim(2).set_bounds(0, 3); - } - -private: - Var x{"x"}, y{"y"}, c{"c"}, k{"k"}; - RDom r; - - Func as_float{"as_float"}, - resized_x{"resized_x"}, - resized_y{"resized_y"}, - unnormalized_kernel_x{"unnormalized_kernel_x"}, - unnormalized_kernel_y{"unnormalized_kernel_y"}, - kernel_x{"kernel_x"}, - kernel_y{"kernel_y"}, - kernel_sum_x{"kernel_sum_x"}, - kernel_sum_y{"kernel_sum_y"}; -}; - -} // namespace - -HALIDE_REGISTER_GENERATOR(Resize, resize) diff --git a/apps/tensorcore_resize/runner.cpp b/apps/tensorcore_resize/runner.cpp deleted file mode 100644 index b96cb7e710de..000000000000 --- a/apps/tensorcore_resize/runner.cpp +++ /dev/null @@ -1,78 +0,0 @@ -#include "Halide.h" -#include "HalideBuffer.h" -#include "HalideRuntimeCuda.h" -#include "halide_benchmark.h" -#include - -#include "resize_cudaonly.h" -#include "resize_tensorcore.h" - -using Halide::float16_t; -using Halide::Runtime::Buffer; -using Halide::Tools::benchmark; - -int main(int argc, char **argv) { - const auto *interface = halide_cuda_device_interface(); - int major, minor; - if (interface->compute_capability(nullptr, &major, &minor) != 0 || - major * 10 + minor < 70) { - printf("[SKIP] Tensor cores require CUDA compute capability 7.0 or above.\n"); - return 0; - } - - const int in_w = 3840, in_h = 2160; - const float scale_factor = 0.25f; - const int out_w = (int)(in_w * scale_factor), out_h = (int)(in_h * scale_factor); - - Buffer input(in_w, in_h, 3); - input.fill([]() { return float16_t((float)rand() / RAND_MAX); }); - - // Both schedules work on whole 16x16 tiles of the output, so round the - // output size up to a multiple of that. - const int tile = 16; - const int buf_w = ((out_w + tile - 1) / tile) * tile; - const int buf_h = ((out_h + tile - 1) / tile) * tile; - - Buffer out_cuda(buf_w, buf_h, 3), out_tensorcore(buf_w, buf_h, 3); - - resize_cudaonly(input, scale_factor, out_cuda); - resize_tensorcore(input, scale_factor, out_tensorcore); - out_cuda.copy_to_host(); - out_tensorcore.copy_to_host(); - - // The two schedules compute the same thing, but accumulate in a different - // order in half precision, so only compare them approximately. - int bad = 0; - for (int c = 0; c < 3; c++) { - for (int y = 0; y < out_h; y++) { - for (int x = 0; x < out_w; x++) { - float a = (float)out_cuda(x, y, c), b = (float)out_tensorcore(x, y, c); - if (std::abs(a - b) > 5e-3f) { - if (bad++ < 10) { - printf("Mismatch at %d %d %d: %f != %f\n", x, y, c, a, b); - } - } - } - } - } - if (bad) { - printf("Failed with %d mismatches\n", bad); - return 1; - } - - double t_cuda = benchmark([&]() { - resize_cudaonly(input, scale_factor, out_cuda); - out_cuda.device_sync(); - }); - double t_tensorcore = benchmark([&]() { - resize_tensorcore(input, scale_factor, out_tensorcore); - out_tensorcore.device_sync(); - }); - - printf("cuda only: %8.3f ms\n", t_cuda * 1e3); - printf("tensor core: %8.3f ms\n", t_tensorcore * 1e3); - printf("speed-up: %8.2fx\n", t_cuda / t_tensorcore); - - printf("Success!\n"); - return 0; -} From 5b2d9a61e1db5cb41e4c24811f985176aa3360ca Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 4 Aug 2026 16:46:40 -0700 Subject: [PATCH 54/59] Fix include order in two files Own header first, then internal headers, then system headers, with a blank line between each group. clang-format sorts within blank-line-separated blocks but does not reorder the blocks, so these had survived reformatting. Co-Authored-By: Claude Opus 5 --- src/CodeGen_PTX_Dev.cpp | 4 ++-- src/ExtractWMMAOperations.cpp | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 6d0bc63dcad8..b74829697a6b 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -1,11 +1,10 @@ -#include +#include "CodeGen_PTX_Dev.h" #include "CSE.h" #include "CanonicalizeGPUVars.h" #include "CodeGen_GPU_Dev.h" #include "CodeGen_Internal.h" #include "CodeGen_LLVM.h" -#include "CodeGen_PTX_Dev.h" #include "ConciseCasts.h" #include "Debug.h" #include "ExprUsesVar.h" @@ -25,6 +24,7 @@ #include "Target.h" #include +#include namespace Halide { namespace Internal { diff --git a/src/ExtractWMMAOperations.cpp b/src/ExtractWMMAOperations.cpp index 8212c922278d..6567c478c32b 100644 --- a/src/ExtractWMMAOperations.cpp +++ b/src/ExtractWMMAOperations.cpp @@ -1,7 +1,5 @@ #include "ExtractWMMAOperations.h" -#include - #include "CanonicalizeGPUVars.h" #include "FindIntrinsics.h" #include "IREquality.h" @@ -12,6 +10,8 @@ #include "Substitute.h" #include "Util.h" +#include + /** \file Support extraction of NVIDIA tensor core (wmma) instructions. * * The wmma instructions are warp-level: the 32 lanes of a warp cooperate to From a0bd1f42b922d5840b8b2b91633b1c8fbf1c25bf Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 4 Aug 2026 16:56:57 -0700 Subject: [PATCH 55/59] Use the tile directives in the existing tests wmma_matmul's two consumer schedules and most of tiled_matmul_errors were still spelled out as reorder, atomic and vectorize. In tiled_matmul_errors that includes the schedule_matmul helper, which is the valid scaffolding nine scenarios share, so it should read the way such a schedule is meant to be written. scenario_no_matmul keeps the desugared spelling. It puts a value that is not a matrix multiply into tile memory on purpose, which is not something the directives can express. Co-Authored-By: Claude Opus 5 --- test/correctness/tiled_matmul_errors.cpp | 68 +++++++----------------- test/correctness/wmma_matmul.cpp | 6 +-- 2 files changed, 22 insertions(+), 52 deletions(-) diff --git a/test/correctness/tiled_matmul_errors.cpp b/test/correctness/tiled_matmul_errors.cpp index ba11bc84da70..9714e2ffe422 100644 --- a/test/correctness/tiled_matmul_errors.cpp +++ b/test/correctness/tiled_matmul_errors.cpp @@ -45,27 +45,21 @@ void schedule_matmul(Func mm, RVar r, int tile_x, int tile_y, int tile_r) { Var x("x"), y("y"), rxi("rxi"), ryi("ryi"); RVar rri("rri"), rro("rro"); mm.compute_at(mm.in(), x) - .store_in(MemoryType::Tile) .update() .tile(x, y, rxi, ryi, tile_x, tile_y, TailStrategy::GuardWithIf) .split(r, rro, rri, tile_r) - .reorder(rri, rxi, ryi, rro, x, y) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi); + .reorder(rro, x, y) + .tile_matmul(rri, rxi, ryi); Var ixi("ixi"), iyi("iyi"); mm.compute_at(mm.in(), x) .tile(x, y, ixi, iyi, tile_x, tile_y) - .vectorize(ixi) - .vectorize(iyi); + .tile_init(ixi, iyi); Var mmxi("mmxi"), mmyi("mmyi"); mm.in() .tile(x, y, mmxi, mmyi, tile_x, tile_y) - .vectorize(mmxi) - .vectorize(mmyi); + .tile_store(mmxi, mmyi); } // A tile too large for an AMX register (rows > 16). Triggers the explicit @@ -226,27 +220,20 @@ void scenario_mismatched_strides() { Var rxi("rxi"), ryi("ryi"); RVar rri("rri"), rro("rro"); - mm.compute_at(mm.in(), x).store_in(MemoryType::Tile); mm.update(0) .tile(x, y, rxi, ryi, 8, 4, TailStrategy::GuardWithIf) .split(r1.x, rro, rri, 8) - .reorder(rri, rxi, ryi, rro, x, y) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi); + .reorder(rro, x, y) + .tile_matmul(rri, rxi, ryi); mm.update(1) .tile(x, y, rxi, ryi, 4, 8, TailStrategy::GuardWithIf) .split(r2.x, rro, rri, 8) - .reorder(rri, rxi, ryi, rro, x, y) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi); + .reorder(rro, x, y) + .tile_matmul(rri, rxi, ryi); Var ixi("ixi"), iyi("iyi"); - mm.compute_at(mm.in(), x).tile(x, y, ixi, iyi, 8, 8).vectorize(ixi).vectorize(iyi); + mm.compute_at(mm.in(), x).tile(x, y, ixi, iyi, 8, 8).tile_init(ixi, iyi); Var mmxi("mmxi"), mmyi("mmyi"); - mm.in().tile(x, y, mmxi, mmyi, 8, 8).vectorize(mmxi).vectorize(mmyi); + mm.in().tile(x, y, mmxi, mmyi, 8, 8).tile_store(mmxi, mmyi); mm.in().compile_jit(amx_target); } @@ -268,35 +255,26 @@ void scenario_inconsistent_tiles() { Var rxi("rxi"), ryi("ryi"); RVar rri("rri"), rro("rro"); - mm.compute_at(mm.in(), x).store_in(MemoryType::Tile); mm.update(0) .tile(x, y, rxi, ryi, 8, 8, TailStrategy::GuardWithIf) .split(r1.x, rro, rri, 8) - .reorder(rri, rxi, ryi, rro, x, y) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi); + .reorder(rro, x, y) + .tile_matmul(rri, rxi, ryi); mm.update(1) .tile(x, y, rxi, ryi, 4, 4, TailStrategy::GuardWithIf) .split(r2.x, rro, rri, 8) - .reorder(rri, rxi, ryi, rro, x, y) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi); + .reorder(rro, x, y) + .tile_matmul(rri, rxi, ryi); Var ixi("ixi"), iyi("iyi"); mm.compute_at(mm.in(), x) .tile(x, y, ixi, iyi, 8, 8) - .vectorize(ixi) - .vectorize(iyi); + .tile_init(ixi, iyi); Var mmxi("mmxi"), mmyi("mmyi"); mm.in() .tile(x, y, mmxi, mmyi, 8, 8) - .vectorize(mmxi) - .vectorize(mmyi); + .tile_store(mmxi, mmyi); mm.in().compile_jit(amx_target); } @@ -318,27 +296,21 @@ void scenario_not_a_matmul_pattern() { Var rxi("rxi"), ryi("ryi"); RVar rri("rri"), rro("rro"); mm.compute_at(mm.in(), x) - .store_in(MemoryType::Tile) .update() .tile(x, y, rxi, ryi, 8, 8, TailStrategy::GuardWithIf) .split(r.x, rro, rri, 8) - .reorder(rri, rxi, ryi, rro, x, y) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi); + .reorder(rro, x, y) + .tile_matmul(rri, rxi, ryi); Var ixi("ixi"), iyi("iyi"); mm.compute_at(mm.in(), x) .tile(x, y, ixi, iyi, 8, 8) - .vectorize(ixi) - .vectorize(iyi); + .tile_init(ixi, iyi); Var mmxi("mmxi"), mmyi("mmyi"); mm.in() .tile(x, y, mmxi, mmyi, 8, 8) - .vectorize(mmxi) - .vectorize(mmyi); + .tile_store(mmxi, mmyi); mm.in().compile_jit(amx_target); } diff --git a/test/correctness/wmma_matmul.cpp b/test/correctness/wmma_matmul.cpp index bdf9c594f8de..cb4779ab865f 100644 --- a/test/correctness/wmma_matmul.cpp +++ b/test/correctness/wmma_matmul.cpp @@ -356,8 +356,7 @@ bool test_staged_operands() { .reorder(xi, yi, x, y) .unroll(xi) .unroll(yi) - .vectorize(mmxi) - .vectorize(mmyi); + .tile_store(mmxi, mmyi); prod.compute_at(out, x) .tile(x, y, rxi, ryi, tile, tile) @@ -449,8 +448,7 @@ bool test_operand_hoisted_out_of_loop() { .reorder(xi, yi, n, x, y) .unroll(xi) .unroll(yi) - .vectorize(mmxi) - .vectorize(mmyi); + .tile_store(mmxi, mmyi); prod.compute_at(out, n) .tile(x, y, rxi, ryi, tile, tile) From 9a51eaa1d272d0d36684b1542e4d6792a933300f Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Wed, 5 Aug 2026 10:40:28 -0700 Subject: [PATCH 56/59] Pass the lane to the wmma intrinsics The value of each of these depends on which lane of the warp it runs in, but nothing in their arguments did, so the dependence was invisible in the IR. wmma_lane_owns was the worst case: its arguments were three constants, so only its impurity stopped it being lifted out of the loop over lanes, and it would have been reasonable to make it pure. The lane goes last, so the argument indices the backend uses are unchanged. Co-Authored-By: Claude Opus 5 --- src/CodeGen_PTX_Dev.cpp | 2 +- src/ExtractWMMAOperations.cpp | 49 ++++++++++++++++++++++------------- src/ExtractWMMAOperations.h | 4 +-- src/IR.h | 16 +++++++++--- 4 files changed, 46 insertions(+), 25 deletions(-) diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index b74829697a6b..afed95155d23 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -669,7 +669,7 @@ void CodeGen_PTX_Dev::codegen_wmma_store(const Store *op) { // and the column-major layout falls out of the index as usual. Expr index; const Call *inflate = peel_store_permutations(op, &index).as(); - internal_assert(inflate && inflate->args.size() == 4); + internal_assert(inflate && inflate->args.size() == 5); Expr predicate = op->predicate; while (const Shuffle *shuffle = predicate.as()) { predicate = shuffle->vectors[0]; diff --git a/src/ExtractWMMAOperations.cpp b/src/ExtractWMMAOperations.cpp index 6567c478c32b..695fe02e6b22 100644 --- a/src/ExtractWMMAOperations.cpp +++ b/src/ExtractWMMAOperations.cpp @@ -188,9 +188,17 @@ bool is_rhs(const Operand &op, Layout *layout, Expr *stride) { // accumulator has to be executed by all 32 lanes of a warp. Nothing in the // schedule says so - it's a consequence of asking for tensor core storage - so // the loop over lanes is introduced here. -Stmt in_lane_loop(Stmt s) { - return For::make(unique_name("wmma_lane") + gpu_thread_name(0), - 0, warp_lanes - 1, ForType::GPULane, Partition::Never, +// The lane of the warp a tensor core operation runs in. Every wmma intrinsic +// takes it, because their values depend on it and nothing else in their +// arguments does. +Expr make_lane(const string &name) { + return Variable::make(Int(32), name); +} + +Stmt in_lane_loop(const Expr &lane, Stmt s) { + const Variable *v = lane.as(); + internal_assert(v) << "the lane of a tensor core operation is not a variable\n"; + return For::make(v->name, 0, warp_lanes - 1, ForType::GPULane, Partition::Never, DeviceAPI::CUDA, std::move(s)); } @@ -219,7 +227,8 @@ Expr make_matrix_address(const string &name, Type element_type, const Expr &base } Expr make_matrix_to_fragment(Role role, const Shape &shape, Layout layout, - const Load *load, const Expr &base, const Expr &stride) { + const Load *load, const Expr &base, const Expr &stride, + const Expr &lane) { int rows, cols; fragment_matrix_shape(role, shape, &rows, &cols); Expr address = make_matrix_address(load->name, load->type.element_of(), base, @@ -227,7 +236,7 @@ Expr make_matrix_to_fragment(Role role, const Shape &shape, Layout layout, Type type = load->type.element_of().with_lanes( elements_per_lane(role, shape, load->type.element_of())); return Call::make(type, intrinsic_for_role(role), - {shape.M, shape.N, shape.K, std::move(address)}, + {shape.M, shape.N, shape.K, std::move(address), lane}, Call::Intrinsic); } @@ -416,12 +425,14 @@ Stmt convert_to_tile_store(const Store *op, const Expr &store_index, Ramp::make(0, 1, accumulator_elements), {}, {}, const_true(accumulator_elements), {}); const int lanes = shape.M * shape.N; + Expr lane = make_lane(unique_name("wmma_lane") + gpu_thread_name(0)); Expr matrix = Call::make(element_type.with_lanes(lanes), Call::wmma_fragment_to_matrix_d, - {shape.M, shape.N, shape.K, std::move(frag)}, + {shape.M, shape.N, shape.K, std::move(frag), lane}, Call::Intrinsic); Expr owned = Call::make(UInt(1, lanes), Call::wmma_lane_owns, - {shape.M, shape.N, shape.K}, Call::Intrinsic); - return in_lane_loop(Store::make(op->name, std::move(matrix), std::move(index), + {shape.M, shape.N, shape.K, lane}, Call::Intrinsic); + return in_lane_loop(lane, + Store::make(op->name, std::move(matrix), std::move(index), op->param, std::move(owned), ModulusRemainder(), op->is_streaming)); } @@ -541,7 +552,7 @@ class ExtractWMMAOperations : public IRMutator { // The value a matrix multiply uses for one of its operands: the fragment it // was staged in, or a load synthesized here if it wasn't staged. Expr operand_value(const Operand &operand, Role role, const Shape &shape, - Layout layout, const Expr &stride) { + Layout layout, const Expr &stride, const Expr &lane) { if (Fragment *f = find_fragment(operand.load->name)) { const int lanes = f->value_type().lanes(); const string name = @@ -550,7 +561,7 @@ class ExtractWMMAOperations : public IRMutator { const_true(lanes), {}); } return make_matrix_to_fragment(role, shape, layout, operand.load, - operand.mr.base, stride); + operand.mr.base, stride, lane); } Stmt convert_to_fill(const Store *op, Fragment *f) { @@ -568,6 +579,7 @@ class ExtractWMMAOperations : public IRMutator { f, make_matrix_index(dest.base, rows, cols, dest.row_major ? Layout::Row : Layout::Col, dest.stride)); const int lanes = f->value_type().lanes(); + Expr lane = make_lane(unique_name("wmma_lane") + gpu_thread_name(0)); Expr value; if (is_const_zero(op->value)) { // Zeroing a fragment is layout-independent, so it doesn't need an @@ -603,17 +615,18 @@ class ExtractWMMAOperations : public IRMutator { } else { value = make_matrix_to_fragment( f->role, f->shape, mem.row_major ? Layout::Row : Layout::Col, - matrix, mem.base, mem.stride); + matrix, mem.base, mem.stride, lane); } } return in_lane_loop( - Store::make(name, std::move(value), Ramp::make(0, 1, lanes), Parameter(), - const_true(lanes), ModulusRemainder())); + lane, Store::make(name, std::move(value), Ramp::make(0, 1, lanes), Parameter(), + const_true(lanes), ModulusRemainder())); } Stmt convert_to_matmul(const Store *op, Fragment *f, const MatmulInfo &info) { - Expr a = operand_value(info.lhs, Role::A, info.shape, info.lhs_layout, info.lda); - Expr b = operand_value(info.rhs, Role::B, info.shape, info.rhs_layout, info.ldb); + Expr lane = make_lane(unique_name("wmma_lane") + gpu_thread_name(0)); + Expr a = operand_value(info.lhs, Role::A, info.shape, info.lhs_layout, info.lda, lane); + Expr b = operand_value(info.rhs, Role::B, info.shape, info.rhs_layout, info.ldb, lane); Type acc_type = info.accumulator_type.with_lanes(accumulator_elements); Expr frag_idx = Ramp::make(0, 1, accumulator_elements); @@ -624,12 +637,12 @@ class ExtractWMMAOperations : public IRMutator { Expr mma = Call::make(acc_type, Call::wmma_mma, {info.shape.M, info.shape.N, info.shape.K, (int)info.lhs_layout, (int)info.rhs_layout, - std::move(a), std::move(b), std::move(c)}, + std::move(a), std::move(b), std::move(c), lane}, Call::Intrinsic); Stmt store = in_lane_loop( - Store::make(name, std::move(mma), frag_idx, Parameter(), - const_true(accumulator_elements), ModulusRemainder())); + lane, Store::make(name, std::move(mma), frag_idx, Parameter(), + const_true(accumulator_elements), ModulusRemainder())); for (const auto &[let_name, v] : reverse_view(info.peeled_lets)) { store = LetStmt::make(let_name, v, store); } diff --git a/src/ExtractWMMAOperations.h b/src/ExtractWMMAOperations.h index 283ea8deb533..78ffb7767c4a 100644 --- a/src/ExtractWMMAOperations.h +++ b/src/ExtractWMMAOperations.h @@ -44,8 +44,8 @@ int wmma_matrix_arg(const Call *op); * produced by the pass above. Such a store writes the whole matrix, with each * lane of the warp writing the entries it holds: * - * out[matrix] = wmma_fragment_to_matrix(M, N, K, fragment) - * with predicate wmma_lane_owns(M, N, K) + * out[matrix] = wmma_fragment_to_matrix(M, N, K, fragment, lane) + * with predicate wmma_lane_owns(M, N, K, lane) * * wmma_fragment_to_matrix_d permutes this lane's fragment up into a whole * matrix, leaving the entries the lane doesn't hold undefined, and diff --git a/src/IR.h b/src/IR.h index de21260403fa..0bdd4914c09e 100644 --- a/src/IR.h +++ b/src/IR.h @@ -884,14 +884,22 @@ struct Call : public ExprNode { widening_shift_left, widening_shift_right, widening_sub, + // Every intrinsic below takes the lane of the warp it runs in, because + // its value depends on which lane that is and nothing else in its + // arguments does. Without it a call whose other arguments are all + // constants, such as wmma_lane_owns, could be hoisted out of the loop + // over lanes. They are Intrinsic rather than PureIntrinsic only to stop + // lets being hoisted out from under them, which makes instruction + // selection in the backend harder. + // // Permute this lane's tensor core accumulator (d) fragment up into a // whole matrix, leaving the entries this lane doesn't hold undefined. - // wmma_fragment_to_matrix_d(M, N, K, fragment) + // wmma_fragment_to_matrix_d(M, N, K, fragment, lane) wmma_fragment_to_matrix_d, // Whether this lane holds each entry of an M x N tensor core // accumulator in its fragment. Used as the predicate of the store that // copies an accumulator out to memory. - // wmma_lane_owns(M, N, K) + // wmma_lane_owns(M, N, K, lane) wmma_lane_owns, // Take this lane's share of a tensor core fragment out of a matrix. // The a operand is M x K, the b operand is K x N, and the accumulator @@ -903,7 +911,7 @@ struct Call : public ExprNode { // rows or columns are recoverable from the index. The hardware can only // do this as part of a memory read, which is why the argument has to be // a Load rather than an arbitrary matrix value. - // wmma_matrix_to_fragment_a(M, N, K, matrix) + // wmma_matrix_to_fragment_a(M, N, K, matrix, lane) // @{ wmma_matrix_to_fragment_a, wmma_matrix_to_fragment_b, @@ -913,7 +921,7 @@ struct Call : public ExprNode { // and b fragments were taken out of their matrices, which changes how // the hardware arranges them in registers. The accumulator's fragment // layout doesn't depend on how it was loaded, so it isn't a parameter. - // wmma_mma(M, N, K, a_layout, b_layout, a, b, c) + // wmma_mma(M, N, K, a_layout, b_layout, a, b, c, lane) wmma_mma, // keep-sorted end IntrinsicOpCount // Sentinel: keep last. From 917e9de08b8c9b32907add12c055b3679193b685 Mon Sep 17 00:00:00 2001 From: "halide-ci[bot]" <266445882+halide-ci[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 18:23:25 +0000 Subject: [PATCH 57/59] Apply pre-commit auto-fixes --- apps/cuda_mat_mul/CMakeLists.txt | 60 +++++++++++++++++++++----------- 1 file changed, 40 insertions(+), 20 deletions(-) diff --git a/apps/cuda_mat_mul/CMakeLists.txt b/apps/cuda_mat_mul/CMakeLists.txt index 5995bbfd6100..734274d1a8c1 100644 --- a/apps/cuda_mat_mul/CMakeLists.txt +++ b/apps/cuda_mat_mul/CMakeLists.txt @@ -30,30 +30,50 @@ add_halide_generator(mat_mul.generator SOURCES mat_mul_generator.cpp) # Filters. The operand type picks the schedule: half precision gets the tensor # cores, which need compute capability 7.0 or above, and float gets a schedule # that accumulates in ordinary registers. -add_halide_library(mat_mul FROM mat_mul.generator - FEATURES cuda cuda_capability_50 - PARAMS size=1024 A.type=float32 B.type=float32 out.type=float32) -add_halide_library(mat_mul_f16 FROM mat_mul.generator - GENERATOR mat_mul - FEATURES cuda cuda_capability_80 - PARAMS size=1024 A.type=float16 B.type=float16 out.type=float32) +add_halide_library( + mat_mul + FROM mat_mul.generator + FEATURES cuda cuda_capability_50 + PARAMS size=1024 A.type=float32 B.type=float32 out.type=float32 +) +add_halide_library( + mat_mul_f16 + FROM mat_mul.generator + GENERATOR mat_mul + FEATURES cuda cuda_capability_80 + PARAMS size=1024 A.type=float16 B.type=float16 out.type=float32 +) -add_halide_library(mat_mul_f16_acc16 FROM mat_mul.generator - GENERATOR mat_mul - FEATURES cuda cuda_capability_80 - PARAMS size=1024 A.type=float16 B.type=float16 out.type=float16) -add_halide_library(mat_mul_bf16 FROM mat_mul.generator - GENERATOR mat_mul - FEATURES cuda cuda_capability_80 - PARAMS size=1024 A.type=bfloat16 B.type=bfloat16 out.type=float32) -add_halide_library(mat_mul_u8 FROM mat_mul.generator - GENERATOR mat_mul - FEATURES cuda cuda_capability_80 - PARAMS size=1024 A.type=uint8 B.type=uint8 out.type=int32) +add_halide_library( + mat_mul_f16_acc16 + FROM mat_mul.generator + GENERATOR mat_mul + FEATURES cuda cuda_capability_80 + PARAMS size=1024 A.type=float16 B.type=float16 out.type=float16 +) +add_halide_library( + mat_mul_bf16 + FROM mat_mul.generator + GENERATOR mat_mul + FEATURES cuda cuda_capability_80 + PARAMS size=1024 A.type=bfloat16 B.type=bfloat16 out.type=float32 +) +add_halide_library( + mat_mul_u8 + FROM mat_mul.generator + GENERATOR mat_mul + FEATURES cuda cuda_capability_80 + PARAMS size=1024 A.type=uint8 B.type=uint8 out.type=int32 +) # Main executable add_executable(runner runner.cpp) -target_link_libraries(runner PRIVATE mat_mul mat_mul_f16 mat_mul_f16_acc16 mat_mul_bf16 mat_mul_u8 Halide::Halide Halide::Tools CUDA::cudart CUDA::cublas) +target_link_libraries( + runner + PRIVATE + mat_mul mat_mul_f16 mat_mul_f16_acc16 mat_mul_bf16 mat_mul_u8 Halide::Halide Halide::Tools + CUDA::cudart CUDA::cublas +) # Test that the app actually works! add_test(NAME mat_mul COMMAND runner) From b06c6931810bf0e13c2e58fd93d2775d09f87257 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Wed, 5 Aug 2026 11:37:37 -0700 Subject: [PATCH 58/59] Use the short form of Load::make in the multiramp test Load::make lost the default for is_streaming when the short forms were added, so the call that left it off no longer compiles. The short form is what it wanted anyway: an unpredicated load from an internal buffer. Co-Authored-By: Claude Opus 5 (cherry picked from commit 3fe9e5414b17936403896417635b33ea2e8cb32c) --- test/correctness/multiramp.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/test/correctness/multiramp.cpp b/test/correctness/multiramp.cpp index 8f325f293c10..da7a37c8cc08 100644 --- a/test/correctness/multiramp.cpp +++ b/test/correctness/multiramp.cpp @@ -712,8 +712,7 @@ void check_reject_gather_shuffle() { // ---- is_load_of_multiramp ------------------------------------------------ Expr make_test_load(int lanes) { - return Load::make(Int(16, lanes), "buf", Ramp::make(Expr(0), Expr(1), lanes), - Buffer<>(), Parameter(), const_true(lanes), ModulusRemainder()); + return Load::make(Int(16, lanes), "buf", Ramp::make(Expr(0), Expr(1), lanes)); } void check_load_under_cast_and_broadcast() { From e6925084dad58bbabd6ed96824736a1f9a9fc47a Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Wed, 5 Aug 2026 11:40:10 -0700 Subject: [PATCH 59/59] Use the short forms of Load::make and Store::make Both lost the default for is_streaming when the short forms were added, so the calls that left it off no longer compile. Most of these are unpredicated accesses to internal buffers, which is what the short form means. The load of a matrix a tensor core instruction takes a fragment from has a real image and parameter, so it keeps the long form and passes is_streaming explicitly. Co-Authored-By: Claude Opus 5 --- src/ExtractWMMAOperations.cpp | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/src/ExtractWMMAOperations.cpp b/src/ExtractWMMAOperations.cpp index 695fe02e6b22..7d15427b2ff1 100644 --- a/src/ExtractWMMAOperations.cpp +++ b/src/ExtractWMMAOperations.cpp @@ -223,7 +223,7 @@ Expr make_matrix_address(const string &name, Type element_type, const Expr &base Expr index = make_matrix_index(base, rows, cols, layout, stride); const int lanes = rows * cols; return Load::make(element_type.with_lanes(lanes), name, index, image, param, - const_true(lanes), ModulusRemainder()); + const_true(lanes), ModulusRemainder(), false); } Expr make_matrix_to_fragment(Role role, const Shape &shape, Layout layout, @@ -422,8 +422,7 @@ Stmt convert_to_tile_store(const Store *op, const Expr &store_index, Expr index = make_matrix_index(mem.base, shape.M, shape.N, layout, mem.stride); Type element_type = op->value.type().element_of(); Expr frag = Load::make(element_type.with_lanes(accumulator_elements), new_name, - Ramp::make(0, 1, accumulator_elements), {}, {}, - const_true(accumulator_elements), {}); + Ramp::make(0, 1, accumulator_elements)); const int lanes = shape.M * shape.N; Expr lane = make_lane(unique_name("wmma_lane") + gpu_thread_name(0)); Expr matrix = Call::make(element_type.with_lanes(lanes), Call::wmma_fragment_to_matrix_d, @@ -557,8 +556,7 @@ class ExtractWMMAOperations : public IRMutator { const int lanes = f->value_type().lanes(); const string name = operand_subtile_name(f, operand.mr.base, role, shape, layout, stride); - return Load::make(f->value_type(), name, Ramp::make(0, 1, lanes), {}, {}, - const_true(lanes), {}); + return Load::make(f->value_type(), name, Ramp::make(0, 1, lanes)); } return make_matrix_to_fragment(role, shape, layout, operand.load, operand.mr.base, stride, lane); @@ -619,8 +617,7 @@ class ExtractWMMAOperations : public IRMutator { } } return in_lane_loop( - lane, Store::make(name, std::move(value), Ramp::make(0, 1, lanes), Parameter(), - const_true(lanes), ModulusRemainder())); + lane, Store::make(name, std::move(value), Ramp::make(0, 1, lanes))); } Stmt convert_to_matmul(const Store *op, Fragment *f, const MatmulInfo &info) { @@ -631,8 +628,7 @@ class ExtractWMMAOperations : public IRMutator { Type acc_type = info.accumulator_type.with_lanes(accumulator_elements); Expr frag_idx = Ramp::make(0, 1, accumulator_elements); const string name = subtile_name(f, op->index); - Expr c = Load::make(acc_type, name, frag_idx, {}, {}, - const_true(accumulator_elements), {}); + Expr c = Load::make(acc_type, name, frag_idx); Expr mma = Call::make(acc_type, Call::wmma_mma, {info.shape.M, info.shape.N, info.shape.K, @@ -641,8 +637,7 @@ class ExtractWMMAOperations : public IRMutator { Call::Intrinsic); Stmt store = in_lane_loop( - lane, Store::make(name, std::move(mma), frag_idx, Parameter(), - const_true(accumulator_elements), ModulusRemainder())); + lane, Store::make(name, std::move(mma), frag_idx)); for (const auto &[let_name, v] : reverse_view(info.peeled_lets)) { store = LetStmt::make(let_name, v, store); }