[mlir][LLVM] Add `BitcastStores` type-consistency pattern

[zero9178]

Summary

Current patterns attempt to immediately create type-consistent stores by bitcasting the value to the type-hint leading to needless code duplication.

This patch extracts that case into its own pattern, allowing other patterns to create type-inconsistent stores and have subsequent pattern applications turn it into a type-consistent store.

[gysit]

mlir/include/mlir/Dialect/LLVMIR/Transforms/TypeConsistency.h

73

nit: contradicts

mlir/test/Dialect/LLVMIR/type-consistency.mlir

365

Can you maybe add one test case that exercises the bitcast pattern in isolation so that it is exercised independent of the other patterns?

[zero9178]

Address review comments

[gysit]

LGTM!

mlir/include/mlir/Dialect/LLVMIR/Transforms/TypeConsistency.h

public:
   SplitStores(MLIRContext *context, unsigned int maxVectorStoreSplitSize)
       : OpRewritePattern(context),
         maxVectorStoreSplitSize(maxVectorStoreSplitSize) {}
 
   LogicalResult matchAndRewrite(StoreOp store,
                                 PatternRewriter &rewrite) const override;
 };
 
+/// Transforms type-inconsistent stores, aka stores where the type hint of
+/// the address contradict the value stored, by inserting a bitcast if
+/// possible.
+class BitcastStores : public OpRewritePattern<StoreOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(StoreOp store,
+                                PatternRewriter &rewriter) const override;
+};
+
 } // namespace LLVM
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LLVMIR_TRANSFORMS_TYPECONSISTENCY_H

mlir/lib/Dialect/LLVMIR/Transforms/TypeConsistency.cpp

static Type isElementTypeInconsistent(Value addr, Type expectedType) {
 
   if (elemType == expectedType)
     return nullptr;
 
   return elemType;
 }
 
 /// Checks that two types are the same or can be bitcast into one another.
-static bool areCastCompatible(DataLayout &layout, Type lhs, Type rhs) {
+static bool areBitcastCompatible(DataLayout &layout, Type lhs, Type rhs) {
   return lhs == rhs || (!isa<LLVMStructType, LLVMArrayType>(lhs) &&
                         !isa<LLVMStructType, LLVMArrayType>(rhs) &&
                         layout.getTypeSize(lhs) == layout.getTypeSize(rhs));
 }
 
 //===----------------------------------------------------------------------===//
 // AddFieldGetterToStructDirectUse
 //===----------------------------------------------------------------------===//
LogicalResult AddFieldGetterToStructDirectUse<LoadOp>::matchAndRewrite(
   Type inconsistentElementType =
       isElementTypeInconsistent(load.getAddr(), load.getType());
   if (!inconsistentElementType)
     return failure();
   Type firstType = getFirstSubelementType(inconsistentElementType);
   if (!firstType)
     return failure();
   DataLayout layout = DataLayout::closest(load);
-  if (!areCastCompatible(layout, firstType, load.getResult().getType()))
+  if (!areBitcastCompatible(layout, firstType, load.getResult().getType()))
     return failure();
 
   insertFieldIndirection<LoadOp>(load, rewriter, inconsistentElementType);
 
   // If the load does not use the first type but a type that can be casted from
   // it, add a bitcast and change the load type.
   if (firstType != load.getResult().getType()) {
     rewriter.setInsertionPointAfterValue(load.getResult());
if (!inconsistentElementType)
     return failure();
   Type firstType = getFirstSubelementType(inconsistentElementType);
   if (!firstType)
     return failure();
 
   DataLayout layout = DataLayout::closest(store);
   // Check that the first field has the right type or can at least be bitcast
   // to the right type.
-  if (!areCastCompatible(layout, firstType, store.getValue().getType()))
+  if (!areBitcastCompatible(layout, firstType, store.getValue().getType()))
     return failure();
 
   insertFieldIndirection<StoreOp>(store, rewriter, inconsistentElementType);
 
-  Value replaceValue = store.getValue();
-  if (firstType != store.getValue().getType()) {
-    rewriter.setInsertionPointAfterValue(store.getValue());
-    replaceValue = rewriter.create<BitcastOp>(store->getLoc(), firstType,
-                                              store.getValue());
-  }
-
   rewriter.updateRootInPlace(
-      store, [&]() { store.getValueMutable().assign(replaceValue); });
+      store, [&]() { store.getValueMutable().assign(store.getValue()); });
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // CanonicalizeAlignedGep
 //===----------------------------------------------------------------------===//
 
auto pos = rewriter.create<ConstantOp>(
         store.getLoc(),
         rewriter.getIntegerAttr(sourceType, (currentOffset - offset) * 8));
 
     auto shrOp = rewriter.create<LShrOp>(store.getLoc(), store.getValue(), pos);
 
     IntegerType fieldIntType = rewriter.getIntegerType(fieldSize * 8);
     Value valueToStore =
         rewriter.create<TruncOp>(store.getLoc(), fieldIntType, shrOp);
-    if (fieldIntType != type) {
-      // Bitcast to the right type. `fieldIntType` was explicitly created
-      // to be of the same size as `type` and must currently be a primitive as
-      // well.
-      valueToStore =
-          rewriter.create<BitcastOp>(store.getLoc(), type, valueToStore);
-    }
 
     // We create an `i8` indexed GEP here as that is the easiest (offset is
     // already known). Other patterns turn this into a type-consistent GEP.
     auto gepOp = rewriter.create<GEPOp>(store.getLoc(), address.getType(),
                                         rewriter.getI8Type(), address,
                                         ArrayRef<GEPArg>{currentOffset});
     rewriter.create<StoreOp>(store.getLoc(), valueToStore, gepOp);
 
     // No need to care about padding here since we already checked previously
     // that no padding exists in this range.
     currentOffset += fieldSize;
   }
 
   rewriter.eraseOp(store);
 
   return success();
 }
 
