[Mlir-commits] [mlir] [mlir][spirv] Mark in-bounds linearized indices no-wrap (PR #215834)

[Igor Wodiany]

================
@@ -1283,6 +1285,48 @@ struct ReturnOpVectorUnroll final : OpRewritePattern<func::ReturnOp> {
   }
 };
 
+static void addNoWrapDecorations(Operation *op,
+                                 spirv::LinearizedIndexNoWrapFlags flags,
+                                 OpBuilder &builder) {
+  if (flags.noSignedWrap)
+    op->setAttr(spirv::getDecorationString(spirv::Decoration::NoSignedWrap),
+                builder.getUnitAttr());
+  if (flags.noUnsignedWrap)
+    op->setAttr(spirv::getDecorationString(spirv::Decoration::NoUnsignedWrap),
+                builder.getUnitAttr());
+}
+
+static std::optional<uint64_t> getMaxLinearizedIndex(ArrayRef<int64_t> shape,
+                                                     ArrayRef<int64_t> strides,
+                                                     int64_t offset) {
+  if (shape.size() != strides.size() || offset < 0)
+    return std::nullopt;
+
+  uint64_t maxLinearIndex = offset;
+  for (auto [dimension, stride] : llvm::zip(shape, strides)) {
+    if (dimension <= 0 || stride < 0)
+      return std::nullopt;
+    std::optional<uint64_t> nextMaxLinearIndex = llvm::checkedMulAddUnsigned(
+        static_cast<uint64_t>(dimension - 1), static_cast<uint64_t>(stride),
+        maxLinearIndex);
+    if (!nextMaxLinearIndex)
+      return std::nullopt;
+    maxLinearIndex = *nextMaxLinearIndex;
+  }
+  return maxLinearIndex;
+}
+
+static spirv::LinearizedIndexNoWrapFlags
+shouldEmitNoWrapDecorations(const SPIRVTypeConverter &typeConverter,
----------------
IgWod wrote:

Sorry, one more nit. I'm not sure about this name. `should` implies to me it's `true`/`false` function, but it returns flags. We could rename it, but maybe it would make more sense to just put the `allows` logic in `getLinearizedIndexNoWrapFlags` and call that directly?

https://github.com/llvm/llvm-project/pull/215834