Skip to content

Commit 61b6bd9

Browse files
bchetiouiGoogle-ML-Automation
authored andcommitted
[XLA] Throw away TilingSpecification in the TransposedDotTiledHloSchedule.
After relaxing the constraints related to the iteration space in a recent change, this is no longer necessary. PiperOrigin-RevId: 820766539
1 parent 1209728 commit 61b6bd9

File tree

2 files changed

+4
-10
lines changed

2 files changed

+4
-10
lines changed

xla/codegen/tiling/tiled_hlo_schedule.cc

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -191,8 +191,7 @@ TransposedDotTiledHloSchedule::Create(
191191
TF_ASSIGN_OR_RETURN(int64_t n_dim_id, tiling_specification.ParameterIndex(
192192
dot, n_local_parameter_index));
193193

194-
return TransposedDotTiledHloSchedule(tiling_specification, m_dim_id,
195-
n_dim_id);
194+
return TransposedDotTiledHloSchedule(m_dim_id, n_dim_id);
196195
}
197196

198197
absl::StatusOr<IndexingMap> TransposedDotTiledHloSchedule::Schedule(

xla/codegen/tiling/tiled_hlo_schedule.h

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -103,14 +103,9 @@ class TransposedDotTiledHloSchedule : public TiledHloSchedule {
103103
const TilingSpecification& tiling_specification);
104104

105105
private:
106-
TransposedDotTiledHloSchedule(const TilingSpecification& tiling_specification,
107-
int64_t m_dim_id, int64_t n_dim_id)
108-
: tiling_specification_(tiling_specification),
109-
m_dim_id_(m_dim_id),
110-
n_dim_id_(n_dim_id) {}
111-
112-
// The `TilingSpecification` used to construct this schedule.
113-
TilingSpecification tiling_specification_;
106+
TransposedDotTiledHloSchedule(int64_t m_dim_id, int64_t n_dim_id)
107+
: m_dim_id_(m_dim_id), n_dim_id_(n_dim_id) {}
108+
114109
// The index of the `m` dimension within the parameter mapping of the
115110
// `TilingSpecification`.
116111
int64_t m_dim_id_;

0 commit comments

Comments
 (0)