mlir-hlo/tests/end2end/broadcast.mlir

// RUN: mlir-hlo-opt %s -chlo-legalize-to-hlo -hlo-legalize-to-lhlo \
// RUN: -buffer-hoisting -buffer-deallocation -copy-removal -canonicalize -cse \
// RUN: -lhlo-legalize-to-linalg -lhlo-fuse-linalg -convert-linalg-to-loops \
// RUN: -canonicalize -cse -convert-linalg-to-llvm -convert-std-to-llvm \
// RUN: | mlir-cpu-runner -e main -entry-point-result=void \
// RUN: -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext \
// RUN: | FileCheck %s --dump-input=always

func @main() -> () {
  call @trivial_broadcast_wrapper() : () -> ()
  call @broadcast_in_X_dim_wrapper() : () -> ()
  call @broadcast_in_Y_dim_wrapper() : () -> ()
  call @broadcast_in_X_dim_transpose_wrapper() : () -> ()
  call @broadcast_in_Y_dim_transpose_wrapper() : () -> ()
  call @broadcast_scalar_1d_wrapper() : () -> ()
  call @broadcast_scalar_2d_wrapper() : () -> ()
  call @broadcast_to_the_same_shape() : () -> ()
  call @broadcast_1d_to_2d() : () -> ()
  call @broadcast_1d_to_2d_with_transpose() : () -> ()
  return
}

func @print_memref_i8(memref<*xi8>) attributes { llvm.emit_c_interface }
func @print_memref_f32(memref<*xf32>) attributes { llvm.emit_c_interface }

func @trivial_broadcast_wrapper() {
  %input_buf = alloc() : memref<3xf32>

  %c1f32 = constant 1.0 : f32
  %c2f32 = constant 2.0 : f32
  %c3f32 = constant 3.0 : f32

  %c0 = constant 0 : index
  %c1 = constant 1 : index
  %c2 = constant 2 : index
  store %c1f32, %input_buf[%c0] : memref<3xf32>
  store %c2f32, %input_buf[%c1] : memref<3xf32>
  store %c3f32, %input_buf[%c2] : memref<3xf32>
  %input = tensor_load %input_buf : memref<3xf32>

  // Test BroadcastInDimOp.
  %output = "mhlo.broadcast_in_dim"(%input) {
    broadcast_dimensions = dense<0> : tensor<1xi64>
  } : (tensor<3xf32>) -> tensor<3x4xf32>

  %output_buf = alloc() : memref<3x4xf32>
  tensor_store %output, %output_buf : memref<3x4xf32>

  %unranked_output = memref_cast %output_buf : memref<3x4xf32> to memref<*xf32>
  call @print_memref_f32(%unranked_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1]
  // CHECK-NEXT: [1,   1,   1,   1]
  // CHECK-NEXT: [2,   2,   2,   2]
  // CHECK-NEXT: [3,   3,   3,   3]

  // Test DynamicBroadcastInDimOp.
  %c3 = constant 3 : index
  %c4 = constant 4 : index
  %shape = tensor_from_elements %c3, %c4 : tensor<2xindex>
  %dyn_output = "mhlo.dynamic_broadcast_in_dim"(%input, %shape) {
    broadcast_dimensions = dense<0> : tensor<1xi64>
  } : (tensor<3xf32>, tensor<2xindex>) -> tensor<3x4xf32>

  %dyn_output_buf = alloc() : memref<3x4xf32>
  tensor_store %dyn_output, %dyn_output_buf : memref<3x4xf32>

  %unranked_dyn_output = memref_cast %dyn_output_buf
    : memref<3x4xf32> to memref<*xf32>
  call @print_memref_f32(%unranked_dyn_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1]
  // CHECK-NEXT: [1,   1,   1,   1]
  // CHECK-NEXT: [2,   2,   2,   2]
  // CHECK-NEXT: [3,   3,   3,   3]
  return
}

func @broadcast_in_X_dim_wrapper() {
  %input_buf = alloc() : memref<1x4xf32>
  %c1f32 = constant 1.0 : f32
  %c0 = constant 0 : index
  store %c1f32, %input_buf[%c0, %c0] : memref<1x4xf32>
  %c2f32 = constant 2.0 : f32
  %c1 = constant 1 : index
  store %c2f32, %input_buf[%c0, %c1] : memref<1x4xf32>
  %c3f32 = constant 3.0 : f32
  %c2 = constant 2 : index
  store %c3f32, %input_buf[%c0, %c2] : memref<1x4xf32>
  %c4f32 = constant 4.0 : f32
  %c3 = constant 3 : index
  store %c4f32, %input_buf[%c0, %c3] : memref<1x4xf32>
  %input = tensor_load %input_buf : memref<1x4xf32>

  // Test BroadcastInDimOp.
  %output = "mhlo.broadcast_in_dim"(%input) {
    broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>
  } : (tensor<1x4xf32>) -> tensor<3x4xf32>

  %output_buf = alloc() : memref<3x4xf32>
  tensor_store %output, %output_buf : memref<3x4xf32>

  %unranked_output = memref_cast %output_buf : memref<3x4xf32> to memref<*xf32>
  call @print_memref_f32(%unranked_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1]
  // CHECK-NEXT: [1,   2,   3,   4]
  // CHECK-NEXT: [1,   2,   3,   4]
  // CHECK-NEXT: [1,   2,   3,   4]

  // Test DynamicBroadcastInDimOp.
  %c4 = constant 4 : index
  %shape = tensor_from_elements %c3, %c4 : tensor<2xindex>
  %dyn_output = "mhlo.dynamic_broadcast_in_dim"(%input, %shape) {
    broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>
  } : (tensor<1x4xf32>, tensor<2xindex>) -> tensor<3x4xf32>

  %dyn_output_buf = alloc() : memref<3x4xf32>
  tensor_store %dyn_output, %dyn_output_buf : memref<3x4xf32>

  %unranked_dyn_output = memref_cast %dyn_output_buf
    : memref<3x4xf32> to memref<*xf32>
  call @print_memref_f32(%unranked_dyn_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1]
  // CHECK-NEXT: [1,   2,   3,   4]
  // CHECK-NEXT: [1,   2,   3,   4]
  // CHECK-NEXT: [1,   2,   3,   4]
  return
}

func @broadcast_in_Y_dim_wrapper() {
  %input_buf = alloc() : memref<3x1xf32>
  %c1f32 = constant 1.0 : f32
  %c0 = constant 0 : index
  store %c1f32, %input_buf[%c0, %c0] : memref<3x1xf32>
  %c2f32 = constant 2.0 : f32
  %c1 = constant 1 : index
  store %c2f32, %input_buf[%c1, %c0] : memref<3x1xf32>
  %c3f32 = constant 3.0 : f32
  %c2 = constant 2 : index
  store %c3f32, %input_buf[%c2, %c0] : memref<3x1xf32>
  %input = tensor_load %input_buf : memref<3x1xf32>

  // Test BroadcastInDimOp.
  %output = "mhlo.broadcast_in_dim"(%input) {
    broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>
  } : (tensor<3x1xf32>) -> tensor<3x4xf32>

  %output_buf = alloc() : memref<3x4xf32>
  tensor_store %output, %output_buf : memref<3x4xf32>

  %unranked_output = memref_cast %output_buf : memref<3x4xf32> to memref<*xf32>
  call @print_memref_f32(%unranked_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1]
  // CHECK-NEXT: [1,   1,   1,   1]
  // CHECK-NEXT: [2,   2,   2,   2]
  // CHECK-NEXT: [3,   3,   3,   3]

  // Test DynamicBroadcastInDimOp.
  %c3 = constant 3 : index
  %c4 = constant 4 : index
  %shape = tensor_from_elements %c3, %c4 : tensor<2xindex>
  %dyn_output = "mhlo.dynamic_broadcast_in_dim"(%input, %shape) {
    broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>
  } : (tensor<3x1xf32>, tensor<2xindex>) -> tensor<3x4xf32>

  %dyn_output_buf = alloc() : memref<3x4xf32>
  tensor_store %dyn_output, %dyn_output_buf : memref<3x4xf32>

  %unranked_dyn_output = memref_cast %dyn_output_buf
    : memref<3x4xf32> to memref<*xf32>
  call @print_memref_f32(%unranked_dyn_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1]
  // CHECK-NEXT: [1,   1,   1,   1]
  // CHECK-NEXT: [2,   2,   2,   2]
  // CHECK-NEXT: [3,   3,   3,   3]
  return
}

