Using MLIR's Affine Dialect With FIR

Rajan Walia, Eric Schweitz, Steve Scalpone

Affine Dialect

Affine Dialect

The Good!

  • Dialect for polyhedral optimizations
  • Provides array, loop and condition operations.
  • Comes with optimizations for loop fusion, tiling, unrolling, vectorization, etc.

Affine Maps

#map0 = affine_map<()[s0, s1] -> (s0 - s1 + 1)>
#map1 = affine_map<(d0, d1)[s0, s1, s2, s3, s4, s5] -> 
           ((d1 - s3) * ((s1 - s0 + 1) * s2) + d0 - s0)>

...
 %5 = affine.apply #map1 (%arg3, %arg4)
                    [%c0, %c20, %c1, %c10, %c20, %c1]
...

Affine loops

  affine.for %i = 1 to 10 {
     ...
  }
  affine.for %arg4 = %22 to affine_map<()[s0] -> (s0 + 1)>()[%24] {
     ...
  }

Affine Constraints

  • Array index can only be specific values
  • No partial indexing for arrays
  • Every affine loop must finish to the end

Affine memory ops

   %12 = affine.load %11[%10] : memref<?xf32>
   affine.store %10, %11[%5] : memref<?xf32>

Affine Dialect

The Bad!

  • Heavily depended on memref type.

    • Array values used in affine operations must be of memref type.

Dummy converts

%7 = fir.convert %3 : (!fir.ref<!fir.array<?x?xf32>>) 
                      -> memref<?xf32>

fortran
optimizations for affine
affine promotion
fir

fir + affine

affine optimizations

affine demotion

fir

llvm ir

Changes to FIR

adding new operations

Operations

  • shape, shape_shift, slice
  • array_coor

shape operations

%sh = fir.shape %row_sz, %col_sz : 
	(index, index) -> !fir.shape<2>
%ss = fir.shape_shift %lo, %extent :
	(index, index) -> !fir.shapeshift<1>
%sl = fir.slice %lo, %hi, %step :
	(index, index, index) -> !fir.slice<1>

array coordinate

%s = fir.shape %n, %m : (index, index) -> !fir.shape<2>
%1 = fir.array_coor %a(%s) %i, %j : 
	(!fir.ref<!fir.array<?x?xf32>>, !fir.shape<2>, index, index) 
    	-> !fir.ref<f32>

Example

subroutine f1dc(a1,a2,ret)
  integer a1(0:4), a2(0:4), ret(0:4)
  integer t1(0:4)

  do i = 0,4
     t1(i) = a1(i) + a1(i)
  end do
  do i = 0,4
     ret(i) = t1(i) * a2(i)
  end do
end subroutine f1dc

Original FIR

