///////////////////////////////////////////////////////////////////////////////
//                                                                           //
// HLMatrixType.cpp                                                          //
// Copyright (C) Microsoft Corporation. All rights reserved.                 //
// This file is distributed under the University of Illinois Open Source     //
// License. See LICENSE.TXT for details.                                     //
//                                                                           //
///////////////////////////////////////////////////////////////////////////////

#include "dxc/HLSL/HLMatrixType.h"
#include "dxc/Support/Global.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Type.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/Value.h"

using namespace llvm;
using namespace hlsl;

HLMatrixType::HLMatrixType(Type *RegReprElemTy, unsigned NumRows, unsigned NumColumns)
  : RegReprElemTy(RegReprElemTy), NumRows(NumRows), NumColumns(NumColumns) {
  DXASSERT(RegReprElemTy != nullptr && (RegReprElemTy->isIntegerTy() || RegReprElemTy->isFloatingPointTy()),
    "Invalid matrix element type.");
  DXASSERT(NumRows >= 1 && NumRows <= 4 && NumColumns >= 1 && NumColumns <= 4,
    "Invalid matrix dimensions.");
}

Type *HLMatrixType::getElementType(bool MemRepr) const {
  // Bool i1s become i32s
  return MemRepr && RegReprElemTy->isIntegerTy(1)
    ? IntegerType::get(RegReprElemTy->getContext(), 32)
    : RegReprElemTy;
}

unsigned HLMatrixType::getRowMajorIndex(unsigned RowIdx, unsigned ColIdx) const {
  return getRowMajorIndex(RowIdx, ColIdx, NumRows, NumColumns);
}

unsigned HLMatrixType::getColumnMajorIndex(unsigned RowIdx, unsigned ColIdx) const {
  return getColumnMajorIndex(RowIdx, ColIdx, NumRows, NumColumns);
}

unsigned HLMatrixType::getRowMajorIndex(unsigned RowIdx, unsigned ColIdx, unsigned NumRows, unsigned NumColumns) {
  DXASSERT_NOMSG(RowIdx < NumRows && ColIdx < NumColumns);
  return RowIdx * NumColumns + ColIdx;
}

unsigned HLMatrixType::getColumnMajorIndex(unsigned RowIdx, unsigned ColIdx, unsigned NumRows, unsigned NumColumns) {
  DXASSERT_NOMSG(RowIdx < NumRows && ColIdx < NumColumns);
  return ColIdx * NumRows + RowIdx;
}

VectorType *HLMatrixType::getLoweredVectorType(bool MemRepr) const {
  return VectorType::get(getElementType(MemRepr), getNumElements());
}

Value *HLMatrixType::emitLoweredMemToReg(Value *Val, IRBuilder<> &Builder) const {
  DXASSERT(Val->getType()->getScalarType() == getElementTypeForMem(), "Lowered matrix type mismatch.");
  if (RegReprElemTy->isIntegerTy(1)) {
    Val = Builder.CreateICmpNE(Val, Constant::getNullValue(Val->getType()), "tobool");
  }
  return Val;
}

Value *HLMatrixType::emitLoweredRegToMem(Value *Val, IRBuilder<> &Builder) const {
  DXASSERT(Val->getType()->getScalarType() == RegReprElemTy, "Lowered matrix type mismatch.");
  if (RegReprElemTy->isIntegerTy(1)) {
    Type *MemReprTy = Val->getType()->isVectorTy() ? getLoweredVectorTypeForMem() : getElementTypeForMem();
    Val = Builder.CreateZExt(Val, MemReprTy, "frombool");
  }
  return Val;
}

Value *HLMatrixType::emitLoweredLoad(Value *Ptr, IRBuilder<> &Builder) const {
  return emitLoweredMemToReg(Builder.CreateLoad(Ptr), Builder);
}

StoreInst *HLMatrixType::emitLoweredStore(Value *Val, Value *Ptr, IRBuilder<> &Builder) const {
  return Builder.CreateStore(emitLoweredRegToMem(Val, Builder), Ptr);
}

Value *HLMatrixType::emitLoweredVectorRowToCol(Value *VecVal, IRBuilder<> &Builder) const {
  DXASSERT(VecVal->getType() == getLoweredVectorTypeForReg(), "Lowered matrix type mismatch.");
  if (NumRows == 1 || NumColumns == 1) return VecVal;

  SmallVector<int, 16> ShuffleIndices;
  for (unsigned ColIdx = 0; ColIdx < NumColumns; ++ColIdx)
    for (unsigned RowIdx = 0; RowIdx < NumRows; ++RowIdx)
      ShuffleIndices.emplace_back((int)getRowMajorIndex(RowIdx, ColIdx));
  return Builder.CreateShuffleVector(VecVal, VecVal, ShuffleIndices, "row2col");
}

