Enable store to PackedA RowMajor to support element wise

This change adds support for store to a PackedA Rewmajor matrix by
adding builtins.  It also fixes a few things in the
JointMatrixResolution pass to make sure slice indexing occur properly
and no cast issue happens.

Author: arnamoy10

IGC/BiFModule/Languages/OpenCL/PreRelease/IBiF_matrix.cl

@@ -258,3 +258,19 @@ INLINE void __builtin_spriv_OpJointMatrixStoreINTEL_Accumulator_RowMajor_8x16_i3
 INLINE void __builtin_spriv_OpJointMatrixStoreINTEL_Accumulator_ColumnMajor_8x8_i32_pi64_v8i8(char *mem, int8 row, int stride) {
     STORE_ACC_COL_MAJOR(mem, stride, 8)
 }
+
+#define STORE_PACKED_A_ROW_MAJOR(mem, stride, element_type, contrib_type, M) \
+    contrib_type *ptr = (contrib_type *)mem; \
+    int slid = get_sub_group_local_id(); \
+    int pack_factor = sizeof (contrib_type) / sizeof (element_type); \
+    stride = stride / pack_factor; \
+    for (int i = 0; i < M; i++) \
+         ptr[slid + i * stride] = row[i]; \
+
+INLINE void __builtin_spriv_OpJointMatrixStoreINTEL_PackedA_RowMajor_SG16_8x16_i16_pi64_v8i8(char *mem, short8 row, int stride) {
+    STORE_PACKED_A_ROW_MAJOR(mem, stride, short, short, 8)
+}
+
+INLINE void __builtin_spriv_OpJointMatrixStoreINTEL_PackedA_RowMajor_SG16_8x32_i8_pi64_v8i8(char *mem, short8 row, int stride) {
+    STORE_PACKED_A_ROW_MAJOR(mem, stride, char, short, 8)
+}

IGC/Compiler/Optimizer/OpenCLPasses/JointMatrixFuncsResolutionPass.cpp

@@ -446,16 +446,23 @@ static Type *getResolvedVectorElementType(Type *matrixType) {
     return ty->getElementType();
 }
 
