mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-02-19 02:44:40 -05:00
6a15fd3844
* take 2 llvm * get_lazybuffers -> get_buffers * llvm tests pass * fix type issues and refactor LLVM
61 lines
1.6 KiB
Python
61 lines
1.6 KiB
Python
#!/usr/bin/env python
|
|
import unittest
|
|
import numpy as np
|
|
from tinygrad.ops import LazyOp, BinaryOps, ReduceOps
|
|
from accel.llvm.ops_llvm import LLVMBuffer
|
|
|
|
class TestLLVM(unittest.TestCase):
|
|
def test_add(self):
|
|
a = LLVMBuffer.fromCPU(np.ones((4,4)))
|
|
b = LLVMBuffer.fromCPU(np.ones((4,4)))
|
|
ast = LazyOp(BinaryOps.ADD, (a,b))
|
|
ret = LLVMBuffer((4,4)).exec_ast(ast)
|
|
print(ret.toCPU())
|
|
|
|
def test_sum(self):
|
|
a = LLVMBuffer.fromCPU(np.ones((4,4)))
|
|
ast = LazyOp(ReduceOps.SUM, (a,), (1,1))
|
|
ret = LLVMBuffer((1,1)).exec_ast(ast)
|
|
print(ret.toCPU())
|
|
|
|
def test_sum_add(self):
|
|
a = LLVMBuffer.fromCPU(np.ones((4,4)))
|
|
b = LLVMBuffer.fromCPU(np.ones((1,1)))
|
|
ast = LazyOp(ReduceOps.SUM, (a,), (1,1))
|
|
ast = LazyOp(BinaryOps.ADD, (ast,b))
|
|
ret = LLVMBuffer((1,1)).exec_ast(ast)
|
|
print(ret.toCPU())
|
|
|
|
def test_add_sum(self):
|
|
a = LLVMBuffer.fromCPU(np.ones((4,4)))
|
|
b = LLVMBuffer.fromCPU(np.ones((4,4)))
|
|
ast = LazyOp(BinaryOps.ADD, (a,b))
|
|
ast = LazyOp(ReduceOps.SUM, (ast,), (1,1))
|
|
ret = LLVMBuffer((1,1)).exec_ast(ast)
|
|
print(ret.toCPU())
|
|
|
|
def test_add_sum_add(self):
|
|
a = LLVMBuffer.fromCPU(np.ones((4,4)))
|
|
b = LLVMBuffer.fromCPU(np.ones((4,4)))
|
|
c = LLVMBuffer.fromCPU(np.ones((1,1)))
|
|
ast = LazyOp(BinaryOps.ADD, (a,b))
|
|
ast = LazyOp(ReduceOps.SUM, (ast,), (1,1))
|
|
ast = LazyOp(BinaryOps.ADD, (ast,c))
|
|
ret = LLVMBuffer((1,1)).exec_ast(ast)
|
|
print(ret.toCPU())
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|
|
|
|
"""
|
|
ast = LazyOp(BinaryOps.ADD, (a,b))
|
|
ast = LazyOp(ReduceOps.SUM, (ast,), (1,1))
|
|
ret = LLVMBuffer((4,4)).exec_ast(ast)
|
|
print(ret.toCPU())
|
|
"""
|
|
|
|
|
|
|