From 20d19c216168a58b4055cdb6c7100463a6973fc9 Mon Sep 17 00:00:00 2001
From: Antoniu Pop <antoniu.pop@zama.ai>
Date: Thu, 4 Aug 2022 08:43:27 +0100
Subject: [PATCH] cleanup(dfr): remove unused key ids and unranked memref
 labels.

---
 .../include/concretelang/Runtime/DFRuntime.hpp  |  3 +--
 .../Runtime/distributed_generic_task_server.hpp | 17 ++++++-----------
 .../RT/Analysis/LowerDataflowTasksToRT.cpp      |  9 ++-------
 3 files changed, 9 insertions(+), 20 deletions(-)

diff --git a/compiler/include/concretelang/Runtime/DFRuntime.hpp b/compiler/include/concretelang/Runtime/DFRuntime.hpp
index 9055ece1d..d6b9103cc 100644
--- a/compiler/include/concretelang/Runtime/DFRuntime.hpp
+++ b/compiler/include/concretelang/Runtime/DFRuntime.hpp
@@ -28,8 +28,7 @@ bool _dfr_use_omp();
 typedef enum _dfr_task_arg_type {
   _DFR_TASK_ARG_BASE = 0,
   _DFR_TASK_ARG_MEMREF = 1,
-  _DFR_TASK_ARG_UNRANKED_MEMREF = 2,
-  _DFR_TASK_ARG_CONTEXT = 3
+  _DFR_TASK_ARG_CONTEXT = 2
 } _dfr_task_arg_type;
 
 static inline _dfr_task_arg_type _dfr_get_arg_type(uint64_t val) {
diff --git a/compiler/include/concretelang/Runtime/distributed_generic_task_server.hpp b/compiler/include/concretelang/Runtime/distributed_generic_task_server.hpp
index 31fd60b0a..8f02782ce 100644
--- a/compiler/include/concretelang/Runtime/distributed_generic_task_server.hpp
+++ b/compiler/include/concretelang/Runtime/distributed_generic_task_server.hpp
@@ -74,15 +74,14 @@ struct OpaqueInputData {
         param_sizes(std::move(_param_sizes)),
         param_types(std::move(_param_types)),
         output_sizes(std::move(_output_sizes)),
-        output_types(std::move(_output_types)), ksk_id(0), bsk_id(0) {}
+        output_types(std::move(_output_types)) {}
 
   OpaqueInputData(const OpaqueInputData &oid)
       : wfn_name(std::move(oid.wfn_name)), params(std::move(oid.params)),
         param_sizes(std::move(oid.param_sizes)),
         param_types(std::move(oid.param_types)),
         output_sizes(std::move(oid.output_sizes)),
-        output_types(std::move(oid.output_types)), ksk_id(oid.ksk_id),
-        bsk_id(oid.bsk_id) {}
+        output_types(std::move(oid.output_types)) {}
 
   friend class hpx::serialization::access;
   template <class Archive> void load(Archive &ar, const unsigned int version) {
@@ -91,7 +90,8 @@ struct OpaqueInputData {
     ar >> output_sizes >> output_types;
     for (size_t p = 0; p < param_sizes.size(); ++p) {
       char *param;
-      _dfr_checked_aligned_alloc((void **)&param, 64, param_sizes[p]);
+      _dfr_checked_aligned_alloc((void **)&param, sizeof(void *),
+                                 param_sizes[p]);
       ar >> hpx::serialization::make_array(param, param_sizes[p]);
       params.push_back((void *)param);
 
@@ -123,7 +123,6 @@ struct OpaqueInputData {
         params[p] =
             (void *)_dfr_node_level_runtime_context_manager->getContext();
       } break;
-      case _DFR_TASK_ARG_UNRANKED_MEMREF:
       default:
         HPX_THROW_EXCEPTION(hpx::no_success, "DFR: OpaqueInputData save",
                             "Error: invalid task argument type.");
@@ -157,7 +156,6 @@ struct OpaqueInputData {
         // Nothing to do now - TODO: pass key ids if these are not
         // unique for a computation.
       } break;
-      case _DFR_TASK_ARG_UNRANKED_MEMREF:
       default:
         HPX_THROW_EXCEPTION(hpx::no_success, "DFR: OpaqueInputData save",
                             "Error: invalid task argument type.");
@@ -172,8 +170,6 @@ struct OpaqueInputData {
   std::vector<uint64_t> param_types;
   std::vector<size_t> output_sizes;
   std::vector<uint64_t> output_types;
-  uint64_t ksk_id;
-  uint64_t bsk_id;
 };
 
 struct OpaqueOutputData {
@@ -193,7 +189,8 @@ struct OpaqueOutputData {
     ar >> output_sizes >> output_types;
     for (size_t p = 0; p < output_sizes.size(); ++p) {
       char *output;
-      _dfr_checked_aligned_alloc((void **)&output, 64, (output_sizes[p]));
+      _dfr_checked_aligned_alloc((void **)&output, sizeof(void *),
+                                 (output_sizes[p]));
 
       ar >> hpx::serialization::make_array(output, output_sizes[p]);
       outputs.push_back((void *)output);
@@ -221,7 +218,6 @@ struct OpaqueOutputData {
       case _DFR_TASK_ARG_CONTEXT: {
 
       } break;
-      case _DFR_TASK_ARG_UNRANKED_MEMREF:
       default:
         HPX_THROW_EXCEPTION(hpx::no_success, "DFR: OpaqueInputData save",
                             "Error: invalid task argument type.");
@@ -251,7 +247,6 @@ struct OpaqueOutputData {
       case _DFR_TASK_ARG_CONTEXT: {
 
       } break;
-      case _DFR_TASK_ARG_UNRANKED_MEMREF:
       default:
         HPX_THROW_EXCEPTION(hpx::no_success, "DFR: OpaqueInputData save",
                             "Error: invalid task argument type.");
diff --git a/compiler/lib/Dialect/RT/Analysis/LowerDataflowTasksToRT.cpp b/compiler/lib/Dialect/RT/Analysis/LowerDataflowTasksToRT.cpp
index 2b5b393de..b96a321d5 100644
--- a/compiler/lib/Dialect/RT/Analysis/LowerDataflowTasksToRT.cpp
+++ b/compiler/lib/Dialect/RT/Analysis/LowerDataflowTasksToRT.cpp
@@ -171,13 +171,8 @@ getSizeInBytes(Value val, Location loc, OpBuilder builder) {
 
   // Unranked memrefs should be lowered to just pointer + size, so we need 16
   // bytes.
-  if (type.isa<mlir::UnrankedMemRefType>()) {
-    Value arg_type = builder.create<arith::ConstantOp>(
-        loc, builder.getI64IntegerAttr(dfr::_DFR_TASK_ARG_UNRANKED_MEMREF));
-    Value result =
-        builder.create<arith::ConstantOp>(loc, builder.getI64IntegerAttr(16));
-    return std::pair<mlir::Value, mlir::Value>(result, arg_type);
-  }
+  assert(!type.isa<mlir::UnrankedMemRefType>() &&
+         "UnrankedMemRefType not currently supported");
 
   Value arg_type = builder.create<arith::ConstantOp>(
       loc, builder.getI64IntegerAttr(dfr::_DFR_TASK_ARG_BASE));