-static int getSliceSize(const JointMatrixTypeDescription *desc) {
+static int getSliceSize(const JointMatrixTypeDescription *desc, Type *matTy) {
+    IGCLLVM::FixedVectorType *ty = dyn_cast<IGCLLVM::FixedVectorType>(matTy);
+    IGC_ASSERT_MESSAGE(ty, "Expecting vector type in calculating slice size");
+
+    IntegerType *vecElemType = dyn_cast<IntegerType>(ty->getElementType());
+    IGC_ASSERT_MESSAGE(vecElemType, "Expecting integer type for vector element.");
+
+    unsigned contribTypeWidth = vecElemType->getBitWidth();
     if (desc->layout == LayoutRowMajor) {
         return desc->rows;
     }
     if (desc->bitWidth != 0) {
         if (desc->layout == LayoutPackedA) {
-            return desc->rows * (32 / desc->bitWidth);
+            return desc->rows * (contribTypeWidth / desc->bitWidth);
         }
         if (desc->layout == LayoutPackedB) {
-            return 8 * (32 / desc->bitWidth);
+            return 8 * (contribTypeWidth / desc->bitWidth);
         }
     }
     IGC_ASSERT_MESSAGE(true, "Unexpected matrix layout.");
@@ -511,7 +518,7 @@ Value *JointMatrixFuncsResolutionPass::ResolveFill(CallInst *CI) {
     Type *matTy = ResolveType(CI->getType(), &desc);
 
     IRBuilder builder(CI);
-    const int sliceSize = getSliceSize(&desc);
+    const int sliceSize = getSliceSize(&desc, matTy);
     const int vectorSize = getResolvedVectorSize(matTy);
     /* Case with packing: */
     if (sliceSize > vectorSize) {
@@ -523,6 +530,14 @@ Value *JointMatrixFuncsResolutionPass::ResolveFill(CallInst *CI) {
         IGC_ASSERT_MESSAGE(false, "Malformed matrix slice.");
     }
 
+    if (fillValue->getType()->isPointerTy())
+    {
+        IntegerType *vectorElementType = dyn_cast<IntegerType>(getResolvedVectorElementType(matTy));
+        PointerType *PT = dyn_cast<PointerType>(fillValue->getType());
+        fillValue = builder.CreateBitCast(fillValue, PointerType::get(vectorElementType, PT->getAddressSpace()));
+        fillValue = builder.CreateLoad(vectorElementType, fillValue);
+    }
+
     Value *slice = UndefValue::get(matTy);
     for (int i = 0; i < vectorSize; i++) {
         slice = builder.CreateInsertElement(slice, fillValue, i);
@@ -534,9 +549,9 @@ Value *JointMatrixFuncsResolutionPass::ResolveFill(CallInst *CI) {
 
 Value *JointMatrixFuncsResolutionPass::ResolveWILength(CallInst *CI) {
     JointMatrixTypeDescription desc;
-    ResolveType(CI->getArgOperand(0)->getType(), &desc);
+    Type *matTy = ResolveType(CI->getArgOperand(0)->getType(), &desc);
 
-    const int sliceSize = getSliceSize(&desc);
+    const int sliceSize = getSliceSize(&desc, matTy);
     Value *lenght = ConstantInt::get(CI->getType(), sliceSize, "matrix.slice.size");
 
     CI->replaceAllUsesWith(lenght);
@@ -546,8 +561,8 @@ Value *JointMatrixFuncsResolutionPass::ResolveWILength(CallInst *CI) {
 
 template <class BuilderT>
 static Value *createSliceExtract
-      (BuilderT *builder, Value *matrix, Value *index, const JointMatrixTypeDescription *desc) {
-    const int sliceSize = getSliceSize(desc);
+      (BuilderT *builder, Value *matrix, Value *index, const JointMatrixTypeDescription *desc, Type *matTy) {
+    const int sliceSize = getSliceSize(desc, matTy);
     const int vectorSize = getResolvedVectorSize(matrix->getType());
     /* Unpacking: */
     if (sliceSize > vectorSize) {
@@ -568,12 +583,12 @@ Value *JointMatrixFuncsResolutionPass::ResolveSliceInsert(CallInst *CI) {
     IGCLLVM::FixedVectorType *matTy = dyn_cast<IGCLLVM::FixedVectorType>(rawMatTy);
 
     IRBuilder builder(CI);
-    const int sliceSize = getSliceSize(&desc);
+    const int sliceSize = getSliceSize(&desc, rawMatTy);
     const int vectorSize = getResolvedVectorSize(matTy);
 
     Value *slice = nullptr;
     if (sliceSize > vectorSize) {
-        Value *element = createSliceExtract(&builder, matrix, index, &desc);
+        Value *element = createSliceExtract(&builder, matrix, index, &desc, rawMatTy);
         if (!isa<IntegerType>(element->getType())) {
             unsigned vecElemSize = matTy->getElementType()->getScalarSizeInBits();
             element = builder.CreateBitCast(element, Type::getIntNTy(builder.getContext(), vecElemSize));
@@ -603,6 +618,10 @@ Value *JointMatrixFuncsResolutionPass::ResolveSliceInsert(CallInst *CI) {
         component = builder.CreateShl(component, offset);
         component = builder.CreateOr(element, component);
     }
+
+    IntegerType *vectorElementType = dyn_cast<IntegerType>(getResolvedVectorElementType(rawMatTy));
+    component = builder.CreateBitCast(component, vectorElementType);
+
     slice = builder.CreateInsertElement(matrix, component, index);
 
     InstsToErase.insert(CI);
@@ -617,9 +636,9 @@ Value *JointMatrixFuncsResolutionPass::ResolveSliceExtract(CallInst *CI) {
     Type *matTy = ResolveType(CI->getArgOperand(0)->getType(), &desc);
 
     IRBuilder builder(CI);
-    Value *element = createSliceExtract(&builder, matrix, index, &desc);
+    Value *element = createSliceExtract(&builder, matrix, index, &desc, matTy);
     /* Unpacking: */
-    const int sliceSize = getSliceSize(&desc);
+    const int sliceSize = getSliceSize(&desc, matTy);
     const int vectorSize = getResolvedVectorSize(matTy);
     if (sliceSize > vectorSize) {
         index = builder.CreateTruncOrBitCast(index, element->getType());
@@ -634,6 +653,10 @@ Value *JointMatrixFuncsResolutionPass::ResolveSliceExtract(CallInst *CI) {
         element = builder.CreateBitCast(element, CI->getType());
     }
 
+    // We need the bitcast, especially for half, as the function call that is
+    // being replaces has a half return type and the vectorElementType is i16
+    element = builder.CreateBitCast(element, CI->getType());
+
     CI->replaceAllUsesWith(element);
     InstsToErase.insert(CI);
     return element;