Yohann/support-int-types

.gitlab/subscribed-pipelines.yml

@@ -30,19 +30,19 @@
 ###
 
 # Dane
-dane-up-check:
-  variables:
-    CI_MACHINE: "dane"
-  extends: [.machine-check]
+# dane-up-check:
+#   variables:
+#     CI_MACHINE: "dane"
+#   extends: [.machine-check]
 
-dane-build-and-test:
-  variables:
-    CI_MACHINE: "dane"
-    JOB_CMD:
-      value: "scripts/gitlab/ci-build-test.sh"
-      expand: false
-  needs: [dane-up-check]
-  extends: [.build-and-test]
+# dane-build-and-test:
+#   variables:
+#     CI_MACHINE: "dane"
+#     JOB_CMD:
+#       value: "scripts/gitlab/ci-build-test.sh"
+#       expand: false
+#   needs: [dane-up-check]
+#   extends: [.build-and-test]
 
 # TIOGA
 tioga-up-check:

examples/ideal_gas/app/eos_ams.cpp

@@ -45,7 +45,7 @@ void AMSEOS<FPType>::Eval(const int length,
   inputs.push_back(
       std::move(AMSTensor::view(density, {length, 1}, {1, 1}, res_)));
   inputs.push_back(
-      std::move(AMSTensor::view(density, {length, 1}, {1, 1}, res_)));
+      std::move(AMSTensor::view(energy, {length, 1}, {1, 1}, res_)));
 
   SmallVector<AMSTensor> inout;
   SmallVector<AMSTensor> outputs;

src/AMSlib/AMSTensor.cpp

@@ -33,6 +33,26 @@ bool AMSTensor::isContiguous(AMSTensor::IntDimType expected_stride) const
   return true;
 }
 
