onnx-mlir/test/mlir/krnl/permute.mlir

// RUN: onnx-mlir-opt --convert-krnl-to-affine %s -split-input-file | FileCheck %s

func @simple_permute() {
  %ii, %jj = krnl.define_loops 2
  krnl.permute(%ii, %jj) [1, 0] : !krnl.loop, !krnl.loop
  krnl.iterate(%ii, %jj) with (%ii -> %i = 0 to 10, %jj -> %j = 0 to 20) {
    %foo = addi %i, %i : index
  }

  // CHECK-LABEL: simple_permute
  // CHECK-NEXT: affine.for [[OUTER_LOOP_IV:%.+]] = 0 to 20 {
  // CHECK-NEXT:   affine.for [[INNER_LOOP_IV:%.+]] = 0 to 10 {
  // CHECK-NEXT:     [[ADD:%.+]] = addi [[INNER_LOOP_IV]], [[INNER_LOOP_IV]] : index
  // CHECK-NEXT:   }
  // CHECK-NEXT: }
  return
}

// -----

func @tiling() {
  %ii, %ij = krnl.define_loops 2
  %ib, %il = krnl.block %ii 5 : (!krnl.loop) -> (!krnl.loop, !krnl.loop)
  %jb, %jl = krnl.block %ij 4 : (!krnl.loop) -> (!krnl.loop, !krnl.loop)
  krnl.permute(%ib, %il, %jb, %jl) [0, 2, 1, 3] : !krnl.loop, !krnl.loop, !krnl.loop, !krnl.loop
  krnl.iterate(%ib, %jb, %il, %jl) with (%ii -> %i = 0 to 10, %ij -> %j = 0 to 20) {
    %foo = addi %i, %i : index
  }

  // CHECK-LABEL: tiling
  // CHECK-NEXT: affine.for [[I_BLOCK_IV:%.+]] = 0 to 10 step 5 {
  // CHECK-NEXT:   affine.for [[J_BLOCK_IV:%.+]] = 0 to 20 step 4 {
  // CHECK-NEXT:     affine.for [[I_LOCAL_IV:%.+]] = #map{{.*}}([[I_BLOCK_IV]]) to #map{{.*}}([[I_BLOCK_IV]]) {
  // CHECK-NEXT:       affine.for [[J_LOCAL_IV:%.+]] = #map{{.*}}([[J_BLOCK_IV]]) to #map{{.*}}([[J_BLOCK_IV]]) {
  // CHECK-NEXT:         %0 = addi %arg2, %arg2 : index
  // CHECK-NEXT:       }
  // CHECK-NEXT:     }
  // CHECK-NEXT:   }
  // CHECK-NEXT: }
  return
}

func @tiling3d() {
  %ii, %jj, %kk = krnl.define_loops 3
  // Blocking each loop by a factor of 4.
  %ib, %il = krnl.block %ii 4 : (!krnl.loop) -> (!krnl.loop, !krnl.loop)
  %jb, %jl = krnl.block %jj 4 : (!krnl.loop) -> (!krnl.loop, !krnl.loop)
  %kb, %kl = krnl.block %kk 4 : (!krnl.loop) -> (!krnl.loop, !krnl.loop)
  // Move iteration over tile coordinates to be the outer loops and iterateion over
  // the inter-tile elements to be the inner loops.
  krnl.permute(%ib, %il, %jb, %jl, %kb, %kl) [0, 3, 1, 4, 2, 5] : !krnl.loop, !krnl.loop, !krnl.loop, !krnl.loop, !krnl.loop, !krnl.loop
  krnl.iterate(%ib, %il, %jb, %jl, %kb, %kl) with (%ii -> %i = 0 to 1024, %jj -> %j = 0 to 2048, %kk -> %k = 0 to 4096)  {
  }

  // CHECK-LABEL: tiling3d
  // CHECK-NEXT:  affine.for [[I_BLOCK_IV:%.+]] = 0 to 1024 step 4 {
  // CHECK-NEXT:    affine.for [[J_BLOCK_IV:%.+]] = 0 to 2048 step 4 {
  // CHECK-NEXT:      affine.for [[K_BLOCK_IV:%.+]] = 0 to 4096 step 4 {
  // CHECK-NEXT:        affine.for [[I_INNER_IV:%.+]] = #map0([[I_BLOCK_IV]]) to #map1([[I_BLOCK_IV]]) {
  // CHECK-NEXT:          affine.for [[J_INNER_IV:%.+]] = #map0([[J_BLOCK_IV]]) to #map1([[J_BLOCK_IV]]) {
  // CHECK-NEXT:            affine.for [[K_INNER_IV:%.+]] = #map0([[K_BLOCK_IV]]) to #map1([[K_BLOCK_IV]]) {
  // CHECK-NEXT:            }
  // CHECK-NEXT:          }
  // CHECK-NEXT:        }
  // CHECK-NEXT:      }
  // CHECK-NEXT:    }
  // CHECK-NEXT:  }
  return
}