func @_QPf1dc(%arg0: !fir.ref<!fir.array<5xi32>>, 
              %arg1: !fir.ref<!fir.array<5xi32>>, 
              %arg2: !fir.ref<!fir.array<5xi32>>) {
  %c0 = constant 0 : index
  %c1 = constant 1 : index
  %c4 = constant 4 : index
  %0 = fir.alloca i32 {name = "i"}
  %1 = fir.alloca !fir.array<5xi32> {name = "t1"}
  %4 = fir.do_loop %arg3 = %c0 to %c4 step %c1 
   iter_args(%arg4 = %2) -> (index) {
    %10 = fir.convert %arg3 : (index) -> i32
    fir.store %10 to %0 : !fir.ref<i32>
    %11 = fir.load %0 : !fir.ref<i32>
    ...
    %14 = subi %12, %13 : i64
    %15 = fir.coordinate_of %1, %14 : 
      		(!fir.ref<!fir.array<5xi32>>, i64) -> !fir.ref<i32>
    %16 = fir.load %0 : !fir.ref<i32>
    ...
    %19 = subi %17, %18 : i64
    %20 = fir.coordinate_of %arg0, %19 : 
    		(!fir.ref<!fir.array<5xi32>>, i64) -> !fir.ref<i32>
    %21 = fir.load %20 : !fir.ref<i32>
    %22 = fir.load %0 : !fir.ref<i32>
    %25 = subi %23, %24 : i64
    %26 = fir.coordinate_of %arg0, %25 : 
    		(!fir.ref<!fir.array<5xi32>>, i64) -> !fir.ref<i32>
    %27 = fir.load %26 : !fir.ref<i32>
    %28 = addi %21, %27 : i32
    fir.store %28 to %15 : !fir.ref<i32>
    %29 = addi %arg3, %c1 : index
    fir.result %29 : index
  }
  %5 = fir.convert %4 : (index) -> i32
  fir.store %5 to %0 : !fir.ref<i32>
  %8 = fir.do_loop %arg3 = %c0 to %c4 step %c1 iter_args(%arg4 = %6) -> (index) {
    %10 = fir.convert %arg3 : (index) -> i32
    fir.store %10 to %0 : !fir.ref<i32>
    %11 = fir.load %0 : !fir.ref<i32>
    %12 = fir.convert %11 : (i32) -> i64
    %13 = fir.convert %c0_1 : (index) -> i64
    %14 = subi %12, %13 : i64
    %15 = fir.coordinate_of %arg2, %14 : (!fir.ref<!fir.array<5xi32>>, i64) -> !fir.ref<i32>
    %16 = fir.load %0 : !fir.ref<i32>
    %17 = fir.convert %16 : (i32) -> i64
    %18 = fir.convert %c0_2 : (index) -> i64
    %19 = subi %17, %18 : i64
    %20 = fir.coordinate_of %1, %19 : (!fir.ref<!fir.array<5xi32>>, i64) -> !fir.ref<i32>
    %21 = fir.load %20 : !fir.ref<i32>
    %22 = fir.load %0 : !fir.ref<i32>
    %23 = fir.convert %22 : (i32) -> i64
    %24 = fir.convert %c0_0 : (index) -> i64
    %25 = subi %23, %24 : i64
    %26 = fir.coordinate_of %arg1, %25 : (!fir.ref<!fir.array<5xi32>>, i64) -> !fir.ref<i32>
    %27 = fir.load %26 : !fir.ref<i32>
    %28 = muli %21, %27 : i32
    fir.store %28 to %15 : !fir.ref<i32>
    %29 = addi %arg3, %c1_5 : index
    fir.result %29 : index
  }
  %9 = fir.convert %8 : (index) -> i32
  fir.store %9 to %0 : !fir.ref<i32>
  return
}

Using new operations

