Skip to content

Polymer with -reg2mem -pluto-opt generates IR with out-of-bounds accesses for imperfectly nested loops #411

Open
@andidr

Description

@andidr

It seems that for imperfect loop nests, polymer-opt can end up generating IR that invokes functions representing statements with parameters that may cause out-of-bounds memory accesses. For example, when invoking polymer-opt -reg2mem -extract-scop-stmt -pluto-opt kernel.cgeist.mlir -o kernel.polymer.mlir --allow-unregistered-dialect on the following IR saved as kernel.mlir:

func.func @kernel(%arg0: memref<153xi64>, %arg1: memref<153x3x7xi64>, %arg2: memref<7xi64>) {
     affine.for %arg3 = 0 to 153 {
       %0 = affine.load %arg0[%arg3] : memref<153xi64>
       affine.for %arg4 = 0 to 3 {
         affine.for %arg5 = 0 to 7 {
           %1 = affine.load %arg1[%arg3, %arg4, %arg5] : memref<153x3x7xi64>
           %2 = arith.muli %1, %0 : i64
           %3 = affine.load %arg2[%arg5] : memref<7xi64>
           %4 = arith.subi %3, %2 : i64
           affine.store %4, %arg2[%arg5] : memref<7xi64>
         }
       }
     }
     return
}

polymer-opt generates:

map = affine_map<(d0) -> (d0 * 2)>                                                                                                                                                           
#map1 = affine_map<(d0) -> (10, d0 * 2 + 3)>                                                                                                                                                  
#map2 = affine_map<(d0, d1) -> (d0 * 6, d1 * 3 - 1)>                                                                                                                                          
#map3 = affine_map<(d0, d1) -> (29, d0 * 6 + 7, d1 * 3 + 4)>                                                                                                                                  
#map4 = affine_map<(d0, d1) -> (d0 * 2 + d1 * 10)>                                                                                                                                            
#map5 = affine_map<(d0, d1) -> (d0 * 32 + 32, d0 * 2 + d1 * 10 + 3)>                                                                                                                          
#map6 = affine_map<(d0) -> (d0 * 32)>                                                                                                                                                         
#map7 = affine_map<(d0, d1) -> (d0 * 6 + d1 * 30 + 7)>                                                                                                                                        
#map8 = affine_map<(d0, d1, d2) -> (d0 * -6 - d1 * 30 + d2)>                                                                                                                                  
#map9 = affine_map<(d0, d1) -> (d0 + d1 * 5)>                                                                                                                                                 
#map10 = affine_map<(d0, d1, d2) -> (d0 * -2 - d1 * 10 + d2)>                                                                                                                                 
#map11 = affine_map<(d0, d1) -> (d0 * 32, d1 * 96 - 6)>                                                                                                                                       
#map12 = affine_map<(d0, d1) -> (d0 * 96 + 1, d1 * 32 + 32)>                                                                                                                                  
#map13 = affine_map<(d0, d1) -> (d0 * -96 + d1 + 6)>                                                                                                                                          
#map14 = affine_map<(d0) -> (d0 * 16 - 1)>                                                                                                                                                    
#map15 = affine_map<(d0, d1) -> ((d0 * 16) ceildiv 3, d1 * 16)>                                                                                                                               
#map16 = affine_map<(d0, d1, d2) -> (153, (d0 * 16) floordiv 3 + 6, d1 * 32 + 32, d2 * 16 + 16)>                                                                                              
#map17 = affine_map<(d0, d1) -> (d0 * 32 + 32, d1 * 2 + 3)>                                                                                                                                   
#map18 = affine_map<(d0) -> (d0 * 6)>                                                                                                                                                         
#map19 = affine_map<(d0, d1) -> (d0 * 32 + 32, d1 * 6 + 7)>                                                                                                                                   
#map20 = affine_map<(d0, d1) -> (d0 * -6 + d1)>                                                                                                                                               
#map21 = affine_map<(d0, d1) -> (d0 * -2 + d1)>                                                                                                                                               
#set = affine_set<(d0, d1) : ((d1 - 1) floordiv 3 - d0 >= 0)>                                                                                                                                 
#set1 = affine_set<(d0, d1, d2) : ((d1 - 1) floordiv 2 - d0 >= 0, d1 - d2 ceildiv 3 >= 0)>                                                                                                    
module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<!llvm.ptr, dense<64> : vector<4xi32>>, #dlti.dl_entry<i8, dense<8> : vector<2xi32>>, #dlti.dl_entry<i1, dense<8> : vector<2xi32
>>, #dlti.dl_entry<i64, dense<64> : vector<2xi32>>, #dlti.dl_entry<f80, dense<128> : vector<2xi32>>, #dlti.dl_entry<f64, dense<64> : vector<2xi32>>, #dlti.dl_entry<!llvm.ptr<270>, dense<32> 
: vector<4xi32>>, #dlti.dl_entry<f128, dense<128> : vector<2xi32>>, #dlti.dl_entry<!llvm.ptr<271>, dense<32> : vector<4xi32>>, #dlti.dl_entry<!llvm.ptr<272>, dense<64> : vector<4xi32>>, #dlt
i.dl_entry<i32, dense<32> : vector<2xi32>>, #dlti.dl_entry<i16, dense<16> : vector<2xi32>>, #dlti.dl_entry<f16, dense<16> : vector<2xi32>>, #dlti.dl_entry<"dlti.stack_alignment", 128 : i32>,
 #dlti.dl_entry<"dlti.endianness", "little">>, llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", llvm.target_triple = "x86_64-unknown-linux-gnu", "
