feat(compiler): Add operation interface for batchable operations

This adds a new operation interface that allows an operation to
specify that a batched version of the operation exists that applies it
on the elements of a flat tensor in parallel.

compiler/include/concretelang/CMakeLists.txt

@@ -1,3 +1,4 @@
 add_subdirectory(Dialect)
 add_subdirectory(Conversion)
 add_subdirectory(Transforms)
+add_subdirectory(Interfaces)

compiler/include/concretelang/Interfaces/BatchableInterface.h

@@ -0,0 +1,14 @@
+// Part of the Concrete Compiler Project, under the BSD3 License with Zama
+// Exceptions. See
+// https://github.com/zama-ai/concrete-compiler-internal/blob/main/LICENSE.txt
+// for license information.
+
+#ifndef CONCRETELANG_INTERFACES_BATCHABLEINTERFACE_H_
+#define CONCRETELANG_INTERFACES_BATCHABLEINTERFACE_H_
+
+#include <mlir/IR/ImplicitLocOpBuilder.h>
+#include <mlir/IR/OpDefinition.h>
+
+#include <concretelang/Interfaces/BatchableInterface.h.inc>
+
+#endif // CONCRETELANG_INTERFACES_BATCHABLEINTERFACE_H_

compiler/include/concretelang/Interfaces/BatchableInterface.td

@@ -0,0 +1,52 @@
+#ifndef CONCRETELANG_INTERFACES_BATCHABLEINTERFACE
+#define CONCRETELANG_INTERFACES_BATCHABLEINTERFACE
+
+include "mlir/IR/OpBase.td"
+
+def BatchableOpInterface : OpInterface<"BatchableOpInterface"> {
+  let description = [{
+      Interface for operations processing a scalar that can be batched
+    if invoked multiple times with different, independent operands.
+  }];
+  let cppNamespace = "::mlir::concretelang";
+
+  let methods = [
+    InterfaceMethod<[{
+        Return the scalar operand that can be batched in a tensor to
+        be passed to the corresponding batched operation.
+      }],
+      /*retTy=*/"::mlir::OpOperand&",
+      /*methodName=*/"getBatchableOperand",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        llvm_unreachable("getBatchableOperand not implemented");
+      }]
+    >,
+    InterfaceMethod<[{
+        Return all operands that cannot be batched.
+      }],
+      /*retTy=*/"::mlir::OperandRange",
+      /*methodName=*/"getNonBatchableOperands",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        llvm_unreachable("getNonBatchableOperands not implemented");
+      }]
+    >,
+    InterfaceMethod<[{
+        Create the batched operation and return it as a value.
+      }],
+      /*retTy=*/"::mlir::Value",
+      /*methodName=*/"createBatchedOperation",
+      /*args=*/(ins "::mlir::ImplicitLocOpBuilder&":$builder,
+                    "::mlir::Value":$batchedOperands),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        llvm_unreachable("createBatchedOperation not implemented");
+      }]
+    >
+  ];
+}
+
+#endif // CONCRETELANG_INTERFACES_BATCHABLEINTERFACE

compiler/include/concretelang/Interfaces/CMakeLists.txt

@@ -0,0 +1 @@
+add_mlir_interface(BatchableInterface)

compiler/lib/CMakeLists.txt

@@ -6,6 +6,7 @@ add_subdirectory(Runtime)
 add_subdirectory(ClientLib)
 add_subdirectory(Bindings)
 add_subdirectory(ServerLib)
+add_subdirectory(Interfaces)
 
 # CAPI needed only for python bindings
 if (CONCRETELANG_BINDINGS_PYTHON_ENABLED)

compiler/lib/Interfaces/BatchableInterface.cpp

@@ -0,0 +1,8 @@
+// Part of the Concrete Compiler Project, under the BSD3 License with Zama
+// Exceptions. See
+// https://github.com/zama-ai/concrete-compiler-internal/blob/main/LICENSE.txt
+// for license information.
+
+#include <concretelang/Interfaces/BatchableInterface.h>
+
+#include <concretelang/Interfaces/BatchableInterface.cpp.inc>

compiler/lib/Interfaces/CMakeLists.txt

@@ -0,0 +1,12 @@
+add_mlir_library(ConcretelangInterfaces
+  BatchableInterface.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${PROJECT_SOURCE_DIR}/concretelang/Interfaces
+
+  DEPENDS
+  mlir-headers
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+)