Refactor layout implementation #491
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/491
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit a8d9218 with merge base 6e7cf71 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
jerryzh168
force-pushed
the
layout
branch
2 times, most recently
from
3ffe60f to
d181a77
Compare
jerryzh168
requested review from
msaroufim,
drisspg,
vayuda,
gau-nernst,
andrewor14 and
HDCharles
| ): | ||
| original_shape = input_float.shape | ||
| if extended_layout == "tensor_core_tiled": | ||
| if isinstance(layout_type, TensorCoreTiledLayoutType): |
There was a problem hiding this comment.
is it possible to move this inside the TensorCoreTiledLayout constructor
There was a problem hiding this comment.
oh, this is padding the input though, not sure how we can do that, but we can move the implementation to TensorCoreTiledLayout I think
There was a problem hiding this comment.
Yea I think that would be a bit cleaner if all layout specific code was kept within the layout class. Otherwise, this pr looks good.
|
|
||
| aten = torch.ops.aten | ||
|
|
||
| @dataclass(frozen=True) | ||
| class PlainLayoutType(LayoutType): |
There was a problem hiding this comment.
comment or error that this shouldnt be instantiated directly
There was a problem hiding this comment.
this can be instantiated I think, are you talking about LayoutType?
There was a problem hiding this comment.
I see I guess I'm a bit thrown off because a data classes primary goal is to store data wheras this class stores nothing and its really just a name
There was a problem hiding this comment.
I would have instead done an enum like this
from enum import Enum
class Operations(Enum):
ADD = (1,)
SUBTRACT = (2,)
MULTIPLY = (3,)
DIVIDE = (4, 'precision')
enums are also a class so you can override __init__ and define a func that only applies on DIVIDE for example
There was a problem hiding this comment.
sorry I don't follow the DIVIDE part, can you elaborate a bit? is this talking about how to support TensorCoreTiledLayoutType that has a inner_k_tiles argument?
|
|
||
| def pad_input(self, input: torch.Tensor) -> torch.Tensor: | ||
| orig_out_features, orig_in_features = input.shape | ||
| in_features = find_multiple(orig_in_features, 1024) |
There was a problem hiding this comment.
where are the in and out numbers coming from? I constants like this were a function of the dtype as well
There was a problem hiding this comment.
this is comes from tinygemm kernel I think, this layout only applies to uint4 dtype
|
|
||
| @dataclass(frozen=True) | ||
| class TensorCoreTiledLayoutType(LayoutType): | ||
| inner_k_tiles: int = 8 |
There was a problem hiding this comment.
@msaroufim see here, we have extra configurable arguments, it's not just a name so I'm not sure how enum would work here
| input_float, | ||
| (0, in_features - orig_in_features, 0, out_features - orig_out_features), | ||
| ) | ||
| input_float = layout_type.pad_input(input_float) |
There was a problem hiding this comment.
This is on the right track but I think we can make it more generic to better serve future usage.
LayoutTypes can implement two functions: pre_process() and post_process() The new workflow would look like this:
input_float = layout.pre_process(input_float)
...
int_data = quantize_affine(...)
int_data = layout.post_process(int_data)
layout_tensor_ctr = get_layout_tensor_constructor(type(layout_type))
layout_tensor = layout_tensor_ctr(int_data, scale, zero_point, layout_type)
To motivate this I can say that I would for sure use the post_process function while integrating my intx work. This allows me to write just a layout_type instead of an entire layout_class (I can reuse the plain layout). Additionally, if you look at #498 they call torch._cslt_compress(int_data) which would also be implemented as a post_process function. I think this also means they can re-use the plain layout since i dont see them use any other arguments in the constructor (@jcaip maybe you can confirm this)
There was a problem hiding this comment.
makes sense, by reusing PlainLayoutType, you mean inherit it and override post_process right?
Summary: TODO Test Plan: TODO Reviewers: Subscribers: Tasks: Tags:
Summary: TODO Test Plan: TODO Reviewers: Subscribers: Tasks: Tags:
* Minimal android app build * Improve script * Detect physical device as well
Summary:
Cleaning up some layout type related arguments, previously we are putting layout specific args like
inner_k_tilesinfrom_float, but now we just uselayout_typewhich can be differentLayoutTypeinstances and can hold different args for each type of LayoutTypeTest Plan:
regression tests:
Reviewers:
Subscribers:
Tasks:
Tags: