llvm / llvm/torch-mlir

[FXImporter] Find a better way to get a buffer representation of a torch.Tensor (support bfloat and complex dtypes when faced with lift_fresh_copy)

Open
#3,653 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
C++
Stars
1.9k
Forks
736
Avg merge
5d 22h
Merged PRs (30d)
15

Description

This issue is being moved from SHARK-Turbine to torch-mlir as the fx importer lives in torch-mlir now.

Currently, we address issues stemming from lift_fresh_copy by creating a tensor literal op (https://github.com/nod-ai/SHARK-Turbine/pull/37), but this is problematic because in order to do so we need a buffer representation of a torch.Tensor. Unfortunately, torch.Tensor does not implement the python array interface fully which precludes us from directly grabbing the representation of the tensor in memory, rather we are forced to use an indirect route through numpy to get a python buffer that can be parsed by MLIR into a tensor literal. This has the unfortunate side effect that we can not support bfloat and complex<> datatypes with this operation because 1) numpy has no bfloat datatype and hence no representation for such a buffer and 2) numpy's buffer format for complex<> datatypes is incompatible with the buffer format that MLIR's DenseElementsAttr expects.

The best solution would be to have a first-class mechanism for getting a memoryview of a torch.Tensor by implementing the python array interface fully for this class. This is an issue tracking this shortcoming in pytorch: https://github.com/pytorch/pytorch/issues/54138

Tracking the implementation of this interface: https://github.com/pytorch/pytorch/issues/58743
Actually the immediately relevant interface is the python buffer interface: https://github.com/pytorch/pytorch/issues/19143

The above issue refers to the following source in torch-mlir:
https://github.com/llvm/torch-mlir/blob/af67f9efb079412d563113f759929876129588d2/python/torch_mlir/extras/fx_importer.py#L195

https://github.com/llvm/torch-mlir/blob/af67f9efb079412d563113f759929876129588d2/python/torch_mlir/extras/fx_importer.py#L2110

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the referenced locations in python/torch_mlir/extras/fx_importer.py, around lines 195 and 2110, and review the linked PyTorch Python array and buffer interface issues. Done means lift_fresh_copy handling can obtain a tensor buffer representation that supports bfloat and complex dtypes without relying on an incompatible NumPy representation.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
compilers
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.