Skip to content

Instantly share code, notes, and snippets.

@ScottTodd
Created July 21, 2021 15:56
Show Gist options
  • Save ScottTodd/f0d3df1d76b28c6c7eaad7cc6ea42e79 to your computer and use it in GitHub Desktop.
Save ScottTodd/f0d3df1d76b28c6c7eaad7cc6ea42e79 to your computer and use it in GitHub Desktop.

Lowering tensor.from_elements

// -----// IR Dump After LinalgDetensorize //----- //
  %6 = tensor.from_elements %5 : tensor<1xi8>
  %7 = linalg.tensor_collapse_shape %6 [] : tensor<1xi8> into tensor<i8>
  %8 = flow.tensor.load %7 : tensor<i8>
// -----// IR Dump After ConvertUpstreamToIREE //----- //
  %6 = flow.tensor.splat %5 : tensor<1xi8>{%c1}
  %7 = linalg.tensor_collapse_shape %6 [] : tensor<1xi8> into tensor<i8>
  %8 = flow.tensor.load %7 : tensor<i8>
// -----// IR Dump After ConvertToFlowTensorOps //----- //
  %6 = flow.tensor.splat %5 : tensor<1xi8>{%c1}
  %7 = flow.tensor.reshape %6 : tensor<1xi8> -> tensor<i8>
  %8 = flow.tensor.load %7 : tensor<i8>

next:

  • canonicalize splat + reshape -> splat
  • canonicalize splat + load -> load
  %6 = flow.tensor.splat %5 : tensor<i8>
  %7 = flow.tensor.load %6 : tensor<i8>
  %6 = flow.tensor.load %5 : tensor<i8>

in context:

  %3 = "mhlo.compare"(%2, %1) {comparison_direction = "LT"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
  %4 = tensor.extract %3[] : tensor<i1>
  %5 = zexti %4 : i1 to i8
  %6 = flow.tensor.load %5 : tensor<i8>
  %7 = trunci %6 : i8 to i1
  cond_br %7, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment