• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1// RUN: mlir-hlo-opt -hlo-legalize-to-lhlo %s | FileCheck %s
2
3// Test that assuming ops propagate tensor types.
4// CHECK-LABEL: func @shape_assuming_tensor
5func @shape_assuming_tensor(%arg0: tensor<?xf16>) -> tensor<?xf16> {
6  %0 = mhlo.constant dense<0.000000e+00> : tensor<f16>
7  %1 = shape.const_witness true
8  // CHECK: shape.assuming %{{.*}} -> (memref<?xf16>)
9  %2 = shape.assuming %1 -> (tensor<?xf16>) {
10    // CHECK: "lmhlo.maximum"(%{{.*}}, %{{.*}}, %{{.*}}) : (memref<?xf16>, memref<?xf16>, memref<?xf16>) -> ()
11    %7 = mhlo.maximum %arg0, %arg0 : tensor<?xf16>
12    // CHECK: shape.assuming_yield %{{.*}} : memref<?xf16>
13    shape.assuming_yield %7 : tensor<?xf16>
14  }
15  return %2 : tensor<?xf16>
16}
17