func @broadcast_in_X_dim_transpose_wrapper() {
  %input_buf = alloc() : memref<4x1xf32>
  %c1f32 = constant 1.0 : f32
  %c0 = constant 0 : index
  store %c1f32, %input_buf[%c0, %c0] : memref<4x1xf32>
  %c2f32 = constant 2.0 : f32
  %c1 = constant 1 : index
  store %c2f32, %input_buf[%c1, %c0] : memref<4x1xf32>
  %c3f32 = constant 3.0 : f32
  %c2 = constant 2 : index
  store %c3f32, %input_buf[%c2, %c0] : memref<4x1xf32>
  %c4f32 = constant 4.0 : f32
  %c3 = constant 3 : index
  store %c4f32, %input_buf[%c3, %c0] : memref<4x1xf32>
  %input = tensor_load %input_buf : memref<4x1xf32>

  // Test BroadcastInDimOp.
  %output = "mhlo.broadcast_in_dim"(%input) {
    broadcast_dimensions = dense<[1, 0]> : tensor<2xi64>
  } : (tensor<4x1xf32>) -> tensor<3x4xf32>

  %output_buf = alloc() : memref<3x4xf32>
  tensor_store %output, %output_buf : memref<3x4xf32>

  %unranked_output = memref_cast %output_buf : memref<3x4xf32> to memref<*xf32>
  call @print_memref_f32(%unranked_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1]
  // CHECK-NEXT: [1,   2,   3,   4]
  // CHECK-NEXT: [1,   2,   3,   4]
  // CHECK-NEXT: [1,   2,   3,   4]

  // Test DynamicBroadcastInDimOp.
  %c4 = constant 4 : index
  %shape = tensor_from_elements %c3, %c4 : tensor<2xindex>
  %dyn_output = "mhlo.dynamic_broadcast_in_dim"(%input, %shape) {
    broadcast_dimensions = dense<[1, 0]> : tensor<2xi64>
  } : (tensor<4x1xf32>, tensor<2xindex>) -> tensor<3x4xf32>

  %dyn_output_buf = alloc() : memref<3x4xf32>
  tensor_store %dyn_output, %dyn_output_buf : memref<3x4xf32>

  %unranked_dyn_output = memref_cast %dyn_output_buf
    : memref<3x4xf32> to memref<*xf32>
  call @print_memref_f32(%unranked_dyn_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1]
  // CHECK-NEXT: [1,   2,   3,   4]
  // CHECK-NEXT: [1,   2,   3,   4]
  // CHECK-NEXT: [1,   2,   3,   4]
  return
}

