Script to generate FileCheck

As we are developing our infrastructure to test ONNX lowering, I wrote a python script that takes mlir output and transform it into FileCheck format, allowing for user to define variable names as well as deriving some names too (e.g. loading from a memref “A” result in a “LOAD_A_MEM” name, renaming loop indices as “I”, …). It is currently in our repo, happy to migrate to MLIR if there is an interest.

Example for a GEMM operation:

mlid-opt ... | mlir2FileCheck.py --args '["A", "B", "C"]' --names '{"cst": "alpha", "cst_0": "beta", "cst_1": "zero"}'

results in the following FileCheck output

// CHECK-LABEL:  func @test_gemm_all_dyn
// CHECK-SAME:   ([[A_:%.+]]: memref<?x?xf32>, [[B_:%.+]]: memref<?x?xf32>, [[C_:%.+]]: memref<?xf32>) -> memref<?x?xf32> {
// CHECK-DAG:       [[ALPHA_:%.+]] = constant 1.000000e+00 : f32
// CHECK-DAG:       [[BETA_:%.+]] = constant 5.000000e+00 : f32
// CHECK-DAG:       [[ZERO_:%.+]] = constant 0.000000e+00 : f32
// CHECK-DAG:       [[CST_1_:%.+]] = constant 1 : index
// CHECK-DAG:       [[CST_0_:%.+]] = constant 0 : index
// CHECK-NOT: separator of consecutive DAGs
// CHECK-DAG:       [[DIM_0_:%.+]] = dim [[A_]], [[CST_1_]] : memref<?x?xf32>
// CHECK-DAG:       [[DIM_1_:%.+]] = dim [[A_]], [[CST_0_]] : memref<?x?xf32>
// CHECK-DAG:       [[DIM_2_:%.+]] = dim [[B_]], [[CST_1_]] : memref<?x?xf32>
// CHECK-DAG:       [[DIM_3_:%.+]] = dim [[C_]], [[CST_0_]] : memref<?xf32>
// CHECK:           [[RES_:%.+]] = alloc([[DIM_0_]], [[DIM_2_]]) : memref<?x?xf32>
// CHECK:           affine.for [[I_0_:%.+]] = 0 to [[DIM_0_]] {
// CHECK:             affine.for [[I_1_:%.+]] = 0 to [[DIM_2_]] {
// CHECK:               affine.store [[ZERO_]], [[RES_]][symbol([[I_0_]]), symbol([[I_1_]])] : memref<?x?xf32>
// CHECK:               affine.for [[I_2_:%.+]] = 0 to [[DIM_1_]] {
// CHECK-DAG:             [[LOAD_A_MEM_:%.+]] = affine.load [[A_]][symbol([[I_2_]]), symbol([[I_0_]])] : memref<?x?xf32>
// CHECK-DAG:             [[LOAD_B_MEM_:%.+]] = affine.load [[B_]][symbol([[I_2_]]), symbol([[I_1_]])] : memref<?x?xf32>
// CHECK-DAG:             [[LOAD_RES_MEM_:%.+]] = affine.load [[RES_]][symbol([[I_0_]]), symbol([[I_1_]])] : memref<?x?xf32>
// CHECK:                 [[VAR_15_:%.+]] = mulf [[LOAD_A_MEM_]], [[LOAD_B_MEM_]] : f32
// CHECK:                 [[VAR_16_:%.+]] = addf [[LOAD_RES_MEM_]], [[VAR_15_]] : f32
// CHECK:                 affine.store [[VAR_16_]], [[RES_]][symbol([[I_0_]]), symbol([[I_1_]])] : memref<?x?xf32>
// CHECK:               }
// CHECK:               [[VAR_5_:%.+]] = cmpi "sgt", [[DIM_3_]], [[CST_1_]] : index
// CHECK-DAG:           [[VAR_6_:%.+]] = select [[VAR_5_]], [[I_1_]], [[CST_0_]] : index
// CHECK-DAG:           [[LOAD_RES_MEM_1_:%.+]] = affine.load [[RES_]][symbol([[I_0_]]), symbol([[I_1_]])] : memref<?x?xf32>
// CHECK-NOT: separator of consecutive DAGs
// CHECK-DAG:           [[VAR_8_:%.+]] = mulf [[ALPHA_]], [[LOAD_RES_MEM_1_]] : f32
// CHECK-DAG:           [[LOAD_C_MEM_:%.+]] = load [[C_]]{{.}}[[VAR_6_]]{{.}} : memref<?xf32>
// CHECK:               [[VAR_10_:%.+]] = mulf [[BETA_]], [[LOAD_C_MEM_]] : f32
// CHECK:               [[VAR_11_:%.+]] = addf [[VAR_8_]], [[VAR_10_]] : f32
// CHECK:               affine.store [[VAR_11_]], [[RES_]][symbol([[I_0_]]), symbol([[I_1_]])] : memref<?x?xf32>
// CHECK:             }
// CHECK:           }
// CHECK:           return [[RES_]] : memref<?x?xf32>
// CHECK:         }

Let me know what you think

2 Likes

Is this intended to replace https://github.com/llvm/llvm-project/blob/master/mlir/utils/generate-test-checks.py ?

Aha, I did not know there was one… will look into it :slight_smile:

Alex, I also have a version of this that I’ve used at various times. see https://reviews.llvm.org/D83924.
I’ve been off on other things, so haven’t had a chance to return to some of the review comments.

Steve