Merge pull request hail-is#2 from tpoterba/tp-mlir-tracking

iris-garden · Nov 3, 2022 · 4589597 · 4589597
2 parents ca3d600 + 53b1bea
commit 4589597
Show file tree

Hide file tree

Showing 7 changed files with 64 additions and 6 deletions.
diff --git a/query/include/Dialect/Sandbox/IR/CMakeLists.txt b/query/include/Dialect/Sandbox/IR/CMakeLists.txt
@@ -1,2 +1,4 @@
 set(LLVM_TARGET_DEFINITIONS SandboxOps.td)
+mlir_tablegen(SandboxOpsEnums.h.inc -gen-enum-decls)
+mlir_tablegen(SandboxOpsEnums.cpp.inc -gen-enum-defs)
 add_mlir_dialect(SandboxOps sb)
diff --git a/query/include/Dialect/Sandbox/IR/Sandbox.h b/query/include/Dialect/Sandbox/IR/Sandbox.h
@@ -1,12 +1,18 @@
 #ifndef DIALECT_SANDBOX_SANDBOX_H
 #define DIALECT_SANDBOX_SANDBOX_H
 
-#include "mlir/IR/BuiltinTypes.h"
+
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/Interfaces/CastInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Builders.h"
+
 #include "Dialect/Sandbox/IR/SandboxOpsDialect.h.inc"
+#include "Dialect/Sandbox/IR/SandboxOpsEnums.h.inc"
 
 #define GET_TYPEDEF_CLASSES
 #include "Dialect/Sandbox/IR/SandboxOpsTypes.h.inc"

diff --git a/query/include/Dialect/Sandbox/IR/SandboxBase.td b/query/include/Dialect/Sandbox/IR/SandboxBase.td
@@ -4,6 +4,9 @@
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/OpBase.td"
 
+include "mlir/IR/EnumAttr.td"
+
+
 def Sandbox_Dialect : Dialect {
   let name = "sb";
   let summary = "Dialect for experimenting with MLIR";
@@ -19,13 +22,34 @@ class Sandbox_Type<string name, string typeMnemonic, list<Trait> traits = []>
 
 // Here is a simple definition of an "integer" type, with a width parameter.
 def Sandbox_Int : Sandbox_Type<"Int", "int"> {
-  let summary = "Hail 32 bit integer type";
+  let summary = "Sandbox 32 bit integer type";
   let description = [{
-    Hail 32 bit integer type
+    Sandbox 32 bit integer type
   }];
 
   /// Indicate that our type will add additional verification to the parameters.
   let genVerifyDecl = 0;
 }
 
+def Sandbox_Bool : Sandbox_Type<"Boolean", "bool"> {
+  let summary = "Sandbox 32 bit integer type";
+  let description = [{
+    Sandbox 32 bit integer type
+  }];
+
+  /// Indicate that our type will add additional verification to the parameters.
+  let genVerifyDecl = 0;
+}
+
+def Sandbox_CmpPredicateAttr : I64EnumAttr<
+    "CmpPredicate", "",
+    [
+      I64EnumAttrCase<"LT", 0, "lt">,
+      I64EnumAttrCase<"LTEQ", 1, "lteq">,
+      I64EnumAttrCase<"GT", 2, "gt">,
+      I64EnumAttrCase<"GTEQ", 3, "gteq">,
+      I64EnumAttrCase<"EQ", 4, "eq">,
+    ]> {
+}
+
 #endif // DIALECT_SANDBOX_SANDBOXBASE
diff --git a/query/include/Dialect/Sandbox/IR/SandboxOps.td b/query/include/Dialect/Sandbox/IR/SandboxOps.td
@@ -25,6 +25,26 @@ def Sandbox_AddIOp : Sandbox_Op<"addi", [NoSideEffect, Commutative]> {
   let hasCanonicalizer = true;
 }
 
+def BoolOp : Sandbox_Op<"bool", [ConstantLike, NoSideEffect]> {
+  let arguments = (ins BoolAttr:$value);
+  let results = (outs Sandbox_Bool:$output);
+  let assemblyFormat = "$value attr-dict `:` type($output)";
+}
+
+
+def ComparisonOp : Sandbox_Op<"compare", [NoSideEffect]> {
+  let arguments = (ins Sandbox_CmpPredicateAttr:$predicate, Sandbox_Int:$lhs, Sandbox_Int:$rhs);
+  let results = (outs Sandbox_Bool:$output);
+  let assemblyFormat = "$predicate `,` $lhs `,` $rhs attr-dict `:` type($output)";
+
+    let extraClassDeclaration = [{
+    static CmpPredicate getPredicateByName(llvm::StringRef name);
+  }];
+
+  let hasFolder = 0;
+  let hasCanonicalizer = 0;
+}
+
 def Sandbox_PrintOp : Sandbox_Op<"print"> {
   let arguments = (ins Sandbox_Int:$value);
   let assemblyFormat = "$value attr-dict";

diff --git a/query/lib/Dialect/Sandbox/IR/SandboxDialect.cpp b/query/lib/Dialect/Sandbox/IR/SandboxDialect.cpp
@@ -19,6 +19,7 @@ mlir::Operation *SandboxDialect::materializeConstant(mlir::OpBuilder &builder, m
 }
 
 #define GET_TYPEDEF_CLASSES
+
 #include "Dialect/Sandbox/IR/SandboxOpsTypes.cpp.inc"
 #include "Dialect/Sandbox/IR/SandboxOpsDialect.cpp.inc"
 

diff --git a/query/lib/Dialect/Sandbox/IR/SandboxOps.cpp b/query/lib/Dialect/Sandbox/IR/SandboxOps.cpp
@@ -7,12 +7,15 @@
 #include "mlir/IR/Location.h"
 #include <cstddef>
 
+#include "Dialect/Sandbox/IR/SandboxOpsEnums.cpp.inc"
+
 #define GET_OP_CLASSES
-#include "Dialect/Sandbox/IR/SandboxOps.cpp.inc"
 
+#include "Dialect/Sandbox/IR/SandboxOps.cpp.inc"
 namespace hail {
 namespace ir{
 
+
 mlir::OpFoldResult ConstantOp::fold(llvm::ArrayRef<mlir::Attribute> operands) {
   return valueAttr();
 }

diff --git a/query/test/current.mlir b/query/test/current.mlir
@@ -1,7 +1,9 @@
-func.func @foo(%in: !sb.int) -> (!sb.int) {
+func.func @foo(%in: !sb.int) -> (!sb.bool) {
     %i1 = sb.constant 5 : !sb.int
     %i2 = sb.constant 7 : !sb.int
     %i3 = sb.addi %in %i1
     %i4 = sb.addi %i3 %i2
-    func.return %i4 : !sb.int
+    %i5 = sb.constant 6 : !sb.int
+    %i6 = sb.compare eq, %i4, %i5 : !sb.bool
+    func.return %i6 : !sb.bool
 }