mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-10 07:28:15 -05:00
49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
#!/usr/bin/env python
|
|
import unittest
|
|
import numpy as np
|
|
from tinygrad.ops import LazyOp, BinaryOps, ReduceOps
|
|
from tinygrad.runtime.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()
|