func @broadcast_in_Y_dim_transpose_wrapper() {
  %input_buf = alloc() : memref<1x3xf32>
  %c1f32 = constant 1.0 : f32
  %c0 = constant 0 : index
  store %c1f32, %input_buf[%c0, %c0] : memref<1x3xf32>
  %c2f32 = constant 2.0 : f32
  %c1 = constant 1 : index
  store %c2f32, %input_buf[%c0, %c1] : memref<1x3xf32>
  %c3f32 = constant 3.0 : f32
  %c2 = constant 2 : index
  store %c3f32, %input_buf[%c0, %c2] : memref<1x3xf32>
  %input = tensor_load %input_buf : memref<1x3xf32>

  // Test BroadcastInDimOp.
  %output = "mhlo.broadcast_in_dim"(%input) {
    broadcast_dimensions = dense<[1, 0]> : tensor<2xi64>
  } : (tensor<1x3xf32>) -> tensor<3x4xf32>

  %output_buf = alloc() : memref<3x4xf32>
  tensor_store %output, %output_buf : memref<3x4xf32>

  %unranked_output = memref_cast %output_buf : memref<3x4xf32> to memref<*xf32>
  call @print_memref_f32(%unranked_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1]
  // CHECK-NEXT-NEXT: [1,   1,   1,   1]
  // CHECK-NEXT-NEXT: [2,   2,   2,   2]
  // CHECK-NEXT-NEXT: [3,   3,   3,   3]

  // Test DynamicBroadcastInDimOp.
  %c3 = constant 3 : index
  %c4 = constant 4 : index
  %shape = tensor_from_elements %c3, %c4 : tensor<2xindex>
  %dyn_output = "mhlo.dynamic_broadcast_in_dim"(%input, %shape) {
    broadcast_dimensions = dense<[1, 0]> : tensor<2xi64>
  } : (tensor<1x3xf32>, tensor<2xindex>) -> tensor<3x4xf32>

  %dyn_output_buf = alloc() : memref<3x4xf32>
  tensor_store %dyn_output, %dyn_output_buf : memref<3x4xf32>

  %unranked_dyn_output = memref_cast %dyn_output_buf
    : memref<3x4xf32> to memref<*xf32>
  call @print_memref_f32(%unranked_dyn_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1]
  // CHECK-NEXT-NEXT: [1,   1,   1,   1]
  // CHECK-NEXT-NEXT: [2,   2,   2,   2]
  // CHECK-NEXT-NEXT: [3,   3,   3,   3]
  return
}

func @broadcast_scalar_1d_wrapper() {
  %input_buf = alloc() : memref<1xf32>
  %c1f32 = constant 1.0 : f32
  %c0 = constant 0 : index
  store %c1f32, %input_buf[%c0] : memref<1xf32>
  %input = tensor_load %input_buf : memref<1xf32>

  // Test BroadcastInDimOp.
  %output = "mhlo.broadcast_in_dim"(%input) {
    broadcast_dimensions = dense<0> : tensor<1xi64>
  } : (tensor<1xf32>) -> tensor<3x4xf32>

  %output_buf = alloc() : memref<3x4xf32>
  tensor_store %output, %output_buf : memref<3x4xf32>

  %unranked_output = memref_cast %output_buf : memref<3x4xf32> to memref<*xf32>
  call @print_memref_f32(%unranked_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1]
  // CHECK-NEXT: [1, 1, 1, 1]
  // CHECK-NEXT: [1, 1, 1, 1]
  // CHECK-NEXT: [1, 1, 1, 1]

  // Test DynamicBroadcastInDimOp.
  %c3 = constant 3 : index
  %c4 = constant 4 : index
  %shape = tensor_from_elements %c3, %c4 : tensor<2xindex>
  %dyn_output = "mhlo.dynamic_broadcast_in_dim"(%input, %shape) {
    broadcast_dimensions = dense<0> : tensor<1xi64>
  } : (tensor<1xf32>, tensor<2xindex>) -> tensor<3x4xf32>

  %dyn_output_buf = alloc() : memref<3x4xf32>
  tensor_store %dyn_output, %dyn_output_buf : memref<3x4xf32>

  %unranked_dyn_output = memref_cast %dyn_output_buf
    : memref<3x4xf32> to memref<*xf32>
  call @print_memref_f32(%unranked_dyn_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1]
  // CHECK-NEXT: [1, 1, 1, 1]
  // CHECK-NEXT: [1, 1, 1, 1]
  // CHECK-NEXT: [1, 1, 1, 1]
  return
}

