experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity
STT-tensorflow/tensorflow/compiler/xla/service/spmd
History
Marcello Maggioni 71fd1ab1d5 [XLA] Make the canonicalize all-gather pass add a new/exclusive all-gather dimension. 
Make the canonicalizer always add a completely new all gather dimension to the shape
of the all-gather operand to help layout assignment.

PiperOrigin-RevId: 360612458
Change-Id: I8299301219f704f14661f6fb2619ef604bb4d7c4
2021-03-03 01:06:58 -08:00
..
BUILD
[XLA] Add option to SPMD partitioner to prefer choosing faster windowed einsums than saving memory.
2021-02-16 21:21:33 -08:00
canonicalize_all_gather_for_cse_test.cc
[XLA] Make the canonicalize all-gather pass add a new/exclusive all-gather dimension.
2021-03-03 01:06:58 -08:00
canonicalize_all_gather_for_cse.cc
[XLA] Make the canonicalize all-gather pass add a new/exclusive all-gather dimension.
2021-03-03 01:06:58 -08:00
canonicalize_all_gather_for_cse.h
[XLA] Extend AllGather CSE to look through reshapes that add degenerate dims.
2021-02-10 18:08:39 -08:00
convolution_handler.cc
[XLA] Add an optional preferred_element_type argument to Dot/Conv builder methods to enable generating HLOs that has wider accumulation type than default shape inference result.
2020-11-20 13:16:56 -08:00
convolution_handler.h
[XLA:SPMD] Support partition convlution input in both spatial dimension and
2020-09-18 23:54:09 -07:00
dot_handler.cc
[XLA] Correctly propagate builder and partition_id for inner_dot_builder.
2021-02-26 17:23:33 -08:00
fft_handler.cc
[XLA:SPMD] Add basic support for SPMD FFT.
2020-10-21 11:19:54 -07:00
gather_scatter_handler.cc
[XLA] Call SetPartitionedHlo for index only gather partitioning.
2021-02-14 20:15:16 -08:00
schedule_aware_all_gather_cse_test.cc
[XLA] Make the canonicalize all-gather pass add a new/exclusive all-gather dimension.
2021-03-03 01:06:58 -08:00
schedule_aware_all_gather_cse.cc
[XLA] Make the canonicalize all-gather pass add a new/exclusive all-gather dimension.
2021-03-03 01:06:58 -08:00
schedule_aware_all_gather_cse.h
[XLA:SPMD] Schedule-aware all-gather CSE for SPMD
2020-08-19 15:58:14 -07:00
spmd_partitioner_test.cc
[XLA:SPMD] Move reduce-scatter creation into creator callback
2021-02-26 15:40:14 -08:00
spmd_partitioner_util.cc
[XLA] Partition with padded shape when creating Group partition shape.
2021-02-23 11:55:12 -08:00
spmd_partitioner_util.h
[XLA] Add sharding propagation support for Gather sharding on parallel dimensions.
2021-01-16 01:56:57 -08:00
spmd_partitioner.cc
[XLA:SPMD] Detect shardable reshape in the other direction
2021-02-26 14:11:23 -08:00
spmd_partitioner.h
[XLA] Add option to SPMD partitioner to prefer choosing faster windowed einsums than saving memory.
2021-02-16 21:21:33 -08:00