Value *HLMatrixType::emitLoweredVectorColToRow(Value *VecVal, IRBuilder<> &Builder) const {
  DXASSERT(VecVal->getType() == getLoweredVectorTypeForReg(), "Lowered matrix type mismatch.");
  if (NumRows == 1 || NumColumns == 1) return VecVal;

  SmallVector<int, 16> ShuffleIndices;
  for (unsigned RowIdx = 0; RowIdx < NumRows; ++RowIdx)
    for (unsigned ColIdx = 0; ColIdx < NumColumns; ++ColIdx)
      ShuffleIndices.emplace_back((int)getColumnMajorIndex(RowIdx, ColIdx));
  return Builder.CreateShuffleVector(VecVal, VecVal, ShuffleIndices, "col2row");
}

bool HLMatrixType::isa(Type *Ty) {
  StructType *StructTy = llvm::dyn_cast<StructType>(Ty);
  return StructTy != nullptr && !StructTy->isLiteral() && StructTy->getName().startswith(StructNamePrefix);
}

bool HLMatrixType::isMatrixPtr(Type *Ty) {
  PointerType *PtrTy = llvm::dyn_cast<PointerType>(Ty);
  return PtrTy != nullptr && isa(PtrTy->getElementType());
}

bool HLMatrixType::isMatrixArray(Type *Ty) {
  ArrayType *ArrayTy = llvm::dyn_cast<ArrayType>(Ty);
  if (ArrayTy == nullptr) return false;
  while (ArrayType *NestedArrayTy = llvm::dyn_cast<ArrayType>(ArrayTy->getElementType()))
    ArrayTy = NestedArrayTy;
  return isa(ArrayTy->getElementType());
}

bool HLMatrixType::isMatrixArrayPtr(Type *Ty) {
  PointerType *PtrTy = llvm::dyn_cast<PointerType>(Ty);
  if (PtrTy == nullptr) return false;
  return isMatrixArray(PtrTy->getElementType());
}

bool HLMatrixType::isMatrixPtrOrArrayPtr(Type *Ty) {
  PointerType *PtrTy = llvm::dyn_cast<PointerType>(Ty);
  if (PtrTy == nullptr) return false;
  Ty = PtrTy->getElementType();
  while (ArrayType *ArrayTy = llvm::dyn_cast<ArrayType>(Ty))
    Ty = Ty->getArrayElementType();
  return isa(Ty);
}

bool HLMatrixType::isMatrixOrPtrOrArrayPtr(Type *Ty) {
  if (PointerType *PtrTy = llvm::dyn_cast<PointerType>(Ty)) Ty = PtrTy->getElementType();
  while (ArrayType *ArrayTy = llvm::dyn_cast<ArrayType>(Ty)) Ty = ArrayTy->getElementType();
  return isa(Ty);
}

// Converts a matrix, matrix pointer, or matrix array pointer type to its lowered equivalent.
// If the type is not matrix-derived, the original type is returned.
// Does not lower struct types containing matrices.
Type *HLMatrixType::getLoweredType(Type *Ty, bool MemRepr) {
  if (PointerType *PtrTy = llvm::dyn_cast<PointerType>(Ty)) {
    // Pointees are always in memory representation
    Type *LoweredElemTy = getLoweredType(PtrTy->getElementType(), /* MemRepr */ true);
    return LoweredElemTy == PtrTy->getElementType()
      ? Ty : PointerType::get(LoweredElemTy, PtrTy->getAddressSpace());
  }
  else if (ArrayType *ArrayTy = llvm::dyn_cast<ArrayType>(Ty)) {
    // Arrays are always in memory and so their elements are in memory representation
    Type *LoweredElemTy = getLoweredType(ArrayTy->getElementType(), /* MemRepr */ true);
    return LoweredElemTy == ArrayTy->getElementType()
      ? Ty : ArrayType::get(LoweredElemTy, ArrayTy->getNumElements());
  }
  else if (HLMatrixType MatrixTy = HLMatrixType::dyn_cast(Ty)) {
    return MatrixTy.getLoweredVectorType(MemRepr);
  }
  else return Ty;
}

HLMatrixType HLMatrixType::cast(Type *Ty) {
  DXASSERT_NOMSG(isa(Ty));
  StructType *StructTy = llvm::cast<StructType>(Ty);
  DXASSERT_NOMSG(Ty->getNumContainedTypes() == 1);
  ArrayType *RowArrayTy = llvm::cast<ArrayType>(StructTy->getElementType(0));
  DXASSERT_NOMSG(RowArrayTy->getNumElements() >= 1 && RowArrayTy->getNumElements() <= 4);
  VectorType *RowTy = llvm::cast<VectorType>(RowArrayTy->getElementType());
  DXASSERT_NOMSG(RowTy->getNumElements() >= 1 && RowTy->getNumElements() <= 4);
  return HLMatrixType(RowTy->getElementType(), RowArrayTy->getNumElements(), RowTy->getNumElements());
}

HLMatrixType HLMatrixType::dyn_cast(Type *Ty) {
  return isa(Ty) ? cast(Ty) : HLMatrixType();
}