Added lit tests

Author: jessicacotturone21

src/enzyme_ad/jax/Dialect/Tessera/Ops.td

@@ -119,6 +119,10 @@ def CallOp : TesseraOp<"call",
 
   let results = (outs Variadic<AnyType>);
 
+  let assemblyFormat = [{
+    $callee `(` $operands `)` attr-dict `:` functional-type($operands, results)
+  }];
+  
   let builders = [
     OpBuilder<(ins "DefineOp":$callee, CArg<"ValueRange", "{}">:$operands), [{
       $_state.addOperands(operands);

src/enzyme_ad/jax/Passes/Tessera/FuncToTessera.cpp

@@ -18,6 +18,7 @@
 #include "src/enzyme_ad/jax/Dialect/Tessera/Dialect.h"
 #include "src/enzyme_ad/jax/Passes/Tessera/Passes.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/IR/BuiltinDialect.h"
 
 using namespace mlir;
 using namespace mlir::enzyme;
@@ -59,6 +60,20 @@ class FuncOpRewrite final : public OpRewritePattern<func::FuncOp> {
       funcOp.getBody().cloneInto(&tesseraDefineOp.getBody(),
                             tesseraDefineOp.getBody().end(),
                             mapper);
+
+      // Now walk through the cloned operations and convert func.return to tessera.return
+      tesseraDefineOp.walk([&](func::ReturnOp returnOp) {
+        rewriter.setInsertionPoint(returnOp);
+        rewriter.replaceOpWithNewOp<tessera::ReturnOp>(returnOp, returnOp.getOperands());
+      });
+      
+      // Convert func.call to tessera.call
+      tesseraDefineOp.walk([&](func::CallOp callOp) {
+        rewriter.setInsertionPoint(callOp);
+        rewriter.replaceOpWithNewOp<tessera::CallOp>(callOp, callOp.getResultTypes(),
+                                                     callOp.getOperands(),
+                                                     callOp->getAttrs());
+      });
     }
 
     rewriter.eraseOp(funcOp);
@@ -81,7 +96,7 @@ class CallOpRewrite final : public OpRewritePattern<func::CallOp> {
     Operation *calleeOp = SymbolTable::lookupSymbolIn(moduleOp, calleeAttr);
 
     // Only convert if the callee is a Tessera DefineOp
-    if (isa<tessera::DefineOp>(calleeOp))
+    if (!isa<tessera::DefineOp>(calleeOp))
       return rewriter.notifyMatchFailure(callOp, "Callee is not a Tessera DefineOp");
 
       rewriter.replaceOpWithNewOp<tessera::CallOp>(callOp, callOp.getResultTypes(),
@@ -122,15 +137,28 @@ namespace mlir::enzyme::tessera {
 struct FuncToTesseraPass
   : public PassWrapper<FuncToTesseraPass, OperationPass<ModuleOp>> {
 
+  StringRef getArgument() const final { return "func-to-tessera"; }
+  StringRef getDescription() const final { return "Convert func dialect to tessera dialect."; }
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<tessera::TesseraDialect>();
+  }
+
   void runOnOperation() override {
   MLIRContext *ctx = &getContext();
+
+  ConversionTarget target(*ctx);
+  target.addLegalDialect<tessera::TesseraDialect>();
+  target.addLegalDialect<BuiltinDialect>();
+  target.addIllegalDialect<func::FuncDialect>();
+
   RewritePatternSet patterns(ctx);
 
   patterns.add<FuncOpRewrite, CallOpRewrite, ReturnOpRewrite>(ctx);
 
-  if (failed(applyPatternsAndFoldGreedily(getOperation(),
-                                          std::move(patterns))))
+  if (failed(applyFullConversion(getOperation(), target, std::move(patterns))))
     signalPassFailure();
+
 }
 };
 

src/enzyme_ad/jax/Passes/Tessera/TesseraToFunc.cpp

@@ -44,7 +44,7 @@ class DefineOpRewrite final : public OpRewritePattern<tessera::DefineOp> {
 
 
     // Create the `func.func` op
-    auto funcOp = rewriter.create<tessera::DefineOp>(
+    auto funcOp = rewriter.create<func::FuncOp>(
       defineOp.getLoc(), defineOp.getName(), fnType);
 
 
@@ -111,6 +111,13 @@ namespace mlir::enzyme::tessera {
 struct TesseraToFuncPass
   : public PassWrapper<TesseraToFuncPass, OperationPass<ModuleOp>> {
 
+  StringRef getArgument() const final { return "tessera-to-func"; }
+  StringRef getDescription() const final { return "Convert tessera dialect to func dialect."; }
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<func::FuncDialect>();
+  }
+  
   void runOnOperation() override {
   MLIRContext *ctx = &getContext();
   RewritePatternSet patterns(ctx);

test/lit_tests/tessera/func-to-tessera.mlir

@@ -0,0 +1,30 @@
+// RUN: enzymexlamlir-opt %s -func-to-tessera | FileCheck %s
+
+// CHECK-LABEL: tessera.define @simple_func
+func.func @simple_func() {
+  // CHECK: tessera.return
+  func.return
+}
+
+// -----
+
+// CHECK-LABEL: tessera.define @func_with_args
+func.func @func_with_args(%arg0: i32, %arg1: f32) -> i32 {
+  // CHECK: tessera.return %arg0 : i32
+  func.return %arg0 : i32
+}
+
+// -----
+
+// CHECK-LABEL: tessera.define @helper
+func.func @helper() {
+  func.return
+}
+
+// CHECK-LABEL: tessera.define @func_with_call
+func.func @func_with_call() {
+  // CHECK: tessera.call @helper() : () -> ()
+  func.call @helper() : () -> ()
+  // CHECK: tessera.return
+  func.return
+}

test/lit_tests/tessera/roundtrip.mlir

@@ -0,0 +1,36 @@
+// RUN: enzymexlamlir-opt %s | FileCheck %s
+
+// CHECK-LABEL: tessera.define @foo
+tessera.define @foo() {
+  // CHECK: tessera.return
+  tessera.return
+}
+
+// -----
+
+// CHECK-LABEL: tessera.define @bar
+tessera.define @bar() -> i32 {
+  %c42_i32 = arith.constant 42 : i32
+  // CHECK: tessera.return %{{.*}} : i32
+  tessera.return %c42_i32 : i32
+}
+
+// -----
+
+// CHECK-LABEL: tessera.define @caller
+tessera.define @caller() {
+  // CHECK: tessera.call @foo() : () -> ()
+  tessera.call @foo() : () -> ()
+  // CHECK: tessera.return
+  tessera.return
+}
+
+// -----
+
+// CHECK-LABEL: tessera.define @with_args
+tessera.define @with_args(%arg0: i32, %arg1: f32) -> i32 {
+  // CHECK: %[[V0:.*]] = tessera.call @bar() : () -> i32
+  %0 = tessera.call @bar() : () -> i32
+  // CHECK: tessera.return %[[V0]] : i32
+  tessera.return %0 : i32
+}

test/lit_tests/tessera/tessera-to-func.mlir

@@ -0,0 +1,30 @@
+// RUN: enzymexlamlir-opt %s -tessera-to-func | FileCheck %s
+
+// CHECK-LABEL: func.func @simple_func
+tessera.define @simple_func() {
+  // CHECK: func.return
+  tessera.return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @func_with_args
+tessera.define @func_with_args(%arg0: i32, %arg1: f32) -> i32 {
+  // CHECK: func.return %arg0 : i32
+  tessera.return %arg0 : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @helper
+tessera.define @helper() {
+  tessera.return
+}
+
+// CHECK-LABEL: func.func @func_with_call
+tessera.define @func_with_call() {
+  // CHECK: func.call @helper() : () -> ()
+  tessera.call @helper() : () -> ()
+  // CHECK: func.return
+  tessera.return
+}