softmax.mlir

// -----// IR Dump After FusionOfTensorOps //----- //
func @softmax__2x2x2x2x2x2x2__f32__uniform(%arg0: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi = "{\22a\22:[[\22ndarray\22,\22f32\22,7,null,null,null,null,null,null,null]],\22r\22:[[\22ndarray\22,\22f32\22,7,null,null,null,null,null,null,null]],\22v\22:1}"}} {
  %c6 = constant 6 : index
  %c5 = constant 5 : index
  %c4 = constant 4 : index
  %c3 = constant 3 : index
  %c2 = constant 2 : index
  %c1 = constant 1 : index
  %c0 = constant 0 : index
  %cst = constant 0.000000e+00 : f32
  %cst_0 = constant 0x7FC00000 : f32
  %cst_1 = constant 0xFF800000 : f32
  %0 = hal.buffer_view.dim %arg0, 0 : index
  %1 = hal.buffer_view.dim %arg0, 1 : index
  %2 = hal.buffer_view.dim %arg0, 2 : index
  %3 = hal.buffer_view.dim %arg0, 3 : index
  %4 = hal.buffer_view.dim %arg0, 4 : index
  %5 = hal.buffer_view.dim %arg0, 5 : index
  %6 = hal.buffer_view.dim %arg0, 6 : index
  %7 = hal.tensor.cast %arg0 : !hal.buffer_view -> tensor<?x?x?x?x?x?x?xf32>{%0, %1, %2, %3, %4, %5, %6}
  %8 = memref.dim %7, %c0 : tensor<?x?x?x?x?x?x?xf32>
  %9 = memref.dim %7, %c1 : tensor<?x?x?x?x?x?x?xf32>
  %10 = memref.dim %7, %c2 : tensor<?x?x?x?x?x?x?xf32>
  %11 = memref.dim %7, %c3 : tensor<?x?x?x?x?x?x?xf32>
  %12 = memref.dim %7, %c4 : tensor<?x?x?x?x?x?x?xf32>
  %13 = memref.dim %7, %c5 : tensor<?x?x?x?x?x?x?xf32>
  %14 = linalg.init_tensor [%8, %9, %10, %11, %12, %13] : tensor<?x?x?x?x?x?xf32>
  %15 = linalg.fill(%14, %cst_1) : tensor<?x?x?x?x?x?xf32>, f32 -> tensor<?x?x?x?x?x?xf32> 
  %16 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4, d5, d6)>, affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4, d5)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "reduction"]} ins(%7 : tensor<?x?x?x?x?x?x?xf32>) outs(%15 : tensor<?x?x?x?x?x?xf32>) {
  ^bb0(%arg1: f32, %arg2: f32):  // no predecessors
    %31 = cmpf ogt, %arg1, %arg2 : f32
    %32 = select %31, %arg1, %arg2 : f32
    %33 = cmpf uno, %arg1, %arg2 : f32
    %34 = select %33, %cst_0, %32 : f32
    linalg.yield %34 : f32
  } -> tensor<?x?x?x?x?x?xf32>
  %17 = memref.dim %7, %c6 : tensor<?x?x?x?x?x?x?xf32>
  %18 = linalg.init_tensor [%8, %9, %10, %11, %12, %13, %17] : tensor<?x?x?x?x?x?x?xf32>
  %19 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4, d5, d6)>, affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4, d5)>, affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4, d5, d6)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%7, %16 : tensor<?x?x?x?x?x?x?xf32>, tensor<?x?x?x?x?x?xf32>) outs(%18 : tensor<?x?x?x?x?x?x?xf32>) {
  ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):  // no predecessors
    %31 = subf %arg1, %arg2 : f32
    %32 = math.exp %31 : f32
    linalg.yield %32 : f32
  } -> tensor<?x?x?x?x?x?x?xf32>
  %20 = linalg.fill(%14, %cst) : tensor<?x?x?x?x?x?xf32>, f32 -> tensor<?x?x?x?x?x?xf32> 
  %21 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4, d5, d6)>, affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4, d5)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "reduction"]} ins(%19 : tensor<?x?x?x?x?x?x?xf32>) outs(%20 : tensor<?x?x?x?x?x?xf32>) {
  ^bb0(%arg1: f32, %arg2: f32):  // no predecessors
    %31 = addf %arg1, %arg2 : f32
    linalg.yield %31 : f32
  } -> tensor<?x?x?x?x?x?xf32>
  %22 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4, d5, d6)>, affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4, d5)>, affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4, d5)>, affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4, d5, d6)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%7, %16, %21 : tensor<?x?x?x?x?x?x?xf32>, tensor<?x?x?x?x?x?xf32>, tensor<?x?x?x?x?x?xf32>) outs(%18 : tensor<?x?x?x?x?x?x?xf32>) {
  ^bb0(%arg1: f32, %arg2: f32, %arg3: f32, %arg4: f32):  // no predecessors
    %31 = subf %arg1, %arg2 : f32
    %32 = math.exp %31 : f32
    %33 = divf %32, %arg3 : f32
    linalg.yield %33 : f32
  } -> tensor<?x?x?x?x?x?x?xf32>
  %23 = memref.dim %22, %c0 : tensor<?x?x?x?x?x?x?xf32>
  %24 = memref.dim %22, %c1 : tensor<?x?x?x?x?x?x?xf32>
  %25 = memref.dim %22, %c2 : tensor<?x?x?x?x?x?x?xf32>
  %26 = memref.dim %22, %c3 : tensor<?x?x?x?x?x?x?xf32>
  %27 = memref.dim %22, %c4 : tensor<?x?x?x?x?x?x?xf32>
  %28 = memref.dim %22, %c5 : tensor<?x?x?x?x?x?x?xf32>
  %29 = memref.dim %22, %c6 : tensor<?x?x?x?x?x?x?xf32>
  %30 = hal.tensor.cast %22 : tensor<?x?x?x?x?x?x?xf32>{%23, %24, %25, %26, %27, %28, %29} -> !hal.buffer_view
  return %30 : !hal.buffer_view
}