[mlir][StandardToSPIRV] Handle i1 case for lowering memref.load/store op

[hanchung]

Summary

This patch unconditionally converts i1 types to i8 types on memrefs. If the
extensions or capabilities are not met, they will be converted to i32. Hence the
logic in IntLoadPattern and IntStorePattern are also updated.

Also added the implementation of SPIRVTypeConverter::getOptions().

[hanchung]

fix debug build

[hanchung]

use build-in function

[hanchung]

I hit a really weird thing... and I can only repro with this PR...

module attributes {
  spv.target_env = #spv.target_env<
    #spv.vce<v1.0, [Int16, StorageBuffer16BitAccess, Shader],
    [SPV_KHR_storage_buffer_storage_class, SPV_KHR_16bit_storage]>, {}>
} {

func @load_i8(%arg0: memref<i8>) -> i8 {
  %0 = memref.load %arg0[] : memref<i8>
  return %0 : i8
}

}

The above one can be lowered to SPIR-V, but the below one can not

module attributes {
  spv.target_env = #spv.target_env<
    #spv.vce<v1.0, [Int16, StorageBuffer16BitAccess, Shader],
    [SPV_KHR_storage_buffer_storage_class, SPV_KHR_16bit_storage]>, {}>
} {

func @load_i1(%arg0: memref<i1>) -> i1 {
  %0 = memref.load %arg0[] : memref<i1>
  return %0 : i1
}

}

They are expected to generate the same IRs. And they do when there are no users.

The i1 case only works if it does not have users, like in the test.

I dumped operands[0] in ReturnOpPattern. i8 case showed the replaced op %15 = spv.ShiftRightArithmetic %14, %13 : i32, i32, and i1 case showed the old one %16 = memref.load <<UNKNOWN SSA VALUE>>[] : memref<i1>. It is weird because it said that the pattern applied successfully, but the users still get the old operand.

I also verified that i1 case and i8 case generate the same log with -debug and they are all marked legal and success.

I'd be appreciate if someone can help on this... It looks like a memref.load was created because the variable name is %16.

To repro, run mlir-opt -convert-std-to-spirv a.mlir

full log for i1: https://gist.github.com/hanhanW/8a4494f337629299959b6e7d59a882f5
full log for i8: https://gist.github.com/hanhanW/56a8fcdd6a8eb25a743a71667754e6e4

[hanchung]

@antiagainst is going to add some support to SPIR-V, and I will rebase on it and rework a bit.

[antiagainst]

I've created D100058 and D100059 to introduce controls over emulation. You can rebase after they are landed. With D100059 I think we can introduce another option there to control the # of bits for boolean types. So the getTypeNumBytes (which should probably be changed to getTypeNumBits) and all those convert*Type functions can properly handle i1.

[hanchung]

fix i1 issues

[hanchung]

Fixed, and enhanced the test that loads i1. (now it will return i1 to make sure there is no such issue again.)

[antiagainst]

Mostly okay to me. Major request to have an option for control.

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.cpp

995

s/isInt1/isBool/

1052

W already have typeConverter at L988.

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

349

s/convertMemrefI1Type/convertBoolMemrefType/

354

s/yte/yet/

382

Can you add an option in SPIRVTypeConverter::Options, like boolNumBits, to control this? (For now we can just support the case == 8 and bail out on type conversion for other cases.)

[hanchung]

address comments

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.cpp

1052

I did not notice it, thanks!

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

382

good point, it is useful

[hanchung]

Also added the implementation of SPIRVTypeConverter::getOptions().

[antiagainst]

Cool, just one final comment.

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

156

Oops, sorry I forgot about this. :) Thanks for fixing it!

370

I think here it should be

c++
if (numBoolBits != 8) {
  LLVM_DEBUG(llvm::dbgs() << "using non-8-bit storage for bool types unimplemented");
  return nullptr;
}

?

[hanchung]

address comments

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

370

good point! I feel that it will work, but did not test it. Let's have the check, and we can turn it back and add tests when we need it.

mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h

struct Options {
     /// runtime to also feed in data with a matched bitwidth and layout for
     /// interface types. The runtime can do that by inspecting the SPIR-V
     /// module.
     ///
     /// If the original scalar type has less than 32-bit, a multiple of its
     /// values will be packed into one 32-bit value to be memory efficient.
     bool emulateNon32BitScalarTypes;
 
+    /// The number of bits to store a boolean value. It is eight bits by
+    /// default.
+    unsigned boolNumBits;
+
     // Note: we need this instead of inline initializers becuase of
     // https://bugs.llvm.org/show_bug.cgi?id=36684
-    Options() : emulateNon32BitScalarTypes(true) {}
+    Options() : emulateNon32BitScalarTypes(true), boolNumBits(8) {}
   };
 
   explicit SPIRVTypeConverter(spirv::TargetEnvAttr targetAttr,
                               Options options = {});
 
   /// Gets the SPIR-V correspondence for the standard index type.
   static Type getIndexType(MLIRContext *context);
 

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.cpp