+LogicalResult BitcastStores::matchAndRewrite(StoreOp store,
+                                             PatternRewriter &rewriter) const {
+  Type sourceType = store.getValue().getType();
+  Type typeHint = isElementTypeInconsistent(store.getAddr(), sourceType);
+  if (!typeHint) {
+    // Nothing to do, since it is already consistent.
+    return failure();
+  }
+
+  auto dataLayout = DataLayout::closest(store);
+  if (!areBitcastCompatible(dataLayout, typeHint, sourceType))
+    return failure();
+
+  auto bitcastOp =
+      rewriter.create<BitcastOp>(store.getLoc(), typeHint, store.getValue());
+  rewriter.updateRootInPlace(
+      store, [&] { store.getValueMutable().assign(bitcastOp); });
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Type consistency pass
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct LLVMTypeConsistencyPass
     : public LLVM::impl::LLVMTypeConsistencyBase<LLVMTypeConsistencyPass> {
   void runOnOperation() override {
     RewritePatternSet rewritePatterns(&getContext());
     rewritePatterns.add<AddFieldGetterToStructDirectUse<LoadOp>>(&getContext());
     rewritePatterns.add<AddFieldGetterToStructDirectUse<StoreOp>>(
         &getContext());
     rewritePatterns.add<CanonicalizeAlignedGep>(&getContext());
     rewritePatterns.add<SplitStores>(&getContext(), maxVectorStoreSplitSize);
+    rewritePatterns.add<BitcastStores>(&getContext());
     FrozenRewritePatternSet frozen(std::move(rewritePatterns));
 
     if (failed(applyPatternsAndFoldGreedily(getOperation(), frozen)))
       signalPassFailure();
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> LLVM::createTypeConsistencyPass() {
   return std::make_unique<LLVMTypeConsistencyPass>();
 }

mlir/test/Dialect/LLVMIR/type-consistency.mlir

llvm.func @coalesced_store_floats(%arg: i64) {
 
   // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (f32, f32)>
   // CHECK: %[[SHR:.*]] = llvm.lshr %[[ARG]], %[[CST0]]
   // CHECK: %[[TRUNC:.*]] = llvm.trunc %[[SHR]] : i64 to i32
   // CHECK: %[[BIT_CAST:.*]] = llvm.bitcast %[[TRUNC]] : i32 to f32
   // CHECK: llvm.store %[[BIT_CAST]], %[[GEP]]
   // CHECK: %[[SHR:.*]] = llvm.lshr %[[ARG]], %[[CST32]] : i64
   // CHECK: %[[TRUNC:.*]] = llvm.trunc %[[SHR]] : i64 to i32
-  // CHECK: %[[BIT_CAST:.*]] = llvm.bitcast %[[TRUNC]] : i32 to f32
   // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 1] : (!llvm.ptr)  -> !llvm.ptr, !llvm.struct<"foo", (f32, f32)>
+  // CHECK: %[[BIT_CAST:.*]] = llvm.bitcast %[[TRUNC]] : i32 to f32
   // CHECK: llvm.store %[[BIT_CAST]], %[[GEP]]
   llvm.store %arg, %1 : i64, !llvm.ptr
   // CHECK-NOT: llvm.store %[[ARG]], %[[ALLOCA]]
   llvm.return
 }
 
 // -----
 

 }
 
 // -----
 
 // CHECK-LABEL: llvm.func @type_consistent_vector_store_other_type
 // CHECK-SAME: %[[ARG:.*]]: vector<4xi32>
 llvm.func @type_consistent_vector_store_other_type(%arg: vector<4xi32>) {
   %0 = llvm.mlir.constant(1 : i32) : i32
-  // CHECK: %[[BIT_CAST:.*]] = llvm.bitcast %[[ARG]] : vector<4xi32> to vector<4xf32>
   // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (vector<4xf32>)>
   %1 = llvm.alloca %0 x !llvm.struct<"foo", (vector<4xf32>)> : (i32) -> !llvm.ptr
   // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (vector<4xf32>)>
+  // CHECK: %[[BIT_CAST:.*]] = llvm.bitcast %[[ARG]] : vector<4xi32> to vector<4xf32>
   // CHECK: llvm.store %[[BIT_CAST]], %[[GEP]]
   llvm.store %arg, %1 : vector<4xi32>, !llvm.ptr
   // CHECK-NOT: llvm.store %[[ARG]], %[[ALLOCA]]
   llvm.return
 }