[Mlir-commits] [mlir] mlir: add an operation to EmitC for function template instantiation (PR #100895)

[Rohan Yadav]

https://github.com/rohany created https://github.com/llvm/llvm-project/pull/100895

This commit adds an `emitc.instantiate_function_template` operation to allow for the expression of function template instantiation. Without this operation, there is no easy way to express a C++ program like:
```
auto x = ...;
auto y = ...;
const void* fptr = &f<decltype(x), decltype(y)>;
```
Doing so is necessary to generate code that interacts with some lower level APIs for launching parallel work into runtime systems.

>From 2372a35e7bd1eba460a5011104aed9bc076cc9cd Mon Sep 17 00:00:00 2001
From: Rohan Yadav <rohany at alumni.cmu.edu>
Date: Sat, 27 Jul 2024 12:37:30 -0700
Subject: [PATCH] mlir: add an operation to EmitC for function template
 instantiation

This commit adds an `emitc.instantiate_function_template` operation to
allow for the expression of function template instantiation. Without
this operation, there is no easy way to express a C++ program like:
```
auto x = ...;
auto y = ...;
const void* fptr = &f<decltype(x), decltype(y)>;
```
Doing so is necessary to generate code that interacts with some lower
level APIs for launching parallel work into runtime systems.
---
 mlir/include/mlir/Dialect/EmitC/IR/EmitC.td | 15 ++++++++++++
 mlir/lib/Target/Cpp/TranslateToCpp.cpp      | 26 +++++++++++++++++++++
 mlir/test/Dialect/EmitC/ops.mlir            |  6 +++++
 3 files changed, 47 insertions(+)

diff --git a/mlir/include/mlir/Dialect/EmitC/IR/EmitC.td b/mlir/include/mlir/Dialect/EmitC/IR/EmitC.td
index 452302c565139..83f2e14d9407d 100644
--- a/mlir/include/mlir/Dialect/EmitC/IR/EmitC.td
+++ b/mlir/include/mlir/Dialect/EmitC/IR/EmitC.td
@@ -1260,5 +1260,20 @@ def EmitC_SubscriptOp : EmitC_Op<"subscript", []> {
   let assemblyFormat = "$value `[` $indices `]` attr-dict `:` functional-type(operands, results)";
 }
 
+def EmitC_InstantiateFunctionTemplateOp : EmitC_Op<"instantiate_function_template", []> {
+  let summary = "Instantiate template operation";
+  let description = [{
+    Instantiate a function template with a given set of types
+    (given by the values as argument to this operation) to obtain
+    a function pointer.
+  }];
+  let arguments = (ins
+    Arg<StrAttr, "the C++ function to instantiate">:$callee,
+    Variadic<EmitCType>:$args
+  );
+  let results = (outs EmitC_PointerType);
+  let assemblyFormat = "$callee `(` $args `)` attr-dict `:` functional-type($args, results)";
+}
+
 
 #endif // MLIR_DIALECT_EMITC_IR_EMITC
diff --git a/mlir/lib/Target/Cpp/TranslateToCpp.cpp b/mlir/lib/Target/Cpp/TranslateToCpp.cpp
index 626638282efe1..60c5c7fb906fb 100644
--- a/mlir/lib/Target/Cpp/TranslateToCpp.cpp
+++ b/mlir/lib/Target/Cpp/TranslateToCpp.cpp
@@ -656,6 +656,30 @@ static LogicalResult printOperation(CppEmitter &emitter,
   return success();
 }
 
+static LogicalResult
+printOperation(CppEmitter &emitter,
+               emitc::InstantiateFunctionTemplateOp instOp) {
+
+  raw_ostream &os = emitter.ostream();
+  Operation &op = *instOp.getOperation();
+
+  if (failed(emitter.emitAssignPrefix(op)))
+    return failure();
+  os << "&" << instOp.getCallee() << "<";
+
+  auto emitArgs = [&](mlir::Value val) -> LogicalResult {
+    os << "decltype(";
+    if (failed(emitter.emitOperand(val)))
+      return failure();
+    os << ")";
+    return success();
+  };
+  if (failed(interleaveCommaWithError(instOp.getArgs(), os, emitArgs)))
+    return failure();
+  os << ">";
+  return success();
+}
+
 static LogicalResult printOperation(CppEmitter &emitter,
                                     emitc::ApplyOp applyOp) {
   raw_ostream &os = emitter.ostream();
@@ -1508,6 +1532,8 @@ LogicalResult CppEmitter::emitOperation(Operation &op, bool trailingSemicolon) {
           .Case<func::CallOp, func::FuncOp, func::ReturnOp>(
               [&](auto op) { return printOperation(*this, op); })
           .Case<emitc::LiteralOp>([&](auto op) { return success(); })
+          .Case<emitc::InstantiateFunctionTemplateOp>(
+              [&](auto op) { return printOperation(*this, op); })
           .Default([&](Operation *) {
             return op.emitOpError("unable to find printer for op");
           });
diff --git a/mlir/test/Dialect/EmitC/ops.mlir b/mlir/test/Dialect/EmitC/ops.mlir
index 51c484a633eec..9ddcb7b539a48 100644
--- a/mlir/test/Dialect/EmitC/ops.mlir
+++ b/mlir/test/Dialect/EmitC/ops.mlir
@@ -224,6 +224,12 @@ func.func @test_subscript(%arg0 : !emitc.array<2x3xf32>, %arg1 : !emitc.ptr<i32>
   return
 }
 
+func.func @test_instantiate_template() {
+  %c1 = "emitc.constant"() <{value = 7 : i32}> : () -> i32
+  %0 = emitc.instantiate_function_template "func_template"(%c1) : (i32) -> !emitc.ptr<!emitc.opaque<"void">>
+  return
+}
+
 emitc.verbatim "#ifdef __cplusplus"
 emitc.verbatim "extern \"C\" {"
 emitc.verbatim "#endif  // __cplusplus"