[Pallas] Set the major-minor layout for inputs and outputs #6826
Conversation
torch_xla/csrc/xla_lower_util.cpp
Outdated
| } | ||
| auto output_shape_impl = MakeTorchTensorLayout( |
There was a problem hiding this comment.
lol can you spell out the xla::Shape directly here? I need to do a search to figure out the return type of MakeTorchTensorLayout . Same for above.
I tried to follow Google's C++ recommendation which states
Type Deduction (including auto)
Use type deduction only if it makes the code clearer to readers who aren't familiar with the project, or if it makes the code safer. Do not use it merely to avoid the inconvenience of writing an explicit type.
| @@ -185,9 +185,7 @@ xla::Shape MakeArrayShapeFromDimensions( | |||
| return MakeShapeWithLayout(type, dimensions, dynamic_dimensions, | |||
| *layout_ptr); | |||
| } | |||
|
|
|||
| bool tpu_layout_env = runtime::sys_util::GetEnvBool("XLA_TPU_LAYOUT", true); | |||
| if (tpu_layout_env && dimensions.size() > 1 && CheckTpuDevice(hw_type)) { | |||
There was a problem hiding this comment.
I don't understand why tpu_layout_env should be deprecated and how is it related to this pr? Are you saying we always use XLA_TPU_LAYOUT ? Does that has a implication on performance on GPU for example?
There was a problem hiding this comment.
well I guess it is default to true anyway.. but still I am not sure why it is deleted in this pr.
There was a problem hiding this comment.
I added that to partially solve this problem. Now I found the completed solution, and therefore the old solution is no longer needed.
|
Thanks, Jack. |
Summary:
Mosaic only accepts major-minor layout for both its inputs and outputs. So we need to enforce those layouts by setting the expected input&output shapes in xla::CustomCallWithLayout. After this change, XLA_TPU_LAYOUT is no longer needed.
Test Plan:
PJRT_DEVICE=TPU python test/test_pallas.py