if (!memrefType.getElementType().isSignlessInteger())
     return failure();
 
   auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
   spirv::AccessChainOp accessChainOp =
       spirv::getElementPtr(typeConverter, memrefType, loadOperands.memref(),
                            loadOperands.indices(), loc, rewriter);
 
   int srcBits = memrefType.getElementType().getIntOrFloatBitWidth();
+  bool isBool = srcBits == 1;
+  if (isBool)
+    srcBits = typeConverter.getOptions().boolNumBits;
   auto dstType = typeConverter.convertType(memrefType)
                      .cast<spirv::PointerType>()
                      .getPointeeType()
                      .cast<spirv::StructType>()
                      .getElementType(0)
                      .cast<spirv::ArrayType>()
                      .getElementType();
   int dstBits = dstType.getIntOrFloatBitWidth();
IntLoadOpPattern::matchAndRewrite(memref::LoadOp loadOp,
   IntegerAttr shiftValueAttr =
       rewriter.getIntegerAttr(dstType, dstBits - srcBits);
   Value shiftValue =
       rewriter.create<spirv::ConstantOp>(loc, dstType, shiftValueAttr);
   result = rewriter.create<spirv::ShiftLeftLogicalOp>(loc, dstType, result,
                                                       shiftValue);
   result = rewriter.create<spirv::ShiftRightArithmeticOp>(loc, dstType, result,
                                                           shiftValue);
+
+  if (isBool) {
+    dstType = typeConverter.convertType(loadOp.getType());
+    mask = spirv::ConstantOp::getOne(result.getType(), loc, rewriter);
+    Value isOne = rewriter.create<spirv::IEqualOp>(loc, result, mask);
+    Value zero = spirv::ConstantOp::getZero(dstType, loc, rewriter);
+    Value one = spirv::ConstantOp::getOne(dstType, loc, rewriter);
+    result = rewriter.create<spirv::SelectOp>(loc, dstType, isOne, one, zero);
+  } else if (result.getType().getIntOrFloatBitWidth() !=
+             static_cast<unsigned>(dstBits)) {
+    result = rewriter.create<spirv::SConvertOp>(loc, dstType, result);
+  }
   rewriter.replaceOp(loadOp, result);
 
   assert(accessChainOp.use_empty());
   rewriter.eraseOp(accessChainOp);
 
   return success();
 }
 
if (!memrefType.getElementType().isSignlessInteger())
     return failure();
 
   auto loc = storeOp.getLoc();
   auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
   spirv::AccessChainOp accessChainOp =
       spirv::getElementPtr(typeConverter, memrefType, storeOperands.memref(),
                            storeOperands.indices(), loc, rewriter);
   int srcBits = memrefType.getElementType().getIntOrFloatBitWidth();
+
+  bool isBool = srcBits == 1;
+  if (isBool)
+    srcBits = typeConverter.getOptions().boolNumBits;
   auto dstType = typeConverter.convertType(memrefType)
                      .cast<spirv::PointerType>()
                      .getPointeeType()
                      .cast<spirv::StructType>()
                      .getElementType(0)
                      .cast<spirv::ArrayType>()
                      .getElementType();
   int dstBits = dstType.getIntOrFloatBitWidth();
IntStoreOpPattern::matchAndRewrite(memref::StoreOp storeOp,
   // Create a mask to clear the destination. E.g., if it is the second i8 in
   // i32, 0xFFFF00FF is created.
   Value mask = rewriter.create<spirv::ConstantOp>(
       loc, dstType, rewriter.getIntegerAttr(dstType, (1 << srcBits) - 1));
   Value clearBitsMask =
       rewriter.create<spirv::ShiftLeftLogicalOp>(loc, dstType, mask, offset);
   clearBitsMask = rewriter.create<spirv::NotOp>(loc, dstType, clearBitsMask);
 
-  Value storeVal =
-      shiftValue(loc, storeOperands.value(), offset, mask, dstBits, rewriter);
+  Value storeVal = storeOperands.value();
+  if (isBool) {
+    Value zero = spirv::ConstantOp::getZero(dstType, loc, rewriter);
+    Value one = spirv::ConstantOp::getOne(dstType, loc, rewriter);
+    storeVal =
+        rewriter.create<spirv::SelectOp>(loc, dstType, storeVal, one, zero);
+  }
+  storeVal = shiftValue(loc, storeVal, offset, mask, dstBits, rewriter);
   Value adjustedPtr = adjustAccessChainForBitwidth(typeConverter, accessChainOp,
                                                    srcBits, dstBits, rewriter);
   Optional<spirv::Scope> scope = getAtomicOpScope(memrefType);
   if (!scope)
     return failure();
   Value result = rewriter.create<spirv::AtomicAndOp>(
       loc, dstType, adjustedPtr, *scope, spirv::MemorySemantics::AcquireRelease,
       clearBitsMask);

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

#define STORAGE_SPACE_MAP_FN(storage, space)                                   \
   switch (space) {
     STORAGE_SPACE_MAP_LIST(STORAGE_SPACE_MAP_FN)
   default:
     return llvm::None;
   }
 #undef STORAGE_SPACE_MAP_FN
 }
 
