Skip to content

Commit a76d5b4

Browse files
pifon2aGoogle-ML-Automation
pifon2a
authored and
Google-ML-Automation
committed
[XLA][IndexAnalysis] Move indexing_analysis to hlo/analysis.
PiperOrigin-RevId: 702221388
1 parent b5604d9 commit a76d5b4

File tree

77 files changed

+286
-306
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

77 files changed

+286
-306
lines changed

‎docs/indexing.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -907,7 +907,7 @@ d2 in [0, 124]
907907

908908
where `s0` refers to the inner-most dimension of the input.
909909

910-
For more examples see [indexing_analysis_test.cc](https://github.com/openxla/xla/blob/main/xla/service/gpu/model/indexing_analysis_test.cc).
910+
For more examples see [indexing_analysis_test.cc](https://github.com/openxla/xla/blob/main/xla/hlo/analysis/indexing_analysis_test.cc).
911911

912912
## Indexing Map Simplifier
913913

@@ -955,4 +955,4 @@ for `d0 in [0, 5]` and `s0 in [1, 3]` are eliminated.
955955
map above.
956956

957957

958-
For more examples see [indexing_map_test.cc](https://github.com/openxla/xla/blob/main/xla/service/gpu/model/indexing_map_test.cc).
958+
For more examples see [indexing_map_test.cc](https://github.com/openxla/xla/blob/main/xla/hlo/analysis/indexing_map_test.cc).

‎xla/hlo/analysis/BUILD‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -514,3 +514,121 @@ xla_cc_test(
514514
"@tsl//tsl/platform:test_main",
515515
],
516516
)
517+
518+
cc_library(
519+
name = "indexing_analysis",
520+
srcs = [
521+
"indexing_analysis.cc",
522+
"indexing_map.cc",
523+
"indexing_map_serialization.cc",
524+
],
525+
hdrs = [
526+
"indexing_analysis.h",
527+
"indexing_map.h",
528+
"indexing_map_serialization.h",
529+
],
530+
deps = [
531+
"//xla:permutation_util",
532+
"//xla:shape_util",
533+
"//xla:util",
534+
"//xla:xla_data_proto_cc",
535+
"//xla/hlo/ir:hlo",
536+
"//xla/hlo/transforms:gather_simplifier",
537+
"//xla/hlo/utils:hlo_traversal",
538+
"//xla/service/gpu:matmul_indexing_utils",
539+
"@com_google_absl//absl/algorithm:container",
540+
"@com_google_absl//absl/container:flat_hash_map",
541+
"@com_google_absl//absl/container:flat_hash_set",
542+
"@com_google_absl//absl/container:inlined_vector",
543+
"@com_google_absl//absl/log:check",
544+
"@com_google_absl//absl/numeric:int128",
545+
"@com_google_absl//absl/strings",
546+
"@com_google_absl//absl/strings:str_format",
547+
"@com_google_absl//absl/types:span",
548+
"@llvm-project//llvm:Support",
549+
"@llvm-project//mlir:AsmParser",
550+
"@llvm-project//mlir:IR",
551+
"@llvm-project//mlir:Support",
552+
"@tsl//tsl/platform:logging",
553+
],
554+
)
555+
556+
xla_cc_test(
557+
name = "indexing_map_test",
558+
srcs = ["indexing_map_test.cc"],
559+
deps = [
560+
":indexing_analysis",
561+
":indexing_test_utils",
562+
"//xla/hlo/testlib:verified_hlo_module",
563+
"//xla/tests:hlo_test_base",
564+
"//xla/tests:verified_hlo_module",
565+
"//xla/tests:xla_internal_test_main",
566+
"@com_google_absl//absl/hash:hash_testing",
567+
"@com_google_absl//absl/strings:string_view",
568+
"@com_google_absl//absl/types:span",
569+
"@com_google_googletest//:gtest",
570+
"@llvm-project//mlir:IR",
571+
"@tsl//tsl/platform:statusor",
572+
"@tsl//tsl/platform:test",
573+
],
574+
)
575+
576+
xla_cc_test(
577+
name = "indexing_map_serialization_test",
578+
srcs = ["indexing_map_serialization_test.cc"],
579+
deps = [
580+
":indexing_analysis",
581+
":indexing_test_utils",
582+
"//xla/tests:hlo_test_base",
583+
"//xla/tests:xla_internal_test_main",
584+
"@com_google_absl//absl/strings:string_view",
585+
"@com_google_googletest//:gtest",
586+
"@llvm-project//mlir:IR",
587+
"@tsl//tsl/platform:test",
588+
],
589+
)
590+
591+
cc_library(
592+
name = "indexing_test_utils",
593+
testonly = True,
594+
srcs = ["indexing_test_utils.cc"],
595+
hdrs = ["indexing_test_utils.h"],
596+
deps = [
597+
":indexing_analysis",
598+
"//xla:status_macros",
599+
"//xla/hlo/ir:hlo",
600+
"//xla/hlo/testlib:verified_hlo_module",
601+
"//xla/tests:hlo_test_base",
602+
"//xla/tests:verified_hlo_module",
603+
"@com_google_absl//absl/container:flat_hash_map",
604+
"@com_google_absl//absl/container:flat_hash_set",
605+
"@com_google_absl//absl/container:inlined_vector",
606+
"@com_google_absl//absl/log",
607+
"@com_google_absl//absl/log:check",
608+
"@com_google_absl//absl/status",
609+
"@com_google_absl//absl/strings",
610+
"@com_google_absl//absl/strings:string_view",
611+
"@com_google_absl//absl/types:span",
612+
"@com_google_googletest//:gtest",
613+
"@llvm-project//llvm:Support",
614+
"@llvm-project//mlir:AsmParser",
615+
"@llvm-project//mlir:IR",
616+
"@llvm-project//mlir:Support",
617+
"@tsl//tsl/platform:errors",
618+
],
619+
)
620+
621+
xla_cc_test(
622+
name = "indexing_analysis_test",
623+
srcs = ["indexing_analysis_test.cc"],
624+
deps = [
625+
":indexing_analysis",
626+
":indexing_test_utils",
627+
"//xla/hlo/ir:hlo",
628+
"//xla/hlo/utils:hlo_traversal",
629+
"//xla/tests:xla_internal_test_main",
630+
"@com_google_absl//absl/strings:string_view",
631+
"@com_google_googletest//:gtest",
632+
"@tsl//tsl/platform:test",
633+
],
634+
)

xla/service/gpu/model/indexing_analysis.cc renamed to xla/hlo/analysis/indexing_analysis.cc

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515

16-
#include "xla/service/gpu/model/indexing_analysis.h"
16+
#include "xla/hlo/analysis/indexing_analysis.h"
1717

1818
#include <algorithm>
1919
#include <cassert>
@@ -40,6 +40,7 @@ limitations under the License.
4040
#include "mlir/IR/AffineMap.h"
4141
#include "mlir/IR/MLIRContext.h"
4242
#include "mlir/Support/LLVM.h"
43+
#include "xla/hlo/analysis/indexing_map.h"
4344
#include "xla/hlo/ir/hlo_casting_utils.h"
4445
#include "xla/hlo/ir/hlo_instruction.h"
4546
#include "xla/hlo/ir/hlo_instructions.h"
@@ -49,14 +50,12 @@ limitations under the License.
4950
#include "xla/layout.h"
5051
#include "xla/permutation_util.h"
5152
#include "xla/service/gpu/matmul_indexing_utils.h"
52-
#include "xla/service/gpu/model/indexing_map.h"
5353
#include "xla/shape.h"
5454
#include "xla/shape_util.h"
5555
#include "xla/util.h"
5656
#include "xla/xla_data.pb.h"
5757

5858
namespace xla {
59-
namespace gpu {
6059
namespace {
6160

6261
using llvm::SmallVector;
@@ -450,8 +449,8 @@ HloInstructionIndexing ComputeOutputToInputDotOpIndexing(
450449
}
451450

452451
// lhs_non_contracting_dims
453-
auto lhs_non_contracting_dims =
454-
GetNonContractingDims(lhs_shape, lhs_batch_dims, lhs_contracting_dims);
452+
auto lhs_non_contracting_dims = gpu::GetNonContractingDims(
453+
lhs_shape, lhs_batch_dims, lhs_contracting_dims);
455454
assert(lhs_non_contracting_dims.ok());
456455

457456
for (int64_t lhs_non_contracting_dim : lhs_non_contracting_dims.value()) {
@@ -460,8 +459,8 @@ HloInstructionIndexing ComputeOutputToInputDotOpIndexing(
460459
}
461460

462461
// rhs_non_contracting_dims
463-
auto rhs_non_contracting_dims =
464-
GetNonContractingDims(rhs_shape, rhs_batch_dims, rhs_contracting_dims);
462+
auto rhs_non_contracting_dims = gpu::GetNonContractingDims(
463+
rhs_shape, rhs_batch_dims, rhs_contracting_dims);
465464
assert(rhs_non_contracting_dims.ok());
466465
for (int64_t rhs_non_contracting_dim : rhs_non_contracting_dims.value()) {
467466
rhs_exprs[rhs_non_contracting_dim] =
@@ -1703,5 +1702,4 @@ IndexingMap ComputeEpilogueInputToOutputIndexing(
17031702
return root_indexing;
17041703
}
17051704

1706-
} // namespace gpu
17071705
} // namespace xla

xla/service/gpu/model/indexing_analysis.h renamed to xla/hlo/analysis/indexing_analysis.h

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,8 @@ See the License for the specific language governing permissions and
1414
limitations under the License.
1515
==============================================================================*/
1616

17-
#ifndef XLA_SERVICE_GPU_MODEL_INDEXING_ANALYSIS_H_
18-
#define XLA_SERVICE_GPU_MODEL_INDEXING_ANALYSIS_H_
17+
#ifndef XLA_HLO_ANALYSIS_INDEXING_ANALYSIS_H_
18+
#define XLA_HLO_ANALYSIS_INDEXING_ANALYSIS_H_
1919

2020
#include <cstdint>
2121
#include <ostream>
@@ -29,13 +29,12 @@ limitations under the License.
2929
#include "mlir/IR/AffineExpr.h"
3030
#include "mlir/IR/AffineMap.h"
3131
#include "mlir/IR/MLIRContext.h"
32+
#include "xla/hlo/analysis/indexing_map.h"
3233
#include "xla/hlo/ir/hlo_instruction.h"
3334
#include "xla/hlo/utils/hlo_traversal.h"
34-
#include "xla/service/gpu/model/indexing_map.h"
3535
#include "xla/shape.h"
3636

3737
namespace xla {
38-
namespace gpu {
3938

4039
using IndexingMapSet = absl::flat_hash_set<IndexingMap>;
4140

@@ -166,7 +165,6 @@ IndexingMap CreateIdentityMap(absl::Span<const int64_t> dimensions,
166165
llvm::SmallVector<mlir::AffineExpr, 4> DelinearizeInBoundsIndex(
167166
mlir::AffineExpr linear, absl::Span<const int64_t> sizes);
168167

169-
} // namespace gpu
170168
} // namespace xla
171169

172-
#endif // XLA_SERVICE_GPU_MODEL_INDEXING_ANALYSIS_H_
170+
#endif // XLA_HLO_ANALYSIS_INDEXING_ANALYSIS_H_

xla/service/gpu/model/indexing_analysis_test.cc renamed to xla/hlo/analysis/indexing_analysis_test.cc

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,15 +13,15 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515

16-
#include "xla/service/gpu/model/indexing_analysis.h"
16+
#include "xla/hlo/analysis/indexing_analysis.h"
1717

1818
#include <gmock/gmock.h>
1919
#include <gtest/gtest.h>
2020
#include "absl/strings/string_view.h"
21+
#include "xla/hlo/analysis/indexing_map_serialization.h"
22+
#include "xla/hlo/analysis/indexing_test_utils.h"
2123
#include "xla/hlo/ir/hlo_instruction.h"
2224
#include "xla/hlo/utils/hlo_traversal.h"
23-
#include "xla/service/gpu/model/indexing_map_serialization.h"
24-
#include "xla/service/gpu/model/indexing_test_utils.h"
2525
#include "tsl/platform/test.h"
2626

2727
namespace xla {

xla/service/gpu/model/indexing_map.cc renamed to xla/hlo/analysis/indexing_map.cc

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515

16-
#include "xla/service/gpu/model/indexing_map.h"
16+
#include "xla/hlo/analysis/indexing_map.h"
1717

1818
#include <algorithm>
1919
#include <cassert>
@@ -48,7 +48,6 @@ limitations under the License.
4848
#include "tsl/platform/logging.h" // IWYU pragma: keep
4949

5050
namespace xla {
51-
namespace gpu {
5251
namespace {
5352

5453
using llvm::ArrayRef;
@@ -1940,5 +1939,4 @@ IndexingMap IndexingMap::ConvertSymbolsToDimensions() const {
19401939
return new_indexing_map;
19411940
}
19421941

1943-
} // namespace gpu
19441942
} // namespace xla

xla/service/gpu/model/indexing_map.h renamed to xla/hlo/analysis/indexing_map.h

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515

16-
#ifndef XLA_SERVICE_GPU_MODEL_INDEXING_MAP_H_
17-
#define XLA_SERVICE_GPU_MODEL_INDEXING_MAP_H_
16+
#ifndef XLA_HLO_ANALYSIS_INDEXING_MAP_H_
17+
#define XLA_HLO_ANALYSIS_INDEXING_MAP_H_
1818

1919
#include <algorithm>
2020
#include <cstddef>
@@ -38,7 +38,6 @@ limitations under the License.
3838
#include "xla/hlo/ir/hlo_instruction.h"
3939

4040
namespace xla {
41-
namespace gpu {
4241

4342
enum class VariableKind : char {
4443
kDefault = 0,
@@ -484,7 +483,6 @@ std::vector<IndexingMap::Variable> DimVarsFromGPUGrid(
484483
std::vector<IndexingMap::Variable> RangeVarsFromTensorSizes(
485484
absl::Span<const int64_t> tensor_sizes);
486485

487-
} // namespace gpu
488486
} // namespace xla
489487

490-
#endif // XLA_SERVICE_GPU_MODEL_INDEXING_MAP_H_
488+
#endif // XLA_HLO_ANALYSIS_INDEXING_MAP_H_

xla/service/gpu/model/indexing_map_serialization.cc renamed to xla/hlo/analysis/indexing_map_serialization.cc

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515

16-
#include "xla/service/gpu/model/indexing_map_serialization.h"
16+
#include "xla/hlo/analysis/indexing_map_serialization.h"
1717

1818
#include <algorithm>
1919
#include <cctype>
@@ -41,10 +41,9 @@ limitations under the License.
4141
#include "mlir/IR/BuiltinAttributes.h"
4242
#include "mlir/IR/MLIRContext.h"
4343
#include "mlir/Support/LLVM.h"
44-
#include "xla/service/gpu/model/indexing_map.h"
44+
#include "xla/hlo/analysis/indexing_map.h"
4545

4646
namespace xla {
47-
namespace gpu {
4847
namespace {
4948

5049
using llvm::SmallVector;
@@ -933,5 +932,4 @@ SmallVector<std::string> GetSymbolVarNames(const IndexingMap& map) {
933932
return symbol_names;
934933
}
935934

936-
} // namespace gpu
937935
} // namespace xla

xla/service/gpu/model/indexing_map_serialization.h renamed to xla/hlo/analysis/indexing_map_serialization.h

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515

16-
#ifndef XLA_SERVICE_GPU_MODEL_INDEXING_MAP_SERIALIZATION_H_
17-
#define XLA_SERVICE_GPU_MODEL_INDEXING_MAP_SERIALIZATION_H_
16+
#ifndef XLA_HLO_ANALYSIS_INDEXING_MAP_SERIALIZATION_H_
17+
#define XLA_HLO_ANALYSIS_INDEXING_MAP_SERIALIZATION_H_
1818

1919
#include <optional>
2020
#include <ostream>
@@ -26,10 +26,9 @@ limitations under the License.
2626
#include "mlir/IR/AffineExpr.h"
2727
#include "mlir/IR/AffineMap.h"
2828
#include "mlir/IR/MLIRContext.h"
29-
#include "xla/service/gpu/model/indexing_map.h"
29+
#include "xla/hlo/analysis/indexing_map.h"
3030

3131
namespace xla {
32-
namespace gpu {
3332

3433
// Parses the given string into an IndexingMap.
3534
std::optional<IndexingMap> ParseIndexingMap(llvm::StringRef input,
@@ -79,7 +78,6 @@ llvm::SmallVector<std::string> GetRTVarNames(const IndexingMap& map);
7978
// Symbol variable names: concatenation of range and runtime variables.
8079
llvm::SmallVector<std::string> GetSymbolVarNames(const IndexingMap& map);
8180

82-
} // namespace gpu
8381
} // namespace xla
8482

85-
#endif // XLA_SERVICE_GPU_MODEL_INDEXING_MAP_SERIALIZATION_H_
83+
#endif // XLA_HLO_ANALYSIS_INDEXING_MAP_SERIALIZATION_H_

xla/service/gpu/model/indexing_map_serialization_test.cc renamed to xla/hlo/analysis/indexing_map_serialization_test.cc

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,20 +13,19 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515

16-
#include "xla/service/gpu/model/indexing_map_serialization.h"
16+
#include "xla/hlo/analysis/indexing_map_serialization.h"
1717

1818
#include <gmock/gmock.h>
1919
#include <gtest/gtest.h>
2020
#include "absl/strings/string_view.h"
2121
#include "mlir/IR/AffineExpr.h"
2222
#include "mlir/IR/AffineMap.h"
2323
#include "mlir/IR/MLIRContext.h"
24-
#include "xla/service/gpu/model/indexing_test_utils.h"
24+
#include "xla/hlo/analysis/indexing_test_utils.h"
2525
#include "xla/tests/hlo_test_base.h"
2626
#include "tsl/platform/test.h"
2727

2828
namespace xla {
29-
namespace gpu {
3029
namespace {
3130

3231
using ::testing::HasSubstr;
@@ -188,6 +187,6 @@ TEST_F(IndexingMapSerializationTest, AffineMapPrinterTest) {
188187
HasSubstr("(d0, d1)[s0, s1, s2, s3] -> "
189188
"(d0 + d1 floordiv 8 - s2 * 64, s0 + s1 mod 16 + s3)"));
190189
}
190+
191191
} // namespace
192-
} // namespace gpu
193192
} // namespace xla

0 commit comments

Comments
 (0)