polygeist.target-cpu" = "x86-64", "polygeist.target-features" = "+cmov,+cx8,+fxsr,+mmx,+sse,+sse2,+x87", "polygeist.tune-cpu" = "generic"} {                                                  
  func.func private @S0(%arg0: memref<1xi64>, %arg1: memref<153xi64>, %arg2: index) attributes {scop.stmt} {                                                                                  
    %0 = affine.load %arg1[symbol(%arg2)] : memref<153xi64>                                                                                                                                   
    affine.store %0, %arg0[0] : memref<1xi64>                                                                                                                                                 
    return                                                                                                                                                                                    
  }                                                                                                                                                                                           
  func.func private @S1(%arg0: memref<7xi64>, %arg1: index, %arg2: memref<1xi64>, %arg3: memref<153x3x7xi64>, %arg4: index, %arg5: index) attributes {scop.stmt} {                            
    %0 = affine.load %arg0[symbol(%arg1)] : memref<7xi64>
    %1 = affine.load %arg3[symbol(%arg4), symbol(%arg5), symbol(%arg1)] : memref<153x3x7xi64>
    %2 = affine.load %arg2[0] : memref<1xi64>
    %3 = arith.muli %1, %2 : i64
    %4 = arith.subi %0, %3 : i64
    affine.store %4, %arg0[symbol(%arg1)] : memref<7xi64>
    return
  }
  func.func @kernel(%arg0: memref<153xi64>, %arg1: memref<153x3x7xi64>, %arg2: memref<7xi64>) attributes {llvm.linkage = #llvm.linkage<external>} {
    %c2 = arith.constant 2 : index
    %alloca = memref.alloca() {scop.scratchpad} : memref<1xi64>
    affine.for %arg3 = 0 to 5 {
      affine.for %arg4 = #map(%arg3) to min #map1(%arg3) {
        affine.for %arg5 = max #map2(%arg3, %arg4) to min #map3(%arg3, %arg4) {
          affine.if #set(%arg4, %arg5) {
            affine.for %arg6 = #map4(%arg4, %arg5) to min #map5(%arg4, %arg5) {
              affine.for %arg7 = #map6(%arg5) to #map7(%arg4, %arg5) {
                %0 = affine.apply #map8(%arg4, %arg5, %arg7)
                %1 = affine.apply #map9(%arg4, %arg5)
                %2 = affine.apply #map10(%arg4, %arg6, %arg5)
                func.call @S1(%arg2, %0, %alloca, %arg1, %1, %2) : (memref<7xi64>, index, memref<1xi64>, memref<153x3x7xi64>, index, index) -> ()
              }
            }
          }
          affine.if #set1(%arg3, %arg4, %arg5) {
            affine.for %arg6 = max #map11(%arg5, %arg4) to min #map12(%arg4, %arg5) {
              %0 = affine.apply #map13(%arg4, %arg6)
              %1 = affine.apply #map14(%arg4)
              func.call @S1(%arg2, %0, %alloca, %arg1, %1, %c2) : (memref<7xi64>, index, memref<1xi64>, memref<153x3x7xi64>, index, index) -> ()
            }
          }
          affine.for %arg6 = max #map15(%arg5, %arg4) to min #map16(%arg5, %arg3, %arg4) {
            func.call @S0(%alloca, %arg0, %arg6) : (memref<1xi64>, memref<153xi64>, index) -> ()
            affine.for %arg7 = #map(%arg6) to min #map17(%arg4, %arg6) {
              affine.for %arg8 = #map18(%arg6) to min #map19(%arg5, %arg6) {
                %0 = affine.apply #map20(%arg6, %arg8)
                %1 = affine.apply #map21(%arg7, %arg6)
                func.call @S1(%arg2, %0, %alloca, %arg1, %arg6, %1) : (memref<7xi64>, index, memref<1xi64>, memref<153x3x7xi64>, index, index) -> ()
              }
            }
          }
        }
      }
    }
    return
  }
}

