[Init] Add identity initializer (#1079)

This is a step towards allowing tests like A * I = A.

As is, it still doesn't do what we want, since like the others, it tried
to initialize all tensors with identity, so any non square tensor will
fail.

The next step is to add options like "--identity-b" together with
"--splat-to-random", so that only the B matrix is replaced by identity.

Also remove some unused initializers.

Author: rengolin

include/TPP/Passes.td

@@ -577,6 +577,9 @@ def TppRunnerWrapper : Pass<"tpp-runner-wrapper", "ModuleOp">{
     Option<"initType", "init-type", "std::string",
             /*default=*/"",
            "Initializer type (const, simple, cont, rand, normal).">,
+    Option<"identity", "identity", "int64_t",
+            /*default=*/"-1",
+           "Identity square argument (-1=none, 0=a, 1=b, ...).">,
   ];
 }
 

include/TPP/Runner/MLIRBench.h

@@ -40,13 +40,14 @@ class FuncOp;
 // pipeline.
 struct MLIRBenchConfig {
   MLIRBenchConfig() = default;
-  MLIRBenchConfig(int seed, TensorInitType initType, std::string backend,
+  MLIRBenchConfig(int seed, TensorInitType initType, int identity, std::string backend,
                   bool offloadToDevice)
-      : seed(seed), initType(initType), backend(backend),
+      : seed(seed), initType(initType), identity(identity), backend(backend),
         offloadToDevice(offloadToDevice) {}
 
   int seed = 0;
   TensorInitType initType = TensorInitType::Auto;
+  int identity = -1;
   std::string backend = "cpu";
   bool offloadToDevice = true;
 };
@@ -85,6 +86,9 @@ class MLIRBench {
   /// Seed for the random tensor filling
   int seed;
 
+  /// Which argument is the identity, if any
+  int identity;
+
   /// Tensor init type
   TensorInitType initType;
 

include/TPP/Transforms/Utils/TensorInit.h

@@ -28,7 +28,8 @@ struct ITensorInit {
   // Returns a dense attribute with a specified shape, initialized
   // with a particular implementation (see derived classes) with
   // a reasonable distribution.
-  virtual mlir::DenseElementsAttr get(mlir::ShapedType shape) = 0;
+  virtual llvm::FailureOr<mlir::DenseElementsAttr>
+  get(mlir::ShapedType shape) = 0;
 };
 
 // Base class.
@@ -39,12 +40,15 @@ template <typename T> struct TensorInit : public ITensorInit {
   // Returns a dense attribute with a specified shape, initialized
   // with a particular implementation (see derived classes) with
   // a reasonable distribution.
-  virtual mlir::DenseElementsAttr get(mlir::ShapedType shape) override {
+  virtual llvm::FailureOr<mlir::DenseElementsAttr>
+  get(mlir::ShapedType shape) override {
+    if (!checkShape(shape))
+      return llvm::failure();
+
+    // Populate the shape
     buffer.clear();
-    size = 1;
-    for (size_t dim = 0, rank = shape.getRank(); dim < rank; dim++)
-      size *= shape.getDimSize(dim);
     fillData();
+
     // For some reason, memref global op needs dense tensor type
     // See: lib/Dialect/MemRef/IR/MemRefOps.cpp :: GlobalOp::verify
     auto tensorType =
@@ -53,11 +57,25 @@ template <typename T> struct TensorInit : public ITensorInit {
   }
 
 protected:
+  // Shape dims
+  std::vector<size_t> dims;
   // Number of elements in the shape
   size_t size;
   // Data pointer
   std::vector<T> buffer;
 
+  // Check the shape and fill the internal structure
+  virtual bool checkShape(mlir::ShapedType shape) {
+    size = 1;
+    dims.clear();
+    for (size_t i = 0, rank = shape.getRank(); i < rank; i++) {
+      auto dim = shape.getDimSize(i);
+      dims.push_back(dim);
+      size *= dim;
+    }
+    return true;
+  }
+
   // Insert element indexed on the buffer
   virtual void insert(size_t index, T value) {
     buffer[index] = value;
@@ -82,10 +100,9 @@ template <typename T> struct TensorInit : public ITensorInit {
 enum class TensorInitType {
   Auto,
   Constant,
-  Simple,
-  Continuous,
   Random,
   Normal,
+  Identity,
   Invalid
 };
 

include/TPP/Transforms/Utils/TensorInitFloat.h

@@ -85,30 +85,14 @@ struct TensorInitFloat : public TensorInit<llvm::APFloat> {
   virtual void fillData() override = 0;
 };
 
-// Constant init (all-ones, do not use!).
+// Constant init (all-ones).
 struct ConstantTensorInitFloat : TensorInitFloat {
   ConstantTensorInitFloat(DataType type) : TensorInitFloat(type) {}
 
   // Return a dense<1.0> repeated throughout the shape.
-  mlir::DenseElementsAttr get(mlir::ShapedType shape) override;
+  mlir::FailureOr<mlir::DenseElementsAttr> get(mlir::ShapedType shape) override;
 
-  void fillData() override;
-};
-
-// Simple init (basic example, not useful).
-struct SimpleTensorInitFloat : TensorInitFloat {
-  SimpleTensorInitFloat(DataType type) : TensorInitFloat(type) {}
-
-  // Return a dense<0.3, 0.6, 0.9> repeated throughout the shape.
-  void fillData() override;
-};
-
-// Continuous init (normalized affine range).
-struct ContinuousTensorInitFloat : TensorInitFloat {
-  ContinuousTensorInitFloat(DataType type) : TensorInitFloat(type) {}
-
-  // Return a dense<0.0 ... 1.0> throughout the shape.
-  void fillData() override;
+  void fillData() override { assert(false && "Should not be called"); }
 };
 
 // Random init (uniform).
@@ -151,4 +135,24 @@ struct NormalTensorInitFloat : TensorInitFloat {
   std::normal_distribution<float> distribution;
 };
 
+// Identity init.
+struct IdentityTensorInitFloat : TensorInitFloat {
+  IdentityTensorInitFloat(DataType type)
+      : TensorInitFloat(type) {}
+
+  // Makes sure the shape is "square"
+  bool checkShape(mlir::ShapedType shape) override {
+    if (!TensorInit::checkShape(shape))
+      return false;
+    // Now the fields are set, compare all dims to be equal, 2D only for now
+    return dims.size() == 2 && dims[0] == dims[1];
+  }
+
+  // Should not be called.
+  float next() { assert(false && "Should not be called"); }
+
+  // Return a diagonal of <1.0>s throughout the shape.
+  void fillData() override;
+};
+
 #endif // TPP_TRANSFORMS_UTILS_TENSORINITFLOAT_H

include/TPP/Transforms/Utils/TensorInitInt.h

@@ -74,37 +74,14 @@ struct TensorInitInt : public TensorInit<llvm::APInt> {
   virtual void fillData() override = 0;
 };
 
-// Constant init (all-ones, do not use!).
+// Constant init (all-ones).
 struct ConstantTensorInitInt : TensorInitInt {
   ConstantTensorInitInt(DataType type) : TensorInitInt(type) {}
 
   // Return a dense<1> repeated throughout the shape.
-  mlir::DenseElementsAttr get(mlir::ShapedType shape) override;
+  mlir::FailureOr<mlir::DenseElementsAttr> get(mlir::ShapedType shape) override;
 
-  void fillData() override;
-};
-
-// Simple init (basic example, not useful).
-struct SimpleTensorInitInt : TensorInitInt {
-  SimpleTensorInitInt(DataType type) : TensorInitInt(type) {}
-
-  // Return a dense<0, 1, 2> repeated throughout the shape.
-  void fillData() override;
-};
-
-// Continuous init (quantized normalized affine range).
-struct ContinuousTensorInitInt : TensorInitInt {
-  ContinuousTensorInitInt(DataType type)
-      : TensorInitInt(type), upperBound(255) {
-    if (type == DataType::I8)
-      upperBound = 127;
-  }
-
-  // Return a dense<0 ... upperBound> throughout the shape.
-  void fillData() override;
-
-  // Upper bound for quantization.
-  int upperBound;
+  void fillData() override { assert(false && "Should not be called"); }
 };
 
 // Random init (uniform).
@@ -152,4 +129,24 @@ struct NormalTensorInitInt : TensorInitInt {
   std::binomial_distribution<uint64_t> distribution;
 };
 
+// Identity init.
+struct IdentityTensorInitInt : TensorInitInt {
+  IdentityTensorInitInt(DataType type)
+      : TensorInitInt(type) {}
+
+  // Makes sure the shape is "square"
+  bool checkShape(mlir::ShapedType shape) override {
+    if (!TensorInit::checkShape(shape))
+      return false;
+    // Now the fields are set, compare all dims to be equal, 2D only for now
+    return dims.size() == 2 && dims[0] == dims[1];
+  }
+
+  // Should not be called.
+  float next() { assert(false && "Should not be called"); }
+
+  // Return a diagonal of <1.0>s throughout the shape.
+  void fillData() override;
+};
+
 #endif // TPP_TRANSFORMS_UTILS_TENSORINITINT_H

lib/TPP/Runner/MLIRBench.cpp

@@ -56,6 +56,7 @@ using namespace mlir;
 MLIRBench::MLIRBench(mlir::Operation *op, const MLIRBenchConfig &config)
     : builder(op->getContext()), unkLoc(builder.getUnknownLoc()) {
   seed = config.seed;
+  identity = config.identity;
   backend = config.backend;
   initType = config.initType;
   offloadToDevice = config.offloadToDevice;
@@ -113,7 +114,7 @@ LogicalResult MLIRBench::replaceSplatWithRandom() {
     return module.emitError("No seed for random init");
 
   // Only replace attribute if it's a dense splat
-  auto replaceSplat = [&](ShapedType shape, Attribute attr) -> Attribute {
+  auto replaceSplat = [&](ShapedType shape, Attribute attr) -> FailureOr<Attribute> {
     // We only change dense attributes that are splat
     auto value = dyn_cast<DenseElementsAttr>(attr);
     if (!value || !value.isSplat())
@@ -145,7 +146,9 @@ LogicalResult MLIRBench::replaceSplatWithRandom() {
     if (!global)
       continue;
     auto newAttr = replaceSplat(global.getType(), global.getInitialValueAttr());
-    global.setInitialValueAttr(newAttr);
+    if (failed(newAttr))
+      return failure();
+    global.setInitialValueAttr(newAttr.value());
   }
 
   // Tensors are arith.constant values
@@ -157,7 +160,9 @@ LogicalResult MLIRBench::replaceSplatWithRandom() {
     if (!cstType)
       continue;
     auto newAttr = replaceSplat(cstType, constant.getValueAttr());
-    constant.setValueAttr(cast<TypedAttr>(newAttr));
+    if (failed(newAttr))
+      return failure();
+    constant.setValueAttr(cast<TypedAttr>(newAttr.value()));
   }
 
   return success();
@@ -212,34 +217,48 @@ LogicalResult MLIRBench::createKernelArgs() {
   auto &mainBody = getMainBlock();
   builder.setInsertionPointToStart(&mainBody);
 
+  int argNum = 0;
   for (auto &ty : kernel.getArgumentTypes()) {
-    auto arg = TypeSwitch<Type, std::optional<Value>>(ty)
-                   .Case<MemRefType>([&](auto memRefTy) {
-                     // Create a memref global
-                     Value data = createDenseMemref(builder, module, initType,
-                                                    memRefTy, seed);
-                     data = registerOnGpu(data, memRefTy);
-                     return data;
-                   })
-                   .Case<TensorType>([&](auto tensorTy) {
-                     // Create a memref global and cast it to a tensor
-                     // to ensure that the buffer is writable and
-                     // bufferization does not insert extra
-                     // allocations + copies
-                     auto memrefType = MemRefType::get(
-                         tensorTy.getShape(), tensorTy.getElementType());
-                     auto data = createDenseMemref(builder, module, initType,
-                                                   memrefType, seed);
-                     data = registerOnGpu(data, memrefType);
-                     return builder.create<bufferization::ToTensorOp>(
-                         unkLoc, tensorTy, data, /*restrict=*/true, /*writable=*/true);
-                   })
-                   .Default([&](auto t) { return std::nullopt; });
+    auto argInitType = initType;
+    // Requested an argument to be identity, must be 2D square
+    if (argNum == identity) {
+      ShapedType shape = dyn_cast<ShapedType>(ty);
+      if (shape && shape.getRank() == 2 &&
+          shape.getDimSize(0) == shape.getDimSize(1)) {
+        argInitType = TensorInitType::Identity;
+      } else {
+        return module.emitError("Invalid shape for identity init");
+      }
+    }
+    auto arg =
+        TypeSwitch<Type, std::optional<Value>>(ty)
+            .Case<MemRefType>([&](auto memRefTy) {
+              // Create a memref global
+              Value data = createDenseMemref(builder, module, argInitType,
+                                             memRefTy, seed);
+              data = registerOnGpu(data, memRefTy);
+              return data;
+            })
+            .Case<TensorType>([&](auto tensorTy) {
+              // Create a memref global and cast it to a tensor
+              // to ensure that the buffer is writable and
+              // bufferization does not insert extra
+              // allocations + copies
+              auto memrefType = MemRefType::get(tensorTy.getShape(),
+                                                tensorTy.getElementType());
+              auto data = createDenseMemref(builder, module, argInitType,
+                                            memrefType, seed);
+              data = registerOnGpu(data, memrefType);
+              return builder.create<bufferization::ToTensorOp>(
+                  unkLoc, tensorTy, data, /*restrict=*/true, /*writable=*/true);
+            })
+            .Default([&](auto t) { return std::nullopt; });
 
     if (!arg)
-      return failure();
+      return module.emitError("Cannot create kernel argument");
 
     kernelArgs.push_back(*arg);
+    argNum++;
   }
 
   return success();

lib/TPP/Runner/TppRunnerWrapper.cpp

@@ -56,7 +56,7 @@ struct TppRunnerWrapper
     }
 
     // Benchmark object.
-    MLIRBenchConfig config(seed, tensorInitType, backend, offloadToDevice);
+    MLIRBenchConfig config(seed, tensorInitType, identity, backend, offloadToDevice);
     MLIRBench bench(module, config);
 
     // Can only either print or run benchmarks, make this clear before we try to

lib/TPP/Transforms/Utils/BuilderUtils.cpp

@@ -79,7 +79,8 @@ Value createDenseTensor(OpBuilder &builder, TensorInitType initType,
   auto unkLoc = builder.getUnknownLoc();
   auto init = getTensorInit(initType, type.getElementType(), seed);
   auto floatInit = init->get(type);
-  return builder.create<arith::ConstantOp>(unkLoc, type, floatInit);
+  assert(!failed(floatInit) && "Invalid dense tensor initializer");
+  return builder.create<arith::ConstantOp>(unkLoc, type, floatInit.value());
 }
 
 Value createDenseMemref(OpBuilder &builder, ModuleOp module,
@@ -103,10 +104,11 @@ Value createDenseMemref(OpBuilder &builder, ModuleOp module,
     auto alignment = builder.getIntegerAttr(builder.getI64Type(), 128);
     auto init = getTensorInit(initType, type.getElementType(), seed);
     auto floatInit = init->get(type);
+    assert(!failed(floatInit) && "Invalid dense tensor initializer");
 
     // Create the global object in the Module's region
     auto global = builder.create<memref::GlobalOp>(
-        unkLoc, StringRef(name), privAttr, type, floatInit,
+        unkLoc, StringRef(name), privAttr, type, floatInit.value(),
         /*constant=*/false, alignment);
     globalName = global.getName();
   }