+const SPIRVTypeConverter::Options &SPIRVTypeConverter::getOptions() const {
+  return options;
+}
+
 #undef STORAGE_SPACE_MAP_LIST
 
 // TODO: This is a utility function that should probably be exposed by the
 // SPIR-V dialect. Keeping it local till the use case arises.
 static Optional<int64_t>
 getTypeNumBytes(const SPIRVTypeConverter::Options &options, Type type) {
   if (type.isa<spirv::ScalarType>()) {
     auto bitWidth = type.getIntOrFloatBitWidth();
if (!arrayElemSize) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot deduce converted element size\n");
     return nullptr;
   }
 
   return spirv::ArrayType::get(arrayElemType, arrayElemCount, *arrayElemSize);
 }
 
+static Type convertBoolMemrefType(const spirv::TargetEnv &targetEnv,
+                                  const SPIRVTypeConverter::Options &options,
+                                  MemRefType type) {
+  if (!type.hasStaticShape()) {
+    LLVM_DEBUG(llvm::dbgs()
+               << type << " dynamic shape on i1 is not supported yet\n");
+    return nullptr;
+  }
+
+  Optional<spirv::StorageClass> storageClass =
+      SPIRVTypeConverter::getStorageClassForMemorySpace(
+          type.getMemorySpaceAsInt());
+  if (!storageClass) {
+    LLVM_DEBUG(llvm::dbgs()
+               << type << " illegal: cannot convert memory space\n");
+    return nullptr;
+  }
+
+  unsigned numBoolBits = options.boolNumBits;
+  if (numBoolBits != 8) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "using non-8-bit storage for bool types unimplemented");
+    return nullptr;
+  }
+  auto elementType = IntegerType::get(type.getContext(), numBoolBits)
+                         .dyn_cast<spirv::ScalarType>();
+  if (!elementType)
+    return nullptr;
+  Type arrayElemType =
+      convertScalarType(targetEnv, options, elementType, storageClass);
+  if (!arrayElemType)
+    return nullptr;
+  Optional<int64_t> arrayElemSize = getTypeNumBytes(options, arrayElemType);
+  if (!arrayElemSize) {
+    LLVM_DEBUG(llvm::dbgs()
+               << type << " illegal: cannot deduce converted element size\n");
+    return nullptr;
+  }
+
+  int64_t memrefSize = (type.getNumElements() * numBoolBits + 7) / 8;
+  auto arrayElemCount = (memrefSize + *arrayElemSize - 1) / *arrayElemSize;
+  auto arrayType =
+      spirv::ArrayType::get(arrayElemType, arrayElemCount, *arrayElemSize);
+
+  // Wrap in a struct to satisfy Vulkan interface requirements. Memrefs with
+  // workgroup storage class do not need the struct to be laid out explicitly.
+  auto structType = *storageClass == spirv::StorageClass::Workgroup
+                        ? spirv::StructType::get(arrayType)
+                        : spirv::StructType::get(arrayType, 0);
+  return spirv::PointerType::get(structType, *storageClass);
+}
+
 static Type convertMemrefType(const spirv::TargetEnv &targetEnv,
                               const SPIRVTypeConverter::Options &options,
                               MemRefType type) {
+  if (type.getElementType().isa<IntegerType>() &&
+      type.getElementTypeBitWidth() == 1) {
+    return convertBoolMemrefType(targetEnv, options, type);
+  }
+
   Optional<spirv::StorageClass> storageClass =
       SPIRVTypeConverter::getStorageClassForMemorySpace(
           type.getMemorySpaceAsInt());
   if (!storageClass) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot convert memory space\n");
     return nullptr;
   }

mlir/test/Conversion/StandardToSPIRV/std-ops-to-spirv.mlir

 }
 
 } // end module
 
 // -----
 
 // Check that access chain indices are properly adjusted if non-32-bit types are
 // emulated via 32-bit types.
-// TODO: Test i1 and i64 types.
+// TODO: Test i64 types.
 module attributes {
   spv.target_env = #spv.target_env<
     #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, {}>
 } {
 
+// CHECK-LABEL: @load_i1
+func @load_i1(%arg0: memref<i1>) -> i1 {
+  //     CHECK: %[[ZERO:.+]] = spv.Constant 0 : i32
+  //     CHECK: %[[FOUR1:.+]] = spv.Constant 4 : i32
+  //     CHECK: %[[QUOTIENT:.+]] = spv.SDiv %[[ZERO]], %[[FOUR1]] : i32
+  //     CHECK: %[[PTR:.+]] = spv.AccessChain %{{.+}}[%[[ZERO]], %[[QUOTIENT]]]
+  //     CHECK: %[[LOAD:.+]] = spv.Load  "StorageBuffer" %[[PTR]]
+  //     CHECK: %[[FOUR2:.+]] = spv.Constant 4 : i32
+  //     CHECK: %[[EIGHT:.+]] = spv.Constant 8 : i32
+  //     CHECK: %[[IDX:.+]] = spv.UMod %[[ZERO]], %[[FOUR2]] : i32
+  //     CHECK: %[[BITS:.+]] = spv.IMul %[[IDX]], %[[EIGHT]] : i32
+  //     CHECK: %[[VALUE:.+]] = spv.ShiftRightArithmetic %[[LOAD]], %[[BITS]] : i32, i32
+  //     CHECK: %[[MASK:.+]] = spv.Constant 255 : i32
+  //     CHECK: %[[T1:.+]] = spv.BitwiseAnd %[[VALUE]], %[[MASK]] : i32
+  //     CHECK: %[[T2:.+]] = spv.Constant 24 : i32
+  //     CHECK: %[[T3:.+]] = spv.ShiftLeftLogical %[[T1]], %[[T2]] : i32, i32
+  //     CHECK: %[[T4:.+]] = spv.ShiftRightArithmetic %[[T3]], %[[T2]] : i32, i32
+  // Convert to i1 type.
+  //     CHECK: %[[ONE:.+]] = spv.Constant 1 : i32
+  //     CHECK: %[[ISONE:.+]]  = spv.IEqual %[[T4]], %[[ONE]] : i32
+  //     CHECK: %[[FALSE:.+]] = spv.Constant false
+  //     CHECK: %[[TRUE:.+]] = spv.Constant true
+  //     CHECK: %[[RES:.+]] = spv.Select %[[ISONE]], %[[TRUE]], %[[FALSE]] : i1, i1
+  //     CHECK: spv.ReturnValue %[[RES]] : i1
+  %0 = memref.load %arg0[] : memref<i1>
+  return %0 : i1
+}
+
 // CHECK-LABEL: @load_i8
 func @load_i8(%arg0: memref<i8>) {
   //     CHECK: %[[ZERO:.+]] = spv.Constant 0 : i32
   //     CHECK: %[[FOUR1:.+]] = spv.Constant 4 : i32
   //     CHECK: %[[QUOTIENT:.+]] = spv.SDiv %[[ZERO]], %[[FOUR1]] : i32
   //     CHECK: %[[PTR:.+]] = spv.AccessChain %{{.+}}[%[[ZERO]], %[[QUOTIENT]]]
   //     CHECK: %[[LOAD:.+]] = spv.Load  "StorageBuffer" %[[PTR]]
   //     CHECK: %[[FOUR2:.+]] = spv.Constant 4 : i32

 func @load_f32(%arg0: memref<f32>) {
   // CHECK-NOT: spv.SDiv
   //     CHECK: spv.Load
   // CHECK-NOT: spv.ShiftRightArithmetic
   %0 = memref.load %arg0[] : memref<f32>
   return
 }
 
+// CHECK-LABEL: @store_i1
+//       CHECK: (%[[ARG0:.+]]: {{.*}}, %[[ARG1:.+]]: i1)
+func @store_i1(%arg0: memref<i1>, %value: i1) {
+  //     CHECK: %[[ZERO:.+]] = spv.Constant 0 : i32
+  //     CHECK: %[[FOUR:.+]] = spv.Constant 4 : i32
+  //     CHECK: %[[EIGHT:.+]] = spv.Constant 8 : i32
+  //     CHECK: %[[IDX:.+]] = spv.UMod %[[ZERO]], %[[FOUR]] : i32
+  //     CHECK: %[[OFFSET:.+]] = spv.IMul %[[IDX]], %[[EIGHT]] : i32
+  //     CHECK: %[[MASK1:.+]] = spv.Constant 255 : i32
+  //     CHECK: %[[TMP1:.+]] = spv.ShiftLeftLogical %[[MASK1]], %[[OFFSET]] : i32, i32
+  //     CHECK: %[[MASK:.+]] = spv.Not %[[TMP1]] : i32
+  //     CHECK: %[[ZERO1:.+]] = spv.Constant 0 : i32
+  //     CHECK: %[[ONE1:.+]] = spv.Constant 1 : i32
+  //     CHECK: %[[CASTED_ARG1:.+]] = spv.Select %[[ARG1]], %[[ONE1]], %[[ZERO1]] : i1, i32
+  //     CHECK: %[[CLAMPED_VAL:.+]] = spv.BitwiseAnd %[[CASTED_ARG1]], %[[MASK1]] : i32
+  //     CHECK: %[[STORE_VAL:.+]] = spv.ShiftLeftLogical %[[CLAMPED_VAL]], %[[OFFSET]] : i32, i32
+  //     CHECK: %[[FOUR2:.+]] = spv.Constant 4 : i32
+  //     CHECK: %[[ACCESS_IDX:.+]] = spv.SDiv %[[ZERO]], %[[FOUR2]] : i32
+  //     CHECK: %[[PTR:.+]] = spv.AccessChain %[[ARG0]][%[[ZERO]], %[[ACCESS_IDX]]]
+  //     CHECK: spv.AtomicAnd "Device" "AcquireRelease" %[[PTR]], %[[MASK]]
+  //     CHECK: spv.AtomicOr "Device" "AcquireRelease" %[[PTR]], %[[STORE_VAL]]
+  memref.store %value, %arg0[] : memref<i1>
+  return
+}
+
 // CHECK-LABEL: @store_i8
 //       CHECK: (%[[ARG0:.+]]: {{.*}}, %[[ARG1:.+]]: i32)
 func @store_i8(%arg0: memref<i8>, %value: i8) {
   //     CHECK: %[[ZERO:.+]] = spv.Constant 0 : i32
   //     CHECK: %[[FOUR:.+]] = spv.Constant 4 : i32
   //     CHECK: %[[EIGHT:.+]] = spv.Constant 8 : i32
   //     CHECK: %[[IDX:.+]] = spv.UMod %[[ZERO]], %[[FOUR]] : i32
   //     CHECK: %[[OFFSET:.+]] = spv.IMul %[[IDX]], %[[EIGHT]] : i32

mlir/test/Conversion/StandardToSPIRV/std-types-to-spirv.mlir

func @memref_mem_space(
     %arg4: memref<4xf32, 5>,
     %arg5: memref<4xf32, 6>
 ) { return }
 
 } // end module
 
 // -----
 
-// Check that boolean memref is not supported at the moment.
-module attributes {
-  spv.target_env = #spv.target_env<#spv.vce<v1.0, [], []>, {}>
-} {
-
-// CHECK-LABEL: func @memref_type({{%.*}}: memref<3xi1>)
-func @memref_type(%arg0: memref<3xi1>) {
-  return
-}
-
-} // end module
-
-// -----
-
 // Check that using non-32-bit scalar types in interface storage classes
 // requires special capability and extension: convert them to 32-bit if not
 // satisfied.
 module attributes {
   spv.target_env = #spv.target_env<#spv.vce<v1.0, [], []>, {}>
 } {
 
+// An i1 is store in 8-bit, so 5xi1 has 40 bits, which is stored in 2xi32.
+// CHECK-LABEL: spv.func @memref_1bit_type
+// CHECK-SAME: !spv.ptr<!spv.struct<(!spv.array<2 x i32, stride=4> [0])>, StorageBuffer>
+func @memref_1bit_type(%arg0: memref<5xi1>) { return }
+
 // CHECK-LABEL: spv.func @memref_8bit_StorageBuffer
 // CHECK-SAME: !spv.ptr<!spv.struct<(!spv.array<16 x i32, stride=4> [0])>, StorageBuffer>
 // NOEMU-LABEL: func @memref_8bit_StorageBuffer
 // NOEMU-SAME: memref<16xi8>
 func @memref_8bit_StorageBuffer(%arg0: memref<16xi8, 0>) { return }
 
 // CHECK-LABEL: spv.func @memref_8bit_Uniform
 // CHECK-SAME: !spv.ptr<!spv.struct<(!spv.array<16 x si32, stride=4> [0])>, Uniform>