The 448th invocation of S1 by this IR is with %arg1 = 2, %arg4 = 21, and %arg5 = 418, which causes %arg3 to be indexed with a negative value, ultimately resulting in a segmentation fault.

This behavior can be reproduced with the code from this repository with a minimal non-working example:

$ git clone https://github.com/andidr/Polygeist-mnwe
$ cd Polygeist-mnwe
$ make mnwe.polymer
cgeist kernel.c --function=kernel -S -I. --memref-fullrank --raise-scf-to-affine -o kernel.cgeist.mlir
polymer-opt kernel.cgeist.mlir -reg2mem -extract-scop-stmt -pluto-opt -o kernel.polymer.mlir --allow-unregistered-dialect >/dev/null 2>&1
mlir-opt \
        -lower-affine \
        -convert-scf-to-cf \
        -cse \
        -canonicalize \
        -convert-func-to-llvm="use-bare-ptr-memref-call-conv" \
        --finalize-memref-to-llvm \
        -canonicalize \
        -o kernel.polymer.low.mlir kernel.polymer.mlir
mlir-translate -mlir-to-llvmir -o kernel.polymer.ll kernel.polymer.low.mlir
llc -opaque-pointers -o kernel.polymer.S kernel.polymer.ll
as -o kernel.polymer.o kernel.polymer.S
gcc -O3 -I. main.c kernel.polymer.o -o mnwe.polymer
$ ./mnwe.polymer 
Segmentation fault

Compiling the kernel with plain cgeist and without any polymer optimizations yields the expected results:

$ make mnwe.cgeist 
cgeist -c kernel.c --function=kernel -I. --memref-fullrank -o kernel.cgeist.o -O3
gcc -O3 -I. -o mnwe.cgeist main.c kernel.cgeist.o
$ ./mnwe.cgeist 
res: 8

which are the same as the results produced by a binary compiled using a plain C compiler:

$ make mnwe.O3
gcc -O3 -I. -c kernel.c -o kernel.O3.o
gcc -O3 -I. main.c kernel.O3.o -o mnwe.O3
$  ./mnwe.O3 
res: 8

I experienced similar issues with imperfect loop nests directly with pluto in the past and this may be related. Could somebody confirm that this is indeed a bug in pluto? Any suggestions for fixes or workarounds are welcome, too.

Thanks!

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions