/*************************************************************************
 * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * See LICENSE for license information.
 ************************************************************************/

#include <optional>
#include <tuple>
#include <utility>
#include <vector>

#include "common.h"
#include "common/common.h"
#include "extensions.h"
#include "pybind.h"
#include "util.h"

namespace transformer_engine {
namespace pytorch {

namespace {

void reset_tensor_data(transformer_engine::TensorWrapper &tensor, bool rowwise, bool columnwise) {
  NVTEShape shape;
  shape.ndim = 1;
  shape.data[0] = 0;
  const transformer_engine::DType dtype = transformer_engine::DType::kFloat32;
  if (rowwise) {
    tensor.set_rowwise_data(nullptr, dtype, shape);
    tensor.set_rowwise_scale_inv(nullptr, dtype, shape);
  }
  if (columnwise) {
    tensor.set_columnwise_data(nullptr, dtype, shape);
    tensor.set_columnwise_scale_inv(nullptr, dtype, shape);
  }
}

bool is_empty_grouped_tensor_param(const NVTEBasicTensor &t) {
  if (t.data_ptr == nullptr) {
    return true;
  }
  return t.shape.ndim == 1 && t.shape.data[0] == 0;
}

}  // namespace

std::tuple<std::optional<at::Tensor>, std::optional<at::Tensor>> swizzle_scales_for_gemm(
    transformer_engine::TensorWrapper &tensor, bool rowwise_usage, bool columnwise_usage) {
  // Return early if scale swizzling is not required
  const auto scaling_mode = tensor.scaling_mode();
  switch (scaling_mode) {
    case NVTE_MXFP8_1D_SCALING:
    case NVTE_NVFP4_1D_SCALING:
      // Tensor format requires scale swizzling
      break;
    case NVTE_INVALID_SCALING:
      NVTE_ERROR("Invalid scaling mode for swizzling scaling factors.");
    default:
      // Tensor format does not require scale swizzling for GEMM
      return {std::nullopt, std::nullopt};
  }

  // Return early if scales are already swizzled
  if (tensor.get_with_gemm_swizzled_scales()) {
    return {std::nullopt, std::nullopt};
  }

  // CUDA stream
  auto stream = at::cuda::getCurrentCUDAStream();

  // Swizzle row-wise scales if needed
  std::optional<at::Tensor> rowwise_scales_pyt;
  if (rowwise_usage) {
    // Buffer for unswizzled scales
    const auto input_scales_nvte = tensor.get_rowwise_scale_inv();
    void *input_scales_dptr = input_scales_nvte.data_ptr;
    const NVTEShape input_scales_shape = input_scales_nvte.shape;
    const auto scales_dtype = static_cast<DType>(input_scales_nvte.dtype);

    // Allocate buffer for swizzled scales
    const NVTEShape output_scales_shape = input_scales_shape;
    rowwise_scales_pyt = allocateSpace(input_scales_shape, scales_dtype, false);
    void *output_scales_dptr = getDataPtr(*rowwise_scales_pyt);

    // Initialize TE tensors with scales
    const auto data_nvte = tensor.get_rowwise_data();
    const auto data_dtype = static_cast<DType>(data_nvte.dtype);
    TensorWrapper input_nvte(scaling_mode);
    input_nvte.set_rowwise_data(nullptr, data_dtype, data_nvte.shape);
    input_nvte.set_rowwise_scale_inv(input_scales_dptr, scales_dtype, input_scales_shape);
    TensorWrapper output_nvte(scaling_mode);
    output_nvte.set_rowwise_data(nullptr, data_dtype, data_nvte.shape);
    output_nvte.set_rowwise_scale_inv(output_scales_dptr, scales_dtype, output_scales_shape);
    output_nvte.set_with_gemm_swizzled_scales(true);

    // Launch kernel
    NVTE_SCOPED_GIL_RELEASE(
        { nvte_swizzle_scaling_factors(input_nvte.data(), output_nvte.data(), stream); });

    // Update tensor with swizzled scales
    tensor.set_rowwise_scale_inv(output_scales_dptr, scales_dtype, output_scales_shape);
  }

  // Swizzle column-wise scales if needed
  std::optional<at::Tensor> columnwise_scales_pyt;
  if (columnwise_usage) {
    // Buffer for unswizzled scales
    const auto input_scales_nvte = tensor.get_columnwise_scale_inv();
    void *input_scales_dptr = input_scales_nvte.data_ptr;
    const NVTEShape input_scales_shape = input_scales_nvte.shape;
    const auto scales_dtype = static_cast<DType>(input_scales_nvte.dtype);

    // Allocate buffer for swizzled scales
    const NVTEShape output_scales_shape = input_scales_shape;
    columnwise_scales_pyt = allocateSpace(input_scales_shape, scales_dtype, false);
    void *output_scales_dptr = getDataPtr(*columnwise_scales_pyt);

    // Initialize TE tensors with scales
    const auto data_nvte = tensor.get_columnwise_data();
    const auto data_dtype = static_cast<DType>(data_nvte.dtype);
    TensorWrapper input_nvte(scaling_mode);
    input_nvte.set_columnwise_data(nullptr, data_dtype, data_nvte.shape);
    input_nvte.set_columnwise_scale_inv(input_scales_dptr, scales_dtype, input_scales_shape);
    TensorWrapper output_nvte(scaling_mode);
    output_nvte.set_columnwise_data(nullptr, data_dtype, data_nvte.shape);
    output_nvte.set_columnwise_scale_inv(output_scales_dptr, scales_dtype, output_scales_shape);
    output_nvte.set_with_gemm_swizzled_scales(true);

    // Launch kernel
    NVTE_SCOPED_GIL_RELEASE(
        { nvte_swizzle_scaling_factors(input_nvte.data(), output_nvte.data(), stream); });

    // Update tensor with swizzled scales
    tensor.set_columnwise_scale_inv(output_scales_dptr, scales_dtype, output_scales_shape);
  }

  // Update tensor
  reset_tensor_data(tensor, !rowwise_usage, !columnwise_usage);
  tensor.set_with_gemm_swizzled_scales(true);

  return {std::move(rowwise_scales_pyt), std::move(columnwise_scales_pyt)};
}

namespace {

std::optional<at::Tensor> multi_tensor_swizzle_scales_for_gemm_impl(
    std::vector<transformer_engine::TensorWrapper> &tensors, bool rowwise_usage,
    bool columnwise_usage, bool check_scale_inv_shapes) {
  // Checks and trivial cases
  NVTE_CHECK(rowwise_usage != columnwise_usage,
             "Expect exactly one of rowwise_usage=", rowwise_usage,
             " and columnwise_usage=", columnwise_usage, ".");
  if (tensors.empty()) {
    return std::nullopt;
  }
  const auto scaling_mode = tensors.front().scaling_mode();
  for (const auto &tensor : tensors) {
    NVTE_CHECK(tensor.scaling_mode() == scaling_mode, "Tensors have different scaling modes");
  }

  // Return early if scale swizzling is not required
  switch (scaling_mode) {
    case NVTE_MXFP8_1D_SCALING:
    case NVTE_NVFP4_1D_SCALING:
      // Tensor format requires scale swizzling
      break;
    case NVTE_INVALID_SCALING:
      NVTE_ERROR("Invalid scaling mode for swizzling scaling factors.");
    default:
      // Tensor format does not require scale swizzling for GEMM
      return std::nullopt;
  }

  // Filter out tensors that already have swizzled scales
  std::vector<TensorWrapper *> tensors_needing_swizzle;
  tensors_needing_swizzle.reserve(tensors.size());
  for (auto &tensor : tensors) {
    if (!tensor.get_with_gemm_swizzled_scales()) {
      tensors_needing_swizzle.push_back(&tensor);
    }
  }
  if (tensors_needing_swizzle.empty()) {
    return std::nullopt;
  }

  // Determine buffer size needed for swizzled scales
  std::vector<size_t> output_scales_offsets;
  output_scales_offsets.reserve(tensors_needing_swizzle.size());
  size_t output_scales_bytes = 0;
  for (auto &tensor : tensors_needing_swizzle) {
    const auto scales_nvte =
        (rowwise_usage ? tensor->get_rowwise_scale_inv() : tensor->get_columnwise_scale_inv());
    const auto &shape = scales_nvte.shape;
    const auto dtype = static_cast<DType>(scales_nvte.dtype);
    const auto dtype_bits = transformer_engine::pytorch::typeToNumBits(dtype);
    const auto size = product(shape, 0, shape.ndim);
    output_scales_bytes = roundup(output_scales_bytes, 16);  // align to 16B
    output_scales_offsets.push_back(output_scales_bytes);
    output_scales_bytes += ceildiv(size * dtype_bits, 8);
  }

  // Allocate buffer for swizzled scales
  auto output_scales_pyt = allocateSpace(std::vector<size_t>{output_scales_bytes},
                                         transformer_engine::DType::kByte, false);
  uint8_t *output_scales_dptr = reinterpret_cast<uint8_t *>(getDataPtr(output_scales_pyt));

  // Allocate input/output NVTETensors as a single batch. The first
  // n_swizzle entries are inputs; the next n_swizzle are outputs.
  const size_t n_swizzle = tensors_needing_swizzle.size();
  MultiTensorWrapper nvte_tensors(2 * n_swizzle, scaling_mode);
  NVTETensor *inputs_nvte = nvte_tensors.data();
  NVTETensor *outputs_nvte = nvte_tensors.data() + n_swizzle;

  auto set_param = [](NVTETensor t, NVTETensorParam param, void *dptr,
                      transformer_engine::DType dtype, const NVTEShape &shape) {
    NVTEBasicTensor data{dptr, static_cast<NVTEDType>(dtype), shape};
    nvte_set_tensor_param_v2(t, param, &data, sizeof(data));
  };

  // Cache output scale dtype/shape per tensor so we can update the
  // source TensorWrappers without re-reading from the output NVTETensors.
  std::vector<transformer_engine::DType> output_scales_dtypes(n_swizzle);
  std::vector<NVTEShape> output_scales_shapes(n_swizzle);

  for (size_t i = 0; i < n_swizzle; ++i) {
    auto &tensor = *tensors_needing_swizzle[i];
    const uint8_t swizzled_flag = 1;
    nvte_set_tensor_param_v2(outputs_nvte[i], kNVTEWithGEMMSwizzledScales, &swizzled_flag,
                             sizeof(swizzled_flag));
    if (rowwise_usage) {
      const auto data_nvte = tensor.get_rowwise_data();
      const auto scales_nvte = tensor.get_rowwise_scale_inv();
      const auto data_dtype = static_cast<transformer_engine::DType>(data_nvte.dtype);
      const auto scales_dtype = static_cast<transformer_engine::DType>(scales_nvte.dtype);
      output_scales_dtypes[i] = scales_dtype;
      output_scales_shapes[i] = scales_nvte.shape;
      set_param(inputs_nvte[i], kNVTERowwiseData, nullptr, data_dtype, data_nvte.shape);
      set_param(inputs_nvte[i], kNVTERowwiseScaleInv, scales_nvte.data_ptr, scales_dtype,
                scales_nvte.shape);
      set_param(outputs_nvte[i], kNVTERowwiseData, nullptr, data_dtype, data_nvte.shape);
      set_param(outputs_nvte[i], kNVTERowwiseScaleInv,
                output_scales_dptr + output_scales_offsets[i], scales_dtype, scales_nvte.shape);
    } else {
      const auto data_nvte = tensor.get_columnwise_data();
      const auto scales_nvte = tensor.get_columnwise_scale_inv();
      const auto data_dtype = static_cast<transformer_engine::DType>(data_nvte.dtype);
      const auto scales_dtype = static_cast<transformer_engine::DType>(scales_nvte.dtype);
      output_scales_dtypes[i] = scales_dtype;
      output_scales_shapes[i] = scales_nvte.shape;
      set_param(inputs_nvte[i], kNVTEColumnwiseData, nullptr, data_dtype, data_nvte.shape);
      set_param(inputs_nvte[i], kNVTEColumnwiseScaleInv, scales_nvte.data_ptr, scales_dtype,
                scales_nvte.shape);
      set_param(outputs_nvte[i], kNVTEColumnwiseData, nullptr, data_dtype, data_nvte.shape);
      set_param(outputs_nvte[i], kNVTEColumnwiseScaleInv,
                output_scales_dptr + output_scales_offsets[i], scales_dtype, scales_nvte.shape);
    }
  }

  // Launch kernel
  NVTE_SCOPED_GIL_RELEASE({
    if (check_scale_inv_shapes) {
      nvte_multi_tensor_swizzle_scaling_factors(inputs_nvte, outputs_nvte, n_swizzle,
                                                at::cuda::getCurrentCUDAStream());
    } else {
      nvte_multi_tensor_swizzle_scaling_factors_unchecked(inputs_nvte, outputs_nvte, n_swizzle,
                                                          at::cuda::getCurrentCUDAStream());
    }
  });

  // Update tensors with swizzled scales
  for (size_t i = 0; i < n_swizzle; ++i) {
    auto &tensor = *tensors_needing_swizzle[i];
    reset_tensor_data(tensor, !rowwise_usage, !columnwise_usage);
    tensor.set_with_gemm_swizzled_scales(true);
    if (rowwise_usage) {
      tensor.set_rowwise_scale_inv(output_scales_dptr + output_scales_offsets[i],
                                   output_scales_dtypes[i], output_scales_shapes[i]);
    } else {
      tensor.set_columnwise_scale_inv(output_scales_dptr + output_scales_offsets[i],
                                      output_scales_dtypes[i], output_scales_shapes[i]);
    }
  }

  return std::move(output_scales_pyt);
}

}  // anonymous namespace

std::optional<at::Tensor> multi_tensor_swizzle_scales_for_gemm(
    std::vector<transformer_engine::TensorWrapper> &tensors, bool rowwise_usage,
    bool columnwise_usage) {
  return multi_tensor_swizzle_scales_for_gemm_impl(tensors, rowwise_usage, columnwise_usage,
                                                   /*check_scale_inv_shapes=*/true);
}

std::optional<at::Tensor> multi_tensor_swizzle_scales_for_gemm_unchecked(
    std::vector<transformer_engine::TensorWrapper> &tensors, bool rowwise_usage,
    bool columnwise_usage) {
  return multi_tensor_swizzle_scales_for_gemm_impl(tensors, rowwise_usage, columnwise_usage,
                                                   /*check_scale_inv_shapes=*/false);
}

at::Tensor convert_block_scaling_to_mxfp8_tensor(transformer_engine::TensorWrapper &input,
                                                 bool rowwise) {
  // Check input tensor
  const NVTEScalingMode scaling_mode = input.scaling_mode();
  NVTE_CHECK(scaling_mode == NVTE_BLOCK_SCALING_1D || scaling_mode == NVTE_BLOCK_SCALING_2D,
             "Input tensor must be a block scaling tensor");

  // Get tensor data
  NVTEBasicTensor data = rowwise ? input.get_rowwise_data() : input.get_columnwise_data();
  const auto [data_flat_first_dim, data_flat_last_dim] = get_2d_dims(data.shape, !rowwise);
  NVTEShape data_shape{};
  data_shape.data[0] = data_flat_first_dim;
  data_shape.data[1] = data_flat_last_dim;
  data_shape.ndim = 2;

  // Recreate input tensor with rowwise usage
  transformer_engine::TensorWrapper input_cu(scaling_mode);
  input_cu.set_rowwise_data(data.data_ptr, input.dtype(), data_shape);
  const NVTEBasicTensor scale_inv =
      rowwise ? input.get_rowwise_scale_inv() : input.get_columnwise_scale_inv();
  input_cu.set_rowwise_scale_inv(
      scale_inv.data_ptr, static_cast<transformer_engine::DType>(scale_inv.dtype), scale_inv.shape);

  // Create output tensor
  transformer_engine::TensorWrapper output_cu(NVTE_MXFP8_1D_SCALING);
  output_cu.set_rowwise_data(data.data_ptr, input.dtype(), data_shape);
  // Output swizzled mxfp8 scaling factor dimensions
  const size_t swizzled_scale_inv_first_dim = ceildiv(data_flat_first_dim, 128) * 128;
  const size_t swizzled_scale_inv_last_dim = ceildiv(data_flat_last_dim, 128) * 4;
  // Allocate memory for swizzled mxfp8 scaling factors
  at::Tensor swizzled_scale_inv =
      allocateSpace(std::vector<size_t>{swizzled_scale_inv_first_dim, swizzled_scale_inv_last_dim},
                    transformer_engine::DType::kByte, false);
  // Set rowwise scaling factors on output
  void *const swizzled_scale_inv_dptr = getDataPtr(swizzled_scale_inv, 0);
  NVTEShape swizzled_scale_inv_shape{};
  swizzled_scale_inv_shape.data[0] = swizzled_scale_inv_first_dim;
  swizzled_scale_inv_shape.data[1] = swizzled_scale_inv_last_dim;
  swizzled_scale_inv_shape.ndim = 2;
  output_cu.set_rowwise_scale_inv(swizzled_scale_inv_dptr, transformer_engine::DType::kFloat8E8M0,
                                  swizzled_scale_inv_shape);
  output_cu.set_with_gemm_swizzled_scales(true);

  // Convert scaling factors from FP8 block scaling GEMM_READY format to mxfp8 swizzled format
  NVTE_SCOPED_GIL_RELEASE({
    nvte_swizzle_block_scaling_to_mxfp8_scaling_factors(input_cu.data(), output_cu.data(),
                                                        at::cuda::getCurrentCUDAStream());
  });

  // Set the input tensor to be the converted mxfp8 tensor and return the swizzled scaling factor
  // for it to be kept alive during the GEMM
  input = std::move(output_cu);
  return swizzled_scale_inv;
}

std::optional<SwizzledGroupedScales> maybe_swizzle_grouped_tensor(GroupedTensorWrapper &input,
                                                                  bool rowwise_usage,
                                                                  bool columnwise_usage) {
  if (input.scaling_mode() != NVTE_MXFP8_1D_SCALING) {
    return std::nullopt;
  }
  if (input.get_with_gemm_swizzled_scales()) {
    return std::nullopt;
  }

  const auto row_scales = input.get_rowwise_scale_inv();
  const auto col_scales = input.get_columnwise_scale_inv();
  const bool swizzle_rowwise = rowwise_usage && !is_empty_grouped_tensor_param(row_scales);
  const bool swizzle_columnwise = columnwise_usage && !is_empty_grouped_tensor_param(col_scales);
  if (!swizzle_rowwise && !swizzle_columnwise) {
    return std::nullopt;
  }

  std::optional<at::Tensor> rowwise_scales_pyt;
  std::optional<at::Tensor> columnwise_scales_pyt;

  GroupedTensorWrapper swizzle_input(input.num_tensors(), input.logical_shape(),
                                     input.scaling_mode());
  GroupedTensorWrapper swizzle_output(input.num_tensors(), input.logical_shape(),
                                      input.scaling_mode());

  const auto tensor_offsets = input.get_tensor_offsets();
  if (tensor_offsets.data_ptr != nullptr) {
    swizzle_input.set_tensor_offsets(
        tensor_offsets.data_ptr, static_cast<DType>(tensor_offsets.dtype), tensor_offsets.shape);
    swizzle_output.set_tensor_offsets(
        tensor_offsets.data_ptr, static_cast<DType>(tensor_offsets.dtype), tensor_offsets.shape);
  }

  // Varying per-tensor dimensions. Leaving these unset declares the grouped tensor uniform,
  // which selects the uniform-shape swizzle kernel.
  const auto first_dims = input.get_first_dims();
  if (first_dims.data_ptr != nullptr) {
    swizzle_input.set_first_dims(first_dims.data_ptr, static_cast<DType>(first_dims.dtype),
                                 first_dims.shape);
    swizzle_output.set_first_dims(first_dims.data_ptr, static_cast<DType>(first_dims.dtype),
                                  first_dims.shape);
  }
  const auto last_dims = input.get_last_dims();
  if (last_dims.data_ptr != nullptr) {
    swizzle_input.set_last_dims(last_dims.data_ptr, static_cast<DType>(last_dims.dtype),
                                last_dims.shape);
    swizzle_output.set_last_dims(last_dims.data_ptr, static_cast<DType>(last_dims.dtype),
                                 last_dims.shape);
  }

  // Per-tensor logical dimensions (uniform-shape grouped tensor).
  const size_t num_tensors = input.num_tensors();
  const auto logical_shape_nvte = input.logical_shape();
  NVTE_CHECK(logical_shape_nvte.ndim >= 2,
             "Grouped GEMM swizzle expects logical_shape with ndim >= 2.");
  const size_t per_tensor_first_dim = logical_shape_nvte.data[0] / num_tensors;
  const size_t per_tensor_last_dim = logical_shape_nvte.data[logical_shape_nvte.ndim - 1];
  constexpr size_t kMxfp8BlockSize = 32;

  // Output is always allocated in the per-tensor padded ("swizzle-ready") layout
  // so the cuDNN grouped GEMM consumer sees the correct stride between experts.
  // The swizzle kernel itself handles converting from the kernel-emitted compact
  // layout (per-tensor first dim is the unpadded value) to this padded layout.
  auto compute_padded_grouped_scale_shape = [&](bool rowwise) {
    const size_t m = rowwise ? per_tensor_first_dim : per_tensor_last_dim;
    const size_t k = rowwise ? per_tensor_last_dim : per_tensor_first_dim;
    const size_t padded_m = ceildiv(m, size_t{128}) * 128;
    const size_t padded_k = ceildiv(ceildiv(k, kMxfp8BlockSize), size_t{4}) * 4;
    return std::vector<size_t>{num_tensors * padded_m, padded_k};
  };

  if (swizzle_rowwise) {
    const auto data = input.get_rowwise_data();
    const auto data_dtype = static_cast<DType>(data.dtype);
    const auto scales_dtype = static_cast<DType>(row_scales.dtype);
    swizzle_input.set_rowwise_data(nullptr, data_dtype, data.shape);
    swizzle_input.set_rowwise_scale_inv(row_scales.data_ptr, scales_dtype, row_scales.shape);
    const auto padded_shape = compute_padded_grouped_scale_shape(/*rowwise=*/true);
    rowwise_scales_pyt = allocateSpace(padded_shape, scales_dtype, false);
    NVTEShape padded_shape_nvte = nvte_make_shape(padded_shape.data(), padded_shape.size());
    swizzle_output.set_rowwise_data(nullptr, data_dtype, data.shape);
    swizzle_output.set_rowwise_scale_inv(getDataPtr(*rowwise_scales_pyt), scales_dtype,
                                         padded_shape_nvte);
  }
  if (swizzle_columnwise) {
    const auto data = input.get_columnwise_data();
    const auto data_dtype = static_cast<DType>(data.dtype);
    const auto scales_dtype = static_cast<DType>(col_scales.dtype);
    swizzle_input.set_columnwise_data(nullptr, data_dtype, data.shape);
    swizzle_input.set_columnwise_scale_inv(col_scales.data_ptr, scales_dtype, col_scales.shape);
    const auto padded_shape = compute_padded_grouped_scale_shape(/*rowwise=*/false);
    columnwise_scales_pyt = allocateSpace(padded_shape, scales_dtype, false);
    NVTEShape padded_shape_nvte = nvte_make_shape(padded_shape.data(), padded_shape.size());
    swizzle_output.set_columnwise_data(nullptr, data_dtype, data.shape);
    swizzle_output.set_columnwise_scale_inv(getDataPtr(*columnwise_scales_pyt), scales_dtype,
                                            padded_shape_nvte);
  }

  swizzle_output.set_with_gemm_swizzled_scales(true);

  NVTE_SCOPED_GIL_RELEASE({
    nvte_swizzle_grouped_scaling_factors(swizzle_input.data(), swizzle_output.data(),
                                         at::cuda::getCurrentCUDAStream());
  });

  if (swizzle_rowwise) {
    const auto scales_dtype = static_cast<DType>(row_scales.dtype);
    input.set_rowwise_scale_inv(getDataPtr(*rowwise_scales_pyt), scales_dtype,
                                getTensorShape(*rowwise_scales_pyt));
  }
  if (swizzle_columnwise) {
    const auto scales_dtype = static_cast<DType>(col_scales.dtype);
    input.set_columnwise_scale_inv(getDataPtr(*columnwise_scales_pyt), scales_dtype,
                                   getTensorShape(*columnwise_scales_pyt));
  }
  input.set_with_gemm_swizzled_scales(true);
  return SwizzledGroupedScales{std::move(rowwise_scales_pyt), std::move(columnwise_scales_pyt)};
}

void grouped_swizzle_for_gemm(py::handle &tensor, bool rowwise, bool columnwise) {
  using namespace transformer_engine::pytorch::detail;

  auto tensor_nvte = GroupedTensorFromPyTorchGroupedTensor(tensor);

  auto result = maybe_swizzle_grouped_tensor(tensor_nvte, rowwise, columnwise);

  if (result.has_value()) {
    if (result->first.has_value()) {
      tensor.attr("scale_inv") = py::cast(*result->first);
    } else {
      tensor.attr("scale_inv") = py::none();
    }
    if (result->second.has_value()) {
      tensor.attr("columnwise_scale_inv") = py::cast(*result->second);
    } else {
      tensor.attr("columnwise_scale_inv") = py::none();
    }
    tensor.attr("_with_gemm_swizzled_scales") = py::cast(true);
  }
}

namespace {

void inplace_multi_tensor_swizzle_scales_for_gemm_impl(std::vector<py::object> &tensors,
                                                       bool rowwise_usage, bool columnwise_usage,
                                                       bool check_scale_inv_shapes) {
  NVTE_CHECK(rowwise_usage != columnwise_usage,
             "Expect exactly one of rowwise_usage and columnwise_usage.");
  if (tensors.empty()) {
    return;
  }

  // Convert Python tensors to TensorWrappers, filtering those that need swizzling
  std::vector<size_t> swizzle_indices;
  std::vector<transformer_engine::TensorWrapper> wrappers_to_swizzle;

  for (size_t i = 0; i < tensors.size(); ++i) {
    auto tw = makeTransformerEngineTensor(tensors[i], py::none());

    if (i == 0) {
      switch (tw.scaling_mode()) {
        case NVTE_MXFP8_1D_SCALING:
        case NVTE_NVFP4_1D_SCALING:
          break;
        case NVTE_INVALID_SCALING:
          NVTE_ERROR("Invalid scaling mode for swizzling scaling factors.");
        default:
          return;
      }
    }

    if (tw.get_with_gemm_swizzled_scales()) {
      continue;
    }
    const auto scales_nvte =
        rowwise_usage ? tw.get_rowwise_scale_inv() : tw.get_columnwise_scale_inv();
    if (scales_nvte.data_ptr == nullptr ||
        (scales_nvte.shape.ndim == 1 && scales_nvte.shape.data[0] == 0)) {
      continue;
    }

    swizzle_indices.push_back(i);
    wrappers_to_swizzle.push_back(std::move(tw));
  }

  if (wrappers_to_swizzle.empty()) {
    return;
  }

  // Delegate to core C++ function
  auto swizzle_fn = check_scale_inv_shapes ? multi_tensor_swizzle_scales_for_gemm
                                           : multi_tensor_swizzle_scales_for_gemm_unchecked;
  auto output_buffer = swizzle_fn(wrappers_to_swizzle, rowwise_usage, columnwise_usage);
  if (!output_buffer.has_value()) {
    return;
  }

  // Update Python objects with properly-shaped views into the contiguous output buffer
  const uint8_t *base = reinterpret_cast<const uint8_t *>(output_buffer->data_ptr());
  for (size_t j = 0; j < wrappers_to_swizzle.size(); ++j) {
    const auto scales_nvte = rowwise_usage ? wrappers_to_swizzle[j].get_rowwise_scale_inv()
                                           : wrappers_to_swizzle[j].get_columnwise_scale_inv();

    const size_t offset = reinterpret_cast<const uint8_t *>(scales_nvte.data_ptr) - base;
    const auto dtype = static_cast<DType>(scales_nvte.dtype);
    const size_t num_elements = product(scales_nvte.shape, 0, scales_nvte.shape.ndim);
    const size_t num_bytes =
        ceildiv(num_elements * transformer_engine::pytorch::typeToNumBits(dtype), size_t(8));

    std::vector<int64_t> torch_shape;
    for (size_t d = 0; d < scales_nvte.shape.ndim; ++d) {
      torch_shape.push_back(static_cast<int64_t>(scales_nvte.shape.data[d]));
    }
    auto scale_view =
        output_buffer->narrow(0, static_cast<int64_t>(offset), static_cast<int64_t>(num_bytes))
            .view(torch_shape);

    if (rowwise_usage) {
      tensors[swizzle_indices[j]].attr("_rowwise_scale_inv") = py::cast(scale_view);
    } else {
      tensors[swizzle_indices[j]].attr("_columnwise_scale_inv") = py::cast(scale_view);
    }
  }
}

}  // anonymous namespace

void inplace_multi_tensor_swizzle_scales_for_gemm(std::vector<py::object> &tensors,
                                                  bool rowwise_usage, bool columnwise_usage) {
  inplace_multi_tensor_swizzle_scales_for_gemm_impl(tensors, rowwise_usage, columnwise_usage,
                                                    /*check_scale_inv_shapes=*/true);
}

void inplace_multi_tensor_swizzle_scales_for_gemm_unchecked(std::vector<py::object> &tensors,
                                                            bool rowwise_usage,
                                                            bool columnwise_usage) {
  inplace_multi_tensor_swizzle_scales_for_gemm_impl(tensors, rowwise_usage, columnwise_usage,
                                                    /*check_scale_inv_shapes=*/false);
}

void inplace_swizzle_scale_for_gemm(py::handle &tensor) {
  // Convert Python tensor to C++ tensor
  auto tensor_nvte = makeTransformerEngineTensor(tensor, py::none());

  // Return early if scale swizzling is not required
  const auto scaling_mode = tensor_nvte.scaling_mode();
  switch (scaling_mode) {
    case NVTE_MXFP8_1D_SCALING:
    case NVTE_NVFP4_1D_SCALING:
      // Tensor format requires scale swizzling
      break;
    case NVTE_INVALID_SCALING:
      NVTE_ERROR("Invalid scaling mode for swizzling scaling factors.");
    default:
      // Tensor format does not require scale swizzling for GEMM
      return;
  }

  // Return early if scales are already swizzled
  if (tensor_nvte.get_with_gemm_swizzled_scales()) {
    return;
  }

  // Check what scaling factors the tensor contains
  auto is_empty = [](const NVTEBasicTensor &t) -> bool {
    return t.shape.ndim == 1 && t.shape.data[0] == 0;
  };
  const bool has_rowwise_scales = !is_empty(tensor_nvte.get_rowwise_scale_inv());
  const bool has_columnwise_scales = !is_empty(tensor_nvte.get_columnwise_scale_inv());

  // Swizzle scaling factors
  auto [rowwise_scales, columnwise_scales] =
      swizzle_scales_for_gemm(tensor_nvte, has_rowwise_scales, has_columnwise_scales);

  // Update Python tensor with swizzled scales
  switch (scaling_mode) {
    case NVTE_MXFP8_1D_SCALING:
      if (has_rowwise_scales) {
        tensor.attr("_rowwise_scale_inv") = rowwise_scales;
      }
      if (has_columnwise_scales) {
        tensor.attr("_columnwise_scale_inv") = columnwise_scales;
      }
      tensor.attr("_with_gemm_swizzled_scales") = true;
      break;
    case NVTE_NVFP4_1D_SCALING:
      if (has_rowwise_scales) {
        tensor.attr("_rowwise_scale_inv") = rowwise_scales;
      }
      if (has_columnwise_scales) {
        tensor.attr("_columnwise_scale_inv") = columnwise_scales;
      }
      tensor.attr("_with_gemm_swizzled_scales") = true;
      break;
    default:
      NVTE_ERROR("Invalid scaling mode for swizzling scaling factors.");
  }
}

}  // namespace pytorch
}  // namespace transformer_engine
