feat(mars): add native Mars operator specializations

src/native/cuda/mars/blas.h

@@ -0,0 +1,51 @@
+#ifndef INFINI_OPS_MARS_BLAS_H_
+#define INFINI_OPS_MARS_BLAS_H_
+
+#include <utility>
+
+// clang-format off
+#include <hcblas/hcblas.h>
+// clang-format on
+
+#include "data_type.h"
+#include "native/cuda/blas.h"
+#include "native/cuda/mars/blas_utils.h"
+#include "native/cuda/mars/runtime_.h"
+
+namespace infini::ops {
+
+template <>
+struct Blas<Device::Type::kMars> : public Runtime<Device::Type::kMars> {
+  using BlasHandle = hcblasHandle_t;
+
+  static constexpr auto BLAS_OP_N = HCBLAS_OP_N;
+
+  static constexpr auto BLAS_OP_T = HCBLAS_OP_T;
+
+  static constexpr auto R_16F = HPCC_R_16F;
+
+  static constexpr auto R_16BF = HPCC_R_16BF;
+
+  static constexpr auto R_32F = HPCC_R_32F;
+
+  static constexpr auto BLAS_COMPUTE_32F = HCBLAS_COMPUTE_32F;
+
+  static constexpr auto BLAS_COMPUTE_32F_FAST_TF32 =
+      HCBLAS_COMPUTE_32F_FAST_TF32;
+
+  static constexpr auto BLAS_GEMM_DEFAULT = HCBLAS_GEMM_DEFAULT;
+
+  static constexpr auto BlasCreate = hcblasCreate;
+
+  static constexpr auto BlasSetStream = hcblasSetStream;
+
+  static constexpr auto BlasDestroy = hcblasDestroy;
+
+  static constexpr auto BlasGemmStridedBatchedEx = [](auto&&... args) {
+    return hcblasGemmStridedBatchedEx(std::forward<decltype(args)>(args)...);
+  };
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/mars/blas_utils.h

@@ -0,0 +1,30 @@
+#ifndef INFINI_OPS_MARS_BLAS_UTILS_H_
+#define INFINI_OPS_MARS_BLAS_UTILS_H_
+
+// clang-format off
+#include <hcblas/hcblas.h>
+// clang-format on
+
+#include "data_type.h"
+#include "native/cuda/blas_utils.h"
+
+namespace infini::ops {
+
+template <>
+struct BlasUtils<Device::Type::kMars> {
+  static auto GetDataType(DataType dtype) {
+    if (dtype == DataType::kFloat16) return HPCC_R_16F;
+    if (dtype == DataType::kBFloat16) return HPCC_R_16BF;
+    return HPCC_R_32F;
+  }
+
+  static auto GetComputeType(DataType dtype) {
+    if (dtype == DataType::kFloat16 || dtype == DataType::kBFloat16)
+      return HCBLAS_COMPUTE_32F;
+    return HCBLAS_COMPUTE_32F_FAST_TF32;
+  }
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/mars/caster.cuh

@@ -0,0 +1,78 @@
+#ifndef INFINI_OPS_MARS_CASTER__H_
+#define INFINI_OPS_MARS_CASTER__H_
+
+#include "native/cuda/caster.cuh"
+#include "native/cuda/mars/data_type_.h"
+
+namespace infini::ops {
+
+namespace detail {
+
+template <>
+struct ToFloat<Device::Type::kMars, __half> {
+  __host__ __device__ float operator()(__half x) { return __half2float(x); }
+};
+
+template <>
+struct ToFloat<Device::Type::kMars, __hpcc_bfloat16> {
+  __host__ __device__ float operator()(__hpcc_bfloat16 x) {
+    return __bfloat162float(x);
+  }
+};
+
+template <>
+struct FromFloat<Device::Type::kMars, __half> {
+  __host__ __device__ __half operator()(float f) { return __float2half(f); }
+};
+
+template <>
+struct FromFloat<Device::Type::kMars, __hpcc_bfloat16> {
+  __host__ __device__ __hpcc_bfloat16 operator()(float f) {
+    return __float2bfloat16(f);
+  }
+};
+
+template <>
+struct HardwareCast<Device::Type::kMars, __hpcc_bfloat16, int> {
+  inline static constexpr bool kSupported = true;
+  __host__ __device__ __hpcc_bfloat16 operator()(int x) {
+    return __int2bfloat16_rn(x);
+  }
+};
+
+template <>
+struct HardwareCast<Device::Type::kMars, __half, int> {
+  inline static constexpr bool kSupported = true;
+  __host__ __device__ __half operator()(int x) { return __int2half_rn(x); }
+};
+
+template <>
+struct HardwareCast<Device::Type::kMars, __hpcc_bfloat16, double> {
+  inline static constexpr bool kSupported = true;
+  __host__ __device__ __hpcc_bfloat16 operator()(double x) {
+    return __double2bfloat16(x);
+  }
+};
+
+template <>
+struct HardwareCast<Device::Type::kMars, __half, double> {
+  inline static constexpr bool kSupported = true;
+  __host__ __device__ __half operator()(double x) { return __double2half(x); }
+};
+
+template <>
+struct HardwareCast<Device::Type::kMars, __half, __hpcc_bfloat16> {
+  inline static constexpr bool kSupported = true;
+  __host__ __device__ __half operator()(__hpcc_bfloat16 x) {
+    return __float2half_rn(__bfloat162float(x));
+  }
+};
+
+}  // namespace detail
+
+template <>
+struct Caster<Device::Type::kMars> : CudaCasterImpl<Device::Type::kMars> {};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/mars/data_type_.h

@@ -0,0 +1,13 @@
+#ifndef INFINI_OPS_MARS_DATA_TYPE__H_
+#define INFINI_OPS_MARS_DATA_TYPE__H_
+
+#include <infini/rt/mars/data_type_.h>
+
+namespace infini::ops {
+
+using infini::rt::cuda_bfloat16;
+using infini::rt::cuda_bfloat162;
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/mars/device_.h

@@ -0,0 +1,6 @@
+#ifndef INFINI_OPS_MARS_DEVICE__H_
+#define INFINI_OPS_MARS_DEVICE__H_
+
+#include <infini/rt/mars/device_.h>
+
+#endif

src/native/cuda/mars/device_property.h

@@ -0,0 +1,11 @@
+#ifndef INFINI_OPS_MARS_DEVICE_PROPERTY_H_
+#define INFINI_OPS_MARS_DEVICE_PROPERTY_H_
+
+namespace infini::ops {
+
+// TODO: Add HCR device properties query for Mars.
+inline int QueryMaxThreadsPerBlock() { return 256; }
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/mars/ops/add/kernel.h

@@ -0,0 +1,21 @@
+#ifndef INFINI_OPS_MARS_ADD_KERNEL_H_
+#define INFINI_OPS_MARS_ADD_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/mars/caster.cuh"
+#include "native/cuda/mars/runtime_.h"
+#include "native/cuda/ops/add/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<Add, Device::Type::kMars>
+    : public CudaAdd<Runtime<Device::Type::kMars>> {
+ public:
+  using CudaAdd<Runtime<Device::Type::kMars>>::CudaAdd;
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/mars/ops/add_rms_norm/kernel.h

@@ -0,0 +1,21 @@
+#ifndef INFINI_OPS_MARS_ADD_RMS_NORM_KERNEL_H_
+#define INFINI_OPS_MARS_ADD_RMS_NORM_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/mars/caster.cuh"
+#include "native/cuda/mars/runtime_.h"
+#include "native/cuda/ops/add_rms_norm/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<AddRmsNorm, Device::Type::kMars>
+    : public CudaAddRmsNorm<Runtime<Device::Type::kMars>> {
+ public:
+  using CudaAddRmsNorm<Runtime<Device::Type::kMars>>::CudaAddRmsNorm;
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/mars/ops/causal_softmax/kernel.h

@@ -0,0 +1,21 @@
+#ifndef INFINI_OPS_MARS_CAUSAL_SOFTMAX_KERNEL_H_
+#define INFINI_OPS_MARS_CAUSAL_SOFTMAX_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/mars/caster.cuh"
+#include "native/cuda/mars/runtime_.h"
+#include "native/cuda/ops/causal_softmax/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<CausalSoftmax, Device::Type::kMars>
+    : public CudaCausalSoftmax<Runtime<Device::Type::kMars>> {
+ public:
+  using CudaCausalSoftmax<Runtime<Device::Type::kMars>>::CudaCausalSoftmax;
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/mars/ops/causal_softmax_infinilm/kernel.h

@@ -0,0 +1,22 @@
+#ifndef INFINI_OPS_MARS_CAUSAL_SOFTMAX_INFINILM_KERNEL_H_
+#define INFINI_OPS_MARS_CAUSAL_SOFTMAX_INFINILM_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/mars/caster.cuh"
+#include "native/cuda/mars/runtime_.h"
+#include "native/cuda/ops/causal_softmax_infinilm/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<CausalSoftmaxInfinilm, Device::Type::kMars>
+    : public CudaCausalSoftmaxInfinilm<Runtime<Device::Type::kMars>> {
+ public:
+  using CudaCausalSoftmaxInfinilm<
+      Runtime<Device::Type::kMars>>::CudaCausalSoftmaxInfinilm;
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/mars/ops/conv1d/kernel.h

@@ -0,0 +1,21 @@
+#ifndef INFINI_OPS_MARS_CONV1D_KERNEL_H_
+#define INFINI_OPS_MARS_CONV1D_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/mars/caster.cuh"
+#include "native/cuda/mars/runtime_.h"
+#include "native/cuda/ops/convolution/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<Conv1d, Device::Type::kMars>
+    : public CudaConv<Runtime<Device::Type::kMars>, Conv1d> {
+ public:
+  using CudaConv<Runtime<Device::Type::kMars>, Conv1d>::CudaConv;
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/mars/ops/conv2d/kernel.h

@@ -0,0 +1,21 @@
+#ifndef INFINI_OPS_MARS_CONV2D_KERNEL_H_
+#define INFINI_OPS_MARS_CONV2D_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/mars/caster.cuh"
+#include "native/cuda/mars/runtime_.h"
+#include "native/cuda/ops/convolution/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<Conv2d, Device::Type::kMars>
+    : public CudaConv<Runtime<Device::Type::kMars>, Conv2d> {
+ public:
+  using CudaConv<Runtime<Device::Type::kMars>, Conv2d>::CudaConv;
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/mars/ops/conv3d/kernel.h

@@ -0,0 +1,21 @@
+#ifndef INFINI_OPS_MARS_CONV3D_KERNEL_H_
+#define INFINI_OPS_MARS_CONV3D_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/mars/caster.cuh"
+#include "native/cuda/mars/runtime_.h"
+#include "native/cuda/ops/convolution/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<Conv3d, Device::Type::kMars>
+    : public CudaConv<Runtime<Device::Type::kMars>, Conv3d> {
+ public:
+  using CudaConv<Runtime<Device::Type::kMars>, Conv3d>::CudaConv;
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/mars/ops/conv_infinilm/kernel.h

@@ -0,0 +1,21 @@
+#ifndef INFINI_OPS_MARS_CONV_INFINILM_KERNEL_H_
+#define INFINI_OPS_MARS_CONV_INFINILM_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/mars/caster.cuh"
+#include "native/cuda/mars/runtime_.h"
+#include "native/cuda/ops/conv_infinilm/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<ConvInfinilm, Device::Type::kMars>
+    : public CudaConvInfinilm<Runtime<Device::Type::kMars>> {
+ public:
+  using CudaConvInfinilm<Runtime<Device::Type::kMars>>::CudaConvInfinilm;
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/mars/ops/convolution/kernel.h

@@ -0,0 +1,21 @@
+#ifndef INFINI_OPS_MARS_CONVOLUTION_KERNEL_H_
+#define INFINI_OPS_MARS_CONVOLUTION_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/mars/caster.cuh"
+#include "native/cuda/mars/runtime_.h"
+#include "native/cuda/ops/convolution/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<Convolution, Device::Type::kMars>
+    : public CudaConv<Runtime<Device::Type::kMars>, Convolution> {
+ public:
+  using CudaConv<Runtime<Device::Type::kMars>, Convolution>::CudaConv;
+};
+
+}  // namespace infini::ops
+
+#endif