+namespace
+{
+template <typename T>
+constexpr AMSDType scalar_to_ams_dtype()
+{
+  using U = std::remove_cv_t<T>;
+  if constexpr (std::is_same_v<U, float>) {
+    return AMS_SINGLE;
+  } else if constexpr (std::is_same_v<U, double>) {
+    return AMS_DOUBLE;
+  } else if constexpr (std::is_same_v<U, int32_t>) {
+    return AMS_INT32;
+  } else if constexpr (std::is_same_v<U, int64_t>) {
+    return AMS_INT64;
+  } else {
+    static_assert(!sizeof(T), "Unsupported AMS scalar type");
+  }
+}
+}  // namespace
+
 AMSTensor::AMSTensor(uint8_t* data,
                      ams::ArrayRef<AMSTensor::IntDimType> shapes,
                      ams::ArrayRef<AMSTensor::IntDimType> strides,
@@ -47,67 +67,43 @@ AMSTensor::AMSTensor(uint8_t* data,
       _location(location),
       _owned(!view)
 {
-  _bytes = _elements * _element_size;
   _elements = computeNumElements(shapes);
+  _bytes = _elements * _element_size;
   if (!_data) {
     throw std::runtime_error("Generating tensor with Null Pointer AMSTensor.");
   }
 }
 
-template <typename FPType, typename>
+template <typename ScalarType>
 AMSTensor AMSTensor::create(ams::ArrayRef<AMSTensor::IntDimType> shapes,
                             ams::ArrayRef<AMSTensor::IntDimType> strides,
                             AMSResourceType location)
 {
   auto numElements = computeNumElements(shapes);
   auto& rm = ams::ResourceManager::getInstance();
-  if constexpr ((std::is_same_v<FPType, float>) ||
-                (std::is_same_v<FPType, const float>)) {
-    float* _data = rm.allocate<float>(numElements, location, sizeof(float));
-    return AMSTensor(reinterpret_cast<uint8_t*>(_data),
-                     shapes,
-                     strides,
-                     AMS_SINGLE,
-                     location);
-  } else if constexpr ((std::is_same_v<FPType, double>) ||
-                       (std::is_same_v<FPType, const double>)) {
-    double* _data = rm.allocate<double>(numElements, location, sizeof(double));
-    return AMSTensor(reinterpret_cast<uint8_t*>(_data),
-                     shapes,
-                     strides,
-                     AMS_DOUBLE,
-                     location);
-  } else {
-    // This should never happen due to the type restriction
-    static_assert(std::is_same_v<FPType, float> ||
-                      std::is_same_v<FPType, double>,
-                  "AMSTensor only supports float or double tensor creation");
-  }
+  using U = std::remove_cv_t<ScalarType>;
+  U* data = rm.allocate<U>(numElements, location, sizeof(U));
+  return AMSTensor(reinterpret_cast<uint8_t*>(data),
+                   shapes,
+                   strides,
+                   scalar_to_ams_dtype<U>(),
+                   location);
 }
 
 
-template <typename FPType, typename>
-AMSTensor AMSTensor::view(FPType* data,
+template <typename ScalarType>
+AMSTensor AMSTensor::view(ScalarType* data,
                           ams::ArrayRef<AMSTensor::IntDimType> shapes,
                           ams::ArrayRef<AMSTensor::IntDimType> strides,
                           AMSResourceType location)
 {
-  if constexpr ((std::is_same_v<FPType, float>) ||
-                (std::is_same_v<FPType, const float>)) {
-    return AMSTensor(
-        (uint8_t*)data, shapes, strides, AMS_SINGLE, location, true);
-  } else if constexpr ((std::is_same_v<FPType, double>) ||
-                       (std::is_same_v<FPType, const double>)) {
-    return AMSTensor(
-        (uint8_t*)data, shapes, strides, AMS_DOUBLE, location, true);
-  } else {
-    static_assert(std::is_same_v<FPType, float> ||
-                      std::is_same_v<FPType, const float> ||
-                      std::is_same_v<FPType, const double> ||
-                      std::is_same_v<FPType, double>,
-                  "AMSTensor only supports float or double tensor view");
-  }
-  throw std::runtime_error("Should never get here\n");
+  using U = std::remove_cv_t<ScalarType>;
+  return AMSTensor(reinterpret_cast<uint8_t*>(const_cast<U*>(data)),
+                   shapes,
+                   strides,
+                   scalar_to_ams_dtype<U>(),
+                   location,
+                   true);
 }
 
 AMSTensor AMSTensor::view(AMSTensor& tensor)
@@ -122,6 +118,16 @@ AMSTensor AMSTensor::view(AMSTensor& tensor)
                            tensor._shape,
                            tensor._strides,
                            tensor._location);
+  else if (tensor._dType == AMS_INT32)
+    return AMSTensor::view((int32_t*)tensor._data,
+                           tensor._shape,
+                           tensor._strides,
+                           tensor._location);
+  else if (tensor._dType == AMS_INT64)
+    return AMSTensor::view((int64_t*)tensor._data,
+                           tensor._shape,
+                           tensor._strides,
+                           tensor._location);
   throw std::runtime_error(
       "Creating view through copying constructor has incorrect dtype");
 }
@@ -193,6 +199,10 @@ AMSTensor AMSTensor::transpose(AMSTensor::IntDimType axis1,
     return view((double*)_data, newShape, newStrides, _location);
   else if (dType() == AMSDType::AMS_SINGLE)
     return view((float*)_data, newShape, newStrides, _location);
+  else if (dType() == AMSDType::AMS_INT32)
+    return view((int32_t*)_data, newShape, newStrides, _location);
+  else if (dType() == AMSDType::AMS_INT64)
+    return view((int64_t*)_data, newShape, newStrides, _location);
   // NOTE: Use defensive programming here and just crash. We can fix a better interface later
   // for error handling.
   throw std::runtime_error("Unknow data type in transpose\n");
@@ -204,6 +214,12 @@ template AMSTensor AMSTensor::create<float>(ams::ArrayRef<IntDimType>,
 template AMSTensor AMSTensor::create<double>(ams::ArrayRef<IntDimType>,
                                              ams::ArrayRef<IntDimType>,
                                              AMSResourceType);
+template AMSTensor AMSTensor::create<int32_t>(ams::ArrayRef<IntDimType>,
+                                              ams::ArrayRef<IntDimType>,
+                                              AMSResourceType);
+template AMSTensor AMSTensor::create<int64_t>(ams::ArrayRef<IntDimType>,
+                                              ams::ArrayRef<IntDimType>,
+                                              AMSResourceType);
 
 template AMSTensor AMSTensor::view<float>(float*,
                                           ams::ArrayRef<IntDimType>,
@@ -213,6 +229,14 @@ template AMSTensor AMSTensor::view<double>(double*,
                                            ams::ArrayRef<IntDimType>,
                                            ams::ArrayRef<IntDimType>,
                                            AMSResourceType);
+template AMSTensor AMSTensor::view<int32_t>(int32_t*,
+                                            ams::ArrayRef<IntDimType>,
+                                            ams::ArrayRef<IntDimType>,
+                                            AMSResourceType);
+template AMSTensor AMSTensor::view<int64_t>(int64_t*,
+                                            ams::ArrayRef<IntDimType>,
+                                            ams::ArrayRef<IntDimType>,
+                                            AMSResourceType);
 
 template AMSTensor AMSTensor::view<const float>(const float*,
                                                 ams::ArrayRef<IntDimType>,
@@ -222,3 +246,11 @@ template AMSTensor AMSTensor::view<const double>(const double*,
                                                  ams::ArrayRef<IntDimType>,
                                                  ams::ArrayRef<IntDimType>,
                                                  AMSResourceType);
+template AMSTensor AMSTensor::view<const int32_t>(const int32_t*,
+                                                  ams::ArrayRef<IntDimType>,
+                                                  ams::ArrayRef<IntDimType>,
+                                                  AMSResourceType);
+template AMSTensor AMSTensor::view<const int64_t>(const int64_t*,
+                                                  ams::ArrayRef<IntDimType>,
+                                                  ams::ArrayRef<IntDimType>,
+                                                  AMSResourceType);

src/AMSlib/include/AMSTensor.hpp

@@ -66,8 +66,7 @@ class AMSTensor
    * @param[in] location The memory location (e.g., CPU, GPU).
    * @return A new AMSTensor with allocated memory.
    */
-  template <typename FPType,
-            typename = std::enable_if_t<std::is_floating_point<FPType>::value>>
+  template <typename ScalarType>
   static AMSTensor create(ams::ArrayRef<IntDimType> shapes,
                           ams::ArrayRef<IntDimType> strides,
                           AMSResourceType location);
@@ -81,9 +80,8 @@ class AMSTensor
    * @param[in] location The memory location (e.g., CPU, GPU).
    * @return A new AMSTensor that acts as a view of the existing data.
    */
-  template <typename FPType,
-            typename = std::enable_if_t<std::is_floating_point<FPType>::value>>
-  static AMSTensor view(FPType* data,
+  template <typename ScalarType>
+  static AMSTensor view(ScalarType* data,
                         ams::ArrayRef<IntDimType> shapes,
                         ams::ArrayRef<IntDimType> strides,
                         AMSResourceType location);
@@ -150,4 +148,12 @@ extern template AMSTensor AMSTensor::create<double>(
     ams::ArrayRef<AMSTensor::IntDimType> shapes,
     ams::ArrayRef<AMSTensor::IntDimType> strides,
     AMSResourceType location);
+extern template AMSTensor AMSTensor::create<int32_t>(
+    ams::ArrayRef<AMSTensor::IntDimType> shapes,
+    ams::ArrayRef<AMSTensor::IntDimType> strides,
+    AMSResourceType location);
+extern template AMSTensor AMSTensor::create<int64_t>(
+    ams::ArrayRef<AMSTensor::IntDimType> shapes,
+    ams::ArrayRef<AMSTensor::IntDimType> strides,
+    AMSResourceType location);
 }  // namespace ams

src/AMSlib/include/AMSTypes.hpp

@@ -2,7 +2,13 @@
 
 namespace ams
 {
-typedef enum { AMS_SINGLE = 0, AMS_DOUBLE, AMS_UNKNOWN_TYPE } AMSDType;
+typedef enum {
+  AMS_SINGLE = 0,
+  AMS_DOUBLE,
+  AMS_INT32,
+  AMS_INT64,
+  AMS_UNKNOWN_TYPE
+} AMSDType;
 
 typedef enum {
   AMS_UNKNOWN = -1,

src/AMSlib/wf/interface.cpp

@@ -33,8 +33,8 @@ static AMSDType torchDTypeToAMSType(torch::Dtype dtype)
       {torch::kFloat, AMSDType::AMS_SINGLE},  // Alias for float32
       {torch::kFloat64, AMSDType::AMS_DOUBLE},
       {torch::kDouble, AMSDType::AMS_DOUBLE},  // Alias for float64
-      {torch::kInt32, AMSDType::AMS_UNKNOWN_TYPE},
-      {torch::kInt64, AMSDType::AMS_UNKNOWN_TYPE},
+      {torch::kInt32, AMSDType::AMS_INT32},
+      {torch::kInt64, AMSDType::AMS_INT64},
       {torch::kBool, AMSDType::AMS_UNKNOWN_TYPE},
       {torch::kUInt8, AMSDType::AMS_UNKNOWN_TYPE},
       {torch::kInt8, AMSDType::AMS_UNKNOWN_TYPE},
@@ -66,6 +66,10 @@ static c10::ScalarType amsToTorchDType(const ams::AMSDType dType)
     return torch::kFloat32;
   else if (dType == ams::AMSDType::AMS_DOUBLE)
     return torch::kFloat64;
+  else if (dType == ams::AMSDType::AMS_INT32)
+    return torch::kInt32;
+  else if (dType == ams::AMSDType::AMS_INT64)
+    return torch::kInt64;
 
   throw std::runtime_error("Unknown ams data type");
   return torch::kHalf;
@@ -91,6 +95,15 @@ static ams::SmallVector<ams::AMSTensor> torchToAMSTensors(
     } else if (dType == AMSDType::AMS_DOUBLE) {
       ams_tensors.push_back(
           AMSTensor::view(tensor.data_ptr<double>(), shapes, strides, rType));
+    } else if (dType == AMSDType::AMS_INT32) {
+      ams_tensors.push_back(
+          AMSTensor::view(tensor.data_ptr<int32_t>(), shapes, strides, rType));
+    } else if (dType == AMSDType::AMS_INT64) {
+      ams_tensors.push_back(
+          AMSTensor::view(tensor.data_ptr<int64_t>(), shapes, strides, rType));
+    } else {
+      throw std::runtime_error(
+          "torchToAMSTensors: unsupported tensor scalar type");
     }
   }
   return ams_tensors;

src/AMSlib/wf/resource_manager.hpp

@@ -155,9 +155,9 @@ class ResourceManager
     std::string pinned_alloc("PINNED");
     if (!RMAllocators[AMSResourceType::AMS_HOST])
       setAllocator(host_alloc, AMSResourceType::AMS_HOST);
-#if defined(__AMS_ENABLE_CUDA__)
+#if defined(__AMS_ENABLE_CUDA__) || defined(__AMS_ENABLE_HIP__)
     if (!RMAllocators[AMSResourceType::AMS_DEVICE])
-      setAllocator(host_alloc, AMSResourceType::AMS_DEVICE);
+      setAllocator(device_alloc, AMSResourceType::AMS_DEVICE);
 
     if (!RMAllocators[AMSResourceType::AMS_PINNED])
       setAllocator(pinned_alloc, AMSResourceType::AMS_PINNED);

src/AMSlib/wf/utils.hpp

@@ -67,6 +67,10 @@ static inline size_t dtype_to_size(ams::AMSDType dType)
       return sizeof(double);
     case ams::AMSDType::AMS_SINGLE:
       return sizeof(float);
+    case ams::AMSDType::AMS_INT64:
+      return sizeof(int64_t);
+    case ams::AMSDType::AMS_INT32:
+      return sizeof(int32_t);
     default:
       throw std::runtime_error("Requesting the size of unknown object");
   }

tests/AMSlib/ams_interface/CMakeLists.txt

@@ -24,3 +24,6 @@ endfunction()
 BUILD_UNIT_TEST(ams_explicit_end_to_end ams_ete.cpp)
 ADD_AMS_UNIT_TEST(AMS_EXPLICIT ams_explicit_end_to_end)
 
+BUILD_UNIT_TEST(int_interface int_interface.cpp)
+ADD_AMS_UNIT_TEST(AMS_INT_INTERFACE int_interface)
+