iree-org / iree-llvm-sandbox

A sandbox for quick iteration and experimentation on projects related to IREE, MLIR, and LLVM
Apache License 2.0
54 stars 31 forks source link

[Indexing] Add runtime infra and runnable, tiled matmul example #744

Closed makslevental closed 8 months ago

makslevental commented 1 year 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

makslevental commented 1 year ago

TODO: doc strings