Skip to content

Add stream_loads and stream_stores scheduling directives #9207

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
alexreinking wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Add stream_loads and stream_stores scheduling directives #9207

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
057fec5
Add stream_loads and stream_stores scheduling directives
alexreinking Jul 14, 2026
8cee1d8
Add missing visitor overloads for StreamingStore and StreamingLoads
alexreinking Jul 16, 2026
21fe7d6
Respond to review comments
alexreinking Jul 16, 2026
a9baf07
Make stream_loads override rather than accumulate
alexreinking Jul 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions python_bindings/src/halide/halide_/PyFunc.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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> &)) & Func::stream_loads, py::arg("funcs"))
.def("stream_stores", &Func::stream_stores)

.def("compile_to", //
[](Func &f, const std::map<OutputFileType, std::string> &output_files, const std::vector<Argument> &args, const std::string &fn_name, const Target &target) {
Expand Down
5 changes: 5 additions & 0 deletions python_bindings/src/halide/halide_/PyStage.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 &(Stage::*)()>(&Stage::stream_loads))
.def("stream_loads", static_cast<Stage &(Stage::*)(const std::vector<Func> &)>(&Stage::stream_loads),
py::arg("funcs"))

.def("rfactor", static_cast<Func (Stage::*)(const std::vector<std::pair<RVar, Var>> &)>(&Stage::rfactor),
py::arg("preserved"))
.def("rfactor", static_cast<Func (Stage::*)(const RVar &, const Var &)>(&Stage::rfactor),
Expand Down
3 changes: 2 additions & 1 deletion src/AddAtomicMutex.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
2 changes: 1 addition & 1 deletion src/AlignLoads.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
2 changes: 1 addition & 1 deletion src/CSE.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
4 changes: 2 additions & 2 deletions src/CodeGen_ARM.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand Down Expand Up @@ -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);
Expand Down
3 changes: 3 additions & 0 deletions src/CodeGen_C.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)) {
user_warning << "The C backend does not implement non-temporal loads/stores.";
rhs << "0";
Comment thread
alexreinking marked this conversation as resolved.
} else if (op->is_intrinsic(Call::strict_fma)) {
internal_assert(op->args.size() == 3)
<< "Wrong number of args for strict_fma: " << op->args.size();
Expand Down
8 changes: 4 additions & 4 deletions src/CodeGen_D3D12Compute_Dev.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand All @@ -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);
}
Expand Down Expand Up @@ -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);
}
Expand All @@ -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);
}
Expand Down
3 changes: 2 additions & 1 deletion src/CodeGen_GPU_Dev.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
8 changes: 4 additions & 4 deletions src/CodeGen_Hexagon.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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);
}
Expand All @@ -304,14 +304,14 @@ class SloppyUnpredicateLoadsAndStores : public IRMutator {
int lanes = value.type().lanes();

if (const Broadcast *scalar_pred = predicate.as<Broadcast>()) {
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);
}
}
};
Expand Down
Loading
Loading