/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*!
 * \file inject_permuted_layout.cc
 * \brief The pass injects permuted layout for shared memory buffers to avoid bank conflicts.
 */
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/op.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt_functor.h>

#include "../../arith/ir_mutator_with_analyzer.h"
#include "../../runtime/thread_storage_scope.h"
#include "../../support/utils.h"
#include "../../tirx/transform/ir_utils.h"

namespace tvm {
namespace s_tir {
using namespace tvm::tirx;

using namespace arith;
using namespace runtime;

namespace {

ffi::Optional<Var> GetBufferDataVar(const ffi::Any& data) {
  if (auto var = data.as<Var>()) {
    return var;
  }
  if (const auto* call = data.as<CallNode>();
      call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) {
    return call->args[0].as<Var>();
  }
  return std::nullopt;
}

}  // namespace

class PermutedLayoutInjector : private IRMutatorWithAnalyzer {
 public:
  static PrimFunc Transform(PrimFunc func) {
    Analyzer analyzer;

    auto new_body = PermutedLayoutInjector(func, analyzer)(func->body);
    auto func_node = func.CopyOnWrite();
    func_node->body = new_body;
    return func;
  }

 private:
  explicit PermutedLayoutInjector(PrimFunc func, const Analyzer& analyzer)
      : IRMutatorWithAnalyzer(analyzer) {
    for (const Var& param : func->params) {
      if (auto buffer = param.as<BufferVar>()) {
        buffer_map_.insert({buffer.value().var(), buffer.value()});
      }
    }
  }

  using IRMutatorWithAnalyzer::VisitExpr_;
  using IRMutatorWithAnalyzer::VisitStmt_;

  ffi::Array<PrimExpr> PermuteIndices(PrimExpr row_idx, PrimExpr col_idx, int row_size) {
    TVM_FFI_ICHECK(permute_);
    // Index after vectorizing by 8
    PrimExpr col_idx_outer = floordiv(col_idx, VECTORIZE_FACTOR),
             col_idx_inner = floormod(col_idx, VECTORIZE_FACTOR);
    PrimExpr new_col_idx_outer;
    if (row_size % 64 == 0) {
      // Use 8 * 8 permuted layout
      // Every number below corresponds to 8 consecutive fp16 number in shared mem, i.e. one read
      // Every row below corresponds to 32 banks
      // 0  1  2  3  4  5  6  7    ==>    0  1  2  3  4  5  6  7
      // 0  1  2  3  4  5  6  7    ==>    1  0  3  2  5  4  7  6
      // 0  1  2  3  4  5  6  7    ==>    2  3  0  1  6  7  4  5
      // 0  1  2  3  4  5  6  7    ==>    3  2  1  0  7  6  5  4
      // 0  1  2  3  4  5  6  7    ==>    4  5  6  7  0  1  2  3
      // 0  1  2  3  4  5  6  7    ==>    5  4  7  6  1  0  3  2
      // 0  1  2  3  4  5  6  7    ==>    6  7  4  5  2  3  0  1
      // 0  1  2  3  4  5  6  7    ==>    7  6  5  4  3  2  1  0
      auto row_idx_sub = floormod(row_idx, 8);
      new_col_idx_outer = col_idx_outer ^ row_idx_sub;
    } else {
      TVM_FFI_ICHECK(row_size % 32 == 0);
      // Use 8 * 4 permuted layout
      // Every number below corresponds to 8 consecutive fp16 number in shared mem, i.e. one read
      // Every row below corresponds to 16 banks
      // 0  1  2  3    ==>    0  1  2  3
      // 0  1  2  3    ==>    0  1  2  3
      // 0  1  2  3    ==>    1  0  3  2
      // 0  1  2  3    ==>    1  0  3  2
      // 0  1  2  3    ==>    2  3  0  1
      // 0  1  2  3    ==>    2  3  0  1
      // 0  1  2  3    ==>    3  2  1  0
      // 0  1  2  3    ==>    3  2  1  0
      // View with 8 elements per row:
      // 0  1  2  3  4  0  1  2  3    ==>    0  1  2  3  0  1  2  3
      // 0  1  2  3  4  0  1  2  3    ==>    1  0  3  2  1  0  3  2
      // 0  1  2  3  4  0  1  2  3    ==>    2  3  0  1  2  3  0  1
      // 0  1  2  3  4  0  1  2  3    ==>    3  2  1  0  3  2  1  0
      auto row_idx_sub = floormod(row_idx, 8);
      new_col_idx_outer = col_idx_outer ^ floordiv(row_idx_sub, 2);
    }
    return {row_idx, analyzer_->Simplify(new_col_idx_outer * 8 + col_idx_inner)};
  }

  static bool CheckAnnotation(const Any& annotation) {
    if (auto opt_str = annotation.as<ffi::String>()) {
      // Support string annotation for backward compatibility
      return *opt_str != "";
    } else if (auto* node = annotation.as<IntImmNode>()) {
      return node->value != 0;
    } else if (auto opt_val = annotation.try_cast<int64_t>()) {
      return *opt_val != 0;
    } else {
      TVM_FFI_THROW(InternalError) << "Invalid permuted layout annotation: " << annotation;
    }
  }

  Stmt VisitStmt_(const SBlockNode* op) final {
    // Record the mapping from buffer identity to buffer for later lookup.
    for (auto buffer : op->alloc_buffers) {
      buffer_map_.insert({buffer.var(), buffer});
    }
    for (auto match_buffer : op->match_buffers) {
      buffer_map_.insert({match_buffer->buffer.var(), match_buffer->buffer});
    }

    if (op->annotations.count("permuted_layout") == 0 ||
        !CheckAnnotation(op->annotations.at("permuted_layout"))) {
      return IRMutatorWithAnalyzer::VisitStmt_(op);
    }

    auto prev_permute = permute_;
    permute_ = true;

    SBlock block = IRMutatorWithAnalyzer::VisitStmt_(op).as_or_throw<SBlock>();

    permute_ = prev_permute;

    // Erase the permuted_layout annotation after the pass
    auto block_node = block.CopyOnWrite();
    block_node->annotations.erase("permuted_layout");
    return block;
  }

  int CheckAndGetBufferRowSize(BufferVar buffer) {
    TVM_FFI_ICHECK(buffer->shape.size() >= 2)
        << "The dimension of BufferVar \"" << buffer.name() << "\" with shape " << buffer->shape
        << " should be at least 2";

    auto dim = buffer->shape.size();
    auto buffer_row_size = buffer->shape[dim - 1].as<IntImmNode>()->value;
    auto buffer_col_size = buffer->shape[dim - 2].as<IntImmNode>()->value;

    if (buffer_row_size % 64 != 0) {
      TVM_FFI_ICHECK(buffer_row_size % 32 == 0)
          << "Permuted SLayout for BufferVar \"" << buffer.name() << "\" with shape "
          << buffer->shape << " is not supported since its second dimension is not divisible by 32";
      TVM_FFI_ICHECK(buffer_col_size % 2 == 0)
          << "Permuted SLayout for BufferVar \"" << buffer.name() << "\" with shape "
          << buffer->shape
          << " is not supported since its first dimension is not divisible by 2 and second "
             "dimension is not divisible by 64";
    }

    return buffer_row_size;
  }

  ffi::Array<PrimExpr> HandleBufferIndices(BufferVar buffer, ffi::Array<PrimExpr> indices) {
    auto buffer_row_size = CheckAndGetBufferRowSize(buffer);

    // Mutate the last two indices
    auto indices_size = indices.size();
    PrimExpr row_idx = indices[indices_size - 2];
    PrimExpr col_idx = indices[indices_size - 1];
    auto new_indices = PermuteIndices(row_idx, col_idx, buffer_row_size);
    indices.Set(indices_size - 2, new_indices[0]);
    indices.Set(indices_size - 1, new_indices[1]);
    return indices;
  }

  Stmt VisitStmt_(const BufferStoreNode* op) final {
    // Rewrite write from global to shared.dyn or shared
    // We assume the shape of the shared memory is [..., row_size, col_size],
    // where row_size is divisible by 64, or divisible by 32 and col_size is divisible by 2.
    auto store = IRMutatorWithAnalyzer::VisitStmt_(op).as_or_throw<BufferStore>();

    if (!permute_ || store->buffer->shape.size() < 2) {
      return store;
    }

    auto scope = StorageScope::Create(store->buffer.scope());
    if (scope.rank != StorageRank::kShared) {
      return store;
    }

    auto store_node = store.CopyOnWrite();
    store_node->indices = HandleBufferIndices(store_node->buffer, store_node->indices);
    return store;
  }

  Expr VisitExpr_(const BufferLoadNode* op) final {
    // Rewrite load from shared or shared.dyn to global
    auto load = IRMutatorWithAnalyzer::VisitExpr_(op).as_or_throw<BufferLoad>();

    if (!permute_ || load->buffer->shape.size() < 2) {
      return load;
    }

    auto scope = StorageScope::Create(load->buffer.scope());
    if (scope.rank != StorageRank::kShared) {
      return load;
    }

    auto load_node = load.CopyOnWrite();
    load_node->indices = HandleBufferIndices(load_node->buffer, load_node->indices);
    return load;
  }