func @_QPf1dc(%arg0: !fir.ref<!fir.array<5xi32>>, %arg1: !fir.ref<!fir.array<5xi32>>, %arg2: !fir.ref<!fir.array<5xi32>>) {
  ...
  %2 = fir.shape_shift %c0, %c5 : (index, index) -> !fir.shapeshift<1>
  %3 = fir.shape_shift %c0, %c5 : ...
  %4 = fir.shape_shift %c0, %c5 : ...
  fir.do_loop %arg3 = %c0 to %c4 step %c1 {
    %8 = fir.array_coor %1(%2) %arg3 : (..., !fir.shapeshift<1>, index) -> !fir.ref<i32>
    %9 = fir.array_coor %arg0(%3) %arg3 : ...
    %10 = fir.load %9 : !fir.ref<i32>
    %11 = fir.array_coor %arg0(%4) %arg3 : 
    		(!fir.ref<!fir.array<5xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
    %12 = fir.load %11 : !fir.ref<i32>
    %13 = addi %10, %12 : i32
    fir.store %13 to %8 : !fir.ref<i32>
  }
  %5 = fir.shape_shift %c0, %c5 : (index, index) -> !fir.shapeshift<1>
  %6 = fir.shape_shift %c0, %c5 : (index, index) -> !fir.shapeshift<1>
  %7 = fir.shape_shift %c0, %c5 : (index, index) -> !fir.shapeshift<1>
  fir.do_loop %arg3 = %c0 to %c4 step %c1 {
    %8 = fir.array_coor %arg2(%5) %arg3 : (!fir.ref<!fir.array<5xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
    %9 = fir.array_coor %1(%6) %arg3 : (!fir.ref<!fir.array<5xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
    %10 = fir.load %9 : !fir.ref<i32>
    %11 = fir.array_coor %arg1(%7) %arg3 : (!fir.ref<!fir.array<5xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
    %12 = fir.load %11 : !fir.ref<i32>
    %13 = muli %10, %12 : i32
    fir.store %13 to %8 : !fir.ref<i32>
  }
  return
}

Affine Promotion

func @_QPf1dc(%arg0: !fir.ref<!fir.array<5xi32>>, %arg1: !fir.ref<!fir.array<5xi32>>, %arg2: !fir.ref<!fir.array<5xi32>>) {
  ...
  affine.for %arg3 = %c0 to affine_map<()[s0] -> (s0 + 1)>()[%c4] {
    %8 = fir.array_coor %1(%2) %arg3 : ...
    %9 = fir.array_coor %arg0(%3) %arg3 : ...
    %10 = affine.apply affine_map<(d0)[s0, s1, s2] -> (d0 * s2 - s0)>(%arg3)[%c0, %c5, %c1_0]
    %11 = fir.convert %arg0 : (!fir.ref<!fir.array<5xi32>>) -> memref<?xi32>
    %12 = affine.load %11[%10] : memref<?xi32>
    %14 = affine.apply affine_map<(d0)[s0, s1, s2] -> (d0 * s2 - s0)>(%arg3)[%c0, %c5, %c1_1]
    %15 = fir.convert %arg0 : (!fir.ref<!fir.array<5xi32>>) -> memref<?xi32>
    %16 = affine.load %15[%14] : memref<?xi32>
    %17 = addi %12, %16 : i32
    %18 = affine.apply affine_map<(d0)[s0, s1, s2] -> (d0 * s2 - s0)>(%arg3)[%c0, %c5, %c1_2]
    affine.store %17, %19[%18] : memref<?xi32>
  }
  affine.for %arg3 = %c0 to affine_map<()[s0] -> (s0 + 1)>()[%c4] {
    ...
    %10 = affine.apply affine_map<(d0)[s0, s1, s2] -> (d0 * s2 - s0)>(%arg3)[%c0, %c5, %c1_0]
    %11 = fir.convert %1 : (!fir.ref<!fir.array<5xi32>>) -> memref<?xi32>
    %12 = affine.load %11[%10] : memref<?xi32>
    %14 = affine.apply affine_map<(d0)[s0, s1, s2] -> (d0 * s2 - s0)>(%arg3)[%c0, %c5, %c1_1]
    %15 = fir.convert %arg1 : (!fir.ref<!fir.array<5xi32>>) -> memref<?xi32>
    %16 = affine.load %15[%14] : memref<?xi32>
    %17 = muli %12, %16 : i32
    %18 = affine.apply affine_map<(d0)[s0, s1, s2] -> (d0 * s2 - s0)>(%arg3)[%c0, %c5, %c1_2]
    %19 = fir.convert %arg2 : (!fir.ref<!fir.array<5xi32>>) -> memref<?xi32>
    affine.store %17, %19[%18] : memref<?xi32>
  }
  return
}

Basic Affine Optimizations

func @_QPf1dc(%arg0: !fir.ref<!fir.array<5xi32>>, %arg1: !fir.ref<!fir.array<5xi32>>, %arg2: !fir.ref<!fir.array<5xi32>>) {
  %0 = fir.alloca i32 {name = "i"}
  %1 = fir.alloca !fir.array<5xi32> {name = "t1"}
  %2 = fir.convert %arg0 : (!fir.ref<!fir.array<5xi32>>) -> memref<?xi32>
  %3 = fir.convert %1 : (!fir.ref<!fir.array<5xi32>>) -> memref<?xi32>
  affine.for %arg3 = 0 to 5 {
    %6 = affine.load %2[%arg3] : memref<?xi32>
    %7 = affine.load %2[%arg3] : memref<?xi32>
    %8 = addi %6, %7 : i32
    affine.store %8, %3[%arg3] : memref<?xi32>
  }
  %4 = fir.convert %arg1 : (!fir.ref<!fir.array<5xi32>>) -> memref<?xi32>
  %5 = fir.convert %arg2 : (!fir.ref<!fir.array<5xi32>>) -> memref<?xi32>
  affine.for %arg3 = 0 to 5 {
    %6 = affine.load %3[%arg3] : memref<?xi32>
    %7 = affine.load %4[%arg3] : memref<?xi32>
    %8 = muli %6, %7 : i32
    affine.store %8, %5[%arg3] : memref<?xi32>
  }
  return
}

Affine Loop Fusion

func @_QPf1dc(%arg0: !fir.ref<!fir.array<5xi32>>, %arg1: !fir.ref<!fir.array<5xi32>>, %arg2: !fir.ref<!fir.array<5xi32>>) {
  %0 = fir.alloca i32 {name = "i"}
  %1 = fir.alloca !fir.array<5xi32> {name = "t1"}
  %2 = fir.convert %arg0 : (!fir.ref<!fir.array<5xi32>>) -> memref<?xi32>
  %3 = fir.convert %1 : (!fir.ref<!fir.array<5xi32>>) -> memref<?xi32>
  %4 = fir.convert %arg1 : (!fir.ref<!fir.array<5xi32>>) -> memref<?xi32>
  %5 = fir.convert %arg2 : (!fir.ref<!fir.array<5xi32>>) -> memref<?xi32>
  affine.for %arg3 = 0 to 5 {
    %6 = affine.load %2[%arg3] : memref<?xi32>
    %7 = affine.load %2[%arg3] : memref<?xi32>
    %8 = addi %6, %7 : i32
    %9 = affine.load %4[%arg3] : memref<?xi32>
    %10 = muli %8, %9 : i32
    affine.store %10, %5[%arg3] : memref<?xi32>
  }
  return
}

Affine Demotion

func @_QPf1dc(%arg0: !fir.ref<!fir.array<5xi32>>, %arg1: !fir.ref<!fir.array<5xi32>>, %arg2: !fir.ref<!fir.array<5xi32>>) {
  %0 = fir.alloca !fir.array<1xi32>
  %1 = fir.alloca i32 {name = "i"}
  %2 = fir.alloca !fir.array<5xi32> {name = "t1"}
  affine.for %arg3 = 0 to 5 {
    %7 = fir.coordinate_of %arg0, %arg3 : (!fir.ref<!fir.array<?xi32>>, index) -> !fir.ref<i32>
    %8 = fir.load %7 : !fir.ref<i32>
    %9 = fir.coordinate_of %arg0, %arg3 : ...
    %10 = fir.load %9 : !fir.ref<i32>
    %11 = addi %8, %10 : i32
    %12 = fir.coordinate_of %0, %c0 : ...
    fir.store %11 to %12 : !fir.ref<i32>
    %13 = fir.coordinate_of %0, %c0 : ...
    %14 = fir.load %13 : !fir.ref<i32>
    %15 = fir.coordinate_of %arg1, %arg3 : ...
    %16 = fir.load %15 : !fir.ref<i32>
    %17 = muli %14, %16 : i32
    %18 = fir.coordinate_of %arg2, %arg3 : ...
    fir.store %17 to %18 : !fir.ref<i32>
  }
  return
}

affine-fir

By Rajan Walia

Made with Slides.com

affine-fir

  • 512

More from Rajan Walia