CodeGen: chose binary operator based on variable type

wpmed92 · Aug 14, 2024 · 5a0cb1b · 5a0cb1b
1 parent f99685d
commit 5a0cb1b
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 5 deletions.
diff --git a/include/AST/Types.h b/include/AST/Types.h
@@ -128,6 +128,18 @@ class Type {
     return other.kind == kind;
   }
 
+  virtual bool isIntLike() {
+    return kind == TypeKind::Integer || kind == TypeKind::UnsignedInteger;
+  }
+
+  virtual bool isUintLike() {
+    return kind == TypeKind::UnsignedInteger;
+  }
+
+  virtual bool isFloatLike() {
+    return kind == TypeKind::Float || kind == TypeKind::Double;
+  }
+
   virtual std::string toString() {
     switch (kind) {
       case TypeKind::Integer:
@@ -217,6 +229,18 @@ class VectorType : public Type {
     return false;
   }
 
+  bool isIntLike() override {
+    return elementType->getKind() == TypeKind::Integer || elementType->getKind() == TypeKind::UnsignedInteger;
+  }
+
+  bool isUintLike() override {
+    return elementType->getKind() == TypeKind::UnsignedInteger;
+  }
+
+  bool isFloatLike() override {
+    return elementType->getKind() == TypeKind::Float || elementType->getKind() == TypeKind::Double;
+  }
+
   std::string toString() override {
     std::string prefix = "";
 

diff --git a/include/CodeGen/MLIRCodeGen.h b/include/CodeGen/MLIRCodeGen.h
@@ -97,6 +97,7 @@ class MLIRCodeGen : public ASTVisitor {
 
   mlir::Value popExpressionStack();
   mlir::Value currentBasePointer;
+  Type* typeContext;
 };
 
 }; // namespace codegen

diff --git a/lib/CodeGen/MLIRCodeGen.cpp b/lib/CodeGen/MLIRCodeGen.cpp
@@ -77,23 +77,49 @@ void MLIRCodeGen::visit(BinaryExpression *binExp) {
 
   switch (binExp->getOp()) {
   case BinaryOperator::Add:
-    val = builder.create<spirv::FAddOp>(loc, lhs, rhs);
+    if (typeContext->isIntLike()) {
+      val = builder.create<spirv::IAddOp>(loc, lhs, rhs);
+    } else {
+      val = builder.create<spirv::FAddOp>(loc, lhs, rhs);
+    }
     expressionStack.push_back(val);
     break;
   case BinaryOperator::Sub:
-    val = builder.create<spirv::FSubOp>(loc, lhs, rhs);
+    if (typeContext->isIntLike()) {
+      val = builder.create<spirv::ISubOp>(loc, lhs, rhs);
+    } else {
+      val = builder.create<spirv::FSubOp>(loc, lhs, rhs);
+    }
+
     expressionStack.push_back(val);
     break;
   case BinaryOperator::Mul:
-    val = builder.create<spirv::FMulOp>(loc, lhs, rhs);
+  if (typeContext->isIntLike()) {
+      val = builder.create<spirv::IMulOp>(loc, lhs, rhs);
+    } else {
+      val = builder.create<spirv::FMulOp>(loc, lhs, rhs);
+    }
+
     expressionStack.push_back(val);
     break;
   case BinaryOperator::Div:
-    val = builder.create<spirv::FDivOp>(loc, lhs, rhs);
+    if (typeContext->isUintLike()) {
+      val = builder.create<spirv::UDivOp>(loc, lhs, rhs);
+    } else if (typeContext->isIntLike()) {
+      val = builder.create<spirv::SDivOp>(loc, lhs, rhs);
+    } else {
+      val = builder.create<spirv::FDivOp>(loc, lhs, rhs);
+    }
+
     expressionStack.push_back(val);
     break;
   case BinaryOperator::Mod:
-    val = builder.create<spirv::FRemOp>(loc, lhs, rhs);
+    if (typeContext->isIntLike()) {
+      val = builder.create<spirv::SRemOp>(loc, lhs, rhs);
+    } else {
+      val = builder.create<spirv::FRemOp>(loc, lhs, rhs);
+    }
+
     expressionStack.push_back(val);
     break;
   case BinaryOperator::ShiftLeft:
@@ -240,6 +266,7 @@ void MLIRCodeGen::visit(VariableDeclarationList *varDeclList) {
 void MLIRCodeGen::createVariable(shaderpulse::Type *type,
                                  VariableDeclaration *varDecl) {
   shaderpulse::Type *varType = (type) ? type : varDecl->getType();
+  typeContext = varType;
 
   if (inGlobalScope) {
     std::cout << "In global scope" << std::endl;