How can I use PDL to match affine code?

MLIR
#1

Hi, I would like to ask for help to match “afine.for” instructions using PDL. I couldn’t find examples of how to match blocks.
I have this gemm input code that uses affine:

#map0 = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<() -> (0)>
#map2 = affine_map<()[s0] -> (s0)>
"builtin.module"() ( {
  "builtin.func"() ( {
  ^bb0(%arg0: i32, %arg1: i32, %arg2: i32, %arg3: f64, %arg4: f64, %arg5: memref<1000x1100xf64>, %arg6: memref<1000x1200xf64>, %arg7: memref<1200x1100xf64>):  // no predecessors
    %0 = "arith.index_cast"(%arg1) : (i32) -> index
    %1 = "arith.index_cast"(%arg2) : (i32) -> index
    %2 = "arith.index_cast"(%arg0) : (i32) -> index
    "affine.for"(%2) ( {
    ^bb0(%arg8: index):  // no predecessors
      "affine.for"(%0) ( {
      ^bb0(%arg9: index):  // no predecessors
        %3 = "affine.load"(%arg5, %arg8, %arg9) {map = #map0} : (memref<1000x1100xf64>, index, index) -> f64
        %4 = "arith.mulf"(%3, %arg4) : (f64, f64) -> f64
        "affine.store"(%4, %arg5, %arg8, %arg9) {map = #map0} : (f64, memref<1000x1100xf64>, index, index) -> ()
        "affine.yield"() : () -> ()
      }) {lower_bound = #map1, step = 1 : index, upper_bound = #map2} : (index) -> ()
      "affine.for"(%1) ( {
      ^bb0(%arg9: index):  // no predecessors
        "affine.for"(%0) ( {
        ^bb0(%arg10: index):  // no predecessors
          %3 = "affine.load"(%arg6, %arg8, %arg9) {map = #map0} : (memref<1000x1200xf64>, index, index) -> f64
          %4 = "arith.mulf"(%arg3, %3) : (f64, f64) -> f64
          %5 = "affine.load"(%arg7, %arg9, %arg10) {map = #map0} : (memref<1200x1100xf64>, index, index) -> f64
          %6 = "arith.mulf"(%4, %5) : (f64, f64) -> f64
          %7 = "affine.load"(%arg5, %arg8, %arg10) {map = #map0} : (memref<1000x1100xf64>, index, index) -> f64
          %8 = "arith.addf"(%7, %6) : (f64, f64) -> f64
          "affine.store"(%8, %arg5, %arg8, %arg10) {map = #map0} : (f64, memref<1000x1100xf64>, index, index) -> ()
          "affine.yield"() : () -> ()
        }) {lower_bound = #map1, step = 1 : index, upper_bound = #map2} : (index) -> ()
        "affine.yield"() : () -> ()
      }) {lower_bound = #map1, step = 1 : index, upper_bound = #map2} : (index) -> ()
      "affine.yield"() : () -> ()
    }) {lower_bound = #map1, step = 1 : index, upper_bound = #map2} : (index) -> ()
    "std.return"() : () -> ()
  }) {sym_name = "kernel_gemm", type = (i32, i32, i32, f64, f64, memref<1000x1100xf64>, memref<1000x1200xf64>, memref<1200x1100xf64>) -> ()} : () -> ()
}) {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"} : () -> ()

Is it possible to match this?

#2

PDL (currently) lacks support for matching regions. You can still match affine.for ops if you don’t care about what’s in the region.

#3

As @Mogball says, PDL itself currently has no support for constraining/building Blocks or Regions. You can still match the affine.for operation, but you won’t be able to constrain any characteristics about the regions/or create an affine.for directly via PDL. It is something that we intend to add in the coming months, but for now any constraints/rewrites that want to operate on control flow (branching or region) will need to go through a native constraint/rewrite.

– River

#4

Thank you both for the answers. I hope to use this feature in the future.