From 4c9f2aca0c8a5fca863c76644377ab75f038a348 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Longeri?= Date: Wed, 22 Nov 2023 03:24:29 -0800 Subject: [PATCH] [Mosaic] C++ apply-vector-layout: don't skip unrecognized operations in `applyLayoutOp` Although the TODO says to return failure, this is actually done at the end of the function (and this way we handle the case for ops without vector args). PiperOrigin-RevId: 584575120 --- jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc | 6 ------ 1 file changed, 6 deletions(-) diff --git a/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc b/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc index 6bcc86c51843..96aad299ad34 100644 --- a/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc +++ b/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc @@ -3664,12 +3664,6 @@ FailureOr relayout(OpBuilder &builder, Value v, VectorLayout src, // For example, we should verify that ops that were supposed to generate // replicated outputs satisfy that requirement. LogicalResult applyLayoutOp(RewriteContext &ctx, Operation &op) { - // TODO(tlongeri): Once we support all ops, return failure instead. - if (!rules().contains(op.getName().getStringRef()) && - !OpTrait::hasElementwiseMappableTraits(&op)) { - return success(); - } - // When an operation does not have any operands, the layout_in tuple is empty. // If one of the operands is not of vector type, the corresponding entry in // the layout_in tuple will be None. The same applies to the results of the