Declare function arguments in symbol table, add function codegen tests

wpmed92 · Aug 25, 2024 · 4077b86 · 4077b86
1 parent d200ffe
commit 4077b86
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 9 deletions.
diff --git a/include/CodeGen/MLIRCodeGen.h b/include/CodeGen/MLIRCodeGen.h
@@ -29,6 +29,8 @@ struct SymbolTableEntry {
   VariableDeclaration* variable = nullptr;
   spirv::PointerType ptrType = nullptr;
   bool isGlobal = false;
+  bool isFunctionParam = false;
+  Type* type = nullptr;
 };
 
 class MLIRCodeGen : public ASTVisitor {
@@ -92,7 +94,7 @@ class MLIRCodeGen : public ASTVisitor {
   SymbolTableScopeT globalScope;
   SmallVector<Attribute, 4> interface;
 
-  void declare(SymbolTableEntry);
+  void declare(StringRef name, SymbolTableEntry entry);
   void createVariable(shaderpulse::Type *, VariableDeclaration *);
   void insertEntryPoint();
 

diff --git a/lib/CodeGen/MLIRCodeGen.cpp b/lib/CodeGen/MLIRCodeGen.cpp
@@ -284,11 +284,12 @@ void MLIRCodeGen::visit(UnaryExpression *unExp) {
   }
 }
 
-void MLIRCodeGen::declare(SymbolTableEntry entry) {
-  if (symbolTable.count(entry.variable->getIdentifierName()))
+void MLIRCodeGen::declare(StringRef name, SymbolTableEntry entry) {
+  if (symbolTable.count(name)) {
     return;
+  }
 
-  symbolTable.insert(entry.variable->getIdentifierName(), entry);
+  symbolTable.insert(name, entry);
 }
 
 void MLIRCodeGen::visit(VariableDeclarationList *varDeclList) {
@@ -335,7 +336,7 @@ void MLIRCodeGen::createVariable(shaderpulse::Type *type,
       varOp->setAttr("location", *locationOpt);
     }
 
-    declare({ mlir::Value(), varDecl, ptrType, /*isGlobal*/ true});
+    declare(varDecl->getIdentifierName(), { mlir::Value(), varDecl, ptrType, /*isGlobal*/ true});
     // Set OpDecorate through attributes
     // example:
     // varOp->setAttr(spirv::stringifyDecoration(spirv::Decoration::Invariant),
@@ -361,7 +362,7 @@ void MLIRCodeGen::createVariable(shaderpulse::Type *type,
       builder.create<spirv::StoreOp>(builder.getUnknownLoc(), var, val);
     }
 
-    declare({ var, varDecl, nullptr, /*isGlobal*/ false});
+    declare(varDecl->getIdentifierName(), { var, varDecl, nullptr, /*isGlobal*/ false});
   }
 }
 
@@ -649,7 +650,9 @@ void MLIRCodeGen::visit(CallExpression *callExp) {
 void MLIRCodeGen::visit(VariableExpression *varExp) {
   auto entry = symbolTable.lookup(varExp->getName());
 
-  if (entry.variable) {
+  if (entry.isFunctionParam) {
+    expressionStack.push_back(std::make_pair(entry.type, entry.value));
+  } else if (entry.variable) {
     Value val;
 
     if (entry.isGlobal) {
@@ -743,7 +746,6 @@ void MLIRCodeGen::visit(DiscardStatement *discardStmt) {}
 
 void MLIRCodeGen::visit(FunctionDeclaration *funcDecl) {
   insideEntryPoint = funcDecl->getName() == "main";
-
   SymbolTableScopeT varScope(symbolTable);
   std::vector<mlir::Type> paramTypes;
 
@@ -759,7 +761,7 @@ void MLIRCodeGen::visit(FunctionDeclaration *funcDecl) {
     }
   }
 
-  auto funcOp = builder.create<spirv::FuncOp>(
+  spirv::FuncOp funcOp = builder.create<spirv::FuncOp>(
       builder.getUnknownLoc(), funcDecl->getName(),
       builder.getFunctionType(
           paramTypes,
@@ -772,6 +774,12 @@ void MLIRCodeGen::visit(FunctionDeclaration *funcDecl) {
   auto entryBlock = funcOp.addEntryBlock();
   builder.setInsertionPointToStart(entryBlock);
 
+  // Declare params as variables in the current scope
+  for (int i = 0; i < funcDecl->getParams().size(); i++) {
+    auto &param = funcDecl->getParams()[i];
+    declare(param->getName(), {funcOp.getArgument(i), nullptr, nullptr, false, true, param->getType()});
+  }
+
   funcDecl->getBody()->accept(this);
 
   // Return insertion for void functions

diff --git a/test/CodeGen/functions.glsl b/test/CodeGen/functions.glsl
@@ -0,0 +1,17 @@
+// CHECK: spirv.func @add(%arg0: si32, %arg1: si32) -> si32 "None" {
+// CHECK-NEXT: %0 = spirv.IAdd %arg0, %arg1 : si32
+// CHECK-NEXT: spirv.ReturnValue %0 : si32
+int add(int a, int b) {
+    return a + b;
+}
+
+// CHECK:  spirv.func @main() "None" {
+void main() {
+    int a = 1;
+    int b = 2;
+
+    // CHECK: %2 = spirv.Load "Function" %0 : si32
+    // CHECK-NEXT: %3 = spirv.Load "Function" %1 : si32
+    // CHECK-NEXT: %4 = spirv.FunctionCall @add(%2, %3) : (si32, si32) -> si32
+    int c = add(a, b);
+}