From 057fec51d28b1f095a08922fcc7683366a21de9a Mon Sep 17 00:00:00 2001 From: Alex Reinking Date: Tue, 14 Jul 2026 14:51:43 -0400 Subject: [PATCH 1/4] Add stream_loads and stream_stores scheduling directives Co-authored-by: GPT 5.6 Sol Co-Authored-By: Claude Sonnet 5 --- python_bindings/src/halide/halide_/PyFunc.cpp | 3 + .../src/halide/halide_/PyStage.cpp | 5 + src/AddAtomicMutex.cpp | 3 +- src/AlignLoads.cpp | 2 +- src/CSE.cpp | 2 +- src/CodeGen_ARM.cpp | 4 +- src/CodeGen_C.cpp | 3 + src/CodeGen_D3D12Compute_Dev.cpp | 8 +- src/CodeGen_GPU_Dev.cpp | 3 +- src/CodeGen_Hexagon.cpp | 8 +- src/CodeGen_LLVM.cpp | 71 +++++- src/CodeGen_LLVM.h | 9 +- src/CodeGen_PTX_Dev.cpp | 8 +- src/CodeGen_X86.cpp | 6 + src/CompilerLogger.cpp | 2 +- src/Deinterleave.cpp | 14 +- src/Deserialization.cpp | 14 +- src/EliminateBoolVectors.cpp | 4 +- src/Expr.h | 2 + src/FindIntrinsics.cpp | 4 +- src/FlattenNestedRamps.cpp | 8 +- src/Func.cpp | 63 ++++++ src/Func.h | 40 ++++ src/FuseGPUThreadLoops.cpp | 12 +- src/FuzzFloatStores.cpp | 2 +- src/HexagonOffload.cpp | 4 +- src/HexagonOptimize.cpp | 2 +- src/IR.cpp | 41 +++- src/IR.h | 51 ++++- src/IREquality.cpp | 10 + src/IRMutator.cpp | 22 +- src/IRPrinter.cpp | 31 +++ src/IRPrinter.h | 2 + src/IRVisitor.cpp | 16 ++ src/Inline.cpp | 12 + src/LLVM_Headers.h | 3 + src/LowerWarpShuffles.cpp | 2 +- src/OptimizeShuffles.cpp | 4 +- src/PartitionLoops.cpp | 2 +- src/RemoveUndef.cpp | 6 +- src/Schedule.cpp | 27 ++- src/Schedule.h | 25 ++- src/ScheduleFunctions.cpp | 42 +++- src/Serialization.cpp | 27 ++- src/Simplify_Exprs.cpp | 8 +- src/Simplify_Internal.h | 2 + src/Simplify_Shuffle.cpp | 2 +- src/Simplify_Stmts.cpp | 31 ++- src/StageStridedLoads.cpp | 14 +- src/StorageFlattening.cpp | 70 +++++- src/VectorizeLoops.cpp | 20 +- src/halide_ir.fbs | 5 + test/correctness/CMakeLists.txt | 7 + test/correctness/simd_op_check.h | 5 +- test/correctness/simd_op_check_arm.cpp | 26 +++ test/correctness/simd_op_check_x86.cpp | 27 +++ test/correctness/streaming.cpp | 193 ++++++++++++++++ test/correctness/streaming_in.cpp | 210 ++++++++++++++++++ test/correctness/streaming_specialize.cpp | 85 +++++++ test/error/CMakeLists.txt | 2 + test/error/stream_loads_illegal_self.cpp | 17 ++ test/error/stream_stores_illegal_rvar.cpp | 20 ++ test/performance/memcpy.cpp | 41 +++- 63 files changed, 1292 insertions(+), 122 deletions(-) create mode 100644 test/correctness/streaming.cpp create mode 100644 test/correctness/streaming_in.cpp create mode 100644 test/correctness/streaming_specialize.cpp create mode 100644 test/error/stream_loads_illegal_self.cpp create mode 100644 test/error/stream_stores_illegal_rvar.cpp diff --git a/python_bindings/src/halide/halide_/PyFunc.cpp b/python_bindings/src/halide/halide_/PyFunc.cpp index a1e222f5c8f0..9daa3fae3e6f 100644 --- a/python_bindings/src/halide/halide_/PyFunc.cpp +++ b/python_bindings/src/halide/halide_/PyFunc.cpp @@ -228,6 +228,9 @@ void define_func(py::module &m) { .def("hoist_storage_root", &Func::hoist_storage_root) .def("store_in", &Func::store_in, py::arg("memory_type")) + .def("stream_loads", (Func & (Func::*)()) & Func::stream_loads) + .def("stream_loads", (Func & (Func::*)(const std::vector &)) & Func::stream_loads, py::arg("funcs")) + .def("stream_stores", &Func::stream_stores) .def("compile_to", // [](Func &f, const std::map &output_files, const std::vector &args, const std::string &fn_name, const Target &target) { diff --git a/python_bindings/src/halide/halide_/PyStage.cpp b/python_bindings/src/halide/halide_/PyStage.cpp index 9a9f6394d7cd..93629beb90a2 100644 --- a/python_bindings/src/halide/halide_/PyStage.cpp +++ b/python_bindings/src/halide/halide_/PyStage.cpp @@ -14,6 +14,11 @@ void define_stage(py::module &m) { .def("dump_argument_list", &Stage::dump_argument_list) .def("name", &Stage::name) + .def("stream_stores", &Stage::stream_stores) + .def("stream_loads", static_cast(&Stage::stream_loads)) + .def("stream_loads", static_cast &)>(&Stage::stream_loads), + py::arg("funcs")) + .def("rfactor", static_cast> &)>(&Stage::rfactor), py::arg("preserved")) .def("rfactor", static_cast(&Stage::rfactor), diff --git a/src/AddAtomicMutex.cpp b/src/AddAtomicMutex.cpp index b2243725d489..c775a2643c62 100644 --- a/src/AddAtomicMutex.cpp +++ b/src/AddAtomicMutex.cpp @@ -266,7 +266,8 @@ class ReplaceStoreIndexWithVar : public IRMutator { var, op->param, std::move(predicate), - op->alignment); + op->alignment, + op->is_streaming); } const std::string &producer_name; diff --git a/src/AlignLoads.cpp b/src/AlignLoads.cpp index 467984923f36..3c48842e76d7 100644 --- a/src/AlignLoads.cpp +++ b/src/AlignLoads.cpp @@ -39,7 +39,7 @@ class AlignLoads : public IRMutator { return mutate(Load::make(load->type.with_lanes(index.type().lanes()), load->name, index, load->image, load->param, const_true(index.type().lanes()), - alignment)); + alignment, load->is_streaming)); } Expr visit(const Load *op) override { diff --git a/src/CSE.cpp b/src/CSE.cpp index 95f61f0214e9..a4dcd41eaea6 100644 --- a/src/CSE.cpp +++ b/src/CSE.cpp @@ -268,7 +268,7 @@ class CSEEveryExprInStmt : public IRMutator { } Stmt s = Store::make(op->name, value, index, - op->param, mutate(op->predicate), op->alignment); + op->param, mutate(op->predicate), op->alignment, op->is_streaming); for (const auto &[var, val] : deferred) { s = LetStmt::make(var, val, s); diff --git a/src/CodeGen_ARM.cpp b/src/CodeGen_ARM.cpp index dbc8458a999e..92bddddecfb5 100644 --- a/src/CodeGen_ARM.cpp +++ b/src/CodeGen_ARM.cpp @@ -1805,7 +1805,7 @@ void CodeGen_ARM::visit(const Store *op) { if (is_float16_and_has_feature(elt)) { Type u16_type = op->value.type().with_code(halide_type_uint); Expr v = reinterpret(u16_type, op->value); - codegen(Store::make(op->name, v, op->index, op->param, op->predicate, op->alignment)); + codegen(Store::make(op->name, v, op->index, op->param, op->predicate, op->alignment, op->is_streaming)); return; } @@ -1958,7 +1958,7 @@ void CodeGen_ARM::visit(const Load *op) { // Rewrite float16 case into load in uint16 and reinterpret, as it is unsupported in LLVM if (is_float16_and_has_feature(op->type)) { Type u16_type = op->type.with_code(halide_type_uint); - Expr equiv = Load::make(u16_type, op->name, op->index, op->image, op->param, op->predicate, op->alignment); + Expr equiv = Load::make(u16_type, op->name, op->index, op->image, op->param, op->predicate, op->alignment, op->is_streaming); equiv = reinterpret(op->type, equiv); equiv = common_subexpression_elimination(equiv); value = codegen(equiv); diff --git a/src/CodeGen_C.cpp b/src/CodeGen_C.cpp index 7b5a9792247e..6c406a449fd6 100644 --- a/src/CodeGen_C.cpp +++ b/src/CodeGen_C.cpp @@ -1863,6 +1863,9 @@ void CodeGen_C::visit(const Call *op) { << " + " << print_expr(base_offset) << "), /*rw*/0, /*locality*/0), 0)"; } else if (op->is_intrinsic(Call::size_of_halide_buffer_t)) { rhs << "(sizeof(halide_buffer_t))"; + } else if (op->is_intrinsic(Call::stream_store_fence)) { + // The C backend does not implement non-temporal loads/stores. + rhs << "0"; } else if (op->is_intrinsic(Call::strict_fma)) { internal_assert(op->args.size() == 3) << "Wrong number of args for strict_fma: " << op->args.size(); diff --git a/src/CodeGen_D3D12Compute_Dev.cpp b/src/CodeGen_D3D12Compute_Dev.cpp index f9c8ab93c7e5..4023fa9dde46 100644 --- a/src/CodeGen_D3D12Compute_Dev.cpp +++ b/src/CodeGen_D3D12Compute_Dev.cpp @@ -1803,7 +1803,7 @@ void CodeGen_D3D12Compute_Dev::CodeGen_D3D12Compute_C::add_kernel(Stmt s, if (it != replacements.end()) { return Load::make(op->type, it->second, mutate(op->index), op->image, op->param, - mutate(op->predicate), op->alignment); + mutate(op->predicate), op->alignment, op->is_streaming); } else { return IRMutator::visit(op); } @@ -1814,7 +1814,7 @@ void CodeGen_D3D12Compute_Dev::CodeGen_D3D12Compute_C::add_kernel(Stmt s, if (it != replacements.end()) { return Store::make(it->second, mutate(op->value), mutate(op->index), op->param, - mutate(op->predicate), op->alignment); + mutate(op->predicate), op->alignment, op->is_streaming); } else { return IRMutator::visit(op); } @@ -1931,7 +1931,7 @@ void CodeGen_D3D12Compute_Dev::CodeGen_D3D12Compute_C::add_kernel(Stmt s, if (it != renames.end()) { return Load::make(op->type, it->second, mutate(op->index), op->image, op->param, - mutate(op->predicate), op->alignment); + mutate(op->predicate), op->alignment, op->is_streaming); } return IRMutator::visit(op); } @@ -1940,7 +1940,7 @@ void CodeGen_D3D12Compute_Dev::CodeGen_D3D12Compute_C::add_kernel(Stmt s, if (it != renames.end()) { return Store::make(it->second, mutate(op->value), mutate(op->index), op->param, - mutate(op->predicate), op->alignment); + mutate(op->predicate), op->alignment, op->is_streaming); } return IRMutator::visit(op); } diff --git a/src/CodeGen_GPU_Dev.cpp b/src/CodeGen_GPU_Dev.cpp index 595dbb82b9c5..35bc210bdc05 100644 --- a/src/CodeGen_GPU_Dev.cpp +++ b/src/CodeGen_GPU_Dev.cpp @@ -118,7 +118,8 @@ class ScalarizePredicatedLoadStore : public IRMutator { op->image, op->param, const_true(), - op->alignment + ln); + op->alignment + ln, + op->is_streaming); lane_values.push_back(Call::make(load_expr.type(), Call::if_then_else, {extract_lane(op->predicate, ln), diff --git a/src/CodeGen_Hexagon.cpp b/src/CodeGen_Hexagon.cpp index c680fdb1a82e..bcf95e66470d 100644 --- a/src/CodeGen_Hexagon.cpp +++ b/src/CodeGen_Hexagon.cpp @@ -275,7 +275,7 @@ class SloppyUnpredicateLoadsAndStores : public IRMutator { } Expr load = Load::make(op->type, op->name, index, op->image, op->param, - const_true(op->type.lanes()), op->alignment); + const_true(op->type.lanes()), op->alignment, op->is_streaming); return Call::make(op->type, Call::if_then_else, {condition, load}, Call::PureIntrinsic); @@ -286,7 +286,7 @@ class SloppyUnpredicateLoadsAndStores : public IRMutator { // introducing a set of runtime functions to do predicated // loads. Expr load = Load::make(op->type, op->name, index, op->image, op->param, - const_true(op->type.lanes()), op->alignment); + const_true(op->type.lanes()), op->alignment, op->is_streaming); return Call::make(op->type, Call::if_then_else, {predicate, load}, Call::PureIntrinsic); } @@ -304,14 +304,14 @@ class SloppyUnpredicateLoadsAndStores : public IRMutator { int lanes = value.type().lanes(); if (const Broadcast *scalar_pred = predicate.as()) { - Stmt unpredicated_store = Store::make(op->name, value, index, op->param, const_true(lanes), op->alignment); + Stmt unpredicated_store = Store::make(op->name, value, index, op->param, const_true(lanes), op->alignment, op->is_streaming); return IfThenElse::make(scalar_pred->value, unpredicated_store); } if (predicate.same_as(op->predicate) && value.same_as(op->value) && index.same_as(op->index)) { return op; } else { - return Store::make(op->name, value, index, op->param, predicate, op->alignment); + return Store::make(op->name, value, index, op->param, predicate, op->alignment, op->is_streaming); } } }; diff --git a/src/CodeGen_LLVM.cpp b/src/CodeGen_LLVM.cpp index 9263aef95f75..2b3f9353265f 100644 --- a/src/CodeGen_LLVM.cpp +++ b/src/CodeGen_LLVM.cpp @@ -134,6 +134,13 @@ using std::vector; namespace { +// AArch64's full-width non-temporal vector accesses are LDNP/STNP pairs of +// 128-bit NEON registers. Keeping streaming slices at least this wide gives +// LLVM a 256-bit access that it can select to those instructions instead of +// splitting it into ordinary 128-bit loads or stores before instruction +// selection. Other backends can still legalize this width as needed. +constexpr int minimum_streaming_vector_bits = 256; + llvm::Value *CreateConstGEP1_32(IRBuilderBase *builder, llvm::Type *gep_type, Value *ptr, unsigned index) { return builder->CreateConstGEP1_32(gep_type, ptr, index); @@ -2063,13 +2070,22 @@ void CodeGen_LLVM::function_does_not_access_memory(llvm::Function *fn) { fn->addFnAttr("memory(none)"); } +void CodeGen_LLVM::emit_streaming_store_fence() { +} + +void CodeGen_LLVM::add_streaming_metadata(llvm::Instruction *inst) { + llvm::MDNode *nontemporal_node = llvm::MDNode::get( + *context, llvm::ConstantAsMetadata::get(ConstantInt::get(i32_t, 1))); + inst->setMetadata(llvm::LLVMContext::MD_nontemporal, nontemporal_node); +} + void CodeGen_LLVM::visit(const Load *op) { // If the type should be stored as some other type, insert a reinterpret cast. Type storage_type = upgrade_type_for_storage(op->type); if (op->type != storage_type) { codegen(reinterpret(op->type, Load::make(storage_type, op->name, op->index, op->image, - op->param, op->predicate, op->alignment))); + op->param, op->predicate, op->alignment, op->is_streaming))); return; } @@ -2085,6 +2101,9 @@ void CodeGen_LLVM::visit(const Load *op) { Value *ptr = codegen_buffer_pointer(op->name, op->type, op->index); LoadInst *load = builder->CreateAlignedLoad(llvm_type_of(op->type), ptr, llvm::Align(op->type.bytes())); add_tbaa_metadata(load, op->name, op->index); + if (op->is_streaming) { + add_streaming_metadata(load); + } value = load; } else { const Ramp *ramp = op->index.as(); @@ -2101,7 +2120,7 @@ void CodeGen_LLVM::visit(const Load *op) { ModulusRemainder align = op->alignment; // Switch to the alignment of the last lane align = align - (ramp->lanes - 1); - Expr flipped_load = Load::make(op->type, op->name, flipped_index, op->image, op->param, op->predicate, align); + Expr flipped_load = Load::make(op->type, op->name, flipped_index, op->image, op->param, op->predicate, align, op->is_streaming); Value *flipped = codegen(flipped_load); @@ -2115,6 +2134,9 @@ void CodeGen_LLVM::visit(const Load *op) { Value *lane = ConstantInt::get(i32_t, i); LoadInst *val = builder->CreateLoad(load_type, ptr); add_tbaa_metadata(val, op->name, op->index); + if (op->is_streaming) { + add_streaming_metadata(val); + } value = builder->CreateInsertElement(value, val, lane); ptr = CreateInBoundsGEP(builder.get(), load_type, ptr, stride); } @@ -2130,6 +2152,9 @@ void CodeGen_LLVM::visit(const Load *op) { Value *ptr = codegen_buffer_pointer(op->name, op->type.element_of(), idx); LoadInst *val = builder->CreateLoad(load_type, ptr); add_tbaa_metadata(val, op->name, op->index); + if (op->is_streaming) { + add_streaming_metadata(val); + } vec = builder->CreateInsertElement(vec, val, ConstantInt::get(i32_t, i)); } value = vec; @@ -2142,6 +2167,9 @@ void CodeGen_LLVM::visit(const Load *op) { Value *ptr = codegen_buffer_pointer(op->name, op->type.element_of(), idx); LoadInst *val = builder->CreateLoad(load_type, ptr); add_tbaa_metadata(val, op->name, op->index); + if (op->is_streaming) { + add_streaming_metadata(val); + } vec = builder->CreateInsertElement(vec, val, ConstantInt::get(i32_t, i)); } value = vec; @@ -2514,6 +2542,9 @@ void CodeGen_LLVM::codegen_predicated_store(const Store *op) { // width, bust them up into native vectors. int store_lanes = value_type.lanes(); int native_lanes = maximum_vector_bits() / value_type.bits(); + if (op->is_streaming) { + native_lanes = std::max(native_lanes, minimum_streaming_vector_bits / value_type.bits()); + } for (int i = 0; i < store_lanes; i += native_lanes) { int slice_lanes = std::min(native_lanes, store_lanes - i); @@ -2538,6 +2569,9 @@ void CodeGen_LLVM::codegen_predicated_store(const Store *op) { store = builder->CreateMaskedStore(slice_val, vec_ptr, llvm::Align(alignment), slice_mask); } add_tbaa_metadata(store, op->name, slice_index); + if (op->is_streaming) { + add_streaming_metadata(store); + } } } else { // It's not dense vector store, we need to scalarize it debug(4) << "Scalarize predicated vector store\n"; @@ -2572,6 +2606,8 @@ void CodeGen_LLVM::codegen_predicated_store(const Store *op) { StoreInst *store = builder->CreateAlignedStore(v, ptr, llvm::Align(value_type.bytes())); if (emit_atomic_stores) { store->setAtomic(AtomicOrdering::Monotonic); + } else if (op->is_streaming) { + add_streaming_metadata(store); } builder->CreateBr(after_bb); @@ -2582,7 +2618,7 @@ void CodeGen_LLVM::codegen_predicated_store(const Store *op) { llvm::Value *CodeGen_LLVM::codegen_vector_load(const Type &type, const std::string &name, const Expr &base, const Buffer<> &image, const Parameter ¶m, const ModulusRemainder &alignment, - llvm::Value *vpred, bool slice_to_native, llvm::Value *stride) { + bool is_streaming, llvm::Value *vpred, bool slice_to_native, llvm::Value *stride) { debug(4) << "Vectorize predicated dense vector load:\n\t" << "(" << type << ")" << name << "[ramp(base, 1, " << type.lanes() << ")]\n"; int align_bytes = type.bytes(); // The size of a single element @@ -2619,6 +2655,9 @@ llvm::Value *CodeGen_LLVM::codegen_vector_load(const Type &type, const std::stri // width, bust them up into native vectors int load_lanes = type.lanes(); int native_lanes = slice_to_native ? std::max(1, maximum_vector_bits() / type.bits()) : load_lanes; + if (slice_to_native && is_streaming) { + native_lanes = std::max(native_lanes, minimum_streaming_vector_bits / type.bits()); + } vector slices; for (int i = 0; i < load_lanes; i += native_lanes) { int slice_lanes = std::min(native_lanes, load_lanes - i); @@ -2664,6 +2703,9 @@ llvm::Value *CodeGen_LLVM::codegen_vector_load(const Type &type, const std::stri } } add_tbaa_metadata(load_inst, name, slice_index); + if (is_streaming) { + add_streaming_metadata(load_inst); + } slices.push_back(load_inst); } value = concat_vectors(slices); @@ -2675,7 +2717,7 @@ Value *CodeGen_LLVM::codegen_dense_vector_load(const Load *load, Value *vpred, b internal_assert(ramp && is_const_one(ramp->stride)) << "Should be dense vector load\n"; return codegen_vector_load(load->type, load->name, ramp->base, load->image, load->param, - load->alignment, vpred, slice_to_native, nullptr); + load->alignment, load->is_streaming, vpred, slice_to_native, nullptr); } void CodeGen_LLVM::codegen_predicated_load(const Load *op) { @@ -2689,7 +2731,7 @@ void CodeGen_LLVM::codegen_predicated_load(const Load *op) { Value *vpred = codegen(op->predicate); Value *llvm_stride = codegen(stride); // Not 1 (dense) as that was caught above. value = codegen_vector_load(op->type, op->name, ramp->base, op->image, op->param, - op->alignment, vpred, true, llvm_stride); + op->alignment, op->is_streaming, vpred, true, llvm_stride); } else if (ramp && stride && stride->value == -1) { debug(4) << "Predicated dense vector load with stride -1\n\t" << Expr(op) << "\n"; @@ -2705,13 +2747,13 @@ void CodeGen_LLVM::codegen_predicated_load(const Load *op) { align = align - (ramp->lanes - 1); Expr flipped_load = Load::make(op->type, op->name, flipped_index, op->image, - op->param, const_true(op->type.lanes()), align); + op->param, const_true(op->type.lanes()), align, op->is_streaming); Value *flipped = codegen_dense_vector_load(flipped_load.as(), vpred); value = reverse_vector(flipped); } else { // It's not dense vector load, we need to scalarize it Expr load_expr = Load::make(op->type, op->name, op->index, op->image, - op->param, const_true(op->type.lanes()), op->alignment); + op->param, const_true(op->type.lanes()), op->alignment, op->is_streaming); debug(4) << "Scalarize predicated vector load\n\t" << load_expr << "\n"; Expr pred_load = Call::make(load_expr.type(), Call::if_then_else, @@ -3519,6 +3561,9 @@ void CodeGen_LLVM::visit(const Call *op) { } else if (op->is_intrinsic(Call::size_of_halide_buffer_t)) { const llvm::DataLayout &d = module->getDataLayout(); value = ConstantInt::get(i32_t, (int)d.getTypeAllocSize(halide_buffer_t_type)); + } else if (op->is_intrinsic(Call::stream_store_fence)) { + emit_streaming_store_fence(); + value = ConstantInt::get(i32_t, 0); } else if (op->is_strict_float_intrinsic()) { // Evaluate the args first outside the strict scope, as they may use // non-strict operations. @@ -4007,7 +4052,7 @@ void CodeGen_LLVM::visit(const Store *op) { // Peel lets off the index to make us more likely to pattern // match a ramp. if (const Let *let = op->index.as()) { - Stmt s = Store::make(op->name, op->value, let->body, op->param, op->predicate, op->alignment); + Stmt s = Store::make(op->name, op->value, let->body, op->param, op->predicate, op->alignment, op->is_streaming); codegen(LetStmt::make(let->name, let->value, s)); return; } @@ -4018,7 +4063,7 @@ void CodeGen_LLVM::visit(const Store *op) { Halide::Type storage_type = upgrade_type_for_storage(value_type); if (value_type != storage_type) { Expr v = reinterpret(storage_type, op->value); - codegen(Store::make(op->name, v, op->index, op->param, op->predicate, op->alignment)); + codegen(Store::make(op->name, v, op->index, op->param, op->predicate, op->alignment, op->is_streaming)); return; } @@ -4044,6 +4089,9 @@ void CodeGen_LLVM::visit(const Store *op) { auto annotate_store = [&](StoreInst *store, const Expr &index) { add_tbaa_metadata(store, op->name, index); + if (op->is_streaming && !emit_atomic_stores) { + add_streaming_metadata(store); + } if (emit_atomic_stores) { store->setAtomic(AtomicOrdering::Monotonic); } @@ -4057,7 +4105,7 @@ void CodeGen_LLVM::visit(const Store *op) { StoreInst *store = builder->CreateAlignedStore(val, ptr, llvm::Align(value_type.bytes())); annotate_store(store, op->index); } else if (const Let *let = op->index.as()) { - Stmt s = Store::make(op->name, op->value, let->body, op->param, op->predicate, op->alignment); + Stmt s = Store::make(op->name, op->value, let->body, op->param, op->predicate, op->alignment, op->is_streaming); codegen(LetStmt::make(let->name, let->value, s)); } else { int alignment = value_type.bytes(); @@ -4093,6 +4141,9 @@ void CodeGen_LLVM::visit(const Store *op) { // width, bust them up into native vectors. int store_lanes = value_type.lanes(); int native_lanes = maximum_vector_bits() / value_type.bits(); + if (op->is_streaming) { + native_lanes = std::max(native_lanes, minimum_streaming_vector_bits / value_type.bits()); + } Expr base = ramp ? ramp->base : 0; Expr stride = ramp ? ramp->stride : 0; diff --git a/src/CodeGen_LLVM.h b/src/CodeGen_LLVM.h index d01ee63b3ce2..7389bc07a0ba 100644 --- a/src/CodeGen_LLVM.h +++ b/src/CodeGen_LLVM.h @@ -179,6 +179,9 @@ class CodeGen_LLVM : public IRVisitor { bool in_strict_float = false; bool any_strict_float = false; + /** Emit any target-specific fence required after non-temporal stores. */ + virtual void emit_streaming_store_fence(); + /** Change floating-point math op emission to use fast flags. */ void set_fast_fp_math(); @@ -708,6 +711,9 @@ class CodeGen_LLVM : public IRVisitor { * vectors. Used by CodeGen_ARM to help with vld2/3/4 emission. */ llvm::Value *codegen_dense_vector_load(const Load *load, llvm::Value *vpred = nullptr, bool slice_to_native = true); + /** Attach LLVM's non-temporal metadata to a memory instruction. */ + void add_streaming_metadata(llvm::Instruction *inst); + /** Warning messages which we want to avoid displaying number of times */ enum class WarningKind { EmulatedFloat16, @@ -764,7 +770,8 @@ class CodeGen_LLVM : public IRVisitor { llvm::Value *codegen_vector_load(const Type &type, const std::string &name, const Expr &base, const Buffer<> &image, const Parameter ¶m, const ModulusRemainder &alignment, - llvm::Value *vpred = nullptr, bool slice_to_native = true, llvm::Value *stride = nullptr); + bool is_streaming, llvm::Value *vpred = nullptr, + bool slice_to_native = true, llvm::Value *stride = nullptr); virtual void codegen_predicated_load(const Load *op); virtual void codegen_predicated_store(const Store *op); diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 0786ade82847..0414b33fd243 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -378,7 +378,7 @@ void CodeGen_PTX_Dev::visit(const Load *op) { if (align.modulus % 4 == 0 && align.remainder % 4 == 0) { Expr index = simplify(r->base / 4); Expr equiv = Load::make(UInt(128), op->name, index, - op->image, op->param, const_true(), align / 4); + op->image, op->param, const_true(), align / 4, op->is_streaming); equiv = reinterpret(op->type, equiv); codegen(equiv); return; @@ -403,7 +403,7 @@ void CodeGen_PTX_Dev::visit(const Store *op) { if (align.modulus % 4 == 0 && align.remainder % 4 == 0) { Expr index = simplify(r->base / 4); Expr value = reinterpret(UInt(128), op->value); - Stmt equiv = Store::make(op->name, value, index, op->param, const_true(), align / 4); + Stmt equiv = Store::make(op->name, value, index, op->param, const_true(), align / 4, op->is_streaming); codegen(equiv); return; } @@ -447,12 +447,12 @@ class RewriteLoadsAs32Bit : public IRMutator { if (op->type.lanes() > sub_lanes) { new_idx = Ramp::make(new_idx, 1, load_lanes); } - Expr new_load = Load::make(Int(32, load_lanes), op->name, new_idx, op->image, op->param, const_true(load_lanes), op->alignment / sub_lanes); + Expr new_load = Load::make(Int(32, load_lanes), op->name, new_idx, op->image, op->param, const_true(load_lanes), op->alignment / sub_lanes, op->is_streaming); return reinterpret(op->type, new_load); } else if (index.same_as(op->index)) { return op; } else { - return Load::make(op->type, op->name, std::move(index), op->image, op->param, op->predicate, op->alignment); + return Load::make(op->type, op->name, std::move(index), op->image, op->param, op->predicate, op->alignment, op->is_streaming); } } }; diff --git a/src/CodeGen_X86.cpp b/src/CodeGen_X86.cpp index 4ff685a7a4e2..eed8cc18fdec 100644 --- a/src/CodeGen_X86.cpp +++ b/src/CodeGen_X86.cpp @@ -93,6 +93,7 @@ class CodeGen_X86 : public CodeGen_CPU { using CodeGen_CPU::visit; void init_module() override; + void emit_streaming_store_fence() override; /** Nodes for which we want to emit specific sse/avx intrinsics */ // @{ @@ -336,6 +337,11 @@ void CodeGen_X86::init_module() { } } +void CodeGen_X86::emit_streaming_store_fence() { + llvm::Function *sfence = llvm::Intrinsic::getOrInsertDeclaration(module.get(), llvm::Intrinsic::x86_sse_sfence); + builder->CreateCall(sfence); +} + // i32(i16_a)*i32(i16_b) +/- i32(i16_c)*i32(i16_d) can be done by // interleaving a, c, and b, d, and then using dot_product. bool should_use_dot_product(const Expr &a, const Expr &b, vector &result) { diff --git a/src/CompilerLogger.cpp b/src/CompilerLogger.cpp index 58fe2f1760da..4698fbdc19f5 100644 --- a/src/CompilerLogger.cpp +++ b/src/CompilerLogger.cpp @@ -49,7 +49,7 @@ class ObfuscateNames : public IRMutator { Expr index = mutate(op->index); Expr predicate = mutate(op->predicate); return Load::make(op->type, name, index, op->image, op->param, - predicate, op->alignment); + predicate, op->alignment, op->is_streaming); } Expr visit(const Variable *op) override { diff --git a/src/Deinterleave.cpp b/src/Deinterleave.cpp index e7ebb55bab2e..3b31675e462e 100644 --- a/src/Deinterleave.cpp +++ b/src/Deinterleave.cpp @@ -246,7 +246,7 @@ class Deinterleaver : public IRGraphMutator { if (starting_lane != 0) { align = ModulusRemainder(); } - return Load::make(t, op->name, mutate(op->index), op->image, op->param, mutate(op->predicate), align); + return Load::make(t, op->name, mutate(op->index), op->image, op->param, mutate(op->predicate), align, op->is_streaming); } } @@ -587,20 +587,20 @@ class Interleaver : public IRMutator { // If we want to deinterleave both the index and predicate // (or the predicate is one), then deinterleave the // resulting load. - expr = Load::make(op->type, op->name, idx, op->image, op->param, predicate, op->alignment); + expr = Load::make(op->type, op->name, idx, op->image, op->param, predicate, op->alignment, op->is_streaming); expr = deinterleave_expr(expr); } else if (should_deinterleave_idx) { // If we only want to deinterleave the index and not the // predicate, deinterleave the index prior to the load. idx = deinterleave_expr(idx); - expr = Load::make(op->type, op->name, idx, op->image, op->param, predicate, op->alignment); + expr = Load::make(op->type, op->name, idx, op->image, op->param, predicate, op->alignment, op->is_streaming); } else if (should_deinterleave_predicate) { // Similarly, deinterleave the predicate prior to the load // if we don't want to deinterleave the index. predicate = deinterleave_expr(predicate); - expr = Load::make(op->type, op->name, idx, op->image, op->param, predicate, op->alignment); + expr = Load::make(op->type, op->name, idx, op->image, op->param, predicate, op->alignment, op->is_streaming); } else if (!idx.same_as(op->index) || !predicate.same_as(op->index)) { - expr = Load::make(op->type, op->name, idx, op->image, op->param, predicate, op->alignment); + expr = Load::make(op->type, op->name, idx, op->image, op->param, predicate, op->alignment, op->is_streaming); } else { expr = op; } @@ -644,7 +644,7 @@ class Interleaver : public IRMutator { predicate = deinterleave_expr(predicate); } - Stmt stmt = Store::make(op->name, value, idx, op->param, predicate, op->alignment); + Stmt stmt = Store::make(op->name, value, idx, op->param, predicate, op->alignment, op->is_streaming); should_deinterleave = old_should_deinterleave; num_lanes = old_num_lanes; @@ -779,7 +779,7 @@ class Interleaver : public IRMutator { Expr index = Ramp::make(base, make_one(base.type()), t.lanes()); Expr value = Shuffle::make_interleave(args); Expr predicate = Shuffle::make_interleave(predicates); - Stmt new_store = Store::make(store->name, value, index, store->param, predicate, ModulusRemainder()); + Stmt new_store = Store::make(store->name, value, index, store->param, predicate, ModulusRemainder(), store->is_streaming); // Rewrap the let statements we pulled off. while (!let_stmts.empty()) { diff --git a/src/Deserialization.cpp b/src/Deserialization.cpp index 8a27cc8c1386..7749fda9a2c8 100644 --- a/src/Deserialization.cpp +++ b/src/Deserialization.cpp @@ -564,7 +564,7 @@ Stmt Deserializer::deserialize_stmt(Serialize::Stmt type_code, const void *stmt) user_error << "unknown parameter used in pipeline '" << param_name << "'\n"; } const auto alignment = deserialize_modulus_remainder(store_stmt->alignment()); - return Store::make(name, value, index, param, predicate, alignment); + return Store::make(name, value, index, param, predicate, alignment, store_stmt->is_streaming()); } case Serialize::Stmt::Provide: { const auto *provide_stmt = (const Serialize::Provide *)stmt; @@ -832,7 +832,7 @@ Expr Deserializer::deserialize_expr(Serialize::Expr type_code, const void *expr) } const auto alignment = deserialize_modulus_remainder(load_expr->alignment()); const auto type = deserialize_type(load_expr->type()); - return Load::make(type, name, index, image, param, predicate, alignment); + return Load::make(type, name, index, image, param, predicate, alignment, load_expr->is_streaming()); } case Serialize::Expr::Ramp: { const auto *ramp_expr = (const Serialize::Ramp *)expr; @@ -1199,8 +1199,16 @@ StageSchedule Deserializer::deserialize_stage_schedule(const Serialize::StageSch const bool allow_race_conditions = stage_schedule->allow_race_conditions(); const bool atomic = stage_schedule->atomic(); const bool override_atomic_associativity_test = stage_schedule->override_atomic_associativity_test(); + const bool stream_stores = stage_schedule->stream_stores(); + std::optional> stream_loads_names; + if (!stage_schedule->stream_loads_all()) { + stream_loads_names = + deserialize_vector(stage_schedule->stream_loads_names(), + &Deserializer::deserialize_string); + } return StageSchedule(rvars, splits, dims, prefetches, fuse_level, fused_pairs, touched, - allow_race_conditions, atomic, override_atomic_associativity_test); + allow_race_conditions, atomic, override_atomic_associativity_test, + stream_stores, stream_loads_names); } BufferConstraint Deserializer::deserialize_buffer_constraint(const Serialize::BufferConstraint *buffer_constraint) { diff --git a/src/EliminateBoolVectors.cpp b/src/EliminateBoolVectors.cpp index b68e1efa2d83..6ee0498a1fd4 100644 --- a/src/EliminateBoolVectors.cpp +++ b/src/EliminateBoolVectors.cpp @@ -157,7 +157,7 @@ class EliminateBoolVectors : public IRMutator { if (predicate.same_as(op->predicate) && value.same_as(op->value) && index.same_as(op->index)) { return op; } else { - return Store::make(op->name, value, index, op->param, predicate, op->alignment); + return Store::make(op->name, value, index, op->param, predicate, op->alignment, op->is_streaming); } } @@ -172,7 +172,7 @@ class EliminateBoolVectors : public IRMutator { } else { return Load::make(op->type, op->name, std::move(index), op->image, op->param, std::move(predicate), - op->alignment); + op->alignment, op->is_streaming); } } diff --git a/src/Expr.h b/src/Expr.h index ca31830e89e8..0f71b860f4b9 100644 --- a/src/Expr.h +++ b/src/Expr.h @@ -76,6 +76,8 @@ class IRVisitor; X(Evaluate) \ X(Prefetch) \ X(Atomic) \ + X(StreamingStore) \ + X(StreamingLoads) \ X(HoistedStorage) #define HALIDE_FOR_EACH_IR_NODE(X) \ diff --git a/src/FindIntrinsics.cpp b/src/FindIntrinsics.cpp index 9d8c2f028d18..2cad41219b70 100644 --- a/src/FindIntrinsics.cpp +++ b/src/FindIntrinsics.cpp @@ -1102,7 +1102,7 @@ class FindIntrinsics : public IRMutator { } else { return Load::make(op->type, op->name, std::move(index), op->image, op->param, std::move(predicate), - op->alignment); + op->alignment, op->is_streaming); } } @@ -1114,7 +1114,7 @@ class FindIntrinsics : public IRMutator { if (predicate.same_as(op->predicate) && value.same_as(op->value) && index.same_as(op->index)) { return op; } else { - return Store::make(op->name, std::move(value), std::move(index), op->param, std::move(predicate), op->alignment); + return Store::make(op->name, std::move(value), std::move(index), op->param, std::move(predicate), op->alignment, op->is_streaming); } } }; diff --git a/src/FlattenNestedRamps.cpp b/src/FlattenNestedRamps.cpp index 7ee062ac5850..6027210463b7 100644 --- a/src/FlattenNestedRamps.cpp +++ b/src/FlattenNestedRamps.cpp @@ -132,7 +132,7 @@ class FlattenRamps : public IRMutator { // It's a load of a broadcast. Convert it to a broadcast of a load Expr load = Load::make(op->type.element_of(), op->name, min_lane, op->image, op->param, - const_true(), ModulusRemainder{}); + const_true(), ModulusRemainder{}, op->is_streaming); return Broadcast::make(load, lanes); } else { // Turn it into a dense load and a shuffle @@ -141,7 +141,7 @@ class FlattenRamps : public IRMutator { Expr dense_load = Load::make(op->type.with_lanes(extent), op->name, dense_index, op->image, op->param, - const_true(extent), ModulusRemainder{}); + const_true(extent), ModulusRemainder{}, op->is_streaming); return Shuffle::make({dense_load}, const_indices); } } @@ -172,7 +172,7 @@ class FlattenRamps : public IRMutator { Expr p = slice_per_inner_ramp(predicate, (int)n, inner_lanes); ModulusRemainder align = (n == 0) ? op->alignment : ModulusRemainder{}; loads.push_back(Load::make(elem_type, op->name, sub_indices[n], - op->image, op->param, p, align)); + op->image, op->param, p, align, op->is_streaming)); } return Shuffle::make_concat(loads); } @@ -221,7 +221,7 @@ class FlattenRamps : public IRMutator { Expr v = slice_per_inner_ramp(value_ref, (int)n, inner_lanes); ModulusRemainder align = (n == 0) ? op->alignment : ModulusRemainder{}; stores.push_back(Store::make(op->name, v, sub_indices[n], - op->param, p, align)); + op->param, p, align, op->is_streaming)); } Stmt result = Block::make(stores); if (!predicate_name.empty()) { diff --git a/src/Func.cpp b/src/Func.cpp index 81c6c7cfc7fa..d5ab769f79f4 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -1654,6 +1654,51 @@ Stage &Stage::atomic(bool override_associativity_test) { return *this; } +Stage &Stage::stream_stores() { + for (const Dim &d : definition.schedule().dims()) { + if (d.is_rvar() && !d.is_pure()) { + user_error << "Can't stream stores for " << name() + << " because it has a reduction variable that Halide " + "cannot prove is safe to parallelize. A self-load in " + "this Stage could observe a value it streamed earlier " + "in the same Stage, before the fence that makes " + "streamed stores visible.\n"; + } + } + definition.schedule().touched() = true; + definition.schedule().stream_stores() = true; + return *this; +} + +Stage &Stage::stream_loads() { + definition.schedule().touched() = true; + definition.schedule().stream_loads_names() = std::nullopt; + return *this; +} + +Stage &Stage::stream_loads(const std::vector &funcs) { + auto &names = definition.schedule().stream_loads_names(); + if (!names) { + names = std::vector{}; + } + for (const Func &f : funcs) { + std::string target_name = f.name(); + if (const Call *call = f.function().is_wrapper(); call && call->param.defined()) { + // A pure wrapper around an ImageParam/Buffer Parameter (e.g. as + // returned by ImageParam's implicit conversion to Func): the + // wrapper itself is inlined away before storage flattening + // ever runs, so match the underlying Parameter directly. + target_name = call->param.name(); + } + user_assert(target_name != function.name()) + << "Can't stream loads of \"" << target_name << "\" in " << name() + << " because a Stage cannot stream its own self-loads.\n"; + names->push_back(target_name); + } + definition.schedule().touched() = true; + return *this; +} + Stage &Stage::serial(const VarOrRVar &var) { set_dim_type(var, ForType::Serial); return *this; @@ -2451,6 +2496,24 @@ Func &Func::store_in(MemoryType t) { return *this; } +Func &Func::stream_loads() { + invalidate_cache(); + Stage(func, func.definition(), 0).stream_loads(); + return *this; +} + +Func &Func::stream_loads(const std::vector &funcs) { + invalidate_cache(); + Stage(func, func.definition(), 0).stream_loads(funcs); + return *this; +} + +Func &Func::stream_stores() { + invalidate_cache(); + Stage(func, func.definition(), 0).stream_stores(); + return *this; +} + Func &Func::async() { invalidate_cache(); func.schedule().async() = true; diff --git a/src/Func.h b/src/Func.h index 0bfb591871c7..000b44561efa 100644 --- a/src/Func.h +++ b/src/Func.h @@ -461,6 +461,27 @@ class Stage { Stage &allow_race_conditions(); Stage &atomic(bool override_associativity_test = false); + /** Use non-temporal (streaming) stores for writes done by this Stage. On + * targets that require it, Halide emits a fence immediately after this + * Stage's production to ensure the streamed values are visible to + * whatever reads them next. Only legal on a Stage all of whose RVars (if + * any) are pure, i.e. already proven safe to parallelize: a Stage with a + * genuine loop-carried self-dependency could otherwise observe data it + * streamed earlier in the same Stage, before the fence. It is a user + * error to call this on a Stage that doesn't meet this condition. */ + Stage &stream_stores(); + + /** Use non-temporal (streaming) loads for every direct read this Stage + * makes of another Func or external buffer (e.g. an ImageParam). This + * is a hint to keep data that is only read once from displacing reusable + * data from the cache. */ + Stage &stream_loads(); + + /** Use non-temporal (streaming) loads for this Stage's direct reads of + * the named Funcs. It is a user error to name this Stage's own Func + * (a self-load can't be streamed). */ + Stage &stream_loads(const std::vector &funcs); + Stage &hexagon(const VarOrRVar &x = Var::outermost()); Stage &prefetch(const Func &f, const VarOrRVar &at, const VarOrRVar &from, Expr offset = 1, @@ -2602,6 +2623,25 @@ class Func { * on MemoryType for more detail. */ Func &store_in(MemoryType memory_type); + /** Use non-temporal (streaming) loads for every direct read this Func's + * pure (initial) definition makes of another Func. Equivalent to calling + * stream_loads() on Stage 0; see \ref Stage::stream_loads. To stream the + * loads of an update definition, call stream_loads() on the Stage returned + * by \ref Func::update instead. */ + Func &stream_loads(); + + /** Use non-temporal (streaming) loads for this Func's pure (initial) + * definition's direct reads of the named Funcs. Equivalent to calling + * stream_loads(funcs) on Stage 0; see \ref Stage::stream_loads. */ + Func &stream_loads(const std::vector &funcs); + + /** Use non-temporal (streaming) stores for writes done by this Func's + * pure (initial) definition. Equivalent to calling stream_stores() on + * Stage 0; see \ref Stage::stream_stores. To stream the stores of an + * update definition, call stream_stores() on the Stage returned by + * \ref Func::update instead. */ + Func &stream_stores(); + /** Trace all loads from this Func by emitting calls to * halide_trace. If the Func is inlined, this has no * effect. */ diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index ec85e5a383f9..84862ceff6a3 100644 --- a/src/FuseGPUThreadLoops.cpp +++ b/src/FuseGPUThreadLoops.cpp @@ -536,7 +536,7 @@ class ExtractSharedAndHeapAllocations : public IRMutator { Expr predicate = mutate(op->predicate); Expr index = mutate_index(alloc, op->index); return Load::make(op->type, alloc->name, - index, op->image, op->param, predicate, op->alignment); + index, op->image, op->param, predicate, op->alignment, op->is_streaming); } else { return IRMutator::visit(op); } @@ -551,7 +551,7 @@ class ExtractSharedAndHeapAllocations : public IRMutator { Expr index = mutate_index(alloc, op->index); Expr value = mutate(op->value); return Store::make(alloc->name, value, index, - op->param, predicate, op->alignment); + op->param, predicate, op->alignment, op->is_streaming); } else { return IRMutator::visit(op); } @@ -886,7 +886,7 @@ class ExtractSharedAndHeapAllocations : public IRMutator { if (op->name == alloc_name) { return Load::make(op->type, cluster_name, mutate(op->index) + offset, op->image, op->param, mutate(op->predicate), - op->alignment); + op->alignment, op->is_streaming); } else { return IRMutator::visit(op); } @@ -895,7 +895,7 @@ class ExtractSharedAndHeapAllocations : public IRMutator { Stmt visit(const Store *op) override { if (op->name == alloc_name) { return Store::make(cluster_name, mutate(op->value), mutate(op->index) + offset, - op->param, mutate(op->predicate), op->alignment); + op->param, mutate(op->predicate), op->alignment, op->is_streaming); } else { return IRMutator::visit(op); } @@ -1150,7 +1150,7 @@ class ExtractRegisterAllocations : public IRMutator { } return Load::make(op->type, *new_name, mutate(op->index), op->image, op->param, mutate(op->predicate), - op->alignment); + op->alignment, op->is_streaming); } Stmt visit(const Store *op) override { @@ -1159,7 +1159,7 @@ class ExtractRegisterAllocations : public IRMutator { new_name = &(op->name); } return Store::make(*new_name, mutate(op->value), mutate(op->index), - op->param, mutate(op->predicate), op->alignment); + op->param, mutate(op->predicate), op->alignment, op->is_streaming); } template diff --git a/src/FuzzFloatStores.cpp b/src/FuzzFloatStores.cpp index 511141889680..74c2c6bde918 100644 --- a/src/FuzzFloatStores.cpp +++ b/src/FuzzFloatStores.cpp @@ -18,7 +18,7 @@ class FuzzFloatStores : public IRMutator { value = reinterpret(mask.type(), value); value = value & ~mask; value = reinterpret(t, value); - return Store::make(op->name, value, op->index, op->param, op->predicate, op->alignment); + return Store::make(op->name, value, op->index, op->param, op->predicate, op->alignment, op->is_streaming); } else { return IRMutator::visit(op); } diff --git a/src/HexagonOffload.cpp b/src/HexagonOffload.cpp index 5cafec409579..c37b99ef9cbd 100644 --- a/src/HexagonOffload.cpp +++ b/src/HexagonOffload.cpp @@ -670,7 +670,7 @@ class ReplaceParams : public IRMutator { auto i = replacements.find(op->name); if (i != replacements.end()) { return Load::make(op->type, op->name, mutate(op->index), op->image, - i->second, mutate(op->predicate), op->alignment); + i->second, mutate(op->predicate), op->alignment, op->is_streaming); } else { return IRMutator::visit(op); } @@ -680,7 +680,7 @@ class ReplaceParams : public IRMutator { auto i = replacements.find(op->name); if (i != replacements.end()) { return Store::make(op->name, mutate(op->value), mutate(op->index), - i->second, mutate(op->predicate), op->alignment); + i->second, mutate(op->predicate), op->alignment, op->is_streaming); } else { return IRMutator::visit(op); } diff --git a/src/HexagonOptimize.cpp b/src/HexagonOptimize.cpp index a24e00ec2dc7..89e98be4fe2f 100644 --- a/src/HexagonOptimize.cpp +++ b/src/HexagonOptimize.cpp @@ -1932,7 +1932,7 @@ class EliminateInterleaves : public IRMutator { if (predicate.same_as(op->predicate) && value.same_as(op->value) && index.same_as(op->index)) { return op; } else { - return Store::make(op->name, value, index, op->param, predicate, op->alignment); + return Store::make(op->name, value, index, op->param, predicate, op->alignment, op->is_streaming); } } diff --git a/src/IR.cpp b/src/IR.cpp index ae0e38a36251..d9b86a132d08 100644 --- a/src/IR.cpp +++ b/src/IR.cpp @@ -247,7 +247,7 @@ Expr Select::make(Expr condition, Expr true_value, Expr false_value) { return node; } -Expr Load::make(Type type, const std::string &name, Expr index, Buffer<> image, Parameter param, Expr predicate, ModulusRemainder alignment) { +Expr Load::make(Type type, const std::string &name, Expr index, Buffer<> image, Parameter param, Expr predicate, ModulusRemainder alignment, bool is_streaming) { internal_assert(predicate.defined()) << "Load with undefined predicate\n"; internal_assert(index.defined()) << "Load of undefined\n"; internal_assert(type.lanes() == index.type().lanes()) << "Vector lanes of Load must match vector lanes of index\n"; @@ -262,6 +262,7 @@ Expr Load::make(Type type, const std::string &name, Expr index, Buffer<> image, node->image = std::move(image); node->param = std::move(param); node->alignment = alignment; + node->is_streaming = is_streaming; return node; } @@ -374,7 +375,7 @@ Stmt Acquire::make(Expr semaphore, Expr count, Stmt body) { return node; } -Stmt Store::make(const std::string &name, Expr value, Expr index, Parameter param, Expr predicate, ModulusRemainder alignment) { +Stmt Store::make(const std::string &name, Expr value, Expr index, Parameter param, Expr predicate, ModulusRemainder alignment, bool is_streaming) { internal_assert(predicate.defined()) << "Store with undefined predicate\n"; internal_assert(value.defined()) << "Store of undefined\n"; internal_assert(index.defined()) << "Store of undefined\n"; @@ -389,6 +390,7 @@ Stmt Store::make(const std::string &name, Expr value, Expr index, Parameter para node->index = std::move(index); node->param = std::move(param); node->alignment = alignment; + node->is_streaming = is_streaming; return node; } @@ -678,6 +680,7 @@ constexpr const char *intrinsic_op_names[] = { "skip_stages_marker", "sliding_window_marker", "sorted_avg", + "stream_store_fence", "strict_add", "strict_cast", "strict_div", @@ -972,6 +975,24 @@ Stmt Atomic::make(const std::string &producer_name, return node; } +Stmt StreamingStore::make(const std::string &producer_name, + Stmt body) { + internal_assert(body.defined()) << "StreamingStore must have a body statement.\n"; + StreamingStore *node = new StreamingStore; + node->producer_name = producer_name; + node->body = std::move(body); + return node; +} + +Stmt StreamingLoads::make(std::optional> names, + Stmt body) { + internal_assert(body.defined()) << "StreamingLoads must have a body statement.\n"; + StreamingLoads *node = new StreamingLoads; + node->names = std::move(names); + node->body = std::move(body); + return node; +} + Stmt HoistedStorage::make(const std::string &name, Stmt body) { internal_assert(body.defined()) << "HoistedStorage must have a body statement.\n"; @@ -1260,6 +1281,14 @@ void StmtNode::accept(IRVisitor *v) const { v->visit((const Atomic *)this); } template<> +void StmtNode::accept(IRVisitor *v) const { + v->visit((const StreamingStore *)this); +} +template<> +void StmtNode::accept(IRVisitor *v) const { + v->visit((const StreamingLoads *)this); +} +template<> void StmtNode::accept(IRVisitor *v) const { v->visit((const HoistedStorage *)this); } @@ -1454,6 +1483,14 @@ Stmt StmtNode::mutate_stmt(IRMutator *v) const { return v->visit((const Atomic *)this); } template<> +Stmt StmtNode::mutate_stmt(IRMutator *v) const { + return v->visit((const StreamingStore *)this); +} +template<> +Stmt StmtNode::mutate_stmt(IRMutator *v) const { + return v->visit((const StreamingLoads *)this); +} +template<> Stmt StmtNode::mutate_stmt(IRMutator *v) const { return v->visit((const HoistedStorage *)this); } diff --git a/src/IR.h b/src/IR.h index 000d8a994851..7f8bcd9e9644 100644 --- a/src/IR.h +++ b/src/IR.h @@ -5,6 +5,7 @@ * Subtypes for Halide expressions (\ref Halide::Expr) and statements (\ref Halide::Internal::Stmt) */ +#include #include #include @@ -309,11 +310,15 @@ struct Load : public ExprNode { // the alignment of the first lane. ModulusRemainder alignment; + // Whether this access should bypass the cache when supported by the target. + bool is_streaming; + static Expr make(Type type, const std::string &name, Expr index, Buffer<> image, Parameter param, Expr predicate, - ModulusRemainder alignment); + ModulusRemainder alignment, + bool is_streaming = false); static const IRNodeType _node_type = IRNodeType::Load; }; @@ -427,8 +432,12 @@ struct Store : public StmtNode { // the alignment of the first lane. ModulusRemainder alignment; + // Whether this access should bypass the cache when supported by the target. + bool is_streaming; + static Stmt make(const std::string &name, Expr value, Expr index, - Parameter param, Expr predicate, ModulusRemainder alignment); + Parameter param, Expr predicate, ModulusRemainder alignment, + bool is_streaming = false); static const IRNodeType _node_type = IRNodeType::Store; }; @@ -731,6 +740,9 @@ struct Call : public ExprNode { sliding_window_marker, // Compute (arg[0] + arg[1]) / 2, assuming arg[0] < arg[1]. sorted_avg, + // Emits a target-specific memory fence after a Stage that + // contains non-temporal (streaming) stores. + stream_store_fence, // strict floating point ops. These are floating point ops that we would // like to optimize around (or let llvm optimize around) by treating // them as reals and ignoring the existence of nan and inf. Using these @@ -1150,6 +1162,41 @@ struct Atomic : public StmtNode { static const IRNodeType _node_type = IRNodeType::Atomic; }; +/** Marks the store(s) produced by the wrapped body as requesting + * non-temporal (streaming) stores, as scheduled via Stage::stream_stores. + * Created directly around a Provide node (or an Atomic node wrapping one) + * during scheduling, and consumed by storage flattening, which uses it to + * set Store::is_streaming on the resulting Store node(s) and then discards + * it -- it should not survive past that point. */ +struct StreamingStore : public StmtNode { + std::string producer_name; + Stmt body; + + static Stmt make(const std::string &producer_name, + Stmt body); + + static const IRNodeType _node_type = IRNodeType::StreamingStore; +}; + +/** Marks the Halide-Func loads made while evaluating the wrapped body as + * requesting non-temporal (streaming) loads, as scheduled via + * Stage::stream_loads. If `names` is nullopt, every direct load of another + * Func is streamed (except a self-load, which is never streamed); + * otherwise only loads of Funcs named in `*names` are streamed. Created + * directly around a Stage's body during scheduling, and consumed by + * storage flattening, which uses it to set Load::is_streaming on the + * resulting Load node(s) and then discards it -- it should not survive + * past that point. */ +struct StreamingLoads : public StmtNode { + std::optional> names; + Stmt body; + + static Stmt make(std::optional> names, + Stmt body); + + static const IRNodeType _node_type = IRNodeType::StreamingLoads; +}; + /** Horizontally reduce a vector to a scalar or narrower vector using * the given commutative and associative binary operator. The reduction * factor is dictated by the number of lanes in the input and output diff --git a/src/IREquality.cpp b/src/IREquality.cpp index 1e7ae422c549..321422d4d2db 100644 --- a/src/IREquality.cpp +++ b/src/IREquality.cpp @@ -380,6 +380,7 @@ struct Comparer { case IRNodeType::Load: cmp(&Load::name); cmp(&Load::alignment); + cmp(&Load::is_streaming); cmp(&Load::index); cmp(&Load::predicate); break; @@ -437,6 +438,7 @@ struct Comparer { case IRNodeType::Store: cmp(&Store::name); cmp(&Store::alignment); + cmp(&Store::is_streaming); cmp(&Store::predicate); cmp(&Store::value); cmp(&Store::index); @@ -494,6 +496,14 @@ struct Comparer { cmp(&Atomic::mutex_name); cmp(&Atomic::body); break; + case IRNodeType::StreamingStore: + cmp(&StreamingStore::producer_name); + cmp(&StreamingStore::body); + break; + case IRNodeType::StreamingLoads: + cmp(&StreamingLoads::names); + cmp(&StreamingLoads::body); + break; case IRNodeType::HoistedStorage: cmp(&HoistedStorage::name); cmp(&HoistedStorage::body); diff --git a/src/IRMutator.cpp b/src/IRMutator.cpp index c5ae276bda62..9850a84bf5a7 100644 --- a/src/IRMutator.cpp +++ b/src/IRMutator.cpp @@ -128,7 +128,7 @@ Expr IRMutator::visit(const Load *op) { } return Load::make(op->type, op->name, std::move(index), op->image, op->param, std::move(predicate), - op->alignment); + op->alignment, op->is_streaming); } Expr IRMutator::visit(const Ramp *op) { @@ -216,7 +216,7 @@ Stmt IRMutator::visit(const Store *op) { if (predicate.same_as(op->predicate) && value.same_as(op->value) && index.same_as(op->index)) { return op; } - return Store::make(op->name, std::move(value), std::move(index), op->param, std::move(predicate), op->alignment); + return Store::make(op->name, std::move(value), std::move(index), op->param, std::move(predicate), op->alignment, op->is_streaming); } Stmt IRMutator::visit(const Provide *op) { @@ -366,6 +366,24 @@ Stmt IRMutator::visit(const Atomic *op) { } } +Stmt IRMutator::visit(const StreamingStore *op) { + Stmt body = mutate(op->body); + if (body.same_as(op->body)) { + return op; + } else { + return StreamingStore::make(op->producer_name, std::move(body)); + } +} + +Stmt IRMutator::visit(const StreamingLoads *op) { + Stmt body = mutate(op->body); + if (body.same_as(op->body)) { + return op; + } else { + return StreamingLoads::make(op->names, std::move(body)); + } +} + Stmt IRMutator::visit(const HoistedStorage *op) { Stmt body = mutate(op->body); if (body.same_as(op->body)) { diff --git a/src/IRPrinter.cpp b/src/IRPrinter.cpp index 331e5de658b2..68102bb21537 100644 --- a/src/IRPrinter.cpp +++ b/src/IRPrinter.cpp @@ -339,6 +339,8 @@ std::ostream &operator<<(std::ostream &stream, IRNodeType type) { CASE(Evaluate) CASE(Prefetch) CASE(Atomic) + CASE(StreamingStore) + CASE(StreamingLoads) CASE(HoistedStorage) } #undef CASE @@ -996,6 +998,9 @@ void IRPrinter::visit(const Load *op) { if (has_pred) { open(); } + if (op->is_streaming) { + stream << kw("streaming "); + } if (!known_type.contains(op->name)) { stream << typep(op->type); } @@ -1206,6 +1211,9 @@ void IRPrinter::visit(const Store *op) { indent++; stream << get_indent(); } + if (op->is_streaming) { + stream << kw("streaming "); + } stream << buf(op->name) << paren("["); print_no_parens(op->index); if (show_alignment) { @@ -1505,6 +1513,29 @@ void IRPrinter::visit(const Atomic *op) { print_braced_stmt(op->body); } +void IRPrinter::visit(const StreamingStore *op) { + stream << get_indent(); + stream << kw("streaming_store (") << op->producer_name << kw(") "); + print_braced_stmt(op->body); +} + +void IRPrinter::visit(const StreamingLoads *op) { + stream << get_indent(); + if (!op->names) { + stream << kw("streaming_loads (") << kw("all") << kw(") "); + } else { + stream << kw("streaming_loads ("); + for (size_t i = 0; i < op->names->size(); i++) { + if (i > 0) { + stream << kw(", "); + } + stream << (*op->names)[i]; + } + stream << kw(") "); + } + print_braced_stmt(op->body); +} + void IRPrinter::visit(const HoistedStorage *op) { if (op->name.empty()) { stream << get_indent() << kw("hoisted_storage "); diff --git a/src/IRPrinter.h b/src/IRPrinter.h index 12ffacb4fa88..571edbc24a1c 100644 --- a/src/IRPrinter.h +++ b/src/IRPrinter.h @@ -288,6 +288,8 @@ class IRPrinter : public IRVisitor { void visit(const VectorReduce *) override; void visit(const Prefetch *) override; void visit(const Atomic *) override; + void visit(const StreamingStore *) override; + void visit(const StreamingLoads *) override; void visit(const HoistedStorage *) override; }; diff --git a/src/IRVisitor.cpp b/src/IRVisitor.cpp index 25fd7e608f27..54c3ceb31fe0 100644 --- a/src/IRVisitor.cpp +++ b/src/IRVisitor.cpp @@ -265,6 +265,14 @@ void IRVisitor::visit(const Atomic *op) { op->body.accept(this); } +void IRVisitor::visit(const StreamingStore *op) { + op->body.accept(this); +} + +void IRVisitor::visit(const StreamingLoads *op) { + op->body.accept(this); +} + void IRVisitor::visit(const HoistedStorage *op) { op->body.accept(this); } @@ -536,6 +544,14 @@ void IRGraphVisitor::visit(const Atomic *op) { include(op->body); } +void IRGraphVisitor::visit(const StreamingStore *op) { + include(op->body); +} + +void IRGraphVisitor::visit(const StreamingLoads *op) { + include(op->body); +} + void IRGraphVisitor::visit(const HoistedStorage *op) { include(op->body); } diff --git a/src/Inline.cpp b/src/Inline.cpp index ce829f1c7326..70e936f07898 100644 --- a/src/Inline.cpp +++ b/src/Inline.cpp @@ -95,6 +95,18 @@ void validate_schedule_inlined_function(Function f) { << b.remainder << "] because the function is scheduled inline.\n"; } } + + if (stage_s.stream_stores()) { + user_warning << "It is meaningless to stream the stores of function " + << f.name() << " because it is scheduled inline: there is " + "no separate Store left to mark non-temporal.\n"; + } + + if (const auto &names = stage_s.stream_loads_names(); !names || !names->empty()) { + user_warning << "It is meaningless to stream the loads made by function " + << f.name() << " because it is scheduled inline: there is " + "no separate Load left to mark non-temporal.\n"; + } } class Inliner : public IRMutator { diff --git a/src/LLVM_Headers.h b/src/LLVM_Headers.h index 6e78b8289116..99cccecd4a99 100644 --- a/src/LLVM_Headers.h +++ b/src/LLVM_Headers.h @@ -53,6 +53,9 @@ #ifdef WITH_HEXAGON #include #endif +#ifdef WITH_X86 +#include +#endif #include #include #include diff --git a/src/LowerWarpShuffles.cpp b/src/LowerWarpShuffles.cpp index 9f244aad2ce0..0eea1fe38318 100644 --- a/src/LowerWarpShuffles.cpp +++ b/src/LowerWarpShuffles.cpp @@ -512,7 +512,7 @@ class LowerWarpShuffles : public IRMutator { // them. Reassembling the result into a flat address gives // the expression below. Expr in_warp_idx = simplify((idx / (warp_size * stride)) * stride + reduce_expr(idx, stride, bounds), bounds); - return Store::make(op->name, value, in_warp_idx, op->param, op->predicate, ModulusRemainder()); + return Store::make(op->name, value, in_warp_idx, op->param, op->predicate, ModulusRemainder(), op->is_streaming); } else { return IRMutator::visit(op); } diff --git a/src/OptimizeShuffles.cpp b/src/OptimizeShuffles.cpp index 83672fa59395..f651bad87d65 100644 --- a/src/OptimizeShuffles.cpp +++ b/src/OptimizeShuffles.cpp @@ -117,7 +117,7 @@ class OptimizeShuffles : public IRMutator { allocations_to_pad.insert(op->name); Expr lut = Load::make(op->type.with_lanes(const_extent), op->name, Ramp::make(base, 1, const_extent), - op->image, op->param, const_true(const_extent), alignment); + op->image, op->param, const_true(const_extent), alignment, op->is_streaming); // We know the size of the LUT is not more than 256, so we // can safely cast the index to 8 bit, which @@ -130,7 +130,7 @@ class OptimizeShuffles : public IRMutator { } } if (!index.same_as(op->index)) { - return Load::make(op->type, op->name, index, op->image, op->param, op->predicate, op->alignment); + return Load::make(op->type, op->name, index, op->image, op->param, op->predicate, op->alignment, op->is_streaming); } else { return op; } diff --git a/src/PartitionLoops.cpp b/src/PartitionLoops.cpp index e31b8b9387b2..34655df2300d 100644 --- a/src/PartitionLoops.cpp +++ b/src/PartitionLoops.cpp @@ -70,7 +70,7 @@ class MarkClampedRampsAsLikely : public IRMutator { if (predicate.same_as(op->predicate) && index.same_as(op->index) && value.same_as(op->value)) { return op; } else { - return Store::make(op->name, value, index, op->param, predicate, op->alignment); + return Store::make(op->name, value, index, op->param, predicate, op->alignment, op->is_streaming); } } diff --git a/src/RemoveUndef.cpp b/src/RemoveUndef.cpp index 512c984427a7..bcb0922eea6f 100644 --- a/src/RemoveUndef.cpp +++ b/src/RemoveUndef.cpp @@ -211,7 +211,7 @@ class RemoveUndef : public IRMutator { if (pred.same_as(op->predicate) && index.same_as(op->index)) { return op; } else { - return Load::make(op->type, op->name, index, op->image, op->param, pred, op->alignment); + return Load::make(op->type, op->name, index, op->image, op->param, pred, op->alignment, op->is_streaming); } } @@ -388,7 +388,7 @@ class RemoveUndef : public IRMutator { if (predicate.defined()) { // This becomes a conditional store - Stmt stmt = IfThenElse::make(predicate, Store::make(op->name, value, index, op->param, pred, op->alignment)); + Stmt stmt = IfThenElse::make(predicate, Store::make(op->name, value, index, op->param, pred, op->alignment, op->is_streaming)); predicate = Expr(); return stmt; } else if (pred.same_as(op->predicate) && @@ -396,7 +396,7 @@ class RemoveUndef : public IRMutator { index.same_as(op->index)) { return op; } else { - return Store::make(op->name, value, index, op->param, pred, op->alignment); + return Store::make(op->name, value, index, op->param, pred, op->alignment, op->is_streaming); } } diff --git a/src/Schedule.cpp b/src/Schedule.cpp index a2583d0fb732..e3e5576e051d 100644 --- a/src/Schedule.cpp +++ b/src/Schedule.cpp @@ -307,6 +307,10 @@ struct StageScheduleContents { bool allow_race_conditions = false; bool atomic = false; bool override_atomic_associativity_test = false; + bool stream_stores = false; + // nullopt = stream all direct loads (except self); present = stream + // only the named Funcs (empty = the default, i.e. none). + std::optional> stream_loads_names = std::vector{}; StageScheduleContents() : fuse_level(FuseLoopLevel()) { @@ -532,7 +536,8 @@ StageSchedule::StageSchedule() StageSchedule::StageSchedule(const std::vector &rvars, const std::vector &splits, const std::vector &dims, const std::vector &prefetches, const FuseLoopLevel &fuse_level, const std::vector &fused_pairs, - bool touched, bool allow_race_conditions, bool atomic, bool override_atomic_associativity_test) + bool touched, bool allow_race_conditions, bool atomic, bool override_atomic_associativity_test, + bool stream_stores, const std::optional> &stream_loads_names) : contents(new StageScheduleContents) { contents->rvars = rvars; contents->splits = splits; @@ -544,6 +549,8 @@ StageSchedule::StageSchedule(const std::vector &rvars, const contents->allow_race_conditions = allow_race_conditions; contents->atomic = atomic; contents->override_atomic_associativity_test = override_atomic_associativity_test; + contents->stream_stores = stream_stores; + contents->stream_loads_names = stream_loads_names; } StageSchedule StageSchedule::get_copy() const { @@ -559,6 +566,8 @@ StageSchedule StageSchedule::get_copy() const { copy.contents->allow_race_conditions = contents->allow_race_conditions; copy.contents->atomic = contents->atomic; copy.contents->override_atomic_associativity_test = contents->override_atomic_associativity_test; + copy.contents->stream_stores = contents->stream_stores; + copy.contents->stream_loads_names = contents->stream_loads_names; return copy; } @@ -642,6 +651,22 @@ bool StageSchedule::override_atomic_associativity_test() const { return contents->override_atomic_associativity_test; } +bool &StageSchedule::stream_stores() { + return contents->stream_stores; +} + +bool StageSchedule::stream_stores() const { + return contents->stream_stores; +} + +std::optional> &StageSchedule::stream_loads_names() { + return contents->stream_loads_names; +} + +const std::optional> &StageSchedule::stream_loads_names() const { + return contents->stream_loads_names; +} + void StageSchedule::accept(IRVisitor *visitor) const { for (const ReductionVariable &r : rvars()) { if (r.min.defined()) { diff --git a/src/Schedule.h b/src/Schedule.h index a72c52bbd59f..8ce415df9497 100644 --- a/src/Schedule.h +++ b/src/Schedule.h @@ -6,6 +6,7 @@ */ #include +#include #include #include #include @@ -696,7 +697,8 @@ class StageSchedule { StageSchedule(const std::vector &rvars, const std::vector &splits, const std::vector &dims, const std::vector &prefetches, const FuseLoopLevel &fuse_level, const std::vector &fused_pairs, - bool touched, bool allow_race_conditions, bool atomic, bool override_atomic_associativity_test); + bool touched, bool allow_race_conditions, bool atomic, bool override_atomic_associativity_test, + bool stream_stores, const std::optional> &stream_loads_names); /** Return a copy of this StageSchedule. */ StageSchedule get_copy() const; @@ -782,6 +784,27 @@ class StageSchedule { bool &override_atomic_associativity_test(); // @} + /** True if stores to this Stage's backing storage should use + * non-temporal (streaming) stores. Only legal on a Stage all of whose + * RVars (if any) are pure, i.e. already proven safe to parallelize; see + * \ref Func::stream_stores. */ + // @{ + bool stream_stores() const; + bool &stream_stores(); + // @} + + /** The stream_loads() request (see \ref Func::stream_loads / \ref + * Stage::stream_loads) for this Stage: nullopt means stream every + * direct load of another Func (except a self-load), set by the + * zero-arg form; otherwise the vector names exactly the Funcs whose + * direct loads by this Stage should be streamed (non-temporal) -- an + * empty vector (the default, if stream_loads() was never called) + * means don't stream any of them. */ + // @{ + const std::optional> &stream_loads_names() const; + std::optional> &stream_loads_names(); + // @} + /** Pass an IRVisitor through to all Exprs referenced in the * Schedule. */ void accept(IRVisitor *) const; diff --git a/src/ScheduleFunctions.cpp b/src/ScheduleFunctions.cpp index a0b4a4c18876..48d35cdea48a 100644 --- a/src/ScheduleFunctions.cpp +++ b/src/ScheduleFunctions.cpp @@ -466,6 +466,25 @@ Stmt build_loop_nest( return stmt; } +// A Stage's Definition may be expanded, by build_provide_loop_nest below, +// into several Provide nodes: one per (possibly recursively nested) +// specialization, plus the base/default Definition. stream_stores() is set +// independently on each of these, so return true if *any* of them (reachable +// via a non-failing specialization) requests it. Used to decide whether a +// fence is needed after this Stage's production; it's safe to over-fence a +// branch that didn't ask for streaming, but not to under-fence one that did. +bool any_specialization_requests_streaming(const Definition &def) { + if (def.schedule().stream_stores()) { + return true; + } + for (const Specialization &s : def.specializations()) { + if (s.failure_message.empty() && any_specialization_requests_streaming(s.definition)) { + return true; + } + } + return false; +} + // Build a loop nest about a provide node using a schedule Stmt build_provide_loop_nest(const map &env, const string &prefix, @@ -514,6 +533,16 @@ Stmt build_provide_loop_nest(const map &env, body = Atomic::make(func.name(), std::string{}, body); } } + if (def.schedule().stream_stores()) { + // Wrapped outside any Atomic node (rather than inside it), so that + // passes between here and storage flattening which pattern-match an + // Atomic node's body as being directly a Provide node (e.g. + // SplitTuples) continue to see that exact shape. + body = StreamingStore::make(func.name(), body); + } + if (const auto &names = def.schedule().stream_loads_names(); !names || !names->empty()) { + body = StreamingLoads::make(names, body); + } // Default schedule/values if there is no specialization Stmt stmt = build_loop_nest(body, prefix, start_fuse, func, def); @@ -1546,6 +1575,15 @@ class InjectFunctionRealization : public IRMutator { add_lets.emplace_back(let->name, let->value); produce = let->body; } + + // Only one branch runs at a time, so a single trailing fence + // after the whole (specialized) production is equivalent to, + // and simpler than, fencing inside each branch individually. + if (any_specialization_requests_streaming(def)) { + Expr fence = Call::make(Int(32), Call::stream_store_fence, {}, Call::Intrinsic); + produce = Block::make(produce, Evaluate::make(fence)); + } + return produce; } @@ -1802,8 +1840,8 @@ class InjectFunctionRealization : public IRMutator { string def_prefix = f.name() + ".s" + std::to_string(func_stage.second) + "."; const auto &def = (func_stage.second == 0) ? f.definition() : f.updates()[func_stage.second - 1]; - const Stmt &produce_def = build_produce_definition(f, def_prefix, def, func_stage.second > 0, - replacements, add_lets, aliases); + Stmt produce_def = build_produce_definition(f, def_prefix, def, func_stage.second > 0, + replacements, add_lets, aliases); producer = inject_stmt(producer, produce_def, def.schedule().fuse_level().level); } diff --git a/src/Serialization.cpp b/src/Serialization.cpp index 33de404edffe..a0de8eae28f9 100644 --- a/src/Serialization.cpp +++ b/src/Serialization.cpp @@ -474,7 +474,8 @@ std::pair> Serializer::serialize_stmt(FlatBufferBu predicate_serialized.first, predicate_serialized.second, value_serialized.first, value_serialized.second, index_serialized.first, index_serialized.second, - param_name_serialized, alignment_serialized) + param_name_serialized, alignment_serialized, + store_stmt->is_streaming) .Union()); } case IRNodeType::Provide: { @@ -841,7 +842,8 @@ std::pair> Serializer::serialize_expr(FlatBufferBu predicate_serialized.first, predicate_serialized.second, index_serialized.first, index_serialized.second, image_name_serialized, param_name_serialized, - alignment_serialized, type_serialized) + alignment_serialized, type_serialized, + load_expr->is_streaming) .Union()); } case IRNodeType::Ramp: { @@ -1124,7 +1126,8 @@ Offset Serializer::serialize_func_schedule(FlatBufferBu builder.CreateVector(bounds_serialized), builder.CreateVector(estimates_serialized), builder.CreateVector(wrappers_serialized), - memory_type, memoized, async, ring_buffer.first, ring_buffer.second, + memory_type, memoized, async, + ring_buffer.first, ring_buffer.second, memoize_eviction_key_serialized.first, memoize_eviction_key_serialized.second); } @@ -1292,6 +1295,16 @@ Offset Serializer::serialize_stage_schedule(FlatBuffer const bool allow_race_conditions = stage_schedule.allow_race_conditions(); const bool atomic = stage_schedule.atomic(); const bool override_atomic_associativity_test = stage_schedule.override_atomic_associativity_test(); + const bool stream_stores = stage_schedule.stream_stores(); + const auto &stream_loads_names_opt = stage_schedule.stream_loads_names(); + const bool stream_loads_all = !stream_loads_names_opt.has_value(); + std::vector> stream_loads_names_serialized; + if (stream_loads_names_opt) { + stream_loads_names_serialized.reserve(stream_loads_names_opt->size()); + for (const auto &name : *stream_loads_names_opt) { + stream_loads_names_serialized.push_back(serialize_string(builder, name)); + } + } return Serialize::CreateStageSchedule(builder, builder.CreateVector(rvars_serialized), builder.CreateVector(splits_serialized), @@ -1300,7 +1313,9 @@ Offset Serializer::serialize_stage_schedule(FlatBuffer fuse_level_serialized, builder.CreateVector(fused_pairs_serialized), touched, allow_race_conditions, atomic, - override_atomic_associativity_test); + override_atomic_associativity_test, + stream_stores, stream_loads_all, + builder.CreateVector(stream_loads_names_serialized)); } Offset Serializer::serialize_buffer_constraint(FlatBufferBuilder &builder, const BufferConstraint &buffer_constraint) { @@ -1338,7 +1353,9 @@ Offset Serializer::serialize_parameter(FlatBufferBuilder & } const auto memory_type_serialized = serialize_memory_type(parameter.memory_type()); return Serialize::CreateParameter(builder, defined, is_buffer, type_serialized, dimensions, name_serialized, host_alignment, - builder.CreateVector(buffer_constraints_serialized), memory_type_serialized); + builder.CreateVector(buffer_constraints_serialized), memory_type_serialized, std::nullopt, + Serialize::Expr::NONE, 0, Serialize::Expr::NONE, 0, + Serialize::Expr::NONE, 0, Serialize::Expr::NONE, 0); } else { static_assert(FLATBUFFERS_USE_STD_OPTIONAL); const auto make_optional_u64 = [](const std::optional &v) -> std::optional { diff --git a/src/Simplify_Exprs.cpp b/src/Simplify_Exprs.cpp index e82eb82acdf5..f7aa1f7288cc 100644 --- a/src/Simplify_Exprs.cpp +++ b/src/Simplify_Exprs.cpp @@ -392,7 +392,7 @@ Expr Simplify::visit(const Load *op, ExprInfo *info) { Expr new_index = b_index->value; int new_lanes = new_index.type().lanes(); Expr load = Load::make(op->type.with_lanes(new_lanes), op->name, b_index->value, - op->image, op->param, const_true(new_lanes, nullptr), align); + op->image, op->param, const_true(new_lanes, nullptr), align, op->is_streaming); return Broadcast::make(load, b_index->lanes); } else if (s_index && (s_index->is_concat() || @@ -410,7 +410,7 @@ Expr Simplify::visit(const Load *op, ExprInfo *info) { predicate_slice = mutate(predicate_slice, nullptr); Expr load = Load::make(op->type.with_lanes(new_lanes), op->name, new_index, - op->image, op->param, predicate_slice, ModulusRemainder{}); + op->image, op->param, predicate_slice, ModulusRemainder{}, op->is_streaming); loaded_vecs.emplace_back(std::move(load)); } return Shuffle::make(loaded_vecs, s_index->indices); @@ -439,13 +439,13 @@ Expr Simplify::visit(const Load *op, ExprInfo *info) { Expr permuted_load = Load::make(op->type, op->name, mr.to_expr(), op->image, - op->param, permuted_predicate, align); + op->param, permuted_predicate, align, op->is_streaming); int B = op->type.lanes() / A; return mutate(Shuffle::make_transpose(permuted_load, B), info); } else if (predicate.same_as(op->predicate) && index.same_as(op->index) && align == op->alignment) { return op; } else { - return Load::make(op->type, op->name, index, op->image, op->param, predicate, align); + return Load::make(op->type, op->name, index, op->image, op->param, predicate, align, op->is_streaming); } } diff --git a/src/Simplify_Internal.h b/src/Simplify_Internal.h index ea1c1e80f386..31074067292e 100644 --- a/src/Simplify_Internal.h +++ b/src/Simplify_Internal.h @@ -540,6 +540,8 @@ class Simplify : public VariadicVisitor { Stmt visit(const Acquire *op); Stmt visit(const Fork *op); Stmt visit(const Atomic *op); + Stmt visit(const StreamingStore *op); + Stmt visit(const StreamingLoads *op); Stmt visit(const HoistedStorage *op); std::pair, bool> mutate_with_changes(const std::vector &old_exprs); diff --git a/src/Simplify_Shuffle.cpp b/src/Simplify_Shuffle.cpp index 9f339dd26614..9080228193f6 100644 --- a/src/Simplify_Shuffle.cpp +++ b/src/Simplify_Shuffle.cpp @@ -86,7 +86,7 @@ Expr Simplify::visit(const Shuffle *op, ExprInfo *info) { t = first_load->type; t = t.with_lanes(op->indices.size()); return Load::make(t, first_load->name, shuffled_index, first_load->image, - first_load->param, shuffled_predicate, alignment); + first_load->param, shuffled_predicate, alignment, first_load->is_streaming); } } } diff --git a/src/Simplify_Stmts.cpp b/src/Simplify_Stmts.cpp index c8f9d3effc6d..b5b255ca5f21 100644 --- a/src/Simplify_Stmts.cpp +++ b/src/Simplify_Stmts.cpp @@ -368,7 +368,7 @@ Stmt Simplify::visit(const Store *op) { return Evaluate::make(0); } else if (scalar_pred && !is_const_one(scalar_pred->value)) { return IfThenElse::make(scalar_pred->value, - Store::make(op->name, value, index, op->param, const_true(value.type().lanes(), nullptr), align)); + Store::make(op->name, value, index, op->param, const_true(value.type().lanes(), nullptr), align, op->is_streaming)); } else if (is_undef(value) || (load && load->name == op->name && equal(load->index, index))) { // foo[x] = foo[x] or foo[x] = undef is a no-op return Evaluate::make(0); @@ -387,7 +387,8 @@ Stmt Simplify::visit(const Store *op) { idx, op->param, Shuffle::make_slice(predicate, lanes, 1, idx.type().lanes()), - ModulusRemainder{})); + ModulusRemainder{}, + op->is_streaming)); lanes += idx.type().lanes(); } Stmt s = Block::make(stores); @@ -420,14 +421,14 @@ Stmt Simplify::visit(const Store *op) { permuted_predicate = Shuffle::make_transpose(predicate, A); } return mutate(Store::make(op->name, permuted_value, mr.to_expr(), - op->param, permuted_predicate, align)); + op->param, permuted_predicate, align, op->is_streaming)); } else if (predicate.same_as(op->predicate) && value.same_as(op->value) && index.same_as(op->index) && align == op->alignment) { return op; } else { - return Store::make(op->name, value, index, op->param, predicate, align); + return Store::make(op->name, value, index, op->param, predicate, align, op->is_streaming); } } @@ -763,6 +764,28 @@ Stmt Simplify::visit(const Atomic *op) { } } +Stmt Simplify::visit(const StreamingStore *op) { + Stmt body = mutate(op->body); + if (is_no_op(body)) { + return Evaluate::make(0); + } else if (body.same_as(op->body)) { + return op; + } else { + return StreamingStore::make(op->producer_name, std::move(body)); + } +} + +Stmt Simplify::visit(const StreamingLoads *op) { + Stmt body = mutate(op->body); + if (is_no_op(body)) { + return Evaluate::make(0); + } else if (body.same_as(op->body)) { + return op; + } else { + return StreamingLoads::make(op->names, std::move(body)); + } +} + Stmt Simplify::visit(const HoistedStorage *op) { Stmt body = mutate(op->body); if (body.same_as(op->body)) { diff --git a/src/StageStridedLoads.cpp b/src/StageStridedLoads.cpp index b4df60a277c9..d19f9f606fcc 100644 --- a/src/StageStridedLoads.cpp +++ b/src/StageStridedLoads.cpp @@ -311,7 +311,7 @@ Stmt stage_strided_loads(const Stmt &stmt, const Target &target) { } else { // Might be a strided load after simplification return Load::make(l->type, l->name, self->mutate(l->index), l->image, l->param, - self->mutate(l->predicate), l->alignment); + self->mutate(l->predicate), l->alignment, l->is_streaming); } }); @@ -367,7 +367,7 @@ Stmt stage_strided_loads(const Stmt &stmt, const Target &target) { } Expr shared_load = Load::make(t, k.buf, idx, op->image, op->param, - const_true(lanes), op->alignment); + const_true(lanes), op->alignment, op->is_streaming); // We now need to pick a site to place our shared dense load. We // can't lift the shared load further out than k.scope, because that @@ -432,7 +432,7 @@ Stmt stage_strided_loads(const Stmt &stmt, const Target &target) { Type t = k.type.with_lanes(lanes); const Load *op = loads[0]; Expr dense_load = Load::make(t, k.buf, idx, op->image, op->param, - const_true(lanes), op->alignment - delta); + const_true(lanes), op->alignment - delta, op->is_streaming); dense_load = common_subexpression_elimination(dense_load); Expr shuf = Shuffle::make_slice(dense_load, delta, k.stride, k.lanes); for (const Load *l : loads) { @@ -470,7 +470,7 @@ Stmt stage_strided_loads(const Stmt &stmt, const Target &target) { Type t = k.type.with_lanes(lanes); const Load *op = loads[0]; Expr dense_load = Load::make(t, k.buf, idx, op->image, op->param, - const_true(lanes), op->alignment); + const_true(lanes), op->alignment, op->is_streaming); dense_load = common_subexpression_elimination(dense_load); Expr shuf = Shuffle::make_slice(dense_load, offset - final_offset, k.stride, k.lanes); for (const Load *l : loads) { @@ -503,7 +503,7 @@ Stmt stage_strided_loads(const Stmt &stmt, const Target &target) { Type t = k.type.with_lanes(lanes); const Load *op = loads[0]; Expr dense_load = Load::make(t, k.buf, idx, op->image, op->param, - const_true(lanes), op->alignment); + const_true(lanes), op->alignment, op->is_streaming); dense_load = common_subexpression_elimination(dense_load); Expr shuf = Shuffle::make_slice(dense_load, offset - first_offset, k.stride, k.lanes); for (const Load *l : loads) { @@ -521,9 +521,9 @@ Stmt stage_strided_loads(const Stmt &stmt, const Target &target) { Type t = k.type.with_lanes(half_lanes); const Load *op = loads[0]; Expr dense_load1 = Load::make(t, k.buf, idx1, op->image, op->param, - const_true(half_lanes), op->alignment); + const_true(half_lanes), op->alignment, op->is_streaming); Expr dense_load2 = Load::make(t, k.buf, idx2, op->image, op->param, - const_true(half_lanes), op->alignment + half_lanes - delta); + const_true(half_lanes), op->alignment + half_lanes - delta, op->is_streaming); dense_load1 = common_subexpression_elimination(dense_load1); dense_load2 = common_subexpression_elimination(dense_load2); Expr shuf1 = Shuffle::make_slice(dense_load1, 0, k.stride, k.lanes / 2); diff --git a/src/StorageFlattening.cpp b/src/StorageFlattening.cpp index da14d1b52530..ae1eca2595ba 100644 --- a/src/StorageFlattening.cpp +++ b/src/StorageFlattening.cpp @@ -12,6 +12,7 @@ #include "Simplify.h" #include "Substitute.h" +#include #include namespace Halide { @@ -44,6 +45,27 @@ class FlattenDimensions : public IRMutator { Scope<> realizations; bool in_gpu = false; + // The name of the Function whose Provide node is currently being + // flattened, if any. Used to avoid streaming self-loads. + string current_provide_name; + + // Whether the Provide node currently being flattened is wrapped in a + // StreamingStore node (see Stage::stream_stores), i.e. whether the Store + // node(s) it becomes should be marked non-temporal. + bool in_streaming_store = false; + + // The stream_loads() request (see Stage::stream_loads) of the + // StreamingLoads node currently wrapping the Provide being flattened. + // The default (present, empty), used whenever no such node is in + // scope, means stream nothing; nullopt (only ever set by an active + // StreamingLoads scope) means every direct load of another Func should + // stream; otherwise only loads of the named Funcs should. Names + // actually matched are recorded in stream_loads_matched (when + // non-null) so we can warn about any requested name that never turned + // out to be a direct load. + std::optional> stream_loads_names = std::set{}; + std::set *stream_loads_matched = nullptr; + Expr make_shape_var(string name, const string &field, size_t dim, const Buffer<> &buf, const Parameter ¶m) { ReductionDomain rdom; @@ -241,6 +263,35 @@ class FlattenDimensions : public IRMutator { return stmt; } + Stmt visit(const StreamingStore *op) override { + ScopedValue old_in_streaming_store(in_streaming_store, true); + return mutate(op->body); + } + + Stmt visit(const StreamingLoads *op) override { + std::optional> requested; + if (op->names) { + requested = std::set(op->names->begin(), op->names->end()); + } + ScopedValue>> old_names(stream_loads_names, requested); + std::set matched; + ScopedValue *> old_matched(stream_loads_matched, op->names ? &matched : nullptr); + + Stmt result = mutate(op->body); + + if (requested) { + for (const std::string &requested_name : *requested) { + if (!matched.count(requested_name)) { + user_warning << "stream_loads({" << requested_name << "}) was requested, " + << "but no direct load of \"" << requested_name + << "\" was found; did you mean to call stream_loads on a " + "different Stage?\n"; + } + } + } + return result; + } + Stmt visit(const Provide *op) override { internal_assert(op->values.size() == 1); @@ -266,6 +317,8 @@ class FlattenDimensions : public IRMutator { } } + ScopedValue old_provide_name(current_provide_name, op->name); + Expr value = mutate(op->values[0]); Expr predicate = mutate(op->predicate); if (in_gpu && textures.count(op->name)) { @@ -288,7 +341,7 @@ class FlattenDimensions : public IRMutator { return result; } else { Expr idx = mutate(flatten_args(op->name, op->args, Buffer<>(), output_buf)); - return Store::make(op->name, value, idx, output_buf, predicate, ModulusRemainder()); + return Store::make(op->name, value, idx, output_buf, predicate, ModulusRemainder(), in_streaming_store); } } @@ -338,8 +391,17 @@ class FlattenDimensions : public IRMutator { op->param); } else { Expr idx = mutate(flatten_args(op->name, op->args, op->image, op->param)); + bool is_streaming = false; + if ((op->call_type == Call::Halide || op->param.defined()) && + op->name != current_provide_name) { + bool matches = !stream_loads_names || stream_loads_names->count(op->name) != 0; + is_streaming |= matches; + if (matches && stream_loads_matched) { + stream_loads_matched->insert(op->name); + } + } return Load::make(op->type, op->name, idx, op->image, op->param, - const_true(op->type.lanes()), ModulusRemainder()); + const_true(op->type.lanes()), ModulusRemainder(), is_streaming); } } else { @@ -592,7 +654,7 @@ class PromoteToMemoryType : public IRMutator { if (t != op->type) { return Cast::make(op->type, Load::make(t, op->name, mutate(op->index), - op->image, op->param, mutate(op->predicate), ModulusRemainder())); + op->image, op->param, mutate(op->predicate), ModulusRemainder(), op->is_streaming)); } else { return IRMutator::visit(op); } @@ -602,7 +664,7 @@ class PromoteToMemoryType : public IRMutator { Type t = upgrade(op->value.type()); if (t != op->value.type()) { return Store::make(op->name, Cast::make(t, mutate(op->value)), mutate(op->index), - op->param, mutate(op->predicate), ModulusRemainder()); + op->param, mutate(op->predicate), ModulusRemainder(), op->is_streaming); } else { return IRMutator::visit(op); } diff --git a/src/VectorizeLoops.cpp b/src/VectorizeLoops.cpp index 945a1eb0545f..eab581801192 100644 --- a/src/VectorizeLoops.cpp +++ b/src/VectorizeLoops.cpp @@ -225,12 +225,12 @@ class RewriteAccessToVectorAlloc : public IRMutator { Expr visit(const Load *op) override { return Load::make(op->type, op->name, mutate_index(op->name, op->index), - op->image, op->param, mutate(op->predicate), mutate_alignment(op->name, op->alignment)); + op->image, op->param, mutate(op->predicate), mutate_alignment(op->name, op->alignment), op->is_streaming); } Stmt visit(const Store *op) override { return Store::make(op->name, mutate(op->value), mutate_index(op->name, op->index), - op->param, mutate(op->predicate), mutate_alignment(op->name, op->alignment)); + op->param, mutate(op->predicate), mutate_alignment(op->name, op->alignment), op->is_streaming); } public: @@ -304,7 +304,7 @@ class PredicateLoadStore : public IRMutator { return op; } vectorized = true; - return Load::make(op->type, op->name, index, op->image, op->param, predicate, op->alignment); + return Load::make(op->type, op->name, index, op->image, op->param, predicate, op->alignment, op->is_streaming); } Stmt visit(const Store *op) override { @@ -335,7 +335,7 @@ class PredicateLoadStore : public IRMutator { return op; } vectorized = true; - return Store::make(op->name, value, index, op->param, predicate, op->alignment); + return Store::make(op->name, value, index, op->param, predicate, op->alignment, op->is_streaming); } Expr visit(const Call *op) override { @@ -530,7 +530,7 @@ class VectorSubs : public IRMutator { int w = index.type().lanes(); predicate = widen(predicate, w); return Load::make(op->type.with_lanes(w), op->name, index, op->image, - op->param, predicate, op->alignment); + op->param, predicate, op->alignment, op->is_streaming); } } @@ -616,7 +616,7 @@ class VectorSubs : public IRMutator { const Load *load = true_value.as(); if (load) { - return Load::make(op->type.with_lanes(max_lanes), load->name, load->index, load->image, load->param, cond, load->alignment); + return Load::make(op->type.with_lanes(max_lanes), load->name, load->index, load->image, load->param, cond, load->alignment, load->is_streaming); } } @@ -705,7 +705,7 @@ class VectorSubs : public IRMutator { value.type().lanes(), index.type().lanes()}); return Store::make(op->name, widen(value, lanes), widen(index, lanes), - op->param, widen(predicate, lanes), op->alignment); + op->param, widen(predicate, lanes), op->alignment, op->is_streaming); } } @@ -1283,7 +1283,7 @@ class VectorSubs : public IRMutator { Expr new_load = Load::make(load_a->type.with_lanes(output_lanes), load_a->name, store_index, load_a->image, load_a->param, const_true(output_lanes), - ModulusRemainder{}); + ModulusRemainder{}, load_a->is_streaming); Expr lhs = cast(b.type().with_lanes(output_lanes), new_load); @@ -1292,7 +1292,7 @@ class VectorSubs : public IRMutator { b = binop(lhs, b); b = cast(new_load.type(), b); s = Store::make(store->name, b, store_index, store->param, - const_true(b.type().lanes()), store->alignment); + const_true(b.type().lanes()), store->alignment, store->is_streaming); } else { // Wrap any containing loops we still need (unrolled). We // enumerate the cartesian product of loop iteration values @@ -1303,7 +1303,7 @@ class VectorSubs : public IRMutator { Stmt store_template = Store::make(store->name, cast(new_load.type(), binop(lhs, b_var)), store_index, store->param, - const_true(output_lanes), ModulusRemainder{}); + const_true(output_lanes), ModulusRemainder{}, store->is_streaming); std::string full_b_var_name = unique_name('b'); Expr full_b_var = Variable::make(b.type(), full_b_var_name); diff --git a/src/halide_ir.fbs b/src/halide_ir.fbs index 7f3492684743..82ce6afb68e2 100644 --- a/src/halide_ir.fbs +++ b/src/halide_ir.fbs @@ -162,6 +162,7 @@ table Store { index: Expr; param_name: string; alignment: ModulusRemainder; + is_streaming: bool; } table Provide { @@ -382,6 +383,7 @@ table Load { param_name: string; alignment: ModulusRemainder; type: Type; + is_streaming: bool; } table Ramp { @@ -606,6 +608,9 @@ table StageSchedule { allow_race_conditions: bool = false; atomic: bool = false; override_atomic_associativity_test: bool = false; + stream_stores: bool = false; + stream_loads_all: bool = false; + stream_loads_names: [string]; } table Definition { diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index 821a62525a2d..55fe27d68b92 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -317,6 +317,9 @@ tests( stmt_to_html.cpp storage_folding.cpp store_in.cpp + streaming.cpp + streaming_in.cpp + streaming_specialize.cpp strict_float.cpp strict_float_bounds.cpp strict_fma.cpp @@ -500,3 +503,7 @@ set_target_properties( PROPERTIES ENABLE_EXPORTS TRUE ) + +if (WITH_SERIALIZATION) + target_compile_definitions(correctness_streaming PRIVATE TEST_WITH_SERIALIZATION) +endif () diff --git a/test/correctness/simd_op_check.h b/test/correctness/simd_op_check.h index 828b16dc969b..7466fd3e5bb9 100644 --- a/test/correctness/simd_op_check.h +++ b/test/correctness/simd_op_check.h @@ -6,6 +6,7 @@ #include "halide_thread_pool.h" #include "test_sharding.h" +#include #include #include @@ -213,12 +214,12 @@ class SimdOpCheckTest { return true; } } while (*str++); - } else if (*p == ' ') { // ignore whitespace in pattern + } else if (std::isspace(static_cast(*p))) { // ignore whitespace in pattern p++; if (wildcard_match(p, str)) { return true; } - } else if (*str == ' ') { // ignore whitespace in string + } else if (std::isspace(static_cast(*str))) { // ignore whitespace in string str++; if (wildcard_match(p, str)) { return true; diff --git a/test/correctness/simd_op_check_arm.cpp b/test/correctness/simd_op_check_arm.cpp index e7b1ca843c28..a5e879622682 100644 --- a/test/correctness/simd_op_check_arm.cpp +++ b/test/correctness/simd_op_check_arm.cpp @@ -22,9 +22,35 @@ class SimdOpCheckARM : public SimdOpCheckTest { void add_tests() override { if (target.arch == Target::ARM) { check_neon_all(); + check_streaming_accesses(); } } + void check_streaming_accesses() { + if (target.bits != 64) { + return; + } + + auto streaming_copy = [&]() { + Func f{"f"}, reader{"reader"}; + + f(x, y) = in_u8(x); + f.compute_root() + .vectorize(x, 32) + .stream_stores(); + + reader(x, y) = f(x, y) + 1; + reader.compute_root() + .vectorize(x, 32) + .stream_loads(); + + return reader(x, y); + }; + + check("ldnp q", 32, streaming_copy()); + check("stnp q", 32, streaming_copy()); + } + void check_neon_all() { Expr f64_1 = in_f64(x), f64_2 = in_f64(x + 16), f64_3 = in_f64(x + 32); Expr f32_1 = in_f32(x), f32_2 = in_f32(x + 16), f32_3 = in_f32(x + 32); diff --git a/test/correctness/simd_op_check_x86.cpp b/test/correctness/simd_op_check_x86.cpp index 421fccb5c9a0..4d23882c0ad5 100644 --- a/test/correctness/simd_op_check_x86.cpp +++ b/test/correctness/simd_op_check_x86.cpp @@ -42,9 +42,36 @@ class SimdOpCheckX86 : public SimdOpCheckTest { // Queue up a bunch of tasks representing each test to run. if (target.arch == Target::X86) { check_sse_and_avx(); + check_streaming_accesses(); } } + void check_streaming_accesses() { + if (!use_avx2) { + return; + } + + auto streaming_copy = [&]() { + Func f{"f"}, reader{"reader"}; + + f(x, y) = in_u8(x); + f.compute_root() + .vectorize(x, 32) + .stream_stores(); + + reader(x, y) = f(x, y) + 1; + reader.compute_root() + .vectorize(x, 32) + .stream_loads(); + + return reader(x, y); + }; + + check("vmovntdqa (", 32, streaming_copy()); + check("vmovntdq %ymm", 32, streaming_copy()); + check("sfence", 32, streaming_copy()); + } + void check_sse_and_avx() { Expr f64_1 = in_f64(x), f64_2 = in_f64(x + 16), f64_3 = in_f64(x + 32); Expr f32_1 = in_f32(x), f32_2 = in_f32(x + 16), f32_3 = in_f32(x + 32); diff --git a/test/correctness/streaming.cpp b/test/correctness/streaming.cpp new file mode 100644 index 000000000000..a5ec60843729 --- /dev/null +++ b/test/correctness/streaming.cpp @@ -0,0 +1,193 @@ +#include "Halide.h" + +#include +#include +#ifdef TEST_WITH_SERIALIZATION +#include +#endif +#include +#include +#include + +using namespace Halide; +using namespace Halide::Internal; + +namespace { + +class CheckStreamingAccesses : public IRMutator { +public: + using IRMutator::mutate; + + Stmt mutate(const Stmt &s) override { + class Finder : public IRVisitor { + public: + // (name, is_streaming) pairs, in the order they're visited (which, + // for a given name, is the order the Stages that touch it run in). + std::vector> loads; + std::vector> stores; + int stream_store_fences = 0; + + protected: + using IRVisitor::visit; + + void visit(const Load *op) override { + loads.emplace_back(op->name, op->is_streaming); + IRVisitor::visit(op); + } + + void visit(const Store *op) override { + stores.emplace_back(op->name, op->is_streaming); + IRVisitor::visit(op); + } + + void visit(const Call *op) override { + if (op->is_intrinsic(Call::stream_store_fence)) { + stream_store_fences++; + } + IRVisitor::visit(op); + } + } finder; + + s.accept(&finder); + + auto has_access = [](const std::vector> &accesses, + const std::string &name, bool streaming) { + for (const auto &a : accesses) { + if (a.first == name && a.second == streaming) { + return true; + } + } + return false; + }; + + auto require = [&](const std::vector> &accesses, + const std::string &name, bool streaming, const char *kind) { + if (!has_access(accesses, name, streaming)) { + std::fprintf(stderr, "Expected a %s access to %s\n", kind, name.c_str()); + std::exit(1); + } + }; + + require(finder.loads, "input_a", true, "streaming load"); + require(finder.loads, "streaming_producer", true, "streaming load"); + require(finder.stores, "streaming_producer", true, "streaming store"); + require(finder.stores, "output", true, "streaming store"); + require(finder.loads, "input_b", false, "ordinary load"); + require(finder.loads, "ordinary_producer", false, "ordinary load"); + require(finder.stores, "ordinary_producer", false, "ordinary store"); + + // self_referencing has two Stages: a pure definition (no self-load) + // followed by an update that self-loads. Only the update Stage was + // scheduled with stream_stores() (see below), so the pure + // definition's store is ordinary. The update Stage has no RVar, so + // it's legal to stream, and since the fence for a streamed Stage is + // emitted immediately after that Stage's own production, the + // update's self-load (of the pure definition's already-visible, + // ordinary data) is unaffected. A downstream consumer's read is not + // a self-load, so it may still stream. + // + // Each Stage's store may appear more than once (e.g. a vectorized + // main loop plus a scalar tail from loop partitioning), so check + // that the sequence is "some ordinary stores, then some streaming + // stores" rather than requiring an exact count. + std::vector self_referencing_stores; + for (const auto &store : finder.stores) { + if (store.first == "self_referencing") { + self_referencing_stores.push_back(store.second); + } + } + bool well_formed = !self_referencing_stores.empty() && + self_referencing_stores.front() == false && + self_referencing_stores.back() == true && + std::is_sorted(self_referencing_stores.begin(), self_referencing_stores.end()); + if (!well_formed) { + std::fprintf(stderr, "Expected self_referencing's pure definition to store ordinarily " + "and its update to store streamingly\n"); + std::exit(1); + } + require(finder.loads, "self_referencing", false, "ordinary load"); // the update's self-load + require(finder.loads, "self_referencing", true, "streaming load"); // output's downstream read + + // One fence each for streaming_producer, self_referencing, and output. + if (finder.stream_store_fences != 3) { + std::fprintf(stderr, "Expected exactly 3 streaming store fences, found %d\n", + finder.stream_store_fences); + std::exit(1); + } + + return s; + } +}; + +} // namespace + +int main(int argc, char **argv) { + Var x; + ImageParam input_a(Int(32), 1, "input_a"); + ImageParam input_b(Int(32), 1, "input_b"); + Func streaming_producer("streaming_producer"); + Func ordinary_producer("ordinary_producer"); + Func self_referencing("self_referencing"); + Func output("output"); + + streaming_producer(x) = input_a(x) + 1; + ordinary_producer(x) = input_b(x) + 2; + + // An update definition with a self-load. + self_referencing(x) = 0; + self_referencing(x) += x; + + output(x) = streaming_producer(x) + ordinary_producer(x) + self_referencing(x); + + streaming_producer.compute_root().stream_stores(); + // An ImageParam can be named directly: it's resolved to the underlying + // Parameter, since the ImageParam's own wrapper Func is inlined away + // before storage flattening ever sees it. + streaming_producer.stream_loads({input_a}); + ordinary_producer.compute_root(); + self_referencing.compute_root(); + self_referencing.update(0).stream_stores(); + output.stream_stores(); + // stream_loads() is a property of the reading Stage, not the Func being + // read: output is the one that directly reads streaming_producer and + // self_referencing, so it's the one that requests streaming those two + // reads (ordinary_producer's read is deliberately left off this list). + output.stream_loads({streaming_producer, self_referencing}); + + streaming_producer.vectorize(x, 8); + ordinary_producer.vectorize(x, 8); + self_referencing.vectorize(x, 8); + self_referencing.update(0).vectorize(x, 8); + output.vectorize(x, 8); + + constexpr int size = 1024; + Buffer a(size), b(size); + for (int i = 0; i < size; i++) { + a(i) = i; + b(i) = 3 * i; + } + input_a.set(a); + input_b.set(b); + + Pipeline pipeline(output); +#ifdef TEST_WITH_SERIALIZATION + // Verify that the access directives survive pipeline serialization as + // well as the normal lowering pipeline. + std::vector serialized; + std::map params; + serialize_pipeline(pipeline, serialized, params); + pipeline = deserialize_pipeline(serialized, params); +#endif + pipeline.add_custom_lowering_pass(new CheckStreamingAccesses); + + Buffer result = pipeline.realize({size}); + for (int i = 0; i < size; i++) { + if (result(i) != 5 * i + 3) { + std::fprintf(stderr, "Incorrect result at %d: %d\n", i, result(i)); + return 1; + } + } + + std::printf("Success!\n"); + return 0; +} diff --git a/test/correctness/streaming_in.cpp b/test/correctness/streaming_in.cpp new file mode 100644 index 000000000000..931259a73556 --- /dev/null +++ b/test/correctness/streaming_in.cpp @@ -0,0 +1,210 @@ +#include +using namespace Halide; +using namespace Halide::Internal; + +// stream_loads() is a Stage property: a Stage requests streaming for its own +// direct reads of named Funcs (or, with no arguments, every direct read +// except a self-load). The request is resolved against the Stage's fully +// inlined body, so it correctly reaches through inlined intermediates -- but +// it can only ever affect a Stage that survives lowering as a real, +// materialized read. This test verifies the design: +// +// 1. h.stream_loads({f}) works even though h only directly calls an +// inlined intermediate g (which in turn calls f): once g inlines into +// h, h's flattened body directly loads f, and that's the load that gets +// marked. +// 2. For the residual-network case that motivated this feature -- +// h(x) = f(x) + g(x), g(x) = ...f(x)..., wanting f streamed for h but +// not for g -- the correct, zero-overhead idiom is h.stream_loads({f}). +// g's independent read of f is left untouched, since stream_loads() only +// ever affects the one Stage that requested it. +// 3. If f.in(h).stream_loads() is used without specifying a compute location, +// it's diagnosed with a user_warning rather than silently dropped. + +namespace { + +struct Finder : public IRVisitor { + using IRVisitor::visit; + + std::string current_producer; + // (producer name, target name) -> (streaming count, ordinary count) + std::map, std::pair> loads; + +protected: + void visit(const ProducerConsumer *op) override { + if (op->is_producer) { + ScopedValue old(current_producer, op->name); + IRVisitor::visit(op); + } else { + IRVisitor::visit(op); + } + } + + void visit(const Load *op) override { + auto &p = loads[{current_producer, op->name}]; + (op->is_streaming ? p.first : p.second)++; + IRVisitor::visit(op); + } +}; + +int streaming_count(Finder &finder, const std::string &producer, const std::string &target) { + auto it = finder.loads.find({producer, target}); + return it == finder.loads.end() ? 0 : it->second.first; +} + +int ordinary_count(Finder &finder, const std::string &producer, const std::string &target) { + auto it = finder.loads.find({producer, target}); + return it == finder.loads.end() ? 0 : it->second.second; +} + +struct WarningCounter final : CompileTimeErrorReporter { + int warnings_occurred = 0; + + void warning(const char *msg) override { + std::fprintf(stderr, "Warning: %s\n", msg); + warnings_occurred++; + } + + [[noreturn]] void error(const char *msg) override { + std::fprintf(stderr, "Unexpected error: %s\n", msg); + exit(1); + } +}; + +} // namespace + +int main(int argc, char **argv) { + // Case 1: h only directly calls the inlined intermediate g, which + // directly calls f. h.stream_loads({f}) still works: it's resolved + // against h's fully-inlined body, where the load of f is now direct. + { + Func f("streaming_in_case1_f"), g("streaming_in_case1_g"), h("streaming_in_case1_h"); + Var x; + f(x) = cast(x) * 3.0f + 1.0f; + g(x) = f(x); // left at the default (inline) schedule + h(x) = g(x) + 1.0f; + + f.compute_root(); + h.compute_root(); + h.stream_loads({f}); + + Finder finder; + Module m = h.compile_to_module({}, h.name()); + for (const auto &fn : m.functions()) { + fn.body.accept(&finder); + } + if (streaming_count(finder, h.name(), f.name()) != 1 || + ordinary_count(finder, h.name(), f.name()) != 0) { + std::fprintf(stderr, "Case 1: expected h.stream_loads({f}) to mark h's " + "(post-inline) direct load of f, found streaming=%d " + "ordinary=%d\n", + streaming_count(finder, h.name(), f.name()), + ordinary_count(finder, h.name(), f.name())); + return 1; + } + } + + // Case 2: the residual-network case, solved directly. h(x) = f(x) + + // g(x), g(x) also depends on f(x). We want f's data streamed when read + // for h, but not when read by g -- with no extra Func and no extra + // buffer: just h.stream_loads({f}). + { + Func f("streaming_in_case2_f"), g("streaming_in_case2_g"), h("streaming_in_case2_h"); + Var x; + f(x) = cast(x) * 3.0f + 1.0f; + g(x) = f(x) * 2.0f; + h(x) = f(x) + g(x); + + f.compute_root(); + g.compute_root(); + h.compute_root(); + h.stream_loads({f}); + + Finder finder; + Module m = h.compile_to_module({}, h.name()); + for (const auto &fn : m.functions()) { + fn.body.accept(&finder); + } + if (streaming_count(finder, h.name(), f.name()) != 1 || + ordinary_count(finder, h.name(), f.name()) != 0) { + std::fprintf(stderr, "Case 2: expected h's own read of f to be streaming, " + "found streaming=%d ordinary=%d\n", + streaming_count(finder, h.name(), f.name()), + ordinary_count(finder, h.name(), f.name())); + return 1; + } + if (streaming_count(finder, g.name(), f.name()) != 0 || + ordinary_count(finder, g.name(), f.name()) != 1) { + std::fprintf(stderr, "Case 2: expected g's independent read of f to remain " + "ordinary, found streaming=%d ordinary=%d\n", + streaming_count(finder, g.name(), f.name()), + ordinary_count(finder, g.name(), f.name())); + return 1; + } + + // Correctness, not just IR shape. + Buffer result = h.realize({16}); + for (int i = 0; i < 16; i++) { + float f_val = i * 3.0f + 1.0f; + float expected = f_val + f_val * 2.0f; + if (result(i) != expected) { + std::fprintf(stderr, "Case 2: incorrect result at %d: %f (expected %f)\n", + i, result(i), expected); + return 1; + } + } + } + + // Case 3: an f.in(h) wrapper, left at its default (inlined) schedule. + // stream_loads() is a complete no-op -- with no materialized Stage of + // its own, there's no distinct read left to mark -- and it's diagnosed + // with a user_warning, matching vectorize/parallel/unroll/split/bound + // on an inlined Func (see validate_schedule_inlined_function). + { + Func f("streaming_in_case3_f"), g("streaming_in_case3_g"), h("streaming_in_case3_h"); + Var x; + f(x) = cast(x) * 3.0f + 1.0f; + g(x) = f(x) * 2.0f; + h(x) = f(x) + g(x); + + Func f_in_h = f.in(h); + f_in_h.stream_loads({f}); // no compute_at/compute_root: stays inlined + + f.compute_root(); + g.compute_root(); + h.compute_root(); + + WarningCounter reporter; + set_custom_compile_time_error_reporter(&reporter); + Module m = h.compile_to_module({}, h.name()); + set_custom_compile_time_error_reporter(nullptr); + if (reporter.warnings_occurred == 0) { + std::fprintf(stderr, "Case 3: expected a user_warning when stream_loads() is " + "requested on a Func that ends up inlined, found none\n"); + return 1; + } + + int streaming_loads_found = 0; + for (const auto &fn : m.functions()) { + visit_with(fn.body, [&](auto *self, const auto *op) { + if constexpr (std::is_same_v) { + ++streaming_loads_found; + } else if constexpr (std::is_same_v) { + if (op->name.rfind(f.name(), 0) == 0 && op->is_streaming) { + ++streaming_loads_found; + } + } + self->visit_base(op); + }); + } + if (streaming_loads_found != 0) { + std::fprintf(stderr, "Case 3: expected stream_loads() on an inlined " + "Func::in() wrapper to be a complete no-op, found %d\n", + streaming_loads_found); + return 1; + } + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/streaming_specialize.cpp b/test/correctness/streaming_specialize.cpp new file mode 100644 index 000000000000..f94ef1352556 --- /dev/null +++ b/test/correctness/streaming_specialize.cpp @@ -0,0 +1,85 @@ +#include +using namespace Halide; +using namespace Halide::Internal; + +// Each specialization is its own Definition and may independently request +// streaming (see Stage::stream_stores); a Stage's Definition can therefore +// be expanded, by lowering, into more than one Provide node (one per +// specialization, plus the base/default case). The decision of which resulting +// Store is streaming -- and whether a fence is needed at all -- must track the +// exact Definition each Provide came from. +int main(int argc, char **argv) { + ImageParam in(Int(32), 1, "in"); + Expr cond = in.width() >= 8; + + // Only the (vectorized) specialization streams; the base case does not. + // The specialization also vectorizes so that the two branches remain + // distinguishable all the way through lowering: if they differed only + // in stream_stores(), which has no effect until Store nodes are built + // during storage flattening, an earlier generic simplification pass + // would see two structurally-identical branches and collapse them into + // one before we could observe them separately. + Func f("stream_specialize"); + Var x; + f(x) = in(x); + Stage spec = f.specialize(cond); + spec.vectorize(x, 8); + spec.stream_stores(); + f.compute_root(); + + // Verify the lowered structure directly (rather than via realize()), + // since binding a concrete buffer to `in` before lowering would let the + // specialization condition be proven statically, collapsing the two + // branches into one before we could observe them separately. + int stores_streaming = 0, stores_ordinary = 0, fences = 0; + Module m = f.compile_to_module({in}, f.name()); + for (const auto &fn : m.functions()) { + visit_with(fn.body, [&](auto *self, const auto *op) { + if constexpr (std::is_same_v) { + if (op->name.rfind("stream_specialize", 0) == 0) { + ++(op->is_streaming ? stores_streaming : stores_ordinary); + } + } else if constexpr (std::is_same_v) { + if (op->is_intrinsic(Call::stream_store_fence)) { + ++fences; + } + } + self->visit_base(op); + }); + } + if (stores_streaming == 0 || stores_ordinary == 0) { + std::fprintf(stderr, "Expected both a streaming store (from the vectorized " + "specialization) and an ordinary store (from the base " + "case) for stream_specialize, found streaming=%d ordinary=%d\n", + stores_streaming, stores_ordinary); + return 1; + } + if (fences != 1) { + std::fprintf(stderr, "Expected exactly 1 streaming store fence for " + "stream_specialize (covering both specialization " + "branches), found %d\n", + fences); + return 1; + } + + // Correctness check: realize with concrete inputs on both sides of the + // specialization condition. + for (int width : {4, 16}) { + Buffer buf(width); + for (int i = 0; i < width; i++) { + buf(i) = i; + } + in.set(buf); + Buffer result = f.realize({width}); + for (int i = 0; i < width; i++) { + if (result(i) != i) { + std::fprintf(stderr, "stream_specialize: incorrect result at %d (width %d): %d\n", + i, width, result(i)); + return 1; + } + } + } + + printf("Success!\n"); + return 0; +} diff --git a/test/error/CMakeLists.txt b/test/error/CMakeLists.txt index 0d18a7a87a3a..32c8a89b574b 100644 --- a/test/error/CMakeLists.txt +++ b/test/error/CMakeLists.txt @@ -118,6 +118,8 @@ tests( split_non_innermost_predicated.cpp split_same_var_names.cpp store_at_without_compute_at.cpp + stream_loads_illegal_self.cpp + stream_stores_illegal_rvar.cpp thread_id_outside_block_id.cpp too_many_args.cpp treat_rvar_as_var.cpp diff --git a/test/error/stream_loads_illegal_self.cpp b/test/error/stream_loads_illegal_self.cpp new file mode 100644 index 000000000000..169cdfe7440d --- /dev/null +++ b/test/error/stream_loads_illegal_self.cpp @@ -0,0 +1,17 @@ +#include "Halide.h" + +using namespace Halide; + +int main(int argc, char **argv) { + Func f("f"); + Var x; + + f(x) = x; + + // A Stage can't stream loads of itself: naming f's own Stage in its own + // stream_loads() request is illegal. + f.stream_loads({f}); + + printf("Success!\n"); + return 0; +} diff --git a/test/error/stream_stores_illegal_rvar.cpp b/test/error/stream_stores_illegal_rvar.cpp new file mode 100644 index 000000000000..77b8ac2e3458 --- /dev/null +++ b/test/error/stream_stores_illegal_rvar.cpp @@ -0,0 +1,20 @@ +#include "Halide.h" + +using namespace Halide; + +int main(int argc, char **argv) { + Func hist("hist"); + Var x; + ImageParam indices(Int(32), 1, "indices"); + RDom r(0, 100); + + hist(x) = 0; + hist(indices(r)) += 1; + + // The scatter above means Halide can't prove r's iterations don't + // collide, so stream_stores() is illegal on this update. + hist.update(0).stream_stores(); + + printf("Success!\n"); + return 0; +} diff --git a/test/performance/memcpy.cpp b/test/performance/memcpy.cpp index 7e604ec59a3e..68f8b17149d5 100644 --- a/test/performance/memcpy.cpp +++ b/test/performance/memcpy.cpp @@ -2,8 +2,10 @@ #include "halide_benchmark.h" #include "halide_test_dirs.h" +#include #include #include +#include using namespace Halide; using namespace Halide::Tools; @@ -20,7 +22,14 @@ int main(int argc, char **argv) { Var x; dst(x) = src(x); - dst.vectorize(x, 32, TailStrategy::GuardWithIf); + // LLVM requires vector-width alignment to select the native streaming + // instructions on some targets. Halide::Buffer allocations satisfy this. + src.set_host_alignment(32); + dst.output_buffer().set_host_alignment(32); + + dst.vectorize(x, 32, TailStrategy::GuardWithIf) + .stream_loads({src}) + .stream_stores(); dst.compile_to_assembly(Internal::get_test_tmp_dir() + "halide_memcpy.s", {src}, "halide_memcpy"); dst.compile_jit(); @@ -32,19 +41,31 @@ int main(int argc, char **argv) { src.set(input); - double t1 = benchmark([&]() { + auto halide_copy = [&]() { dst.realize(output); - }); - - double t2 = benchmark([&]() { + }; + auto system_copy = [&]() { memcpy(output.data(), input.data(), input.width()); - }); + }; + + // Warm up, then alternate the benchmarks to reduce the effect of thermal + // and frequency drift. + halide_copy(); + system_copy(); + + double t_halide = std::numeric_limits::infinity(); + double t_memcpy = std::numeric_limits::infinity(); + for (int i = 0; i < 5; i++) { + t_halide = std::min(t_halide, (double)benchmark(halide_copy)); + t_memcpy = std::min(t_memcpy, (double)benchmark(system_copy)); + } - printf("system memcpy: %.3e byte/s\n", buffer_size / t2); - printf("halide memcpy: %.3e byte/s\n", buffer_size / t1); + printf("system memcpy: %.2f GB/s\n", (buffer_size / t_memcpy) / 1e9); + printf("halide memcpy: %.2f GB/s\n", (buffer_size / t_halide) / 1e9); - // memcpy will win by a little bit for large inputs because it uses streaming stores - if (t1 > t2 * 3) { + // Streaming directives should bring the generated copy close to the + // platform memcpy implementation for a large contiguous transfer. + if (t_halide > t_memcpy * 1.2) { printf("Halide memcpy is slower than it should be.\n"); return 1; } From 8cee1d82f5e77e9d466df42dbce08584433f164f Mon Sep 17 00:00:00 2001 From: Alex Reinking Date: Wed, 15 Jul 2026 23:22:04 -0400 Subject: [PATCH 2/4] Add missing visitor overloads for StreamingStore and StreamingLoads --- src/Derivative.cpp | 6 +++ src/ModulusRemainder.cpp | 10 +++++ src/Monotonic.cpp | 8 ++++ src/StmtToHTML.cpp | 83 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 107 insertions(+) diff --git a/src/Derivative.cpp b/src/Derivative.cpp index a7e9ade253fe..041fb721513a 100644 --- a/src/Derivative.cpp +++ b/src/Derivative.cpp @@ -139,6 +139,12 @@ class ReverseAccumulationVisitor : public IRVisitor { void visit(const Atomic *op) override { internal_error << "Encounter unexpected statement \"Atomic\" when differentiating."; } + void visit(const StreamingStore *op) override { + internal_error << "Encounter unexpected statement \"StreamingStore\" when differentiating."; + } + void visit(const StreamingLoads *op) override { + internal_error << "Encounter unexpected statement \"StreamingLoads\" when differentiating."; + } void visit(const HoistedStorage *op) override { internal_error << "Encounter unexpected statement \"HoistedStorage\" when differentiating."; } diff --git a/src/ModulusRemainder.cpp b/src/ModulusRemainder.cpp index 5cb490a083de..f4f9752823e2 100644 --- a/src/ModulusRemainder.cpp +++ b/src/ModulusRemainder.cpp @@ -77,6 +77,8 @@ class ComputeModulusRemainder : public IRVisitor { void visit(const VectorReduce *) override; void visit(const Prefetch *) override; void visit(const Atomic *) override; + void visit(const StreamingStore *) override; + void visit(const StreamingLoads *) override; void visit(const HoistedStorage *) override; }; @@ -283,6 +285,14 @@ void ComputeModulusRemainder::visit(const Atomic *) { internal_error << "modulus_remainder of statement\n"; } +void ComputeModulusRemainder::visit(const StreamingStore *) { + internal_error << "modulus_remainder of statement\n"; +} + +void ComputeModulusRemainder::visit(const StreamingLoads *) { + internal_error << "modulus_remainder of statement\n"; +} + void ComputeModulusRemainder::visit(const HoistedStorage *) { internal_error << "modulus_remainder of statement\n"; } diff --git a/src/Monotonic.cpp b/src/Monotonic.cpp index 82934a31de5e..63e4f01f2ef1 100644 --- a/src/Monotonic.cpp +++ b/src/Monotonic.cpp @@ -489,6 +489,14 @@ class DerivativeBounds : public IRVisitor { internal_error << "Monotonic of statement\n"; } + void visit(const StreamingStore *op) override { + internal_error << "Monotonic of statement\n"; + } + + void visit(const StreamingLoads *op) override { + internal_error << "Monotonic of statement\n"; + } + void visit(const HoistedStorage *op) override { internal_error << "Monotonic of statement\n"; } diff --git a/src/StmtToHTML.cpp b/src/StmtToHTML.cpp index a40f3bfe51a6..ef8d719cb855 100644 --- a/src/StmtToHTML.cpp +++ b/src/StmtToHTML.cpp @@ -2241,6 +2241,89 @@ class HTMLCodePrinter : public IRVisitor { print_ln(); } + void visit(const StreamingStore *op) override { + // Give this node a unique id + int id = gen_unique_id(); + + // Start a dive to hold code for this streaming store + print_opening_tag("div", "StreamingStore"); + + // Generate the show hide icon/text buttons + print_show_hide_btn_begin(id); + + // -- print text + print_html_element("span", "matched keyword", "streaming_store"); + print_html_element("span", "matched", "("); + print_html_element("span", "Symbol", op->producer_name); + print_html_element("span", "matched", ")"); + + // Open code block to hold streaming store body + print_opening_brace(); + print_show_hide_btn_end(op); + + // Open indented div to hold streaming store code + print_opening_tag("div", "indent StreamingStoreBody", "", id); + + // Print streaming store body + print(op->body); + + // Close indented div holding body code + print_closing_tag("div"); + + // Close code block holding fork body + print_html_element("span", "matched ClosingBrace cb-" + std::to_string(id), "}"); + + // Close div holding this streaming store + print_closing_tag("div"); + print_ln(); + } + + void visit(const StreamingLoads *op) override { + // Give this node a unique id + int id = gen_unique_id(); + + // Start a dive to hold code for this streaming loads + print_opening_tag("div", "StreamingLoads"); + + // Generate the show hide icon/text buttons + print_show_hide_btn_begin(id); + + // -- print text + print_html_element("span", "matched keyword", "streaming_loads"); + print_html_element("span", "matched", "("); + if (!op->names) { + print_html_element("span", "Symbol", "all"); + } else { + for (size_t i = 0; i < op->names->size(); i++) { + if (i > 0) { + print_html_element("span", "matched", ", "); + } + print_html_element("span", "Symbol", (*op->names)[i]); + } + } + print_html_element("span", "matched", ")"); + + // Open code block to hold streaming loads body + print_opening_brace(); + print_show_hide_btn_end(op); + + // Open indented div to hold streaming loads code + print_opening_tag("div", "indent StreamingLoadsBody", "", id); + + // Print streaming loads body + print(op->body); + + // Close indented div holding body code + print_closing_tag("div"); + + // Close code block holding fork body + print_html_element("span", "matched ClosingBrace cb-" + std::to_string(id), "}"); + + // Close div holding this streaming loads + print_closing_tag("div"); + print_ln(); + } + void visit(const HoistedStorage *op) override { // Give this node a unique id int id = gen_unique_id(); From 21fe7d66b0e009abf45fadfc8052b19bf06aaff6 Mon Sep 17 00:00:00 2001 From: Alex Reinking Date: Thu, 16 Jul 2026 17:15:39 -0400 Subject: [PATCH 3/4] Respond to review comments --- src/CodeGen_C.cpp | 2 +- src/IR.h | 15 +++++++-------- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/src/CodeGen_C.cpp b/src/CodeGen_C.cpp index 6c406a449fd6..a6cc4dd3cac9 100644 --- a/src/CodeGen_C.cpp +++ b/src/CodeGen_C.cpp @@ -1864,7 +1864,7 @@ void CodeGen_C::visit(const Call *op) { } else if (op->is_intrinsic(Call::size_of_halide_buffer_t)) { rhs << "(sizeof(halide_buffer_t))"; } else if (op->is_intrinsic(Call::stream_store_fence)) { - // The C backend does not implement non-temporal loads/stores. + user_warning << "The C backend does not implement non-temporal loads/stores."; rhs << "0"; } else if (op->is_intrinsic(Call::strict_fma)) { internal_assert(op->args.size() == 3) diff --git a/src/IR.h b/src/IR.h index 7f8bcd9e9644..9c7f307f7cde 100644 --- a/src/IR.h +++ b/src/IR.h @@ -1164,10 +1164,10 @@ struct Atomic : public StmtNode { /** Marks the store(s) produced by the wrapped body as requesting * non-temporal (streaming) stores, as scheduled via Stage::stream_stores. - * Created directly around a Provide node (or an Atomic node wrapping one) - * during scheduling, and consumed by storage flattening, which uses it to - * set Store::is_streaming on the resulting Store node(s) and then discards - * it -- it should not survive past that point. */ + * Created directly around a Provide node (or an Atomic node wrapping + * one) during scheduling and is never nested. Consumed by storage flattening, + * which uses it to set Store::is_streaming on the resulting Store node(s) and + * then discards it. */ struct StreamingStore : public StmtNode { std::string producer_name; Stmt body; @@ -1183,10 +1183,9 @@ struct StreamingStore : public StmtNode { * Stage::stream_loads. If `names` is nullopt, every direct load of another * Func is streamed (except a self-load, which is never streamed); * otherwise only loads of Funcs named in `*names` are streamed. Created - * directly around a Stage's body during scheduling, and consumed by - * storage flattening, which uses it to set Load::is_streaming on the - * resulting Load node(s) and then discards it -- it should not survive - * past that point. */ + * directly around a Provide node during scheduling and is never nested. + * Consumed by storage flattening, which uses it to set Load::is_streaming + * on the resulting Load node(s) and then discards it. */ struct StreamingLoads : public StmtNode { std::optional> names; Stmt body; From a9baf07b5886a80b6d3f6b46994a47fbb5860432 Mon Sep 17 00:00:00 2001 From: Alex Reinking Date: Thu, 16 Jul 2026 17:21:58 -0400 Subject: [PATCH 4/4] Make stream_loads override rather than accumulate --- src/Func.cpp | 8 +++----- test/correctness/streaming_in.cpp | 1 - 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/src/Func.cpp b/src/Func.cpp index d5ab769f79f4..daa3f7ac0e02 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -1677,10 +1677,7 @@ Stage &Stage::stream_loads() { } Stage &Stage::stream_loads(const std::vector &funcs) { - auto &names = definition.schedule().stream_loads_names(); - if (!names) { - names = std::vector{}; - } + std::vector names; for (const Func &f : funcs) { std::string target_name = f.name(); if (const Call *call = f.function().is_wrapper(); call && call->param.defined()) { @@ -1693,8 +1690,9 @@ Stage &Stage::stream_loads(const std::vector &funcs) { user_assert(target_name != function.name()) << "Can't stream loads of \"" << target_name << "\" in " << name() << " because a Stage cannot stream its own self-loads.\n"; - names->push_back(target_name); + names.push_back(target_name); } + definition.schedule().stream_loads_names() = std::move(names); definition.schedule().touched() = true; return *this; } diff --git a/test/correctness/streaming_in.cpp b/test/correctness/streaming_in.cpp index 931259a73556..039fc1797f6d 100644 --- a/test/correctness/streaming_in.cpp +++ b/test/correctness/streaming_in.cpp @@ -61,7 +61,6 @@ struct WarningCounter final : CompileTimeErrorReporter { int warnings_occurred = 0; void warning(const char *msg) override { - std::fprintf(stderr, "Warning: %s\n", msg); warnings_occurred++; }