Fix element type of target attributes in oneToOneRewrite when converting to llvm

Author: MengmengSun

mlir/lib/Conversion/LLVMCommon/Pattern.cpp

@@ -331,10 +331,35 @@ LogicalResult LLVM::detail::oneToOneRewrite(
       return failure();
   }
 
+  // If the targetAttrs contains DenseElementsAttr,
+  // and the element type of the DenseElementsAttr and result type is
+  // inconsistent after the conversion of result types, we need to convert the
+  // element type of the DenseElementsAttr to the target type by creating a new
+  // DenseElementsAttr with the converted element type, and use the new
+  // DenseElementsAttr to replace the old one in the targetAttrs
+  SmallVector<NamedAttribute> convertedAttrs;
+  for (auto attr : targetAttrs) {
+    if (auto denseAttr = dyn_cast<DenseElementsAttr>(attr.getValue())) {
+      VectorType vectorType = dyn_cast<VectorType>(denseAttr.getType());
+      if (vectorType) {
+        auto convertedElementType =
+            typeConverter.convertType(vectorType.getElementType());
+        VectorType convertedVectorType =
+            VectorType::get(vectorType.getShape(), convertedElementType,
+                            vectorType.getScalableDims());
+        convertedAttrs.emplace_back(
+            attr.getName(), DenseElementsAttr::getFromRawBuffer(
+                                convertedVectorType, denseAttr.getRawData()));
+      }
+    } else {
+      convertedAttrs.push_back(attr);
+    }
+  }
+
   // Create the operation through state since we don't know its C++ type.
   Operation *newOp =
       rewriter.create(op->getLoc(), rewriter.getStringAttr(targetOp), operands,
-                      resultTypes, targetAttrs);
+                      resultTypes, convertedAttrs);
 
   setNativeProperties(newOp, overflowFlags);
 

mlir/test/Conversion/ArithToLLVM/arith-to-llvm.mlir

@@ -428,7 +428,7 @@ func.func @fcmp(f32, f32) -> () {
 
 // CHECK-LABEL: @index_vector
 func.func @index_vector(%arg0: vector<4xindex>) {
-  // CHECK: %[[CST:.*]] = llvm.mlir.constant(dense<[0, 1, 2, 3]> : vector<4xindex>) : vector<4xi64>
+  // CHECK: %[[CST:.*]] = llvm.mlir.constant(dense<[0, 1, 2, 3]> : vector<4xi64>) : vector<4xi64>
   %0 = arith.constant dense<[0, 1, 2, 3]> : vector<4xindex>
   // CHECK: %[[V:.*]] = llvm.add %{{.*}}, %[[CST]] : vector<4xi64>
   %1 = arith.addi %arg0, %0 : vector<4xindex>
@@ -437,6 +437,21 @@ func.func @index_vector(%arg0: vector<4xindex>) {
 
 // -----
 
+// CHECK-LABEL: @f8E4M3FN_vector
+func.func @f8E4M3FN_vector() -> vector<4xf8E4M3FN> {
+  // CHECK: %[[CST0:.*]] = llvm.mlir.constant(dense<0> : vector<4xi8>) : vector<4xi8>
+  %0 = arith.constant dense<0.000000e+00> : vector<4xf8E4M3FN>
+  // CHECK: %[[CST1:.*]] = llvm.mlir.constant(dense<[56, 64, 68, 72]> : vector<4xi8>) : vector<4xi8>
+  %1 = arith.constant dense<[1.0, 2.0, 3.0, 4.0]> : vector<4xf8E4M3FN>
+  // CHECK: %[[V:.*]] = llvm.mlir.constant(dense<[56, 64, 68, 72]> : vector<4xi8>) : vector<4xi8>
+  %2 = arith.addf %0, %1 : vector<4xf8E4M3FN>
+  // CHECK: %[[RES:.*]] = builtin.unrealized_conversion_cast %[[V]] : vector<4xi8> to vector<4xf8E4M3FN>
+  // CHECK-NEXT: return %[[RES]] : vector<4xf8E4M3FN>
+  func.return %2 : vector<4xf8E4M3FN>
+}
+
+// -----
+
 // CHECK-LABEL: @bitcast_1d
 func.func @bitcast_1d(%arg0: vector<2xf32>) {
   // CHECK: llvm.bitcast %{{.*}} : vector<2xf32> to vector<2xi32>