  Expr HandleAccessPtrAndOffset(Expr access_ptr, ffi::Optional<PrimExpr> offset = std::nullopt) {
    // The 2th arg of T.tvm_access_ptr call is offset, we set it to 0 and accumulate it to
    // smem_offset
    TVM_FFI_ICHECK(access_ptr->IsInstance<CallNode>())
        << "Invalid access ptr for permuted layout: " << access_ptr;
    auto access_ptr_call = access_ptr.as_or_throw<Call>();
    TVM_FFI_ICHECK(access_ptr_call->op.same_as(builtin::tvm_access_ptr()))
        << "Invalid access ptr for permuted layout: " << access_ptr;

    auto data_var = GetBufferDataVar(access_ptr_call->args[1]);
    TVM_FFI_ICHECK(data_var.has_value())
        << "Expected a buffer data expression, but received " << access_ptr_call->args[1];
    auto buffer_map_iter = buffer_map_.find(data_var.value());
    TVM_FFI_ICHECK(buffer_map_iter != buffer_map_.end())
        << "The buffer corresponding to data Var " << access_ptr_call->args[1] << " is not found";
    int buffer_row_size = CheckAndGetBufferRowSize(buffer_map_iter->second);

    PrimExpr smem_offset = access_ptr_call->args[2].as_or_throw<PrimExpr>() +
                           (offset.has_value() ? offset.value() : 0);

    // Convert offset to 2-dimension, reindex it and convert it back
    PrimExpr row_idx = floordiv(smem_offset, buffer_row_size);
    PrimExpr col_idx = floormod(smem_offset, buffer_row_size);

    auto new_indices = PermuteIndices(row_idx, col_idx, buffer_row_size);
    auto new_offset = analyzer_->Simplify(new_indices[0] * buffer_row_size + new_indices[1]);

    auto new_access_ptr = access_ptr_call.CopyOnWrite();
    new_access_ptr->args.Set(2, new_offset);
    return access_ptr_call;
  }

  Expr VisitExpr_(const CallNode* op) final {
    // Rewrite from/to shared or shared.dyn to/from local
    auto call = IRMutatorWithAnalyzer::VisitExpr_(op).as_or_throw<Call>();

    if (!permute_) {
      return call;
    }

    static const Op& ptx_ldmatrix_op = Op::Get("tirx.ptx_legacy.ldmatrix");
    static const Op& mma_store_op = Op::Get("tirx.mma_store_legacy");
    if (!call->op.same_as(ptx_ldmatrix_op) && !call->op.same_as(mma_store_op)) {
      return call;
    }

    if (call->op.same_as(ptx_ldmatrix_op)) {
      // form: T.ptx_legacy.ldmatrix(..., smem_ptr, smem_offset)
      // smem_ptr: T.tvm_access_ptr(ptype, data, offset, extent, rw_mask)
      Expr access_ptr = call->args[5];
      PrimExpr smem_offset = call->args[6].as_or_throw<PrimExpr>();
      auto new_access_ptr = HandleAccessPtrAndOffset(access_ptr, smem_offset);
      auto new_call = call.CopyOnWrite();
      new_call->args.Set(5, new_access_ptr);
      new_call->args.Set(6, IntImm(smem_offset.ty(), 0));
      return call;
    } else if (call->op.same_as(mma_store_op)) {
      // TODO(yixin): mma_store is not fully tested yet
      // because we will directly store result to BufferVar instead of calling mma_store now
      Expr access_ptr = call->args[2];
      auto new_access_ptr = HandleAccessPtrAndOffset(access_ptr);
      auto new_call = call.CopyOnWrite();
      new_call->args.Set(2, new_access_ptr);
      return call;
    } else {
      TVM_FFI_THROW(InternalError) << "Invalid call node: " << call;
    }
  }

  static constexpr size_t VECTORIZE_FACTOR = 8;
  static constexpr size_t BANK_SIZE_BYTES = 128;

  // Mapping from data Var of a BufferVar to BufferVar, for lookup
  std::unordered_map<Var, BufferVar> buffer_map_;
  bool permute_ = false;
};

namespace transform {

Pass InjectPermutedLayout() {
  auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
    return PermutedLayoutInjector::Transform(std::move(f));
  };
  return CreatePrimFuncPass(pass_func, 0, "s_tir.InjectPermutedLayout", {});
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("s_tir.transform.InjectPermutedLayout", InjectPermutedLayout);
}

}  // namespace transform

}  // namespace s_tir
}  // namespace tvm
