Use force_create to allocate return buffer
MODIFY
accel/llvm/ops_llvm.py
accel/llvm/ops_llvm.py
@@ -165,7 +165,7 @@ class LLVMBuffer(ExplicitExecAST):
ReduceOps.MAX: ir.Constant(ir.FloatType(), -math.inf)
}
- def __init__(self, shape:Union[ShapeTracker, Tuple[int, ...]], hostbuf=None):
+ def __init__(self, shape:Union[ShapeTracker, Tuple[int, ...]], hostbuf=None, force_create=False):
super().__init__(shape, hostbuf)
# TODO: force alignment?
self._buf = (ctypes.c_float * (prod(self.shape)))() if hostbuf is None else hostbuf._buf

MODIFY
tinygrad/ast.py
tinygrad/ast.py
@@ -46,8 +46,7 @@ class ASTKernel:
self.ast = ast
# create the buffer we are returning (as the same type as the input buffers) and add it as the first buffer
- self.ret = type(self.bufs[0])(output_shape if output_shape else self.info.shape)
- if hasattr(self.ret, "cl"): self.ret.cl # does the allocation of unbacked buffer, pylint: disable=W0104
+ self.ret = type(self.bufs[0])(output_shape if output_shape else self.info.shape, force_create=True)
self.bufs = [type(self.ret)(self.info.shape, hostbuf=self.ret)] + self.bufs
self.buftokens = [Token(f"data{i}", Types.FLOAT, ptr=True) for i in range(len(self.bufs))]
self.group_for_reduce : List[int] = []

Use force_create to allocate return buffer (#491)