Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add pytorch interface to ATen Dialect #30

Merged
merged 1 commit into from
Aug 21, 2020

Conversation

stephenneuendorffer
Copy link
Contributor

This patch adds a pytorch interface to npcomp. This interface is modeled
after pytorch_xla and exposes the MLIR-based flow as a virtual device (similar
to a gpu device or an xla backend). Usage is something like:

dev = torch_mlir.mlir_device()
t0 = torch.randn((4,4), device=dev)
t1 = torch.randn((4,4), device=dev)
t2 = t0 + t1
t2_mlir = torch_mlir.get_mlir( t2 )
t2_cpu = t2.to('cpu')

In this case t2_cpu contains the result of the computation, and t2_mlir
contains the mlir description of the computation. Note that this also
properly returns backward paths synthesized by pytorch. There are roughly
speaking three parts of this:

  1. A tensor type (implemented by Tensor.* and TensorImpl.*)
  2. The device modeling (ATenMLIRBridge, ATenMLIRDevice, ATenMLIRType)
  3. a temporary IR (implemented by ir.cpp)
  4. The driver that uses the IR to generate MLIR, run Passes and compile the
    result using mlir::ExecutionEngine (implemented by jit.cpp and
    MLIRGenerator.cpp)
  5. A runtime library implemented by lib/aten_ops.cpp

Particular feedback that would be useful include:

  1. how do we merge the runtime library and jit infrastructure with what is
    being done for numpy?
  2. Suggestions about naming, or how to merge this into the existing directory
    hierarchy. Currently it seems like a bit of a blob with everything depending
    on everything else. There's also some function definitions that don't live in
    the file corresponding to their declaration.
  3. It's unclear to me how much of the 'IR' is necessary, or whether we should
    just create MLIR on the fly.

Some aspects of this are known to be less than optimal, in particular:

  1. Some files don't follow the LLVM naming convention. pragma once is used
    everywhere (pending file renaming)
  2. It's unclear how much of the complexity inherited from pytorch_xla is
    necessary, is something simpler sufficient, or should pytorch_xla be extended,
    or should we keep the complexity and support backend devices like pytorch_xla,
    or will there be a different model about how backend devices are handled.
  3. More aspects of this (e.g. the IR) seem like they should be automatically
    generated.
  4. The tests should get moved around and live with the existing tests.

Copy link
Contributor

@silvasean silvasean left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

First round of comments.

torch_mlir/csrc/ATenMLIRType.h Outdated Show resolved Hide resolved
torch_mlir/csrc/jit.cpp Outdated Show resolved Hide resolved
torch_mlir/csrc/Tensor.cpp Outdated Show resolved Hide resolved
torch_mlir/csrc/MLIRGenerator.h Outdated Show resolved Hide resolved
torch_mlir/lib/aten_ops.cpp Outdated Show resolved Hide resolved
torch_mlir/lib/aten_ops.cpp Outdated Show resolved Hide resolved
torch_mlir/lib/aten_ops.cpp Outdated Show resolved Hide resolved
torch_mlir/csrc/InitPythonBindings.cpp Outdated Show resolved Hide resolved
torch_mlir/test/test_export_ResA.py Outdated Show resolved Hide resolved
torch_mlir/test/test_jit_add_views.py Outdated Show resolved Hide resolved
stellaraccident added a commit that referenced this pull request Aug 18, 2020
* Adds/updates readmes with some notes about code organization and direction.
* Meant to prepare a space for upcoming integration of #30.
stellaraccident added a commit that referenced this pull request Aug 18, 2020
* Adds/updates readmes with some notes about code organization and direction.
* Meant to prepare a space for upcoming integration of #30.
frontends/pytorch/csrc/aten_mlir_bridge.cpp Outdated Show resolved Hide resolved
frontends/pytorch/csrc/aten_mlir_bridge.cpp Outdated Show resolved Hide resolved
frontends/pytorch/csrc/aten_mlir_bridge.cpp Outdated Show resolved Hide resolved
frontends/pytorch/csrc/aten_mlir_bridge.cpp Outdated Show resolved Hide resolved
frontends/pytorch/csrc/aten_mlir_type.cpp Outdated Show resolved Hide resolved
frontends/pytorch/test/test_export_batchnorm.py Outdated Show resolved Hide resolved
python/npcomp/frontends/pytorch/__init__.py Outdated Show resolved Hide resolved
frontends/pytorch/README.md Show resolved Hide resolved
@stephenneuendorffer stephenneuendorffer force-pushed the aten branch 2 times, most recently from 8f4e4ec to 30533c0 Compare August 21, 2020 07:04
This patch adds a pytorch interface to npcomp.  This interface is modeled
after pytorch_xla and exposes the MLIR-based flow as a virtual device (similar
to a gpu device or the xla backend).  Usage is intended to be something like:

  dev = torch_mlir.mlir_device()
  t0 = torch.randn((4,4), device=dev)
  t1 = torch.randn((4,4), device=dev)
  t2 = t0 + t1
  t2_mlir = torch_mlir.get_mlir( t2 )
  t2_cpu = t2.to('cpu')

In this case t2_cpu would contain the result of the computation, and t2_mlir
contains the mlir description of the computation.  Note that this also
properly returns backward paths synthesized by pytorch.  There are several
parts of this:

1) A tensor type (implemented by tensor.* and tensor_impl.*)
2) The device modeling (aten_mlir_bridge.*, aten_mlir_device.*, aten_mlir_type*)
3) a temporary IR (implemented by ir.cpp)

There is also a reference lowering directly from the ATen dialect to C
function calls consisting of two parts:

1) The driver that uses the IR to generate MLIR, run Passes and compile the
result using mlir::ExecutionEngine (implemented by jit.cpp and
mlir_gen.cpp)
2) A runtime library implemented by lib/aten_ops.cpp.  Most of the operations
are implemented by callbacks into the torch C++ libraries.

Some aspects of this are known to be less than optimal, in particular:
1) There's some function definitions that don't live in the file corresponding
to their declaration.
2) More aspects of this (e.g. the IR) seem like they should be automatically
generated.
3) It's unclear to me how much of the 'IR' is actually necessary, or whether
MLIR could be created on the fly.

Note that this code is licensed in a way similar to pytorch, with the
intention that eventually (when npcomp reaches some maturity) it should be
pushed there.  (see frontends/pytorch/LICENSE)  The code is also structured
much closer to the pytorch coding style than the LLVM coding style.
@stephenneuendorffer stephenneuendorffer merged commit 31b3041 into llvm:master Aug 21, 2020
qedawkins pushed a commit to nod-ai/torch-mlir that referenced this pull request Oct 3, 2022
* Import initialized tensor as dense attribute

* Import all initialize tensors as dense constants

* Remove unintentional code

* Fix value attribute format in shape inference tests of reshape

* Readd rank check for reshape's shape inference

* Remove a redundant variable

Co-authored-by: Gheorghe-Teodor Bercea <gt.bercea@gmail.com>
@renxida renxida mentioned this pull request Mar 1, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants