[mlir][spirv] Unify resources of different vector sizes

[antiagainst]

Summary

This commit extends UnifyAliasedResourcePass to handle the case
where aliased resources have different vector sizes. (It still
requires all scalar types to be of the same bitwidth.) This is
effectively reusing the code for handling different-bitwidth
scalar types.

Depends on D130669

mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp

 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/Dialect/SPIRV/Transforms/Passes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/SymbolTable.h"
-#include "mlir/Pass/AnalysisManager.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Debug.h"
 #include <algorithm>
 #include <iterator>
 
 #define DEBUG_TYPE "spirv-unify-aliased-resource"
static Type getRuntimeArrayElementType(Type type) {
 
   return rtArrayType.getElementType();
 }
 
 /// Given a list of resource element `types`, returns the index of the canonical
 /// resource that all resources should be unified into. Returns llvm::None if
 /// unable to unify.
 static Optional<int> deduceCanonicalResource(ArrayRef<spirv::SPIRVType> types) {
-  SmallVector<int> scalarNumBits, totalNumBits;
+  // scalarNumBits: contains all resources' scalar types' bit counts.
+  // vectorNumBits: only contains resources whose element types are vectors.
+  SmallVector<int> scalarNumBits, vectorNumBits, vectorIndices;
   scalarNumBits.reserve(types.size());
-  totalNumBits.reserve(types.size());
-  bool hasVector = false;
+  vectorNumBits.reserve(types.size());
+  vectorIndices.reserve(types.size());
 
-  for (spirv::SPIRVType type : types) {
+  for (const auto &indexedTypes : llvm::enumerate(types)) {
+    spirv::SPIRVType type = indexedTypes.value();
     assert(type.isScalarOrVector());
     if (auto vectorType = type.dyn_cast<VectorType>()) {
       if (vectorType.getNumElements() % 2 != 0)
         return llvm::None; // Odd-sized vector has special layout requirements.
 
       Optional<int64_t> numBytes = type.getSizeInBytes();
       if (!numBytes)
         return llvm::None;
 
       scalarNumBits.push_back(
           vectorType.getElementType().getIntOrFloatBitWidth());
-      totalNumBits.push_back(*numBytes * 8);
-      hasVector = true;
+      vectorNumBits.push_back(*numBytes * 8);
+      vectorIndices.push_back(indexedTypes.index());
     } else {
       scalarNumBits.push_back(type.getIntOrFloatBitWidth());
-      totalNumBits.push_back(scalarNumBits.back());
     }
   }
 
-  if (hasVector) {
+  if (!vectorNumBits.empty()) {
     // If there are vector types, require all element types to be the same for
     // now to simplify the transformation.
     if (!llvm::is_splat(scalarNumBits))
       return llvm::None;
 
-    // Choose the one with the largest bitwidth as the canonical resource, so
-    // that we can still keep vectorized load/store.
-    auto *maxVal = std::max_element(totalNumBits.begin(), totalNumBits.end());
+    // Choose the *vector* with the smallest bitwidth as the canonical resource,
+    // so that we can still keep vectorized load/store and avoid partial updates
+    // to large vectors.
+    auto *minVal = std::min_element(vectorNumBits.begin(), vectorNumBits.end());
     // Make sure that the canonical resource's bitwidth is divisible by others.
     // With out this, we cannot properly adjust the index later.
-    if (llvm::any_of(totalNumBits,
-                     [maxVal](int64_t bits) { return *maxVal % bits != 0; }))
+    if (llvm::any_of(vectorNumBits,
+                     [minVal](int64_t bits) { return bits % *minVal != 0; }))
       return llvm::None;
 
-    return std::distance(totalNumBits.begin(), maxVal);
+    return vectorIndices[std::distance(vectorNumBits.begin(), minVal)];
   }
 
   // All element types are scalars. Then choose the smallest bitwidth as the
   // cannonical resource to avoid subcomponent load/store.
   auto *minVal = std::min_element(scalarNumBits.begin(), scalarNumBits.end());
   if (llvm::any_of(scalarNumBits,
                    [minVal](int64_t bit) { return bit % *minVal != 0; }))
     return llvm::None;
if (srcElemType.isIntOrFloat() && dstElemType.isa<VectorType>()) {
       indices.push_back(
           rewriter.create<spirv::SModOp>(loc, i32Type, oldIndex, ratioValue));
 
       rewriter.replaceOpWithNewOp<spirv::AccessChainOp>(
           acOp, adaptor.base_ptr(), indices);
       return success();
     }
 
-    if (srcElemType.isIntOrFloat() && dstElemType.isIntOrFloat()) {
-      // The source indices are for a buffer with larger bitwidth scalar element
-      // types. Rewrite them into a buffer with smaller bitwidth element types.
-      // We only need to scale the last index.
+    if ((srcElemType.isIntOrFloat() && dstElemType.isIntOrFloat()) ||
+        (srcElemType.isa<VectorType>() && dstElemType.isa<VectorType>())) {
+      // The source indices are for a buffer with larger bitwidth scalar/vector
+      // element types. Rewrite them into a buffer with smaller bitwidth element
+      // types. We only need to scale the last index.
       int srcNumBits = *srcElemType.getSizeInBytes();
       int dstNumBits = *dstElemType.getSizeInBytes();
       assert(srcNumBits > dstNumBits && srcNumBits % dstNumBits == 0);
       int ratio = srcNumBits / dstNumBits;
       auto ratioValue = rewriter.create<spirv::ConstantOp>(
           loc, i32Type, rewriter.getI32IntegerAttr(ratio));
 
       auto indices = llvm::to_vector<4>(acOp.indices());
       Value oldIndex = indices.back();
       indices.back() =
           rewriter.create<spirv::IMulOp>(loc, i32Type, oldIndex, ratioValue);
 
       rewriter.replaceOpWithNewOp<spirv::AccessChainOp>(
           acOp, adaptor.base_ptr(), indices);
       return success();
     }
 
-    return rewriter.notifyMatchFailure(acOp, "unsupported src/dst types");
+    return rewriter.notifyMatchFailure(
+        acOp, "unsupported src/dst types for spv.AccessChain");
   }
 };
 
 struct ConvertLoad : public ConvertAliasResource<spirv::LoadOp> {
   using ConvertAliasResource::ConvertAliasResource;
 
   LogicalResult
   matchAndRewrite(spirv::LoadOp loadOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    auto srcElemType =
-        loadOp.ptr().getType().cast<spirv::PointerType>().getPointeeType();
-    auto dstElemType =
-        adaptor.ptr().getType().cast<spirv::PointerType>().getPointeeType();
-    if (!srcElemType.isIntOrFloat() || !dstElemType.isIntOrFloat())
-      return rewriter.notifyMatchFailure(loadOp, "not scalar type");
+    auto srcPtrType = loadOp.ptr().getType().cast<spirv::PointerType>();
+    auto srcElemType = srcPtrType.getPointeeType().cast<spirv::SPIRVType>();
+    auto dstPtrType = adaptor.ptr().getType().cast<spirv::PointerType>();
+    auto dstElemType = dstPtrType.getPointeeType().cast<spirv::SPIRVType>();
 
     Location loc = loadOp.getLoc();
     auto newLoadOp = rewriter.create<spirv::LoadOp>(loc, adaptor.ptr());
     if (srcElemType == dstElemType) {
       rewriter.replaceOp(loadOp, newLoadOp->getResults());
       return success();
     }
 
     if (areSameBitwidthScalarType(srcElemType, dstElemType)) {
       auto castOp = rewriter.create<spirv::BitcastOp>(loc, srcElemType,
                                                       newLoadOp.value());
       rewriter.replaceOp(loadOp, castOp->getResults());
 
       return success();
     }
 
-    // The source and destination have scalar types of different bitwidths.
-    // For such cases, we need to load multiple smaller bitwidth values and
-    // construct a larger bitwidth one.
+    if ((srcElemType.isIntOrFloat() && dstElemType.isIntOrFloat()) ||
+        (srcElemType.isa<VectorType>() && dstElemType.isa<VectorType>())) {
+      // The source and destination have scalar types of different bitwidths, or
+      // vector types of different component counts. For such cases, we load
+      // multiple smaller bitwidth values and construct a larger bitwidth one.
 
-    int srcNumBits = srcElemType.getIntOrFloatBitWidth();
-    int dstNumBits = dstElemType.getIntOrFloatBitWidth();
+      int srcNumBits = *srcElemType.getSizeInBytes() * 8;
+      int dstNumBits = *dstElemType.getSizeInBytes() * 8;
       assert(srcNumBits > dstNumBits && srcNumBits % dstNumBits == 0);
       int ratio = srcNumBits / dstNumBits;
       if (ratio > 4)
         return rewriter.notifyMatchFailure(loadOp, "more than 4 components");
 
       SmallVector<Value> components;
       components.reserve(ratio);
       components.push_back(newLoadOp);
 
       auto acOp = adaptor.ptr().getDefiningOp<spirv::AccessChainOp>();
       if (!acOp)
         return rewriter.notifyMatchFailure(loadOp, "ptr not spv.AccessChain");
 
       auto i32Type = rewriter.getI32Type();
       Value oneValue = spirv::ConstantOp::getOne(i32Type, loc, rewriter);
       auto indices = llvm::to_vector<4>(acOp.indices());
       for (int i = 1; i < ratio; ++i) {
         // Load all subsequent components belonging to this element.
-      indices.back() = rewriter.create<spirv::IAddOp>(loc, i32Type,
-                                                      indices.back(), oneValue);
-      auto componentAcOp =
-          rewriter.create<spirv::AccessChainOp>(loc, acOp.base_ptr(), indices);
-      // Assuming little endian, this reads lower-ordered bits of the number to
-      // lower-numbered components of the vector.
-      components.push_back(rewriter.create<spirv::LoadOp>(loc, componentAcOp));
+        indices.back() = rewriter.create<spirv::IAddOp>(
+            loc, i32Type, indices.back(), oneValue);
+        auto componentAcOp = rewriter.create<spirv::AccessChainOp>(
+            loc, acOp.base_ptr(), indices);
+        // Assuming little endian, this reads lower-ordered bits of the number
+        // to lower-numbered components of the vector.
+        components.push_back(
+            rewriter.create<spirv::LoadOp>(loc, componentAcOp));
       }
 
       // Create a vector of the components and then cast back to the larger
-    // bitwidth element type. For spv.bitcast, the lower-numbered components of
-    // the vector map to lower-ordered bits of the larger bitwidth element type.
-    auto vectorType = VectorType::get({ratio}, dstElemType);
+      // bitwidth element type. For spv.bitcast, the lower-numbered components
+      // of the vector map to lower-ordered bits of the larger bitwidth element
+      // type.
+      Type vectorType = srcElemType;
+      if (!srcElemType.isa<VectorType>())
+        vectorType = VectorType::get({ratio}, dstElemType);
       Value vectorValue = rewriter.create<spirv::CompositeConstructOp>(
           loc, vectorType, components);
-    rewriter.replaceOpWithNewOp<spirv::BitcastOp>(loadOp, srcElemType,
-                                                  vectorValue);
+      if (!srcElemType.isa<VectorType>())
+        vectorValue =
+            rewriter.create<spirv::BitcastOp>(loc, srcElemType, vectorValue);
+      rewriter.replaceOp(loadOp, vectorValue);
       return success();
     }
+
+    return rewriter.notifyMatchFailure(
+        loadOp, "unsupported src/dst types for spv.Load");
+  }
 };
 
 struct ConvertStore : public ConvertAliasResource<spirv::StoreOp> {
   using ConvertAliasResource::ConvertAliasResource;
 
   LogicalResult
   matchAndRewrite(spirv::StoreOp storeOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {

mlir/test/Dialect/SPIRV/Transforms/unify-aliased-resource.mlir

spv.func @store_different_scalar_bitwidth(%i0: i32, %i1: i32) "None" {
     %addr1 = spv.mlir.addressof @var01s_i64 : !spv.ptr<!spv.struct<(!spv.rtarray<i64, stride=4> [0])>, StorageBuffer>
     %ac1 = spv.AccessChain %addr1[%c0, %i1] : !spv.ptr<!spv.struct<(!spv.rtarray<i64, stride=4> [0])>, StorageBuffer>, i32, i32
     // expected-error@+1 {{failed to legalize operation 'spv.Store'}}
     spv.Store "StorageBuffer" %ac1, %i64val : i64
 
     spv.Return
   }
 }
+
+// -----
+
+spv.module Logical GLSL450 {
+  spv.GlobalVariable @var01_scalar bind(0, 1) {aliased} : !spv.ptr<!spv.struct<(!spv.rtarray<f32, stride=4> [0])>, StorageBuffer>
+  spv.GlobalVariable @var01_vec2 bind(0, 1) {aliased} : !spv.ptr<!spv.struct<(!spv.rtarray<vector<2xf32>, stride=8> [0])>, StorageBuffer>
+  spv.GlobalVariable @var01_vec4 bind(0, 1) {aliased} : !spv.ptr<!spv.struct<(!spv.rtarray<vector<4xf32>, stride=16> [0])>, StorageBuffer>
+
+  spv.func @load_different_vector_sizes(%i0: i32) -> vector<4xf32> "None" {
+    %c0 = spv.Constant 0 : i32
+
+    %addr0 = spv.mlir.addressof @var01_vec4 : !spv.ptr<!spv.struct<(!spv.rtarray<vector<4xf32>, stride=16> [0])>, StorageBuffer>
+    %ac0 = spv.AccessChain %addr0[%c0, %i0] : !spv.ptr<!spv.struct<(!spv.rtarray<vector<4xf32>, stride=16> [0])>, StorageBuffer>, i32, i32
+    %vec4val = spv.Load "StorageBuffer" %ac0 : vector<4xf32>
+
+    %addr1 = spv.mlir.addressof @var01_scalar : !spv.ptr<!spv.struct<(!spv.rtarray<f32, stride=4> [0])>, StorageBuffer>
+    %ac1 = spv.AccessChain %addr1[%c0, %i0] : !spv.ptr<!spv.struct<(!spv.rtarray<f32, stride=4> [0])>, StorageBuffer>, i32, i32
+    %scalarval = spv.Load "StorageBuffer" %ac1 : f32
+
+    %val = spv.CompositeInsert %scalarval, %vec4val[0 : i32] : f32 into vector<4xf32>
+    spv.ReturnValue %val : vector<4xf32>
+  }
+}
+
+// CHECK-LABEL: spv.module
+
+// CHECK-NOT: @var01_scalar
+// CHECK-NOT: @var01_vec4
+//     CHECK: spv.GlobalVariable @var01_vec2 bind(0, 1) : !spv.ptr<{{.+}}>
+// CHECK-NOT: @var01_scalar
+// CHECK-NOT: @var01_vec4
+
+//     CHECK: spv.func @load_different_vector_sizes(%[[IDX:.+]]: i32)
+//     CHECK:   %[[ZERO:.+]] = spv.Constant 0 : i32
+//     CHECK:   %[[ADDR:.+]] = spv.mlir.addressof @var01_vec2
+//     CHECK:   %[[TWO:.+]] = spv.Constant 2 : i32
+//     CHECK:   %[[IDX0:.+]] = spv.IMul %[[IDX]], %[[TWO]] : i32
+//     CHECK:   %[[AC0:.+]] = spv.AccessChain %[[ADDR]][%[[ZERO]], %[[IDX0]]]
+//     CHECK:   %[[LD0:.+]] = spv.Load "StorageBuffer" %[[AC0]] : vector<2xf32>
+//     CHECK:   %[[ONE:.+]] = spv.Constant 1 : i32
+//     CHECK:   %[[IDX1:.+]] = spv.IAdd %0, %[[ONE]] : i32
+//     CHECK:   %[[AC1:.+]] = spv.AccessChain %[[ADDR]][%[[ZERO]], %[[IDX1]]]
+//     CHECK:   %[[LD1:.+]] = spv.Load "StorageBuffer" %[[AC1]] : vector<2xf32>
+//     CHECK:   spv.CompositeConstruct %[[LD0]], %[[LD1]] : (vector<2xf32>, vector<2xf32>) -> vector<4xf32>
+
+//     CHECK:   %[[ADDR:.+]] = spv.mlir.addressof @var01_vec2
+//     CHECK:   %[[TWO:.+]] = spv.Constant 2 : i32
+//     CHECK:   %[[DIV:.+]] = spv.SDiv %[[IDX]], %[[TWO]] : i32
+//     CHECK:   %[[MOD:.+]] = spv.SMod %[[IDX]], %[[TWO]] : i32
+//     CHECK:   %[[AC:.+]] = spv.AccessChain %[[ADDR]][%[[ZERO]], %[[DIV]], %[[MOD]]]
+//     CHECK:   %[[LD:.+]] = spv.Load "StorageBuffer" %[[AC]] : f32