Python-binding: Add method to convert from Func to pure stage? #6702
Description
In the C++ Func API there is an operator to convert from Func to its pure stage:
Halide::Func::operator Stage ( ) const
You can cast a Func to its pure stage for the purposes of scheduling it.
I could not find the same conversion in the Python wrapper. I was looking for it because I wanted to use compute_with, in order to calculate 3 functions in a Pipeline in the same loop. I was able to work around the issue by making a LoopLevel object.
Reproducer:
from halide import Pipeline, Func, Var, LoopLevel
def gpu_tile_in_place(func):
TILE_SIZE = 16
if func.dimensions() == 3:
func.reorder(img, col, row)
func.compute_root()
func.gpu_tile(col, row, col_outer, row_outer, col_inner, row_inner, TILE_SIZE, TILE_SIZE)
col = Var("col")
row = Var("row")
img = Var("img")
col_inner = Var("col_inner")
col_outer = Var("col_outer")
row_inner = Var("row_inner")
row_outer = Var("row_outer")
output_A = Func("output_A")
output_A[col, row, img] = img * 2
output_B = Func("output_B")
output_B[col, row, img] = img + 1
output_C = Func("output_C")
output_C[col, row] = (col + row) / 2
gpu_tile_in_place(output_A)
gpu_tile_in_place(output_B)
gpu_tile_in_place(output_C)
output_pipeline = Pipeline([output_A, output_B, output_C])
# These don't work:
#output_B.compute_with(output_A, img)
#output_C.compute_with(output_A, col_inner)
# They give the following error:
"""
TypeError: compute_with(): incompatible function arguments. The following argument types are supported:
1. (self: halide.Func, loop_level: halide.LoopLevel, align: List[Tuple[halide.VarOrRVar, halide.LoopAlignStrategy]]) -> halide.Func
2. (self: halide.Func, loop_level: halide.LoopLevel, align: halide.LoopAlignStrategy = <LoopAlignStrategy.Auto: 3>) -> halide.Func
3. (self: halide.Func, stage: Halide::Stage, var: halide.VarOrRVar, align: List[Tuple[halide.VarOrRVar, halide.LoopAlignStrategy]]) -> halide.Func
4. (self: halide.Func, stage: Halide::Stage, var: halide.VarOrRVar, align: halide.LoopAlignStrategy = <LoopAlignStrategy.Auto: 3>) -> halide.Func
Invoked with: <halide.Func 'output_B'>, <halide.Func 'output_A'>, <halide.Var 'img'>
"""
# These however work fine and make the expected loop nest:
output_B.compute_with(LoopLevel(output_A, img, stage_index=0))
output_C.compute_with(LoopLevel(output_A, col_inner, stage_index=0))
output_pipeline.print_loop_nest()
Output:
Injecting realization of { output_B, output_C, output_A }
produce output_A:
produce output_C:
produce output_B:
gpu_block row.fused.row_outer<Default_GPU>:
gpu_block col.fused.col_outer<Default_GPU>:
gpu_thread row.fused.row_inner in [0, 15]<Default_GPU>:
gpu_thread col.fused.col_inner in [0, 15]<Default_GPU>:
for fused.img:
output_A(...) = ...
output_B(...) = ...
output_C(...) = ...
Metadata
Metadata
Assignees
Labels
No labels