From 5d96c825cedb11f121cbd7d243e3f26c113c5a3c Mon Sep 17 00:00:00 2001
From: "rizwan.ashraf" <rizwan.ashraf@knights.ucf.edu>
Date: Mon, 28 Nov 2022 18:22:06 -0800
Subject: [PATCH] WIP: code-gen of for-loop and support in ta-dialect #2

---
 frontends/comet_dsl/mlir/MLIRGen.cpp          | 48 +++++++++++++++++++
 .../comet/Dialect/TensorAlgebra/IR/TAOps.td   | 19 ++++++++
 2 files changed, 67 insertions(+)

diff --git a/frontends/comet_dsl/mlir/MLIRGen.cpp b/frontends/comet_dsl/mlir/MLIRGen.cpp
index b26cf6ec..0cd12a44 100644
--- a/frontends/comet_dsl/mlir/MLIRGen.cpp
+++ b/frontends/comet_dsl/mlir/MLIRGen.cpp
@@ -1662,6 +1662,40 @@ namespace
       return t;
     }
 
+    /// Codegen for-loop 
+    mlir::LogicalResult mlirGen(ForLoopExprAST &forLoop)
+    {
+      comet_debug() << "codegen: ForLoopExprAST \n";
+
+      mlir::Value lo = builder.create<mlir::ConstantIndexOp>(
+          loc(forLoop.loc()), forLoop.getBegin());
+      mlir::Value hi = builder.create<mlir::ConstantIndexOp>(loc(forLoop.loc()),
+                                                             forLoop.getEnd());
+      mlir::Value step = builder.create<mlir::ConstantIndexOp>(
+          loc(forLoop.loc()), forLoop.getIncrement());
+
+      mlir::Value value = builder.create<IndexLabelStaticOp>(loc(forLoop.loc()), lo, hi, step);
+    
+      builder.create<ForLoopStartOp>(loc(forLoop.loc()), value, forLoop.getName());
+
+      if (failed(declare(forLoop.getName(), value)))
+        return mlir::failure();
+
+      comet_debug() << "codegen: ForLoopExprAST done \n";
+      return mlir::success();
+    }
+
+    /// Codegen for-loop 
+    mlir::LogicalResult mlirGen(ForLoopEndExprAST &forLoopEnd)
+    {
+      comet_debug() << "codegen: ForLoopEndExprAST \n";
+
+      builder.create<ForLoopEndOp>(loc(forLoopEnd.loc()));
+
+      comet_debug() << "codegen: ForLoopEndExprAST done \n";
+      return mlir::success();
+    }
+
     /// Codegen a list of expression, return failure if one of them hit an error.
     mlir::LogicalResult mlirGen(ExprASTList &blockAST)
     {
@@ -1727,6 +1761,20 @@ namespace
           continue;
         }
 
+        if (auto *forLoopStart = dyn_cast<ForLoopExprAST>(expr.get()))
+        {
+          if (mlir::failed(mlirGen(*forLoopStart)))
+            return mlir::success();
+          continue;
+        }
+
+        if (auto *forLoopEnd = dyn_cast<ForLoopEndExprAST>(expr.get()))
+        {
+          if (mlir::failed(mlirGen(*forLoopEnd)))
+            return mlir::success();
+          continue;
+        }
+
         if (auto *tensor_op = dyn_cast<TensorOpExprAST>(expr.get()))
         {
           comet_debug() << " generate ops for TensorOpExprAST\n";
diff --git a/include/comet/Dialect/TensorAlgebra/IR/TAOps.td b/include/comet/Dialect/TensorAlgebra/IR/TAOps.td
index 087098db..de118606 100644
--- a/include/comet/Dialect/TensorAlgebra/IR/TAOps.td
+++ b/include/comet/Dialect/TensorAlgebra/IR/TAOps.td
@@ -441,6 +441,25 @@ def PrintElapsedTimeOp : TA_Op<"print_elapsed_time"> {
   let arguments = (ins F64:$start, F64:$end);
 }
 
+def ForLoopStartOp : TA_Op<"start_for_loop">
+{
+  let summary = "identifies the start of a for-loop body";
+  let description = [{
+    The "start_for_loop" identifies the start of a for-loop body
+  }];
+
+  // the inputs are 1) index_label_static (index start, index end, index step)
+  // and 2) string-attr with info about the iterator
+  let arguments = (ins Range:$label, StrAttr:$iterator);
+}
+
+def ForLoopEndOp : TA_Op<"end_for_loop"> {
+  let summary = "identifies the end of a for-loop body";
+  let description = [{
+    The "end_for_loop" identifies the end of a for-loop body
+  }];
+}
+
 def TAReturnOp : TA_Op<"return", [Terminator, HasParent<"FuncOp">]> {
   let summary = "return operation";
   let description = [{