Closed makslevental closed 8 months ago
Depends on https://github.com/iree-org/iree-llvm-sandbox/pull/744
This PR adds the infra necessary for generating target code (through LLVM JIT/ExecutionEngine) and passing buffers between; e.g. this now passes:
@func.FuncOp.from_py_func(A.type, B.type, C.type) def test_matmul_no_fill(A, B, C): A = Tensor(A) B = Tensor(B) C = Tensor(C) # tile size ts = 5 for i, r1 in scf_range(0, 10, ts, iter_args=[C]): for j, r2 in scf_range(0, 30, ts, iter_args=[C]): for k, r3 in scf_range(0, 20, ts, iter_args=[C]): a = A[i:i + ts, k:k + ts] b = B[k:k + ts, j:j + ts] C[i:i + ts, j:j + ts] += a @ b scf_yield(C) scf_yield(r3) scf_yield(r2) return r1 backend = LLVMJITBackend() module = backend.compile( module, kernel_name="test_matmul_no_fill", pipeline= "builtin.module(...)", ) A = np.random.randint(low=0, high=10, size=(M, K)).astype(np.float32) B = np.random.randint(low=0, high=10, size=(K, N)).astype(np.float32) C = np.zeros((M, N)).astype(np.float32) ... invoker = backend.load(module, consume_return_func=callback) invoker.test_matmul_no_fill(A, B, C) assert result is not None and np.allclose(A @ B, result)
cc @kylechard @ianfoster
TODO: doc strings
Depends on https://github.com/iree-org/iree-llvm-sandbox/pull/744
This PR adds the infra necessary for generating target code (through LLVM JIT/ExecutionEngine) and passing buffers between; e.g. this now passes:
cc @kylechard @ianfoster