[MOSAIC] apply_vector_layout C++ rewrite (3) applyLayoutOp and relayout

PiperOrigin-RevId: 563556815
google · Sep 7, 2023 · a6eed40 · a6eed40
1 parent 52ee404
commit a6eed40
Show file tree

Hide file tree

Showing 3 changed files with 768 additions and 2 deletions.
diff --git a/jaxlib/mosaic/BUILD b/jaxlib/mosaic/BUILD
@@ -53,6 +53,7 @@ cc_library(
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:ControlFlowDialect",
+        "@llvm-project//mlir:Dialect",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:LinalgTransforms",
@@ -63,10 +64,15 @@ cc_library(
         "@llvm-project//mlir:TransformUtils",
         "@llvm-project//mlir:VectorDialect",
         "@llvm-project//mlir:VectorTransforms",
+        "@xla//xla:array",
         "@xla//xla:shape_util",
+        "@xla//xla:util",
+        "@tsl//tsl/platform:statusor",
         "@com_google_absl//absl/hash",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/types:span",
     ],
 )