Merge branch 'main' into extract_vec

Author: abhishek-kaushik22

clang/lib/CodeGen/CodeGenModule.cpp

@@ -4562,15 +4562,17 @@ llvm::Constant *CodeGenModule::GetOrCreateMultiVersionResolver(GlobalDecl GD) {
     ResolverName += ".resolver";
   }
 
+  bool ShouldReturnIFunc =
+      getTarget().supportsIFunc() && !FD->isCPUSpecificMultiVersion();
+
   // If the resolver has already been created, just return it. This lookup may
   // yield a function declaration instead of a resolver on AArch64. That is
   // because we didn't know whether a resolver will be generated when we first
   // encountered a use of the symbol named after this resolver. Therefore,
   // targets which support ifuncs should not return here unless we actually
   // found an ifunc.
   llvm::GlobalValue *ResolverGV = GetGlobalValue(ResolverName);
-  if (ResolverGV &&
-      (isa<llvm::GlobalIFunc>(ResolverGV) || !getTarget().supportsIFunc()))
+  if (ResolverGV && (isa<llvm::GlobalIFunc>(ResolverGV) || !ShouldReturnIFunc))
     return ResolverGV;
 
   const CGFunctionInfo &FI = getTypes().arrangeGlobalDeclaration(GD);
@@ -4583,7 +4585,7 @@ llvm::Constant *CodeGenModule::GetOrCreateMultiVersionResolver(GlobalDecl GD) {
 
   // For cpu_specific, don't create an ifunc yet because we don't know if the
   // cpu_dispatch will be emitted in this translation unit.
-  if (getTarget().supportsIFunc() && !FD->isCPUSpecificMultiVersion()) {
+  if (ShouldReturnIFunc) {
     unsigned AS = getTypes().getTargetAddressSpace(FD->getType());
     llvm::Type *ResolverType =
         llvm::FunctionType::get(llvm::PointerType::get(DeclTy, AS), false);
@@ -4602,11 +4604,9 @@ llvm::Constant *CodeGenModule::GetOrCreateMultiVersionResolver(GlobalDecl GD) {
 
   llvm::Constant *Resolver = GetOrCreateLLVMFunction(
       ResolverName, DeclTy, GlobalDecl{}, /*ForVTable=*/false);
-  assert(isa<llvm::GlobalValue>(Resolver) &&
+  assert(isa<llvm::GlobalValue>(Resolver) && !ResolverGV &&
          "Resolver should be created for the first time");
   SetCommonAttributes(FD, cast<llvm::GlobalValue>(Resolver));
-  if (ResolverGV)
-    replaceDeclarationWith(ResolverGV, Resolver);
   return Resolver;
 }
 

clang/test/CodeGen/attr-cpuspecific.c

@@ -154,6 +154,12 @@ void usages(void) {
   CpuSpecificNoDispatch();
   // LINUX: @CpuSpecificNoDispatch.ifunc()
   // WINDOWS: @CpuSpecificNoDispatch()
+  //
+  // Adding another use of CpuSpecificNoDispatch reproduces the
+  // crash in https://github.com/llvm/llvm-project/issues/115299
+  CpuSpecificNoDispatch();
+  // LINUX: @CpuSpecificNoDispatch.ifunc()
+  // WINDOWS: @CpuSpecificNoDispatch()
   OrderDispatchUsageSpecific();
   // LINUX: @OrderDispatchUsageSpecific.ifunc()
   // WINDOWS: @OrderDispatchUsageSpecific()

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

@@ -52,6 +52,17 @@ def ApplyDecomposeTensorPackUnpackPatternsOp
   let assemblyFormat = "attr-dict";
 }
 
+def ApplyDecomposeTensorPadPatternsOp
+    : Op<Transform_Dialect, "apply_patterns.linalg.decompose_pad",
+         [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+  let description = [{
+    Collect patterns to decompose tensor.pad into e.g. tensor::EmptyOp,
+    linalg::FillOp and tensor::InsertSliceOp.
+  }];
+
+  let assemblyFormat = "attr-dict";
+}
+
 def ApplyFoldUnitExtentDimsViaReshapesPatternsOp : Op<Transform_Dialect,
     "apply_patterns.linalg.fold_unit_extent_dims_via_reshapes",
     [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

@@ -1503,8 +1503,8 @@ using OptimizeCopyFn =
 
 /// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp and
 /// InsertSliceOp. For now, only constant padding values are supported.
-struct GeneralizePadOpPattern : public OpRewritePattern<tensor::PadOp> {
-  GeneralizePadOpPattern(MLIRContext *context, PatternBenefit benefit = 1)
+struct DecomposePadOpPattern : public OpRewritePattern<tensor::PadOp> {
+  DecomposePadOpPattern(MLIRContext *context, PatternBenefit benefit = 1)
       : OpRewritePattern<tensor::PadOp>(context, benefit) {}
   LogicalResult matchAndRewrite(tensor::PadOp padOp,
                                 PatternRewriter &rewriter) const override;
@@ -1688,6 +1688,10 @@ void populateDecomposeConvolutionPatterns(RewritePatternSet &patterns,
 /// outer dims to be unit.
 void populateDecomposePackUnpackPatterns(RewritePatternSet &patterns);
 
+/// Populates patterns to decompose tensor.pad into e.g.
+/// tensor.empty, linalg.fill, tensor.insert_slice.
+void populateDecomposePadPatterns(RewritePatternSet &patterns);
+
 /// Populates patterns to transform linalg.conv_2d_xxx operations into
 /// linalg.generic (for img2col packing) and linalg.matmul.
 /// \see rewriteInIm2Col for more details.

mlir/lib/Conversion/TensorToLinalg/TensorToLinalg.cpp

@@ -25,5 +25,7 @@ using namespace mlir;
 //===----------------------------------------------------------------------===//
 
 void mlir::populateTensorToLinalgPatterns(RewritePatternSet &patterns) {
-  patterns.add<mlir::linalg::GeneralizePadOpPattern>(patterns.getContext());
+  // TODO: Add the remaining patterns, e.g. to decompose Pack/Unpack Ops.
+  // Alternatively, delete this file.
+  patterns.add<mlir::linalg::DecomposePadOpPattern>(patterns.getContext());
 }

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

@@ -234,6 +234,11 @@ void transform::ApplyDecomposeTensorPackUnpackPatternsOp::populatePatterns(
   linalg::populateDecomposePackUnpackPatterns(patterns);
 }
 
+void transform::ApplyDecomposeTensorPadPatternsOp::populatePatterns(
+    RewritePatternSet &patterns) {
+  linalg::populateDecomposePadPatterns(patterns);
+}
+
 void transform::ApplyFoldUnitExtentDimsViaReshapesPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   linalg::ControlDropUnitDims options;
@@ -3491,8 +3496,12 @@ transform::VectorizeChildrenAndApplyPatternsOp::applyToOne(
   // Add misc. vectorization patterns (e.g. for tensor.insert_slice)
   linalg::populateInsertSliceVectorizationPatterns(patterns);
 
-  if (getVectorizePadding())
+  if (getVectorizePadding()) {
     linalg::populatePadOpVectorizationPatterns(patterns);
+    // This creates an alternative path for lowering tensor.pad - by
+    // decomposing it into e.g. linalg.fill.
+    linalg::populateDecomposePadPatterns(patterns);
+  }
   vector::populateVectorStepLoweringPatterns(patterns);
 
   TrackingListener listener(state, *this);

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

@@ -921,7 +921,7 @@ LogicalResult mlir::linalg::CopyVectorizationPattern::matchAndRewrite(
 
 /// Filling `dest` using FillOp constant padding value if possible.
 /// Otherwise, generate a tensor::GenerateOp.
-Value GeneralizePadOpPattern::createFillOrGenerateOp(
+Value DecomposePadOpPattern::createFillOrGenerateOp(
     RewriterBase &rewriter, tensor::PadOp padOp, Value dest,
     const SmallVector<Value> &dynSizes) const {
   auto padValue = padOp.getConstantPaddingValue();
@@ -938,8 +938,8 @@ Value GeneralizePadOpPattern::createFillOrGenerateOp(
 }
 
 LogicalResult
-GeneralizePadOpPattern::matchAndRewrite(tensor::PadOp padOp,
-                                        PatternRewriter &rewriter) const {
+DecomposePadOpPattern::matchAndRewrite(tensor::PadOp padOp,
+                                       PatternRewriter &rewriter) const {
   // Given an OpFoldResult, return an index-typed value.
   auto getIdxValue = [&](OpFoldResult ofr) {
     if (auto val = llvm::dyn_cast_if_present<Value>(ofr))
@@ -1623,3 +1623,7 @@ void linalg::populateDecomposePackUnpackPatterns(RewritePatternSet &patterns) {
   // TODO: Add and test patterns for tensor.unpack
   patterns.add<DecomposeOuterUnitDimsPackOpPattern>(patterns.getContext());
 }
+
+void linalg::populateDecomposePadPatterns(RewritePatternSet &patterns) {
+  patterns.add<DecomposePadOpPattern>(patterns.getContext());
+}

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

@@ -2770,12 +2770,6 @@ void mlir::linalg::populateInsertSliceVectorizationPatterns(
 
 void mlir::linalg::populatePadOpVectorizationPatterns(
     RewritePatternSet &patterns, PatternBenefit baseBenefit) {
-  // TODO: The following pattern implements "decomposition" and
-  // optional "vectorization". Seperate "decomposition" into a sepereate
-  // pre-processing pattern group.
-  patterns.add<GeneralizePadOpPattern>(patterns.getContext(), baseBenefit);
-
-  // Try these specialized patterns first before resorting to the generic one.
   patterns.add<PadOpVectorizationWithTransferReadPattern,
                PadOpVectorizationWithTransferWritePattern,
                PadOpVectorizationWithInsertSlicePattern>(

mlir/test/Dialect/Linalg/generalize-pad-tensor.mlir renamed to mlir/test/Dialect/Linalg/decompose-pad-tensor.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt -split-input-file --test-linalg-transform-patterns="test-generalize-pad-tensor"  %s | FileCheck %s
+// RUN: mlir-opt -split-input-file --test-linalg-transform-patterns="test-decompose-pad-tensor"  %s | FileCheck %s
 
 // CHECK-LABEL:   func @generalize_pad_tensor_static_shape(
 // CHECK-SAME:                                             %[[IN:.*]]: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> {

mlir/test/Dialect/Linalg/vectorization-pad-patterns.mlir

@@ -202,6 +202,8 @@ module attributes {transform.with_named_sequence} {
     %func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func">
 
     transform.apply_patterns to %func_op {
+      // TODO: Split into two tests, one for each pattern
+      transform.apply_patterns.linalg.decompose_pad
       transform.apply_patterns.linalg.pad_vectorization
     } : !transform.op<"func.func">
     transform.yield
@@ -236,6 +238,8 @@ module attributes {transform.with_named_sequence} {
     %func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func">
 
     transform.apply_patterns to %func_op {
+      // TODO: Split into two tests, one for each pattern
+      transform.apply_patterns.linalg.decompose_pad
       transform.apply_patterns.linalg.pad_vectorization
     } : !transform.op<"func.func">
     transform.yield
@@ -270,6 +274,8 @@ module attributes {transform.with_named_sequence} {
     %func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func">
 
     transform.apply_patterns to %func_op {
+      // TODO: Split into two tests, one for each pattern
+      transform.apply_patterns.linalg.decompose_pad
       transform.apply_patterns.linalg.pad_vectorization
     } : !transform.op<"func.func">
     transform.yield