search

minisparrow / aboutAI

keep track of the state of art AI
1 stars 0 forks source link

[ONNX-MLIR]onnx.conv to krnl.affine #12

Open minisparrow opened 2 years ago

minisparrow commented 2 years ago

input mlir

`func private @test_conv_no_bias_no_pad(%arg0 : tensor<1x2x32x64xf32>, %arg1 : tensor<5x2x6x7xf32>) -> tensor<*xf32> {
%cst = constant unit

%0 = "onnx.Conv"(%arg0, %arg1, %cst) {auto_pad = "NOTSET", group = 1 : si64} : (tensor<1x2x32x64xf32>, tensor<5x2x6x7xf32>, none) -> tensor<*xf32>

"std.return"(%0) : (tensor<*xf32>) -> ()
} `

minisparrow commented 2 years ago

onnx-mlir-opt demo_conv.mlir --shape-inference --convert-onnx-to-krnl --canonicalize --convert-krnl-to-affine

map0 = affine_map<(d0, d1) -> (d0 * 5 + d1)>

map1 = affine_map<(d0) -> (-d0, 0)>

map2 = affine_map<(d0) -> (-d0 + 32, 6)>

map3 = affine_map<(d0, d1) -> (-d1, 0)>

map4 = affine_map<(d0, d1) -> (-d1 + 64, 7)>

map5 = affine_map<(d0, d1) -> (d0 + d1 * 2)>

map6 = affine_map<(d0, d1) -> (d0 + d1)>

module { func private @test_conv_no_bias_no_pad(%arg0: memref<1x2x32x64xf32>, %arg1: memref<5x2x6x7xf32>) -> memref<1x5x27x58xf32> { %cst = arith.constant 0.000000e+00 : f32 %0 = memref.alloc() {alignment = 16 : i64} : memref<1x5x27x58xf32> affine.for %arg2 = 0 to 1 { affine.for %arg3 = 0 to 1 { affine.for %arg4 = 0 to 5 { %1 = affine.apply #map0(%arg3, %arg4) affine.for %arg5 = 0 to 27 { affine.for %arg6 = 0 to 58 { %2 = memref.alloca() : memref affine.store %cst, %2[] : memref affine.for %arg7 = 0 to 2 { affine.for %arg8 = max #map1(%arg5) to min #map2(%arg5) { affine.for %arg9 = max #map3(%arg5, %arg6) to min #map4(%arg5, %arg6) { %4 = affine.apply #map5(%arg7, %arg3) %5 = affine.apply #map6(%arg8, %arg5) %6 = affine.apply #map6(%arg9, %arg6) %7 = affine.load %arg0[%arg2, %4, %5, %6] : memref<1x2x32x64xf32> %8 = affine.load %arg1[%1, %arg7, %arg8, %arg9] : memref<5x2x6x7xf32> %9 = affine.load %2[] : memref %10 = arith.mulf %7, %8 : f32 %11 = arith.addf %9, %10 : f32 affine.store %11, %2[] : memref } } } %3 = affine.load %2[] : memref affine.store %3, %0[%arg2, %1, %arg5, %arg6] : memref<1x5x27x58xf32> } } } } } return %0 : memref<1x5x27x58xf32> } }