update unit test

Author: yifeizh2

lib/gc/Transforms/Pipeline.cpp

@@ -104,6 +104,7 @@ void populateCPURuntimePasses(mlir::PassManager &pm) {
 }
 
 void populateLoweringToLLVMPasses(mlir::PassManager &pm) {
+  pm.addPass(createLowerAffinePass());
   pm.addPass(createFinalizeMemRefToLLVMConversionPass());
   pm.addPass(createConvertSCFToCFPass());
   pm.addPass(cpuruntime::createCPURuntimeToLLVM());

test/gc/Transform/flashAttention.mlir

@@ -1,7 +1,31 @@
-// RUN: gc-opt --split-input-file --flash-attention-conversion %s
+// RUN: gc-opt --split-input-file --flash-attention-conversion --gc-cpu-pipeline %s | gc-cpu-runner -e main -entry-point-result=void
 
 func.func @flash_attention(%arg0: tensor<1x16x384x64xf32>, %arg1: tensor<1x16x384x64xf32>, %arg2: tensor<1x16x384x64xf32>, %arg3: tensor<1x16x384x384xf32>) -> tensor<1x16x384x64xf32> {
     %0 = tensor.empty() : tensor<1x16x384x64xf32>
     %1 = linalgx.scaled_dot_product_attention ins(%arg0, %arg1, %arg2, %arg3: tensor<1x16x384x64xf32>, tensor<1x16x384x64xf32>, tensor<1x16x384x64xf32>, tensor<1x16x384x384xf32>) outs(%0 : tensor<1x16x384x64xf32>)  -> tensor<1x16x384x64xf32>
     return %1 : tensor<1x16x384x64xf32>
 }
+
+func.func @main() {
+  %cst = arith.constant 1.000000e+00 : f32
+
+  %QKVShape  = tensor.empty() : tensor<1x16x384x64xf32>
+  %maskShape  = tensor.empty() : tensor<1x16x384x384xf32>
+
+  %Q = linalg.fill ins(%cst : f32) outs(%QKVShape : tensor<1x16x384x64xf32>) -> tensor<1x16x384x64xf32>
+  %K = linalg.fill ins(%cst : f32) outs(%QKVShape : tensor<1x16x384x64xf32>) -> tensor<1x16x384x64xf32>
+  %V = linalg.fill ins(%cst : f32) outs(%QKVShape : tensor<1x16x384x64xf32>) -> tensor<1x16x384x64xf32>
+  %mask = linalg.fill ins(%cst : f32) outs(%maskShape : tensor<1x16x384x384xf32>) -> tensor<1x16x384x384xf32>
+
+  %out = func.call @flash_attention(%Q, %K, %V, %mask) :
+  (tensor<1x16x384x64xf32>, tensor<1x16x384x64xf32>, tensor<1x16x384x64xf32>, tensor<1x16x384x384xf32>)
+    -> (tensor<1x16x384x64xf32>)
+
+  %idx = arith.constant 0 : index
+  %val = tensor.extract %out[%idx, %idx, %idx, %idx] : tensor<1x16x384x64xf32>
+  cpuruntime.printf "output[0, 0, 0]: %f\n" %val : f32
+
+  return
+}
+// CHECK: output[0, 0, 0]: 1.0
+