[SYCL-MLIR] Use VectorType for ext_vector_type (#7342)

whitneywhtsang · web-flow · commit 48f9e5be4d11 · 2022-11-10T12:42:21.000-05:00
Globals of `ext_vector_type` type used to represent as `memref<3xi64>`, which lowers to array type pointer in LLVM. This can cause link failures when files are compiled from different compilers. This PR changes globals of `ext_vector_type` type to be represented as `memref<vector<3xi64>>`, which lowers to vector type pointer, the same type as directly compile from `clang`. Note: After this PR, `parallel_for.cpp` can run successfully. Test case added in intel/llvm-test-suite#1374. Signed-off-by: Tsang, Whitney <whitney.tsang@intel.com>
diff --git a/polygeist/tools/cgeist/Lib/CGExpr.cc b/polygeist/tools/cgeist/Lib/CGExpr.cc
@@ -29,23 +29,16 @@ MLIRScanner::VisitExtVectorElementExpr(clang::ExtVectorElementExpr *expr) {
   expr->getEncodedElementAccess(indices);
   assert(indices.size() == 1 &&
          "The support for higher dimensions to be implemented.");
-  auto idx = castToIndex(getMLIRLocation(expr->getAccessorLoc()),
-                         builder.create<ConstantIntOp>(loc, indices[0], 32));
   assert(base.isReference);
-  base.isReference = false;
-  auto mt = base.val.getType().cast<MemRefType>();
-  auto shape = std::vector<int64_t>(mt.getShape());
-  if (shape.size() == 1) {
-    shape[0] = -1;
-  } else {
-    shape.erase(shape.begin());
-  }
-  auto mt0 =
-      mlir::MemRefType::get(shape, mt.getElementType(),
-                            MemRefLayoutAttrInterface(), mt.getMemorySpace());
-  base.val = builder.create<polygeist::SubIndexOp>(loc, mt0, base.val,
-                                                   getConstantIndex(0));
-  return CommonArrayLookup(base, idx, base.isReference);
+  assert(base.val.getType().isa<MemRefType>() &&
+         "Expecting ExtVectorElementExpr to have memref type");
+  auto MT = base.val.getType().cast<MemRefType>();
+  assert(MT.getElementType().isa<mlir::VectorType>() &&
+         "Expecting ExtVectorElementExpr to have memref of vector elements");
+  auto Idx = builder.create<ConstantIntOp>(loc, indices[0], 64);
+  mlir::Value Val = base.getValue(builder);
+  return ValueCategory(builder.create<LLVM::ExtractElementOp>(loc, Val, Idx),
+                       /*IsReference*/ false);
 }
 
 ValueCategory MLIRScanner::VisitConstantExpr(clang::ConstantExpr *expr) {
@@ -1690,7 +1683,6 @@ ValueCategory MLIRScanner::VisitCastExpr(CastExpr *E) {
       lres.dump();
     }
 #endif
-    assert(prev.isReference);
     return ValueCategory(lres, /*isReference*/ false);
   }
   case clang::CastKind::CK_IntegralToFloating: {
diff --git a/polygeist/tools/cgeist/Lib/CodeGenTypes.cc b/polygeist/tools/cgeist/Lib/CodeGenTypes.cc
@@ -1453,6 +1453,8 @@ mlir::Type CodeGenTypes::getMLIRType(clang::QualType QT, bool *ImplicitRef,
     bool SubRef = false;
     auto ET = getMLIRType(AT->getElementType(), &SubRef, AllowMerge);
     int64_t Size = AT->getNumElements();
+    if (isa<clang::ExtVectorType>(T))
+      return mlir::VectorType::get(Size, ET);
     if (MemRefABI && SubRef) {
       auto MT = ET.cast<MemRefType>();
       auto Shape2 = std::vector<int64_t>(MT.getShape());
@@ -1530,6 +1532,10 @@ mlir::Type CodeGenTypes::getMLIRType(clang::QualType QT, bool *ImplicitRef,
     }
 
     if (isa<clang::VectorType>(PTT) || isa<clang::ComplexType>(PTT)) {
+      if (auto VT = SubType.dyn_cast<mlir::VectorType>())
+        // FIXME: We should create memref of rank 0.
+        // Details: https://github.com/intel/llvm/issues/7354
+        return mlir::MemRefType::get(Outer, SubType);
       if (SubType.isa<MemRefType>()) {
         assert(SubRef);
         auto MT = SubType.cast<MemRefType>();
diff --git a/polygeist/tools/cgeist/Lib/clang-mlir.cc b/polygeist/tools/cgeist/Lib/clang-mlir.cc
@@ -46,6 +46,8 @@
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
+#define DEBUG_TYPE "cgeist"
+
 using namespace clang;
 using namespace llvm;
 using namespace clang::driver;
@@ -2086,19 +2088,19 @@ MLIRASTConsumer::GetOrCreateGlobal(const ValueDecl *FD, std::string prefix,
   mlir::Type rt = getTypes().getMLIRType(FD->getType());
   auto *VD = dyn_cast<VarDecl>(FD);
   LLVM_DEBUG({
-    if (!VD)
-      FD->dump();
+    if (!VD) {
+      llvm::dbgs() << "GetOrCreateGlobal ";
+      VD->dump(llvm::dbgs());
+    }
   });
   VD = VD->getCanonicalDecl();
   unsigned memspace = VD ? CGM.getContext().getTargetAddressSpace(
                                CGM.GetGlobalVarAddressSpace(VD))
                          : CGM.getDataLayout().getDefaultGlobalsAddressSpace();
   bool isArray = isa<clang::ArrayType>(FD->getType());
-  bool isExtVectorType =
-      isa<clang::ExtVectorType>(FD->getType()->getUnqualifiedDesugaredType());
 
   mlir::MemRefType mr;
-  if (!isArray && !isExtVectorType) {
+  if (!isArray) {
     mr = mlir::MemRefType::get({}, rt, {}, memspace);
   } else {
     auto mt = rt.cast<mlir::MemRefType>();
diff --git a/polygeist/tools/cgeist/Test/Verification/ext_vector_type.cpp b/polygeist/tools/cgeist/Test/Verification/ext_vector_type.cpp
@@ -1,4 +1,7 @@
 // RUN: cgeist %s --function=* -S | FileCheck %s
+// RUN: cgeist %s --function=* -S -emit-llvm | FileCheck %s --check-prefix=LLVM
+
+#include <cstddef>
 
 typedef size_t size_t_vec __attribute__((ext_vector_type(3)));
 
@@ -12,13 +15,30 @@ size_t evt2() {
   return stv.x;
 }
 
-// CHECK:   func.func @_Z3evtv() -> i32 attributes {llvm.linkage = #llvm.linkage<external>} {
-// CHECK-NEXT:    %alloca = memref.alloca() : memref<1x3xi32>
-// CHECK-NEXT:    %0 = affine.load %alloca[0, 0] : memref<1x3xi32>
-// CHECK-NEXT:    return %0 : i32
+// CHECK:   func.func @_Z3evtv() -> i64 attributes {llvm.linkage = #llvm.linkage<external>} {
+// CHECK-NEXT:    %c0_i64 = arith.constant 0 : i64
+// CHECK-NEXT:    %alloca = memref.alloca() : memref<1xvector<3xi64>>
+// CHECK-NEXT:    %0 = affine.load %alloca[0] : memref<1xvector<3xi64>>
+// CHECK-NEXT:    %1 = llvm.extractelement %0[%c0_i64 : i64] : vector<3xi64>
+// CHECK-NEXT:    return %1 : i64
 // CHECK-NEXT:    }
-// CHECK:   func.func @_Z4evt2v() -> i32 attributes {llvm.linkage = #llvm.linkage<external>} {
-// CHECK-NEXT:     %0 = memref.get_global @stv : memref<3xi32>
-// CHECK-NEXT:     %1 = affine.load %0[0] : memref<3xi32>
-// CHECK-NEXT:     return %1 : i32
+// CHECK:   func.func @_Z4evt2v() -> i64 attributes {llvm.linkage = #llvm.linkage<external>} {
+// CHECK-NEXT:     %c0_i64 = arith.constant 0 : i64
+// CHECK-NEXT:     %0 = memref.get_global @stv : memref<vector<3xi64>>
+// CHECK-NEXT:     %alloca = memref.alloca() : memref<1xindex>
+// CHECK-NEXT:     %reshape = memref.reshape %0(%alloca) : (memref<vector<3xi64>>, memref<1xindex>) -> memref<1xvector<3xi64>>
+// CHECK-NEXT:     %1 = affine.load %reshape[0] : memref<1xvector<3xi64>>
+// CHECK-NEXT:     %2 = llvm.extractelement %1[%c0_i64 : i64] : vector<3xi64>
+// CHECK-NEXT:     return %2 : i64
 // CHECK-NEXT:     }
+
+// LLVM:       @stv = external global <3 x i64>
+// LLVM-LABEL: define i64 @_Z3evtv() !dbg !3 {
+// LLVM-NEXT:   %1 = alloca <3 x i64>, align 32, !dbg !7
+// LLVM-NEXT:   %2 = load <3 x i64>, <3 x i64>* %1, align 32
+// LLVM-NEXT:   %3 = extractelement <3 x i64> %2, i64 0
+// LLVM-NEXT:   ret i64 %3
+// LLVM-LABEL: define i64 @_Z4evt2v() !dbg !9 {
+// LLVM-NEXT:   %1 = load <3 x i64>, <3 x i64>* @stv, align 32
+// LLVM-NEXT:   %2 = extractelement <3 x i64> %1, i64 0
+// LLVM-NEXT:   ret i64 %2
diff --git a/polygeist/tools/cgeist/Test/Verification/sycl/constructors.cpp b/polygeist/tools/cgeist/Test/Verification/sycl/constructors.cpp
@@ -14,13 +14,26 @@
 // CHECK-DAG: memref.global @__spirv_BuiltInNumSubgroups : memref<i32, 1>
 // CHECK-DAG: memref.global @__spirv_BuiltInSubgroupMaxSize : memref<i32, 1>
 // CHECK-DAG: memref.global @__spirv_BuiltInSubgroupSize : memref<i32, 1>
-// CHECK-DAG: memref.global @__spirv_BuiltInLocalInvocationId : memref<3xi64, 1>
-// CHECK-DAG: memref.global @__spirv_BuiltInWorkgroupId : memref<3xi64, 1>
-// CHECK-DAG: memref.global @__spirv_BuiltInWorkgroupSize : memref<3xi64, 1>
-// CHECK-DAG: memref.global @__spirv_BuiltInNumWorkgroups : memref<3xi64, 1>
-// CHECK-DAG: memref.global @__spirv_BuiltInGlobalOffset : memref<3xi64, 1>
-// CHECK-DAG: memref.global @__spirv_BuiltInGlobalSize : memref<3xi64, 1>
-// CHECK-DAG: memref.global @__spirv_BuiltInGlobalInvocationId : memref<3xi64, 1>
+// CHECK-DAG: memref.global @__spirv_BuiltInLocalInvocationId : memref<vector<3xi64>, 1>
+// CHECK-DAG: memref.global @__spirv_BuiltInWorkgroupId : memref<vector<3xi64>, 1>
+// CHECK-DAG: memref.global @__spirv_BuiltInWorkgroupSize : memref<vector<3xi64>, 1>
+// CHECK-DAG: memref.global @__spirv_BuiltInNumWorkgroups : memref<vector<3xi64>, 1>
+// CHECK-DAG: memref.global @__spirv_BuiltInGlobalOffset : memref<vector<3xi64>, 1>
+// CHECK-DAG: memref.global @__spirv_BuiltInGlobalSize : memref<vector<3xi64>, 1>
+// CHECK-DAG: memref.global @__spirv_BuiltInGlobalInvocationId : memref<vector<3xi64>, 1>
+
+// CHECK-LLVM-DAG: @__spirv_BuiltInSubgroupLocalInvocationId = external addrspace(1) global i32
+// CHECK-LLVM-DAG: @__spirv_BuiltInSubgroupId = external addrspace(1) global i32
+// CHECK-LLVM-DAG: @__spirv_BuiltInNumSubgroups = external addrspace(1) global i32
+// CHECK-LLVM-DAG: @__spirv_BuiltInSubgroupMaxSize = external addrspace(1) global i32
+// CHECK-LLVM-DAG: @__spirv_BuiltInSubgroupSize = external addrspace(1) global i32
+// CHECK-LLVM-DAG: @__spirv_BuiltInLocalInvocationId = external addrspace(1) global <3 x i64>
+// CHECK-LLVM-DAG: @__spirv_BuiltInWorkgroupId = external addrspace(1) global <3 x i64>
+// CHECK-LLVM-DAG: @__spirv_BuiltInWorkgroupSize = external addrspace(1) global <3 x i64>
+// CHECK-LLVM-DAG: @__spirv_BuiltInNumWorkgroups = external addrspace(1) global <3 x i64>
+// CHECK-LLVM-DAG: @__spirv_BuiltInGlobalOffset = external addrspace(1) global <3 x i64>
+// CHECK-LLVM-DAG: @__spirv_BuiltInGlobalSize = external addrspace(1) global <3 x i64>
+// CHECK-LLVM-DAG: @__spirv_BuiltInGlobalInvocationId = external addrspace(1) global <3 x i64>
 
 // Ensure the spirv functions that reference these globals are not filtered out
 // CHECK-DAG: func.func @_Z28__spirv_GlobalInvocationId_xv()