func @broadcast_scalar_2d_wrapper() {
  %input_buf = alloc() : memref<1x1xf32>
  %c1f32 = constant 1.0 : f32
  %c0 = constant 0 : index
  store %c1f32, %input_buf[%c0, %c0] : memref<1x1xf32>
  %input = tensor_load %input_buf : memref<1x1xf32>

  // Test BroadcastInDimOp.
  %output = "mhlo.broadcast_in_dim"(%input) {
    broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>
  } : (tensor<1x1xf32>) -> tensor<3x4xf32>

  %output_buf = alloc() : memref<3x4xf32>
  tensor_store %output, %output_buf : memref<3x4xf32>

  %unranked_output = memref_cast %output_buf : memref<3x4xf32> to memref<*xf32>
  call @print_memref_f32(%unranked_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1]
  // CHECK-NEXT: [1, 1, 1, 1]
  // CHECK-NEXT: [1, 1, 1, 1]
  // CHECK-NEXT: [1, 1, 1, 1]

  // Test DynamicBroadcastInDimOp.
  %c3 = constant 3 : index
  %c4 = constant 4 : index
  %shape = tensor_from_elements %c3, %c4 : tensor<2xindex>
  %dyn_output = "mhlo.dynamic_broadcast_in_dim"(%input, %shape) {
    broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>
  } : (tensor<1x1xf32>, tensor<2xindex>) -> tensor<3x4xf32>

  %dyn_output_buf = alloc() : memref<3x4xf32>
  tensor_store %dyn_output, %dyn_output_buf : memref<3x4xf32>

  %unranked_dyn_output = memref_cast %dyn_output_buf
    : memref<3x4xf32> to memref<*xf32>
  call @print_memref_f32(%unranked_dyn_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1]
  // CHECK-NEXT: [1, 1, 1, 1]
  // CHECK-NEXT: [1, 1, 1, 1]
  // CHECK-NEXT: [1, 1, 1, 1]
  return
}

func @broadcast_to_the_same_shape() {
  %input_buf = alloc() : memref<2x3xf32>

  %c1f32 = constant 1.0 : f32
  %c2f32 = constant 2.0 : f32
  %c3f32 = constant 3.0 : f32

  %c0 = constant 0 : index
  %c1 = constant 1 : index
  %c2 = constant 2 : index
  %c3 = constant 3 : index
  store %c1f32, %input_buf[%c0, %c0] : memref<2x3xf32>
  store %c1f32, %input_buf[%c1, %c0] : memref<2x3xf32>
  store %c2f32, %input_buf[%c0, %c1] : memref<2x3xf32>
  store %c2f32, %input_buf[%c1, %c1] : memref<2x3xf32>
  store %c3f32, %input_buf[%c0, %c2] : memref<2x3xf32>
  store %c3f32, %input_buf[%c1, %c2] : memref<2x3xf32>
  %input = tensor_load %input_buf : memref<2x3xf32>

  // Test BroadcastInDimOp.
  %output = "mhlo.broadcast_in_dim"(%input) {
    broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>
  } : (tensor<2x3xf32>) -> tensor<2x3xf32>

  %output_buf = alloc() : memref<2x3xf32>
  tensor_store %output, %output_buf : memref<2x3xf32>

  %unraked_output = memref_cast %output_buf : memref<2x3xf32> to memref<*xf32>
  call @print_memref_f32(%unraked_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1]
  // CHECK-NEXT: [1,   2,   3]
  // CHECK-NEXT: [1,   2,   3]

  // Test DynamicBroadcastInDimOp.
  %shape = tensor_from_elements %c2, %c3 : tensor<2xindex>
  %dyn_output = "mhlo.dynamic_broadcast_in_dim"(%input, %shape) {
    broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>
  } : (tensor<2x3xf32>, tensor<2xindex>) -> tensor<2x3xf32>

  %dyn_output_buf = alloc() : memref<2x3xf32>
  tensor_store %dyn_output, %dyn_output_buf : memref<2x3xf32>

  %unranked_dyn_output = memref_cast %dyn_output_buf
    : memref<2x3xf32> to memref<*xf32>
  call @print_memref_f32(%unranked_dyn_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1]
  // CHECK-NEXT: [1,   2,   3]
  // CHECK-NEXT: [1,   2,   3]
  return
}

func @broadcast_1d_to_2d() {
  %input_buf = alloc() : memref<3xf32>

  %c1f32 = constant 1.0 : f32
  %c2f32 = constant 2.0 : f32
  %c3f32 = constant 3.0 : f32

  %c0 = constant 0 : index
  %c1 = constant 1 : index
  %c2 = constant 2 : index
  store %c1f32, %input_buf[%c0] : memref<3xf32>
  store %c2f32, %input_buf[%c1] : memref<3xf32>
  store %c3f32, %input_buf[%c2] : memref<3xf32>
  %input = tensor_load %input_buf : memref<3xf32>

  // Test BroadcastInDimOp.
  %output = "mhlo.broadcast_in_dim"(%input) {
    broadcast_dimensions = dense<0> : tensor<1xi64>
  } : (tensor<3xf32>) -> tensor<3x3xf32>

  %output_buf = alloc() : memref<3x3xf32>
  tensor_store %output, %output_buf : memref<3x3xf32>

  %unraked_output = memref_cast %output_buf : memref<3x3xf32> to memref<*xf32>
  call @print_memref_f32(%unraked_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [3, 3] strides = [3, 1]
  // CHECK-NEXT: [1,   1,   1]
  // CHECK-NEXT: [2,   2,   2]
  // CHECK-NEXT: [3,   3,   3]

  // Test DynamicBroadcastInDimOp.
  %c3 = constant 3 : index
  %c4 = constant 3 : index
  %shape = tensor_from_elements %c3, %c4 : tensor<2xindex>
  %dyn_output = "mhlo.dynamic_broadcast_in_dim"(%input, %shape) {
    broadcast_dimensions = dense<0> : tensor<1xi64>
  } : (tensor<3xf32>, tensor<2xindex>) -> tensor<3x3xf32>

  %dyn_output_buf = alloc() : memref<3x3xf32>
  tensor_store %dyn_output, %dyn_output_buf : memref<3x3xf32>

  %unranked_dyn_output = memref_cast %dyn_output_buf
    : memref<3x3xf32> to memref<*xf32>
  call @print_memref_f32(%unranked_dyn_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [3, 3] strides = [3, 1]
  // CHECK-NEXT: [1,   1,   1]
  // CHECK-NEXT: [2,   2,   2]
  // CHECK-NEXT: [3,   3,   3]
  return
}

func @broadcast_1d_to_2d_with_transpose() {
  %input_buf = alloc() : memref<3xf32>

  %c1f32 = constant 1.0 : f32
  %c2f32 = constant 2.0 : f32
  %c3f32 = constant 3.0 : f32

  %c0 = constant 0 : index
  %c1 = constant 1 : index
  %c2 = constant 2 : index
  store %c1f32, %input_buf[%c0] : memref<3xf32>
  store %c2f32, %input_buf[%c1] : memref<3xf32>
  store %c3f32, %input_buf[%c2] : memref<3xf32>
  %input = tensor_load %input_buf : memref<3xf32>

  // Test BroadcastInDimOp.
  %output = "mhlo.broadcast_in_dim"(%input) {
    broadcast_dimensions = dense<1> : tensor<1xi64>
  } : (tensor<3xf32>) -> tensor<3x3xf32>

  %output_buf = alloc() : memref<3x3xf32>
  tensor_store %output, %output_buf : memref<3x3xf32>

  %unraked_output = memref_cast %output_buf : memref<3x3xf32> to memref<*xf32>
  call @print_memref_f32(%unraked_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [3, 3] strides = [3, 1]
  // CHECK-NEXT: [1,   2,   3]
  // CHECK-NEXT: [1,   2,   3]
  // CHECK-NEXT: [1,   2,   3]

  // Test DynamicBroadcastInDimOp.
  %c3 = constant 3 : index
  %shape = tensor_from_elements %c3, %c3 : tensor<2xindex>
  %dyn_output = "mhlo.dynamic_broadcast_in_dim"(%input, %shape) {
    broadcast_dimensions = dense<1> : tensor<1xi64>
  } : (tensor<3xf32>, tensor<2xindex>) -> tensor<3x3xf32>

  %dyn_output_buf = alloc() : memref<3x3xf32>
  tensor_store %dyn_output, %dyn_output_buf : memref<3x3xf32>

  %unranked_dyn_output = memref_cast %dyn_output_buf
    : memref<3x3xf32> to memref<*xf32>
  call @print_memref_f32(%unranked_dyn_output) : (memref<*xf32>) -> ()
  // CHECK: rank = 2 offset = 0 sizes = [3, 3] strides = [3, 1]
  // CHECK-NEXT: [1,   2,   3]
  // CHECK-NEXT: [1,   2,   3]
  // CHECK-NEXT: [1,   2,   3]
  return
}