[MLIR][GPU] Drop mgpuMemHostRegisterMemRef's dependence on LLVM Support

[bondhugula]

Summary

Drop mgpuMemHostRegisterMemRef's dependence on LLVM Support. This
method is the only one in CUDA runtime wrappers library that creates
a dependence on libLLVMSupport due to its use of SmallVector and
ArrayRef. The code can be as easily/compactly written without those ADT.
The dependence on LLVMSupport adds a significant amount of additional
complexity for external things that want to link this library in (both
statically or as a shared object) since libLLVMSupport includes numerous
other objects that are sensitive to C++ compiler version and ABI.

[bondhugula]

Use uint64_t for size.

[herhut]

Thanks, makes sense to keep the runtime as self contained as possible. Should this also reflect in cmake somewhere?

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

14–15

Is this still needed?

[bondhugula]

Remove unused headers.

[bondhugula]

In D108684#2966816, @herhut wrote:

Thanks, makes sense to keep the runtime as self contained as possible. Should this also reflect in cmake somewhere?

I checked the cmake file and I couldn't find anything to remove for mlir_cuda_runtime.

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

14–15

Thanks, removed.

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

 //===- CudaRuntimeWrappers.cpp - MLIR CUDA API wrapper library ------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Implements C wrappers around the CUDA library for easy linking in ORC jit.
 // Also adds some debugging helpers that are helpful when writing MLIR code to
 // run on GPUs.
 //
 //===----------------------------------------------------------------------===//
 
-#include <cassert>
-#include <numeric>
-
 #include "mlir/ExecutionEngine/CRunnerUtils.h"
-#include "llvm/ADT/ArrayRef.h"
 
 #include "cuda.h"
 
 #ifdef _WIN32
 #define MLIR_CUDA_WRAPPERS_EXPORT __declspec(dllexport)
 #else
 #define MLIR_CUDA_WRAPPERS_EXPORT
 #endif // _WIN32

 // Allows to register byte array with the CUDA runtime. Helpful until we have
 // transfer functions implemented.
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
 mgpuMemHostRegister(void *ptr, uint64_t sizeBytes) {
   ScopedContext scopedContext;
   CUDA_REPORT_IF_ERROR(cuMemHostRegister(ptr, sizeBytes, /*flags=*/0));
 }
 
-// Allows to register a MemRef with the CUDA runtime. Helpful until we have
-// transfer functions implemented.
+/// Registers a memref with the CUDA runtime. `descriptor` is a pointer to a
+/// ranked memref descriptor struct of rank `rank`. Helpful until we have
+/// transfer functions implemented.
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
 mgpuMemHostRegisterMemRef(int64_t rank, StridedMemRefType<char, 1> *descriptor,
                           int64_t elementSizeBytes) {
-
-  llvm::SmallVector<int64_t, 4> denseStrides(rank);
-  llvm::ArrayRef<int64_t> sizes(descriptor->sizes, rank);
-  llvm::ArrayRef<int64_t> strides(sizes.end(), rank);
-
-  std::partial_sum(sizes.rbegin(), sizes.rend(), denseStrides.rbegin(),
-                   std::multiplies<int64_t>());
-  auto sizeBytes = denseStrides.front() * elementSizeBytes;
-
   // Only densely packed tensors are currently supported.
-  std::rotate(denseStrides.begin(), denseStrides.begin() + 1,
-              denseStrides.end());
-  denseStrides.back() = 1;
-  assert(strides == llvm::makeArrayRef(denseStrides));
+  int64_t *denseStrides = (int64_t *)alloca(rank * sizeof(int64_t));
+  int64_t *sizes = descriptor->sizes;
+  for (int64_t i = rank - 1, runningStride = 1; i >= 0; i--) {
+    denseStrides[i] = runningStride;
+    runningStride *= sizes[i];
+  }
+  uint64_t sizeBytes = sizes[0] * denseStrides[0] * elementSizeBytes;
+  int64_t *strides = &sizes[rank];
+  for (unsigned i = 0; i < rank; ++i)
+    assert(strides[i] == denseStrides[i] &&
+           "Mismatch in computed dense strides");
 
-  auto ptr = descriptor->data + descriptor->offset * elementSizeBytes;
+  auto *ptr = descriptor->data + descriptor->offset * elementSizeBytes;
   mgpuMemHostRegister(ptr, sizeBytes);
 }