[XLA:CPU][XTile] Create lowering for Iota.

PiperOrigin-RevId: 825568861

Author: WillFroom

xla/backends/cpu/codegen/tiled/transforms/passes.td

@@ -46,11 +46,12 @@ def ShloToVectorPass : Pass<"xtile-cpu-shlo-to-vector", "mlir::ModuleOp"> {
   let constructor = "CreateShloToVectorPass()";
 
   let dependentDialects = [
+    "mlir::arith::ArithDialect",
+    "mlir::memref::MemRefDialect",
+    "mlir::scf::SCFDialect",
+    "mlir::stablehlo::StablehloDialect",
     "mlir::tensor::TensorDialect",
     "mlir::vector::VectorDialect",
-    "mlir::stablehlo::StablehloDialect",
-    "mlir::scf::SCFDialect",
-    "mlir::memref::MemRefDialect",
   ];
 }
 

xla/backends/cpu/codegen/tiled/transforms/shlo_to_vector.cc

@@ -317,6 +317,35 @@ struct LowerReshape : mlir::OpRewritePattern<mlir::stablehlo::ReshapeOp> {
   }
 };
 
+struct LowerIota : mlir::OpRewritePattern<mlir::stablehlo::IotaOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::stablehlo::IotaOp op,
+      mlir::PatternRewriter& rewriter) const override {
+    if (op.getType().getRank() != 1) {
+      return rewriter.notifyMatchFailure(
+          op, "iota op with rank != 1 is not supported");
+    }
+
+    auto result_vector_type = GetVectorType(op.getType());
+    auto element_type = result_vector_type.getElementType();
+    int64_t iota_size = result_vector_type.getNumElements();
+
+    llvm::SmallVector<mlir::Attribute> iota_values(iota_size);
+    for (int idx = 0; idx != iota_size; ++idx) {
+      iota_values[idx] = rewriter.getIntegerAttr(element_type, idx);
+    }
+
+    mlir::Value iota_const = mlir::arith::ConstantOp::create(
+        rewriter, op->getLoc(),
+        mlir::DenseElementsAttr::get(result_vector_type, iota_values));
+
+    rewriter.replaceOp(op, CastToTensor(rewriter, iota_const));
+    return mlir::success();
+  }
+};
+
 class ShloToVectorPass : public impl::ShloToVectorPassBase<ShloToVectorPass> {
  public:
   using ShloToVectorPassBase::ShloToVectorPassBase;
@@ -325,7 +354,7 @@ class ShloToVectorPass : public impl::ShloToVectorPassBase<ShloToVectorPass> {
     mlir::MLIRContext* context = &getContext();
     mlir::RewritePatternSet patterns(context);
     patterns.add<LowerTranspose, LowerDotGeneral, LowerReduce,
-                 LowerBroadcastInDim, LowerReshape>(context);
+                 LowerBroadcastInDim, LowerReshape, LowerIota>(context);
     if (mlir::failed(
             mlir::applyPatternsGreedily(getOperation(), std::move(patterns)))) {
       signalPassFailure();

xla/backends/cpu/codegen/tiled/transforms/tests/shlo_to_vector.mlir

@@ -185,3 +185,13 @@ func.func @reshape(%input : tensor<4xf32>) -> tensor<2x1x2xf32> {
 // CHECK-LABEL: @reshape
 // CHECK:vector.shape_cast {{.*}} : vector<4xf32> to vector<2x1x2xf32>
 
+// -----
+
+func.func @iota() -> tensor<4xi32> {
+  %result = stablehlo.iota dim = 0 : tensor<4xi32>
+  return %result : tensor<4xi32>
+}
+
+// CHECK-LABEL: @iota
+// CHECK: arith.constant dense<[0, 1, 2, 3]> : vector<4xi32>
+