From b44ab8ad49195f3a36b7c0fde6faeb1c42e9de65 Mon Sep 17 00:00:00 2001 From: Hanhan Wang Date: Tue, 15 Jun 2021 10:57:15 -0700 Subject: [PATCH] Add support for lowering DataMovementOp ops to Linalg on unsigned types. PiperOrigin-RevId: 379527360 --- .../mhlo/transforms/legalize_to_linalg.cc | 2 ++ tests/hlo-legalize-to-linalg.mlir | 18 ++++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc b/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc index 71a5d25..2478ec5 100644 --- a/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc +++ b/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc @@ -406,6 +406,8 @@ class DataMovementOpConverter : public OpConversionPattern { ConversionPatternRewriter& rewriter) const final { if (!VerifyHloOpBufferOrTensorSemantics(op)) return failure(); auto result_type = GetHloOpResultType(op); + result_type = this->typeConverter->convertType(result_type) + .template cast(); SmallVector indexing_maps = Derived::getIndexingMaps(op, &rewriter); diff --git a/tests/hlo-legalize-to-linalg.mlir b/tests/hlo-legalize-to-linalg.mlir index fd528cc..409a648 100644 --- a/tests/hlo-legalize-to-linalg.mlir +++ b/tests/hlo-legalize-to-linalg.mlir @@ -493,6 +493,24 @@ func @broadcast_in_dim(%operand: tensor<5x7x1xf32>) -> tensor<7x10x6x4x5xf32> { // ----- +// CHECK-DAG: #[[OPERAND_MAP:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d4, d0, 0)> +// CHECK-DAG: #[[RESULT_MAP:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)> +// CHECK-LABEL: func @broadcast_in_dim_ui32 +func @broadcast_in_dim_ui32(%operand: tensor<5x7x1xui32>) -> tensor<7x10x6x4x5xui32> { + %0 = "mhlo.broadcast_in_dim"(%operand) + {broadcast_dimensions = dense<[4,0,2]> : tensor<3xi64>} + : (tensor<5x7x1xui32>) -> tensor<7x10x6x4x5xui32> + return %0 : tensor<7x10x6x4x5xui32> +} +// CHECK: unrealized_conversion_cast %{{.*}} : tensor<5x7x1xui32> to tensor<5x7x1xi32> +// CHECK: linalg.init_tensor [7, 10, 6, 4, 5] : tensor<7x10x6x4x5xi32> +// CHECK: %[[RES:.*]] = linalg.generic {{{.*}}indexing_maps = [#[[OPERAND_MAP]], #[[RESULT_MAP]]] +// CHECK-NEXT: ^bb0(%[[OPERAND:.*]]: i32, %{{.*}}: i32): +// CHECK-NEXT: linalg.yield %[[OPERAND]] : i32 +// CHECK: unrealized_conversion_cast %[[RES]] : tensor<7x10x6x4x5xi32> to tensor<7x10x6x4x5xui32> + +// ----- + // CHECK-DAG: #[[OPERAND_MAP:.+]] = affine_map<(d0, d1) -> (d0)> // CHECK-DAG: #[[RESULT_MAP:.+]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-LABEL: func @broadcast_in_dim_with_one_to_one