[EmitHLSCpp] support AffineIfOp with explicit AffineYieldOp as terminator for returning results

lib/EmitHLSCpp/EmitHLSCpp.cpp

@@ -125,6 +125,7 @@ public:
             // Affine statements.
             AffineForOp, AffineIfOp, AffineParallelOp, AffineApplyOp,
             AffineMaxOp, AffineMinOp, AffineLoadOp, AffineStoreOp,
+            AffineYieldOp,
             // Memref-related statements.
             AllocOp, LoadOp, StoreOp,
             // Unary expressions.
@@ -137,10 +138,9 @@ public:
             UnsignedDivIOp, UnsignedRemIOp, XOrOp, AndOp, OrOp, ShiftLeftOp,
             SignedShiftRightOp, UnsignedShiftRightOp,
             // Special operations.
-            AffineYieldOp, ConstantOp, ReturnOp>(
-            [&](auto opNode) -> ResultType {
-              return thisCast->visitOp(opNode, args...);
-            })
+            ConstantOp, ReturnOp>([&](auto opNode) -> ResultType {
+          return thisCast->visitOp(opNode, args...);
+        })
         .Default([&](auto opNode) -> ResultType {
           return thisCast->visitInvalidOp(op, args...);
         });
@@ -172,6 +172,7 @@ public:
   HANDLE(AffineMinOp);
   HANDLE(AffineLoadOp);
   HANDLE(AffineStoreOp);
+  HANDLE(AffineYieldOp);
 
   // Memref-related statements.
   HANDLE(AllocOp);
@@ -218,7 +219,6 @@ public:
   HANDLE(UnsignedShiftRightOp);
 
   // Special operations.
-  HANDLE(AffineYieldOp);
   HANDLE(ConstantOp);
   HANDLE(ReturnOp);
 #undef HANDLE
@@ -244,6 +244,7 @@ public:
   void emitAffineFor(AffineForOp *op);
   void emitAffineIf(AffineIfOp *op);
   void emitAffineParallel(AffineParallelOp *op);
+  void emitAffineYield(AffineYieldOp *op);
 
   /// Memref-related statement emitters.
   void emitAlloc(AllocOp *op);
@@ -344,13 +345,12 @@ public:
   bool visitOp(AffineParallelOp op) {
     return emitter.emitAffineParallel(&op), true;
   }
-
-  /// Affine statements (without region).
   bool visitOp(AffineApplyOp op) { return true; }
   bool visitOp(AffineMaxOp op) { return true; }
   bool visitOp(AffineMinOp op) { return true; }
   bool visitOp(AffineLoadOp op) { return true; }
   bool visitOp(AffineStoreOp op) { return true; }
+  bool visitOp(AffineYieldOp op) { return emitter.emitAffineYield(&op), true; }
 
   /// Memref related statements.
   bool visitOp(AllocOp op) { return emitter.emitAlloc(&op), true; }
@@ -358,7 +358,6 @@ public:
   bool visitOp(StoreOp op) { return emitter.emitStore(&op), true; }
 
   /// Special operations.
-  bool visitOp(AffineYieldOp op) { return true; }
   bool visitOp(ConstantOp op) { return true; }
   bool visitOp(ReturnOp op) { return true; }
 
@@ -532,7 +531,15 @@ void ModuleEmitter::emitAffineFor(AffineForOp *op) {
   os << "}\n";
 }
 
-void ModuleEmitter::emitAffineIf(mlir::AffineIfOp *op) {
+void ModuleEmitter::emitAffineIf(AffineIfOp *op) {
+  // Declare all values returned by AffineYieldOp. They will be further handled
+  // by the AffineYieldOp emitter.
+  for (auto result : op->getResults()) {
+    indent();
+    emitValue(result);
+    os << ";\n";
+  }
+
   indent();
   os << "if (";
   auto constrSet = op->getIntegerSet();
@@ -568,6 +575,21 @@ void ModuleEmitter::emitAffineIf(mlir::AffineIfOp *op) {
 
 void ModuleEmitter::emitAffineParallel(AffineParallelOp *op) { return; }
 
+void ModuleEmitter::emitAffineYield(AffineYieldOp *op) {
+  // For now, only AffineIfOp may yield values.
+  if (auto affineIf = dyn_cast<AffineIfOp>(op->getParentOp())) {
+    unsigned resultIdx = 0;
+    for (auto result : affineIf.getResults()) {
+      indent();
+      emitValue(result);
+      os << " = ";
+      emitValue(op->getOperand(resultIdx));
+      os << ";\n";
+      resultIdx += 1;
+    }
+  }
+}
+
 /// Memref-related statement emitters.
 void ModuleEmitter::emitAlloc(AllocOp *op) {
   indent();

test/EmitHLSCpp/test_affine.mlir

@@ -11,7 +11,7 @@
 func @test_affine(%arg0: i32, %arg1: memref<16xi32>, %arg2: index) -> () {
   %c11 = constant 11 : index
   // CHECK: for (int val[[INT4:.*]] = 0; val[[INT4:.*]] < min(min((val[[INT3:.*]] + 11), val[[INT3:.*]]), (val[[INT3:.*]] + (11 * (-1)))); val[[INT4:.*]] += 1) {
-  affine.for %i = 0 to min #map0(%arg2)[%c11] {
+  affine.for %i = 0 to min #map0 (%arg2)[%c11] {
     // CHECK: for (int val[[INT5:.*]] = 0; val[[INT5:.*]] < 16; val[[INT5:.*]] += 2) {
     affine.for %j = 0 to 16 step 2 {
       //  ap_int<32> val[[INT6:.*]] = val[[INT2:.*]][val[[INT4:.*]]];
@@ -20,8 +20,15 @@ func @test_affine(%arg0: i32, %arg1: memref<16xi32>, %arg2: index) -> () {
       %1 = addi %arg0, %0 : i32
       // CHECK: val[[INT2:.*]][val[[INT5:.*]]] = val[[INT7:.*]];
       store %1, %arg1[%j] : memref<16xi32>
-      affine.if #set0(%i)[%c11] {
+      %2:2 = affine.if #set0 (%i)[%c11] -> (i32, i32) {
         store %0, %arg1[%j] : memref<16xi32>
+        %3 = muli %arg0, %1 : i32
+        %4 = subi %arg0, %1 : i32
+        affine.yield %3, %4 : i32, i32
+      } else {
+        %5 = shift_left %arg0, %1 : i32
+        %6 = divi_signed %arg0, %1 : i32
+        affine.yield %5, %6 : i32, i32
       }
     // CHECK: }
     }