[Substrait] Add folder for project op with empty expressions.

Signed-off-by: Ingo Müller <[email protected]>
iree-org · May 31, 2024 · 51e3db6 · 51e3db6
1 parent 54029b4
commit 51e3db6
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 0 deletions.
diff --git a/include/structured/Dialect/Substrait/IR/SubstraitOps.td b/include/structured/Dialect/Substrait/IR/SubstraitOps.td
@@ -395,6 +395,7 @@ def Substrait_ProjectOp : Substrait_RelOp<"project", [
     $input attr-dict `:` type($input) `->` type($result) $expressions
   }];
   let hasRegionVerifier = 1;
+  let hasFolder = 1;
   let extraClassDefinition = [{
     /// Implement OpAsmOpInterface.
     ::llvm::StringRef $cppClass::getDefaultDialect() {

diff --git a/lib/Dialect/Substrait/IR/Substrait.cpp b/lib/Dialect/Substrait/IR/Substrait.cpp
@@ -339,6 +339,17 @@ LogicalResult PlanRelOp::verifyRegions() {
   return verifyNamedStruct(getOperation(), fieldNames, tupleType);
 }
 
+OpFoldResult ProjectOp::fold(FoldAdaptor adaptor) {
+  Operation *terminator = adaptor.getExpressions().front().getTerminator();
+
+  // If the region does not yield any values, the the `project` has no effect.
+  if (terminator->getNumOperands() == 0) {
+    return getInput();
+  }
+
+  return {};
+}
+
 LogicalResult ProjectOp::verifyRegions() {
   // Verify that the expression block has a matching argument type.
   auto inputTupleType = llvm::cast<TupleType>(getInput().getType());

diff --git a/test/Dialect/Substrait/canonicalize.mlir b/test/Dialect/Substrait/canonicalize.mlir
@@ -72,3 +72,22 @@ substrait.plan version 0 : 42 : 1 {
     yield %5 : tuple<si1, si32>
   }
 }
+
+// -----
+
+// Check that empty `project` folded.
+
+// CHECK-LABEL: substrait.plan
+// CHECK-NEXT:    relation
+// CHECK-NEXT:      %[[V0:.*]] = named_table
+// CHECK-NEXT:      yield %[[V0]]
+
+substrait.plan version 0 : 42 : 1 {
+  relation {
+    %0 = named_table @t1 as ["a"] : tuple<si32>
+    %1 = project %0 : tuple<si32> -> tuple<si32> {
+    ^bb0(%arg0: tuple<si32>):
+    }
+    yield %1 : tuple<si32>
+  }
+}